From d539de1c5d1a120ba3281d62c3bcb0e19a36348d Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:56:07 +0100 Subject: [PATCH] [Docs] Add missing return annotations flagged by griffe (#54980) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- vllm/_aiter_ops.py | 8 ++++++-- vllm/_custom_ops.py | 8 ++++---- vllm/model_executor/determinism/batch_invariant.py | 2 +- .../layers/fused_moe/experts/rocm_aiter_moe.py | 7 +++++++ 4 files changed, 18 insertions(+), 7 deletions(-) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 9e1fb6fd88..e865ea01b9 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -4,6 +4,7 @@ import ctypes import functools import os from collections.abc import Callable +from typing import TYPE_CHECKING import torch from torch._ops import OpOverload @@ -18,6 +19,9 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( rocm_aiter_sparse_attn_indexer_fake, ) +if TYPE_CHECKING: + from aiter import ActivationType, QuantType + logger = init_logger(__name__) try: @@ -1790,7 +1794,7 @@ class rocm_aiter_ops: cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM @staticmethod - def get_aiter_activation_type(activation_str: str): + def get_aiter_activation_type(activation_str: str) -> "ActivationType | None": """ Given an activation type as a string, returns the corresponding aiter ActivationType enum. Supported activation types: "no", "none", "silu", "gelu", "swiglu". @@ -1823,7 +1827,7 @@ class rocm_aiter_ops: return mapping.get(name) @staticmethod - def get_aiter_quant_type(quant_type_str: str): + def get_aiter_quant_type(quant_type_str: str) -> "QuantType | None": """ Given a quantization type as a string, returns the corresponding aiter QuantType enum. Supported quantization types: "no", "per_tensor", "per_token", "per_1x32", "per_1x128", "per_128x128". diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 8c8d28b04b..04b8b829f5 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -1508,7 +1508,7 @@ def cutlass_w4a8_moe_mm( c_strides: torch.Tensor, group_scale_strides: torch.Tensor, maybe_schedule: str | None = None, -): +) -> None: """ Executes the CUTLASS-based fused-MoE grouped matrix multiplication for the W4A8 quantization scheme. Uses group-wise quantization (INT4 -> FP8) @@ -2065,7 +2065,7 @@ def scaled_int8_quant( symmetric: Whether to use symmetric quantization (scale only, azp ignored). Returns: - tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp. + Output int8 tensor, scales, and optionally azp. """ if current_platform.is_xpu(): # XPU has no _C int8 quant op; use the torch.compile reference. @@ -3860,7 +3860,7 @@ def onednn_scaled_int8_quant( scale: torch.Tensor | None = None, azp: torch.Tensor | None = None, symmetric: bool = True, -): +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: """ Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp. @@ -3873,7 +3873,7 @@ def onednn_scaled_int8_quant( symmetric: Whether to use symmetric quantization (scale only, azp ignored). Returns: - tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp. + Output int8 tensor, scales, and optionally azp. """ output = torch.empty_like(input, dtype=torch.int8) token_num = input.numel() // input.shape[-1] diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py index e2579f30b3..6c701cfa16 100644 --- a/vllm/model_executor/determinism/batch_invariant.py +++ b/vllm/model_executor/determinism/batch_invariant.py @@ -221,7 +221,7 @@ def matmul_kernel_descriptor_persistent( def matmul_descriptor_persistent( a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None -): +) -> torch.Tensor: """Persistent matmul using tensor descriptors (Intel XPU fast path). Args: diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index a433b18487..066aba97cd 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import IntEnum from functools import lru_cache +from typing import TYPE_CHECKING import torch @@ -29,6 +30,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kMxfp4Static, ) +if TYPE_CHECKING: + from aiter import ActivationType + class QuantMethod(IntEnum): # This allows interfacing with AITER QuantType Enum @@ -253,6 +257,7 @@ def rocm_aiter_fused_experts( # Gate/up interleave hint; only the SWIGLUOAI activations override it. activation_interleave = None + activation_method: ActivationMethod | ActivationType if activation == MoEActivation.SILU: activation_method = ActivationMethod.SILU elif activation == MoEActivation.GELU: @@ -266,6 +271,8 @@ def rocm_aiter_fused_experts( activation_method = rocm_aiter_ops.get_aiter_activation_type("situ") else: raise ValueError(f"Unsupported activation: {activation}") + if activation_method is None: + raise ValueError(f"AITER does not support activation: {activation}") # All AITER Fused MoE kernels are expecting the following datatypes topk_weights = topk_weights.to(torch.float32)