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

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, 02:15 UTC

[Misc] ModularKernel : Perform WeightAndReduce inside TritonExperts & DeepGemmExperts (#20725)

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>

Varun Sundar Rabindranath committed Jul 15, 2025 at 01:17 UTC c0569dbc82b5e945a77878190114d1b68027828b

9 files changed +203 -157

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

@@ -260,6 +260,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

260 hidden_states: torch.Tensor,

261 w1: torch.Tensor,

262 w2: torch.Tensor,

263 + topk_weights: torch.Tensor,

264 topk_ids: torch.Tensor,

265 activation: str,

266 global_num_experts: int,

@@ -273,6 +274,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

274 workspace13: torch.Tensor,

275 workspace2: torch.Tensor,

276 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

277 + apply_router_weight_on_input: bool,

278 ):

279 assert expert_tokens_meta is not None

280 expert_num_tokens = expert_tokens_meta.expert_num_tokens

vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py

+16 -24

@@ -129,30 +129,22 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

129 return self.batched_triton_experts.workspace_shapes(

130 a, aq, M, N, K, topk, global_num_experts, local_num_experts)

131

132 - def apply(

133 - self,

134 - output: torch.Tensor,

135 - hidden_states: torch.Tensor,

136 - w1: torch.Tensor,

137 - w2: torch.Tensor,

138 - topk_ids: torch.Tensor,

139 - activation: str,

140 - global_num_experts: int,

141 - expert_map: Optional[torch.Tensor],

142 - w1_scale: Optional[torch.Tensor],

143 - w2_scale: Optional[torch.Tensor],

144 - w1_zp: Optional[torch.Tensor],

145 - w2_zp: Optional[torch.Tensor],

146 - a1q_scale: Optional[torch.Tensor],

147 - a2_scale: Optional[torch.Tensor],

148 - workspace13: torch.Tensor,

149 - workspace2: torch.Tensor,

150 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

151 - ):

132 + def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,

133 + w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,

134 + topk_ids: torch.Tensor, activation: str, global_num_experts: int,

135 + expert_map: Optional[torch.Tensor],

136 + w1_scale: Optional[torch.Tensor],

137 + w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],

138 + w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],

139 + a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,

140 + workspace2: torch.Tensor,

141 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

142 + apply_router_weight_on_input: bool):

143 experts = (self.batched_deep_gemm_experts

144 if self.allow_deep_gemm else self.batched_triton_experts)

145 assert experts is not None

155 - experts.apply(output, hidden_states, w1, w2, topk_ids, activation,

156 - global_num_experts, expert_map, w1_scale, w2_scale,

157 - w1_zp, w2_zp, a1q_scale, a2_scale, workspace13,

158 - workspace2, expert_tokens_meta)

146 + experts.apply(output, hidden_states, w1, w2, topk_weights, topk_ids,

147 + activation, global_num_experts, expert_map, w1_scale,

148 + w2_scale, w1_zp, w2_zp, a1q_scale, a2_scale, workspace13,

149 + workspace2, expert_tokens_meta,

150 + apply_router_weight_on_input)

vllm/model_executor/layers/fused_moe/cutlass_moe.py

+11 -20

@@ -291,26 +291,17 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):

291 return (workspace1, workspace2, output,

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

293

294 - def apply(

295 - self,

296 - output: torch.Tensor,

297 - hidden_states: torch.Tensor,

298 - w1: torch.Tensor,

299 - w2: torch.Tensor,

300 - topk_ids: torch.Tensor,

301 - activation: str,

302 - global_num_experts: int,

303 - expert_map: Optional[torch.Tensor],

304 - w1_scale: Optional[torch.Tensor],

305 - w2_scale: Optional[torch.Tensor],

306 - w1_zp: Optional[torch.Tensor],

307 - w2_zp: Optional[torch.Tensor],

308 - a1q_scale: Optional[torch.Tensor],

309 - a2_scale: Optional[torch.Tensor],

310 - workspace13: torch.Tensor,

311 - workspace2: torch.Tensor,

312 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

313 - ):

294 + def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,

295 + w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,

296 + topk_ids: torch.Tensor, activation: str, global_num_experts: int,

297 + expert_map: Optional[torch.Tensor],

298 + w1_scale: Optional[torch.Tensor],

299 + w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],

300 + w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],

301 + a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,

302 + workspace2: torch.Tensor,

303 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

304 + apply_router_weight_on_input: bool):

305 assert w1_zp is None, "w1_zp is not supported in CUTLASS MoE"

306 assert w2_zp is None, "w2_zp is not supported in CUTLASS MoE"

307

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+21 -10

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

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

14 MoEPrepareAndFinalizeNoEP)

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

16 - TopKWeightAndReduceDelegate)

16 + TopKWeightAndReduceContiguous, TopKWeightAndReduceNoOP)

17 from vllm.model_executor.layers.fused_moe.utils import _resize_cache

18 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

19 per_token_group_quant_fp8)

@@ -90,8 +90,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

90 return True

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

93 - # Let PrepareAndFinalize::finalize() decide the impl.

94 - return TopKWeightAndReduceDelegate()

93 + return TopKWeightAndReduceNoOP()

95 def workspace_shapes(

96 self, a: torch.Tensor, aq: torch.Tensor, M: int, N: int, K: int,

@@ -104,9 +103,9 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

103 block_m = self.block_shape[0]

104 M_sum = (M * topk) + num_experts * (block_m - 1)

105 M_sum = round_up(M_sum, block_m)

107 - workspace1 = (M_sum, max(N * 2, K))

106 + workspace1 = (M_sum, max(N // 2, K))

107 workspace2 = (M_sum, max(N, K))

109 - output = (M, topk, K)

108 + output = (M, K)

109 return (workspace1, workspace2, output, a.dtype)

110

111 def apply(

@@ -115,6 +114,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

114 hidden_states: torch.Tensor,

115 w1: torch.Tensor,

116 w2: torch.Tensor,

117 + topk_weights: torch.Tensor,

118 topk_ids: torch.Tensor,

119 activation: str,

120 global_num_experts: int,

@@ -128,11 +128,14 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

128 workspace13: torch.Tensor,

129 workspace2: torch.Tensor,

130 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

131 + apply_router_weight_on_input: bool,

132 ):

133 assert self.block_shape is not None

134

135 a1q = hidden_states

136 _, N, K = w1.size()

137 + M, _ = output.size()

138 + num_topk = topk_ids.size(1)

139

140 if global_num_experts == -1:

141 global_num_experts = w1.size(0)

@@ -159,11 +162,12 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

162 # Note: M_sum is different than the pre-permuted shape of a1q.

163 M_sum = a1q.size(0)

164

162 - mm1_out = _resize_cache(workspace13, (M_sum, N))

163 - act_out = _resize_cache(workspace2, (M_sum, N // 2))

164 - quant_out = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn),

165 + mm1_out = _resize_cache(workspace2, (M_sum, N))

166 + act_out = _resize_cache(workspace13, (M_sum, N // 2))

167 + quant_out = _resize_cache(workspace2.view(dtype=torch.float8_e4m3fn),

168 (M_sum, N // 2))

166 - mm2_out = _resize_cache(workspace2, (M_sum, K))

169 + mm2_out = _resize_cache(workspace13, (M_sum, K))

170 + perm_out = _resize_cache(workspace2, (M * num_topk, K))

171

172 m_grouped_fp8_gemm_nt_contiguous((a1q, a1q_scale), (w1, w1_scale),

173 mm1_out, expert_ids)

@@ -179,7 +183,14 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

183 m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, w2_scale),

184 mm2_out, expert_ids)

185

182 - torch.index_select(mm2_out, 0, inv_perm, out=output.view((-1, K)))

186 + torch.index_select(mm2_out, 0, inv_perm, out=perm_out)

187 +

188 + TopKWeightAndReduceContiguous().apply(

189 + output=output,

190 + fused_expert_output=perm_out,

191 + topk_weights=topk_weights,

192 + topk_ids=topk_ids,

193 + apply_router_weight_on_input=apply_router_weight_on_input)

194

195

196 def deep_gemm_moe_fp8(

vllm/model_executor/layers/fused_moe/fused_batched_moe.py

+8 -6

@@ -696,15 +696,16 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

696 return t.to(f32) * group_broadcast(scale, t.shape)

697

698 def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,

699 - w1: torch.Tensor, w2: torch.Tensor, topk_ids: torch.Tensor,

700 - activation: str, global_num_experts: int,

699 + w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,

700 + topk_ids: torch.Tensor, activation: str, global_num_experts: int,

701 expert_map: Optional[torch.Tensor],

702 w1_scale: Optional[torch.Tensor],

703 w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],

704 w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],

705 a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,

706 workspace2: torch.Tensor,

707 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata]):

707 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

708 + apply_router_weight_on_input: bool):

709 assert hidden_states.dim() == 3

710 assert expert_tokens_meta is not None

711 expert_num_tokens = expert_tokens_meta.expert_num_tokens

@@ -899,15 +900,16 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

900 return (workspace13, workspace2, output, a.dtype)

901

902 def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,

902 - w1: torch.Tensor, w2: torch.Tensor, topk_ids: torch.Tensor,

903 - activation: str, global_num_experts: int,

903 + w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,

904 + topk_ids: torch.Tensor, activation: str, global_num_experts: int,

905 expert_map: Optional[torch.Tensor],

906 w1_scale: Optional[torch.Tensor],

907 w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],

908 w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],

909 a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,

910 workspace2: torch.Tensor,

910 - expert_tokens_meta: Optional[mk.ExpertTokensMetadata]):

911 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

912 + apply_router_weight_on_input: bool):

913 # Check constraints.

914 if self.use_int4_w4a16:

915 assert hidden_states.size(-1) // 2 == w1.size(2), (

vllm/model_executor/layers/fused_moe/fused_moe.py

+38 -33

@@ -26,7 +26,7 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (

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

27 MoEPrepareAndFinalizeNoEP)

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

29 - TopKWeightAndReduceDelegate)

29 + TopKWeightAndReduceNoOP)

30 from vllm.model_executor.layers.fused_moe.utils import (

31 _resize_cache, moe_kernel_quantize_input)

32 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (

@@ -1606,8 +1606,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1606 return True

1607

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

1609 - # Let PrepareAndFinalize::finalize() decide the impl.

1610 - return TopKWeightAndReduceDelegate()

1609 + return TopKWeightAndReduceNoOP()

1610

1611 def workspace_shapes(

1612 self,

@@ -1620,9 +1619,9 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1619 global_num_experts: int,

1620 local_num_experts: int,

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

1623 - workspace1 = (M, topk, max(N * 2, K))

1624 - workspace2 = (M, topk, N)

1625 - output = (M, topk, K)

1622 + workspace1 = (M, topk, max(N // 2, K))

1623 + workspace2 = (M, topk, max(N, K))

1624 + output = (M, K)

1625 return (workspace1, workspace2, output, a.dtype)

1626

1627 def apply(

@@ -1631,6 +1630,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1630 hidden_states: torch.Tensor,

1631 w1: torch.Tensor,

1632 w2: torch.Tensor,

1633 + topk_weights: torch.Tensor,

1634 topk_ids: torch.Tensor,

1635 activation: str,

1636 global_num_experts: int,

@@ -1644,6 +1644,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1644 workspace13: torch.Tensor,

1645 workspace2: torch.Tensor,

1646 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

1647 + apply_router_weight_on_input: bool,

1648 ):

1649 # Check constraints.

1650 if self.use_int4_w4a16:

@@ -1696,37 +1697,39 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1697 raise ValueError(

1698 f"Unsupported compute_type: {hidden_states.dtype}")

1699

1699 - # We can reuse the memory between these because by the time we need

1700 - # cache3, we're done with cache1

1701 - intermediate_cache1 = _resize_cache(workspace13,

1700 + # Note that the output tensor might be in workspace1

1701 + intermediate_cache1 = _resize_cache(workspace2,

1702 (num_tokens, top_k_num, N))

1703 - intermediate_cache2 = _resize_cache(workspace2,

1703 + intermediate_cache2 = _resize_cache(workspace13,

1704 (num_tokens * top_k_num, N // 2))

1705 + intermediate_cache3 = _resize_cache(workspace2,

1706 + (num_tokens, top_k_num, K))

1707

1708 sorted_token_ids, expert_ids, num_tokens_post_padded = (

1709 moe_align_block_size(topk_ids, config['BLOCK_SIZE_M'],

1710 global_num_experts, expert_map))

1711

1710 - invoke_fused_moe_kernel(hidden_states,

1711 - w1,

1712 - intermediate_cache1,

1713 - a1q_scale,

1714 - w1_scale,

1715 - w1_zp,

1716 - None,

1717 - sorted_token_ids,

1718 - expert_ids,

1719 - num_tokens_post_padded,

1720 - False,

1721 - top_k_num,

1722 - config,

1723 - compute_type=compute_type,

1724 - use_fp8_w8a8=self.use_fp8_w8a8,

1725 - use_int8_w8a8=self.use_int8_w8a8,

1726 - use_int8_w8a16=self.use_int8_w8a16,

1727 - use_int4_w4a16=self.use_int4_w4a16,

1728 - per_channel_quant=self.per_act_token_quant,

1729 - block_shape=self.block_shape)

1712 + invoke_fused_moe_kernel(

1713 + hidden_states,

1714 + w1,

1715 + intermediate_cache1,

1716 + a1q_scale,

1717 + w1_scale,

1718 + w1_zp,

1719 + None, # topk_weights

1720 + sorted_token_ids,

1721 + expert_ids,

1722 + num_tokens_post_padded,

1723 + False, # mul_routed_weights

1724 + top_k_num,

1725 + config,

1726 + compute_type=compute_type,

1727 + use_fp8_w8a8=self.use_fp8_w8a8,

1728 + use_int8_w8a8=self.use_int8_w8a8,

1729 + use_int8_w8a16=self.use_int8_w8a16,

1730 + use_int4_w4a16=self.use_int4_w4a16,

1731 + per_channel_quant=self.per_act_token_quant,

1732 + block_shape=self.block_shape)

1733

1734 self.activation(activation, intermediate_cache2,

1735 intermediate_cache1.view(-1, N))

@@ -1739,15 +1742,15 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1742

1743 invoke_fused_moe_kernel(qintermediate_cache2,

1744 w2,

1742 - output,

1745 + intermediate_cache3,

1746 a2q_scale,

1747 w2_scale,

1748 w2_zp,

1746 - None,

1749 + topk_weights,

1750 sorted_token_ids,

1751 expert_ids,

1752 num_tokens_post_padded,

1750 - False,

1753 + not apply_router_weight_on_input,

1754 1,

1755 config,

1756 compute_type=compute_type,

@@ -1758,6 +1761,8 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1761 per_channel_quant=self.per_act_token_quant,

1762 block_shape=self.block_shape)

1763

1764 + ops.moe_sum(intermediate_cache3, output)

1765 +

1766

1767 def modular_triton_fused_moe(

1768 use_fp8_w8a8: bool,

vllm/model_executor/layers/fused_moe/modular_kernel.py

+91 -59

@@ -360,6 +360,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

360 hidden_states: torch.Tensor,

361 w1: torch.Tensor,

362 w2: torch.Tensor,

363 + topk_weights: torch.Tensor,

364 topk_ids: torch.Tensor,

365 activation: str,

366 global_num_experts: int,

@@ -373,6 +374,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

374 workspace13: torch.Tensor,

375 workspace2: torch.Tensor,

376 expert_tokens_meta: Optional[ExpertTokensMetadata],

377 + apply_router_weight_on_input: bool,

378 ):

379 """

380 This function computes the intermediate result of a Mixture of Experts

@@ -384,6 +386,8 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

386 layer.

387 - w1 (torch.Tensor): The first set of expert weights.

388 - w2 (torch.Tensor): The second set of expert weights.

389 + - topk_weights: A map of row to expert weights. Some implementations

390 + choose to do weight application.

391 - topk_ids (torch.Tensor): A map of row to expert id.

392 - activation (str): The activation function to apply after the first

393 MoE layer.

@@ -409,6 +413,9 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

413 ExpertTokensMetadata object containing gpu/cpu tensors

414 as big as the number of local experts with the information about the

415 number of tokens assigned to each local expert.

416 + - apply_router_weight_on_input: True if router weights are already

417 + applied on the input. This is relevant if the implementation

418 + chooses to do weight application.

419 """

420 raise NotImplementedError

421

@@ -452,17 +459,21 @@ class FusedMoEModularKernel(torch.nn.Module):

459 f"{fused_experts.__class__.__name__}."

460 f"{fused_experts.activation_formats[0]}")

461

455 - def _do_fused_experts(

456 - self, fused_out: Optional[torch.Tensor], a1: torch.Tensor,

457 - a1q: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor,

458 - topk_ids: torch.Tensor, activation: str, global_num_experts: int,

459 - local_num_experts: int, expert_map: Optional[torch.Tensor],

460 - w1_scale: Optional[torch.Tensor], w2_scale: Optional[torch.Tensor],

461 - w1_zp: Optional[torch.Tensor], w2_zp: Optional[torch.Tensor],

462 - a1q_scale: Optional[torch.Tensor],

463 - a2_scale: Optional[torch.Tensor],

464 - expert_tokens_meta: Optional[ExpertTokensMetadata]

465 - ) -> torch.Tensor:

462 + def _do_fused_experts(self, fused_out: Optional[torch.Tensor],

463 + a1: torch.Tensor, a1q: torch.Tensor,

464 + w1: torch.Tensor, w2: torch.Tensor,

465 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,

466 + activation: str, global_num_experts: int,

467 + local_num_experts: int,

468 + expert_map: Optional[torch.Tensor],

469 + w1_scale: Optional[torch.Tensor],

470 + w2_scale: Optional[torch.Tensor],

471 + w1_zp: Optional[torch.Tensor],

472 + w2_zp: Optional[torch.Tensor],

473 + a1q_scale: Optional[torch.Tensor],

474 + a2_scale: Optional[torch.Tensor],

475 + expert_tokens_meta: Optional[ExpertTokensMetadata],

476 + apply_router_weight_on_input: bool) -> torch.Tensor:

477

478 _, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)

479

@@ -485,36 +496,49 @@ class FusedMoEModularKernel(torch.nn.Module):

496 # reuse workspace13 for the output

497 fused_out = _resize_cache(workspace13, fused_out_shape)

498

488 - self.fused_experts.apply(fused_out,

489 - a1q,

490 - w1,

491 - w2,

492 - topk_ids=topk_ids,

493 - activation=activation,

494 - global_num_experts=global_num_experts,

495 - expert_map=expert_map,

496 - w1_scale=w1_scale,

497 - w2_scale=w2_scale,

498 - w1_zp=w1_zp,

499 - w2_zp=w2_zp,

500 - a1q_scale=a1q_scale,

501 - a2_scale=a2_scale,

502 - workspace13=workspace13,

503 - workspace2=workspace2,

504 - expert_tokens_meta=expert_tokens_meta)

499 + self.fused_experts.apply(

500 + fused_out,

501 + a1q,

502 + w1,

503 + w2,

504 + topk_weights=topk_weights,

505 + topk_ids=topk_ids,

506 + activation=activation,

507 + global_num_experts=global_num_experts,

508 + expert_map=expert_map,

509 + w1_scale=w1_scale,

510 + w2_scale=w2_scale,

511 + w1_zp=w1_zp,

512 + w2_zp=w2_zp,

513 + a1q_scale=a1q_scale,

514 + a2_scale=a2_scale,

515 + workspace13=workspace13,

516 + workspace2=workspace2,

517 + expert_tokens_meta=expert_tokens_meta,

518 + apply_router_weight_on_input=apply_router_weight_on_input)

519

520 return fused_out

521

522 def _maybe_chunk_fused_experts(

509 - self, a1: torch.Tensor, a1q: torch.Tensor, w1: torch.Tensor,

510 - w2: torch.Tensor, topk_ids: torch.Tensor, activation: str,

511 - global_num_experts: int, local_num_experts: int,

512 - expert_map: Optional[torch.Tensor],

513 - w1_scale: Optional[torch.Tensor], w2_scale: Optional[torch.Tensor],

514 - w1_zp: Optional[torch.Tensor], w2_zp: Optional[torch.Tensor],

515 - a1q_scale: Optional[torch.Tensor],

516 - a2_scale: Optional[torch.Tensor],

517 - expert_tokens_meta: Optional[ExpertTokensMetadata]

523 + self,

524 + a1: torch.Tensor,

525 + a1q: torch.Tensor,

526 + w1: torch.Tensor,

527 + w2: torch.Tensor,

528 + topk_weights: torch.Tensor,

529 + topk_ids: torch.Tensor,

530 + activation: str,

531 + global_num_experts: int,

532 + local_num_experts: int,

533 + expert_map: Optional[torch.Tensor],

534 + w1_scale: Optional[torch.Tensor],

535 + w2_scale: Optional[torch.Tensor],

536 + w1_zp: Optional[torch.Tensor],

537 + w2_zp: Optional[torch.Tensor],

538 + a1q_scale: Optional[torch.Tensor],

539 + a2_scale: Optional[torch.Tensor],

540 + expert_tokens_meta: Optional[ExpertTokensMetadata],

541 + apply_router_weight_on_input: bool,

542 ) -> torch.Tensor:

543

544 _, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)

@@ -529,6 +553,7 @@ class FusedMoEModularKernel(torch.nn.Module):

553 a1q=a1q,

554 w1=w1,

555 w2=w2,

556 + topk_weights=topk_weights,

557 topk_ids=topk_ids,

558 activation=activation,

559 global_num_experts=global_num_experts,

@@ -540,7 +565,8 @@ class FusedMoEModularKernel(torch.nn.Module):

565 w2_zp=w2_zp,

566 a1q_scale=a1q_scale,

567 a2_scale=a2_scale,

543 - expert_tokens_meta=expert_tokens_meta)

568 + expert_tokens_meta=expert_tokens_meta,

569 + apply_router_weight_on_input=apply_router_weight_on_input)

570

571 # Chunking required case

572 assert num_chunks > 1

@@ -557,11 +583,12 @@ class FusedMoEModularKernel(torch.nn.Module):

583 def slice_input_tensors(

584 chunk_idx: int

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

560 - Optional[torch.Tensor], torch.Tensor]:

586 + Optional[torch.Tensor], torch.Tensor, torch.Tensor]:

587 s = chunk_idx * CHUNK_SIZE

588 e = min(s + CHUNK_SIZE, M)

589 return (a1q[s:e], _chunk_scales(a1q_scale, s, e),

564 - _chunk_scales(a2_scale, s, e), topk_ids[s:e])

590 + _chunk_scales(a2_scale, s,

591 + e), topk_ids[s:e], topk_weights[s:e])

592

593 def slice_output_tensor(chunk_idx: int) -> torch.Tensor:

594 assert fused_out.size(0) % M == 0, (

@@ -594,7 +621,7 @@ class FusedMoEModularKernel(torch.nn.Module):

621 expert_num_tokens_cpu=c_expert_num_tokens_cpu)

622

623 for chunk_idx in range(num_chunks):

597 - c_a1q, c_a1q_scale, c_a2_scale, c_topk_ids = (

624 + c_a1q, c_a1q_scale, c_a2_scale, c_topk_ids, c_topk_weights = (

625 slice_input_tensors(chunk_idx))

626

627 c_expert_tokens_meta = None

@@ -603,23 +630,26 @@ class FusedMoEModularKernel(torch.nn.Module):

630 expert_tokens_meta, c_topk_ids, local_num_experts,

631 expert_map)

632

606 - self._do_fused_experts(fused_out=slice_output_tensor(chunk_idx),

607 - a1=a1,

608 - a1q=c_a1q,

609 - w1=w1,

610 - w2=w2,

611 - topk_ids=c_topk_ids,

612 - activation=activation,

613 - global_num_experts=global_num_experts,

614 - local_num_experts=local_num_experts,

615 - expert_map=expert_map,

616 - w1_scale=w1_scale,

617 - w2_scale=w2_scale,

618 - w1_zp=w1_zp,

619 - w2_zp=w2_zp,

620 - a1q_scale=c_a1q_scale,

621 - a2_scale=c_a2_scale,

622 - expert_tokens_meta=c_expert_tokens_meta)

633 + self._do_fused_experts(

634 + fused_out=slice_output_tensor(chunk_idx),

635 + a1=a1,

636 + a1q=c_a1q,

637 + w1=w1,

638 + w2=w2,

639 + topk_weights=c_topk_weights,

640 + topk_ids=c_topk_ids,

641 + activation=activation,

642 + global_num_experts=global_num_experts,

643 + local_num_experts=local_num_experts,

644 + expert_map=expert_map,

645 + w1_scale=w1_scale,

646 + w2_scale=w2_scale,

647 + w1_zp=w1_zp,

648 + w2_zp=w2_zp,

649 + a1q_scale=c_a1q_scale,

650 + a2_scale=c_a2_scale,

651 + expert_tokens_meta=c_expert_tokens_meta,

652 + apply_router_weight_on_input=apply_router_weight_on_input)

653

654 return fused_out

655

@@ -719,6 +749,7 @@ class FusedMoEModularKernel(torch.nn.Module):

749 a1q=a1q,

750 w1=w1,

751 w2=w2,

752 + topk_weights=topk_weights,

753 topk_ids=topk_ids,

754 activation=activation,

755 global_num_experts=global_num_experts,

@@ -730,7 +761,8 @@ class FusedMoEModularKernel(torch.nn.Module):

761 w2_zp=w2_zp,

762 a1q_scale=a1q_scale,

763 a2_scale=a2_scale,

733 - expert_tokens_meta=expert_tokens_meta)

764 + expert_tokens_meta=expert_tokens_meta,

765 + apply_router_weight_on_input=apply_router_weight_on_input)

766

767 self.prepare_finalize.finalize(

768 output, fused_out, topk_weights, topk_ids,

vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py

+12 -5

@@ -48,11 +48,18 @@ class TopKWeightAndReduceNoOP(mk.TopKWeightAndReduce):

48 fused_expert_output: torch.Tensor, topk_weights: torch.Tensor,

49 topk_ids: torch.Tensor,

50 apply_router_weight_on_input: bool) -> torch.Tensor:

51 - # Relax this if an explicit copy is necessary. Note that,

52 - # if a copy is employed we have to make sure that the

53 - # tensors don't overlap

54 - assert output is None

55 - return fused_expert_output

51 + # Weight application and reduction operations are already done.

52 + if output is None:

53 + return fused_expert_output

54 +

55 + # MoEPrepareAndFinalizeNoEP needs the output to be in the `output`

56 + # tensor.

57 + assert output.size() == fused_expert_output.size(), (

58 + "output shape is expected to match the fused_expert_output shape. "

59 + f"But got output={output.size()}, "

60 + f"used_expert_output={fused_expert_output.size()}")

61 + output.copy_(fused_expert_output, non_blocking=True)

62 + return output

65 class TopKWeightAndReduceContiguous(mk.TopKWeightAndReduce):

vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py

@@ -122,6 +122,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

122 hidden_states: torch.Tensor,

123 w1: torch.Tensor,

124 w2: torch.Tensor,

125 + topk_weights: torch.Tensor,

126 topk_ids: torch.Tensor,

127 activation: str,

128 global_num_experts: int,

@@ -135,6 +136,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

136 workspace13: torch.Tensor,

137 workspace2: torch.Tensor,

138 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

139 + apply_router_weight_on_input: bool,

140 ):

141 use_deep_gemm = (self.allow_deep_gemm

142 and (_valid_deep_gemm(hidden_states, w1, w2)

@@ -148,6 +150,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

150 hidden_states,

151 w1,

152 w2,

153 + topk_weights,

154 topk_ids,

155 activation,

156 global_num_experts,

@@ -161,4 +164,5 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

164 workspace13,

165 workspace2,

166 expert_tokens_meta,

167 + apply_router_weight_on_input,

168 )