[bastianhagedorn]

@coderabbitai summary

Description

tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py, TmemCorrResource.load_corr. The correction warps of the Prim-TS 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), and the same fix as #21849 for the correction-record load of the Blackwell MLA decode kernels (mla_decode_fp16.py / mla_decode_fp8.py). No other behaviour changes.

Vendored source

tensorrt_llm/_torch/attention/backends/prims_ts is a vendored copy of FlashInfer's flashinfer/attention/prims_ts (lock entry flashinfer-prims-ts in 3rdparty/vendor_sources.lock.yaml). Following 3rdparty/vendor-sources.md, the change was exported with scripts/vendor/manage.py export to a temporary branch of my FlashInfer fork at the locked commit (bastianhagedorn/flashinfer branch trtllm-prims-ts-mla-decode-corr-load-wait-20261110, commit 8846717c, one commit on top of the canonical trtllm-prims-ts-dev-20261107 tip 911d6a9d, which is the live tip of that branch) and the lock is pinned to that commit (manage.py pin); the compatibility patch is unchanged and manage.py check passes offline and against the source checkout. The paired upstream PR is flashinfer-ai/flashinfer#9319.

schema_version: 1
vendors:
  flashinfer-prims-ts:
    upstream_prs:
      - https://github.com/flashinfer-ai/flashinfer/pull/9319

Hardware verification

The FlashInfer kernel (byte-identical to this vendored file at the site) was run on a B200 (compute capability 10.0) with the repository's own test inputs and fp32 reference check: batch 64, 128 heads, latent 512, rope 64, page 64, KV lengths 4111..8192 (also 256..455, 1024..2023 and 8192 everywhere), bf16, 2-CTA clusters, non-persistent and CLC persistent schedules. 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, no instruction is added (same SASS size), and the output is bitwise identical to the unpatched kernel over 2000 launches. A sentinel experiment (per-tile marker written into correction_factor by the softmax warps, the factor actually used recorded by the correction warps, extra TMEM loads queued in front of the record load, 10-100 us stalls between the release and the first use) read 211.7 M records over 11 800 launches of the unpatched kernel 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 measurable cost.

Test Coverage

The kernel is exercised by tests/unittest/_torch/attention/test_prims_ts_attention_backend.py (test_prims_ts_deepseek_v3_lite_mla_generation runs MLA decode through the Prim-TS backend; the 2-CTA throughput policy is the default selection for dense decode shapes). No new test: memory-model conformance, no observable failure to test for with the shipped pipeline depths (see above). The change is limited to the position of one existing prims.tcgen05_wait call.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

[yuxianq]

Vendor validation: flashinfer-prims-ts

@bastianhagedorn — Vendor promotion needs maintainer action:

The base lock now uses refresh branch trtllm-prims-ts-dev-20261107, but this monitor still publishes to trtllm-prims-ts-dev-20261012.

Fix: A maintainer must stop this monitor and migrate its workdir with --canonical-branch trtllm-prims-ts-dev-20261107 --migrate-canonical-branch-from trtllm-prims-ts-dev-20261012. Verify a dry run before resuming publication. Do not promote back to the old branch or change contributor metadata to bypass this cutover.

The bot will recheck automatically on a later poll.

Checked PR revision: cb9927309da5. This is a snapshot; newer results are posted as new comments. Observed dev/upstream tip SHAs may be stale: tip movement alone does not post a new comment when the result is unchanged. Fetch the live branch before rebasing or choosing a new pin.