| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Bit-exact kernel equivalence for MLA decode/write kernels on the |
| 4 | cross-layer (block-major) KV cache layout. |
| 5 | |
| 6 | The cross-layer layout carves each layer's per-block page out of a single |
| 7 | unified slot, so the per-layer view has an inflated ``stride(0)`` (the full |
| 8 | unified slot) and a non-zero storage offset. These tests confirm the MLA |
| 9 | kernels behind the backends that opt in to the layout (FlashMLA dense, |
| 10 | FlashInfer MLA dense, FlashMLA fp8 sparse, plus the ``concat_and_cache_mla`` |
| 11 | write) honor that strided view bit-identically to a contiguous per-layer |
| 12 | cache, and that writes do not bleed into neighbouring layers' segments. |
| 13 | """ |
| 14 | |
| 15 | import pytest |
| 16 | import torch |
| 17 | |
| 18 | pytestmark = pytest.mark.skipif( |
| 19 | not torch.cuda.is_available(), reason="MLA cache kernels require CUDA" |
| 20 | ) |
| 21 | |
| 22 | |
| 23 | def test_concat_and_cache_mla_into_unified_slot_view(): |
| 24 | """concat_and_cache_mla must write correctly into a per-layer view whose |
| 25 | block stride is the full unified slot (block-major), with zero bleed into |
| 26 | the other layers' segments of the same slot.""" |
| 27 | from vllm import _custom_ops as ops |
| 28 | |
| 29 | torch.manual_seed(0) |
| 30 | dev = "cuda" |
| 31 | kv_lora_rank = 512 |
| 32 | pe = 64 |
| 33 | entry = kv_lora_rank + pe |
| 34 | page = 64 |
| 35 | num_blocks = 32 |
| 36 | ntok = 200 |
| 37 | |
| 38 | kv_c = torch.randn(ntok, kv_lora_rank, device=dev, dtype=torch.bfloat16) |
| 39 | k_pe = torch.randn(ntok, pe, device=dev, dtype=torch.bfloat16) |
| 40 | slot = torch.randperm(num_blocks * page, device=dev, dtype=torch.int64)[:ntok] |
| 41 | scale = torch.tensor(1.0, device=dev) |
| 42 | |
| 43 | def write(cache): |
| 44 | ops.concat_and_cache_mla(kv_c, k_pe, cache, slot, "auto", scale) |
| 45 | |
| 46 | # Contiguous per-layer reference: (num_blocks, page, entry). |
| 47 | ref = torch.zeros(num_blocks, page, entry, device=dev, dtype=torch.bfloat16) |
| 48 | write(ref) |
| 49 | |
| 50 | # Unified slot holding three layer pages per block. Carve the middle |
| 51 | # layer's view (non-zero offset, block stride == full unified slot). |
| 52 | layer_page_elems = page * entry |
| 53 | n_layers = 3 |
| 54 | unified_slot_elems = n_layers * layer_page_elems |
| 55 | big = torch.zeros(num_blocks, unified_slot_elems, device=dev, dtype=torch.bfloat16) |
| 56 | flat = big.view(-1) |
| 57 | offset = layer_page_elems # middle layer |
| 58 | view = torch.as_strided( |
| 59 | flat, |
| 60 | size=(num_blocks, page, entry), |
| 61 | stride=(unified_slot_elems, entry, 1), |
| 62 | storage_offset=offset, |
| 63 | ) |
| 64 | assert not view.is_contiguous() |
| 65 | assert view.stride(0) == unified_slot_elems |
| 66 | write(view) |
| 67 | |
| 68 | # Bit-exact equivalence and zero bleed into the neighbour segments. |
| 69 | max_diff = (ref.float() - view.float()).abs().max().item() |
| 70 | assert max_diff == 0.0, f"max|Δ| = {max_diff}" |
| 71 | |
| 72 | neighbour_lo = torch.as_strided( |
| 73 | flat, (num_blocks, layer_page_elems), (unified_slot_elems, 1), 0 |
| 74 | ) |
| 75 | neighbour_hi = torch.as_strided( |
| 76 | flat, |
| 77 | (num_blocks, layer_page_elems), |
| 78 | (unified_slot_elems, 1), |
| 79 | 2 * layer_page_elems, |
| 80 | ) |
| 81 | assert neighbour_lo.abs().max().item() == 0.0 |
| 82 | assert neighbour_hi.abs().max().item() == 0.0 |
| 83 | |
| 84 | |
| 85 | def test_flashmla_dense_decode_unified_slot_view(): |
| 86 | """FlashMLA dense decode (FLASHMLA backend, e.g. Kimi-K2-style dense MLA |
| 87 | on Hopper) must read a unified-slot block-major view bit-identically to a |
| 88 | contiguous per-layer cache.""" |
| 89 | import vllm.v1.attention.ops.flashmla as fm |
| 90 | |
| 91 | ok, reason = fm.is_flashmla_dense_supported() |
| 92 | if not ok: |
| 93 | pytest.skip(reason) |
| 94 | |
| 95 | torch.manual_seed(0) |
| 96 | dev = "cuda" |
| 97 | dt = torch.bfloat16 |
| 98 | head_dim = 576 |
| 99 | hdv = 512 |
| 100 | h_q = 128 |
| 101 | page = 64 |
| 102 | num_blocks = 64 |
| 103 | bs = 4 |
| 104 | n_layers = 3 |
| 105 | layer = 1 |
| 106 | |
| 107 | q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=dt) * 0.1 |
| 108 | kv_data = torch.randn(num_blocks, page, 1, head_dim, device=dev, dtype=dt) * 0.1 |
| 109 | |
| 110 | # (A) contiguous per-layer reference. |
| 111 | cache_contiguous = kv_data.clone().contiguous() |
| 112 | |
| 113 | # (B) unified slot: view one layer -> inflated stride(0), non-zero offset. |
| 114 | unified = ( |
| 115 | torch.randn(num_blocks, n_layers, page, 1, head_dim, device=dev, dtype=dt) * 0.1 |
| 116 | ) |
| 117 | unified[:, layer].copy_(kv_data) |
| 118 | cache_view = unified[:, layer] |
| 119 | assert not cache_view.is_contiguous() |
| 120 | assert cache_view.stride(0) == n_layers * page * 1 * head_dim |
| 121 | |
| 122 | max_blk = num_blocks // bs |
| 123 | block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view( |
| 124 | bs, max_blk |
| 125 | ) |
| 126 | cache_seqlens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32) |
| 127 | |
| 128 | def run(kc): |
| 129 | meta, num_splits = fm.get_mla_metadata() |
| 130 | out, _ = fm.flash_mla_with_kvcache( |
| 131 | q=q, |
| 132 | k_cache=kc, |
| 133 | block_table=block_table, |
| 134 | cache_seqlens=cache_seqlens, |
| 135 | head_dim_v=hdv, |
| 136 | tile_scheduler_metadata=meta, |
| 137 | num_splits=num_splits, |
| 138 | softmax_scale=head_dim**-0.5, |
| 139 | causal=True, |
| 140 | ) |
| 141 | return out.clone().float() |
| 142 | |
| 143 | out_ref = run(cache_contiguous) |
| 144 | out_view = run(cache_view) |
| 145 | assert torch.isfinite(out_ref).all() |
| 146 | assert out_ref.abs().max().item() > 0.0 |
| 147 | assert (out_ref - out_view).abs().max().item() == 0.0 |
| 148 | |
| 149 | |
| 150 | def test_flashinfer_mla_dense_decode_unified_slot_view(): |
| 151 | """FlashInfer MLA dense decode must read a unified-slot block-major view |
| 152 | (inflated stride(0), non-zero storage offset) bit-identically to a |
| 153 | contiguous per-layer cache.""" |
| 154 | try: |
| 155 | from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla |
| 156 | except ImportError: |
| 157 | pytest.skip("flashinfer is not available") |
| 158 | from vllm.platforms import current_platform |
| 159 | |
| 160 | if not current_platform.is_device_capability_family(100): |
| 161 | pytest.skip("FlashInfer trtllm-gen MLA requires sm100") |
| 162 | |
| 163 | torch.manual_seed(0) |
| 164 | dev = "cuda" |
| 165 | dt = torch.bfloat16 |
| 166 | kv_lora_rank = 512 |
| 167 | qk_rope_head_dim = 64 |
| 168 | qk_nope_head_dim = 128 |
| 169 | head_dim = kv_lora_rank + qk_rope_head_dim # 576 |
| 170 | num_qo_heads = 128 |
| 171 | page = 64 |
| 172 | num_blocks = 64 |
| 173 | bs = 4 |
| 174 | n_layers = 3 # >1 so the per-layer view's block stride is inflated. |
| 175 | layer = 1 |
| 176 | |
| 177 | q = torch.randn(bs, 1, num_qo_heads, head_dim, device=dev, dtype=dt) |
| 178 | kv_data = torch.randn(num_blocks, 1, page, head_dim, device=dev, dtype=dt) |
| 179 | |
| 180 | # (A) contiguous per-layer reference. |
| 181 | kv_contiguous = kv_data.clone().contiguous() |
| 182 | |
| 183 | # (B) unified slot: block b of every layer packed together; view one layer |
| 184 | # -> stride(0) is n_layers x larger and storage offset is non-zero. |
| 185 | unified = torch.randn(num_blocks, n_layers, 1, page, head_dim, device=dev, dtype=dt) |
| 186 | unified[:, layer].copy_(kv_data) |
| 187 | kv_view = unified[:, layer] |
| 188 | assert not kv_view.is_contiguous() |
| 189 | assert kv_view.stride(0) == n_layers * 1 * page * head_dim |
| 190 | |
| 191 | max_blk = num_blocks // bs |
| 192 | block_tables = torch.arange(num_blocks, device=dev, dtype=torch.int32).view( |
| 193 | bs, max_blk |
| 194 | ) |
| 195 | seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32) |
| 196 | ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=dev) |
| 197 | scale = head_dim**-0.5 |
| 198 | |
| 199 | def run(kv): |
| 200 | return trtllm_batch_decode_with_kv_cache_mla( |
| 201 | query=q, |
| 202 | kv_cache=kv, |
| 203 | workspace_buffer=ws, |
| 204 | qk_nope_head_dim=qk_nope_head_dim, |
| 205 | kv_lora_rank=kv_lora_rank, |
| 206 | qk_rope_head_dim=qk_rope_head_dim, |
| 207 | block_tables=block_tables, |
| 208 | seq_lens=seq_lens, |
| 209 | max_seq_len=int(seq_lens.max().item()), |
| 210 | bmm1_scale=scale, |
| 211 | bmm2_scale=1.0, |
| 212 | ).clone() |
| 213 | |
| 214 | out_ref = run(kv_contiguous).float() |
| 215 | out_view = run(kv_view).float() |
| 216 | assert torch.isfinite(out_ref).all() |
| 217 | assert (out_ref - out_view).abs().max().item() == 0.0 |
| 218 | |
| 219 | |
| 220 | @pytest.mark.parametrize("next_n", [1, 2, 4]) |
| 221 | def test_flashmla_fp8_sparse_decode_unified_slot_view(next_n: int): |
| 222 | """FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a |
| 223 | unified-slot block-major view bit-identically to a contiguous fp8_ds_mla |
| 224 | cache, with finite nonzero output.""" |
| 225 | import vllm.v1.attention.ops.flashmla as fm |
| 226 | |
| 227 | ok, reason = fm.is_flashmla_sparse_supported() |
| 228 | if not ok: |
| 229 | pytest.skip(reason) |
| 230 | |
| 231 | torch.manual_seed(0) |
| 232 | dev = "cuda" |
| 233 | entry = 656 # fp8_ds_mla bytes per token |
| 234 | page = 64 |
| 235 | num_blocks = 32 |
| 236 | h_q = 128 |
| 237 | head_dim = 576 |
| 238 | hdv = 512 |
| 239 | batch = 2 |
| 240 | topk = 128 |
| 241 | n_layers = 3 |
| 242 | layer = 1 |
| 243 | |
| 244 | q = ( |
| 245 | torch.randn(batch, next_n, h_q, head_dim, device=dev, dtype=torch.bfloat16) |
| 246 | * 0.1 |
| 247 | ) |
| 248 | |
| 249 | # Structurally valid fp8 ds_mla payload: 512B fp8 + 16B f32 scales + 128B |
| 250 | # bf16 rope (random bytes corrupt the scale region and yield NaNs). |
| 251 | nope = (torch.randn(num_blocks, page, 1, 512, device=dev) * 0.1).to( |
| 252 | torch.float8_e4m3fn |
| 253 | ) |
| 254 | scales = torch.ones(num_blocks, page, 1, 4, device=dev, dtype=torch.float32) |
| 255 | rope = (torch.randn(num_blocks, page, 1, 64, device=dev) * 0.1).to(torch.bfloat16) |
| 256 | payload = torch.cat( |
| 257 | [ |
| 258 | nope.view(torch.uint8).view(num_blocks, page, 1, 512), |
| 259 | scales.view(torch.uint8).view(num_blocks, page, 1, 16), |
| 260 | rope.view(torch.uint8).view(num_blocks, page, 1, 128), |
| 261 | ], |
| 262 | dim=-1, |
| 263 | ).contiguous() |
| 264 | assert payload.shape[-1] == entry and payload.dtype == torch.uint8 |
| 265 | |
| 266 | # (A) contiguous reference. |
| 267 | cache_contiguous = payload.clone().contiguous() |
| 268 | |
| 269 | # (B) unified slot: view one layer -> inflated stride(0), non-zero offset. |
| 270 | unified = torch.randint( |
| 271 | 0, 256, (num_blocks, n_layers, page, 1, entry), device=dev, dtype=torch.uint8 |
| 272 | ) |
| 273 | unified[:, layer].copy_(payload) |
| 274 | cache_view = unified[:, layer] |
| 275 | assert not cache_view.is_contiguous() |
| 276 | assert cache_view.stride(0) == n_layers * page * 1 * entry |
| 277 | |
| 278 | # Sparse indices: each batch uses its own disjoint blocks. |
| 279 | blocks_per_batch = num_blocks // batch |
| 280 | idx = torch.full((batch, next_n, topk), -1, device=dev, dtype=torch.int32) |
| 281 | for b in range(batch): |
| 282 | slots: list[int] = [] |
| 283 | for blk in range(b * blocks_per_batch, (b + 1) * blocks_per_batch): |
| 284 | slots.extend(blk * page + off for off in range(page)) |
| 285 | slots_t = torch.tensor(slots[:topk], device=dev, dtype=torch.int32) |
| 286 | idx[b, :, : slots_t.numel()] = slots_t |
| 287 | |
| 288 | def run(kc): |
| 289 | meta, num_splits = fm.get_mla_metadata() |
| 290 | out, _ = fm.flash_mla_with_kvcache( |
| 291 | q=q, |
| 292 | k_cache=kc, |
| 293 | block_table=None, |
| 294 | cache_seqlens=None, |
| 295 | head_dim_v=hdv, |
| 296 | tile_scheduler_metadata=meta, |
| 297 | is_fp8_kvcache=True, |
| 298 | indices=idx, |
| 299 | softmax_scale=head_dim**-0.5, |
| 300 | ) |
| 301 | return out.clone().float() |
| 302 | |
| 303 | out_ref = run(cache_contiguous) |
| 304 | out_view = run(cache_view) |
| 305 | assert torch.isfinite(out_ref).all() |
| 306 | assert out_ref.abs().max().item() > 0.0 |
| 307 | assert (out_ref - out_view).abs().max().item() == 0.0 |
| 308 | |
| 309 | |
| 310 | def test_indexer_k_quant_and_cache_into_unified_slot_view(): |
| 311 | """indexer_k_quant_and_cache (DeepSeek V3.2/V4 DSA indexer K write) must |
| 312 | write correctly into a per-layer view whose block stride is the full |
| 313 | unified slot, with zero bleed into the other layers' segments.""" |
| 314 | from vllm import _custom_ops as ops |
| 315 | |
| 316 | torch.manual_seed(0) |
| 317 | dev = "cuda" |
| 318 | head_dim = 128 |
| 319 | quant_block_size = 128 |
| 320 | block_size = 64 |
| 321 | num_blocks = 16 |
| 322 | ntok = 100 |
| 323 | # Indexer cache layout per token: head_dim fp8 bytes followed by |
| 324 | # head_dim * 4 / quant_block_size scale bytes. |
| 325 | cache_stride = head_dim + head_dim * 4 // quant_block_size |
| 326 | |
| 327 | k = torch.randn(ntok, head_dim, device=dev, dtype=torch.bfloat16) |
| 328 | slot = torch.randperm(num_blocks * block_size, device=dev, dtype=torch.int64)[:ntok] |
| 329 | |
| 330 | def write(cache): |
| 331 | ops.indexer_k_quant_and_cache(k, cache, slot, quant_block_size, "ue8m0") |
| 332 | |
| 333 | # Contiguous per-layer reference. |
| 334 | ref = torch.zeros( |
| 335 | num_blocks, block_size, cache_stride, device=dev, dtype=torch.uint8 |
| 336 | ) |
| 337 | write(ref) |
| 338 | |
| 339 | # Unified slot holding three layer pages per block; carve the middle one. |
| 340 | n_layers = 3 |
| 341 | layer = 1 |
| 342 | unified = torch.zeros( |
| 343 | num_blocks, n_layers, block_size, cache_stride, device=dev, dtype=torch.uint8 |
| 344 | ) |
| 345 | view = unified[:, layer] |
| 346 | assert not view.is_contiguous() |
| 347 | assert view.stride(0) == n_layers * block_size * cache_stride |
| 348 | write(view) |
| 349 | |
| 350 | assert torch.equal(ref, view.contiguous()) |
| 351 | # Zero bleed into the neighbour layers' segments. |
| 352 | assert unified[:, 0].abs().max().item() == 0 |
| 353 | assert unified[:, 2].abs().max().item() == 0 |
| 354 | |
| 355 | |
| 356 | def test_flashattn_mla_dense_decode_unified_slot_view(): |
| 357 | """FA3 decode (FLASH_ATTN_MLA backend) must read a unified-slot |
| 358 | block-major view bit-identically to a contiguous per-layer cache.""" |
| 359 | try: |
| 360 | from vllm.vllm_flash_attn import flash_attn_varlen_func |
| 361 | except ImportError: |
| 362 | pytest.skip("vllm_flash_attn is not available") |
| 363 | from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla |
| 364 | |
| 365 | if not flash_attn_supports_mla(): |
| 366 | pytest.skip("FA3 MLA requires a Hopper device") |
| 367 | |
| 368 | torch.manual_seed(0) |
| 369 | dev = "cuda" |
| 370 | dt = torch.bfloat16 |
| 371 | kv_lora_rank = 512 |
| 372 | rope_dim = 64 |
| 373 | entry = kv_lora_rank + rope_dim # 576 |
| 374 | h_q = 16 |
| 375 | page = 64 |
| 376 | num_blocks = 64 |
| 377 | bs = 4 |
| 378 | n_layers = 3 |
| 379 | layer = 1 |
| 380 | |
| 381 | q_pe = torch.randn(bs, h_q, rope_dim, device=dev, dtype=dt) * 0.1 |
| 382 | q_nope = torch.randn(bs, h_q, kv_lora_rank, device=dev, dtype=dt) * 0.1 |
| 383 | kv_data = torch.randn(num_blocks, page, entry, device=dev, dtype=dt) * 0.1 |
| 384 | |
| 385 | # (A) contiguous per-layer reference. |
| 386 | cache_contiguous = kv_data.clone().contiguous() |
| 387 | |
| 388 | # (B) unified slot: view one layer -> inflated stride(0), non-zero offset. |
| 389 | unified = torch.randn(num_blocks, n_layers, page, entry, device=dev, dtype=dt) * 0.1 |
| 390 | unified[:, layer].copy_(kv_data) |
| 391 | cache_view = unified[:, layer] |
| 392 | assert not cache_view.is_contiguous() |
| 393 | assert cache_view.stride(0) == n_layers * page * entry |
| 394 | |
| 395 | max_blk = num_blocks // bs |
| 396 | block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view( |
| 397 | bs, max_blk |
| 398 | ) |
| 399 | seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32) |
| 400 | cu_seqlens_q = torch.arange(bs + 1, device=dev, dtype=torch.int32) |
| 401 | |
| 402 | def run(cache): |
| 403 | kv_c_cache = cache[..., :kv_lora_rank] |
| 404 | k_pe_cache = cache[..., kv_lora_rank:] |
| 405 | out = flash_attn_varlen_func( |
| 406 | q=q_pe, |
| 407 | k=k_pe_cache.unsqueeze(-2), # Add head dim of 1 |
| 408 | v=kv_c_cache.unsqueeze(-2), # Add head dim of 1 |
| 409 | q_v=q_nope, |
| 410 | max_seqlen_q=1, |
| 411 | cu_seqlens_q=cu_seqlens_q, |
| 412 | max_seqlen_k=int(seq_lens.max().item()), |
| 413 | seqused_k=seq_lens, |
| 414 | block_table=block_table, |
| 415 | softmax_scale=entry**-0.5, |
| 416 | causal=True, |
| 417 | fa_version=3, |
| 418 | ) |
| 419 | return out.clone().float() |
| 420 | |
| 421 | out_ref = run(cache_contiguous) |
| 422 | out_view = run(cache_view) |
| 423 | assert torch.isfinite(out_ref).all() |
| 424 | assert out_ref.abs().max().item() > 0.0 |
| 425 | assert (out_ref - out_view).abs().max().item() == 0.0 |
| 426 | |
| 427 | |
| 428 | def test_flashmla_dense_fp8_decode_unified_slot_view(): |
| 429 | """FlashMLA dense fp8 decode (FLASHMLA backend with quantized KV cache) |
| 430 | must read a unified-slot block-major view bit-identically to a contiguous |
| 431 | per-layer fp8 cache.""" |
| 432 | import vllm.v1.attention.ops.flashmla as fm |
| 433 | |
| 434 | ok, reason = fm.is_flashmla_dense_supported() |
| 435 | if not ok: |
| 436 | pytest.skip(reason) |
| 437 | |
| 438 | torch.manual_seed(0) |
| 439 | dev = "cuda" |
| 440 | head_dim = 576 |
| 441 | hdv = 512 |
| 442 | h_q = 128 |
| 443 | page = 64 |
| 444 | num_blocks = 64 |
| 445 | bs = 4 |
| 446 | n_layers = 3 |
| 447 | layer = 1 |
| 448 | |
| 449 | # With a quantized KV cache the decode query is quantized to fp8 as well; |
| 450 | # the fp8 dense MLA kernel requires q already in fp8. |
| 451 | q = (torch.randn(bs, 1, h_q, head_dim, device=dev) * 0.1).to(torch.float8_e4m3fn) |
| 452 | kv_data = (torch.randn(num_blocks, page, head_dim, device=dev) * 0.1).to( |
| 453 | torch.float8_e4m3fn |
| 454 | ) |
| 455 | |
| 456 | # (A) contiguous per-layer reference. |
| 457 | cache_contiguous = kv_data.clone().contiguous() |
| 458 | |
| 459 | # (B) unified slot: view one layer -> inflated stride(0), non-zero offset. |
| 460 | unified = (torch.randn(num_blocks, n_layers, page, head_dim, device=dev) * 0.1).to( |
| 461 | torch.float8_e4m3fn |
| 462 | ) |
| 463 | unified[:, layer].copy_(kv_data) |
| 464 | cache_view = unified[:, layer] |
| 465 | assert not cache_view.is_contiguous() |
| 466 | assert cache_view.stride(0) == n_layers * page * head_dim |
| 467 | |
| 468 | max_blk = num_blocks // bs |
| 469 | block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view( |
| 470 | bs, max_blk |
| 471 | ) |
| 472 | cache_seqlens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32) |
| 473 | descale = torch.ones(1, device=dev, dtype=torch.float32) |
| 474 | |
| 475 | def run(kc): |
| 476 | tile_md, num_splits = fm.get_mla_metadata_dense_fp8(cache_seqlens, h_q, 1) |
| 477 | out, _ = fm.flash_mla_with_kvcache_fp8( |
| 478 | q=q, |
| 479 | k_cache=kc.unsqueeze(-2), # Add head dim of 1 |
| 480 | block_table=block_table, |
| 481 | cache_seqlens=cache_seqlens, |
| 482 | head_dim_v=hdv, |
| 483 | tile_scheduler_metadata=tile_md, |
| 484 | num_splits=num_splits, |
| 485 | softmax_scale=head_dim**-0.5, |
| 486 | causal=True, |
| 487 | descale_q=descale, |
| 488 | descale_k=descale, |
| 489 | ) |
| 490 | return out.clone().float() |
| 491 | |
| 492 | out_ref = run(cache_contiguous) |
| 493 | out_view = run(cache_view) |
| 494 | assert torch.isfinite(out_ref).all() |
| 495 | assert out_ref.abs().max().item() > 0.0 |
| 496 | assert (out_ref - out_view).abs().max().item() == 0.0 |
| 497 | |
| 498 | |
| 499 | def test_flashinfer_mla_dense_fp8_decode_unified_slot_view(): |
| 500 | """FlashInfer MLA dense decode with an fp8 KV cache must read a |
| 501 | unified-slot block-major view bit-identically to a contiguous per-layer |
| 502 | cache.""" |
| 503 | try: |
| 504 | from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla |
| 505 | except ImportError: |
| 506 | pytest.skip("flashinfer is not available") |
| 507 | from vllm.platforms import current_platform |
| 508 | |
| 509 | if not current_platform.is_device_capability_family(100): |
| 510 | pytest.skip("FlashInfer trtllm-gen MLA requires sm100") |
| 511 | |
| 512 | torch.manual_seed(0) |
| 513 | dev = "cuda" |
| 514 | kv_lora_rank = 512 |
| 515 | qk_rope_head_dim = 64 |
| 516 | qk_nope_head_dim = 128 |
| 517 | head_dim = kv_lora_rank + qk_rope_head_dim # 576 |
| 518 | num_qo_heads = 128 |
| 519 | page = 64 |
| 520 | num_blocks = 64 |
| 521 | bs = 4 |
| 522 | n_layers = 3 |
| 523 | layer = 1 |
| 524 | |
| 525 | # With a quantized KV cache the decode query is quantized to fp8 as well |
| 526 | # (trtllm-gen has no bf16-query x fp8-cache decode kernel). |
| 527 | q = (torch.randn(bs, 1, num_qo_heads, head_dim, device=dev) * 0.1).to( |
| 528 | torch.float8_e4m3fn |
| 529 | ) |
| 530 | kv_data = (torch.randn(num_blocks, 1, page, head_dim, device=dev) * 0.1).to( |
| 531 | torch.float8_e4m3fn |
| 532 | ) |
| 533 | |
| 534 | # (A) contiguous per-layer reference. |
| 535 | kv_contiguous = kv_data.clone().contiguous() |
| 536 | |
| 537 | # (B) unified slot: view one layer -> inflated stride(0), non-zero offset. |
| 538 | unified = ( |
| 539 | torch.randn(num_blocks, n_layers, 1, page, head_dim, device=dev) * 0.1 |
| 540 | ).to(torch.float8_e4m3fn) |
| 541 | unified[:, layer].copy_(kv_data) |
| 542 | kv_view = unified[:, layer] |
| 543 | assert not kv_view.is_contiguous() |
| 544 | assert kv_view.stride(0) == n_layers * 1 * page * head_dim |
| 545 | |
| 546 | max_blk = num_blocks // bs |
| 547 | block_tables = torch.arange(num_blocks, device=dev, dtype=torch.int32).view( |
| 548 | bs, max_blk |
| 549 | ) |
| 550 | seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32) |
| 551 | ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=dev) |
| 552 | scale = head_dim**-0.5 |
| 553 | |
| 554 | def run(kv): |
| 555 | return trtllm_batch_decode_with_kv_cache_mla( |
| 556 | query=q, |
| 557 | kv_cache=kv, |
| 558 | workspace_buffer=ws, |
| 559 | qk_nope_head_dim=qk_nope_head_dim, |
| 560 | kv_lora_rank=kv_lora_rank, |
| 561 | qk_rope_head_dim=qk_rope_head_dim, |
| 562 | block_tables=block_tables, |
| 563 | seq_lens=seq_lens, |
| 564 | max_seq_len=int(seq_lens.max().item()), |
| 565 | bmm1_scale=scale, |
| 566 | bmm2_scale=1.0, |
| 567 | ).clone() |
| 568 | |
| 569 | out_ref = run(kv_contiguous).float() |
| 570 | out_view = run(kv_view).float() |
| 571 | assert torch.isfinite(out_ref).all() |
| 572 | assert (out_ref - out_view).abs().max().item() == 0.0 |