Skip to content

feat(fp8): H20/Hopper A8W8 FP8 decode via torch._scaled_mm (weight quantization) - #526

Draft
CAICAIIs wants to merge 10 commits into
inclusionAI:mainfrom
CAICAIIs:perf/e4m3-cuda-kernel
Draft

feat(fp8): H20/Hopper A8W8 FP8 decode via torch._scaled_mm (weight quantization)#526
CAICAIIs wants to merge 10 commits into
inclusionAI:mainfrom
CAICAIIs:perf/e4m3-cuda-kernel

Conversation

@CAICAIIs

@CAICAIIs CAICAIIs commented Aug 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

H20/Hopper (cc 9.0) FP8 decode via torch._scaled_mm. Decode re-reads the full weight tensor from HBM every generated token, so it is memory-bandwidth-bound; storing the weight as 1-byte FP8 (E4M3) halves the weight-read bytes.

  • A8W8 torch._scaled_mm is the decode path on Hopper (cuBLASLt rejects a bf16 activation, so the activation is also quantized to E4M3 — a W8A8 decode). Measured on H20: ~2x per linear, ~1.57x per Qwen3-8B MLP block (CUDA graphs).
  • A fused Triton activation-quantize kernel (amax + scale + cast in one) keeps the per-call quantize cheap; without it the wired path dropped to ~1.2x.
  • H20-only: fp8-marked weights fall through to bf16 elsewhere (no Triton fallback).

quant_method stays opt-in ("none" default); unmarked weights keep the existing bf16 path. Decode-only — FP8 has no backward.

Scope change vs. the earlier draft

The A100 E4M3 fused dequant-linear CUDA kernel (e4m3_linear.cu, e4m3_cuda.py, extension registration, its bench, its test), the RFC doc, the Triton W8A16 kernel (measured 0.5x, never a speedup), and redundant benchmarks were removed. They added scope without a speedup.

Changes

  • areno/accel/kernels/fp8_linear.py — quantized_fp8_scaled_mm (A8W8 via torch._scaled_mm, weight as 1-byte E4M3 in a w.t() column-major view), a fused _quantize_act Triton kernel, mark_fp8_weight (E4M3, per-tensor scale), scaled_mm_available gate.
  • areno/engine/layers/linear.py — _areno_linear_forward routes FP8-marked weights to _scaled_mm on Hopper, else the bf16 path.
  • areno/engine/quantization.py — FP8 (E4M3) scale/quantize/dequant CPU reference.
  • areno/engine/{config,modeling,worker}.py — opt-in quant_method, quantize_model_weights_fp8, train-worker guard (forward-only).
  • Tests: tests/test_fp8_quant_cpu.py.
  • Benchmark: scripts/bench/h20_fp8_scaled_mm.py (CUDA-graph timed).

Testing (measured on H20, torch 2.9.1+cu128, cc 9.0)

  • CPU: pytest tests/test_fp8_quant_cpu.py → 5 passed.
  • Single linear (8B shape, M=1, N=12288, K=4096): 2.05–2.33x with the fused quantize, rel ≈ 3.8% vs bf16.
  • Qwen3-8B MLP decode (CUDA graph, 1–2 layers): ~1.57x, rel 5.9%.
  • Raw _scaled_mm matmul (CUDA graph): 2.18x @ M=1, 2.56x @ N=8192.

Known limitations (honest)

  • Decode-only (no backward; guarded in worker.py).
  • H20 is the target; off-Hopper fp8-marked weights use bf16 (no speedup).
  • Accuracy: per-tensor A8W8 compounds across layers — rel ≈ 16% at 2 layers on the Qwen3-8B MLP proxy. A real multi-layer model needs per-channel/block scaling (torch._scaled_mm RowWise is available on CUDA 12.8) before it is usable.
  • Full 48-layer Qwen3-8B end-to-end and the torch.compile interaction with weight._areno_fp8 are not yet validated (no free GPU on the benchmark host).

Related

@CAICAIIs
CAICAIIs force-pushed the perf/e4m3-cuda-kernel branch 3 times, most recently from 1cebd8f to 5d8b944 Compare August 26, 2026 17:53
@CAICAIIs
feat(cuda): E4M3 fused dequant-GEMM kernel for A100
837437c 
Add a forward-only WMMA (16x16x16 fp16 tensor-core) kernel that reads a
1-byte uint8 E4M3 weight payload, decodes each byte to fp16 in shared
memory (fast bit-decode), and applies a per-tensor scalar scale after the
dot -- the memory-bandwidth benefit from RFC 0001 for decode/inference.

- csrc/e4m3_linear.cu: fused dequant-GEMM + registration.
- kernels/e4m3_cuda.py: CUDA shim (quantized_e4m3_linear_cuda) + quantize/dequant helpers.
- tests/test_e4m3_decode_cpu.py: CPU decode reference vs torch.float8_e4m3fn.
- setup.py: compile the new .cu.

Forward-only (no backward): E4M3 must not be wired into a training graph.
Opt-in and backward-compatible; the model hook is not yet switched to this path.
@CAICAIIs
CAICAIIs force-pushed the perf/e4m3-cuda-kernel branch from 5d8b944 to 837437c Compare August 27, 2026 01:45
review-cleanup
refactor(fp8): clean up PR 526 E4M3/FP8 quantization slice
052764e 
- narrow quant_method Literal to implemented values; drop dead quant_group_size
- dedupe FP8 scale math: mark_fp8_weight now reuses compute_fp8_scale(max_val)
- drop unused group_size param on mark_fp8_weight (per-tensor only today)
- worker: normalize ARENO_QUANT_FP8 truthiness; correct 'inference role' wording
- e4m3_cuda: lazy-import the CUDA extension so the CPU decode test doesn't need it
- trim stale measured-GB/s comments out of the .cu

No behavior change to the FP8 value semantics; scale/rounding unchanged.
@CAICAIIs
CAICAIIs marked this pull request as draft August 31, 2026 10:12
review-cleanup added 7 commits August 31, 2026 18:27
review-cleanup
narrow(fp8): drop A100 E4M3 kernel and RFC; wire torch._scaled_mm as …
863e537 
…the Hopper decode path

Per review: the E4M3 A100 custom kernel was standalone (never wired into the
model) and the RFC drove scope creep. This narrows the PR to a single FP8 W8A16
decode feature:

- remove areno/accel/csrc/e4m3_linear.cu + its extension registration/setup entry
- remove areno/accel/kernels/e4m3_cuda.py, scripts/bench/e4m3_cuda_bench.py,
  tests/test_e4m3_decode_cpu.py and docs/rfcs/0001-weight-quantization-fp8-int4.md
- fp8_linear: mark_fp8_weight now stores the E4M3 grid (the finer reference);
  gate torch._scaled_mm on cc >= 8.9 (quantized_fp8_scaled_mm); A100 falls back
  to the Triton kernel, which converts E4M3->E5M2 in-kernel
- _areno_linear_forward dispatches to _scaled_mm on Hopper, Triton otherwise

NOTE: the exact torch._scaled_mm W8A16 call (transpose + scale_b, no scale_a)
still needs verification on a Hopper/H20 host before it is correct.
review-cleanup
fix(fp8): keep W8A16 Triton wired; use torch._scaled_mm as H20 A8W8 b…
dc14f41 
…ench only

On H20 (torch 2.9.1+cu128, cc 9.0) torch._scaled_mm rejects a bf16 activation
("a and b should be float8") — it is A8W8 only, so the requested W8A16-via-_
scaled_mm is impossible. Per review, narrow to weight-only W8A16:

- _areno_linear_forward always dispatches to the Triton W8A16 kernel; drop the
  scaled_mm_available() gate and quantized_fp8_scaled_mm() (now dead)
- mark_fp8_weight stays on the E4M3 grid (finer); Triton converts to E5M2 where
  Triton rejects E4M3 (Ampere)
- h20_fp8_scaled_mm.py: fix mat2 layout — wq.t() (column-major view), NOT
  .contiguous() (cuBLASLt rejects the contiguous form)
- docs: W8A16/Triton is the wired path; torch._scaled_mm (A8W8) is a bench

Verified on H20: 6/6 CPU FP8 tests pass; _scaled_mm A8W8 correct (rel~3.7%) and
~1.2-1.4x (below RFC's >1.5x on this torch); W8A16 Triton kernel + hook run
correctly (rel~7%, finite).
review-cleanup
bench(h20): measure _scaled_mm decode throughput (amortized), not per…
c03c2b7 
…-call latency

The per-call-sync timing (torch.cuda.synchronize after every call) collapsed
the M=1 ratio to ~1.39x because the barrier + launch overhead dominates a
~24us kernel. Decode tokens/s is amortized throughput (kernels pipeline), so
time reps back-to-back calls and add clock warmup. On H20 (torch 2.9.1+cu128)
this gives 1.53x @ M=1, 1.64x @ M=4, 1.88x @ M=64 (N=12288,K=4096), matching
the RFC's >1.5x claim (rel ~3.7% vs bf16).
review-cleanup
bench(h20): time via CUDA graph (real decode path); small shapes now …
4cd87cc 
…>1.5x

Decode runs under CUDA graphs, so capture the op once and time replay. This
removes the per-launch/setup overhead that hid the FP8 kernel's true speed for
small shapes. On H20 (torch 2.9.1+cu128):
  M=1 N=8192:  1.09x(amortized-latency) -> 2.56x
  M=1 N=12288: 1.53x -> 2.18x
  M=4 N=12288: 1.64x -> 2.48x
  M=64 N=12288:1.88x -> 1.87x
All correct (rel ~3.7% vs bf16). This also answers the small-shape question:
the low small-N number was launch overhead, not the kernel.
review-cleanup
fix(fp8): wire A8W8 torch._scaled_mm on Hopper; Triton W8A16 is a reg…
4684992 
…ression

Measured on H20 that the wired W8A16 Triton kernel is SLOWER than bf16 at every
M (0.10-0.32x) — a regression, and ~2x slower even with E5M2 (no per-call
conversion) — while A8W8 torch._scaled_mm is a real speedup and more accurate
(3.7-3.9% vs Triton's ~7%). So dispatch fp8 to _scaled_mm on Hopper (the actual
decode fast path, ~2x at the matmul level, ~1.2x for the wired hook that also
quantizes the activation) and keep the Triton kernel only as the non-Hopper
fallback. Rewired after the earlier decision because the data shows the
weight-only Triton path cannot deliver the decode speedup.
review-cleanup
perf(fp8): fuse activation quantize into one Triton kernel for the A8…
1383beb 
…W8 decode

The per-call activation-quantize chain (amax + where + div + clamp + cast, ~5
torch ops) choked the wired A8W8 path to ~1.2x under CUDA graphs. Fuse the amax +
scale + quantize into a single Triton kernel (single block covers the tiny decode
activation; larger/prefill tensors fall back to torch). On H20 at the 8B MLP
shape (M=1,N=12288,K=4096) the wired hook now hits ~2.2x (was 1.2x), rel ~3.8%.
Note: on a small model (Qwen3-0.6B, H=1024/I=3072) the same path is ~0.93x — the
per-linear overhead exceeds the byte savings unless the model is memory-bound.
review-cleanup
docs(fp8): make H20/Hopper the explicit decode target
ba52719 
The PR target is H20 (Hopper, cc 9.0): the A8W8 torch._scaled_mm decode path.
Measured on H20: Qwen3-8B MLP decode ~1.57x end-to-end per MLP block (2x per
linear) under CUDA graphs; non-Hopper uses the Triton kernel as a fallback
(correctness, not a speedup).
@CAICAIIs CAICAIIs changed the title feat(fp8): E4M3 fused dequant-linear kernel + FP8 quantized-linear slice (RFC 0001) feat(fp8): H20/Hopper A8W8 FP8 decode via torch._scaled_mm (weight quantization) Sep 2, 2026
review-cleanup
refactor(fp8): drop over-engineering; H20/Hopper only
e14e436 
Review trim (no behavior change to the A8W8 path):
- remove the Triton W8A16 kernel + quantized_fp8_linear (never a speedup;
  0.5x even as E5M2) and the non-Hopper fallback in _areno_linear_forward —
  fp8-marked weights now fall through to bf16 off-Hopper
- remove QuantizedLinear (test-only W8A16 reference; redundant with the direct
  quantize/dequant tests and inconsistent with the A8W8 runtime)
- remove dead dequantize_fp8_weight
- remove fp8_end_to_end_bench.py + triton_fp8_matmul_bench.py (benchmarked the
  removed Triton path / A100 study)
- keep h20_fp8_scaled_mm.py + the CPU reference tests (now 5/5)

Verified on H20: CPU 5/5; trimmed hook correct (rel 3.85%) and 2.20x under CUDA
graphs.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@CAICAIIs