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 
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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
    @property
    def activation_format(self) -> mk.FusedMoEActivationFormat:
        return mk.FusedMoEActivationFormat.Standard

    def max_num_tokens_per_rank(self) -> int | None:
        return None

    def topk_indices_dtype(self) -> torch.dtype | None:
        return None

    def num_dispatchers(self) -> int:
        return 1

    def output_is_reduced(self) -> bool:
        return False

    def prepare(
        self,
        a1: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        quant_config: FusedMoEQuantConfig,
    ) -> mk.PrepareResultType:
        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,
            quant_config.a1_scale,
            quant_config.quant_dtype,
            quant_config.per_act_token_quant,
            quant_config.block_shape,
        )

        return a1q, a1q_scale, None, None, 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,
        weight_and_reduce_impl: mk.TopKWeightAndReduce,
    ) -> None:
        if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
            weight_and_reduce_impl = TopKWeightAndReduceContiguous()
        weight_and_reduce_impl.apply(
            output=output,
            fused_expert_output=fused_expert_output,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            apply_router_weight_on_input=apply_router_weight_on_input,
        )

activation_format property 

activation_format: FusedMoEActivationFormat

finalize 

finalize(
    output: Tensor,
    fused_expert_output: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    apply_router_weight_on_input: bool,
    weight_and_reduce_impl: TopKWeightAndReduce,
) -> None
Source code in vllm/model_executor/layers/fused_moe/prepare_finalize.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
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,
    weight_and_reduce_impl: mk.TopKWeightAndReduce,
) -> None:
    if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
        weight_and_reduce_impl = TopKWeightAndReduceContiguous()
    weight_and_reduce_impl.apply(
        output=output,
        fused_expert_output=fused_expert_output,
        topk_weights=topk_weights,
        topk_ids=topk_ids,
        apply_router_weight_on_input=apply_router_weight_on_input,
    )

max_num_tokens_per_rank 

max_num_tokens_per_rank() -> int | None
Source code in vllm/model_executor/layers/fused_moe/prepare_finalize.py 
20
21
def max_num_tokens_per_rank(self) -> int | None:
    return None

num_dispatchers 

num_dispatchers() -> int
Source code in vllm/model_executor/layers/fused_moe/prepare_finalize.py 
26
27
def num_dispatchers(self) -> int:
    return 1

output_is_reduced 

output_is_reduced() -> bool
Source code in vllm/model_executor/layers/fused_moe/prepare_finalize.py 
29
30
def output_is_reduced(self) -> bool:
    return False

prepare 

prepare(
    a1: Tensor,
    topk_weights: Tensor,
    topk_ids: Tensor,
    num_experts: int,
    expert_map: Tensor | None,
    apply_router_weight_on_input: bool,
    quant_config: FusedMoEQuantConfig,
) -> PrepareResultType
Source code in vllm/model_executor/layers/fused_moe/prepare_finalize.py 
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
def prepare(
    self,
    a1: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    num_experts: int,
    expert_map: torch.Tensor | None,
    apply_router_weight_on_input: bool,
    quant_config: FusedMoEQuantConfig,
) -> mk.PrepareResultType:
    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,
        quant_config.a1_scale,
        quant_config.quant_dtype,
        quant_config.per_act_token_quant,
        quant_config.block_shape,
    )

    return a1q, a1q_scale, None, None, None

topk_indices_dtype 

topk_indices_dtype() -> dtype | None
Source code in vllm/model_executor/layers/fused_moe/prepare_finalize.py 
23
24
def topk_indices_dtype(self) -> torch.dtype | None:
    return None