Repository navigation
Support for Zyphra zaya1 model - #26347
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
|
Prefix cache and Cuda Graph issue fixed |
alexnails
left a comment
alexnails
left a comment
There was a problem hiding this comment.
Have left a mix of perf + correctness comments with some validation + suggestions by myself and Claude. Let me know if you have any questions!
| mod_out = tensor_model_parallel_all_reduce(mod_out) | ||
| experts_out = tensor_model_parallel_all_reduce(experts_out) | ||
|
|
||
| mod_mask = (indices != self.num_moe_experts).to(experts_out.dtype) |
There was a problem hiding this comment.
[Claude]
mod_out = hidden_states * probs is replicated on every rank. All-reducing a replicated tensor multiplies it by tp_size, so the skip-expert blend is tp_size× too large whenever MOD is enabled at TP>1. The PR was only tested at TP=1, so this didn't surface.
Math check: sum_r(mod_mask·partial_r) + (1-mod_mask)·mod_out = mod_mask·full_experts + (1-mod_mask)·mod_out. Correct, and one all-reduce instead of two.
[Me]
mod_mask = (indices != self.num_moe_experts).to(experts_out.dtype)
masked_experts = mod_mask * experts_out
if self.tp_size > 1:
masked_experts = tensor_model_parallel_all_reduce(masked_experts)
hidden_out = masked_experts + (1.0 - mod_mask) * mod_outThere was a problem hiding this comment.
Good catch — fixed in dd99074b6. The mask is now applied to the per-rank expert output before the all-reduce, and the replicated skip-path term is folded back after the reduce, so the single all-reduce yields mask·experts_full + (1-mask)·mod_out with no tp_size× inflation. The two steps are extracted as pure helpers mod_premask_experts / mod_blend (zaya.py:1021/1044) and called from ZayaBlock.forward (zaya.py:1151). Added test/registered/unit/models/test_zaya_mod_tp.py, which checks tp_size∈{2,4,8} against a single-rank reference and keeps a regression case proving the old (×tp_size) formula fails.
| self.clamp_temp = bool(getattr(config, "clamp_temp", False)) | ||
|
|
||
| bias = bool(getattr(config, "attention_bias", False)) | ||
| self.linear_q = ReplicatedLinear( |
There was a problem hiding this comment.
both the grouped-mean step and the second conv_qk stage's groups=num_k_heads+num_q_heads can be split on head and parallelized?
this change could be some more work so let me know if you want this to be a follow up
There was a problem hiding this comment.
Done in this PR (dd99074b6) rather than a follow-up. CCA now runs head-parallel TP: both Q heads and K heads are split evenly across ranks, so grouped-mean and conv_qk.1 stay head-local and need no collective inside the QKV projection (groups uses the per-rank head count, zaya.py:303). V is computed full then sliced to the rank's K-head range (_slice_v_per_rank, zaya.py:447), and o_proj is a RowParallelLinear with one all-reduce. The grouped conv weights/bias and per-K-head temp get TP-aware weight_loaders (zaya.py:325). Constraints: both head counts must be divisible by tp_size; KV-replication TP (tp_size > num_k_heads) and DP attention are rejected with explicit asserts (zaya.py:217, zaya.py:773). Covered by TestZayaCCATensorParallel in test/registered/unit/models/test_zaya_cca.py.
| start = 0 | ||
| for i, seq_len in enumerate(extend_seq_lens_cpu): | ||
| end = start + int(seq_len) | ||
| mamba_idx = int(mamba_indices[i].item()) |
There was a problem hiding this comment.
GPU<->CPU sync per request per layer.
Instead we can do one sync, outside the loop:
conv_state, prev_hs_state, mamba_indices = self._get_pool_state(forward_batch)
+ mamba_idx_cpu = mamba_indices.tolist()
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu or []
extend_prefix_lens_cpu = forward_batch.extend_prefix_lens_cpu or []
start = 0
for i, seq_len in enumerate(extend_seq_lens_cpu):
end = start + int(seq_len)
- mamba_idx = int(mamba_indices[i].item())
+ mamba_idx = mamba_idx_cpu[i]There was a problem hiding this comment.
Adopted in 9bc04264b. mamba_indices.tolist() is now done once before the loop (zaya.py:535) and the loop uses mamba_idx_cpu[i], removing the per-request .item() sync. The decode path uses index_select/index_copy_, so it has no per-element sync either.
| else: | ||
| hidden_states = norm(residual) | ||
| return hidden_states.to(target_dtype) | ||
| return norm(residual.to(target_dtype)) |
There was a problem hiding this comment.
[Claude]
High-impact perf — eager-mode RMSNorm on every layer.
residual here is always fp32 (produced by ResidualScaling.forward and the residual.float() + hidden_states.float() accumulations in the decoder layers), and norm.weight is bf16, so the residual.dtype != norm.weight.dtype branch always fires. That routes to RMSNorm.forward_native (layernorm.py:386) which does cast → pow → mean → rsqrt → mul → cast in eager mode — 5+ kernel launches per call, on every one of 120 norm calls per step.
The fp32 residual stream is preserved across layers by ResidualScaling independently of what we do here; only the norm's own input dtype matters. Simplify to:
def _apply_norm_with_fp32_residual(norm, residual, target_dtype):
# The fp32 residual stream is preserved by the caller (still holds the
# fp32 residual tensor for the next accumulation). The only precision
# concession is bf16-rounding the norm input — same as every other SGLang
# model — which lets us hit the fused sgl_kernel rmsnorm path.
return norm(residual.to(target_dtype))Alternative if MMLU drifts: use torch.nn.functional.rms_norm(residual, (norm.hidden_size,), norm.weight.data, norm.variance_epsilon) which is fused at the ATen level and accepts fp32 input with a lower-precision weight.
There was a problem hiding this comment.
Fixed and verified accuracy no loss.
| extend_prefix_lens_cpu = forward_batch.extend_prefix_lens_cpu or [] | ||
|
|
||
| start = 0 | ||
| for i, seq_len in enumerate(extend_seq_lens_cpu): |
There was a problem hiding this comment.
[Claude]
Medium-impact perf — Python-driven Conv1d launches.
Each iteration calls _conv_qk_run on a tiny [1, C, S_cur+2] tensor. For B requests × 60 attention layers that's 60·B kernel launches per prefill step.
The all-no-prefix common case (extend_prefix_lens_cpu all zero) batches into one call by packing [1, C, sum(S_cur) + B·total_padding] with zero left-pads inserted at each request boundary, running one conv, then slicing back out. The resumed-prefill case is harder (each request has its own cached left-pad) but rarer.
Easy wrapper:
if not any(int(p) > 0 for p in extend_prefix_lens_cpu):
return self._forward_extend_fresh(...)There was a problem hiding this comment.
Implemented in 9bc04264b. Added an all_fresh fast path (zaya.py:547): when no request carries a prefix, all requests are packed into a single [1, C, Σ(S_i)+B·total_padding] buffer with zero left-pads at each boundary, run through one conv_qk, then sliced back out (and each request's tail conv-state written back). The rarer resumed-prefill case (each request has its own cached left-pad) keeps the per-request loop. Same dispatch idea as your _forward_extend_fresh suggestion, kept inline to share the buffers.
| ): | ||
| hs = hs + prev_router_hidden_states * self.router_states_scale | ||
|
|
||
| router_hidden_states_next = hs.clone() |
There was a problem hiding this comment.
Do we actually need a clone here?
- router_hidden_states_next = hs.clone()
+ router_hidden_states_next = hsThere was a problem hiding this comment.
Removed in 9bc04264b (zaya.py:986)
| query_out = query_out_grouped.reshape(query_pre.shape[0], -1, query_pre.shape[-1]) | ||
|
|
||
| query_pre_mean = torch.mean( | ||
| query_pre_grouped.float(), dim=-2, dtype=torch.float32 |
There was a problem hiding this comment.
query_pre_grouped.float() is computed twice (line 308 inside the sum, line 314 here inside the mean). Cache once; also avoid the view_as(...).float() round-trip:
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_fp32There was a problem hiding this comment.
Applied verbatim in 9bc04264b (zaya.py:429). query_pre_grouped_fp32 is computed once and reused for both query_out_grouped and query_pre_mean, and the view_as(...).float() round-trip is gone.
| sqrt_head_dim = float(self.sqrt_head_dim) | ||
| query_fp32 = query.to(torch.float32) | ||
| q_norm = torch.linalg.vector_norm(query_fp32, ord=2, dim=-1, keepdim=True) | ||
| query_fp32 = query_fp32 * torch.rsqrt(q_norm * q_norm + eps) * sqrt_head_dim |
There was a problem hiding this comment.
torch.linalg.vector_norm(..., ord=2) already computes sqrt(Σx²), then the next line squares it back and rsqrts. Compute rsqrt(Σx² + eps) directly which is mathematically equivalent:
inv_q = torch.rsqrt(query_fp32.pow(2).sum(-1, keepdim=True) + eps) * sqrt_head_dim
query_fp32 = query_fp32 * inv_q
inv_k = torch.rsqrt(key_fp32.pow(2).sum(-1, keepdim=True) + eps) * sqrt_head_dim
key_fp32 = key_fp32 * inv_kThere was a problem hiding this comment.
Done in 9bc04264b (zaya.py:396). Replaced the vector_norm → square → rsqrt round-trip with a direct rsqrt(Σx² + eps), mathematically equivalent.
|
|
||
| conv_state, prev_hs_state, mamba_indices = self._get_pool_state(forward_batch) | ||
|
|
||
| extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu or [] |
There was a problem hiding this comment.
Claude suggesting throwing an assert like the following code, but I am not sure that is best solution, open to your thoughts!
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
extend_prefix_lens_cpu = forward_batch.extend_prefix_lens_cpu
assert (
extend_seq_lens_cpu is not None and extend_prefix_lens_cpu is not None
), "CCA._forward_extend requires extend_seq_lens_cpu and extend_prefix_lens_cpu"There was a problem hiding this comment.
I left the assert out for now. _forward_extend only runs for the EXTEND/MIXED/DRAFT_EXTEND/DLLM_EXTEND modes, where SGLang always populates extend_seq_lens_cpu/extend_prefix_lens_cpu, and the only empty case (T==0 idle) is already short-circuited at the top of forward (zaya.py:721). The fresh-prefill guard also uses bool(extend_seq_lens_cpu) defensively. So an assert would duplicate an existing contract — but if you'd prefer an explicit fail-fast there, I'm happy to add it; it's cheap.
| torch.cat([v1, v2], dim=-1) | ||
| .view(T, self.num_k_heads, self.head_dim) | ||
| .to(torch.float32) | ||
| ) |
There was a problem hiding this comment.
No need for this fp32 cast?
value = torch.cat([v1, v2], dim=-1).view(T, self.num_k_heads, self.head_dim)Also: can get change fp32 cast on the empty-batch early-return at L531 to match (return in input dtype, not fp32)
There was a problem hiding this comment.
Both done in 9bc04264b. Dropped the .to(torch.float32) on value in all three build paths (no-state / extend / decode); value now stays in the input dtype since the caller casts before rotary + attention anyway. The empty-batch early-return matches: q/k stay fp32 (the normalize step keeps fp32 for stability) and v is returned in the input dtype, no longer cast to fp32 (zaya.py:721).
alexnails
left a comment
alexnails
left a comment
There was a problem hiding this comment.
can you try fixing the per layer mamba sync ?
[Claude fix below]
Edit 1 — replace _get_pool_state (zaya.py:398–407)
Replace this:
def _get_pool_state(self, forward_batch: ForwardBatch):
"""Retrieve per-request CCA state from the centralized MambaPool."""
req_to_token_pool = get_req_to_token_pool()
layer_cache = req_to_token_pool.mamba2_layer_cache(self.layer_id)
conv_state = layer_cache.conv[0]
prev_hs_state = layer_cache.conv[1]
mamba_indices = req_to_token_pool.get_mamba_indices(
forward_batch.req_pool_indices
).to(torch.long)
return conv_state, prev_hs_state, mamba_indiceswith this:
@staticmethod
def _mamba_indices(forward_batch: ForwardBatch, req_to_token_pool) -> torch.Tensor:
"""Gather the int64 mamba indices once per forward, shared across layers.
``get_mamba_indices`` depends only on ``req_pool_indices`` and is
therefore identical for every CCA layer in a forward, and
``req_index_to_mamba_index_mapping`` is int32 so the ``.to(torch.long)``
allocates on every call. Cache the result on ``forward_batch`` keyed by
the ``req_pool_indices`` tensor identity (rebuilt per forward, so the
cache self-invalidates) to run the gather + cast once instead of once
per attention layer (60x for ZAYA1-base).
"""
cache = getattr(forward_batch, "_zaya_mamba_indices", None)
if cache is not None and cache[0] is forward_batch.req_pool_indices:
return cache[1]
mamba_indices = req_to_token_pool.get_mamba_indices(
forward_batch.req_pool_indices
).to(torch.long)
forward_batch._zaya_mamba_indices = (
forward_batch.req_pool_indices,
mamba_indices,
)
return mamba_indices
@staticmethod
def _mamba_indices_cpu(
forward_batch: ForwardBatch, mamba_indices: torch.Tensor
) -> list[int]:
"""Materialize the mamba indices on CPU once per forward (extend only).
The prefill loop indexes the state pool with Python ints, so it needs a
CPU list; ``.tolist()`` forces a GPU->CPU sync. Cache it the same way as
:meth:`_mamba_indices` so the sync happens once per forward instead of
once per attention layer (60 syncs/prefill-step -> 1 for ZAYA1-base).
The decode path never calls this, so decode stays sync-free.
"""
cache = getattr(forward_batch, "_zaya_mamba_indices_cpu", None)
if cache is not None and cache[0] is forward_batch.req_pool_indices:
return cache[1]
mamba_idx_cpu = mamba_indices.tolist()
forward_batch._zaya_mamba_indices_cpu = (
forward_batch.req_pool_indices,
mamba_idx_cpu,
)
return mamba_idx_cpu
def _get_pool_state(self, forward_batch: ForwardBatch):
"""Retrieve per-request CCA state from the centralized MambaPool."""
req_to_token_pool = get_req_to_token_pool()
layer_cache = req_to_token_pool.mamba2_layer_cache(self.layer_id)
conv_state = layer_cache.conv[0]
prev_hs_state = layer_cache.conv[1]
mamba_indices = self._mamba_indices(forward_batch, req_to_token_pool)
return conv_state, prev_hs_state, mamba_indicesEdit 2 — use the cached CPU list in _forward_extend (zaya.py:554–557)
Replace this:
conv_state, prev_hs_state, mamba_indices = self._get_pool_state(forward_batch)
# Materialize mamba_indices on CPU once instead of once per loop iter
# (each ``.item()`` is a GPU->CPU sync; B*60 syncs/step otherwise).
mamba_idx_cpu = mamba_indices.tolist()with this:
conv_state, prev_hs_state, mamba_indices = self._get_pool_state(forward_batch)
# Materialize mamba_indices on CPU once per forward (cached across the 60
# CCA layers) instead of once per layer; each ``.tolist()`` is a
# GPU->CPU sync, so this turns 60 syncs/prefill-step into 1.
mamba_idx_cpu = self._mamba_indices_cpu(forward_batch, mamba_indices)
Fixed in the latest commit |
alexnails
left a comment
alexnails
left a comment
There was a problem hiding this comment.
seems there is small conflict. approving!
|
/tag-and-rerun-ci |
12ee549 to
33b8812
Compare
github-actions
Bot
added
the
documentation
The `Sync LMSYS SGLang blogs` workflow pushed directly to `main` on a
schedule, which had two failure modes:
1. On `sgl-project/sglang`, `main` is protected, so every scheduled run
failed with GH013 ("Changes must be made through a pull request") and
the blog cards in `docs_new/index.mdx` never updated.
2. Any fork (or non-main branch) that ran this workflow pushed the
generated commit to its own branch, and those "docs: sync LMSYS
SGLang blog cards" commits then leaked into pull requests (e.g. sgl-project#26347).
Fix:
- Guard the job to `github.repository == 'sgl-project/sglang'` AND
`github.ref == 'refs/heads/main'`, so it never runs on a fork or on a
non-main branch (a stray manual workflow_dispatch included).
- Propose the change through a pull request opened with the GitHub CLI
(`gh`, pre-installed on the runner -- no third-party action) instead of
pushing to the protected branch. The step is a no-op when the cards are
unchanged and reuses a single `auto/lmsys-blog-cards` branch/PR.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
/rerun-test test/registered/models/test_zaya.py |
|
Results for 🚀 |
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
## Summary - Add FP8 training for ZAYA1-8B (`fp8: hybrid`). CCA Q/K/V/O projections and expert SwiGLU GEMMs run as Transformer Engine linears (`parallel_mode="duplicated"`) under one FP8 autocast for the whole stack. The router, depthwise convolutions, RMSNorm, attention, embeddings, and the tied logits projection stay in BF16/FP32. - Pad the token axis to a multiple of 8 before each TE GEMM. A top-1 route can hand one expert as few as 4 tokens, and TE FP8 needs the leading-dimension product divisible by 8. Reject FP8 GEMMs whose in or out features are not multiples of 16. The router stays BF16 because its last projection scores 17 classes. - Add `zaya1_8B-FP8-pretrain.yaml` for MI300X and MI355X with `fp8_recipe: tensorwise`. Current scaling takes each scale from the tensor being cast, so FP8 is accurate from step 1 with no amax history. Delayed scaling starts with an empty history; on the released checkpoint a cold delayed-scaling forward scored SGLang rollouts at Pearson r 0.72. - Switch the BF16 and FP8 ZAYA1 recipes to mock data so they run without a dataset mount. `vocab_size` stays 262272 from `zaya1.yaml`. - Update `docs/04-technical-guides/zaya1/validation.md` and `docs/06-developer-guide/model-support-matrix.md`. ## Validation Run on 1x MI300X node (tw042), `rocm/primus:v26.8`, TE 2.17.1, Megatron-LM `d3528a21`. Log-prob parity against SGLang ([sgl-project/sglang#26347](sgl-project/sglang#26347)): the same 4 GSM8K greedy rollouts used for BF16, 1,024 tokens each, scored by Primus with the `Zyphra/ZAYA1-base` weights and no FP8 warm-up. | Rollout tokens (n=4096) | Pearson r | MAE | RMSE | |---|---|---|---| | BF16 | 0.978 | 0.010 | 0.055 | | FP8 tensorwise, cold | 0.935 | 0.026 | 0.101 | | FP8 delayed, cold | 0.720 | 0.155 | — | | FP8 delayed, after 5 backward calibration passes | 0.893 | 0.033 | 0.142 | GSM8K pretraining from scratch with the FP8 recipe (seq 128, global batch 8, lr 1e-5). Early losses track the BF16 validation run: | Step | BF16 | FP8 tensorwise | |---|---|---| | 1 | 12.820 | 12.873 | | 10 | 11.229 | 10.911 | | 30 | 7.597 | 7.791 | No skipped or NaN iterations. The 30k-step run was not completed: at 3.2 s/step on 4 GPUs it needs about 27 h, beyond the 8 h job limit. ## Test plan - [x] `pytest tests/unit_tests/megatron/transformer/zaya1` (20 passed on MI300X, including `test_cold_tensorwise_fp8_matches_high_precision`) - [x] `pytest tests/trainer/test_megatron_trainer.py -k zaya1` (BF16 and FP8, 4 stages, 3 iterations, mock data) - [x] `pre-commit run` on the changed files - [ ] MI355X run of `zaya1_8B-FP8-pretrain.yaml` (not yet exercised on gfx950) - [ ] Full 30k-step FP8 GSM8K run ---------
Motivation
Add SGLang inference support for the Zyphra ZAYA1 family of hybrid models
(
Zyphra/ZAYA1-base, 8.4 B total / 760 M active parameters). ZAYA1 alternatesCompressed Convolutional Attention (CCA, arXiv:2510.04476)
attention layers with top-1 Mixture-of-Experts FFN layers and currently
has no upstream implementation in SGLang. After this PR, users can serve
ZAYA1 checkpoints on both NVIDIA and AMD GPUs through the normal
sglang.serveentry point.
ZAYA1 architecture summary (defaults match
Zyphra/ZAYA1-base/config.json):model_typezayahead_dim=128,partial_rotary_factor=0.5,rope_theta=1e6,max_position_embeddings=32768ResidualScaling(scale_residual_merge=True)tie_word_embeddings=True; fp32 projection whenzaya_high_prec=TrueModifications
New code
python/sglang/srt/configs/zaya.py—ZayaConfig(model_type="zaya") withAutoConfig.registerregistration, including derivation of HF-stylerope_parametersfromrope_theta/partial_rotary_factor.python/sglang/srt/configs/__init__.py— exportZayaConfig.python/sglang/srt/models/zaya.py— full model implementation:CCA— fused QKV projection, two-stage depthwise + groupednn.Conv1d(
cca_time0 = cca_time1 = 2), grouped-mean QK correction, fp32 RMSNorm(no learnable scale,
eps=1e-12), per-head learnable temperature on K,and a
val_proj1/val_proj2pair that consumes the current and theright-shifted hidden state.
ZayaAttention— wraps CCA outputs into SGLang'sRadixAttentionso theactual softmax MHA reuses the existing FlashAttention3 / Triton / AITER
backends.
ZayaRouter— 3-layer MLP + RMSNorm + softmax with optional EDA(depth-wise averaging) and MOD (expert skipping).
ZayaBlock— combinesZayaRouterwithFusedMoE.ZayaDecoderATTLayer/ZayaDecoderMLPLayer— inputRMSNormin fp32plus per-layer
ResidualScalingon the residual branch.ZayaModel+ZayaForCausalLM(EntryClass = ZayaForCausalLM),VocabParallelEmbedding,ParallelLMHeadwith fp32 projection path,HF safetensors weight loader matching the upstream key layout.
Per-request CCA state (rationale for not using
HybridLinearAttnBackend)CCA needs two small per-request, per-layer states:
conv_states(latent_q + latent_k, 2)conv_qkacross prefill / decode boundaryprev_hs(hidden_size,)val_proj2during decodeThe post-CCA attention is still vanilla softmax MHA, so reusing
HybridLinearAttnBackendwould force a new backend class and a layer-id ↔backend dispatch change for what is effectively a tiny pre-projection cache.
Instead, the CCA module owns its own pool:
Indexed by
forward_batch.req_pool_indices, lazily sized fromreq_to_token_pool.size, written back at the end of each forward, and clearedwhen the request pool is reset. Memory cost is ≈ 60 MB for 80 attention layers
× 80 concurrent requests at bf16 — negligible compared to KV cache.
Parallelism
Conv1duseReplicatedLinear: each TP rankholds a full copy of the (small) CCA weights and state pool.
q / k / vare sliced across heads by TP rank, then handed toRadixAttention, so attention compute itself scales with TP as usual.FusedMoE, supporting both intra-expert TP andinter-expert EP.
Tests and docs
test/registered/unit/configs/test_zaya_config.py— CPU unit tests fordefault fields,
rope_parametersderivation, and AutoConfig registration.test/registered/unit/models/test_zaya_cca.py— CPU numerical tests forCCA: single-chunk vs reference, prefill-then-decode equivalence, state-pool
lazy resize, two-request state isolation, and batched-decode vs
single-decode parity (specifically guards against a bf16-accumulation
regression that was observed during development).
test/registered/models/test_zaya.py— E2E server test (gated behindRUN_ZAYA_E2E=1so the standard registered suite does not download the~17 GB checkpoint). Registered into the CUDA
extra-astage on the1-gpu-largerunner and the AMDstage-b-test-1-gpu-large-amdsuite.docs/supported_models/text_generation/zaya_design.md— architecture,state-cache strategy, parallelism plan, NVIDIA/AMD compatibility notes.
docs/supported_models/text_generation/zaya_testing.md— repro commands,expected outputs, and measured memory / throughput numbers.
Accuracy Tests
Unit tests (CPU,
pytest test/registered/unit/{configs,models}/test_zaya_*.py -v)E2E (AMD MI300X, ROCm 7.2,
Zyphra/ZAYA1-base)Launch command used by
test/registered/models/test_zaya.py:MMLU 32-sample slice on the pretrain checkpoint (
max_tokens=1024is requiredbecause ZAYA1-base emits a long
<think>…</think>chain-of-thought before thefinal letter — instruction-tuned variants will score substantially higher):
Manual three-prompt batched smoke (greedy,
max_new_tokens=16) — outputsmatch the single-prompt path token-for-token, confirming that the
per-request
conv_states/prev_hspool isolates state correctly:Speed Tests and Profiling
Single MI300X,
mem-fraction-static=0.5,tritonattention backend,disable-radix-cache,disable-cuda-graph:ZayaForCausalLM)A dedicated Triton/CUDA-graph optimization pass is intentionally out of scope
for this PR (see "Known limitations" below); the numbers above are the
unoptimized baseline that future work can compare against.
Known limitations (filed as follow-up PRs)
--disable-radix-cacherequired)conv_statecross-request reuse needs a separate design--disable-cuda-graphrequired)swa_layers = Noneexercisedtorch.*ops, nosgl_kernel/ HIP-only calls); CUDA-sideregister_cuda_ciis already wired but needs a real NVIDIA runner runChecklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ⏳ Run #27252960731
Latest PR Test (Extra): ❌ Run #27252960577