Repository navigation
mtp use decode_update_mla_metadata_v1 instead of get_mla_metadata_v1 - #287
Merged
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates the MTP (Multi-Token Prediction) speculative decoding flow so that, after the initial MLA metadata setup, subsequent metadata updates during decode steps use the more efficient decode_update_mla_metadata_v1 function instead of the full get_mla_metadata_v1 recomputation. For prefill-phase steps, the original get_mla_metadata_v1 is still used.
Changes:
- Introduces a
do_attn_metadata_updateflag ineagle.pyto differentiate between decode-phase (update-only) and prefill-phase (full recompute) MTP iterations. - Extends
set_mla_persistent_worker_buffersandprepare_mtp_decodeinaiter_mla.pywithonly_updateandnum_reject_tokensparameters to select betweendecode_update_mla_metadata_v1andget_mla_metadata_v1. - Adds
decode_update_mla_metadata_v1to theaiterimports.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
atom/spec_decode/eagle.py |
Computes do_attn_metadata_update flag and forwards num_reject_tokens (only on first iteration) to prepare_mtp_decode |
atom/model_ops/attentions/aiter_mla.py |
Adds decode_update_mla_metadata_v1 import, extends set_mla_persistent_worker_buffers and prepare_mtp_decode with only_update/num_reject_tokens parameters, and branches between the two metadata functions |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
valarLip
requested changes
Mar 10, 2026
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
valarLip
approved these changes
Mar 10, 2026
Jasen2201
pushed a commit
to Jasen2201/ATOM
that referenced
this pull request
Apr 10, 2026
…OCm#287) * mtp use decode_update_mla_metadata_v1 instead of get_mla_metadata_v1 * update * use q_dtype kv_dtype both from config.kv_cache_dtype * add dtype_q to metadata
qichu-yun
added a commit
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
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
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 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.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist