/test/registered/unit/models/test_zaya_mod_tp.py

"""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()