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

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 14, 05:30 UTC

[Kernel] CUTLASS MoE FP8: Integrate cuda moe permute/unpermute (#23045)

Signed-off-by: Shixian Cui <shixian@amazon.com>

shixianc committed Aug 20, 2025 at 07:35 UTC b17109beeafbf9577c319ab61530810943a7fc4b

15 files changed +369 -121

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

+14 -19

@@ -45,8 +45,6 @@ void moe_permute(

45 auto copy_topk_ids = topk_ids.clone(); // copy topk_ids for preprocess

46 auto permuted_experts_id = torch::empty_like(topk_ids);

47 auto sorted_row_idx = torch::empty_like(inv_permuted_idx);

48 - auto align_expert_first_token_offset =

49 - torch::zeros_like(expert_first_token_offset);

49 CubKeyValueSorter sorter{};

50 int64_t* valid_num_ptr = nullptr;

@@ -85,12 +83,14 @@ void moe_permute(

83 });

85 // get m_indices and update expert_first_token_offset with align block

88 - getMIndices(get_ptr<int64_t>(expert_first_token_offset),

89 - get_ptr<int64_t>(align_expert_first_token_offset),

90 - get_ptr<int>(m_indices), n_local_expert, align_block_size_value,

91 - stream);

86 + // this is only required for DeepGemm and not required for CUTLASS group gemm

87 if (align_block_size.has_value()) {

93 - // update align_expert_first_token_offset

88 + auto align_expert_first_token_offset =

89 + torch::zeros_like(expert_first_token_offset);

90 + getMIndices(get_ptr<int64_t>(expert_first_token_offset),

91 + get_ptr<int64_t>(align_expert_first_token_offset),

92 + get_ptr<int>(m_indices), n_local_expert, align_block_size_value,

93 + stream);

94 expert_first_token_offset.copy_(align_expert_first_token_offset);

95 }

96 }

@@ -195,19 +195,14 @@ void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_weights,

195 torch::Tensor& expert_first_token_offset,

196 torch::Tensor& src_row_id2dst_row_id_map,

197 torch::Tensor& m_indices) {

198 - TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0");

198 + TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0");

199 }

200

201 -void moe_unpermute(const torch::Tensor& input,

202 - const torch::Tensor& topk_weights, torch::Tensor& topk_ids,

203 - const torch::Tensor& token_expert_indices,

204 - const std::optional<torch::Tensor>& expert_map,

205 - int64_t n_expert, int64_t n_local_expert, int64_t topk,

206 - const std::optional<int64_t>& align_block_size,

207 - torch::Tensor& permuted_input,

208 - torch::Tensor& expert_first_token_offset,

209 - torch::Tensor& src_row_id2dst_row_id_map,

210 - torch::Tensor& m_indices) {

201 +void moe_unpermute(

202 + const torch::Tensor& permuted_hidden_states,

203 + const torch::Tensor& topk_weights, const torch::Tensor& inv_permuted_idx,

204 + const std::optional<torch::Tensor>& expert_first_token_offset, int64_t topk,

205 + torch::Tensor& hidden_states) {

206 TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0");

207 }

208

@@ -224,4 +219,4 @@ bool moe_permute_unpermute_supported() {

219 TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {

220 m.impl("moe_permute", &moe_permute);

221 m.impl("moe_unpermute", &moe_unpermute);

227 -}

222 +}

\ No newline at end of file

csrc/ops.h

@@ -229,6 +229,11 @@ void get_cutlass_moe_mm_data(

229 const int64_t num_experts, const int64_t n, const int64_t k,

230 const std::optional<torch::Tensor>& blockscale_offsets);

231

232 +void get_cutlass_moe_mm_problem_sizes(

233 + const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,

234 + torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,

235 + const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets);

236 +

237 void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,

238 torch::Tensor& problem_sizes1,

239 torch::Tensor& problem_sizes2,

csrc/quantization/cutlass_w8a8/moe/get_group_starts.cuh

+4 -2

@@ -10,7 +10,7 @@

11 template <typename ElementAB, typename ElementC, typename ElementAccumulator>

12 __global__ void get_group_gemm_starts(

13 - int32_t* expert_offsets, ElementAB** a_offsets, ElementAB** b_offsets,

13 + int64_t* expert_offsets, ElementAB** a_offsets, ElementAB** b_offsets,

14 ElementC** out_offsets, ElementAccumulator** a_scales_offsets,

15 ElementAccumulator** b_scales_offsets, ElementAB* a_base_as_int,

16 ElementAB* b_base_as_int, ElementC* out_base_as_int,

@@ -34,7 +34,7 @@ __global__ void get_group_gemm_starts(

34 else if (out_tensors.dtype() == TENSOR_C_TYPE) { \

35 get_group_gemm_starts<cutlass::float_e4m3_t, C_TYPE, float> \

36 <<<1, num_experts, 0, stream>>>( \

37 - static_cast<int32_t*>(expert_offsets.data_ptr()), \

37 + static_cast<int64_t*>(expert_offsets.data_ptr()), \

38 static_cast<cutlass::float_e4m3_t**>(a_ptrs.data_ptr()), \

39 static_cast<cutlass::float_e4m3_t**>(b_ptrs.data_ptr()), \

40 static_cast<C_TYPE**>(out_ptrs.data_ptr()), \

@@ -61,6 +61,8 @@ void run_get_group_gemm_starts(

61 TORCH_CHECK(b_tensors.dtype() == torch::kFloat8_e4m3fn);

62 TORCH_CHECK(a_scales.dtype() == torch::kFloat32);

63 TORCH_CHECK(b_scales.dtype() == torch::kFloat32);

64 + // expect int64_t to avoid overflow during offset calculations

65 + TORCH_CHECK(expert_offsets.dtype() == torch::kInt64);

67 int num_experts = static_cast<int>(expert_offsets.size(0));

68 bool per_act_token = a_scales.numel() != 1;

csrc/quantization/cutlass_w8a8/moe/moe_data.cu

+50 -15

@@ -104,6 +104,53 @@ __global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids,

104 }

105 }

106

107 +namespace {

108 +inline void launch_compute_problem_sizes(const torch::Tensor& topk_ids,

109 + torch::Tensor& problem_sizes1,

110 + torch::Tensor& problem_sizes2,

111 + torch::Tensor& atomic_buffer,

112 + int64_t num_experts, int64_t n,

113 + int64_t k, cudaStream_t stream,

114 + const bool swap_ab) {

115 + int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());

116 +

117 + const int32_t* topk_ptr = static_cast<const int32_t*>(topk_ids.data_ptr());

118 + int32_t* ps1_ptr = static_cast<int32_t*>(problem_sizes1.data_ptr());

119 + int32_t* ps2_ptr = static_cast<int32_t*>(problem_sizes2.data_ptr());

120 + int32_t* atomic_ptr = static_cast<int32_t*>(atomic_buffer.data_ptr());

121 +

122 + if (swap_ab) {

123 + compute_problem_sizes<true><<<num_experts, num_threads, 0, stream>>>(

124 + topk_ptr, ps1_ptr, ps2_ptr, atomic_ptr,

125 + static_cast<int>(topk_ids.numel()), static_cast<int>(n),

126 + static_cast<int>(k));

127 + } else {

128 + compute_problem_sizes<false><<<num_experts, num_threads, 0, stream>>>(

129 + topk_ptr, ps1_ptr, ps2_ptr, atomic_ptr,

130 + static_cast<int>(topk_ids.numel()), static_cast<int>(n),

131 + static_cast<int>(k));

132 + }

133 +}

134 +} // namespace

135 +

136 +void get_cutlass_moe_mm_problem_sizes_caller(

137 + const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,

138 + torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,

139 + const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets) {

140 + auto stream = at::cuda::getCurrentCUDAStream(topk_ids.device().index());

141 + auto options_int32 =

142 + torch::TensorOptions().dtype(torch::kInt32).device(topk_ids.device());

143 + torch::Tensor atomic_buffer = torch::zeros(num_experts, options_int32);

144 +

145 + // Swap-AB should be disabled for FP4 path

146 + bool may_swap_ab = (!blockscale_offsets.has_value()) &&

147 + (topk_ids.numel() <= SWAP_AB_THRESHOLD);

148 +

149 + launch_compute_problem_sizes(topk_ids, problem_sizes1, problem_sizes2,

150 + atomic_buffer, num_experts, n, k, stream,

151 + may_swap_ab);

152 +}

153 +

154 void get_cutlass_moe_mm_data_caller(

155 const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,

156 torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,

@@ -121,21 +168,9 @@ void get_cutlass_moe_mm_data_caller(

168 bool may_swap_ab = (!blockscale_offsets.has_value()) &&

169 (topk_ids.numel() <= SWAP_AB_THRESHOLD);

170

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 {

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 - static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(), n,

137 - k);

138 - }

171 + launch_compute_problem_sizes(topk_ids, problem_sizes1, problem_sizes2,

172 + atomic_buffer, num_experts, n, k, stream,

173 + may_swap_ab);

174

175 if (blockscale_offsets.has_value()) {

176 // fp4 path

csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu

+24

@@ -76,6 +76,11 @@ void get_cutlass_moe_mm_data_caller(

76 const int64_t num_experts, const int64_t n, const int64_t k,

77 const std::optional<torch::Tensor>& blockscale_offsets);

79 +void get_cutlass_moe_mm_problem_sizes_caller(

80 + const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,

81 + torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,

82 + const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets);

83 +

84 void get_cutlass_pplx_moe_mm_data_caller(torch::Tensor& expert_offsets,

85 torch::Tensor& problem_sizes1,

86 torch::Tensor& problem_sizes2,

@@ -293,6 +298,25 @@ void get_cutlass_moe_mm_data(

298 version_num, ". Required capability: 90 or 100");

299 }

300

301 +void get_cutlass_moe_mm_problem_sizes(

302 + const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,

303 + torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,

304 + const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets) {

305 + int32_t version_num = get_sm_version_num();

306 +#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \

307 + (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100)

308 + get_cutlass_moe_mm_problem_sizes_caller(topk_ids, problem_sizes1,

309 + problem_sizes2, num_experts, n, k,

310 + blockscale_offsets);

311 + return;

312 +#endif

313 + TORCH_CHECK_NOT_IMPLEMENTED(

314 + false,

315 + "No compiled get_cutlass_moe_mm_problem_sizes: no cutlass_scaled_mm "

316 + "kernel for CUDA device capability: ",

317 + version_num, ". Required capability: 90 or 100");

318 +}

319 +

320 void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,

321 torch::Tensor& problem_sizes1,

322 torch::Tensor& problem_sizes2,

csrc/torch_bindings.cpp

+13

@@ -440,6 +440,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

440 {stride_tag});

441 ops.impl("get_cutlass_moe_mm_data", torch::kCUDA, &get_cutlass_moe_mm_data);

442

443 + // A function that computes problem sizes for each expert's multiplication

444 + // used by the two mms called from fused MoE operation. It takes topk_ids as

445 + // an input, and computes problem_sizes1 and problem_sizes2 only.

446 + ops.def(

447 + "get_cutlass_moe_mm_problem_sizes(Tensor topk_ids, "

448 + " Tensor! problem_sizes1, "

449 + " Tensor! problem_sizes2, "

450 + " int num_experts, int n, int k, "

451 + " Tensor? blockscale_offsets) -> ()",

452 + {stride_tag});

453 + ops.impl("get_cutlass_moe_mm_problem_sizes", torch::kCUDA,

454 + &get_cutlass_moe_mm_problem_sizes);

455 +

456 // A function that computes data required to run fused MoE with w8a8 grouped

457 // GEMM and PPLX. It takes expert_num_tokens and non_zero_expert_idxs

458 // as an input, and computes expert_offsets (token start indices of each

tests/kernels/moe/test_cutlass_moe.py

+14 -4

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

207 'topk_ids': topk_ids,

208 'w1_scale': moe_tensors.w1_scale,

209 'w2_scale': moe_tensors.w2_scale,

210 + 'ab_strides1': moe_tensors.ab_strides1,

211 + 'ab_strides2': moe_tensors.ab_strides2,

212 + 'c_strides1': moe_tensors.c_strides1,

213 + 'c_strides2': moe_tensors.c_strides2,

214 'per_act_token': per_act_token,

215 'a1_scale': None #moe_tensors.a_scale

216 }

@@ -424,8 +428,8 @@ def test_run_cutlass_moe_fp8(

428 topk_ids[0][1] = 1

429

430 workspace13_shape = (m * topk, max(2 * n, k))

427 - workspace2_shape = (m * topk, n)

428 - output_shape = (m * topk, k)

431 + workspace2_shape = (m * topk, max(n, k))

432 + output_shape = (m, k)

433

434 workspace13 = torch.empty(prod(workspace13_shape),

435 device="cuda",

@@ -440,6 +444,11 @@ def test_run_cutlass_moe_fp8(

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

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

446

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

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

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

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

451 +

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

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

454 torch.float8_e4m3fn,

@@ -448,8 +457,9 @@ def test_run_cutlass_moe_fp8(

457 func = lambda output: run_cutlass_moe_fp8(

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

459 global_num_experts, expert_map, mt.w1_scale, mt.w2_scale,

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

452 - per_act_token, per_out_channel, False)

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

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

462 + per_out_channel, False, topk_weights)

463

464 workspace13.random_()

465 output_random_workspace = torch.empty(output_shape,

tests/kernels/moe/test_moe_permute_unpermute.py

+5 -1

@@ -238,7 +238,11 @@ def test_moe_permute_unpermute(n_token: int, n_hidden: int, topk: int,

238 atol=0,

239 rtol=0)

240 # check mindice

241 - torch.testing.assert_close(gold_m_indices, m_indices, atol=0, rtol=0)

241 + # current kernel usage assumes deepgemm requires align_block_size

242 + # when it's not provided then we don't compute m_indices (for cutlass)

243 + if align_block_size is not None:

244 + torch.testing.assert_close(gold_m_indices, m_indices, atol=0, rtol=0)

245 +

246 # check permuted_hidden_states, only valid token

247 torch.testing.assert_close(gold_permuted_hidden_states[valid_row_idx],

248 permuted_hidden_states[valid_row_idx],

tests/kernels/moe/test_pplx_cutlass_moe.py

+21 -1

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

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

78 num_tokens, hidden_dim = a.shape

79 + intermediate_dim = w2.shape[2]

80 num_experts = w1.shape[0]

81 block_size = hidden_dim # TODO support more cases

82 device = pgi.device

@@ -124,8 +125,27 @@ def pplx_cutlass_moe(

125 num_local_experts=num_local_experts,

126 num_dispatchers=num_dispatchers)

127

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

129 + hidden_dim,

130 + device="cuda",

131 + dtype=torch.int64)

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

133 + intermediate_dim,

134 + device="cuda",

135 + dtype=torch.int64)

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

137 + 2 * intermediate_dim,

138 + device="cuda",

139 + dtype=torch.int64)

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

141 + hidden_dim,

142 + device="cuda",

143 + dtype=torch.int64)

144 +

145 experts = CutlassBatchedExpertsFp8(num_local_experts, num_dispatchers,

128 - out_dtype, per_act_token, per_out_ch)

146 + out_dtype, per_act_token, per_out_ch,

147 + ab_strides1, ab_strides2, c_strides1,

148 + c_strides2)

149

150 fused_cutlass_experts = FusedMoEModularKernel(

151 prepare_finalize,

tests/kernels/quantization/test_cutlass_scaled_mm.py

+1 -1

@@ -535,7 +535,7 @@ def test_cutlass_fp8_group_gemm(num_experts: int, per_act_token: bool,

535

536 expert_offsets = torch.zeros((num_experts + 1),

537 device=device,

538 - dtype=torch.int32)

538 + dtype=torch.int64)

539

540 problem_sizes = torch.zeros((num_experts, 3),

541 device=device,

vllm/_custom_ops.py

+22

@@ -844,6 +844,28 @@ def get_cutlass_moe_mm_data(topk_ids: torch.Tensor,

844 blockscale_offsets)

845

846

847 +def get_cutlass_moe_mm_problem_sizes(

848 + topk_ids: torch.Tensor,

849 + problem_sizes1: torch.Tensor,

850 + problem_sizes2: torch.Tensor,

851 + num_experts: int,

852 + n: int,

853 + k: int,

854 + blockscale_offsets: Optional[torch.Tensor] = None):

855 + """

856 + Compute only the per-expert problem sizes needed by the two grouped matrix

857 + multiplications used in CUTLASS-based fused MoE.

858 +

859 + The function takes in topk_ids (token→expert mapping) and computes:

860 + - problem_sizes1, problem_sizes2: M×N×K sizes of each expert's

861 + multiplication for the two grouped MMs

862 + used in the fused MoE operation.

863 + """

864 + return torch.ops._C.get_cutlass_moe_mm_problem_sizes(

865 + topk_ids, problem_sizes1, problem_sizes2, num_experts, n, k,

866 + blockscale_offsets)

867 +

868 +

869 def shuffle_rows(input_tensor: torch.Tensor, dst2src_map: torch.Tensor):

870 """

871 Shuffle and expand the input tensor according to the dst2src_map and store the result in output_tensor.

vllm/model_executor/layers/fused_moe/cutlass_moe.py

+111 -68

@@ -9,12 +9,13 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 from vllm import _custom_ops as ops

10 from vllm.logger import init_logger

11 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

12 +from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (

13 + moe_permute, moe_unpermute)

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

15 MoEPrepareAndFinalizeNoEP)

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

17 TopKWeightAndReduceDelegate, TopKWeightAndReduceNoOP)

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

17 - _fp8_quantize,

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

19 _resize_cache)

20 from vllm.scalar_type import scalar_types

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

35 w2_scale: Optional[torch.Tensor],

36 a1q_scale: Optional[torch.Tensor],

37 a2_scale: Optional[torch.Tensor],

38 + ab_strides1: torch.Tensor,

39 + ab_strides2: torch.Tensor,

40 + c_strides1: torch.Tensor,

41 + c_strides2: torch.Tensor,

42 workspace13: torch.Tensor,

43 workspace2: torch.Tensor,

44 expert_num_tokens: Optional[torch.Tensor],

@@ -41,6 +46,7 @@ def run_cutlass_moe_fp8(

46 per_act_token: bool,

47 per_out_ch: bool,

48 use_batched_format: bool,

49 + topk_weights: Optional[torch.Tensor],

50 ):

51 a1q = hidden_states

@@ -99,6 +105,22 @@ def run_cutlass_moe_fp8(

105 topk = local_topk_ids.size(1)

106 local_E = w1.size(0)

107

108 + if use_batched_format:

109 + mm1_out = _resize_cache(workspace13, (local_E * padded_M, N * 2))

110 + act_out = _resize_cache(workspace2, (local_E * padded_M, N))

111 + quant_out = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn),

112 + (local_E * padded_M, N))

113 + mm2_out = _resize_cache(workspace2, (local_E * padded_M, K))

114 + else:

115 + a1q_perm = _resize_cache(workspace2.view(dtype=torch.float8_e4m3fn),

116 + (M * topk, K))

117 + mm1_out = _resize_cache(workspace13, (M * topk, N * 2))

118 + act_out = _resize_cache(workspace2, (M * topk, N))

119 + # original workspace are based on input hidden_states dtype (bf16)

120 + quant_out = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn),

121 + (M * topk, N))

122 + mm2_out = _resize_cache(workspace2, (M * topk, K))

123 +

124 if use_batched_format:

125 assert expert_num_tokens is not None

126

@@ -120,11 +142,10 @@ def run_cutlass_moe_fp8(

142 w2_scale = w2_scale.reshape(w2_scale.size(0), -1)

143 a1q = a1q.reshape(-1, a1q.size(2))

144 a1q_scale = a1q_scale.reshape(-1, a1q_scale.size(2)).contiguous()

123 -

145 + # c3x get_group_gemm_starts expects int64 to avoid overflow

146 + # during offset calculations

147 + expert_offsets = expert_offsets.to(torch.int64)

148 else:

125 - expert_offsets = torch.empty((global_num_experts + 1),

126 - dtype=torch.int32,

127 - device=device)

149 problem_sizes1 = torch.empty((global_num_experts, 3),

150 dtype=torch.int32,

151 device=device)

@@ -132,84 +153,57 @@ def run_cutlass_moe_fp8(

153 dtype=torch.int32,

154 device=device)

155

135 - # With expert_map each Rank processes only a subset of experts. As

136 - # a result not all of a_map and c2 tensors are filled. We fill it

137 - # zeros for correctness.

138 - if expert_map is not None:

139 - a_map = torch.zeros((local_topk_ids.numel()),

140 - dtype=torch.int32,

141 - device=device)

142 - else:

143 - a_map = torch.empty((local_topk_ids.numel()),

144 - dtype=torch.int32,

145 - device=device)

146 -

147 - c_map = torch.empty((local_topk_ids.numel()),

148 - dtype=torch.int32,

149 - device=device)

150 -

151 - ops.get_cutlass_moe_mm_data(local_topk_ids, expert_offsets,

152 - problem_sizes1, problem_sizes2, a_map,

153 - c_map, global_num_experts, N, K)

154 -

155 - a1q = _fp8_perm(a1q, a_map)

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

156 + num_expert = global_num_experts if expert_map is None \

157 + else expert_map.size(0)

158 + # permuted a1q reuses workspace2

159 + a1q, a1q_scale, expert_offsets, inv_perm, _ = moe_permute(

160 + a1q,

161 + a1q_scale,

162 + topk_ids,

163 + num_expert,

164 + local_E,

165 + expert_map,

166 + permuted_hidden_states=a1q_perm)

167 expert_offsets = expert_offsets[:-1]

168

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 -

176 - if use_batched_format:

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

178 - c2 = _resize_cache(workspace2, (local_E * padded_M, N))

179 - c3 = _resize_cache(workspace13, (local_E * padded_M, K))

180 - else:

181 - c1 = _resize_cache(workspace13, (M * topk, N * 2))

182 - c2 = _resize_cache(workspace2, (M * topk, N))

183 - c3 = _resize_cache(workspace13, (M * topk, K))

169 + ops.get_cutlass_moe_mm_problem_sizes(local_topk_ids, problem_sizes1,

170 + problem_sizes2,

171 + global_num_experts, N, K)

172

173 if not per_act_token and (expert_map is not None or use_batched_format):

174 # this is necessary to avoid imprecise scale calculation caused by

175 # random data in the unused workspace. The workspace is unused when

176 # this rank handles only partial tokens, or when it is batched .

189 - c1.fill_(0)

177 + mm1_out.fill_(0)

178

191 - ops.cutlass_moe_mm(c1, a1q, w1, a1q_scale, w1_scale, expert_offsets,

179 + ops.cutlass_moe_mm(mm1_out, a1q, w1, a1q_scale, w1_scale, expert_offsets,

180 problem_sizes1, ab_strides1, ab_strides1, c_strides1,

181 per_act_token, per_out_ch)

182

195 - activation_callable(c2, c1)

183 + activation_callable(act_out, mm1_out)

184

185 a2q, a2q_scale = ops.scaled_fp8_quant(

198 - c2, a2_scale, use_per_token_if_dynamic=per_act_token)

186 + act_out,

187 + a2_scale,

188 + use_per_token_if_dynamic=per_act_token,

189 + output=quant_out)

190

191 if expert_map is not None:

201 - c3.fill_(0)

192 + mm2_out.fill_(0)

193

203 - ops.cutlass_moe_mm(c3, a2q, w2, a2q_scale, w2_scale, expert_offsets,

194 + ops.cutlass_moe_mm(mm2_out, a2q, w2, a2q_scale, w2_scale, expert_offsets,

195 problem_sizes2, ab_strides2, ab_strides2, c_strides2,

196 per_act_token, per_out_ch)

197

198 if use_batched_format:

208 - output.copy_(c3.reshape(local_E, padded_M, K), non_blocking=True)

199 + output.copy_(mm2_out.reshape(local_E, padded_M, K), non_blocking=True)

200 else:

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

211 - # as c3.

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

201 + # for non-chunking mode the output is resized from workspace13

202 + # so we need to make sure mm2_out uses workspace2.

203 + moe_unpermute(out=output,

204 + permuted_hidden_states=mm2_out,

205 + topk_weights=topk_weights,

206 + inv_permuted_idx=inv_perm)

207

208

209 class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

@@ -219,6 +213,10 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

213 out_dtype: Optional[torch.dtype],

214 per_act_token_quant: bool,

215 per_out_ch_quant: bool,

216 + ab_strides1: torch.Tensor,

217 + ab_strides2: torch.Tensor,

218 + c_strides1: torch.Tensor,

219 + c_strides2: torch.Tensor,

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

221 ):

222 super().__init__(

@@ -229,6 +227,10 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

227 block_shape=block_shape,

228 ))

229 self.out_dtype = out_dtype

230 + self.ab_strides1 = ab_strides1

231 + self.ab_strides2 = ab_strides2

232 + self.c_strides1 = c_strides1

233 + self.c_strides2 = c_strides2

234

235 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

236 # Let PrepareAndFinalize::finalize() decide the impl.

@@ -272,10 +274,11 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

274 run_cutlass_moe_fp8(

275 output, hidden_states, w1, w2, topk_ids, activation_callable,

276 global_num_experts, expert_map, w1_scale, w2_scale, a1q_scale,

275 - a2_scale, workspace13, workspace2, expert_num_tokens,

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

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

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

280 self.per_act_token_quant, self.per_out_ch_quant,

278 - use_batched_format)

281 + use_batched_format, topk_weights)

282

283

284 class CutlassExpertsFp8(CutlassExpertsFp8Base):

@@ -285,12 +288,20 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):

288 out_dtype: Optional[torch.dtype],

289 per_act_token_quant: bool,

290 per_out_ch_quant: bool,

291 + ab_strides1: torch.Tensor,

292 + ab_strides2: torch.Tensor,

293 + c_strides1: torch.Tensor,

294 + c_strides2: torch.Tensor,

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

296 ):

297 super().__init__(

298 out_dtype,

299 per_act_token_quant,

300 per_out_ch_quant,

301 + ab_strides1,

302 + ab_strides2,

303 + c_strides1,

304 + c_strides2,

305 block_shape,

306 )

307

@@ -307,6 +318,10 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):

318 def supports_expert_map(self) -> bool:

319 return True

320

321 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

322 + # topk weights and reduction are fused in moe_unpermute cuda kernel

323 + return TopKWeightAndReduceNoOP()

324 +

325 def workspace_shapes(

326 self,

327 a: torch.Tensor,

@@ -320,8 +335,8 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):

335 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

336 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:

337 workspace1 = (M * topk, max(N, K))

323 - workspace2 = (M * topk, N // 2)

324 - output = (M * topk, K)

338 + workspace2 = (M * topk, max(N // 2, K))

339 + output = (M, K)

340 return (workspace1, workspace2, output,

341 self.out_dtype if self.out_dtype is not None else a.dtype)

342

@@ -335,12 +350,20 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):

350 out_dtype: Optional[torch.dtype],

351 per_act_token_quant: bool,

352 per_out_ch_quant: bool,

353 + ab_strides1: torch.Tensor,

354 + ab_strides2: torch.Tensor,

355 + c_strides1: torch.Tensor,

356 + c_strides2: torch.Tensor,

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

358 ):

359 super().__init__(

360 out_dtype,

361 per_act_token_quant,

362 per_out_ch_quant,

363 + ab_strides1,

364 + ab_strides2,

365 + c_strides1,

366 + c_strides2,

367 block_shape,

368 )

369 assert max_experts_per_worker > 0

@@ -378,7 +401,8 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):

401 assert num_dp is not None

402 workspace1 = (self.max_experts_per_worker, padded_M * num_dp,

403 max(N, K))

381 - workspace2 = (self.max_experts_per_worker, padded_M * num_dp, (N // 2))

404 + workspace2 = (self.max_experts_per_worker, padded_M * num_dp,

405 + max(N // 2, K))

406 output = (self.max_experts_per_worker, padded_M, K)

407 return (workspace1, workspace2, output,

408 self.out_dtype if self.out_dtype is not None else a.dtype)

@@ -392,6 +416,10 @@ def cutlass_moe_fp8(

416 topk_ids: torch.Tensor,

417 w1_scale: torch.Tensor,

418 w2_scale: torch.Tensor,

419 + ab_strides1: torch.Tensor,

420 + ab_strides2: torch.Tensor,

421 + c_strides1: torch.Tensor,

422 + c_strides2: torch.Tensor,

423 per_act_token: Optional[bool] = None,

424 activation: str = "silu",

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

@@ -419,6 +447,17 @@ def cutlass_moe_fp8(

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

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

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

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

451 + Shape: [num_experts]

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

453 + Shape: [num_experts]

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

455 + Shape: [num_experts]

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

457 + Shape: [num_experts]

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

459 + per-tensor.

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

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

462 Shape: scalar or [M]

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

@@ -450,6 +489,10 @@ def cutlass_moe_fp8(

489 out_dtype=a.dtype,

490 per_act_token_quant=per_act_token,

491 per_out_ch_quant=per_out_ch,

492 + ab_strides1=ab_strides1,

493 + ab_strides2=ab_strides2,

494 + c_strides1=c_strides1,

495 + c_strides2=c_strides2,

496 ),

497 )

498

vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py

+20 -9

@@ -82,7 +82,8 @@ def moe_permute(

82 n_local_expert: int = -1,

83 expert_map: Optional[torch.Tensor] = None,

84 align_block_size: Optional[int] = None,

85 - fill_invalid_expert: int = -1

85 + fill_invalid_expert: int = -1,

86 + permuted_hidden_states: Optional[torch.Tensor] = None,

87 ) -> tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor,

88 torch.Tensor]:

89 """

@@ -95,14 +96,17 @@ def moe_permute(

96 - n_expert (int): The number of expert.

97 - n_local_expert (int): The number of expert in current EP rank.

98 - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices

98 - from the global expert space to the local expert space of the expert

99 + from the global expert space to the local expert space of the expert

100 parallel shard.

101 - align_block_size (Optional[int]): align group gemm block size for deepgemm

102 - fill_invalid_expert(int): fill expert id in m_indices for invalid expert

103 to workaround DeepGemm unsupported -1 in m_indices

104 + - permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor.

105 + If None, the output tensor will be created in this function.

106 Returns:

107 - permuted_hidden_states (torch.Tensor): permuted activation.

105 - - a1q_scale (Optional[torch.Tensor]): quant scale for hidden_states

108 + - a1q_scale (Optional[torch.Tensor]): permuted quant scale for hidden_states

109 + if original scale not per-tensor scaling

110 - expert_first_token_offset (torch.Tensor): offset of the first token

111 of each expert for standard grouped gemm. if enable 'align_block_size'

112 expert_first_token_offset will align up to 'align_block_size'.

@@ -122,11 +126,16 @@ def moe_permute(

126 1) // align_block_size * align_block_size

127 if n_local_expert == -1:

128 n_local_expert = n_expert

125 - permuted_hidden_states = torch.empty(

126 - (permuted_row_size, n_hidden),

127 - dtype=hidden_states.dtype,

128 - device=hidden_states.device,

129 - )

129 + if permuted_hidden_states is None:

130 + permuted_hidden_states = torch.empty(

131 + (permuted_row_size, n_hidden),

132 + dtype=hidden_states.dtype,

133 + device=hidden_states.device,

134 + )

135 + assert permuted_hidden_states.size() == (permuted_row_size, n_hidden), (

136 + f"Expected permuted hidden states to be {(permuted_row_size, n_hidden)}"

137 + f" but got {permuted_hidden_states.size()}")

138 +

139 token_expert_indices = torch.arange(0,

140 n_token * topk,

141 dtype=torch.int32,

@@ -153,7 +162,8 @@ def moe_permute(

162 align_block_size, permuted_hidden_states,

163 expert_first_token_offset, inv_permuted_idx,

164 permuted_idx, m_indices)

156 - if a1q_scale is not None:

165 +

166 + if a1q_scale is not None and a1q_scale.dim() > 1:

167 a1q_scale = a1q_scale[permuted_idx.clamp(max=n_token * topk - 1) //

168 topk]

169 return (permuted_hidden_states, a1q_scale, expert_first_token_offset,

@@ -185,6 +195,7 @@ def moe_unpermute(

195 n_hidden = permuted_hidden_states.size(-1)

196 assert (n_hidden * permuted_hidden_states.element_size()

197 ) % 16 == 0, "unpermue kernel need hidden dim align to 16B"

198 +

199 torch.ops._moe_C.moe_unpermute(permuted_hidden_states, topk_weights,

200 inv_permuted_idx, expert_first_token_offset,

201 topk, out)

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

+31

@@ -669,6 +669,25 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

669 from vllm.model_executor.layers.fused_moe import fused_experts

670 self.fused_experts_func = fused_experts

671

672 + if self.use_cutlass:

673 + device = layer.w13_weight.device

674 + # ab_strides1 and c_strides2 are the same

675 + self.ab_strides1_c_strides2 = torch.full(

676 + (layer.local_num_experts, ),

677 + layer.hidden_size,

678 + device=device,

679 + dtype=torch.int64)

680 + self.ab_strides2 = torch.full(

681 + (layer.local_num_experts, ),

682 + layer.intermediate_size_per_partition,

683 + device=device,

684 + dtype=torch.int64)

685 + self.c_strides1 = torch.full(

686 + (layer.local_num_experts, ),

687 + 2 * layer.intermediate_size_per_partition,

688 + device=device,

689 + dtype=torch.int64)

690 +

691 def select_gemm_impl(

692 self,

693 prepare_finalize: FusedMoEPrepareAndFinalize,

@@ -693,6 +712,10 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

712 moe.in_dtype,

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

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

715 + ab_strides1=self.ab_strides1_c_strides2,

716 + ab_strides2=self.ab_strides2,

717 + c_strides1=self.c_strides1,

718 + c_strides2=self.ab_strides1_c_strides2,

719 )

720 else:

721 logger.debug("CutlassExpertsFp8(%s)", self.__class__.__name__)

@@ -700,6 +723,10 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

723 moe.in_dtype,

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

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

726 + ab_strides1=self.ab_strides1_c_strides2,

727 + ab_strides2=self.ab_strides2,

728 + c_strides1=self.c_strides1,

729 + c_strides2=self.ab_strides1_c_strides2,

730 )

731

732 self.disable_expert_map = (num_dispatchers > 1

@@ -822,6 +849,10 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

849 expert_map=None if self.disable_expert_map else expert_map,

850 w1_scale=layer.w13_weight_scale,

851 w2_scale=layer.w2_weight_scale,

852 + ab_strides1=self.ab_strides1_c_strides2,

853 + ab_strides2=self.ab_strides2,

854 + c_strides1=self.c_strides1,

855 + c_strides2=self.ab_strides1_c_strides2,

856 a1_scale=layer.w13_input_scale,

857 a2_scale=layer.w2_input_scale,

858 )