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 7, 00:15 UTC
[Attention][Perf][Kernel] Replace torch.cat with vectorized CUDA kernel MLA query concat - DeepSeek-V3.2 (#34917)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>
Roberto L. Castro committed
Mar 9, 2026 at 17:50 UTC
580864d81eb03d9fb1383e1782636ff6a9425fa2
10 files changed
+415
-15
.buildkite/test_areas/kernels.yaml
+2
-1
@@ -8,8 +8,9 @@ steps:
8
- csrc/
9
- tests/kernels/core
10
- tests/kernels/test_top_k_per_row.py
11
+ - tests/kernels/test_concat_mla_q.py
12
commands:
12
- - pytest -v -s kernels/core kernels/test_top_k_per_row.py
13
+ - pytest -v -s kernels/core kernels/test_top_k_per_row.py kernels/test_concat_mla_q.py
14
15
- label: Kernels Attention Test %N
16
timeout_in_minutes: 35
benchmarks/kernels/bench_concat_mla_q.py
new
+98
@@ -0,0 +1,98 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import argparse
5
+
6
+import torch
7
+
8
+from vllm import _custom_ops as ops
9
+from vllm.triton_utils import triton
10
+
11
+# DeepSeek V3 dimensions
12
+NOPE_DIM = 512
13
+ROPE_DIM = 64
14
+NUM_HEADS = 128
15
+
16
+NUM_TOKENS = [8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
17
+
18
+
19
+def get_configs():
20
+ return NUM_TOKENS
21
+
22
+
23
+def make_inputs(num_tokens, dtype):
24
+ """Create inputs matching the real code path.
25
+
26
+ Args:
27
+ contiguous_nope: If False, simulate the transposed BMM output
28
+ (non-contiguous nope with stride pattern from
29
+ [N,B,L].transpose(0,1)).
30
+ """
31
+ # Simulate: bmm output [N, B, L].transpose(0, 1) -> [B, N, L]
32
+ raw = torch.randn(NUM_HEADS, num_tokens, NOPE_DIM, dtype=dtype, device="cuda")
33
+ ql_nope = raw.transpose(0, 1)
34
+
35
+ q_pe = torch.randn(num_tokens, NUM_HEADS, ROPE_DIM, dtype=dtype, device="cuda")
36
+ return ql_nope, q_pe
37
+
38
+
39
+# ---- Non-contiguous nope benchmark (real code path) ----
40
+@triton.testing.perf_report(
41
+ triton.testing.Benchmark(
42
+ x_names=["num_tokens"],
43
+ x_vals=get_configs(),
44
+ line_arg="provider",
45
+ line_vals=["torch_cat", "concat_mla_q"],
46
+ line_names=["torch.cat", "concat_mla_q (v8)"],
47
+ styles=[("blue", "--"), ("green", "-")],
48
+ ylabel="Latency (us)",
49
+ plot_name="concat_mla_q-transposed",
50
+ args={},
51
+ )
52
+)
53
+def bench_transposed(num_tokens, provider):
54
+ dtype = torch.bfloat16
55
+ ql_nope, q_pe = make_inputs(num_tokens, dtype)
56
+
57
+ q_out = torch.empty(
58
+ num_tokens, NUM_HEADS, NOPE_DIM + ROPE_DIM, dtype=dtype, device="cuda"
59
+ )
60
+
61
+ quantiles = [0.5, 0.2, 0.8]
62
+
63
+ if provider == "torch_cat":
64
+ ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
65
+ lambda: torch.cat((ql_nope, q_pe), dim=-1), quantiles=quantiles, rep=500
66
+ )
67
+ else:
68
+ ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
69
+ lambda: ops.concat_mla_q(ql_nope, q_pe, q_out), quantiles=quantiles, rep=500
70
+ )
71
+
72
+ return ms * 1000, max_ms * 1000, min_ms * 1000 # us
73
+
74
+
75
+if __name__ == "__main__":
76
+ parser = argparse.ArgumentParser(description="Benchmark concat_mla_q vs torch.cat")
77
+ parser.add_argument(
78
+ "--save-path", type=str, default=None, help="Path to save benchmark results"
79
+ )
80
+ args = parser.parse_args()
81
+
82
+ print("\n" + "=" * 70)
83
+ print("CONCAT MLA Q KERNEL BENCHMARKS")
84
+ print("=" * 70)
85
+ print(f"Dimensions: nope={NOPE_DIM}, rope={ROPE_DIM}, heads={NUM_HEADS}")
86
+ print(
87
+ f"Per-head output: {NOPE_DIM + ROPE_DIM} bf16 = "
88
+ f"{(NOPE_DIM + ROPE_DIM) * 2} bytes"
89
+ )
90
+ print(f"num_tokens (decode=batch_size, prefill=chunk_size): {NUM_TOKENS}")
91
+ print("=" * 70)
92
+
93
+ print("\n--- Non-contiguous nope inputs (transposed BMM output) ---")
94
+ bench_transposed.run(print_data=True, save_path=args.save_path)
95
+
96
+ print("\n" + "=" * 70)
97
+ print("Benchmarking complete!")
98
+ print("=" * 70)
csrc/cache.h
+6
@@ -74,6 +74,12 @@ void indexer_k_quant_and_cache(
74
int64_t quant_block_size, // quantization block size
75
const std::string& scale_fmt);
76
77
+// Concatenate query nope and rope for MLA/DSA attention
78
+void concat_mla_q(
79
+ torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
80
+ torch::Tensor& q_pe, // [num_tokens, num_heads, rope_dim]
81
+ torch::Tensor& q_out); // [num_tokens, num_heads, nope_dim + rope_dim]
82
+
83
// Extract function to gather quantized K cache
84
void cp_gather_indexer_k_quant_cache(
85
const torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride]
csrc/cache_kernels.cu
+41
@@ -8,6 +8,7 @@
8
#include "cuda_compat.h"
9
#include "dispatch_utils.h"
10
#include "quantization/vectorization_utils.cuh"
11
+#include "concat_mla_q.cuh"
12
13
#ifdef USE_ROCM
14
#include "quantization/w8a8/fp8/amd/quant_utils.cuh"
@@ -1358,3 +1359,43 @@ void cp_gather_indexer_k_quant_cache(
1359
CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(32);
1360
}
1361
}
1362
+
1363
+// Concatenate ql_nope and q_pe into a contiguous q_out tensor for MLA/DSA.
1364
+// Replaces torch.cat((ql_nope, q_pe), dim=-1).
1365
+void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
1366
+ torch::Tensor& q_pe, // [num_tokens, num_heads, rope_dim]
1367
+ torch::Tensor& q_out // [num_tokens, num_heads, nope_dim +
1368
+ // rope_dim]
1369
+) {
1370
+ const int num_tokens = ql_nope.size(0);
1371
+ const int num_heads = ql_nope.size(1);
1372
+ const int nope_dim = ql_nope.size(2);
1373
+ const int rope_dim = q_pe.size(2);
1374
+
1375
+ TORCH_CHECK(nope_dim % 512 == 0, "nope_dim must be a multiple of 512, got ",
1376
+ nope_dim);
1377
+ TORCH_CHECK(rope_dim == 64, "rope_dim must be 64, got ", rope_dim);
1378
+ TORCH_CHECK(q_out.size(2) == nope_dim + rope_dim);
1379
+
1380
+ TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2");
1381
+ TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
1382
+ TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
1383
+
1384
+ if (num_tokens == 0) return;
1385
+
1386
+ constexpr int warps_per_block = 8;
1387
+ const int total_warps = num_tokens * num_heads;
1388
+ const int grid_size = (total_warps + warps_per_block - 1) / warps_per_block;
1389
+ const int block_size = warps_per_block * 32;
1390
+
1391
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope));
1392
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1393
+
1394
+ VLLM_DISPATCH_FLOATING_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
1395
+ vllm::ConcatMLAQKernel<scalar_t, 512><<<grid_size, block_size, 0, stream>>>(
1396
+ q_out.data_ptr<scalar_t>(), ql_nope.data_ptr<scalar_t>(),
1397
+ q_pe.data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),
1398
+ q_out.stride(1), ql_nope.stride(0), ql_nope.stride(1), q_pe.stride(0),
1399
+ q_pe.stride(1));
1400
+ });
1401
+}
csrc/concat_mla_q.cuh
new
+60
@@ -0,0 +1,60 @@
1
+#ifndef CONCAT_MLA_Q_CUH_
2
+#define CONCAT_MLA_Q_CUH_
3
+
4
+#include <cuda_bf16.h>
5
+#include <cuda_fp16.h>
6
+
7
+#include "cuda_vec_utils.cuh"
8
+
9
+namespace vllm {
10
+
11
+// Concatenates ql_nope [num_tokens, num_heads, NOPE_DIM] and
12
+// q_pe [num_tokens, num_heads, 64]
13
+// into q_out [num_tokens, num_heads, NOPE_DIM+64].
14
+// Currently instantiated only for NOPE_DIM=512.
15
+// Rope dim is hardcoded to 64 (DeepSeek V3.2 MLA)
16
+template <typename DType, int NOPE_DIM>
17
+__global__ void ConcatMLAQKernel(
18
+ DType* __restrict__ q_out, const DType* __restrict__ ql_nope,
19
+ const DType* __restrict__ q_pe, const int num_tokens, const int num_heads,
20
+ const int64_t out_stride_0, const int64_t out_stride_1,
21
+ const int64_t nope_stride_0, const int64_t nope_stride_1,
22
+ const int64_t pe_stride_0, const int64_t pe_stride_1) {
23
+ const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
24
+ if (flat_warp_id >= num_tokens * num_heads) return;
25
+
26
+ const int token_id = flat_warp_id / num_heads;
27
+ const int head_id = flat_warp_id % num_heads;
28
+ const int lane_id = threadIdx.x & 31;
29
+
30
+ constexpr bool use_256b = VLLM_256B_PTX_ENABLED;
31
+ constexpr int nope_vec_loads =
32
+ NOPE_DIM * sizeof(DType) / (VecTraits<use_256b>::ARCH_MAX_VEC_SIZE * 32);
33
+
34
+ const DType* nope_src =
35
+ ql_nope + token_id * nope_stride_0 + head_id * nope_stride_1;
36
+ DType* nope_dst = q_out + token_id * out_stride_0 + head_id * out_stride_1;
37
+
38
+#pragma unroll
39
+ for (int i = 0; i < nope_vec_loads; i++) {
40
+ const int offset = i * 32 + lane_id;
41
+ if constexpr (use_256b) {
42
+ st256_cs(reinterpret_cast<u32x8_t*>(nope_dst) + offset,
43
+ ld256_cs(reinterpret_cast<const u32x8_t*>(nope_src) + offset));
44
+ } else {
45
+ st128_cs(reinterpret_cast<int4*>(nope_dst) + offset,
46
+ ld128_cs(reinterpret_cast<const int4*>(nope_src) + offset));
47
+ }
48
+ }
49
+
50
+ const int* rope_src = reinterpret_cast<const int*>(
51
+ q_pe + token_id * pe_stride_0 + head_id * pe_stride_1);
52
+ int* rope_dst = reinterpret_cast<int*>(q_out + token_id * out_stride_0 +
53
+ head_id * out_stride_1 + NOPE_DIM);
54
+
55
+ st32_cs(rope_dst + lane_id, ld32_cs(rope_src + lane_id));
56
+}
57
+
58
+} // namespace vllm
59
+
60
+#endif // CONCAT_MLA_Q_CUH_
csrc/cuda_vec_utils.cuh
+37
-10
@@ -196,7 +196,6 @@ __forceinline__ __device__ u32x8_t ld256_cs(const u32x8_t* addr) {
196
return val;
197
#else
198
assert(false && "ld256_cs requires SM100+ with CUDA 12.9+");
199
- return {};
199
#endif
200
}
201
@@ -211,23 +210,51 @@ __forceinline__ __device__ void st256_cs(u32x8_t* addr, u32x8_t val) {
210
#endif
211
}
212
214
-// 32-bit cache-streaming (.cs) load / store — SM100+ only.
213
+// 32-bit load / store.
214
+__device__ __forceinline__ int ld32(const int* addr) { return __ldg(addr); }
215
+
216
+__device__ __forceinline__ void st32(int* addr, int val) { *addr = val; }
217
+
218
+// 32-bit cache-streaming (.cs) load / store.
219
+// Falls back to ld32/st32 on ROCm (no .cs hint).
220
__forceinline__ __device__ int ld32_cs(const int* addr) {
216
-#if VLLM_256B_PTX_ENABLED
221
int val;
222
+#ifndef USE_ROCM
223
asm volatile("ld.global.cs.b32 %0, [%1];" : "=r"(val) : "l"(addr));
219
- return val;
224
#else
221
- assert(false && "ld32_cs requires SM100+ with CUDA 12.9+");
222
- return 0;
225
+ val = ld32(addr);
226
#endif
227
+ return val;
228
}
229
230
__forceinline__ __device__ void st32_cs(int* addr, int val) {
227
-#if VLLM_256B_PTX_ENABLED
231
+#ifndef USE_ROCM
232
asm volatile("st.global.cs.b32 [%0], %1;" ::"l"(addr), "r"(val));
233
#else
230
- assert(false && "st32_cs requires SM100+ with CUDA 12.9+");
234
+ st32(addr, val);
235
+#endif
236
+}
237
+
238
+// 128-bit cache-streaming (.cs) load / store.
239
+// Falls back to ld128/st128 on ROCm (no .cs hint).
240
+__forceinline__ __device__ int4 ld128_cs(const int4* addr) {
241
+ int4 val;
242
+#ifndef USE_ROCM
243
+ asm volatile("ld.global.cs.v4.u32 {%0,%1,%2,%3}, [%4];"
244
+ : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
245
+ : "l"(addr));
246
+#else
247
+ ld128(val, addr);
248
+#endif
249
+ return val;
250
+}
251
+
252
+__forceinline__ __device__ void st128_cs(int4* addr, int4 val) {
253
+#ifndef USE_ROCM
254
+ asm volatile("st.global.cs.v4.u32 [%0], {%1,%2,%3,%4};" ::"l"(addr),
255
+ "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w));
256
+#else
257
+ st128(val, addr);
258
#endif
259
}
260
@@ -260,7 +287,7 @@ __device__ __forceinline__ void ld256_cg_or_zero(u32x8_t& val, const void* ptr,
287
288
__device__ __forceinline__ void ld128_cg_or_zero(uint4& val, const void* ptr,
289
bool pred) {
263
-#if VLLM_256B_PTX_ENABLED
290
+#ifndef USE_ROCM
291
uint32_t r0, r1, r2, r3;
292
293
asm volatile(
@@ -278,7 +305,7 @@ __device__ __forceinline__ void ld128_cg_or_zero(uint4& val, const void* ptr,
305
306
val = uint4{r0, r1, r2, r3};
307
#else
281
- assert(false && "ld128_cg_or_zero requires SM100+ with CUDA 12.9+");
308
+ assert(false && "ld128_cg_or_zero is not supported on ROCm");
309
#endif
310
}
311
csrc/torch_bindings.cpp
+4
@@ -802,6 +802,10 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
802
cache_ops.impl("indexer_k_quant_and_cache", torch::kCUDA,
803
&indexer_k_quant_and_cache);
804
805
+ cache_ops.def(
806
+ "concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()");
807
+ cache_ops.impl("concat_mla_q", torch::kCUDA, &concat_mla_q);
808
+
809
cache_ops.def(
810
"cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
811
"dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
tests/kernels/test_concat_mla_q.py
new
+139
@@ -0,0 +1,139 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm import _custom_ops as ops
8
+
9
+NUM_TOKENS = [1, 4, 16, 64, 128]
10
+NUM_HEADS = [128]
11
+NOPE_DIM = [512]
12
+ROPE_DIM = [64]
13
+DTYPES = [torch.bfloat16, torch.float16]
14
+
15
+
16
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
17
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
18
+@pytest.mark.parametrize("nope_dim", NOPE_DIM)
19
+@pytest.mark.parametrize("rope_dim", ROPE_DIM)
20
+@pytest.mark.parametrize("dtype", DTYPES)
21
+def test_concat_mla_q_contiguous(num_tokens, num_heads, nope_dim, rope_dim, dtype):
22
+ """Test with contiguous inputs (standard layout)."""
23
+ torch.manual_seed(42)
24
+ ql_nope = torch.randn(num_tokens, num_heads, nope_dim, dtype=dtype, device="cuda")
25
+ q_pe = torch.randn(num_tokens, num_heads, rope_dim, dtype=dtype, device="cuda")
26
+
27
+ ref = torch.cat((ql_nope, q_pe), dim=-1)
28
+
29
+ q_out = torch.empty(
30
+ num_tokens, num_heads, nope_dim + rope_dim, dtype=dtype, device="cuda"
31
+ )
32
+ ops.concat_mla_q(ql_nope, q_pe, q_out)
33
+
34
+ torch.testing.assert_close(q_out, ref, atol=0, rtol=0)
35
+
36
+
37
+@pytest.mark.parametrize("num_tokens", [t for t in NUM_TOKENS if t > 1])
38
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
39
+@pytest.mark.parametrize("nope_dim", NOPE_DIM)
40
+@pytest.mark.parametrize("rope_dim", ROPE_DIM)
41
+@pytest.mark.parametrize("dtype", DTYPES)
42
+def test_concat_mla_q_transposed_nope(num_tokens, num_heads, nope_dim, rope_dim, dtype):
43
+ """Test with transposed nope input (simulates BMM output after transpose).
44
+
45
+ In the real code path, mqa_ql_nope is the result of:
46
+ torch.bmm(q_nope, W_UK_T) # [N, B, L]
47
+ .transpose(0, 1) # [B, N, L] — non-contiguous!
48
+ """
49
+ torch.manual_seed(42)
50
+ nope_raw = torch.randn(num_heads, num_tokens, nope_dim, dtype=dtype, device="cuda")
51
+ ql_nope = nope_raw.transpose(0, 1) # [B, N, L], non-contiguous
52
+ assert not ql_nope.is_contiguous()
53
+
54
+ q_pe = torch.randn(num_tokens, num_heads, rope_dim, dtype=dtype, device="cuda")
55
+
56
+ ref = torch.cat((ql_nope, q_pe), dim=-1)
57
+
58
+ q_out = torch.empty(
59
+ num_tokens, num_heads, nope_dim + rope_dim, dtype=dtype, device="cuda"
60
+ )
61
+ ops.concat_mla_q(ql_nope, q_pe, q_out)
62
+
63
+ torch.testing.assert_close(q_out, ref, atol=0, rtol=0)
64
+
65
+
66
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
67
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
68
+@pytest.mark.parametrize("dtype", DTYPES)
69
+def test_concat_mla_q_split_rope(num_tokens, num_heads, dtype):
70
+ """Test with rope from a split (simulates the actual code path).
71
+
72
+ In the real code path, q_pe comes from:
73
+ mqa_q.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
74
+ which creates a non-contiguous view with stride(1) != rope_dim.
75
+ """
76
+ torch.manual_seed(42)
77
+ nope_dim = 512
78
+ rope_dim = 64
79
+ orig_dim = 128 + 64 # original q before absorption: [B, N, 192]
80
+
81
+ # Simulate split from original q tensor
82
+ q_orig = torch.randn(num_tokens, num_heads, orig_dim, dtype=dtype, device="cuda")
83
+ q_nope_orig, q_pe = q_orig.split([128, 64], dim=-1)
84
+
85
+ # q_pe is non-contiguous: stride(1) = 192, not 64
86
+ assert q_pe.stride(1) == orig_dim
87
+ assert q_pe.stride(2) == 1 # but innermost is fine
88
+
89
+ # Simulate absorbed nope (contiguous, different size)
90
+ ql_nope = torch.randn(num_tokens, num_heads, nope_dim, dtype=dtype, device="cuda")
91
+
92
+ ref = torch.cat((ql_nope, q_pe), dim=-1)
93
+
94
+ q_out = torch.empty(
95
+ num_tokens, num_heads, nope_dim + rope_dim, dtype=dtype, device="cuda"
96
+ )
97
+ ops.concat_mla_q(ql_nope, q_pe, q_out)
98
+
99
+ torch.testing.assert_close(q_out, ref, atol=0, rtol=0)
100
+
101
+
102
+def test_concat_mla_q_zero_tokens():
103
+ """Test with zero tokens (edge case)."""
104
+ ql_nope = torch.empty(0, 128, 512, dtype=torch.bfloat16, device="cuda")
105
+ q_pe = torch.empty(0, 128, 64, dtype=torch.bfloat16, device="cuda")
106
+ q_out = torch.empty(0, 128, 576, dtype=torch.bfloat16, device="cuda")
107
+
108
+ ops.concat_mla_q(ql_nope, q_pe, q_out)
109
+
110
+
111
+@pytest.mark.parametrize("num_tokens", [1, 64])
112
+def test_concat_mla_q_values_preserved(num_tokens):
113
+ """Verify exact bit-level preservation (no computation, pure copy).
114
+
115
+ Compares raw int16 bits to avoid NaN != NaN issues from IEEE 754.
116
+ """
117
+ nope_dim, rope_dim = 512, 64
118
+
119
+ # Use specific bit patterns (stay in int16 for bit-exact comparison)
120
+ ql_nope_bits = torch.arange(
121
+ num_tokens * 128 * nope_dim, dtype=torch.int16, device="cuda"
122
+ ).view(num_tokens, 128, nope_dim)
123
+ q_pe_bits = torch.arange(
124
+ num_tokens * 128 * rope_dim, dtype=torch.int16, device="cuda"
125
+ ).view(num_tokens, 128, rope_dim)
126
+
127
+ ql_nope = ql_nope_bits.view(torch.bfloat16)
128
+ q_pe = q_pe_bits.view(torch.bfloat16)
129
+
130
+ q_out = torch.empty(
131
+ num_tokens, 128, nope_dim + rope_dim, dtype=torch.bfloat16, device="cuda"
132
+ )
133
+ ops.concat_mla_q(ql_nope, q_pe, q_out)
134
+
135
+ out_bits = q_out.view(torch.int16)
136
+
137
+ assert torch.equal(out_bits[..., :nope_dim], ql_nope_bits)
138
+
139
+ assert torch.equal(out_bits[..., nope_dim:], q_pe_bits)
vllm/_custom_ops.py
+15
@@ -2672,6 +2672,21 @@ def cp_gather_and_upconvert_fp8_kv_cache(
2672
)
2673
2674
2675
+def concat_mla_q(
2676
+ ql_nope: torch.Tensor,
2677
+ q_pe: torch.Tensor,
2678
+ q_out: torch.Tensor,
2679
+) -> None:
2680
+ """Concatenate query nope and rope for MLA/DSA attention.
2681
+
2682
+ Args:
2683
+ ql_nope: Query nope component [num_tokens, num_heads, nope_dim]
2684
+ q_pe: Query rope component [num_tokens, num_heads, rope_dim]
2685
+ q_out: Output tensor [num_tokens, num_heads, nope_dim + rope_dim]
2686
+ """
2687
+ torch.ops._C_cache_ops.concat_mla_q(ql_nope, q_pe, q_out)
2688
+
2689
+
2690
def indexer_k_quant_and_cache(
2691
k: torch.Tensor,
2692
kv_cache: torch.Tensor,
vllm/v1/attention/backends/mla/flashmla_sparse.py
+13
-4
@@ -568,6 +568,9 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
568
)
569
self.fp8_decode_padded_heads = self._compute_fp8_decode_padded_heads(num_heads)
570
571
+ vllm_config = get_current_vllm_config()
572
+ max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
573
+ q_concat_shape = (max_tokens, num_heads, head_size)
574
if kv_cache_dtype.startswith("fp8"):
575
assert kv_cache_dtype == "fp8_ds_mla", (
576
"FlashMLA Sparse Attention backend fp8 only supports "
@@ -576,17 +579,21 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
579
580
if kv_cache_dtype == "fp8_ds_mla":
581
# Reserve workspace during initialization
579
- vllm_config = get_current_vllm_config()
582
assert vllm_config is not None and vllm_config.model_config is not None
583
prefill_workspace_size = get_prefill_workspace_size(
584
vllm_config.model_config.max_model_len
585
)
586
self.prefill_workspace_shape = (prefill_workspace_size, head_size)
585
- (self.prefill_bf16_workspace,) = (
587
+ self.q_concat_buffer, self.prefill_bf16_workspace = (
588
current_workspace_manager().get_simultaneous(
587
- (self.prefill_workspace_shape, torch.bfloat16)
589
+ (q_concat_shape, torch.bfloat16),
590
+ (self.prefill_workspace_shape, torch.bfloat16),
591
)
592
)
593
+ else:
594
+ (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous(
595
+ (q_concat_shape, torch.bfloat16),
596
+ )
597
598
def _forward_bf16_kv(
599
self,
@@ -828,7 +835,9 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
835
836
# Concatenate q if it's a tuple (ql_nope, q_pe)
837
if isinstance(q, tuple):
831
- q = torch.cat(q, dim=-1)
838
+ ql_nope, q_pe = q
839
+ q = self.q_concat_buffer[: ql_nope.shape[0]]
840
+ ops.concat_mla_q(ql_nope, q_pe, q)
841
842
num_actual_toks = q.shape[0]
843