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
[backends][short_conv] CUDA graph piecewise edits (#24215)
Signed-off-by: Paul Pak <paulpak58@gmail.com>
Paul Pak committed
Oct 3, 2025 at 21:59 UTC
5f42fc53b6448e79e63559aefe3481909dc971f2
2 files changed
+21
-21
vllm/model_executor/layers/mamba/short_conv.py
+1
-1
@@ -115,7 +115,7 @@ class ShortConv(MambaBase, CustomOp):
115
self_kv_cache = self.kv_cache[forward_context.virtual_engine]
116
conv_state = self_kv_cache[0].transpose(-1, -2)
117
state_indices_tensor = attn_metadata.state_indices_tensor
118
- has_initial_states_p = attn_metadata.has_initial_states
118
+ has_initial_states_p = attn_metadata.has_initial_states_p
119
120
BCx, _ = self.in_proj(hidden_states)
121
vllm/v1/attention/backends/short_conv_attn.py
+20
-20
@@ -6,12 +6,12 @@ from typing import Optional
6
import torch
7
8
from vllm.attention.backends.abstract import AttentionBackend
9
-from vllm.config import VllmConfig
10
-from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
9
+from vllm.v1.attention.backends.mamba_attn import (
10
+ BaseMambaAttentionMetadataBuilder)
11
+from vllm.v1.attention.backends.utils import (PAD_SLOT_ID,
12
CommonAttentionMetadata,
13
compute_causal_conv1d_metadata,
14
split_decodes_and_prefills)
14
-from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
15
16
17
class ShortConvAttentionBackend(AttentionBackend):
@@ -29,8 +29,8 @@ class ShortConvAttentionMetadata:
29
num_decode_tokens: int
30
31
query_start_loc: torch.Tensor
32
- has_initial_states: torch.Tensor
33
- state_indices_tensor: torch.Tensor # shape: [batch,]
32
+ state_indices_tensor: torch.Tensor
33
+ has_initial_states_p: Optional[torch.Tensor]
34
35
# For causal_conv1d
36
nums_dict: Optional[dict] = None
@@ -39,14 +39,7 @@ class ShortConvAttentionMetadata:
39
40
41
class ShortConvAttentionMetadataBuilder(
42
- AttentionMetadataBuilder[ShortConvAttentionMetadata]):
43
-
44
- reorder_batch_threshold: int = 1
45
-
46
- def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
47
- vllm_config: VllmConfig, device: torch.device):
48
- super().__init__(kv_cache_spec, layer_names, vllm_config, device)
49
- assert isinstance(kv_cache_spec, MambaSpec)
42
+ BaseMambaAttentionMetadataBuilder[ShortConvAttentionMetadata]):
43
44
def build(self,
45
common_prefix_len: int,
@@ -54,7 +47,6 @@ class ShortConvAttentionMetadataBuilder(
47
fast_build: bool = False) -> ShortConvAttentionMetadata:
48
num_reqs = common_attn_metadata.num_reqs
49
query_start_loc = common_attn_metadata.query_start_loc
57
-
50
state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
51
52
# for causal_conv1d
@@ -64,13 +56,13 @@ class ShortConvAttentionMetadataBuilder(
56
split_decodes_and_prefills(
57
common_attn_metadata,
58
decode_threshold=self.reorder_batch_threshold))
67
- has_initial_states = None
59
+
60
+ has_initial_states_p = None
61
if num_prefills > 0:
69
- #[batch,]
62
has_initial_states_cpu = (
63
common_attn_metadata.
64
num_computed_tokens_cpu[num_reqs - num_prefills:num_reqs] > 0)
73
- has_initial_states = has_initial_states_cpu.to(
65
+ has_initial_states_p = has_initial_states_cpu.to(
66
query_start_loc.device)
67
68
query_start_loc_p = common_attn_metadata.query_start_loc[
@@ -79,14 +71,22 @@ class ShortConvAttentionMetadataBuilder(
71
nums_dict, batch_ptr, token_chunk_offset_ptr = \
72
compute_causal_conv1d_metadata(query_start_loc_p)
73
74
+ elif (num_decodes > 0 and num_decodes <= self.decode_cudagraph_max_bs
75
+ and self.compilation_config.full_cuda_graph):
76
+ num_input_tokens = self.vllm_config.pad_for_cudagraph(num_decodes)
77
+ self.state_indices_tensor[:num_decodes].copy_(state_indices_tensor,
78
+ non_blocking=True)
79
+ state_indices_tensor = self.state_indices_tensor[:num_input_tokens]
80
+ state_indices_tensor[num_decodes:] = PAD_SLOT_ID
81
+
82
attn_metadata = ShortConvAttentionMetadata(
83
+ query_start_loc=query_start_loc,
84
+ state_indices_tensor=state_indices_tensor,
85
+ has_initial_states_p=has_initial_states_p,
86
num_prefills=num_prefills,
87
num_prefill_tokens=num_prefill_tokens,
88
num_decodes=num_decodes,
89
num_decode_tokens=num_decode_tokens,
87
- query_start_loc=query_start_loc,
88
- has_initial_states=has_initial_states,
89
- state_indices_tensor=state_indices_tensor,
90
nums_dict=nums_dict,
91
batch_ptr=batch_ptr,
92
token_chunk_offset_ptr=token_chunk_offset_ptr,