[FlyDSL] Cache compiled FlyDSL launchers on host paths by coderfeli · Pull Request #6946 · ROCm/aiter

coderfeli · 2026-10-26T12:10:23Z

Summary

Cache compiled FlyDSL launchers on live host paths using _run_compiled, including batched GEMM v2, grouped MoE routing/quantization, MQA, QSA, HC mix, MLA, gfx1250 ptpc GEMM/projection and MegaMoE compact-plan/dispatch/combine stages.
Build grouped GEMM launch arguments and specialization keys from the same configuration tuples. Include varying Constexpr configuration in other shared launchers' keys.
Normalize tensor arguments used as raw buffers that accept varying ranks into flat aliases. Keep page/block tables in their existing 2D layout. Their compiled ABI stays valid across changes in buffer rank, while runtime row counts continue to reuse the same compiled entry. Preserve native tensor layouts where the kernel indexes them directly.
Honor COMPILE_ONLY after a callable has been warmed and retain the execution entry for subsequent runtime use.
Pass typed pointers to grouped MoE route-map and scatter launchers. GPU kernel implementations are unchanged.
Exclude already obsolete legacy batched paths; the supported v2 path remains covered.

Host measurements

Measured on gfx950 with FlyDSL 0.4.0 and PyTorch 2.10.0 / ROCm 7.2.4. The comparison loads the exact parent source (41d4d8d7e) and the PR source in one process with the same input tensors. torch.set_default_device(None) disables PyTorch's default-device interception during timing; all tensor allocations specify their device explicitly.

Public call	Parent	PR	Host-time reduction
MXFP4 quantization, M=2, N=1024	55.7 µs	12.5 µs	77.6%
QSA, 300 rows, 80 columns	89.6 µs	18.4 µs	79.4%
Contiguous prefix sum, 8 experts	50.1 µs	13.4 µs	73.2%
FP4 MQA decode, batch=2, context=512	67.8 µs	9.8 µs	85.5%
HC mix, M=1, K=10240, R=320	78.5 µs	21.7 µs	72.4%

Method: warmed real GPU launches; medians of 15 alternating-order batches of 200 calls. Synchronization runs before/after each batch outside the timed host interval. Public timings include allocations, views, conversions and factory/configuration work. MQA uses precomputed scheduling and preallocated output. Cold compilation and model throughput are not measured.

The earlier 24–26 µs figures included torch.set_default_device("cuda"), installed by imported test modules. That PyTorch mode intercepts even tensor metadata calls; setting the default device to "cpu" also installs it. Those figures measured the full Python API, not the native launch.

Quantization breakdown with the final compile-only guard (500 calls × 15 batches, host time):

Layer	Default-device mode disabled	Default-device mode set to CUDA
Prepacked native launch	3.3 µs	3.4 µs
Tensor ABI filling, without a launch	1.4 µs	7.0 µs
Compiled callable, including ABI filling and launch	5.1 µs	10.4 µs
`_run_compiled`	5.9 µs	11.1 µs
Full quantization API	12.3 µs	23.4 µs

The old JIT already cached machine code. The saving comes from avoiding repeated signature binding, global/compile-hint checks and full argument cache-key construction. The full Python API still allocates tensors, constructs views/stream arguments and packs runtime metadata.

Validation

Local review probes passed: scale-sort raw buffer ranks; QSA metadata ranks and two divisors; compile-only after warmup with both single-entry and specialization caches, for both 1 and true. Tests check correctness and compiled-entry reuse.
All 52 HC mix tests passed, including graph/stream isolation, fused reduction and specialization variants.
FP4 MQA decode and prefill/variable-query-length correctness passed after raw-buffer normalization.
Exact parent/PR outputs matched for quantization, QSA, contiguous prefix sum, HC mix and FP4 MQA decode.
black --check, ruff check, python -m compileall and git diff --check passed.

The gfx1250-only GEMM and MLA paths remain untested on hardware because only gfx950 is available.

Global dispatch audit

A scan of all 887 Python files under aiter/, with manual tracing of launcher factories and custom caches, identified six further direct JIT dispatch paths. All six are now wrapped: gfx1250 a8w8 bpreshuffle GEMM, MLA gather/projection, and MegaMoE compact plan, TDM dispatch, combine barrier and combine reduce.

The shared GEMM key includes all 18 Constexpr parameters; projection specializes by head count, scale source and cluster size. MegaMoE caches per existing factory/configuration. Runtime sizes, pointers and streams remain runtime arguments.

All 13 gfx1250 compile-only code-generation configurations passed: GEMM output dtype/split-K, projection head count/scale source/cluster, compact-plan slots, TDM plain/compact dispatch, and bf16/FP8/FP4 combine.
Standalone host probes compared 23 pre/post wrapper argument and stream-order cases against native launcher signatures. With simulated compilation/dispatch, they exercised the real cache helper and passed six warmed compile-only/resume checks.
All 14 FlyDSL GPU/JIT definitions in the changed PR files match the exact parent AST; no extra unit-test file is added.

For cross-compilation from a gfx950 host, use FLYDSL_GPU_ARCH=gfx1250 ARCH=gfx1250 MORI_GPU_ARCHS=gfx1250 COMPILE_ONLY=1. The Mori setting prevents linking CCO bitcode built for the host architecture.

Available hardware is gfx950. gfx1250 output correctness, graph replay and multi-GPU communication execution remain unverified; compile-only and host probes do not establish those results.

github-actions · 2026-10-26T12:10:51Z

🏷️ 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
`ci:windows`	Windows JIT op smoke test on the gfx1201 ROCm runner
`ci:multi-gpu`	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: 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 6946 --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.

…mode

yadaish · 2026-10-27T02:58:20Z

LGTM

… wrappers Builds on ROCm#6946 (compiled FlyDSL launcher cache) and keeps only the host-side savings it does not cover: - batched mxfp8_32 GEMM: memoise the untuned-shape heuristic per M class (ceil(M / tile_m) per candidate tile, saturated past the M-block limit, plus M parity when an A-preshuffle candidate exists) instead of per exact M, and cache kernel-name parsing. Prefill hands every linear a new M each step; the pick is identical to the per-M heuristic. - get_mxscale_bpreshuffle_config: answer an N/K with no tuned row at any M without walking the padded-M levels, and log an untuned weight shape once per (N, K, w_scale_block) instead of once per M. - gemm_a8w8_blockscale -> gemm_a8w8_blockscale_bpreshuffle and tuned_gemm.gemm_a16w16 / gemm_a16w16_out -> triton gemm_a16w16_: the outer custom op's body calls the nested op's body directly instead of a second torch.ops dispatch. torch_compile_guard exposes the body it registered as ``_eager_body``; a rebound module attribute is still honoured. - triton gemm_a16w16 gluon path: memoise the WMMA/shared layouts per tile config. - grouped_gemm_gfx1250_a8w4: build debug strings only when AITER_GROUPED_DEBUG is on and skip the arch-override env reads on gfx1250. op_tests/test_gemm_host_dispatch_cache.py checks each cache against the uncached answer (CPU contracts plus bitwise GPU comparisons). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… wrappers Complements ROCm#6946 (compiled FlyDSL launcher cache, merged) with the host-side savings it does not cover: - batched mxfp8_32 GEMM: memoise the untuned-shape heuristic per M class (ceil(M / tile_m) per candidate tile, saturated past the M-block limit, plus M parity when an A-preshuffle candidate exists) instead of per exact M, and cache kernel-name parsing. Prefill hands every linear a new M each step; the pick is identical to the per-M heuristic. - get_mxscale_bpreshuffle_config: answer an N/K with no tuned row at any M without walking the padded-M levels, and log an untuned weight shape once per (N, K, w_scale_block) instead of once per M. - gemm_a8w8_blockscale -> gemm_a8w8_blockscale_bpreshuffle and tuned_gemm.gemm_a16w16 / gemm_a16w16_out -> triton gemm_a16w16_: the outer custom op's body calls the nested op's body directly instead of a second torch.ops dispatch. torch_compile_guard exposes the body it registered as ``_eager_body``; a rebound module attribute is still honoured. - triton gemm_a16w16 gluon path: memoise the WMMA/shared layouts per tile config. - grouped_gemm_gfx1250_a8w4: build debug strings only when AITER_GROUPED_DEBUG is on and skip the arch-override env reads on gfx1250. op_tests/test_gemm_host_dispatch_cache.py checks each cache against the uncached answer (CPU contracts plus bitwise GPU comparisons). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

) ROCm#6946 flattens the grouped GEMM's scale arguments with .view(-1), which a non-contiguous tensor cannot do. MegaMoE's compact rows hand over exactly that -- each row's e8m0 scales sit inside its wire row -- so every stage1_fused MegaMoE step raises. The kernel only takes the base address and walks the rows by a_scale_row_stride_bytes, so a strided scale view becomes a one-dimensional view over the same base, as it was before ROCm#6946. Temporary: drop this commit once main fixes it. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…ppers (#7006) FlyDSL-only split of #6948 (the Triton gemm_a16w16 changes are dropped). Complements #6946 (compiled FlyDSL launcher cache). - batched mxfp8_32 GEMM: memoise the untuned-shape heuristic per M class (ceil(M / tile_m) per candidate tile, saturated past the M-block limit, plus M parity when an A-preshuffle candidate exists) instead of per exact M, and cache kernel-name parsing. Prefill hands every linear a new M each step; the pick is identical to the per-M heuristic. - get_mxscale_bpreshuffle_config: answer an N/K with no tuned row at any M without walking the padded-M levels, and log an untuned weight shape once per (N, K, w_scale_block) instead of once per M. - gemm_a8w8_blockscale -> gemm_a8w8_blockscale_bpreshuffle: the outer custom op's body calls the nested op's body directly instead of a second torch.ops dispatch. torch_compile_guard exposes the body it registered as ``_eager_body``; a rebound module attribute is still honoured. - grouped_gemm_gfx1250_a8w4: build debug strings only when AITER_GROUPED_DEBUG is on and skip the arch-override env reads on gfx1250. op_tests/test_gemm_host_dispatch_cache.py checks each cache against the uncached answer (CPU contracts plus bitwise GPU comparisons). Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

The runners rely on three properties of AITER's _preload_compiled and _run_compiled: a preload never runs the kernel, the first real call runs it once, and later calls skip FlyDSL's JIT dispatch. ROCm/aiter#6946 already changed _run_compiled, so a later AITER pin bump can break them without any other test failing. The preload runs on a real buffer first, so a dispatching preload fails an assertion instead of faulting the GPU at compile_only's zero addresses.

## Motivation Without CUDA graphs, every launch of the DeepSeek-V4.1, GLM-5.3, GLM-5.3-Flash and MiniMax-M3 runners goes through FlyDSL's JIT dispatch (`JitFunction.__call__`): it binds the signature, checks the globals and builds the full cache key at each launch. For GLM-5.3-Flash's launchers (50 integer arguments), that costs 182 µs of host time a launch, so an eager decode step is host-bound. This PR reuses AITER's compiled-dispatch helper: it keeps the first real launch's `CompiledFunction` and calls it after that. A replayed CUDA graph does not call the launchers, so decode with CUDA graphs (the default) does not change. ## Modifications - These launches use `_run_compiled` from `aiter.ops.flydsl.kernels.tensor_shim`, with the runtime addresses, scalars and stream passed positionally: - DeepSeek-V4.1: K1, K2, FFN and K2i. - GLM-5.3: the layer. - GLM-5.3-Flash: the layer and the head. - MiniMax-M3: the launch. MiniMax-M3.1 already caches its compiled FFN callable. - `compile_only` in the four runners uses AITER's `_preload_compiled`. AITER handles the no-dispatch preload and restores the environment. - The GLM runners import both helpers where they use them. AITER probes the GPU when it is imported, and `test/unit` imports these modules on CPU. - `test/kernels/test_compiled_dispatch.py` checks the dispatch that the runners rely on: - a preload never runs the kernel; - the first real call runs it once; - later calls skip JIT dispatch, also inside a graph. - AITER requirement: `_preload_compiled`, from ROCm/aiter#5346 (`354e278e5`) onwards. ROCm/aiter#6946 is not needed, because this repository never calls `_run_compiled` while `COMPILE_ONLY` is set. ## Accuracy Tests The kernels and their arguments do not change, and the outputs are bitwise equal to main's: - GLM-5.3-Flash: the eager step of layers 3-43 at TP4, 4 and 8 rows, on every rank. - DeepSeek-V4.1 K1 through `DSV41MonoRunner.k1`: at TP2; widths 1, 6, 24 and 48; fold, settled and index. The outputs match after the first and third eager call and after a graph replay, in 24/24 cases. E2E on MI355X: - `test_glm53_flash_server.py` (TP4): 5/5. GSM8K past the largest width scores 0.995. - `test_glm53_server.py` (FP8, TP8): 8/8. GSM8K scores 0.99 for decode and 0.99 with MTP, with an accept length of 3.3-3.4. - `test_minimax_m3_server.py` (TP4, with sgl-project/sglang#41998): 7/9. GSM8K scores 0.925 with EAGLE3 and 0.885 for decode. The 2 failing cases also fail on main: the test expects the spec name `minimax_m3_hip`, but the spec has been named `minimax_m3_family_hip` since #56. After the GLM import change, CPU-only `test/unit` gives the same result as on main, and GLM-5.3-Flash stays bitwise equal to main at TP4. ## Speed Tests and Profiling MI355X at TP4, with a fresh server for each arm and two rounds. The load is random 2048 tokens in and 256 out. The table gives median TPOT in ms, main -> this PR: | Model | CUDA graphs, c1 / c4 | `--disable-cuda-graph`, c1 / c4 | | --- | --- | --- | | GLM-5.3-Flash | 2.520 -> 2.523 / 3.032 -> 3.029 | 11.28 -> 2.72 / 11.75 -> 3.19 | | GLM-5.3 (FP8) | 6.812 -> 6.812 / 9.351 -> 9.342 | 20.06 -> 7.05 / 20.43 -> 9.39 | | MiniMax-M3 + EAGLE3 | 1.430 -> 1.426 / 4.118 -> 4.174 | 6.29 -> 3.31 / 9.58 -> 6.38 | - **With CUDA graphs**, TPOT does not change (within ±1.4% at every covered width). The gain applies only to steps that run without a graph. - **Host time** of a GLM-5.3-Flash launch goes from 182 µs to 14.7 µs. An eager step of layers 3-43 goes from 7.50 ms to 2.54 ms, which equals the step's graph replay time. Setup: image `rocm/sgl-dev:v0.5.22-rocm10-mi35x-20261026`, with SGLang `5ef99013ac`, AITER `285579b62` and FlyDSL 0.4.0. ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change. - The dispatch: `test_compiled_dispatch.py`. - DeepSeek-V4.1 FFN and K2: `test_deepseek_v41_kernels.py`. K2i: `test_deepseek_v41_k2i.py`. - GLM-5.3-Flash: `test_glm53_flash_kernels.py`. - MiniMax-M3 `compile_only`: `test_minimax_m3_compile.py`. - Only the e2e tests reach the GLM-5.3 and MiniMax-M3 launches. - No committed test reaches `DSV41MonoRunner.k1`, because `check_k1.py` calls the K1 launcher directly. K1 was checked bitwise against main, as described above. - [x] The SGLang commit or PR this was tested against is named above. - [x] Accuracy and speed results are attached. - [ ] A new model adds one line to `.github/CODEOWNERS` that names its owner. (No new model.) --------- Co-authored-by: Jin Pan <jin.pan@amd.com>

#6362, which this branch was stacked on, landed squashed as fc0a013. Every file except the 12 of the replay commit is therefore taken as main has it, which also leaves the PR's diff against main at those 12. Conflicts: - aiter/aot/flydsl/gemm.py: #6060 taught the hand-written AOT the gfx950 1x32 BMM (mxscale_bmm_groups_of_block) and to read a blank or 0 w_scale_block cell as the runtime's loaders do (mxscale_w_scale_block_cell). The replay calls the public entries, which take the scale block from the caller's w_scale, so none of that code returns. The cell reading does: _blocks() reads a cell through mxscale_w_scale_block_cell, 1x32 for the mxfp8 table and 128x128 for the others, as their loaders default. A 0 cell made it raise before. Main also brought #6060's BMM changes and #6946's launcher cache, which this branch had not been validated against. Ten rows from the tables they touch (gfx950 1x32, 32x32 and 128x128 through the 2D and batched entries, gfx950 mxfp8, gfx1250 batched and a8w8 bpreshuffle) replay into a cache that --check accepts, and on gfx950 the real calls of those rows and of the sizes between them (9 calls) all run from it under FLYDSL_RUNTIME_RUN_ONLY=1. Co-authored-by: Cursor <cursoragent@cursor.com>

Replaces the temporary workaround for the cached launcher (ROCm#6946), which flattens the grouped GEMM's scale arguments with .view(-1). A non-contiguous tensor cannot do that, and MegaMoE's expert-grouped dispatch hands over exactly one: each row's e8m0 scales sit inside its wire row, read by a_scale_row_stride_bytes. The kernel reads these arguments only through their base pointer. tensor_shim gains flat_arg(), beside ptr_arg(): a contiguous tensor is .view(-1), as before, so no other caller changes; a strided one becomes the run of elements from its first to its last, at the same base. tests/test_grouped_gemm_strided_scales.py checks flat_arg and what the launcher hands the compiled kernel, without a GPU. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

flydsl_grouped_gemm_a8w4_masked normalizes a_scales with view(torch.int32).view(-1) since #6946. On the compact mega_moe path (stage1_fused, prequantized fp4 wire) the GEMM1 activation scales are a strided view into the received wire rows (payload | e8m0 | pad), so view(-1) raises "view size is not compatible with input tensor's size and stride" at the first forward. The kernel reads the scales by base address and takes the row pitch as a compile-time argument, so keep the address with a rank-1 as_strided view instead of view(-1) (or reshape, which would copy). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ets (#7095) (backported onto the 7061cbd-based k3-455-tp8-1029: the TP route call keeps this base's direct launcher call; upstream #6946's _dispatch_compiled is not in it)

…tches, mori combine path, current mori EPv2 (#6430) * MegaMoE: stage1_fused/stage2_fused switches, mori combine path * bench_mega_moe: fix kernel table on multi-stream graphs * MegaMoE: follow mori tokoff-ext and selfFirst defaults * bench_mega_moe: --fused_shared_experts as ATOM DP+EP runs them * MegaMoE: follow mori#726 and size mori's fp4 combine flags Four fixes to the mori dispatch/combine path (stage1_fused/stage2_fused off), each from driving mori's plans without its op layer: - dispatch: since ROCm/mori#726 the plan takes tok_off_ext and derives selfFirst itself; self_first_enabled() is gone, so importing it fails. - tuning lookup: key it on the dispatch wire the way mori's op layer does ("fp4_disp_bf16_comb" for fp4). Under the hard-coded "bf16" key topk 9 matched no entry and both dispatch and the fp4 combine fell back to the single-shot default. - comb_inp: add the overlapped fp4 combine's u64 flag per landing row (four ranks up). Its 24-warp geometries write them past the region, i.e. past the arena: a page fault at 16384 tokens a rank. - close(): synchronize first, so a launch still in flight does not read a window that has been freed. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: zero-copy hand-off from fused_moe to mori's combine With stage2_fused off, fused_moe's rows went into a fresh tensor and then into mori's staging rows (comb_inp): the bf16 combine copied them in its own kernel, the fp4 one through a host-issued copy whenever the recv bound was below max_recv -- every decode step under ATOM. Have fused_moe write them there directly. fused_moe(output=) already existed but the gfx1250 grouped path ignored it and copied at the end; thread it down to the gather-reduce. mori's bf16 combine then sees its input at the staging base and skips the copy, and the fp4 send reads the staging rows in place. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: follow mori#745/#746 on the mori dispatch/combine path mori's gfx1250 dispatch now reserves its slots once per destination (the aggregated reservation, ROCm/mori#745) and hands the push combine its landing offsets through the dispatch. MegaMoE drives the plans directly, so it has to size and pass what mori's op layer does: - dest_pe_token_counter holds 2 * world_size ints: send counts, then per destination the reservation base and the push combine's landing offset. It held world_size; the dispatch wrote the second half into allocator padding. - recv_num signals are 8 bytes a source. The region held 4 bytes a source; the arena's 256-byte alignment hid the overrun. - The combine reads its landing offsets from dest_pe_token_counter, which MegaMoE did not pass, so the mxfp4 combine (stage2_fused=False) got a null pointer there. - MORI_EP_TOKOFF_EXT=1 turns the aggregated reservation off, so a mori push combine is refused with it, as mori's op layer refuses it. - The push landing rows come from mori's push_land_rows(): one per (token, destination rank), max_tokens * min(world_size, topk) rather than world_size * max_recv. - The mori dispatch launcher drops total_recv.zero_(): gfx1250's dispatch overwrites the count (since ROCm/mori#599), and mori's op layer skips the same zero there. Against a mori from before #745 the landing rows keep the old count and the two larger buffers have unused tails. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: bind mori's dispatch staging buffer (ROCm/mori#703) mori#703 moves the gfx1250 dispatch's staging arrays out of each kernel module's .bss into one buffer the caller allocates and binds as staging_base; an unbound plan is refused on the host. MegaMoE drives EpDispatchPlan directly, so it allocates ep_staging_bytes() once and binds it to every dispatch plan (they never run at the same time), and drops it after the plans on close(). A mori without ep_staging_bytes keeps the .bss arrays and is left as is. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: keep accepting dispatch_backend stage1_fused now picks the dispatch, and the dispatch_backend argument and config field went away with that. Callers still pass it (ATOM main passes dispatch_backend="mori" on a quantizing wire; ROCm/ATOM#2384 passes "flydsl" or "mori" alongside stage1_fused) and log mega._config.dispatch_backend, so both would fail on this aiter. Accept the argument again and check it against stage1_fused: "mori" is stage1_fused=False, "flydsl" is the compact plan. A "flydsl" without stage1_fused asked for the FlyDSL token-major dispatch, which no longer exists, so it is refused rather than run on mori. MegaMoEConfig gets a read-only dispatch_backend derived from stage1_fused. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * ci: run the MegaMoE bench on mori's dispatch stage1_fused=false with stage2_fused=true is mori's fp4 dispatch with the fused combine (the bench's auto wire is fp4 for a4w4_mxfp4), the configuration DeepSeek-R1 serves with. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: gfx1250 behaviour tests ahead of the refactor test_mega_moe_gfx1250.py pins what a change to MegaMoEGfx1250's dispatch/combine wiring must keep, through the constructor and members its callers use today: - every dispatch/combine pairing matches the fp32 reference, and a captured graph replays bit for bit what eager produced; - the pairings MegaMoE rejects fail before the symmetric arena is allocated; - mori's combine reads fused_moe's rows in place (zero-copy staging); - close() releases the plans before the window, the window before its memory; - the members ATOM's Triton experts reach (_dispatch, _combine, _arena, _config). It spawns its own ranks and skips (exit 0) off gfx1250, below two GPUs or without mori, so CI's MULTIGPU job can run it on gfx950. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: build from a MegaMoERecipe and a MegaMoEGeometry A caller now says what to run once, as a recipe, and MegaMoE builds it: recipe = MegaMoERecipe(dispatch=DispatchKind.MORI, combine=CombineKind.FUSED_SCATTER, dispatch_wire="fp4", combine_wires=("bf16", "fp4"), activation=..., quant_type=...) geometry = MegaMoEGeometry(rank, world_size, model_dim, inter_dim, experts, topk, max_tokens_per_rank) mega = MegaMoEGfx1250.from_recipe(recipe, geometry, communicator=comm) - MegaMoERecipe holds the transport (dispatch and combine kind, dispatch wire, the combine wires to build) and the expert-GEMM recipe; frozen and hashable, so a caller can key its cache of instances on it. MegaMoEGeometry holds the model and EP shape. - MegaMoERecipe.from_legacy is the one reader of stage1_fused, stage2_fused, dispatch_backend and the $MEGA_DISPATCH_WIRE / $MEGA_COMBINE_QUANT defaults. The keyword constructor goes through it, so both entries build the same thing. - MegaMoEConfig is now that recipe resolved against a geometry: frozen, with the same properties as before (stage1_fused, stage2_fused and dispatch_backend included, for the callers that read them). The dispatch schedule it used to have written into it lives on the instance; the dispatch_block_num / dispatch_warp_num_per_block fallback it carried was never reached. - forward() takes combine_wire= (bf16 | fp8 | fp4), the recipe's vocabulary, alongside combine_quant=. Every check and its message is unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: probe mori once, and decide everything before the arena exists _MoriCapabilities gathers what MegaMoE needs to know about the installed mori -- the plan classes, the tuning lookup, and each helper a mori PR added (scale_stride_bytes #593, push_wire_nbytes #719, push_land_rows #745, ep_staging_bytes #703, TokOffExt.wanted #726, self_first_enabled) -- into one cached probe, and the decisions that hang off them into its methods: the tuned schedules, the landing-row sizes, the staging size, where the slot word lives. The ImportError/hasattr branches are gone from the pipeline; the RuntimeError messages for an mori that is too old are unchanged. _plan_transport runs all of it, plus the fused reduce's tiles and mori combine's block/epoch-slot check, before the symmetric arena is allocated. A configuration that cannot run -- mori missing, MORI_EP_TOKOFF_EXT=1 with mori's fp4 combine, no reduce tile, too many combine blocks -- now fails with nothing allocated instead of leaking the arena. TokOffExt.wanted, collective when MORI_EP_TOKOFF_EXT is on, moves ahead of the arena's own collective allocation on every rank alike. The compact dispatch with the fused combine still never loads mori's ops. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: one arena layout, and required_vmm_bytes from it The arena's regions now come from one formula, _arena_layout: the dispatch says what it writes (_dispatch_regions: the routing both kinds share, then its own rows), the combine what it stages through (_combine_regions), and _ArenaLayout.merge lays them down in the order the pipeline always used, so every pairing keeps the offsets it had. SymmetricArena takes the layout (or a region list, as before) and frees its memory if the window registration fails. MegaMoEGfx1250.required_vmm_bytes(recipe, geometry) reports the same layout, rounded the way ccoMemAlloc rounds an allocation -- the per_rank_vmm a caller needs for this MegaMoE, before any communicator exists, and exactly what the constructor then allocates. It is MegaMoE's share only; a communicator that also carries mori's own EP ops needs theirs on top. required_vmm_bytes(world_size=, hidden_dim=, max_tokens_per_rank=, experts=, topk=, stage1_fused=), the keyword form #6525 proposed and ATOM's yanbo/mega already calls, still works: it names neither wires nor combine nor inter_dim, so it returns the largest any recipe with that dispatch takes. _compact_gemm_align_m reads activation, quant type and inter_dim off the config instead of taking them separately. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: dispatch and combine stages under one dispatcher; linear forward The transport is now two stages the recipe picks, over one shared arena: _MoriDispatch / _CompactDispatch tokens out to their experts' ranks _MoriCombine / _FusedScatterCombine results back, summed per token Each declares its arena regions (regions(), pure -- what _arena_layout and required_vmm_bytes merge), plans in __init__ with nothing allocated, and allocates in build(). They meet only through _EpResources (the arena, the routing state both ends read, TokOffExt), the _DispatchBatch the dispatch hands the expert GEMM, and the _ComputeTarget the combine has it write into. _EpDispatcher picks the two from tables, runs every check before the arena is allocated, and frees what was built when a later step fails -- the arena and plans used to leak. close() goes: synchronize, combine plans, dispatch plans, staging, arena, TokOffExt. forward() is now: pick the step's combine wire, check the inputs, dispatch, prepare, compute, finish. The compact plan's machinery -- slots, side stream, per-step geometry (_CompactStep, now immutable) -- lives in _CompactDispatch, and the next layer's plan launches at the end of its dispatch, the same place on the GPU streams as before. The recv views are made once at build instead of per step. The TDM dispatch keeps only its compact arm; the token-major arms were unreachable once stage1_fused=False went through mori. Kept for the callers that reach past forward(), ATOM's Triton experts among them: _config, _arena, _dispatch() and its 5-tuple, _combine(), _recv_dispatch_scales(), prefetch_compact_plan(), warmup_compact_plan(). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: size mori's regions from describe(); the bench sizes cco from the recipe On a mori with Plan.describe() (ROCm/mori#844), the mori stages ask it for their arena regions and the dispatch's staging bytes -- the numbers mori's own kernels and its bind-time size check use -- instead of keeping formulas that mirror them. Each stage builds its plan request in one place (request()), which both describe() and the plan take. describe() compiles nothing and resolves against gfx1250 by name, so required_vmm_bytes stays off the GPU; the slot word's kwargs stay out of it, since where that word lives sizes nothing and asking would be collective. An older mori keeps the formulas. SymmetricArena.size() lets a #844 plan check its regions at bind. bench_mega_moe builds from a MegaMoERecipe and sizes its communicator with required_vmm_bytes(recipe, geometry) instead of a fixed 16 GiB. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: CPU tests for the recipe and the arena it sizes tests/test_mega_moe_gfx1250_recipe.py needs no GPU and skips where MegaMoE does not import. It checks how the legacy flags map to a recipe and that the rejections keep their wording, and lays the arena out over a grid of geometries and every pairing against the region list #6430 built inline -- through a fake mori both with and without Plan.describe(), the describe one following mori's rules, which add no overlapped-reduce flags on a wide EP. required_vmm_bytes is held to the layout as cco rounds it, and its keyword form to at least that. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: bind mori's dispatch staging as a tensor A mori with ROCm/mori#844 checks a staging tensor against the plan's stagingBytes, device and alignment at bind, and holds it; an address is taken on trust. An older mori reads the tensor's address at each launch, as it did the int. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * TEMP: grouped GEMM takes a strided scale view again (fixup for #6946) #6946 flattens the grouped GEMM's scale arguments with .view(-1), which a non-contiguous tensor cannot do. MegaMoE's compact rows hand over exactly that -- each row's e8m0 scales sit inside its wire row -- so every stage1_fused MegaMoE step raises. The kernel only takes the base address and walks the rows by a_scale_row_stride_bytes, so a strided scale view becomes a one-dimensional view over the same base, as it was before #6946. Temporary: drop this commit once main fixes it. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * bench_mega_moe: take out the kernel-table and --fused_shared_experts changes They are bench fixes rather than part of the MegaMoE recipe work, and move to #7074. This reverts "bench_mega_moe: fix kernel table on multi-stream graphs" and "bench_mega_moe: --fused_shared_experts as ATOM DP+EP runs them". Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * FlyDSL grouped GEMM: pass a strided scale view at its base address Replaces the temporary workaround for the cached launcher (#6946), which flattens the grouped GEMM's scale arguments with .view(-1). A non-contiguous tensor cannot do that, and MegaMoE's expert-grouped dispatch hands over exactly one: each row's e8m0 scales sit inside its wire row, read by a_scale_row_stride_bytes. The kernel reads these arguments only through their base pointer. tensor_shim gains flat_arg(), beside ptr_arg(): a contiguous tensor is .view(-1), as before, so no other caller changes; a strided one becomes the run of elements from its first to its last, at the same base. tests/test_grouped_gemm_strided_scales.py checks flat_arg and what the launcher hands the compiled kernel, without a GPU. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: stage contracts, mori sizing from describe() only, from_recipe first The recipe - The expert GEMM's fields move into ExpertComputeConfig, nested as MegaMoERecipe.compute; the recipe's own fields are the transport's. - Which pairings run is declared by the stages: a dispatch names the wires it carries and the BatchLayout it delivers (token-major or expert-grouped), a combine the layouts it accepts and the wires it offers. The recipe checks against those, and its errors name recipe fields (dispatch=, combine=, dispatch_wire=, combine_wires=). Each stage's geometry check (check()) runs from MegaMoEConfig. - from_recipe is the entry point. The keyword constructor warns DeprecationWarning; the underscored members ATOM's Triton experts use are documented as kept for them. The facade reads geometry and compute fields through properties rather than copying them. The stages - DispatchStage and CombineStage (Protocols) say what each provides; every stage has regions(config, caps) and check(config), and is built as Stage(config, caps). The compact- and mori-only attributes are gone from the shared base. - Each dispatch stage owns the stride its e8m0 rows arrive at (scale_stride), replacing MegaMoEConfig.dispatch_scale_dst_nbytes. - The batch carries its layout explicitly, and an expert-grouped one an ExpertGrouping; the combine hands the grouped GEMM its dispatch context in the _ComputeTarget, so the expert GEMM does not branch on the layout. - Whether the combine needs mori's aggregated reservation and whether the dispatch makes it are each stage's to say; the dispatcher refuses the pair before the arena exists. - One schedule pick (_pick) serves both mori plans and the compact kernels; _quantize returns the scale rows' address instead of leaving it in state. - _EpDispatcher has dispatch/prepare/combine/prefetch/warmup and the recv views, so forward() does not reach into the stages, and close() releases the rest even when one part raises. mori - _MoriCapabilities needs Plan.describe() (ROCm/mori#844). Regions, staging bytes and the scale stride come from there, so the copies of mori's formulas and the older-mori branches are gone; such a mori also overwrites total_recv in its dispatch, which is why nothing zeroes it. - _XDB_FLAG_SLOTS, TokOffExt and the plan constructors are reached only through _MoriCapabilities. Checks - An out-of-range recv_token_bound (and next_recv_token_bound) raises on the compact dispatch as on mori's, in dispatch, prefetch and warmup. - a1_scale on a quantizing dispatch wire raises ValueError (was an assert). - The keyword form of required_vmm_bytes raises when no candidate recipe runs at the geometry, instead of returning 0. - _MoriCombine.prepare documents why writing fused_moe's rows straight into comb_inp is safe: the next dispatch does not finish until every peer has dispatched, which each does only after its own combine. Tests: the CPU test fakes mori from describe() alone and covers the stage contract messages, the bound checks and the keyword-bound refusal; the multi-GPU test builds through from_recipe and checks the deprecated keyword constructor against it. bench_moe_gemm_a4w4_cudagraph builds through from_recipe. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: split mega_moe.py into recipe, mori_caps, arena, stages, transport A move, no change in behaviour: each block of mega_moe.py lands in one module as it was, and mega_moe.py keeps only MegaMoEGfx1250. - recipe.py: the wires, DispatchKind / CombineKind, ExpertComputeConfig, MegaMoERecipe, MegaMoEGeometry, MegaMoEConfig. - mori_caps.py: _MoriCapabilities and mori's region names. - arena.py: SymmetricArena, the region order, _ArenaLayout, cco's rounding. - stages/base.py: DispatchStage / CombineStage and what the stages share (routing, resources, batch, compute target, schedule pick). - stages/dispatch_mori.py, dispatch_compact.py, combine_fused.py, combine_mori.py: one stage each; stages/__init__.py maps each kind to its stage. - transport.py: _EpDispatcher and the one arena layout. The only edit: recipe.py looks the stages up when asked for (the stages import recipe.py). mega_moe.py still exports the package's public names, and Routing, which a bench imports from there. The CPU test imports the private names it checks from the modules that define them. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: compact dispatch bound and plan fixes, wire-picked GEMM A dtype - Compact dispatch: with no recv_token_bound the step is sized for the arena's capacity; a plan started for another alignment is refused; the plan's segment and wave limits are checked before the arena exists; plans are compiled when created, ahead of any graph capture. - A quantizing dispatch wire sets fused_moe's quant_dtype_a. - bench_mega_moe runs on gfx1250 only. - Tests: fp8 dispatch and fp8 combine cases, rank-local failures raised on every rank, CPU tests skip only without flydsl. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE CPU test: skip without mori as well - The package's kernel modules import mori.cco, so the test skips when mori is not installed instead of failing collection. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: yanboshao <yashao@amd.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

#6362, which this branch was stacked on, landed squashed as fc0a013. Every file except the 12 of the replay commit is therefore taken as main has it, which also leaves the PR's diff against main at those 12. Conflicts: - aiter/aot/flydsl/gemm.py: #6060 taught the hand-written AOT the gfx950 1x32 BMM (mxscale_bmm_groups_of_block) and to read a blank or 0 w_scale_block cell as the runtime's loaders do (mxscale_w_scale_block_cell). The replay calls the public entries, which take the scale block from the caller's w_scale, so none of that code returns. The cell reading does: _blocks() reads a cell through mxscale_w_scale_block_cell, 1x32 for the mxfp8 table and 128x128 for the others, as their loaders default. A 0 cell made it raise before. Main also brought #6060's BMM changes and #6946's launcher cache, which this branch had not been validated against. Ten rows from the tables they touch (gfx950 1x32, 32x32 and 128x128 through the 2D and batched entries, gfx950 mxfp8, gfx1250 batched and a8w8 bpreshuffle) replay into a cache that --check accepts, and on gfx950 the real calls of those rows and of the sizes between them (9 calls) all run from it under FLYDSL_RUNTIME_RUN_ONLY=1. Co-authored-by: Cursor <cursoragent@cursor.com>

coderfeli requested a review from a team October 26, 2026 12:10

github-actions Bot changed the title ~~Cache compiled FlyDSL launchers on host paths~~ [FlyDSL] Cache compiled FlyDSL launchers on host paths Oct 26, 2026

github-actions Bot added the FlyDSL label Oct 26, 2026

coderfeli force-pushed the codex/flyc-compile-host-overhead branch from 568687a to 8f75570 Compare October 26, 2026 14:36

Cache compiled FlyDSL launchers on production host paths

efcceb4

coderfeli force-pushed the codex/flyc-compile-host-overhead branch from 8f75570 to efcceb4 Compare October 26, 2026 14:48

coderfeli added 2 commits October 26, 2026 15:02

Keep host-launch caching scoped to live FlyDSL paths

a8261c1

Preserve compiled launch safety across buffer ranks and compile-only …

318e92c

…mode

zufayu requested a review from yadaish October 27, 2026 01:13

zufayu assigned yadaish Oct 27, 2026

coderfeli added 2 commits October 27, 2026 02:08

Remove extra dispatch tests and redundant table flattening

1920c18

Cache remaining gfx1250 GEMM projection and MegaMoE launchers

73db860

yhl-amd mentioned this pull request Oct 27, 2026

[gfx1250] JitFunction.__call__ cached fast path ROCm/FlyDSL#1331

Closed

coderfeli merged commit 2bafed4 into main Oct 27, 2026
63 checks passed

coderfeli deleted the codex/flyc-compile-host-overhead branch October 27, 2026 06:50

This was referenced Oct 27, 2026

perf(mono): reuse AITER compiled FlyDSL dispatch ROCm/ATOM#2741

Merged

[Perf] mono: Reuse AITER compiled FlyDSL dispatch sgl-project/sgl-amd-mono#66

Merged

yhl-amd mentioned this pull request Oct 27, 2026

[Triton/Gluon] [FlyDSL] [gfx1250] Cut remaining host overhead in GEMM wrappers (v2 of #6948, on top of #6946) yhl-amd/aiter#10

Closed

This was referenced Oct 27, 2026

[Triton/Gluon] [FlyDSL] [gfx1250] Cut remaining host overhead in GEMM wrappers #6948

Closed

[FlyDSL] [JIT] [gfx1250] Cut remaining host overhead in FlyDSL GEMM/MoE wrappers #7006

Merged

JiaoliangYu mentioned this pull request Oct 27, 2026

fix(gfx1250): honor real token counts in TP grouped MoE routing #6989

Closed

jhchouuu mentioned this pull request Oct 27, 2026

[Triton/Gluon] [FlyDSL] [CI] [MoE] MegaMoE gfx1250: stage1/stage2 switches, mori combine path, current mori EPv2 #6430

Merged

amd-ruitang3 mentioned this pull request Oct 28, 2026

bench_mega_moe: fix kernel table on multi-stream graphs, add --fused_shared_experts #7073

Closed

yanboshao mentioned this pull request Oct 29, 2026

[FlyDSL] fix(mega_moe): strided A scales on stage1_fused grouped GEMM #7143

Closed

1 task done

This was referenced Oct 31, 2026

[CI Monitor] Daily Report - 2026-10-31 bingxche/sglang-ci-bot#263

Open

[CI Monitor] Daily Report - 2026-11-01 bingxche/sglang-ci-bot#265

Open

coderfeli mentioned this pull request Nov 1, 2026

[FlyDSL] Add flyc.compile to skills. and cache HERD host launches #7271

Merged

bingxche mentioned this pull request Nov 2, 2026

[CI Monitor] Daily Report - 2026-11-02 bingxche/sglang-ci-bot#266

Open

bingxche mentioned this pull request Nov 3, 2026

[CI Monitor] Daily Report - 2026-11-03 bingxche/sglang-ci-bot#269

Open

This was referenced Nov 4, 2026

[CI Monitor] Daily Report - 2026-11-04 bingxche/sglang-ci-bot#271

Open

[CI Monitor] Daily Report - 2026-11-05 bingxche/sglang-ci-bot#273

Open

Navigation Menu

Repository navigation

[FlyDSL] Cache compiled FlyDSL launchers on host paths - #6946

[FlyDSL] Cache compiled FlyDSL launchers on host paths#6946
coderfeli merged 5 commits into
mainfrom
codex/flyc-compile-host-overhead

coderfeli commented Oct 26, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 26, 2026

Uh oh!

yadaish commented Oct 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

coderfeli commented Oct 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Host measurements

Validation

Global dispatch audit

Uh oh!

github-actions Bot commented Oct 26, 2026

🏷️ CI Guide

Uh oh!

yadaish commented Oct 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

coderfeli commented Oct 26, 2026 •

edited

Loading