chore: import upstream snapshot with attribution

python/sglang/srt/models/hrm_text.py

@@ -0,0 +1,477 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Inference-only HRM-Text (Hierarchical Reasoning Model -- Text) model.
+
+Reference: transformers/models/hrm_text (transformers >= 5.9.0).
+
+HRM-Text runs a hierarchical recurrent forward over two transformer stacks
+(``H`` slow, ``L`` fast) in nested loops. Each recurrence step gets its own KV
+cache slot via a unique ``RadixAttention(layer_id=...)``; the global index for
+``(step, layer)`` is ``step * num_layers_per_stack + layer``. The total slot
+count ``num_layers_per_stack * H_cycles * (L_cycles + 1)`` equals the HF config
+``num_hidden_layers`` after ``__post_init__`` inflation, exposed by
+``ModelConfig`` as ``num_attention_layers``.
+
+PrefixLM (prompt bidirectional at prefill, causal at decode) uses
+``AttentionType.DECODER_BIDIRECTIONAL``, which only the Triton backend honors
+and only with cuda graph / chunked prefill / radix cache off --
+``ModelRunner.model_specific_adjustment`` forces those for this model.
+
+On-disk ``attn.gqkv_proj.weight`` is fused ``[gate | q | k | v]`` rows and
+``mlp.gate_up_proj`` is ``[gate | up]``; both load directly via
+``MergedColumnParallelLinear``'s fused-on-disk auto-split path.
+"""
+
+import logging
+from typing import Iterable, Optional, Tuple
+
+import torch
+from torch import nn
+from transformers import PretrainedConfig
+
+from sglang.srt.layers.activation import SiluAndMul
+from sglang.srt.layers.layernorm import RMSNorm
+from sglang.srt.layers.linear import (
+    MergedColumnParallelLinear,
+    RowParallelLinear,
+)
+from sglang.srt.layers.logits_processor import LogitsProcessor
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.radix_attention import AttentionType, RadixAttention
+from sglang.srt.layers.rotary_embedding import get_rope
+from sglang.srt.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils import add_prefix
+
+logger = logging.getLogger(__name__)
+
+
+def _num_layers_per_stack(config: PretrainedConfig) -> int:
+    """Layers in one (H or L) stack.
+
+    Native configs store this in ``num_layers_per_stack`` after ``__post_init__``
+    rewrites ``num_hidden_layers`` to the inflated total; fall back to deriving
+    it for non-native configs.
+    """
+    nlps = getattr(config, "num_layers_per_stack", None)
+    if nlps is not None:
+        return int(nlps)
+    return config.num_hidden_layers // (config.H_cycles * (config.L_cycles + 1))
+
+
+def _steps_used(config: PretrainedConfig, stack_kind: str) -> list[int]:
+    """Recurrence steps at which a stack runs.
+
+    L runs at ``h*(L+1)+l`` (``0<=h<H, 0<=l<L``); H runs at the trailing
+    ``h*(L+1)+L``. Disjoint, so each ``(step, layer)`` maps to a unique KV index.
+    """
+    H_cycles = config.H_cycles
+    L_cycles = config.L_cycles
+    if stack_kind == "L":
+        return [
+            h * (L_cycles + 1) + low_idx
+            for h in range(H_cycles)
+            for low_idx in range(L_cycles)
+        ]
+    return [h * (L_cycles + 1) + L_cycles for h in range(H_cycles)]
+
+
+class HrmTextMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: Optional[QuantizationConfig] = 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=False,
+            quant_config=quant_config,
+            prefix=add_prefix("gate_up_proj", prefix),
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("down_proj", prefix),
+        )
+        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):
+    """Self-attention block; projection weights are shared across recurrence
+    steps, while per-step KV slots come from weightless ``RadixAttention``
+    instances keyed by step in ``self.attn``."""
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        layer_idx_in_stack: int,
+        stack_kind: str,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        tp_size = get_parallel().tp_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 hardcodes MHA (kv heads == q heads); no GQA.
+        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
+
+        # Fused [gate | q | k | v] on disk; MHA only (GQA would need
+        # QKVParallelLinear's q/k/v shard replication).
+        per_head_size = self.total_num_heads * self.head_dim
+        self.gqkv_proj = MergedColumnParallelLinear(
+            self.hidden_size,
+            [per_head_size] * 4,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("gqkv_proj", prefix),
+        )
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            self.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("o_proj", prefix),
+        )
+
+        # rope_parameters (HF 5.9.0) or flat rope_theta for older configs.
+        rope_parameters = getattr(config, "rope_parameters", None) or {}
+        rope_theta = rope_parameters.get("rope_theta", None)
+        if rope_theta is None:
+            rope_theta = getattr(config, "rope_theta", 10000.0)
+        rope_type = rope_parameters.get("rope_type", "default")
+        # "default" rope = no scaling; pass the dict through otherwise.
+        rope_scaling = None if rope_type in ("default", None) else rope_parameters
+        self.rotary_emb = get_rope(
+            head_size=self.head_dim,
+            rotary_dim=self.head_dim,
+            max_position=config.max_position_embeddings,
+            base=rope_theta,
+            is_neox_style=True,
+            rope_scaling=rope_scaling,
+        )
+
+        # One weightless RadixAttention per step, each with a unique layer_id
+        # (= global KV slot) so the recurrent forward writes disjoint slots.
+        num_layers_per_stack = _num_layers_per_stack(config)
+        self.attn = nn.ModuleDict()
+        for step in _steps_used(config, stack_kind):
+            global_idx = step * num_layers_per_stack + layer_idx_in_stack
+            self.attn[str(step)] = RadixAttention(
+                num_heads=self.num_heads,
+                head_dim=self.head_dim,
+                scaling=self.scaling,
+                num_kv_heads=self.num_kv_heads,
+                layer_id=global_idx,
+                attn_type=AttentionType.DECODER_BIDIRECTIONAL,
+                quant_config=quant_config,
+                prefix=add_prefix(f"attn.{step}", prefix),
+            )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+        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[str(current_step)](q, k, v, forward_batch)
+        # Sigmoid gate (HrmText / Qwen3Next style).
+        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,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.self_attn = HrmTextAttention(
+            config=config,
+            layer_idx_in_stack=layer_idx_in_stack,
+            stack_kind=stack_kind,
+            quant_config=quant_config,
+            prefix=add_prefix("self_attn", prefix),
+        )
+        self.mlp = HrmTextMLP(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            prefix=add_prefix("mlp", prefix),
+        )
+        # Parameterless RMSNorm (HF 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,
+        forward_batch: ForwardBatch,
+        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,
+            forward_batch=forward_batch,
+            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 -- instantiated twice (H and L)."""
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        stack_kind: str,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        num_layers_per_stack = _num_layers_per_stack(config)
+        self.layers = nn.ModuleList(
+            [
+                HrmTextDecoderLayer(
+                    config=config,
+                    layer_idx_in_stack=i,
+                    stack_kind=stack_kind,
+                    quant_config=quant_config,
+                    prefix=add_prefix(f"layers.{i}", prefix),
+                )
+                for i in range(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,
+        forward_batch: ForwardBatch,
+        current_step_base: int,
+    ) -> torch.Tensor:
+        for layer in self.layers:
+            hidden_states = layer(
+                positions=positions,
+                hidden_states=hidden_states,
+                forward_batch=forward_batch,
+                current_step=current_step_base,
+            )
+        return self.final_norm(hidden_states)
+
+
+class HrmTextModel(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=add_prefix("embed_tokens", prefix),
+        )
+        self.L_module = HrmTextStack(
+            config=config,
+            stack_kind="L",
+            quant_config=quant_config,
+            prefix=add_prefix("L_module", prefix),
+        )
+        self.H_module = HrmTextStack(
+            config=config,
+            stack_kind="H",
+            quant_config=quant_config,
+            prefix=add_prefix("H_module", prefix),
+        )
+        # Frozen learned initial low-cycle state (disk key `model.z_L_init`).
+        self.z_L_init = nn.Parameter(
+            torch.zeros(config.hidden_size), requires_grad=False
+        )
+
+        # HF uses config.embedding_scale (= 1 / initializer_range), NOT
+        # sqrt(hidden_size).
+        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 forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        input_embeds: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if input_embeds is not None:
+            hidden_states_high_cycle = input_embeds
+        else:
+            hidden_states_high_cycle = self.embed_tokens(input_ids)
+        hidden_states_high_cycle = hidden_states_high_cycle * self.embedding_scale
+
+        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)
+            .contiguous()
+        )
+
+        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,
+                    forward_batch=forward_batch,
+                    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,
+                forward_batch=forward_batch,
+                current_step_base=step,
+            )
+
+        return hidden_states_high_cycle
+
+
+class HrmTextForCausalLM(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.quant_config = quant_config
+
+        self.model = HrmTextModel(
+            config=config,
+            quant_config=quant_config,
+            prefix=add_prefix("model", prefix),
+        )
+
+        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=add_prefix("lm_head", prefix),
+            )
+        self.logits_processor = LogitsProcessor(config)
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        input_embeds: Optional[torch.Tensor] = None,
+    ):
+        hidden_states = self.model(input_ids, positions, forward_batch, input_embeds)
+        return self.logits_processor(
+            input_ids, hidden_states, self.lm_head, forward_batch
+        )
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        # Disk keys use `.attn.`; rename to our `.self_attn.`. The per-step
+        # RadixAttention modules hold no params, and disk tensors are already
+        # fused so no stacked_params_mapping is needed.
+        params_dict = dict(self.named_parameters())
+        for name, loaded_weight in weights:
+            if ".attn." in name:
+                name = name.replace(".attn.", ".self_attn.", 1)
+            if name not in params_dict:
+                continue
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+
+
+EntryClass = HrmTextForCausalLM