From a52205bccfc7c270dd445cf253f18deb3778e334 Mon Sep 17 00:00:00 2001 From: yifei wu <50608184+abcd1927@users.noreply.github.com> Date: Tue, 16 Jun 2026 22:41:41 +0800 Subject: [PATCH] =?UTF-8?q?[Model]=20Add=20HrmTextForCausalLM=20(Hierarchi?= =?UTF-8?q?cal=20Reasoning=20Model=20=E2=80=94=20Text)=20(#43098)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Wuyifei 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 | 1 + tests/models/registry.py | 4 + .../layers/attention/__init__.py | 4 + .../attention/prefill_prefix_lm_attention.py | 86 +++ vllm/model_executor/models/hrm_text.py | 527 ++++++++++++++++++ vllm/model_executor/models/registry.py | 1 + vllm/v1/engine/core.py | 22 + vllm/v1/kv_cache_interface.py | 13 + 8 files changed, 658 insertions(+) create mode 100644 vllm/model_executor/layers/attention/prefill_prefix_lm_attention.py create mode 100644 vllm/model_executor/models/hrm_text.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 21e801a4232..3b9705db4dc 100644 --- a/docs/models/supported_models.md +++ b/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` | ✅︎ | ✅︎ | diff --git a/tests/models/registry.py b/tests/models/registry.py index 50a8ec76952..d0af28368ee 100644 --- a/tests/models/registry.py +++ b/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 diff --git a/vllm/model_executor/layers/attention/__init__.py b/vllm/model_executor/layers/attention/__init__.py index 1be9f77427d..ca3574164d5 100644 --- a/vllm/model_executor/layers/attention/__init__.py +++ b/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", ] diff --git a/vllm/model_executor/layers/attention/prefill_prefix_lm_attention.py b/vllm/model_executor/layers/attention/prefill_prefix_lm_attention.py new file mode 100644 index 00000000000..475184fa6a5 --- /dev/null +++ b/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 diff --git a/vllm/model_executor/models/hrm_text.py b/vllm/model_executor/models/hrm_text.py new file mode 100644 index 00000000000..a7546b0cc44 --- /dev/null +++ b/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) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index ecd31f2dc01..2d278973fd7 100644 --- a/vllm/model_executor/models/registry.py +++ b/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"), diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index bf89f3e9d5c..0803c98294b 100644 --- a/vllm/v1/engine/core.py +++ b/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: diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 2f8048c7966..9528fb65af1 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/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):