[FLYDSL] Support prefill GDN K5 fp32 chunk states - #4732
Open
huizzhan wants to merge 24 commits into
Open
Conversation
Align the FlyDSL mfma16/HIP K5 fork with the HIP kernel's snapshot-dtype policy: `snapshot_dtype` defaults to `k.dtype` and is independent of `state_dtype`, so fp32 per-chunk snapshots become an opt-in that no longer has to be rejected on the FlyDSL dispatch path. fp32 snapshots store the f32 accumulator cells directly, the same trade-off the HIP `SNAPSHOT_BF16=false` specialization makes. Routing them through the existing [V][K] LDS transpose buffer is not an option: an fp32 buffer adds 8-16 KB and the bf16 layout already sits at the 64 KB limit at BV=64. The bf16 specialization keeps its LDS transpose and coalesced b128 store unchanged, and the fp32 one reclaims that buffer. Co-authored-by: Cursor <cursoragent@cursor.com>
`PrefillGroup` / `PrefillArgs` gain a `snapshot_dtype` field (None -> k.dtype, i.e. the bf16 specialization), passed to the FlyDSL, HIP and Triton calls in the correctness sweep and the perf comparison. A group now selects fp32 snapshots with one field, the pytest id gains a `_snapFP32` suffix, and the perf table gains a `snap` column so both specializations are distinguishable in one run. The varlen-64k qwen-ptpc rows are duplicated into bf16/fp32 snapshot pairs (and renamed to name the model they came from) to track the fp32 store path's cost against the bf16 one. Co-authored-by: Cursor <cursoragent@cursor.com>
The K5 BV selector needs total_chunks and max_seq_chunks and took them from chunk_offsets.tolist() on every forward. That blocking D2H sits between the K1..K4 launches and K5's, so the host waited for the whole front end to retire and the GPU then idled through the remaining host work for K5 and K6 rather than the host running ahead. The copy itself is 4us; the bubble it opened was 66us of a 738us block. Prefer the host-side counts prebuilt metadata already carries (GatedDeltaRuleChunkSchedule.total_chunks / .max_seq_chunks, derived from seq_lens_cpu). The old call sat outside the prefill_metadata branch, so passing metadata did not avoid the transfer -- the case #4532 missed here. Without metadata, cache the counts on chunk_offsets itself, the same tensor-attribute pattern as _as_int32 and _resolve_prologue: it comes from the tensor_cache'd prologue, so its identity is stable across forwards and the counts cannot change underneath the cache. varlen-64k-qwen-ptpc-ali T=8192 (Hg=2, H=8, packed varlen with state I/O, one 80-CU gfx942): the block goes 738.1us -> 667.5us, its launch gap 66.1 -> 0.7us, and Memcpy DtoH disappears from the profile. The 79 non-perf cases in test_flydsl_linear_attention_prefill.py are unchanged. Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit 24b6f80)
…them The mfma16_hip launcher built its int32 placeholders by casting the float32 one -- `dummy.to(torch.int32)` for the state_indices slot on every forward that does not use a state pool, and again for cu/chunk_offsets on the dense path. The kernel reads nothing but the base pointer of a disabled slot, so each cast launched an elementwise kernel to copy the single element it never looks at. torch.empty gives the same valid non-null address with no launch. Profiled at ~5.2us per forward on gfx942, all of it launch overhead rather than the 4 bytes: 1% of a T=8192 varlen GQA block and constant in T. The block bench loses a kernel from every FlyDSL K5 path -- 670.0 -> 665.9us with the Triton prepare pair, 510.8 -> 508.2us with the fused prepare. Tests: op_tests/flydsl_tests/test_flydsl_linear_attention_prefill.py 79 passed (the state-pool paths pass a real index array and are untouched), op_tests/test_gated_delta_rule.py 142 passed. Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit 723bd83)
The mfma16_hip fork already narrows cu_seqlens and chunk_offsets through _as_int32, which caches the result on the source tensor, and takes its null slots from an allocation. The baseline wrapper still cast on every call: two copy kernels per forward on a varlen batch and two more for the dense placeholders, for slots whose base pointer is all the kernel reads. Measured on the mfma16_hip fork, the same placeholder cast cost ~5.2us per forward, all launch overhead. Not measured here: this wrapper's kernel uses mfma.f32.16x16x32.bf16, so it cannot be built on gfx942 (LLVM cannot select the intrinsic) and it has no in-tree caller, which leaves it outside the test suite. The change is a static equivalence -- _as_int32 returns an already-int32 tensor untouched and otherwise caches the narrowing, and both spellings of the placeholder are a one-element int32 tensor -- but it wants a run on gfx950 to confirm. The aiter HIP K5 path is unaffected, as it does not enter this module: use_chunk_hip=True profiles identically before and after across dense and varlen, same kernels and bit-identical outputs. Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit 79adfdd)
… buffer fp32 snapshots stored straight from the f32 accumulators, mirroring the HIP kernel. That layout is the problem: an MFMA C fragment owns 4 contiguous K per lane and 16 K (64 B) of a V row per wave, and the neighbouring 64 B belongs to another wave, so one buffer_store_dwordx4 turns into 16 half-line requests where a coalesced store issues 8 full ones. Re-pointing the store at addresses that cover whole 128 B lines (wrong data placement, perf probe only) took the kernel 630 -> 508us, so 122 of the 146us fp32 sat above bf16 was the request count, not the 2x bytes; 8x128 B and 2x512 B per wave measured identical, so reaching a full line is all that matters. Ablating the snapshot store entirely lands at 400us, and moving it later in the chunk (638us) or turning SCHED_GFX942 off (630us, i.e. its 27us GEMM1 win erased) does not help. Route fp32 through the same [V][K] transpose buffer bf16 uses. The f32 tile is twice the bytes and the bf16 layout already sits exactly at the 64 KB LDS limit, so it goes out in two half-BV rounds. Both rounds stage while the accumulators still hold the chunk-start state, and each is placed where the loop already has the barrier it needs -- round 0 before the top-of-loop publish barrier, round 1 after GEMM1's WAR barrier (which also retires round 0's reads) -- so the split costs no extra barrier. Each flush follows a GEMM so the MFMA chain covers the ds_read latency, worth another 21us over flushing both rounds back-to-back at the top. LDS returns to the bf16 layout's 64 KB, occupancy unchanged at 1 workgroup/CU. varlen-qwen-ali-tp1 T=8192 (Hg=16, H=32, one gfx942): fp32 631 -> 512.9us against 483.2us for bf16, +30% -> +6%; HIP fp32 is 609.6us, so fly/hip goes 0.97x -> 1.19x. Over both snapshot perf groups (8 batched-token lengths each) the fp32 overhead is 3-7% where HIP pays 19-22%, and fly/hip is 1.18-1.31x. bf16 is byte-identical and measures unchanged (483.2 vs 483.8us). SCHED_GFX942 matters again for fp32 too: 514.5 vs 546.5us with it off. The 110 non-perf cases in test_flydsl_linear_attention_prefill.py pass, as do the 71 snapshot/fp32 cases with FLYDSL_K5_MFMA16HIP_BV forced to 32 (two rounds, one cell each) and 16 (single round). Co-authored-by: Cursor <cursoragent@cursor.com>
The four snapshot groups sweep max_num_batched_tokens from 8192 to 65536, so "varlen-64k-" named only the last row. Use the model name alone. Co-authored-by: Cursor <cursoragent@cursor.com>
The harness let every backend rediscover the chunk schedule. FlyDSL caches that host-side result on the cu_seqlens tensor, so it pays it once for the whole sweep; the HIP wrapper rebuilds it per forward and reads chunk_offsets back with a blocking D2H, and the triton path builds its own. The reported number is K5 device self time, so this mostly leaks in as launch-stream stalls and a few schedule-building kernels rather than as K5 time, but it makes the backends differ in host caching on top of the kernel work the table is meant to compare -- and it is not what a serving stack does, which builds one GatedDeltaRulePrefillMetadata per forward and reuses it across the layer stack. Build the schedule once per shape and pass it to all three backends (the file's HIP adapter grows the pass-through argument). Dense shapes get None: without cu_seqlens the wrappers take the layout from the tensor shapes and there is no schedule to prebuild. Every backend now also runs metadata.validate() on the same object, so the sweep covers that they agree on one schedule. varlen-qwen-ali-tp1 T=8192 fp32 is unchanged within noise (FlyDSL 510.2us, HIP 607.0us); at block level the same change takes HIP's host gap from 169 to 1.7us. Co-authored-by: Cursor <cursoragent@cursor.com>
…dtype bench_gated_delta_rule_snapshot_dtype.py times K5 and K6 in isolation, and the K5 test harness reports K5 device self time. Neither answers what the fp32 snapshot costs the layer a serving stack actually runs: K6 reads the snapshot back, so its 2x bytes land there too, and a wrapper that stalls the launch stream shows up in neither. Run chunk_gated_delta_rule_opt_vk and split the result into wall time (median of CUDA-event timings), profiler self time bucketed into K1..K6, and the wall minus device gap that isolates launch overhead. Shapes come from the K5 suite's PrefillArgs cases, selected by their pytest ids (--list-cases), so a block row and a K5 row describe the same workload instead of a second copy of the shape table that can drift. The block entry always cumsums a gate, so the suite's no-g cases (which cover K5's padding masking) are skipped. The block is timed with a prebuilt metadata by default, the way a serving stack reuses one schedule across the layer stack; --without-metadata adds the None column that exposes the chunk-schedule D2H as host gap. varlen-qwen-ali-tp1 T=8192 fp32 snapshots, one gfx942: stage (us) flydsl flydsl no-meta hip hip no-meta K5 state scan 506.4 506.9 609.4 608.0 K6 output 842.5 843.2 842.0 839.4 wall total 2235.7 2233.3 2335.1 2510.8 host gap 2.8 0.9 1.7 169.4 so with the schedule prebuilt the block gap is the ~100us of K5, and HIP's remaining 169us is its per-forward chunk_offsets readback, which the FlyDSL wrapper already caches. Co-authored-by: Cursor <cursoragent@cursor.com>
The example -k filter named varlen-64k-qwen-ptpc-ali, which the group rename removed, so copying the line selected nothing. Use one of the current snapshot groups. Co-authored-by: Cursor <cursoragent@cursor.com>
… apart Sweeping several PrefillGroups printed one stage table per case, which does not read across cases, and the host gap alone was ambiguous: the TP8 T=8192 rows showed 100-400us of it, which looks like the wrapper stall the metadata=None runs expose but is not. Add a cross-case summary once more than one case runs: cases down (one row per snapshot dtype, so mixing bf16 and fp32 groups keeps the columns dense), backend x metadata variant across, one table per metric, plus a speedup column against HIP under the same variant, >1 meaning the other backend is faster. --summary-metrics picks the metrics, --summary-only drops the per-case tables. Add `launch only`, the median python-side cost of one call measured without syncing, which is what reads the host gap. It is a floor on wall time: at TP8 T=8192 the block needs ~750us of host per call against 678us of device work, so wall lands at ~800us for both backends and the gap is that shortfall, not a stall. With `prefill_metadata=None` the HIP wrapper's blocking chunk_offsets readback instead takes `launch only` from 780 to 2504us, i.e. up to wall, which is what a stall looks like. Co-authored-by: Cursor <cursoragent@cursor.com>
The mfma16_hip fork sizes its value tile from a runtime heuristic alone, so a shape the heuristic reads wrong cannot be corrected without patching the rule. Let the wrapper consult a tuned table first and fall back to the rule for everything the table does not cover. The key is built from host-side quantities only, so the probe stays one dict lookup (~0.09us against ~46us of wrapper host time) and never forces a device-to-host copy; the table is read once at import with stdlib csv rather than lazily, keeping the file read out of a live server's first prefill. The seeded rows come from the varlen-qwen3.5-397b-ptpc-ali and varlen-qwen-ali-tp1 sweeps, measured by the accompanying tune script. On those the rule already picks the optimum for all 32 combinations, so the table changes no selection today -- it pins that result against future regressions and records the K5 time each tile achieved. Co-authored-by: Cursor <cursoragent@cursor.com>
aiter.aot.flydsl.chunk_gdn_h covers the baseline chunk_gdn_fwd_h_flydsl_vk, but use_chunk_flydsl dispatches the mfma16_hip fork -- a separate compiled product with no AOT coverage at all. Every process therefore JITs it on the first prefill (2.3s cold), and JITs again mid-serving whenever a batch shifts the selected BV or a caller switches snapshot/state dtype. Add an AOT kind driven by the same tuned table the runtime reads. The table's BV column deliberately does not narrow the fan-out: shapes it does not cover fall back to the CU/LDS rule and can land on any tile, so every legal BV is built, as are both dtype specializations. That is 24 jobs for today's two shapes, 2.4s of build time, and it takes the first call down to ~0.1-0.3s with FLYDSL_RUNTIME_RUN_ONLY=1 passing on off-table shapes too. Co-authored-by: Cursor <cursoragent@cursor.com>
Extends the K5 table with the third case family, which unlike the two already in it batches ragged segments (33/7/4 sequences under a fixed 32k token budget), exercising the max_seq_chunks half of the lookup key. BV=64 wins all three, again matching what the rule already picks. Co-authored-by: Cursor <cursoragent@cursor.com>
The two Qwen3.5-* prefill groups run dense with output_final_state=False, which the tuned table could not express: use_h0/store_fs were neither part of the lookup key nor of the AOT fan-out, so a dense no-final-state row would have collided with an existing varlen row and its kernel would still have hit the JIT. Add both switches as columns, key on them, and fan the AOT jobs over the values each row asks for (24 -> 84 jobs). Qwen3.5-35B at TP2/T2500 is the first shape where the CU/LDS rule is wrong: BV=64 measures 144.3us against 147-150us for the rule's BV=32, repeatable across three 200-iteration runs. The other 19 shapes confirm the rule. Rows can now describe the same physical batch shape under two model names (varlen-32k-qwen at TP8/T8192 is the same 4x8192 batch as the 397b-ptpc mnbt32768 row), so warn instead of silently overwriting when two rows disagree on BV. Co-authored-by: Cursor <cursoragent@cursor.com>
…er it in AOT chunk.py rejected initial_state_indices / inplace_final_state whenever use_chunk_flydsl was set, even though the K5 wrapper implements both: it gathers each sequence's slot on read and writes the final state back into the pool in place. Drop the guard and forward both arguments, matching what the HIP branch already did. That makes USE_STATE_INDICES reachable, so AOT fans out over it on the rows that permit a pool (use_h0 and store_fs both set; the wrapper rejects the rest). Also fan out G_HEAD_MAJOR: chunk.py pins it True, but False is the wrapper's own default, so a caller reaching the wrapper directly was the one combination AOT skipped. 84 -> 240 jobs, 8.1s. Remaining pinned switches are unreachable through the production dispatch: g is always passed, gk never is, v_new is always saved, use_exp2 pre-scaling is on, wu_contig is hardcoded at the call site, and bf16_convert_trunc only flips in a unit test that checks RNE rounding. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
- EXE001: the new AOT module carries a shebang like every other module in aiter/aot/flydsl, so give it the executable bit those have (chunk_gdn_h.py, common.py, moe.py, fhmoe.py). gemm.py and grouped_moe.py miss it too, but they are outside this PR's diff. - I001: let ruff sort the tune script's imports, which splits the aliased k5 import into its own statement. - C408: dict() -> a literal in the tune script's kernel-argument builder. Co-authored-by: Cursor <cursoragent@cursor.com>
huizzhan
changed the title
huizzhan
changed the title
The sweep harness is a local tool, not something the repo needs to carry, so drop it from the PR. The csv and the AOT README described the data's origin by naming its path; say what the measurement is instead, and note that resolving the ~2% gaps takes a few hundred iterations, which is the part worth knowing when reproducing a row. Co-authored-by: Cursor <cursoragent@cursor.com>
Keep only what the code cannot show (the LDS budget forcing two snapshot rounds, the barriers round 1 rides on, why the host metadata is cached) and restore the wording that was already there for the rest. Co-authored-by: Cursor <cursoragent@cursor.com>
Both K5 compiled products now live in one module: chunk_gdn_h.py grows a
_KERNEL_SPECS table that drives parse/compile/print per kernel, plus a
--kernel {vk,mfma16_hip,all} switch defaulting to all, so one command
builds both. chunk_gdn_h_mfma16_hip stays behind as a re-export, keeping
OpKind.CHUNK_GDN_H_MFMA16_HIP and its documented python -m entry point
working with no change to common.py or the README.
The job sets are byte-identical to before the merge (vk 24, mfma16_hip
240), and all 240 mfma16_hip configs load their AOT artifact under
FLYDSL_RUNTIME_RUN_ONLY=1, so no batch shape can fall back to the JIT.
Also corrects a stale doc line carried over from the old module:
chunk_gdn_h_tuned.csv is an AOT seed list only -- runtime BV comes from
the _heuristic_bv rule, which never reads the file.
Co-authored-by: Cursor <cursoragent@cursor.com>
chunk_gdn_h.py already carries both K5 compiled products, so the mfma16_hip module was nothing but a re-export. Delete it and point common.py's OpKind dispatch straight at the chunk_gdn_h symbols; the enum member stays, since it still buckets jobs and build logs. The only entry point is now `python -m aiter.aot.flydsl.chunk_gdn_h [--kernel vk|mfma16_hip|all]`, which is why the README goes back to its mainline text and the tuned csv header points at the surviving module. Job counts and artifacts are unchanged: vk 24, mfma16_hip 240, and all 240 still load from cache under FLYDSL_RUNTIME_RUN_ONLY=1. Co-authored-by: Cursor <cursoragent@cursor.com>
The block benchmark only ever ran dense per-sequence states, so the indexed pool a serving stack actually uses went unmeasured. Passing --with-state-pool adds it as a variant next to the dense one, the way --without-metadata already works, so a single run reads the cost of the slot gather and the in-place write-back. The pool is oversized and the slots scattered, since an identity mapping would hide the gather. This also corrects chunk_gated_delta_rule_opt_vk's docstring, which still claimed the pool was unsupported on the FlyDSL path; all three K5 paths forward the indices, and the benchmark now exercises that. Co-authored-by: Cursor <cursoragent@cursor.com>
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
1 participant
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.
Summary
fp32 chunk snapshots (K5 mfma16_hip fork)
snapshot_dtypedecoupled fromstate_dtype, defaulting tok.dtypelike the HIP pathf32 tile stays within the LDS budget; the bf16 specialization is byte-identical
Host overhead on the K5 wrapper
chunk_offsetsevery forward (the.tolist()D2H was stalling thelaunch stream: 169-384us host gap -> 1.7us when prefill metadata is passed through)
AOT coverage for the kernel production actually dispatches
aiter/aot/flydsl/chunk_gdn_h_mfma16_hip.py: the previous module only covered thebaseline
chunk_gdn_fwd_h_flydsl_vk, so the mfma16_hip fork K5 dispatches had zerocache overlap and still paid a 2.3s cold JIT
csv edit; each shape fans out over every legal BV, both snapshot/state dtypes, both
g-layouts, and state-pool indices where the row allows one (240 jobs, 8.1s)
Performance
K5 device kernel time (us) via
torch.profiler, one gfx942,SeqLen=8192,sweeping
max_num_batched_tokens.fly/hip > 1means FlyDSL is faster.varlen-qwen-ali-tp1 (TP1, Hg=16, H=32)
varlen-qwen3.5-397b-ptpc-ali (TP8, Hg=2, H=8)
FlyDSL pays 3-7% for fp32 snapshots where HIP pays 19-22%. Triton is 2.0-3.6x
slower than HIP on both dtypes and is omitted from the tables. The two TP8 rows
where HIP fp32 lands at or below its own bf16 (T=8192, 16384) are noise on
those points; every other length shows the same ~20% HIP fp32 cost.
Block level
bench_gated_delta_rule_block.pyruns the same shapes throughchunk_gated_delta_rule_opt_vk, so the snapshot policy is also visible in K6(it reads the snapshot back) and in the wall-minus-device gap. One gfx942,
varlen-qwen-ali-tp1 T=8192, fp32 snapshots, wall = median of 50 CUDA-event
timings, device = 20-iter profiler self time:
With one prebuilt metadata the whole 100us block gap is K5, i.e. the store
coalescing; K6 is identical because both backends hand it the same fp32
snapshot (fp32 costs K6 ~23% over bf16 on this shape, for either backend).
The
no-metacolumns are the same run withprefill_metadata=None: the HIPwrapper then rebuilds the chunk schedule per forward and reads
chunk_offsetsback with a blocking D2H, worth 169us of host gap plus ~12us of schedule
kernels, while the FlyDSL wrapper caches that result on the
cu_seqlenstensor. That gap is wrapper-side and orthogonal to this PR's kernel work; the
same caching still wants to land on the HIP wrapper.
Host-side commits, measured on the same shape before the store work: caching
the chunk counts instead of reading
chunk_offsets.tolist()per forward closeda 66us launch gap (block 738 -> 667us,
Memcpy DtoHgone from the profile),and allocating the null slots instead of casting into them dropped two copy
kernels per launch.
Test plan
op_tests/flydsl_tests/test_flydsl_linear_attention_prefill.py(correctness + e2e dispatch, 111 passed)op_tests/test_gated_delta_rule.py -k "indexed or flydsl or state_pool"(19 passed)FLYDSL_RUNTIME_RUN_ONLY=1to prove none of them fall back to the JIT