diff --git a/CMakeLists.txt b/CMakeLists.txt index 706f390948f..5a3f95faee0 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -305,10 +305,6 @@ endif() # set(VLLM_EXT_SRC - "csrc/cache_kernels.cu" - "csrc/cache_kernels_fused.cu" - "csrc/attention/paged_attention_v1.cu" - "csrc/attention/paged_attention_v2.cu" "csrc/cuda_view.cu" "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu" "csrc/quantization/activation_kernels.cu" @@ -647,7 +643,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/attention/merge_attn_states.cu" "csrc/libtorch_stable/sampler.cu" "csrc/libtorch_stable/topk.cu" - "csrc/libtorch_stable/mamba/selective_scan_fwd.cu") + "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" + "csrc/libtorch_stable/attention/paged_attention_v1.cu" + "csrc/libtorch_stable/attention/paged_attention_v2.cu" + "csrc/libtorch_stable/cache_kernels.cu" + "csrc/libtorch_stable/cache_kernels_fused.cu") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_STABLE_EXT_SRC @@ -924,13 +924,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") SRCS "${SRCS}" CUDA_ARCHS "${FP4_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") - # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly - # from cache_kernels.cu, so it belongs in _C rather than _C_stable. - set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu") + set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") set_gencode_flags_for_srcs( SRCS "${NVFP4_KV_SRC}" CUDA_ARCHS "${FP4_ARCHS}") - target_sources(_C PRIVATE ${NVFP4_KV_SRC}) + list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}") target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1) list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1") @@ -960,11 +958,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") SRCS "${SRCS}" CUDA_ARCHS "${FP4_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") - set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu") + set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") set_gencode_flags_for_srcs( SRCS "${NVFP4_KV_SRC}" CUDA_ARCHS "${FP4_ARCHS}") - target_sources(_C PRIVATE ${NVFP4_KV_SRC}) + list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}") target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1) list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1") diff --git a/csrc/libtorch_stable/activation_kernels.cu b/csrc/libtorch_stable/activation_kernels.cu index 28fdce5c305..cdab456348e 100644 --- a/csrc/libtorch_stable/activation_kernels.cu +++ b/csrc/libtorch_stable/activation_kernels.cu @@ -4,7 +4,7 @@ #include #include "../cuda_compat.h" -#include "../cuda_vec_utils.cuh" +#include "cuda_vec_utils.cuh" #include "dispatch_utils.h" #include "torch_utils.h" diff --git a/csrc/attention/attention_kernels.cuh b/csrc/libtorch_stable/attention/attention_kernels.cuh similarity index 99% rename from csrc/attention/attention_kernels.cuh rename to csrc/libtorch_stable/attention/attention_kernels.cuh index 052ff168cec..c5f9a9876c3 100644 --- a/csrc/attention/attention_kernels.cuh +++ b/csrc/libtorch_stable/attention/attention_kernels.cuh @@ -17,21 +17,18 @@ * limitations under the License. */ -#include -#include -#include #include -#include "attention_dtypes.h" +#include "../../attention/attention_dtypes.h" #include "attention_utils.cuh" -#include "../cuda_compat.h" +#include "../../cuda_compat.h" #ifdef USE_ROCM #include - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "../../quantization/w8a8/fp8/amd/quant_utils.cuh" typedef __hip_bfloat16 __nv_bfloat16; #else - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "../../quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif #define MAX(a, b) ((a) > (b) ? (a) : (b)) diff --git a/csrc/attention/attention_utils.cuh b/csrc/libtorch_stable/attention/attention_utils.cuh similarity index 95% rename from csrc/attention/attention_utils.cuh rename to csrc/libtorch_stable/attention/attention_utils.cuh index 826b0edffae..29783cb0e5f 100644 --- a/csrc/attention/attention_utils.cuh +++ b/csrc/libtorch_stable/attention/attention_utils.cuh @@ -18,8 +18,8 @@ */ #pragma once -#include "../cuda_compat.h" -#include "attention_dtypes.h" +#include "../../cuda_compat.h" +#include "../../attention/attention_dtypes.h" #include #include diff --git a/csrc/libtorch_stable/attention/merge_attn_states.cu b/csrc/libtorch_stable/attention/merge_attn_states.cu index 88e7723ad2f..b132e82e253 100644 --- a/csrc/libtorch_stable/attention/merge_attn_states.cu +++ b/csrc/libtorch_stable/attention/merge_attn_states.cu @@ -7,7 +7,7 @@ #include #include "../../attention/attention_dtypes.h" -#include "../../attention/attention_utils.cuh" +#include "attention_utils.cuh" #include "../../quantization/w8a8/fp8/common.cuh" namespace vllm { diff --git a/csrc/attention/paged_attention_v1.cu b/csrc/libtorch_stable/attention/paged_attention_v1.cu similarity index 80% rename from csrc/attention/paged_attention_v1.cu rename to csrc/libtorch_stable/attention/paged_attention_v1.cu index 307300e5566..8fa41791593 100644 --- a/csrc/attention/paged_attention_v1.cu +++ b/csrc/libtorch_stable/attention/paged_attention_v1.cu @@ -16,8 +16,9 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +#include "../torch_utils.h" #include "attention_kernels.cuh" -#include "../cuda_compat.h" +#include "../../cuda_compat.h" #define MAX(a, b) ((a) > (b) ? (a) : (b)) #define MIN(a, b) ((a) < (b) ? (a) : (b)) @@ -44,13 +45,15 @@ template void paged_attention_v1_launcher( - torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache, - torch::Tensor& value_cache, int num_kv_heads, float scale, - torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len, - const std::optional& alibi_slopes, torch::Tensor& k_scale, - torch::Tensor& v_scale, const int tp_rank, - const int blocksparse_local_blocks, const int blocksparse_vert_stride, - const int blocksparse_block_size, const int blocksparse_head_sliding_step) { + torch::stable::Tensor& out, torch::stable::Tensor& query, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + int num_kv_heads, float scale, torch::stable::Tensor& block_tables, + torch::stable::Tensor& seq_lens, int max_seq_len, + const std::optional& alibi_slopes, + torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale, + const int tp_rank, const int blocksparse_local_blocks, + const int blocksparse_vert_stride, const int blocksparse_block_size, + const int blocksparse_head_sliding_step) { int num_seqs = query.size(0); int num_heads = query.size(1); int head_size = query.size(2); @@ -69,8 +72,8 @@ void paged_attention_v1_launcher( T* query_ptr = reinterpret_cast(query.data_ptr()); CACHE_T* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); CACHE_T* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); - int* block_tables_ptr = block_tables.data_ptr(); - int* seq_lens_ptr = seq_lens.data_ptr(); + int* block_tables_ptr = block_tables.mutable_data_ptr(); + int* seq_lens_ptr = seq_lens.mutable_data_ptr(); const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); @@ -85,8 +88,9 @@ void paged_attention_v1_launcher( dim3 grid(num_heads, num_seqs, 1); dim3 block(NUM_THREADS); - const at::cuda::OptionalCUDAGuard device_guard(device_of(query)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + query.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); switch (head_size) { // NOTE(woosuk): To reduce the compilation time, we only compile for the // head sizes that we use in the model. However, we can easily extend this @@ -119,7 +123,7 @@ void paged_attention_v1_launcher( LAUNCH_PAGED_ATTENTION_V1(256); break; default: - TORCH_CHECK(false, "Unsupported head size: ", head_size); + STD_TORCH_CHECK(false, "Unsupported head size: ", head_size); break; } } @@ -141,43 +145,43 @@ void paged_attention_v1_launcher( // NOTE(woosuk): To reduce the compilation time, we omitted block sizes // 1, 2, 4, 64, 128, 256. -#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ - switch (block_size) { \ - case 8: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ - break; \ - case 16: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ - break; \ - case 32: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported block size: ", block_size); \ - break; \ +#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ + switch (block_size) { \ + case 8: \ + CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ + break; \ + case 16: \ + CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ + break; \ + case 32: \ + CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ + break; \ + default: \ + STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \ + break; \ } void paged_attention_v1( - torch::Tensor& out, // [num_seqs, num_heads, head_size] - torch::Tensor& query, // [num_seqs, num_heads, head_size] - torch::Tensor& + torch::stable::Tensor& out, // [num_seqs, num_heads, head_size] + torch::stable::Tensor& query, // [num_seqs, num_heads, head_size] + torch::stable::Tensor& key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] - torch::Tensor& + torch::stable::Tensor& value_cache, // [num_blocks, num_heads, head_size, block_size] int64_t num_kv_heads, // [num_heads] double scale, - torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] - torch::Tensor& seq_lens, // [num_seqs] + torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] + torch::stable::Tensor& seq_lens, // [num_seqs] int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::Tensor& k_scale, - torch::Tensor& v_scale, const int64_t tp_rank, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, const int64_t blocksparse_local_blocks, const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, const int64_t blocksparse_head_sliding_step) { const bool is_block_sparse = (blocksparse_vert_stride > 1); - DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype, CALL_V1_LAUNCHER_BLOCK_SIZE) } diff --git a/csrc/attention/paged_attention_v2.cu b/csrc/libtorch_stable/attention/paged_attention_v2.cu similarity index 78% rename from csrc/attention/paged_attention_v2.cu rename to csrc/libtorch_stable/attention/paged_attention_v2.cu index eb9b4feb4a8..4e8e56ae05c 100644 --- a/csrc/attention/paged_attention_v2.cu +++ b/csrc/libtorch_stable/attention/paged_attention_v2.cu @@ -16,8 +16,9 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +#include "../torch_utils.h" #include "attention_kernels.cuh" -#include "../cuda_compat.h" +#include "../../cuda_compat.h" #define MAX(a, b) ((a) > (b) ? (a) : (b)) #define MIN(a, b) ((a) < (b) ? (a) : (b)) @@ -44,14 +45,16 @@ template void paged_attention_v2_launcher( - torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, - torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, - torch::Tensor& value_cache, int num_kv_heads, float scale, - torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len, - const std::optional& alibi_slopes, torch::Tensor& k_scale, - torch::Tensor& v_scale, const int tp_rank, - const int blocksparse_local_blocks, const int blocksparse_vert_stride, - const int blocksparse_block_size, const int blocksparse_head_sliding_step) { + torch::stable::Tensor& out, torch::stable::Tensor& exp_sums, + torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out, + torch::stable::Tensor& query, torch::stable::Tensor& key_cache, + torch::stable::Tensor& value_cache, int num_kv_heads, float scale, + torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens, + int max_seq_len, const std::optional& alibi_slopes, + torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale, + const int tp_rank, const int blocksparse_local_blocks, + const int blocksparse_vert_stride, const int blocksparse_block_size, + const int blocksparse_head_sliding_step) { int num_seqs = query.size(0); int num_heads = query.size(1); int head_size = query.size(2); @@ -73,8 +76,8 @@ void paged_attention_v2_launcher( T* query_ptr = reinterpret_cast(query.data_ptr()); CACHE_T* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); CACHE_T* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); - int* block_tables_ptr = block_tables.data_ptr(); - int* seq_lens_ptr = seq_lens.data_ptr(); + int* block_tables_ptr = block_tables.mutable_data_ptr(); + int* seq_lens_ptr = seq_lens.mutable_data_ptr(); const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); @@ -91,8 +94,9 @@ void paged_attention_v2_launcher( int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float); dim3 block(NUM_THREADS); - const at::cuda::OptionalCUDAGuard device_guard(device_of(query)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + query.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); switch (head_size) { // NOTE(woosuk): To reduce the compilation time, we only compile for the // head sizes that we use in the model. However, we can easily extend this @@ -125,7 +129,7 @@ void paged_attention_v2_launcher( LAUNCH_PAGED_ATTENTION_V2(256); break; default: - TORCH_CHECK(false, "Unsupported head size: ", head_size); + STD_TORCH_CHECK(false, "Unsupported head size: ", head_size); break; } } @@ -148,46 +152,48 @@ void paged_attention_v2_launcher( // NOTE(woosuk): To reduce the compilation time, we omitted block sizes // 1, 2, 4, 64, 128, 256. -#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ - switch (block_size) { \ - case 8: \ - CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ - break; \ - case 16: \ - CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ - break; \ - case 32: \ - CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported block size: ", block_size); \ - break; \ +#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ + switch (block_size) { \ + case 8: \ + CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ + break; \ + case 16: \ + CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ + break; \ + case 32: \ + CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ + break; \ + default: \ + STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \ + break; \ } void paged_attention_v2( - torch::Tensor& out, // [num_seqs, num_heads, head_size] - torch::Tensor& exp_sums, // [num_seqs, num_heads, max_num_partitions] - torch::Tensor& max_logits, // [num_seqs, num_heads, max_num_partitions] - torch::Tensor& + torch::stable::Tensor& out, // [num_seqs, num_heads, head_size] + torch::stable::Tensor& + exp_sums, // [num_seqs, num_heads, max_num_partitions] + torch::stable::Tensor& + max_logits, // [num_seqs, num_heads, max_num_partitions] + torch::stable::Tensor& tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size] - torch::Tensor& query, // [num_seqs, num_heads, head_size] - torch::Tensor& + torch::stable::Tensor& query, // [num_seqs, num_heads, head_size] + torch::stable::Tensor& key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] - torch::Tensor& + torch::stable::Tensor& value_cache, // [num_blocks, num_heads, head_size, block_size] int64_t num_kv_heads, // [num_heads] double scale, - torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] - torch::Tensor& seq_lens, // [num_seqs] + torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] + torch::stable::Tensor& seq_lens, // [num_seqs] int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::Tensor& k_scale, - torch::Tensor& v_scale, const int64_t tp_rank, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, const int64_t blocksparse_local_blocks, const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, const int64_t blocksparse_head_sliding_step) { const bool is_block_sparse = (blocksparse_vert_stride > 1); - DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype, CALL_V2_LAUNCHER_BLOCK_SIZE) } diff --git a/csrc/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu similarity index 73% rename from csrc/cache_kernels.cu rename to csrc/libtorch_stable/cache_kernels.cu index 9130dd2ccae..eac93ac9a9f 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -1,20 +1,16 @@ -#include -#include -#include -#include -#include - -#include "cuda_utils.h" -#include "cuda_compat.h" +#include "torch_utils.h" #include "dispatch_utils.h" -#include "libtorch_stable/quantization/vectorization_utils.cuh" +#include "../cuda_utils.h" +#include "../cuda_compat.h" + +#include "quantization/vectorization_utils.cuh" #include "concat_mla_q.cuh" #ifdef USE_ROCM - #include "quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif #include @@ -34,40 +30,45 @@ constexpr float kFp8ScaleDivisor = 224.f; constexpr float kFp8ScaleDivisor = 448.f; #endif -void swap_blocks(torch::Tensor& src, torch::Tensor& dst, +void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, int64_t block_size_in_bytes, - const torch::Tensor& block_mapping) { - torch::Device src_device = src.device(); - torch::Device dst_device = dst.device(); + const torch::stable::Tensor& block_mapping) { + torch::stable::Device src_device = src.device(); + torch::stable::Device dst_device = dst.device(); cudaMemcpyKind memcpy_type; if (src_device.is_cuda() && dst_device.is_cuda()) { - TORCH_CHECK(src_device.index() == dst_device.index(), - "src and dst must be on the same GPU"); + STD_TORCH_CHECK(src_device.index() == dst_device.index(), + "src and dst must be on the same GPU"); memcpy_type = cudaMemcpyDeviceToDevice; } else if (src_device.is_cuda() && dst_device.is_cpu()) { memcpy_type = cudaMemcpyDeviceToHost; } else if (src_device.is_cpu() && dst_device.is_cuda()) { memcpy_type = cudaMemcpyHostToDevice; } else { - TORCH_CHECK(false, "Invalid device combination"); + STD_TORCH_CHECK(false, "Invalid device combination"); } // NOTE(youkaichao): keep in mind that `block_mapping` should be // a cpu tensor, otherwise every `item` call will require a gpu-cpu // synchronization. - TORCH_CHECK(block_mapping.device().is_cpu(), "block_mapping must be on CPU"); + STD_TORCH_CHECK(block_mapping.device().is_cpu(), + "block_mapping must be on CPU"); char* src_ptr = static_cast(src.data_ptr()); char* dst_ptr = static_cast(dst.data_ptr()); - const at::cuda::OptionalCUDAGuard device_guard( - src_device.is_cuda() ? src_device : dst_device); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + auto guard_device = src_device.is_cuda() ? src_device : dst_device; + const torch::stable::accelerator::DeviceGuard device_guard( + guard_device.index()); + const cudaStream_t stream = get_current_cuda_stream(); // NOTE(woosuk): This can be slow if the number of blocks is large. const int64_t num_blocks = block_mapping.size(0); + const int64_t* bm_ptr = block_mapping.const_data_ptr(); + const int64_t bm_stride0 = block_mapping.stride(0); + const int64_t bm_stride1 = block_mapping.stride(1); for (size_t i = 0; i < num_blocks; i++) { - int64_t src_block_number = block_mapping[i][0].item(); - int64_t dst_block_number = block_mapping[i][1].item(); + int64_t src_block_number = bm_ptr[i * bm_stride0]; + int64_t dst_block_number = bm_ptr[i * bm_stride0 + bm_stride1]; int64_t src_offset = src_block_number * block_size_in_bytes; int64_t dst_offset = dst_block_number * block_size_in_bytes; cudaMemcpyAsync(dst_ptr + dst_offset, src_ptr + src_offset, @@ -75,20 +76,23 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst, } } -void swap_blocks_batch(const torch::Tensor& src_ptrs, - const torch::Tensor& dst_ptrs, - const torch::Tensor& sizes, +void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, + const torch::stable::Tensor& dst_ptrs, + const torch::stable::Tensor& sizes, bool is_src_access_order_any) { - TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU"); - TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU"); - TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU"); - TORCH_CHECK(src_ptrs.dtype() == torch::kInt64, "src_ptrs must be int64"); - TORCH_CHECK(dst_ptrs.dtype() == torch::kInt64, "dst_ptrs must be int64"); - TORCH_CHECK(sizes.dtype() == torch::kInt64, "sizes must be int64"); + STD_TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU"); + STD_TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU"); + STD_TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU"); + STD_TORCH_CHECK(src_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, + "src_ptrs must be int64"); + STD_TORCH_CHECK(dst_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, + "dst_ptrs must be int64"); + STD_TORCH_CHECK(sizes.scalar_type() == torch::headeronly::ScalarType::Long, + "sizes must be int64"); const int64_t n = src_ptrs.size(0); - TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs"); - TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs"); + STD_TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs"); + STD_TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs"); if (n == 0) return; @@ -96,7 +100,7 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs, int64_t* dst_data = dst_ptrs.mutable_data_ptr(); int64_t* size_data = sizes.mutable_data_ptr(); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const cudaStream_t stream = get_current_cuda_stream(); // Use cuMemcpyBatchAsync / hipMemcpyBatchAsync to submit all copies in a // single driver call, amortizing per-copy submission overhead. int64_t @@ -138,8 +142,9 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs, reinterpret_cast(size_data), static_cast(n), &attr, &attrs_idx, 1, &fail_idx, static_cast(stream)); - TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ", - fail_idx, " with error ", result); + STD_TORCH_CHECK(result == CUDA_SUCCESS, + "cuMemcpyBatchAsync failed at index ", fail_idx, + " with error ", result); return; } #elif defined(USE_ROCM) && defined(HIP_VERSION) && HIP_VERSION >= 70100000 @@ -155,8 +160,9 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs, reinterpret_cast(dst_data), reinterpret_cast(src_data), reinterpret_cast(size_data), static_cast(n), &attr, &attrs_idx, 0, &fail_idx, static_cast(stream)); - TORCH_CHECK(result == hipSuccess, "hipMemcpyBatchAsync failed at index ", - fail_idx, " with error ", result); + STD_TORCH_CHECK(result == hipSuccess, + "hipMemcpyBatchAsync failed at index ", fail_idx, + " with error ", result); return; } #endif @@ -675,28 +681,28 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel( // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. // KV_DTYPE is the real data type of kv-cache. -#define CALL_RESHAPE_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \ - vllm::reshape_and_cache_kernel \ - <<>>( \ - reinterpret_cast(key.data_ptr()), \ - reinterpret_cast(value.data_ptr()), \ - reinterpret_cast(key_cache.data_ptr()), \ - reinterpret_cast(value_cache.data_ptr()), \ - slot_mapping.data_ptr(), key_stride, value_stride, \ - num_heads, head_size, block_size, x, \ - reinterpret_cast(k_scale.data_ptr()), \ +#define CALL_RESHAPE_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \ + vllm::reshape_and_cache_kernel \ + <<>>( \ + reinterpret_cast(key.data_ptr()), \ + reinterpret_cast(value.data_ptr()), \ + reinterpret_cast(key_cache.data_ptr()), \ + reinterpret_cast(value_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), key_stride, value_stride, \ + num_heads, head_size, block_size, x, \ + reinterpret_cast(k_scale.data_ptr()), \ reinterpret_cast(v_scale.data_ptr())); void reshape_and_cache( - torch::Tensor& key, // [num_tokens, num_heads, head_size] - torch::Tensor& value, // [num_tokens, num_heads, head_size] - torch::Tensor& + torch::stable::Tensor& key, // [num_tokens, num_heads, head_size] + torch::stable::Tensor& value, // [num_tokens, num_heads, head_size] + torch::stable::Tensor& key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] - torch::Tensor& + torch::stable::Tensor& value_cache, // [num_blocks, num_heads, head_size, block_size] - torch::Tensor& slot_mapping, // [num_tokens] - const std::string& kv_cache_dtype, torch::Tensor& k_scale, - torch::Tensor& v_scale) { + torch::stable::Tensor& slot_mapping, // [num_tokens] + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale) { int num_tokens = slot_mapping.size(0); int num_heads = key.size(1); int head_size = key.size(2); @@ -709,39 +715,41 @@ void reshape_and_cache( dim3 grid(num_tokens); dim3 block(std::min(num_heads * head_div_x, 512)); - const at::cuda::OptionalCUDAGuard device_guard(device_of(key)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + key.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); - DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, CALL_RESHAPE_AND_CACHE); } // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. // KV_DTYPE is the real data type of kv-cache. -#define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE) \ - vllm::reshape_and_cache_flash_kernel \ - <<>>( \ - reinterpret_cast(key.data_ptr()), \ - reinterpret_cast(value.data_ptr()), \ - reinterpret_cast(key_cache.data_ptr()), \ - reinterpret_cast(value_cache.data_ptr()), \ - slot_mapping.data_ptr(), block_stride, page_stride, \ - head_stride, key_stride, value_stride, num_heads, head_size, \ - block_size, reinterpret_cast(k_scale.data_ptr()), \ - reinterpret_cast(v_scale.data_ptr()), \ +#define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE) \ + vllm::reshape_and_cache_flash_kernel \ + <<>>( \ + reinterpret_cast(key.data_ptr()), \ + reinterpret_cast(value.data_ptr()), \ + reinterpret_cast(key_cache.data_ptr()), \ + reinterpret_cast(value_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), block_stride, page_stride, \ + head_stride, key_stride, value_stride, num_heads, head_size, \ + block_size, reinterpret_cast(k_scale.data_ptr()), \ + reinterpret_cast(v_scale.data_ptr()), \ kv_scale_stride); void reshape_and_cache_flash( - torch::Tensor& key, // [num_tokens, num_heads, head_size] - torch::Tensor& value, // [num_tokens, num_heads, head_size] - torch::Tensor& key_cache, // [num_blocks, block_size, num_heads, head_size] - torch::Tensor& + torch::stable::Tensor& key, // [num_tokens, num_heads, head_size] + torch::stable::Tensor& value, // [num_tokens, num_heads, head_size] + torch::stable::Tensor& + key_cache, // [num_blocks, block_size, num_heads, head_size] + torch::stable::Tensor& value_cache, // [num_blocks, block_size, num_heads, head_size] - torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] + torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] const std::string& kv_cache_dtype, - torch::Tensor& k_scale, // [1] or [num_heads] - torch::Tensor& v_scale) { // [1] or [num_heads] + torch::stable::Tensor& k_scale, // [1] or [num_heads] + torch::stable::Tensor& v_scale) { // [1] or [num_heads] // NOTE(woosuk): In vLLM V1, key.size(0) can be different from // slot_mapping.size(0) because of padding for CUDA graphs. // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because @@ -756,23 +764,26 @@ void reshape_and_cache_flash( int num_heads = key.size(1); int head_size = key.size(2); - const at::cuda::OptionalCUDAGuard device_guard(device_of(key)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + key.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); if (kv_cache_dtype == "nvfp4") { #if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120) // NVFP4 dispatch is compiled separately for SM100+. extern void reshape_and_cache_nvfp4_dispatch( - torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache, - torch::Tensor & value_cache, torch::Tensor & slot_mapping, - torch::Tensor & k_scale, torch::Tensor & v_scale); + torch::stable::Tensor & key, torch::stable::Tensor & value, + torch::stable::Tensor & key_cache, torch::stable::Tensor & value_cache, + torch::stable::Tensor & slot_mapping, torch::stable::Tensor & k_scale, + torch::stable::Tensor & v_scale); reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache, slot_mapping, k_scale, v_scale); return; #else - TORCH_CHECK(false, - "NVFP4 KV cache requires SM100+ (Blackwell). " - "Please rebuild vllm with a Blackwell-compatible CUDA target."); + STD_TORCH_CHECK( + false, + "NVFP4 KV cache requires SM100+ (Blackwell). " + "Please rebuild vllm with a Blackwell-compatible CUDA target."); #endif } @@ -784,53 +795,53 @@ void reshape_and_cache_flash( int64_t block_stride = key_cache.stride(0); int64_t page_stride = key_cache.stride(1); int64_t head_stride = key_cache.stride(2); - TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0)); + STD_TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0)); - TORCH_CHECK(k_scale.sizes() == v_scale.sizes(), - "k_scale and v_scale must have the same shape"); - TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads, - "k_scale and v_scale must be of shape [1] or [num_heads]"); + STD_TORCH_CHECK(k_scale.sizes().equals(v_scale.sizes()), + "k_scale and v_scale must have the same shape"); + STD_TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads, + "k_scale and v_scale must be of shape [1] or [num_heads]"); int kv_scale_stride = (k_scale.numel() > 1) ? 1 : 0; dim3 grid(num_tokens); dim3 block(std::min(num_heads * head_size, 512)); - DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, CALL_RESHAPE_AND_CACHE_FLASH); } // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. // KV_DTYPE is the real data type of kv-cache. -#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \ - vllm::concat_and_cache_mla_kernel \ - <<>>( \ - reinterpret_cast(kv_c.data_ptr()), \ - reinterpret_cast(k_pe.data_ptr()), \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.data_ptr(), block_stride, entry_stride, \ - kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ +#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \ + vllm::concat_and_cache_mla_kernel \ + <<>>( \ + reinterpret_cast(kv_c.data_ptr()), \ + reinterpret_cast(k_pe.data_ptr()), \ + reinterpret_cast(kv_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), block_stride, entry_stride, \ + kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ reinterpret_cast(scale.data_ptr())); // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. -#define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \ - vllm::concat_and_cache_ds_mla_kernel \ - <<>>( \ - reinterpret_cast(kv_c.data_ptr()), \ - reinterpret_cast(k_pe.data_ptr()), \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.data_ptr(), block_stride, entry_stride, \ - kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ +#define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \ + vllm::concat_and_cache_ds_mla_kernel \ + <<>>( \ + reinterpret_cast(kv_c.data_ptr()), \ + reinterpret_cast(k_pe.data_ptr()), \ + reinterpret_cast(kv_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), block_stride, entry_stride, \ + kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ reinterpret_cast(scale.data_ptr())); void concat_and_cache_mla( - torch::Tensor& kv_c, // [num_tokens, kv_lora_rank] - torch::Tensor& k_pe, // [num_tokens, pe_dim] - torch::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank + - // pe_dim)] - torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] - const std::string& kv_cache_dtype, torch::Tensor& scale) { + torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank] + torch::stable::Tensor& k_pe, // [num_tokens, pe_dim] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank + // + pe_dim)] + torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] + const std::string& kv_cache_dtype, torch::stable::Tensor& scale) { // NOTE(woosuk): In vLLM V1, key.size(0) can be different from // slot_mapping.size(0) because of padding for CUDA graphs. // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because @@ -847,16 +858,17 @@ void concat_and_cache_mla( int block_size = kv_cache.size(1); if (kv_cache_dtype == "fp8_ds_mla") { - TORCH_CHECK(kv_lora_rank == 512, "kv_lora_rank must be 512 for fp8_ds_mla"); - TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for fp8_ds_mla"); - TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.itemsize(), - "kv_cache.size(2) must be 656 bytes for fp8_ds_mla"); - TORCH_CHECK(kv_c.itemsize() == 2, - "kv_c.itemsize() must be 2 for fp8_ds_mla"); - TORCH_CHECK(k_pe.itemsize() == 2, - "k_pe.itemsize() must be 2 for fp8_ds_mla"); + STD_TORCH_CHECK(kv_lora_rank == 512, + "kv_lora_rank must be 512 for fp8_ds_mla"); + STD_TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for fp8_ds_mla"); + STD_TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.element_size(), + "kv_cache.size(2) must be 656 bytes for fp8_ds_mla"); + STD_TORCH_CHECK(kv_c.element_size() == 2, + "kv_c.element_size() must be 2 for fp8_ds_mla"); + STD_TORCH_CHECK(k_pe.element_size() == 2, + "k_pe.element_size() must be 2 for fp8_ds_mla"); } else { - TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim); + STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim); } int kv_c_stride = kv_c.stride(0); @@ -864,8 +876,9 @@ void concat_and_cache_mla( int block_stride = kv_cache.stride(0); int entry_stride = kv_cache.stride(1); - const at::cuda::OptionalCUDAGuard device_guard(device_of(kv_c)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + kv_c.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); if (kv_cache_dtype == "fp8_ds_mla") { dim3 grid(num_tokens); @@ -875,12 +888,12 @@ void concat_and_cache_mla( // The RoPE part (last 64 elements) is handled by another 1 warp (32 // threads). So in total, we use 3 warps (96 threads) per block. dim3 block(96); - DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, CALL_CONCAT_AND_CACHE_DS_MLA); } else { dim3 grid(num_tokens); dim3 block(std::min(kv_lora_rank, 512)); - DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, CALL_CONCAT_AND_CACHE_MLA); } } @@ -908,55 +921,62 @@ __global__ void convert_fp8_kernel(const Tin* __restrict__ src_cache, reinterpret_cast(dst_cache.data_ptr()), scale, block_stride); // Only for testing. -void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache, - const double scale, const std::string& kv_cache_dtype) { - torch::Device src_device = src_cache.device(); - torch::Device dst_device = dst_cache.device(); - TORCH_CHECK(src_device.is_cuda(), "src must be on a GPU") - TORCH_CHECK(dst_device.is_cuda(), "dst must be on a GPU") - TORCH_CHECK(src_device.index() == dst_device.index(), - "src and dst must be on the same GPU"); - at::cuda::OptionalCUDAGuard device_guard(src_device); +void convert_fp8(torch::stable::Tensor& dst_cache, + torch::stable::Tensor& src_cache, const double scale, + const std::string& kv_cache_dtype) { + torch::stable::Device src_device = src_cache.device(); + torch::stable::Device dst_device = dst_cache.device(); + STD_TORCH_CHECK(src_device.is_cuda(), "src must be on a GPU") + STD_TORCH_CHECK(dst_device.is_cuda(), "dst must be on a GPU") + STD_TORCH_CHECK(src_device.index() == dst_device.index(), + "src and dst must be on the same GPU"); + torch::stable::accelerator::DeviceGuard device_guard(src_device.index()); int64_t num_blocks = src_cache.size(0); int64_t block_stride = src_cache.stride(0); dim3 grid(num_blocks); dim3 block(std::min(block_stride, int64_t(512))); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const cudaStream_t stream = get_current_cuda_stream(); if (kv_cache_dtype == "auto") { - if (src_cache.dtype() == at::ScalarType::Float) { + if (src_cache.scalar_type() == torch::headeronly::ScalarType::Float) { CALL_CONVERT_FP8(uint8_t, float, vllm::Fp8KVCacheDataType::kAuto); - } else if (src_cache.dtype() == at::ScalarType::Half) { + } else if (src_cache.scalar_type() == torch::headeronly::ScalarType::Half) { CALL_CONVERT_FP8(uint8_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); - } else if (src_cache.dtype() == at::ScalarType::BFloat16) { + } else if (src_cache.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { CALL_CONVERT_FP8(uint8_t, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); - } else if (dst_cache.dtype() == at::ScalarType::Float) { + } else if (dst_cache.scalar_type() == + torch::headeronly::ScalarType::Float) { CALL_CONVERT_FP8(float, uint8_t, vllm::Fp8KVCacheDataType::kAuto); - } else if (dst_cache.dtype() == at::ScalarType::Half) { + } else if (dst_cache.scalar_type() == torch::headeronly::ScalarType::Half) { CALL_CONVERT_FP8(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kAuto); - } else if (dst_cache.dtype() == at::ScalarType::BFloat16) { + } else if (dst_cache.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { CALL_CONVERT_FP8(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kAuto); } } else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") { - if (src_cache.dtype() == at::ScalarType::Float) { + if (src_cache.scalar_type() == torch::headeronly::ScalarType::Float) { CALL_CONVERT_FP8(uint8_t, float, vllm::Fp8KVCacheDataType::kFp8E4M3); - } else if (src_cache.dtype() == at::ScalarType::Half) { + } else if (src_cache.scalar_type() == torch::headeronly::ScalarType::Half) { CALL_CONVERT_FP8(uint8_t, uint16_t, vllm::Fp8KVCacheDataType::kFp8E4M3); - } else if (src_cache.dtype() == at::ScalarType::BFloat16) { + } else if (src_cache.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { CALL_CONVERT_FP8(uint8_t, __nv_bfloat16, vllm::Fp8KVCacheDataType::kFp8E4M3); - } else if (dst_cache.dtype() == at::ScalarType::Float) { + } else if (dst_cache.scalar_type() == + torch::headeronly::ScalarType::Float) { CALL_CONVERT_FP8(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); - } else if (dst_cache.dtype() == at::ScalarType::Half) { + } else if (dst_cache.scalar_type() == torch::headeronly::ScalarType::Half) { CALL_CONVERT_FP8(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); - } else if (dst_cache.dtype() == at::ScalarType::BFloat16) { + } else if (dst_cache.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { CALL_CONVERT_FP8(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); } } else { - TORCH_CHECK(false, "Unsupported data type: ", kv_cache_dtype); + STD_TORCH_CHECK(false, "Unsupported data type: ", kv_cache_dtype); } } @@ -1053,8 +1073,9 @@ __global__ void gather_and_maybe_dequant_cache( <<>>( \ reinterpret_cast(src_cache.data_ptr()), \ reinterpret_cast(dst.data_ptr()), \ - block_table.data_ptr(), cu_seq_lens.data_ptr(), \ - token_to_seq.data_ptr(), num_tokens, block_size, \ + block_table.const_data_ptr(), \ + cu_seq_lens.const_data_ptr(), \ + token_to_seq.const_data_ptr(), num_tokens, block_size, \ block_table_stride, cache_block_stride, cache_entry_stride, \ dst_entry_stride, reinterpret_cast(scale.data_ptr()), \ seq_starts_ptr); @@ -1072,42 +1093,47 @@ __global__ void gather_and_maybe_dequant_cache( // - Optionally, seq_starts (if provided) offsets the starting block index by // (seq_starts[bid] / page_size) void gather_and_maybe_dequant_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] - torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& cu_seq_lens, // [BATCH+1] - torch::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] + torch::stable::Tensor const& + src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] int64_t num_tokens, const std::string& kv_cache_dtype, - torch::Tensor const& scale, - std::optional seq_starts = std::nullopt) { - at::cuda::OptionalCUDAGuard device_guard(src_cache.device()); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + torch::stable::Tensor const& scale, + std::optional seq_starts = std::nullopt) { + torch::stable::accelerator::DeviceGuard device_guard( + src_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); int32_t block_size = src_cache.size(1); int32_t head_dim = dst.size(-1); - TORCH_CHECK(block_table.dtype() == torch::kInt32, - "block_table must be int32"); - TORCH_CHECK(cu_seq_lens.dtype() == torch::kInt32, - "cu_seq_lens must be int32"); + STD_TORCH_CHECK( + block_table.scalar_type() == torch::headeronly::ScalarType::Int, + "block_table must be int32"); + STD_TORCH_CHECK( + cu_seq_lens.scalar_type() == torch::headeronly::ScalarType::Int, + "cu_seq_lens must be int32"); if (seq_starts.has_value()) { - TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32, - "seq_starts must be int32"); + STD_TORCH_CHECK( + seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int, + "seq_starts must be int32"); } - TORCH_CHECK( + STD_TORCH_CHECK( head_dim == 320 || head_dim == 576, "gather_and_maybe_dequant_cache only support the head_dim to 320 or 576 " "for better performance") - TORCH_CHECK(src_cache.device() == dst.device(), - "src_cache and dst must be on the same device"); - TORCH_CHECK(src_cache.device() == block_table.device(), - "src_cache and block_table must be on the same device"); - TORCH_CHECK(src_cache.device() == cu_seq_lens.device(), - "src_cache and cu_seq_lens must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == dst.device(), + "src_cache and dst must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == block_table.device(), + "src_cache and block_table must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == cu_seq_lens.device(), + "src_cache and cu_seq_lens must be on the same device"); if (seq_starts.has_value()) { - TORCH_CHECK(src_cache.device() == seq_starts.value().device(), - "src_cache and seq_starts must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), + "src_cache and seq_starts must be on the same device"); } int64_t block_table_stride = block_table.stride(0); @@ -1120,13 +1146,14 @@ void gather_and_maybe_dequant_cache( dim3 block(thread_block_size); const int32_t* seq_starts_ptr = - seq_starts.has_value() ? seq_starts.value().data_ptr() : nullptr; + seq_starts.has_value() ? seq_starts.value().const_data_ptr() + : nullptr; if (head_dim == 576) { - DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype, CALL_GATHER_CACHE_576); } else { - DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype, CALL_GATHER_CACHE_320); } } @@ -1267,13 +1294,14 @@ __global__ void cp_gather_cache( } // namespace vllm // Macro to dispatch the kernel based on the data type. -#define CALL_CP_GATHER_CACHE(CPY_DTYPE) \ - vllm::cp_gather_cache<<>>( \ - reinterpret_cast(src_cache.data_ptr()), \ - reinterpret_cast(dst.data_ptr()), \ - block_table.data_ptr(), cu_seq_lens.data_ptr(), \ - block_size, entry_size, block_table_stride, cache_block_stride, \ - cache_entry_stride, dst_entry_stride, seq_starts_ptr); +#define CALL_CP_GATHER_CACHE(CPY_DTYPE) \ + vllm::cp_gather_cache<<>>( \ + reinterpret_cast(src_cache.data_ptr()), \ + reinterpret_cast(dst.data_ptr()), \ + block_table.const_data_ptr(), \ + cu_seq_lens.const_data_ptr(), block_size, entry_size, \ + block_table_stride, cache_block_stride, cache_entry_stride, \ + dst_entry_stride, seq_starts_ptr); // Gather sequences from the cache into the destination tensor. // - cu_seq_lens contains the cumulative sequence lengths for each batch @@ -1281,36 +1309,41 @@ __global__ void cp_gather_cache( // - Optionally, seq_starts (if provided) offsets the starting slot index by // seq_starts[bid] void cp_gather_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] - torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& cu_seq_lens, // [BATCH+1] + torch::stable::Tensor const& + src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] int64_t batch_size, - std::optional seq_starts = std::nullopt) { - at::cuda::OptionalCUDAGuard device_guard(src_cache.device()); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + std::optional seq_starts = std::nullopt) { + torch::stable::accelerator::DeviceGuard device_guard( + src_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); int32_t block_size = src_cache.size(1); - int32_t entry_size = src_cache.flatten(2, -1).size(2); + int32_t entry_size = torch::stable::flatten(src_cache, 2, -1).size(2); - TORCH_CHECK(block_table.dtype() == torch::kInt32, - "block_table must be int32"); - TORCH_CHECK(cu_seq_lens.dtype() == torch::kInt32, - "cu_seq_lens must be int32"); + STD_TORCH_CHECK( + block_table.scalar_type() == torch::headeronly::ScalarType::Int, + "block_table must be int32"); + STD_TORCH_CHECK( + cu_seq_lens.scalar_type() == torch::headeronly::ScalarType::Int, + "cu_seq_lens must be int32"); if (seq_starts.has_value()) { - TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32, - "seq_starts must be int32"); + STD_TORCH_CHECK( + seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int, + "seq_starts must be int32"); } - TORCH_CHECK(src_cache.device() == dst.device(), - "src_cache and dst must be on the same device"); - TORCH_CHECK(src_cache.device() == block_table.device(), - "src_cache and block_table must be on the same device"); - TORCH_CHECK(src_cache.device() == cu_seq_lens.device(), - "src_cache and cu_seq_lens must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == dst.device(), + "src_cache and dst must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == block_table.device(), + "src_cache and block_table must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == cu_seq_lens.device(), + "src_cache and cu_seq_lens must be on the same device"); if (seq_starts.has_value()) { - TORCH_CHECK(src_cache.device() == seq_starts.value().device(), - "src_cache and seq_starts must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), + "src_cache and seq_starts must be on the same device"); } int64_t block_table_stride = block_table.stride(0); @@ -1323,12 +1356,13 @@ void cp_gather_cache( dim3 grid(batch_size, num_splits); dim3 block(1024); - TORCH_CHECK(src_cache.dtype() == dst.dtype(), - "src_cache and dst must have the same dtype"); + STD_TORCH_CHECK(src_cache.scalar_type() == dst.scalar_type(), + "src_cache and dst must have the same dtype"); const int dtype_bits = src_cache.element_size() * 8; const int32_t* seq_starts_ptr = - seq_starts.has_value() ? seq_starts.value().data_ptr() : nullptr; + seq_starts.has_value() ? seq_starts.value().const_data_ptr() + : nullptr; if (dtype_bits == 32) { CALL_CP_GATHER_CACHE(uint32_t); @@ -1337,46 +1371,51 @@ void cp_gather_cache( } else if (dtype_bits == 8) { CALL_CP_GATHER_CACHE(uint8_t); } else { - TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits); + STD_TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits); } } void cp_gather_and_upconvert_fp8_kv_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] - torch::Tensor const& dst, // [TOT_TOKENS, 576] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& seq_lens, // [BATCH] - torch::Tensor const& workspace_starts, // [BATCH] + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& seq_lens, // [BATCH] + torch::stable::Tensor const& workspace_starts, // [BATCH] int64_t batch_size) { - at::cuda::OptionalCUDAGuard device_guard(src_cache.device()); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + torch::stable::accelerator::DeviceGuard device_guard( + src_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); int32_t block_size = src_cache.size(1); int32_t head_dim = dst.size(1); - TORCH_CHECK(block_table.dtype() == torch::kInt32, - "block_table must be int32"); - TORCH_CHECK(seq_lens.dtype() == torch::kInt32, "seq_lens must be int32"); - TORCH_CHECK(workspace_starts.dtype() == torch::kInt32, - "workspace_starts must be int32"); + STD_TORCH_CHECK( + block_table.scalar_type() == torch::headeronly::ScalarType::Int, + "block_table must be int32"); + STD_TORCH_CHECK(seq_lens.scalar_type() == torch::headeronly::ScalarType::Int, + "seq_lens must be int32"); + STD_TORCH_CHECK( + workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int, + "workspace_starts must be int32"); - TORCH_CHECK(src_cache.device() == dst.device(), - "src_cache and dst must be on the same device"); - TORCH_CHECK(src_cache.device() == block_table.device(), - "src_cache and block_table must be on the same device"); - TORCH_CHECK(src_cache.device() == seq_lens.device(), - "src_cache and seq_lens must be on the same device"); - TORCH_CHECK(src_cache.device() == workspace_starts.device(), - "src_cache and workspace_starts must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == dst.device(), + "src_cache and dst must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == block_table.device(), + "src_cache and block_table must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == seq_lens.device(), + "src_cache and seq_lens must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(), + "src_cache and workspace_starts must be on the same device"); auto dtype = src_cache.scalar_type(); - TORCH_CHECK( - dtype == at::ScalarType::Byte || // uint8 - dtype == at::ScalarType::Float8_e4m3fn || // fp8 e4m3 - dtype == at::ScalarType::Float8_e5m2, // fp8 e5m2 + STD_TORCH_CHECK( + dtype == torch::headeronly::ScalarType::Byte || // uint8 + dtype == torch::headeronly::ScalarType::Float8_e4m3fn || // fp8 e4m3 + dtype == torch::headeronly::ScalarType::Float8_e5m2, // fp8 e5m2 "src_cache must be uint8, float8_e4m3fn, or float8_e5m2, but got ", - src_cache.dtype()); - TORCH_CHECK(dst.dtype() == torch::kBFloat16, "dst must be bfloat16"); - TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA"); + src_cache.scalar_type()); + STD_TORCH_CHECK(dst.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "dst must be bfloat16"); + STD_TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA"); int64_t block_table_stride = block_table.stride(0); int64_t cache_block_stride = src_cache.stride(0); @@ -1384,8 +1423,8 @@ void cp_gather_and_upconvert_fp8_kv_cache( int64_t dst_entry_stride = dst.stride(0); const uint8_t* src_ptr = nullptr; - if (dtype == at::ScalarType::Byte) { - src_ptr = src_cache.data_ptr(); + if (dtype == torch::headeronly::ScalarType::Byte) { + src_ptr = src_cache.const_data_ptr(); } else { // float8_e4m3fn or float8_e5m2 src_ptr = reinterpret_cast(src_cache.data_ptr()); @@ -1399,26 +1438,27 @@ void cp_gather_and_upconvert_fp8_kv_cache( vllm::cp_gather_and_upconvert_fp8_kv_cache<<>>( src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()), - block_table.data_ptr(), workspace_starts.data_ptr(), + block_table.const_data_ptr(), + workspace_starts.const_data_ptr(), static_cast(batch_size), block_size, total_tokens, block_table_stride, cache_block_stride, cache_entry_stride, dst_entry_stride); } // Macro to dispatch the kernel based on the data type. -#define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \ - vllm::indexer_k_quant_and_cache_kernel \ - <<>>( \ - reinterpret_cast(k.data_ptr()), \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.data_ptr(), head_dim, quant_block_size, \ +#define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \ + vllm::indexer_k_quant_and_cache_kernel \ + <<>>( \ + reinterpret_cast(k.data_ptr()), \ + reinterpret_cast(kv_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), head_dim, quant_block_size, \ cache_block_size, cache_stride, use_ue8m0); void indexer_k_quant_and_cache( - torch::Tensor& k, // [num_tokens, head_dim] - torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] - torch::Tensor& slot_mapping, // [num_tokens] - int64_t quant_block_size, // quantization block size + torch::stable::Tensor& k, // [num_tokens, head_dim] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] + torch::stable::Tensor& slot_mapping, // [num_tokens] + int64_t quant_block_size, // quantization block size const std::string& scale_fmt) { int num_tokens = k.size(0); int head_dim = k.size(1); @@ -1426,65 +1466,70 @@ void indexer_k_quant_and_cache( int cache_stride = kv_cache.size(2); bool use_ue8m0 = scale_fmt == "ue8m0"; - TORCH_CHECK(k.device() == kv_cache.device(), - "k and kv_cache must be on the same device"); - TORCH_CHECK(k.device() == slot_mapping.device(), - "k and slot_mapping must be on the same device"); - TORCH_CHECK(head_dim % quant_block_size == 0, - "head_dim must be divisible by quant_block_size"); + STD_TORCH_CHECK(k.device() == kv_cache.device(), + "k and kv_cache must be on the same device"); + STD_TORCH_CHECK(k.device() == slot_mapping.device(), + "k and slot_mapping must be on the same device"); + STD_TORCH_CHECK(head_dim % quant_block_size == 0, + "head_dim must be divisible by quant_block_size"); constexpr int vec_size = 4; dim3 grid(num_tokens, (head_dim + quant_block_size * vec_size - 1) / (quant_block_size * vec_size)); dim3 block(32, vec_size); - const at::cuda::OptionalCUDAGuard device_guard(device_of(k)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + k.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); static const std::string kv_cache_dtype = "fp8_e4m3"; - DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(k.scalar_type(), kv_cache_dtype, CALL_INDEXER_K_QUANT_AND_CACHE); } // Macro to dispatch the kernel based on the data amount. -#define CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(BLOCK_Y_SIZE) \ - vllm::cp_gather_indexer_k_quant_cache_kernel \ - <<>>( \ - reinterpret_cast(kv_cache.data_ptr()), \ - reinterpret_cast(dst_k.data_ptr()), \ - reinterpret_cast(dst_scale.data_ptr()), \ - block_table.data_ptr(), cu_seq_lens.data_ptr(), \ - batch_size, dst_k.stride(0), dst_k.size(1), kv_cache.stride(0), \ - kv_cache.stride(1), kv_cache.size(1), block_table.size(1), \ - num_tokens, quant_block_size); +#define CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(BLOCK_Y_SIZE) \ + vllm::cp_gather_indexer_k_quant_cache_kernel \ + <<>>( \ + reinterpret_cast(kv_cache.data_ptr()), \ + reinterpret_cast(dst_k.data_ptr()), \ + reinterpret_cast(dst_scale.data_ptr()), \ + block_table.const_data_ptr(), \ + cu_seq_lens.const_data_ptr(), batch_size, dst_k.stride(0), \ + dst_k.size(1), kv_cache.stride(0), kv_cache.stride(1), \ + kv_cache.size(1), block_table.size(1), num_tokens, \ + quant_block_size); void cp_gather_indexer_k_quant_cache( - const torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] - torch::Tensor& dst_k, // [num_tokens, head_dim] - torch::Tensor& dst_scale, // [num_tokens, head_dim / quant_block_size * 4] - const torch::Tensor& block_table, // [batch_size, num_blocks] - const torch::Tensor& cu_seq_lens // [batch_size + 1] + const torch::stable::Tensor& + kv_cache, // [num_blocks, block_size, cache_stride] + torch::stable::Tensor& dst_k, // [num_tokens, head_dim] + torch::stable::Tensor& + dst_scale, // [num_tokens, head_dim / quant_block_size * 4] + const torch::stable::Tensor& block_table, // [batch_size, num_blocks] + const torch::stable::Tensor& cu_seq_lens // [batch_size + 1] ) { int batch_size = block_table.size(0); int num_tokens = dst_k.size(0); int head_dim = dst_k.size(1); int quant_block_size = head_dim * 4 / dst_scale.size(1); - TORCH_CHECK(kv_cache.device() == dst_k.device(), - "kv_cache and dst_k must be on the same device"); - TORCH_CHECK(kv_cache.device() == dst_scale.device(), - "kv_cache and dst_scale must be on the same device"); - TORCH_CHECK(kv_cache.device() == block_table.device(), - "kv_cache and block_table must be on the same device"); - TORCH_CHECK(kv_cache.device() == cu_seq_lens.device(), - "kv_cache and cu_seq_lens must be on the same device"); - TORCH_CHECK(head_dim % quant_block_size == 0, - "head_dim must be divisible by quant_block_size"); + STD_TORCH_CHECK(kv_cache.device() == dst_k.device(), + "kv_cache and dst_k must be on the same device"); + STD_TORCH_CHECK(kv_cache.device() == dst_scale.device(), + "kv_cache and dst_scale must be on the same device"); + STD_TORCH_CHECK(kv_cache.device() == block_table.device(), + "kv_cache and block_table must be on the same device"); + STD_TORCH_CHECK(kv_cache.device() == cu_seq_lens.device(), + "kv_cache and cu_seq_lens must be on the same device"); + STD_TORCH_CHECK(head_dim % quant_block_size == 0, + "head_dim must be divisible by quant_block_size"); constexpr int vec_size = 16; - const at::cuda::OptionalCUDAGuard device_guard(device_of(kv_cache)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + kv_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); if (num_tokens < 32) { CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(1); @@ -1503,27 +1548,30 @@ void cp_gather_indexer_k_quant_cache( // Concatenate ql_nope and q_pe into a contiguous q_out tensor for MLA/DSA. // Replaces torch.cat((ql_nope, q_pe), dim=-1). -void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] - torch::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] - torch::Tensor& q_out // [num_tokens, num_heads, nope_dim + - // rope_dim] +void concat_mla_q( + torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] + torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] + torch::stable::Tensor& q_out // [num_tokens, num_heads, nope_dim + + // rope_dim] ) { const int num_tokens = ql_nope.size(0); const int num_heads = ql_nope.size(1); const int nope_dim = ql_nope.size(2); const int rope_dim = q_pe.size(2); - TORCH_CHECK(nope_dim % 512 == 0, "nope_dim must be a multiple of 512, got ", - nope_dim); - TORCH_CHECK(rope_dim == 64, "rope_dim must be 64, got ", rope_dim); - TORCH_CHECK(q_out.size(2) == nope_dim + rope_dim); + STD_TORCH_CHECK(nope_dim % 512 == 0, + "nope_dim must be a multiple of 512, got ", nope_dim); + STD_TORCH_CHECK(rope_dim == 64, "rope_dim must be 64, got ", rope_dim); + STD_TORCH_CHECK(q_out.size(2) == nope_dim + rope_dim); - TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2"); - TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2"); - TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2"); - TORCH_CHECK(ql_nope.scalar_type() == at::ScalarType::Half || - ql_nope.scalar_type() == at::ScalarType::BFloat16, - "ql_nope must be float16 or bfloat16 dtype"); + STD_TORCH_CHECK(ql_nope.stride(2) == 1, + "ql_nope must have stride 1 in dim 2"); + STD_TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2"); + STD_TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2"); + STD_TORCH_CHECK( + ql_nope.scalar_type() == torch::headeronly::ScalarType::Half || + ql_nope.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "ql_nope must be float16 or bfloat16 dtype"); if (num_tokens == 0) return; @@ -1532,13 +1580,14 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] const int grid_size = (total_warps + warps_per_block - 1) / warps_per_block; const int block_size = warps_per_block * 32; - const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + ql_nope.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); - VLLM_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] { + VLLM_STABLE_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] { vllm::ConcatMLAQKernel<<>>( - q_out.data_ptr(), ql_nope.data_ptr(), - q_pe.data_ptr(), num_tokens, num_heads, q_out.stride(0), + q_out.mutable_data_ptr(), ql_nope.const_data_ptr(), + q_pe.const_data_ptr(), num_tokens, num_heads, q_out.stride(0), q_out.stride(1), ql_nope.stride(0), ql_nope.stride(1), q_pe.stride(0), q_pe.stride(1)); }); diff --git a/csrc/cache_kernels_fused.cu b/csrc/libtorch_stable/cache_kernels_fused.cu similarity index 53% rename from csrc/cache_kernels_fused.cu rename to csrc/libtorch_stable/cache_kernels_fused.cu index 8687ebe1f14..744e3a2dcdc 100644 --- a/csrc/cache_kernels_fused.cu +++ b/csrc/libtorch_stable/cache_kernels_fused.cu @@ -1,15 +1,13 @@ -#include -#include -#include - -#include "cuda_compat.h" +#include "torch_utils.h" #include "dispatch_utils.h" -#include "quantization/w8a8/fp8/common.cuh" +#include "../cuda_compat.h" + +#include "../quantization/w8a8/fp8/common.cuh" #ifdef USE_ROCM - #include "quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif #ifdef USE_ROCM @@ -164,43 +162,52 @@ __global__ void concat_and_cache_mla_rope_fused_kernel( } // namespace vllm -#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \ - do { \ - VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \ - using qk_t = scalar_t; \ - VLLM_DISPATCH_FLOATING_TYPES( \ - rope_cos_sin_cache.scalar_type(), "rope_cos_sin_cache_scalar_type", \ - [&] { \ - using cos_sin_t = scalar_t; \ - if (rope_is_neox) { \ - vllm::concat_and_cache_mla_rope_fused_kernel< \ - qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \ - <<>>( \ - positions.data_ptr(), q_pe.data_ptr(), \ - k_pe.data_ptr(), kv_c.data_ptr(), \ - rope_cos_sin_cache.data_ptr(), rot_dim, \ - q_pe_stride_token, q_pe_stride_head, k_pe_stride, \ - kv_c_stride, num_q_heads, \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.data_ptr(), block_stride, \ - entry_stride, kv_lora_rank, block_size, \ - kv_cache_quant_scale.data_ptr()); \ - } else { \ - vllm::concat_and_cache_mla_rope_fused_kernel< \ - qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \ - <<>>( \ - positions.data_ptr(), q_pe.data_ptr(), \ - k_pe.data_ptr(), kv_c.data_ptr(), \ - rope_cos_sin_cache.data_ptr(), rot_dim, \ - q_pe_stride_token, q_pe_stride_head, k_pe_stride, \ - kv_c_stride, num_q_heads, \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.data_ptr(), block_stride, \ - entry_stride, kv_lora_rank, block_size, \ - kv_cache_quant_scale.data_ptr()); \ - } \ - }); \ - }); \ +#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \ + do { \ + VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ + q_pe.scalar_type(), "qk_scalar_type", [&] { \ + using qk_t = scalar_t; \ + VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ + rope_cos_sin_cache.scalar_type(), \ + "rope_cos_sin_cache_scalar_type", [&] { \ + using cos_sin_t = scalar_t; \ + if (rope_is_neox) { \ + vllm::concat_and_cache_mla_rope_fused_kernel< \ + qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \ + <<>>( \ + positions.const_data_ptr(), \ + q_pe.mutable_data_ptr(), \ + k_pe.mutable_data_ptr(), \ + kv_c.const_data_ptr(), \ + rope_cos_sin_cache.const_data_ptr(), \ + rot_dim, q_pe_stride_token, q_pe_stride_head, \ + k_pe_stride, kv_c_stride, num_q_heads, \ + reinterpret_cast( \ + kv_cache.mutable_data_ptr()), \ + slot_mapping.const_data_ptr(), \ + block_stride, entry_stride, kv_lora_rank, \ + block_size, \ + kv_cache_quant_scale.const_data_ptr()); \ + } else { \ + vllm::concat_and_cache_mla_rope_fused_kernel< \ + qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \ + <<>>( \ + positions.const_data_ptr(), \ + q_pe.mutable_data_ptr(), \ + k_pe.mutable_data_ptr(), \ + kv_c.const_data_ptr(), \ + rope_cos_sin_cache.const_data_ptr(), \ + rot_dim, q_pe_stride_token, q_pe_stride_head, \ + k_pe_stride, kv_c_stride, num_q_heads, \ + reinterpret_cast( \ + kv_cache.mutable_data_ptr()), \ + slot_mapping.const_data_ptr(), \ + block_stride, entry_stride, kv_lora_rank, \ + block_size, \ + kv_cache_quant_scale.const_data_ptr()); \ + } \ + }); \ + }); \ } while (false) // Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache. @@ -208,64 +215,69 @@ __global__ void concat_and_cache_mla_rope_fused_kernel( // Replaces DeepseekScalingRotaryEmbedding.self.rotary_emb and // concat_and_cache_mla. void concat_and_cache_mla_rope_fused( - torch::Tensor& positions, // [num_tokens] - torch::Tensor& q_pe, // [num_tokens, num_q_heads, rot_dim] - torch::Tensor& k_pe, // [num_tokens, rot_dim] - torch::Tensor& kv_c, // [num_tokens, kv_lora_rank] - torch::Tensor& rope_cos_sin_cache, // [max_position, rot_dim] + torch::stable::Tensor& positions, // [num_tokens] + torch::stable::Tensor& q_pe, // [num_tokens, num_q_heads, rot_dim] + torch::stable::Tensor& k_pe, // [num_tokens, rot_dim] + torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank] + torch::stable::Tensor& rope_cos_sin_cache, // [max_position, rot_dim] bool rope_is_neox, - torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] - torch::Tensor& + torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank + rot_dim)] - const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) { + const std::string& kv_cache_dtype, + torch::stable::Tensor& kv_cache_quant_scale) { // NOTE(woosuk): In vLLM V1, query/key/position.size(0) can be different from // slot_mapping.size(0) because of padding for CUDA graphs. - // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because - // both include padding. - // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0) - // since key includes padding for CUDA graphs, while slot_mapping does not. - // In this case, slot_mapping.size(0) represents the actual number of tokens + // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) + // because both include padding. + // In vLLM V1, however, key.size(0) can be larger than + // slot_mapping.size(0) since key includes padding for CUDA graphs, + // while slot_mapping does not. In this case, + // slot_mapping.size(0) represents the actual number of tokens // before padding. - // For compatibility with both cases, we use slot_mapping.size(0) as the - // number of tokens. - int num_tokens = slot_mapping.size(0); - int num_padded_tokens = q_pe.size(0); - TORCH_CHECK_GE(num_padded_tokens, num_tokens); + // For compatibility with both cases, we use slot_mapping.size(0) as + // the number of tokens. + const int64_t num_tokens = slot_mapping.size(0); + const int64_t num_padded_tokens = q_pe.size(0); + STD_TORCH_CHECK(num_padded_tokens >= num_tokens); const int num_q_heads = q_pe.size(1); const int rot_dim = q_pe.size(2); const int kv_lora_rank = kv_c.size(1); - TORCH_CHECK_EQ(positions.size(0), num_padded_tokens); - TORCH_CHECK_EQ(positions.dim(), 1); - TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long); + STD_TORCH_CHECK(positions.size(0) == num_padded_tokens); + STD_TORCH_CHECK(positions.dim() == 1); + STD_TORCH_CHECK(positions.scalar_type() == + torch::headeronly::ScalarType::Long); - TORCH_CHECK_EQ(q_pe.dim(), 3); - TORCH_CHECK_EQ(q_pe.size(0), num_padded_tokens); - TORCH_CHECK_EQ(q_pe.size(1), num_q_heads); - TORCH_CHECK_EQ(q_pe.size(2), rot_dim); + STD_TORCH_CHECK(q_pe.dim() == 3); + STD_TORCH_CHECK(q_pe.size(0) == num_padded_tokens); + STD_TORCH_CHECK(q_pe.size(1) == num_q_heads); + STD_TORCH_CHECK(q_pe.size(2) == rot_dim); - TORCH_CHECK_EQ(k_pe.dim(), 2); - TORCH_CHECK_EQ(k_pe.size(0), num_padded_tokens); - TORCH_CHECK_EQ(k_pe.size(1), rot_dim); - TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type()); + STD_TORCH_CHECK(k_pe.dim() == 2); + STD_TORCH_CHECK(k_pe.size(0) == num_padded_tokens); + STD_TORCH_CHECK(k_pe.size(1) == rot_dim); + STD_TORCH_CHECK(k_pe.scalar_type() == q_pe.scalar_type()); - TORCH_CHECK_EQ(kv_c.dim(), 2); - TORCH_CHECK_EQ(kv_c.size(0), num_padded_tokens); - TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank); - TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type()); - TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype()); + STD_TORCH_CHECK(kv_c.dim() == 2); + STD_TORCH_CHECK(kv_c.size(0) == num_padded_tokens); + STD_TORCH_CHECK(kv_c.size(1) == kv_lora_rank); + STD_TORCH_CHECK(kv_c.scalar_type() == q_pe.scalar_type()); - TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim); + STD_TORCH_CHECK(rope_cos_sin_cache.size(1) == rot_dim); + STD_TORCH_CHECK(rope_cos_sin_cache.scalar_type() == q_pe.scalar_type()); - TORCH_CHECK_EQ(slot_mapping.size(0), num_tokens); - TORCH_CHECK_EQ(slot_mapping.scalar_type(), c10::ScalarType::Long); + STD_TORCH_CHECK(slot_mapping.size(0) == num_tokens); + STD_TORCH_CHECK(slot_mapping.scalar_type() == + torch::headeronly::ScalarType::Long); - TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim); - TORCH_CHECK_EQ(kv_cache.dim(), 3); + STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + rot_dim); + STD_TORCH_CHECK(kv_cache.dim() == 3); - TORCH_CHECK_EQ(kv_cache_quant_scale.numel(), 1); - TORCH_CHECK_EQ(kv_cache_quant_scale.scalar_type(), c10::ScalarType::Float); + STD_TORCH_CHECK(kv_cache_quant_scale.numel() == 1); + STD_TORCH_CHECK(kv_cache_quant_scale.scalar_type() == + torch::headeronly::ScalarType::Float); int64_t q_pe_stride_token = q_pe.stride(0); int64_t q_pe_stride_head = q_pe.stride(1); @@ -286,9 +298,10 @@ void concat_and_cache_mla_rope_fused( dim3 grid(num_tokens, 1, 1); dim3 block(thread_block_size, 1, 1); - const at::cuda::OptionalCUDAGuard device_guard(device_of(positions)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + positions.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); - DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED); } diff --git a/csrc/concat_mla_q.cuh b/csrc/libtorch_stable/concat_mla_q.cuh similarity index 100% rename from csrc/concat_mla_q.cuh rename to csrc/libtorch_stable/concat_mla_q.cuh diff --git a/csrc/cuda_vec_utils.cuh b/csrc/libtorch_stable/cuda_vec_utils.cuh similarity index 100% rename from csrc/cuda_vec_utils.cuh rename to csrc/libtorch_stable/cuda_vec_utils.cuh diff --git a/csrc/nvfp4_kv_cache_kernels.cu b/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu similarity index 81% rename from csrc/nvfp4_kv_cache_kernels.cu rename to csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu index d6aa715c203..693d55f231e 100644 --- a/csrc/nvfp4_kv_cache_kernels.cu +++ b/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu @@ -17,11 +17,8 @@ #define NVFP4_ENABLE_ELTS16 1 #include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh" -#include -#include -#include - -#include "dispatch_utils.h" +#include "libtorch_stable/dispatch_utils.h" +#include "libtorch_stable/torch_utils.h" namespace vllm { @@ -184,12 +181,13 @@ __global__ void reshape_and_cache_nvfp4_kernel( // Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1]. // Each KV side contains both data and scale: // page = [K_data | K_scale | V_data | V_scale] -void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value, - torch::Tensor& key_cache, - torch::Tensor& value_cache, - torch::Tensor& slot_mapping, - torch::Tensor& k_scale, - torch::Tensor& v_scale) { +void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key, + torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, + torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, + torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale) { int num_tokens = slot_mapping.size(0); int num_heads = key.size(1); int head_size = key.size(2); @@ -200,17 +198,18 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value, // key_cache is kv_cache[:, 0] with shape // [num_blocks, block_size, num_heads, full_dim] in logical order. // Strides encode the physical layout (HND or NHD). - TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D"); - TORCH_CHECK(key_cache.size(3) == full_dim, - "key_cache last dim must be data_dim + scale_dim, got ", - key_cache.size(3), " expected ", full_dim); + STD_TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D"); + STD_TORCH_CHECK(key_cache.size(3) == full_dim, + "key_cache last dim must be data_dim + scale_dim, got ", + key_cache.size(3), " expected ", full_dim); int block_size = key_cache.size(1); - TORCH_CHECK(head_size % 16 == 0, - "head_size must be divisible by 16 for NVFP4 KV cache"); - TORCH_CHECK(block_size % 4 == 0, - "block_size must be divisible by 4 for NVFP4 KV cache swizzle"); + STD_TORCH_CHECK(head_size % 16 == 0, + "head_size must be divisible by 16 for NVFP4 KV cache"); + STD_TORCH_CHECK(block_size % 4 == 0, + "block_size must be divisible by 4 for NVFP4 KV cache " + "swizzle"); // Detect physical layout from strides (based on full_dim). // HND: head stride > block_offset stride. @@ -230,8 +229,9 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value, // Scale follows data within each KV side. int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim; - uint8_t* key_scale_ptr = key_cache.data_ptr() + data_per_kv; - uint8_t* value_scale_ptr = value_cache.data_ptr() + data_per_kv; + uint8_t* key_scale_ptr = key_cache.mutable_data_ptr() + data_per_kv; + uint8_t* value_scale_ptr = + value_cache.mutable_data_ptr() + data_per_kv; // Scale strides: same page stride, inner strides from layout. int64_t scale_block_stride = data_block_stride; @@ -244,8 +244,8 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value, scale_block_offset_stride = (int64_t)num_heads * scale_dim; } - const float* k_scale_ptr = k_scale.data_ptr(); - const float* v_scale_ptr = v_scale.data_ptr(); + const float* k_scale_ptr = k_scale.const_data_ptr(); + const float* v_scale_ptr = v_scale.const_data_ptr(); int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE; int total_groups = num_heads * groups_per_head; @@ -256,20 +256,22 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value, dim3 grid(num_tokens); dim3 block(num_threads); - const at::cuda::OptionalCUDAGuard device_guard(device_of(key)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + key.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); - AT_DISPATCH_REDUCED_FLOATING_TYPES( + VLLM_STABLE_DISPATCH_HALF_TYPES( key.scalar_type(), "reshape_and_cache_nvfp4", [&] { vllm::reshape_and_cache_nvfp4_kernel <<>>( - key.data_ptr(), value.data_ptr(), - key_cache.data_ptr(), value_cache.data_ptr(), - key_scale_ptr, value_scale_ptr, - slot_mapping.data_ptr(), k_scale_ptr, v_scale_ptr, - key.stride(0), value.stride(0), num_heads, head_size, - block_size, data_block_stride, data_head_stride, - data_block_offset_stride, scale_block_stride, scale_head_stride, - scale_block_offset_stride); + key.const_data_ptr(), + value.const_data_ptr(), + key_cache.mutable_data_ptr(), + value_cache.mutable_data_ptr(), key_scale_ptr, + value_scale_ptr, slot_mapping.const_data_ptr(), + k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0), + num_heads, head_size, block_size, data_block_stride, + data_head_stride, data_block_offset_stride, scale_block_stride, + scale_head_stride, scale_block_offset_stride); }); } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 90082ac3d29..2b64cab597a 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -355,3 +355,132 @@ torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X, int64_t tokens); int64_t ggml_moe_get_block_size(int64_t type); + +void paged_attention_v1( + torch::stable::Tensor& out, torch::stable::Tensor& query, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables, + torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step); + +void paged_attention_v2( + torch::stable::Tensor& out, torch::stable::Tensor& exp_sums, + torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out, + torch::stable::Tensor& query, torch::stable::Tensor& key_cache, + torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale, + torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens, + int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step); + +// Cache ops (shared CUDA/ROCm) +void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, + int64_t block_size_in_bytes, + const torch::stable::Tensor& block_mapping); + +// Batch swap: submit all block copies in a single driver call. +void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, + const torch::stable::Tensor& dst_ptrs, + const torch::stable::Tensor& sizes, + bool is_src_access_order_any); + +void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, + torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, + const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale); + +void reshape_and_cache_flash( + torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale); + +void concat_and_cache_mla(torch::stable::Tensor& kv_c, + torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache, + torch::stable::Tensor& slot_mapping, + const std::string& kv_cache_dtype, + torch::stable::Tensor& scale); + +// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla +void concat_and_cache_mla_rope_fused( + torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, + torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c, + torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox, + torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache, + const std::string& kv_cache_dtype, + torch::stable::Tensor& kv_cache_quant_scale); + +// Just for unittest +void convert_fp8(torch::stable::Tensor& dst_cache, + torch::stable::Tensor& src_cache, const double scale, + const std::string& kv_cache_dtype); + +void gather_and_maybe_dequant_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] + int64_t num_tokens, const std::string& kv_cache_dtype, + torch::stable::Tensor const& scale, + std::optional seq_starts = std::nullopt); + +// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. +void cp_gather_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + int64_t batch_size, + std::optional seq_starts = std::nullopt); + +// Gather and upconvert FP8 KV cache to BF16 workspace +void cp_gather_and_upconvert_fp8_kv_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // 656] + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& seq_lens, // [BATCH] + torch::stable::Tensor const& workspace_starts, // [BATCH] + int64_t batch_size); + +// Indexer K quantization and cache function +void indexer_k_quant_and_cache( + torch::stable::Tensor& k, // [num_tokens, head_dim] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, + // cache_stride] + torch::stable::Tensor& slot_mapping, // [num_tokens] + int64_t quant_block_size, // quantization block size + const std::string& scale_fmt); + +// Concatenate query nope and rope for MLA/DSA attention +void concat_mla_q( + torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] + torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] + torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim + + // rope_dim] + +// Extract function to gather quantized K cache +void cp_gather_indexer_k_quant_cache( + const torch::stable::Tensor& kv_cache, // [num_blocks, block_size, + // cache_stride] + torch::stable::Tensor& dst_k, // [num_tokens, head_dim] + torch::stable::Tensor& dst_scale, // [num_tokens, head_dim / + // quant_block_size * 4] + const torch::stable::Tensor& block_table, // [batch_size, num_blocks] + const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1] diff --git a/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu b/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu index 49f2944f3fd..d0cafa0c6df 100644 --- a/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu +++ b/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu @@ -17,7 +17,7 @@ #include #include "libtorch_stable/torch_utils.h" #include "libtorch_stable/dispatch_utils.h" -#include "cuda_vec_utils.cuh" +#include "../../cuda_vec_utils.cuh" #include #include diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu index 78e4eda0c01..ae4f4829d93 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu @@ -27,7 +27,7 @@ #include #include "libtorch_stable/torch_utils.h" #include "libtorch_stable/dispatch_utils.h" -#include "cuda_vec_utils.cuh" +#include "../../cuda_vec_utils.cuh" #include "cuda_utils.h" #include "nvfp4_utils.cuh" diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu index 744ae4f7311..8f2b4405cf6 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu @@ -17,7 +17,7 @@ #include #include "libtorch_stable/torch_utils.h" #include "libtorch_stable/dispatch_utils.h" -#include "cuda_vec_utils.cuh" +#include "../../cuda_vec_utils.cuh" #include #include diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu index e1d101f9d86..ccefe294022 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu @@ -23,7 +23,7 @@ #include "libtorch_stable/torch_utils.h" #include "libtorch_stable/dispatch_utils.h" -#include "cuda_vec_utils.cuh" +#include "../../cuda_vec_utils.cuh" #include "cuda_utils.h" #include "launch_bounds_utils.h" diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh b/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh index 590e4c06b62..0c04f010888 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh @@ -20,7 +20,7 @@ #include #include -#include "cuda_vec_utils.cuh" +#include "../../cuda_vec_utils.cuh" #if defined(NVFP4_ENABLE_ELTS16) && defined(CUDA_VERSION) && \ CUDA_VERSION >= 12090 diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 23934193230..152391f9c20 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -474,6 +474,33 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor? initial_state_idx," "Tensor? cu_chunk_seqlen," "Tensor? last_chunk_indices) -> ()"); + + // Attention ops + // Compute the attention between an input query and the cached + // keys/values using PagedAttention. + ops.def( + "paged_attention_v1(" + " Tensor! out, Tensor query, Tensor key_cache," + " Tensor value_cache, int num_kv_heads, float scale," + " Tensor block_tables, Tensor seq_lens, int block_size," + " int max_seq_len, Tensor? alibi_slopes," + " str kv_cache_dtype, Tensor k_scale, Tensor v_scale," + " int tp_rank, int blocksparse_local_blocks," + " int blocksparse_vert_stride, int blocksparse_block_size," + " int blocksparse_head_sliding_step) -> ()"); + + // PagedAttention V2. + ops.def( + "paged_attention_v2(" + " Tensor! out, Tensor! exp_sums, Tensor! max_logits," + " Tensor! tmp_out, Tensor query, Tensor key_cache," + " Tensor value_cache, int num_kv_heads, float scale," + " Tensor block_tables, Tensor seq_lens, int block_size," + " int max_seq_len, Tensor? alibi_slopes," + " str kv_cache_dtype, Tensor k_scale, Tensor v_scale," + " int tp_rank, int blocksparse_local_blocks," + " int blocksparse_vert_stride, int blocksparse_block_size," + " int blocksparse_head_sliding_step) -> ()"); } STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { @@ -581,6 +608,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8)); ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec)); ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); + + ops.impl("paged_attention_v1", TORCH_BOX(&paged_attention_v1)); + ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2)); } // These capability-check functions take only primitive args (no tensors), so @@ -603,4 +633,115 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { ops.impl("ggml_moe_get_block_size", TORCH_BOX(&ggml_moe_get_block_size)); } +// Cache ops +STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { + // Swap in (out) the cache blocks from src to dst. + ops.def( + "swap_blocks(Tensor src, Tensor! dst," + " int block_size_in_bytes, Tensor block_mapping) -> ()"); + + // Batch swap: submit all block copies in a single driver call. + ops.def( + "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," + " Tensor sizes," + " bool is_src_access_order_any=False) -> ()"); + + // Reshape the key and value tensors and cache them. + ops.def( + "reshape_and_cache(Tensor key, Tensor value," + " Tensor! key_cache, Tensor! value_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor k_scale, Tensor v_scale) -> ()"); + + // Reshape the key and value tensors and cache them. + ops.def( + "reshape_and_cache_flash(Tensor key, Tensor value," + " Tensor! key_cache," + " Tensor! value_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor k_scale, Tensor v_scale) -> ()"); + + // Concat kv_c and k_pe and cache them. + ops.def( + "concat_and_cache_mla(Tensor kv_c, Tensor k_pe," + " Tensor! kv_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor scale) -> ()"); + + // Rotate Q and K, then write to kv cache for MLA + ops.def( + "concat_and_cache_mla_rope_fused(" + " Tensor positions," + " Tensor! q_pe," + " Tensor! k_pe," + " Tensor kv_c," + " Tensor cos_sin_cache," + " bool is_neox," + " Tensor slot_mapping," + " Tensor! kv_cache," + " str kv_cache_dtype," + " Tensor kv_cache_scale) -> ()"); + + // Convert the key and value cache to fp8 data type. + ops.def( + "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, " + "str kv_cache_dtype) -> ()"); + + // Gather cache blocks from src_cache to dst, dequantizing from + // src_cache's dtype to dst's dtype if necessary. + ops.def( + "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, " + " Tensor block_table, Tensor cu_seq_lens, " + " Tensor token_to_seq, " + " int num_tokens, " + " str kv_cache_dtype, " + " Tensor scale, Tensor? seq_starts) -> ()"); + + ops.def( + "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, " + "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()"); + + ops.def( + "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " + "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int " + "batch_size) -> ()"); + + ops.def( + "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " + "slot_mapping, " + "int quant_block_size, str kv_cache_dtype) -> ()"); + + ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()"); + + ops.def( + "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! " + "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) { + ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { + ops.impl("swap_blocks", TORCH_BOX(&swap_blocks)); + ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); + ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); + ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); + ops.impl("concat_and_cache_mla_rope_fused", + TORCH_BOX(&concat_and_cache_mla_rope_fused)); + ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); + ops.impl("gather_and_maybe_dequant_cache", + TORCH_BOX(&gather_and_maybe_dequant_cache)); + ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache)); + ops.impl("cp_gather_and_upconvert_fp8_kv_cache", + TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache)); + ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache)); + ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q)); + ops.impl("cp_gather_indexer_k_quant_cache", + TORCH_BOX(&cp_gather_indexer_k_quant_cache)); +} + REGISTER_EXTENSION(_C_stable_libtorch) diff --git a/csrc/ops.h b/csrc/ops.h index 34da56fec4b..f458f79d6f4 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -31,29 +31,6 @@ torch::Tensor weak_ref_tensor(torch::Tensor& tensor) { return new_tensor; } -void paged_attention_v1( - torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache, - torch::Tensor& value_cache, int64_t num_kv_heads, double scale, - torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size, - int64_t max_seq_len, const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::Tensor& k_scale, - torch::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step); - -void paged_attention_v2( - torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, - torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, - torch::Tensor& value_cache, int64_t num_kv_heads, double scale, - torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size, - int64_t max_seq_len, const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::Tensor& k_scale, - torch::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step); - // rms_norm and fused_add_rms_norm declarations also exist in // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here // because the CPU build still uses these torch::Tensor declarations. diff --git a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh index 7ae644d81d4..58d19de7349 100644 --- a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh @@ -6,6 +6,7 @@ #include #include "../../../../attention/attention_dtypes.h" +#include namespace vllm { #ifdef USE_ROCM @@ -642,27 +643,29 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \ vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \ if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ + if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \ FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \ FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \ FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); \ } else { \ - TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + STD_TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ + if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \ FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \ FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \ FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ } else { \ - TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + STD_TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ } else { \ - TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ + STD_TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ } } // namespace fp8 diff --git a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh index 3b7e25dc56b..6e95d72ed5e 100644 --- a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh @@ -1,6 +1,7 @@ #pragma once #include "../../../../attention/attention_dtypes.h" +#include #include #include #include @@ -546,37 +547,40 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \ vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \ if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ + if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \ FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \ FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \ FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); \ } else { \ - TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + STD_TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ + if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \ FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \ FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \ FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ } else { \ - TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + STD_TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E5M2) { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ + if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \ FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \ FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \ FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ } else { \ - TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + STD_TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ } else { \ - TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ + STD_TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ } } // namespace fp8 diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index a2a5d809745..01869474e0f 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -1,4 +1,7 @@ -#include "cache.h" +// Provides torch::Tensor for ops.h (previously included transitively via +// cache.h, which is no longer included here after cache ops moved to +// _C_stable_libtorch). +#include #include "cuda_utils.h" #include "ops.h" #include "core/registration.h" @@ -33,35 +36,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU, &get_cuda_view_from_cpu_tensor); - // Attention ops - // Compute the attention between an input query and the cached - // keys/values using PagedAttention. - ops.def( - "paged_attention_v1(" - " Tensor! out, Tensor query, Tensor key_cache," - " Tensor value_cache, int num_kv_heads, float scale," - " Tensor block_tables, Tensor seq_lens, int block_size," - " int max_seq_len, Tensor? alibi_slopes," - " str kv_cache_dtype, Tensor k_scale, Tensor v_scale," - " int tp_rank, int blocksparse_local_blocks," - " int blocksparse_vert_stride, int blocksparse_block_size," - " int blocksparse_head_sliding_step) -> ()"); - ops.impl("paged_attention_v1", torch::kCUDA, &paged_attention_v1); - - // PagedAttention V2. - ops.def( - "paged_attention_v2(" - " Tensor! out, Tensor! exp_sums, Tensor! max_logits," - " Tensor! tmp_out, Tensor query, Tensor key_cache," - " Tensor value_cache, int num_kv_heads, float scale," - " Tensor block_tables, Tensor seq_lens, int block_size," - " int max_seq_len, Tensor? alibi_slopes," - " str kv_cache_dtype, Tensor k_scale, Tensor v_scale," - " int tp_rank, int blocksparse_local_blocks," - " int blocksparse_vert_stride, int blocksparse_block_size," - " int blocksparse_head_sliding_step) -> ()"); - ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2); - // Activation ops (quantized only — basic ops moved to _C_stable_libtorch) ops.def( "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); @@ -217,114 +191,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { #endif } -TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) { - // Cache ops - // Swap in (out) the cache blocks from src to dst. - cache_ops.def( - "swap_blocks(Tensor src, Tensor! dst," - " int block_size_in_bytes, Tensor block_mapping) -> ()"); - cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks); - - // Batch swap: submit all block copies in a single driver call. - cache_ops.def( - "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," - " Tensor sizes," - " bool is_src_access_order_any=False) -> ()"); - cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch); - - // Reshape the key and value tensors and cache them. - cache_ops.def( - "reshape_and_cache(Tensor key, Tensor value," - " Tensor! key_cache, Tensor! value_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); - cache_ops.impl("reshape_and_cache", torch::kCUDA, &reshape_and_cache); - - // Reshape the key and value tensors and cache them. - cache_ops.def( - "reshape_and_cache_flash(Tensor key, Tensor value," - " Tensor! key_cache," - " Tensor! value_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); - cache_ops.impl("reshape_and_cache_flash", torch::kCUDA, - &reshape_and_cache_flash); - - // Concat kv_c and k_pe and cache them. - cache_ops.def( - "concat_and_cache_mla(Tensor kv_c, Tensor k_pe," - " Tensor! kv_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor scale) -> ()"); - cache_ops.impl("concat_and_cache_mla", torch::kCUDA, &concat_and_cache_mla); - - // Rotate Q and K, then write to kv cache for MLA - cache_ops.def( - "concat_and_cache_mla_rope_fused(" - " Tensor positions," - " Tensor! q_pe," - " Tensor! k_pe," - " Tensor kv_c," - " Tensor cos_sin_cache," - " bool is_neox," - " Tensor slot_mapping," - " Tensor! kv_cache," - " str kv_cache_dtype," - " Tensor kv_cache_scale) -> ()"); - cache_ops.impl("concat_and_cache_mla_rope_fused", torch::kCUDA, - &concat_and_cache_mla_rope_fused); - - // Convert the key and value cache to fp8 data type. - cache_ops.def( - "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, " - "str kv_cache_dtype) -> ()"); - cache_ops.impl("convert_fp8", torch::kCUDA, &convert_fp8); - - // Gather cache blocks from src_cache to dst, dequantizing from - // src_cache's dtype to dst's dtype if necessary. - cache_ops.def( - "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, " - " Tensor block_table, Tensor cu_seq_lens, " - " Tensor token_to_seq, " - " int num_tokens, " - " str kv_cache_dtype, " - " Tensor scale, Tensor? seq_starts) -> ()"); - cache_ops.impl("gather_and_maybe_dequant_cache", torch::kCUDA, - &gather_and_maybe_dequant_cache); - - cache_ops.def( - "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, " - "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()"); - cache_ops.impl("cp_gather_cache", torch::kCUDA, &cp_gather_cache); - - cache_ops.def( - "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " - "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int " - "batch_size) -> ()"); - cache_ops.impl("cp_gather_and_upconvert_fp8_kv_cache", torch::kCUDA, - &cp_gather_and_upconvert_fp8_kv_cache); - - cache_ops.def( - "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " - "slot_mapping, " - "int quant_block_size, str kv_cache_dtype) -> ()"); - cache_ops.impl("indexer_k_quant_and_cache", torch::kCUDA, - &indexer_k_quant_and_cache); - - cache_ops.def( - "concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()"); - cache_ops.impl("concat_mla_q", torch::kCUDA, &concat_mla_q); - - cache_ops.def( - "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! " - "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); - cache_ops.impl("cp_gather_indexer_k_quant_cache", torch::kCUDA, - &cp_gather_indexer_k_quant_cache); -} - TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) { // Cuda utils