[FlyDSL] Add full-expert DeepSeek V4.1 Flash MegaMoE tuning for gfx950

[Fangzhou-Ai]

Summary

Enable DeepSeek V4.1 Flash MegaMoEV2 on gfx950 for full expert parallelism: H=5120, I=2304, E=384, top-k=6, with 192 complete experts/rank at EP2 or 96 at EP4. Fuse peer dispatch with W13, clamp-10 SwiGLU and FP8 intermediate quantization, followed by MXFP4 W2 and weighted peer combine. Shared experts remain separate.

This extends the now-merged #6127 and remains a draft for implementation and tuning review. AI assistance (Codex) was used for implementation, benchmarking and this description.

Duplicate checks: searched open ROCm/aiter PRs for MegaMoE/Flash and related EP work. #6127 supplies the underlying MegaMoEV2 improvements. #5937 targets TP A4W4; #5562 supplies standalone EP A8W4 GEMM tuning with a different top-k; #5967 supplies CK TP2 tuning. This PR adds full-expert Flash dispatch/W13/W2/combine integration and a distinct per-M configuration policy.

Changes and tuning

• Extend fixed-slot dispatch to Flash's 192/96 complete local experts. Workspace-capacity limits are 512 at EP2 and 1024 at EP4; larger workspaces use compact dispatch. Removed AITER_MEGA_FIXED_SLOT_MAX_MTPR and the standalone tuning file from the submitted diff.
• Add a deduplicated, measured gfx950 Flash configuration table. Cover every M=6*c for c=1..192, including the fixed/compact boundary and M1152, at both EP sizes. M is the common maximum source-token count/rank, not received expert routes; balanced DP can have a smaller per-rank M. Include M1..64 and prefill powers through the serving workspace capacity.
• Compact transport is tuned separately for capacities 8192 and 16384. Flash preserves those capacity classes in configuration lookup, rather than collapsing them to the generic 32768 class. A check of every measured entry verifies that the public selector and compiled bundle select the submitted configuration. Other models retain their existing bucket/capacity policy; arbitrary Flash M uses the nearest supported source bucket, with ties selecting the larger bucket.
• Search W13 M16/32/64/128, N64/128/256/512, 2/4/8 waves, MFMA orientation, async copies, weight pipelining, A swizzle, B cache policy, dispatch CU count, grid multiplier, work shards and waves-per-EU. Search W2 M/N tiles, persistent CU count, non-temporal loads, strided scheduling, B hoisting, scale prefetch, spatial partitioning and BF16 LDS staging. Also search combine block/wave counts and compact transport chunk size/preparation CU count. Check every timed candidate against the reference configuration.
• An additional exploratory EP2 compact-grid search rejected register spills and grids exceeding the HIP-reported resident workgroup budget on any rank. No unconfirmed selections from that exploration were retained. All submitted selections were compared against inherited configurations in independent reversed timing passes. The configurations are the best measured within the completed finite search space; they are not asserted to be a global optimum.
• Support BM16 scale access in both GEMMs and dense BM16 scale access in W2 when sort/GEMM tiles differ. Let a shared workspace bind the current layer's prepared weights without copying them. config_tokens supplies one source count on all peers, including empty ranks.
• Benchmark graphs place the peer barrier before timing events. Cache flushing and synchronization remain outside timing. Isolated W2/combine measurements regenerate Stage1 before every replay because its peer handshake protects reuse of the combine buffers.

Complete performance and occupancy data

644 source-M/workspace performance cases · All kernel/rank occupancy and resource records · Measurement methodology

The attached CSVs include every selected table entry, deduplicated config IDs, full configurations, independent timings, actual launch sizes, per-rank VGPR/AGPR/SGPR/LDS and spill counts, theoretical residency and achieved rank-0 occupancy.

EP	Workspace	Path	Measured source shapes	Unique configs	Speedup range
2	512	fixed	142	93	1.018–1.268x
2	8192	compact	111	81	1.011–1.309x
2	16384	compact	112	90	1.018–1.314x
4	1024	fixed	228	143	1.141–2.380x
4	8192	compact	25	25	1.437–1.877x
4	16384	compact	26	23	1.511–1.892x

Performance measurements: 8×MI355X/gfx950, isolated EP groups, BF16 inputs, MXFP4 weights, top-6 routing across 384 experts. HIP graph replay with a 512 MiB cache flush and five warmups before sampling. For each sample take the maximum time across ranks, then its median. MegaMoE uses two reversed/shuffled passes of 30 samples; the MoRI+AITER baseline uses 60. The table reports the mean of the two pass medians. Source/received token counts are recorded rather than assumed equal.

Decode comparisons use MoRI's 64-block/4-wave geometry and static uniform-DP receive trim; prefill powers M≥2048 use its 128-block/16-wave geometry and full receive capacity. The CSV explicitly records the baseline phase. These kernel comparisons exclude the caller's final output copy and shared-expert computation. Hardware-counter timings are excluded from performance results.

Occupancy and spills

Resource occupancy comes from hipModuleOccupancyMaxActiveBlocksPerMultiprocessor queried on exported HSACO using each kernel's actual workgroup size and zero dynamic LDS. MI355X has 256 CUs and 32 wave slots/CU. The launch ceiling additionally limits residency by the launched workgroup count. These are ceilings, not useful MFMA activity. AMD occupancy guide.

Achieved occupancy uses ROCprofiler SDK's gfx950 OccupancyPercent counter on rank 0 of each EP job. Report the median of ten cold-cache replays: five measured replays following five warmups in each of two passes. Capture-time executions and warmups are excluded. Counter collection serializes dispatches and can alter overlap; its occupancy is diagnostic rather than an uninstrumented serving measurement. Every rank has separate static metadata/residency checks; achieved runtime percentages are measured only on rank 0 and labeled as such in the attachment.

All 842 distinct selected kernel/rank combinations have 0 VGPR spills. 30 compact W2 kernel/rank combinations report 2 SGPR spills; these were retained based on measured performance across ranks, rather than assuming that a spill-free variant is always faster. Maximum selected VGPR count: 394. The attachment reports the actual register, LDS and spill counts for every rank/configuration; no blanket zero-spill claim is made.

Rejected layouts include W13 M32/N512/4 waves (11 spilled VGPRs), M64/N512/4 waves (100), and M128/N512/8 waves (66, with 151,600 B LDS). Raising waves-per-EU to 4 also introduced spills and was excluded from selection. Higher occupancy alone was not used as the selection criterion.

Representative speculative-decode source bounds are below. For the compact cases this view uses capacity 8192; capacity 16384 and every intermediate M are included in the attachments. W13 and W2 occupancy columns show resource ceiling / launch ceiling / achieved rank-0 median (%).

EP	c	M=6c/rank	Path	Mega (µs)	MoRI+AITER (µs)	Speedup	W13 occupancy	W2 occupancy
2	1	6	fixed	217.4	258.3	1.188x	37.5/14.1/11.6	25.0/25.0/15.9
2	4	24	fixed	442.6	499.1	1.128x	37.5/18.8/12.3	62.5/62.5/51.6
2	16	96	fixed	690.7	786.0	1.138x	37.5/18.8/12.7	62.5/62.5/54.0
2	64	384	fixed	884.2	944.8	1.069x	25.0/25.0/23.8	50.0/50.0/43.1
2	128	768	compact	1035.0	1069.0	1.033x	25.0/25.0/23.5	37.5/37.5/29.7
2	192	1152	compact	1191.4	1544.8	1.297x	25.0/25.0/23.3	12.5/12.5/10.0
4	1	6	fixed	201.0	246.7	1.228x	37.5/14.1/10.7	50.0/50.0/29.2
4	4	24	fixed	356.6	497.9	1.396x	37.5/18.8/12.2	62.5/62.5/49.2
4	16	96	fixed	433.0	601.0	1.388x	37.5/18.8/12.5	50.0/50.0/37.5
4	64	384	fixed	549.5	851.1	1.549x	25.0/25.0/22.7	50.0/50.0/40.5
4	128	768	fixed	721.7	891.1	1.235x	25.0/25.0/22.8	25.0/25.0/19.9
4	192	1152	compact	929.1	1429.6	1.539x	25.0/25.0/23.0	12.5/12.5/10.3

Large-M prefill comparisons and occupancy at both serving capacities

EP	Capacity	M/rank	Mega (µs)	MoRI+AITER prefill (µs)	Speedup	W13 occupancy	W2 occupancy
2	8192	2048	1711.4	1952.1	1.141x	25.0/25.0/22.9	12.5/12.5/9.7
2	8192	4096	2805.8	3098.6	1.104x	25.0/25.0/22.9	12.5/12.5/9.5
2	8192	8192	5428.1	5487.6	1.011x	25.0/25.0/23.5	12.5/12.5/9.6
2	16384	2048	1710.2	2177.8	1.273x	25.0/25.0/22.7	12.5/12.5/9.6
2	16384	4096	2796.1	3339.7	1.194x	25.0/25.0/22.9	12.5/12.5/9.5
2	16384	8192	5430.8	5747.2	1.058x	25.0/25.0/23.8	12.5/12.5/9.6
2	16384	16384	10734.4	10930.9	1.018x	25.0/25.0/24.1	12.5/12.5/9.6
4	8192	2048	1245.4	1875.8	1.506x	25.0/25.0/23.3	37.5/37.5/31.6
4	8192	4096	2049.6	2975.8	1.452x	25.0/25.0/23.6	12.5/12.5/10.5
4	8192	8192	3677.3	5282.5	1.437x	25.0/25.0/23.9	12.5/12.5/10.5
4	16384	2048	1250.4	2366.1	1.892x	25.0/25.0/23.5	37.5/37.5/31.4
4	16384	4096	2059.2	3450.6	1.676x	25.0/25.0/23.7	12.5/12.5/10.5
4	16384	8192	3667.7	5804.1	1.582x	25.0/25.0/23.8	12.5/12.5/10.5
4	16384	16384	6959.7	10518.7	1.511x	25.0/25.0/24.0	12.5/12.5/11.2

Correctness and model evaluation

• All 644 source-M/workspace cases pass the independent dequantized BF16 MXFP4 reference with clamp-10 SwiGLU and weighted routing, and HIP graph replay matches eager output. Relative L2 threshold is 0.10 for the quantized path; tuning candidates use a stricter 0.001 reference-config threshold.
• Per-layer zero-W2 substitution produces zero output; restoring weights restores the earlier result exactly. Boundary fan-out, expert/sort padding boundaries, invalid routes and invalid whole-row tails, hot-rank routing, unequal local batches and empty ranks were checked. Extended skew/padding checks are included in the final validation logs.
• EP2 repeated-shape testing previously passed 345 transitions, 10,000 eager forwards and 10,000 graph forwards after correcting the isolated-stage benchmark's missing Stage1 preparation.
• Actual Flash EP2 serving with real DSpark rejection: 1286/1319 GSM8K correct = 97.498%, measured before the additional M6×/compact policy tuning. This is not an accuracy rerun of the extended policy. The new configurations retain the same arithmetic and pass the independent BF16/graph checks. Evaluation uses the published InferenceX five-shot multiturn chat prompts and strict final-answer filter, concurrency 64, max output 12288, temperature 0, top-p 1, context 1,048,576, capture 1024, batch limit 8192 and five probabilistically sampled draft tokens with real block rejection. Published TP reference: 1281/1319 = 97.119%. The earlier 0.908 run used the wrong protocol and is excluded. Reference evaluation job.
• Companion vLLM full-local-expert loader tests: 4 passed. The serving integration is currently local and needs separate review.

Reproduction and checks

The exhaustive tuning/resource-export harness was run locally and is excluded from this PR. The submitted benchmark and correctness suite reproduce the selected configurations. In an environment with ROCm PyTorch, FlyDSL, MoRI and AITER native dependencies:

export PYTHONPATH="$PWD"
export AITER_JIT_DIR=/tmp/aiter-mega-jit
export CK_DIR=/usr/local/lib/python3.12/dist-packages/aiter_meta/3rdparty/composable_kernel
export OMP_NUM_THREADS=1 TORCH_NCCL_BLOCKING_WAIT=1

# Speculative decode example; use four GPUs and mtpr=1024 for EP4.
HIP_VISIBLE_DEVICES=0,1 /home/fai/vllm/.venv/bin/python -m torch.distributed.run \
  --standalone --nproc-per-node=2 op_tests/multigpu_tests/bench_mega_moe_v2.py \
  --network v4_1_flash --tokens 384 --mtpr 512 --phase decode \
  --iters 60 --cold-cache-mb 512

# Prefill example using the actual serving workspace and prefill baseline.
HIP_VISIBLE_DEVICES=0,1 /home/fai/vllm/.venv/bin/python -m torch.distributed.run \
  --standalone --nproc-per-node=2 op_tests/multigpu_tests/bench_mega_moe_v2.py \
  --network v4_1_flash --tokens 8192 --mtpr 8192 --phase prefill \
  --iters 60 --cold-cache-mb 512

HIP_VISIBLE_DEVICES=0,1 /home/fai/vllm/.venv/bin/python -m torch.distributed.run \
  --standalone --nproc-per-node=2 op_tests/multigpu_tests/test_mega_moe_v2.py \
  --network v4_1_flash --bs-list 6,24,96,384,510,512 \
  --max-tok-per-rank 512 --accuracy-max-bs 512

All submitted Python changes pass Black and Ruff; git diff --check passes. Final validation covers every measured table entry, and exported HSACO/resource and launch metadata is cross-checked as described in the attachments.

EP2 serving sweep versus published TP2

AIPerf AgentX MVP, 393 trace entries, 3600-second profiling duration plus the reference warmup, four simultaneous two-GPU servers at concurrency 1/4/16/64. Same attempt-3 workload and serving recipe, with TP2 replaced by TP1+DP2+EP2 and MegaMoE. Synthetic DSpark acceptance is 3.51 in performance runs; GSM8K above uses real rejection. Benchmark and aggregation code use the exact reference revisions. Modified local vLLM/AITER sources are used rather than claiming the unchanged reference container image.

All rates are tokens/s for a complete two-GPU server. p90 interactivity = 1/p90 TPOT. Total throughput counts input plus output tokens; output throughput is provided separately. Changes are EP2/TP2−1.

Concurrency	TP2 p90 intvty	EP2 p90 intvty	Change	TP2 output tput	EP2 output tput	Change	TP2 total tput	EP2 total tput	Change
1	308.63	240.26	-22.2%	153.58	135.89	-11.5%	22509.16	18920.91	-15.9%
4	268.80	181.24	-32.6%	228.53	212.13	-7.2%	33173.11	29947.14	-9.7%
16	134.72	62.18	-53.8%	686.17	593.20	-13.5%	108214.84	83799.91	-22.6%
64	31.51	13.71	-56.5%	1998.44	1333.35	-33.3%	256473.59	160710.55	-37.3%

The reference aggregation excludes warmup and error records. Error accounting is retained below, including the published TP2 error counts.

Concurrency	Profiled requests	Measured span (s)	EP2 warmup errors	EP2 profiling errors	TP2 warmup errors	TP2 profiling errors
1	258	3562.2	0	0	0	0
4	598	3603.9	0	0	0	0
16	2231	3624.3	0	0	0	1
64	5041	3626.8	2	0	2	2

Serving outcome and additional accounting

All four EP2 runs trail the published TP2 result on p90 interactivity, output throughput and total throughput. The isolated MoE speedups do not establish an end-to-end serving win. The vLLM/AITER build differs from the published reference image; a same-build local TP2 run and serving instrumentation would help isolate the cause. Prefill/decode interference, DP synchronization and graph-replay fallback are hypotheses, not measured explanations.

The benchmark's final drain cancelled 1/0/0/10 unfinished profiling requests at concurrency 1/4/16/64 respectively. These cancellations do not have successful profiling records and are excluded from aggregation. All four runs report zero profiling request errors. C64 additionally has two warmup InvalidInferenceResultError records, as does the published TP2 C64 warmup.

DSpark settings were checked against the exported TP2 reference configuration and actual startup logs for all four EP2 servers: five draft tokens, probabilistic drafting, synthetic rejection with acceptance length 3.51, and adaptive verification disabled. Exported counters show exactly five draft tokens per draft. Mean acceptance including the bonus token is TP2 C1 3.510170; EP2 C1/C4/C16/C64 3.510312/3.509701/3.510071/3.509569. Thus neither a DSpark configuration mismatch nor a measured acceptance-rate mismatch explains the gap. These synthetic performance counters are separate from the real-rejection GSM8K evaluation.

No KV-cache preemptions occurred in any EP2 sweep; peak observed usage is below 20% per engine at C16 and approximately 20.7% at C64. These observations exclude KV capacity exhaustion for these runs; they do not establish the remaining bottleneck.

Complete serving comparison CSV · Phase completion/cancellation accounting · DSpark configuration check · Runtime acceptance counters

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6154 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.