@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d539de1c5d

[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>

Harry Mellor committed Sep 2, 2026 at 17:56 UTC d539de1c5d1a120ba3281d62c3bcb0e19a36348d

4 files changed +18 -7

vllm/_aiter_ops.py

+6 -2

@@ -4,6 +4,7 @@ import ctypes

4 import functools

5 import os

6 from collections.abc import Callable

7 +from typing import TYPE_CHECKING

9 import torch

10 from torch._ops import OpOverload

@@ -18,6 +19,9 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (

19 rocm_aiter_sparse_attn_indexer_fake,

20 )

22 +if TYPE_CHECKING:

23 + from aiter import ActivationType, QuantType

24 +

25 logger = init_logger(__name__)

27 try:

@@ -1790,7 +1794,7 @@ class rocm_aiter_ops:

1794 cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM

1795

1796 @staticmethod

1793 - def get_aiter_activation_type(activation_str: str):

1797 + def get_aiter_activation_type(activation_str: str) -> "ActivationType | None":

1798 """

1799 Given an activation type as a string, returns the corresponding aiter ActivationType enum.

1800 Supported activation types: "no", "none", "silu", "gelu", "swiglu".

@@ -1823,7 +1827,7 @@ class rocm_aiter_ops:

1827 return mapping.get(name)

1828

1829 @staticmethod

1826 - def get_aiter_quant_type(quant_type_str: str):

1830 + def get_aiter_quant_type(quant_type_str: str) -> "QuantType | None":

1831 """

1832 Given a quantization type as a string, returns the corresponding aiter QuantType enum.

1833 Supported quantization types: "no", "per_tensor", "per_token", "per_1x32", "per_1x128", "per_128x128".

vllm/_custom_ops.py

+4 -4

@@ -1508,7 +1508,7 @@ def cutlass_w4a8_moe_mm(

1508 c_strides: torch.Tensor,

1509 group_scale_strides: torch.Tensor,

1510 maybe_schedule: str | None = None,

1511 -):

1511 +) -> None:

1512 """

1513 Executes the CUTLASS-based fused-MoE grouped matrix multiplication for the

1514 W4A8 quantization scheme. Uses group-wise quantization (INT4 -> FP8)

@@ -2065,7 +2065,7 @@ def scaled_int8_quant(

2065 symmetric: Whether to use symmetric quantization (scale only, azp ignored).

2066

2067 Returns:

2068 - tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp.

2068 + Output int8 tensor, scales, and optionally azp.

2069 """

2070 if current_platform.is_xpu():

2071 # XPU has no _C int8 quant op; use the torch.compile reference.

@@ -3860,7 +3860,7 @@ def onednn_scaled_int8_quant(

3860 scale: torch.Tensor | None = None,

3861 azp: torch.Tensor | None = None,

3862 symmetric: bool = True,

3863 -):

3863 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

3864 """

3865 Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.

3866

@@ -3873,7 +3873,7 @@ def onednn_scaled_int8_quant(

3873 symmetric: Whether to use symmetric quantization (scale only, azp ignored).

3874

3875 Returns:

3876 - tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp.

3876 + Output int8 tensor, scales, and optionally azp.

3877 """

3878 output = torch.empty_like(input, dtype=torch.int8)

3879 token_num = input.numel() // input.shape[-1]

vllm/model_executor/determinism/batch_invariant.py

+1 -1

@@ -221,7 +221,7 @@ def matmul_kernel_descriptor_persistent(

221

222 def matmul_descriptor_persistent(

223 a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None

224 -):

224 +) -> torch.Tensor:

225 """Persistent matmul using tensor descriptors (Intel XPU fast path).

226

227 Args:

vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from enum import IntEnum

4 from functools import lru_cache

5 +from typing import TYPE_CHECKING

7 import torch

@@ -29,6 +30,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

30 kMxfp4Static,

31 )

33 +if TYPE_CHECKING:

34 + from aiter import ActivationType

35 +

37 class QuantMethod(IntEnum):

38 # This allows interfacing with AITER QuantType Enum

@@ -253,6 +257,7 @@ def rocm_aiter_fused_experts(

257

258 # Gate/up interleave hint; only the SWIGLUOAI activations override it.

259 activation_interleave = None

260 + activation_method: ActivationMethod | ActivationType

261 if activation == MoEActivation.SILU:

262 activation_method = ActivationMethod.SILU

263 elif activation == MoEActivation.GELU:

@@ -266,6 +271,8 @@ def rocm_aiter_fused_experts(

271 activation_method = rocm_aiter_ops.get_aiter_activation_type("situ")

272 else:

273 raise ValueError(f"Unsupported activation: {activation}")

274 + if activation_method is None:

275 + raise ValueError(f"AITER does not support activation: {activation}")

276

277 # All AITER Fused MoE kernels are expecting the following datatypes

278 topk_weights = topk_weights.to(torch.float32)