Repository navigation
perf(kernel): split triton decode attention by context length and sm count - #729
Conversation
- stage 2 loads SPLIT_TILE splits per step and merges them with one online-softmax update instead of one split per iteration - SPLIT_TILE is the launch's split cap rounded up to a power of two, at most 8192 fp32 elements and 64 splits; num_warps scales with it
|
Pushed 2bd6c33: stage 2 now merges a tile of splits per step instead of one split per iteration. WhyI measured this PR on four consumer and workstation cards: RTX PRO 6000 Blackwell Server (sm_120, 188 SMs), RTX 5090 (170), RTX 4090 (128) and RTX 3080 Ti (80). On all four, stage 1 stops getting faster at 16 to 32 splits. Stage 2 runs one program per (request, query head), and that program loads one split per iteration. Each iteration waits for the previous load to return, so at batch 1 the merge costs about 0.21 us per split. With 64 splits it takes 13.5 us on the PRO 6000, which is half of stage 1's time at 1 x 8k. flashinfer's merge kernel handles 188 chunks in 5 to 9 us, and its attention kernel takes about the same time as this PR's stage 1. So the gap to flashinfer was the merge alone. Change
Kernel results
The worst case, 1.023, is the 32/8/128 shape at 1 x 300 on the 4090, which is 0.2 us slower. Qwen3.6 full-attention shape against flashinfer, as the
Before this commit, triton was 11 to 23% slower than flashinfer at 1 x 8k on the PRO 6000, 5090 and 4090. Now it is faster at batch 1 on all four cards and within 4% at batch 4 and above. The fa and trtllm backends do not run on these GPUs. End-to-endQwen3.6-35B-A3B-NVFP4 with
Other checks
Setup: driver 595.91 (PRO 6000), 595.71 (5090, 3080 Ti) and 580.76 (4090); torch 2.14.1+cu130, Triton 3.8.0, flashinfer 0.6.18.post1. Not in this commitTwo things still differ from the best config in the sweep, on all four cards:
One rule fits both on all four cards: about 2 programs per SM, at most 32 splits, |
RTX 4060 Laptop (sm_89, 24 SMs)Kernel only. Qwen3.6 full-attention shape 16/2/256, bf16, batch 1, us, one CUDA graph replay after an L2 flush, median.
|
Upstream 0bca88f..11c1adb: GGUF on vendored ggml (FlashML-org#690 and the per-arch loaders), MoE slot pools (FlashML-org#689, FlashML-org#700), GGUF experts on the cpu/hybrid executor (FlashML-org#704), the NVFP4 MoE prefill GEMM rewrite (FlashML-org#723), triton decode split-k by context length (FlashML-org#729), graph upload at capture (FlashML-org#697), transformers 5.18. Resolution: - GGUF: upstream's stack replaces ours (dense llama/qwen2/qwen3, qwen3moe, qwen35moe loaders, q4_0/gguf expert banks, our ggml kernels and their tests). Re-ported on top: rope-only --hf-overrides (YaRN) for the llama/qwen2/qwen3/qwen3moe GGUF parsers, and the scalar head_count_kv of Gemma-4 GGUFs (upstream PR FlashML-org#190). Gemma-4 E2B/E4B GGUF is not re-ported. - offload cache: upstream's pooled cache, with our disk tier (one pool only), fill_slots preload per pool, sync-free prefill invalidation for a contiguous staging range, hit-D2D-by-default wording and the HIP graph per-bank copy. - NVFP4 prefill: upstream's GEMM and per-GPU tuning supersede our sm_89 tile config and JOIN_B. - triton decode attention: upstream's split-k over our fp8/nvfp4 KV paths and HIP head padding; DFlash verify graphs are uploaded at capture too. - embedding: upstream's embedding quant method, with --embed-device cpu in front and the tied head's kernel choice kept for unquantized tables. - engine: DFlash, --kv-reserve-share, shrink-first rebuild, startup hybrid split and the WSL pin probe combined with the slot pools. GPU-checked on the RTX 4090: full suite green; Qwen3.6-35B-A3B NVFP4 prefill ~13.2k / decode ~131 tok/s (baseline 13.6k / 131), Qwen3-30B-A3B Q4_K_M GGUF 14.6k / 259 (was 8.3k / 225). Assisted-by: Claude Opus 5.5
…r RTX), PR FlashML-org#729 (split decode attention), PR FlashML-org#726, PR FlashML-org#700 and GGUF updates
1 participant
Summary
Decode attention reads the whole KV cache for one query per request, so it is bound by memory bandwidth and needs enough parallel programs to keep every SM busy. The Triton kernel gets that parallelism by splitting each request's keys into chunks that run in parallel and merging them afterwards (split-K), but it always used 8 chunks. Layers with few KV heads then ran on a small part of the GPU (16 programs on a 132-SM H100 for Gemma 4's full-attention layers at batch 1), and decode slowed down sharply with context: Gemma 4 26B-A4B decodes at 204 tok/s at 1k but 77 tok/s at 128k.
This PR sizes the chunks from the request length and the GPU. Batch-1 decode at 128k goes from 77 to 174 tok/s on Gemma 4, 90 to 213 on gpt-oss-20b and 67 to 196 on Qwen3.6-35B-A3B with
--attention-backend triton, which puts triton within 2% of flashinfer and FA3 on Qwen3.6 from 1k to 192k. At 1k all three stay within 1.5% of main.Design
The kernel is flash decoding: stage 1 cuts each request's keys into chunks and computes, in parallel, a partial attention output and its log-sum-exp for each chunk; stage 2 merges the chunks. This PR changes how many chunks a request gets, taking the smaller of two counts:
To work under CUDA graphs, which fix the grid at capture, the host only sizes the grid for the most chunks a request can take at that batch size. Each program then reads its request's current length from the graph's input buffers, computes the request's chunk count and chunk length, and exits if its chunk is past the end; stage 2 runs the same formula to merge only those chunks. One captured graph therefore serves every context length, and short requests in a batch stay coarse. For Qwen3.6's full-attention layers at batch 1 on H100, a 32k request now takes 64 chunks of 512 keys instead of 8 chunks of 4096.
A batch of four requests of different lengths on the same layer (batch 4, so the grid has 64 chunk slots per request;
#is a chunk that runs,.a program that exits at once):With few programs per SM, each program has to keep more memory requests in flight: the key loop pipelines 3 steps instead of 2, and when the grid has at most one program per SM, each step reads 64 keys instead of 32 (skipped where it would not fit in shared memory).
Stage 2 merges a tile of up to 64 chunks per step instead of one chunk per loop iteration, so its loads no longer wait on each other one chunk at a time. The tile size follows the launch's chunk cap.
The backend sizes its split-K scratch to match, and under CUDA graphs all batch sizes share one buffer. A new test covers long requests next to short ones and sliding windows that span several chunks.
Tested on
NVIDIA H100 80GB HBM3 (132 SMs), Intel Xeon Platinum 8480+, driver 580.95.05, CUDA 13.0, torch 2.14.1+cu130, Triton 3.8.0, flashinfer 0.6.18.post1.
mainis bc0e2fe, the base of this branch.End-to-end results
Models: nvidia/Gemma-4-26B-A4B-NVFP4 and openai/gpt-oss-20b, which use Triton attention by default, and nvidia/Qwen3.6-35B-A3B-NVFP4 with
--attention-backend triton.The client sends one streamed
/v1/chat/completionsrequest at a time: a prompt of about 1k to 192k tokens (source code, cut with each model's tokenizer), 256 output tokens withignore_eos, temperature 1.0, top_p 0.95, top_k 64. Decode tok/s is (output tokens - 1) / (last token time - first token time). Median of 3 runs. TTFT is the same on both sides: these prompts are prefilled by the extend kernel, which this PR does not touch.These runs are from 855aedd, before stage 2 merged a tile of chunks per step. The RTX results comment has an end-to-end run of the final commit on an RTX PRO 6000.
Decode tok/s, batch 1, main -> this PR:
Qwen3.6 against the other attention backends, decode tok/s, batch 1:
Kernel results
decode_paged_attentionalone, bf16, one CUDA graph replay after an L2 flush, median us (main -> this PR). Shapes are query heads / KV heads / head_dim.Qwen3.6 full-attention shape against the other backends, median us:
Other checks
PYTHONPATH=python pytest tests/kernels/test_triton_attention.py: 41 passed; withCUDA_VISIBLE_DEVICES=, 11 passed and 30 skipped. Same counts on the final commit, andpytest tests/attentionpasses (15).