| 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 tests.kernels.quantization.nvfp4_utils import ( |
| 8 | dequant_nvfp4_kv_cache, |
| 9 | dequantize_nvfp4_to_dtype, |
| 10 | get_nvfp4_global_scale, |
| 11 | ) |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.utils.math_utils import round_up |
| 14 | from vllm.utils.torch_utils import ( |
| 15 | nvfp4_kv_cache_full_dim, |
| 16 | nvfp4_split_data_scale, |
| 17 | set_random_seed, |
| 18 | ) |
| 19 | |
| 20 | if not current_platform.is_device_capability_family(100): |
| 21 | pytest.skip( |
| 22 | "This TRTLLM kernel requires NVIDIA Blackwell.", allow_module_level=True |
| 23 | ) |
| 24 | else: |
| 25 | import flashinfer |
| 26 | |
| 27 | FLOAT32_BYTES = torch.finfo(torch.float).bits // 8 |
| 28 | FP8_DTYPE = current_platform.fp8_dtype() |
| 29 | FP4_DTYPE = torch.uint8 |
| 30 | |
| 31 | |
| 32 | def to_float8(x, dtype=torch.float8_e4m3fn): |
| 33 | finfo = torch.finfo(dtype) |
| 34 | min_val, max_val = x.aminmax() |
| 35 | amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) |
| 36 | scale = finfo.max / amax * 0.1 |
| 37 | x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max) |
| 38 | return x_scl_sat.to(dtype), scale.float().reciprocal() |
| 39 | |
| 40 | |
| 41 | def build_paged_kv_metadata( |
| 42 | seq_lens: torch.Tensor, |
| 43 | block_tables: torch.Tensor, |
| 44 | block_size: int, |
| 45 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 46 | """Build paged-KV indptr/indices/last_page_lens from seq_lens + block_tables.""" |
| 47 | kv_indptr = [0] |
| 48 | kv_indices = [] |
| 49 | kv_last_page_lens = [] |
| 50 | for i in range(len(seq_lens)): |
| 51 | sl = int(seq_lens[i]) |
| 52 | assert sl > 0 |
| 53 | nb = (sl + block_size - 1) // block_size |
| 54 | kv_indices.extend(block_tables[i, :nb].tolist()) |
| 55 | kv_indptr.append(kv_indptr[-1] + nb) |
| 56 | kv_last_page_lens.append(sl % block_size or block_size) |
| 57 | return ( |
| 58 | torch.tensor(kv_indptr, dtype=torch.int32), |
| 59 | torch.tensor(kv_indices, dtype=torch.int32), |
| 60 | torch.tensor(kv_last_page_lens, dtype=torch.int32), |
| 61 | ) |
| 62 | |
| 63 | |
| 64 | def make_nvfp4_kv_cache( |
| 65 | kv_bf16_hnd: torch.Tensor, block_size: int, head_size: int |
| 66 | ) -> tuple: |
| 67 | """Quantize bf16 KV cache to nvfp4 via reshape_and_cache_flash. |
| 68 | |
| 69 | Returns (k_data, v_data), (k_scales, v_scales), kv_scale, ref_kv_bf16. |
| 70 | """ |
| 71 | num_blocks, _, num_kv_heads, _, _ = kv_bf16_hnd.shape |
| 72 | kv_scale_val = (kv_bf16_hnd.abs().amax() / 448.0).item() |
| 73 | kv_scale_tensor = torch.tensor( |
| 74 | kv_scale_val, dtype=torch.float32, device=kv_bf16_hnd.device |
| 75 | ) |
| 76 | |
| 77 | # layout: (B, 2*H, N, full_dim) |
| 78 | # where K heads occupy the first H heads and V heads occupy the second H heads. |
| 79 | full_dim = nvfp4_kv_cache_full_dim(head_size) |
| 80 | kv_cache_hnd = torch.zeros( |
| 81 | (num_blocks, 2 * num_kv_heads, block_size, full_dim), |
| 82 | dtype=torch.uint8, |
| 83 | device=kv_bf16_hnd.device, |
| 84 | ) |
| 85 | kv_cache_nhd = kv_cache_hnd.permute(0, 2, 1, 3) |
| 86 | k_view_nhd, v_view_nhd = kv_cache_nhd.split(num_kv_heads, dim=-2) |
| 87 | |
| 88 | # Flatten input KV → token tensors [B*N, H, head_size] for the kernel. |
| 89 | num_tokens = num_blocks * block_size |
| 90 | k_tokens = ( |
| 91 | kv_bf16_hnd[:, 0] |
| 92 | .permute(0, 2, 1, 3) |
| 93 | .reshape(num_tokens, num_kv_heads, head_size) |
| 94 | ) |
| 95 | v_tokens = ( |
| 96 | kv_bf16_hnd[:, 1] |
| 97 | .permute(0, 2, 1, 3) |
| 98 | .reshape(num_tokens, num_kv_heads, head_size) |
| 99 | ) |
| 100 | slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=kv_bf16_hnd.device) |
| 101 | |
| 102 | torch.ops._C_cache_ops.reshape_and_cache_flash( |
| 103 | k_tokens, |
| 104 | v_tokens, |
| 105 | k_view_nhd, |
| 106 | v_view_nhd, |
| 107 | slot_mapping, |
| 108 | "nvfp4", |
| 109 | kv_scale_tensor, |
| 110 | kv_scale_tensor, |
| 111 | ) |
| 112 | |
| 113 | # Split into data/scale views in HNC order for trtllm kernel. |
| 114 | k_cache_hnc, v_cache_hnc = kv_cache_hnd.split(num_kv_heads, dim=1) |
| 115 | k_data, k_scales = nvfp4_split_data_scale(k_cache_hnc) |
| 116 | v_data, v_scales = nvfp4_split_data_scale(v_cache_hnc) |
| 117 | |
| 118 | # Dequantize for the FA2 reference baseline. |
| 119 | ref_k = dequant_nvfp4_kv_cache( |
| 120 | k_data, k_scales, kv_scale_val, head_size, block_size, swizzled_scales=False |
| 121 | ).to(torch.bfloat16) |
| 122 | ref_v = dequant_nvfp4_kv_cache( |
| 123 | v_data, v_scales, kv_scale_val, head_size, block_size |
| 124 | ).to(torch.bfloat16) |
| 125 | ref_kv_bf16 = torch.stack([ref_k, ref_v], dim=1) # [N, 2, H, T, D] |
| 126 | |
| 127 | return (k_data, v_data), (k_scales, v_scales), kv_scale_val, ref_kv_bf16 |
| 128 | |
| 129 | |
| 130 | def make_quantized_kv_cache( |
| 131 | kv_cache: torch.Tensor, |
| 132 | kv_quant_dtype: torch.dtype, |
| 133 | block_size: int, |
| 134 | head_size: int, |
| 135 | ) -> tuple: |
| 136 | """Quantize kv_cache based on dtype. Returns (kv_cache, kv_cache_sf, |
| 137 | kv_scale, ref_kv_cache, is_nvfp4_kv).""" |
| 138 | is_nvfp4_kv = kv_quant_dtype == FP4_DTYPE |
| 139 | if is_nvfp4_kv: |
| 140 | data, scales, kv_scale, ref = make_nvfp4_kv_cache( |
| 141 | kv_cache, block_size, head_size |
| 142 | ) |
| 143 | return data, scales, kv_scale, ref, True |
| 144 | elif kv_quant_dtype == FP8_DTYPE: |
| 145 | kv_fp8, kv_scale = to_float8(kv_cache) |
| 146 | ref = kv_fp8.to(kv_cache.dtype) * kv_scale |
| 147 | return kv_fp8, None, kv_scale, ref, False |
| 148 | else: |
| 149 | return kv_cache, None, 1.0, kv_cache, False |
| 150 | |
| 151 | |
| 152 | DTYPE = [torch.bfloat16] |
| 153 | QUANT_DTYPES = [ |
| 154 | # (q_quant_dtype, kv_quant_dtype, o_quant_dtype) |
| 155 | (None, None, None), |
| 156 | (None, FP8_DTYPE, None), |
| 157 | (FP8_DTYPE, FP8_DTYPE, None), |
| 158 | (FP8_DTYPE, FP8_DTYPE, FP8_DTYPE), |
| 159 | (FP8_DTYPE, FP8_DTYPE, FP4_DTYPE), |
| 160 | (FP8_DTYPE, FP4_DTYPE, FP8_DTYPE), # nvfp4 KV cache |
| 161 | ] |
| 162 | BATCH_SIZE = [4, 12] |
| 163 | MAX_SEQ_LENS = [(1024, 4096)] |
| 164 | NUM_HEADS = [(64, 8), (40, 8)] |
| 165 | HEAD_SIZE = [128] |
| 166 | KV_LAYOUT = ["HND"] # currently only HND is supported |
| 167 | BLOCK_SIZE = [16] |
| 168 | WINDOW_LEFT = [-1, 127] |
| 169 | SOFT_CAP = [None, 50.0] |
| 170 | HAS_SINKS = [True, False] |
| 171 | |
| 172 | NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation. |
| 173 | |
| 174 | |
| 175 | @pytest.mark.parametrize("dtype", DTYPE) |
| 176 | @pytest.mark.parametrize("quant_dtypes", QUANT_DTYPES) |
| 177 | @pytest.mark.parametrize("batch_size", BATCH_SIZE) |
| 178 | @pytest.mark.parametrize("max_seq_lens", MAX_SEQ_LENS) |
| 179 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 180 | @pytest.mark.parametrize("head_size", HEAD_SIZE) |
| 181 | @pytest.mark.parametrize("kv_layout", KV_LAYOUT) |
| 182 | @pytest.mark.parametrize("block_size", BLOCK_SIZE) |
| 183 | @pytest.mark.parametrize("window_left", WINDOW_LEFT) |
| 184 | @pytest.mark.parametrize("soft_cap", SOFT_CAP) |
| 185 | @pytest.mark.parametrize("has_sinks", HAS_SINKS) |
| 186 | @torch.inference_mode |
| 187 | def test_flashinfer_trtllm_decode_with_baseline( |
| 188 | dtype: torch.dtype, |
| 189 | quant_dtypes: tuple[torch.dtype | None, torch.dtype | None, torch.dtype | None], |
| 190 | batch_size: int, |
| 191 | max_seq_lens: tuple[int, int], |
| 192 | num_heads: tuple[int, int], |
| 193 | head_size: int, |
| 194 | kv_layout: str, |
| 195 | block_size: int, |
| 196 | window_left: int, |
| 197 | soft_cap: float | None, |
| 198 | has_sinks: bool, |
| 199 | ) -> None: |
| 200 | torch.set_default_device("cuda") |
| 201 | set_random_seed(42) |
| 202 | |
| 203 | q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtypes |
| 204 | q_quant_dtype = q_quant_dtype or dtype |
| 205 | kv_quant_dtype = kv_quant_dtype or dtype |
| 206 | o_quant_dtype = o_quant_dtype or dtype |
| 207 | |
| 208 | _, max_kv_len = max_seq_lens |
| 209 | |
| 210 | num_qo_heads, num_kv_heads = num_heads |
| 211 | assert num_qo_heads % num_kv_heads == 0 |
| 212 | |
| 213 | sm_scale = float(1.0 / (head_size**0.5)) |
| 214 | |
| 215 | kv_cache_shape = None |
| 216 | if kv_layout == "NHD": |
| 217 | kv_cache_shape = (NUM_BLOCKS, 2, block_size, num_kv_heads, head_size) |
| 218 | elif kv_layout == "HND": |
| 219 | kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size) |
| 220 | else: |
| 221 | raise ValueError(f"Invalid kv_layout: {kv_layout}") |
| 222 | |
| 223 | # max_q_len = 1 |
| 224 | q_lens = torch.ones((batch_size,), dtype=torch.int32) |
| 225 | q_indptr = torch.cat( |
| 226 | [ |
| 227 | torch.tensor([0], dtype=torch.int32), |
| 228 | torch.cumsum(q_lens, dim=0, dtype=torch.int32), |
| 229 | ] |
| 230 | ) |
| 231 | |
| 232 | query = torch.randn(torch.sum(q_lens).item(), num_qo_heads, head_size, dtype=dtype) |
| 233 | if q_quant_dtype == FP8_DTYPE: |
| 234 | query, q_scale = to_float8(query) |
| 235 | ref_query = query.to(dtype) * q_scale |
| 236 | else: |
| 237 | q_scale = 1.0 |
| 238 | ref_query = query |
| 239 | |
| 240 | kv_lens = torch.randint(1, max_kv_len, (batch_size,), dtype=torch.int32) |
| 241 | kv_lens[-1] = max_kv_len |
| 242 | |
| 243 | seq_lens = kv_lens + q_lens |
| 244 | max_seq_len = torch.max(seq_lens).item() |
| 245 | |
| 246 | kv_cache = torch.randn(kv_cache_shape, dtype=dtype) |
| 247 | kv_cache, kv_cache_sf, kv_scale, ref_kv_cache, is_nvfp4_kv = ( |
| 248 | make_quantized_kv_cache(kv_cache, kv_quant_dtype, block_size, head_size) |
| 249 | ) |
| 250 | |
| 251 | k_scale = v_scale = kv_scale |
| 252 | |
| 253 | max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size |
| 254 | block_tables = torch.randint( |
| 255 | 0, NUM_BLOCKS, (batch_size, max_num_blocks_per_seq), dtype=torch.int32 |
| 256 | ) |
| 257 | kv_indptr, kv_indices, kv_last_page_lens = build_paged_kv_metadata( |
| 258 | seq_lens, block_tables, block_size |
| 259 | ) |
| 260 | workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8) |
| 261 | |
| 262 | # Baseline Decode |
| 263 | if has_sinks: |
| 264 | sinks = torch.rand(num_qo_heads, dtype=torch.float32) * 5 |
| 265 | wrapper = flashinfer.BatchAttentionWithAttentionSinkWrapper( |
| 266 | float_workspace_buffer=workspace_buffer, kv_layout=kv_layout, backend="fa2" |
| 267 | ) |
| 268 | else: |
| 269 | sinks = None |
| 270 | wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( |
| 271 | float_workspace_buffer=workspace_buffer, kv_layout=kv_layout, backend="fa2" |
| 272 | ) |
| 273 | |
| 274 | wrapper.plan( |
| 275 | qo_indptr=q_indptr, |
| 276 | paged_kv_indptr=kv_indptr, |
| 277 | paged_kv_indices=kv_indices, |
| 278 | paged_kv_last_page_len=kv_last_page_lens, |
| 279 | num_qo_heads=num_qo_heads, |
| 280 | num_kv_heads=num_kv_heads, |
| 281 | head_dim_qk=head_size, |
| 282 | page_size=block_size, |
| 283 | causal=True, |
| 284 | sm_scale=sm_scale, |
| 285 | window_left=window_left, |
| 286 | logits_soft_cap=soft_cap, |
| 287 | q_data_type=dtype, |
| 288 | kv_data_type=dtype, |
| 289 | ) |
| 290 | output = torch.empty(ref_query.shape, dtype=dtype) |
| 291 | wrapper.run(ref_query, ref_kv_cache, sinks, sm_scale, out=output) |
| 292 | |
| 293 | o_scale = 1.0 |
| 294 | o_sf_scale_float = None |
| 295 | if o_quant_dtype == FP8_DTYPE: |
| 296 | _, o_scale = to_float8(output) |
| 297 | elif o_quant_dtype == FP4_DTYPE: |
| 298 | o_sf_scale = get_nvfp4_global_scale(output) |
| 299 | o_sf_scale_float = o_sf_scale.item() |
| 300 | |
| 301 | # TRTLLM Decode |
| 302 | if o_quant_dtype == FP4_DTYPE: |
| 303 | output_trtllm = flashinfer.utils.FP4Tensor( |
| 304 | torch.empty(query.shape[:-1] + (query.shape[-1] // 2,), dtype=torch.uint8), |
| 305 | torch.empty( |
| 306 | ( |
| 307 | round_up(query.shape[0], 128), |
| 308 | round_up(query.shape[1] * query.shape[2] // 16, 4), |
| 309 | ), |
| 310 | dtype=torch.float8_e4m3fn, |
| 311 | ), |
| 312 | ) |
| 313 | else: |
| 314 | output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype) |
| 315 | |
| 316 | flashinfer.decode.trtllm_batch_decode_with_kv_cache( |
| 317 | query=query, |
| 318 | kv_cache=kv_cache, |
| 319 | workspace_buffer=workspace_buffer, |
| 320 | block_tables=block_tables, |
| 321 | seq_lens=seq_lens, |
| 322 | max_seq_len=max_seq_len, |
| 323 | bmm1_scale=q_scale * k_scale * sm_scale, |
| 324 | bmm2_scale=v_scale / o_scale, |
| 325 | window_left=window_left, |
| 326 | sinks=sinks, |
| 327 | o_sf_scale=o_sf_scale_float, |
| 328 | out=output_trtllm, |
| 329 | kv_cache_sf=kv_cache_sf, |
| 330 | ) |
| 331 | if o_quant_dtype == FP8_DTYPE: |
| 332 | output_trtllm = output_trtllm.to(dtype) * o_scale |
| 333 | elif o_quant_dtype == FP4_DTYPE: |
| 334 | output_trtllm.data = output_trtllm.data.reshape( |
| 335 | -1, query.shape[1] * query.shape[2] // 2 |
| 336 | ) |
| 337 | output_trtllm = dequantize_nvfp4_to_dtype( |
| 338 | output_trtllm.data, output_trtllm.scale, o_sf_scale, dtype, query.device |
| 339 | ) |
| 340 | output_trtllm = output_trtllm.reshape(-1, query.shape[1], query.shape[2]) |
| 341 | |
| 342 | if is_nvfp4_kv: |
| 343 | rtol, atol = 1.0, 1.0 # nvfp4 has higher quantization error |
| 344 | elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE: |
| 345 | rtol, atol = 7e-2, 9e-2 |
| 346 | elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE: |
| 347 | rtol, atol = 3e-2, 4e-2 |
| 348 | elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype: |
| 349 | rtol, atol = 2e-2, 2e-2 |
| 350 | elif kv_quant_dtype == FP8_DTYPE: |
| 351 | rtol, atol = 4e-2, 6e-2 |
| 352 | else: |
| 353 | rtol, atol = 1e-2, 1e-2 |
| 354 | |
| 355 | ( |
| 356 | torch.testing.assert_close(output, output_trtllm, atol=atol, rtol=rtol), |
| 357 | f"{torch.max(torch.abs(output - output_trtllm))}", |
| 358 | ) |
| 359 | |
| 360 | |
| 361 | @pytest.mark.parametrize("quant_dtypes", [(None, None, None), (None, FP8_DTYPE, None)]) |
| 362 | @pytest.mark.parametrize("q_lens_cfg", [[1, 3, 2, 4], [2, 2, 1, 5, 3]]) |
| 363 | @torch.inference_mode |
| 364 | def test_flashinfer_trtllm_decode_varlen( |
| 365 | quant_dtypes: tuple[torch.dtype | None, torch.dtype | None, torch.dtype | None], |
| 366 | q_lens_cfg: list[int], |
| 367 | ) -> None: |
| 368 | """Ragged decode via cum_seq_lens_q must match the varlen baseline.""" |
| 369 | dtype = torch.bfloat16 |
| 370 | num_qo_heads, num_kv_heads = 40, 8 |
| 371 | head_size = 128 |
| 372 | block_size = 16 |
| 373 | max_kv_len = 4096 |
| 374 | torch.set_default_device("cuda") |
| 375 | set_random_seed(42) |
| 376 | |
| 377 | _, kv_quant_dtype, _ = quant_dtypes |
| 378 | kv_quant_dtype = kv_quant_dtype or dtype |
| 379 | sm_scale = float(1.0 / (head_size**0.5)) |
| 380 | batch_size = len(q_lens_cfg) |
| 381 | |
| 382 | q_lens = torch.tensor(q_lens_cfg, dtype=torch.int32) |
| 383 | q_indptr = torch.cat( |
| 384 | [ |
| 385 | torch.tensor([0], dtype=torch.int32), |
| 386 | torch.cumsum(q_lens, dim=0, dtype=torch.int32), |
| 387 | ] |
| 388 | ) |
| 389 | query = torch.randn(int(q_lens.sum()), num_qo_heads, head_size, dtype=dtype) |
| 390 | |
| 391 | kv_lens = torch.randint(1, max_kv_len, (batch_size,), dtype=torch.int32) |
| 392 | kv_lens[-1] = max_kv_len |
| 393 | seq_lens = kv_lens + q_lens |
| 394 | max_seq_len = int(seq_lens.max()) |
| 395 | |
| 396 | kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size) |
| 397 | kv_cache = torch.randn(kv_cache_shape, dtype=dtype) |
| 398 | kv_cache, kv_cache_sf, kv_scale, ref_kv_cache, _ = make_quantized_kv_cache( |
| 399 | kv_cache, kv_quant_dtype, block_size, head_size |
| 400 | ) |
| 401 | |
| 402 | max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size |
| 403 | block_tables = torch.randint( |
| 404 | 0, NUM_BLOCKS, (batch_size, max_num_blocks_per_seq), dtype=torch.int32 |
| 405 | ) |
| 406 | kv_indptr, kv_indices, kv_last_page_lens = build_paged_kv_metadata( |
| 407 | seq_lens, block_tables, block_size |
| 408 | ) |
| 409 | workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8) |
| 410 | |
| 411 | wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( |
| 412 | float_workspace_buffer=workspace_buffer, kv_layout="HND", backend="fa2" |
| 413 | ) |
| 414 | wrapper.plan( |
| 415 | qo_indptr=q_indptr, |
| 416 | paged_kv_indptr=kv_indptr, |
| 417 | paged_kv_indices=kv_indices, |
| 418 | paged_kv_last_page_len=kv_last_page_lens, |
| 419 | num_qo_heads=num_qo_heads, |
| 420 | num_kv_heads=num_kv_heads, |
| 421 | head_dim_qk=head_size, |
| 422 | page_size=block_size, |
| 423 | causal=True, |
| 424 | sm_scale=sm_scale, |
| 425 | window_left=-1, |
| 426 | q_data_type=dtype, |
| 427 | kv_data_type=dtype, |
| 428 | ) |
| 429 | output = torch.empty(query.shape, dtype=dtype) |
| 430 | wrapper.run(query, ref_kv_cache, out=output) |
| 431 | |
| 432 | output_trtllm = torch.empty(query.shape, dtype=dtype) |
| 433 | flashinfer.decode.trtllm_batch_decode_with_kv_cache( |
| 434 | query=query, |
| 435 | kv_cache=kv_cache, |
| 436 | workspace_buffer=workspace_buffer, |
| 437 | block_tables=block_tables, |
| 438 | seq_lens=seq_lens, |
| 439 | max_seq_len=max_seq_len, |
| 440 | bmm1_scale=kv_scale * sm_scale, |
| 441 | bmm2_scale=kv_scale, |
| 442 | window_left=-1, |
| 443 | out=output_trtllm, |
| 444 | q_len_per_req=None, |
| 445 | max_q_len=int(q_lens.max()), |
| 446 | cum_seq_lens_q=q_indptr, |
| 447 | kv_cache_sf=kv_cache_sf, |
| 448 | ) |
| 449 | |
| 450 | rtol, atol = (4e-2, 6e-2) if kv_quant_dtype == FP8_DTYPE else (1e-2, 1e-2) |
| 451 | torch.testing.assert_close(output, output_trtllm, atol=atol, rtol=rtol) |
| 452 | |
| 453 | |
| 454 | @pytest.mark.parametrize("dtype", DTYPE) |
| 455 | @pytest.mark.parametrize("quant_dtypes", QUANT_DTYPES) |
| 456 | @pytest.mark.parametrize("batch_size", BATCH_SIZE) |
| 457 | @pytest.mark.parametrize("max_seq_lens", MAX_SEQ_LENS) |
| 458 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 459 | @pytest.mark.parametrize("head_size", HEAD_SIZE) |
| 460 | @pytest.mark.parametrize("kv_layout", KV_LAYOUT) |
| 461 | @pytest.mark.parametrize("block_size", BLOCK_SIZE) |
| 462 | @pytest.mark.parametrize("window_left", WINDOW_LEFT) |
| 463 | @pytest.mark.parametrize("soft_cap", [None]) |
| 464 | @pytest.mark.parametrize("has_sinks", HAS_SINKS) |
| 465 | @torch.inference_mode |
| 466 | def test_flashinfer_trtllm_prefill_with_baseline( |
| 467 | dtype: torch.dtype, |
| 468 | quant_dtypes: tuple[torch.dtype | None, torch.dtype | None, torch.dtype | None], |
| 469 | batch_size: int, |
| 470 | max_seq_lens: tuple[int, int], |
| 471 | num_heads: tuple[int, int], |
| 472 | head_size: int, |
| 473 | kv_layout: str, |
| 474 | block_size: int, |
| 475 | window_left: int, |
| 476 | soft_cap: float | None, |
| 477 | has_sinks: bool, |
| 478 | ) -> None: |
| 479 | torch.set_default_device("cuda") |
| 480 | set_random_seed(42) |
| 481 | |
| 482 | q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtypes |
| 483 | q_quant_dtype = q_quant_dtype or dtype |
| 484 | kv_quant_dtype = kv_quant_dtype or dtype |
| 485 | o_quant_dtype = o_quant_dtype or dtype |
| 486 | |
| 487 | # FP8 Q + nvfp4 KV is the required combination for the nvfp4 KV path. |
| 488 | # All other mixed Q/KV dtype combinations are unsupported. |
| 489 | is_nvfp4_kv = kv_quant_dtype == FP4_DTYPE |
| 490 | if q_quant_dtype != kv_quant_dtype and not ( |
| 491 | q_quant_dtype == FP8_DTYPE and is_nvfp4_kv |
| 492 | ): |
| 493 | pytest.skip("Skipped mixed QKV dtypes for prefill") |
| 494 | |
| 495 | max_q_len, max_kv_len = max_seq_lens |
| 496 | |
| 497 | num_qo_heads, num_kv_heads = num_heads |
| 498 | assert num_qo_heads % num_kv_heads == 0 |
| 499 | |
| 500 | sm_scale = float(1.0 / (head_size**0.5)) |
| 501 | |
| 502 | kv_cache_shape = None |
| 503 | if kv_layout == "NHD": |
| 504 | kv_cache_shape = (NUM_BLOCKS, 2, block_size, num_kv_heads, head_size) |
| 505 | elif kv_layout == "HND": |
| 506 | kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size) |
| 507 | else: |
| 508 | raise ValueError(f"Invalid kv_layout: {kv_layout}") |
| 509 | |
| 510 | q_lens = torch.randint(1, max_q_len, (batch_size,), dtype=torch.int32) |
| 511 | q_lens[-1] = max_q_len |
| 512 | q_indptr = torch.cat( |
| 513 | [ |
| 514 | torch.tensor([0], dtype=torch.int32), |
| 515 | torch.cumsum(q_lens, dim=0, dtype=torch.int32), |
| 516 | ] |
| 517 | ) |
| 518 | |
| 519 | query = torch.randn(torch.sum(q_lens).item(), num_qo_heads, head_size, dtype=dtype) |
| 520 | if q_quant_dtype == FP8_DTYPE: |
| 521 | query, q_scale = to_float8(query) |
| 522 | ref_query = query.to(dtype) * q_scale |
| 523 | else: |
| 524 | q_scale = 1.0 |
| 525 | ref_query = query |
| 526 | |
| 527 | kv_lens = torch.randint(1, max_kv_len, (batch_size,), dtype=torch.int32) |
| 528 | kv_lens[-1] = max_kv_len |
| 529 | |
| 530 | seq_lens = kv_lens + q_lens |
| 531 | max_seq_len = torch.max(seq_lens).item() |
| 532 | |
| 533 | kv_cache = torch.randn(kv_cache_shape, dtype=dtype) |
| 534 | kv_cache, kv_cache_sf, kv_scale, ref_kv_cache, is_nvfp4_kv = ( |
| 535 | make_quantized_kv_cache(kv_cache, kv_quant_dtype, block_size, head_size) |
| 536 | ) |
| 537 | |
| 538 | k_scale = v_scale = kv_scale |
| 539 | |
| 540 | max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size |
| 541 | block_tables = torch.randint( |
| 542 | 0, NUM_BLOCKS, (batch_size, max_num_blocks_per_seq), dtype=torch.int32 |
| 543 | ) |
| 544 | kv_indptr, kv_indices, kv_last_page_lens = build_paged_kv_metadata( |
| 545 | seq_lens, block_tables, block_size |
| 546 | ) |
| 547 | workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8) |
| 548 | |
| 549 | # Baseline Prefill |
| 550 | if has_sinks: |
| 551 | sinks = torch.rand(num_qo_heads, dtype=torch.float32) * 5 |
| 552 | wrapper = flashinfer.BatchAttentionWithAttentionSinkWrapper( |
| 553 | float_workspace_buffer=workspace_buffer, kv_layout=kv_layout, backend="fa2" |
| 554 | ) |
| 555 | else: |
| 556 | sinks = None |
| 557 | wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( |
| 558 | float_workspace_buffer=workspace_buffer, kv_layout=kv_layout, backend="fa2" |
| 559 | ) |
| 560 | |
| 561 | wrapper.plan( |
| 562 | qo_indptr=q_indptr, |
| 563 | paged_kv_indptr=kv_indptr, |
| 564 | paged_kv_indices=kv_indices, |
| 565 | paged_kv_last_page_len=kv_last_page_lens, |
| 566 | num_qo_heads=num_qo_heads, |
| 567 | num_kv_heads=num_kv_heads, |
| 568 | head_dim_qk=head_size, |
| 569 | page_size=block_size, |
| 570 | causal=True, |
| 571 | sm_scale=sm_scale, |
| 572 | window_left=window_left, |
| 573 | logits_soft_cap=soft_cap, |
| 574 | q_data_type=dtype, |
| 575 | kv_data_type=dtype, |
| 576 | ) |
| 577 | output = torch.empty(ref_query.shape, dtype=dtype) |
| 578 | wrapper.run(ref_query, ref_kv_cache, sinks, sm_scale, out=output) |
| 579 | |
| 580 | o_scale = 1.0 |
| 581 | o_sf_scale_float = None |
| 582 | if o_quant_dtype == FP8_DTYPE: |
| 583 | _, o_scale = to_float8(output) |
| 584 | elif o_quant_dtype == FP4_DTYPE: |
| 585 | o_sf_scale = get_nvfp4_global_scale(output) |
| 586 | o_sf_scale_float = o_sf_scale.item() |
| 587 | |
| 588 | # TRTLLM Prefill |
| 589 | if o_quant_dtype == FP4_DTYPE: |
| 590 | output_trtllm = flashinfer.utils.FP4Tensor( |
| 591 | torch.empty(query.shape[:-1] + (query.shape[-1] // 2,), dtype=torch.uint8), |
| 592 | torch.empty( |
| 593 | ( |
| 594 | round_up(query.shape[0], 128), |
| 595 | round_up(query.shape[1] * query.shape[2] // 16, 4), |
| 596 | ), |
| 597 | dtype=torch.float8_e4m3fn, |
| 598 | ), |
| 599 | ) |
| 600 | else: |
| 601 | output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype) |
| 602 | |
| 603 | flashinfer.prefill.trtllm_batch_context_with_kv_cache( |
| 604 | query=query, |
| 605 | kv_cache=kv_cache, |
| 606 | workspace_buffer=workspace_buffer, |
| 607 | block_tables=block_tables, |
| 608 | seq_lens=seq_lens, |
| 609 | max_q_len=max_q_len, |
| 610 | max_kv_len=max_seq_len, |
| 611 | bmm1_scale=q_scale * k_scale * sm_scale, |
| 612 | bmm2_scale=v_scale / o_scale, |
| 613 | batch_size=batch_size, |
| 614 | cum_seq_lens_q=q_indptr, |
| 615 | cum_seq_lens_kv=kv_indptr, |
| 616 | window_left=window_left, |
| 617 | sinks=sinks, |
| 618 | o_sf_scale=o_sf_scale_float, |
| 619 | out=output_trtllm, |
| 620 | kv_cache_sf=kv_cache_sf, |
| 621 | ) |
| 622 | if o_quant_dtype == FP8_DTYPE: |
| 623 | output_trtllm = output_trtllm.to(dtype) * o_scale |
| 624 | elif o_quant_dtype == FP4_DTYPE: |
| 625 | output_trtllm.data = output_trtllm.data.reshape( |
| 626 | -1, query.shape[1] * query.shape[2] // 2 |
| 627 | ) |
| 628 | output_trtllm = dequantize_nvfp4_to_dtype( |
| 629 | output_trtllm.data, output_trtllm.scale, o_sf_scale, dtype, query.device |
| 630 | ) |
| 631 | output_trtllm = output_trtllm.reshape(-1, query.shape[1], query.shape[2]) |
| 632 | |
| 633 | if is_nvfp4_kv: |
| 634 | rtol, atol = 1.0, 1.5 # nvfp4 has higher quantization error |
| 635 | elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE: |
| 636 | rtol, atol = 3e-1, 4e-1 |
| 637 | elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE: |
| 638 | rtol, atol = 4e-2, 6e-2 |
| 639 | elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype: |
| 640 | rtol, atol = 2e-2, 3e-2 |
| 641 | else: |
| 642 | rtol, atol = 1e-2, 1e-2 |
| 643 | |
| 644 | ( |
| 645 | torch.testing.assert_close(output, output_trtllm, atol=atol, rtol=rtol), |
| 646 | f"{torch.max(torch.abs(output - output_trtllm))}", |
| 647 | ) |