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
[PERF] Add `conv1d` metadata to GDN attn (#25105)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Vadim Gimpelson committed
Sep 18, 2025 at 18:27 UTC
072d7e53e534d337b41262dd44ded9b44aa699ef
5 files changed
+24
-8
vllm/model_executor/layers/mamba/mamba2_metadata.py
+5
-3
@@ -11,6 +11,7 @@ from vllm.attention.backends.placeholder_attn import (
11
PlaceholderAttentionMetadata)
12
from vllm.attention.backends.utils import PAD_SLOT_ID
13
from vllm.platforms import current_platform
14
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
15
from vllm.v1.attention.backends.mamba2_attn import (
16
Mamba2AttentionMetadata, _query_start_loc_to_chunk_indices_offsets)
17
@@ -45,8 +46,8 @@ class Mamba2Metadata:
46
"""
47
nums_dict: Optional[dict] = None
48
cu_seqlen: Optional[int] = None
48
- batch_ptr: Optional[torch.tensor] = None
49
- token_chunk_offset_ptr: Optional[torch.tensor] = None
49
+ batch_ptr: Optional[torch.Tensor] = None
50
+ token_chunk_offset_ptr: Optional[torch.Tensor] = None
51
52
53
def get_platform_metadata_classes() -> tuple[type[AttentionMetadata], ...]:
@@ -117,7 +118,8 @@ def prepare_mamba2_metadata(
118
119
def update_metadata(x: torch.Tensor, query_start_loc: torch.Tensor,
120
mamba2_metadata: Union[Mamba2Metadata,
120
- Mamba2AttentionMetadata]):
121
+ Mamba2AttentionMetadata,
122
+ GDNAttentionMetadata]):
123
"""
124
this is triggered upon handling a new input at the first layer
125
"""
vllm/model_executor/models/qwen3_next.py
+9
-1
@@ -35,6 +35,7 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
35
RowParallelLinear)
36
from vllm.model_executor.layers.logits_processor import LogitsProcessor
37
from vllm.model_executor.layers.mamba.abstract import MambaBase
38
+from vllm.model_executor.layers.mamba.mamba2_metadata import update_metadata
39
from vllm.model_executor.layers.mamba.mamba_mixer2 import (
40
mamba_v2_sharded_weight_loader)
41
from vllm.model_executor.layers.mamba.mamba_utils import (
@@ -414,6 +415,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
415
416
assert isinstance(attn_metadata, dict)
417
attn_metadata = attn_metadata[self.prefix]
418
+ conv_metadata = attn_metadata
419
assert isinstance(attn_metadata, GDNAttentionMetadata)
420
has_initial_state = attn_metadata.has_initial_state
421
spec_query_start_loc = attn_metadata.spec_query_start_loc
@@ -475,10 +477,15 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
477
478
# 2.2: process the remaining part
479
if attn_metadata.num_prefills > 0:
480
+ mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1)
481
+ if conv_metadata.cu_seqlen is None:
482
+ conv_metadata = update_metadata(mixed_qkv_non_spec_T,
483
+ non_spec_query_start_loc,
484
+ conv_metadata)
485
# - "cache_indices" updates the conv_state cache in positions
486
# pointed to by "mamba_cache_params.state_indices_tensor"
487
mixed_qkv_non_spec = causal_conv1d_fn(
481
- mixed_qkv_non_spec.transpose(0, 1),
488
+ mixed_qkv_non_spec_T,
489
conv_weights,
490
self.conv1d.bias,
491
activation=self.activation,
@@ -486,6 +493,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
493
has_initial_state=has_initial_state,
494
cache_indices=non_spec_state_indices_tensor,
495
query_start_loc=non_spec_query_start_loc,
496
+ metadata=conv_metadata,
497
).transpose(0, 1)
498
elif attn_metadata.num_decodes > 0:
499
mixed_qkv_non_spec = causal_conv1d_update(
vllm/v1/attention/backends/gdn_attn.py
+6
@@ -50,6 +50,12 @@ class GDNAttentionMetadata:
50
Tensor] = None # shape: [num_prefill_tokens + num_decode_tokens,]
51
num_accepted_tokens: Optional[torch.Tensor] = None # shape: [batch,]
52
53
+ # The following attributes are for triton implementation of causal_conv1d
54
+ nums_dict: Optional[dict] = None
55
+ cu_seqlen: Optional[int] = None
56
+ batch_ptr: Optional[torch.Tensor] = None
57
+ token_chunk_offset_ptr: Optional[torch.Tensor] = None
58
+
59
60
class GDNAttentionMetadataBuilder(
61
AttentionMetadataBuilder[GDNAttentionMetadata]):
vllm/v1/attention/backends/mamba2_attn.py
+2
-2
@@ -132,8 +132,8 @@ class Mamba2AttentionMetadata:
132
# The following attributes are for triton implementation of causal_conv1d
133
nums_dict: Optional[dict] = None
134
cu_seqlen: Optional[int] = None
135
- batch_ptr: Optional[torch.tensor] = None
136
- token_chunk_offset_ptr: Optional[torch.tensor] = None
135
+ batch_ptr: Optional[torch.Tensor] = None
136
+ token_chunk_offset_ptr: Optional[torch.Tensor] = None
137
138
139
class Mamba2AttentionMetadataBuilder(
vllm/v1/attention/backends/short_conv_attn.py
+2
-2
@@ -34,8 +34,8 @@ class ShortConvAttentionMetadata:
34
# For causal_conv1d
35
nums_dict: Optional[dict] = None
36
cu_seqlen: Optional[int] = None
37
- batch_ptr: Optional[torch.tensor] = None
38
- token_chunk_offset_ptr: Optional[torch.tensor] = None
37
+ batch_ptr: Optional[torch.Tensor] = None
38
+ token_chunk_offset_ptr: Optional[torch.Tensor] = None
39
40
41
class ShortConvAttentionMetadataBuilder(