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 15, 03:30 UTC
[Bugfix] Fix topk_ids indices_type for CUTLASS w8a8 FP8 MoE (#20166)
Signed-off-by: Ming Yang <yming@meta.com>
Ming Yang committed
Jul 8, 2025 at 16:10 UTC
c438183e99ae8a7a85f3de0180c762c197c1dd76
3 files changed
+17
-11
csrc/quantization/cutlass_w8a8/moe/moe_data.cu
+4
-4
@@ -7,7 +7,7 @@
7
8
constexpr uint64_t THREADS_PER_EXPERT = 512;
9
10
-__global__ void compute_problem_sizes(const uint32_t* __restrict__ topk_ids,
10
+__global__ void compute_problem_sizes(const int32_t* __restrict__ topk_ids,
11
int32_t* problem_sizes1,
12
int32_t* problem_sizes2,
13
int32_t* atomic_buffer,
@@ -62,7 +62,7 @@ __global__ void compute_expert_blockscale_offsets(
62
}
63
}
64
65
-__global__ void compute_arg_sorts(const uint32_t* __restrict__ topk_ids,
65
+__global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids,
66
const int32_t* __restrict__ expert_offsets,
67
int32_t* input_permutation,
68
int32_t* output_permutation,
@@ -103,7 +103,7 @@ void get_cutlass_moe_mm_data_caller(
103
104
int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());
105
compute_problem_sizes<<<num_experts, num_threads, 0, stream>>>(
106
- static_cast<const uint32_t*>(topk_ids.data_ptr()),
106
+ static_cast<const int32_t*>(topk_ids.data_ptr()),
107
static_cast<int32_t*>(problem_sizes1.data_ptr()),
108
static_cast<int32_t*>(problem_sizes2.data_ptr()),
109
static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(), n, k);
@@ -120,7 +120,7 @@ void get_cutlass_moe_mm_data_caller(
120
static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts);
121
}
122
compute_arg_sorts<<<num_experts, num_threads, 0, stream>>>(
123
- static_cast<const uint32_t*>(topk_ids.data_ptr()),
123
+ static_cast<const int32_t*>(topk_ids.data_ptr()),
124
static_cast<const int32_t*>(expert_offsets.data_ptr()),
125
static_cast<int32_t*>(input_permutation.data_ptr()),
126
static_cast<int32_t*>(output_permutation.data_ptr()),
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
+4
-2
@@ -78,7 +78,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
78
return self.max_num_tokens
79
80
def topk_indices_dtype(self) -> Optional[torch.dtype]:
81
- return torch.uint32
81
+ return torch.int32
82
83
def num_dispatchers(self) -> int:
84
return self.num_dispatchers_
@@ -100,7 +100,9 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
100
hidden_dim = a1.size(-1) # K
101
102
assert topk_ids.size(0) == num_tokens
103
- # assert expert_map is None, "NYI"
103
+ assert expert_map is None, """with expert map, -1 id is used for
104
+ non-local token; this causes error when casting ids to the
105
+ topk_indices_dtype() uint32"""
106
107
# Is this always going to be a1.device?
108
device = a1.device
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+9
-5
@@ -929,9 +929,12 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
929
num_expert_group=num_expert_group,
930
custom_routing_function=custom_routing_function,
931
scoring_func=scoring_func,
932
- e_score_correction_bias=e_score_correction_bias,
933
- indices_type=self.topk_indices_dtype,
934
- )
932
+ e_score_correction_bias=e_score_correction_bias)
933
+
934
+ a1_scale = layer.w13_input_scale
935
+ a2_scale = layer.w2_input_scale
936
+ per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
937
+ a2_scale.numel() != 1 if a2_scale is not None else False)
938
939
return self.fused_experts(
940
x,
@@ -939,13 +942,14 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
942
layer.w2_weight,
943
topk_weights,
944
topk_ids,
945
+ per_act_token=per_act_token,
946
activation=activation,
947
global_num_experts=global_num_experts,
948
expert_map=None if self.disable_expert_map else expert_map,
949
w1_scale=layer.w13_weight_scale,
950
w2_scale=layer.w2_weight_scale,
947
- a1_scale=layer.w13_input_scale,
948
- a2_scale=layer.w2_input_scale,
951
+ a1_scale=a1_scale,
952
+ a2_scale=a2_scale,
953
)
954
955