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 14, 23:45 UTC
[Bug] Fix Compressed Tensor NVFP4 `cutlass_fp4_group_mm` illegal memory access (#21465)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 24, 2025 at 11:13 UTC
e8cb0d04950c9e0b39d00f3744806c19dff4b4bc
1 file changed
+15
-12
csrc/quantization/cutlass_w8a8/moe/moe_data.cu
+15
-12
@@ -47,13 +47,12 @@ __global__ void compute_problem_sizes(const int32_t* __restrict__ topk_ids,
47
48
__global__ void compute_expert_offsets(
49
const int32_t* __restrict__ problem_sizes1, int32_t* expert_offsets,
50
- int32_t* atomic_buffer, const int num_experts, const int topk_length) {
50
+ int32_t* atomic_buffer, const int num_experts, const bool swap_ab) {
51
int32_t tot_offset = 0;
52
expert_offsets[0] = 0;
53
for (int i = 0; i < num_experts; ++i) {
54
atomic_buffer[i] = tot_offset;
55
- tot_offset += topk_length > SWAP_AB_THRESHOLD ? problem_sizes1[i * 3]
56
- : problem_sizes1[i * 3 + 1];
55
+ tot_offset += swap_ab ? problem_sizes1[i * 3 + 1] : problem_sizes1[i * 3];
56
expert_offsets[i + 1] = tot_offset;
57
}
58
}
@@ -61,15 +60,14 @@ __global__ void compute_expert_offsets(
60
__global__ void compute_expert_blockscale_offsets(
61
const int32_t* __restrict__ problem_sizes1, int32_t* expert_offsets,
62
int32_t* blockscale_offsets, int32_t* atomic_buffer, const int num_experts,
64
- const int topk_length) {
63
+ const bool swap_ab) {
64
int32_t tot_offset = 0;
65
int32_t tot_offset_round = 0;
66
expert_offsets[0] = 0;
67
blockscale_offsets[0] = 0;
68
for (int i = 0; i < num_experts; ++i) {
70
- int32_t cur_offset = topk_length > SWAP_AB_THRESHOLD
71
- ? problem_sizes1[i * 3]
72
- : problem_sizes1[i * 3 + 1];
69
+ int32_t cur_offset =
70
+ swap_ab ? problem_sizes1[i * 3 + 1] : problem_sizes1[i * 3];
71
atomic_buffer[i] = tot_offset;
72
tot_offset += cur_offset;
73
expert_offsets[i + 1] = tot_offset;
@@ -119,15 +117,19 @@ void get_cutlass_moe_mm_data_caller(
117
118
int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());
119
122
- if (topk_ids.numel() > SWAP_AB_THRESHOLD) {
123
- compute_problem_sizes<false><<<num_experts, num_threads, 0, stream>>>(
120
+ // Swap-AB should be disabled for FP4 path
121
+ bool may_swap_ab = (!blockscale_offsets.has_value()) &&
122
+ (topk_ids.numel() <= SWAP_AB_THRESHOLD);
123
+
124
+ if (may_swap_ab) {
125
+ compute_problem_sizes<true><<<num_experts, num_threads, 0, stream>>>(
126
static_cast<const int32_t*>(topk_ids.data_ptr()),
127
static_cast<int32_t*>(problem_sizes1.data_ptr()),
128
static_cast<int32_t*>(problem_sizes2.data_ptr()),
129
static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(), n,
130
k);
131
} else {
130
- compute_problem_sizes<true><<<num_experts, num_threads, 0, stream>>>(
132
+ compute_problem_sizes<false><<<num_experts, num_threads, 0, stream>>>(
133
static_cast<const int32_t*>(topk_ids.data_ptr()),
134
static_cast<int32_t*>(problem_sizes1.data_ptr()),
135
static_cast<int32_t*>(problem_sizes2.data_ptr()),
@@ -136,18 +138,19 @@ void get_cutlass_moe_mm_data_caller(
138
}
139
140
if (blockscale_offsets.has_value()) {
141
+ // fp4 path
142
compute_expert_blockscale_offsets<<<1, 1, 0, stream>>>(
143
static_cast<const int32_t*>(problem_sizes1.data_ptr()),
144
static_cast<int32_t*>(expert_offsets.data_ptr()),
145
static_cast<int32_t*>(blockscale_offsets.value().data_ptr()),
146
static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts,
144
- topk_ids.numel());
147
+ may_swap_ab);
148
} else {
149
compute_expert_offsets<<<1, 1, 0, stream>>>(
150
static_cast<const int32_t*>(problem_sizes1.data_ptr()),
151
static_cast<int32_t*>(expert_offsets.data_ptr()),
152
static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts,
150
- topk_ids.numel());
153
+ may_swap_ab);
154
}
155
compute_arg_sorts<<<num_experts, num_threads, 0, stream>>>(
156
static_cast<const int32_t*>(topk_ids.data_ptr()),