diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index 30bedf4ff..7f29c481b 100644 --- a/python/freetoken/attention/triton.py +++ b/python/freetoken/attention/triton.py @@ -5,6 +5,7 @@ import torch from freetoken.core import Batch, get_global_ctx +from freetoken.kernel.triton.attention import decode_kv_split_cap from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata from .utils import BaseCaptureData @@ -16,9 +17,9 @@ @dataclass class TritonCaptureData(BaseCaptureData): q_to_req: torch.Tensor + # flat decode split scratch: each graph batch size views it with its own split count attn_logits: torch.Tensor attn_lse: torch.Tensor - num_kv_splits: torch.Tensor swa_page_table: torch.Tensor | None = None @classmethod @@ -30,7 +31,7 @@ def create( *, num_q_heads: int, max_head_dim: int, - max_kv_splits: int, + scratch_rows: int, **kwargs, ): return cls( @@ -40,22 +41,8 @@ def create( cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device), page_table=torch.zeros((max_bs, max_seq_len), dtype=torch.int32, device=device), q_to_req=torch.arange(max_bs, dtype=torch.int32, device=device), - attn_logits=torch.empty( - (max_bs, num_q_heads, max_kv_splits, max_head_dim), - dtype=torch.float32, - device=device, - ), - attn_lse=torch.empty( - (max_bs, num_q_heads, max_kv_splits), - dtype=torch.float32, - device=device, - ), - num_kv_splits=torch.full( - (max_bs,), - max_kv_splits, - dtype=torch.int32, - device=device, - ), + attn_logits=torch.empty((scratch_rows * num_q_heads * max_head_dim,), dtype=torch.float32, device=device), + attn_lse=torch.empty((scratch_rows * num_q_heads,), dtype=torch.float32, device=device), **kwargs, ) @@ -72,7 +59,6 @@ class TritonMetadata(BaseAttnMetadata): max_q_len: int attn_logits: torch.Tensor | None = None attn_lse: torch.Tensor | None = None - num_kv_splits: torch.Tensor | None = None swa_indices: torch.Tensor | None = None def get_last_indices(self, bs: int) -> torch.Tensor: @@ -87,7 +73,6 @@ def __init__(self, config: ModelConfig): self.capture: TritonCaptureData | None = None self.capture_bs: List[int] = [] self.max_graph_bs = 0 - self.max_kv_splits = 8 self.prefill_tile_min_q = 128 self.num_q_heads = int(getattr(config, "num_qo_heads", 1)) kv_groups = getattr(config, "kv_cache_group_specs", lambda: ())() @@ -95,6 +80,11 @@ def __init__(self, config: ModelConfig): (group.head_dim for group in kv_groups), default=int(getattr(config, "head_dim", 1)), ) + self.kv_heads = sorted({group.num_kv_heads for group in kv_groups}) or [int(getattr(config, "num_kv_heads", 1))] + + def _split_cap(self, bs: int) -> int: + """The most decode KV splits any layer group takes at this batch size.""" + return max(decode_kv_split_cap(bs, max(1, self.num_q_heads), heads, self.device.index) for heads in self.kv_heads) def _ensure_decode_scratch( self, @@ -103,32 +93,27 @@ def _ensure_decode_scratch( num_q_heads: int, head_dim: int, ) -> None: + splits = self._split_cap(bs) if ( metadata.attn_logits is not None and metadata.attn_lse is not None - and metadata.num_kv_splits is not None and metadata.attn_logits.shape[0] >= bs and metadata.attn_logits.shape[1] >= num_q_heads + and metadata.attn_logits.shape[2] >= splits and metadata.attn_logits.shape[3] >= head_dim ): return scratch_head_dim = max(self.max_head_dim, head_dim) metadata.attn_logits = torch.empty( - (bs, num_q_heads, self.max_kv_splits, scratch_head_dim), + (bs, num_q_heads, splits, scratch_head_dim), dtype=torch.float32, device=self.device, ) metadata.attn_lse = torch.empty( - (bs, num_q_heads, self.max_kv_splits), + (bs, num_q_heads, splits), dtype=torch.float32, device=self.device, ) - metadata.num_kv_splits = torch.full( - (bs,), - self.max_kv_splits, - dtype=torch.int32, - device=self.device, - ) def forward( self, @@ -167,7 +152,6 @@ def forward( self._ensure_decode_scratch(metadata, bs, q.shape[1], q.shape[-1]) assert metadata.attn_logits is not None assert metadata.attn_lse is not None - assert metadata.num_kv_splits is not None return decode_paged_attention( q=q, k_cache=k_cache, @@ -177,8 +161,6 @@ def forward( q_positions=metadata.q_positions, attn_logits=metadata.attn_logits[:bs], attn_lse=metadata.attn_lse[:bs], - num_kv_splits=metadata.num_kv_splits[:bs], - max_kv_splits=self.max_kv_splits, sm_scale=scale, sliding_window=spec.sliding_window, sinks=spec.sinks, @@ -280,7 +262,7 @@ def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: self.device, num_q_heads=max(1, self.num_q_heads), max_head_dim=max(1, self.max_head_dim), - max_kv_splits=self.max_kv_splits, + scratch_rows=max(bs * self._split_cap(bs) for bs in bs_list), ) if self._swa_capture_enabled(): self.capture.swa_page_table = torch.zeros( @@ -324,14 +306,20 @@ def _point_to_capture(self, metadata: TritonMetadata, bs: int) -> None: metadata.indices = indices metadata.q_to_req = self.capture.q_to_req[: metadata.q_to_req.numel()] metadata.q_positions = self.capture.positions[:q_tokens] - metadata.attn_logits = self.capture.attn_logits[:bs] - metadata.attn_lse = self.capture.attn_lse[:bs] - metadata.num_kv_splits = self.capture.num_kv_splits[:bs] + metadata.attn_logits, metadata.attn_lse = self._capture_scratch(bs) + + 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) + return logits, lse def prepare_for_capture(self, batch: Batch) -> None: bs = batch.size assert bs in self.capture_bs and self.capture is not None capture = self.capture + 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], @@ -341,9 +329,8 @@ def prepare_for_capture(self, batch: Batch) -> None: is_decode=True, prefix_lens=capture.positions[:bs], max_q_len=1, - attn_logits=capture.attn_logits[:bs], - attn_lse=capture.attn_lse[:bs], - num_kv_splits=capture.num_kv_splits[:bs], + attn_logits=attn_logits, + attn_lse=attn_lse, swa_indices=self._capture_swa_indices(), ) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index c9087e03c..d19457cff 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -7,10 +7,28 @@ import triton.language as tl -_MAX_KV_SPLITS = 8 _MIN_BLOCK_KV = 32 +# a long decode request takes splits of at least this many keys, so the partial result each one writes stays small next to the keys it reads +_MIN_SPLIT_LEN = 128 +# 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 +# stage 2 merges at most this many fp32 partial elements per step +_STAGE2_TILE_ELEMS = 8192 + + +def _stage2_launch(max_kv_splits: int, block_dv: int) -> tuple[int, int]: + """(splits merged per step, num_warps) for stage 2; the tile is at most the split cap rounded up to a power of two.""" + tile = min(triton.next_power_of_2(max_kv_splits), max(4, _STAGE2_TILE_ELEMS // block_dv), 64) + return tile, max(1, min(8, tile * block_dv // 2048)) + + +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) + + @functools.lru_cache(maxsize=None) def _optin_smem_bytes(device_index: int) -> int: """Per-block opt-in shared-memory budget for a CUDA device (0 if unavailable).""" @@ -23,6 +41,12 @@ def _num_sms(device_index: int) -> int: return int(torch.cuda.get_device_properties(device_index).multi_processor_count) +def decode_kv_split_cap(batch: int, num_q_heads: int, num_kv_heads: int, device_index: int) -> int: + """Most KV splits per request for a decode batch: about four programs per SM, within [_MIN_SPLITS, _MAX_KV_SPLITS].""" + head_blocks = triton.cdiv(num_q_heads, min(16, max(1, num_q_heads // num_kv_heads))) + return max(_MIN_SPLITS, min(_MAX_KV_SPLITS, triton.cdiv(4 * _num_sms(device_index), batch * head_blocks))) + + def _select_extend_tile(head_dim: int, block_d: int, smem_optin: int) -> tuple[int, int]: """Pick ``(BLOCK_M, BLOCK_N)`` for the extend/prefill kernel, shared-memory aware. @@ -150,6 +174,13 @@ def _paged_attention_kernel( @triton.jit +def _decode_num_splits(effective_len, max_kv_splits, MIN_SPLIT_LEN: tl.constexpr, MIN_SPLITS: tl.constexpr, MIN_BLOCK_KV: tl.constexpr): + # same rule as _max_request_splits: long requests split by length, short ones still take up to MIN_SPLITS small splits + by_len = tl.maximum(tl.cdiv(effective_len, MIN_SPLIT_LEN), tl.minimum(MIN_SPLITS, tl.cdiv(effective_len, MIN_BLOCK_KV))) + return tl.minimum(tl.maximum(by_len, 1), max_kv_splits) + + +@triton.jit(do_not_specialize=["max_kv_splits"]) def _decode_grouped_stage1_kernel( q_ptr, k_ptr, @@ -160,7 +191,7 @@ def _decode_grouped_stage1_kernel( q_pos_ptr, mid_o_ptr, mid_lse_ptr, - num_kv_splits_ptr, + max_kv_splits, stride_qt, stride_qh, stride_ks, @@ -181,6 +212,8 @@ def _decode_grouped_stage1_kernel( BLOCK_H: tl.constexpr, VALID_BLOCK_H: tl.constexpr, MIN_BLOCK_KV: tl.constexpr, + MIN_SPLIT_LEN: tl.constexpr, + MIN_SPLITS: tl.constexpr, D: tl.constexpr, DV: tl.constexpr, SLIDING_WINDOW: tl.constexpr, @@ -212,7 +245,7 @@ def _decode_grouped_stage1_kernel( effective_start = tl.maximum(0, q_pos - SLIDING_WINDOW + 1) effective_len = tl.maximum(0, effective_end - effective_start) - kv_splits = tl.load(num_kv_splits_ptr + batch_id) + kv_splits = _decode_num_splits(effective_len, max_kv_splits, MIN_SPLIT_LEN, MIN_SPLITS, MIN_BLOCK_KV) kv_len_per_split = ( tl.cdiv(tl.cdiv(effective_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV ) @@ -275,14 +308,14 @@ def _decode_grouped_stage1_kernel( tl.store(mid_lse_ptr + lse_offsets, m_i + tl.log(l_i), mask=mask_h) -@triton.jit +@triton.jit(do_not_specialize=["max_kv_splits"]) def _decode_stage2_kernel( mid_o_ptr, mid_lse_ptr, o_ptr, indptr_ptr, q_pos_ptr, - num_kv_splits_ptr, + max_kv_splits, sinks_ptr, stride_mid_ob, stride_mid_oh, @@ -292,12 +325,14 @@ def _decode_stage2_kernel( stride_lse_s, stride_ot, stride_oh, - MAX_KV_SPLITS: tl.constexpr, MIN_BLOCK_KV: tl.constexpr, + MIN_SPLIT_LEN: tl.constexpr, + MIN_SPLITS: tl.constexpr, BLOCK_DV: tl.constexpr, DV: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + SPLIT_TILE: tl.constexpr, ): batch_id = tl.program_id(0) q_head = tl.program_id(1) @@ -310,13 +345,14 @@ def _decode_stage2_kernel( effective_start = tl.maximum(0, q_pos - SLIDING_WINDOW + 1) effective_len = tl.maximum(0, effective_end - effective_start) - kv_splits = tl.load(num_kv_splits_ptr + batch_id) + kv_splits = _decode_num_splits(effective_len, max_kv_splits, MIN_SPLIT_LEN, MIN_SPLITS, MIN_BLOCK_KV) kv_len_per_split = ( tl.cdiv(tl.cdiv(effective_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV ) offs_d = tl.arange(0, BLOCK_DV) mask_d = offs_d < DV + offs_s = tl.arange(0, SPLIT_TILE) if HAS_SINKS: m_i = tl.load(sinks_ptr + q_head).to(tl.float32) l_i = 1.0 @@ -325,26 +361,27 @@ 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 + offs_d + mid_base = batch_id * stride_mid_ob + q_head * stride_mid_oh lse_base = batch_id * stride_lse_b + q_head * stride_lse_h - for split_id in tl.range(0, MAX_KV_SPLITS, num_stages=2): - split_start = kv_len_per_split * split_id - split_end = tl.minimum(split_start + kv_len_per_split, effective_len) - - if split_end > split_start: - partial = tl.load( - mid_o_ptr + mid_base + split_id * stride_mid_os, - mask=mask_d, - other=0.0, - ) - partial_lse = tl.load(mid_lse_ptr + lse_base + split_id * stride_lse_s) - m_new = tl.maximum(partial_lse, m_i) - alpha = tl.exp(m_i - m_new) - beta = tl.exp(partial_lse - m_new) - acc = acc * alpha + partial * beta - l_i = l_i * alpha + beta - m_i = m_new + # one load round trip per SPLIT_TILE splits + for s0 in range(0, kv_splits, SPLIT_TILE): + split_id = s0 + offs_s + valid = (split_id < kv_splits) & (kv_len_per_split * split_id < effective_len) + partial_lse = tl.load(mid_lse_ptr + lse_base + split_id * stride_lse_s, mask=valid, other=-float("inf")) + partial = tl.load( + mid_o_ptr + mid_base + split_id[:, None] * stride_mid_os + offs_d[None, :], + mask=valid[:, None] & mask_d[None, :], + other=0.0, + ) + m_new = tl.maximum(m_i, tl.max(partial_lse, axis=0)) + # a tile of empty splits with no sink leaves m at -inf, where exp(-inf - -inf) is nan + m_safe = tl.where(m_new == -float("inf"), 0.0, m_new) + alpha = tl.exp(m_i - m_safe) + beta = tl.exp(partial_lse - m_safe) + acc = acc * alpha + tl.sum(partial * beta[:, None], axis=0) + l_i = l_i * alpha + tl.sum(beta, axis=0) + m_i = m_new out = tl.where(l_i == 0.0, 0.0, acc / l_i) tl.store( @@ -363,14 +400,17 @@ def decode_paged_attention( q_positions: torch.Tensor, attn_logits: torch.Tensor, attn_lse: torch.Tensor, - num_kv_splits: torch.Tensor, - max_kv_splits: int, sm_scale: float, sliding_window: int | None = None, sinks: torch.Tensor | None = None, out: torch.Tensor | None = None, + max_kv_splits: int | None = None, ) -> torch.Tensor: - """SGLang-style split-k grouped decode attention for one query per request.""" + """SGLang-style split-k grouped decode attention for one query 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``. + 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 @@ -382,11 +422,18 @@ def decode_paged_attention( assert num_q_heads % num_kv_heads == 0 assert attn_logits.shape[0] >= batch assert attn_logits.shape[1] >= num_q_heads - assert attn_logits.shape[2] >= max_kv_splits assert attn_logits.shape[3] >= head_dim assert attn_lse.shape[0] >= batch assert attn_lse.shape[1] >= num_q_heads - assert attn_lse.shape[2] >= max_kv_splits + if sliding_window is not None and sliding_window <= 0: + # the kernels treat a window <= 0 as no window + sliding_window = None + 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)) + max_kv_splits = min(max_kv_splits, attn_logits.shape[2], attn_lse.shape[2]) if sinks is not None: assert sinks.is_cuda assert sinks.dim() == 1 @@ -404,12 +451,17 @@ def decode_paged_attention( block_d = triton.next_power_of_2(head_dim) block_dv = triton.next_power_of_2(head_dim) grid = (batch, triton.cdiv(num_q_heads, valid_block_h), max_kv_splits) - block_n = 32 - if sliding_window and head_dim <= 256: + # a grid that fits in one wave runs about one program per SM, which then needs wider tiles to keep enough loads in flight + block_n = 64 if grid[0] * grid[1] * grid[2] <= _num_sms(q.device.index) else _MIN_BLOCK_KV + num_stages = 3 + # stage 1 keeps one K tile and one V tile in shared memory + if 2 * block_n * block_d * k_cache.element_size() > 0.8 * _optin_smem_bytes(q.device.index): + block_n = _MIN_BLOCK_KV + if sliding_window: split_len = triton.cdiv(triton.cdiv(sliding_window, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV - # A window split of several tiles in a single-wave grid is latency bound; 64-key tiles halve its serial steps. - if split_len >= 128 and grid[0] * grid[1] * grid[2] <= _num_sms(q.device.index): - block_n = 64 + if split_len <= block_n: + # a full window's split fits in one tile, where a deeper pipeline only adds overhead + block_n, num_stages = split_len, 2 _decode_grouped_stage1_kernel[grid]( q, @@ -421,7 +473,7 @@ def decode_paged_attention( q_positions, attn_logits, attn_lse, - num_kv_splits, + max_kv_splits, q.stride(0), q.stride(1), k_cache.stride(0), @@ -442,19 +494,22 @@ def decode_paged_attention( BLOCK_H=block_h, VALID_BLOCK_H=valid_block_h, MIN_BLOCK_KV=_MIN_BLOCK_KV, + MIN_SPLIT_LEN=_MIN_SPLIT_LEN, + MIN_SPLITS=_MIN_SPLITS, D=head_dim, DV=head_dim, SLIDING_WINDOW=sliding_window or 0, num_warps=4, - num_stages=2, + num_stages=num_stages, ) + split_tile, stage2_warps = _stage2_launch(max_kv_splits, block_dv) _decode_stage2_kernel[(batch, num_q_heads)]( attn_logits, attn_lse, o, indptr, q_positions, - num_kv_splits, + max_kv_splits, sinks_arg, attn_logits.stride(0), attn_logits.stride(1), @@ -464,14 +519,15 @@ def decode_paged_attention( attn_lse.stride(2), o.stride(0), o.stride(1), - MAX_KV_SPLITS=max_kv_splits, MIN_BLOCK_KV=_MIN_BLOCK_KV, + MIN_SPLIT_LEN=_MIN_SPLIT_LEN, + MIN_SPLITS=_MIN_SPLITS, BLOCK_DV=block_dv, DV=head_dim, SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, - num_warps=4, - num_stages=2, + SPLIT_TILE=split_tile, + num_warps=stage2_warps, ) return o diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index a29f621dd..60a5e2747 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -286,7 +286,6 @@ def test_decode_triton_attention_matches_reference( device=device, ) attn_lse = torch.empty(batch, num_q_heads, max_kv_splits, dtype=torch.float32, device=device) - num_kv_splits = torch.full((batch,), max_kv_splits, dtype=torch.int32, device=device) actual = decode_paged_attention( q, @@ -297,8 +296,6 @@ def test_decode_triton_attention_matches_reference( q_positions, attn_logits, attn_lse, - num_kv_splits, - max_kv_splits, sm_scale, sliding_window=sliding_window, ) @@ -344,11 +341,10 @@ def test_decode_triton_attention_non_pow2_group(num_q_heads: int, num_kv_heads: batch, num_q_heads, max_kv_splits, head_dim, dtype=torch.float32, device=device ) attn_lse = torch.empty(batch, num_q_heads, max_kv_splits, dtype=torch.float32, device=device) - num_kv_splits = torch.full((batch,), max_kv_splits, dtype=torch.int32, device=device) actual = decode_paged_attention( q, k_cache, v_cache, indptr, indices, q_positions, - attn_logits, attn_lse, num_kv_splits, max_kv_splits, sm_scale, + attn_logits, attn_lse, sm_scale, ) expected = _reference_paged_attention( q, k_cache, v_cache, indptr, indices, q_to_req, q_positions, sm_scale, None, @@ -392,7 +388,6 @@ def test_decode_triton_attention_with_sinks_matches_reference(): device=device, ) attn_lse = torch.empty(batch, num_q_heads, max_kv_splits, dtype=torch.float32, device=device) - num_kv_splits = torch.full((batch,), max_kv_splits, dtype=torch.int32, device=device) actual = decode_paged_attention( q, @@ -403,8 +398,6 @@ def test_decode_triton_attention_with_sinks_matches_reference(): q_positions, attn_logits, attn_lse, - num_kv_splits, - max_kv_splits, sm_scale, sinks=sinks, ) @@ -424,6 +417,54 @@ def test_decode_triton_attention_with_sinks_matches_reference(): 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( + ("seq_lens", "num_kv_heads", "sliding_window", "scratch_splits", "use_sinks"), + [ + ([1, 129, 1000, 8193], 2, None, 64, False), + ([1, 129, 1000, 8193], 2, None, 8, False), + ([300, 2000, 70], 8, 300, 64, True), + ], +) +def test_decode_triton_attention_splits_long_requests( + seq_lens: list[int], + num_kv_heads: int, + sliding_window: int | None, + scratch_splits: int, + use_sinks: bool, +): + """Requests of many splits next to short ones, a scratch smaller than the split cap, and windows that span several splits.""" + from freetoken.kernel.triton.attention import decode_paged_attention + + torch.manual_seed(5) + device = torch.device("cuda") + batch = len(seq_lens) + num_q_heads = 16 + head_dim = 256 + total_kv = sum(seq_lens) + q = torch.randn(batch, 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, *seq_lens], dtype=torch.int32, device=device).cumsum(0).to(torch.int32) + indices = torch.randperm(total_kv, device=device).to(torch.int32) + q_positions = torch.tensor([n - 1 for n in seq_lens], dtype=torch.int64, device=device) + q_to_req = torch.arange(batch, dtype=torch.int32, 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(batch, num_q_heads, scratch_splits, head_dim, dtype=torch.float32, device=device) + attn_lse = torch.empty(batch, 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, + ) + expected = _reference_paged_attention( + q, k_cache, v_cache, indptr, indices, q_to_req, q_positions, sm_scale, sliding_window, sinks=sinks, + ) + 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( @@ -837,10 +878,8 @@ def __init__(self, device: torch.device): assert replay_metadata.q_positions.data_ptr() == backend.capture.positions.data_ptr() assert replay_metadata.attn_logits is not None assert replay_metadata.attn_lse is not None - assert replay_metadata.num_kv_splits is not None assert replay_metadata.attn_logits.data_ptr() == backend.capture.attn_logits.data_ptr() assert replay_metadata.attn_lse.data_ptr() == backend.capture.attn_lse.data_ptr() - assert replay_metadata.num_kv_splits.data_ptr() == backend.capture.num_kv_splits.data_ptr() torch.testing.assert_close(backend.capture.cu_seqlens_k[:3], expected_indptr) torch.testing.assert_close(backend.capture.page_table.view(-1)[:8], expected_indices)