You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[Triton/Gluon] KDA: faster FlashKDA and token-parallel intra forward on gfx950 - #6177
Faster KDA forward on gfx950: the FlashKDA kernels, and the default pipeline's token-parallel intra kernel.
Technical Details
FlashKDA K2: XCD remap. Program ids are renumbered so that all V blocks of one (segment, head) run on the same XCD and share its L2 for re-reads of the chunk workspace. Applies to the Triton segment kernel and the Gluon pass A.
intra_token_parallel:
Token i's q/k/g are loaded once instead of once per (j, K block), and one block spans all of K.
The default config goes from BH=1 / 4 warps to BH=4 / 1 warp.
This kernel only runs on the default pipeline with safe_gate=False.
Segment length: a segment is capped at 64 chunks. Before, a long sequence with few heads could get 256-chunk segments, and their serial walk set the runtime.
Gluon K1:
Sigmoids go through exp2.
The cumulative gate is computed in-thread, in a column-per-thread layout.
Barriers drop from 38 to 10.
The gate load is now masked, since buffer loads past the end of the tensor are not clamped.
Launches with waves_per_eu=3.
New ws_beta workspace: when K2 runs on Gluon, both K1s (Gluon and Triton) write sigmoid(beta) to it, so K2 can copy it into LDS along with the chunk's other tiles. It adds 128 B per (chunk, head) tile to a 29 KiB workspace (0.44%), and is neither allocated nor written on the Triton K2 route.
Gluon pass A:
While chunk j computes, chunk j+1's tiles are async-copied into a double-buffered LDS stage.
One barrier per chunk.
The published schedules set waves_per_eu=2.
New Gluon pass C (k2_c_gluon):
Uses the same LDS prefetch as pass A.
The output projection runs before the state update, and the output is stored straight from the MFMA layout.
Supports the paged state_cache read and write-back, like the Triton pass C.
Tile width is BW=64 unsegmented and BW=128 segmented.
Used for every segmented shape. Unsegmented, it is used only above a published MIN_BLOCKS_PER_CU floor (gfx950: 0.3), the same mechanism pass A uses. Below the floor, i.e. with few (sequence, head) pairs, the serial chunk walk sets the time and the Triton pass C's BW=16 tile is faster. Measured on MI355X: Triton is 1.08–1.18x faster at 32 pairs or fewer; Gluon ties or wins from 40 pairs up, by as much as 2.2x.
Kernel names: the Gluon K2 kernels take a NUM_WARPS constexpr, statically asserted to equal gl.num_warps(), so the warp count shows up in the kernel name.
bench_chunk_delta_attn.py --no-safe-gate: new flag that runs without safe_gate, the only path that reaches the token-parallel intra kernel.
Test Plan
op_tests/triton_tests/kimi_delta_attn/ on gfx950 (MI355X). Test changes:
test_routes_agree and test_cases_reach_the_gluon_k2 pin pass C to Gluon, and check that every case reaches it (segmented cases reach pass A too).
New test checking that pass C is routed by size.
test_paged_cache_matches_dense runs on both the Triton and Gluon K2 routes.
The schedule test covers the pass C tile widths.
test_fast_launch covers the pass C launcher.
The repo's benchmark scripts, run on main and on this branch on an idle MI355X:
Outputs and final states compared against main across 10 FlashKDA shapes, including varlen, and 6 default-pipeline cases.
black and ruff check (ruff 0.16.0) on the changed files.
Test Result
Tests: 138 passed, 1 skipped.
Numerics:
FlashKDA output is within 4.6e-4 relative of main and the final state within 3.5e-5. The difference comes from K1's l2norm and gate summation order.
Default pipeline with safe_gate=False: output within 2.2e-5. Only the summation order of Aqk/Akk changes.
Default pipeline with safe_gate=True: bitwise identical.
bench_flash_kda.py (ms):
B×T×H (K=V=128)
main
this PR
Speedup
1×256×12
0.0292
0.0247
1.18x
1×512×12
0.0385
0.0331
1.16x
1×1664×12
0.0839
0.0806
1.04x
1×2048×12
0.1030
0.0713
1.44x
1×8192×12
0.2770
0.1414
1.96x
1×16384×12
0.5154
0.2395
2.15x
1×16384×64
2.4625
0.9872
2.49x
1×16384×96
2.8588
1.3068
2.19x
1×8192×32
0.6569
0.2754
2.39x
1×8192×64
1.2858
0.4933
2.61x
2×2664×12
0.2007
0.1203
1.67x
4×4096×16
0.6641
0.2689
2.47x
8×2048×16
0.4906
0.2068
2.37x
bench_chunk_delta_attn.py --no-safe-gate (ms):
B×T×H, K=V
main
this PR
Speedup
1×512×16, 64
0.1886
0.1853
1.02x
1×1024×16, 64
0.1793
0.1840
0.97x
1×2048×16, 64
0.2597
0.1819
1.43x
1×4096×16, 64
0.4682
0.2207
2.12x
1×8192×16, 64
0.8517
0.3769
2.26x
1×16384×16, 64
1.7663
0.8361
2.11x
2×2048×16, 64
0.4316
0.1907
2.26x
2×4096×16, 64
0.7806
0.3033
2.57x
4×2048×16, 64
0.7434
0.2663
2.79x
4×4096×16, 64
1.5173
0.5487
2.77x
1×8192×32, 128
2.8719
1.2039
2.39x
1×8192×64, 128
5.4671
2.1663
2.52x
1×8192×96, 128
8.0437
2.9434
2.73x
1×8192×128, 128
10.7907
3.9281
2.75x
1×16384×32, 128
5.7695
2.4195
2.38x
1×16384×64, 128
10.9629
4.1234
2.66x
1×16384×96, 128
15.9610
5.7415
2.78x
1×16384×128, 128
21.4549
7.7804
2.76x
The two shortest shapes run at the same speed on both, within noise.
bench_chunk_delta_attn.py with its default safe_gate=True stays on the sub-chunk intra kernel, which this PR doesn't touch. All 18 of its shapes are within ±4% of main (noise).
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
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
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).
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
_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.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
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.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
All standard extended tests (excludes ci:atom_full)
Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6177 --add-label <label>
One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.
PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.
… 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.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The new 64-chunk cap has no assertion in the scheduling tests, and the existing numerical tests would still pass if this max/min ordering regressed because segmentation is mathematically equivalent. Add a focused heuristic test covering a long, few-head shape where the length cap must override _SEG_MAX_SEGMENTS, and assert the returned segment length is at most 64.
This gfx950-calibrated default is selected whenever autotuning is disabled on every architecture, so the PR also changes the gfx942 launch from BH=1/4 warps to BH=4/1 warp. The repository’s Python-side config rule keeps architecture-specific tuning in configs/<arch>/...; publish this default in the gfx950 Triton JSON and resolve it through the family config loader while retaining a portable fallback for other architectures.
…ch 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).
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This adds a hardcoded fallback tile in Python, contrary to the Triton config rule that tuning values and fallback schedules live in the family JSON rather than kernel modules. Keeping BW/warp choices here creates a second source of truth beside configs/<arch>/gluon/attention/chunk_delta_attn/DEFAULT.json; please route the default through the family config loader instead. See coding guideline 1000002.
Move tuning search space and defaults into family config
The updated BH/warp/stage search space and default are hardcoded in the kernel module. Repository policy requires tuning values—including default configs—to come from the family JSON so retuning does not require a code change; autotune_configs only pins a supplied config and does not make these inline values configuration-driven. Move this shortlist/default into the chunk-delta-attention config loader/JSON. See coding guideline 1000002.
…fx950 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).
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Shared recurrence duplication and the independently bundled token-parallel optimization need resolution under the repository’s reuse and PR-scope rules.
…ute 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.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…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.
…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
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.
omuhamma
left a comment
Motivation
Faster KDA forward on gfx950: the FlashKDA kernels, and the default pipeline's token-parallel intra kernel.
Technical Details
intra_token_parallel:BH=1/ 4 warps toBH=4/ 1 warp.safe_gate=False.exp2.waves_per_eu=3.ws_betaworkspace: when K2 runs on Gluon, both K1s (Gluon and Triton) writesigmoid(beta)to it, so K2 can copy it into LDS along with the chunk's other tiles. It adds 128 B per (chunk, head) tile to a 29 KiB workspace (0.44%), and is neither allocated nor written on the Triton K2 route.waves_per_eu=2.k2_c_gluon):state_cacheread and write-back, like the Triton pass C.MIN_BLOCKS_PER_CUfloor (gfx950: 0.3), the same mechanism pass A uses. Below the floor, i.e. with few (sequence, head) pairs, the serial chunk walk sets the time and the Triton pass C's BW=16 tile is faster. Measured on MI355X: Triton is 1.08–1.18x faster at 32 pairs or fewer; Gluon ties or wins from 40 pairs up, by as much as 2.2x.NUM_WARPSconstexpr, statically asserted to equalgl.num_warps(), so the warp count shows up in the kernel name.bench_chunk_delta_attn.py --no-safe-gate: new flag that runs withoutsafe_gate, the only path that reaches the token-parallel intra kernel.Test Plan
op_tests/triton_tests/kimi_delta_attn/on gfx950 (MI355X). Test changes:test_routes_agreeandtest_cases_reach_the_gluon_k2pin pass C to Gluon, and check that every case reaches it (segmented cases reach pass A too).test_paged_cache_matches_denseruns on both the Triton and Gluon K2 routes.test_fast_launchcovers the pass C launcher.mainand on this branch on an idle MI355X:python op_tests/op_benchmarks/triton/bench_flash_kda.pypython op_tests/op_benchmarks/triton/bench_chunk_delta_attn.pypython op_tests/op_benchmarks/triton/bench_chunk_delta_attn.py --no-safe-gatemainacross 10 FlashKDA shapes, including varlen, and 6 default-pipeline cases.blackandruff check(ruff 0.16.0) on the changed files.Test Result
mainand the final state within 3.5e-5. The difference comes from K1's l2norm and gate summation order.safe_gate=False: output within 2.2e-5. Only the summation order ofAqk/Akkchanges.safe_gate=True: bitwise identical.bench_flash_kda.py(ms):bench_chunk_delta_attn.py --no-safe-gate(ms):The two shortest shapes run at the same speed on both, within noise.
bench_chunk_delta_attn.pywith its defaultsafe_gate=Truestays on the sub-chunk intra kernel, which this PR doesn't touch. All 18 of its shapes are within ±4% ofmain(noise).Submission Checklist
🤖 Generated with Claude Code