"""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", ]