@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 51878e5b6e

[2/N][KV-Cache Layout Refactor] Pack K/V into the content dim across attention backends (#44455)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>

Lucas Wilkinson committed Jul 11, 2026 at 11:11 UTC 51878e5b6e4cf27352c427b8508b4f338ba39c9c
31 files changed +664 -604
csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
+28 -29
@@ -288,16 +288,15 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
288 int64_t const* __restrict__ positions, // [N] i64
289 int64_t const* __restrict__ slot_mapping, // main K/V slots or nullptr
290 int64_t const* __restrict__ index_slot_mapping, // index K slots/nullptr
291 - cache_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
291 + cache_t* __restrict__ kv_cache, // [nb,nkv,bs,2*128] or nullptr
292 out_idx_t* __restrict__ index_cache, // [nb*bs, 128]; scalar_t or e4m3 byte
293 float const eps, int const rotary_dim, int const num_tokens, int const nq,
294 int const nkv, int const niq, int const block_size,
295 - // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
296 - // The head_dim (last) dim is always innermost-contiguous (stride 1), so the
297 - // NHD/HND layout choice is fully captured by these four strides: NHD keeps
298 - // s_token < s_head, HND swaps them. dim_base addresses head_dim directly.
299 - int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
300 - int64_t const kv_s_head) {
295 + // kv_cache strides (in elements) for logical shape [nb, nkv, bs, 2*128].
296 + // The content (last) dim is always innermost-contiguous (stride 1), so the
297 + // NHD/HND layout choice is captured by the head/token strides.
298 + int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
299 + int64_t const kv_s_dim) {
300 #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
301 // _typeConvert<BFloat16> is unavailable on pre-Ampere; the M3 kernel only
302 // runs with bf16/fp16 inputs in practice. Discard the bf16 body there.
@@ -438,16 +437,16 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
437 storeElems<scalar_t>(index_cache + sm * kHeadDim + dim_base, elems);
438 }
439 } else if (isK || isV) {
441 - // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
440 + // kv_cache logical shape [num_blocks, nkv, block_size, 2*head_dim].
441 // Paging is logical (block = sm/block_size, token = sm%block_size);
442 // the physical NHD/HND layout is honoured via the passed strides.
443 int64_t const b = sm / block_size;
444 int64_t const t = sm % block_size;
445 int const kv = isK ? 0 : 1;
447 - int64_t const off =
448 - b * kv_s_block + kv * kv_s_kv + t * kv_s_token + head * kv_s_head;
449 - storeCacheElems<scalar_t, cache_t, kv_dt>(kv_cache + off + dim_base,
450 - elems);
446 + int64_t const off = b * kv_s_block + head * kv_s_head +
447 + t * kv_s_token +
448 + (kv * kHeadDim + dim_base) * kv_s_dim;
449 + storeCacheElems<scalar_t, cache_t, kv_dt>(kv_cache + off, elems);
450 }
451 }
452 }
@@ -474,8 +473,8 @@ void launchFusedMiniMaxM3(
473 int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache,
474 float const eps, int const rotary_dim, int const num_tokens, int const nq,
475 int const nkv, int const niq, int const block_size,
477 - int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
478 - int64_t const kv_s_head, bool const has_index, bool const insert_kv,
476 + int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
477 + int64_t const kv_s_dim, bool const has_index, bool const insert_kv,
478 bool const fp8_idx, cudaStream_t stream) {
479 // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
480 // void* pointers per instantiation in the LAUNCH macro.
@@ -517,20 +516,20 @@ void launchFusedMiniMaxM3(
516 iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping, \
517 index_slot_mapping, kv_cache, reinterpret_cast<OUT_T*>(index_cache), \
518 eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
520 - kv_s_kv, kv_s_token, kv_s_head)
519 + kv_s_head, kv_s_token, kv_s_dim)
520 #else
521 // ROCm: standard kernel launch syntax (no PDL/stream serialization).
522 // clang-format off
523 #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T) \
525 - fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
524 + fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
525 IS_SPARSE, INSERT, FP8> \
526 <<<grid, kBlockSize, 0, stream>>>( \
527 qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, \
528 k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
529 slot_mapping, index_slot_mapping, kv_cache, \
530 reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim, \
532 - num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv, \
533 - kv_s_token, kv_s_head)
531 + num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_head, \
532 + kv_s_token, kv_s_dim)
533 // clang-format on
534 #endif
535
@@ -584,8 +583,8 @@ void launchFusedMiniMaxM3(
583 ? reinterpret_cast<void*>(index_cache->data_ptr()) \
584 : nullptr, \
585 static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq, \
587 - nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv, kv_s_token, \
588 - kv_s_head, has_index, insert_kv, fp8_idx, stream)
586 + nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_head, \
587 + kv_s_token, kv_s_dim, has_index, insert_kv, fp8_idx, stream)
588
589 // ────────────────────────────────────────────────────────────────────────────
590 // Torch op wrapper
@@ -602,7 +601,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
601 int64_t num_index_heads, // niq; 0 => dense
602 std::optional<torch::stable::Tensor> slot_mapping, // [N] i64
603 std::optional<torch::stable::Tensor> index_slot_mapping, // [N] i64
605 - std::optional<torch::stable::Tensor> kv_cache, // [nb,2,bs,nkv,128]
604 + std::optional<torch::stable::Tensor> kv_cache, // [nb,nkv,bs,2*128]
605 std::optional<torch::stable::Tensor> index_cache, // [nb,bs,128]
606 int64_t block_size,
607 std::optional<torch::stable::Tensor> q_out, // [N, nq*128] contiguous
@@ -673,11 +672,11 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
672 index_k_norm_weight->numel() == kHeadDim,
673 "index norm weights must have 128 elements");
674 }
676 - // kv_cache strides (logical shape [nb, 2, bs, nkv, head_dim]). Read straight
675 + // kv_cache strides (logical shape [nb, nkv, bs, 2*head_dim]). Read straight
676 // off the tensor so the kernel honours whatever physical layout the attention
678 - // backend allocated (NHD: stride order (0,1,2,3,4); HND: (0,1,3,2,4)). No new
677 + // backend allocated (NHD: stride order (0,2,1,3); HND: (0,1,2,3)). No new
678 // op argument is needed -- the strides ride along with the tensor itself.
680 - int64_t kv_s_block = 0, kv_s_kv = 0, kv_s_token = 0, kv_s_head = 0;
679 + int64_t kv_s_block = 0, kv_s_head = 0, kv_s_token = 0, kv_s_dim = 0;
680 torch::stable::Tensor const* effective_index_slot_mapping = nullptr;
681 if (insert_kv) {
682 STD_TORCH_CHECK(
@@ -707,13 +706,13 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
706 index_cache->scalar_type() ==
707 torch::headeronly::ScalarType::Float8_e4m3fn),
708 "insert mode requires index_cache matching qkv dtype or fp8 e4m3");
710 - STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
711 - "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
712 - "head_dim (stride(4)==1)");
709 + STD_TORCH_CHECK(kv_cache->dim() == 4 && kv_cache->stride(3) == 1,
710 + "kv_cache must be [nb,nkv,bs,2*head_dim] with contiguous "
711 + "content dim (stride(3)==1)");
712 kv_s_block = kv_cache->stride(0);
714 - kv_s_kv = kv_cache->stride(1);
713 + kv_s_head = kv_cache->stride(1);
714 kv_s_token = kv_cache->stride(2);
716 - kv_s_head = kv_cache->stride(3);
715 + kv_s_dim = kv_cache->stride(3);
716 effective_index_slot_mapping = index_slot_mapping.has_value()
717 ? &index_slot_mapping.value()
718 : &slot_mapping.value();
tests/compile/passes/test_fusion_attn.py
+6 -3
@@ -111,7 +111,11 @@ class AttentionQuantPatternModel(torch.nn.Module):
111 # Fetch the attention backend and kv cache shape and stride order
112 attn_backend = self.attn.attn_backend
113 kv_cache_shape = attn_backend.get_kv_cache_shape(
114 - num_blocks, self.block_size, self.num_kv_heads, self.head_size
114 + num_blocks,
115 + self.block_size,
116 + self.num_kv_heads,
117 + self.head_size,
118 + cache_dtype_str=self.attn.kv_cache_dtype,
119 )
120 try:
121 kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
@@ -125,11 +129,10 @@ class AttentionQuantPatternModel(torch.nn.Module):
129
130 # Create dummy KV cache
131 raw_tensor = torch.zeros(
128 - 2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
132 + kv_cache_shape,
133 dtype=self.attn.kv_cache_torch_dtype,
134 device=self.device,
135 )
132 - raw_tensor = raw_tensor.view(kv_cache_shape)
136 kv_cache = raw_tensor.permute(*inv_order)
137
138 self.attn.kv_cache = kv_cache
tests/kernels/attention/test_cache.py
+3 -5
@@ -10,7 +10,7 @@ from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
10 from vllm import _custom_ops as ops
11 from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_dequantize
12 from vllm.platforms import current_platform
13 -from vllm.utils.torch_utils import nvfp4_kv_cache_split_views, set_random_seed
13 +from vllm.utils.torch_utils import nvfp4_split_data_scale, set_random_seed
14
15 COPYING_DIRECTION = [("cuda", "cpu"), ("cuda", "cuda"), ("cpu", "cuda")]
16 DTYPES = [torch.bfloat16, torch.float]
@@ -255,10 +255,8 @@ def test_reshape_and_cache_flash(
255 nvfp4_key_data = None
256 nvfp4_value_data = None
257 if kv_cache_dtype == "nvfp4":
258 - (nvfp4_key_data,), (key_scale_cache,) = nvfp4_kv_cache_split_views(key_cache)
259 - (nvfp4_value_data,), (value_scale_cache,) = nvfp4_kv_cache_split_views(
260 - value_cache
261 - )
258 + nvfp4_key_data, key_scale_cache = nvfp4_split_data_scale(key_cache)
259 + nvfp4_value_data, value_scale_cache = nvfp4_split_data_scale(value_cache)
260
261 if kv_cache_dtype == "nvfp4":
262 # Global scale = amax / 448 (per-tensor)
tests/kernels/attention/test_flashinfer_trtllm_attention.py
+15 -15
@@ -13,7 +13,7 @@ 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_kv_cache_split_views,
16 + nvfp4_split_data_scale,
17 set_random_seed,
18 )
19
@@ -74,17 +74,18 @@ def make_nvfp4_kv_cache(
74 kv_scale_val, dtype=torch.float32, device=kv_bf16_hnd.device
75 )
76
77 - # Allocate in HND physical order, permute to NHD logical order.
78 - # hnd_order swaps dims 2↔3; it is its own inverse.
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 - hnd_order = (0, 1, 3, 2, 4)
81 - kv_cache = torch.zeros(
82 - (num_blocks, 2, num_kv_heads, block_size, full_dim),
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,
85 - ).permute(*hnd_order)
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
87 - # Flatten NHD [N, T, H, D] → token tensors [N*T, H, D] for the kernel.
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]
@@ -98,22 +99,21 @@ def make_nvfp4_kv_cache(
99 )
100 slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=kv_bf16_hnd.device)
101
101 - # reshape_and_cache_flash: kernel receives kv_cache[:, 0] and [:, 1]
102 - # (full K/V buffers containing both data and scale).
102 torch.ops._C_cache_ops.reshape_and_cache_flash(
103 k_tokens,
104 v_tokens,
106 - kv_cache[:, 0],
107 - kv_cache[:, 1],
105 + k_view_nhd,
106 + v_view_nhd,
107 slot_mapping,
108 "nvfp4",
109 kv_scale_tensor,
110 kv_scale_tensor,
111 )
112
114 - # Split in HND order for trtllm kernel (expects HND numTokensPerPage).
115 - kv_cache_hnd = kv_cache.permute(*hnd_order)
116 - (k_data, v_data), (k_scales, v_scales) = nvfp4_kv_cache_split_views(kv_cache_hnd)
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(
tests/kernels/attention/test_minimax_m3.py
+40 -37
@@ -662,8 +662,9 @@ def _reference_sparse_attn(
662 positions = torch.arange(seq_len, device="cuda")
663 pages = block_table[req_id, positions // BLOCK_SIZE]
664 rows = positions % BLOCK_SIZE
665 - k_req = kv_cache[pages, 0, rows]
666 - v_req = kv_cache[pages, 1, rows].float()
665 + kv_req = kv_cache[pages, :, rows]
666 + k_req = kv_req[..., :HEAD_DIM]
667 + v_req = kv_req[..., HEAD_DIM:].float()
668
669 q_pos = prefix_len + torch.arange(q_len, device="cuda")
670 key_blocks = positions // BLOCK_SIZE
@@ -791,15 +792,15 @@ def test_main_backend_layout_contract():
792 flash_attn-style stride order for each layout."""
793 nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
794 logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
794 - assert logical == (nb, 2, bs, h, d)
795 - # The old HND-ordered shape is no longer the logical shape.
796 - assert logical != (nb, 2, h, bs, d)
795 + assert logical == (nb, h, bs, 2 * d)
796 + # The old separate K/V-axis shape is no longer the logical shape.
797 + assert logical != (nb, 2, bs, h, d)
798
799 try:
800 set_kv_cache_layout("HND")
800 - assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 3, 2, 4)
801 + assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3)
802 set_kv_cache_layout("NHD")
802 - assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3, 4)
803 + assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 2, 1, 3)
804 finally:
805 set_kv_cache_layout(None)
806
@@ -829,7 +830,7 @@ def test_main_backend_unknown_layout_raises(monkeypatch):
830
831
832 def test_indexer_backend_stride_order_is_identity():
832 - """The 3-dim indexer cache must not inherit the parent's 5-element stride
833 + """The 3-dim indexer cache must not inherit the parent's 4-element stride
834 order; it overrides to the 3-element identity so the allocator keeps the
835 contiguous layout."""
836 assert MiniMaxM3IndexerBackend.get_kv_cache_stride_order() == (0, 1, 2)
@@ -848,9 +849,9 @@ def test_indexer_backend_stride_order_is_identity():
849 assert _stride_order_for(MiniMaxM3IndexerBackend, len(indexer_shape)) == (0, 1, 2)
850
851
851 -def test_hnd_allocation_is_byte_identical_to_transpose():
852 - """Under HND the backend-visible logical view is byte-identical to the
853 - pre-change allocate-HND-then-transpose(2, 3) workaround."""
852 +def test_hnd_allocation_is_packed_head_major():
853 + """Under HND the backend-visible logical view is the packed head-major
854 + physical allocation."""
855 nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
856 logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
857 try:
@@ -860,21 +861,22 @@ def test_hnd_allocation_is_byte_identical_to_transpose():
861 set_kv_cache_layout(None)
862
863 physical_shape = tuple(logical[i] for i in stride_order)
863 - # The physical (permuted) shape equals the old hardcoded HND shape.
864 - assert physical_shape == (nb, 2, h, bs, d)
864 + assert physical_shape == (nb, h, bs, 2 * d)
865
866 inv_order = [stride_order.index(i) for i in range(len(stride_order))]
867 raw = torch.empty(physical_shape, device="cuda", dtype=DTYPE)
868 view = raw.permute(*inv_order)
869 - expected = raw.view((nb, 2, h, bs, d)).transpose(2, 3)
869 + expected = raw.view((nb, h, bs, 2 * d))
870
871 assert view.shape == expected.shape
872 assert view.stride() == expected.stride()
873 assert view.storage_offset() == expected.storage_offset()
874
875 - # Negative: the identity (wrong) stride order under HND does not reproduce
876 - # the transpose view.
877 - wrong_view = raw.view(logical)
875 + # Negative: the NHD stride order under HND does not reproduce the
876 + # head-major view.
877 + wrong_order = (0, 2, 1, 3)
878 + wrong_inv = [wrong_order.index(i) for i in range(len(wrong_order))]
879 + wrong_view = raw.view(tuple(logical[i] for i in wrong_order)).permute(*wrong_inv)
880 assert wrong_view.stride() != expected.stride()
881
882
@@ -1037,14 +1039,18 @@ def test_decode_wrong_layout_breaks_parity():
1039 seq_lens_list = (130, 257)
1040 q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(seq_lens_list)
1041
1040 - # Physical HND storage [blocks, 2, heads, block, dim].
1042 + # Physical HND storage [blocks, heads, block, packed_kv_dim].
1043 phys = torch.randn(
1042 - (num_pages, 2, NUM_KV_HEADS, BLOCK_SIZE, HEAD_DIM), device="cuda", dtype=DTYPE
1044 + (num_pages, NUM_KV_HEADS, BLOCK_SIZE, 2 * HEAD_DIM),
1045 + device="cuda",
1046 + dtype=DTYPE,
1047 + )
1048 + # Correct logical packed-HND view vs. the same bytes mislabeled as NHD
1049 + # physical storage and then exposed as a logical cache.
1050 + correct = phys
1051 + wrong = phys.reshape(num_pages, BLOCK_SIZE, NUM_KV_HEADS, 2 * HEAD_DIM).permute(
1052 + 0, 2, 1, 3
1053 )
1044 - # Correct logical-NHD view (strided) vs. the same bytes mislabeled as a
1045 - # contiguous-NHD cache — same shape, different content mapping.
1046 - correct = phys.permute(0, 1, 3, 2, 4)
1047 - wrong = phys.reshape(num_pages, 2, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM)
1054
1055 q_lens_t = torch.ones(len(seq_lens_list), device="cuda", dtype=torch.int32)
1056 prefix_lens = seq_lens - q_lens_t
@@ -1072,9 +1078,9 @@ def _make_attn_group(backend, spec):
1078 def test_main_cache_byte_identical_through_production_allocator():
1079 """AC-2: drive the real allocator (`_reshape_kv_cache`) for the M3 main
1080 `FullAttentionSpec` under HND and assert the backend-visible view has the
1075 - same shape, stride, and storage offset as the pre-change
1076 - allocate-HND-then-transpose path; the indexer `MLAAttentionSpec` allocates
1077 - through the same path to its 3-dim shape."""
1081 + same shape, stride, and storage offset as the packed-HND allocation; the
1082 + indexer `MLAAttentionSpec` allocates through the same path to its 3-dim
1083 + shape."""
1084 nb = 4
1085 spec = FullAttentionSpec(
1086 block_size=BLOCK_SIZE,
@@ -1092,8 +1098,7 @@ def test_main_cache_byte_identical_through_production_allocator():
1098 set_kv_cache_layout(None)
1099 view = kv_caches["main"]
1100
1095 - oracle = raw.view(DTYPE).view((nb, 2, NUM_KV_HEADS, BLOCK_SIZE, HEAD_DIM))
1096 - oracle = oracle.transpose(2, 3)
1101 + oracle = raw.view(DTYPE).view((nb, NUM_KV_HEADS, BLOCK_SIZE, 2 * HEAD_DIM))
1102 assert tuple(view.shape) == tuple(oracle.shape)
1103 assert view.stride() == oracle.stride()
1104 assert view.storage_offset() == oracle.storage_offset()
@@ -1119,13 +1124,13 @@ def test_main_cache_byte_identical_through_production_allocator():
1124
1125
1126 def test_indexer_inherited_stride_order_trips_allocator_assert():
1122 - """AC-4 negative: without the indexer override, the inherited 5-element
1127 + """AC-4 negative: without the indexer override, the inherited 4-element
1128 stride order trips the allocator's `len(stride_order) == len(shape)` assert
1129 for the 3-dim indexer shape; the `AssertionError` is NOT swallowed by the
1130 allocator's `(AttributeError, NotImplementedError)` fallback."""
1131
1132 class _BrokenIndexerBackend(MiniMaxM3IndexerBackend):
1128 - # Simulate inheriting the parent's 5-element stride order.
1133 + # Simulate inheriting the parent's 4-element stride order.
1134 get_kv_cache_stride_order = staticmethod(
1135 MiniMaxM3SparseBackend.get_kv_cache_stride_order
1136 )
@@ -1192,10 +1197,8 @@ def test_padded_main_cache_is_flagged():
1197 @pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
1198 def test_reshape_and_cache_flash_write_persists(kv_layout: str):
1199 """AC-5 write path: the `reshape_and_cache_flash` write site now consumes
1195 - `self.kv_cache.unbind(1)` directly. Writing through those views must persist
1196 - into the bound storage (read back through an independent logical view) under
1197 - both layouts — a `.contiguous()` copy of the unbind slice would leave the
1198 - bound storage unchanged."""
1200 + packed-content K/V split views. Writing through those views must persist
1201 + into the bound storage under both layouts."""
1202 torch.manual_seed(0)
1203 num_pages = 4
1204 kv_cache = _allocate_main_kv_via_contract(num_pages)
@@ -1203,7 +1206,7 @@ def test_reshape_and_cache_flash_write_persists(kv_layout: str):
1206 kv_cache.zero_()
1207
1208 # Exactly the production write-site code under test.
1206 - key_cache, value_cache = kv_cache.unbind(1)
1209 + key_cache, value_cache = kv_cache.transpose(1, 2).split(HEAD_DIM, dim=-1)
1210
1211 num_tokens = 12
1212 slot_mapping = torch.randperm(num_pages * BLOCK_SIZE, device="cuda")[
@@ -1222,5 +1225,5 @@ def test_reshape_and_cache_flash_write_persists(kv_layout: str):
1225 for t in range(num_tokens):
1226 slot = int(slot_mapping[t].item())
1227 blk, intra = divmod(slot, BLOCK_SIZE)
1225 - torch.testing.assert_close(kv_cache[blk, 0, intra], key[t])
1226 - torch.testing.assert_close(kv_cache[blk, 1, intra], value[t])
1228 + torch.testing.assert_close(kv_cache[blk, :, intra, :HEAD_DIM], key[t])
1229 + torch.testing.assert_close(kv_cache[blk, :, intra, HEAD_DIM:], value[t])
tests/v1/attention/test_attention_backends.py
+9 -13
@@ -140,12 +140,10 @@ def create_and_prepopulate_kv_cache(
140 block_table = common_attn_metadata.block_table_tensor
141 slot_mapping = common_attn_metadata.slot_mapping
142
143 - # Create KV cache and populate in (2, num_blocks, ...) layout for easy
144 - # flat indexing, then transpose to (num_blocks, 2, ...) layout.
143 kv_cache = torch.zeros(
146 - 2, num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
144 + num_blocks, block_size, num_kv_heads, 2 * head_size, dtype=dtype, device=device
145 )
148 - kv_cache_flat = kv_cache.view(2, -1, num_kv_heads, head_size)
146 + kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size)
147
148 # Populate the cache with the context tokens
149 # Start from block_id=1 since block_id=0 is considered the null block
@@ -154,15 +152,12 @@ def create_and_prepopulate_kv_cache(
152 k_context, v_context = k_contexts[i], v_contexts[i]
153 start = start_block_idx * block_size
154 end = start + k_context.shape[0]
157 - kv_cache_flat[0, start:end, ...] = k_context
158 - kv_cache_flat[1, start:end, ...] = v_context
155 + kv_cache_flat[start:end, :, :head_size] = k_context
156 + kv_cache_flat[start:end, :, head_size:] = v_context
157
158 # Stay block aligned and allocate enough blocks for the new tokens
159 start_block_idx += cdiv(int(seq_lens[i]), block_size)
160
163 - # Transpose to (num_blocks, 2, ...) layout
164 - kv_cache = kv_cache.transpose(0, 1).contiguous()
165 -
161 blocks_end = start_block_idx
162
163 # Permute the context blocks (excluding block 0 which is null)
@@ -199,7 +194,8 @@ def create_and_prepopulate_kv_cache(
194 i, block_indices
195 ] * block_size + token_inter_block_offsets.to(device)
196
202 - return kv_cache
197 + # Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs)
198 + return kv_cache.transpose(1, 2).contiguous()
199
200
201 class MockAttentionLayer:
@@ -496,9 +492,6 @@ def _test_backend_correctness(
492 set_kv_cache_layout("HND")
493 reset_kv_cache_layout = True
494
499 - # Apply stride order like runtime does in
500 - # _reshape_kv_cache (attn_utils.py:182-210): permute to physical
501 - # layout, make contiguous, then permute to logical layout.
495 kv_cache_for_backend = kv_cache
496 if backend_cls is not None:
497 try:
@@ -506,6 +499,9 @@ def _test_backend_correctness(
499 except (AttributeError, NotImplementedError):
500 stride_order = tuple(range(kv_cache.ndim))
501 if stride_order != tuple(range(kv_cache.ndim)):
502 + # Apply stride order like runtime does in
503 + # _reshape_kv_cache (attn_utils.py:182-210): permute to physical
504 + # layout, make contiguous, then permute to logical layout.
505 inv_order = [stride_order.index(i) for i in range(len(stride_order))]
506 kv_cache_for_backend = (
507 kv_cache.permute(*stride_order).contiguous().permute(*inv_order)
tests/v1/attention/test_trtllm_attention_integration.py
+43 -39
@@ -97,13 +97,13 @@ def _create_hnd_kv_cache(
97 device,
98 num_blocks,
99 common_attn_metadata,
100 + kv_in_head_dim=False,
101 ):
101 - """Create and populate a KV cache with HND-compatible strides.
102 + """Create and populate a packed KV cache with HND-compatible strides.
103
103 - The returned tensor has logical shape
104 - (num_blocks, 2, block_size, num_kv_heads, head_size) but is physically
105 - laid out as (num_blocks, 2, num_kv_heads, block_size, head_size) so that
106 - ``kv_cache.permute(0, 1, 3, 2, 4)`` yields a contiguous HND view.
104 + When kv_in_head_dim=False (default), returns (B, H, N, 2*hs) with K/V
105 + packed in the content dim. When kv_in_head_dim=True, returns
106 + (B, 2*H, N, hs) with K/V as separate head groups.
107 """
108 seq_lens = common_attn_metadata.seq_lens.cpu()
109 query_lens = (
@@ -114,26 +114,34 @@ def _create_hnd_kv_cache(
114 slot_mapping = common_attn_metadata.slot_mapping
115 batch_size = len(k_contexts)
116
117 - # Build cache in (2, num_blocks, block_size, num_kv_heads, head_size)
118 - # then convert to HND format (same approach as test_attention_backends.py).
119 - kv_cache_raw = torch.zeros(
120 - 2,
117 + # kv_in_head_dim: (B, N, 2*H, hs) — K/V as separate head groups
118 + # else: (B, N, H, 2*hs) — K/V packed in content dim
119 + n_heads, content = (
120 + (2 * num_kv_heads, head_size)
121 + if kv_in_head_dim
122 + else (num_kv_heads, 2 * head_size)
123 + )
124 + kv_cache = torch.zeros(
125 num_blocks,
126 block_size,
123 - num_kv_heads,
124 - head_size,
127 + n_heads,
128 + content,
129 dtype=dtype,
130 device=device,
131 )
128 - kv_cache_flat = kv_cache_raw.view(2, -1, num_kv_heads, head_size)
132 + kv_cache_flat = kv_cache.view(-1, n_heads, content)
133
134 start_block_idx = 1
135 for i in range(batch_size):
136 k_ctx, v_ctx = k_contexts[i], v_contexts[i]
137 start = start_block_idx * block_size
138 end = start + k_ctx.shape[0]
135 - kv_cache_flat[0, start:end] = k_ctx
136 - kv_cache_flat[1, start:end] = v_ctx
139 + if kv_in_head_dim:
140 + kv_cache_flat[start:end, :num_kv_heads] = k_ctx
141 + kv_cache_flat[start:end, num_kv_heads:] = v_ctx
142 + else:
143 + kv_cache_flat[start:end, :, :head_size] = k_ctx
144 + kv_cache_flat[start:end, :, head_size:] = v_ctx
145 start_block_idx += cdiv(int(seq_lens[i]), block_size)
146
147 blocks_end = start_block_idx
@@ -142,7 +150,7 @@ def _create_hnd_kv_cache(
150 perm = torch.randperm(blocks_end - 1) + 1
151 inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
152 inv_perm[1:] = torch.argsort(perm) + 1
145 - kv_cache_raw[:, 1:blocks_end] = kv_cache_raw[:, perm]
153 + kv_cache[1:blocks_end] = kv_cache[perm]
154
155 # Build block table.
156 start_block_idx = 1
@@ -165,10 +173,8 @@ def _create_hnd_kv_cache(
173 i, block_indices
174 ] * block_size + intra_block_offsets.to(device)
175
168 - # Transpose to FlashInfer logical shape then make HND-strided.
169 - kv_cache = kv_cache_raw.transpose(0, 1)
170 - kv_cache = kv_cache.transpose(2, 3).contiguous().transpose(2, 3)
171 - return kv_cache
176 + # Transpose to canonical: (B, H, N, 2*hs) or (B, 2*H, N, hs)
177 + return kv_cache.transpose(1, 2).contiguous()
178
179
180 def _create_nvfp4_hnd_kv_cache(
@@ -187,20 +193,15 @@ def _create_nvfp4_hnd_kv_cache(
193 reshape_and_cache_flash, using the same block-table layout as
194 _create_hnd_kv_cache.
195
190 - The returned tensor is dtype ``uint8`` with shape
191 - ``(num_blocks, 2, block_size, num_kv_heads, full_dim)`` in logical
192 - (NHD) order, but physically permuted to HND layout via stride order
193 - ``(0, 1, 3, 2, 4)`` (i.e. ``num_kv_heads`` before ``block_size``).
194 -
195 - The last dimension ``full_dim = head_size // 2 + head_size // 16``
196 - packs two regions contiguously:
196 + The returned tensor is dtype ``uint8`` with head-group layout
197 + ``(num_blocks, 2 * num_kv_heads, block_size, full_dim)``
198 + where K heads occupy the first ``num_kv_heads`` heads and V heads the second.
199 + Each ``full_dim = head_size // 2 + head_size // 16`` block packs two regions:
200 - **FP4 data** (``head_size // 2`` bytes): pairs of E2M1 values,
201 two per byte.
202 - **FP8 block scales** (``head_size // 16`` bytes): one E4M3
203 scale per 16-element block.
204
202 - Dimension 1 indexes K (``[:, 0]``) and V (``[:, 1]``).
203 -
205 Args:
206 k_contexts: List of key context tensors, one per sequence.
207 v_contexts: List of value context tensors, one per sequence.
@@ -219,6 +220,7 @@ def _create_nvfp4_hnd_kv_cache(
220 ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, HND-strided).
221 """
222 # First create a bf16 HND cache so block tables are populated.
223 + # Use kv_in_head_dim=True so K/V are separate head groups (B, 2*H, N, hs).
224 bf16_cache = _create_hnd_kv_cache(
225 k_contexts,
226 v_contexts,
@@ -229,20 +231,20 @@ def _create_nvfp4_hnd_kv_cache(
231 device,
232 num_blocks,
233 common_attn_metadata,
234 + kv_in_head_dim=True,
235 )
236
234 - # Allocate nvfp4 cache: same shape but with full_dim (data + scale).
237 + # (num_blocks, 2 * num_kv_heads, block_size, full_dim) — K heads first, then V heads
238 full_dim = nvfp4_kv_cache_full_dim(head_size)
236 - hnd_order = (0, 1, 3, 2, 4)
239 nvfp4_cache = torch.zeros(
238 - (num_blocks, 2, num_kv_heads, block_size, full_dim),
240 + (num_blocks, 2 * num_kv_heads, block_size, full_dim),
241 dtype=torch.uint8,
242 device=device,
241 - ).permute(*hnd_order)
243 + )
244 + k_cache, v_cache = nvfp4_cache.split(num_kv_heads, dim=1)
245
246 # Flatten bf16 context into tokens and quantize via reshape_and_cache_flash.
244 - # bf16_cache is (num_blocks, 2, block_size, num_kv_heads, head_size) logical
245 - # with HND physical strides.
247 + # bf16_cache is (B, 2*H, N, hs); split K/V on head dim.
248 block_table = common_attn_metadata.block_table_tensor
249 seq_lens = common_attn_metadata.seq_lens.cpu()
250 query_lens = (
@@ -258,19 +260,21 @@ def _create_nvfp4_hnd_kv_cache(
260 # Gather context tokens from the bf16 cache using block table.
261 n_ctx_blocks = (ctx_len + block_size - 1) // block_size
262 blocks = block_table[i, :n_ctx_blocks]
261 - # bf16_cache[:, kv_idx] is (num_blocks, block_size, num_kv_heads, head_size)
262 - k_ctx = bf16_cache[blocks, 0].reshape(-1, num_kv_heads, head_size)[:ctx_len]
263 - v_ctx = bf16_cache[blocks, 1].reshape(-1, num_kv_heads, head_size)[:ctx_len]
263 + # bf16_cache is (B, 2*H, N, hs); split K and V head groups.
264 + k_bf16, v_bf16 = bf16_cache[blocks].split(num_kv_heads, dim=1)
265 + k_ctx = k_bf16.transpose(1, 2).reshape(-1, num_kv_heads, head_size)[:ctx_len]
266 + v_ctx = v_bf16.transpose(1, 2).reshape(-1, num_kv_heads, head_size)[:ctx_len]
267 # Build slot mapping for these context tokens.
268 token_offsets = torch.arange(ctx_len, device=device)
269 block_indices = token_offsets // block_size
270 intra_offsets = token_offsets % block_size
271 slots = block_table[i, block_indices] * block_size + intra_offsets
272 + # reshape_and_cache_flash expects (B, N, H, D) cache views.
273 torch.ops._C_cache_ops.reshape_and_cache_flash(
274 k_ctx,
275 v_ctx,
272 - nvfp4_cache[:, 0],
273 - nvfp4_cache[:, 1],
276 + k_cache.transpose(1, 2),
277 + v_cache.transpose(1, 2),
278 slots,
279 "nvfp4",
280 kv_scale_t,
tests/v1/kv_connector/unit/test_mooncake_connector.py
+37 -67
@@ -289,10 +289,9 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
289 prefill_worker = prefill_connector.connector_worker
290
291 block_len = 4096
292 - kv_half = block_len // 2
292 prefill_worker.kv_caches_base_addr = [0x1000]
293 prefill_worker.block_len_per_layer = [block_len]
295 - prefill_worker.kv_block_len_per_layer = [kv_half]
294 + prefill_worker.kv_block_len_per_layer = [block_len]
295 prefill_worker.registered_layer_names = ["model.layers.1.self_attn"]
296 prefill_worker.registered_layer_indices = [1]
297
@@ -321,7 +320,7 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
320 req_blocks={"d-req-layer-align": (transfer_id, [[20]])},
321 kv_caches_base_addr=[0xA000, 0xB000],
322 block_lens=[block_len, block_len],
324 - kv_block_lens=[kv_half, kv_half],
323 + kv_block_lens=[block_len, block_len],
324 registered_layer_names=[
325 "model.layers.0.self_attn",
326 "model.layers.1.self_attn",
@@ -338,15 +337,9 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
337 await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
338
339 src_ptrs, dst_ptrs, lengths = mock_send_blocks.call_args[0][1:]
341 - assert src_ptrs == [
342 - 0x1000 + 10 * block_len,
343 - 0x1000 + 10 * block_len + kv_half,
344 - ]
345 - assert dst_ptrs == [
346 - 0xB000 + 20 * block_len,
347 - 0xB000 + 20 * block_len + kv_half,
348 - ]
349 - assert lengths == [kv_half, kv_half]
340 + assert src_ptrs == [0x1000 + 10 * block_len]
341 + assert dst_ptrs == [0xB000 + 20 * block_len]
342 + assert lengths == [block_len]
343
344 sent_identity, sent_payload = mock_socket.send_multipart.call_args[0][0]
345 assert sent_identity == identity
@@ -832,9 +825,8 @@ async def test_kv_producer(monkeypatch):
825 prefill_worker = prefill_connector.connector_worker
826 prefill_worker.kv_caches_base_addr = [0x1000]
827 block_len = 4096
835 - kv_half = block_len // 2
828 prefill_worker.block_len_per_layer = [block_len]
837 - prefill_worker.kv_block_len_per_layer = [kv_half]
829 + prefill_worker.kv_block_len_per_layer = [block_len]
830 prefill_worker.registered_layer_names = ["model.layers.0.self_attn"]
831 prefill_worker.registered_layer_indices = [0]
832
@@ -862,7 +854,7 @@ async def test_kv_producer(monkeypatch):
854 req_blocks={"d-req-1": (transfer_id, [[20, 21]])},
855 kv_caches_base_addr=[0x2000],
856 block_lens=[block_len],
865 - kv_block_lens=[kv_half],
857 + kv_block_lens=[block_len],
858 registered_layer_names=["model.layers.0.self_attn"],
859 registered_layer_indices=[0],
860 )
@@ -874,24 +866,18 @@ async def test_kv_producer(monkeypatch):
866 with patch.object(
867 prefill_worker, "_send_blocks", return_value=0
868 ) as mock_send_blocks:
877 - # With blocks-first layout, each block is virtually split
878 - # into K and V halves, producing non-coalesced transfers.
879 - def expected_split_transfers(src_base, dst_base, src_blocks, dst_blocks):
880 - """Build expected (src_ptrs, dst_ptrs, lengths) for
881 - virtual-split K/V transfers."""
882 - src_ptrs, dst_ptrs, lengths = [], [], []
883 - for kv_offset in (0, kv_half):
884 - for sb, db in zip(src_blocks, dst_blocks):
885 - src_ptrs.append(src_base + sb * block_len + kv_offset)
886 - dst_ptrs.append(dst_base + db * block_len + kv_offset)
887 - lengths.append(kv_half)
888 - return src_ptrs, dst_ptrs, lengths
869 +
870 + def expected_transfers(src_base, dst_base, src_blocks, dst_blocks):
871 + n = len(src_blocks)
872 + return (
873 + [src_base + src_blocks[0] * block_len],
874 + [dst_base + dst_blocks[0] * block_len],
875 + [n * block_len],
876 + )
877
878 # Normal case: 2 blocks to 2 blocks
879 await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
892 - src, dst, lens = expected_split_transfers(
893 - 0x1000, 0x2000, [10, 11], [20, 21]
894 - )
880 + src, dst, lens = expected_transfers(0x1000, 0x2000, [10, 11], [20, 21])
881 mock_send_blocks.assert_called_once_with(
882 "consumer-host:54321",
883 src,
@@ -923,7 +909,7 @@ async def test_kv_producer(monkeypatch):
909 # Worker processes the consumer's request
910 await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
911 # Verify transfer parameters are correct: 11 to 20
926 - src, dst, lens = expected_split_transfers(0x1000, 0x2000, [11], [20])
912 + src, dst, lens = expected_transfers(0x1000, 0x2000, [11], [20])
913 mock_send_blocks.assert_called_once_with(
914 "consumer-host:54321",
915 src,
@@ -1266,7 +1252,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
1252
1253 prefill_worker.kv_caches_base_addr = [0x1000]
1254 prefill_worker.block_len_per_layer = [local_block_len]
1269 - prefill_worker.kv_block_len_per_layer = [local_block_len // 2]
1255 + prefill_worker.kv_block_len_per_layer = [local_block_len]
1256 prefill_worker.registered_layer_names = ["model.layers.0.self_attn"]
1257 prefill_worker.registered_layer_indices = [0]
1258
@@ -1314,7 +1300,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
1300 },
1301 kv_caches_base_addr=[0x2000],
1302 block_lens=[remote_block_len],
1317 - kv_block_lens=[remote_block_len // 2],
1303 + kv_block_lens=[remote_block_len],
1304 registered_layer_names=["model.layers.0.self_attn"],
1305 registered_layer_indices=[0],
1306 )
@@ -1336,47 +1322,31 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
1322 flat_remote = [b for g in remote_block_ids for b in g]
1323 num_blocks = len(flat_local)
1324
1339 - # With blocks-first layout, virtual split halves block
1340 - # lengths and doubles transfer regions (K + V).
1341 - local_kv_block_len = local_block_len // 2
1342 - remote_kv_block_len = remote_block_len // 2
1343 -
1344 - assert len(src_ptrs) == 2 * num_blocks
1345 - assert len(dst_ptrs) == 2 * num_blocks
1346 - assert len(lengths) == 2 * num_blocks
1325 + assert len(src_ptrs) == num_blocks
1326 + assert len(dst_ptrs) == num_blocks
1327 + assert len(lengths) == num_blocks
1328
1348 - # Compute expected offsets using kv_block_len
1329 if d_tp_size <= P_TP_SIZE:
1330 tp_ratio = P_TP_SIZE // d_tp_size
1331 expected_src_off = 0
1352 - expected_dst_off = (P_TP_RANK % tp_ratio) * local_kv_block_len
1353 - expected_xfer_len = local_kv_block_len
1332 + expected_dst_off = (P_TP_RANK % tp_ratio) * local_block_len
1333 + expected_xfer_len = local_block_len
1334 else:
1335 ratio_abs = d_tp_size // P_TP_SIZE
1356 - expected_src_off = (d_rank % ratio_abs) * remote_kv_block_len
1336 + expected_src_off = (d_rank % ratio_abs) * remote_block_len
1337 expected_dst_off = 0
1358 - expected_xfer_len = remote_kv_block_len
1359 -
1360 - # First num_blocks entries are K region,
1361 - # next num_blocks are V region.
1362 - for region_idx in range(2):
1363 - local_region_base = 0x1000 + region_idx * local_kv_block_len
1364 - remote_region_base = 0x2000 + region_idx * remote_kv_block_len
1365 - for blk_idx, (lblk, rblk) in enumerate(
1366 - zip(flat_local, flat_remote)
1367 - ):
1368 - idx = region_idx * num_blocks + blk_idx
1369 - assert src_ptrs[idx] == (
1370 - local_region_base
1371 - + lblk * local_block_len
1372 - + expected_src_off
1373 - )
1374 - assert dst_ptrs[idx] == (
1375 - remote_region_base
1376 - + rblk * remote_block_len
1377 - + expected_dst_off
1378 - )
1379 - assert lengths[idx] == expected_xfer_len
1338 + expected_xfer_len = remote_block_len
1339 +
1340 + local_region_base = 0x1000
1341 + remote_region_base = 0x2000
1342 + for blk_idx, (lblk, rblk) in enumerate(zip(flat_local, flat_remote)):
1343 + assert src_ptrs[blk_idx] == (
1344 + local_region_base + lblk * local_block_len + expected_src_off
1345 + )
1346 + assert dst_ptrs[blk_idx] == (
1347 + remote_region_base + rblk * remote_block_len + expected_dst_off
1348 + )
1349 + assert lengths[blk_idx] == expected_xfer_len
1350
1351 # Verify successful response sent back to consumer
1352 mock_socket.send_multipart.assert_called_once()
tests/v1/kv_connector/unit/test_nixl_connector.py
+64 -15
@@ -1068,6 +1068,65 @@ class TestNixlHandshake:
1068 # whole block is moved.
1069 worker.add_remote_agent(meta, remote_tp_size=1)
1070
1071 + @patch(
1072 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1073 + FakeNixlWrapper,
1074 + )
1075 + def test_hybrid_mamba_attention_remote_descs_use_packed_head_slices(
1076 + self, default_vllm_config, dist_init
1077 + ):
1078 + worker = FakeNixlConnectorWorker(
1079 + create_vllm_config(), "engine", hand_shake_latency=0
1080 + )
1081 +
1082 + remote_block_len = 2048
1083 + local_block_len = remote_block_len // 2
1084 + worker.block_len_per_layer = [local_block_len]
1085 + worker._region_is_mla = [False]
1086 + worker.num_blocks = 1
1087 + worker.num_regions = 1
1088 + worker._has_mamba = True
1089 + worker._mamba_ssm_size = (128, 256)
1090 + worker.transfer_topo = TransferTopology(
1091 + tp_rank=1,
1092 + tp_size=2,
1093 + block_size=worker.block_size,
1094 + engine_id=worker.engine_id,
1095 + is_mla=False,
1096 + is_mamba=True,
1097 + total_num_kv_heads=2,
1098 + attn_backends=worker.attn_backends,
1099 + tensor_shape=None,
1100 + )
1101 + assert worker.transfer_topo.virtually_split_kv_in_blocks
1102 +
1103 + plan = MagicMock(
1104 + source_ranks_per_group=((0,), (0,)),
1105 + rank_offset_factor=1,
1106 + )
1107 + meta = MagicMock(
1108 + kv_caches_base_addr=[0x1000],
1109 + device_id=0,
1110 + num_blocks=1,
1111 + block_lens=[remote_block_len],
1112 + )
1113 +
1114 + assert worker.get_backend_aware_kv_block_len(0, mamba_view=False) == (
1115 + local_block_len
1116 + )
1117 + assert (
1118 + worker.get_backend_aware_kv_block_len(0, first_split=True, mamba_view=True)
1119 + == worker._mamba_ssm_size[0]
1120 + )
1121 + assert (
1122 + worker.get_backend_aware_kv_block_len(0, first_split=False, mamba_view=True)
1123 + == worker._mamba_ssm_size[1]
1124 + )
1125 +
1126 + assert worker._build_fa_remote(plan, meta, block_size_ratio=1) == [
1127 + (0x1000 + local_block_len, local_block_len, 0)
1128 + ]
1129 +
1130 @patch(
1131 "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1132 FakeNixlWrapper,
@@ -1691,13 +1750,6 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
1750 reason="Attention backend FLASH_ATTN is not supported on ROCm",
1751 ),
1752 ),
1694 - pytest.param(
1695 - "ROCM_ATTN",
1696 - marks=pytest.mark.skipif(
1697 - not current_platform.is_rocm(),
1698 - reason="Attention backend ROCM_ATTN is only supported on ROCm",
1699 - ),
1700 - ),
1753 "TRITON_ATTN",
1754 ],
1755 )
@@ -1824,8 +1876,7 @@ def test_register_kv_caches(
1876 test_shape = backend_cls.get_kv_cache_shape(
1877 num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
1878 )
1827 - is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
1828 - virtually_split = is_blocks_first and not connector.prefer_cross_layer_blocks
1879 + is_blocks_first = len(test_shape) == 4 and test_shape[0] == 1
1880
1881 if connector.prefer_cross_layer_blocks:
1882 with set_current_vllm_config(vllm_config):
@@ -1856,7 +1907,7 @@ def test_register_kv_caches(
1907 ]
1908 expected_num_entries = 1
1909
1859 - expected_blocks_count = num_blocks * (2 if virtually_split else 1)
1910 + expected_blocks_count = num_blocks
1911
1912 kv_caches = {"all-layers": cross_layers_kv_cache}
1913 else:
@@ -1885,6 +1936,7 @@ def test_register_kv_caches(
1936 unique_tensor.data_ptr(),
1937 ]
1938 expected_num_entries = 2
1939 + expected_blocks_count = kv_cache_config.num_blocks * 2
1940 else:
1941 expected_tensor_size = (
1942 shared_tensor[0].element_size() * shared_tensor[0].numel()
@@ -1896,7 +1948,7 @@ def test_register_kv_caches(
1948 unique_tensor[1].data_ptr(),
1949 ]
1950 expected_num_entries = 4
1899 - expected_blocks_count = kv_cache_config.num_blocks * 4
1951 + expected_blocks_count = kv_cache_config.num_blocks * 4
1952
1953 # Execute register_kv_caches
1954 connector.register_kv_caches(kv_caches)
@@ -1930,10 +1982,7 @@ def test_register_kv_caches(
1982 else:
1983 num_blocks = kv_cache_config.num_blocks
1984
1933 - if virtually_split:
1934 - expected_block_len = expected_tensor_size // num_blocks // 2
1935 - else:
1936 - expected_block_len = expected_tensor_size // num_blocks
1985 + expected_block_len = expected_tensor_size // num_blocks
1986
1987 for i, block_entry in enumerate(blocks_data):
1988 block_start_addr, block_len, tp_rank = block_entry
tests/v1/kv_connector/unit/test_transfer_topology_sharded.py
+2 -2
@@ -18,8 +18,8 @@ class _FakeAttentionBackend:
18 block_size: int,
19 num_kv_heads: int,
20 head_size: int,
21 - ) -> tuple[int, int, int, int, int]:
22 - return (2, num_blocks, num_kv_heads, block_size, head_size)
21 + ) -> tuple[int, int, int, int]:
22 + return (num_blocks, num_kv_heads, block_size, 2 * head_size)
23
24
25 def _make_topology(
tests/v1/worker/test_gpu_model_runner.py
+7 -5
@@ -795,9 +795,10 @@ def test_kv_cache_stride_order(monkeypatch, model_runner):
795 )
796
797 # TODO mla test
798 - default_stride = tuple(range(5))
798 + default_stride = tuple(range(len(expected_kv_cache_shape)))
799 + non_default_stride = (*default_stride[1:], default_stride[0])
800 # Permutation that gets you back to expected kv shape
800 - for test_stride in ((1, 4, 0, 2, 3), (0, 1, 2, 3, 4)):
801 + for test_stride in (non_default_stride, default_stride):
802
803 def rnd_stride_order(
804 include_num_layers_dimension: bool = False, test_stride=test_stride
@@ -1300,9 +1301,10 @@ def test_hybrid_attention_mamba_tensor_shapes():
1301 actual_kv = vllm_ctx[layer].kv_cache[kernel_block, :]
1302 expected = attn_blocks_constant[i]
1303
1303 - # Check K and V separately
1304 - assert torch.equal(actual_kv[0], expected)
1305 - assert torch.equal(actual_kv[1], expected)
1304 + # Packed layout: (num_kv_heads, block_size, 2*head_size). Every
1305 + # head in the block was filled with the same constant.
1306 + for head_idx in range(actual_kv.shape[0]):
1307 + assert torch.equal(actual_kv[head_idx], expected)
1308
1309 for layer in [layer_2, layer_3, layer_4, layer_5]:
1310 for i, kv_block in enumerate(kv_blocks_for_mamba):
vllm/distributed/kv_transfer/kv_connector/utils.py
+35 -34
@@ -10,7 +10,12 @@ from typing import TYPE_CHECKING, Any, Literal, cast
10
11 import torch
12
13 -from vllm.config import VllmConfig, get_current_vllm_config, get_layers_from_vllm_config
13 +from vllm.config import (
14 + VllmConfig,
15 + get_current_vllm_config,
16 + get_layers_from_vllm_config,
17 + set_current_vllm_config,
18 +)
19 from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
20 from vllm.logger import init_logger
21 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
@@ -344,14 +349,15 @@ def get_current_attn_backends(
349 )
350 from vllm.v1.attention.selector import get_attn_backend
351
347 - return [
348 - get_attn_backend(
349 - head_size=vllm_config.model_config.get_head_size(),
350 - dtype=vllm_config.model_config.dtype,
351 - kv_cache_dtype=vllm_config.cache_config.cache_dtype,
352 - use_mla=vllm_config.model_config.use_mla,
353 - )
354 - ]
352 + with set_current_vllm_config(vllm_config):
353 + return [
354 + get_attn_backend(
355 + head_size=vllm_config.model_config.get_head_size(),
356 + dtype=vllm_config.model_config.dtype,
357 + kv_cache_dtype=vllm_config.cache_config.cache_dtype,
358 + use_mla=vllm_config.model_config.use_mla,
359 + )
360 + ]
361
362
363 def get_current_attn_backend(
@@ -426,12 +432,16 @@ class TransferTopology:
432 head_size=1,
433 )
434 logger.debug("Test kv_cache_shape: %s", kv_cache_shape)
429 - # Non-MLA backends caches have 5 dims [num_blocks, 2, H,N,D],
430 - # we just mock num_blocks to 1 for the dimension check below.
431 - # Hybrid SSM models assume a single blocks_first layout
432 - self._is_kv_layout_blocks_first = self.is_mamba or (
433 - len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1
434 - )
435 + assert kv_cache_shape[0] == 1, (
436 + "KV cache layout must be blocks-first; expected mocked "
437 + f"num_blocks=1 in leading dim, got shape {kv_cache_shape}."
438 + )
439 + if not self.is_mla:
440 + assert len(kv_cache_shape) == 4, (
441 + "Attention KV cache layout must be standardized as "
442 + "[num_blocks, num_kv_heads, block_size, content_size], "
443 + f"got shape {kv_cache_shape}."
444 + )
445
446 self._cross_layers_blocks = False
447 if self.tensor_shape is not None:
@@ -490,29 +500,19 @@ class TransferTopology:
500 # Layout properties
501 # ============================================================
502
493 - @property
494 - def is_kv_layout_blocks_first(self) -> bool:
495 - return self._is_kv_layout_blocks_first
496 -
503 @property
504 def cross_layers_blocks(self) -> bool:
505 return self._cross_layers_blocks
506
507 @property
508 def virtually_split_kv_in_blocks(self) -> bool:
503 - # Whether to logically split each block into K and V halves.
504 - # Applies when K/V are interleaved within each block (blocks-first),
505 - # but NOT when cross-layer blocks are used — cross-layer blocks have
506 - # per-layer K/V interleaving (L0_K, L0_V, L1_K, L1_V, ...) so a
507 - # simple half-split does not separate K from V.
508 - return self._is_kv_layout_blocks_first and not self._cross_layers_blocks
509 -
510 - @property
511 - def split_k_and_v(self) -> bool:
512 - # Whether to register regions for K and V separately (when present).
513 - return not (
514 - self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first
515 - )
509 + # Whether to logically split each block into two separately-indexable
510 + # sub-regions. With K and V packed into the content dim, an attention
511 + # block transfers as a single unit — no K/V sub-split is needed. Only
512 + # Mamba still needs this, to index its two state regions (conv/ssm)
513 + # separately. Not applicable to cross-layer blocks (per-layer
514 + # interleaving means a simple half-split does not separate the parts).
515 + return self.is_mamba and not self._cross_layers_blocks
516
517 # ============================================================
518 # Common methods
@@ -616,8 +616,9 @@ class TransferTopology:
616 # Swap [2<>num_blocks] dims for hybrid SSM layout.
617 cache = cache.transpose(0, 1)
618
619 - # Regular case: backends like FA register K/V in separate regions
620 - return cache if self.split_k_and_v else [cache]
619 + # K and V are packed into one tensor (content dim), so each layer
620 + # registers as a single region.
621 + return [cache]
622
623 def describe(self, remote_engine_id: EngineId, remote_pp_rank: int = 0) -> str:
624 """One-line summary of transfer config for logging."""
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py
+30
@@ -7,6 +7,7 @@ from typing import NamedTuple
7 import torch
8
9 from vllm.v1.kv_cache_interface import (
10 + AttentionSpec,
11 KVCacheConfig,
12 KVCacheSpec,
13 MLAAttentionSpec,
@@ -56,6 +57,11 @@ def is_mla_cache_layer(
57 return isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))
58
59
60 +def _content_packed_dim(spec: AttentionSpec) -> int:
61 + head_size_v = getattr(spec, "head_size_v", spec.head_size)
62 + return spec.head_size + head_size_v
63 +
64 +
65 def _spec_dim_matches(value: int, expected: int | None) -> bool:
66 return expected is None or value == expected
67
@@ -223,6 +229,30 @@ def get_layer_transfer_geometry(
229 split_kv_regions=False,
230 )
231
232 + if (
233 + not is_mla_cache
234 + and isinstance(spec, AttentionSpec)
235 + and len(shape) == 4
236 + and shape[1] == spec.num_kv_heads
237 + and shape[2] == spec.block_size
238 + and shape[3] == _content_packed_dim(spec)
239 + ):
240 + num_blocks, num_kv_heads, block_size, packed_dim = shape
241 + slot_size_bytes = num_kv_heads * packed_dim * element_size
242 + block_len = block_size * slot_size_bytes
243 + return LayerTransferGeometry(
244 + num_blocks=num_blocks,
245 + block_size=block_size,
246 + block_len=block_len,
247 + slot_size_bytes=slot_size_bytes,
248 + block_stride=stride[0],
249 + local_kv_stride=None,
250 + remote_kv_stride=None,
251 + transfers_per_block=1,
252 + regions_per_block=1,
253 + split_kv_regions=False,
254 + )
255 +
256 cache_kind = "MLA" if is_mla_cache else "K/V"
257 raise ValueError(
258 f"Unsupported MoRIIO {cache_kind} cache shape for layer "
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+33 -91
@@ -216,12 +216,10 @@ class NixlBaseConnectorWorker:
216 if n_regions == 0 or self.num_regions == 0:
217 return [False] * num_fa_descs
218 nblk = num_fa_descs // self.num_regions
219 - virtually_split = self.transfer_topo.virtually_split_kv_in_blocks
219 flags: list[bool] = []
220 for i in range(n_regions):
221 replicated = self._is_region_replicated(i)
223 - num_streams = 1 if replicated or not virtually_split else 2
224 - flags.extend([replicated] * (num_streams * nblk))
222 + flags.extend([replicated] * nblk)
223 assert len(flags) == num_fa_descs, (
224 f"FA desc flags {len(flags)} != num_fa_descs {num_fa_descs}"
225 )
@@ -712,31 +710,31 @@ class NixlBaseConnectorWorker:
710 NOT directly supported by NIXL (e.g., tpu)
711 """
712 xfer_buffers: dict[str, torch.Tensor] = {}
715 - inv_order = [0, 1, 3, 2, 4]
713 try:
714 for layer_name, kv_cache in kv_caches.items():
715 kv_shape = kv_cache.shape
716 kv_dtype = kv_cache.dtype
717 permute_shape = False
721 - if (
722 - self.kv_cache_layout == "NHD"
723 - and self.vllm_config.kv_transfer_config is not None
724 - and self.vllm_config.kv_transfer_config.enable_permute_local_kv
725 - ):
726 - logger.info_once(
727 - "'enable_permute_local_kv' flag is enabled while "
728 - "device KV Layout is NHD. Init host buffer with"
729 - " HND to better support Decode/Prefill TP_ratio > 1."
730 - )
731 - # Since NHD will not support Decode/Prefill TP_ratio > 1,
732 - # we can leverage host_buffer for permute
733 - self.host_buffer_kv_cache_layout = "HND"
734 - kv_shape = (
735 - tuple(kv_shape[i] for i in inv_order)
736 - if not self.use_mla
737 - else kv_shape
738 - )
739 - permute_shape = not self.use_mla
718 + inv_order = (0, 2, 1, 3)
719 + if not self.use_mla:
720 + assert kv_cache.ndim == 4
721 +
722 + if self.kv_cache_layout == "NHD":
723 + if self.kv_transfer_config.enable_permute_local_kv:
724 + logger.info_once(
725 + "'enable_permute_local_kv' flag is enabled while "
726 + "device KV Layout is NHD. Init host buffer with"
727 + " HND to better support Decode/Prefill TP_ratio > 1."
728 + )
729 + # Since NHD will not support Decode/Prefill TP_ratio > 1,
730 + # we can leverage host_buffer for permute.
731 + self.host_buffer_kv_cache_layout = "HND"
732 + else:
733 + # Packed KV layout is logical (B, H, N, 2*D). Allocate
734 + # (B, N, H, 2*D) and view it as logical (B, H, N, 2*D)
735 + # so raw NIXL transfers see NHD physical strides.
736 + kv_shape = tuple(kv_shape[i] for i in inv_order)
737 + permute_shape = True
738
739 xfer_buffers[layer_name] = torch.empty(
740 kv_shape, dtype=kv_dtype, device="cpu"
@@ -1151,9 +1149,7 @@ class NixlBaseConnectorWorker:
1149 f"backend={self.backend_name}, "
1150 "all_backends="
1151 f"{[backend.get_name() for backend in self.attn_backends]}, "
1154 - f"kv_cache_layout={self.kv_cache_layout}, "
1155 - "blocks_first="
1156 - f"{self.transfer_topo.is_kv_layout_blocks_first}"
1152 + f"kv_cache_layout={self.kv_cache_layout}"
1153 )
1154
1155 # Need to make sure the device ID is non-negative for NIXL,
@@ -1184,22 +1180,6 @@ class NixlBaseConnectorWorker:
1180 regions_per_layer = self.num_regions // num_local_layers
1181 self._remote_region_offset = regions_per_layer * start_layer
1182
1187 - if self.transfer_topo.virtually_split_kv_in_blocks:
1188 - # NOTE (NickLucche) When FlashInfer is used, memory is registered
1189 - # with joint KV for each block. This minimizes the overhead in
1190 - # registerMem allowing faster descs queries. In order to be able to
1191 - # split on kv_heads dim as required by heterogeneous TP, one must
1192 - # be able to index K/V separately. Hence we double the number
1193 - # of 'virtual' regions here and halve `block_len` below.
1194 - # Similarly for Mamba layers, we register SSM+Conv as a single region and
1195 - # then duplicate it logically to be able to index SSM/Conv separately.
1196 - # Exception: key-only REPLICATE regions (MLA) have no V half, so
1197 - # they contribute a single desc stream and are not doubled.
1198 - self.num_regions = sum(
1199 - 1 if self._is_region_replicated(i) else 2
1200 - for i in range(len(self._region_is_mla))
1201 - )
1202 -
1183 # Total local FA descriptors (boundary between FA and mamba descs).
1184 self.num_descs = self.num_regions * self.num_blocks
1185
@@ -1363,22 +1343,6 @@ class NixlBaseConnectorWorker:
1343 block_offset = block_id * page_stride
1344 addr = base_addr + block_offset
1345 result.append((addr, kv_block_len, self.device_id))
1366 -
1367 - if (
1368 - self.transfer_topo.virtually_split_kv_in_blocks
1369 - and not self._is_region_replicated(i)
1370 - ):
1371 - # Separate and interleave K/V regions to maintain the same
1372 - # descs ordering. This is needed for selecting contiguous heads
1373 - # when split across TP ranks. (Skipped for key-only REPLICATE.)
1374 - second_split = self.get_backend_aware_kv_block_len(
1375 - layer_idx=i, first_split=False, mamba_view=False
1376 - )
1377 - for block_id in range(num_blocks):
1378 - block_offset = block_id * page_stride
1379 - addr = base_addr + block_offset
1380 - v_addr = addr + kv_block_len
1381 - result.append((v_addr, second_split, self.device_id))
1346 return result
1347
1348 def _build_fa_remote(
@@ -1423,20 +1387,6 @@ class NixlBaseConnectorWorker:
1387 # tp rank of size local_block_len.
1388 addr = base_addr + block_offset + rank_offset
1389 result.append((addr, local_block_len, nixl_agent_meta.device_id))
1426 -
1427 - emits_v = self.transfer_topo.virtually_split_kv_in_blocks and not replicated
1428 - if emits_v:
1429 - # With FlashInfer index V separately to allow head splitting.
1430 - second_split = self.get_backend_aware_kv_block_len(
1431 - layer_idx=i, first_split=False, mamba_view=False
1432 - )
1433 - second_split = second_split // num_reads
1434 - for block_id in range(num_blocks):
1435 - block_offset = block_id * page_size
1436 - addr = base_addr + block_offset + rank_offset
1437 - # Hop over the first split of remote page, K, to read V.
1438 - v_addr = addr + nixl_agent_meta.block_lens[i] // 2
1439 - result.append((v_addr, second_split, nixl_agent_meta.device_id))
1390 return result
1391
1392 def register_local_xfer_handler(
@@ -1927,24 +1877,18 @@ class NixlBaseConnectorWorker:
1877 block_size_ratio,
1878 )
1879
1930 - split_k_and_v = self.transfer_topo.split_k_and_v
1931 -
1880 for block_ids in block_ids_list:
1881 indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
1882
1935 - for _, cache_or_caches in self.device_kv_caches.items():
1936 - cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
1937 - for cache in cache_list:
1938 - if self.enable_permute_local_kv and block_size_ratio > 1:
1939 - kv_postprocess_blksize_and_layout_on_receive(
1940 - cache, indices, block_size_ratio
1941 - )
1942 - elif self.enable_permute_local_kv:
1943 - kv_postprocess_layout_on_receive(cache, indices)
1944 - else:
1945 - kv_postprocess_blksize_on_receive(
1946 - cache, indices, block_size_ratio
1947 - )
1883 + for cache in self.device_kv_caches.values():
1884 + if self.enable_permute_local_kv and block_size_ratio > 1:
1885 + kv_postprocess_blksize_and_layout_on_receive(
1886 + cache, indices, block_size_ratio
1887 + )
1888 + elif self.enable_permute_local_kv:
1889 + kv_postprocess_layout_on_receive(cache, indices)
1890 + else:
1891 + kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio)
1892
1893 def post_process_device_kv_on_receive_heterogeneous_attn(
1894 self, block_ids: list[int]
@@ -2398,12 +2342,10 @@ class NixlBaseConnectorWorker:
2342 |1st_split-2nd_split| |1st_split-2nd_split |
2343 """
2344 assert self.transfer_topo is not None
2401 - virtually_split = self.transfer_topo.virtually_split_kv_in_blocks
2402 - if virtually_split and mamba_view:
2345 + if self.transfer_topo.virtually_split_kv_in_blocks and mamba_view:
2346 block_len = self._mamba_ssm_size[not first_split]
2347 else:
2405 - half_block = virtually_split and not self._is_region_replicated(layer_idx)
2406 - block_len = self.block_len_per_layer[layer_idx] // (2 if half_block else 1)
2348 + block_len = self.block_len_per_layer[layer_idx]
2349 return block_len
2350
2351 def get_kv_connector_stats(self) -> KVConnectorStats | None:
vllm/models/minimax_m3/common/ops/sparse_attn.py
+14 -21
@@ -8,7 +8,8 @@ equal the sparse block size (128), so one selected block maps to exactly one
8 page.
9
10 Main K/V cache layout (vLLM):
11 - ``(num_blocks, 2, 128, num_kv_heads, head_dim)`` K=[:,0] V=[:,1]
11 + ``(num_blocks, num_kv_heads, 128, 2 * head_dim)``
12 + K=[..., :head_dim] V=[..., head_dim:]
13
14 Only the paths MiniMax M3 uses are implemented: no attention sink, base-2
15 (exp2/log2) softmax. The decode kernels use split-K (flash-decoding) over the
@@ -50,7 +51,7 @@ _FP8_DTYPES = (
51 @triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"])
52 def _gqa_sparse_fwd_kernel(
53 q_ptr, # [total_q, num_heads, head_dim]
53 - kv_cache_ptr, # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim]
54 + kv_cache_ptr, # main cache: [num_blocks, num_kv_heads, 128, 2*head_dim]
55 t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
56 o_ptr, # [total_q, num_heads, head_dim]
57 block_table_ptr, # [num_reqs, max_blocks]
@@ -68,9 +69,8 @@ def _gqa_sparse_fwd_kernel(
69 stride_qh,
70 stride_qd,
71 stride_kv_blk,
71 - stride_kv_kv,
72 - stride_kv_pos,
72 stride_kv_h,
73 + stride_kv_pos,
74 stride_kv_d,
75 stride_th,
76 stride_tn,
@@ -140,9 +140,8 @@ def _gqa_sparse_fwd_kernel(
140 k = tl.load(
141 kv_cache_ptr
142 + page * stride_kv_blk
143 - + 0 * stride_kv_kv
144 - + off_n[None, :] * stride_kv_pos
143 + pid_kh * stride_kv_h
144 + + off_n[None, :] * stride_kv_pos
145 + off_d[:, None] * stride_kv_d,
146 mask=d_mask[:, None] & pos_mask[None, :],
147 other=0.0,
@@ -162,10 +161,9 @@ def _gqa_sparse_fwd_kernel(
161 v = tl.load(
162 kv_cache_ptr
163 + page * stride_kv_blk
165 - + 1 * stride_kv_kv
166 - + off_n[:, None] * stride_kv_pos
164 + pid_kh * stride_kv_h
168 - + off_d[None, :] * stride_kv_d,
165 + + off_n[:, None] * stride_kv_pos
166 + + (head_dim + off_d[None, :]) * stride_kv_d,
167 mask=pos_mask[:, None] & d_mask[None, :],
168 other=0.0,
169 )
@@ -206,7 +204,7 @@ def _gqa_sparse_fwd_kernel(
204 @triton.jit(do_not_specialize=["decode_query_len"])
205 def _gqa_sparse_decode_kernel(
206 q_ptr, # [total_q, num_heads, head_dim]
209 - kv_cache_ptr, # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim]
207 + kv_cache_ptr, # main cache: [num_blocks, num_kv_heads, 128, 2*head_dim]
208 t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
209 o_ptr, # partial out: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim]
210 lse_ptr, # partial lse (log2): [NUM_TOPK_CHUNKS, total_q, num_heads]
@@ -222,9 +220,8 @@ def _gqa_sparse_decode_kernel(
220 stride_qh,
221 stride_qd,
222 stride_kv_blk,
225 - stride_kv_kv,
226 - stride_kv_pos,
223 stride_kv_h,
224 + stride_kv_pos,
225 stride_kv_d,
226 stride_th,
227 stride_tn,
@@ -300,9 +297,8 @@ def _gqa_sparse_decode_kernel(
297 k = tl.load(
298 kv_cache_ptr
299 + page * stride_kv_blk
303 - + 0 * stride_kv_kv
304 - + off_n[None, :] * stride_kv_pos
300 + pid_kh * stride_kv_h
301 + + off_n[None, :] * stride_kv_pos
302 + off_d[:, None] * stride_kv_d,
303 mask=d_mask[:, None] & pos_mask[None, :],
304 other=0.0,
@@ -319,10 +315,9 @@ def _gqa_sparse_decode_kernel(
315 v = tl.load(
316 kv_cache_ptr
317 + page * stride_kv_blk
322 - + 1 * stride_kv_kv
323 - + off_n[:, None] * stride_kv_pos
318 + pid_kh * stride_kv_h
325 - + off_d[None, :] * stride_kv_d,
319 + + off_n[:, None] * stride_kv_pos
320 + + (head_dim + off_d[None, :]) * stride_kv_d,
321 mask=pos_mask[:, None] & d_mask[None, :],
322 other=0.0,
323 )
@@ -418,7 +413,7 @@ def _merge_topk_attn_out_kernel(
413 @torch.no_grad()
414 def minimax_m3_sparse_attn(
415 q: torch.Tensor, # [total_q, num_heads, head_dim]
421 - kv_cache: torch.Tensor, # [num_blocks, 2, 128, num_kv_heads, head_dim]
416 + kv_cache: torch.Tensor, # [num_blocks, num_kv_heads, 128, 2*head_dim]
417 topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk]
418 block_table: torch.Tensor, # [batch, max_blocks]
419 cu_seqlens_q: torch.Tensor, # [batch+1] int32
@@ -459,7 +454,6 @@ def minimax_m3_sparse_attn(
454 kv_cache.stride(1),
455 kv_cache.stride(2),
456 kv_cache.stride(3),
462 - kv_cache.stride(4),
457 topk_idx.stride(0),
458 topk_idx.stride(1),
459 topk_idx.stride(2),
@@ -476,7 +470,7 @@ def minimax_m3_sparse_attn(
470 @torch.no_grad()
471 def minimax_m3_sparse_attn_decode(
472 q: torch.Tensor, # [total_q, num_heads, head_dim]
479 - kv_cache: torch.Tensor, # [num_blocks, 2, 128, num_kv_heads, head_dim]
473 + kv_cache: torch.Tensor, # [num_blocks, num_kv_heads, 128, 2*head_dim]
474 topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk]
475 block_table: torch.Tensor, # [num_reqs, max_blocks]
476 seq_lens: torch.Tensor, # [num_reqs] int32
@@ -529,7 +523,6 @@ def minimax_m3_sparse_attn_decode(
523 kv_cache.stride(1),
524 kv_cache.stride(2),
525 kv_cache.stride(3),
532 - kv_cache.stride(4),
526 topk_idx.stride(0),
527 topk_idx.stride(1),
528 topk_idx.stride(2),
vllm/models/minimax_m3/common/sparse_attention.py
+9 -4
@@ -104,20 +104,25 @@ class MiniMaxM3SparseBackend(AttentionBackend):
104 head_size: int,
105 cache_dtype_str: str = "auto",
106 ) -> tuple[int, ...]:
107 - return (num_blocks, 2, block_size, num_kv_heads, head_size)
107 + # K and V are packed into the content dim: logical (B, H, N, 2*hs).
108 + return (num_blocks, num_kv_heads, block_size, 2 * head_size)
109
110 @staticmethod
111 def get_kv_cache_stride_order(
112 include_num_layers_dimension: bool = False,
113 ) -> tuple[int, ...]:
113 - # Permutation from get_kv_cache_shape to the actual memory layout.
114 + # `stride_order` indicates the permutation that gets us from
115 + # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory
116 + # layout we want.
117 if include_num_layers_dimension:
118 raise NotImplementedError # no cross-layer KV blocks in M3
119 cache_layout = get_kv_cache_layout()
120 if cache_layout == "NHD":
118 - stride_order = (0, 1, 2, 3, 4)
121 + # (num_blocks, block_size, num_kv_heads, 2*head_size)
122 + stride_order = (0, 2, 1, 3)
123 elif cache_layout == "HND":
120 - stride_order = (0, 1, 3, 2, 4)
124 + # (num_blocks, num_kv_heads, block_size, 2*head_size)
125 + stride_order = (0, 1, 2, 3)
126 else:
127 raise ValueError(f"Unknown cache layout format {cache_layout}.")
128 return stride_order
vllm/models/minimax_m3/nvidia/sparse_attention_msa.py
+1 -2
@@ -79,8 +79,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
79 # strided view directly (topK stays innermost-contiguous).
80 prefill_topk = topk[nd:num_tokens].transpose(0, 1)
81 qp = q[nd:]
82 - k_cache = kv_cache[:, 0].transpose(1, 2)
83 - v_cache = kv_cache[:, 1].transpose(1, 2)
82 + k_cache, v_cache = kv_cache.split(self.head_size, dim=-1)
83 k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr(
84 prefill_topk,
85 p.cu_seqlens_q,
vllm/platforms/rocm.py
+16 -2
@@ -427,8 +427,8 @@ def _get_backend_priorities(
427 ]
428
429 backends = []
430 - # ROCM_ATTN uses (2, num_blocks, ...) KV cache layout which is
431 - # incompatible with KV connectors that require blocks-first layout.
430 + # Keep ROCM_ATTN disabled for KV connectors until connector transfer
431 + # semantics are validated for its asymmetric native K/V cache views.
432 if not use_kv_connector:
433 backends.append(AttentionBackendEnum.ROCM_ATTN)
434 if rocm_aiter_ops.is_mha_enabled():
@@ -511,6 +511,20 @@ class RocmPlatform(Platform):
511 attn_selector_config.use_sparse,
512 attn_selector_config.use_kv_connector,
513 )
514 + from vllm.config import get_current_vllm_config_or_none
515 +
516 + vllm_config = get_current_vllm_config_or_none()
517 + is_encoder_decoder = (
518 + getattr(getattr(vllm_config, "model_config", None), "attn_type", None)
519 + == "encoder_decoder"
520 + )
521 + # ROCM_ATTN still uses a legacy attention layout (KV is the outer
522 + # dimension) that is incompatible with the encoder backend layouts. The
523 + # encoder and decoder need the layouts to match. This is currently
524 + # enforced implicitly.
525 + # TODO: Make this explicit in the selector in a future PR.
526 + if is_encoder_decoder and AttentionBackendEnum.ROCM_ATTN in backend_priorities:
527 + backend_priorities.remove(AttentionBackendEnum.ROCM_ATTN)
528 for priority, backend in enumerate(backend_priorities):
529 try:
530 backend_class = backend.get_class()
vllm/utils/torch_utils.py
+11 -45
@@ -416,26 +416,24 @@ def nvfp4_kv_cache_full_dim(head_size: int) -> int:
416 return head_size // 2 + head_size // 16
417
418
419 -def _nvfp4_split_data_scale(
419 +def nvfp4_split_data_scale(
420 kv_side: torch.Tensor,
421 ) -> tuple[torch.Tensor, torch.Tensor]:
422 - """Split a single NVFP4 KV-side buffer into data and scale views.
422 + """Split one side (K or V) of an NVFP4 KV cache into data and scale.
423
424 - The input is a 4D tensor for one KV side (K or V) whose last
425 - dimension is ``full_dim = data_dim + scale_dim``. The physical
426 - layout within each side is [data | scale], both packed contiguously.
424 + The input is a 4D uint8 tensor whose last dimension is
425 + ``full_dim = data_dim + scale_dim``. The physical layout within each
426 + side is ``[data | scale]``, both packed contiguously.
427 +
428 + The caller is responsible for slicing K and V from the combined cache
429 + first (e.g. ``kv_cache.split(num_kv_heads, dim=1)``).
430
431 Args:
429 - kv_side: 4D uint8 tensor with shape
430 - ``(num_pages, dim_1, dim_2, full_dim)``.
431 - May be in any permutation order (NHD or HND).
432 + kv_side: 4D uint8 tensor ``(B, H, N, full_dim)``.
433
434 Returns:
434 - ``(data, scale)`` where
435 - ``data`` is a uint8 view with shape
436 - ``(num_pages, dim_1, dim_2, data_dim)``.
437 - ``scale`` is a float8_e4m3fn view with shape
438 - ``(num_pages, dim_1, dim_2, scale_dim)``.
435 + ``(data, scale)`` where *data* is uint8 and *scale* is
436 + float8_e4m3fn, both views of the same storage.
437 """
438 num_pages = kv_side.shape[0]
439 dim_1, dim_2 = kv_side.shape[1], kv_side.shape[2]
@@ -468,38 +466,6 @@ def _nvfp4_split_data_scale(
466 return data, scale
467
468
471 -def nvfp4_kv_cache_split_views(kv_cache: torch.Tensor) -> tuple[tuple, tuple]:
472 - """Split an NVFP4 KV cache tensor into data and scale views.
473 -
474 - Accepts either a 5D tensor ``(num_pages, 2, dim_2, dim_3, full_dim)``
475 - or a 4D single-side tensor ``(num_pages, dim_2, dim_3, full_dim)``.
476 -
477 - Per-page layout: [K_data | K_scale | V_data | V_scale].
478 - Each KV side is self-contained (data followed by its scale), so the
479 - 5D case simply splits each side independently.
480 -
481 - The returned views are in the same dim order as the input (NHD or
482 - HND), so callers get views matching whichever order they passed in.
483 -
484 - Args:
485 - kv_cache: 5D or 4D uint8 tensor where the last dimension is
486 - ``full_dim = data_dim + scale_dim = 9 * head_size / 16``.
487 -
488 - Returns:
489 - For 5D input:
490 - ``(k_data, v_data), (k_scale, v_scale)``
491 - For 4D input (single KV side):
492 - ``(data,), (scale,)``
493 - """
494 - if kv_cache.dim() == 4:
495 - data, scale = _nvfp4_split_data_scale(kv_cache)
496 - return (data,), (scale,)
497 -
498 - k_data, k_scale = _nvfp4_split_data_scale(kv_cache[:, 0])
499 - v_data, v_scale = _nvfp4_split_data_scale(kv_cache[:, 1])
500 - return (k_data, v_data), (k_scale, v_scale)
501 -
502 -
469 def create_kv_caches_with_random_flash(
470 num_blocks: int,
471 block_size: int,
vllm/v1/attention/backend.py
+6 -6
@@ -122,11 +122,11 @@ class AttentionBackend(ABC):
122 ) -> tuple[int, ...]:
123 """
124 Get the physical (memory layout) ordering of the kv cache dimensions.
125 - e.g. if the KV cache shape is
126 - [2, num_blocks, block_size, num_heads, head_size],
127 - and get_kv_cache_stride_order returns (1, 3, 0, 2, 4) then the physical
125 + Standard attention backends pack K and V into the content dim, giving
126 + the logical shape [num_blocks, num_heads, block_size, 2 * head_size].
127 + e.g. if get_kv_cache_stride_order returns (0, 2, 1, 3) then the physical
128 ordering of dimensions is
129 - [num_blocks, num_heads, 2, block_size, head_size].
129 + [num_blocks, block_size, num_heads, 2 * head_size].
130
131 If this function is unimplemented / raises NotImplementedError,
132 the physical layout of the KV cache will match the logical shape.
@@ -135,9 +135,9 @@ class AttentionBackend(ABC):
135 include_num_layers_dimension: if True, includes an additional
136 num_layers dimension, which is assumed to be prepended
137 to the logical KV cache shape.
138 - With the above example, a return value (2, 4, 0, 1, 3, 5)
138 + With the above example, a return value (1, 0, 3, 2, 4)
139 corresponds to
140 - [num_blocks, num_heads, num_layers, 2, block_size, head_size].
140 + [num_blocks, num_layers, block_size, num_heads, 2 * head_size].
141
142 If an additional dimension is NOT included in the returned
143 tuple, the physical layout will not include a layers dimension.
vllm/v1/attention/backends/flash_attn.py
+18 -13
@@ -133,25 +133,29 @@ class FlashAttentionBackend(AttentionBackend):
133 ) -> tuple[int, ...]:
134 if block_size % 16 != 0:
135 raise ValueError("Block size must be a multiple of 16.")
136 - return (num_blocks, 2, block_size, num_kv_heads, head_size)
136 + # K and V are packed into the content dim: logical (B, H, N, 2*D).
137 + return (num_blocks, num_kv_heads, block_size, 2 * head_size)
138
139 @staticmethod
140 def get_kv_cache_stride_order(
141 include_num_layers_dimension: bool = False,
142 ) -> tuple[int, ...]:
142 - # `stride_order` indicates the permutation that gets
143 - # us from `get_kv_cache_shape` to the actual memory layout we want.
143 + # `stride_order` indicates the permutation that gets us from
144 + # `get_kv_cache_shape` (logical (B, H, N, 2*D)) to the actual memory
145 + # layout we want.
146 cache_layout = get_kv_cache_layout()
147 if cache_layout == "NHD" and include_num_layers_dimension:
146 - # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size)
147 - return (1, 0, 2, 3, 4, 5)
148 + # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size)
149 + return (1, 0, 3, 2, 4)
150 elif cache_layout == "NHD":
149 - stride_order = (0, 1, 2, 3, 4)
151 + # (num_blocks, block_size, num_kv_heads, 2*head_size)
152 + stride_order = (0, 2, 1, 3)
153 elif cache_layout == "HND" and include_num_layers_dimension:
151 - # (num_blocks, num_kv_heads, num_layers, 2, block_size, head_size)
152 - return (1, 4, 0, 2, 3, 5)
154 + # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size)
155 + return (1, 2, 0, 3, 4)
156 elif cache_layout == "HND":
154 - stride_order = (0, 1, 3, 2, 4)
157 + # (num_blocks, num_kv_heads, block_size, 2*head_size)
158 + stride_order = (0, 1, 2, 3)
159 else:
160 raise ValueError(f"Unknown cache layout format {cache_layout}.")
161 return stride_order
@@ -819,7 +823,7 @@ class FlashAttentionImpl(AttentionImpl):
823 key: shape = [num_tokens, num_kv_heads, head_size]
824 value: shape = [num_tokens, num_kv_heads, head_size]
825 kv_cache: shape =
822 - [num_blocks, 2, block_size, num_kv_heads, head_size]
826 + [num_blocks, num_kv_heads, block_size, 2 * head_size]
827 attn_metadata: Metadata for attention.
828 Returns:
829 shape = [num_tokens, num_heads * head_size]
@@ -866,8 +870,8 @@ class FlashAttentionImpl(AttentionImpl):
870 layer,
871 )
872
869 - # For decoder and cross-attention, use KV cache as before
870 - key_cache, value_cache = kv_cache.unbind(1)
873 + # (B, H, N, 2*D) -> ((B, N, H, D), (B, N, H, D))
874 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
875 # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP).
876 # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment.
877 # See vllm.utils.torch_utils.canonicalize_singleton_dim_strides.
@@ -1075,7 +1079,8 @@ class FlashAttentionImpl(AttentionImpl):
1079
1080 # Scatter write into the KV cache using slot_mapping indices.
1081 # No TMA kernel is invoked here, so stride canonicalization is not needed.
1078 - key_cache, value_cache = kv_cache.unbind(1)
1082 + # (B, H, N, 2*D) -> ((B, N, H, D), (B, N, H, D))
1083 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
1084
1085 # Reshape the input keys and values and store them in the cache.
1086 # Skip this if sharing KV cache with an earlier attention layer.
vllm/v1/attention/backends/flash_attn_diffkv.py
+18 -24
@@ -83,10 +83,12 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend):
83 ) -> tuple[int, ...]:
84 if block_size % 16 != 0:
85 raise ValueError("Block size must be a multiple of 16.")
86 + # Logical (blocks-first, head-major) layout: K and V (with their
87 + # different head sizes) packed in the content dim.
88 return (
89 num_blocks,
88 - block_size,
90 num_kv_heads,
91 + block_size,
92 head_size + FlashAttentionDiffKVBackend.head_size_v,
93 )
94
@@ -94,21 +96,22 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend):
96 def get_kv_cache_stride_order(
97 include_num_layers_dimension: bool = False,
98 ) -> tuple[int, ...]:
97 - # `stride_order` indicates the permutation that gets
98 - # us from `get_kv_cache_shape` to the actual memory layout we want.
99 + # `stride_order` indicates the permutation that gets us from
100 + # `get_kv_cache_shape` (logical (B, H, N, C_k+C_v)) to the actual
101 + # memory layout we want.
102 cache_layout = get_kv_cache_layout()
103 if cache_layout == "NHD" and include_num_layers_dimension:
101 - # (num_blocks, num_layers, block_size,
102 - # num_kv_heads, head_size + head_size_v)
103 - return (1, 0, 2, 3, 4)
104 + # (num_blocks, num_layers, block_size, num_kv_heads, C_k+C_v)
105 + return (1, 0, 3, 2, 4)
106 elif cache_layout == "NHD":
105 - stride_order = (0, 1, 2, 3)
107 + # (num_blocks, block_size, num_kv_heads, C_k+C_v)
108 + stride_order = (0, 2, 1, 3)
109 elif cache_layout == "HND" and include_num_layers_dimension:
107 - # (num_blocks, num_kv_heads, num_layers,
108 - # block_size, head_size + head_size_v)
109 - return (1, 3, 0, 2, 4)
110 + # (num_blocks, num_kv_heads, num_layers, block_size, C_k+C_v)
111 + return (1, 2, 0, 3, 4)
112 elif cache_layout == "HND":
111 - stride_order = (0, 2, 1, 3)
113 + # (num_blocks, num_kv_heads, block_size, C_k+C_v)
114 + stride_order = (0, 1, 2, 3)
115 else:
116 raise ValueError(f"Unknown cache layout format {cache_layout}.")
117 return stride_order
@@ -142,21 +145,14 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):
145 # we use direct Q, K, V tensors without caching
146 return
147
145 - # Unlike standard FlashAttn which splits kv_cache via unbind(0),
148 # DiffKV packs K and V into a single tensor along the last dim:
149 # kv_cache shape: [num_blocks, block_size, num_kv_heads,
150 # head_size_k + head_size_v]
149 - # The triton kernel handles this combined layout directly.
150 - #
151 - # NOTE(woosuk): key and value are padded while slot_mapping is
152 - # not padded. However, we don't need to do key[:num_actual_tokens]
153 - # and value[:num_actual_tokens] because the reshape_and_cache_flash
154 - # op uses the slot_mapping's shape to determine the number of
155 - # actual tokens.
151 + # (B, H, N, C) -> (B, N, H, C) for kernel compatibility.
152 triton_reshape_and_cache_flash_diffkv(
153 key,
154 value,
159 - kv_cache,
155 + kv_cache.transpose(1, 2),
156 slot_mapping,
157 self.kv_cache_dtype,
158 layer._k_scale,
@@ -229,10 +225,8 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):
225 layer,
226 )
227
232 - # For decoder and cross-attention, use KV cache as before
233 - # Different head_size for K and V
234 - key_cache = kv_cache[..., : self.head_size]
235 - value_cache = kv_cache[..., self.head_size :]
228 + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
229 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
230 # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP).
231 # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment.
232 # See vllm.utils.torch_utils.canonicalize_singleton_dim_strides.
vllm/v1/attention/backends/flashinfer.py
+97 -43
@@ -51,7 +51,7 @@ from vllm.utils.torch_utils import (
51 is_quantized_kv_cache,
52 is_strictly_contiguous,
53 nvfp4_kv_cache_full_dim,
54 - nvfp4_kv_cache_split_views,
54 + nvfp4_split_data_scale,
55 )
56 from vllm.v1.attention.backend import (
57 AttentionBackend,
@@ -111,9 +111,12 @@ def _trtllm_prefill_attn_kvfp8_dequant(
111 src_stride_page,
112 src_stride_kv,
113 src_stride_head,
114 + src_stride_block,
115 + src_stride_head_size,
116 DST_K_CACHE_STRIDE: tl.constexpr,
117 DST_KV_CACHE_STRIDE: tl.constexpr,
118 HEAD_STRIDE: tl.constexpr,
119 + HEAD_SIZE: tl.constexpr,
120 NUM_KV_HEADS: tl.constexpr,
121 ):
122 batch_idx = tl.program_id(0).to(tl.int64)
@@ -129,18 +132,25 @@ def _trtllm_prefill_attn_kvfp8_dequant(
132 v_scale_val = tl.load(v_scale_ptr)
133
134 mock_page_idx = batch_idx * block_table_stride + mock_block_table_idx + 1
132 - head_offsets = tl.arange(0, HEAD_STRIDE)
135 + logical_offsets = tl.arange(0, HEAD_STRIDE)
136 + block_offsets = logical_offsets // HEAD_SIZE
137 + head_size_offsets = logical_offsets % HEAD_SIZE
138
139 for h in range(NUM_KV_HEADS):
140 h_off = tl.cast(h, tl.int64)
141
142 # Read K from source (supports non-contiguous page/kv/head strides)
138 - src_k = orig_page_num * src_stride_page + h_off * src_stride_head + head_offsets
143 + src_k = (
144 + orig_page_num * src_stride_page
145 + + h_off * src_stride_head
146 + + block_offsets * src_stride_block
147 + + head_size_offsets * src_stride_head_size
148 + )
149 fp8_k = tl.load(kv_cache_ptr + src_k)
150 dequant_k = (fp8_k.to(tl.float32) * k_scale_val).to(dequant_dtype)
151
152 # Write K to contiguous mock cache
143 - dst_k = mock_page_idx * DST_KV_CACHE_STRIDE + h * HEAD_STRIDE + head_offsets
153 + dst_k = mock_page_idx * DST_KV_CACHE_STRIDE + h * HEAD_STRIDE + logical_offsets
154 tl.store(mock_kv_cache_ptr + dst_k, dequant_k)
155
156 # Read V from source (offset by src_stride_kv for the V half)
@@ -148,7 +158,8 @@ def _trtllm_prefill_attn_kvfp8_dequant(
158 orig_page_num * src_stride_page
159 + src_stride_kv
160 + h_off * src_stride_head
151 - + head_offsets
161 + + block_offsets * src_stride_block
162 + + head_size_offsets * src_stride_head_size
163 )
164 fp8_v = tl.load(kv_cache_ptr + src_v)
165 dequant_v = (fp8_v.to(tl.float32) * v_scale_val).to(dequant_dtype)
@@ -158,7 +169,7 @@ def _trtllm_prefill_attn_kvfp8_dequant(
169 mock_page_idx * DST_KV_CACHE_STRIDE
170 + DST_K_CACHE_STRIDE
171 + h * HEAD_STRIDE
161 - + head_offsets
172 + + logical_offsets
173 )
174 tl.store(mock_kv_cache_ptr + dst_v, dequant_v)
175
@@ -175,17 +186,14 @@ def trtllm_prefill_attn_kvfp8_dequant(
186 assert s[1] == 2
187 assert dequant_dtype in (torch.bfloat16, torch.float16)
188
189 + # Logical source layout is (B, 2, H, N, D). The tensor may be a
190 + # non-contiguous view, so the Triton kernel indexes it with actual strides.
191 + strides = kv_cache.stride()
192 num_kv_heads, block_size, head_size = s[2], s[3], s[4]
193 head_stride = block_size * head_size
194 k_cache_stride = num_kv_heads * head_stride
195 kv_cache_stride = k_cache_stride * s[1]
196
183 - strides = kv_cache.stride()
184 - assert strides[3] == head_size and strides[4] == 1, (
185 - "For kv cache layouts, (block_size, head_size) "
186 - f"dimensions must be contiguous, got strides {strides}"
187 - )
188 -
197 new_s = (batch_size * num_of_page_per_token + 1, s[1], s[2], s[3], s[4])
198 # mock kv cache contains just the pages needed by this prefill
199 mock_kv_cache = torch.empty(new_s, dtype=dequant_dtype, device=kv_cache.device)
@@ -207,9 +215,12 @@ def trtllm_prefill_attn_kvfp8_dequant(
215 strides[0],
216 strides[1],
217 strides[2],
218 + strides[3],
219 + strides[4],
220 k_cache_stride,
221 kv_cache_stride,
222 head_stride,
223 + head_size,
224 num_kv_heads,
225 )
226 return mock_kv_cache, mock_block_table
@@ -286,7 +297,7 @@ class BatchDCPPrefillWrapper:
297 self,
298 layer: torch.nn.Module,
299 prefill_query: torch.Tensor,
289 - kv_cache_permute: torch.Tensor,
300 + kv_cache_tuple: tuple[torch.Tensor, torch.Tensor],
301 key: torch.Tensor,
302 value: torch.Tensor,
303 out: torch.Tensor,
@@ -296,7 +307,7 @@ class BatchDCPPrefillWrapper:
307 )
308 output_context_tmp, lse_context_tmp = self._context.run(
309 prefill_query_across_dcp,
299 - kv_cache_permute,
310 + kv_cache_tuple,
311 k_scale=layer._k_scale_float,
312 v_scale=layer._v_scale_float,
313 return_lse=True,
@@ -386,28 +397,31 @@ class FlashInferBackend(AttentionBackend):
397 cache_dtype_str: str = "auto",
398 ) -> tuple[int, ...]:
399 if cache_dtype_str == "nvfp4":
389 - # Packed layout: fp4 data + fp8 block scales in last dim
390 - last_dim = nvfp4_kv_cache_full_dim(head_size)
391 - return (num_blocks, 2, block_size, num_kv_heads, last_dim)
392 - return (num_blocks, 2, block_size, num_kv_heads, head_size)
400 + full_dim = nvfp4_kv_cache_full_dim(head_size)
401 + return (num_blocks, 2 * num_kv_heads, block_size, full_dim)
402 + # Pack K and V in the content dim (B, H, N, 2*hs).
403 + return (num_blocks, num_kv_heads, block_size, 2 * head_size)
404
405 @staticmethod
406 def get_kv_cache_stride_order(
407 include_num_layers_dimension: bool = False,
408 ) -> tuple[int, ...]:
409 # `stride_order` indicates the permutation that gets us from
399 - # `get_kv_cache_shape` to the actual memory layout we want.
410 + # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory
411 + # layout we want.
412 cache_layout = get_kv_cache_layout()
413 if cache_layout == "NHD" and include_num_layers_dimension:
402 - # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size)
403 - return (1, 0, 2, 3, 4, 5)
414 + # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size)
415 + return (1, 0, 3, 2, 4)
416 elif cache_layout == "NHD":
405 - stride_order = (0, 1, 2, 3, 4)
417 + # (num_blocks, block_size, num_kv_heads, 2*head_size)
418 + stride_order = (0, 2, 1, 3)
419 elif cache_layout == "HND" and include_num_layers_dimension:
407 - # (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
408 - return (1, 2, 4, 0, 3, 5)
420 + # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size)
421 + return (1, 2, 0, 3, 4)
422 elif cache_layout == "HND":
410 - stride_order = (0, 1, 3, 2, 4)
423 + # (num_blocks, num_kv_heads, block_size, 2*head_size)
424 + stride_order = (0, 1, 2, 3)
425 else:
426 raise ValueError(f"Unknown cache layout format {cache_layout}.")
427 return stride_order
@@ -1714,7 +1728,19 @@ class FlashInferImpl(AttentionImpl):
1728 if attn_metadata.use_cascade:
1729 # Cascade attention (rare case).
1730 assert attn_metadata.cascade_wrapper is not None
1717 - output.copy_(attn_metadata.cascade_wrapper.run(query, kv_cache))
1731 + stride_order = FlashInferBackend.get_kv_cache_stride_order()
1732 + if self.is_kvcache_nvfp4:
1733 + kv_cache_views = tuple(
1734 + cache.permute(*stride_order)
1735 + for cache in kv_cache.split(self.num_kv_heads, dim=1)
1736 + )
1737 + else:
1738 + kv_perm = kv_cache.permute(*stride_order)
1739 + kv_cache_views = kv_perm.split(self.head_size, dim=-1)
1740 + kv_tuple = tuple(
1741 + canonicalize_singleton_dim_strides(cache) for cache in kv_cache_views
1742 + )
1743 + output.copy_(attn_metadata.cascade_wrapper.run(query, kv_tuple))
1744 return output
1745
1746 # When using spec decoding, num_decodes can be < num_decode_tokens
@@ -1740,14 +1766,23 @@ class FlashInferImpl(AttentionImpl):
1766 )
1767 kv_cache_permute = fixed
1768
1743 - # For NVFP4, the kv_cache last dim is full_dim (data + scale packed).
1744 - # Split into correctly-strided data and scale views.
1769 + # Split K/V — zero-copy views. NVFP4 stores K/V as separate head
1770 + # groups; other dtypes pack K/V in the content dim.
1771 + hs = self.head_size
1772 nvfp4_kv_data = None
1773 nvfp4_kv_block_scales = None
1774 if self.is_kvcache_nvfp4:
1748 - nvfp4_kv_data, nvfp4_kv_block_scales = nvfp4_kv_cache_split_views(
1749 - kv_cache_permute
1775 + k_cache, v_cache = kv_cache.split(self.num_kv_heads, dim=1)
1776 + kv_cache_tuple = (
1777 + canonicalize_singleton_dim_strides(k_cache.permute(*stride_order)),
1778 + canonicalize_singleton_dim_strides(v_cache.permute(*stride_order)),
1779 )
1780 + k_data, k_sf = nvfp4_split_data_scale(kv_cache_tuple[0])
1781 + v_data, v_sf = nvfp4_split_data_scale(kv_cache_tuple[1])
1782 + nvfp4_kv_data = (k_data, v_data)
1783 + nvfp4_kv_block_scales = (k_sf, v_sf)
1784 + else:
1785 + kv_cache_tuple = kv_cache_permute.split(hs, dim=-1)
1786
1787 use_dcp = self.dcp_world_size > 1
1788
@@ -1786,7 +1821,7 @@ class FlashInferImpl(AttentionImpl):
1821 prefill_wrapper.run(
1822 layer,
1823 prefill_query,
1789 - kv_cache_permute,
1824 + kv_cache_tuple,
1825 key[num_decode_tokens:],
1826 value[num_decode_tokens:],
1827 out=output[num_decode_tokens:],
@@ -1803,7 +1838,9 @@ class FlashInferImpl(AttentionImpl):
1838 assert prefill_wrapper._causal == attn_metadata.causal
1839
1840 if self.is_kvcache_nvfp4:
1806 - kv_cache_permute = nvfp4_kv_data
1841 + kv_cache_for_fi = nvfp4_kv_data
1842 + else:
1843 + kv_cache_for_fi = kv_cache_tuple
1844 kv_cache_sf = (
1845 nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None
1846 )
@@ -1821,7 +1858,7 @@ class FlashInferImpl(AttentionImpl):
1858
1859 prefill_wrapper.run(
1860 prefill_query,
1824 - kv_cache_permute,
1861 + kv_cache_for_fi,
1862 q_scale=layer._q_scale_float,
1863 k_scale=layer._k_scale_float,
1864 v_scale=layer._v_scale_float,
@@ -1888,8 +1925,7 @@ class FlashInferImpl(AttentionImpl):
1925 # TRTLLM prefill attention does not support BF16 Q
1926 # and fp8 kv cache. So to enable prefill attention
1927 # with fp8 kv cache, we can construct a mock block
1891 - # and mock kv cache with BF16 KV involved in the prefill
1892 - #
1928 + # and mock kv cache with BF16 KV involved in the prefill.
1929 kv_cache_permute = canonicalize_singleton_dim_strides(
1930 kv_cache_permute
1931 )
@@ -1901,15 +1937,21 @@ class FlashInferImpl(AttentionImpl):
1937 "KV cache inner dims (block_size, head_size) must be "
1938 f"contiguous, got strides {kv_strides}"
1939 )
1940 + # fp8 uses (B, H, N, 2*hs); reshape to (B, 2, H, N, hs)
1941 + # for the dequant kernel — zero-copy view. The dequant
1942 + # kernel handles the interleaved K/V block stride.
1943 + B_kv, H_kv, N_kv = kv_cache_permute.shape[:3]
1944 + kv_cache_5d = kv_cache_permute.view(B_kv, H_kv, N_kv, 2, hs)
1945 + kv_cache_5d = kv_cache_5d.permute(0, 3, 1, 2, 4)
1946 mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
1905 - kv_cache_permute,
1947 + kv_cache_5d,
1948 block_tables_prefill,
1949 layer._k_scale,
1950 layer._v_scale,
1951 attn_metadata.q_data_type_prefill,
1952 )
1953 else:
1912 - mock_kv_cache = kv_cache_permute
1954 + mock_kv_cache = kv_cache_tuple
1955 mock_block_table = block_tables_prefill
1956
1957 trtllm_batch_context_with_kv_cache(
@@ -1957,7 +1999,9 @@ class FlashInferImpl(AttentionImpl):
1999 assert decode_wrapper._sm_scale == self.scale
2000
2001 if self.is_kvcache_nvfp4:
1960 - kv_cache_permute = nvfp4_kv_data
2002 + kv_cache_for_fi = nvfp4_kv_data
2003 + else:
2004 + kv_cache_for_fi = kv_cache_tuple
2005 kv_cache_sf = nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None
2006
2007 # NVFP4 kernel only supports FP8 output.
@@ -1980,7 +2024,7 @@ class FlashInferImpl(AttentionImpl):
2024 )
2025 decode_wrapper.run(
2026 decode_query,
1983 - kv_cache_permute,
2027 + kv_cache_for_fi,
2028 q_scale=layer._q_scale_float,
2029 k_scale=layer._k_scale_float,
2030 v_scale=layer._v_scale_float,
@@ -1997,7 +2041,7 @@ class FlashInferImpl(AttentionImpl):
2041 else:
2042 decode_wrapper.run(
2043 decode_query,
2000 - kv_cache_permute,
2044 + kv_cache_for_fi,
2045 q_scale=layer._q_scale_float,
2046 k_scale=layer._k_scale_float,
2047 v_scale=layer._v_scale_float,
@@ -2079,7 +2123,7 @@ class FlashInferImpl(AttentionImpl):
2123 trtllm_batch_decode_with_kv_cache(
2124 query=decode_query,
2125 kv_cache=(
2082 - nvfp4_kv_data if self.is_kvcache_nvfp4 else kv_cache_permute
2126 + nvfp4_kv_data if self.is_kvcache_nvfp4 else kv_cache_tuple
2127 ),
2128 workspace_buffer=workspace_buffer,
2129 block_tables=block_tables_decode,
@@ -2119,8 +2163,18 @@ class FlashInferImpl(AttentionImpl):
2163 # and value[:num_actual_tokens] because the reshape_and_cache_flash
2164 # op uses the slot_mapping's shape to determine the number of
2165 # actual tokens.
2122 - k_cache = kv_cache[:, 0]
2123 - v_cache = kv_cache[:, 1]
2166 + if self.is_kvcache_nvfp4:
2167 + # (B, 2*H, N, full_dim) -> ((B, N, H, full_dim),
2168 + # (B, N, H, full_dim));
2169 + # K heads first, then V heads.
2170 + k_cache, v_cache = kv_cache.transpose(1, 2).split(
2171 + self.num_kv_heads, dim=-2
2172 + )
2173 + else:
2174 + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
2175 + k_cache, v_cache = kv_cache.transpose(1, 2).split(
2176 + self.head_size, dim=-1
2177 + )
2178 torch.ops._C_cache_ops.reshape_and_cache_flash(
2179 key,
2180 value,
vllm/v1/attention/backends/flex_attention.py
+11 -7
@@ -130,15 +130,16 @@ class FlexAttentionBackend(AttentionBackend):
130 head_size: int,
131 cache_dtype_str: str = "auto",
132 ) -> tuple[int, ...]:
133 - return (num_blocks, 2, block_size, num_kv_heads, head_size)
133 + # K and V are packed into the content dim: logical (B, H, N, 2*hs).
134 + return (num_blocks, num_kv_heads, block_size, 2 * head_size)
135
136 @staticmethod
137 def get_kv_cache_stride_order(
138 include_num_layers_dimension: bool = False,
139 ) -> tuple[int, ...]:
140 if include_num_layers_dimension:
140 - return (1, 0, 3, 2, 4, 5)
141 - return (0, 2, 1, 3, 4)
141 + return (1, 0, 3, 2, 4)
142 + return (0, 2, 1, 3)
143
144 @staticmethod
145 def get_builder_cls() -> type["FlexAttentionMetadataBuilder"]:
@@ -1242,7 +1243,8 @@ class FlexAttentionImpl(AttentionImpl):
1243 if self.attn_type == AttentionType.ENCODER_ONLY:
1244 return
1245
1245 - key_cache, value_cache = kv_cache.unbind(1)
1246 + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
1247 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
1248 torch.ops._C_cache_ops.reshape_and_cache_flash(
1249 key,
1250 value,
@@ -1273,7 +1275,7 @@ class FlexAttentionImpl(AttentionImpl):
1275 key: shape = [num_tokens, num_kv_heads, head_size]
1276 value: shape = [num_tokens, num_kv_heads, head_size]
1277 kv_cache: shape =
1276 - [num_blocks, 2, block_size, num_kv_heads, head_size]
1278 + [num_blocks, num_kv_heads, block_size, 2 * head_size]
1279 attn_metadata: Metadata for attention.
1280 Returns:
1281 shape = [num_tokens, num_heads * head_size]
@@ -1347,9 +1349,11 @@ class FlexAttentionImpl(AttentionImpl):
1349
1350 else:
1351 assert self.attn_type == AttentionType.DECODER
1350 - key_cache, value_cache = kv_cache.unbind(1)
1352 + kv_cache = kv_cache.transpose(1, 2)
1353 + hs = self.head_size
1354 + key_cache, value_cache = kv_cache.split(hs, dim=-1)
1355
1352 - # Flatten (num_blocks, block_size) into a single token dim
1356 + # Flatten (num_blocks, block_size) into a single token dim.
1357 key_cache = key_cache.view(-1, self.num_kv_heads, self.head_size)
1358 value_cache = value_cache.view(-1, self.num_kv_heads, self.head_size)
1359 query, key_tensor, value_tensor = map(
vllm/v1/attention/backends/rocm_aiter_fa.py
+8 -4
@@ -768,7 +768,8 @@ class AiterFlashAttentionBackend(AttentionBackend):
768 ) -> tuple[int, ...]:
769 if block_size % 16 != 0:
770 raise ValueError("Block size must be a multiple of 16.")
771 - return (num_blocks, 2, block_size, num_kv_heads, head_size)
771 + # K and V are packed into the content dim: logical (B, H, N, 2*hs).
772 + return (num_blocks, num_kv_heads, block_size, 2 * head_size)
773
774 @classmethod
775 def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
@@ -1061,7 +1062,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
1062 # Whenever making a change in this method, please benchmark the
1063 # performance to make sure it does not introduce any overhead.
1064 num_actual_tokens = attn_metadata.num_actual_tokens
1064 - key_cache, value_cache = kv_cache.unbind(1)
1065 + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
1066 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
1067
1068 if is_quantized_kv_cache(self.kv_cache_dtype):
1069 key_cache = key_cache.view(current_platform.fp8_dtype())
@@ -1388,7 +1390,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
1390 kv_cache: torch.Tensor,
1391 slot_mapping: torch.Tensor,
1392 ):
1391 - key_cache, value_cache = kv_cache.unbind(1)
1393 + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
1394 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
1395
1396 # key and value may be None in the case of cross attention. They are
1397 # calculated once based on the output from the encoder and then cached
@@ -1454,7 +1457,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
1457 kv_cache: torch.Tensor,
1458 layer_slot_mapping: torch.Tensor,
1459 ):
1457 - key_cache, value_cache = kv_cache.unbind(1)
1460 + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
1461 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
1462 flash_layout = True
1463
1464 is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+4 -5
@@ -87,7 +87,8 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
87 ) -> tuple[int, ...]:
88 if block_size % 16 != 0:
89 raise ValueError("Block size must be a multiple of 16.")
90 - return (num_blocks, 2, block_size, num_kv_heads, head_size)
90 + # K and V are packed into the content dim: logical (B, H, N, 2*hs).
91 + return (num_blocks, num_kv_heads, block_size, 2 * head_size)
92
93 @staticmethod
94 def use_cascade_attention(*args, **kwargs) -> bool:
@@ -150,10 +151,8 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
151 def _split_kv_cache(
152 self, kv_cache: torch.Tensor
153 ) -> tuple[torch.Tensor, torch.Tensor]:
153 - # Blocks-first ``(num_blocks, 2, ...)``. The model runner normalizes any
154 - # shared decoder/cross-attention allocation to this layout, so no
155 - # per-backend restriding is needed here.
156 - return kv_cache.unbind(1)
154 + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
155 + return kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
156
157 def forward(
158 self,
vllm/v1/attention/backends/triton_attn.py
+60 -42
@@ -324,6 +324,7 @@ class TritonAttentionBackend(AttentionBackend):
324 ) -> tuple[int, ...]:
325 if block_size % 16 != 0:
326 raise ValueError("Block size must be a multiple of 16.")
327 + # K and V are packed into the content dim: logical (B, H, N, 2*hs).
328 if kv_cache_uses_per_token_head_scales(cache_dtype_str):
329 # Pad the head dim by sizeof(float32)/sizeof(cache_dtype) so the
330 # per-(token, head) scale fits inline after the quantized data;
@@ -341,26 +342,30 @@ class TritonAttentionBackend(AttentionBackend):
342 data_head_size = head_size // 2
343 else:
344 data_head_size = head_size
344 - return (num_blocks, 2, block_size, num_kv_heads, data_head_size + scale_pad)
345 - return (num_blocks, 2, block_size, num_kv_heads, head_size)
345 + padded_hs = data_head_size + scale_pad
346 + return (num_blocks, num_kv_heads, block_size, 2 * padded_hs)
347 + return (num_blocks, num_kv_heads, block_size, 2 * head_size)
348
349 @staticmethod
350 def get_kv_cache_stride_order(
351 include_num_layers_dimension: bool = False,
352 ) -> tuple[int, ...]:
351 - # `stride_order` indicates the permutation that gets
352 - # us from `get_kv_cache_shape` to the actual memory layout we want.
353 + # `stride_order` indicates the permutation that gets us from
354 + # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory
355 + # layout we want.
356 cache_layout = get_kv_cache_layout()
357 if cache_layout == "NHD" and include_num_layers_dimension:
355 - # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size)
356 - return (1, 0, 2, 3, 4, 5)
358 + # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size)
359 + return (1, 0, 3, 2, 4)
360 elif cache_layout == "NHD":
358 - stride_order = (0, 1, 2, 3, 4)
361 + # (num_blocks, block_size, num_kv_heads, 2*head_size)
362 + stride_order = (0, 2, 1, 3)
363 elif cache_layout == "HND" and include_num_layers_dimension:
360 - # (num_blocks, num_kv_heads, num_layers, 2, block_size, head_size)
361 - return (1, 4, 0, 2, 3, 5)
364 + # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size)
365 + return (1, 2, 0, 3, 4)
366 elif cache_layout == "HND":
363 - stride_order = (0, 1, 3, 2, 4)
367 + # (num_blocks, num_kv_heads, block_size, 2*head_size)
368 + stride_order = (0, 1, 2, 3)
369 else:
370 raise ValueError(f"Unknown cache layout: {cache_layout}")
371 return stride_order
@@ -410,12 +415,16 @@ class TritonAttentionImpl(AttentionImpl):
415 _v_scale_cache: torch.Tensor | None = None
416
417 def _ensure_scale_caches(self, kv_cache: torch.Tensor) -> None:
413 - """Extract per-head scale views from the padded head dimension.
418 + """Extract per-head scale views from the padded content dimension.
419
415 - The KV cache shape is ``(num_blocks, 2, block_size, nkv, hs+pad)``
416 - where ``pad = sizeof(float32) / sizeof(cache_dtype)``. The last
417 - ``pad`` elements of each head hold one float32 scale. We create
418 - strided float32 views over those bytes.
420 + The KV cache is packed as logical shape
421 + ``(num_blocks, nkv, block_size, 2 * (hs + pad))`` where
422 + ``pad = sizeof(float32) / sizeof(cache_dtype)``. The content dim holds
423 + ``[K(hs) | K_scale(pad) | V(hs) | V_scale(pad)]`` per (head, slot); the
424 + last ``pad`` elements of each half hold one float32 scale. We create
425 + strided float32 views over those bytes. ``kv_cache`` must be the
426 + packed logical tensor (call before any transpose), but may have HND or
427 + NHD physical strides.
428
429 Scale shape: ``(num_blocks, block_size, num_kv_heads)``
430 """
@@ -423,9 +432,10 @@ class TritonAttentionImpl(AttentionImpl):
432 return
433 from vllm.utils.torch_utils import get_dtype_size
434
426 - num_blocks, _, block_size, nkv, padded_hs = kv_cache.shape
435 + num_blocks, nkv, block_size, content = kv_cache.shape
436 dtype_sz = kv_cache.element_size()
437 scale_pad = get_dtype_size(torch.float32) // dtype_sz # e.g. 4
438 + padded_hs = content // 2
439 hs = padded_hs - scale_pad
440
441 raw = kv_cache.untyped_storage()
@@ -433,31 +443,37 @@ class TritonAttentionImpl(AttentionImpl):
443 raw
444 )
445
436 - # In the raw bytes, each (block, kv_half, slot, head) occupies
437 - # padded_hs * dtype_sz bytes. The scale float32 sits at byte
438 - # offset hs * dtype_sz within that region.
439 - kv_half_bytes = block_size * nkv * padded_hs * dtype_sz
440 - full_block_f32 = 2 * kv_half_bytes // 4 # stride between blocks
441 - slot_f32 = nkv * padded_hs * dtype_sz // 4 # stride between slots
442 - head_f32 = padded_hs * dtype_sz // 4 # stride between heads
443 - scale_off_f32 = hs * dtype_sz // 4 # offset to scale within head
444 -
445 - # K scales: kv_half=0
446 + def to_f32_units(elements: int) -> int:
447 + nbytes = elements * dtype_sz
448 + assert nbytes % 4 == 0
449 + return nbytes // 4
450 +
451 + # Actual strides (in float32 units) from the tensor. The logical cache
452 + # may be physically NHD, so do not assume C-contiguous HND layout.
453 + strides = kv_cache.stride()
454 + block_f32 = to_f32_units(strides[0])
455 + head_f32 = to_f32_units(strides[1])
456 + slot_f32 = to_f32_units(strides[2])
457 + # Scale sits at byte offset hs within each (K, then V) content half.
458 + base_off_f32 = to_f32_units(kv_cache.storage_offset())
459 + k_scale_off_f32 = base_off_f32 + to_f32_units(hs)
460 + v_scale_off_f32 = base_off_f32 + to_f32_units(padded_hs + hs)
461 +
462 + # K scales (first content half)
463 self._k_scale_cache = torch.as_strided(
464 base_f32,
465 size=(num_blocks, block_size, nkv),
449 - stride=(full_block_f32, slot_f32, head_f32),
450 - storage_offset=scale_off_f32,
466 + stride=(block_f32, slot_f32, head_f32),
467 + storage_offset=k_scale_off_f32,
468 )
469 self._k_scale_cache.fill_(1.0)
470
454 - # V scales: kv_half=1, offset by kv_half_bytes
455 - v_base_f32 = kv_half_bytes // 4
471 + # V scales (second content half)
472 self._v_scale_cache = torch.as_strided(
473 base_f32,
474 size=(num_blocks, block_size, nkv),
459 - stride=(full_block_f32, slot_f32, head_f32),
460 - storage_offset=v_base_f32 + scale_off_f32,
475 + stride=(block_f32, slot_f32, head_f32),
476 + storage_offset=v_scale_off_f32,
477 )
478 self._v_scale_cache.fill_(1.0)
479
@@ -576,7 +592,7 @@ class TritonAttentionImpl(AttentionImpl):
592 key: shape = [num_tokens, num_kv_heads, head_size]
593 value: shape = [num_tokens, num_kv_heads, head_size]
594 kv_cache: shape =
579 - [num_blocks, 2, block_size, num_kv_heads, head_size]
595 + [num_blocks, num_kv_heads, block_size, 2 * head_size]
596 attn_metadata: Metadata for attention.
597 Returns:
598 shape = [num_tokens, num_heads * head_size]
@@ -617,6 +633,7 @@ class TritonAttentionImpl(AttentionImpl):
633 layer,
634 )
635
636 + # KV cache arrives in logical (B, H, N, 2*hs) order.
637 # Per-token-head quantized KV cache: handled by the core unified
638 # kernel, which dequantizes per-(token, head) inline via constexpr
639 # branches (INT8 / FP8) and dispatches to the packed INT4 kernel.
@@ -627,7 +644,9 @@ class TritonAttentionImpl(AttentionImpl):
644 q_descale = k_descale = v_descale = None
645 # FP8 per-tensor / auto path (original flow).
646 else:
630 - key_cache, value_cache = kv_cache.unbind(1)
647 + kv_cache = kv_cache.transpose(1, 2)
648 + hs = self.head_size
649 + key_cache, value_cache = kv_cache.split(hs, dim=-1)
650 if (
651 is_quantized_kv_cache(self.kv_cache_dtype)
652 and key_cache.dtype != self.fp8_dtype
@@ -711,7 +730,8 @@ class TritonAttentionImpl(AttentionImpl):
730 ) -> tuple[torch.Tensor, torch.Tensor]:
731 """Per-token-head K/V cache views (ensures scale caches; FP8 retyped)."""
732 self._ensure_scale_caches(kv_cache)
714 - key_cache, value_cache = kv_cache.unbind(1)
733 + padded_hs = kv_cache.shape[-1] // 2
734 + key_cache, value_cache = kv_cache.transpose(1, 2).split(padded_hs, dim=-1)
735 if self._kv_quant_mode == KVQuantMode.FP8_PER_TOKEN_HEAD:
736 key_cache = key_cache.view(self.fp8_dtype)
737 value_cache = value_cache.view(self.fp8_dtype)
@@ -778,13 +798,9 @@ class TritonAttentionImpl(AttentionImpl):
798 return
799 # Reshape the input keys and values and store them in the cache.
800 if self._is_per_token_head_quant:
781 - self._ensure_scale_caches(kv_cache)
782 - key_cache, value_cache = kv_cache.unbind(1)
801 + key_cache, value_cache = self._pth_key_value_caches(kv_cache)
802 k_scale_cache = self._k_scale_cache
803 v_scale_cache = self._v_scale_cache
785 - if self._kv_quant_mode == KVQuantMode.FP8_PER_TOKEN_HEAD:
786 - key_cache = key_cache.view(self.fp8_dtype)
787 - value_cache = value_cache.view(self.fp8_dtype)
804 triton_reshape_and_cache_flash_per_token_head_quant(
805 key,
806 value,
@@ -797,7 +813,8 @@ class TritonAttentionImpl(AttentionImpl):
813 )
814 return
815 # For decoder and cross-attention, use KV cache as before.
800 - key_cache, value_cache = kv_cache.unbind(1)
816 + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
817 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
818 if is_quantized_kv_cache(self.kv_cache_dtype):
819 key_cache = key_cache.view(self.fp8_dtype)
820 value_cache = value_cache.view(self.fp8_dtype)
@@ -829,7 +846,8 @@ class TritonAttentionImpl(AttentionImpl):
846 kv_cache: torch.Tensor,
847 layer_slot_mapping: torch.Tensor,
848 ):
832 - key_cache, value_cache = kv_cache.unbind(1)
849 + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
850 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
851 flash_layout = True
852
853 is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)
vllm/v1/attention/backends/triton_attn_diffkv.py
+23 -22
@@ -2,12 +2,9 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Triton attention backend with different K/V head dimensions (DiffKV).
4
5 -The KV cache layout is identical to ``FlashAttentionDiffKVBackend`` — K
6 -and V are packed along the last dim:
7 -
8 - [num_blocks, block_size, num_kv_heads, head_size_qk + head_size_v]
9 -
10 -so existing helpers (``triton_reshape_and_cache_flash_diffkv``) are reused.
5 +The KV cache layout is identical to ``FlashAttentionDiffKVBackend``: K and V
6 +are packed along the last dim in the logical shape
7 +``[num_blocks, num_kv_heads, block_size, head_size_qk + head_size_v]``.
8 """
9
10 from typing import ClassVar
@@ -106,10 +103,12 @@ class TritonAttentionDiffKVBackend(TritonAttentionBackend):
103 ) -> tuple[int, ...]:
104 if block_size % 16 != 0:
105 raise ValueError("Block size must be a multiple of 16.")
106 + # Logical (blocks-first, head-major) layout: K and V (with their
107 + # different head sizes) packed in the content dim.
108 return (
109 num_blocks,
111 - block_size,
110 num_kv_heads,
111 + block_size,
112 head_size + TritonAttentionDiffKVBackend.head_size_v,
113 )
114
@@ -117,19 +116,22 @@ class TritonAttentionDiffKVBackend(TritonAttentionBackend):
116 def get_kv_cache_stride_order(
117 include_num_layers_dimension: bool = False,
118 ) -> tuple[int, ...]:
119 + # `stride_order` indicates the permutation that gets us from
120 + # `get_kv_cache_shape` (logical (B, H, N, C_k+C_v)) to the actual
121 + # memory layout we want.
122 cache_layout = get_kv_cache_layout()
123 if cache_layout == "NHD" and include_num_layers_dimension:
122 - # (num_blocks, num_layers, block_size,
123 - # num_kv_heads, head_size + head_size_v)
124 - return (1, 0, 2, 3, 4)
124 + # (num_blocks, num_layers, block_size, num_kv_heads, C_k+C_v)
125 + return (1, 0, 3, 2, 4)
126 elif cache_layout == "NHD":
126 - return (0, 1, 2, 3)
127 + # (num_blocks, block_size, num_kv_heads, C_k+C_v)
128 + return (0, 2, 1, 3)
129 elif cache_layout == "HND" and include_num_layers_dimension:
128 - # (num_blocks, num_kv_heads, num_layers,
129 - # block_size, head_size + head_size_v)
130 - return (1, 3, 0, 2, 4)
130 + # (num_blocks, num_kv_heads, num_layers, block_size, C_k+C_v)
131 + return (1, 2, 0, 3, 4)
132 elif cache_layout == "HND":
132 - return (0, 2, 1, 3)
133 + # (num_blocks, num_kv_heads, block_size, C_k+C_v)
134 + return (0, 1, 2, 3)
135 else:
136 raise ValueError(f"Unknown cache layout format {cache_layout}.")
137
@@ -175,13 +177,12 @@ class TritonAttentionDiffKVImpl(TritonAttentionImpl):
177 kv_cache: torch.Tensor,
178 slot_mapping: torch.Tensor,
179 ) -> None:
178 - # Cache is packed [..., head_size_qk + head_size_v]; the diffkv
179 - # reshape kernel writes K to [..., :head_size_qk] and V to
180 - # [..., head_size_qk:hqk+hv].
180 + # Cache is logical (B, H, N, C); the diffkv reshape kernel expects
181 + # (B, N, H, C).
182 triton_reshape_and_cache_flash_diffkv(
183 key,
184 value,
184 - kv_cache,
185 + kv_cache.transpose(1, 2),
186 slot_mapping,
187 self.kv_cache_dtype,
188 layer._k_scale,
@@ -210,7 +211,7 @@ class TritonAttentionDiffKVImpl(TritonAttentionImpl):
211 query: [num_tokens, num_heads, head_size_qk]
212 key: [num_tokens, num_kv_heads, head_size_qk]
213 value: [num_tokens, num_kv_heads, head_size_v]
213 - kv_cache: [num_blocks, block_size, num_kv_heads,
214 + kv_cache: [num_blocks, num_kv_heads, block_size,
215 head_size_qk + head_size_v]
216 output: [num_tokens, num_heads, head_size_v]
217 """
@@ -231,8 +232,8 @@ class TritonAttentionDiffKVImpl(TritonAttentionImpl):
232 head_size_qk = self.head_size
233 head_size_v = TritonAttentionDiffKVBackend.head_size_v
234
234 - # Slice the packed cache into K / V views. Strides on dims 0/1/2
235 - # match the original cache; dim 3 stays contiguous (stride 1).
235 + # Triton DiffKV kernels consume (B, N, H, D) cache views.
236 + kv_cache = kv_cache.transpose(1, 2)
237 key_cache = kv_cache[..., :head_size_qk]
238 value_cache = kv_cache[..., head_size_qk : head_size_qk + head_size_v]
239
vllm/v1/attention/backends/turboquant_attn.py
+9 -3
@@ -141,9 +141,10 @@ class TurboQuantAttentionBackend(AttentionBackend):
141
142 Standard attention backends use (2, num_blocks, block_size, num_kv_heads,
143 head_dim) with a leading 2 to separate K and V. TurboQuant packs K+V
144 - into a single interleaved slot per head per position, so the cache is:
144 + into a single interleaved slot per head per position. The logical
145 + (blocks-first, head-major) shape is:
146
146 - (num_blocks, block_size, num_kv_heads, slot_size_aligned)
147 + (num_blocks, num_kv_heads, block_size, slot_size_aligned)
148
149 Each slot = [key_packed | value_packed | padding].
150 This is safe because TQ has its own get_kv_cache_shape override and
@@ -159,7 +160,7 @@ class TurboQuantAttentionBackend(AttentionBackend):
160 )
161
162 tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype_str, head_size)
162 - return (num_blocks, block_size, num_kv_heads, tq_config.slot_size_aligned)
163 + return (num_blocks, num_kv_heads, block_size, tq_config.slot_size_aligned)
164
165 @classmethod
166 def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
@@ -422,6 +423,8 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
423
424 k = key[:N].view(N, self.num_kv_heads, self.head_size)
425 v = value[:N].view(N, self.num_kv_heads, self.head_size)
426 + # (B, H, N, C) -> (B, N, H, C) for TQ kernels
427 + kv_cache = kv_cache.transpose(1, 2)
428 self._store_kv(k, v, kv_cache, slot_mapping, layer)
429
430 def forward(
@@ -449,6 +452,9 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
452 if attn_metadata is None:
453 return output.fill_(0)
454
455 + # (B, H, N, C) -> (B, N, H, C) for TQ kernels
456 + kv_cache = kv_cache.transpose(1, 2)
457 +
458 # Slice to actual tokens
459 N = attn_metadata.num_actual_tokens
460 if N <= 0:
vllm/v1/attention/ops/triton_reshape_and_cache_flash.py
+5 -4
@@ -473,6 +473,7 @@ def reshape_and_cache_kernel_flash_diffkv(
473 value_stride: tl.int64,
474 block_stride: tl.int64,
475 page_stride: tl.int64,
476 + head_stride: tl.int64,
477 num_heads: tl.constexpr,
478 head_size_k: tl.constexpr,
479 head_size_v: tl.constexpr,
@@ -498,9 +499,7 @@ def reshape_and_cache_kernel_flash_diffkv(
499 src_value_idx = token_idx * value_stride + tile_i * head_size_v
500
501 tgt_idx = (
501 - block_idx * block_stride
502 - + block_offset * page_stride
503 - + tile_i * (head_size_k + head_size_v)
502 + block_idx * block_stride + block_offset * page_stride + tile_i * head_stride
503 )
504
505 # [TILE_SIZE]
@@ -542,7 +541,7 @@ def reshape_and_cache_kernel_flash_diffkv(
541 def triton_reshape_and_cache_flash_diffkv(
542 key: torch.Tensor, # [num_tokens, num_heads, head_size]
543 value: torch.Tensor, # [num_tokens, num_heads, head_size_v]
545 - # [num_blocks, block_size, num_heads, head_size + head_size_v]
544 + # Strided [num_blocks, block_size, num_heads, head_size + head_size_v].
545 kv_cache: torch.Tensor,
546 slot_mapping: torch.Tensor, # [num_tokens]
547 kv_cache_dtype: str, # "auto", "fp8"
@@ -558,6 +557,7 @@ def triton_reshape_and_cache_flash_diffkv(
557 v_stride = value.stride()[0]
558 block_stride = kv_cache.stride()[0]
559 page_stride = kv_cache.stride()[1]
560 + head_stride = kv_cache.stride()[2]
561
562 kv_cache_torch_dtype = (
563 current_platform.fp8_dtype()
@@ -599,6 +599,7 @@ def triton_reshape_and_cache_flash_diffkv(
599 value_stride=v_stride,
600 block_stride=block_stride,
601 page_stride=page_stride,
602 + head_stride=head_stride,
603 num_heads=num_heads,
604 head_size_k=head_size_k,
605 head_size_v=head_size_v,
vllm/v1/attention/ops/triton_turboquant_store.py
+2 -2
@@ -389,7 +389,7 @@ def triton_turboquant_store(
389 _tq_fused_store_fp8[grid](
390 k_flat,
391 v_flat,
392 - kv_cache.view(-1),
392 + kv_cache,
393 slot_mapping,
394 stride_cache_block=stride_block,
395 stride_cache_pos=stride_pos,
@@ -425,7 +425,7 @@ def triton_turboquant_store(
425 norms.squeeze(1),
426 v_flat,
427 midpoints,
428 - kv_cache.view(-1),
428 + kv_cache,
429 slot_mapping,
430 stride_cache_block=stride_block,
431 stride_cache_pos=stride_pos,