[bastianhagedorn]

📌 Description

flashinfer/attention/prims_ts/kernels/mla_decode/throughput_2cta/resources.py, TmemCorrResource.load_corr. The correction warps of the 2-CTA MLA decode kernel read the 4-word correction record of the current stage (row_sum, row_max, correction_factor, no_correction) from tensor memory with one prims.tcgen05_ld(..., num=4). TmemCorrResource is a 2-stage ring written by the softmax warps (store_corr: tcgen05_st after the runtime's producer acquire) and consumed by the correction warps; the task runtime releases the stage (an mbarrier.arrive on the stage's empty barrier) as soon as load_corr returns, and the loaded words are consumed only afterwards: rescale_o votes on no_correction and multiplies O by correction_factor, epilogue_store reads row_sum / row_max from final_row_stats.

The function issued prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) before the load. tcgen05.ld writes its destination registers asynchronously; the PTX ISA (tcgen05 instructions, "Asynchronous operations", section 9.7.18.6.4.5 in the 9.x documents) says that the load completes only at a tcgen05.wait::ld in the issuing thread and that a register dependency orders only the dependent instruction, "not any other form of ordering". The arrive that releases the stage has no register dependency on the loaded values, so nothing orders it after the load: a softmax warp parked on that barrier may acquire the stage and overwrite the same four columns while the load is still reading them. If that happened, tile t's O rescale would use tile t+2's correction_factor / no_correction (a skipped or misapplied rescale), and with a single KV tile the only record feeds the epilogue's normalisation. The wait that was there is not useless, it covers the previous iteration's O loads, which is why it stays, just on the other side of the load.

Change: one line, prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) moved from before the record load to right after it, with a short comment. This is what the 1-CTA kernel of the same backend already does after each of its loads (throughput_latency_1cta/resources/tmem_corr.py). No other behaviour changes.

Observed behaviour (B200, compute capability 10.0). The compiled stock kernel shows the gap: LDTM.x4 of the record, then the releasing SYNCS.ARRIVE with an empty wait mask, then the first consumer (ISETP on no_correction, the FMUL2s by correction_factor) waiting the load's scoreboard after the arrive. With the fix the arrive itself waits the scoreboard (watdb=SB0), no instruction is added (same SASS size), and the output is bitwise identical to the unpatched kernel. A sentinel experiment (the softmax warps write a per-tile marker into correction_factor, the correction warps record the factor they actually multiply with, plus extra TMEM loads queued in front of the record load and a 10-100 us stall between the release and the first use) read 211.7 M records over 11 800 launches of the unpatched kernel (4 KV-length ranges, non-persistent and CLC persistent schedules) with 0 stale records: in this kernel the refill of a released stage is ordered behind the consumer's first use through the O and P pipelines (the softmax warp acquires P before it stores the record, P is released by the MMA warp after PV, which waits on O, and O is released by the correction warps after they have used the record), at least 2.2 us behind in every measurement. So the race is not observable with the shipped pipeline depths; it would be with a deeper O or P ring or a per-N O pipeline. The fix makes the ordering architectural at no cost.

🔍 Related Issues

Same release-before-use shape as #8787 and #8854 (TMEM loads ordered before the release / dealloc in the GDN, MQA-logits and TGV kernels) and NVIDIA/TensorRT-LLM#21849 (the same correction-record load in the Blackwell MLA decode kernels). The vendored copy of this tree in TensorRT-LLM (tensorrt_llm/_torch/attention/backends/prims_ts) is updated by the paired NVIDIA/TensorRT-LLM#21987, which pins this change through the vendor lock.

🚀 Pull Request Checklist

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues. (flashinfer/attention/prims_ts/kernels/ is excluded from ruff by .pre-commit-config.yaml; the file compiles.)

🧪 Tests

• Tests have been added or updated as needed. (No new test: the race needs a deeper O/P pipeline or an injected re-ordering of the refill to become observable; 211.7 M records on a B200 without a stale read, see above.)

• All tests are passing (unittest, etc.). (Not yet run on CI. The kernel compiles with the change; the patched and unpatched kernels produce bitwise identical output on a B200.)

• This PR requires full-scope CI, not targeted tests (expensive; justify below).

  Suggested CI test targets: tests/attention/test_attention_ts_mla_decode.py

Reviewer Notes

The change is the position of one existing prims.tcgen05_wait(LOAD) call. A library-level alternative would be for the task runtime's consumer release to issue tcgen05.wait::ld before releasing a tensor-memory-backed resource; this PR takes the local fix the 1-CTA kernel already uses.

🤖 Generated with Claude Code

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

• X
• Mastodon
• Reddit
• LinkedIn