A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[ROCM] enable aiter fused moe kernel for llama4 bf16 checkpoints (#16674)
Sijia(Jackson) Chen committed
Apr 17, 2025 at 11:44 UTC
92edf35826e7067bfcba9ae2f0ae2cacd3109cb2
1 file changed
+13
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+13
@@ -26,6 +26,7 @@ def rocm_aiter_fused_experts(
26
topk_weights: torch.Tensor,
27
topk_ids: torch.Tensor,
28
use_fp8_w8a8: bool = False,
29
+ apply_router_weight_on_input: bool = False,
30
w1_scale: Optional[torch.Tensor] = None,
31
w2_scale: Optional[torch.Tensor] = None,
32
block_shape: Optional[List[int]] = None,
@@ -39,6 +40,18 @@ def rocm_aiter_fused_experts(
40
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
41
per_token_group_quant_fp8)
42
43
+ if apply_router_weight_on_input:
44
+ assert (topk_weights.dim() == 2
45
+ ), "`topk_weights` should be in shape (num_tokens, topk)"
46
+ _, topk = topk_weights.shape
47
+ assert (
48
+ topk == 1
49
+ ), "Only support topk=1 when `apply_router_weight_on_input` is True"
50
+
51
+ hidden_states = hidden_states * topk_weights.to(hidden_states.dtype)
52
+ topk_ids = topk_ids.to(torch.int32)
53
+ topk_weights = torch.ones_like(topk_weights, dtype=torch.float32)
54
+
55
if envs.VLLM_ROCM_USE_AITER_FP8_BLOCK_SCALED_MOE and use_fp8_w8a8:
56
assert w1_scale is not None
57
assert w2_scale is not None