Repository navigation
Conversation
… decode The gfx1250 Gluon 3d kernel read only the first element of k_descale and v_descale. A per-token KV cache (one scale per cached token and KV head, [num_blocks, num_kv_heads, block_size]) then used one scale for all tokens. Add PER_TOKEN_KV_SCALE. Each tile is one page and loads the K and V scales of its tokens. The K scales multiply the QK columns. The V scales multiply the P columns after the row sum. Slots past seq_len get scale 0. The wrapper finds per-token scales by their page shape. It raises NotImplementedError on paths that cannot read them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
P, scaled by the per-token V scales, was truncated to bf16 before P @ V. The truncation biases the output low by about 0.3% against the fp32 row sum. Split P into a bf16 high part and a bf16 low part and run two WMMAs. The error then matches an fp32 PV, and decode time changes by 0-3%. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ernel On gfx1250 every batch that was not all-decode took the 2d kernel. The 2d grid has one program per (KV head, query block), so a speculative verify of a few requests over a long context leaves most CUs idle: 3 requests x 4 tokens at 227K context take 16.5 ms per layer. Route queries of at most 4 tokens to the 3d split-K kernel when each query block holds one query position (BLOCK_Q == 1). A query block that ends before a segment now stores an empty partial, so the reduce weighs it 0. The same shape takes 75 us (bf16 KV) or 177 us (per-token fp8 KV). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
Rule 18 of aiter/ops/triton/README.md allows one or two sentences in a comment. Shorten the docstrings of compute_pv_per_token and _gfx1250_split_short_queries. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
There was a problem hiding this comment.
🟡 Changes recommended
Short-query routing admits unsupported dtype pairs, unsafe query blocks and strided outputs that the direct store cannot handle.
3 open issues
What changed in this PR
Adds per-token fp8 KV scales and short-query split-K execution to gfx1250 Gluon unified attention.
Changes:
- Applies per-token scales with higher-precision probability accumulation.
- Routes eligible short queries through the 3d kernel.
- Adds accuracy, boundary and unsupported-path tests.
| File | Description |
|---|---|
op_tests/triton_tests/attention/test_unified_attention.py |
Tests per-token scales and short-query execution. |
aiter/ops/triton/attention/unified_attention.py |
Validates scales and selects short-query routing. |
aiter/ops/triton/_gluon_kernels/gfx1250/attention/unified_attention_3d.py |
Applies token scales and handles empty segments. |
Review effort: Balanced
| backend != "gluon" | ||
| or not params.shuffled_kv_cache |
| backend != "gluon" | ||
| or not params.shuffled_kv_cache | ||
| or params.max_seqlen_q > _GFX1250_SPLIT_MAX_QUERY_LEN | ||
| or not is_3d_gluon_available(params, backend) |
| return False | ||
| config = get_unified_attention_config("attn_3d", params, backend="gluon") | ||
| block_m = max(config["BLOCK_M"], triton.next_power_of_2(params.num_queries_per_kv)) | ||
| return block_m // params.num_queries_per_kv == 1 |
|
Closing. MiniMax-M3 on gfx1250 now takes the FlyDSL paged decode from #6462 (merged), which reads per-token FP8 KV scales itself, so unified attention no longer needs them. The ATOM side is on branch |
|
Not needed: MiniMax-M3 decode on gfx1250 now runs on FlyDSL paged decode (#6462), see ROCm/ATOM#2710. |
Draft. Ready for review after the ATOM side lands; no blocking issues in this PR.
Motivation
MiniMax-M3 on gfx1250 decodes through the Gluon unified attention kernels with a page-shuffled fp8 KV cache that has one K and one V scale per cached token (LLM-Booster #211). Two gaps block it:
k_descaleandv_descale. With a per-token cache the output is wrong without an error.Technical details
PER_TOKEN_KV_SCALE). The scales use the cache page layout[num_blocks, num_kv_heads, block_size], float32, contiguous. Each tile is one page and loads the scales of its tokens. The K scales multiply the QK columns. The V scales multiply the P columns after the softmax row sum. Slots pastseq_lenget scale 0.BLOCK_Q == 1), takes the 3d split-K kernel. A query block that ends before a segment stores an empty partial (M = -inf), so the reduce gives it weight 0. Other batches keep the 2d kernel.NotImplementedErroron paths that do not read them. Per-tensor scales use the same code as before.Test plan
test_unified_attention_per_token_kv_scale: decode, GQA 16 and 8, head 128, page 16 and 128, context 1 to 4096, token magnitudes that differ by up to 20x, NaN in unused scale slots. Checks against the dequantized reference, the unquantized reference, and the output bias.test_unified_attention_short_query_split: 2 to 4 tokens per request, bf16 and per-token fp8 KV, tails that end 1 to 4 tokens into a page, context up to 65537.test_unified_attention_per_token_kv_scale_rejects_prefill.test_triton_unified_attn_3d.Test results
gfx1250 (MI455), this branch on main @ 84077b3:
test_unified_attention.py: no test changes outcome against the same tree without the routing commit. Main already has failures intest_triton_unified_attnand a process abort in that file; both runs show them.Kernel accuracy and time (GQA 16, head 128, CUDA graph timing)
Accuracy against an fp32 reference on the same fp8 cache:
Decode time, per-token fp8 vs bf16 KV:
Speculative verify, 3 requests x 4 tokens:
🤖 Generated with Claude Code