Add no_combine support to cutlass_moe_fp4 (#25688)
Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com>
This commit is contained in:
co-authored by
Hanming Lu
parent
878e6b8886
commit
745abd6cc0
@@ -359,6 +359,7 @@ def cutlass_moe_fp4(
|
|||||||
topk_ids: torch.Tensor,
|
topk_ids: torch.Tensor,
|
||||||
params: CutlassMoEParams,
|
params: CutlassMoEParams,
|
||||||
apply_router_weight_on_input: bool = False,
|
apply_router_weight_on_input: bool = False,
|
||||||
|
no_combine: bool = False,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
MoE implementation for FP4 Inputs
|
MoE implementation for FP4 Inputs
|
||||||
@@ -492,6 +493,8 @@ def cutlass_moe_fp4(
|
|||||||
del int_fp4, int_blockscale
|
del int_fp4, int_blockscale
|
||||||
c2 = shuffle_rows(c2, c_map, (m_a * num_topk, params.hidden_size))
|
c2 = shuffle_rows(c2, c_map, (m_a * num_topk, params.hidden_size))
|
||||||
c2 = c2.view(m_a, num_topk, params.hidden_size)
|
c2 = c2.view(m_a, num_topk, params.hidden_size)
|
||||||
|
if no_combine:
|
||||||
|
return c2.to(out_dtype)
|
||||||
if not apply_router_weight_on_input:
|
if not apply_router_weight_on_input:
|
||||||
c2 = c2 * topk_weights.view(m_a, num_topk, 1).to(out_dtype)
|
c2 = c2 * topk_weights.view(m_a, num_topk, 1).to(out_dtype)
|
||||||
return c2.sum(dim=1).to(out_dtype)
|
return c2.sum(dim=1).to(out_dtype)
|
||||||
|
|||||||
@@ -2175,6 +2175,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
|
|||||||
topk_ids=topk_ids,
|
topk_ids=topk_ids,
|
||||||
params=layer.cutlass_moe_params,
|
params=layer.cutlass_moe_params,
|
||||||
apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input,
|
apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input,
|
||||||
|
no_combine=moe_runner_config.no_combine,
|
||||||
).to(x.dtype)
|
).to(x.dtype)
|
||||||
# Scale by routed_scaling_factor is fused into select_experts.
|
# Scale by routed_scaling_factor is fused into select_experts.
|
||||||
return StandardCombineInput(hidden_states=output)
|
return StandardCombineInput(hidden_states=output)
|
||||||
|
|||||||
Reference in New Issue
Block a user