| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | #pragma once |
| 4 | |
| 5 | #include <cuda_runtime_api.h> |
| 6 | #include <cstdint> |
| 7 | |
| 8 | // Launchers for the nvfp4_ds_mla KV cache kernels. Those kernels live in |
| 9 | // their own TU (nvfp4_ds_mla_cache_kernels.cu) so that only that TU is built |
| 10 | // for the arch-conditional Blackwell targets their e2m1 conversions require; |
| 11 | // these declarations are what keeps cache_kernels.cu generic. |
| 12 | // |
| 13 | // Only defined when the build has an SM100 target (-DENABLE_NVFP4_SM100=1); |
| 14 | // callers must guard on the same macro. |
| 15 | namespace vllm { |
| 16 | |
| 17 | // concat_and_cache_mla for the nvfp4_ds_mla layout. |
| 18 | // kv_c: bf16 [num_tokens, 512] |
| 19 | // k_pe: bf16 [num_tokens, 64] |
| 20 | // kv_cache: uint8 [num_blocks, block_size, 352] |
| 21 | void launch_concat_and_cache_nvfp4_ds_mla(const void* kv_c, const void* k_pe, |
| 22 | void* kv_cache, |
| 23 | const int64_t* slot_mapping, |
| 24 | int block_stride, int entry_stride, |
| 25 | int kv_c_stride, int k_pe_stride, |
| 26 | int block_size, int num_tokens, |
| 27 | cudaStream_t stream); |
| 28 | |
| 29 | // Gather an nvfp4_ds_mla cache into a bf16 [total_tokens, 576] workspace. |
| 30 | void launch_cp_gather_and_upconvert_nvfp4_kv_cache( |
| 31 | const uint8_t* src_cache, void* dst, const int32_t* block_table, |
| 32 | const int32_t* workspace_starts, int32_t num_reqs, int32_t block_size, |
| 33 | int32_t total_tokens, int64_t block_table_stride, |
| 34 | int64_t cache_block_stride, int64_t cache_entry_stride, |
| 35 | int64_t dst_entry_stride, cudaStream_t stream); |
| 36 | |
| 37 | } // namespace vllm |