@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cu 2,056 lines 89.8 KB
Raw
1 #include "torch_utils.h"
2 #include "dispatch_utils.h"
3
4 #include "../cuda_utils.h"
5 #include "../cuda_compat.h"
6
7 #include "quantization/vectorization_utils.cuh"
8 #include "concat_mla_q.cuh"
9
10 #if !defined(USE_ROCM) && defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
11 #include "nvfp4_ds_mla_cache.h"
12 #endif
13
14 #ifdef USE_ROCM
15 #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
16 #else
17 #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
18 #endif
19
20 #include <algorithm>
21 #include <cassert>
22 #include <cstdlib>
23
24 #ifdef USE_ROCM
25 #include <hip/hip_bf16.h>
26 typedef __hip_bfloat16 __nv_bfloat16;
27 #else
28 #include <cuda.h>
29 #endif
30
31 #if defined(__gfx942__)
32 constexpr float kFp8ScaleDivisor = 224.f;
33 #else
34 constexpr float kFp8ScaleDivisor = 448.f;
35 #endif
36
37 void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
38 int64_t block_size_in_bytes,
39 const torch::stable::Tensor& block_mapping) {
40 torch::stable::Device src_device = src.device();
41 torch::stable::Device dst_device = dst.device();
42 cudaMemcpyKind memcpy_type;
43 if (src_device.is_cuda() && dst_device.is_cuda()) {
44 STD_TORCH_CHECK(src_device.index() == dst_device.index(),
45 "src and dst must be on the same GPU");
46 memcpy_type = cudaMemcpyDeviceToDevice;
47 } else if (src_device.is_cuda() && dst_device.is_cpu()) {
48 memcpy_type = cudaMemcpyDeviceToHost;
49 } else if (src_device.is_cpu() && dst_device.is_cuda()) {
50 memcpy_type = cudaMemcpyHostToDevice;
51 } else {
52 STD_TORCH_CHECK(false, "Invalid device combination");
53 }
54
55 // NOTE(youkaichao): keep in mind that `block_mapping` should be
56 // a cpu tensor, otherwise every `item` call will require a gpu-cpu
57 // synchronization.
58 STD_TORCH_CHECK(block_mapping.device().is_cpu(),
59 "block_mapping must be on CPU");
60
61 char* src_ptr = static_cast<char*>(src.data_ptr());
62 char* dst_ptr = static_cast<char*>(dst.data_ptr());
63
64 auto guard_device = src_device.is_cuda() ? src_device : dst_device;
65 const torch::stable::accelerator::DeviceGuard device_guard(
66 guard_device.index());
67 const cudaStream_t stream = get_current_cuda_stream();
68 // NOTE(woosuk): This can be slow if the number of blocks is large.
69 const int64_t num_blocks = block_mapping.size(0);
70 const int64_t* bm_ptr = block_mapping.const_data_ptr<int64_t>();
71 const int64_t bm_stride0 = block_mapping.stride(0);
72 const int64_t bm_stride1 = block_mapping.stride(1);
73 for (size_t i = 0; i < num_blocks; i++) {
74 int64_t src_block_number = bm_ptr[i * bm_stride0];
75 int64_t dst_block_number = bm_ptr[i * bm_stride0 + bm_stride1];
76 int64_t src_offset = src_block_number * block_size_in_bytes;
77 int64_t dst_offset = dst_block_number * block_size_in_bytes;
78 cudaMemcpyAsync(dst_ptr + dst_offset, src_ptr + src_offset,
79 block_size_in_bytes, memcpy_type, stream);
80 }
81 }
82
83 namespace {
84 // ROCm hipMemcpyBatchAsync faults for count > 8192 (MI350X/gfx950), so chunk
85 // at that ceiling.
86 #if defined(USE_ROCM)
87 constexpr int64_t kRocmDefaultMaxBatchDescriptors = 8192;
88 #endif
89
90 constexpr const char* kMaxBatchDescriptorsEnv =
91 "VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS";
92
93 // Max copy descriptors per batch-memcpy call (0 = unlimited). ROCm caps and
94 // chunks; CUDA is uncapped. The env var (>0) overrides on any platform.
95 int64_t resolve_max_batch_descriptors() {
96 static const int64_t cached = []() -> int64_t {
97 const char* val = std::getenv(kMaxBatchDescriptorsEnv);
98 const int64_t override_val = val ? std::atoll(val) : 0;
99 if (override_val > 0) return override_val;
100 #if defined(USE_ROCM)
101 return kRocmDefaultMaxBatchDescriptors;
102 #else
103 return 0;
104 #endif
105 }();
106 return cached;
107 }
108 } // namespace
109
110 void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
111 const torch::stable::Tensor& dst_ptrs,
112 const torch::stable::Tensor& sizes,
113 bool is_src_access_order_any) {
114 STD_TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU");
115 STD_TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU");
116 STD_TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU");
117 STD_TORCH_CHECK(src_ptrs.scalar_type() == torch::headeronly::ScalarType::Long,
118 "src_ptrs must be int64");
119 STD_TORCH_CHECK(dst_ptrs.scalar_type() == torch::headeronly::ScalarType::Long,
120 "dst_ptrs must be int64");
121 STD_TORCH_CHECK(sizes.scalar_type() == torch::headeronly::ScalarType::Long,
122 "sizes must be int64");
123
124 const int64_t n = src_ptrs.size(0);
125 STD_TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs");
126 STD_TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs");
127
128 if (n == 0) return;
129
130 int64_t* src_data = src_ptrs.mutable_data_ptr<int64_t>();
131 int64_t* dst_data = dst_ptrs.mutable_data_ptr<int64_t>();
132 int64_t* size_data = sizes.mutable_data_ptr<int64_t>();
133
134 const cudaStream_t stream = get_current_cuda_stream();
135
136 // Use cuMemcpyBatchAsync / hipMemcpyBatchAsync to submit all copies in a
137 // single driver call, amortizing per-copy submission overhead. int64_t
138 // and CUdeviceptr/void*/size_t are all 8 bytes on 64-bit platforms, so we
139 // reinterpret_cast the tensor data directly to avoid copies.
140 static_assert(sizeof(size_t) == sizeof(int64_t));
141 #if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080
142 static_assert(sizeof(CUdeviceptr) == sizeof(int64_t));
143 // Resolve cuMemcpyBatchAsync at runtime via cuGetProcAddress so that
144 // binaries compiled with CUDA 12.8+ still work on older drivers, and
145 // we avoid the CUDA 13.0 header remapping (#define to _v2 signature).
146 // The function pointer is cached after the first call.
147 using BatchFn =
148 CUresult (*)(CUdeviceptr*, CUdeviceptr*, size_t*, size_t,
149 CUmemcpyAttributes*, size_t*, size_t, size_t*, CUstream);
150 static BatchFn batch_fn = []() -> BatchFn {
151 CUdriverProcAddressQueryResult sym_status;
152 void* fn_ptr = nullptr;
153 CUresult res = cuGetProcAddress("cuMemcpyBatchAsync", &fn_ptr, 12080,
154 CU_GET_PROC_ADDRESS_DEFAULT, &sym_status);
155 if (res != CUDA_SUCCESS || fn_ptr == nullptr) {
156 return nullptr;
157 }
158 return reinterpret_cast<BatchFn>(fn_ptr);
159 }();
160
161 // cuMemcpyBatchAsync rejects the legacy default stream (handle 0 /
162 // cudaStreamLegacy) with CUDA_ERROR_INVALID_VALUE; route it to the per-copy
163 // fallback below, which is correct on any stream. Real and per-thread-default
164 // streams take the batch fast path.
165 const bool usable_stream = stream != nullptr && stream != cudaStreamLegacy;
166 if (batch_fn != nullptr && usable_stream) {
167 CUmemcpyAttributes attr = {};
168 // ANY lets the DMA engine prefetch source bytes out of stream order,
169 // which is only safe when no GPU stream is concurrently writing the
170 // source.
171 attr.srcAccessOrder = is_src_access_order_any
172 ? CU_MEMCPY_SRC_ACCESS_ORDER_ANY
173 : CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
174 size_t attrs_idx = 0;
175 size_t fail_idx = 0;
176 // Uncapped on CUDA (max_desc == 0 -> single call) unless overridden by
177 // VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS; chunk to honor the override.
178 const int64_t max_desc = resolve_max_batch_descriptors();
179 const int64_t step = max_desc <= 0 ? n : max_desc;
180 for (int64_t off = 0; off < n; off += step) {
181 const int64_t cnt = std::min(step, n - off);
182 CUresult result = batch_fn(reinterpret_cast<CUdeviceptr*>(dst_data + off),
183 reinterpret_cast<CUdeviceptr*>(src_data + off),
184 reinterpret_cast<size_t*>(size_data + off),
185 static_cast<size_t>(cnt), &attr, &attrs_idx, 1,
186 &fail_idx, static_cast<CUstream>(stream));
187 STD_TORCH_CHECK(result == CUDA_SUCCESS,
188 "cuMemcpyBatchAsync failed at index ", fail_idx,
189 " with error ", result);
190 }
191 return;
192 }
193 #elif defined(USE_ROCM) && defined(HIP_VERSION) && HIP_VERSION >= 70100000
194 // ROCm 7.1+ exposes hipMemcpyBatchAsync. ROCm 7.2.1-7.2.3 early-return
195 // hipErrorNotSupported whenever numAttrs > 0 (see ROCm/clr @ rocm-7.2.1
196 // hipamd/src/hip_memory.cpp:2819-2822), so those releases must call with
197 // numAttrs=0. ROCm 7.13+ accepts numAttrs > 0.
198 // rocm-7.14+ has better performance.
199 {
200 hipMemcpyAttributes attr = {};
201 size_t attrs_idx = 0;
202 size_t fail_idx = 0;
203 size_t num_attrs = 0;
204 #if HIP_VERSION >= 71300000
205 static const bool runtime_accepts_attrs = []() {
206 int runtime_version = 0;
207 return hipRuntimeGetVersion(&runtime_version) == hipSuccess &&
208 runtime_version >= 71300000;
209 }();
210 if (runtime_accepts_attrs) {
211 attr.srcAccessOrder = is_src_access_order_any
212 ? hipMemcpySrcAccessOrderAny
213 : hipMemcpySrcAccessOrderStream;
214 num_attrs = 1;
215 }
216 #endif
217 // hipMemcpyBatchAsync faults (GPU memory access fault) above 8192
218 // descriptors/call on ROCm 7.15, so chunk the batch at the resolved cap.
219 const int64_t max_desc = resolve_max_batch_descriptors();
220 const int64_t step = max_desc <= 0 ? n : max_desc;
221 for (int64_t off = 0; off < n; off += step) {
222 const int64_t cnt = std::min(step, n - off);
223 hipError_t result = hipMemcpyBatchAsync(
224 reinterpret_cast<void**>(dst_data + off),
225 reinterpret_cast<void**>(src_data + off),
226 reinterpret_cast<size_t*>(size_data + off), static_cast<size_t>(cnt),
227 &attr, &attrs_idx, num_attrs, &fail_idx,
228 static_cast<hipStream_t>(stream));
229 STD_TORCH_CHECK(result == hipSuccess,
230 "hipMemcpyBatchAsync failed at index ", fail_idx,
231 " with error ", result);
232 }
233 return;
234 }
235 #endif
236 {
237 // Fallback for CUDA < 12.8, older CUDA drivers, and ROCm < 7.1:
238 // individual async copies. cudaMemcpyDefault lets the driver infer
239 // direction from pointer types.
240 for (int64_t i = 0; i < n; i++) {
241 cudaMemcpyAsync(reinterpret_cast<void*>(dst_data[i]),
242 reinterpret_cast<void*>(src_data[i]),
243 static_cast<size_t>(size_data[i]), cudaMemcpyDefault,
244 stream);
245 }
246 }
247 }
248
249 namespace vllm {
250
251 // Used to copy/convert one element
252 template <typename OutT, typename InT, Fp8KVCacheDataType kv_dt>
253 struct CopyWithScaleOp {
254 float scale;
255
256 __device__ __forceinline__ void operator()(OutT& dst, const InT src) const {
257 if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
258 dst = static_cast<OutT>(src);
259 } else {
260 dst = fp8::scaled_convert<OutT, InT, kv_dt>(src, scale);
261 }
262 }
263 };
264
265 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
266 __global__ void reshape_and_cache_kernel(
267 const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size]
268 const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size]
269 cache_t* __restrict__ key_cache, // [num_blocks, num_heads, head_size/x,
270 // block_size, x]
271 cache_t* __restrict__ value_cache, // [num_blocks, num_heads, head_size,
272 // block_size]
273 const int64_t* __restrict__ slot_mapping, // [num_tokens]
274 const int key_stride, const int value_stride, const int num_heads,
275 const int head_size, const int block_size, const int x,
276 const float* k_scale, const float* v_scale) {
277 const int64_t token_idx = blockIdx.x;
278 const int64_t slot_idx = slot_mapping[token_idx];
279 if (slot_idx < 0) {
280 return;
281 }
282
283 const int64_t block_idx = slot_idx / block_size;
284 const int64_t block_offset = slot_idx % block_size;
285 const int h_block_count = head_size / x; // head_size//x
286
287 const int h_block_idx = threadIdx.x;
288 if (h_block_idx >= num_heads * h_block_count) {
289 return;
290 }
291
292 const int head_idx = h_block_idx / h_block_count;
293 const int h_block = h_block_idx % h_block_count;
294
295 const scalar_t* __restrict__ key_src =
296 key + token_idx * key_stride + head_idx * head_size + h_block * x;
297 const int64_t src_value_start =
298 token_idx * value_stride + head_idx * head_size + h_block * x;
299
300 cache_t* __restrict__ key_dst =
301 key_cache + block_idx * num_heads * h_block_count * block_size * x +
302 head_idx * h_block_count * block_size * x + h_block * block_size * x +
303 block_offset * x;
304 const int64_t tgt_value_start =
305 block_idx * num_heads * h_block_count * x * block_size +
306 head_idx * h_block_count * x * block_size + h_block * x * block_size +
307 block_offset;
308
309 constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4;
310 float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale;
311 CopyWithScaleOp<cache_t, scalar_t, kv_dt> k_op{k_scale_val};
312 float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *v_scale;
313 CopyWithScaleOp<cache_t, scalar_t, kv_dt> v_op{v_scale_val};
314
315 vectorize_with_alignment<VEC_SIZE>(key_src, key_dst, x, 0, 1, k_op);
316
317 const scalar_t* __restrict__ value_src = value + src_value_start;
318 cache_t* __restrict__ value_dst = value_cache + tgt_value_start;
319 #pragma unroll
320 for (int i = 0; i < x; i++) {
321 v_op(value_dst[i * block_size], value_src[i]);
322 }
323 }
324
325 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
326 __global__ void reshape_and_cache_flash_kernel(
327 const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size]
328 const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size]
329 cache_t* __restrict__ key_cache, // NHD or HND, shape see comments below
330 cache_t* __restrict__ value_cache, // same above
331 const int64_t* __restrict__ slot_mapping, // [num_tokens]
332 const int64_t block_stride, const int64_t page_stride,
333 const int64_t head_stride, const int64_t key_stride,
334 const int64_t value_stride, const int num_heads, const int head_size,
335 const int block_size, const float* k_scale, const float* v_scale,
336 const int kv_scale_stride) {
337 const int64_t token_idx = blockIdx.x;
338 const int64_t slot_idx = slot_mapping[token_idx];
339 // NOTE: slot_idx can be -1 if the token is padded
340 if (slot_idx < 0) {
341 return;
342 }
343 const int64_t block_idx = slot_idx / block_size;
344 const int64_t block_offset = slot_idx % block_size;
345 const int n_elems = num_heads * head_size;
346
347 // pointers to the beginning of the source row for this token.
348 const scalar_t* __restrict__ key_src = key + token_idx * key_stride;
349 const scalar_t* __restrict__ value_src = value + token_idx * value_stride;
350
351 // find the start position inside the kv-cache for this token.
352 cache_t* __restrict__ key_dst =
353 key_cache + block_idx * block_stride + block_offset * page_stride;
354 cache_t* __restrict__ value_dst =
355 value_cache + block_idx * block_stride + block_offset * page_stride;
356
357 // this is true for the NHD layout where `head_stride == head_size`
358 const bool is_contiguous_heads = (head_stride == head_size);
359
360 constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4;
361
362 if (is_contiguous_heads && kv_scale_stride == 0) {
363 // NHD layout and k/v_scales are [1] (i.e. single scale for all heads)
364 // kv cache: [num_blocks, block_size, num_heads, head_size]
365 float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale;
366 float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *v_scale;
367
368 CopyWithScaleOp<cache_t, scalar_t, kv_dt> k_op{k_scale_val};
369 CopyWithScaleOp<cache_t, scalar_t, kv_dt> v_op{v_scale_val};
370
371 vectorize_with_alignment<VEC_SIZE>(key_src, key_dst, n_elems, threadIdx.x,
372 blockDim.x, k_op);
373 vectorize_with_alignment<VEC_SIZE>(value_src, value_dst, n_elems,
374 threadIdx.x, blockDim.x, v_op);
375 } else {
376 // HND layout OR k/v_scales are [num_heads] (i.e. per-attn-head)
377 // HND layout: heads are strided, but each head_size segment is contiguous
378 // kv cache: [num_blocks, num_heads, block_size, head_size]
379 const int lane = threadIdx.x & 31; // 0..31 within warp
380 const int warp_id = threadIdx.x >> 5; // warp index within block
381 const int warps_per_block = blockDim.x >> 5;
382
383 for (int head = warp_id; head < num_heads; head += warps_per_block) {
384 const scalar_t* __restrict__ k_src_h = key_src + head * head_size;
385 const scalar_t* __restrict__ v_src_h = value_src + head * head_size;
386
387 cache_t* __restrict__ k_dst_h =
388 key_dst + static_cast<int64_t>(head) * head_stride;
389 cache_t* __restrict__ v_dst_h =
390 value_dst + static_cast<int64_t>(head) * head_stride;
391
392 float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto)
393 ? 0.f
394 : k_scale[head * kv_scale_stride];
395 float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto)
396 ? 0.f
397 : v_scale[head * kv_scale_stride];
398
399 CopyWithScaleOp<cache_t, scalar_t, kv_dt> k_op{k_scale_val};
400 CopyWithScaleOp<cache_t, scalar_t, kv_dt> v_op{v_scale_val};
401
402 // within each head, let the 32 threads of the warp perform the vector
403 // copy
404 vectorize_with_alignment<VEC_SIZE>(k_src_h, k_dst_h, head_size, lane, 32,
405 k_op);
406
407 vectorize_with_alignment<VEC_SIZE>(v_src_h, v_dst_h, head_size, lane, 32,
408 v_op);
409 }
410 }
411 }
412
413 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
414 __global__ void concat_and_cache_mla_kernel(
415 const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank]
416 const scalar_t* __restrict__ k_pe, // [num_tokens, pe_dim]
417 cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank
418 // + pe_dim)]
419 const int64_t* __restrict__ slot_mapping, // [num_tokens]
420 const int block_stride, //
421 const int entry_stride, //
422 const int kv_c_stride, //
423 const int k_pe_stride, //
424 const int kv_lora_rank, //
425 const int pe_dim, //
426 const int block_size, //
427 const float* scale //
428 ) {
429 const int64_t token_idx = blockIdx.x;
430 const int64_t slot_idx = slot_mapping[token_idx];
431 // NOTE: slot_idx can be -1 if the token is padded
432 if (slot_idx < 0) {
433 return;
434 }
435 const int64_t block_idx = slot_idx / block_size;
436 const int64_t block_offset = slot_idx % block_size;
437
438 auto copy = [&](const scalar_t* __restrict__ src, cache_t* __restrict__ dst,
439 int src_stride, int dst_stride, int size, int offset) {
440 for (int i = threadIdx.x; i < size; i += blockDim.x) {
441 const int64_t src_idx = token_idx * src_stride + i;
442 const int64_t dst_idx =
443 block_idx * block_stride + block_offset * entry_stride + i + offset;
444 if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
445 dst[dst_idx] = src[src_idx];
446 } else {
447 dst[dst_idx] =
448 fp8::scaled_convert<cache_t, scalar_t, kv_dt>(src[src_idx], *scale);
449 }
450 }
451 };
452
453 copy(kv_c, kv_cache, kv_c_stride, block_stride, kv_lora_rank, 0);
454 copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank);
455 }
456
457 // Grouped variant of concat_and_cache_mla: inserts the context K/V for every
458 // draft layer in a single launch. Grid is (num_tokens, num_layers); each layer
459 // reads its own cache base pointer from kv_cache_ptrs (same pointer-array
460 // pattern as copy_blocks_kernel).
461 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
462 __global__ void concat_and_cache_mla_grouped_kernel(
463 const scalar_t* __restrict__ kv_c, // [num_layers, num_tokens,
464 // kv_lora_rank]
465 const scalar_t* __restrict__ k_pe, // [num_layers, num_tokens, pe_dim]
466 const int64_t* __restrict__ kv_cache_ptrs, // [num_layers]
467 const float* __restrict__ kv_scales, // [num_layers] or nullptr
468 const int64_t* __restrict__ slot_mapping, // [num_layers, num_tokens]
469 const int64_t kv_c_layer_stride, const int64_t kv_c_token_stride,
470 const int64_t k_pe_layer_stride, const int64_t k_pe_token_stride,
471 const int64_t slot_layer_stride, const int64_t block_stride,
472 const int64_t entry_stride, const int kv_lora_rank, const int pe_dim,
473 const int block_size) {
474 const int64_t token_idx = blockIdx.x;
475 const int64_t layer_idx = blockIdx.y;
476 const int64_t slot_idx =
477 slot_mapping[layer_idx * slot_layer_stride + token_idx];
478 // NOTE: slot_idx can be -1 if the token is padded
479 if (slot_idx < 0) {
480 return;
481 }
482 const int64_t block_idx = slot_idx / block_size;
483 const int64_t block_offset = slot_idx % block_size;
484
485 cache_t* __restrict__ kv_cache =
486 reinterpret_cast<cache_t*>(kv_cache_ptrs[layer_idx]);
487 const scalar_t* __restrict__ kv_c_layer =
488 kv_c + layer_idx * kv_c_layer_stride;
489 const scalar_t* __restrict__ k_pe_layer =
490 k_pe + layer_idx * k_pe_layer_stride;
491 float scale = 0.0f;
492 if constexpr (kv_dt != Fp8KVCacheDataType::kAuto) {
493 scale = kv_scales[layer_idx];
494 }
495
496 auto copy = [&](const scalar_t* __restrict__ src, int64_t src_token_stride,
497 int size, int offset) {
498 for (int i = threadIdx.x; i < size; i += blockDim.x) {
499 const int64_t src_idx = token_idx * src_token_stride + i;
500 const int64_t dst_idx =
501 block_idx * block_stride + block_offset * entry_stride + i + offset;
502 if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
503 kv_cache[dst_idx] = src[src_idx];
504 } else {
505 kv_cache[dst_idx] =
506 fp8::scaled_convert<cache_t, scalar_t, kv_dt>(src[src_idx], scale);
507 }
508 }
509 };
510
511 copy(kv_c_layer, kv_c_token_stride, kv_lora_rank, 0);
512 copy(k_pe_layer, k_pe_token_stride, pe_dim, kv_lora_rank);
513 }
514
515 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
516 __global__ void concat_and_cache_ds_mla_kernel(
517 const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank]
518 const scalar_t* __restrict__ k_pe, // [num_tokens, pe_dim]
519 cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank
520 // + pe_dim)]
521 const int64_t* __restrict__ slot_mapping, // [num_tokens]
522 const int block_stride, //
523 const int entry_stride, //
524 const int kv_c_stride, //
525 const int k_pe_stride, //
526 const int kv_lora_rank, //
527 const int pe_dim, //
528 const int block_size, //
529 const float* scale //
530 ) {
531 const int64_t token_idx = blockIdx.x;
532 const int64_t slot_idx = slot_mapping[token_idx];
533 // NOTE: slot_idx can be -1 if the token is padded
534 if (slot_idx < 0) {
535 return;
536 }
537 const int64_t block_idx = slot_idx / block_size;
538 const int64_t block_offset = slot_idx % block_size;
539 const int64_t dst_idx_start =
540 block_idx * block_stride + block_offset * entry_stride;
541
542 // For the NoPE part, each tile of 128 elements is handled by half of one warp
543 // (16 threads). There are 4 total tiles, so 2 warps (64 threads).
544 // Lanes 0 and 16 of each warp write the scale values for that warp's tiles.
545 // The RoPE part (last 64 elements) is handled by another 1 warp (32 threads).
546 // So in total, we use 3 warps (96 threads) per block.
547
548 // Cast kv_cache to 16_bit for RoPE values
549 scalar_t* kv_cache_16bit =
550 reinterpret_cast<scalar_t*>(&kv_cache[dst_idx_start]);
551
552 // Zero the reserved RoPE tail for NoPE rows.
553 if (threadIdx.x >= 64) {
554 // Each thread handles two elements of RoPE
555 const int8_t pe_idx_start = (threadIdx.x - 64) * 2;
556 // RoPE values start after the packed 8-bit NoPE values and the
557 // 32-bit scales
558 const int64_t dst_idx = kv_lora_rank / 2 + 8 + pe_idx_start;
559 if (pe_dim == 0) {
560 *reinterpret_cast<int32_t*>(&kv_cache_16bit[dst_idx]) = 0;
561 return;
562 }
563 const int64_t src_idx = token_idx * k_pe_stride + pe_idx_start;
564 // Vectorized load of two 16-bit values, performed as one 32-bit load
565 const int32_t vals = *reinterpret_cast<const int32_t*>(&k_pe[src_idx]);
566 // Vectorized store of two 16-bit values, performed as one 32-bit store
567 *reinterpret_cast<int32_t*>(&kv_cache_16bit[dst_idx]) = vals;
568 return;
569 }
570
571 // The first two warps handle the NoPE part
572 const int8_t warp_idx = threadIdx.x >> 5;
573 const int8_t lane_idx = threadIdx.x & 31;
574 const int8_t tile_idx = warp_idx * 2 + (lane_idx >> 4);
575
576 // Each thread handles 8 elements of NoPE
577 // Load the NoPE elements for this thread into registers
578 const int64_t src_idx_start = token_idx * kv_c_stride + (threadIdx.x * 8);
579 // Vectorized load of eight 16-bit values, performed as an int4 load
580 const int4 vals_i4 = *reinterpret_cast<const int4*>(&kv_c[src_idx_start]);
581 const scalar_t* vals = reinterpret_cast<const scalar_t*>(&vals_i4);
582
583 // Max absolute value of this thread's elements
584 float max_abs = fmaxf(fmaxf(fmaxf(fabsf(vals[0]), fabsf(vals[1])),
585 fmaxf(fabsf(vals[2]), fabsf(vals[3]))),
586 fmaxf(fmaxf(fabsf(vals[4]), fabsf(vals[5])),
587 fmaxf(fabsf(vals[6]), fabsf(vals[7]))));
588
589 // Warp-level reduction to find the max absolute value in each half-warp
590 #pragma unroll
591 for (int offset = 8; offset > 0; offset /= 2) {
592 max_abs = fmaxf(max_abs, VLLM_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 16));
593 }
594
595 // Both SM90 and SM100 readers preserve power-of-two fp32 scales exactly.
596 float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, 1e-4f);
597 tile_scale = exp2f(ceilf(log2f(tile_scale)));
598
599 // The first lane of each half-warp writes the scale to kv_cache
600 if ((lane_idx == 0) || (lane_idx == 16)) {
601 float* kv_cache_32bit = reinterpret_cast<float*>(&kv_cache[dst_idx_start]);
602 const uint64_t dst_idx = kv_lora_rank / 4 + tile_idx;
603 kv_cache_32bit[dst_idx] = tile_scale;
604 }
605
606 // Now all threads in the block scale and write their elements
607 // NoPE data occupies the first kv_lora_rank bytes (512 bytes).
608 const int64_t dst_idx_base = dst_idx_start + (threadIdx.x * 8);
609
610 uint8_t result[8];
611 #pragma unroll
612 for (int i = 0; i < 8; i++) {
613 result[i] =
614 fp8::scaled_convert<uint8_t, scalar_t, Fp8KVCacheDataType::kFp8E4M3>(
615 vals[i], tile_scale);
616 }
617
618 // Store as aligned 64-bit writes
619 *reinterpret_cast<uint64_t*>(&kv_cache[dst_idx_base]) =
620 *reinterpret_cast<const uint64_t*>(result);
621 }
622
623 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
624 __global__ void indexer_k_quant_and_cache_kernel(
625 const scalar_t* __restrict__ k, // [num_tokens, head_dim]
626 cache_t* __restrict__ kv_cache, // [num_blocks, block_size, cache_stride]
627 const int64_t* __restrict__ slot_mapping, // [num_tokens]
628 const int head_dim, // dimension of each head
629 const int quant_block_size, // quantization block size
630 const int cache_block_size, // cache block size
631 const int64_t cache_block_stride, // stride for each block in kv_cache
632
633 const bool use_ue8m0 // use ue8m0 scale format
634 ) {
635 constexpr int VEC_SIZE = 4;
636 const int64_t token_idx = blockIdx.x;
637 const int64_t head_dim_idx = (blockIdx.y * blockDim.y * blockDim.x +
638 threadIdx.y * blockDim.x + threadIdx.x) *
639 VEC_SIZE;
640 const int64_t slot_idx = slot_mapping[token_idx];
641 const int64_t block_idx = slot_idx / cache_block_size;
642 const int64_t block_offset = slot_idx % cache_block_size;
643
644 // NOTE: slot_idx can be -1 if the token is padded
645 if (slot_idx < 0 || (head_dim_idx >= head_dim)) {
646 return;
647 }
648
649 float2 k_val = (reinterpret_cast<const float2*>(
650 k))[(token_idx * head_dim + head_dim_idx) / VEC_SIZE];
651 scalar_t* k_val_ptr = reinterpret_cast<scalar_t*>(&k_val);
652 float amax = 0.0f;
653 for (int i = 0; i < VEC_SIZE; i++) {
654 amax = fmaxf(amax, fabsf(float(k_val_ptr[i])));
655 }
656
657 // Reduced amax
658 for (int mask = 16; mask > 0; mask /= 2) {
659 #ifdef USE_ROCM
660 amax = fmaxf(amax, __shfl_xor_sync(uint64_t(-1), amax, mask));
661 #else
662 amax = fmaxf(amax, __shfl_xor_sync(unsigned(-1), amax, mask));
663 #endif
664 }
665
666 float scale = fmaxf(amax, 1e-4) / kFp8ScaleDivisor;
667
668 if (use_ue8m0) {
669 scale = exp2f(ceilf(log2f(scale)));
670 }
671
672 const int64_t dst_offset =
673 block_idx * cache_block_stride + block_offset * head_dim + head_dim_idx;
674 for (int i = 0; i < VEC_SIZE; i++) {
675 kv_cache[dst_offset + i] =
676 fp8::scaled_convert<cache_t, scalar_t, kv_dt>(k_val_ptr[i], scale);
677 }
678 if (threadIdx.x == 0) {
679 const int64_t dst_scale_idx =
680 block_idx * cache_block_stride + cache_block_size * head_dim +
681 (block_offset * head_dim + head_dim_idx) * 4 / quant_block_size;
682 reinterpret_cast<float*>(kv_cache)[dst_scale_idx / 4] = scale;
683 }
684 }
685
686 template <int BLOCK_Y_SIZE>
687 __global__ void cp_gather_indexer_k_quant_cache_kernel(
688 const char* __restrict__ kv_cache, // [num_blocks, block_size,
689 // cache_stride]
690 char* __restrict__ dst_k, // [num_tokens, head_dim]
691 char* __restrict__ dst_scale, // [num_tokens, head_dim / quant_block_size *
692 // 4]
693 const int* __restrict__ block_table, // [batch_size, num_blocks]
694 const int* __restrict__ cu_seq_lens, // [batch_size + 1]
695 const int batch_size, // batch size
696 const int64_t token_stride, // stride for each token in dst_k
697 const int64_t head_dim, // dimension of each head
698 const int64_t block_stride, // stride for each block in kv_cache
699 const int64_t cache_token_stride, // stride for each token in kv_cache
700 const int64_t cache_block_size, // num_tokens for each block in kv_cache
701 const int num_blocks, // number of blocks
702 const int num_tokens, // number of tokens
703 const int quant_block_size // quantization block size
704 ) {
705 constexpr int VEC_SIZE = sizeof(float4) / sizeof(char);
706 const int token_idx = blockIdx.x * blockDim.y + threadIdx.y;
707 const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE;
708 // Find batch index within a block
709 __shared__ int batch_idx[BLOCK_Y_SIZE];
710 if (threadIdx.x == 0) {
711 batch_idx[threadIdx.y] = -1;
712 }
713 __syncthreads();
714
715 for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x));
716 iter++) {
717 int tid = iter * blockDim.x + threadIdx.x;
718 if (tid < batch_size) {
719 const int seq_start = cu_seq_lens[tid];
720 const int seq_end = cu_seq_lens[tid + 1];
721 if (token_idx >= seq_start && token_idx < seq_end) {
722 batch_idx[threadIdx.y] = tid;
723 }
724 }
725 }
726
727 __syncthreads();
728
729 // num_tokens may be an allocation upper bound when Python avoids a D2H sync.
730 // Only tokens covered by the exact device-side cu_seq_lens are valid to
731 // gather.
732 const int batch = batch_idx[threadIdx.y];
733 if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) {
734 return;
735 }
736 const int inbatch_seq_idx = token_idx - cu_seq_lens[batch];
737 const int block_idx =
738 block_table[batch * num_blocks + inbatch_seq_idx / cache_block_size];
739 const int64_t src_block_offset = block_idx * block_stride;
740 const int64_t cache_inblock_offset =
741 (inbatch_seq_idx % cache_block_size) * head_dim + head_idx;
742 const int64_t src_inblock_offset = src_block_offset + cache_inblock_offset;
743 const int64_t dst_inblock_offset = token_idx * token_stride + head_idx;
744
745 reinterpret_cast<float4*>(dst_k)[dst_inblock_offset / VEC_SIZE] =
746 reinterpret_cast<const float4*>(kv_cache)[src_inblock_offset / VEC_SIZE];
747 ;
748 if (threadIdx.x == 0) {
749 const int64_t src_scale_offset =
750 src_block_offset + cache_block_size * head_dim +
751 cache_inblock_offset * 4 / quant_block_size;
752 reinterpret_cast<float*>(dst_scale)[dst_inblock_offset / quant_block_size] =
753 reinterpret_cast<const float*>(kv_cache)[src_scale_offset / 4];
754 }
755 }
756
757 } // namespace vllm
758
759 // KV_T is the data type of key and value tensors.
760 // CACHE_T is the stored data type of kv-cache.
761 // KV_DTYPE is the real data type of kv-cache.
762 #define CALL_RESHAPE_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \
763 vllm::reshape_and_cache_kernel<KV_T, CACHE_T, KV_DTYPE> \
764 <<<grid, block, 0, stream>>>( \
765 reinterpret_cast<KV_T*>(key.data_ptr()), \
766 reinterpret_cast<KV_T*>(value.data_ptr()), \
767 reinterpret_cast<CACHE_T*>(key_cache.data_ptr()), \
768 reinterpret_cast<CACHE_T*>(value_cache.data_ptr()), \
769 slot_mapping.const_data_ptr<int64_t>(), key_stride, value_stride, \
770 num_heads, head_size, block_size, x, \
771 reinterpret_cast<const float*>(k_scale.data_ptr()), \
772 reinterpret_cast<const float*>(v_scale.data_ptr()));
773
774 void reshape_and_cache(
775 torch::stable::Tensor& key, // [num_tokens, num_heads, head_size]
776 torch::stable::Tensor& value, // [num_tokens, num_heads, head_size]
777 torch::stable::Tensor&
778 key_cache, // [num_blocks, num_heads, head_size/x, block_size, x]
779 torch::stable::Tensor&
780 value_cache, // [num_blocks, num_heads, head_size, block_size]
781 torch::stable::Tensor& slot_mapping, // [num_tokens]
782 const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
783 torch::stable::Tensor& v_scale) {
784 int num_tokens = slot_mapping.size(0);
785 int num_heads = key.size(1);
786 int head_size = key.size(2);
787 int block_size = key_cache.size(3);
788 int x = key_cache.size(4);
789
790 int key_stride = key.stride(0);
791 int value_stride = value.stride(0);
792 int head_div_x = head_size / x;
793
794 dim3 grid(num_tokens);
795 dim3 block(std::min(num_heads * head_div_x, 512));
796 const torch::stable::accelerator::DeviceGuard device_guard(
797 key.get_device_index());
798 const cudaStream_t stream = get_current_cuda_stream();
799
800 DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype,
801 CALL_RESHAPE_AND_CACHE);
802 }
803
804 // KV_T is the data type of key and value tensors.
805 // CACHE_T is the stored data type of kv-cache.
806 // KV_DTYPE is the real data type of kv-cache.
807 #define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE) \
808 vllm::reshape_and_cache_flash_kernel<KV_T, CACHE_T, KV_DTYPE> \
809 <<<grid, block, 0, stream>>>( \
810 reinterpret_cast<KV_T*>(key.data_ptr()), \
811 reinterpret_cast<KV_T*>(value.data_ptr()), \
812 reinterpret_cast<CACHE_T*>(key_cache.data_ptr()), \
813 reinterpret_cast<CACHE_T*>(value_cache.data_ptr()), \
814 slot_mapping.const_data_ptr<int64_t>(), block_stride, page_stride, \
815 head_stride, key_stride, value_stride, num_heads, head_size, \
816 block_size, reinterpret_cast<const float*>(k_scale.data_ptr()), \
817 reinterpret_cast<const float*>(v_scale.data_ptr()), \
818 kv_scale_stride);
819
820 void reshape_and_cache_flash(
821 torch::stable::Tensor& key, // [num_tokens, num_heads, head_size]
822 torch::stable::Tensor& value, // [num_tokens, num_heads, head_size]
823 torch::stable::Tensor&
824 key_cache, // [num_blocks, block_size, num_heads, head_size]
825 torch::stable::Tensor&
826 value_cache, // [num_blocks, block_size, num_heads, head_size]
827 torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
828 const std::string& kv_cache_dtype,
829 torch::stable::Tensor& k_scale, // [1] or [num_heads]
830 torch::stable::Tensor& v_scale) { // [1] or [num_heads]
831 // NOTE(woosuk): In vLLM V1, key.size(0) can be different from
832 // slot_mapping.size(0) because of padding for CUDA graphs.
833 // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
834 // both include padding.
835 // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
836 // since key includes padding for CUDA graphs, while slot_mapping does not.
837 // In this case, slot_mapping.size(0) represents the actual number of tokens
838 // before padding.
839 // For compatibility with both cases, we use slot_mapping.size(0) as the
840 // number of tokens.
841 int num_tokens = slot_mapping.size(0);
842 int num_heads = key.size(1);
843 int head_size = key.size(2);
844
845 const torch::stable::accelerator::DeviceGuard device_guard(
846 key.get_device_index());
847 const cudaStream_t stream = get_current_cuda_stream();
848
849 if (kv_cache_dtype == "nvfp4" || kv_cache_dtype == "nvfp4_4over6") {
850 #if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
851 // NVFP4 dispatch is compiled separately for SM100+.
852 extern void reshape_and_cache_nvfp4_dispatch(
853 torch::stable::Tensor & key, torch::stable::Tensor & value,
854 torch::stable::Tensor & key_cache, torch::stable::Tensor & value_cache,
855 torch::stable::Tensor & slot_mapping, torch::stable::Tensor & k_scale,
856 torch::stable::Tensor & v_scale, const std::string& kv_cache_dtype);
857 reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache,
858 slot_mapping, k_scale, v_scale,
859 kv_cache_dtype);
860 return;
861 #else
862 STD_TORCH_CHECK(
863 false,
864 "NVFP4 KV cache requires SM100+ (Blackwell). "
865 "Please rebuild vllm with a Blackwell-compatible CUDA target.");
866 #endif
867 }
868
869 // Original FP8/auto path.
870 int block_size = key_cache.size(1);
871
872 int64_t key_stride = key.stride(0);
873 int64_t value_stride = value.stride(0);
874 int64_t block_stride = key_cache.stride(0);
875 int64_t page_stride = key_cache.stride(1);
876 int64_t head_stride = key_cache.stride(2);
877 STD_TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0));
878
879 STD_TORCH_CHECK(k_scale.sizes().equals(v_scale.sizes()),
880 "k_scale and v_scale must have the same shape");
881 STD_TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads,
882 "k_scale and v_scale must be of shape [1] or [num_heads]");
883 int kv_scale_stride = (k_scale.numel() > 1) ? 1 : 0;
884
885 dim3 grid(num_tokens);
886 dim3 block(std::min(num_heads * head_size, 512));
887
888 DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype,
889 CALL_RESHAPE_AND_CACHE_FLASH);
890 }
891
892 // KV_T is the data type of key and value tensors.
893 // CACHE_T is the stored data type of kv-cache.
894 // KV_DTYPE is the real data type of kv-cache.
895 #define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \
896 vllm::concat_and_cache_mla_kernel<KV_T, CACHE_T, KV_DTYPE> \
897 <<<grid, block, 0, stream>>>( \
898 reinterpret_cast<KV_T*>(kv_c.data_ptr()), \
899 reinterpret_cast<KV_T*>(k_pe.data_ptr()), \
900 reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
901 slot_mapping.const_data_ptr<int64_t>(), block_stride, entry_stride, \
902 kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
903 reinterpret_cast<const float*>(scale.data_ptr()));
904
905 // KV_T is the data type of key and value tensors.
906 // CACHE_T is the stored data type of kv-cache.
907 #define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \
908 vllm::concat_and_cache_ds_mla_kernel<KV_T, CACHE_T, KV_DTYPE> \
909 <<<grid, block, 0, stream>>>( \
910 reinterpret_cast<KV_T*>(kv_c.data_ptr()), \
911 reinterpret_cast<KV_T*>(k_pe.data_ptr()), \
912 reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
913 slot_mapping.const_data_ptr<int64_t>(), block_stride, entry_stride, \
914 kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
915 reinterpret_cast<const float*>(scale.data_ptr()));
916
917 void concat_and_cache_mla(
918 torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank]
919 torch::stable::Tensor& k_pe, // [num_tokens, pe_dim]
920 torch::stable::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank
921 // + pe_dim)]
922 torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
923 const std::string& kv_cache_dtype, torch::stable::Tensor& scale) {
924 // NOTE(woosuk): In vLLM V1, key.size(0) can be different from
925 // slot_mapping.size(0) because of padding for CUDA graphs.
926 // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
927 // both include padding.
928 // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
929 // since key includes padding for CUDA graphs, while slot_mapping does not.
930 // In this case, slot_mapping.size(0) represents the actual number of tokens
931 // before padding.
932 // For compatibility with both cases, we use slot_mapping.size(0) as the
933 // number of tokens.
934 int num_tokens = slot_mapping.size(0);
935 int kv_lora_rank = kv_c.size(1);
936 int pe_dim = k_pe.size(1);
937 int block_size = kv_cache.size(1);
938
939 const bool is_nvfp4_ds_mla = kv_cache_dtype == "nvfp4_ds_mla";
940 if (kv_cache_dtype == "fp8_ds_mla") {
941 STD_TORCH_CHECK(kv_lora_rank == 512,
942 "kv_lora_rank must be 512 for fp8_ds_mla");
943 STD_TORCH_CHECK(pe_dim == 64 || pe_dim == 0,
944 "pe_dim must be 64 or 0 for fp8_ds_mla");
945 STD_TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.element_size(),
946 "kv_cache.size(2) must be 656 bytes for fp8_ds_mla");
947 STD_TORCH_CHECK(kv_c.element_size() == 2,
948 "kv_c.element_size() must be 2 for fp8_ds_mla");
949 STD_TORCH_CHECK(k_pe.element_size() == 2,
950 "k_pe.element_size() must be 2 for fp8_ds_mla");
951 } else if (is_nvfp4_ds_mla) {
952 constexpr int bytes_per_token = 352;
953 STD_TORCH_CHECK(kv_lora_rank == 512, "kv_lora_rank must be 512 for ",
954 kv_cache_dtype);
955 STD_TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for ", kv_cache_dtype);
956 STD_TORCH_CHECK(
957 kv_cache.size(2) == bytes_per_token / kv_cache.element_size(),
958 "kv_cache.size(2) must be ", bytes_per_token, " bytes for ",
959 kv_cache_dtype);
960 STD_TORCH_CHECK(
961 kv_c.scalar_type() == torch::headeronly::ScalarType::BFloat16,
962 "kv_c must be bfloat16 for ", kv_cache_dtype);
963 STD_TORCH_CHECK(
964 k_pe.scalar_type() == torch::headeronly::ScalarType::BFloat16,
965 "k_pe must be bfloat16 for ", kv_cache_dtype);
966 } else {
967 STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim);
968 }
969
970 int kv_c_stride = kv_c.stride(0);
971 int k_pe_stride = k_pe.stride(0);
972 int block_stride = kv_cache.stride(0);
973 int entry_stride = kv_cache.stride(1);
974
975 const torch::stable::accelerator::DeviceGuard device_guard(
976 kv_c.get_device_index());
977 const cudaStream_t stream = get_current_cuda_stream();
978
979 if (kv_cache_dtype == "fp8_ds_mla") {
980 dim3 grid(num_tokens);
981 // For the NoPE part, each tile of 128 elements is handled by half of one
982 // warp (16 threads). There are 4 total tiles, so 2 warps (64 threads).
983 // Lanes 0 and 16 of each warp write the scale values for that warp's tiles.
984 // The RoPE part (last 64 elements) is handled by another 1 warp (32
985 // threads). So in total, we use 3 warps (96 threads) per block.
986 dim3 block(96);
987 DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype,
988 CALL_CONCAT_AND_CACHE_DS_MLA);
989 } else if (is_nvfp4_ds_mla) {
990 // The kernels live in nvfp4_ds_mla_cache_kernels.cu, which is only built
991 // when the target list contains an SM100 arch.
992 #if !defined(USE_ROCM) && defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
993 const cudaDeviceProp* props = get_device_prop();
994 STD_TORCH_CHECK(props->major == 10,
995 "nvfp4_ds_mla KV cache requires SM100 (Blackwell); "
996 "got SM",
997 props->major, props->minor);
998 vllm::launch_concat_and_cache_nvfp4_ds_mla(
999 kv_c.data_ptr(), k_pe.data_ptr(), kv_cache.data_ptr(),
1000 slot_mapping.const_data_ptr<int64_t>(), block_stride, entry_stride,
1001 kv_c_stride, k_pe_stride, block_size, num_tokens, stream);
1002 #else
1003 STD_TORCH_CHECK(false,
1004 "the nvfp4 ds_mla kv-cache format requires a CUDA build "
1005 "targeting SM100 (Blackwell)");
1006 #endif
1007 } else {
1008 dim3 grid(num_tokens);
1009 dim3 block(std::min(kv_lora_rank, 512));
1010 DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype,
1011 CALL_CONCAT_AND_CACHE_MLA);
1012 }
1013 }
1014
1015 void concat_and_cache_mla_grouped(
1016 torch::stable::Tensor& kv_c, // [num_layers, num_tokens, kv_lora_rank]
1017 torch::stable::Tensor& k_pe, // [num_layers, num_tokens, pe_dim]
1018 torch::stable::Tensor& kv_cache_ptrs, // [num_layers] int64, on device
1019 torch::stable::Tensor& slot_mapping, // [num_layers, num_tokens] int64
1020 int64_t block_size, int64_t block_stride, int64_t entry_stride,
1021 std::optional<torch::stable::Tensor> kv_scales, // [num_layers] or None
1022 const std::string& kv_cache_dtype) {
1023 const bool use_fp8 = kv_cache_dtype == "fp8" ||
1024 kv_cache_dtype == "fp8_e4m3" ||
1025 kv_cache_dtype == "fp8_e5m2";
1026 #ifdef USE_ROCM
1027 STD_TORCH_CHECK(kv_cache_dtype != "fp8_e5m2",
1028 "concat_and_cache_mla_grouped does not support fp8_e5m2 "
1029 "KV cache on ROCm");
1030 #endif
1031 STD_TORCH_CHECK(
1032 use_fp8 || kv_cache_dtype == "auto" || kv_cache_dtype == "bfloat16",
1033 "concat_and_cache_mla_grouped only supports BF16 and plain "
1034 "FP8 KV cache; got ",
1035 kv_cache_dtype);
1036
1037 STD_TORCH_CHECK(
1038 kv_c.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
1039 k_pe.scalar_type() == torch::headeronly::ScalarType::BFloat16,
1040 "concat_and_cache_mla_grouped requires BF16 inputs; got kv_c=",
1041 kv_c.scalar_type(), ", k_pe=", k_pe.scalar_type());
1042 STD_TORCH_CHECK(
1043 kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long &&
1044 slot_mapping.scalar_type() == torch::headeronly::ScalarType::Long,
1045 "cache pointers and slot mapping must be int64");
1046
1047 const int num_layers = kv_c.size(0);
1048 const int num_tokens = kv_c.size(1);
1049 const int kv_lora_rank = kv_c.size(2);
1050 const int pe_dim = k_pe.size(2);
1051 const float* kv_scales_ptr = nullptr;
1052 if (use_fp8) {
1053 STD_TORCH_CHECK(kv_scales.has_value(),
1054 "FP8 grouped cache insert requires kv_scales");
1055 STD_TORCH_CHECK(
1056 kv_scales->scalar_type() == torch::headeronly::ScalarType::Float,
1057 "kv_scales must be float32");
1058 STD_TORCH_CHECK(kv_scales->numel() == num_layers,
1059 "kv_scales must contain one scale per layer");
1060 STD_TORCH_CHECK(kv_scales->is_cuda() && kv_scales->get_device_index() ==
1061 kv_c.get_device_index(),
1062 "kv_scales must be on the same CUDA device as kv_c");
1063 STD_TORCH_CHECK(kv_scales->is_contiguous(), "kv_scales must be contiguous");
1064 kv_scales_ptr = kv_scales->const_data_ptr<float>();
1065 } else {
1066 STD_TORCH_CHECK(!kv_scales.has_value(),
1067 "BF16 grouped cache insert does not use kv_scales");
1068 }
1069
1070 if (num_tokens == 0 || num_layers == 0) {
1071 return;
1072 }
1073
1074 const int64_t kv_c_layer_stride = kv_c.stride(0);
1075 const int64_t kv_c_token_stride = kv_c.stride(1);
1076 const int64_t k_pe_layer_stride = k_pe.stride(0);
1077 const int64_t k_pe_token_stride = k_pe.stride(1);
1078 const int64_t slot_layer_stride = slot_mapping.stride(0);
1079
1080 const torch::stable::accelerator::DeviceGuard device_guard(
1081 kv_c.get_device_index());
1082 const cudaStream_t stream = get_current_cuda_stream();
1083
1084 const dim3 grid(num_tokens, num_layers);
1085 const dim3 block(std::min(kv_lora_rank, 512));
1086
1087 if (!use_fp8) {
1088 vllm::concat_and_cache_mla_grouped_kernel<uint16_t, uint16_t,
1089 vllm::Fp8KVCacheDataType::kAuto>
1090 <<<grid, block, 0, stream>>>(
1091 reinterpret_cast<const uint16_t*>(kv_c.data_ptr()),
1092 reinterpret_cast<const uint16_t*>(k_pe.data_ptr()),
1093 kv_cache_ptrs.const_data_ptr<int64_t>(), nullptr,
1094 slot_mapping.const_data_ptr<int64_t>(), kv_c_layer_stride,
1095 kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride,
1096 slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim,
1097 block_size);
1098 return;
1099 }
1100
1101 #define LAUNCH_GROUPED_FP8(KV_DTYPE) \
1102 vllm::concat_and_cache_mla_grouped_kernel<__nv_bfloat16, uint8_t, KV_DTYPE> \
1103 <<<grid, block, 0, stream>>>( \
1104 reinterpret_cast<const __nv_bfloat16*>(kv_c.data_ptr()), \
1105 reinterpret_cast<const __nv_bfloat16*>(k_pe.data_ptr()), \
1106 kv_cache_ptrs.const_data_ptr<int64_t>(), kv_scales_ptr, \
1107 slot_mapping.const_data_ptr<int64_t>(), kv_c_layer_stride, \
1108 kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride, \
1109 slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim, \
1110 block_size)
1111
1112 if (kv_cache_dtype == "fp8_e5m2") {
1113 LAUNCH_GROUPED_FP8(vllm::Fp8KVCacheDataType::kFp8E5M2);
1114 } else {
1115 LAUNCH_GROUPED_FP8(vllm::Fp8KVCacheDataType::kFp8E4M3);
1116 }
1117 #undef LAUNCH_GROUPED_FP8
1118 }
1119
1120 namespace vllm {
1121
1122 template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
1123 __global__ void convert_fp8_kernel(const Tin* __restrict__ src_cache,
1124 Tout* __restrict__ dst_cache,
1125 const float scale,
1126 const int64_t block_stride) {
1127 const int64_t block_idx = blockIdx.x;
1128 for (int i = threadIdx.x; i < block_stride; i += blockDim.x) {
1129 int64_t idx = block_idx * block_stride + i;
1130 dst_cache[idx] =
1131 fp8::scaled_convert<Tout, Tin, kv_dt>(src_cache[idx], scale);
1132 }
1133 }
1134
1135 } // namespace vllm
1136
1137 #define CALL_CONVERT_FP8(Tout, Tin, KV_DTYPE) \
1138 vllm::convert_fp8_kernel<Tout, Tin, KV_DTYPE><<<grid, block, 0, stream>>>( \
1139 reinterpret_cast<Tin*>(src_cache.data_ptr()), \
1140 reinterpret_cast<Tout*>(dst_cache.data_ptr()), scale, block_stride);
1141
1142 // Only for testing.
1143 void convert_fp8(torch::stable::Tensor& dst_cache,
1144 torch::stable::Tensor& src_cache, const double scale,
1145 const std::string& kv_cache_dtype) {
1146 torch::stable::Device src_device = src_cache.device();
1147 torch::stable::Device dst_device = dst_cache.device();
1148 STD_TORCH_CHECK(src_device.is_cuda(), "src must be on a GPU")
1149 STD_TORCH_CHECK(dst_device.is_cuda(), "dst must be on a GPU")
1150 STD_TORCH_CHECK(src_device.index() == dst_device.index(),
1151 "src and dst must be on the same GPU");
1152 torch::stable::accelerator::DeviceGuard device_guard(src_device.index());
1153
1154 int64_t num_blocks = src_cache.size(0);
1155 int64_t block_stride = src_cache.stride(0);
1156
1157 dim3 grid(num_blocks);
1158 dim3 block(std::min(block_stride, int64_t(512)));
1159 const cudaStream_t stream = get_current_cuda_stream();
1160
1161 if (kv_cache_dtype == "auto") {
1162 if (src_cache.scalar_type() == torch::headeronly::ScalarType::Float) {
1163 CALL_CONVERT_FP8(uint8_t, float, vllm::Fp8KVCacheDataType::kAuto);
1164 } else if (src_cache.scalar_type() == torch::headeronly::ScalarType::Half) {
1165 CALL_CONVERT_FP8(uint8_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto);
1166 } else if (src_cache.scalar_type() ==
1167 torch::headeronly::ScalarType::BFloat16) {
1168 CALL_CONVERT_FP8(uint8_t, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto);
1169 } else if (dst_cache.scalar_type() ==
1170 torch::headeronly::ScalarType::Float) {
1171 CALL_CONVERT_FP8(float, uint8_t, vllm::Fp8KVCacheDataType::kAuto);
1172 } else if (dst_cache.scalar_type() == torch::headeronly::ScalarType::Half) {
1173 CALL_CONVERT_FP8(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kAuto);
1174 } else if (dst_cache.scalar_type() ==
1175 torch::headeronly::ScalarType::BFloat16) {
1176 CALL_CONVERT_FP8(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kAuto);
1177 }
1178 } else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") {
1179 if (src_cache.scalar_type() == torch::headeronly::ScalarType::Float) {
1180 CALL_CONVERT_FP8(uint8_t, float, vllm::Fp8KVCacheDataType::kFp8E4M3);
1181 } else if (src_cache.scalar_type() == torch::headeronly::ScalarType::Half) {
1182 CALL_CONVERT_FP8(uint8_t, uint16_t, vllm::Fp8KVCacheDataType::kFp8E4M3);
1183 } else if (src_cache.scalar_type() ==
1184 torch::headeronly::ScalarType::BFloat16) {
1185 CALL_CONVERT_FP8(uint8_t, __nv_bfloat16,
1186 vllm::Fp8KVCacheDataType::kFp8E4M3);
1187 } else if (dst_cache.scalar_type() ==
1188 torch::headeronly::ScalarType::Float) {
1189 CALL_CONVERT_FP8(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);
1190 } else if (dst_cache.scalar_type() == torch::headeronly::ScalarType::Half) {
1191 CALL_CONVERT_FP8(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);
1192 } else if (dst_cache.scalar_type() ==
1193 torch::headeronly::ScalarType::BFloat16) {
1194 CALL_CONVERT_FP8(__nv_bfloat16, uint8_t,
1195 vllm::Fp8KVCacheDataType::kFp8E4M3);
1196 }
1197 } else {
1198 STD_TORCH_CHECK(false, "Unsupported data type: ", kv_cache_dtype);
1199 }
1200 }
1201
1202 namespace vllm {
1203
1204 struct GatherPageTask {
1205 int32_t req_id;
1206 int32_t logical_block;
1207 int32_t page_token_begin;
1208 int32_t page_token_end;
1209 int32_t output_token_begin;
1210 };
1211
1212 template <bool has_terminal_start>
1213 __device__ __forceinline__ bool map_gather_page_task(
1214 int32_t task, const int32_t* __restrict__ output_starts, int32_t num_reqs,
1215 int32_t total_tokens, int32_t block_size,
1216 const int32_t* __restrict__ seq_starts, GatherPageTask& page) {
1217 int32_t relative_page = task;
1218 for (int32_t req_id = 0; req_id < num_reqs; ++req_id) {
1219 const int32_t output_begin = min(output_starts[req_id], total_tokens);
1220 int32_t output_end;
1221 if constexpr (has_terminal_start) {
1222 output_end = min(output_starts[req_id + 1], total_tokens);
1223 } else {
1224 output_end =
1225 min(req_id + 1 < num_reqs ? output_starts[req_id + 1] : total_tokens,
1226 total_tokens);
1227 }
1228 const int32_t seq_len = max(output_end - output_begin, 0);
1229 const int32_t source_begin = seq_starts == nullptr ? 0 : seq_starts[req_id];
1230 const int32_t first_block = source_begin / block_size;
1231 const int32_t num_pages =
1232 cuda_utils::ceil_div(source_begin + seq_len, block_size) - first_block;
1233 if (relative_page < num_pages) {
1234 page.req_id = req_id;
1235 page.logical_block = first_block + relative_page;
1236 const int32_t page_begin = page.logical_block * block_size;
1237 const int32_t copy_begin = max(source_begin, page_begin);
1238 const int32_t copy_end =
1239 min(source_begin + seq_len, page_begin + block_size);
1240 page.page_token_begin = copy_begin - page_begin;
1241 page.page_token_end = copy_end - page_begin;
1242 page.output_token_begin = output_begin + copy_begin - source_begin;
1243 return true;
1244 }
1245 relative_page -= num_pages;
1246 }
1247 return false;
1248 }
1249
1250 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
1251 int ENTRY_SIZE>
1252 __global__ void gather_and_maybe_dequant_cache_page(
1253 const cache_t* __restrict__ src_cache, scalar_t* __restrict__ dst,
1254 const int32_t* __restrict__ block_table,
1255 const int32_t* __restrict__ cu_seq_lens, const int32_t num_reqs,
1256 const int32_t num_tokens, const int32_t block_size,
1257 const int64_t block_table_stride, const int64_t cache_block_stride,
1258 const int64_t cache_entry_stride, const int64_t dst_entry_stride,
1259 const float* __restrict__ scale, const int32_t* __restrict__ seq_starts) {
1260 constexpr int32_t vec_size = sizeof(float4) / sizeof(scalar_t);
1261 constexpr int32_t vec_iter_cnt = ENTRY_SIZE / vec_size;
1262 static_assert(ENTRY_SIZE % vec_size == 0);
1263 using ltype = vllm::vec_n_t<cache_t, vec_size>;
1264 using stype = vllm::vec_n_t<scalar_t, vec_size>;
1265
1266 __shared__ GatherPageTask page;
1267 __shared__ int32_t physical_block;
1268 __shared__ bool has_task;
1269 __shared__ bool copy_task;
1270 __shared__ float scale_value;
1271
1272 if (threadIdx.x == 0) {
1273 if constexpr (kv_dt != Fp8KVCacheDataType::kAuto) {
1274 scale_value = *scale;
1275 }
1276 }
1277 __syncthreads();
1278
1279 if (threadIdx.x == 0) {
1280 has_task =
1281 map_gather_page_task<true>(blockIdx.x, cu_seq_lens, num_reqs,
1282 num_tokens, block_size, seq_starts, page);
1283 copy_task = has_task && page.logical_block < block_table_stride;
1284 if (copy_task) {
1285 physical_block =
1286 block_table[page.req_id * block_table_stride + page.logical_block];
1287 }
1288 }
1289 __syncthreads();
1290 if (!has_task) {
1291 return;
1292 }
1293
1294 if (copy_task) {
1295 const int32_t page_vectors =
1296 (page.page_token_end - page.page_token_begin) * vec_iter_cnt;
1297 for (int32_t flat_idx = threadIdx.x; flat_idx < page_vectors;
1298 flat_idx += blockDim.x) {
1299 const int32_t token_offset = flat_idx / vec_iter_cnt;
1300 const int32_t idx = flat_idx - token_offset * vec_iter_cnt;
1301 const int32_t page_token = page.page_token_begin + token_offset;
1302 const int32_t output_token = page.output_token_begin + token_offset;
1303 const cache_t* src = src_cache + physical_block * cache_block_stride +
1304 page_token * cache_entry_stride;
1305 scalar_t* output = dst + output_token * dst_entry_stride;
1306
1307 if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
1308 reinterpret_cast<stype*>(output)[idx] =
1309 static_cast<stype>(reinterpret_cast<const ltype*>(src)[idx]);
1310 } else {
1311 const ltype loaded = reinterpret_cast<const ltype*>(src)[idx];
1312 stype converted;
1313 #pragma unroll
1314 for (int32_t j = 0; j < vec_size; ++j) {
1315 converted.val[j] = fp8::scaled_convert<scalar_t, cache_t, kv_dt>(
1316 loaded.val[j], scale_value);
1317 }
1318 reinterpret_cast<stype*>(output)[idx] = converted;
1319 }
1320 }
1321 }
1322 }
1323
1324 } // namespace vllm
1325
1326 // Macro to dispatch the kernel based on the data type.
1327 // SCALAR_T is the data type of the destination tensor.
1328 // CACHE_T is the stored data type of kv-cache.
1329 // KV_DTYPE is the real data type of kv-cache.
1330 #define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, ENTRY_SZ) \
1331 vllm::gather_and_maybe_dequant_cache_page<SCALAR_T, CACHE_T, KV_DTYPE, \
1332 ENTRY_SZ> \
1333 <<<grid, block, 0, stream>>>( \
1334 reinterpret_cast<CACHE_T*>(src_cache.data_ptr()), \
1335 reinterpret_cast<SCALAR_T*>(dst.data_ptr()), \
1336 block_table.const_data_ptr<int32_t>(), \
1337 cu_seq_lens.const_data_ptr<int32_t>(), num_reqs, \
1338 static_cast<int32_t>(num_tokens), block_size, block_table_stride, \
1339 cache_block_stride, cache_entry_stride, dst_entry_stride, \
1340 reinterpret_cast<const float*>(scale.data_ptr()), seq_starts_ptr)
1341
1342 #define CALL_GATHER_CACHE_576(SCALAR_T, CACHE_T, KV_DTYPE) \
1343 CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 576)
1344
1345 #define CALL_GATHER_CACHE_512(SCALAR_T, CACHE_T, KV_DTYPE) \
1346 CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 512)
1347
1348 #define CALL_GATHER_CACHE_320(SCALAR_T, CACHE_T, KV_DTYPE) \
1349 CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 320)
1350
1351 // Gather sequences from the cache into the destination tensor.
1352 // - cu_seq_lens contains the cumulative sequence lengths for each batch
1353 // - block_table contains the cache block indices for each sequence
1354 // - Optionally, seq_starts (if provided) offsets the starting block index by
1355 // (seq_starts[bid] / page_size)
1356 void gather_and_maybe_dequant_cache(
1357 torch::stable::Tensor const&
1358 src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
1359 torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
1360 torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
1361 torch::stable::Tensor const& cu_seq_lens, // [BATCH+1]
1362 torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS]
1363 int64_t num_tokens, const std::string& kv_cache_dtype,
1364 torch::stable::Tensor const& scale,
1365 std::optional<torch::stable::Tensor> seq_starts = std::nullopt) {
1366 torch::stable::accelerator::DeviceGuard device_guard(
1367 src_cache.get_device_index());
1368 const cudaStream_t stream = get_current_cuda_stream();
1369 (void)token_to_seq;
1370
1371 int32_t block_size = src_cache.size(1);
1372 int32_t head_dim = dst.size(-1);
1373
1374 STD_TORCH_CHECK(
1375 block_table.scalar_type() == torch::headeronly::ScalarType::Int,
1376 "block_table must be int32");
1377 STD_TORCH_CHECK(
1378 cu_seq_lens.scalar_type() == torch::headeronly::ScalarType::Int,
1379 "cu_seq_lens must be int32");
1380 if (seq_starts.has_value()) {
1381 STD_TORCH_CHECK(
1382 seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int,
1383 "seq_starts must be int32");
1384 }
1385 STD_TORCH_CHECK(head_dim == 320 || head_dim == 512 || head_dim == 576,
1386 "gather_and_maybe_dequant_cache only support the head_dim to "
1387 "320 or 512 or 576 for better performance")
1388
1389 STD_TORCH_CHECK(src_cache.device() == dst.device(),
1390 "src_cache and dst must be on the same device");
1391 STD_TORCH_CHECK(src_cache.device() == block_table.device(),
1392 "src_cache and block_table must be on the same device");
1393 STD_TORCH_CHECK(src_cache.device() == cu_seq_lens.device(),
1394 "src_cache and cu_seq_lens must be on the same device");
1395 if (seq_starts.has_value()) {
1396 STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
1397 "src_cache and seq_starts must be on the same device");
1398 }
1399
1400 if (num_tokens == 0) {
1401 return;
1402 }
1403
1404 const int32_t num_reqs = cu_seq_lens.size(0) - 1;
1405 const int64_t block_table_stride = block_table.stride(0);
1406 const int64_t cache_block_stride = src_cache.stride(0);
1407 const int64_t cache_entry_stride = src_cache.stride(1);
1408 const int64_t dst_entry_stride = dst.stride(0);
1409 const int32_t page_threads = num_tokens >= (1 << 20) ? 128 : 256;
1410 const int32_t required_blocks =
1411 cuda_utils::ceil_div(static_cast<int32_t>(num_tokens), block_size) +
1412 2 * num_reqs;
1413 const dim3 grid(required_blocks);
1414 const dim3 block(page_threads);
1415
1416 const int32_t* seq_starts_ptr =
1417 seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
1418 : nullptr;
1419
1420 if (head_dim == 576) {
1421 DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype,
1422 CALL_GATHER_CACHE_576);
1423 } else if (head_dim == 512) {
1424 DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype,
1425 CALL_GATHER_CACHE_512);
1426 } else {
1427 DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype,
1428 CALL_GATHER_CACHE_320);
1429 }
1430 }
1431
1432 namespace vllm {
1433
1434 __device__ __forceinline__ void gather_and_upconvert_fp8_token(
1435 const uint8_t* __restrict__ token_ptr, __nv_bfloat16* __restrict__ dst_ptr,
1436 int32_t lane_id) {
1437 const uint2* fp8_src = reinterpret_cast<const uint2*>(token_ptr);
1438 const float* scales = reinterpret_cast<const float*>(token_ptr + 512);
1439 int4* nope_dst = reinterpret_cast<int4*>(dst_ptr);
1440
1441 #pragma unroll
1442 for (int32_t phase = 0; phase < 2; ++phase) {
1443 const int32_t chunk = phase * 32 + lane_id;
1444 const uint2 fp8_data = fp8_src[chunk];
1445 const float scale = scales[chunk >> 4];
1446 #ifdef USE_ROCM
1447 const bf16_8_t bf16_data =
1448 fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_data, scale);
1449 #else
1450 const bf16_8_t bf16_data =
1451 fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_data, scale, __NV_E4M3);
1452 #endif
1453 nope_dst[chunk] = *reinterpret_cast<const int4*>(&bf16_data);
1454 }
1455
1456 const int* rope_src = reinterpret_cast<const int*>(token_ptr + 528);
1457 int* rope_dst = reinterpret_cast<int*>(dst_ptr + 512);
1458 rope_dst[lane_id] = rope_src[lane_id];
1459 }
1460
1461 __global__ void cp_gather_and_upconvert_fp8_kv_cache_page(
1462 const uint8_t* __restrict__ src_cache, __nv_bfloat16* __restrict__ dst,
1463 const int32_t* __restrict__ block_table,
1464 const int32_t* __restrict__ workspace_starts, const int32_t num_reqs,
1465 const int32_t block_size, const int32_t total_tokens,
1466 const int64_t block_table_stride, const int64_t cache_block_stride,
1467 const int64_t cache_entry_stride, const int64_t dst_entry_stride,
1468 const int32_t* __restrict__ seq_starts,
1469 const uint8_t* __restrict__ host_cache,
1470 const int32_t* __restrict__ host_row_ids,
1471 const int32_t* __restrict__ device_row_ids,
1472 const int64_t host_entry_stride) {
1473 constexpr int32_t warps_per_cta = 16;
1474 __shared__ GatherPageTask page;
1475 __shared__ int32_t physical_block;
1476 __shared__ bool has_task;
1477
1478 if (threadIdx.x == 0) {
1479 has_task =
1480 map_gather_page_task<false>(blockIdx.x, workspace_starts, num_reqs,
1481 total_tokens, block_size, seq_starts, page);
1482 if (has_task) {
1483 physical_block =
1484 block_table[page.req_id * block_table_stride + page.logical_block];
1485 }
1486 }
1487 __syncthreads();
1488 if (!has_task) {
1489 return;
1490 }
1491
1492 const int32_t warp_id = threadIdx.x >> 5;
1493 const int32_t lane_id = threadIdx.x & 31;
1494 for (int32_t page_token = page.page_token_begin + warp_id;
1495 page_token < page.page_token_end; page_token += warps_per_cta) {
1496 const int32_t output_token =
1497 page.output_token_begin + page_token - page.page_token_begin;
1498 const int32_t source_row = physical_block * block_size + page_token;
1499 const uint8_t* token_ptr;
1500 if (device_row_ids == nullptr) {
1501 token_ptr = src_cache + physical_block * cache_block_stride +
1502 page_token * cache_entry_stride;
1503 } else {
1504 const int32_t device_row = device_row_ids[source_row];
1505 if (device_row >= 0) {
1506 token_ptr = src_cache + (device_row / block_size) * cache_block_stride +
1507 (device_row % block_size) * cache_entry_stride;
1508 } else {
1509 token_ptr =
1510 host_cache +
1511 static_cast<int64_t>(host_row_ids[source_row]) * host_entry_stride;
1512 }
1513 }
1514 __nv_bfloat16* dst_ptr = dst + output_token * dst_entry_stride;
1515 gather_and_upconvert_fp8_token(token_ptr, dst_ptr, lane_id);
1516 }
1517 }
1518
1519 template <bool contiguous_entries, bool vectorized>
1520 __global__ void cp_gather_cache_page(
1521 const uint8_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
1522 // ENTRY_SIZE_BYTES]
1523 uint8_t* __restrict__ dst, // [TOT_TOKENS, ENTRY_SIZE_BYTES]
1524 const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES]
1525 const int32_t* __restrict__ cu_seq_lens, // [BATCH+1]
1526 const int32_t num_reqs, const int32_t block_size,
1527 const int32_t entry_size_bytes, const int32_t total_tokens,
1528 const int64_t block_table_stride, const int64_t cache_block_stride,
1529 const int64_t cache_entry_stride, const int64_t dst_entry_stride,
1530 const int32_t* __restrict__ seq_starts) {
1531 constexpr int32_t threads_per_token = 32;
1532 constexpr int32_t tokens_per_cta = 256 / threads_per_token;
1533 __shared__ GatherPageTask page;
1534 __shared__ int32_t physical_block;
1535 __shared__ bool has_task;
1536
1537 if (threadIdx.x == 0) {
1538 has_task =
1539 map_gather_page_task<true>(blockIdx.x, cu_seq_lens, num_reqs,
1540 total_tokens, block_size, seq_starts, page);
1541 if (has_task) {
1542 physical_block =
1543 block_table[page.req_id * block_table_stride + page.logical_block];
1544 }
1545 }
1546 __syncthreads();
1547 if (!has_task) {
1548 return;
1549 }
1550
1551 if constexpr (contiguous_entries) {
1552 const int64_t src_offset = physical_block * cache_block_stride +
1553 page.page_token_begin * entry_size_bytes;
1554 const int64_t dst_offset = page.output_token_begin * entry_size_bytes;
1555 const int32_t copy_bytes =
1556 (page.page_token_end - page.page_token_begin) * entry_size_bytes;
1557 if constexpr (vectorized) {
1558 const int4* src = reinterpret_cast<const int4*>(src_cache + src_offset);
1559 int4* output = reinterpret_cast<int4*>(dst + dst_offset);
1560 const int32_t num_vectors = copy_bytes / sizeof(int4);
1561 int32_t idx = threadIdx.x;
1562 for (; idx + blockDim.x < num_vectors; idx += 2 * blockDim.x) {
1563 const int4 first = src[idx];
1564 const int4 second = src[idx + blockDim.x];
1565 output[idx] = first;
1566 output[idx + blockDim.x] = second;
1567 }
1568 for (; idx < num_vectors; idx += blockDim.x) {
1569 output[idx] = src[idx];
1570 }
1571 } else {
1572 const uint8_t* src = src_cache + src_offset;
1573 uint8_t* output = dst + dst_offset;
1574 for (int32_t idx = threadIdx.x; idx < copy_bytes; idx += blockDim.x) {
1575 output[idx] = src[idx];
1576 }
1577 }
1578 } else {
1579 const int32_t token_group = threadIdx.x / threads_per_token;
1580 const int32_t lane_id = threadIdx.x % threads_per_token;
1581 for (int32_t page_token = page.page_token_begin + token_group;
1582 page_token < page.page_token_end; page_token += tokens_per_cta) {
1583 const int32_t output_token =
1584 page.output_token_begin + page_token - page.page_token_begin;
1585 const uint8_t* src = src_cache + physical_block * cache_block_stride +
1586 page_token * cache_entry_stride;
1587 uint8_t* output = dst + output_token * dst_entry_stride;
1588 if constexpr (vectorized) {
1589 const int4* src_vec = reinterpret_cast<const int4*>(src);
1590 int4* dst_vec = reinterpret_cast<int4*>(output);
1591 const int32_t num_vectors = entry_size_bytes / sizeof(int4);
1592 for (int32_t idx = lane_id; idx < num_vectors;
1593 idx += threads_per_token) {
1594 dst_vec[idx] = src_vec[idx];
1595 }
1596 } else {
1597 for (int32_t idx = lane_id; idx < entry_size_bytes;
1598 idx += threads_per_token) {
1599 output[idx] = src[idx];
1600 }
1601 }
1602 }
1603 }
1604 }
1605 } // namespace vllm
1606
1607 // Gather sequences from the cache into the destination tensor.
1608 // - cu_seq_lens contains the cumulative sequence lengths for each batch
1609 // - block_table contains the cache block indices for each sequence
1610 // - Optionally, seq_starts (if provided) offsets the starting slot index by
1611 // seq_starts[bid]
1612 void cp_gather_cache(
1613 torch::stable::Tensor const&
1614 src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
1615 torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
1616 torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
1617 torch::stable::Tensor const& cu_seq_lens, // [BATCH+1]
1618 int64_t batch_size,
1619 std::optional<torch::stable::Tensor> seq_starts = std::nullopt) {
1620 torch::stable::accelerator::DeviceGuard device_guard(
1621 src_cache.get_device_index());
1622 const cudaStream_t stream = get_current_cuda_stream();
1623
1624 int32_t block_size = src_cache.size(1);
1625 int32_t entry_size = torch::stable::flatten(src_cache, 2, -1).size(2);
1626
1627 STD_TORCH_CHECK(
1628 block_table.scalar_type() == torch::headeronly::ScalarType::Int,
1629 "block_table must be int32");
1630 STD_TORCH_CHECK(
1631 cu_seq_lens.scalar_type() == torch::headeronly::ScalarType::Int,
1632 "cu_seq_lens must be int32");
1633 if (seq_starts.has_value()) {
1634 STD_TORCH_CHECK(
1635 seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int,
1636 "seq_starts must be int32");
1637 }
1638
1639 STD_TORCH_CHECK(src_cache.device() == dst.device(),
1640 "src_cache and dst must be on the same device");
1641 STD_TORCH_CHECK(src_cache.device() == block_table.device(),
1642 "src_cache and block_table must be on the same device");
1643 STD_TORCH_CHECK(src_cache.device() == cu_seq_lens.device(),
1644 "src_cache and cu_seq_lens must be on the same device");
1645 if (seq_starts.has_value()) {
1646 STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
1647 "src_cache and seq_starts must be on the same device");
1648 }
1649
1650 STD_TORCH_CHECK(src_cache.scalar_type() == dst.scalar_type(),
1651 "src_cache and dst must have the same dtype");
1652 const int32_t element_size = src_cache.element_size();
1653 STD_TORCH_CHECK(element_size == 1 || element_size == 2 || element_size == 4,
1654 "Unsupported data type width: ", element_size * 8);
1655 STD_TORCH_CHECK(batch_size == cu_seq_lens.size(0) - 1,
1656 "batch_size must match cu_seq_lens");
1657
1658 const int32_t total_tokens = dst.size(0);
1659 if (total_tokens == 0) {
1660 return;
1661 }
1662
1663 const int32_t entry_size_bytes = entry_size * element_size;
1664 const int64_t block_table_stride = block_table.stride(0);
1665 const int64_t cache_block_stride = src_cache.stride(0) * element_size;
1666 const int64_t cache_entry_stride = src_cache.stride(1) * element_size;
1667 const int64_t dst_entry_stride = dst.stride(0) * element_size;
1668 const int32_t* seq_starts_ptr =
1669 seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
1670 : nullptr;
1671
1672 constexpr std::uintptr_t vector_alignment = alignof(int4);
1673 const bool vectorized =
1674 reinterpret_cast<std::uintptr_t>(src_cache.data_ptr()) %
1675 vector_alignment ==
1676 0 &&
1677 reinterpret_cast<std::uintptr_t>(dst.data_ptr()) % vector_alignment ==
1678 0 &&
1679 entry_size_bytes % vector_alignment == 0 &&
1680 cache_block_stride % vector_alignment == 0 &&
1681 cache_entry_stride % vector_alignment == 0 &&
1682 dst_entry_stride % vector_alignment == 0;
1683 const bool contiguous_entries = cache_entry_stride == entry_size_bytes &&
1684 dst_entry_stride == entry_size_bytes;
1685
1686 const int32_t num_reqs = static_cast<int32_t>(batch_size);
1687 const int32_t required_blocks =
1688 cuda_utils::ceil_div(total_tokens, block_size) + 2 * num_reqs;
1689 const dim3 grid(required_blocks);
1690 const dim3 block(256);
1691
1692 #define CALL_CP_GATHER_CACHE(CONTIGUOUS, VECTORIZED) \
1693 vllm::cp_gather_cache_page<CONTIGUOUS, VECTORIZED> \
1694 <<<grid, block, 0, stream>>>( \
1695 reinterpret_cast<const uint8_t*>(src_cache.data_ptr()), \
1696 reinterpret_cast<uint8_t*>(dst.data_ptr()), \
1697 block_table.const_data_ptr<int32_t>(), \
1698 cu_seq_lens.const_data_ptr<int32_t>(), num_reqs, block_size, \
1699 entry_size_bytes, total_tokens, block_table_stride, \
1700 cache_block_stride, cache_entry_stride, dst_entry_stride, \
1701 seq_starts_ptr)
1702
1703 if (contiguous_entries && vectorized) {
1704 CALL_CP_GATHER_CACHE(true, true);
1705 } else if (contiguous_entries) {
1706 CALL_CP_GATHER_CACHE(true, false);
1707 } else if (vectorized) {
1708 CALL_CP_GATHER_CACHE(false, true);
1709 } else {
1710 CALL_CP_GATHER_CACHE(false, false);
1711 }
1712
1713 #undef CALL_CP_GATHER_CACHE
1714 }
1715
1716 void cp_gather_and_upconvert_fp8_kv_cache(
1717 torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]
1718 torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
1719 torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
1720 torch::stable::Tensor const& workspace_starts, // [BATCH]
1721 int64_t batch_size,
1722 std::optional<torch::stable::Tensor> seq_starts = std::nullopt,
1723 std::optional<torch::stable::Tensor> host_cache = std::nullopt,
1724 std::optional<torch::stable::Tensor> host_row_ids = std::nullopt,
1725 std::optional<torch::stable::Tensor> device_row_ids = std::nullopt) {
1726 torch::stable::accelerator::DeviceGuard device_guard(
1727 src_cache.get_device_index());
1728 const cudaStream_t stream = get_current_cuda_stream();
1729
1730 int32_t block_size = src_cache.size(1);
1731 int32_t head_dim = dst.size(1);
1732
1733 STD_TORCH_CHECK(
1734 block_table.scalar_type() == torch::headeronly::ScalarType::Int,
1735 "block_table must be int32");
1736 STD_TORCH_CHECK(
1737 workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int,
1738 "workspace_starts must be int32");
1739 if (seq_starts.has_value()) {
1740 STD_TORCH_CHECK(
1741 seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int,
1742 "seq_starts must be int32");
1743 }
1744
1745 const bool has_host_rows = host_cache.has_value();
1746 STD_TORCH_CHECK(
1747 has_host_rows == host_row_ids.has_value() &&
1748 has_host_rows == device_row_ids.has_value(),
1749 "host_cache, host_row_ids, and device_row_ids must be provided together");
1750
1751 STD_TORCH_CHECK(src_cache.device() == dst.device(),
1752 "src_cache and dst must be on the same device");
1753 STD_TORCH_CHECK(src_cache.device() == block_table.device(),
1754 "src_cache and block_table must be on the same device");
1755 STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(),
1756 "src_cache and workspace_starts must be on the same device");
1757 if (seq_starts.has_value()) {
1758 STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
1759 "src_cache and seq_starts must be on the same device");
1760 }
1761 const uint8_t* host_cache_ptr = nullptr;
1762 const int32_t* host_row_ids_ptr = nullptr;
1763 const int32_t* device_row_ids_ptr = nullptr;
1764 int64_t host_entry_stride = 0;
1765 if (has_host_rows) {
1766 auto const& host = host_cache.value();
1767 auto const& host_rows = host_row_ids.value();
1768 auto const& device_rows = device_row_ids.value();
1769 STD_TORCH_CHECK(
1770 host.device().is_cpu() &&
1771 host.scalar_type() == torch::headeronly::ScalarType::Byte &&
1772 host.dim() == 2 && host.is_contiguous() &&
1773 host.size(1) == src_cache.size(2),
1774 "host_cache must be a contiguous uint8 CPU row matrix matching "
1775 "src_cache");
1776 cudaPointerAttributes attributes{};
1777 const cudaError_t pointer_status =
1778 cudaPointerGetAttributes(&attributes, host.const_data_ptr());
1779 if (pointer_status != cudaSuccess) {
1780 cudaGetLastError();
1781 }
1782 STD_TORCH_CHECK(
1783 pointer_status == cudaSuccess && attributes.type == cudaMemoryTypeHost,
1784 "host_cache must be CUDA-accessible pinned memory");
1785 STD_TORCH_CHECK(
1786 host_rows.is_cuda() && device_rows.is_cuda() &&
1787 host_rows.scalar_type() == torch::headeronly::ScalarType::Int &&
1788 device_rows.scalar_type() == torch::headeronly::ScalarType::Int &&
1789 host_rows.is_contiguous() && device_rows.is_contiguous() &&
1790 host_rows.numel() == device_rows.numel() && host_rows.dim() <= 2 &&
1791 device_rows.dim() <= 2,
1792 "host/device row maps must be matching contiguous CUDA int32 tensors");
1793 host_cache_ptr = reinterpret_cast<const uint8_t*>(host.const_data_ptr());
1794 host_row_ids_ptr = host_rows.const_data_ptr<int32_t>();
1795 device_row_ids_ptr = device_rows.const_data_ptr<int32_t>();
1796 host_entry_stride = host.stride(0);
1797 }
1798 auto dtype = src_cache.scalar_type();
1799 STD_TORCH_CHECK(
1800 dtype == torch::headeronly::ScalarType::Byte || // uint8
1801 dtype == torch::headeronly::ScalarType::Float8_e4m3fn || // fp8 e4m3
1802 dtype == torch::headeronly::ScalarType::Float8_e5m2, // fp8 e5m2
1803 "src_cache must be uint8, float8_e4m3fn, or float8_e5m2, but got ",
1804 src_cache.scalar_type());
1805 STD_TORCH_CHECK(dst.scalar_type() == torch::headeronly::ScalarType::BFloat16,
1806 "dst must be bfloat16");
1807 STD_TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA");
1808
1809 int64_t block_table_stride = block_table.stride(0);
1810 int64_t cache_block_stride = src_cache.stride(0);
1811 int64_t cache_entry_stride = src_cache.stride(1);
1812 int64_t dst_entry_stride = dst.stride(0);
1813
1814 const uint8_t* src_ptr = nullptr;
1815 if (dtype == torch::headeronly::ScalarType::Byte) {
1816 src_ptr = src_cache.const_data_ptr<uint8_t>();
1817 } else {
1818 // float8_e4m3fn or float8_e5m2
1819 src_ptr = reinterpret_cast<const uint8_t*>(src_cache.data_ptr());
1820 }
1821
1822 const int total_tokens = dst.size(0);
1823 if (total_tokens == 0) {
1824 return;
1825 }
1826 constexpr int warps_per_block = 16;
1827 const int block_size_threads = warps_per_block * 32;
1828 const int32_t* seq_starts_ptr =
1829 seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
1830 : nullptr;
1831 const int required_blocks = cuda_utils::ceil_div(total_tokens, block_size) +
1832 2 * static_cast<int32_t>(batch_size);
1833 vllm::cp_gather_and_upconvert_fp8_kv_cache_page<<<
1834 required_blocks, block_size_threads, 0, stream>>>(
1835 src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
1836 block_table.const_data_ptr<int32_t>(),
1837 workspace_starts.const_data_ptr<int32_t>(),
1838 static_cast<int32_t>(batch_size), block_size, total_tokens,
1839 block_table_stride, cache_block_stride, cache_entry_stride,
1840 dst_entry_stride, seq_starts_ptr, host_cache_ptr, host_row_ids_ptr,
1841 device_row_ids_ptr, host_entry_stride);
1842 }
1843
1844 void cp_gather_and_upconvert_nvfp4_kv_cache(
1845 torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
1846 // 352]
1847 torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
1848 torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
1849 torch::stable::Tensor const& workspace_starts, // [BATCH]
1850 int64_t batch_size) {
1851 // The kernel lives in nvfp4_ds_mla_cache_kernels.cu, which is only built when
1852 // the target list contains an SM100 arch.
1853 #if !defined(USE_ROCM) && defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
1854 torch::stable::accelerator::DeviceGuard device_guard(
1855 src_cache.get_device_index());
1856 const cudaStream_t stream = get_current_cuda_stream();
1857
1858 const cudaDeviceProp* props = get_device_prop();
1859 STD_TORCH_CHECK(props->major == 10,
1860 "nvfp4_ds_mla KV cache requires SM100 (Blackwell); "
1861 "got SM",
1862 props->major, props->minor);
1863
1864 int32_t block_size = src_cache.size(1);
1865 int32_t head_dim = dst.size(1);
1866
1867 STD_TORCH_CHECK(
1868 block_table.scalar_type() == torch::headeronly::ScalarType::Int,
1869 "block_table must be int32");
1870 STD_TORCH_CHECK(
1871 workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int,
1872 "workspace_starts must be int32");
1873 STD_TORCH_CHECK(src_cache.device() == dst.device(),
1874 "src_cache and dst must be on the same device");
1875 STD_TORCH_CHECK(src_cache.device() == block_table.device(),
1876 "src_cache and block_table must be on the same device");
1877 STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(),
1878 "src_cache and workspace_starts must be on the same device");
1879 STD_TORCH_CHECK(
1880 src_cache.scalar_type() == torch::headeronly::ScalarType::Byte,
1881 "src_cache must be uint8 for the nvfp4 ds_mla format, but got ",
1882 src_cache.scalar_type());
1883 STD_TORCH_CHECK(dst.scalar_type() == torch::headeronly::ScalarType::BFloat16,
1884 "dst must be bfloat16");
1885 STD_TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA");
1886 STD_TORCH_CHECK(src_cache.size(2) == 352,
1887 "src_cache.size(2) must be 352 for the nvfp4 ds_mla format");
1888
1889 int64_t block_table_stride = block_table.stride(0);
1890 int64_t cache_block_stride = src_cache.stride(0);
1891 int64_t cache_entry_stride = src_cache.stride(1);
1892 int64_t dst_entry_stride = dst.stride(0);
1893
1894 const int total_tokens = dst.size(0);
1895
1896 vllm::launch_cp_gather_and_upconvert_nvfp4_kv_cache(
1897 src_cache.const_data_ptr<uint8_t>(), dst.data_ptr(),
1898 block_table.const_data_ptr<int32_t>(),
1899 workspace_starts.const_data_ptr<int32_t>(),
1900 static_cast<int32_t>(batch_size), block_size, total_tokens,
1901 block_table_stride, cache_block_stride, cache_entry_stride,
1902 dst_entry_stride, stream);
1903 #else
1904 STD_TORCH_CHECK(false,
1905 "the nvfp4 ds_mla kv-cache format requires a CUDA build "
1906 "targeting SM100 (Blackwell)");
1907 #endif
1908 }
1909
1910 // Macro to dispatch the kernel based on the data type.
1911 #define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \
1912 vllm::indexer_k_quant_and_cache_kernel<KV_T, CACHE_T, KV_DTYPE> \
1913 <<<grid, block, 0, stream>>>( \
1914 reinterpret_cast<KV_T*>(k.data_ptr()), \
1915 reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
1916 slot_mapping.const_data_ptr<int64_t>(), head_dim, quant_block_size, \
1917 cache_block_size, cache_block_stride, use_ue8m0);
1918
1919 void indexer_k_quant_and_cache(
1920 torch::stable::Tensor& k, // [num_tokens, head_dim]
1921 torch::stable::Tensor& kv_cache, // [num_blocks, block_size, cache_stride]
1922 torch::stable::Tensor& slot_mapping, // [num_tokens]
1923 int64_t quant_block_size, // quantization block size
1924 const std::string& scale_fmt) {
1925 int num_tokens = k.size(0);
1926 int head_dim = k.size(1);
1927 int cache_block_size = kv_cache.size(1);
1928 int64_t cache_block_stride = kv_cache.stride(0);
1929 bool use_ue8m0 = scale_fmt == "ue8m0";
1930
1931 STD_TORCH_CHECK(k.device() == kv_cache.device(),
1932 "k and kv_cache must be on the same device");
1933 STD_TORCH_CHECK(k.device() == slot_mapping.device(),
1934 "k and slot_mapping must be on the same device");
1935 STD_TORCH_CHECK(head_dim % quant_block_size == 0,
1936 "head_dim must be divisible by quant_block_size");
1937
1938 constexpr int vec_size = 4;
1939 dim3 grid(num_tokens, (head_dim + quant_block_size * vec_size - 1) /
1940 (quant_block_size * vec_size));
1941 dim3 block(32, vec_size);
1942 const torch::stable::accelerator::DeviceGuard device_guard(
1943 k.get_device_index());
1944 const cudaStream_t stream = get_current_cuda_stream();
1945
1946 static const std::string kv_cache_dtype = "fp8_e4m3";
1947 DISPATCH_BY_KV_CACHE_DTYPE(k.scalar_type(), kv_cache_dtype,
1948 CALL_INDEXER_K_QUANT_AND_CACHE);
1949 }
1950
1951 // Macro to dispatch the kernel based on the data amount.
1952 #define CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(BLOCK_Y_SIZE) \
1953 vllm::cp_gather_indexer_k_quant_cache_kernel<BLOCK_Y_SIZE> \
1954 <<<dim3((num_tokens + BLOCK_Y_SIZE - 1) / BLOCK_Y_SIZE, \
1955 (head_dim + 8 * vec_size - 1) / (8 * vec_size)), \
1956 dim3(8, BLOCK_Y_SIZE), 0, stream>>>( \
1957 reinterpret_cast<char*>(kv_cache.data_ptr()), \
1958 reinterpret_cast<char*>(dst_k.data_ptr()), \
1959 reinterpret_cast<char*>(dst_scale.data_ptr()), \
1960 block_table.const_data_ptr<int32_t>(), \
1961 cu_seq_lens.const_data_ptr<int32_t>(), batch_size, dst_k.stride(0), \
1962 dst_k.size(1), kv_cache.stride(0), kv_cache.stride(1), \
1963 kv_cache.size(1), block_table.size(1), num_tokens, \
1964 quant_block_size);
1965
1966 void cp_gather_indexer_k_quant_cache(
1967 const torch::stable::Tensor&
1968 kv_cache, // [num_blocks, block_size, cache_stride]
1969 torch::stable::Tensor& dst_k, // [num_tokens, head_dim]
1970 torch::stable::Tensor&
1971 dst_scale, // [num_tokens, head_dim / quant_block_size * 4]
1972 const torch::stable::Tensor& block_table, // [batch_size, num_blocks]
1973 const torch::stable::Tensor& cu_seq_lens // [batch_size + 1]
1974 ) {
1975 int batch_size = block_table.size(0);
1976 int num_tokens = dst_k.size(0);
1977 int head_dim = dst_k.size(1);
1978 int quant_block_size = head_dim * 4 / dst_scale.size(1);
1979
1980 STD_TORCH_CHECK(kv_cache.device() == dst_k.device(),
1981 "kv_cache and dst_k must be on the same device");
1982 STD_TORCH_CHECK(kv_cache.device() == dst_scale.device(),
1983 "kv_cache and dst_scale must be on the same device");
1984 STD_TORCH_CHECK(kv_cache.device() == block_table.device(),
1985 "kv_cache and block_table must be on the same device");
1986 STD_TORCH_CHECK(kv_cache.device() == cu_seq_lens.device(),
1987 "kv_cache and cu_seq_lens must be on the same device");
1988 STD_TORCH_CHECK(head_dim % quant_block_size == 0,
1989 "head_dim must be divisible by quant_block_size");
1990
1991 constexpr int vec_size = 16;
1992 const torch::stable::accelerator::DeviceGuard device_guard(
1993 kv_cache.get_device_index());
1994 const cudaStream_t stream = get_current_cuda_stream();
1995
1996 if (num_tokens < 32) {
1997 CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(1);
1998 } else if (num_tokens < 64) {
1999 CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(2);
2000 } else if (num_tokens < 128) {
2001 CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(4);
2002 } else if (num_tokens < 256) {
2003 CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(8);
2004 } else if (num_tokens < 512) {
2005 CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(16);
2006 } else {
2007 CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(32);
2008 }
2009 }
2010
2011 // Concatenate ql_nope and q_pe into a contiguous q_out tensor for MLA/DSA.
2012 // Replaces torch.cat((ql_nope, q_pe), dim=-1).
2013 void concat_mla_q(
2014 torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
2015 torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim]
2016 torch::stable::Tensor& q_out // [num_tokens, num_heads, nope_dim +
2017 // rope_dim]
2018 ) {
2019 const int num_tokens = ql_nope.size(0);
2020 const int num_heads = ql_nope.size(1);
2021 const int nope_dim = ql_nope.size(2);
2022 const int rope_dim = q_pe.size(2);
2023
2024 STD_TORCH_CHECK(nope_dim % 512 == 0,
2025 "nope_dim must be a multiple of 512, got ", nope_dim);
2026 STD_TORCH_CHECK(rope_dim == 64, "rope_dim must be 64, got ", rope_dim);
2027 STD_TORCH_CHECK(q_out.size(2) == nope_dim + rope_dim);
2028
2029 STD_TORCH_CHECK(ql_nope.stride(2) == 1,
2030 "ql_nope must have stride 1 in dim 2");
2031 STD_TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
2032 STD_TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
2033 STD_TORCH_CHECK(
2034 ql_nope.scalar_type() == torch::headeronly::ScalarType::Half ||
2035 ql_nope.scalar_type() == torch::headeronly::ScalarType::BFloat16,
2036 "ql_nope must be float16 or bfloat16 dtype");
2037
2038 if (num_tokens == 0) return;
2039
2040 constexpr int warps_per_block = 8;
2041 const int total_warps = num_tokens * num_heads;
2042 const int grid_size = (total_warps + warps_per_block - 1) / warps_per_block;
2043 const int block_size = warps_per_block * 32;
2044
2045 const torch::stable::accelerator::DeviceGuard device_guard(
2046 ql_nope.get_device_index());
2047 const cudaStream_t stream = get_current_cuda_stream();
2048
2049 VLLM_STABLE_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
2050 vllm::ConcatMLAQKernel<scalar_t, 512><<<grid_size, block_size, 0, stream>>>(
2051 q_out.mutable_data_ptr<scalar_t>(), ql_nope.const_data_ptr<scalar_t>(),
2052 q_pe.const_data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),
2053 q_out.stride(1), ql_nope.stride(0), ql_nope.stride(1), q_pe.stride(0),
2054 q_pe.stride(1));
2055 });
2056 }