| 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._custom_ops import ( |
| 8 | merge_attn_states as merge_attn_states_native, |
| 9 | ) |
| 10 | from vllm._custom_ops import ( |
| 11 | scaled_fp8_quant, |
| 12 | ) |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.v1.attention.ops.triton_merge_attn_states import ( |
| 15 | merge_attn_states as merge_attn_states_triton, |
| 16 | ) |
| 17 | |
| 18 | DEVICE = current_platform.device_type |
| 19 | |
| 20 | pytestmark = [ |
| 21 | pytest.mark.skip_global_cleanup, |
| 22 | pytest.mark.skipif( |
| 23 | not (current_platform.is_cuda_alike() or current_platform.is_xpu()), |
| 24 | reason="merge_attn_states kernels require CUDA, ROCm or XPU.", |
| 25 | ), |
| 26 | ] |
| 27 | |
| 28 | |
| 29 | # Naive PyTorch Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 |
| 30 | # can be used to combine partial attention results (in the split-KV case) |
| 31 | def merge_attn_states_torch( |
| 32 | output: torch.Tensor, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] |
| 33 | prefix_output: torch.Tensor, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] |
| 34 | prefix_lse: torch.Tensor, # [NUM_HEADS, NUM_TOKENS] |
| 35 | suffix_output: torch.Tensor, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] |
| 36 | suffix_lse: torch.Tensor, # [NUM_HEADS, NUM_TOKENS] |
| 37 | output_lse: torch.Tensor | None = None, # [NUM_HEADS, NUM_TOKENS] |
| 38 | prefill_tokens_with_context: int | None = None, |
| 39 | output_scale: torch.Tensor | None = None, # scalar, per-tensor FP8 scale |
| 40 | ): |
| 41 | # Apply prefill_tokens_with_context mask if needed |
| 42 | if prefill_tokens_with_context is None: |
| 43 | prefill_tokens_with_context = output.shape[0] |
| 44 | p_lse = prefix_lse |
| 45 | s_lse = suffix_lse |
| 46 | # inf -> -inf |
| 47 | p_lse[p_lse == torch.inf] = -torch.inf |
| 48 | s_lse[s_lse == torch.inf] = -torch.inf |
| 49 | # max_lse [NUM_HEADS, NUM_TOKENS] |
| 50 | max_lse = torch.maximum(p_lse, s_lse) |
| 51 | |
| 52 | mask = torch.ones((prefix_lse.shape[1], 1, 1), device=p_lse.device) |
| 53 | mask[prefill_tokens_with_context:].fill_(0) |
| 54 | p_lse = p_lse - max_lse |
| 55 | s_lse = s_lse - max_lse |
| 56 | p_lse_exp = torch.exp(p_lse) |
| 57 | s_lse_exp = torch.exp(s_lse) |
| 58 | out_se = p_lse_exp + s_lse_exp |
| 59 | if output_lse is not None: |
| 60 | output_lse = torch.log(out_se) + max_lse |
| 61 | output_lse[prefill_tokens_with_context:] = suffix_lse[ |
| 62 | prefill_tokens_with_context: |
| 63 | ] |
| 64 | p_scale = p_lse_exp / out_se # [NUM_HEADS, NUM_TOKENS] |
| 65 | s_scale = s_lse_exp / out_se # [NUM_HEADS, NUM_TOKENS] |
| 66 | p_scale = torch.transpose(p_scale, 0, 1).unsqueeze(2) # [NUM_TOKENS, NUM_HEADS, 1] |
| 67 | s_scale = torch.transpose(s_scale, 0, 1).unsqueeze(2) # [NUM_TOKENS, NUM_HEADS, 1] |
| 68 | output = prefix_output * p_scale * mask + suffix_output * ( |
| 69 | s_scale * mask + (1 - mask) |
| 70 | ) |
| 71 | if output_scale is not None: |
| 72 | shape = output.shape |
| 73 | output, _ = scaled_fp8_quant(output.float().view(-1, shape[-1]), output_scale) |
| 74 | output = output.view(shape) |
| 75 | return output, output_lse |
| 76 | |
| 77 | |
| 78 | NUM_BATCH_TOKENS = [256, 512, 613, 1024, 1536, 4096] |
| 79 | NUM_QUERY_HEADS = [4, 8, 16, 32, 48, 64] |
| 80 | HEAD_SIZES = [32, 48, 64, 96, 128, 256] |
| 81 | DTYPES = [torch.float32, torch.half, torch.bfloat16] |
| 82 | |
| 83 | all_case_info: list[tuple] = [] |
| 84 | |
| 85 | |
| 86 | @pytest.mark.parametrize( |
| 87 | "merge_fn", [merge_attn_states_native, merge_attn_states_triton] |
| 88 | ) |
| 89 | @pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16]) |
| 90 | def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None: |
| 91 | """When a token is empty on both sides (both LSE -inf), the 0/0 softmax |
| 92 | scales must not surface as NaN in the merged output.""" |
| 93 | num_tokens, num_heads, head_size = 6, 8, 128 |
| 94 | prefix_output = torch.zeros( |
| 95 | num_tokens, num_heads, head_size, device=DEVICE, dtype=output_dtype |
| 96 | ) |
| 97 | prefix_lse = torch.randn(num_heads, num_tokens, device=DEVICE) |
| 98 | suffix_output = torch.zeros( |
| 99 | num_tokens, num_heads, head_size, device=DEVICE, dtype=output_dtype |
| 100 | ) |
| 101 | suffix_lse = torch.randn(num_heads, num_tokens, device=DEVICE) |
| 102 | |
| 103 | # Tokens 2 and 3 are empty on both sides. |
| 104 | empty = slice(2, 4) |
| 105 | prefix_lse[:, empty] = float("-inf") |
| 106 | suffix_lse[:, empty] = float("-inf") |
| 107 | |
| 108 | output = torch.empty_like(prefix_output) |
| 109 | merge_fn(output, prefix_output, prefix_lse, suffix_output, suffix_lse) |
| 110 | |
| 111 | assert not output.isnan().any() |
| 112 | |
| 113 | |
| 114 | def generate_markdown_table(): |
| 115 | global all_case_info |
| 116 | table_header = ( |
| 117 | "| tokens | heads | headsize | dtype " |
| 118 | "| device | torch | triton | native | speedup |" |
| 119 | ) |
| 120 | table_separator = "| --- | --- | --- | --- | --- | --- | --- | --- | --- |" |
| 121 | |
| 122 | def shortly_dtype(dtype: torch.dtype) -> str: |
| 123 | return str(dtype).removeprefix("torch.") |
| 124 | |
| 125 | def shortly_device(device: str) -> str: |
| 126 | return device.removeprefix("NVIDIA").strip() |
| 127 | |
| 128 | print(table_header) |
| 129 | print(table_separator) |
| 130 | for info in all_case_info: |
| 131 | ( |
| 132 | num_tokens, |
| 133 | num_heads, |
| 134 | head_size, |
| 135 | dtype, |
| 136 | device, |
| 137 | avg_time_torch_kernel, |
| 138 | avg_time_triton_kernel, |
| 139 | avg_time_native_kernel, |
| 140 | performance_improved, |
| 141 | ) = info |
| 142 | dtype = shortly_dtype(dtype) |
| 143 | device = shortly_device(device) |
| 144 | print( |
| 145 | f"| {num_tokens} | {num_heads} | {head_size} " |
| 146 | f"| {dtype} | {device} | {avg_time_torch_kernel:.5f}ms " |
| 147 | f"| {avg_time_triton_kernel:.5f}ms " |
| 148 | f"| {avg_time_native_kernel:.5f}ms " |
| 149 | f"| {performance_improved:.4f}x |" |
| 150 | ) |
| 151 | |
| 152 | |
| 153 | @pytest.mark.parametrize("use_fp8", [False, True]) |
| 154 | @pytest.mark.parametrize("prefill_tokens_with_context", [None, 128]) |
| 155 | @pytest.mark.parametrize("num_tokens", NUM_BATCH_TOKENS) |
| 156 | @pytest.mark.parametrize("num_query_heads", NUM_QUERY_HEADS) |
| 157 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 158 | @pytest.mark.parametrize("input_dtype", DTYPES) |
| 159 | @torch.inference_mode() |
| 160 | def test_merge_attn_states( |
| 161 | prefill_tokens_with_context: int | None, |
| 162 | num_tokens: int, |
| 163 | num_query_heads: int, |
| 164 | head_size: int, |
| 165 | input_dtype: torch.dtype, |
| 166 | use_fp8: bool, |
| 167 | ): |
| 168 | NUM_TOKENS = num_tokens |
| 169 | NUM_HEADS = num_query_heads |
| 170 | HEAD_SIZE = head_size |
| 171 | |
| 172 | # When use_fp8 is set, inputs stay as input_dtype (bf16/fp16/fp32) |
| 173 | # and output becomes FP8. |
| 174 | output_dtype = input_dtype |
| 175 | output_scale = None |
| 176 | if use_fp8: |
| 177 | output_dtype = current_platform.fp8_dtype() |
| 178 | output_scale = torch.tensor([0.05], dtype=torch.float32, device=DEVICE) |
| 179 | |
| 180 | print( |
| 181 | f"\nNUM_TOKENS:{NUM_TOKENS}, NUM_HEADS:{NUM_HEADS}, " |
| 182 | f"HEAD_SIZE:{HEAD_SIZE}, input_dtype: {input_dtype}, " |
| 183 | f"output_dtype: {output_dtype}, use_fp8: {use_fp8}, " |
| 184 | f"prefill_tokens_with_context: {prefill_tokens_with_context}, " |
| 185 | f"Device: {current_platform.get_device_name()}" |
| 186 | ) |
| 187 | |
| 188 | # prefix_lse and suffix_lse contain inf and normal values |
| 189 | prefix_lse = torch.randn(NUM_HEADS, NUM_TOKENS, dtype=torch.float32, device=DEVICE) |
| 190 | suffix_lse = torch.randn(NUM_HEADS, NUM_TOKENS, dtype=torch.float32, device=DEVICE) |
| 191 | |
| 192 | # Generate boolean masks |
| 193 | mask_prefix = torch.rand(NUM_HEADS, NUM_TOKENS) < 0.1 |
| 194 | mask_suffix = torch.rand(NUM_HEADS, NUM_TOKENS) < 0.1 |
| 195 | # Ensure that the same position is not True at the same time |
| 196 | combined_mask = torch.logical_and(mask_prefix, mask_suffix) |
| 197 | mask_prefix = torch.logical_and(mask_prefix, ~combined_mask) |
| 198 | mask_suffix = torch.logical_and(mask_suffix, ~combined_mask) |
| 199 | |
| 200 | prefix_lse[mask_prefix] = float("inf") |
| 201 | suffix_lse[mask_suffix] = float("inf") |
| 202 | |
| 203 | # Other input tensors (need to be initialized but |
| 204 | # no actual calculation needed) |
| 205 | output = torch.zeros( |
| 206 | (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device=DEVICE |
| 207 | ) |
| 208 | output_lse = torch.zeros( |
| 209 | (NUM_HEADS, NUM_TOKENS), dtype=torch.float32, device=DEVICE |
| 210 | ) |
| 211 | prefix_output = torch.randn( |
| 212 | (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device=DEVICE |
| 213 | ) |
| 214 | suffix_output = torch.randn( |
| 215 | (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device=DEVICE |
| 216 | ) |
| 217 | |
| 218 | warmup_times = 2 |
| 219 | repeat_times = 20 |
| 220 | |
| 221 | output_torch = output.clone() |
| 222 | output_lse_torch = output_lse.clone() |
| 223 | total_time_torch_kernel = 0 |
| 224 | start = torch.Event(enable_timing=True) |
| 225 | end = torch.Event(enable_timing=True) |
| 226 | |
| 227 | # 0. Run the Torch kernel |
| 228 | prefix_lse_torch = prefix_lse.clone() |
| 229 | suffix_lse_torch = suffix_lse.clone() |
| 230 | for _ in range(warmup_times): |
| 231 | output_torch, output_lse_torch = merge_attn_states_torch( |
| 232 | output_torch, |
| 233 | prefix_output, |
| 234 | prefix_lse_torch, |
| 235 | suffix_output, |
| 236 | suffix_lse_torch, |
| 237 | output_lse_torch, |
| 238 | prefill_tokens_with_context, |
| 239 | output_scale, |
| 240 | ) |
| 241 | torch.accelerator.synchronize() |
| 242 | |
| 243 | for _ in range(repeat_times): |
| 244 | start.record() |
| 245 | output_torch, output_lse_torch = merge_attn_states_torch( |
| 246 | output_torch, |
| 247 | prefix_output, |
| 248 | prefix_lse_torch, |
| 249 | suffix_output, |
| 250 | suffix_lse_torch, |
| 251 | output_lse_torch, |
| 252 | prefill_tokens_with_context, |
| 253 | output_scale, |
| 254 | ) |
| 255 | end.record() |
| 256 | torch.accelerator.synchronize() |
| 257 | total_time_torch_kernel += start.elapsed_time(end) |
| 258 | |
| 259 | avg_time_torch_kernel = total_time_torch_kernel / repeat_times |
| 260 | |
| 261 | # 1. Run the Triton kernel |
| 262 | output_ref_triton = output.clone() |
| 263 | output_lse_ref_triton = output_lse.clone() |
| 264 | |
| 265 | total_time_triton_kernel = 0 |
| 266 | start = torch.Event(enable_timing=True) |
| 267 | end = torch.Event(enable_timing=True) |
| 268 | |
| 269 | for _ in range(warmup_times): |
| 270 | merge_attn_states_triton( |
| 271 | output_ref_triton, |
| 272 | prefix_output, |
| 273 | prefix_lse, |
| 274 | suffix_output, |
| 275 | suffix_lse, |
| 276 | output_lse_ref_triton, |
| 277 | prefill_tokens_with_context, |
| 278 | output_scale, |
| 279 | ) |
| 280 | torch.accelerator.synchronize() |
| 281 | |
| 282 | for _ in range(repeat_times): |
| 283 | start.record() |
| 284 | merge_attn_states_triton( |
| 285 | output_ref_triton, |
| 286 | prefix_output, |
| 287 | prefix_lse, |
| 288 | suffix_output, |
| 289 | suffix_lse, |
| 290 | output_lse_ref_triton, |
| 291 | prefill_tokens_with_context, |
| 292 | output_scale, |
| 293 | ) |
| 294 | end.record() |
| 295 | torch.accelerator.synchronize() |
| 296 | total_time_triton_kernel += start.elapsed_time(end) |
| 297 | |
| 298 | avg_time_triton_kernel = total_time_triton_kernel / repeat_times |
| 299 | |
| 300 | # 2. Run the native (CUDA/XPU) kernel |
| 301 | total_time_native_kernel = 0 |
| 302 | output_native = output.clone() |
| 303 | output_lse_native = output_lse.clone() |
| 304 | |
| 305 | for _ in range(warmup_times): |
| 306 | merge_attn_states_native( |
| 307 | output_native, |
| 308 | prefix_output, |
| 309 | prefix_lse, |
| 310 | suffix_output, |
| 311 | suffix_lse, |
| 312 | output_lse_native, |
| 313 | prefill_tokens_with_context, |
| 314 | output_scale, |
| 315 | ) |
| 316 | torch.accelerator.synchronize() |
| 317 | |
| 318 | for _ in range(repeat_times): |
| 319 | start.record() |
| 320 | merge_attn_states_native( |
| 321 | output_native, |
| 322 | prefix_output, |
| 323 | prefix_lse, |
| 324 | suffix_output, |
| 325 | suffix_lse, |
| 326 | output_lse_native, |
| 327 | prefill_tokens_with_context, |
| 328 | output_scale, |
| 329 | ) |
| 330 | end.record() |
| 331 | torch.accelerator.synchronize() |
| 332 | total_time_native_kernel += start.elapsed_time(end) |
| 333 | |
| 334 | avg_time_native_kernel = total_time_native_kernel / repeat_times |
| 335 | |
| 336 | # 3. Performance compare |
| 337 | performance_improved = avg_time_triton_kernel / avg_time_native_kernel |
| 338 | print(f" Torch time: {avg_time_torch_kernel:.6f}ms") |
| 339 | print(f"Triton time: {avg_time_triton_kernel:.6f}ms") |
| 340 | print( |
| 341 | f"Native time: {avg_time_native_kernel:.6f}ms, " |
| 342 | f"Performance: {performance_improved:.5f}x" |
| 343 | ) |
| 344 | print("-" * 100) |
| 345 | |
| 346 | # 4. Correctness compare |
| 347 | # Liger Kernel: Efficient Triton Kernels for LLM Training |
| 348 | # https://arxiv.org/pdf/2410.10989, 3.3 Correctness |
| 349 | # use rtol = 1e-2 for bfloat16. |
| 350 | if use_fp8: |
| 351 | # Compare in dequantized space (multiply back by scale) so that |
| 352 | # absolute differences reflect real precision, not amplified FP8 |
| 353 | # quantization steps. |
| 354 | atol, rtol = 1e-1, 1e-1 |
| 355 | assert output_scale is not None |
| 356 | scale = output_scale.item() |
| 357 | elif output_dtype == torch.bfloat16: |
| 358 | atol, rtol = 1e-3, 1e-2 |
| 359 | scale = 1.0 |
| 360 | else: |
| 361 | atol, rtol = 1e-3, 1e-3 |
| 362 | scale = 1.0 |
| 363 | |
| 364 | def diff(a: torch.Tensor, b: torch.Tensor): |
| 365 | max_diff = torch.max(torch.abs(a.float() - b.float())) |
| 366 | return max_diff |
| 367 | |
| 368 | # Use Triton output as reference because we want to replace |
| 369 | # the Triton kernel with the custom native kernel for merge attn |
| 370 | # states operation. |
| 371 | output_ref = output_ref_triton |
| 372 | output_lse_ref = output_lse_ref_triton |
| 373 | torch.testing.assert_close( |
| 374 | output_native.float() * scale, |
| 375 | output_ref.float() * scale, |
| 376 | atol=atol, |
| 377 | rtol=rtol, |
| 378 | ) |
| 379 | print( |
| 380 | "Output all match, max abs diff (dequantized):" |
| 381 | if use_fp8 |
| 382 | else "Output all match, max abs diff:" |
| 383 | ) |
| 384 | _diff = diff(output_ref.float() * scale, output_torch.float() * scale) |
| 385 | print(f"(Triton vs Torch) : {_diff}") |
| 386 | _diff = diff(output_torch.float() * scale, output_native.float() * scale) |
| 387 | print(f"(Native vs Torch) : {_diff}") |
| 388 | _diff = diff(output_ref.float() * scale, output_native.float() * scale) |
| 389 | print(f"(Native vs Triton): {_diff}") |
| 390 | print("-" * 100) |
| 391 | |
| 392 | torch.testing.assert_close( |
| 393 | output_lse_native.float(), output_lse_ref.float(), atol=atol, rtol=rtol |
| 394 | ) |
| 395 | print("Output LSE all match, max abs diff:") |
| 396 | print(f"(Triton vs Torch) : {diff(output_lse_torch, output_lse_ref)}") |
| 397 | print(f"(Native vs Torch) : {diff(output_lse_torch, output_lse_native)}") |
| 398 | print(f"(Native vs Triton): {diff(output_lse_ref, output_lse_native)}") |
| 399 | print("-" * 100) |
| 400 | |
| 401 | print( |
| 402 | "All output values test passed! All inf values " |
| 403 | "are correctly replaced with -inf." |
| 404 | ) |
| 405 | print("-" * 100) |
| 406 | |
| 407 | device = current_platform.get_device_name() |
| 408 | all_case_info.append( |
| 409 | ( |
| 410 | NUM_TOKENS, |
| 411 | NUM_HEADS, |
| 412 | HEAD_SIZE, |
| 413 | output_dtype, |
| 414 | device, |
| 415 | avg_time_torch_kernel, |
| 416 | avg_time_triton_kernel, |
| 417 | avg_time_native_kernel, |
| 418 | performance_improved, |
| 419 | ) |
| 420 | ) |
| 421 | if len(all_case_info) == ( |
| 422 | len(NUM_BATCH_TOKENS) * len(HEAD_SIZES) * len(NUM_QUERY_HEADS) * len(DTYPES) |
| 423 | ): |
| 424 | generate_markdown_table() |