@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a52205bccf
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Model] Add HrmTextForCausalLM (Hierarchical Reasoning Model — Text) (#43098)

Signed-off-by: Wuyifei <wuyifei@me.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

yifei wu committed Jun 16, 2026 at 22:41 UTC a52205bccfc7c270dd445cf253f18deb3778e334
8 files changed +658
docs/models/supported_models.md
+1
@@ -419,6 +419,7 @@ th {
419 | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ |
420 | `Grok1ModelForCausalLM` | Grok1 | `hpcai-tech/grok-1`. | ✅︎ | ✅︎ |
421 | `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
422 +| `HrmTextForCausalLM` | HRM-Text | `sapientinc/HRM-Text-1B`, etc. | | |
423 | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
424 | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
425 | `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ |
tests/models/registry.py
+4
@@ -335,6 +335,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
335 "hpcai-tech/grok-1", trust_remote_code=True
336 ),
337 "Grok1ForCausalLM": _HfExamplesInfo("xai-org/grok-2", trust_remote_code=True),
338 + "HrmTextForCausalLM": _HfExamplesInfo(
339 + "sapientinc/HRM-Text-1B",
340 + min_transformers_version="5.9.0",
341 + ),
342 "HunYuanDenseV1ForCausalLM": _HfExamplesInfo("tencent/Hunyuan-7B-Instruct"),
343 "HunYuanMoEV1ForCausalLM": _HfExamplesInfo(
344 "tencent/Hunyuan-A13B-Instruct", trust_remote_code=True
vllm/model_executor/layers/attention/__init__.py
+4
@@ -11,6 +11,9 @@ from vllm.model_executor.layers.attention.encoder_only_attention import (
11 )
12 from vllm.model_executor.layers.attention.mla_attention import MLAAttention
13 from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
14 +from vllm.model_executor.layers.attention.prefill_prefix_lm_attention import (
15 + PrefillPrefixLMAttention,
16 +)
17 from vllm.model_executor.layers.attention.static_sink_attention import (
18 StaticSinkAttention,
19 )
@@ -22,5 +25,6 @@ __all__ = [
25 "EncoderOnlyAttention",
26 "MLAAttention",
27 "MMEncoderAttention",
28 + "PrefillPrefixLMAttention",
29 "StaticSinkAttention",
30 ]
vllm/model_executor/layers/attention/prefill_prefix_lm_attention.py new
+86
@@ -0,0 +1,86 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from dataclasses import replace
4 +
5 +import torch
6 +
7 +from vllm.config import CacheConfig, VllmConfig
8 +from vllm.model_executor.layers.attention import Attention
9 +from vllm.model_executor.layers.attention.encoder_only_attention import (
10 + create_encoder_only_attention_backend,
11 +)
12 +from vllm.v1.attention.backend import AttentionType
13 +from vllm.v1.attention.selector import get_attn_backend
14 +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec
15 +
16 +
17 +class PrefillPrefixLMAttention(Attention):
18 + """Decoder attention that runs non-causally (Prefix LM).
19 +
20 + This reuses the encoder-only backend wrapper, which forces
21 + ``causal=False`` on *every* metadata build (prefill and decode alike),
22 + while keeping ``attn_type=DECODER`` so a KV cache is still allocated.
23 +
24 + Effect by phase:
25 + - Prefill: query tokens attend to each other bidirectionally -- this is
26 + where the Prefix LM (non-causal) behavior actually takes effect.
27 + - Single-token decode: ``causal=False`` is a no-op. The one new query
28 + attends to the whole (frozen) KV cache exactly as a causal decode
29 + would, and cached tokens cannot attend back to it, so the output is
30 + identical to a causal decoder.
31 + """
32 +
33 + def __init__(
34 + self,
35 + num_heads: int,
36 + head_size: int,
37 + scale: float,
38 + cache_config: CacheConfig | None = None,
39 + attn_type: str | None = None,
40 + **kwargs,
41 + ):
42 + dtype = torch.get_default_dtype()
43 +
44 + if cache_config is not None:
45 + kv_cache_dtype = cache_config.cache_dtype
46 + else:
47 + kv_cache_dtype = "auto"
48 +
49 + underlying_attn_backend = get_attn_backend(
50 + head_size,
51 + dtype,
52 + kv_cache_dtype,
53 + attn_type=AttentionType.DECODER,
54 + )
55 +
56 + attn_backend = create_encoder_only_attention_backend(underlying_attn_backend)
57 +
58 + if attn_type is not None:
59 + assert attn_type == AttentionType.DECODER, (
60 + "PrefillPrefixLMAttention only supports AttentionType.DECODER"
61 + )
62 +
63 + super().__init__(
64 + num_heads=num_heads,
65 + head_size=head_size,
66 + scale=scale,
67 + cache_config=cache_config,
68 + attn_backend=attn_backend,
69 + attn_type=AttentionType.DECODER,
70 + **kwargs,
71 + )
72 +
73 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
74 + """Tag the KV cache spec as non-causal.
75 +
76 + The layout is identical to a regular decoder full-attention layer, so
77 + we reuse the base spec and only flip ``non_causal=True``. The engine
78 + core reads this flag (across the worker/engine process boundary, via
79 + the pickled spec) to disable scheduling features that assume causal
80 + attention -- chunked prefill and prefix caching -- which would
81 + otherwise corrupt the bidirectional prefill of a Prefix LM.
82 + """
83 + spec = super().get_kv_cache_spec(vllm_config)
84 + if isinstance(spec, FullAttentionSpec):
85 + return replace(spec, non_causal=True)
86 + return spec
vllm/model_executor/models/hrm_text.py new
+527
@@ -0,0 +1,527 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +HRM-Text: Hierarchical Reasoning Model — Text variant.
5 +
6 +Reference Hugging Face implementation:
7 + src/transformers/models/hrm_text/modeling_hrm_text.py
8 +
9 +The model performs a hierarchical recurrent forward over two transformer
10 +stacks (``H`` slow, ``L`` fast) inside nested loops. Each recurrence step
11 +gets its own KV cache slot via a unique vLLM-visible layer index. The
12 +PrefixLM attention pattern (prompt bidirectional, response causal) is
13 +realized by reusing ``EncoderOnlyAttention`` (which sets ``causal=False``
14 +unconditionally on every metadata build) but with ``attn_type=DECODER``
15 +so the KV cache is allocated; see ``HrmTextAttention`` for usage.
16 +
17 +The on-disk ``attn.gqkv_proj.weight`` (rows concatenated as
18 +``[gate | q | k | v]``) is loaded by a single
19 +``MergedColumnParallelLinear`` with four equal-sized output partitions;
20 +its weight loader auto-splits the fused tensor along the output dim by
21 +``output_sizes`` (the same path used by Phi-3's fused gate_up_proj).
22 +"""
23 +
24 +from collections.abc import Iterable
25 +from typing import Literal
26 +
27 +import torch
28 +from torch import nn
29 +from transformers import PretrainedConfig
30 +
31 +from vllm.compilation.decorators import support_torch_compile
32 +from vllm.config import CacheConfig, VllmConfig
33 +from vllm.distributed import get_tensor_model_parallel_world_size
34 +from vllm.model_executor.layers.activation import SiluAndMul
35 +from vllm.model_executor.layers.attention import PrefillPrefixLMAttention
36 +from vllm.model_executor.layers.layernorm import RMSNorm
37 +from vllm.model_executor.layers.linear import (
38 + MergedColumnParallelLinear,
39 + RowParallelLinear,
40 +)
41 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
42 +from vllm.model_executor.layers.quantization import QuantizationConfig
43 +from vllm.model_executor.layers.rotary_embedding import get_rope
44 +from vllm.model_executor.layers.vocab_parallel_embedding import (
45 + ParallelLMHead,
46 + VocabParallelEmbedding,
47 +)
48 +from vllm.sequence import IntermediateTensors
49 +
50 +from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
51 +
52 +
53 +class HrmTextMLP(nn.Module):
54 + def __init__(
55 + self,
56 + hidden_size: int,
57 + intermediate_size: int,
58 + hidden_act: str,
59 + bias: bool = False,
60 + quant_config: QuantizationConfig | None = None,
61 + prefix: str = "",
62 + ) -> None:
63 + super().__init__()
64 + if hidden_act != "silu":
65 + raise ValueError(
66 + f"HrmTextMLP only supports hidden_act='silu', got {hidden_act!r}"
67 + )
68 + self.gate_up_proj = MergedColumnParallelLinear(
69 + hidden_size,
70 + [intermediate_size] * 2,
71 + bias=bias,
72 + quant_config=quant_config,
73 + prefix=f"{prefix}.gate_up_proj",
74 + )
75 + self.down_proj = RowParallelLinear(
76 + intermediate_size,
77 + hidden_size,
78 + bias=bias,
79 + quant_config=quant_config,
80 + prefix=f"{prefix}.down_proj",
81 + )
82 + self.act_fn = SiluAndMul()
83 +
84 + def forward(self, x: torch.Tensor) -> torch.Tensor:
85 + gate_up, _ = self.gate_up_proj(x)
86 + x = self.act_fn(gate_up)
87 + x, _ = self.down_proj(x)
88 + return x
89 +
90 +
91 +class HrmTextAttention(nn.Module):
92 + """One self-attention block; weights shared across recurrence steps.
93 +
94 + HF transformers writes a single fused ``attn.gqkv_proj.weight`` on
95 + disk (per ``transformers/conversion_mapping.py`` ``"hrm_text"``
96 + mapping; rows are concatenated as ``[gate | q | k | v]`` along
97 + ``dim=0``). We mirror that on the model side with a single
98 + ``MergedColumnParallelLinear`` whose four equal output partitions
99 + are sharded along the head axis under TP; its weight loader
100 + auto-splits the fused tensor (same path used by Phi-3's fused
101 + gate_up_proj). HF's runtime config currently hardcodes MHA
102 + (``num_key_value_groups=1``); GQA would require ``QKVParallelLinear``
103 + semantics for q/k/v shard replication and is left for a follow-up
104 + if/when HF adds it.
105 +
106 + Holds:
107 + - parameters: gqkv_proj, o_proj, rotary_emb (shared across cycles).
108 + - ``attn_per_step``: a ``nn.ModuleDict`` keyed by recurrence step
109 + (as a string), each value an ``EncoderOnlyAttention`` (with
110 + ``attn_type=DECODER`` so the KV cache is allocated; the
111 + ``EncoderOnlyAttention`` wrapper sets ``causal=False`` on every
112 + metadata build). The L stack steps are
113 + ``[high_cycle_idx*(L_cycles+1)+low_cycle_idx]`` and the H stack
114 + steps are ``[high_cycle_idx*(L_cycles+1)+L_cycles]``; the two
115 + ranges are disjoint so each instance registers a unique vLLM
116 + ``layer_name``
117 + (``model.{H,L}_module.layers.{global_idx}.self_attn``) and gets
118 + its own KV cache slot. The global layer index per recurrence step
119 + is ``step * num_layers_per_stack + layer_idx_in_stack``, matching
120 + the HF transformers ``cycle_offset`` formula in
121 + ``modeling_hrm_text.py``.
122 + """
123 +
124 + def __init__(
125 + self,
126 + config: PretrainedConfig,
127 + layer_idx_in_stack: int,
128 + stack_kind: Literal["L", "H"],
129 + cache_config: CacheConfig | None = None,
130 + quant_config: QuantizationConfig | None = None,
131 + prefix: str = "",
132 + ) -> None:
133 + super().__init__()
134 + self.config = config
135 + self.hidden_size = config.hidden_size
136 + tp_size = get_tensor_model_parallel_world_size()
137 +
138 + self.total_num_heads = config.num_attention_heads
139 + assert self.total_num_heads % tp_size == 0, (
140 + f"num_attention_heads={self.total_num_heads} must be divisible "
141 + f"by tp_size={tp_size}"
142 + )
143 + # HF main hardcodes MHA (num_key_value_groups=1). We follow.
144 + self.total_num_kv_heads = config.num_attention_heads
145 + self.num_heads = self.total_num_heads // tp_size
146 + self.num_kv_heads = self.total_num_kv_heads // tp_size
147 + self.head_dim = getattr(
148 + config, "head_dim", self.hidden_size // self.total_num_heads
149 + )
150 + self.q_size = self.num_heads * self.head_dim
151 + self.kv_size = self.num_kv_heads * self.head_dim
152 + self.scaling = self.head_dim**-0.5
153 +
154 + bias = getattr(config, "attention_bias", False)
155 +
156 + # gqkv_proj: 4-way fused [gate | q | k | v] matching the on-disk
157 + # `attn.gqkv_proj.weight` row layout. MergedColumnParallelLinear's
158 + # weight_loader auto-splits the fused disk tensor along the output
159 + # dim by `output_sizes` (Phi-3's fused gate_up_proj path). MHA
160 + # only: GQA (num_kv_heads != num_heads) would need
161 + # QKVParallelLinear semantics for q/k/v shard replication.
162 + per_head_size = self.total_num_heads * self.head_dim
163 + self.gqkv_proj = MergedColumnParallelLinear(
164 + input_size=self.hidden_size,
165 + output_sizes=[per_head_size] * 4,
166 + bias=bias,
167 + quant_config=quant_config,
168 + prefix=f"{prefix}.gqkv_proj",
169 + )
170 + self.o_proj = RowParallelLinear(
171 + input_size=self.total_num_heads * self.head_dim,
172 + output_size=self.hidden_size,
173 + bias=bias,
174 + quant_config=quant_config,
175 + prefix=f"{prefix}.o_proj",
176 + )
177 +
178 + # vllm get_rope accepts ``rope_parameters`` directly, matching
179 + # the dict-shaped HF config field.
180 + self.rotary_emb = get_rope(
181 + head_size=self.head_dim,
182 + max_position=config.max_position_embeddings,
183 + rope_parameters=config.rope_parameters,
184 + )
185 +
186 + # Create one Attention instance per recurrence step actually used
187 + # by this stack. L runs at steps {h*(L+1)+l : 0 <= l < L_cycles},
188 + # H at steps {h*(L+1)+L : 0 <= h < H_cycles}; the sets are
189 + # disjoint, so one global index per (step, layer_in_stack) gives
190 + # each Attention its own ``layer_name`` and KV cache slot.
191 + H_cycles = config.H_cycles
192 + L_cycles = config.L_cycles
193 + num_layers_per_stack = config.num_layers_per_stack
194 + if stack_kind == "L":
195 + steps_used = [
196 + high_cycle_idx * (L_cycles + 1) + low_cycle_idx
197 + for high_cycle_idx in range(H_cycles)
198 + for low_cycle_idx in range(L_cycles)
199 + ]
200 + else: # "H"
201 + steps_used = [
202 + high_cycle_idx * (L_cycles + 1) + L_cycles
203 + for high_cycle_idx in range(H_cycles)
204 + ]
205 +
206 + # `PrefillPrefixLMAttention` forces `causal=False` on every metadata
207 + # build, so the prompt attends bidirectionally during prefill (matching
208 + # the HRM-Text training distribution), while `attn_type=DECODER` keeps
209 + # the KV cache allocation needed by the recurrent forward. At
210 + # single-token decode `causal=False` is a no-op. See
211 + # `PrefillPrefixLMAttention`.
212 + self.attn_per_step = nn.ModuleDict()
213 + for step in steps_used:
214 + global_idx = step * num_layers_per_stack + layer_idx_in_stack
215 + unique_prefix = prefix.replace(
216 + f"layers.{layer_idx_in_stack}", f"layers.{global_idx}"
217 + )
218 + self.attn_per_step[str(step)] = PrefillPrefixLMAttention(
219 + num_heads=self.num_heads,
220 + head_size=self.head_dim,
221 + scale=self.scaling,
222 + num_kv_heads=self.num_kv_heads,
223 + cache_config=cache_config,
224 + quant_config=quant_config,
225 + prefix=f"{unique_prefix}.attn",
226 + )
227 +
228 + def forward(
229 + self,
230 + positions: torch.Tensor,
231 + hidden_states: torch.Tensor,
232 + current_step: int,
233 + ) -> torch.Tensor:
234 + gqkv, _ = self.gqkv_proj(hidden_states)
235 + g, q, k, v = gqkv.split(
236 + [self.q_size, self.q_size, self.kv_size, self.kv_size], dim=-1
237 + )
238 + q, k = self.rotary_emb(positions, q, k)
239 + attn_out = self.attn_per_step[str(current_step)](q, k, v)
240 + # Sigmoid gate. Shapes: attn_out is (..., q_size); g is (..., q_size).
241 + attn_out = torch.sigmoid(g) * attn_out
242 + out, _ = self.o_proj(attn_out)
243 + return out
244 +
245 +
246 +class HrmTextDecoderLayer(nn.Module):
247 + def __init__(
248 + self,
249 + config: PretrainedConfig,
250 + layer_idx_in_stack: int,
251 + stack_kind: str,
252 + cache_config: CacheConfig | None = None,
253 + quant_config: QuantizationConfig | None = None,
254 + prefix: str = "",
255 + ) -> None:
256 + super().__init__()
257 + self.hidden_size = config.hidden_size
258 + # Attribute name `self_attn` matches HF's model class. The on-disk
259 + # `attn.{gqkv_proj,o_proj}.weight` keys are renamed to
260 + # `self_attn.{gqkv_proj,o_proj}.weight` by the `WeightsMapper` in
261 + # `HrmTextForCausalLM` so vLLM's standard `AutoWeightsLoader`
262 + # handles the rest.
263 + self.self_attn = HrmTextAttention(
264 + config=config,
265 + layer_idx_in_stack=layer_idx_in_stack,
266 + stack_kind=stack_kind,
267 + cache_config=cache_config,
268 + quant_config=quant_config,
269 + prefix=f"{prefix}.self_attn",
270 + )
271 + self.mlp = HrmTextMLP(
272 + hidden_size=config.hidden_size,
273 + intermediate_size=config.intermediate_size,
274 + hidden_act=config.hidden_act,
275 + bias=getattr(config, "mlp_bias", False),
276 + quant_config=quant_config,
277 + prefix=f"{prefix}.mlp",
278 + )
279 + # Parameterless RMSNorm (HF main: HrmTextRMSNorm has no weight).
280 + self.input_layernorm = RMSNorm(
281 + config.hidden_size, eps=config.rms_norm_eps, has_weight=False
282 + )
283 + self.post_attention_layernorm = RMSNorm(
284 + config.hidden_size, eps=config.rms_norm_eps, has_weight=False
285 + )
286 +
287 + def forward(
288 + self,
289 + positions: torch.Tensor,
290 + hidden_states: torch.Tensor,
291 + current_step: int,
292 + ) -> torch.Tensor:
293 + residual = hidden_states
294 + hidden_states = self.input_layernorm(hidden_states)
295 + hidden_states = self.self_attn(
296 + positions=positions,
297 + hidden_states=hidden_states,
298 + current_step=current_step,
299 + )
300 + hidden_states = residual + hidden_states
301 +
302 + residual = hidden_states
303 + hidden_states = self.post_attention_layernorm(hidden_states)
304 + hidden_states = self.mlp(hidden_states)
305 + hidden_states = residual + hidden_states
306 + return hidden_states
307 +
308 +
309 +class HrmTextStack(nn.Module):
310 + """A single transformer stack — used twice (H and L)."""
311 +
312 + def __init__(
313 + self,
314 + config: PretrainedConfig,
315 + stack_kind: str,
316 + cache_config: CacheConfig | None = None,
317 + quant_config: QuantizationConfig | None = None,
318 + prefix: str = "",
319 + ) -> None:
320 + super().__init__()
321 + self.layers = nn.ModuleList(
322 + [
323 + HrmTextDecoderLayer(
324 + config=config,
325 + layer_idx_in_stack=i,
326 + stack_kind=stack_kind,
327 + cache_config=cache_config,
328 + quant_config=quant_config,
329 + prefix=f"{prefix}.layers.{i}",
330 + )
331 + for i in range(config.num_layers_per_stack)
332 + ]
333 + )
334 + self.final_norm = RMSNorm(
335 + config.hidden_size, eps=config.rms_norm_eps, has_weight=False
336 + )
337 +
338 + def forward(
339 + self,
340 + positions: torch.Tensor,
341 + hidden_states: torch.Tensor,
342 + current_step_base: int,
343 + ) -> torch.Tensor:
344 + for layer in self.layers:
345 + hidden_states = layer(
346 + positions=positions,
347 + hidden_states=hidden_states,
348 + current_step=current_step_base,
349 + )
350 + return self.final_norm(hidden_states)
351 +
352 +
353 +@support_torch_compile
354 +class HrmTextModel(nn.Module):
355 + """Hierarchical recurrent transformer body.
356 +
357 + Forward (matches HF main exactly,
358 + src/transformers/models/hrm_text/modeling_hrm_text.py:495-547):
359 +
360 + hidden_states_high_cycle = embed(input_ids) * embedding_scale
361 + hidden_states_low_cycle = z_L_init.expand_as(hidden_states_high_cycle)
362 + for high_cycle_idx in range(H_cycles):
363 + for low_cycle_idx in range(L_cycles):
364 + step = high_cycle_idx * (L_cycles + 1) + low_cycle_idx
365 + hidden_states_low_cycle = L_module(
366 + hidden_states_low_cycle + hidden_states_high_cycle,
367 + current_step=step,
368 + )
369 + step = high_cycle_idx * (L_cycles + 1) + L_cycles
370 + hidden_states_high_cycle = H_module(
371 + hidden_states_high_cycle + hidden_states_low_cycle,
372 + current_step=step,
373 + )
374 + return hidden_states_high_cycle
375 + """
376 +
377 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
378 + super().__init__()
379 + config = vllm_config.model_config.hf_config
380 + cache_config = vllm_config.cache_config
381 + quant_config = vllm_config.quant_config
382 + self.config = config
383 +
384 + self.embed_tokens = VocabParallelEmbedding(
385 + config.vocab_size,
386 + config.hidden_size,
387 + quant_config=quant_config,
388 + prefix=f"{prefix}.embed_tokens",
389 + )
390 + self.L_module = HrmTextStack(
391 + config=config,
392 + stack_kind="L",
393 + cache_config=cache_config,
394 + quant_config=quant_config,
395 + prefix=f"{prefix}.L_module",
396 + )
397 + self.H_module = HrmTextStack(
398 + config=config,
399 + stack_kind="H",
400 + cache_config=cache_config,
401 + quant_config=quant_config,
402 + prefix=f"{prefix}.H_module",
403 + )
404 + # Frozen learned initial L state. HF inits to zeros and sets
405 + # requires_grad_(False); for inference we just load the tensor.
406 + self.z_L_init = nn.Parameter(
407 + torch.zeros(config.hidden_size), requires_grad=False
408 + )
409 +
410 + # Embedding scale: HF uses config.embedding_scale (default
411 + # 1 / initializer_range = 50.0 when initializer_range=0.02). NOT
412 + # sqrt(hidden_size) like Gemma.
413 + self.embedding_scale = getattr(config, "embedding_scale", None)
414 + if self.embedding_scale is None:
415 + init_range = getattr(config, "initializer_range", 0.02)
416 + self.embedding_scale = 1.0 / init_range
417 +
418 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
419 + return self.embed_tokens(input_ids) * self.embedding_scale
420 +
421 + def forward(
422 + self,
423 + input_ids: torch.Tensor | None,
424 + positions: torch.Tensor,
425 + intermediate_tensors: IntermediateTensors | None = None,
426 + inputs_embeds: torch.Tensor | None = None,
427 + ) -> torch.Tensor:
428 + if inputs_embeds is None:
429 + assert input_ids is not None
430 + inputs_embeds = self.embed_input_ids(input_ids)
431 +
432 + hidden_states_high_cycle = inputs_embeds
433 + hidden_states_low_cycle = self.z_L_init.to(
434 + dtype=hidden_states_high_cycle.dtype,
435 + device=hidden_states_high_cycle.device,
436 + ).expand_as(hidden_states_high_cycle)
437 +
438 + H_cycles = self.config.H_cycles
439 + L_cycles = self.config.L_cycles
440 + for high_cycle_idx in range(H_cycles):
441 + for low_cycle_idx in range(L_cycles):
442 + step = high_cycle_idx * (L_cycles + 1) + low_cycle_idx
443 + hidden_states_low_cycle = self.L_module(
444 + positions=positions,
445 + hidden_states=hidden_states_low_cycle + hidden_states_high_cycle,
446 + current_step_base=step,
447 + )
448 + step = high_cycle_idx * (L_cycles + 1) + L_cycles
449 + hidden_states_high_cycle = self.H_module(
450 + positions=positions,
451 + hidden_states=hidden_states_high_cycle + hidden_states_low_cycle,
452 + current_step_base=step,
453 + )
454 +
455 + return hidden_states_high_cycle
456 +
457 +
458 +class HrmTextForCausalLM(nn.Module):
459 + """Hierarchical Reasoning Model — Text variant, causal LM.
460 +
461 + Reference: src/transformers/models/hrm_text/modeling_hrm_text.py
462 + """
463 +
464 + # On-disk weight key remap: HF stores attention weights as
465 + # `attn.{gqkv_proj,o_proj}.weight`; our model uses `self_attn.*`
466 + # (matching HF's runtime model class). Both `gqkv_proj` (4-way fused
467 + # gate/q/k/v) and `mlp.gate_up_proj` (2-way fused gate/up) are loaded
468 + # directly via MergedColumnParallelLinear's fused-on-disk path; no
469 + # packed_modules_mapping entries are needed.
470 + hf_to_vllm_mapper = WeightsMapper(
471 + orig_to_new_substr={".attn.": ".self_attn."},
472 + )
473 +
474 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
475 + super().__init__()
476 + config = vllm_config.model_config.hf_config
477 + quant_config = vllm_config.quant_config
478 + self.config = config
479 + self.quant_config = quant_config
480 +
481 + if vllm_config.parallel_config.pipeline_parallel_size > 1:
482 + raise ValueError(
483 + "HrmTextForCausalLM does not support pipeline parallelism."
484 + )
485 +
486 + self.model = HrmTextModel(
487 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
488 + )
489 +
490 + if config.tie_word_embeddings:
491 + self.lm_head = self.model.embed_tokens
492 + else:
493 + self.lm_head = ParallelLMHead(
494 + config.vocab_size,
495 + config.hidden_size,
496 + quant_config=quant_config,
497 + prefix=maybe_prefix(prefix, "lm_head"),
498 + )
499 + self.logits_processor = LogitsProcessor(config.vocab_size)
500 +
501 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
502 + return self.model.embed_input_ids(input_ids)
503 +
504 + def forward(
505 + self,
506 + input_ids: torch.Tensor | None,
507 + positions: torch.Tensor,
508 + intermediate_tensors: IntermediateTensors | None = None,
509 + inputs_embeds: torch.Tensor | None = None,
510 + ) -> torch.Tensor:
511 + return self.model(
512 + input_ids=input_ids,
513 + positions=positions,
514 + intermediate_tensors=intermediate_tensors,
515 + inputs_embeds=inputs_embeds,
516 + )
517 +
518 + def compute_logits(
519 + self,
520 + hidden_states: torch.Tensor,
521 + ) -> torch.Tensor | None:
522 + return self.logits_processor(self.lm_head, hidden_states)
523 +
524 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
525 + skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
526 + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
527 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/registry.py
+1
@@ -134,6 +134,7 @@ _TEXT_GENERATION_MODELS = {
134 "GritLM": ("gritlm", "GritLM"),
135 "Grok1ModelForCausalLM": ("grok1", "GrokForCausalLM"),
136 "Grok1ForCausalLM": ("grok1", "GrokForCausalLM"),
137 + "HrmTextForCausalLM": ("hrm_text", "HrmTextForCausalLM"),
138 "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"),
139 "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"),
140 "HYV3ForCausalLM": ("hy_v3", "HYV3ForCausalLM"),
vllm/v1/engine/core.py
+22
@@ -246,6 +246,28 @@ class EngineCore:
246 # Get all kv cache needed by the model
247 kv_cache_specs = self.model_executor.get_kv_cache_specs()
248
249 + # Some layers (e.g. Prefix LM attention) run non-causally and tag their
250 + # KV cache spec with ``non_causal=True``. The specs are collected here in
251 + # the engine-core process (the same process that builds the scheduler),
252 + # so this is the multiproc-safe place to translate that layer-level
253 + # signal into a scheduling policy: chunked prefill and prefix caching
254 + # both assume causal attention and would corrupt non-causal prefill.
255 + if any(
256 + getattr(spec, "non_causal", False)
257 + for worker_specs in kv_cache_specs
258 + for spec in worker_specs.values()
259 + ):
260 + if vllm_config.scheduler_config.enable_chunked_prefill:
261 + logger.info(
262 + "Disabling chunked prefill: model has non-causal attention layers."
263 + )
264 + vllm_config.scheduler_config.enable_chunked_prefill = False
265 + if vllm_config.cache_config.enable_prefix_caching:
266 + logger.info(
267 + "Disabling prefix caching: model has non-causal attention layers."
268 + )
269 + vllm_config.cache_config.enable_prefix_caching = False
270 +
271 has_kv_cache = any(kv_cache_spec for kv_cache_spec in kv_cache_specs)
272 if has_kv_cache:
273 if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
vllm/v1/kv_cache_interface.py
+13
@@ -219,6 +219,15 @@ class FullAttentionSpec(AttentionSpec):
219 """
220 attention_chunk_size: int | None = None
221
222 + non_causal: bool = False
223 + """
224 + Whether the layer attends non-causally (e.g. Prefix LM). Carried on the
225 + spec so the engine core, which collects specs from all workers before the
226 + scheduler is built, can adjust scheduling policy (chunked prefill / prefix
227 + caching) regardless of tensor-parallel layout. It does not affect the KV
228 + cache layout itself.
229 + """
230 +
231 def __post_init__(self):
232 if self.head_size_v is None:
233 object.__setattr__(self, "head_size_v", self.head_size)
@@ -276,6 +285,9 @@ class FullAttentionSpec(AttentionSpec):
285 page_size_padded=specs[0].page_size_padded,
286 sliding_window=cls.merge_window_sizes(sliding_window),
287 attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
288 + # If any layer in the group is non-causal, treat the group as
289 + # non-causal so the engine core disables incompatible scheduling.
290 + non_causal=any(spec.non_causal for spec in specs),
291 )
292 for spec in specs:
293 for f in fields(AttentionSpec):
@@ -701,6 +713,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
713 page_size_padded=specs[0].page_size_padded,
714 sliding_window=cls.merge_window_sizes(sliding_window),
715 attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
716 + non_causal=any(spec.non_causal for spec in specs),
717 )
718 for spec in specs:
719 for f in fields(AttentionSpec):