A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 14:00 UTC
[Perf] Fuse stride preparation for NVFP4 cutlass_moe (#31837)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jan 7, 2026 at 13:31 UTC
f347ac6c346f8e84a75ec975618c5aa79c1b9db5
1 file changed
+30
-18
csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu
+30
-18
@@ -62,7 +62,9 @@ __global__ void __get_group_gemm_starts(
62
ElementSF* a_scales_base_as_int, ElementSF* b_scales_base_as_int,
63
ElementAccumulator* alphas_base_as_int, const int32_t* expert_offsets,
64
const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes,
65
- const int K, const int N) {
65
+ int64_t* a_strides, int64_t* b_strides, int64_t* c_strides,
66
+ const int64_t a_stride_val, const int64_t b_stride_val,
67
+ const int64_t c_stride_val, const int K, const int N) {
68
int64_t expert_id = threadIdx.x;
69
if (expert_id >= gridDim.x * blockDim.x) {
70
return;
@@ -103,6 +105,11 @@ __global__ void __get_group_gemm_starts(
105
// Shape of alpha = [E]
106
alpha_offsets[expert_id] = alphas_base_as_int + expert_id;
107
108
+ // Initialize strides (constant across all experts, avoids separate kernels)
109
+ a_strides[expert_id] = a_stride_val;
110
+ b_strides[expert_id] = b_stride_val;
111
+ c_strides[expert_id] = c_stride_val;
112
+
113
LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
114
LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
115
@@ -135,7 +142,11 @@ __global__ void __get_group_gemm_starts(
142
static_cast<float*>(alphas.data_ptr()), \
143
static_cast<int32_t*>(expert_offsets.data_ptr()), \
144
static_cast<int32_t*>(sf_offsets.data_ptr()), \
138
- static_cast<int32_t*>(problem_sizes.data_ptr()), K, N); \
145
+ static_cast<int32_t*>(problem_sizes.data_ptr()), \
146
+ static_cast<int64_t*>(a_strides.data_ptr()), \
147
+ static_cast<int64_t*>(b_strides.data_ptr()), \
148
+ static_cast<int64_t*>(c_strides.data_ptr()), a_stride_val, \
149
+ b_stride_val, c_stride_val, K, N); \
150
}
151
152
template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
@@ -144,6 +155,9 @@ void run_get_group_gemm_starts(
155
const torch::Tensor& out_starts, const torch::Tensor& a_scales_starts,
156
const torch::Tensor& b_scales_starts, const torch::Tensor& alpha_starts,
157
const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb,
158
+ const torch::Tensor& a_strides, const torch::Tensor& b_strides,
159
+ const torch::Tensor& c_strides, int64_t a_stride_val, int64_t b_stride_val,
160
+ int64_t c_stride_val,
161
/*these are used for their base addresses*/
162
torch::Tensor const& a_tensors, torch::Tensor const& b_tensors,
163
torch::Tensor const& out_tensors, torch::Tensor const& a_scales,
@@ -269,17 +283,16 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
283
torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
284
torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
285
torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
272
- torch::Tensor c_strides1 =
273
- torch::full({num_experts}, output.stride(0), options_int);
274
- torch::Tensor a_strides1 =
275
- torch::full({num_experts}, a.stride(0) * 2, options_int);
276
- torch::Tensor b_strides1 =
277
- torch::full({num_experts}, b.stride(1) * 2, options_int);
286
+ torch::Tensor a_strides1 = torch::empty(num_experts, options_int);
287
+ torch::Tensor b_strides1 = torch::empty(num_experts, options_int);
288
+ torch::Tensor c_strides1 = torch::empty(num_experts, options_int);
289
290
run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
291
a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs,
281
- layout_sfa, layout_sfb, a, b, output, a_blockscale, b_blockscales, alphas,
282
- expert_offsets, sf_offsets, problem_sizes, M, N, K);
292
+ layout_sfa, layout_sfb, a_strides1, b_strides1, c_strides1,
293
+ a.stride(0) * 2, b.stride(1) * 2, output.stride(0), a, b, output,
294
+ a_blockscale, b_blockscales, alphas, expert_offsets, sf_offsets,
295
+ problem_sizes, M, N, K);
296
297
// Create an instance of the GEMM
298
Gemm gemm_op;
@@ -444,17 +457,16 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
457
torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
458
torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
459
torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
447
- torch::Tensor c_strides1 =
448
- torch::full({num_experts}, output.stride(0), options_int);
449
- torch::Tensor a_strides1 =
450
- torch::full({num_experts}, a.stride(0) * 2, options_int);
451
- torch::Tensor b_strides1 =
452
- torch::full({num_experts}, b.stride(1) * 2, options_int);
460
+ torch::Tensor a_strides1 = torch::empty(num_experts, options_int);
461
+ torch::Tensor b_strides1 = torch::empty(num_experts, options_int);
462
+ torch::Tensor c_strides1 = torch::empty(num_experts, options_int);
463
464
run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
465
a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs,
456
- layout_sfa, layout_sfb, a, b, output, a_blockscale, b_blockscales, alphas,
457
- expert_offsets, sf_offsets, problem_sizes, M, N, K);
466
+ layout_sfa, layout_sfb, a_strides1, b_strides1, c_strides1,
467
+ a.stride(0) * 2, b.stride(1) * 2, output.stride(0), a, b, output,
468
+ a_blockscale, b_blockscales, alphas, expert_offsets, sf_offsets,
469
+ problem_sizes, M, N, K);
470
471
// Create an instance of the GEMM
472
Gemm gemm_op;