[test] Re-enable a pruned Inkling LoRA unit-test set (68 -> 9 cases) (#33752)

2026-08-05 15:06:28 -07:00
parent 2d27133fcf
commit b9d572ee02
9 changed files with 778 additions and 3223 deletions
@@ -1,252 +0,0 @@
 """CUDA graph tests for multi-LoRA merged alignment."""
 from __future__ import annotations
 import ast
 import sys
 import types
 from pathlib import Path
 import pytest
 import torch
 from sglang.test.ci.ci_register import register_cuda_ci
 register_cuda_ci(
    est_time=15,
    stage="base-b",
    runner_config="1-gpu-small",
    disabled="new inkling LoRA test; disabled on CI",
 )
 ALIGN_PATH = (
    Path(__file__).resolve().parents[4]
    / "python/sglang/kernels/ops/moe/trtllm_lora_temp/virtual_experts.py"
 )
 def _load_align_function():
    tree = ast.parse(ALIGN_PATH.read_text())
    function = next(
        node
        for node in tree.body
        if isinstance(node, ast.FunctionDef) and node.name == "_align_block_size_jit"
    )
    module = ast.fix_missing_locations(ast.Module(body=[function], type_ignores=[]))
    namespace = {
        "torch": torch,
        "jit_moe_align_block_size": sys.modules[
            "sglang.kernels.ops.moe.moe_align"
        ].moe_align_block_size,
    }
    exec(compile(module, str(ALIGN_PATH), "exec"), namespace)
    return namespace["_align_block_size_jit"]
 def test_experimental_alignment_geometry_and_empty_input(monkeypatch):
    calls = []
    def fake_jit_align(*args):
        (
            topk_ids,
            num_buckets,
            block_size,
            sorted_ids,
            expert_ids,
            total,
            cumsum,
            flag,
        ) = args
        calls.append((num_buckets, cumsum.numel(), flag))
        sorted_ids.fill_(topk_ids.numel())
        expert_ids[:2] = torch.tensor([-1, num_buckets - 2])
        total.fill_(2 * block_size)
    monkeypatch.setitem(
        sys.modules,
        "sglang.kernels.ops.moe.moe_align",
        types.SimpleNamespace(moe_align_block_size=fake_jit_align),
    )
    align = _load_align_function()
    topk_ids = torch.tensor([[-1, 383]], dtype=torch.int32)
    sorted_ids, expert_ids, num_tokens_post_pad = align(topk_ids, 5, 384)
    assert sorted_ids.numel() == 12  # int4-safe capacity above logical 10.
    assert expert_ids.numel() == 3
    assert calls == [(385, 386, True)]
    assert expert_ids[:2].tolist() == [-1, 383]
    assert num_tokens_post_pad.item() == 10
    outputs = align(torch.empty((0, 6), dtype=torch.int32), 16, 384)
    assert [tensor.numel() for tensor in outputs] == [0, 0, 1]
    assert outputs[2].item() == 0
    assert len(calls) == 1
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
 def test_experimental_alignment_cuda_sentinel_and_max_expert():
    from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import (
        _align_block_size_jit,
    )
    topk_ids = torch.tensor([[-1, 383], [0, 0]], device="cuda", dtype=torch.int32)
    _, expert_ids, num_tokens_post_pad = _align_block_size_jit(topk_ids, 16, 384)
    active_experts = expert_ids[: num_tokens_post_pad.item() // 16].cpu().tolist()
    assert sorted(active_experts) == [-1, 0, 383]
 def _assert_shared_outer_merged_align_semantics(
    outputs: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int],
    token_lora_mapping: torch.Tensor,
    *,
    topk: int,
    block_size: int,
    num_slots: int,
 ) -> None:
    """Validate routing without depending on atomic scatter order."""
    (
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        token_lora_mask,
        virtual_num_experts,
    ) = outputs
    mapping = token_lora_mapping.cpu()
    num_routes = mapping.numel() * topk
    total_padded = int(num_tokens_post_padded.item())
    sorted_cpu = sorted_token_ids[:total_padded].cpu()
    experts_cpu = expert_ids[: total_padded // block_size].cpu()
    assert virtual_num_experts == num_slots
    assert torch.equal(token_lora_mask.cpu(), mapping >= 0)
    assert (mapping == -1).any()  # -1 is the runtime base/no-adapter sentinel.
    assert (mapping == 0).any()  # Slot 0 remains a valid adapter slot.
    expected_routes = []
    expected_total_padded = 0
    for slot in range(num_slots):
        slot_tokens = torch.nonzero(mapping == slot, as_tuple=False).flatten()
        slot_routes = (
            slot_tokens[:, None] * topk + torch.arange(topk)[None, :]
        ).flatten()
        expected_routes.extend(slot_routes.tolist())
        route_count = slot_routes.numel()
        expected_total_padded += (
            (route_count + block_size - 1) // block_size
        ) * block_size
    assert total_padded == expected_total_padded
    observed_routes = []
    for block, slot in enumerate(experts_cpu.tolist()):
        assert 0 <= slot < num_slots
        block_routes = sorted_cpu[block * block_size : (block + 1) * block_size]
        real_routes = block_routes[block_routes < num_routes].to(torch.long)
        if real_routes.numel():
            routed_tokens = torch.div(real_routes, topk, rounding_mode="floor")
            assert torch.all(mapping[routed_tokens] == slot)
            observed_routes.extend(real_routes.tolist())
        assert torch.all((block_routes >= 0) & (block_routes <= num_routes))
    assert sorted(observed_routes) == sorted(expected_routes)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
@pytest.mark.parametrize("num_slots", [2, 3, 4])
 def test_multi_slot_shared_outer_merged_align_cuda_graph_parity(num_slots):
    """The fused hot path must replay with current multi-LoRA routing data."""
    from sglang.kernels.ops.moe.trtllm_lora_temp.moe_lora_merged_align import (
        moe_lora_merged_align,
    )
    device = torch.device("cuda")
    num_tokens = 37  # Exercise the multi-LoRA prefill-size routing contract.
    topk = 6
    block_size = 16
    num_experts = 384
    generator = torch.Generator(device=device).manual_seed(9000 + num_slots)
    topk_ids = torch.randint(
        0,
        num_experts,
        (num_tokens, topk),
        device=device,
        dtype=torch.int32,
        generator=generator,
    )
    token_lora_mapping = torch.arange(
        num_tokens, device=device, dtype=torch.int32
    ).remainder(num_slots)
    token_lora_mapping[0] = -1
    token_lora_mapping[-1] = -1
    def invoke(fuse_scatter: bool):
        return moe_lora_merged_align(
            topk_ids,
            token_lora_mapping,
            num_experts,
            shared_outer=True,
            max_loras=num_slots,
            block_size=block_size,
            do_skip=True,
            fuse_scatter=fuse_scatter,
        )
    # Compile both real kernel variants and initialize CUDA state off the
    # capture stream. Production selects the fused variant for this geometry.
    warmup_stream = torch.cuda.Stream()
    warmup_stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(warmup_stream):
        for _ in range(2):
            invoke(fuse_scatter=True)
            invoke(fuse_scatter=False)
    torch.cuda.current_stream().wait_stream(warmup_stream)
    torch.cuda.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        fused_outputs = invoke(fuse_scatter=True)
        split_outputs = invoke(fuse_scatter=False)
    stable_outputs = (*fused_outputs[:4], *split_outputs[:4])
    stable_addresses = tuple(tensor.data_ptr() for tensor in stable_outputs)
    for replay in range(3):
        if replay:
            next_mapping = (
                torch.arange(num_tokens, device=device, dtype=torch.int32)
                .add_(replay)
                .remainder_(num_slots)
            )
            # Move the base/no-adapter rows on every replay while retaining a
            # valid adapter in slot 0.
            next_mapping[replay] = -1
            next_mapping[-replay - 1] = -1
            token_lora_mapping.copy_(next_mapping)
            topk_ids.copy_(torch.roll(topk_ids, shifts=1, dims=1))
        for outputs in (fused_outputs, split_outputs):
            outputs[0].fill_(-12345)
            outputs[1].fill_(-12345)
            outputs[2].fill_(-1)
            outputs[3].fill_(False)
        graph.replay()
        torch.cuda.synchronize()
        assert tuple(tensor.data_ptr() for tensor in stable_outputs) == stable_addresses
        for outputs in (fused_outputs, split_outputs):
            _assert_shared_outer_merged_align_semantics(
                outputs,
                token_lora_mapping,
                topk=topk,
                block_size=block_size,
                num_slots=num_slots,
            )
 if __name__ == "__main__":
    import sys
    sys.exit(pytest.main([__file__, "-v"]))
@@ -1,270 +0,0 @@
 """CUDA graph and numerical tests for direct Inkling decode LoRA kernels."""
 from __future__ import annotations
 import pytest
 import torch
 from sglang.test.ci.ci_register import register_cuda_ci
 register_cuda_ci(est_time=20, stage="base-c", runner_config="4-gpu-b200")
 # Skipped on CI: newly-added inkling LoRA test, disabled pending stabilization.
 pytestmark = pytest.mark.skip(reason="new inkling LoRA test; disabled on CI")
 _B200_AVAILABLE = bool(
    torch.cuda.is_available()
    and torch.version.hip is None
    and torch.cuda.get_device_capability()[0] == 10
 )
 E = 256
 TOPK = 6
 RANK = 32
 DTYPE = torch.bfloat16
 def _make_topk_ids(
    num_tokens: int, *, device: torch.device, offset: int
 ) -> torch.Tensor:
    tokens = torch.arange(num_tokens, device=device, dtype=torch.int32)[:, None]
    routes = torch.arange(TOPK, device=device, dtype=torch.int32)[None, :]
    topk_ids = (tokens * 11 + routes * 17 + offset).remainder(E - 1)
    topk_ids[:, 0] = E - 1
    return topk_ids.contiguous()
 def _gate_reference(
    shared: torch.Tensor,
    gate_b: torch.Tensor,
    topk_ids: torch.Tensor,
    token_lora_mapping: torch.Tensor,
 ) -> torch.Tensor:
    intermediate = gate_b.shape[2] // 2
    gate_width = gate_b.shape[2]
    flat_ids = topk_ids.reshape(-1).to(torch.long)
    flat_slots = token_lora_mapping[:, None].expand(-1, TOPK).reshape(-1).to(torch.long)
    active = flat_slots >= 0
    routed_b = gate_b[flat_slots.clamp_min(0), flat_ids].to(torch.float32)
    if shared.ndim == 3:
        token = torch.arange(shared.shape[1], device=shared.device)
        selected_shared = shared[token_lora_mapping.clamp_min(0).long(), token]
    else:
        selected_shared = shared
    routed_shared = (
        selected_shared[:, None, :].expand(-1, TOPK, -1).reshape(-1, 2 * RANK)
    ).float()
    gate = torch.bmm(routed_b[:, :intermediate], routed_shared[:, :RANK, None]).squeeze(
        -1
    )
    up = torch.bmm(routed_b[:, intermediate:], routed_shared[:, RANK:, None]).squeeze(
        -1
    )
    result = torch.cat((gate, up), dim=1)
    result[~active] = 0
    return result.view(topk_ids.shape[0], TOPK, gate_width)
 def _down_reference(
    activation: torch.Tensor,
    down_a: torch.Tensor,
    topk_ids: torch.Tensor,
    token_lora_mapping: torch.Tensor,
 ) -> torch.Tensor:
    flat_ids = topk_ids.reshape(-1).to(torch.long)
    flat_slots = token_lora_mapping[:, None].expand(-1, TOPK).reshape(-1).to(torch.long)
    active = flat_slots >= 0
    routed_a = down_a[flat_slots.clamp_min(0), flat_ids].to(torch.float32)
    result = torch.bmm(routed_a, activation.to(torch.float32).unsqueeze(-1)).squeeze(-1)
    result[~active] = 0
    return result.view(topk_ids.shape[0], TOPK, RANK)
 def _make_operands(
    num_tokens: int, intermediate: int, num_slots: int, device: torch.device
 ):
    gate_width = 2 * intermediate
    generator = torch.Generator(device=device).manual_seed(9000 + num_tokens)
    shared_shape = (
        (num_slots, num_tokens, RANK) if num_slots > 1 else (num_tokens, RANK)
    )
    gate_half = torch.randn(
        shared_shape, device=device, dtype=DTYPE, generator=generator
    )
    # Deliberately unrelated halves regression-protect the gated split.
    up_half = (
        torch.randn(shared_shape, device=device, dtype=DTYPE, generator=generator)
        * -0.75
        + 0.25
    )
    shared = torch.cat((gate_half, up_half), dim=-1).contiguous()
    gate_b = (
        torch.randn(
            (num_slots, E, gate_width, RANK),
            device=device,
            dtype=DTYPE,
            generator=generator,
        )
        / RANK**0.5
    ).contiguous()
    activation = torch.randn(
        (num_tokens * TOPK, intermediate),
        device=device,
        dtype=DTYPE,
        generator=generator,
    )
    down_a = (
        torch.randn(
            (num_slots, E, RANK, intermediate),
            device=device,
            dtype=DTYPE,
            generator=generator,
        )
        / intermediate**0.5
    ).contiguous()
    topk_ids = _make_topk_ids(num_tokens, device=device, offset=0)
    token_lora_mapping = torch.arange(
        num_tokens, device=device, dtype=torch.int32
    ).remainder(num_slots)
    if num_tokens > 1:
        token_lora_mapping[-1] = -1
    gate_output = torch.empty(
        (num_tokens, TOPK, gate_width), device=device, dtype=DTYPE
    )
    down_output = torch.empty((num_tokens, TOPK, RANK), device=device, dtype=DTYPE)
    return (
        shared,
        gate_b,
        activation,
        down_a,
        topk_ids,
        token_lora_mapping,
        gate_output,
        down_output,
    )
@pytest.mark.skipif(
    not _B200_AVAILABLE,
    reason="direct Inkling decode kernels are currently gated to B200",
 )
@pytest.mark.parametrize(
    ("num_slots", "num_tokens", "intermediate"),
    [(1, 1, 384), (2, 4, 768), (3, 4, 384), (4, 32, 768)],
 )
 def test_direct_decode_cuda_graph_replay_and_base_weights(
    num_tokens: int, intermediate: int, num_slots: int
 ):
    from sglang.srt.lora.marlin_lora_temp.direct_decode import (
        direct_decode_down_shrink,
        direct_decode_gate_expand,
    )
    device = torch.device("cuda")
    (
        shared,
        gate_b,
        activation,
        down_a,
        topk_ids,
        token_lora_mapping,
        gate_output,
        down_output,
    ) = _make_operands(num_tokens, intermediate, num_slots, device)
    def invoke() -> None:
        direct_decode_gate_expand(
            shared, gate_b, topk_ids, token_lora_mapping, gate_output
        )
        direct_decode_down_shrink(
            activation, down_a, topk_ids, token_lora_mapping, down_output
        )
    # Compile and initialize CUDA state outside capture.
    warmup_stream = torch.cuda.Stream()
    warmup_stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(warmup_stream):
        for _ in range(3):
            invoke()
    torch.cuda.current_stream().wait_stream(warmup_stream)
    torch.cuda.synchronize()
    assert int(topk_ids[0, 0]) == 255
    torch.testing.assert_close(
        gate_output.float(),
        _gate_reference(shared, gate_b, topk_ids, token_lora_mapping),
        rtol=0.03,
        atol=0.01,
    )
    torch.testing.assert_close(
        down_output.float(),
        _down_reference(activation, down_a, topk_ids, token_lora_mapping),
        rtol=0.03,
        atol=0.01,
    )
    stable_tensors = (
        shared,
        gate_b,
        activation,
        down_a,
        topk_ids,
        token_lora_mapping,
        gate_output,
        down_output,
    )
    stable_addresses = tuple(tensor.data_ptr() for tensor in stable_tensors)
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        invoke()
    # Mutate every captured input in place. Replay must follow stable addresses
    # and read the new expert ids and values, including expert 255.
    shared.mul_(-0.5).add_(0.125)
    gate_b.mul_(0.75).add_(0.001)
    activation.mul_(0.625).sub_(0.03125)
    down_a.mul_(-0.875).add_(0.0005)
    topk_ids.copy_(_make_topk_ids(num_tokens, device=device, offset=29))
    token_lora_mapping.copy_((token_lora_mapping + 1).remainder(num_slots))
    if num_tokens > 1:
        token_lora_mapping[0] = -1
    gate_output.fill_(float("nan"))
    down_output.fill_(float("nan"))
    graph.replay()
    torch.cuda.synchronize()
    assert tuple(tensor.data_ptr() for tensor in stable_tensors) == stable_addresses
    assert int(topk_ids[0, 0]) == 255
    torch.testing.assert_close(
        gate_output.float(),
        _gate_reference(shared, gate_b, topk_ids, token_lora_mapping),
        rtol=0.03,
        atol=0.01,
    )
    torch.testing.assert_close(
        down_output.float(),
        _down_reference(activation, down_a, topk_ids, token_lora_mapping),
        rtol=0.03,
        atol=0.01,
    )
    # Base/None replay retains the same captured pointers and zeroes the loaded
    # adapter weights in place. Both kernels must fully overwrite their output.
    gate_b.zero_()
    down_a.zero_()
    gate_output.fill_(float("nan"))
    down_output.fill_(float("nan"))
    graph.replay()
    torch.cuda.synchronize()
    assert tuple(tensor.data_ptr() for tensor in stable_tensors) == stable_addresses
    assert torch.count_nonzero(gate_output).item() == 0
    assert torch.count_nonzero(down_output).item() == 0
    assert torch.isfinite(gate_output).all().item()
    assert torch.isfinite(down_output).all().item()
 if __name__ == "__main__":
    import sys
    sys.exit(pytest.main([__file__, "-v"]))
@@ -1,4 +1,6 @@
-"""CUDA parity for the multi-slot shared-outer Marlin prefill factorization."""
+"""CUDA parity for the multi-slot shared-outer Marlin prefill factorization: the
 only numeric guard for the factored prefill schedule under CUDA-graph replay and
 for the gated ``gate_up`` expand gate/up column split; no e2e test selects it."""
 from __future__ import annotations
@@ -7,17 +9,7 @@ import torch
 from sglang.test.ci.ci_register import register_cuda_ci
-register_cuda_ci(
+register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-small")
    est_time=45,
    stage="base-b",
    runner_config="1-gpu-small",
    disabled="fused MoE LoRA-add kernel needs more opt-in shared memory than the ",
 )
 # The fused MoE LoRA-add kernel's shared-memory footprint exceeds the opt-in
 # ceiling of the small-GPU CI runner (~99 KiB on L4) at rank=128, so the
 # generic-fallback parity case OOMs there. Skip this file on CI rather than
 # shrink the production kernel to a small-GPU block config.
 _CUDA_BF16_AVAILABLE = bool(
@@ -44,7 +36,8 @@ def _reference_gate(
    topk_ids: torch.Tensor,
    mapping: torch.Tensor,
 ) -> torch.Tensor:
-    """Materialize shrink/expand with the same BF16 stage boundary."""
+    """Materialize shrink/expand with the same BF16 stage boundary: the gate half
    contracts shrink columns ``[0:rank]``, the up half ``[rank:2*rank]``."""
    active = mapping >= 0
    slots = mapping.clamp_min(0).long()
@@ -108,52 +101,6 @@ def _reference_down(
    return output
 def _reference_generic_delta(
    hidden_states: torch.Tensor,
    lora_a: torch.Tensor,
    lora_b: torch.Tensor,
    topk_ids: torch.Tensor,
    topk_weights: torch.Tensor,
    mapping: torch.Tensor,
    *,
    shared_a: bool,
    shared_b: bool,
    mul_routed_weight: bool,
 ) -> torch.Tensor:
    num_tokens, topk = topk_ids.shape
    output = hidden_states.new_zeros(num_tokens, topk, lora_b.shape[2])
    routed_inputs = hidden_states.shape[0] == topk_ids.numel()
    for token in range(num_tokens):
        slot = int(mapping[token])
        if slot < 0:
            continue
        for route in range(topk):
            expert = int(topk_ids[token, route])
            if expert < 0:
                continue
            row = token * topk + route if routed_inputs else token
            a = lora_a[slot, 0 if shared_a else expert]
            b = lora_b[slot, 0 if shared_b else expert]
            shrink = torch.mv(a.float(), hidden_states[row].float()).to(
                hidden_states.dtype
            )
            rank = b.shape[1]
            if shrink.numel() == 2 * rank:
                output_half = b.shape[0] // 2
                delta = torch.cat(
                    (
                        torch.mv(b[:output_half].float(), shrink[:rank].float()),
                        torch.mv(b[output_half:].float(), shrink[rank:].float()),
                    )
                ).to(hidden_states.dtype)
            else:
                delta = torch.mv(b.float(), shrink.float()).to(hidden_states.dtype)
            if mul_routed_weight:
                delta.mul_(topk_weights[token, route])
            output[token, route] = delta
    return output
 def _run_factored_pipeline(
    *,
    hidden_states: torch.Tensor,
@@ -184,9 +131,8 @@ def _run_factored_pipeline(
    collapsed_ids = mapping.view(num_tokens, 1)
    collapsed_weights = topk_weights[:, :1]
-    # Capture the same four routing domains as the production schedule.  The
+    # Full and collapsed top-k need distinct routing caches: the cache key
-    # dictionaries must remain distinct because full and collapsed top-k have
+    # (num_experts, shared_outer, block_m) does not encode the token domain.
    # different flattened token domains.
    merged_experts_fused_moe_lora_add(
        output=gate_output,
        hidden_states=hidden_states,
@@ -350,13 +296,15 @@ def _run_factored_pipeline(
    not _CUDA_BF16_AVAILABLE,
    reason="multi-prefill parity requires a CUDA GPU with BF16 tensor cores",
 )
-@pytest.mark.parametrize(
+# 64 tokens is a whole number of shrink-stage token blocks; 65 leaves a ragged
-    ("num_slots", "num_tokens"), [(2, 33), (3, 64), (4, 65), (5, 33), (8, 64), (16, 65)]
+# final block so the token mask on the padded tail is actually exercised.
-)
+@pytest.mark.parametrize(("num_slots", "num_tokens"), [(3, 64), (4, 65)])
 def test_multi_shared_outer_prefill_cuda_graph_parity(
    num_slots: int, num_tokens: int
 ) -> None:
-    """Replay full+collapsed routing while adapter selections change in place."""
+    """Guards the factored prefill under CUDA-graph replay with in-place adapter
    reselection; reds when the gated expand reads gate-shrink columns for the up
    half, the broadcast mis-maps route rows to tokens, or slot 0/-1 leaks."""
    device = torch.device("cuda")
    dtype = torch.bfloat16
@@ -451,22 +399,6 @@ def test_multi_shared_outer_prefill_cuda_graph_parity(
        )
    assert full_routing_cache and collapsed_routing_cache
    routing_tensors = tuple(
        tensor
        for cache in (full_routing_cache, collapsed_routing_cache)
        for value in cache.values()
        for tensor in value
    )
    stable_tensors = (
        mapping,
        gate_rank,
        gate_output,
        down_routed_rank,
        down_rank_sum,
        down_output,
        *routing_tensors,
    )
    stable_addresses = tuple(tensor.data_ptr() for tensor in stable_tensors)
    for offset in (1, 3):
        _set_mapping(mapping, num_slots, offset)
@@ -494,7 +426,6 @@ def test_multi_shared_outer_prefill_cuda_graph_parity(
        graph.replay()
        torch.cuda.synchronize()
        assert tuple(tensor.data_ptr() for tensor in stable_tensors) == stable_addresses
        torch.testing.assert_close(gate_output, expected_gate, rtol=0.025, atol=0.025)
        torch.testing.assert_close(down_output, expected_down, rtol=0.025, atol=0.025)
        base_rows = mapping <= 0
@@ -502,151 +433,6 @@ def test_multi_shared_outer_prefill_cuda_graph_parity(
        torch.testing.assert_close(
            down_output[base_rows], base_output[base_rows], rtol=0, atol=0
        )
        assert torch.isfinite(gate_output).all().item()
        assert torch.isfinite(down_output).all().item()
@pytest.mark.skipif(
    not _CUDA_BF16_AVAILABLE,
    reason="generic fallback parity requires a CUDA GPU with BF16 tensor cores",
 )
@pytest.mark.parametrize(
    ("num_slots", "rank", "shared_outer", "ep"),
    [(5, 32, True, False), (8, 128, True, True), (16, 128, False, True)],
 )
 def test_generic_fallback_cuda_graph_parity(
    num_slots: int, rank: int, shared_outer: bool, ep: bool
 ) -> None:
    from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import (
        merged_experts_fused_moe_lora_add,
    )
    device = torch.device("cuda")
    dtype = torch.bfloat16
    num_tokens, topk = 8, 2
    num_experts = 2 if ep else 4
    hidden_size = intermediate_size = 64
    generator = torch.Generator(device=device).manual_seed(3100 + num_slots + rank)
    def randn(*shape: int) -> torch.Tensor:
        return (
            torch.randn(*shape, device=device, dtype=dtype, generator=generator) * 0.03
        )
    gate_a = randn(num_slots, 1 if shared_outer else num_experts, 2 * rank, hidden_size)
    gate_b = randn(num_slots, num_experts, 2 * intermediate_size, rank)
    down_a = randn(num_slots, num_experts, rank, intermediate_size)
    down_b = randn(num_slots, 1 if shared_outer else num_experts, hidden_size, rank)
    hidden_states = randn(num_tokens, hidden_size)
    activation = randn(num_tokens * topk, intermediate_size)
    topk_ids = (
        torch.arange(num_tokens * topk, device=device, dtype=torch.int32)
        .remainder(num_experts)
        .view(num_tokens, topk)
    )
    if ep:
        topk_ids[::2, 1] = -1
    topk_weights = torch.tensor([0.4, 0.6], device=device, dtype=torch.float32).expand(
        num_tokens, -1
    )
    mapping = torch.arange(num_tokens, device=device, dtype=torch.int32).remainder(
        num_slots
    )
    mapping[::5] = -1
    gate_output = torch.empty(
        num_tokens, topk, 2 * intermediate_size, device=device, dtype=dtype
    )
    down_output = torch.empty(num_tokens, topk, hidden_size, device=device, dtype=dtype)
    def run(gate_cache: dict, down_cache: dict) -> None:
        gate_output.zero_()
        merged_experts_fused_moe_lora_add(
            output=gate_output,
            hidden_states=hidden_states,
            lora_a=gate_a,
            lora_b=gate_b,
            topk_ids=topk_ids,
            topk_weights=topk_weights,
            token_lora_mapping=mapping,
            mul_routed_weight=False,
            experts_shared_outer_loras_a=shared_outer,
            experts_shared_outer_loras_b=False,
            routing_cache=gate_cache,
            fuse_add_to_output=False,
            use_direct_expand_add=True,
            local_num_experts=num_experts,
        )
        down_output.zero_()
        merged_experts_fused_moe_lora_add(
            output=down_output,
            hidden_states=activation,
            lora_a=down_a,
            lora_b=down_b,
            topk_ids=topk_ids,
            topk_weights=topk_weights,
            token_lora_mapping=mapping,
            mul_routed_weight=True,
            experts_shared_outer_loras_a=False,
            experts_shared_outer_loras_b=shared_outer,
            routing_cache=down_cache,
            use_direct_expand_add=False,
            local_num_experts=num_experts,
            zero_intermediate=ep and shared_outer,
        )
    def assert_expected() -> None:
        torch.testing.assert_close(
            gate_output,
            _reference_generic_delta(
                hidden_states,
                gate_a,
                gate_b,
                topk_ids,
                topk_weights,
                mapping,
                shared_a=shared_outer,
                shared_b=False,
                mul_routed_weight=False,
            ),
            rtol=2e-2,
            atol=2e-2,
        )
        torch.testing.assert_close(
            down_output,
            _reference_generic_delta(
                activation,
                down_a,
                down_b,
                topk_ids,
                topk_weights,
                mapping,
                shared_a=False,
                shared_b=shared_outer,
                mul_routed_weight=True,
            ),
            rtol=2e-2,
            atol=2e-2,
        )
    run({}, {})
    torch.cuda.synchronize()
    assert_expected()
    gate_cache: dict = {}
    down_cache: dict = {}
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        run(gate_cache, down_cache)
    for offset in (0, 1):
        mapping.copy_(
            torch.arange(num_tokens, device=device, dtype=torch.int32)
            .add(offset)
            .remainder(num_slots)
        )
        mapping[(torch.arange(num_tokens, device=device) + offset) % 5 == 0] = -1
        graph.replay()
        assert_expected()
 if __name__ == "__main__":
@@ -1,125 +1,51 @@
-"""CPU-only tests for experimental_sgl_marlin's correctness contract."""
+"""Guards the LoRA + non-trivial expert-placement startup rejection; reds when a
-
+placement term leaves the or-chain, or when a `--lora-paths`-only launch stops
-from __future__ import annotations
+counting as LoRA, silently landing adapter deltas on remapped experts."""
 import types
 import pytest
 from sglang.srt.lora.marlin_lora_temp.policy import (
    use_post_reduce_down_delta,
    validate_experimental_sgl_marlin_contract,
    validate_experimental_sgl_marlin_server_args,
 )
-from sglang.srt.lora.trtllm_lora_temp.specialized_expand import _get_gated_a_half
+from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.ci.ci_register import register_cuda_ci
-register_cuda_ci(
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
-    est_time=5,
+
-    stage="base-b",
+
-    runner_config="1-gpu-small",
+# Both spellings of "this server serves adapters". The validator must apply the
-    disabled="new inkling LoRA test; disabled on CI",
+# same guards to each; only the second one exercises the tri-state.
 LORA_LAUNCH_FORMS = (
    {"enable_lora": True, "lora_paths": []},
    {"enable_lora": None, "lora_paths": ["adapter=/tmp/adapter"]},
 )
 def _config(**overrides):
    values = dict(
        activation="silu",
        is_gated=True,
        gemm1_alpha=None,
        gemm1_clamp_limit=None,
        swiglu_limit=None,
        apply_router_weight_on_input=False,
        no_combine=False,
        num_experts=256,
        num_local_experts=256,
    )
    values.update(overrides)
    return types.SimpleNamespace(**values)
 def _validate(config=None, **overrides):
    values = dict(
        runner_config=config or _config(),
        moe_ep_size=1,
        device_capability=(9, 0),
    )
    values.update(overrides)
    return validate_experimental_sgl_marlin_contract(**values)
 def _validate_server(**overrides):
-    ep_size = overrides.pop("ep_size", 4)
+    """Runs the real validator on a minimal stand-in; the resolved view is fixed
-    moe_a2a_backend = overrides.pop("moe_a2a_backend", "none")
+    at `ep_size=4` / `moe_a2a_backend="none"` so the placement chain is reached."""
-    values = dict(
+    server_args = dict(
        enable_lora=True,
        lora_paths=[],
        lora_use_virtual_experts=True,
        lora_backend="triton",
        init_expert_location="trivial",
        ep_num_redundant_experts=0,
        enable_eplb=False,
        elastic_ep_backend=None,
        enable_elastic_expert_backup=False,
        elastic_ep_rejoin=False,
        experts_shared_outer_loras=False,
        max_lora_rank=64,
        lora_backend="triton",
    )
-    values.update(overrides)
+    server_args.update(overrides)
    return validate_experimental_sgl_marlin_server_args(
-        types.SimpleNamespace(**values),
+        types.SimpleNamespace(**server_args),
-        types.SimpleNamespace(
+        types.SimpleNamespace(ep_size=4, moe_a2a_backend="none"),
            ep_size=ep_size,
            moe_a2a_backend=moe_a2a_backend,
        ),
    )
 def test_supported_contract_passes():
    _validate()
 def test_supported_ep_contract_passes():
    _validate(config=_config(num_local_experts=64), moe_ep_size=4)
@pytest.mark.parametrize(
-    ("enable_lora", "lora_paths"),
+    "placement",
    [(None, []), (False, []), (False, ["ignored=/tmp/adapter"])],
 )
 def test_base_only_ep_preserves_stock_marlin_placement_support(enable_lora, lora_paths):
    _validate_server(
        enable_lora=enable_lora,
        lora_paths=lora_paths,
        init_expert_location="random",
        ep_num_redundant_experts=1,
        enable_eplb=True,
        elastic_ep_backend="mooncake",
        enable_elastic_expert_backup=True,
        elastic_ep_rejoin=True,
    )
 def test_base_only_ep_rejects_unsupported_a2a():
    with pytest.raises(ValueError, match="moe-a2a-backend none"):
        _validate_server(enable_lora=False, moe_a2a_backend="deepep")
 def test_lora_rejects_non_triton_backend():
    with pytest.raises(ValueError, match="requires --lora-backend triton"):
        _validate_server(lora_backend="csgmv")
    # base-only servers are free to pick any dense backend
    _validate_server(enable_lora=False, lora_backend="csgmv")
@pytest.mark.parametrize("ep_size", [1, 4])
 def test_lora_requires_virtual_experts(ep_size):
    with pytest.raises(ValueError, match="lora-use-virtual-experts"):
        _validate_server(ep_size=ep_size, lora_use_virtual_experts=False)
@pytest.mark.parametrize(
    "setting",
    [
        {"init_expert_location": "random"},
        {"ep_num_redundant_experts": 1},
@@ -128,100 +54,22 @@ def test_lora_requires_virtual_experts(ep_size):
        {"enable_elastic_expert_backup": True},
        {"elastic_ep_rejoin": True},
    ],
-)
+    ids=[
-def test_lora_ep_rejects_nontrivial_placement_features(setting):
+        "init_expert_location",
-    with pytest.raises(ValueError, match="trivial expert placement"):
+        "ep_num_redundant_experts",
-        _validate_server(**setting)
+        "enable_eplb",
-
+        "elastic_ep_backend",
-
+        "enable_elastic_expert_backup",
-def test_adapter_paths_implicitly_enable_lora_ep_validation():
+        "elastic_ep_rejoin",
    with pytest.raises(ValueError, match="trivial expert placement"):
        _validate_server(
            enable_lora=None,
            lora_paths=["adapter=/tmp/adapter"],
            enable_eplb=True,
        )
 def test_supported_lora_ep_passes():
    _validate_server(experts_shared_outer_loras=True, max_lora_rank=64)
@pytest.mark.parametrize("rank", [65, 128, 256])
 def test_shared_outer_lora_ep_allows_generic_rank_fallback(rank):
    _validate_server(experts_shared_outer_loras=True, max_lora_rank=rank)
@pytest.mark.parametrize(
    ("config", "message"),
    [
        (_config(activation="relu2"), "activation must be 'silu'"),
        (_config(is_gated=False), "only gated SwiGLU"),
        (_config(gemm1_alpha=1.0), "gemm1_alpha"),
        (_config(gemm1_clamp_limit=7.0), "gemm1_clamp_limit"),
        (_config(swiglu_limit=7.0), "swiglu_limit"),
        (
            _config(apply_router_weight_on_input=True),
            "apply_router_weight_on_input",
        ),
        (_config(no_combine=True), "no_combine"),
    ],
 )
-def test_rejects_unimplemented_activation_semantics(config, message):
+def test_lora_ep_placement_validation(placement):
-    with pytest.raises(ValueError, match=message):
+    """Guards the or-chain: reds when a placement term stops rejecting, or -- if
-        _validate(config)
+    only the `lora_paths` form reds -- when the implicit-enable tri-state
-
+    collapsed and adapter-path launches bypass the LoRA guards."""
-
+    for launch in LORA_LAUNCH_FORMS:
-@pytest.mark.parametrize(
+        with pytest.raises(ValueError, match="trivial expert placement"):
-    "overrides",
+            _validate_server(**launch, **placement)
    [
        {"config": _config(num_local_experts=128)},
        {"config": _config(num_local_experts=63), "moe_ep_size": 4},
        {"config": _config(num_experts=255, num_local_experts=64), "moe_ep_size": 4},
        {"moe_ep_size": 0},
    ],
 )
 def test_rejects_incoherent_expert_parallelism(overrides):
    config = overrides.get("config")
    validation_overrides = {k: v for k, v in overrides.items() if k != "config"}
    with pytest.raises(ValueError, match="moe_ep_size|num_local_experts"):
        _validate(config, **validation_overrides)
 def test_rejects_pre_hopper_gpu():
    with pytest.raises(ValueError, match="compute capability 9.0 or newer"):
        _validate(device_capability=(8, 0))
 def test_direct_expand_always_splits_gated_gate_up_a():
    assert _get_gated_a_half(intermediate_width=64, rank=32, output_width=768) == 384
    assert _get_gated_a_half(intermediate_width=32, rank=32, output_width=6144) == 0
 def test_direct_expand_rejects_invalid_intermediate_width():
    with pytest.raises(ValueError, match="intermediate width"):
        _get_gated_a_half(intermediate_width=48, rank=32, output_width=768)
@pytest.mark.parametrize(
    ("run_lora", "scale", "num_tokens", "expected"),
    [
        (True, 1.0, 1, True),
        (True, 1.0, 2048, True),
        (True, 1.0, 2049, False),
        (True, 0.5, 32, False),
        (False, 1.0, 32, False),
    ],
 )
 def test_post_reduce_down_policy(run_lora, scale, num_tokens, expected):
    assert (
        use_post_reduce_down_delta(
            run_lora=run_lora,
            routed_scaling_factor=scale,
            num_tokens=num_tokens,
        )
        is expected
    )
 if __name__ == "__main__":
@@ -1,4 +1,6 @@
-"""CUDA parity tests for the fused shared-outer Marlin decode reduction."""
+"""Guards the fused shared-outer Marlin decode reduction against the three-launch
 fallback it replaces; reds when down-B orientation, router weighting, the routing
 scale, a ragged-tile mask, or capture-time value baking is wrong."""
 from __future__ import annotations
@@ -7,12 +9,7 @@ import torch
 from sglang.test.ci.ci_register import register_cuda_ci
-register_cuda_ci(
+register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
    est_time=20,
    stage="base-b",
    runner_config="1-gpu-small",
    disabled="new inkling LoRA test; disabled on CI",
 )
 _CUDA_BF16_AVAILABLE = bool(
@@ -21,6 +18,10 @@ _CUDA_BF16_AVAILABLE = bool(
    and torch.cuda.get_device_capability()[0] >= 8
 )
 # Two BF16 ulps (2**-7 each), plus an absolute floor for cancelled outputs.
 _RTOL = 1.6e-2
 _ATOL = 1e-3
 def _reference_reduce(
    routed_base: torch.Tensor,
@@ -56,6 +57,8 @@ def _mapped_reference_reduce(
    token_lora_mapping: torch.Tensor,
    routed_scaling_factor: float,
 ) -> torch.Tensor:
    """Same reduction with a per-token adapter slot; ``-1`` rows stay base-only."""
    dtype = routed_base.dtype
    base_sum = routed_base.float().sum(dim=1).mul(routed_scaling_factor).to(dtype)
    rank_sum = (
@@ -81,11 +84,11 @@ def _mapped_reference_reduce(
@pytest.mark.parametrize(
    ("num_tokens", "rank", "routed_scaling_factor", "hidden_width"),
    [
-        (1, 16, 1.0, 128),
+        # Dominant decode shape (BLOCK_M=1) with a ragged trailing column tile.
-        (2, 32, 1.75, 137),
+        (1, 32, 1.75, 137),
-        (4, 64, 1.0, 128),
+        # BLOCK_M=8 with 7 out-of-range rows -- the only param reaching the token
-        (32, 16, 1.75, 137),
+        # mask; scale 1.0 pairs with 1.75 so neither value can be hard-coded.
-        (512, 64, 1.0, 137),
+        (65, 64, 1.0, 137),
    ],
 )
 def test_fused_base_shared_lora_reduce_cuda_graph_parity(
@@ -146,6 +149,7 @@ def test_fused_base_shared_lora_reduce_cuda_graph_parity(
    )
    output = torch.empty((num_tokens, hidden_width), device=device, dtype=dtype)
    # Production launch geometry, so the tuned tiles and their ragged tails run.
    block_m, block_k = fused_base_shared_lora_reduce_config(num_tokens)
    def invoke() -> None:
@@ -160,8 +164,7 @@ def test_fused_base_shared_lora_reduce_cuda_graph_parity(
            block_k=block_k,
        )
-    # Compile the rank/block specialization and initialize CUDA state away from
+    # Specialize and initialize CUDA state off the capture stream.
    # the capture stream.
    warmup_stream = torch.cuda.Stream()
    warmup_stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(warmup_stream):
@@ -170,16 +173,12 @@ def test_fused_base_shared_lora_reduce_cuda_graph_parity(
    torch.cuda.current_stream().wait_stream(warmup_stream)
    torch.cuda.synchronize()
    stable_tensors = (routed_base, routed_rank, topk_weights, shared_b, output)
    stable_addresses = tuple(tensor.data_ptr() for tensor in stable_tensors)
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        invoke()
    for replay in range(2):
-        # Mutate every captured operand in place so replay proves that the graph
+        # In-place mutation: replay must follow pointers, not capture-time values.
        # follows stable addresses rather than values observed during capture.
        routed_base.mul_(0.75).add_(0.002 * (replay + 1))
        routed_rank.mul_(-0.5).add_(0.001 * (replay + 1))
        topk_weights.copy_(torch.roll(topk_weights, shifts=1, dims=1))
@@ -189,7 +188,6 @@ def test_fused_base_shared_lora_reduce_cuda_graph_parity(
        graph.replay()
        torch.cuda.synchronize()
        assert tuple(tensor.data_ptr() for tensor in stable_tensors) == stable_addresses
        expected = _reference_reduce(
            routed_base,
            routed_rank,
@@ -197,8 +195,7 @@ def test_fused_base_shared_lora_reduce_cuda_graph_parity(
            shared_b,
            routed_scaling_factor,
        )
-        torch.testing.assert_close(output, expected, rtol=0.03, atol=0.004)
+        torch.testing.assert_close(output, expected, rtol=_RTOL, atol=_ATOL)
        assert torch.isfinite(output).all().item()
@pytest.mark.skipif(
@@ -210,8 +207,12 @@ def test_fused_base_shared_lora_reduce_cuda_graph_parity(
    [(2, 1, 128), (3, 32, 137)],
 )
 def test_fused_base_mapped_shared_lora_reduce_cuda_graph_parity(
-    num_tokens: int, num_slots: int, hidden_width: int
+    num_slots: int, num_tokens: int, hidden_width: int
 ):
    """Guards the device-side per-token slot lookup; reds when the adapted-token
    predicate degrades to always-true (``-1`` rows gain a delta) or the lookup
    moves host-side (replay serves the capture-time mapping)."""
    from sglang.srt.lora.marlin_lora_temp.shared_outer import (
        fused_base_mapped_shared_lora_reduce,
    )
@@ -298,8 +299,7 @@ def test_fused_base_mapped_shared_lora_reduce_cuda_graph_parity(
        token_lora_mapping,
        1.75,
    )
-    torch.testing.assert_close(output, expected, rtol=0.03, atol=0.004)
+    torch.testing.assert_close(output, expected, rtol=_RTOL, atol=_ATOL)
    assert torch.isfinite(output).all().item()
 if __name__ == "__main__":
@@ -1,122 +0,0 @@
 """Hermetic regression tests for Inkling shared-sink LoRA normalization."""
 from __future__ import annotations
 import ast
 from pathlib import Path
 from types import SimpleNamespace
 import pytest
 import torch
 from sglang.test.ci.ci_register import register_cuda_ci
 register_cuda_ci(
    est_time=1,
    stage="base-b",
    runner_config="1-gpu-small",
    disabled="new inkling LoRA test; disabled on CI",
 )
 REPO_ROOT = Path(__file__).resolve().parents[4]
 LORA_PATH = REPO_ROOT / "python/sglang/srt/lora/lora.py"
 def _load_normalizer_class():
    tree = ast.parse(LORA_PATH.read_text())
    source_class = next(
        node
        for node in tree.body
        if isinstance(node, ast.ClassDef) and node.name == "LoRAAdapter"
    )
    method_names = {"_normalize_shared_expert_moe", "normalize_gate_up_proj"}
    methods = [
        node
        for node in source_class.body
        if isinstance(node, ast.FunctionDef) and node.name in method_names
    ]
    assert {method.name for method in methods} == method_names
    test_class = ast.ClassDef(
        name="_NormalizerUnderTest",
        bases=[],
        keywords=[],
        body=methods,
        decorator_list=[],
    )
    namespace = {"Dict": dict, "re": __import__("re"), "torch": torch}
    exec(
        compile(
            ast.fix_missing_locations(ast.Module(body=[test_class], type_ignores=[])),
            str(LORA_PATH),
            "exec",
        ),
        namespace,
    )
    return namespace[test_class.name]
 def _normalizer(num_shared: int = 2, *, text_config: bool = False):
    normalizer = _load_normalizer_class()()
    config = SimpleNamespace(
        architectures=None if text_config else ["InklingForConditionalGeneration"],
        model_type="inkling_text" if text_config else "inkling",
        n_shared_experts=num_shared,
    )
    normalizer.base_hf_config = config
    return normalizer
@pytest.mark.parametrize("text_config", [False, True])
 def test_proj_named_shared_sink_factors_gain_the_expert_axis(text_config):
    n, rank, hidden, intermediate = 2, 3, 5, 7
    prefix = "model.layers.0.mlp.shared_experts"
    gate_a = torch.arange(rank * hidden).reshape(rank, hidden)
    gate_b = torch.arange(n * 2 * intermediate * rank).reshape(
        n * 2 * intermediate, rank
    )
    down_a = torch.arange(rank * n * intermediate).reshape(rank, n * intermediate)
    down_b = torch.arange(hidden * rank).reshape(hidden, rank)
    weights = {
        f"{prefix}.gate_up_proj.lora_A.weight": gate_a,
        f"{prefix}.gate_up_proj.lora_B.weight": gate_b,
        f"{prefix}.down_proj.lora_A.weight": down_a,
        f"{prefix}.down_proj.lora_B.weight": down_b,
    }
    normalizer = _normalizer(n, text_config=text_config)
    normalizer._normalize_shared_expert_moe(weights)
    normalizer.normalize_gate_up_proj(list(weights), weights)
    torch.testing.assert_close(
        weights[f"{prefix}.gate_up_proj.lora_A.weight"],
        gate_a.unsqueeze(0).repeat(1, 2, 1),
    )
    torch.testing.assert_close(
        weights[f"{prefix}.gate_up_proj.lora_B.weight"],
        gate_b.reshape(n, 2 * intermediate, rank),
    )
    torch.testing.assert_close(
        weights[f"{prefix}.down_proj.lora_A.weight"],
        down_a.reshape(rank, n, intermediate).transpose(0, 1).contiguous(),
    )
    torch.testing.assert_close(
        weights[f"{prefix}.down_proj.lora_B.weight"], down_b.unsqueeze(0)
    )
 def test_named_per_expert_outer_factor_is_not_collapsed_to_shared_outer():
    name = "model.layers.0.mlp.shared_experts.1.gate_up_proj.lora_A.weight"
    weight = torch.arange(15).reshape(3, 5)
    weights = {name: weight}
    _normalizer()._normalize_shared_expert_moe(weights)
    torch.testing.assert_close(weights[name], weight)
    assert weights[name].dim() == 2
 if __name__ == "__main__":
    import sys
    sys.exit(pytest.main([__file__, "-v"]))
@@ -1,313 +0,0 @@
 """Hermetic stream-order checks for Inkling shared/routed overlap."""
 from __future__ import annotations
 import ast
 import sys
 import types
 from pathlib import Path
 from types import SimpleNamespace
 import pytest
 from sglang.test.ci.ci_register import register_cuda_ci
 register_cuda_ci(
    est_time=5,
    stage="base-b",
    runner_config="1-gpu-small",
    disabled="refactor-fragile source-parsing unit test; skipped on CI",
 )
 # Skipped on CI: this hermetic check re-parses the InklingMoE forward source and
 # pins its exact stream-order, so it breaks on unrelated refactors of that
 REPO_ROOT = Path(__file__).resolve().parents[4]
 MOE_PATH = REPO_ROOT / "python/sglang/srt/models/inkling_common/moe.py"
 INKLING_DENSE_PATH = (
    REPO_ROOT / "python/sglang/srt/lora/trtllm_lora_temp/inkling_dense.py"
 )
 class _Flag:
    def __init__(self, value: bool):
        self.value = value
    def get(self) -> bool:
        return self.value
 class _Stream:
    def __init__(self, name: str, events: list[str]):
        self.name = name
        self.events = events
    def wait_stream(self, other: _Stream) -> None:
        self.events.append(f"{self.name}.wait({other.name})")
 class _StreamContext:
    def __init__(self, cuda, stream: _Stream):
        self.cuda = cuda
        self.stream = stream
        self.previous = None
    def __enter__(self):
        self.previous = self.cuda.current
        self.cuda.current = self.stream
        self.cuda.events.append(f"enter({self.stream.name})")
    def __exit__(self, *_):
        self.cuda.events.append(f"exit({self.stream.name})")
        self.cuda.current = self.previous
 class _Cuda:
    def __init__(self, events: list[str]):
        self.events = events
        self.current = _Stream("main", events)
    def current_stream(self) -> _Stream:
        return self.current
    def stream(self, stream: _Stream) -> _StreamContext:
        return _StreamContext(self, stream)
 class _Tensor:
    def __init__(self, name: str, events: list[str], *, tokens: int = 1):
        self.name = name
        self.events = events
        self.shape = (tokens, 8)
        self.dtype = "bf16"
        self.is_cuda = True
    def record_stream(self, stream: _Stream) -> None:
        self.events.append(f"{self.name}.record({stream.name})")
    def __add__(self, other: _Tensor) -> _Tensor:
        self.events.append(f"add({self.name},{other.name})")
        return _Tensor("sum", self.events)
 def _load_forward(fake_torch, capture: bool = False):
    tree = ast.parse(MOE_PATH.read_text())
    source_class = next(
        node
        for node in tree.body
        if isinstance(node, ast.ClassDef) and node.name == "InklingMoE"
    )
    forward = next(
        node
        for node in source_class.body
        if isinstance(node, ast.FunctionDef) and node.name == "forward"
    )
    test_class = ast.ClassDef(
        name="_InklingMoEForwardUnderTest",
        bases=[],
        keywords=[],
        body=[forward],
        decorator_list=[],
    )
    namespace = {
        "ForwardBatch": object,
        "envs": SimpleNamespace(
            SGLANG_OPT_USE_INKLING_MULTI_STREAM_OVERLAP=_Flag(True)
        ),
        # capture gating: overlap only inside cuda-graph capture
        "get_is_capture_mode": lambda: capture,
        "get_ar_buffer": lambda *_: None,
        "get_tensor_model_parallel_group": lambda: SimpleNamespace(world_size=1),
        "lora_compatible_layout_enabled": lambda: True,
        "torch": fake_torch,
    }
    exec(
        compile(
            ast.fix_missing_locations(ast.Module(body=[test_class], type_ignores=[])),
            str(MOE_PATH),
            "exec",
        ),
        namespace,
    )
    return namespace[test_class.name]
 def _load_lora_overlap_policy():
    tree = ast.parse(INKLING_DENSE_PATH.read_text())
    policy = next(
        node
        for node in tree.body
        if isinstance(node, ast.FunctionDef)
        and node.name == "allow_inkling_moe_two_stream"
    )
    namespace = {}
    exec(
        compile(
            ast.fix_missing_locations(ast.Module(body=[policy], type_ignores=[])),
            str(INKLING_DENSE_PATH),
            "exec",
        ),
        namespace,
    )
    return namespace[policy.name]
 def _make_moe(events: list[str], cuda: _Cuda, capture: bool = False):
    fake_torch = SimpleNamespace(Tensor=_Tensor, cuda=cuda)
    moe = _load_forward(fake_torch, capture)()
    moe.alt_stream = _Stream("alt", events)
    moe.shared_experts = SimpleNamespace(
        lora_backend=SimpleNamespace(batch_info=SimpleNamespace(has_active_lora=True))
    )
    moe.experts = SimpleNamespace()
    moe._clone_fused_sink_input = False
    moe._fused_ar_shared = False
    moe.gate = lambda x: (
        _Tensor("topk_weights", events),
        _Tensor("topk_ids", events),
        _Tensor("gammas", events),
        None,
    )
    def forward_shared(x, gammas):
        events.append(f"shared({cuda.current.name})")
        return _Tensor("shared_out", events)
    def forward_routed(*_):
        assert cuda.current.name == "main"
        events.append("routed(main)")
        return _Tensor("routed_out", events)
    moe._forward_shared = forward_shared
    moe._forward_routed = forward_routed
    return moe
 def _install_lora_policy(monkeypatch, *, main_alloc: bool, capture: bool = False):
    lora_envs = SimpleNamespace(SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC=_Flag(main_alloc))
    monkeypatch.setitem(
        sys.modules,
        "sglang.srt.lora.trtllm_lora_temp.environ",
        types.SimpleNamespace(lora_envs=lora_envs),
    )
    monkeypatch.setitem(
        sys.modules,
        "sglang.srt.model_executor.runner_utils.capture_mode",
        types.SimpleNamespace(get_is_capture_mode=lambda: capture),
    )
    monkeypatch.setitem(
        sys.modules,
        "sglang.srt.lora.trtllm_lora_temp.inkling_dense",
        types.SimpleNamespace(allow_inkling_moe_two_stream=_load_lora_overlap_policy()),
    )
@pytest.mark.parametrize("tokens", [1, 32])
 def test_direct_sink_keeps_decode_overlap(monkeypatch, tokens):
    events: list[str] = []
    cuda = _Cuda(events)
    _install_lora_policy(monkeypatch, main_alloc=True, capture=True)
    moe = _make_moe(events, cuda, capture=True)
    moe.forward(_Tensor("x", events, tokens=tokens), reduce=False)
    assert events == [
        "x.record(alt)",
        "gammas.record(alt)",
        "alt.wait(main)",
        "enter(alt)",
        "shared(alt)",
        "exit(alt)",
        "routed(main)",
        "main.wait(alt)",
        "shared_out.record(main)",
        "add(routed_out,shared_out)",
    ]
 def test_lora_prefill_stays_serial(monkeypatch):
    # Even when capture would allow overlap, the M>32 LoRA policy forces serial.
    events: list[str] = []
    cuda = _Cuda(events)
    _install_lora_policy(monkeypatch, main_alloc=True, capture=True)
    moe = _make_moe(events, cuda, capture=True)
    moe.forward(_Tensor("x", events, tokens=33), reduce=False)
    assert events == [
        "routed(main)",
        "shared(main)",
        "add(routed_out,shared_out)",
    ]
 def test_captured_prefill_stays_serial_even_base_only(monkeypatch):
    # Capture forces has_lora_work (one schedule for every replay), so
    # prefill-sized batches (>32 tokens) are serial even with no live adapter.
    events: list[str] = []
    cuda = _Cuda(events)
    _install_lora_policy(monkeypatch, main_alloc=True, capture=True)
    moe = _make_moe(events, cuda, capture=True)
    moe.shared_experts.lora_backend.batch_info.has_active_lora = False
    moe.forward(_Tensor("x", events, tokens=33), reduce=False)
    assert events == [
        "routed(main)",
        "shared(main)",
        "add(routed_out,shared_out)",
    ]
 def test_eager_forward_stays_serial_even_base_only(monkeypatch):
    # overlap is gated on cuda-graph capture; eager forwards are serial.
    events: list[str] = []
    cuda = _Cuda(events)
    _install_lora_policy(monkeypatch, main_alloc=False, capture=False)
    moe = _make_moe(events, cuda, capture=False)
    moe.shared_experts.lora_backend.batch_info.has_active_lora = False
    moe.forward(_Tensor("x", events, tokens=33), reduce=False)
    assert events == [
        "routed(main)",
        "shared(main)",
        "add(routed_out,shared_out)",
    ]
 def test_lora_overlap_stays_serial_without_main_alloc(monkeypatch):
    events: list[str] = []
    cuda = _Cuda(events)
    _install_lora_policy(monkeypatch, main_alloc=False, capture=True)
    moe = _make_moe(events, cuda, capture=True)
    moe.forward(_Tensor("x", events), reduce=False)
    assert events == [
        "routed(main)",
        "shared(main)",
        "add(routed_out,shared_out)",
    ]
 def test_capture_keeps_lora_schedule_without_active_adapter(monkeypatch):
    events: list[str] = []
    cuda = _Cuda(events)
    _install_lora_policy(monkeypatch, main_alloc=False, capture=True)
    moe = _make_moe(events, cuda, capture=True)
    moe.shared_experts.lora_backend.batch_info.has_active_lora = False
    moe.forward(_Tensor("x", events), reduce=False)
    assert events == [
        "routed(main)",
        "shared(main)",
        "add(routed_out,shared_out)",
    ]
 if __name__ == "__main__":
    import sys
    sys.exit(pytest.main([__file__, "-v"]))