Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jul 14, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c9f2001
do reduction in experts
Jul 11, 2025
4d7e07b
fix workspace overallocation
Jul 12, 2025
2961f53
TritonExperts opt
Jul 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,7 @@ def apply(
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
activation: str,
global_num_experts: int,
Expand All @@ -273,6 +274,7 @@ def apply(
workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
apply_router_weight_on_input: bool,
):
assert expert_tokens_meta is not None
expert_num_tokens = expert_tokens_meta.expert_num_tokens
Expand Down
40 changes: 16 additions & 24 deletions vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,30 +129,22 @@ def workspace_shapes(
return self.batched_triton_experts.workspace_shapes(
a, aq, M, N, K, topk, global_num_experts, local_num_experts)

def apply(
self,
output: torch.Tensor,
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_ids: torch.Tensor,
activation: str,
global_num_experts: int,
expert_map: Optional[torch.Tensor],
w1_scale: Optional[torch.Tensor],
w2_scale: Optional[torch.Tensor],
w1_zp: Optional[torch.Tensor],
w2_zp: Optional[torch.Tensor],
a1q_scale: Optional[torch.Tensor],
a2_scale: Optional[torch.Tensor],
workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
):
def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
topk_ids: torch.Tensor, activation: str, global_num_experts: int,
expert_map: Optional[torch.Tensor],
w1_scale: Optional[torch.Tensor],
w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],
w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
apply_router_weight_on_input: bool):
experts = (self.batched_deep_gemm_experts
if self.allow_deep_gemm else self.batched_triton_experts)
assert experts is not None
experts.apply(output, hidden_states, w1, w2, topk_ids, activation,
global_num_experts, expert_map, w1_scale, w2_scale,
w1_zp, w2_zp, a1q_scale, a2_scale, workspace13,
workspace2, expert_tokens_meta)
experts.apply(output, hidden_states, w1, w2, topk_weights, topk_ids,
activation, global_num_experts, expert_map, w1_scale,
w2_scale, w1_zp, w2_zp, a1q_scale, a2_scale, workspace13,
workspace2, expert_tokens_meta,
apply_router_weight_on_input)
31 changes: 11 additions & 20 deletions vllm/model_executor/layers/fused_moe/cutlass_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -291,26 +291,17 @@ def workspace_shapes(
return (workspace1, workspace2, output,
self.out_dtype if self.out_dtype is not None else a.dtype)

def apply(
self,
output: torch.Tensor,
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_ids: torch.Tensor,
activation: str,
global_num_experts: int,
expert_map: Optional[torch.Tensor],
w1_scale: Optional[torch.Tensor],
w2_scale: Optional[torch.Tensor],
w1_zp: Optional[torch.Tensor],
w2_zp: Optional[torch.Tensor],
a1q_scale: Optional[torch.Tensor],
a2_scale: Optional[torch.Tensor],
workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
):
def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
topk_ids: torch.Tensor, activation: str, global_num_experts: int,
expert_map: Optional[torch.Tensor],
w1_scale: Optional[torch.Tensor],
w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],
w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
apply_router_weight_on_input: bool):
assert w1_zp is None, "w1_zp is not supported in CUTLASS MoE"
assert w2_zp is None, "w2_zp is not supported in CUTLASS MoE"

Expand Down
31 changes: 21 additions & 10 deletions vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP)
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
TopKWeightAndReduceDelegate)
TopKWeightAndReduceContiguous, TopKWeightAndReduceNoOP)
from vllm.model_executor.layers.fused_moe.utils import _resize_cache
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
per_token_group_quant_fp8)
Expand Down Expand Up @@ -90,8 +90,7 @@ def supports_expert_map(self) -> bool:
return True

def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
# Let PrepareAndFinalize::finalize() decide the impl.
return TopKWeightAndReduceDelegate()
return TopKWeightAndReduceNoOP()

def workspace_shapes(
self, a: torch.Tensor, aq: torch.Tensor, M: int, N: int, K: int,
Expand All @@ -104,9 +103,9 @@ def workspace_shapes(
block_m = self.block_shape[0]
M_sum = (M * topk) + num_experts * (block_m - 1)
M_sum = round_up(M_sum, block_m)
workspace1 = (M_sum, max(N * 2, K))
workspace1 = (M_sum, max(N // 2, K))
workspace2 = (M_sum, max(N, K))
output = (M, topk, K)
output = (M, K)
return (workspace1, workspace2, output, a.dtype)

def apply(
Expand All @@ -115,6 +114,7 @@ def apply(
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
activation: str,
global_num_experts: int,
Expand All @@ -128,11 +128,14 @@ def apply(
workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
apply_router_weight_on_input: bool,
):
assert self.block_shape is not None

a1q = hidden_states
_, N, K = w1.size()
M, _ = output.size()
num_topk = topk_ids.size(1)

if global_num_experts == -1:
global_num_experts = w1.size(0)
Expand All @@ -159,11 +162,12 @@ def apply(
# Note: M_sum is different than the pre-permuted shape of a1q.
M_sum = a1q.size(0)

mm1_out = _resize_cache(workspace13, (M_sum, N))
act_out = _resize_cache(workspace2, (M_sum, N // 2))
quant_out = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn),
mm1_out = _resize_cache(workspace2, (M_sum, N))
act_out = _resize_cache(workspace13, (M_sum, N // 2))
quant_out = _resize_cache(workspace2.view(dtype=torch.float8_e4m3fn),
(M_sum, N // 2))
mm2_out = _resize_cache(workspace2, (M_sum, K))
mm2_out = _resize_cache(workspace13, (M_sum, K))
perm_out = _resize_cache(workspace2, (M * num_topk, K))
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rearrage how workspaces are used to make space for perm_out - note that perm_out cannot use workspace13 as workspace13 may be used as the output tensor (

vllm/vllm/model_executor/layers/fused_moe/modular_kernel.py

Line 486 in 5923ab9

fused_out = _resize_cache(workspace13, fused_out_shape)
)


m_grouped_fp8_gemm_nt_contiguous((a1q, a1q_scale), (w1, w1_scale),
mm1_out, expert_ids)
Expand All @@ -179,7 +183,14 @@ def apply(
m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, w2_scale),
mm2_out, expert_ids)

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

TopKWeightAndReduceContiguous().apply(
output=output,
fused_expert_output=perm_out,
topk_weights=topk_weights,
topk_ids=topk_ids,
apply_router_weight_on_input=apply_router_weight_on_input)


def deep_gemm_moe_fp8(
Expand Down
14 changes: 8 additions & 6 deletions vllm/model_executor/layers/fused_moe/fused_batched_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -696,15 +696,16 @@ def dequant(self, t: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
return t.to(f32) * group_broadcast(scale, t.shape)

def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
w1: torch.Tensor, w2: torch.Tensor, topk_ids: torch.Tensor,
activation: str, global_num_experts: int,
w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
topk_ids: torch.Tensor, activation: str, global_num_experts: int,
expert_map: Optional[torch.Tensor],
w1_scale: Optional[torch.Tensor],
w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],
w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: Optional[mk.ExpertTokensMetadata]):
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
apply_router_weight_on_input: bool):
assert hidden_states.dim() == 3
assert expert_tokens_meta is not None
expert_num_tokens = expert_tokens_meta.expert_num_tokens
Expand Down Expand Up @@ -899,15 +900,16 @@ def workspace_shapes(
return (workspace13, workspace2, output, a.dtype)

def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
w1: torch.Tensor, w2: torch.Tensor, topk_ids: torch.Tensor,
activation: str, global_num_experts: int,
w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
topk_ids: torch.Tensor, activation: str, global_num_experts: int,
expert_map: Optional[torch.Tensor],
w1_scale: Optional[torch.Tensor],
w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],
w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: Optional[mk.ExpertTokensMetadata]):
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
apply_router_weight_on_input: bool):
# Check constraints.
if self.use_int4_w4a16:
assert hidden_states.size(-1) // 2 == w1.size(2), (
Expand Down
71 changes: 38 additions & 33 deletions vllm/model_executor/layers/fused_moe/fused_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP)
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
TopKWeightAndReduceDelegate)
TopKWeightAndReduceNoOP)
from vllm.model_executor.layers.fused_moe.utils import (
_resize_cache, moe_kernel_quantize_input)
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
Expand Down Expand Up @@ -1606,8 +1606,7 @@ def supports_expert_map(self) -> bool:
return True

def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
# Let PrepareAndFinalize::finalize() decide the impl.
return TopKWeightAndReduceDelegate()
return TopKWeightAndReduceNoOP()

def workspace_shapes(
self,
Expand All @@ -1620,9 +1619,9 @@ def workspace_shapes(
global_num_experts: int,
local_num_experts: int,
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
workspace1 = (M, topk, max(N * 2, K))
workspace2 = (M, topk, N)
output = (M, topk, K)
workspace1 = (M, topk, max(N // 2, K))
workspace2 = (M, topk, max(N, K))
output = (M, K)
return (workspace1, workspace2, output, a.dtype)

def apply(
Expand All @@ -1631,6 +1630,7 @@ def apply(
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
activation: str,
global_num_experts: int,
Expand All @@ -1644,6 +1644,7 @@ def apply(
workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
apply_router_weight_on_input: bool,
):
# Check constraints.
if self.use_int4_w4a16:
Expand Down Expand Up @@ -1696,37 +1697,39 @@ def apply(
raise ValueError(
f"Unsupported compute_type: {hidden_states.dtype}")

# We can reuse the memory between these because by the time we need
# cache3, we're done with cache1
intermediate_cache1 = _resize_cache(workspace13,
# Note that the output tensor might be in workspace1
intermediate_cache1 = _resize_cache(workspace2,
(num_tokens, top_k_num, N))
intermediate_cache2 = _resize_cache(workspace2,
intermediate_cache2 = _resize_cache(workspace13,
(num_tokens * top_k_num, N // 2))
intermediate_cache3 = _resize_cache(workspace2,
(num_tokens, top_k_num, K))
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rearrage how workspaces are used to make space for intermediate_cache3 - note that intermediate_cache3 cannot use workspace13 as workspace13 may be used as the output tensor


sorted_token_ids, expert_ids, num_tokens_post_padded = (
moe_align_block_size(topk_ids, config['BLOCK_SIZE_M'],
global_num_experts, expert_map))

invoke_fused_moe_kernel(hidden_states,
w1,
intermediate_cache1,
a1q_scale,
w1_scale,
w1_zp,
None,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
False,
top_k_num,
config,
compute_type=compute_type,
use_fp8_w8a8=self.use_fp8_w8a8,
use_int8_w8a8=self.use_int8_w8a8,
use_int8_w8a16=self.use_int8_w8a16,
use_int4_w4a16=self.use_int4_w4a16,
per_channel_quant=self.per_act_token_quant,
block_shape=self.block_shape)
invoke_fused_moe_kernel(
hidden_states,
w1,
intermediate_cache1,
a1q_scale,
w1_scale,
w1_zp,
None, # topk_weights
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
False, # mul_routed_weights
top_k_num,
config,
compute_type=compute_type,
use_fp8_w8a8=self.use_fp8_w8a8,
use_int8_w8a8=self.use_int8_w8a8,
use_int8_w8a16=self.use_int8_w8a16,
use_int4_w4a16=self.use_int4_w4a16,
per_channel_quant=self.per_act_token_quant,
block_shape=self.block_shape)

self.activation(activation, intermediate_cache2,
intermediate_cache1.view(-1, N))
Expand All @@ -1739,15 +1742,15 @@ def apply(

invoke_fused_moe_kernel(qintermediate_cache2,
w2,
output,
intermediate_cache3,
a2q_scale,
w2_scale,
w2_zp,
None,
topk_weights,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
False,
not apply_router_weight_on_input,
1,
config,
compute_type=compute_type,
Expand All @@ -1758,6 +1761,8 @@ def apply(
per_channel_quant=self.per_act_token_quant,
block_shape=self.block_shape)

ops.moe_sum(intermediate_cache3, output)


def modular_triton_fused_moe(
use_fp8_w8a8: bool,
Expand Down
Loading