From e3c4e4ab49db72c3e4f3c3d4bbcc11ec3a15f28d Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Fri, 13 Nov 2026 01:00:57 +0000 Subject: [PATCH 1/2] perf(attention): run up to 8 query rows per request in the triton decode kernel Speculative verify (MTP, DFlash) and the MTP draft extend sent q_len > 1 to the extend kernel, which streams the whole context per (request, query head) with no KV split: 6.6-60x the single-row decode time at 32K-128K on sm_120. - the grouped decode packs query rows x GQA heads in the M axis, limits each row to its own position and window, and keeps the KV split - the triton backend sends batches of at most 8 rows per request without multimodal blocks to it - a program's query heads divide the GQA group, which also fixes wrong output for groups above 16 that are not a multiple of 16 (e.g. 48/2) --- python/freetoken/attention/triton.py | 31 ++-- python/freetoken/kernel/triton/attention.py | 163 ++++++++++++++++---- tests/kernels/test_triton_attention.py | 66 +++++++- 3 files changed, 216 insertions(+), 44 deletions(-) diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index 179c57b5c..9739d53b0 100644 --- a/python/freetoken/attention/triton.py +++ b/python/freetoken/attention/triton.py @@ -5,7 +5,7 @@ import torch from freetoken.core import Batch, get_global_ctx -from freetoken.kernel.triton.attention import decode_kv_split_cap +from freetoken.kernel.triton.attention import MAX_DECODE_ROWS, decode_kv_split_cap from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata from .utils import BaseCaptureData @@ -92,6 +92,7 @@ def _ensure_decode_scratch( self, metadata: TritonMetadata, bs: int, + rows: int, num_q_heads: int, head_dim: int, ) -> None: @@ -99,7 +100,7 @@ def _ensure_decode_scratch( if ( metadata.attn_logits is not None and metadata.attn_lse is not None - and metadata.attn_logits.shape[0] >= bs + and metadata.attn_logits.shape[0] >= rows and metadata.attn_logits.shape[1] >= num_q_heads and metadata.attn_logits.shape[2] >= splits and metadata.attn_logits.shape[3] >= head_dim @@ -107,12 +108,12 @@ def _ensure_decode_scratch( return scratch_head_dim = max(self.max_head_dim, head_dim) metadata.attn_logits = torch.empty( - (bs, num_q_heads, splits, scratch_head_dim), + (rows, num_q_heads, splits, scratch_head_dim), dtype=torch.float32, device=self.device, ) metadata.attn_lse = torch.empty( - (bs, num_q_heads, splits), + (rows, num_q_heads, splits), dtype=torch.float32, device=self.device, ) @@ -154,9 +155,10 @@ def forward( if spec.sliding_window is not None and metadata.swa_indices is not None: indices = metadata.swa_indices scale = spec.sm_scale if spec.sm_scale is not None else q.shape[-1] ** -0.5 - if metadata.is_decode and q.dtype in (torch.float16, torch.bfloat16): + if metadata.max_q_len <= MAX_DECODE_ROWS and block_ends is None and q.dtype in (torch.float16, torch.bfloat16): bs = metadata.indptr.numel() - 1 - self._ensure_decode_scratch(metadata, bs, q.shape[1], q.shape[-1]) + rows = q.shape[0] + self._ensure_decode_scratch(metadata, bs, rows, q.shape[1], q.shape[-1]) assert metadata.attn_logits is not None assert metadata.attn_lse is not None return decode_paged_attention( @@ -166,11 +168,13 @@ def forward( indptr=metadata.indptr, indices=indices, q_positions=metadata.q_positions, - attn_logits=metadata.attn_logits[:bs], - attn_lse=metadata.attn_lse[:bs], + attn_logits=metadata.attn_logits[:rows], + attn_lse=metadata.attn_lse[:rows], sm_scale=scale, sliding_window=spec.sliding_window, sinks=spec.sinks, + qo_indptr=metadata.cu_seqlens_q_gpu, + max_q_len=metadata.max_q_len, **kv_args, ) if ( @@ -272,7 +276,7 @@ def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: self.device, num_q_heads=max(1, self.num_q_heads), max_head_dim=max(1, self.max_head_dim), - scratch_rows=max(bs * self._split_cap(bs) for bs in bs_list), + scratch_rows=max(bs * rows_per_req * self._split_cap(bs) for bs in bs_list), rows_per_req=rows_per_req, ) self.capture_rows_per_req = rows_per_req @@ -327,8 +331,9 @@ def _point_to_capture(self, metadata: TritonMetadata, bs: int) -> None: def _capture_scratch(self, bs: int) -> tuple[torch.Tensor, torch.Tensor]: assert self.capture is not None splits, heads, dim = self._split_cap(bs), max(1, self.num_q_heads), max(1, self.max_head_dim) - logits = self.capture.attn_logits[: bs * heads * splits * dim].view(bs, heads, splits, dim) - lse = self.capture.attn_lse[: bs * heads * splits].view(bs, heads, splits) + rows = bs * self.capture_rows_per_req + logits = self.capture.attn_logits[: rows * heads * splits * dim].view(rows, heads, splits, dim) + lse = self.capture.attn_lse[: rows * heads * splits].view(rows, heads, splits) return logits, lse def prepare_for_capture(self, batch: Batch) -> None: @@ -336,6 +341,7 @@ def prepare_for_capture(self, batch: Batch) -> None: assert bs in self.capture_bs and self.capture is not None capture = self.capture rows_per_req = self.capture_rows_per_req + attn_logits, attn_lse = self._capture_scratch(bs) if rows_per_req > 1: n = bs * rows_per_req # placeholders for a fresh request (kv rows == query rows, no prefix); replay overwrites both @@ -350,10 +356,11 @@ def prepare_for_capture(self, batch: Batch) -> None: is_decode=False, prefix_lens=capture.seq_lens[:bs], max_q_len=rows_per_req, + attn_logits=attn_logits, + attn_lse=attn_lse, swa_indices=self._capture_swa_indices(), ) return - attn_logits, attn_lse = self._capture_scratch(bs) batch.attn_metadata = TritonMetadata( cu_seqlens_q_gpu=capture.cu_seqlens_q[: bs + 1], indptr=capture.cu_seqlens_k[: bs + 1], diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index b1a4455a4..57ddaf71b 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -18,6 +18,8 @@ # a short request still takes up to this many splits of _MIN_BLOCK_KV keys, so tiny requests stay parallel _MIN_SPLITS = 8 _MAX_KV_SPLITS = 64 +# requests with more query rows take the extend kernel +MAX_DECODE_ROWS = 8 # stage 2 merges at most this many fp32 partial elements per step _STAGE2_TILE_ELEMS = 8192 @@ -49,6 +51,36 @@ def _quantized_decode_launch(block_d: int, block_n: int, num_stages: int) -> tup return block_n, num_stages, 4 +def _multi_row_decode_launch(m: int, block_d: int, quantized: bool, q_bytes: int, kv_bytes: int, smem: int) -> tuple[int, int, int]: + """(block_n, num_stages, num_warps) of stage 1 for ``m`` > 1 query rows x heads per program. + + On sm_120 quantized K/V spill at 4 warps and compute-dtype K/V at 8 (ptxas stops at 128 registers there); 2 stages never win. + """ + if quantized: + block_n, num_warps = (32 if block_d >= 512 else 64), 8 + else: + block_n, num_warps = (32 if block_d >= 512 or (block_d <= 128 and m >= 32) else 64), 4 + # the Q tile stays in the compute dtype, quantized K and V tiles take one byte per element + while smem and block_n > 16 and m * block_d * q_bytes + 2 * block_n * block_d * kv_bytes > smem: + block_n //= 2 + return block_n, 3, num_warps + + +def _decode_heads_per_program(group: int, rows: int, head_dim: int, quantized: bool) -> int: + """Query heads per decode program: the whole GQA group within the cap, else its largest power-of-two divisor within the cap. + + A program never spans two kv heads and a split group wastes no tile lanes. The cap is 16 heads and, for several rows, an M budget of rows x heads per ``head_dim``. + """ + cap = 16 + if rows > 1: + # quantized head_dim 512 spills above 8 on sm_120 + max_m = 64 if head_dim <= 128 else 32 if head_dim <= 256 else 8 if quantized else 16 + cap = min(16, max(1, max_m // rows)) + if group <= cap: + return group + return max(h for h in (1, 2, 4, 8, 16) if h <= cap and group % h == 0) + + def _max_request_splits(keys: int) -> int: return max(triton.cdiv(keys, _MIN_SPLIT_LEN), min(_MIN_SPLITS, triton.cdiv(keys, _MIN_BLOCK_KV)), 1) @@ -229,6 +261,7 @@ def _decode_grouped_stage1_kernel( sm_scale, indptr_ptr, indices_ptr, + qo_indptr_ptr, q_pos_ptr, mid_o_ptr, mid_lse_ptr, @@ -249,6 +282,7 @@ def _decode_grouped_stage1_kernel( stride_lse_s, GROUP: tl.constexpr, NUM_Q_HEADS: tl.constexpr, + ROWS: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_DV: tl.constexpr, BLOCK_N: tl.constexpr, @@ -267,14 +301,33 @@ def _decode_grouped_stage1_kernel( head_block_id = tl.program_id(1) split_id = tl.program_id(2) - # VALID_BLOCK_H == min(cap, GROUP) is the number of query heads actually handled per program; + # VALID_BLOCK_H, a divisor of GROUP of at most 16, is the number of query heads actually handled per program; # BLOCK_H is the power-of-two tile size for the head axis (tl.arange requires a power of two), # so a non-power-of-two GQA group (e.g. 24/4 == 6) rounds the tile up and masks the extra # lanes. Each kv head spans cdiv(GROUP, VALID_BLOCK_H) head blocks. + # The M axis is ROWS query rows x BLOCK_H heads, row-major. kv_head = head_block_id // tl.cdiv(GROUP, VALID_BLOCK_H) - q_heads = head_block_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H) - mask_h = q_heads < (head_block_id + 1) * VALID_BLOCK_H - mask_h = mask_h & (q_heads < NUM_Q_HEADS) + # one row keeps the head and row offsets in their single-query form: vector rows raise its register use + if ROWS == 1: + q_heads = head_block_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = q_heads < (head_block_id + 1) * VALID_BLOCK_H + mask_h = mask_h & (q_heads < NUM_Q_HEADS) + q_start = batch_id + q_row = batch_id * stride_qt + mid_row = batch_id * stride_mid_ob + lse_row = batch_id * stride_lse_b + else: + offs_m = tl.arange(0, ROWS * BLOCK_H) + rows = offs_m // BLOCK_H + lanes = offs_m % BLOCK_H + q_heads = head_block_id * VALID_BLOCK_H + lanes + q_start = tl.load(qo_indptr_ptr + batch_id) + q_len = tl.load(qo_indptr_ptr + batch_id + 1) - q_start + mask_h = (lanes < VALID_BLOCK_H) & (rows < q_len) + tokens = q_start + rows + q_row = tokens[:, None] * stride_qt + mid_row = tokens[:, None] * stride_mid_ob + lse_row = tokens * stride_lse_b offs_d = tl.arange(0, BLOCK_D) offs_dv = tl.arange(0, BLOCK_DV) @@ -283,8 +336,13 @@ def _decode_grouped_stage1_kernel( kv_start = tl.load(indptr_ptr + batch_id) kv_len = tl.load(indptr_ptr + batch_id + 1) - kv_start - q_pos = tl.load(q_pos_ptr + batch_id) - effective_end = tl.minimum(kv_len, q_pos + 1) + q_pos = tl.load(q_pos_ptr + q_start) + if ROWS == 1: + last_pos = q_pos + else: + last_pos = tl.load(q_pos_ptr + q_start + q_len - 1) + row_pos = tl.load(q_pos_ptr + q_start + rows, mask=rows < q_len, other=0) + effective_end = tl.minimum(kv_len, last_pos + 1) effective_start = 0 if SLIDING_WINDOW > 0: effective_start = tl.maximum(0, q_pos - SLIDING_WINDOW + 1) @@ -297,11 +355,11 @@ def _decode_grouped_stage1_kernel( split_start = kv_len_per_split * split_id split_end = tl.minimum(split_start + kv_len_per_split, effective_len) - m_i = tl.zeros((BLOCK_H,), dtype=tl.float32) - float("inf") - l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) - acc = tl.zeros((BLOCK_H, BLOCK_DV), dtype=tl.float32) + m_i = tl.zeros((ROWS * BLOCK_H,), dtype=tl.float32) - float("inf") + l_i = tl.zeros((ROWS * BLOCK_H,), dtype=tl.float32) + acc = tl.zeros((ROWS * BLOCK_H, BLOCK_DV), dtype=tl.float32) - q_offsets = batch_id * stride_qt + q_heads[:, None] * stride_qh + offs_d[None, :] + q_offsets = q_row + q_heads[:, None] * stride_qh + offs_d[None, :] k_base_offsets = kv_head * stride_kh + offs_d[:, None] v_base_offsets = kv_head * stride_vh + offs_dv[None, :] @@ -319,12 +377,22 @@ def _decode_grouped_stage1_kernel( k = kv_tile(k_ptr + slots[None, :] * stride_ks + k_base_offsets, mask_n[None, :] & mask_d[:, None], KV_FMT) scale_k = kv_row_scale(k_scale_ptr + slots * stride_kss + kv_head, mask_n, KV_FMT) scores = tl.dot(q, k) * sm_scale * scale_k[None, :] - scores = tl.where(mask_h[:, None] & mask_n[None, :], scores, -float("inf")) + mask = mask_h[:, None] & mask_n[None, :] + if ROWS > 1: + # the span covers every row: each row sees only keys up to its own position and inside its own window + mask = mask & (logical_offs[None, :] <= row_pos[:, None]) + if SLIDING_WINDOW > 0: + mask = mask & ((logical_offs[None, :] + SLIDING_WINDOW) > row_pos[:, None]) + scores = tl.where(mask, scores, -float("inf")) v = kv_tile(v_ptr + slots[:, None] * stride_vs + v_base_offsets, mask_n[:, None] & mask_dv[None, :], KV_FMT) scale_v = kv_row_scale(v_scale_ptr + slots * stride_vss + kv_head, mask_n, KV_FMT) - m_new = tl.maximum(tl.max(scores, axis=1), m_i) + row_max = tl.max(scores, axis=1) + if ROWS > 1: + # a row can see no key of a tile: -1e20 keeps exp() of its scores at 0 instead of nan + row_max = tl.where(row_max == -float("inf"), -1e20, row_max) + m_new = tl.maximum(row_max, m_i) alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) acc = acc * alpha[:, None] + tl.dot((p * scale_v[None, :]).to(v.dtype), v) @@ -332,8 +400,11 @@ def _decode_grouped_stage1_kernel( m_i = m_new out = acc / l_i[:, None] + if ROWS > 1: + # a row that sees no key of this split writes 0 with lse -inf, which stage 2 weighs by 0 + out = tl.where(l_i[:, None] > 0.0, out, 0.0) mid_offsets = ( - batch_id * stride_mid_ob + mid_row + q_heads[:, None] * stride_mid_oh + split_id * stride_mid_os + offs_dv[None, :] @@ -341,7 +412,7 @@ def _decode_grouped_stage1_kernel( tl.store(mid_o_ptr + mid_offsets, out, mask=mask_h[:, None] & mask_dv[None, :]) lse_offsets = ( - batch_id * stride_lse_b + lse_row + q_heads * stride_lse_h + split_id * stride_lse_s ) @@ -354,6 +425,7 @@ def _decode_stage2_kernel( mid_lse_ptr, o_ptr, indptr_ptr, + qo_indptr_ptr, q_pos_ptr, max_kv_splits, sinks_ptr, @@ -365,6 +437,7 @@ def _decode_stage2_kernel( stride_lse_s, stride_ot, stride_oh, + ROWS: tl.constexpr, MIN_BLOCK_KV: tl.constexpr, MIN_SPLIT_LEN: tl.constexpr, MIN_SPLITS: tl.constexpr, @@ -377,9 +450,22 @@ def _decode_stage2_kernel( batch_id = tl.program_id(0) q_head = tl.program_id(1) + if ROWS == 1: + token = batch_id + q_pos = tl.load(q_pos_ptr + batch_id) + last_pos = q_pos + else: + q_start = tl.load(qo_indptr_ptr + batch_id) + q_len = tl.load(qo_indptr_ptr + batch_id + 1) - q_start + row = tl.program_id(2) + if row >= q_len: + return + token = q_start + row + q_pos = tl.load(q_pos_ptr + q_start) + last_pos = tl.load(q_pos_ptr + q_start + q_len - 1) + kv_len = tl.load(indptr_ptr + batch_id + 1) - tl.load(indptr_ptr + batch_id) - q_pos = tl.load(q_pos_ptr + batch_id) - effective_end = tl.minimum(kv_len, q_pos + 1) + effective_end = tl.minimum(kv_len, last_pos + 1) effective_start = 0 if SLIDING_WINDOW > 0: effective_start = tl.maximum(0, q_pos - SLIDING_WINDOW + 1) @@ -401,8 +487,8 @@ def _decode_stage2_kernel( l_i = 0.0 acc = tl.zeros((BLOCK_DV,), dtype=tl.float32) - mid_base = batch_id * stride_mid_ob + q_head * stride_mid_oh - lse_base = batch_id * stride_lse_b + q_head * stride_lse_h + mid_base = token * stride_mid_ob + q_head * stride_mid_oh + lse_base = token * stride_lse_b + q_head * stride_lse_h # one load round trip per SPLIT_TILE splits for s0 in range(0, kv_splits, SPLIT_TILE): @@ -425,7 +511,7 @@ def _decode_stage2_kernel( out = tl.where(l_i == 0.0, 0.0, acc / l_i) tl.store( - o_ptr + batch_id * stride_ot + q_head * stride_oh + offs_d, + o_ptr + token * stride_ot + q_head * stride_oh + offs_d, out.to(o_ptr.dtype.element_ty), mask=mask_d, ) @@ -448,25 +534,33 @@ def decode_paged_attention( kv_format: KVFormat = NONE, k_scale: torch.Tensor | None = None, v_scale: torch.Tensor | None = None, + qo_indptr: torch.Tensor | None = None, + max_q_len: int = 1, ) -> torch.Tensor: - """SGLang-style split-k grouped decode attention for one query per request. + """SGLang-style split-k grouped decode attention for up to MAX_DECODE_ROWS causal query rows per request. - Each request takes ``max(ceil(keys / 128), min(8, ceil(keys / 32)))`` splits, capped by ``max_kv_splits`` (default :func:`decode_kv_split_cap`) and by the split dim of ``attn_logits`` / ``attn_lse``. + Request ``i`` owns query rows ``qo_indptr[i]:qo_indptr[i + 1]`` (one row each when ``qo_indptr`` is None); ``q_positions`` holds each row's key position. + Each request takes ``max(ceil(keys / 128), min(8, ceil(keys / 32)))`` splits, capped by ``max_kv_splits`` (default :func:`decode_kv_split_cap`) and by the split dim of ``attn_logits`` / ``attn_lse``, which hold one entry per query row. The kernels compute the split count on the device, so a captured graph serves any context length. """ assert q.is_cuda and k_cache.is_cuda and v_cache.is_cuda assert q.dim() == 3 and k_cache.dim() == 3 and v_cache.dim() == 3 - batch, num_q_heads, head_dim = q.shape + num_rows, num_q_heads, head_dim = q.shape + batch = indptr.numel() - 1 num_kv_heads = k_cache.shape[1] - assert batch == indptr.numel() - 1 + if qo_indptr is None: + assert num_rows == batch and max_q_len == 1 + else: + assert qo_indptr.numel() == batch + 1 + assert 1 <= max_q_len <= MAX_DECODE_ROWS assert v_cache.shape[1] == num_kv_heads assert k_cache.shape[-1] == head_dim and v_cache.shape[-1] == head_dim assert num_q_heads % num_kv_heads == 0 - assert attn_logits.shape[0] >= batch + assert attn_logits.shape[0] >= num_rows assert attn_logits.shape[1] >= num_q_heads assert attn_logits.shape[3] >= head_dim - assert attn_lse.shape[0] >= batch + assert attn_lse.shape[0] >= num_rows assert attn_lse.shape[1] >= num_q_heads if sliding_window is not None and sliding_window <= 0: # the kernels treat a window <= 0 as no window @@ -474,8 +568,9 @@ def decode_paged_attention( if max_kv_splits is None: max_kv_splits = decode_kv_split_cap(batch, num_q_heads, num_kv_heads, q.device.index) if sliding_window: - # no request attends more than the window, so programs past its splits would only exit - max_kv_splits = min(max_kv_splits, _max_request_splits(sliding_window)) + # no request attends more than the window plus its later rows, so programs past its splits would only exit + span = sliding_window + max_q_len - 1 + max_kv_splits = min(max_kv_splits, _max_request_splits(span)) max_kv_splits = min(max_kv_splits, attn_logits.shape[2], attn_lse.shape[2]) if sinks is not None: assert sinks.is_cuda @@ -485,11 +580,13 @@ def decode_paged_attention( o = out if out is not None else torch.empty_like(q) sinks_arg = sinks if sinks is not None else q + qo_indptr_arg = qo_indptr if qo_indptr is not None else indptr + rows = triton.next_power_of_2(max_q_len) group = num_q_heads // num_kv_heads # valid_block_h = heads computed per program (drives the grid + head indexing); block_h = # power-of-two tile size for tl.arange. They differ only for non-power-of-two GQA groups # (e.g. 6), where block_h rounds up and the kernel masks the extra lanes. - valid_block_h = min(16, group) + valid_block_h = _decode_heads_per_program(group, rows, head_dim, kv_format.quantized) block_h = triton.next_power_of_2(valid_block_h) block_d = triton.next_power_of_2(head_dim) block_dv = triton.next_power_of_2(head_dim) @@ -514,6 +611,10 @@ def decode_paged_attention( num_warps = 4 if kv_format.quantized: block_n, num_stages, num_warps = _quantized_decode_launch(block_d, block_n, num_stages) + if rows > 1: + block_n, num_stages, num_warps = _multi_row_decode_launch( + rows * block_h, block_d, kv_format.quantized, elem, k_cache.element_size(), smem + ) k_scale_arg, v_scale_arg, stride_kss, stride_vss = kv_scale_args(kv_format, k_cache, v_cache, k_scale, v_scale, q) _decode_grouped_stage1_kernel[grid]( @@ -525,6 +626,7 @@ def decode_paged_attention( sm_scale, indptr, indices, + qo_indptr_arg, q_positions, attn_logits, attn_lse, @@ -547,6 +649,7 @@ def decode_paged_attention( KV_FMT=kv_format.name, GROUP=group, NUM_Q_HEADS=num_q_heads, + ROWS=rows, BLOCK_D=block_d, BLOCK_DV=block_dv, BLOCK_N=block_n, @@ -562,11 +665,12 @@ def decode_paged_attention( num_stages=num_stages, ) split_tile, stage2_warps = _stage2_launch(max_kv_splits, block_dv) - _decode_stage2_kernel[(batch, num_q_heads)]( + _decode_stage2_kernel[(batch, num_q_heads, rows)]( attn_logits, attn_lse, o, indptr, + qo_indptr_arg, q_positions, max_kv_splits, sinks_arg, @@ -578,6 +682,7 @@ def decode_paged_attention( attn_lse.stride(2), o.stride(0), o.stride(1), + ROWS=rows, MIN_BLOCK_KV=_MIN_BLOCK_KV, MIN_SPLIT_LEN=_MIN_SPLIT_LEN, MIN_SPLITS=_MIN_SPLITS, diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index ad1120181..841bf4533 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -167,6 +167,13 @@ def _attend(path: str, q, k_cache, v_cache, indices, kv_args: dict, k_extend=Non lse = torch.empty(1, num_q_heads, 64, device=device, dtype=torch.float32) return decode_paged_attention(q, k_cache, v_cache, indptr, indices, q_positions, mid, lse, sm_scale, **kv_args) qo_indptr = torch.tensor([0, q_len], device=device, dtype=torch.int32) + if path == "decode_rows": + mid = torch.empty(q_len, num_q_heads, 64, head_dim, device=device, dtype=torch.float32) + lse = torch.empty(q_len, num_q_heads, 64, device=device, dtype=torch.float32) + return decode_paged_attention( + q, k_cache, v_cache, indptr, indices, q_positions, mid, lse, sm_scale, + qo_indptr=qo_indptr, max_q_len=q_len, **kv_args, + ) prefix_lens = torch.tensor([count - q_len], device=device, dtype=torch.int32) return extend_paged_attention( q, k_cache, v_cache, qo_indptr, indptr, indices, prefix_lens, q_len, sm_scale, @@ -206,14 +213,14 @@ def test_store_kv_rows_round_trips_within_the_format_error_bound(fmt_name: str): @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") @pytest.mark.parametrize("fmt_name", KV_QUANT_FORMATS) -@pytest.mark.parametrize("path", ["paged", "decode", "extend", "split_extend"]) +@pytest.mark.parametrize("path", ["paged", "decode", "decode_rows", "extend", "split_extend"]) @pytest.mark.parametrize("head_dim", [64, 128, 256, 512]) def test_quantized_kv_attention_matches_dequantized_reference(fmt_name: str, path: str, head_dim: int): torch.manual_seed(11) device = torch.device("cuda") num_kv_heads, num_q_heads = 2, 4 count, capacity = 66, 200 - q_len = 2 if "extend" in path else 1 + q_len = {"paged": 1, "decode": 1, "decode_rows": 3}.get(path, 2) row = num_kv_heads * head_dim # rows and heads of different magnitude, so a scale read from the wrong slot or head changes the result head_gain = torch.ones(row, device=device) @@ -564,7 +571,7 @@ def test_decode_triton_attention_matches_reference( @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") -@pytest.mark.parametrize(("num_q_heads", "num_kv_heads"), [(24, 4), (20, 4), (28, 4)]) +@pytest.mark.parametrize(("num_q_heads", "num_kv_heads"), [(24, 4), (20, 4), (28, 4), (48, 2)]) def test_decode_triton_attention_non_pow2_group(num_q_heads: int, num_kv_heads: int): """GQA groups that are not a power of two (e.g. Qwen3.6-27B's 24/4 == 6). The grouped decode tiles the head axis to a power of two (tl.arange constraint) and masks the extra @@ -714,6 +721,59 @@ def test_decode_triton_attention_splits_long_requests( torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") +@pytest.mark.parametrize( + ("num_q_heads", "num_kv_heads", "head_dim", "cached_lens", "q_lens", "sliding_window", "scratch_splits", "use_sinks"), + [ + (16, 2, 256, [0, 129, 1017, 8193], [4, 1, 8, 3], None, 64, False), + (16, 2, 256, [0, 129, 1017, 8193], [4, 1, 8, 3], None, 8, False), + (16, 2, 256, [300, 2000, 70], [4, 8, 2], 3, 64, True), + (24, 4, 256, [500, 33], [5, 2], 48, 64, False), + (8, 1, 512, [700, 64], [8, 8], None, 64, True), + ], +) +def test_decode_triton_attention_multi_row_matches_reference( + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + cached_lens: list[int], + q_lens: list[int], + sliding_window: int | None, + scratch_splits: int, + use_sinks: bool, +): + """Several causal query rows per request: rows that see no key of a split (1017 + 8 keys leave one key in the last split), a window shorter than the rows, short requests next to many-split ones, and fewer heads per program at head_dim 512.""" + from freetoken.kernel.triton.attention import decode_paged_attention + + torch.manual_seed(7) + device = torch.device("cuda") + kv_lens = [c + r for c, r in zip(cached_lens, q_lens)] + total_kv = sum(kv_lens) + num_rows = sum(q_lens) + q = torch.randn(num_rows, num_q_heads, head_dim, device=device, dtype=torch.bfloat16) + k_cache = torch.randn(total_kv, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16) + v_cache = torch.randn(total_kv, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16) + indptr = torch.tensor([0, *kv_lens], device=device).cumsum(0).to(torch.int32) + qo_indptr = torch.tensor([0, *q_lens], device=device).cumsum(0).to(torch.int32) + indices = torch.randperm(total_kv, device=device).to(torch.int32) + q_positions = torch.cat([torch.arange(c, c + r, device=device) for c, r in zip(cached_lens, q_lens)]) + q_to_req = torch.repeat_interleave(torch.arange(len(q_lens), device=device, dtype=torch.int32), torch.tensor(q_lens, device=device)) + sinks = torch.randn(num_q_heads, dtype=torch.float32, device=device) if use_sinks else None + sm_scale = head_dim**-0.5 + attn_logits = torch.empty(num_rows, num_q_heads, scratch_splits, head_dim, dtype=torch.float32, device=device) + attn_lse = torch.empty(num_rows, num_q_heads, scratch_splits, dtype=torch.float32, device=device) + + actual = decode_paged_attention( + q, k_cache, v_cache, indptr, indices, q_positions, attn_logits, attn_lse, sm_scale, + sliding_window=sliding_window, sinks=sinks, qo_indptr=qo_indptr, max_q_len=max(q_lens), + ) + expected = _reference_paged_attention( + q, k_cache, v_cache, indptr, indices, q_to_req, q_positions, sm_scale, sliding_window, sinks=sinks, + ) + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") @pytest.mark.parametrize("use_split_inputs", [False, True]) @pytest.mark.parametrize( From 7e353d11cd370e55bc5abbb07587f52a8ac12a6c Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Fri, 13 Nov 2026 01:52:17 +0000 Subject: [PATCH 2/2] fix(attention): keep multi-row triton decode free of register spills on sm_80 to sm_120 - pass maxnreg=255 to multi-row stage-1 launches on CUDA: without .maxnreg ptxas holds these tiles at 80-128 registers for occupancy and spills 8-16 B on sm_80 to sm_120 - pick block_n and warps from a launch sweep under maxnreg, among the variants that do not spill on sm_80, 86, 89, 90, 100 and 120; quantized head_dim 512 takes 16 rows x heads again --- python/freetoken/kernel/triton/attention.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index 57ddaf71b..f6a84692e 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -8,6 +8,7 @@ from freetoken.kernel.triton.kv_rows import kv_row_scale, kv_scale_args, kv_tile, kv_tile_staged from freetoken.kvcache.format import NONE, KVFormat +from freetoken.utils.arch import is_rocm _MIN_BLOCK_KV = 32 @@ -54,27 +55,26 @@ def _quantized_decode_launch(block_d: int, block_n: int, num_stages: int) -> tup def _multi_row_decode_launch(m: int, block_d: int, quantized: bool, q_bytes: int, kv_bytes: int, smem: int) -> tuple[int, int, int]: """(block_n, num_stages, num_warps) of stage 1 for ``m`` > 1 query rows x heads per program. - On sm_120 quantized K/V spill at 4 warps and compute-dtype K/V at 8 (ptxas stops at 128 registers there); 2 stages never win. + Picked from a launch sweep on sm_120 among the variants that do not spill on sm_80 to sm_120 (with maxnreg); 2 stages never win. """ - if quantized: - block_n, num_warps = (32 if block_d >= 512 else 64), 8 + if quantized and block_d >= 512: + block_n, num_warps = 32, 8 else: - block_n, num_warps = (32 if block_d >= 512 or (block_d <= 128 and m >= 32) else 64), 4 + block_n, num_warps = (64 if m * block_d <= 4096 and not (quantized and block_d >= 256) else 32), 4 # the Q tile stays in the compute dtype, quantized K and V tiles take one byte per element while smem and block_n > 16 and m * block_d * q_bytes + 2 * block_n * block_d * kv_bytes > smem: block_n //= 2 return block_n, 3, num_warps -def _decode_heads_per_program(group: int, rows: int, head_dim: int, quantized: bool) -> int: +def _decode_heads_per_program(group: int, rows: int, head_dim: int) -> int: """Query heads per decode program: the whole GQA group within the cap, else its largest power-of-two divisor within the cap. A program never spans two kv heads and a split group wastes no tile lanes. The cap is 16 heads and, for several rows, an M budget of rows x heads per ``head_dim``. """ cap = 16 if rows > 1: - # quantized head_dim 512 spills above 8 on sm_120 - max_m = 64 if head_dim <= 128 else 32 if head_dim <= 256 else 8 if quantized else 16 + max_m = 64 if head_dim <= 128 else 32 if head_dim <= 256 else 16 cap = min(16, max(1, max_m // rows)) if group <= cap: return group @@ -586,7 +586,7 @@ def decode_paged_attention( # valid_block_h = heads computed per program (drives the grid + head indexing); block_h = # power-of-two tile size for tl.arange. They differ only for non-power-of-two GQA groups # (e.g. 6), where block_h rounds up and the kernel masks the extra lanes. - valid_block_h = _decode_heads_per_program(group, rows, head_dim, kv_format.quantized) + valid_block_h = _decode_heads_per_program(group, rows, head_dim) block_h = triton.next_power_of_2(valid_block_h) block_d = triton.next_power_of_2(head_dim) block_dv = triton.next_power_of_2(head_dim) @@ -663,6 +663,8 @@ def decode_paged_attention( SLIDING_WINDOW=sliding_window or 0, num_warps=num_warps, num_stages=num_stages, + # without .maxnreg ptxas holds multi-row tiles below their register need for occupancy and spills a few bytes + **({} if rows == 1 or is_rocm() else {"maxnreg": 255}), ) split_tile, stage2_warps = _stage2_launch(max_kv_splits, block_dv) _decode_stage2_kernel[(batch, num_q_heads, rows)](