auto: pre-test commit
This commit is contained in:
@@ -84,4 +84,34 @@ But PyTorch's `.to(float8_e4m3fn)` rounds 250.0 to 256.0 (exp=15, mant=0). This
|
||||
|
||||
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:
|
||||
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:
|
||||
Reference in New Issue
Block a user