You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[FlyDSL] [Feature] Precompile MoE deployments (EP and TP), described by model and deployment knobs - #7063
Follows #6362, which builds the MoE and MegaMoE AOT caches by replaying the runtime on CPU.
Summary
A caller behind an expert-parallel dispatch passes fused_moe an expert_mask, which brings in FlyDSL launches no tuned row reaches: the gfx950 stage-2 reduction, compiled for the mask's length (moe_reduction_..._m1e<len>), and the gfx1250 route (moe_g2l_lut, moe_route_g2l_lds) and bf16 gather-reduce. The mask's length follows from the deployment (global experts, EPLB replicas, fused shared experts), which no tuned row records, so the AOT build never compiled these kernels and EP serving JIT-compiled them on the request path.
This PR describes those deployments in aiter/aot/flydsl/moe_deployments.json: each MoE model's architecture once, as its config.json gives it, and each deployment by its knobs (GPU, quant, EP size, EPLB redundant experts, shared-expert fusion, token range). aiter/aot/flydsl/deployments.py validates the file and derives what each rank hands the runtime, with the convention ATOM and vLLM share. From it:
the MoE replay (moe.py) precompiles every fused_moe deployment's masked calls;
the MegaMoE replay (mega_moe.py) takes its DeepSeek-V4-Pro and Kimi-K3 deployments from the same file instead of Python constants;
op_tests/test_moe_ep.py -t aot_ep_profiles runs each deployment through the EP fused_moe call in FlyDSL run-only mode, so a kernel missing from the cache fails the test.
The third commit extends the file to tensor-parallel callers. They pass no mask, but their call still depends on what no tuned row records: how the framework runs the shared expert (fused, apart, or through FHMoE). It registers DeepSeek-V4-Pro's two TP8 calls on gfx950 as the first TP deployments.
After a merge of main, two commits address review: a deployment's replay passes its model's activation clamp, the EP check routes and checks the call as serving makes it, a MegaMoE deployment must be on MI355X (gfx950, 256 CUs), and an EP check whose explicit deployment selection runs no case fails instead of passing.
The last commit trims tests/test_aot_drivers.py to the checks no other test makes, and drops the replay of a SiTUv2 model's own betas, which no registered model reaches: Kimi-K3's (4, 25) are among the betas every SiTUv2 row replays. The default deployments expand to the same replay jobs.
Motivation
EP serving compiled on the request path. ATOM and vLLM pass expert_mask on gfx950 and gfx1250 whenever expert parallelism is on. On this head, the 13 default deployments need 98 FlyDSL artifacts the tuned-row replay does not build.
Per-rank numbers are deployment facts, not model facts. The same model at the same EP size passes different shapes depending on EPLB redundancy (DeepSeek-V4-Pro EP8: 48/52/56 local experts with a 385/417/449-entry mask) and shared-expert fusion (GLM-5.3-Flash EP4: 73 experts, 10 routing columns and a 290-entry mask when fused; 72/8/289 when not), and each value is compiled into some kernel. Recording the knobs and deriving the numbers keeps them consistent and reviewable.
One place for model facts. DeepSeek-V4-Pro's and Kimi-K3's dims, top-k, swiglu_limit and SiTUv2 betas were Python constants in mega_moe.py; both replays and the test now read the same entries.
Tuned tables keep deciding which kernels the runtime picks. A deployment is replayed against the tables serving reads: it uses a tuned row where one matches and the runtime's default heuristic where none does, so it needs no tuned row.
aiter/aot/flydsl/deployments.py (new): loads and validates deployment files, and derives each rank's call. With P the routed plus redundant experts and S the fused shared experts, a rank passes w1.shape[0] = P / EP + S, topk_ids.shape[1] = top-k (top-k + S + 1 with fused shared experts, the last column a sentinel that masks them on other ranks' tokens), and a mask of P + S + 1 entries. Unknown or missing fields, experts that do not split over EP, fusing a model without shared experts, or an id listed twice fail the load. Files given together share their models, so a file of your own can name the default file's models. python -m aiter.aot.flydsl.deployments lists what the files expand to.
aiter/aot/flydsl/moe.py: replays each fused_moe deployment with its mask at both ends of every get_padded_M bucket up to max_tokens (bisecting between them, as for rows), with and without num_local_tokens (MORI dispatch passes it, the all-gather path does not), with the activation dtype pinned to the deployment's quant as serving pins it. --deployments selects deployment files; with neither --csv nor --deployments, both defaults are replayed, as run_aot (the wheel build) does.
aiter/aot/flydsl/mega_moe.py: jobs come from the mega_moe deployments, and the job list is unchanged (192 jobs). --mtpr, --quant and --experts-per-rank now narrow the deployments down (--experts-per-rank 48 keeps working), --deployments takes other files, and the shape flags (--topk, --model-dim, --inter-dim, --world-size, --act, --situ-beta, --situ-linear-beta, --swiglu-limit) are removed in favor of a deployment file.
op_tests/test_moe_ep.py:
-t aot_ep_profiles runs every fused_moe deployment of the GPU's arch at each -m token count up to its max_tokens, in FlyDSL run-only mode. Every case runs checkAllclose on fp32 outputs, and an AOT or gfx1250 case with mismatches must reach logits_diff < 0.01 (test_moe_2stage.py's strict rule). It runs first among the defaults, so kernels other tests JIT-compile in the same process cannot satisfy its check. --deployments and --deployment-id <glob> select what it runs; a selection given that way that runs no case on the GPU fails.
The MXFP4 cases pass quant_dtype_a instead of setting env switches, accept SiTUv2 (with the model's betas), and compare against a stage-wise reference that quantizes a4w4 activations as the kernels do; gfx950 a4w4 logits_diff drops from about 0.022 to 1e-6.
tests/test_aot_drivers.py: bucket boundaries against get_padded_M, the derived rank calls (EP with EPLB, fused shared expert, TP, FHMoE), and the load errors that would otherwise derive a wrong call.
aiter/aot/flydsl/README.md: the deployment file, how rank calls are derived, and how to add a deployment.
Tensor-parallel deployments (third commit):
deployments.py: a fused_moe deployment takes tp (default 1; TP or EP, not both) and shared_expert_mode: "fhmoe". Under TP a rank passes every routed expert, a column per shared expert it runs with no sentinel, no mask, and an intermediate dim of inter_dim / TP; with fhmoe, shared_expert_id is the routed expert count. Ids read {model}/{gfx}/{quant}/tp{N}[-fused|-fhmoe]. No listed deployment pads its weights (ATOM pads DeepSeek-V4 to 128, so 3072 / TP8 = 384 stays), so padding is not a knob: when a deployment that pads is registered, it should be derived from its framework's rule. TP with EP, redundant experts without EP, FHMoE with EP or with other than one shared expert, and an inter_dim TP does not split fail the load.
moe.py: TP deployments replay like EP ones, without num_local_tokens (an EP dispatch's count); an FHMoE deployment carries the caller's gate/up layout in its row. Jobs mark their deployment explicitly instead of inferring it from the mask; without the new entries the job list is unchanged.
The two DeepSeek-V4-Pro TP8 entries: ATOM's recipe runs -tp 8 with ATOM_MOE_GU_ITLV=1, its DeepSeek-V4 model pads the MoE to 128 rather than 512 (inter 384) and keeps the block-FP8 shared expert apart from the MXFP4 routed experts; vLLM's DeepSeek-V4 runs the shared expert through FHMoE on gfx950 TP8 while aiter.fhmoe.supports_dsv4_i384_fhmoe(M) holds, which tuned_fhmoe.csv makes true up to M 2048.
test_moe_ep.py -t aot_ep_profiles keeps to EP deployments.
Review fixes (2b4316884 and fe0162770, after merging main in 9cb6ecdfd):
moe.py: deployment jobs pass the model's swiglu_limit to fused_moe and FHMoE. MXMOE kernels compile the clamp in; the kernels the registered deployments reach take it at launch, so their cache is unchanged.
test_moe_ep.py: real mode always appended shared_E + 1 columns, so a deployment whose shared experts run outside fused_moe routed top-k − 1 experts plus a masked one (DeepSeek-V4-Pro EP 5 of 6, GLM-5.3-Flash EP4 7 of 8). sentinel_col=False now routes the top-k columns alone, and deployments pass it unless they fuse shared experts. The check also forwards swiglu_limit to fused_moe and the reference, and takes the derived per-rank intermediate dim.
deployments.py: a mega_moe deployment off 256 CUs fails the load, as one off gfx950 does: mega_moe.py builds MI355X's bundles.
test_moe_ep.py (fe0162770): -t aot_ep_profiles counts the cases it runs. With none, a --deployments or --deployment-id given on the command line raises instead of exiting 0 (an id that matches nothing, another arch's ids, -m counts outside the deployment); the defaults print a skip, since they do not cover every GPU. CI runs the test with neither, so its run does not change.
Testing
On this head (88ab02584): the last commit leaves the 1195 replay jobs of the 15 default fused_moe deployments unchanged and touches no GPU test; tests/test_aot_drivers.py (9 tests, 122 subtests), ruff 0.16.0, black.
On fe0162770, gfx950 GPU (MI355X), run-only on a cache built from the deployments alone: deepseek-v4-pro/gfx950/a8w4/ep8-r0 (no such id), .../ep8 -m 4, deepseek-v4-pro/gfx1250/* and an explicit --deployments under CU_NUM=80 exit 1; the defaults under CU_NUM=80 skip with exit 0; .../ep8 -m 128 runs 1 case and the defaults run all 9 EP deployments, all passing. The commit changes only the test and the README, so the results below still hold.
Full MoE build: 1827 rows, 15 deployments, 12845 jobs, 0 failures, 2513 artifacts; --check of that cache passes.
The review fixes change no job other than adding swiglu_limit=10 to DeepSeek-V4-Pro's 419 deployment jobs, and the deployment-only build leaves the same 202 launch directories before and after.
gfx950 GPU (MI355X), run-only, on a cache built from the deployments alone: -t aot_ep_profiles passes all 9 gfx950 EP deployments at M 16/128/1024/8192 in real mode (36 cases) and at M 16/256/4096 in fake mode (27 cases), max logits_diff 0.0015, now routing DeepSeek-V4-Pro's 6 experts under its clamp. test_moe_2stage.py on DeepSeek-V4-Pro's 17 TP8 tuned rows and test_fhmoe.py's DeepSeek-V4 serving check at 15 token counts pass, and so does the default test_moe_ep.py.
Full MoE build (tuned FMoE, FHMoE and grouped rows plus deployments; gfx942, gfx950, gfx1250): 1812 rows, 13 EP deployments, 12624 jobs and 1084 bisection replays (gfx1250 722, gfx950 362), 0 failures, 2719 artifacts; --check of that cache passes. The tuned rows alone give 2621 artifacts, all of them in the full build, so the deployments add 98.
gfx950 GPU (MI355X), run-only: on that cache, -t aot_ep_profiles passes all 9 gfx950 deployments at M 16/128/1024/8192 in real mode (36 cases) and at M 16/1024 in fake mode (18 cases), max logits_diff 0.0012. The full default test_moe_ep.py passes.
The schema changes no job: the fused_moe jobs are those of the flat profile table the first commit introduces (1142 jobs, 618 with --row-tokens-only), labels aside; the tuned-row parsing is untouched; the MegaMoE jobs are those of the previous production_jobs() (192), and a subset (Kimi-K3 a4w4, MTPR 8192, 8 ranks) compiled with the previous and the new mega_moe.py leaves the same 281 artifacts.
On the third commit (9e3ced880; the fourth, a658033a1, drops the pad_align knob and leaves the job list identical):
Full MoE build: 1812 rows, 15 deployments, 12677 jobs, 0 failures; --check passes. The cache holds the same kernels as before: every kernel the TP8 calls need was already compiled from tuned rows, so the deployments add none and now state which calls need them.
gfx950 GPU, run-only, on a cache built from the two TP deployments alone: op_tests/test_moe_2stage.py on DeepSeek-V4-Pro's 17 TP8 tuned rows (expert 384, topk 6, inter 384; M 1 to 131072) passes; test_fhmoe.py's DeepSeek-V4 production check (fp8 activations, tuned stage 2) passes at 15 token counts from 1 to 2048, with every fused_moe call run-only except the oracle's single-expert reference.
No gfx1250 GPU was available: the gfx1250 deployments are validated by the CPU replay and --check only, and -t ep_csv (gfx1250 only) was not run.
gfx1250 MegaMoE's stage-2 scatter (ep_fused grouped rows) still has no AOT.
Main's [Config] [Perf] Tune gfx1250 grouped MoE for DeepSeek-V4.1 MegaMoE EP… #7068 added 71 gfx1250 EP-only grouped rows for DeepSeek-V4.1 (model_dim 5120, inter_dim 2304, topk -1). The row replay skips topk -1 rows and no registered deployment matches them, so they have no AOT until DeepSeek-V4.1's gfx1250 EP deployments are registered.
Only DeepSeek-V4-Pro's TP8 calls are registered. Other models' TP deployments are still covered through tuned rows, and so are DeepSeek-V4-Pro's remaining TP rows from [Triton/Gluon] [GFX950] Sparse paged attention and Sparse MLA #4919 (fused shared expert at inter 384/512/768/1536, inter 512, TP4, TP2), for which I found no framework call; they need their owners to confirm.
test_fhmoe.py::test_heterogeneous_moe_matches_precision_oracles fails at M 4 and 8 with AITER_HETERO_MOE_DSV4=1 on this machine, with or without an AOT cache (an accuracy assertion against the separately scheduled kernels); unrelated to this PR.
Kimi-K3's shared_experts is a placeholder (1, unused while no deployment fuses it), and two gfx950 deployments cite their tuned tables as their source.
…hem in test_moe_ep
A caller behind an expert-parallel dispatch passes fused_moe an expert mask,
and the launches it masks compile for the mask's length: the gfx950 masked
stage-2 reduction and route-guarded GEMM2, the gfx1250 g2l route and bf16
gather-reduce. No tuned row records that length, so the AOT build never
compiled them and EP serving JIT-compiled them on the request path.
aiter/aot/flydsl/moe_ep_profiles.csv lists EP callers as what one rank passes
fused_moe: local experts, routing columns, expert-mask length, largest M.
moe.py replays each profile with the mask at both ends of every get_padded_M
bucket up to that M (bisecting between them as for rows), with and without
num_local_tokens, against the tuned tables serving reads, so a profile needs
no tuned row. --ep-profile selects profile tables; without --csv and
--ep-profile both defaults are replayed, and run_aot replays both.
The default table covers gfx950 Kimi-K3 EP8, DeepSeek-V4-Pro EP8 (EPLB r0,
r32, r64) and EP16, GLM-5.3-Flash EP8/EP4/EP2, and gfx1250 Kimi-K3 EP8/EP16
and DeepSeek-V4-Pro EP4, from the EP rows tuned in #5804, #6928, #6435,
op_tests/test_moe_ep.py -t aot_ep_profiles, first in the defaults, runs every
profile of the GPU's arch through the EP fused_moe call in FlyDSL run-only
mode and requires logits_diff < 0.01. The MXFP4 cases now pin quant_dtype_a
instead of setting env switches, accept SiTUv2, and compare against a
stage-wise reference that quantizes a4w4 activations as the kernels do.
Co-authored-by: Cursor <cursoragent@cursor.com>
…all from them
The EP profile table recorded what one rank passes fused_moe (local experts,
routing columns, mask length) as hand-computed numbers, and mega_moe.py kept
DeepSeek-V4-Pro and Kimi-K3 as Python constants. aiter/aot/flydsl/
moe_deployments.json now records each MoE model's architecture once, as its
config.json gives it, and each deployment by its knobs: GPU, quant, EP size,
EPLB redundant experts, shared-expert fusion and token range. deployments.py
validates the files (unknown keys, experts that do not split over EP, fusing
a model without shared experts, an id listed twice) and derives each rank's
call with the convention ATOM and vLLM share. python -m
aiter.aot.flydsl.deployments lists what the files expand to.
moe.py replays the fused_moe deployments; --deployments replaces --ep-profile.
mega_moe.py builds the mega_moe deployments; --mtpr, --quant and
--experts-per-rank narrow them down, --deployments takes other files, and
the shape flags are gone. Both job lists are those of before, labels aside:
1142 EP jobs, 11394 tuned-row jobs, 192 MegaMoE jobs. --check passes on the
cache the previous build filled.
op_tests/test_moe_ep.py -t aot_ep_profiles takes the global experts, fused
shared experts and EP size from the deployment instead of searching for a
layout that fits the table's numbers, uses the model's SiTUv2 betas, and
picks deployments by id with --deployment-id.
Co-authored-by: Cursor <cursoragent@cursor.com>
ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm
vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
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 7063 --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.
github-actionsBot
changed the title
[FlyDSL] [JIT] [Feature] Precompile expert-parallel MoE deployments, described by model and deployment knobs
[FlyDSL] [Feature] Precompile expert-parallel MoE deployments, described by model and deployment knobs
Oct 28, 2026
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Unfused routing drops a valid expert route, DeepSeek’s activation clamp is left untested, deployment CU counts are ignored, and new AOT cases omit required elementwise correctness checks.
Extends FlyDSL’s CPU-based ahead-of-time compilation to expert-parallel MoE deployments, covering masked runtime calls that tuned-table replay alone does not describe.
Changes:
Adds a shared model/deployment schema with validation and derived per-rank shapes.
Integrates deployments into MoE and MegaMoE replay and the default AOT build.
Adds run-only deployment checks, quantization-aware references, unit tests, and documentation.
File
Description
tests/test_aot_drivers.py
Tests deployment validation, derived shapes, token buckets, and replay jobs.
op_tests/test_moe_ep.py
Adds deployment-driven cache checks and updates quantized references.
aiter/aot/flydsl/README.md
Documents deployment configuration, replay, and cache validation.
aiter/aot/flydsl/moe.py
Generates masked deployment replays across token buckets.
aiter/aot/flydsl/moe_deployments.json
Defines shared model facts and default deployment profiles.
aiter/aot/flydsl/mega_moe.py
Replaces hardcoded profiles with deployment-derived jobs.
aiter/aot/flydsl/deployments.py
Validates deployment files and derives per-rank runtime shapes.
aiter/aot/flydsl/common.py
Includes default MoE deployments in the AOT build.
… TP8
moe_deployments.json could only describe callers behind an EP dispatch. A
tensor-parallel caller passes no expert mask, but its call still depends on
knobs no tuned row records: how its framework pads the split intermediate dim
(intermediate_pad, which the FlyDSL GEMMs compile in) and how it runs the
shared expert (fused into the routed weights, apart, or through FHMoE).
A fused_moe deployment now takes tp (default 1; TP or EP, not both),
pad_align (ATOM's FusedMoE pad_align, default 1) and shared_expert_mode
"fhmoe" (one shared expert in its own weights, TP only, as vLLM's DeepSeek-V4
runs it). Under TP a rank passes every routed expert, the shared experts it
runs without a sentinel column, no mask, and inter_dim / TP padded to
pad_align; with fhmoe, shared_expert_id is the routed expert count. ids read
{model}/{gfx}/{quant}/tp{N}[-pad{align}][-fused|-fhmoe]. The loader refuses TP
and EP together, redundant experts without EP, FHMoE with EP or with other
than one shared expert, and an inter_dim TP does not split.
moe.py replays TP deployments like EP ones, without num_local_tokens (an EP
dispatch's count), against the tables serving reads; an FHMoE deployment takes
the caller's gate/up layout into its row. Jobs mark their deployment
explicitly instead of inferring it from the mask, and with the new entries
left out the job list is the previous one.
The default file gains DeepSeek-V4-Pro on gfx950 a8w4 at TP8: the ATOM recipe
call (expert 384, topk 6, inter 384: ATOM pads DeepSeek-V4 to 128, not 512,
and keeps the block-FP8 shared expert apart) and vLLM's heterogeneous call
(FHMoE, expert 385, topk 7, shared slot 384, up to the M tuned_fhmoe.csv
covers). Every kernel they need was already compiled from tuned rows, so the
default build is unchanged; the deployments now say which calls need them.
op_tests/test_moe_ep.py -t aot_ep_profiles keeps to EP deployments.
Co-authored-by: Cursor <cursoragent@cursor.com>
zhiding512
changed the title
[FlyDSL] [Feature] Precompile expert-parallel MoE deployments, described by model and deployment knobs
[FlyDSL] [JIT] [Feature] Precompile MoE deployments (EP and TP), described by model and deployment knobs
Oct 28, 2026
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Deployment replay and tests fail to preserve declared activation parameters, routing, padding, and CU counts, while the new checks bypass elementwise validation.
Deployment jobs drop activation parameters during replay
aiter/aot/flydsl/moe.py:355
Deployment jobs discard model.swiglu_limit and replace the model's SiTUv2 betas with generic _variants() defaults. A model declaring betas (2, 17) therefore never replays those values, and DeepSeek's declared clamp of 10 is not forwarded, so --check can validate a different caller from the registered deployment. Carry the model's activation parameters through job construction and _replay(), with regression coverage for non-default betas and a nonzero clamp.
No listed deployment pads its expert weights: every rank's intermediate dim
is already a multiple of what its framework pads MXFP4 weights to (ATOM pads
DeepSeek-V4 to 128, and 3072 / TP8 = 384). pad_align only restated a framework
constant per deployment, so a rank's dims now follow from the model and TP
alone. The job list is unchanged.
Padding is the framework's rule (ATOM's FusedMoE pad_align with per-model
overrides, vLLM's per-backend round-up), so when a deployment that pads is
registered, its padding should be derived from that rule rather than written
into each deployment.
Co-authored-by: Cursor <cursoragent@cursor.com>
github-actionsBot
changed the title
[FlyDSL] [JIT] [Feature] Precompile MoE deployments (EP and TP), described by model and deployment knobs
[FlyDSL] [Feature] Precompile MoE deployments (EP and TP), described by model and deployment knobs
Oct 28, 2026
Preserve custom activation parameters when replaying deployments
aiter/aot/flydsl/moe.py:353
Deployment jobs drop model.swiglu_limit and replace the declared SiTUv2 betas with _variants()'s fixed choices, so a custom deployment with betas such as (2, 3) is never replayed with its requested activation parameters. Paths that compile these constants can therefore still require JIT compilation when serving that deployment. Preserve the model's activation parameters in deployment jobs and forward them through _replay (including the clamp on the FHMoE path), leaving generic activation variants for CSV rows.
No conflicts, but main changed code the deployment replay runs through:
- #7020: the MXFP4 inline sort takes num_local_tokens and skips expert
ids outside [0, E), as the sentinel ids of EP callers are.
- #7067: the a16w4 stage-2 reduce mode in fused_moe.py and Kimi-K3's
a16w4 table.
- #7068, #6668: gfx1250 grouped tables for DeepSeek-V4.1 MegaMoE EP and
MiniMax-M3.
- #7055: the OPUS CO tests read local kernel IDs; this branch's 1-GPU CI
shards failed on the global ones.
Co-authored-by: Cursor <cursoragent@cursor.com>
Review of #7063 found places where a deployment's replay or its EP check
did not make the call serving makes.
- moe.py: deployment jobs pass the model's swiglu_limit to fused_moe and
FHMoE, and a SiTUv2 model's own betas join the generic variants when
they are not among them. MXMOE kernels compile the clamp in; the FlyDSL
v1, FHMoE and gfx1250 grouped kernels the registered deployments reach
take it at launch, so their cache is unchanged (the deployment-only
build leaves the same 202 launch directories before and after).
- test_moe_ep.py's real mode always appended shared_E + 1 columns, so a
deployment whose shared experts run outside fused_moe routed top-k - 1
experts plus a masked one (DeepSeek-V4-Pro EP 5 of 6, GLM-5.3-Flash EP4
7 of 8). sentinel_col=False now routes the top-k columns alone;
deployments pass it unless they fuse shared experts, and other callers
keep their layout.
- The EP check forwards swiglu_limit to fused_moe and the reference,
takes the derived per-rank intermediate dim, and runs checkAllclose on
fp32 outputs for every case. AOT and gfx1250 cases follow
test_moe_2stage.py's strict rule, under which a case with mismatches
must reach logits_diff < 0.01: against these references 30-50% of the
outputs miss atol/rtol 5e-2 at logits_diff <= 0.0015, so a bound on the
mismatch ratio would be loose or flaky.
- deployments.py refuses a mega_moe deployment off 256 CUs, as it refuses
one off gfx950: mega_moe.py builds MI355X's bundles.
Co-authored-by: Cursor <cursoragent@cursor.com>
Review of #7063 found places where a deployment's replay or its EP check
did not make the call serving makes.
- moe.py: deployment jobs pass the model's swiglu_limit to fused_moe and
FHMoE, and a SiTUv2 model's own betas join the generic variants when
they are not among them. MXMOE kernels compile the clamp in; the FlyDSL
v1, FHMoE and gfx1250 grouped kernels the registered deployments reach
take it at launch, so their cache is unchanged (the deployment-only
build leaves the same 202 launch directories before and after).
- test_moe_ep.py's real mode always appended shared_E + 1 columns, so a
deployment whose shared experts run outside fused_moe routed top-k - 1
experts plus a masked one (DeepSeek-V4-Pro EP 5 of 6, GLM-5.3-Flash EP4
7 of 8). sentinel_col=False now routes the top-k columns alone;
deployments pass it unless they fuse shared experts, and other callers
keep their layout.
- The EP check forwards swiglu_limit to fused_moe and the reference,
takes the derived per-rank intermediate dim, and runs checkAllclose on
fp32 outputs for every case. AOT and gfx1250 cases follow
test_moe_2stage.py's strict rule, under which a case with mismatches
must reach logits_diff < 0.01: against these references 30-50% of the
outputs miss atol/rtol 5e-2 at logits_diff <= 0.0015, so a bound on the
mismatch ratio would be loose or flaky.
- deployments.py refuses a mega_moe deployment off 256 CUs, as it refuses
one off gfx950: mega_moe.py builds MI355X's bundles.
Co-authored-by: Cursor <cursoragent@cursor.com>
Re the activation-parameter finding in the review summary (moe.py deployment jobs), fixed in 2b43168:
Deployment jobs now pass the model's swiglu_limit to fused_moe and FHMoE in the replay.
If a SiTUv2 model's declared betas aren't already among the generic variants, they're added. The generic variants stay, since tuned rows enumerate them too.
Tuned-row jobs are unchanged.
test_deployment_jobs_pass_model_activation covers non-default betas (2/17), DeepSeek-V4-Pro's clamp, and the forwarding through _replay.
The cache for the registered deployments is unchanged. The deployment-only build leaves the same 202 launch directories before and after, because the FlyDSL v1, FHMoE and gfx1250 grouped kernels those deployments reach take the clamp and betas as launch arguments. MXMOE kernels compile the clamp in, so a deployment that reaches them now gets the clamped kernel serving asks for.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
An explicitly supplied deployment filter can match no cases and still exit successfully, producing a false-green AOT-cache check without performing any cache lookup.
With --deployment-id or --deployments given, aot_ep_profiles could run
no case and exit 0: an id that matches nothing (the zero-redundancy
DeepSeek-V4-Pro EP8 id has no -r0 suffix), ids of another arch or CU
count, or -m counts all above max_tokens or below ep in real mode. That
reported an AOT cache check that looked nothing up as passing.
It now counts the cases it runs. With none, a selection given on the
command line raises; the defaults print a skip, since the default
deployments do not cover every GPU. Other archs than gfx950 and gfx1250
still skip up front. CI runs the test without either option, so its run
does not change.
Co-authored-by: Cursor <cursoragent@cursor.com>
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Explicitly supplying the default deployments path can bypass the empty-selection error, allowing real-mode runs to exit successfully without checking the cache.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Deployment replay omits quantized EP inputs, risking AOT cache misses, while explicitly selecting the default deployment file can incorrectly report success without checking any cases.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
An explicitly supplied default deployment-file path can skip every real-mode EP case yet exit successfully because the code compares list contents rather than distinguishing explicit input from the parser default.
tests/test_aot_drivers.py keeps three of its eight deployment tests:
- the replay's token buckets against get_padded_M, since moe.py keeps its
own copy of the padded-M tiers;
- the rank call of one deployment per branch of the derivation (EP with
EPLB replicas, EP with a fused shared expert, TP, TP through FHMoE)
against numbers worked out by hand. aot_ep_profiles checks EP calls
only, on a GPU with a built cache, and takes the expert counts from
deployments.py;
- the load errors without which a deployment would derive a wrong call:
experts or inter_dim that do not split, TP with EP, a misspelled knob.
The others restated what deployment_row copies, derived again what the
kept rows derive, counted the default MegaMoE jobs, or built a SiTUv2
model no deployment registers. Each validation case now loads a file of
its own that names the default models, which a separate test used to
check.
That SiTUv2 model was all that reached the replay of a model's own betas,
since Kimi-K3's are among the betas every SiTUv2 row replays, so the
branch goes too. The default deployments expand to the same 1195 jobs.
Co-authored-by: Cursor <cursoragent@cursor.com>
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Job generation omits custom SiTUv2 beta pairs, leaving valid model-specific MXMOE specializations out of the deployment cache; the declared pair must be included and covered by a non-default-pair test.
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
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.
github-actions
Bot
changed the title
Follows #6362, which builds the MoE and MegaMoE AOT caches by replaying the runtime on CPU.
Summary
A caller behind an expert-parallel dispatch passes
fused_moeanexpert_mask, which brings in FlyDSL launches no tuned row reaches: the gfx950 stage-2 reduction, compiled for the mask's length (moe_reduction_..._m1e<len>), and the gfx1250 route (moe_g2l_lut,moe_route_g2l_lds) and bf16 gather-reduce. The mask's length follows from the deployment (global experts, EPLB replicas, fused shared experts), which no tuned row records, so the AOT build never compiled these kernels and EP serving JIT-compiled them on the request path.This PR describes those deployments in
aiter/aot/flydsl/moe_deployments.json: each MoE model's architecture once, as itsconfig.jsongives it, and each deployment by its knobs (GPU, quant, EP size, EPLB redundant experts, shared-expert fusion, token range).aiter/aot/flydsl/deployments.pyvalidates the file and derives what each rank hands the runtime, with the convention ATOM and vLLM share. From it:moe.py) precompiles everyfused_moedeployment's masked calls;mega_moe.py) takes its DeepSeek-V4-Pro and Kimi-K3 deployments from the same file instead of Python constants;op_tests/test_moe_ep.py -t aot_ep_profilesruns each deployment through the EPfused_moecall in FlyDSL run-only mode, so a kernel missing from the cache fails the test.The third commit extends the file to tensor-parallel callers. They pass no mask, but their call still depends on what no tuned row records: how the framework runs the shared expert (fused, apart, or through FHMoE). It registers DeepSeek-V4-Pro's two TP8 calls on gfx950 as the first TP deployments.
After a merge of main, two commits address review: a deployment's replay passes its model's activation clamp, the EP check routes and checks the call as serving makes it, a MegaMoE deployment must be on MI355X (gfx950, 256 CUs), and an EP check whose explicit deployment selection runs no case fails instead of passing.
The last commit trims
tests/test_aot_drivers.pyto the checks no other test makes, and drops the replay of a SiTUv2 model's own betas, which no registered model reaches: Kimi-K3's (4, 25) are among the betas every SiTUv2 row replays. The default deployments expand to the same replay jobs.Motivation
expert_maskon gfx950 and gfx1250 whenever expert parallelism is on. On this head, the 13 default deployments need 98 FlyDSL artifacts the tuned-row replay does not build.swiglu_limitand SiTUv2 betas were Python constants inmega_moe.py; both replays and the test now read the same entries.Tuned tables keep deciding which kernels the runtime picks. A deployment is replayed against the tables serving reads: it uses a tuned row where one matches and the runtime's default heuristic where none does, so it needs no tuned row.
Changes
aiter/aot/flydsl/moe_deployments.json(new): 3 models and 11 deployment entries, which expand to 15fused_moerank calls and 8 MegaMoE profiles. The EP entries come from the rows tuned for EP in [FlyDSL] fix(flydsl): support kmk3 ep moe & opt for ep moe gemm2 #5804, [Config] [Perf] Retune DSV4 EP8 a8w4 fused-MoE rows for M 2-8 and 1024-8192 on gfx950 #6928, [HIP] [GLM-5.3-Flash] MXFP4 fused-MoE tuned configs for MI355X #6435, [Config] [Tuning] Add gfx1250 A4W4 grouped-MoE tuning for the Kimi-K3 WideEP #6438 and [CK] [FlyDSL] [FEAT] support fp8/fp4 quant in mega stage2 combine #5176; the TP entries from the ATOM and vLLM code noted below.aiter/aot/flydsl/deployments.py(new): loads and validates deployment files, and derives each rank's call. With P the routed plus redundant experts and S the fused shared experts, a rank passesw1.shape[0]= P / EP + S,topk_ids.shape[1]= top-k (top-k + S + 1 with fused shared experts, the last column a sentinel that masks them on other ranks' tokens), and a mask of P + S + 1 entries. Unknown or missing fields, experts that do not split over EP, fusing a model without shared experts, or an id listed twice fail the load. Files given together share their models, so a file of your own can name the default file's models.python -m aiter.aot.flydsl.deploymentslists what the files expand to.aiter/aot/flydsl/moe.py: replays eachfused_moedeployment with its mask at both ends of everyget_padded_Mbucket up tomax_tokens(bisecting between them, as for rows), with and withoutnum_local_tokens(MORI dispatch passes it, the all-gather path does not), with the activation dtype pinned to the deployment's quant as serving pins it.--deploymentsselects deployment files; with neither--csvnor--deployments, both defaults are replayed, asrun_aot(the wheel build) does.aiter/aot/flydsl/mega_moe.py: jobs come from themega_moedeployments, and the job list is unchanged (192 jobs).--mtpr,--quantand--experts-per-ranknow narrow the deployments down (--experts-per-rank 48keeps working),--deploymentstakes other files, and the shape flags (--topk,--model-dim,--inter-dim,--world-size,--act,--situ-beta,--situ-linear-beta,--swiglu-limit) are removed in favor of a deployment file.op_tests/test_moe_ep.py:-t aot_ep_profilesruns everyfused_moedeployment of the GPU's arch at each-mtoken count up to itsmax_tokens, in FlyDSL run-only mode. Every case runscheckAllcloseon fp32 outputs, and an AOT or gfx1250 case with mismatches must reachlogits_diff < 0.01(test_moe_2stage.py's strict rule). It runs first among the defaults, so kernels other tests JIT-compile in the same process cannot satisfy its check.--deploymentsand--deployment-id <glob>select what it runs; a selection given that way that runs no case on the GPU fails.quant_dtype_ainstead of setting env switches, accept SiTUv2 (with the model's betas), and compare against a stage-wise reference that quantizes a4w4 activations as the kernels do; gfx950 a4w4logits_diffdrops from about 0.022 to 1e-6.-t ep_csvis kept and runs through the same test function.tests/test_aot_drivers.py: bucket boundaries againstget_padded_M, the derived rank calls (EP with EPLB, fused shared expert, TP, FHMoE), and the load errors that would otherwise derive a wrong call.aiter/aot/flydsl/README.md: the deployment file, how rank calls are derived, and how to add a deployment.Tensor-parallel deployments (third commit):
deployments.py: afused_moedeployment takestp(default 1; TP or EP, not both) andshared_expert_mode: "fhmoe". Under TP a rank passes every routed expert, a column per shared expert it runs with no sentinel, no mask, and an intermediate dim ofinter_dim/ TP; withfhmoe,shared_expert_idis the routed expert count. Ids read{model}/{gfx}/{quant}/tp{N}[-fused|-fhmoe]. No listed deployment pads its weights (ATOM pads DeepSeek-V4 to 128, so 3072 / TP8 = 384 stays), so padding is not a knob: when a deployment that pads is registered, it should be derived from its framework's rule. TP with EP, redundant experts without EP, FHMoE with EP or with other than one shared expert, and aninter_dimTP does not split fail the load.moe.py: TP deployments replay like EP ones, withoutnum_local_tokens(an EP dispatch's count); an FHMoE deployment carries the caller's gate/up layout in its row. Jobs mark their deployment explicitly instead of inferring it from the mask; without the new entries the job list is unchanged.-tp 8withATOM_MOE_GU_ITLV=1, its DeepSeek-V4 model pads the MoE to 128 rather than 512 (inter 384) and keeps the block-FP8 shared expert apart from the MXFP4 routed experts; vLLM's DeepSeek-V4 runs the shared expert through FHMoE on gfx950 TP8 whileaiter.fhmoe.supports_dsv4_i384_fhmoe(M)holds, whichtuned_fhmoe.csvmakes true up to M 2048.test_moe_ep.py -t aot_ep_profileskeeps to EP deployments.Review fixes (
2b4316884andfe0162770, after merging main in9cb6ecdfd):moe.py: deployment jobs pass the model'sswiglu_limittofused_moeand FHMoE. MXMOE kernels compile the clamp in; the kernels the registered deployments reach take it at launch, so their cache is unchanged.test_moe_ep.py: real mode always appendedshared_E + 1columns, so a deployment whose shared experts run outsidefused_moerouted top-k − 1 experts plus a masked one (DeepSeek-V4-Pro EP 5 of 6, GLM-5.3-Flash EP4 7 of 8).sentinel_col=Falsenow routes the top-k columns alone, and deployments pass it unless they fuse shared experts. The check also forwardsswiglu_limittofused_moeand the reference, and takes the derived per-rank intermediate dim.deployments.py: amega_moedeployment off 256 CUs fails the load, as one off gfx950 does:mega_moe.pybuilds MI355X's bundles.test_moe_ep.py(fe0162770):-t aot_ep_profilescounts the cases it runs. With none, a--deploymentsor--deployment-idgiven on the command line raises instead of exiting 0 (an id that matches nothing, another arch's ids,-mcounts outside the deployment); the defaults print a skip, since they do not cover every GPU. CI runs the test with neither, so its run does not change.Testing
On this head (
88ab02584): the last commit leaves the 1195 replay jobs of the 15 defaultfused_moedeployments unchanged and touches no GPU test;tests/test_aot_drivers.py(9 tests, 122 subtests), ruff 0.16.0, black.On
fe0162770, gfx950 GPU (MI355X), run-only on a cache built from the deployments alone:deepseek-v4-pro/gfx950/a8w4/ep8-r0(no such id),.../ep8 -m 4,deepseek-v4-pro/gfx1250/*and an explicit--deploymentsunderCU_NUM=80exit 1; the defaults underCU_NUM=80skip with exit 0;.../ep8 -m 128runs 1 case and the defaults run all 9 EP deployments, all passing. The commit changes only the test and the README, so the results below still hold.On
2b4316884(main through #6519), FlyDSL 0.4.0:--checkof that cache passes.swiglu_limit=10to DeepSeek-V4-Pro's 419 deployment jobs, and the deployment-only build leaves the same 202 launch directories before and after.-t aot_ep_profilespasses all 9 gfx950 EP deployments at M 16/128/1024/8192 in real mode (36 cases) and at M 16/256/4096 in fake mode (27 cases), maxlogits_diff0.0015, now routing DeepSeek-V4-Pro's 6 experts under its clamp.test_moe_2stage.pyon DeepSeek-V4-Pro's 17 TP8 tuned rows andtest_fhmoe.py's DeepSeek-V4 serving check at 15 token counts pass, and so does the defaulttest_moe_ep.py.tests/test_aot_drivers.py(14 tests, 152 subtests), ruff 0.16.0, black.Earlier, on main through #6704:
--checkof that cache passes. The tuned rows alone give 2621 artifacts, all of them in the full build, so the deployments add 98.-t aot_ep_profilespasses all 9 gfx950 deployments at M 16/128/1024/8192 in real mode (36 cases) and at M 16/1024 in fake mode (18 cases), maxlogits_diff0.0012. The full defaulttest_moe_ep.pypasses.fused_moejobs are those of the flat profile table the first commit introduces (1142 jobs, 618 with--row-tokens-only), labels aside; the tuned-row parsing is untouched; the MegaMoE jobs are those of the previousproduction_jobs()(192), and a subset (Kimi-K3 a4w4, MTPR 8192, 8 ranks) compiled with the previous and the newmega_moe.pyleaves the same 281 artifacts.tests/test_aot_drivers.py(12 tests), ruff, black.On the third commit (
9e3ced880; the fourth,a658033a1, drops thepad_alignknob and leaves the job list identical):--checkpasses. The cache holds the same kernels as before: every kernel the TP8 calls need was already compiled from tuned rows, so the deployments add none and now state which calls need them.op_tests/test_moe_2stage.pyon DeepSeek-V4-Pro's 17 TP8 tuned rows (expert 384, topk 6, inter 384; M 1 to 131072) passes;test_fhmoe.py's DeepSeek-V4 production check (fp8 activations, tuned stage 2) passes at 15 token counts from 1 to 2048, with everyfused_moecall run-only except the oracle's single-expert reference.tests/test_aot_drivers.py(13 tests, 152 subtests), ruff, black.Not covered
--checkonly, and-t ep_csv(gfx1250 only) was not run.ep_fusedgrouped rows) still has no AOT.test_fhmoe.py::test_heterogeneous_moe_matches_precision_oraclesfails at M 4 and 8 withAITER_HETERO_MOE_DSV4=1on this machine, with or without an AOT cache (an accuracy assertion against the separately scheduled kernels); unrelated to this PR.shared_expertsis a placeholder (1, unused while no deployment fuses it), and two gfx950 deployments cite their tuned tables as their source.