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 10, 12:45 UTC
[0/N][Attention] Fix miscellaneous pre-commit issues (#31924)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Jan 7, 2026 at 20:15 UTC
0d7667419f738e44ce7f9bf311987e15b01970b0
8 files changed
+35
-25
vllm/attention/layers/static_sink_attention.py
+1
-1
@@ -140,7 +140,7 @@ class StaticSinkAttention(Attention, CustomOp):
140
head_size, dtype, kv_cache_dtype, block_size
141
)
142
attn_backend = create_static_sink_attention_backend(
143
- underlying_attn_backend,
143
+ underlying_attn_backend, # type: ignore[arg-type]
144
sink_len=sink_len,
145
)
146
Attention.__init__(
vllm/attention/ops/flashmla.py
+5
-2
@@ -55,7 +55,7 @@ def is_flashmla_dense_supported() -> tuple[bool, str | None]:
55
is_availble, maybe_reason = _is_flashmla_available()
56
if not is_availble:
57
return False, maybe_reason
58
- if current_platform.get_device_capability()[0] != 9:
58
+ if not current_platform.is_device_capability_family(90):
59
return False, "FlashMLA Dense is only supported on Hopper devices."
60
return True, None
61
@@ -67,7 +67,10 @@ def is_flashmla_sparse_supported() -> tuple[bool, str | None]:
67
is_availble, maybe_reason = _is_flashmla_available()
68
if not is_availble:
69
return False, maybe_reason
70
- if current_platform.get_device_capability()[0] not in (9, 10):
70
+ if not (
71
+ current_platform.is_device_capability_family(90)
72
+ or current_platform.is_device_capability_family(100)
73
+ ):
74
return (
75
False,
76
"FlashMLA Sparse is only supported on Hopper and Blackwell devices.",
vllm/attention/ops/paged_attn.py
+6
-2
@@ -7,9 +7,13 @@ import torch
7
from vllm.platforms import current_platform
8
9
if current_platform.is_cuda_alike():
10
- from vllm import _custom_ops as ops
10
+ from vllm import _custom_ops
11
+
12
+ ops = _custom_ops
13
elif current_platform.is_xpu():
12
- from vllm._ipex_ops import ipex_ops as ops
14
+ from vllm._ipex_ops import ipex_ops
15
+
16
+ ops = ipex_ops
17
18
19
class PagedAttention:
vllm/attention/ops/prefix_prefill.py
+2
-2
@@ -754,8 +754,8 @@ def context_attention_fwd(
754
if current_platform.is_rocm():
755
extra_kargs = {"kpack": 1, "waves_per_eu": 2}
756
757
- grid = lambda META: (batch, head, triton.cdiv(max_input_len, META["BLOCK_M"]))
758
- _fwd_kernel[grid](
757
+ grid_fn = lambda META: (batch, head, triton.cdiv(max_input_len, META["BLOCK_M"]))
758
+ _fwd_kernel[grid_fn](
759
q,
760
k,
761
v,
vllm/attention/ops/rocm_aiter_mla_sparse.py
+3
-3
@@ -37,9 +37,9 @@ def fp8_mqa_logits_torch(
37
Returns:
38
Logits tensor of shape [M, N], dtype `torch.float32`.
39
"""
40
- kv, scale = kv
41
- seq_len_kv = kv.shape[0]
42
- k = kv.to(torch.bfloat16)
40
+ k_fp8, scale = kv
41
+ seq_len_kv = k_fp8.shape[0]
42
+ k = k_fp8.to(torch.bfloat16)
43
q = q.to(torch.bfloat16)
44
45
mask_lo = (
vllm/attention/ops/triton_decode_attention.py
+1
-4
@@ -282,10 +282,7 @@ def _fwd_grouped_kernel_stage1(
282
cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H)
283
split_kv_id = tl.program_id(2)
284
285
- if kv_group_num > BLOCK_H:
286
- VALID_BLOCK_H: tl.constexpr = BLOCK_H
287
- else:
288
- VALID_BLOCK_H: tl.constexpr = kv_group_num
285
+ VALID_BLOCK_H: tl.constexpr = BLOCK_H if kv_group_num > BLOCK_H else kv_group_num
286
cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
287
mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H
288
mask_h = mask_h & (cur_head < q_head_num)
vllm/attention/ops/triton_prefill_attention.py
+3
-3
@@ -202,9 +202,9 @@ def _fwd_kernel(
202
def get_block_size(dtype: torch.dtype) -> int:
203
if dtype == torch.float32:
204
return 32
205
- elif (
206
- current_platform.is_cuda_alike()
207
- ) and current_platform.get_device_capability().major > 8:
205
+ elif current_platform.is_cuda_alike() and current_platform.has_device_capability(
206
+ 80
207
+ ):
208
return 128
209
else:
210
return 64
vllm/attention/utils/fa_utils.py
+14
-8
@@ -7,16 +7,23 @@ from vllm.platforms import current_platform
7
logger = init_logger(__name__)
8
9
if current_platform.is_cuda():
10
- from vllm import _custom_ops as ops
10
+ from vllm import _custom_ops
11
12
+ ops = _custom_ops
13
reshape_and_cache_flash = ops.reshape_and_cache_flash
13
- from vllm.vllm_flash_attn import flash_attn_varlen_func, get_scheduler_metadata
14
+ from vllm.vllm_flash_attn import ( # type: ignore[attr-defined]
15
+ flash_attn_varlen_func,
16
+ get_scheduler_metadata,
17
+ )
18
+
19
elif current_platform.is_xpu():
15
- from vllm._ipex_ops import ipex_ops as ops
20
+ from vllm._ipex_ops import ipex_ops
21
22
+ ops = ipex_ops
23
reshape_and_cache_flash = ops.reshape_and_cache_flash
24
flash_attn_varlen_func = ops.flash_attn_varlen_func
25
get_scheduler_metadata = ops.get_scheduler_metadata
26
+
27
elif current_platform.is_rocm():
28
try:
29
from flash_attn import flash_attn_varlen_func # noqa: F401
@@ -85,7 +92,7 @@ def get_flash_attn_version(requires_alibi: bool = False) -> int | None:
92
def flash_attn_supports_fp8() -> bool:
93
return (
94
get_flash_attn_version() == 3
88
- and current_platform.get_device_capability().major == 9
95
+ and current_platform.is_device_capability_family(90)
96
)
97
98
@@ -105,10 +112,9 @@ def flash_attn_supports_mla():
112
is_fa_version_supported,
113
)
114
108
- return (
109
- is_fa_version_supported(3)
110
- and current_platform.get_device_capability()[0] == 9
111
- )
115
+ return is_fa_version_supported(
116
+ 3
117
+ ) and current_platform.is_device_capability_family(90)
118
except (ImportError, AssertionError):
119
pass
120
return False