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 5, 09:45 UTC
[CI] Fix mypy for `vllm/v1/ops` (#39219)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Apr 8, 2026 at 23:06 UTC
aec18492d0c0d4144ab47b5b7cdbcbc10e6861a6
4 files changed
+28
-25
tools/pre_commit/mypy.py
-2
@@ -36,8 +36,6 @@ SEPARATE_GROUPS = [
36
EXCLUDE = [
37
"vllm/model_executor/models",
38
"vllm/model_executor/layers/fla/ops",
39
- # Ignore triton kernels in ops.
40
- "vllm/v1/attention/ops",
39
# TODO: Remove these entries after fixing mypy errors.
40
"vllm/benchmarks",
41
]
vllm/v1/attention/ops/prefix_prefill.py
+3
-1
@@ -4,6 +4,8 @@
4
# The kernels in this file are adapted from LightLLM's context_attention_fwd:
5
# https://github.com/ModelTC/lightllm/blob/main/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py
6
7
+from typing import Any
8
+
9
import torch
10
11
from vllm.platforms import current_platform
@@ -780,7 +782,7 @@ def context_attention_fwd(
782
return
783
784
max_seq_len = 0 if max_seq_len is None else max_seq_len
783
- extra_kargs = {}
785
+ extra_kargs: dict[str, Any] = {}
786
if current_platform.is_rocm():
787
extra_kargs = {}
788
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+22
-22
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import functools
4
import importlib
5
+from importlib.util import find_spec
6
7
import torch
8
@@ -276,11 +277,9 @@ def fp8_paged_mqa_logits_torch(
277
@functools.lru_cache
278
def paged_mqa_logits_module():
279
paged_mqa_logits_module_path = None
279
- if importlib.util.find_spec("aiter.ops.triton.pa_mqa_logits") is not None:
280
+ if find_spec("aiter.ops.triton.pa_mqa_logits") is not None:
281
paged_mqa_logits_module_path = "aiter.ops.triton.pa_mqa_logits"
281
- elif (
282
- importlib.util.find_spec("aiter.ops.triton.attention.pa_mqa_logits") is not None
283
- ):
282
+ elif find_spec("aiter.ops.triton.attention.pa_mqa_logits") is not None:
283
paged_mqa_logits_module_path = "aiter.ops.triton.attention.pa_mqa_logits"
284
285
if paged_mqa_logits_module_path is not None:
@@ -380,9 +379,9 @@ def fp8_mqa_logits_torch(
379
Returns:
380
Logits tensor of shape [M, N], dtype `torch.float32`.
381
"""
383
- kv, scale = kv
384
- seq_len_kv = kv.shape[0]
385
- k = kv.to(torch.bfloat16)
382
+ k_fp8, scale = kv
383
+ seq_len_kv = k_fp8.shape[0]
384
+ k = k_fp8.to(torch.bfloat16)
385
q = q.to(torch.bfloat16)
386
387
mask_lo = (
@@ -403,12 +402,9 @@ def fp8_mqa_logits_torch(
402
@functools.lru_cache
403
def mqa_logits_module():
404
mqa_logits_module_path = None
406
- if importlib.util.find_spec("aiter.ops.triton.fp8_mqa_logits") is not None:
405
+ if find_spec("aiter.ops.triton.fp8_mqa_logits") is not None:
406
mqa_logits_module_path = "aiter.ops.triton.fp8_mqa_logits"
408
- elif (
409
- importlib.util.find_spec("aiter.ops.triton.attention.fp8_mqa_logits")
410
- is not None
411
- ):
407
+ elif find_spec("aiter.ops.triton.attention.fp8_mqa_logits") is not None:
408
mqa_logits_module_path = "aiter.ops.triton.attention.fp8_mqa_logits"
409
410
if mqa_logits_module_path is not None:
@@ -455,8 +451,8 @@ def rocm_fp8_mqa_logits(
451
452
if aiter_mqa_logits_module is not None:
453
fp8_mqa_logits = aiter_mqa_logits_module.fp8_mqa_logits
458
- kv, scale = kv
459
- return fp8_mqa_logits(q, kv, scale, weights, cu_seqlen_ks, cu_seqlen_ke)
454
+ k_fp8, scale = kv
455
+ return fp8_mqa_logits(q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke)
456
else:
457
return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)
458
@@ -523,12 +519,14 @@ def rocm_aiter_sparse_attn_indexer(
519
total_seq_lens,
520
topk_indices_buffer,
521
)
526
- attn_metadata = attn_metadata[k_cache_prefix]
527
- assert isinstance(attn_metadata, DeepseekV32IndexerMetadata)
528
- slot_mapping = attn_metadata.slot_mapping
529
- has_decode = attn_metadata.num_decodes > 0
530
- has_prefill = attn_metadata.num_prefills > 0
531
- num_decode_tokens = attn_metadata.num_decode_tokens
522
+ layer_attn_metadata = attn_metadata[k_cache_prefix]
523
+ assert isinstance(layer_attn_metadata, DeepseekV32IndexerMetadata)
524
+ assert topk_indices_buffer is not None
525
+ assert scale_fmt is not None
526
+ slot_mapping = layer_attn_metadata.slot_mapping
527
+ has_decode = layer_attn_metadata.num_decodes > 0
528
+ has_prefill = layer_attn_metadata.num_prefills > 0
529
+ num_decode_tokens = layer_attn_metadata.num_decode_tokens
530
531
ops.indexer_k_quant_and_cache(
532
k,
@@ -540,7 +538,8 @@ def rocm_aiter_sparse_attn_indexer(
538
539
topk_indices_buffer[: hidden_states.shape[0]] = -1
540
if has_prefill:
543
- prefill_metadata = attn_metadata.prefill
541
+ prefill_metadata = layer_attn_metadata.prefill
542
+ assert prefill_metadata is not None
543
for chunk in prefill_metadata.chunks:
544
k_fp8 = torch.empty(
545
[chunk.total_seq_lens, head_dim],
@@ -585,7 +584,8 @@ def rocm_aiter_sparse_attn_indexer(
584
)
585
586
if has_decode:
588
- decode_metadata = attn_metadata.decode
587
+ decode_metadata = layer_attn_metadata.decode
588
+ assert decode_metadata is not None
589
# kv_cache size requirement [num_block, block_size, n_head, head_dim],
590
# we only have [num_block, block_size, head_dim],
591
kv_cache = kv_cache.unsqueeze(-2)
vllm/v1/attention/ops/vit_attn_wrappers.py
+3
@@ -292,6 +292,9 @@ def flashinfer_wrapper(
292
# RoPE has already made q and k contiguous.
293
q, k = q.contiguous(), k.contiguous()
294
295
+ assert cu_seqlens is not None
296
+ assert max_seqlen is not None
297
+ assert sequence_lengths is not None
298
assert len(cu_seqlens) % 2 == 0, "cu_seqlens must be divisible by 2"
299
cu_seqlength = len(cu_seqlens) // 2
300
batch_offsets_qko = cu_seqlens[:cu_seqlength].view(-1, 1, 1, 1)