| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Gates zentorch CPU linear dispatch on platform/op availability.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import logging |
| 8 | from typing import TYPE_CHECKING |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from vllm.platforms import current_platform |
| 13 | |
| 14 | if TYPE_CHECKING: |
| 15 | from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig |
| 16 | |
| 17 | __all__ = [ |
| 18 | "has_zentorch_op", |
| 19 | "is_zentorch_moe_config_supported", |
| 20 | "is_zentorch_moe_supported", |
| 21 | "_ZENTORCH_MOE_ACTIVATIONS", |
| 22 | ] |
| 23 | |
| 24 | |
| 25 | _ZENTORCH_MOE_ACTIVATIONS = frozenset({"gelu", "gelu_tanh", "silu", "swigluoai"}) |
| 26 | |
| 27 | |
| 28 | def _moe_activation_to_str(activation: object) -> str: |
| 29 | """Normalize activation to a lowercase string (enum-safe).""" |
| 30 | if hasattr(activation, "value"): |
| 31 | return str(activation.value).lower() |
| 32 | return str(activation).lower() |
| 33 | |
| 34 | |
| 35 | def has_zentorch_op(op_names: list[str]) -> bool: |
| 36 | """Return ``True`` when running on Zen CPU with all named ops registered.""" |
| 37 | if not op_names: |
| 38 | raise ValueError("has_zentorch_op requires at least one op name") |
| 39 | if not current_platform.is_zen_cpu(): |
| 40 | return False |
| 41 | ns = getattr(torch.ops, "zentorch", None) |
| 42 | if ns is None: |
| 43 | return False |
| 44 | return all(hasattr(ns, op_name) for op_name in op_names) |
| 45 | |
| 46 | |
| 47 | def is_zentorch_moe_config_supported(moe_config: FusedMoEConfig) -> bool: |
| 48 | if not has_zentorch_op(["zentorch_fused_moe"]): |
| 49 | return False |
| 50 | if not moe_config.is_act_and_mul: |
| 51 | return False |
| 52 | act = _moe_activation_to_str(moe_config.activation) |
| 53 | return act in _ZENTORCH_MOE_ACTIVATIONS |
| 54 | |
| 55 | |
| 56 | def is_zentorch_moe_supported(layer: torch.nn.Module) -> bool: |
| 57 | if not has_zentorch_op(["zentorch_fused_moe"]): |
| 58 | logging.debug( |
| 59 | "Skipping zentorch fused-MoE: not a Zen CPU or " |
| 60 | "zentorch not loaded; using default MoE." |
| 61 | ) |
| 62 | return False |
| 63 | moe_config = getattr(layer, "moe_config", None) |
| 64 | if moe_config is not None and not moe_config.is_act_and_mul: |
| 65 | logging.debug( |
| 66 | "Skipping zentorch fused-MoE: layer is not a gated " |
| 67 | "(act-and-mul, e.g. SwiGLU) MLP, the only structure supported." |
| 68 | ) |
| 69 | return False |
| 70 | activation = getattr(layer, "activation", None) |
| 71 | if activation is None: |
| 72 | logging.debug( |
| 73 | "Skipping zentorch fused-MoE: layer has no 'activation' " |
| 74 | "attribute, so the activation can't be verified." |
| 75 | ) |
| 76 | return False |
| 77 | act = str(getattr(activation, "value", activation)).lower() |
| 78 | if act not in _ZENTORCH_MOE_ACTIVATIONS: |
| 79 | logging.debug( |
| 80 | "Skipping zentorch fused-MoE: activation %r unsupported (supported: %s).", |
| 81 | act, |
| 82 | _ZENTORCH_MOE_ACTIVATIONS, |
| 83 | ) |
| 84 | return False |
| 85 | return True |