Repository navigation
[Triton/Gluon] Gfx1250 flash kda gluon - #6609
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
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Positional calls still pack projection views, and new fallback tuning bypasses the required JSON configuration.
2 open findings
What changed in this PR
Ports FlashKDA to gfx1250 Gluon kernels, adding architecture-specific acceleration within the existing Kimi attention pipeline.
Changes:
- Adds TDM loads, wave32 WMMA computation, and per-pass K2 schedules.
- Enables supported strided inputs without packing.
- Extends layout, routing, and launch-cache tests.
| File | Description |
|---|---|
op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py |
Tests routing, schedules, and strided inputs. |
op_tests/triton_tests/kimi_delta_attn/test_fast_launch.py |
Selects architecture-specific launchers. |
aiter/ops/triton/kimi_delta_attn/chunk_delta_attn.py |
Defers operand packing to pipeline dispatch. |
aiter/ops/triton/configs/gfx1250/gluon/attention/chunk_delta_attn/DEFAULT.json |
Publishes K2 schedules. |
aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py |
Adds gfx1250 dispatch and stride handling. |
aiter/ops/triton/_triton_kernels/kimi_delta_attn/chunk_fwd.py |
Preserves packing for the default pipeline. |
aiter/ops/triton/_gluon_kernels/gfx1250/kimi_delta_attn/flash_kda_k2.py |
Implements chunk recurrence and output passes. |
aiter/ops/triton/_gluon_kernels/gfx1250/kimi_delta_attn/flash_kda_k1.py |
Implements workspace preparation. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
amd-ruitang3
force-pushed
the
gfx1250-flash-kda-gluon
branch
2 times, most recently
from
October 16, 2026 14:11
1ab7278 to
de7eff6
Compare
Liang-jianhao97
force-pushed
the
gfx1250-flash-kda-gluon
branch
2 times, most recently
from
October 18, 2026 12:39
c9482a9 to
8865724
Compare
Port the Gluon FlashKDA kernels to gfx1250, in the gfx950 layout, and route flash_kda_fwd to them there. - _gluon_kernels/gfx1250/kimi_delta_attn/flash_kda_k1.py: the prepare kernel. q, k and g arrive by one fused TDM load, and the [C, C] products and the (I - L)^-1 solve run on wave32 WMMA. Same workspace contract and the same gluon_k1_prepare entry point as the gfx950 K1. - _gluon_kernels/gfx1250/kimi_delta_attn/flash_kda_k2.py: the chunk walk. Pass A computes b_seg and A_seg in one launch; pass C computes the outputs and the final state, paged state_cache included. The state stays transposed in a WMMA accumulator, the workspace tiles and v arrive by fused TDM loads into a double buffer, and the outputs leave LDS by async copy. - flash_kda.py: on gfx1250, K1 and both K2 passes go to Gluon. Routing on gfx950 is unchanged; pass C stays on Triton there. _k2_gluon_schedule now takes the pass name, so each pass has its own wide and narrow entries. - configs/gfx1250/gluon/attention/chunk_delta_attn/DEFAULT.json: BW 128 and BW 64 at four warps. The wide floors were calibrated on the bench_chunk_kda and bench_flash_kda shapes: 2 blocks per CU for pass A, 0.3 for pass C. - test_flash_kda.py: the schedule and routing tests also cover pass C on archs where it has a Gluon kernel, and import that arch's K2 module. Tests on gfx1250: op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py, 76 passed, 1 skipped (this device has no K2 autotune candidates to switch between). Performance on gfx1250 (one GPU, 256 CUs), upstream 428af1c -> this commit, both runs with the same environment. bench_chunk_kda.py, the FlashKDA line ("triton"), us: H B T do_bench cudagraph 24 1 1024 178.8 -> 62.3 2.87x 176.9 -> 37.1 4.77x 24 1 4096 486.0 -> 99.7 4.87x 487.9 -> 97.8 4.99x 24 3 1024 190.8 -> 67.0 2.85x 184.8 -> 54.6 3.38x 24 3 4096 448.1 -> 170.1 2.63x 449.7 -> 168.7 2.67x 24 8 1024 295.1 -> 102.7 2.87x 302.2 -> 101.1 2.99x 24 8 4096 925.2 -> 340.5 2.72x 924.1 -> 343.2 2.69x 12 1 1024 171.4 -> 61.9 2.77x 199.4 -> 34.6 5.76x 12 1 4096 410.8 -> 106.0 3.87x 398.5 -> 98.7 4.04x 12 3 1024 220.4 -> 65.9 3.34x 170.5 -> 40.9 4.17x 12 3 4096 561.2 -> 126.1 4.45x 546.9 -> 134.1 4.08x 12 8 1024 207.6 -> 69.8 2.97x 205.6 -> 65.1 3.16x 12 8 4096 542.5 -> 205.7 2.64x 534.6 -> 207.3 2.58x The FlashKDA line is now faster than the chunk_kda Gluon line ("gluon") on every shape in both timings. Upstream, it was slower on every shape. bench_flash_kda.py, ms: B T H upstream this speedup 1 256 12 0.0479 0.0417 1.15x 1 512 12 0.0753 0.0396 1.90x 1 1664 12 0.1119 0.0418 2.68x 1 2048 12 0.3429 0.0700 4.90x 1 8192 12 0.4866 0.1050 4.63x 1 16384 12 0.6712 0.1489 4.51x 1 16384 64 2.0335 0.5721 3.55x 1 16384 96 1.7003 0.7673 2.22x 1 8192 32 0.7352 0.1615 4.55x 1 8192 64 1.2016 0.3074 3.91x 2 2664 12 0.3962 0.0839 4.72x 4 4096 16 0.7366 0.1587 4.64x 8 2048 16 0.3784 0.1278 2.96x Tests on gfx1250: test_flash_kda.py and test_fast_launch.py, 107 passed, 2 skipped; test_chunk_delta_attn_fwd.py, 39 passed;
Liang-jianhao97
force-pushed
the
gfx1250-flash-kda-gluon
branch
from
October 18, 2026 12:47
8865724 to
b1b335d
Compare
omuhamma
reviewed
Oct 22, 2026
added 2 commits
October 23, 2026 12:20
…JSON Address the review of the gfx1250 Gluon FlashKDA kernels: share the Gluon helpers, and read the K1 and K2 launch values from the published config. - _gluon_kernels/common/utils.py: sigmoid_log2, fast_sigmoid, sum_combine and l2norm, in the form that the FlashKDA kernels used. The sigmoid that is already in this file uses tanh and gives different results. - _gluon_kernels/gfx1250/kimi_delta_attn/flash_kda_k1.py: import these helpers and call gl.exp2 directly. gluon_k1_prepare takes num_warps and waves_per_eu; the kernel still requires four warps. - _gluon_kernels/gfx1250/kimi_delta_attn/flash_kda_k2.py: import fast_sigmoid. Both launchers take CPL, the number of chunks in one TDM load, in place of the module constant _CPL. - flash_kda.py: _k1_gluon_schedule reads the k1_prepare_gluon entry, and _k2_gluon_schedule returns the whole entry, so pass A and pass C also get CPL. The fallbacks do not depend on the arch. - configs/gfx1250/gluon/attention/chunk_delta_attn/DEFAULT.json: a k1_prepare_gluon entry with 4 warps and waves_per_eu 5, and CPL 2 in the four K2 entries. These are the values that were in Python. - test_flash_kda.py: read BW and num_warps from the returned config. The outputs, the final states and the ISA of all K1 and K2 variants do not change. Tests on gfx1250: test_flash_kda.py and test_fast_launch.py, 108 passed, 1 skipped (this device has no K2 autotune candidates to switch between).
Pass B, the cross-segment scan, is a chain of dependent updates, and _flash_kda_seg_scan_kernel loads b_seg in the update that adds it: two exposed memory round trips per update. _flash_kda_seg_scan_prefetch_kernel does two updates per loop trip and loads the next trip's b_seg, two to three updates ahead, while the software pipeliner prefetches A_seg. It also skips the last segment's update, which nothing reads. It is pass B on gfx1250 (AITER_FDA_SCAN_PREFETCH=0 goes back to the old kernel) at BV=16, 4 warps and the default 2 stages; other archs keep the old kernel and its schedule. On H=12 forwards of 1 to 3 sequences, 13 to 20 segments deep, pass B is 1.8-2.5x faster and the forward's kernel time 6-17% lower. h_in matches the old kernel to 2.4e-7 with dense, no and paged initial states.
omuhamma
reviewed
Oct 23, 2026
omuhamma
left a comment
omuhamma
left a comment
Contributor
There was a problem hiding this comment.
Please add a benchmark file in op_tests/op_benchmarks so others can repro your results
1 task
Micky774
added a commit
that referenced
this pull request
Oct 23, 2026
…per with #6609 - Move sigmoid_log2/fast_sigmoid/sum_combine/l2norm into _gluon_kernels/common/utils.py (same definitions as #6609); k1 uses them. - Align flash_kda.py with #6609: _k2_gluon_schedule(..., kernel) returns a Config, use_gluon_k2a/k2c split, _GLUON_K2C_ARCH_OK, unpaged_args; gfx950 k2 gets gluon_k2_ab_fused/gluon_k2_c launchers mirroring gfx1250's. - Tests: arch-generic Gluon module imports, pass-C checks gated on arch. - Fix remap_xcd import path (moved to _triton_kernels/common in #6505). - Trim verbose comments. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Time FlashKDA on the calls that Kimi-K3 makes, and chunk_kda on the same inputs. - --k3-tp8: the 26 prefill forwards of Kimi-K3 TP8+DCP8 rank 0 (H=12), in order. Each forward is one varlen call. A sequence that continues a cached prefix reads its initial state from a paged state cache; the others start from zero. The table and the CSV gain a seqlens column. - --qkv-bands: pass q, k and v as views of one fused [B, T, 3 * H * 128] projection, as the model does, instead of three packed tensors. - --impl chunk_kda: time the Gluon chunk_kda on the inputs that FlashKDA gets; a batch becomes one varlen call. --seg-sweep needs --impl flash_kda. - --shape can repeat. It and --k3-tp8 exclude each other. The default sweep and its inputs do not change. Run on gfx1250 with --k3-tp8, --k3-tp8 --impl chunk_kda, --qkv-bands with repeated --shape, and the default sweep.
Contributor
Author
@omuhamma Added in 5cf2e0f. op_tests/op_benchmarks/triton/bench_flash_kda.py now takes --k3-tp8 (the 26 prefill forwards of Kimi-K3 TP8+DCP8 rank 0), --qkv-bands (q, k and v as bands of one fused projection) and --impl chunk_kda (chunk_kda on the same inputs). Every timing in the description now comes from bench_chunk_kda.py or bench_flash_kda.py, and each table gives its command. |
added 2 commits
October 26, 2026 03:08
bench_flash_kda.py conflicted: main builds the parser with get_parser (#6347), next to the imports and the --shape/--k3-tp8 group this branch added. Keep both.
get_num_xcds (#6327) now imports aiter.jit.utils.chip_info on every call. Under AITER_TRITON_ONLY=1, which bench_flash_kda.py sets, that import fails and is retried each time, about 2 ms. Pass A and pass C call it on every launch, so after the merge of main bench_flash_kda.py read 2-4 ms for every FlashKDA call. Cache the count, as build_layouts already is.
Dewei-Wang-sh
approved these changes
Oct 26, 2026
eppaneamd
added a commit
that referenced
this pull request
Oct 27, 2026
…ens-cpu Conflicts with #6609 (gfx1250 FlashKDA, ffddbaf), resolved by keeping both sides: - chunk_fwd.py: main packs q, k, v and g after the FlashKDA return; the default pipeline then builds chunk_indices from cu_seqlens_cpu. - chunk_delta_attn.py: cu_seqlens_cpu is accepted (no longer listed as rejected); main's contiguity note, with the gfx1250 exception, is kept. - test_flash_kda.py: both test blocks kept. #6609's test_positional_call_keeps_skipped_storage lists every flash_kda_fwd parameter, so it gains cu_seqlens_cpu. - bench_flash_kda.py: --varlen's column is named "bounds", since --k3-tp8's "seqlens" column holds the lengths. --varlen rejects --k3-tp8 (its inputs carry their own cu_seqlens) and --impl chunk_kda (it would time flash_kda_fwd). The gfx1250 K1 and K2 read chunk_indices and the segment descriptors as the gfx950 kernels do, so they take the int32 metadata from _build_metadata unchanged. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Rony Leppänen <rony.leppanen@amd.com>
azaidy
pushed a commit
that referenced
this pull request
Oct 27, 2026
…on gfx950 (#6177) * [Triton/Gluon] FlashKDA K2: keep a (segment, head)'s V blocks on one XCD The K2 grid is (cdiv(W, BW), num_segs * H), and every V block of one (segment, head) re-reads that segment's whole chunk workspace. Workgroups reach the XCDs round-robin in launch order, x fastest, so with the tuned BW=16 the eight V blocks of each (segment, head) landed on eight different XCDs and each pulled the workspace through its own L2. Measured on MI355X (rocprofv3 FETCH_SIZE/WRITE_SIZE): passes A and C moved 6-8x their compulsory bytes past L2, pinned at ~7 TB/s on every shape -- they were bound by L2-miss traffic, not by the recurrence. Renumbering the flat program id with the existing remap_xcd, so each XCD owns a contiguous run of whole (segment, head) groups, brings that traffic to 1.01-1.02x compulsory and halves those kernels' time at an unchanged grid and tile. Applied to the Triton segment kernel (passes A and C) and the Gluon pass A, whose BW=32/64 tiles split V the same way. The remap only changes which program computes which tile, so results are bit-identical. FlashKDA end to end, MI355X, bf16, K=V=128: gfx950 default route (Gluon K1 + pass A, Triton pass C): 1.24-1.45x Triton only (non-gfx950 route): 1.32-1.61x 1x512x12: 1.01x * [Triton/Gluon] KDA: load token i once in intra_token_parallel The token-parallel intra kernel (the default pipeline's path when safe_gate is off) reloaded token i's q/k/g for every j and every K block. Load them once outside the j loop, span all of K in one block (BK from K, no longer tuned) and fold beta into k, following fla. The default config moves from BH=1 / 4 warps, which idles most of each warp, to BH=4 / 1 warp: within 1.24x of the per-shape best on gfx950. Default pipeline forward with safe_gate=False on MI355X: 2.0-2.7x (1x8192x32 2862 -> 1221 us). Aqk/Akk change only in summation order (output within 2.2e-5 relative). * [Triton/Gluon] FlashKDA: cap segment length at 64 chunks _choose_chunks_per_seg capped the segment count at 16 over both of its reasons, so a long few-head sequence got 256-chunk segments whose serial walk bounds the recurrence. Cap the occupancy target alone at 16 segments and then cut segments to at most 64 chunks however many that takes. On gfx950, 64-chunk segments beat longer ones by up to 1.9x on 1x131072x4. * [Triton/Gluon] FlashKDA: LDS-prefetching Gluon K2 and a Gluon pass C Speeds up the Gluon FlashKDA kernels on gfx950 and moves pass C to Gluon: - K1: sigmoids through exp2 (one v_exp_f32 instead of exp's denormal-safe sequence); the gate in a column-per-thread layout, so its cumulative sum is in-thread, with tail rows given a zero gate (the last and pivot rows are then fixed rows); tail-row masks on the workspace dropped; warps split the rows of the [C, K] tiles, so the l2norm sums stay in a warp; operands staged through LDS once per product (10 barriers instead of 38); the gate load is masked, since buffer loads past the tensor's end are not clamped. Launches with waves_per_eu=3 (170 -> 168 VGPRs, 3 waves/SIMD). - K1 (Gluon and Triton) writes sigmoid(beta) to a new ws_beta workspace, which K2 copies into LDS with the chunk's other tiles. - Pass A: chunk j+1's tiles are async-copied into a double-buffered LDS stage while chunk j computes; one barrier per chunk. The kernel needs ~280 VGPRs, so the published schedules carry waves_per_eu=2. - Pass C: new k2_c_gluon, used for every shape on gfx950. Same prefetch; the output projection runs before the state update and is stored from the MFMA layout; seg_seq / seg_is_last are read before the loop. It supports the paged state_cache read and write-back of the Triton pass C. BW=64 unsegmented, BW=128 segmented. FlashKDA forward on MI355X vs main: 1.43-3.02x (1x8192x12 279 -> 142 us, 1x8192x96 1458 -> 690 us, 1x131072x4 1634 -> 541 us, varlen 1.70-1.99x), including the two commits before this one. Output within 4.5e-4 relative of main (K1's l2norm and gate summation order), final state within 3.5e-5. Tests: test_cases_reach_the_gluon_k2 now checks pass C as well as pass A, the paged-cache test runs on both K2 routes, and the schedule test covers the pass C tiles. * [Triton/Gluon] FlashKDA: skip ws_beta off the Gluon K2 route; name K2 kernels by config - ws_beta is only read by the Gluon K2, so it is now allocated only when K2 takes that route, and both K1s store it only then (STORE_BETA). The K2 route is resolved before K1 launches for this. - The Gluon pass A and pass C names now carry their warp count, through a NUM_WARPS constexpr that must equal the launch's num_warps (num_warps itself is not a constexpr and would name as NONE), and pass C's also carries STATE_V_FIRST. - test_fast_launch covers the pass C launcher too. * [Triton/Gluon] FlashKDA: keep small unsegmented pass C on Triton; bench without safe_gate - On unsegmented shapes with few (sequence, head) pairs, the serial walk over the chunks sets pass C's time, and the Triton kernel's lighter iteration at BW=16 beats the Gluon one, which cannot go narrower than BW=32: 1.08-1.18x for Triton at <= 32 pairs on MI355X, while Gluon wins from 40 pairs up (to 2.2x). The gfx950 k2_c_gluon schedule now publishes a MIN_BLOCKS_PER_CU floor, as pass A's does, below which pass C runs on Triton; ws_beta is then skipped too. Segmented shapes keep the Gluon pass C. - bench_chunk_delta_attn.py gains --no-safe-gate, which takes the token-parallel intra kernel. bench_flash_kda.py default sweep vs main on MI355X: 1.04-2.61x (1x256x12 and 1x1664x12 were 0.91x and 0.94x before this change). * [Triton/Gluon] FlashKDA: publish the Gluon K1 occupancy hint in the gfx950 config K1's waves_per_eu=3 was pinned in Python while pass A's occupancy hints are published in configs/gfx950/gluon/.../chunk_delta_attn/DEFAULT.json. Move it there (k1_prepare_gluon) and load it through chunk_delta_attn_tuned_config, so retuning it is a config change. The value holds K1 to ~168 VGPRs, i.e. 3 waves per SIMD; left alone it lands at 170 and drops to 2 (0.85x). The warp count stays fixed by K1's layouts: the published entry restates it, and an entry naming another count is rejected rather than launched against them. Devices without an entry get no hint. Compiled K1 on gfx950 is unchanged (waves_per_eu=3, 164 VGPRs). * [Triton/Gluon] KDA: address review — shared Gluon helpers, align wrapper with #6609 - Move sigmoid_log2/fast_sigmoid/sum_combine/l2norm into _gluon_kernels/common/utils.py (same definitions as #6609); k1 uses them. - Align flash_kda.py with #6609: _k2_gluon_schedule(..., kernel) returns a Config, use_gluon_k2a/k2c split, _GLUON_K2C_ARCH_OK, unpaged_args; gfx950 k2 gets gluon_k2_ab_fused/gluon_k2_c launchers mirroring gfx1250's. - Tests: arch-generic Gluon module imports, pass-C checks gated on arch. - Fix remap_xcd import path (moved to _triton_kernels/common in #6505). - Trim verbose comments. * [Triton/Gluon] KDA: address review — published schedules, resolved-route log, launch-level route tests - K1 and pass C read their Gluon schedules through _published_gluon_schedule and fall back to Triton when none is published (no Python fallback). The K1 schedule moves to the wrapper; gluon_k1_prepare rejects a wrong warp count. - Segment cap moves to the gfx950 Triton config (flash_kda_segments: MAX_CHUNKS=64, HARD_CAP). Other archs keep the previous 32-chunk soft cap. - intra_attn token-parallel default (BH=4, 1 warp) moves to the gfx950 Triton config; other archs keep the previous BH=1, 4-warp default. - Use get_num_sms() (honors CU_NUM) for routing and segmentation. - Log the resolved K1, pass A and pass C routes after pass C is settled. - Pass A and pass C share _solve_u/_update (AMDGCN unchanged). - Tests: count actual pass-C launches for size routing, cover the segment cap and unpublished-schedule fallback, test K1 warp validation via flash_kda_fwd.
vpietila-amd
pushed a commit
that referenced
this pull request
Oct 28, 2026
…on gfx950 (#6177) * [Triton/Gluon] FlashKDA K2: keep a (segment, head)'s V blocks on one XCD The K2 grid is (cdiv(W, BW), num_segs * H), and every V block of one (segment, head) re-reads that segment's whole chunk workspace. Workgroups reach the XCDs round-robin in launch order, x fastest, so with the tuned BW=16 the eight V blocks of each (segment, head) landed on eight different XCDs and each pulled the workspace through its own L2. Measured on MI355X (rocprofv3 FETCH_SIZE/WRITE_SIZE): passes A and C moved 6-8x their compulsory bytes past L2, pinned at ~7 TB/s on every shape -- they were bound by L2-miss traffic, not by the recurrence. Renumbering the flat program id with the existing remap_xcd, so each XCD owns a contiguous run of whole (segment, head) groups, brings that traffic to 1.01-1.02x compulsory and halves those kernels' time at an unchanged grid and tile. Applied to the Triton segment kernel (passes A and C) and the Gluon pass A, whose BW=32/64 tiles split V the same way. The remap only changes which program computes which tile, so results are bit-identical. FlashKDA end to end, MI355X, bf16, K=V=128: gfx950 default route (Gluon K1 + pass A, Triton pass C): 1.24-1.45x Triton only (non-gfx950 route): 1.32-1.61x 1x512x12: 1.01x * [Triton/Gluon] KDA: load token i once in intra_token_parallel The token-parallel intra kernel (the default pipeline's path when safe_gate is off) reloaded token i's q/k/g for every j and every K block. Load them once outside the j loop, span all of K in one block (BK from K, no longer tuned) and fold beta into k, following fla. The default config moves from BH=1 / 4 warps, which idles most of each warp, to BH=4 / 1 warp: within 1.24x of the per-shape best on gfx950. Default pipeline forward with safe_gate=False on MI355X: 2.0-2.7x (1x8192x32 2862 -> 1221 us). Aqk/Akk change only in summation order (output within 2.2e-5 relative). * [Triton/Gluon] FlashKDA: cap segment length at 64 chunks _choose_chunks_per_seg capped the segment count at 16 over both of its reasons, so a long few-head sequence got 256-chunk segments whose serial walk bounds the recurrence. Cap the occupancy target alone at 16 segments and then cut segments to at most 64 chunks however many that takes. On gfx950, 64-chunk segments beat longer ones by up to 1.9x on 1x131072x4. * [Triton/Gluon] FlashKDA: LDS-prefetching Gluon K2 and a Gluon pass C Speeds up the Gluon FlashKDA kernels on gfx950 and moves pass C to Gluon: - K1: sigmoids through exp2 (one v_exp_f32 instead of exp's denormal-safe sequence); the gate in a column-per-thread layout, so its cumulative sum is in-thread, with tail rows given a zero gate (the last and pivot rows are then fixed rows); tail-row masks on the workspace dropped; warps split the rows of the [C, K] tiles, so the l2norm sums stay in a warp; operands staged through LDS once per product (10 barriers instead of 38); the gate load is masked, since buffer loads past the tensor's end are not clamped. Launches with waves_per_eu=3 (170 -> 168 VGPRs, 3 waves/SIMD). - K1 (Gluon and Triton) writes sigmoid(beta) to a new ws_beta workspace, which K2 copies into LDS with the chunk's other tiles. - Pass A: chunk j+1's tiles are async-copied into a double-buffered LDS stage while chunk j computes; one barrier per chunk. The kernel needs ~280 VGPRs, so the published schedules carry waves_per_eu=2. - Pass C: new k2_c_gluon, used for every shape on gfx950. Same prefetch; the output projection runs before the state update and is stored from the MFMA layout; seg_seq / seg_is_last are read before the loop. It supports the paged state_cache read and write-back of the Triton pass C. BW=64 unsegmented, BW=128 segmented. FlashKDA forward on MI355X vs main: 1.43-3.02x (1x8192x12 279 -> 142 us, 1x8192x96 1458 -> 690 us, 1x131072x4 1634 -> 541 us, varlen 1.70-1.99x), including the two commits before this one. Output within 4.5e-4 relative of main (K1's l2norm and gate summation order), final state within 3.5e-5. Tests: test_cases_reach_the_gluon_k2 now checks pass C as well as pass A, the paged-cache test runs on both K2 routes, and the schedule test covers the pass C tiles. * [Triton/Gluon] FlashKDA: skip ws_beta off the Gluon K2 route; name K2 kernels by config - ws_beta is only read by the Gluon K2, so it is now allocated only when K2 takes that route, and both K1s store it only then (STORE_BETA). The K2 route is resolved before K1 launches for this. - The Gluon pass A and pass C names now carry their warp count, through a NUM_WARPS constexpr that must equal the launch's num_warps (num_warps itself is not a constexpr and would name as NONE), and pass C's also carries STATE_V_FIRST. - test_fast_launch covers the pass C launcher too. * [Triton/Gluon] FlashKDA: keep small unsegmented pass C on Triton; bench without safe_gate - On unsegmented shapes with few (sequence, head) pairs, the serial walk over the chunks sets pass C's time, and the Triton kernel's lighter iteration at BW=16 beats the Gluon one, which cannot go narrower than BW=32: 1.08-1.18x for Triton at <= 32 pairs on MI355X, while Gluon wins from 40 pairs up (to 2.2x). The gfx950 k2_c_gluon schedule now publishes a MIN_BLOCKS_PER_CU floor, as pass A's does, below which pass C runs on Triton; ws_beta is then skipped too. Segmented shapes keep the Gluon pass C. - bench_chunk_delta_attn.py gains --no-safe-gate, which takes the token-parallel intra kernel. bench_flash_kda.py default sweep vs main on MI355X: 1.04-2.61x (1x256x12 and 1x1664x12 were 0.91x and 0.94x before this change). * [Triton/Gluon] FlashKDA: publish the Gluon K1 occupancy hint in the gfx950 config K1's waves_per_eu=3 was pinned in Python while pass A's occupancy hints are published in configs/gfx950/gluon/.../chunk_delta_attn/DEFAULT.json. Move it there (k1_prepare_gluon) and load it through chunk_delta_attn_tuned_config, so retuning it is a config change. The value holds K1 to ~168 VGPRs, i.e. 3 waves per SIMD; left alone it lands at 170 and drops to 2 (0.85x). The warp count stays fixed by K1's layouts: the published entry restates it, and an entry naming another count is rejected rather than launched against them. Devices without an entry get no hint. Compiled K1 on gfx950 is unchanged (waves_per_eu=3, 164 VGPRs). * [Triton/Gluon] KDA: address review — shared Gluon helpers, align wrapper with #6609 - Move sigmoid_log2/fast_sigmoid/sum_combine/l2norm into _gluon_kernels/common/utils.py (same definitions as #6609); k1 uses them. - Align flash_kda.py with #6609: _k2_gluon_schedule(..., kernel) returns a Config, use_gluon_k2a/k2c split, _GLUON_K2C_ARCH_OK, unpaged_args; gfx950 k2 gets gluon_k2_ab_fused/gluon_k2_c launchers mirroring gfx1250's. - Tests: arch-generic Gluon module imports, pass-C checks gated on arch. - Fix remap_xcd import path (moved to _triton_kernels/common in #6505). - Trim verbose comments. * [Triton/Gluon] KDA: address review — published schedules, resolved-route log, launch-level route tests - K1 and pass C read their Gluon schedules through _published_gluon_schedule and fall back to Triton when none is published (no Python fallback). The K1 schedule moves to the wrapper; gluon_k1_prepare rejects a wrong warp count. - Segment cap moves to the gfx950 Triton config (flash_kda_segments: MAX_CHUNKS=64, HARD_CAP). Other archs keep the previous 32-chunk soft cap. - intra_attn token-parallel default (BH=4, 1 warp) moves to the gfx950 Triton config; other archs keep the previous BH=1, 4-warp default. - Use get_num_sms() (honors CU_NUM) for routing and segmentation. - Log the resolved K1, pass A and pass C routes after pass C is settled. - Pass A and pass C share _solve_u/_update (AMDGCN unchanged). - Tests: count actual pass-C launches for size routing, cover the segment cap and unpublished-schedule fallback, test K1 warp validation via flash_kda_fwd.
solinzby1
pushed a commit
that referenced
this pull request
Nov 2, 2026
…on gfx950 (#6177) * [Triton/Gluon] FlashKDA K2: keep a (segment, head)'s V blocks on one XCD The K2 grid is (cdiv(W, BW), num_segs * H), and every V block of one (segment, head) re-reads that segment's whole chunk workspace. Workgroups reach the XCDs round-robin in launch order, x fastest, so with the tuned BW=16 the eight V blocks of each (segment, head) landed on eight different XCDs and each pulled the workspace through its own L2. Measured on MI355X (rocprofv3 FETCH_SIZE/WRITE_SIZE): passes A and C moved 6-8x their compulsory bytes past L2, pinned at ~7 TB/s on every shape -- they were bound by L2-miss traffic, not by the recurrence. Renumbering the flat program id with the existing remap_xcd, so each XCD owns a contiguous run of whole (segment, head) groups, brings that traffic to 1.01-1.02x compulsory and halves those kernels' time at an unchanged grid and tile. Applied to the Triton segment kernel (passes A and C) and the Gluon pass A, whose BW=32/64 tiles split V the same way. The remap only changes which program computes which tile, so results are bit-identical. FlashKDA end to end, MI355X, bf16, K=V=128: gfx950 default route (Gluon K1 + pass A, Triton pass C): 1.24-1.45x Triton only (non-gfx950 route): 1.32-1.61x 1x512x12: 1.01x * [Triton/Gluon] KDA: load token i once in intra_token_parallel The token-parallel intra kernel (the default pipeline's path when safe_gate is off) reloaded token i's q/k/g for every j and every K block. Load them once outside the j loop, span all of K in one block (BK from K, no longer tuned) and fold beta into k, following fla. The default config moves from BH=1 / 4 warps, which idles most of each warp, to BH=4 / 1 warp: within 1.24x of the per-shape best on gfx950. Default pipeline forward with safe_gate=False on MI355X: 2.0-2.7x (1x8192x32 2862 -> 1221 us). Aqk/Akk change only in summation order (output within 2.2e-5 relative). * [Triton/Gluon] FlashKDA: cap segment length at 64 chunks _choose_chunks_per_seg capped the segment count at 16 over both of its reasons, so a long few-head sequence got 256-chunk segments whose serial walk bounds the recurrence. Cap the occupancy target alone at 16 segments and then cut segments to at most 64 chunks however many that takes. On gfx950, 64-chunk segments beat longer ones by up to 1.9x on 1x131072x4. * [Triton/Gluon] FlashKDA: LDS-prefetching Gluon K2 and a Gluon pass C Speeds up the Gluon FlashKDA kernels on gfx950 and moves pass C to Gluon: - K1: sigmoids through exp2 (one v_exp_f32 instead of exp's denormal-safe sequence); the gate in a column-per-thread layout, so its cumulative sum is in-thread, with tail rows given a zero gate (the last and pivot rows are then fixed rows); tail-row masks on the workspace dropped; warps split the rows of the [C, K] tiles, so the l2norm sums stay in a warp; operands staged through LDS once per product (10 barriers instead of 38); the gate load is masked, since buffer loads past the tensor's end are not clamped. Launches with waves_per_eu=3 (170 -> 168 VGPRs, 3 waves/SIMD). - K1 (Gluon and Triton) writes sigmoid(beta) to a new ws_beta workspace, which K2 copies into LDS with the chunk's other tiles. - Pass A: chunk j+1's tiles are async-copied into a double-buffered LDS stage while chunk j computes; one barrier per chunk. The kernel needs ~280 VGPRs, so the published schedules carry waves_per_eu=2. - Pass C: new k2_c_gluon, used for every shape on gfx950. Same prefetch; the output projection runs before the state update and is stored from the MFMA layout; seg_seq / seg_is_last are read before the loop. It supports the paged state_cache read and write-back of the Triton pass C. BW=64 unsegmented, BW=128 segmented. FlashKDA forward on MI355X vs main: 1.43-3.02x (1x8192x12 279 -> 142 us, 1x8192x96 1458 -> 690 us, 1x131072x4 1634 -> 541 us, varlen 1.70-1.99x), including the two commits before this one. Output within 4.5e-4 relative of main (K1's l2norm and gate summation order), final state within 3.5e-5. Tests: test_cases_reach_the_gluon_k2 now checks pass C as well as pass A, the paged-cache test runs on both K2 routes, and the schedule test covers the pass C tiles. * [Triton/Gluon] FlashKDA: skip ws_beta off the Gluon K2 route; name K2 kernels by config - ws_beta is only read by the Gluon K2, so it is now allocated only when K2 takes that route, and both K1s store it only then (STORE_BETA). The K2 route is resolved before K1 launches for this. - The Gluon pass A and pass C names now carry their warp count, through a NUM_WARPS constexpr that must equal the launch's num_warps (num_warps itself is not a constexpr and would name as NONE), and pass C's also carries STATE_V_FIRST. - test_fast_launch covers the pass C launcher too. * [Triton/Gluon] FlashKDA: keep small unsegmented pass C on Triton; bench without safe_gate - On unsegmented shapes with few (sequence, head) pairs, the serial walk over the chunks sets pass C's time, and the Triton kernel's lighter iteration at BW=16 beats the Gluon one, which cannot go narrower than BW=32: 1.08-1.18x for Triton at <= 32 pairs on MI355X, while Gluon wins from 40 pairs up (to 2.2x). The gfx950 k2_c_gluon schedule now publishes a MIN_BLOCKS_PER_CU floor, as pass A's does, below which pass C runs on Triton; ws_beta is then skipped too. Segmented shapes keep the Gluon pass C. - bench_chunk_delta_attn.py gains --no-safe-gate, which takes the token-parallel intra kernel. bench_flash_kda.py default sweep vs main on MI355X: 1.04-2.61x (1x256x12 and 1x1664x12 were 0.91x and 0.94x before this change). * [Triton/Gluon] FlashKDA: publish the Gluon K1 occupancy hint in the gfx950 config K1's waves_per_eu=3 was pinned in Python while pass A's occupancy hints are published in configs/gfx950/gluon/.../chunk_delta_attn/DEFAULT.json. Move it there (k1_prepare_gluon) and load it through chunk_delta_attn_tuned_config, so retuning it is a config change. The value holds K1 to ~168 VGPRs, i.e. 3 waves per SIMD; left alone it lands at 170 and drops to 2 (0.85x). The warp count stays fixed by K1's layouts: the published entry restates it, and an entry naming another count is rejected rather than launched against them. Devices without an entry get no hint. Compiled K1 on gfx950 is unchanged (waves_per_eu=3, 164 VGPRs). * [Triton/Gluon] KDA: address review — shared Gluon helpers, align wrapper with #6609 - Move sigmoid_log2/fast_sigmoid/sum_combine/l2norm into _gluon_kernels/common/utils.py (same definitions as #6609); k1 uses them. - Align flash_kda.py with #6609: _k2_gluon_schedule(..., kernel) returns a Config, use_gluon_k2a/k2c split, _GLUON_K2C_ARCH_OK, unpaged_args; gfx950 k2 gets gluon_k2_ab_fused/gluon_k2_c launchers mirroring gfx1250's. - Tests: arch-generic Gluon module imports, pass-C checks gated on arch. - Fix remap_xcd import path (moved to _triton_kernels/common in #6505). - Trim verbose comments. * [Triton/Gluon] KDA: address review — published schedules, resolved-route log, launch-level route tests - K1 and pass C read their Gluon schedules through _published_gluon_schedule and fall back to Triton when none is published (no Python fallback). The K1 schedule moves to the wrapper; gluon_k1_prepare rejects a wrong warp count. - Segment cap moves to the gfx950 Triton config (flash_kda_segments: MAX_CHUNKS=64, HARD_CAP). Other archs keep the previous 32-chunk soft cap. - intra_attn token-parallel default (BH=4, 1 warp) moves to the gfx950 Triton config; other archs keep the previous BH=1, 4-warp default. - Use get_num_sms() (honors CU_NUM) for routing and segmentation. - Log the resolved K1, pass A and pass C routes after pass C is settled. - Pass A and pass C share _solve_u/_update (AMDGCN unchanged). - Tests: count actual pass-C launches for size routing, cover the segment cap and unpublished-schedule fallback, test K1 warp validation via flash_kda_fwd.
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.
Motivation
On gfx1250, FlashKDA (flash_kda_fwd, invoked by chunk_kimi_delta_attn) uses pointer‑based
tl.loadfor loads, which Triton does not lower to TDM. Furthermore, its performance is worse than MI355, which is unexpected.Technical Details
gfx1250 Gluon K1 and K2
aiter/ops/triton/_gluon_kernels/gfx1250/kimi_delta_attn/flash_kda_k1.py: the prepare kernel. q, k and g arrive by one fused TDM load, and the [C, C] products and the (I - L)^-1 solve run on wave32 WMMA. Same workspace contract andgluon_k1_prepareentry point as the gfx950 K1, plus the token strides,num_warpsandwaves_per_eu.aiter/ops/triton/_gluon_kernels/gfx1250/kimi_delta_attn/flash_kda_k2.py: the chunk walk. Pass A computesb_segandA_segin one launch, where Triton takes two; pass C computes the outputs and the final state, pagedstate_cacheincluded. The state stays transposed in a WMMA accumulator, the workspace tiles and v arrive by fused TDM loads into a double buffer, and the outputs leave LDS by async copy.aiter/ops/triton/_gluon_kernels/common/utils.py:sigmoid_log2,fast_sigmoid,sum_combineandl2norm, which the gfx1250 K1 and K2 import. Thesigmoidalready in this file uses tanh and gives different results.aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py: on gfx1250, K1 and both K2 passes go to Gluon.AITER_FDA_USE_GLUONnow applies on gfx1250 too: 1, the default, routes K1 and K2 to Gluon,k1ork2routes only that one, and 0 keeps FlashKDA on Triton. Routing on gfx950 is unchanged; pass C stays on Triton there._k1_gluon_schedulereads K1'snum_warpsandwaves_per_eufrom the published config._k2_gluon_scheduletakes the pass name and returns its whole entry, so each pass has its own wide and narrow entries, CPL included.aiter/ops/triton/configs/gfx1250/gluon/attention/chunk_delta_attn/DEFAULT.json: K1 at four warps andwaves_per_eu5. Each K2 pass has BW 128 (wide) and BW 64 (narrow) at four warps, with CPL 2, the number of chunks in one TDM load. A pass runs wide when that still gives at least 2 blocks per CU for pass A, or 0.3 for pass C; both floors were calibrated on thebench_chunk_kda.pyandbench_flash_kda.pyshapes.Pass B on gfx1250
flash_kda.py: pass B, the cross-segment scan, stays on Triton. It is a chain of dependent updates, and_flash_kda_seg_scan_kernelloadsb_segin the update that adds it: two exposed memory round trips per update. On gfx1250, pass B runs_flash_kda_seg_scan_prefetch_kernelat BV=16 and four warps. It does two updates per loop trip and loads the next trip'sb_segtwo to three updates ahead, while the software pipeliner prefetchesA_seg. It also skips the last segment's update, which nothing reads.h_inmatches the old kernel to 2.4e-7 with dense, no and paged initial states.AITER_FDA_SCAN_PREFETCH=0goes back to the old kernel; other archs keep the old kernel and its schedule.Strided q, k, v and g, read in place
flash_kda.py: q, k, v and g are left out ofinput_guard. A view that a gfx1250 Gluon kernel reads goes in as it is when the descriptor can walk it (the heads of each token packed, and B == 1); anything else is packed as before. Packed inputs only pay for the contiguity checks.aiter/ops/triton/kimi_delta_attn/chunk_delta_attn.py,aiter/ops/triton/_triton_kernels/kimi_delta_attn/chunk_fwd.py: the public wrapper no longer packs q, k, v and g ahead ofchunk_delta_attn_fwd, which packs them after the FlashKDA branch for the kernels below it.Benchmark
op_tests/op_benchmarks/triton/bench_flash_kda.py:--k3-tp8runs the 26 prefill forwards of Kimi-K3 TP8+DCP8 rank 0. Each forward is one varlen call, and a sequence that continues a cached prefix reads its initial state from a paged state cache.--qkv-bandspasses q, k and v as bands of one fused projection, as the model does.--impl chunk_kdatimes the Gluon chunk_kda on the same inputs, and--shapecan repeat. The default sweep and its inputs do not change.Tests
op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py: the schedule and routing tests also cover pass C on archs where it has a Gluon kernel, and read BW and num_warps from the returned config.op_tests/triton_tests/kimi_delta_attn/test_fast_launch.py: the fast-launch tests take the Gluon kernels of the device's arch, pass C on gfx1250 included, and the prefetching pass B kernel.Test Plan
HIP_VISIBLE_DEVICES=0). Imagerocm/fw-bringup-mlperf:gfx1250-atom-20261008: ROCm SDK 10.1.0a20260811, PyTorch 2.11.0, Triton 3.8.0. Upstream ismainat 8779c30, the merge base of this PR.op_tests/triton_tests/kimi_delta_attn/test_flash_kda.pyandtest_fast_launch.pyon gfx1250, withAITER_FDA_SCAN_PREFETCH=1(the default) and0.op_tests/op_benchmarks/triton/bench_chunk_kda.py(--timing do_benchand--timing cudagraph) andbench_flash_kda.py, upstream and this PR interleaved in one session. Upstream runs this PR'sbench_flash_kda.py, as--k3-tp8,--qkv-bandsand--implare new. Each result gives its command and its number of rounds; times are means over the rounds. NoAITER_FDA_*variable is set unless a result says so.Test Result
test_flash_kda.pyandtest_fast_launch.py: 107 passed, 2 skipped, with and without the pass B prefetch (gfx1250 has no Triton K2 autotune candidates to switch between; one test needs a second GPU).bench_chunk_kda.py --timing do_bench, us, two rounds. H=24 and H=12 are Kimi-K3's 96 heads sharded TP4 and TP8. The script'stritonline callschunk_kimi_delta_attn, which runs FlashKDA (flash_kda_fwdinflash_kda.py); itsgluonline runs chunk_kda. On this PR thetritonline runs the Gluon K1 and K2 by default, which is This PR, andAITER_FDA_USE_GLUON=0 AITER_FDA_SCAN_PREFETCH=0puts it back on upstream's Triton kernels and schedules, which is upstream (triton). Thegluonline reads neither variable and chunk_kda is the same code on both trees, so it is upstream (chunk_kda gluon) in either run. Run from the root of this PR's checkout;PYTHONPATH=$PWDkeeps an installed aiter from shadowing it.The upstream columns were measured on main at 8779c30. On three of the shapes, one of them segmented, this PR with both variables off launches the same kernels as main, with the same grids and blocks, and gives bit-identical outputs. Over four interleaved rounds on a second GPU of the same machine, its
tritonline is 0.90-1.10x of main's per shape, geomean 0.99x; main's own rounds differ by up to 39% at a shape, as its Triton K2 autotunes in each process.This PR is 2.60-4.49x faster than upstream (triton), geomean 3.23x, and 1.32-2.87x faster than upstream (chunk_kda gluon), geomean 1.59x. With
--timing cudagraph(one round on upstream, two on this PR): 2.46-5.70x and 1.25-3.55x.bench_flash_kda.py --k3-tp8, us, four rounds; the chunk_kda gluon column is--impl chunk_kdaon upstream. These are the 26 prefill forwards of Kimi-K3 TP8+DCP8 rank 0 (H=12) in order, and step 13 is the headline case.seqlenslists the sequences of each forward; only the sequences that continue a cached prefix start from a carried state.This PR is 3.69-4.72x faster than upstream (triton), geomean 4.20x, and 3.66-7.19x faster than upstream (chunk_kda gluon), geomean 5.13x. Upstream's Triton K2 autotunes in each process, so its times move by up to 16% between rounds; this PR's move by less than 5%.
bench_flash_kda.py --qkv-bands --shape 1 4096 96 128 128 --shape 1 13312 96 128 128 --shape 1 16384 96 128 128, ms, four rounds; the chunk_kda gluon column is--impl chunk_kdaon upstream. H=96, with q, k and v as bands of one fused projection.This PR is 2.78-3.00x faster than upstream (triton) and 1.44-1.46x faster than upstream (chunk_kda gluon). Upstream packs the bands before its kernels, which costs 0.14-0.45 ms per call over packed inputs (the same command without
--qkv-bands; one round on upstream, two on this PR). This PR reads the bands in place and stays within 3 us of packed inputs.Pass B's prefetch, the two commands above with
AITER_FDA_SCAN_PREFETCH=0against the default, two rounds: the TP8 forwards are 1.06-1.22x faster with it, geomean 1.09x; H=96 does not change (0.997-1.001x).bench_flash_kda.py(default sweep; one round on upstream, two on this PR): 2.13-5.65x over upstream on 12 of its 13 shapes. B=1, T=256, H=12 is host-bound: a call takes about 44 us on both trees (0.98x). The in-place reads add about 0.9 us of host time per call (four more kernel arguments); GPU time on packed inputs is unchanged.Submission Checklist