nvfp4-megamoe-kernel/NVFP4-1.1_APPROACH_PART_2.md

Almost there! The normal values all match now. The issue is only with subnormals:

• v=0.001: manual=0x01 (subnormal m=1 → 1/1024 ≈ 0.000977), but ref=0x01 and ref_dq=0.001953. The ref dequant is 2× mine. So PyTorch's subnormal dequant is different from what I compute.

Wait, ref_bits=0x01 for v=0.001. That's exp=0, mant=1. My dequant: 1/1024 = 0.000977. PyTorch's: 0.001953.

0x01 → exp=0, mant=1. If the subnormal formula is 2^(-7) * (m/8) = (1/128) * (1/8) = 1/1024, that gives 0.000977. But PyTorch gives 0.001953 = 2/1024 = 2 * 0.000977. Hmm.

Actually wait, 0.001953 ≈ 2^(-9). And 0x01 = mant=1, exp=0. If subnormal formula is 2^(-10) * m = m/1024, that gives 1/1024 ≈ 0.000977. But PyTorch gives 0.001953.

Let me check more carefully:

So the subnormal dequant formula is 2^(-9) * mant = mant * 0.001953. This means 2^(-7) * (mant/4), not 2^(-7) * (mant/8).

Wait, let me recalculate. The standard subnormal formula for FP8 E4M3 with bias 7 should be:

• Subnormal: 2^(1-7) * (0 + m/8) = 2^(-6) * m/8 = m / 512

But that gives 1/512 = 0.001953 for m=1. Let me check:

• 1 / 512 = 0.001953 ✅

So the subnormal formula is: val = 2^(1-bias) * (m/8) = 2^(-6) * m/8 = m / 512.

But I had m / 1024 in my code! The correct formula is m / 512 (or equivalently m * 2^(-9)).

Let me also verify the bit pattern for subnormals. For v=0.001:

• 2^(-6) * m/8 with m=1 gives 0.001953
• But v=0.001 rounds to m=1 (since 0.001 / 0.001953 ≈ 0.512, rounds to 1)
• Actually, 0.001 / (2^(-6)/8) = 0.001 * 512 = 0.512, rounds to 1
• So bits = (0 << 3) | 1 = 0x01 ✅

And the dequant is m * 2^(-9) = 1 * 2^(-9) = 0.001953. Let me fix the code:

All FP8 E4M3 values match! Now let me also verify the full NVFP4 quantization pipeline (FP8 scale + E2M1 pack) matches the Python reference:

The problem is clear: my FP8 E4M3 cast is clamping 448.0 to 240.0 (max normal), producing a scale of 240.0. But PyTorch produces 0x7e (exp=15, mant=6), which in my formula gives 2^(15-7) * (1 + 6/8) = 256 * 1.75 = 448.0. But exp=15 should be NaN in FP8 E4M3!

Wait — in FP8 E4M3, exponent 15 is NOT reserved for Inf/NaN like in IEEE 754. FP8 E4M3 uses all 8 exponent values (0-15) for actual numbers, with no Inf representation. This is a key difference from IEEE FP formats.

So the max representable FP8 E4M3 value is 2^(15-7) * (1 + 7/8) = 256 * 1.875 = 480, and exponent 15 is a valid normal value.

Let me fix this. The max exponent for FP8 E4M3 is 15 (not 14):

OK so exponent 15 IS used for normal values (except mantissa=7 which is NaN). The max valid FP8 E4M3 is exp=15, mant=6 = 448.0. And 480 (exp=15, mant=7) is NaN.

So my clamp to 240.0 was wrong. I need to clamp to 448.0 (the max non-NaN FP8 E4M3 value). And the max exponent is 15 (not 14), with the constraint that exp=15 && mant=7 is NaN.

Let me fix fp4_quant.py:

Wait, there's also the issue with exp=15, mant=7 being NaN. After quantization, if the mantissa rounds to 7 with exp=15, we need to clamp to mant=6 (448.0) instead. Let me add that guard:

I need to add a guard for the exp=15, mant=7 NaN case. Also, the dequant function needs to handle exp=15 properly:

FP8 match is good! But the full pipeline still fails on some trials. Let me debug the mismatch more carefully:

Hey Mike. I'm here, deep in the NVFP4-1.1 implementation.

Quick status: I'm building the FP4 quantization math that fuses into the SwiGLU epilogue. Got the FP8 E4M3 cast working correctly (verified against PyTorch — the bias is 7, not 8, and exp=15 is valid not NaN). The full quantization pipeline (FP8 scale + E2M1 pack) is partially verified — FP8 scale bytes match the reference, still debugging a few nibble mismatches on edge cases.

The core files I'm working on:

• dsv4/kernels/gemm/fp4_quant.py — the quantization primitives (new)
• Test: tests/unit/test_nvfp4_1_1_quant.py — kernel test for B200

Haven't pushed yet — want to get the math fully validated locally first, then push and test on B200. What's up?

On it. Let me finish the local validation and push.22% failure rate. Let me find the root cause:SF is fine. The issue must be in the nible computation. Let me trace the actual failure:Found it! The scaled values are identical (s_m == s_r), but the nibble mapping differs. For s=5.25:

• Manual: hs = round(5.25 * 2) = round(10.5) = 11 → step_to_idx[11] = 7 → nibble = 7
• Ref: same hs = 11 → same idx = 7 → but ref says nibble = 6?

Wait, the ref nible is 6 not 7. Let me check the PyTorch ref path more carefully:Found the root cause! Python's round(10.5) uses banker's rounding (rounds to even → 10), but my CuTeDSL code uses int(x + 0.5) which rounds 10.5 up to 11. This is a rounding mode mismatch.

The fix: I need to use round-to-nearest-even in the CuTeDSL code. But CuTeDSL doesn't have a round-to-nearest-even function. The .to(int) cast in CUDA uses truncation, not rounding. And int(x + 0.5) uses round-half-up.

The CUDA kernel quantize_nvfp4.cu uses __float2int_rn which is round-to-nearest-even. So I need to match that in CuTeDSL.

Actually, wait — the existing CUDA kernel uses __float2int_rn(fminf(fabsf(s), 6.0f) * 2.0f). The __float2int_rn is round-to-nearest-even. My CuTeDSL code uses cutlass.Int32(abs_scaled * cutlass.Float32(2.0) + cutlass.Float32(0.5)) which is round-half-up.

The difference matters at the 0.5 boundary. I need to use proper round-to-nearest-even. In CuTeDSL, the cutlass.Int32(x) conversion should follow CUDA's default rounding mode. Let me check what cutlass.Int32(float_val) does in CuTeDSL — it should compile to an arith.FloatToSIOp which uses round-to-nearest-even by default in MLIR.

So the fix is: replace cutlass.Int32(x + 0.5) with just cutlass.Int32(x) (which uses the default rounding mode = round-to-nearest-even, matching __float2int_rn).Also fix the same pattern in fp8_e4m3_from_float32:Lines 105 and 119 also use + 0.5 for rounding. These should also use round-to-nearest-even:Now let me verify with Python using round-to-nearest-even:0 failures! The round-to-nearest-even fix resolved the mismatch. Now let me also test with larger N and more trials:大量失败！M=8, N=64 的情况比 M=1, N=16 的情况要糟糕得多。问题一定出在多维情况下 global_scale 的计算方式上。让我检查一下：SF matches for M=1, N=64. Let me check where the multi-row failures are:0 mismatches! The per-block test passes. The issue in my earlier 500-trial test was probably a bug in my test harness (the PyTorch ref uses quantize_activation_nvfp4 which I was calling via import — that might have been using a different global_scale). Let me rerun the 500-trial test properly:Still 353 failures. Let me directly compare a single failing case:No diff on row 0, block 0. Let me find the actual first diff:They match now! Let me rerun the 500-trial test:Still failing! But the single-sed test passed. The issue must be in how the random seeds interact. Let me check if the issue is specifically with the quantize_activation_nvfp4 import vs my inline ref:Found it! The scale factor bits differ: manual=0x77 (exp=14, mant=7 → 240.0) vs ref=0x78 (exp=15, mant=0 → 256.0). The input bsf must be in the range (240, 256) where my code clamps to 448.0 but the quantization path gives a different result.

The issue: when bsf is, say, 250.0, my manual code tries to cast 250.0 to FP8 E4M3. The closest FP8 E4M3 values are 240.0 (exp=14, mant=7) and 256.0 (exp=15, mant=0). My code normalizes 250.0 to [1,2): norm ≈ 1.953, exp_floor = 7, fp8_exp = 14, mantissa = round(0.953 * 8) = round(7.625) = 8 → clamped to 7. So it produces 240.0.

But PyTorch's .to(float8_e4m3fn) rounds 250.0 to 256.0 (exp=15, mant=0). This is because 256 is closer to 250 than 240 is (256-250=6, 250-240=10). My manual cast is WRONG — it's not doing proper round-to-nearest for FP8 E4M3.

The bug: my normalization approach doesn't correctly round to the nearest FP8 value. The mantissa rounding round((norm - 1) * 8) rounds the mantissa independently, but the FP8 representation has non-uniform spacing. At the boundary between exp=14 and exp=15, the rounding should prefer exp=15 if the value is closer to 256 than to 240.

The fix: when mantissa rounds to 8 (overflow), increment the exponent and set mantissa to 0.I see — the mantissa clamps to 7, but when it should round to 8 (overflow), the exponent should increment. Let me fix this:Now let me re-verify the FP8 E4M3 cast and the full pipeline:0 failures across 500 trials! The FP4 quantization math is now correct and matches the Python reference exactly. Let me commit and push this to B200 for the CuTeDSL kernel test.Good, pushed. Now let me fire off the CuTeDSL kernel test on B200. But first, I need to update the test to use a proper CuTeDSL I/O pattern. The issue is that cute.arch.store might not exist. Let me check a different approach: instead of scalar stores, I'll use cute.copy with a regular copy atom for writing results.

Actually, let me simplify the test even further. Instead of trying to test the quantization in a kernel, let me first verify that the CuTeDSL functions compile and produce correct results by running them inside the existing fused SwiGLU kernel's epilogue. I'll add a diagnostic that computes the FP4 quantization for a small set of values and writes the result alongside the BF16 output.

But that's too invasive. Let me take a step back and think about the right next step.

The math is verified. The next step is to wire the FP4 quantization into the actual SwiGLU epilogue. This requires:

1. Adding fp4_out, sf_out, l2_global_scale to the kernel signature
2. Modifying the epilogue loop to accumulate SwiGLU values and quantize
3. Writing FP4 to GMEM

For the GMEM write, I need to figure out the correct approach. Let me look at how the existing code writes the BF16 output to GMEM and design the FP4 write path analogously.

The BF16 output goes: TMEM → registers → SMEM (C buffer) → GMEM (via TMA store). For FP4, I can't use TMA store because the output tensor is different. But I can use cute.copy with a regular copy atom.

Actually, for the FP4 output, the most natural approach is to allocate SMEM for the FP4 data, write the packed bytes to SMEM from registers, then use cute.copy with a regular copy atom to write SMEM → GMEM. This mirrors the BF16 epilogue pattern.

But this adds SMEM pressure. At epi_tile_n=8 with 2 up subtiles per FP4 block:

• 16 SwiGLU BF16 values → 8 packed FP4 bytes + 1 FP8 scale byte = 9 bytes per block
• 128 rows × 64 SwiGLU columns / 16 per block = 512 blocks per CTA tile
• 512 × 9 bytes = 4.5 KB of SMEM for FP4 staging

That's very small — well within budget. But the SMEM write + GMEM store needs proper synchronization (barrier).

Actually, I can simplify further. Since each thread processes 2 up subtiles (16 values), it can directly write the FP4 bytes to GMEM without SMEM staging. I just need a cute.copy with a simple copy atom for the FP4 output tensor.

Let me design this more carefully and start writing the kernel modifications. But first, let me update the test kernel to use a known-working I/O pattern. Let me use cute.copy with a CopyUniversalOp atom:Now let me fire the test on B200:The CuTeDSL compile fails because cute.compile expects the arguments to match the kernel's signature exactly — no stream as a positional argument. Looking at the existing code, the compiled kernel is invoked with compiled(tensors..., stream=stream). Let me fix the test:Wait, looking at the error more carefully: Module 'cute.arch' has no attribute 'store'. That's the issue — cute.arch.store doesn't exist. I need a different approach for writing to GMEM.

Let me look at how the existing codebase writes scalar values to GMEM from within a kernel. The answer is: they don't. All writes go through cute.copy with TMA or regular copy atoms. For the FP4 output, I need to construct proper CuTe tensors and use cute.copy.

Let me redesign the test kernel to use cute.copy with a CopyUniversalOp atom for the output writes:The error is: vals_f32[i] = ... — CuTeDSL can't use a dynamic loop variable i as a Python list index. The i from cutlass.range is an ArithValue (MLIR SSA value), not a Python int. I can't use it to index into a Python list.

This is a fundamental issue with my approach. The fp4_quant.py functions use Python lists like vals_f32 = [cutlass.Float32(0.0)] * 16 and index them with vals_f32[i]. But CuTeDSL's i is an MLIR value, not a Python int.

I need to use CuTe register tensors instead of Python lists. Or I need to use cute.make_rmem_tensor for the accumulator buffer.

Actually, looking at the existing kernel code, register tensors are used with cute.make_rmem_tensor. The SwiGLU epilogue uses silu_gate_buf = cute.make_rmem_tensor(tiled_copy_r2s.retile(tTR_rAcc).shape, self.c_dtype). I should use the same pattern for the SwiGLU accumulation buffer.

But for the test kernel, I need a simpler approach. Instead of accumulating in a Python list, I should process the 16 elements one at a time in the loop, computing amax, then doing a second pass for E2M1 quantization.

Let me rewrite the test kernel to avoid Python list indexing with CuTeDSL variables:OK, I'm going in circles on the test kernel's I/O. Let me step back and think about this properly.

The real question is: how does the actual fused_swiglu kernel write to GMEM? It uses cute.copy(tma_atom_c, bSG_sC, bSG_gC) with a TMA store atom. The BF16 data goes SMEM → GMEM via TMA.

For the FP4 output, I need a similar path: registers → SMEM → GMEM. The FP4 data is 9 bytes per 16-element block (8 FP4 + 1 SF). In the actual kernel, I'd write to SMEM and then do a TMA store or a regular GMEM store.

But for the test, I don't need TMA. I just need to write 9 Int32 values (8 FP4 bytes + 1 SF byte) to GMEM. The simplest approach that WORKS in CuTeDSL:

1. Create a register tensor with the 8 FP4 bytes and 1 SF byte
2. Copy it to GMEM using cute.copy with a simple copy atom

Let me look at how cute.copy is used with CopyUniversalOp in the codebase:OK so the pattern is: cute.copy(copy_atom, rmem_tensor, gmem_tensor). The copy atom uses CopyUniversalOp() with num_bits_per_copy=128 for 4 × Int32 (128 bits). For my test, I need 8 × Int32 for FP4 + 1 × Int32 for SF.

But actually, I realize there's a much simpler approach. Since the FP4 data will be written to GMEM in the actual kernel, I should design the GMEM write pattern to match how the L2 GEMM reads it. The L2 GEM reads padded_activated_fp4 which is (tokens, intermediate//2) in float4_e2m1fn_x2 format with scales (tokens, intermediate//16) in float8_e4m3fn.

For the test, I just need to verify the math works on GPU. Let me use a completely different approach: compute the FP4 quantization in the kernel, write the result back as BF16 (just the dequantized values), and compare with the original BF16 after round-trip. This avoids the FP4 byte write entirely.

Actually no, that doesn't verify the bit-exact FP4 output. Let me try yet another approach: write the FP4 bytes into a BF16 output tensor (packing 2 bytes per BF16 value) and unpack on the CPU side.

OK, I'm overcomplicating this. Let me just check if cute.arch.store exists by trying a different way. The existing codebase uses cute.arch.load(ptr, dtype) for scalar loads. Maybe there's an equivalent cute.arch.store or cute.arch.global_store that I can discover on B200.

Actually, looking at the CUTLASS Python source more carefully, there should be a cute.arch.store — it's the counterpart of cute.arch.load. Let me check the CUTLASS Python documentation or source:The docs don't show cute.arch.store — they only show thread_idx, barriers, SMEM/TMEM management, etc. The cute.arch.load used in the codebase must be from a different module or a custom wrapper. Let me check:So cute.arch.load(ptr, dtype) loads from a GMEM pointer. There must be a corresponding cute.arch.store. Let me search the CUTLASS GitHub:There's a GitHub issue about cute.arch.load with BFloat16. And from issue #3266, cute.arch.load exists but has issues with BFloat16. Let me check if cute.arch.store exists:Let me just skip the CuTeDSL kernel test for now and instead verify the math works by running it within the actual fused SwiGLU kernel. The FP4 quantization math is already verified in Python (0/500 failures). The CuTeDSL kernel test is just to verify that the CuTeDSL implementation compiles and produces the same results on GPU.

Instead of fighting with cute.arch.store, let me move directly to integrating the FP4 quantization into the actual SwiGLU epilogue. I'll write the FP4 bytes to SMEM and then to GMEM using the same patterns the MoE kernel already uses.

Let me focus on the kernel modification now:I see — the kernel signature doesn't have fp4_out, sf_out, l2_global_scale yet. These were added to __call__ but never wired to kernel(). I need to add them.

This is a significant change to the kernel. Let me plan it carefully:

Changes needed to fused_swiglu.py:

1. Add fp4_out, sf_out, l2_global_scale parameters to kernel()
2. In the epilogue warp section, when fp4_mode=True:
   • For gate subtiles: compute SiLU(gate), save to silu_gate_buf. Skip SMEM write and TMA store.
   • For up subtiles: compute SwiGLU result. Save to a swiglu_buf register buffer (2 consecutive up subtiles = 16 values).
   • After every 2 up subtiles: compute amax + FP8 E4M3 scale + E2M1 pack. Write FP4 bytes to SMEM, then to GMEM (via TMA store to the FP4 output tensor).
3. Write FP8 scale factors similarly.
4. When fp4_mode=False: keep existing BF16 path unchanged.

The biggest challenge is the GMEM write for FP4. Let me look at how TMA stores work for the C tensor:The TMA store for C is:

1. tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(c_tma_op, c_gemm, epi_smem_layout, epi_tile) — create TMA atom for C
2. bSG_sC, bSG_gC = cpasync.tma_partition(...) — partition for TMA store
3. cute.copy(tma_atom_c, bSG_sC, bSG_gC, tma_desc_ptr=desc_ptr_c) — do the TMA store

For the FP4 output, I'd need a similar TMA setup. But the FP4 output tensor has a different shape and dtype than the BF16 C tensor. This is complex and requires careful layout design.

Let me take a more pragmatic (but still correct) approach: write the FP4 bytes directly from registers to GMEM using cute.copy with a non-TMA copy atom. The FP4 output doesn't need TMA because:

1. The write is small (9 bytes per 16-element block)
2. The FP4 output tensor is contiguous and simple
3. TMA setup for FP4 is complex and may not be needed for correctness

The simplest write pattern: each thread computes its FP4 bytes and writes them to GMEM via a CopyUniversalOp copy atom. But I need to figure out how to construct the GMEM tensor for the FP4 output.

Actually, let me look at how the cute.arch.load call works — it uses raw pointer arithmetic. For the FP4 output, I can do the same: compute the GMEM address from the work tile info and write using a similar pattern.

But cute.arch.store doesn't exist. Let me check if there's a cute.copy that works with raw pointers:

Wait — actually, the cute.arch.load in the codebase is used for scalar loads from GMEM. If there's a cute.arch.load, there should be a cute.arch.store — they're symmetric operations. The fact that it's not in the docs doesn't mean it doesn't exist. It might just be undocumented.

Let me just try it on B200. If cute.arch.store doesn't exist, I'll adapt. But I need to first get the test kernel to compile. The previous attempt failed because cute.arch.store was used. Let me try a different approach for the test: write the results to the output tensor via cute.copy with a simple copy atom, using proper CuTe tensor construction.

Actually, I just realized I can look at the B200's CUTLASS Python installation to check:cute.arch.store EXISTS on the B200! The test just failed because of the Python list indexing issue, not because cute.arch.store doesn't exist. Let me fix the test kernel properly and rerun:Two errors found:

1. cute.math.fmax doesn't exist — the correct API is cute.math.fmin and cute.arch.fmax (based on the README's CuTeDSL constraint #4)
2. The fp4_quant.py functions also use cute.math.fmax which won't work

Looking at the README:

cute.arch.fmax is impure for the vectorizer. Use it inside plain range(), never inside vectorize=True.

So it's cute.arch.fmax, not cute.math.fmax! And for fmin, it's cute.math.fmin. Let me check the codebase for the exact APIs:So cute.arch.fmax exists, and cute.math.fmin exists. But cute.math.fmax does NOT exist. Let me fix fp4_quant.py to use cute.arch.fmax instead of cute.math.fmax:Good, all cute.math.fmax replaced with cute.arch.fmax. Now fix the test kernel too: