Repository navigation
[Triton/Gluon] Improving GFX950 Unified Attention for spec. decode and reduce opimizations - #6336
Merged
Merged
Conversation
…oken batches, chunked segment reduce Fixes the gfx950 Gluon slowdowns reported in #6221 (speculative-decode verify and decode with few KV heads). - Batches of short queries (max_seqlen_q * queries_per_kv <= 128, no sliding window) take the split-KV grid when the kv_split config splits them. The query block is widened to hold every query of a sequence, so the split ranges match what reduce_segments assumes. Before, every such batch took the unsplit 2d grid: at batch 1 with 2 KV heads that is 2 programs walking the whole context. - store_partial writes -inf as the max of a row that saw no key in its split. - reduce_segments takes SEGMENT_BLOCK and merges the segments in chunks of that size with an online softmax. The single [NUM_SEGMENTS, head] tile spills at head size 256 and 512. 0 (default) keeps the single-tile code. - gfx950 configs: kv_split caps the split at 128 for decode and 64 otherwise (vLLM full-graph capture passes max_model_len as max_seqlen_k, so the uncapped fp8 split reached 512); D >= 256 decode uses the KV double buffer; the Triton reduce section sets SEGMENT_BLOCK 32 for D >= 256.
…e single-tile reduce as on main - gfx942 Triton reduce config: SEGMENT_BLOCK 32 for D > 128, as on gfx950. A D_LEQ_128 entry with the previous values keeps D <= 128 on the single tile. - reduce_segments: the single-tile branch keeps main's op order, so an arch without SEGMENT_BLOCK compiles to the same instructions as before.
…56 entry, leave any as on main Replaces the D_LEQ_128 entry and the SEGMENT_BLOCK on any with a D_GEQ_256 reduce entry (the previous any values plus SEGMENT_BLOCK 32), so head sizes below 256 keep the single-tile reduce from main.
…block bound into _gfx950_split_grid, drop the new tests
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 new routing and reduction paths need targeted coverage, and the kernel representation omits the new tuning key.
3 open findings
What changed in this PR
Improves unified attention for speculative decoding on GFX950 by refining split-KV routing and reducing register pressure.
Changes:
- Routes undersubscribed multi-token Gluon workloads through split-KV.
- Adds chunked Triton segment reduction.
- Retunes GFX950 configurations and extends GFX942 reduction tuning.
| File | Description |
|---|---|
configs/gfx950/triton/attention/unified_attention/DEFAULT.json |
Enables chunked reduction. |
configs/gfx950/gluon/attention/unified_attention/DEFAULT.json |
Retunes kernels and split heuristics. |
configs/gfx942/triton/attention/unified_attention/DEFAULT.json |
Enables chunking for larger heads. |
attention/unified_attention.py |
Adds GFX950 split-grid routing. |
_triton_kernels/attention/unified_attention.py |
Implements chunked segment reduction. |
_gluon_kernels/gfx950/attention/unified_attention.py |
Aligns split boundaries and handles empty partials. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Contributor
Author
|
Comparing #6339 vs this PR( #6336): Cold cache: L2/MALL flushed each call, profiler device time
|
vgokhale
approved these changes
Oct 12, 2026
This was referenced Oct 12, 2026
This was referenced Oct 19, 2026
a-sidorova
pushed a commit
to sshlyapn/aiter
that referenced
this pull request
Oct 21, 2026
…d reduce opimizations (ROCm#6336)
4 of 7 tasks
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
Addressing #6221 and beyond.
Main improvements:
(can be enabled via SEGMENT_BLOCK from the config for the reduce, 0 is the default full head size chunking)
Technical Details
Results
Issue shape: 32/2 heads, D=256, page 64, batch 1, fp8 (q/k/v descales)
All 96 runs checked against the torch reference: max abs error 6.07e-03.
Against Triton main, Gluon on the branch is faster or level everywhere except q >= 2 at 200 tokens with live max_seqlen_k. There the batch stays on the unsplit 2d grid, same as main: 15.7 vs 9.8 us. Against Triton on the branch, which now has the faster reduce, Gluon is within about 10 % up to 2k context and 9-19 % faster for verify from 12k up.
Gluon backend, main vs branch (per shape/dtype over B 1-128, q 1/2/4, ctx 2k/30k, live/graph)