[refactor] replace scattered business kwargs with typed request objects and explicit stage boundaries (#7024)
### What this PR does / why we need it? Refactor `vllm_ascend/ops/fused_moe` to replace scattered MoE business `**kwargs` with typed request objects and explicit stage boundaries. - Prepare, dispatch, MLP, and quant stages now have clearer ownership. - Main MoE path no longer depends on business `kwargs.get(...)` lookups. - Comm and dispatcher interfaces are request-only on the main path. - UTs can assert stage-level fields directly instead of inferring behavior indirectly. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? CI passed. --------- Signed-off-by: linfeng-yuan <1102311262@qq.com>
This commit is contained in:
@@ -18,6 +18,8 @@
|
||||
import torch
|
||||
import torch_npu
|
||||
|
||||
from vllm_ascend.ops.fused_moe.moe_runtime_args import MoEMlpComputeInput
|
||||
|
||||
|
||||
def quant_apply_mlp(
|
||||
hidden_states: torch.Tensor,
|
||||
@@ -66,17 +68,20 @@ def unquant_apply_mlp(
|
||||
return hidden_states
|
||||
|
||||
|
||||
def unified_apply_mlp(
|
||||
hidden_states: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
group_list: torch.Tensor,
|
||||
w1_scale: torch.Tensor | None = None,
|
||||
w2_scale: torch.Tensor | None = None,
|
||||
group_list_type: int = 1,
|
||||
with_quant: bool = False,
|
||||
) -> torch.Tensor:
|
||||
if with_quant:
|
||||
def unified_apply_mlp(*, mlp_compute_input: MoEMlpComputeInput) -> torch.Tensor:
|
||||
hidden_states = mlp_compute_input.hidden_states
|
||||
w1 = mlp_compute_input.weights.w1
|
||||
w2 = mlp_compute_input.weights.w2
|
||||
w1_scale = mlp_compute_input.weights.w1_scale
|
||||
w2_scale = mlp_compute_input.weights.w2_scale
|
||||
group_list = mlp_compute_input.group_list
|
||||
group_list_type = mlp_compute_input.group_list_type
|
||||
assert isinstance(w1, torch.Tensor)
|
||||
assert isinstance(w2, torch.Tensor)
|
||||
|
||||
if mlp_compute_input.quant.is_quant:
|
||||
assert isinstance(w1_scale, torch.Tensor)
|
||||
assert isinstance(w2_scale, torch.Tensor)
|
||||
assert w1_scale is not None and w2_scale is not None
|
||||
return quant_apply_mlp(
|
||||
hidden_states=hidden_states,
|
||||
@@ -87,7 +92,11 @@ def unified_apply_mlp(
|
||||
group_list=group_list,
|
||||
group_list_type=group_list_type,
|
||||
)
|
||||
else:
|
||||
return unquant_apply_mlp(
|
||||
hidden_states=hidden_states, w1=w1, w2=w2, group_list=group_list, group_list_type=group_list_type
|
||||
)
|
||||
|
||||
return unquant_apply_mlp(
|
||||
hidden_states=hidden_states,
|
||||
w1=w1,
|
||||
w2=w2,
|
||||
group_list=group_list,
|
||||
group_list_type=group_list_type,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user