4fb767ff0a · tduncan/Nvidia-Optimization-Stack-VLLM

tests/kernels/attention/test_attention_selector.py

+2 -6

@@ -848,7 +848,7 @@ def hopper_selection():

848 def test_hopper_mm_prefix_selects_triton_flash_attn(

849 use_mm_prefix, flash_attn_version, hopper_selection

850 ):

851 - """Hopper selects the composite only when its causal route resolves FA4."""

851 + """Hopper selects the composite when its causal route resolves FA3 or FA4."""

852 from vllm.engine.arg_utils import EngineArgs

853

854 config = EngineArgs(

@@ -860,11 +860,7 @@ def test_hopper_mm_prefix_selects_triton_flash_attn(

860 backend = get_attn_backend(

861 256, torch.bfloat16, None, use_mm_prefix=use_mm_prefix

862 )

863 - fa4_resolved = config.attention_config.flash_attn_version == 4

864 - if use_mm_prefix:

865 - expected = "TRITON_FLASH_ATTN" if fa4_resolved else "TRITON_ATTN"

866 - else:

867 - expected = "FLASH_ATTN"

863 + expected = "TRITON_FLASH_ATTN" if use_mm_prefix else "FLASH_ATTN"

864 assert backend.get_name() == expected

865

866

tests/v1/attention/test_backend_per_kind.py

+1 -1

@@ -29,7 +29,7 @@ from vllm.v1.kv_cache_interface import KVCacheSpecKind

29 def test_largest_kernel_block_within(supported_sizes, expected):

30 class Backend:

31 @staticmethod

32 - def get_supported_kernel_block_sizes():

32 + def get_supported_kernel_block_sizes(kv_cache_spec=None):

33 return supported_sizes

34

35 assert _largest_kernel_block_within(Backend, 1024, 1024 * 1536, 2048) == expected

vllm/model_executor/layers/attention/attention.py

+14 -5

@@ -8,7 +8,10 @@ import torch.nn as nn

8

9 import vllm.envs as envs

10 from vllm.compilation.breakable_cudagraph import eager_break_during_capture

11 -from vllm.config import CacheConfig, get_current_vllm_config

11 +from vllm.config import (

12 + CacheConfig,

13 + get_current_vllm_config,

14 +)

15 from vllm.config.vllm import VllmConfig

16 from vllm.forward_context import ForwardContext, get_forward_context

17 from vllm.logger import init_logger

@@ -100,6 +103,7 @@ def _largest_kernel_block_within(

103 per_token_bytes: int,

104 page_budget: int,

105 fallback: int,

106 + kv_cache_spec: KVCacheSpec | None = None,

107 ) -> int:

108 """Largest supported kernel block size whose page fits in ``page_budget``.

109

@@ -111,7 +115,7 @@ def _largest_kernel_block_within(

115 """

116 from vllm.v1.attention.backend import MultipleOf

117

114 - sizes = attn_backend.get_supported_kernel_block_sizes()

118 + sizes = attn_backend.get_supported_kernel_block_sizes(kv_cache_spec)

119 max_block_size = page_budget // per_token_bytes

120 candidates = [s for s in sizes if isinstance(s, int)]

121 candidates.extend(

@@ -632,7 +636,7 @@ class Attention(nn.Module, AttentionLayerBase):

636 # ``unify`` scales it up by an integer ratio.

637 shared_page = vllm_config.cache_config.skip_page_size_padded

638 # The backend owns its packing

635 - sw_per_token = self.attn_backend.customize_spec(

639 + kv_cache_spec = self.attn_backend.customize_spec(

640 SlidingWindowSpec(

641 block_size=1,

642 num_kv_heads=self.num_kv_heads,

@@ -642,10 +646,15 @@ class Attention(nn.Module, AttentionLayerBase):

646 kv_quant_mode=quant_mode,

647 sliding_window=self.sliding_window,

648 )

645 - ).real_page_size_bytes

649 + )

650 + sw_per_token = kv_cache_spec.real_page_size_bytes

651 page_budget = shared_page or sw_per_token * block_size

652 sw_block_size = _largest_kernel_block_within(

648 - self.attn_backend, sw_per_token, page_budget, block_size

653 + self.attn_backend,

654 + sw_per_token,

655 + page_budget,

656 + block_size,

657 + kv_cache_spec,

658 )

659 return SlidingWindowSpec(

660 block_size=sw_block_size,

vllm/model_executor/models/config.py

+26 -12

@@ -212,12 +212,15 @@ class Gemma4Config(VerifyAndUpdateConfig):

212 """Configure attention for heterogeneous head dimensions.

213

214 Gemma4 uses different head dimensions for sliding window vs full attention

215 - layers. The default FA3 on Hopper cannot handle head_dim > 256, which causes

216 - mixed backend selection and numerical divergence.

217 -

218 - When FA4 is available we force it for ALL layers, giving a uniform kernel path

219 - and avoiding the mixed FA3+FA4 penalty. When FA4 is not available we fall back

220 - to Triton.

215 + layers. The default FA3 on Hopper cannot handle head_dim > 256.

216 +

217 + On SM90 with FP8 KV cache, use FA3 for supported layers and let the generic

218 + FlashAttention selector upgrade larger head dimensions to FA4.

219 + The multimodal-prefix composite routes image masks to Triton and causal

220 + requests to this per-layer FA3/FA4 selection. For other configurations,

221 + force FA4 for all layers to avoid the mixed

222 + FA3+FA4 penalty.

223 + When FA4 is not available, fall back to Triton.

224 """

225 model_config = vllm_config.model_config

226 arch_config = model_config.model_arch_config

@@ -230,6 +233,7 @@ class Gemma4Config(VerifyAndUpdateConfig):

233 if len(set(head_dims.values())) <= 1:

234 return

235

236 + from vllm.platforms import current_platform

237 from vllm.v1.attention.backends.fa_utils import is_fa_version_supported

238 from vllm.v1.attention.backends.registry import AttentionBackendEnum

239

@@ -241,12 +245,22 @@ class Gemma4Config(VerifyAndUpdateConfig):

245 and vllm_config.attention_config.backend

246 in (None, AttentionBackendEnum.FLASH_ATTN)

247 ):

244 - vllm_config.attention_config.flash_attn_version = 4

245 - logger.info(

246 - "Gemma4 model has heterogeneous head dimensions %s. Using FA4 for "

247 - "all layers to avoid mixed FA3/FA4 penalty.",

248 - head_dims,

249 - )

248 + use_per_layer_fa = current_platform.is_device_capability_family(

249 + 90

250 + ) and vllm_config.cache_config.cache_dtype.startswith("fp8")

251 + if use_per_layer_fa:

252 + logger.info(

253 + "Gemma4 model has heterogeneous head dimensions %s. Using "

254 + "per-layer FA3/FA4 selection for FP8 KV cache on SM90.",

255 + head_dims,

256 + )

257 + else:

258 + vllm_config.attention_config.flash_attn_version = 4

259 + logger.info(

260 + "Gemma4 model has heterogeneous head dimensions %s. Using FA4 "

261 + "for all layers to avoid mixed FA3/FA4 penalty.",

262 + head_dims,

263 + )

264 elif vllm_config.attention_config.backend is None:

265 vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN

266 logger.info(

vllm/models/deepseek_v4/compressor.py

+1 -1

@@ -66,7 +66,7 @@ class CompressorBackend(AttentionBackend):

66 return "CompressorBackend"

67

68 @staticmethod

69 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

69 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

70 return [MultipleOf(1)]

71

72 @classmethod

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

+1 -1

@@ -111,7 +111,7 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):

111 ]

112

113 @staticmethod

114 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

114 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

115 return [256]

116

117 @staticmethod

vllm/models/deepseek_v4/sparse_mla.py

+1 -1

@@ -57,7 +57,7 @@ class DeepseekV4SparseMLABackend(AttentionBackend):

57 ]

58

59 @staticmethod

60 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

60 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

61 return [256]

62

63 @staticmethod

vllm/models/deepseek_v41/compressor.py

+1 -1

@@ -37,7 +37,7 @@ class CompressorBackend(AttentionBackend):

37 return "CompressorBackend"

38

39 @staticmethod

40 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

40 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

41 return [MultipleOf(1)]

42

43 @classmethod

vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py

+1 -1

@@ -113,7 +113,7 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):

113 ]

114

115 @staticmethod

116 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

116 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

117 return [128]

118

119 @staticmethod

vllm/models/deepseek_v41/sparse_mla.py

+2 -2

@@ -86,7 +86,7 @@ class DeepseekV4SparseMLABackend(AttentionBackend):

86 ]

87

88 @staticmethod

89 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

89 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

90 return [64 if current_platform.is_device_capability_family(90) else 128]

91

92 @staticmethod

@@ -255,7 +255,7 @@ class FlashMLAMegaAttnBackend(DeepseekV4FlashMLABackend):

255 return "FLASHMLA_MEGA_ATTN_DSV41"

256

257 @staticmethod

258 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

258 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

259 return [128]

260

261 @classmethod

vllm/models/minimax_m3/common/indexer.py

+1 -1

@@ -91,7 +91,7 @@ class MiniMaxM3IndexerBackend(AttentionBackend):

91 return [128]

92

93 @staticmethod

94 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

94 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

95 return [128]

96

97 @classmethod

vllm/models/minimax_m3/common/sparse_attention.py

+1 -1

@@ -159,7 +159,7 @@ class MiniMaxM3SparseBackend(AttentionBackend):

159 return [128]

160

161 @staticmethod

162 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

162 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

163 # Page size == sparse block size (one sparse block per KV page).

164 return [128]

165

vllm/models/qwen4_exp/amd/qsa.py

+1 -1

@@ -75,7 +75,7 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):

75 return "QWEN4_EXP_QSA_TRITON"

76

77 @staticmethod

78 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

78 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

79 # QSA consumes manager pages directly and does not use FA4 paged attention.

80 return [MultipleOf(16)]

81

vllm/models/qwen4_exp/nvidia/qsa.py

+1 -1

@@ -103,7 +103,7 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):

103 return "QWEN4_EXP_QSA_TRITON"

104

105 @staticmethod

106 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

106 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

107 # QSA consumes manager pages directly and does not use FA4 paged attention.

108 return [MultipleOf(16)]

109

vllm/v1/attention/backend.py

+3 -1

@@ -69,7 +69,9 @@ class AttentionBackend(ABC):

69 forward_includes_kv_cache_update: bool = True

70

71 @staticmethod

72 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

72 + def get_supported_kernel_block_sizes(

73 + kv_cache_spec: "KVCacheSpec | None" = None,

74 + ) -> list[int | MultipleOf]:

75 return [MultipleOf(1)]

76

77 @staticmethod

vllm/v1/attention/backends/b12x.py

+1 -1

@@ -183,7 +183,7 @@ class B12xPagedAttentionBackend(AttentionBackend):

183 return B12xPagedMetadataBuilder

184

185 @staticmethod

186 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

186 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

187 return list(_B12X_SUPPORTED_PAGE_SIZES)

188

189 @classmethod

vllm/v1/attention/backends/composite.py

+11 -5

@@ -116,14 +116,16 @@ class MMPrefixAttentionRouting(CompositeAttentionRouting):

116

117

118 def _intersect_kernel_block_sizes(

119 - first: type[AttentionBackend], second: type[AttentionBackend]

119 + first: type[AttentionBackend],

120 + second: type[AttentionBackend],

121 + kv_cache_spec: KVCacheSpec | None = None,

122 ):

123 from vllm.v1.attention.backend import MultipleOf

124

125 result: list[int | MultipleOf] = []

126 seen: set[tuple[type, int]] = set()

125 - for lhs in first.get_supported_kernel_block_sizes():

126 - for rhs in second.get_supported_kernel_block_sizes():

127 + for lhs in first.get_supported_kernel_block_sizes(kv_cache_spec):

128 + for rhs in second.get_supported_kernel_block_sizes(kv_cache_spec):

129 candidate: int | MultipleOf | None = None

130 if isinstance(lhs, MultipleOf) and isinstance(rhs, MultipleOf):

131 candidate = MultipleOf(math.lcm(lhs.base, rhs.base))

@@ -378,7 +380,9 @@ def create_composite_attention_backend(

380 return CompositeAttentionMetadataBuilder

381

382 @staticmethod

381 - def get_supported_kernel_block_sizes():

383 + def get_supported_kernel_block_sizes(

384 + kv_cache_spec: KVCacheSpec | None = None,

385 + ):

386 if kernel_block_sizes:

387 return [

388 size

@@ -388,7 +392,9 @@ def create_composite_attention_backend(

392 for backend in (general_backend, causal_backend)

393 )

394 ]

391 - return _intersect_kernel_block_sizes(general_backend, causal_backend)

395 + return _intersect_kernel_block_sizes(

396 + general_backend, causal_backend, kv_cache_spec

397 + )

398

399 @classmethod

400 def supports_block_size(cls, block_size):

vllm/v1/attention/backends/cpu_attn.py

+1 -1

@@ -57,7 +57,7 @@ class CPUAttentionBackend(AttentionBackend):

57 ]

58

59 @staticmethod

60 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

60 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

61 return [MultipleOf(32)]

62

63 @classmethod

vllm/v1/attention/backends/fa_utils.py

+5 -1

@@ -303,7 +303,11 @@ def flash_attn_supports_kv_cache_dtype(

303 kv_cache_block_size=kv_cache_block_size,

304 supports_fa4_hd256=supports_fa4_hd256,

305 )

306 - return (fa_version == 3 and current_platform.is_device_capability_family(90)) or (

306 + is_sm90 = current_platform.is_device_capability_family(90)

307 + sm90_fp8_kv_supported = fa_version == 3 or (

308 + fa_version == 4 and head_size in (None, 512)

309 + )

310 + return (sm90_fp8_kv_supported and is_sm90) or (

311 fa_version == 4 and current_platform.is_device_capability_family(100)

312 )

313

vllm/v1/attention/backends/flash_attn.py

+81 -16

@@ -291,18 +291,56 @@ class FlashAttentionBackend(AttentionBackend):

291 ]

292 head_size_v: int | None = None

293

294 + @staticmethod

295 + def _get_sm90_fa4_fp8_kv_block_size(

296 + kv_cache_spec: KVCacheSpec | None = None,

297 + ) -> int | None:

298 + if kv_cache_spec is None:

299 + vllm_config = get_current_vllm_config_or_none()

300 + if vllm_config is None or vllm_config.model_config is None:

301 + return None

302 + head_size = vllm_config.model_config.get_head_size()

303 + uses_fp8_kv = vllm_config.cache_config.cache_dtype in (

304 + "fp8",

305 + "fp8_e4m3",

306 + )

307 + else:

308 + if not isinstance(kv_cache_spec, AttentionSpec):

309 + return None

310 + head_size = kv_cache_spec.head_size

311 + uses_fp8_kv = kv_cache_spec.kv_quant_mode == KVQuantMode.FP8_PER_TENSOR

312 +

313 + if (

314 + current_platform.is_device_capability_family(90)

315 + and uses_fp8_kv

316 + and head_size == 512

317 + and get_flash_attn_version(head_size=head_size) == 4

318 + ):

319 + # The SM90 FP8-KV-dequant kernel uses a 64-token TMA tile/page.

320 + return 64

321 + return None

322 +

323 @classmethod

295 - def _get_fa4_hd256_block_size(cls) -> int | None:

296 - vllm_config = get_current_vllm_config_or_none()

297 - if vllm_config is None or vllm_config.model_config is None:

298 - return None

324 + def _get_fa4_hd256_block_size(

325 + cls, kv_cache_spec: KVCacheSpec | None = None

326 + ) -> int | None:

327 + if kv_cache_spec is None:

328 + vllm_config = get_current_vllm_config_or_none()

329 + if vllm_config is None or vllm_config.model_config is None:

330 + return None

331 + head_size = vllm_config.model_config.get_head_size()

332 + head_size_v = cls.head_size_v

333 + else:

334 + if not isinstance(kv_cache_spec, AttentionSpec):

335 + return None

336 + head_size = kv_cache_spec.head_size

337 + head_size_v = kv_cache_spec.head_size_v

338

300 - head_size = vllm_config.model_config.get_head_size()

339 if (

302 - uses_fa4_hd256_kernel(head_size, cls.head_size_v)

340 + uses_fa4_hd256_kernel(head_size, head_size_v)

341 and get_flash_attn_version(

342 head_size=head_size,

305 - head_size_v=cls.head_size_v,

343 + head_size_v=head_size_v,

344 supports_fa4_hd256=True,

345 )

346 == 4

@@ -311,8 +349,12 @@ class FlashAttentionBackend(AttentionBackend):

349 return None

350

351 @classmethod

314 - def get_supported_kernel_block_sizes(cls) -> list[int | MultipleOf]:

315 - if block_size := cls._get_fa4_hd256_block_size():

352 + def get_supported_kernel_block_sizes(

353 + cls, kv_cache_spec: KVCacheSpec | None = None

354 + ) -> list[int | MultipleOf]:

355 + if block_size := cls._get_sm90_fa4_fp8_kv_block_size(kv_cache_spec):

356 + return [block_size]

357 + if block_size := cls._get_fa4_hd256_block_size(kv_cache_spec):

358 # Sliding-window specs select the smallest advertised size.

359 return [block_size]

360 return [MultipleOf(16)]

@@ -321,6 +363,8 @@ class FlashAttentionBackend(AttentionBackend):

363

364 @classmethod

365 def get_preferred_block_size(cls, default_block_size: int) -> int:

366 + if block_size := cls._get_sm90_fa4_fp8_kv_block_size():

367 + return max(default_block_size, block_size)

368 if block_size := cls._get_fa4_hd256_block_size():

369 return max(default_block_size, block_size)

370 if current_platform.is_xpu():

@@ -419,6 +463,13 @@ class FlashAttentionBackend(AttentionBackend):

463 ) -> str | None:

464 if has_sink and device_capability < DeviceCapability(9, 0):

465 return "sink not supported on compute capability < 9.0"

466 + if (

467 + use_mm_prefix

468 + and kv_cache_dtype is not None

469 + and is_quantized_kv_cache(kv_cache_dtype)

470 + and device_capability == DeviceCapability(9, 0)

471 + ):

472 + return "SM90 FP8 KV with mm_prefix requires Triton"

473 if (

474 kv_cache_dtype is not None

475 and is_quantized_kv_cache(kv_cache_dtype)

@@ -431,7 +482,10 @@ class FlashAttentionBackend(AttentionBackend):

482 supports_fa4_hd256=True,

483 )

484 ):

434 - return "FP8 KV cache requires FA3 on SM90 or FA4 on SM100"

485 + return (

486 + "FP8 KV cache requires FA3 on SM90, FA4 with head_size=512 "

487 + "on SM90, or FA4 on SM100"

488 + )

489 if (

490 use_mm_prefix

491 and get_flash_attn_version(

@@ -514,15 +568,16 @@ def _get_sliding_window_configs(

568 ) -> set[tuple[int, int] | None]:

569 """Get the set of all sliding window configs used in the model.

570

517 - Only inspects FlashAttentionImpl layers. Other backends (e.g.

518 - TurboQuant, MLA) use their own metadata builders and are skipped.

571 + Only inspects FlashAttentionImpl layers, including implementations wrapped

572 + by a composite backend. Other backends (e.g. TurboQuant, MLA) use their own

573 + metadata builders and are skipped.

574 """

575 sliding_window_configs: set[tuple[int, int] | None] = set()

576 layers = get_layers_from_vllm_config(vllm_config, Attention)

577 for layer in layers.values():

523 - if not isinstance(layer.impl, FlashAttentionImpl):

524 - continue

525 - sliding_window_configs.add(layer.impl.sliding_window)

578 + for impl in layer.impl.get_impl_variants():

579 + if isinstance(impl, FlashAttentionImpl):

580 + sliding_window_configs.add(impl.sliding_window)

581 return sliding_window_configs

582

583

@@ -1145,7 +1200,17 @@ class FlashAttentionImpl(AttentionImpl):

1200 "heads in the layer"

1201 )

1202

1148 - self.supports_quant_query_input = flash_attn_supports_quant_query_input()

1203 + # FA4's SM90 FP8-KV path consumes native FP16/BF16 Q and dequantizes

1204 + # FP8 K/V in-kernel. Other FA4 paths (notably SM100) still require Q,

1205 + # K, and V to have the same FP8 dtype.

1206 + uses_sm90_fa4_fp8_kv_dequant = (

1207 + self.vllm_flash_attn_version == 4

1208 + and current_platform.is_device_capability_family(90)

1209 + and self.kv_cache_dtype in ("fp8", "fp8_e4m3")

1210 + )

1211 + self.supports_quant_query_input = flash_attn_supports_quant_query_input() and (

1212 + not uses_sm90_fa4_fp8_kv_dequant

1213 + )

1214

1215 dcp_a2a = (

1216 vllm_config is not None

vllm/v1/attention/backends/flashinfer.py

+1 -1

@@ -423,7 +423,7 @@ class FlashInferBackend(AttentionBackend):

423 ]

424

425 @staticmethod

426 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

426 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

427 # Page sizes >= 128 only run on the trtllm-gen dynamic kernel (GQA/MQA

428 # on Blackwell); advertise them only when usable so selection never

429 # picks a large kernel block we cannot serve.

vllm/v1/attention/backends/flex_attention.py

+1 -1

@@ -154,7 +154,7 @@ class FlexAttentionBackend(AttentionBackend):

154 return []

155

156 @staticmethod

157 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

157 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

158 return [MultipleOf(16)]

159

160

vllm/v1/attention/backends/hpc_attn.py

+1 -1

@@ -288,7 +288,7 @@ class HpcAttentionBackend(AttentionBackend):

288 forward_includes_kv_cache_update: bool = True

289

290 @staticmethod

291 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

291 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

292 return [64]

293

294 @staticmethod

vllm/v1/attention/backends/mla/amx_mla.py

+1 -1

@@ -130,7 +130,7 @@ class AMXMLABackend(MLACommonBackend):

130 return [576]

131

132 @staticmethod

133 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

133 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

134 return [MultipleOf(32)]

135

136 @classmethod

vllm/v1/attention/backends/mla/cpu_mla.py

+1 -1

@@ -64,7 +64,7 @@ class CPUMLABackend(MLACommonBackend):

64 return [576]

65

66 @staticmethod

67 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

67 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

68 return [16]

69

70 @classmethod

vllm/v1/attention/backends/mla/cutlass_mla.py

+1 -1

@@ -46,7 +46,7 @@ class CutlassMLABackend(MLACommonBackend):

46 ]

47

48 @staticmethod

49 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

49 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

50 return [128]

51

52 @staticmethod

vllm/v1/attention/backends/mla/flashattn_mla.py

+1 -1

@@ -49,7 +49,7 @@ class FlashAttnMLABackend(MLACommonBackend):

49 ]

50

51 @staticmethod

52 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

52 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

53 return [MultipleOf(16)]

54

55 @staticmethod

vllm/v1/attention/backends/mla/flashattn_mla_sparse.py

+1 -1

@@ -40,7 +40,7 @@ class FlashAttnMLASparseBackend(AttentionBackend):

40 ]

41

42 @staticmethod

43 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

43 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

44 return [64]

45

46 @staticmethod

vllm/v1/attention/backends/mla/flashinfer_mla.py

+1 -1

@@ -201,7 +201,7 @@ class FlashInferMLABackend(MLACommonBackend):

201 ]

202

203 @staticmethod

204 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

204 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

205 return [32, 64]

206

207 @staticmethod

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

+2 -2

@@ -83,7 +83,7 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):

83 ]

84

85 @staticmethod

86 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

86 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

87 return [32, 64]

88

89 @staticmethod

@@ -172,7 +172,7 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):

172 return "FLASHINFER_MLA_SPARSE_SM120"

173

174 @staticmethod

175 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

175 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

176 return [64, 256]

177

178 @staticmethod

vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py

+1 -1

@@ -74,7 +74,7 @@ class FlashInferMLASparseSM90Backend(AttentionBackend):

74 ]

75

76 @staticmethod

77 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

77 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

78 return [MultipleOf(64)]

79

80 @staticmethod

vllm/v1/attention/backends/mla/flashmla.py

+1 -1

@@ -55,7 +55,7 @@ class FlashMLABackend(MLACommonBackend):

55 ]

56

57 @staticmethod

58 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

58 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

59 return [64]

60

61 @staticmethod

vllm/v1/attention/backends/mla/flashmla_sparse.py

+1 -1

@@ -139,7 +139,7 @@ class FlashMLASparseBackend(AttentionBackend):

139 ]

140

141 @staticmethod

142 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

142 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

143 return [64]

144

145 @staticmethod

vllm/v1/attention/backends/mla/indexer.py

+4 -4

@@ -200,7 +200,7 @@ class DeepseekV32IndexerBackend(AttentionBackend):

200 return "DEEPSEEK_V32_INDEXER"

201

202 @staticmethod

203 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

203 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

204 return [1, MultipleOf(16)] if current_platform.is_rocm() else [64]

205

206 @classmethod

@@ -228,7 +228,7 @@ class KpoolTailBackend(DeepseekV32IndexerBackend):

228 return []

229

230 @staticmethod

231 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

231 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

232 return [MultipleOf(1)]

233

234 @staticmethod

@@ -257,7 +257,7 @@ class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):

257 return (KVCacheLayout.BLHNC, KVCacheLayout.BLNHC)

258

259 @staticmethod

260 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

260 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

261 # Block sizes count uncompressed tokens: C4 indexer pages hold 64 rows.

262 return [256]

263

@@ -274,7 +274,7 @@ class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend):

274 return "DEEPSEEK_V41_INDEXER"

275

276 @staticmethod

277 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

277 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

278 return [64 if current_platform.is_device_capability_family(90) else 128]

279

280

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

+1 -1

@@ -282,7 +282,7 @@ class AiterMLABackend(MLACommonBackend):

282 return []

283

284 @staticmethod

285 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

285 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

286 # The aiter MLA decode kernel always operates with page_size=1

287 # internally (the wrapper flattens kv_buffer via .view(-1, 1, 1, H)).

288 # We support any kernel_block_size by expanding block-level indices

vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py

+1 -1

@@ -287,7 +287,7 @@ class ROCMAiterMLASparseBackend(AttentionBackend):

287 ]

288

289 @staticmethod

290 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

290 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

291 return [1, MultipleOf(16)]

292

293 @staticmethod

vllm/v1/attention/backends/mla/sparse_swa.py

+1 -1

@@ -148,7 +148,7 @@ class DeepseekSparseSWABackend(AttentionBackend):

148 return "DEEPSEEK_SPARSE_SWA"

149

150 @staticmethod

151 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

151 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

152 return [MultipleOf(32)]

153

154 @classmethod

vllm/v1/attention/backends/mla/tokenspeed_mla.py

+1 -1

@@ -87,7 +87,7 @@ class TokenspeedMLABackend(MLACommonBackend):

87 ]

88

89 @staticmethod

90 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

90 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

91 return [32, 64]

92

93 @staticmethod

vllm/v1/attention/backends/mla/triton_mla.py

+1 -1

@@ -107,7 +107,7 @@ class TritonMLABackend(MLACommonBackend):

107 return []

108

109 @staticmethod

110 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

110 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

111 return [MultipleOf(16)]

112

113 @classmethod

vllm/v1/attention/backends/rocm_aiter_fa.py

+1 -1

@@ -852,7 +852,7 @@ class AiterFlashAttentionBackend(AttentionBackend):

852 return attn_type in (AttentionType.DECODER,)

853

854 @staticmethod

855 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

855 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

856 if not rocm_aiter_ops.is_shuffle_kv_cache_enabled():

857 return [16, 32]

858 # Only gluon serves 128-token pages; the pa_fwd_asm/ll4mi decode

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

+1 -1

@@ -64,7 +64,7 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):

64 ]

65

66 @staticmethod

67 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

67 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

68 return [MultipleOf(16)]

69

70 @classmethod

vllm/v1/attention/backends/rocm_attn.py

+1 -1

@@ -176,7 +176,7 @@ class RocmAttentionBackend(AttentionBackend):

176 ]

177

178 @staticmethod

179 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

179 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

180 # ROCM paged attention native C++ kernel only supports block sizes 16 and 32

181 # due to shared memory (LDS) constraints on AMD GPUs.

182 # See csrc/rocm/attention.cu CALL_CUSTOM_LAUNCHER_BLK macro.

vllm/v1/attention/backends/triton_attn.py

+1 -1

@@ -304,7 +304,7 @@ class TritonAttentionBackend(AttentionBackend):

304 ]

305

306 @staticmethod

307 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

307 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

308 return [MultipleOf(16)]

309

310 @classmethod

vllm/v1/attention/backends/triton_flash_attn.py

+10 -12

@@ -15,7 +15,7 @@ from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend

15 from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend

16

17

18 -class _FA4Backend(FlashAttentionBackend):

18 +class _FABackend(FlashAttentionBackend):

19 @classmethod

20 def supports_combination(

21 cls,

@@ -42,22 +42,20 @@ class _FA4Backend(FlashAttentionBackend):

42 )

43 if reason is not None:

44 return reason

45 - if (

46 - get_flash_attn_version(

47 - head_size=head_size,

48 - has_sinks=has_sink,

49 - kv_cache_block_size=block_size,

50 - supports_fa4_hd256=True,

51 - )

52 - != 4

53 - ):

54 - return "causal route requires FlashAttention v4"

45 + fa_version = get_flash_attn_version(

46 + head_size=head_size,

47 + has_sinks=has_sink,

48 + kv_cache_block_size=block_size,

49 + supports_fa4_hd256=True,

50 + )

51 + if fa_version not in (3, 4):

52 + return "causal route requires FlashAttention v3 or v4"

53 return None

54

55

56 TritonFlashAttentionBackend = create_composite_attention_backend(

57 TritonAttentionBackend,

60 - _FA4Backend,

58 + _FABackend,

59 name="TritonFlashAttentionBackend",

60 backend_name="TRITON_FLASH_ATTN",

61 module=__name__,

vllm/v1/attention/backends/turboquant_attn.py

+1 -1

@@ -161,7 +161,7 @@ class TurboQuantAttentionBackend(AttentionBackend):

161 return (KVCacheLayout.LBNHC,)

162

163 @staticmethod

164 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

164 + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

165 return [16, 32, 64, 128]

166

167 @classmethod

vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py

+37 -2

@@ -24,6 +24,27 @@ from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (

24 logger = init_logger(__name__)

25

26

27 +def _copy_target_kv_scales(attn: nn.Module, target_attn: nn.Module) -> None:

28 + """Copy target KV scales while preserving their tensor representation.

29 +

30 + Default attention scales are scalar buffers, while some quantization

31 + methods replace them with length-one or per-head parameters. Preserve the

32 + draft layer's registration type so the shared KV cache is interpreted with

33 + the target's values and shapes without aliasing target parameters.

34 + """

35 + for scale_name in ("_k_scale", "_v_scale"):

36 + draft_scale = getattr(attn, scale_name)

37 + target_scale = getattr(target_attn, scale_name)

38 + scale = target_scale.detach().clone()

39 + if isinstance(draft_scale, nn.Parameter):

40 + scale = nn.Parameter(scale, requires_grad=False)

41 + setattr(attn, scale_name, scale)

42 + for scale_name in ("_k_scale_float", "_v_scale_float"):

43 + setattr(attn, scale_name, getattr(target_attn, scale_name))

44 + for scale_name in ("_k_scale_cpu", "_v_scale_cpu"):

45 + getattr(attn, scale_name).copy_(getattr(target_attn, scale_name))

46 +

47 +

48 class Gemma4Speculator(AutoRegressiveSpeculator):

49 @property

50 def advance_draft_positions(self) -> bool:

@@ -77,7 +98,7 @@ class Gemma4Speculator(AutoRegressiveSpeculator):

98 model: nn.Module,

99 target_attn_layer_names: set[str],

100 ) -> None:

80 - """Wire draft layers to share KV with the target model.

101 + """Wire draft layers to share KV and KV scales with the target model.

102

103 Each draft decoder layer is mapped to the last non-KV-shared

104 target layer of the same attention type (sliding or full).

@@ -92,7 +113,10 @@ class Gemma4Speculator(AutoRegressiveSpeculator):

113 return

114

115 target_num_kv_shared = getattr(target_text_config, "num_kv_shared_layers", 0)

95 - num_non_shared = len(target_layer_types) - target_num_kv_shared

116 + target_num_layers = getattr(

117 + target_text_config, "num_hidden_layers", len(target_layer_types)

118 + )

119 + num_non_shared = target_num_layers - target_num_kv_shared

120 type_to_target_indices: dict[str, list[int]] = defaultdict(list)

121 for idx, lt in enumerate(target_layer_types[:num_non_shared]):

122 type_to_target_indices[lt].append(idx)

@@ -128,6 +152,17 @@ class Gemma4Speculator(AutoRegressiveSpeculator):

152 target_idx = candidates[-1]

153 target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn"

154 attn.kv_sharing_target_layer_name = target_layer_name

155 +

156 + # KV-cache sharing aliases the cache tensor during allocation, but

157 + # the quantization scales live on the Attention modules themselves.

158 + # The BF16 draft model has no quantization config, so its K/V scales

159 + # otherwise remain at the default 1.0 while it reads the target's

160 + # calibrated FP8 cache.

161 + target_attn = self.vllm_config.compilation_config.static_forward_context[

162 + target_layer_name

163 + ]

164 + _copy_target_kv_scales(attn, target_attn)

165 +

166 logger.info(

167 "Gemma4 MTP: draft layer %d (%s) -> %s",

168 draft_idx,

vllm/vllm_flash_attn/flash_attn_interface.py

+10 -1

@@ -434,6 +434,14 @@ def flash_attn_varlen_func(

434

435 from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd

436

437 + # SM90 FA4 fp8-KV path: fp8 e4m3 paged K/V dequantized and the K/V descale folded

438 + # in-kernel; accepts bf16/fp16 Q and writes O in its native dtype (no Q cast, no

439 + # output copy). Only the (batch, num_kv_heads) f32 K/V descales are forwarded.

440 + fa4_fp8_kv_dequant = (

441 + k.dtype == torch.float8_e4m3fn

442 + and torch.cuda.get_device_capability()[0] == 9

443 + )

444 +

445 out, softmax_lse, _, _ = _flash_attn_fwd(

446 q,

447 k,

@@ -458,10 +466,11 @@ def flash_attn_varlen_func(

466 block_sparse_tensors=block_sparse_tensors,

467 aux_tensors=aux_tensors,

468 aux_tensor_leading_dims=aux_tensor_leading_dims,

461 - q_descale=q_descale,

469 + q_descale=None if fa4_fp8_kv_dequant else q_descale,

470 k_descale=k_descale,

471 v_descale=v_descale,

472 output_scale=output_scale,

473 + fp8_kv_dequant=fa4_fp8_kv_dequant,

474 )

475 else:

476 raise ValueError(f"Unsupported FA version: {fa_version}")

[Kernel] Resubmit PR 48666 - Gemma4 FP8 KV FA4 head dim 512 backend selection (#53175)