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

[abcd1927]

HRM-Text shipped in transformers 5.9.0 (released 2026-05-20, merged to main via huggingface/transformers#46025). The model performs a hierarchical recurrent forward over two transformer stacks (H slow, L fast) inside nested H/L cycle loops; each recurrence step consumes a distinct KV cache slot, matching transformers' cycle_offset formula. Attention uses a sigmoid gate (Qwen3Next-style) and a fused gate+q+k+v projection on disk.

This PR is fully scoped to the new model file, the registry, and the supported-models doc — no changes to attention backends, the AttentionType enum, or Attention itself. Pipeline parallelism and LoRA are deferred to follow-up PRs (each currently raises at construction).

Architecture

• HrmTextAttention keeps gate_proj (ColumnParallelLinear) separate from qkv_proj (QKVParallelLinear) so both shard cleanly along the head axis under TP. The HF on-disk schema fuses them into a single attn.gqkv_proj.weight (rows concatenated as [gate | q | k | v]); HrmTextForCausalLM.load_weights splits that tensor at load time before passing it to AutoWeightsLoader. Each recurrence step gets its own Attention instance under a nn.ModuleDict, with a unique prefix (...layers.{global_idx}.attn), so vLLM allocates a distinct KV cache slot per cycle step. Total slots match H_cycles * (L_cycles + 1) * num_layers_per_stack, equal to config.num_hidden_layers after HF's __post_init__ inflation.

• _create_hrm_attention_backend wraps an attention backend's metadata builder so causal=False is set unconditionally on every build. Mirrors EncoderOnlyAttention's subclass_attention_backend pattern, but keeps the KV cache (the recurrent forward needs to reuse it).

• HrmTextModel.forward matches HF main exactly: nested H-outer L-inner loop, z_L = L(z_L + z_H) then z_H = H(z_H + z_L). The Python loop is unrolled by @support_torch_compile since cycle counts are config constants.

Why unconditional causal=False

vLLM v1 schedules continuous batching. Even without chunked prefill, a single build() call can see a mixed batch where some requests are decoding and others are entering prefill on this step. Gating on is_prefilling.all() keeps causal=True in that mixed case, silently running the newly-prefilling requests as pure causal — which diverges from HRM-Text's PrefixLM training distribution.

Unconditional causal=False is correct because:

• Prefill rows (query_len=N): bidirectional, matching HF main with token_type_ids=ones_like(input_ids).
• Decode rows (query_len=1): causal=True and causal=False are mathematically identical (no future tokens to mask).
• FlashAttention varlen kernels apply causal per sub-sequence via cu_seqlens_q, so a single global flag does not cause cross-prompt contamination.

Empirical validation

sapientinc/HRM-Text-1B, vLLM matches HF transformers main with token_type_ids=ones_like(input_ids) to within statistical noise.

Benchmark	HF (token_type_ids=1)	vLLM (this PR)
GSM8K, first 50	82.0% (41/50)	82.0% (41/50)
GSM8K, full (1319)	—	84.99%
GSM8K, tp=1 vs tp=2	—	identical (both 82.0% on first 50)
DROP, first 100 EM	85.0%	85.0%, 100/100 normalized text match
DROP, full (9536)	—	79.98% EM

On-disk weight schema

HF transformers writes HrmText checkpoints with these keys (per the "hrm_text" entry in transformers/conversion_mapping.py):

model.embed_tokens.weight
model.z_L_init
model.{H,L}_module.layers.{i}.attn.gqkv_proj.weight  ([gate|q|k|v] dim 0)
model.{H,L}_module.layers.{i}.attn.o_proj.weight
model.{H,L}_module.layers.{i}.mlp.{gate_up,down}_proj.weight

load_weights splits attn.gqkv_proj.weight into attn.gate_proj.weight (first num_heads * head_dim rows) and attn.qkv_proj.weight (remaining [q | k | v] rows) at load time, so vLLM's standard QKVParallelLinear weight loader handles TP partitioning of q/k/v without unfusing the gate.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[mergify]

Documentation preview: https://vllm--43098.org.readthedocs.build/en/43098/

[gemini-code-assist]

Code Review

This pull request implements support for the HRM-Text (Hierarchical Reasoning Model) architecture, featuring a hierarchical recurrent structure with dual transformer stacks and a PrefixLM attention pattern. The changes include the core model implementation, registration in the model executor, documentation updates, and new tests for generation and batching. A review comment identifies that the min_transformers_version in the test registry is set to a placeholder value ('5.9') and suggests updating it to a current version to ensure tests are executed correctly.

[hmellor]

General comment, can we be more explicit with variable names? In the HF src the names are much easier to understand high_cycle_idx & low_cycle_idx for example.

[abcd1927]

Done. Renamed loop vars and hidden states to match HF. Kept H_cycles/L_cycles (config fields) and H_module/L_module (module names) since HF uses the same.

[Isotr0py]

Actually, we have multimodal prefix-lm (partial bidi) support already:

vllm/vllm/v1/worker/gpu_model_runner.py

Lines 2453 to 2479 in 309385a

	if self.is_mm_prefix_lm:
	req_doc_ranges = {}
	# Gemma4 bidi: skip ranges that exceed the sliding
	# window. When image tokens > sliding_window, bidi causes
	# early image tokens to attend to the entire image
	# (e.g. 6 → 1092 targets), degrading spatial precision.
	# Per-range filtering keeps bidi for small images/video
	# frames while skipping oversized images.
	hf_text_config = self.model_config.hf_text_config
	_bidi_sw = getattr(hf_text_config, "sliding_window", None)
	for req_id in self.input_batch.req_ids:
	image_doc_ranges = []
	req_state = self.requests[req_id]
	for mm_feature in req_state.mm_features:
	pos_info = mm_feature.mm_position
	img_doc_range = pos_info.extract_embeds_range()
	for r in img_doc_range:
	if _bidi_sw is not None and (r[1] - r[0] + 1) > _bidi_sw:
	continue
	image_doc_ranges.append(r)
	req_idx = self.input_batch.req_id_to_index[req_id]
	req_doc_ranges[req_idx] = image_doc_ranges
	# Set mm_prefix_range for all attention metadata
	self._set_mm_prefix_range_for_metadata(attn_metadata, req_doc_ranges)

Perhaps you can reuse this code path with some modifications for text-only models. (Setting doc_ranges to prefill range)

[mergify]

Hi @abcd1927, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[abcd1927]

Implemented via the KV cache spec channel as discussed. PrefixLMAttention.get_kv_cache_spec tags its spec with non_causal=True, and the engine core scans the specs it already collects from all workers (before the scheduler is built) and disables chunked prefill + prefix caching when any layer is non-causal. TP-safe — the signal crosses the worker→engine boundary on the pickled spec, applied to the same config object the scheduler reads.

Verified on H100: GSM8K (50 samples) gives identical accuracy at TP=1 and TP=2, with the auto-disable confirmed in the engine-core logs at both sizes.

@Isotr0py would appreciate your review on the spec-channel approach.

[Isotr0py]

I prefer to rename this class as PrefillPrefixLMAttention, because this will only work for models that treat the whole prefill sequence as bidi prefix, while PrefixLMAttention will cause confusion between Paligemma/Gemma4's multimodal prefix bidi attention.

[abcd1927]

Done — renamed to PrefillPrefixLMAttention (and the file to prefill_prefix_lm_attention.py to match). Thanks!