Skip to content

Commit ffb3069

Browse files
apstenku123apstenku123
committed
feat(nvfp4): wire+verify REAL forward NVFP4 GEMM on gb10 (TE/cuBLASLt), fail-loud broken backward
- Add nvfp4_te_cuda_capability() + nvfp4_te_gemm_probe() to scripts/_nvfp4_route.py: real CUDA NVFP4 GEMM via TransformerEngine NVFP4BlockScaling + cuBLASLt. - VERIFIED on gb10 (sm_121): forward NVFP4 GEMM works, rel_err vs bf16 = 0.147. - Backward NVFP4 GEMM (default RHT+SR recipe) RAISES Nvfp4CudaKernelMiscompiled: installed TE built sm_120 PLAIN (arch 75;80;89;90;100;120), FP4-cvt/RHT PTX is arch-specific; on desktop/consumer Blackwell the fix is family target compute_120f (CUDA13), NOT sm_120a (segfaults; cf. NVIDIA/cutlass#3096). - RULE #1: no silent bf16 fallback; probe defeats TE's async-error masking (fresh module + CUDA_LAUNCH_BLOCKING) so it cannot accept garbage gradients. - Correct all enablement strings sm_120a/121a -> compute_120f. - Add tests/test_nvfp4_route.py (10 tests): Mac 7p/3s, gb10 9p/1s; CUDA fwd matches bf16, CUDA bwd fails loud. - Add docs/NVFP4-TRAINING-KERNELS.md: kernels/versions, gb10 evidence, e2m1 block-scale math, wired vs fail-loud ops, test results, sources.
1 parent 66932dc commit ffb3069

3 files changed

Lines changed: 815 additions & 21 deletions

File tree

‎docs/NVFP4-TRAINING-KERNELS.md‎

Lines changed: 293 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,293 @@
1+
# NVFP4 (e2m1 + block-scale) TRAINING kernels — what exists, what runs on gb10, what fails loud
2+
3+
Status as of 2026-06-01. Author: David Gornshtein (davidgornshtein@gmail.com).
4+
5+
This document is the authoritative, evidence-backed record of the NVFP4
6+
4-bit **training** kernel situation for cppmega.mlx, covering: the production
7+
NVFP4 training recipes/kernels that exist upstream, exactly what is available on
8+
our gb10 (NVIDIA GB10, sm_121) box vs. what is NOT, the e2m1 + block-scale math,
9+
which ops we wired with REAL kernels, which ops fail loud (and how to enable
10+
them), and the test we added with its measured results.
11+
12+
It obeys RULE #1 (no silent fallbacks): every nvfp4 op goes through ONE clear
13+
path; where a kernel is missing or mis-compiled the route RAISES a precise,
14+
actionable error naming WHERE + WHAT failed and the enablement — it never
15+
silently downcasts to bf16.
16+
17+
---
18+
19+
## 1. What NVFP4 is (format + training recipe)
20+
21+
NVFP4 is NVIDIA's native Blackwell 4-bit **training** format. Two-level block
22+
scaling:
23+
24+
* **Elements**: `E2M1` (1 sign, 2 exponent, 1 mantissa). Value set:
25+
`{±0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6}`.
26+
* **Level-1 block scale**: every **16 consecutive elements** share one scale of
27+
type `E4M3` (FP8, 4 exp / 3 mantissa) — NOT power-of-two. (Our repo's existing
28+
Metal codec uses fp32 block scales over the same 16-element block geometry.)
29+
* **Level-2 scale**: a single per-tensor `FP32` scale for the whole tensor.
30+
31+
cuBLASLt exposes this directly as `CUDA_R_4F_E2M1` matmuls with `CUDA_R_UE4M3`
32+
scales over 16-element blocks (cuBLAS 12.9+).
33+
34+
**Training recipe** (NVIDIA TransformerEngine `NVFP4BlockScaling`, verified live
35+
on gb10 — see §3). The three GEMMs of a linear layer use NVFP4 with these
36+
per-operand quant params (exact recipe dump from gb10 TE 2.16.0.dev0):
37+
38+
| Operand cast | RHT | Stochastic rounding | 2D block quant |
39+
|-------------------------|------|---------------------|----------------|
40+
| `fp4_quant_fwd_inp` | True | False | False |
41+
| `fp4_quant_fwd_weight` | False| False | True |
42+
| `fp4_quant_bwd_grad` | True | **True** | False |
43+
44+
* **Random Hadamard Transform (RHT)** is applied to GEMM inputs (fwd input and
45+
bwd grad) to reshape the distribution toward Gaussian and smooth outliers.
46+
* **Stochastic rounding (SR)** is applied when quantizing gradients to NVFP4, to
47+
remove the bias quantization would otherwise introduce.
48+
* **Weights** use 2D (selective) block quantization.
49+
* On NVIDIA datacenter Blackwell, FP4 GEMMs run ~4x BF16 (GB200) / ~6x (GB300).
50+
51+
Sources: §7.
52+
53+
---
54+
55+
## 2. Which NVFP4 training kernels exist upstream (with URLs/versions)
56+
57+
* **NVIDIA TransformerEngine** — `NVFP4BlockScaling` recipe + `te.fp8_autocast`
58+
over `te.Linear`/`te.LayerNormLinear`. Python API verbatim:
59+
`transformer_engine.common.recipe.NVFP4BlockScaling`,
60+
`transformer_engine.pytorch.NVFP4Quantizer`,
61+
`transformer_engine.pytorch.NVFP4Tensor`,
62+
`transformer_engine.pytorch.is_nvfp4_available()`,
63+
`transformer_engine.pytorch.fp8.check_nvfp4_support()`.
64+
C API: `kNVTEFloat4E2M1`, `nvte_quantize()`,
65+
`kNVTEQuantizationConfigStochasticRounding`.
66+
Provides fwd + bwd (dgrad/wgrad) NVFP4 GEMM via cuBLASLt. Version on gb10:
67+
`2.16.0.dev0+a46079cb` (built from source at `/home/dave/TransformerEngine`).
68+
* **cuBLASLt** — native block-scaled FP4 GEMM (`CUDA_R_4F_E2M1` + `CUDA_R_UE4M3`,
69+
16-elem blocks) since cuBLAS 12.9. gb10 has `nvidia-cublas 13.4.0.1`.
70+
* **CUTLASS** — `examples/72_blackwell_narrow_precision_gemm/
71+
72b_blackwell_nvfp4_nvfp4_gemm.cu` block-scaled FP4 GEMM. gb10 has
72+
`nvidia-cutlass-dsl 4.4.2`.
73+
* **FlashInfer** — SM120 NVFP4 MoE grouped GEMM (with the SM120 patches /
74+
`compute_120f`). gb10 has `flashinfer-python 0.6.8`.
75+
* **Community references**: `github.com/vuiseng9/fp4-training` (cuBLASLt +
76+
Microxcaling mxfp8/nvfp4 training, concept→implementation);
77+
NVIDIA's 12B hybrid Mamba-Transformer NVFP4 pretraining (10T tokens).
78+
79+
---
80+
81+
## 3. What is available on gb10 (sm_121) — measured, not assumed
82+
83+
Box: `ssh gb10` → NVIDIA GB10, **compute_cap 12.1 (sm_121)**, driver 595.71.05,
84+
CUDA 13.2. Venv `/home/dave/cppmega-venv`: torch `2.13.0.dev20260417+cu132`
85+
(cuda 13.2), TE `2.16.0.dev0+a46079cb`, cublas `13.4.0.1`, cutlass-dsl `4.4.2`,
86+
flashinfer `0.6.8`.
87+
88+
Capability probe (`scripts/_nvfp4_route.nvfp4_te_cuda_capability()`) on gb10:
89+
90+
```json
91+
{ "cuda_device": "NVIDIA GB10", "compute_capability": "sm_121",
92+
"te_version": "2.16.0.dev0+a46079cb",
93+
"te_check_nvfp4_support": [true, ""], "te_is_nvfp4_available": true }
94+
```
95+
96+
### 3a. FORWARD NVFP4 GEMM — WORKS on gb10 (real, measured)
97+
98+
`te.Linear` under `te.fp8_autocast(NVFP4BlockScaling())` forward = cuBLASLt
99+
`CUDA_R_4F_E2M1` + E4M3 block-scale matmul. Measured forward output rel-RMSE vs
100+
the bf16 reference:
101+
102+
```
103+
forward_nvfp4_gemm_te_cublaslt_cuda_sm121: rel_rmse_vs_bf16 = 0.147, finite=True
104+
```
105+
106+
0.147 is the expected FP4 forward error (no RHT). This is a REAL nvfp4 op on
107+
gb10 and we wire/test it (§5, §6).
108+
109+
### 3b. BACKWARD NVFP4 GEMM — BROKEN on gb10 (root cause identified)
110+
111+
Running the **default production recipe** (RHT on fwd-inp + bwd, SR on bwd)
112+
backward on gb10 raises a hard CUDA error from the Random-Hadamard fused quant
113+
kernel:
114+
115+
```
116+
RuntimeError: .../hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu:1200
117+
in function row_col_rht_gemm_ntt_w_sfc: CUDA Error: invalid argument
118+
```
119+
120+
With RHT disabled, the backward stochastic-rounding FP4 cast asserts (one per
121+
thread):
122+
123+
```
124+
.../util/ptx.cuh:935 in function mul_cvt_bf16_to_fp4_8x_stochastic_rounding:
125+
FP4 cvt PTX instructions are architecture-specific.
126+
Try recompiling with sm_XXXa instead of sm_XXX.
127+
```
128+
129+
**Root cause (proven via cuobjdump on the TE .so):** the installed
130+
`libtransformer_engine.so` embeds SASS for:
131+
132+
```
133+
sm_75, sm_80, sm_89, sm_90, sm_90a, sm_100, sm_100a, sm_103a, sm_120 (PLAIN)
134+
```
135+
136+
i.e. TE's default CUDA-13 arch list `75;80;89;90;100;120` — `sm_120` is built
137+
**PLAIN**, and the architecture-specific `a` variants exist only for datacenter
138+
Blackwell (`sm_100a`, `sm_103a`). GB10 is `sm_121`, runs the `sm_120` plain
139+
SASS, which lacks the arch-specific FP4 cvt / RHT instructions → the asserts
140+
above.
141+
142+
**The `a` variant is NOT the fix on desktop/consumer Blackwell.** Per
143+
NVIDIA/cutlass#3096 and flashinfer-ai/flashinfer#2723, forcing `compute_120a`
144+
gencode SEGFAULTS on desktop Blackwell ("RTX PRO 6000 reports compute capability
145+
12.0, not 12.0a; the `a`-specific instructions are not available on desktop
146+
Blackwell"). The working target is the **family-specific `compute_120f`** added
147+
in CUDA 13.0, which "enables the full SM120 feature set with working TMA WS
148+
grouped-GEMM tactics" (compute_120f took a CUTLASS NVFP4 MoE from 14.6 → 39.0
149+
tok/s vs compute_120a). So the enablement for gb10 backward is to rebuild TE/
150+
cuBLASLt-path FP4 kernels under CUDA ≥ 13 targeting `120f`, once TE's build
151+
emits the `f` family target.
152+
153+
**Danger note (why we fail loud):** the backward FP4 error is ASYNC. Without
154+
`CUDA_LAUNCH_BLOCKING=1` AND a fresh `te.Linear`, TE quantizer-workspace state
155+
can mask the failure and hand back gradients that *look* plausible — a silent
156+
mis-execution. Our probe defeats this by (a) using a fresh module for the
157+
default-recipe backward and (b) treating the arch-specific assertion as a hard
158+
RAISE. We never accept the masked gradients.
159+
160+
### 3c. MXFP8 — also unavailable on gb10 (for completeness)
161+
162+
`check_mxfp8_support()` → `(False, 'MXFP8 (for all gemm layouts) is not supported
163+
on 12.0+ architectures yet.')`. Not used by the nvfp4 route; recorded so the
164+
arch picture is complete.
165+
166+
---
167+
168+
## 4. The e2m1 + block-scale math (as implemented)
169+
170+
For a vector `v` split into 16-element blocks `b_j`:
171+
172+
1. `amax_j = max(|b_j|)`; level-1 scale `s_j = amax_j / E2M1_MAX` (E2M1_MAX = 6),
173+
stored as `E4M3` (TE) or `fp32` (repo Metal codec).
174+
2. Each element `e = round(b_j[i] / s_j)` snapped to the nearest E2M1 value;
175+
gradients additionally use **stochastic rounding** (round up/down with
176+
probability proportional to distance between the two representable values).
177+
3. A per-tensor `FP32` level-2 scale captures global dynamic range.
178+
4. Optional **RHT** is applied to fwd-input and grad tensors before step 1 to
179+
Gaussianize the distribution.
180+
5. Dequant: `b_j[i] ≈ s_j * e * level2_scale`.
181+
182+
GEMM: NVFP4(A) × NVFP4(B) accumulates in FP32 on the tensor cores (cuBLASLt
183+
`CUDA_R_4F_E2M1`), de-scaling per-block by the E4M3 scales.
184+
185+
---
186+
187+
## 5. What we WIRED (real kernels) vs FAIL-LOUD (and how to enable)
188+
189+
Route module: `scripts/_nvfp4_route.py`. Wiring into the training step:
190+
`scripts/m04_train_step.py` (imports + RULE #1 gate at `_run_existing_training`
191+
and a defensive guard at `run_local_gb10_quarter_training`). `--dtype nvfp4` is
192+
an accepted route (no argparse rejection); the dry-run path emits the route
193+
metadata; the real training step RAISES the precise blocker.
194+
195+
**REAL nvfp4 ops wired (`NVFP4_SUPPORTED_OPS`):**
196+
197+
| Op | Backend | Kernel | Status |
198+
|----|---------|--------|--------|
199+
| `forward_gemm_operands_e2m1_block_scale_metal_m4` | Metal/M4 | `cppmega_mlx.nn._tilelang.mxfp4_matmul_path_c` over `quantize_mxfp4_blockwise` | real (M4); exercised by `nvfp4_gemm_smoke` |
200+
| `forward_nvfp4_gemm_te_cublaslt_cuda_sm121` | CUDA/gb10 | TE `NVFP4BlockScaling` + cuBLASLt FP4 | **VERIFIED on gb10, rel_err 0.147**; exercised by `nvfp4_te_gemm_probe(run_backward=False)` |
201+
202+
**FAIL-LOUD ops (`NVFP4_UNSUPPORTED_TRAINING_OPS`)** — the route RAISES naming
203+
each; it does NOT bf16-fallback:
204+
205+
| Op | Why | Enablement |
206+
|----|-----|-----------|
207+
| `nvfp4_gemm_backward_vjp_te_cuda_sm121` | TE FP4-cvt/RHT PTX is arch-specific; TE built `sm_120` plain (not the `f` family) → asserts / CUDA invalid-argument on gb10 | Rebuild TE under CUDA ≥ 13 with family target `compute_120f` (NOT `sm_120a`, which SEGFAULTS on desktop Blackwell; cf. cutlass#3096) |
208+
| `cuda_blackwell_nvfp4_gemm_metal_only` | The Metal fwd GEMM is Metal-only; it RAISES on CUDA | Use the TE/cuBLASLt CUDA forward path above |
209+
| `gemm_backward_vjp` | No nvfp4 grad kernel for the Metal e2m1 GEMM | Implement an nvfp4 VJP for `mxfp4_matmul_path_c` |
210+
| `optimizer_state_and_update` | AdamW moments are fp32-only; no nvfp4 optimizer state | Implement nvfp4 optimizer-state kernels |
211+
| `rmsnorm`, `swiglu_ffn`, `softmax`, `sparse_mla_attention`, `mamba3_selective_scan`, `m2rnn_recurrence`, `embedding_and_lm_head`, `residual_add` | No nvfp4 kernel for these graph ops | Implement per-op nvfp4 kernels |
212+
213+
Because a full HybridTinyLM step needs ALL of the above, `--dtype nvfp4`
214+
delivers the honest partial: real forward nvfp4 GEMM, fail-loud for the rest.
215+
For a full training step TODAY use `--dtype fp8_path_c` or `--dtype bfloat16`.
216+
217+
**Fail-loud guards (RULE #1):**
218+
219+
* `raise_if_nvfp4_training_unsupported(args)` — on the training critical path,
220+
RAISES `Nvfp4TrainingRouteUnavailable` before any op would run in bf16.
221+
* `_run_existing_training` emits a `blocked_receipt(... NVFP4_E2E_TRAINING_
222+
BLOCKER_TYPE)` with the precise reason.
223+
* `nvfp4_te_gemm_probe(run_backward=True)` RAISES
224+
`Nvfp4CudaKernelMiscompiled` on the broken gb10 backward (arch-specific PTX),
225+
rather than returning the masked/garbage gradients TE would otherwise hand
226+
back. No silent fallback exists anywhere in the route.
227+
228+
---
229+
230+
## 6. The test we added and its results
231+
232+
Test: `tests/test_nvfp4_route.py` (10 tests). Backend-aware: always-on metadata
233+
+ fail-loud assertions; Metal-only assertions skip when MLX/Metal absent; CUDA
234+
assertions run when torch+TE+CUDA+NVFP4 present.
235+
236+
* `test_route_constants_well_formed`, `test_route_requested_predicate`,
237+
`test_payload_advertises_both_forward_backends_and_no_e2e`,
238+
`test_reason_message_is_actionable`,
239+
`test_capability_probe_never_crashes_without_torch` — route metadata + probe.
240+
* `test_fail_loud_guard_raises_for_nvfp4_and_names_missing_ops`,
241+
`test_fail_loud_guard_is_noop_for_other_dtypes` — RULE #1 gate.
242+
* `test_metal_forward_nvfp4_gemm_matches_bf16` — real Metal fwd GEMM < 0.5 rel.
243+
* `test_cuda_forward_nvfp4_gemm_matches_bf16` — **real gb10 NVFP4 fwd GEMM
244+
< 0.3 rel** (measured 0.147).
245+
* `test_cuda_backward_nvfp4_gemm_fails_loud_on_miscompiled_te` — the broken gb10
246+
backward RAISES `Nvfp4CudaKernelMiscompiled` naming the `compute_120f`
247+
enablement. (Skips, expecting pass, when `NVFP4_BACKWARD_FIXED=1` once TE is
248+
rebuilt with `120f`.)
249+
250+
**Results:**
251+
252+
* Mac (this checkout): `7 passed, 3 skipped` — Metal skipped because the
253+
in-tree tilelang dev build (`merge/upstream-codegen-reorg`) has a tvm_ffi
254+
circular-import (pre-existing infra gap, unrelated to nvfp4); CUDA skipped (no
255+
CUDA on Mac). The Metal assertion is gated by a real tiny-GEMM smoke so it
256+
skips cleanly on a broken-tilelang host and runs where tilelang is healthy.
257+
* gb10 (`/home/dave/cppmega-venv`): **`9 passed, 1 skipped`** — Metal skipped (no
258+
Apple GPU); the CUDA forward test PASSES (real nvfp4 GEMM, 0.147 rel) and the
259+
CUDA backward fail-loud test PASSES (broken backward RAISES as designed).
260+
261+
Reproduce on gb10:
262+
263+
```bash
264+
# copy scripts/_nvfp4_route.py + tests/test_nvfp4_route.py to a dir with a
265+
# scripts/ package, then:
266+
cd <dir> && /home/dave/cppmega-venv/bin/python -m pytest tests/test_nvfp4_route.py -v
267+
```
268+
269+
---
270+
271+
## 7. Sources
272+
273+
* TransformerEngine NVFP4 feature doc —
274+
https://nvidia.github.io/TransformerEngine/features/low_precision_training/nvfp4/nvfp4.html
275+
* TransformerEngine Common API (NVFP4BlockScaling, kNVTEFloat4E2M1) —
276+
https://nvidia.github.io/TransformerEngine/api/common.html
277+
* "NVFP4 Trains with Precision of 16-Bit and Speed/Efficiency of 4-Bit" (NVIDIA
278+
blog) —
279+
https://developer.nvidia.com/blog/nvfp4-trains-with-precision-of-16-bit-and-speed-and-efficiency-of-4-bit/
280+
* NVIDIA 4-bit pretraining (12B hybrid Mamba-Transformer, 10T tokens) —
281+
https://www.marktechpost.com/2026/05/18/nvidia-introduces-a-4-bit-pretraining-methodology-using-nvfp4-validated-on-a-12b-hybrid-mamba-transformer-at-10t-token-horizon/
282+
* TransformerEngine repo —
283+
https://github.com/NVIDIA/TransformerEngine
284+
* cuBLAS 12.9 block-scaled FP4 (CUDA_R_4F_E2M1 + CUDA_R_UE4M3, 16-elem blocks) —
285+
https://developer.nvidia.com/blog/boosting-matrix-multiplication-speed-and-flexibility-with-nvidia-cublas-12-9/
286+
* CUTLASS SM120 NVFP4: sm_120a SEGFAULT, compute_120f fix (CUDA 13.0) —
287+
https://github.com/NVIDIA/cutlass/issues/3096
288+
* FlashInfer SM120 NVFP4 grouped GEMM patching —
289+
https://github.com/flashinfer-ai/flashinfer/issues/2723
290+
* CUTLASS NVFP4 GEMM example —
291+
https://github.com/NVIDIA/cutlass/blob/main/examples/72_blackwell_narrow_precision_gemm/72b_blackwell_nvfp4_nvfp4_gemm.cu
292+
* Community nvfp4 training (cuBLASLt + Microxcaling) —
293+
https://github.com/vuiseng9/fp4-training

0 commit comments

Comments
 (0)