@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9fb2d22032

Fork 0

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

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

Synced Nov 15, 01:00 UTC

[Performance] Performance improvements in non-blockwise fp8 CUTLASS MoE (#20762)

Signed-off-by: ElizaWszola <ewszola@redhat.com>

ElizaWszola committed Jul 17, 2025 at 15:56 UTC 9fb2d22032cee577a189f8c4cddd88a3c190cb0c

6 files changed +174 -38

benchmarks/kernels/benchmark_grouped_gemm_cutlass.py

+34 -1

@@ -80,6 +80,11 @@ def bench_run(

80 a, score, topk, renormalize=False

81 )

83 + ab_strides1 = torch.full((num_experts,), k, device="cuda", dtype=torch.int64)

84 + ab_strides2 = torch.full((num_experts,), n, device="cuda", dtype=torch.int64)

85 + c_strides1 = torch.full((num_experts,), 2 * n, device="cuda", dtype=torch.int64)

86 + c_strides2 = torch.full((num_experts,), k, device="cuda", dtype=torch.int64)

87 +

88 def run_triton_moe(

89 a: torch.Tensor,

90 w1: torch.Tensor,

@@ -111,6 +116,10 @@ def bench_run(

116 w2: torch.Tensor,

117 w1_scale: torch.Tensor,

118 w2_scale: torch.Tensor,

119 + ab_strides1: torch.Tensor,

120 + ab_strides2: torch.Tensor,

121 + c_strides1: torch.Tensor,

122 + c_strides2: torch.Tensor,

123 topk_weights: torch.Tensor,

124 topk_ids: torch.Tensor,

125 per_act_token: bool,

@@ -125,6 +134,10 @@ def bench_run(

134 topk_ids,

135 w1_scale,

136 w2_scale,

137 + ab_strides1,

138 + ab_strides2,

139 + c_strides1,

140 + c_strides2,

141 per_act_token,

142 a1_scale=None,

143 )

@@ -136,6 +149,10 @@ def bench_run(

149 w2_q: torch.Tensor,

150 w1_scale: torch.Tensor,

151 w2_scale: torch.Tensor,

152 + ab_strides1: torch.Tensor,

153 + ab_strides2: torch.Tensor,

154 + c_strides1: torch.Tensor,

155 + c_strides2: torch.Tensor,

156 topk_weights: torch.Tensor,

157 topk_ids: torch.Tensor,

158 ):

@@ -150,6 +167,10 @@ def bench_run(

167 topk_ids,

168 w1_scale,

169 w2_scale,

170 + ab_strides1,

171 + ab_strides2,

172 + c_strides1,

173 + c_strides2,

174 per_act_token,

175 a1_scale=None,

176 )

@@ -194,6 +215,10 @@ def bench_run(

215 w2_q,

216 w1_scale,

217 w2_scale,

218 + ab_strides1,

219 + ab_strides2,

220 + c_strides1,

221 + c_strides2,

222 topk_weights,

223 topk_ids,

224 )

@@ -231,6 +256,10 @@ def bench_run(

256 "w1_scale": w1_scale,

257 "w2_scale": w2_scale,

258 "per_act_token": per_act_token,

259 + "ab_strides1": ab_strides1,

260 + "ab_strides2": ab_strides2,

261 + "c_strides1": c_strides1,

262 + "c_strides2": c_strides2,

263 # cuda graph params

264 "cutlass_graph": cutlass_graph,

265 "triton_graph": triton_graph,

@@ -289,6 +318,10 @@ def bench_run(

318 w2_q,

319 w1_scale,

320 w2_scale,

321 + ab_strides1,

322 + ab_strides2,

323 + c_strides1,

324 + c_strides2,

325 topk_weights,

326 topk_ids,

327 per_act_token,

@@ -297,7 +330,7 @@ def bench_run(

330

331 results.append(

332 benchmark.Timer(

300 - stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, per_act_token, num_runs)", # noqa: E501

333 + stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, ab_strides1, ab_strides2, c_strides1, c_strides2, topk_weights, topk_ids, per_act_token, num_runs)", # noqa: E501

334 globals=globals,

335 label=label,

336 sub_label=sub_label,

csrc/moe/moe_permute_unpermute_op.cu

+42 -11

@@ -160,6 +160,30 @@ __global__ void shuffleInputRowsKernel(const T* input,

160 }

161 }

162

163 +template <typename T>

164 +__global__ void shuffleInputRowsKernelSlow(const T* input,

165 + const int32_t* dst2src_map,

166 + T* output, int64_t num_src_rows,

167 + int64_t num_dst_rows,

168 + int64_t num_cols) {

169 + int64_t dest_row_idx = blockIdx.x;

170 + int64_t const source_row_idx = dst2src_map[dest_row_idx];

171 +

172 + if (blockIdx.x < num_dst_rows) {

173 + // Duplicate and permute rows

174 + auto const* source_row_ptr = input + source_row_idx * num_cols;

175 + auto* dest_row_ptr = output + dest_row_idx * num_cols;

176 +

177 + int64_t const start_offset = threadIdx.x;

178 + int64_t const stride = blockDim.x;

179 +

180 + for (int elem_index = start_offset; elem_index < num_cols;

181 + elem_index += stride) {

182 + dest_row_ptr[elem_index] = source_row_ptr[elem_index];

183 + }

184 + }

185 +}

186 +

187 void shuffle_rows(const torch::Tensor& input_tensor,

188 const torch::Tensor& dst2src_map,

189 torch::Tensor& output_tensor) {

@@ -173,17 +197,24 @@ void shuffle_rows(const torch::Tensor& input_tensor,

197 int64_t const num_src_rows = input_tensor.size(0);

198 int64_t const num_cols = input_tensor.size(1);

199

176 - TORCH_CHECK(!(num_cols % (128 / sizeof(input_tensor.scalar_type()) / 8)),

177 - "num_cols must be divisible by 128 / "

178 - "sizeof(input_tensor.scalar_type()) / 8");

179 -

180 - MOE_DISPATCH(input_tensor.scalar_type(), [&] {

181 - shuffleInputRowsKernel<scalar_t><<<blocks, threads, 0, stream>>>(

182 - reinterpret_cast<scalar_t*>(input_tensor.data_ptr()),

183 - dst2src_map.data_ptr<int32_t>(),

184 - reinterpret_cast<scalar_t*>(output_tensor.data_ptr()), num_src_rows,

185 - num_dest_rows, num_cols);

186 - });

200 + if (num_cols % (128 / sizeof(input_tensor.scalar_type()) / 8)) {

201 + // use slow kernel if num_cols can't be aligned to 128 bits

202 + MOE_DISPATCH(input_tensor.scalar_type(), [&] {

203 + shuffleInputRowsKernelSlow<scalar_t><<<blocks, threads, 0, stream>>>(

204 + reinterpret_cast<scalar_t*>(input_tensor.data_ptr()),

205 + dst2src_map.data_ptr<int32_t>(),

206 + reinterpret_cast<scalar_t*>(output_tensor.data_ptr()), num_src_rows,

207 + num_dest_rows, num_cols);

208 + });

209 + } else {

210 + MOE_DISPATCH(input_tensor.scalar_type(), [&] {

211 + shuffleInputRowsKernel<scalar_t><<<blocks, threads, 0, stream>>>(

212 + reinterpret_cast<scalar_t*>(input_tensor.data_ptr()),

213 + dst2src_map.data_ptr<int32_t>(),

214 + reinterpret_cast<scalar_t*>(output_tensor.data_ptr()), num_src_rows,

215 + num_dest_rows, num_cols);

216 + });

217 + }

218 }

219

220 #else

tests/kernels/moe/test_cutlass_moe.py

+12 -2

@@ -206,6 +206,10 @@ def run_8_bit(moe_tensors: MOETensors8Bit,

206 'topk_ids': topk_ids,

207 'w1_scale': moe_tensors.w1_scale,

208 'w2_scale': moe_tensors.w2_scale,

209 + 'ab_strides1': moe_tensors.ab_strides1,

210 + 'ab_strides2': moe_tensors.ab_strides2,

211 + 'c_strides1': moe_tensors.c_strides1,

212 + 'c_strides2': moe_tensors.c_strides2,

213 'per_act_token': per_act_token,

214 'a1_scale': None #moe_tensors.a_scale

215 }

@@ -439,6 +443,11 @@ def test_run_cutlass_moe_fp8(

443 expert_map[start:end] = list(range(num_local_experts))

444 expert_map = torch.tensor(expert_map, dtype=torch.int32, device="cuda")

445

446 + ab_strides1 = torch.full((e, ), k, device="cuda", dtype=torch.int64)

447 + ab_strides2 = torch.full((e, ), n, device="cuda", dtype=torch.int64)

448 + c_strides1 = torch.full((e, ), 2 * n, device="cuda", dtype=torch.int64)

449 + c_strides2 = torch.full((e, ), k, device="cuda", dtype=torch.int64)

450 +

451 activation = lambda o, i: torch.ops._C.silu_and_mul(o, i)

452 a1q, a1q_scale = moe_kernel_quantize_input(mt.a, mt.a_scale,

453 torch.float8_e4m3fn,

@@ -447,8 +456,9 @@ def test_run_cutlass_moe_fp8(

456 func = lambda output: run_cutlass_moe_fp8(

457 output, a1q, mt.w1_q, mt.w2_q, topk_ids, activation,

458 global_num_experts, expert_map, mt.w1_scale, mt.w2_scale,

450 - a1q_scale, None, workspace13, workspace2, None, mt.a.dtype,

451 - per_act_token, per_out_channel, False)

459 + a1q_scale, None, ab_strides1, ab_strides2, c_strides1, c_strides2,

460 + workspace13, workspace2, None, mt.a.dtype, per_act_token,

461 + per_out_channel, False)

462

463 workspace13.random_()

464 output_random_workspace = torch.empty(output_shape,

tests/kernels/moe/test_pplx_cutlass_moe.py

+22

@@ -75,6 +75,7 @@ def pplx_cutlass_moe(

75 assert torch.cuda.current_device() == pgi.local_rank

77 num_tokens, hidden_dim = a.shape

78 + intermediate_dim = w2.shape[2]

79 num_experts = w1.shape[0]

80 block_size = hidden_dim # TODO support more cases

81 device = pgi.device

@@ -123,10 +124,31 @@ def pplx_cutlass_moe(

124 num_local_experts=num_local_experts,

125 num_dispatchers=num_dispatchers)

126

127 + ab_strides1 = torch.full((num_local_experts, ),

128 + hidden_dim,

129 + device="cuda",

130 + dtype=torch.int64)

131 + ab_strides2 = torch.full((num_local_experts, ),

132 + intermediate_dim,

133 + device="cuda",

134 + dtype=torch.int64)

135 + c_strides1 = torch.full((num_local_experts, ),

136 + 2 * intermediate_dim,

137 + device="cuda",

138 + dtype=torch.int64)

139 + c_strides2 = torch.full((num_local_experts, ),

140 + hidden_dim,

141 + device="cuda",

142 + dtype=torch.int64)

143 +

144 experts = CutlassExpertsFp8(num_local_experts,

145 out_dtype,

146 per_act_token,

147 per_out_ch,

148 + ab_strides1,

149 + ab_strides2,

150 + c_strides1,

151 + c_strides2,

152 num_dispatchers=num_dispatchers,

153 use_batched_format=True)

154

vllm/model_executor/layers/fused_moe/cutlass_moe.py

+39 -23

@@ -13,8 +13,7 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import (

13 MoEPrepareAndFinalizeNoEP)

14 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

15 TopKWeightAndReduceDelegate)

16 -from vllm.model_executor.layers.fused_moe.utils import (_fp8_perm,

17 - _fp8_quantize,

16 +from vllm.model_executor.layers.fused_moe.utils import (_fp8_quantize,

17 _resize_cache)

18 from vllm.scalar_type import scalar_types

@@ -34,6 +33,10 @@ def run_cutlass_moe_fp8(

33 w2_scale: Optional[torch.Tensor],

34 a1q_scale: Optional[torch.Tensor],

35 a2_scale: Optional[torch.Tensor],

36 + ab_strides1: torch.Tensor,

37 + ab_strides2: torch.Tensor,

38 + c_strides1: torch.Tensor,

39 + c_strides2: torch.Tensor,

40 workspace13: torch.Tensor,

41 workspace2: torch.Tensor,

42 expert_num_tokens: Optional[torch.Tensor],

@@ -152,27 +155,11 @@ def run_cutlass_moe_fp8(

155 problem_sizes1, problem_sizes2, a_map,

156 c_map, global_num_experts, N, K)

157

155 - a1q = _fp8_perm(a1q, a_map)

156 - a1q_scale = a1q_scale[a_map] if per_act_token else a1q_scale

158 + a1q = ops.shuffle_rows(a1q, a_map)

159 + a1q_scale = (ops.shuffle_rows(a1q_scale, a_map)

160 + if per_act_token else a1q_scale)

161 expert_offsets = expert_offsets[:-1]

162

159 - ab_strides1 = torch.full((w1.size(0), ),

160 - K,

161 - device=device,

162 - dtype=torch.int64)

163 - c_strides1 = torch.full((w1.size(0), ),

164 - 2 * N,

165 - device=device,

166 - dtype=torch.int64)

167 - ab_strides2 = torch.full((w1.size(0), ),

168 - N,

169 - device=device,

170 - dtype=torch.int64)

171 - c_strides2 = torch.full((w1.size(0), ),

172 - K,

173 - device=device,

174 - dtype=torch.int64)

175 -

163 if use_batched_format:

164 c1 = _resize_cache(workspace13, (local_E * padded_M, N * 2))

165 c2 = _resize_cache(workspace2, (local_E * padded_M, N))

@@ -209,7 +196,8 @@ def run_cutlass_moe_fp8(

196 else:

197 # We can't do this inplace because output may point to the same tensor

198 # as c3.

212 - output.copy_(c3[c_map].view(M * topk, K), non_blocking=True)

199 + output.copy_(ops.shuffle_rows(c3, c_map).view(M * topk, K),

200 + non_blocking=True)

201

202

203 # TODO (bnell): split class batched vs. non-batched?

@@ -222,6 +210,10 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):

210 out_dtype: Optional[torch.dtype],

211 per_act_token_quant: bool,

212 per_out_ch_quant: bool,

213 + ab_strides1: torch.Tensor,

214 + ab_strides2: torch.Tensor,

215 + c_strides1: torch.Tensor,

216 + c_strides2: torch.Tensor,

217 block_shape: Optional[list[int]] = None,

218 num_dispatchers: Optional[int] = None,

219 use_batched_format: bool = False,

@@ -238,6 +230,10 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):

230 self.max_experts_per_worker = max_experts_per_worker

231 self.num_dispatchers = num_dispatchers

232 self.out_dtype = out_dtype

233 + self.ab_strides1 = ab_strides1

234 + self.ab_strides2 = ab_strides2

235 + self.c_strides1 = c_strides1

236 + self.c_strides2 = c_strides2

237 self.use_batched_format = use_batched_format

238

239 @property

@@ -316,7 +312,8 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):

312 run_cutlass_moe_fp8(

313 output, hidden_states, w1, w2, topk_ids, activation_callable,

314 global_num_experts, expert_map, w1_scale, w2_scale, a1q_scale,

319 - a2_scale, workspace13, workspace2, expert_num_tokens,

315 + a2_scale, self.ab_strides1, self.ab_strides2, self.c_strides1,

316 + self.c_strides2, workspace13, workspace2, expert_num_tokens,

317 self.out_dtype if self.out_dtype is not None else in_dtype,

318 self.per_act_token_quant, self.per_out_ch_quant,

319 self.use_batched_format)

@@ -330,6 +327,10 @@ def cutlass_moe_fp8(

327 topk_ids: torch.Tensor,

328 w1_scale: torch.Tensor,

329 w2_scale: torch.Tensor,

330 + ab_strides1: torch.Tensor,

331 + ab_strides2: torch.Tensor,

332 + c_strides1: torch.Tensor,

333 + c_strides2: torch.Tensor,

334 per_act_token: Optional[bool] = None,

335 activation: str = "silu",

336 a1_scale: Optional[torch.Tensor] = None,

@@ -357,6 +358,17 @@ def cutlass_moe_fp8(

358 Shape: [num_experts] or [num_experts, 2N]

359 - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.

360 Shape: [num_experts] or [num_experts, K]

361 + - ab_strides1 (torch.Tensor): The input/weight strides for the first gemm.

362 + Shape: [num_experts]

363 + - ab_strides2 (torch.Tensor): The input/weight strides for the second gemm.

364 + Shape: [num_experts]

365 + - c_strides1 (torch.Tensor): The output strides for the first gemm.

366 + Shape: [num_experts]

367 + - c_strides2 (torch.Tensor): The output strides for the second gemm.

368 + Shape: [num_experts]

369 + - per_act_token (Optional[bool]): Whether the scale is per-token or

370 + per-tensor.

371 + - activation (str): The activation function to use.

372 - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.

373 Shape: scalar or [M]

374 - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to

@@ -389,6 +401,10 @@ def cutlass_moe_fp8(

401 out_dtype=a.dtype,

402 per_act_token_quant=per_act_token,

403 per_out_ch_quant=per_out_ch,

404 + ab_strides1=ab_strides1,

405 + ab_strides2=ab_strides2,

406 + c_strides1=c_strides1,

407 + c_strides2=c_strides2,

408 use_batched_format=False,

409 ),

410 )

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

+25 -1

@@ -859,6 +859,21 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):

859 layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,

860 requires_grad=False)

861

862 + device = layer.w13_weight.device

863 + # ab_strides1 and c_strides2 are the same

864 + self.ab_strides1_c_strides2 = torch.full((layer.local_num_experts, ),

865 + layer.hidden_size,

866 + device=device,

867 + dtype=torch.int64)

868 + self.ab_strides2 = torch.full((layer.local_num_experts, ),

869 + layer.intermediate_size_per_partition,

870 + device=device,

871 + dtype=torch.int64)

872 + self.c_strides1 = torch.full((layer.local_num_experts, ),

873 + 2 * layer.intermediate_size_per_partition,

874 + device=device,

875 + dtype=torch.int64)

876 +

877 def select_gemm_impl(

878 self,

879 prepare_finalize: FusedMoEPrepareAndFinalize,

@@ -881,6 +896,10 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):

896 moe.in_dtype,

897 self.input_quant.strategy == QuantizationStrategy.TOKEN,

898 self.weight_quant.strategy == QuantizationStrategy.CHANNEL,

899 + ab_strides1=self.ab_strides1_c_strides2,

900 + ab_strides2=self.ab_strides2,

901 + c_strides1=self.c_strides1,

902 + c_strides2=self.ab_strides1_c_strides2,

903 num_dispatchers=num_dispatchers,

904 use_batched_format=use_batched_format,

905 )

@@ -927,7 +946,8 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):

946 num_expert_group=num_expert_group,

947 custom_routing_function=custom_routing_function,

948 scoring_func=scoring_func,

930 - e_score_correction_bias=e_score_correction_bias)

949 + e_score_correction_bias=e_score_correction_bias,

950 + indices_type=self.topk_indices_dtype)

951

952 per_act_token = (

953 self.input_quant.strategy == QuantizationStrategy.TOKEN)

@@ -948,6 +968,10 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):

968 expert_map=None if self.disable_expert_map else expert_map,

969 w1_scale=layer.w13_weight_scale,

970 w2_scale=layer.w2_weight_scale,

971 + ab_strides1=self.ab_strides1_c_strides2,

972 + ab_strides2=self.ab_strides2,

973 + c_strides1=self.c_strides1,

974 + c_strides2=self.ab_strides1_c_strides2,

975 a1_scale=layer.w13_input_scale,

976 a2_scale=layer.w2_input_scale,

977 )