Repository navigation
[ROCm][DSv4] AITER MegaMoEV2 Integration For DeepSeek V4 - #59685
Conversation
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
|
Great job @micah-wil The code LGTM. and we absolutely need some UTs to gate test these changes. |
|
Does this kernel support both TP and EP settings? |
|
@Fangzhou-Ai Right now it requires TP=1, but we could work on TP>1 in a follow-up. |
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
|
✅ Triggered Buildkite CI #93167 for commit |
CI selector (shadow): 75 test steps (137 jobs) instead of 87 (128 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (75)
Would skip (today's rules run them) (41)
Would add (today's rules do not run them) (29)
AMD mirrors: would skip (0)none AMD mirrors: would add (82)
10 changed files · base |
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
|
/ci run |
|
✅ Triggered Buildkite CI #93193 for commit |
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
|
/ci run |
|
✅ Triggered Buildkite CI #93205 for commit |
…nabled Signed-off-by: Micah Williamson <micah.williamson@amd.com>
|
/ci run |
|
✅ Triggered Buildkite CI #93219 for commit |
|
hi @AndreasKaratzas @shen-shanshan @tjtanaa can we merge this mega kernel PR? |
tjtanaa
left a comment
There was a problem hiding this comment.
LGTM. Thanks for addressing the comments that we discussed in the sig.
- Update the image from nightly-rocm100-ac9126e58aa7bbab1856ba6593ba4d5003fea516 to nightly-rocm100-7d0b4e57aac4c4323225b0eee9ae960ce75ffffe, the first ROCm 10.0 nightly with the DeepSeek-V4.1 mono decode kernels of vllm-project/vllm#60397, and update model.container to match. The image also brings vllm-project/vllm#59685, #56638, #60159 and #60389. - Enable the mono decode kernels (VLLM_ROCM_MONO_DECODE=1) at TP2 and TP4: a decode step of up to 48 rows, i.e. concurrency 1-8 with five DSpark drafts, runs each backbone layer as two persistent launches (30 of 40 layers whole, the other 10 the FFN launch only). Larger steps take the regular path, so C16-C64 are expected unchanged. - Enable VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4 (vllm-project/recipes#1061) for tensor-parallel all-reduces above the quick-reduce size threshold. - Restore --max-num-batched-tokens 16384 at TP2 c64: vllm-project/vllm#58014 removed the profiling over-reservation the 8192 cap worked around. The prefill chunk is now 16384 at every point and the ladder stays at concurrency 1-64. - No data-type or precision change to the DSpark draft: vllm-project/vllm#60397 skips the draft layers (layer id >= num_hidden_layers), so the bundled DSpark draft of deepseek-ai/DeepSeek-V4.1-Flash runs on vLLM's regular path as shipped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…D][MI355X] DSv4.1 Flash:启用 vLLM mono decode 内核(镜像待定) (#3690) * perf(dsv41flash): repin MI355X to the 2026-10-06 ROCm 10.0 nightly Repin the MI355X DeepSeek-V4.1-Flash DSpark recipe to nightly-rocm100-21d93d0d and retune the ladder around it. - Leave VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 unset. The a4w4 opt-in faults during graph capture at TP=2; all nine TP=2 points failed and all nine TP=4 points passed in run 37424787697. vllm-project/vllm#60273 tracks the fix, so the experts stay on the generic AITER CK a8w4 selector. - Enable VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4 so tensor-parallel all-reduces above the quick-reduce size threshold use the INT4 codec. - Hold both TP rows at concurrency 64. Run 37510565585 measured TP=2 c128 at 150,683 tok/s/chip against 151,896 at c64 with prompt tokens within 0.5%, so the extra concurrency bought real prefill rather than scored tokens. - Set VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD=1024. Upstream caps the shared-expert overlap at 256 tokens, which predates speculative decoding: five DSpark drafts make a decode step CONC x 6 tokens, so c64 submits 384 and the overlap never engages above c42. 1024 covers the ladder and matches the ATOM default, and stays below the prefill chunk so memory profiling never takes the overlap. - Restore the upstream 16384 prefill chunk at TP2 c64, making the chunk uniform across all 14 points. That point had capped it at 8192 to buy KV room, but the room was taken by a profiling bug rather than by the chunk: the ROCm paged MXFP4 indexer sizes its decode logits workspace as (rows, max-model-len), and with no ambient vLLM config in scope the row count fell back to the batched-token count, so at this recipe's 1M context the reservation scaled linearly with the chunk. vllm-project/vllm#58014 bounds rows by max-num-seqs x (1 + 5 drafts), and the repinned image carries that fix. Rebased onto main: the configuration-procedure guides this branch had edited were removed by #3789, so those edits are dropped. Co-authored-by: Cursor Agent <cursoragent@cursor.com> * perf(dsv41flash): enable mono decode on MI355X; image TBD - Set VLLM_ROCM_MONO_DECODE=1 (vllm-project/vllm#60397): two persistent launches a layer for decode steps of up to 48 rows; larger steps take the regular path. - Image TBD until the first ROCm 10.0 nightly with #60397 is published. - Drop the a4w4 notes and shorten the comments and the changelog entry. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(dsv41flash): drop the shared-expert stream threshold Keep the upstream default for simplicity. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(dsv41flash): pin the ROCm 10.0 nightly that carries mono decode Replace the TBD placeholder with nightly-rocm100-7d0b4e57, the first ROCm 10.0 nightly built after vllm-project/vllm#60397 merged, so the VLLM_ROCM_MONO_DECODE=1 setting on this recipe has kernels to run. 将 TBD 占位符替换为 nightly-rocm100-7d0b4e57,这是 vllm-project/vllm#60397 合并后构建的首个 ROCm 10.0 nightly,使本 recipe 中的 VLLM_ROCM_MONO_DECODE=1 设置有对应的 kernel 可用。 Co-authored-by: Cursor <cursoragent@cursor.com> * docs(dsv41flash): align MI355X vLLM changelog entry with the #3690 diff - Update the image from nightly-rocm100-ac9126e58aa7bbab1856ba6593ba4d5003fea516 to nightly-rocm100-7d0b4e57aac4c4323225b0eee9ae960ce75ffffe, the first ROCm 10.0 nightly with the DeepSeek-V4.1 mono decode kernels of vllm-project/vllm#60397, and update model.container to match. The image also brings vllm-project/vllm#59685, #56638, #60159 and #60389. - Enable the mono decode kernels (VLLM_ROCM_MONO_DECODE=1) at TP2 and TP4: a decode step of up to 48 rows, i.e. concurrency 1-8 with five DSpark drafts, runs each backbone layer as two persistent launches (30 of 40 layers whole, the other 10 the FFN launch only). Larger steps take the regular path, so C16-C64 are expected unchanged. - Enable VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4 (vllm-project/recipes#1061) for tensor-parallel all-reduces above the quick-reduce size threshold. - Restore --max-num-batched-tokens 16384 at TP2 c64: vllm-project/vllm#58014 removed the profiling over-reservation the 8192 cap worked around. The prefill chunk is now 16384 at every point and the ladder stays at concurrency 1-64. - No data-type or precision change to the DSpark draft: vllm-project/vllm#60397 skips the draft layers (layer id >= num_hidden_layers), so the bundled DSpark draft of deepseek-ai/DeepSeek-V4.1-Flash runs on vLLM's regular path as shipped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Cursor Agent <cursoragent@cursor.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: seungrokj <144636725+seungrokj@users.noreply.github.com> Co-authored-by: seungrokj <seungrok.jung@amd.com>
| @@ -0,0 +1,154 @@ | |||
| # SPDX-License-Identifier: Apache-2.0 | |||
There was a problem hiding this comment.
I think we should strive to use vLLM's linear/MOE oracles, and avoid adding model-specific MOE backends per-model.
This could very well benefit other models, and may not be deepseek-specific.
If necessary, the oracle could be extended to have per-model supported backends.
There was a problem hiding this comment.
Hi @fxmarty-amd, yeah you are right. This could also benefit GLM 5.2. I'm working on pulling it out into the oracle here main...ROCm:vllm:micah/aiter-mega-moe-oracle. Just needs more cleanup and testing. I used the initial DSV4-specific impl of FlashInfer MegaMoE as an example for this one #49636, but am now looking to copy the generalized/oracle-based version of FI MoE that landed a couple weeks ago #54049. Once it's ready I'll send it your way if you want to take a look. Thanks!
Bump dsv4-fp4-mi355x-vllm-agentic-mtp to vllm/vllm-openai-rocm:nightly-fe4df3df0eb703862c686b256e9fe44fa94c8945, which includes the AITER MegaMoEV2 integration for DeepSeek-V4 (vllm-project/vllm#59685). Add DEP8 points with --moe-backend aiter_mega_moe, expert parallelism and the MoRI high-throughput all2all at concurrency 32-256, beside the unchanged TP8 and EP1 DEP8 bands. MegaMoE shared-expert fusion stays off because it requantizes the FP8 shared experts of the DSpark draft layers to MXFP4. 将 dsv4-fp4-mi355x-vllm-agentic-mtp 的镜像更新为 vllm/vllm-openai-rocm:nightly-fe4df3df0eb703862c686b256e9fe44fa94c8945,该镜像包含 DeepSeek-V4 的 AITER MegaMoEV2 集成(vllm-project/vllm#59685)。新增并发 32-256 的 DEP8 测试点,使用 --moe-backend aiter_mega_moe、专家并行与 MoRI high-throughput all2all,原有 TP8 与 EP1 DEP8 测试点保持不变。MegaMoE 的 shared-expert fusion 保持关闭,因为它会将 DSpark draft 层的 FP8 shared experts 重新量化为 MXFP4。 Co-authored-by: Cursor <cursoragent@cursor.com>
Description
This PR integrates AITER's MegaMoEV2 for DeepSeek-V4. The new backend is opt-in via

--moe-backend aiter_mega_moe. MegaMoEV2 is a FlyDSL implementation that fuses the routed experts (EP dispatch, MXFP4 GEMM1 + SwiGLU, GEMM2, combine). In this example config I profiled, it boils down into four kernels per layermegamoe_prepare,megamoe_stage1,megamoe_stage2,ep_combine_intranode.The current implementation of
--moe-backend aiterinstead had 10 kernels plus an extra__amd_rocclr_copyBufferin the config that I profiled (MoRI EP) for comparison.From my testing, MegaMoEV2 requires >4GB for
MORI_SHMEM_HEAP_SIZE, so I have set the default to 8G in this PR. MoRI has a descriptive error message if you override this to a smaller value.Usage
Requirements (enforced in the code):
--enable-expert-parallel, TP=1, PP=1)Example server:
VLLM_ROCM_USE_AITER=1 MORI_SHMEM_HEAP_SIZE=8G vllm serve /models/DeepSeek-V4-Pro --data-parallel-size 8 --enable-expert-parallel --all2all-backend mori_high_throughput --moe-backend aiter_mega_moePerformance and Accuracy
I've evaluated lm_eval scores as well as performance for several shapes with and without DSpark. All tests used the following serve commands (with
--moe-backendvaried accordingly):No DSpark:
VLLM_ROCM_USE_AITER=1 MORI_SHMEM_HEAP_SIZE=8G vllm serve /models/DeepSeek-V4-Pro --tokenizer-mode deepseek_v4 --data-parallel-size 8 --enable-expert-parallel --all2all-backend mori_high_throughput --no-enable-prefix-caching --port 8000 --moe-backend {aiter, aiter_mega_moe}DSpark (
num_speculative_tokens=3):lm_eval:
No Dspark:
DSpark:
Performance
Bench command for all perf runs:
No spec decode (Baseline vs. MegaMoE)
DSpark w/ 3 spec tokens (Baseline vs. MegaMoE)