diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index a0d0ce61c..b1b787b7b 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -915,6 +915,19 @@ class ModelConfig: if "IQuestLoopCoderForCausalLM" in self.hf_config.architectures: loop_num = getattr(self.hf_text_config, "loop_num", 1) self.num_attention_layers = int(self.num_hidden_layers * int(loop_num)) + if "HrmTextForCausalLM" in self.hf_config.architectures: + # Compute KV slot count explicitly: native 5.9.0 configs inflate + # num_hidden_layers to this in __post_init__, but non-native ones + # may carry the raw per-stack count. + H_cycles = self.hf_text_config.H_cycles + L_cycles = self.hf_text_config.L_cycles + num_layers_per_stack = ( + getattr(self.hf_text_config, "num_layers_per_stack", None) + or self.num_hidden_layers + ) + self.num_attention_layers = ( + int(num_layers_per_stack) * H_cycles * (L_cycles + 1) + ) if "WhisperForConditionalGeneration" in self.hf_config.architectures: # Whisper has unique layer ID scheme: # - Encoder self-attention: 0 to encoder_layers-1 (no KV cache) @@ -1449,6 +1462,18 @@ class ModelConfig: # while for GLM-4.6v, it is 'glm4v_moe_vision'. ) needs_tf_v5 = is_glm_46vmoe + # Older transformers lacks the native hrm_text config, so it silently + # falls back to TransformersForCausalLM and loads fused weights as junk. + architectures = getattr(self.hf_config, "architectures", []) or [] + is_hrm_text = getattr(self.hf_config, "model_type", None) == "hrm_text" or ( + "HrmTextForCausalLM" in architectures + ) + if is_hrm_text and version.parse(tf_version_str) < version.parse("5.9.0"): + raise ValueError( + f"HRM-Text (model type {self.hf_config.model_type!r}) requires " + f"transformers >= 5.9.0, but {tf_version_str} is installed. " + "Please upgrade transformers." + ) tf_version = version.parse(tf_version_str) required_version = version.parse("5.0.0dev0") diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 08b6a550f..bff9cec8b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1105,6 +1105,39 @@ class ModelRunner(ModelRunnerKVCacheMixin): def model_specific_adjustment(self): server_args = self.server_args + # HRM-Text needs bidirectional prompt attention (prefill), which only the + # Triton backend honors and only with cuda graph / chunked prefill off + # (TritonAttnBackend.allow_bidirectional_attention_in_extend). Radix cache + # is also unsafe: the recurrent forward writes direction-dependent KV + # across many slots. + hf_config = self.model_config.hf_config + is_hrm_text = getattr( + hf_config, "model_type", None + ) == "hrm_text" or "HrmTextForCausalLM" in getattr( + hf_config, "architectures", [] + ) + # prefix_lm defaults to True upstream; defaulting False would skip the + # bidirectional-attention forcing and silently produce junk output. + is_prefix_lm_recurrent = is_hrm_text and getattr(hf_config, "prefix_lm", True) + if is_prefix_lm_recurrent: + if server_args.attention_backend not in (None, "triton"): + logger.warning( + f"Overriding --attention-backend " + f"{server_args.attention_backend!r} -> 'triton': only the " + "Triton backend supports HRM-Text's bidirectional prefix " + "attention." + ) + server_args.attention_backend = "triton" + server_args.chunked_prefill_size = -1 + server_args.disable_radix_cache = True + server_args.disable_cuda_graph = True + logger.warning( + "HRM-Text (prefix_lm) detected: forcing --attention-backend " + "triton, --chunked-prefill-size -1, --disable-radix-cache, and " + "--disable-cuda-graph for correctness of the bidirectional " + "prompt attention." + ) + if self.is_multimodal: if not self.is_multimodal_chunked_prefill_supported: server_args.chunked_prefill_size = -1 diff --git a/python/sglang/srt/models/hrm_text.py b/python/sglang/srt/models/hrm_text.py new file mode 100644 index 000000000..83010507d --- /dev/null +++ b/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.distributed import get_tensor_model_parallel_world_size +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.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 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_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 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