Fix/vllm runtime cudagraph mode - #1792
Merged
XiaobingSuper merged 2 commits intoAug 4, 2026
Merged
XiaobingSuper merged 2 commits into
XiaobingSuper merged 2 commits into
Conversation
Use the frontend's per-forward graph decision so FULL decode can retain overlap while PIECEWISE capture remains safely single-stream. Co-authored-by: Cursor <cursoragent@cursor.com>
Match the repository's Ruff import ordering rules. Co-authored-by: Cursor <cursoragent@cursor.com>
XiaobingSuper
merged commit Aug 4, 2026
540b1d1
into
guanbao/k3_rmsnorm_quant_fusion
20 of 21 checks passed
valarLip
pushed a commit
that referenced
this pull request
Aug 6, 2026
* feat(kimi-k3): fuse RMSNorm activation quant into following projections Port the deepseek-style norm+quant fusion onto the true-MLA base so the standalone activation-quant op emitted before online-quantized (ptpc_fp8) Linear layers is folded into the immediately-preceding RMSNorm. atom/model_ops/linear.py: - Skip post-load work for empty fused shells (weight.numel()==0), e.g. KDA b_proj/f_a_proj after their weights are concatenated into in_proj. - Pad a8w8 (per_Token fp8) preshuffle output up to the CK N-tile when the fused N is not tile-aligned, slicing the padding back off in forward. atom/models/kimi_k3.py: - Add _effective_layer_quant + _RMS_FUSABLE_QUANT_TYPES; route QK-norm through deepseek _fuse_rmsnorm_quant (degrades to bf16 when unquantized). - Thread (fp8, scale) tuples through KimiMLP, KimiFullAttention (q_scale into MLA) and KimiKDAAttention (scale carried across the splitting custom op). - Wire KimiDecoderLayer input_layernorm / post_attention_layernorm fusion, gated per layer so non-fusable schemes fall back to today's bf16 path. Every fusion is self-gated; validated by CPU unit tests. fp8 numerics still need one gsm8k run on GPU. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix(kimi-k3): enable online quant for compressed-tensors checkpoints QuantizationConfig only sets online_quant=True when quant_method is in a fixed allow-list. Kimi-K3's checkpoint declares quant_method "compressed-tensors", which was absent, so --online_quant_config was silently ignored: ptpc_fp8 never activated and every target Linear ran bf16 (gemm_a8w8_bpreshuffle was never called). Add "compressed-tensors" to the allow-list so online fp8 quant engages as intended. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * test(kimi-k3): realign fusion tests with _effective_layer_quant The fusion implementation was renamed to _effective_layer_quant / _kda_packed_modules_mapping, but the tests still referenced the old _layer_effectively_consumes_per_token_fp8 / _kimi_packed_modules_mapping helpers, leaving the branch red (10 failures on HEAD). Update the tests to the current helpers and per-consumer fusion gating, and modernize a few Optional[X] annotations to X | None in kimi_k3.py (+ exec noqa in the linear alignment test). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * feat(kimi-k3): dual-stream MoE overlapping shared-expert GEMMs Overlap shared-expert GEMMs (on alt_stream) with the routed-expert path on the main stream. Both TP all-reduces stay serial on the main stream to avoid concurrent collectives on the single TP communicator. GPU-validated on MI355 TP8 at gsm8k exact_match 0.9555 (no regression vs fp8 baseline). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * opt k3 * Update atom/model_ops/linear.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * rm local test files * update recipe * make ar parallel on two streams * small update * change ar launch order to ensure better overlap * fix online quant with plugin * Fix/vllm runtime cudagraph mode (#1792) * fix(plugin): gate dual-stream MoE by runtime graph mode Use the frontend's per-forward graph decision so FULL decode can retain overlap while PIECEWISE capture remains safely single-stream. Co-authored-by: Cursor <cursoragent@cursor.com> * style: sort forward context imports Match the repository's Ruff import ordering rules. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: xytpai <xytpai@foxmail.com> Co-authored-by: Cursor <cursoragent@cursor.com> * correct rope_max_position * update linear.py * clean comments * update vllm plugin recipe * fix format --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: XiaobingZhang <xiaobingzhangupc@gmail.com> Co-authored-by: xytpai <xytpai@foxmail.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: ganyi <ygan@amd.com>
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
2 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