ISSUE — Lightning Indexer FP4 dequant decodes E2M1 wrong

Status: FIXED ✅ — E2M1 LUT fix landed in both dsv4/kernels/indexer/indexer_score_topk.cu and dsv4/kernels/cuda/indexer_score_topk.cu. Crossed off the list. Severity: Was HIGH. Corrupts top-k selection, which is the whole job of the indexer. Scope: dsv4/kernels/indexer/indexer_score_topk.cu and the duplicate dsv4/kernels/cuda/indexer_score_topk.cu. Does NOT touch FMHA, MoE, or the GEMM stack.

TL;DR

The indexer dequantizes FP4 keys by treating the 3-bit nibble as a linear integer magnitude (0–6). NVFP4 is E2M1 — the nibble is an index into the level set [0, 0.5, 1, 1.5, 2, 3, 4, 6]. The quantizer stores the index; the dequant never inverts it. Every nonzero indexer key is read at the wrong value.

This is the textbook "scalar math written to make it compile" trap. It runs, it produces numbers, the cosine on a smoke test looks plausible because the map is monotonic — and it silently mis-ranks blocks. The paper's indexer targets 99.7% recall (§5.2.1). A wrong decode burns that, and because the bad base is being used as the "known-correct FP32 oracle," it poisons every downstream comparison.

The bug (FIXED)

Was in dsv4/kernels/indexer/indexer_score_topk.cu, dequant_fp4_scalar (~line 33). Fixed: replaced (float)mag * scale with kE2M1[nibble & 0x07] * scale using __constant__ LUT. Both copies fixed. Deduplication (single #include) still TODO.

Why it's wrong — read it off the actual quantizer, do not assume

dsv4/ops/quantize.py:15

E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]

quantize.py rounds to half-steps, then maps half-step → E2M1 index via _get_step_to_idx_lut, and stores that index (sign → +8) as the nibble. The forward decode that the GEMM/quantize side already documents lives in dsv4/kernels/cuda/deinterleave_quantize.cu:13 (half_step_to_e2m1, "Matches Python step_to_idx LUT"). The indexer is the only consumer that skips the LUT.

Concrete error (nibble index → what it should be → what the code returns):

index	E2M1 value	code returns	wrong by
1	0.5	1.0	2×
2	1.0	2.0	2×
5	3.0	5.0	1.67×
6	4.0	6.0	1.5×
7	6.0	7.0	1.17×

Per-element it's monotonic (so naive cosine looks "fine"), but I[t,s] = Σ_h w_h · ReLU(q·K) is a nonlinear reduction across head_dim and heads — the distortion is per-element and uneven, so the argsort over blocks changes. That is exactly the quantity top-k depends on.

Fix (interim, correct FP32 base)

Replace the scalar arithmetic with the actual E2M1 level set. Use __constant__ memory, not a recomputed branch ladder.

__constant__ float kE2M1[8] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f};

__device__ __forceinline__ float dequant_fp4_e2m1(
    uint8_t packed, int lane, float group_scale, float global_scale
) {
    int nibble = (lane == 0) ? (packed & 0x0F) : (packed >> 4);
    int sign   = (nibble >> 3) & 1;
    float mag  = kE2M1[nibble & 0x07];          // INDEX -> VALUE via LUT
    float val  = mag * group_scale * global_scale;
    return sign ? -val : val;
}

Fix both copies. Better: delete the duplicate and #include one source so this can never drift again.

Fix (production target — where this should actually go)

The interim fix removes the correctness bug but is still circa-2010 scalar dot products. Paper §5.2.1 is explicit: the indexer QK path is cached, loaded, and multiplied entirely in FP4 on tensor cores. The production indexer is:

FP4 keys + FP4 indexer queries → tcgen05.mma FP4 (mxf4nvf4 kind), hardware decodes E2M1 natively (the LUT bug literally cannot exist on this path).
index scores accumulated in the tensor core, ReLU + per-head weight in the epilogue, warp-level top-k (shfl/ballot), not a serial shared-mem heap merge.

Land the LUT fix first to get a trustworthy oracle, then build the FP4 MMA path against it.

Test plan — measure, don't eyeball

Unit (exactness): quantize a known vector with quantize.py, dequant in the kernel, assert bitwise-equal to E2M1_MAGNITUDES[idx] * group * global for all 8 indices and both signs. This alone catches the regression.
Recall (the metric that matters): build I[t,:] three ways — (a) FP32 reference (no quant), (b) kernel before fix, (c) kernel after fix. Report |TopK_a ∩ TopK_b| / k and … ∩ TopK_c / k at k = 512 (Flash) and 1024 (Pro). Expect (b) well below paper's 99.7%, (c) at/near it.
End-to-end: CSA attention output cosine with selection from (a) vs (c). Regression gate: cos ≥ 0.999.
Run at a realistic compressed length (≥ 8k tokens / ≥ 2k compressed blocks), not a 128-token toy — ranking corruption only shows up once blocks compete.

ENGINEERING DOCTRINE (applies to every issue in this repo)

These are the rules an agent must not violate while "making it work." This bug is what happens when they are violated.

1. A wall in CuTeDSL/CUTLASS means raw CUDA C++ — NOT Python.

When the DSL toolchain can't express something (MLIR can't lower float→int, MLIR optimizer hangs at hd=512, TMEM atom layouts don't pair), the fallback is a raw CUDA C++ kernel, not a host-side Python loop / multi-launch merge. Python merges and per-head launches are orchestration scaffolding, not a kernel. They are the single biggest threat to "production grade" in this codebase. If you hit a DSL wall: drop to .cu/.cuh, keep the math on the device, keep it one launch.

2. Raw CUDA C++ does NOT mean circa-2010 scalar math.

Falling back to CUDA is not license to write scalar FMA loops. On Blackwell the floor is: tcgen05.mma / UMMA for any matmul, TMEM accumulators, TMA loads, vectorized SMEM access, warp-level reductions (__shfl_xor, ballot), __constant__ LUTs over branch ladders. A scalar dot-product is acceptable ONLY as a temporary correctness oracle that is explicitly labeled as such and has a tensor-core replacement tracked. "It compiles and the cosine is 0.97" is not done.

3. Do NOT guess layouts and shapes. Print, document, then code to the data.

Almost every multi-day rabbit hole here started as a guessed layout (P (128,128) accumulation, V MN/K swap, TMEM column mapping, sf_dtype E8M0 vs E4M3). The discipline:

Print the actual thing — dtype, shape, stride, sf layout, TMEM offset, MMA instruction shape — at construction and at the kernel boundary.
Write down what you observed in the issue file (the table above is an example: indices read off quantize.py, not assumed).
Code against the observed data, then re-print to confirm. A guessed layout that happens to pass a toy test is a landmine. The E2M1 comment "magnitude (0-6)" was a guess; the data said "index into 8 levels." The data wins.

6.9 KiB Raw Blame History Unescape Escape