| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Correctness tests for the fused QSA prepare kernel.""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.models.qwen4_exp.common.qsa_cache import ( |
| 9 | canonical_qsa_rope_positions, |
| 10 | circular_qsa_slot_mapping, |
| 11 | compressed_qsa_slot_mapping, |
| 12 | ) |
| 13 | from vllm.models.qwen4_exp.nvidia.indexer_qsa import apply_qsa_rope |
| 14 | from vllm.models.qwen4_exp.nvidia.ops.qsa import ( |
| 15 | qsa_compress_groups_with_ratio, |
| 16 | qsa_store_cache_rows, |
| 17 | ) |
| 18 | from vllm.models.qwen4_exp.nvidia.ops.qsa_prepare import qsa_prepare |
| 19 | from vllm.platforms import current_platform |
| 20 | from vllm.triton_utils import HAS_TRITON |
| 21 | |
| 22 | requires_qsa_kernels = pytest.mark.skipif( |
| 23 | not current_platform.is_cuda() or not HAS_TRITON, |
| 24 | reason="QSA kernels require CUDA and Triton", |
| 25 | ) |
| 26 | |
| 27 | HQ, D = 4, 128 |
| 28 | CR = 4 |
| 29 | MROPE_SECTION = (11, 11, 10) |
| 30 | EPS = 1e-6 |
| 31 | BLOCK_SIZE = 16 |
| 32 | COMP_PAGE = BLOCK_SIZE // CR |
| 33 | ROPE_POS_OFFSET = D |
| 34 | RTOL = 1.6e-2 |
| 35 | ATOL = 1e-2 |
| 36 | MIXED_BATCH = ([260, 259, 138], [1, 1, 37], [8, 8, 8]) |
| 37 | MAIN_HQ, MAIN_HK, MAIN_D, MAIN_PAGE = 6, 2, 256, 16 |
| 38 | |
| 39 | |
| 40 | def _make_block_table(block_counts): |
| 41 | num_blocks = sum(block_counts) |
| 42 | table = torch.full((len(block_counts), max(block_counts)), -1, dtype=torch.int32) |
| 43 | physical_blocks = torch.randperm(num_blocks) |
| 44 | offset = 0 |
| 45 | for request, count in enumerate(block_counts): |
| 46 | table[request, :count] = physical_blocks[offset : offset + count] |
| 47 | offset += count |
| 48 | return table, num_blocks |
| 49 | |
| 50 | |
| 51 | def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> None: |
| 52 | # e4m3 is sign-magnitude, so within a sign the uint8 code order matches the |
| 53 | # value order and one ulp is one code step. The two paths' intermediates |
| 54 | # differ in the pooling accumulation order, which at denormal magnitudes |
| 55 | # (absolute grid step 2^-9) shows up as up to 2 code steps. |
| 56 | code_diff = ( |
| 57 | actual.view(torch.uint8).int() - expected.view(torch.uint8).int() |
| 58 | ).abs() |
| 59 | abs_diff = (actual.float() - expected.float()).abs() |
| 60 | assert bool(((code_diff <= 1) | (abs_diff <= 2**-8)).all()) |
| 61 | |
| 62 | |
| 63 | def _make_main_inputs(num_tokens: int, fp8_cache: bool) -> dict: |
| 64 | """Random main-attention arguments with a paged cache in vLLM's layout.""" |
| 65 | num_blocks = num_tokens // MAIN_PAGE + 2 |
| 66 | cache = torch.zeros( |
| 67 | num_blocks, |
| 68 | MAIN_HK, |
| 69 | MAIN_PAGE, |
| 70 | 2 * MAIN_D, |
| 71 | dtype=torch.uint8 if fp8_cache else torch.bfloat16, |
| 72 | device="cuda", |
| 73 | ).transpose(1, 2) |
| 74 | slots = torch.randperm(num_blocks * MAIN_PAGE, device="cuda")[:num_tokens] |
| 75 | slots[-1] = -1 |
| 76 | norm_weights = torch.randn(2, MAIN_D, dtype=torch.bfloat16, device="cuda") * 0.2 |
| 77 | return dict( |
| 78 | main_qkv=torch.randn( |
| 79 | num_tokens, |
| 80 | 2 * (MAIN_HQ + MAIN_HK) * MAIN_D, |
| 81 | dtype=torch.bfloat16, |
| 82 | device="cuda", |
| 83 | ), |
| 84 | main_q_norm_weight=norm_weights[0], |
| 85 | main_k_norm_weight=norm_weights[1], |
| 86 | main_eps=EPS, |
| 87 | main_kv_cache=cache.view(torch.float8_e4m3fn) if fp8_cache else cache, |
| 88 | main_slot_mapping=slots, |
| 89 | main_k_scale=0.5, |
| 90 | main_v_scale=2.0, |
| 91 | ) |
| 92 | |
| 93 | |
| 94 | def _check_main_outputs(main: dict, q_out, gate_out, rope, positions) -> None: |
| 95 | """Compare with fused_qk_rmsnorm_rope_gate + reshape_and_cache_flash.""" |
| 96 | from vllm._custom_ops import reshape_and_cache_flash |
| 97 | from vllm.model_executor.layers.fused_qk_norm_rope import ( |
| 98 | fused_qk_rmsnorm_rope_gate, |
| 99 | ) |
| 100 | |
| 101 | q_gate, k, v = main["main_qkv"].split( |
| 102 | [2 * MAIN_HQ * MAIN_D, MAIN_HK * MAIN_D, MAIN_HK * MAIN_D], dim=-1 |
| 103 | ) |
| 104 | q, k, gate = fused_qk_rmsnorm_rope_gate( |
| 105 | q_gate, |
| 106 | k, |
| 107 | main["main_q_norm_weight"], |
| 108 | main["main_k_norm_weight"], |
| 109 | rope.cos_sin_cache, |
| 110 | positions, |
| 111 | EPS, |
| 112 | MAIN_HQ, |
| 113 | MAIN_HK, |
| 114 | MAIN_D, |
| 115 | rope.rotary_dim, |
| 116 | mrope_section=MROPE_SECTION if positions.ndim == 2 else None, |
| 117 | norm_beta=1.0, |
| 118 | ) |
| 119 | kv_cache = main["main_kv_cache"] |
| 120 | fp8_cache = kv_cache.dtype == torch.float8_e4m3fn |
| 121 | cache = torch.zeros_like(kv_cache.view(torch.uint8) if fp8_cache else kv_cache) |
| 122 | key_cache, value_cache = cache.split(MAIN_D, dim=-1) |
| 123 | reshape_and_cache_flash( |
| 124 | k.view(-1, MAIN_HK, MAIN_D), |
| 125 | v.view(-1, MAIN_HK, MAIN_D), |
| 126 | key_cache, |
| 127 | value_cache, |
| 128 | main["main_slot_mapping"], |
| 129 | "fp8" if fp8_cache else "auto", |
| 130 | torch.tensor(main["main_k_scale"], device="cuda"), |
| 131 | torch.tensor(main["main_v_scale"], device="cuda"), |
| 132 | ) |
| 133 | torch.testing.assert_close(q_out.flatten(1), q, rtol=RTOL, atol=ATOL) |
| 134 | assert torch.equal(gate_out.flatten(1), gate) |
| 135 | if fp8_cache: |
| 136 | assert_fp8_within_one_ulp(kv_cache, cache.view(torch.float8_e4m3fn)) |
| 137 | else: |
| 138 | torch.testing.assert_close(kv_cache, cache, rtol=RTOL, atol=ATOL) |
| 139 | |
| 140 | |
| 141 | @requires_qsa_kernels |
| 142 | @pytest.mark.usefixtures("default_vllm_config") |
| 143 | @pytest.mark.parametrize("indexer_dtype", [torch.bfloat16, torch.float8_e4m3fn]) |
| 144 | @pytest.mark.parametrize( |
| 145 | "mrope,is_2d_positions,cache_rope_positions,state_size,seq_lens,query_lens,history_lens", |
| 146 | [ |
| 147 | pytest.param(True, True, True, 4, *MIXED_BATCH, id="mrope"), |
| 148 | pytest.param(False, False, False, 4, *MIXED_BATCH, id="text"), |
| 149 | pytest.param(True, False, True, 8, *MIXED_BATCH, id="mrope-model-1d"), |
| 150 | pytest.param(True, True, False, 4, *MIXED_BATCH, id="mrope-no-position-cache"), |
| 151 | pytest.param( |
| 152 | True, |
| 153 | False, |
| 154 | False, |
| 155 | 8, |
| 156 | *MIXED_BATCH, |
| 157 | id="mrope-model-1d-no-position-cache", |
| 158 | ), |
| 159 | pytest.param( |
| 160 | False, False, True, 4, *MIXED_BATCH, id="text-with-position-cache" |
| 161 | ), |
| 162 | pytest.param(True, True, True, 4, [37], [37], [0], id="fresh"), |
| 163 | pytest.param(True, True, True, 4, [4097], [4097], [0], id="tiled"), |
| 164 | ], |
| 165 | ) |
| 166 | def test_qsa_fused_prepare_matches_unfused( |
| 167 | indexer_dtype, |
| 168 | mrope, |
| 169 | is_2d_positions, |
| 170 | cache_rope_positions, |
| 171 | state_size, |
| 172 | seq_lens, |
| 173 | query_lens, |
| 174 | history_lens, |
| 175 | ) -> None: |
| 176 | from flashinfer.norm import gemma_rmsnorm |
| 177 | |
| 178 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 179 | |
| 180 | device = "cuda" |
| 181 | rope_params = { |
| 182 | "partial_rotary_factor": 0.25, |
| 183 | "rope_theta": 10000000, |
| 184 | "rope_type": "default", |
| 185 | } |
| 186 | if mrope: |
| 187 | rope_params["mrope_interleaved"] = True |
| 188 | rope_params["mrope_section"] = list(MROPE_SECTION) |
| 189 | with torch.device(device): |
| 190 | rope = get_rope( |
| 191 | head_size=256, |
| 192 | max_position=32768, |
| 193 | rope_parameters=rope_params, |
| 194 | dtype=torch.bfloat16, |
| 195 | ) |
| 196 | |
| 197 | token_to_req = torch.cat( |
| 198 | [ |
| 199 | torch.full((length,), request, dtype=torch.int32) |
| 200 | for request, length in enumerate(query_lens) |
| 201 | ] |
| 202 | ).to(device) |
| 203 | logical_positions = torch.cat( |
| 204 | [ |
| 205 | torch.arange(seq_len - query_len, seq_len, dtype=torch.int64) |
| 206 | for seq_len, query_len in zip(seq_lens, query_lens) |
| 207 | ] |
| 208 | ).to(device) |
| 209 | num_tokens = logical_positions.numel() |
| 210 | query_start_loc = torch.tensor( |
| 211 | [0, *torch.tensor(query_lens).cumsum(0).tolist()], dtype=torch.int32 |
| 212 | ).to(device) |
| 213 | positions = ( |
| 214 | torch.stack( |
| 215 | [ |
| 216 | logical_positions, |
| 217 | logical_positions // 7 + 3, |
| 218 | logical_positions // 13 + 11, |
| 219 | ] |
| 220 | ) |
| 221 | if is_2d_positions |
| 222 | else logical_positions |
| 223 | ) |
| 224 | position_rows = ( |
| 225 | canonical_qsa_rope_positions(positions) |
| 226 | if cache_rope_positions |
| 227 | else logical_positions.view(-1, 1, 1).expand(-1, 1, 3) |
| 228 | ) |
| 229 | |
| 230 | compressed_block_counts = [ |
| 231 | (seq_len // CR + COMP_PAGE - 1) // COMP_PAGE for seq_len in seq_lens |
| 232 | ] |
| 233 | raw_block_table, num_raw_blocks = _make_block_table([1] * len(seq_lens)) |
| 234 | compressed_block_table, num_compressed_blocks = _make_block_table( |
| 235 | compressed_block_counts |
| 236 | ) |
| 237 | raw_block_table = raw_block_table.to(device) |
| 238 | raw_slots = circular_qsa_slot_mapping( |
| 239 | raw_block_table, |
| 240 | token_to_req, |
| 241 | logical_positions, |
| 242 | state_size, |
| 243 | query_start_loc, |
| 244 | ) |
| 245 | compressed_slots = compressed_qsa_slot_mapping( |
| 246 | compressed_block_table.to(device), |
| 247 | token_to_req, |
| 248 | logical_positions, |
| 249 | COMP_PAGE, |
| 250 | CR, |
| 251 | ) |
| 252 | group_counts = torch.tensor( |
| 253 | [ |
| 254 | seq_len // CR - (seq_len - query_len) // CR |
| 255 | for seq_len, query_len in zip(seq_lens, query_lens) |
| 256 | ], |
| 257 | dtype=torch.int32, |
| 258 | device=device, |
| 259 | ) |
| 260 | k_work_counts = torch.maximum(group_counts, torch.ones_like(group_counts)) |
| 261 | k_start_loc = torch.cat([k_work_counts.new_zeros(1), k_work_counts.cumsum(0)]) |
| 262 | work_requests = torch.repeat_interleave( |
| 263 | torch.arange(len(query_lens), dtype=torch.int32, device=device), |
| 264 | k_work_counts, |
| 265 | ) |
| 266 | local_work = torch.arange( |
| 267 | int(k_start_loc[-1]), dtype=torch.int32, device=device |
| 268 | ) - torch.repeat_interleave(k_start_loc[:-1], k_work_counts) |
| 269 | max_k_work = (num_tokens + (CR - 1) * len(query_lens)) // CR |
| 270 | k_work_metadata = torch.full((max_k_work, 2), -1, dtype=torch.int32, device=device) |
| 271 | k_work_metadata[: work_requests.numel()] = torch.stack( |
| 272 | (work_requests, local_work), dim=1 |
| 273 | ) |
| 274 | |
| 275 | raw_width = D + 12 if cache_rope_positions else D |
| 276 | # Match vLLM's padded-page cache layout: rows are contiguous, while physical |
| 277 | # blocks have a larger stride than their logical contents. |
| 278 | raw_page_elements = state_size * raw_width |
| 279 | fused_raw_storage = torch.zeros( |
| 280 | num_raw_blocks, |
| 281 | raw_page_elements + 16, |
| 282 | dtype=torch.bfloat16, |
| 283 | device=device, |
| 284 | ) |
| 285 | fused_raw = torch.as_strided( |
| 286 | fused_raw_storage, |
| 287 | (num_raw_blocks, state_size, 1, raw_width), |
| 288 | (raw_page_elements + 16, raw_width, raw_width, 1), |
| 289 | ) |
| 290 | for request, history_len in enumerate(history_lens): |
| 291 | history_end = seq_lens[request] - query_lens[request] |
| 292 | for position in range(history_end - history_len, history_end): |
| 293 | block = int(raw_block_table[request, 0]) |
| 294 | row = fused_raw[block, position % state_size, 0] |
| 295 | row[:D] = torch.randn(D, dtype=torch.bfloat16, device=device) |
| 296 | if cache_rope_positions: |
| 297 | row[ROPE_POS_OFFSET:].view(torch.int64).copy_( |
| 298 | torch.tensor( |
| 299 | [position, position // 7 + 3, position // 13 + 11], |
| 300 | dtype=torch.int64, |
| 301 | device=device, |
| 302 | ) |
| 303 | ) |
| 304 | unfused_raw = fused_raw.clone() |
| 305 | compressed_page_elements = COMP_PAGE * D |
| 306 | fused_compressed_storage = torch.zeros( |
| 307 | num_compressed_blocks, |
| 308 | compressed_page_elements + 16, |
| 309 | dtype=indexer_dtype, |
| 310 | device=device, |
| 311 | ) |
| 312 | fused_compressed = torch.as_strided( |
| 313 | fused_compressed_storage, |
| 314 | (num_compressed_blocks, COMP_PAGE, 1, D), |
| 315 | (compressed_page_elements + 16, D, D, 1), |
| 316 | ) |
| 317 | unfused_compressed = fused_compressed.clone() |
| 318 | |
| 319 | projected_qk = torch.randn( |
| 320 | num_tokens, (HQ + 1) * D, dtype=torch.bfloat16, device=device |
| 321 | ) |
| 322 | q_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2 |
| 323 | k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2 |
| 324 | |
| 325 | fused_query = torch.empty(num_tokens, HQ, D, dtype=indexer_dtype, device=device) |
| 326 | # The main-attention cache follows the recipe's pairing with the indexer. |
| 327 | main = _make_main_inputs(num_tokens, indexer_dtype == torch.float8_e4m3fn) |
| 328 | main_q_out, main_gate_out = qsa_prepare( |
| 329 | projected_qk[:, : HQ * D], |
| 330 | projected_qk[:, HQ * D :], |
| 331 | positions, |
| 332 | rope.cos_sin_cache, |
| 333 | q_weight, |
| 334 | k_weight, |
| 335 | EPS, |
| 336 | fused_query, |
| 337 | fused_raw, |
| 338 | raw_slots, |
| 339 | raw_block_table, |
| 340 | query_start_loc, |
| 341 | logical_positions, |
| 342 | fused_compressed, |
| 343 | compressed_slots, |
| 344 | k_work_metadata, |
| 345 | compress_ratio=CR, |
| 346 | mrope_section=MROPE_SECTION if mrope else None, |
| 347 | rope_pos_offset=ROPE_POS_OFFSET if cache_rope_positions else None, |
| 348 | **main, |
| 349 | ) |
| 350 | _check_main_outputs(main, main_q_out, main_gate_out, rope, positions) |
| 351 | |
| 352 | unfused_query = projected_qk[:, : HQ * D].reshape(num_tokens, HQ, D) |
| 353 | unfused_query = gemma_rmsnorm( |
| 354 | unfused_query.reshape(-1, D), q_weight, EPS |
| 355 | ).reshape_as(unfused_query) |
| 356 | unfused_query = apply_qsa_rope(rope, positions, unfused_query) |
| 357 | |
| 358 | raw_keys = unfused_raw[..., :D] |
| 359 | rope_positions = ( |
| 360 | unfused_raw[..., ROPE_POS_OFFSET:].view(torch.int64) |
| 361 | if cache_rope_positions |
| 362 | else None |
| 363 | ) |
| 364 | pooled, first_positions = qsa_compress_groups_with_ratio( |
| 365 | projected_qk[:, HQ * D :].reshape(-1, 1, D), |
| 366 | position_rows, |
| 367 | raw_keys, |
| 368 | raw_block_table, |
| 369 | token_to_req, |
| 370 | query_start_loc, |
| 371 | logical_positions, |
| 372 | compressed_slots, |
| 373 | CR, |
| 374 | rope_positions, |
| 375 | ) |
| 376 | compressed_rows = gemma_rmsnorm(pooled.reshape(-1, D), k_weight, EPS).reshape( |
| 377 | -1, 1, D |
| 378 | ) |
| 379 | group_positions = ( |
| 380 | first_positions.transpose(0, 1) if mrope else first_positions[:, 0] |
| 381 | ) |
| 382 | compressed_rows = apply_qsa_rope(rope, group_positions, compressed_rows) |
| 383 | qsa_store_cache_rows(unfused_compressed, compressed_slots, compressed_rows) |
| 384 | qsa_store_cache_rows(raw_keys, raw_slots, projected_qk[:, HQ * D :]) |
| 385 | if rope_positions is not None: |
| 386 | qsa_store_cache_rows(rope_positions, raw_slots, position_rows) |
| 387 | |
| 388 | if indexer_dtype == torch.float8_e4m3fn: |
| 389 | # Both paths round the same ~bf16 intermediates to e4m3. Bitwise |
| 390 | # equality (the dsv4 indexer test's bar) does not hold here: the 1D |
| 391 | # reference rope (forward_cuda) differs from _norm_rope by up to 1 |
| 392 | # bf16 ulp, and even the MRoPE pair flips a code occasionally. |
| 393 | unfused_query = unfused_query.to(indexer_dtype) |
| 394 | assert_fp8_within_one_ulp(fused_query, unfused_query) |
| 395 | else: |
| 396 | torch.testing.assert_close(fused_query, unfused_query, rtol=RTOL, atol=ATOL) |
| 397 | assert torch.equal(fused_raw.view(torch.int16), unfused_raw.view(torch.int16)) |
| 398 | if indexer_dtype == torch.float8_e4m3fn: |
| 399 | assert_fp8_within_one_ulp(fused_compressed, unfused_compressed) |
| 400 | else: |
| 401 | torch.testing.assert_close( |
| 402 | fused_compressed, unfused_compressed, rtol=RTOL, atol=ATOL |
| 403 | ) |