Skip to content

vllm.model_executor.layers.fused_moe.prepare_finalize

MoEPrepareAndFinalizeNoEP

Bases: FusedMoEPrepareAndFinalize

Source code in vllm/model_executor/layers/fused_moe/prepare_finalize.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

    def __init__(
        self,
        quant_dtype: Optional[torch.dtype] = None,
        per_channel_quant: bool = False,
        block_shape: Optional[list[int]] = None,
    ):
        super().__init__()
        self.per_channel_quant = per_channel_quant
        self.block_shape = block_shape
        self.quant_dtype = quant_dtype

    def prepare(
        self,
        a1: torch.Tensor,
        a1_scale: Optional[torch.Tensor],
        a2_scale: Optional[torch.Tensor],
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        num_experts: int,
        expert_map: Optional[torch.Tensor],
        apply_router_weight_on_input: bool = False,
    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
        if apply_router_weight_on_input:
            topk = topk_ids.size(1)
            # TODO: this only works for topK=1, will need to update for topK>1
            assert topk == 1, \
                "apply_router_weight_on_input is only implemented for topk=1"
            a1.mul_(topk_weights.to(a1.dtype))

        a1q, a1q_scale = moe_kernel_quantize_input(a1, a1_scale,
                                                   self.quant_dtype,
                                                   self.per_channel_quant,
                                                   self.block_shape)

        return a1q, a1q_scale, None

    def finalize(
        self,
        output: torch.Tensor,
        fused_expert_output: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        apply_router_weight_on_input: bool,
    ) -> None:
        _moe_unpermute_and_reduce(output, fused_expert_output, None,
                                  topk_weights, apply_router_weight_on_input)

block_shape instance-attribute 

block_shape = block_shape

per_channel_quant instance-attribute 

per_channel_quant = per_channel_quant

quant_dtype instance-attribute 

quant_dtype = quant_dtype

__init__ 

__init__(
    quant_dtype: Optional[dtype] = None,
    per_channel_quant: bool = False,
    block_shape: Optional[list[int]] = None,
)
Source code in vllm/model_executor/layers/fused_moe/prepare_finalize.py 
15
16
17
18
19
20
21
22
23
24
def __init__(
    self,
    quant_dtype: Optional[torch.dtype] = None,
    per_channel_quant: bool = False,
    block_shape: Optional[list[int]] = None,
):
    super().__init__()
    self.per_channel_quant = per_channel_quant
    self.block_shape = block_shape
    self.quant_dtype = quant_dtype

finalize 

finalize(
    output: Tensor,
    fused_expert_output: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    apply_router_weight_on_input: bool,
) -> None
Source code in vllm/model_executor/layers/fused_moe/prepare_finalize.py 
51
52
53
54
55
56
57
58
59
60
def finalize(
    self,
    output: torch.Tensor,
    fused_expert_output: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    apply_router_weight_on_input: bool,
) -> None:
    _moe_unpermute_and_reduce(output, fused_expert_output, None,
                              topk_weights, apply_router_weight_on_input)

prepare 

prepare(
    a1: Tensor,
    a1_scale: Optional[Tensor],
    a2_scale: Optional[Tensor],
    topk_weights: Tensor,
    topk_ids: Tensor,
    num_experts: int,
    expert_map: Optional[Tensor],
    apply_router_weight_on_input: bool = False,
) -> tuple[Tensor, Optional[Tensor], Optional[Tensor]]
Source code in vllm/model_executor/layers/fused_moe/prepare_finalize.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def prepare(
    self,
    a1: torch.Tensor,
    a1_scale: Optional[torch.Tensor],
    a2_scale: Optional[torch.Tensor],
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    num_experts: int,
    expert_map: Optional[torch.Tensor],
    apply_router_weight_on_input: bool = False,
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
    if apply_router_weight_on_input:
        topk = topk_ids.size(1)
        # TODO: this only works for topK=1, will need to update for topK>1
        assert topk == 1, \
            "apply_router_weight_on_input is only implemented for topk=1"
        a1.mul_(topk_weights.to(a1.dtype))

    a1q, a1q_scale = moe_kernel_quantize_input(a1, a1_scale,
                                               self.quant_dtype,
                                               self.per_channel_quant,
                                               self.block_shape)

    return a1q, a1q_scale, None