CUDA TYPINGS

## Problem

Running short completions (like `max_tokens=1` health check probes) can
finish so fast that `mlx_lm`'s internal `generation_time` rounds to
zero. When that happens, `BatchGenerator.stats()` in
`mlx_lm/generate.py` divides `generation_tokens / generation_time` and
throws a `ZeroDivisionError`, which kills the runner process.

EXO already handles this on its side — lines 289-293 in
`batch_generate.py` guard the TPS calculation with `if
generation_elapsed > 0`. But the call to `self._exo_gen.stats()` on line
294 goes into mlx_lm's *separate* timing code, which doesn't have the
same guard. Two different timers, only one is protected.

In my case, this was triggered by health check probes (content: `"a"`,
`max_tokens=1`). The generation completed in sub-microsecond time,
`generation_time` was exactly `0`, and the runner crashed. Since the
health check command stays in the queue and retries after recovery, it
created an infinite crash loop — every ~15 seconds the runner would load
the model, get the same health check, and die again.

## Fix

Wrap `self._exo_gen.stats()` in a `try/except ZeroDivisionError`. If it
throws, set `mlx_stats` to `None` and fall back to `0.0` for
`prompt_tps`. The only fields EXO reads from `mlx_stats` are
`prompt_tps` and `prompt_time` — losing them on a sub-microsecond
generation has no practical impact.

## Traceback

```
File "batch_generate.py", line 294, in step
    mlx_stats = self._exo_gen.stats()
File "mlx_lm/generate.py", line 1224, in stats
    self._stats.generation_tokens / self._stats.generation_time
ZeroDivisionError: division by zero
```

.cuda_typings/openai_harmony/__init__.pyi

@@ -0,0 +1,20 @@
+from enum import Enum
+
+class HarmonyEncodingName(Enum):
+    HARMONY_GPT_OSS = ...
+
+class HarmonyEncoding: ...
+class HarmonyError(Exception): ...
+
+class Role(Enum):
+    ASSISTANT = ...
+
+class StreamableParser:
+    last_content_delta: str
+    current_channel: str | None
+    current_recipient: str | None
+
+    def __init__(self, encoding: HarmonyEncoding, role: Role = ...) -> None: ...
+    def process(self, token_id: int) -> None: ...
+
+def load_harmony_encoding(name: HarmonyEncodingName) -> HarmonyEncoding: ...

.cuda_typings/torch/__init__.pyi

@@ -0,0 +1,8 @@
+from torch import backends as backends
+from torch import cuda as cuda
+from torch import distributed as distributed
+
+__version__: str
+
+class version:
+    cuda: str

.cuda_typings/torch/backends/__init__.pyi

@@ -0,0 +1 @@
+from torch.backends import cuda as cuda

.cuda_typings/torch/backends/cuda/__init__.pyi

@@ -0,0 +1 @@
+def is_built() -> bool: ...

.cuda_typings/torch/cuda/__init__.pyi

@@ -0,0 +1,8 @@
+class _DeviceProperties:
+    total_memory: int
+
+def is_available() -> bool: ...
+def get_device_name(device: int) -> str: ...
+def get_device_properties(device: int) -> _DeviceProperties: ...
+def empty_cache() -> None: ...
+def mem_get_info() -> tuple[int, int]: ...

.cuda_typings/torch/distributed/__init__.pyi

@@ -0,0 +1,2 @@
+def is_initialized() -> bool: ...
+def destroy_process_group() -> None: ...

.cuda_typings/vllm/__init__.pyi

@@ -0,0 +1,56 @@
+from .version import __version__ as __version__, __version_tuple__ as __version_tuple__
+from vllm.engine.arg_utils import (
+    AsyncEngineArgs as AsyncEngineArgs,
+    EngineArgs as EngineArgs,
+)
+from vllm.engine.async_llm_engine import AsyncLLMEngine as AsyncLLMEngine
+from vllm.engine.llm_engine import LLMEngine as LLMEngine
+from vllm.entrypoints.llm import LLM as LLM
+from vllm.inputs import (
+    PromptType as PromptType,
+    TextPrompt as TextPrompt,
+    TokensPrompt as TokensPrompt,
+)
+from vllm.model_executor.models import ModelRegistry as ModelRegistry
+from vllm.outputs import (
+    ClassificationOutput as ClassificationOutput,
+    ClassificationRequestOutput as ClassificationRequestOutput,
+    CompletionOutput as CompletionOutput,
+    EmbeddingOutput as EmbeddingOutput,
+    EmbeddingRequestOutput as EmbeddingRequestOutput,
+    PoolingOutput as PoolingOutput,
+    PoolingRequestOutput as PoolingRequestOutput,
+    RequestOutput as RequestOutput,
+    ScoringOutput as ScoringOutput,
+    ScoringRequestOutput as ScoringRequestOutput,
+)
+from vllm.pooling_params import PoolingParams as PoolingParams
+from vllm.sampling_params import SamplingParams as SamplingParams
+from vllm.v1.executor.ray_utils import initialize_ray_cluster as initialize_ray_cluster
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "LLM",
+    "ModelRegistry",
+    "PromptType",
+    "TextPrompt",
+    "TokensPrompt",
+    "SamplingParams",
+    "RequestOutput",
+    "CompletionOutput",
+    "PoolingOutput",
+    "PoolingRequestOutput",
+    "EmbeddingOutput",
+    "EmbeddingRequestOutput",
+    "ClassificationOutput",
+    "ClassificationRequestOutput",
+    "ScoringOutput",
+    "ScoringRequestOutput",
+    "LLMEngine",
+    "EngineArgs",
+    "AsyncLLMEngine",
+    "AsyncEngineArgs",
+    "initialize_ray_cluster",
+    "PoolingParams",
+]

.cuda_typings/vllm/_aiter_ops.pyi

@@ -0,0 +1,341 @@
+import torch
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable
+from torch._ops import OpOverload as OpOverload
+from vllm.platforms import current_platform as current_platform
+from vllm.utils.torch_utils import (
+    direct_register_custom_op as direct_register_custom_op,
+)
+from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+    rocm_aiter_sparse_attn_indexer as rocm_aiter_sparse_attn_indexer,
+    rocm_aiter_sparse_attn_indexer_fake as rocm_aiter_sparse_attn_indexer_fake,
+)
+
+FP8_DTYPE: Incomplete
+
+def is_aiter_found() -> bool: ...
+
+IS_AITER_FOUND: Incomplete
+
+def is_aiter_found_and_supported() -> bool: ...
+def if_aiter_supported(func: Callable) -> Callable: ...
+
+class rocm_aiter_ops:
+    @classmethod
+    def refresh_env_variables(cls) -> None: ...
+    @staticmethod
+    def get_aiter_activation_type(activation_str: str): ...
+    @staticmethod
+    def get_aiter_quant_type(quant_type_str: str): ...
+    @classmethod
+    @if_aiter_supported
+    def is_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_linear_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_linear_fp8_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_rmsnorm_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_fused_moe_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_fusion_moe_shared_experts_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_mla_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_mha_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_shuffle_kv_cache_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_triton_unified_attn_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_fp8bmm_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_fp4bmm_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_asm_fp4_gemm_dynamic_quant_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_triton_rotary_embed_enabled(cls) -> bool: ...
+    @classmethod
+    @if_aiter_supported
+    def is_triton_gemm_enabled(cls) -> bool: ...
+    @staticmethod
+    @if_aiter_supported
+    def register_ops_once() -> None: ...
+    @staticmethod
+    def get_rmsnorm_fused_add_op() -> OpOverload: ...
+    @staticmethod
+    def get_rmsnorm_op() -> OpOverload: ...
+    @staticmethod
+    def get_rmsnorm_fused_add_dynamic_quant_op() -> OpOverload: ...
+    @staticmethod
+    def get_rmsnorm_fused_dynamic_quant_op() -> OpOverload: ...
+    @staticmethod
+    def get_rmsnorm_group_fused_quant_op() -> OpOverload: ...
+    @staticmethod
+    def get_rmsnorm_group_add_fused_quant_op() -> OpOverload: ...
+    @staticmethod
+    def get_per_token_quant_op() -> OpOverload: ...
+    @staticmethod
+    def get_group_quant_op() -> OpOverload: ...
+    @staticmethod
+    def get_act_mul_fused_fp8_group_quant_op() -> OpOverload: ...
+    @staticmethod
+    def get_triton_add_rmsnorm_pad_op() -> OpOverload: ...
+    @staticmethod
+    def get_triton_rotary_embedding_op() -> OpOverload: ...
+    @staticmethod
+    def rms_norm(
+        x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
+    ) -> torch.Tensor: ...
+    @staticmethod
+    def rms_norm2d_with_add(
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        weight: torch.Tensor,
+        variance_epsilon: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+    @staticmethod
+    def gemm_a8w8(
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None = None,
+        output_dtype: torch.dtype = ...,
+    ) -> torch.Tensor: ...
+    @staticmethod
+    def triton_gemm_a8w8_blockscale(
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        block_size: list[int],
+        output_dtype: torch.dtype = ...,
+    ) -> torch.Tensor: ...
+    @staticmethod
+    def gemm_a8w8_blockscale(
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        block_size: list[int],
+        output_dtype: torch.dtype = ...,
+    ) -> torch.Tensor: ...
+    @staticmethod
+    def fused_moe(
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weight: torch.Tensor,
+        topk_ids: torch.Tensor,
+        expert_mask: torch.Tensor | None = None,
+        activation_method: int = 0,
+        quant_method: int = 0,
+        doweight_stage1: bool = False,
+        w1_scale: torch.Tensor | None = None,
+        w2_scale: torch.Tensor | None = None,
+        a1_scale: torch.Tensor | None = None,
+        a2_scale: torch.Tensor | None = None,
+        num_local_tokens: torch.Tensor | None = None,
+        output_dtype: torch.dtype | None = None,
+        hidden_pad: int = 0,
+        intermediate_pad: int = 0,
+        bias1: torch.Tensor | None = None,
+        bias2: torch.Tensor | None = None,
+    ) -> torch.Tensor: ...
+    @staticmethod
+    def asm_moe_tkw1(
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        fc1_scale: torch.Tensor | None = None,
+        fc2_scale: torch.Tensor | None = None,
+        fc1_smooth_scale: torch.Tensor | None = None,
+        fc2_smooth_scale: torch.Tensor | None = None,
+        a16: bool = False,
+        per_tensor_quant_scale: torch.Tensor | None = None,
+        expert_mask: torch.Tensor | None = None,
+        activation_method: int = 0,
+    ) -> torch.Tensor: ...
+    @staticmethod
+    def topk_softmax(
+        topk_weights: torch.Tensor,
+        topk_indices: torch.Tensor,
+        token_expert_indices: torch.Tensor,
+        gating_output: torch.Tensor,
+        renormalize: bool,
+    ) -> tuple[torch.Tensor, ...]: ...
+    @staticmethod
+    def topk_sigmoid(
+        topk_weights: torch.Tensor,
+        topk_indices: torch.Tensor,
+        token_expert_indices: torch.Tensor,
+        gating_output: torch.Tensor,
+        renormalize: bool,
+    ) -> tuple[torch.Tensor, ...]: ...
+    @staticmethod
+    def biased_grouped_topk(
+        gating_output: torch.Tensor,
+        correction_bias: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        num_expert_group: int,
+        topk_group: int,
+        need_renorm: bool,
+        routed_scaling_factor: float = 1.0,
+    ) -> None: ...
+    @staticmethod
+    def grouped_topk(
+        gating_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        num_expert_group: int,
+        topk_group: int,
+        need_renorm: bool,
+        scoring_func: str = "softmax",
+        routed_scaling_factor: float = 1.0,
+    ) -> None: ...
+    @staticmethod
+    def fused_topk(
+        x: torch.Tensor, router_logits: torch.Tensor, top_k: int, gate_up: bool
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+    @staticmethod
+    def mla_decode_fwd(
+        q: torch.Tensor,
+        kv_buffer: torch.Tensor,
+        o: torch.Tensor,
+        sm_scale: float,
+        qo_indptr: torch.Tensor,
+        max_seqlen_qo: int,
+        kv_indptr: torch.Tensor | None = None,
+        kv_indices: torch.Tensor | None = None,
+        kv_last_page_lens: torch.Tensor | None = None,
+        logit_cap: float = 0.0,
+        q_scale: torch.Tensor | None = None,
+        kv_scale: torch.Tensor | None = None,
+    ): ...
+    @staticmethod
+    def per_tensor_quant(
+        x: torch.Tensor, quant_dtype: torch.dtype, scale: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+    @staticmethod
+    def per_token_quant(
+        x: torch.Tensor, quant_dtype: torch.dtype, scale: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+    @staticmethod
+    def triton_fp4_gemm_dynamic_qaunt(
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        weight_scale: torch.Tensor,
+        out_dtype: torch.dtype | None = ...,
+        x_scales: torch.Tensor | None = None,
+    ) -> torch.Tensor: ...
+    @staticmethod
+    def triton_rope_and_cache(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        positions: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        is_neox: bool,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        layer_slot_mapping: torch.Tensor,
+        k_scale: torch.Tensor,
+        v_scale: torch.Tensor,
+        flash_layout: bool,
+        apply_scale: bool,
+    ): ...
+    @staticmethod
+    def batched_gemm_a16wfp4(
+        X: torch.Tensor,
+        W: torch.Tensor,
+        w_scale: torch.Tensor,
+        Y: torch.Tensor,
+        transpose_bm: bool | None = False,
+        prequant: bool | None = False,
+        y_scale: torch.Tensor | None = None,
+    ) -> torch.Tensor: ...
+    @staticmethod
+    def triton_fp8_bmm(
+        X: torch.Tensor,
+        WQ: torch.Tensor,
+        w_scale: torch.Tensor,
+        group_size: int = 128,
+        bias: torch.Tensor | None = None,
+        dtype: torch.dtype | None = ...,
+        splitK: int | None = None,
+        YQ: torch.Tensor | None = None,
+        transpose_bm: bool | None = False,
+        config: dict | None = None,
+    ) -> torch.Tensor: ...
+    @staticmethod
+    def group_fp8_quant(
+        input_2d: torch.Tensor, group_size: int = 128
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+    @staticmethod
+    def is_triton_gemm_w8a8_tuned(n: int, k: int) -> bool: ...
+    @staticmethod
+    def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k: int) -> bool: ...
+    @staticmethod
+    def shuffle_weight(
+        self, tensor: torch.Tensor, layout: tuple[int, int] = (16, 16)
+    ) -> torch.Tensor: ...
+    @staticmethod
+    def shuffle_weight_a16w4(
+        tensor: torch.Tensor, nLane: int, gate_up: bool
+    ) -> torch.Tensor: ...
+    @staticmethod
+    def shuffle_scale_a16w4(
+        tensor: torch.Tensor, num_experts: int, gate_up: bool
+    ) -> torch.Tensor: ...
+    @staticmethod
+    def shuffle_weights(
+        *tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)
+    ) -> tuple[torch.Tensor, ...]: ...
+    @staticmethod
+    def flash_attn_varlen_func(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        min_seqlen_q: int | None = None,
+        dropout_p: float = 0.0,
+        softmax_scale: float | None = None,
+        causal: bool = False,
+        window_size: tuple[int, int] | None = None,
+        alibi_slopes: torch.Tensor | None = None,
+        return_lse: bool = False,
+        out: torch.Tensor | None = None,
+    ): ...
+    @staticmethod
+    def pa_fwd_asm(
+        Q: torch.Tensor,
+        K: torch.Tensor,
+        V: torch.Tensor,
+        block_tables: torch.Tensor,
+        context_lens: torch.Tensor,
+        block_tables_stride0: int,
+        K_QScale: torch.Tensor,
+        V_QScale: torch.Tensor,
+        out_: torch.Tensor,
+    ): ...

.cuda_typings/vllm/_oink_ops.pyi

@@ -0,0 +1,14 @@
+import torch
+from collections.abc import Callable as Callable
+from torch._dynamo import disable as _dynamo_disable
+
+@_dynamo_disable
+def is_oink_available_for_device(device_index: int) -> bool: ...
+def has_fused_add_rms_norm() -> bool: ...
+def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: ...
+def fused_add_rms_norm_(
+    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
+) -> None: ...
+def fused_add_rms_norm(
+    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
+) -> tuple[torch.Tensor, torch.Tensor]: ...

.cuda_typings/vllm/_version.pyi

@@ -0,0 +1,17 @@
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+VERSION_TUPLE = tuple[int | str, ...]
+COMMIT_ID = str | None
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID

.cuda_typings/vllm/_xpu_ops.pyi

@@ -0,0 +1,59 @@
+import torch
+from _typeshed import Incomplete
+from vllm.logger import init_logger as init_logger
+
+logger: Incomplete
+
+def register_fake(fn): ...
+
+class xpu_ops:
+    @staticmethod
+    def flash_attn_varlen_func(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        softmax_scale: float | None = None,
+        causal: bool = False,
+        out: torch.Tensor | None = None,
+        block_table: torch.Tensor | None = None,
+        alibi_slopes: torch.Tensor | None = None,
+        window_size: list[int] | None = None,
+        softcap: float | None = 0.0,
+        seqused_k: torch.Tensor | None = None,
+        cu_seqlens_k: torch.Tensor | None = None,
+        dropout_p: float = 0.0,
+        scheduler_metadata=None,
+        fa_version: int = 2,
+        q_descale=None,
+        k_descale=None,
+        v_descale=None,
+        num_splits: int = 0,
+        return_softmax_lse: bool | None = False,
+        s_aux: torch.Tensor | None = None,
+    ): ...
+    @staticmethod
+    def get_scheduler_metadata(
+        batch_size,
+        max_seqlen_q,
+        max_seqlen_k,
+        num_heads_q,
+        num_heads_kv,
+        headdim,
+        cache_seqlens: torch.Tensor,
+        qkv_dtype=...,
+        headdim_v=None,
+        cu_seqlens_q: torch.Tensor | None = None,
+        cu_seqlens_k_new: torch.Tensor | None = None,
+        cache_leftpad: torch.Tensor | None = None,
+        page_size: int | None = None,
+        max_seqlen_k_new: int = 0,
+        causal: bool = False,
+        window_size=(-1, -1),
+        has_softcap: bool = False,
+        num_splits: int = 0,
+        pack_gqa=None,
+        sm_margin: int = 0,
+    ) -> None: ...

.cuda_typings/vllm/assets/audio.pyi

@@ -0,0 +1,23 @@
+import numpy.typing as npt
+from .base import (
+    VLLM_S3_BUCKET_URL as VLLM_S3_BUCKET_URL,
+    get_vllm_public_assets as get_vllm_public_assets,
+)
+from _typeshed import Incomplete
+from dataclasses import dataclass
+from pathlib import Path
+from vllm.utils.import_utils import PlaceholderModule as PlaceholderModule
+
+ASSET_DIR: str
+AudioAssetName: Incomplete
+
+@dataclass(frozen=True)
+class AudioAsset:
+    name: AudioAssetName
+    @property
+    def filename(self) -> str: ...
+    @property
+    def audio_and_sample_rate(self) -> tuple[npt.NDArray, float]: ...
+    def get_local_path(self) -> Path: ...
+    @property
+    def url(self) -> str: ...

.cuda_typings/vllm/assets/base.pyi

@@ -0,0 +1,9 @@
+from functools import lru_cache
+from pathlib import Path
+from vllm.connections import global_http_connection as global_http_connection
+
+VLLM_S3_BUCKET_URL: str
+
+def get_cache_dir() -> Path: ...
+@lru_cache
+def get_vllm_public_assets(filename: str, s3_prefix: str | None = None) -> Path: ...

.cuda_typings/vllm/assets/image.pyi

@@ -0,0 +1,20 @@
+import torch
+from .base import get_vllm_public_assets as get_vllm_public_assets
+from PIL import Image
+from _typeshed import Incomplete
+from dataclasses import dataclass
+from pathlib import Path
+
+VLM_IMAGES_DIR: str
+ImageAssetName: Incomplete
+
+@dataclass(frozen=True)
+class ImageAsset:
+    name: ImageAssetName
+    def get_path(self, ext: str) -> Path: ...
+    @property
+    def pil_image(self) -> Image.Image: ...
+    def pil_image_ext(self, ext: str) -> Image.Image: ...
+    @property
+    def image_embeds(self) -> torch.Tensor: ...
+    def read_bytes(self, ext: str) -> bytes: ...

.cuda_typings/vllm/assets/video.pyi

@@ -0,0 +1,32 @@
+import numpy.typing as npt
+from .base import get_cache_dir as get_cache_dir
+from PIL import Image
+from _typeshed import Incomplete
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any
+from vllm.utils.import_utils import PlaceholderModule as PlaceholderModule
+
+@lru_cache
+def download_video_asset(filename: str) -> str: ...
+def video_to_ndarrays(path: str, num_frames: int = -1) -> npt.NDArray: ...
+def video_to_pil_images_list(path: str, num_frames: int = -1) -> list[Image.Image]: ...
+def video_get_metadata(path: str, num_frames: int = -1) -> dict[str, Any]: ...
+
+VideoAssetName: Incomplete
+
+@dataclass(frozen=True)
+class VideoAsset:
+    name: VideoAssetName
+    num_frames: int = ...
+    @property
+    def filename(self) -> str: ...
+    @property
+    def video_path(self) -> str: ...
+    @property
+    def pil_images(self) -> list[Image.Image]: ...
+    @property
+    def np_ndarrays(self) -> npt.NDArray: ...
+    @property
+    def metadata(self) -> dict[str, Any]: ...
+    def get_audio(self, sampling_rate: float | None = None) -> npt.NDArray: ...

.cuda_typings/vllm/beam_search.pyi

@@ -0,0 +1,48 @@
+from dataclasses import dataclass
+from vllm.inputs import (
+    EncoderDecoderInputs as EncoderDecoderInputs,
+    TokenInputs as TokenInputs,
+    token_inputs as token_inputs,
+)
+from vllm.inputs.data import DecoderInputs as DecoderInputs
+from vllm.logprobs import Logprob as Logprob
+from vllm.lora.request import LoRARequest as LoRARequest
+from vllm.multimodal.inputs import (
+    MultiModalInputs as MultiModalInputs,
+    mm_inputs as mm_inputs,
+)
+
+@dataclass
+class BeamSearchSequence:
+    orig_prompt: TokenInputs | MultiModalInputs | EncoderDecoderInputs
+    tokens: list[int]
+    logprobs: list[dict[int, Logprob]]
+    lora_request: LoRARequest | None = ...
+    cum_logprob: float = ...
+    text: str | None = ...
+    finish_reason: str | None = ...
+    stop_reason: int | str | None = ...
+    def get_prompt(self): ...
+
+@dataclass
+class BeamSearchOutput:
+    sequences: list[BeamSearchSequence]
+
+class BeamSearchInstance:
+    beams: list[BeamSearchSequence]
+    completed: list[BeamSearchSequence]
+    def __init__(
+        self,
+        prompt: TokenInputs | MultiModalInputs | EncoderDecoderInputs,
+        lora_request: LoRARequest | None = None,
+        logprobs: list[dict[int, Logprob]] | None = None,
+        **kwargs,
+    ) -> None: ...
+
+def get_beam_search_score(
+    tokens: list[int],
+    cumulative_logprob: float,
+    eos_token_id: int,
+    length_penalty: float = 1.0,
+) -> float: ...
+def create_sort_beams_key_function(eos_token_id: int, length_penalty: float): ...

.cuda_typings/vllm/benchmarks/datasets.pyi

@@ -0,0 +1,516 @@
+import abc
+import argparse
+import numpy as np
+from PIL import Image
+from _typeshed import Incomplete
+from abc import ABC, abstractmethod
+from collections.abc import Callable as Callable, Iterator, Mapping
+from dataclasses import dataclass
+from functools import cache
+from typing import Any
+from vllm.lora.request import LoRARequest as LoRARequest
+from vllm.lora.utils import get_adapter_absolute_path as get_adapter_absolute_path
+from vllm.multimodal import MultiModalDataDict as MultiModalDataDict
+from vllm.multimodal.image import convert_image_mode as convert_image_mode
+from vllm.tokenizers import TokenizerLike as TokenizerLike
+from vllm.utils.argparse_utils import FlexibleArgumentParser as FlexibleArgumentParser
+from vllm.utils.import_utils import PlaceholderModule as PlaceholderModule
+
+datasets: Incomplete
+logger: Incomplete
+DEFAULT_NUM_PROMPTS: int
+
+@dataclass
+class SampleRequest:
+    prompt: str | list[str]
+    prompt_len: int
+    expected_output_len: int
+    multi_modal_data: MultiModalDataDict | dict | list[dict] | None = ...
+    lora_request: LoRARequest | None = ...
+    request_id: str | None = ...
+
+class BenchmarkDataset(ABC, metaclass=abc.ABCMeta):
+    DEFAULT_SEED: int
+    IS_MULTIMODAL: bool
+    dataset_path: Incomplete
+    random_seed: Incomplete
+    disable_shuffle: Incomplete
+    data: Incomplete
+    def __init__(
+        self,
+        dataset_path: str | None = None,
+        random_seed: int = ...,
+        disable_shuffle: bool = False,
+        **kwargs,
+    ) -> None: ...
+    def apply_multimodal_chat_transformation(
+        self,
+        prompt: str,
+        mm_content: MultiModalDataDict | dict | list[dict] | None = None,
+    ) -> list[dict]: ...
+    def load_data(self) -> None: ...
+    def get_random_lora_request(
+        self, max_loras: int | None = None, lora_path: str | None = None
+    ) -> LoRARequest | None: ...
+    @abstractmethod
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+    ) -> list[SampleRequest]: ...
+    def maybe_oversample_requests(
+        self,
+        requests: list[SampleRequest],
+        num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+    ) -> None: ...
+
+def is_valid_sequence(
+    prompt_len: int,
+    output_len: int,
+    min_len: int = 4,
+    max_prompt_len: int = 1024,
+    max_total_len: int = 2048,
+    skip_min_output_len_check: bool = False,
+) -> bool: ...
+@cache
+def lora_path_on_disk(lora_path: str) -> str: ...
+
+lora_tokenizer_cache: dict[int, TokenizerLike]
+
+def process_image(image: Any) -> Mapping[str, Any]: ...
+def process_video(video: Any) -> Mapping[str, Any]: ...
+def gen_prompt_decode_to_target_len(
+    tokenizer: TokenizerLike,
+    token_sequence: list[int],
+    target_token_len: int,
+    max_retry: int = 10,
+    add_special_tokens: bool = False,
+    rng: np.random.Generator | None = None,
+) -> tuple[str, list[int], int]: ...
+
+class RandomDataset(BenchmarkDataset):
+    DEFAULT_PREFIX_LEN: int
+    DEFAULT_RANGE_RATIO: float
+    DEFAULT_INPUT_LEN: int
+    DEFAULT_OUTPUT_LEN: int
+    def __init__(self, **kwargs) -> None: ...
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        prefix_len: int = ...,
+        range_ratio: float = ...,
+        input_len: int = ...,
+        output_len: int = ...,
+        batchsize: int = 1,
+        **kwargs,
+    ) -> list[SampleRequest]: ...
+    def get_prefix(
+        self, tokenizer: TokenizerLike, allowed_tokens: np.ndarray, prefix_len: int
+    ) -> list[int]: ...
+    def get_sampling_params(
+        self,
+        num_requests: int,
+        range_ratio: float,
+        input_len: int,
+        output_len: int,
+        tokenizer: TokenizerLike,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ...
+    def generate_token_sequence(
+        self,
+        *,
+        tokenizer: TokenizerLike,
+        prefix_token_ids: list[int],
+        prefix_len: int,
+        vocab_size: int,
+        input_len: int,
+        offset: int,
+        index: int,
+        allowed_tokens: np.ndarray,
+    ) -> tuple[str, int, int]: ...
+
+class RandomDatasetForReranking(RandomDataset):
+    def __init__(self, **kwargs) -> None: ...
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        request_id_prefix: str = "",
+        range_ratio: float = ...,
+        input_len: int = ...,
+        batchsize: int = 1,
+        is_reranker: bool = True,
+        **kwargs,
+    ) -> list[SampleRequest]: ...
+
+class RandomMultiModalDataset(RandomDataset):
+    IS_MULTIMODAL: bool
+    DEFAULT_LIMIT_MM_PER_PROMPT: Incomplete
+    DEFAULT_BASE_ITEMS_PER_REQUEST: int
+    DEFAULT_NUM_MM_ITEMS_RANGE_RATIO: float
+    DEFAULT_MM_ITEM_BUCKET_CONFIG: Incomplete
+    DEFAULT_ENABLE_MULTIMODAL_CHAT: bool
+    def __init__(self, **kwargs) -> None: ...
+    def generate_synthetic_image(self, width: int, height: int) -> Image.Image: ...
+    def generate_synthetic_video(
+        self, width: int, height: int, num_frames: int
+    ) -> dict: ...
+    def map_config_to_modality(self, config: tuple[int, int, int]) -> str: ...
+    def normalize_bucket_config(
+        self, bucket_config: dict[tuple[int, int, int], float]
+    ) -> dict[tuple[int, int, int], float]: ...
+    def generate_mm_item(
+        self, mm_item_config: tuple[int, int, int]
+    ) -> Mapping[str, Any]: ...
+    def get_mm_item_sampling_params(
+        self,
+        base_items_per_request: int,
+        num_mm_items_range_ratio: float,
+        limit_mm_per_prompt: dict[str, int],
+        bucket_config: dict[tuple[int, int, int], float],
+    ) -> tuple[int, int, dict[str, int], dict[tuple[int, int, int], float]]: ...
+    def get_mm_item_iterator(
+        self,
+        min_num_mm_items: int,
+        max_num_mm_items: int,
+        bucket_config: dict[tuple[int, int, int], float],
+        limit_mm_per_prompt: dict[str, int],
+    ) -> Iterator[tuple[int, int, int]]: ...
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        prefix_len: int = ...,
+        range_ratio: float = ...,
+        input_len: int = ...,
+        output_len: int = ...,
+        limit_mm_per_prompt: dict[str, int] = ...,
+        base_items_per_request: int = ...,
+        num_mm_items_range_ratio: float = ...,
+        bucket_config: dict[tuple[int, int, int], float] = ...,
+        enable_multimodal_chat: bool = ...,
+        **kwargs,
+    ) -> list[SampleRequest]: ...
+
+class ShareGPTDataset(BenchmarkDataset):
+    def __init__(self, **kwargs) -> None: ...
+    data: Incomplete
+    def load_data(self) -> None: ...
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        lora_path: str | None = None,
+        max_loras: int | None = None,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list: ...
+
+class _ValidateDatasetArgs(argparse.Action):
+    def __call__(self, parser, namespace, values, option_string=None) -> None: ...
+
+def add_dataset_parser(parser: FlexibleArgumentParser): ...
+def add_random_dataset_base_args(
+    parser_or_group: FlexibleArgumentParser | argparse._ArgumentGroup,
+) -> None: ...
+def add_random_multimodal_dataset_args(
+    parser_or_group: FlexibleArgumentParser | argparse._ArgumentGroup,
+) -> None: ...
+def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: ...
+
+class CustomDataset(BenchmarkDataset):
+    def __init__(self, **kwargs) -> None: ...
+    data: Incomplete
+    def load_data(self) -> None: ...
+    num_available_samples: Incomplete
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        lora_path: str | None = None,
+        max_loras: int | None = None,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
+        skip_chat_template: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list: ...
+
+class CustomMMDataset(CustomDataset):
+    IS_MULTIMODAL: bool
+    num_available_samples: Incomplete
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list: ...
+
+class SpecBench(CustomDataset):
+    category: Incomplete
+    def __init__(self, **kwargs) -> None: ...
+    data: Incomplete
+    def load_data(self) -> None: ...
+    def sample(self, **kwargs) -> list: ...
+
+class SonnetDataset(BenchmarkDataset):
+    DEFAULT_PREFIX_LEN: int
+    DEFAULT_INPUT_LEN: int
+    DEFAULT_OUTPUT_LEN: int
+    def __init__(self, **kwargs) -> None: ...
+    data: Incomplete
+    def load_data(self) -> None: ...
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        prefix_len: int = ...,
+        input_len: int = ...,
+        output_len: int = ...,
+        return_prompt_formatted: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list: ...
+
+class BurstGPTDataset(BenchmarkDataset):
+    def __init__(self, **kwargs) -> None: ...
+    data: Incomplete
+    def load_data(self) -> None: ...
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        max_loras: int | None = None,
+        lora_path: str | None = None,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list[SampleRequest]: ...
+
+class HuggingFaceDataset(BenchmarkDataset, metaclass=abc.ABCMeta):
+    SUPPORTED_DATASET_PATHS: set[str] | dict[str, Callable]
+    dataset_split: Incomplete
+    dataset_subset: Incomplete
+    load_stream: Incomplete
+    hf_name: Incomplete
+    trust_remote_code: Incomplete
+    def __init__(
+        self,
+        dataset_path: str,
+        dataset_split: str,
+        no_stream: bool = False,
+        dataset_subset: str | None = None,
+        hf_name: str | None = None,
+        trust_remote_code: bool = False,
+        **kwargs,
+    ) -> None: ...
+    data: Incomplete
+    def load_data(self) -> None: ...
+
+class ConversationDataset(HuggingFaceDataset):
+    SUPPORTED_DATASET_PATHS: Incomplete
+    IS_MULTIMODAL: bool
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list: ...
+
+class MultiModalConversationDataset(HuggingFaceDataset):
+    SUPPORTED_DATASET_PATHS: Incomplete
+    IS_MULTIMODAL: bool
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list: ...
+
+class VisionArenaDataset(HuggingFaceDataset):
+    DEFAULT_OUTPUT_LEN: int
+    SUPPORTED_DATASET_PATHS: Incomplete
+    IS_MULTIMODAL: bool
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list: ...
+
+class MMVUDataset(HuggingFaceDataset):
+    DEFAULT_OUTPUT_LEN: int
+    SUPPORTED_DATASET_PATHS: Incomplete
+    def __init__(self, **kwargs) -> None: ...
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list: ...
+
+class InstructCoderDataset(HuggingFaceDataset):
+    DEFAULT_OUTPUT_LEN: int
+    SUPPORTED_DATASET_PATHS: Incomplete
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
+        skip_chat_template: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list[SampleRequest]: ...
+    def sample_prompts(self, n: int) -> Iterator[str]: ...
+
+class MTBenchDataset(HuggingFaceDataset):
+    DEFAULT_OUTPUT_LEN: int
+    SUPPORTED_DATASET_PATHS: Incomplete
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
+        skip_chat_template: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list: ...
+
+class BlazeditDataset(HuggingFaceDataset):
+    DEFAULT_OUTPUT_LEN: int
+    SUPPORTED_DATASET_PATHS: Incomplete
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        output_len: int | None = None,
+        skip_chat_template: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        min_distance: float = 0.0,
+        max_distance: float = 1.0,
+        **kwargs,
+    ) -> list: ...
+
+class AIMODataset(HuggingFaceDataset):
+    SUPPORTED_DATASET_PATHS: Incomplete
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        output_len: int | None = None,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list: ...
+
+zeta_prompt: str
+
+class NextEditPredictionDataset(HuggingFaceDataset):
+    SUPPORTED_DATASET_PATHS: Incomplete
+    MAPPING_PROMPT_FUNCS: Incomplete
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ): ...
+
+class ASRDataset(HuggingFaceDataset):
+    SUPPORTED_DATASET_PATHS: Incomplete
+    DEFAULT_OUTPUT_LEN: int
+    IS_MULTIMODAL: bool
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        output_len: int | None = None,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list: ...
+
+class MLPerfDataset(HuggingFaceDataset):
+    SUPPORTED_DATASET_PATHS: Incomplete
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        output_len: int | None = None,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list[SampleRequest]: ...
+
+class PrefixRepetitionRandomDataset(BenchmarkDataset):
+    DEFAULT_PREFIX_LEN: int
+    DEFAULT_SUFFIX_LEN: int
+    DEFAULT_NUM_PREFIXES: int
+    DEFAULT_OUTPUT_LEN: int
+    def __init__(self, **kwargs) -> None: ...
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        prefix_len: int = ...,
+        suffix_len: int = ...,
+        num_prefixes: int = ...,
+        output_len: int = ...,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list[SampleRequest]: ...
+
+class MMStarDataset(HuggingFaceDataset):
+    DEFAULT_OUTPUT_LEN: int
+    SUPPORTED_DATASET_PATHS: Incomplete
+    IS_MULTIMODAL: bool
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        **kwargs,
+    ) -> list[SampleRequest]: ...

.cuda_typings/vllm/benchmarks/latency.pyi

@@ -0,0 +1,15 @@
+import argparse
+from typing import Any
+from vllm.benchmarks.lib.utils import (
+    convert_to_pytorch_benchmark_format as convert_to_pytorch_benchmark_format,
+    write_to_json as write_to_json,
+)
+from vllm.engine.arg_utils import EngineArgs as EngineArgs
+from vllm.inputs import PromptType as PromptType
+from vllm.sampling_params import BeamSearchParams as BeamSearchParams
+
+def save_to_pytorch_benchmark_format(
+    args: argparse.Namespace, results: dict[str, Any]
+) -> None: ...
+def add_cli_args(parser: argparse.ArgumentParser): ...
+def main(args: argparse.Namespace): ...

.cuda_typings/vllm/benchmarks/lib/endpoint_request_func.pyi

@@ -0,0 +1,113 @@
+import aiohttp
+from _typeshed import Incomplete
+from collections.abc import Awaitable
+from dataclasses import dataclass, field
+from tqdm.asyncio import tqdm as tqdm
+from typing import Literal, Protocol
+
+AIOHTTP_TIMEOUT: Incomplete
+
+class StreamedResponseHandler:
+    buffer: str
+    def __init__(self) -> None: ...
+    def add_chunk(self, chunk_bytes: bytes) -> list[str]: ...
+
+@dataclass
+class RequestFuncInput:
+    prompt: str | list[str]
+    api_url: str
+    prompt_len: int
+    output_len: int
+    model: str
+    model_name: str | None = ...
+    logprobs: int | None = ...
+    extra_headers: dict | None = ...
+    extra_body: dict | None = ...
+    multi_modal_content: dict | list[dict] | None = ...
+    ignore_eos: bool = ...
+    language: str | None = ...
+    request_id: str | None = ...
+
+@dataclass
+class RequestFuncOutput:
+    generated_text: str = ...
+    success: bool = ...
+    latency: float = ...
+    output_tokens: int = ...
+    ttft: float = ...
+    itl: list[float] = field(default_factory=list)
+    tpot: float = ...
+    prompt_len: int = ...
+    error: str = ...
+    start_time: float = ...
+    input_audio_duration: float = ...
+
+class RequestFunc(Protocol):
+    def __call__(
+        self,
+        request_func_input: RequestFuncInput,
+        session: aiohttp.ClientSession,
+        pbar: tqdm | None = None,
+    ) -> Awaitable[RequestFuncOutput]: ...
+
+async def async_request_openai_completions(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput: ...
+async def async_request_openai_chat_completions(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+    mm_position: Literal["first", "last"] = "last",
+) -> RequestFuncOutput: ...
+async def async_request_openai_audio(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput: ...
+async def async_request_openai_embeddings(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput: ...
+async def async_request_vllm_rerank(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput: ...
+async def async_request_openai_embeddings_chat(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+    mm_position: Literal["first", "last"] = "last",
+) -> RequestFuncOutput: ...
+async def async_request_openai_embeddings_clip(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput: ...
+async def async_request_openai_embeddings_vlm2vec(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput: ...
+async def async_request_infinity_embeddings(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput: ...
+async def async_request_infinity_embeddings_clip(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput: ...
+async def async_request_vllm_pooling(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput: ...
+
+ASYNC_REQUEST_FUNCS: dict[str, RequestFunc]
+POOLING_BACKENDS: Incomplete
+OPENAI_COMPATIBLE_BACKENDS: Incomplete

.cuda_typings/vllm/benchmarks/lib/ready_checker.pyi

@@ -0,0 +1,18 @@
+import aiohttp
+from .endpoint_request_func import (
+    RequestFunc as RequestFunc,
+    RequestFuncInput as RequestFuncInput,
+    RequestFuncOutput as RequestFuncOutput,
+)
+from _typeshed import Incomplete
+from vllm.logger import init_logger as init_logger
+
+logger: Incomplete
+
+async def wait_for_endpoint(
+    request_func: RequestFunc,
+    test_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    timeout_seconds: int = 600,
+    retry_interval: int = 5,
+) -> RequestFuncOutput: ...

.cuda_typings/vllm/benchmarks/lib/utils.pyi

@@ -0,0 +1,21 @@
+import argparse
+import json
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+def extract_field(
+    args: argparse.Namespace, extra_info: dict[str, Any], field_name: str
+) -> str: ...
+def use_compile(args: argparse.Namespace, extra_info: dict[str, Any]) -> bool: ...
+def convert_to_pytorch_benchmark_format(
+    args: argparse.Namespace, metrics: dict[str, list], extra_info: dict[str, Any]
+) -> list: ...
+
+class InfEncoder(json.JSONEncoder):
+    def clear_inf(self, o: Any): ...
+    def iterencode(self, o: Any, *args, **kwargs) -> Any: ...
+
+def write_to_json(filename: str, records: list) -> None: ...
+@contextmanager
+def default_vllm_config() -> Generator[None]: ...

.cuda_typings/vllm/benchmarks/mm_processor.pyi

@@ -0,0 +1,26 @@
+import argparse
+from typing import Any, Literal
+from vllm.benchmarks.datasets import (
+    MultiModalConversationDataset as MultiModalConversationDataset,
+    VisionArenaDataset as VisionArenaDataset,
+)
+from vllm.benchmarks.throughput import get_requests as get_requests
+from vllm.engine.arg_utils import EngineArgs as EngineArgs
+from vllm.utils.gc_utils import freeze_gc_heap as freeze_gc_heap
+from vllm.utils.import_utils import PlaceholderModule as PlaceholderModule
+from vllm.v1.engine.llm_engine import LLMEngine as LLMEngine
+
+def get_timing_stats_from_engine(
+    llm_engine: LLMEngine,
+) -> dict[str, dict[str, float]]: ...
+def collect_mm_processor_stats(llm_engine: LLMEngine) -> dict[str, list[float]]: ...
+def calculate_mm_processor_metrics(
+    stats_by_stage: dict[str, list[float]],
+    selected_percentiles: list[float],
+    *,
+    unit: Literal["us", "ms", "s"] = "ms",
+) -> dict[str, dict[str, float]]: ...
+def validate_args(args) -> None: ...
+def benchmark_multimodal_processor(args: argparse.Namespace) -> dict[str, Any]: ...
+def add_cli_args(parser: argparse.ArgumentParser) -> None: ...
+def main(args: argparse.Namespace) -> None: ...

.cuda_typings/vllm/benchmarks/plot.pyi

@@ -0,0 +1,17 @@
+from pathlib import Path
+from typing import Any
+from vllm.utils.import_utils import PlaceholderModule as PlaceholderModule
+
+def generate_timeline_plot(
+    results: list[dict[str, Any]],
+    output_path: Path,
+    colors: list[str] | None = None,
+    itl_thresholds: list[float] | None = None,
+    labels: list[str] | None = None,
+) -> None: ...
+def construct_timeline_data(
+    requests_data: list[dict[str, Any]], itl_thresholds: list[float], labels: list[str]
+) -> list[dict[str, Any]]: ...
+def generate_dataset_stats_plot(
+    results: list[dict[str, Any]], output_path: Path
+) -> None: ...

.cuda_typings/vllm/benchmarks/serve.pyi

@@ -0,0 +1,156 @@
+import aiohttp
+import argparse
+import ssl
+from _typeshed import Incomplete
+from collections.abc import AsyncGenerator, Iterable
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Literal
+from vllm.benchmarks.datasets import (
+    SampleRequest as SampleRequest,
+    add_dataset_parser as add_dataset_parser,
+    get_samples as get_samples,
+)
+from vllm.benchmarks.lib.endpoint_request_func import (
+    ASYNC_REQUEST_FUNCS as ASYNC_REQUEST_FUNCS,
+    OPENAI_COMPATIBLE_BACKENDS as OPENAI_COMPATIBLE_BACKENDS,
+    POOLING_BACKENDS as POOLING_BACKENDS,
+    RequestFuncInput as RequestFuncInput,
+    RequestFuncOutput as RequestFuncOutput,
+)
+from vllm.benchmarks.lib.ready_checker import wait_for_endpoint as wait_for_endpoint
+from vllm.benchmarks.lib.utils import (
+    convert_to_pytorch_benchmark_format as convert_to_pytorch_benchmark_format,
+    write_to_json as write_to_json,
+)
+from vllm.tokenizers import (
+    TokenizerLike as TokenizerLike,
+    get_tokenizer as get_tokenizer,
+)
+from vllm.utils.gc_utils import freeze_gc_heap as freeze_gc_heap
+from vllm.utils.network_utils import join_host_port as join_host_port
+
+MILLISECONDS_TO_SECONDS_CONVERSION: int
+TERM_PLOTLIB_AVAILABLE: Incomplete
+
+async def get_first_model_from_server(
+    base_url: str,
+    headers: dict | None = None,
+    ssl_context: ssl.SSLContext | bool | None = None,
+) -> tuple[str, str]: ...
+@dataclass
+class SpecDecodeMetrics:
+    num_drafts: int
+    num_draft_tokens: int
+    num_accepted_tokens: int
+    accepted_per_pos: dict[int, int]
+
+async def fetch_spec_decode_metrics(
+    base_url: str, session: aiohttp.ClientSession
+) -> SpecDecodeMetrics | None: ...
+
+class TaskType(Enum):
+    GENERATION = "generation"
+    POOLING = "pooling"
+
+@dataclass
+class BenchmarkMetrics:
+    completed: int
+    failed: int
+    total_input: int
+    total_output: int
+    request_throughput: float
+    request_goodput: float
+    output_throughput: float
+    total_token_throughput: float
+    mean_ttft_ms: float
+    median_ttft_ms: float
+    std_ttft_ms: float
+    percentiles_ttft_ms: list[tuple[float, float]]
+    mean_tpot_ms: float
+    median_tpot_ms: float
+    std_tpot_ms: float
+    percentiles_tpot_ms: list[tuple[float, float]]
+    mean_itl_ms: float
+    median_itl_ms: float
+    std_itl_ms: float
+    percentiles_itl_ms: list[tuple[float, float]]
+    mean_e2el_ms: float
+    median_e2el_ms: float
+    std_e2el_ms: float
+    percentiles_e2el_ms: list[tuple[float, float]]
+    max_output_tokens_per_s: float
+    max_concurrent_requests: int
+    rtfx: float = ...
+
+@dataclass
+class EmbedBenchmarkMetrics:
+    completed: int
+    failed: int
+    total_input: int
+    request_throughput: float
+    total_token_throughput: float
+    mean_e2el_ms: float
+    std_e2el_ms: float
+    median_e2el_ms: float
+    percentiles_e2el_ms: float
+
+async def get_request(
+    input_requests: list[SampleRequest],
+    request_rate: float,
+    burstiness: float = 1.0,
+    ramp_up_strategy: Literal["linear", "exponential"] | None = None,
+    ramp_up_start_rps: int | None = None,
+    ramp_up_end_rps: int | None = None,
+) -> AsyncGenerator[tuple[SampleRequest, float], None]: ...
+def calculate_metrics_for_embeddings(
+    outputs: list[RequestFuncOutput], dur_s: float, selected_percentiles: list[float]
+) -> EmbedBenchmarkMetrics: ...
+def calculate_metrics(
+    input_requests: list[SampleRequest],
+    outputs: list[RequestFuncOutput],
+    dur_s: float,
+    tokenizer: TokenizerLike,
+    selected_percentiles: list[float],
+    goodput_config_dict: dict[str, float],
+) -> tuple[BenchmarkMetrics, list[int]]: ...
+async def benchmark(
+    task_type: TaskType,
+    endpoint_type: str,
+    api_url: str,
+    base_url: str,
+    model_id: str,
+    model_name: str,
+    tokenizer: TokenizerLike,
+    input_requests: list[SampleRequest],
+    logprobs: int | None,
+    request_rate: float,
+    burstiness: float,
+    disable_tqdm: bool,
+    num_warmups: int,
+    profile: bool,
+    selected_percentile_metrics: list[str],
+    selected_percentiles: list[float],
+    ignore_eos: bool,
+    goodput_config_dict: dict[str, float],
+    max_concurrency: int | None,
+    lora_modules: Iterable[str] | None,
+    extra_headers: dict | None,
+    extra_body: dict | None,
+    ramp_up_strategy: Literal["linear", "exponential"] | None = None,
+    ramp_up_start_rps: int | None = None,
+    ramp_up_end_rps: int | None = None,
+    ready_check_timeout_sec: int = 600,
+    ssl_context: ssl.SSLContext | bool | None = None,
+): ...
+def check_goodput_args(args): ...
+def parse_goodput(slo_pairs): ...
+def save_to_pytorch_benchmark_format(
+    args: argparse.Namespace, results: dict[str, Any], file_name: str
+) -> None: ...
+def compute_result_filename(
+    args: argparse.Namespace, model_id: str, label: str, current_dt: str
+) -> str | None: ...
+def add_cli_args(parser: argparse.ArgumentParser): ...
+def main(args: argparse.Namespace) -> dict[str, Any]: ...
+async def main_async(args: argparse.Namespace) -> dict[str, Any]: ...

.cuda_typings/vllm/benchmarks/startup.pyi

@@ -0,0 +1,18 @@
+import argparse
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+from vllm.benchmarks.lib.utils import (
+    convert_to_pytorch_benchmark_format as convert_to_pytorch_benchmark_format,
+    write_to_json as write_to_json,
+)
+from vllm.engine.arg_utils import EngineArgs as EngineArgs
+
+@contextmanager
+def cold_startup() -> Generator[None]: ...
+def run_startup_in_subprocess(engine_args, result_queue) -> None: ...
+def save_to_pytorch_benchmark_format(
+    args: argparse.Namespace, results: dict[str, Any]
+) -> None: ...
+def add_cli_args(parser: argparse.ArgumentParser): ...
+def main(args: argparse.Namespace): ...

.cuda_typings/vllm/benchmarks/sweep/cli.pyi

@@ -0,0 +1,15 @@
+import argparse
+from .plot import SweepPlotArgs as SweepPlotArgs
+from .plot_pareto import SweepPlotParetoArgs as SweepPlotParetoArgs
+from .serve import SweepServeArgs as SweepServeArgs
+from .serve_workload import SweepServeWorkloadArgs as SweepServeWorkloadArgs
+from .startup import SweepStartupArgs as SweepStartupArgs
+from _typeshed import Incomplete
+from vllm.entrypoints.utils import (
+    VLLM_SUBCMD_PARSER_EPILOG as VLLM_SUBCMD_PARSER_EPILOG,
+)
+
+SUBCOMMANDS: Incomplete
+
+def add_cli_args(parser: argparse.ArgumentParser): ...
+def main(args: argparse.Namespace): ...

.cuda_typings/vllm/benchmarks/sweep/param_sweep.pyi

@@ -0,0 +1,20 @@
+import os
+from typing import Any
+
+class ParameterSweep(list["ParameterSweepItem"]):
+    @classmethod
+    def read_json(cls, filepath: os.PathLike): ...
+    @classmethod
+    def read_from_dict(cls, data: dict[str, dict[str, object]]): ...
+    @classmethod
+    def from_records(cls, records: list[dict[str, object]]): ...
+
+class ParameterSweepItem(dict[str, object]):
+    @classmethod
+    def from_record(cls, record: dict[str, object]): ...
+    def __or__(self, other: dict[str, Any]): ...
+    @property
+    def name(self) -> str: ...
+    def has_param(self, param_key: str) -> bool: ...
+    def apply_to_cmd(self, cmd: list[str]) -> list[str]: ...
+    def as_text(self, sep: str = ", ") -> str: ...

.cuda_typings/vllm/benchmarks/sweep/plot.pyi

@@ -0,0 +1,137 @@
+import abc
+import argparse
+import pandas as pd
+from .utils import sanitize_filename as sanitize_filename
+from _typeshed import Incomplete
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from types import TracebackType
+from typing import ClassVar
+from typing_extensions import Self, override
+from vllm.utils.collection_utils import full_groupby as full_groupby
+from vllm.utils.import_utils import PlaceholderModule as PlaceholderModule
+
+seaborn: Incomplete
+
+@dataclass
+class PlotFilterBase(ABC, metaclass=abc.ABCMeta):
+    var: str
+    target: str
+    @classmethod
+    def parse_str(cls, s: str): ...
+    @abstractmethod
+    def apply(self, df: pd.DataFrame) -> pd.DataFrame: ...
+
+@dataclass
+class PlotEqualTo(PlotFilterBase):
+    @override
+    def apply(self, df: pd.DataFrame) -> pd.DataFrame: ...
+
+@dataclass
+class PlotNotEqualTo(PlotFilterBase):
+    @override
+    def apply(self, df: pd.DataFrame) -> pd.DataFrame: ...
+
+@dataclass
+class PlotLessThan(PlotFilterBase):
+    @override
+    def apply(self, df: pd.DataFrame) -> pd.DataFrame: ...
+
+@dataclass
+class PlotLessThanOrEqualTo(PlotFilterBase):
+    @override
+    def apply(self, df: pd.DataFrame) -> pd.DataFrame: ...
+
+@dataclass
+class PlotGreaterThan(PlotFilterBase):
+    @override
+    def apply(self, df: pd.DataFrame) -> pd.DataFrame: ...
+
+@dataclass
+class PlotGreaterThanOrEqualTo(PlotFilterBase):
+    @override
+    def apply(self, df: pd.DataFrame) -> pd.DataFrame: ...
+
+PLOT_FILTERS: dict[str, type[PlotFilterBase]]
+
+class PlotFilters(list[PlotFilterBase]):
+    @classmethod
+    def parse_str(cls, s: str): ...
+    def apply(self, df: pd.DataFrame) -> pd.DataFrame: ...
+
+@dataclass
+class PlotBinner:
+    var: str
+    bin_size: float
+    @classmethod
+    def parse_str(cls, s: str): ...
+    def apply(self, df: pd.DataFrame) -> pd.DataFrame: ...
+
+PLOT_BINNERS: dict[str, type[PlotBinner]]
+
+class PlotBinners(list[PlotBinner]):
+    @classmethod
+    def parse_str(cls, s: str): ...
+    def apply(self, df: pd.DataFrame) -> pd.DataFrame: ...
+
+class DummyExecutor:
+    map = map
+    def __enter__(self) -> Self: ...
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        exc_traceback: TracebackType | None,
+    ) -> None: ...
+
+def plot(
+    output_dir: Path,
+    fig_dir: Path,
+    fig_by: list[str],
+    row_by: list[str],
+    col_by: list[str],
+    curve_by: list[str],
+    *,
+    var_x: str,
+    var_y: str,
+    filter_by: PlotFilters,
+    bin_by: PlotBinners,
+    scale_x: str | None,
+    scale_y: str | None,
+    dry_run: bool,
+    fig_name: str = "FIGURE",
+    error_bars: bool = True,
+    fig_height: float = 6.4,
+    fig_dpi: int = 300,
+): ...
+@dataclass
+class SweepPlotArgs:
+    output_dir: Path
+    fig_dir: Path
+    fig_by: list[str]
+    row_by: list[str]
+    col_by: list[str]
+    curve_by: list[str]
+    var_x: str
+    var_y: str
+    filter_by: PlotFilters
+    bin_by: PlotBinners
+    scale_x: str | None
+    scale_y: str | None
+    dry_run: bool
+    fig_name: str = ...
+    error_bars: bool = ...
+    fig_height: float = ...
+    fig_dpi: int = ...
+    parser_name: ClassVar[str] = ...
+    parser_help: ClassVar[str] = ...
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace): ...
+    @classmethod
+    def add_cli_args(
+        cls, parser: argparse.ArgumentParser
+    ) -> argparse.ArgumentParser: ...
+
+def run_main(args: SweepPlotArgs): ...
+def main(args: argparse.Namespace): ...

.cuda_typings/vllm/benchmarks/sweep/plot_pareto.pyi

@@ -0,0 +1,36 @@
+import argparse
+from .plot import DummyExecutor as DummyExecutor
+from .utils import sanitize_filename as sanitize_filename
+from _typeshed import Incomplete
+from dataclasses import dataclass
+from pathlib import Path
+from typing import ClassVar
+from vllm.utils.collection_utils import full_groupby as full_groupby
+from vllm.utils.import_utils import PlaceholderModule as PlaceholderModule
+
+seaborn: Incomplete
+
+def plot_pareto(
+    output_dir: Path,
+    user_count_var: str | None,
+    gpu_count_var: str | None,
+    label_by: list[str],
+    *,
+    dry_run: bool,
+): ...
+@dataclass
+class SweepPlotParetoArgs:
+    output_dir: Path
+    user_count_var: str | None
+    gpu_count_var: str | None
+    label_by: list[str]
+    dry_run: bool
+    parser_name: ClassVar[str] = ...
+    parser_help: ClassVar[str] = ...
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace): ...
+    @classmethod
+    def add_cli_args(cls, parser: argparse.ArgumentParser): ...
+
+def run_main(args: SweepPlotParetoArgs): ...
+def main(args: argparse.Namespace): ...

.cuda_typings/vllm/benchmarks/sweep/serve.pyi

@@ -0,0 +1,101 @@
+import argparse
+import contextlib
+from .param_sweep import (
+    ParameterSweep as ParameterSweep,
+    ParameterSweepItem as ParameterSweepItem,
+)
+from .server import ServerProcess as ServerProcess
+from .utils import sanitize_filename as sanitize_filename
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import ClassVar
+from vllm.utils.import_utils import PlaceholderModule as PlaceholderModule
+
+@contextlib.contextmanager
+def run_server(
+    serve_cmd: list[str],
+    after_bench_cmd: list[str],
+    *,
+    show_stdout: bool,
+    serve_overrides: ParameterSweepItem,
+    dry_run: bool,
+    server_ready_timeout: int = 300,
+): ...
+def run_benchmark(
+    server: ServerProcess | None,
+    bench_cmd: list[str],
+    *,
+    serve_overrides: ParameterSweepItem,
+    bench_overrides: ParameterSweepItem,
+    run_number: int,
+    output_path: Path,
+    dry_run: bool,
+): ...
+def server_ctx(
+    serve_cmd: list[str],
+    after_bench_cmd: list[str],
+    *,
+    show_stdout: bool,
+    serve_comb: ParameterSweepItem,
+    bench_params: ParameterSweep,
+    experiment_dir: Path,
+    dry_run: bool,
+    server_ready_timeout: int = 300,
+): ...
+def run_comb(
+    server: ServerProcess | None,
+    bench_cmd: list[str],
+    *,
+    serve_comb: ParameterSweepItem,
+    bench_comb: ParameterSweepItem,
+    link_vars: list[tuple[str, str]],
+    base_path: Path,
+    num_runs: int,
+    dry_run: bool,
+): ...
+def run_combs(
+    serve_cmd: list[str],
+    bench_cmd: list[str],
+    after_bench_cmd: list[str],
+    *,
+    show_stdout: bool,
+    server_ready_timeout: int,
+    serve_params: ParameterSweep,
+    bench_params: ParameterSweep,
+    link_vars: list[tuple[str, str]],
+    experiment_dir: Path,
+    num_runs: int,
+    dry_run: bool,
+): ...
+@dataclass
+class SweepServeArgs:
+    serve_cmd: list[str]
+    bench_cmd: list[str]
+    after_bench_cmd: list[str]
+    show_stdout: bool
+    server_ready_timeout: int
+    serve_params: ParameterSweep
+    bench_params: ParameterSweep
+    link_vars: list[tuple[str, str]]
+    output_dir: Path
+    experiment_name: str
+    num_runs: int
+    dry_run: bool
+    resume: bool
+    parser_name: ClassVar[str] = ...
+    parser_help: ClassVar[str] = ...
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace): ...
+    @classmethod
+    def add_cli_args(
+        cls, parser: argparse.ArgumentParser
+    ) -> argparse.ArgumentParser: ...
+    @staticmethod
+    def parse_link_vars(s: str) -> list[tuple[str, str]]: ...
+    def resolve_experiment_dir(self) -> Path: ...
+    @contextmanager
+    def run_ctx(self, experiment_dir: Path): ...
+
+def run_main(args: SweepServeArgs): ...
+def main(args: argparse.Namespace): ...

.cuda_typings/vllm/benchmarks/sweep/serve_workload.pyi

@@ -0,0 +1,77 @@
+import argparse
+from .param_sweep import (
+    ParameterSweep as ParameterSweep,
+    ParameterSweepItem as ParameterSweepItem,
+)
+from .serve import (
+    SweepServeArgs as SweepServeArgs,
+    run_comb as run_comb,
+    server_ctx as server_ctx,
+)
+from .server import ServerProcess as ServerProcess
+from _typeshed import Incomplete
+from dataclasses import dataclass
+from pathlib import Path
+from typing import ClassVar
+from vllm.benchmarks.datasets import DEFAULT_NUM_PROMPTS as DEFAULT_NUM_PROMPTS
+from vllm.utils.import_utils import PlaceholderModule as PlaceholderModule
+
+WorkloadVariable: Incomplete
+
+def run_comb_workload(
+    server: ServerProcess | None,
+    bench_cmd: list[str],
+    *,
+    serve_comb: ParameterSweepItem,
+    bench_comb: ParameterSweepItem,
+    link_vars: list[tuple[str, str]],
+    experiment_dir: Path,
+    num_runs: int,
+    dry_run: bool,
+    workload_var: WorkloadVariable,
+    workload_value: int,
+) -> list[dict[str, object]] | None: ...
+def explore_comb_workloads(
+    server: ServerProcess | None,
+    bench_cmd: list[str],
+    *,
+    serve_comb: ParameterSweepItem,
+    bench_comb: ParameterSweepItem,
+    link_vars: list[tuple[str, str]],
+    workload_var: WorkloadVariable,
+    workload_iters: int,
+    experiment_dir: Path,
+    num_runs: int,
+    dry_run: bool,
+): ...
+def explore_combs_workloads(
+    serve_cmd: list[str],
+    bench_cmd: list[str],
+    after_bench_cmd: list[str],
+    *,
+    show_stdout: bool,
+    server_ready_timeout: int,
+    serve_params: ParameterSweep,
+    bench_params: ParameterSweep,
+    link_vars: list[tuple[str, str]],
+    workload_var: WorkloadVariable,
+    workload_iters: int,
+    experiment_dir: Path,
+    num_runs: int,
+    dry_run: bool,
+): ...
+@dataclass
+class SweepServeWorkloadArgs(SweepServeArgs):
+    workload_var: WorkloadVariable
+    workload_iters: int
+    parser_name: ClassVar[str] = ...
+    parser_help: ClassVar[str] = ...
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace): ...
+    @classmethod
+    def add_cli_args(
+        cls, parser: argparse.ArgumentParser
+    ) -> argparse.ArgumentParser: ...
+
+def run_main(args: SweepServeWorkloadArgs): ...
+def main(args: argparse.Namespace): ...

.cuda_typings/vllm/benchmarks/sweep/server.pyi

@@ -0,0 +1,26 @@
+from _typeshed import Incomplete
+from types import TracebackType
+from typing_extensions import Self
+
+class ServerProcess:
+    VLLM_RESET_CACHE_ENDPOINTS: Incomplete
+    server_cmd: Incomplete
+    after_bench_cmd: Incomplete
+    show_stdout: Incomplete
+    def __init__(
+        self, server_cmd: list[str], after_bench_cmd: list[str], *, show_stdout: bool
+    ) -> None: ...
+    def __enter__(self) -> Self: ...
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        exc_traceback: TracebackType | None,
+    ) -> None: ...
+    def start(self) -> None: ...
+    def stop(self) -> None: ...
+    def run_subcommand(self, cmd: list[str]): ...
+    def after_bench(self) -> None: ...
+    def is_server_ready(self) -> bool: ...
+    def wait_until_ready(self, timeout: int) -> None: ...
+    def reset_caches(self) -> None: ...

.cuda_typings/vllm/benchmarks/sweep/startup.pyi

@@ -0,0 +1,69 @@
+import argparse
+import pandas as pd
+from .param_sweep import (
+    ParameterSweep as ParameterSweep,
+    ParameterSweepItem as ParameterSweepItem,
+)
+from .utils import sanitize_filename as sanitize_filename
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import ClassVar
+from vllm.utils.argparse_utils import FlexibleArgumentParser as FlexibleArgumentParser
+from vllm.utils.import_utils import PlaceholderModule as PlaceholderModule
+
+def run_benchmark(
+    startup_cmd: list[str],
+    *,
+    serve_overrides: ParameterSweepItem,
+    startup_overrides: ParameterSweepItem,
+    run_number: int,
+    output_path: Path,
+    show_stdout: bool,
+    dry_run: bool,
+) -> dict[str, object] | None: ...
+def run_comb(
+    startup_cmd: list[str],
+    *,
+    serve_comb: ParameterSweepItem,
+    startup_comb: ParameterSweepItem,
+    base_path: Path,
+    num_runs: int,
+    show_stdout: bool,
+    dry_run: bool,
+) -> list[dict[str, object]] | None: ...
+def run_combs(
+    startup_cmd: list[str],
+    *,
+    serve_params: ParameterSweep,
+    startup_params: ParameterSweep,
+    experiment_dir: Path,
+    num_runs: int,
+    show_stdout: bool,
+    dry_run: bool,
+) -> pd.DataFrame | None: ...
+@dataclass
+class SweepStartupArgs:
+    startup_cmd: list[str]
+    serve_params: ParameterSweep
+    startup_params: ParameterSweep
+    output_dir: Path
+    experiment_name: str
+    num_runs: int
+    show_stdout: bool
+    dry_run: bool
+    resume: bool
+    parser_name: ClassVar[str] = ...
+    parser_help: ClassVar[str] = ...
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace): ...
+    @classmethod
+    def add_cli_args(
+        cls, parser: argparse.ArgumentParser
+    ) -> argparse.ArgumentParser: ...
+    def resolve_experiment_dir(self) -> Path: ...
+    @contextmanager
+    def run_ctx(self, experiment_dir: Path): ...
+
+def run_main(args: SweepStartupArgs): ...
+def main(args: argparse.Namespace): ...

.cuda_typings/vllm/benchmarks/sweep/utils.pyi

@@ -0,0 +1 @@
+def sanitize_filename(filename: str) -> str: ...

.cuda_typings/vllm/benchmarks/throughput.pyi

@@ -0,0 +1,80 @@
+import argparse
+import torch
+from typing import Any
+from vllm.benchmarks.datasets import (
+    AIMODataset as AIMODataset,
+    BurstGPTDataset as BurstGPTDataset,
+    ConversationDataset as ConversationDataset,
+    InstructCoderDataset as InstructCoderDataset,
+    MultiModalConversationDataset as MultiModalConversationDataset,
+    PrefixRepetitionRandomDataset as PrefixRepetitionRandomDataset,
+    RandomDataset as RandomDataset,
+    RandomDatasetForReranking as RandomDatasetForReranking,
+    RandomMultiModalDataset as RandomMultiModalDataset,
+    SampleRequest as SampleRequest,
+    ShareGPTDataset as ShareGPTDataset,
+    SonnetDataset as SonnetDataset,
+    VisionArenaDataset as VisionArenaDataset,
+    add_random_dataset_base_args as add_random_dataset_base_args,
+    add_random_multimodal_dataset_args as add_random_multimodal_dataset_args,
+)
+from vllm.benchmarks.lib.utils import (
+    convert_to_pytorch_benchmark_format as convert_to_pytorch_benchmark_format,
+    write_to_json as write_to_json,
+)
+from vllm.engine.arg_utils import (
+    AsyncEngineArgs as AsyncEngineArgs,
+    EngineArgs as EngineArgs,
+)
+from vllm.inputs import TextPrompt as TextPrompt, TokensPrompt as TokensPrompt
+from vllm.lora.request import LoRARequest as LoRARequest
+from vllm.outputs import RequestOutput as RequestOutput
+from vllm.platforms import current_platform as current_platform
+from vllm.sampling_params import BeamSearchParams as BeamSearchParams
+from vllm.tokenizers import (
+    TokenizerLike as TokenizerLike,
+    get_tokenizer as get_tokenizer,
+)
+from vllm.utils.async_utils import merge_async_iterators as merge_async_iterators
+
+def run_vllm(
+    requests: list[SampleRequest],
+    n: int,
+    engine_args: EngineArgs,
+    do_profile: bool,
+    disable_detokenize: bool = False,
+) -> tuple[float, list[RequestOutput] | None]: ...
+def run_vllm_chat(
+    requests: list[SampleRequest],
+    n: int,
+    engine_args: EngineArgs,
+    do_profile: bool,
+    disable_detokenize: bool = False,
+) -> tuple[float, list[RequestOutput]]: ...
+async def run_vllm_async(
+    requests: list[SampleRequest],
+    n: int,
+    engine_args: AsyncEngineArgs,
+    do_profile: bool,
+    disable_frontend_multiprocessing: bool = False,
+    disable_detokenize: bool = False,
+) -> float: ...
+def run_hf(
+    requests: list[SampleRequest],
+    model: str,
+    tokenizer: TokenizerLike,
+    n: int,
+    max_batch_size: int,
+    trust_remote_code: bool,
+    disable_detokenize: bool = False,
+    dtype: torch.dtype | None = ...,
+    enable_torch_compile: bool = False,
+) -> float: ...
+def save_to_pytorch_benchmark_format(
+    args: argparse.Namespace, results: dict[str, Any]
+) -> None: ...
+def get_requests(args, tokenizer): ...
+def filter_requests_for_dp(requests, data_parallel_size): ...
+def validate_args(args) -> None: ...
+def add_cli_args(parser: argparse.ArgumentParser): ...
+def main(args: argparse.Namespace): ...

.cuda_typings/vllm/collect_env.pyi

@@ -0,0 +1,79 @@
+from _typeshed import Incomplete
+from typing import NamedTuple
+from vllm.envs import environment_variables as environment_variables
+
+TORCH_AVAILABLE: bool
+
+class SystemEnv(NamedTuple):
+    torch_version: Incomplete
+    is_debug_build: Incomplete
+    cuda_compiled_version: Incomplete
+    gcc_version: Incomplete
+    clang_version: Incomplete
+    cmake_version: Incomplete
+    os: Incomplete
+    libc_version: Incomplete
+    python_version: Incomplete
+    python_platform: Incomplete
+    is_cuda_available: Incomplete
+    cuda_runtime_version: Incomplete
+    cuda_module_loading: Incomplete
+    nvidia_driver_version: Incomplete
+    nvidia_gpu_models: Incomplete
+    cudnn_version: Incomplete
+    pip_version: Incomplete
+    pip_packages: Incomplete
+    conda_packages: Incomplete
+    hip_compiled_version: Incomplete
+    hip_runtime_version: Incomplete
+    miopen_runtime_version: Incomplete
+    caching_allocator_config: Incomplete
+    is_xnnpack_available: Incomplete
+    cpu_info: Incomplete
+    rocm_version: Incomplete
+    vllm_version: Incomplete
+    vllm_build_flags: Incomplete
+    gpu_topo: Incomplete
+    env_vars: Incomplete
+
+DEFAULT_CONDA_PATTERNS: Incomplete
+DEFAULT_PIP_PATTERNS: Incomplete
+
+def run(command): ...
+def run_and_read_all(run_lambda, command): ...
+def run_and_parse_first_match(run_lambda, command, regex): ...
+def get_conda_packages(run_lambda, patterns=None): ...
+def get_gcc_version(run_lambda): ...
+def get_clang_version(run_lambda): ...
+def get_cmake_version(run_lambda): ...
+def get_nvidia_driver_version(run_lambda): ...
+def get_gpu_info(run_lambda): ...
+def get_running_cuda_version(run_lambda): ...
+def get_cudnn_version(run_lambda): ...
+def get_nvidia_smi(): ...
+def get_rocm_version(run_lambda): ...
+def get_vllm_version(): ...
+def summarize_vllm_build_flags(): ...
+def get_gpu_topo(run_lambda): ...
+def get_cpu_info(run_lambda): ...
+def get_platform(): ...
+def get_mac_version(run_lambda): ...
+def get_windows_version(run_lambda): ...
+def get_lsb_version(run_lambda): ...
+def check_release_file(run_lambda): ...
+def get_os(run_lambda): ...
+def get_python_platform(): ...
+def get_libc_version(): ...
+def is_uv_venv(): ...
+def get_pip_packages(run_lambda, patterns=None): ...
+def get_cachingallocator_config(): ...
+def get_cuda_module_loading_config(): ...
+def is_xnnpack_available(): ...
+def get_env_vars(): ...
+def get_env_info(): ...
+
+env_info_fmt: Incomplete
+
+def pretty_str(envinfo): ...
+def get_pretty_env_info(): ...
+def main() -> None: ...

.cuda_typings/vllm/compilation/backends.pyi

@@ -0,0 +1,158 @@
+import dataclasses
+import torch
+import torch.fx as fx
+from .compiler_interface import (
+    CompilerInterface as CompilerInterface,
+    EagerAdaptor as EagerAdaptor,
+    InductorAdaptor as InductorAdaptor,
+    InductorStandaloneAdaptor as InductorStandaloneAdaptor,
+    is_compile_cache_enabled as is_compile_cache_enabled,
+)
+from .counter import compilation_counter as compilation_counter
+from .partition_rules import (
+    inductor_partition_rule_context as inductor_partition_rule_context,
+    should_split as should_split,
+)
+from .passes.inductor_pass import (
+    InductorPass as InductorPass,
+    pass_context as pass_context,
+)
+from .passes.pass_manager import PostGradPassManager as PostGradPassManager
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable, Generator, Sequence
+from contextlib import contextmanager
+from typing import Any
+from vllm.config import (
+    CUDAGraphMode as CUDAGraphMode,
+    CompilationConfig as CompilationConfig,
+    VllmConfig as VllmConfig,
+)
+from vllm.config.compilation import DynamicShapesType as DynamicShapesType
+from vllm.config.utils import Range as Range, hash_factors as hash_factors
+from vllm.logger import init_logger as init_logger
+from vllm.logging_utils import lazy as lazy
+from vllm.platforms import current_platform as current_platform
+from vllm.tracing import (
+    instrument as instrument,
+    instrument_manual as instrument_manual,
+)
+from vllm.utils.import_utils import resolve_obj_by_qualname as resolve_obj_by_qualname
+
+logger: Incomplete
+
+def make_copy_and_call(
+    sym_tensor_indices: list[int],
+    input_buffers: list[torch.Tensor | None],
+    callable_fn: Callable[..., Any],
+) -> Callable[..., Any]: ...
+def make_compiler(compilation_config: CompilationConfig) -> CompilerInterface: ...
+
+class CompilerManager:
+    cache: dict[tuple[Range, int, str], Any]
+    is_cache_updated: bool
+    compilation_config: Incomplete
+    compiler: Incomplete
+    loaded_artifacts: dict[str, Any]
+    def __init__(self, compilation_config: CompilationConfig) -> None: ...
+    def compute_hash(self, vllm_config: VllmConfig) -> str: ...
+    @contextmanager
+    def compile_context(self, compile_range: Range) -> Generator[None, None, None]: ...
+    disable_cache: Incomplete
+    cache_dir: Incomplete
+    cache_file_path: Incomplete
+    def initialize_cache(
+        self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
+    ) -> None: ...
+    def save_to_file(self) -> None: ...
+    def load(
+        self,
+        graph: fx.GraphModule,
+        example_inputs: list[Any],
+        graph_index: int,
+        compile_range: Range,
+    ) -> Callable[..., Any] | None: ...
+    def compile(
+        self,
+        graph: fx.GraphModule,
+        example_inputs: list[Any],
+        additional_inductor_config: dict[str, Any],
+        compilation_config: CompilationConfig,
+        compile_range: Range,
+        graph_index: int = 0,
+        num_graphs: int = 1,
+    ) -> Any: ...
+
+class StopCompiling(BaseException): ...
+
+@dataclasses.dataclass
+class SplitItem:
+    submod_name: str
+    graph_id: int
+    is_splitting_graph: bool
+    graph: fx.GraphModule
+
+def split_graph(
+    graph: fx.GraphModule, splitting_ops: list[str]
+) -> tuple[fx.GraphModule, list[SplitItem]]: ...
+
+compilation_start_time: float
+
+def wrap_with_cudagraph_if_needed(
+    piecewise_backend: Any,
+    vllm_config: VllmConfig,
+    compilation_config: CompilationConfig,
+    is_first_graph: bool,
+    is_last_graph: bool,
+) -> Any: ...
+
+class PiecewiseCompileInterpreter(torch.fx.Interpreter):
+    compile_submod_names: Incomplete
+    compilation_config: Incomplete
+    vllm_config: Incomplete
+    vllm_backend: Incomplete
+    extra_traceback: bool
+    def __init__(
+        self,
+        module: torch.fx.GraphModule,
+        compile_submod_names: list[str],
+        vllm_config: VllmConfig,
+        vllm_backend: VllmBackend,
+    ) -> None: ...
+    def run(self, *args: Any) -> Any: ...
+    def call_module(
+        self,
+        target: torch.fx.node.Target,
+        args: tuple[torch.fx.node.Argument, ...],
+        kwargs: dict[str, Any],
+    ) -> Any: ...
+
+model_tag: str
+model_is_encoder: bool
+
+@contextmanager
+def set_model_tag(
+    tag: str, is_encoder: bool = False
+) -> Generator[None, None, None]: ...
+
+class VllmBackend:
+    vllm_config: VllmConfig
+    compilation_config: CompilationConfig
+    graph: fx.GraphModule
+    split_gm: fx.GraphModule
+    piecewise_graphs: list[SplitItem]
+    returned_callable: Callable[..., Any]
+    post_grad_passes: Sequence[Callable[..., Any]]
+    compiler_manager: CompilerManager
+    inductor_config: dict[str, Any]
+    prefix: Incomplete
+    is_encoder: Incomplete
+    pass_manager: Incomplete
+    pass_key: Incomplete
+    def __init__(
+        self, vllm_config: VllmConfig, prefix: str = "", is_encoder: bool = False
+    ) -> None: ...
+    def collect_standalone_compile_artifacts(
+        self,
+    ) -> tuple[Any, dict[str, list[int]] | None, dict[str, bool] | None]: ...
+    def configure_post_pass(self) -> None: ...
+    def __call__(self, graph: fx.GraphModule, example_inputs: Sequence[Any]) -> Any: ...

.cuda_typings/vllm/compilation/base_static_graph.pyi

@@ -0,0 +1,13 @@
+from collections.abc import Callable as Callable
+from typing import Any, Protocol
+from vllm.config import CUDAGraphMode as CUDAGraphMode, VllmConfig as VllmConfig
+
+class AbstractStaticGraphWrapper(Protocol):
+    def __init__(
+        self,
+        runnable: Callable[..., Any],
+        vllm_config: VllmConfig,
+        runtime_mode: CUDAGraphMode,
+        **kwargs: Any,
+    ) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...

.cuda_typings/vllm/compilation/caching.pyi

@@ -0,0 +1,77 @@
+import contextlib
+import torch
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable, Generator, Sequence
+from torch._dynamo.aot_compile import SerializableCallable
+from typing import Any, Literal
+from vllm.compilation.compiler_interface import (
+    get_inductor_factors as get_inductor_factors,
+)
+from vllm.config import (
+    VllmConfig as VllmConfig,
+    get_current_vllm_config as get_current_vllm_config,
+)
+from vllm.config.utils import hash_factors as hash_factors
+from vllm.logger import init_logger as init_logger
+from vllm.utils.hashing import safe_hash as safe_hash
+
+SerializableCallable = object
+logger: Incomplete
+
+class StandaloneCompiledArtifacts:
+    submodule_bytes: dict[str, str]
+    submodule_bytes_store: dict[str, bytes]
+    loaded_submodule_store: dict[str, Any]
+    def __init__(self) -> None: ...
+    def insert(self, submod_name: str, shape: str, entry: bytes) -> None: ...
+    def get(self, submod_name: str, shape: str) -> bytes: ...
+    def get_loaded(self, submod_name: str, shape: str) -> Any: ...
+    def size_bytes(self) -> int: ...
+    def num_artifacts(self) -> int: ...
+    def num_entries(self) -> int: ...
+    def submodule_names(self) -> list[str]: ...
+    def load_all(self) -> None: ...
+
+@contextlib.contextmanager
+def patch_pytree_map_over_slice() -> Generator[None, None, Incomplete]: ...
+
+class VllmSerializableFunction(SerializableCallable):
+    graph_module: Incomplete
+    example_inputs: Incomplete
+    prefix: Incomplete
+    optimized_call: Incomplete
+    is_encoder: Incomplete
+    shape_env: Incomplete
+    vllm_backend: Incomplete
+    sym_tensor_indices: Incomplete
+    aot_autograd_config: Incomplete
+    def __init__(
+        self,
+        graph_module: torch.fx.GraphModule,
+        example_inputs: Sequence[Any],
+        prefix: str,
+        optimized_call: Callable[..., Any],
+        is_encoder: bool = False,
+        vllm_backend: Any | None = None,
+        sym_tensor_indices: list[int] | None = None,
+        aot_autograd_config: dict[str, Any] | None = None,
+    ) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+    @classmethod
+    def serialize_compile_artifacts(
+        cls, compiled_fn: VllmSerializableFunction
+    ) -> bytes: ...
+    @classmethod
+    def deserialize_compile_artifacts(cls, data: bytes) -> VllmSerializableFunction: ...
+    def finalize_loading(self, vllm_config: VllmConfig) -> None: ...
+    @property
+    def co_name(self) -> Literal["VllmSerializableFunction"]: ...
+
+def reconstruct_serializable_fn_from_mega_artifact(
+    state: dict[str, Any],
+    standalone_compile_artifacts: StandaloneCompiledArtifacts,
+    vllm_config: VllmConfig,
+    sym_shape_indices_map: dict[str, list[int]],
+    returns_tuple_map: dict[str, bool],
+) -> VllmSerializableFunction: ...
+def aot_compile_hash_factors(vllm_config: VllmConfig) -> list[str]: ...

.cuda_typings/vllm/compilation/compiler_interface.pyi

@@ -0,0 +1,117 @@
+import contextlib
+import torch.fx as fx
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable
+from typing import Any, Literal
+from vllm.compilation.counter import compilation_counter as compilation_counter
+from vllm.config import VllmConfig as VllmConfig
+from vllm.config.utils import Range as Range
+from vllm.logger import init_logger as init_logger
+from vllm.utils.hashing import safe_hash as safe_hash
+from vllm.utils.torch_utils import is_torch_equal_or_newer as is_torch_equal_or_newer
+
+logger: Incomplete
+
+class CompilerInterface:
+    name: str
+    def initialize_cache(
+        self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
+    ) -> None: ...
+    def compute_hash(self, vllm_config: VllmConfig) -> str: ...
+    def compile(
+        self,
+        graph: fx.GraphModule,
+        example_inputs: list[Any],
+        compiler_config: dict[str, Any],
+        compile_range: Range,
+        key: str | None = None,
+    ) -> tuple[Callable[..., Any] | None, Any | None]: ...
+    def load(
+        self,
+        handle: Any,
+        graph: fx.GraphModule,
+        example_inputs: list[Any],
+        graph_index: int,
+        compile_range: Range,
+    ) -> Callable[..., Any]: ...
+
+class AlwaysHitShapeEnv:
+    guards: list[Any]
+    def __init__(self) -> None: ...
+    def evaluate_guards_expression(
+        self, *args: Any, **kwargs: Any
+    ) -> Literal[True]: ...
+    def get_pruned_guards(self, *args: Any, **kwargs: Any) -> list[Any]: ...
+    def produce_guards_expression(self, *args: Any, **kwargs: Any) -> Literal[""]: ...
+
+def get_inductor_factors() -> list[Any]: ...
+def is_compile_cache_enabled(
+    vllm_additional_inductor_config: dict[str, Any],
+) -> bool: ...
+
+class InductorStandaloneAdaptor(CompilerInterface):
+    name: str
+    save_format: Incomplete
+    def __init__(self, save_format: Literal["binary", "unpacked"]) -> None: ...
+    def compute_hash(self, vllm_config: VllmConfig) -> str: ...
+    cache_dir: Incomplete
+    def initialize_cache(
+        self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
+    ) -> None: ...
+    def compile(
+        self,
+        graph: fx.GraphModule,
+        example_inputs: list[Any],
+        compiler_config: dict[str, Any],
+        compile_range: Range,
+        key: str | None = None,
+    ) -> tuple[Callable[..., Any] | None, Any | None]: ...
+    def load(
+        self,
+        handle: Any,
+        graph: fx.GraphModule,
+        example_inputs: list[Any],
+        graph_index: int,
+        compile_range: Range,
+    ) -> Callable[..., Any]: ...
+
+class InductorAdaptor(CompilerInterface):
+    name: str
+    def compute_hash(self, vllm_config: VllmConfig) -> str: ...
+    cache_dir: Incomplete
+    prefix: Incomplete
+    base_cache_dir: Incomplete
+    def initialize_cache(
+        self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
+    ) -> None: ...
+    def compile(
+        self,
+        graph: fx.GraphModule,
+        example_inputs: list[Any],
+        compiler_config: dict[str, Any],
+        compile_range: Range,
+        key: str | None = None,
+    ) -> tuple[Callable[..., Any] | None, Any | None]: ...
+    def load(
+        self,
+        handle: Any,
+        graph: fx.GraphModule,
+        example_inputs: list[Any],
+        graph_index: int,
+        compile_range: Range,
+    ) -> Callable[..., Any]: ...
+    def metrics_context(self) -> contextlib.AbstractContextManager[Any]: ...
+
+def set_inductor_config(config: dict[str, Any], compile_range: Range) -> None: ...
+def set_functorch_config() -> None: ...
+
+class EagerAdaptor(CompilerInterface):
+    name: str
+    def compile(
+        self,
+        graph: fx.GraphModule,
+        example_inputs: list[Any],
+        compiler_config: dict[str, Any],
+        compile_range: Range,
+        key: str | None = None,
+    ) -> tuple[Callable[..., Any] | None, Any | None]: ...

.cuda_typings/vllm/compilation/counter.pyi

@@ -0,0 +1,26 @@
+import dataclasses
+from _typeshed import Incomplete
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+@dataclasses.dataclass
+class CompilationCounter:
+    num_models_seen: int = ...
+    num_graphs_seen: int = ...
+    num_piecewise_graphs_seen: int = ...
+    num_piecewise_capturable_graphs_seen: int = ...
+    num_backend_compilations: int = ...
+    num_gpu_runner_capture_triggers: int = ...
+    num_cudagraph_captured: int = ...
+    num_inductor_compiles: int = ...
+    num_eager_compiles: int = ...
+    num_cache_entries_updated: int = ...
+    num_compiled_artifacts_saved: int = ...
+    num_compiled_artifacts_loaded: int = ...
+    stock_torch_compile_count: int = ...
+    def clone(self) -> CompilationCounter: ...
+    @contextmanager
+    def expect(self, **kwargs: Any) -> Generator[None, None, None]: ...
+
+compilation_counter: Incomplete

.cuda_typings/vllm/compilation/cuda_graph.pyi

@@ -0,0 +1,86 @@
+import dataclasses
+import torch
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable
+from typing import Any
+from vllm.compilation.counter import compilation_counter as compilation_counter
+from vllm.compilation.monitor import (
+    validate_cudagraph_capturing_enabled as validate_cudagraph_capturing_enabled,
+)
+from vllm.config import CUDAGraphMode as CUDAGraphMode, VllmConfig as VllmConfig
+from vllm.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id as set_graph_pool_id,
+)
+from vllm.forward_context import (
+    BatchDescriptor as BatchDescriptor,
+    get_forward_context as get_forward_context,
+)
+from vllm.logger import init_logger as init_logger
+from vllm.model_executor.offloader.base import get_offloader as get_offloader
+from vllm.platforms import current_platform as current_platform
+from vllm.utils.torch_utils import (
+    current_stream as current_stream,
+    weak_ref_tensors as weak_ref_tensors,
+)
+
+logger: Incomplete
+
+@dataclasses.dataclass(frozen=True)
+class CUDAGraphStat:
+    num_unpadded_tokens: int
+    num_padded_tokens: int
+    num_paddings: int
+    runtime_mode: str
+
+class CUDAGraphLogging:
+    COLUMN_HEADERS: Incomplete
+    cg_mode: Incomplete
+    cg_capture_sizes: Incomplete
+    settings_header: Incomplete
+    def __init__(
+        self, cg_mode: CUDAGraphMode, cg_capture_sizes: list[int] | None
+    ) -> None: ...
+    stats: list[CUDAGraphStat]
+    def reset(self) -> None: ...
+    def observe(self, cudagraph_stat: CUDAGraphStat) -> None: ...
+    def generate_metric_table(self) -> str: ...
+    def log(self, log_fn: Callable[..., Any] = ...) -> None: ...
+
+@dataclasses.dataclass
+class CUDAGraphEntry:
+    batch_descriptor: BatchDescriptor
+    cudagraph: torch.cuda.CUDAGraph | None = ...
+    output: Any | None = ...
+    input_addresses: list[int] | None = ...
+
+@dataclasses.dataclass
+class CUDAGraphOptions:
+    debug_log_enable: bool = ...
+    gc_disable: bool = ...
+    weak_ref_output: bool = ...
+
+class CUDAGraphWrapper:
+    @classmethod
+    def clear_all_graphs(cls) -> None: ...
+    runnable: Incomplete
+    vllm_config: Incomplete
+    runtime_mode: Incomplete
+    compilation_config: Incomplete
+    first_run_finished: bool
+    is_debugging_mode: Incomplete
+    graph_pool: Incomplete
+    cudagraph_options: Incomplete
+    concrete_cudagraph_entries: dict[BatchDescriptor, CUDAGraphEntry]
+    def __init__(
+        self,
+        runnable: Callable[..., Any],
+        vllm_config: VllmConfig,
+        runtime_mode: CUDAGraphMode,
+        cudagraph_options: CUDAGraphOptions | None = None,
+    ) -> None: ...
+    def __getattr__(self, key: str) -> Any: ...
+    def unwrap(self) -> Callable[..., Any]: ...
+    @property
+    def cudagraph_wrapper(self) -> CUDAGraphWrapper: ...
+    def clear_graphs(self) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> Any | None: ...

.cuda_typings/vllm/compilation/decorators.pyi

@@ -0,0 +1,58 @@
+import contextlib
+from .monitor import (
+    monitor_profiling_run as monitor_profiling_run,
+    monitor_torch_compile as monitor_torch_compile,
+)
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable, Generator
+from typing import Any, overload
+from vllm.compilation.counter import compilation_counter as compilation_counter
+from vllm.compilation.wrapper import (
+    TorchCompileWithNoGuardsWrapper as TorchCompileWithNoGuardsWrapper,
+)
+from vllm.config import (
+    CompilationMode as CompilationMode,
+    VllmConfig as VllmConfig,
+    get_current_vllm_config as get_current_vllm_config,
+    set_current_vllm_config as set_current_vllm_config,
+)
+from vllm.config.compilation import DynamicShapesType as DynamicShapesType
+from vllm.forward_context import (
+    get_forward_context as get_forward_context,
+    is_forward_context_available as is_forward_context_available,
+)
+from vllm.logger import init_logger as init_logger
+from vllm.sequence import IntermediateTensors as IntermediateTensors
+from vllm.utils.import_utils import resolve_obj_by_qualname as resolve_obj_by_qualname
+from vllm.utils.torch_utils import is_torch_equal_or_newer as is_torch_equal_or_newer
+
+SourceInfo = Any
+logger: Incomplete
+IGNORE_COMPILE_KEY: str
+
+def should_torch_compile_mm_encoder(vllm_config: VllmConfig) -> bool: ...
+def ignore_torch_compile(cls) -> type[_T]: ...
+@overload
+def support_torch_compile(
+    *, enable_if: Callable[[VllmConfig], bool] | None = None
+) -> Callable[[type[_T]], type[_T]]: ...
+@overload
+def support_torch_compile(
+    *, dynamic_arg_dims: dict[str, int | list[int]] | None
+) -> Callable[[type[_T]], type[_T]]: ...
+@overload
+def support_torch_compile(
+    *, mark_unbacked_dims: dict[str, int | list[int]] | None
+) -> Callable[[type[_T]], type[_T]]: ...
+@overload
+def support_torch_compile(
+    *,
+    dynamic_arg_dims: dict[str, int | list[int]] | None,
+    mark_unbacked_dims: dict[str, int | list[int]] | None,
+) -> Callable[[type[_T]], type[_T]]: ...
+@overload
+def support_torch_compile(cls) -> type[_T]: ...
+@contextlib.contextmanager
+def maybe_use_cudagraph_partition_wrapper(
+    vllm_config: VllmConfig,
+) -> Generator[None, None, None]: ...

.cuda_typings/vllm/compilation/monitor.pyi

@@ -0,0 +1,20 @@
+import contextlib
+from _typeshed import Incomplete
+from collections.abc import Generator
+from vllm.config import CompilationMode as CompilationMode, VllmConfig as VllmConfig
+from vllm.logger import init_logger as init_logger
+
+logger: Incomplete
+torch_compile_start_time: float
+
+@contextlib.contextmanager
+def monitor_torch_compile(
+    vllm_config: VllmConfig, message: str = "torch.compile took %.2f s in total"
+) -> Generator[None, None, None]: ...
+@contextlib.contextmanager
+def monitor_profiling_run() -> Generator[None, None, None]: ...
+
+cudagraph_capturing_enabled: bool
+
+def validate_cudagraph_capturing_enabled() -> None: ...
+def set_cudagraph_capturing_enabled(enabled: bool) -> None: ...

.cuda_typings/vllm/compilation/partition_rules.pyi

@@ -0,0 +1,13 @@
+import contextlib
+import torch
+from _typeshed import Incomplete
+from collections.abc import Generator
+from vllm.logger import init_logger as init_logger
+
+logger: Incomplete
+
+def should_split(node: torch.fx.Node, splitting_ops: list[str]) -> bool: ...
+@contextlib.contextmanager
+def inductor_partition_rule_context(
+    splitting_ops: list[str] | None,
+) -> Generator[None, None, None]: ...

.cuda_typings/vllm/compilation/passes/fusion/act_quant_fusion.pyi

@@ -0,0 +1,68 @@
+import abc
+import torch
+from ..inductor_pass import enable_fake_mode as enable_fake_mode
+from ..vllm_inductor_pass import (
+    VllmInductorPass as VllmInductorPass,
+    VllmPatternMatcherPass as VllmPatternMatcherPass,
+)
+from .matcher_utils import (
+    MatcherQuantFP8 as MatcherQuantFP8,
+    MatcherSiluAndMul as MatcherSiluAndMul,
+)
+from .rms_quant_fusion import (
+    QUANT_OPS as QUANT_OPS,
+    empty_bf16 as empty_bf16,
+    empty_fp32 as empty_fp32,
+    empty_i32 as empty_i32,
+)
+from _typeshed import Incomplete
+from abc import ABC, abstractmethod
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from torch._ops import OpOverload as OpOverload
+from typing import Any
+from vllm.config import VllmConfig as VllmConfig
+from vllm.logger import init_logger as init_logger
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey as QuantKey,
+    kFp8StaticTensorSym as kFp8StaticTensorSym,
+    kNvfp4Dynamic as kNvfp4Dynamic,
+)
+from vllm.platforms import current_platform as current_platform
+
+logger: Incomplete
+FP8_DTYPE: Incomplete
+FP4_DTYPE: Incomplete
+SILU_MUL_OP: Incomplete
+FUSED_OPS: dict[QuantKey, OpOverload]
+silu_and_mul_nvfp4_quant_supported: Incomplete
+
+class ActivationQuantPattern(ABC, metaclass=abc.ABCMeta):
+    quant_key: Incomplete
+    quant_dtype: Incomplete
+    QUANT_OP: Incomplete
+    FUSED_OP: Incomplete
+    silu_and_mul_matcher: Incomplete
+    def __init__(self, quant_key: QuantKey) -> None: ...
+    def empty_quant(self, *args: Any, **kwargs: Any) -> torch.Tensor: ...
+    @abstractmethod
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class SiluMulFp8StaticQuantPattern(ActivationQuantPattern):
+    quant_matcher: Incomplete
+    def __init__(self) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
+    def __init__(self) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class ActivationQuantFusionPass(VllmPatternMatcherPass):
+    patterns: PatternMatcherPass
+    @enable_fake_mode
+    def __init__(self, config: VllmConfig) -> None: ...
+    matched_count: Incomplete
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: torch.fx.Graph) -> None: ...
+    def uuid(self) -> str: ...

.cuda_typings/vllm/compilation/passes/fusion/allreduce_rms_fusion.pyi

@@ -0,0 +1,201 @@
+import torch
+import torch.fx as fx
+from ..inductor_pass import enable_fake_mode as enable_fake_mode
+from ..vllm_inductor_pass import (
+    VllmInductorPass as VllmInductorPass,
+    VllmPatternMatcherPass as VllmPatternMatcherPass,
+)
+from .matcher_utils import (
+    MatcherFusedAddRMSNorm as MatcherFusedAddRMSNorm,
+    MatcherQuantFP8 as MatcherQuantFP8,
+    MatcherRMSNorm as MatcherRMSNorm,
+)
+from _typeshed import Incomplete
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from types import ModuleType
+from vllm.config import VllmConfig as VllmConfig
+from vllm.config.utils import Range as Range
+from vllm.distributed import (
+    get_tp_group as get_tp_group,
+    tensor_model_parallel_all_reduce as tensor_model_parallel_all_reduce,
+)
+from vllm.distributed.device_communicators.flashinfer_all_reduce import (
+    destroy_fi_ar_workspace as destroy_fi_ar_workspace,
+    get_fi_ar_quant_workspace as get_fi_ar_quant_workspace,
+    get_fi_ar_workspace as get_fi_ar_workspace,
+    initialize_fi_ar_quant_workspace as initialize_fi_ar_quant_workspace,
+    initialize_fi_ar_workspace as initialize_fi_ar_workspace,
+)
+from vllm.distributed.parallel_state import (
+    get_tensor_model_parallel_rank as get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size as get_tensor_model_parallel_world_size,
+)
+from vllm.logger import init_logger as init_logger
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kFp8StaticTensorSym as kFp8StaticTensorSym,
+)
+from vllm.platforms import current_platform as current_platform
+from vllm.utils.torch_utils import (
+    direct_register_custom_op as direct_register_custom_op,
+)
+
+FP8_DTYPE: Incomplete
+logger: Incomplete
+flashinfer_comm: ModuleType | None
+STATIC_FP4_QUANT_OP: Incomplete
+FI_ALLREDUCE_FUSION_MAX_SIZE_MB: dict[int, dict[int, float]]
+ar_fusion_patterns: Incomplete
+MiB: Incomplete
+
+def call_trtllm_fused_allreduce_norm(
+    allreduce_in: torch.Tensor,
+    residual: torch.Tensor,
+    rms_gamma: torch.Tensor,
+    rms_eps: float,
+    world_size: int,
+    launch_with_pdl: bool,
+    fp32_acc: bool,
+    max_token_num: int,
+    pattern_code: int,
+    norm_out: torch.Tensor | None = None,
+    quant_out: torch.Tensor | None = None,
+    scale_out: torch.Tensor | None = None,
+    scale_factor: torch.Tensor | None = None,
+) -> None: ...
+def call_trtllm_fused_allreduce_norm_fake(
+    allreduce_in: torch.Tensor,
+    residual: torch.Tensor,
+    rms_gamma: torch.Tensor,
+    rms_eps: float,
+    world_size: int,
+    launch_with_pdl: bool,
+    fp32_acc: bool,
+    max_token_num: int,
+    pattern_code: int,
+    norm_out: torch.Tensor | None = None,
+    quant_out: torch.Tensor | None = None,
+    scale_out: torch.Tensor | None = None,
+    scale_factor: torch.Tensor | None = None,
+) -> None: ...
+
+flashinfer_trtllm_fused_allreduce_norm: Incomplete
+
+class FlashInferFusedAllReduceParams:
+    world_size: Incomplete
+    launch_with_pdl: bool
+    fp32_acc: bool
+    max_token_num: Incomplete
+    def __init__(self, world_size: int, max_token_num: int = 1024) -> None: ...
+    def get_trtllm_fused_allreduce_kwargs(self) -> dict[str, bool | int]: ...
+
+class BasePattern:
+    dtype: Incomplete
+    device: Incomplete
+    tp: Incomplete
+    tp_size: Incomplete
+    def __init__(self, dtype: torch.dtype, device: str | None) -> None: ...
+
+class AllReduceRMSNormPattern(BasePattern):
+    epsilon: Incomplete
+    allreduce_params: Incomplete
+    rmsnorm_matcher: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        dtype: torch.dtype,
+        device: str | None,
+        allreduce_params: FlashInferFusedAllReduceParams,
+    ) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AllReduceFusedAddRMSNormPattern(BasePattern):
+    epsilon: Incomplete
+    allreduce_params: Incomplete
+    rmsnorm_matcher: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        dtype: torch.dtype,
+        device: str | None,
+        allreduce_params: FlashInferFusedAllReduceParams,
+    ) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):
+    epsilon: Incomplete
+    allreduce_params: Incomplete
+    quant_dtype: Incomplete
+    rmsnorm_matcher: Incomplete
+    quant_matcher: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        dtype: torch.dtype,
+        device: str | None,
+        allreduce_params: FlashInferFusedAllReduceParams,
+    ) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
+    epsilon: Incomplete
+    allreduce_params: Incomplete
+    quant_dtype: Incomplete
+    rmsnorm_matcher: Incomplete
+    quant_matcher: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        dtype: torch.dtype,
+        device: str | None,
+        allreduce_params: FlashInferFusedAllReduceParams,
+    ) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
+    epsilon: Incomplete
+    allreduce_params: Incomplete
+    rmsnorm_matcher: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        dtype: torch.dtype,
+        device: str | None,
+        allreduce_params: FlashInferFusedAllReduceParams,
+    ) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
+    epsilon: Incomplete
+    allreduce_params: Incomplete
+    rmsnorm_matcher: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        dtype: torch.dtype,
+        device: str | None,
+        allreduce_params: FlashInferFusedAllReduceParams,
+    ) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AllReduceFusionPass(VllmPatternMatcherPass):
+    disabled: bool
+    tp_size: Incomplete
+    patterns: PatternMatcherPass
+    hidden_dim: Incomplete
+    group: Incomplete
+    max_token_num: Incomplete
+    allreduce_params: Incomplete
+    def __init__(self, config: VllmConfig) -> None: ...
+    @enable_fake_mode
+    def register_patterns(self) -> None: ...
+    def is_applicable_for_range(self, compile_range: Range) -> bool: ...
+    matched_count: Incomplete
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None: ...
+    def __del__(self) -> None: ...

.cuda_typings/vllm/compilation/passes/fusion/attn_quant_fusion.pyi

@@ -0,0 +1,84 @@
+import abc
+import torch
+from ..fx_utils import is_func as is_func
+from ..inductor_pass import enable_fake_mode as enable_fake_mode
+from ..vllm_inductor_pass import (
+    VllmInductorPass as VllmInductorPass,
+    VllmPatternMatcherPass as VllmPatternMatcherPass,
+)
+from .matcher_utils import MatcherQuantFP8 as MatcherQuantFP8
+from .rms_quant_fusion import (
+    QUANT_OPS as QUANT_OPS,
+    empty_bf16 as empty_bf16,
+    empty_fp32 as empty_fp32,
+    empty_i32 as empty_i32,
+)
+from _typeshed import Incomplete
+from abc import ABC
+from collections.abc import Callable as Callable
+from torch import fx as fx
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from typing import Any, ParamSpec
+from vllm.config import (
+    VllmConfig as VllmConfig,
+    get_layers_from_vllm_config as get_layers_from_vllm_config,
+)
+from vllm.logger import init_logger as init_logger
+from vllm.model_executor.layers.attention import Attention as Attention
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey as QuantKey,
+    kNvfp4Dynamic as kNvfp4Dynamic,
+    kStaticTensorScale as kStaticTensorScale,
+)
+from vllm.platforms import current_platform as current_platform
+from vllm.utils.math_utils import round_up as round_up
+
+logger: Incomplete
+P = ParamSpec("P")
+FP8_DTYPE: Incomplete
+FP4_DTYPE: Incomplete
+ATTN_OP: Incomplete
+RESHAPE_OP: Incomplete
+
+class AttentionQuantPattern(ABC, metaclass=abc.ABCMeta):
+    layer: Incomplete
+    layer_name: Incomplete
+    num_heads: Incomplete
+    head_size: Incomplete
+    quant_key: Incomplete
+    quant_dtype: Incomplete
+    dtype: Incomplete
+    QUANT_OP: Incomplete
+    def __init__(
+        self, layer: Attention, quant_key: QuantKey, dtype: torch.dtype
+    ) -> None: ...
+    def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor: ...
+    def empty_quant(self, *args: Any, **kwargs: Any) -> torch.Tensor: ...
+    @staticmethod
+    def wrap_trace_fn(
+        trace_fn: Callable[P, fx.GraphModule],
+        *process_fx_fns: Callable[[fx.GraphModule], None],
+    ) -> Callable[P, fx.GraphModule]: ...
+    @staticmethod
+    def fx_view_to_reshape(gm: torch.fx.GraphModule) -> None: ...
+    @staticmethod
+    def remove_noop_permutes(gm: torch.fx.GraphModule) -> None: ...
+    def register_if_supported(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
+    quant_matcher: Incomplete
+    def __init__(
+        self, layer: Attention, dtype: torch.dtype, symmetric: bool = True
+    ) -> None: ...
+
+class AttentionNvfp4QuantPattern(AttentionQuantPattern):
+    def __init__(self, layer: Attention, dtype: torch.dtype) -> None: ...
+
+class AttnFusionPass(VllmPatternMatcherPass):
+    patterns: Incomplete
+    @enable_fake_mode
+    def __init__(self, config: VllmConfig) -> None: ...
+    matched_count: Incomplete
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: torch.fx.graph.Graph) -> None: ...
+    def uuid(self) -> str: ...

.cuda_typings/vllm/compilation/passes/fusion/collective_fusion.pyi

@@ -0,0 +1,60 @@
+import torch
+import torch.fx as fx
+from ..inductor_pass import enable_fake_mode as enable_fake_mode
+from ..vllm_inductor_pass import (
+    VllmInductorPass as VllmInductorPass,
+    VllmPatternMatcherPass as VllmPatternMatcherPass,
+)
+from _typeshed import Incomplete
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from vllm.config import VllmConfig as VllmConfig
+from vllm.config.utils import Range as Range
+from vllm.distributed import get_tp_group as get_tp_group
+from vllm.distributed.parallel_state import (
+    get_tensor_model_parallel_world_size as get_tensor_model_parallel_world_size,
+)
+from vllm.logger import init_logger as init_logger
+from vllm.platforms import current_platform as current_platform
+
+FP8_DTYPE: Incomplete
+logger: Incomplete
+
+class BasePattern:
+    dtype: Incomplete
+    device: Incomplete
+    tp: Incomplete
+    tp_size: Incomplete
+    def __init__(self, dtype: torch.dtype, device: str | None) -> None: ...
+
+class GEMMReduceScatterPattern(BasePattern):
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AllGatherGEMMPattern(BasePattern):
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class ScaledMMReduceScatterPattern(BasePattern):
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AllGatherScaledMMPattern(BasePattern):
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class CutlassScaledMMReduceScatterPattern(BasePattern):
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AllGatherCutlassScaledMMPattern(BasePattern):
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AsyncTPPass(VllmPatternMatcherPass):
+    patterns: PatternMatcherPass
+    @enable_fake_mode
+    def __init__(self, config: VllmConfig) -> None: ...
+    def is_applicable_for_range(self, compile_range: Range) -> bool: ...
+    matched_count: Incomplete
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None: ...

.cuda_typings/vllm/compilation/passes/fusion/matcher_utils.pyi

@@ -0,0 +1,160 @@
+import abc
+import torch
+from _typeshed import Incomplete
+from abc import ABC, abstractmethod
+from torch._ops import OpOverload as OpOverload
+from typing import Any
+from vllm._aiter_ops import rocm_aiter_ops as rocm_aiter_ops
+from vllm.config import get_current_vllm_config as get_current_vllm_config
+from vllm.model_executor.layers.activation import SiluAndMul as SiluAndMul
+from vllm.model_executor.layers.layernorm import RMSNorm as RMSNorm
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 as QuantFP8
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape as GroupShape,
+    QuantKey as QuantKey,
+    kFp8Dynamic128Sym as kFp8Dynamic128Sym,
+    kFp8Dynamic64Sym as kFp8Dynamic64Sym,
+    kFp8DynamicTensorSym as kFp8DynamicTensorSym,
+    kFp8DynamicTokenSym as kFp8DynamicTokenSym,
+    kFp8StaticTensorSym as kFp8StaticTensorSym,
+    kNvfp4Dynamic as kNvfp4Dynamic,
+)
+from vllm.model_executor.layers.rotary_embedding import (
+    RotaryEmbedding as RotaryEmbedding,
+)
+from vllm.platforms import current_platform as current_platform
+
+RMS_OP: Incomplete
+RMS_ADD_OP: Incomplete
+ROTARY_OP: Incomplete
+FLASHINFER_ROTARY_OP: Incomplete
+QUANT_OPS: dict[QuantKey, OpOverload]
+SILU_MUL_OP: Incomplete
+
+class MatcherCustomOp(ABC, metaclass=abc.ABCMeta):
+    model_dtype: Incomplete
+    device: Incomplete
+    enabled: Incomplete
+    forward: Incomplete
+    def __init__(self, enabled: bool) -> None: ...
+    @abstractmethod
+    def forward_custom(self, *args: Any, **kwargs: Any) -> Any: ...
+    @abstractmethod
+    def forward_native(self, *args: Any, **kwargs: Any) -> Any: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+    def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor: ...
+    def empty_int64(self, *args: Any, **kwargs: Any) -> torch.Tensor: ...
+    def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor: ...
+    def inputs(self) -> list[torch.Tensor]: ...
+
+class MatcherRotaryEmbedding(MatcherCustomOp):
+    is_neox: Incomplete
+    head_size: Incomplete
+    num_heads: Incomplete
+    num_kv_heads: Incomplete
+    q_size: Incomplete
+    kv_size: Incomplete
+    rotary_dim: Incomplete
+    rotary_op: Incomplete
+    def __init__(
+        self,
+        is_neox: bool,
+        head_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        use_flashinfer: bool = False,
+        match_rocm_aiter: bool | None = None,
+        enabled: bool | None = None,
+    ) -> None: ...
+    def inputs(self) -> list[torch.Tensor]: ...
+    def forward_custom(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor | None,
+        cos_sin_cache: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]: ...
+    def forward_native(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor | None,
+        cos_sin_cache: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]: ...
+
+class MatcherRMSNorm(MatcherCustomOp):
+    epsilon: Incomplete
+    match_rocm_aiter: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        enabled: bool | None = None,
+        match_rocm_aiter: bool = False,
+    ) -> None: ...
+    def inputs(self) -> list[torch.Tensor]: ...
+    def forward_rocm_aiter(
+        self, input: torch.Tensor, weight: torch.Tensor
+    ) -> torch.Tensor: ...
+    def forward_custom(
+        self, input: torch.Tensor, weight: torch.Tensor
+    ) -> torch.Tensor: ...
+    def forward_native(
+        self, input: torch.Tensor, weight: torch.Tensor
+    ) -> torch.Tensor: ...
+
+class MatcherFusedAddRMSNorm(MatcherCustomOp):
+    epsilon: Incomplete
+    match_rocm_aiter: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        enabled: bool | None = None,
+        match_rocm_aiter: bool = False,
+    ) -> None: ...
+    def inputs(self) -> list[torch.Tensor]: ...
+    def forward_rocm_aiter(
+        self, input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+    def forward_custom(
+        self, input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+    def forward_native(
+        self, input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+
+class MatcherQuantFP8(MatcherCustomOp):
+    quant_key: Incomplete
+    has_col_major_scales: Incomplete
+    is_e8m0: Incomplete
+    match_rocm_aiter: Incomplete
+    is_tma_aligned: Incomplete
+    QUANT_OP: Incomplete
+    quant_fp8: Incomplete
+    def __init__(
+        self,
+        quant_key: QuantKey,
+        enabled: bool | None = None,
+        has_col_major_scales: bool = False,
+        is_e8m0: bool = False,
+        match_rocm_aiter: bool = False,
+        is_tma_aligned: bool = False,
+    ) -> None: ...
+    def forward_rocm_aiter(
+        self, input: torch.Tensor, scale: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+    def forward_custom(
+        self, input: torch.Tensor, scale: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+    def forward_native(
+        self, input: torch.Tensor, scale: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+    def make_scale(
+        self, input: torch.Tensor, transposed: bool = False
+    ) -> torch.Tensor: ...
+    def inputs(self) -> list[torch.Tensor]: ...
+
+class MatcherSiluAndMul(MatcherCustomOp):
+    def __init__(self, enabled: bool | None = None) -> None: ...
+    def inputs(self) -> list[torch.Tensor]: ...
+    def forward_custom(self, x: torch.Tensor) -> torch.Tensor: ...
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor: ...

.cuda_typings/vllm/compilation/passes/fusion/qk_norm_rope_fusion.pyi

@@ -0,0 +1,72 @@
+import torch
+from ..inductor_pass import enable_fake_mode as enable_fake_mode
+from ..vllm_inductor_pass import (
+    VllmInductorPass as VllmInductorPass,
+    VllmPatternMatcherPass as VllmPatternMatcherPass,
+)
+from .matcher_utils import (
+    MatcherRMSNorm as MatcherRMSNorm,
+    MatcherRotaryEmbedding as MatcherRotaryEmbedding,
+)
+from .rms_quant_fusion import (
+    empty_bf16 as empty_bf16,
+    empty_fp32 as empty_fp32,
+    empty_i64 as empty_i64,
+)
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable
+from torch import fx as fx
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from typing import ParamSpec
+from vllm.config import (
+    VllmConfig as VllmConfig,
+    get_layers_from_vllm_config as get_layers_from_vllm_config,
+)
+from vllm.logger import init_logger as init_logger
+from vllm.model_executor.layers.attention import Attention as Attention
+from vllm.model_executor.layers.rotary_embedding import (
+    RotaryEmbedding as RotaryEmbedding,
+)
+
+logger: Incomplete
+FUSED_QK_ROPE_OP: Incomplete
+P = ParamSpec("P")
+
+class QkNormRopePattern:
+    num_heads: Incomplete
+    num_kv_heads: Incomplete
+    head_dim: Incomplete
+    q_size: Incomplete
+    kv_size: Incomplete
+    eps: Incomplete
+    rmsnorm_matcher: Incomplete
+    is_neox: Incomplete
+    rope_flashinfer: Incomplete
+    rope_matcher: Incomplete
+    def __init__(
+        self,
+        head_dim: int,
+        num_heads: int,
+        num_kv_heads: int,
+        eps: float,
+        is_neox: bool,
+        rope_flashinfer: bool = False,
+    ) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    @staticmethod
+    def wrap_trace_fn(
+        trace_fn: Callable[P, fx.GraphModule],
+        *process_fx_fns: Callable[[fx.GraphModule], None],
+    ) -> Callable[P, fx.GraphModule]: ...
+    @staticmethod
+    def fx_view_to_reshape(gm: torch.fx.GraphModule) -> None: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class QKNormRoPEFusionPass(VllmPatternMatcherPass):
+    patterns: PatternMatcherPass
+    @enable_fake_mode
+    def __init__(self, config: VllmConfig) -> None: ...
+    matched_count: Incomplete
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None: ...
+    def uuid(self) -> str: ...

.cuda_typings/vllm/compilation/passes/fusion/rms_quant_fusion.pyi

@@ -0,0 +1,143 @@
+import torch
+from ..inductor_pass import enable_fake_mode as enable_fake_mode
+from ..vllm_inductor_pass import (
+    VllmInductorPass as VllmInductorPass,
+    VllmPatternMatcherPass as VllmPatternMatcherPass,
+)
+from .matcher_utils import (
+    MatcherFusedAddRMSNorm as MatcherFusedAddRMSNorm,
+    MatcherQuantFP8 as MatcherQuantFP8,
+    MatcherRMSNorm as MatcherRMSNorm,
+)
+from _typeshed import Incomplete
+from torch import fx as fx
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from torch._ops import OpOverload as OpOverload
+from typing import Any, NamedTuple
+from vllm.config import (
+    VllmConfig as VllmConfig,
+    get_current_vllm_config as get_current_vllm_config,
+)
+from vllm.logger import init_logger as init_logger
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape as GroupShape,
+    QuantKey as QuantKey,
+    ScaleDesc as ScaleDesc,
+    kFp8Dynamic128Sym as kFp8Dynamic128Sym,
+    kFp8Dynamic64Sym as kFp8Dynamic64Sym,
+    kFp8DynamicTensorSym as kFp8DynamicTensorSym,
+    kFp8DynamicTokenSym as kFp8DynamicTokenSym,
+    kFp8StaticTensorSym as kFp8StaticTensorSym,
+    kNvfp4Dynamic as kNvfp4Dynamic,
+    kStaticTensorScale as kStaticTensorScale,
+)
+from vllm.platforms import current_platform as current_platform
+
+logger: Incomplete
+FP8_DTYPE: Incomplete
+FP4_DTYPE: Incomplete
+
+def empty_bf16(*args: Any, **kwargs: Any) -> torch.Tensor: ...
+def empty_fp32(*args: Any, **kwargs: Any) -> torch.Tensor: ...
+def empty_i32(*args: Any, **kwargs: Any) -> torch.Tensor: ...
+def empty_i64(*args: Any, **kwargs: Any) -> torch.Tensor: ...
+
+RMS_OP: Incomplete
+RMS_ADD_OP: Incomplete
+QUANT_OPS: dict[QuantKey, OpOverload]
+
+class FusedRMSQuantKey(NamedTuple):
+    quant: QuantKey
+    fused_add: bool
+
+FUSED_OPS: dict[FusedRMSQuantKey, OpOverload]
+
+class RMSNormQuantPattern:
+    epsilon: Incomplete
+    quant_dtype: Incomplete
+    model_dtype: Incomplete
+    FUSED_OP: Incomplete
+    rmsnorm_matcher: Incomplete
+    quant_matcher: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        key: FusedRMSQuantKey,
+        has_col_major_scales: bool = False,
+        is_e8m0: bool = False,
+        is_tma_aligned: bool = False,
+    ) -> None: ...
+
+class RMSNormStaticQuantPattern(RMSNormQuantPattern):
+    def __init__(
+        self, epsilon: float, quant_dtype: torch.dtype, symmetric: bool = True
+    ) -> None: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
+    def __init__(
+        self, epsilon: float, quant_dtype: torch.dtype, symmetric: bool = True
+    ) -> None: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
+    group_shape: Incomplete
+    is_e8m0: Incomplete
+    has_col_major_scales: Incomplete
+    is_tma_aligned: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        quant_dtype: torch.dtype,
+        group_shape: GroupShape,
+        symmetric: bool = True,
+        is_e8m0: bool = False,
+        has_col_major_scales: bool = True,
+        is_tma_aligned: bool = True,
+    ) -> None: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class RMSNormGroupQuantPattern(RMSNormQuantPattern):
+    group_shape: Incomplete
+    has_col_major_scales: Incomplete
+    is_tma_aligned: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        quant_dtype: torch.dtype,
+        group_shape: GroupShape,
+        symmetric: bool = True,
+        is_e8m0: bool = False,
+        has_col_major_scales: bool = True,
+        is_tma_aligned: bool = True,
+    ) -> None: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
+    def __init__(
+        self,
+        epsilon: float,
+        quant_dtype: torch.dtype,
+        group_shape: GroupShape = ...,
+        symmetric: bool = True,
+    ) -> None: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
+    def __init__(
+        self,
+        epsilon: float,
+        quant_dtype: torch.dtype,
+        group_shape: GroupShape = ...,
+        symmetric: bool = True,
+    ) -> None: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class RMSNormQuantFusionPass(VllmPatternMatcherPass):
+    patterns: PatternMatcherPass
+    @enable_fake_mode
+    def __init__(self, config: VllmConfig) -> None: ...
+    matched_count: Incomplete
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None: ...
+    def uuid(self) -> str: ...

.cuda_typings/vllm/compilation/passes/fusion/rocm_aiter_fusion.pyi

@@ -0,0 +1,133 @@
+import torch
+from ..inductor_pass import enable_fake_mode as enable_fake_mode
+from ..vllm_inductor_pass import (
+    VllmInductorPass as VllmInductorPass,
+    VllmPatternMatcherPass as VllmPatternMatcherPass,
+)
+from .act_quant_fusion import ActivationQuantPattern as ActivationQuantPattern
+from .matcher_utils import (
+    MatcherFusedAddRMSNorm as MatcherFusedAddRMSNorm,
+    MatcherQuantFP8 as MatcherQuantFP8,
+    MatcherRMSNorm as MatcherRMSNorm,
+    MatcherSiluAndMul as MatcherSiluAndMul,
+)
+from .rms_quant_fusion import FusedRMSQuantKey as FusedRMSQuantKey
+from _typeshed import Incomplete
+from torch import fx as fx
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from vllm._aiter_ops import rocm_aiter_ops as rocm_aiter_ops
+from vllm.config import VllmConfig as VllmConfig
+from vllm.logger import init_logger as init_logger
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape as GroupShape,
+    QuantKey as QuantKey,
+    ScaleDesc as ScaleDesc,
+    kFp8Dynamic128Sym as kFp8Dynamic128Sym,
+)
+from vllm.platforms import current_platform as current_platform
+
+logger: Incomplete
+FP8_DTYPE: Incomplete
+
+class AiterRMSNormQuantPattern:
+    epsilon: Incomplete
+    quant_dtype: Incomplete
+    rmsnorm_matcher: Incomplete
+    quant_matcher: Incomplete
+    def __init__(
+        self, epsilon: float, key: FusedRMSQuantKey, match_aiter_quant: bool = True
+    ) -> None: ...
+
+class AiterRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
+    FUSED_OP: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        quant_dtype: torch.dtype,
+        match_aiter_quant: bool = True,
+        group_shape: GroupShape = ...,
+        symmetric: bool = True,
+    ) -> None: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AiterFusedAddRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
+    FUSED_OP: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        quant_dtype: torch.dtype,
+        match_aiter_quant: bool = True,
+        group_shape: GroupShape = ...,
+        symmetric: bool = True,
+    ) -> None: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
+    FUSED_OP: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        quant_dtype: torch.dtype,
+        group_shape: GroupShape,
+        match_aiter_quant: bool = True,
+        symmetric: bool = True,
+    ) -> None: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
+    FUSED_OP: Incomplete
+    def __init__(
+        self,
+        epsilon: float,
+        quant_dtype: torch.dtype,
+        group_shape: GroupShape,
+        match_aiter_quant: bool = True,
+        symmetric: bool = True,
+    ) -> None: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
+    patterns: PatternMatcherPass
+    @enable_fake_mode
+    def __init__(self, config: VllmConfig) -> None: ...
+    matched_count: Incomplete
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None: ...
+    def uuid(self) -> str: ...
+
+class AiterSiluMulFp8GroupQuantPattern(ActivationQuantPattern):
+    FUSED_SILU_MUL_QUANT_OP: Incomplete
+    silu_and_mul_matcher: Incomplete
+    quant_matcher: Incomplete
+    def __init__(self) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):
+    patterns: PatternMatcherPass
+    @enable_fake_mode
+    def __init__(self, config: VllmConfig) -> None: ...
+    matched_count: Incomplete
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: torch.fx.Graph) -> None: ...
+    def uuid(self) -> str: ...
+
+class AddAiterRMSNormPadPattern:
+    AITER_TRITON_ADD_RMSNORM_PAD_OP: Incomplete
+    epsilon: Incomplete
+    hidden_size: Incomplete
+    x_pad_to_multiple: Incomplete
+    rmsnorm_matcher: Incomplete
+    def __init__(
+        self, epsilon: float, hidden_size: int, x_pad_to_multiple: int
+    ) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class RocmAiterTritonAddRMSNormPadFusionPass(VllmPatternMatcherPass):
+    patterns: PatternMatcherPass
+    def __init__(self, config: VllmConfig) -> None: ...
+    matched_count: Incomplete
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: torch.fx.Graph) -> None: ...
+    def uuid(self) -> str: ...

.cuda_typings/vllm/compilation/passes/fusion/rope_kvcache_fusion.pyi

@@ -0,0 +1,72 @@
+import torch
+from ..inductor_pass import enable_fake_mode as enable_fake_mode
+from ..vllm_inductor_pass import (
+    VllmInductorPass as VllmInductorPass,
+    VllmPatternMatcherPass as VllmPatternMatcherPass,
+)
+from .matcher_utils import MatcherRotaryEmbedding as MatcherRotaryEmbedding
+from .rms_quant_fusion import empty_bf16 as empty_bf16, empty_i64 as empty_i64
+from _typeshed import Incomplete
+from torch import fx as fx
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from vllm.config import (
+    VllmConfig as VllmConfig,
+    get_layers_from_vllm_config as get_layers_from_vllm_config,
+)
+from vllm.config.utils import Range as Range
+from vllm.logger import init_logger as init_logger
+from vllm.model_executor.layers.attention.attention import (
+    Attention as Attention,
+    get_attention_context as get_attention_context,
+)
+from vllm.utils.torch_utils import (
+    direct_register_custom_op as direct_register_custom_op,
+)
+
+logger: Incomplete
+
+def fused_rope_and_unified_kv_cache_update_impl(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+    layer_name: str = "",
+) -> torch.Tensor: ...
+def fused_rope_and_unified_kv_cache_update_fake(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+    layer_name: str = "",
+) -> torch.Tensor: ...
+
+class RopeReshapeKVCachePattern:
+    FUSED_OP: Incomplete
+    layer_name: Incomplete
+    num_heads: Incomplete
+    num_kv_heads: Incomplete
+    head_size: Incomplete
+    head_size_v: Incomplete
+    is_neox: Incomplete
+    q_size: Incomplete
+    k_size: Incomplete
+    v_size: Incomplete
+    rope_matcher: Incomplete
+    def __init__(self, layer: Attention, is_neox: bool) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class RopeKVCacheFusionPass(VllmPatternMatcherPass):
+    patterns: PatternMatcherPass
+    max_token_num: Incomplete
+    @enable_fake_mode
+    def __init__(self, config: VllmConfig) -> None: ...
+    matched_count: Incomplete
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None: ...
+    def is_applicable_for_range(self, compile_range: Range) -> bool: ...
+    def uuid(self) -> str: ...

.cuda_typings/vllm/compilation/passes/fusion/sequence_parallelism.pyi

@@ -0,0 +1,95 @@
+import torch
+import torch.fx as fx
+from ..inductor_pass import enable_fake_mode as enable_fake_mode
+from ..utility.noop_elimination import NoOpEliminationPass as NoOpEliminationPass
+from ..vllm_inductor_pass import (
+    VllmInductorPass as VllmInductorPass,
+    VllmPatternMatcherPass as VllmPatternMatcherPass,
+)
+from .matcher_utils import (
+    MatcherFusedAddRMSNorm as MatcherFusedAddRMSNorm,
+    MatcherQuantFP8 as MatcherQuantFP8,
+    MatcherRMSNorm as MatcherRMSNorm,
+)
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable, Sequence
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from vllm.config import VllmConfig as VllmConfig
+from vllm.config.utils import Range as Range
+from vllm.distributed import (
+    get_tp_group as get_tp_group,
+    tensor_model_parallel_all_reduce as tensor_model_parallel_all_reduce,
+)
+from vllm.distributed.parallel_state import (
+    get_tensor_model_parallel_world_size as get_tensor_model_parallel_world_size,
+)
+from vllm.logger import init_logger as init_logger
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kFp8StaticTensorSym as kFp8StaticTensorSym,
+)
+
+logger: Incomplete
+SP_MIN_HIDDEN_SIZE: dict[int, int]
+SP_MIN_PER_GPU_SIZE_MB: dict[int, float]
+
+def get_sequence_parallelism_threshold(
+    hidden_size: int, tp_size: int, element_size: int
+) -> int | None: ...
+def get_first_out_wrapper(
+    fn: Callable[..., Sequence[torch.Tensor]],
+) -> Callable[..., torch.Tensor]: ...
+
+class _SequenceParallelPatternHelper:
+    epsilon: Incomplete
+    dtype: Incomplete
+    device: Incomplete
+    tp_group: Incomplete
+    tp_size: Incomplete
+    def __init__(
+        self, epsilon: float, dtype: torch.dtype, device: str | None
+    ) -> None: ...
+
+class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
+    rmsnorm_matcher: Incomplete
+    def __init__(
+        self, epsilon: float, dtype: torch.dtype, device: str | None
+    ) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
+    rmsnorm_matcher: Incomplete
+    def __init__(
+        self, epsilon: float, dtype: torch.dtype, device: str | None
+    ) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
+    rmsnorm_matcher: Incomplete
+    quant_matcher: Incomplete
+    def __init__(
+        self, epsilon: float, dtype: torch.dtype, device: str | None
+    ) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
+    rmsnorm_matcher: Incomplete
+    quant_matcher: Incomplete
+    def __init__(
+        self, epsilon: float, dtype: torch.dtype, device: str | None
+    ) -> None: ...
+    def get_inputs(self) -> list[torch.Tensor]: ...
+    def register(self, pm_pass: PatternMatcherPass) -> None: ...
+
+class SequenceParallelismPass(VllmPatternMatcherPass):
+    min_token_num: Incomplete
+    noop_cleanup: Incomplete
+    patterns: PatternMatcherPass
+    @enable_fake_mode
+    def __init__(self, config: VllmConfig) -> None: ...
+    def is_applicable_for_range(self, compile_range: Range) -> bool: ...
+    matched_count: Incomplete
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None: ...

.cuda_typings/vllm/compilation/passes/fx_utils.pyi

@@ -0,0 +1,15 @@
+from collections.abc import Iterable, Iterator
+from torch import fx as fx
+from torch._ops import OpOverload, OpOverloadPacket
+from torch.fx.node import Target as Target
+
+def is_func(node: fx.Node, target: Target) -> bool: ...
+def is_auto_func(node: fx.Node, op: OpOverload) -> bool: ...
+def find_auto_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node | None: ...
+def find_auto_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node: ...
+def find_getitem_maybe(node: fx.Node, idx: int) -> fx.Node | None: ...
+def find_getitem(node: fx.Node, idx: int) -> fx.Node: ...
+def find_op_nodes(
+    op: OpOverload | OpOverloadPacket, graph: fx.Graph
+) -> Iterator[fx.Node]: ...
+def get_only_user(node: fx.Node) -> fx.Node: ...

.cuda_typings/vllm/compilation/passes/inductor_pass.pyi

@@ -0,0 +1,37 @@
+import torch
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable, Generator
+from contextlib import contextmanager
+from torch import fx as fx
+from torch._inductor.custom_graph_pass import CustomGraphPass
+from typing import Any, ParamSpec, TypeVar
+from vllm.config.utils import Range as Range
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+class PassContext:
+    compile_range: Range
+    def __init__(self, compile_range: Range) -> None: ...
+
+def get_pass_context() -> PassContext: ...
+@contextmanager
+def pass_context(compile_range: Range) -> Generator[None, None, None]: ...
+
+class InductorPass(CustomGraphPass):
+    def uuid(self) -> str: ...
+    @staticmethod
+    def hash_source(*srcs: str | Any) -> str: ...
+    @staticmethod
+    def hash_dict(dict_: dict[Any, Any]) -> str: ...
+    def is_applicable_for_range(self, compile_range: Range) -> bool: ...
+
+class CallableInductorPass(InductorPass):
+    callable: Incomplete
+    def __init__(
+        self, callable: Callable[[fx.Graph], None], uuid: Any | None = None
+    ) -> None: ...
+    def __call__(self, graph: torch.fx.Graph) -> None: ...
+    def uuid(self) -> Any: ...
+
+def enable_fake_mode(fn: Callable[P, R]) -> Callable[P, R]: ...

.cuda_typings/vllm/compilation/passes/pass_manager.pyi

@@ -0,0 +1,65 @@
+from .fusion.act_quant_fusion import (
+    ActivationQuantFusionPass as ActivationQuantFusionPass,
+)
+from .fusion.allreduce_rms_fusion import AllReduceFusionPass as AllReduceFusionPass
+from .fusion.attn_quant_fusion import AttnFusionPass as AttnFusionPass
+from .fusion.collective_fusion import AsyncTPPass as AsyncTPPass
+from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass as QKNormRoPEFusionPass
+from .fusion.rms_quant_fusion import RMSNormQuantFusionPass as RMSNormQuantFusionPass
+from .fusion.rocm_aiter_fusion import (
+    RocmAiterRMSNormQuantFusionPass as RocmAiterRMSNormQuantFusionPass,
+    RocmAiterSiluMulFp8GroupQuantFusionPass as RocmAiterSiluMulFp8GroupQuantFusionPass,
+    RocmAiterTritonAddRMSNormPadFusionPass as RocmAiterTritonAddRMSNormPadFusionPass,
+)
+from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass as RopeKVCacheFusionPass
+from .fusion.sequence_parallelism import (
+    SequenceParallelismPass as SequenceParallelismPass,
+)
+from .inductor_pass import (
+    CustomGraphPass as CustomGraphPass,
+    InductorPass as InductorPass,
+    get_pass_context as get_pass_context,
+)
+from .utility.fix_functionalization import (
+    FixFunctionalizationPass as FixFunctionalizationPass,
+)
+from .utility.noop_elimination import NoOpEliminationPass as NoOpEliminationPass
+from .utility.scatter_split_replace import (
+    ScatterSplitReplacementPass as ScatterSplitReplacementPass,
+)
+from .utility.split_coalescing import SplitCoalescingPass as SplitCoalescingPass
+from .vllm_inductor_pass import VllmInductorPass as VllmInductorPass
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable
+from torch import fx as fx
+from typing import ParamSpec, TypeVar
+from vllm import envs as envs
+from vllm._aiter_ops import rocm_aiter_ops as rocm_aiter_ops
+from vllm.compilation.passes.utility.post_cleanup import (
+    PostCleanupPass as PostCleanupPass,
+)
+from vllm.config import (
+    VllmConfig as VllmConfig,
+    set_current_vllm_config as set_current_vllm_config,
+)
+from vllm.logger import init_logger as init_logger
+from vllm.platforms import current_platform as current_platform
+from vllm.utils.system_utils import set_env_var as set_env_var
+
+logger: Incomplete
+P = ParamSpec("P")
+R = TypeVar("R")
+
+def with_pattern_match_debug(fn: Callable[P, R]) -> Callable[P, R]: ...
+
+class PostGradPassManager(CustomGraphPass):
+    passes: list[InductorPass]
+    def __init__(self) -> None: ...
+    @with_pattern_match_debug
+    def __call__(self, graph: fx.Graph) -> None: ...
+    pass_config: Incomplete
+    post_cleanup: Incomplete
+    fix_functionalization: Incomplete
+    def configure(self, config: VllmConfig) -> None: ...
+    def add(self, pass_: InductorPass) -> None: ...
+    def uuid(self) -> str: ...

.cuda_typings/vllm/compilation/passes/utility/fix_functionalization.pyi

@@ -0,0 +1,30 @@
+import torch
+from ..fx_utils import is_func as is_func
+from ..vllm_inductor_pass import VllmInductorPass as VllmInductorPass
+from _typeshed import Incomplete
+from vllm.logger import init_logger as init_logger
+from vllm.platforms import current_platform as current_platform
+
+logger: Incomplete
+
+class FixFunctionalizationPass(VllmInductorPass):
+    nodes_to_remove: list[torch.fx.Node]
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: torch.fx.Graph) -> None: ...
+    def defunctionalize(
+        self,
+        graph: torch.fx.Graph,
+        node: torch.fx.Node,
+        mutated_args: dict[int, torch.fx.Node | str],
+        args: tuple[torch.fx.Node | str, ...] | None = None,
+    ) -> None: ...
+    def replace_users_with_mutated_args(
+        self, node: torch.fx.Node, mutated_args: dict[int, torch.fx.Node | str]
+    ) -> None: ...
+    def getitem_users(self, node: torch.fx.Node) -> dict[int, torch.fx.Node]: ...
+    def insert_defunctionalized(
+        self,
+        graph: torch.fx.Graph,
+        node: torch.fx.Node,
+        args: tuple[torch.fx.Node | str, ...] | None = None,
+    ) -> None: ...

.cuda_typings/vllm/compilation/passes/utility/noop_elimination.pyi

@@ -0,0 +1,17 @@
+import torch.fx
+from ..fx_utils import is_func as is_func
+from ..vllm_inductor_pass import VllmInductorPass as VllmInductorPass
+from _typeshed import Incomplete
+from collections.abc import Iterable
+from torch import SymInt as SymInt
+from vllm.logger import init_logger as init_logger
+
+logger: Incomplete
+
+class NoOpEliminationPass(VllmInductorPass):
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: torch.fx.Graph) -> None: ...
+    def dims_equivalent(self, dim: int | SymInt, i_dim: int | SymInt) -> bool: ...
+    def all_dims_equivalent(
+        self, dims: Iterable[int | SymInt], i_dims: Iterable[int | SymInt]
+    ) -> bool: ...

.cuda_typings/vllm/compilation/passes/utility/post_cleanup.pyi

@@ -0,0 +1,6 @@
+from ..vllm_inductor_pass import VllmInductorPass as VllmInductorPass
+from torch import fx as fx
+
+class PostCleanupPass(VllmInductorPass):
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None: ...

.cuda_typings/vllm/compilation/passes/utility/scatter_split_replace.pyi

@@ -0,0 +1,11 @@
+from ..fx_utils import is_func as is_func
+from ..vllm_inductor_pass import VllmInductorPass as VllmInductorPass
+from _typeshed import Incomplete
+from torch import fx as fx
+from vllm.logger import init_logger as init_logger
+
+logger: Incomplete
+
+class ScatterSplitReplacementPass(VllmInductorPass):
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None: ...

.cuda_typings/vllm/compilation/passes/utility/split_coalescing.pyi

@@ -0,0 +1,11 @@
+from ..fx_utils import is_func as is_func
+from ..vllm_inductor_pass import VllmInductorPass as VllmInductorPass
+from _typeshed import Incomplete
+from torch import fx as fx
+from vllm.logger import init_logger as init_logger
+
+logger: Incomplete
+
+class SplitCoalescingPass(VllmInductorPass):
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None: ...

.cuda_typings/vllm/compilation/passes/vllm_inductor_pass.pyi

@@ -0,0 +1,43 @@
+import torch
+from .inductor_pass import InductorPass as InductorPass
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable
+from dataclasses import dataclass
+from torch._inductor.pattern_matcher import PatternMatcherPass as PatternMatcherPass
+from typing import ClassVar
+from vllm.config import VllmConfig as VllmConfig
+from vllm.logger import init_logger as init_logger
+
+logger: Incomplete
+
+@dataclass
+class InductorCompilationConfig:
+    splitting_ops: list[str] | None = ...
+    use_inductor_graph_partition: bool = ...
+
+class VllmInductorPass(InductorPass):
+    dump_prefix: ClassVar[int | None]
+    compilation_config: Incomplete
+    pass_config: Incomplete
+    model_dtype: Incomplete
+    device: str | None
+    pass_name: Incomplete
+    def __init__(self, config: VllmConfig) -> None: ...
+    @staticmethod
+    def time_and_log(
+        call_fn: Callable[[VllmInductorPass, torch.fx.Graph], None],
+    ) -> Callable[[VllmInductorPass, torch.fx.Graph], None]: ...
+    def dump_graph(self, graph: torch.fx.Graph, stage: str) -> None: ...
+    def begin(self) -> None: ...
+    def end_and_log(self) -> None: ...
+
+class VllmPatternMatcherPass(VllmInductorPass):
+    matched_count: int
+    def dump_patterns(
+        self, config: VllmConfig, pm_pass: PatternMatcherPass
+    ) -> None: ...
+
+class PrinterInductorPass(VllmInductorPass):
+    name: Incomplete
+    def __init__(self, name: str, config: VllmConfig) -> None: ...
+    def __call__(self, graph: torch.fx.Graph) -> None: ...

.cuda_typings/vllm/compilation/piecewise_backend.pyi

@@ -0,0 +1,57 @@
+import dataclasses
+import torch.fx as fx
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable
+from typing import Any
+from vllm.compilation.backends import VllmBackend as VllmBackend
+from vllm.config import VllmConfig as VllmConfig
+from vllm.config.utils import Range as Range
+from vllm.logger import init_logger as init_logger
+
+logger: Incomplete
+
+def get_fake_args_from_graph(graph: fx.GraphModule) -> list[Any]: ...
+def create_concrete_args(graph: fx.GraphModule, size: int) -> list[Any]: ...
+@dataclasses.dataclass
+class RangeEntry:
+    compile_range: Range
+    compiled: bool = ...
+    runnable: Callable[..., Any] = ...
+
+class PiecewiseBackend:
+    graph: Incomplete
+    vllm_config: Incomplete
+    compilation_config: Incomplete
+    piecewise_compile_index: Incomplete
+    total_piecewise_compiles: Incomplete
+    vllm_backend: Incomplete
+    compiled_runnables: Incomplete
+    submod_name: Incomplete
+    is_first_graph: Incomplete
+    is_last_graph: Incomplete
+    is_full_graph: Incomplete
+    is_encoder_compilation: Incomplete
+    compile_ranges: Incomplete
+    compile_sizes: Incomplete
+    sym_shape_indices: Incomplete
+    returns_tuple: Incomplete
+    range_entries: dict[Range, RangeEntry]
+    def __init__(
+        self,
+        graph: fx.GraphModule | None,
+        vllm_config: VllmConfig,
+        piecewise_compile_index: int,
+        total_piecewise_compiles: int,
+        sym_shape_indices: list[int],
+        vllm_backend: VllmBackend,
+        returns_tuple: bool,
+        compiled_runnables: dict[str, Callable[..., Any]] | None = None,
+        submod_name: str = "",
+    ) -> None: ...
+    def get_compiled_graph_wrapper(
+        self, compiled_graph: Callable[..., Any]
+    ) -> Callable[..., Any]: ...
+    def to_bytes(self) -> dict[str, bytes]: ...
+    def compile_all_ranges(self) -> None: ...
+    def load_all_ranges(self) -> None: ...
+    def __call__(self, *args: Any) -> Any: ...

.cuda_typings/vllm/compilation/wrapper.pyi

@@ -0,0 +1,38 @@
+import abc
+import torch
+from _typeshed import Incomplete
+from abc import abstractmethod
+from collections.abc import Callable as Callable
+from types import CodeType
+from typing import Any, ParamSpec, TypeVar
+from vllm.config import (
+    CUDAGraphMode as CUDAGraphMode,
+    CompilationMode as CompilationMode,
+    get_current_vllm_config as get_current_vllm_config,
+)
+from vllm.config.compilation import DynamicShapesType as DynamicShapesType
+from vllm.logger import init_logger as init_logger
+from vllm.utils.nvtx_pytorch_hooks import (
+    layerwise_nvtx_marker_context as layerwise_nvtx_marker_context,
+)
+
+logger: Incomplete
+R = TypeVar("R")
+P = ParamSpec("P")
+
+class TorchCompileWithNoGuardsWrapper(metaclass=abc.ABCMeta):
+    def check_invariants_and_forward(self, *args: Any, **kwargs: Any) -> Any: ...
+    compiled: bool
+    vllm_config: Incomplete
+    layerwise_nvtx_tracing_enabled: Incomplete
+    first_compile: bool
+    evaluate_guards: Incomplete
+    def __init__(self) -> None: ...
+    def aot_compile(self, *args: Any, **kwargs: Any) -> Any: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+    @abstractmethod
+    def forward(self, *args: Any, **kwargs: Any) -> Any: ...
+    def original_code_object(self) -> CodeType: ...
+    def bytecode_hook(self, old_code: CodeType, new_code: CodeType) -> None: ...
+
+def reset_compile_wrapper(model: torch.nn.Module) -> None: ...

.cuda_typings/vllm/config/__init__.pyi

@@ -0,0 +1,107 @@
+from vllm.config.attention import AttentionConfig as AttentionConfig
+from vllm.config.cache import CacheConfig as CacheConfig
+from vllm.config.compilation import (
+    CUDAGraphMode as CUDAGraphMode,
+    CompilationConfig as CompilationConfig,
+    CompilationMode as CompilationMode,
+    PassConfig as PassConfig,
+)
+from vllm.config.device import DeviceConfig as DeviceConfig
+from vllm.config.ec_transfer import ECTransferConfig as ECTransferConfig
+from vllm.config.kernel import KernelConfig as KernelConfig
+from vllm.config.kv_events import KVEventsConfig as KVEventsConfig
+from vllm.config.kv_transfer import KVTransferConfig as KVTransferConfig
+from vllm.config.load import LoadConfig as LoadConfig
+from vllm.config.lora import LoRAConfig as LoRAConfig
+from vllm.config.model import (
+    ModelConfig as ModelConfig,
+    iter_architecture_defaults as iter_architecture_defaults,
+    str_dtype_to_torch_dtype as str_dtype_to_torch_dtype,
+    try_match_architecture_defaults as try_match_architecture_defaults,
+)
+from vllm.config.multimodal import MultiModalConfig as MultiModalConfig
+from vllm.config.observability import ObservabilityConfig as ObservabilityConfig
+from vllm.config.offload import (
+    OffloadBackend as OffloadBackend,
+    OffloadConfig as OffloadConfig,
+    PrefetchOffloadConfig as PrefetchOffloadConfig,
+    UVAOffloadConfig as UVAOffloadConfig,
+)
+from vllm.config.parallel import (
+    EPLBConfig as EPLBConfig,
+    ParallelConfig as ParallelConfig,
+)
+from vllm.config.pooler import PoolerConfig as PoolerConfig
+from vllm.config.profiler import ProfilerConfig as ProfilerConfig
+from vllm.config.scheduler import SchedulerConfig as SchedulerConfig
+from vllm.config.speculative import SpeculativeConfig as SpeculativeConfig
+from vllm.config.speech_to_text import SpeechToTextConfig as SpeechToTextConfig
+from vllm.config.structured_outputs import (
+    StructuredOutputsConfig as StructuredOutputsConfig,
+)
+from vllm.config.utils import (
+    ConfigType as ConfigType,
+    SupportsMetricsInfo as SupportsMetricsInfo,
+    config as config,
+    get_attr_docs as get_attr_docs,
+    is_init_field as is_init_field,
+    replace as replace,
+    update_config as update_config,
+)
+from vllm.config.vllm import (
+    VllmConfig as VllmConfig,
+    get_cached_compilation_config as get_cached_compilation_config,
+    get_current_vllm_config as get_current_vllm_config,
+    get_current_vllm_config_or_none as get_current_vllm_config_or_none,
+    get_layers_from_vllm_config as get_layers_from_vllm_config,
+    set_current_vllm_config as set_current_vllm_config,
+)
+from vllm.config.weight_transfer import WeightTransferConfig as WeightTransferConfig
+
+__all__ = [
+    "AttentionConfig",
+    "CacheConfig",
+    "CompilationConfig",
+    "CompilationMode",
+    "CUDAGraphMode",
+    "PassConfig",
+    "DeviceConfig",
+    "ECTransferConfig",
+    "KernelConfig",
+    "KVEventsConfig",
+    "KVTransferConfig",
+    "LoadConfig",
+    "LoRAConfig",
+    "ModelConfig",
+    "iter_architecture_defaults",
+    "str_dtype_to_torch_dtype",
+    "try_match_architecture_defaults",
+    "MultiModalConfig",
+    "ObservabilityConfig",
+    "OffloadBackend",
+    "OffloadConfig",
+    "PrefetchOffloadConfig",
+    "UVAOffloadConfig",
+    "EPLBConfig",
+    "ParallelConfig",
+    "PoolerConfig",
+    "SchedulerConfig",
+    "SpeculativeConfig",
+    "SpeechToTextConfig",
+    "StructuredOutputsConfig",
+    "ProfilerConfig",
+    "ConfigType",
+    "SupportsMetricsInfo",
+    "config",
+    "get_attr_docs",
+    "is_init_field",
+    "replace",
+    "update_config",
+    "VllmConfig",
+    "get_cached_compilation_config",
+    "get_current_vllm_config",
+    "get_current_vllm_config_or_none",
+    "set_current_vllm_config",
+    "get_layers_from_vllm_config",
+    "WeightTransferConfig",
+]

.cuda_typings/vllm/config/attention.pyi

@@ -0,0 +1,21 @@
+from typing import Any, Literal
+from vllm.config.utils import config as config
+from vllm.v1.attention.backends.registry import (
+    AttentionBackendEnum as AttentionBackendEnum,
+)
+
+@config
+class AttentionConfig:
+    backend: AttentionBackendEnum | None = ...
+    flash_attn_version: Literal[2, 3, 4] | None = ...
+    use_prefill_decode_attention: bool = ...
+    flash_attn_max_num_splits_for_cuda_graph: int = ...
+    use_cudnn_prefill: bool = ...
+    use_trtllm_ragged_deepseek_prefill: bool = ...
+    use_trtllm_attention: bool | None = ...
+    disable_flashinfer_prefill: bool = ...
+    disable_flashinfer_q_quantization: bool = ...
+    use_prefill_query_quantization: bool = ...
+    def compute_hash(self) -> str: ...
+    @classmethod
+    def validate_backend_before(cls, value: Any) -> Any: ...

.cuda_typings/vllm/config/cache.pyi

@@ -0,0 +1,40 @@
+from _typeshed import Incomplete
+from pydantic import SkipValidation as SkipValidation
+from typing import ClassVar
+from vllm.config.utils import config as config
+from vllm.logger import init_logger as init_logger
+
+logger: Incomplete
+CacheDType: Incomplete
+MambaDType: Incomplete
+MambaCacheMode: Incomplete
+PrefixCachingHashAlgo: Incomplete
+KVOffloadingBackend: Incomplete
+
+@config
+class CacheConfig:
+    DEFAULT_BLOCK_SIZE: ClassVar[int] = ...
+    block_size: SkipValidation[int] = ...
+    user_specified_block_size: bool = ...
+    gpu_memory_utilization: float = ...
+    cache_dtype: CacheDType = ...
+    is_attention_free: bool = ...
+    num_gpu_blocks_override: int | None = ...
+    sliding_window: int | None = ...
+    enable_prefix_caching: bool = ...
+    prefix_caching_hash_algo: PrefixCachingHashAlgo = ...
+    calculate_kv_scales: bool = ...
+    cpu_kvcache_space_bytes: int | None = ...
+    mamba_page_size_padded: int | None = ...
+    mamba_block_size: int | None = ...
+    mamba_cache_dtype: MambaDType = ...
+    mamba_ssm_cache_dtype: MambaDType = ...
+    mamba_cache_mode: MambaCacheMode = ...
+    num_gpu_blocks: int | None = ...
+    num_cpu_blocks: int | None = ...
+    kv_sharing_fast_prefill: bool = ...
+    kv_cache_memory_bytes: int | None = ...
+    kv_offloading_size: float | None = ...
+    kv_offloading_backend: KVOffloadingBackend = ...
+    def compute_hash(self) -> str: ...
+    def metrics_info(self): ...

.cuda_typings/vllm/config/compilation.pyi

@@ -0,0 +1,142 @@
+import enum
+from _typeshed import Incomplete
+from collections import Counter
+from collections.abc import Callable as Callable
+from pathlib import Path
+from typing import Any, Literal
+from vllm.compilation.passes.inductor_pass import (
+    CallableInductorPass as CallableInductorPass,
+    InductorPass as InductorPass,
+)
+from vllm.config import VllmConfig as VllmConfig
+from vllm.config.utils import (
+    Range as Range,
+    config as config,
+    get_hash_factors as get_hash_factors,
+    hash_factors as hash_factors,
+)
+from vllm.logger import init_logger as init_logger
+from vllm.platforms import current_platform as current_platform
+from vllm.utils.import_utils import resolve_obj_by_qualname as resolve_obj_by_qualname
+from vllm.utils.math_utils import round_up as round_up
+from vllm.utils.torch_utils import is_torch_equal_or_newer as is_torch_equal_or_newer
+
+logger: Incomplete
+
+class CompilationMode(enum.IntEnum):
+    NONE = 0
+    STOCK_TORCH_COMPILE = 1
+    DYNAMO_TRACE_ONCE = 2
+    VLLM_COMPILE = 3
+
+class CUDAGraphMode(enum.Enum):
+    NONE = 0
+    PIECEWISE = 1
+    FULL = 2
+    FULL_DECODE_ONLY = ...
+    FULL_AND_PIECEWISE = ...
+    def decode_mode(self) -> CUDAGraphMode: ...
+    def mixed_mode(self) -> CUDAGraphMode: ...
+    def has_mode(self, mode: CUDAGraphMode) -> bool: ...
+    def requires_piecewise_compilation(self) -> bool: ...
+    def max_cudagraph_mode(self) -> CUDAGraphMode: ...
+    def has_full_cudagraphs(self) -> bool: ...
+    def has_piecewise_cudagraphs(self) -> bool: ...
+    def separate_routine(self) -> bool: ...
+    @classmethod
+    def valid_runtime_modes(cls) -> frozenset["CUDAGraphMode"]: ...
+    def is_valid_runtime_mode(self) -> bool: ...
+    def __bool__(self) -> bool: ...
+
+@config
+class PassConfig:
+    fuse_norm_quant: bool = ...
+    fuse_act_quant: bool = ...
+    fuse_attn_quant: bool = ...
+    eliminate_noops: bool = ...
+    enable_sp: bool = ...
+    fuse_gemm_comms: bool = ...
+    fuse_allreduce_rms: bool = ...
+    enable_qk_norm_rope_fusion: bool = ...
+    fuse_act_padding: bool = ...
+    fuse_rope_kvcache: bool = ...
+    rope_kvcache_fusion_max_token_num: int = ...
+    fi_allreduce_fusion_max_size_mb: float | None = ...
+    sp_min_token_num: int | None = ...
+    def flashinfer_max_size(self, world_size: int) -> int | None: ...
+    @staticmethod
+    def default_fi_allreduce_fusion_max_size_mb() -> dict[int, float]: ...
+    def compute_hash(self) -> str: ...
+    def __post_init__(self) -> None: ...
+    def log_enabled_passes(self) -> None: ...
+
+class DynamicShapesType(str, enum.Enum):
+    BACKED = "backed"
+    UNBACKED = "unbacked"
+    BACKED_SIZE_OBLIVIOUS = "backed_size_oblivious"
+
+@config
+class DynamicShapesConfig:
+    type: DynamicShapesType = ...
+    evaluate_guards: bool = ...
+    assume_32_bit_indexing: bool = ...
+    def compute_hash(self) -> str: ...
+
+@config
+class CompilationConfig:
+    mode: CompilationMode = ...
+    debug_dump_path: Path | None = ...
+    cache_dir: str = ...
+    compile_cache_save_format: Literal["binary", "unpacked"] = ...
+    backend: str = ...
+    custom_ops: list[str] = ...
+    splitting_ops: list[str] | None = ...
+    compile_mm_encoder: bool = ...
+    compile_sizes: list[int | str] | None = ...
+    compile_ranges_endpoints: list[int] | None = ...
+    inductor_compile_config: dict = ...
+    inductor_passes: dict[str, str] = ...
+    cudagraph_mode: CUDAGraphMode = ...
+    cudagraph_num_of_warmups: int = ...
+    cudagraph_capture_sizes: list[int] | None = ...
+    cudagraph_copy_inputs: bool = ...
+    cudagraph_specialize_lora: bool = ...
+    use_inductor_graph_partition: bool = ...
+    pass_config: PassConfig = ...
+    max_cudagraph_capture_size: int = ...
+    dynamic_shapes_config: DynamicShapesConfig = ...
+    local_cache_dir: str = ...
+    fast_moe_cold_start: bool | None = ...
+    enabled_custom_ops: Counter[str] = ...
+    disabled_custom_ops: Counter[str] = ...
+    traced_files: set[str] = ...
+    compilation_time: float = ...
+    static_forward_context: dict[str, Any] = ...
+    static_all_moe_layers: list[str] = ...
+    def compute_hash(self) -> str: ...
+    @classmethod
+    def validate_mode_before(cls, value: Any) -> Any: ...
+    @classmethod
+    def validate_cudagraph_mode_before(cls, value: Any) -> Any: ...
+    @classmethod
+    def validate_pass_config_before(cls, value: Any) -> Any: ...
+    @classmethod
+    def validate_compile_cache_save_format(cls, value: str) -> str: ...
+    def __post_init__(self) -> None: ...
+    def init_backend(self, vllm_config: VllmConfig) -> str | Callable: ...
+    def post_init_cudagraph_sizes(self) -> None: ...
+    def set_splitting_ops_for_v1(
+        self, all2all_backend: str, data_parallel_size: int = 1
+    ): ...
+    def set_splitting_ops_for_attn_fusion(self) -> None: ...
+    def splitting_ops_contain_attention(self) -> bool: ...
+    def is_attention_compiled_piecewise(self) -> bool: ...
+    def custom_op_log_check(self) -> None: ...
+    def is_custom_op_enabled(self, op: str) -> bool: ...
+    def adjust_cudagraph_sizes_for_spec_decode(
+        self, uniform_decode_query_len: int, tensor_parallel_size: int
+    ): ...
+    def adjust_cudagraph_sizes_for_mamba_cache(
+        self, num_mamba_cache_blocks: int
+    ) -> None: ...
+    def get_compile_ranges(self) -> list[Range]: ...

.cuda_typings/vllm/config/device.pyi

@@ -0,0 +1,14 @@
+import torch
+from _typeshed import Incomplete
+from pydantic import ConfigDict, SkipValidation as SkipValidation
+from vllm.config.utils import config as config
+from vllm.utils.hashing import safe_hash as safe_hash
+
+Device: Incomplete
+
+@config(config=ConfigDict(arbitrary_types_allowed=True))
+class DeviceConfig:
+    device: SkipValidation[Device | torch.device | None] = ...
+    device_type: str = ...
+    def compute_hash(self) -> str: ...
+    def __post_init__(self) -> None: ...

.cuda_typings/vllm/config/ec_transfer.pyi

@@ -0,0 +1,30 @@
+from _typeshed import Incomplete
+from typing import Any, Literal
+from vllm.config.utils import config as config
+
+ECProducer: Incomplete
+ECConsumer: Incomplete
+ECRole = Literal[ECProducer, ECConsumer]
+
+@config
+class ECTransferConfig:
+    ec_connector: str | None = ...
+    engine_id: str | None = ...
+    ec_buffer_device: str | None = ...
+    ec_buffer_size: float = ...
+    ec_role: ECRole | None = ...
+    ec_rank: int | None = ...
+    ec_parallel_size: int = ...
+    ec_ip: str = ...
+    ec_port: int = ...
+    ec_connector_extra_config: dict[str, Any] = ...
+    ec_connector_module_path: str | None = ...
+    def compute_hash(self) -> str: ...
+    def __post_init__(self) -> None: ...
+    @property
+    def is_ec_transfer_instance(self) -> bool: ...
+    @property
+    def is_ec_producer(self) -> bool: ...
+    @property
+    def is_ec_consumer(self) -> bool: ...
+    def get_from_extra_config(self, key, default) -> Any: ...

.cuda_typings/vllm/config/kernel.pyi

@@ -0,0 +1,12 @@
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable
+from vllm.config.utils import config as config
+from vllm.utils.hashing import safe_hash as safe_hash
+
+MoEBackend: Incomplete
+
+@config
+class KernelConfig:
+    enable_flashinfer_autotune: bool = ...
+    moe_backend: MoEBackend = ...
+    def compute_hash(self) -> str: ...

.cuda_typings/vllm/config/kv_events.pyi

@@ -0,0 +1,14 @@
+from typing import Literal
+from vllm.config.utils import config as config
+
+@config
+class KVEventsConfig:
+    enable_kv_cache_events: bool = ...
+    publisher: Literal["null", "zmq"] = ...
+    endpoint: str = ...
+    replay_endpoint: str | None = ...
+    buffer_steps: int = ...
+    hwm: int = ...
+    max_queue_size: int = ...
+    topic: str = ...
+    def __post_init__(self) -> None: ...

.cuda_typings/vllm/config/kv_transfer.pyi

@@ -0,0 +1,33 @@
+from _typeshed import Incomplete
+from typing import Any, Literal
+from vllm.config.utils import config as config
+from vllm.utils.hashing import safe_hash as safe_hash
+
+KVProducer: Incomplete
+KVConsumer: Incomplete
+KVRole = Literal[KVProducer, KVConsumer]
+
+@config
+class KVTransferConfig:
+    kv_connector: str | None = ...
+    engine_id: str | None = ...
+    kv_buffer_device: str | None = ...
+    kv_buffer_size: float = ...
+    kv_role: KVRole | None = ...
+    kv_rank: int | None = ...
+    kv_parallel_size: int = ...
+    kv_ip: str = ...
+    kv_port: int = ...
+    kv_connector_extra_config: dict[str, Any] = ...
+    kv_connector_module_path: str | None = ...
+    enable_permute_local_kv: bool = ...
+    kv_load_failure_policy: Literal["recompute", "fail"] = ...
+    def compute_hash(self) -> str: ...
+    def __post_init__(self) -> None: ...
+    @property
+    def is_kv_transfer_instance(self) -> bool: ...
+    @property
+    def is_kv_producer(self) -> bool: ...
+    @property
+    def is_kv_consumer(self) -> bool: ...
+    def get_from_extra_config(self, key, default) -> Any: ...

.cuda_typings/vllm/config/load.pyi

@@ -0,0 +1,22 @@
+from _typeshed import Incomplete
+from vllm.config.utils import config as config
+from vllm.logger import init_logger as init_logger
+from vllm.model_executor.model_loader import LoadFormats as LoadFormats
+from vllm.model_executor.model_loader.tensorizer import (
+    TensorizerConfig as TensorizerConfig,
+)
+from vllm.utils.hashing import safe_hash as safe_hash
+
+logger: Incomplete
+
+@config
+class LoadConfig:
+    load_format: str | LoadFormats = ...
+    download_dir: str | None = ...
+    safetensors_load_strategy: str = ...
+    model_loader_extra_config: dict | TensorizerConfig = ...
+    device: str | None = ...
+    ignore_patterns: list[str] | str = ...
+    use_tqdm_on_load: bool = ...
+    pt_load_map_location: str | dict[str, str] = ...
+    def compute_hash(self) -> str: ...

.cuda_typings/vllm/config/lora.pyi

@@ -0,0 +1,26 @@
+import torch
+from _typeshed import Incomplete
+from pydantic import ConfigDict
+from vllm.config import ModelConfig as ModelConfig
+from vllm.config.cache import CacheConfig as CacheConfig
+from vllm.config.utils import config as config
+from vllm.logger import init_logger as init_logger
+from vllm.utils.hashing import safe_hash as safe_hash
+
+logger: Incomplete
+LoRADType: Incomplete
+MaxLoRARanks: Incomplete
+LoRAExtraVocabSize: Incomplete
+
+@config(config=ConfigDict(arbitrary_types_allowed=True))
+class LoRAConfig:
+    max_lora_rank: MaxLoRARanks = ...
+    max_loras: int = ...
+    fully_sharded_loras: bool = ...
+    max_cpu_loras: int | None = ...
+    lora_dtype: torch.dtype | LoRADType = ...
+    default_mm_loras: dict[str, str] | None = ...
+    enable_tower_connector_lora: bool = ...
+    specialize_active_lora: bool = ...
+    def compute_hash(self) -> str: ...
+    def verify_with_model_config(self, model_config: ModelConfig): ...

.cuda_typings/vllm/config/model.pyi

@@ -0,0 +1,260 @@
+import torch
+from _typeshed import Incomplete
+from collections.abc import Callable, Generator
+from dataclasses import InitVar
+from functools import cached_property as cached_property
+from pydantic import ConfigDict
+from transformers import PretrainedConfig
+from typing import Any
+from vllm.config.load import LoadConfig as LoadConfig
+from vllm.config.model_arch import ModelArchitectureConfig as ModelArchitectureConfig
+from vllm.config.multimodal import (
+    MMCacheType as MMCacheType,
+    MMEncoderTPMode as MMEncoderTPMode,
+    MultiModalConfig as MultiModalConfig,
+)
+from vllm.config.parallel import ParallelConfig as ParallelConfig
+from vllm.config.pooler import PoolerConfig as PoolerConfig
+from vllm.config.scheduler import RunnerType as RunnerType
+from vllm.config.utils import config as config, getattr_iter as getattr_iter
+from vllm.logger import init_logger as init_logger
+from vllm.model_executor.layers.quantization import (
+    QuantizationMethods as QuantizationMethods,
+)
+from vllm.platforms import current_platform as current_platform
+from vllm.tasks import ScoreType as ScoreType
+from vllm.transformers_utils.config import (
+    ConfigFormat as ConfigFormat,
+    get_config as get_config,
+    get_hf_image_processor_config as get_hf_image_processor_config,
+    get_hf_text_config as get_hf_text_config,
+    get_pooling_config as get_pooling_config,
+    get_sentence_transformer_tokenizer_config as get_sentence_transformer_tokenizer_config,
+    is_encoder_decoder as is_encoder_decoder,
+    is_rope_parameters_nested as is_rope_parameters_nested,
+    try_get_dense_modules as try_get_dense_modules,
+    try_get_generation_config as try_get_generation_config,
+    try_get_tokenizer_config as try_get_tokenizer_config,
+    uses_mrope as uses_mrope,
+    uses_xdrope_dim as uses_xdrope_dim,
+)
+from vllm.transformers_utils.gguf_utils import (
+    is_gguf as is_gguf,
+    is_remote_gguf as is_remote_gguf,
+    maybe_patch_hf_config_from_gguf as maybe_patch_hf_config_from_gguf,
+    split_remote_gguf as split_remote_gguf,
+)
+from vllm.transformers_utils.model_arch_config_convertor import (
+    MODEL_ARCH_CONFIG_CONVERTORS as MODEL_ARCH_CONFIG_CONVERTORS,
+    ModelArchConfigConvertorBase as ModelArchConfigConvertorBase,
+)
+from vllm.transformers_utils.runai_utils import (
+    ObjectStorageModel as ObjectStorageModel,
+    is_runai_obj_uri as is_runai_obj_uri,
+)
+from vllm.transformers_utils.utils import maybe_model_redirect as maybe_model_redirect
+from vllm.utils.import_utils import LazyLoader as LazyLoader
+from vllm.v1.attention.backends.registry import (
+    AttentionBackendEnum as AttentionBackendEnum,
+)
+from vllm.v1.sample.logits_processor import LogitsProcessor as LogitsProcessor
+
+logger: Incomplete
+RunnerOption: Incomplete
+ConvertType: Incomplete
+ConvertOption: Incomplete
+TokenizerMode: Incomplete
+ModelDType: Incomplete
+LogprobsMode: Incomplete
+HfOverrides = dict[str, Any] | Callable[[PretrainedConfig], PretrainedConfig]
+ModelImpl: Incomplete
+LayerBlockType: Incomplete
+AttnTypeStr: Incomplete
+
+@config(config=ConfigDict(arbitrary_types_allowed=True))
+class ModelConfig:
+    model: str = ...
+    model_weights: str = ...
+    runner: RunnerOption = ...
+    convert: ConvertOption = ...
+    tokenizer: str = ...
+    tokenizer_mode: TokenizerMode | str = ...
+    trust_remote_code: bool = ...
+    dtype: ModelDType | torch.dtype = ...
+    seed: int = ...
+    hf_config: PretrainedConfig = ...
+    hf_text_config: PretrainedConfig = ...
+    hf_config_path: str | None = ...
+    allowed_local_media_path: str = ...
+    allowed_media_domains: list[str] | None = ...
+    revision: str | None = ...
+    code_revision: str | None = ...
+    tokenizer_revision: str | None = ...
+    max_model_len: int = ...
+    spec_target_max_model_len: int | None = ...
+    quantization: QuantizationMethods | str | None = ...
+    allow_deprecated_quantization: bool = ...
+    enforce_eager: bool = ...
+    enable_return_routed_experts: bool = ...
+    max_logprobs: int = ...
+    logprobs_mode: LogprobsMode = ...
+    disable_sliding_window: bool = ...
+    disable_cascade_attn: bool = ...
+    skip_tokenizer_init: bool = ...
+    enable_prompt_embeds: bool = ...
+    served_model_name: str | list[str] | None = ...
+    config_format: str | ConfigFormat = ...
+    hf_token: bool | str | None = ...
+    hf_overrides: HfOverrides = ...
+    generation_config: str = ...
+    override_generation_config: dict[str, Any] = ...
+    enable_sleep_mode: bool = ...
+    model_impl: str | ModelImpl = ...
+    override_attention_dtype: str | None = ...
+    logits_processors: list[str | type[LogitsProcessor]] | None = ...
+    io_processor_plugin: str | None = ...
+    pooler_config: PoolerConfig | None = ...
+    multimodal_config: MultiModalConfig | None = ...
+    language_model_only: InitVar[bool] = ...
+    limit_mm_per_prompt: InitVar[dict[str, int | dict[str, int]] | None] = ...
+    enable_mm_embeds: InitVar[bool | None] = ...
+    media_io_kwargs: InitVar[dict[str, dict[str, Any]] | None] = ...
+    mm_processor_kwargs: InitVar[dict[str, Any] | None] = ...
+    mm_processor_cache_gb: InitVar[float | None] = ...
+    mm_processor_cache_type: InitVar[MMCacheType | None] = ...
+    mm_shm_cache_max_object_size_mb: InitVar[int | None] = ...
+    mm_encoder_only: InitVar[bool | None] = ...
+    mm_encoder_tp_mode: InitVar[MMEncoderTPMode | None] = ...
+    mm_encoder_attn_backend: InitVar[AttentionBackendEnum | str | None] = ...
+    interleave_mm_strings: InitVar[bool | None] = ...
+    skip_mm_profiling: InitVar[bool | None] = ...
+    video_pruning_rate: InitVar[float | None] = ...
+    def compute_hash(self) -> str: ...
+    attention_chunk_size = ...
+    encoder_config = ...
+    hf_image_processor_config = ...
+    model_arch_config = ...
+    runner_type = ...
+    convert_type = ...
+    original_max_model_len = ...
+    config_updated = ...
+    def __post_init__(
+        self,
+        language_model_only: bool,
+        limit_mm_per_prompt: dict[str, int | dict[str, int]] | None,
+        enable_mm_embeds: bool | None,
+        media_io_kwargs: dict[str, dict[str, Any]] | None,
+        mm_processor_kwargs: dict[str, Any] | None,
+        mm_processor_cache_gb: float | None,
+        mm_processor_cache_type: MMCacheType | None,
+        mm_shm_cache_max_object_size_mb: int | None,
+        mm_encoder_only: bool | None,
+        mm_encoder_tp_mode: MMEncoderTPMode | None,
+        mm_encoder_attn_backend: AttentionBackendEnum | str | None,
+        interleave_mm_strings: bool | None,
+        skip_mm_profiling: bool | None,
+        video_pruning_rate: float | None,
+    ) -> None: ...
+    def get_model_arch_config(self) -> ModelArchitectureConfig: ...
+    @classmethod
+    def validate_quantization_before(cls, value: Any) -> Any: ...
+    def validate_model_config_after(self) -> ModelConfig: ...
+    def using_transformers_backend(self) -> bool: ...
+    @property
+    def registry(self): ...
+    @property
+    def architectures(self) -> list[str]: ...
+    @property
+    def architecture(self) -> str: ...
+    def maybe_pull_model_tokenizer_for_runai(
+        self, model: str, tokenizer: str
+    ) -> None: ...
+    def verify_dual_chunk_attention_config(self, load_config: LoadConfig) -> None: ...
+    def verify_with_parallel_config(self, parallel_config: ParallelConfig) -> None: ...
+    def get_sliding_window(self) -> int | None: ...
+    def get_vocab_size(self) -> int: ...
+    def get_hidden_size(self) -> int: ...
+    def get_inputs_embeds_size(self) -> int: ...
+    @property
+    def is_deepseek_mla(self) -> bool: ...
+    @cached_property
+    def is_mm_prefix_lm(self) -> bool: ...
+    def get_head_size(self) -> int: ...
+    def get_total_num_kv_heads(self) -> int: ...
+    def get_num_kv_heads(self, parallel_config: ParallelConfig) -> int: ...
+    def get_num_attention_heads(self, parallel_config: ParallelConfig) -> int: ...
+    def get_num_experts(self) -> int: ...
+    def get_total_num_hidden_layers(self) -> int: ...
+    def get_layers_start_end_indices(
+        self, parallel_config: ParallelConfig
+    ) -> tuple[int, int]: ...
+    def get_num_layers(self, parallel_config: ParallelConfig) -> int: ...
+    def get_num_layers_by_block_type(
+        self, parallel_config: ParallelConfig, block_type: LayerBlockType = "attention"
+    ) -> int: ...
+    def get_mamba_chunk_size(self) -> int | None: ...
+    def get_multimodal_config(self) -> MultiModalConfig: ...
+    def try_get_generation_config(self) -> dict[str, Any]: ...
+    def get_diff_sampling_param(self) -> dict[str, Any]: ...
+    @cached_property
+    def is_encoder_decoder(self) -> bool: ...
+    @property
+    def uses_alibi(self) -> bool: ...
+    @property
+    def uses_mrope(self) -> bool: ...
+    @property
+    def uses_xdrope_dim(self) -> int: ...
+    @property
+    def is_multimodal_model(self) -> bool: ...
+    @property
+    def is_multimodal_raw_input_only_model(self) -> bool: ...
+    @property
+    def requires_raw_input_tokens(self) -> bool: ...
+    @property
+    def score_type(self) -> ScoreType: ...
+    @property
+    def is_pp_supported(self) -> bool: ...
+    @property
+    def is_attention_free(self) -> bool: ...
+    @property
+    def is_hybrid(self) -> bool: ...
+    @property
+    def has_noops(self) -> bool: ...
+    @property
+    def has_inner_state(self): ...
+    @property
+    def supports_mamba_prefix_caching(self) -> bool: ...
+    @property
+    def use_mla(self) -> bool: ...
+    @property
+    def is_matryoshka(self) -> bool: ...
+    @property
+    def matryoshka_dimensions(self): ...
+    @property
+    def use_sep_token(self) -> bool: ...
+    @property
+    def head_dtype(self) -> torch.dtype: ...
+    @property
+    def embedding_size(self): ...
+    def get_and_verify_max_len(self, max_model_len: int): ...
+    @property
+    def attn_type(self) -> AttnTypeStr: ...
+    @property
+    def is_chunked_prefill_supported(self) -> bool: ...
+    @property
+    def is_prefix_caching_supported(self) -> bool: ...
+    @property
+    def is_moe(self) -> bool: ...
+    @property
+    def is_quantized(self) -> bool: ...
+    def is_nvfp4_quantized(self) -> bool: ...
+
+def get_served_model_name(model: str, served_model_name: str | list[str] | None): ...
+def iter_architecture_defaults() -> Generator[Incomplete, Incomplete]: ...
+def try_match_architecture_defaults(
+    architecture: str,
+    *,
+    runner_type: RunnerType | None = None,
+    convert_type: ConvertType | None = None,
+) -> tuple[str, tuple[RunnerType, ConvertType]] | None: ...
+def str_dtype_to_torch_dtype(type: str): ...

.cuda_typings/vllm/config/model_arch.pyi

@@ -0,0 +1,20 @@
+from _typeshed import Incomplete
+from typing import Any
+from vllm.logger import init_logger as init_logger
+
+logger: Incomplete
+
+class ModelArchitectureConfig:
+    architectures: list[str] | None
+    model_type: str
+    text_model_type: str | None
+    hidden_size: int
+    total_num_hidden_layers: int
+    total_num_attention_heads: int
+    head_size: int
+    vocab_size: int
+    total_num_kv_heads: int
+    num_experts: int
+    quantization_config: dict[str, Any] | None
+    is_deepseek_mla: bool
+    derived_max_model_len_and_key: tuple[float, str | None]

.cuda_typings/vllm/config/multimodal.pyi

@@ -0,0 +1,55 @@
+from _typeshed import Incomplete
+from collections.abc import Mapping
+from typing import Any, TypeAlias, TypedDict
+from vllm.config.utils import config as config
+from vllm.utils.hashing import safe_hash as safe_hash
+from vllm.v1.attention.backends.registry import (
+    AttentionBackendEnum as AttentionBackendEnum,
+)
+
+class BaseDummyOptions:
+    count: int
+
+class VideoDummyOptions(BaseDummyOptions):
+    num_frames: int | None
+    width: int | None
+    height: int | None
+
+class ImageDummyOptions(BaseDummyOptions):
+    width: int | None
+    height: int | None
+
+class AudioDummyOptions(BaseDummyOptions):
+    length: int | None
+
+class MultiModalDummyOptionsBuiltins(TypedDict, total=False):
+    image: ImageDummyOptions
+    video: VideoDummyOptions
+    audio: AudioDummyOptions
+
+MMEncoderTPMode: Incomplete
+MMCacheType: Incomplete
+MMDummyOptions: TypeAlias = dict[str, BaseDummyOptions]
+
+@config
+class MultiModalConfig:
+    language_model_only: bool = ...
+    limit_per_prompt: MMDummyOptions = ...
+    enable_mm_embeds: bool = ...
+    media_io_kwargs: dict[str, dict[str, Any]] = ...
+    mm_processor_kwargs: dict[str, object] | None = ...
+    mm_processor_cache_gb: float = ...
+    mm_processor_cache_type: MMCacheType = ...
+    mm_shm_cache_max_object_size_mb: int = ...
+    mm_encoder_only: bool = ...
+    mm_encoder_tp_mode: MMEncoderTPMode = ...
+    mm_encoder_attn_backend: AttentionBackendEnum | None = ...
+    interleave_mm_strings: bool = ...
+    skip_mm_profiling: bool = ...
+    video_pruning_rate: float | None = ...
+    def compute_hash(self) -> str: ...
+    def get_limit_per_prompt(self, modality: str) -> int: ...
+    def merge_mm_processor_kwargs(
+        self, inference_kwargs: Mapping[str, object]
+    ) -> dict[str, object]: ...
+    def is_multimodal_pruning_enabled(self): ...

.cuda_typings/vllm/config/observability.pyi

@@ -0,0 +1,27 @@
+from _typeshed import Incomplete
+from functools import cached_property as cached_property
+from vllm import version as version
+from vllm.config.utils import config as config
+from vllm.utils.hashing import safe_hash as safe_hash
+
+DetailedTraceModules: Incomplete
+
+@config
+class ObservabilityConfig:
+    show_hidden_metrics_for_version: str | None = ...
+    @cached_property
+    def show_hidden_metrics(self) -> bool: ...
+    otlp_traces_endpoint: str | None = ...
+    collect_detailed_traces: list[DetailedTraceModules] | None = ...
+    kv_cache_metrics: bool = ...
+    kv_cache_metrics_sample: float = ...
+    cudagraph_metrics: bool = ...
+    enable_layerwise_nvtx_tracing: bool = ...
+    enable_mfu_metrics: bool = ...
+    enable_mm_processor_stats: bool = ...
+    enable_logging_iteration_details: bool = ...
+    @cached_property
+    def collect_model_forward_time(self) -> bool: ...
+    @cached_property
+    def collect_model_execute_time(self) -> bool: ...
+    def compute_hash(self) -> str: ...

.cuda_typings/vllm/config/offload.pyi

@@ -0,0 +1,24 @@
+from _typeshed import Incomplete
+from vllm.config.utils import config as config
+
+OffloadBackend: Incomplete
+
+@config
+class UVAOffloadConfig:
+    cpu_offload_gb: float = ...
+    cpu_offload_params: set[str] = ...
+
+@config
+class PrefetchOffloadConfig:
+    offload_group_size: int = ...
+    offload_num_in_group: int = ...
+    offload_prefetch_step: int = ...
+    offload_params: set[str] = ...
+
+@config
+class OffloadConfig:
+    offload_backend: OffloadBackend = ...
+    uva: UVAOffloadConfig = ...
+    prefetch: PrefetchOffloadConfig = ...
+    def validate_offload_config(self) -> OffloadConfig: ...
+    def compute_hash(self) -> str: ...

.cuda_typings/vllm/config/parallel.pyi

@@ -0,0 +1,126 @@
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable
+from ray.runtime_env import RuntimeEnv as RuntimeEnv
+from ray.util.placement_group import PlacementGroup as PlacementGroup
+from torch.distributed import ProcessGroup as ProcessGroup, Store as Store
+from typing import Literal, overload
+from vllm.config.utils import config as config
+from vllm.logger import init_logger as init_logger
+from vllm.model_executor.layers.batch_invariant import (
+    vllm_is_batch_invariant as vllm_is_batch_invariant,
+)
+from vllm.platforms import current_platform as current_platform
+from vllm.utils.network_utils import get_open_ports_list as get_open_ports_list
+from vllm.utils.torch_utils import (
+    cuda_device_count_stateless as cuda_device_count_stateless,
+)
+from vllm.v1.executor import Executor as Executor
+
+logger: Incomplete
+ExpertPlacementStrategy: Incomplete
+DistributedExecutorBackend: Incomplete
+DataParallelBackend: Incomplete
+EPLBPolicyOption: Incomplete
+DCPCommBackend: Incomplete
+All2AllBackend: Incomplete
+
+@config
+class EPLBConfig:
+    window_size: int = ...
+    step_interval: int = ...
+    num_redundant_experts: int = ...
+    log_balancedness: bool = ...
+    log_balancedness_interval: int = ...
+    use_async: bool = ...
+    policy: EPLBPolicyOption = ...
+
+@config
+class ParallelConfig:
+    pipeline_parallel_size: int = ...
+    tensor_parallel_size: int = ...
+    prefill_context_parallel_size: int = ...
+    data_parallel_size: int = ...
+    data_parallel_size_local: int = ...
+    data_parallel_rank: int = ...
+    data_parallel_rank_local: int | None = ...
+    data_parallel_master_ip: str = ...
+    data_parallel_rpc_port: int = ...
+    data_parallel_master_port: int = ...
+    data_parallel_backend: DataParallelBackend = ...
+    data_parallel_external_lb: bool = ...
+    data_parallel_hybrid_lb: bool = ...
+    is_moe_model: bool | None = ...
+    enable_expert_parallel: bool = ...
+    enable_eplb: bool = ...
+    eplb_config: EPLBConfig = ...
+    expert_placement_strategy: ExpertPlacementStrategy = ...
+    all2all_backend: All2AllBackend = ...
+    max_parallel_loading_workers: int | None = ...
+    disable_custom_all_reduce: bool = ...
+    enable_elastic_ep: bool = ...
+    enable_dbo: bool = ...
+    ubatch_size: int = ...
+    dbo_decode_token_threshold: int = ...
+    dbo_prefill_token_threshold: int = ...
+    disable_nccl_for_dp_synchronization: bool | None = ...
+    ray_workers_use_nsight: bool = ...
+    ray_runtime_env: RuntimeEnv | None = ...
+    placement_group: PlacementGroup | None = ...
+    distributed_executor_backend: (
+        str | DistributedExecutorBackend | type[Executor] | None
+    ) = ...
+    worker_cls: str = ...
+    sd_worker_cls: str = ...
+    worker_extension_cls: str = ...
+    master_addr: str = ...
+    master_port: int = ...
+    node_rank: int = ...
+    nnodes: int = ...
+    distributed_timeout_seconds: int | None = ...
+    world_size: int = ...
+    rank: int = ...
+    decode_context_parallel_size: int = ...
+    dcp_kv_cache_interleave_size: int = ...
+    dcp_comm_backend: DCPCommBackend = ...
+    cp_kv_cache_interleave_size: int = ...
+    data_parallel_index: int = ...
+    @property
+    def world_size_across_dp(self) -> int: ...
+    @property
+    def use_ubatching(self) -> bool: ...
+    @property
+    def num_ubatches(self) -> int: ...
+    @property
+    def local_engines_only(self) -> bool: ...
+    def get_next_dp_init_port(self) -> int: ...
+    def allocate_elastic_ep_ports(self) -> None: ...
+    def get_next_stateless_world_group_port(self) -> list[int]: ...
+    def get_next_stateless_dp_group_port(self) -> list[int]: ...
+    def get_next_stateless_ep_group_port(self) -> list[int]: ...
+    def get_next_stateless_eplb_group_port(self) -> list[int]: ...
+    @overload
+    def stateless_init_dp_group(
+        self, return_store: Literal[False] = ...
+    ) -> ProcessGroup: ...
+    @overload
+    def stateless_init_dp_group(
+        self, return_store: Literal[True] = ...
+    ) -> tuple[ProcessGroup, Store]: ...
+    @property
+    def use_sequence_parallel_moe(self) -> bool: ...
+    @property
+    def node_rank_within_dp(self) -> int: ...
+    @property
+    def nnodes_within_dp(self) -> int: ...
+    @property
+    def local_world_size(self) -> int: ...
+    @staticmethod
+    def has_unfinished_dp(dp_group: ProcessGroup, has_unfinished: bool) -> bool: ...
+    @staticmethod
+    def sync_kv_cache_memory_size(
+        dp_group: ProcessGroup, kv_cache_memory: int
+    ) -> int: ...
+    def compute_hash(self): ...
+    def __post_init__(self) -> None: ...
+    @property
+    def use_ray(self) -> bool: ...

.cuda_typings/vllm/config/pooler.pyi

@@ -0,0 +1,27 @@
+from _typeshed import Incomplete
+from vllm.config.utils import config as config
+from vllm.logger import init_logger as init_logger
+from vllm.utils.hashing import safe_hash as safe_hash
+
+logger: Incomplete
+SequencePoolingType: Incomplete
+SEQ_POOLING_TYPES: tuple[SequencePoolingType, ...]
+TokenPoolingType: Incomplete
+TOK_POOLING_TYPES: tuple[TokenPoolingType, ...]
+
+@config
+class PoolerConfig:
+    pooling_type: SequencePoolingType | TokenPoolingType | None = ...
+    seq_pooling_type: SequencePoolingType | None = ...
+    tok_pooling_type: TokenPoolingType | None = ...
+    use_activation: bool | None = ...
+    dimensions: int | None = ...
+    enable_chunked_processing: bool = ...
+    max_embed_len: int | None = ...
+    logit_bias: float | None = ...
+    step_tag_id: int | None = ...
+    returned_token_ids: list[int] | None = ...
+    def __post_init__(self) -> None: ...
+    def get_seq_pooling_type(self) -> SequencePoolingType: ...
+    def get_tok_pooling_type(self) -> TokenPoolingType: ...
+    def compute_hash(self) -> str: ...

.cuda_typings/vllm/config/profiler.pyi

@@ -0,0 +1,25 @@
+from _typeshed import Incomplete
+from vllm.config.utils import config as config
+from vllm.logger import init_logger as init_logger
+from vllm.utils.hashing import safe_hash as safe_hash
+
+logger: Incomplete
+ProfilerKind: Incomplete
+
+@config
+class ProfilerConfig:
+    profiler: ProfilerKind | None = ...
+    torch_profiler_dir: str = ...
+    torch_profiler_with_stack: bool = ...
+    torch_profiler_with_flops: bool = ...
+    torch_profiler_use_gzip: bool = ...
+    torch_profiler_dump_cuda_time_total: bool = ...
+    torch_profiler_record_shapes: bool = ...
+    torch_profiler_with_memory: bool = ...
+    ignore_frontend: bool = ...
+    delay_iterations: int = ...
+    max_iterations: int = ...
+    warmup_iterations: int = ...
+    active_iterations: int = ...
+    wait_iterations: int = ...
+    def compute_hash(self) -> str: ...

.cuda_typings/vllm/config/scheduler.pyi

@@ -0,0 +1,44 @@
+from _typeshed import Incomplete
+from collections.abc import Callable as Callable
+from dataclasses import InitVar
+from typing import ClassVar
+from typing_extensions import Self
+from vllm.config.utils import config as config
+from vllm.logger import init_logger as init_logger
+from vllm.utils.hashing import safe_hash as safe_hash
+from vllm.utils.import_utils import resolve_obj_by_qualname as resolve_obj_by_qualname
+from vllm.v1.core.sched.interface import SchedulerInterface as SchedulerInterface
+
+logger: Incomplete
+RunnerType: Incomplete
+SchedulerPolicy: Incomplete
+
+@config
+class SchedulerConfig:
+    max_model_len: InitVar[int]
+    is_encoder_decoder: InitVar[bool]
+    DEFAULT_MAX_NUM_BATCHED_TOKENS: ClassVar[int] = ...
+    DEFAULT_MAX_NUM_SEQS: ClassVar[int] = ...
+    runner_type: RunnerType = ...
+    max_num_batched_tokens: int = ...
+    max_num_scheduled_tokens: int | None = ...
+    max_num_seqs: int = ...
+    max_num_partial_prefills: int = ...
+    max_long_partial_prefills: int = ...
+    long_prefill_token_threshold: int = ...
+    enable_chunked_prefill: bool = ...
+    is_multimodal_model: bool = ...
+    max_num_encoder_input_tokens: int = ...
+    encoder_cache_size: int = ...
+    policy: SchedulerPolicy = ...
+    disable_chunked_mm_input: bool = ...
+    scheduler_cls: str | type[object] | None = ...
+    disable_hybrid_kv_cache_manager: bool | None = ...
+    async_scheduling: bool | None = ...
+    stream_interval: int = ...
+    @staticmethod
+    def default_factory(**kwargs): ...
+    def get_scheduler_cls(self) -> type["SchedulerInterface"]: ...
+    def compute_hash(self) -> str: ...
+    def __post_init__(self, max_model_len: int, is_encoder_decoder: bool) -> None: ...
+    def verify_max_model_len(self, max_model_len: int) -> Self: ...

.cuda_typings/vllm/config/speculative.pyi

@@ -0,0 +1,66 @@
+import vllm.model_executor.layers.quantization as me_quant
+from _typeshed import Incomplete
+from pydantic import SkipValidation as SkipValidation
+from transformers import PretrainedConfig as PretrainedConfig
+from vllm.config import LoadConfig as LoadConfig
+from vllm.config.model import ModelConfig as ModelConfig
+from vllm.config.parallel import ParallelConfig as ParallelConfig
+from vllm.config.utils import config as config
+from vllm.logger import init_logger as init_logger
+from vllm.transformers_utils.config import get_hf_text_config as get_hf_text_config
+from vllm.utils.hashing import safe_hash as safe_hash
+from vllm.utils.import_utils import (
+    LazyLoader as LazyLoader,
+    has_arctic_inference as has_arctic_inference,
+)
+
+logger: Incomplete
+MTPModelTypes: Incomplete
+EagleModelTypes: Incomplete
+NgramGPUTypes: Incomplete
+SpeculativeMethod: Incomplete
+
+@config
+class SpeculativeConfig:
+    enforce_eager: bool | None = ...
+    num_speculative_tokens: int = ...
+    model: str | None = ...
+    method: SpeculativeMethod | None = ...
+    draft_tensor_parallel_size: int | None = ...
+    tensor_parallel_size: int | None = ...
+    quantization: me_quant.QuantizationMethods | None = ...
+    max_model_len: int | None = ...
+    revision: str | None = ...
+    code_revision: str | None = ...
+    disable_padded_drafter_batch: bool = ...
+    use_local_argmax_reduction: bool = ...
+    prompt_lookup_max: int | None = ...
+    prompt_lookup_min: int | None = ...
+    speculative_token_tree: str | None = ...
+    parallel_drafting: bool = ...
+    target_model_config: SkipValidation[ModelConfig] = ...
+    target_parallel_config: SkipValidation[ParallelConfig] = ...
+    draft_model_config: SkipValidation[ModelConfig] = ...
+    draft_parallel_config: SkipValidation[ParallelConfig] = ...
+    suffix_decoding_max_tree_depth: int = ...
+    suffix_decoding_max_cached_requests: int = ...
+    suffix_decoding_max_spec_factor: float = ...
+    suffix_decoding_min_token_prob: float = ...
+    draft_load_config: LoadConfig | None = ...
+    def compute_hash(self) -> str: ...
+    @staticmethod
+    def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: ...
+    def __post_init__(self): ...
+    def update_arch_(self) -> None: ...
+    @staticmethod
+    def create_draft_parallel_config(
+        target_parallel_config: ParallelConfig,
+        speculative_draft_tensor_parallel_size: int,
+    ) -> ParallelConfig: ...
+    def verify_equal_vocab_size_if_draft_model(self) -> None: ...
+    @property
+    def max_num_new_slots_for_drafting(self) -> int: ...
+    def use_eagle(self) -> bool: ...
+    def uses_draft_model(self) -> bool: ...
+    def uses_extract_hidden_states(self) -> bool: ...
+    def use_ngram_gpu(self) -> bool: ...

.cuda_typings/vllm/config/speech_to_text.pyi

@@ -0,0 +1,10 @@
+from vllm.config.utils import config as config
+
+@config
+class SpeechToTextConfig:
+    sample_rate: float = ...
+    max_audio_clip_s: int | None = ...
+    overlap_chunk_second: int = ...
+    min_energy_split_window_size: int | None = ...
+    @property
+    def allow_audio_chunking(self) -> bool: ...

.cuda_typings/vllm/config/structured_outputs.pyi

@@ -0,0 +1,15 @@
+from _typeshed import Incomplete
+from vllm.config.utils import config as config
+from vllm.utils.hashing import safe_hash as safe_hash
+
+StructuredOutputsBackend: Incomplete
+
+@config
+class StructuredOutputsConfig:
+    backend: StructuredOutputsBackend = ...
+    disable_any_whitespace: bool = ...
+    disable_additional_properties: bool = ...
+    reasoning_parser: str = ...
+    reasoning_parser_plugin: str = ...
+    enable_in_reasoning: bool = ...
+    def compute_hash(self) -> str: ...