A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 21:45 UTC
Add 320 dimension size support to MLA (#36161)
Signed-off-by: Julien Denize <julien.denize@mistral.ai>
Julien Denize committed
Mar 11, 2026 at 18:21 UTC
a5d06dc557f9b04685e10793d3182358a47f7ba6
3 files changed
+25
-9
csrc/cache_kernels.cu
+19
-6
@@ -919,8 +919,8 @@ __global__ void gather_and_maybe_dequant_cache(
919
// SCALAR_T is the data type of the destination tensor.
920
// CACHE_T is the stored data type of kv-cache.
921
// KV_DTYPE is the real data type of kv-cache.
922
-#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE) \
923
- vllm::gather_and_maybe_dequant_cache<SCALAR_T, CACHE_T, KV_DTYPE, 576, \
922
+#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, ENTRY_SZ) \
923
+ vllm::gather_and_maybe_dequant_cache<SCALAR_T, CACHE_T, KV_DTYPE, ENTRY_SZ, \
924
thread_block_size> \
925
<<<grid, block, 0, stream>>>( \
926
reinterpret_cast<CACHE_T*>(src_cache.data_ptr()), \
@@ -931,6 +931,12 @@ __global__ void gather_and_maybe_dequant_cache(
931
dst_entry_stride, reinterpret_cast<const float*>(scale.data_ptr()), \
932
seq_starts_ptr);
933
934
+#define CALL_GATHER_CACHE_576(SCALAR_T, CACHE_T, KV_DTYPE) \
935
+ CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 576)
936
+
937
+#define CALL_GATHER_CACHE_320(SCALAR_T, CACHE_T, KV_DTYPE) \
938
+ CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 320)
939
+
940
// Gather sequences from the cache into the destination tensor.
941
// - cu_seq_lens contains the cumulative sequence lengths for each batch
942
// - block_table contains the cache block indices for each sequence
@@ -960,9 +966,10 @@ void gather_and_maybe_dequant_cache(
966
TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32,
967
"seq_starts must be int32");
968
}
963
- TORCH_CHECK(head_dim == 576,
964
- "gather_and_maybe_dequant_cache only support the head_dim to 576 "
965
- "for better performance")
969
+ TORCH_CHECK(
970
+ head_dim == 320 || head_dim == 576,
971
+ "gather_and_maybe_dequant_cache only support the head_dim to 320 or 576 "
972
+ "for better performance")
973
974
TORCH_CHECK(src_cache.device() == dst.device(),
975
"src_cache and dst must be on the same device");
@@ -987,7 +994,13 @@ void gather_and_maybe_dequant_cache(
994
const int32_t* seq_starts_ptr =
995
seq_starts.has_value() ? seq_starts.value().data_ptr<int32_t>() : nullptr;
996
990
- DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype, CALL_GATHER_CACHE);
997
+ if (head_dim == 576) {
998
+ DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype,
999
+ CALL_GATHER_CACHE_576);
1000
+ } else {
1001
+ DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype,
1002
+ CALL_GATHER_CACHE_320);
1003
+ }
1004
}
1005
1006
namespace vllm {
tests/kernels/attention/test_cache.py
+5
-2
@@ -23,7 +23,7 @@ CACHE_LAYOUTS = ["NHD", "HND"]
23
KV_SCALE_TYPES = ["tensor", "attn_head"]
24
25
# Parameters for MLA tests.
26
-KV_LORA_RANKS = [512]
26
+KV_LORA_RANKS = [256, 512]
27
QK_ROPE_HEAD_DIMS = [64]
28
NUM_TOKENS_MLA = [42]
29
BLOCK_SIZES_MLA = [16]
@@ -627,6 +627,8 @@ def test_concat_and_cache_ds_mla(
627
pytest.skip("concat_and_cache_mla doesn't support fp8_ds_mla on ROCm")
628
if dtype.itemsize != 2:
629
pytest.skip("ds_mla only supports 16-bit input")
630
+ if kv_lora_rank != 512:
631
+ pytest.skip("fp8_ds_mla requires kv_lora_rank == 512")
632
kv_cache_dtype = "fp8_ds_mla"
633
set_random_seed(seed)
634
torch.set_default_device(device)
@@ -663,7 +665,8 @@ def test_concat_and_cache_ds_mla(
665
ref_cache_32bit = ref_cache_slice.view(torch.float32)
666
667
kv_c_data = kv_c[i]
666
- for tile_idx in range(4):
668
+ num_tiles = kv_lora_rank // 128
669
+ for tile_idx in range(num_tiles):
670
tile_start = tile_idx * 128
671
tile_end = (tile_idx + 1) * 128
672
tile_data[:] = kv_c_data[tile_start:tile_end]
vllm/model_executor/layers/attention/mla_attention.py
+1
-1
@@ -1148,7 +1148,7 @@ class MLACommonBackend(AttentionBackend):
1148
1149
@classmethod
1150
def get_supported_head_sizes(cls) -> list[int]:
1151
- return [576]
1151
+ return [320, 576]
1152
1153
@classmethod
1154
def is_mla(cls) -> bool: