@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 205 lines 6.93 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit tests for the Triton DiffKV unified-attention kernel."""
4
5 import pytest
6 import torch
7
8 from tests.kernels.attention.test_triton_unified_attention import ref_paged_attn
9 from vllm.platforms import current_platform
10 from vllm.utils.math_utils import next_power_of_2
11 from vllm.utils.torch_utils import (
12 canonicalize_singleton_dim_strides,
13 set_random_seed,
14 )
15 from vllm.v1.attention.backends.fa_utils import (
16 get_flash_attn_version,
17 is_flash_attn_varlen_func_available,
18 )
19 from vllm.v1.attention.ops.triton_unified_attention_diffkv import (
20 unified_attention_diffkv,
21 )
22
23 pytestmark = pytest.mark.skip_global_cleanup
24
25 DEVICE_TYPE = current_platform.device_type
26
27 # (num_query_heads, num_kv_heads): MHA, GQA, and the num_kv_heads==1
28 # (degenerate-stride) case.
29 NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
30 # (head_size_qk, head_size_v). (192, 128) is the canonical asymmetric
31 # DiffKV shape; FA4 on Blackwell only supports head_size>128 when it is
32 # 192, and FA3 on Hopper supports it too -- so this pair is runnable on
33 # both. (128, 128) keeps the equal-dim path covered through the DiffKV
34 # kernel.
35 HEAD_SIZES = [(128, 128), (192, 128)]
36 BLOCK_SIZES = [16]
37 DTYPES = [torch.bfloat16]
38
39 NUM_BLOCKS = 2048
40
41 # 0: 2D decode kernel; 8: 3D (split-KV) decode kernel.
42 SEQ_THRESHOLD_3D_VALUES = [0, 8]
43
44 NUM_PAR_SOFTMAX_SEGMENTS = 16
45
46
47 def _alloc_segm_buffers(seq_threshold_3D: int, num_query_heads: int, head_size_v: int):
48 """Allocate the split-KV softmax scratch (last dim == head_size_v)."""
49 head_size_v_padded = next_power_of_2(head_size_v)
50 segm_output = torch.empty(
51 (
52 seq_threshold_3D,
53 num_query_heads,
54 NUM_PAR_SOFTMAX_SEGMENTS,
55 head_size_v_padded,
56 ),
57 dtype=torch.float32,
58 )
59 segm_max = torch.empty(
60 (seq_threshold_3D, num_query_heads, NUM_PAR_SOFTMAX_SEGMENTS),
61 dtype=torch.float32,
62 )
63 segm_expsum = torch.empty(
64 (seq_threshold_3D, num_query_heads, NUM_PAR_SOFTMAX_SEGMENTS),
65 dtype=torch.float32,
66 )
67 return segm_output, segm_max, segm_expsum
68
69
70 @pytest.mark.parametrize(
71 "seq_lens",
72 [
73 [(1, 1328), (5, 18), (129, 463)], # mixed prefill + decode
74 [(1, 523), (1, 37), (1, 2011)], # decode-only (exercises 3D path)
75 ],
76 )
77 @pytest.mark.parametrize("num_heads", NUM_HEADS)
78 @pytest.mark.parametrize("head_sizes", HEAD_SIZES)
79 @pytest.mark.parametrize("block_size", BLOCK_SIZES)
80 @pytest.mark.parametrize("sliding_window", [None, 128])
81 @pytest.mark.parametrize("soft_cap", [None, 50.0])
82 @pytest.mark.parametrize("dtype", DTYPES)
83 @pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
84 @torch.inference_mode()
85 def test_triton_unified_attn_diffkv_vs_reference(
86 seq_lens: list[tuple[int, int]],
87 num_heads: tuple[int, int],
88 head_sizes: tuple[int, int],
89 sliding_window: int | None,
90 soft_cap: float | None,
91 dtype: torch.dtype,
92 block_size: int,
93 seq_threshold_3D: int,
94 ) -> None:
95 head_size_qk, head_size_v = head_sizes
96
97 # Keep the FA3/FA4 comparison on NVIDIA; ROCm uses the PyTorch oracle.
98 if not current_platform.is_rocm():
99 fa_version = get_flash_attn_version(
100 head_size=head_size_qk, head_size_v=head_size_v
101 )
102 if not is_flash_attn_varlen_func_available() or fa_version not in (3, 4):
103 pytest.skip(f"FA DiffKV needs FA3/FA4 (got version {fa_version}).")
104
105 torch.set_default_device(DEVICE_TYPE)
106 set_random_seed(0)
107
108 num_seqs = len(seq_lens)
109 query_lens = [x[0] for x in seq_lens]
110 kv_lens = [x[1] for x in seq_lens]
111 num_query_heads, num_kv_heads = num_heads
112 assert num_query_heads % num_kv_heads == 0
113 max_query_len = max(query_lens)
114 max_kv_len = max(kv_lens)
115 window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
116 scale = head_size_qk**-0.5
117
118 query = torch.randn(sum(query_lens), num_query_heads, head_size_qk, dtype=dtype)
119 # Packed KV cache: [num_blocks, block_size, num_kv_heads, hqk + hv].
120 kv_cache = torch.randn(
121 NUM_BLOCKS,
122 block_size,
123 num_kv_heads,
124 head_size_qk + head_size_v,
125 dtype=dtype,
126 )
127 key_cache = kv_cache[..., :head_size_qk]
128 value_cache = kv_cache[..., head_size_qk:]
129
130 cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
131 dim=0, dtype=torch.int32
132 )
133 kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32)
134
135 max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
136 block_tables = torch.randint(
137 0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
138 )
139
140 if current_platform.is_rocm():
141 # FP32 also keeps the helper's in-place scaling off the kernel input.
142 ref_out = ref_paged_attn(
143 query.float(),
144 key_cache.float(),
145 value_cache.float(),
146 query_lens,
147 kv_lens,
148 block_tables,
149 scale,
150 sliding_window=sliding_window,
151 soft_cap=soft_cap,
152 ).to(dtype)
153 else:
154 from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
155
156 assert fa_version is not None
157 # FA's TMA path needs aligned singleton strides (num_kv_heads == 1).
158 fa_k = canonicalize_singleton_dim_strides(key_cache)
159 fa_v = canonicalize_singleton_dim_strides(value_cache)
160 ref_out = torch.empty(
161 sum(query_lens), num_query_heads, head_size_v, dtype=dtype
162 )
163 flash_attn_varlen_func(
164 q=query,
165 k=fa_k,
166 v=fa_v,
167 out=ref_out,
168 cu_seqlens_q=cu_query_lens,
169 max_seqlen_q=max_query_len,
170 seqused_k=kv_lens_t,
171 max_seqlen_k=max_kv_len,
172 softmax_scale=scale,
173 causal=True,
174 window_size=list(window_size),
175 block_table=block_tables,
176 softcap=soft_cap if soft_cap is not None else 0,
177 fa_version=fa_version,
178 )
179
180 # ---- Triton DiffKV --------------------------------------------------
181 segm_output, segm_max, segm_expsum = _alloc_segm_buffers(
182 seq_threshold_3D, num_query_heads, head_size_v
183 )
184 triton_out = torch.empty(sum(query_lens), num_query_heads, head_size_v, dtype=dtype)
185 unified_attention_diffkv(
186 q=query,
187 k=key_cache,
188 v=value_cache,
189 out=triton_out,
190 cu_seqlens_q=cu_query_lens,
191 seqused_k=kv_lens_t,
192 softmax_scale=scale,
193 causal=True,
194 window_size=window_size,
195 block_table=block_tables,
196 softcap=soft_cap if soft_cap is not None else 0,
197 max_seqlen_q=max_query_len,
198 seq_threshold_3D=seq_threshold_3D,
199 num_par_softmax_segments=NUM_PAR_SOFTMAX_SEGMENTS,
200 softmax_segm_output=segm_output,
201 softmax_segm_max=segm_max,
202 softmax_segm_expsum=segm_expsum,
203 )
204
205 torch.testing.assert_close(triton_out, ref_out, atol=2e-2, rtol=2e-2)