Guanbao/fix atom mtp memory fault#1571
Draft
gbyu-amd wants to merge 12 commits into
Draft
Conversation
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
…legal memory access (#1560) * [plugin][fix] DSA indexer: single-source decode_lens to fix random illegal memory access _build_indexer's multi-token-decode branch feeds torch.repeat_interleave with repeats=decode_lens and output_size=actual_expanded. With output_size given, repeat_interleave skips the device sync and trusts output_size == repeats.sum(); a violation makes its internal gather index run past the source rows, so the underlying index_select over-reads and triggers an intermittent illegal memory access (faults only when the over-read lands on an unmapped page -> random, data/timing dependent; surfaced late by async_copy_ready_event.synchronize() under async scheduling). The two operands came from different sources: decode_lens was diffed from the GPU query_start_loc, while actual_expanded came from the CPU query_start_loc_cpu. Under async scheduling the GPU/CPU copies can momentarily disagree, so repeats and output_size could mismatch. Derive decode_lens from the same CPU query_start_loc_cpu (copied into the persistent GPU buffer) so decode_lens.sum() == actual_expanded by construction. Same single-source-of-truth principle as the native-path fix in #1559. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * [fix] ReplicatedEmbedding: mask out-of-range ids to fix random illegal memory access Under async scheduling + MTP spec-decode, input_ids can transiently carry the optimistic placeholder token -1 (an unresolved "assumed-accepted" draft/bonus slot: produced in gpu_model_runner as output_token_ids.extend([-1]*...), read back via the proposer's prepare_next_token_ids_padded backup before the deferred correction lands, during batch churn). It flows into BOTH the target input_ids and the shared draft input_ids (prev_sampled_token_ids = next_token_ids). GLM-5.2 uses ReplicatedEmbedding whose forward did a raw F.embedding, so -1 read the row before the [vocab, hidden] table -> random illegal memory access (BFloat16 indexSelectSmallIndex; faults only when the -1 row lands on an unmapped page -> rare, ~52 min). vLLM-native VocabParallelEmbedding masks the same -1 to 0 (get_masked_input_and_mask), which is why native never crashes; ReplicatedEmbedding dropped that safety net. Route replicated_embedding through the existing masked kernel with the full-table range [0, num_rows): out-of-range ids get a zero vector via a masked load (never a raw gather). Bit-identical to F.embedding for every valid token, and matches native's zero for -1; the unverified -1 slots are discarded/corrected by async spec-decode, so no accuracy change. Stays a torch_compile_guard custom op (opaque to inductor). Fixes both the target and the shared draft embedding. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
…stop OOB KV read Sparse (DSA) MLA decode packs KV per token at page_size=1, so last_page_len must be 1 for every seq. The decode / draft / cudagraph-capture paths were feeding the dense per-block kv_last_page_lens into get_mla_metadata_v1 and mla_decode_fwd, which makes the asm kernel compute a per-seq KV extent up to block_size-1 pages past the written sparse-index region -> illegal memory access once the context exceeds index_topk (dense >> sparse). - aiter_mla.py: build sparse work metadata (sparse_kv_indptr + all-1s sparse_kv_last_page_lens) whenever is_sparse and (max_q_len==1 or sparse_decode). The MTP draft does a FRESH sparse build (only_update=False, max_seqlen_qo=1) instead of rebasing the target's dense work_info onto sparse seq lengths (which drove kv_start negative / kv_end past the sparse region). Expose sparse_kv_last_page_lens in prepare_decode and the cudagraph capture paths. - attention_mla.py: non-MTP sparse decode passes sparse_kv_last_page_lens to mla_decode_fwd. - eagle.py: the draft re-points attn_metadata.sparse_kv_last_page_lens to the per-seq all-1s slice (it reuses the target metadata but drops to max_seqlen_q=1). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gbyu-amd
force-pushed
the
guanbao/fix_atom_mtp_memory_fault
branch
from
0dfbac0 to
40de778
Compare
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
3 participants
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