[FlyDSL] [JIT] Build the GEMM, conv, chunk_gdn_h and FP8 FMHA AOT caches by replaying the runtime

aiter/aot/flydsl/README.md

@@ -5,17 +5,18 @@ FlyDSL kernels. Each module derives the runtime-reachable jobs from aiter's
 tuned CSV configs—including ABI/cache variants that share a kernel name—and
 compiles them up front so the runtime JIT path hits the cache.
 
-| Module | OpKind | Description |
-| --- | --- | --- |
-| `moe.py` | `MOE` | Every FlyDSL kernel the tuned fused-MoE / FHMoE / gfx1250 grouped rows reach, by replaying `fused_moe` itself (see below) |
-| `gemm.py` | `GEMM` | GEMM kernels |
-| `conv.py` | `CONV` | conv3d implicit-GEMM kernels, plus the NCDHW->NDHWC pre-transpose each NCDHW caller reaches first |
-| `chunk_gdn_h.py` | `CHUNK_GDN_H` | chunk-gdn-h opt (K5) kernels |
-| `mega_moe.py` | `MEGA_MOE` | MegaMoE A8W4 and A4W4 bundles of every deployment profile, by preloading `MegaMoEV2` itself (see below) |
-| `fmha_fp8.py` | `FMHA_FP8` | gfx950 FP8 flash-attention forward, per head shape in `DEFAULT_SHAPES` |
-| `mxfp4_moe.py` | — | Not part of the AOT build; the MXFP4 tuner's `--fast-scan` precompiles its candidates with it |
-| `common.py` | — | Shared job collection, the deadlock-free fork pool, and cache-hit checking logic |
-| `dry_run.py` | — | Runs aiter's real host code on CPU (FakeTensors, `COMPILE_ONLY=1`, GPU runtime forbidden) |
+| Module | Description |
+| --- | --- |
+| `moe.py` | Every FlyDSL kernel the tuned fused-MoE / FHMoE / gfx1250 grouped rows reach, by replaying `fused_moe` itself (see below) |
+| `gemm.py` | Every FlyDSL kernel the tuned GEMM tables reach, by replaying each table's public GEMM entry (see below) |
+| `conv.py` | conv3d implicit-GEMM kernels and the NCDHW->NDHWC pre-transpose, by replaying `flydsl_conv_implicit` |
+| `chunk_gdn_h.py` | chunk-gdn-h opt (K5) kernels, by replaying `chunk_gated_delta_rule_fwd_h_flydsl_opt` |
+| `mega_moe.py` | MegaMoE A8W4 and A4W4 bundles of every deployment profile, by preloading `MegaMoEV2` itself (see below) |
+| `fmha_fp8.py` | gfx950 FP8 flash-attention forward, by replaying `flydsl_flash_attn_fp8_func` per head shape in `DEFAULT_SHAPES` |
+| `mxfp4_moe.py` | Not part of the AOT build; the MXFP4 tuner's `--fast-scan` precompiles its candidates with it |
+| `common.py` | Job collection, the deadlock-free fork pool, `run_aot` (what `setup.py` runs), and cache-hit checking logic |
+| `replay.py` | The shared driver of the replaying modules other than MoE/MegaMoE: per-arch processes, worker batches, size scans and bisection, `--check` |
+| `dry_run.py` | Runs aiter's real host code on CPU (FakeTensors, `COMPILE_ONLY=1`, GPU runtime forbidden) |
 
 ### How `moe.py` works
 
@@ -66,9 +67,11 @@ need no AOT change.
   arch and `AITER_AOT_REPLAY=1`. Under that flag `get_gfx_runtime()` and
   `dtypes.fp8` answer the `GPU_ARCHS` arch instead of probing a device, so a
   visible GPU of another arch cannot leak in; the GPUs stay hidden all the
-  same. `python -m aiter.aot.flydsl.moe` starts those processes side by side
-  and splits the worker budget (`AITER_MAX_JOBS`, below) between them by job
-  count. Beforehand, it builds aiter's C++ core module and the one
+  same. `python -m aiter.aot.flydsl.moe` starts those processes side by side,
+  as many as the worker budget (`AITER_MAX_JOBS`, below) gives a worker each,
+  and splits the budget between them by job count; with fewer workers than
+  archs, the rest wait for a later round. Beforehand, it builds aiter's C++
+  core module and the one
   behind the two-stage host-side query once, for the arch with the most
   two-stage jobs and with the whole budget: the replays run only their host
   code, and not every toolchain builds every arch (ROCm 7.2 cannot build the
@@ -98,6 +101,68 @@ variant of the profile's bundle plan.
   `python -m aiter.aot.flydsl.mega_moe` and `setup.py` start it that way. It
   is skipped when `GPU_ARCHS` names targets without gfx950.
 
+### How `gemm.py`, `conv.py`, `chunk_gdn_h.py` and `fmha_fp8.py` work
+
+They follow `moe.py`: none restates launch logic. Each turns a tuned row (or,
+for `fmha_fp8.py`, a configured head shape) into what a caller passes the op's
+public entry, as FakeTensors, and `replay.py` calls that entry inside
+`dry_run`. The runtime looks the row up, picks and parameterizes its kernels
+and launches them; every FlyDSL launch is compiled into the cache instead.
+
+- `gemm.py` replays one public entry per tuned table (`TABLES`):
+  `gemm_a8w8_bpreshuffle`, `gemm_a8w8_blockscale` and
+  `gemm_a8w8_blockscale_bpreshuffle` with FP32 scales (each with bf16 and
+  fp16 output, which their rows leave open),
+  `gemm_a8w8_blockscale_abpreshuffle`, `gemm_a8w8_mxfp8_bpreshuffle`,
+  `gemm_a8w8_blockscale_bpreshuffle` with e8m0 scales,
+  `batched_gemm_a8w8_mxscale_bpreshuffle`, `gemm_a16w16` and `gemm_a4w4`.
+  A row tuned at M also serves the M down to the previous tuned M
+  of its shape. On gfx1250 (`SCAN_ARCHS`) every one of those M is identified
+  first: replayed with each FlyDSL launch stopped at its cache key, so nothing
+  compiles. Then one M per set of kernels the row's own replay did not build
+  is replayed for real, reported as `[SCAN]`. Bisecting between the two ends
+  would miss kernels there, since the gfx1250 GEMMs fit cluster_m to the
+  number of M tiles and specialize a ragged last tile: M=5120 and M=6144 can
+  request one kernel and M=5632 another. Rows of other archs are replayed at
+  both ends and bisected between them where they request different kernels
+  (`[BISECT]`). Either way this reaches the FlyDSL heuristic kernels the
+  runtime falls back to below a shape's smallest tuned M; the scan also
+  reaches them past a smallest M the runtime refuses. The a4w4 table's
+  non-FlyDSL rows are replayed too, since M below them falls back to FlyDSL.
+- `conv.py` replays each row with both output layouts (part of the compile
+  key) from an NCDHW input, which reaches the pre-transpose as well.
+- `chunk_gdn_h.py` replays each row at its batch shape, with `g` head- and
+  token-major and a dense and indexed initial state, and each of those layouts
+  over batch sizes the table never measured, where the BV rule decides.
+- `fmha_fp8.py` replays a grid of batch sizes and sequence lengths per head
+  shape, bisecting between grid points where the wrapper switches kernels.
+  Between grid points of the batch and query length the wrapper can still
+  pick a launcher no grid point does (causal cross attention at batch 1,
+  query length 5675 and key/value length 66157 splits 11 ways), so each
+  launcher it can pick at all is replayed too, with its tile and split count
+  pinned through the wrapper's own overrides.
+- Ops with a Python body behind `torch_compile_guard` run that body under
+  `dry_run(python_ops=True)`, so a replay can enter at a public op; C++ ops
+  still dispatch to their fakes.
+- A call the runtime refuses (`NotImplementedError`) is reported as `[SKIP]`.
+  A call the runtime hands to another backend (Triton, asm, CK) compiles no
+  FlyDSL kernel; a row that never reaches FlyDSL is reported as
+  `[UNREACHED]`. Reaching the GPU runtime from host code fails the build.
+- At a size inside a row's range the runtime may meet a limit of its own,
+  which the module names in `REFUSALS` (`gemm.py`: for some untuned M,
+  `gemm_a8w8_blockscale_abpreshuffle` has no FlyDSL kernel that fits). A size
+  that fails with one is reported as `[REFUSED]`; any other error there fails
+  the build, as it does at the row itself.
+- Each module runs one process per arch its jobs compile for, with
+  `GPU_ARCHS` pinned and `AITER_AOT_REPLAY=1`, as `moe.py` does. `--csv` points
+  the runtime's tuned-config lookups of those processes at the given tables.
+  A GEMM CSV belongs to the table whose `AITER_CONFIG_*` variable names it,
+  under any file name, else to the table its file name holds; any other file
+  is an error, not a table to skip.
+- `--check` compiles nothing and fails on any kernel the replay would need but
+  the cache lacks. Set `AITER_FLYDSL_AOT_TRACEBACK=1` for a failing replay's
+  traceback.
+
 ---
 
 ## 0. Set up the environment
@@ -116,9 +181,8 @@ directory of your checkout).
 ## 1. Run AOT pre-compilation (compile smoke test)
 
 The most direct "test" is to run each module as a `python -m` entry point and
-confirm every kernel compiles. Each module prints `Compiled: N ok, M failed` at
-the end and exits 0 when all succeed, 1 on any failure — so it plugs straight
-into CI.
+confirm every replay compiles. Each module prints a per-arch summary at the end
+and exits 0 when all succeed, 1 on any failure — so it plugs straight into CI.
 
 ```bash
 # MoE / Mixed-MoE (default CSVs)
@@ -145,13 +209,14 @@ python -m aiter.aot.flydsl.mega_moe --experts-per-rank 48
 
 The FP8 flash-attention kernel picks its tile, rescale threshold and split-K
 factor per call, so a serving process JIT-compiles each new combination on the
-request path. `fmha_fp8` enumerates every combination the wrapper's heuristics
-can reach for each head shape in its `DEFAULT_SHAPES` table. Add a model by
+request path. `fmha_fp8` replays the wrapper over a grid of batch sizes and
+sequence lengths, and once per launcher it can pick, for each head shape in
+its `DEFAULT_SHAPES` table. Add a model by
 adding an entry (`num_heads, num_kv_heads, head_dim, head_dim_v, layout`, heads
 per rank, layout `varlen_cross` for the chunked-prefill serving path).
 
 ```bash
-# every configured shape (Kimi-K3 at TP8: 12:12:192:128 -> 92 kernels)
+# every configured shape
 python -m aiter.aot.flydsl.fmha_fp8
 # ad hoc shape, without editing DEFAULT_SHAPES
 python -m aiter.aot.flydsl.fmha_fp8 --shape 12:12:192:128
@@ -229,20 +294,20 @@ correctness nor precompiled replay acceptance.
 | `AITER_AOT_IMPORT` | Set to `1` so `import aiter` only loads the lightweight JIT core and skips the full top-level op namespace — faster and avoids heavy import side effects during AOT compilation (this is what `setup.py` sets while pre-compiling). | `0` |
 | `FLYDSL_RUNTIME_CACHE_DIR` | Cache directory | `~/.flydsl/cache` |
 | `AITER_MAX_JOBS` | Optional AITER-local ceiling for concurrent CPU compilation workers. If it is unset, standalone FlyDSL AOT CLIs adopt a valid positive legacy `MAX_JOBS` with a `FutureWarning`; AITER imports and runtime JIT do not. Sizing and precedence follow the shared policy in the [top-level README](../../../README.md#build-parallelism); invalid values use automatic sizing and non-positive values impose a one-worker ceiling. | auto |
-| `AITER_FLYDSL_AOT_TIMEOUT` | Per-kernel wall-clock cap (seconds). A worker stuck *alive* past this is killed (and retried); `0` disables. | `1200` |
+| `AITER_FLYDSL_AOT_TIMEOUT` | Wall-clock cap (seconds) of one worker: one kernel, or for the replaying modules the batch of jobs it replays in turn. A worker stuck *alive* past this is killed (and retried); `0` disables. | `1200` |
 | `AITER_FLYDSL_AOT_MAX_RETRIES` | Retries for a worker that **died abnormally** (OOM-kill / segfault / timeout-kill). A clean compile error is never retried. `0` disables. | `2` |
 | `AITER_CONFIGS` | Resolves the default CSV lookup path (same as the runtime JIT) | repo built-in |
-| `ARCH` / `GPU_ARCHS` | Selects which jobs to build, not what arch a job compiles *for* (that comes from the CSV's `gfx`, with legacy `cu_num` fallback). `conv.py` applies it inside `parse_csv`, and `moe.py` and `mega_moe.py` (`GPU_ARCHS` only) before they start their replays, so both `python -m` and the `setup.py` path (`run_aot`) honour it. `gemm.py` still filters in `main()` only, so `run_aot` builds all of its archs. | auto-detect |
+| `ARCH` / `GPU_ARCHS` | Selects which jobs to build, not what arch a job compiles *for* (that comes from the CSV's `gfx`, with legacy `cu_num` fallback). Every module applies it (`GPU_ARCHS` only) before it starts its per-arch replays, so both `python -m` and the `setup.py` path (`run_aot`) honour it. | auto-detect |
 
 > **About the compile target arch.** The arch each kernel is actually compiled
 > for is taken from the CSV's `gfx` column, with `cu_num_to_arch(...)` retained
 > as a fallback for legacy rows, and applied internally via `FLYDSL_GPU_ARCH`.
 > That internal var is overwritten
 > for every job, so setting `ARCH` / `GPU_ARCHS` / `FLYDSL_GPU_ARCH` in your shell
 > does **not** change what gets built. To cross-compile, use the target `gfx` and
-> `cu_num` values in the CSV. `moe.py` and `mega_moe.py` additionally run each
-> target arch in its own process with `GPU_ARCHS` pinned and `AITER_AOT_REPLAY=1`,
-> so they build the same kernels on a machine with or without GPUs.
+> `cu_num` values in the CSV. Every module additionally runs each target arch in
+> its own process with `GPU_ARCHS` pinned and `AITER_AOT_REPLAY=1`, so it builds
+> the same kernels on a machine with or without GPUs.
 
 Example:
 
@@ -288,10 +353,12 @@ python op_tests/test_moe_2stage.py
 
 - **`CSV file not found`**: check the `--csv` path, or whether `AITER_CONFIGS`
   points at a valid config directory.
-- **Lots of `[FAIL]` prints + exit code 1**: an individual kernel failed to
-  compile; stdout has per-kernel diagnostics. The exception message inlines at
-  most 10 entries (`_MAX_ERRORS_IN_MSG` in `common.py`), the rest are elided as
-  `(... N more)`.
+- **Lots of `[FAIL]` prints + exit code 1**: an individual replay failed;
+  stdout names the job and the error. Rerun the module with
+  `AITER_FLYDSL_AOT_TRACEBACK=1` (and `--csv` narrowed to the failing table)
+  for the traceback. `FlyDSL AOT dry run reached the GPU runtime` means host
+  code on the replayed path queries the device in a way `dry_run.py` does not
+  stand in for yet.
 - **Every kernel fails with the same error** (e.g. `'ArithValue' object has no
   attribute 'ir_value'`): this is a **FlyDSL version mismatch**, not a per-kernel
   problem. Check the *imported* FlyDSL:
@@ -304,6 +371,13 @@ python op_tests/test_moe_2stage.py
   your local FlyDSL checkout (`scripts/build.sh`, after `pip install
   nanobind==2.12.0` if CMake reports it missing) so the on-`PYTHONPATH` build dir
   is refreshed to the right version.
+- **`ReplayHookBroken`, or an `AttributeError` on `JitFunction`, after a FlyDSL
+  upgrade**: the replay learns each launch's cache key by hooking two private
+  methods of FlyDSL's `JitFunction`, `_compile_or_execute` and
+  `_build_full_cache_key` (`replay.requested_kernels`). A release that renames
+  either fails every replay with an `AttributeError`; one whose launches stop
+  deriving their key in `_build_full_cache_key` fails them with
+  `ReplayHookBroken`. Update `requested_kernels` for that release.
 - **Worker OOM / killed (exitcode -9)**: abnormal exits are auto-retried
   (`AITER_FLYDSL_AOT_MAX_RETRIES`) and the default worker count is already
   bounded by the shared CPU and memory policy. If it still happens,