@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 620db1fc58

Fork 0

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 14, 05:30 UTC

[Attention] FlashAttention MLA cudagraph support (#23958)

Signed-off-by: Matthew Bonanni <mbonanni001@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

Matthew Bonanni committed Sep 8, 2025 at 15:05 UTC 620db1fc5879329dae62367238f97d446904a3ee

7 files changed +118 -29

tests/compile/piecewise/test_full_cudagraph.py

+11 -1

@@ -61,6 +61,16 @@ backend_configs = {

61 "cudagraph_mode": "FULL_AND_PIECEWISE",

62 },

63 specific_gpu_arch=(9, 0)),

64 + # FlashAttention MLA on Hopper

65 + "FlashAttentionMLA":

66 + BackendConfig(name="FlashAttentionMLA",

67 + env_vars={

68 + "VLLM_ATTENTION_BACKEND": "FLASH_ATTN_MLA",

69 + },

70 + comp_config={

71 + "cudagraph_mode": "FULL_DECODE_ONLY",

72 + },

73 + specific_gpu_arch=(9, 0)),

74 # Cutlass MLA on Blackwell

75 "CutlassMLA":

76 BackendConfig(

@@ -102,7 +112,7 @@ backend_configs = {

112 test_params_full_cudagraph = []

113

114 # deepseek-ai/DeepSeek-V2-Lite with MLA

105 -MLA_backends = ["FlashMLA", "CutlassMLA"]

115 +MLA_backends = ["FlashMLA", "FlashAttentionMLA", "CutlassMLA"]

116 for mla_backend in MLA_backends:

117 test_params_full_cudagraph.append(

118 pytest.param(

tests/v1/attention/test_mla_backends.py

-5

@@ -73,7 +73,6 @@ def create_and_prepopulate_kv_cache(

73 kv_c_contexts: list[torch.Tensor],

74 k_pe_contexts: list[torch.Tensor],

75 block_size: int,

76 - num_kv_heads: int,

76 head_size: int,

77 dtype: torch.dtype,

78 device: torch.device,

@@ -87,7 +86,6 @@ def create_and_prepopulate_kv_cache(

86 k_pe_contexts: List of key positional embedding context tensors

87 for each sequence

88 block_size: Size of each block

90 - num_kv_heads: Number of KV heads (should be 1 for MLA)

89 head_size: Size of each head (latent dimension)

90 dtype: Data type for the cache

91 device: Device to create the cache on

@@ -285,8 +283,6 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):

283 query_lens = batch_spec.query_lens

284 num_q_heads = vllm_config.model_config.get_num_attention_heads(

285 vllm_config.parallel_config)

288 - num_kv_heads = vllm_config.model_config.get_num_kv_heads(

289 - vllm_config.parallel_config)

286 head_size = vllm_config.model_config.get_head_size()

287 dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)

288 block_size = vllm_config.cache_config.block_size

@@ -476,7 +472,6 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):

472 kv_c_contexts=kv_c_contexts,

473 k_pe_contexts=k_pe_contexts,

474 block_size=block_size,

479 - num_kv_heads=num_kv_heads,

475 head_size=head_size,

476 dtype=dtype,

477 device=device,

tests/v1/cudagraph/test_cudagraph_mode.py

+10

@@ -62,6 +62,16 @@ backend_configs = {

62 "cudagraph_mode": "FULL_AND_PIECEWISE",

63 },

64 specific_gpu_arch=(9, 0)),

65 + # FlashAttention MLA on Hopper

66 + "FlashAttentionMLA":

67 + BackendConfig(name="FlashAttentionMLA",

68 + env_vars={

69 + "VLLM_ATTENTION_BACKEND": "FLASH_ATTN_MLA",

70 + },

71 + comp_config={

72 + "cudagraph_mode": "FULL_DECODE_ONLY",

73 + },

74 + specific_gpu_arch=(9, 0)),

75 # FA2

76 "FA2":

77 BackendConfig(name="FA2",

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

+14 -8

@@ -443,11 +443,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

443 self.metadata_cls = metadata_cls \

444 if metadata_cls is not None else MLACommonMetadata

445 self.kv_cache_spec = kv_cache_spec

446 - self.device = device

446 scheduler_config = vllm_config.scheduler_config

447 self.model_config = vllm_config.model_config

449 - cache_config = vllm_config.cache_config

448 parallel_config = vllm_config.parallel_config

449 + cache_config = vllm_config.cache_config

450 + self.compilation_config = vllm_config.compilation_config

451 + self.device = device

452 +

453 self.num_heads = self.model_config.get_num_attention_heads(

454 parallel_config)

455 self.mla_dims = get_mla_dims(self.model_config)

@@ -608,10 +610,12 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

610 prefill.prefill_main = self._fi_prefill_main

611 prefill.prefill_chunks = self._fi_prefill_chunks

612

611 - def _build_decode(

612 - self, block_table_tensor: torch.Tensor, seq_lens_cpu: torch.Tensor,

613 - seq_lens_device: torch.Tensor, query_start_loc_cpu: torch.Tensor,

614 - query_start_loc_device: torch.Tensor) -> MLACommonDecodeMetadata:

613 + def _build_decode(self, block_table_tensor: torch.Tensor,

614 + seq_lens_cpu: torch.Tensor,

615 + seq_lens_device: torch.Tensor,

616 + query_start_loc_cpu: torch.Tensor,

617 + query_start_loc_device: torch.Tensor,

618 + num_decode_tokens: int) -> MLACommonDecodeMetadata:

619 return MLACommonDecodeMetadata(

620 block_table=block_table_tensor,

621 seq_lens=seq_lens_device,

@@ -624,11 +628,12 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

628 Currently, only decode is supported for full cudagraphs with MLA.

629 """

630 m = common_attn_metadata

627 - assert m.num_reqs == m.num_actual_tokens, \

631 + assert m.num_reqs <= (m.num_actual_tokens *

632 + self.reorder_batch_threshold), \

633 "MLA only supports decode-only full CUDAGraph capture. " \

634 "Make sure all cudagraph capture sizes <= max_num_seq."

635

631 - assert m.max_query_len == 1 # decode-only

636 + assert m.max_query_len <= self.reorder_batch_threshold # decode only

637

638 return self.build(0, m)

639

@@ -819,6 +824,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

824 seq_lens_device=seq_lens[:num_decodes],

825 query_start_loc_cpu=query_start_loc_cpu[:num_decodes + 1],

826 query_start_loc_device=query_start_loc[:num_decodes + 1],

827 + num_decode_tokens=num_decode_tokens,

828 )

829

830 attn_metadata = self.metadata_cls(

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

+71 -6

@@ -17,11 +17,16 @@ from vllm.v1.attention.backends.mla.common import (MLACommonBackend,

17 MLACommonImpl,

18 MLACommonMetadata,

19 MLACommonMetadataBuilder)

20 +from vllm.v1.attention.backends.utils import AttentionCGSupport

21 from vllm.v1.kv_cache_interface import AttentionSpec

22 from vllm.vllm_flash_attn import flash_attn_varlen_func, get_scheduler_metadata

24 logger = init_logger(__name__)

26 +# NOTE(matt): This is an arbitrary number, copied from

27 +# woosuk's implementation in standard FlashAttention backend

28 +_DEFAULT_MAX_NUM_SPLITS_FOR_CUDA_GRAPH = 16

29 +

31 class FlashAttnMLABackend(MLACommonBackend):

@@ -48,6 +53,7 @@ class FlashAttnMLADecodeMetadata(MLACommonDecodeMetadata):

53 max_query_len: int

54 max_seq_len: int

55 scheduler_metadata: Optional[torch.Tensor] = None

56 + max_num_splits: int = 0

59 @dataclass

@@ -57,14 +63,41 @@ class FlashAttnMLAMetadata(MLACommonMetadata[FlashAttnMLADecodeMetadata]):

64 class FlashAttnMLAMetadataBuilder(

65 MLACommonMetadataBuilder[FlashAttnMLAMetadata]):

66 + cudagraph_support: ClassVar[AttentionCGSupport] = \

67 + AttentionCGSupport.UNIFORM_BATCH

68 +

69 reorder_batch_threshold: ClassVar[int] = 512

71 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

72 vllm_config: VllmConfig, device: torch.device):

73 super().__init__(kv_cache_spec, layer_names, vllm_config, device,

74 FlashAttnMLAMetadata)

75 + self.max_num_splits = 0 # No upper bound on the number of splits.

76 self.fa_aot_schedule = (get_flash_attn_version() == 3)

78 + self.use_full_cuda_graph = \

79 + self.compilation_config.cudagraph_mode.has_full_cudagraphs()

80 +

81 + if self.use_full_cuda_graph and self.fa_aot_schedule:

82 + self.max_cudagraph_size = self.compilation_config.max_capture_size

83 +

84 + if self.max_cudagraph_size > 992:

85 + # This condition derives from FA3's internal heuristic.

86 + # TODO(woosuk): Support larger cudagraph sizes.

87 + raise ValueError(

88 + "Capture size larger than 992 is not supported for "

89 + "full cuda graph.")

90 +

91 + self.scheduler_metadata = torch.zeros(

92 + vllm_config.scheduler_config.max_num_seqs + 1,

93 + dtype=torch.int32,

94 + device=self.device,

95 + )

96 + # When using cuda graph, we need to set the upper bound of the

97 + # number of splits so that large enough intermediate buffers are

98 + # pre-allocated during capture.

99 + self.max_num_splits = _DEFAULT_MAX_NUM_SPLITS_FOR_CUDA_GRAPH

100 +

101 def _schedule_decode(self, num_reqs, cu_query_lens, max_query_len, seqlens,

102 max_seq_len, causal):

103 if self.fa_aot_schedule:

@@ -81,14 +114,16 @@ class FlashAttnMLAMetadataBuilder(

114 page_size=self.page_size,

115 cu_seqlens_q=cu_query_lens,

116 causal=causal,

117 + num_splits=self.max_num_splits,

118 )

119 return None

120

87 - def _build_decode(

88 - self, block_table_tensor: torch.Tensor, seq_lens_cpu: torch.Tensor,

89 - seq_lens_device: torch.Tensor, query_start_loc_cpu: torch.Tensor,

90 - query_start_loc_device: torch.Tensor

91 - ) -> FlashAttnMLADecodeMetadata:

121 + def _build_decode(self, block_table_tensor: torch.Tensor,

122 + seq_lens_cpu: torch.Tensor,

123 + seq_lens_device: torch.Tensor,

124 + query_start_loc_cpu: torch.Tensor,

125 + query_start_loc_device: torch.Tensor,

126 + num_decode_tokens: int) -> FlashAttnMLADecodeMetadata:

127 query_lens_cpu = (query_start_loc_cpu[1:] - query_start_loc_cpu[:-1])

128 max_query_len = query_lens_cpu.max().item()

129 max_seq_len = seq_lens_cpu.max().item()

@@ -102,6 +137,29 @@ class FlashAttnMLAMetadataBuilder(

137 causal=True,

138 )

139

140 + # For FA3 + full cudagraph

141 + max_num_splits = 0

142 + if self.use_full_cuda_graph and scheduler_metadata is not None:

143 + n = scheduler_metadata.shape[0]

144 + # Ensure the persistent buffer is large enough

145 + assert n <= self.scheduler_metadata.shape[0], \

146 + f"Scheduler metadata size {n} exceeds buffer size " + \

147 + f"{self.scheduler_metadata.shape[0]}"

148 + self.scheduler_metadata[:n] = scheduler_metadata

149 + # NOTE(woosuk): We should zero out the rest of the scheduler

150 + # metadata to guarantee the correctness. Otherwise, some thread

151 + # blocks may use the invalid scheduler metadata and overwrite the

152 + # output buffer.

153 + self.scheduler_metadata[n:] = 0

154 + scheduler_metadata = self.scheduler_metadata[:n]

155 +

156 + if num_decode_tokens <= self.max_cudagraph_size:

157 + # NOTE(woosuk): Setting num_splits > 1 may increase the memory

158 + # usage, because the intermediate buffers of size [num_splits,

159 + # num_heads, num_tokens, head_size] are allocated. Therefore,

160 + # we only set num_splits when using cuda graphs.

161 + max_num_splits = self.max_num_splits

162 +

163 return FlashAttnMLADecodeMetadata(

164 block_table=block_table_tensor,

165 seq_lens=seq_lens_device,

@@ -109,6 +167,7 @@ class FlashAttnMLAMetadataBuilder(

167 max_query_len=max_query_len,

168 max_seq_len=max_seq_len,

169 scheduler_metadata=scheduler_metadata,

170 + max_num_splits=max_num_splits,

171 )

172

173

@@ -175,12 +234,17 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):

234 kv_c_cache = kv_c_and_k_pe_cache[..., :self.kv_lora_rank]

235 k_pe_cache = kv_c_and_k_pe_cache[..., self.kv_lora_rank:]

236

237 + # NOTE(matt): During CUDA graph capture, max_query_len can be 0, but the

238 + # kernel uses this to calculate grid dimensions. Ensure it's at least 1

239 + # to prevent invalid grid configuration during graph capture.

240 + max_seqlen_q = max(attn_metadata.decode.max_query_len, 1)

241 +

242 o = flash_attn_varlen_func(

243 q=q_pe,

244 k=k_pe_cache.unsqueeze(-2), # Add head dim of 1

245 v=kv_c_cache.unsqueeze(-2), # Add head dim of 1

246 q_v=q_nope,

183 - max_seqlen_q=attn_metadata.decode.max_query_len,

247 + max_seqlen_q=max_seqlen_q,

248 cu_seqlens_q=attn_metadata.decode.query_start_loc,

249 max_seqlen_k=attn_metadata.decode.max_seq_len,

250 seqused_k=attn_metadata.decode.seq_lens,

@@ -189,6 +253,7 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):

253 causal=True,

254 fa_version=3, # only version 3 is supported

255 scheduler_metadata=attn_metadata.decode.scheduler_metadata,

256 + num_splits=attn_metadata.decode.max_num_splits,

257 )

258

259 return self._v_up_proj(o)

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

+6 -5

@@ -62,7 +62,6 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

62 super().__init__(kv_cache_spec, layer_names, vllm_config, device,

63 FlashMLAMetadata)

65 - self.compilation_config = vllm_config.compilation_config

65 self.num_q_heads = vllm_config.model_config.get_num_attention_heads(

66 vllm_config.parallel_config)

@@ -85,10 +84,12 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

84 device=self.device,

85 dtype=torch.int32)

88 - def _build_decode(

89 - self, block_table_tensor: torch.Tensor, seq_lens_cpu: torch.Tensor,

90 - seq_lens_device: torch.Tensor, query_start_loc_cpu: torch.Tensor,

91 - query_start_loc_device: torch.Tensor) -> FlashMLADecodeMetadata:

87 + def _build_decode(self, block_table_tensor: torch.Tensor,

88 + seq_lens_cpu: torch.Tensor,

89 + seq_lens_device: torch.Tensor,

90 + query_start_loc_cpu: torch.Tensor,

91 + query_start_loc_device: torch.Tensor,

92 + num_decode_tokens: int) -> FlashMLADecodeMetadata:

93 tile_scheduler_metadata, num_splits = \

94 get_mla_metadata(

95 seq_lens_device,

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

+6 -4

@@ -104,10 +104,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

104 dtype=torch.int32,

105 device=device)

106

107 - def _build_decode(

108 - self, block_table_tensor: torch.Tensor, seq_lens_cpu: torch.Tensor,

109 - seq_lens_device: torch.Tensor, query_start_loc_cpu: torch.Tensor,

110 - query_start_loc_device: torch.Tensor) -> AiterMLADecodeMetadata:

107 + def _build_decode(self, block_table_tensor: torch.Tensor,

108 + seq_lens_cpu: torch.Tensor,

109 + seq_lens_device: torch.Tensor,

110 + query_start_loc_cpu: torch.Tensor,

111 + query_start_loc_device: torch.Tensor,

112 + num_decode_tokens: int) -> AiterMLADecodeMetadata:

113 page_size = self.kv_cache_spec.block_size

114 block_table_bounds = (seq_lens_device + page_size - 1) // page_size

115 device = self.device