# SPDX-License-Identifier: Apache-2.0 # Copyright 2023-2024 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Inference-only Zyphra ZAYA1 (CCA attention + MoE) model implementation. Architecture summary (see docs/supported_models/text_generation/zaya_design.md for the full design notes): - Even-indexed layers run :class:`ZayaAttention`, which feeds hidden states to the :class:`CCA` (Compressed Convolutional Attention) projection. CCA emits q/k/v via two small (``kernel_size=2``) depthwise + grouped 1D convolutions over the time axis plus a learnable per-K-head temperature. The conv needs a two-token left padding that is sourced from a per-request state cache owned by the CCA module itself. The q/k/v then go through partial rotary embedding (``partial_rotary_factor=0.5``) and SGLang's :class:`RadixAttention` for the softmax MHA. The implementation only uses ``torch`` / ``torch.nn`` ops, so the same code runs on NVIDIA and AMD GPUs. - Odd-indexed layers run :class:`ZayaBlock`, an MoE mixer built around SGLang's :class:`FusedMoE`. Expert routing uses a 3-layer MLP with EDA (depth-wise averaging across MoE layers) and MOD (mixture-of-depths skip expert). - Per-layer :class:`ResidualScaling` keeps the residual stream in fp32 with affine scale/bias both on the residual and on the post-mixer hidden states. - Per-request CCA state (``conv_state`` + ``prev_hs``) lives in SGLang's centralized ``MambaPool`` inside ``HybridReqToTokenPool``. The per-request state plumbing (slot indices, prefix mask, cuda-graph buffers) is owned by ``ShortConvAttnBackend`` and reached via ``get_attn_backend().conv_state_metadata()``, so the model holds no pool access; CCA runs its own conv (:func:`cca_extend` / :func:`cca_decode`) against the returned handle. """ from __future__ import annotations import logging import re from collections.abc import Iterable from typing import List, Optional, Tuple import torch import torch.nn.functional as F from torch import nn from sglang.srt.configs.zaya import ZayaConfig from sglang.srt.distributed import ( get_pp_group, tensor_model_parallel_all_reduce, ) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( ColumnParallelLinear, ReplicatedLinear, RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.topk import StandardTopKOutput from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.rotary_embedding import get_rope from sglang.srt.layers.utils import PPMissingLayer from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix, make_layers, set_weight_attrs logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Residual scaling # --------------------------------------------------------------------------- class ResidualScaling(nn.Module): """Affine fp32 scaling applied to the residual / hidden_states streams. Layer 0 has no incoming residual stream, so its checkpoint omits ``residual_scale`` / ``residual_bias`` and ``has_residual`` stays False. """ def __init__(self, config: ZayaConfig, layer_n: int) -> None: super().__init__() self.hidden_size = config.hidden_size self.has_residual = layer_n != 0 self.hidden_states_scale = nn.Parameter(torch.ones(self.hidden_size)) self.hidden_states_bias = nn.Parameter(torch.zeros(self.hidden_size)) if self.has_residual: self.residual_scale = nn.Parameter(torch.ones(self.hidden_size)) self.residual_bias = nn.Parameter(torch.zeros(self.hidden_size)) def forward( self, residual: Optional[torch.Tensor], hidden_states: torch.Tensor, ) -> tuple[Optional[torch.Tensor], torch.Tensor]: hs_scale = self.hidden_states_scale.to(torch.float32) hs_bias = self.hidden_states_bias.to(torch.float32) hidden_states = (hidden_states.float() + hs_bias) * hs_scale if self.has_residual and residual is not None: res_scale = self.residual_scale.to(torch.float32) res_bias = self.residual_bias.to(torch.float32) residual = (residual.float() + res_bias) * res_scale return residual, hidden_states def _apply_norm_with_fp32_residual( norm: nn.Module, residual: torch.Tensor, target_dtype: torch.dtype, ) -> torch.Tensor: """Normalize ``residual`` (typically fp32) and cast back to ``target_dtype``. The fp32 residual stream is preserved by the caller (the residual tensor is kept around for the next accumulation), so the norm itself can run at ``target_dtype`` -- this lets us hit the fused sgl_kernel rmsnorm path instead of the eager ``forward_native`` fallback (5+ kernel launches per call, ×120 norms per step). """ return norm(residual.to(target_dtype)) # --------------------------------------------------------------------------- # CCA conv-state kernels (v1 torch) # # ZAYA1-specific conv step: the CCA conv is a causal two-stage conv over # ``qk = [W_q hs || W_k hs]`` plus a one-token ``prev_hs`` lag for val_proj2. # The per-request conv state lives in the centralized MambaPool; the backend # (ShortConvAttnBackend) hands out the slot indices + prefix flags and CCA runs # these functions against them. ``conv_qk`` is the module's two-stage conv; # both functions mutate ``conv_state`` / ``prev_hs_state`` in place and return # ``(qk_out, v2_input)`` -- the conv output ``[T, in_out_ch]`` and the (shifted) # ``val_proj2`` input ``[T, hidden_size]``. # --------------------------------------------------------------------------- def cca_extend( qk: torch.Tensor, hidden_states: torch.Tensor, conv_qk: nn.Module, conv_state: torch.Tensor, prev_hs_state: torch.Tensor, slot_ids: List[int], has_prefix: List[bool], extend_seq_lens_cpu: List[int], total_padding: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Prefill / extend conv-state step (v1, pure torch). Walks each request in the batch, applies ``conv_qk`` with the request's own initial state (zeros on a fresh first chunk, the cached ``conv_state`` slot otherwise), writes the updated ``conv_state`` / ``prev_hs_state`` back, and returns the concatenated ``(qk_out, v2_input)`` in the original token layout. ``slot_ids`` is the host mirror of the per-request MambaPool slot indices and ``has_prefix[i]`` is ``True`` when request ``i`` resumes a cached prefix. The Triton swap (:func:`cca_conv1d_fn`) removes this per-request loop. """ dtype = hidden_states.dtype if total_padding is None: total_padding = conv_state.shape[-1] in_out_ch = qk.shape[-1] hidden_size = hidden_states.shape[-1] qk_out = torch.empty_like(qk) v2_input = torch.empty_like(hidden_states) # Fresh-prefill fast path: when no request has a cached prefix the per-request # convs can be coalesced into a single packed convolution. Each request's # segment is laid out as ``[total_padding zeros, S_i tokens]``. all_fresh = bool(extend_seq_lens_cpu) and not any(has_prefix) if all_fresh: seq_lens = [int(s) for s in extend_seq_lens_cpu] pad = total_padding offsets_in = [0] for s in seq_lens: offsets_in.append(offsets_in[-1] + s + pad) packed = qk.new_zeros((1, in_out_ch, offsets_in[-1])) start = 0 for i, s in enumerate(seq_lens): end = start + s packed[0, :, offsets_in[i] + pad : offsets_in[i + 1]] = qk[ start:end ].transpose(0, 1) start = end packed_out = conv_qk(packed) # [1, C, offsets_in[-1] - pad] start = 0 for i, s in enumerate(seq_lens): end = start + s a_i = offsets_in[i] qk_out[start:end] = packed_out[0, :, a_i : a_i + s].transpose(0, 1) new_state = packed[0, :, a_i + s : a_i + s + pad] conv_state[slot_ids[i]] = new_state.to(conv_state.dtype) hs_cur = hidden_states[start:end] first = hidden_states.new_zeros((1, hidden_size)) v2_input[start:end] = torch.cat([first, hs_cur[:-1]], dim=0) prev_hs_state[slot_ids[i]] = ( hs_cur[-1].unsqueeze(-1).to(prev_hs_state.dtype) ) start = end else: start = 0 for i, seq_len in enumerate(extend_seq_lens_cpu): end = start + int(seq_len) slot = slot_ids[i] prefix = bool(has_prefix[i]) qk_cur = qk[start:end].transpose(0, 1).unsqueeze(0) # [1, C, S_cur] if prefix: left_pad = conv_state[slot].unsqueeze(0).to(dtype) else: left_pad = qk_cur.new_zeros((1, in_out_ch, total_padding)) padded = torch.cat([left_pad, qk_cur], dim=-1) out = conv_qk(padded) # [1, C, S_cur] qk_out[start:end] = out.squeeze(0).transpose(0, 1) new_state = padded[..., -total_padding:] conv_state[slot] = new_state.squeeze(0).to(conv_state.dtype) hs_cur = hidden_states[start:end] if prefix: first = prev_hs_state[slot].squeeze(-1).to(dtype).unsqueeze(0) else: first = hidden_states.new_zeros((1, hidden_size)) v2_input[start:end] = torch.cat([first, hs_cur[:-1]], dim=0) prev_hs_state[slot] = hs_cur[-1].unsqueeze(-1).to(prev_hs_state.dtype) start = end return qk_out, v2_input def cca_decode( qk: torch.Tensor, hidden_states: torch.Tensor, conv_qk: nn.Module, conv_state: torch.Tensor, prev_hs_state: torch.Tensor, mamba_indices: torch.Tensor, total_padding: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Single-token decode conv-state step (v1, pure torch). Gathers each request's cached ``conv_state`` / ``prev_hs_state`` via ``index_select``, runs ``conv_qk`` on the ``[T, C, total_padding + 1]`` window, and scatters the updated state back with ``index_copy_``. All ops are on-device (``mamba_indices`` is a device ``long`` tensor), so this stays CUDA-graph capturable. Returns ``(qk_out, prev_hs)`` where ``prev_hs`` is the previous hidden state feeding ``val_proj2``. The Triton swap is :func:`cca_conv1d_update`. """ dtype = hidden_states.dtype if total_padding is None: total_padding = conv_state.shape[-1] left_pad = conv_state.index_select(0, mamba_indices).to(dtype) cur = qk.unsqueeze(-1) # [T, C, 1] padded = torch.cat([left_pad, cur], dim=-1) # [T, C, total_padding + 1] out = conv_qk(padded) # [T, C, 1] qk_out = out.squeeze(-1) # [T, C] new_state = padded[..., -total_padding:] conv_state.index_copy_(0, mamba_indices, new_state.to(conv_state.dtype)) # Read the previous hidden state (val_proj2 input) BEFORE overwriting the # slot with the current token. prev_hs = prev_hs_state.index_select(0, mamba_indices).squeeze(-1).to(dtype) prev_hs_state.index_copy_( 0, mamba_indices, hidden_states.unsqueeze(-1).to(prev_hs_state.dtype) ) return qk_out, prev_hs # Fused kernel seam (TODO) -- perf swap for the v1 torch paths above. These # mirror the ``causal_conv1d_fn`` / ``causal_conv1d_update`` contract but for # CCA's two-stage *grouped* conv (conv_qk[0] depthwise + conv_qk[1] grouped # per-head), which the stock depthwise ``causal_conv1d`` cannot express. Once # implemented they replace the per-request loop in ``cca_extend`` and the # separate gather/conv/scatter launches in ``cca_decode`` with a single # index-driven kernel. Same ``(qk_out, v2_input)`` return contract. def cca_conv1d_fn(*args, **kwargs): raise NotImplementedError( "Fused CCA prefill conv-with-state kernel not implemented yet; " "the model uses cca_extend (v1 torch) in the meantime." ) def cca_conv1d_update(*args, **kwargs): raise NotImplementedError( "Fused CCA decode conv-with-state kernel not implemented yet; " "the model uses cca_decode (v1 torch) in the meantime." ) # --------------------------------------------------------------------------- # CCA: Compressed Convolutional Attention QKV projection # --------------------------------------------------------------------------- class CCA(nn.Module): """Compressed Convolutional Attention QKV projection. Given hidden states ``hs`` of shape ``[S, H]`` this layer produces ``(q, k, v)`` where: q = (W_q hs + Conv(W_q hs ‖ W_k hs)_q) / 2 + mean_group(W_k hs) / 2 (fp32, RMSNorm'd) k = (W_k hs + Conv(W_q hs ‖ W_k hs)_k) / 2 + mean_group(W_q hs) / 2, scaled by per-head temperature v = concat(W_{v1} hs, W_{v2} hs_prev_shifted) The two-stage conv on ``(W_q hs ‖ W_k hs)`` needs ``total_padding = (cca_time0 - 1) + (cca_time1 - 1)`` tokens of left padding. For the first prefill chunk of a request the padding is zero; for a resumed prefill or for decode it is read from a per-request cache that this module maintains internally. Parallelism: when ``tp_size > 1`` the CCA is head-parallel. Both the grouped-mean step and the second ``conv_qk`` stage with ``groups=num_q_heads+num_k_heads`` are head-local (each GQA group lives on a single rank), so the entire QKV projection runs without any cross-rank collective. The QKV projections become ``ColumnParallelLinear`` and the two ``nn.Conv1d`` layers are sized per-rank with custom weight loaders that slice the HF checkpoint rows into ``[rank's q heads, rank's k heads]``. """ def __init__( self, config: ZayaConfig, cca_num_k_heads: int, cca_num_q_heads: int, hidden_size: int, head_dim: int, cca_time0: int, cca_time1: int, layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", tp_rank: Optional[int] = None, tp_size: Optional[int] = None, ) -> None: super().__init__() self.config = config self.layer_id = layer_id self.hidden_size = int(hidden_size) self.head_dim = int(head_dim) self.cca_time0 = int(cca_time0) self.cca_time1 = int(cca_time1) self.padding0 = self.cca_time0 - 1 self.padding1 = self.cca_time1 - 1 self.total_padding = self.padding0 + self.padding1 if tp_rank is None: tp_rank = get_parallel().tp_rank if tp_size is None: tp_size = get_parallel().tp_size self.tp_rank = int(tp_rank) self.tp_size = int(tp_size) # Full (global) head counts retained for weight loading and shape asserts. self.num_q_heads_full = int(cca_num_q_heads) self.num_k_heads_full = int(cca_num_k_heads) assert ( self.num_q_heads_full % self.num_k_heads_full == 0 ), "num_q_heads must be a multiple of num_k_heads" self.gqa_groups = self.num_q_heads_full // self.num_k_heads_full # Head-parallel TP requires both head counts to be divisible by tp_size. # KV-replication-style TP (tp_size > num_k_heads) is not yet supported. assert self.num_q_heads_full % self.tp_size == 0, ( f"num_q_heads ({self.num_q_heads_full}) must be divisible by " f"tp_size ({self.tp_size}) for ZAYA1 head-parallel CCA" ) assert self.num_k_heads_full % self.tp_size == 0, ( f"num_k_heads ({self.num_k_heads_full}) must be divisible by " f"tp_size ({self.tp_size}); KV-replication TP is not supported " "for ZAYA1 because both grouped-mean and conv_qk.1 are per-head" ) # Per-rank head counts. self.num_q_heads = self.num_q_heads_full // self.tp_size self.num_k_heads = self.num_k_heads_full // self.tp_size # Per-rank channel layout. self.latent_q_dim_full = self.num_q_heads_full * self.head_dim self.latent_k_dim_full = self.num_k_heads_full * self.head_dim self.in_out_ch_full = self.latent_q_dim_full + self.latent_k_dim_full self.latent_q_dim = self.num_q_heads * self.head_dim self.latent_k_dim = self.num_k_heads * self.head_dim self.in_out_ch = self.latent_q_dim + self.latent_k_dim self.sqrt_head_dim = float(self.head_dim) ** 0.5 self.clamp_temp = bool(getattr(config, "clamp_temp", False)) bias = bool(getattr(config, "attention_bias", False)) # ``linear_q`` / ``linear_k`` outputs are laid out as a contiguous head # sequence in the HF checkpoint, so the natural ColumnParallel shard # (``tp_rank * shard``) lands rank ``r`` on the head set # ``[r * heads_per_rank, (r+1) * heads_per_rank)``. # # At ``tp_size == 1`` there is nothing to shard, and on ROCm/aiter the # ColumnParallelLinear path selects a slower GEMM for the large-M prefill # (1.6-2.25x slower than ReplicatedLinear in bench_one_batch), so the # single-GPU case uses ReplicatedLinear. ``tp_size > 1`` keeps # ColumnParallelLinear for the per-rank head shard. if self.tp_size > 1: self.linear_q = ColumnParallelLinear( self.hidden_size, self.latent_q_dim_full, bias=bias, gather_output=False, quant_config=quant_config, prefix=add_prefix("linear_q", prefix), tp_rank=self.tp_rank, tp_size=self.tp_size, ) self.linear_k = ColumnParallelLinear( self.hidden_size, self.latent_k_dim_full, bias=bias, gather_output=False, quant_config=quant_config, prefix=add_prefix("linear_k", prefix), tp_rank=self.tp_rank, tp_size=self.tp_size, ) else: self.linear_q = ReplicatedLinear( self.hidden_size, self.latent_q_dim_full, bias=bias, quant_config=quant_config, prefix=add_prefix("linear_q", prefix), ) self.linear_k = ReplicatedLinear( self.hidden_size, self.latent_k_dim_full, bias=bias, quant_config=quant_config, prefix=add_prefix("linear_k", prefix), ) # The HF V-projection layout maps val_proj1 to the FIRST half of K # heads and val_proj2 to the SECOND half (after ``cat([v1, v2]).view( # T, num_k_heads_full, head_dim)``). That doesn't align with a simple # output-dim ColumnParallel shard, so val_proj1 / val_proj2 are kept # Replicated and the per-rank K-head slice is taken in the forward # passes after ``cat + view``. The replicated weight memory is small # (~0.5 MB / layer) and the wasted compute is negligible compared to # linear_q / linear_k / o_proj. self.val_proj1 = ReplicatedLinear( self.hidden_size, self.latent_k_dim_full // 2, bias=bias, quant_config=quant_config, prefix=add_prefix("val_proj1", prefix), ) self.val_proj2 = ReplicatedLinear( self.hidden_size, self.latent_k_dim_full // 2, bias=bias, quant_config=quant_config, prefix=add_prefix("val_proj2", prefix), ) # Per-rank K head range, used for slicing the replicated v tensor. self.k_head_start = self.tp_rank * self.num_k_heads self.k_head_end = self.k_head_start + self.num_k_heads # Two-stage depthwise + grouped conv along the time axis, sized for # this rank's head subset. Wrapping the two nn.Conv1d modules in # nn.Sequential makes the HF checkpoint keys ``conv_qk.{0,1}.weight`` # / ``conv_qk.{0,1}.bias`` map onto submodules 1:1, with TP slicing # handled by the custom weight_loader attached below. self.conv_qk = nn.Sequential( nn.Conv1d( in_channels=self.in_out_ch, out_channels=self.in_out_ch, kernel_size=self.cca_time0, groups=self.in_out_ch, padding=0, stride=1, ), nn.Conv1d( in_channels=self.in_out_ch, out_channels=self.in_out_ch, kernel_size=self.cca_time1, groups=(self.num_k_heads + self.num_q_heads), padding=0, stride=1, ), ) # Per-K-head learnable temperature scalar (per-rank slice). self.temp = nn.Parameter(torch.zeros(self.num_k_heads)) # Attach TP-aware weight loaders to conv_qk weights/biases and ``temp`` # so the existing ``load_weights`` dispatch (``getattr(param, # "weight_loader", default_weight_loader)``) automatically slices the # HF checkpoint into rank-local rows. if self.tp_size > 1: self._install_tp_weight_loaders() # ----- TP weight loaders ---------------------------------------------- def _install_tp_weight_loaders(self) -> None: """Attach TP-aware ``weight_loader`` attributes to parameters whose full-tensor → per-rank slicing cannot be expressed by a generic ColumnParallelLinear loader: the two ``conv_qk`` Conv1d weights and biases (where the per-rank "row" set is the discontiguous union of this rank's q heads and this rank's k heads) and the per-K-head ``temp`` parameter. """ head_dim = self.head_dim latent_q_dim_full = self.latent_q_dim_full num_q_heads_per_rank = self.num_q_heads num_k_heads_per_rank = self.num_k_heads tp_rank = self.tp_rank q_start = tp_rank * num_q_heads_per_rank * head_dim q_end = q_start + num_q_heads_per_rank * head_dim k_start = latent_q_dim_full + tp_rank * num_k_heads_per_rank * head_dim k_end = k_start + num_k_heads_per_rank * head_dim k_temp_start = tp_rank * num_k_heads_per_rank k_temp_end = k_temp_start + num_k_heads_per_rank def conv_row_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: # Both Conv1d.weight ([C_out, in_per_group, K]) and Conv1d.bias # ([C_out]) slice along the leading (output channel) dim. The # per-rank rows are the rank's q heads (contiguous) followed by # the rank's k heads (contiguous in the second half of the full # tensor). sliced = torch.cat( [loaded_weight[q_start:q_end], loaded_weight[k_start:k_end]], dim=0, ) assert ( sliced.shape == param.data.shape ), f"conv shard shape mismatch: {sliced.shape} vs {param.data.shape}" param.data.copy_(sliced) def temp_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: sliced = loaded_weight[k_temp_start:k_temp_end] assert ( sliced.shape == param.data.shape ), f"temp shard shape mismatch: {sliced.shape} vs {param.data.shape}" param.data.copy_(sliced) set_weight_attrs(self.conv_qk[0].weight, {"weight_loader": conv_row_loader}) set_weight_attrs(self.conv_qk[0].bias, {"weight_loader": conv_row_loader}) set_weight_attrs(self.conv_qk[1].weight, {"weight_loader": conv_row_loader}) set_weight_attrs(self.conv_qk[1].bias, {"weight_loader": conv_row_loader}) set_weight_attrs(self.temp, {"weight_loader": temp_loader}) # ----- helpers --------------------------------------------------------- def _normalize_qk( self, query: torch.Tensor, key: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: """RMSNorm (no learnable weight) + sqrt(head_dim) scaling on q and k, plus per-K-head temperature on k. Computed in fp32 for stability. """ eps = 1e-12 sqrt_head_dim = float(self.sqrt_head_dim) query_fp32 = query.to(torch.float32) inv_q = ( torch.rsqrt(query_fp32.pow(2).sum(-1, keepdim=True) + eps) * sqrt_head_dim ) query_fp32 = query_fp32 * inv_q key_fp32 = key.to(torch.float32) inv_k = torch.rsqrt(key_fp32.pow(2).sum(-1, keepdim=True) + eps) * sqrt_head_dim key_fp32 = key_fp32 * inv_k temp = self.temp.to(torch.float32).view(1, self.num_k_heads, 1) if self.clamp_temp: temp = torch.exp(torch.clamp(temp, 1e-7, 2.0)) key_fp32 = key_fp32 * temp return query_fp32, key_fp32 def _add_grouped_qk_means( self, query_conv: torch.Tensor, key_conv: torch.Tensor, query_pre: torch.Tensor, key_base: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Blend the post-conv q/k with the per-GQA-group mean of the pre-conv (raw projection) q/k, matching the ZAYA1 training formula. Shapes (T = num_tokens): query_conv : [T, num_q_heads, head_dim] (fp32, post conv) key_conv : [T, num_k_heads, head_dim] (fp32, post conv) query_pre : [T, num_q_heads, head_dim] (raw W_q hs) key_base : [T, num_k_heads, head_dim] (raw W_k hs) """ num_k_heads = key_base.shape[-2] key_base_fp32 = key_base.to(torch.float32) query_pre_grouped = query_pre.view( query_pre.shape[0], num_k_heads, self.gqa_groups, query_pre.shape[-1] ) query_pre_grouped_fp32 = query_pre_grouped.to(torch.float32) query_out_grouped = ( query_conv.view_as(query_pre_grouped).to(torch.float32) + 0.5 * query_pre_grouped_fp32 + 0.5 * key_base_fp32.unsqueeze(-2) ) query_out = query_out_grouped.reshape( query_pre.shape[0], -1, query_pre.shape[-1] ) query_pre_mean = query_pre_grouped_fp32.mean(dim=-2, dtype=torch.float32) key_out = ( key_conv.to(torch.float32) + 0.5 * query_pre_mean + 0.5 * key_base_fp32 ) return query_out, key_out def _conv_qk_run(self, padded: torch.Tensor) -> torch.Tensor: """Run ``conv_qk`` on ``[N, C, S + total_padding]`` → ``[N, C, S]``.""" return self.conv_qk(padded) # ----- forward modes --------------------------------------------------- def _slice_v_per_rank(self, value_full: torch.Tensor) -> torch.Tensor: """Take this rank's K-head slice of the full ``value`` tensor. Returns a no-op view when ``tp_size == 1``. For ``tp_size > 1`` the full V tensor is computed on every rank (see the comment on ``val_proj1`` / ``val_proj2``) and the rank's contiguous K-head range is selected here, leaving the downstream RadixAttention call with a per-rank shape ``[T, num_k_heads_per_rank, head_dim]``. """ if self.tp_size == 1: return value_full return value_full[:, self.k_head_start : self.k_head_end, :].contiguous() def _forward_no_state( self, hs: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Reference path: process the entire ``hs`` of shape ``[S, H]`` with a zero initial conv state and a zero ``prev_hs``. Exercised by the CCA unit tests so the prefill / decode paths can be compared against a single-shot torch reference, and used as a fallback for profile / warmup runs where no state cache is meaningful. """ S = hs.shape[0] hs_3d = hs.unsqueeze(1) # [S, 1, H] q_raw, _ = self.linear_q(hs_3d) # [S, 1, latent_q_dim_per_rank] k_raw, _ = self.linear_k(hs_3d) # [S, 1, latent_k_dim_per_rank] qk = torch.cat([q_raw, k_raw], dim=-1) # [S, 1, in_out_ch_per_rank] query_pre = q_raw.view(S, self.num_q_heads, self.head_dim) key_base = k_raw.view(S, self.num_k_heads, self.head_dim) # [1, C, S+pad] -> [1, C, S] qk_perm = qk.permute(1, 2, 0) qk_pad = F.pad(qk_perm, (self.total_padding, 0)) qk_out = self._conv_qk_run(qk_pad).permute(2, 0, 1).squeeze(1) # [S, C] query_conv = qk_out[:, : self.latent_q_dim].view( S, self.num_q_heads, self.head_dim ) key_conv = qk_out[:, self.latent_q_dim :].view( S, self.num_k_heads, self.head_dim ) query, key = self._add_grouped_qk_means( query_conv, key_conv, query_pre, key_base ) query, key = self._normalize_qk(query, key) # val_proj1 / val_proj2 are replicated; compute the full V tensor and # then take this rank's K-head slice. # val_proj2 uses a right-shifted hidden_state. First val_proj2 input is 0. hs_shifted = F.pad(hs_3d[:-1], (0, 0, 0, 0, 1, 0)) # [S, 1, H] v1, _ = self.val_proj1(hs_3d) v2, _ = self.val_proj2(hs_shifted) value_full = ( torch.cat([v1, v2], dim=-1) .squeeze(1) .view(S, self.num_k_heads_full, self.head_dim) ) value = self._slice_v_per_rank(value_full) return query, key, value def forward( self, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Project ``hidden_states`` into ``(q, k, v)`` honoring per-request state. The per-request conv-state plumbing (slot gather/scatter, prefix mask, cuda-graph buffers) is owned by :class:`ShortConvAttnBackend `, reached via ``get_attn_backend().conv_state_metadata``; CCA runs its own two-stage grouped conv (:func:`cca_extend` / :func:`cca_decode`) against that handle, so this module holds no pool access. Those functions return the conv output ``qk_out`` and the ``val_proj2`` input ``v2_input`` (the shifted / previous hidden state), updating the ``conv_state`` / ``prev_hs`` pool slots in place. ``q`` / ``k`` are returned in fp32 (the normalize step keeps fp32 for stability); ``v`` is returned in the input dtype since the caller casts everything back to ``hidden_states.dtype`` before rotary + attention anyway. Shapes:: q : [T, num_q_heads, head_dim] k : [T, num_k_heads, head_dim] v : [T, num_k_heads, head_dim] """ if hidden_states.shape[0] == 0: zero = hidden_states.new_zeros((0,)) return ( zero.view(0, self.num_q_heads, self.head_dim).to(torch.float32), zero.view(0, self.num_k_heads, self.head_dim).to(torch.float32), zero.view(0, self.num_k_heads, self.head_dim), ) T = hidden_states.shape[0] q_raw, _ = self.linear_q(hidden_states) # [T, latent_q] k_raw, _ = self.linear_k(hidden_states) qk = torch.cat([q_raw, k_raw], dim=-1) # [T, in_out_ch] query_pre = q_raw.view(T, self.num_q_heads, self.head_dim) key_base = k_raw.view(T, self.num_k_heads, self.head_dim) # The backend hands out the per-request conv-state handle (slot indices, # prefix mask, cuda-graph buffers); CCA runs its own two-stage grouped # conv against it and gets back the conv output + val_proj2 input, with # the conv_state / prev_hs pool slots updated in place. meta = get_attn_backend().conv_state_metadata(self.layer_id, forward_batch) conv_state = meta.layer_cache.conv[0] prev_hs_state = meta.layer_cache.conv[1] if forward_batch.forward_mode.is_decode_or_idle(): qk_out, v2_input = cca_decode( qk, hidden_states, self.conv_qk, conv_state, prev_hs_state, meta.cache_indices, self.total_padding, ) else: qk_out, v2_input = cca_extend( qk, hidden_states, self.conv_qk, conv_state, prev_hs_state, meta.slot_ids_cpu, meta.has_prefix_cpu, forward_batch.extend_seq_lens_cpu, self.total_padding, ) query_conv = qk_out[:, : self.latent_q_dim].view( T, self.num_q_heads, self.head_dim ) key_conv = qk_out[:, self.latent_q_dim :].view( T, self.num_k_heads, self.head_dim ) query, key = self._add_grouped_qk_means( query_conv, key_conv, query_pre, key_base ) query, key = self._normalize_qk(query, key) v1, _ = self.val_proj1(hidden_states) v2, _ = self.val_proj2(v2_input) value_full = torch.cat([v1, v2], dim=-1).view( T, self.num_k_heads_full, self.head_dim ) value = self._slice_v_per_rank(value_full) return query, key, value # --------------------------------------------------------------------------- # Attention layer (CCA QKV + rotary + RadixAttention) # --------------------------------------------------------------------------- class ZayaAttention(nn.Module): def __init__( self, config: ZayaConfig, layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config self.layer_id = layer_id self.hidden_size = config.hidden_size self.num_q_heads_full = config.num_attention_heads self.num_k_heads_full = config.num_query_groups self.head_dim = config.head_dim # Head-parallel TP: split both Q and KV heads across ranks. Since the # grouped-mean and conv_qk.1 are head-local, no cross-rank collective # is required inside the QKV projection. Both head counts must be # divisible by tp_size; the KV-replicated GQA-TP variant (tp_size > # num_k_heads) is intentionally rejected with a clear error message # because both per-K-head paths assume each rank holds whole K heads. self.tp_rank = get_parallel().tp_rank self.tp_size = get_parallel().tp_size # The head split, the ``o_proj`` RowParallel all-reduce, and the # RadixAttention KV cache are all organized on the *global* TP group, # and ``ZayaConfig.mamba2_cache_params`` sizes the conv-state cache on # that same group. DP attention would run attention on the smaller # attention-TP group (and ``o_proj`` would need # ``use_dp_attention_reduce``), which this model does not wire up, so # require the two groups to coincide and fail fast instead of silently # mis-sizing the conv-state cache. attn_tp_size = get_parallel().attn_tp_size assert attn_tp_size == self.tp_size, ( f"ZAYA1 head-parallel attention requires the attention TP group " f"({attn_tp_size}) to equal the global TP group ({self.tp_size}); " "DP attention (enable_dp_attention) is not supported for ZAYA1." ) assert self.num_q_heads_full % self.tp_size == 0, ( f"num_attention_heads ({self.num_q_heads_full}) must be divisible " f"by tp_size ({self.tp_size}) for ZAYA1 head-parallel attention" ) assert self.num_k_heads_full % self.tp_size == 0, ( f"num_query_groups ({self.num_k_heads_full}) must be divisible by " f"tp_size ({self.tp_size}); set tp_size <= num_k_heads to keep " "both grouped-mean and conv_qk.1 head-local on each rank" ) self.num_q_heads = self.num_q_heads_full // self.tp_size self.num_k_heads = self.num_k_heads_full // self.tp_size self.q_dim_full = self.num_q_heads_full * self.head_dim self.scale = self.head_dim**-0.5 # The HF checkpoint stores the CCA QKV projection under # ``self_attn.qkv.*``, so the CCA submodule is registered with that # exact name to keep weight loading a 1:1 key mapping. self.qkv = CCA( config=config, cca_num_k_heads=self.num_k_heads_full, cca_num_q_heads=self.num_q_heads_full, hidden_size=self.hidden_size, head_dim=self.head_dim, cca_time0=config.cca_time0, cca_time1=config.cca_time1, layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("qkv", prefix), tp_rank=self.tp_rank, tp_size=self.tp_size, ) # RowParallel o_proj: per-rank input is the rank's q heads, full # output is replicated via the end-of-forward all-reduce. self.o_proj = RowParallelLinear( self.q_dim_full, self.hidden_size, bias=bool(getattr(config, "attention_bias", False)), input_is_parallel=True, reduce_results=True, quant_config=quant_config, prefix=add_prefix("o_proj", prefix), tp_rank=self.tp_rank, tp_size=self.tp_size, ) rope_theta = float(getattr(config, "rope_theta", 1_000_000.0)) partial_rotary_factor = float(getattr(config, "partial_rotary_factor", 0.5)) self.rotary_emb = get_rope( head_size=self.head_dim, rotary_dim=self.head_dim, max_position=int(config.max_position_embeddings), base=int(rope_theta), is_neox_style=True, partial_rotary_factor=partial_rotary_factor, ) self.attn = RadixAttention( num_heads=self.num_q_heads, head_dim=self.head_dim, scaling=self.scale, num_kv_heads=self.num_k_heads, layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("attn", prefix), ) def forward( self, hidden_states: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: # CCA returns fp32 q/k and input-dtype v as ``[T, heads, head_dim]`` # tensors; flatten the head dim and cast all to the model dtype before # rotary + RadixAttention. q, k, v = self.qkv(hidden_states, forward_batch) target_dtype = hidden_states.dtype q = q.reshape(q.shape[0], -1).to(target_dtype) k = k.reshape(k.shape[0], -1).to(target_dtype) v = v.reshape(v.shape[0], -1).to(target_dtype) q, k = self.rotary_emb(positions, q, k) # Some rotary backends (notably AITER on ROCm) hand back tensors with # a different stride than the input. RadixAttention's KV-store kernel # asserts contiguous layout, so normalize q/k/v before the attention. q = q.contiguous() k = k.contiguous() v = v.contiguous() attn_output = self.attn(q, k, v, forward_batch) output, _ = self.o_proj(attn_output) return output # --------------------------------------------------------------------------- # Router (EDA + MOD) and MoE block # --------------------------------------------------------------------------- class ZayaRouter(nn.Module): """ZAYA1 expert router: 3-layer MLP with optional EDA and MOD. EDA (Exponential Decay Averaging) adds a scaled copy of the previous MoE layer's router hidden_state to the current layer's input, threading state across MoE layers. MOD (Mixture of Depths) reserves the last expert slot as a "skip" expert whose contribution to the residual stream is just the routing probability times the unprocessed hidden_state, letting individual tokens bypass the MoE entirely when the router scores the skip expert highest. """ def __init__( self, config: ZayaConfig, layer_id: int, num_moe_experts: int, moe_router_topk: int, mlp_expansion: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config self.layer_id = layer_id self.hidden_size = config.hidden_size self.router_softmax_fp32 = bool(getattr(config, "zaya_high_prec", False)) self.use_mod = bool(getattr(config, "zaya_use_mod", False)) self.num_experts = (num_moe_experts + 1) if self.use_mod else num_moe_experts self.topk = int(moe_router_topk) self.mlp_expansion = int(mlp_expansion) self.down_proj = ReplicatedLinear( self.hidden_size, self.mlp_expansion, bias=True, quant_config=quant_config, prefix=add_prefix("down_proj", prefix), ) # EDA threads router state from the previous MoE layer through # ``router_states_scale``. The first MoE layer in the model has no # previous state; whether to fold it in is decided at call time based on # ``prev_router_hidden_states``. ln_eps = float(getattr(config, "norm_epsilon", 1e-5)) self.use_eda = bool(getattr(config, "zaya_use_eda", False)) self.rmsnorm_eda = RMSNorm(self.mlp_expansion, eps=ln_eps) if self.use_eda: self.router_states_scale = nn.Parameter(torch.ones(self.mlp_expansion)) self.non_linearity = nn.GELU() self.router_mlp = nn.Sequential( ReplicatedLinear( self.mlp_expansion, self.mlp_expansion, bias=True, quant_config=quant_config, prefix=add_prefix("router_mlp.0", prefix), ), self.non_linearity, ReplicatedLinear( self.mlp_expansion, self.mlp_expansion, bias=True, quant_config=quant_config, prefix=add_prefix("router_mlp.2", prefix), ), self.non_linearity, ReplicatedLinear( self.mlp_expansion, self.num_experts, bias=False, quant_config=quant_config, prefix=add_prefix("router_mlp.4", prefix), ), ) self.register_buffer( "balancing_biases", torch.zeros(self.num_experts, dtype=torch.float32), persistent=True, ) if self.use_mod: with torch.no_grad(): self.balancing_biases[-1] = -1.0 def forward( self, hidden_states: torch.Tensor, prev_router_hidden_states: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: # ``hidden_states`` is ``[T, H]``. hs, _ = self.down_proj(hidden_states) if ( self.use_eda and prev_router_hidden_states is not None and hasattr(self, "router_states_scale") ): hs = hs + prev_router_hidden_states * self.router_states_scale # ``hs`` is a freshly-allocated tensor (output of ``down_proj`` or the # EDA add above) and ``rmsnorm_eda`` is non-residual / out-of-place, # so we can hand the same buffer to the next layer without cloning. router_hidden_states_next = hs hs_norm = self.rmsnorm_eda(hs) # Step through the Sequential manually so the ``(tensor, bias)`` tuple # returned by each ReplicatedLinear is unpacked correctly. out = hs_norm for stage in self.router_mlp: if isinstance(stage, ReplicatedLinear): out, _ = stage(out) else: out = stage(out) logits = out if self.router_softmax_fp32: expert_prob = torch.softmax(logits, dim=-1, dtype=torch.float32) else: expert_prob = torch.softmax(logits, dim=-1) biased = expert_prob.detach().to(torch.float32) + self.balancing_biases _, expert_choice = torch.topk(biased, self.topk, dim=-1) if self.topk > 1 and self.use_mod: skip_idx = self.num_experts - 1 n_mask = expert_choice == skip_idx cumsum_mask = torch.cumsum(n_mask, dim=-1) expert_choice = expert_choice.masked_fill(cumsum_mask > 0, skip_idx) route_prob = torch.gather(expert_prob, dim=1, index=expert_choice) if route_prob.dtype != hidden_states.dtype: route_prob = route_prob.to(hidden_states.dtype) return route_prob, expert_choice, router_hidden_states_next def mod_premask_experts( experts_out: torch.Tensor, indices: torch.Tensor, num_moe_experts: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Mask the (per-rank, pre-all-reduce) expert output for the MOD skip path. Returns ``(mod_mask, masked_experts)`` where ``mod_mask`` is ``1`` for tokens routed to a real expert and ``0`` for tokens routed to the skip slot (``indices == num_moe_experts``), and ``masked_experts = mod_mask * experts_out``. The masking is applied *before* the cross-rank all-reduce so the single reduction yields ``mask · sum_r(partial_r) = mask · experts_out_full`` without the replicated ``mod_out`` term being summed ``tp_size`` times. Pairs with :func:`mod_blend`, which adds the skip-path term back after the reduce. Kept as a free function so the MOD math is unit-testable without a live ``torch.distributed`` group. """ mod_mask = (indices != num_moe_experts).to(experts_out.dtype) return mod_mask, mod_mask * experts_out def mod_blend( masked_experts_reduced: torch.Tensor, mod_mask: torch.Tensor, mod_out: torch.Tensor, ) -> torch.Tensor: """Combine the already-all-reduced masked expert output with the skip path. ``mod_out`` (the skip-expert residual, ``hidden_states * prob``) is replicated on every rank, so it is folded in here -- after the reduce of ``masked_experts`` -- weighted by ``(1 - mod_mask)``. See :func:`mod_premask_experts`. """ return masked_experts_reduced + (1.0 - mod_mask) * mod_out class ZayaBlock(nn.Module): """ZAYA1 MoE mixer: ZayaRouter feeding FusedMoE, with optional MOD residual blend.""" def __init__( self, config: ZayaConfig, layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config self.layer_id = layer_id self.num_moe_experts = int(config.num_experts) self.mlp_expansion = int(config.zaya_mlp_expansion) self.topk = int(getattr(config, "moe_router_topk", 1)) self.tp_size = get_parallel().tp_size if self.tp_size > self.num_moe_experts: raise ValueError( f"Tensor parallel size {self.tp_size} is greater than the " f"number of experts {self.num_moe_experts}" ) assert ( config.activation_func == "swiglu" ), "ZayaBlock only supports SwiGLU activation" assert config.gated_linear_unit, "ZayaBlock requires gated_linear_unit=True" self.router = ZayaRouter( config=config, layer_id=layer_id, num_moe_experts=self.num_moe_experts, moe_router_topk=self.topk, mlp_expansion=self.mlp_expansion, quant_config=quant_config, prefix=add_prefix("router", prefix), ) # ffn_hidden_size is the merged (gate+up) hidden dim; the per-side # intermediate is half. intermediate = int(config.ffn_hidden_size) // 2 self.experts = get_moe_impl_class(quant_config)( num_experts=self.num_moe_experts, top_k=self.topk, hidden_size=config.hidden_size, intermediate_size=intermediate, layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("experts", prefix), activation="silu", reduce_results=False, ) def forward( self, hidden_states: torch.Tensor, prev_router_hidden_states: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor]: if hidden_states.shape[0] == 0: return hidden_states, hidden_states.new_zeros((0, self.mlp_expansion)) probs, indices, router_hs_next = self.router( hidden_states, prev_router_hidden_states ) topk_out = StandardTopKOutput( topk_weights=probs.to(hidden_states.dtype), topk_ids=indices.to(torch.int32), router_logits=probs.to(hidden_states.dtype), ) if self.config.zaya_use_mod: # MOD: clamp the "skip expert" id (== num_moe_experts) into the # valid expert range so FusedMoE never indexes out of bounds; the # mask below decides per-token whether to actually use experts or # the skip path. clamped_ids = torch.clamp(indices, min=0, max=self.num_moe_experts - 1).to( torch.int32 ) topk_out = topk_out._replace(topk_ids=clamped_ids) experts_out = self.experts(hidden_states, topk_out) # ``mod_out`` is computed identically on every TP rank (both # ``hidden_states`` and ``probs`` are replicated). Fold the skip # mask into the per-rank partial experts output *before* # all-reduce so the single reduction yields: # sum_r(mask · partial_r) + (1 - mask) · mod_out # = mask · experts_out_full + (1 - mask) · mod_out # without double-counting ``mod_out`` by tp_size. The two steps are # ``mod_premask_experts`` / ``mod_blend`` so the math is testable # without a live distributed group. mod_out = hidden_states * probs mod_mask, masked_experts = mod_premask_experts( experts_out, indices, self.num_moe_experts ) if self.tp_size > 1: masked_experts = tensor_model_parallel_all_reduce(masked_experts) hidden_out = mod_blend(masked_experts, mod_mask, mod_out) else: hidden_out = self.experts(hidden_states, topk_out) if self.tp_size > 1: hidden_out = tensor_model_parallel_all_reduce(hidden_out) return hidden_out, router_hs_next # --------------------------------------------------------------------------- # Decoder layers # --------------------------------------------------------------------------- class ZayaDecoderATTLayer(nn.Module): """Attention decoder layer: ``res_scale → input_norm → ZayaAttention``.""" def __init__( self, config: ZayaConfig, layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config self.layer_id = layer_id self.self_attn = ZayaAttention( config=config, layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("self_attn", prefix), ) self.input_norm = self._build_norm(config) if config.scale_residual_merge: self.res_scale = ResidualScaling(config, layer_id) else: self.res_scale = None @staticmethod def _build_norm(config: ZayaConfig) -> nn.Module: if config.normalization == "RMSNorm": return RMSNorm(config.hidden_size, eps=config.norm_epsilon) if config.normalization == "LayerNorm": return nn.LayerNorm(config.hidden_size, eps=config.norm_epsilon) raise ValueError(f"Unsupported normalization: {config.normalization}") def forward( self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor], positions: torch.Tensor, forward_batch: ForwardBatch, prev_router_hidden_states: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: target_dtype = ( self.input_norm.weight.dtype if isinstance(self.input_norm, RMSNorm) else hidden_states.dtype ) if self.res_scale is not None: residual, hidden_states = self.res_scale(residual, hidden_states) if residual is not None: residual = residual.float() + hidden_states.float() else: residual = hidden_states.float() hidden_states = _apply_norm_with_fp32_residual( self.input_norm, residual, target_dtype ) hidden_states = self.self_attn(hidden_states, positions, forward_batch) return hidden_states, residual, prev_router_hidden_states class ZayaDecoderMLPLayer(nn.Module): """MoE decoder layer: ``res_scale → input_norm → ZayaBlock``.""" def __init__( self, config: ZayaConfig, layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config self.layer_id = layer_id self.zaya_block = ZayaBlock( config=config, layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("zaya_block", prefix), ) self.input_norm = ZayaDecoderATTLayer._build_norm(config) if config.scale_residual_merge: self.res_scale = ResidualScaling(config, layer_id) else: self.res_scale = None def forward( self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor], positions: torch.Tensor, forward_batch: ForwardBatch, prev_router_hidden_states: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: target_dtype = ( self.input_norm.weight.dtype if isinstance(self.input_norm, RMSNorm) else hidden_states.dtype ) if self.res_scale is not None: residual, hidden_states = self.res_scale(residual, hidden_states) if residual is not None: residual = residual.float() + hidden_states.float() else: residual = hidden_states.float() hidden_states = _apply_norm_with_fp32_residual( self.input_norm, residual, target_dtype ) hidden_states, prev_router_hidden_states = self.zaya_block( hidden_states, prev_router_hidden_states ) return hidden_states, residual, prev_router_hidden_states # --------------------------------------------------------------------------- # Top-level model # --------------------------------------------------------------------------- def _build_layer( layer_id: int, config: ZayaConfig, quant_config: Optional[QuantizationConfig], prefix: str, ) -> nn.Module: # Even layer ids are attention, odd layer ids are MoE. This matches the HF # checkpoint keys: ``model.layers.<2k>.self_attn.*`` (CCA) versus # ``model.layers.<2k+1>.zaya_block.*`` (MoE). if layer_id % 2 == 0: return ZayaDecoderATTLayer( config=config, layer_id=layer_id, quant_config=quant_config, prefix=prefix, ) return ZayaDecoderMLPLayer( config=config, layer_id=layer_id, quant_config=quant_config, prefix=prefix, ) class ZayaModel(nn.Module): def __init__( self, config: ZayaConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.pp_group = get_pp_group() if self.pp_group.is_first_rank: self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, org_num_embeddings=config.vocab_size, prefix=add_prefix("embed_tokens", prefix), ) else: self.embed_tokens = PPMissingLayer() self.layers, self.start_layer, self.end_layer = make_layers( config.num_hidden_layers, lambda idx, prefix: _build_layer( layer_id=idx, config=config, quant_config=quant_config, prefix=prefix, ), pp_rank=self.pp_group.rank_in_group, pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: self.final_norm = ZayaDecoderATTLayer._build_norm(config) if config.scale_residual_merge: self.res_scale = ResidualScaling(config, config.num_hidden_layers) else: self.res_scale = None else: self.final_norm = PPMissingLayer() self.res_scale = None def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, inputs_embeds: Optional[torch.Tensor] = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: if self.pp_group.is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds else: hidden_states = self.embed_tokens(input_ids) residual = None else: assert pp_proxy_tensors is not None hidden_states = pp_proxy_tensors["hidden_states"] residual = pp_proxy_tensors["residual"] prev_router_hidden_states: Optional[torch.Tensor] = None for i in range(self.start_layer, self.end_layer): layer = self.layers[i] hidden_states, residual, prev_router_hidden_states = layer( hidden_states=hidden_states, residual=residual, positions=positions, forward_batch=forward_batch, prev_router_hidden_states=prev_router_hidden_states, ) if not self.pp_group.is_last_rank: return PPProxyTensors( { "hidden_states": hidden_states, "residual": residual, } ) if self.res_scale is not None: residual, hidden_states = self.res_scale(residual, hidden_states) target_dtype = ( self.final_norm.weight.dtype if isinstance(self.final_norm, RMSNorm) else hidden_states.dtype ) if residual is not None: merged = hidden_states.float() + residual.float() else: merged = hidden_states.float() hidden_states = _apply_norm_with_fp32_residual( self.final_norm, merged, target_dtype ) return hidden_states class ZayaForCausalLM(nn.Module): def __init__( self, config: ZayaConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config self.quant_config = quant_config self.pp_group = get_pp_group() self.model = ZayaModel( config=config, quant_config=quant_config, prefix=add_prefix("model", prefix), ) if self.pp_group.is_last_rank: self.lm_head = ParallelLMHead( config.vocab_size, config.hidden_size, org_num_embeddings=config.vocab_size, bias=bool(getattr(config, "lm_head_bias", False)), quant_config=None, prefix=add_prefix("lm_head", prefix), ) if config.tie_word_embeddings: self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config) @torch.no_grad() def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, inputs_embeds: Optional[torch.Tensor] = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: hidden_states = self.model( input_ids=input_ids, positions=positions, forward_batch=forward_batch, inputs_embeds=inputs_embeds, pp_proxy_tensors=pp_proxy_tensors, ) if not self.pp_group.is_last_rank: return hidden_states return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch ) # ---------------- weight loading ---------------- _EXPERT_RE = re.compile( r"^(.*\.zaya_block\.experts)\.local_experts\.(\d+)\.(linear_fc1|linear_fc2)\.weight$" ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load an HF ZAYA1 safetensors checkpoint into the SGLang module tree. Most keys map 1:1 because the module names already mirror the HF checkpoint layout. Two cases need rewriting: 1. ``self_attn.qkv.{linear_q, linear_k, conv_qk.{0,1}, val_proj{1,2}, temp}`` lands directly on the registered ``CCA`` submodule (which is named ``qkv`` exactly to keep this mapping trivial). 2. ``zaya_block.experts.local_experts..linear_fc1.weight`` (gate and up projections concatenated along dim 0) is split and routed to FusedMoE shards ``w1`` (first half) and ``w3`` (second half); ``linear_fc2.weight`` becomes the FusedMoE ``w2`` shard. """ params_dict = dict(self.named_parameters()) buffers_dict = dict(self.named_buffers()) # ``balancing_biases`` is a persistent buffer; FusedMoE may also expose # buffers. Expose them all through ``params_dict`` so that the regular # ``default_weight_loader`` can write to them. for key, buf in buffers_dict.items(): params_dict.setdefault(key, buf) fused_moe_modules: dict[str, nn.Module] = {} for name, module in self.named_modules(): if module.__class__.__name__ == "FusedMoE" or hasattr(module, "w13_weight"): fused_moe_modules[name] = module loaded_params: set[str] = set() for ckpt_name, loaded_weight in weights: # Skip keys that have no runtime counterpart in this model. if ckpt_name.startswith("lm_head") and self.config.tie_word_embeddings: continue if "rotary_emb" in ckpt_name: continue match = self._EXPERT_RE.match(ckpt_name) if match is not None: experts_prefix = match.group( 1 ) # e.g. model.layers.1.zaya_block.experts expert_id = int(match.group(2)) kind = match.group(3) moe_module = fused_moe_modules.get(experts_prefix) if moe_module is None: logger.warning( "FusedMoE module %s not found; skipping %s", experts_prefix, ckpt_name, ) continue weight_loader = moe_module.weight_loader if kind == "linear_fc1": param_name = f"{experts_prefix}.w13_weight" param = params_dict.get(param_name) if param is None: logger.warning("No param %s for %s", param_name, ckpt_name) continue half = loaded_weight.shape[0] // 2 weight_loader( param, loaded_weight[:half], ckpt_name, shard_id="w1", expert_id=expert_id, ) weight_loader( param, loaded_weight[half:], ckpt_name, shard_id="w3", expert_id=expert_id, ) loaded_params.add(param_name) else: # linear_fc2 param_name = f"{experts_prefix}.w2_weight" param = params_dict.get(param_name) if param is None: logger.warning("No param %s for %s", param_name, ckpt_name) continue weight_loader( param, loaded_weight, ckpt_name, shard_id="w2", expert_id=expert_id, ) loaded_params.add(param_name) continue # HF stores CCA tensors under ``self_attn.qkv.*``, which already # matches our submodule registration, so no rename is needed. if ckpt_name not in params_dict: # ``conv_qk`` is an ``nn.Sequential`` of two ``nn.Conv1d``, # whose keys end in ``.0.{weight,bias}`` / ``.1.{weight,bias}`` # and are exposed through ``named_parameters()`` automatically. # Anything else is genuinely unknown – warn and skip. logger.warning( "WARNING: checkpoint key %s has no matching parameter; skipping", ckpt_name, ) continue param = params_dict[ckpt_name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(ckpt_name) return loaded_params EntryClass = ZayaForCausalLM