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, 17:15 UTC
[TPU] Avoid Triton Import (#15589)
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
Robert Shaw committed
Mar 27, 2025 at 02:43 UTC
e1e0fd7543d6759ac15615717bd904f64e7137ae
2 files changed
+8
-6
vllm/model_executor/layers/fused_moe/layer.py
+3
-3
@@ -16,8 +16,6 @@ from vllm.distributed import (get_dp_group, get_tensor_model_parallel_rank,
16
from vllm.forward_context import ForwardContext, get_forward_context
17
from vllm.logger import init_logger
18
from vllm.model_executor.custom_op import CustomOp
19
-from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
20
- is_rocm_aiter_moe_enabled, shuffle_weights)
19
from vllm.model_executor.layers.quantization.base_config import (
20
QuantizationConfig, QuantizeMethodBase)
21
from vllm.model_executor.utils import set_weight_attrs
@@ -119,7 +117,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
117
layer.w2_weight = torch.nn.Parameter(self._maybe_pad_weight(
118
layer.w2_weight.data),
119
requires_grad=False)
122
-
120
+ # Lazy import to avoid importing triton.
121
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
122
+ is_rocm_aiter_moe_enabled, shuffle_weights)
123
if is_rocm_aiter_moe_enabled():
124
# reshaping weights is required for aiter moe kernel.
125
shuffled_w13, shuffled_w2 = shuffle_weights(
vllm/model_executor/layers/quantization/fp8.py
+5
-3
@@ -13,9 +13,6 @@ from vllm.distributed import get_tensor_model_parallel_world_size
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEMethodBase,
15
FusedMoeWeightScaleSupported)
16
-from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
17
- expand_weights, is_rocm_aiter_block_scaled_moe_enabled,
18
- is_rocm_aiter_moe_enabled, shuffle_weights)
16
from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
17
UnquantizedLinearMethod)
18
from vllm.model_executor.layers.quantization.base_config import (
@@ -532,6 +529,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):
529
layer.w2_input_scale = None
530
531
def process_weights_after_loading(self, layer: Module) -> None:
532
+ # Lazy import to avoid importing triton too early.
533
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
534
+ expand_weights, is_rocm_aiter_block_scaled_moe_enabled,
535
+ is_rocm_aiter_moe_enabled, shuffle_weights)
536
+
537
# TODO (rob): refactor block quant into separate class.
538
if self.block_quant:
539
assert self.quant_config.activation_scheme == "dynamic"