@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1d93f11675
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 12, 04:45 UTC

[Attention][CUDAGraph] Remove CG padding from attention backends (#29352)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed Dec 2, 2025 at 13:48 UTC 1d93f116754f6e81acb9287ebcca0d1d1170a944
5 files changed +20 -46
vllm/model_executor/layers/mamba/mamba_mixer.py
+8 -9
@@ -252,7 +252,6 @@ class MambaMixer(MambaBase, CustomOp):
252 conv_state = self_kv_cache[0].transpose(-1, -2)
253 ssm_state = self_kv_cache[1]
254 has_initial_states_p = attn_metadata.has_initial_states_p
255 - num_padded_decodes = attn_metadata.num_padded_decodes
255
256 # 1. Gated MLP's linear projection
257 projected_states = self.in_proj(hidden_states)[0].transpose(-2, -1)
@@ -281,7 +280,7 @@ class MambaMixer(MambaBase, CustomOp):
280 state_indices_tensor,
281 num_prefill_tokens,
282 num_prefills,
284 - num_padded_decodes,
283 + num_decode_tokens,
284 )
285 hidden_states_BC_p = prefill_decode_split.hidden_states_BC_p
286 hidden_states_BC_d = prefill_decode_split.hidden_states_BC_d
@@ -470,24 +469,24 @@ def split_batch_to_prefill_and_decode(
469 state_indices_tensor: torch.Tensor,
470 num_prefill_tokens: int,
471 num_prefills: int,
473 - num_padded_decodes: int,
472 + num_decode_tokens: int,
473 ) -> PrefillDecodeSplit:
475 - num_actual_tokens = num_prefill_tokens + num_padded_decodes
474 + num_actual_tokens = num_prefill_tokens + num_decode_tokens
475
476 # In v1, decode tokens come first, then prefill tokens.
477 hidden_states_BC_d, hidden_states_BC_p = torch.split(
478 hidden_states_BC[..., :num_actual_tokens],
480 - [num_padded_decodes, num_prefill_tokens],
479 + [num_decode_tokens, num_prefill_tokens],
480 dim=-1,
481 )
482 gate_d, gate_p = torch.split(
484 - gate[..., :num_actual_tokens], [num_padded_decodes, num_prefill_tokens], dim=-1
483 + gate[..., :num_actual_tokens], [num_decode_tokens, num_prefill_tokens], dim=-1
484 )
485
487 - # num_padded_decodes accounts for CUDA graph padding when applicable
486 + # num_decode_tokens accounts for CUDA graph padding when applicable
487 state_indices_tensor_d, state_indices_tensor_p = torch.split(
489 - state_indices_tensor[: num_padded_decodes + num_prefills],
490 - [num_padded_decodes, num_prefills],
488 + state_indices_tensor[: num_decode_tokens + num_prefills],
489 + [num_decode_tokens, num_prefills],
490 dim=0,
491 )
492
vllm/v1/attention/backends/gdn_attn.py
+5 -17
@@ -254,17 +254,11 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
254 )
255 else:
256 has_initial_state = None
257 - num_actual_tokens = (
258 - num_prefill_tokens + num_decode_tokens + num_spec_decode_tokens
259 - )
257
261 - # prepare tensors for cudagraph
262 - #
263 - # With speculative decoding, the xgrammar backend may rollback tokens
264 - # and causing some sequences has less draft tokens than self.num_spec.
265 - #
266 - # In above cases, the max possible batch size for n tokens, can be
267 - # min(n, cudagraph_max_bs).
258 + # Prepare tensors for cudagraph
259 + # Note: m.num_actual_tokens is already padded by the model runner for CUDAGraph
260 + batch_size = m.num_actual_tokens
261 +
262 if (
263 self.use_full_cuda_graph
264 and num_prefills == 0
@@ -272,9 +266,6 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
266 and num_spec_decodes <= self.decode_cudagraph_max_bs
267 and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
268 ):
275 - num_actual_tokens = self.vllm_config.pad_for_cudagraph(m.num_actual_tokens)
276 - batch_size = min(self.decode_cudagraph_max_bs, num_actual_tokens)
277 -
269 self.spec_state_indices_tensor[:num_spec_decodes].copy_(
270 spec_state_indices_tensor, non_blocking=True
271 )
@@ -319,9 +310,6 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
310 and num_spec_decodes == 0
311 and num_decodes <= self.decode_cudagraph_max_bs
312 ):
322 - num_actual_tokens = self.vllm_config.pad_for_cudagraph(m.num_actual_tokens)
323 - batch_size = num_actual_tokens
324 -
313 self.non_spec_state_indices_tensor[:num_decodes].copy_(
314 non_spec_state_indices_tensor, non_blocking=True
315 )
@@ -344,7 +332,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
332 num_decode_tokens=num_decode_tokens,
333 num_spec_decodes=num_spec_decodes,
334 num_spec_decode_tokens=num_spec_decode_tokens,
347 - num_actual_tokens=num_actual_tokens,
335 + num_actual_tokens=m.num_actual_tokens,
336 has_initial_state=has_initial_state,
337 spec_query_start_loc=spec_query_start_loc,
338 non_spec_query_start_loc=non_spec_query_start_loc,
vllm/v1/attention/backends/mamba1_attn.py
+3 -9
@@ -31,7 +31,6 @@ class Mamba1AttentionMetadata:
31 num_prefill_tokens: int
32 num_decodes: int
33 num_decode_tokens: int
34 - num_padded_decodes: int
34
35 block_idx_last_scheduled_token: torch.Tensor # shape: [batch,]
36 block_idx_first_scheduled_token_p: torch.Tensor # shape: [batch,]
@@ -68,7 +67,6 @@ class Mamba1AttentionMetadataBuilder(
67
68 has_initial_states_p = None
69 query_start_loc_p = None
71 - padded_decodes = num_decodes
70 num_computed_tokens, num_computed_tokens_p = None, None
71 block_idx_first_scheduled_token = None
72 block_idx_first_scheduled_token_p = None
@@ -125,11 +123,10 @@ class Mamba1AttentionMetadataBuilder(
123 and num_decodes <= self.decode_cudagraph_max_bs
124 and self.compilation_config.cudagraph_mode.has_full_cudagraphs()
125 ):
128 - padded_decodes = self.vllm_config.pad_for_cudagraph(num_decodes)
126 self.state_indices_tensor[:num_decodes].copy_(
127 state_indices_tensor, non_blocking=True
128 )
132 - state_indices_tensor = self.state_indices_tensor[:padded_decodes]
129 + state_indices_tensor = self.state_indices_tensor[:num_decode_tokens]
130 state_indices_tensor[num_decodes:] = PAD_SLOT_ID
131
132 if self.vllm_config.cache_config.enable_prefix_caching:
@@ -137,17 +134,15 @@ class Mamba1AttentionMetadataBuilder(
134 block_idx_last_scheduled_token, non_blocking=True
135 )
136 block_idx_last_scheduled_token = self.block_idx_last_scheduled_token[
140 - :padded_decodes
137 + :num_decode_tokens
138 ]
142 - block_idx_last_scheduled_token[num_decodes:] = 0
139
140 self.block_idx_last_computed_token[:num_decodes].copy_(
141 block_idx_last_computed_token, non_blocking=True
142 )
143 block_idx_last_computed_token = self.block_idx_last_computed_token[
148 - :padded_decodes
144 + :num_decode_tokens
145 ]
150 - block_idx_last_computed_token[num_decodes:] = 0
146
147 return Mamba1AttentionMetadata(
148 query_start_loc_p=query_start_loc_p,
@@ -157,7 +152,6 @@ class Mamba1AttentionMetadataBuilder(
152 num_prefill_tokens=num_prefill_tokens,
153 num_decodes=num_decodes,
154 num_decode_tokens=num_decode_tokens,
160 - num_padded_decodes=padded_decodes,
155 block_idx_last_scheduled_token=block_idx_last_scheduled_token,
156 block_idx_first_scheduled_token_p=block_idx_first_scheduled_token_p,
157 block_idx_last_computed_token=block_idx_last_computed_token,
vllm/v1/attention/backends/mamba2_attn.py
+3 -9
@@ -10,7 +10,6 @@ from vllm.config import VllmConfig
10 from vllm.utils.math_utils import cdiv
11 from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadataBuilder
12 from vllm.v1.attention.backends.utils import (
13 - PAD_SLOT_ID,
13 CommonAttentionMetadata,
14 compute_causal_conv1d_metadata,
15 split_decodes_and_prefills,
@@ -304,30 +303,25 @@ class Mamba2AttentionMetadataBuilder(
303 num_decodes <= self.decode_cudagraph_max_bs
304 and self.compilation_config.cudagraph_mode.has_full_cudagraphs()
305 ):
307 - # Pad state tensor for CUDA graph
308 - num_input_tokens = self.vllm_config.pad_for_cudagraph(num_decodes)
306 self.state_indices_tensor[:num_decodes].copy_(
307 state_indices_tensor, non_blocking=True
308 )
312 - state_indices_tensor = self.state_indices_tensor[:num_input_tokens]
313 - state_indices_tensor[num_decodes:] = PAD_SLOT_ID
309 + state_indices_tensor = self.state_indices_tensor[:num_decode_tokens]
310
311 if self.vllm_config.cache_config.enable_prefix_caching:
312 self.block_idx_last_scheduled_token[:num_decodes].copy_(
313 block_idx_last_scheduled_token, non_blocking=True
314 )
315 block_idx_last_scheduled_token = self.block_idx_last_scheduled_token[
320 - :num_input_tokens
316 + :num_decode_tokens
317 ]
322 - block_idx_last_scheduled_token[num_decodes:] = 0
318
319 self.block_idx_last_computed_token[:num_decodes].copy_(
320 block_idx_last_computed_token, non_blocking=True
321 )
322 block_idx_last_computed_token = self.block_idx_last_computed_token[
328 - :num_input_tokens
323 + :num_decode_tokens
324 ]
330 - block_idx_last_computed_token[num_decodes:] = 0
325
326 attn_metadata = Mamba2AttentionMetadata(
327 num_prefills=num_prefills,
vllm/v1/attention/backends/short_conv_attn.py
+1 -2
@@ -83,11 +83,10 @@ class ShortConvAttentionMetadataBuilder(
83 and num_decodes <= self.decode_cudagraph_max_bs
84 and self.compilation_config.cudagraph_mode.has_full_cudagraphs()
85 ):
86 - num_input_tokens = self.vllm_config.pad_for_cudagraph(num_decodes)
86 self.state_indices_tensor[:num_decodes].copy_(
87 state_indices_tensor, non_blocking=True
88 )
90 - state_indices_tensor = self.state_indices_tensor[:num_input_tokens]
89 + state_indices_tensor = self.state_indices_tensor[:num_decode_tokens]
90 state_indices_tensor[num_decodes:] = PAD_SLOT_ID
91
92 attn_metadata = ShortConvAttentionMetadata(