@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7e1b45a092
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 3, 01:45 UTC

[Attention] Mamba attention module refactor (#41126)

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>

wangxiyuan committed May 22, 2026 at 17:13 UTC 7e1b45a09252a5b513cd83116aa7a2f310220c34
10 files changed +765 -774
vllm/config/compilation.py
+1 -1
@@ -752,7 +752,7 @@ class CompilationConfig:
752 "vllm::short_conv",
753 "vllm::linear_attention",
754 "vllm::plamo2_mamba_mixer",
755 - "vllm::gdn_attention_core",
755 + "vllm::qwen_gdn_attention_core",
756 "vllm::gdn_attention_core_xpu",
757 "vllm::olmo_hybrid_gdn_full_forward",
758 "vllm::kda_attention",
vllm/model_executor/layers/mamba/gdn/__init__.py
vllm/model_executor/layers/mamba/gdn/base.py new
+58
@@ -0,0 +1,58 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +from transformers import PretrainedConfig
5 +
6 +from vllm.config import (
7 + VllmConfig,
8 +)
9 +from vllm.distributed import (
10 + get_tensor_model_parallel_rank,
11 + get_tensor_model_parallel_world_size,
12 +)
13 +from vllm.model_executor.custom_op import PluggableLayer
14 +from vllm.model_executor.layers.mamba.abstract import MambaBase
15 +from vllm.model_executor.layers.mamba.mamba_utils import (
16 + MambaStateDtypeCalculator,
17 +)
18 +from vllm.model_executor.models.utils import extract_layer_index
19 +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
20 +
21 +
22 +class GatedDeltaNetAttention(PluggableLayer, MambaBase):
23 + """Base class for GatedDeltaNet attention layer."""
24 +
25 + def __init__(
26 + self,
27 + config: PretrainedConfig,
28 + vllm_config: VllmConfig,
29 + prefix: str = "",
30 + ) -> None:
31 + super().__init__()
32 + self.prefix = prefix
33 + self.tp_size = get_tensor_model_parallel_world_size()
34 + self.tp_rank = get_tensor_model_parallel_rank()
35 + self.layer_idx = extract_layer_index(prefix)
36 + self.hidden_size = config.hidden_size
37 + self.activation = config.hidden_act
38 + self.layer_norm_epsilon = config.rms_norm_eps
39 + self.model_config = vllm_config.model_config
40 + self.cache_config = vllm_config.cache_config
41 + self.quant_config = vllm_config.quant_config
42 + self.speculative_config = vllm_config.speculative_config
43 + self.num_spec = (
44 + self.speculative_config.num_speculative_tokens
45 + if self.speculative_config
46 + else 0
47 + )
48 +
49 + @property
50 + def mamba_type(self) -> MambaAttentionBackendEnum:
51 + return MambaAttentionBackendEnum.GDN_ATTN
52 +
53 + def get_state_dtype(self) -> tuple[torch.dtype, ...]:
54 + return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
55 + self.model_config.dtype,
56 + self.cache_config.mamba_cache_dtype,
57 + self.cache_config.mamba_ssm_cache_dtype,
58 + )
vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py renamed
+26 -45
@@ -5,39 +5,37 @@ import torch
5 from einops import rearrange
6 from torch import nn
7
8 -from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
8 +from vllm.config import VllmConfig, get_current_vllm_config
9 from vllm.distributed import (
10 divide,
11 - get_tensor_model_parallel_rank,
12 - get_tensor_model_parallel_world_size,
11 )
12 from vllm.forward_context import ForwardContext, get_forward_context
13 from vllm.logger import init_logger
14 +from vllm.model_executor.custom_op import PluggableLayer
15 +from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
16 from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
17 from vllm.model_executor.utils import set_weight_attrs
18 +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
19 from vllm.utils.torch_utils import direct_register_custom_op
20 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
20 -from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
21
22 -from .fla.ops.kda import (
22 +from ...fla.ops.kda import (
23 FusedRMSNormGated,
24 chunk_kda,
25 fused_kda_gate,
26 fused_recurrent_kda,
27 )
28 -from .linear import (
28 +from ...linear import (
29 ColumnParallelLinear,
30 ReplicatedLinear,
31 RowParallelLinear,
32 )
33 -from .mamba.abstract import MambaBase
34 -from .mamba.mamba_utils import (
33 +from ..mamba_utils import (
34 MambaStateDtypeCalculator,
35 MambaStateShapeCalculator,
36 is_conv_state_dim_first,
37 )
39 -from .mamba.ops.causal_conv1d import causal_conv1d_fn, causal_conv1d_update
40 -from .quantization.base_config import QuantizationConfig
38 +from ..ops.causal_conv1d import causal_conv1d_fn, causal_conv1d_update
39
40 logger = init_logger(__name__)
41
@@ -83,11 +81,8 @@ direct_register_custom_op(
81 )
82
83
86 -class KimiDeltaAttention(nn.Module, MambaBase):
87 - @property
88 - def mamba_type(self) -> MambaAttentionBackendEnum:
89 - return MambaAttentionBackendEnum.GDN_ATTN
90 -
84 +@PluggableLayer.register("kimi_gated_delta_net_attention")
85 +class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
86 def get_state_dtype(
87 self,
88 ) -> tuple[torch.dtype, torch.dtype, torch.dtype, torch.dtype]:
@@ -106,28 +101,16 @@ class KimiDeltaAttention(nn.Module, MambaBase):
101
102 def __init__(
103 self,
109 - layer_idx: int,
110 - hidden_size: int,
111 - quant_config: QuantizationConfig | None = None,
112 - cache_config: CacheConfig | None = None,
113 - model_config: ModelConfig | None = None,
114 - rms_norm_eps: float = 1e-5,
104 + config: KimiLinearConfig,
105 + vllm_config: VllmConfig,
106 prefix: str = "",
116 - **kwargs,
107 ) -> None:
118 - super().__init__()
119 - self.tp_size = get_tensor_model_parallel_world_size()
120 - self.tp_rank = get_tensor_model_parallel_rank()
121 - self.hidden_size = hidden_size
122 - self.model_config = model_config
123 - self.cache_config = cache_config
124 - if model_config is None:
125 - raise ValueError("model_config must be provided")
126 - kda_config = model_config.linear_attn_config # type: ignore[attr-defined]
108 + super().__init__(config, vllm_config, prefix)
109 +
110 + kda_config = config.linear_attn_config # type: ignore[attr-defined]
111 + assert kda_config is not None, "linear_attn_config must be set"
112 self.head_dim = kda_config["head_dim"]
113 self.num_heads = kda_config["num_heads"]
129 - self.layer_idx = layer_idx
130 - self.prefix = prefix
114 assert self.num_heads % self.tp_size == 0
115 self.local_num_heads = divide(self.num_heads, self.tp_size)
116
@@ -138,21 +121,21 @@ class KimiDeltaAttention(nn.Module, MambaBase):
121 self.hidden_size,
122 projection_size,
123 bias=False,
141 - quant_config=quant_config,
124 + quant_config=self.quant_config,
125 prefix=f"{prefix}.q_proj",
126 )
127 self.k_proj = ColumnParallelLinear(
128 self.hidden_size,
129 projection_size,
130 bias=False,
148 - quant_config=quant_config,
131 + quant_config=self.quant_config,
132 prefix=f"{prefix}.k_proj",
133 )
134 self.v_proj = ColumnParallelLinear(
135 self.hidden_size,
136 projection_size,
137 bias=False,
155 - quant_config=quant_config,
138 + quant_config=self.quant_config,
139 prefix=f"{prefix}.v_proj",
140 )
141
@@ -160,7 +143,7 @@ class KimiDeltaAttention(nn.Module, MambaBase):
143 self.hidden_size,
144 self.head_dim,
145 bias=False,
163 - quant_config=quant_config,
146 + quant_config=self.quant_config,
147 prefix=f"{prefix}.f_a_proj",
148 )
149
@@ -168,7 +151,7 @@ class KimiDeltaAttention(nn.Module, MambaBase):
151 self.head_dim,
152 projection_size,
153 bias=False,
171 - quant_config=quant_config,
154 + quant_config=self.quant_config,
155 prefix=f"{prefix}.f_b_proj",
156 )
157 self.dt_bias = nn.Parameter(
@@ -181,7 +164,7 @@ class KimiDeltaAttention(nn.Module, MambaBase):
164 self.hidden_size,
165 self.num_heads,
166 bias=False,
184 - quant_config=quant_config,
167 + quant_config=self.quant_config,
168 prefix=f"{prefix}.b_proj",
169 )
170
@@ -223,24 +206,22 @@ class KimiDeltaAttention(nn.Module, MambaBase):
206 self.hidden_size,
207 self.head_dim,
208 bias=False,
226 - quant_config=quant_config,
209 + quant_config=self.quant_config,
210 prefix=f"{prefix}.g_a_proj",
211 )
212 self.g_b_proj = ColumnParallelLinear(
213 self.head_dim,
214 projection_size,
215 bias=False,
233 - quant_config=quant_config,
216 + quant_config=self.quant_config,
217 prefix=f"{prefix}.g_b_proj",
218 )
236 - self.o_norm = FusedRMSNormGated(
237 - self.head_dim, eps=rms_norm_eps, activation="sigmoid"
238 - )
219 + self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
220 self.o_proj = RowParallelLinear(
221 projection_size,
222 self.hidden_size,
223 bias=False,
243 - quant_config=quant_config,
224 + quant_config=self.quant_config,
225 prefix=f"{prefix}.o_proj",
226 )
227
vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py new
+634
@@ -0,0 +1,634 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +from einops import rearrange
5 +from torch import nn
6 +
7 +from vllm.config import (
8 + VllmConfig,
9 + get_current_vllm_config,
10 +)
11 +from vllm.distributed import (
12 + divide,
13 +)
14 +from vllm.forward_context import ForwardContext, get_forward_context
15 +from vllm.model_executor.custom_op import PluggableLayer
16 +from vllm.model_executor.layers.fla.ops import (
17 + chunk_gated_delta_rule,
18 + fused_recurrent_gated_delta_rule,
19 +)
20 +from vllm.model_executor.layers.layernorm import RMSNormGated
21 +from vllm.model_executor.layers.linear import (
22 + ColumnParallelLinear,
23 + MergedColumnParallelLinear,
24 + RowParallelLinear,
25 +)
26 +from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
27 +from vllm.model_executor.layers.mamba.mamba_utils import (
28 + MambaStateShapeCalculator,
29 + is_conv_state_dim_first,
30 +)
31 +from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
32 + causal_conv1d_fn,
33 + causal_conv1d_update,
34 +)
35 +from vllm.model_executor.model_loader.weight_utils import (
36 + sharded_weight_loader,
37 +)
38 +from vllm.model_executor.utils import set_weight_attrs
39 +from vllm.platforms import current_platform
40 +from vllm.triton_utils import tl, triton
41 +from vllm.triton_utils.allocation import set_triton_allocator
42 +from vllm.utils.torch_utils import direct_register_custom_op
43 +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
44 +
45 +
46 +@PluggableLayer.register("olmo_hybrid_gated_delta_net_attention")
47 +class OlmoHybridGatedDeltaNetAttention(GatedDeltaNetAttention):
48 + """
49 + Gated DeltaNet linear attention layer for OLMo Hybrid.
50 +
51 + This implements the linear attention mechanism that replaces sliding window
52 + attention in the hybrid architecture.
53 + """
54 +
55 + def get_state_shape(
56 + self,
57 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
58 + return MambaStateShapeCalculator.gated_delta_net_state_shape(
59 + self.tp_size,
60 + self.num_k_heads,
61 + self.num_v_heads,
62 + self.head_k_dim,
63 + self.head_v_dim,
64 + self.conv_kernel_size,
65 + self.num_spec,
66 + )
67 +
68 + def __init__(
69 + self,
70 + config,
71 + vllm_config: VllmConfig,
72 + prefix: str = "",
73 + ) -> None:
74 + super().__init__(config, vllm_config, prefix=prefix)
75 +
76 + assert getattr(config, "linear_use_gate", True), (
77 + "OlmoHybridGatedDeltaNet requires linear_use_gate=True"
78 + )
79 + self.num_k_heads = config.linear_num_key_heads
80 + self.num_v_heads = config.linear_num_value_heads
81 + self.head_k_dim = config.linear_key_head_dim
82 + self.head_v_dim = config.linear_value_head_dim
83 + self.conv_kernel_size = config.linear_conv_kernel_dim
84 + self.key_dim = self.head_k_dim * self.num_k_heads
85 + self.value_dim = self.head_v_dim * self.num_v_heads
86 + self.allow_neg_eigval = getattr(config, "linear_allow_neg_eigval", False)
87 +
88 + # Fused QKVG projection: 1 matmul instead of 4
89 + self.in_proj_qkvg = MergedColumnParallelLinear(
90 + input_size=self.hidden_size,
91 + output_sizes=[self.key_dim, self.key_dim, self.value_dim, self.value_dim],
92 + bias=False,
93 + quant_config=self.quant_config,
94 + prefix=f"{prefix}.in_proj_qkvg",
95 + )
96 +
97 + # Separate B and A projections to preserve numerical precision.
98 + # Fusing these into one matmul changes FP accumulation order for the
99 + # gating scalars, which compounds through the GDN recurrent state.
100 + self.b_proj = ColumnParallelLinear(
101 + input_size=self.hidden_size,
102 + output_size=self.num_v_heads,
103 + bias=False,
104 + quant_config=self.quant_config,
105 + prefix=f"{prefix}.b_proj",
106 + )
107 + self.a_proj = ColumnParallelLinear(
108 + input_size=self.hidden_size,
109 + output_size=self.num_v_heads,
110 + bias=False,
111 + quant_config=self.quant_config,
112 + prefix=f"{prefix}.a_proj",
113 + )
114 +
115 + # Fused conv1d: single parameter instead of 3
116 + self.conv_dim = self.key_dim * 2 + self.value_dim
117 + self.conv1d = ColumnParallelLinear(
118 + input_size=self.conv_kernel_size,
119 + output_size=self.conv_dim,
120 + bias=False,
121 + prefix=f"{prefix}.conv1d",
122 + )
123 + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
124 + delattr(self.conv1d.weight, "weight_loader")
125 + set_weight_attrs(
126 + self.conv1d.weight,
127 + {
128 + "weight_loader": _make_fused_conv1d_weight_loader(
129 + [self.key_dim, self.key_dim, self.value_dim],
130 + self.tp_size,
131 + self.tp_rank,
132 + )
133 + },
134 + )
135 +
136 + self.dt_bias = nn.Parameter(
137 + torch.ones(self.num_v_heads // self.tp_size),
138 + )
139 + self.A_log = nn.Parameter(
140 + torch.empty(
141 + divide(self.num_v_heads, self.tp_size),
142 + )
143 + )
144 +
145 + set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
146 + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
147 +
148 + # use eps=1e-5 to match FLA's FusedRMSNormGated
149 + self.o_norm = RMSNormGated(
150 + self.head_v_dim,
151 + eps=1e-5,
152 + group_size=None,
153 + norm_before_gate=True,
154 + device=current_platform.current_device(),
155 + dtype=config.torch_dtype if hasattr(config, "torch_dtype") else None,
156 + )
157 +
158 + self.o_proj = RowParallelLinear(
159 + self.value_dim,
160 + self.hidden_size,
161 + bias=False,
162 + input_is_parallel=True,
163 + quant_config=self.quant_config,
164 + prefix=f"{prefix}.o_proj",
165 + )
166 +
167 + # FLA triton kernels need a PyTorch-backed allocator for scratch
168 + # memory (required by triton >= 3.x autotuner). Set once at init.
169 + set_triton_allocator(current_platform.current_device())
170 +
171 + compilation_config = get_current_vllm_config().compilation_config
172 + if prefix in compilation_config.static_forward_context:
173 + raise ValueError(f"Duplicate layer name: {prefix}")
174 + compilation_config.static_forward_context[prefix] = self
175 +
176 + def rearrange_mixed_qkv(self, mixed_qkv):
177 + if mixed_qkv is None:
178 + return None, None, None
179 + query, key, value = torch.split(
180 + mixed_qkv,
181 + [
182 + self.key_dim // self.tp_size,
183 + self.key_dim // self.tp_size,
184 + self.value_dim // self.tp_size,
185 + ],
186 + dim=-1,
187 + )
188 +
189 + num_k_heads = self.num_k_heads // self.tp_size
190 + num_v_heads = self.num_v_heads // self.tp_size
191 +
192 + query = rearrange(query, "l (h d) -> 1 l h d", h=num_k_heads, d=self.head_k_dim)
193 + key = rearrange(key, "l (h d) -> 1 l h d", h=num_k_heads, d=self.head_k_dim)
194 + value = rearrange(value, "l (h d) -> 1 l h d", h=num_v_heads, d=self.head_v_dim)
195 +
196 + # GQA expansion if needed
197 + if num_v_heads > num_k_heads:
198 + expand_ratio = num_v_heads // num_k_heads
199 + query = query.unsqueeze(3).expand(-1, -1, -1, expand_ratio, -1)
200 + query = query.reshape(1, query.shape[1], num_v_heads, self.head_k_dim)
201 + key = key.unsqueeze(3).expand(-1, -1, -1, expand_ratio, -1)
202 + key = key.reshape(1, key.shape[1], num_v_heads, self.head_k_dim)
203 +
204 + return query.contiguous(), key.contiguous(), value.contiguous()
205 +
206 + def forward(
207 + self,
208 + hidden_states: torch.Tensor,
209 + output: torch.Tensor,
210 + ):
211 + # NOTE: We wrap the ENTIRE linear attention forward (projections +
212 + # core recurrence + output norm + output projection) in a single
213 + # custom op, rather than just wrapping the recurrent core like
214 + # other GDN models (e.g. Qwen3Next) do.
215 + #
216 + # Why: torch.compile with inductor generates fused kernels for
217 + # matmuls and pointwise ops. These fused kernels can differ in
218 + # floating-point accumulation order from eager-mode cuBLAS,
219 + # introducing small numerical differences (~1e-7 per op). For
220 + # standard transformer attention this is harmless because each
221 + # position is computed independently. But for the GDN recurrent
222 + # state, these tiny input differences compound at every timestep
223 + # across the full sequence length, causing severe logprob
224 + # divergence (e.g. ~15% top-1 agreement with eager baseline).
225 + #
226 + # By making the full forward opaque to inductor, the projections
227 + # and output norm run with eager-mode kernels (cuBLAS, triton),
228 + # preserving numerical consistency. The tradeoff is reduced
229 + # compilation speedup (~1.5x vs ~3x), but logprob agreement
230 + # improves from ~15% to ~83% top-1 vs eager.
231 + #
232 + # The remaining ~17% divergence comes from inductor compiling
233 + # the MLP and transformer attention layers that are NOT wrapped
234 + # in custom ops -- their small precision differences propagate
235 + # as inputs to the GDN layers from outside.
236 + torch.ops.vllm.olmo_hybrid_gdn_full_forward(
237 + hidden_states,
238 + output,
239 + self.prefix,
240 + )
241 +
242 + def _full_forward(
243 + self,
244 + hidden_states: torch.Tensor,
245 + output: torch.Tensor,
246 + ):
247 + num_tokens = hidden_states.size(0)
248 +
249 + # ============================================================
250 + # Part 1: Input Projection (2 fused matmuls instead of 6)
251 + # ============================================================
252 + projected_qkvg, _ = self.in_proj_qkvg(hidden_states)
253 + conv_dim_sharded = (self.key_dim * 2 + self.value_dim) // self.tp_size
254 + mixed_qkv = projected_qkvg[..., :conv_dim_sharded]
255 + gate = projected_qkvg[..., conv_dim_sharded:]
256 +
257 + b, _ = self.b_proj(hidden_states)
258 + a, _ = self.a_proj(hidden_states)
259 +
260 + # ============================================================
261 + # Part 2: Core Attention
262 + # ============================================================
263 + core_attn_out = torch.zeros(
264 + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
265 + dtype=hidden_states.dtype,
266 + device=hidden_states.device,
267 + )
268 +
269 + self._forward_core(
270 + mixed_qkv=mixed_qkv,
271 + b=b,
272 + a=a,
273 + core_attn_out=core_attn_out,
274 + )
275 +
276 + # ============================================================
277 + # Part 3: Output Projection
278 + # ============================================================
279 + gate = gate.view(num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim)
280 + core_attn_out_flat = core_attn_out.reshape(-1, core_attn_out.shape[-1])
281 + gate_flat = gate.reshape(-1, gate.shape[-1])
282 + core_attn_out_normed = self.o_norm(core_attn_out_flat, gate_flat)
283 + core_attn_out = core_attn_out_normed.view(
284 + num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim
285 + )
286 +
287 + core_attn_out = rearrange(core_attn_out, "l h d -> l (h d)")
288 + output[:num_tokens], _ = self.o_proj(core_attn_out)
289 +
290 + def _forward_core(
291 + self,
292 + mixed_qkv: torch.Tensor,
293 + b: torch.Tensor,
294 + a: torch.Tensor,
295 + core_attn_out: torch.Tensor,
296 + ):
297 + """
298 + Core attention computation (called by custom op).
299 + """
300 + forward_context = get_forward_context()
301 + attn_metadata = forward_context.attn_metadata
302 +
303 + if attn_metadata is None:
304 + # V1 profile run
305 + return
306 +
307 + assert isinstance(attn_metadata, dict)
308 + attn_metadata = attn_metadata[self.prefix] # type: ignore[assignment]
309 + assert isinstance(attn_metadata, GDNAttentionMetadata)
310 + has_initial_state = attn_metadata.has_initial_state
311 + spec_query_start_loc = attn_metadata.spec_query_start_loc
312 + non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
313 + spec_sequence_masks = attn_metadata.spec_sequence_masks
314 + spec_token_indx = attn_metadata.spec_token_indx
315 + non_spec_token_indx = attn_metadata.non_spec_token_indx
316 + spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor
317 + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor
318 + self_kv_cache = self.kv_cache
319 + # conv_state must be (..., dim, width-1) for the conv kernels.
320 + # DS layout stores it that way directly; SD layout needs a transpose.
321 + conv_state = (
322 + self_kv_cache[0]
323 + if is_conv_state_dim_first()
324 + else self_kv_cache[0].transpose(-1, -2)
325 + )
326 + ssm_state = self_kv_cache[1]
327 + num_actual_tokens = attn_metadata.num_actual_tokens
328 + num_accepted_tokens = attn_metadata.num_accepted_tokens
329 +
330 + mixed_qkv = mixed_qkv[:num_actual_tokens]
331 + b = b[:num_actual_tokens]
332 + a = a[:num_actual_tokens]
333 +
334 + conv_weights = self.conv1d.weight.view(
335 + self.conv1d.weight.size(0), self.conv1d.weight.size(2)
336 + )
337 +
338 + if spec_sequence_masks is not None:
339 + if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:
340 + mixed_qkv_spec = mixed_qkv
341 + mixed_qkv_non_spec = None
342 + else:
343 + mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
344 + mixed_qkv_non_spec = mixed_qkv.index_select(0, non_spec_token_indx)
345 + else:
346 + mixed_qkv_spec = None
347 + mixed_qkv_non_spec = mixed_qkv
348 +
349 + if spec_sequence_masks is not None:
350 + assert spec_query_start_loc is not None
351 + assert spec_state_indices_tensor is not None
352 + assert num_accepted_tokens is not None
353 + mixed_qkv_spec = causal_conv1d_update(
354 + mixed_qkv_spec,
355 + conv_state,
356 + conv_weights,
357 + None, # no bias
358 + self.activation,
359 + conv_state_indices=spec_state_indices_tensor[:, 0][
360 + : attn_metadata.num_spec_decodes
361 + ],
362 + num_accepted_tokens=num_accepted_tokens,
363 + query_start_loc=spec_query_start_loc,
364 + max_query_len=spec_state_indices_tensor.size(-1),
365 + validate_data=False,
366 + )
367 +
368 + if attn_metadata.num_prefills > 0:
369 + assert mixed_qkv_non_spec is not None
370 + mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1)
371 + mixed_qkv_non_spec = causal_conv1d_fn(
372 + mixed_qkv_non_spec_T,
373 + conv_weights,
374 + None,
375 + activation=self.activation,
376 + conv_states=conv_state,
377 + has_initial_state=has_initial_state,
378 + cache_indices=non_spec_state_indices_tensor,
379 + query_start_loc=non_spec_query_start_loc,
380 + metadata=attn_metadata,
381 + ).transpose(0, 1)
382 + elif attn_metadata.num_decodes > 0:
383 + assert non_spec_state_indices_tensor is not None
384 + mixed_qkv_non_spec = causal_conv1d_update(
385 + mixed_qkv_non_spec,
386 + conv_state,
387 + conv_weights,
388 + None,
389 + self.activation,
390 + conv_state_indices=non_spec_state_indices_tensor[
391 + : attn_metadata.num_decodes
392 + ],
393 + validate_data=True,
394 + )
395 + else:
396 + mixed_qkv_non_spec = None
397 +
398 + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)
399 + query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv(
400 + mixed_qkv_non_spec
401 + )
402 +
403 + g, beta = fused_olmo_hybrid_gdn_gating(
404 + self.A_log, a, b, self.dt_bias, self.allow_neg_eigval
405 + )
406 +
407 + if spec_sequence_masks is not None:
408 + assert spec_token_indx is not None
409 + assert non_spec_token_indx is not None
410 + if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:
411 + g_spec = g
412 + beta_spec = beta
413 + g_non_spec = None
414 + beta_non_spec = None
415 + else:
416 + g_spec = g.index_select(1, spec_token_indx)
417 + beta_spec = beta.index_select(1, spec_token_indx)
418 + g_non_spec = g.index_select(1, non_spec_token_indx)
419 + beta_non_spec = beta.index_select(1, non_spec_token_indx)
420 + else:
421 + g_spec = None
422 + beta_spec = None
423 + g_non_spec = g
424 + beta_non_spec = beta
425 +
426 + if spec_sequence_masks is not None:
427 + assert spec_query_start_loc is not None
428 + assert spec_state_indices_tensor is not None
429 + assert num_accepted_tokens is not None
430 + core_attn_out_spec, last_recurrent_state = fused_recurrent_gated_delta_rule(
431 + q=query_spec,
432 + k=key_spec,
433 + v=value_spec,
434 + g=g_spec,
435 + beta=beta_spec,
436 + initial_state=ssm_state,
437 + inplace_final_state=True,
438 + cu_seqlens=spec_query_start_loc[: attn_metadata.num_spec_decodes + 1],
439 + ssm_state_indices=spec_state_indices_tensor,
440 + num_accepted_tokens=num_accepted_tokens,
441 + use_qk_l2norm_in_kernel=True,
442 + )
443 + else:
444 + core_attn_out_spec, last_recurrent_state = None, None
445 +
446 + if attn_metadata.num_prefills > 0:
447 + assert non_spec_state_indices_tensor is not None
448 + assert has_initial_state is not None
449 + assert non_spec_query_start_loc is not None
450 + initial_state = ssm_state[non_spec_state_indices_tensor].contiguous()
451 + initial_state[~has_initial_state, ...] = 0
452 + (
453 + core_attn_out_non_spec,
454 + last_recurrent_state,
455 + ) = chunk_gated_delta_rule(
456 + q=query_non_spec,
457 + k=key_non_spec,
458 + v=value_non_spec,
459 + g=g_non_spec,
460 + beta=beta_non_spec,
461 + initial_state=initial_state,
462 + output_final_state=True,
463 + cu_seqlens=non_spec_query_start_loc,
464 + use_qk_l2norm_in_kernel=True,
465 + )
466 + ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to(
467 + ssm_state.dtype
468 + )
469 + elif attn_metadata.num_decodes > 0:
470 + assert non_spec_query_start_loc is not None
471 + assert non_spec_state_indices_tensor is not None
472 + core_attn_out_non_spec, last_recurrent_state = (
473 + fused_recurrent_gated_delta_rule(
474 + q=query_non_spec,
475 + k=key_non_spec,
476 + v=value_non_spec,
477 + g=g_non_spec,
478 + beta=beta_non_spec,
479 + initial_state=ssm_state,
480 + inplace_final_state=True,
481 + cu_seqlens=non_spec_query_start_loc[
482 + : attn_metadata.num_decodes + 1
483 + ],
484 + ssm_state_indices=non_spec_state_indices_tensor,
485 + use_qk_l2norm_in_kernel=True,
486 + )
487 + )
488 + else:
489 + core_attn_out_non_spec, last_recurrent_state = None, None
490 +
491 + if spec_sequence_masks is not None and core_attn_out_non_spec is not None:
492 + merged_out = torch.empty(
493 + (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
494 + dtype=core_attn_out_non_spec.dtype,
495 + device=core_attn_out_non_spec.device,
496 + )
497 + merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
498 + merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
499 + core_attn_out[:num_actual_tokens] = merged_out.squeeze(0)
500 + elif spec_sequence_masks is not None:
501 + core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
502 + else:
503 + core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)
504 +
505 +
506 +def _make_fused_conv1d_weight_loader(dims, tp_size, tp_rank):
507 + """Weight loader for loading separate HF conv weights into a fused conv1d.
508 +
509 + dims: list of original (un-sharded) dims per section,
510 + e.g. [key_dim, key_dim, value_dim]
511 + """
512 + sharded_dims = [d // tp_size for d in dims]
513 +
514 + def weight_loader(param, loaded_weight, loaded_shard_id=None):
515 + if loaded_weight.dim() == 2:
516 + loaded_weight = loaded_weight.unsqueeze(1)
517 + dim = dims[loaded_shard_id]
518 + shard_size = dim // tp_size
519 + tp_start = tp_rank * shard_size
520 + sharded_weight = loaded_weight[tp_start : tp_start + shard_size]
521 + offset = sum(sharded_dims[:loaded_shard_id])
522 + param.data[offset : offset + shard_size].copy_(sharded_weight)
523 +
524 + return weight_loader
525 +
526 +
527 +def olmo_hybrid_gdn_full_forward(
528 + hidden_states: torch.Tensor,
529 + output: torch.Tensor,
530 + layer_name: str,
531 +) -> None:
532 + """Full linear attention forward wrapped as a custom op.
533 +
534 + Prevents inductor from compiling the projections around the GDN core,
535 + which would introduce numerical divergence that compounds through
536 + the recurrent state.
537 + """
538 + forward_context: ForwardContext = get_forward_context()
539 + self = forward_context.no_compile_layers[layer_name]
540 + self._full_forward(
541 + hidden_states=hidden_states,
542 + output=output,
543 + )
544 +
545 +
546 +def olmo_hybrid_gdn_full_forward_fake(
547 + hidden_states: torch.Tensor,
548 + output: torch.Tensor,
549 + layer_name: str,
550 +) -> None:
551 + """Fake implementation for torch.compile."""
552 + return
553 +
554 +
555 +direct_register_custom_op(
556 + op_name="olmo_hybrid_gdn_full_forward",
557 + op_func=olmo_hybrid_gdn_full_forward,
558 + mutates_args=["output"],
559 + fake_impl=olmo_hybrid_gdn_full_forward_fake,
560 +)
561 +
562 +
563 +@triton.jit
564 +def fused_olmo_hybrid_gdn_gating_kernel(
565 + g,
566 + beta_output,
567 + A_log,
568 + a,
569 + b,
570 + dt_bias,
571 + seq_len,
572 + allow_neg_eigval: tl.constexpr,
573 + NUM_HEADS: tl.constexpr,
574 + beta: tl.constexpr,
575 + threshold: tl.constexpr,
576 + BLK_HEADS: tl.constexpr,
577 +):
578 + i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2)
579 + head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS)
580 + off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off
581 + mask = head_off < NUM_HEADS
582 + blk_A_log = tl.load(A_log + head_off, mask=mask)
583 + blk_a = tl.load(a + off, mask=mask)
584 + blk_b = tl.load(b + off, mask=mask)
585 + blk_bias = tl.load(dt_bias + head_off, mask=mask)
586 +
587 + # g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
588 + x = blk_a.to(tl.float32) + blk_bias.to(tl.float32)
589 + softplus_x = tl.where(
590 + beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x
591 + )
592 + blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x
593 + tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask)
594 +
595 + # beta = self.b_proj(hidden_states).sigmoid()
596 + # if self.allow_neg_eigval: beta = beta * 2.0
597 + blk_beta_output = tl.sigmoid(blk_b.to(tl.float32))
598 + if allow_neg_eigval:
599 + blk_beta_output = blk_beta_output * 2.0
600 + tl.store(
601 + beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask
602 + )
603 +
604 +
605 +def fused_olmo_hybrid_gdn_gating(
606 + A_log: torch.Tensor,
607 + a: torch.Tensor,
608 + b: torch.Tensor,
609 + dt_bias: torch.Tensor,
610 + allow_neg_eigval: bool = False,
611 + beta: float = 1.0,
612 + threshold: float = 20.0,
613 +) -> tuple[torch.Tensor, torch.Tensor]:
614 + batch, num_heads = a.shape
615 + seq_len = 1
616 + grid = (batch, seq_len, triton.cdiv(num_heads, 8))
617 + g = torch.empty(1, batch, num_heads, dtype=torch.float32, device=a.device)
618 + beta_output = torch.empty(1, batch, num_heads, dtype=torch.float32, device=b.device)
619 + fused_olmo_hybrid_gdn_gating_kernel[grid](
620 + g,
621 + beta_output,
622 + A_log,
623 + a,
624 + b,
625 + dt_bias,
626 + seq_len,
627 + allow_neg_eigval,
628 + num_heads,
629 + beta,
630 + threshold,
631 + 8,
632 + num_warps=1,
633 + )
634 + return g, beta_output
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py renamed
+19 -52
@@ -5,7 +5,6 @@
5 import torch
6 from einops import rearrange
7 from torch import nn
8 -from transformers.activations import ACT2FN
8
9 from vllm import envs
10 from vllm._aiter_ops import rocm_aiter_ops
@@ -15,8 +14,6 @@ from vllm.config import (
14 )
15 from vllm.distributed import (
16 divide,
18 - get_tensor_model_parallel_rank,
19 - get_tensor_model_parallel_world_size,
17 )
18 from vllm.forward_context import ForwardContext, get_forward_context
19 from vllm.logger import init_logger
@@ -37,10 +34,9 @@ from vllm.model_executor.layers.linear import (
34 MergedColumnParallelLinear,
35 RowParallelLinear,
36 )
40 -from vllm.model_executor.layers.mamba.abstract import MambaBase
37 +from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
38 from vllm.model_executor.layers.mamba.mamba_mixer2 import mamba_v2_sharded_weight_loader
39 from vllm.model_executor.layers.mamba.mamba_utils import (
43 - MambaStateDtypeCalculator,
40 MambaStateShapeCalculator,
41 is_conv_state_dim_first,
42 )
@@ -55,7 +51,6 @@ from vllm.model_executor.layers.quantization.inc import INCConfig
51 from vllm.model_executor.model_loader.weight_utils import (
52 sharded_weight_loader,
53 )
58 -from vllm.model_executor.models.utils import extract_layer_index
54 from vllm.model_executor.utils import set_weight_attrs
55 from vllm.platforms import current_platform
56 from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
@@ -67,7 +62,6 @@ from vllm.utils.torch_utils import (
62 direct_register_custom_op,
63 )
64 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
70 -from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
65
66 # Optional ROCm AITER Triton kernels for the GDN decode fast-path.
67 # Availability is checked centrally via rocm_aiter_ops; the actual function
@@ -263,20 +257,11 @@ class ChunkGatedDeltaRule(CustomOp):
257 )
258
259
266 -@PluggableLayer.register("gated_delta_net_attention")
267 -class GatedDeltaNetAttention(PluggableLayer, MambaBase):
268 - @property
269 - def mamba_type(self) -> MambaAttentionBackendEnum:
270 - return MambaAttentionBackendEnum.GDN_ATTN
271 -
272 - def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
273 - return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
274 - self.model_config.dtype,
275 - self.cache_config.mamba_cache_dtype,
276 - self.cache_config.mamba_ssm_cache_dtype,
277 - )
278 -
279 - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
260 +@PluggableLayer.register("qwen_gated_delta_net_attention")
261 +class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
262 + def get_state_shape(
263 + self,
264 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
265 return MambaStateShapeCalculator.gated_delta_net_state_shape(
266 self.tp_size,
267 self.num_k_heads,
@@ -294,33 +279,15 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
279 prefix: str = "",
280 gqa_interleaved_layout=False,
281 ) -> None:
297 - super().__init__()
298 - self.tp_size = get_tensor_model_parallel_world_size()
299 - self.tp_rank = get_tensor_model_parallel_rank()
300 - self.hidden_size = config.hidden_size
301 - self.num_v_heads = config.linear_num_value_heads
282 + super().__init__(config, vllm_config, prefix)
283 +
284 self.num_k_heads = config.linear_num_key_heads
285 + self.num_v_heads = config.linear_num_value_heads
286 self.head_k_dim = config.linear_key_head_dim
287 self.head_v_dim = config.linear_value_head_dim
288 + self.conv_kernel_size = config.linear_conv_kernel_dim
289 self.key_dim = self.head_k_dim * self.num_k_heads
290 self.value_dim = self.head_v_dim * self.num_v_heads
307 -
308 - self.conv_kernel_size = config.linear_conv_kernel_dim
309 - self.layer_idx = extract_layer_index(prefix)
310 - self.activation = config.hidden_act
311 - self.act = ACT2FN[config.hidden_act]
312 - self.layer_norm_epsilon = config.rms_norm_eps
313 - self.prefix = prefix
314 - self.config = config
315 - self.model_config = vllm_config.model_config
316 - self.cache_config = vllm_config.cache_config
317 - quant_config = vllm_config.quant_config
318 - self.speculative_config = vllm_config.speculative_config
319 - self.num_spec = (
320 - self.speculative_config.num_speculative_tokens
321 - if self.speculative_config
322 - else 0
323 - )
291 self.gqa_interleaved_layout = gqa_interleaved_layout
292 if current_platform.is_xpu():
293 self._forward_method = self.forward_xpu
@@ -355,7 +322,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
322 hidden_size=self.hidden_size,
323 key_dim=self.key_dim,
324 value_dim=self.value_dim,
358 - quant_config=quant_config,
325 + quant_config=self.quant_config,
326 prefix=f"{prefix}.in_proj_qkvz",
327 )
328
@@ -365,10 +332,10 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
332 self.in_proj_ba = self.create_ba_proj(
333 hidden_size=self.hidden_size,
334 num_v_heads=self.num_v_heads,
368 - quant_config=quant_config,
335 + quant_config=self.quant_config,
336 prefix=f"{prefix}.in_proj_ba",
337 )
371 - self.disable_tp_for_ba_proj = self.maybe_disable_tp(quant_config)
338 + self.disable_tp_for_ba_proj = self.maybe_disable_tp(self.quant_config)
339
340 query_key_settings = (self.key_dim, 0, False)
341 value_settings = (self.value_dim, 0, False)
@@ -421,7 +388,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
388 self.hidden_size,
389 bias=False,
390 input_is_parallel=True,
424 - quant_config=quant_config,
391 + quant_config=self.quant_config,
392 prefix=f"{prefix}.out_proj",
393 )
394
@@ -765,7 +732,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
732 device=projected_states_qkvz.device,
733 )
734
768 - torch.ops.vllm.gdn_attention_core(
735 + torch.ops.vllm.qwen_gdn_attention_core(
736 projected_states_qkvz,
737 projected_states_ba,
738 z,
@@ -826,7 +793,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
793 device=hidden_states.device,
794 )
795
829 - torch.ops.vllm.gdn_attention_core(
796 + torch.ops.vllm.qwen_gdn_attention_core(
797 mixed_qkv,
798 b,
799 a,
@@ -1492,7 +1459,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
1459 return
1460
1461
1495 -def gdn_attention_core(
1462 +def qwen_gdn_attention_core(
1463 qkv_or_qkvz: torch.Tensor,
1464 b_or_ba: torch.Tensor,
1465 a_or_z_out: torch.Tensor,
@@ -1545,8 +1512,8 @@ def gdn_attention_core_fake(
1512
1513
1514 direct_register_custom_op(
1548 - op_name="gdn_attention_core",
1549 - op_func=gdn_attention_core,
1515 + op_name="qwen_gdn_attention_core",
1516 + op_func=qwen_gdn_attention_core,
1517 mutates_args=["a_or_z_out", "core_attn_out"],
1518 fake_impl=gdn_attention_core_fake,
1519 )
vllm/model_executor/models/kimi_linear.py
+13 -27
@@ -7,7 +7,7 @@ import torch
7 from torch import nn
8
9 from vllm.compilation.decorators import support_torch_compile
10 -from vllm.config import CacheConfig, ModelConfig, ParallelConfig, VllmConfig
10 +from vllm.config import CacheConfig, VllmConfig
11 from vllm.distributed import (
12 get_pp_group,
13 get_tensor_model_parallel_world_size,
@@ -18,7 +18,6 @@ from vllm.model_executor.layers.fused_moe import (
18 FusedMoE,
19 fused_moe_make_expert_params_mapping,
20 )
21 -from vllm.model_executor.layers.kda import KimiDeltaAttention
21 from vllm.model_executor.layers.layernorm import RMSNorm
22 from vllm.model_executor.layers.linear import (
23 ColumnParallelLinear,
@@ -27,6 +26,9 @@ from vllm.model_executor.layers.linear import (
26 RowParallelLinear,
27 )
28 from vllm.model_executor.layers.logits_processor import LogitsProcessor
29 +from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import (
30 + KimiGatedDeltaNetAttention,
31 +)
32 from vllm.model_executor.layers.mamba.mamba_utils import (
33 MambaStateCopyFunc,
34 MambaStateCopyFuncCalculator,
@@ -286,26 +288,22 @@ class KimiDecoderLayer(nn.Module):
288 def __init__(
289 self,
290 config: KimiLinearConfig,
289 - layer_idx: int,
290 - cache_config: CacheConfig | None = None,
291 - quant_config: QuantizationConfig | None = None,
292 - parallel_config: ParallelConfig | None = None,
293 - model_config: ModelConfig | None = None,
291 + vllm_config: VllmConfig,
292 prefix: str = "",
295 - **kwargs,
293 ) -> None:
294 super().__init__()
295 self.hidden_size = config.hidden_size
296
297 self.is_moe = config.is_moe
298 + layer_idx = int(prefix.rsplit(".", 1)[1])
299 + model_config = vllm_config.model_config
300 + cache_config = vllm_config.cache_config
301 + quant_config = vllm_config.quant_config
302
303 if config.is_kda_layer(layer_idx):
303 - self.self_attn = KimiDeltaAttention(
304 - layer_idx=layer_idx,
305 - hidden_size=config.hidden_size,
306 - quant_config=quant_config,
307 - cache_config=cache_config,
308 - model_config=config,
304 + self.self_attn = KimiGatedDeltaNetAttention(
305 + config,
306 + vllm_config,
307 prefix=f"{prefix}.self_attn",
308 )
309 else:
@@ -385,10 +383,6 @@ class KimiLinearModel(nn.Module):
383 super().__init__()
384
385 config = vllm_config.model_config.hf_text_config
388 - model_config = vllm_config.model_config
389 - cache_config = vllm_config.cache_config
390 - quant_config = vllm_config.quant_config
391 - parallel_config = vllm_config.parallel_config
386 self.config = config
387
388 self.vocab_size = config.vocab_size
@@ -402,19 +396,11 @@ class KimiLinearModel(nn.Module):
396 else:
397 self.embed_tokens = PPMissingLayer()
398
405 - extra_kwargs = {}
406 -
399 def get_layer(prefix: str):
408 - layer_idx = int(prefix.rsplit(".", 1)[1])
400 return KimiDecoderLayer(
401 config,
411 - layer_idx,
412 - cache_config,
413 - quant_config,
414 - parallel_config,
415 - model_config,
402 + vllm_config,
403 prefix,
417 - **extra_kwargs,
404 )
405
406 self.start_layer, self.end_layer, self.layers = make_layers(
vllm/model_executor/models/olmo_hybrid.py
+6 -645
@@ -26,55 +26,38 @@ from functools import partial
26 from itertools import islice
27
28 import torch
29 -from einops import rearrange
29 from torch import nn
31 -from transformers.activations import ACT2FN
30
31 from vllm.compilation.decorators import support_torch_compile
32 from vllm.config import (
35 - CacheConfig,
36 - ModelConfig,
37 - SpeculativeConfig,
33 VllmConfig,
39 - get_current_vllm_config,
34 )
35 from vllm.distributed import (
42 - divide,
36 get_pp_group,
37 get_tensor_model_parallel_rank,
38 get_tensor_model_parallel_world_size,
39 tensor_model_parallel_all_gather,
40 )
41 from vllm.distributed.utils import split_tensor_along_last_dim
49 -from vllm.forward_context import ForwardContext, get_forward_context
42 from vllm.logger import init_logger
43 from vllm.model_executor.layers.activation import SiluAndMul
44 from vllm.model_executor.layers.attention import Attention
53 -from vllm.model_executor.layers.fla.ops import (
54 - chunk_gated_delta_rule,
55 - fused_recurrent_gated_delta_rule,
56 -)
57 -from vllm.model_executor.layers.layernorm import RMSNorm, RMSNormGated
45 +from vllm.model_executor.layers.layernorm import RMSNorm
46 from vllm.model_executor.layers.linear import (
59 - ColumnParallelLinear,
47 MergedColumnParallelLinear,
48 QKVParallelLinear,
49 RowParallelLinear,
50 )
51 from vllm.model_executor.layers.logits_processor import LogitsProcessor
65 -from vllm.model_executor.layers.mamba.abstract import MambaBase
52 +from vllm.model_executor.layers.mamba.gdn.olmo_gdn_linear_attn import (
53 + OlmoHybridGatedDeltaNetAttention,
54 +)
55 from vllm.model_executor.layers.mamba.mamba_utils import (
56 MambaStateCopyFunc,
57 MambaStateCopyFuncCalculator,
58 MambaStateDtypeCalculator,
59 MambaStateShapeCalculator,
71 - is_conv_state_dim_first,
60 )
73 -from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
74 - causal_conv1d_fn,
75 - causal_conv1d_update,
76 -)
77 -from vllm.model_executor.layers.quantization import QuantizationConfig
61 from vllm.model_executor.layers.rotary_embedding import get_rope
62 from vllm.model_executor.layers.vocab_parallel_embedding import (
63 ParallelLMHead,
@@ -82,17 +65,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
65 )
66 from vllm.model_executor.model_loader.weight_utils import (
67 default_weight_loader,
85 - sharded_weight_loader,
68 )
87 -from vllm.model_executor.utils import set_weight_attrs
88 -from vllm.platforms import current_platform
69 from vllm.sequence import IntermediateTensors
90 -from vllm.triton_utils import tl, triton
91 -from vllm.triton_utils.allocation import set_triton_allocator
92 -from vllm.utils.torch_utils import direct_register_custom_op
93 -from vllm.v1.attention.backend import AttentionMetadata
94 -from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
95 -from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
70
71 from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP
72 from .utils import (
@@ -107,502 +81,6 @@ from .utils import (
81 logger = init_logger(__name__)
82
83
110 -def _make_fused_conv1d_weight_loader(dims, tp_size, tp_rank):
111 - """Weight loader for loading separate HF conv weights into a fused conv1d.
112 -
113 - dims: list of original (un-sharded) dims per section,
114 - e.g. [key_dim, key_dim, value_dim]
115 - """
116 - sharded_dims = [d // tp_size for d in dims]
117 -
118 - def weight_loader(param, loaded_weight, loaded_shard_id=None):
119 - if loaded_weight.dim() == 2:
120 - loaded_weight = loaded_weight.unsqueeze(1)
121 - dim = dims[loaded_shard_id]
122 - shard_size = dim // tp_size
123 - tp_start = tp_rank * shard_size
124 - sharded_weight = loaded_weight[tp_start : tp_start + shard_size]
125 - offset = sum(sharded_dims[:loaded_shard_id])
126 - param.data[offset : offset + shard_size].copy_(sharded_weight)
127 -
128 - return weight_loader
129 -
130 -
131 -class OlmoHybridGatedDeltaNet(nn.Module, MambaBase):
132 - """
133 - Gated DeltaNet linear attention layer for OLMo Hybrid.
134 -
135 - This implements the linear attention mechanism that replaces sliding window
136 - attention in the hybrid architecture.
137 - """
138 -
139 - @property
140 - def mamba_type(self) -> MambaAttentionBackendEnum:
141 - return MambaAttentionBackendEnum.GDN_ATTN
142 -
143 - def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
144 - return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
145 - self.model_config.dtype,
146 - self.cache_config.mamba_cache_dtype,
147 - self.cache_config.mamba_ssm_cache_dtype,
148 - )
149 -
150 - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
151 - return MambaStateShapeCalculator.gated_delta_net_state_shape(
152 - self.tp_size,
153 - self.num_k_heads,
154 - self.num_v_heads,
155 - self.head_k_dim,
156 - self.head_v_dim,
157 - self.conv_kernel_size,
158 - self.num_spec,
159 - )
160 -
161 - def __init__(
162 - self,
163 - config,
164 - model_config: ModelConfig | None = None,
165 - cache_config: CacheConfig | None = None,
166 - quant_config: QuantizationConfig | None = None,
167 - speculative_config: SpeculativeConfig | None = None,
168 - prefix: str = "",
169 - ) -> None:
170 - super().__init__()
171 - self.tp_size = get_tensor_model_parallel_world_size()
172 - self.tp_rank = get_tensor_model_parallel_rank()
173 - self.hidden_size = config.hidden_size
174 - self.num_v_heads = config.linear_num_value_heads
175 - self.num_k_heads = config.linear_num_key_heads
176 - self.head_k_dim = config.linear_key_head_dim
177 - self.head_v_dim = config.linear_value_head_dim
178 - self.key_dim = self.head_k_dim * self.num_k_heads
179 - self.value_dim = self.head_v_dim * self.num_v_heads
180 -
181 - self.conv_kernel_size = config.linear_conv_kernel_dim
182 - self.layer_idx = extract_layer_index(prefix)
183 - self.activation = config.hidden_act
184 - self.act = ACT2FN[config.hidden_act]
185 - self.layer_norm_epsilon = config.rms_norm_eps
186 - assert getattr(config, "linear_use_gate", True), (
187 - "OlmoHybridGatedDeltaNet requires linear_use_gate=True"
188 - )
189 - self.allow_neg_eigval = getattr(config, "linear_allow_neg_eigval", False)
190 - self.prefix = prefix
191 -
192 - self.config = config
193 - self.model_config = model_config
194 - self.cache_config = cache_config
195 - self.quant_config = quant_config
196 - self.speculative_config = speculative_config
197 - self.num_spec = (
198 - self.speculative_config.num_speculative_tokens
199 - if self.speculative_config
200 - else 0
201 - )
202 -
203 - # Fused QKVG projection: 1 matmul instead of 4
204 - self.in_proj_qkvg = MergedColumnParallelLinear(
205 - input_size=self.hidden_size,
206 - output_sizes=[self.key_dim, self.key_dim, self.value_dim, self.value_dim],
207 - bias=False,
208 - quant_config=quant_config,
209 - prefix=f"{prefix}.in_proj_qkvg",
210 - )
211 -
212 - # Separate B and A projections to preserve numerical precision.
213 - # Fusing these into one matmul changes FP accumulation order for the
214 - # gating scalars, which compounds through the GDN recurrent state.
215 - self.b_proj = ColumnParallelLinear(
216 - input_size=self.hidden_size,
217 - output_size=self.num_v_heads,
218 - bias=False,
219 - quant_config=quant_config,
220 - prefix=f"{prefix}.b_proj",
221 - )
222 - self.a_proj = ColumnParallelLinear(
223 - input_size=self.hidden_size,
224 - output_size=self.num_v_heads,
225 - bias=False,
226 - quant_config=quant_config,
227 - prefix=f"{prefix}.a_proj",
228 - )
229 -
230 - # Fused conv1d: single parameter instead of 3
231 - self.conv_dim = self.key_dim * 2 + self.value_dim
232 - self.conv1d = ColumnParallelLinear(
233 - input_size=self.conv_kernel_size,
234 - output_size=self.conv_dim,
235 - bias=False,
236 - prefix=f"{prefix}.conv1d",
237 - )
238 - self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
239 - delattr(self.conv1d.weight, "weight_loader")
240 - set_weight_attrs(
241 - self.conv1d.weight,
242 - {
243 - "weight_loader": _make_fused_conv1d_weight_loader(
244 - [self.key_dim, self.key_dim, self.value_dim],
245 - self.tp_size,
246 - self.tp_rank,
247 - )
248 - },
249 - )
250 -
251 - self.dt_bias = nn.Parameter(
252 - torch.ones(self.num_v_heads // self.tp_size),
253 - )
254 - self.A_log = nn.Parameter(
255 - torch.empty(
256 - divide(self.num_v_heads, self.tp_size),
257 - )
258 - )
259 -
260 - set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
261 - set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
262 -
263 - # use eps=1e-5 to match FLA's FusedRMSNormGated
264 - self.o_norm = RMSNormGated(
265 - self.head_v_dim,
266 - eps=1e-5,
267 - group_size=None,
268 - norm_before_gate=True,
269 - device=current_platform.current_device(),
270 - dtype=config.torch_dtype if hasattr(config, "torch_dtype") else None,
271 - )
272 -
273 - self.o_proj = RowParallelLinear(
274 - self.value_dim,
275 - self.hidden_size,
276 - bias=False,
277 - input_is_parallel=True,
278 - quant_config=quant_config,
279 - prefix=f"{prefix}.o_proj",
280 - )
281 -
282 - # FLA triton kernels need a PyTorch-backed allocator for scratch
283 - # memory (required by triton >= 3.x autotuner). Set once at init.
284 - set_triton_allocator(current_platform.current_device())
285 -
286 - compilation_config = get_current_vllm_config().compilation_config
287 - if prefix in compilation_config.static_forward_context:
288 - raise ValueError(f"Duplicate layer name: {prefix}")
289 - compilation_config.static_forward_context[prefix] = self
290 -
291 - def rearrange_mixed_qkv(self, mixed_qkv):
292 - if mixed_qkv is None:
293 - return None, None, None
294 - query, key, value = torch.split(
295 - mixed_qkv,
296 - [
297 - self.key_dim // self.tp_size,
298 - self.key_dim // self.tp_size,
299 - self.value_dim // self.tp_size,
300 - ],
301 - dim=-1,
302 - )
303 -
304 - num_k_heads = self.num_k_heads // self.tp_size
305 - num_v_heads = self.num_v_heads // self.tp_size
306 -
307 - query = rearrange(query, "l (h d) -> 1 l h d", h=num_k_heads, d=self.head_k_dim)
308 - key = rearrange(key, "l (h d) -> 1 l h d", h=num_k_heads, d=self.head_k_dim)
309 - value = rearrange(value, "l (h d) -> 1 l h d", h=num_v_heads, d=self.head_v_dim)
310 -
311 - # GQA expansion if needed
312 - if num_v_heads > num_k_heads:
313 - expand_ratio = num_v_heads // num_k_heads
314 - query = query.unsqueeze(3).expand(-1, -1, -1, expand_ratio, -1)
315 - query = query.reshape(1, query.shape[1], num_v_heads, self.head_k_dim)
316 - key = key.unsqueeze(3).expand(-1, -1, -1, expand_ratio, -1)
317 - key = key.reshape(1, key.shape[1], num_v_heads, self.head_k_dim)
318 -
319 - return query.contiguous(), key.contiguous(), value.contiguous()
320 -
321 - def forward(
322 - self,
323 - hidden_states: torch.Tensor,
324 - output: torch.Tensor,
325 - ):
326 - # NOTE: We wrap the ENTIRE linear attention forward (projections +
327 - # core recurrence + output norm + output projection) in a single
328 - # custom op, rather than just wrapping the recurrent core like
329 - # other GDN models (e.g. Qwen3Next) do.
330 - #
331 - # Why: torch.compile with inductor generates fused kernels for
332 - # matmuls and pointwise ops. These fused kernels can differ in
333 - # floating-point accumulation order from eager-mode cuBLAS,
334 - # introducing small numerical differences (~1e-7 per op). For
335 - # standard transformer attention this is harmless because each
336 - # position is computed independently. But for the GDN recurrent
337 - # state, these tiny input differences compound at every timestep
338 - # across the full sequence length, causing severe logprob
339 - # divergence (e.g. ~15% top-1 agreement with eager baseline).
340 - #
341 - # By making the full forward opaque to inductor, the projections
342 - # and output norm run with eager-mode kernels (cuBLAS, triton),
343 - # preserving numerical consistency. The tradeoff is reduced
344 - # compilation speedup (~1.5x vs ~3x), but logprob agreement
345 - # improves from ~15% to ~83% top-1 vs eager.
346 - #
347 - # The remaining ~17% divergence comes from inductor compiling
348 - # the MLP and transformer attention layers that are NOT wrapped
349 - # in custom ops -- their small precision differences propagate
350 - # as inputs to the GDN layers from outside.
351 - torch.ops.vllm.olmo_hybrid_gdn_full_forward(
352 - hidden_states,
353 - output,
354 - self.prefix,
355 - )
356 -
357 - def _full_forward(
358 - self,
359 - hidden_states: torch.Tensor,
360 - output: torch.Tensor,
361 - ):
362 - num_tokens = hidden_states.size(0)
363 -
364 - # ============================================================
365 - # Part 1: Input Projection (2 fused matmuls instead of 6)
366 - # ============================================================
367 - projected_qkvg, _ = self.in_proj_qkvg(hidden_states)
368 - conv_dim_sharded = (self.key_dim * 2 + self.value_dim) // self.tp_size
369 - mixed_qkv = projected_qkvg[..., :conv_dim_sharded]
370 - gate = projected_qkvg[..., conv_dim_sharded:]
371 -
372 - b, _ = self.b_proj(hidden_states)
373 - a, _ = self.a_proj(hidden_states)
374 -
375 - # ============================================================
376 - # Part 2: Core Attention
377 - # ============================================================
378 - core_attn_out = torch.zeros(
379 - (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
380 - dtype=hidden_states.dtype,
381 - device=hidden_states.device,
382 - )
383 -
384 - self._forward_core(
385 - mixed_qkv=mixed_qkv,
386 - b=b,
387 - a=a,
388 - core_attn_out=core_attn_out,
389 - )
390 -
391 - # ============================================================
392 - # Part 3: Output Projection
393 - # ============================================================
394 - gate = gate.view(num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim)
395 - core_attn_out_flat = core_attn_out.reshape(-1, core_attn_out.shape[-1])
396 - gate_flat = gate.reshape(-1, gate.shape[-1])
397 - core_attn_out_normed = self.o_norm(core_attn_out_flat, gate_flat)
398 - core_attn_out = core_attn_out_normed.view(
399 - num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim
400 - )
401 -
402 - core_attn_out = rearrange(core_attn_out, "l h d -> l (h d)")
403 - output[:num_tokens], _ = self.o_proj(core_attn_out)
404 -
405 - def _forward_core(
406 - self,
407 - mixed_qkv: torch.Tensor,
408 - b: torch.Tensor,
409 - a: torch.Tensor,
410 - core_attn_out: torch.Tensor,
411 - ):
412 - """
413 - Core attention computation (called by custom op).
414 - """
415 - forward_context = get_forward_context()
416 - attn_metadata: AttentionMetadata = forward_context.attn_metadata
417 -
418 - if attn_metadata is None:
419 - # V1 profile run
420 - return
421 -
422 - assert isinstance(attn_metadata, dict)
423 - attn_metadata = attn_metadata[self.prefix]
424 - assert isinstance(attn_metadata, GDNAttentionMetadata)
425 - has_initial_state = attn_metadata.has_initial_state
426 - spec_query_start_loc = attn_metadata.spec_query_start_loc
427 - non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
428 - spec_sequence_masks = attn_metadata.spec_sequence_masks
429 - spec_token_indx = attn_metadata.spec_token_indx
430 - non_spec_token_indx = attn_metadata.non_spec_token_indx
431 - spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor
432 - non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor
433 - self_kv_cache = self.kv_cache
434 - # conv_state must be (..., dim, width-1) for the conv kernels.
435 - # DS layout stores it that way directly; SD layout needs a transpose.
436 - conv_state = (
437 - self_kv_cache[0]
438 - if is_conv_state_dim_first()
439 - else self_kv_cache[0].transpose(-1, -2)
440 - )
441 - ssm_state = self_kv_cache[1]
442 - num_actual_tokens = attn_metadata.num_actual_tokens
443 - num_accepted_tokens = attn_metadata.num_accepted_tokens
444 -
445 - mixed_qkv = mixed_qkv[:num_actual_tokens]
446 - b = b[:num_actual_tokens]
447 - a = a[:num_actual_tokens]
448 -
449 - conv_weights = self.conv1d.weight.view(
450 - self.conv1d.weight.size(0), self.conv1d.weight.size(2)
451 - )
452 -
453 - if spec_sequence_masks is not None:
454 - if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:
455 - mixed_qkv_spec = mixed_qkv
456 - mixed_qkv_non_spec = None
457 - else:
458 - mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
459 - mixed_qkv_non_spec = mixed_qkv.index_select(0, non_spec_token_indx)
460 - else:
461 - mixed_qkv_spec = None
462 - mixed_qkv_non_spec = mixed_qkv
463 -
464 - if spec_sequence_masks is not None:
465 - mixed_qkv_spec = causal_conv1d_update(
466 - mixed_qkv_spec,
467 - conv_state,
468 - conv_weights,
469 - None, # no bias
470 - self.activation,
471 - conv_state_indices=spec_state_indices_tensor[:, 0][
472 - : attn_metadata.num_spec_decodes
473 - ],
474 - num_accepted_tokens=num_accepted_tokens,
475 - query_start_loc=spec_query_start_loc,
476 - max_query_len=spec_state_indices_tensor.size(-1),
477 - validate_data=False,
478 - )
479 -
480 - if attn_metadata.num_prefills > 0:
481 - mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1)
482 - mixed_qkv_non_spec = causal_conv1d_fn(
483 - mixed_qkv_non_spec_T,
484 - conv_weights,
485 - None,
486 - activation=self.activation,
487 - conv_states=conv_state,
488 - has_initial_state=has_initial_state,
489 - cache_indices=non_spec_state_indices_tensor,
490 - query_start_loc=non_spec_query_start_loc,
491 - metadata=attn_metadata,
492 - ).transpose(0, 1)
493 - elif attn_metadata.num_decodes > 0:
494 - mixed_qkv_non_spec = causal_conv1d_update(
495 - mixed_qkv_non_spec,
496 - conv_state,
497 - conv_weights,
498 - None,
499 - self.activation,
500 - conv_state_indices=non_spec_state_indices_tensor[
501 - : attn_metadata.num_decodes
502 - ],
503 - validate_data=True,
504 - )
505 - else:
506 - mixed_qkv_non_spec = None
507 -
508 - query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)
509 - query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv(
510 - mixed_qkv_non_spec
511 - )
512 -
513 - g, beta = fused_olmo_hybrid_gdn_gating(
514 - self.A_log, a, b, self.dt_bias, self.allow_neg_eigval
515 - )
516 -
517 - if spec_sequence_masks is not None:
518 - if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:
519 - g_spec = g
520 - beta_spec = beta
521 - g_non_spec = None
522 - beta_non_spec = None
523 - else:
524 - g_spec = g.index_select(1, spec_token_indx)
525 - beta_spec = beta.index_select(1, spec_token_indx)
526 - g_non_spec = g.index_select(1, non_spec_token_indx)
527 - beta_non_spec = beta.index_select(1, non_spec_token_indx)
528 - else:
529 - g_spec = None
530 - beta_spec = None
531 - g_non_spec = g
532 - beta_non_spec = beta
533 -
534 - if spec_sequence_masks is not None:
535 - core_attn_out_spec, last_recurrent_state = fused_recurrent_gated_delta_rule(
536 - q=query_spec,
537 - k=key_spec,
538 - v=value_spec,
539 - g=g_spec,
540 - beta=beta_spec,
541 - initial_state=ssm_state,
542 - inplace_final_state=True,
543 - cu_seqlens=spec_query_start_loc[: attn_metadata.num_spec_decodes + 1],
544 - ssm_state_indices=spec_state_indices_tensor,
545 - num_accepted_tokens=num_accepted_tokens,
546 - use_qk_l2norm_in_kernel=True,
547 - )
548 - else:
549 - core_attn_out_spec, last_recurrent_state = None, None
550 -
551 - if attn_metadata.num_prefills > 0:
552 - initial_state = ssm_state[non_spec_state_indices_tensor].contiguous()
553 - initial_state[~has_initial_state, ...] = 0
554 - (
555 - core_attn_out_non_spec,
556 - last_recurrent_state,
557 - ) = chunk_gated_delta_rule(
558 - q=query_non_spec,
559 - k=key_non_spec,
560 - v=value_non_spec,
561 - g=g_non_spec,
562 - beta=beta_non_spec,
563 - initial_state=initial_state,
564 - output_final_state=True,
565 - cu_seqlens=non_spec_query_start_loc,
566 - use_qk_l2norm_in_kernel=True,
567 - )
568 - ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to(
569 - ssm_state.dtype
570 - )
571 - elif attn_metadata.num_decodes > 0:
572 - core_attn_out_non_spec, last_recurrent_state = (
573 - fused_recurrent_gated_delta_rule(
574 - q=query_non_spec,
575 - k=key_non_spec,
576 - v=value_non_spec,
577 - g=g_non_spec,
578 - beta=beta_non_spec,
579 - initial_state=ssm_state,
580 - inplace_final_state=True,
581 - cu_seqlens=non_spec_query_start_loc[
582 - : attn_metadata.num_decodes + 1
583 - ],
584 - ssm_state_indices=non_spec_state_indices_tensor,
585 - use_qk_l2norm_in_kernel=True,
586 - )
587 - )
588 - else:
589 - core_attn_out_non_spec, last_recurrent_state = None, None
590 -
591 - if spec_sequence_masks is not None and core_attn_out_non_spec is not None:
592 - merged_out = torch.empty(
593 - (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
594 - dtype=core_attn_out_non_spec.dtype,
595 - device=core_attn_out_non_spec.device,
596 - )
597 - merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
598 - merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
599 - core_attn_out[:num_actual_tokens] = merged_out.squeeze(0)
600 - elif spec_sequence_masks is not None:
601 - core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
602 - else:
603 - core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)
604 -
605 -
84 class OlmoHybridAttention(nn.Module):
85 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
86 super().__init__()
@@ -750,22 +228,15 @@ class OlmoHybridDecoderLayer(nn.Module):
228 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
229 super().__init__()
230 config = vllm_config.model_config.hf_config
753 - model_config = vllm_config.model_config
754 - cache_config = vllm_config.cache_config
755 - quant_config = vllm_config.quant_config
756 - speculative_config = vllm_config.speculative_config
231
232 layer_idx = extract_layer_index(prefix)
233 self.layer_type = config.layer_types[layer_idx]
234 self.layer_idx = layer_idx
235
236 if self.layer_type == "linear_attention":
763 - self.linear_attn = OlmoHybridGatedDeltaNet(
237 + self.linear_attn = OlmoHybridGatedDeltaNetAttention(
238 config,
765 - model_config=model_config,
766 - cache_config=cache_config,
767 - quant_config=quant_config,
768 - speculative_config=speculative_config,
239 + vllm_config,
240 prefix=f"{prefix}.linear_attn",
241 )
242 self.input_layernorm = RMSNorm(
@@ -1068,113 +539,3 @@ class OlmoHybridForCausalLM(
539 ),
540 )
541 return loader.load_weights(weights)
1071 -
1072 -
1073 -def olmo_hybrid_gdn_full_forward(
1074 - hidden_states: torch.Tensor,
1075 - output: torch.Tensor,
1076 - layer_name: str,
1077 -) -> None:
1078 - """Full linear attention forward wrapped as a custom op.
1079 -
1080 - Prevents inductor from compiling the projections around the GDN core,
1081 - which would introduce numerical divergence that compounds through
1082 - the recurrent state.
1083 - """
1084 - forward_context: ForwardContext = get_forward_context()
1085 - self = forward_context.no_compile_layers[layer_name]
1086 - self._full_forward(
1087 - hidden_states=hidden_states,
1088 - output=output,
1089 - )
1090 -
1091 -
1092 -def olmo_hybrid_gdn_full_forward_fake(
1093 - hidden_states: torch.Tensor,
1094 - output: torch.Tensor,
1095 - layer_name: str,
1096 -) -> None:
1097 - """Fake implementation for torch.compile."""
1098 - return
1099 -
1100 -
1101 -direct_register_custom_op(
1102 - op_name="olmo_hybrid_gdn_full_forward",
1103 - op_func=olmo_hybrid_gdn_full_forward,
1104 - mutates_args=["output"],
1105 - fake_impl=olmo_hybrid_gdn_full_forward_fake,
1106 -)
1107 -
1108 -
1109 -@triton.jit
1110 -def fused_olmo_hybrid_gdn_gating_kernel(
1111 - g,
1112 - beta_output,
1113 - A_log,
1114 - a,
1115 - b,
1116 - dt_bias,
1117 - seq_len,
1118 - allow_neg_eigval: tl.constexpr,
1119 - NUM_HEADS: tl.constexpr,
1120 - beta: tl.constexpr,
1121 - threshold: tl.constexpr,
1122 - BLK_HEADS: tl.constexpr,
1123 -):
1124 - i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2)
1125 - head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS)
1126 - off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off
1127 - mask = head_off < NUM_HEADS
1128 - blk_A_log = tl.load(A_log + head_off, mask=mask)
1129 - blk_a = tl.load(a + off, mask=mask)
1130 - blk_b = tl.load(b + off, mask=mask)
1131 - blk_bias = tl.load(dt_bias + head_off, mask=mask)
1132 -
1133 - # g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
1134 - x = blk_a.to(tl.float32) + blk_bias.to(tl.float32)
1135 - softplus_x = tl.where(
1136 - beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x
1137 - )
1138 - blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x
1139 - tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask)
1140 -
1141 - # beta = self.b_proj(hidden_states).sigmoid()
1142 - # if self.allow_neg_eigval: beta = beta * 2.0
1143 - blk_beta_output = tl.sigmoid(blk_b.to(tl.float32))
1144 - if allow_neg_eigval:
1145 - blk_beta_output = blk_beta_output * 2.0
1146 - tl.store(
1147 - beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask
1148 - )
1149 -
1150 -
1151 -def fused_olmo_hybrid_gdn_gating(
1152 - A_log: torch.Tensor,
1153 - a: torch.Tensor,
1154 - b: torch.Tensor,
1155 - dt_bias: torch.Tensor,
1156 - allow_neg_eigval: bool = False,
1157 - beta: float = 1.0,
1158 - threshold: float = 20.0,
1159 -) -> tuple[torch.Tensor, torch.Tensor]:
1160 - batch, num_heads = a.shape
1161 - seq_len = 1
1162 - grid = (batch, seq_len, triton.cdiv(num_heads, 8))
1163 - g = torch.empty(1, batch, num_heads, dtype=torch.float32, device=a.device)
1164 - beta_output = torch.empty(1, batch, num_heads, dtype=torch.float32, device=b.device)
1165 - fused_olmo_hybrid_gdn_gating_kernel[grid](
1166 - g,
1167 - beta_output,
1168 - A_log,
1169 - a,
1170 - b,
1171 - dt_bias,
1172 - seq_len,
1173 - allow_neg_eigval,
1174 - num_heads,
1175 - beta,
1176 - threshold,
1177 - 8,
1178 - num_warps=1,
1179 - )
1180 - return g, beta_output
vllm/model_executor/models/qwen3_5.py
+4 -2
@@ -40,7 +40,9 @@ from vllm.model_executor.layers.layernorm import (
40 GemmaRMSNorm as Qwen3_5RMSNorm,
41 )
42 from vllm.model_executor.layers.logits_processor import LogitsProcessor
43 -from vllm.model_executor.layers.mamba.gdn_linear_attn import GatedDeltaNetAttention
43 +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
44 + QwenGatedDeltaNetAttention,
45 +)
46 from vllm.model_executor.layers.mamba.mamba_utils import (
47 MambaStateCopyFunc,
48 MambaStateCopyFuncCalculator,
@@ -133,7 +135,7 @@ class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):
135 self.layer_idx = extract_layer_index(prefix)
136
137 if self.layer_type == "linear_attention":
136 - self.linear_attn = GatedDeltaNetAttention(
138 + self.linear_attn = QwenGatedDeltaNetAttention(
139 config=config,
140 vllm_config=vllm_config,
141 prefix=f"{prefix}.linear_attn",
vllm/model_executor/models/qwen3_next.py
+4 -2
@@ -37,7 +37,9 @@ from vllm.model_executor.layers.linear import (
37 RowParallelLinear,
38 )
39 from vllm.model_executor.layers.logits_processor import LogitsProcessor
40 -from vllm.model_executor.layers.mamba.gdn_linear_attn import GatedDeltaNetAttention
40 +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
41 + QwenGatedDeltaNetAttention,
42 +)
43 from vllm.model_executor.layers.mamba.mamba_utils import (
44 MambaStateCopyFunc,
45 MambaStateCopyFuncCalculator,
@@ -337,7 +339,7 @@ class Qwen3NextDecoderLayer(nn.Module):
339 self.layer_idx = extract_layer_index(prefix)
340
341 if self.layer_type == "linear_attention":
340 - self.linear_attn = GatedDeltaNetAttention(
342 + self.linear_attn = QwenGatedDeltaNetAttention(
343 config,
344 vllm_config=vllm_config,
345 prefix=f"{prefix}.linear_attn",