Optimize AMD GQA verify attention path

[fanxingran]

Add a HIP few-query GQA packed split-K attention path for EAGLE/MTP verify workloads, reducing repeated KV loads and improving long-context verify performance on AMD GPUs.

Motivation

SGLang's Triton extend_attention_fwd is currently a bottleneck for speculative decoding verify workloads on AMD ROCm, especially for EAGLE3 / MTP style shapes where a small number of draft tokens attend over a long paged KV prefix.

For GQA models, the existing kernel launches one program per query head. This causes every query head in the same GQA group to re-read the same KV head from HBM, creating up to kv_group_num times redundant memory traffic. The default few-query tiling is also suboptimal on HIP: BLOCK_M=64 wastes most rows when the verify length is only a few tokens, and waves_per_eu=1 does not hide memory latency well on long-prefix workloads.

This PR optimizes the AMD few-query GQA verify path by reducing redundant KV reads and improving GPU occupancy. On MI355X with MiniMax-M2.5, fp8 KV, and EAGLE3 verify shapes, the optimized path improves the extend-attention microbenchmark from about 1598us to 331us per call, and reduces long-context E2E median ITL from about 32.6ms to 14.2ms, while preserving the existing behavior for general extend/prefill paths.

Modifications

This PR updates python/sglang/srt/layers/attention/triton_ops/extend_attention.py.

The main changes are:

• Add HIP-specific few-query tuning for extend_attention_fwd. When running on HIP with max_len_extend <= 16, the kernel uses BLOCK_M=16, BLOCK_N=128, num_warps=4, and waves_per_eu=4, which better matches EAGLE3 / MTP verify shapes.

• Add a GQA-grouped split-K extend-attention path for AMD few-query verify workloads. The new _fwd_kernel_gqa_split kernel computes prefix attention by grouping all query heads that share the same KV head, so each KV head is loaded once instead of once per query head. The long prefix is split across multiple KV splits to preserve GPU occupancy.

• Add _fwd_kernel_gqa_combine to reduce the per-split online-softmax states and then compute the small stage-2 extend-token triangle before storing the final output.

• Add host-side dispatch and intermediate buffer allocation for the optimized GQA path. The new path is enabled only for HIP few-query GQA shapes without sliding window attention, xai temperature, or prefix custom masks. Other shapes continue to use the existing _fwd_kernel.

• Add environment knobs for benchmarking and retuning:

  • EXT_GQA_SPLITS to tune the number of KV splits
  • EXT_GQA_BLOCK_N to tune the prefix block size
  • EXT_GQA_WPE to tune HIP waves_per_eu

The CUDA path and unsupported attention feature combinations are left unchanged by the gating logic.

Accuracy Tests

Speed Tests and Profiling

off eagle

on eagle

on eagle + Optimize AMD GQA verify attention path

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #26392134415
Latest PR Test (Extra): ❌ Run #26392134359

[gemini-code-assist]

Code Review

This pull request introduces a GQA-grouped, split-K optimization for the extend_attention Triton kernels, specifically targeting 'few query token' verify shapes on HIP-based systems. The implementation adds two new kernels, _fwd_kernel_gqa_split and _fwd_kernel_gqa_combine, to improve HBM utilization. Review feedback highlights the need to align split boundaries to BLOCK_N to avoid undefined behavior, ensure numerical precision by using float32 accumulation in dot products, and optimize performance by using fixed block sizes to prevent kernel recompilation and pre-allocating intermediate buffers.

[gemini-code-assist]

The current calculation of split_len does not guarantee that split_start is a multiple of BLOCK_N. This violates the tl.multiple_of(start_n, BLOCK_N) hint on line 667, which can lead to undefined behavior, incorrect memory access, or compilation errors. The split boundaries should be aligned to BLOCK_N boundaries.

Suggested change

	split_len = tl.cdiv(cur_seq_len_prefix, NUM_KV_SPLITS)
	split_start = cur_split * split_len
	split_end = tl.minimum(split_start + split_len, cur_seq_len_prefix)
	num_blocks = tl.cdiv(cur_seq_len_prefix, BLOCK_N)
	blocks_per_split = tl.cdiv(num_blocks, NUM_KV_SPLITS)
	split_start = cur_split * blocks_per_split * BLOCK_N
	split_end = tl.minimum(split_start + blocks_per_split * BLOCK_N, cur_seq_len_prefix)

[gemini-code-assist]

The tl.dot call is missing out_dtype=tl.float32. For attention score computation, especially when using fp8 or fp16 inputs, accumulating in float32 is critical for precision and numerical stability. This is also inconsistent with the _fwd_kernel_gqa_combine kernel at line 891 which explicitly uses out_dtype=tl.float32.

Suggested change

	qk = tl.dot(q.to(k.dtype), k)
	qk = tl.dot(q.to(k.dtype), k, out_dtype=tl.float32)

[gemini-code-assist]

Similar to the QK dot product, the accumulation of acc should explicitly use out_dtype=tl.float32 to maintain precision during the reduction, especially when inputs are low-precision types like fp8 or fp16.

Suggested change

	acc = acc * re_scale[:, None] + tl.dot(p, v) * v_scale
	acc = acc * re_scale[:, None] + tl.dot(p, v, out_dtype=tl.float32) * v_scale

[gemini-code-assist]

BLOCK_M_PACKED depends on max_len_extend, which can vary between batches. Since BLOCK_M is a tl.constexpr in the Triton kernel, this will trigger a kernel recompilation whenever the maximum sequence length in a batch changes to a new power-of-2 bucket. Consider using a fixed BLOCK_M (e.g., 128) and masking out unused rows to avoid recompilation overhead during inference.

[gemini-code-assist]

Allocating a large intermediate buffer mid_o using torch.empty on every call to extend_attention_fwd can introduce significant overhead and memory fragmentation, especially in high-throughput verify workloads. It is highly recommended to use a pre-allocated workspace buffer managed by the attention backend to avoid repeated allocations.

[github-actions]

Thanks @fanxingran. Closing this because it has had no updates in 104 days.

Reopen it if the work is still relevant.

Some directories moved recently, so an older branch may need retargeting:
sgl-kernel/ -> python/sglang/kernels/aot/, python/sglang/jit_kernel/
-> python/sglang/kernels/jit/, docs/ -> docs/docs/ (.mdx),
bench_serving.py -> benchmark/serving.py, test/srt/ -> test/registered/.