Repository navigation
[FlyDSL] [JIT] Build the GEMM, conv, chunk_gdn_h and FP8 FMHA AOT caches by replaying the runtime - #6956
Merged
Conversation
_gen_opus_moe_stage1_a8w4_fake_tensors and _gen_opus_moe_stage2_a8w4_decode_fake_tensors lacked the scalar parameters of the ops they are registered for, so any call under FakeTensorMode raised TypeError (10 of 18 and 9 of 14 positional arguments accepted). Signed-off-by: zhimding <zhimding@amd.com>
The MoE AOT scripts (moe.py, mxfp4_moe.py, fhmoe.py) restated the runtime's kernel selection and launch parameters, so every kernel or wrapper change needed a matching AOT change, and drift between the two left kernels the runtime uses uncompiled. moe.py now builds what a caller passes fused_moe for each tuned row as FakeTensors and calls the runtime's own _fused_moe_impl (_fhmoe_impl for shared-expert rows) inside dry_run, which sets COMPILE_ONLY=1 and forbids the GPU runtime; every FlyDSL launch on that path is compiled into the cache instead of run. The AOT side only enumerates what a row does not fix: expert bias, gate/up layout and SiTUv2 betas on the caller side, and the M values get_padded_M maps to the row. Each target arch replays in its own process with the GPUs hidden, and --check verifies a cache without compiling. Runtime-side support: - torch_guard: scalar-returning custom ops can run their implementation under FakeTensorMode, whose fakes carry no value (off by default). - chip_info: with no GPU to detect, a single explicit GPU_ARCHS stands in for the live GPU. - fhmoe: fhmoe_ is split into a plain _fhmoe_impl and the decorated op. mxfp4_moe.py and fhmoe.py are removed; moe.py replays their rows. Signed-off-by: zhimding <zhimding@amd.com>
A from-scratch build forked one worker per job (9903) from the full-import arch process, so it was bound by the parent's fork rate, and the variants of one row ran concurrently and waited on each other's compile locks. A worker now replays up to 12 jobs of one row in turn, equal-sized batches are shuffled so that rows of one model do not start together, and the arch processes run side by side on a split worker budget. Fresh full build at 96 workers: 537 s -> 187 s; rebuild with a complete cache: 157 s -> 21 s. The compiled artifact set is unchanged. Signed-off-by: zhimding <zhimding@amd.com>
mega_moe.py restated MegaMoEV2's buffer geometry, placeholder tensors, Stage2 variant expansion and combine geometry lookup, so MegaMoE kernel changes needed matching AOT changes. It had also drifted: it compiled only the mask_topk_ids=False combine kernels, so a caller passing mask_invalid_slots=True compiles online. Each job now builds MegaMoEV2 for one rank of a deployment profile inside dry_run, with FakeTensor weights and a mori shmem that hands out FakeTensors and fake peer addresses, and calls its preload_aot_bundles() and quantize(). The module keeps only the deployment profiles, and one job covers both stages, so 192 jobs become 96. dry_run takes an optional device name, since the combine launch geometry is tuned per GPU model, and stands in for torch.cuda Stream and Event. Signed-off-by: zhimding <zhimding@amd.com>
Merging read each table with type inference, so a column with a blank cell turned float and every integer in it was written back as N.0: 6 of the 15 merged tables carry such cells, among them tuned_fmoe and tuned_grouped_fmoe. The gfx1250 grouped MoE runtime parses integers from that text, so since #6090 added a second grouped table it raises ValueError on every row of the merged one. Read the tables as text so the merge reproduces each cell. Shapes are still compared on numeric values, so 256 and 256.0 stay one shape, and the fastest duplicate is still picked on a numeric us. Signed-off-by: zhimding <zhimding@amd.com>
grouped_moe.py compiled only the auxiliary kernels around the grouped GEMM: the runtime has since moved the GEMM to the TDM batched kernel, which got no AOT coverage and compiled on first use, and most of the auxiliary variants it did compile were for the removed non-TDM path. On gfx1250 _fused_moe_impl hands the grouped table's rows to FlyDSL's grouped GEMM, so moe.py now replays those rows as well, with the gate/up-interleaved weights and n32k4 scales gfx1250 callers pass. Rows tuned for EP only (topk -1, or ep_fused) are skipped: the runtime picks them only under a caller's expert mask or MegaMoE's stage-2 scatter, and the table records neither the global expert count nor the real topk. The two-stage host query is no longer preloaded when an arch replays only grouped rows, since that path never makes it. grouped_moe.py and OpKind.GROUPED_MOE are removed. Signed-off-by: zhimding <zhimding@amd.com>
A row serves a bucket of M, and the replay compiled the bucket's two ends, which covers any choice that switches once. The gfx1250 stage-1 quant's K split switches several times in one bucket (7168x768, E384, topk 6, token 2048: K split 7, then 4, then the compact quant), so M in the middle asked for a kernel the cache did not have. Where a variant's two ends request different FlyDSL kernels, identified by FlyDSL's own cache key, the worker now bisects the bucket to every M at which they change and replays it. The default build replays 775 more M (573 on gfx1250, 202 on gfx950) and compiles nothing new on gfx942/gfx950; a densely sampled RUN_ONLY replay of every bucket now finds no missing kernel on any of the three archs. Signed-off-by: zhimding <zhimding@amd.com>
mxfp4_moe.py stays deleted. #5804 taught it the use_valid_token_count GEMM2 variant, which the runtime selects when the caller passes num_local_tokens; the next commit has the replay enumerate that caller choice instead. Signed-off-by: zhimding <zhimding@amd.com>
Callers behind an EP dispatch pass the received token count as num_local_tokens. Since #5804 the runtime then bounds the layout-v2 GEMM2 atomics by it, a compile-time choice the old mxfp4_moe.py had to be taught by hand; on inline-sort rows the runtime instead drops the tuned config for heuristic FlyDSL kernels that no AOT compiled. Each row is now replayed with num_local_tokens too. Of the old script's use_valid_token_count kernels, the replay has every one the runtime reaches; the 227 it lacks come from bias variants of Silu/SiTUv2 rows (the runtime enables bias for Swiglu only since #6376), from inline-sort rows the runtime abandons under num_local_tokens, and from a token=3 row that get_padded_M never selects. Signed-off-by: zhimding <zhimding@amd.com>
…rch replays The wheel build's AOT step failed both replays: - The per-arch MoE replays run side by side and share module_moe_mxfp4_aux, whose host query they call. The first to build it built it for its own arch. The others then refused it, since the runtime rebuilds a module whose device code lacks the running arch, and with the build lock just released did not rebuild it either, so the gfx950 replay died. dry_run now skips that check: only host code runs under it. - MegaMoE replayed in setup.py's own process, which loads aiter/jit/core.py a second time as jit.core. The first load of a C++ module there is recorded in the other copy, so every job raised KeyError: 'module_aiter_core'. It now runs in a process of its own, like the MoE replay, with GPU_ARCHS=gfx950 and the GPUs hidden. The environment of those processes moves to common.replay_env. Signed-off-by: zhimding <zhimding@amd.com>
Conflicts: - aiter/aot/flydsl/mxfp4_moe.py: kept. The MXFP4 tuner's --fast-scan (#6405) precompiles its candidates with jobs_from_rows and compile_job_batch. It stays out of the AOT build, which replays these rows through fused_moe; its docstring and the README say so. - aiter/jit/core.py: #6365 moved the duplicate-shape check into _raise_on_duplicate_shapes, shared with the single-file path. Comparing key columns as numbers, so 256 and 256.0 stay one shape, and ordering on a numeric `us` move there with it, and the single file is read as text like the merged ones, so a dedup rewrite keeps every cell as written. Signed-off-by: zhimding <zhimding@amd.com>
On a host without mori, _fake_shmem installs a stand-in module. A second
call in the same process then asked importlib.util.find_spec("mori"), which
raises ValueError for a module whose __spec__ is None. Reuse the stand-in
instead: it stays, since modules imported under it keep referring to it.
The AOT build runs each job in a process of its own, so it was not affected.
Signed-off-by: zhimding <zhimding@amd.com>
Conflict in aiter/aot/flydsl/mega_moe.py: #4980 added A4W4 MegaMoE and taught the old script its kernels (quant mode, FP4 output and scale shapes). The replay needs none of that, since MegaMoEV2 picks them itself; it only adds "a4w4" to DEFAULT_QUANTS, which takes the default build from 96 jobs to 192. The A4W4 replay passes all 96 jobs. Every one of the 3218 artifacts the old script builds for A4W4 is in its cache: 3082 byte-identical, 136 differing only in v_add3_u32 vs v_add_u32 + v_sub_u32 (2-7 asm lines), the FlyDSL/LLVM run-to-run variation also seen for A8W4. The other 144 are the mask_topk_ids=True combine kernels the old script never built. Signed-off-by: zhimding <zhimding@amd.com>
The per-arch replays import aiter side by side, and the first to take the build lock built module_aiter_core with its own arch's flags. When that was the gfx1250 replay, the build turned on cluster launch, which the HIP of ROCm 7.2 (CI's toolchain) does not declare, and all three replays failed. The module has no device code, so build it once, for the arch with the most jobs, before the replays start. Signed-off-by: zhimding <zhimding@amd.com>
Signed-off-by: zhimding <zhimding@amd.com>
The test cast fused_topk's route weights to bf16, which the gfx1250 gather-reduce required until it learned to read fp32 weights natively (#4150). Serving passes fp32 route weights, and so does the replay, so the cache holds only the fp32 gather-reduce and the test, which runs run-only by default, missed launch_moe_gather_reduce on every row. The bf16 and f16 weights the EP route kernel writes for gather-reduce are still exercised, now by the gather-reduce sweep of test_flydsl_moe_aux.py, which until now read only fp32 weights. Signed-off-by: zhimding <zhimding@amd.com>
Signed-off-by: zhimding <zhimding@amd.com>
#6589 made get_dtype_fp8() return uint8 whenever torch sees no device. The MoE replays hide the GPUs and name their target in GPU_ARCHS, so dtypes.fp8 became uint8 in every replay: fp8 activations failed `q_dtype_a == dtypes.fp8` and were quantized to fp4, MXFP8 rows found no input recipe, and the gfx1250 grouped rows compiled other kernels than the runtime requests without raising anything. Without a visible device, a single explicit GPU_ARCHS target now stands in for the GPU, as it already does in chip_info.get_gfx_runtime(); the runtime detector is still not imported. Signed-off-by: zhimding <zhimding@amd.com>
Since #6478 hands a caller's stage-1 bias to the CK-Tile split-k route, MXFP4 callers with gate/up-interleaved weights and a bias reach its interleaved epilogue, which cannot apply one. That is a missing feature, not a bad argument, so raise NotImplementedError like fused_moe's other unsupported combinations; the MoE replay records such caller-side variants as unsupported rather than failing on them. Signed-off-by: zhimding <zhimding@amd.com>
The Qwen3 8x2400B fp4 row tuned at token 32768 stages a 5.4 GB route-reduce output, beyond the 32-bit buffer addressing #6050 now guards: the runtime raises NotImplementedError before it launches anything, and op_tests/test_moe_2stage.py skips the case. The replay failed the build on it although there is nothing to compile. When the runtime serves a row's default caller at some of the row's token counts but refuses it at others, the refused token counts are now reported as [REFUSED]. Refused at every token count, the row still fails, as does any other error. Signed-off-by: zhimding <zhimding@amd.com>
A replay process learned its target arch only indirectly: replay_env hid the GPUs and set GPU_ARCHS, and get_gfx_runtime() and get_dtype_fp8() fell back to a single GPU_ARCHS entry whenever no GPU was visible. A replay process that could see a GPU took the machine's arch instead (gfx950's fp8 type for gfx942 kernels), and the fallback changed both functions for every GPU-less process with one GPU_ARCHS entry, which main's test_linux_runtime_errors_are_not_fallbacks rejects. replay_env now also sets AITER_AOT_REPLAY=1. Only under that flag do get_gfx_runtime() and dtypes.fp8 answer the one GPU_ARCHS arch, without probing a device; everywhere else they behave as on main. A replay process checks at entry that it was started that way and that dtypes.fp8 is its arch's type. The parent replays only the archs GPU_ARCHS names, as setup.py's wheel build sets it, intersected with the tables' gfx column; unset or native, every arch the tables have. CI's gfx942;gfx950 build thus no longer compiles gfx1250 MoE kernels, and MegaMoE runs only when gfx950 is listed. Signed-off-by: zhimding <zhimding@amd.com>
The wheel build replays only the archs GPU_ARCHS names, so a gfx942;gfx950 build leaves the gfx1250 rows out; GPU_ARCHS=gfx1250 builds them on their own. That build imports aiter and so needs its C++ core module, which ROCm 7.2.4 cannot build for gfx1250 (its HIP lacks the cluster launch API); the module has host code only, so one built for a CDNA arch serves. Signed-off-by: zhimding <zhimding@amd.com>
Silu/SiTUv2 never pass expert bias at serve time, and INTERLEAVE x bias hits CK-Tile's unsupported split-k epilogue, so keep those out of the replay instead of widening the runtime to report them as unsupported. Signed-off-by: zhimding <zhimding@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
2f85842 dropped every gate_mode=interleave x bias variant from the replay to keep CK-Tile's interleaved split-k epilogue, which cannot apply a bias, out of it. The same commit limits bias to Swiglu + MXFP4 rows, and since then no row the replay parses reaches that epilogue: the only Swiglu MXFP4 rows with a cktile_ stage 1 are the 14 gptoss_fp8fp4 rows tagged flydsl_fallback, which parse_csv skips. What the filter did drop is the path GPT-OSS a8w4 serves on gfx950 (gate/up-interleaved weights with expert bias), and every biased gfx1250 caller, which takes only the interleaved layout. The runtime serves all of those jobs and requests kernels the cache then lacked: CI's MI35X shard failed op_tests/test_moe_2stage.py on the gptoss_fp8fp4 rows with "no usable AOT cache for launch_mixed_moe_gemm1". A full replay of every arch now has 9440 jobs (324 more) and no failure, and its cache gains 31 kernel directories and loses none. The extra [SKIP]s are MiniMax-M3 a4w4 interleave x bias combinations at inter_dim 384, which the runtime refuses for want of a safe MXFP4 backend, as it already does for their other variants. Signed-off-by: zhimding <zhimding@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
A replay process imports aiter for the arch it replays, so its dtypes.fp8 matches the table's fp8 flavour. The pytest process imports aiter for its host instead: on gfx942 dtypes.fp8 is float8_e4m3fnuz, _caller_inputs finds no recipe for the HY4 float8_e4m3fn row, and both cases of test_fhmoe_aot_replay_forwards_caller_choices failed on CI's MI300X shard. Pin the alias as _mock_gfx950_fhmoe_metadata does. Signed-off-by: zhimding <zhimding@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Reading tuned tables as text keeps every cell as written, but _raise_on_duplicate_shapes still sorted on `comparable`, where a column with one unparseable written cell stays text. One us written as nan, on any row, made the lowest-us pick a text sort: "10.5" sorts before "9.5", so the slower duplicate was kept and the faster row deleted from its source table, where main, parsing nan as NaN, keeps 9.5. Coerce us on its own. op_tests/tuning_tests/test_config_shape_collision.py also pins the two behaviours the text merge exists for: integer cells of a column with a blank stay integers in the merged table, and 256 and 256.0 stay one shape. Signed-off-by: zhimding <zhimding@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
…-moe-dryrun-20261009
…-moe-dryrun-20261009
Conflicts: - aiter/aot/flydsl/gemm.py: #5969 taught the hand-written AOT the gfx950 FP32-scale blockscale kernels (flydsl_blockscale_8w_*), which it compiled for every flydsl row of a8w8_blockscale_tuned_gemm and a8w8_blockscale_bpreshuffle_tuned_gemm, with bf16 and fp16 output. The replay keeps no per-family code; a later commit replays those two tables through gemm_a8w8_blockscale and gemm_a8w8_blockscale_bpreshuffle instead. Co-authored-by: Cursor <cursoragent@cursor.com>
gemm_a8w8_blockscale_ck, gemm_a8w8_blockscale_cktile and their bpreshuffle counterparts shared one fake that takes only the five tensors. The tuned dispatch passes each of them kernelName, and the non-preshuffled ones splitK, so under FakeTensorMode, as in torch.compile or the AOT replay's dry run, all four raised TypeError on such a call. Each op now has a fake of its own signature, as the file's other ops do. A later commit replays the FP32-scale blockscale tables, which reaches gemm_a8w8_blockscale_bpreshuffle_cktile at 17 untuned sizes inside FlyDSL rows' ranges, where the runtime's padded M finds a CK-tile row. Co-authored-by: Cursor <cursoragent@cursor.com>
Since #7219 the FlyDSL MoE route, quant and grouped GEMM launches pass the raw stream handle (tensor_shim.current_stream_arg), taken from the torch._C query tensor_shim resolves at import rather than through torch.cuda. With the GPUs hidden that query raises AcceleratorError (no ROCm-capable device), so every gfx1250 MoE replay job failed, on main too: a sample of 298 jobs in 47 batches failed all 298. The dry run now answers the query with stream 0, as it does the ctypes path's getter, and the same sample replays its 298 rows and 404 sizes between them. A gfx950 sample (246 jobs) does not reach the query and replays as before. Co-authored-by: Cursor <cursoragent@cursor.com>
Since #5969, gemm_a8w8_blockscale and gemm_a8w8_blockscale_bpreshuffle serve gfx950 FlyDSL rows with FP32 128x128 scales (flydsl_blockscale_8w_*): 6,513 rows of a8w8_blockscale_tuned_gemm in 63 shapes and 670 of a8w8_blockscale_bpreshuffle_tuned_gemm in 76. Main's hand-written AOT compiled their kernels; the merge dropped that code with the rest. Both tables are now replayed through their public entries, each row with bf16 and fp16 output: the rows leave it open and the kernels specialize on it, as main's AOT compiled both. That is 14,366 of the default build's 22,878 gfx950 GEMM jobs. For GLM-5.2's two tables (27 and 6 FlyDSL rows) the replay builds the same 8 artifacts main's AOT builds, under the same launcher directories and cache keys, and a rebuild over them writes none. On gfx950 the real calls at each row's M, the smallest M it serves and one between, with bf16 and fp16 output (182 FlyDSL calls), run from that cache under FLYDSL_RUNTIME_RUN_ONLY=1 and match a torch reference; against an empty cache all 182 miss. An identify-only walk of every default GEMM job replays all 22,878 gfx950 and 1,500 gfx1250 rows with no failure. Co-authored-by: Cursor <cursoragent@cursor.com>
…back paths" This reverts commit ce0b847, so the replay no longer changes the gfx1250 asm gate: aiter/jit/utils/asm_guard.py stays as on main, and the dry run's Triton stand-in raises a plain RuntimeError again. replay_job tells a call another backend takes by the traceback frames in Triton's code and asm_guard.py, as before, still only when the call requested no FlyDSL kernel, did not pass through FlyDSL code, and raised neither GpuRuntimeReached nor ReplayHookBroken. The raw stream query test db4ae7e added next to the reverted tests stays. Co-authored-by: Cursor <cursoragent@cursor.com>
op_tests/tuning_tests/test_fast_tuning.py is back as on main, without the class's three cases (the per-arch worker waves, a GEMM table found by its AITER_CONFIG_* variable under any file name, and the FP8 FMHA launchers pinned through the wrapper's overrides) and the shutil and tempfile imports only they used. Co-authored-by: Cursor <cursoragent@cursor.com>
Conflicts: - aiter/aot/flydsl/gemm.py: #7309 taught the hand-written AOT the gfx950 MX preshuffle prefetch and b_nt kernels (_pf, _nt) and both l2_reduce grids of a fused split-K row, and #6767 let its gfx1250 A16W16 path take gfx1250-strict rows. The replay keeps no per-family code: run_gemm_a8w8_mxfp8_gfx950, which serves every table #7309 retuned, reads prefetch and b_nt off the tuned kernelName itself, and flydsl_mxscale_preshuffle_gemm picks l2_reduce from M, so it switches at most once in a row's range, where the replay bisects. No tuned table has a gfx1250-strict row. Co-authored-by: Cursor <cursoragent@cursor.com>
coderfeli
approved these changes
Nov 3, 2026
zhiding512
force-pushed
the
zhiding512/flydsl-aot-replay-ops-20261026
branch
from
November 3, 2026 09:02
d7f84a2 to
2ce2911
Compare
2 of 6 tasks
zufayu
added a commit
that referenced
this pull request
Nov 4, 2026
main's #6956 moved the GEMM/conv/chunk_gdn_h/FP8-FMHA AOT drivers into a shared aiter/aot/flydsl/replay.py. Re-apply the gfx1250-strict family-match + retarget there (one place) instead of the now-removed per-entry filters; take main's version of the four refactored entries. moe.py and mxfp4_moe.py keep their own drivers and retain the strict fix; common.py keeps resolve_target_arch / assert_targets_covered. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
fsx950223
added a commit
that referenced
this pull request
Nov 4, 2026
Resolve conflicts with #6956 (AOT caches built by replaying the runtime) and #7079 (gfx1250 PA decode partition merges). - aiter/aot/flydsl/pa_decode_plan.py now follows the replay protocol instead of the removed OpKind/compile_one_config dispatch: it replays plan_pa_decode on FakeTensors at the largest batch of each bucket and at each capacity where the block size can change, dense and windowed. Each arch replays with the most CUs of its devices, so the set equals the previous plan_variants minus the variants no device can reach (e.g. bucket 1 at 1024 threads). common.py registers it in AOT_MODULES; arch selection moves from --arch to GPU_ARCHS, as for every module. - README: main's module table and GPU_ARCHS wording, plus the planner entry and section. - op_tests: #7079 dropped test_pa_decode_plan; keep this PR's version alongside main's new tests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
4 of 10 tasks
IzacharyI
added a commit
to IzacharyI/aiter
that referenced
this pull request
Nov 4, 2026
The AOT GEMM module now replays each tuned row through the public gemm_a16w16 entry (ROCm#6956), so it no longer dispatches on a per-kernel kind. Drop the skinny registration it needed: tuned_gemm already routes a psk row into flydsl_hgemm, which the replay reaches on its own.
IzacharyI
added a commit
to IzacharyI/aiter
that referenced
this pull request
Nov 4, 2026
The AOT GEMM module now replays each tuned row through the public gemm_a16w16 entry (ROCm#6956), so it no longer dispatches on a per-kernel kind. Drop the skinny registration it needed: tuned_gemm already routes a psk row into flydsl_hgemm, which the replay reaches on its own.
huizzhan
added a commit
that referenced
this pull request
Nov 4, 2026
main gained #6956, which rewrote aiter/aot/flydsl/conv.py to build the AOT cache by replaying the runtime. That deleted the hand-enumerated job and per-kernel compile machinery this branch had extended for gfx1250, so all six hunks conflicted; take main's file outright. The gfx1250 compile keys now come from replaying flydsl_conv_implicit, which already owns the channel pad (CGP_ALIGN), the split-K resolver and the dyn_hw gate, so this branch's _conv_dedupe_key, _probe/_conv_probe_args and _compile_conv3d_gfx1250_to_cache paths have no remaining caller. Checked on gfx1250 (one card): replay builds 10 artifacts for a 5-row tuned table, --check passes, and every entry point of those rows (2 input x 2 output layouts, bias omitted where the table has none) hits under FLYDSL_RUNTIME_RUN_ONLY=1. Numerical accuracy was not re-validated here: this container's toolchain reproduces the same op-test failures as the pre-merge head, so the merge is behaviour-neutral but the head's accuracy is still open.
yuychang
added a commit
to yuychang/aiter
that referenced
this pull request
Nov 4, 2026
- Resolve aiter/aot/flydsl/gemm.py to main: ROCm#6956 replays the runtime's gemm_a16w16, so AOT already uses int64 semaphores - Drop alloc_split_k_buffers: no AOT caller is left, so gemm_a16w16_gfx950.py matches main again - Drop the cache-key test of the removed _compile_hgemm_to_cache; key equality is now by construction - Rework the RUN_ONLY test: AOT-build one split-K A16W16 row through the replay driver, then run gemm_a16w16 from it
3 of 10 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Applies the design of #6362, which builds the MoE and MegaMoE AOT caches by replaying the runtime on CPU, to every other FlyDSL AOT module.
Motivation
gemm.py,conv.py,chunk_gdn_h.pyandfmha_fp8.pyre-derived, by hand, the kernel names, compile-time parameters and launch arguments the runtime would use. Any drift between that copy and the runtime makes the AOT artifact unreachable, and nothing reports it. One such drift is live on main today. Since #6494 (2026-10-15) the gfx950 bf16 hgemm runtime allocates its split-K semaphore asint64, while the AOT still builds it asint32. The tensor dtype is part of the FlyDSL cache key, so every gfx950flydsl_hgemmrow (2,543 tuned rows, 780 kernels) misses the AOT cache and JIT-compiles on the serving path, while the AOT log reports those kernels as compiled.What changes
As in #6362, each tuned row (or configured head shape) becomes the arguments a caller passes the op's public entry, as FakeTensors. That entry runs inside
dry_run, where the runtime does its own lookup, kernel selection, sizing and launch, and every FlyDSL launch compiles into the cache instead of running.replay.py(new): the shared driver.replay_env,GPU_ARCHShonoured), with jobs batched per worker. The per-arch processes get the command line as given.SCAN_ARCHS) every M of that range is identified without compiling instead.requested_kernels).--check: compile nothing and fail on any kernel the cache lacks.[SKIP]forNotImplementedError;[REFUSED]at in-between sizes, only for a limit the module declares inREFUSALS;[UNREACHED]for rows the runtime hands to Triton or asm, told by the error's traceback passing through Triton's code or the asm gate. Any other error fails the build.gemm.py: about 1,340 lines of per-family compile code become one caller per tuned table (TABLES):gemm_a8w8_bpreshuffle;gemm_a8w8_blockscaleandgemm_a8w8_blockscale_bpreshufflewith FP32 scales ([CK] [FlyDSL] Introduce FP8 a8w8 gemm kernel with fp32 blockscale to uplift the MI355 blockwise gemm perf #5969'sflydsl_blockscale_8w_*, each row with bf16 and fp16 output);gemm_a8w8_blockscale_abpreshuffle;gemm_a8w8_mxfp8_bpreshuffle;gemm_a8w8_blockscale_bpreshufflewith e8m0 scales;batched_gemm_a8w8_mxscale_bpreshuffle;gemm_a16w16; andgemm_a4w4.--csvpoints the runtime's lookups at the given tables.conv.py: replaysflydsl_conv_implicitwith both output layouts; an NCDHW input also reaches the pre-transpose.chunk_gdn_h.py: replayschunk_gated_delta_rule_fwd_h_flydsl_optat each row's batch shape, and over batch shapes the table never measured, where the BV rule decides.fmha_fp8.py: replaysflydsl_flash_attn_fp8_funcover a batch × seqlen grid, bisecting between grid points, and each launcher the wrapper can pick, pinned through its own overrides.common.run_aot: now runs every module's replay. The pooled per-kernel path,OpKindandJobLabelare removed.torch_guard.run_python_ops_for_real/dry_run(python_ops=True): atorch_compile_guardop with a Python body runs that body under FakeTensorMode, while C++ ops still dispatch to their fakes. Outside a dry run this costs one global check in the existing wrapper.jit/core.pyctypes path reads its stream/device getters per call from_pybind_develop_hooks. It used to capturetorch.cuda.current_deviceat decoration time, soget_padded_mreached the GPU in a dry run.dry_runstands in fortorch.cuda.device, MemPools,record_stream, Triton's driver queries and the raw stream handle FlyDSL launches pass since [Perf] Avoid torch Stream construction across grouped MoE launches #7219. Reaching the GPU runtime raisesGpuRuntimeReached, which always fails the build.kernelName, raisedTypeError.moe.py: sharesrequested_kernelswith the driver; its behaviour is unchanged.Non-public interfaces the replay depends on
The replay runs the runtime's own host code without a GPU, so it stands in for some interfaces that are not public. Each is in
dry_run.pyorreplay.py, and a change to any of them fails the build instead of passing silently:JitFunction._compile_or_execute,_build_full_cache_keystop, ends it thereReplayHookBrokenCompiledArtifact._ensure_enginetorch.cuda._lazy_init; thetorch.cudastream, device and property queries and MemPools; the Tensor methods whose bindings take a device guard (__getitem__,contiguous,data_ptr,record_stream)_lazy_initraisesGpuRuntimeReached; the rest are stood in forRuntimeError[UNREACHED]; anything else fails_pybind_develop_hooks_cache,tensor_shim._get_current_raw_stream,graph_alloc._persistent_pool,chip_info._detect_nativeThe stream stand-in is registered through FlyDSL's public
JitArgumentRegistry, and the per-arch processes get the command line as given rather than rebuilt through argparse internals.tests/test_flydsl_aot_replay.pycovers the FlyDSL hooks, the[REFUSED]classification and the raw stream stand-in on CPU, and CI's wheel build replays the default tables for gfx942 and gfx950.Validation
This revision (
d7f84a2, merged with main at6d8c8e9c2) isc9fba78below with ce0b847 reverted, so the PR no longer changesaiter/jit/utils/asm_guard.pyand a hand-off to Triton or asm is told by its traceback again, and withTestFlydslAotReplaydropped fromop_tests/tuning_tests/test_fast_tuning.py. ce0b847 had found both ways of telling a hand-off give the same outcomes on the tables where the build meets Triton and asm. The merge drops #7309's and #6767's edits to the hand-written GEMM AOT; the replay reaches what they added through the runtime (see the merge commit). On this revisiontests/test_flydsl_aot_replay.pypasses (5), and black and ruff at CI's versions are clean on every changed file; thec9fba78results below were not rerun.c9fba78(merged with main at3febddf6a, FlyDSL 0.4.1, gfx950 box with 2 CPUs):[REFUSED]). conv 152 rows, chunk_gdn_h 1,344 rows and 876 in-between sizes, fmha_fp8 3,112 rows (760[SKIP]: tensors of 2^31 elements or more), 0 failed. A sample of the MoE replay (298 gfx1250 and 246 gfx950 jobs) has 0 failures.FLYDSL_RUNTIME_RUN_ONLY=1and match a torch reference; against an empty cache all 182 miss.pytest tests/: 210 passed, 7 of them intests/test_flydsl_aot_replay.py. The AOT tests ofop_tests/test_gemm_group32_interface.py(3),op_tests/tuning_tests/test_fast_tuning.py::TestFlydslAotReplay(3) andop_tests/test_fhmoe.py -k aot(8) pass. black and ruff at CI's versions are clean.Earlier revisions (
c759a80and before, FlyDSL 0.4.0, both caches built from empty):--checkof the runtime replay against the cache main's AOT builds:Every kernel the replay built compiled, and
--checkof every module against that cache found 0 misses. On real GPUs withFLYDSL_RUNTIME_RUN_ONLY=1against the replay-built cache: gfx950 conv 152/152 calls, GEMM sample rows 80/80, fmha 326/326; gfx942 chunk_gdn_h and gfx1250 GEMM and MoE in the comments below.Notes for reviewers
gemm_a8w8_blockscale_abpreshufflehas no FlyDSL kernel that fits ("no FlyDSL config"). It is the one limitgemm.pydeclares inREFUSALS, and is reported as[REFUSED]; a runtime gap, not something this PR changes.PR Self-Review
Reviewed revision
d7f84a2
Behavior and intent
The FlyDSL AOT build of the GEMM, conv3d, chunk_gdn_h and FP8 FMHA caches no longer restates the runtime's kernel selection: it calls each op's public entry on FakeTensors in a GPU-less dry run and compiles exactly the launches the runtime makes. The hand-written AOT drifts from the runtime without any report; today every gfx950 bf16 hgemm row misses the cache main builds (int32 vs int64 split-K semaphore), and the 20 gfx1250 batched mxscale rows have no AOT at all. Runtime behaviour is unchanged except that the four blockscale CK/CK-tile fakes now accept their ops' arguments under FakeTensorMode.
Validation and unverified scope
Commands and results are in Validation above. On this revision:
tests/test_flydsl_aot_replay.py(5 passed), and black and ruff at CI's versions on every changed file. Atc9fba78, before the revert of ce0b847 and the merge with main: an identify-only replay of every default job of the four modules (0 failures);python -m aiter.aot.flydsl.gemm --csvwith GLM-5.2's FP32 blockscale tables, from an empty cache and again over it (the 8 artifacts main's AOT builds; the rebuild writes none); aFLYDSL_RUNTIME_RUN_ONLY=1run of 182 real FlyDSL calls on gfx950 (all hit; all miss on an empty cache);pytest tests/(210 passed) and the AOT op tests. Unverified on this revision: thosec9fba78results, rerun after the revert and the merge, in particular the replay of the four tables #7309 retuned, which CI's wheel build replays for gfx950; a compiling build of every default table (the box has 2 CPUs; the identify-only walk shows which kernels are requested, not that each compiles, which CI's wheel build does for gfx942 and gfx950), real-GPU runs on gfx942 and gfx1250 (done at c759a80, see the comments), the gfx1250 MoE replay beyond a 298-job sample, and MegaMoE, which needs 8-GPU EP with mori.PR-specific question
No dependency changes. The cache holds what the runtime requests, keyed as it looks them up: at c759a80 main's build had 1,239 artifacts no runtime path requests (780 of them hgemm under the stale key) and lacked the kernels of 2,543 gfx950 and 21 gfx1250 rows. The replay relies on FlyDSL's private
JitFunction._compile_or_executeand_build_full_cache_keyat the pinned 0.4.1 (a launch that stops deriving its key there fails withReplayHookBroken) and forbidsCompiledArtifact._ensure_engine. Clean build: every module built from an empty cache, and--checkpassed on that cache and failed on an empty one (c759a80); atc9fba78GLM-5.2's FP32 blockscale tables built from empty into the 8 artifacts main builds. Existing cache: a rebuild over that complete cache writes no artifact, andFLYDSL_RUNTIME_RUN_ONLY=1calls run from it on gfx950.Changed op tests
op_tests/test_gemm_group32_interface.py: its three AOT tests asserted on the hand-written AOT this PR removes (parse_csv's per-kind jobs,_compile_bmm_mfma_to_cache), and three of them failed under pytest. They now replay the rows underdry_runand compare the kernels the replay identifies with those of a call shaped as a caller makes it, which no other test covers. They also clear the_mxscale_bpreshuffle_tuned_nkcache that leaked between tests.