Repository navigation
[Triton/Gluon] [CI] Move common kernels into common/ - #6505
Merged
Merged
Conversation
Device helpers reused across kernel folders lived in two places that do not
own them: the activation helpers inside _triton_kernels/activation.py, and
pid preprocessing / fp8 scaling / moe zero-write helpers under
utils/_triton/. Move them to _triton_kernels/common/, next to
splitk_reduce.py and segment_tile.py.
- Split _triton_kernels/activation.py: the activation helpers (_silu,
_sigmoid, _gelu, _tanh, _relu, ..., _get_activation_from_str,
_apply_activation_from_str) move to common/activation.py; the four
kernels launched by activation.py stay where they are.
- Move utils/_triton/{pid_preprocessing,mha_kernel_utils,moe_common}.py to
_triton_kernels/common/ (pure renames, no code changes).
- Update every importer and the mha_kernel_utils _BACKWARD_COMPAT_MAP
entry.
- Update README.md and aiter-ops-triton.instructions.md to point at the new
locations.
No kernel code changes; function bodies are byte-identical.
The four kernels left in _triton_kernels/activation.py (fused_silu_mul_kernel, _silu_and_mul_backward_kernel, _act_mul_and_dynamic_mxfp4_quant_kernel, _act_mul_and_dynamic_fp8_group_quant_kernel) are activation+mul fusions, so move them, and the wrapper that launches them, into fusions/: - aiter/ops/triton/activation.py -> aiter/ops/triton/fusions/activation.py - _triton_kernels/activation.py -> _triton_kernels/fusions/activation.py - test_activation.py and test_fused_silu_mul.py move to triton_tests/fusions/ (the split-test-times entries follow them) Both old module paths keep working through _BACKWARD_COMPAT_MAP: aiter.ops.triton.activation and aiter.ops.triton._triton_kernels.activation. In-tree callers now import from the categorized path.
vgokhale
requested review from
a team,
Boss2002n,
azaidy,
brunomazzottiamd and
rahulbatra85
October 13, 2026 22:33
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
The compatibility alias omits previously available activation helpers, and the README introduces a contradictory portability rule.
1 open finding
What changed in this PR
Moves shared Triton helpers into _triton_kernels/common/ and activation fusions into categorized folders, updating consumers and documentation accordingly.
Changes:
- Relocates activation, PID-remapping, MHA, and MoE helpers.
- Updates Triton/Gluon imports, tests, benchmarks, and compatibility mappings.
- Documents the revised kernel layout.
| File | Description |
|---|---|
op_tests/triton_tests/fusions/test_silu_and_mul_backward.py |
Updates wrapper import. |
op_tests/triton_tests/fusions/test_fused_silu_mul.py |
Updates wrapper import. |
op_tests/triton_tests/fusions/test_activation.py |
Updates wrapper import. |
op_tests/op_benchmarks/triton/bench_silu_and_mul_backward.py |
Updates benchmark import. |
aiter/ops/triton/README.md |
Documents helper placement. |
aiter/ops/triton/quant/fused_mxfp4_quant.py |
Updates activation helper path. |
aiter/ops/triton/quant/fused_fp8_quant.py |
Updates activation helper path. |
aiter/ops/triton/gemm/fused/fused_gemm_a16w16_quant_x.py |
Updates activation helper path. |
aiter/ops/triton/gemm/feed_forward/ff_a16w16_fused_ungated.py |
Updates activation helper path. |
aiter/ops/triton/gemm/feed_forward/ff_a16w16_fused_gated.py |
Updates activation helper path. |
aiter/ops/triton/gemm/basic/gemm_a16w16.py |
Updates activation helper path. |
aiter/ops/triton/gemm/basic/gemm_a16w16_gated.py |
Updates activation helper path. |
aiter/ops/triton/fusions/activation.py |
Uses relocated fusion kernels. |
aiter/ops/triton/_triton_kernels/quant/sage_attention_quant.py |
Uses common PID helpers. |
aiter/ops/triton/_triton_kernels/moe/moe_wgrad.py |
Uses common XCD remapping. |
aiter/ops/triton/_triton_kernels/moe/moe_situ_epilogue.py |
Uses common activation helper. |
aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_int8_smoothquant.py |
Uses common PID helper. |
aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a8w8.py |
Uses common PID helper. |
aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a8w8_blockscale.py |
Uses common PID helper. |
aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a8w4.py |
Uses common PID helper. |
aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a4w4.py |
Uses common PID helper. |
aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a16w4.py |
Uses common PID helper. |
aiter/ops/triton/_triton_kernels/gemm/grouped/gmm.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_afp4wfp4_split_cat.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_afp4wfp4_mul_add.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_afp4wfp4_a16w16.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a8w8_blockscale_split_cat.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a8w8_blockscale_mul_add.py |
Uses common PID helper. |
aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a8w8_blockscale_a16w16.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a16w16_quant_x.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/feed_forward/ff_a16w16_fused_ungated.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/feed_forward/ff_a16w16_fused_gated.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/batched/batched_gemm_afp4wfp4.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/batched/batched_gemm_a16wfp4.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/basic/gemm_afp8wfp8.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/basic/gemm_afp4wfp4.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a8wfp4.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a8w8.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a8w8_per_token_scale.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a8w8_blockscale.py |
Uses common PID helper. |
aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16wfp4.py |
Uses common PID helper. |
aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w8_blockscale.py |
Uses common PID helper. |
aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16_persistent.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16_gated.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16_atomic.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/fusions/fused_sigmoid_mul.py |
Uses common activation helper. |
aiter/ops/triton/_triton_kernels/fusions/fused_clamp_act_mul.py |
Uses common activation dispatch. |
aiter/ops/triton/_triton_kernels/fusions/fused_bmm_rope_kv_cache.py |
Uses common PID helper. |
aiter/ops/triton/_triton_kernels/fusions/activation.py |
Separates fusion kernels from shared activations. |
aiter/ops/triton/_triton_kernels/conv/winograd_f4x3_common.py |
Uses common activation dispatch. |
aiter/ops/triton/_triton_kernels/conv/conv3d_general.py |
Uses common activation dispatch. |
aiter/ops/triton/_triton_kernels/conv/conv3d_3x3x3.py |
Uses common activation dispatch. |
aiter/ops/triton/_triton_kernels/conv/conv3d_1x1x1.py |
Uses common activation dispatch. |
aiter/ops/triton/_triton_kernels/conv/conv_general.py |
Uses common activation dispatch. |
aiter/ops/triton/_triton_kernels/conv/conv_3x3.py |
Uses common activation dispatch. |
aiter/ops/triton/_triton_kernels/conv/conv_1x1.py |
Uses common activation dispatch. |
aiter/ops/triton/_triton_kernels/common/pid_preprocessing.py |
Houses shared PID/XCD helpers. |
aiter/ops/triton/_triton_kernels/common/moe_common.py |
Houses shared MoE zero-fill helper. |
aiter/ops/triton/_triton_kernels/common/mha_kernel_utils.py |
Houses shared FP8 scaling helper. |
aiter/ops/triton/_triton_kernels/common/activation.py |
Houses shared activation helpers. |
aiter/ops/triton/_triton_kernels/attention/mla_decode_rope.py |
Updates shared-helper imports. |
aiter/ops/triton/_triton_kernels/attention/mha.py |
Updates MHA and PID helper imports. |
aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py |
Updates MHA helper import. |
aiter/ops/triton/_triton_kernels/attention/mha_fused_bwd.py |
Updates MHA and XCD imports. |
aiter/ops/triton/_triton_kernels/attention/extend_attention.py |
Updates shared-helper imports. |
aiter/ops/triton/_gluon_kernels/gfx950/moe/moe_op_gemm_a16w4.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_gluon_kernels/gfx950/gemm/basic/gemm_afp4wfp4.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_gluon_kernels/gfx950/gemm/basic/gemm_a8w8.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_gluon_kernels/gfx950/gemm/basic/gemm_a8w8_blockscale.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_gluon_kernels/gfx950/gated_delta_net/fused_gdn_decode_qkvz.py |
Uses common sigmoid helper. |
aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_mla.py |
Uses common XCD remapping. |
aiter/ops/triton/_gluon_kernels/gfx950/attention/mha.py |
Updates MHA and XCD helper imports. |
aiter/ops/triton/_gluon_kernels/gfx1250/moe/moe_op_gemm_a8w4.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_gluon_kernels/gfx1250/moe/moe_op_gemm_a4w4.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_gluon_kernels/gfx1250/moe/moe_op_gemm_a16w4.py |
Uses common PID/XCD helpers. |
aiter/ops/triton/_gluon_kernels/gfx1250/gemm/basic/gemm_mxfp8.py |
Uses common PID helper. |
aiter/ops/triton/_gluon_kernels/gfx1250/gemm/basic/gemm_a8w8_blockscale.py |
Uses common PID helper. |
aiter/ops/triton/_gluon_kernels/gfx1250/gemm/basic/gemm_a16w16_persistent.py |
Uses common XCD remapping. |
aiter/ops/triton/_gluon_kernels/gfx1250/fusions/fused_clamp_act_mul.py |
Uses common activation dispatch. |
aiter/ops/triton/__init__.py |
Adds relocation compatibility mappings. |
.github/split-test-times/triton.tsv |
Updates categorized test paths. |
.github/instructions/aiter-ops-triton.instructions.md |
Updates repository layout guidance. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Boss2002n
added a commit
that referenced
this pull request
Oct 14, 2026
Conflicts: - aiter/ops/triton/README.md and the Copilot review rules: kept the short versions from this branch and added the changes of #6505. The shared @triton.jit device helpers are now in _triton_kernels/common/, and activation.py is not a flat wrapper any more. - aiter/ops/triton/utils/_triton/tuning/README.md: took the new version from #6128.
yuychang
added a commit
to yuychang/sglang
that referenced
this pull request
Oct 22, 2026
AITER main moved aiter.ops.triton.utils._triton.pid_preprocessing to aiter.ops.triton._triton_kernels.common.pid_preprocessing. Try the new path first and fall back to the old one for older AITER builds.
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>
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.
yuychang
added a commit
to yuychang/sglang
that referenced
this pull request
Oct 29, 2026
…pilogue - Add batched_gemm_a16wfp4_flatten_mxfp4_quant in sglang.kernels.ops.gemm.mla_value_bmm_mxfp4_hip (Triton, gfx95) - Quantize the accumulator to flattened per-1x32 MXFP4 in the epilogue so the BF16 intermediate never reaches HBM - Round to BF16 before quantizing so values and scales match the split GEMM + fused_flatten_mxfp4_quant bit for bit - Add can_fuse_mla_value_bmm_mxfp4_quant to probe dtypes and the pinned Kimi N=128, K=512 value shape - Fall back to the split path when AITER has no non-split config for the shape - Register gemm.batched_gemm_a16wfp4_flatten_mxfp4_quant (Triton, HIP) in the gemm kernel inventory - Support eager mode, torch.compile (custom op with a fake impl) and HIP graph capture - Import pid_preprocessing from its ROCm/aiter#6505 location, falling back to the old path on older AITER - Add an MI35x parity test: token counts, torch.compile, graph replay and the config-lookup fallback - Add a manual benchmark that replays rotated 1 GiB inputs so reads are cold in L2 and the Infinity Cache - No call sites yet; the MLA decode path adopts the kernel in the next commit
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.
This is mainly just cleanup. It moves all common functions into the common folder. And it moves a few fusions from activation.py into fusions/