support qwen 3.8 flash next (#37500)

Co-authored-by: ch-wan <54331508+ch-wan@users.noreply.github.com>
Co-authored-by: ispobock <26454835+ispobock@users.noreply.github.com>
Co-authored-by: JustinTong0323 <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: samuellees <26428561+samuellees@users.noreply.github.com>
Co-authored-by: YAMY1234 <74099316+YAMY1234@users.noreply.github.com>
Co-authored-by: yhyang201 <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: yizhang2077 <25844240+yizhang2077@users.noreply.github.com>
Co-authored-by: zijiexia <37504505+zijiexia@users.noreply.github.com>
Co-authored-by: Shinto C V <cshintov@gmail.com>
Co-authored-by: Julian Huang <huangzhilin.hzl@antgroup.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
Co-authored-by: yhyang201 <yhyang201@gmail.com>

test/registered/e2e/models/test_qwen4_exp_models.py

@@ -0,0 +1,90 @@
+"""Qwen3.8-Flash-Next (Qwen4-Exp) E2E on B200; the plain-serving case is kept:
+MTP's verify widths never exercise the QSA sparse-decode path."""
+
+import unittest
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
+from sglang.test.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+    try_cached_model,
+)
+
+register_cuda_ci(est_time=1500, stage="base-c", runner_config="4-gpu-b200")
+
+MODEL = "RadixArk/Qwen3.8-Flash-Next-NVFP4"
+
+SERVER_LAUNCH_TIMEOUT = 3600
+GSM8K_SCORE_THRESHOLD = 0.94
+
+BASE_ARGS = [
+    "--tp-size",
+    "4",
+    "--mem-fraction-static",
+    "0.85",
+    "--chunked-prefill-size",
+    "8192",
+    "--linear-attn-prefill-backend",
+    "flashinfer",
+    "--linear-attn-decode-backend",
+    "flashinfer",
+    "--mamba-ssm-dtype",
+    "bfloat16",
+    "--reasoning-parser",
+    "qwen3-thinking",
+]
+
+
+class _Qwen4ExpServer:
+    speculative_args: list[str] = []
+    model = try_cached_model(MODEL)
+    base_url = DEFAULT_URL_FOR_TEST
+    gsm8k_backend = "sgl_eval"
+    gsm8k_thinking = True
+    gsm8k_num_examples = 200
+    gsm8k_num_threads = 32
+    gsm8k_max_tokens = 16384
+    gsm8k_score_threshold = GSM8K_SCORE_THRESHOLD
+
+    @classmethod
+    def setUpClass(cls):
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=BASE_ARGS + cls.speculative_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+
+class TestQwen4ExpBase(_Qwen4ExpServer, GSM8KMixin, CustomTestCase):
+    """Normal autoregressive serving."""
+
+
+class TestQwen4ExpMTP(_Qwen4ExpServer, GSM8KMixin, CustomTestCase):
+    """NEXTN MTP serving (3 steps, topk 1, 4 draft tokens)."""
+
+    # GSM8K accept length measured at 3.02-3.03 (max 4.0 with 3 steps);
+    # 2.9 leaves noise margin while still failing on a real drop.
+    gsm8k_accept_length_thres = 2.9
+    speculative_args = [
+        "--speculative-algorithm",
+        "NEXTN",
+        "--speculative-num-steps",
+        "3",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "4",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernel/embeddings/test_qwen4_ple_offload.py

@@ -0,0 +1,195 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from sglang.srt.layers.quantization.unquant import UnquantizedEmbeddingMethod
+from sglang.srt.layers.vocab_parallel_embedding import (
+    VocabParallelEmbeddingShardIndices,
+)
+from sglang.srt.models import qwen4_exp as qwen4_exp_module
+from sglang.srt.models.qwen4_exp import (
+    Qwen4ExpPinnedHostEmbedding,
+    Qwen4ExpPLELayer,
+)
+from sglang.srt.utils import set_weight_attrs
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required for this test."
+)
+
+
+def _make_source_embedding(
+    *,
+    dtype=torch.bfloat16,
+    embedding_dim=7,
+    vocab_start=0,
+    vocab_end=8,
+    org_vocab_size=8,
+    tp_size=1,
+    num_added_embeddings=0,
+):
+    local_rows = vocab_end - vocab_start
+    weight = nn.Parameter(
+        torch.empty((local_rows, embedding_dim), dtype=dtype, device="cuda"),
+        requires_grad=False,
+    )
+    set_weight_attrs(
+        weight,
+        {
+            "input_dim": 1,
+            "output_dim": 0,
+            "weight_loader": lambda *_args, **_kwargs: None,
+        },
+    )
+    shard_indices = VocabParallelEmbeddingShardIndices(
+        padded_org_vocab_start_index=vocab_start,
+        padded_org_vocab_end_index=vocab_end,
+        padded_added_vocab_start_index=org_vocab_size,
+        padded_added_vocab_end_index=org_vocab_size,
+        org_vocab_start_index=vocab_start,
+        org_vocab_end_index=vocab_end,
+        added_vocab_start_index=org_vocab_size,
+        added_vocab_end_index=org_vocab_size,
+    )
+    return SimpleNamespace(
+        weight=weight,
+        quant_config=None,
+        enable_tp=True,
+        use_attn_tp_group=False,
+        tp_size=tp_size,
+        num_embeddings=org_vocab_size + num_added_embeddings,
+        org_vocab_size=org_vocab_size,
+        padding_size=1,
+        num_added_embeddings=num_added_embeddings,
+        use_presharded_weights=False,
+        org_vocab_size_padded=org_vocab_size,
+        num_embeddings_padded=org_vocab_size + num_added_embeddings,
+        shard_indices=shard_indices,
+        embedding_dim=embedding_dim,
+        weight_scale=None,
+        quant_method=UnquantizedEmbeddingMethod(),
+        num_embeddings_per_partition=local_rows,
+        num_org_embeddings_per_partition=local_rows,
+        num_added_embeddings_per_partition=0,
+    )
+
+
+def _load_rows(offloaded, rows):
+    pointer = offloaded.weight.data_ptr()
+    offloaded.weight_loader(offloaded.weight, rows)
+    assert offloaded.weight.data_ptr() == pointer
+    assert offloaded.weight.is_pinned()
+    assert offloaded.weight.weight_loader.__self__ is offloaded
+    assert offloaded.quant_method is None
+
+
+@pytest.mark.parametrize("input_dtype", [torch.int32, torch.int64])
+@pytest.mark.parametrize("embedding_dim", [7, 64, 257])
+def test_qwen4_ple_pinned_gather_tp1(input_dtype, embedding_dim):
+    source = _make_source_embedding(embedding_dim=embedding_dim)
+    offloaded = Qwen4ExpPinnedHostEmbedding(source)
+    rows = torch.arange(8 * embedding_dim, dtype=torch.bfloat16, device="cuda").reshape(
+        8, embedding_dim
+    )
+    _load_rows(offloaded, rows)
+
+    ids = torch.tensor([[0, 7, 3], [4, 1, 6]], dtype=input_dtype, device="cuda")
+    expected = rows.index_select(0, ids.long().flatten()).reshape(
+        *ids.shape, embedding_dim
+    )
+    actual = offloaded(ids)
+
+    assert actual.shape == expected.shape
+    assert actual.is_contiguous()
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_qwen4_ple_pinned_gather_shard_boundaries_and_out_buffer():
+    embedding_dim = 13
+    source = _make_source_embedding(
+        embedding_dim=embedding_dim,
+        vocab_start=4,
+        vocab_end=8,
+        org_vocab_size=8,
+        tp_size=2,
+    )
+    offloaded = Qwen4ExpPinnedHostEmbedding(source)
+    rows = torch.arange(8 * embedding_dim, dtype=torch.bfloat16, device="cuda").reshape(
+        8, embedding_dim
+    )
+    _load_rows(offloaded, rows)
+
+    ids = torch.tensor([[-1, 3, 4], [7, 8, 100]], device="cuda")
+    output = torch.full(
+        (*ids.shape, embedding_dim),
+        torch.nan,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    actual = offloaded.gather(ids, out=output)
+    expected = torch.zeros_like(output)
+    expected[0, 2] = rows[4]
+    expected[1, 0] = rows[7]
+
+    assert actual.data_ptr() == output.data_ptr()
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_qwen4_ple_pinned_gather_empty_input():
+    offloaded = Qwen4ExpPinnedHostEmbedding(_make_source_embedding())
+    _load_rows(offloaded, torch.zeros((8, 7), dtype=torch.bfloat16, device="cuda"))
+    ids = torch.empty((0, 3), dtype=torch.int64, device="cuda")
+    actual = offloaded.gather(ids)
+    assert actual.shape == (0, 3, 7)
+    assert actual.numel() == 0
+
+
+def test_qwen4_ple_pinned_embedding_rejects_unsupported_weights():
+    with pytest.raises(TypeError, match="requires bfloat16"):
+        Qwen4ExpPinnedHostEmbedding(_make_source_embedding(dtype=torch.float16))
+    with pytest.raises(NotImplementedError, match="added vocabulary"):
+        Qwen4ExpPinnedHostEmbedding(_make_source_embedding(num_added_embeddings=1))
+
+
+def test_qwen4_ple_prefetch_buffer_lifecycle(monkeypatch):
+    layer = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
+    nn.Module.__init__(layer)
+    layer.ple_embed_dim = 7
+    layer.ple_embedding = SimpleNamespace(
+        ngram_embedding=Qwen4ExpPinnedHostEmbedding(
+            _make_source_embedding(embedding_dim=layer.ple_embed_dim)
+        )
+    )
+    layer._graph_prefetch_buffers = {}
+    layer._eager_prefetch_buffer = None
+    lookup_ids = torch.empty((0,), dtype=torch.int64, device="cuda")
+
+    monkeypatch.setattr(qwen4_exp_module, "get_is_capture_mode", lambda: False)
+    eager_large = layer._get_prefetch_buffer(8, lookup_ids)
+    eager_small = layer._get_prefetch_buffer(3, lookup_ids)
+    assert eager_small.data_ptr() == eager_large.data_ptr()
+    assert layer._eager_prefetch_buffer.shape == (8, layer.ple_embed_dim)
+
+    eager_grown = layer._get_prefetch_buffer(12, lookup_ids)
+    eager_grown_small = layer._get_prefetch_buffer(4, lookup_ids)
+    assert eager_grown_small.data_ptr() == eager_grown.data_ptr()
+    assert layer._eager_prefetch_buffer.shape == (12, layer.ple_embed_dim)
+
+    monkeypatch.setattr(qwen4_exp_module, "get_is_capture_mode", lambda: True)
+    graph_three = layer._get_prefetch_buffer(3, lookup_ids)
+    graph_five = layer._get_prefetch_buffer(5, lookup_ids)
+    graph_three_reused = layer._get_prefetch_buffer(3, lookup_ids)
+    assert graph_three_reused.data_ptr() == graph_three.data_ptr()
+    assert graph_five.data_ptr() != graph_three.data_ptr()
+    assert set(layer._graph_prefetch_buffers) == {3, 5}
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernel/hyperconnection/test_hc_mix_triton.py

@@ -0,0 +1,84 @@
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from sglang.srt.layers.hc_mix_triton import (
+    _FUSED_MIX_MAX_ROWS,
+    fused_hc_mix,
+    fused_hc_mix_supported,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+HC_COUNT = 4
+HIDDEN_SIZE = 2560
+LOWRANK = 320
+
+
+def _reference_mix(
+    hyper_input_normed: torch.Tensor,
+    w_down: torch.Tensor,
+    w_up: torch.Tensor,
+    hc: int,
+    hs: int,
+    compute_dtype: torch.dtype = torch.float64,
+) -> torch.Tensor:
+    """Mirrors GatedResidual._mix_compute in hyperconnection.py."""
+    x = hyper_input_normed.to(compute_dtype)
+    t = F.silu(F.linear(x, w_down.to(compute_dtype)) / hc)
+    u = torch.sigmoid(F.linear(t, w_up.to(compute_dtype)))
+    return (u.unflatten(-1, (hc, hs)) * x.unflatten(-1, (hc, hs))).mean(dim=-2)
+
+
+def _make_inputs(num_tokens: int, dtype: torch.dtype):
+    torch.manual_seed(0)
+    x = torch.randn(num_tokens, HC_COUNT * HIDDEN_SIZE, dtype=dtype, device="cuda")
+    w_down = (
+        torch.randn(LOWRANK, HC_COUNT * HIDDEN_SIZE, dtype=dtype, device="cuda") * 0.02
+    )
+    w_up = (
+        torch.randn(HC_COUNT * HIDDEN_SIZE, LOWRANK, dtype=dtype, device="cuda") * 0.02
+    )
+    return x, w_down, w_up
+
+
+_TOLERANCES = {
+    torch.bfloat16: dict(rtol=1e-2, atol=5e-3),
+    torch.float16: dict(rtol=2e-3, atol=1e-3),
+}
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("num_tokens", [1, 4, 7, _FUSED_MIX_MAX_ROWS])
+def test_fused_hc_mix_matches_reference(dtype, num_tokens):
+    x, w_down, w_up = _make_inputs(num_tokens, dtype)
+    assert fused_hc_mix_supported(x, w_down, w_up)
+    out = fused_hc_mix(x, w_down, w_up, HC_COUNT, HIDDEN_SIZE)
+    ref = _reference_mix(x, w_down, w_up, HC_COUNT, HIDDEN_SIZE)
+    torch.testing.assert_close(out.to(torch.float64), ref, **_TOLERANCES[dtype])
+
+
+def test_fused_hc_mix_no_less_accurate_than_eager():
+    """The fused kernel (fp32 accumulation throughout) must not be farther
+    from the fp64 reference than the eager bf16 chain it replaces."""
+    x, w_down, w_up = _make_inputs(8, torch.bfloat16)
+    ref = _reference_mix(x, w_down, w_up, HC_COUNT, HIDDEN_SIZE)
+    fused = fused_hc_mix(x, w_down, w_up, HC_COUNT, HIDDEN_SIZE)
+    eager = _reference_mix(
+        x, w_down, w_up, HC_COUNT, HIDDEN_SIZE, compute_dtype=torch.bfloat16
+    )
+    fused_err = (fused.to(torch.float64) - ref).abs().max()
+    eager_err = (eager.to(torch.float64) - ref).abs().max()
+    assert fused_err <= eager_err * 1.5 + 1e-6
+
+
+def test_fused_hc_mix_gate_rejects_prefill_rows():
+    x, w_down, w_up = _make_inputs(_FUSED_MIX_MAX_ROWS + 1, torch.bfloat16)
+    assert not fused_hc_mix_supported(x, w_down, w_up)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernel/jit/test_fast_topk.py

@@ -0,0 +1,141 @@
+import pytest
+import torch
+
+from sglang.kernels.ops.elementwise.fast_topk import fast_topk
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+def _check_topk_values(score, lengths, indices, topk, row_starts):
+    """fast_topk leaves order and tie-breaking unspecified,
+    so compare the sorted top-k values rather than index sets."""
+    for b in range(score.shape[0]):
+        start = int(row_starts[b]) if row_starts is not None else 0
+        length = int(lengths[b])
+        section = score[b, start : start + length]
+        row = indices[b]
+        if length <= topk:
+            # naive path: identity indices, then -1 fill
+            assert torch.equal(
+                row[:length].cpu(), torch.arange(length, dtype=torch.int32)
+            )
+            assert (row[length:] == -1).all()
+            continue
+        assert (row >= 0).all(), "long rows must fill every slot"
+        picked = section[row.long()]
+        expected = torch.topk(section, topk).values
+        assert torch.equal(
+            picked.sort(descending=True).values, expected.sort(descending=True).values
+        ), f"row {b}: top-{topk} value multiset mismatch"
+
+
+@pytest.mark.parametrize("topk", [512, 2048])
+@pytest.mark.parametrize(
+    "batch,length",
+    [
+        (1, 4096),
+        (7, 3000),
+        (33, 32768),
+        (128, 2050),
+    ],
+)
+def test_fast_topk_long_rows(topk, batch, length):
+    torch.manual_seed(0)
+    score = torch.randn(batch, length, dtype=torch.float32, device="cuda")
+    lengths = torch.full((batch,), length, dtype=torch.int32, device="cuda")
+
+    indices = fast_topk(score, lengths, topk)
+    _check_topk_values(score, lengths, indices, topk, None)
+
+
+@pytest.mark.parametrize("topk", [512, 2048])
+def test_fast_topk_short_and_mixed_rows(topk):
+    torch.manual_seed(0)
+    max_len = topk + 128
+    batch = 8
+    score = torch.randn(batch, max_len, dtype=torch.float32, device="cuda")
+    # rows shorter than k (naive path), exactly k, and longer than k
+    lens = [1, topk // 3, topk - 1, topk, topk + 1, topk + 7, 17, max_len]
+    lengths = torch.tensor(lens[:batch], dtype=torch.int32, device="cuda")
+
+    indices = fast_topk(score, lengths, topk)
+    _check_topk_values(score, lengths, indices, topk, None)
+
+
+@pytest.mark.parametrize("topk", [512, 2048])
+def test_fast_topk_ragged_with_row_starts(topk):
+    torch.manual_seed(0)
+    batch, width = 16, 8192
+    score = torch.randn(batch, width, dtype=torch.float32, device="cuda")
+    row_starts = torch.randint(0, 2048, (batch,), dtype=torch.int32, device="cuda")
+    lengths = torch.randint(1, 2048, (batch,), dtype=torch.int32, device="cuda")
+    lengths = torch.minimum(lengths, width - row_starts).to(torch.int32)
+    # ensure some rows are longer than k
+    lengths[0] = min(width - int(row_starts[0]), topk + 100)
+
+    indices = fast_topk(score, lengths, topk, row_starts=row_starts)
+    _check_topk_values(score, lengths, indices, topk, row_starts)
+
+
+@pytest.mark.parametrize("topk", [512, 2048])
+def test_fast_topk_row_stride(topk):
+    torch.manual_seed(0)
+    batch, length = 8, 4096
+    base = torch.randn(batch, 2 * length, dtype=torch.float32, device="cuda")
+    score = base[:, :length]  # stride(0) == 2*length, stride(1) == 1
+    lengths = torch.full((batch,), length, dtype=torch.int32, device="cuda")
+
+    indices = fast_topk(score, lengths, topk)
+    _check_topk_values(score, lengths, indices, topk, None)
+
+
+@pytest.mark.parametrize("topk", [512, 2048])
+@pytest.mark.parametrize(
+    "fill",
+    [
+        "binary",  # only 0s and 1s: extreme duplication at the threshold bin
+        "few_levels",  # a handful of distinct levels incl. negatives
+        "constant",  # whole rows of one value
+    ],
+)
+def test_fast_topk_duplicate_heavy(topk, fill):
+    torch.manual_seed(0)
+    batch, length = 16, 8192
+    if fill == "binary":
+        score = torch.randint(0, 2, (batch, length), dtype=torch.float32, device="cuda")
+    elif fill == "few_levels":
+        levels = torch.tensor([-5.0, -1.0, 0.0, 0.5, 2.0], device="cuda")
+        score = levels[torch.randint(0, 5, (batch, length), device="cuda")]
+    else:
+        score = torch.full((batch, length), 3.25, dtype=torch.float32, device="cuda")
+    lengths = torch.full((batch,), length, dtype=torch.int32, device="cuda")
+
+    indices = fast_topk(score, lengths, topk)
+    _check_topk_values(score, lengths, indices, topk, None)
+
+
+@pytest.mark.parametrize("topk", [512, 2048])
+def test_fast_topk_negative_and_zero(topk):
+    torch.manual_seed(0)
+    batch, length = 8, 16384
+    score = torch.randn(batch, length, dtype=torch.float32, device="cuda") * 100
+    score[:, : length // 3] = 0.0  # long zero prefix
+    score[:, length // 3 : length // 2] = -1e30  # very negative block
+    lengths = torch.full((batch,), length, dtype=torch.int32, device="cuda")
+
+    indices = fast_topk(score, lengths, topk)
+    _check_topk_values(score, lengths, indices, topk, None)
+
+
+def test_fast_topk_unsupported_k():
+    score = torch.randn(2, 4096, dtype=torch.float32, device="cuda")
+    lengths = torch.full((2,), 4096, dtype=torch.int32, device="cuda")
+    with pytest.raises(RuntimeError, match="topk"):
+        fast_topk(score, lengths, 1024)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernel/jit/test_grouped_gemma_rmsnorm.py

@@ -0,0 +1,87 @@
+import pytest
+import torch
+
+from sglang.kernels.ops.layernorm.grouped_gemma_rmsnorm import grouped_gemma_rmsnorm
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+def _reference_grouped_gemma_rmsnorm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    group_size: int,
+    eps: float,
+    compute_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Mirrors GroupedGemmaRMSNorm.forward (hyperconnection.py); keep them in sync."""
+    x_float = x.to(compute_dtype)
+    hidden = x_float.shape[-1]
+    x_grouped = x_float.reshape(*x_float.shape[:-1], hidden // group_size, group_size)
+    variance = x_grouped.pow(2).mean(dim=-1, keepdim=True)
+    x_norm = (x_grouped * torch.rsqrt(variance + eps)).flatten(-2)
+    return x_norm * (1.0 + weight.to(compute_dtype))
+
+
+# Tolerances are at the output-dtype quantization floor, measured against the
+# fp64 reference on 4xB300 (sm103): bf16 max rel err 3.9e-3 (1 ulp), fp16
+# 4.9e-4 (0.5 ulp). The kernel computes in fp32 like the eager reference.
+_TOLERANCES = {
+    torch.bfloat16: dict(rtol=5e-3, atol=5e-3),
+    torch.float16: dict(rtol=1e-3, atol=1e-3),
+}
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize(
+    "num_tokens,hidden_size,group_size",
+    [
+        (1, 10240, 2560),  # production shape (HC 4 x 2560)
+        (7, 10240, 2560),
+        (128, 10240, 2560),
+        (33, 1024, 512),
+        (5, 512, 512),  # single group == plain gemma rmsnorm
+        (1024, 2048, 1024),
+    ],
+)
+@pytest.mark.parametrize("eps", [1e-6, 1e-5])
+def test_grouped_gemma_rmsnorm_correctness(
+    dtype, num_tokens, hidden_size, group_size, eps
+):
+    torch.manual_seed(0)
+    x = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda")
+    weight = torch.randn(hidden_size, dtype=dtype, device="cuda") * 0.2
+
+    out = grouped_gemma_rmsnorm(x, weight, group_size, eps)
+    expected = _reference_grouped_gemma_rmsnorm(
+        x, weight, group_size, eps, compute_dtype=torch.float64
+    ).to(dtype)
+
+    torch.testing.assert_close(out, expected, **_TOLERANCES[dtype])
+
+
+def test_grouped_gemma_rmsnorm_out_param():
+    x = torch.randn(64, 10240, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(10240, dtype=torch.bfloat16, device="cuda") * 0.2
+    out = torch.empty_like(x)
+
+    result = grouped_gemma_rmsnorm(x, weight, 2560, 1e-6, out=out)
+    expected = _reference_grouped_gemma_rmsnorm(
+        x, weight, 2560, 1e-6, compute_dtype=torch.float64
+    ).to(x.dtype)
+
+    assert result.data_ptr() == out.data_ptr()
+    torch.testing.assert_close(result, expected, **_TOLERANCES[x.dtype])
+
+
+def test_grouped_gemma_rmsnorm_bad_group_size():
+    x = torch.randn(4, 10240, dtype=torch.bfloat16, device="cuda")
+    weight = torch.zeros(10240, dtype=torch.bfloat16, device="cuda")
+    with pytest.raises(RuntimeError, match="group_size"):
+        grouped_gemma_rmsnorm(x, weight, 1000, 1e-6)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernel/jit/test_hc_combine.py

@@ -0,0 +1,161 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+from sglang.kernels.ops.elementwise.hc_combine import hc_combine
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+HC_COUNT = 4
+HIDDEN_SIZE = 2560
+
+
+def _reference_hc_combine(
+    block_output: torch.Tensor,
+    residual: torch.Tensor,
+    normed_residual: torch.Tensor,
+    inject_weight: torch.Tensor,
+    hc: int,
+    hs: int,
+    compute_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Eager reference mirroring ``GatedResidual._combine_compute``;
+    ``compute_dtype=torch.float64`` is the near-exact reference."""
+    R = residual.to(compute_dtype).unflatten(-1, (hc, hs))
+    gates = 2 * torch.sigmoid(
+        F.linear(normed_residual.to(compute_dtype), inject_weight.to(compute_dtype))
+        / hc
+    )
+    injection = block_output.to(compute_dtype).unsqueeze(-2) * gates.unsqueeze(-1)
+    return (R + injection).flatten(-2)
+
+
+def _make_inputs(
+    num_tokens: int, dtype: torch.dtype, hc: int = HC_COUNT, hs: int = HIDDEN_SIZE
+):
+    torch.manual_seed(0)
+    block_output = torch.randn(num_tokens, hs, dtype=dtype, device="cuda")
+    residual = torch.randn(num_tokens, hc * hs, dtype=dtype, device="cuda")
+    normed_residual = torch.randn(num_tokens, hc * hs, dtype=dtype, device="cuda")
+    inject_weight = torch.randn(hc, hc * hs, dtype=dtype, device="cuda") * 0.02
+    return block_output, residual, normed_residual, inject_weight
+
+
+# Worst case over M in {1, 7, 128, 8192} against the fp64 reference on B300 (sm103):
+# bf16 max rel err 7.8e-3 (1 ulp at a binade edge), fp16 below 1e-3 (1 ulp = 9.8e-4);
+# the residual is fp32 reordering flipping the final rounding.
+_TOLERANCES = {
+    torch.bfloat16: dict(rtol=1e-2, atol=5e-3),
+    torch.float16: dict(rtol=1e-3, atol=1e-3),
+}
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("num_tokens", [1, 7, 128, 8192])
+def test_hc_combine_correctness(dtype, num_tokens):
+    block_output, residual, normed_residual, inject_weight = _make_inputs(
+        num_tokens, dtype
+    )
+
+    out = hc_combine(
+        block_output,
+        residual,
+        normed_residual,
+        inject_weight,
+        HC_COUNT,
+        HIDDEN_SIZE,
+    )
+    expected = _reference_hc_combine(
+        block_output,
+        residual,
+        normed_residual,
+        inject_weight,
+        HC_COUNT,
+        HIDDEN_SIZE,
+        compute_dtype=torch.float64,
+    ).to(dtype)
+
+    torch.testing.assert_close(out, expected, **_TOLERANCES[dtype])
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_hc_combine_out_param(dtype):
+    block_output, residual, normed_residual, inject_weight = _make_inputs(64, dtype)
+    out = torch.empty_like(residual)
+
+    result = hc_combine(
+        block_output,
+        residual,
+        normed_residual,
+        inject_weight,
+        HC_COUNT,
+        HIDDEN_SIZE,
+        out=out,
+    )
+    expected = _reference_hc_combine(
+        block_output,
+        residual,
+        normed_residual,
+        inject_weight,
+        HC_COUNT,
+        HIDDEN_SIZE,
+        compute_dtype=torch.float64,
+    ).to(dtype)
+
+    assert result.data_ptr() == out.data_ptr()
+    torch.testing.assert_close(result, expected, **_TOLERANCES[dtype])
+
+
+def test_hc_combine_3d_input():
+    dtype = torch.bfloat16
+    block_output, residual, normed_residual, inject_weight = _make_inputs(32, dtype)
+    block_output = block_output.reshape(4, 8, HIDDEN_SIZE)
+    residual = residual.reshape(4, 8, HC_COUNT * HIDDEN_SIZE)
+    normed_residual = normed_residual.reshape(4, 8, HC_COUNT * HIDDEN_SIZE)
+
+    out = hc_combine(
+        block_output,
+        residual,
+        normed_residual,
+        inject_weight,
+        HC_COUNT,
+        HIDDEN_SIZE,
+    )
+    expected = _reference_hc_combine(
+        block_output,
+        residual,
+        normed_residual,
+        inject_weight,
+        HC_COUNT,
+        HIDDEN_SIZE,
+        compute_dtype=torch.float64,
+    ).to(dtype)
+
+    assert out.shape == residual.shape
+    torch.testing.assert_close(out, expected, **_TOLERANCES[dtype])
+
+
+def test_hc_combine_bad_hidden_size():
+    dtype = torch.bfloat16
+    block_output, residual, normed_residual, inject_weight = _make_inputs(
+        4,
+        dtype,
+        hc=4,
+        hs=1000,  # 4 * 1000 = 4000, not a multiple of 2048
+    )
+    with pytest.raises(RuntimeError, match="2048"):
+        hc_combine(
+            block_output,
+            residual,
+            normed_residual,
+            inject_weight,
+            4,
+            1000,
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernel/qsa/test_qsa_indexer.py

@@ -0,0 +1,273 @@
+"""Fused QSA indexer-prep kernels must match the eager indexer path bit-for-bit,
+up to rare last-ulp RMSNorm flips (see assert_bit_comparable)."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+from sglang.srt.layers.attention.qsa.kernel import (
+    average_pool_qsa_keys,
+    expand_qsa_block_indices,
+    torch_expand_qsa_block_indices,
+)
+from sglang.srt.layers.attention.qsa.qsa_indexer import QSAIndexer
+from sglang.srt.layers.rotary_embedding.mrope import MRotaryEmbedding
+
+# MRotaryEmbedding reads the exec config bag at init; publish a minimal
+# process context for the bare pytest process.
+from sglang.srt.runtime_context import publish
+from sglang.srt.server_args import ServerArgs
+
+publish(ServerArgs(model_path="dummy"), role="test")
+
+HEAD_DIM = 128
+NUM_Q_HEADS = 4
+RATIO = 4
+HIDDEN = 2560
+EPS = 1e-6
+
+
+def _make_config():
+    return SimpleNamespace(
+        indexer_n_heads=NUM_Q_HEADS,
+        indexer_kv_heads=1,
+        indexer_head_dim=HEAD_DIM,
+        indexer_budget=2048,
+        indexer_compress_ratio=RATIO,
+        hidden_size=HIDDEN,
+        rms_norm_eps=EPS,
+    )
+
+
+def _make_rotary(mrope_section, mrope_interleaved, device, dtype=torch.bfloat16):
+    return MRotaryEmbedding(
+        head_size=HEAD_DIM,
+        rotary_dim=HEAD_DIM,
+        max_position_embeddings=32768,
+        base=1000000,
+        is_neox_style=True,
+        dtype=dtype,
+        mrope_section=mrope_section,
+        mrope_interleaved=mrope_interleaved,
+    )
+
+
+def _make_indexer(rotary, device, dtype=torch.bfloat16):
+    # Build under the model dtype like ModelRunner does; device-only .to()
+    # afterwards so the fp32 cos_sin_cache buffer keeps its dtype.
+    prev_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        indexer = QSAIndexer(
+            _make_config(), layer_id=0, quant_config=None, rotary_emb=rotary
+        )
+        indexer.to(device=device)
+    finally:
+        torch.set_default_dtype(prev_dtype)
+    with torch.no_grad():
+        out_features = (NUM_Q_HEADS + 1) * HEAD_DIM
+        indexer.index_qk_proj.weight.data.copy_(
+            torch.randn(out_features, HIDDEN, device=device, dtype=dtype) * 0.02
+        )
+        for norm in (indexer.q_layernorm, indexer.k_layernorm):
+            w = torch.randn(HEAD_DIM, device=device, dtype=dtype) * 0.1
+            norm._weight_loader(norm.weight, w)
+    return indexer
+
+
+class FakePool:
+    """Minimal stand-in for the QSA KV pool buffers used by the indexer."""
+
+    index_state_dtype = torch.bfloat16
+
+    def __init__(self, num_slots, num_compressed, device, dtype=torch.bfloat16):
+        self.key_state = torch.zeros(num_slots, 1, HEAD_DIM, dtype=dtype, device=device)
+        self.qsa_rope_position_buffer = torch.zeros(
+            num_slots, 3, dtype=torch.int64, device=device
+        )
+        self.compressed = torch.zeros(
+            num_compressed, 1, HEAD_DIM, dtype=dtype, device=device
+        )
+
+    def get_qsa_key_state_buffer(self, layer_id):
+        return self.key_state
+
+    def set_qsa_key_state_buffer(self, layer_id, loc, token_k):
+        self.key_state[loc.long()] = token_k.to(self.key_state.dtype)
+
+    def set_qsa_rope_position_buffer(self, loc, positions):
+        positions = positions.long()
+        if positions.ndim == 1:
+            positions = positions.unsqueeze(0).expand(3, -1)
+        self.qsa_rope_position_buffer[loc.long()] = positions.transpose(0, 1)
+
+    def get_qsa_rope_position_buffer(self, loc):
+        return self.qsa_rope_position_buffer[loc.long()]
+
+    def get_qsa_compressed_k_buffer(self, layer_id):
+        return self.compressed
+
+    def set_qsa_compressed_k_buffer(self, layer_id, loc, compressed_k):
+        self.compressed[loc.long()] = compressed_k.to(self.compressed.dtype)
+
+
+def assert_bit_comparable(actual, expected, max_frac=1e-5, max_abs=0.02):
+    """Eager RMSNorm (flashinfer CuTe DSL) reduces in an unreproducible order,
+    so ~1 row in 30k flips by 1-2 bf16 ulp; max_frac and max_abs bound that."""
+    diff = (actual.float() - expected.float()).abs()
+    mismatches = int((diff > 0).sum())
+    allowed = max(16, int(max_frac * actual.numel()))
+    assert mismatches <= allowed, f"{mismatches} mismatched elements"
+    if mismatches:
+        peak = diff.max().item()
+        assert peak <= max_abs, f"largest deviation {peak} exceeds {max_abs}"
+
+
+def _eager_compress_reference(indexer, pool, group_locs, write_locs):
+    """The pre-fusion compression chain, via the indexer's own helpers."""
+    key_groups = pool.get_qsa_key_state_buffer(0)[group_locs.long()]
+    pooled = average_pool_qsa_keys(key_groups)
+    rope_positions = indexer._rope_from_matrix(
+        pool.get_qsa_rope_position_buffer(group_locs[:, 0])
+    )
+    normalized = indexer.normalize_compressed_keys(pooled, rope_positions)
+    pool.set_qsa_compressed_k_buffer(0, write_locs, normalized)
+
+
+@pytest.mark.parametrize("num_groups", [1, 5, 2000])
+@pytest.mark.parametrize(
+    "mrope_section, mrope_interleaved",
+    [([24, 20, 20], True), ([24, 20, 20], False), (None, False)],
+)
+def test_fused_compress_matches_eager(num_groups, mrope_section, mrope_interleaved):
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    torch.manual_seed(num_groups)
+    rotary = _make_rotary(mrope_section, mrope_interleaved, device, dtype)
+    indexer = _make_indexer(rotary, device, dtype)
+
+    pool_ref = FakePool(8192, 4096, device, dtype)
+    pool_new = FakePool(8192, 4096, device, dtype)
+    pool_new.key_state.copy_(
+        pool_ref.key_state.copy_(
+            torch.randn(8192, 1, HEAD_DIM, device=device, dtype=dtype)
+        )
+    )
+    positions = torch.randint(0, 30000, (8192, 3), device=device)
+    pool_new.qsa_rope_position_buffer.copy_(positions)
+    pool_ref.qsa_rope_position_buffer.copy_(positions)
+
+    # Random groups; slot 0 doubles as the CUDA-graph dummy write target, so
+    # allow repeats there too.
+    group_locs = torch.randint(0, 8192, (num_groups, RATIO), device=device).to(
+        torch.int32
+    )
+    write_locs = torch.randperm(4096, device=device)[:num_groups].to(torch.int32)
+
+    _eager_compress_reference(indexer, pool_ref, group_locs, write_locs)
+    indexer._fused_compress_store(pool_new, group_locs, write_locs)
+
+    assert_bit_comparable(pool_new.compressed, pool_ref.compressed)
+
+
+@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
+def test_expand_block_indices_int_inputs(dtype):
+    device = torch.device("cuda")
+    torch.manual_seed(0)
+    rows, block_topk, token_topk, ratio = 37, 512, 2048, 4
+    query_positions = torch.randint(0, 8000, (rows,), dtype=dtype, device=device)
+    sequence_lengths = (
+        query_positions + torch.randint(1, 9, (rows,), dtype=dtype, device=device)
+    ).to(dtype)
+    # Production contract: top-k only selects blocks inside [0, seq_len//4),
+    # so no selected block ever masks out against sequence_lengths.
+    counts = torch.randint(0, block_topk + 1, (rows,))
+    block_indices = torch.full((rows, block_topk), -1, dtype=torch.int32)
+    seq_lens_host = sequence_lengths.cpu()
+    for r in range(rows):
+        limit = max(int(seq_lens_host[r]) // ratio, 1)
+        count = min(int(counts[r]), limit)
+        if count:
+            block_indices[r, :count] = torch.randperm(limit)[:count].to(torch.int32)
+    block_indices = block_indices.to(device)
+    out = expand_qsa_block_indices(
+        block_indices, query_positions, sequence_lengths, ratio, token_topk
+    )
+    ref = torch_expand_qsa_block_indices(
+        block_indices.cpu(),
+        query_positions.cpu(),
+        sequence_lengths.cpu(),
+        ratio,
+        token_topk,
+    )
+    assert torch.equal(out.cpu(), ref)
+
+
+def test_decode_selection_equivalent():
+    """Last-ulp norm flips must not change the selected blocks:
+    scores are fp32 sums of 128-dim dots, so a 1-ulp flip only matters on exact ties."""
+    from sglang.srt.layers.attention.qsa.kernel import qsa_fast_topk
+    from sglang.srt.layers.attention.qsa.mqa import torch_qsa_mqa_decode
+
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    torch.manual_seed(7)
+    rotary = _make_rotary([24, 20, 20], True, device, dtype)
+    indexer = _make_indexer(rotary, device, dtype)
+
+    batch, max_pages, page_size = 4, 32, 64
+    max_model_len = max_pages * page_size
+    hidden = torch.randn(batch, HIDDEN, device=device, dtype=dtype)
+    positions = (
+        torch.arange(8000, 8000 + batch, device=device)
+        .unsqueeze(0)
+        .expand(3, -1)
+        .contiguous()
+    )
+    qk, _ = indexer.index_qk_proj(hidden)
+
+    # Eager index q.
+    q_ref = indexer.q_layernorm(qk[:, : NUM_Q_HEADS * HEAD_DIM].reshape(-1, HEAD_DIM))
+    q_ref = q_ref.reshape(batch, NUM_Q_HEADS, HEAD_DIM)
+    q_ref = indexer.apply_rope(positions, q_ref)
+
+    # Fused index q.
+    pool = FakePool(64, 4096, device, dtype)
+    cache_loc = torch.arange(1, batch + 1, device=device)
+    q_new, _, stored = indexer.project_qk(
+        hidden, positions, pool=pool, cache_loc=cache_loc
+    )
+    assert stored
+
+    compressed_cache = torch.randn(
+        64, page_size, 1, HEAD_DIM, device=device, dtype=dtype
+    )
+    page_table = torch.arange(max_pages, dtype=torch.int32, device=device).repeat(
+        batch, 1
+    )
+    context_lens = torch.full((batch,), 1500, dtype=torch.int32, device=device)
+
+    def select(q):
+        logits = torch_qsa_mqa_decode(
+            q, compressed_cache, page_table, context_lens, max_model_len
+        )
+        row_starts = torch.zeros_like(context_lens)
+        return qsa_fast_topk(logits, row_starts, context_lens, topk=512)
+
+    idx_ref = select(q_ref)
+    idx_new = select(q_new[:, :NUM_Q_HEADS].contiguous())
+    for row in range(batch):
+        ref_set = set(idx_ref[row][idx_ref[row] >= 0].tolist())
+        new_set = set(idx_new[row][idx_new[row] >= 0].tolist())
+        assert ref_set == new_set, f"row {row}: selection mismatch"
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernel/speculative/test_verify_commit_triton.py

@@ -0,0 +1,81 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
+    fused_commit_track_indices,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+def _reference(accept_index, accept_lens, seq_lens, draft_token_num, track_interval):
+    """Mirrors the eager branch of spec_utils._verify_commit_step_indices."""
+    bs = accept_lens.shape[0]
+    offset = torch.arange(
+        0,
+        bs * draft_token_num,
+        step=draft_token_num,
+        dtype=accept_lens.dtype,
+        device=accept_lens.device,
+    )
+    req_idx = torch.arange(bs, dtype=torch.int64, device=accept_lens.device)
+    last = accept_index[req_idx, (accept_lens - 1).to(torch.int64)] - offset
+    if track_interval <= 0:
+        return last, None
+    pre = seq_lens
+    post = seq_lens + accept_lens
+    mask = pre // track_interval != post // track_interval
+    point = post // track_interval * track_interval
+    ith = torch.clamp(point - pre - 1, min=0).to(torch.int64)
+    cand = accept_index[req_idx, ith] - offset
+    track = torch.where(mask, cand, torch.full_like(cand, -1))
+    return last, track
+
+
+@pytest.mark.parametrize("bs", [1, 3, 48, 257])
+@pytest.mark.parametrize("track_interval", [0, 64])
+@pytest.mark.parametrize("tree_depth", [4, 3])
+def test_verify_commit_steps_matches_eager(bs, track_interval, tree_depth):
+    """The fused kernel must match eager on both outputs near tracking boundaries
+    and when accept_index rows (max_tree_depth) are narrower than draft_token_num."""
+    if not torch.cuda.is_available():
+        pytest.skip("needs CUDA")
+    torch.manual_seed(bs + track_interval + tree_depth)
+    device = "cuda"
+    draft_token_num = 4
+    accept_lens = torch.randint(
+        1, tree_depth + 1, (bs,), device=device, dtype=torch.int32
+    )
+    tree_nodes = torch.argsort(torch.rand(bs, draft_token_num, device=device), dim=1)[
+        :, :tree_depth
+    ]
+    accept_index = (
+        torch.arange(bs, device=device, dtype=torch.int64).unsqueeze(1)
+        * draft_token_num
+        + tree_nodes
+    ).to(torch.int32)
+    # Cluster seq lens around tracking boundaries to exercise the crossing.
+    seq_lens = torch.randint(60, 70, (bs,), device=device, dtype=torch.int64)
+
+    exp_last, exp_track = _reference(
+        accept_index, accept_lens, seq_lens, draft_token_num, track_interval
+    )
+    got_last, got_track = fused_commit_track_indices(
+        accept_index,
+        accept_lens,
+        seq_lens if track_interval > 0 else None,
+        draft_token_num,
+        track_interval,
+    )
+    assert torch.equal(got_last, exp_last)
+    if track_interval > 0:
+        assert torch.equal(got_track, exp_track)
+    else:
+        assert got_track is None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/configs/test_model_config.py

@@ -7,7 +7,9 @@ from sglang.srt.configs.model_config import (
     ModelConfig,
     get_hybrid_layer_ids,
     is_embedding_gemma,
+    resolve_spec_hidden_size,
 )
+from sglang.srt.configs.qwen4_exp import Qwen4ExpTextConfig
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -69,6 +71,27 @@ class TestDraftModelConfig(CustomTestCase):
         self.assertEqual(config.hf_config.num_nextn_predict_layers, 1)
         self.assertEqual(config.hf_text_config.num_nextn_predict_layers, 1)
 
+    def test_qwen4_exp_spec_hidden_size_keeps_hc_width(self):
+        """Qwen4-Exp's MTP draft consumes the hc-flattened target stream,
+        so spec_hidden_size must stay hidden_size * hc_mult; hy_v4 collapses first."""
+        hidden_size, hc_mult = 2560, 4
+        self.assertEqual(Qwen4ExpTextConfig(hc_count=hc_mult).hc_mult, hc_mult)
+        for arch in ("Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLMMTP"):
+            hf_config = SimpleNamespace(architectures=[arch])
+            self.assertEqual(
+                resolve_spec_hidden_size(
+                    hf_config=hf_config, hidden_size=hidden_size, hc_mult=hc_mult
+                ),
+                (hidden_size * hc_mult, hidden_size * hc_mult),
+            )
+        hy_v4 = SimpleNamespace(architectures=["HYV4ForCausalLM"])
+        self.assertEqual(
+            resolve_spec_hidden_size(
+                hf_config=hy_v4, hidden_size=hidden_size, hc_mult=hc_mult
+            ),
+            (hidden_size, None),
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_mamba_unittest.py

@@ -631,6 +631,157 @@ class TestMamba(unittest.TestCase):
 
         return tree, allocator, req_to_token_pool, make_dummy_req
 
+    # Qwen4-Exp's PLE N-gram window is 2 wide (ngram_size=3) and its "no history"
+    # sentinel is the eos id; pick a recognisable one for the tests.
+    NGRAM_CONTEXT_LEN = 2
+    NGRAM_EOS = 248044
+
+    def _setup_pool_with_ngram(self, ngram_context_len: int = NGRAM_CONTEXT_LEN):
+        server_args = ServerArgs(model_path="dummy", page_size=1)
+        server_args._mamba_cache_chunk_size = FLA_CHUNK_SIZE
+        set_global_server_args_for_scheduler(server_args)
+        with envs.SGLANG_MAMBA_SSM_DTYPE.override("bfloat16"):
+            shape = Mamba2StateShape.create(
+                tp_world_size=1,
+                intermediate_size=4096,
+                n_groups=16,
+                num_heads=32,
+                head_dim=128,
+                state_size=128,
+                conv_kernel=4,
+            )
+            cache_params = Mamba2CacheParams(shape=shape, layers=[0])
+        return HybridReqToTokenPool(
+            size=10,
+            mamba_size=20,
+            mamba_spec_state_size=10,
+            max_context_len=128,
+            device=get_device(),
+            enable_memory_saver=False,
+            cache_params=cache_params,
+            mamba_layer_ids=[0],
+            enable_mamba_extra_buffer=False,
+            speculative_num_draft_tokens=3,
+            ngram_context_len=ngram_context_len,
+            ngram_eos_token_id=self.NGRAM_EOS,
+        )
+
+    # Slot-sibling parity: each test below pins one way a mamba slot changes owner.
+
+    def test_slot_siblings_registered(self):
+        """Enabled PLE side states register on the pool that owns the slots;
+        disabled ones stay off so the host-offload payload keeps its legacy shape."""
+        _, _, base_pool, _ = self._setup_tree_and_allocator()
+        # The default hybrid setup has no PLE config: no siblings ride along.
+        self.assertEqual(len(base_pool.mamba_pool._slot_siblings), 0)
+        pool = self._setup_pool_with_ngram()
+        self.assertEqual(len(pool.mamba_pool._slot_siblings), 1)
+
+    def test_ngram_clear_slots_resets_window(self):
+        """A recycled slot must not carry its previous owner's N-gram window;
+        the sibling reset must ride the same deferred ``clear_slots`` call."""
+        pool = self._setup_pool_with_ngram()
+        mamba_pool = pool.mamba_pool
+        ngram = pool.ngram_pool
+
+        victim = pool.mamba_allocator.alloc(1)
+        ngram.context[victim.long()] = 777  # poison, as a real request's history
+        mamba_pool.clear_slots(victim)
+        self.assertTrue(
+            torch.all(ngram.context[victim.long()] == self.NGRAM_EOS),
+            f"clear_slots left a dirty N-gram row: {ngram.context[victim.long()]}",
+        )
+
+    def test_ngram_copy_from_copies_window(self):
+        """copy_from carries the window, so radix cow gets the cached prefix's state."""
+        pool = self._setup_pool_with_ngram()
+        mamba_pool = pool.mamba_pool
+        ngram = pool.ngram_pool
+
+        src = pool.mamba_allocator.alloc(1)
+        dst = pool.mamba_allocator.alloc(1)
+        window = torch.tensor(
+            [[55, 66]], dtype=ngram.context.dtype, device=ngram.context.device
+        )
+        ngram.context[src.long()] = window
+
+        mamba_pool.copy_from(src, dst)
+        self.assertTrue(
+            torch.equal(ngram.context[dst.long()], window),
+            f"copy_from lost the N-gram window: got {ngram.context[dst.long()]}",
+        )
+
+    def test_ngram_cpu_offload_roundtrip(self):
+        """The window survives a host offload round-trip along with mamba state."""
+        pool = self._setup_pool_with_ngram()
+        mamba_pool = pool.mamba_pool
+        ngram = pool.ngram_pool
+
+        indices = pool.mamba_allocator.alloc(2)
+        window = torch.tensor(
+            [[11, 12], [13, 14]],
+            dtype=ngram.context.dtype,
+            device=ngram.context.device,
+        )
+        ngram.context[indices.long()] = window
+
+        saved = mamba_pool.get_cpu_copy(indices)
+        ngram.context[indices.long()] = self.NGRAM_EOS  # simulate slot reuse
+        mamba_pool.load_cpu_copy(saved, indices)
+
+        self.assertTrue(
+            torch.equal(ngram.context[indices.long()], window),
+            f"offload round-trip lost the window: got {ngram.context[indices.long()]}",
+        )
+
+    def test_ngram_pool_absent_keeps_legacy_offload_shape(self):
+        """Disabled pool stays inert: legacy 2-tuple offload payload, no sibling."""
+        pool = self._setup_pool_with_ngram(ngram_context_len=0)
+        self.assertIsNone(pool.ngram_pool.context)
+        self.assertEqual(len(pool.mamba_pool._slot_siblings), 0)
+
+        src = pool.mamba_allocator.alloc(1)
+        payload = pool.mamba_pool.get_cpu_copy(src)
+        self.assertEqual(len(payload), 2)
+        pool.mamba_pool.load_cpu_copy(payload, src)
+
+    def test_mamba_track_aligned_lens_math(self):
+        """Floor division must swallow the scheduler's `aligned + 1` (_force_track_h),
+        or the PLE side states snapshot one token past the mamba state."""
+        from types import SimpleNamespace
+
+        from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+        def aligned_for(chunk_size, track_seqlens, prefix_lens):
+            server_args = ServerArgs(model_path="dummy", page_size=1)
+            server_args._mamba_cache_chunk_size = chunk_size
+            set_global_server_args_for_scheduler(server_args)
+            fake = SimpleNamespace(
+                mamba_track_mask=torch.tensor([True] * len(track_seqlens)),
+                mamba_track_seqlens=torch.tensor(track_seqlens, dtype=torch.int64),
+                extend_prefix_lens=torch.tensor(prefix_lens, dtype=torch.int64),
+            )
+            return ForwardBatch.mamba_track_aligned_lens(fake).tolist()
+
+        # normal: track_seqlens = prefix + extend_input_len
+        self.assertEqual(
+            aligned_for(64, [100 + 64, 100 + 100, 100 + 127], [100, 100, 100]),
+            [64, 64, 64],
+        )
+        # _force_track_h with chunk > 64: track_seqlens = aligned + 1
+        self.assertEqual(aligned_for(128, [100 + 128 + 1], [100]), [128])
+        self.assertEqual(aligned_for(128, [100 + 256 + 1], [100]), [256])
+        # branching point inside the chunk, also handed over as +1
+        self.assertEqual(aligned_for(64, [100 + 64 + 1], [100]), [64])
+        # a masked-off row carries -1 and must come out non-positive, so the
+        # caller's clamp(min=0) routes it harmlessly
+        self.assertLessEqual(aligned_for(64, [-1], [100])[0], 0)
+
+        # restore the chunk size the rest of the suite expects
+        server_args = ServerArgs(model_path="dummy", page_size=1)
+        server_args._mamba_cache_chunk_size = FLA_CHUNK_SIZE
+        set_global_server_args_for_scheduler(server_args)
+
     def test_mamba_pool_cpu_offload(self):
         """MambaPool.get_cpu_copy / load_cpu_copy round-trips conv and temporal state."""
         _, _, req_to_token_pool, _ = self._setup_tree_and_allocator()

test/registered/unit/server_args/test_server_args.py

@@ -105,6 +105,23 @@ _mock_device.start()
 
 
 class TestPrepareServerArgs(CustomTestCase):
+    def test_ple_embedding_offload_rejects_generic_weight_offload(self):
+        for generic_offload in (
+            {"cpu_offload_gb": 1},
+            {"offload_group_size": 1},
+        ):
+            with (
+                self.subTest(generic_offload=generic_offload),
+                self.assertRaisesRegex(
+                    ValueError, "ple-offload-embedding cannot be combined"
+                ),
+            ):
+                ServerArgs(
+                    model_path="dummy",
+                    ple_offload_embedding=True,
+                    **generic_offload,
+                ).resolve_once()
+
     def test_weight_cache_daemon_allows_static_eplb(self):
         args = ServerArgs(
             model_path="dummy",

test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py

@@ -252,7 +252,10 @@ class TestEagleWorkerV2BackendFallback(CustomTestCase):
         existing_backend = object()
         decode_backend = object()
         worker.server_args = _fake_server_args()
-        worker.draft_runner = SimpleNamespace(attn_backend=existing_backend)
+        worker.draft_runner = SimpleNamespace(
+            attn_backend=existing_backend,
+            model_config=SimpleNamespace(hf_config=SimpleNamespace()),
+        )
         worker.topk = 1
         worker.speculative_num_steps = 2
         worker.seed_dsa_topk_from_draft_extend = False
@@ -274,7 +277,10 @@ class TestEagleWorkerV2BackendFallback(CustomTestCase):
         decode_backend = object()
         draft_extend_backend = object()
         worker.server_args = _fake_server_args()
-        worker.draft_runner = SimpleNamespace(attn_backend=existing_backend)
+        worker.draft_runner = SimpleNamespace(
+            attn_backend=existing_backend,
+            model_config=SimpleNamespace(hf_config=SimpleNamespace()),
+        )
         worker.topk = 1
         worker.speculative_num_steps = 2
         worker.seed_dsa_topk_from_draft_extend = True

test/registered/unit/test_model_overrides.py

@@ -88,6 +88,7 @@ class TestModelOverridableWhitelist(CustomTestCase):
                     "uses_mamba_radix_cache",
                     "mamba_radix_cache_strategy",
                     "mamba_full_memory_ratio",
+                    "ple_offload_embedding",
                     "speculative_moe_runner_backend",
                     "speculative_moe_a2a_backend",
                     "disable_shared_experts_fusion",
@@ -605,12 +606,48 @@ class TestGoldenModelOverrides(_IsolatedPublish):
     def test_control_arch_keeps_pristine_dtype(self):
         sa = self._construct("LlamaForCausalLM", "llama")
         self.assertEqual(self._resolved(sa, "dtype"), "auto")
+        self.assertIsNone(self._resolved(sa, "ple_offload_embedding"))
         declared = {f for _s, d in sa._resolved_overrides for f in d}
         self.assertNotIn("dtype", declared)  # no arch declaration for Llama
         # publish still projects the whitelisted leaf with the pristine
         # value: readers only ever read flags.
         self.assertEqual((self._publish(sa), self._leaf("dtype"))[1], "auto")
 
+    def test_qwen4_rejects_pd_and_unified_memory(self):
+        qwen4 = ("Qwen4ExpForConditionalGeneration", "qwen4_exp")
+        for kwargs, message in (
+            ({"disaggregation_mode": "prefill"}, "PD disaggregation"),
+            ({"disaggregation_mode": "decode"}, "PD disaggregation"),
+            ({"enable_unified_memory": True}, "enable-unified-memory"),
+        ):
+            with self.subTest(**kwargs):
+                with self.assertRaisesRegex(ValueError, message):
+                    self._construct(*qwen4, **kwargs)
+
+    def test_qwen4_ple_offload_default(self):
+        qwen4 = ("Qwen4ExpForConditionalGeneration", "qwen4_exp")
+        with override_platform(is_cuda=True):
+            for kwargs, expected in (
+                ({}, True),
+                ({"dtype": "float16"}, False),
+                ({"ple_offload_embedding": False}, False),
+                ({"ple_offload_embedding": False, "cpu_offload_gb": 1}, False),
+            ):
+                with self.subTest(kwargs=kwargs):
+                    self.assertEqual(
+                        self._resolved(
+                            self._construct(*qwen4, **kwargs),
+                            "ple_offload_embedding",
+                        ),
+                        expected,
+                    )
+            with self.assertRaisesRegex(ValueError, "cannot be combined"):
+                self._construct(*qwen4, cpu_offload_gb=1)
+        with override_platform(is_cuda=False, is_hip=True):
+            self.assertFalse(
+                self._resolved(self._construct(*qwen4), "ple_offload_embedding")
+            )
+
     def test_minimax_m2_enables_tf32_matmul(self):
         sa = self._construct("MiniMaxM2ForCausalLM", "llama")
         self.assertTrue(self._resolved(sa, "enable_tf32_matmul"))