Repository navigation
[Fix] Skip empty-work gathers in gfx1250 MLA PS1 - #7089
Draft
XiaobingSuper wants to merge 2 commits into
Draft
XiaobingSuper wants to merge 2 commits into
XiaobingSuper wants to merge 2 commits into
Conversation
Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
Co-authored-by: Cursor <cursoragent@cursor.com>
qichu-yun
added a commit
to ROCm/ATOM
that referenced
this pull request
Oct 29, 2026
The previous commit planned an MTP verify step's sparse per-token work over the real tokens and attributed the GPU memory fault to the padded tokens' empty work items reading the sparse-index buffer past what the step wrote. That was a guess, and it was wrong. What actually runs past an end: the MTP drafter's step 0 runs the target's token stream through its own sparse MLA layer reusing the target's verify metadata, and that stream is not padded (eagle_proposer: "Step 0 carries the target's stream and pads nothing"), so its q and o hold scheduled_tokens rows. The persistent MLA kernel reads q and writes o for every work item, zero-KV padding included, so a plan over the cudagraph-padded running_tokens read q and wrote o past their ends. Whether that faulted depended on whether the page after q happened to be mapped, so the fault was intermittent; when it did not fault, it silently overwrote whatever followed o. Measured on MI455 (gfx1250, FlyDSL PS1 decode), GLM-5.2 MTP K3, 39 requests padded to 40, LLM-Booster #287 / #342: - the four GPUs' fault addresses all lay 20-24 KB past the end of a 140-row q (35 x 4 tokens), with the next page unmapped; - an eager-call check saw q with 156 rows against a plan reaching row 158; - a model-free test with q/o sized to the real tokens and a sentinel tail shows 3 work items past q and 3 rows written past o at step 0, with the original aiter kernel and with ROCm/aiter#7089 alike, and none when planned over the real tokens. So the code change of the previous commit is the fix; this commit corrects its rationale, records the planned row count in the metadata (sparse_mtp_num_tokens), and has MLAAttention._forward_decode refuse a sparse MTP plan larger than its q -- host ints only, no sync, skipped under DCP, which rebuilds the plan per layer. Tests: tests/model_ops/test_mla_sparse_mtp_plan_rows.py (the refusal, and a plan within q reaching the kernel for both the drafter's unpadded and the target's padded rows); the builder test now also checks the recorded row count.
qichu-yun
added a commit
to ROCm/ATOM
that referenced
this pull request
Nov 5, 2026
The previous commit planned an MTP verify step's sparse per-token work over the real tokens and attributed the GPU memory fault to the padded tokens' empty work items reading the sparse-index buffer past what the step wrote. That was a guess, and it was wrong. What actually runs past an end: the MTP drafter's step 0 runs the target's token stream through its own sparse MLA layer reusing the target's verify metadata, and that stream is not padded (eagle_proposer: "Step 0 carries the target's stream and pads nothing"), so its q and o hold scheduled_tokens rows. The persistent MLA kernel reads q and writes o for every work item, zero-KV padding included, so a plan over the cudagraph-padded running_tokens read q and wrote o past their ends. Whether that faulted depended on whether the page after q happened to be mapped, so the fault was intermittent; when it did not fault, it silently overwrote whatever followed o. Measured on MI455 (gfx1250, FlyDSL PS1 decode), GLM-5.2 MTP K3, 39 requests padded to 40, LLM-Booster #287 / #342: - the four GPUs' fault addresses all lay 20-24 KB past the end of a 140-row q (35 x 4 tokens), with the next page unmapped; - an eager-call check saw q with 156 rows against a plan reaching row 158; - a model-free test with q/o sized to the real tokens and a sentinel tail shows 3 work items past q and 3 rows written past o at step 0, with the original aiter kernel and with ROCm/aiter#7089 alike, and none when planned over the real tokens. So the code change of the previous commit is the fix; this commit corrects its rationale, records the planned row count in the metadata (sparse_mtp_num_tokens), and has MLAAttention._forward_decode refuse a sparse MTP plan larger than its q -- host ints only, no sync, skipped under DCP, which rebuilds the plan per layer. Tests: tests/model_ops/test_mla_sparse_mtp_plan_rows.py (the refusal, and a plan within q reaching the kernel for both the drafter's unpadded and the target's padded rows); the builder test now also checks the recorded row count.
valarLip
pushed a commit
to ROCm/ATOM
that referenced
this pull request
Nov 9, 2026
#2715) * [MLA] Plan a padded sparse MTP verify's work over the real tokens only prepare_decode built the sparse per-token persistent metadata of an MTP verify step over running_tokens, so a cudagraph-padded batch planned its padded tail tokens too. Those have zero sparse KV, but get_mla_metadata_v1 still gives every row a work item, starting at sparse_kv_indptr[-1] -- the end of what the step wrote to the sparse-index buffer. With those items planned, GLM-5.2 MXFP4 with MTP K=3 on MI455 (gfx1250, FlyDSL PS1 MLA decode) memory-faults in the decode kernel some tens of steps into any padded batch -- presumably by reading indices past that end -- which ends every GSM8K full run near its tail. Plan over scheduled_tokens instead, as the cprr branch already does for its requests and as the capture path does. The persistent kernel is work_indptr-driven, so the padded rows are left unplanned. Adds a regression test: a padded sparse MTP verify plans 2*q tokens for 2 real requests padded to 4. * [MLA] Explain the padded sparse MTP plan fix: the drafter's step 0 rows The previous commit planned an MTP verify step's sparse per-token work over the real tokens and attributed the GPU memory fault to the padded tokens' empty work items reading the sparse-index buffer past what the step wrote. That was a guess, and it was wrong. What actually runs past an end: the MTP drafter's step 0 runs the target's token stream through its own sparse MLA layer reusing the target's verify metadata, and that stream is not padded (eagle_proposer: "Step 0 carries the target's stream and pads nothing"), so its q and o hold scheduled_tokens rows. The persistent MLA kernel reads q and writes o for every work item, zero-KV padding included, so a plan over the cudagraph-padded running_tokens read q and wrote o past their ends. Whether that faulted depended on whether the page after q happened to be mapped, so the fault was intermittent; when it did not fault, it silently overwrote whatever followed o. Measured on MI455 (gfx1250, FlyDSL PS1 decode), GLM-5.2 MTP K3, 39 requests padded to 40, LLM-Booster #287 / #342: - the four GPUs' fault addresses all lay 20-24 KB past the end of a 140-row q (35 x 4 tokens), with the next page unmapped; - an eager-call check saw q with 156 rows against a plan reaching row 158; - a model-free test with q/o sized to the real tokens and a sentinel tail shows 3 work items past q and 3 rows written past o at step 0, with the original aiter kernel and with ROCm/aiter#7089 alike, and none when planned over the real tokens. So the code change of the previous commit is the fix; this commit corrects its rationale, records the planned row count in the metadata (sparse_mtp_num_tokens), and has MLAAttention._forward_decode refuse a sparse MTP plan larger than its q -- host ints only, no sync, skipped under DCP, which rebuilds the plan per layer. Tests: tests/model_ops/test_mla_sparse_mtp_plan_rows.py (the refusal, and a plan within q reaching the kernel for both the drafter's unpadded and the target's padded rows); the builder test now also checks the recorded row count. * [MLA] Cut the sparse MTP verify inputs to q's rows for the DCP rebuild Under DCP the sparse_mtp_ plan is rebuilt per layer from paged_cu_seqlens_q, paged_kv_indptr and paged_kv_last_page_lens. paged_kv_indptr was already cut to q's rows (B + 1), but the other two kept the padded running_tokens width, so the drafter's unpadded step 0 could still be planned past q and o. Cut them to B + 1 / B as the non-MTP branch does; the target's padded verify, where B == running_tokens, is unchanged. * [MLA] Scope the sparse_mtp_num_tokens comment to where the check runs The assertion in MLAAttention._forward_decode is skipped under DCP and, being a host check, runs on eager steps and at capture but not on graph replay.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Problem and root-cause candidate
The production fault was reproduced on aiter
3145069b1: 39 real requests padded to 40 fault after about 43 decode steps. Planning only the real sparse-MTP tokens removes the fault, while changing dense metadata does not. Planner dumps showed no negative range: each padding token intentionally receives an empty work item whosekv_start == kv_endis the end of the sparse-index region written in that step. See https://github.com/wuhuikx/LLM-Booster/issues/287#issuecomment-6410896538 and ROCm/ATOM#2715.In that faulting kernel, TDM gathers were already guarded by
has_tokens, but seed setup still unconditionally calledload_page_indices(kv_start)for an empty item. Current main retains that seed load and additionally preloads prologue page indices before checkinghas_tokens. These speculative scalar reads are the only page-index accesses on an empty item and address no live index entry. Metadata rewrites can make the address still less safe on other paths.This patch maps only the physical scalar-load address to index 0 when
kv_start >= kv_end, while preserving the logical empty range. It also avoids current main’s unnecessary empty-item prologue gathers. Non-empty work follows the original instruction path.Validation
batch=1,ctx=64,nhead=16,q_seq=1, error=0)git diff --checkpassedValidation limitation
The available
gfx1250-atom-20261026-dsv41stack does not reproduce the fault with its baseline kernel either. The PS1 implementation has changed substantially since3145069b1, so the 300-step pass is a regression/smoke result, not causal A/B proof. The exact older image is unavailable locally; this PR remains Draft pending the original 39→40 model replay or equivalent validation on the faulting stack.Tracking issue: https://github.com/wuhuikx/LLM-Booster/issues/342