A high-throughput and memory-efficient inference and serving engine for LLMs
[10c/n] Migrate MoE kernels to torch stable ABI (#44565)
Signed-off-by: Chris Leonard <chleonar@redhat.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
Chris Leonard committed
Jun 11, 2026 at 02:02 UTC
6e64c1bab1875a5c096860dff9b7250f7d484094
34 files changed
+1300
-1102
.gitignore
+1
-1
@@ -233,7 +233,7 @@ actionlint
233
shellcheck*/
234
235
# Ignore moe/marlin_moe gen code
236
-csrc/moe/marlin_moe_wna16/kernel_*
236
+csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_*
237
238
# Ignore ep_kernels_workspace folder
239
ep_kernels_workspace/
.pre-commit-config.yaml
+1
-1
@@ -21,7 +21,7 @@ repos:
21
rev: v21.1.2
22
hooks:
23
- id: clang-format
24
- exclude: 'csrc/(moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
24
+ exclude: 'csrc/(libtorch_stable/moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
25
types_or: [c++, cuda]
26
args: [--style=file, --verbose]
27
- repo: https://github.com/DavidAnson/markdownlint-cli2
CMakeLists.txt
+61
-20
@@ -1115,25 +1115,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1115
endif()
1116
1117
#
1118
-# _moe_C extension
1118
+# _moe_C_stable_libtorch extension
1119
#
1120
1121
set(VLLM_MOE_EXT_SRC
1122
- "csrc/moe/torch_bindings.cpp"
1123
- "csrc/moe/moe_align_sum_kernels.cu"
1124
- "csrc/moe/topk_softmax_kernels.cu"
1125
- "csrc/moe/topk_softplus_sqrt_kernels.cu")
1122
+ "csrc/libtorch_stable/moe/torch_bindings.cpp"
1123
+ "csrc/libtorch_stable/moe/moe_align_sum_kernels.cu"
1124
+ "csrc/libtorch_stable/moe/topk_softmax_kernels.cu"
1125
+ "csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu")
1126
1127
if(VLLM_GPU_LANG STREQUAL "CUDA")
1128
list(APPEND VLLM_MOE_EXT_SRC
1129
- "csrc/moe/moe_wna16.cu"
1130
- "csrc/moe/grouped_topk_kernels.cu")
1129
+ "csrc/libtorch_stable/moe/moe_wna16.cu"
1130
+ "csrc/libtorch_stable/moe/grouped_topk_kernels.cu")
1131
endif()
1132
1133
if(VLLM_GPU_LANG STREQUAL "CUDA")
1134
set(MOE_PERMUTE_SRC
1135
- "csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu"
1136
- "csrc/moe/moe_permute_unpermute_op.cu")
1135
+ "csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu"
1136
+ "csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu")
1137
1138
list(APPEND VLLM_MOE_EXT_SRC "${MOE_PERMUTE_SRC}")
1139
endif()
@@ -1144,7 +1144,7 @@ set_gencode_flags_for_srcs(
1144
1145
if(VLLM_GPU_LANG STREQUAL "CUDA")
1146
set(VLLM_MOE_WNA16_SRC
1147
- "csrc/moe/moe_wna16.cu")
1147
+ "csrc/libtorch_stable/moe/moe_wna16.cu")
1148
1149
set_gencode_flags_for_srcs(
1150
SRCS "${VLLM_MOE_WNA16_SRC}"
@@ -1175,7 +1175,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1175
# preselected input type pairs and schedules.
1176
# Generate sources:
1177
set(MOE_MARLIN_GEN_SCRIPT
1178
- ${CMAKE_CURRENT_SOURCE_DIR}/csrc/moe/marlin_moe_wna16/generate_kernels.py)
1178
+ ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py)
1179
file(MD5 ${MOE_MARLIN_GEN_SCRIPT} MOE_MARLIN_GEN_SCRIPT_HASH)
1180
list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
1181
set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MOE_MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")
@@ -1210,7 +1210,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1210
endif()
1211
1212
if (MARLIN_MOE_ARCHS)
1213
- file(GLOB MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/sm80_kernel_*.cu")
1213
+ file(GLOB MARLIN_MOE_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_*.cu")
1214
set_gencode_flags_for_srcs(
1215
SRCS "${MARLIN_MOE_SRC}"
1216
CUDA_ARCHS "${MARLIN_MOE_ARCHS}")
@@ -1222,7 +1222,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1222
endif()
1223
1224
if (MARLIN_MOE_SM75_ARCHS)
1225
- file(GLOB MARLIN_MOE_SM75_SRC "csrc/moe/marlin_moe_wna16/sm75_kernel_*.cu")
1225
+ file(GLOB MARLIN_MOE_SM75_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm75_kernel_*.cu")
1226
set_gencode_flags_for_srcs(
1227
SRCS "${MARLIN_MOE_SM75_SRC}"
1228
CUDA_ARCHS "${MARLIN_MOE_SM75_ARCHS}")
@@ -1234,7 +1234,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1234
endif()
1235
1236
if (MARLIN_MOE_FP8_ARCHS)
1237
- file(GLOB MARLIN_MOE_FP8_SRC "csrc/moe/marlin_moe_wna16/sm89_kernel_*.cu")
1237
+ file(GLOB MARLIN_MOE_FP8_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm89_kernel_*.cu")
1238
set_gencode_flags_for_srcs(
1239
SRCS "${MARLIN_MOE_FP8_SRC}"
1240
CUDA_ARCHS "${MARLIN_MOE_FP8_ARCHS}")
@@ -1245,7 +1245,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1245
list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_FP8_SRC})
1246
endif()
1247
1248
- set(MARLIN_MOE_OTHER_SRC "csrc/moe/marlin_moe_wna16/ops.cu")
1248
+ set(MARLIN_MOE_OTHER_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu")
1249
set_gencode_flags_for_srcs(
1250
SRCS "${MARLIN_MOE_OTHER_SRC}"
1251
CUDA_ARCHS "${MARLIN_MOE_OTHER_ARCHS}")
@@ -1266,9 +1266,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1266
cuda_archs_sm90plus(SM90PLUS_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
1267
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SM90PLUS_ROUTER_GEMM_ARCHS)
1268
set(DSV3_ROUTER_GEMM_SRC
1269
- "csrc/moe/dsv3_router_gemm_entry.cu"
1270
- "csrc/moe/dsv3_router_gemm_float_out.cu"
1271
- "csrc/moe/dsv3_router_gemm_bf16_out.cu")
1269
+ "csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu"
1270
+ "csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu"
1271
+ "csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu")
1272
set_gencode_flags_for_srcs(
1273
SRCS "${DSV3_ROUTER_GEMM_SRC}"
1274
CUDA_ARCHS "${SM90PLUS_ROUTER_GEMM_ARCHS}")
@@ -1281,9 +1281,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1281
endif()
1282
endif()
1283
1284
-message(STATUS "Enabling moe extension.")
1284
+message(STATUS "Enabling MoE C_stable extension.")
1285
define_extension_target(
1286
- _moe_C
1286
+ _moe_C_stable_libtorch
1287
DESTINATION vllm
1288
LANGUAGE ${VLLM_GPU_LANG}
1289
SOURCES ${VLLM_MOE_EXT_SRC}
@@ -1294,6 +1294,47 @@ define_extension_target(
1294
USE_SABI 3
1295
WITH_SOABI)
1296
1297
+# Needed to use cuda/hip APIs from C-shim
1298
+if(VLLM_GPU_LANG STREQUAL "CUDA")
1299
+ # Set TORCH_TARGET_VERSION for stable ABI compatibility.
1300
+ # This ensures we only use C-shim APIs available in PyTorch 2.11.
1301
+ # _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1302
+ # which is currently set to 2.11.
1303
+ target_compile_definitions(_moe_C_stable_libtorch PRIVATE
1304
+ TORCH_TARGET_VERSION=0x020B000000000000ULL)
1305
+ target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_CUDA)
1306
+ # Needed by CUTLASS kernels
1307
+ target_compile_definitions(_moe_C_stable_libtorch PRIVATE
1308
+ CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
1309
+elseif(VLLM_GPU_LANG STREQUAL "HIP")
1310
+ # Set TORCH_TARGET_VERSION for stable ABI compatibility.
1311
+ # This ensures we only use C-shim APIs available in PyTorch 2.10.
1312
+ # _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1313
+ # which is currently set to 2.10.
1314
+ target_compile_definitions(_moe_C_stable_libtorch PRIVATE
1315
+ TORCH_TARGET_VERSION=0x020A000000000000ULL)
1316
+ target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_ROCM)
1317
+endif()
1318
+
1319
+# On ROCm, _moe_C_stable_libtorch calls raw HIP APIs (e.g. hipGetDevice in
1320
+# get_device_prop()) which must resolve to the same libamdhip64.so that
1321
+# PyTorch uses. When PyTorch bundles its own copy (pip/conda wheels),
1322
+# the raw HIP calls would otherwise resolve to the system ROCm copy,
1323
+# initializing a second HIP runtime that corrupts device state (wrong
1324
+# device on DeviceGuard, core dumps on multi-GPU tests).
1325
+#
1326
+# If PyTorch doesn't bundle libamdhip64 (built from source against system
1327
+# ROCm), there is only one copy in the process and no action is needed —
1328
+# the HIP compiler already links the system libamdhip64 automatically.
1329
+if(VLLM_GPU_LANG STREQUAL "HIP")
1330
+ find_library(_MOE_STABLE_TORCH_AMDHIP64 amdhip64
1331
+ PATHS "${TORCH_INSTALL_PREFIX}/lib" NO_DEFAULT_PATH)
1332
+ if(_MOE_STABLE_TORCH_AMDHIP64)
1333
+ message(STATUS "Found PyTorch-bundled libamdhip64 for _moe_C_stable_libtorch at ${_MOE_STABLE_TORCH_AMDHIP64}")
1334
+ target_link_libraries(_moe_C_stable_libtorch PRIVATE ${_MOE_STABLE_TORCH_AMDHIP64})
1335
+ endif()
1336
+endif()
1337
+
1338
if(VLLM_GPU_LANG STREQUAL "HIP")
1339
#
1340
# _rocm_C extension
csrc/libtorch_stable/dispatch_utils.h
+22
@@ -30,6 +30,28 @@
30
THO_DISPATCH_SWITCH(TYPE, NAME, \
31
VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
32
33
+#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(...) \
34
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__) \
35
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) \
36
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Short, __VA_ARGS__) \
37
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Int, __VA_ARGS__) \
38
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Long, __VA_ARGS__)
39
+
40
+#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \
41
+ VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__) \
42
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt16, __VA_ARGS__) \
43
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt32, __VA_ARGS__) \
44
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt64, __VA_ARGS__)
45
+
46
+#define VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \
47
+ THO_DISPATCH_SWITCH(TYPE, NAME, \
48
+ VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__))
49
+
50
+#define VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \
51
+ THO_DISPATCH_SWITCH( \
52
+ TYPE, NAME, \
53
+ VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__))
54
+
55
// FP8 type dispatch - ROCm uses FNUZ format, CUDA uses OCP format
56
#ifdef USE_ROCM
57
#define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \
csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu
renamed
+1
-4
@@ -18,14 +18,11 @@
18
* limitations under the License.
19
*/
20
21
-#include <ATen/ATen.h>
22
-#include <ATen/cuda/CUDAContext.h>
21
+#include <cstdint>
22
23
#include <cuda_bf16.h>
24
#include <cuda_runtime.h>
25
27
-#include "dsv3_router_gemm_utils.h"
28
-
26
// Custom FMA implementation using PTX assembly instructions
27
__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b,
28
float2 const& c) {
csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu
renamed
+50
-32
@@ -18,15 +18,25 @@
18
* limitations under the License.
19
*/
20
21
-#include <ATen/ATen.h>
22
-#include <ATen/cuda/CUDAContext.h>
23
-#include <torch/all.h>
21
+#include <torch/csrc/stable/library.h>
22
+#include <torch/csrc/stable/tensor.h>
23
+#include <torch/headeronly/core/ScalarType.h>
24
+
25
+#include "libtorch_stable/torch_utils.h"
26
27
#include <cuda_bf16.h>
28
#include <cuda_runtime.h>
29
28
-#include "core/registration.h"
29
-#include "dsv3_router_gemm_utils.h"
30
+#include <stdexcept>
31
+
32
+namespace {
33
+
34
+inline int getSMVersion() {
35
+ auto* props = get_device_prop();
36
+ return props->major * 10 + props->minor;
37
+}
38
+
39
+} // namespace
40
41
static constexpr int DEFAULT_NUM_EXPERTS = 256;
42
static constexpr int KIMI_K2_NUM_EXPERTS = 384;
@@ -98,40 +108,48 @@ struct LoopUnroller<kEnd, kEnd, kNumExperts, kHiddenDim> {
108
}
109
};
110
101
-void dsv3_router_gemm(at::Tensor& output, // [num_tokens, num_experts]
102
- const at::Tensor& mat_a, // [num_tokens, hidden_dim]
103
- const at::Tensor& mat_b // [num_experts, hidden_dim]
111
+void dsv3_router_gemm(
112
+ torch::stable::Tensor& output, // [num_tokens, num_experts]
113
+ torch::stable::Tensor const& mat_a, // [num_tokens, hidden_dim]
114
+ torch::stable::Tensor const& mat_b // [num_experts, hidden_dim]
115
) {
105
- TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
116
+ STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
117
118
const int num_tokens = mat_a.size(0);
119
const int num_experts = mat_b.size(0);
120
const int hidden_dim = mat_a.size(1);
121
111
- TORCH_CHECK(mat_a.size(1) == mat_b.size(1),
112
- "mat_a and mat_b must have the same hidden_dim");
113
- TORCH_CHECK(hidden_dim == DEFAULT_HIDDEN_DIM,
114
- "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
115
- ", but got hidden_dim=", hidden_dim);
116
- TORCH_CHECK(
122
+ STD_TORCH_CHECK(mat_a.size(1) == mat_b.size(1),
123
+ "mat_a and mat_b must have the same hidden_dim");
124
+ STD_TORCH_CHECK(hidden_dim == DEFAULT_HIDDEN_DIM,
125
+ "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
126
+ ", but got hidden_dim=", hidden_dim);
127
+ STD_TORCH_CHECK(
128
num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS,
129
"Expected num_experts=", DEFAULT_NUM_EXPERTS,
130
" or num_experts=", KIMI_K2_NUM_EXPERTS,
131
", but got num_experts=", num_experts);
121
- TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
122
- "currently num_tokens must be less than or equal to 16 for "
123
- "router_gemm");
124
- TORCH_CHECK(mat_a.dtype() == at::kBFloat16, "mat_a must be bf16");
125
- TORCH_CHECK(mat_b.dtype() == at::kBFloat16, "mat_b must be bf16");
126
- TORCH_CHECK(output.dtype() == at::kFloat || output.dtype() == at::kBFloat16,
127
- "output must be float32 or bf16");
128
-
129
- auto const sm = getSMVersion();
130
- TORCH_CHECK(sm >= 90 && sm <= 103, "required SM_103 >= CUDA ARCH >= SM_90");
131
-
132
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
133
-
134
- if (output.dtype() == at::kFloat) {
132
+ STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
133
+ "currently num_tokens must be less than or equal to 16 for "
134
+ "router_gemm");
135
+ STD_TORCH_CHECK(
136
+ mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16,
137
+ "mat_a must be bf16");
138
+ STD_TORCH_CHECK(
139
+ mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16,
140
+ "mat_b must be bf16");
141
+ STD_TORCH_CHECK(
142
+ output.scalar_type() == torch::headeronly::ScalarType::Float ||
143
+ output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
144
+ "output must be float32 or bf16");
145
+
146
+ const int sm = getSMVersion();
147
+ STD_TORCH_CHECK(sm >= 90 && sm <= 103,
148
+ "required SM_103 >= CUDA ARCH >= SM_90");
149
+
150
+ const cudaStream_t stream = get_current_cuda_stream(mat_a.get_device_index());
151
+
152
+ if (output.scalar_type() == torch::headeronly::ScalarType::Float) {
153
if (num_experts == DEFAULT_NUM_EXPERTS) {
154
LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
155
unroll_float_output(
@@ -145,7 +163,7 @@ void dsv3_router_gemm(at::Tensor& output, // [num_tokens, num_experts]
163
reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
164
reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
165
}
148
- } else if (output.dtype() == at::kBFloat16) {
166
+ } else if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
167
if (num_experts == DEFAULT_NUM_EXPERTS) {
168
LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
169
unroll_bf16_output(
@@ -164,6 +182,6 @@ void dsv3_router_gemm(at::Tensor& output, // [num_tokens, num_experts]
182
}
183
}
184
167
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
168
- m.impl("dsv3_router_gemm", &dsv3_router_gemm);
185
+STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
186
+ m.impl("dsv3_router_gemm", TORCH_BOX(&dsv3_router_gemm));
187
}
csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu
renamed
+1
-4
@@ -18,14 +18,11 @@
18
* limitations under the License.
19
*/
20
21
-#include <ATen/ATen.h>
22
-#include <ATen/cuda/CUDAContext.h>
21
+#include <cstdint>
22
23
#include <cuda_bf16.h>
24
#include <cuda_runtime.h>
25
27
-#include "dsv3_router_gemm_utils.h"
28
-
26
// Custom FMA implementation using PTX assembly instructions
27
__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b,
28
float2 const& c) {
csrc/libtorch_stable/moe/grouped_topk_kernels.cu
renamed
+48
-40
@@ -18,9 +18,14 @@
18
* limitations under the License.
19
*/
20
#include "moeTopKFuncs.cuh"
21
-#include <c10/cuda/CUDAStream.h>
22
-#include <torch/all.h>
21
+
22
+#include <torch/csrc/stable/tensor.h>
23
+#include <torch/headeronly/core/ScalarType.h>
24
+
25
+#include "libtorch_stable/torch_utils.h"
26
+
27
#include <cmath>
28
+#include <tuple>
29
#include <cuda_fp16.h>
30
#include <cuda_bf16.h>
31
#include <cuda/std/limits>
@@ -1001,38 +1006,40 @@ INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_NONE);
1006
} // end namespace moe
1007
} // namespace vllm
1008
1004
-std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
1005
- torch::Tensor const& scores, int64_t n_group, int64_t topk_group,
1009
+std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
1010
+ torch::stable::Tensor const& scores, int64_t n_group, int64_t topk_group,
1011
int64_t topk, bool renormalize, double routed_scaling_factor,
1007
- torch::Tensor const& bias, int64_t scoring_func = 0) {
1008
- auto data_type = scores.scalar_type();
1009
- auto bias_type = bias.scalar_type();
1010
- auto input_size = scores.sizes();
1011
- int64_t num_tokens = input_size[0];
1012
- int64_t num_experts = input_size[1];
1013
- TORCH_CHECK(input_size.size() == 2, "scores must be a 2D Tensor");
1014
- TORCH_CHECK(n_group > 0, "n_group must be positive");
1015
- TORCH_CHECK(topk > 0, "topk must be positive");
1016
- TORCH_CHECK(topk_group > 0, "topk_group must be positive");
1017
- TORCH_CHECK(topk_group <= n_group, "topk_group must be <= n_group");
1018
- TORCH_CHECK(num_experts % n_group == 0,
1019
- "num_experts should be divisible by n_group");
1020
- TORCH_CHECK(n_group <= 32,
1021
- "n_group should be smaller than or equal to 32 for now");
1022
- TORCH_CHECK(topk <= 32, "topk should be smaller than or equal to 32 for now");
1023
- TORCH_CHECK(topk <= topk_group * (num_experts / n_group),
1024
- "topk must be <= topk_group * (num_experts / n_group)");
1025
- TORCH_CHECK(scoring_func == vllm::moe::SCORING_NONE ||
1026
- scoring_func == vllm::moe::SCORING_SIGMOID,
1027
- "scoring_func must be SCORING_NONE (0) or SCORING_SIGMOID (1)");
1012
+ torch::stable::Tensor const& bias, int64_t scoring_func = 0) {
1013
+ const auto data_type = scores.scalar_type();
1014
+ const auto bias_type = bias.scalar_type();
1015
+ STD_TORCH_CHECK(scores.dim() == 2, "scores must be a 2D Tensor");
1016
+ const int64_t num_tokens = scores.size(0);
1017
+ const int64_t num_experts = scores.size(1);
1018
+ STD_TORCH_CHECK(n_group > 0, "n_group must be positive");
1019
+ STD_TORCH_CHECK(topk > 0, "topk must be positive");
1020
+ STD_TORCH_CHECK(topk_group > 0, "topk_group must be positive");
1021
+ STD_TORCH_CHECK(topk_group <= n_group, "topk_group must be <= n_group");
1022
+ STD_TORCH_CHECK(num_experts % n_group == 0,
1023
+ "num_experts should be divisible by n_group");
1024
+ STD_TORCH_CHECK(n_group <= 32,
1025
+ "n_group should be smaller than or equal to 32 for now");
1026
+ STD_TORCH_CHECK(topk <= 32,
1027
+ "topk should be smaller than or equal to 32 for now");
1028
+ STD_TORCH_CHECK(topk <= topk_group * (num_experts / n_group),
1029
+ "topk must be <= topk_group * (num_experts / n_group)");
1030
+ STD_TORCH_CHECK(
1031
+ scoring_func == vllm::moe::SCORING_NONE ||
1032
+ scoring_func == vllm::moe::SCORING_SIGMOID,
1033
+ "scoring_func must be SCORING_NONE (0) or SCORING_SIGMOID (1)");
1034
1035
// Always output float32 for topk_values (eliminates Python-side conversion)
1030
- torch::Tensor topk_values = torch::empty(
1031
- {num_tokens, topk}, torch::dtype(torch::kFloat32).device(torch::kCUDA));
1032
- torch::Tensor topk_indices = torch::empty(
1033
- {num_tokens, topk}, torch::dtype(torch::kInt32).device(torch::kCUDA));
1036
+ auto topk_values = torch::stable::new_empty(
1037
+ scores, {num_tokens, topk}, torch::headeronly::ScalarType::Float);
1038
+ auto topk_indices = torch::stable::new_empty(
1039
+ scores, {num_tokens, topk}, torch::headeronly::ScalarType::Int);
1040
1035
- auto stream = c10::cuda::getCurrentCUDAStream(scores.get_device());
1041
+ const cudaStream_t stream =
1042
+ get_current_cuda_stream(scores.get_device_index());
1043
auto const sf = static_cast<vllm::moe::ScoringFunc>(scoring_func);
1044
1045
#define LAUNCH_KERNEL_SF(T, BiasT, IdxT) \
@@ -1057,7 +1064,7 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
1064
routed_scaling_factor, false, stream); \
1065
break; \
1066
default: \
1060
- throw std::invalid_argument("Unsupported scoring_func"); \
1067
+ STD_TORCH_CHECK(false, "Unsupported scoring_func"); \
1068
break; \
1069
} \
1070
} while (0)
@@ -1065,17 +1072,18 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
1072
#define LAUNCH_KERNEL(T, IdxT) \
1073
do { \
1074
switch (bias_type) { \
1068
- case torch::kFloat16: \
1075
+ case torch::headeronly::ScalarType::Half: \
1076
LAUNCH_KERNEL_SF(T, half, IdxT); \
1077
break; \
1071
- case torch::kFloat32: \
1078
+ case torch::headeronly::ScalarType::Float: \
1079
LAUNCH_KERNEL_SF(T, float, IdxT); \
1080
break; \
1074
- case torch::kBFloat16: \
1081
+ case torch::headeronly::ScalarType::BFloat16: \
1082
LAUNCH_KERNEL_SF(T, __nv_bfloat16, IdxT); \
1083
break; \
1084
default: \
1078
- throw std::invalid_argument( \
1085
+ STD_TORCH_CHECK( \
1086
+ false, \
1087
"Invalid bias dtype, only supports float16, float32, and " \
1088
"bfloat16"); \
1089
break; \
@@ -1083,22 +1091,22 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
1091
} while (0)
1092
1093
switch (data_type) {
1086
- case torch::kFloat16:
1094
+ case torch::headeronly::ScalarType::Half:
1095
// Handle Float16
1096
LAUNCH_KERNEL(half, int32_t);
1097
break;
1090
- case torch::kFloat32:
1098
+ case torch::headeronly::ScalarType::Float:
1099
// Handle Float32
1100
LAUNCH_KERNEL(float, int32_t);
1101
break;
1094
- case torch::kBFloat16:
1102
+ case torch::headeronly::ScalarType::BFloat16:
1103
// Handle BFloat16
1104
LAUNCH_KERNEL(__nv_bfloat16, int32_t);
1105
break;
1106
default:
1107
// Handle other data types
1100
- throw std::invalid_argument(
1101
- "Invalid dtype, only supports float16, float32, and bfloat16");
1108
+ STD_TORCH_CHECK(
1109
+ false, "Invalid dtype, only supports float16, float32, and bfloat16");
1110
break;
1111
}
1112
#undef LAUNCH_KERNEL
csrc/libtorch_stable/moe/marlin_moe_wna16/.gitignore
renamed
csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py
renamed
+1
-1
@@ -302,7 +302,7 @@ def generate_new_kernels():
302
if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT:
303
kernel_selector_str += (
304
"else if (a_type == vllm::kFE4M3fn)\n"
305
- " TORCH_CHECK(false, "
305
+ " STD_TORCH_CHECK(false, "
306
'"marlin kernel with fp8 activation is not built.");'
307
)
308
csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h
renamed
csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h
renamed
csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu
renamed
+232
-208
@@ -24,7 +24,15 @@
24
#endif
25
26
#include "kernel.h"
27
-#include "core/registration.h"
27
+
28
+#include <torch/csrc/stable/accelerator.h>
29
+#include <torch/csrc/stable/library.h>
30
+#include <torch/csrc/stable/ops.h>
31
+#include <torch/csrc/stable/tensor.h>
32
+#include <torch/headeronly/core/ScalarType.h>
33
+#include <torch/headeronly/util/Exception.h>
34
+
35
+#include "libtorch_stable/torch_utils.h"
36
37
#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
38
static_assert(std::is_same<scalar_t, half>::value || \
@@ -350,18 +358,18 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
358
bool m_block_size_8 = moe_block_size == 8;
359
bool is_a_8bit = a_type.size_bits() == 8;
360
353
- TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
354
- ", ", prob_n, ", ", prob_k, "]");
361
+ STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [",
362
+ prob_m, ", ", prob_n, ", ", prob_k, "]");
363
364
int group_blocks = 0;
365
if (has_act_order) {
366
if (is_k_full) {
359
- TORCH_CHECK(group_size != -1);
367
+ STD_TORCH_CHECK(group_size != -1);
368
group_blocks = group_size / 16;
361
- TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
362
- " is not divisible by group_blocks = ", group_blocks);
369
+ STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
370
+ " is not divisible by group_blocks = ", group_blocks);
371
} else {
364
- TORCH_CHECK(group_size == 0);
372
+ STD_TORCH_CHECK(group_size == 0);
373
group_blocks = 0;
374
}
375
} else {
@@ -369,8 +377,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
377
group_blocks = -1;
378
} else {
379
group_blocks = group_size / 16;
372
- TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
373
- " is not divisible by group_blocks = ", group_blocks);
380
+ STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
381
+ " is not divisible by group_blocks = ", group_blocks);
382
}
383
}
384
@@ -407,7 +415,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
415
else if (moe_block_size == 64)
416
kernel = permute_cols_kernel<64>;
417
else
410
- TORCH_CHECK(false, "unsupported moe_block_size ", moe_block_size);
418
+ STD_TORCH_CHECK(false, "unsupported moe_block_size ", moe_block_size);
419
420
// avoid ">>>" being formatted to "> > >"
421
// clang-format off
@@ -428,25 +436,25 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
436
int max_shared_mem = 0;
437
cudaDeviceGetAttribute(&max_shared_mem,
438
cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
431
- TORCH_CHECK(max_shared_mem > 0);
439
+ STD_TORCH_CHECK(max_shared_mem > 0);
440
441
int major_capability, minor_capability;
442
cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor,
443
dev);
444
cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor,
445
dev);
438
- TORCH_CHECK(major_capability * 10 + minor_capability >= 75,
439
- "marlin kernel only support Turing or newer GPUs.");
446
+ STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 75,
447
+ "marlin kernel only support Turing or newer GPUs.");
448
int stages = 4;
449
if (major_capability == 7 && minor_capability == 5) {
450
stages = 2;
443
- TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8,
444
- "Turing only support FP16 or INT8 activation.");
451
+ STD_TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8,
452
+ "Turing only support FP16 or INT8 activation.");
453
}
454
if (a_type == vllm::kFE4M3fn) {
447
- TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
448
- "FP8 only support Ada Lovelace or newer GPUs.");
449
- TORCH_CHECK(
455
+ STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
456
+ "FP8 only support Ada Lovelace or newer GPUs.");
457
+ STD_TORCH_CHECK(
458
major_capability * 10 + minor_capability == 89 ||
459
major_capability == 12,
460
"Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than "
@@ -460,10 +468,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
468
thread_tfg = thread_config_t{thread_k, thread_n, thread_k * thread_n / 64};
469
if (blocks_per_sm == -1) blocks_per_sm = 1;
470
exec_cfg = exec_config_t{blocks_per_sm, thread_tfg};
463
- TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
464
- " is not divisible by thread_n = ", thread_n);
465
- TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
466
- " is not divisible by thread_k = ", thread_k);
471
+ STD_TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
472
+ " is not divisible by thread_n = ", thread_n);
473
+ STD_TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
474
+ " is not divisible by thread_k = ", thread_k);
475
} else {
476
// Auto config
477
exec_cfg = determine_exec_config(
@@ -484,19 +492,19 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
492
int thread_k_blocks = thread_k / 16;
493
int thread_n_blocks = thread_n / 16;
494
487
- TORCH_CHECK(is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks,
488
- prob_m, prob_n, prob_k, num_bits, group_size,
489
- has_act_order, is_k_full, has_zp, is_zp_float,
490
- is_a_8bit, stages, max_shared_mem),
491
- "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
492
- ", thread_k = ", thread_tfg.thread_k,
493
- ", thread_n = ", thread_tfg.thread_n,
494
- ", num_threads = ", thread_tfg.num_threads, " for MKN = [",
495
- prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
496
- ", group_size = ", group_size,
497
- ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full,
498
- ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float,
499
- ", max_shared_mem = ", max_shared_mem);
495
+ STD_TORCH_CHECK(
496
+ is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
497
+ prob_n, prob_k, num_bits, group_size, has_act_order,
498
+ is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
499
+ max_shared_mem),
500
+ "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
501
+ ", thread_k = ", thread_tfg.thread_k,
502
+ ", thread_n = ", thread_tfg.thread_n,
503
+ ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ",
504
+ prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
505
+ ", group_size = ", group_size, ", has_act_order = ", has_act_order,
506
+ ", is_k_full = ", is_k_full, ", has_zp = ", has_zp,
507
+ ", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem);
508
509
int sh_cache_size =
510
get_kernel_cache_size(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
@@ -509,13 +517,13 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
517
num_threads, is_zp_float, stages);
518
519
if (kernel == MarlinDefault) {
512
- TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
513
- ", ", prob_k, "]", ", has_act_order = ", has_act_order,
514
- ", num_groups = ", num_groups, ", group_size = ", group_size,
515
- ", thread_m_blocks = ", thread_m_blocks,
516
- ", thread_n_blocks = ", thread_n_blocks,
517
- ", thread_k_blocks = ", thread_k_blocks,
518
- ", num_bits = ", num_bits);
520
+ STD_TORCH_CHECK(
521
+ false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ",
522
+ prob_k, "]", ", has_act_order = ", has_act_order,
523
+ ", num_groups = ", num_groups, ", group_size = ", group_size,
524
+ ", thread_m_blocks = ", thread_m_blocks,
525
+ ", thread_n_blocks = ", thread_n_blocks,
526
+ ", thread_k_blocks = ", thread_k_blocks, ", num_bits = ", num_bits);
527
}
528
529
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
@@ -532,75 +540,81 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
540
541
} // namespace MARLIN_NAMESPACE_NAME
542
535
-torch::Tensor moe_wna16_marlin_gemm(
536
- torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
537
- torch::Tensor& b_q_weight,
538
- std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
539
- std::optional<torch::Tensor> const& a_scales_or_none,
540
- std::optional<torch::Tensor> const& global_scale_or_none,
541
- std::optional<torch::Tensor> const& b_zeros_or_none,
542
- std::optional<torch::Tensor> const& g_idx_or_none,
543
- std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
544
- torch::Tensor& sorted_token_ids, torch::Tensor& expert_ids,
545
- torch::Tensor& num_tokens_past_padded, torch::Tensor& topk_weights,
546
- int64_t moe_block_size, int64_t top_k, bool mul_topk_weights,
547
- vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n,
548
- int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
549
- bool is_zp_float, int64_t thread_k, int64_t thread_n,
543
+torch::stable::Tensor moe_wna16_marlin_gemm(
544
+ torch::stable::Tensor& a, std::optional<torch::stable::Tensor> c_or_none,
545
+ torch::stable::Tensor& b_q_weight,
546
+ std::optional<torch::stable::Tensor> const& b_bias_or_none,
547
+ torch::stable::Tensor& b_scales,
548
+ std::optional<torch::stable::Tensor> const& a_scales_or_none,
549
+ std::optional<torch::stable::Tensor> const& global_scale_or_none,
550
+ std::optional<torch::stable::Tensor> const& b_zeros_or_none,
551
+ std::optional<torch::stable::Tensor> const& g_idx_or_none,
552
+ std::optional<torch::stable::Tensor> const& perm_or_none,
553
+ torch::stable::Tensor& workspace, torch::stable::Tensor& sorted_token_ids,
554
+ torch::stable::Tensor& expert_ids,
555
+ torch::stable::Tensor& num_tokens_past_padded,
556
+ torch::stable::Tensor& topk_weights, int64_t moe_block_size, int64_t top_k,
557
+ bool mul_topk_weights, vllm::ScalarTypeId const& b_type_id, int64_t size_m,
558
+ int64_t size_n, int64_t size_k, bool is_k_full, bool use_atomic_add,
559
+ bool use_fp32_reduce, bool is_zp_float, int64_t thread_k, int64_t thread_n,
560
int64_t blocks_per_sm) {
561
vllm::ScalarTypeId a_type_id, c_type_id, s_type_id;
562
553
- auto c_dtype = a.dtype();
554
- if (a.scalar_type() == at::ScalarType::Half) {
563
+ auto c_dtype = a.scalar_type();
564
+ if (a.scalar_type() == torch::headeronly::ScalarType::Half) {
565
a_type_id = vllm::kFloat16.id();
566
c_type_id = vllm::kFloat16.id();
557
- } else if (a.scalar_type() == at::ScalarType::BFloat16) {
567
+ } else if (a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
568
a_type_id = vllm::kBFloat16.id();
569
c_type_id = vllm::kBFloat16.id();
570
} else {
561
- c_dtype = b_scales.dtype();
562
- if (b_scales.scalar_type() == at::ScalarType::Half) {
571
+ c_dtype = b_scales.scalar_type();
572
+ if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) {
573
c_type_id = vllm::kFloat16.id();
564
- } else if (b_scales.scalar_type() == at::ScalarType::BFloat16) {
574
+ } else if (b_scales.scalar_type() ==
575
+ torch::headeronly::ScalarType::BFloat16) {
576
c_type_id = vllm::kBFloat16.id();
577
} else {
578
c_type_id = vllm::kBFloat16.id();
579
569
- TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4");
570
- torch::Tensor c = c_or_none.value();
571
- c_dtype = c.dtype();
580
+ STD_TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4");
581
+ torch::stable::Tensor c = c_or_none.value();
582
+ c_dtype = c.scalar_type();
583
573
- if (c.scalar_type() == at::ScalarType::Half) {
584
+ if (c.scalar_type() == torch::headeronly::ScalarType::Half) {
585
c_type_id = vllm::kFloat16.id();
575
- } else if (c.scalar_type() == at::ScalarType::BFloat16) {
586
+ } else if (c.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
587
c_type_id = vllm::kBFloat16.id();
588
} else {
578
- TORCH_CHECK(false, "unsupported c dtype");
589
+ STD_TORCH_CHECK(false, "unsupported c dtype");
590
}
591
}
592
582
- if (a.scalar_type() == at::ScalarType::Float8_e4m3fn) {
593
+ if (a.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn) {
594
a_type_id = vllm::kFE4M3fn.id();
584
- } else if (a.scalar_type() == at::ScalarType::Char) {
595
+ } else if (a.scalar_type() == torch::headeronly::ScalarType::Char) {
596
a_type_id = vllm::kS8.id();
597
} else {
587
- TORCH_CHECK(false, "unsupported `a` scalar_type");
598
+ STD_TORCH_CHECK(false, "unsupported `a` scalar_type");
599
}
600
}
601
602
s_type_id = c_type_id;
603
if (b_type_id == vllm::kFE2M1f.id()) {
593
- if (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn) {
604
+ if (b_scales.scalar_type() ==
605
+ torch::headeronly::ScalarType::Float8_e4m3fn) {
606
s_type_id = vllm::kFE4M3fn.id();
595
- } else if (b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
607
+ } else if (b_scales.scalar_type() ==
608
+ torch::headeronly::ScalarType::Float8_e8m0fnu) {
609
s_type_id = vllm::kFE8M0fnu.id();
610
} else {
598
- TORCH_CHECK(false,
599
- "When b_type = float4_e2m1f, b_scale scalar type must be",
600
- "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
611
+ STD_TORCH_CHECK(
612
+ false, "When b_type = float4_e2m1f, b_scale scalar type must be",
613
+ "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
614
}
615
} else if (b_type_id == vllm::kFE4M3fn.id() &&
603
- b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
616
+ b_scales.scalar_type() ==
617
+ torch::headeronly::ScalarType::Float8_e8m0fnu) {
618
s_type_id = vllm::kFE8M0fnu.id();
619
}
620
@@ -613,58 +627,60 @@ torch::Tensor moe_wna16_marlin_gemm(
627
int num_experts = b_q_weight.size(0);
628
629
if (moe_block_size != 8) {
616
- TORCH_CHECK(moe_block_size % 16 == 0,
617
- "unsupported moe_block_size=", moe_block_size);
618
- TORCH_CHECK(moe_block_size >= 16 && moe_block_size <= 64,
619
- "unsupported moe_block_size=", moe_block_size);
630
+ STD_TORCH_CHECK(moe_block_size % 16 == 0,
631
+ "unsupported moe_block_size=", moe_block_size);
632
+ STD_TORCH_CHECK(moe_block_size >= 16 && moe_block_size <= 64,
633
+ "unsupported moe_block_size=", moe_block_size);
634
}
635
636
// Verify A
623
- TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0),
624
- ", size_m = ", size_m);
625
- TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1),
626
- ", size_k = ", size_k);
637
+ STD_TORCH_CHECK(a.size(0) == size_m,
638
+ "Shape mismatch: a.size(0) = ", a.size(0),
639
+ ", size_m = ", size_m);
640
+ STD_TORCH_CHECK(a.size(1) == size_k,
641
+ "Shape mismatch: a.size(1) = ", a.size(1),
642
+ ", size_k = ", size_k);
643
644
// Verify B
629
- TORCH_CHECK(
645
+ STD_TORCH_CHECK(
646
size_k % MARLIN_NAMESPACE_NAME::tile_size == 0, "size_k = ", size_k,
647
" is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
632
- TORCH_CHECK((size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(1),
633
- "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1),
634
- ", size_k = ", size_k,
635
- ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
636
- TORCH_CHECK(
648
+ STD_TORCH_CHECK(
649
+ (size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(1),
650
+ "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1),
651
+ ", size_k = ", size_k,
652
+ ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
653
+ STD_TORCH_CHECK(
654
b_q_weight.size(2) % MARLIN_NAMESPACE_NAME::tile_size == 0,
655
"b_q_weight.size(2) = ", b_q_weight.size(2),
656
" is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
657
int actual_size_n =
658
(b_q_weight.size(2) / MARLIN_NAMESPACE_NAME::tile_size) * pack_factor;
642
- TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n,
643
- ", actual_size_n = ", actual_size_n);
659
+ STD_TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n,
660
+ ", actual_size_n = ", actual_size_n);
661
662
// Verify device and strides
646
- TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
647
- TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
663
+ STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
664
+ STD_TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
665
649
- TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
650
- TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
666
+ STD_TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
667
+ STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
668
652
- TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
653
- TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
669
+ STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
670
+ STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
671
655
- torch::Tensor a_scales;
656
- auto options = torch::TensorOptions().dtype(c_dtype).device(a.device());
657
- auto options_fp32 =
658
- torch::TensorOptions().dtype(at::kFloat).device(a.device());
672
+ torch::stable::Tensor a_scales;
673
+ constexpr auto kFloat = torch::headeronly::ScalarType::Float;
674
675
if (a_scales_or_none.has_value()) {
676
a_scales = a_scales_or_none.value();
662
- TORCH_CHECK(a_type.size_bits() == 8,
663
- "a_scales can only be used for 8bit activation.");
677
+ STD_TORCH_CHECK(a_type.size_bits() == 8,
678
+ "a_scales can only be used for 8bit activation.");
679
} else {
665
- a_scales = torch::empty({0}, options_fp32);
666
- TORCH_CHECK(a_type.size_bits() != 8,
667
- "the a_scales parameter must be passed for 8bit activation.");
680
+ a_scales = torch::stable::new_empty(a, {0}, kFloat);
681
+ STD_TORCH_CHECK(
682
+ a_type.size_bits() != 8,
683
+ "the a_scales parameter must be passed for 8bit activation.");
684
}
685
686
// sms: number of SMs to use for the kernel
@@ -672,82 +688,84 @@ torch::Tensor moe_wna16_marlin_gemm(
688
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device());
689
690
// Alloc buffers
675
- const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
676
- torch::Tensor c;
691
+ torch::stable::accelerator::DeviceGuard device_guard(a.get_device_index());
692
+ torch::stable::Tensor c;
693
if (c_or_none.has_value()) {
694
c = c_or_none.value();
679
- TORCH_CHECK(c.device().is_cuda(), "c is not on GPU");
680
- TORCH_CHECK(c.is_contiguous(), "c is not contiguous");
681
- TORCH_CHECK(c.size(0) == size_m * top_k,
682
- "Shape mismatch: c.size(0) = ", c.size(0),
683
- ", size_m * topk = ", size_m * top_k);
684
- TORCH_CHECK(c.size(1) == size_n, "Shape mismatch: c.size(1) = ", c.size(1),
685
- ", size_n = ", size_n);
695
+ STD_TORCH_CHECK(c.device().is_cuda(), "c is not on GPU");
696
+ STD_TORCH_CHECK(c.is_contiguous(), "c is not contiguous");
697
+ STD_TORCH_CHECK(c.size(0) == size_m * top_k,
698
+ "Shape mismatch: c.size(0) = ", c.size(0),
699
+ ", size_m * topk = ", size_m * top_k);
700
+ STD_TORCH_CHECK(c.size(1) == size_n,
701
+ "Shape mismatch: c.size(1) = ", c.size(1),
702
+ ", size_n = ", size_n);
703
} else {
687
- c = torch::empty({size_m * top_k, size_n}, options);
704
+ c = torch::stable::new_empty(a, {size_m * top_k, size_n}, c_dtype);
705
}
706
707
// Alloc C tmp buffer that is going to be used for the global reduce
691
- torch::Tensor c_tmp;
708
+ torch::stable::Tensor c_tmp;
709
if (use_fp32_reduce && !use_atomic_add) {
710
// max num of threadblocks is sms * 4
711
long max_c_tmp_size = min(
712
(long)size_n * sorted_token_ids.size(0),
713
(long)sms * 4 * moe_block_size * MARLIN_NAMESPACE_NAME::max_thread_n);
714
if (moe_block_size == 8) max_c_tmp_size *= 2;
698
- c_tmp = torch::empty({max_c_tmp_size}, options_fp32);
715
+ c_tmp = torch::stable::new_empty(a, {max_c_tmp_size}, kFloat);
716
} else {
700
- c_tmp = torch::empty({0}, options_fp32);
717
+ c_tmp = torch::stable::new_empty(a, {0}, kFloat);
718
}
719
720
// Detect groupsize and act_order
721
int num_groups = -1;
722
int group_size = -1;
723
707
- int rank = b_scales.sizes().size();
708
- TORCH_CHECK(rank == 3, "b_scales rank = ", rank, " is not 3");
709
- TORCH_CHECK(b_scales.size(2) == size_n, "b_scales dim 2 = ", b_scales.size(2),
710
- " is not size_n = ", size_n);
724
+ int rank = b_scales.dim();
725
+ STD_TORCH_CHECK(rank == 3, "b_scales rank = ", rank, " is not 3");
726
+ STD_TORCH_CHECK(b_scales.size(2) == size_n,
727
+ "b_scales dim 2 = ", b_scales.size(2),
728
+ " is not size_n = ", size_n);
729
num_groups = b_scales.size(1);
730
713
- torch::Tensor g_idx, perm, a_tmp;
731
+ torch::stable::Tensor g_idx, perm, a_tmp;
732
if (g_idx_or_none.has_value() && perm_or_none.has_value()) {
733
g_idx = g_idx_or_none.value();
734
perm = perm_or_none.value();
735
718
- TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU");
719
- TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous");
720
- TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
721
- TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
736
+ STD_TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU");
737
+ STD_TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous");
738
+ STD_TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
739
+ STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
740
741
// Verify g_idx and perm
724
- TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) ||
725
- (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
726
- "Unexpected g_idx.size(-1) = ", g_idx.size(-1),
727
- " and perm.size(-1) = ", perm.size(-1),
728
- ", where size_k = ", size_k);
742
+ STD_TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) ||
743
+ (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
744
+ "Unexpected g_idx.size(-1) = ", g_idx.size(-1),
745
+ " and perm.size(-1) = ", perm.size(-1),
746
+ ", where size_k = ", size_k);
747
} else {
730
- g_idx = torch::empty({0}, options);
731
- perm = torch::empty({0}, options);
732
- a_tmp = torch::empty({0}, options);
748
+ g_idx = torch::stable::new_empty(a, {0}, c_dtype);
749
+ perm = torch::stable::new_empty(a, {0}, c_dtype);
750
+ a_tmp = torch::stable::new_empty(a, {0}, c_dtype);
751
}
752
bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0;
753
754
if (has_act_order) {
737
- a_tmp = torch::empty({size_m * top_k, size_k}, options);
755
+ a_tmp = torch::stable::new_empty(a, {size_m * top_k, size_k}, c_dtype);
756
if (is_k_full) {
739
- TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1");
740
- TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k,
741
- ", is not divisible by num_groups = ", num_groups);
757
+ STD_TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1");
758
+ STD_TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k,
759
+ ", is not divisible by num_groups = ", num_groups);
760
group_size = size_k / num_groups;
761
} else {
762
group_size = 0;
763
}
764
765
} else {
748
- a_tmp = torch::empty({0}, options);
766
+ a_tmp = torch::stable::new_empty(a, {0}, c_dtype);
767
if (num_groups > 1) {
750
- TORCH_CHECK(
768
+ STD_TORCH_CHECK(
769
size_k % num_groups == 0, "size_k = ", size_k,
770
", is not divisible by b_scales.size(1) = ", b_scales.size(1));
771
group_size = size_k / num_groups;
@@ -756,119 +774,125 @@ torch::Tensor moe_wna16_marlin_gemm(
774
}
775
}
776
759
- torch::Tensor global_scale;
777
+ torch::stable::Tensor global_scale;
778
if (global_scale_or_none.has_value()) {
779
global_scale = global_scale_or_none.value();
762
- TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
763
- "global_scale can only be used for nvfp4 format.");
780
+ STD_TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
781
+ "global_scale can only be used for nvfp4 format.");
782
} else {
765
- global_scale = torch::empty({0}, options_fp32);
766
- TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
767
- "the global_scale parameter must be passed for nvfp4 format.");
783
+ global_scale = torch::stable::new_empty(a, {0}, kFloat);
784
+ STD_TORCH_CHECK(
785
+ !(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
786
+ "the global_scale parameter must be passed for nvfp4 format.");
787
}
788
789
bool has_bias = b_bias_or_none.has_value();
771
- torch::Tensor b_bias;
790
+ torch::stable::Tensor b_bias;
791
if (has_bias) {
792
b_bias = b_bias_or_none.value();
774
- TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU");
775
- TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous");
776
- TORCH_CHECK(b_bias.size(1) == size_n, "b_bias.size(1) != size_n");
777
- TORCH_CHECK(b_bias.stride(1) == 1, "b_bias.stride(1) != 1");
793
+ STD_TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU");
794
+ STD_TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous");
795
+ STD_TORCH_CHECK(b_bias.size(1) == size_n, "b_bias.size(1) != size_n");
796
+ STD_TORCH_CHECK(b_bias.stride(1) == 1, "b_bias.stride(1) != 1");
797
} else {
779
- b_bias = torch::empty({0}, options);
798
+ b_bias = torch::stable::new_empty(a, {0}, c_dtype);
799
}
800
782
- torch::Tensor b_zeros;
801
+ torch::stable::Tensor b_zeros;
802
if (b_zeros_or_none.has_value()) {
803
b_zeros = b_zeros_or_none.value();
785
- TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU");
786
- TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous");
804
+ STD_TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU");
805
+ STD_TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous");
806
} else {
788
- b_zeros = torch::empty({0}, options);
807
+ b_zeros = torch::stable::new_empty(a, {0}, c_dtype);
808
}
809
bool has_zp = b_zeros.size(-1) > 0;
810
if (has_zp) {
792
- TORCH_CHECK(
811
+ STD_TORCH_CHECK(
812
b_type == vllm::kU4 || b_type == vllm::kU8,
813
"b_type must be u4 or u8 when has_zp = True. Got = ", b_type.str());
814
} else {
796
- TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 ||
797
- b_type == vllm::kS4 || b_type == vllm::kS8 ||
798
- b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f,
799
- "b_type must be uint4b8, uint8b128, int4, int8, "
800
- "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ",
801
- b_type.str());
815
+ STD_TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 ||
816
+ b_type == vllm::kS4 || b_type == vllm::kS8 ||
817
+ b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f,
818
+ "b_type must be uint4b8, uint8b128, int4, int8, "
819
+ "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ",
820
+ b_type.str());
821
}
822
823
if (has_zp && is_zp_float) {
805
- TORCH_CHECK(a.scalar_type() == at::ScalarType::Half,
806
- "Computation type must be float16 (half) when using float zero "
807
- "points.");
824
+ STD_TORCH_CHECK(
825
+ a.scalar_type() == torch::headeronly::ScalarType::Half,
826
+ "Computation type must be float16 (half) when using float zero "
827
+ "points.");
828
}
829
830
// Verify b_zeros
831
if (has_zp) {
812
- int rank = b_zeros.sizes().size();
813
- TORCH_CHECK(rank == 3, "b_zeros rank = ", rank, " is not 3");
832
+ int rank = b_zeros.dim();
833
+ STD_TORCH_CHECK(rank == 3, "b_zeros rank = ", rank, " is not 3");
834
if (is_zp_float) {
815
- TORCH_CHECK(b_zeros.size(2) == size_n,
816
- "b_zeros dim 2 = ", b_zeros.size(2),
817
- " is not size_n = ", size_n);
818
- TORCH_CHECK(num_groups == b_zeros.size(1),
819
- "b_zeros dim 1 = ", b_zeros.size(1),
820
- " is not num_groups = ", num_groups);
821
- TORCH_CHECK(num_groups != -1, "num_groups must be != -1");
835
+ STD_TORCH_CHECK(b_zeros.size(2) == size_n,
836
+ "b_zeros dim 2 = ", b_zeros.size(2),
837
+ " is not size_n = ", size_n);
838
+ STD_TORCH_CHECK(num_groups == b_zeros.size(1),
839
+ "b_zeros dim 1 = ", b_zeros.size(1),
840
+ " is not num_groups = ", num_groups);
841
+ STD_TORCH_CHECK(num_groups != -1, "num_groups must be != -1");
842
} else {
823
- TORCH_CHECK(b_zeros.size(1) == num_groups,
824
- "b_zeros dim 1 = ", b_zeros.size(1),
825
- " is not num_groups = ", num_groups);
826
- TORCH_CHECK(b_zeros.size(2) == size_n / pack_factor,
827
- "b_zeros dim 2 = ", b_zeros.size(2),
828
- " is not size_n / pack_factor = ", size_n / pack_factor);
843
+ STD_TORCH_CHECK(b_zeros.size(1) == num_groups,
844
+ "b_zeros dim 1 = ", b_zeros.size(1),
845
+ " is not num_groups = ", num_groups);
846
+ STD_TORCH_CHECK(b_zeros.size(2) == size_n / pack_factor,
847
+ "b_zeros dim 2 = ", b_zeros.size(2),
848
+ " is not size_n / pack_factor = ", size_n / pack_factor);
849
}
850
}
851
852
// Verify workspace size
833
- TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0,
834
- "size_n = ", size_n, ", is not divisible by min_thread_n = ",
835
- MARLIN_NAMESPACE_NAME::min_thread_n);
853
+ STD_TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0,
854
+ "size_n = ", size_n, ", is not divisible by min_thread_n = ",
855
+ MARLIN_NAMESPACE_NAME::min_thread_n);
856
857
int max_n_tiles = size_n / MARLIN_NAMESPACE_NAME::min_thread_n;
858
int min_workspace_size = min(
859
max_n_tiles * (int)(sorted_token_ids.size(0) / moe_block_size), sms * 4);
840
- TORCH_CHECK(workspace.numel() >= min_workspace_size,
841
- "workspace.numel = ", workspace.numel(),
842
- " is below min_workspace_size = ", min_workspace_size);
860
+ STD_TORCH_CHECK(workspace.numel() >= min_workspace_size,
861
+ "workspace.numel = ", workspace.numel(),
862
+ " is below min_workspace_size = ", min_workspace_size);
863
864
int dev = a.get_device();
865
846
- TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
847
- "scalar type of a_scales must be float");
848
- TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float,
849
- "scalar type of global_scale must be float");
866
+ STD_TORCH_CHECK(
867
+ a_scales.scalar_type() == torch::headeronly::ScalarType::Float,
868
+ "scalar type of a_scales must be float");
869
+ STD_TORCH_CHECK(
870
+ global_scale.scalar_type() == torch::headeronly::ScalarType::Float,
871
+ "scalar type of global_scale must be float");
872
if (a_type.size_bits() == 16) {
851
- TORCH_CHECK(
873
+ STD_TORCH_CHECK(
874
a.scalar_type() == c.scalar_type(),
875
"scalar type of a must be the same with c for 16 bit activation");
876
}
877
878
MARLIN_NAMESPACE_NAME::marlin_mm(
857
- a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(),
858
- b_bias.data_ptr(), a_scales.data_ptr(), b_scales.data_ptr(),
859
- global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(),
860
- perm.data_ptr(), a_tmp.data_ptr(), sorted_token_ids.data_ptr(),
861
- expert_ids.data_ptr(), num_tokens_past_padded.data_ptr(),
862
- topk_weights.data_ptr(), moe_block_size, num_experts, top_k,
863
- mul_topk_weights, size_m, size_n, size_k, workspace.data_ptr(), a_type,
864
- b_type, c_type, s_type, has_bias, has_act_order, is_k_full, has_zp,
865
- num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
879
+ a.const_data_ptr(), b_q_weight.const_data_ptr(), c.mutable_data_ptr(),
880
+ c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(),
881
+ a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(),
882
+ global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(),
883
+ g_idx.mutable_data_ptr(), perm.mutable_data_ptr(),
884
+ a_tmp.mutable_data_ptr(), sorted_token_ids.mutable_data_ptr(),
885
+ expert_ids.mutable_data_ptr(), num_tokens_past_padded.mutable_data_ptr(),
886
+ topk_weights.mutable_data_ptr(), moe_block_size, num_experts, top_k,
887
+ mul_topk_weights, size_m, size_n, size_k, workspace.mutable_data_ptr(),
888
+ a_type, b_type, c_type, s_type, has_bias, has_act_order, is_k_full,
889
+ has_zp, num_groups, group_size, dev, get_current_cuda_stream(dev),
890
thread_k, thread_n, sms, blocks_per_sm, use_atomic_add, use_fp32_reduce,
891
is_zp_float);
892
893
return c;
894
}
895
872
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
873
- m.impl("moe_wna16_marlin_gemm", &moe_wna16_marlin_gemm);
896
+STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
897
+ m.impl("moe_wna16_marlin_gemm", TORCH_BOX(&moe_wna16_marlin_gemm));
898
}
csrc/libtorch_stable/moe/moeTopKFuncs.cuh
renamed
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
renamed
+154
-117
@@ -1,14 +1,17 @@
1
-#include <torch/all.h>
2
-#include <ATen/cuda/CUDAContext.h>
3
-#include <c10/cuda/CUDAGuard.h>
1
+#include <array>
2
#include <cub/cub.cuh>
3
6
-#include <ATen/ATen.h>
7
-#include <ATen/cuda/Atomic.cuh>
4
+#include <cuda_runtime.h>
5
+#include <torch/csrc/stable/macros.h>
6
+#include <torch/csrc/stable/accelerator.h>
7
+#include <torch/csrc/stable/ops.h>
8
+#include <torch/csrc/stable/tensor.h>
9
+#include <torch/headeronly/core/ScalarType.h>
10
9
-#include "../cuda_compat.h"
10
-#include "../dispatch_utils.h"
11
+#include "../../cuda_compat.h"
12
#include "core/math.hpp"
13
+#include "libtorch_stable/dispatch_utils.h"
14
+#include "libtorch_stable/torch_utils.h"
15
16
#define CEILDIV(x, y) (((x) + (y) - 1) / (y))
17
@@ -492,12 +495,13 @@ __global__ void moe_lora_align_block_size_small_batch_expert_kernel(
495
496
// taken from
497
// https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc
495
-void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
496
- int64_t block_size, torch::Tensor sorted_token_ids,
497
- torch::Tensor experts_ids,
498
- torch::Tensor num_tokens_post_pad,
499
- std::optional<torch::Tensor> maybe_expert_map) {
500
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
498
+void moe_align_block_size(
499
+ torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
500
+ torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
501
+ torch::stable::Tensor num_tokens_post_pad,
502
+ std::optional<torch::stable::Tensor> maybe_expert_map) {
503
+ const cudaStream_t stream =
504
+ get_current_cuda_stream(topk_ids.get_device_index());
505
506
int64_t padded_num_experts =
507
((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
@@ -506,19 +510,18 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
510
threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
511
512
// BlockScan uses 1024 threads and assigns one thread per expert.
509
- TORCH_CHECK(padded_num_experts < 1024,
510
- "padded_num_experts must be less than 1024");
511
- auto options_int =
512
- torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
513
+ STD_TORCH_CHECK(padded_num_experts < 1024,
514
+ "padded_num_experts must be less than 1024");
515
bool has_expert_map = maybe_expert_map.has_value();
514
- torch::Tensor expert_map;
516
+ torch::stable::Tensor expert_map;
517
if (has_expert_map) {
518
expert_map = maybe_expert_map.value();
519
} else {
518
- expert_map = torch::empty({0}, options_int);
520
+ expert_map = torch::stable::new_empty(topk_ids, {0},
521
+ torch::headeronly::ScalarType::Int);
522
}
523
521
- VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
524
+ VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
525
topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
526
// calc needed amount of shared mem for `cumsum` tensors
527
bool small_batch_expert_mode =
@@ -538,16 +541,17 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
541
scalar_t, fill_threads>;
542
small_batch_expert_kernel<<<1, fill_threads + threads,
543
shared_mem_size, stream>>>(
541
- topk_ids.data_ptr<scalar_t>(),
542
- sorted_token_ids.data_ptr<int32_t>(),
543
- experts_ids.data_ptr<int32_t>(),
544
- num_tokens_post_pad.data_ptr<int32_t>(),
545
- expert_map.data_ptr<int32_t>(), num_experts, block_size,
546
- topk_ids.numel(), sorted_token_ids.size(0), topk_ids.size(1),
547
- has_expert_map);
544
+ reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
545
+ reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
546
+ reinterpret_cast<int32_t*>(experts_ids.mutable_data_ptr()),
547
+ reinterpret_cast<int32_t*>(
548
+ num_tokens_post_pad.mutable_data_ptr()),
549
+ reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
550
+ num_experts, block_size, topk_ids.numel(),
551
+ sorted_token_ids.size(0), topk_ids.size(1), has_expert_map);
552
} else {
549
- torch::Tensor cumsum_buffer =
550
- torch::empty({num_experts + 1}, options_int);
553
+ torch::stable::Tensor cumsum_buffer = torch::stable::new_empty(
554
+ topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int);
555
auto align_kernel = vllm::moe::moe_align_block_size_kernel<scalar_t>;
556
557
size_t num_warps = CEILDIV(padded_num_experts, experts_per_warp);
@@ -558,14 +562,16 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
562
// blockIdx.x == 0: counting experts and aligning
563
// blockIdx.x == 1: filling sorted_token_ids
564
align_kernel<<<2, threads, shared_mem_size, stream>>>(
561
- topk_ids.data_ptr<scalar_t>(),
562
- sorted_token_ids.data_ptr<int32_t>(),
563
- experts_ids.data_ptr<int32_t>(),
564
- num_tokens_post_pad.data_ptr<int32_t>(),
565
- expert_map.data_ptr<int32_t>(), num_experts, padded_num_experts,
566
- experts_per_warp, block_size, topk_ids.numel(),
567
- cumsum_buffer.data_ptr<int32_t>(), sorted_token_ids.size(0),
568
- topk_ids.size(1), has_expert_map);
565
+ reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
566
+ reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
567
+ reinterpret_cast<int32_t*>(experts_ids.mutable_data_ptr()),
568
+ reinterpret_cast<int32_t*>(
569
+ num_tokens_post_pad.mutable_data_ptr()),
570
+ reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
571
+ num_experts, padded_num_experts, experts_per_warp, block_size,
572
+ topk_ids.numel(),
573
+ reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()),
574
+ sorted_token_ids.size(0), topk_ids.size(1), has_expert_map);
575
576
const int block_threads = std::min(256, (int)threads);
577
const int num_blocks =
@@ -577,9 +583,10 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
583
auto sort_kernel =
584
vllm::moe::count_and_sort_expert_tokens_kernel<scalar_t>;
585
sort_kernel<<<gridDims, block_threads, 0, stream>>>(
580
- topk_ids.data_ptr<scalar_t>(),
581
- sorted_token_ids.data_ptr<int32_t>(),
582
- cumsum_buffer.data_ptr<int32_t>(), expert_map.data_ptr<int32_t>(),
586
+ reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
587
+ reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
588
+ reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()),
589
+ reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
590
topk_ids.numel(), num_experts, sorted_token_ids.size(0),
591
topk_ids.size(1), has_expert_map);
592
}
@@ -588,33 +595,36 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
595
596
void batched_moe_align_block_size(int64_t max_tokens_per_batch,
597
int64_t block_size,
591
- torch::Tensor const& batch_num_tokens,
592
- torch::Tensor sorted_ids,
593
- torch::Tensor batch_ids,
594
- torch::Tensor num_tokens_post_pad) {
598
+ const torch::stable::Tensor& batch_num_tokens,
599
+ torch::stable::Tensor sorted_ids,
600
+ torch::stable::Tensor batch_ids,
601
+ torch::stable::Tensor num_tokens_post_pad) {
602
namespace batched_kernel = vllm::moe::batched_moe_align_block_size;
603
597
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
604
+ const cudaStream_t stream =
605
+ get_current_cuda_stream(batch_num_tokens.get_device_index());
606
int32_t const B = batch_num_tokens.size(0);
607
int32_t const num_blocks_per_batch =
608
round_to_next_multiple_of(max_tokens_per_batch, block_size) / block_size;
609
int32_t const num_blocks = num_blocks_per_batch * B;
610
int64_t const sorted_ids_size = num_blocks * block_size;
611
604
- TORCH_CHECK(sorted_ids.size(0) == sorted_ids_size);
605
- TORCH_CHECK(batch_ids.size(0) == sorted_ids_size / block_size);
606
- TORCH_CHECK(num_tokens_post_pad.size(0) == 1);
607
- TORCH_CHECK(B <= batched_kernel::num_threads);
612
+ STD_TORCH_CHECK(sorted_ids.size(0) == sorted_ids_size);
613
+ STD_TORCH_CHECK(batch_ids.size(0) == sorted_ids_size / block_size);
614
+ STD_TORCH_CHECK(num_tokens_post_pad.size(0) == 1);
615
+ STD_TORCH_CHECK(B <= batched_kernel::num_threads);
616
617
batched_kernel::batched_moe_align_block_size_kernel<<<
618
batched_kernel::num_blocks, batched_kernel::num_threads, 0, stream>>>(
611
- B, max_tokens_per_batch, block_size, batch_num_tokens.data_ptr<int32_t>(),
612
- sorted_ids.data_ptr<int32_t>(), batch_ids.data_ptr<int32_t>(),
613
- num_tokens_post_pad.data_ptr<int32_t>());
619
+ B, max_tokens_per_batch, block_size,
620
+ reinterpret_cast<const int32_t*>(batch_num_tokens.const_data_ptr()),
621
+ reinterpret_cast<int32_t*>(sorted_ids.mutable_data_ptr()),
622
+ reinterpret_cast<int32_t*>(batch_ids.mutable_data_ptr()),
623
+ reinterpret_cast<int32_t*>(num_tokens_post_pad.mutable_data_ptr()));
624
}
625
616
-void moe_sum(torch::Tensor& input, // [num_tokens, topk, hidden_size]
617
- torch::Tensor& output) // [num_tokens, hidden_size]
626
+void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size]
627
+ torch::stable::Tensor& output) // [num_tokens, hidden_size]
628
{
629
const int hidden_size = input.size(-1);
630
const auto num_tokens = output.numel() / hidden_size;
@@ -622,77 +632,86 @@ void moe_sum(torch::Tensor& input, // [num_tokens, topk, hidden_size]
632
633
dim3 grid(num_tokens);
634
dim3 block(std::min(hidden_size, 1024));
625
- const at::cuda::OptionalCUDAGuard device_guard(device_of(output));
626
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
635
+ const torch::stable::accelerator::DeviceGuard device_guard(
636
+ output.get_device_index());
637
+ const cudaStream_t stream =
638
+ get_current_cuda_stream(output.get_device_index());
639
640
switch (topk) {
641
case 2:
630
- VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum_kernel", [&] {
631
- vllm::moe::moe_sum_kernel<scalar_t, 2><<<grid, block, 0, stream>>>(
632
- output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
633
- hidden_size);
634
- });
642
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
643
+ input.scalar_type(), "moe_sum_kernel", [&] {
644
+ vllm::moe::moe_sum_kernel<scalar_t, 2><<<grid, block, 0, stream>>>(
645
+ reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
646
+ reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
647
+ hidden_size);
648
+ });
649
break;
650
651
case 3:
638
- VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum_kernel", [&] {
639
- vllm::moe::moe_sum_kernel<scalar_t, 3><<<grid, block, 0, stream>>>(
640
- output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
641
- hidden_size);
642
- });
652
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
653
+ input.scalar_type(), "moe_sum_kernel", [&] {
654
+ vllm::moe::moe_sum_kernel<scalar_t, 3><<<grid, block, 0, stream>>>(
655
+ reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
656
+ reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
657
+ hidden_size);
658
+ });
659
break;
660
661
case 4:
646
- VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum_kernel", [&] {
647
- vllm::moe::moe_sum_kernel<scalar_t, 4><<<grid, block, 0, stream>>>(
648
- output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
649
- hidden_size);
650
- });
662
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
663
+ input.scalar_type(), "moe_sum_kernel", [&] {
664
+ vllm::moe::moe_sum_kernel<scalar_t, 4><<<grid, block, 0, stream>>>(
665
+ reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
666
+ reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
667
+ hidden_size);
668
+ });
669
break;
670
671
default:
654
- at::sum_out(output, input, 1);
672
+ torch::stable::sum_out(output, input, std::array<int64_t, 1>{1});
673
break;
674
}
675
}
676
677
void moe_lora_align_block_size(
660
- torch::Tensor topk_ids, torch::Tensor token_lora_mapping,
678
+ torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping,
679
int64_t num_experts, int64_t block_size, int64_t max_loras,
680
int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
663
- torch::Tensor sorted_token_ids, torch::Tensor expert_ids,
664
- torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled,
665
- torch::Tensor lora_ids, std::optional<torch::Tensor> maybe_expert_map) {
681
+ torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
682
+ torch::stable::Tensor num_tokens_post_pad,
683
+ torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids,
684
+ std::optional<torch::stable::Tensor> maybe_expert_map) {
685
const int topk_num = topk_ids.size(1);
686
668
- TORCH_CHECK(block_size > 0, "block_size should be greater than 0. ");
687
+ STD_TORCH_CHECK(block_size > 0, "block_size should be greater than 0. ");
688
689
int device_max_shared_mem;
671
- auto dev = topk_ids.get_device();
690
+ int dev = topk_ids.get_device_index();
691
cudaDeviceGetAttribute(&device_max_shared_mem,
692
cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
674
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
693
+ const cudaStream_t stream = get_current_cuda_stream(dev);
694
695
int64_t padded_num_experts =
696
((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
697
698
// BlockScan uses 1024 threads and assigns one thread per expert.
680
- TORCH_CHECK(padded_num_experts < 1024,
681
- "padded_num_experts must be less than 1024");
699
+ STD_TORCH_CHECK(padded_num_experts < 1024,
700
+ "padded_num_experts must be less than 1024");
701
683
- auto options_int =
684
- torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
685
- torch::Tensor token_mask =
686
- torch::empty({max_loras * topk_ids.size(0)}, options_int);
702
+ torch::stable::Tensor token_mask =
703
+ torch::stable::new_empty(topk_ids, {max_loras * topk_ids.size(0)},
704
+ torch::headeronly::ScalarType::Int);
705
bool has_expert_map = maybe_expert_map.has_value();
688
- torch::Tensor expert_map;
706
+ torch::stable::Tensor expert_map;
707
if (has_expert_map) {
708
expert_map = maybe_expert_map.value();
709
} else {
692
- expert_map = torch::empty({0}, options_int);
710
+ expert_map = torch::stable::new_empty(topk_ids, {0},
711
+ torch::headeronly::ScalarType::Int);
712
}
713
695
- VLLM_DISPATCH_INTEGRAL_TYPES(
714
+ VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(
715
topk_ids.scalar_type(), "moe_lora_align_sum_kernel", [&] {
716
bool small_batch_expert_mode =
717
(topk_ids.numel() < 1024) && (num_experts <= 64);
@@ -703,7 +722,7 @@ void moe_lora_align_block_size(
722
(num_thread + 1) * num_experts * sizeof(int32_t) +
723
(num_experts + 1) * sizeof(int32_t);
724
if (shared_mem > device_max_shared_mem) {
706
- TORCH_CHECK(false, "Shared memory usage exceeds device limit.");
725
+ STD_TORCH_CHECK(false, "Shared memory usage exceeds device limit.");
726
}
727
728
// threadIdx.x >= fill_threads: counting experts and aligning
@@ -714,7 +733,7 @@ void moe_lora_align_block_size(
733
auto kernel =
734
vllm::moe::moe_lora_align_block_size_small_batch_expert_kernel<
735
scalar_t, fill_threads>;
717
- AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
736
+ STD_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
737
(void*)kernel, shared_mem));
738
// Grid size is (max_loras + 1) because active_lora_ids has length
739
// max_loras + 1: sorted-unique values of token_lora_mapping, which
@@ -725,15 +744,21 @@ void moe_lora_align_block_size(
744
// MoE-LoRA kernels. This mirrors the fix made for the Triton
745
// _fused_moe_lora_kernel grid in vllm-project/vllm#32277.
746
kernel<<<max_loras + 1, blockDim, shared_mem, stream>>>(
728
- topk_ids.data_ptr<scalar_t>(),
729
- token_lora_mapping.data_ptr<int32_t>(), block_size,
730
- expert_map.data_ptr<int32_t>(), num_experts, max_loras,
731
- topk_ids.numel(), max_num_tokens_padded, max_num_m_blocks,
732
- sorted_token_ids.data_ptr<int32_t>(),
733
- expert_ids.data_ptr<int32_t>(), topk_num,
734
- num_tokens_post_pad.data_ptr<int32_t>(),
735
- adapter_enabled.data_ptr<int32_t>(), lora_ids.data_ptr<int32_t>(),
736
- token_mask.data_ptr<int32_t>(), has_expert_map);
747
+ reinterpret_cast<scalar_t*>(topk_ids.mutable_data_ptr()),
748
+ reinterpret_cast<int32_t*>(token_lora_mapping.mutable_data_ptr()),
749
+ block_size,
750
+ reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
751
+ num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded,
752
+ max_num_m_blocks,
753
+ reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
754
+ reinterpret_cast<int32_t*>(expert_ids.mutable_data_ptr()),
755
+ topk_num,
756
+ reinterpret_cast<int32_t*>(
757
+ num_tokens_post_pad.mutable_data_ptr()),
758
+ reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
759
+ reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
760
+ reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()),
761
+ has_expert_map);
762
} else {
763
int num_thread = 1024;
764
dim3 blockDim(num_thread);
@@ -742,8 +767,9 @@ void moe_lora_align_block_size(
767
size_t shared_mem_size = num_warps * WARP_SIZE * sizeof(int32_t);
768
769
// cumsum buffer
745
- torch::Tensor cumsum =
746
- torch::zeros({max_loras * (num_experts + 1)}, options_int);
770
+ torch::stable::Tensor cumsum = torch::stable::new_zeros(
771
+ topk_ids, {max_loras * (num_experts + 1)},
772
+ torch::headeronly::ScalarType::Int);
773
774
auto align_kernel =
775
vllm::moe::moe_lora_align_block_size_kernel<scalar_t>;
@@ -759,16 +785,23 @@ void moe_lora_align_block_size(
785
// blockIdx.x % 2 == 1: filling sorted_token_ids
786
align_kernel<<<(max_loras + 1) * 2, blockDim, shared_mem_size,
787
stream>>>(
762
- topk_ids.data_ptr<scalar_t>(),
763
- token_lora_mapping.data_ptr<int32_t>(), block_size,
764
- expert_map.data_ptr<int32_t>(), num_experts, max_loras,
765
- topk_ids.numel(), max_num_tokens_padded, max_num_m_blocks,
766
- sorted_token_ids.data_ptr<int32_t>(),
767
- expert_ids.data_ptr<int32_t>(), topk_num,
768
- num_tokens_post_pad.data_ptr<int32_t>(),
769
- adapter_enabled.data_ptr<int32_t>(), cumsum.data_ptr<int32_t>(),
770
- WARP_SIZE, padded_num_experts, lora_ids.data_ptr<int32_t>(),
771
- token_mask.data_ptr<int32_t>(), has_expert_map);
788
+ reinterpret_cast<scalar_t*>(topk_ids.mutable_data_ptr()),
789
+ reinterpret_cast<int32_t*>(token_lora_mapping.mutable_data_ptr()),
790
+ block_size,
791
+ reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
792
+ num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded,
793
+ max_num_m_blocks,
794
+ reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
795
+ reinterpret_cast<int32_t*>(expert_ids.mutable_data_ptr()),
796
+ topk_num,
797
+ reinterpret_cast<int32_t*>(
798
+ num_tokens_post_pad.mutable_data_ptr()),
799
+ reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
800
+ reinterpret_cast<int32_t*>(cumsum.mutable_data_ptr()), WARP_SIZE,
801
+ padded_num_experts,
802
+ reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
803
+ reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()),
804
+ has_expert_map);
805
806
const int block_threads = std::min(256, (int)num_thread);
807
const int num_blocks =
@@ -785,12 +818,16 @@ void moe_lora_align_block_size(
818
vllm::moe::lora_count_and_sort_expert_tokens_kernel<scalar_t>;
819
820
sort_kernel<<<gridDims, block_threads, 0, stream>>>(
788
- topk_ids.data_ptr<scalar_t>(),
789
- sorted_token_ids.data_ptr<int32_t>(), cumsum.data_ptr<int32_t>(),
790
- expert_map.data_ptr<int32_t>(), topk_ids.numel(), num_experts,
791
- max_num_tokens_padded, topk_num, token_mask.data_ptr<int32_t>(),
792
- max_loras, lora_ids.data_ptr<int32_t>(),
793
- adapter_enabled.data_ptr<int32_t>(), has_expert_map);
821
+ reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
822
+ reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
823
+ reinterpret_cast<int32_t*>(cumsum.mutable_data_ptr()),
824
+ reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
825
+ topk_ids.numel(), num_experts, max_num_tokens_padded, topk_num,
826
+ reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()),
827
+ max_loras,
828
+ reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
829
+ reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
830
+ has_expert_map);
831
}
832
});
833
}
\ No newline at end of file
csrc/libtorch_stable/moe/moe_ops.h
new
+87
@@ -0,0 +1,87 @@
1
+#pragma once
2
+
3
+#include <torch/csrc/stable/tensor.h>
4
+
5
+#include <optional>
6
+#include <tuple>
7
+
8
+void topk_softmax(torch::stable::Tensor& topk_weights,
9
+ torch::stable::Tensor& topk_indices,
10
+ torch::stable::Tensor& token_expert_indices,
11
+ torch::stable::Tensor& gating_output, bool renormalize,
12
+ std::optional<torch::stable::Tensor> bias);
13
+
14
+void topk_sigmoid(torch::stable::Tensor& topk_weights,
15
+ torch::stable::Tensor& topk_indices,
16
+ torch::stable::Tensor& token_expert_indices,
17
+ torch::stable::Tensor& gating_output, bool renormalize,
18
+ std::optional<torch::stable::Tensor> bias);
19
+
20
+void topk_softplus_sqrt(
21
+ torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
22
+ torch::stable::Tensor& token_expert_indices,
23
+ torch::stable::Tensor& gating_output, bool renormalize,
24
+ double routed_scaling_factor,
25
+ const std::optional<torch::stable::Tensor>& correction_bias,
26
+ const std::optional<torch::stable::Tensor>& input_ids,
27
+ const std::optional<torch::stable::Tensor>& tid2eid);
28
+
29
+void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output);
30
+
31
+void moe_align_block_size(
32
+ torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
33
+ torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
34
+ torch::stable::Tensor num_tokens_post_pad,
35
+ std::optional<torch::stable::Tensor> maybe_expert_map);
36
+
37
+void batched_moe_align_block_size(
38
+ int64_t max_tokens_per_batch, int64_t block_size,
39
+ const torch::stable::Tensor& expert_num_tokens,
40
+ torch::stable::Tensor sorted_ids, torch::stable::Tensor expert_ids,
41
+ torch::stable::Tensor num_tokens_post_pad);
42
+
43
+void moe_lora_align_block_size(
44
+ torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping,
45
+ int64_t num_experts, int64_t block_size, int64_t max_loras,
46
+ int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
47
+ torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
48
+ torch::stable::Tensor num_tokens_post_pad,
49
+ torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids,
50
+ std::optional<torch::stable::Tensor> maybe_expert_map);
51
+#ifndef USE_ROCM
52
+torch::stable::Tensor moe_wna16_gemm(
53
+ torch::stable::Tensor input, torch::stable::Tensor output,
54
+ torch::stable::Tensor b_qweight, torch::stable::Tensor b_scales,
55
+ std::optional<torch::stable::Tensor> b_qzeros,
56
+ std::optional<torch::stable::Tensor> topk_weights,
57
+ torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
58
+ torch::stable::Tensor num_tokens_post_pad, int64_t top_k,
59
+ int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, int64_t BLOCK_SIZE_K,
60
+ int64_t bit);
61
+
62
+std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
63
+ const torch::stable::Tensor& scores, int64_t n_group, int64_t topk_group,
64
+ int64_t topk, bool renormalize, double routed_scaling_factor,
65
+ const torch::stable::Tensor& bias, int64_t scoring_func);
66
+#endif
67
+
68
+bool moe_permute_unpermute_supported();
69
+
70
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
71
+ int64_t num_expert);
72
+
73
+void shuffle_rows(const torch::stable::Tensor& input_tensor,
74
+ const torch::stable::Tensor& dst2src_map,
75
+ torch::stable::Tensor& output_tensor);
76
+
77
+#ifndef USE_ROCM
78
+// DeepSeek V3 optimized router GEMM kernel for SM90+
79
+// Computes output = mat_a @ mat_b.T where:
80
+// mat_a: [num_tokens, hidden_dim] in bf16
81
+// mat_b: [num_experts, hidden_dim] in bf16
82
+// output: [num_tokens, num_experts] in bf16 or fp32
83
+// Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
84
+void dsv3_router_gemm(torch::stable::Tensor& output,
85
+ const torch::stable::Tensor& mat_a,
86
+ const torch::stable::Tensor& mat_b);
87
+#endif
csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu
new
+319
@@ -0,0 +1,319 @@
1
+#include <cuda.h>
2
+#include <cuda_runtime.h>
3
+#include <torch/csrc/stable/accelerator.h>
4
+#include <torch/csrc/stable/ops.h>
5
+#include <torch/csrc/stable/tensor.h>
6
+#include <torch/headeronly/core/ScalarType.h>
7
+#include <torch/headeronly/util/Exception.h>
8
+
9
+#include "core/registration.h"
10
+#include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h"
11
+#include "libtorch_stable/torch_utils.h"
12
+
13
+#include <torch/csrc/stable/library.h>
14
+
15
+// moe_permute kernels require at least CUDA 12.0
16
+#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
17
+
18
+namespace {
19
+
20
+int64_t product_integers(torch::headeronly::IntHeaderOnlyArrayRef sizes) {
21
+ int64_t numel = 1;
22
+ for (int64_t s : sizes) {
23
+ numel *= s;
24
+ }
25
+ return numel;
26
+}
27
+
28
+torch::stable::Tensor maybe_allocate_tensor(
29
+ const std::optional<torch::stable::Tensor>& maybe_tensor,
30
+ torch::headeronly::IntHeaderOnlyArrayRef expected_sizes,
31
+ torch::headeronly::ScalarType dtype, torch::stable::Device device,
32
+ char const* name) {
33
+ auto expected_numel = product_integers(expected_sizes);
34
+ if (maybe_tensor.has_value()) {
35
+ auto tensor = maybe_tensor.value();
36
+ STD_TORCH_CHECK(tensor.device() == device, name,
37
+ " must be on the same device");
38
+ STD_TORCH_CHECK(tensor.scalar_type() == dtype, name,
39
+ " has incorrect dtype");
40
+ STD_TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
41
+ STD_TORCH_CHECK(tensor.numel() >= expected_numel, name,
42
+ " is too small for the requested shape");
43
+ auto flat_tensor = torch::stable::view(tensor, {tensor.numel()});
44
+ return torch::stable::view(
45
+ torch::stable::narrow(flat_tensor, 0, 0, expected_numel),
46
+ expected_sizes);
47
+ }
48
+ return torch::stable::empty(expected_sizes, dtype, std::nullopt, device);
49
+}
50
+
51
+} // namespace
52
+
53
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
54
+ int64_t n_expert) {
55
+ return static_cast<int64_t>(
56
+ CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert));
57
+}
58
+
59
+void moe_permute_impl(
60
+ const torch::stable::Tensor& input, // [n_token, hidden]
61
+ const torch::stable::Tensor& topk_ids, // [n_token, topk]
62
+ const torch::stable::Tensor& token_expert_indices, // [n_token, topk]
63
+ const std::optional<torch::stable::Tensor>& expert_map, // [n_expert]
64
+ int64_t n_expert, int64_t n_local_expert, int64_t topk,
65
+ torch::stable::Tensor& permuted_input, // [permuted_size, hidden]
66
+ torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1]
67
+ torch::stable::Tensor& inv_permuted_idx, // [n_token, topk]
68
+ torch::stable::Tensor& permuted_idx, // [permute_size]
69
+ const std::optional<torch::stable::Tensor>& maybe_sort_workspace,
70
+ const std::optional<torch::stable::Tensor>& maybe_permuted_experts_id,
71
+ const std::optional<torch::stable::Tensor>& maybe_sorted_row_idx,
72
+ const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort) {
73
+ STD_TORCH_CHECK(expert_first_token_offset.scalar_type() ==
74
+ torch::headeronly::ScalarType::Long,
75
+ "expert_first_token_offset must be int64");
76
+ STD_TORCH_CHECK(topk_ids.scalar_type() == torch::headeronly::ScalarType::Int,
77
+ "topk_ids must be int32");
78
+ STD_TORCH_CHECK(
79
+ token_expert_indices.scalar_type() == torch::headeronly::ScalarType::Int,
80
+ "token_expert_indices must be int32");
81
+ STD_TORCH_CHECK(
82
+ inv_permuted_idx.scalar_type() == torch::headeronly::ScalarType::Int,
83
+ "inv_permuted_idx must be int32");
84
+ STD_TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1,
85
+ "expert_first_token_offset shape != n_local_expert+1");
86
+ STD_TORCH_CHECK(
87
+ inv_permuted_idx.sizes().equals(token_expert_indices.sizes()),
88
+ "token_expert_indices shape must be same as inv_permuted_idx");
89
+
90
+ auto device = input.device();
91
+ auto n_token = input.sizes()[0];
92
+ auto n_hidden = input.sizes()[1];
93
+ auto expanded_rows = n_token * topk;
94
+ auto stream = get_current_cuda_stream(input.get_device_index());
95
+
96
+ auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert);
97
+ auto sort_workspace = maybe_allocate_tensor(
98
+ maybe_sort_workspace, {sorter_size}, torch::headeronly::ScalarType::Char,
99
+ device, "sort_workspace");
100
+ auto permuted_experts_id = maybe_allocate_tensor(
101
+ maybe_permuted_experts_id, topk_ids.sizes(),
102
+ torch::headeronly::ScalarType::Int, device, "permuted_experts_id");
103
+ auto sorted_row_idx = maybe_allocate_tensor(
104
+ maybe_sorted_row_idx, inv_permuted_idx.sizes(),
105
+ torch::headeronly::ScalarType::Int, device, "sorted_row_idx");
106
+
107
+ CubKeyValueSorter sorter{};
108
+ int64_t* valid_num_ptr = nullptr;
109
+ torch::stable::Tensor topk_ids_for_sort = topk_ids;
110
+
111
+ if (expert_map.has_value()) {
112
+ const int* expert_map_ptr = get_ptr<int>(expert_map.value());
113
+ valid_num_ptr =
114
+ get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
115
+ topk_ids_for_sort = maybe_allocate_tensor(
116
+ maybe_topk_ids_for_sort, topk_ids.sizes(),
117
+ torch::headeronly::ScalarType::Int, device, "topk_ids_for_sort");
118
+ torch::stable::copy_(topk_ids_for_sort, topk_ids);
119
+ preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
120
+ expert_map_ptr, n_expert, stream);
121
+ }
122
+
123
+ sortAndScanExpert(
124
+ get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
125
+ get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
126
+ get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
127
+ n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
128
+
129
+ MOE_DISPATCH(input.scalar_type(), [&] {
130
+ expandInputRowsKernelLauncher<scalar_t>(
131
+ get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
132
+ get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
133
+ get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
134
+ n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
135
+ });
136
+}
137
+
138
+void moe_permute(
139
+ const torch::stable::Tensor& input, // [n_token, hidden]
140
+ const torch::stable::Tensor& topk_ids, // [n_token, topk]
141
+ const torch::stable::Tensor& token_expert_indices, // [n_token, topk]
142
+ const std::optional<torch::stable::Tensor>& expert_map, // [n_expert]
143
+ int64_t n_expert, int64_t n_local_expert, int64_t topk,
144
+ torch::stable::Tensor& permuted_input, // [permuted_size, hidden]
145
+ torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1]
146
+ torch::stable::Tensor& inv_permuted_idx, // [n_token, topk]
147
+ torch::stable::Tensor& permuted_idx) { // [permute_size]
148
+ moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
149
+ n_local_expert, topk, permuted_input,
150
+ expert_first_token_offset, inv_permuted_idx, permuted_idx,
151
+ std::nullopt, std::nullopt, std::nullopt, std::nullopt);
152
+}
153
+
154
+void moe_permute_with_scratch(
155
+ const torch::stable::Tensor& input, const torch::stable::Tensor& topk_ids,
156
+ const torch::stable::Tensor& token_expert_indices,
157
+ const std::optional<torch::stable::Tensor>& expert_map, int64_t n_expert,
158
+ int64_t n_local_expert, int64_t topk, torch::stable::Tensor& permuted_input,
159
+ torch::stable::Tensor& expert_first_token_offset,
160
+ torch::stable::Tensor& inv_permuted_idx,
161
+ torch::stable::Tensor& permuted_idx, torch::stable::Tensor& sort_workspace,
162
+ torch::stable::Tensor& permuted_experts_id,
163
+ torch::stable::Tensor& sorted_row_idx,
164
+ torch::stable::Tensor& topk_ids_for_sort) {
165
+ moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
166
+ n_local_expert, topk, permuted_input,
167
+ expert_first_token_offset, inv_permuted_idx, permuted_idx,
168
+ sort_workspace, permuted_experts_id, sorted_row_idx,
169
+ topk_ids_for_sort);
170
+}
171
+
172
+void moe_unpermute(
173
+ const torch::stable::Tensor&
174
+ permuted_hidden_states, // [n_token * topk, hidden]
175
+ const torch::stable::Tensor& topk_weights, // [n_token, topk]
176
+ const torch::stable::Tensor& inv_permuted_idx, // [n_token, topk]
177
+ const std::optional<torch::stable::Tensor>&
178
+ expert_first_token_offset, // [n_local_expert+1]
179
+ int64_t topk,
180
+ torch::stable::Tensor& hidden_states) { // [n_token, hidden]
181
+ STD_TORCH_CHECK(
182
+ permuted_hidden_states.scalar_type() == hidden_states.scalar_type(),
183
+ "permuted_hidden_states dtype must be same as hidden_states");
184
+
185
+ auto n_token = hidden_states.size(0);
186
+ auto n_hidden = hidden_states.size(1);
187
+ auto stream = get_current_cuda_stream(hidden_states.get_device_index());
188
+
189
+ int64_t const* valid_ptr = nullptr;
190
+ if (expert_first_token_offset.has_value()) {
191
+ int n_local_expert = expert_first_token_offset.value().size(0) - 1;
192
+ valid_ptr =
193
+ get_ptr<int64_t>(expert_first_token_offset.value()) + n_local_expert;
194
+ }
195
+
196
+ MOE_DISPATCH(hidden_states.scalar_type(), [&] {
197
+ finalizeMoeRoutingKernelLauncher<scalar_t, scalar_t>(
198
+ get_ptr<scalar_t>(permuted_hidden_states),
199
+ get_ptr<scalar_t>(hidden_states), get_ptr<float>(topk_weights),
200
+ get_ptr<int>(inv_permuted_idx), n_token, n_hidden, topk, valid_ptr,
201
+ stream);
202
+ });
203
+}
204
+
205
+template <typename T>
206
+__global__ void shuffleInputRowsKernel(const T* input,
207
+ const int32_t* dst2src_map, T* output,
208
+ int64_t num_src_rows,
209
+ int64_t num_dst_rows, int64_t num_cols) {
210
+ int64_t dest_row_idx = blockIdx.x;
211
+ int64_t const source_row_idx = dst2src_map[dest_row_idx];
212
+
213
+ if (blockIdx.x < num_dst_rows) {
214
+ // Load 128-bits per thread
215
+ constexpr int64_t ELEM_PER_THREAD = 128 / sizeof(T) / 8;
216
+ using DataElem = cutlass::Array<T, ELEM_PER_THREAD>;
217
+
218
+ // Duplicate and permute rows
219
+ auto const* source_row_ptr =
220
+ reinterpret_cast<DataElem const*>(input + source_row_idx * num_cols);
221
+ auto* dest_row_ptr =
222
+ reinterpret_cast<DataElem*>(output + dest_row_idx * num_cols);
223
+
224
+ int64_t const start_offset = threadIdx.x;
225
+ int64_t const stride = blockDim.x;
226
+ int64_t const num_elems_in_col = num_cols / ELEM_PER_THREAD;
227
+
228
+ for (int elem_index = start_offset; elem_index < num_elems_in_col;
229
+ elem_index += stride) {
230
+ dest_row_ptr[elem_index] = source_row_ptr[elem_index];
231
+ }
232
+ }
233
+}
234
+
235
+void shuffle_rows(const torch::stable::Tensor& input_tensor,
236
+ const torch::stable::Tensor& dst2src_map,
237
+ torch::stable::Tensor& output_tensor) {
238
+ STD_TORCH_CHECK(input_tensor.scalar_type() == output_tensor.scalar_type(),
239
+ "Input and output tensors must have the same data type");
240
+
241
+ auto stream = get_current_cuda_stream(output_tensor.get_device_index());
242
+ const int64_t blocks = output_tensor.size(0);
243
+ const int64_t threads = 256;
244
+ const int64_t num_dest_rows = output_tensor.size(0);
245
+ const int64_t num_src_rows = input_tensor.size(0);
246
+ const int64_t num_cols = input_tensor.size(1);
247
+
248
+ STD_TORCH_CHECK(!(num_cols % (128 / input_tensor.element_size() / 8)),
249
+ "num_cols must be divisible by 128 / "
250
+ "input_tensor.element_size() / 8");
251
+
252
+ MOE_DISPATCH(input_tensor.scalar_type(), [&] {
253
+ shuffleInputRowsKernel<scalar_t><<<blocks, threads, 0, stream>>>(
254
+ reinterpret_cast<const scalar_t*>(input_tensor.const_data_ptr()),
255
+ reinterpret_cast<const int32_t*>(dst2src_map.const_data_ptr()),
256
+ reinterpret_cast<scalar_t*>(output_tensor.mutable_data_ptr()),
257
+ num_src_rows, num_dest_rows, num_cols);
258
+ });
259
+}
260
+
261
+#else
262
+
263
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
264
+ int64_t n_expert) {
265
+ STD_TORCH_CHECK(
266
+ false, "moe_permute_sort_workspace_size is not supported on CUDA < 12.0");
267
+}
268
+
269
+void moe_permute(const torch::stable::Tensor& input,
270
+ const torch::stable::Tensor& topk_ids,
271
+ const torch::stable::Tensor& token_expert_indices,
272
+ const std::optional<torch::stable::Tensor>& expert_map,
273
+ int64_t n_expert, int64_t n_local_expert, int64_t topk,
274
+ torch::stable::Tensor& permuted_input,
275
+ torch::stable::Tensor& expert_first_token_offset,
276
+ torch::stable::Tensor& inv_permuted_idx,
277
+ torch::stable::Tensor& permuted_idx) {
278
+ STD_TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0");
279
+}
280
+
281
+void moe_permute_with_scratch(
282
+ const torch::stable::Tensor& input, const torch::stable::Tensor& topk_ids,
283
+ const torch::stable::Tensor& token_expert_indices,
284
+ const std::optional<torch::stable::Tensor>& expert_map, int64_t n_expert,
285
+ int64_t n_local_expert, int64_t topk, torch::stable::Tensor& permuted_input,
286
+ torch::stable::Tensor& expert_first_token_offset,
287
+ torch::stable::Tensor& inv_permuted_idx,
288
+ torch::stable::Tensor& permuted_idx, torch::stable::Tensor& sort_workspace,
289
+ torch::stable::Tensor& permuted_experts_id,
290
+ torch::stable::Tensor& sorted_row_idx,
291
+ torch::stable::Tensor& topk_ids_for_sort) {
292
+ STD_TORCH_CHECK(false,
293
+ "moe_permute_with_scratch is not supported on CUDA < 12.0");
294
+}
295
+
296
+void moe_unpermute(
297
+ const torch::stable::Tensor& permuted_hidden_states,
298
+ const torch::stable::Tensor& topk_weights,
299
+ const torch::stable::Tensor& inv_permuted_idx,
300
+ const std::optional<torch::stable::Tensor>& expert_first_token_offset,
301
+ int64_t topk, torch::stable::Tensor& hidden_states) {
302
+ STD_TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0");
303
+}
304
+
305
+#endif
306
+
307
+bool moe_permute_unpermute_supported() {
308
+#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
309
+ return true;
310
+#else
311
+ return false;
312
+#endif
313
+}
314
+
315
+STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
316
+ m.impl("moe_permute", TORCH_BOX(&moe_permute));
317
+ m.impl("moe_permute_with_scratch", TORCH_BOX(&moe_permute_with_scratch));
318
+ m.impl("moe_unpermute", TORCH_BOX(&moe_unpermute));
319
+}
\ No newline at end of file
csrc/libtorch_stable/moe/moe_wna16.cu
renamed
+57
-49
@@ -1,11 +1,14 @@
1
+#include <algorithm>
2
2
-#include <torch/all.h>
3
-#include <c10/cuda/CUDAGuard.h>
4
-#include <ATen/cuda/CUDAContext.h>
3
#include <cuda_runtime.h>
4
+#include <torch/csrc/stable/accelerator.h>
5
+#include <torch/csrc/stable/ops.h>
6
+#include <torch/csrc/stable/tensor.h>
7
+#include <torch/headeronly/core/ScalarType.h>
8
9
#include <cuda_fp16.h>
10
#include <cuda_bf16.h>
11
+#include "libtorch_stable/torch_utils.h"
12
#include "moe_wna16_utils.h"
13
14
#define DIVIDE(x, size) (((x) + (size) - 1) / (size))
@@ -263,7 +266,7 @@ void run_moe_wna16_gemm(const scalar_t* input, scalar_t* output,
266
}
267
268
const int shared_mem_size = BLOCK_SIZE_M * BLOCK_SIZE_K * 2;
266
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
269
+ const cudaStream_t stream = get_current_cuda_stream();
270
kernel<<<gridDim, blockDim, shared_mem_size, stream>>>(
271
input, output, b_qweight, b_scales, b_qzeros, topk_weights,
272
sorted_token_ids, expert_ids, num_tokens_post_pad, num_experts,
@@ -271,17 +274,18 @@ void run_moe_wna16_gemm(const scalar_t* input, scalar_t* output,
274
BLOCK_SIZE_K, has_zp, mul_topk_weight);
275
}
276
274
-torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
275
- torch::Tensor b_qweight, torch::Tensor b_scales,
276
- std::optional<torch::Tensor> b_qzeros,
277
- std::optional<torch::Tensor> topk_weights,
278
- torch::Tensor sorted_token_ids,
279
- torch::Tensor expert_ids,
280
- torch::Tensor num_tokens_post_pad, int64_t top_k,
281
- int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
282
- int64_t BLOCK_SIZE_K, int64_t bit) {
283
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
284
- output.zero_();
277
+torch::stable::Tensor moe_wna16_gemm(
278
+ torch::stable::Tensor input, torch::stable::Tensor output,
279
+ torch::stable::Tensor b_qweight, torch::stable::Tensor b_scales,
280
+ std::optional<torch::stable::Tensor> b_qzeros,
281
+ std::optional<torch::stable::Tensor> topk_weights,
282
+ torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
283
+ torch::stable::Tensor num_tokens_post_pad, int64_t top_k,
284
+ int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, int64_t BLOCK_SIZE_K,
285
+ int64_t bit) {
286
+ const torch::stable::accelerator::DeviceGuard device_guard(
287
+ input.get_device_index());
288
+ torch::stable::zero_(output);
289
290
const int num_experts = b_qweight.size(0);
291
const int size_m = input.size(0);
@@ -291,52 +295,56 @@ torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
295
296
int64_t EM = sorted_token_ids.size(0);
297
if (size_m <= BLOCK_SIZE_M) {
294
- EM = min(EM, size_m * BLOCK_SIZE_M * top_k);
298
+ EM = std::min(EM, size_m * BLOCK_SIZE_M * top_k);
299
}
300
const int num_token_blocks = (EM + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M;
301
302
const uint32_t* b_qzeros_ptr;
303
if (b_qzeros.has_value())
300
- b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
304
+ b_qzeros_ptr = (const uint32_t*)b_qzeros.value().const_data_ptr<uint8_t>();
305
const float* topk_weights_ptr = nullptr;
306
if (topk_weights.has_value())
303
- topk_weights_ptr = (const float*)topk_weights.value().data_ptr<float>();
307
+ topk_weights_ptr =
308
+ (const float*)topk_weights.value().const_data_ptr<float>();
309
310
int groups_per_block_row = BLOCK_SIZE_K / group_size;
306
- TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
307
- TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
308
- "size_k must divisible by BLOCK_SIZE_K");
309
- TORCH_CHECK(BLOCK_SIZE_K % group_size == 0,
310
- "BLOCK_SIZE_K must divisible by group_size");
311
- TORCH_CHECK(BLOCK_SIZE_M <= 64, "BLOCK_SIZE_M must less or equal to 64");
312
- TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
313
- groups_per_block_row == 4 || groups_per_block_row == 8,
314
- "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
315
-
316
- if (input.scalar_type() == at::ScalarType::Half) {
311
+ STD_TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
312
+ STD_TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
313
+ "size_k must divisible by BLOCK_SIZE_K");
314
+ STD_TORCH_CHECK(BLOCK_SIZE_K % group_size == 0,
315
+ "BLOCK_SIZE_K must divisible by group_size");
316
+ STD_TORCH_CHECK(BLOCK_SIZE_M <= 64, "BLOCK_SIZE_M must less or equal to 64");
317
+ STD_TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
318
+ groups_per_block_row == 4 || groups_per_block_row == 8,
319
+ "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
320
+
321
+ if (input.scalar_type() == torch::headeronly::ScalarType::Half) {
322
run_moe_wna16_gemm<half>(
318
- (const half*)input.data_ptr<at::Half>(),
319
- (half*)output.data_ptr<at::Half>(),
320
- (const uint32_t*)b_qweight.data_ptr<uint8_t>(),
321
- (const half*)b_scales.data_ptr<at::Half>(), b_qzeros_ptr,
322
- topk_weights_ptr, sorted_token_ids.data_ptr<int32_t>(),
323
- expert_ids.data_ptr<int32_t>(), num_tokens_post_pad.data_ptr<int32_t>(),
324
- num_experts, group_size, num_token_blocks, top_k, size_m, size_n,
325
- size_k, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, bit,
326
- b_qzeros.has_value(), topk_weights.has_value());
327
- } else if (input.scalar_type() == at::ScalarType::BFloat16) {
323
+ reinterpret_cast<const half*>(input.const_data_ptr()),
324
+ reinterpret_cast<half*>(output.mutable_data_ptr()),
325
+ (const uint32_t*)b_qweight.const_data_ptr<uint8_t>(),
326
+ reinterpret_cast<const half*>(b_scales.const_data_ptr()), b_qzeros_ptr,
327
+ topk_weights_ptr, sorted_token_ids.const_data_ptr<int32_t>(),
328
+ expert_ids.const_data_ptr<int32_t>(),
329
+ num_tokens_post_pad.const_data_ptr<int32_t>(), num_experts, group_size,
330
+ num_token_blocks, top_k, size_m, size_n, size_k, BLOCK_SIZE_M,
331
+ BLOCK_SIZE_N, BLOCK_SIZE_K, bit, b_qzeros.has_value(),
332
+ topk_weights.has_value());
333
+ } else if (input.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
334
run_moe_wna16_gemm<nv_bfloat16>(
329
- (const nv_bfloat16*)input.data_ptr<at::BFloat16>(),
330
- (nv_bfloat16*)output.data_ptr<at::BFloat16>(),
331
- (const uint32_t*)b_qweight.data_ptr<uint8_t>(),
332
- (const nv_bfloat16*)b_scales.data_ptr<at::BFloat16>(), b_qzeros_ptr,
333
- topk_weights_ptr, sorted_token_ids.data_ptr<int32_t>(),
334
- expert_ids.data_ptr<int32_t>(), num_tokens_post_pad.data_ptr<int32_t>(),
335
- num_experts, group_size, num_token_blocks, top_k, size_m, size_n,
336
- size_k, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, bit,
337
- b_qzeros.has_value(), topk_weights.has_value());
335
+ reinterpret_cast<const nv_bfloat16*>(input.const_data_ptr()),
336
+ reinterpret_cast<nv_bfloat16*>(output.mutable_data_ptr()),
337
+ (const uint32_t*)b_qweight.const_data_ptr<uint8_t>(),
338
+ reinterpret_cast<const nv_bfloat16*>(b_scales.const_data_ptr()),
339
+ b_qzeros_ptr, topk_weights_ptr,
340
+ sorted_token_ids.const_data_ptr<int32_t>(),
341
+ expert_ids.const_data_ptr<int32_t>(),
342
+ num_tokens_post_pad.const_data_ptr<int32_t>(), num_experts, group_size,
343
+ num_token_blocks, top_k, size_m, size_n, size_k, BLOCK_SIZE_M,
344
+ BLOCK_SIZE_N, BLOCK_SIZE_K, bit, b_qzeros.has_value(),
345
+ topk_weights.has_value());
346
} else {
339
- TORCH_CHECK(false, "moe_wna16_gemm only supports bfloat16 and float16");
347
+ STD_TORCH_CHECK(false, "moe_wna16_gemm only supports bfloat16 and float16");
348
}
349
return output;
350
}
csrc/libtorch_stable/moe/moe_wna16_utils.h
renamed
csrc/libtorch_stable/moe/permute_unpermute_kernels/dispatch.h
new
+60
@@ -0,0 +1,60 @@
1
+#pragma once
2
+
3
+#include <cuda_fp8.h>
4
+#include <torch/headeronly/core/ScalarType.h>
5
+#include <torch/headeronly/util/Exception.h>
6
+
7
+#define MOE_SWITCH(TYPE, ...) \
8
+ const auto _st = (TYPE); \
9
+ switch (_st) { \
10
+ __VA_ARGS__ \
11
+ default: \
12
+ STD_TORCH_CHECK(false, "[moe permute]data type dispatch fail!") \
13
+ }
14
+
15
+#define MOE_DISPATCH_CASE(enum_type, ...) \
16
+ case enum_type: { \
17
+ using scalar_t = ScalarType2CudaType<enum_type>::type; \
18
+ __VA_ARGS__(); \
19
+ break; \
20
+ }
21
+
22
+#define MOE_DISPATCH_FLOAT_CASE(...) \
23
+ MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \
24
+ MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \
25
+ MOE_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) \
26
+ MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e5m2, __VA_ARGS__) \
27
+ MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
28
+ MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__)
29
+
30
+#define MOE_DISPATCH(TYPE, ...) \
31
+ MOE_SWITCH(TYPE, MOE_DISPATCH_FLOAT_CASE(__VA_ARGS__))
32
+
33
+template <torch::headeronly::ScalarType type>
34
+struct ScalarType2CudaType;
35
+
36
+template <>
37
+struct ScalarType2CudaType<torch::headeronly::ScalarType::Float> {
38
+ using type = float;
39
+};
40
+template <>
41
+struct ScalarType2CudaType<torch::headeronly::ScalarType::Half> {
42
+ using type = half;
43
+};
44
+template <>
45
+struct ScalarType2CudaType<torch::headeronly::ScalarType::BFloat16> {
46
+ using type = __nv_bfloat16;
47
+};
48
+// uint8 for packed fp4
49
+template <>
50
+struct ScalarType2CudaType<torch::headeronly::ScalarType::Byte> {
51
+ using type = uint8_t;
52
+};
53
+template <>
54
+struct ScalarType2CudaType<torch::headeronly::ScalarType::Float8_e5m2> {
55
+ using type = __nv_fp8_e5m2;
56
+};
57
+template <>
58
+struct ScalarType2CudaType<torch::headeronly::ScalarType::Float8_e4m3fn> {
59
+ using type = __nv_fp8_e4m3;
60
+};
\ No newline at end of file
csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
renamed
+7
-4
@@ -1,5 +1,7 @@
1
+#include <cuda.h>
2
+#include <torch/headeronly/util/Exception.h>
3
2
-#include "moe_permute_unpermute_kernel.h"
4
+#include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h"
5
6
// moe_permute kernels require at least CUDA 12.0
7
#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
@@ -48,9 +50,10 @@ void CubKeyValueSorter::run(void* workspace, size_t const workspace_size,
50
size_t expected_ws_size = getWorkspaceSize(num_key_value_pairs, num_experts_);
51
size_t actual_ws_size = workspace_size;
52
51
- TORCH_CHECK(expected_ws_size <= workspace_size,
52
- "[CubKeyValueSorter::run] The allocated workspace is too small "
53
- "to run this problem.");
53
+ STD_TORCH_CHECK(
54
+ expected_ws_size <= workspace_size,
55
+ "[CubKeyValueSorter::run] The allocated workspace is too small "
56
+ "to run this problem.");
57
cub::DeviceRadixSort::SortPairs(workspace, actual_ws_size, keys_in, keys_out,
58
values_in, values_out, num_key_value_pairs, 0,
59
num_bits_, stream);
csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
renamed
+9
-8
@@ -2,23 +2,24 @@
2
// reference from tensorrt_llm moe kernel implementation archive in
3
// https://github.com/BBuf/tensorrt-llm-moe/tree/master
4
5
-#include <c10/core/ScalarType.h>
6
-#include <torch/all.h>
7
-#include "dispatch.h"
5
+#include <torch/csrc/stable/tensor.h>
6
+
7
#include <cub/cub.cuh>
8
#include <cub/device/device_radix_sort.cuh>
9
#include <cub/util_type.cuh>
11
-#include "cutlass/numeric_size.h"
10
+
11
#include "cutlass/array.h"
12
+#include "cutlass/numeric_size.h"
13
+#include "libtorch_stable/moe/permute_unpermute_kernels/dispatch.h"
14
15
template <typename T>
15
-inline T* get_ptr(torch::Tensor& t) {
16
- return reinterpret_cast<T*>(t.data_ptr());
16
+inline T* get_ptr(torch::stable::Tensor& t) {
17
+ return reinterpret_cast<T*>(t.mutable_data_ptr());
18
}
19
20
template <typename T>
20
-inline const T* get_ptr(const torch::Tensor& t) {
21
- return reinterpret_cast<const T*>(t.data_ptr());
21
+inline const T* get_ptr(const torch::stable::Tensor& t) {
22
+ return reinterpret_cast<const T*>(t.const_data_ptr());
23
}
24
25
class CubKeyValueSorter {
csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl
renamed
csrc/libtorch_stable/moe/topk_softmax_kernels.cu
renamed
+72
-63
@@ -17,11 +17,16 @@
17
* limitations under the License.
18
*/
19
#include <type_traits>
20
-#include <torch/all.h>
21
-#include <ATen/cuda/CUDAContext.h>
22
-#include <c10/cuda/CUDAGuard.h>
23
-#include "../cuda_compat.h"
24
-#include "../cub_helpers.h"
20
+
21
+#include <cuda_runtime.h>
22
+#include <torch/csrc/stable/accelerator.h>
23
+#include <torch/csrc/stable/tensor.h>
24
+#include <torch/headeronly/core/ScalarType.h>
25
+#include <torch/headeronly/util/Exception.h>
26
+
27
+#include "../../cuda_compat.h"
28
+#include "../../cub_helpers.h"
29
+#include "libtorch_stable/torch_utils.h"
30
31
#ifndef USE_ROCM
32
#include <cuda_bf16.h>
@@ -713,7 +718,7 @@ void topkGatingKernelLauncher(
718
break;
719
#endif
720
default: {
716
- TORCH_CHECK(workspace != nullptr,
721
+ STD_TORCH_CHECK(workspace != nullptr,
722
"workspace must be provided for num_experts that are not a power of 2 or multiple of 64.");
723
static constexpr int TPB = 256;
724
if constexpr (SF == SCORING_SOFTMAX) {
@@ -723,7 +728,7 @@ void topkGatingKernelLauncher(
728
moeSigmoid<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
729
gating_output, nullptr, workspace, num_experts);
730
} else {
726
- TORCH_CHECK(false, "Unsupported scoring func");
731
+ STD_TORCH_CHECK(false, "Unsupported scoring func");
732
}
733
moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
734
workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
@@ -738,63 +743,65 @@ void topkGatingKernelLauncher(
743
744
template<typename ComputeType, vllm::moe::ScoringFunc SF>
745
void dispatch_topk_launch(
741
- torch::Tensor& gating_output,
742
- torch::Tensor& topk_weights,
743
- torch::Tensor& topk_indices,
744
- torch::Tensor& token_expert_indices,
745
- torch::Tensor& softmax_workspace,
746
+ torch::stable::Tensor& gating_output,
747
+ torch::stable::Tensor& topk_weights,
748
+ torch::stable::Tensor& topk_indices,
749
+ torch::stable::Tensor& token_expert_indices,
750
+ torch::stable::Tensor& softmax_workspace,
751
int num_tokens, int num_experts, int topk, bool renormalize,
747
- std::optional<torch::Tensor> bias,
752
+ std::optional<torch::stable::Tensor> bias,
753
cudaStream_t stream)
754
{
755
const float* bias_ptr = nullptr;
756
if (bias.has_value()) {
752
- const torch::Tensor& bias_tensor = bias.value();
753
- TORCH_CHECK(bias_tensor.scalar_type() == at::ScalarType::Float, "bias tensor must be float32");
754
- TORCH_CHECK(bias_tensor.dim() == 1, "bias tensor must be 1D");
755
- TORCH_CHECK(bias_tensor.size(0) == num_experts, "bias size mismatch, expected: ", num_experts);
756
- TORCH_CHECK(bias_tensor.is_contiguous(), "bias tensor must be contiguous");
757
- bias_ptr = bias_tensor.data_ptr<float>();
757
+ const torch::stable::Tensor& bias_tensor = bias.value();
758
+ STD_TORCH_CHECK(bias_tensor.scalar_type() == torch::headeronly::ScalarType::Float,
759
+ "bias tensor must be float32");
760
+ STD_TORCH_CHECK(bias_tensor.dim() == 1, "bias tensor must be 1D");
761
+ STD_TORCH_CHECK(bias_tensor.size(0) == num_experts,
762
+ "bias size mismatch, expected: ", num_experts);
763
+ STD_TORCH_CHECK(bias_tensor.is_contiguous(), "bias tensor must be contiguous");
764
+ bias_ptr = bias_tensor.const_data_ptr<float>();
765
}
766
760
- if (topk_indices.scalar_type() == at::ScalarType::Int) {
767
+ if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
768
vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
762
- reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
763
- topk_weights.data_ptr<float>(),
764
- topk_indices.data_ptr<int>(),
765
- token_expert_indices.data_ptr<int>(),
766
- softmax_workspace.data_ptr<float>(),
769
+ reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
770
+ topk_weights.mutable_data_ptr<float>(),
771
+ topk_indices.mutable_data_ptr<int>(),
772
+ token_expert_indices.mutable_data_ptr<int>(),
773
+ softmax_workspace.mutable_data_ptr<float>(),
774
num_tokens, num_experts, topk, renormalize,
775
bias_ptr, stream);
769
- } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
776
+ } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
777
vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
771
- reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
772
- topk_weights.data_ptr<float>(),
773
- topk_indices.data_ptr<uint32_t>(),
774
- token_expert_indices.data_ptr<int>(),
775
- softmax_workspace.data_ptr<float>(),
778
+ reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
779
+ topk_weights.mutable_data_ptr<float>(),
780
+ topk_indices.mutable_data_ptr<uint32_t>(),
781
+ token_expert_indices.mutable_data_ptr<int>(),
782
+ softmax_workspace.mutable_data_ptr<float>(),
783
num_tokens, num_experts, topk, renormalize,
784
bias_ptr, stream);
785
} else {
779
- TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
786
+ STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
787
vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
781
- reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
782
- topk_weights.data_ptr<float>(),
783
- topk_indices.data_ptr<int64_t>(),
784
- token_expert_indices.data_ptr<int>(),
785
- softmax_workspace.data_ptr<float>(),
788
+ reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
789
+ topk_weights.mutable_data_ptr<float>(),
790
+ topk_indices.mutable_data_ptr<int64_t>(),
791
+ token_expert_indices.mutable_data_ptr<int>(),
792
+ softmax_workspace.mutable_data_ptr<float>(),
793
num_tokens, num_experts, topk, renormalize,
794
bias_ptr, stream);
795
}
796
}
797
798
void topk_softmax(
792
- torch::Tensor& topk_weights, // [num_tokens, topk]
793
- torch::Tensor& topk_indices, // [num_tokens, topk]
794
- torch::Tensor& token_expert_indices, // [num_tokens, topk]
795
- torch::Tensor& gating_output, // [num_tokens, num_experts]
799
+ torch::stable::Tensor& topk_weights, // [num_tokens, topk]
800
+ torch::stable::Tensor& topk_indices, // [num_tokens, topk]
801
+ torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
802
+ torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
803
bool renormalize,
797
- std::optional<torch::Tensor> bias)
804
+ std::optional<torch::stable::Tensor> bias)
805
{
806
const int num_experts = gating_output.size(-1);
807
const auto num_tokens = gating_output.numel() / num_experts;
@@ -804,35 +811,36 @@ void topk_softmax(
811
const bool needs_workspace = !is_pow_2 || num_experts > 256;
812
const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
813
807
- const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
808
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
809
- const auto workspace_options = gating_output.options().dtype(at::ScalarType::Float);
810
- torch::Tensor softmax_workspace = torch::empty({workspace_size}, workspace_options);
814
+ torch::stable::accelerator::DeviceGuard guard(gating_output.get_device_index());
815
+ const cudaStream_t stream =
816
+ get_current_cuda_stream(gating_output.get_device_index());
817
+ auto softmax_workspace = torch::stable::new_empty(
818
+ gating_output, {workspace_size}, torch::headeronly::ScalarType::Float);
819
812
- if (gating_output.scalar_type() == at::ScalarType::Float) {
820
+ if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
821
dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
822
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
823
bias, stream);
816
- } else if (gating_output.scalar_type() == at::ScalarType::Half) {
824
+ } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
825
dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
826
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
827
bias, stream);
820
- } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
828
+ } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
829
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
830
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
831
bias, stream);
832
} else {
825
- TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
833
+ STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
834
}
835
}
836
837
void topk_sigmoid(
830
- torch::Tensor& topk_weights, // [num_tokens, topk]
831
- torch::Tensor& topk_indices, // [num_tokens, topk]
832
- torch::Tensor& token_expert_indices, // [num_tokens, topk]
833
- torch::Tensor& gating_output, // [num_tokens, num_experts]
838
+ torch::stable::Tensor& topk_weights, // [num_tokens, topk]
839
+ torch::stable::Tensor& topk_indices, // [num_tokens, topk]
840
+ torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
841
+ torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
842
bool renormalize,
835
- std::optional<torch::Tensor> bias)
843
+ std::optional<torch::stable::Tensor> bias)
844
{
845
const int num_experts = gating_output.size(-1);
846
const auto num_tokens = gating_output.numel() / num_experts;
@@ -842,24 +850,25 @@ void topk_sigmoid(
850
const bool needs_workspace = !is_pow_2 || num_experts > 256;
851
const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
852
845
- const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
846
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
847
- const auto workspace_options = gating_output.options().dtype(at::ScalarType::Float);
848
- torch::Tensor workspace = torch::empty({workspace_size}, workspace_options);
853
+ torch::stable::accelerator::DeviceGuard guard(gating_output.get_device_index());
854
+ const cudaStream_t stream =
855
+ get_current_cuda_stream(gating_output.get_device_index());
856
+ auto workspace = torch::stable::new_empty(
857
+ gating_output, {workspace_size}, torch::headeronly::ScalarType::Float);
858
850
- if (gating_output.scalar_type() == at::ScalarType::Float) {
859
+ if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
860
dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
861
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
862
bias, stream);
854
- } else if (gating_output.scalar_type() == at::ScalarType::Half) {
863
+ } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
864
dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
865
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
866
bias, stream);
858
- } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
867
+ } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
868
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
869
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
870
bias, stream);
871
} else {
863
- TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
872
+ STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
873
}
874
}
csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu
renamed
+68
-54
@@ -18,11 +18,16 @@
18
* limitations under the License.
19
*/
20
#include <type_traits>
21
-#include <torch/all.h>
22
-#include <ATen/cuda/CUDAContext.h>
23
-#include <c10/cuda/CUDAGuard.h>
24
-#include "../cuda_compat.h"
25
-#include "../cub_helpers.h"
21
+
22
+#include <cuda_runtime.h>
23
+#include <torch/csrc/stable/accelerator.h>
24
+#include <torch/csrc/stable/tensor.h>
25
+#include <torch/headeronly/core/ScalarType.h>
26
+#include <torch/headeronly/util/Exception.h>
27
+
28
+#include "../../cuda_compat.h"
29
+#include "../../cub_helpers.h"
30
+#include "libtorch_stable/torch_utils.h"
31
#ifndef USE_ROCM
32
#include <cuda_bf16.h>
33
#include <cuda_fp16.h>
@@ -618,7 +623,7 @@ void topkGatingSoftplusSqrtKernelLauncher(
623
LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
624
break;
625
default: {
621
- TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
626
+ STD_TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
627
}
628
}
629
}
@@ -628,100 +633,109 @@ void topkGatingSoftplusSqrtKernelLauncher(
633
634
template <typename ComputeType>
635
void dispatch_topk_softplus_sqrt_launch(
631
- const ComputeType* gating_output, torch::Tensor& topk_weights,
632
- torch::Tensor& topk_indices, torch::Tensor& token_expert_indices,
633
- int num_tokens, int num_experts, int topk, bool renormalize,
634
- double routed_scaling_factor,
635
- const c10::optional<torch::Tensor>& correction_bias,
636
- const c10::optional<torch::Tensor>& input_ids,
637
- const c10::optional<torch::Tensor>& tid2eid, cudaStream_t stream) {
636
+ const ComputeType* gating_output, torch::stable::Tensor& topk_weights,
637
+ torch::stable::Tensor& topk_indices,
638
+ torch::stable::Tensor& token_expert_indices, int num_tokens,
639
+ int num_experts, int topk, bool renormalize, double routed_scaling_factor,
640
+ const std::optional<torch::stable::Tensor>& correction_bias,
641
+ const std::optional<torch::stable::Tensor>& input_ids,
642
+ const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {
643
const float* bias_ptr = nullptr;
644
if (correction_bias.has_value()) {
640
- bias_ptr = correction_bias.value().data_ptr<float>();
645
+ bias_ptr = correction_bias.value().const_data_ptr<float>();
646
}
647
bool use_hash = false;
648
if (tid2eid.has_value()) {
644
- TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE");
649
+ STD_TORCH_CHECK(input_ids.has_value(),
650
+ "input_ids is required for hash MoE");
651
use_hash = true;
652
}
647
- if (topk_indices.scalar_type() == at::ScalarType::Int) {
653
+ if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
654
const int* input_ids_ptr = nullptr;
655
const int* tid2eid_ptr = nullptr;
656
if (tid2eid.has_value()) {
651
- input_ids_ptr = input_ids.value().data_ptr<int>();
652
- tid2eid_ptr = tid2eid.value().data_ptr<int>();
657
+ input_ids_ptr = input_ids.value().const_data_ptr<int>();
658
+ tid2eid_ptr = tid2eid.value().const_data_ptr<int>();
659
}
660
661
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int, ComputeType>(
656
- gating_output, topk_weights.data_ptr<float>(),
657
- topk_indices.data_ptr<int>(), token_expert_indices.data_ptr<int>(),
658
- num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
659
- bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
660
- } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
662
+ gating_output, topk_weights.mutable_data_ptr<float>(),
663
+ topk_indices.mutable_data_ptr<int>(),
664
+ token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
665
+ topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
666
+ input_ids_ptr, tid2eid_ptr, stream);
667
+ } else if (topk_indices.scalar_type() ==
668
+ torch::headeronly::ScalarType::UInt32) {
669
const uint32_t* input_ids_ptr = nullptr;
670
const uint32_t* tid2eid_ptr = nullptr;
671
if (tid2eid.has_value()) {
664
- input_ids_ptr = input_ids.value().data_ptr<uint32_t>();
665
- tid2eid_ptr = tid2eid.value().data_ptr<uint32_t>();
672
+ input_ids_ptr = input_ids.value().const_data_ptr<uint32_t>();
673
+ tid2eid_ptr = tid2eid.value().const_data_ptr<uint32_t>();
674
}
675
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<uint32_t, ComputeType>(
668
- gating_output, topk_weights.data_ptr<float>(),
669
- topk_indices.data_ptr<uint32_t>(), token_expert_indices.data_ptr<int>(),
670
- num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
671
- bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
676
+ gating_output, topk_weights.mutable_data_ptr<float>(),
677
+ topk_indices.mutable_data_ptr<uint32_t>(),
678
+ token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
679
+ topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
680
+ input_ids_ptr, tid2eid_ptr, stream);
681
} else {
673
- TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
682
+ STD_TORCH_CHECK(topk_indices.scalar_type() ==
683
+ torch::headeronly::ScalarType::Long);
684
685
const int64_t* input_ids_ptr = nullptr;
686
const int64_t* tid2eid_ptr = nullptr;
687
if (tid2eid.has_value()) {
678
- input_ids_ptr = input_ids.value().data_ptr<int64_t>();
679
- tid2eid_ptr = tid2eid.value().data_ptr<int64_t>();
688
+ input_ids_ptr = input_ids.value().const_data_ptr<int64_t>();
689
+ tid2eid_ptr = tid2eid.value().const_data_ptr<int64_t>();
690
}
691
692
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int64_t, ComputeType>(
683
- gating_output, topk_weights.data_ptr<float>(),
684
- topk_indices.data_ptr<int64_t>(), token_expert_indices.data_ptr<int>(),
685
- num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
686
- bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
693
+ gating_output, topk_weights.mutable_data_ptr<float>(),
694
+ topk_indices.mutable_data_ptr<int64_t>(),
695
+ token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
696
+ topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
697
+ input_ids_ptr, tid2eid_ptr, stream);
698
}
699
}
700
701
void topk_softplus_sqrt(
691
- torch::Tensor& topk_weights, // [num_tokens, topk]
692
- torch::Tensor& topk_indices, // [num_tokens, topk]
693
- torch::Tensor& token_expert_indices, // [num_tokens, topk]
694
- torch::Tensor& gating_output, // [num_tokens, num_experts]
702
+ torch::stable::Tensor& topk_weights, // [num_tokens, topk]
703
+ torch::stable::Tensor& topk_indices, // [num_tokens, topk]
704
+ torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
705
+ torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
706
bool renormalize, double routed_scaling_factor,
696
- const c10::optional<torch::Tensor>& correction_bias,
697
- const c10::optional<torch::Tensor>& input_ids,
698
- const c10::optional<torch::Tensor>& tid2eid) {
707
+ const std::optional<torch::stable::Tensor>& correction_bias,
708
+ const std::optional<torch::stable::Tensor>& input_ids,
709
+ const std::optional<torch::stable::Tensor>& tid2eid) {
710
const int num_experts = gating_output.size(-1);
711
const auto num_tokens = gating_output.numel() / num_experts;
712
const int topk = topk_weights.size(-1);
702
- const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
703
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
713
+ const torch::stable::accelerator::DeviceGuard guard(
714
+ gating_output.get_device_index());
715
+ const cudaStream_t stream =
716
+ get_current_cuda_stream(gating_output.get_device_index());
717
705
- if (gating_output.scalar_type() == at::ScalarType::Float) {
718
+ if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
719
dispatch_topk_softplus_sqrt_launch<float>(
707
- gating_output.data_ptr<float>(), topk_weights, topk_indices,
720
+ gating_output.const_data_ptr<float>(), topk_weights, topk_indices,
721
token_expert_indices, num_tokens, num_experts, topk, renormalize,
722
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
710
- } else if (gating_output.scalar_type() == at::ScalarType::Half) {
723
+ } else if (gating_output.scalar_type() ==
724
+ torch::headeronly::ScalarType::Half) {
725
dispatch_topk_softplus_sqrt_launch<__half>(
712
- reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
726
+ reinterpret_cast<const __half*>(gating_output.const_data_ptr()),
727
topk_weights, topk_indices, token_expert_indices, num_tokens,
728
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
729
input_ids, tid2eid, stream);
716
- } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
730
+ } else if (gating_output.scalar_type() ==
731
+ torch::headeronly::ScalarType::BFloat16) {
732
dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
718
- reinterpret_cast<const __nv_bfloat16*>(
719
- gating_output.data_ptr<at::BFloat16>()),
733
+ reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()),
734
topk_weights, topk_indices, token_expert_indices, num_tokens,
735
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
736
input_ids, tid2eid, stream);
737
} else {
724
- TORCH_CHECK(false, "Unsupported gating_output data type: ",
725
- gating_output.scalar_type());
738
+ STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
739
+ gating_output.scalar_type());
740
}
741
}
\ No newline at end of file
csrc/libtorch_stable/moe/torch_bindings.cpp
renamed
+30
-17
@@ -1,32 +1,30 @@
1
#include "core/registration.h"
2
#include "moe_ops.h"
3
4
-TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
4
+#include <torch/csrc/stable/library.h>
5
+
6
+STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
7
// Apply topk softmax to the gating outputs.
8
m.def(
9
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
10
"token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
11
"bias) -> ()");
10
- m.impl("topk_softmax", torch::kCUDA, &topk_softmax);
12
13
// Apply topk sigmoid to the gating outputs.
14
m.def(
15
"topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
16
"token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
17
"bias) -> ()");
17
- m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
18
19
m.def(
20
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
21
"token_expert_indices, Tensor gating_output, bool renormalize, float "
22
"routed_scaling_factor, Tensor? "
23
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
24
- m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
24
25
// Calculate the result of moe by summing up the partial results
26
// from all selected experts.
27
m.def("moe_sum(Tensor input, Tensor! output) -> ()");
29
- m.impl("moe_sum", torch::kCUDA, &moe_sum);
28
29
// Aligning the number of tokens to be processed by each expert such
30
// that it is divisible by the block size.
@@ -36,7 +34,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
34
" Tensor! experts_ids,"
35
" Tensor! num_tokens_post_pad,"
36
" Tensor? maybe_expert_map) -> ()");
39
- m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
37
38
// Aligning the number of tokens to be processed by each expert such
39
// that it is divisible by the block size, but for the batched case.
@@ -46,8 +43,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
43
" Tensor! sorted_token_ids,"
44
" Tensor! experts_ids,"
45
" Tensor! num_tokens_post_pad) -> ()");
49
- m.impl("batched_moe_align_block_size", torch::kCUDA,
50
- &batched_moe_align_block_size);
46
47
// Aligning the number of tokens to be processed by each expert such
48
// that it is divisible by the block size.
@@ -64,8 +59,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
59
" Tensor !adapter_enabled,"
60
" Tensor !lora_ids,"
61
" Tensor? maybe_expert_map) -> () ");
67
- m.impl("moe_lora_align_block_size", torch::kCUDA, &moe_lora_align_block_size);
68
-
62
#ifndef USE_ROCM
63
m.def(
64
"moe_wna16_gemm(Tensor input, Tensor! output, Tensor b_qweight, "
@@ -75,8 +68,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
68
"int top_k, int BLOCK_SIZE_M, int BLOCK_SIZE_N, int BLOCK_SIZE_K, "
69
"int bit) -> Tensor");
70
78
- m.impl("moe_wna16_gemm", torch::kCUDA, &moe_wna16_gemm);
79
-
71
m.def(
72
"moe_wna16_marlin_gemm(Tensor! a, Tensor? c_or_none,"
73
"Tensor! b_q_weight, Tensor? b_bias_or_none,"
@@ -118,14 +109,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
109
m.def(
110
"moe_permute_sort_workspace_size(int num_expanded_rows, int n_expert) -> "
111
"int");
121
- m.impl("moe_permute_unpermute_supported", &moe_permute_unpermute_supported);
122
- m.impl("moe_permute_sort_workspace_size", &moe_permute_sort_workspace_size);
112
113
// Row shuffle for MoE
114
m.def(
115
"shuffle_rows(Tensor input_tensor, Tensor dst2src_map, Tensor! "
116
"output_tensor) -> ()");
128
- m.impl("shuffle_rows", torch::kCUDA, &shuffle_rows);
117
118
// Apply grouped topk routing to select experts.
119
m.def(
@@ -133,7 +121,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
121
"topk_group, int topk, bool renormalize, float "
122
"routed_scaling_factor, Tensor bias, int scoring_func) -> (Tensor, "
123
"Tensor)");
136
- m.impl("grouped_topk", torch::kCUDA, &grouped_topk);
124
125
// DeepSeek V3 optimized router GEMM for SM90+
126
m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
@@ -141,4 +128,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
128
#endif
129
}
130
144
-REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
131
+STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
132
+ m.impl("topk_softmax", TORCH_BOX(&topk_softmax));
133
+ m.impl("topk_sigmoid", TORCH_BOX(&topk_sigmoid));
134
+ m.impl("topk_softplus_sqrt", TORCH_BOX(&topk_softplus_sqrt));
135
+ m.impl("moe_sum", TORCH_BOX(&moe_sum));
136
+ m.impl("moe_align_block_size", TORCH_BOX(&moe_align_block_size));
137
+ m.impl("batched_moe_align_block_size",
138
+ TORCH_BOX(&batched_moe_align_block_size));
139
+ m.impl("moe_lora_align_block_size", TORCH_BOX(&moe_lora_align_block_size));
140
+#ifndef USE_ROCM
141
+ m.impl("moe_wna16_gemm", TORCH_BOX(&moe_wna16_gemm));
142
+ m.impl("shuffle_rows", TORCH_BOX(&shuffle_rows));
143
+ m.impl("grouped_topk", TORCH_BOX(&grouped_topk));
144
+#endif
145
+}
146
+
147
+#ifndef USE_ROCM
148
+// Primitive-only ops have no tensor to dispatch on.
149
+STABLE_TORCH_LIBRARY_IMPL(_moe_C, CompositeExplicitAutograd, m) {
150
+ m.impl("moe_permute_unpermute_supported",
151
+ TORCH_BOX(&moe_permute_unpermute_supported));
152
+ m.impl("moe_permute_sort_workspace_size",
153
+ TORCH_BOX(&moe_permute_sort_workspace_size));
154
+}
155
+#endif
156
+
157
+REGISTER_EXTENSION(_moe_C_stable_libtorch)
csrc/libtorch_stable/torch_bindings.cpp
+16
@@ -246,6 +246,22 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
246
"awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
247
"Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
248
249
+ // Expert-specialization mxfp8 blockscaled grouped quantization (SM100+).
250
+ ops.def(
251
+ "mxfp8_experts_quant("
252
+ " Tensor input, Tensor problem_sizes, Tensor expert_offsets,"
253
+ " Tensor blockscale_offsets, Tensor! quant_output, Tensor! scale_factor)"
254
+ " -> ()");
255
+ // conditionally compiled so impl registration is in source file
256
+
257
+ // Expert-specialization mxfp8 blockscaled grouped GEMM (SM100+).
258
+ ops.def(
259
+ "cutlass_mxfp8_grouped_mm("
260
+ " Tensor a, Tensor b, Tensor sfa, Tensor sfb, Tensor! out,"
261
+ " Tensor problem_sizes, Tensor expert_offsets, Tensor blockscale_offsets)"
262
+ " -> ()");
263
+ // conditionally compiled so impl registration is in source file
264
+
265
// DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
266
// conditionally compiled so impl registration is in source file
267
ops.def(
csrc/moe/dsv3_router_gemm_utils.h
deleted
-31
@@ -1,31 +0,0 @@
1
-/*
2
- * Adapted from SGLang's sgl-kernel implementation, which was adapted from
3
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
4
- * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
5
- *
6
- * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved.
7
- *
8
- * Licensed under the Apache License, Version 2.0 (the "License");
9
- * you may not use this file except in compliance with the License.
10
- * You may obtain a copy of the License at
11
- *
12
- * http://www.apache.org/licenses/LICENSE-2.0
13
- *
14
- * Unless required by applicable law or agreed to in writing, software
15
- * distributed under the License is distributed on an "AS IS" BASIS,
16
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17
- * See the License for the specific language governing permissions and
18
- * limitations under the License.
19
- */
20
-
21
-#pragma once
22
-
23
-#include <ATen/cuda/CUDAContext.h>
24
-
25
-#include <cstdlib>
26
-#include <mutex>
27
-
28
-inline int getSMVersion() {
29
- auto* props = at::cuda::getCurrentDeviceProperties();
30
- return props->major * 10 + props->minor;
31
-}
csrc/moe/moe_ops.h
deleted
-81
@@ -1,81 +0,0 @@
1
-#pragma once
2
-
3
-#include <torch/all.h>
4
-
5
-void topk_softmax(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
6
- torch::Tensor& token_expert_indices,
7
- torch::Tensor& gating_output, bool renormalize,
8
- std::optional<torch::Tensor> bias);
9
-
10
-void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
11
- torch::Tensor& token_expert_indices,
12
- torch::Tensor& gating_output, bool renormalize,
13
- std::optional<torch::Tensor> bias);
14
-
15
-void topk_softplus_sqrt(torch::Tensor& topk_weights,
16
- torch::Tensor& topk_indices,
17
- torch::Tensor& token_expert_indices,
18
- torch::Tensor& gating_output, bool renormalize,
19
- double routed_scaling_factor,
20
- const c10::optional<torch::Tensor>& correction_bias,
21
- const c10::optional<torch::Tensor>& input_ids,
22
- const c10::optional<torch::Tensor>& tid2eid);
23
-
24
-void moe_sum(torch::Tensor& input, torch::Tensor& output);
25
-
26
-void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
27
- int64_t block_size, torch::Tensor sorted_token_ids,
28
- torch::Tensor experts_ids,
29
- torch::Tensor num_tokens_post_pad,
30
- std::optional<torch::Tensor> maybe_expert_map);
31
-
32
-void batched_moe_align_block_size(int64_t max_tokens_per_batch,
33
- int64_t block_size,
34
- torch::Tensor const& expert_num_tokens,
35
- torch::Tensor sorted_ids,
36
- torch::Tensor expert_ids,
37
- torch::Tensor num_tokens_post_pad);
38
-
39
-void moe_lora_align_block_size(
40
- torch::Tensor topk_ids, torch::Tensor token_lora_mapping,
41
- int64_t num_experts, int64_t block_size, int64_t max_loras,
42
- int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
43
- torch::Tensor sorted_token_ids, torch::Tensor expert_ids,
44
- torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled,
45
- torch::Tensor lora_ids, std::optional<torch::Tensor> maybe_expert_map);
46
-#ifndef USE_ROCM
47
-torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
48
- torch::Tensor b_qweight, torch::Tensor b_scales,
49
- std::optional<torch::Tensor> b_qzeros,
50
- std::optional<torch::Tensor> topk_weights,
51
- torch::Tensor sorted_token_ids,
52
- torch::Tensor expert_ids,
53
- torch::Tensor num_tokens_post_pad, int64_t top_k,
54
- int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
55
- int64_t BLOCK_SIZE_K, int64_t bit);
56
-
57
-std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
58
- torch::Tensor const& scores, int64_t n_group, int64_t topk_group,
59
- int64_t topk, bool renormalize, double routed_scaling_factor,
60
- torch::Tensor const& bias, int64_t scoring_func);
61
-#endif
62
-
63
-bool moe_permute_unpermute_supported();
64
-
65
-int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
66
- int64_t num_experts);
67
-
68
-void shuffle_rows(const torch::Tensor& input_tensor,
69
- const torch::Tensor& dst2src_map,
70
- torch::Tensor& output_tensor);
71
-
72
-#ifndef USE_ROCM
73
-// DeepSeek V3 optimized router GEMM kernel for SM90+
74
-// Computes output = mat_a @ mat_b.T where:
75
-// mat_a: [num_tokens, hidden_dim] in bf16
76
-// mat_b: [num_experts, hidden_dim] in bf16
77
-// output: [num_tokens, num_experts] in bf16 or fp32
78
-// Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
79
-void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a,
80
- const torch::Tensor& mat_b);
81
-#endif
csrc/moe/moe_permute_unpermute_op.cu
deleted
-286
@@ -1,286 +0,0 @@
1
-#include <c10/core/ScalarType.h>
2
-#include <torch/all.h>
3
-#include <ATen/cuda/CUDAContext.h>
4
-#include "permute_unpermute_kernels/moe_permute_unpermute_kernel.h"
5
-#include "permute_unpermute_kernels/dispatch.h"
6
-#include "core/registration.h"
7
-
8
-// moe_permute kernels require at least CUDA 12.0
9
-#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
10
-
11
-namespace {
12
-
13
-torch::Tensor maybe_allocate_tensor(
14
- const std::optional<torch::Tensor>& maybe_tensor,
15
- at::IntArrayRef expected_sizes, torch::ScalarType dtype, c10::Device device,
16
- char const* name) {
17
- auto expected_numel = c10::multiply_integers(expected_sizes);
18
- if (maybe_tensor.has_value()) {
19
- auto tensor = maybe_tensor.value();
20
- TORCH_CHECK(tensor.device() == device, name, " must be on the same device");
21
- TORCH_CHECK(tensor.scalar_type() == dtype, name, " has incorrect dtype");
22
- TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
23
- TORCH_CHECK(tensor.numel() >= expected_numel, name,
24
- " is too small for the requested shape");
25
- auto flat_tensor = tensor.view({tensor.numel()});
26
- return flat_tensor.narrow(0, 0, expected_numel).view(expected_sizes);
27
- }
28
- return torch::empty(expected_sizes, torch::dtype(dtype).device(device));
29
-}
30
-
31
-} // namespace
32
-
33
-int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
34
- int64_t n_expert) {
35
- return static_cast<int64_t>(
36
- CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert));
37
-}
38
-
39
-void moe_permute_impl(
40
- const torch::Tensor& input, // [n_token, hidden]
41
- const torch::Tensor& topk_ids, // [n_token, topk]
42
- const torch::Tensor& token_expert_indices, // [n_token, topk]
43
- const std::optional<torch::Tensor>& expert_map, // [n_expert]
44
- int64_t n_expert, int64_t n_local_expert, int64_t topk,
45
- torch::Tensor& permuted_input, // [permuted_size, hidden]
46
- torch::Tensor& expert_first_token_offset, // [n_local_expert + 1]
47
- torch::Tensor& inv_permuted_idx, // [n_token, topk]
48
- torch::Tensor& permuted_idx, // [permute_size]
49
- const std::optional<torch::Tensor>& maybe_sort_workspace,
50
- const std::optional<torch::Tensor>& maybe_permuted_experts_id,
51
- const std::optional<torch::Tensor>& maybe_sorted_row_idx,
52
- const std::optional<torch::Tensor>& maybe_topk_ids_for_sort) {
53
- TORCH_CHECK(expert_first_token_offset.scalar_type() == at::ScalarType::Long,
54
- "expert_first_token_offset must be int64");
55
- TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
56
- "topk_ids must be int32");
57
- TORCH_CHECK(token_expert_indices.scalar_type() == at::ScalarType::Int,
58
- "token_expert_indices must be int32");
59
- TORCH_CHECK(inv_permuted_idx.scalar_type() == at::ScalarType::Int,
60
- "inv_permuted_idx must be int32");
61
- TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1,
62
- "expert_first_token_offset shape != n_local_expert+1");
63
- TORCH_CHECK(inv_permuted_idx.sizes() == token_expert_indices.sizes(),
64
- "token_expert_indices shape must be same as inv_permuted_idx");
65
- auto device = input.device();
66
- auto n_token = input.sizes()[0];
67
- auto n_hidden = input.sizes()[1];
68
- auto expanded_rows = n_token * topk;
69
- auto stream = at::cuda::getCurrentCUDAStream().stream();
70
-
71
- auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert);
72
- auto sort_workspace =
73
- maybe_allocate_tensor(maybe_sort_workspace, {sorter_size}, torch::kInt8,
74
- device, "sort_workspace");
75
- auto permuted_experts_id =
76
- maybe_allocate_tensor(maybe_permuted_experts_id, topk_ids.sizes(),
77
- at::ScalarType::Int, device, "permuted_experts_id");
78
- auto sorted_row_idx =
79
- maybe_allocate_tensor(maybe_sorted_row_idx, inv_permuted_idx.sizes(),
80
- at::ScalarType::Int, device, "sorted_row_idx");
81
-
82
- CubKeyValueSorter sorter{};
83
- int64_t* valid_num_ptr = nullptr;
84
- torch::Tensor topk_ids_for_sort = topk_ids;
85
-
86
- if (expert_map.has_value()) {
87
- const int* expert_map_ptr = get_ptr<int>(expert_map.value());
88
- valid_num_ptr =
89
- get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
90
- topk_ids_for_sort =
91
- maybe_allocate_tensor(maybe_topk_ids_for_sort, topk_ids.sizes(),
92
- at::ScalarType::Int, device, "topk_ids_for_sort");
93
- topk_ids_for_sort.copy_(topk_ids);
94
- preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
95
- expert_map_ptr, n_expert, stream);
96
- }
97
-
98
- sortAndScanExpert(
99
- get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
100
- get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
101
- get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
102
- n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
103
-
104
- MOE_DISPATCH(input.scalar_type(), [&] {
105
- expandInputRowsKernelLauncher<scalar_t>(
106
- get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
107
- get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
108
- get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
109
- n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
110
- });
111
-}
112
-
113
-void moe_permute(
114
- const torch::Tensor& input, // [n_token, hidden]
115
- const torch::Tensor& topk_ids, // [n_token, topk]
116
- const torch::Tensor& token_expert_indices, // [n_token, topk]
117
- const std::optional<torch::Tensor>& expert_map, // [n_expert]
118
- int64_t n_expert, int64_t n_local_expert, int64_t topk,
119
- torch::Tensor& permuted_input, // [permuted_size, hidden]
120
- torch::Tensor& expert_first_token_offset, // [n_local_expert + 1]
121
- torch::Tensor& inv_permuted_idx, // [n_token, topk]
122
- torch::Tensor& permuted_idx) { // [permute_size]
123
- moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
124
- n_local_expert, topk, permuted_input,
125
- expert_first_token_offset, inv_permuted_idx, permuted_idx,
126
- std::nullopt, std::nullopt, std::nullopt, std::nullopt);
127
-}
128
-
129
-void moe_permute_with_scratch(
130
- const torch::Tensor& input, const torch::Tensor& topk_ids,
131
- const torch::Tensor& token_expert_indices,
132
- const std::optional<torch::Tensor>& expert_map, int64_t n_expert,
133
- int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input,
134
- torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx,
135
- torch::Tensor& permuted_idx, torch::Tensor& sort_workspace,
136
- torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx,
137
- torch::Tensor& topk_ids_for_sort) {
138
- moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
139
- n_local_expert, topk, permuted_input,
140
- expert_first_token_offset, inv_permuted_idx, permuted_idx,
141
- sort_workspace, permuted_experts_id, sorted_row_idx,
142
- topk_ids_for_sort);
143
-}
144
-
145
-void moe_unpermute(
146
- const torch::Tensor& permuted_hidden_states, // [n_token * topk, hidden]
147
- const torch::Tensor& topk_weights, // [n_token, topk]
148
- const torch::Tensor& inv_permuted_idx, // [n_token, topk]
149
- const std::optional<torch::Tensor>&
150
- expert_first_token_offset, // [n_local_expert+1]
151
- int64_t topk,
152
- torch::Tensor& hidden_states // [n_token, hidden]
153
-) {
154
- TORCH_CHECK(
155
- permuted_hidden_states.scalar_type() == hidden_states.scalar_type(),
156
- "permuted_hidden_states dtype must be same as hidden_states");
157
- auto n_token = hidden_states.size(0);
158
- auto n_hidden = hidden_states.size(1);
159
- auto stream = at::cuda::getCurrentCUDAStream().stream();
160
-
161
- int64_t const* valid_ptr = nullptr;
162
- if (expert_first_token_offset.has_value()) {
163
- int n_local_expert = expert_first_token_offset.value().size(0) - 1;
164
- valid_ptr =
165
- get_ptr<int64_t>(expert_first_token_offset.value()) + n_local_expert;
166
- }
167
-
168
- MOE_DISPATCH(hidden_states.scalar_type(), [&] {
169
- finalizeMoeRoutingKernelLauncher<scalar_t, scalar_t>(
170
- get_ptr<scalar_t>(permuted_hidden_states),
171
- get_ptr<scalar_t>(hidden_states), get_ptr<float>(topk_weights),
172
- get_ptr<int>(inv_permuted_idx), n_token, n_hidden, topk, valid_ptr,
173
- stream);
174
- });
175
-}
176
-
177
-template <typename T>
178
-__global__ void shuffleInputRowsKernel(const T* input,
179
- const int32_t* dst2src_map, T* output,
180
- int64_t num_src_rows,
181
- int64_t num_dst_rows, int64_t num_cols) {
182
- int64_t dest_row_idx = blockIdx.x;
183
- int64_t const source_row_idx = dst2src_map[dest_row_idx];
184
-
185
- if (blockIdx.x < num_dst_rows) {
186
- // Load 128-bits per thread
187
- constexpr int64_t ELEM_PER_THREAD = 128 / sizeof(T) / 8;
188
- using DataElem = cutlass::Array<T, ELEM_PER_THREAD>;
189
-
190
- // Duplicate and permute rows
191
- auto const* source_row_ptr =
192
- reinterpret_cast<DataElem const*>(input + source_row_idx * num_cols);
193
- auto* dest_row_ptr =
194
- reinterpret_cast<DataElem*>(output + dest_row_idx * num_cols);
195
-
196
- int64_t const start_offset = threadIdx.x;
197
- int64_t const stride = blockDim.x;
198
- int64_t const num_elems_in_col = num_cols / ELEM_PER_THREAD;
199
-
200
- for (int elem_index = start_offset; elem_index < num_elems_in_col;
201
- elem_index += stride) {
202
- dest_row_ptr[elem_index] = source_row_ptr[elem_index];
203
- }
204
- }
205
-}
206
-
207
-void shuffle_rows(const torch::Tensor& input_tensor,
208
- const torch::Tensor& dst2src_map,
209
- torch::Tensor& output_tensor) {
210
- TORCH_CHECK(input_tensor.scalar_type() == output_tensor.scalar_type(),
211
- "Input and output tensors must have the same data type");
212
-
213
- auto stream = at::cuda::getCurrentCUDAStream().stream();
214
- int64_t const blocks = output_tensor.size(0);
215
- int64_t const threads = 256;
216
- int64_t const num_dest_rows = output_tensor.size(0);
217
- int64_t const num_src_rows = input_tensor.size(0);
218
- int64_t const num_cols = input_tensor.size(1);
219
-
220
- TORCH_CHECK(!(num_cols % (128 / sizeof(input_tensor.scalar_type()) / 8)),
221
- "num_cols must be divisible by 128 / "
222
- "sizeof(input_tensor.scalar_type()) / 8");
223
-
224
- MOE_DISPATCH(input_tensor.scalar_type(), [&] {
225
- shuffleInputRowsKernel<scalar_t><<<blocks, threads, 0, stream>>>(
226
- reinterpret_cast<scalar_t*>(input_tensor.data_ptr()),
227
- dst2src_map.data_ptr<int32_t>(),
228
- reinterpret_cast<scalar_t*>(output_tensor.data_ptr()), num_src_rows,
229
- num_dest_rows, num_cols);
230
- });
231
-}
232
-
233
-#else
234
-
235
-int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
236
- int64_t n_expert) {
237
- TORCH_CHECK(
238
- false, "moe_permute_sort_workspace_size is not supported on CUDA < 12.0");
239
-}
240
-
241
-void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids,
242
- const torch::Tensor& token_expert_indices,
243
- const std::optional<torch::Tensor>& expert_map,
244
- int64_t n_expert, int64_t n_local_expert, int64_t topk,
245
- torch::Tensor& permuted_input,
246
- torch::Tensor& expert_first_token_offset,
247
- torch::Tensor& inv_permuted_idx, torch::Tensor& permuted_idx) {
248
- TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0");
249
-}
250
-
251
-void moe_permute_with_scratch(
252
- const torch::Tensor& input, const torch::Tensor& topk_ids,
253
- const torch::Tensor& token_expert_indices,
254
- const std::optional<torch::Tensor>& expert_map, int64_t n_expert,
255
- int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input,
256
- torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx,
257
- torch::Tensor& permuted_idx, torch::Tensor& sort_workspace,
258
- torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx,
259
- torch::Tensor& topk_ids_for_sort) {
260
- TORCH_CHECK(false,
261
- "moe_permute_with_scratch is not supported on CUDA < 12.0");
262
-}
263
-
264
-void moe_unpermute(
265
- const torch::Tensor& permuted_hidden_states,
266
- const torch::Tensor& topk_weights, const torch::Tensor& inv_permuted_idx,
267
- const std::optional<torch::Tensor>& expert_first_token_offset, int64_t topk,
268
- torch::Tensor& hidden_states) {
269
- TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0");
270
-}
271
-
272
-#endif
273
-
274
-bool moe_permute_unpermute_supported() {
275
-#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
276
- return true;
277
-#else
278
- return false;
279
-#endif
280
-}
281
-
282
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
283
- m.impl("moe_permute", &moe_permute);
284
- m.impl("moe_permute_with_scratch", &moe_permute_with_scratch);
285
- m.impl("moe_unpermute", &moe_unpermute);
286
-}
\ No newline at end of file
csrc/moe/permute_unpermute_kernels/dispatch.h
deleted
-59
@@ -1,59 +0,0 @@
1
-#pragma once
2
-#include <cuda_fp8.h>
3
-#define MOE_SWITCH(TYPE, ...) \
4
- at::ScalarType _st = ::detail::scalar_type(TYPE); \
5
- switch (_st) { \
6
- __VA_ARGS__ \
7
- default: \
8
- TORCH_CHECK(false, "[moe permute]data type dispatch fail!") \
9
- }
10
-
11
-#define MOE_DISPATCH_CASE(enum_type, ...) \
12
- case enum_type: { \
13
- using scalar_t = ScalarType2CudaType<enum_type>::type; \
14
- __VA_ARGS__(); \
15
- break; \
16
- }
17
-#define MOE_DISPATCH_FLOAT_CASE(...) \
18
- MOE_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
19
- MOE_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
20
- MOE_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
21
- MOE_DISPATCH_CASE(at::ScalarType::Float8_e5m2, __VA_ARGS__) \
22
- MOE_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
23
- MOE_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__)
24
-
25
-#define MOE_DISPATCH(TYPE, ...) \
26
- MOE_SWITCH(TYPE, MOE_DISPATCH_FLOAT_CASE(__VA_ARGS__))
27
-
28
-template <at::ScalarType type>
29
-struct ScalarType2CudaType;
30
-
31
-template <>
32
-struct ScalarType2CudaType<at::ScalarType::Float> {
33
- using type = float;
34
-};
35
-template <>
36
-struct ScalarType2CudaType<at::ScalarType::Half> {
37
- using type = half;
38
-};
39
-template <>
40
-struct ScalarType2CudaType<at::ScalarType::BFloat16> {
41
- using type = __nv_bfloat16;
42
-};
43
-// uint8 for packed fp4
44
-template <>
45
-struct ScalarType2CudaType<at::ScalarType::Byte> {
46
- using type = uint8_t;
47
-};
48
-
49
-// #if __CUDA_ARCH__ >= 890
50
-// fp8
51
-template <>
52
-struct ScalarType2CudaType<at::ScalarType::Float8_e5m2> {
53
- using type = __nv_fp8_e5m2;
54
-};
55
-template <>
56
-struct ScalarType2CudaType<at::ScalarType::Float8_e4m3fn> {
57
- using type = __nv_fp8_e4m3;
58
-};
59
-// #endif
\ No newline at end of file
csrc/torch_bindings.cpp
-19
@@ -130,25 +130,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
130
"Tensor? qzeros_or_none, bool inplace) -> Tensor");
131
// conditionally compiled so impl registrations are in source file
132
133
-#endif
134
-
135
-#ifndef USE_ROCM
136
- // Expert-specialization mxfp8 blockscaled grouped quantization (SM100+).
137
- ops.def(
138
- "mxfp8_experts_quant("
139
- " Tensor input, Tensor problem_sizes, Tensor expert_offsets,"
140
- " Tensor blockscale_offsets, Tensor! quant_output, Tensor! scale_factor)"
141
- " -> ()");
142
- // conditionally compiled so impl registration is in source file
143
-
144
- // Expert-specialization mxfp8 blockscaled grouped GEMM (SM100+).
145
- ops.def(
146
- "cutlass_mxfp8_grouped_mm("
147
- " Tensor a, Tensor b, Tensor sfa, Tensor sfb, Tensor! out,"
148
- " Tensor problem_sizes, Tensor expert_offsets, Tensor blockscale_offsets)"
149
- " -> ()");
150
- // conditionally compiled so impl registration is in source file
151
-
133
#endif
134
}
135
setup.py
+2
-2
@@ -755,7 +755,7 @@ class precompiled_wheel_utils:
755
{
756
"vllm/_C.abi3.so",
757
"vllm/_C_stable_libtorch.abi3.so",
758
- "vllm/_moe_C.abi3.so",
758
+ "vllm/_moe_C_stable_libtorch.abi3.so",
759
"vllm/_flashmla_C.abi3.so",
760
"vllm/_flashmla_extension_C.abi3.so",
761
"vllm/_sparse_flashmla_C.abi3.so",
@@ -1081,7 +1081,6 @@ def get_requirements() -> list[str]:
1081
ext_modules = []
1082
1083
if _is_cuda() or _is_hip():
1084
- ext_modules.append(CMakeExtension(name="vllm._moe_C"))
1084
ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
1085
# Optional since this doesn't get built (produce an .so file). This is just
1086
# copying the relevant .py files from the source repository.
@@ -1135,6 +1134,7 @@ if _build_custom_ops():
1134
ext_modules.append(CMakeExtension(name="vllm._C"))
1135
if _is_cuda() or _is_hip():
1136
ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
1137
+ ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch"))
1138
1139
package_data = {
1140
"vllm": [
vllm/platforms/interface.py
+1
-1
@@ -254,7 +254,7 @@ class Platform:
254
except ImportError as e:
255
logger.warning("Failed to import from vllm._C: %r", e)
256
with contextlib.suppress(ImportError):
257
- import vllm._moe_C # noqa: F401
257
+ import vllm._moe_C_stable_libtorch # noqa: F401
258
259
@classmethod
260
def get_attn_backend_cls(