Repository navigation
[Perf] Avoid torch Stream construction across grouped MoE launches - #7219
Merged
Merged
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
github-actions
Bot
changed the title
NOT READY FOR HUMAN REVIEWPR Self-ReviewRevision Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision. Questions to answer in the PR description
Edit the Missing or invalid answers
Shadow mode reports what would block but keeps the status successful. Set |
Contributor
Author
Model ABBA: combined inference-anchor + grouped-GEMM stream changesBoth arms include the prior MoE dispatch change (#7210), existing Gluon layout cache (#7090), ATOM #2761 and the standalone-cache fix (ROCm/ATOM#2858). The candidate adds #7217 + #7219 together. This is a combination test, not attribution to either PR separately. M3 MXFP4, TP4/gfx1250, Simple, FP8 KV, EAGLE3; fixed synthetic acceptance 0.5933. Performance only, not generated-text accuracy. Same recorded prompt hashes across all four fresh processes. Source imports checked in all worker processes. Per process: 20 measured short requests, 20 measured long-prefix requests, six measured C2 rounds. cProfile diagnostics run after all timing samples and are excluded. No PyTorch profiler in timing.
Descriptive candidate-vs-baseline comparison using the mean of per-process medians:
This four-process experiment does not establish a robust end-to-end speedup. The raw order/spread and any regressions above are retained. Do not sum the independent operator microbenchmark savings or interpret this combination run as per-PR throughput evidence. Artifacts: f01 |
xytpai
changed the title
Contributor
Author
Expanded stream paths: model ABBAReference already contains original #7219, #7217, #7210, #7090, and ATOM cache fix #2858. The sole incremental change is the six extra MoE helper stream paths.
Mean-of-process-medians descriptive changes (not a significance test):
Input hashes and worker import paths checked for all arms. 20 short +20 long measured requests and 6 C2 rounds per process. Same model/TP4/Simple/FP8 KV/EAGLE3 settings. All warmups and post-timing cProfile diagnostic requests excluded. The sample is small and process-to-process differences remain visible; do not extrapolate operator savings or guarantee full-model speedups. Synthetic acceptance 0.5933: output text is not accuracy evidence. All raw timings, cache counters and source manifest are retained in the experiment directory. Observed compile-event totals were not uniform (baseline A recorded Triton misses; expanded B recorded two FlyDSL compiles). The logger's ordinary-forward label does not identify timing boundaries, so these totals alone cannot prove zero compilation inside every measured sample. Warmup requests are excluded, but we retain this instrumentation limitation rather than claim a statistically isolated end-to-end win. This is the expanded #7219 versus the original #7219, not the earlier #7217+#7219 combined run. Both reports remain valid for their own configurations. |
valarLip
approved these changes
Nov 2, 2026
XiaobingSuper
added a commit
that referenced
this pull request
Nov 2, 2026
XiaobingSuper
added a commit
that referenced
this pull request
Nov 2, 2026
XiaobingSuper
added a commit
that referenced
this pull request
Nov 2, 2026
current_stream_arg (#7219) reads the stream through torch._C._cuda_getCurrentRawStream, which the dry run's torch.cuda.current_stream stub does not cover. On a host without a GPU, such as a wheel build, every grouped MoE replay then fails with 'no ROCm-capable device is detected', so PREBUILD_KERNELS leaves the gfx1250 grouped MoE kernels out of the FlyDSL cache. The dry run now stubs that query too. Without a GPU, run_aot replays all 1160 gfx1250 grouped MoE jobs (0 before) and compiles the 16359 pooled kernels as before.
2 of 6 tasks
zhiding512
added a commit
that referenced
this pull request
Nov 2, 2026
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>
zhiding512
added a commit
that referenced
this pull request
Nov 3, 2026
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>
vgokhale
added a commit
that referenced
this pull request
Nov 5, 2026
* CI: auto-update split test times (#7290)
Co-authored-by: gyohuangxin <42127654+gyohuangxin@users.noreply.github.com>
* [HIP] [FlyDSL] perf(mega_moe): fuse the gfx1250 compact dispatch and GEMM1 into one stage-1 kernel (#6564)
* [ASM] [HIP] [OPUS] Add gfx950 MLA v4.1 a8w4 prefill kernels for an NVFP4 prefix KV (#7124)
* Add a gfx950 MLA v4.1 a8w4 prefill kernel for an NVFP4 prefix KV
* Add a gfx950 MLA v4.1 a8w4 16mx4_64nx1 prefill kernel and run the NVFP4 prefix in f16
* Run the dsa_v41 a8w4 NVFP4 prefix in bf16 to keep MXFP8 Q exact
* [Perf] Avoid torch Stream construction across grouped MoE launches (#7219)
* Avoid constructing a torch Stream for default grouped GEMM launches
* Reuse raw stream arguments through MoE route, quant and reduction wrappers
* [Bugfix] MLA decode: merge only the splits gfx950 bf16 gqa32 stage 1 writes (#7182)
On gfx950, mla_a16w16_qh16_m16x4_n32x1_coex0_mask1 (bf16, gqa 32,
non-persistent) deals KV to splits round-robin in 32-token tiles and
exits without writing a split that gets no tile. mla_decode_fwd used
mgc = 16 for nhead 32, so stage 2 merged min(splits, cdiv(kv, 16))
splits and read up to twice as many as stage 1 wrote, from torch.empty
buffers. Rows with less KV than splits * 32 came back wrong, or NaN
when that memory held NaN. Use mgc = 32 for this kernel, as the fp8
qh64 kernel already does with mgc = 64. gfx942's gqa32 kernel uses
16-token tiles and keeps mgc = 16.
The new op test runs that case with NaN-filled split buffers.
Signed-off-by: Batzorig Zorigoo <baz@fractalyze.io>
* [Config] [gfx1250] Tune MiniMax-M3 TP4 BF16, indexer-CP QKV and MoE token 1024 (#7209)
* [Triton/Gluon] Window-aware split-KV for sliding-window verify on gfx950 (#6889)
* [Triton/Gluon] Split window-trimmed sliding-window verify on gfx950
On gfx950 a sliding-window call with q > 1 never takes the split grid
(_gfx950_split_grid returns False, and kv_split["SW"] caps MAX_SEGMENTS
at 1), so a spec-decode verify on a sliding layer runs num_kv_heads
programs per sequence over the whole window. Splitting it only pays once
the KV range is cut to the window: the split carves seq_len, so over a
full context the window lands in one or two segments and the rest are
empty (31.9-37.6 us vs 29 us unsplit when forced).
Split only when the KV range is bounded by the window. Row i of a q-row
sequence of length L reads keys from L - q + i - (SW - 1) on, so a block
table trimmed to whole pages holds at most SW + q - 1 + block_size - 1
keys. The new SWB config axis is sliding_window_bounds_kv(): q > 1 and
max_seqlen_k <= sliding_window + max_seqlen_q + block_size - 2. A batch
whose contexts are all that short is bounded as well, trimmed or not, and
takes the same path. _gfx950_split_grid lets verify-sized SWB calls
(max_seqlen_q <= 16 and max_seqlen_q * num_queries_per_kv <= 128)
through and then decides like any other batch (NUM_SEGMENTS > 1).
Longer queries (prefill chunks) keep the 2d grid: with q rows per
sequence in the reduce, MHA q = 128 at window 256-512 ran up to 1.33x
slower split.
The gfx950 gluon kv_split table gains an "SW.SWB" entry, the SW entry
without its segment cap, plus:
- SPLIT_FIT_BUDGET (new): the largest power-of-two split that keeps the
launch within SEGMENTS_PER_CU (1) programs per CU, sized by the real
one-block-per-sequence query grid (_gfx950_split_params). The default
ceil-rounded share splits 2 ways up to B = 51 and loses 30-40% from
B = 36 (39.7-43.0 us vs 30.4-30.8 us on the 2d grid).
- SPLIT_KV_WINDOW (new): count the KV tiles over the sliding_window +
max_seqlen_q - 1 keys the rows can see, not the page slack of a trimmed
table, so SPLIT_MIN_TILES 4 keeps window 128 unsplit at block 128 too
(64 q / 8 kv heads, q = 4: 8.2 us split vs 6.7 us 2d).
- SPLIT_SEQS_PER_TILE 1.5 (new): no split once num_seqs exceeds 1.5 x the
KV tiles. The triton reduce_segments pass that every gfx950 split pays
costs more with the batch (7.6 us at B = 32, 2 segments), while the 2d
kernel's time follows the KV length, so at windows 256-1024 a split at
B = 12-32 ran up to 1.87x slower than the 2d grid.
Kernel: a split whose tiles all fall before the window start stores an
empty partial, and split 0 is the only one that counts the attention
sink, so a window starting past tile 0 (block size > 64, or a shorter
row in an untrimmed bounded batch) dropped the sink from the softmax.
The empty-partial path of split 0 now stores M = sink, L = 1.
Unchanged dispatch: SW calls with a context past the bound, SW decode
(q == 1 keeps the cap), q > 16, window 128, full attention, gfx1250 and
the triton backend. Of 328,800 no-launch shapes (window 128-32768, q
1-128, B 1-128, 5 head layouts, D 64-256, bf16/fp8, block 16-256,
ctx bound -300..+5000) only gluon SW calls with 1 < q <= 16 inside the
bound change grid (22,392).
MI355X, 16 q / 2 kv heads, head 64, bf16, block 16, window 2048, q = 8,
ctx 32768 (us per call, CUDA-graph replay, natural dispatch):
B 1 2 4 8 16 32 48 64
full 29.0 29.0 29.1 29.5 30.0 30.5 30.6 30.9
trimmed 6.7 7.2 8.6 10.9 15.8 25.1 30.6 30.8
segments 64 32 16 8 4 2 2d 2d
Full-attention floor at 2056 keys: 6.6 / 6.8 / 8.1 / 9.4 / 14.4 / 21.9.
The natural pick matches the fastest forced count (1-128) at every B.
Shorter windows, trimmed, same heads (base 2d -> split, us):
W 512: B4 9.78 -> 6.79, B8 9.79 -> 7.61, B12 9.82 -> 9.19, B16+ 2d
W 1024: B4 16.30 -> 7.66, B8 16.35 -> 8.61, B16 16.71 -> 12.27, B32 2d
W 256: B4 6.54 -> 6.30, B8+ 2d
Known cost: a caller that captures a graph at max_seqlen_k = window bound
and replays it for batches whose real contexts are a few hundred keys
pays the split + reduce where the 2d kernel was cheaper (B >= 12, ctx <=
1000). A cheaper gfx950 reduce would remove that.
Test: test_triton_unified_attn_sw_verify_split (full and trimmed tables
vs reference with a dominant sink, exact segment counts: trimmed verify
splits, full context stays 2d, SW decode stays 1 segment, q = 17 and
q = 128 prefill chunks stay 2d, block 128 and an untrimmed bounded batch
leave split 0 empty) and test_unified_attn_sw_kv_split_config
(CPU-only kv_split picks at the bound edge, untrimmed, decode, window 128,
and the batch sizes where the split stops).
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* [Triton/Gluon] Drop SPLIT_KV_WINDOW; clamp every window-bounded split
Review follow-up. The SPLIT_KV_WINDOW knob was set only in the SW.SWB
entry, so it configured nothing. compute_segment_params now counts KV
tiles over the sliding_window + max_seqlen_q - 1 keys the rows can see
whenever sliding_window_bounds_kv() holds, with no config switch.
The clamp is gated on the bound, not on sliding_window > 0. The split
carves the whole seq_len, so on an untrimmed long context a window-sized
count leaves most segments idle. No unbounded sliding-window call reaches
a split lookup today, so the two gates pick the same grid.
CPU dispatch sweep vs the previous commit, 2,048,256 shapes per arch on
gfx950 (gluon and triton), gfx1250, gfx1100 and gfx1201: no 2d/3d,
NUM_SEGMENTS or TILE_SIZE change. 1,184 gfx950 gluon shapes change only
the triton reduce num_warps (head 256, block >= 64, window + q - 1 of
about 257-512), since the reduce lookup now sees the same clamped length.
test_unified_attn_sw_kv_split_config: 11 passed on CPU; removing the clamp
fails the window 128, block 128 case.
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* [Triton/Gluon] Key the SWA split on window size: numeric SW axis
Review follow-up. The SW axis was a bool, so a table could not target a
window size. SW is now "opt_num": the keys a row can see,
min(sliding_window, max_seqlen_k), floored at 1, and 0 with no window.
SW_EQ/LEQ/GEQ_<n> keys pick by size, a bare SW matches any window, and
no window matches only "any". Tables that only use a bare SW key (gfx950
gluon and triton, gfx1250 gluon, gfx1201 triton) resolve as before.
_index rejects a table that mixes SW_LEQ and SW_GEQ keys, because every
matching LEQ key would outrank every GEQ key. It is cached on the key
names, so the check runs once per table.
The gfx950 gluon kv_split entry SW.SWB becomes SW_GEQ_256.SWB. Bare SW,
with MAX_SEGMENTS 1, stays the fallback for every other window and for SW
decode. 256 keys is 4 tiles of 64, the smallest window measured to win.
Window 128 now falls back through the threshold, not SPLIT_MIN_TILES.
SWB stays: an untrimmed long context also gives SW = window, and
splitting that table was measured slower.
CPU dispatch sweep vs the previous commit, 3,692,160 shapes per arch
(grid, NUM_SEGMENTS, reduce config, and attn_2d/attn_3d configs for both
backends) on gfx950, gfx1250, gfx1100 and gfx1201: the only change is
6,300 gfx950 gluon shapes at window 255 (q 4-16, just under the
threshold), split -> 2d. Windows above 2048 still split; their perf is
not measured.
Tests: test_unified_attn_sw_kv_split_config gains window 255 vs 256 at
the threshold and window 256 at block 128 (the clamp case).
test_unified_attn_sw_axis_lookup covers the candidate order, the axis
value and its floor, and the LEQ/GEQ guard. 15 passed on CPU. Each of
these mutants fails at least one test: threshold, clamp, guard, bare-SW
fallback, the no-window check, min(), and the floor.
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* [Triton/Gluon] Test the SWA split without sinks on long pages
Review follow-up (Copilot). Without sinks, a sliding-window split over a
bounded table whose page is much longer than a tile can leave 32 or more
empty leading segments. One example is head 256, q 8, KV 10246, window
2048, block 8192, 64 segments. Before #6888 that made the first
SEGMENT_BLOCK chunk of reduce_segments all -inf, and the merge returned
NaN. #6888's two-pass merge takes the global max first, so those chunks
now weigh 0. Under TRITON_INTERPRET=1, main's reduce_segments on the
partials both shapes leave gives no NaN and matches the single-chunk
reduce within 1.2e-7.
test_triton_unified_attn_sw_split_no_sinks runs both of Copilot's shapes
on gfx950 with head 256 and no sinks, against the reference.
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* [Triton/Gluon] SWA split tests: check accuracy only
Review follow-up. Tests here target accuracy, not configs or routing.
- Remove test_unified_attn_sw_kv_split_config and
test_unified_attn_sw_axis_lookup. They pinned kv_split picks and lookup
order, which couples the tests to the tuned configs and heuristics.
- test_triton_unified_attn_sw_verify_split and
test_triton_unified_attn_sw_split_no_sinks no longer record or assert
the grid. Both still compare full and trimmed tables, or the no-sink
long-page shapes, against the reference.
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* [Triton/Gluon] Window-aware split-KV for sliding-window attention
With a sliding window, every split-KV launch carved [0, seq_len) into
NUM_SEGMENTS tile ranges, so over an untrimmed long context only the few
segments that overlap the window did work. That is why this PR so far
split only tables bounded by the window.
Carve only the window instead. For SLIDING_WINDOW > 0 each sequence
splits the tiles [tile_lo, cdiv(seq_len, TILE_SIZE)), where tile_lo is
the tile of row 0's first visible key, max(context_len - SLIDING_WINDOW
+ 1, 0) // TILE_SIZE. It is per sequence, so every query block and the
reduce agree. The split launch already puts a sequence's spec-decode rows
in one query block when q * num_queries_per_kv <= 128. The window's tiles
are now what gets split across the segment workgroups.
- The gfx950 gluon split, kernel_unified_attention_3d and reduce_segments
carve this way. The triton pair shares split_window_tiles, and
reduce_segments gets SLIDING_WINDOW from the launch. The gfx1250 3d
wrapper asserts no window: its split still carves seq_len, and
use_2d_kernel never sends a window there. With no window, all three
kernels compile to the same gfx950 instructions as before.
- Dispatch: an untrimmed table now splits like a trimmed one.
_gfx950_split_grid drops the bound check and keeps the q <= 16 and
q * num_queries_per_kv <= 128 gates. The kv_len clamp applies to every
sliding-window call. The SWB axis and sliding_window_bounds_kv are gone,
and the split entry is SW_GEQ_256.Q_GEQ_2, so sliding-window decode
stays on the capped bare SW entry. CPU dispatch sweep vs the previous
commit, 3,692,160 shapes per arch: on gfx950 the only change is 75,220
gluon sliding-window calls with 1 < q <= 16 past the old bound going
2d -> split. Each gets exactly the dispatch the bounded table already
had. gfx1250 and gfx1100 do not change.
The existing accuracy tests now run the full-table rows through the
window carve: test_triton_unified_attn_sw_verify_split and
test_triton_unified_attn_sw_split_no_sinks.
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---------
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Satya Nikhil Kodukula <nikhil.kodukula@gmail.com>
* [FlyDSL] gfx1201 (RDNA4) FlyDSL BF16 attention optimizations & FP8 attention implementation (#4188)
* feat(flydsl): add per-tensor fp8 gfx1201 flash-attention kernel
* feat(flydsl): generalize the gfx1201 bf16 flash-attention kernel and wire fp8 dispatch
* feat(flydsl): fully-FlyDSL fp8 per-tensor quant+rotation for gfx1201
* refactor(flydsl): share and harden gfx1201 attention
* test(flydsl): cover gfx1201 bf16 and fp8 attention
* bench(flydsl): add gfx1201 fmha benchmark
* merge(main): resolve gfx1201 FlyDSL attention conflict
* fix(flydsl): validate attention softmax scale
* fix(flydsl): address gfx1201 attention review
* refactor(flydsl): centralize gfx1201 attention build parameters
* test(flydsl): execute gfx1201 pytest suite in op-test shards
* refactor(flydsl): share gfx1201 attention helpers and localize FP8 intrinsics
* refactor(flydsl): clean up gfx1201 kernels
* refactor(flydsl): organize gfx1201 FMHA modules
* refactor(flydsl): localize gfx1201 stream readiness
* refactor(flydsl): use typed pointers in fp8 quant
* refactor(flydsl): use public gfx1201 kernel attributes
* refactor(flydsl): export gfx1201 fp8 quant directly
* test(flydsl): keep gfx1201 FMHA public regressions
* feat(flydsl): add native gfx1201 1 FP8 quantization for more shapes
* flydsl: simplify gfx1201 FMHA stream/quant integration
* op_tests: modernize gfx1201 FMHA validation and benchmark
* flydsl: harden gfx1201 FMHA stream validation and multi-device tests
* flydsl: remove empty gfx1201 FMHA package init
* flydsl: migrate gfx1201 FP8 flash-attn to FlyDSL buffer-view/copy API
* flydsl: use static layouts and coordinate helpers for gfx1201 FP8 KV loads
* flydsl: validate gfx1201 FP8 buffer indices against signed Int32 limits
* flydsl: replace raw arith intrinsics with typed FlyDSL scalar/vector ops in gfx1201 FP8 FMHA
* refactor(flydsl): use typed exp2 in gfx1201 BF16 attention
* test(flydsl): tighten gfx1201 FMHA validation
* refactor: modernize gfx1201 FlyDSL FP8 paths
* test: consolidate FlyDSL gfx1201 op tests
* fix: support wide FlyDSL FP8 quantizer loads
* fix: validate FlyDSL quant offsets and stream coverage
* fix(flydsl): harden gfx1201 FMHA correctness
* fix(flydsl): align FP8 softmax normalization to prevent long-sequence normalization bias
---------
Co-authored-by: vivienfanghuagood <vivienfanghua@163.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Triton/Gluon] Enable and tune pipelined unshuffled MLA decode on gfx1250 (#6545)
* Allow the pipelined version of the kernel to be used
* Add tuning config
* More tuning changes
* Pad fp8 rows
* Update reduce waves_per_eu
* Fit 64 heads into two waves
* Add unit test
* Formatting
* Attend to Copilot review comments
* Test fixes
* Update tuning params
* Attend to Copilot review comments
* Fix formatting issues
* Add extra guards
* Add params to desc
* Fix arch name for test
* [Triton/Gluon] [Kernel] Add dual-layout H16 MXFP4 quantization (#6458)
* [Triton/Gluon] [Kernel] Add dual-layout H16 MXFP4 quantization
Fuse row-wise MXFP4 quantization with signed-H16 quantization of the
transposed input. Support independent stochastic rounding and GEMM
layouts on gfx950, with tests, benchmarks and a tuning configuration.
Signed-off-by: DaiXindi-AMD <xdai@amd.com>
* [Misc] Rename fused MXFP4 tuning configuration directory
Signed-off-by: DaiXindi-AMD <xdai@amd.com>
* [Bugfix] Address fused transpose Hadamard MXFP4 review feedback
Use the specific fused API and config key, relocate the four unchanged device helpers to quant.py, and document sign validation and physical scale shapes. Fix module-object monkeypatching after the API/module rename.
Validated all eight changed Python files with Ruff 0.16.0, Black 26.3.0 and 26.10.1, compile and diff checks. Kernel and helper AST equivalence plus isolated CPU helper checks passed. Full package import and GPU tests were not rerun: prior import/collection is blocked by missing Ninja, and no GPU allocation is authorized.
Signed-off-by: DaiXindi-AMD <xdai@amd.com>
* [Docs] Use GPU graph capture terminology for H16 signs
Signed-off-by: DaiXindi-AMD <xdai@amd.com>
* [Misc] Keep MXFP4 layout constants in the host utility
Move the five unchanged layout constants to utils/mxfp4.py and remove the constants-only utils/_triton/mxfp4.py module. The public wrapper still passes the same values as constexpr arguments; device kernel and helper code remain unchanged.
Signed-off-by: DaiXindi-AMD <xdai@amd.com>
---------
Signed-off-by: DaiXindi-AMD <xdai@amd.com>
* [Triton/Gluon] [Perf] fused_reduce_act_mul_and_mxfp4_quant: wider launch tile, add round_to_input_dtype and transpose_scale (#6024)
* [Triton/Gluon] [Perf] fused_reduce_act_mul_and_mxfp4_quant: wider launch tile, add round_to_input_dtype and transpose_scale
For N_half > 1024 (no shuffle, no scale padding, no split-K) launch one warp
per row over 256 (M <= 64) or 512 columns instead of a 1x128 tile on 4 warps.
On MI355X thats 1.02-3.77x faster over N_half 2048..28672 and M 1..8192
(9.85 -> 4.29 us at M=64, N_half=14336), output bytes unchanged.
round_to_input_dtype rounds act(x) * x_mul to the input dtype before the
quant, so the result is bitwise equal to act_mul in that dtype followed by
dynamic_mxfp4_quant. transpose_scale (also added to fused_rms_mxfp4_quant)
returns the scales column-major, the layout dynamic_mxfp4_quant returns.
Both default to False.
Signed-off-by: adlashab <adlashab@amd.com>
* Fix a typo in the act_mul launch tile comment
Signed-off-by: adlashab <adlashab@amd.com>
* [Triton/Gluon] fused_reduce_act_mul_and_mxfp4_quant: honor scaling_mode="ceil"
scaling_mode was accepted but the act_mul kernel always used the "even"
scale. Pass it through to _mxfp4_quant_op, which already has "ceil" as
SCALING_MODE 1: 2^ceil(log2(amax / 6)), the RoundUp scale
per_1x32_f4_quant_hip uses by default. With round_to_input_dtype the output
is then the same fp4 bytes and the same scales as silu_mul followed by
per_1x32_f4_quant_hip, in the same scale layout (plain, padded or shuffled),
so it can stand in for that quant in front of gemm_afp4wfp4_preshuffle /
gemm_a4w4. That is what vLLM's MXFP4 linears on the ASM layout run. One
difference: a 32-block thats all zeros gets scale byte 1 here and 92 from the
HIP kernel. All its fp4 values are 0, so the GEMM output is the same.
Also let the wider launch tile run with scale_shuffle_padding when shuffle is
off. The padding only changes the allocation, the scale store is the same.
New test against per_1x32_f4_quant_hip for plain, padded and shuffled scales.
Signed-off-by: adlashab <adlashab@amd.com>
* [Triton/Gluon] [Test] Check fused act_mul + MXFP4 quant against a torch reference
Add a plain torch reference for the per-1x32 MXFP4 quant and compare
fused_reduce_act_mul_and_mxfp4_quant against it bit for bit. The reference
uses the shared CPU helpers from aiter.utility.fp4_utils: the e8m0 scale from
f32_to_mx_e8m0_scale (Even for "even", RoundUp for "ceil"), then x / scale
rounded to nearest even on the E2M1 grid with f32_to_mxfp4.
The test covers both scaling modes, round_to_input_dtype on and off, bf16 and
fp16 input, plain / transposed / padded / shuffled scales, odd M, and an
all-zero 32-block at each end of the output. In "ceil" mode the Triton scale
clamps amax to 6 * 2^-126, so an all-zero block gets scale 1 where RoundUp
gives 0. The fp4 values of that block are 0 either way, the test expects 1.
Also use device="cuda" in the tests this PR adds, like the rest of the file
since #6783.
Signed-off-by: adlashab <adlashab@amd.com>
---------
Signed-off-by: adlashab <adlashab@amd.com>
Co-authored-by: Shao-Chun Lee <Shao-Chun.Lee@amd.com>
* [CI] Run the selected Triton tests on each merge to main (#7229)
Since #7029, main runs Triton Test only once a night, as the full
suite. A merge that breaks a Triton test then shows up only the next
day, and the nightly cannot say which merge it was.
A push to main now runs Triton Test again, with the same test selection
as a PR: the tests that the merged commit affects, against its parent on
main. MI300X runs only when the selection asks for it, as on a PR. Each
push gets a concurrency group of its own, so a merge no longer cancels
the run of the one before. A push run does not wait for the pre-checks
and does not post on the nightly status issue.
The nightly and manual runs still run the full suite on MI35X and
MI300X.
Co-authored-by: Xin Huang <Xin.Huang@amd.com>
* [FlyDSL] [gfx1250] mega_moe S1F: reachable stage-1 hook, plan precompile, tuned defaults (on by default) (#7315)
* [FlyDSL] [gfx1250] mega_moe S1F: reachable stage-1 hook, plan precompile, opt-in tuning knobs
On top of the stage-1 fusion (#6564):
- Fix: swap the stage-1 hook into dispatch_context as well as scatter. The
grouped GEMM reads it off dispatch_context, so replacing scatter alone
left the probe / fused stage-1 path unreachable on the compact path.
- Precompile every stage-1 compact-plan variant at setup (compile only,
AITER_S1F_PRECOMPILE=0 skips). Eager widths otherwise JIT-compiled inside
the serving loop: 250-470 ms prefill steps for the first ~100 s of a run.
- AITER_S1F_MAX_RECV=N: steps whose receive bound exceeds N run the whole
S1F=0 path (plan, GEMM1, no probe), e.g. fuse decode only.
- AITER_S1F_ANY_TILE=1: fuse when the tuned GEMM1's tile_m matches; the
stage-1 kernel runs its own 256x256 n/k tiles.
- AITER_S1F_PRODUCERS / AITER_S1F_BLOCK_RUNS: per tokens-per-rank bucket
overrides of the stage-1 tables.
- AITER_TDM_COMPACT_ROUTE_UPTO / _ROUTE_SPEC / _ROUTE_PARALLEL: keep the
route-parallel compact walk up to N tokens per rank.
All knobs default to the current behaviour.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* [FlyDSL] [gfx1250] mega_moe S1F: make S1F and its tuned settings the default
Make S1F and its tuned settings the default on gfx1250; env vars become
optional overrides. With nothing set, gfx1250 MegaMoE now runs the
configuration measured on DSV4.1-Flash DEP4 (MI455X, MLPerf Endpoints
agentic c512, +2.1% gen tok/s/GPU over S1F off):
- AITER_MEGA_S1F: on by default; =0 turns the stage-1 fusion off. Configs
or steps that do not fit still fall back to dispatch + GEMM1; the
fallback logs at info level unless AITER_MEGA_S1F=1 was set explicitly.
- S1F_BLOCK_RUNS table: (4, 2, 8, 8). AITER_S1F_BLOCK_RUNS /
AITER_S1F_PRODUCERS still override the tables.
- AITER_S1F_MAX_RECV: 4096 by default; 0 means no limit.
- AITER_S1F_ANY_TILE: on by default; =0 requires all GEMM1 tiles to match.
- AITER_TDM_COMPACT_ROUTE_UPTO / _ROUTE_SPEC: 1024 / 384,16 by default
(config._COMPACT_ROUTE_UPTO / _SPEC); UPTO=0 keeps _DISPATCH_COMPACT.
- AITER_S1F_PRECOMPILE stays on by default.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Read the new S1F/route env vars through the shared env module
Register the env vars this PR adds or changes (AITER_MEGA_S1F,
AITER_S1F_PRECOMPILE / MAX_RECV / ANY_TILE / PRODUCERS / BLOCK_RUNS,
AITER_TDM_COMPACT_ROUTE_UPTO / SPEC / PARALLEL) in mega_moe/envs.py, with
their defaults and parsing, and read them as envs.<NAME> instead of
os.environ. Env vars that were already on main are left as they are.
No behavior change.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Drop the new env vars; the tuned S1F/route settings are fixed defaults
This PR no longer adds any environment variable. mega_moe/envs.py is back
to main, and the settings tuned on DSV4.1-Flash DEP4 (MI455X, MLPerf
Endpoints agentic c512) are now plain defaults:
- stage-1 plans are always precompiled when S1F is on;
- _S1F_MAX_RECV = 4096 (s1f_ctl.py);
- GEMM1 fuses whenever its tile_m matches (stage-1 runs its own n/k tiles);
- S1F_PRODUCERS / S1F_BLOCK_RUNS come from the s1f_stage1.py tables
(block runs 4, 2, 8, 8);
- the compact dispatch always walks routes up to _COMPACT_ROUTE_UPTO = 1024
tokens/rank with _COMPACT_ROUTE_SPEC = (384, 16) (config.py); the
AITER_TDM_COMPACT_BLOCKS/WARPS override behaves as on main.
AITER_MEGA_S1F keeps its original os.environ read; only its default
changes to on (0 turns it off). S1F fallbacks log at info level.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
* [CI] Reuse FlyDSL AOT cache and reduce A8W8 GEMM test overhead (#7316)
* [CI] Reduce A8W8 GEMM benchmark repetition and preserve scale-mode coverage
Signed-off-by: yifehuan <278623243+yifehuan@users.noreply.github.com>
* Use supported MXFP8 row sizes in focused numerical tests
Signed-off-by: yifehuan <yifehuan@local>
* [CI] Reuse FlyDSL AOT cache and sample A8W8 benchmark timing
Signed-off-by: yifehuan <yifehuan@amd.com>
---------
Signed-off-by: yifehuan <278623243+yifehuan@users.noreply.github.com>
Signed-off-by: yifehuan <yifehuan@local>
Signed-off-by: yifehuan <yifehuan@amd.com>
Co-authored-by: yifehuan <278623243+yifehuan@users.noreply.github.com>
Co-authored-by: yifehuan <yifehuan@local>
* [FlyDSL] Fix grouped GEMM C ABI overflow for large EP buffers (#7299)
Pass address-only grouped GEMM operands as typed pointers instead of flattened tensors. FlyDSL packs runtime tensor dimensions as int32, so buffers with more than 2^31-1 elements fail in the generated C ABI fill function. The EP arena handle remains a constexpr and is already excluded from runtime packing.
Remove redundant scale-view normalization and replace its tests with host ABI coverage for large buffers, storage offsets, pitched layouts, cached dispatch, 48-bit addresses, and changed arena specializations.
Signed-off-by: Felix Li <felix.li@amd.com>
* [Triton/Gluon] [HIP] [OPUS] [AITERKER-209] Add gfx1250-strict architecture plumbing (#6767)
* [JIT] Add gfx1250-strict architecture plumbing
* [JIT] Fix ruff F821 and black formatting in gfx1250-strict plumbing
Add the missing build_targets._parse_gpu_archs_env import in core.py
(F821 flagged by CI ruff) and apply black to core.py and
test_build_targets_architecture.py.
Co-Authored-By: Claude Code <noreply@anthropic.com>
* [JIT] Add gfx1250-strict kernel inventory and WS2 build driver
Skeleton inventory for every JIT/FlyDSL/Triton path plus the driver that
fills its strict-build column from a strict-capable container.
Driver verified offline on a non-strict host (module enumeration, child
plumbing, skip classification, --update-md folding):
- CHILD called os.chdir without importing os, so every module would have
failed with NameError and the whole matrix would have read "fail".
- optCompilerConfig lists 135 modules but get_args_of_build("all") yields
123. The other 12 are not failures: 10 end in "tune" (skipped by design)
and 2 are experimental (gated by AITER_ENABLE_EXPERIMENTAL). They are now
recorded as skip-tune / skip-experimental instead of "missing", so the
matrix keeps an explicit result for all 135.
Actual pass/fail rows still require the strict toolchain; stock ROCm 7.0.0
rejects gfx1250-strict as an invalid target ID.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* [JIT] Fix empty GPU_ARCHS raising out of build_module
The gfx1250-strict plumbing routed GPU_ARCHS through
build_targets._parse_gpu_archs_env, which raises when the value holds no
architecture names. GPU_ARCHS set but empty -- what a Dockerfile's
`ENV GPU_ARCHS=` or a bare `export GPU_ARCHS=` produces -- previously just
meant "no cluster-launch flag, build normally"; it now raised RuntimeError
inside build_module and failed every JIT build. Blank, whitespace-only and
";" values had the same problem.
Fall back to the detected arch instead, matching aiter/ops/opus/_arch.py,
which treats an empty or "native" GPU_ARCHS as "probe the device" and
never raises. The arch test moves into _build_archs_from_env /
_cluster_launch_target so it is reachable from a CPU-only test.
Covered by TestClusterLaunchBuildArchs: blank/unset/"native" fall back to
detection, explicit archs win, and -DAITER_ENABLE_CLUSTER_LAUNCH still
follows gfx1250 or gfx1250-strict anywhere in a multi-arch GPU_ARCHS.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* [JIT] Inventory the gfx1250 dispatch sites a strict device silently misses
The infra PR keeps gfx1250-strict distinct through parsing, build selection
and the asm guard, but Python dispatch still compares against the exact
literal "gfx1250". On A0 get_gfx() returns "gfx1250-strict", so those
comparisons are False and the constant arch allowlists do not contain it:
the gfx1250 path is skipped with no error, which reads as a kernel-selection
change rather than a failure.
ws2_gfx1250_dispatch_sites.py enumerates them -- 150 uses across 65 files,
132 runtime comparisons and 18 allowlist entries -- and records where the
compared string comes from. get_gfx/get_gfx_runtime sites keep the
hyphenated spelling and are known to miss; triton-arch sites read Triton's
own target name, whose strict spelling stays unverified until the toolchain
lands.
The table is generated into the inventory between HTML markers and the
rewrite is idempotent, so it can be refreshed as owners decide each site.
No dispatch behaviour is changed here: every row is an explicit pending
decision, which is what workstream 2 asks for.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* [JIT] Pin the offload-target parser against hyphenated target names
_so_offload_archs backs the guard that keeps a gfx1250 and a gfx1250-strict
.so from standing in for each other, so the shapes it must parse deserve a
test of their own: a feature list after ':', the strict suffix, both
together, and a ROCm generic target.
The previous gfx\w+ pattern truncated at the first hyphen. That reported a
strict .so as plain gfx1250 -- the substitution this guard exists to
prevent -- and also read gfx9-4-generic as gfx9, which matches no running
arch and so rebuilt such a module on every import.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* [JIT] Inventory the arch-keyed tables that have no strict entry
The dispatch-site scan only saw comparisons against the "gfx1250" literal,
which left a blind spot: tables keyed by architecture. A strict device is
simply absent from those, so it takes whatever the lookup does with a
missing key.
Four of them carry a gfx1250 entry and no strict entry, and they do not
behave alike:
- radix_topk_one_block._LONG_RADIX_BY_ARCH and ._ONE_BLOCK_LDS_BUDGET_BYTES
read with a default, so A0 quietly falls back to the generic radix layout
and a zero LDS budget, losing the 320 KiB one-block path with nothing
logged.
- pa_decode's inline fp8 dtype table reads with .get() and raises
NotImplementedError naming the arch -- the "fail clearly" behaviour the
enablement plan asks for.
- mhc.MHC_FUSED_POST_PRE_M_UPPER_BOUND reads with a default of 1024, which
happens to equal the gfx1250 entry, so A0 is correct today only by
coincidence.
The generated table records the lookup style per row so the difference
between a silent fallback and a clear failure is visible to the owner
making the call. Behaviour is unchanged.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* [JIT] Extend the gfx1250 site inventory to C++ and skip handled sites
get_gpu_arch() in aiter_hip_common.h cuts the target id at ':' and keeps
everything before it, so a strict device reports "gfx1250-strict" in C++
too and the arch_id == "gfx1250" branches miss exactly as the Python ones
do. 27 of them across 14 files, 11 in mha_bwd.cu alone, where several pick
launch geometry rather than a code path -- a miss changes a block size
instead of raising.
Worth recording for the Opus owner: opus_co_launch_gfx1250.cuh builds its
code-object path from get_gpu_arch(), so A0 looks for gen_co/gfx1250-strict
and fails with a clear AITER_CHECK rather than loading a B0 object. That is
what the plan asks for, so it is called out in the table's preamble instead
of being listed as a gap.
The scan now skips any site with a strict spelling within four lines, so
the asm gate's own early return no longer pads the table and a site drops
out once an owner has handled it.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* [JIT] Count the tuned rows a strict device cannot use
Tuned GEMM and MoE configs carry a gfx column and are looked up through
get_gfx(), which reports gfx1250-strict on A0. 1620 rows across 18 config
files are keyed to gfx1250 and none to gfx1250-strict, so on A0 every shape
they cover falls back to the default path.
That is what the plan asks for -- strict tuning has to be explicit and
measured, not a filename alias -- so the generated table is the workstream
3 worklist rather than a defect list: each file needs either measured
strict rows or a recorded decision that the B0 rows hold on A0.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* [JIT] Generate the module table's gfx1250 signal column
The column was hand-filled and followed no single rule: 13 of 135 rows
disagreed with any consistent reading of the sources, in both directions.
module_custom_all_reduce_gfx1250 was marked by its module name while
module_cache was marked `guard` although its sources also compare against
the "gfx1250" string literal. Owners use this column to decide where to
look, so wrong values send them to the wrong file.
--refresh-signals now computes it from the sources that
core.get_args_of_build actually resolves, falling back to the raw config
expressions for the tune and experimental modules that call skips. The
markers are explicit and can combine: `name` for the module name, `str` for
the "gfx1250" string literal, `macro` for the __gfx1250__ guard.
The distinction matters for A0 risk. `macro` rows keep compiling, because
the strict target defines both __gfx1250_strict__ and __gfx1250__. `str`
rows are the runtime decisions a strict device silently misses. The refresh
changed 27 cells and is idempotent.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* [JIT] Tie the three arch tables together with a test
Adding a target means editing GFX_MAP, GFX_CU_NUM_MAP and core.py's
allowed_archs, and nothing tied them together. allowed_archs carried a
duplicate "gfx1250" entry until this series removed it, which is what that
drift looks like in practice.
Two invariants hold today and are now pinned: every arch with an offline CU
default is accepted by validate_and_update_archs and maps back through
GFX_MAP, and gfx1250-strict does the same. Deleting gfx1250-strict from
allowed_archs fails both tests, so a half-landed target is caught in CPU-only
CI rather than at build time.
The check is behavioural -- it drives validate_and_update_archs and
get_gfx_custom_op_core with GPU_ARCHS set -- so it also covers the arch
reaching hipcc unchanged as --offload-arch, hyphen included.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* [JIT] Record CK modules as skipped when the submodule is not checked out
aiter's CK-free auto-exclusion tests os.path.isdir(CK_3RDPARTY_DIR), which
is True for an uninitialised submodule: the directory exists, it is simply
empty. A checkout without `git submodule update` therefore still tries to
build the CK modules, and they fail on missing sources.
core._get_ck_exclude_modules() reports 69 of the 135 modules as
CK-dependent, so on such a checkout more than half the build matrix would
fill with failures that say nothing about the strict target. The driver now
treats an empty composable_kernel directory as CK-free and records those
modules as skip-no-ck, keeping pass/fail for results that are actually
about gfx1250-strict.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* [JIT] Pin PYTORCH_ROCM_ARCH in the build driver
Setting GPU_ARCHS alone does not reach the compiler. torch's cpp_extension
takes --offload-arch from PYTORCH_ROCM_ARCH, so with GPU_ARCHS=gfx1250 the
modules still compiled for torch's own default list -- observed in a ROCm
10.0.0 container as
--offload-arch=gfx90a ... --offload-arch=gfx1151
with no gfx1250 anywhere. Every row of the matrix would have been a result
for the wrong architectures while looking like a validated one. The driver
now derives PYTORCH_ROCM_ARCH from GPU_ARCHS; with that fixed,
module_aiter_unary builds for gfx1250 and build.ninja carries exactly one
--offload-arch.
Failures also kept the wrong 15 lines: ninja echoes the whole compile
command, which pushes the diagnostic out of a plain tail, so a failing
module recorded only the Python traceback. The recorded tail now leads with
up to five "error:" lines.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* [JIT] Record what ROCm 10.0.0 does and does not provide for strict
Checked in rocm/atom-dev:latest-rocm10.0 -- ROCm 10.0.0, AMD clang
23.0.0git (ROCm/llvm-project 8f497e0992fb), HIP 7.15.26333. hipcc compiles
gfx942 and gfx1250 and rejects gfx1250-strict as an unsupported HIP gpu
architecture. The processor list is gfx1250 and gfx1251, and neither
libLLVM nor amdclang contains the string gfx1250-strict or
__gfx1250_strict__ at all.
The only A0/B0 notion in that toolchain is the AMDGPU subtarget feature
gfx1250-b0-specific, and it is not reachable from the driver: not a
processor, rejected inside a target id, and ignored as a clang target
feature.
So the build matrix cannot be filled from the public ROCm 10 container --
the plan's toolchain dependency is unmet rather than satisfied -- and there
is a question worth settling before the identity model is locked in. If the
mechanism that ships is a subtarget feature on gfx1250 rather than a
separate processor name, an A0 and a B0 object both carry the target id
gfx1250, which is exactly what _so_offload_archs reads, so that guard would
need a different discriminator. The survey records the evidence and the
question without presuming the answer.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* [JIT] Lead the recorded failure tail with the real compiler diagnostic
The error capture grabbed the first five lines containing "error:", which
on a failing build are usually the Python RuntimeError and
CalledProcessError wrapper lines, not the compiler's own message. Prefer
"file:line:col: error:" frontend diagnostics when present, fall back to any
"error:" line otherwise (backend errors such as the gfx1250 DPP rejection
have no file:line prefix), and de-duplicate -- clang repeats each
diagnostic once per offload arch.
Verified against a real ROCm 10 matrix run: the ck_tile "extraneous closing
brace" / "template exec" and ck "operator[] in Tuple<>" errors now lead
their tails instead of the wrapper, and the opus DPP "broadcasts are not
supported on GFX10+" backend error still surfaces via the fallback.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* [JIT] Record the measured gfx1250 build matrix on ROCm 10.0.0
Ran the full 135-module matrix in rocm/atom-dev:latest-rocm10.0 with
GPU_ARCHS/PYTORCH_ROCM_ARCH=gfx1250: 102 pass, 21 fail, 12 skip. Strict
itself does not compile in this toolchain, so this is the gfx1250 baseline
-- anything that fails for gfx1250 cannot build for strict, which is gfx1250
plus an A0 code-object identity.
Every failure was rebuilt for gfx942 as a control, which corrected my first
read of them. They are not a ROCm-10/clang-23 regression: 18 of 21 build
cleanly on gfx942 and fail only on gfx1250, and the 3 FMHA-bwd modules that
timed out on gfx942 were compiling without error (module_mha_bwd from the
same family finished and passed). Every failure is a genuine gfx1250
porting gap:
- DPP broadcast dpp_ctrl modes GFX12 removed (moe_sorting_opus,
moe_mxfp4_aux)
- MFMA / buffer-load-LDS intrinsics GFX12 lacks (iq2r_moe,
mha_fwd_native_splitkv)
- a wavefrontsize64 request on a wave32 part (chunk_gdr_fwd_h)
- warp-size and gfx950-trait static assertions (topk_plain, moe_opus)
- CK's ck_tile TDM path, selected on gfx1250 by
CK_TILE_USE_BUFFER_ADDRESSING_BUILTIN=0 (the cktile/blockscale GEMMs,
activation, sample, the ck_tile FMHA fwd/prefill)
- fmha-v3-bwd asm configs not generated for gfx1250 (the mha_bwd family)
The inventory now carries the result table with the gfx942 control column
and the per-module cause, as the workstream-2 porting worklist.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* [JIT] Correct the strict macro model and inventory __gfx1250__ guards
The plan claimed the strict target defines both __gfx1250_strict__ and
__gfx1250__, so existing __gfx1250__ guards keep compiling. ROCm 10.1's
actual behaviour is the opposite: strict defines __gfx1250_strict__ and
does NOT define __gfx1250__ (ROCm/rocm-systems#12068). A __gfx1250__-only
device guard therefore compiles its body away on a strict agent -- a silent
empty kernel, the upstream TDM_Basic_load_2d failure, with no build error.
JAX hit the matching compiler-side fact and dropped strict from its ROCm
10.0 defaults (jax-ml/jax#41187), and ROCm 10.1 adds gfx1250-strict as a
distinct --offload-arch target across the stack (RCCL, CK, rocSHMEM), which
confirms this PR's distinct-name identity model and the .so guard that keys
on it.
Fixes, docs only -- no source guard is touched, because the change must not
break the shipping ROCm 7 build, whose compiler knows only __gfx1250__ and
has no strict target:
- plan and inventory legend now state the real macro behaviour and the
backward-compatible fix, `defined(__gfx1250__) || defined(__gfx1250_strict__)`,
whose strict arm is simply never taken on ROCm 7.
- the ROCm 10.0 survey's open "name vs feature" question is resolved: 10.1
ships a distinct target name; the 10.0 gfx1250-b0-specific subtarget
feature was an early-LLVM artifact.
- ws2_gfx1250_dispatch_sites.py now also inventories the __gfx1250__
preprocessor guards (95 across 24 csrc files) as a generated review table,
flagging which already combine with ||/&& to read first. This is the class
the earlier string-comparison scan did not cover.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* [JIT] Make every __gfx1250__ device guard also fire on strict
Non-asm operators are meant to run on both MI450 SKUs (A0 = gfx1250-strict,
B0 = gfx1250); only the shipped asm is B0-only. But the strict compiler
defines __gfx1250_strict__ and not __gfx1250__, so every
`#if defined(__gfx1250__)` device guard took its #else branch on a strict
build -- a silent wrong path that still compiles. Read on A0 this is:
get_smem_size returning 64 KB instead of 320 KB, mhc falling back to fp32
instead of the WMMA path, make_gmem returning the 0xffffffff sentinel so
stores drop, and the TDM sync counters using the legacy lgkmcnt/vmcnt
instead of the split dscnt/loadcnt.
Rewrite all 97 `defined(__gfx1250__)` occurrences across 24 csrc files to
`(defined(__gfx1250__) || defined(__gfx1250_strict__))`. The parentheses
keep it correct in every context: positive guards, `!(...)` exclusion
lists (strict is excluded exactly as gfx1250 is), and `||`/`&&` family
lists. The strict arm is never taken on the shipping ROCm 7 build, which
has no strict target, so that build is unaffected.
Compile-verified on ROCm 10.1 (clang 24.0.0git) for both gfx1250-strict and
gfx1250 on module_mhc and module_top_k_per_row; both pass for both targets
at low parallelism (the earlier full-matrix failures for these were -j179
contention, not real). Runtime correctness on A0 is left to the daily A0
op-test once the modules are built as gfx1250-strict.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* [JIT] Add strict to the compiled-path gfx1250 dispatch comparisons
Non-asm operators run on both MI450 SKUs, and strict is gfx1250 (A0)
hardware, so a `get_gfx() == "gfx1250"` dispatch that selects a compiled
path must also fire on gfx1250-strict -- otherwise A0 silently takes the
generic fallback and, where the gfx1250 branch picks an arch-specific
layout (wave32 tiling, LDS budget), gets a wrong result, not just a slow
one.
Rewrite the 101 compiled-path comparisons across 51 files from
`== "gfx1250"` / `in ("gfx1250",)` to `in ("gfx1250", "gfx1250-strict")`.
Deliberately NOT changed: the one asm-selecting dispatch,
gemm_op_a4w4.py `if gfx_arch in ["gfx1250"]`. Adding strict there would
route A0 into the f4gemm asm path, which the asm gate then throws on --
and the policy is that asm is simply unavailable on A0, so strict must keep
missing that comparison and fall to the compiled path. The classifier that
split compiled from asm sites flagged exactly this one site.
py_compile, ruff and black clean (topk.py's pre-existing fmt:skip parse
issue aside); the CPU-only arch/dispatch tests still pass. Multi-element
arch allowlists (gluon_archs, SUPPORTED_ARCHS) are left for a follow-up --
those are a capability/feature gap on strict, not the correctness-level
layout divergence this commit closes.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* [JIT] Cover != "gfx1250" and post-merge dispatch sites for strict
The first dispatch pass only rewrote == and single-element `in (...)`
comparisons, so `!= "gfx1250"` negations were left behind -- on strict
`arch != "gfx1250"` is True and the kernel takes the non-gfx1250 branch,
the same silent-wrong-path class the == sites had. The main merge also
added fresh gfx1250 comparisons (flydsl pa_decode, triton quant/mla,
mla backend plan).
Fix the 30 compiled-path sites: `!= "gfx1250"` ->
`not in ("gfx1250", "gfx1250-strict")`, `== "gfx1250"` -> `in (...)`.
Three asm-selecting sites are deliberately left (attention.py:1752,
mha.py:3779, mla_sparse_prefill.py:100) so strict keeps missing them and
A0 errors on asm rather than loading it, per the policy that asm is
unavailable on A0.
py_compile, ruff and black clean; CPU-only arch tests pass. Only the three
asm sites now compare against the bare "gfx1250" literal.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* [JIT] Add strict to the C++ arch_id == "gfx1250" dispatch comparisons
get_gpu_arch() in aiter_hip_common.h keeps the full target id, so a strict
device reports "gfx1250-strict" and every `arch_id == "gfx1250"` branch in
csrc misses -- the same silent-wrong-path class the Python and __gfx1250__
guards had. has_tdm_arch() returned false on strict even though A0 has TDM;
mha_bwd picked the wrong launch geometry; and so on.
Rewrite the 28 comparisons across 14 files: `X == "gfx1250"` ->
`(X == "gfx1250" || X == "gfx1250-strict")`, `X != "gfx1250"` ->
`(X != "gfx1250" && X != "gfx1250-strict")`, where X is a plain identifier
or get_gpu_arch(). The asm-dispatch sites are included too: the actual asm
load still funnels through require_gfx1250_asm_or_throw, which throws on A0,
so computing the dispatch as gfx1250 is safe and A0 still cannot load asm.
The already-strict-aware asm gate in aiter_hip_common.h is left untouched.
Paren balance verified; the gfx942/gfx950 C++ build in CI compiles it.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* [JIT] Add strict to the multi-element arch allowlists
The capability allowlists -- gluon_archs, SUPPORTED_ARCHS, _SUPPORTED_GFX,
DEVICE_ARCH tuples, the triton tuning kernel decorators -- list gfx1250 as
a supported arch but not gfx1250-strict, so strict (which is gfx1250
hardware) was treated as unsupported: no gluon path, excluded from the
tuned kernel sets, etc.
Add "gfx1250-strict" next to each "gfx1250" element in those arch tuples
and lists (41 across 22 files), skipping the substring-match helper
(quick_all_reduce's `gfx in arch`, where "gfx1250" already matches
"gfx1250-strict") and any asm-selecting context. With the gemm_a4w4 path
main restructured, the gfx1250 branch now routes plain MXFP4 to FlyDSL and
the rest to asm, so strict correctly gets the FlyDSL path and the asm gate
still handles A0.
ruff and black clean; CPU-only arch tests pass. This closes the last
aiter-side strict dispatch gap; what remains is A0 runtime validation and
the Opus .co rebuild.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* [JIT] Pin flydsl to 0.4.0.dev973 for gfx1250-strict LDS support
FlyDSL's arch handling is otherwise strict-ready because it matches by
prefix (arch.startswith("gfx1250") / "gfx12"), which gfx1250-strict
satisfies -- WMMA instruction selection, warp_size (32) and the compile
target all come out right. The one gap was the LDS-capacity map, keyed on
the exact "gfx1250" string, so on strict the lookup missed and
check_smem_capacity silently skipped the overflow guard.
0.4.0.dev973 adds "gfx1250-strict": 327680 to SMEM_CAPACITY_MAP (now in
flydsl.runtime.device), closing that gap. Bump the pin in requirements.txt
and setup.py (the version assertion and auto-install both derive from
FLYDSL_VERSION). Revert to the default release once the fix ships there.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* [JIT] Clean up: drop investigation tooling, tidy strict macro guards
Self-review cleanup, no behaviour change:
- Remove the ws2_* enumeration/matrix scripts and the generated
kernel-inventory doc. Those were one-off investigation flow (build
matrix, dispatch-site enumeration) and do not belong in the repo; nothing
in aiter imports them.
- Drop the dangling "former hotswap path" bullet from the enablement plan;
the term exists nowhere in the tree.
- Flatten the redundant parentheses the strict sweep left behind: the
mechanical `defined(__gfx1250__)` -> `(defined(__gfx1250__) ||
defined(__gfx1250_strict__))` wrapped the pair even inside OR chains.
Unwrap it where the context is pure-OR (90 sites); keep the parentheses
only where they are load-bearing -- negation `!(...)` and `&&` operands
(7 sites). Compile-verified on ROCm 10.1 for both gfx1250-strict and
gfx1250 (module_mhc, module_top_k_per_row); CPU-only arch tests pass.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix: pin flydsl build-system dep to 0.4.0.dev973 to match setup.py
pyproject.toml build-system.requires still pinned flydsl==0.4.0 while
setup.py (FLYDSL_VERSION) and requirements.txt use 0.4.0.dev973. Under
build isolation pip installed 0.4.0 as the build dep, then setup.py's
version assertion rejected it and the build failed. Align all three.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix: asm-gate error message covers both A0 detection routes
The gate blocks B0 asm on A0 via two routes: gcnArchName == gfx1250-strict,
or gcnArchName == gfx1250 with asicRevision 0. The message hardcoded
'gfx1250-strict/A0', which misleads when A0 is caught by the asicRevision
route (device name is literally gfx1250, rev 0). State both forms.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix: strict Opus codegen + 4 CI-red CPU tests
gen_instances: GPU_ARCHS=gfx1250-strict looked up KERNEL_REGISTRY by the
strict name, which has no registry key, so Opus codegen emitted an empty
header and module_deepgemm_opus failed to build (the only strict-specific
failure of the strict build matrix). Fold the registry key to the gfx1250
family (mirroring chip_info._detect_native) while keeping the gfx1250-strict
build target, so hipcc still compiles for mach 0xeb. Strict now resolves to
256 CU, same as gfx1250.
Tests (Draft CI had skipped these):
- test_opus_co_profile load_chip_parsers: keep the _GFX_ARCH_PATTERN constant
in the extracted AST namespace (the strict PR added it; parsers NameError'd).
- test_build_targets preserve->normalize: _detect_native now folds strict to
gfx1250 for dispatch; assert the new behavior.
- test_launcher_does_not_pin_the_first_device: use the first declared local
kid instead of hard-coded global kid 21016 (removed by #6937).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix: strict Opus codegen emits the gfx1250 family (routes>0), not an empty module
The earlier fix only folded the target list; the real route generation still
queried the registry by the strict name and produced routes=0 (empty
module_deepgemm_opus). Apply the gfx1250 family name at every registry/namespace
lookup -- available's 'if gfx in arches', get_registry, default_compiled_keys --
and at the OPUS_BUILD_HAS_* macro (C checks #ifdef OPUS_BUILD_HAS_GFX1250, and a
hyphenated STRICT macro is invalid C). KernelKey and the build target keep the
gfx1250-strict spelling, so hipcc still offloads for mach 0xeb.
Verified: GPU_ARCHS=gfx1250-strict now yields routes=228 (== gfx1250) with
OPUS_BUILD_HAS_GFX1250; CPU tests stay green.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* [JIT] Opus codegen: normalize gfx1250-strict targets to the gfx1250 family at the source
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix: bump flydsl pin to 0.4.0 to match main (dev973 would downgrade it on merge)
main pins flydsl==0.4.0; this PR still carried the 0.4.0.dev973 prerelease,
so merging would silently downgrade the whole repo. 0.4.0 is a superset of
dev973 (contains the gfx1250-strict LDS fix) and is what main already uses.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* test: drop redundant test_strict_resolves_through_all_three_tables
test_every_offline_target_resolves_through_all_three_tables already loops
over GFX_CU_NUM_MAP, which contains gfx1250-strict, so the strict-specific
case was an exact duplicate of one loop iteration.
* fix: avoid 'in' tuple test in gluon kernel constexpr (unsupported by gluon)
The strict-allowlist sweep turned 'ARCH == "gfx1250"' into
'ARCH in ("gfx1250", "gfx1250-strict")' inside gluon kernel bodies/constexpr.
Triton's gluon compiler does not implement the 'In' AST operator, so these
kernels failed to compile ('AST comparison operator In is not implemented'),
breaking test_pa_mqa_logits_2d_context_lens[gluon*] on gfx942/gfx950. Use
'(x == a or x == b)', which gluon supports, in kernel scope. Host-side Python
'in' checks are unaffected and left as-is.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Preserve exact gfx1250-strict identity; map to gfx1250 family only at dispatch
Per the ROCr rollout contract (all SW consumes the exact GFX ID ROCr reports),
global normalization in _detect_native() is the wrong boundary. Once ROCr
reports gfx1250-strict for A0, folding it to gfx1250 there makes
_needs_arch_rebuild() treat a strict .so as not covering a strict device and
rebuild in a loop; it also remixes A0/B0 in tuned configs, caches, and logs.
- chip_info._detect_native: drop the strict->gfx1250 fold; return the exact
ROCr-reported identity. get_gfx_runtime()/get_gfx() now report gfx1250-strict
on A0.
- build_targets.py: add the single source of truth for family folding --
GFX_BASE_ARCH, base_gfx_arch(), is_gfx1250_family(). Torch-free, so build
scripts and tests import it without side effects.
- Opus: replace its two private strict->gfx1250 folds (gen_instances
_registry_arch, opus_gemm_common _get_registry_records/get_registry/
default_compiled_keys) with base_gfx_arch(), so there is no second mapping.
hipcc still gets --offload-arch=gfx1250-strict and the .co stays mach-0xeb.
- mega_moe _require_gfx1250: restore the gfx1250-family guard (via
is_gfx1250_family) that main's #6430 refactor collapsed to bare gfx1250.
- Build/cache/tuning/.so/.co/ASM keep the exact identity: _needs_arch_rebuild
and _so_offload_archs already key on it, so A0 never loads a B0 object.
- Pre-existing gfx1250-only dispatch sites authored by other PRs stay exact;
with the fold gone they correctly distinguish strict (fall back) instead of
silently reusing a B0 path.
- Tests: flip the native-detection test into a regression lock that strict
identity is preserved; add family-helper and 4-way code-object rebuild-guard
tests (runtime/built strict x gfx1250 -> accept/reject).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Apply ruff import sort and black formatting to the identity refactor
Fixes the pre-checks (ruff I001 import order in mega_moe, black reformat of
mega_moe and test_build_targets_architecture) that gated check-signal and
skipped the Triton/Aiter test runs. No logic change.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* black: wrap the long gfx1250/gfx1250-strict dispatch conditions in the gluon kernels
The in->or fix in the pa_mqa_logits / unified_attention_2d gluon kernels left
two over-length boolean conditions that Black wraps. Formatting only; the
gluon compiler still sees the same (x == a or x == b) expression.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Define Opus base_gfx_arch locally instead of importing build_targets
The cross-package `from build_targets import base_gfx_arch` (via a
__file__-relative sys.path insert) breaks the Aiter wheel build:
gen_instances.py runs from a staging tree (aiter_meta/csrc/opus_gemm) where
../../aiter/jit/utils does not exist, so module_deepgemm_opus failed with
ModuleNotFoundError: No module named 'build_targets'. Duplicate the one-line
helper locally in opus_gemm_common (documented to mirror
build_targets.base_gfx_arch) so the standalone Opus codegen has no cross-tree
import dependency. gen_instances keeps importing it from opus_gemm_common.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* test(jit_cache): seed _cluster_launch_target in the build_module sandbox
test_jit_cache_transaction AST-extracts build_module into a sandbox namespace
with CPU deps injected. Since the gfx1250-strict work added a
_cluster_launch_target(get_gfx) call inside MainFunc, the sandbox hit
NameError: name '_cluster_launch_target' is not defined (25 TestBuildPublication
errors on MI35X). This surfaced only now because the earlier Opus wheel-build
failure skipped the standard tests. Inject the dependency like the others; for
the gfx942 test context the real function returns False, so the mock matches.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: zufayu <8138883+zufayu@users.noreply.github.com>
Co-authored-by: Claude Code <noreply@anthropic.com>
Co-authored-by: root <root@hjbog-srdc-24.amd.com>
* [HIP] [FlyDSL] optimize dsv41 perf for decode shape (#7309)
* optimize the mxfp8 op for decode shape
* retune mxfp8 tables with optimize decode shape
---------
Co-authored-by: solin <bingzhou@amd.com>
* [HIP] [CK] [JIT] [Tuner] Rank accuracy-qualified kernels by package energy (#6970)
* [Tuner] Rank accuracy-qualified kernels by package energy
Add --select power so shared tuners keep the lowest runtime times PPT watts. The accuracy gate is unchanged, and latency remains the default.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tuner] Satisfy black and ruff on the power-selection change
Format the tuner edits and replace the NaN self-comparisons and bare exception handlers that the pre-check jobs reject.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tuner] Keep latency order when package power is missing
Sort power candidates by energy and then latency so an unread sensor cannot drop the fastest kernel. A single AITER_POWER_PATH is rejected when more than one tuner process is running.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tuner] Keep power-tuned rows apart from latency tables
Power tunes of shipped tables now write to aiter/configs/power_configs/, which loads only with AITER_SELECT=power and replaces only the shapes it covers. Each power row records perf level, power cap, and SCLK range, and the config merge refuses to auto-resolve power rows against latency rows or rows tuned under different conditions. Power mode is marked experimental.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tuner] Isolate power-kernel selection from the tuner and JIT path
Ranking, package-power measurement, and power_configs loading each have one owner. Latency CSV columns stay unchanged unless --select power.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tuner] Measure power with the same call the timer uses
The replay was passing run_perftest keywords into the kernel, so watts came back NaN and the row was still stamped select=power. A batch with no finite watts stays a latency result, and --select power is rejected on tuners that do not record it.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tuner] Give the MXFP4 power retime test a supported shape
Main now rejects an MXFP4 row before the scan unless dtype and quant fields match the port. The energy-ranking fixture was missing those fields, so it never reached the retime.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tuner] Format the power-selection files with Black
CI skipped the GPU tests because Black rejected these three files.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tuner] Keep the wheel build from importing the aiter package
setup.py loads jit.core before aiter exists. Importing the power overlay from that path ran aiter/__init__.py while module_aiter_core was still being registered.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tuner] Decide latency or power in one policy and measure in one place
A tuner builds a SelectionPolicy from --select and asks it for the
candidate rank, frame order, pair energy, accurate-pass count, and the
power sample. Tuners no longer branch on the select mode or write
AITER_SELECT; that variable only picks the tables jit/core loads.
mp_tuner reads the policy TunerCommon.run activates.
package_power.measure_kernel is the one power read. mp_tuner hands it
the kernel launch with the run_perftest options removed, so the window
replays the kernel back to back. Replaying run_perftest itself left the
GPU idle for most of the window.
The duplicate-shape dedup in jit/core sorts by lowest us without
importing the aiter package when setup.py loads it as jit.core.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tuner] Fix power-tune paths found by end-to-end runs on MI355X
- TunerCommon.run wrapped its body in _run, which IQ2R already defines
as its kernel launcher, so every IQ2R tune failed. Use _run_selected.
- The a16w16 tuner built gradlib's Gemm without the policy, so hipBLASLt
was never power-measured and always ranked last. Pass the policy, and
default Gemm and the K5 sweep to the active policy.
- The power window queued launches without bound. IQ2R graph replays
crash inside the HIP runtime when hundreds are queued. Keep a few
events in flight, one per ~200 us of work, so the GPU never drains.
- IQ2R times a captured CUDA graph; sample power from that graph too.
- Staged MXFP4 logs each pair's watts like the other loops.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tuner] Correct the power sample for host launch gaps
A kernel whose Python launch takes longer than its GPU time leaves the
GPU idle between launches, so the window read close to idle and the
kernel looked efficient (split-K FlyDSL GEMMs at 256x4096x4096 read
~370 W against ~990 W replayed from a CUDA graph, and won power mode).
Every measurement site now passes the call's timed us. When the loop
never waited on the GPU, the window's busy fraction is launches * us /
window, and the reading is scaled to active power over the between-launch
baseline (a one-element kernel loop, ~300 W on MI355X). A window that
waited on the GPU kept it saturated and stands as read. Under 20% busy
the candidate stays unmeasured.
Co-authored-by: Cursor <cursoragent@cursor.com>
* [Tuner] Keep latency profiles as main and write only measured power rows
- MoE latency tunes write the -o2 profile and profile_fmoe.csv with
main's columns and row order. Stage watts, rank, energy, and
tflops_per_w appear only with --select power.
- A power tune drops rows whose watts were not measured instead of
writ…
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.
Summary
Use a shared typed FlyDSL stream argument to avoid constructing a Python
torch.cuda.Streamjust to recover its handle. Query current device and stream every invocation; no stream/pointer/tensor caching.The initial revision covered grouped GEMM only. This update also covers six audited MoE wrapper paths: ordinary/EP route rows, fused small-route prefix scan, contiguous prefix-sum remap, fused quant/preshuffle, scale scatter/preshuffle and gather/reduce. Only stream-argument construction changes; kernel math, routing policy, launch order and specialization keys are unchanged.
current_stream_arg()intensor_shim.py; raw API availability resolved once._cuda_getCurrentRawStreamkeeps the publictorch.cuda.current_stream()fallback.stream=is forwarded unchanged.fx.Streamretained so integers are not accidentally specialized as scalar kernel parameters.Validation of the expanded revision
Incremental performance beyond the original grouped-GEMM-only revision
Frozen reference: initial PR head
2fd43d5845f5a965200b8d9807a454a24a1265d1. The reference already includes the grouped-GEMM stream optimization. The table is the additional benefit of extending the same stream handling to the six other MoE helpers.gfx1250, M3 TP4 per-rank operator shape
E129 / top5 / K6144 / I768, SwiGLU. Real routing, quantization, GEMMs and output reduction; same inputs/kernels. One GPU, not distributed TP4.Final release-file run: 31 randomized paired samples ×150 eager calls; 100 warmups per arm; synchronize before and after each sample. WARNING log level. No profiler in timing. Old control functions are extracted unchanged from the frozen original source; candidate functions come from the final PR working tree.
Independent prior processes also show positive incremental savings (roughly 11–18µs). All raw paired samples and completion timings retained. These are host-limited eager-call savings, not faster GPU kernels; graph-only replay need not improve.
Model limitation / negative result retained
The earlier full-model ABBA combining #7217 + the original grouped-GEMM-only version of this PR did not establish a robust end-to-end speedup: short latency +1.11%, long-prefix latency −2.67%, C2 throughput −0.04%, with substantial run-to-run spread. Full raw table: #7219 (comment) . That result is not superseded by operator measurements. The expanded revision needs its own model comparison; no whole-model gain is claimed.
Known pre-existing device limitation
Launching one already-compiled FlyDSL MoE executable on GPU0 then GPU1 in a single process fails with
hipErrorInvalidDevicein both the original and candidate. Separate processes starting on GPU1 pass. This separate runtime executable/device-ownership issue is not fixed here. Numerical tests use one device per process; the multi-device stream-handle test itself does not reuse a compiled kernel across devices.Artifacts: f01
~/bench/m3_raw_stream_extended_20261030/(frozen reference7219, final tests/JUnit,logs/bench_release.json, independent repeats, kernel identity and EP negative controls); original PR evidence remains under~/bench/m3_raw_stream_20261030/.