[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>

docs/models/supported_models.md

@@ -419,6 +419,7 @@ th {
 | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ |
 | `Grok1ModelForCausalLM` | Grok1 | `hpcai-tech/grok-1`. | ✅︎ | ✅︎ |
 | `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
+| `HrmTextForCausalLM` | HRM-Text | `sapientinc/HRM-Text-1B`, etc. | | |
 | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
 | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
 | `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ |

tests/models/registry.py

@@ -335,6 +335,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         "hpcai-tech/grok-1", trust_remote_code=True
     ),
     "Grok1ForCausalLM": _HfExamplesInfo("xai-org/grok-2", trust_remote_code=True),
+    "HrmTextForCausalLM": _HfExamplesInfo(
+        "sapientinc/HRM-Text-1B",
+        min_transformers_version="5.9.0",
+    ),
     "HunYuanDenseV1ForCausalLM": _HfExamplesInfo("tencent/Hunyuan-7B-Instruct"),
     "HunYuanMoEV1ForCausalLM": _HfExamplesInfo(
         "tencent/Hunyuan-A13B-Instruct", trust_remote_code=True

vllm/model_executor/layers/attention/__init__.py

@@ -11,6 +11,9 @@ from vllm.model_executor.layers.attention.encoder_only_attention import (
 )
 from vllm.model_executor.layers.attention.mla_attention import MLAAttention
 from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
+from vllm.model_executor.layers.attention.prefill_prefix_lm_attention import (
+    PrefillPrefixLMAttention,
+)
 from vllm.model_executor.layers.attention.static_sink_attention import (
     StaticSinkAttention,
 )
@@ -22,5 +25,6 @@ __all__ = [
     "EncoderOnlyAttention",
     "MLAAttention",
     "MMEncoderAttention",
+    "PrefillPrefixLMAttention",
     "StaticSinkAttention",
 ]

vllm/model_executor/layers/attention/prefill_prefix_lm_attention.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import replace
+
+import torch
+
+from vllm.config import CacheConfig, VllmConfig
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.attention.encoder_only_attention import (
+    create_encoder_only_attention_backend,
+)
+from vllm.v1.attention.backend import AttentionType
+from vllm.v1.attention.selector import get_attn_backend
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec
+
+
+class PrefillPrefixLMAttention(Attention):
+    """Decoder attention that runs non-causally (Prefix LM).
+
+    This reuses the encoder-only backend wrapper, which forces
+    ``causal=False`` on *every* metadata build (prefill and decode alike),
+    while keeping ``attn_type=DECODER`` so a KV cache is still allocated.
+
+    Effect by phase:
+    - Prefill: query tokens attend to each other bidirectionally -- this is
+      where the Prefix LM (non-causal) behavior actually takes effect.
+    - Single-token decode: ``causal=False`` is a no-op. The one new query
+      attends to the whole (frozen) KV cache exactly as a causal decode
+      would, and cached tokens cannot attend back to it, so the output is
+      identical to a causal decoder.
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        cache_config: CacheConfig | None = None,
+        attn_type: str | None = None,
+        **kwargs,
+    ):
+        dtype = torch.get_default_dtype()
+
+        if cache_config is not None:
+            kv_cache_dtype = cache_config.cache_dtype
+        else:
+            kv_cache_dtype = "auto"
+
+        underlying_attn_backend = get_attn_backend(
+            head_size,
+            dtype,
+            kv_cache_dtype,
+            attn_type=AttentionType.DECODER,
+        )
+
+        attn_backend = create_encoder_only_attention_backend(underlying_attn_backend)
+
+        if attn_type is not None:
+            assert attn_type == AttentionType.DECODER, (
+                "PrefillPrefixLMAttention only supports AttentionType.DECODER"
+            )
+
+        super().__init__(
+            num_heads=num_heads,
+            head_size=head_size,
+            scale=scale,
+            cache_config=cache_config,
+            attn_backend=attn_backend,
+            attn_type=AttentionType.DECODER,
+            **kwargs,
+        )
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
+        """Tag the KV cache spec as non-causal.
+
+        The layout is identical to a regular decoder full-attention layer, so
+        we reuse the base spec and only flip ``non_causal=True``. The engine
+        core reads this flag (across the worker/engine process boundary, via
+        the pickled spec) to disable scheduling features that assume causal
+        attention -- chunked prefill and prefix caching -- which would
+        otherwise corrupt the bidirectional prefill of a Prefix LM.
+        """
+        spec = super().get_kv_cache_spec(vllm_config)
+        if isinstance(spec, FullAttentionSpec):
+            return replace(spec, non_causal=True)
+        return spec

vllm/model_executor/models/hrm_text.py

@@ -0,0 +1,527 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+HRM-Text: Hierarchical Reasoning Model — Text variant.
+
+Reference Hugging Face implementation:
+    src/transformers/models/hrm_text/modeling_hrm_text.py
+
+The model performs a hierarchical recurrent forward over two transformer
+stacks (``H`` slow, ``L`` fast) inside nested loops. Each recurrence step
+gets its own KV cache slot via a unique vLLM-visible layer index. The
+PrefixLM attention pattern (prompt bidirectional, response causal) is
+realized by reusing ``EncoderOnlyAttention`` (which sets ``causal=False``
+unconditionally on every metadata build) but with ``attn_type=DECODER``
+so the KV cache is allocated; see ``HrmTextAttention`` for usage.
+
+The on-disk ``attn.gqkv_proj.weight`` (rows concatenated as
+``[gate | q | k | v]``) is loaded by a single
+``MergedColumnParallelLinear`` with four equal-sized output partitions;
+its weight loader auto-splits the fused tensor along the output dim by
+``output_sizes`` (the same path used by Phi-3's fused gate_up_proj).
+"""
+
+from collections.abc import Iterable
+from typing import Literal
+
+import torch
+from torch import nn
+from transformers import PretrainedConfig
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.attention import PrefillPrefixLMAttention
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.sequence import IntermediateTensors
+
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
+
+
+class HrmTextMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        bias: bool = False,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        if hidden_act != "silu":
+            raise ValueError(
+                f"HrmTextMLP only supports hidden_act='silu', got {hidden_act!r}"
+            )
+        self.gate_up_proj = MergedColumnParallelLinear(
+            hidden_size,
+            [intermediate_size] * 2,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class HrmTextAttention(nn.Module):
+    """One self-attention block; weights shared across recurrence steps.
+
+    HF transformers writes a single fused ``attn.gqkv_proj.weight`` on
+    disk (per ``transformers/conversion_mapping.py`` ``"hrm_text"``
+    mapping; rows are concatenated as ``[gate | q | k | v]`` along
+    ``dim=0``). We mirror that on the model side with a single
+    ``MergedColumnParallelLinear`` whose four equal output partitions
+    are sharded along the head axis under TP; its weight loader
+    auto-splits the fused tensor (same path used by Phi-3's fused
+    gate_up_proj). HF's runtime config currently hardcodes MHA
+    (``num_key_value_groups=1``); GQA would require ``QKVParallelLinear``
+    semantics for q/k/v shard replication and is left for a follow-up
+    if/when HF adds it.
+
+    Holds:
+      - parameters: gqkv_proj, o_proj, rotary_emb (shared across cycles).
+      - ``attn_per_step``: a ``nn.ModuleDict`` keyed by recurrence step
+        (as a string), each value an ``EncoderOnlyAttention`` (with
+        ``attn_type=DECODER`` so the KV cache is allocated; the
+        ``EncoderOnlyAttention`` wrapper sets ``causal=False`` on every
+        metadata build). The L stack steps are
+        ``[high_cycle_idx*(L_cycles+1)+low_cycle_idx]`` and the H stack
+        steps are ``[high_cycle_idx*(L_cycles+1)+L_cycles]``; the two
+        ranges are disjoint so each instance registers a unique vLLM
+        ``layer_name``
+        (``model.{H,L}_module.layers.{global_idx}.self_attn``) and gets
+        its own KV cache slot. The global layer index per recurrence step
+        is ``step * num_layers_per_stack + layer_idx_in_stack``, matching
+        the HF transformers ``cycle_offset`` formula in
+        ``modeling_hrm_text.py``.
+    """
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        layer_idx_in_stack: int,
+        stack_kind: Literal["L", "H"],
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+
+        self.total_num_heads = config.num_attention_heads
+        assert self.total_num_heads % tp_size == 0, (
+            f"num_attention_heads={self.total_num_heads} must be divisible "
+            f"by tp_size={tp_size}"
+        )
+        # HF main hardcodes MHA (num_key_value_groups=1). We follow.
+        self.total_num_kv_heads = config.num_attention_heads
+        self.num_heads = self.total_num_heads // tp_size
+        self.num_kv_heads = self.total_num_kv_heads // tp_size
+        self.head_dim = getattr(
+            config, "head_dim", self.hidden_size // self.total_num_heads
+        )
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+
+        bias = getattr(config, "attention_bias", False)
+
+        # gqkv_proj: 4-way fused [gate | q | k | v] matching the on-disk
+        # `attn.gqkv_proj.weight` row layout. MergedColumnParallelLinear's
+        # weight_loader auto-splits the fused disk tensor along the output
+        # dim by `output_sizes` (Phi-3's fused gate_up_proj path). MHA
+        # only: GQA (num_kv_heads != num_heads) would need
+        # QKVParallelLinear semantics for q/k/v shard replication.
+        per_head_size = self.total_num_heads * self.head_dim
+        self.gqkv_proj = MergedColumnParallelLinear(
+            input_size=self.hidden_size,
+            output_sizes=[per_head_size] * 4,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gqkv_proj",
+        )
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=self.hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        # vllm get_rope accepts ``rope_parameters`` directly, matching
+        # the dict-shaped HF config field.
+        self.rotary_emb = get_rope(
+            head_size=self.head_dim,
+            max_position=config.max_position_embeddings,
+            rope_parameters=config.rope_parameters,
+        )
+
+        # Create one Attention instance per recurrence step actually used
+        # by this stack. L runs at steps {h*(L+1)+l : 0 <= l < L_cycles},
+        # H at steps {h*(L+1)+L : 0 <= h < H_cycles}; the sets are
+        # disjoint, so one global index per (step, layer_in_stack) gives
+        # each Attention its own ``layer_name`` and KV cache slot.
+        H_cycles = config.H_cycles
+        L_cycles = config.L_cycles
+        num_layers_per_stack = config.num_layers_per_stack
+        if stack_kind == "L":
+            steps_used = [
+                high_cycle_idx * (L_cycles + 1) + low_cycle_idx
+                for high_cycle_idx in range(H_cycles)
+                for low_cycle_idx in range(L_cycles)
+            ]
+        else:  # "H"
+            steps_used = [
+                high_cycle_idx * (L_cycles + 1) + L_cycles
+                for high_cycle_idx in range(H_cycles)
+            ]
+
+        # `PrefillPrefixLMAttention` forces `causal=False` on every metadata
+        # build, so the prompt attends bidirectionally during prefill (matching
+        # the HRM-Text training distribution), while `attn_type=DECODER` keeps
+        # the KV cache allocation needed by the recurrent forward. At
+        # single-token decode `causal=False` is a no-op. See
+        # `PrefillPrefixLMAttention`.
+        self.attn_per_step = nn.ModuleDict()
+        for step in steps_used:
+            global_idx = step * num_layers_per_stack + layer_idx_in_stack
+            unique_prefix = prefix.replace(
+                f"layers.{layer_idx_in_stack}", f"layers.{global_idx}"
+            )
+            self.attn_per_step[str(step)] = PrefillPrefixLMAttention(
+                num_heads=self.num_heads,
+                head_size=self.head_dim,
+                scale=self.scaling,
+                num_kv_heads=self.num_kv_heads,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=f"{unique_prefix}.attn",
+            )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        current_step: int,
+    ) -> torch.Tensor:
+        gqkv, _ = self.gqkv_proj(hidden_states)
+        g, q, k, v = gqkv.split(
+            [self.q_size, self.q_size, self.kv_size, self.kv_size], dim=-1
+        )
+        q, k = self.rotary_emb(positions, q, k)
+        attn_out = self.attn_per_step[str(current_step)](q, k, v)
+        # Sigmoid gate. Shapes: attn_out is (..., q_size); g is (..., q_size).
+        attn_out = torch.sigmoid(g) * attn_out
+        out, _ = self.o_proj(attn_out)
+        return out
+
+
+class HrmTextDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        layer_idx_in_stack: int,
+        stack_kind: str,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        # Attribute name `self_attn` matches HF's model class. The on-disk
+        # `attn.{gqkv_proj,o_proj}.weight` keys are renamed to
+        # `self_attn.{gqkv_proj,o_proj}.weight` by the `WeightsMapper` in
+        # `HrmTextForCausalLM` so vLLM's standard `AutoWeightsLoader`
+        # handles the rest.
+        self.self_attn = HrmTextAttention(
+            config=config,
+            layer_idx_in_stack=layer_idx_in_stack,
+            stack_kind=stack_kind,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.mlp = HrmTextMLP(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            bias=getattr(config, "mlp_bias", False),
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+        # Parameterless RMSNorm (HF main: HrmTextRMSNorm has no weight).
+        self.input_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, has_weight=False
+        )
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, has_weight=False
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        current_step: int,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            current_step=current_step,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class HrmTextStack(nn.Module):
+    """A single transformer stack — used twice (H and L)."""
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        stack_kind: str,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                HrmTextDecoderLayer(
+                    config=config,
+                    layer_idx_in_stack=i,
+                    stack_kind=stack_kind,
+                    cache_config=cache_config,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.layers.{i}",
+                )
+                for i in range(config.num_layers_per_stack)
+            ]
+        )
+        self.final_norm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, has_weight=False
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        current_step_base: int,
+    ) -> torch.Tensor:
+        for layer in self.layers:
+            hidden_states = layer(
+                positions=positions,
+                hidden_states=hidden_states,
+                current_step=current_step_base,
+            )
+        return self.final_norm(hidden_states)
+
+
+@support_torch_compile
+class HrmTextModel(nn.Module):
+    """Hierarchical recurrent transformer body.
+
+    Forward (matches HF main exactly,
+    src/transformers/models/hrm_text/modeling_hrm_text.py:495-547):
+
+        hidden_states_high_cycle = embed(input_ids) * embedding_scale
+        hidden_states_low_cycle = z_L_init.expand_as(hidden_states_high_cycle)
+        for high_cycle_idx in range(H_cycles):
+            for low_cycle_idx in range(L_cycles):
+                step = high_cycle_idx * (L_cycles + 1) + low_cycle_idx
+                hidden_states_low_cycle = L_module(
+                    hidden_states_low_cycle + hidden_states_high_cycle,
+                    current_step=step,
+                )
+            step = high_cycle_idx * (L_cycles + 1) + L_cycles
+            hidden_states_high_cycle = H_module(
+                hidden_states_high_cycle + hidden_states_low_cycle,
+                current_step=step,
+            )
+        return hidden_states_high_cycle
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        self.config = config
+
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.embed_tokens",
+        )
+        self.L_module = HrmTextStack(
+            config=config,
+            stack_kind="L",
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.L_module",
+        )
+        self.H_module = HrmTextStack(
+            config=config,
+            stack_kind="H",
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.H_module",
+        )
+        # Frozen learned initial L state. HF inits to zeros and sets
+        # requires_grad_(False); for inference we just load the tensor.
+        self.z_L_init = nn.Parameter(
+            torch.zeros(config.hidden_size), requires_grad=False
+        )
+
+        # Embedding scale: HF uses config.embedding_scale (default
+        # 1 / initializer_range = 50.0 when initializer_range=0.02). NOT
+        # sqrt(hidden_size) like Gemma.
+        self.embedding_scale = getattr(config, "embedding_scale", None)
+        if self.embedding_scale is None:
+            init_range = getattr(config, "initializer_range", 0.02)
+            self.embedding_scale = 1.0 / init_range
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids) * self.embedding_scale
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if inputs_embeds is None:
+            assert input_ids is not None
+            inputs_embeds = self.embed_input_ids(input_ids)
+
+        hidden_states_high_cycle = inputs_embeds
+        hidden_states_low_cycle = self.z_L_init.to(
+            dtype=hidden_states_high_cycle.dtype,
+            device=hidden_states_high_cycle.device,
+        ).expand_as(hidden_states_high_cycle)
+
+        H_cycles = self.config.H_cycles
+        L_cycles = self.config.L_cycles
+        for high_cycle_idx in range(H_cycles):
+            for low_cycle_idx in range(L_cycles):
+                step = high_cycle_idx * (L_cycles + 1) + low_cycle_idx
+                hidden_states_low_cycle = self.L_module(
+                    positions=positions,
+                    hidden_states=hidden_states_low_cycle + hidden_states_high_cycle,
+                    current_step_base=step,
+                )
+            step = high_cycle_idx * (L_cycles + 1) + L_cycles
+            hidden_states_high_cycle = self.H_module(
+                positions=positions,
+                hidden_states=hidden_states_high_cycle + hidden_states_low_cycle,
+                current_step_base=step,
+            )
+
+        return hidden_states_high_cycle
+
+
+class HrmTextForCausalLM(nn.Module):
+    """Hierarchical Reasoning Model — Text variant, causal LM.
+
+    Reference: src/transformers/models/hrm_text/modeling_hrm_text.py
+    """
+
+    # On-disk weight key remap: HF stores attention weights as
+    # `attn.{gqkv_proj,o_proj}.weight`; our model uses `self_attn.*`
+    # (matching HF's runtime model class). Both `gqkv_proj` (4-way fused
+    # gate/q/k/v) and `mlp.gate_up_proj` (2-way fused gate/up) are loaded
+    # directly via MergedColumnParallelLinear's fused-on-disk path; no
+    # packed_modules_mapping entries are needed.
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_substr={".attn.": ".self_attn."},
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        self.config = config
+        self.quant_config = quant_config
+
+        if vllm_config.parallel_config.pipeline_parallel_size > 1:
+            raise ValueError(
+                "HrmTextForCausalLM does not support pipeline parallelism."
+            )
+
+        self.model = HrmTextModel(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+
+        if config.tie_word_embeddings:
+            self.lm_head = self.model.embed_tokens
+        else:
+            self.lm_head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
+        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/registry.py

@@ -134,6 +134,7 @@ _TEXT_GENERATION_MODELS = {
     "GritLM": ("gritlm", "GritLM"),
     "Grok1ModelForCausalLM": ("grok1", "GrokForCausalLM"),
     "Grok1ForCausalLM": ("grok1", "GrokForCausalLM"),
+    "HrmTextForCausalLM": ("hrm_text", "HrmTextForCausalLM"),
     "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"),
     "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"),
     "HYV3ForCausalLM": ("hy_v3", "HYV3ForCausalLM"),

vllm/v1/engine/core.py

@@ -246,6 +246,28 @@ class EngineCore:
         # Get all kv cache needed by the model
         kv_cache_specs = self.model_executor.get_kv_cache_specs()
 
+        # Some layers (e.g. Prefix LM attention) run non-causally and tag their
+        # KV cache spec with ``non_causal=True``. The specs are collected here in
+        # the engine-core process (the same process that builds the scheduler),
+        # so this is the multiproc-safe place to translate that layer-level
+        # signal into a scheduling policy: chunked prefill and prefix caching
+        # both assume causal attention and would corrupt non-causal prefill.
+        if any(
+            getattr(spec, "non_causal", False)
+            for worker_specs in kv_cache_specs
+            for spec in worker_specs.values()
+        ):
+            if vllm_config.scheduler_config.enable_chunked_prefill:
+                logger.info(
+                    "Disabling chunked prefill: model has non-causal attention layers."
+                )
+                vllm_config.scheduler_config.enable_chunked_prefill = False
+            if vllm_config.cache_config.enable_prefix_caching:
+                logger.info(
+                    "Disabling prefix caching: model has non-causal attention layers."
+                )
+                vllm_config.cache_config.enable_prefix_caching = False
+
         has_kv_cache = any(kv_cache_spec for kv_cache_spec in kv_cache_specs)
         if has_kv_cache:
             if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:

vllm/v1/kv_cache_interface.py

@@ -219,6 +219,15 @@ class FullAttentionSpec(AttentionSpec):
     """
     attention_chunk_size: int | None = None
 
+    non_causal: bool = False
+    """
+    Whether the layer attends non-causally (e.g. Prefix LM). Carried on the
+    spec so the engine core, which collects specs from all workers before the
+    scheduler is built, can adjust scheduling policy (chunked prefill / prefix
+    caching) regardless of tensor-parallel layout. It does not affect the KV
+    cache layout itself.
+    """
+
     def __post_init__(self):
         if self.head_size_v is None:
             object.__setattr__(self, "head_size_v", self.head_size)
@@ -276,6 +285,9 @@ class FullAttentionSpec(AttentionSpec):
             page_size_padded=specs[0].page_size_padded,
             sliding_window=cls.merge_window_sizes(sliding_window),
             attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
+            # If any layer in the group is non-causal, treat the group as
+            # non-causal so the engine core disables incompatible scheduling.
+            non_causal=any(spec.non_causal for spec in specs),
         )
         for spec in specs:
             for f in fields(AttentionSpec):
@@ -701,6 +713,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
             page_size_padded=specs[0].page_size_padded,
             sliding_window=cls.merge_window_sizes(sliding_window),
             attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
+            non_causal=any(spec.non_causal for spec in specs),
         )
         for spec in specs:
             for f in fields(AttentionSpec):