Fix SM100 V32 sparse decode FP8 handling

[jianglan89]

Summary

• fix the SM100 V32 sparse decode path to read per-token FP8 scales as the stored float32 values
• keep the SM100 sparse index validity rule consistent with the reference path by masking positive indices outside the physical KV cache
• add an SM100 sparse decode regression test with non-power-of-two V32 scales and physical-OOB indices

Background

The real B300 precision issue is the V32 FP8 scale handling in the SM100 sparse decode head64 path. V32 stores four float32 scale values per token. The old SM100 code converted those float32 scales through e8m0 before dequantization, which changes the scale values and produces wrong FP8 dequantization.

SM9x and SM10x do not share this sparse decode implementation. SM9x was already correct; the bug is in the SM100 path.

flowchart TD
    A[V32 FP8 sparse decode] --> B{GPU path}
    B -->|SM9x / H20| C[SM90 path]
    C --> D[reads V32 scales correctly]
    D --> E[matches PyTorch reference]

    B -->|SM10x / B300| F[SM100 head64 path]
    F --> G[BUG: float32 scale -> e8m0]
    G --> H[wrong dequantization]
    H --> I[precision mismatch]

    J[fix] --> K[read float32 scale directly]
    K --> L[SM10x matches PyTorch reference]

Loading

The V32 FP8 KV cache stores each token as 656 bytes:

• bytes [0, 512): 512 e4m3 FP8 values for the latent/nope part
• bytes [512, 528): four float32 scale values, one per 128 FP8 values
• bytes [528, 656): 64 bf16 values for the RoPE part

This patch also masks positive sparse indices outside the physical KV cache. That is a related correctness fix for partial/local KV cache inputs, but it is not the main source of the B300 numeric drift. With OOB indices manually masked before calling the old SM100 kernel, the old path still differs from the PyTorch reference by about 28%-32% relative mean error on the H20 dump layers. That remaining error comes from the V32 scale conversion through e8m0.

Before / After

All numbers below compare flash_mla against the exact PyTorch dequantized sparse decode reference on the same V32 FP8 H20 dump, running on CUDA 12.9 / sm103.

layer	before max_abs	before rel_mean	after max_abs	after rel_mean
00	0.00933838	0.359332	0.00012207	0.000801996
01	0.00540161	0.311769	0.00006104	0.00165462
02	0.00473022	0.343794	0.00006104	0.00164364
03	0.02017210	0.364807	0.00012207	0.00219231
04	0.00662231	0.363873	0.00006104	0.00204405

To isolate the root cause, I also manually masked physical-OOB indices before calling the old SM100 kernel. The old kernel still had large error, which points to the V32 scale conversion as the primary cause.

layer	old SM100 after manual OOB mask max_abs	old SM100 after manual OOB mask rel_mean
00	0.00891113	0.323777
01	0.00540161	0.293805
02	0.00473022	0.303478
03	0.01782227	0.307482
04	0.00564575	0.283648

Testing

• PYTHONPATH=tests CUDA_LAUNCH_BLOCKING=1 python tests/test_flash_mla_sm100_sparse_decode.py
• H20 dump compare on CUDA 12.9 / sm103 after this patch: overall PASS; flash_mla vs torch_dequant_ref max_abs <= 0.00012207 across layers 00-04