[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) <noreply@anthropic.com>

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".

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]

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:

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)