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

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

[Model] Add support for OLMo Hybrid (#32550)

Yanhong Li committed Mar 5, 2026 at 11:51 UTC a911f4dd20d0a0fcfee362f096e9c6fd23d59590

10 files changed +1520 -53

docs/models/supported_models.md

@@ -448,6 +448,7 @@ th {

448 | `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ |

449 | `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ |

450 | `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ |

451 +| `OlmoHybridForCausalLM` | OLMo Hybrid | `allenai/Olmo-Hybrid-7B` | ✅︎ | ✅︎ |

453 | `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | ✅︎ | ✅︎ |

tests/models/registry.py

@@ -420,6 +420,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

420 "OlmoForCausalLM": _HfExamplesInfo("allenai/OLMo-1B-hf"),

421 "Olmo2ForCausalLM": _HfExamplesInfo("allenai/OLMo-2-0425-1B"),

422 "Olmo3ForCausalLM": _HfExamplesInfo("allenai/Olmo-3-7B-Instruct"),

423 + "OlmoHybridForCausalLM": _HfExamplesInfo("allenai/Olmo-Hybrid-7B"),

424 "OlmoeForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924-Instruct"),

425 "OPTForCausalLM": _HfExamplesInfo(

426 "facebook/opt-125m", {"1b": "facebook/opt-iml-max-1.3b"}

vllm/config/compilation.py

@@ -666,6 +666,7 @@ class CompilationConfig:

666 "vllm::linear_attention",

667 "vllm::plamo2_mamba_mixer",

668 "vllm::gdn_attention_core",

669 + "vllm::olmo_hybrid_gdn_full_forward",

670 "vllm::kda_attention",

671 "vllm::sparse_attn_indexer",

672 "vllm::rocm_aiter_sparse_attn_indexer",

vllm/model_executor/layers/fla/ops/l2norm.py

+10 -5

@@ -76,16 +76,20 @@ def l2norm_fwd_kernel(

78 @triton.jit

79 -def l2norm_fwd_kernel2(X, Y, eps, M, N: tl.constexpr, MBLOCK: tl.constexpr):

79 +def l2norm_fwd_kernel2(

80 + X, Y, eps, M, N: tl.constexpr, BD: tl.constexpr, MBLOCK: tl.constexpr

81 +):

82 xoffset = tl.program_id(0) * MBLOCK

83 row_idx = xoffset + tl.arange(0, MBLOCK)[:, None]

84 xmask = row_idx < M

83 - rindex = tl.arange(0, N)[None, :]

84 - xs = tl.load(X + (rindex + N * row_idx), xmask).to(tl.float32)

85 - square = tl.broadcast_to(xs * xs, [MBLOCK, N])

85 + rindex = tl.arange(0, BD)[None, :]

86 + cmask = rindex < N

87 + mask = xmask & cmask

88 + xs = tl.load(X + (rindex + N * row_idx), mask, other=0.0).to(tl.float32)

89 + square = tl.broadcast_to(xs * xs, [MBLOCK, BD])

90 square_sum = tl.sum(tl.where(xmask, square, 0), 1)[:, None]

91 rsqrt = tl.rsqrt(square_sum + eps)

88 - tl.store(Y + (rindex + N * row_idx), xs * rsqrt, xmask)

92 + tl.store(Y + (rindex + N * row_idx), xs * rsqrt, mask)

95 def l2norm_fwd(

@@ -116,6 +120,7 @@ def l2norm_fwd(

120 eps,

121 T,

122 D,

123 + BD,

124 MBLOCK,

125 )

126 else:

vllm/model_executor/layers/fla/ops/layernorm_guard.py

+47 -48

@@ -250,57 +250,55 @@ def layer_norm_fwd(

250 return out, mean, rstd

251

252

253 -class LayerNormFn(torch.autograd.Function):

254 - @input_guard

255 - @staticmethod

256 - def forward(

257 - ctx,

253 +def _layer_norm_fn_impl(

254 + x,

255 + weight,

256 + bias,

257 + z=None,

258 + eps=1e-6,

259 + group_size=None,

260 + norm_before_gate=True,

261 + is_rms_norm=False,

262 + activation: str = "swish",

263 +):

264 + """Triton layer/RMS norm with optional gating.

265 +

266 + If z is not None, computes norm(x) * silu(z) when norm_before_gate,

267 + else norm(x * silu(z)).

268 +

269 + This calls the triton kernel directly. The original code wrapped this

270 + in a torch.autograd.Function (LayerNormFn) to save tensors for a

271 + backward pass, but vLLM is inference-only so there is no backward pass.

272 + The autograd wrapper also prevented torch.compile/dynamo from tracing

273 + through the function due to its @staticmethod forward.

274 + """

275 + x_shape_og = x.shape

276 + x = x.reshape(-1, x.shape[-1])

277 + if x.stride(-1) != 1:

278 + x = x.contiguous()

279 + if z is not None:

280 + assert z.shape == x_shape_og

281 + z = z.reshape(-1, z.shape[-1])

282 + if z.stride(-1) != 1:

283 + z = z.contiguous()

284 + weight = weight.contiguous()

285 + if bias is not None:

286 + bias = bias.contiguous()

287 + y, _, _ = layer_norm_fwd(

288 x,

289 weight,

290 bias,

261 - z=None,

262 - eps=1e-6,

263 - group_size=None,

264 - norm_before_gate=True,

265 - is_rms_norm=False,

266 - activation: str = "swish",

267 - ):

268 - """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))"""

269 -

270 - x_shape_og = x.shape

271 - # reshape input data into 2D tensor

272 - x = x.reshape(-1, x.shape[-1])

273 - if x.stride(-1) != 1:

274 - x = x.contiguous()

275 - if z is not None:

276 - assert z.shape == x_shape_og

277 - z = z.reshape(-1, z.shape[-1])

278 - if z.stride(-1) != 1:

279 - z = z.contiguous()

280 - weight = weight.contiguous()

281 - if bias is not None:

282 - bias = bias.contiguous()

283 - y, mean, rstd = layer_norm_fwd(

284 - x,

285 - weight,

286 - bias,

287 - eps,

288 - z=z,

289 - group_size=group_size,

290 - norm_before_gate=norm_before_gate,

291 - is_rms_norm=is_rms_norm,

292 - activation=activation,

293 - )

294 - ctx.save_for_backward(x, weight, bias, mean, rstd, z)

295 - ctx.x_shape_og = x_shape_og

296 - ctx.eps = eps

297 - ctx.group_size = group_size

298 - ctx.norm_before_gate = norm_before_gate

299 - ctx.is_rms_norm = is_rms_norm

300 - ctx.activation = activation

301 - return y.reshape(x_shape_og)

291 + eps,

292 + z=z,

293 + group_size=group_size,

294 + norm_before_gate=norm_before_gate,

295 + is_rms_norm=is_rms_norm,

296 + activation=activation,

297 + )

298 + return y.reshape(x_shape_og)

299

300

301 +@input_guard

302 def layernorm_fn(

303 x,

304 weight,

@@ -312,11 +310,12 @@ def layernorm_fn(

310 is_rms_norm=False,

311 activation: str = "swish",

312 ):

315 - return LayerNormFn.apply(

313 + return _layer_norm_fn_impl(

314 x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm, activation

315 )

316

317

318 +@input_guard

319 def rmsnorm_fn(

320 x,

321 weight,

@@ -327,7 +326,7 @@ def rmsnorm_fn(

326 norm_before_gate=True,

327 activation: str = "swish",

328 ):

330 - return LayerNormFn.apply(

329 + return _layer_norm_fn_impl(

330 x, weight, bias, z, eps, group_size, norm_before_gate, True, activation

331 )

332

vllm/model_executor/models/olmo_hybrid.py new

+1172

@@ -0,0 +1,1172 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +# Adapted from:

5 +# https://github.com/huggingface/transformers/blob/main/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py

7 +#

8 +# This code combines OLMo2/OLMo3 attention with Gated DeltaNet linear attention

9 +# for the OLMo Hybrid architecture.

10 +#

11 +# Licensed under the Apache License, Version 2.0 (the "License");

12 +# you may not use this file except in compliance with the License.

13 +# You may obtain a copy of the License at

14 +#

15 +# http://www.apache.org/licenses/LICENSE-2.0

16 +#

17 +# Unless required by applicable law or agreed to in writing, software

18 +# distributed under the License is distributed on an "AS IS" BASIS,

19 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

20 +# See the License for the specific language governing permissions and

21 +# limitations under the License.

22 +"""Inference-only OLMo Hybrid model compatible with HuggingFace weights."""

23 +

24 +from collections.abc import Iterable

25 +from functools import partial

26 +from itertools import islice

27 +

28 +import torch

29 +from einops import rearrange

30 +from torch import nn

31 +from transformers.activations import ACT2FN

32 +

33 +from vllm.compilation.decorators import support_torch_compile

34 +from vllm.config import (

35 + CacheConfig,

36 + ModelConfig,

37 + SpeculativeConfig,

38 + VllmConfig,

39 + get_current_vllm_config,

40 +)

41 +from vllm.distributed import (

42 + divide,

43 + get_pp_group,

44 + get_tensor_model_parallel_rank,

45 + get_tensor_model_parallel_world_size,

46 + tensor_model_parallel_all_gather,

47 +)

48 +from vllm.distributed.utils import split_tensor_along_last_dim

49 +from vllm.forward_context import ForwardContext, get_forward_context

50 +from vllm.logger import init_logger

51 +from vllm.model_executor.layers.activation import SiluAndMul

52 +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

58 +from vllm.model_executor.layers.linear import (

59 + ColumnParallelLinear,

60 + MergedColumnParallelLinear,

61 + QKVParallelLinear,

62 + RowParallelLinear,

63 +)

64 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

65 +from vllm.model_executor.layers.mamba.abstract import MambaBase

66 +from vllm.model_executor.layers.mamba.mamba_utils import (

67 + MambaStateCopyFunc,

68 + MambaStateCopyFuncCalculator,

69 + MambaStateDtypeCalculator,

70 + MambaStateShapeCalculator,

71 +)

72 +from vllm.model_executor.layers.mamba.ops.causal_conv1d import (

73 + causal_conv1d_fn,

74 + causal_conv1d_update,

75 +)

76 +from vllm.model_executor.layers.quantization import QuantizationConfig

77 +from vllm.model_executor.layers.rotary_embedding import get_rope

78 +from vllm.model_executor.layers.vocab_parallel_embedding import (

79 + ParallelLMHead,

80 + VocabParallelEmbedding,

81 +)

82 +from vllm.model_executor.model_loader.weight_utils import (

83 + default_weight_loader,

84 + sharded_weight_loader,

85 +)

86 +from vllm.model_executor.utils import set_weight_attrs

87 +from vllm.platforms import current_platform

88 +from vllm.sequence import IntermediateTensors

89 +from vllm.triton_utils import tl, triton

90 +from vllm.triton_utils.allocation import set_triton_allocator

91 +from vllm.utils.torch_utils import direct_register_custom_op

92 +from vllm.v1.attention.backend import AttentionMetadata

93 +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata

94 +

95 +from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP

96 +from .utils import (

97 + AutoWeightsLoader,

98 + extract_layer_index,

99 + is_pp_missing_parameter,

100 + make_empty_intermediate_tensors_factory,

101 + make_layers,

102 + maybe_prefix,

103 +)

104 +

105 +logger = init_logger(__name__)

106 +

107 +

108 +def _make_fused_conv1d_weight_loader(dims, tp_size, tp_rank):

109 + """Weight loader for loading separate HF conv weights into a fused conv1d.

110 +

111 + dims: list of original (un-sharded) dims per section,

112 + e.g. [key_dim, key_dim, value_dim]

113 + """

114 + sharded_dims = [d // tp_size for d in dims]

115 +

116 + def weight_loader(param, loaded_weight, loaded_shard_id=None):

117 + if loaded_weight.dim() == 2:

118 + loaded_weight = loaded_weight.unsqueeze(1)

119 + dim = dims[loaded_shard_id]

120 + shard_size = dim // tp_size

121 + tp_start = tp_rank * shard_size

122 + sharded_weight = loaded_weight[tp_start : tp_start + shard_size]

123 + offset = sum(sharded_dims[:loaded_shard_id])

124 + param.data[offset : offset + shard_size].copy_(sharded_weight)

125 +

126 + return weight_loader

127 +

128 +

129 +class OlmoHybridGatedDeltaNet(nn.Module, MambaBase):

130 + """

131 + Gated DeltaNet linear attention layer for OLMo Hybrid.

132 +

133 + This implements the linear attention mechanism that replaces sliding window

134 + attention in the hybrid architecture.

135 + """

136 +

137 + @property

138 + def mamba_type(self) -> str:

139 + return "gdn_attention"

140 +

141 + def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:

142 + return MambaStateDtypeCalculator.gated_delta_net_state_dtype(

143 + self.model_config.dtype,

144 + self.cache_config.mamba_cache_dtype,

145 + self.cache_config.mamba_ssm_cache_dtype,

146 + )

147 +

148 + def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:

149 + return MambaStateShapeCalculator.gated_delta_net_state_shape(

150 + self.tp_size,

151 + self.num_k_heads,

152 + self.num_v_heads,

153 + self.head_k_dim,

154 + self.head_v_dim,

155 + self.conv_kernel_size,

156 + self.num_spec,

157 + )

158 +

159 + def __init__(

160 + self,

161 + config,

162 + model_config: ModelConfig | None = None,

163 + cache_config: CacheConfig | None = None,

164 + quant_config: QuantizationConfig | None = None,

165 + speculative_config: SpeculativeConfig | None = None,

166 + prefix: str = "",

167 + ) -> None:

168 + super().__init__()

169 + self.tp_size = get_tensor_model_parallel_world_size()

170 + self.tp_rank = get_tensor_model_parallel_rank()

171 + self.hidden_size = config.hidden_size

172 + self.num_v_heads = config.linear_num_value_heads

173 + self.num_k_heads = config.linear_num_key_heads

174 + self.head_k_dim = config.linear_key_head_dim

175 + self.head_v_dim = config.linear_value_head_dim

176 + self.key_dim = self.head_k_dim * self.num_k_heads

177 + self.value_dim = self.head_v_dim * self.num_v_heads

178 +

179 + self.conv_kernel_size = config.linear_conv_kernel_dim

180 + self.layer_idx = extract_layer_index(prefix)

181 + self.activation = config.hidden_act

182 + self.act = ACT2FN[config.hidden_act]

183 + self.layer_norm_epsilon = config.rms_norm_eps

184 + assert getattr(config, "linear_use_gate", True), (

185 + "OlmoHybridGatedDeltaNet requires linear_use_gate=True"

186 + )

187 + self.allow_neg_eigval = getattr(config, "linear_allow_neg_eigval", False)

188 + self.prefix = prefix

189 +

190 + self.config = config

191 + self.model_config = model_config

192 + self.cache_config = cache_config

193 + self.quant_config = quant_config

194 + self.speculative_config = speculative_config

195 + self.num_spec = (

196 + self.speculative_config.num_speculative_tokens

197 + if self.speculative_config

198 + else 0

199 + )

200 +

201 + # Fused QKVG projection: 1 matmul instead of 4

202 + self.in_proj_qkvg = MergedColumnParallelLinear(

203 + input_size=self.hidden_size,

204 + output_sizes=[self.key_dim, self.key_dim, self.value_dim, self.value_dim],

205 + bias=False,

206 + quant_config=quant_config,

207 + prefix=f"{prefix}.in_proj_qkvg",

208 + )

209 +

210 + # Separate B and A projections to preserve numerical precision.

211 + # Fusing these into one matmul changes FP accumulation order for the

212 + # gating scalars, which compounds through the GDN recurrent state.

213 + self.b_proj = ColumnParallelLinear(

214 + input_size=self.hidden_size,

215 + output_size=self.num_v_heads,

216 + bias=False,

217 + quant_config=quant_config,

218 + prefix=f"{prefix}.b_proj",

219 + )

220 + self.a_proj = ColumnParallelLinear(

221 + input_size=self.hidden_size,

222 + output_size=self.num_v_heads,

223 + bias=False,

224 + quant_config=quant_config,

225 + prefix=f"{prefix}.a_proj",

226 + )

227 +

228 + # Fused conv1d: single parameter instead of 3

229 + self.conv_dim = self.key_dim * 2 + self.value_dim

230 + self.conv1d = ColumnParallelLinear(

231 + input_size=self.conv_kernel_size,

232 + output_size=self.conv_dim,

233 + bias=False,

234 + prefix=f"{prefix}.conv1d",

235 + )

236 + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)

237 + delattr(self.conv1d.weight, "weight_loader")

238 + set_weight_attrs(

239 + self.conv1d.weight,

240 + {

241 + "weight_loader": _make_fused_conv1d_weight_loader(

242 + [self.key_dim, self.key_dim, self.value_dim],

243 + self.tp_size,

244 + self.tp_rank,

245 + )

246 + },

247 + )

248 +

249 + self.dt_bias = nn.Parameter(

250 + torch.ones(self.num_v_heads // self.tp_size),

251 + )

252 + self.A_log = nn.Parameter(

253 + torch.empty(

254 + divide(self.num_v_heads, self.tp_size),

255 + )

256 + )

257 +

258 + set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})

259 + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})

260 +

261 + # use eps=1e-5 to match FLA's FusedRMSNormGated

262 + self.o_norm = RMSNormGated(

263 + self.head_v_dim,

264 + eps=1e-5,

265 + group_size=None,

266 + norm_before_gate=True,

267 + device=current_platform.current_device(),

268 + dtype=config.torch_dtype if hasattr(config, "torch_dtype") else None,

269 + )

270 +

271 + self.o_proj = RowParallelLinear(

272 + self.value_dim,

273 + self.hidden_size,

274 + bias=False,

275 + input_is_parallel=True,

276 + quant_config=quant_config,

277 + prefix=f"{prefix}.o_proj",

278 + )

279 +

280 + # FLA triton kernels need a PyTorch-backed allocator for scratch

281 + # memory (required by triton >= 3.x autotuner). Set once at init.

282 + set_triton_allocator(current_platform.current_device())

283 +

284 + compilation_config = get_current_vllm_config().compilation_config

285 + if prefix in compilation_config.static_forward_context:

286 + raise ValueError(f"Duplicate layer name: {prefix}")

287 + compilation_config.static_forward_context[prefix] = self

288 +

289 + def rearrange_mixed_qkv(self, mixed_qkv):

290 + if mixed_qkv is None:

291 + return None, None, None

292 + query, key, value = torch.split(

293 + mixed_qkv,

294 + [

295 + self.key_dim // self.tp_size,

296 + self.key_dim // self.tp_size,

297 + self.value_dim // self.tp_size,

298 + ],

299 + dim=-1,

300 + )

301 +

302 + num_k_heads = self.num_k_heads // self.tp_size

303 + num_v_heads = self.num_v_heads // self.tp_size

304 +

305 + query = rearrange(query, "l (h d) -> 1 l h d", h=num_k_heads, d=self.head_k_dim)

306 + key = rearrange(key, "l (h d) -> 1 l h d", h=num_k_heads, d=self.head_k_dim)

307 + value = rearrange(value, "l (h d) -> 1 l h d", h=num_v_heads, d=self.head_v_dim)

308 +

309 + # GQA expansion if needed

310 + if num_v_heads > num_k_heads:

311 + expand_ratio = num_v_heads // num_k_heads

312 + query = query.unsqueeze(3).expand(-1, -1, -1, expand_ratio, -1)

313 + query = query.reshape(1, query.shape[1], num_v_heads, self.head_k_dim)

314 + key = key.unsqueeze(3).expand(-1, -1, -1, expand_ratio, -1)

315 + key = key.reshape(1, key.shape[1], num_v_heads, self.head_k_dim)

316 +

317 + return query.contiguous(), key.contiguous(), value.contiguous()

318 +

319 + def forward(

320 + self,

321 + hidden_states: torch.Tensor,

322 + output: torch.Tensor,

323 + ):

324 + # NOTE: We wrap the ENTIRE linear attention forward (projections +

325 + # core recurrence + output norm + output projection) in a single

326 + # custom op, rather than just wrapping the recurrent core like

327 + # other GDN models (e.g. Qwen3Next) do.

328 + #

329 + # Why: torch.compile with inductor generates fused kernels for

330 + # matmuls and pointwise ops. These fused kernels can differ in

331 + # floating-point accumulation order from eager-mode cuBLAS,

332 + # introducing small numerical differences (~1e-7 per op). For

333 + # standard transformer attention this is harmless because each

334 + # position is computed independently. But for the GDN recurrent

335 + # state, these tiny input differences compound at every timestep

336 + # across the full sequence length, causing severe logprob

337 + # divergence (e.g. ~15% top-1 agreement with eager baseline).

338 + #

339 + # By making the full forward opaque to inductor, the projections

340 + # and output norm run with eager-mode kernels (cuBLAS, triton),

341 + # preserving numerical consistency. The tradeoff is reduced

342 + # compilation speedup (~1.5x vs ~3x), but logprob agreement

343 + # improves from ~15% to ~83% top-1 vs eager.

344 + #

345 + # The remaining ~17% divergence comes from inductor compiling

346 + # the MLP and transformer attention layers that are NOT wrapped

347 + # in custom ops -- their small precision differences propagate

348 + # as inputs to the GDN layers from outside.

349 + torch.ops.vllm.olmo_hybrid_gdn_full_forward(

350 + hidden_states,

351 + output,

352 + self.prefix,

353 + )

354 +

355 + def _full_forward(

356 + self,

357 + hidden_states: torch.Tensor,

358 + output: torch.Tensor,

359 + ):

360 + num_tokens = hidden_states.size(0)

361 +

362 + # ============================================================

363 + # Part 1: Input Projection (2 fused matmuls instead of 6)

364 + # ============================================================

365 + projected_qkvg, _ = self.in_proj_qkvg(hidden_states)

366 + conv_dim_sharded = (self.key_dim * 2 + self.value_dim) // self.tp_size

367 + mixed_qkv = projected_qkvg[..., :conv_dim_sharded]

368 + gate = projected_qkvg[..., conv_dim_sharded:]

369 +

370 + b, _ = self.b_proj(hidden_states)

371 + a, _ = self.a_proj(hidden_states)

372 +

373 + # ============================================================

374 + # Part 2: Core Attention

375 + # ============================================================

376 + core_attn_out = torch.zeros(

377 + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),

378 + dtype=hidden_states.dtype,

379 + device=hidden_states.device,

380 + )

381 +

382 + self._forward_core(

383 + mixed_qkv=mixed_qkv,

384 + b=b,

385 + a=a,

386 + core_attn_out=core_attn_out,

387 + )

388 +

389 + # ============================================================

390 + # Part 3: Output Projection

391 + # ============================================================

392 + gate = gate.view(num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim)

393 + core_attn_out_flat = core_attn_out.reshape(-1, core_attn_out.shape[-1])

394 + gate_flat = gate.reshape(-1, gate.shape[-1])

395 + core_attn_out_normed = self.o_norm(core_attn_out_flat, gate_flat)

396 + core_attn_out = core_attn_out_normed.view(

397 + num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim

398 + )

399 +

400 + core_attn_out = rearrange(core_attn_out, "l h d -> l (h d)")

401 + output[:num_tokens], _ = self.o_proj(core_attn_out)

402 +

403 + def _forward_core(

404 + self,

405 + mixed_qkv: torch.Tensor,

406 + b: torch.Tensor,

407 + a: torch.Tensor,

408 + core_attn_out: torch.Tensor,

409 + ):

410 + """

411 + Core attention computation (called by custom op).

412 + """

413 + forward_context = get_forward_context()

414 + attn_metadata: AttentionMetadata = forward_context.attn_metadata

415 +

416 + if attn_metadata is None:

417 + # V1 profile run

418 + return

419 +

420 + assert isinstance(attn_metadata, dict)

421 + attn_metadata = attn_metadata[self.prefix]

422 + assert isinstance(attn_metadata, GDNAttentionMetadata)

423 + has_initial_state = attn_metadata.has_initial_state

424 + spec_query_start_loc = attn_metadata.spec_query_start_loc

425 + non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc

426 + spec_sequence_masks = attn_metadata.spec_sequence_masks

427 + spec_token_indx = attn_metadata.spec_token_indx

428 + non_spec_token_indx = attn_metadata.non_spec_token_indx

429 + spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor

430 + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor

431 + self_kv_cache = self.kv_cache[forward_context.virtual_engine]

432 + conv_state = self_kv_cache[0].transpose(-1, -2)

433 + ssm_state = self_kv_cache[1]

434 + num_actual_tokens = attn_metadata.num_actual_tokens

435 + num_accepted_tokens = attn_metadata.num_accepted_tokens

436 +

437 + mixed_qkv = mixed_qkv[:num_actual_tokens]

438 + b = b[:num_actual_tokens]

439 + a = a[:num_actual_tokens]

440 +

441 + conv_weights = self.conv1d.weight.view(

442 + self.conv1d.weight.size(0), self.conv1d.weight.size(2)

443 + )

444 +

445 + if spec_sequence_masks is not None:

446 + if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:

447 + mixed_qkv_spec = mixed_qkv

448 + mixed_qkv_non_spec = None

449 + else:

450 + mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)

451 + mixed_qkv_non_spec = mixed_qkv.index_select(0, non_spec_token_indx)

452 + else:

453 + mixed_qkv_spec = None

454 + mixed_qkv_non_spec = mixed_qkv

455 +

456 + if spec_sequence_masks is not None:

457 + mixed_qkv_spec = causal_conv1d_update(

458 + mixed_qkv_spec,

459 + conv_state,

460 + conv_weights,

461 + None, # no bias

462 + self.activation,

463 + conv_state_indices=spec_state_indices_tensor[:, 0][

464 + : attn_metadata.num_spec_decodes

465 + ],

466 + num_accepted_tokens=num_accepted_tokens,

467 + query_start_loc=spec_query_start_loc,

468 + max_query_len=spec_state_indices_tensor.size(-1),

469 + validate_data=False,

470 + )

471 +

472 + if attn_metadata.num_prefills > 0:

473 + mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1)

474 + mixed_qkv_non_spec = causal_conv1d_fn(

475 + mixed_qkv_non_spec_T,

476 + conv_weights,

477 + None,

478 + activation=self.activation,

479 + conv_states=conv_state,

480 + has_initial_state=has_initial_state,

481 + cache_indices=non_spec_state_indices_tensor,

482 + query_start_loc=non_spec_query_start_loc,

483 + metadata=attn_metadata,

484 + ).transpose(0, 1)

485 + elif attn_metadata.num_decodes > 0:

486 + mixed_qkv_non_spec = causal_conv1d_update(

487 + mixed_qkv_non_spec,

488 + conv_state,

489 + conv_weights,

490 + None,

491 + self.activation,

492 + conv_state_indices=non_spec_state_indices_tensor[

493 + : attn_metadata.num_decodes

494 + ],

495 + validate_data=True,

496 + )

497 + else:

498 + mixed_qkv_non_spec = None

499 +

500 + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)

501 + query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv(

502 + mixed_qkv_non_spec

503 + )

504 +

505 + g, beta = fused_olmo_hybrid_gdn_gating(

506 + self.A_log, a, b, self.dt_bias, self.allow_neg_eigval

507 + )

508 +

509 + if spec_sequence_masks is not None:

510 + if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:

511 + g_spec = g

512 + beta_spec = beta

513 + g_non_spec = None

514 + beta_non_spec = None

515 + else:

516 + g_spec = g.index_select(1, spec_token_indx)

517 + beta_spec = beta.index_select(1, spec_token_indx)

518 + g_non_spec = g.index_select(1, non_spec_token_indx)

519 + beta_non_spec = beta.index_select(1, non_spec_token_indx)

520 + else:

521 + g_spec = None

522 + beta_spec = None

523 + g_non_spec = g

524 + beta_non_spec = beta

525 +

526 + if spec_sequence_masks is not None:

527 + core_attn_out_spec, last_recurrent_state = fused_recurrent_gated_delta_rule(

528 + q=query_spec,

529 + k=key_spec,

530 + v=value_spec,

531 + g=g_spec,

532 + beta=beta_spec,

533 + initial_state=ssm_state,

534 + inplace_final_state=True,

535 + cu_seqlens=spec_query_start_loc[: attn_metadata.num_spec_decodes + 1],

536 + ssm_state_indices=spec_state_indices_tensor,

537 + num_accepted_tokens=num_accepted_tokens,

538 + use_qk_l2norm_in_kernel=True,

539 + )

540 + else:

541 + core_attn_out_spec, last_recurrent_state = None, None

542 +

543 + if attn_metadata.num_prefills > 0:

544 + initial_state = ssm_state[non_spec_state_indices_tensor].contiguous()

545 + initial_state[~has_initial_state, ...] = 0

546 + (

547 + core_attn_out_non_spec,

548 + last_recurrent_state,

549 + ) = chunk_gated_delta_rule(

550 + q=query_non_spec,

551 + k=key_non_spec,

552 + v=value_non_spec,

553 + g=g_non_spec,

554 + beta=beta_non_spec,

555 + initial_state=initial_state,

556 + output_final_state=True,

557 + cu_seqlens=non_spec_query_start_loc,

558 + use_qk_l2norm_in_kernel=True,

559 + )

560 + ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to(

561 + ssm_state.dtype

562 + )

563 + elif attn_metadata.num_decodes > 0:

564 + core_attn_out_non_spec, last_recurrent_state = (

565 + fused_recurrent_gated_delta_rule(

566 + q=query_non_spec,

567 + k=key_non_spec,

568 + v=value_non_spec,

569 + g=g_non_spec,

570 + beta=beta_non_spec,

571 + initial_state=ssm_state,

572 + inplace_final_state=True,

573 + cu_seqlens=non_spec_query_start_loc[

574 + : attn_metadata.num_decodes + 1

575 + ],

576 + ssm_state_indices=non_spec_state_indices_tensor,

577 + use_qk_l2norm_in_kernel=True,

578 + )

579 + )

580 + else:

581 + core_attn_out_non_spec, last_recurrent_state = None, None

582 +

583 + if spec_sequence_masks is not None and core_attn_out_non_spec is not None:

584 + merged_out = torch.empty(

585 + (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),

586 + dtype=core_attn_out_non_spec.dtype,

587 + device=core_attn_out_non_spec.device,

588 + )

589 + merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec)

590 + merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)

591 + core_attn_out[:num_actual_tokens] = merged_out.squeeze(0)

592 + elif spec_sequence_masks is not None:

593 + core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)

594 + else:

595 + core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)

596 +

597 +

598 +class OlmoHybridAttention(nn.Module):

599 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

600 + super().__init__()

601 + self.config = vllm_config.model_config.hf_config

602 +

603 + hidden_size = self.config.hidden_size

604 + self.tp_size = get_tensor_model_parallel_world_size()

605 + self.total_num_heads = self.config.num_attention_heads

606 +

607 + assert hidden_size % self.total_num_heads == 0

608 + assert self.total_num_heads % self.tp_size == 0

609 +

610 + self.num_heads = self.total_num_heads // self.tp_size

611 + self.total_num_kv_heads = (

612 + self.config.num_key_value_heads or self.total_num_heads

613 + )

614 + if self.total_num_kv_heads >= self.tp_size:

615 + assert self.total_num_kv_heads % self.tp_size == 0

616 + else:

617 + assert self.tp_size % self.total_num_kv_heads == 0

618 +

619 + self.num_kv_heads = max(1, self.total_num_kv_heads // self.tp_size)

620 + self.head_dim = hidden_size // self.total_num_heads

621 + self.q_size = self.num_heads * self.head_dim

622 + self.kv_size = self.num_kv_heads * self.head_dim

623 + self.max_position_embeddings = self.config.max_position_embeddings

624 +

625 + self.qkv_proj = QKVParallelLinear(

626 + hidden_size,

627 + self.head_dim,

628 + self.total_num_heads,

629 + self.total_num_kv_heads,

630 + bias=False,

631 + quant_config=vllm_config.quant_config,

632 + prefix=f"{prefix}.qkv_proj",

633 + )

634 +

635 + self.tp_rank = get_tensor_model_parallel_rank()

636 +

637 + self.k_norm = RMSNorm(

638 + self.total_num_kv_heads * self.head_dim,

639 + eps=self.config.rms_norm_eps,

640 + )

641 + self.q_norm = RMSNorm(

642 + self.config.hidden_size,

643 + eps=self.config.rms_norm_eps,

644 + )

645 +

646 + self.scaling = self.head_dim**-0.5

647 +

648 + self.attn = Attention(

649 + self.num_heads,

650 + self.head_dim,

651 + self.scaling,

652 + num_kv_heads=self.num_kv_heads,

653 + cache_config=vllm_config.cache_config,

654 + quant_config=vllm_config.quant_config,

655 + prefix=f"{prefix}.attn",

656 + )

657 +

658 + rope_parameters = getattr(self.config, "rope_parameters", None)

659 + self._use_rope = (rope_parameters is not None) and (

660 + rope_parameters["rope_theta"] is not None

661 + )

662 +

663 + if self._use_rope:

664 + self.rotary_emb = get_rope(

665 + self.head_dim,

666 + max_position=self.max_position_embeddings,

667 + rope_parameters=rope_parameters,

668 + )

669 + else:

670 + self.rotary_emb = None

671 +

672 + self.o_proj = RowParallelLinear(

673 + self.total_num_heads * self.head_dim,

674 + hidden_size,

675 + bias=False,

676 + quant_config=vllm_config.quant_config,

677 + prefix=f"{prefix}.o_proj",

678 + )

679 +

680 + def _apply_qk_norm(

681 + self, q: torch.Tensor, k: torch.Tensor

682 + ) -> tuple[torch.Tensor, torch.Tensor]:

683 + if self.tp_size > 1:

684 + q = tensor_model_parallel_all_gather(q.contiguous())

685 + k = tensor_model_parallel_all_gather(k.contiguous())

686 + q = self.q_norm(q)

687 + k = self.k_norm(k)

688 + if self.tp_size > 1:

689 + splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size)

690 + q = splitter(q)[self.tp_rank]

691 + k = splitter(k)[self.tp_rank]

692 + return q, k

693 +

694 + def forward(

695 + self,

696 + positions: torch.Tensor,

697 + hidden_states: torch.Tensor,

698 + ) -> torch.Tensor:

699 + qkv, _ = self.qkv_proj(hidden_states)

700 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

701 + q, k = self._apply_qk_norm(q, k)

702 + if self._use_rope:

703 + q, k = self.rotary_emb(positions, q, k)

704 + attn_output = self.attn(q, k, v)

705 + output, _ = self.o_proj(attn_output)

706 + return output

707 +

708 +

709 +class OlmoHybridMLP(nn.Module):

710 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

711 + super().__init__()

712 + config = vllm_config.model_config.hf_config

713 + hidden_size = config.hidden_size

714 + intermediate_size = config.intermediate_size

715 +

716 + self.gate_up_proj = MergedColumnParallelLinear(

717 + hidden_size,

718 + [intermediate_size] * 2,

719 + bias=False,

720 + quant_config=vllm_config.quant_config,

721 + prefix=f"{prefix}.gate_up_proj",

722 + )

723 +

724 + self.act_fn = SiluAndMul()

725 +

726 + self.down_proj = RowParallelLinear(

727 + intermediate_size,

728 + hidden_size,

729 + bias=False,

730 + quant_config=vllm_config.quant_config,

731 + prefix=f"{prefix}.down_proj",

732 + )

733 +

734 + def forward(self, x: torch.Tensor) -> torch.Tensor:

735 + gate_up, _ = self.gate_up_proj(x)

736 + x = self.act_fn(gate_up)

737 + x, _ = self.down_proj(x)

738 + return x

739 +

740 +

741 +class OlmoHybridDecoderLayer(nn.Module):

742 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

743 + super().__init__()

744 + config = vllm_config.model_config.hf_config

745 + model_config = vllm_config.model_config

746 + cache_config = vllm_config.cache_config

747 + quant_config = vllm_config.quant_config

748 + speculative_config = vllm_config.speculative_config

749 +

750 + layer_idx = extract_layer_index(prefix)

751 + self.layer_type = config.layer_types[layer_idx]

752 + self.layer_idx = layer_idx

753 +

754 + if self.layer_type == "linear_attention":

755 + self.linear_attn = OlmoHybridGatedDeltaNet(

756 + config,

757 + model_config=model_config,

758 + cache_config=cache_config,

759 + quant_config=quant_config,

760 + speculative_config=speculative_config,

761 + prefix=f"{prefix}.linear_attn",

762 + )

763 + self.input_layernorm = RMSNorm(

764 + config.hidden_size,

765 + eps=config.rms_norm_eps,

766 + )

767 + self.post_attention_layernorm = RMSNorm(

768 + config.hidden_size,

769 + eps=config.rms_norm_eps,

770 + )

771 + else:

772 + self.self_attn = OlmoHybridAttention(

773 + vllm_config=vllm_config,

774 + prefix=f"{prefix}.self_attn",

775 + )

776 + # Attention layers use these norm names

777 + self.post_attention_layernorm = RMSNorm(

778 + config.hidden_size,

779 + eps=config.rms_norm_eps,

780 + )

781 + self.post_feedforward_layernorm = RMSNorm(

782 + config.hidden_size,

783 + eps=config.rms_norm_eps,

784 + )

785 +

786 + self.mlp = OlmoHybridMLP(

787 + vllm_config=vllm_config,

788 + prefix=f"{prefix}.mlp",

789 + )

790 +

791 + def forward(

792 + self,

793 + positions: torch.Tensor,

794 + hidden_states: torch.Tensor,

795 + ) -> torch.Tensor:

796 + if self.layer_type == "linear_attention":

797 + residual = hidden_states

798 + hidden_states = self.input_layernorm(hidden_states)

799 +

800 + attn_output = torch.empty_like(hidden_states)

801 + self.linear_attn(

802 + hidden_states=hidden_states,

803 + output=attn_output,

804 + )

805 + hidden_states = residual + attn_output

806 +

807 + residual = hidden_states

808 + hidden_states = self.post_attention_layernorm(hidden_states)

809 + hidden_states = self.mlp(hidden_states)

810 + hidden_states = residual + hidden_states

811 + else:

812 + residual = hidden_states

813 + hidden_states = self.self_attn(positions, hidden_states)

814 + hidden_states = self.post_attention_layernorm(hidden_states)

815 + hidden_states = residual + hidden_states

816 +

817 + residual = hidden_states

818 + hidden_states = self.mlp(hidden_states)

819 + hidden_states = self.post_feedforward_layernorm(hidden_states)

820 + hidden_states = residual + hidden_states

821 + return hidden_states

822 +

823 +

824 +@support_torch_compile

825 +class OlmoHybridModel(nn.Module):

826 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

827 + super().__init__()

828 + self.config = vllm_config.model_config.hf_config

829 +

830 + self.embed_tokens = VocabParallelEmbedding(

831 + self.config.vocab_size,

832 + self.config.hidden_size,

833 + prefix=f"{prefix}.embed_tokens",

834 + )

835 +

836 + self.start_layer, self.end_layer, self.layers = make_layers(

837 + self.config.num_hidden_layers,

838 + lambda prefix: OlmoHybridDecoderLayer(

839 + vllm_config=vllm_config, prefix=prefix

840 + ),

841 + prefix=f"{prefix}.layers",

842 + )

843 +

844 + self.norm = RMSNorm(

845 + self.config.hidden_size,

846 + eps=self.config.rms_norm_eps,

847 + )

848 +

849 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

850 + ["hidden_states"], self.config.hidden_size

851 + )

852 +

853 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

854 + return self.embed_tokens(input_ids)

855 +

856 + def forward(

857 + self,

858 + input_ids: torch.Tensor,

859 + positions: torch.Tensor,

860 + intermediate_tensors: IntermediateTensors | None,

861 + inputs_embeds: torch.Tensor | None = None,

862 + ) -> torch.Tensor | IntermediateTensors:

863 + if get_pp_group().is_first_rank:

864 + if inputs_embeds is not None:

865 + hidden_states = inputs_embeds

866 + else:

867 + hidden_states = self.embed_tokens(input_ids)

868 + else:

869 + assert intermediate_tensors is not None

870 + hidden_states = intermediate_tensors["hidden_states"]

871 + assert isinstance(hidden_states, torch.Tensor)

872 +

873 + for layer in islice(self.layers, self.start_layer, self.end_layer):

874 + hidden_states = layer(positions, hidden_states)

875 +

876 + if not get_pp_group().is_last_rank:

877 + return IntermediateTensors({"hidden_states": hidden_states})

878 +

879 + hidden_states = self.norm(hidden_states)

880 + return hidden_states

881 +

882 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

883 + stacked_params_mapping = [

884 + ("qkv_proj", "q_proj", "q"),

885 + ("qkv_proj", "k_proj", "k"),

886 + ("qkv_proj", "v_proj", "v"),

887 + ("gate_up_proj", "gate_proj", 0),

888 + ("gate_up_proj", "up_proj", 1),

889 + ]

890 +

891 + linear_attn_stacked_params_mapping = [

892 + ("in_proj_qkvg", "q_proj", 0),

893 + ("in_proj_qkvg", "k_proj", 1),

894 + ("in_proj_qkvg", "v_proj", 2),

895 + ("in_proj_qkvg", "g_proj", 3),

896 + ("conv1d", "q_conv1d", 0),

897 + ("conv1d", "k_conv1d", 1),

898 + ("conv1d", "v_conv1d", 2),

899 + ]

900 +

901 + params_dict = dict(self.named_parameters(remove_duplicate=False))

902 + loaded_params: set[str] = set()

903 +

904 + for name, loaded_weight in weights:

905 + if is_pp_missing_parameter(name, self):

906 + continue

907 +

908 + handled = False

909 +

910 + if "linear_attn" in name:

911 + for (

912 + param_name,

913 + weight_name,

914 + shard_id,

915 + ) in linear_attn_stacked_params_mapping:

916 + if weight_name not in name:

917 + continue

918 + mapped_name = name.replace(weight_name, param_name)

919 + if mapped_name.endswith(".bias") and (

920 + mapped_name not in params_dict

921 + ):

922 + continue

923 + if mapped_name not in params_dict:

924 + continue

925 + param = params_dict[mapped_name]

926 + weight_loader = param.weight_loader

927 + weight_loader(param, loaded_weight, shard_id)

928 + name = mapped_name

929 + handled = True

930 + break

931 + else:

932 + for param_name, weight_name, shard_id in stacked_params_mapping:

933 + if weight_name not in name:

934 + continue

935 + name = name.replace(weight_name, param_name)

936 + if name.endswith(".bias") and name not in params_dict:

937 + continue

938 + if name not in params_dict:

939 + continue

940 + param = params_dict[name]

941 + weight_loader = param.weight_loader

942 + weight_loader(param, loaded_weight, shard_id)

943 + handled = True

944 + break

945 +

946 + if not handled:

947 + if name.endswith(".bias") and name not in params_dict:

948 + continue

949 + if name not in params_dict:

950 + continue

951 + param = params_dict[name]

952 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

953 + weight_loader(param, loaded_weight)

954 + loaded_params.add(name)

955 + return loaded_params

956 +

957 +

958 +class OlmoHybridForCausalLM(

959 + nn.Module, HasInnerState, SupportsPP, SupportsLoRA, IsHybrid

960 +):

961 + packed_modules_mapping = {

962 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

963 + "gate_up_proj": ["gate_proj", "up_proj"],

964 + "in_proj_qkvg": ["q_proj", "k_proj", "v_proj", "g_proj"],

965 + }

966 +

967 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

968 + super().__init__()

969 + config = vllm_config.model_config.hf_config

970 + self.config = config

971 + self.vllm_config = vllm_config

972 + self.model_config = vllm_config.model_config

973 +

974 + self.model = OlmoHybridModel(

975 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

976 + )

977 +

978 + if config.tie_word_embeddings:

979 + self.lm_head = self.model.embed_tokens

980 + else:

981 + self.lm_head = ParallelLMHead(

982 + config.vocab_size,

983 + config.hidden_size,

984 + quant_config=vllm_config.quant_config,

985 + prefix=maybe_prefix(prefix, "lm_head"),

986 + )

987 +

988 + self.logits_processor = LogitsProcessor(config.vocab_size)

989 + self.make_empty_intermediate_tensors = (

990 + self.model.make_empty_intermediate_tensors

991 + )

992 +

993 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

994 + return self.model.embed_input_ids(input_ids)

995 +

996 + def forward(

997 + self,

998 + input_ids: torch.Tensor,

999 + positions: torch.Tensor,

1000 + intermediate_tensors: IntermediateTensors | None = None,

1001 + inputs_embeds: torch.Tensor | None = None,

1002 + ) -> torch.Tensor | IntermediateTensors:

1003 + hidden_states = self.model(

1004 + input_ids=input_ids,

1005 + positions=positions,

1006 + intermediate_tensors=intermediate_tensors,

1007 + inputs_embeds=inputs_embeds,

1008 + )

1009 + return hidden_states

1010 +

1011 + def compute_logits(

1012 + self,

1013 + hidden_states: torch.Tensor,

1014 + ) -> torch.Tensor | None:

1015 + logits = self.logits_processor(self.lm_head, hidden_states)

1016 + return logits

1017 +

1018 + @classmethod

1019 + def get_mamba_state_dtype_from_config(

1020 + cls,

1021 + vllm_config: "VllmConfig",

1022 + ) -> tuple[torch.dtype, torch.dtype]:

1023 + return MambaStateDtypeCalculator.gated_delta_net_state_dtype(

1024 + vllm_config.model_config.dtype,

1025 + vllm_config.cache_config.mamba_cache_dtype,

1026 + vllm_config.cache_config.mamba_ssm_cache_dtype,

1027 + )

1028 +

1029 + @classmethod

1030 + def get_mamba_state_shape_from_config(

1031 + cls, vllm_config: "VllmConfig"

1032 + ) -> tuple[tuple[int, int], tuple[int, int]]:

1033 + parallel_config = vllm_config.parallel_config

1034 + hf_config = vllm_config.model_config.hf_config

1035 + tp_size = parallel_config.tensor_parallel_size

1036 + num_spec = (

1037 + vllm_config.speculative_config.num_speculative_tokens

1038 + if vllm_config.speculative_config

1039 + else 0

1040 + )

1041 + return MambaStateShapeCalculator.gated_delta_net_state_shape(

1042 + tp_size,

1043 + hf_config.linear_num_key_heads,

1044 + hf_config.linear_num_value_heads,

1045 + hf_config.linear_key_head_dim,

1046 + hf_config.linear_value_head_dim,

1047 + hf_config.linear_conv_kernel_dim,

1048 + num_spec,

1049 + )

1050 +

1051 + @classmethod

1052 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:

1053 + return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()

1054 +

1055 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

1056 + loader = AutoWeightsLoader(

1057 + self,

1058 + skip_prefixes=(

1059 + ["lm_head.weight"] if self.config.tie_word_embeddings else None

1060 + ),

1061 + )

1062 + return loader.load_weights(weights)

1063 +

1064 +

1065 +def olmo_hybrid_gdn_full_forward(

1066 + hidden_states: torch.Tensor,

1067 + output: torch.Tensor,

1068 + layer_name: str,

1069 +) -> None:

1070 + """Full linear attention forward wrapped as a custom op.

1071 +

1072 + Prevents inductor from compiling the projections around the GDN core,

1073 + which would introduce numerical divergence that compounds through

1074 + the recurrent state.

1075 + """

1076 + forward_context: ForwardContext = get_forward_context()

1077 + self = forward_context.no_compile_layers[layer_name]

1078 + self._full_forward(

1079 + hidden_states=hidden_states,

1080 + output=output,

1081 + )

1082 +

1083 +

1084 +def olmo_hybrid_gdn_full_forward_fake(

1085 + hidden_states: torch.Tensor,

1086 + output: torch.Tensor,

1087 + layer_name: str,

1088 +) -> None:

1089 + """Fake implementation for torch.compile."""

1090 + return

1091 +

1092 +

1093 +direct_register_custom_op(

1094 + op_name="olmo_hybrid_gdn_full_forward",

1095 + op_func=olmo_hybrid_gdn_full_forward,

1096 + mutates_args=["output"],

1097 + fake_impl=olmo_hybrid_gdn_full_forward_fake,

1098 +)

1099 +

1100 +

1101 +@triton.jit

1102 +def fused_olmo_hybrid_gdn_gating_kernel(

1103 + g,

1104 + beta_output,

1105 + A_log,

1106 + a,

1107 + b,

1108 + dt_bias,

1109 + seq_len,

1110 + allow_neg_eigval: tl.constexpr,

1111 + NUM_HEADS: tl.constexpr,

1112 + beta: tl.constexpr,

1113 + threshold: tl.constexpr,

1114 + BLK_HEADS: tl.constexpr,

1115 +):

1116 + i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2)

1117 + head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS)

1118 + off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off

1119 + mask = head_off < NUM_HEADS

1120 + blk_A_log = tl.load(A_log + head_off, mask=mask)

1121 + blk_a = tl.load(a + off, mask=mask)

1122 + blk_b = tl.load(b + off, mask=mask)

1123 + blk_bias = tl.load(dt_bias + head_off, mask=mask)

1124 +

1125 + # g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)

1126 + x = blk_a.to(tl.float32) + blk_bias.to(tl.float32)

1127 + softplus_x = tl.where(

1128 + beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x

1129 + )

1130 + blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x

1131 + tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask)

1132 +

1133 + # beta = self.b_proj(hidden_states).sigmoid()

1134 + # if self.allow_neg_eigval: beta = beta * 2.0

1135 + blk_beta_output = tl.sigmoid(blk_b.to(tl.float32))

1136 + if allow_neg_eigval:

1137 + blk_beta_output = blk_beta_output * 2.0

1138 + tl.store(

1139 + beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask

1140 + )

1141 +

1142 +

1143 +def fused_olmo_hybrid_gdn_gating(

1144 + A_log: torch.Tensor,

1145 + a: torch.Tensor,

1146 + b: torch.Tensor,

1147 + dt_bias: torch.Tensor,

1148 + allow_neg_eigval: bool = False,

1149 + beta: float = 1.0,

1150 + threshold: float = 20.0,

1151 +) -> tuple[torch.Tensor, torch.Tensor]:

1152 + batch, num_heads = a.shape

1153 + seq_len = 1

1154 + grid = (batch, seq_len, triton.cdiv(num_heads, 8))

1155 + g = torch.empty(1, batch, num_heads, dtype=torch.float32, device=a.device)

1156 + beta_output = torch.empty(1, batch, num_heads, dtype=torch.float32, device=b.device)

1157 + fused_olmo_hybrid_gdn_gating_kernel[grid](

1158 + g,

1159 + beta_output,

1160 + A_log,

1161 + a,

1162 + b,

1163 + dt_bias,

1164 + seq_len,

1165 + allow_neg_eigval,

1166 + num_heads,

1167 + beta,

1168 + threshold,

1169 + 8,

1170 + num_warps=1,

1171 + )

1172 + return g, beta_output

vllm/model_executor/models/registry.py

@@ -171,6 +171,7 @@ _TEXT_GENERATION_MODELS = {

171 "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),

172 "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),

173 "Olmo3ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),

174 + "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"),

175 "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),

176 "OPTForCausalLM": ("opt", "OPTForCausalLM"),

177 "OrionForCausalLM": ("orion", "OrionForCausalLM"),

vllm/transformers_utils/config.py

@@ -97,6 +97,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(

97 speculators="SpeculatorsConfig",

98 nemotron="NemotronConfig",

99 olmo3="Olmo3Config",

100 + olmo_hybrid="OlmoHybridConfig",

101 ovis="OvisConfig",

102 ultravox="UltravoxConfig",

103 step3_vl="Step3VLConfig",

vllm/transformers_utils/configs/__init__.py

@@ -49,6 +49,7 @@ _CLASS_TO_MODULE: dict[str, str] = {

49 "NemotronConfig": "vllm.transformers_utils.configs.nemotron",

50 "NemotronHConfig": "vllm.transformers_utils.configs.nemotron_h",

51 "Olmo3Config": "vllm.transformers_utils.configs.olmo3",

52 + "OlmoHybridConfig": "vllm.transformers_utils.configs.olmo_hybrid",

53 "OvisConfig": "vllm.transformers_utils.configs.ovis",

54 "PixelShuffleSiglip2VisionConfig": "vllm.transformers_utils.configs.isaac",

55 "RadioConfig": "vllm.transformers_utils.configs.radio",

@@ -102,6 +103,7 @@ __all__ = [

103 "NemotronConfig",

104 "NemotronHConfig",

105 "Olmo3Config",

106 + "OlmoHybridConfig",

107 "OvisConfig",

108 "PixelShuffleSiglip2VisionConfig",

109 "RadioConfig",

vllm/transformers_utils/configs/olmo_hybrid.py new

+284

@@ -0,0 +1,284 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +

5 +from transformers.configuration_utils import PretrainedConfig, layer_type_validation

6 +

7 +

8 +class OlmoHybridConfig(PretrainedConfig):

9 + r"""

10 + Configuration class for [`OlmoHybridModel`]. It is used to

11 + instantiate an OLMo Hybrid model according to the specified

12 + arguments, defining the model architecture. Instantiating a

13 + configuration with the defaults will yield a similar

14 + configuration to that of the

15 + [allenai/Olmo-Hybrid-7B](https://huggingface.co/allenai/Olmo-Hybrid-7B)

16 + model.

17 +

18 + Configuration objects inherit from [`PreTrainedConfig`] and

19 + can be used to control the model outputs. Read the

20 + documentation from [`PreTrainedConfig`] for more information.

21 +

22 + Args:

23 + vocab_size (`int`, *optional*, defaults to 100352):

24 + Vocabulary size of the OlmoHybrid model. Defines

25 + the number of different tokens that can be

26 + represented by the `inputs_ids` passed when

27 + calling [`OlmoHybridModel`].

28 + hidden_size (`int`, *optional*, defaults to 3840):

29 + Dimension of the hidden representations.

30 + intermediate_size (`int`, *optional*,

31 + defaults to 11008):

32 + Dimension of the MLP representations.

33 + num_hidden_layers (`int`, *optional*,

34 + defaults to 32):

35 + Number of hidden layers in the Transformer

36 + decoder.

37 + num_attention_heads (`int`, *optional*,

38 + defaults to 30):

39 + Number of attention heads for each attention

40 + layer in the Transformer decoder.

41 + num_key_value_heads (`int`, *optional*):

42 + This is the number of key_value heads that

43 + should be used to implement Grouped Query

44 + Attention. If

45 + `num_key_value_heads=num_attention_heads`,

46 + the model will use Multi Head Attention (MHA),

47 + if `num_key_value_heads=1` the model will use

48 + Multi Query Attention (MQA) otherwise GQA is

49 + used. When converting a multi-head checkpoint

50 + to a GQA checkpoint, each group key and value

51 + head should be constructed by meanpooling all

52 + the original heads within that group. For more

53 + details, check out

54 + [this paper](https://huggingface.co/papers/2305.13245).

55 + If it is not specified, will default to

56 + `num_attention_heads`.

57 + hidden_act (`str` or `function`, *optional*,

58 + defaults to `"silu"`):

59 + The non-linear activation function (function

60 + or string) in the decoder.

61 + max_position_embeddings (`int`, *optional*,

62 + defaults to 65536):

63 + The maximum sequence length that this model

64 + might ever be used with.

65 + initializer_range (`float`, *optional*,

66 + defaults to 0.02):

67 + The standard deviation of the

68 + truncated_normal_initializer for initializing

69 + all weight matrices.

70 + use_cache (`bool`, *optional*, defaults to `True`):

71 + Whether or not the model should return the last

72 + key/values attentions (not used by all models).

73 + Only relevant if `config.is_decoder=True`.

74 + pad_token_id (`int`, *optional*,

75 + defaults to 100277):

76 + Padding token id.

77 + bos_token_id (`int`, *optional*):

78 + Beginning of stream token id.

79 + eos_token_id (`int`, *optional*,

80 + defaults to 100257):

81 + End of stream token id.

82 + tie_word_embeddings (`bool`, *optional*,

83 + defaults to `False`):

84 + Whether to tie weight embeddings.

85 + rope_parameters (`RopeParameters`, *optional*):

86 + Dictionary containing the configuration

87 + parameters for the RoPE embeddings. Can be

88 + `None` to disable RoPE.

89 + attention_bias (`bool`, *optional*,

90 + defaults to `False`):

91 + Whether to use a bias in the query, key, value

92 + and output projection layers during

93 + self-attention.

94 + attention_dropout (`float`, *optional*,

95 + defaults to 0.0):

96 + The dropout ratio for the attention

97 + probabilities.

98 + rms_norm_eps (`float`, *optional*,

99 + defaults to 1e-06):

100 + The epsilon used by the rms normalization

101 + layers.

102 + layer_types (`list`, *optional*):

103 + Attention pattern for each layer. Can contain

104 + `"full_attention"` or `"linear_attention"`.

105 + Defaults to linear attention for most layers

106 + with full attention for every 4th layer.

107 + linear_num_key_heads (`int`, *optional*):

108 + Number of key heads for the linear attention

109 + layers. Defaults to `num_attention_heads`.

110 + linear_num_value_heads (`int`, *optional*):

111 + Number of value heads for the linear attention

112 + layers. Defaults to `num_attention_heads`.

113 + linear_key_head_dim (`int`, *optional*):

114 + Dimension of each key head in linear attention

115 + layers. Defaults to

116 + `0.75 * hidden_size / linear_num_key_heads`.

117 + linear_value_head_dim (`int`, *optional*):

118 + Dimension of each value head in linear

119 + attention layers. Defaults to

120 + `2 * linear_key_head_dim`.

121 + linear_a_log_min (`float`, *optional*,

122 + defaults to 0.0):

123 + Minimum value for uniform initialization of

124 + A_log in GatedDeltaNet layers.

125 + linear_a_log_max (`float`, *optional*,

126 + defaults to 16.0):

127 + Maximum value for uniform initialization of

128 + A_log in GatedDeltaNet layers.

129 + linear_dt_min (`float`, *optional*,

130 + defaults to 0.001):

131 + Minimum value for dt initialization in

132 + GatedDeltaNet layers.

133 + linear_dt_max (`float`, *optional*,

134 + defaults to 0.1):

135 + Maximum value for dt initialization in

136 + GatedDeltaNet layers.

137 + linear_dt_init_floor (`float`, *optional*,

138 + defaults to 0.0001):

139 + Floor value for clamping dt during

140 + initialization in GatedDeltaNet layers.

141 + linear_conv_kernel_dim (`int`, *optional*,

142 + defaults to 4):

143 + Kernel size for the short convolution applied

144 + to queries, keys, and values in linear

145 + attention layers.

146 + linear_allow_neg_eigval (`bool`, *optional*,

147 + defaults to `True`):

148 + Whether to allow negative eigenvalues in the

149 + GatedDeltaNet recurrence. When `True`, the

150 + beta parameter is scaled by 2.0 to allow

151 + values in range [0, 2] instead of [0, 1].

152 + ```python

153 + >>> from transformers import (

154 + ... OlmoHybridModel,

155 + ... OlmoHybridConfig,

156 + ... )

157 +

158 + >>> configuration = OlmoHybridConfig()

159 + >>> model = OlmoHybridModel(configuration)

160 + >>> configuration = model.config

161 + ```

162 + """

163 +

164 + model_type = "olmo_hybrid"

165 + keys_to_ignore_at_inference = ["past_key_values"]

166 + base_model_tp_plan = {

167 + "layers.*.self_attn.q_proj": "colwise_gather_output",

168 + "layers.*.self_attn.k_proj": "colwise_gather_output",

169 + "layers.*.self_attn.v_proj": "colwise_gather_output",

170 + "layers.*.self_attn.o_proj": "rowwise_split_input",

171 + "layers.*.mlp.gate_proj": "colwise",

172 + "layers.*.mlp.up_proj": "colwise",

173 + "layers.*.mlp.down_proj": "rowwise",

174 + }

175 + base_model_pp_plan = {

176 + "embed_tokens": (["input_ids"], ["inputs_embeds"]),

177 + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),

178 + "norm": (["hidden_states"], ["hidden_states"]),

179 + }

180 +

181 + def __init__(

182 + self,

183 + vocab_size: int | None = 100352,

184 + hidden_size: int | None = 3840,

185 + intermediate_size: int | None = 11008,

186 + num_hidden_layers: int | None = 32,

187 + num_attention_heads: int | None = 30,

188 + num_key_value_heads: int | None = None,

189 + hidden_act: str | None = "silu",

190 + max_position_embeddings: int | None = 65536,

191 + initializer_range: float | None = 0.02,

192 + use_cache: bool | None = True,

193 + pad_token_id: int | None = 100277,

194 + bos_token_id: int | None = None,

195 + eos_token_id: int | None = 100257,

196 + tie_word_embeddings: bool | None = False,

197 + rope_parameters=None,

198 + attention_bias: bool | None = False,

199 + attention_dropout: float | None = 0.0,

200 + rms_norm_eps: float | None = 1e-06,

201 + layer_types: list[str] | None = None,

202 + linear_num_key_heads: int | None = None,

203 + linear_num_value_heads: int | None = None,

204 + linear_key_head_dim: int | None = None,

205 + linear_value_head_dim: int | None = None,

206 + linear_a_log_min: float = 0.0,

207 + linear_a_log_max: float = 16.0,

208 + linear_dt_min: float = 0.001,

209 + linear_dt_max: float = 0.1,

210 + linear_dt_init_floor: float = 1e-4,

211 + linear_conv_kernel_dim: int = 4,

212 + linear_allow_neg_eigval: bool = True,

213 + **kwargs,

214 + ):

215 + super().__init__(**kwargs)

216 +

217 + assert num_hidden_layers is not None

218 + assert hidden_size is not None

219 + assert num_attention_heads is not None

220 +

221 + if layer_types is None:

222 + # Default: linear attention for most layers, full attention every 4th layer

223 + layer_types = ["linear_attention"] * int(num_hidden_layers)

224 + for i in range(int(num_hidden_layers)):

225 + if i % 4 == 3:

226 + layer_types[i] = "full_attention"

227 + # Ensure at least one full attention layer for small num_hidden_layers

228 + if "full_attention" not in layer_types:

229 + layer_types[-1] = "full_attention"

230 +

231 + layer_type_validation(layer_types, num_hidden_layers)

232 + if "linear_attention" not in layer_types:

233 + raise ValueError(

234 + "OLMoHybrid expects at least one 'linear_attention' layer."

235 + )

236 + if all(t == "linear_attention" for t in layer_types):

237 + raise ValueError("OLMoHybrid expects at least one attention layer.")

238 +

239 + self.layer_types = layer_types

240 +

241 + if linear_num_key_heads is None:

242 + linear_num_key_heads = num_attention_heads

243 + if linear_num_value_heads is None:

244 + linear_num_value_heads = num_attention_heads

245 + if linear_key_head_dim is None:

246 + linear_key_head_dim = int(0.75 * hidden_size / linear_num_key_heads)

247 + if linear_value_head_dim is None:

248 + linear_value_head_dim = 2 * linear_key_head_dim

249 +

250 + self.linear_num_key_heads = linear_num_key_heads

251 + self.linear_num_value_heads = linear_num_value_heads

252 + self.linear_key_head_dim = linear_key_head_dim

253 + self.linear_value_head_dim = linear_value_head_dim

254 + self.linear_a_log_min = linear_a_log_min

255 + self.linear_a_log_max = linear_a_log_max

256 + self.linear_dt_min = linear_dt_min

257 + self.linear_dt_max = linear_dt_max

258 + self.linear_dt_init_floor = linear_dt_init_floor

259 + self.linear_conv_kernel_dim = linear_conv_kernel_dim

260 + self.linear_allow_neg_eigval = linear_allow_neg_eigval

261 + self.vocab_size = vocab_size

262 + self.max_position_embeddings = max_position_embeddings

263 + self.hidden_size = hidden_size

264 + self.intermediate_size = intermediate_size

265 + self.num_hidden_layers = num_hidden_layers

266 + self.num_attention_heads = num_attention_heads

267 +

268 + # for backward compatibility

269 + if num_key_value_heads is None:

270 + num_key_value_heads = num_attention_heads

271 +

272 + self.num_key_value_heads = num_key_value_heads

273 + self.hidden_act = hidden_act

274 + self.initializer_range = initializer_range

275 + self.rms_norm_eps = rms_norm_eps

276 + self.use_cache = use_cache

277 + self.attention_bias = attention_bias

278 + self.attention_dropout = attention_dropout

279 + self.rope_parameters = rope_parameters

280 +

281 + self.tie_word_embeddings = tie_word_embeddings

282 + self.pad_token_id = pad_token_id

283 + self.bos_token_id = bos_token_id

284 + self.eos_token_id = eos_token_id