@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 827a2af806

[Kernel] Gemma-4 FA4 FP8 Kernel (#48666)

Signed-off-by: Qixiang Lin <qixiangl@nvidia.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com> Co-authored-by: Qixiang Lin <qixiangl@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>

Jhao-Ting Chen committed Aug 13, 2026 at 21:42 UTC 827a2af806c4e4ea7bcc280f57f793e6a5fcc676

8 files changed +152 -29

cmake/external_projects/vllm_flash_attn.cmake

+1 -1

@@ -39,7 +39,7 @@ else()

39 FetchContent_Declare(

40 vllm-flash-attn

41 GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git

42 - GIT_TAG f3e1a4f74c99145c0717709860bf765de1703779

42 + GIT_TAG 617264c1c7955c9e84817654ebeedff069f3c5f1

43 GIT_PROGRESS TRUE

44 # Don't share the vllm-flash-attn build between build types

45 BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

vllm/model_executor/layers/attention/attention.py

+11 -3

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

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

@@ -94,6 +97,7 @@ def should_load_quant_weights(quant_method: QuantizeMethodBase | None) -> bool:

98 def _largest_kernel_block_within(

99 attn_backend: "type[AttentionBackend]",

100 + vllm_config: VllmConfig,

101 per_token_bytes: int,

102 page_budget: int | None,

103 fallback: int,

@@ -108,7 +112,7 @@ def _largest_kernel_block_within(

112 """

113 from vllm.v1.attention.backend import MultipleOf

114

111 - sizes = attn_backend.get_supported_kernel_block_sizes()

115 + sizes = attn_backend.get_supported_kernel_block_sizes_for_config(vllm_config)

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

117 if not candidates:

118 candidates = [s.base for s in sizes if isinstance(s, MultipleOf)]

@@ -629,7 +633,11 @@ class Attention(nn.Module, AttentionLayerBase):

633 sliding_window=self.sliding_window,

634 ).real_page_size_bytes

635 sw_block_size = _largest_kernel_block_within(

632 - self.attn_backend, sw_per_token, shared_page, block_size

636 + self.attn_backend,

637 + vllm_config,

638 + sw_per_token,

639 + shared_page,

640 + block_size,

641 )

642 return SlidingWindowSpec(

643 block_size=sw_block_size,

vllm/platforms/interface.py

+3 -5

@@ -612,7 +612,6 @@ class Platform:

612 For hybrid models, also aligns block_size with mamba page sizes.

613 """

614 from vllm.config.cache import CacheConfig

615 - from vllm.config.vllm import set_current_vllm_config

615

616 cache_config = vllm_config.cache_config

617 model_config = vllm_config.model_config

@@ -627,10 +626,9 @@ class Platform:

626

627 # Phase 1: Pick block size from backend (skip if user set --block-size)

628 if not cache_config.user_specified_block_size:

630 - with set_current_vllm_config(vllm_config):

631 - preferred = backend_cls.get_preferred_block_size(

632 - CacheConfig.DEFAULT_BLOCK_SIZE

633 - )

629 + preferred = backend_cls.get_preferred_block_size_for_config(

630 + CacheConfig.DEFAULT_BLOCK_SIZE, vllm_config

631 + )

632 if preferred != CacheConfig.DEFAULT_BLOCK_SIZE:

633 logger.info(

634 "Setting kv cache block size to %d for %s backend.",

vllm/v1/attention/backend.py

+14

@@ -69,6 +69,13 @@ class AttentionBackend(ABC):

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

70 return [MultipleOf(1)]

72 + @classmethod

73 + def get_supported_kernel_block_sizes_for_config(

74 + cls, vllm_config: "VllmConfig"

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

76 + """Return kernel block sizes for a concrete engine configuration."""

77 + return cls.get_supported_kernel_block_sizes()

78 +

79 @staticmethod

80 @abstractmethod

81 def get_name() -> str:

@@ -201,6 +208,13 @@ class AttentionBackend(ABC):

208

209 return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes)

210

211 + @classmethod

212 + def get_preferred_block_size_for_config(

213 + cls, default_block_size: int, vllm_config: "VllmConfig"

214 + ) -> int:

215 + """Return the preferred block size for a concrete engine config."""

216 + return cls.get_preferred_block_size(default_block_size)

217 +

218 @classmethod

219 def indexes_kv_by_block_stride(cls) -> bool:

220 """Whether the backend reads KV pages by the runtime block stride.

vllm/v1/attention/backends/fa_utils.py

+3 -3

@@ -241,9 +241,9 @@ def flash_attn_supports_kv_cache_dtype(

241 head_size_v=head_size_v,

242 has_sinks=has_sinks,

243 )

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

245 - fa_version == 4 and current_platform.is_device_capability_family(100)

246 - )

244 + return (

245 + fa_version in (3, 4) and current_platform.is_device_capability_family(90)

246 + ) or (fa_version == 4 and current_platform.is_device_capability_family(100))

247

248

249 def flash_attn_supports_quant_query_input() -> bool:

vllm/v1/attention/backends/flash_attn.py

+57 -2

@@ -85,13 +85,58 @@ class FlashAttentionBackend(AttentionBackend):

85 ]

87 @staticmethod

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

88 + def _get_sm90_fa4_fp8_kv_block_size(

89 + vllm_config: VllmConfig | None = None,

90 + ) -> int | None:

91 + if vllm_config is None:

92 + vllm_config = get_current_vllm_config_or_none()

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

94 + return None

95 +

96 + head_size = vllm_config.model_config.get_head_size()

97 + if (

98 + current_platform.is_device_capability_family(90)

99 + and vllm_config.cache_config.cache_dtype in ("fp8", "fp8_e4m3")

100 + and head_size == 512

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

102 + ):

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

104 + return 64

105 + return None

106 +

107 + @classmethod

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

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

110 + # Sliding-window cache specs select the smallest advertised size.

111 + # Report the kernel's exact page-size contract instead of the

112 + # generic FlashAttention multiple-of-16 capability.

113 + return [block_size]

114 + return [MultipleOf(16)]

115 +

116 + @classmethod

117 + def get_supported_kernel_block_sizes_for_config(

118 + cls, vllm_config: VllmConfig

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

120 + if block_size := cls._get_sm90_fa4_fp8_kv_block_size(vllm_config):

121 + return [block_size]

122 return [MultipleOf(16)]

123

124 forward_includes_kv_cache_update: bool = False

125

126 @classmethod

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

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

129 + return max(default_block_size, block_size)

130 + if current_platform.is_xpu():

131 + return max(default_block_size, 64)

132 + return super().get_preferred_block_size(default_block_size)

133 +

134 + @classmethod

135 + def get_preferred_block_size_for_config(

136 + cls, default_block_size: int, vllm_config: VllmConfig

137 + ) -> int:

138 + if block_size := cls._get_sm90_fa4_fp8_kv_block_size(vllm_config):

139 + return max(default_block_size, block_size)

140 if current_platform.is_xpu():

141 return max(default_block_size, 64)

142 return super().get_preferred_block_size(default_block_size)

@@ -895,7 +940,17 @@ class FlashAttentionImpl(AttentionImpl):

940 "heads in the layer"

941 )

942

898 - self.supports_quant_query_input = flash_attn_supports_quant_query_input()

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

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

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

946 + uses_sm90_fa4_fp8_kv_dequant = (

947 + self.vllm_flash_attn_version == 4

948 + and current_platform.is_device_capability_family(90)

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

950 + )

951 + self.supports_quant_query_input = flash_attn_supports_quant_query_input() and (

952 + not uses_sm90_fa4_fp8_kv_dequant

953 + )

954

955 vllm_config = get_current_vllm_config_or_none()

956 dcp_a2a = (

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

+43 -12

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

23 logger = init_logger(__name__)

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

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

28 +

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

30 + methods replace them with length-one or per-head parameters. Re-register

31 + cloned buffers on the draft layer so the shared KV cache is interpreted

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

33 + """

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

35 + target_scale = getattr(target_attn, scale_name)

36 + attn.register_buffer(scale_name, target_scale.detach().clone())

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

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

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

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

41 +

42 +

43 class Gemma4Speculator(AutoRegressiveSpeculator):

44 @property

45 def advance_draft_positions(self) -> bool:

@@ -76,7 +93,7 @@ class Gemma4Speculator(AutoRegressiveSpeculator):

93 model: nn.Module,

94 target_attn_layer_names: set[str],

95 ) -> None:

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

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

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

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

@@ -92,15 +109,19 @@ class Gemma4Speculator(AutoRegressiveSpeculator):

109

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

111 num_non_shared = len(target_layer_types) - target_num_kv_shared

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

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

97 - type_to_target_indices[lt].append(idx)

98 -

99 - target_prefix = "model.layers"

112 + target_names_by_index: dict[int, str] = {}

113 for name in target_attn_layer_names:

101 - if ".layers." in name:

102 - target_prefix = name.split(".layers.")[0] + ".layers"

103 - break

114 + _, separator, layer_suffix = name.partition(".layers.")

115 + if not separator:

116 + continue

117 + layer_index, _, _ = layer_suffix.partition(".")

118 + if layer_index.isdigit():

119 + target_names_by_index[int(layer_index)] = name

120 +

121 + type_to_target_names: dict[str, list[str]] = defaultdict(list)

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

123 + if target_name := target_names_by_index.get(idx):

124 + type_to_target_names[lt].append(target_name)

125

126 draft_layer_types = getattr(draft_text_config, "layer_types", [])

127 for draft_idx, layer in enumerate(model.model.layers):

@@ -115,7 +136,7 @@ class Gemma4Speculator(AutoRegressiveSpeculator):

136 if draft_idx < len(draft_layer_types)

137 else "full_attention"

138 )

118 - candidates = type_to_target_indices.get(draft_layer_type, [])

139 + candidates = type_to_target_names.get(draft_layer_type, [])

140 if not candidates:

141 logger.warning(

142 "No target layer of type '%s' for draft layer %d",

@@ -124,9 +145,19 @@ class Gemma4Speculator(AutoRegressiveSpeculator):

145 )

146 continue

147

127 - target_idx = candidates[-1]

128 - target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn"

148 + target_layer_name = candidates[-1]

149 attn.kv_sharing_target_layer_name = target_layer_name

150 +

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

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

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

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

155 + # calibrated FP8 cache.

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

157 + target_layer_name

158 + ]

159 + _copy_target_kv_scales(attn, target_attn)

160 +

161 logger.info(

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

163 draft_idx,

vllm/vllm_flash_attn/flash_attn_interface.py

+20 -3

@@ -404,6 +404,22 @@ def flash_attn_varlen_func(

404

405 from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd

406

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

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

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

410 + fa4_fp8_kv_dequant = (

411 + k.dtype == torch.float8_e4m3fn

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

413 + )

414 + if fa4_fp8_kv_dequant:

415 + fa4_q_descale = None

416 + fa4_k_descale = k_descale

417 + fa4_v_descale = v_descale

418 + else:

419 + fa4_q_descale = None

420 + fa4_k_descale = None

421 + fa4_v_descale = None

422 +

423 out, softmax_lse, _, _ = _flash_attn_fwd(

424 q,

425 k,

@@ -428,10 +444,11 @@ def flash_attn_varlen_func(

444 block_sparse_tensors=block_sparse_tensors,

445 aux_tensors=aux_tensors,

446 aux_tensor_leading_dims=aux_tensor_leading_dims,

431 - q_descale=q_descale,

432 - k_descale=k_descale,

433 - v_descale=v_descale,

447 + q_descale=fa4_q_descale,

448 + k_descale=fa4_k_descale,

449 + v_descale=fa4_v_descale,

450 output_scale=output_scale,

451 + fp8_kv_dequant=fa4_fp8_kv_dequant,

452 )

453 else:

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