@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 59d0236193
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 01:00 UTC

[10b/n] Migrate custom all-reduce, DeepSeek V4 fused MLA, MiniMax reduce-RMS, and MXFP8 MoE to libtorch stable ABI (#44365)

Signed-off-by: Chris Leonard <chleonar@redhat.com> Signed-off-by: Shengqi Chen <harry-chen@outlook.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>

Chris Leonard committed Jun 3, 2026 at 12:29 UTC 59d0236193a1cb5e88bee9cd4a3692a5587c9558
18 files changed +568 -481
CMakeLists.txt
+13 -12
@@ -311,14 +311,9 @@ set(VLLM_EXT_SRC
311 "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
312 "csrc/quantization/activation_kernels.cu"
313 "csrc/cuda_utils_kernels.cu"
314 - "csrc/custom_all_reduce.cu"
315 - "csrc/torch_bindings.cpp"
316 - "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
314 + "csrc/torch_bindings.cpp")
315
316 if(VLLM_GPU_LANG STREQUAL "CUDA")
319 - list(APPEND VLLM_EXT_SRC
320 - "csrc/minimax_reduce_rms_kernel.cu")
321 -
317 SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
318
319 # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
@@ -505,12 +500,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
500 endif()
501 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND ES_MXFP8_GROUPED_MM_ARCHS)
502 set(SRCS
508 - "csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
509 - "csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu")
503 + "csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
504 + "csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu")
505 set_gencode_flags_for_srcs(
506 SRCS "${SRCS}"
507 CUDA_ARCHS "${ES_MXFP8_GROUPED_MM_ARCHS}")
513 - list(APPEND VLLM_EXT_SRC "${SRCS}")
508 + list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
509 list(APPEND VLLM_GPU_FLAGS "-DENABLE_ES_MXFP8_GROUPED_MM_SM100=1")
510 message(STATUS "Building ES MXFP8 grouped kernels for archs: ${ES_MXFP8_GROUPED_MM_ARCHS}")
511 else()
@@ -600,7 +595,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
595 endif()
596
597 if (VLLM_GPU_LANG STREQUAL "HIP")
603 - # Add QuickReduce kernels
598 + # Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
599 list(APPEND VLLM_EXT_SRC
600 "csrc/custom_quickreduce.cu"
601 )
@@ -651,7 +646,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
646 "csrc/libtorch_stable/attention/paged_attention_v1.cu"
647 "csrc/libtorch_stable/attention/paged_attention_v2.cu"
648 "csrc/libtorch_stable/cache_kernels.cu"
654 - "csrc/libtorch_stable/cache_kernels_fused.cu")
649 + "csrc/libtorch_stable/cache_kernels.cu"
650 + "csrc/libtorch_stable/cache_kernels_fused.cu"
651 + "csrc/libtorch_stable/custom_all_reduce.cu"
652 + "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
653
654 if(VLLM_GPU_LANG STREQUAL "CUDA")
655 list(APPEND VLLM_STABLE_EXT_SRC
@@ -661,7 +659,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
659 "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
660 "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
661 "csrc/libtorch_stable/permute_cols.cu"
664 - "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu")
662 + "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
663 + "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
664 + "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
665 + "csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
666
667 set_gencode_flags_for_srcs(
668 SRCS "${VLLM_STABLE_EXT_SRC}"
csrc/libtorch_stable/custom_all_reduce.cu renamed
+65 -52
@@ -1,7 +1,11 @@
1 -#include <ATen/cuda/Exceptions.h>
2 -#include <c10/cuda/CUDAGuard.h>
3 -#include <c10/cuda/CUDAStream.h>
4 -#include <torch/all.h>
1 +#include "torch_utils.h"
2 +
3 +#include <torch/csrc/stable/macros.h>
4 +#include <torch/csrc/stable/accelerator.h>
5 +#include <torch/csrc/stable/tensor.h>
6 +#include <torch/csrc/stable/ops.h>
7 +#include <torch/headeronly/core/ScalarType.h>
8 +#include <torch/csrc/stable/device.h>
9
10 #include "custom_all_reduce.cuh"
11
@@ -11,7 +15,7 @@ using fptr_t = int64_t;
15 static_assert(sizeof(void*) == sizeof(fptr_t));
16
17 fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
14 - torch::Tensor& rank_data, int64_t rank,
18 + torch::stable::Tensor& rank_data, int64_t rank,
19 bool fully_connected) {
20 int world_size = fake_ipc_ptrs.size();
21 if (world_size > 8)
@@ -25,9 +29,9 @@ fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
29 for (int i = 0; i < world_size; i++) {
30 ipc_ptrs[i] = reinterpret_cast<vllm::Signal*>(fake_ipc_ptrs[i]);
31 }
28 - return (fptr_t) new vllm::CustomAllreduce(ipc_ptrs, rank_data.data_ptr(),
29 - rank_data.numel(), rank, world_size,
30 - fully_connected);
32 + return (fptr_t) new vllm::CustomAllreduce(
33 + ipc_ptrs, rank_data.mutable_data_ptr(), rank_data.numel(), rank,
34 + world_size, fully_connected);
35 }
36
37 /**
@@ -46,10 +50,14 @@ fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
50 * 5. A[None].expand(2, -1, -1, -1): Not OK
51 * 6. A[:, 1:, 1:]: Not OK
52 */
49 -bool _is_weak_contiguous(torch::Tensor& t) {
50 - return t.is_contiguous() ||
51 - (t.storage().nbytes() - t.storage_offset() * t.element_size() ==
52 - t.numel() * t.element_size());
53 +bool _is_weak_contiguous(torch::stable::Tensor& t) {
54 + if (t.is_contiguous()) {
55 + return true;
56 + }
57 + int64_t storage_nbytes = 0;
58 + TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_size(t.get(), &storage_nbytes));
59 + return storage_nbytes - t.storage_offset() * t.element_size() ==
60 + static_cast<int64_t>(t.numel() * t.element_size());
61 }
62
63 /**
@@ -59,42 +67,45 @@ bool _is_weak_contiguous(torch::Tensor& t) {
67 * Otherwise, _reg_buffer is assumed to be IPC-registered and inp is first
68 * copied into _reg_buffer.
69 */
62 -void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
63 - fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes) {
70 +void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
71 + torch::stable::Tensor& out, fptr_t _reg_buffer,
72 + int64_t reg_buffer_sz_bytes) {
73 auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
65 - const at::cuda::OptionalCUDAGuard device_guard(device_of(inp));
66 - auto stream = c10::cuda::getCurrentCUDAStream().stream();
67 -
68 - TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
69 - TORCH_CHECK_EQ(inp.numel(), out.numel());
70 - TORCH_CHECK(_is_weak_contiguous(out));
71 - TORCH_CHECK(_is_weak_contiguous(inp));
74 + const torch::stable::accelerator::DeviceGuard device_guard(
75 + inp.get_device_index());
76 + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
77 +
78 + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
79 + STD_TORCH_CHECK((inp.numel()) == (out.numel()));
80 + STD_TORCH_CHECK(_is_weak_contiguous(out));
81 + STD_TORCH_CHECK(_is_weak_contiguous(inp));
82 auto input_size = inp.numel() * inp.element_size();
83 auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
84 if (reg_buffer) {
75 - TORCH_CHECK_LE(input_size, reg_buffer_sz_bytes);
76 - AT_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.data_ptr(), input_size,
77 - cudaMemcpyDeviceToDevice, stream));
85 + STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes));
86 + STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size,
87 + cudaMemcpyDeviceToDevice, stream));
88 } else {
79 - reg_buffer = inp.data_ptr();
89 + reg_buffer = inp.mutable_data_ptr();
90 }
91 switch (out.scalar_type()) {
82 - case at::ScalarType::Float: {
92 + case torch::headeronly::ScalarType::Float: {
93 fa->allreduce<float>(stream, reinterpret_cast<float*>(reg_buffer),
84 - reinterpret_cast<float*>(out.data_ptr()),
94 + reinterpret_cast<float*>(out.mutable_data_ptr()),
95 out.numel());
96 break;
97 }
88 - case at::ScalarType::Half: {
98 + case torch::headeronly::ScalarType::Half: {
99 fa->allreduce<half>(stream, reinterpret_cast<half*>(reg_buffer),
90 - reinterpret_cast<half*>(out.data_ptr()), out.numel());
100 + reinterpret_cast<half*>(out.mutable_data_ptr()),
101 + out.numel());
102 break;
103 }
104 #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
94 - case at::ScalarType::BFloat16: {
105 + case torch::headeronly::ScalarType::BFloat16: {
106 fa->allreduce<nv_bfloat16>(
107 stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
97 - reinterpret_cast<nv_bfloat16*>(out.data_ptr()), out.numel());
108 + reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), out.numel());
109 break;
110 }
111 #endif
@@ -112,7 +123,7 @@ int64_t meta_size() { return sizeof(vllm::Signal); }
123
124 void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
125 auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
115 - TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
126 + STD_TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
127 void* ipc_ptrs[8];
128 for (int i = 0; i < fake_ipc_ptrs.size(); i++) {
129 ipc_ptrs[i] = reinterpret_cast<void*>(fake_ipc_ptrs[i]);
@@ -143,47 +154,49 @@ void register_graph_buffers(fptr_t _fa,
154 fa->register_graph_buffers(bytes, offsets);
155 }
156
146 -std::tuple<fptr_t, torch::Tensor> allocate_shared_buffer_and_handle(
157 +std::tuple<fptr_t, torch::stable::Tensor> allocate_shared_buffer_and_handle(
158 int64_t size) {
148 - auto device_index = c10::cuda::current_device();
149 - at::DeviceGuard device_guard(at::Device(at::DeviceType::CUDA, device_index));
159 + int device_index;
160 + STD_CUDA_CHECK(cudaGetDevice(&device_index));
161 + const torch::stable::accelerator::DeviceGuard device_guard(device_index);
162 void* buffer;
163 cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
152 - auto stream = c10::cuda::getCurrentCUDAStream().stream();
153 - AT_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
164 + const cudaStream_t stream = get_current_cuda_stream(device_index);
165 + STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
166
167 // Allocate buffer
168 #if defined(USE_ROCM)
169 // data buffers need to be "uncached" for signal on MI200
158 - AT_CUDA_CHECK(
170 + STD_CUDA_CHECK(
171 hipExtMallocWithFlags((void**)&buffer, size, hipDeviceMallocUncached));
172 #else
161 - AT_CUDA_CHECK(cudaMalloc((void**)&buffer, size));
173 + STD_CUDA_CHECK(cudaMalloc((void**)&buffer, size));
174 #endif
163 - AT_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream));
164 - AT_CUDA_CHECK(cudaStreamSynchronize(stream));
165 - AT_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
175 + STD_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream));
176 + STD_CUDA_CHECK(cudaStreamSynchronize(stream));
177 + STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
178
179 // Create IPC memhandle for the allocated buffer.
180 // Will use it in open_mem_handle.
169 - auto options =
170 - torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCPU);
171 - auto handle =
172 - torch::empty({static_cast<int64_t>(sizeof(cudaIpcMemHandle_t))}, options);
173 - AT_CUDA_CHECK(
174 - cudaIpcGetMemHandle((cudaIpcMemHandle_t*)handle.data_ptr(), buffer));
181 + auto handle = torch::stable::empty(
182 + {static_cast<int64_t>(sizeof(cudaIpcMemHandle_t))},
183 + torch::headeronly::ScalarType::Byte, std::nullopt,
184 + torch::stable::Device(torch::stable::DeviceType::CPU));
185 + STD_CUDA_CHECK(cudaIpcGetMemHandle(
186 + (cudaIpcMemHandle_t*)handle.mutable_data_ptr(), buffer));
187
188 return std::make_tuple(reinterpret_cast<fptr_t>(buffer), handle);
189 }
190
179 -fptr_t open_mem_handle(torch::Tensor& mem_handle) {
191 +fptr_t open_mem_handle(torch::stable::Tensor& mem_handle) {
192 void* ipc_ptr;
181 - AT_CUDA_CHECK(cudaIpcOpenMemHandle(
182 - (void**)&ipc_ptr, *((const cudaIpcMemHandle_t*)mem_handle.data_ptr()),
193 + STD_CUDA_CHECK(cudaIpcOpenMemHandle(
194 + (void**)&ipc_ptr,
195 + *((const cudaIpcMemHandle_t*)mem_handle.const_data_ptr()),
196 cudaIpcMemLazyEnablePeerAccess));
197 return reinterpret_cast<fptr_t>(ipc_ptr);
198 }
199
200 void free_shared_buffer(fptr_t buffer) {
188 - AT_CUDA_CHECK(cudaFree(reinterpret_cast<void*>(buffer)));
201 + STD_CUDA_CHECK(cudaFree(reinterpret_cast<void*>(buffer)));
202 }
csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu renamed
+70 -56
@@ -28,7 +28,20 @@
28 * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token
29 */
30
31 +#include "torch_utils.h"
32 +
33 +#include <torch/csrc/stable/macros.h>
34 +#include <torch/csrc/stable/accelerator.h>
35 +#include <torch/csrc/stable/tensor.h>
36 +#include <torch/csrc/stable/ops.h>
37 +#include <torch/headeronly/core/ScalarType.h>
38 +#include <torch/csrc/stable/device.h>
39 +
40 #include <cmath>
41 +#include "cuda_compat.h"
42 +#include "dispatch_utils.h"
43 +#include "type_convert.cuh"
44 +
45 #ifndef USE_ROCM
46 #include <cuda_fp8.h>
47 #else
@@ -37,14 +50,6 @@
50 #include <cuda_runtime.h>
51 #include <type_traits>
52
40 -#include <ATen/cuda/CUDAContext.h>
41 -#include <c10/cuda/CUDAGuard.h>
42 -#include <torch/cuda.h>
43 -
44 -#include "cuda_compat.h"
45 -#include "dispatch_utils.h"
46 -#include "type_convert.cuh"
47 -
53 #ifndef FINAL_MASK
54 #ifdef USE_ROCM
55 #define FINAL_MASK 0xffffffffffffffffULL
@@ -70,7 +75,7 @@ namespace deepseek_v4_fused_ops {
75
76 namespace {
77 inline int getSMVersion() {
73 - auto* props = at::cuda::getCurrentDeviceProperties();
78 + auto* props = get_device_prop();
79 return props->major * 10 + props->minor;
80 }
81 } // namespace
@@ -564,7 +569,7 @@ static void launchFusedDeepseekV4Templated(
569 // bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
570 // unavailable there. Refuse the launch loudly instead of silently
571 // skipping the work.
567 - TORCH_CHECK(
572 + STD_TORCH_CHECK(
573 sm_version >= 80,
574 "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ "
575 "(Ampere or newer); got sm_",
@@ -635,7 +640,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
640 DISPATCH(64)
641 DISPATCH(128)
642 default:
638 - TORCH_CHECK(false,
643 + STD_TORCH_CHECK(false,
644 "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: "
645 "unsupported num_heads_q_padded=",
646 num_heads_q_padded,
@@ -650,71 +655,80 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
655 // ────────────────────────────────────────────────────────────────────────────
656 // Torch op wrapper
657 // ────────────────────────────────────────────────────────────────────────────
653 -torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
654 - torch::Tensor const& q_in, // [N, num_heads_q, 512] bf16
655 - torch::Tensor const& kv, // [N, 512] bf16 (read-only)
656 - torch::Tensor& k_cache, // [num_blocks, block_bytes] uint8
657 - torch::Tensor const& slot_mapping, // [N] int64
658 - torch::Tensor const& position_ids, // [N] int64
659 - torch::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16
660 - int64_t q_head_padded, // padded Q head count for output
658 +torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
659 + torch::stable::Tensor const& q_in, // [N, num_heads_q, 512] bf16
660 + torch::stable::Tensor const& kv, // [N, 512] bf16 (read-only)
661 + torch::stable::Tensor& k_cache, // [num_blocks, block_bytes] uint8
662 + torch::stable::Tensor const& slot_mapping, // [N] int64
663 + torch::stable::Tensor const& position_ids, // [N] int64
664 + torch::stable::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16
665 + int64_t q_head_padded, // padded Q head count for output
666 double eps, int64_t cache_block_size) {
662 - TORCH_CHECK(q_in.is_cuda() && q_in.is_contiguous(),
663 - "q_in must be contiguous CUDA");
664 - TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
665 - TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
666 - TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
667 - "slot_mapping must be int64 CUDA");
668 - TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
669 - "position_ids must be int64 CUDA");
670 - TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
671 - TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512,
672 - "q_in shape [N, num_heads_q, 512]");
673 - TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
674 - TORCH_CHECK(q_in.dtype() == kv.dtype(), "q_in and kv dtype must match");
675 - TORCH_CHECK(q_head_padded >= q_in.size(1),
676 - "q_head_padded must be >= q_in.size(1) (num_heads_q)");
677 - TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
678 - TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
679 - "cos_sin_cache shape [max_pos, 64]");
680 - TORCH_CHECK(cos_sin_cache.dtype() == torch::kFloat32,
681 - "cos_sin_cache must be float32");
667 + STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(),
668 + "q_in must be contiguous CUDA");
669 + STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
670 + "kv must be contiguous CUDA");
671 + STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
672 + STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
673 + slot_mapping.scalar_type() ==
674 + torch::headeronly::ScalarType::Long,
675 + "slot_mapping must be int64 CUDA");
676 + STD_TORCH_CHECK(position_ids.device().is_cuda() &&
677 + position_ids.scalar_type() ==
678 + torch::headeronly::ScalarType::Long,
679 + "position_ids must be int64 CUDA");
680 + STD_TORCH_CHECK(cos_sin_cache.device().is_cuda(), "cos_sin_cache must be CUDA");
681 + STD_TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512,
682 + "q_in shape [N, num_heads_q, 512]");
683 + STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
684 + STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(),
685 + "q_in and kv dtype must match");
686 + STD_TORCH_CHECK(q_head_padded >= q_in.size(1),
687 + "q_head_padded must be >= q_in.size(1) (num_heads_q)");
688 + STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte,
689 + "k_cache must be uint8");
690 + STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
691 + "cos_sin_cache shape [max_pos, 64]");
692 + STD_TORCH_CHECK(cos_sin_cache.scalar_type() ==
693 + torch::headeronly::ScalarType::Float,
694 + "cos_sin_cache must be float32");
695
696 // With DP padding, slot_mapping can be shorter than q/kv/positions.
697 // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
698 // KV quant+insert runs only on the first slot_mapping.size(0) rows.
699 int const num_tokens_full = static_cast<int>(q_in.size(0));
700 int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
688 - TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
689 - static_cast<int>(position_ids.size(0)) == num_tokens_full,
690 - "q/kv/position_ids row counts must match");
691 - TORCH_CHECK(num_tokens_insert <= num_tokens_full,
692 - "slot_mapping must not exceed q row count");
701 + STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
702 + static_cast<int>(position_ids.size(0)) == num_tokens_full,
703 + "q/kv/position_ids row counts must match");
704 + STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
705 + "slot_mapping must not exceed q row count");
706 int const num_heads_q = static_cast<int>(q_in.size(1));
707 int const num_heads_q_padded = static_cast<int>(q_head_padded);
708 int const cache_block_size_i = static_cast<int>(cache_block_size);
709 int const kv_block_stride = static_cast<int>(k_cache.stride(0));
710
698 - at::cuda::OptionalCUDAGuard device_guard(device_of(q_in));
699 - auto stream = at::cuda::getCurrentCUDAStream();
711 + const torch::stable::accelerator::DeviceGuard device_guard(
712 + q_in.get_device_index());
713 + const cudaStream_t stream = get_current_cuda_stream(q_in.get_device_index());
714
715 // Allocate the padded q output. The kernel writes every element (live
716 // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
703 - torch::Tensor q_out = torch::empty(
704 - {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.options());
717 + auto q_out = torch::stable::new_empty(
718 + q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
719
706 - VLLM_DISPATCH_HALF_TYPES(
720 + VLLM_STABLE_DISPATCH_HALF_TYPES(
721 q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
722 using qkv_scalar_t = scalar_t;
723 vllm::deepseek_v4_fused_ops::
724 launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
711 - reinterpret_cast<qkv_scalar_t const*>(q_in.data_ptr()),
712 - reinterpret_cast<qkv_scalar_t*>(q_out.data_ptr()),
713 - reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
714 - reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
715 - reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
716 - reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
717 - cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
725 + reinterpret_cast<qkv_scalar_t const*>(q_in.const_data_ptr()),
726 + reinterpret_cast<qkv_scalar_t*>(q_out.mutable_data_ptr()),
727 + reinterpret_cast<qkv_scalar_t const*>(kv.const_data_ptr()),
728 + reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
729 + slot_mapping.const_data_ptr<int64_t>(),
730 + position_ids.const_data_ptr<int64_t>(),
731 + cos_sin_cache.const_data_ptr<float>(), static_cast<float>(eps),
732 num_tokens_full, num_tokens_insert, num_heads_q,
733 num_heads_q_padded, cache_block_size_i, kv_block_stride,
734 stream);
csrc/libtorch_stable/minimax_reduce_rms_kernel.cu renamed
+71 -58
@@ -15,16 +15,19 @@
15 * limitations under the License.
16 */
17
18 +#include "torch_utils.h"
19 +
20 +#include <torch/csrc/stable/macros.h>
21 +#include <torch/csrc/stable/accelerator.h>
22 +#include <torch/csrc/stable/tensor.h>
23 +#include <torch/csrc/stable/ops.h>
24 +#include <torch/headeronly/core/ScalarType.h>
25 +#include <torch/csrc/stable/device.h>
26 +
27 #include <cooperative_groups.h>
28 #include <cuda_runtime.h>
29
21 -#include <torch/cuda.h>
22 -#include <ATen/cuda/CUDAContext.h>
23 -#include <c10/cuda/CUDAGuard.h>
24 -
30 #include "cuda_compat.h"
26 -#include "cuda_utils.h"
27 -#include "core/registration.h"
31 #include "minimax_reduce_rms_kernel.h"
32
33 #include <algorithm>
@@ -611,7 +614,7 @@ int get_sm_count() {
614 static int sm_count = 0;
615 if (sm_count == 0) {
616 int device_id;
614 - CUDA_CHECK(cudaGetDevice(&device_id));
617 + STD_CUDA_CHECK(cudaGetDevice(&device_id));
618 cudaDeviceProp device_prop;
619 cudaGetDeviceProperties(&device_prop, device_id);
620 sm_count = device_prop.multiProcessorCount;
@@ -621,13 +624,13 @@ int get_sm_count() {
624
625 inline int getSMVersion(bool queryRealSmArch = false) {
626 int device{-1};
624 - CUDA_CHECK(cudaGetDevice(&device));
627 + STD_CUDA_CHECK(cudaGetDevice(&device));
628 int sm_major = 0;
629 int sm_minor = 0;
627 - CUDA_CHECK(cudaDeviceGetAttribute(&sm_major,
628 - cudaDevAttrComputeCapabilityMajor, device));
629 - CUDA_CHECK(cudaDeviceGetAttribute(&sm_minor,
630 - cudaDevAttrComputeCapabilityMinor, device));
630 + STD_CUDA_CHECK(cudaDeviceGetAttribute(
631 + &sm_major, cudaDevAttrComputeCapabilityMajor, device));
632 + STD_CUDA_CHECK(cudaDeviceGetAttribute(
633 + &sm_minor, cudaDevAttrComputeCapabilityMinor, device));
634 int sm = sm_major * 10 + sm_minor;
635 if (sm == 121 && !queryRealSmArch) {
636 return 120;
@@ -639,7 +642,7 @@ template <typename KernelFunc>
642 int get_max_active_blocks(KernelFunc kernel, int block_size,
643 int dynamic_smem = 0) {
644 int max_active = 0;
642 - CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
645 + STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
646 &max_active, kernel, block_size, dynamic_smem));
647 return std::max(max_active, 1);
648 }
@@ -678,27 +681,27 @@ void minimax_reduce_rms_kernel_launcher(MiniMaxReduceRMSParams const& params) {
681 cfg.attrs = attribute;
682 cfg.numAttrs = SM >= 90 ? 2 : 0;
683
681 - CUDA_CHECK(cudaLaunchKernelEx(
684 + STD_CUDA_CHECK(cudaLaunchKernelEx(
685 &cfg, minimax_reduce_rms_kernel_lamport<DType, NRanks>, params));
686 }
687
688 template <typename DType, int NRanks, int OriginQDim, int OriginKDim>
689 void minimax_reduce_rms_kernel_launcher_float4(
690 MiniMaxReduceRMSParams const& params) {
688 - TORCH_CHECK(params.size_q % params.hidden_dim == 0);
689 - TORCH_CHECK(params.hidden_dim % kElemsPerAccess<DType> == 0);
691 + STD_TORCH_CHECK(params.size_q % params.hidden_dim == 0);
692 + STD_TORCH_CHECK(params.hidden_dim % kElemsPerAccess<DType> == 0);
693 if (params.stride_q > 0) {
691 - TORCH_CHECK(params.stride_q % kElemsPerAccess<DType> == 0);
694 + STD_TORCH_CHECK(params.stride_q % kElemsPerAccess<DType> == 0);
695 }
693 - TORCH_CHECK(params.allreduce_in_k != nullptr,
694 - "float4 QK kernel requires K input");
695 - TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k);
696 - TORCH_CHECK(params.size_k % params.hidden_dim_k == 0);
697 - TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess<DType> == 0);
698 - TORCH_CHECK(params.size_q / params.hidden_dim ==
699 - params.size_k / params.hidden_dim_k);
696 + STD_TORCH_CHECK(params.allreduce_in_k != nullptr,
697 + "float4 QK kernel requires K input");
698 + STD_TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k);
699 + STD_TORCH_CHECK(params.size_k % params.hidden_dim_k == 0);
700 + STD_TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess<DType> == 0);
701 + STD_TORCH_CHECK(params.size_q / params.hidden_dim ==
702 + params.size_k / params.hidden_dim_k);
703 if (params.stride_k > 0) {
701 - TORCH_CHECK(params.stride_k % kElemsPerAccess<DType> == 0);
704 + STD_TORCH_CHECK(params.stride_k % kElemsPerAccess<DType> == 0);
705 }
706
707 int token_num = params.size_q / params.hidden_dim;
@@ -746,7 +749,7 @@ void minimax_reduce_rms_kernel_launcher_float4(
749 cfg.attrs = attribute;
750 cfg.numAttrs = SM >= 90 ? 2 : 0;
751
749 - CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params));
752 + STD_CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params));
753 }
754
755 template <int NRanks>
@@ -759,21 +762,21 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) {
762 (params.hidden_dim * params.nranks == 6144) &&
763 (params.hidden_dim_k * params.nranks == 1024);
764
762 - if (params.dtype == at::ScalarType::Half) {
765 + if (params.dtype == torch::headeronly::ScalarType::Half) {
766 if (use_float4) {
767 minimax_reduce_rms_kernel_launcher_float4<half, NRanks, 6144, 1024>(
768 params);
769 } else {
770 minimax_reduce_rms_kernel_launcher<half, NRanks>(params);
771 }
769 - } else if (params.dtype == at::ScalarType::BFloat16) {
772 + } else if (params.dtype == torch::headeronly::ScalarType::BFloat16) {
773 if (use_float4) {
774 minimax_reduce_rms_kernel_launcher_float4<__nv_bfloat16, NRanks, 6144,
775 1024>(params);
776 } else {
777 minimax_reduce_rms_kernel_launcher<__nv_bfloat16, NRanks>(params);
778 }
776 - } else if (params.dtype == at::ScalarType::Float) {
779 + } else if (params.dtype == torch::headeronly::ScalarType::Float) {
780 if (use_float4) {
781 minimax_reduce_rms_kernel_launcher_float4<float, NRanks, 6144, 1024>(
782 params);
@@ -781,7 +784,7 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) {
784 minimax_reduce_rms_kernel_launcher<float, NRanks>(params);
785 }
786 } else {
784 - TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op");
787 + STD_TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op");
788 }
789 }
790
@@ -795,16 +798,18 @@ void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) {
798 } else if (params.nranks == 16) {
799 dispatch_dtype<16>(params);
800 } else {
798 - TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!");
801 + STD_TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!");
802 }
803 }
804 } // namespace tensorrt_llm
805 } // namespace vllm
806
804 -torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
805 - torch::Tensor const& norm_weight,
806 - torch::Tensor workspace, int64_t const rank,
807 - int64_t const nranks, double const eps) {
807 +torch::stable::Tensor minimax_allreduce_rms(
808 + torch::stable::Tensor const& input,
809 + torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
810 + int64_t const rank, int64_t const nranks, double const eps) {
811 + const torch::stable::accelerator::DeviceGuard device_guard(
812 + input.get_device_index());
813 auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
814
815 allreduce_params.nranks = static_cast<int>(nranks);
@@ -815,12 +820,12 @@ torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
820 allreduce_params.stride_q = allreduce_params.hidden_dim;
821 allreduce_params.workspace =
822 reinterpret_cast<void**>(workspace.mutable_data_ptr());
818 - allreduce_params.allreduce_in = input.data_ptr();
819 - allreduce_params.rms_gamma = norm_weight.data_ptr();
823 + allreduce_params.allreduce_in = const_cast<void*>(input.const_data_ptr());
824 + allreduce_params.rms_gamma = const_cast<void*>(norm_weight.const_data_ptr());
825 allreduce_params.rms_eps = static_cast<float>(eps);
821 - allreduce_params.stream = at::cuda::getCurrentCUDAStream(input.get_device());
826 + allreduce_params.stream = get_current_cuda_stream(input.get_device_index());
827
823 - torch::Tensor rms_norm_out = torch::empty_like(input);
828 + torch::stable::Tensor rms_norm_out = torch::stable::empty_like(input);
829 allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr();
830
831 vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params);
@@ -828,26 +833,33 @@ torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
833 return rms_norm_out;
834 }
835
831 -std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
832 - torch::Tensor qkv, torch::Tensor const& norm_weight_q,
833 - torch::Tensor const& norm_weight_k, torch::Tensor workspace,
834 - int64_t const q_size, int64_t const kv_size, int64_t const rank,
835 - int64_t const nranks, double const eps) {
836 - TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D");
837 - TORCH_CHECK(qkv.is_contiguous(),
838 - "minimax_allreduce_rms_qk: qkv must be contiguous");
836 +std::tuple<torch::stable::Tensor, torch::stable::Tensor>
837 +minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
838 + torch::stable::Tensor const& norm_weight_q,
839 + torch::stable::Tensor const& norm_weight_k,
840 + torch::stable::Tensor workspace, int64_t const q_size,
841 + int64_t const kv_size, int64_t const rank,
842 + int64_t const nranks, double const eps) {
843 + STD_TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D");
844 + STD_TORCH_CHECK(qkv.is_contiguous(),
845 + "minimax_allreduce_rms_qk: qkv must be contiguous");
846 int64_t qkv_dim = qkv.size(-1);
840 - TORCH_CHECK(qkv_dim == q_size + 2 * kv_size,
841 - "minimax_allreduce_rms_qk: qkv last dim must equal "
842 - "q_size + 2 * kv_size");
843 - TORCH_CHECK(rank < nranks,
844 - "minimax_allreduce_rms_qk: rank must be less than nranks");
847 + STD_TORCH_CHECK(qkv_dim == q_size + 2 * kv_size,
848 + "minimax_allreduce_rms_qk: qkv last dim must equal "
849 + "q_size + 2 * kv_size");
850 + STD_TORCH_CHECK(rank < nranks,
851 + "minimax_allreduce_rms_qk: rank must be less than nranks");
852 +
853 + const torch::stable::accelerator::DeviceGuard device_guard(
854 + qkv.get_device_index());
855
856 int64_t num_tokens = qkv.size(0);
857 int elem_bytes = qkv.element_size();
858
849 - torch::Tensor q_out = torch::empty({num_tokens, q_size}, qkv.options());
850 - torch::Tensor k_out = torch::empty({num_tokens, kv_size}, qkv.options());
859 + torch::stable::Tensor q_out =
860 + torch::stable::new_empty(qkv, {num_tokens, q_size}, qkv.scalar_type());
861 + torch::stable::Tensor k_out =
862 + torch::stable::new_empty(qkv, {num_tokens, kv_size}, qkv.scalar_type());
863
864 auto params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
865 params.nranks = static_cast<int>(nranks);
@@ -863,13 +875,14 @@ std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
875 params.stride_k_out = 0; // k_out is contiguous; kernel uses hidden_dim_k
876 params.workspace = reinterpret_cast<void**>(workspace.mutable_data_ptr());
877
866 - uint8_t* base = static_cast<uint8_t*>(qkv.data_ptr());
878 + uint8_t* base =
879 + const_cast<uint8_t*>(static_cast<const uint8_t*>(qkv.const_data_ptr()));
880 params.allreduce_in = base;
881 params.allreduce_in_k = base + q_size * elem_bytes;
869 - params.rms_gamma = norm_weight_q.data_ptr();
870 - params.rms_gamma_k = norm_weight_k.data_ptr();
882 + params.rms_gamma = const_cast<void*>(norm_weight_q.const_data_ptr());
883 + params.rms_gamma_k = const_cast<void*>(norm_weight_k.const_data_ptr());
884 params.rms_eps = static_cast<float>(eps);
872 - params.stream = at::cuda::getCurrentCUDAStream(qkv.get_device());
885 + params.stream = get_current_cuda_stream(qkv.get_device_index());
886
887 params.rms_norm_out = q_out.mutable_data_ptr();
888 params.rms_norm_out_k = k_out.mutable_data_ptr();
csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu new
+69
@@ -0,0 +1,69 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +// Adapted from SGLang:
4 +// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu
5 +
6 +#include <torch/csrc/stable/library.h>
7 +#include <torch/csrc/stable/tensor.h>
8 +#include "libtorch_stable/torch_utils.h"
9 +
10 +#include "cutlass_mxfp8_grouped_mm_launcher.cuh"
11 +
12 +void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a,
13 + const torch::stable::Tensor& b,
14 + const torch::stable::Tensor& sfa,
15 + const torch::stable::Tensor& sfb,
16 + torch::stable::Tensor& d,
17 + const torch::stable::Tensor& problem_sizes,
18 + const torch::stable::Tensor& expert_offsets,
19 + const torch::stable::Tensor& blockscale_offsets) {
20 +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
21 + STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
22 + STD_TORCH_CHECK(problem_sizes.size(1) == 3,
23 + "problem_sizes must have shape (num_experts, 3)");
24 + STD_TORCH_CHECK(
25 + problem_sizes.size(0) == expert_offsets.size(0),
26 + "Number of experts in problem_sizes must match expert_offsets");
27 + STD_TORCH_CHECK(
28 + problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
29 + "problem_sizes must be int32");
30 + STD_TORCH_CHECK(
31 + expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
32 + "expert_offsets must be int32");
33 + STD_TORCH_CHECK(
34 + blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
35 + "blockscale_offsets must be int32");
36 + STD_TORCH_CHECK(a.dim() == 2,
37 + "a must be a 2D tensor of shape (num_tokens, k)");
38 + STD_TORCH_CHECK(b.dim() == 3,
39 + "b must be a 3D tensor of shape (num_experts, k, n)");
40 + STD_TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0,
41 + "k should align 128");
42 + STD_TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128");
43 + STD_TORCH_CHECK(a.stride(1) == 1, "a must be row major");
44 + STD_TORCH_CHECK(b.stride(1) == 1, "b must be column major");
45 +
46 + const torch::stable::accelerator::DeviceGuard device_guard(
47 + a.get_device_index());
48 + auto stream = get_current_cuda_stream(a.get_device_index());
49 + if (d.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
50 + expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
51 + cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
52 + blockscale_offsets, stream);
53 + } else if (d.scalar_type() == torch::headeronly::ScalarType::Half) {
54 + expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
55 + cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
56 + blockscale_offsets, stream);
57 + } else {
58 + STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
59 + }
60 +#else
61 + STD_TORCH_CHECK(false,
62 + "No implemented cutlass_mxfp8_grouped_mm for "
63 + "current device");
64 +#endif
65 +}
66 +
67 +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
68 + m.impl("cutlass_mxfp8_grouped_mm", TORCH_BOX(&cutlass_mxfp8_grouped_mm));
69 +}
csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_functor.cuh renamed
csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh renamed
+73 -54
@@ -4,9 +4,9 @@
4 // https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_launcher.cuh
5
6 #pragma once
7 -#include <ATen/cuda/CUDAContext.h>
8 -#include <c10/cuda/CUDAGuard.h>
9 -#include <torch/all.h>
7 +
8 +#include <torch/csrc/stable/tensor.h>
9 +#include <torch/headeronly/util/Exception.h>
10
11 #include <cassert>
12 #include <iostream>
@@ -15,18 +15,22 @@
15 #include "cute/tensor.hpp"
16 #include "cutlass_mxfp8_grouped_mm_functor.cuh"
17 #include "cutlass_mxfp8_grouped_mm_traits.cuh"
18 +#include "libtorch_stable/torch_utils.h"
19
20 namespace expert_specialization {
21
22 template <typename GemmTraits>
23 void cutlass_mxfp8_grouped_mm_pre_compute(
23 - torch::Tensor& a_ptrs, torch::Tensor& b_ptrs, torch::Tensor& sfa_ptrs,
24 - torch::Tensor& sfb_ptrs, torch::Tensor& d_ptrs, torch::Tensor& stride_a,
25 - torch::Tensor& stride_b, torch::Tensor& stride_d, torch::Tensor& layout_sfa,
26 - torch::Tensor& layout_sfb, const torch::Tensor& a, const torch::Tensor& b,
27 - const torch::Tensor& sfa, const torch::Tensor& sfb, const torch::Tensor& d,
28 - const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets,
29 - const torch::Tensor& blockscale_offsets, cudaStream_t stream) {
24 + torch::stable::Tensor& a_ptrs, torch::stable::Tensor& b_ptrs,
25 + torch::stable::Tensor& sfa_ptrs, torch::stable::Tensor& sfb_ptrs,
26 + torch::stable::Tensor& d_ptrs, torch::stable::Tensor& stride_a,
27 + torch::stable::Tensor& stride_b, torch::stable::Tensor& stride_d,
28 + torch::stable::Tensor& layout_sfa, torch::stable::Tensor& layout_sfb,
29 + const torch::stable::Tensor& a, const torch::stable::Tensor& b,
30 + const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb,
31 + const torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes,
32 + const torch::stable::Tensor& expert_offsets,
33 + const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) {
34 using OffsetFunctor = CutlassMxfp8GroupedMmOffsetFunctor<GemmTraits>;
35 using ElementA = typename OffsetFunctor::ElementA;
36 using ElementB = typename OffsetFunctor::ElementB;
@@ -42,10 +46,10 @@ void cutlass_mxfp8_grouped_mm_pre_compute(
46 using StrideB = typename StrideFunctor::StrideB;
47 using StrideD = typename StrideFunctor::StrideD;
48
45 - int num_experts = (int)expert_offsets.size(0);
46 - TORCH_CHECK(num_experts <= 1024,
47 - "Number of experts cannot exceed 1024, the maximum number of "
48 - "threads per block.");
49 + int num_experts = static_cast<int>(expert_offsets.size(0));
50 + STD_TORCH_CHECK(num_experts <= 1024,
51 + "Number of experts cannot exceed 1024, the maximum number of "
52 + "threads per block.");
53
54 OffsetFunctor offset_functor(
55 reinterpret_cast<int*>(expert_offsets.data_ptr()),
@@ -72,13 +76,18 @@ void cutlass_mxfp8_grouped_mm_pre_compute(
76 }
77
78 template <typename GemmTraits>
75 -void cutlass_mxfp8_grouped_mm(
76 - const torch::Tensor& a_ptrs, const torch::Tensor& b_ptrs,
77 - const torch::Tensor& sfa_ptrs, const torch::Tensor& sfb_ptrs,
78 - const torch::Tensor& d_ptrs, const torch::Tensor& stride_a,
79 - const torch::Tensor& stride_b, const torch::Tensor& stride_d,
80 - const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb,
81 - const torch::Tensor& problem_sizes, cudaStream_t stream) {
79 +void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a_ptrs,
80 + const torch::stable::Tensor& b_ptrs,
81 + const torch::stable::Tensor& sfa_ptrs,
82 + const torch::stable::Tensor& sfb_ptrs,
83 + const torch::stable::Tensor& d_ptrs,
84 + const torch::stable::Tensor& stride_a,
85 + const torch::stable::Tensor& stride_b,
86 + const torch::stable::Tensor& stride_d,
87 + const torch::stable::Tensor& layout_sfa,
88 + const torch::stable::Tensor& layout_sfb,
89 + const torch::stable::Tensor& problem_sizes,
90 + cudaStream_t stream) {
91 using Gemm = typename GemmTraits::Gemm;
92 using ElementA = typename Gemm::ElementA;
93 using ElementB = typename Gemm::ElementB;
@@ -93,13 +102,12 @@ void cutlass_mxfp8_grouped_mm(
102 typename GemmTraits::ProblemShape::UnderlyingProblemShape;
103
104 cutlass::KernelHardwareInfo hw_info;
96 - hw_info.device_id = c10::cuda::current_device();
97 - hw_info.sm_count =
98 - at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
105 + hw_info.device_id = d_ptrs.get_device_index();
106 + hw_info.sm_count = get_device_prop()->multiProcessorCount;
107 hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster;
108 hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster;
109
102 - int num_experts = (int)problem_sizes.size(0);
110 + int num_experts = static_cast<int>(problem_sizes.size(0));
111
112 UnderlyingProblemShape* underlying_problem_shape =
113 reinterpret_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
@@ -127,44 +135,55 @@ void cutlass_mxfp8_grouped_mm(
135 Gemm gemm;
136
137 auto can_implement_status = gemm.can_implement(arguments);
130 - TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
131 - "Failed to implement GEMM");
138 + STD_TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
139 + "Failed to implement GEMM");
140
133 - torch::TensorOptions options_uint8 =
134 - torch::TensorOptions().dtype(torch::kUInt8).device(d_ptrs.device());
141 size_t workspace_size = gemm.get_workspace_size(arguments);
136 - torch::Tensor workspace = torch::empty(workspace_size, options_uint8);
142 + torch::stable::Tensor workspace = torch::stable::empty(
143 + {static_cast<int64_t>(workspace_size)},
144 + torch::headeronly::ScalarType::Byte, std::nullopt, d_ptrs.device());
145
146 auto status = gemm.initialize(arguments, workspace.data_ptr(), stream);
139 - TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");
147 + STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
148 + "Failed to initialize GEMM");
149
150 status = gemm.run(stream, nullptr, true); // Enable PDL
142 - TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
151 + STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
152 }
153
154 template <typename OutType>
155 void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
147 - const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& sfa,
148 - const torch::Tensor& sfb, torch::Tensor& d,
149 - const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets,
150 - const torch::Tensor& blockscale_offsets, cudaStream_t stream) {
151 - int num_experts = (int)problem_sizes.size(0);
152 - torch::TensorOptions options_int64 =
153 - torch::TensorOptions().dtype(torch::kInt64).device(a.device());
154 - torch::TensorOptions options_int32 =
155 - torch::TensorOptions().dtype(torch::kInt32).device(a.device());
156 -
157 - torch::Tensor a_ptrs = torch::empty(num_experts, options_int64);
158 - torch::Tensor b_ptrs = torch::empty(num_experts, options_int64);
159 - torch::Tensor sfa_ptrs = torch::empty(num_experts, options_int64);
160 - torch::Tensor sfb_ptrs = torch::empty(num_experts, options_int64);
161 - torch::Tensor d_ptrs = torch::empty(num_experts, options_int64);
162 -
163 - torch::Tensor stride_a = torch::empty(num_experts, options_int64);
164 - torch::Tensor stride_b = torch::empty(num_experts, options_int64);
165 - torch::Tensor stride_d = torch::empty(num_experts, options_int64);
166 - torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int32);
167 - torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int32);
156 + const torch::stable::Tensor& a, const torch::stable::Tensor& b,
157 + const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb,
158 + torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes,
159 + const torch::stable::Tensor& expert_offsets,
160 + const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) {
161 + int num_experts = static_cast<int>(problem_sizes.size(0));
162 + auto device = a.device();
163 +
164 + torch::stable::Tensor a_ptrs = torch::stable::empty(
165 + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
166 + torch::stable::Tensor b_ptrs = torch::stable::empty(
167 + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
168 + torch::stable::Tensor sfa_ptrs = torch::stable::empty(
169 + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
170 + torch::stable::Tensor sfb_ptrs = torch::stable::empty(
171 + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
172 + torch::stable::Tensor d_ptrs = torch::stable::empty(
173 + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
174 +
175 + torch::stable::Tensor stride_a = torch::stable::empty(
176 + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
177 + torch::stable::Tensor stride_b = torch::stable::empty(
178 + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
179 + torch::stable::Tensor stride_d = torch::stable::empty(
180 + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
181 + torch::stable::Tensor layout_sfa =
182 + torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int,
183 + std::nullopt, device);
184 + torch::stable::Tensor layout_sfb =
185 + torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int,
186 + std::nullopt, device);
187
188 using GemmTraits = CutlassMxfp8GroupedMmGemmTraits<MMA1SMConfig, OutType>;
189 cutlass_mxfp8_grouped_mm_pre_compute<GemmTraits>(
@@ -176,4 +195,4 @@ void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
195 layout_sfa, layout_sfb, problem_sizes, stream);
196 }
197
179 -} // namespace expert_specialization
\ No newline at end of file
198 +} // namespace expert_specialization
csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_traits.cuh renamed
csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu new
+66
@@ -0,0 +1,66 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +// Adapted from SGLang:
4 +// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu
5 +
6 +#include <torch/csrc/stable/library.h>
7 +#include <torch/csrc/stable/tensor.h>
8 +#include "libtorch_stable/torch_utils.h"
9 +
10 +#include "mxfp8_experts_quant.cuh"
11 +
12 +void mxfp8_experts_quant(const torch::stable::Tensor& input,
13 + const torch::stable::Tensor& problem_sizes,
14 + const torch::stable::Tensor& expert_offsets,
15 + const torch::stable::Tensor& blockscale_offsets,
16 + torch::stable::Tensor& quant_output,
17 + torch::stable::Tensor& scale_factor) {
18 +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
19 + STD_TORCH_CHECK(input.dim() == 2, "input must be 2D tensor");
20 + STD_TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128");
21 + STD_TORCH_CHECK(input.stride(1) == 1, "input must be row major");
22 + STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
23 + STD_TORCH_CHECK(
24 + problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
25 + "problem_sizes must be int32");
26 + STD_TORCH_CHECK(
27 + expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
28 + "expert_offsets must be int32");
29 + STD_TORCH_CHECK(
30 + blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
31 + "blockscale_offsets must be int32");
32 +
33 + auto groups = problem_sizes.size(0);
34 + STD_TORCH_CHECK(
35 + expert_offsets.dim() == 1 && expert_offsets.size(0) == groups,
36 + "expert_offsets must be 1D and have size equal to the number of groups");
37 + STD_TORCH_CHECK(
38 + blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups,
39 + "blockscale_offsets must be 1D and have size equal to the number of "
40 + "groups");
41 +
42 + const torch::stable::accelerator::DeviceGuard device_guard(
43 + input.get_device_index());
44 + if (input.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
45 + expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>(
46 + input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
47 + scale_factor);
48 + } else if (input.scalar_type() == torch::headeronly::ScalarType::Half) {
49 + expert_specialization::launch_mxfp8_experts_quant<__half>(
50 + input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
51 + scale_factor);
52 + } else {
53 + STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
54 + }
55 +#else
56 + STD_TORCH_CHECK(false,
57 + "No implemented mxfp8_experts_quant for "
58 + "current device");
59 +#endif
60 +}
61 +
62 +// Registered here (not torch_bindings.cpp) because ENABLE_ES_MXFP8_GROUPED_MM
63 +// is applied only under COMPILE_LANGUAGE:CUDA.
64 +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
65 + m.impl("mxfp8_experts_quant", TORCH_BOX(&mxfp8_experts_quant));
66 +}
csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cuh renamed
+16 -14
@@ -4,16 +4,19 @@
4 // https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh
5
6 #pragma once
7 -#include <ATen/cuda/CUDAContext.h>
8 -#include <c10/cuda/CUDAGuard.h>
7 #include <cuda.h>
8 #include <cuda_bf16.h>
9 #include <cuda_fp16.h>
12 -#include <torch/all.h>
10 +
11 +#include <torch/csrc/inductor/aoti_torch/c/shim.h>
12 +#include <torch/csrc/stable/macros.h>
13 +#include <torch/csrc/stable/tensor.h>
14 +#include <torch/headeronly/util/Exception.h>
15
16 #include <cuda/ptx>
17
18 #include "cute/tensor.hpp"
19 +#include "libtorch_stable/torch_utils.h"
20
21 namespace expert_specialization {
22
@@ -356,12 +359,12 @@ __global__ void mxfp8_experts_quant_kernel(
359 }
360
361 template <typename T_IN>
359 -void launch_mxfp8_experts_quant(const torch::Tensor& input,
360 - const torch::Tensor& problem_sizes,
361 - const torch::Tensor& expert_offsets,
362 - const torch::Tensor& blockscale_offsets,
363 - torch::Tensor& quant_output,
364 - torch::Tensor& scale_factor) {
362 +void launch_mxfp8_experts_quant(const torch::stable::Tensor& input,
363 + const torch::stable::Tensor& problem_sizes,
364 + const torch::stable::Tensor& expert_offsets,
365 + const torch::stable::Tensor& blockscale_offsets,
366 + torch::stable::Tensor& quant_output,
367 + torch::stable::Tensor& scale_factor) {
368 ThrLayout thr_layout{};
369 ValLayout val_layout{};
370 SfR2SThrLayout r2s_thr_layout{};
@@ -386,19 +389,18 @@ void launch_mxfp8_experts_quant(const torch::Tensor& input,
389 CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout); // Tiler_MN: (16, 4)
390
391 int max_active_blocks_per_sm = -1;
389 - AT_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
392 + STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
393 &max_active_blocks_per_sm,
394 mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
395 decltype(tiled_copy_r2g),
396 decltype(tiled_copy_r2s)>,
397 THREAD_BLOCK_SIZE, 0));
398
396 - dim3 grid(at::cuda::getCurrentDeviceProperties()->multiProcessorCount *
397 - max_active_blocks_per_sm,
399 + dim3 grid(get_device_prop()->multiProcessorCount * max_active_blocks_per_sm,
400 1, 1);
401 dim3 block(THREAD_BLOCK_SIZE, 1, 1);
400 - int num_experts = (int)problem_sizes.size(0);
401 - auto stream = at::cuda::getCurrentCUDAStream();
402 + int num_experts = static_cast<int>(problem_sizes.size(0));
403 + auto stream = get_current_cuda_stream(input.get_device_index());
404 mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
405 decltype(tiled_copy_r2g), decltype(tiled_copy_r2s)>
406 <<<grid, block, 0, stream>>>(
csrc/libtorch_stable/ops.h
+41
@@ -231,6 +231,27 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
231 torch::stable::Tensor& position_ids,
232 int64_t forced_token_heads_per_warp);
233
234 +torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
235 + torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
236 + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
237 + torch::stable::Tensor const& position_ids,
238 + torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
239 + double eps, int64_t cache_block_size);
240 +
241 +#ifndef USE_ROCM
242 +torch::stable::Tensor minimax_allreduce_rms(
243 + torch::stable::Tensor const& input,
244 + torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
245 + int64_t const rank, int64_t const nranks, double const eps);
246 +std::tuple<torch::stable::Tensor, torch::stable::Tensor>
247 +minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
248 + torch::stable::Tensor const& norm_weight_q,
249 + torch::stable::Tensor const& norm_weight_k,
250 + torch::stable::Tensor workspace, int64_t const q_size,
251 + int64_t const kv_size, int64_t const rank,
252 + int64_t const nranks, double const eps);
253 +#endif
254 +
255 // Sampler kernels (shared CUDA/ROCm)
256 void apply_repetition_penalties_(
257 torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -273,6 +294,26 @@ void selective_scan_fwd(
294 const std::optional<torch::stable::Tensor>& cu_chunk_seqlen,
295 const std::optional<torch::stable::Tensor>& last_chunk_indices);
296
297 +using fptr_t = int64_t;
298 +fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
299 + torch::stable::Tensor& rank_data, int64_t rank,
300 + bool fully_connected);
301 +void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
302 + torch::stable::Tensor& out, fptr_t reg_buffer,
303 + int64_t reg_buffer_sz_bytes);
304 +void dispose(fptr_t _fa);
305 +int64_t meta_size();
306 +void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
307 +std::tuple<std::vector<int64_t>, std::vector<int64_t>>
308 +get_graph_buffer_ipc_meta(fptr_t _fa);
309 +void register_graph_buffers(fptr_t _fa,
310 + const std::vector<std::vector<int64_t>>& handles,
311 + const std::vector<std::vector<int64_t>>& offsets);
312 +std::tuple<int64_t, torch::stable::Tensor> allocate_shared_buffer_and_handle(
313 + int64_t size);
314 +int64_t open_mem_handle(torch::stable::Tensor& mem_handle);
315 +void free_shared_buffer(int64_t buffer);
316 +
317 // Activation kernels (shared CUDA/ROCm)
318 void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
319 void silu_and_mul_clamp(torch::stable::Tensor& out,
csrc/libtorch_stable/torch_bindings.cpp
+63
@@ -337,6 +337,24 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
337 "bool is_neox, Tensor position_ids, "
338 "int forced_token_heads_per_warp=-1) -> ()");
339
340 + ops.def(
341 + "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
342 + "Tensor q_in, Tensor kv, Tensor! k_cache, "
343 + "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
344 + "int q_head_padded, float eps, int cache_block_size) -> Tensor");
345 +
346 +#ifndef USE_ROCM
347 + ops.def(
348 + "minimax_allreduce_rms("
349 + "Tensor input, Tensor norm_weight, Tensor workspace, "
350 + "int rank, int nranks, float eps) -> Tensor");
351 + ops.def(
352 + "minimax_allreduce_rms_qk("
353 + "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, "
354 + "Tensor workspace, int q_size, int kv_size, int rank, int nranks, "
355 + "float eps) -> (Tensor, Tensor)");
356 +#endif
357 +
358 // Apply repetition penalties to logits in-place.
359 ops.def(
360 "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -571,6 +589,12 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
589 // Positional encoding kernels (shared CUDA/ROCm)
590 ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
591 ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
592 + ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
593 + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert));
594 +#ifndef USE_ROCM
595 + ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
596 + ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
597 +#endif
598
599 // Sampler kernels (shared CUDA/ROCm)
600 ops.impl("apply_repetition_penalties_",
@@ -725,6 +749,45 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
749 "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
750 }
751
752 +STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) {
753 + custom_ar.def(
754 + "init_custom_ar(int[] ipc_tensors, Tensor rank_data, "
755 + "int rank, bool fully_connected) -> int");
756 + custom_ar.def(
757 + "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, "
758 + "int reg_buffer_sz_bytes) -> ()");
759 + custom_ar.def("dispose(int fa) -> ()");
760 + custom_ar.def("meta_size() -> int");
761 + custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()");
762 + custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])");
763 + custom_ar.def(
764 + "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()");
765 + custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)");
766 + custom_ar.def("open_mem_handle(Tensor mem_handle) -> int");
767 + custom_ar.def("free_shared_buffer(int ptr) -> ()");
768 +}
769 +
770 +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ar) {
771 + custom_ar.impl("init_custom_ar", TORCH_BOX(&init_custom_ar));
772 + custom_ar.impl("all_reduce", TORCH_BOX(&all_reduce));
773 +}
774 +
775 +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) {
776 + custom_ar.impl("open_mem_handle", TORCH_BOX(&open_mem_handle));
777 +}
778 +
779 +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) {
780 + custom_ar.impl("dispose", TORCH_BOX(&dispose));
781 + custom_ar.impl("meta_size", TORCH_BOX(&meta_size));
782 + custom_ar.impl("register_buffer", TORCH_BOX(&register_buffer));
783 + custom_ar.impl("get_graph_buffer_ipc_meta",
784 + TORCH_BOX(&get_graph_buffer_ipc_meta));
785 + custom_ar.impl("register_graph_buffers", TORCH_BOX(&register_graph_buffers));
786 + custom_ar.impl("allocate_shared_buffer_and_handle",
787 + TORCH_BOX(&allocate_shared_buffer_and_handle));
788 + custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer));
789 +}
790 +
791 STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) {
792 ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch));
793 }
csrc/minimax_reduce_rms_kernel.h
+2 -2
@@ -19,7 +19,7 @@
19 #include <cuda_bf16.h>
20 #include <cuda_fp16.h>
21
22 -#include <torch/types.h>
22 +#include <torch/headeronly/core/ScalarType.h>
23
24 namespace vllm {
25 namespace tensorrt_llm {
@@ -51,7 +51,7 @@ static constexpr int kElemsPerAccess = ElemsPerAccess<DType>::value;
51 struct MiniMaxReduceRMSParams {
52 int nranks{};
53 int rank{};
54 - at::ScalarType dtype{at::ScalarType::Undefined};
54 + torch::headeronly::ScalarType dtype{torch::headeronly::ScalarType::Undefined};
55 int size_q{};
56 int hidden_dim{};
57 int size_k{};
csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu deleted
-60
@@ -1,60 +0,0 @@
1 -// SPDX-License-Identifier: Apache-2.0
2 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -// Adapted from SGLang:
4 -// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu
5 -
6 -#include <torch/all.h>
7 -
8 -#include "cutlass_mxfp8_grouped_mm_launcher.cuh"
9 -
10 -void cutlass_mxfp8_grouped_mm(const torch::Tensor& a, const torch::Tensor& b,
11 - const torch::Tensor& sfa,
12 - const torch::Tensor& sfb, torch::Tensor& d,
13 - const torch::Tensor& problem_sizes,
14 - const torch::Tensor& expert_offsets,
15 - const torch::Tensor& blockscale_offsets) {
16 -#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
17 - TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
18 - TORCH_CHECK(problem_sizes.size(1) == 3,
19 - "problem_sizes must have shape (num_experts, 3)");
20 - TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0),
21 - "Number of experts in problem_sizes must match expert_offsets");
22 - TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
23 - "problem_sizes must be int32");
24 - TORCH_CHECK(expert_offsets.dtype() == torch::kInt32,
25 - "expert_offsets must be int32");
26 - TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32,
27 - "blockscale_offsets must be int32");
28 - TORCH_CHECK(a.dim() == 2, "a must be a 2D tensor of shape (num_tokens, k)");
29 - TORCH_CHECK(b.dim() == 3,
30 - "b must be a 3D tensor of shape (num_experts, k, n)");
31 - TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0,
32 - "k should align 128");
33 - TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128");
34 - TORCH_CHECK(a.strides()[1] == 1, "a must be row major");
35 - TORCH_CHECK(b.strides()[1] == 1, "b must be column major");
36 -
37 - auto stream = at::cuda::getCurrentCUDAStream();
38 - if (d.dtype() == torch::kBFloat16) {
39 - expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
40 - cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
41 - blockscale_offsets, stream);
42 - } else if (d.dtype() == torch::kFloat16) {
43 - expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
44 - cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
45 - blockscale_offsets, stream);
46 - } else {
47 - TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
48 - }
49 -#else
50 - TORCH_CHECK(false,
51 - "No implemented cutlass_mxfp8_grouped_mm for "
52 - "current device");
53 -#endif
54 -}
55 -
56 -#include "core/registration.h"
57 -
58 -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
59 - m.impl("cutlass_mxfp8_grouped_mm", cutlass_mxfp8_grouped_mm);
60 -}
\ No newline at end of file
csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu deleted
-60
@@ -1,60 +0,0 @@
1 -// SPDX-License-Identifier: Apache-2.0
2 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -// Adapted from SGLang:
4 -// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu
5 -
6 -#include <torch/all.h>
7 -
8 -#include "mxfp8_experts_quant.cuh"
9 -
10 -void mxfp8_experts_quant(const torch::Tensor& input,
11 - const torch::Tensor& problem_sizes,
12 - const torch::Tensor& expert_offsets,
13 - const torch::Tensor& blockscale_offsets,
14 - torch::Tensor& quant_output,
15 - torch::Tensor& scale_factor) {
16 -#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
17 - TORCH_CHECK(input.dim() == 2, "input must be 2D tensor");
18 - TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128");
19 - TORCH_CHECK(input.strides()[1] == 1, "input must be row major");
20 - TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
21 - TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
22 - "problem_sizes must be int32");
23 - TORCH_CHECK(expert_offsets.dtype() == torch::kInt32,
24 - "expert_offsets must be int32");
25 - TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32,
26 - "blockscale_offsets must be int32");
27 -
28 - auto groups = problem_sizes.size(0);
29 - TORCH_CHECK(
30 - expert_offsets.dim() == 1 && expert_offsets.size(0) == groups,
31 - "expert_offsets must be 1D and have size equal to the number of groups");
32 - TORCH_CHECK(
33 - blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups,
34 - "blockscale_offsets must be 1D and have size equal to the number of "
35 - "groups");
36 -
37 - auto stream = at::cuda::getCurrentCUDAStream();
38 - if (input.dtype() == torch::kBFloat16) {
39 - expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>(
40 - input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
41 - scale_factor);
42 - } else if (input.dtype() == torch::kFloat16) {
43 - expert_specialization::launch_mxfp8_experts_quant<__half>(
44 - input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
45 - scale_factor);
46 - } else {
47 - TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
48 - }
49 -#else
50 - TORCH_CHECK(false,
51 - "No implemented mxfp8_experts_quant for "
52 - "current device");
53 -#endif
54 -}
55 -
56 -#include "core/registration.h"
57 -
58 -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
59 - m.impl("mxfp8_experts_quant", mxfp8_experts_quant);
60 -}
\ No newline at end of file
csrc/ops.h
-36
@@ -40,12 +40,6 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
40 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
41 torch::Tensor& weight, double epsilon);
42
43 -torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
44 - torch::Tensor const& q_in, torch::Tensor const& kv, torch::Tensor& k_cache,
45 - torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
46 - torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
47 - int64_t cache_block_size);
48 -
43 void silu_and_mul_per_block_quant(torch::Tensor& out,
44 torch::Tensor const& input,
45 torch::Tensor& scales, int64_t group_size,
@@ -107,24 +101,6 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
101 int64_t activation_kind);
102
103 using fptr_t = int64_t;
110 -fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
111 - torch::Tensor& rank_data, int64_t rank,
112 - bool fully_connected);
113 -void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
114 - fptr_t reg_buffer, int64_t reg_buffer_sz_bytes);
115 -void dispose(fptr_t _fa);
116 -int64_t meta_size();
117 -void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
118 -std::tuple<std::vector<int64_t>, std::vector<int64_t>>
119 -get_graph_buffer_ipc_meta(fptr_t _fa);
120 -void register_graph_buffers(fptr_t _fa,
121 - const std::vector<std::vector<int64_t>>& handles,
122 - const std::vector<std::vector<int64_t>>& offsets);
123 -std::tuple<int64_t, torch::Tensor> allocate_shared_buffer_and_handle(
124 - int64_t size);
125 -int64_t open_mem_handle(torch::Tensor& mem_handle);
126 -void free_shared_buffer(int64_t buffer);
127 -
104 #ifdef USE_ROCM
105 fptr_t init_custom_qr(int64_t rank, int64_t world_size,
106 std::optional<int64_t> qr_max_size = std::nullopt);
@@ -135,15 +111,3 @@ void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
111 int64_t quant_level, bool cast_bf2half = false);
112 int64_t qr_max_size();
113 #endif
138 -
139 -#ifndef USE_ROCM
140 -torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
141 - torch::Tensor const& norm_weight,
142 - torch::Tensor workspace, int64_t const rank,
143 - int64_t const nranks, double const eps);
144 -std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
145 - torch::Tensor qkv, torch::Tensor const& norm_weight_q,
146 - torch::Tensor const& norm_weight_k, torch::Tensor workspace,
147 - int64_t const q_size, int64_t const kv_size, int64_t const rank,
148 - int64_t const nranks, double const eps);
149 -#endif
csrc/torch_bindings.cpp
+17 -75
@@ -55,14 +55,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
55
56 // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
57 // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
58 - // kernel launch.
59 - ops.def(
60 - "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
61 - "Tensor q_in, Tensor kv, Tensor! k_cache, "
62 - "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
63 - "int q_head_padded, float eps, int cache_block_size) -> Tensor");
64 - ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
65 - &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
58 + // kernel launch. Registered in _C_stable_libtorch.
59
60 // Quantization ops
61 #ifndef USE_ROCM
@@ -163,74 +156,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
156 // conditionally compiled so impl registration is in source file
157
158 #endif
166 -
167 -#ifndef USE_ROCM
168 - ops.def(
169 - "minimax_allreduce_rms("
170 - "Tensor input,"
171 - "Tensor norm_weight,"
172 - "Tensor workspace,"
173 - "int rank,"
174 - "int nranks,"
175 - "float eps) -> Tensor");
176 - ops.impl("minimax_allreduce_rms", torch::kCUDA, &minimax_allreduce_rms);
177 - ops.def(
178 - "minimax_allreduce_rms_qk("
179 - "Tensor qkv,"
180 - "Tensor norm_weight_q,"
181 - "Tensor norm_weight_k,"
182 - "Tensor workspace,"
183 - "int q_size,"
184 - "int kv_size,"
185 - "int rank,"
186 - "int nranks,"
187 - "float eps) -> (Tensor, Tensor)");
188 - ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk);
189 -
190 - // conditionally compiled so impl in source file
191 -#endif
159 }
160
194 -TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
195 - // Cuda utils
196 -
197 - // Gets the specified device attribute.
198 - cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
199 - cuda_utils.impl("get_device_attribute", &get_device_attribute);
200 -
201 - // Gets the maximum shared memory per block device attribute.
202 - cuda_utils.def(
203 - "get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
204 - cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
205 - &get_max_shared_memory_per_block_device_attribute);
206 -}
207 -
208 -TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
209 - // Custom all-reduce kernels
210 - custom_ar.def(
211 - "init_custom_ar(int[] ipc_tensors, Tensor rank_data, "
212 - "int rank, bool fully_connected) -> int");
213 - custom_ar.impl("init_custom_ar", torch::kCUDA, &init_custom_ar);
214 - custom_ar.def(
215 - "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, "
216 - "int reg_buffer_sz_bytes) -> ()");
217 - custom_ar.impl("all_reduce", torch::kCUDA, &all_reduce);
218 -
219 - custom_ar.def("dispose", &dispose);
220 - custom_ar.def("meta_size", &meta_size);
221 -
222 - custom_ar.def("register_buffer", &register_buffer);
223 - custom_ar.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta);
224 - custom_ar.def("register_graph_buffers", &register_graph_buffers);
225 -
226 - custom_ar.def("allocate_shared_buffer_and_handle",
227 - &allocate_shared_buffer_and_handle);
228 - custom_ar.def("open_mem_handle(Tensor mem_handle) -> int", &open_mem_handle);
229 - custom_ar.impl("open_mem_handle", torch::kCPU, &open_mem_handle);
230 -
231 - custom_ar.def("free_shared_buffer", &free_shared_buffer);
161 #ifdef USE_ROCM
233 - // Quick Reduce all-reduce kernels
162 +TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
163 + // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C).
164 custom_ar.def(
165 "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool "
166 "cast_bf2half) -> ()");
@@ -238,15 +168,27 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
168
169 custom_ar.def("init_custom_qr", &init_custom_qr);
170 custom_ar.def("qr_destroy", &qr_destroy);
241 -
171 custom_ar.def("qr_get_handle", &qr_get_handle);
172
173 custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()");
174 custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles);
175
247 - // Max input size in bytes
176 custom_ar.def("qr_max_size", &qr_max_size);
177 +}
178 #endif
179 +
180 +TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
181 + // Cuda utils
182 +
183 + // Gets the specified device attribute.
184 + cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
185 + cuda_utils.impl("get_device_attribute", &get_device_attribute);
186 +
187 + // Gets the maximum shared memory per block device attribute.
188 + cuda_utils.def(
189 + "get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
190 + cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
191 + &get_max_shared_memory_per_block_device_attribute);
192 }
193
194 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
csrc/type_convert.cuh
+2 -2
@@ -50,7 +50,7 @@ struct _typeConvert<float> {
50 #if defined(USE_ROCM) || (defined(CUDA_VERSION) && (CUDA_VERSION >= 12000))
51 // CUDA < 12.0 runs into issues with packed type conversion
52 template <>
53 -struct _typeConvert<c10::Half> {
53 +struct _typeConvert<torch::headeronly::Half> {
54 static constexpr bool exists = true;
55 using hip_type = __half;
56 using packed_hip_type = __half2;
@@ -73,7 +73,7 @@ struct _typeConvert<c10::Half> {
73 // CUDA_ARCH < 800 does not have BF16 support
74 // ROCm 7.0+ supports bfloat16
75 template <>
76 -struct _typeConvert<c10::BFloat16> {
76 +struct _typeConvert<torch::headeronly::BFloat16> {
77 static constexpr bool exists = true;
78 using hip_type = __nv_bfloat16;
79 using packed_hip_type = __nv_bfloat162;