triton attention: sync-free SWA metadata, per-shape launch-config table, split-KV policy (reorg 5/9) - #5
Open
mike-char wants to merge 5 commits into
Open
Conversation
update_sliding_window_buffer read GPU scalars per step (torch.empty sized by window_kv_indptr[-1] + a gpu-scalar slice bound), serializing step N+1's metadata prep behind step N's kernels and collapsing the overlap schedule. Thread the window totals from seq_lens_cpu / extend_prefix_lens_cpu instead; hot paths never read GPU scalars. (cherry picked from commit 864dca9fc, triton_backend portion)
Add SGLANG_G4_SPLITKV_POLICY (default OFF, byte-for-byte current behavior when unset). When on, get_num_kv_splits sizes each request's splits from its own KV length via a chunk chosen to fill ~3x the CU count with split workgroups, and the per-request split cap is raised to 64 for shapes with <=4 KV heads (Gemma-4 GLOBAL: 2 KV heads on 26b, 4 on 31b at head_dim 512). max_kv_splits is raised once at init so attn_logits/attn_lse/cuda_graph_* buffers and the kernel grid split-dim are all sized to the new cap at capture time; the policy only writes values in [1, max_kv_splits] into the existing num_kv_splits buffer, so no tensor shape or grid dim changes at graph replay. Feeds both the triton grouped kernel and the g4 HIP shim unchanged. Deferred to for MLA, deterministic, and fixed-split-tile modes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> (cherry picked from commit aa1115900d90553f5c0517489426d780de1daecd)
self.num_kv_head is the model's standard num_key_value_heads == the SWA-layer count for hybrid Gemma-4 (8 on 26b, 16 on 31b), so keying the few-KV-head cap raise on it alone missed the GLOBAL layers (2/4 KV heads) that are the whole point -- per-layer KV-head count is only known at dispatch (k_buffer.shape[-2]). Also raise the cap for any sliding-window model, which has a low-KV-head global layer by construction. SWA layers share the mid-buffer but their window-capped KV yields few splits, so the larger cap costs only buffer bytes, not wrong math. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> (cherry picked from commit e039b2f8ec4a624a39624d30872ee7e00d164883)
Measured (MI325X): the cost-model split policy REGRESSES the full/global shape (stage-2 combine + partial-output traffic outweigh the occupancy gain) and WINS ~6% on the window-capped SWA shape. Make the policy tunable without editing the file so the winning regime can be found by sweep: - SGLANG_G4_SPLITKV_CU_MULT (default 3): target split-workgroups per CU. - SGLANG_G4_SPLITKV_MIN_BLOCK (default 32): min/rounding tile for the chunk. - SGLANG_G4_SPLITKV_MAX_SPLITS_FEWKV (default 64): raised few-KV-head cap. - SGLANG_G4_SPLITKV_SCOPE = all|swa|full (default all): which decode shapes the per-request policy applies to. "swa" is the measured winner. Knobs are resolved once into _SplitKVConfig at init and logged (scope + values + resolved max_kv_splits) so bench logs are self-describing. Scope is decided at the policy writer from an is_window flag threaded from the two window call sites (the cleanest signal there: full seq_len vs window-capped length). swa scope also skips the few-KV-head cap raise, since window-capped KV never exceeds the stock cap and the larger mid-buffer would be pure waste. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> (cherry picked from commit fd1869d861ce60e5644520bb79e8a13815756b41)
Replace hardcoded launch constants in the grouped decode (stage1/stage2) and extend kernels with a bucketed lookup keyed by (kernel, head_dim, kv_group, kv_heads, bs/ntok bucket, ctx/prefix bucket). Resolution: env JSON (SGLANG_TRITON_ATTN_TUNING, sweep injection) > bundled table > hardware defaults that reproduce exactly the values the kernels launched with before this module existed — zero behavior change when no entry matches, and zero CUDA change. Measured tables ship separately as optional gemma4_mi325x_tuning.py (guarded import). Squashed final state of cab84c2a8 + 42d209ada + the tuning_table portions of c4cf5ecca/219af68f4 (original history on archive branches).
mike-char
force-pushed
the
cai-pr4-generic-features
branch
from
August 27, 2026 00:27
98438b9 to
2e2f367
Compare
mike-char
force-pushed
the
cai-pr5-triton-attn-tuning
branch
from
August 27, 2026 00:27
364e93e to
25b01af
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part 5 of 9.
Headline measured result (gemma-4 shapes, MI325X): waves_per_eu=0 beats the hardcoded 1 by 1.5-1.7x on the 512/GQA-8 decode shape. Upstream candidate: mechanism yes; needs GPU arch added to the table key first.
🤖 Generated with Claude Code