Skip to content

[None][fix] Wait for the Prim-TS MLA decode correction-record load before releasing its stage #21987

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
bastianhagedorn wants to merge 1 commit into
base: main
Choose a base branch
from
+8 −5
Draft

[None][fix] Wait for the Prim-TS MLA decode correction-record load before releasing its stage #21987

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions 3rdparty/vendor_sources.lock.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,16 +13,16 @@ vendors:
patch_digest: sha256:42e8306de30fef448378066321977aeb7b42c3405dc883f9527e256590a58b46
digest: sha256-tree-v1:12eb5c011c3fef1831aeba02691df8e2f5bff323f9da9c58c0f31e80883a9f06
flashinfer-prims-ts:
url: https://github.com/yuxianq/flashinfer.git
branch: trtllm-prims-ts-dev-20261107
commit: 911d6a9d62308b6c8a5249e8dc4aabc7518020f1
url: https://github.com/bastianhagedorn/flashinfer.git
branch: trtllm-prims-ts-mla-decode-corr-load-wait-20261110
commit: 8846717c966bdd53653b218ee563711a4235291a
source: flashinfer/attention/prims_ts
destination: tensorrt_llm/_torch/attention/backends/prims_ts
include:
- '**/*.py'
patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch
patch_digest: sha256:f35ef4d096314a21c710b849b0a1f731cf1edbf50d32cdee46b635f45fe6b47a
digest: sha256-tree-v1:bf74c7664256a1c7064163ca417a6c589f8ab8f47bb478444bb1b3c402c856ee
digest: sha256-tree-v1:cbe1befea1424d9630170d1fc7c89a46d1aeb42cedcc15a94fee1bb032d92734
openengine-proto:
url: https://github.com/ai-dynamo/openengine.git
tag: v0.1.0
Expand Down
5 changes: 4 additions & 1 deletion ...rt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3695,12 +3695,15 @@ def load_corr(self, stage_info: StageInfo):
tmem_raw_addr = (tmem_warp_row_id << 16) | col_offset
tmem_ptr_arr = prims.make_tmem_ptr(tmem_raw_addr, Float32)

prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD)
loaded = prims.tcgen05_ld(
TCGEN05_32B_SHAPE,
tmem_ptr_arr,
num=4,
)
# The runtime releases this stage to the softmax warps as soon as the
# task body returns, and the loaded words are consumed later (rescale_o,
# epilogue), so the load must be complete before the release.
prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD)

self.final_row_stats[0] = loaded[0]
self.final_row_stats[1] = loaded[1]
Expand Down
Loading