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

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

Browse Source
This commit is contained in:
Yanbin Jiang
2026-08-05 15:06:28 -07:00
committed by GitHub
parent 2d27133fcf
commit b9d572ee02
9 changed files with 778 additions and 3223 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/lora/test_experimental_sgl_marlin_alignment.py
-252
View File
@@ -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"]))
test/registered/unit/lora/test_experimental_sgl_marlin_direct_decode.py
-270
View File
@@ -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"]))
test/registered/unit/lora/test_experimental_sgl_marlin_multi_prefill.py
+14 -228
View File
@@ -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(
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.
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-small")
_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
# dictionaries must remain distinct because full and collapsed top-k have
# different flattened token domains.
# Full and collapsed top-k need distinct routing caches: the cache key
# (num_experts, shared_outer, block_m) does not encode the token domain.
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(
("num_slots", "num_tokens"), [(2, 33), (3, 64), (4, 65), (5, 33), (8, 64), (16, 65)]
)
# 64 tokens is a whole number of shrink-stage token blocks; 65 leaves a ragged
# 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__":
test/registered/unit/lora/test_experimental_sgl_marlin_policy.py
+34 -186
View File
@@ -1,125 +1,51 @@
"""CPU-only tests for experimental_sgl_marlin's correctness contract."""
from __future__ import annotations
"""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
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_cuda_ci
from sglang.test.ci.ci_register import register_cpu_ci
register_cuda_ci(
est_time=5,
stage="base-b",
runner_config="1-gpu-small",
disabled="new inkling LoRA test; disabled on CI",
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
# Both spellings of "this server serves adapters". The validator must apply the
# 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)
moe_a2a_backend = overrides.pop("moe_a2a_backend", "none")
values = dict(
"""Runs the real validator on a minimal stand-in; the resolved view is fixed
at `ep_size=4` / `moe_a2a_backend="none"` so the placement chain is reached."""
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(
ep_size=ep_size,
moe_a2a_backend=moe_a2a_backend,
),
types.SimpleNamespace(**server_args),
types.SimpleNamespace(ep_size=4, moe_a2a_backend="none"),
)
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"),
[(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",
"placement",
[
{"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=[
"init_expert_location",
"ep_num_redundant_experts",
"enable_eplb",
"elastic_ep_backend",
"enable_elastic_expert_backup",
"elastic_ep_rejoin",
],
)
def test_lora_ep_rejects_nontrivial_placement_features(setting):
def test_lora_ep_placement_validation(placement):
"""Guards the or-chain: reds when a placement term stops rejecting, or -- if
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:
with pytest.raises(ValueError, match="trivial expert placement"):
_validate_server(**setting)
def test_adapter_paths_implicitly_enable_lora_ep_validation():
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):
with pytest.raises(ValueError, match=message):
_validate(config)
@pytest.mark.parametrize(
"overrides",
[
{"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
)
_validate_server(**launch, **placement)
if __name__ == "__main__":
test/registered/unit/lora/test_experimental_sgl_marlin_runtime_unit.py
+236 -694
View File
File diff suppressed because it is too large Load Diff
test/registered/unit/lora/test_experimental_sgl_marlin_shared_outer_reduce.py
+25 -25
View File
@@ -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(
est_time=20,
stage="base-b",
runner_config="1-gpu-small",
disabled="new inkling LoRA test; disabled on CI",
)
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
_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),
(2, 32, 1.75, 137),
(4, 64, 1.0, 128),
(32, 16, 1.75, 137),
(512, 64, 1.0, 137),
# Dominant decode shape (BLOCK_M=1) with a ragged trailing column tile.
(1, 32, 1.75, 137),
# BLOCK_M=8 with 7 out-of-range rows -- the only param reaching the token
# mask; scale 1.0 pairs with 1.75 so neither value can be hard-coded.
(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
# the capture stream.
# Specialize and initialize CUDA state off 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
# follows stable addresses rather than values observed during capture.
# In-place mutation: replay must follow pointers, not capture-time values.
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)
assert torch.isfinite(output).all().item()
torch.testing.assert_close(output, expected, rtol=_RTOL, atol=_ATOL)
@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)
assert torch.isfinite(output).all().item()
torch.testing.assert_close(output, expected, rtol=_RTOL, atol=_ATOL)
if __name__ == "__main__":
test/registered/unit/lora/test_inkling_linearized_lora_unit.py
+461 -1125
View File
File diff suppressed because it is too large Load Diff
test/registered/unit/lora/test_inkling_lora_normalization_unit.py
-122
View File
@@ -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"]))
test/registered/unit/lora/test_inkling_moe_lora_overlap_unit.py
-313
View File
@@ -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"]))