@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 268 lines 8.58 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 import torch
6
7 from vllm.model_executor.layers.lightning_attn import linear_decode_forward_triton
8 from vllm.platforms import current_platform
9 from vllm.utils.torch_utils import set_random_seed
10
11 DEVICE = current_platform.device_type
12
13 pytestmark = pytest.mark.skipif(
14 not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
15 reason="Lightning attention Triton kernels require CUDA/ROCm or XPU.",
16 )
17
18 NUM_HEADS = [4, 8]
19 HEAD_SIZES = [64]
20 BATCH_SIZES = [1, 2]
21 SEQ_LENGTHS = [16]
22 DTYPES = [torch.float32]
23
24
25 def reference_lightning_attention(q, k, v, ed, block_size, kv_history):
26 """Reference implementation of lightning attention core algorithm.
27
28 The difference from the main implementation is that this processes
29 each step sequentially, instead of using parallelized triton kernels
30 """
31 B, H, S, D = q.shape
32 E = v.shape[-1]
33 dtype = q.dtype
34 output = torch.zeros((B, H, S, E), dtype=dtype, device=q.device)
35
36 # Use clone() to ensure an independent copy
37 if kv_history is None:
38 kv_cache = torch.zeros((B, H, D, E), dtype=dtype, device=q.device)
39 else:
40 kv_cache = kv_history.clone()
41
42 # More efficient implementation
43 # Convert decay factors to matrix form
44 decay = torch.exp(-ed).view(1, -1, 1, 1) if ed.dim() == 1 else torch.exp(-ed)
45
46 for b in range(B):
47 for step in range(S):
48 # Process all heads at once for this position
49 q_bs = q[b, :, step] # [H, D]
50 k_bs = k[b, :, step] # [H, D]
51 v_bs = v[b, :, step] # [H, E]
52
53 # Calculate KV outer products for all heads
54 for h in range(H):
55 # Calculate KV outer product
56 kv_outer = torch.outer(k_bs[h], v_bs[h])
57
58 # Update KV cache with decay
59 # Note: Using the same order as in the Triton kernel
60 kv_cache[b, h] = decay[0, h, 0, 0] * kv_cache[b, h] + kv_outer
61
62 # Calculate attention output
63 output[b, h, step] = torch.matmul(q_bs[h], kv_cache[b, h])
64
65 # Match the shape returned by the actual implementation
66 # The actual implementation returns a tensor of shape [B, H, 2, D, E]
67 # where dimension 2 contains both KV and KV history
68 kv_reshaped = kv_cache.unsqueeze(2) # [B, H, 1, D, E]
69 final_kv_cache = torch.cat([kv_reshaped, kv_reshaped], dim=2) # [B, H, 2, D, E]
70
71 return output, final_kv_cache
72
73
74 def reference_linear_decode(q, k, v, kv_caches, slope_rate, slot_idx):
75 """Reference implementation: linear attention decode function."""
76 B, H, _, D = q.shape
77 output = torch.zeros(B, H * D, dtype=q.dtype, device=q.device)
78
79 # Calculate decay factors once (more efficient)
80 decay = torch.exp(-slope_rate).view(-1, 1, 1) # [H, 1, 1]
81
82 # Process each batch
83 for b in range(B):
84 slot_id = slot_idx[b].item()
85
86 # Skip padding positions
87 if slot_id == -1:
88 continue
89
90 # Process all heads at once for this batch
91 q_b = q[b, :, 0] # [H, D]
92 k_b = k[b, :, 0] # [H, D]
93 v_b = v[b, :, 0] # [H, D]
94
95 # Process each attention head
96 for h in range(H):
97 # Get current query, key and value
98 q_bh = q_b[h]
99 k_bh = k_b[h]
100 v_bh = v_b[h]
101
102 # Get cache
103 kv_cache_old = kv_caches[b, h]
104
105 # Calculate new key-value outer product
106 kv_outer = torch.outer(k_bh, v_bh)
107
108 # Apply decay and update cache
109 kv_new = kv_outer + decay[h, 0, 0] * kv_cache_old
110
111 # Calculate output
112 out_h = torch.matmul(q_bh, kv_new)
113
114 # Update output and cache
115 output[b, h * D : (h + 1) * D] = out_h
116 kv_caches[b, h] = kv_new
117
118 return output
119
120
121 @pytest.mark.parametrize("batch_size", BATCH_SIZES)
122 @pytest.mark.parametrize("num_heads", NUM_HEADS)
123 @pytest.mark.parametrize("head_size", HEAD_SIZES)
124 @pytest.mark.parametrize("dtype", DTYPES)
125 @torch.inference_mode()
126 def test_linear_decode_forward_triton(
127 batch_size: int,
128 num_heads: int,
129 head_size: int,
130 dtype: torch.dtype,
131 ):
132 torch.set_default_device(DEVICE)
133 set_random_seed(42)
134 base = 0.01
135 q = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
136 k = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
137 v = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
138
139 kv_caches = base * torch.randn(
140 batch_size, num_heads, head_size, head_size, dtype=dtype, device=DEVICE
141 )
142
143 kv_caches_copy = kv_caches.clone()
144
145 slope_rate = torch.zeros(num_heads, device=DEVICE)
146 for h in range(num_heads):
147 slope_rate[h] = 0.1 * (h + 1)
148
149 slot_idx = torch.arange(batch_size, device=DEVICE)
150
151 triton_output = linear_decode_forward_triton(
152 q, k, v, kv_caches, slope_rate, slot_idx
153 )
154
155 reference_output = reference_linear_decode(
156 q, k, v, kv_caches_copy, slope_rate, slot_idx
157 )
158 torch.testing.assert_close(triton_output, reference_output, rtol=1e-1, atol=1e-1)
159 torch.testing.assert_close(kv_caches, kv_caches_copy, rtol=1e-1, atol=1e-1)
160
161 assert triton_output.shape == (batch_size, num_heads * head_size)
162
163
164 @pytest.mark.parametrize("num_heads", NUM_HEADS)
165 @pytest.mark.parametrize("head_size", HEAD_SIZES)
166 @pytest.mark.parametrize("dtype", DTYPES)
167 @torch.inference_mode()
168 def test_linear_decode_forward_triton_with_padding(
169 num_heads: int,
170 head_size: int,
171 dtype: torch.dtype,
172 ):
173 torch.set_default_device(DEVICE)
174 set_random_seed(42)
175
176 batch_size = 4
177 base = 0.01
178 q = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
179 k = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
180 v = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
181
182 kv_caches = base * torch.randn(
183 batch_size, num_heads, head_size, head_size, dtype=dtype, device=DEVICE
184 )
185
186 kv_caches_copy = kv_caches.clone()
187
188 slope_rate = torch.zeros(num_heads, device=DEVICE)
189 for h in range(num_heads):
190 slope_rate[h] = 0.1 * (h + 1)
191
192 slot_idx = torch.tensor([0, 1, -1, 2], device=DEVICE)
193
194 triton_output = linear_decode_forward_triton(
195 q, k, v, kv_caches, slope_rate, slot_idx
196 )
197
198 reference_output = reference_linear_decode(
199 q, k, v, kv_caches_copy, slope_rate, slot_idx
200 )
201
202 padding_mask = (slot_idx != -1).unsqueeze(1).expand(-1, num_heads * head_size)
203
204 triton_masked = triton_output[padding_mask]
205 reference_masked = reference_output[padding_mask]
206
207 atol, rtol = 1.5e-1, 1.5e-1
208
209 valid_indices = slot_idx != -1
210
211 for i in range(batch_size):
212 if valid_indices[i] > 0:
213 torch.testing.assert_close(
214 kv_caches[i], kv_caches_copy[i], rtol=rtol, atol=atol
215 )
216
217 torch.testing.assert_close(triton_masked, reference_masked, rtol=rtol, atol=atol)
218
219 assert triton_output.shape == (batch_size, num_heads * head_size)
220
221
222 @pytest.mark.parametrize("batch_size", BATCH_SIZES)
223 @pytest.mark.parametrize("num_heads", NUM_HEADS)
224 @pytest.mark.parametrize("head_size", HEAD_SIZES)
225 @pytest.mark.parametrize("seq_len", SEQ_LENGTHS)
226 @pytest.mark.parametrize("dtype", DTYPES)
227 @torch.inference_mode()
228 def test_lightning_attention_reference(
229 batch_size: int,
230 num_heads: int,
231 head_size: int,
232 seq_len: int,
233 dtype: torch.dtype,
234 ):
235 torch.set_default_device(DEVICE)
236 set_random_seed(42)
237
238 base = 0.01
239 q = base * torch.randn(batch_size, num_heads, seq_len, head_size, dtype=dtype)
240 k = base * torch.randn(batch_size, num_heads, seq_len, head_size, dtype=dtype)
241 v = base * torch.randn(batch_size, num_heads, seq_len, head_size, dtype=dtype)
242
243 ed = torch.zeros(num_heads, device=DEVICE)
244 for h in range(num_heads):
245 ed[h] = 0.1 * (h + 1)
246
247 kv_history = base * torch.randn(
248 batch_size, num_heads, head_size, head_size, dtype=dtype, device=DEVICE
249 )
250
251 kv_history_clone = kv_history.clone()
252
253 ref_output, ref_kv_cache = reference_lightning_attention(
254 q, k, v, ed, 256, kv_history
255 )
256
257 from vllm.model_executor.layers.lightning_attn import lightning_attention
258
259 actual_output, actual_kv_cache = lightning_attention(
260 q, k, v, ed, 256, kv_history_clone
261 )
262
263 atol, rtol = 1.5e-1, 1.5e-1
264 torch.testing.assert_close(ref_output, actual_output, rtol=rtol, atol=atol)
265 torch.testing.assert_close(ref_kv_cache, actual_kv_cache, rtol=rtol, atol=atol)
266
267 assert ref_output.shape == (batch_size, num_heads, seq_len, head_size)
268 assert ref_kv_cache.shape == actual_kv_cache.shape