@tduncan / Nvidia-Optimization-Stack-VLLM / commits / aa4990a9a2

[Attention] Re-enable cross-layer KV cache layout for MLA via stride-aware kernels (#45111)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>

Yifan Qiao committed Jun 22, 2026 at 06:57 UTC aa4990a9a2024b3f93f1f26f828931f7301daa15
13 files changed +847 -21
csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu
+5 -1
@@ -136,8 +136,12 @@ typename T::Fmha::Arguments args_from_options(
136 StrideQ stride_Q_pe = cute::make_tuple(
137 static_cast<int64_t>(q_pe.stride(1)), _1{}, static_cast<int64_t>(q_pe.stride(0)));
138
139 + // Read the token and page strides from the cache tensor instead of assuming
140 + // packed pages, so strided views (e.g. per-layer views into a cross-layer
141 + // block-major cache) are addressed correctly.
142 StrideK stride_C = cute::make_tuple(
140 - static_cast<int64_t>(0 + D_latent + D_rope), _1{}, static_cast<int64_t>(page_size * (D_latent + D_rope)));
143 + static_cast<int64_t>(kv_c_and_k_pe_cache.stride(1)), _1{},
144 + static_cast<int64_t>(kv_c_and_k_pe_cache.stride(0)));
145 StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq);
146 StrideLSE stride_LSE = cute::make_tuple(_1{}, 0 + H);
147 StrideO stride_O = cute::make_tuple(static_cast<int64_t>(0 + D_latent), _1{}, static_cast<int64_t>(0 + H * D_latent));
csrc/libtorch_stable/cache_kernels.cu
+6 -7
@@ -549,7 +549,7 @@ __global__ void indexer_k_quant_and_cache_kernel(
549 const int head_dim, // dimension of each head
550 const int quant_block_size, // quantization block size
551 const int cache_block_size, // cache block size
552 - const int cache_stride, // stride for each token in kv_cache
552 + const int64_t cache_block_stride, // stride for each block in kv_cache
553
554 const bool use_ue8m0 // use ue8m0 scale format
555 ) {
@@ -590,16 +590,15 @@ __global__ void indexer_k_quant_and_cache_kernel(
590 scale = exp2f(ceilf(log2f(scale)));
591 }
592
593 - const int64_t dst_offset = block_idx * cache_block_size * cache_stride +
594 - block_offset * head_dim + head_dim_idx;
593 + const int64_t dst_offset =
594 + block_idx * cache_block_stride + block_offset * head_dim + head_dim_idx;
595 for (int i = 0; i < VEC_SIZE; i++) {
596 kv_cache[dst_offset + i] =
597 fp8::scaled_convert<cache_t, scalar_t, kv_dt>(k_val_ptr[i], scale);
598 }
599 if (threadIdx.x == 0) {
600 const int64_t dst_scale_idx =
601 - block_idx * cache_block_size * cache_stride +
602 - cache_block_size * head_dim +
601 + block_idx * cache_block_stride + cache_block_size * head_dim +
602 (block_offset * head_dim + head_dim_idx) * 4 / quant_block_size;
603 reinterpret_cast<float*>(kv_cache)[dst_scale_idx / 4] = scale;
604 }
@@ -1452,7 +1451,7 @@ void cp_gather_and_upconvert_fp8_kv_cache(
1451 reinterpret_cast<KV_T*>(k.data_ptr()), \
1452 reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
1453 slot_mapping.const_data_ptr<int64_t>(), head_dim, quant_block_size, \
1455 - cache_block_size, cache_stride, use_ue8m0);
1454 + cache_block_size, cache_block_stride, use_ue8m0);
1455
1456 void indexer_k_quant_and_cache(
1457 torch::stable::Tensor& k, // [num_tokens, head_dim]
@@ -1463,7 +1462,7 @@ void indexer_k_quant_and_cache(
1462 int num_tokens = k.size(0);
1463 int head_dim = k.size(1);
1464 int cache_block_size = kv_cache.size(1);
1466 - int cache_stride = kv_cache.size(2);
1465 + int64_t cache_block_stride = kv_cache.stride(0);
1466 bool use_ue8m0 = scale_fmt == "ue8m0";
1467
1468 STD_TORCH_CHECK(k.device() == kv_cache.device(),
tests/kernels/attention/test_cutlass_mla_decode.py
+66
@@ -212,3 +212,69 @@ def test_cutlass_mla_decode(
212 print(
213 f"{t:.3f} ms, {FLOPS / 10**9 / t:.0f} TFLOPS,", f"{bytes / 10**6 / t:.0f} GB/s"
214 )
215 +
216 +
217 +@pytest.mark.skipif(
218 + not current_platform.has_device_capability(100),
219 + reason=CUTLASS_MLA_UNSUPPORTED_REASON,
220 +)
221 +@torch.inference_mode()
222 +def test_cutlass_mla_decode_cross_layer_view():
223 + """The kernel must read the cache's page-dim stride instead of assuming
224 + pages are packed back-to-back. A per-layer view into a cross-layer
225 + (block-major) cache has stride(0) inflated by num_layers; outputs must
226 + match a contiguous cache holding the same data exactly."""
227 + device = torch.device("cuda:0")
228 + torch.set_default_dtype(torch.bfloat16)
229 + torch.set_default_device(device)
230 + torch.manual_seed(42)
231 +
232 + b, mean_sk, d, dv, block_size = 4, 512, 576, 512, 64
233 + num_layers, layer_idx = 3, 1
234 + scale = math.sqrt(d) ** (-1)
235 +
236 + num_pages = b * (mean_sk // block_size)
237 + cache_seqlens = torch.full((b,), mean_sk, dtype=torch.int32)
238 + block_table = torch.arange(num_pages, dtype=torch.int32).view(
239 + b, mean_sk // block_size
240 + )
241 +
242 + kv_contig = torch.randn(num_pages, block_size, d)
243 + # Neighbor layers hold random data so packed-pages addressing reads
244 + # garbage rather than zeros.
245 + kv_cross_layer = torch.randn(num_pages, num_layers, block_size, d)
246 + kv_view = kv_cross_layer[:, layer_idx]
247 + kv_view.copy_(kv_contig)
248 + assert kv_view.stride(0) == num_layers * block_size * d
249 +
250 + q_nope = torch.randn(b, 128, dv)
251 + q_pe = torch.randn(b, 128, d - dv)
252 + sm_count = num_compute_units(device.index)
253 + workspace_size = ops.sm100_cutlass_mla_get_workspace_size(
254 + mean_sk, b, sm_count, num_kv_splits=1
255 + )
256 + workspace = torch.empty(workspace_size, dtype=torch.uint8)
257 +
258 + def run(cache):
259 + out = torch.empty(b, 128, dv)
260 + lse = torch.empty(b, 128, dtype=torch.float32)
261 + ops.sm100_cutlass_mla_decode(
262 + out,
263 + lse,
264 + q_nope,
265 + q_pe,
266 + cache,
267 + cache_seqlens,
268 + block_table,
269 + workspace,
270 + scale,
271 + 1,
272 + )
273 + return out, lse
274 +
275 + out_contig, lse_contig = run(kv_contig)
276 + out_view, lse_view = run(kv_view)
277 +
278 + # Same data and same compute order; only addressing differs.
279 + assert torch.equal(out_contig, out_view)
280 + assert torch.equal(lse_contig, lse_view)
tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py new
+566
@@ -0,0 +1,566 @@
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 +def test_flashmla_fp8_sparse_decode_unified_slot_view():
221 + """FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a
222 + unified-slot block-major view bit-identically to a contiguous fp8_ds_mla
223 + cache, with finite nonzero output."""
224 + import vllm.v1.attention.ops.flashmla as fm
225 +
226 + ok, reason = fm.is_flashmla_sparse_supported()
227 + if not ok:
228 + pytest.skip(reason)
229 +
230 + torch.manual_seed(0)
231 + dev = "cuda"
232 + entry = 656 # fp8_ds_mla bytes per token
233 + page = 64
234 + num_blocks = 32
235 + h_q = 128
236 + head_dim = 576
237 + hdv = 512
238 + batch = 2
239 + topk = 128
240 + n_layers = 3
241 + layer = 1
242 +
243 + q = torch.randn(batch, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
244 +
245 + # Structurally valid fp8 ds_mla payload: 512B fp8 + 16B f32 scales + 128B
246 + # bf16 rope (random bytes corrupt the scale region and yield NaNs).
247 + nope = (torch.randn(num_blocks, page, 1, 512, device=dev) * 0.1).to(
248 + torch.float8_e4m3fn
249 + )
250 + scales = torch.ones(num_blocks, page, 1, 4, device=dev, dtype=torch.float32)
251 + rope = (torch.randn(num_blocks, page, 1, 64, device=dev) * 0.1).to(torch.bfloat16)
252 + payload = torch.cat(
253 + [
254 + nope.view(torch.uint8).view(num_blocks, page, 1, 512),
255 + scales.view(torch.uint8).view(num_blocks, page, 1, 16),
256 + rope.view(torch.uint8).view(num_blocks, page, 1, 128),
257 + ],
258 + dim=-1,
259 + ).contiguous()
260 + assert payload.shape[-1] == entry and payload.dtype == torch.uint8
261 +
262 + # (A) contiguous reference.
263 + cache_contiguous = payload.clone().contiguous()
264 +
265 + # (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
266 + unified = torch.randint(
267 + 0, 256, (num_blocks, n_layers, page, 1, entry), device=dev, dtype=torch.uint8
268 + )
269 + unified[:, layer].copy_(payload)
270 + cache_view = unified[:, layer]
271 + assert not cache_view.is_contiguous()
272 + assert cache_view.stride(0) == n_layers * page * 1 * entry
273 +
274 + # Sparse indices: each batch uses its own disjoint blocks.
275 + blocks_per_batch = num_blocks // batch
276 + idx = torch.full((batch, 1, topk), -1, device=dev, dtype=torch.int32)
277 + for b in range(batch):
278 + slots: list[int] = []
279 + for blk in range(b * blocks_per_batch, (b + 1) * blocks_per_batch):
280 + slots.extend(blk * page + off for off in range(page))
281 + slots_t = torch.tensor(slots[:topk], device=dev, dtype=torch.int32)
282 + idx[b, 0, : slots_t.numel()] = slots_t
283 +
284 + def run(kc):
285 + meta, num_splits = fm.get_mla_metadata()
286 + out, _ = fm.flash_mla_with_kvcache(
287 + q=q,
288 + k_cache=kc,
289 + block_table=None,
290 + cache_seqlens=None,
291 + head_dim_v=hdv,
292 + tile_scheduler_metadata=meta,
293 + is_fp8_kvcache=True,
294 + indices=idx,
295 + softmax_scale=head_dim**-0.5,
296 + )
297 + return out.clone().float()
298 +
299 + out_ref = run(cache_contiguous)
300 + out_view = run(cache_view)
301 + assert torch.isfinite(out_ref).all()
302 + assert out_ref.abs().max().item() > 0.0
303 + assert (out_ref - out_view).abs().max().item() == 0.0
304 +
305 +
306 +def test_indexer_k_quant_and_cache_into_unified_slot_view():
307 + """indexer_k_quant_and_cache (DeepSeek V3.2/V4 DSA indexer K write) must
308 + write correctly into a per-layer view whose block stride is the full
309 + unified slot, with zero bleed into the other layers' segments."""
310 + from vllm import _custom_ops as ops
311 +
312 + torch.manual_seed(0)
313 + dev = "cuda"
314 + head_dim = 128
315 + quant_block_size = 128
316 + block_size = 64
317 + num_blocks = 16
318 + ntok = 100
319 + # Indexer cache layout per token: head_dim fp8 bytes followed by
320 + # head_dim * 4 / quant_block_size scale bytes.
321 + cache_stride = head_dim + head_dim * 4 // quant_block_size
322 +
323 + k = torch.randn(ntok, head_dim, device=dev, dtype=torch.bfloat16)
324 + slot = torch.randperm(num_blocks * block_size, device=dev, dtype=torch.int64)[:ntok]
325 +
326 + def write(cache):
327 + ops.indexer_k_quant_and_cache(k, cache, slot, quant_block_size, "ue8m0")
328 +
329 + # Contiguous per-layer reference.
330 + ref = torch.zeros(
331 + num_blocks, block_size, cache_stride, device=dev, dtype=torch.uint8
332 + )
333 + write(ref)
334 +
335 + # Unified slot holding three layer pages per block; carve the middle one.
336 + n_layers = 3
337 + layer = 1
338 + unified = torch.zeros(
339 + num_blocks, n_layers, block_size, cache_stride, device=dev, dtype=torch.uint8
340 + )
341 + view = unified[:, layer]
342 + assert not view.is_contiguous()
343 + assert view.stride(0) == n_layers * block_size * cache_stride
344 + write(view)
345 +
346 + assert torch.equal(ref, view.contiguous())
347 + # Zero bleed into the neighbour layers' segments.
348 + assert unified[:, 0].abs().max().item() == 0
349 + assert unified[:, 2].abs().max().item() == 0
350 +
351 +
352 +def test_flashattn_mla_dense_decode_unified_slot_view():
353 + """FA3 decode (FLASH_ATTN_MLA backend) must read a unified-slot
354 + block-major view bit-identically to a contiguous per-layer cache."""
355 + try:
356 + from vllm.vllm_flash_attn import flash_attn_varlen_func
357 + except ImportError:
358 + pytest.skip("vllm_flash_attn is not available")
359 + from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
360 +
361 + if not flash_attn_supports_mla():
362 + pytest.skip("FA3 MLA requires a Hopper device")
363 +
364 + torch.manual_seed(0)
365 + dev = "cuda"
366 + dt = torch.bfloat16
367 + kv_lora_rank = 512
368 + rope_dim = 64
369 + entry = kv_lora_rank + rope_dim # 576
370 + h_q = 16
371 + page = 64
372 + num_blocks = 64
373 + bs = 4
374 + n_layers = 3
375 + layer = 1
376 +
377 + q_pe = torch.randn(bs, h_q, rope_dim, device=dev, dtype=dt) * 0.1
378 + q_nope = torch.randn(bs, h_q, kv_lora_rank, device=dev, dtype=dt) * 0.1
379 + kv_data = torch.randn(num_blocks, page, entry, device=dev, dtype=dt) * 0.1
380 +
381 + # (A) contiguous per-layer reference.
382 + cache_contiguous = kv_data.clone().contiguous()
383 +
384 + # (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
385 + unified = torch.randn(num_blocks, n_layers, page, entry, device=dev, dtype=dt) * 0.1
386 + unified[:, layer].copy_(kv_data)
387 + cache_view = unified[:, layer]
388 + assert not cache_view.is_contiguous()
389 + assert cache_view.stride(0) == n_layers * page * entry
390 +
391 + max_blk = num_blocks // bs
392 + block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
393 + bs, max_blk
394 + )
395 + seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
396 + cu_seqlens_q = torch.arange(bs + 1, device=dev, dtype=torch.int32)
397 +
398 + def run(cache):
399 + kv_c_cache = cache[..., :kv_lora_rank]
400 + k_pe_cache = cache[..., kv_lora_rank:]
401 + out = flash_attn_varlen_func(
402 + q=q_pe,
403 + k=k_pe_cache.unsqueeze(-2), # Add head dim of 1
404 + v=kv_c_cache.unsqueeze(-2), # Add head dim of 1
405 + q_v=q_nope,
406 + max_seqlen_q=1,
407 + cu_seqlens_q=cu_seqlens_q,
408 + max_seqlen_k=int(seq_lens.max().item()),
409 + seqused_k=seq_lens,
410 + block_table=block_table,
411 + softmax_scale=entry**-0.5,
412 + causal=True,
413 + fa_version=3,
414 + )
415 + return out.clone().float()
416 +
417 + out_ref = run(cache_contiguous)
418 + out_view = run(cache_view)
419 + assert torch.isfinite(out_ref).all()
420 + assert out_ref.abs().max().item() > 0.0
421 + assert (out_ref - out_view).abs().max().item() == 0.0
422 +
423 +
424 +def test_flashmla_dense_fp8_decode_unified_slot_view():
425 + """FlashMLA dense fp8 decode (FLASHMLA backend with quantized KV cache)
426 + must read a unified-slot block-major view bit-identically to a contiguous
427 + per-layer fp8 cache."""
428 + import vllm.v1.attention.ops.flashmla as fm
429 +
430 + ok, reason = fm.is_flashmla_dense_supported()
431 + if not ok:
432 + pytest.skip(reason)
433 +
434 + torch.manual_seed(0)
435 + dev = "cuda"
436 + head_dim = 576
437 + hdv = 512
438 + h_q = 128
439 + page = 64
440 + num_blocks = 64
441 + bs = 4
442 + n_layers = 3
443 + layer = 1
444 +
445 + q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
446 + kv_data = (torch.randn(num_blocks, page, head_dim, device=dev) * 0.1).to(
447 + torch.float8_e4m3fn
448 + )
449 +
450 + # (A) contiguous per-layer reference.
451 + cache_contiguous = kv_data.clone().contiguous()
452 +
453 + # (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
454 + unified = (torch.randn(num_blocks, n_layers, page, head_dim, device=dev) * 0.1).to(
455 + torch.float8_e4m3fn
456 + )
457 + unified[:, layer].copy_(kv_data)
458 + cache_view = unified[:, layer]
459 + assert not cache_view.is_contiguous()
460 + assert cache_view.stride(0) == n_layers * page * head_dim
461 +
462 + max_blk = num_blocks // bs
463 + block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
464 + bs, max_blk
465 + )
466 + cache_seqlens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
467 + descale = torch.ones(1, device=dev, dtype=torch.float32)
468 +
469 + def run(kc):
470 + tile_md, num_splits = fm.get_mla_metadata_dense_fp8(cache_seqlens, h_q, 1)
471 + out, _ = fm.flash_mla_with_kvcache_fp8(
472 + q=q,
473 + k_cache=kc.unsqueeze(-2), # Add head dim of 1
474 + block_table=block_table,
475 + cache_seqlens=cache_seqlens,
476 + head_dim_v=hdv,
477 + tile_scheduler_metadata=tile_md,
478 + num_splits=num_splits,
479 + softmax_scale=head_dim**-0.5,
480 + causal=True,
481 + descale_q=descale,
482 + descale_k=descale,
483 + )
484 + return out.clone().float()
485 +
486 + out_ref = run(cache_contiguous)
487 + out_view = run(cache_view)
488 + assert torch.isfinite(out_ref).all()
489 + assert out_ref.abs().max().item() > 0.0
490 + assert (out_ref - out_view).abs().max().item() == 0.0
491 +
492 +
493 +def test_flashinfer_mla_dense_fp8_decode_unified_slot_view():
494 + """FlashInfer MLA dense decode with an fp8 KV cache must read a
495 + unified-slot block-major view bit-identically to a contiguous per-layer
496 + cache."""
497 + try:
498 + from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
499 + except ImportError:
500 + pytest.skip("flashinfer is not available")
501 + from vllm.platforms import current_platform
502 +
503 + if not current_platform.is_device_capability_family(100):
504 + pytest.skip("FlashInfer trtllm-gen MLA requires sm100")
505 +
506 + torch.manual_seed(0)
507 + dev = "cuda"
508 + kv_lora_rank = 512
509 + qk_rope_head_dim = 64
510 + qk_nope_head_dim = 128
511 + head_dim = kv_lora_rank + qk_rope_head_dim # 576
512 + num_qo_heads = 128
513 + page = 64
514 + num_blocks = 64
515 + bs = 4
516 + n_layers = 3
517 + layer = 1
518 +
519 + # With a quantized KV cache the decode query is quantized to fp8 as well
520 + # (trtllm-gen has no bf16-query x fp8-cache decode kernel).
521 + q = (torch.randn(bs, 1, num_qo_heads, head_dim, device=dev) * 0.1).to(
522 + torch.float8_e4m3fn
523 + )
524 + kv_data = (torch.randn(num_blocks, 1, page, head_dim, device=dev) * 0.1).to(
525 + torch.float8_e4m3fn
526 + )
527 +
528 + # (A) contiguous per-layer reference.
529 + kv_contiguous = kv_data.clone().contiguous()
530 +
531 + # (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
532 + unified = (
533 + torch.randn(num_blocks, n_layers, 1, page, head_dim, device=dev) * 0.1
534 + ).to(torch.float8_e4m3fn)
535 + unified[:, layer].copy_(kv_data)
536 + kv_view = unified[:, layer]
537 + assert not kv_view.is_contiguous()
538 + assert kv_view.stride(0) == n_layers * 1 * page * head_dim
539 +
540 + max_blk = num_blocks // bs
541 + block_tables = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
542 + bs, max_blk
543 + )
544 + seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
545 + ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=dev)
546 + scale = head_dim**-0.5
547 +
548 + def run(kv):
549 + return trtllm_batch_decode_with_kv_cache_mla(
550 + query=q,
551 + kv_cache=kv,
552 + workspace_buffer=ws,
553 + qk_nope_head_dim=qk_nope_head_dim,
554 + kv_lora_rank=kv_lora_rank,
555 + qk_rope_head_dim=qk_rope_head_dim,
556 + block_tables=block_tables,
557 + seq_lens=seq_lens,
558 + max_seq_len=int(seq_lens.max().item()),
559 + bmm1_scale=scale,
560 + bmm2_scale=1.0,
561 + ).clone()
562 +
563 + out_ref = run(kv_contiguous).float()
564 + out_view = run(kv_view).float()
565 + assert torch.isfinite(out_ref).all()
566 + assert (out_ref - out_view).abs().max().item() == 0.0
tests/kernels/attention/test_triton_decode_attention.py
+92
@@ -231,3 +231,95 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
231
232 # FP8 tolerances match test_mla_backends.py test_backend_correctness.
233 torch.testing.assert_close(o_ref, o_fp8, atol=5e-1, rtol=1e-2)
234 +
235 +
236 +@pytest.mark.parametrize(
237 + "H_Q,H_KV,D_QK,D_V,is_mla",
238 + [
239 + (16, 1, 576, 512, True), # MLA path (grouped kernel, v = trans(k))
240 + (32, 8, 128, 128, False), # GQA path (grouped kernel)
241 + (32, 32, 128, 128, False), # MHA path (normal kernel)
242 + ],
243 +)
244 +@pytest.mark.parametrize("PAGE_SIZE", [16])
245 +def test_decode_attention_cross_layer_view(H_Q, H_KV, D_QK, D_V, is_mla, PAGE_SIZE):
246 + """The kernel must honor the cache's page-dim stride, not assume pages are
247 + packed back-to-back. A per-layer view into a cross-layer (block-major)
248 + cache has stride(0) inflated by num_layers; outputs must match a
249 + contiguous cache holding the same data exactly."""
250 + B = 3
251 + seq_len = 1027
252 + CACHE_SIZE = 16384
253 + NUM_LAYERS = 3
254 + LAYER_IDX = 1
255 + dtype = torch.bfloat16
256 + sm_scale = 1.0 / (D_QK**0.5)
257 + num_kv_splits = 8
258 + num_pages = CACHE_SIZE // PAGE_SIZE
259 +
260 + num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
261 + req_to_page = torch.randint(
262 + 0, num_pages, (B, num_pages_per_batch), device=DEVICE_TYPE
263 + )
264 +
265 + q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
266 + b_seq_len = torch.full((B,), seq_len, device=DEVICE_TYPE)
267 +
268 + # Reference: contiguous paged cache.
269 + k_ref = torch.randn(
270 + num_pages, PAGE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE
271 + )
272 + if is_mla:
273 + v_ref = k_ref[..., :D_V]
274 + else:
275 + v_ref = torch.randn(
276 + num_pages, PAGE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE
277 + )
278 +
279 + # Cross-layer cache: all layers' pages for a block are adjacent. The
280 + # per-layer view has the same shape as the contiguous cache but
281 + # stride(0) is NUM_LAYERS x larger. Neighbor layers hold random data so
282 + # any packed-pages addressing reads garbage rather than zeros.
283 + k_xl = torch.randn(
284 + num_pages, NUM_LAYERS, PAGE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE
285 + )
286 + k_view = k_xl[:, LAYER_IDX]
287 + k_view.copy_(k_ref)
288 + assert k_view.stride(0) == NUM_LAYERS * PAGE_SIZE * H_KV * D_QK
289 + if is_mla:
290 + v_view = k_view[..., :D_V]
291 + else:
292 + v_xl = torch.randn(
293 + num_pages, NUM_LAYERS, PAGE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE
294 + )
295 + v_view = v_xl[:, LAYER_IDX]
296 + v_view.copy_(v_ref)
297 +
298 + def run(k_buffer, v_buffer):
299 + o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
300 + lse = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
301 + attn_logits = torch.empty(
302 + (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
303 + )
304 + decode_attention_fwd(
305 + q,
306 + k_buffer,
307 + v_buffer,
308 + o,
309 + lse,
310 + req_to_page,
311 + b_seq_len,
312 + attn_logits,
313 + num_kv_splits,
314 + sm_scale,
315 + PAGE_SIZE,
316 + is_mla=is_mla,
317 + )
318 + return o, lse
319 +
320 + o_ref, lse_ref = run(k_ref, v_ref)
321 + o_xl, lse_xl = run(k_view, v_view)
322 +
323 + # Same data and same compute order; only addressing differs.
324 + assert torch.equal(o_ref, o_xl)
325 + assert torch.equal(lse_ref, lse_xl)
tests/v1/kv_connector/unit/test_kv_cache_layout.py
+35 -3
@@ -1,10 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import pytest
5
5 -def test_mla_backend_rejects_cross_layer_kv_cache():
6 - """MLA backends return identity permutation (layers dim first)
7 - to signal cross-layer KV cache is unsupported."""
6 +
7 +def test_mla_common_backend_rejects_cross_layer_kv_cache():
8 + """MLACommonBackend defaults to the identity permutation (layers dim
9 + first) so MLA backends whose decode kernels are not verified to honor
10 + the cache's block-dim stride stay opted out of cross-layer KV cache."""
11 from vllm.model_executor.layers.attention.mla_attention import (
12 MLACommonBackend,
13 )
@@ -19,6 +22,35 @@ def test_mla_backend_rejects_cross_layer_kv_cache():
22 ) == (0, 1, 2)
23
24
25 +@pytest.mark.parametrize(
26 + "backend_path",
27 + [
28 + "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
29 + "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend",
30 + "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend",
31 + "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend",
32 + "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend",
33 + ],
34 +)
35 +def test_verified_mla_backends_support_cross_layer_kv_cache(backend_path):
36 + """Backends whose decode kernels honor the cache's block-dim stride opt
37 + in to the cross-layer layout with a non-identity permutation placing
38 + num_blocks first in physical layout."""
39 + module_path, name = backend_path.rsplit(".", 1)
40 + backend = getattr(
41 + pytest.importorskip(module_path, reason="backend deps unavailable"), name
42 + )
43 +
44 + stride_order = backend.get_kv_cache_stride_order(include_num_layers_dimension=True)
45 + assert stride_order == (1, 0, 2, 3)
46 + assert stride_order[0] != 0 # num_blocks first => cross-layer supported
47 + assert backend.get_kv_cache_stride_order(include_num_layers_dimension=False) == (
48 + 0,
49 + 1,
50 + 2,
51 + )
52 +
53 +
54 def test_deepseek_v32_indexer_rejects_cross_layer_kv_cache():
55 """DeepseekV32Indexer returns identity permutation (layers dim first)
56 to signal cross-layer KV cache is unsupported."""
vllm/model_executor/layers/attention/mla_attention.py
+3 -3
@@ -1212,9 +1212,9 @@ class MLACommonBackend(AttentionBackend):
1212 include_num_layers_dimension: bool = False,
1213 ) -> tuple[int, ...]:
1214 if include_num_layers_dimension:
1215 - # MLA kernels require contiguous per-layer KV cache views.
1216 - # Identity permutation keeps num_layers first in physical
1217 - # layout, signaling cross-layer allocation is unsupported.
1215 + # Default to identity permutation to signal cross-layer allocation
1216 + # is unsupported. Each MLA backend must opt in to support cross-layer
1217 + # allocation by overriding this method.
1218 return (0, 1, 2, 3)
1219 return (0, 1, 2)
1220
vllm/v1/attention/backends/mla/cutlass_mla.py
+8
@@ -49,6 +49,14 @@ class CutlassMLABackend(MLACommonBackend):
49 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
50 return [128]
51
52 + @staticmethod
53 + def get_kv_cache_stride_order(
54 + include_num_layers_dimension: bool = False,
55 + ) -> tuple[int, ...]:
56 + if include_num_layers_dimension:
57 + return (1, 0, 2, 3)
58 + return (0, 1, 2)
59 +
60 @staticmethod
61 def get_name() -> str:
62 return "CUTLASS_MLA"
vllm/v1/attention/backends/mla/flashattn_mla.py
+8
@@ -52,6 +52,14 @@ class FlashAttnMLABackend(MLACommonBackend):
52 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
53 return [MultipleOf(16)]
54
55 + @staticmethod
56 + def get_kv_cache_stride_order(
57 + include_num_layers_dimension: bool = False,
58 + ) -> tuple[int, ...]:
59 + if include_num_layers_dimension:
60 + return (1, 0, 2, 3)
61 + return (0, 1, 2)
62 +
63 @staticmethod
64 def get_name() -> str:
65 return "FLASH_ATTN_MLA"
vllm/v1/attention/backends/mla/flashinfer_mla.py
+8
@@ -49,6 +49,14 @@ class FlashInferMLABackend(MLACommonBackend):
49 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
50 return [32, 64]
51
52 + @staticmethod
53 + def get_kv_cache_stride_order(
54 + include_num_layers_dimension: bool = False,
55 + ) -> tuple[int, ...]:
56 + if include_num_layers_dimension:
57 + return (1, 0, 2, 3)
58 + return (0, 1, 2)
59 +
60 @staticmethod
61 def get_name() -> str:
62 return "FLASHINFER_MLA"
vllm/v1/attention/backends/mla/flashmla.py
+8
@@ -58,6 +58,14 @@ class FlashMLABackend(MLACommonBackend):
58 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
59 return [64]
60
61 + @staticmethod
62 + def get_kv_cache_stride_order(
63 + include_num_layers_dimension: bool = False,
64 + ) -> tuple[int, ...]:
65 + if include_num_layers_dimension:
66 + return (1, 0, 2, 3)
67 + return (0, 1, 2)
68 +
69 @staticmethod
70 def get_name() -> str:
71 return "FLASHMLA"
vllm/v1/attention/backends/mla/triton_mla.py
+8
@@ -57,6 +57,14 @@ class TritonMLABackend(MLACommonBackend):
57 return True
58 return block_size % 16 == 0
59
60 + @staticmethod
61 + def get_kv_cache_stride_order(
62 + include_num_layers_dimension: bool = False,
63 + ) -> tuple[int, ...]:
64 + if include_num_layers_dimension:
65 + return (1, 0, 2, 3)
66 + return (0, 1, 2)
67 +
68 @staticmethod
69 def get_name() -> str:
70 return "TRITON_MLA"
vllm/v1/attention/ops/triton_decode_attention.py
+34 -7
@@ -56,6 +56,15 @@ def tanh(x):
56 return 2 * tl.sigmoid(2 * x) - 1
57
58
59 +def _page_stride(buf, page_size):
60 + # Stride between pages. 4D buffers have a page dim; 3D buffers pack pages
61 + # along the token dim, so split it out first. Read the real stride (a
62 + # cross-layer view has gaps), don't assume PAGE_SIZE * token stride.
63 + if buf.ndim == 3:
64 + buf = buf.unflatten(-3, (-1, page_size))
65 + return buf.stride(-4)
66 +
67 +
68 @triton.jit
69 def _fwd_kernel_stage1(
70 Q,
@@ -68,8 +77,10 @@ def _fwd_kernel_stage1(
77 stride_req_to_tokens_b,
78 stride_qbs,
79 stride_qh,
80 + stride_buf_kpbs,
81 stride_buf_kbs,
82 stride_buf_kh,
83 + stride_buf_vpbs,
84 stride_buf_vbs,
85 stride_buf_vh,
86 stride_mid_ob,
@@ -123,9 +134,11 @@ def _fwd_kernel_stage1(
134 mask=offs_n < split_kv_end,
135 other=0,
136 )
126 - kv_loc = kv_page_number * PAGE_SIZE + offs_n % PAGE_SIZE
137 + kv_in_page = offs_n % PAGE_SIZE
138 offs_buf_k = (
128 - kv_loc[:, None] * stride_buf_kbs
139 + (kv_page_number * stride_buf_kpbs + kv_in_page * stride_buf_kbs)[
140 + :, None
141 + ]
142 + cur_kv_head * stride_buf_kh
143 + offs_d[None, :]
144 )
@@ -145,7 +158,9 @@ def _fwd_kernel_stage1(
158 qk = tl.where(offs_n < split_kv_end, qk, float("-inf"))
159
160 offs_buf_v = (
148 - kv_loc[:, None] * stride_buf_vbs
161 + (kv_page_number * stride_buf_vpbs + kv_in_page * stride_buf_vbs)[
162 + :, None
163 + ]
164 + cur_kv_head * stride_buf_vh
165 + offs_dv[None, :]
166 )
@@ -235,8 +250,10 @@ def _decode_att_m_fwd(
250 Req_to_tokens.stride(0),
251 q.stride(0),
252 q.stride(1),
253 + _page_stride(k_buffer, page_size),
254 k_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
255 k_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
256 + _page_stride(v_buffer, page_size),
257 v_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
258 v_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
259 att_out.stride(0),
@@ -270,8 +287,10 @@ def _fwd_grouped_kernel_stage1(
287 stride_req_to_tokens_b,
288 stride_qbs,
289 stride_qh,
290 + stride_buf_kpbs,
291 stride_buf_kbs,
292 stride_buf_kh,
293 + stride_buf_vpbs,
294 stride_buf_vbs,
295 stride_buf_vh,
296 stride_mid_ob,
@@ -357,10 +376,12 @@ def _fwd_grouped_kernel_stage1(
376 other=0,
377 cache_modifier=".ca",
378 )
360 - kv_loc = kv_page_number * PAGE_SIZE + offs_n % PAGE_SIZE
379 + kv_off_k = (
380 + kv_page_number * stride_buf_kpbs + (offs_n % PAGE_SIZE) * stride_buf_kbs
381 + )
382
383 # explicitly facilitate overlapping load/compute
363 - offs_buf_k = kv_loc[None, :] * stride_buf_kbs + base_offs_k
384 + offs_buf_k = kv_off_k[None, :] + base_offs_k
385 k = tl.load(
386 K_Buffer + offs_buf_k,
387 mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]),
@@ -372,7 +393,7 @@ def _fwd_grouped_kernel_stage1(
393 k = (k.to(tl.float32) * ks).to(q.dtype)
394 qk = tl.dot(q, k.to(q.dtype))
395 if BLOCK_DPE > 0:
375 - offs_buf_kpe = kv_loc[None, :] * stride_buf_kbs + base_offs_kpe
396 + offs_buf_kpe = kv_off_k[None, :] + base_offs_kpe
397 kpe = tl.load(
398 K_Buffer + offs_buf_kpe,
399 mask=(offs_n[None, :] < split_kv_end) & (mask_dpe[:, None]),
@@ -392,7 +413,11 @@ def _fwd_grouped_kernel_stage1(
413 )
414
415 if not IS_MLA:
395 - offs_buf_v = kv_loc[:, None] * stride_buf_vbs + base_offs_v
416 + kv_off_v = (
417 + kv_page_number * stride_buf_vpbs
418 + + (offs_n % PAGE_SIZE) * stride_buf_vbs
419 + )
420 + offs_buf_v = kv_off_v[:, None] + base_offs_v
421 v = tl.load(
422 V_Buffer + offs_buf_v,
423 mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]),
@@ -517,8 +542,10 @@ def _decode_grouped_att_m_fwd(
542 Req_to_tokens.stride(0),
543 q.stride(0),
544 q.stride(1),
545 + _page_stride(k_buffer, page_size),
546 k_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
547 k_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
548 + _page_stride(v_buffer, page_size),
549 v_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
550 v_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
551 att_out.stride(0),