A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 21:15 UTC
[MyPy] Enable mypy for `vllm/model_executor/layers/` (#40159)
Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>
Martin Hickey committed
Apr 22, 2026 at 04:15 UTC
3951d3eacde3a3addb81e09f7569972f6d6150cb
28 files changed
+243
-146
tools/pre_commit/mypy.py
-1
@@ -29,7 +29,6 @@ SEPARATE_GROUPS = [
29
"tests",
30
# v0 related
31
"vllm/lora",
32
- "vllm/model_executor/layers",
32
]
33
34
# TODO(woosuk): Include the code from Megatron and HuggingFace.
vllm/model_executor/layers/activation.py
+15
-12
@@ -666,16 +666,7 @@ _ACTIVATION_REGISTRY = LazyDict(
666
"gelu": lambda: GELU(),
667
"gelu_fast": lambda: FastGELU(),
668
"gelu_new": lambda: NewGELU(),
669
- "gelu_pytorch_tanh": lambda: (
670
- # TODO:[ROCm] PyTorch native GELU with tanh is unstable with torch.compile
671
- logger.warning_once(
672
- "[ROCm] PyTorch's native GELU with tanh approximation is unstable. "
673
- "Falling back to GELU(approximate='none')."
674
- ),
675
- nn.GELU(approximate="none"),
676
- )[1]
677
- if current_platform.is_rocm()
678
- else nn.GELU(approximate="tanh"),
669
+ "gelu_pytorch_tanh": lambda: _get_gelu_pytorch_tanh(),
670
"relu": lambda: nn.ReLU(),
671
"relu2": lambda: ReLUSquaredActivation(),
672
"silu": lambda: nn.SiLU(),
@@ -687,6 +678,18 @@ _ACTIVATION_REGISTRY = LazyDict(
678
)
679
680
681
+def _get_gelu_pytorch_tanh() -> nn.Module:
682
+ """Get PyTorch GELU with tanh approximation, with ROCm fallback."""
683
+ if current_platform.is_rocm():
684
+ # TODO:[ROCm] PyTorch native GELU with tanh is unstable with torch.compile
685
+ logger.warning_once(
686
+ "[ROCm] PyTorch's native GELU with tanh approximation is unstable. "
687
+ "Falling back to GELU(approximate='none')."
688
+ )
689
+ return nn.GELU(approximate="none")
690
+ return nn.GELU(approximate="tanh")
691
+
692
+
693
def get_act_fn(act_fn_name: str) -> nn.Module:
694
"""Get an activation function by name."""
695
act_fn_name = act_fn_name.lower()
@@ -703,12 +706,12 @@ def get_act_fn(act_fn_name: str) -> nn.Module:
706
return _ACTIVATION_REGISTRY[act_fn_name]
707
708
706
-_ACTIVATION_AND_MUL_REGISTRY = LazyDict(
709
+_ACTIVATION_AND_MUL_REGISTRY: LazyDict[nn.Module] = LazyDict(
710
{
711
"gelu": lambda: GeluAndMul(),
712
"silu": lambda: SiluAndMul(),
713
"geglu": lambda: GeluAndMul(),
711
- "swigluoai": lambda *args, **kwargs: SwigluOAIAndMul(*args, **kwargs),
714
+ "swigluoai": lambda: SwigluOAIAndMul(),
715
}
716
)
717
vllm/model_executor/layers/attention/attention.py
+15
-6
@@ -33,6 +33,7 @@ from vllm.utils.torch_utils import (
33
)
34
from vllm.v1.attention.backend import (
35
AttentionBackend,
36
+ AttentionMetadata,
37
AttentionType,
38
)
39
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -209,6 +210,7 @@ class Attention(nn.Module, AttentionLayerBase):
210
`self.kv_cache`.
211
"""
212
super().__init__()
213
+ sliding_window: int | None
214
if per_layer_sliding_window is not None:
215
# per-layer sliding window
216
sliding_window = per_layer_sliding_window
@@ -335,7 +337,7 @@ class Attention(nn.Module, AttentionLayerBase):
337
cache_config.enable_prefix_caching = False
338
339
impl_cls = self.attn_backend.get_impl_cls()
338
- self.impl = impl_cls(
340
+ self.impl = impl_cls( # type: ignore[assignment] # impl_cls always returns an AttentionImpl subclass
341
num_heads,
342
head_size,
343
scale,
@@ -576,7 +578,7 @@ class Attention(nn.Module, AttentionLayerBase):
578
def get_attn_backend(self) -> type[AttentionBackend]:
579
return self.attn_backend
580
579
- def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
581
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
582
# Block size may get updated after model loading, refresh it
583
block_size = vllm_config.cache_config.block_size
584
# Should not be called for enc-dec or encoder-only attention.
@@ -680,9 +682,16 @@ def get_attention_context(
682
extracted from the forward context.
683
"""
684
forward_context: ForwardContext = get_forward_context()
683
- attn_metadata = forward_context.attn_metadata
684
- if isinstance(attn_metadata, dict):
685
- attn_metadata = attn_metadata[layer_name]
685
+ attn_metadata_raw = forward_context.attn_metadata
686
+ attn_metadata: AttentionMetadata
687
+ if isinstance(attn_metadata_raw, dict):
688
+ attn_metadata = attn_metadata_raw[layer_name]
689
+ elif isinstance(attn_metadata_raw, list):
690
+ # list[dict[str, AttentionMetadata]]: used in speculative decoding
691
+ # where [0] is the base-model (non-speculative) metadata dict.
692
+ attn_metadata = attn_metadata_raw[0][layer_name]
693
+ else:
694
+ attn_metadata = attn_metadata_raw
695
attn_layer: Attention | MLAAttention = forward_context.no_compile_layers[layer_name]
696
kv_cache = attn_layer.kv_cache
697
slot_mapping = forward_context.slot_mapping
@@ -708,7 +717,7 @@ def unified_kv_cache_update(
717
assert hasattr(attn_layer.impl, "do_kv_cache_update"), (
718
f"{attn_layer.impl.__class__.__name__} does not support kv cache update"
719
)
711
- attn_layer.impl.do_kv_cache_update(
720
+ attn_layer.impl.do_kv_cache_update( # type: ignore[attr-defined]
721
attn_layer,
722
key,
723
value,
vllm/model_executor/layers/attention/chunked_local_attention.py
+1
-1
@@ -29,7 +29,7 @@ from vllm.v1.kv_cache_interface import (
29
30
@functools.lru_cache
31
def create_chunked_local_attention_backend(
32
- underlying_attn_backend: AttentionBackend,
32
+ underlying_attn_backend: type[AttentionBackend],
33
attention_chunk_size: int,
34
) -> type[AttentionBackend]:
35
prefix = f"ChunkedLocalAttention_{attention_chunk_size}_"
vllm/model_executor/layers/attention/cross_attention.py
+9
-4
@@ -72,7 +72,7 @@ def _get_cross_slot_mapping(
72
73
@functools.lru_cache
74
def create_cross_attention_backend(
75
- underlying_attn_backend: AttentionBackend,
75
+ underlying_attn_backend: type[AttentionBackend],
76
) -> type[AttentionBackend]:
77
prefix = "CrossAttention_"
78
underlying_builder = underlying_attn_backend.get_builder_cls()
@@ -87,6 +87,7 @@ def create_cross_attention_backend(
87
) -> AttentionMetadata:
88
new_metadata = copy(common_attn_metadata)
89
new_metadata.causal = False
90
+ assert new_metadata.encoder_seq_lens_cpu is not None
91
max_encoder_len = int(new_metadata.encoder_seq_lens_cpu.max())
92
new_metadata.max_seq_len = max_encoder_len
93
# Any computed tokens indicated decode step>1 (no chunked prefill)
@@ -118,7 +119,7 @@ def create_cross_attention_backend(
119
self.device,
120
)
121
attn_metadata = super().build(common_prefix_len, new_metadata, fast_build)
121
- attn_metadata.slot_mapping = slot_mapping
122
+ attn_metadata.slot_mapping = slot_mapping # type: ignore[attr-defined]
123
return attn_metadata
124
125
# NOTE(Lucas): we need a custom impl so we can use the slot-mapping computed by
@@ -144,8 +145,12 @@ def create_cross_attention_backend(
145
and key is not None
146
and value is not None
147
):
147
- self.do_kv_cache_update(
148
- layer, key, value, kv_cache, attn_metadata.slot_mapping
148
+ self.do_kv_cache_update( # type: ignore[attr-defined]
149
+ layer,
150
+ key,
151
+ value,
152
+ kv_cache,
153
+ attn_metadata.slot_mapping, # type: ignore[attr-defined]
154
)
155
156
return super().forward(
vllm/model_executor/layers/attention/encoder_only_attention.py
+2
-2
@@ -21,7 +21,7 @@ from vllm.v1.kv_cache_interface import KVCacheSpec
21
22
@functools.lru_cache
23
def create_encoder_only_attention_backend(
24
- underlying_attn_backend: AttentionBackend,
24
+ underlying_attn_backend: type[AttentionBackend],
25
) -> type[AttentionBackend]:
26
prefix = "EncoderOnlyAttention_"
27
underlying_builder = underlying_attn_backend.get_builder_cls()
@@ -93,6 +93,6 @@ class EncoderOnlyAttention(Attention):
93
**kwargs,
94
)
95
96
- def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
96
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
97
# Does not need KV cache
98
return None
vllm/model_executor/layers/attention/mla_attention.py
+17
-10
@@ -389,7 +389,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
389
cache_config.enable_prefix_caching = False
390
391
impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
392
- self.impl = impl_cls(
392
+ self.impl = impl_cls( # type: ignore[assignment] # impl_cls always returns an MLAAttentionImpl subclass
393
num_heads=self.num_heads,
394
head_size=self.head_size,
395
scale=self.scale,
@@ -485,16 +485,23 @@ class MLAAttention(nn.Module, AttentionLayerBase):
485
486
if self.use_direct_call:
487
forward_context: ForwardContext = get_forward_context()
488
- attn_metadata = forward_context.attn_metadata
489
- if isinstance(attn_metadata, dict):
490
- attn_metadata = attn_metadata[self.layer_name]
488
+ attn_metadata_raw = forward_context.attn_metadata
489
+ attn_metadata: MLACommonMetadata
490
+ if isinstance(attn_metadata_raw, dict):
491
+ attn_metadata = attn_metadata_raw[self.layer_name] # type: ignore[assignment]
492
+ elif isinstance(attn_metadata_raw, list):
493
+ # list[dict[str, AttentionMetadata]]: used in speculative decoding
494
+ # where [0] is the base-model (non-speculative) metadata dict.
495
+ attn_metadata = attn_metadata_raw[0][self.layer_name] # type: ignore[assignment]
496
+ else:
497
+ attn_metadata = attn_metadata_raw
498
self_kv_cache = self.kv_cache
499
slot_mapping = forward_context.slot_mapping
500
501
assert isinstance(slot_mapping, dict), (
502
f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
503
)
497
- self.impl.do_kv_cache_update(
504
+ self.impl.do_kv_cache_update( # type: ignore[attr-defined]
505
kv_c_normed,
506
k_pe,
507
self_kv_cache,
@@ -612,7 +619,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
619
num_mha_tokens = q.size(0) - num_mqa_tokens
620
621
if num_mha_tokens > 0:
615
- self.impl.forward_mha(
622
+ self.impl.forward_mha( # type: ignore[attr-defined]
623
q[num_mqa_tokens:],
624
k_c_normed[num_mqa_tokens:],
625
k_pe[num_mqa_tokens:],
@@ -695,7 +702,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
702
# call decode attn
703
if not is_sparse_impl:
704
assert attn_metadata.decode is not None
698
- attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)
705
+ attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # type: ignore[attr-defined]
706
707
# correct dcp attn_out with lse.
708
if self.impl.dcp_world_size > 1:
@@ -1053,9 +1060,9 @@ except ImportError:
1060
"AITER_MLA backends use aiter kernels instead."
1061
)
1062
elif current_platform.is_xpu():
1056
- from vllm._xpu_ops import xpu_ops as ops
1063
+ from vllm._xpu_ops import xpu_ops
1064
1058
- flash_attn_varlen_func = ops.flash_attn_varlen_func # type: ignore[no-redef]
1065
+ flash_attn_varlen_func = xpu_ops.flash_attn_varlen_func # type: ignore[no-redef,attr-defined,assignment]
1066
1067
1068
def dynamic_per_batched_tensor_quant(
@@ -1988,7 +1995,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1995
assert isinstance(attn_metadata.prefill, FlashInferPrefillMetadata)
1996
self._build_fi_prefill_wrappers(attn_metadata.prefill)
1997
1991
- return attn_metadata
1998
+ return attn_metadata # type: ignore[return-value]
1999
2000
2001
def reorg_kvcache(
vllm/model_executor/layers/fused_moe/all2all_utils.py
+7
-4
@@ -117,17 +117,20 @@ def maybe_make_prepare_finalize(
117
"Detected DP deployment with no --enable-expert-parallel. "
118
"Falling back to AllGather+ReduceScatter dispatch/combine."
119
)
120
+ device_communicator = get_ep_group().device_communicator
121
+ assert device_communicator is not None
122
+ assert device_communicator.all2all_manager is not None
123
return make_moe_prepare_and_finalize_naive_dp_ep(
124
is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel,
122
- num_dispatchers=(
123
- get_ep_group().device_communicator.all2all_manager.world_size
124
- ),
125
+ num_dispatchers=(device_communicator.all2all_manager.world_size),
126
use_monolithic=use_monolithic,
127
)
128
else:
129
return make_moe_prepare_and_finalize_no_dp_ep(use_monolithic)
130
130
- all2all_manager = get_ep_group().device_communicator.all2all_manager
131
+ device_communicator = get_ep_group().device_communicator
132
+ assert device_communicator is not None
133
+ all2all_manager = device_communicator.all2all_manager
134
assert all2all_manager is not None
135
136
prepare_finalize: FusedMoEPrepareAndFinalize | None = None
vllm/model_executor/layers/fused_moe/config.py
+2
-1
@@ -7,6 +7,7 @@ from typing import Union
7
import torch
8
9
from vllm.config import ParallelConfig, SchedulerConfig
10
+from vllm.config.kernel import MoEBackend
11
from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
@@ -1192,7 +1193,7 @@ class FusedMoEConfig:
1193
# Defaults to intermediate_size_per_partition if not specified.
1194
intermediate_size_per_partition_unpadded: int | None = None
1195
1195
- moe_backend: str = "auto"
1196
+ moe_backend: MoEBackend = "auto"
1197
max_num_tokens: int = SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP
1198
has_bias: bool = False
1199
is_act_and_mul: bool = True
vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py
+3
-3
@@ -210,9 +210,9 @@ def persistent_masked_m_silu_mul_quant(
210
DeepGemmQuantScaleFMT.UE8M0,
211
]
212
213
- cuda_arch = current_platform.get_device_capability(
214
- device_id=y.device.index
215
- ).to_int()
213
+ device_capability = current_platform.get_device_capability(device_id=y.device.index)
214
+ assert device_capability is not None
215
+ cuda_arch = device_capability.to_int()
216
217
if current_platform.is_cuda() and cuda_arch >= 80:
218
torch.ops._C.persistent_masked_m_silu_mul_quant(
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+8
-5
@@ -7,6 +7,7 @@ import torch
7
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
from vllm import envs
10
+from vllm.config.kernel import MoEBackend
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.fused_moe import (
13
FusedMoEConfig,
@@ -146,7 +147,7 @@ def backend_to_kernel_cls(
147
raise ValueError(f"Unknown MXFP4 MoE backend: {backend.value}")
148
149
149
-def map_mxfp4_backend(runner_backend: str) -> Mxfp4MoeBackend:
150
+def map_mxfp4_backend(runner_backend: MoEBackend) -> Mxfp4MoeBackend:
151
"""Map user's moe_backend string to Mxfp4MoeBackend."""
152
mapping: dict[str, Mxfp4MoeBackend] = {
153
"flashinfer_trtllm": Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
@@ -201,10 +202,12 @@ def select_gpt_oss_mxfp4_moe_backend(
202
Select the primary MXFP4 MoE backend.
203
Note: Shape-specific fallbacks may still occur at runtime.
204
"""
204
- triton_kernels_supported = has_triton_kernels() and (
205
- 9,
206
- 0,
207
- ) <= current_platform.get_device_capability() < (11, 0)
205
+ device_capability = current_platform.get_device_capability()
206
+ triton_kernels_supported = (
207
+ has_triton_kernels()
208
+ and device_capability is not None
209
+ and (9, 0) <= device_capability < (11, 0)
210
+ )
211
212
# LoRA: separate experts backend path
213
if config.is_lora_enabled:
vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py
+18
-6
@@ -4,6 +4,9 @@ import torch
4
5
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
6
from vllm.distributed import get_ep_group
7
+from vllm.distributed.device_communicators.base_device_communicator import (
8
+ All2AllManagerBase,
9
+)
10
from vllm.forward_context import get_forward_context
11
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
@@ -11,12 +14,16 @@ from vllm.utils.flashinfer import nvfp4_block_scale_interleave
14
15
16
def get_local_sizes():
14
- return get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()
17
+ dp_metadata = get_forward_context().dp_metadata
18
+ assert dp_metadata is not None
19
+ return dp_metadata.get_chunk_sizes_across_dp_rank()
20
21
22
class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
23
"""FlashInfer implementation using the Moe AlltoAll kernel."""
24
25
+ all2all_manager: All2AllManagerBase
26
+
27
def __init__(
28
self,
29
max_num_tokens: int,
@@ -32,8 +39,12 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
39
self.hidden_size = hidden_size
40
self.num_dispatchers_ = num_dispatchers
41
35
- self.all2all_manager = get_ep_group().device_communicator.all2all_manager
36
- self.all2all_manager.initialize(
42
+ device_communicator = get_ep_group().device_communicator
43
+ assert device_communicator is not None
44
+ all2all_manager = device_communicator.all2all_manager
45
+ assert all2all_manager is not None
46
+ self.all2all_manager = all2all_manager
47
+ self.all2all_manager.initialize( # type: ignore[attr-defined]
48
max_num_tokens=self.max_num_tokens,
49
top_k=self.top_k,
50
num_experts=self.num_experts,
@@ -97,7 +108,8 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
108
payloads.append(topk_ids)
109
payloads.append(topk_weights)
110
100
- recv_payloads = self.all2all_manager.moe_alltoall.dispatch(
111
+ assert self.all2all_manager.moe_alltoall is not None # type: ignore[attr-defined]
112
+ recv_payloads = self.all2all_manager.moe_alltoall.dispatch( # type: ignore[attr-defined]
113
token_selected_experts=topk_ids,
114
input_payloads=payloads,
115
runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
@@ -131,7 +143,7 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
143
apply_router_weight_on_input: bool,
144
weight_and_reduce_impl: mk.TopKWeightAndReduce,
145
) -> None:
134
- assert self.all2all_manager.moe_alltoall is not None
146
+ assert self.all2all_manager.moe_alltoall is not None # type: ignore[attr-defined]
147
148
ep_size = self.all2all_manager.world_size
149
hidden_size = fused_expert_output.shape[-1]
@@ -139,7 +151,7 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
151
ep_size, self.runtime_max_tokens_per_rank, hidden_size
152
)
153
142
- combined_output = self.all2all_manager.moe_alltoall.combine(
154
+ combined_output = self.all2all_manager.moe_alltoall.combine( # type: ignore[attr-defined]
155
payload=fused_expert_output,
156
runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
157
)
vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_two_sided.py
+16
-9
@@ -15,19 +15,26 @@ from vllm.utils.flashinfer import nvfp4_block_scale_interleave
15
16
17
def get_local_sizes():
18
- return get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()
18
+ dp_metadata = get_forward_context().dp_metadata
19
+ assert dp_metadata is not None
20
+ return dp_metadata.get_chunk_sizes_across_dp_rank()
21
22
23
class FlashInferNVLinkTwoSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
24
"""Base class for FlashInfer MoE prepare and finalize operations."""
25
26
+ all2all_manager: All2AllManagerBase
27
+
28
def __init__(
29
self,
30
num_dispatchers: int = 1,
31
):
32
super().__init__()
33
self.num_dispatchers_ = num_dispatchers
30
- self.all2all_manager = get_ep_group().device_communicator.all2all_manager
34
+ device_communicator = get_ep_group().device_communicator
35
+ assert device_communicator is not None
36
+ assert device_communicator.all2all_manager is not None
37
+ self.all2all_manager = device_communicator.all2all_manager
38
39
@property
40
def activation_format(self) -> mk.FusedMoEActivationFormat:
@@ -129,7 +136,7 @@ def flashinfer_alltoall_dispatch(
136
):
137
from flashinfer.comm.trtllm_alltoall import MnnvlMoe
138
132
- assert all2all_manager.ensure_alltoall_workspace_initialized(), (
139
+ assert all2all_manager.ensure_alltoall_workspace_initialized(), ( # type: ignore[attr-defined]
140
"FlashInfer AllToAll workspace not available"
141
)
142
@@ -144,7 +151,7 @@ def flashinfer_alltoall_dispatch(
151
topk_ids,
152
topk_weights,
153
None,
147
- all2all_manager.prepare_workspace_tensor,
154
+ all2all_manager.prepare_workspace_tensor, # type: ignore[attr-defined]
155
max_num_token,
156
ep_rank,
157
ep_size,
@@ -172,7 +179,7 @@ def flashinfer_alltoall_dispatch(
179
x = MnnvlMoe.mnnvl_moe_alltoallv(
180
x,
181
alltoall_info,
175
- all2all_manager.workspace_tensor,
182
+ all2all_manager.workspace_tensor, # type: ignore[attr-defined]
183
ep_rank,
184
ep_size,
185
)
@@ -180,7 +187,7 @@ def flashinfer_alltoall_dispatch(
187
x_sf = MnnvlMoe.mnnvl_moe_alltoallv(
188
x_sf,
189
alltoall_info,
183
- all2all_manager.workspace_tensor,
190
+ all2all_manager.workspace_tensor, # type: ignore[attr-defined]
191
ep_rank,
192
ep_size,
193
)
@@ -196,7 +203,7 @@ def flashinfer_alltoall_dispatch(
203
x = MnnvlMoe.mnnvl_moe_alltoallv(
204
x,
205
alltoall_info,
199
- all2all_manager.workspace_tensor,
206
+ all2all_manager.workspace_tensor, # type: ignore[attr-defined]
207
ep_rank,
208
ep_size,
209
)
@@ -212,13 +219,13 @@ def flashinfer_alltoall_combine(
219
):
220
from flashinfer.comm.trtllm_alltoall import MnnvlMoe
221
215
- assert all2all_manager.ensure_alltoall_workspace_initialized(), (
222
+ assert all2all_manager.ensure_alltoall_workspace_initialized(), ( # type: ignore[attr-defined]
223
"FlashInfer AllToAll workspace not available"
224
)
225
return MnnvlMoe.mnnvl_moe_alltoallv_combine(
226
output,
227
alltoall_info,
221
- all2all_manager.workspace_tensor,
228
+ all2all_manager.workspace_tensor, # type: ignore[attr-defined]
229
ep_rank=all2all_manager.rank,
230
ep_size=all2all_manager.world_size,
231
top_k=top_k,
vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py
+4
@@ -132,9 +132,11 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular
132
)
133
134
if scales is None:
135
+ assert len(res) == 3
136
a1q, topk_weights, topk_ids = res
137
a1q_scale = None
138
else:
139
+ assert len(res) == 4
140
a1q, topk_weights, topk_ids, scales = res
141
a1q_scale = _unwrap_scale_and_prepare_for_moe(scales, quant_config)
142
@@ -217,9 +219,11 @@ class MoEPrepareAndFinalizeNaiveDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMono
219
)
220
221
if scales is None:
222
+ assert len(res) == 2
223
a1q, router_logits = res
224
a1q_scale = None
225
else:
226
+ assert len(res) == 3
227
a1q, router_logits, scales = res
228
a1q_scale = _unwrap_scale_and_prepare_for_moe(scales, quant_config)
229
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py
+3
-1
@@ -54,11 +54,13 @@ class DefaultMoERunner(MoERunnerBase):
54
# NOTE: this will be removed once all kernels are migrated into the
55
# MoEKernel framework.
56
if self.do_naive_dispatch_combine:
57
- hidden_states, router_logits = get_ep_group().dispatch_router_logits(
57
+ res = get_ep_group().dispatch_router_logits(
58
hidden_states,
59
router_logits,
60
self.moe_config.is_sequence_parallel,
61
)
62
+ assert len(res) == 2
63
+ hidden_states, router_logits = res
64
65
# NOTE: Similar with DP, PCP also needs dispatch and combine. For
66
# simplicity, AgRsAll2All was added separately for PCP here. Maybe
vllm/model_executor/layers/kda.py
+25
-18
@@ -16,7 +16,6 @@ from vllm.logger import init_logger
16
from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
17
from vllm.model_executor.utils import set_weight_attrs
18
from vllm.utils.torch_utils import direct_register_custom_op
19
-from vllm.v1.attention.backend import AttentionMetadata
19
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
20
21
from .fla.ops.kda import (
@@ -123,7 +122,7 @@ class KimiDeltaAttention(nn.Module, MambaBase):
122
self.cache_config = cache_config
123
if model_config is None:
124
raise ValueError("model_config must be provided")
126
- kda_config = model_config.linear_attn_config
125
+ kda_config = model_config.linear_attn_config # type: ignore[attr-defined]
126
self.head_dim = kda_config["head_dim"]
127
self.num_heads = kda_config["num_heads"]
128
self.layer_idx = layer_idx
@@ -297,19 +296,21 @@ class KimiDeltaAttention(nn.Module, MambaBase):
296
core_attn_out: torch.Tensor,
297
) -> None:
298
forward_context = get_forward_context()
300
- attn_metadata: AttentionMetadata = forward_context.attn_metadata
299
+ attn_metadata_raw = forward_context.attn_metadata
300
302
- if attn_metadata is None:
301
+ if attn_metadata_raw is None:
302
# # V1 profile run
303
return
304
306
- assert isinstance(attn_metadata, dict)
307
- attn_metadata = attn_metadata[self.prefix]
308
- assert isinstance(attn_metadata, GDNAttentionMetadata)
309
- has_initial_state = attn_metadata.has_initial_state
310
- non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
311
- non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
312
- num_actual_tokens = attn_metadata.num_actual_tokens
305
+ assert isinstance(attn_metadata_raw, dict)
306
+ attn_metadata_narrowed = attn_metadata_raw[self.prefix]
307
+ assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata)
308
+ has_initial_state = attn_metadata_narrowed.has_initial_state
309
+ non_spec_query_start_loc = attn_metadata_narrowed.non_spec_query_start_loc
310
+ non_spec_state_indices_tensor = (
311
+ attn_metadata_narrowed.non_spec_state_indices_tensor
312
+ ) # noqa: E501
313
+ num_actual_tokens = attn_metadata_narrowed.num_actual_tokens
314
constant_caches = self.kv_cache
315
316
q_proj_states = q_proj_states[:num_actual_tokens]
@@ -335,7 +336,7 @@ class KimiDeltaAttention(nn.Module, MambaBase):
336
v_conv_weights = self.v_conv1d.weight.view(
337
self.v_conv1d.weight.size(0), self.v_conv1d.weight.size(2)
338
)
338
- if attn_metadata.num_prefills > 0:
339
+ if attn_metadata_narrowed.num_prefills > 0:
340
q_proj_states = q_proj_states.transpose(0, 1)
341
k_proj_states = k_proj_states.transpose(0, 1)
342
v_proj_states = v_proj_states.transpose(0, 1)
@@ -348,7 +349,7 @@ class KimiDeltaAttention(nn.Module, MambaBase):
349
has_initial_state=has_initial_state,
350
cache_indices=non_spec_state_indices_tensor,
351
query_start_loc=non_spec_query_start_loc,
351
- metadata=attn_metadata,
352
+ metadata=attn_metadata_narrowed,
353
).transpose(0, 1)
354
k = causal_conv1d_fn(
355
k_proj_states,
@@ -359,7 +360,7 @@ class KimiDeltaAttention(nn.Module, MambaBase):
360
has_initial_state=has_initial_state,
361
cache_indices=non_spec_state_indices_tensor,
362
query_start_loc=non_spec_query_start_loc,
362
- metadata=attn_metadata,
363
+ metadata=attn_metadata_narrowed,
364
).transpose(0, 1)
365
v = causal_conv1d_fn(
366
v_proj_states,
@@ -370,11 +371,12 @@ class KimiDeltaAttention(nn.Module, MambaBase):
371
has_initial_state=has_initial_state,
372
cache_indices=non_spec_state_indices_tensor,
373
query_start_loc=non_spec_query_start_loc,
373
- metadata=attn_metadata,
374
+ metadata=attn_metadata_narrowed,
375
).transpose(0, 1)
376
else:
377
+ assert non_spec_state_indices_tensor is not None
378
decode_conv_indices = non_spec_state_indices_tensor[
377
- : attn_metadata.num_actual_tokens
379
+ : attn_metadata_narrowed.num_actual_tokens
380
]
381
q = causal_conv1d_update(
382
q_proj_states,
@@ -408,7 +410,9 @@ class KimiDeltaAttention(nn.Module, MambaBase):
410
lambda x: rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim), (q, k, v)
411
)
412
411
- if attn_metadata.num_prefills > 0:
413
+ if attn_metadata_narrowed.num_prefills > 0:
414
+ assert non_spec_state_indices_tensor is not None
415
+ assert has_initial_state is not None
416
zero_idx = non_spec_state_indices_tensor[~has_initial_state]
417
recurrent_state[zero_idx] = 0
418
initial_state = recurrent_state[non_spec_state_indices_tensor].contiguous()
@@ -429,6 +433,7 @@ class KimiDeltaAttention(nn.Module, MambaBase):
433
# Init cache
434
recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state
435
else:
436
+ assert non_spec_query_start_loc is not None
437
(
438
core_attn_out_non_spec,
439
last_recurrent_state,
@@ -440,7 +445,9 @@ class KimiDeltaAttention(nn.Module, MambaBase):
445
beta=beta,
446
initial_state=recurrent_state,
447
use_qk_l2norm_in_kernel=True,
443
- cu_seqlens=non_spec_query_start_loc[: attn_metadata.num_decodes + 1],
448
+ cu_seqlens=non_spec_query_start_loc[
449
+ : attn_metadata_narrowed.num_decodes + 1
450
+ ],
451
ssm_state_indices=non_spec_state_indices_tensor,
452
)
453
core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
vllm/model_executor/layers/layernorm.py
+1
-1
@@ -76,7 +76,7 @@ def poly_norm(
76
from vllm import _custom_ops as ops
77
78
out = torch.empty_like(x)
79
- ops.poly_norm(
79
+ ops.poly_norm( # type: ignore[attr-defined]
80
out,
81
x,
82
weight,
vllm/model_executor/layers/mamba/abstract.py
+2
-1
@@ -42,9 +42,10 @@ class MambaBase(AttentionLayerBase):
42
43
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
44
mamba_block_size = vllm_config.cache_config.mamba_block_size
45
+ assert mamba_block_size is not None
46
page_size_padded = vllm_config.cache_config.mamba_page_size_padded
47
return MambaSpec(
47
- shapes=self.get_state_shape(),
48
+ shapes=tuple(self.get_state_shape()),
49
dtypes=self.get_state_dtype(),
50
block_size=mamba_block_size,
51
page_size_padded=page_size_padded,
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+36
-30
@@ -62,7 +62,6 @@ from vllm.utils.torch_utils import (
62
_resolve_layer_name,
63
direct_register_custom_op,
64
)
65
-from vllm.v1.attention.backend import AttentionMetadata
65
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
66
67
logger = init_logger(__name__)
@@ -121,9 +120,9 @@ def fi_chunk_gated_delta_rule(
120
class ChunkGatedDeltaRule(CustomOp):
121
def __init__(self) -> None:
122
super().__init__()
124
- backend_cfg = get_current_vllm_config().additional_config.get(
125
- "gdn_prefill_backend", "auto"
126
- )
123
+ additional_config = get_current_vllm_config().additional_config
124
+ assert isinstance(additional_config, dict)
125
+ backend_cfg = additional_config.get("gdn_prefill_backend", "auto")
126
backend = str(backend_cfg).strip().lower()
127
128
supports_flashinfer = (
@@ -621,18 +620,19 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
620
# Part 2: Core Attention
621
# ============================================================
622
forward_context = get_forward_context()
624
- attn_metadata: AttentionMetadata = forward_context.attn_metadata
623
+ attn_metadata_raw = forward_context.attn_metadata
624
core_attn_out = torch.zeros(
625
(num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
626
dtype=hidden_states.dtype,
627
device=hidden_states.device,
628
)
629
z = torch.empty_like(core_attn_out)
631
- if attn_metadata is not None:
632
- attn_metadata = attn_metadata[self.prefix]
630
+ if attn_metadata_raw is not None:
631
+ assert isinstance(attn_metadata_raw, dict)
632
+ attn_metadata = attn_metadata_raw[self.prefix]
633
634
# TODO: xpu does not support this param yet
635
- spec_sequence_masks = attn_metadata.spec_sequence_masks
635
+ spec_sequence_masks = attn_metadata.spec_sequence_masks # type: ignore[attr-defined]
636
assert spec_sequence_masks is None
637
638
conv_weights = self.conv1d.weight.view(
@@ -658,12 +658,12 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
658
activation=self.activation,
659
A_log=self.A_log,
660
dt_bias=self.dt_bias,
661
- num_prefills=attn_metadata.num_prefills,
662
- num_decodes=attn_metadata.num_decodes,
663
- has_initial_state=attn_metadata.has_initial_state,
664
- non_spec_query_start_loc=attn_metadata.non_spec_query_start_loc,
665
- non_spec_state_indices_tensor=attn_metadata.non_spec_state_indices_tensor,
666
- num_actual_tokens=attn_metadata.num_actual_tokens,
661
+ num_prefills=attn_metadata.num_prefills, # type: ignore[attr-defined]
662
+ num_decodes=attn_metadata.num_decodes, # type: ignore[attr-defined]
663
+ has_initial_state=attn_metadata.has_initial_state, # type: ignore[attr-defined]
664
+ non_spec_query_start_loc=attn_metadata.non_spec_query_start_loc, # type: ignore[attr-defined]
665
+ non_spec_state_indices_tensor=attn_metadata.non_spec_state_indices_tensor, # type: ignore[attr-defined]
666
+ num_actual_tokens=attn_metadata.num_actual_tokens, # type: ignore[attr-defined]
667
tp_size=self.tp_size,
668
reorder_input=not self.gqa_interleaved_layout,
669
)
@@ -792,16 +792,16 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
792
core_attn_out: torch.Tensor,
793
):
794
forward_context = get_forward_context()
795
- attn_metadata: AttentionMetadata = forward_context.attn_metadata
795
+ attn_metadata_raw = forward_context.attn_metadata
796
797
- if attn_metadata is None:
797
+ if attn_metadata_raw is None:
798
# V1 profile run — warm up prefill kernels so that
799
# autotuning completes before KV cache allocation.
800
self._warmup_prefill_kernels(mixed_qkv)
801
return
802
803
- assert isinstance(attn_metadata, dict)
804
- attn_metadata = attn_metadata[self.prefix]
803
+ assert isinstance(attn_metadata_raw, dict)
804
+ attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index]
805
assert isinstance(attn_metadata, GDNAttentionMetadata)
806
807
if (
@@ -860,14 +860,16 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
860
861
# 1.1: Process the multi-query part
862
if spec_sequence_masks is not None:
863
+ # spec_state_indices_tensor is always set when spec_sequence_masks is set
864
+ assert spec_state_indices_tensor is not None
865
mixed_qkv_spec = causal_conv1d_update(
866
mixed_qkv_spec,
867
conv_state,
868
conv_weights,
869
self.conv1d.bias,
870
self.activation,
869
- conv_state_indices=spec_state_indices_tensor[:, 0][
870
- : attn_metadata.num_spec_decodes
871
+ conv_state_indices=spec_state_indices_tensor[:, 0][ # type: ignore[index]
872
+ : attn_metadata.num_spec_decodes # type: ignore[attr-defined]
873
],
874
num_accepted_tokens=num_accepted_tokens,
875
query_start_loc=spec_query_start_loc,
@@ -900,8 +902,8 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
902
conv_weights,
903
self.conv1d.bias,
904
self.activation,
903
- conv_state_indices=non_spec_state_indices_tensor[
904
- : attn_metadata.num_actual_tokens
905
+ conv_state_indices=non_spec_state_indices_tensor[ # type: ignore[index]
906
+ : attn_metadata.num_actual_tokens # type: ignore[attr-defined]
907
],
908
validate_data=True,
909
)
@@ -965,8 +967,9 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
967
v=value_spec,
968
initial_state=ssm_state,
969
inplace_final_state=True,
968
- cu_seqlens=spec_query_start_loc[
969
- : attn_metadata.num_spec_decodes + 1
970
+ cu_seqlens=spec_query_start_loc[ # type: ignore[index]
971
+ : attn_metadata.num_spec_decodes
972
+ + 1 # type: ignore[attr-defined]
973
],
974
ssm_state_indices=spec_state_indices_tensor,
975
num_accepted_tokens=num_accepted_tokens,
@@ -978,8 +981,10 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
981
982
# 2.2: Process the remaining part
983
if attn_metadata.num_prefills > 0:
981
- initial_state = ssm_state[non_spec_state_indices_tensor].contiguous()
982
- initial_state[~has_initial_state, ...] = 0
984
+ assert non_spec_state_indices_tensor is not None
985
+ initial_state = ssm_state[non_spec_state_indices_tensor].contiguous() # type: ignore[index]
986
+ assert has_initial_state is not None
987
+ initial_state[~has_initial_state, ...] = 0 # type: ignore[operator]
988
(
989
core_attn_out_non_spec,
990
last_recurrent_state,
@@ -1012,8 +1017,9 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
1017
v=value_non_spec,
1018
initial_state=ssm_state,
1019
inplace_final_state=True,
1015
- cu_seqlens=non_spec_query_start_loc[
1016
- : attn_metadata.num_decodes + 1
1020
+ cu_seqlens=non_spec_query_start_loc[ # type: ignore[index]
1021
+ : attn_metadata.num_decodes
1022
+ + 1 # type: ignore[attr-defined]
1023
],
1024
ssm_state_indices=non_spec_state_indices_tensor,
1025
use_qk_l2norm_in_kernel=True,
@@ -1073,7 +1079,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
1079
conv_weights,
1080
self.conv1d.bias,
1081
self.activation,
1076
- conv_state_indices=non_spec_state_indices_tensor[:num_actual_tokens],
1082
+ conv_state_indices=non_spec_state_indices_tensor[:num_actual_tokens], # type: ignore[index]
1083
validate_data=False,
1084
)
1085
out_buf = core_attn_out[:num_actual_tokens].unsqueeze(1)
@@ -1086,7 +1092,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
1092
scale=self.head_k_dim**-0.5,
1093
initial_state=ssm_state,
1094
out=out_buf,
1089
- ssm_state_indices=non_spec_state_indices_tensor[:num_actual_tokens],
1095
+ ssm_state_indices=non_spec_state_indices_tensor[:num_actual_tokens], # type: ignore[index]
1096
use_qk_l2norm_in_kernel=True,
1097
)
1098
return
vllm/model_executor/layers/mamba/linear_attn.py
+5
-4
@@ -396,10 +396,11 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
396
self, hidden_states: torch.Tensor, output: torch.Tensor, positions: torch.Tensor
397
) -> None:
398
forward_context = get_forward_context()
399
- attn_metadata: AttentionMetadata = forward_context.attn_metadata
400
- if attn_metadata is not None:
401
- assert isinstance(attn_metadata, dict)
402
- attn_metadata = attn_metadata[self.prefix]
399
+ attn_metadata_raw = forward_context.attn_metadata
400
+ attn_metadata: AttentionMetadata | None = None
401
+ if attn_metadata_raw is not None:
402
+ assert isinstance(attn_metadata_raw, dict)
403
+ attn_metadata = attn_metadata_raw[self.prefix]
404
assert isinstance(attn_metadata, LinearAttentionMetadata)
405
num_actual_tokens = (
406
attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens
vllm/model_executor/layers/mamba/mamba_mixer.py
+9
-4
@@ -40,6 +40,7 @@ from vllm.utils.torch_utils import (
40
_resolve_layer_name,
41
direct_register_custom_op,
42
)
43
+from vllm.v1.attention.backend import AttentionMetadata
44
from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionMetadata
45
46
@@ -258,15 +259,16 @@ class MambaMixer(MambaBase, PluggableLayer):
259
"""
260
261
forward_context: ForwardContext = get_forward_context()
261
- attn_metadata = forward_context.attn_metadata
262
+ attn_metadata_raw = forward_context.attn_metadata
263
264
assert self.cache_config is not None
265
mamba_block_size = self.cache_config.mamba_block_size
266
is_mamba_cache_all = self.cache_config.mamba_cache_mode == "all"
267
267
- if attn_metadata is not None:
268
- assert isinstance(attn_metadata, dict)
269
- attn_metadata = attn_metadata[self.prefix]
268
+ attn_metadata: AttentionMetadata | None = None
269
+ if attn_metadata_raw is not None:
270
+ assert isinstance(attn_metadata_raw, dict)
271
+ attn_metadata = attn_metadata_raw[self.prefix]
272
assert isinstance(attn_metadata, Mamba1AttentionMetadata)
273
query_start_loc_p = attn_metadata.query_start_loc_p
274
state_indices_tensor_p = attn_metadata.state_indices_tensor_p
@@ -391,6 +393,9 @@ class MambaMixer(MambaBase, PluggableLayer):
393
ssm_outputs.append(scan_out_p)
394
395
if has_decode:
396
+ # state_indices_tensor_d is assigned when attn_metadata is not None,
397
+ # and has_decode is only True when attn_metadata is not None
398
+ assert state_indices_tensor_d is not None
399
if is_mamba_cache_all:
400
state_indices_tensor_d_input = state_indices_tensor_d.gather(
401
1, block_idx_last_computed_token_d.unsqueeze(1)
vllm/model_executor/layers/mamba/mamba_mixer2.py
+17
-4
@@ -572,14 +572,16 @@ class MambaMixer2(MambaBase, PluggableLayer):
572
# kernels to operate in continuous batching and in chunked prefill
573
# modes; they are computed at top-level model forward since they
574
# stay the same and reused for all mamba layers in the same iteration
575
- attn_metadata: AttentionMetadata = forward_context.attn_metadata
575
+ attn_metadata_raw = forward_context.attn_metadata
576
577
assert self.cache_config is not None
578
mamba_block_size = self.cache_config.mamba_block_size
579
is_mamba_cache_all = self.cache_config.mamba_cache_mode == "all"
580
- if attn_metadata is not None:
581
- assert isinstance(attn_metadata, dict)
582
- attn_metadata = attn_metadata[self.prefix]
580
+
581
+ attn_metadata: AttentionMetadata | None = None
582
+ if attn_metadata_raw is not None:
583
+ assert isinstance(attn_metadata_raw, dict)
584
+ attn_metadata = attn_metadata_raw[self.prefix]
585
assert isinstance(attn_metadata, Mamba2AttentionMetadata)
586
# conv_state must be (..., dim, width-1) for the conv kernels.
587
# DS layout stores it that way directly; SD layout needs a
@@ -708,6 +710,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
710
# 3. State Space Model sequence transformation
711
initial_states = None
712
if has_initial_states_p is not None and prep_initial_states:
713
+ assert state_indices_tensor_p is not None
714
kernel_ssm_indices = state_indices_tensor_p
715
if is_mamba_cache_all:
716
kernel_ssm_indices = state_indices_tensor_p.gather(
@@ -746,6 +749,13 @@ class MambaMixer2(MambaBase, PluggableLayer):
749
)
750
751
if is_mamba_cache_all:
752
+ assert mamba_block_size is not None
753
+ assert state_indices_tensor_p is not None
754
+ assert block_idx_first_scheduled_token_p is not None
755
+ assert block_idx_last_scheduled_token_p is not None
756
+ assert last_chunk_indices_p is not None
757
+ assert num_computed_tokens_p is not None
758
+
759
# The chunk_stride is the number of chunks per mamba block
760
# e.g., if mamba_block_size = 512 and chunk_size = 256,
761
# then chunk_stride = 2
@@ -810,6 +820,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
820
ssm_state[cache_blocks_to_fill] = from_where
821
822
# For all seqs, store the last state (note: might be partial):
823
+ assert state_indices_tensor_p is not None
824
ssm_state[
825
state_indices_tensor_p.gather(
826
1, block_idx_last_scheduled_token_p.unsqueeze(1)
@@ -820,10 +831,12 @@ class MambaMixer2(MambaBase, PluggableLayer):
831
# update ssm states
832
# - varlen state is a (num_prefills, nheads, headdim, dstate)
833
# tensor
834
+ assert state_indices_tensor_p is not None
835
ssm_state[state_indices_tensor_p] = varlen_states
836
837
# Process decode requests
838
if has_decode:
839
+ assert state_indices_tensor_d is not None
840
if is_mamba_cache_all:
841
state_indices_tensor_d_input = state_indices_tensor_d.gather(
842
1, block_idx_last_computed_token_d.unsqueeze(1)
vllm/model_executor/layers/mamba/short_conv.py
+5
-4
@@ -113,10 +113,11 @@ class ShortConv(MambaBase, CustomOp):
113
# chunked prefill modes; they are computed at top-level model forward
114
# since they stay the same and reused for all mamba layers in the same
115
# iteration.
116
- attn_metadata: AttentionMetadata = forward_context.attn_metadata
117
- if attn_metadata is not None:
118
- assert isinstance(attn_metadata, dict)
119
- attn_metadata = attn_metadata[self.prefix]
116
+ attn_metadata_raw = forward_context.attn_metadata
117
+ attn_metadata: AttentionMetadata | None = None
118
+ if attn_metadata_raw is not None:
119
+ assert isinstance(attn_metadata_raw, dict)
120
+ attn_metadata = attn_metadata_raw[self.prefix]
121
assert isinstance(attn_metadata, ShortConvAttentionMetadata)
122
conv_state = (
123
self.kv_cache[0]
vllm/model_executor/layers/pooler/seqwise/poolers.py
+1
@@ -115,6 +115,7 @@ def pooler_for_classify(
115
116
vllm_config = get_current_vllm_config()
117
model_config = vllm_config.model_config
118
+ assert model_config.pooler_config is not None
119
head = ClassifierPoolerHead(
120
head_dtype=model_config.head_dtype,
121
classifier=classifier,
vllm/model_executor/layers/pooler/tokwise/poolers.py
+1
@@ -124,6 +124,7 @@ def pooler_for_token_classify(
124
125
vllm_config = get_current_vllm_config()
126
model_config = vllm_config.model_config
127
+ assert model_config.pooler_config is not None
128
head = TokenClassifierPoolerHead(
129
head_dtype=model_config.head_dtype,
130
classifier=classifier,
vllm/model_executor/layers/quantization/fp_quant.py
+6
-2
@@ -3,7 +3,7 @@
3
4
# Supports FP-Quant compression, see https://arxiv.org/abs/2509.23202
5
6
-from typing import Any
6
+from typing import Any, Literal, cast
7
8
import torch
9
from torch.nn.parameter import Parameter
@@ -251,7 +251,11 @@ class FPQuantLinearMethod(LinearMethodBase):
251
def fused_quantize_mx(
252
x_flat: torch.Tensor, hadamard_matrix: torch.Tensor, forward_method: str
253
) -> tuple[torch.Tensor, torch.Tensor]:
254
- return fusedQuantizeMx(x_flat, hadamard_matrix, method=forward_method)
254
+ return fusedQuantizeMx(
255
+ x_flat,
256
+ hadamard_matrix,
257
+ method=cast(Literal["quest", "abs_max"], forward_method),
258
+ )
259
260
261
def fused_quantize_mx_fake(x_flat, hadamard_matrix, forward_method):
vllm/model_executor/layers/quantization/quark/quark.py
+1
-1
@@ -114,7 +114,7 @@ class QuarkConfig(QuantizationConfig):
114
:param hf_to_vllm_mapper: maps from hf model structure (the assumed
115
structure of the qconfig) to vllm model structure
116
"""
117
- quant_config_with_hf_to_vllm_mapper = {}
117
+ quant_config_with_hf_to_vllm_mapper: dict[str, Any] = {}
118
119
for k, v in self.quant_config.items():
120
if isinstance(v, list):
vllm/model_executor/layers/sparse_attn_indexer.py
+14
-12
@@ -26,7 +26,7 @@ from vllm.v1.worker.workspace import current_workspace_manager
26
if current_platform.is_cuda_alike():
27
from vllm import _custom_ops as ops
28
elif current_platform.is_xpu():
29
- from vllm._xpu_ops import xpu_ops as ops
29
+ from vllm._xpu_ops import xpu_ops
30
31
logger = init_logger(__name__)
32
@@ -84,12 +84,12 @@ def sparse_attn_indexer(
84
total_seq_lens,
85
topk_indices_buffer,
86
)
87
- attn_metadata = attn_metadata[k_cache_prefix]
88
- assert isinstance(attn_metadata, DeepseekV32IndexerMetadata)
89
- slot_mapping = attn_metadata.slot_mapping
90
- has_decode = attn_metadata.num_decodes > 0
91
- has_prefill = attn_metadata.num_prefills > 0
92
- num_decode_tokens = attn_metadata.num_decode_tokens
87
+ attn_metadata_narrowed = attn_metadata[k_cache_prefix]
88
+ assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata)
89
+ slot_mapping = attn_metadata_narrowed.slot_mapping
90
+ has_decode = attn_metadata_narrowed.num_decodes > 0
91
+ has_prefill = attn_metadata_narrowed.num_prefills > 0
92
+ num_decode_tokens = attn_metadata_narrowed.num_decode_tokens
93
94
# During speculative decoding, k may be padded to the CUDA graph batch
95
# size while slot_mapping only covers actual tokens. Truncate k to avoid
@@ -97,6 +97,8 @@ def sparse_attn_indexer(
97
num_tokens = slot_mapping.shape[0]
98
k = k[:num_tokens]
99
100
+ # scale_fmt can be None, but the function expects str
101
+ assert scale_fmt is not None
102
ops.indexer_k_quant_and_cache(
103
k,
104
kv_cache,
@@ -107,7 +109,7 @@ def sparse_attn_indexer(
109
110
topk_indices_buffer[: hidden_states.shape[0]] = -1
111
if has_prefill:
110
- prefill_metadata = attn_metadata.prefill
112
+ prefill_metadata = attn_metadata_narrowed.prefill
113
assert prefill_metadata is not None
114
115
# Get the full shared workspace buffers once (will allocate on first use)
@@ -144,7 +146,7 @@ def sparse_attn_indexer(
146
]
147
148
if current_platform.is_xpu():
147
- ops.top_k_per_row_prefill(
149
+ xpu_ops.top_k_per_row_prefill( # type: ignore[attr-defined]
150
logits,
151
chunk.cu_seqlen_ks,
152
chunk.cu_seqlen_ke,
@@ -167,7 +169,7 @@ def sparse_attn_indexer(
169
)
170
171
if has_decode:
170
- decode_metadata = attn_metadata.decode
172
+ decode_metadata = attn_metadata_narrowed.decode
173
assert decode_metadata is not None
174
# kv_cache shape [
175
# kv_cache size requirement [num_block, block_size, n_head, head_dim],
@@ -217,11 +219,11 @@ def sparse_attn_indexer(
219
topk_indices,
220
topk_workspace,
221
topk_tokens,
220
- attn_metadata.max_seq_len,
222
+ attn_metadata_narrowed.max_seq_len,
223
)
224
else:
225
if current_platform.is_xpu():
224
- ops.top_k_per_row_decode(
226
+ xpu_ops.top_k_per_row_decode( # type: ignore[attr-defined]
227
logits,
228
next_n,
229
seq_lens,