Sparse MLA decode is not batch-invariant: a row's output depends on its batchmates (1-2 ULP), flipping argmax at near-tie tokens

Sparse MLA decode is not batch-invariant: a row's output depends on its batchmates (1–2 ULP), flipping argmax at near-tie tokens

Hi again — following up on our earlier benchmark report (#192), this time about correctness rather than throughput.

Summary: sparse_decode_fwd is deterministic (identical inputs → bit-exact identical output) but not batch-invariant: the same query row produces a different result depending on how many other rows share the batch. The difference is 1–2 ULP of bf16, which is normally invisible — but at a decode token whose top-2 logit margin is near zero it flips the argmax, and the sampled token changes.

Downstream this shows up as an LLM returning different completions for identical temperature=0 requests, at a rate that scales with concurrency. We traced it from that symptom down to this kernel.

Measurement

Instrumented at the flash_mla_with_kvcache call site inside vLLM's DeepSeek-V4 decode path. On sampled calls we do two things and compare both against the output the real call produced:

1. Re-run with the identical tensors → tests non-determinism.
2. Solo re-run of row 0 alone — q[:1], correspondingly sliced indices / topk_length / extra_indices_in_kvcache / extra_topk_length, and a fresh FlashMLASchedMeta → tests batch-shape dependence.

max_absdiff on differing rows: 0.00195–0.0039 — 1–2 ULP of bf16 near 1.0.

The re-run column is the important control: with the batch shape held fixed the kernel is bit-exact reproducible, so this is not a race, not atomics, and not non-determinism in the usual sense. Only changing batch composition changes the answer.

All kernel inputs were independently verified well-formed at the time of the differing calls (indexer top-k selections identical across rows at equal positions; KV content at every fetched slot byte-stable; topk_length equal to the valid-entry count of the index array in 57,616/57,616 rows). The indices and data are right — the arithmetic order differs.

Where it comes from

We believe the mechanism is visible in the source.

csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu:

int payload = cutlass::ceil_div(total_num_blocks, num_sm_parts) + fixed_overhead_num_blocks;

total_num_blocks is accumulated across the entire batch (the loop above it sums num_blocks + fixed_overhead_num_blocks for every request, then a warp reduction). So payload — the block quota per SM part — is a function of every other row's length. The greedy packing loop that follows then decides where each request is split (now_n_split_idx, is_first_req_splitted, is_last_req_splitted) from that batch-wide quota. A given row's KV range therefore gets divided into a different number of splits depending on which rows it happens to share the batch with.

Those splits are then recombined in csrc/smxx/decode/combine/combine.cu:

for (int split = 0; split < my_num_splits; ++split) {
    float lse_scale = smem_buf[warp_idx][split];
    result[i].x += lse_scale * datas[i].x;
    ...
}

a sequential, non-compensated fp32 accumulation whose length and per-split rescale factors both come from the split decision. Different split count → different addends in a different order → different rounding. That matches the observed 1–2 ULP delta.

(For what it's worth, the recent 15f13e5 — extending the SWITCH_MAX_SPLITS ladder to 256 because "on GPUs with high SM counts the decode scheduler can request more than 160 splits" — is independent evidence that the split count varies widely with batch and hardware.)

Why it matters downstream

For most tokens a 1-ULP perturbation changes nothing. It matters only where the top-2 logits are already nearly tied — and there it changes which token is emitted. In our end-to-end measurements the top-2 margin at the divergent token collapses from ~18 logprobs at concurrency 1 to 0.1–2.0 at concurrency 32. The result is a single wrong token, after which the completion continues plausibly from the wrong prefix.

Practical consequences for serving:

• greedy decoding is not reproducible when batch composition varies
• an identical request can succeed or fail depending on unrelated concurrent traffic
• the rate rises with concurrency, so it is worst under production load

Question

Is batch-invariant splitting something you would consider supporting — e.g. an opt-in mode where the split count for a row derives only from that row's own sequence length (and fixed hardware parameters), rather than from batch-wide totals?

We understand the current design is deliberate: batch-wide balancing is what keeps SMs busy on ragged batches, and a per-row split rule will cost occupancy. So this may be better as an opt-in flag than a default. We're happy to prototype it and measure the throughput cost on B300 with the same benchmark harness from #192, if you think that's a direction you'd accept.

Environment

• NVIDIA B300 SXM6 (compute capability 10.3), driver 610.43.02
• CUDA 13.0, PyTorch 2.11.0+cu130
• FlashMLA as vendored by vLLM 0.27.1 (vllm-project/FlashMLA @ a8f794d1)
• DeepSeek-V4-Flash (NVFP4), TP=1, sparse decode path with extra_k_cache
• Downstream tracking issue: https://github.com/vllm-project/vllm/issues/53257