minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support for Zyphra zaya1 model (#26347)

Browse Source 
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
This commit is contained in:
ChengYao-amd
2026-06-10 02:44:47 -07:00
committed by GitHub
co-authored by github-actions[bot]
parent b40f365732
commit 255843d454
9 changed files with 3167 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/__init__.py
+2
View File
@@ -39,6 +39,7 @@ from sglang.srt.configs.step3_vl import (
)
from sglang.srt.configs.step3p5 import Step3p5Config
from sglang.srt.configs.step3p7 import Step3p7Config
from sglang.srt.configs.zaya import ZayaConfig
__all__ = [
"AfmoeConfig",
@@ -80,4 +81,5 @@ __all__ = [
"Step3p5Config",
"Step3p7Config",
"Qwen3ASRConfig",
"ZayaConfig",
]
python/sglang/srt/configs/zaya.py
+326
View File
@@ -0,0 +1,326 @@
# 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.
# ==============================================================================
"""Configuration class for Zyphra ZAYA1 series models."""
from __future__ import annotations
from typing import TYPE_CHECKING, List, Optional
from transformers.configuration_utils import PretrainedConfig
if TYPE_CHECKING:
from sglang.srt.configs.mamba_utils import Mamba2CacheParams
class ZayaConfig(PretrainedConfig):
"""HuggingFace configuration for ZAYA1 hybrid (CCA attention + MoE) models.
Mirrors the field set used by Zyphra/ZAYA1-base/config.json. Most fields
are surfaced as constructor arguments so the same class can be instantiated
either from a published checkpoint via ``AutoConfig.from_pretrained`` or
programmatically in unit tests.
"""
model_type = "zaya"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
cca: bool = True,
num_query_groups: int = 2,
use_cache: bool = True,
attention_bias: bool = False,
lm_head_bias: bool = False,
vocab_size: int = 262272,
hidden_size: int = 2048,
ffn_hidden_size: int = 4096,
num_hidden_layers: int = 80,
num_experts: int = 16,
num_attention_heads: int = 8,
head_dim: int = 128,
activation_func: str = "swiglu",
max_position_embeddings: int = 32768,
norm_epsilon: float = 1e-5,
pad_token_id: int = 0,
bos_token_id: int = 2,
eos_token_id: int = 1,
tie_word_embeddings: bool = True,
rope_theta: float = 1_000_000.0,
attention_dropout: float = 0.0,
moe_router_topk: int = 1,
normalization: str = "RMSNorm",
zaya_mlp_expansion=256,
zaya_use_mod: bool = True,
zaya_high_prec: bool = True,
zaya_use_eda: bool = True,
add_bias_linear: bool = False,
gated_linear_unit: bool = True,
scale_residual_merge: bool = True,
fused_add_norm: bool = False,
residual_in_fp32: bool = True,
apply_rope_fusion: bool = True,
bias_activation_fusion: bool = True,
activation_func_fp8_input_store: bool = False,
sliding_window=None,
rope_scaling=None,
rope_parameters=None,
partial_rotary_factor: float = 0.5,
num_key_value_heads: int = 2,
clamp_temp: bool = False,
cca_time0: int = 2,
cca_time1: int = 2,
swa_layers=None,
swa_rotary_base=None,
zaya_layers=None,
cca_num_q_heads=None,
num_query_groups_list=None,
ffn_hidden_size_list=None,
kv_channels=None,
_attn_implementation: str = "eager",
**kwargs,
):
self.cca = cca
self.use_cache = use_cache
self.attention_bias = attention_bias
self.lm_head_bias = lm_head_bias
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.num_experts = num_experts
# ZAYA1-base ships a ``zaya_layers`` list whose entries are either the
# literal string ``"a"`` (attention layer) or an integer (number of
# experts in a MoE layer). When present it is the source of truth for
# both the total layer count and the per-layer placement. The HF
# config also carries a scalar ``num_hidden_layers`` that can disagree
# with ``len(zaya_layers)`` for historical reasons, so the list takes
# precedence whenever it is non-empty.
self.zaya_layers = list(zaya_layers) if zaya_layers else None
if self.zaya_layers:
self.num_hidden_layers = len(self.zaya_layers)
else:
self.num_hidden_layers = num_hidden_layers
# When the per-layer lists are present, derive each active scalar
# field from the first non-zero entry of the corresponding list.
# This matches ZAYA1-base in practice: every attention layer shares
# the same ``cca_num_q_heads`` (e.g. 8) and ``num_query_groups``
# (e.g. 2), and every MoE layer shares the same ``ffn_hidden_size``
# (e.g. 4096) and ``zaya_mlp_expansion`` (e.g. 256). When no list is
# provided, the constructor argument is used unchanged.
self.cca_num_q_heads_list = list(cca_num_q_heads) if cca_num_q_heads else None
self.num_query_groups_list = (
list(num_query_groups_list) if num_query_groups_list else None
)
self.ffn_hidden_size_list = (
list(ffn_hidden_size_list) if ffn_hidden_size_list else None
)
if isinstance(zaya_mlp_expansion, (list, tuple)):
self.zaya_mlp_expansion_list = list(zaya_mlp_expansion)
zaya_mlp_expansion_scalar = next(
(v for v in self.zaya_mlp_expansion_list if v), 256
)
else:
self.zaya_mlp_expansion_list = None
zaya_mlp_expansion_scalar = int(zaya_mlp_expansion)
if self.cca_num_q_heads_list:
self.num_attention_heads = next(
(v for v in self.cca_num_q_heads_list if v), num_attention_heads
)
else:
self.num_attention_heads = num_attention_heads
if self.num_query_groups_list:
self.num_query_groups = next(
(v for v in self.num_query_groups_list if v), num_query_groups
)
else:
self.num_query_groups = num_query_groups
if self.ffn_hidden_size_list:
self.ffn_hidden_size = next(
(v for v in self.ffn_hidden_size_list if v), ffn_hidden_size
)
else:
self.ffn_hidden_size = ffn_hidden_size
self.zaya_mlp_expansion = zaya_mlp_expansion_scalar
# The HF config exposes the per-head dim as ``kv_channels``; accept
# either spelling and keep both attributes in sync for downstream code.
if head_dim is None and kv_channels is not None:
head_dim = int(kv_channels)
self.head_dim = head_dim
self.kv_channels = kv_channels if kv_channels is not None else head_dim
assert self.head_dim is not None, "head_dim is required for ZayaConfig"
assert (
self.num_query_groups == num_key_value_heads
), "num_query_groups must equal num_key_value_heads for ZAYA1 checkpoints"
self.num_key_value_heads = num_key_value_heads
self.activation_func = activation_func
self.max_position_embeddings = max_position_embeddings
self.norm_epsilon = norm_epsilon
self.normalization = normalization
self.pad_token_id = pad_token_id
self.bos_token_id = bos_token_id
self.eos_token_id = eos_token_id
self.tie_word_embeddings = tie_word_embeddings
self.attention_dropout = attention_dropout
self.moe_router_topk = moe_router_topk
self.zaya_use_mod = zaya_use_mod
self.zaya_high_prec = zaya_high_prec
self.zaya_use_eda = zaya_use_eda
self.add_bias_linear = add_bias_linear
self.gated_linear_unit = gated_linear_unit
self.scale_residual_merge = scale_residual_merge
self.fused_add_norm = fused_add_norm
self.residual_in_fp32 = residual_in_fp32
self.apply_rope_fusion = apply_rope_fusion
self.bias_activation_fusion = bias_activation_fusion
self.activation_func_fp8_input_store = activation_func_fp8_input_store
self.sliding_window = sliding_window
self.partial_rotary_factor = partial_rotary_factor
self.rope_theta = rope_theta
self.rope_scaling = rope_scaling
if isinstance(rope_parameters, dict):
rope_parameters_dict = dict(rope_parameters)
elif isinstance(rope_scaling, dict):
rope_parameters_dict = dict(rope_scaling)
else:
rope_parameters_dict = {"rope_type": "default"}
if "type" in rope_parameters_dict:
rope_parameters_dict.setdefault(
"rope_type", rope_parameters_dict.pop("type")
)
rope_parameters_dict.setdefault("rope_theta", rope_theta)
rope_parameters_dict.setdefault("partial_rotary_factor", partial_rotary_factor)
self.rope_parameters = rope_parameters_dict
self.clamp_temp = clamp_temp
self.cca_time0 = cca_time0
self.cca_time1 = cca_time1
self.swa_layers = swa_layers
self.swa_rotary_base = swa_rotary_base
self._attn_implementation = _attn_implementation
super().__init__(
pad_token_id=pad_token_id,
bos_token_id=bos_token_id,
eos_token_id=eos_token_id,
tie_word_embeddings=self.tie_word_embeddings,
**kwargs,
)
# -- Hybrid model interface (HybridReqToTokenPool / MambaPool) ----------
@property
def full_attention_layer_ids(self) -> List[int]:
if self.zaya_layers:
return [i for i, lt in enumerate(self.zaya_layers) if lt == "a"]
return [i for i in range(self.num_hidden_layers) if i % 2 == 0]
@property
def linear_layer_ids(self) -> List[int]:
return self.full_attention_layer_ids
@property
def mamba_chunk_size(self) -> int:
return 1
@property
def mamba2_cache_params(self) -> Optional[Mamba2CacheParams]:
from sglang.srt.configs.mamba_utils import (
Mamba2CacheParams,
Mamba2StateShape,
mamba2_state_dtype,
)
attn_layer_ids = self.linear_layer_ids
if not attn_layer_ids:
return None
# ``conv[0]`` (conv_qk left padding) is sized per TP rank because CCA
# is head-parallel. ``conv[1]`` (prev_hs) carries the full hidden_state
# and feeds the replicated val_proj1 / val_proj2, so it stays at full
# ``hidden_size`` on every rank.
#
# Use the *global* TP world size -- the same accessor that
# ``ZayaAttention`` / ``CCA`` use to split heads and over which
# ``o_proj`` all-reduces -- so the cache shape and the per-rank
# ``in_out_ch`` stay in lockstep. ZAYA1 asserts the attention-TP group
# equals the global TP group (DP attention is unsupported), so the two
# are always identical in practice.
try:
from sglang.srt.distributed import (
get_tensor_model_parallel_world_size,
)
tp_size = get_tensor_model_parallel_world_size()
except (AssertionError, RuntimeError):
tp_size = 1
in_out_ch_full = (
self.num_attention_heads + self.num_key_value_heads
) * self.head_dim
assert in_out_ch_full % tp_size == 0, (
f"CCA channels ({in_out_ch_full}) must be divisible by TP size "
f"({tp_size}); both num_attention_heads and num_query_groups must "
"be divisible by tp_size for ZAYA1 head-parallel attention."
)
in_out_ch_per_rank = in_out_ch_full // tp_size
total_padding = (self.cca_time0 - 1) + (self.cca_time1 - 1)
shape = Mamba2StateShape(
conv=[
(in_out_ch_per_rank, total_padding),
(self.hidden_size, 1),
],
temporal=(1, 1, 0),
intermediate_size=in_out_ch_per_rank,
conv_dim=in_out_ch_per_rank,
ssm_state_size=0,
num_heads=1,
head_dim=1,
state_size=0,
conv_kernel=total_padding + 1,
)
return Mamba2CacheParams(
shape=shape,
layers=attn_layer_ids,
dtype=mamba2_state_dtype(self),
)
def register_zaya_config() -> None:
"""Register :class:`ZayaConfig` with HuggingFace ``AutoConfig``.
Safe to call multiple times. ``AutoConfig.register`` raises ``ValueError``
on duplicate registration, which is suppressed so importing this module
stays idempotent.
"""
try:
from transformers import AutoConfig
AutoConfig.register(ZayaConfig.model_type, ZayaConfig)
except (ValueError, ImportError):
# Either the installed ``transformers`` does not expose
# ``AutoConfig.register``, or the "zaya" model type is already
# registered – nothing to do in either case.
pass
register_zaya_config()
python/sglang/srt/model_executor/model_runner.py
+3 -1
View File
@@ -52,6 +52,7 @@ from sglang.srt.configs import (
Qwen3_5Config,
Qwen3_5MoeConfig,
Qwen3NextConfig,
ZayaConfig,
)
from sglang.srt.configs.device_config import DeviceConfig
from sglang.srt.configs.linear_attn_model_registry import get_linear_attn_config
@@ -2181,7 +2182,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
| NemotronHConfig
| Lfm2Config
| Lfm2MoeConfig
| Lfm2VlConfig,
| Lfm2VlConfig
| ZayaConfig,
):
return config
if isinstance(config, NemotronH_Nano_VL_V2_Config):
python/sglang/srt/models/zaya.py
+1658
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/server_args.py
+7
View File
@@ -2745,6 +2745,13 @@ class ServerArgs:
"as the first layer might not be an attention layer"
)
elif model_arch in ["ZayaForCausalLM"]:
self._handle_mamba_radix_cache(
model_arch=model_arch,
support_mamba_cache=True,
support_mamba_cache_extra_buffer=False,
)
if (
model_arch in ["Qwen3VLForConditionalGeneration"]
and is_hip()
test/registered/models/test_zaya.py
+127
View File
@@ -0,0 +1,127 @@
"""End-to-end server test for Zyphra ZAYA1 (hybrid CCA attention + MoE).
This test boots a real ``Zyphra/ZAYA1-base`` SGLang server via
``popen_launch_server``, sends a handful of completions through the HTTP API,
and finishes with a small MMLU sanity slice.
The test is gated behind ``RUN_ZAYA_E2E=1`` so the registered suite does not
have to download the full ZAYA1-base checkpoint (≈17 GB) on every run; the CI
job that owns this test sets the variable explicitly.
"""
import os
import unittest
from types import SimpleNamespace
from sglang.srt.utils import is_hip, kill_process_tree
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
# ZAYA1-base is a heavyweight launch (≈120 transformer layers with MoE), so
# the estimated time is set generously to keep the CI scheduler from preempting
# the job before the server finishes warming up.
register_cuda_ci(est_time=420, stage="extra-a", runner_config="1-gpu-large")
register_amd_ci(est_time=420, suite="stage-b-test-1-gpu-large-amd")
_MODEL_PATH = os.environ.get("ZAYA_MODEL_PATH", "Zyphra/ZAYA1-base")
def _zaya_enabled() -> bool:
return os.environ.get("RUN_ZAYA_E2E", "0") == "1"
@unittest.skipUnless(
_zaya_enabled(),
"Set RUN_ZAYA_E2E=1 to enable the ZAYA1 end-to-end server test "
"(requires downloading the model weights).",
)
class TestZayaServer(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = _MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--mem-fraction-static",
"0.5",
"--max-running-requests",
"8",
]
if is_hip():
other_args += ["--attention-backend", "triton"]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
if getattr(cls, "process", None) is not None:
kill_process_tree(cls.process.pid)
def test_generation_basic(self):
"""Send three prompts through the ``/generate`` endpoint and require
non-empty completions for each."""
import requests
prompts = [
"The capital of France is",
"1 + 2 + 3 + 4 + 5 =",
"Write a haiku about silicon:",
]
for prompt in prompts:
resp = requests.post(
f"{self.base_url}/generate",
json={
"text": prompt,
"sampling_params": {
"temperature": 0.0,
"max_new_tokens": 16,
},
},
timeout=60,
)
self.assertEqual(resp.status_code, 200, resp.text)
data = resp.json()
self.assertIn("text", data, data)
self.assertGreater(len(data["text"].strip()), 0, data)
def test_mmlu_sanity(self):
"""32-example MMLU sanity slice.
ZAYA1-base is a pretrained (non instruction-tuned) checkpoint that
emits long ``<think>…</think>`` reasoning blocks before settling on a
final letter, so ``max_tokens`` must be large enough for the evaluator
to see the chosen answer. The threshold sits just above chance: it is
a regression sanity check rather than a production-quality gate. An
instruction-tuned ZAYA1 checkpoint scores meaningfully higher and
should raise this bound when wired in.
"""
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="mmlu",
num_examples=32,
num_threads=8,
max_tokens=1024,
)
metrics = run_eval(args)
self.assertGreaterEqual(
metrics["score"],
0.30,
f"MMLU sanity below threshold: {metrics}",
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/configs/test_zaya_config.py
+101
View File
@@ -0,0 +1,101 @@
"""Unit tests for ``sglang.srt.configs.zaya.ZayaConfig``."""
import unittest
from transformers import AutoConfig
from sglang.srt.configs.zaya import ZayaConfig, register_zaya_config
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
class TestZayaConfig(CustomTestCase):
def test_default_fields_match_zaya1_base(self):
"""Defaults reflect Zyphra/ZAYA1-base reference checkpoint."""
cfg = ZayaConfig()
self.assertEqual(cfg.model_type, "zaya")
self.assertEqual(cfg.hidden_size, 2048)
self.assertEqual(cfg.head_dim, 128)
self.assertEqual(cfg.num_attention_heads, 8)
self.assertEqual(cfg.num_query_groups, 2)
self.assertEqual(cfg.num_key_value_heads, 2)
self.assertEqual(cfg.num_experts, 16)
self.assertEqual(cfg.moe_router_topk, 1)
self.assertEqual(cfg.ffn_hidden_size, 4096)
self.assertEqual(cfg.zaya_mlp_expansion, 256)
self.assertEqual(cfg.cca_time0, 2)
self.assertEqual(cfg.cca_time1, 2)
self.assertTrue(cfg.tie_word_embeddings)
self.assertTrue(cfg.zaya_use_eda)
self.assertTrue(cfg.zaya_use_mod)
self.assertTrue(cfg.scale_residual_merge)
self.assertEqual(cfg.partial_rotary_factor, 0.5)
self.assertEqual(cfg.rope_theta, 1_000_000.0)
def test_rope_parameters_auto_derived(self):
"""When neither ``rope_scaling`` nor ``rope_parameters`` is supplied,
both ``rope_theta`` and ``partial_rotary_factor`` should still appear
inside ``rope_parameters`` together with a default ``rope_type``.
"""
cfg = ZayaConfig()
rp = cfg.rope_parameters
self.assertEqual(rp["rope_type"], "default")
self.assertEqual(rp["rope_theta"], 1_000_000.0)
self.assertEqual(rp["partial_rotary_factor"], 0.5)
def test_rope_parameters_explicit_takes_priority(self):
cfg = ZayaConfig(rope_parameters={"type": "linear", "factor": 4.0})
rp = cfg.rope_parameters
# ``type`` is normalized to ``rope_type``.
self.assertEqual(rp["rope_type"], "linear")
self.assertEqual(rp["factor"], 4.0)
# Defaults are still merged in.
self.assertEqual(rp["rope_theta"], 1_000_000.0)
def test_head_dim_required(self):
with self.assertRaises(AssertionError):
ZayaConfig(head_dim=None)
def test_num_query_groups_must_equal_kv_heads(self):
with self.assertRaises(AssertionError):
ZayaConfig(num_query_groups=4, num_key_value_heads=2)
def test_hybrid_model_properties(self):
"""Verify properties required for HybridReqToTokenPool integration."""
cfg = ZayaConfig()
# Default 80 layers: even layers are attention, odd are MoE
self.assertEqual(cfg.full_attention_layer_ids, list(range(0, 80, 2)))
self.assertEqual(cfg.linear_layer_ids, cfg.full_attention_layer_ids)
self.assertEqual(cfg.mamba_chunk_size, 1)
params = cfg.mamba2_cache_params
self.assertIsNotNone(params)
# conv[0] = conv_state: (in_out_ch, total_padding)
in_out_ch = (cfg.num_attention_heads + cfg.num_key_value_heads) * cfg.head_dim
total_padding = (cfg.cca_time0 - 1) + (cfg.cca_time1 - 1)
self.assertEqual(params.shape.conv[0], (in_out_ch, total_padding))
# conv[1] = prev_hs: (hidden_size, 1)
self.assertEqual(params.shape.conv[1], (cfg.hidden_size, 1))
self.assertEqual(params.layers, cfg.linear_layer_ids)
def test_hybrid_model_properties_with_zaya_layers(self):
"""When zaya_layers is provided, layer IDs derive from the list."""
cfg = ZayaConfig(zaya_layers=["a", 16, "a", 16])
self.assertEqual(cfg.num_hidden_layers, 4)
self.assertEqual(cfg.full_attention_layer_ids, [0, 2])
self.assertEqual(cfg.linear_layer_ids, [0, 2])
def test_auto_config_registration_is_idempotent(self):
# Calling the helper twice must not raise even though importing the
# module already registered the model type.
register_zaya_config()
register_zaya_config()
# ``AutoConfig.for_model`` now resolves to ``ZayaConfig``.
cfg = AutoConfig.for_model("zaya")
self.assertIsInstance(cfg, ZayaConfig)
if __name__ == "__main__":
unittest.main()
test/registered/unit/models/test_zaya_cca.py
+765
View File
@@ -0,0 +1,765 @@
"""Numerical and state-cache correctness tests for the ZAYA1 CCA module.
The CCA per-request conv-state cache must satisfy the following invariants,
which are each exercised by a dedicated test case:
1. A single-chunk extend forward (no prefix) is numerically equivalent to the
reference torch implementation that processes the whole sequence at once.
2. Splitting a sequence into one prefill of ``S0`` tokens and ``S1`` single-
token decode steps produces the same q / k / v tensors as the equivalent
single-chunk run.
3. A batched two-request decode for request 0 yields identical q / k / v to a
single-request decode of request 0 at the same step.
4. Multi-request prefills update only the conv state and ``prev_hs`` slots for
each request and leave unused slots zero.
5. A simulated tensor-parallel (TP=2) CCA produces per-rank q / k / v slices
that match the corresponding head slices of a TP=1 reference, both for
prefill (``_forward_extend``) and for decode (``_forward_decode``).
All tests run on CPU with a tiny configuration so they stay fast and have no
GPU dependency. State is stored in a mock centralized pool that mirrors the
``HybridReqToTokenPool`` / ``MambaPool`` interface used at serving time.
"""
import os
import unittest
from contextlib import contextmanager
from dataclasses import dataclass
from types import SimpleNamespace
from typing import List, Optional
import torch
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=30, suite="base-a-test-cpu")
def _ensure_dist_initialized() -> None:
"""Set up a minimal single-rank gloo distributed environment plus the
SGLang model-parallel groups (TP=1, PP=1, EP=1). The CCA module reads
``get_tensor_model_parallel_rank()`` / ``get_tensor_model_parallel_world_size()``
inside ``__init__`` to size its head-parallel projections, so the world
group and model parallel groups must both be initialized before any CCA
construction.
"""
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", "29632")
os.environ.setdefault("RANK", "0")
os.environ.setdefault("WORLD_SIZE", "1")
os.environ.setdefault("LOCAL_RANK", "0")
from sglang.srt.distributed.parallel_state import (
init_distributed_environment,
initialize_model_parallel,
model_parallel_is_initialized,
)
if not torch.distributed.is_initialized():
init_distributed_environment(
world_size=1,
rank=0,
local_rank=0,
backend="gloo",
)
if not model_parallel_is_initialized():
# Pass arguments as kwargs because ``ensure_model_parallel_initialized``
# forwards positional ``backend`` into the ``attention_data_parallel_size``
# slot of ``initialize_model_parallel``, which then explodes on
# ``int // str``. Using kwargs avoids that footgun.
initialize_model_parallel(
tensor_model_parallel_size=1,
expert_model_parallel_size=1,
pipeline_model_parallel_size=1,
backend="gloo",
)
# ---------------------------------------------------------------------------
# Mock centralized pool
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class _MockLayerCache:
conv: List[torch.Tensor]
temporal: torch.Tensor
class _MockReqToTokenPool:
"""Minimal stand-in for ``HybridReqToTokenPool`` providing the two methods
that CCA calls: ``mamba2_layer_cache`` and ``get_mamba_indices``.
For TP-aware tests, ``tp_size`` controls the per-rank ``in_out_ch`` of the
``conv[0]`` state. ``conv[1]`` (prev_hs) is replicated and stays at full
``hidden_size``.
"""
def __init__(self, pool_size: int, cca_config, tp_size: int = 1):
in_out_ch_full = (
cca_config.num_attention_heads + cca_config.num_key_value_heads
) * cca_config.head_dim
assert in_out_ch_full % tp_size == 0
in_out_ch_per_rank = in_out_ch_full // tp_size
total_padding = (cca_config.cca_time0 - 1) + (cca_config.cca_time1 - 1)
num_layers = len(cca_config.linear_layer_ids)
self.conv_state = torch.zeros(
num_layers, pool_size + 1, in_out_ch_per_rank, total_padding
)
self.prev_hs_state = torch.zeros(
num_layers, pool_size + 1, cca_config.hidden_size, 1
)
self.temporal = torch.zeros(num_layers, pool_size + 1, 1, 1, 0)
self._layer_map = {lid: i for i, lid in enumerate(cca_config.linear_layer_ids)}
self._identity_map = torch.arange(pool_size + 1, dtype=torch.int32)
def mamba2_layer_cache(self, layer_id: int):
idx = self._layer_map[layer_id]
return _MockLayerCache(
conv=[self.conv_state[idx], self.prev_hs_state[idx]],
temporal=self.temporal[idx],
)
def get_mamba_indices(self, req_pool_indices: torch.Tensor) -> torch.Tensor:
return req_pool_indices.to(torch.int32)
@contextmanager
def _mock_pool_context(pool: _MockReqToTokenPool):
"""Install a mock ``ForwardContext`` whose ``req_to_token_pool`` is ``pool``."""
from sglang.srt.model_executor.forward_context import (
ForwardContext,
set_forward_context,
)
backend = SimpleNamespace(req_to_token_pool=pool, token_to_kv_pool=None)
ctx = ForwardContext(attn_backend=backend)
prev = set_forward_context(ctx)
try:
yield pool
finally:
set_forward_context(prev)
# ---------------------------------------------------------------------------
# Helper factories
# ---------------------------------------------------------------------------
def _make_forward_batch(
*,
is_decode: bool,
extend_seq_lens_cpu,
extend_prefix_lens_cpu,
req_pool_indices,
input_ids: torch.Tensor,
):
from sglang.srt.model_executor.forward_batch_info import ForwardMode
mode = ForwardMode.DECODE if is_decode else ForwardMode.EXTEND
forward_batch = SimpleNamespace()
forward_batch.forward_mode = mode
forward_batch.input_ids = input_ids
forward_batch.req_pool_indices = torch.as_tensor(
req_pool_indices, dtype=torch.int32
)
forward_batch.extend_seq_lens_cpu = list(extend_seq_lens_cpu)
forward_batch.extend_prefix_lens_cpu = list(extend_prefix_lens_cpu)
return forward_batch
def _make_tiny_config(num_hidden_layers: int = 2):
from sglang.srt.configs.zaya import ZayaConfig
return ZayaConfig(
hidden_size=16,
ffn_hidden_size=32,
num_hidden_layers=num_hidden_layers,
num_experts=2,
num_attention_heads=4,
num_query_groups=2,
num_key_value_heads=2,
head_dim=8,
cca_time0=2,
cca_time1=2,
max_position_embeddings=64,
moe_router_topk=1,
zaya_mlp_expansion=8,
attention_bias=False,
)
def _make_tiny_cca(
seed: int = 0,
tp_rank: Optional[int] = None,
tp_size: Optional[int] = None,
layer_id: int = 0,
config=None,
):
from sglang.srt.models.zaya import CCA
if config is None:
config = _make_tiny_config()
torch.manual_seed(seed)
cca = CCA(
config=config,
cca_num_k_heads=config.num_query_groups,
cca_num_q_heads=config.num_attention_heads,
hidden_size=config.hidden_size,
head_dim=config.head_dim,
cca_time0=config.cca_time0,
cca_time1=config.cca_time1,
layer_id=layer_id,
tp_rank=tp_rank,
tp_size=tp_size,
)
cca.eval()
with torch.no_grad():
for p in cca.parameters():
p.data.normal_(mean=0.0, std=0.05)
cca.temp.data.zero_()
return cca, config
class TestZayaCCA(CustomTestCase):
@classmethod
def setUpClass(cls) -> None:
_ensure_dist_initialized()
def test_single_chunk_matches_reference(self):
"""A single-chunk extend with empty prefix matches the no-state path."""
cca, config = _make_tiny_cca(seed=1)
cca_ref, _ = _make_tiny_cca(seed=1)
with torch.no_grad():
cca_ref.load_state_dict(cca.state_dict())
S = 5
hs = torch.randn(S, cca.hidden_size, dtype=torch.float32) * 0.1
q_ref, k_ref, v_ref = cca_ref._forward_no_state(hs)
pool = _MockReqToTokenPool(pool_size=8, cca_config=config)
fb = _make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[S],
extend_prefix_lens_cpu=[0],
req_pool_indices=[0],
input_ids=torch.arange(S, dtype=torch.int64),
)
with _mock_pool_context(pool):
q, k, v = cca.forward(hs, fb)
torch.testing.assert_close(q, q_ref, atol=1e-5, rtol=1e-5)
torch.testing.assert_close(k, k_ref, atol=1e-5, rtol=1e-5)
torch.testing.assert_close(v, v_ref, atol=1e-5, rtol=1e-5)
def test_prefill_then_decode_matches_full_sequence(self):
"""Prefill(S0) followed by ``S1`` single-token decode steps matches a
one-shot reference over ``S0 + S1`` tokens."""
cca, config = _make_tiny_cca(seed=2)
cca_ref, _ = _make_tiny_cca(seed=2)
with torch.no_grad():
cca_ref.load_state_dict(cca.state_dict())
S0, S1 = 4, 2
S_total = S0 + S1
torch.manual_seed(77)
hs = torch.randn(S_total, cca.hidden_size, dtype=torch.float32) * 0.1
q_ref, k_ref, v_ref = cca_ref._forward_no_state(hs)
pool = _MockReqToTokenPool(pool_size=8, cca_config=config)
with _mock_pool_context(pool):
fb_prefill = _make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[S0],
extend_prefix_lens_cpu=[0],
req_pool_indices=[0],
input_ids=torch.arange(S0, dtype=torch.int64),
)
q0, k0, v0 = cca.forward(hs[:S0], fb_prefill)
q_decodes = [q0]
k_decodes = [k0]
v_decodes = [v0]
for t in range(S1):
fb_decode = _make_forward_batch(
is_decode=True,
extend_seq_lens_cpu=[],
extend_prefix_lens_cpu=[],
req_pool_indices=[0],
input_ids=torch.tensor([0], dtype=torch.int64),
)
qd, kd, vd = cca.forward(hs[S0 + t : S0 + t + 1], fb_decode)
q_decodes.append(qd)
k_decodes.append(kd)
v_decodes.append(vd)
q_cat = torch.cat(q_decodes, dim=0)
k_cat = torch.cat(k_decodes, dim=0)
v_cat = torch.cat(v_decodes, dim=0)
torch.testing.assert_close(q_cat, q_ref, atol=1e-4, rtol=1e-4)
torch.testing.assert_close(k_cat, k_ref, atol=1e-4, rtol=1e-4)
torch.testing.assert_close(v_cat, v_ref, atol=1e-4, rtol=1e-4)
def test_batched_decode_matches_single_decode(self):
"""A two-request batched decode of request 0 must produce the same
q / k / v tensors as a single-request decode of request 0."""
cca_single, config = _make_tiny_cca(seed=11)
cca_batched, _ = _make_tiny_cca(seed=11)
with torch.no_grad():
cca_batched.load_state_dict(cca_single.state_dict())
S0 = 4
torch.manual_seed(202)
hs0 = torch.randn(S0, cca_single.hidden_size, dtype=torch.float32) * 0.1
hs1 = torch.randn(S0, cca_single.hidden_size, dtype=torch.float32) * 0.1
decode0 = torch.randn(cca_single.hidden_size, dtype=torch.float32) * 0.1
decode1 = torch.randn(cca_single.hidden_size, dtype=torch.float32) * 0.1
pool_single = _MockReqToTokenPool(pool_size=8, cca_config=config)
with _mock_pool_context(pool_single):
cca_single.forward(
hs0,
_make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[S0],
extend_prefix_lens_cpu=[0],
req_pool_indices=[0],
input_ids=torch.arange(S0, dtype=torch.int64),
),
)
q_solo, k_solo, v_solo = cca_single.forward(
decode0.unsqueeze(0),
_make_forward_batch(
is_decode=True,
extend_seq_lens_cpu=[],
extend_prefix_lens_cpu=[],
req_pool_indices=[0],
input_ids=torch.tensor([0], dtype=torch.int64),
),
)
pool_batched = _MockReqToTokenPool(pool_size=8, cca_config=config)
with _mock_pool_context(pool_batched):
cca_batched.forward(
torch.cat([hs0, hs1], dim=0),
_make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[S0, S0],
extend_prefix_lens_cpu=[0, 0],
req_pool_indices=[0, 1],
input_ids=torch.arange(2 * S0, dtype=torch.int64),
),
)
q_batch, k_batch, v_batch = cca_batched.forward(
torch.stack([decode0, decode1], dim=0),
_make_forward_batch(
is_decode=True,
extend_seq_lens_cpu=[],
extend_prefix_lens_cpu=[],
req_pool_indices=[0, 1],
input_ids=torch.tensor([0, 1], dtype=torch.int64),
),
)
torch.testing.assert_close(q_batch[0:1], q_solo, atol=1e-5, rtol=1e-5)
torch.testing.assert_close(k_batch[0:1], k_solo, atol=1e-5, rtol=1e-5)
torch.testing.assert_close(v_batch[0:1], v_solo, atol=1e-5, rtol=1e-5)
def test_two_requests_state_isolation(self):
"""A batched prefill of two requests must update only the requests'
own slots in the centralized pool."""
cca, config = _make_tiny_cca(seed=4)
S0, S1 = 3, 2
hs0 = torch.randn(S0, cca.hidden_size, dtype=torch.float32) * 0.1
hs1 = torch.randn(S1, cca.hidden_size, dtype=torch.float32) * 0.1
hs = torch.cat([hs0, hs1], dim=0)
pool = _MockReqToTokenPool(pool_size=8, cca_config=config)
fb = _make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[S0, S1],
extend_prefix_lens_cpu=[0, 0],
req_pool_indices=[2, 5],
input_ids=torch.arange(S0 + S1, dtype=torch.int64),
)
with _mock_pool_context(pool):
cca.forward(hs, fb)
layer_cache = pool.mamba2_layer_cache(0)
conv_state = layer_cache.conv[0]
prev_hs_state = layer_cache.conv[1]
self.assertTrue(torch.any(conv_state[2] != 0))
self.assertTrue(torch.any(conv_state[5] != 0))
torch.testing.assert_close(
prev_hs_state[2].squeeze(-1).to(torch.float32),
hs0[-1].to(torch.float32),
atol=1e-5,
rtol=1e-5,
)
torch.testing.assert_close(
prev_hs_state[5].squeeze(-1).to(torch.float32),
hs1[-1].to(torch.float32),
atol=1e-5,
rtol=1e-5,
)
for idx in (0, 1, 3, 4):
self.assertTrue(torch.all(conv_state[idx] == 0))
self.assertTrue(torch.all(prev_hs_state[idx] == 0))
def test_mamba_indices_resolved_once_per_forward_step(self):
"""The req -> MambaPool-slot mapping is identical for every CCA layer in
a step, so it (and its GPU->CPU ``.tolist()`` sync) must be resolved once
per forward step and shared across layers, not recomputed per layer.
Regression guard for the per-layer mamba-sync fix: two CCA layers driven
by a single ForwardBatch must trigger exactly one ``get_mamba_indices``
lookup and one host materialization for the whole step.
"""
class _CountingPool(_MockReqToTokenPool):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.get_mamba_indices_calls = 0
def get_mamba_indices(self, req_pool_indices):
self.get_mamba_indices_calls += 1
return super().get_mamba_indices(req_pool_indices)
# num_hidden_layers=4 -> CCA (even) layers live at ids 0 and 2.
config = _make_tiny_config(num_hidden_layers=4)
self.assertEqual(config.linear_layer_ids, [0, 2])
cca0, _ = _make_tiny_cca(seed=5, layer_id=0, config=config)
cca2, _ = _make_tiny_cca(seed=6, layer_id=2, config=config)
S = 4
hs = torch.randn(S, config.hidden_size, dtype=torch.float32) * 0.1
def _fresh_fb():
return _make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[S],
extend_prefix_lens_cpu=[0],
req_pool_indices=[0],
input_ids=torch.arange(S, dtype=torch.int64),
)
pool = _CountingPool(pool_size=8, cca_config=config)
with _mock_pool_context(pool):
fb = _fresh_fb()
cca0.forward(hs, fb)
cca2.forward(hs, fb)
# Two CCA layers, one forward step -> one shared lookup, both the
# device tensor and its host mirror memoized on the ForwardBatch.
self.assertEqual(pool.get_mamba_indices_calls, 1)
self.assertTrue(hasattr(fb, "_zaya_mamba_indices"))
self.assertTrue(hasattr(fb, "_zaya_mamba_indices_cpu"))
self.assertEqual(fb._zaya_mamba_indices_cpu, [0])
# A new forward step (fresh ForwardBatch) resolves the mapping again.
cca0.forward(hs, _fresh_fb())
self.assertEqual(pool.get_mamba_indices_calls, 2)
def test_decode_path_does_not_sync_indices_to_host(self):
"""The decode path indexes the pool entirely on-device, so it must not
populate the host-side index cache (keeping it CUDA-graph friendly)."""
cca, config = _make_tiny_cca(seed=7)
pool = _MockReqToTokenPool(pool_size=8, cca_config=config)
with _mock_pool_context(pool):
cca.forward(
torch.randn(3, config.hidden_size, dtype=torch.float32) * 0.1,
_make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[3],
extend_prefix_lens_cpu=[0],
req_pool_indices=[0],
input_ids=torch.arange(3, dtype=torch.int64),
),
)
fb_decode = _make_forward_batch(
is_decode=True,
extend_seq_lens_cpu=[],
extend_prefix_lens_cpu=[],
req_pool_indices=[0],
input_ids=torch.tensor([0], dtype=torch.int64),
)
cca.forward(
torch.randn(1, config.hidden_size, dtype=torch.float32) * 0.1,
fb_decode,
)
# Device indices are memoized, but the host ``.tolist()`` mirror is only
# built by the extend path.
self.assertTrue(hasattr(fb_decode, "_zaya_mamba_indices"))
self.assertFalse(hasattr(fb_decode, "_zaya_mamba_indices_cpu"))
class TestZayaCCATensorParallel(CustomTestCase):
"""Head-parallel TP equivalence:
For each TP rank, the CCA's q / k / v output must equal the head slice of
the TP=1 reference's output that corresponds to that rank's heads. This
verifies that the grouped-mean step and ``conv_qk.1`` (groups = num_q_heads
+ num_k_heads) are correctly partitioned across heads with no cross-rank
leakage.
"""
TP_SIZE = 2
@classmethod
def setUpClass(cls) -> None:
_ensure_dist_initialized()
def _slice_full_state_dict_into_rank(self, ref_cca, tp_cca, tp_rank: int):
"""Copy the reference's full weights into the per-rank CCA, using the
per-parameter ``weight_loader`` that the CCA installs on its own
parameters during ``__init__``. This mirrors what
``ZayaForCausalLM.load_weights`` does at serving time and is the
only way TP correctness is exercised end-to-end.
"""
ref_state = dict(ref_cca.state_dict())
from sglang.srt.model_loader.weight_utils import default_weight_loader
with torch.no_grad():
for name, param in tp_cca.named_parameters():
full_weight = ref_state[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, full_weight)
def _check_per_rank_outputs(
self,
full_q: torch.Tensor,
full_k: torch.Tensor,
full_v: torch.Tensor,
rank_q: torch.Tensor,
rank_k: torch.Tensor,
rank_v: torch.Tensor,
tp_rank: int,
cfg,
):
"""Compare a TP=2 rank's output against the corresponding head slice
of the TP=1 reference output. Q heads and K heads are partitioned
contiguously across ranks: rank ``r`` owns
``[r*Q_per_rank, (r+1)*Q_per_rank)`` for Q and similarly for K.
"""
q_heads_per_rank = cfg.num_attention_heads // self.TP_SIZE
k_heads_per_rank = cfg.num_query_groups // self.TP_SIZE
q_lo, q_hi = tp_rank * q_heads_per_rank, (tp_rank + 1) * q_heads_per_rank
k_lo, k_hi = tp_rank * k_heads_per_rank, (tp_rank + 1) * k_heads_per_rank
torch.testing.assert_close(
rank_q, full_q[:, q_lo:q_hi, :], atol=1e-5, rtol=1e-5
)
torch.testing.assert_close(
rank_k, full_k[:, k_lo:k_hi, :], atol=1e-5, rtol=1e-5
)
torch.testing.assert_close(
rank_v, full_v[:, k_lo:k_hi, :], atol=1e-5, rtol=1e-5
)
def test_tp2_extend_matches_full(self):
"""Single-chunk extend with TP=2 produces the same q / k / v slices
as a TP=1 reference, verified rank-by-rank.
"""
ref_cca, cfg = _make_tiny_cca(seed=21, tp_rank=0, tp_size=1)
S = 6
torch.manual_seed(901)
hs = torch.randn(S, ref_cca.hidden_size, dtype=torch.float32) * 0.1
ref_pool = _MockReqToTokenPool(pool_size=8, cca_config=cfg, tp_size=1)
ref_fb = _make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[S],
extend_prefix_lens_cpu=[0],
req_pool_indices=[0],
input_ids=torch.arange(S, dtype=torch.int64),
)
with _mock_pool_context(ref_pool):
full_q, full_k, full_v = ref_cca.forward(hs, ref_fb)
for tp_rank in range(self.TP_SIZE):
rank_cca, _ = _make_tiny_cca(
seed=21 + tp_rank, tp_rank=tp_rank, tp_size=self.TP_SIZE
)
self._slice_full_state_dict_into_rank(ref_cca, rank_cca, tp_rank)
rank_pool = _MockReqToTokenPool(
pool_size=8, cca_config=cfg, tp_size=self.TP_SIZE
)
rank_fb = _make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[S],
extend_prefix_lens_cpu=[0],
req_pool_indices=[0],
input_ids=torch.arange(S, dtype=torch.int64),
)
with _mock_pool_context(rank_pool):
rank_q, rank_k, rank_v = rank_cca.forward(hs, rank_fb)
self._check_per_rank_outputs(
full_q, full_k, full_v, rank_q, rank_k, rank_v, tp_rank, cfg
)
def test_tp2_decode_matches_full(self):
"""Prefill(S0) + decode(1 token) with TP=2 produces the same q / k / v
slices as a TP=1 reference, verifying that the per-rank conv state
and prev_hs cache (which is replicated on every rank) agree.
"""
ref_cca, cfg = _make_tiny_cca(seed=22, tp_rank=0, tp_size=1)
S0 = 5
torch.manual_seed(902)
hs_prefill = torch.randn(S0, ref_cca.hidden_size, dtype=torch.float32) * 0.1
hs_decode = torch.randn(1, ref_cca.hidden_size, dtype=torch.float32) * 0.1
ref_pool = _MockReqToTokenPool(pool_size=8, cca_config=cfg, tp_size=1)
with _mock_pool_context(ref_pool):
ref_cca.forward(
hs_prefill,
_make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[S0],
extend_prefix_lens_cpu=[0],
req_pool_indices=[0],
input_ids=torch.arange(S0, dtype=torch.int64),
),
)
full_q, full_k, full_v = ref_cca.forward(
hs_decode,
_make_forward_batch(
is_decode=True,
extend_seq_lens_cpu=[],
extend_prefix_lens_cpu=[],
req_pool_indices=[0],
input_ids=torch.tensor([0], dtype=torch.int64),
),
)
for tp_rank in range(self.TP_SIZE):
rank_cca, _ = _make_tiny_cca(
seed=22 + tp_rank, tp_rank=tp_rank, tp_size=self.TP_SIZE
)
self._slice_full_state_dict_into_rank(ref_cca, rank_cca, tp_rank)
rank_pool = _MockReqToTokenPool(
pool_size=8, cca_config=cfg, tp_size=self.TP_SIZE
)
with _mock_pool_context(rank_pool):
rank_cca.forward(
hs_prefill,
_make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[S0],
extend_prefix_lens_cpu=[0],
req_pool_indices=[0],
input_ids=torch.arange(S0, dtype=torch.int64),
),
)
rank_q, rank_k, rank_v = rank_cca.forward(
hs_decode,
_make_forward_batch(
is_decode=True,
extend_seq_lens_cpu=[],
extend_prefix_lens_cpu=[],
req_pool_indices=[0],
input_ids=torch.tensor([0], dtype=torch.int64),
),
)
self._check_per_rank_outputs(
full_q, full_k, full_v, rank_q, rank_k, rank_v, tp_rank, cfg
)
def test_tp2_conv_state_is_per_rank_sliced(self):
"""After a TP=2 prefill, each rank's conv state must equal the head
slice of the TP=1 conv state corresponding to that rank's heads.
"""
ref_cca, cfg = _make_tiny_cca(seed=23, tp_rank=0, tp_size=1)
S = 4
torch.manual_seed(903)
hs = torch.randn(S, ref_cca.hidden_size, dtype=torch.float32) * 0.1
ref_pool = _MockReqToTokenPool(pool_size=4, cca_config=cfg, tp_size=1)
with _mock_pool_context(ref_pool):
ref_cca.forward(
hs,
_make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[S],
extend_prefix_lens_cpu=[0],
req_pool_indices=[0],
input_ids=torch.arange(S, dtype=torch.int64),
),
)
full_state = ref_pool.mamba2_layer_cache(0).conv[0][0] # [in_out_ch_full, pad]
head_dim = cfg.head_dim
num_q_heads_full = cfg.num_attention_heads
num_k_heads_full = cfg.num_query_groups
latent_q_full = num_q_heads_full * head_dim
q_per_rank = num_q_heads_full // self.TP_SIZE
k_per_rank = num_k_heads_full // self.TP_SIZE
for tp_rank in range(self.TP_SIZE):
rank_cca, _ = _make_tiny_cca(
seed=23 + tp_rank, tp_rank=tp_rank, tp_size=self.TP_SIZE
)
self._slice_full_state_dict_into_rank(ref_cca, rank_cca, tp_rank)
rank_pool = _MockReqToTokenPool(
pool_size=4, cca_config=cfg, tp_size=self.TP_SIZE
)
with _mock_pool_context(rank_pool):
rank_cca.forward(
hs,
_make_forward_batch(
is_decode=False,
extend_seq_lens_cpu=[S],
extend_prefix_lens_cpu=[0],
req_pool_indices=[0],
input_ids=torch.arange(S, dtype=torch.int64),
),
)
rank_state = rank_pool.mamba2_layer_cache(0).conv[0][0]
q_lo = tp_rank * q_per_rank * head_dim
q_hi = q_lo + q_per_rank * head_dim
k_lo = latent_q_full + tp_rank * k_per_rank * head_dim
k_hi = k_lo + k_per_rank * head_dim
expected = torch.cat([full_state[q_lo:q_hi], full_state[k_lo:k_hi]], dim=0)
torch.testing.assert_close(rank_state, expected, atol=1e-5, rtol=1e-5)
def test_tp_assertions_reject_indivisible_head_counts(self):
"""The CCA constructor must reject TP sizes that don't evenly divide
both num_q_heads and num_k_heads, since both grouped-mean and
conv_qk.1 require each rank to hold whole K-head groups.
"""
from sglang.srt.models.zaya import CCA
cfg = _make_tiny_config()
# tiny config has num_query_groups=2; TP=4 cannot divide it cleanly.
with self.assertRaises(AssertionError):
CCA(
config=cfg,
cca_num_k_heads=cfg.num_query_groups,
cca_num_q_heads=cfg.num_attention_heads,
hidden_size=cfg.hidden_size,
head_dim=cfg.head_dim,
cca_time0=cfg.cca_time0,
cca_time1=cfg.cca_time1,
layer_id=0,
tp_rank=0,
tp_size=4,
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/models/test_zaya_mod_tp.py
+178
View File
@@ -0,0 +1,178 @@
"""Numerical correctness test for ZAYA1 MoE + MOD under TP>1.
Background: the MOD (mixture-of-depths) skip-expert residual blend must be
combined with the experts output on the correct side of the cross-rank
all-reduce. ``mod_out = hidden_states * prob`` is replicated on every TP rank,
so all-reducing it would multiply it by ``tp_size``. The model therefore masks
the *per-rank partial* experts output before the reduce and only adds the
replicated ``mod_out`` afterwards:
sum_r(mask · partial_r) + (1 - mask) · mod_out
= mask · experts_out_full + (1 - mask) · mod_out
This test drives the *real* helpers used by ``ZayaBlock.forward`` --
``mod_premask_experts`` and ``mod_blend`` -- so a regression in that math is
caught. The cross-rank all-reduce is simulated by summing the per-rank partials
(the masks are replicated, so the sum is exact), which keeps the test runnable
on CPU CI without a live ``torch.distributed`` group.
"""
import unittest
import torch
from sglang.srt.models.zaya import mod_blend, mod_premask_experts
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _reference_blend(
hidden_states: torch.Tensor, # [T, H]
probs: torch.Tensor, # [T, 1]
indices: torch.Tensor, # [T, 1]
experts_out_full: torch.Tensor, # [T, H] -- already-reduced full experts output
num_moe_experts: int,
) -> torch.Tensor:
"""Reference: apply the MOD mask to the *full* (already-reduced) experts
output, then add the skip path. Mirrors the intended algebra directly.
"""
mod_mask = (indices != num_moe_experts).to(experts_out_full.dtype)
mod_out = hidden_states * probs
return mod_mask * experts_out_full + (1.0 - mod_mask) * mod_out
def _real_tp_blend(
hidden_states: torch.Tensor,
probs: torch.Tensor,
indices: torch.Tensor,
partial_experts_per_rank: list[torch.Tensor], # one [T, H] per rank
num_moe_experts: int,
) -> torch.Tensor:
"""Production path: ``mod_premask_experts`` per rank -> simulated all-reduce
(sum) -> ``mod_blend``. Uses the exact helpers ``ZayaBlock.forward`` calls.
"""
mod_out = hidden_states * probs
reduced = None
mod_mask = None
for partial in partial_experts_per_rank:
mask, masked = mod_premask_experts(partial, indices, num_moe_experts)
mod_mask = mask
reduced = masked if reduced is None else reduced + masked
return mod_blend(reduced, mod_mask, mod_out)
def _buggy_old_tp_blend(
hidden_states: torch.Tensor,
probs: torch.Tensor,
indices: torch.Tensor,
partial_experts_per_rank: list[torch.Tensor],
num_moe_experts: int,
) -> torch.Tensor:
"""Old, broken sequence: all-reduce the replicated ``mod_out`` (so it gets
scaled by ``tp_size``) then mix. Proves the test catches a regression.
"""
tp_size = len(partial_experts_per_rank)
mod_out_replicated = hidden_states * probs
mod_out_after_allreduce = mod_out_replicated * tp_size # all-reduce of replicated
experts_out_full = torch.stack(partial_experts_per_rank, dim=0).sum(dim=0)
mod_mask = (indices != num_moe_experts).to(experts_out_full.dtype)
return mod_mask * experts_out_full + (1.0 - mod_mask) * mod_out_after_allreduce
class TestZayaMODUnderTP(CustomTestCase):
def _make_partials(self, T: int, H: int, tp_size: int):
torch.manual_seed(31)
experts_out_full = torch.randn(T, H, dtype=torch.float32) * 0.1
# Split into ``tp_size`` random partial tensors that sum to the full output.
partials = []
remaining = experts_out_full.clone()
for _ in range(tp_size - 1):
p = torch.randn_like(remaining) * 0.05
partials.append(p)
remaining = remaining - p
partials.append(remaining)
return experts_out_full, partials
def _make_inputs(self, T: int, H: int, num_experts: int, frac_skip: float):
torch.manual_seed(7)
hidden_states = torch.randn(T, H, dtype=torch.float32)
probs = torch.rand(T, 1, dtype=torch.float32)
# Build indices: with probability ``frac_skip`` mark token as skip-expert.
skip_id = num_experts # MOD uses ``num_moe_experts`` as the skip slot
rand = torch.rand(T, 1)
real = torch.randint(0, num_experts, (T, 1))
indices = torch.where(rand < frac_skip, torch.full_like(real, skip_id), real)
return hidden_states, probs, indices
def test_real_helpers_match_reference_for_tp(self):
"""The real ``mod_premask_experts`` / ``mod_blend`` path must equal the
reference blend for any TP size and any skip fraction.
"""
T, H = 8, 16
num_experts = 4
for tp_size in (2, 4, 8):
for frac_skip in (0.0, 0.5, 1.0):
hidden_states, probs, indices = self._make_inputs(
T, H, num_experts, frac_skip
)
full, partials = self._make_partials(T, H, tp_size)
ref = _reference_blend(hidden_states, probs, indices, full, num_experts)
real = _real_tp_blend(
hidden_states, probs, indices, partials, num_experts
)
torch.testing.assert_close(
real,
ref,
atol=1e-5,
rtol=1e-5,
msg=f"tp_size={tp_size} frac_skip={frac_skip}",
)
def test_premask_zeroes_skip_tokens(self):
"""``mod_premask_experts`` must zero the experts contribution exactly on
skip-routed tokens and pass through real-expert tokens unchanged.
"""
T, H = 6, 8
num_experts = 4
experts_out = torch.randn(T, H, dtype=torch.float32)
# Alternate skip / real tokens.
indices = torch.tensor(
[[num_experts], [0], [num_experts], [1], [num_experts], [2]],
dtype=torch.long,
)
mod_mask, masked = mod_premask_experts(experts_out, indices, num_experts)
skip_rows = indices.squeeze(-1) == num_experts
self.assertTrue(torch.all(masked[skip_rows] == 0))
torch.testing.assert_close(masked[~skip_rows], experts_out[~skip_rows])
# mask is 0 on skip rows, 1 elsewhere.
self.assertTrue(torch.all(mod_mask.squeeze(-1)[skip_rows] == 0))
self.assertTrue(torch.all(mod_mask.squeeze(-1)[~skip_rows] == 1))
def test_old_blend_is_wrong_when_skip_used(self):
"""Sanity: confirm the old (all-reduce mod_out) formula diverges from the
reference so a regression to that behavior would be caught.
"""
T, H = 8, 16
num_experts = 4
tp_size = 4
hidden_states, probs, indices = self._make_inputs(
T, H, num_experts, frac_skip=0.5
)
full, partials = self._make_partials(T, H, tp_size)
ref = _reference_blend(hidden_states, probs, indices, full, num_experts)
buggy = _buggy_old_tp_blend(
hidden_states, probs, indices, partials, num_experts
)
with self.assertRaises(AssertionError):
torch.testing.assert_close(buggy, ref, atol=1e-3, rtol=1e-3)
if __name__ == "__main__":
unittest.main()