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 4, 17:30 UTC
Support only half types for concat_mla_q kernel (#37892)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Apr 23, 2026 at 23:51 UTC
cf8a613a87264183058801309868722f9013e101
1 file changed
+4
-1
csrc/cache_kernels.cu
+4
-1
@@ -1490,6 +1490,9 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
1490
TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2");
1491
TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
1492
TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
1493
+ TORCH_CHECK(ql_nope.scalar_type() == at::ScalarType::Half ||
1494
+ ql_nope.scalar_type() == at::ScalarType::BFloat16,
1495
+ "ql_nope must be float16 or bfloat16 dtype");
1496
1497
if (num_tokens == 0) return;
1498
@@ -1501,7 +1504,7 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
1504
const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope));
1505
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1506
1504
- VLLM_DISPATCH_FLOATING_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
1507
+ VLLM_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
1508
vllm::ConcatMLAQKernel<scalar_t, 512><<<grid_size, block_size, 0, stream>>>(
1509
q_out.data_ptr<scalar_t>(), ql_nope.data_ptr<scalar_t>(),
1510
q_pe.data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),