Skip to content

perf(m3): enable prefill CUDA graph + EAGLE3 draft-loop kernel cleanup - #837

Open
syuoni wants to merge 7 commits into
lightseekorg:mainfrom
syuoni:m3-pdl
Open

perf(m3): enable prefill CUDA graph + EAGLE3 draft-loop kernel cleanup#837
syuoni wants to merge 7 commits into
lightseekorg:mainfrom
syuoni:m3-pdl

Conversation

@syuoni

@syuoni syuoni commented Jul 29, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Summary

  • Enable prefill breakable CUDA graph for M3 (22aaf8b): drop --disable-prefill-graph from all M3 recipes/CI configs, and add a tqdm progress bar to prefill bucket capture, matching decode graphs.
  • Replicate M3 embed_tokens (51701c6): M3 overrides resolve_embed to a non-sharded embedding — removes the vocab-mask kernel chain and the embedding allreduce from every forward (same pattern as perf(K2.5): disable embedding TP and skip concat for EAGEL3 drafter loop #94).
  • Share replicated embed correctly with the EAGLE3 draft (e3088de): set_embed_and_head now rebuilds the draft's embedding module (on meta, zero allocation) when the shared target weight's layout differs from the draft's vocab-parallel default; mismatched shapes raise instead of silently corrupting lookups.
  • EAGLE3 norm fusions (50c15c5, 96c0c78): gate fuse_embed_reduce on embed_tokens.tp_size; fuse fc_norm pairs and the midlayer embed/hidden norms via FusedRMSNorm writing into preallocated slices (perf(K2.5): fuse embeds and hidden norm in MLA eagle3 #78 pattern) — eliminates the concat kernels and one standalone RMSNorm per draft step. Applied to llama_eagle3 and the deepseek_v3 Eagle3 model. RMSNorm.forward(inplace=...) is replaced by out= (callers pass out=x).
  • PDL args for M3 qk-norm/RoPE (09197b2).

Notes for reviewers

  • Memory cost: the replicated embedding is ~2.46GB/rank vs 0.61GB sharded at attn-tp4 — the KV pool shrinks accordingly at the same --gpu-memory-utilization.
  • Breaking API: RMSNorm.forward no longer accepts inplace; the only in-tree caller was updated, but out-of-tree callers must switch to out=x.

syuoni and others added 6 commits July 29, 2026 08:16
@syuoni
fix pdl
09197b2 
Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
@syuoni @claude
feat(m3): enable prefill CUDA graph for MiniMax-M3
22aaf8b 
Add tqdm progress bar to PrefillGraph._capture_all_buckets mirroring
the decode graph capture UX. Drop --disable-prefill-graph from all M3
server commands; the flag was required to work around a D2H sync
(.item() call) inside the CuTe MSA extend wrapper that crashed capture,
fixed in lightseekorg#819 (perf(m3): avoid per-layer D2H sync in MSA prefill).

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
@syuoni @claude
perf(m3): replicate embed_tokens to avoid vocab-parallel allreduce
51701c6 
Override resolve_embed in MiniMaxM3Model to use a non-sharded
VocabParallelEmbedding (tp_size=1), eliminating the allreduce on every
embedding lookup. Each rank holds the full 200K-entry table (~2.4 GB vs
~0.6 GB sharded); weight-sharing with the EAGLE3 draft via
set_embed_and_head is unaffected.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
@syuoni @claude
perf(eagle3): gate fuse_embed_reduce on embed_tokens.tp_size
50c15c5 
Replace the mapping.attn.tp_size > 1 check with embed_tokens.tp_size > 1.
The former is redundant (embed tp_size derives from the mapping), and with
non-sharded embeddings (tp_size == 1, e.g. M3 after the resolve_embed
override) it caused an unnecessary allreduce + FillFunctor that scaled
embeds by the TP degree before input_layernorm -- harmless due to RMSNorm
scale-invariance but wasteful (~9 us per draft step on B200 TP4).

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
@syuoni @claude
perf(eagle3): fuse fc_norm and embed/hidden norms, drop RMSNorm inpla…
96c0c78 
…ce arg

Apply the PR lightseekorg#78 FusedRMSNorm pattern to the remaining eagle3 norm+concat
sites:

- llama_eagle3 fc_norm (step 0): chunk pairs run through the fused
  parallel kernel and all chunks write into slices of one preallocated
  buffer -- 3 norm launches + concat become 1 fused + 1 single, no concat.
- llama_eagle3 midlayer (every draft step): input_layernorm(embeds) +
  hidden_norm(hidden) + concat collapse into one FusedRMSNorm launch on
  the non-allreduce-fused arm; the fuse_embed_reduce arm is unchanged.
- deepseek_v3 Eagle3MlaModel fc_norm: same treatment (its midlayer was
  already fused by lightseekorg#78).

RMSNorm.forward gains an out= buffer arg (both backing kernels already
accept strided outputs) and loses inplace -- out=x expresses it; the one
caller (kv_a_layernorm in deepseek_v3) migrated. Zero-token guards skip
the fused launches on idle forwards, matching RMSNorm's internal guard.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
@syuoni @claude
fix(eagle3): rebuild draft embedding to match shared target embed layout
e3088de 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
@syuoni syuoni self-assigned this Jul 29, 2026
@syuoni
syuoni marked this pull request as ready for review July 29, 2026 14:10
@syuoni
syuoni requested a review from a team as a code owner July 29, 2026 14:10
@syuoni syuoni changed the title perf(m3): enable prefill-graph, pdl, disable embed sharding perf(m3): enable prefill CUDA graph + EAGLE3 draft-loop kernel cleanup Jul 29, 2026
chatgpt-codex-connector[bot]
chatgpt-codex-connector Bot reviewed Jul 29, 2026
View reviewed changes

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e3088dea80

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread python/tokenspeed/runtime/models/minimax_m3.py
@syuoni @claude
fix(runtime): gate fuse_embed_reduce on embed_tokens.tp_size
f4e9e01 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@chatgpt-codex-connector chatgpt-codex-connector[bot] chatgpt-codex-connector[bot] left review comments
@yweng0828 yweng0828 Awaiting requested review from yweng0828
@FlamingoPg FlamingoPg Awaiting requested review from FlamingoPg
@wtsung wtsung Awaiting requested review from wtsung
@zhyncs zhyncs Awaiting requested review from zhyncs

At least 1 approving review is required to merge this pull request.

Assignees

@syuoni syuoni

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@syuoni