diff --git "a/inference/model.py" "b/inference/model.py" new file mode 100644--- /dev/null +++ "b/inference/model.py" @@ -0,0 +1,2500 @@ +"""Step-5 standalone text and vision inference. + +Includes configuration, packed-token DSA attention, TP/EP layers and vision +prefill. Compute kernels live in kernel.py; sliding attention uses PyTorch SDPA. +""" + +# Includes StepFun code licensed under the Apache License, Version 2.0. + +from __future__ import annotations + +import threading +from typing import Any, Sequence + +import torch +import torch.distributed as dist +from torch import nn +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_outputs import CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel + +from vision_encoder import Step5VisionConfig, Step5VisionEncoder + +from kernel import ( + DSAGeometry, + DSALayerCache, + Step5SparseIndexer, + build_rope_cache, + decode_metadata, + prefill_metadata, + update_summaries_decode, + update_summaries_prefill, + sparse_attention_decode, + sparse_attention_prefill, + fused_qknorm_rope, + linear_fp8_or_bf16, + clamped_swiglu as triton_clamped_swiglu, + weighted_topk_gather, +) + + +# Configuration and checkpoint serialization. + + +# Incremented whenever the on-disk TP tensor layout changes incompatibly. Version 2 keeps +# v1's corrected GQA/provider-group mapping and additionally TP-shards the shared expert. +# The latter is required because the deployed MoE combines the local shared and routed +# branches in FP32 and performs one all-reduce; replicating the shared expert would sum TP +# identical copies. +STEP5_TP_LAYOUT_VERSION = "gqa-provider-shared-tp-v2" + + +def _normalize_config_identity(kwargs, *, model_type, architecture): + """Read compatible checkpoint identities and serialize this implementation.""" + legacy_type = {"step5": "step4", "step5_vl": "step4_vl"}[model_type] + supplied_type = kwargs.pop("model_type", model_type) + if supplied_type not in (model_type, legacy_type): + raise ValueError(f"unsupported model_type {supplied_type!r} for {model_type}") + kwargs["architectures"] = [architecture] + kwargs["auto_map"] = { + "AutoConfig": "model.Step5VLConfig" if model_type == "step5_vl" else "model.Step5Config", + "AutoModelForCausalLM": f"model.{architecture}", + } + if model_type == "step5_vl": + kwargs["auto_map"]["AutoModelForImageTextToText"] = f"model.{architecture}" + + +class Step5SparseConfig: + """The ``sparse_config`` block, as a plain attribute holder. + + Kept as its own object rather than flattened onto the parent because the DSA code reads + it as a unit, and because the checkpoint nests it -- flattening would require rewriting + ``config.json`` on the way in and out. + """ + + def __init__( + self, + enabled: bool = True, + proxy_dim: int = 256, + sparse_indexer_rope_dim: int = 32, + sparse_indexer_use_rope: bool = True, + sparse_indexer_num_heads: int = 16, + sparse_indexer_num_k_heads: int = 1, + sparse_indexer_q_norm_type: str = "rmsnorm", + sparse_indexer_k_norm_type: str = "layernorm", + sparse_indexer_csa_z_norm_type: str = "none", + sparse_indexer_softmax_variant: str = "ssmax", + sparse_indexer_ssmax_s_granularity: str = "q_head", + num_provider_groups: int = 4, + topk: int = 512, + region_block_size: int = 8, + compression_method: str = "csa_block_compress", + attention_impl: str = "sparse_gqa", + apply_to_layer_types: list[str] | None = None, + **kwargs: Any, + ) -> None: + self.enabled = enabled + self.proxy_dim = proxy_dim + self.sparse_indexer_rope_dim = sparse_indexer_rope_dim + self.sparse_indexer_use_rope = sparse_indexer_use_rope + self.sparse_indexer_num_heads = sparse_indexer_num_heads + self.sparse_indexer_num_k_heads = sparse_indexer_num_k_heads + self.sparse_indexer_q_norm_type = sparse_indexer_q_norm_type + self.sparse_indexer_k_norm_type = sparse_indexer_k_norm_type + self.sparse_indexer_csa_z_norm_type = sparse_indexer_csa_z_norm_type + # Preserve checkpoint metadata for round trips. The indexer uses weighted + # ReLU scores; these softmax settings do not change its arithmetic. + self.sparse_indexer_softmax_variant = sparse_indexer_softmax_variant + self.sparse_indexer_ssmax_s_granularity = sparse_indexer_ssmax_s_granularity + self.num_provider_groups = num_provider_groups + self.topk = topk + self.region_block_size = region_block_size + self.compression_method = compression_method + self.attention_impl = attention_impl + self.apply_to_layer_types = apply_to_layer_types or ["full_attention"] + for key, value in kwargs.items(): + setattr(self, key, value) + + def to_dict(self) -> dict[str, Any]: + return dict(self.__dict__) + + +class Step5Config(PretrainedConfig): + """Checkpoint configuration for the packed Step-5 text decoder. + + Per-layer arrays include trailing MTP entries; ordinary generation builds + only num_hidden_layers. Attention type and MoE placement are independent. + RoPE parameters are per layer, with scaling gated by yarn_only_types. + """ + + model_type = "step5" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + hidden_size: int = 4096, + intermediate_size: int = 13824, + num_hidden_layers: int = 92, + vocab_size: int = 128896, + num_attention_heads: int = 64, + num_attention_groups: int = 4, + head_dim: int = 192, + rms_norm_eps: float = 1e-5, + norm_dtype: str = "float32", + fp32_residual_connection: bool = True, + use_head_wise_attn_gate: bool = True, + rope_theta: float | list[float] = 10000.0, + partial_rotary_factors: list[float] | None = None, + layer_types: list[str] | None = None, + sliding_window: int | None = 512, + rope_scaling: dict[str, Any] | None = None, + yarn_only_types: list[str] | None = None, + max_position_embeddings: int = 524288, + use_moe: bool = True, + moe_intermediate_size: int = 1536, + share_expert_dim: int = 1536, + moe_num_experts: int = 352, + moe_top_k: int = 8, + moe_layer_list: list[int] | None = None, + use_moe_router_bias: bool = True, + need_fp32_gate: bool = True, + norm_expert_weight: bool = True, + moe_router_scaling_factor: float = 3.0, + swiglu_limits: list[float] | None = None, + swiglu_limits_shared: list[float] | None = None, + num_nextn_predict_layers: int = 1, + sparse_config: dict[str, Any] | Step5SparseConfig | None = None, + tie_word_embeddings: bool = False, + bos_token_id: int = 0, + eos_token_id: int | list[int] | None = None, + tp_size: int = 1, + tp_layout_version: str | None = None, + **kwargs: Any, + ) -> None: + if tp_size < 1: + raise ValueError(f"tp_size must be positive, got {tp_size}") + self.tp_size = tp_size + self.tp_layout_version = tp_layout_version + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.vocab_size = vocab_size + self.num_attention_heads = num_attention_heads + self.num_attention_groups = num_attention_groups + self.head_dim = head_dim + self.rms_norm_eps = rms_norm_eps + self.norm_dtype = norm_dtype + self.fp32_residual_connection = fp32_residual_connection + self.use_head_wise_attn_gate = use_head_wise_attn_gate + self.max_position_embeddings = max_position_embeddings + self.sliding_window = sliding_window + self.rope_scaling = rope_scaling + self.yarn_only_types = yarn_only_types or [] + + # The checkpoint has one entry per layer *including* the MTP layers. Transformers 5 + # validates ``layer_types`` as a main-stack-only field, so retain the trailing MTP + # entries separately and join them again in ``to_dict``. Runtime access continues + # through ``layer_type`` for both the main stack and MTP layers. + total_layers = num_hidden_layers + num_nextn_predict_layers + serialized_layer_types = ( + list(layer_types) + if layer_types is not None + else ["full_attention"] * total_layers + ) + if len(serialized_layer_types) != total_layers: + raise ValueError( + "layer_types must contain one entry per main and next-token " + f"prediction layer: expected {total_layers}, got " + f"{len(serialized_layer_types)}" + ) + self.layer_types = serialized_layer_types[:num_hidden_layers] + self._mtp_layer_types = serialized_layer_types[num_hidden_layers:] + self.rope_theta = ( + list(rope_theta) + if isinstance(rope_theta, (list, tuple)) + else [float(rope_theta)] * total_layers + ) + self.partial_rotary_factors = partial_rotary_factors or [1.0] * total_layers + + self.use_moe = use_moe + self.moe_intermediate_size = moe_intermediate_size + self.share_expert_dim = share_expert_dim + self.moe_num_experts = moe_num_experts + self.moe_top_k = moe_top_k + self.moe_layer_list = moe_layer_list or [] + self.use_moe_router_bias = use_moe_router_bias + self.need_fp32_gate = need_fp32_gate + self.norm_expert_weight = norm_expert_weight + self.moe_router_scaling_factor = moe_router_scaling_factor + self.swiglu_limits = swiglu_limits or [] + self.swiglu_limits_shared = swiglu_limits_shared or [] + self.num_nextn_predict_layers = num_nextn_predict_layers + + if isinstance(sparse_config, Step5SparseConfig): + self.sparse_config = sparse_config + elif sparse_config: + self.sparse_config = Step5SparseConfig(**sparse_config) + else: + self.sparse_config = None + + self._validate_parallel_geometry() + + # ``num_experts_per_tok`` and ``moe_top_k`` are the same quantity under two names; + # the checkpoint carries both. Trust ``moe_top_k`` and drop the duplicate rather + # than storing two fields that can disagree. + kwargs.pop("num_experts_per_tok", None) + self._consume_quantization_config(kwargs) + _normalize_config_identity( + kwargs, model_type=self.model_type, architecture="Step5ForCausalLM", + ) + super().__init__( + tie_word_embeddings=tie_word_embeddings, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id if eos_token_id is not None else [1, 2], + **kwargs, + ) + + @property + def num_key_value_heads(self) -> int: + """HuggingFace's name for ``num_attention_groups``.""" + return self.num_attention_groups + + @property + def num_experts_per_tok(self) -> int: + return self.moe_top_k + + @property + def tp(self) -> int: + """Active tensor-parallel degree (defaults to 1, i.e. single-GPU).""" + return self.tp_size + + @property + def num_attention_heads_per_rank(self) -> int: + """Heads owned by each rank. ``q_proj``/``o_proj`` slice the head dim by this.""" + return self.num_attention_heads // self.tp + + @property + def heads_per_group_per_rank(self) -> int: + """Local attention Q heads served by each local KV group.""" + return self.num_attention_heads_per_rank // self.num_kv_groups_per_rank + + @property + def num_kv_groups_per_rank(self) -> int: + """KV heads needed by one rank's contiguous Q-head slice. + + With step5's 64 Q heads, 4 KV heads and TP=8 this is one, not four. Ranks + ``0,1`` share global KV0, ranks ``2,3`` share KV1, and so on. Replicating all four + KV heads and interpreting eight local Q heads as four two-head groups changes the + GQA function for 48 of the 64 heads. + """ + return self._groups_per_rank(self.num_attention_groups) + + @property + def num_provider_groups_per_rank(self) -> int: + """DSA provider groups aligned with the local KV groups.""" + if self.sparse_config is None: + return 0 + return self._groups_per_rank(self.sparse_config.num_provider_groups) + + @property + def sparse_indexer_heads_per_provider_group(self) -> int: + """The checkpoint carries four indexer Q/W heads per provider group.""" + if self.sparse_config is None: + return 0 + return ( + self.sparse_config.sparse_indexer_num_heads + // self.sparse_config.num_provider_groups + ) + + @property + def sparse_indexer_num_heads_per_rank(self) -> int: + """Indexer Q/W heads stored on this rank (four for the production TP=8 layout).""" + return ( + self.num_provider_groups_per_rank + * self.sparse_indexer_heads_per_provider_group + ) + + @property + def intermediate_per_rank(self) -> int: + """Dense MLP gate/up output dim and down input dim, per rank.""" + return self.intermediate_size // self.tp + + @property + def moe_intermediate_per_rank(self) -> int: + return self.moe_intermediate_size // self.tp + + @property + def share_expert_per_rank(self) -> int: + return self.share_expert_dim // self.tp + + @property + def vocab_per_rank(self) -> int: + return self.vocab_size // self.tp + + def _groups_per_rank(self, total_groups: int) -> int: + """Number of contiguous GQA/provider groups intersecting one rank's Q slice. + + Supported TP layouts either partition groups across ranks (``tp <= groups``) or + duplicate one group over several ranks (``tp > groups``). Requiring divisibility + in one direction prevents a rank's Q slice from straddling a fractional group. + """ + if total_groups % self.tp == 0: + return total_groups // self.tp + if self.tp % total_groups == 0: + return 1 + raise ValueError( + f"tp_size={self.tp} and group count {total_groups} are incompatible: " + "one must divide the other" + ) + + def _validate_parallel_geometry(self) -> None: + if self.num_attention_heads % self.tp: + raise ValueError( + f"num_attention_heads={self.num_attention_heads} is not divisible by " + f"tp_size={self.tp}" + ) + if self.num_attention_heads % self.num_attention_groups: + raise ValueError( + f"num_attention_heads={self.num_attention_heads} is not divisible by " + f"num_attention_groups={self.num_attention_groups}" + ) + + local_kv = self.num_kv_groups_per_rank + local_q = self.num_attention_heads_per_rank + if local_q % local_kv: + raise ValueError( + f"{local_q} local Q heads do not divide into {local_kv} local KV groups" + ) + + sparse = self.sparse_config + if sparse is not None: + if sparse.num_provider_groups != self.num_attention_groups: + raise ValueError( + "DSA provider groups must align with attention KV groups: " + f"{sparse.num_provider_groups} != {self.num_attention_groups}" + ) + if sparse.sparse_indexer_num_heads % sparse.num_provider_groups: + raise ValueError( + f"{sparse.sparse_indexer_num_heads} sparse indexer heads do not divide " + f"into {sparse.num_provider_groups} provider groups" + ) + if self.num_provider_groups_per_rank != local_kv: + raise ValueError( + "local DSA provider groups must align with local KV groups: " + f"{self.num_provider_groups_per_rank} != {local_kv}" + ) + + if self.tp > 1 and self.tp_layout_version != STEP5_TP_LAYOUT_VERSION: + raise ValueError( + "incompatible or unversioned Step5 TP shards: expected " + f"tp_layout_version={STEP5_TP_LAYOUT_VERSION!r}, got " + f"{self.tp_layout_version!r}. Older layouts either replicated all " + "KV/provider groups or replicated the shared expert; both are numerically " + "wrong. Regenerate every rank with " + "convert.py (do not resume old shards)." + ) + + def layer_type(self, layer_idx: int) -> str: + if layer_idx < 0: + raise IndexError(f"layer index must be non-negative, got {layer_idx}") + if layer_idx < self.num_hidden_layers: + return self.layer_types[layer_idx] + mtp_idx = layer_idx - self.num_hidden_layers + if mtp_idx >= len(self._mtp_layer_types): + raise IndexError( + f"layer index {layer_idx} is outside the " + f"{self.num_hidden_layers + len(self._mtp_layer_types)} configured layers" + ) + return self._mtp_layer_types[mtp_idx] + + def is_sparse_layer(self, layer_idx: int) -> bool: + """Whether DSA runs on this layer. + + Sparse attention is opt-in per *layer type*, not per layer index, and it is also + gated on the sparse block being present and enabled. Both halves matter: dropping + the type check would put an indexer on every layer, and dropping the enabled + check would ignore a config that asks for the dense path. + """ + sparse = self.sparse_config + if sparse is None or not sparse.enabled: + return False + return self.layer_type(layer_idx) in sparse.apply_to_layer_types + + def is_moe_layer(self, layer_idx: int) -> bool: + return self.use_moe and layer_idx in self.moe_layer_list + + def _consume_quantization_config(self, kwargs: dict[str, Any]) -> None: + """Preserve expert quantization metadata without invoking HF's quantizer. + + Experts use stacked [experts, out, in] weights and kernel.py's block-scaled + GEMM. Exposing quantization_config to HF would incorrectly quantize BF16 + attention, shared-expert and vision projections as well. Retain it + privately for to_dict. Unquantized checkpoints use BF16 experts; FP8 + checkpoints select BF16 expert layers through their exclusions. + """ + quantization = kwargs.pop("quantization_config", None) + if quantization is not None and not isinstance(quantization, dict): + quantization = ( + quantization.to_dict() if hasattr(quantization, "to_dict") else dict(quantization) + ) + if quantization and quantization.get("quant_method") != "fp8": + raise ValueError("Step5 experts support only unquantized BF16 or quant_method='fp8'") + self._quantization_config = quantization + + # Entries name tensors in the *serving* framework's vocabulary, not the checkpoint's. + # The layer index is the only part that survives a rename, so it is the only part read. + excluded = (quantization or {}).get("modules_to_not_convert") or () + layers: set[int] = set() if quantization else set(self.moe_layer_list) + for name in excluded: + # Vision block numbers and MTP blocks are not language MoE layer ids. + name = str(name) + if not (name.endswith(".moe") or any(part in name for part in ( + ".moe.experts.", ".moe.gate_proj", ".moe.up_proj", ".moe.down_proj", + ))): + continue + for part in str(name).split("."): + if part.isdigit(): + layers.add(int(part)) + break + kwargs.pop("bf16_expert_layers", None) + self.bf16_expert_layers = sorted(layers) + + def is_fp8_expert_layer(self, layer_idx: int) -> bool: + """Whether this layer's routed experts are block-scaled fp8 rather than bf16.""" + return self.is_moe_layer(layer_idx) and layer_idx not in self.bf16_expert_layers + + def sliding_window_for(self, layer_idx: int) -> int | None: + """The window this layer attends over, or ``None`` for unbounded. + + A layer is windowed because its *type* says so. ``sliding_window`` in the config is + the size, not the switch -- it stays set even on full-attention layers, where the + DSA metadata builder ignores it entirely. + """ + if self.layer_type(layer_idx) != "sliding_attention": + return None + return self.sliding_window + + def rope_scaling_for(self, layer_idx: int) -> dict[str, Any] | None: + if self.rope_scaling is None: + return None + if self.yarn_only_types and self.layer_type(layer_idx) not in self.yarn_only_types: + return None + return self.rope_scaling + + def to_dict(self) -> dict[str, Any]: + output = super().to_dict() + mtp_layer_types = output.pop("_mtp_layer_types", []) + output["layer_types"] = list(self.layer_types) + list(mtp_layer_types) + if isinstance(output.get("sparse_config"), Step5SparseConfig): + output["sparse_config"] = output["sparse_config"].to_dict() + # Re-attached on the way out only. See :meth:`_consume_quantization_config` for why it + # must not be an attribute while the model is being built. + output.pop("_quantization_config", None) + if self._quantization_config is not None: + output["quantization_config"] = self._quantization_config + return output + + +class Step5VLConfig(PretrainedConfig): + """Released step5_vl envelope around the existing Step-5 text config.""" + + model_type = "step5_vl" + sub_configs = {"text_config": Step5Config, "vision_config": Step5VisionConfig} + + def __init__( + self, + text_config=None, + vision_config=None, + image_token_id=128001, + image_token_len=169, + patch_token_len=81, + projector_bias=False, + vision_select_layer=-1, + tp_size=None, + tp_layout_version=None, + expert_parallel_size=None, + **kwargs, + ): + quantization = kwargs.pop("quantization_config", None) + text = text_config.to_dict() if isinstance(text_config, Step5Config) else dict(text_config or {}) + if quantization is not None: + text["quantization_config"] = quantization + for key, value in (("tp_size", tp_size), ("tp_layout_version", tp_layout_version), + ("expert_parallel_size", expert_parallel_size)): + if value is not None: + text[key] = value + self.text_config = Step5Config(**text) + self.vision_config = (vision_config if isinstance(vision_config, Step5VisionConfig) + else Step5VisionConfig(**(vision_config or {}))) + self.image_token_id = image_token_id + self.image_token_len = image_token_len + self.patch_token_len = patch_token_len + self.projector_bias = projector_bias + self.vision_select_layer = vision_select_layer + if vision_select_layer != -1: + raise ValueError("Only vision_select_layer=-1 (all encoder blocks) is supported") + self.tp_size = self.text_config.tp_size + self.tp_layout_version = self.text_config.tp_layout_version + self.expert_parallel_size = getattr(self.text_config, "expert_parallel_size", self.tp_size) + _normalize_config_identity( + kwargs, model_type=self.model_type, architecture="Step5ForConditionalGeneration", + ) + kwargs.setdefault("bos_token_id", self.text_config.bos_token_id) + kwargs.setdefault("eos_token_id", self.text_config.eos_token_id) + kwargs.setdefault("tie_word_embeddings", False) + super().__init__(**kwargs) + + @property + def tp(self): + return self.text_config.tp + + def to_dict(self): + output = super().to_dict() + output["text_config"] = self.text_config.to_dict() + quantization = output["text_config"].pop("quantization_config", None) + if quantization is not None: + output["quantization_config"] = quantization + output["vision_config"] = self.vision_config.to_dict() + return output + + +def config_from_dict(payload): + """Select text or multimodal configuration without dropping the VL envelope.""" + if payload.get("model_type") in ("step5_vl", "step4_vl") or "vision_config" in payload: + return Step5VLConfig(**payload) + if "text_config" in payload: + raise ValueError("Nested text_config requires a step5_vl/vision_config envelope") + return Step5Config(**payload) + + +def model_class_for_config(config): + return Step5ForConditionalGeneration if isinstance(config, Step5VLConfig) else Step5ForCausalLM + + +# Tensor-parallel layers and collectives. Shards use model-r{rank}.safetensors. + + +_TP_GROUP: dist.ProcessGroup | None = None +_TP_GATHER_MAX_SCRATCH_BYTES = 64 * 1024 * 1024 +_TP_GATHER_CACHE_MAX_BYTES = 16 * 1024 * 1024 +_TP_GATHER_SCRATCH: dict[ + tuple[torch.device, torch.dtype, int, int, int, int, object], + tuple[torch.Tensor, torch.Tensor], +] = {} +_TP_GATHER_SCRATCH_BYTES = 0 +_TP_GATHER_SCRATCH_LOCK = threading.Lock() + + +def set_tp_group(group: dist.ProcessGroup | None) -> None: + """Override the collective target. Defaults to the world group when None.""" + global _TP_GROUP + _TP_GROUP = group + + +def tp_group() -> dist.ProcessGroup | None: + if not (dist.is_available() and dist.is_initialized()): + return None + return _TP_GROUP + + +def tp_rank() -> int: + if not (dist.is_available() and dist.is_initialized()): + return 0 + return dist.get_rank(_TP_GROUP) + + +def tp_size() -> int: + if not (dist.is_available() and dist.is_initialized()): + return 1 + return dist.get_world_size(_TP_GROUP) + + +def _local_size(total: int, tp: int) -> int: + if tp < 1: + raise ValueError(f"tp_size must be positive, got {tp}") + if total % tp: + raise ValueError(f"{total} not divisible by tp_size={tp}") + return total // tp + + +def _all_reduce(x: torch.Tensor) -> torch.Tensor: + """Sum TP partials in a fixed rank order with FP32 accumulation. + + NCCL does not promise a fixed reduction tree. On the supported H200 + runtimes, that made otherwise identical batch rows differ after a BF16 + all-reduce. Greedy decoding can amplify that last-bit difference into a + different token sequence. Small decode tensors use one bit-preserving + all-gather followed by local rank-ordered addition. Large prefill tensors + use rank-ordered broadcasts to keep scratch memory bounded. + """ + if tp_size() <= 1: + return x + if x.dtype not in (torch.bfloat16, torch.float32): + raise TypeError( + "fixed-order TP reduction supports only bfloat16 and float32, " + f"got {x.dtype}" + ) + if x.numel() == 0: + return x + group = _TP_GROUP + if group is None: + group = dist.group.WORLD + rank = dist.get_rank(group) + world_size = dist.get_world_size(group) + local = x.contiguous() + gather_scratch_bytes = local.numel() * ( + world_size * local.element_size() + 4 + ) + gather = getattr(dist, "all_gather_single", None) + if not callable(gather): + gather = getattr(dist, "all_gather_into_tensor", None) + if ( + gather_scratch_bytes <= _TP_GATHER_MAX_SCRATCH_BYTES + and callable(gather) + ): + global _TP_GATHER_SCRATCH_BYTES + stream_id = ( + int(torch.cuda.current_stream(local.device).cuda_stream) + if local.is_cuda + else 0 + ) + key = ( + local.device, + local.dtype, + world_size, + local.numel(), + threading.get_ident(), + stream_id, + group, + ) + with _TP_GATHER_SCRATCH_LOCK: + scratch = _TP_GATHER_SCRATCH.get(key) + if scratch is None: + gathered = torch.empty( + (world_size, local.numel()), + device=local.device, + dtype=local.dtype, + ) + reduced = torch.empty( + (local.numel(),), + device=local.device, + dtype=torch.float32, + ) + scratch = (gathered, reduced) + if ( + _TP_GATHER_SCRATCH_BYTES + gather_scratch_bytes + <= _TP_GATHER_CACHE_MAX_BYTES + ): + _TP_GATHER_SCRATCH[key] = scratch + _TP_GATHER_SCRATCH_BYTES += gather_scratch_bytes + gathered, reduced = scratch + gather( + gathered.view(-1), + local.view(-1), + group=group, + ) + reduced.copy_(gathered[0]) + for source in range(1, world_size): + reduced.add_(gathered[source]) + local.copy_(reduced.view_as(local)) + if local is not x: + x.copy_(local) + return x + + partial = torch.empty_like(local) + reduced = torch.empty_like(local, dtype=torch.float32) + for source in range(world_size): + source_tensor = local if rank == source else partial + dist.broadcast( + source_tensor, + src=dist.get_global_rank(group, source), + group=group, + ) + if source == 0: + reduced.copy_(source_tensor) + else: + reduced.add_(source_tensor) + local.copy_(reduced) + if local is not x: + x.copy_(local) + return x + + +class ColumnParallelLinear(nn.Module): + """Output-sliced linear, no internal reduce. + + Weight shape per rank: ``[out/tp, in]``. Forward computes ``x @ W.T`` locally and + returns the partial output of shape ``[..., out/tp]``. The caller is responsible for + ensuring the upstream produced a full input and the downstream handles partial output. + """ + + def __init__( + self, + in_features: int, + out_features: int, + *, + tp_size: int = 1, + bias: bool = True, + device=None, + dtype=None, + gather_output: bool = False, + ) -> None: + super().__init__() + self.tp_size = tp_size + out_local = _local_size(out_features, tp_size) + factory = {"device": device, "dtype": dtype} + self.weight = nn.Parameter( + torch.empty((out_local, in_features), **factory), requires_grad=False + ) + self.bias = ( + nn.Parameter(torch.zeros(out_local, **factory), requires_grad=False) + if bias + else None + ) + self.gather_output = gather_output + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out = torch.nn.functional.linear(x, self.weight, self.bias) + if self.gather_output and tp_size() > 1: + out = _all_gather(out) + return out + + +def _all_gather(x: torch.Tensor) -> torch.Tensor: + group = _TP_GROUP + if group is None: + group = dist.group.WORLD + n = dist.get_world_size(group) + out = [torch.empty_like(x) for _ in range(n)] + torch.distributed.all_gather(out, x, group=group) + return torch.cat(out, dim=-1) + + +def vocab_parallel_argmax( + local_logits: torch.Tensor, + *, + expected_tp_size: int, +) -> torch.Tensor: + """Return greedy token ids after reconstructing the full TP vocabulary. + + :class:`VocabParallelLinear` intentionally leaves the vocabulary sharded: rank + ``r`` returns the contiguous window ``[r * vocab_local, (r + 1) * + vocab_local)``. Taking ``argmax`` on that tensor directly produces a *local + column index*, not a token id, and (worse) lets every rank feed a different token + into the next decode collective. + + All-gather vocabulary windows in rank order, concatenate them, then take the + argmax. The + returned tensor has shape ``local_logits.shape[:-1]`` and is identical on every + TP rank. For TP=1 it is a plain local argmax and does not require a process + group. + + ``expected_tp_size`` is explicit rather than inferred so a TP checkpoint cannot + silently run generation with an uninitialised or wrong-sized process group. + """ + if local_logits.ndim == 0: + raise ValueError("vocab logits must have at least one dimension") + if expected_tp_size < 1: + raise ValueError(f"expected_tp_size must be positive, got {expected_tp_size}") + if expected_tp_size == 1: + return local_logits.argmax(dim=-1) + + if not (dist.is_available() and dist.is_initialized()): + raise RuntimeError( + f"TP={expected_tp_size} vocab argmax requires an initialized " + "torch.distributed process group" + ) + group = _TP_GROUP if _TP_GROUP is not None else dist.group.WORLD + actual_tp_size = dist.get_world_size(group) + if actual_tp_size != expected_tp_size: + raise RuntimeError( + "vocab argmax process-group size does not match the checkpoint: " + f"expected TP={expected_tp_size}, got {actual_tp_size}" + ) + + full_logits = _all_gather(local_logits.contiguous()) + return full_logits.argmax(dim=-1) + + +class RowParallelLinear(nn.Module): + """Input-sliced linear with an optional internal all-reduce. + + Weight shape per rank: ``[out, in/tp]``. The input is the partial output of an + upstream :class:`ColumnParallelLinear`; each rank multiplies its slice and returns + the full output (all-reduced across ranks) by default. With + ``reduce_results=False`` it returns the local GEMM result instead. + """ + + def __init__( + self, + in_features: int, + out_features: int, + *, + tp_size: int = 1, + bias: bool = True, + device=None, + dtype=None, + reduce_results: bool = True, + ) -> None: + super().__init__() + self.tp_size = tp_size + self.reduce_results = reduce_results + in_local = _local_size(in_features, tp_size) + factory = {"device": device, "dtype": dtype} + self.weight = nn.Parameter( + torch.empty((out_features, in_local), **factory), requires_grad=False + ) + self.bias = ( + nn.Parameter(torch.zeros(out_features, **factory), requires_grad=False) + if bias + else None + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out = torch.nn.functional.linear(x, self.weight, self.bias) + if self.reduce_results and self.tp_size > 1: + return _all_reduce(out) + return out + + +class VocabParallelLinear(nn.Module): + """Output-vocab-sliced linear, no internal reduce (no all-gather). + + Used for ``lm_head`` in forward-only numeric check. Each rank produces logits for + its vocab slice; a serving path would all-gather, but the gate only needs each + rank's slice to be finite. + + The conversion step ships the **full** ``[vocab, in]`` weight on every rank. The + parameter therefore starts full-sized so both the legacy Transformers loader and the + direct tensor-injection loader in Transformers 5 see an exact checkpoint shape. + :meth:`shard_` narrows it to this rank's contiguous vocabulary window immediately + after loading and before the model is moved to the GPU. + """ + + def __init__( + self, + in_features: int, + out_features: int, + *, + tp_size: int = 1, + bias: bool = True, + device=None, + dtype=None, + ) -> None: + super().__init__() + self.tp_size = tp_size + self.out_features = out_features + self._out_features_per_rank = _local_size(out_features, tp_size) + self._vocab_sharded = False + factory = {"device": device, "dtype": dtype} + self.weight = nn.Parameter( + torch.empty((out_features, in_features), **factory), requires_grad=False + ) + self.bias = ( + nn.Parameter(torch.zeros(out_features, **factory), requires_grad=False) + if bias + else None + ) + + def shard_(self, rank: int | None = None) -> "VocabParallelLinear": + """Keep only this rank's vocabulary window after checkpoint loading.""" + if self._vocab_sharded: + return self + local = self._out_features_per_rank + rank = tp_rank() if rank is None else rank + if rank < 0 or rank >= self.tp_size: + raise ValueError(f"rank {rank} is outside TP range [0, {self.tp_size})") + start = rank * local + self.weight = nn.Parameter( + self.weight.narrow(0, start, local).contiguous(), + requires_grad=False, + ) + if self.bias is not None: + self.bias = nn.Parameter( + self.bias.narrow(0, start, local).contiguous(), + requires_grad=False, + ) + self._vocab_sharded = True + return self + + def _load_from_state_dict( + self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) -> None: + # Keep reloading an already-sharded module compatible with the legacy + # state-dict path. Normal from_pretrained loading starts full-sized and + # is narrowed exactly once by shard_ after the checkpoint is loaded. + rank = tp_rank() if self.tp_size > 1 else 0 + for param_name in ("weight", "bias"): + key = prefix + param_name + if key not in state_dict: + continue + loaded = state_dict[key] + is_bias = param_name == "bias" + param = self.bias if is_bias else self.weight + if param is None: + continue + if loaded.shape[0] > param.shape[0]: + state_dict[key] = loaded.narrow(0, rank * param.shape[0], param.shape[0]) + super()._load_from_state_dict( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.linear(x, self.weight, self.bias) + + +# Packed attention and caches: sliding layers use ring buffers; DSA layers +# retain full K/V and completed-region summaries. Geometry and RoPE are per layer. + + +# Queries per chunk in the sliding-window prefill. Bounds the attention mask to +# ``chunk x (chunk + window)`` so a 256k prefill never materialises a square mask. +SLIDING_PREFILL_CHUNK = 2048 + + +class Step5RMSNorm(nn.Module): + """Zero-centered RMSNorm: ``x / rms(x) * (1 + w)``, reduced in fp32. + + The ``1 +`` is not a stylistic variant -- the checkpoint's weights are centred on zero, + so reading them as a conventional RMSNorm scale multiplies every activation by roughly + zero. It fails loudly at the first layer, which is the good case; what makes it worth a + comment is that the same weights are a *valid* conventional scale, so nothing in the + tensor shapes or dtypes objects. + """ + + def __init__(self, hidden_size: int, eps: float = 1e-5) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=torch.float32)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + dtype = hidden_states.dtype + hidden_states = hidden_states.float() + variance = hidden_states.pow(2).mean(-1, keepdim=True) + normed = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return (normed * (self.weight.float() + 1.0)).to(dtype) + + +class DenseLayerCache: + """Ring buffer of the last ``window`` tokens for one sliding-attention layer.""" + + def __init__( + self, + *, + batch: int, + window: int, + num_kv_groups: int, + head_dim: int, + device: torch.device, + dtype: torch.dtype, + ) -> None: + self.window = window + self.key = torch.zeros( + (batch, window, num_kv_groups, head_dim), device=device, dtype=dtype + ) + self.value = torch.zeros_like(self.key) + + def write(self, request: int, start: int, key: torch.Tensor, value: torch.Tensor) -> None: + """Store one request's tokens, keeping only the most recent ``window`` of them.""" + count = key.shape[0] + if count > self.window: + key, value = key[-self.window :], value[-self.window :] + start += count - self.window + count = self.window + slots = (torch.arange(count, device=key.device) + start) % self.window + self.key[request, slots] = key + self.value[request, slots] = value + + def valid_mask(self, lengths: torch.Tensor) -> torch.Tensor: + """``[batch, window]`` bool: which slots hold a live token. + + ``min(length, window)`` covers both regimes with one expression -- below the window + only the first ``length`` slots have been written, and at or above it every slot has. + """ + slots = torch.arange(self.window, device=lengths.device) + return slots.view(1, -1) < lengths.clamp_max(self.window).view(-1, 1) + + +class Step5Cache: + """Per-layer KV state for a batch of sequences, plus the lengths both paths read. + + One object rather than a list of independent caches because ``lengths`` is shared: every + layer advances together, and duplicating the counter per layer is how a decode step ends + up scoring against a context one token stale on some layers and not others. + """ + + def __init__(self, config, *, batch: int, max_tokens: int, device: torch.device, dtype: torch.dtype) -> None: + sparse = config.sparse_config + geometry = DSAGeometry( + proxy_dim=sparse.proxy_dim, + topk=sparse.topk, + region_size=sparse.region_block_size, + ) + self.config = config + self.batch = batch + self.lengths = torch.zeros((batch,), device=device, dtype=torch.int32) + self.layers: list[DSALayerCache | DenseLayerCache] = [] + # Q heads are column-sharded, so each rank only caches the KV groups intersecting + # its contiguous Q-head slice. For production TP=8 this is one group: ranks 0/1 + # cache KV0, 2/3 cache KV1, etc. The pre-sharder duplicates that one group across + # the two ranks that consume it. + local_kv_groups = config.num_kv_groups_per_rank + for layer_idx in range(config.num_hidden_layers): + if config.is_sparse_layer(layer_idx): + self.layers.append( + DSALayerCache( + batch=batch, + max_tokens=max_tokens, + num_kv_groups=local_kv_groups, + head_dim=config.head_dim, + geometry=geometry, + device=device, + dtype=dtype, + ) + ) + else: + # A full-attention layer only reaches here when the sparse path is disabled, + # and then its history is genuinely unbounded -- reusing ``sliding_window`` + # would silently truncate it to 512 and look like a quality regression rather + # than a configuration error. + window = config.sliding_window_for(layer_idx) or max_tokens + self.layers.append( + DenseLayerCache( + batch=batch, + window=window, + num_kv_groups=local_kv_groups, + head_dim=config.head_dim, + device=device, + dtype=dtype, + ) + ) + + def advance(self, seq_lens: torch.Tensor) -> torch.Tensor: + """Record ``seq_lens`` new tokens per request and return the lengths from before.""" + past = self.lengths.clone() + self.lengths = self.lengths + seq_lens.to(self.lengths.dtype) + for layer in self.layers: + if isinstance(layer, DSALayerCache): + layer.lengths = self.lengths + return past + + def compact(self, active_indices: Sequence[int] | torch.Tensor) -> None: + """Drop finished requests while preserving the remaining cache rows. + + Generation normally keeps one cache row per request for the whole decode. Once a + request emits EOS, feeding that row through subsequent decode steps is both wasted + work and a source of batch-shape-dependent numerical variation in fused kernels. + Compacting is done only at an EOS boundary, so the active rows retain their original + context lengths and logical order. Only live tokens, completed summaries, and the + pending region tail are moved; the unused capacity can be very large at long context. + """ + if isinstance(active_indices, torch.Tensor): + indices = active_indices.to(device=self.lengths.device, dtype=torch.long) + else: + indices = torch.tensor( + list(active_indices), device=self.lengths.device, dtype=torch.long + ) + if indices.ndim != 1: + raise ValueError( + f"active_indices must be one-dimensional, got {tuple(indices.shape)}" + ) + if indices.numel() == 0: + raise ValueError("cannot compact a cache to zero requests") + if bool((indices < 0).any().item()) or bool( + (indices >= self.batch).any().item() + ): + raise IndexError( + f"active cache indices must be in [0, {self.batch}), " + f"got {indices.tolist()}" + ) + if indices.numel() > 1 and bool((indices[1:] <= indices[:-1]).any().item()): + raise ValueError( + "active cache indices must be strictly increasing so rows can " + f"be compacted in place, got {indices.tolist()}" + ) + if indices.numel() == self.batch and torch.equal( + indices, torch.arange(self.batch, device=indices.device) + ): + return + + new_batch = int(indices.numel()) + index_values = [int(index) for index in indices.tolist()] + new_lengths = self.lengths.index_select(0, indices).contiguous() + length_values = [int(length) for length in new_lengths.tolist()] + if any(length < 0 for length in length_values): + raise ValueError( + f"cache lengths must be non-negative, got {length_values}" + ) + for layer in self.layers: + if isinstance(layer, DSALayerCache) and any( + length > layer.max_tokens for length in length_values + ): + raise ValueError( + f"cache lengths {length_values} exceed DSA capacity " + f"{layer.max_tokens}" + ) + + for layer in self.layers: + if isinstance(layer, DSALayerCache): + region_size = layer.geometry.region_size + pages = layer.block_table.shape[1] + expected_block_table = torch.arange( + self.batch * pages, + device=layer.block_table.device, + dtype=layer.block_table.dtype, + ).view(self.batch, pages) + if not torch.equal(layer.block_table, expected_block_table): + raise ValueError( + "cache compaction requires the current contiguous " + "identity DSA block table" + ) + for target, (source, length) in enumerate( + zip(index_values, length_values) + ): + if target == source: + continue + + target_token = target * layer.max_tokens + source_token = source * layer.max_tokens + if length: + layer.key[ + target_token : target_token + length + ].copy_(layer.key[source_token : source_token + length]) + layer.value[ + target_token : target_token + length + ].copy_(layer.value[source_token : source_token + length]) + + complete_regions = length // region_size + if complete_regions: + target_region = target * layer.max_regions + source_region = source * layer.max_regions + layer.summary[ + target_region : target_region + complete_regions + ].copy_( + layer.summary[ + source_region : source_region + complete_regions + ] + ) + + pending = length % region_size + if pending: + layer.pending_key[target, :pending].copy_( + layer.pending_key[source, :pending] + ) + layer.pending_z[target, :pending].copy_( + layer.pending_z[source, :pending] + ) + + layer.key = layer.key.narrow( + 0, 0, new_batch * layer.max_tokens + ) + layer.value = layer.value.narrow( + 0, 0, new_batch * layer.max_tokens + ) + layer.summary = layer.summary.narrow( + 0, 0, new_batch * layer.max_regions + ) + layer.pending_key = layer.pending_key.narrow(0, 0, new_batch) + layer.pending_z = layer.pending_z.narrow(0, 0, new_batch) + layer.block_table = layer.block_table.narrow(0, 0, new_batch) + torch.arange( + new_batch * pages, + device=layer.block_table.device, + dtype=layer.block_table.dtype, + out=layer.block_table.view(-1), + ) + layer.batch = new_batch + else: + for target, (source, length) in enumerate( + zip(index_values, length_values) + ): + if target == source: + continue + live_slots = min(length, layer.window) + if live_slots: + layer.key[target, :live_slots].copy_( + layer.key[source, :live_slots] + ) + layer.value[target, :live_slots].copy_( + layer.value[source, :live_slots] + ) + layer.key = layer.key.narrow(0, 0, new_batch) + layer.value = layer.value.narrow(0, 0, new_batch) + + self.lengths = new_lengths + self.batch = new_batch + for layer in self.layers: + if isinstance(layer, DSALayerCache): + layer.lengths = self.lengths + + +class Step5Attention(nn.Module): + """One attention layer. Sparse or windowed according to the layer's type. + + The indexer's parameters live directly on this module (``sparse_indexer_q`` and friends) + rather than inside a submodule, because that is where the checkpoint puts them. The + logic that drives them is in :class:`kernel.Step5SparseIndexer`, which + borrows them -- see its docstring. + """ + + def __init__(self, config, layer_idx: int, device=None, dtype=None) -> None: + super().__init__() + factory = {"device": device, "dtype": dtype} + self.config = config + self.layer_idx = layer_idx + self.num_heads = config.num_attention_heads_per_rank + self.num_kv_groups = config.num_kv_groups_per_rank + self.head_dim = config.head_dim + if self.num_heads % self.num_kv_groups: + raise ValueError( + f"{self.num_heads} local Q heads do not divide into " + f"{self.num_kv_groups} local KV groups" + ) + self.heads_per_group = self.num_heads // self.num_kv_groups + self.scaling = self.head_dim**-0.5 + self.sliding_window = config.sliding_window_for(layer_idx) + self.is_sparse = config.is_sparse_layer(layer_idx) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_groups * self.head_dim + + # Keep the three checkpoint-facing names while loading, then pack them into the + # single [Q_local, K_local, V_local] matrix for the fused projection. + # Separate GEMM shapes are not numerically interchangeable in BF16: on H200 the + # Q result happens to agree, while sparse K/V elements cross rounding boundaries. + # ``pack_qkv_`` rebinds these parameters to views of the non-persistent packed + # buffer, so matching the deployed fused GEMM costs no steady-state duplicate. + self.q_proj = ColumnParallelLinear( + config.hidden_size, config.num_attention_heads * config.head_dim, + tp_size=config.tp, bias=False, **factory, + ) + self.k_proj = nn.Linear( + config.hidden_size, self.kv_size, bias=False, **factory + ) + self.v_proj = nn.Linear( + config.hidden_size, self.kv_size, bias=False, **factory + ) + self.register_buffer( + "_qkv_weight", + torch.empty(0, **factory), + persistent=False, + ) + self.o_proj = RowParallelLinear( + config.num_attention_heads * config.head_dim, config.hidden_size, + tp_size=config.tp, bias=False, **factory, + ) + self.q_norm = Step5RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Step5RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + self.use_head_wise_attn_gate = config.use_head_wise_attn_gate + if self.use_head_wise_attn_gate: + self.g_proj = ColumnParallelLinear( + config.hidden_size, config.num_attention_heads, + tp_size=config.tp, bias=False, **factory, + ) + + rotary_span = int(round(self.head_dim * config.partial_rotary_factors[layer_idx])) + self.rotary_span = rotary_span - rotary_span % 2 + self.register_buffer("rope_cos", torch.empty(0), persistent=False) + self.register_buffer("rope_sin", torch.empty(0), persistent=False) + + if self.is_sparse: + self._build_indexer(config, layer_idx, factory) + else: + self.indexer = None + + @property + def qkv_is_packed(self) -> bool: + """Whether the runtime fused QKV matrix has been materialised.""" + return self._qkv_weight.numel() != 0 + + def _invalidate_packed_qkv(self) -> None: + # Preserve device/dtype (including meta during low-memory construction) without + # retaining the packed storage. The three named parameters keep their own + # references until a subsequent load or repack replaces them. + self._qkv_weight = self.q_proj.weight.new_empty(0) + + def _validate_packed_qkv(self) -> None: + expected_shape = ( + self.q_size + 2 * self.kv_size, + self.config.hidden_size, + ) + if tuple(self._qkv_weight.shape) != expected_shape: + raise RuntimeError( + f"layer {self.layer_idx}: packed QKV shape " + f"{tuple(self._qkv_weight.shape)} != {expected_shape}" + ) + + parameters = (self.q_proj.weight, self.k_proj.weight, self.v_proj.weight) + expected_rows = (self.q_size, self.kv_size, self.kv_size) + expected_offset = 0 + packed_storage = self._qkv_weight.untyped_storage().data_ptr() + for name, weight, rows in zip(("q", "k", "v"), parameters, expected_rows): + expected = (rows, self.config.hidden_size) + if tuple(weight.shape) != expected: + raise RuntimeError( + f"layer {self.layer_idx}: {name}_proj shape " + f"{tuple(weight.shape)} != {expected}" + ) + if weight.device != self._qkv_weight.device or weight.dtype != self._qkv_weight.dtype: + raise RuntimeError( + f"layer {self.layer_idx}: {name}_proj device/dtype does not " + "match packed QKV" + ) + if weight.untyped_storage().data_ptr() != packed_storage: + raise RuntimeError( + f"layer {self.layer_idx}: {name}_proj no longer shares packed " + "QKV storage; repack after moving or loading the model" + ) + if weight.storage_offset() != expected_offset: + raise RuntimeError( + f"layer {self.layer_idx}: {name}_proj storage offset " + f"{weight.storage_offset()} != {expected_offset}" + ) + expected_offset += weight.numel() + + def pack_qkv_(self) -> Step5Attention: + """Prepack Q/K/V weights for one fused projection, without steady-state copies. + + The checkpoint intentionally retains the separate ``q_proj``/``k_proj``/ + ``v_proj`` keys. Call this after loading and placing the model on its execution + device. The first forward also calls it lazily for compatibility with direct + model construction. + """ + if self.qkv_is_packed: + self._validate_packed_qkv() + return self + + weights = (self.q_proj.weight, self.k_proj.weight, self.v_proj.weight) + if any(weight.is_meta for weight in weights): + raise RuntimeError( + f"layer {self.layer_idx}: cannot pack QKV weights on the meta device; " + "load the checkpoint first" + ) + devices = {weight.device for weight in weights} + dtypes = {weight.dtype for weight in weights} + if len(devices) != 1 or len(dtypes) != 1: + raise RuntimeError( + f"layer {self.layer_idx}: Q/K/V device or dtype mismatch: " + f"devices={devices}, dtypes={dtypes}" + ) + + # Packing can run under inference_mode. Temporarily disable it so these + # parameters remain ordinary tensors that a later model.to(...) can transform. + with torch.inference_mode(False), torch.no_grad(): + packed = torch.cat(weights, dim=0).contiguous() + self._qkv_weight = packed + + offset = 0 + for projection, rows in ( + (self.q_proj, self.q_size), + (self.k_proj, self.kv_size), + (self.v_proj, self.kv_size), + ): + view = packed.narrow(0, offset, rows) + projection.weight = nn.Parameter(view, requires_grad=False) + offset += rows + + self._validate_packed_qkv() + return self + + def _apply(self, fn, recurse: bool = True): + """Keep the packed buffer and checkpoint-facing parameter views coalesced.""" + was_packed = self.qkv_is_packed + result = super()._apply(fn, recurse=recurse) + if was_packed: + # Module._apply transforms the buffer and each registered parameter + # independently. Repack immediately to restore one shared storage. + self._invalidate_packed_qkv() + self.pack_qkv_() + return result + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) -> None: + # Loading after a previous inference run must not leave a stale packed buffer. + # Child q/k/v modules are populated after this module-level callback returns. + if self.qkv_is_packed: + self._invalidate_packed_qkv() + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def _build_indexer(self, config, layer_idx: int, factory: dict) -> None: + sparse = config.sparse_config + proxy_dim = sparse.proxy_dim + # DSA provider groups follow the same partition/duplication as KV groups. The + # production checkpoint has 16 indexer heads = 4 provider groups x 4 heads; TP=8 + # therefore constructs only four Q/W heads on each rank. K and Z stay single-head + # MQA projections and are replicated by the pre-sharder. + heads = config.sparse_indexer_num_heads_per_rank + self.sparse_indexer_q = nn.Linear(config.hidden_size, heads * proxy_dim, bias=False, **factory) + self.sparse_indexer_k = nn.Linear(config.hidden_size, proxy_dim, bias=False, **factory) + self.sparse_indexer_z = nn.Linear(config.hidden_size, proxy_dim, bias=False, **factory) + self.sparse_indexer_w = nn.Linear(config.hidden_size, heads, bias=False, dtype=torch.float32, device=factory["device"]) + self.sparse_indexer_q_norm = Step5RMSNorm(proxy_dim, eps=config.rms_norm_eps) + self.sparse_indexer_k_norm = nn.LayerNorm(proxy_dim, eps=config.rms_norm_eps, dtype=torch.float32, device=factory["device"]) + # Present in the checkpoint at its initialisation value on every head of every + # layer, and unused by the deployed kernels: the indexer scores with a weighted + # ReLU, which has no softmax for a scalable-softmax scale to act on. Registered so + # loading is clean and the tensor is not silently dropped. Shape is the full head + # count (the convert step ships ssmax_s as a replicate, see ``_SHARD_TABLE``); the + # full shape is retained for checkpoint compatibility. + self.register_buffer( + "ssmax_s", torch.zeros(config.num_attention_heads, dtype=torch.float32, device=factory["device"]) + ) + self.register_buffer("sparse_indexer_rope_cos", torch.empty(0), persistent=False) + self.register_buffer("sparse_indexer_rope_sin", torch.empty(0), persistent=False) + self.indexer = Step5SparseIndexer( + self, + geometry=DSAGeometry( + proxy_dim=proxy_dim, topk=sparse.topk, region_size=sparse.region_block_size + ), + num_kv_groups=self.num_kv_groups, + ) + + def build_rope( + self, + *, + max_position: int, + device: torch.device, + shared_cache: dict[tuple[object, ...], tuple[torch.Tensor, torch.Tensor]] + | None = None, + ) -> None: + """Materialise this layer's cos/sin tables, optionally sharing equal tables. + + Step5 has 92 layers but only a few distinct RoPE geometries. At the production + 524288-token capacity, independently materialising every table costs many GiB per + rank. ``Step5Model.build_rope`` passes a model-local + cache here to do the same without introducing process-global state. + """ + config = self.config + theta = config.rope_theta[self.layer_idx] + scaling = config.rope_scaling_for(self.layer_idx) + common_key = ( + float(theta), + int(max_position), + str(device), + repr(scaling), + ) + main_key = ("attention", self.rotary_span, torch.bfloat16, *common_key) + if self.rope_cos.numel() == 0: + cached = shared_cache.get(main_key) if shared_cache is not None else None + if cached is None: + cached = build_rope_cache( + rotary_span=self.rotary_span, + theta=theta, + max_position=max_position, + device=device, + # Construct frequencies/trigonometry in FP32, then + # stores the cache in the model's default BF16 dtype. + dtype=torch.bfloat16, + scaling=scaling, + ) + if shared_cache is not None: + shared_cache[main_key] = cached + cos, sin = cached + self.rope_cos, self.rope_sin = cos, sin + elif shared_cache is not None: + shared_cache.setdefault(main_key, (self.rope_cos, self.rope_sin)) + if self.is_sparse and self.sparse_indexer_rope_cos.numel() == 0: + indexer_span = self.config.sparse_config.sparse_indexer_rope_dim + indexer_key = ( + "sparse_indexer", + indexer_span, + torch.bfloat16, + *common_key, + ) + cached = ( + shared_cache.get(indexer_key) + if shared_cache is not None + else None + ) + if cached is None: + cached = build_rope_cache( + rotary_span=indexer_span, + theta=theta, + max_position=max_position, + device=device, + dtype=torch.bfloat16, + scaling=scaling, + ) + if shared_cache is not None: + shared_cache[indexer_key] = cached + cos, sin = cached + self.sparse_indexer_rope_cos, self.sparse_indexer_rope_sin = cos, sin + elif self.is_sparse and shared_cache is not None: + indexer_span = self.config.sparse_config.sparse_indexer_rope_dim + indexer_key = ( + "sparse_indexer", + indexer_span, + torch.bfloat16, + *common_key, + ) + shared_cache.setdefault( + indexer_key, + (self.sparse_indexer_rope_cos, self.sparse_indexer_rope_sin), + ) + + def _apply_rope(self, tensor: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Rotate the leading ``rotary_span`` dims of every head, NEOX-paired. + + Dims past the span are returned untouched, which is what "partial" means here: on + full-attention layers two thirds of each head carries no positional signal at all. + + The deployed fused QK-norm/RoPE path has a precise rounding boundary: norm first + produces BF16 Q/K; the FP32-built, BF16-stored cache and Q/K are then widened for + the RoPE arithmetic, and only the rotated result is rounded back to BF16. Performing + the multiply itself in BF16 drifts from the oracle at every layer. + """ + half = self.rotary_span // 2 + cos = self.rope_cos[positions].unsqueeze(1).float() + sin = self.rope_sin[positions].unsqueeze(1).float() + real = tensor[..., :half].float() + imaginary = tensor[..., half : 2 * half].float() + rotated = torch.cat( + (real * cos - imaginary * sin, real * sin + imaginary * cos), dim=-1 + ).to(tensor.dtype) + return torch.cat((rotated, tensor[..., 2 * half :]), dim=-1) + + def _project_packed_qkv(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Run the deployed single QKV GEMM and retain its packed row layout.""" + if not self.qkv_is_packed: + self.pack_qkv_() + return torch.nn.functional.linear(hidden_states, self._qkv_weight) + + def _split_packed_qkv( + self, qkv: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return TP-local Q/K/V head views without materialising another copy.""" + tokens = qkv.shape[0] + query, key, value = qkv.split( + (self.q_size, self.kv_size, self.kv_size), dim=-1 + ) + query = query.view(tokens, self.num_heads, self.head_dim) + key = key.view(tokens, self.num_kv_groups, self.head_dim) + value = value.view(tokens, self.num_kv_groups, self.head_dim) + return query, key, value + + def _project( + self, hidden_states: torch.Tensor, positions: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + qkv = self._project_packed_qkv(hidden_states) + if qkv.is_cuda and self.head_dim == 192: + # Keep the projection packed all the way into the Triton kernel. Besides + # avoiding three launches and an intermediate concat, this reproduces the + # deployed norm -> BF16 materialisation -> RoPE rounding boundary. + query, key, value = fused_qknorm_rope( + qkv, + self.q_norm.weight, + self.k_norm.weight, + self.rope_cos, + self.rope_sin, + positions.contiguous(), + head_dim=self.head_dim, + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_groups, + rotary_pairs=self.rotary_span // 2, + eps=self.q_norm.variance_epsilon, + norm_weight_bias=1.0, + ) + tokens = qkv.shape[0] + return ( + query.view(tokens, self.num_heads, self.head_dim), + key.view(tokens, self.num_kv_groups, self.head_dim), + value.view(tokens, self.num_kv_groups, self.head_dim), + ) + + # CPU and non-Step5 geometries retain a transparent reference fallback for + # configuration tooling and the small assembled-model tests. + query, key, value = self._split_packed_qkv(qkv) + query = self._apply_rope(self.q_norm(query), positions) + key = self._apply_rope(self.k_norm(key), positions) + return query, key, value + + def _gate_and_project(self, attn_output: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: + """Apply the per-head sigmoid gate, then ``o_proj``. + + The gate is applied *before* the output projection and per head, so it can suppress + an entire head's contribution. Applying it after ``o_proj`` would need a different + weight shape and is not the same function. + """ + if self.use_head_wise_attn_gate: + gate = self.g_proj(hidden_states).sigmoid().unsqueeze(-1) + attn_output = attn_output * gate.to(attn_output.dtype) + return self.o_proj(attn_output.reshape(attn_output.shape[0], -1)) + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + cache: Step5Cache, + seq_lens: torch.Tensor, + past_lens: torch.Tensor, + *, + is_decode: bool, + ) -> torch.Tensor: + """Attend over a packed batch of tokens. + + Args: + hidden_states: ``[total_tokens, hidden]``, requests concatenated. + positions: ``[total_tokens]`` int32, absolute position of each token. + seq_lens: ``[batch]`` int32, tokens contributed by each request now. + past_lens: ``[batch]`` int32, tokens already in the cache. + is_decode: exactly one token per request. A separate flag rather than + ``seq_lens.max() == 1`` because a one-token prefill is a real case and takes + the prefill path. + """ + query, key, value = self._project(hidden_states, positions) + layer_cache = cache.layers[self.layer_idx] + + if not is_decode and bool((past_lens > 0).any()): + # Chunked prefill is unsupported on *both* paths, and each fails differently. + # The dense path would need the ring buffer's live span joined to the chunk's own + # keys, and the join is where the window bound gets lost. The sparse path is + # worse: ``update_summaries_prefill`` numbers regions from zero for whatever it is + # given, so a second chunk would overwrite the first chunk's summaries with its + # own tokens rather than appending -- corruption, not staleness, and silent. + raise NotImplementedError( + "prefill on top of an existing cache is not implemented; prefill once, then " + "decode one token at a time" + ) + + if self.is_sparse: + attn_output = self._sparse_attention( + query, key, value, layer_cache, hidden_states, positions, + seq_lens, past_lens, is_decode=is_decode, + ) + else: + attn_output = self._sliding_attention( + query, key, value, layer_cache, seq_lens, past_lens, is_decode=is_decode + ) + return self._gate_and_project(attn_output, hidden_states) + + def _sparse_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + layer_cache: DSALayerCache, + hidden_states: torch.Tensor, + positions: torch.Tensor, + seq_lens: torch.Tensor, + past_lens: torch.Tensor, + *, + is_decode: bool, + ) -> torch.Tensor: + index_q, index_k, index_z, weights = self.indexer.project(hidden_states, positions) + + offset = 0 + for request in range(layer_cache.batch): + length = int(seq_lens[request]) + layer_cache.write_kv( + request, int(past_lens[request]), key[offset : offset + length], value[offset : offset + length] + ) + offset += length + + if is_decode: + update_summaries_decode(layer_cache, index_k, index_z) + packed, counts = decode_metadata(self.indexer, layer_cache, index_q, weights) + out, _ = sparse_attention_decode( + query, + layer_cache.key, + layer_cache.value, + packed, + counts, + layer_cache.lengths, + num_kv_groups=self.num_kv_groups, + region_size=layer_cache.geometry.region_size, + softmax_scale=self.scaling, + ) + return out + + update_summaries_prefill(layer_cache, index_k, index_z, seq_lens) + packed, counts = prefill_metadata( + self.indexer, layer_cache, index_q, weights, seq_lens, past_lens + ) + out, _ = sparse_attention_prefill( + query, + layer_cache.key, + layer_cache.value, + packed, + counts, + num_kv_groups=self.num_kv_groups, + region_size=layer_cache.geometry.region_size, + softmax_scale=self.scaling, + ) + return out + + def _sliding_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + layer_cache: DenseLayerCache, + seq_lens: torch.Tensor, + past_lens: torch.Tensor, + *, + is_decode: bool, + ) -> torch.Tensor: + window = layer_cache.window + if is_decode: + lengths = past_lens + 1 + batch = layer_cache.key.shape[0] + for request in range(batch): + layer_cache.write( + request, + int(past_lens[request]), + key[request : request + 1], + value[request : request + 1], + ) + out = torch.nn.functional.scaled_dot_product_attention( + query.unsqueeze(2), + layer_cache.key.transpose(1, 2), + layer_cache.value.transpose(1, 2), + attn_mask=layer_cache.valid_mask(lengths).view( + batch, 1, 1, -1 + ), + scale=self.scaling, + enable_gqa=True, + ) + return out.reshape(batch, self.num_heads, self.head_dim) + + out = torch.empty_like(query) + offset = 0 + for request in range(layer_cache.key.shape[0]): + length = int(seq_lens[request]) + past = int(past_lens[request]) + if past: + raise NotImplementedError( + "dense prefill on top of an existing cache is not implemented; chunked " + "prefill would need the ring buffer's live span joined to the chunk's own " + "keys, and the join is where the window bound gets lost" + ) + span = slice(offset, offset + length) + out[span] = self._sliding_prefill_chunked( + query[span], key[span], value[span], window + ) + layer_cache.write(request, 0, key[span], value[span]) + offset += length + return out + + def _sliding_prefill_chunked( + self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, window: int + ) -> torch.Tensor: + """Causal attention limited to the last ``window`` tokens, one query chunk at a time. + + Chunking bounds the mask rather than the compute: a 256k prefill with a 512-token + window would otherwise build a 256k x 256k mask to express a band 512 wide. + """ + length = query.shape[0] + out = torch.empty_like(query) + positions = torch.arange(length, device=query.device) + + for start in range(0, length, SLIDING_PREFILL_CHUNK): + stop = min(start + SLIDING_PREFILL_CHUNK, length) + key_start = max(0, start - window + 1) + q_positions = positions[start:stop].view(-1, 1) + k_positions = positions[key_start:stop].view(1, -1) + distance = q_positions - k_positions + mask = (distance >= 0) & (distance < window) + chunk = torch.nn.functional.scaled_dot_product_attention( + query[start:stop].transpose(0, 1).unsqueeze(0), + key[key_start:stop].transpose(0, 1).unsqueeze(0), + value[key_start:stop].transpose(0, 1).unsqueeze(0), + attn_mask=mask.view(1, 1, stop - start, -1), + scale=self.scaling, + enable_gqa=True, + ) + out[start:stop] = chunk.squeeze(0).transpose(0, 1) + return out + + +# Decoder: FP32 residual accumulation, asymmetric SwiGLU clamps and +# sigmoid routing with a selection-only bias. + + +# The renormalisation guard from the deployed router kernel. Small enough to be a no-op on +# any real top-k sum, and reproduced rather than replaced so the arithmetic matches. +ROUTER_RENORM_EPS = 1e-20 + + +class Step5MLP(nn.Module): + """Clamped SwiGLU. Used for the dense layers and for every MoE layer's shared expert. + + ``tp`` defaults to ``config.tp`` for the dense FFN (ColumnParallel/RowParallel, TP-sliced + on the intermediate dim). The shared expert uses the same TP slicing but sets + ``reduce_results=False``: its local down-projection is added to the local EP-routed + output in FP32 and those two branches are all-reduced exactly once. + """ + + def __init__( + self, + config: Step5Config, + intermediate_size: int, + limit: float, + *, + tp: int | None = None, + reduce_results: bool = True, + device=None, + dtype=None, + ) -> None: + super().__init__() + tp = config.tp if tp is None else tp + factory = {"device": device, "dtype": dtype} + # ``tp`` is the slice the ColumnParallel / RowParallel wrappers apply. Pre-slicing + # ``intermediate_size`` here and then passing ``tp_size=tp`` would double-slice -- + # each rank ends up with ``inter / tp^2`` instead of ``inter / tp`` and the + # checkpoint shape check fails. ``intermediate_size`` is the full out dim; the + # wrapper does the projection onto this rank's output slice. + self.gate_proj = ColumnParallelLinear(config.hidden_size, intermediate_size, tp_size=tp, bias=False, **factory) + self.up_proj = ColumnParallelLinear(config.hidden_size, intermediate_size, tp_size=tp, bias=False, **factory) + self.down_proj = RowParallelLinear( + intermediate_size, + config.hidden_size, + tp_size=tp, + bias=False, + reduce_results=reduce_results, + **factory, + ) + self.limit = limit + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + gate = self.gate_proj(hidden_states) + up = self.up_proj(hidden_states) + return self.down_proj(clamped_swiglu(gate, up, self.limit)) + + +def clamped_swiglu(gate: torch.Tensor, up: torch.Tensor, limit: float) -> torch.Tensor: + """``silu(gate).clamp(max=limit) * up.clamp(-limit, limit)``, computed in fp32. + + The clamps bound the activation's magnitude so the fp8 expert GEMMs downstream keep their + dynamic range. The gate has no lower clamp because ``silu`` already bounds it from below at + about -0.28; adding one would be inert on real values and wrong on the boundary. + + Gate and up projections are widened to FP32 before SiLU, clamping and + multiplication. Only the final product is rounded back to the input dtype. + """ + if gate.is_cuda: + return triton_clamped_swiglu( + gate.contiguous(), up.contiguous(), limit=limit + ) + + # CPU fallback keeps shape-only/unit tests usable. Production inference is + # CUDA/Hopper and therefore always takes the pure-Triton path above. + activated = torch.nn.functional.silu(gate.float()).clamp(max=limit) + bounded = up.float().clamp(-limit, limit) + return (activated * bounded).to(gate.dtype) + + +class _StackedExpertWeight(nn.Module): + """One ``[n_local_experts, out, in]`` weight plus its block scale, addressed one expert at a time. + + Named to match the checkpoint (``moe.gate_proj.weight``), so these are direct children of + the MoE block rather than of an ``experts`` container -- an extra level of nesting here + would rename every expert tensor. + + ``weight_scale_inv`` is declared only for the fp8 layers. The alternative, always + declaring it and leaving it unset for bf16 layers, means the dispatcher's "is the scale + present" test becomes "is the scale meaningful", and there is no value that makes an + unquantised weight's scale meaningful. + + EP, not TP-of-inner-dim. Each rank owns ``n_experts // tp_size`` contiguous experts and + keeps each expert's *full* inner dim. That is the only layout under which the fp8 block + scale's K-block grid (``ceil(in_features / 128)``) stays aligned with the GEMM's K-iters: + the pre-shard step slices the experts dim only, leaving each expert's ``(N, K)`` weight + and ``(N_blocks, K_blocks)`` scale intact. Slicing the inner dim would put a single + K-iter across a weight-block boundary -- forcing a per-tile scale gather the kernel does + not do -- and break the block-scaling assumption. + """ + + def __init__( + self, + n_local_experts: int, + out_features: int, + in_features: int, + *, + quantized: bool, + block: int = 128, + device=None, + ) -> None: + super().__init__() + self.n_local_experts = n_local_experts + self.block = block + weight_dtype = torch.float8_e4m3fn if quantized else torch.bfloat16 + self.weight = nn.Parameter( + torch.empty( + (n_local_experts, out_features, in_features), + device=device, + dtype=weight_dtype, + ), + requires_grad=False, + ) + if quantized: + self.weight_scale_inv = nn.Parameter( + torch.empty( + ( + n_local_experts, + -(-out_features // block), + -(-in_features // block), + ), + device=device, + dtype=torch.float32, + ), + requires_grad=False, + ) + else: + self.weight_scale_inv = None + + def apply_expert(self, expert_idx: int, hidden_states: torch.Tensor) -> torch.Tensor: + # ``hidden_states`` is the routed tokens for this expert on this rank; it is not + # sliced. Each rank owns its experts' full weights, so the GEMM runs over the full + # in dim and the per-rank output is one expert's worth of contribution. There is no + # inner-dim TP slice to sum across ranks -- all-reduce is the responsibility of the + # caller, summing only the routed-expert outputs (one row per token) across ranks. + scale = None if self.weight_scale_inv is None else self.weight_scale_inv[expert_idx] + return linear_fp8_or_bf16(hidden_states, self.weight[expert_idx], scale) + + +class Step5MoE(nn.Module): + """Sigmoid router with a selection-only bias, over 352 stacked experts. + + The shared expert is *not* part of this block. The checkpoint stores it as a sibling + (``layers.N.share_expert``), so the decoder layer owns it and adds its output; owning it + here would nest it under ``moe`` and rename all three of its tensors. + + EP, not TP-of-inner-dim: each rank owns ``moe_num_experts // tp`` contiguous experts + with the full inner dim per expert, so the pre-shard's experts-dim slice is load-bearing + here (see :class:`_StackedExpertWeight`). The router weight and bias are replicated -- + every rank scores all experts and selects the same top-k -- so the routed tokens are the + same everywhere and ``(expert_ids == global_idx)`` matches the same rows on every rank. + """ + + def __init__(self, config: Step5Config, layer_idx: int, device=None, dtype=None) -> None: + super().__init__() + self.top_k = config.moe_top_k + self.scaling_factor = config.moe_router_scaling_factor + self.renormalize = config.norm_expert_weight + self.limit = config.swiglu_limits[layer_idx] if config.swiglu_limits else 7.0 + self.tp = config.tp + # ``tp_rank()`` is 0 when dist is not initialised (the single-GPU path), which is + # correct: with tp=1 every rank owns every expert and start_idx/end_idx span [0, N). + self.rank = tp_rank() + + n_routed = config.moe_num_experts + n_local = n_routed // self.tp + assert n_routed % self.tp == 0, f"{n_routed} not divisible by tp={self.tp}" + self.n_local_experts = n_local + self.experts_start_idx = self.rank * n_local + self.experts_end_idx = self.experts_start_idx + n_local + + # Router weight is replicated: every rank computes all ``moe_num_experts`` logits per token + # and picks the same top-k, so a single rank could in principle make the decision. Keeping + # it full across ranks means the routed tokens are the same everywhere -- which is the + # precondition for expert slicing: the same ``(rows_for_expert_e)`` on every rank. + self.gate = nn.Linear( + config.hidden_size, config.moe_num_experts, bias=False, device=device, dtype=dtype + ) + self.router_bias = nn.Parameter( + torch.zeros(config.moe_num_experts, dtype=torch.float32, device=device), + requires_grad=False, + ) + + quantized = config.is_fp8_expert_layer(layer_idx) + self.gate_proj = _StackedExpertWeight(n_local, config.moe_intermediate_size, config.hidden_size, quantized=quantized, device=device) + self.up_proj = _StackedExpertWeight(n_local, config.moe_intermediate_size, config.hidden_size, quantized=quantized, device=device) + self.down_proj = _StackedExpertWeight(n_local, config.hidden_size, config.moe_intermediate_size, quantized=quantized, device=device) + + def route(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Pick ``top_k`` experts per token and weight them.""" + logits = torch.nn.functional.linear(hidden_states.float(), self.gate.weight.float()) + gate_prob = torch.sigmoid(logits) + scores = gate_prob + self.router_bias.view(1, -1) + top_scores, expert_ids = torch.topk(scores, self.top_k, dim=-1) + + weights = top_scores - self.router_bias[expert_ids] + if self.renormalize: + weights = weights / (weights.sum(dim=-1, keepdim=True) + ROUTER_RENORM_EPS) + return weights * self.scaling_factor, expert_ids + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Return this rank's BF16 routed-expert contribution without a collective. + + Iterates over this rank's local experts (the contiguous index slice + ``[start, end)``) and gathers, for each, the tokens routed to that expert's global + index. The residual stream is full on every rank, so the routed tokens are the same + everywhere. Expert GEMMs are grouped by expert for efficiency, but the deployed + ``ep_gather`` accumulates their BF16 outputs in top-k *slot* order using an FP32 + accumulator and only then rounds the local result to BF16. Keep that order here; + the decoder combines this local result with its local shared-expert result and + performs one FP32 all-reduce. + """ + weights, expert_ids = self.route(hidden_states) + slot_output = torch.zeros( + (*hidden_states.shape[:-1], self.top_k, hidden_states.shape[-1]), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + for local in range(self.n_local_experts): + global_idx = self.experts_start_idx + local + rows, slots = (expert_ids == global_idx).nonzero(as_tuple=True) + if rows.numel() == 0: + continue + tokens = hidden_states[rows] + hidden = clamped_swiglu( + self.gate_proj.apply_expert(local, tokens), + self.up_proj.apply_expert(local, tokens), + self.limit, + ) + contribution = self.down_proj.apply_expert(local, hidden) + slot_output[rows, slots] = contribution.to(hidden_states.dtype) + + if slot_output.is_cuda: + return weighted_topk_gather( + slot_output.contiguous(), weights.contiguous() + ) + + # CPU fallback mirrors the Triton kernel's logical slot order. + out = torch.zeros( + hidden_states.shape, device=hidden_states.device, dtype=torch.float32 + ) + for slot in range(self.top_k): + out = out + ( + slot_output[:, slot].float() * weights[:, slot].unsqueeze(-1) + ) + return out.to(hidden_states.dtype) + + +class Step5DecoderLayer(nn.Module): + def __init__(self, config: Step5Config, layer_idx: int, device=None, dtype=None) -> None: + super().__init__() + self.layer_idx = layer_idx + self.fp32_residual = config.fp32_residual_connection + self.self_attn = Step5Attention(config, layer_idx, device=device, dtype=dtype) + self.input_layernorm = Step5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Step5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + if config.is_moe_layer(layer_idx): + self.moe = Step5MoE(config, layer_idx, device=device, dtype=dtype) + shared_limit = ( + config.swiglu_limits_shared[layer_idx] if config.swiglu_limits_shared else 7.0 + ) + # The shared expert is TP-sharded. Its RowParallel down projection deliberately + # leaves the local partial unreduced so it can be combined with the local + # expert-parallel routed output before one FP32 collective. + self.share_expert = Step5MLP( + config, + config.share_expert_dim, + shared_limit, + tp=config.tp, + reduce_results=False, + device=device, + dtype=dtype, + ) + self.mlp = None + else: + limit = config.swiglu_limits[layer_idx] if config.swiglu_limits else 7.0 + self.mlp = Step5MLP(config, config.intermediate_size, limit, device=device, dtype=dtype) + self.moe = None + self.share_expert = None + + def _feed_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.moe is None: + return self.mlp(hidden_states) + # Both branches are local BF16 values. The NV path widens each independently, + # adds in FP32, then performs a single TP all-reduce. Reducing either branch first + # changes both the function (for a replicated shared expert) and rounding order. + combined = self.moe(hidden_states).float() + combined = combined + self.share_expert(hidden_states).float() + if self.moe.tp > 1: + combined = _all_reduce(combined) + return combined + + def forward( + self, + residual: torch.Tensor, + positions: torch.Tensor, + cache: Step5Cache, + seq_lens: torch.Tensor, + past_lens: torch.Tensor, + *, + is_decode: bool, + ) -> torch.Tensor: + """Take and return the fp32 residual stream, with bf16 into the matmuls.""" + hidden = self.input_layernorm(residual).to(self.self_attn.q_proj.weight.dtype) + attn_out = self.self_attn( + hidden, positions, cache, seq_lens, past_lens, is_decode=is_decode + ) + residual = residual + attn_out.float() + + hidden = self.post_attention_layernorm(residual).to(attn_out.dtype) + return residual + self._feed_forward(hidden).float() + + +class Step5PreTrainedModel(PreTrainedModel): + config_class = Step5Config + base_model_prefix = "model" + supports_gradient_checkpointing = False + _no_split_modules = ["Step5DecoderLayer"] + _supports_sdpa = True + + +class Step5Model(Step5PreTrainedModel): + def __init__(self, config: Step5Config) -> None: + super().__init__(config) + dtype = torch.bfloat16 + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, dtype=dtype) + self.layers = nn.ModuleList( + Step5DecoderLayer(config, layer_idx, dtype=dtype) + for layer_idx in range(config.num_hidden_layers) + ) + self.norm = Step5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def build_rope(self, device: torch.device) -> None: + shared_cache: dict[ + tuple[object, ...], tuple[torch.Tensor, torch.Tensor] + ] = {} + for layer in self.layers: + layer.self_attn.build_rope( + max_position=self.config.max_position_embeddings, + device=device, + shared_cache=shared_cache, + ) + + def pack_qkv_(self) -> Step5Model: + """Prepack every layer's TP-local Q/K/V weights one layer at a time.""" + for layer in self.layers: + layer.self_attn.pack_qkv_() + return self + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + cache: Step5Cache, + seq_lens: torch.Tensor, + past_lens: torch.Tensor, + *, + is_decode: bool, + ) -> torch.Tensor: + residual = hidden_states.float() + for layer in self.layers: + residual = layer( + residual, positions, cache, seq_lens, past_lens, is_decode=is_decode + ) + # Narrow the FP32 residual to the parameter dtype before final RMSNorm. + # This rounding boundary is part of the Step-5 inference arithmetic. + norm_input = ( + residual.to(self.embed_tokens.weight.dtype) + if self.config.fp32_residual_connection + else residual + ) + return self.norm(norm_input) + + +class Step5ForCausalLM(Step5PreTrainedModel): + """step5 with a language-modelling head. + + Input is a packed batch: ``[total_tokens, ...]`` with per-request lengths, rather than a + padded ``[batch, seq]``. That is the layout the DSA kernels want, and converting at the + boundary means padding never reaches the cache -- a padded token written into the region + summaries would be scored and selected like a real one. + """ + + _tied_weights_keys: list[str] = [] + + def __init__(self, config: Step5Config) -> None: + super().__init__(config) + self.model = Step5Model(config) + self.lm_head = VocabParallelLinear( + config.hidden_size, config.vocab_size, tp_size=config.tp, bias=False, dtype=torch.bfloat16 + ) + self.post_init() + + def allocate_cache(self, *, batch: int, max_tokens: int, device: torch.device) -> Step5Cache: + self.model.pack_qkv_() + self.model.build_rope(device) + return Step5Cache( + self.config.get_text_config(), batch=batch, max_tokens=max_tokens, device=device, dtype=torch.bfloat16 + ) + + def pack_qkv_(self) -> Step5ForCausalLM: + """Materialise the inference-only fused QKV layout after loading/device move.""" + self.model.pack_qkv_() + return self + + def forward( + self, + input_ids: torch.Tensor, + cache: Step5Cache, + seq_lens: torch.Tensor, + *, + is_decode: bool = False, + return_all_logits: bool = False, + inputs_embeds: torch.Tensor | None = None, + ) -> CausalLMOutputWithPast: + """Run one prefill or one decode step over a packed batch. + + Args: + input_ids: ``[total_tokens]`` int64, requests concatenated. + seq_lens: ``[batch]`` int32, tokens contributed by each request. + return_all_logits: emit a logit row per token instead of one per request. Off by + default because at a 128896-wide vocabulary the prefill logits are larger + than the activations that produced them. + """ + hidden = self.model.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + if hidden.shape != (input_ids.numel(), self.model.config.hidden_size): + raise ValueError("inputs_embeds must have shape [total_tokens, hidden_size]") + past_lens = cache.advance(seq_lens) + positions = _packed_positions(seq_lens, past_lens) + hidden = self.model( + hidden, positions, cache, seq_lens, past_lens, is_decode=is_decode + ) + + if not return_all_logits: + hidden = hidden[_last_token_indices(seq_lens)] + logits = self.lm_head(hidden.to(self.lm_head.weight.dtype)) + return CausalLMOutputWithPast(logits=logits.float(), past_key_values=cache) + + @torch.inference_mode() + def generate_greedy( + self, + prompt_ids: list[list[int]], + *, + max_new_tokens: int, + device: torch.device, + eos_token_ids: set[int] | None = None, + **multimodal_inputs, + ) -> list[list[int]]: + """Minimal greedy decode, so the model can be exercised without a serving stack. + + Deliberately not ``transformers.generate``: that expects a padded batch and a + ``Cache`` it can slice, and adapting the DSA cache to that interface would mean + pretending its region summaries are sliceable per token. + """ + batch = len(prompt_ids) + lengths = [len(ids) for ids in prompt_ids] + cache = self.allocate_cache( + batch=batch, max_tokens=max(lengths) + max_new_tokens, device=device + ) + eos = eos_token_ids or set( + self.config.eos_token_id + if isinstance(self.config.eos_token_id, (list, tuple)) + else [self.config.eos_token_id] + ) + + flat = torch.tensor([t for ids in prompt_ids for t in ids], device=device, dtype=torch.long) + seq_lens = torch.tensor(lengths, device=device, dtype=torch.int32) + logits = self(flat, cache, seq_lens, **multimodal_inputs).logits + + outputs: list[list[int]] = [[] for _ in range(batch)] + active_indices = list(range(batch)) + for step in range(max_new_tokens): + # ``lm_head`` owns only a contiguous vocab slice on each TP rank. + # A local argmax is merely an index into that slice and makes every + # rank feed a different token into the next collective. Reconstruct + # the full vocabulary so every rank advances the cache with the + # same global token id. + tokens = vocab_parallel_argmax( + logits, expected_tp_size=self.config.tp + ) + next_active_indices: list[int] = [] + next_active_rows: list[int] = [] + for active_row, request in enumerate(active_indices): + token = int(tokens[active_row]) + outputs[request].append(token) + if token not in eos: + next_active_indices.append(request) + next_active_rows.append(active_row) + # ``logits`` already produced this iteration's token. A decode + # forward is needed only when another sampling iteration follows; + # running it after the final requested token mutates the cache and + # launches TP collectives whose result is never consumed. + if not next_active_indices or step + 1 == max_new_tokens: + break + if len(next_active_indices) != len(active_indices): + keep_rows = torch.tensor( + next_active_rows, + device=device, + dtype=torch.long, + ) + cache.compact(keep_rows) + tokens = tokens.index_select(0, keep_rows) + active_indices = next_active_indices + ones = torch.ones( + (len(active_indices),), device=device, dtype=torch.int32 + ) + logits = self(tokens, cache, ones, is_decode=True).logits + return outputs + + +class Step5ForConditionalGeneration(Step5ForCausalLM): + """Vision prefill followed by the same packed DSA text decoder. + + Vision and projector weights are replicated across TP ranks. Their checkpoint + keys remain ``vision_model.*`` and ``vit_large_projector.*``; language keys + remain ``model.*`` and ``lm_head.*``. No Transformers cache adapter is needed. + """ + + config_class = Step5VLConfig + _no_split_modules = ["Step5DecoderLayer", "EncoderVisionBlock"] + + def __init__(self, config: Step5VLConfig): + Step5PreTrainedModel.__init__(self, config) + text = config.text_config + self.model = Step5Model(text) + self.lm_head = VocabParallelLinear( + text.hidden_size, text.vocab_size, tp_size=text.tp, bias=False, dtype=torch.bfloat16 + ) + self.vision_model = Step5VisionEncoder(config.vision_config).to(dtype=torch.bfloat16) + self.vit_large_projector = nn.Linear( + config.vision_config.width * 4, text.hidden_size, + bias=config.projector_bias, dtype=torch.bfloat16, + ) + self.post_init() + + def _encode_images(self, pixels): + if pixels.ndim < 4: + raise ValueError("pixel_values must have at least four dimensions") + pixels = pixels.reshape(-1, *pixels.shape[-3:]) + weight = self.vision_model.conv1.weight + pixels = pixels.to(device=weight.device, dtype=weight.dtype) + features = self.vision_model(pixels) + grid_h, grid_w = (side // self.vision_model.patch_size for side in pixels.shape[-2:]) + features = features.transpose(1, 2).reshape(pixels.shape[0], -1, grid_h, grid_w) + features = self.vision_model.vit_downsampler1(features) + features = self.vision_model.vit_downsampler2(features) + return self.vit_large_projector(features.flatten(2).transpose(1, 2)) + + def get_multimodal_embeddings( + self, *, pixel_values=None, patch_pixel_values=None, num_patches=None, image_embeds=None, + ): + """Return final projected tokens in image order, crops before each overview. + + ``image_embeds`` are already projected [..., text_hidden_size] tensors. + Pixel inputs must use the checkpoint's external image preprocessing. + ``num_patches`` has one crop count per overview image, across packed requests. + """ + if image_embeds is not None: + if any(value is not None for value in (pixel_values, patch_pixel_values, num_patches)): + raise ValueError("image_embeds and pixel/crop inputs are mutually exclusive") + if image_embeds.ndim < 2 or image_embeds.shape[-1] != self.config.text_config.hidden_size: + raise ValueError("image_embeds must end in text hidden_size") + return image_embeds.reshape(-1, image_embeds.shape[-1]) + if pixel_values is None: + if patch_pixel_values is not None or num_patches is not None: + raise ValueError("Crop inputs require pixel_values overview images") + return None + if pixel_values.ndim < 4: + raise ValueError("pixel_values must have at least four dimensions") + pixels = pixel_values.reshape(-1, *pixel_values.shape[-3:]) + counts = [0] * pixels.shape[0] if num_patches is None else ( + num_patches.tolist() if isinstance(num_patches, torch.Tensor) else list(num_patches) + ) + if len(counts) != pixels.shape[0] or any(type(n) is not int or n < 0 for n in counts): + raise ValueError("num_patches must contain one nonnegative integer per overview image") + patches = None + if patch_pixel_values is not None: + if patch_pixel_values.ndim < 4: + raise ValueError("patch_pixel_values must have at least four dimensions") + patches = patch_pixel_values.reshape(-1, *patch_pixel_values.shape[-3:]) + if sum(counts) != (0 if patches is None else patches.shape[0]): + raise ValueError("sum(num_patches) must equal the number of crop images") + overview = self._encode_images(pixels) + crops = self._encode_images(patches) if patches is not None and patches.shape[0] else None + merged, offset = [], 0 + for i, count in enumerate(counts): + if count: + merged.append(crops[offset:offset + count].flatten(0, 1)) + merged.append(overview[i]) + offset += count + return torch.cat(merged, dim=0) + + def forward( + self, input_ids, cache, seq_lens, *, is_decode=False, return_all_logits=False, + inputs_embeds=None, pixel_values=None, patch_pixel_values=None, + num_patches=None, image_embeds=None, + ): + supplied_images = any(value is not None for value in ( + pixel_values, patch_pixel_values, num_patches, image_embeds, + )) + if inputs_embeds is not None and supplied_images: + raise ValueError("inputs_embeds and image inputs are mutually exclusive") + if is_decode and supplied_images: + raise ValueError("Image inputs are only accepted during prefill") + if supplied_images: + features = self.get_multimodal_embeddings( + pixel_values=pixel_values, patch_pixel_values=patch_pixel_values, + num_patches=num_patches, image_embeds=image_embeds, + ) + mask = input_ids == self.config.image_token_id + count = int(mask.sum().item()) + if features.shape[0] != count: + raise ValueError(f"Image features contain {features.shape[0]} tokens but prompt has {count} image placeholders") + # The placeholder need not be an in-vocabulary embedding-table entry. + inputs_embeds = self.model.embed_tokens(input_ids.masked_fill(mask, 0)) + inputs_embeds[mask] = features.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + elif inputs_embeds is None and not is_decode and torch.any(input_ids == self.config.image_token_id): + raise ValueError("Image placeholders require image inputs during prefill") + return super().forward( + input_ids, cache, seq_lens, is_decode=is_decode, + return_all_logits=return_all_logits, inputs_embeds=inputs_embeds, + ) + + +def _packed_positions(seq_lens: torch.Tensor, past_lens: torch.Tensor) -> torch.Tensor: + """Absolute position of every token in a packed batch.""" + device = seq_lens.device + spans = [ + torch.arange(int(past), int(past) + int(length), device=device, dtype=torch.int32) + for past, length in zip(past_lens.tolist(), seq_lens.tolist()) + ] + return torch.cat(spans) + + +def _last_token_indices(seq_lens: torch.Tensor) -> torch.Tensor: + ends = torch.cumsum(seq_lens.to(torch.long), dim=0) + return ends - 1 + + +__all__ = [ + "STEP5_TP_LAYOUT_VERSION", + "Step5VisionConfig", + "Step5VisionEncoder", + "Step5VLConfig", + "Step5ForConditionalGeneration", + "config_from_dict", + "model_class_for_config", + "Step5SparseConfig", + "Step5Config", + "set_tp_group", + "tp_group", + "tp_rank", + "tp_size", + "ColumnParallelLinear", + "RowParallelLinear", + "VocabParallelLinear", + "vocab_parallel_argmax", + "DenseLayerCache", + "Step5Attention", + "Step5Cache", + "Step5RMSNorm", + "Step5DecoderLayer", + "Step5ForCausalLM", + "Step5MLP", + "Step5MoE", + "Step5Model", + "clamped_swiglu", +]