From b9d572ee0245f222ddcde362ea6ab4fc2149f180 Mon Sep 17 00:00:00 2001 From: Yanbin Jiang Date: Wed, 5 Aug 2026 15:06:28 -0700 Subject: [PATCH] [test] Re-enable a pruned Inkling LoRA unit-test set (68 -> 9 cases) (#33752) --- .../test_experimental_sgl_marlin_alignment.py | 252 --- ...t_experimental_sgl_marlin_direct_decode.py | 270 --- ...t_experimental_sgl_marlin_multi_prefill.py | 242 +-- .../test_experimental_sgl_marlin_policy.py | 220 +-- ...st_experimental_sgl_marlin_runtime_unit.py | 932 +++------- ...rimental_sgl_marlin_shared_outer_reduce.py | 50 +- .../lora/test_inkling_linearized_lora_unit.py | 1600 +++++------------ .../test_inkling_lora_normalization_unit.py | 122 -- .../test_inkling_moe_lora_overlap_unit.py | 313 ---- 9 files changed, 778 insertions(+), 3223 deletions(-) delete mode 100644 test/registered/unit/lora/test_experimental_sgl_marlin_alignment.py delete mode 100644 test/registered/unit/lora/test_experimental_sgl_marlin_direct_decode.py delete mode 100644 test/registered/unit/lora/test_inkling_lora_normalization_unit.py delete mode 100644 test/registered/unit/lora/test_inkling_moe_lora_overlap_unit.py diff --git a/test/registered/unit/lora/test_experimental_sgl_marlin_alignment.py b/test/registered/unit/lora/test_experimental_sgl_marlin_alignment.py deleted file mode 100644 index 3dcd9d5c0..000000000 --- a/test/registered/unit/lora/test_experimental_sgl_marlin_alignment.py +++ /dev/null @@ -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"])) diff --git a/test/registered/unit/lora/test_experimental_sgl_marlin_direct_decode.py b/test/registered/unit/lora/test_experimental_sgl_marlin_direct_decode.py deleted file mode 100644 index bf3991157..000000000 --- a/test/registered/unit/lora/test_experimental_sgl_marlin_direct_decode.py +++ /dev/null @@ -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"])) diff --git a/test/registered/unit/lora/test_experimental_sgl_marlin_multi_prefill.py b/test/registered/unit/lora/test_experimental_sgl_marlin_multi_prefill.py index b1ba5da79..7bb98b355 100644 --- a/test/registered/unit/lora/test_experimental_sgl_marlin_multi_prefill.py +++ b/test/registered/unit/lora/test_experimental_sgl_marlin_multi_prefill.py @@ -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__": diff --git a/test/registered/unit/lora/test_experimental_sgl_marlin_policy.py b/test/registered/unit/lora/test_experimental_sgl_marlin_policy.py index 70eda0568..3927edba3 100644 --- a/test/registered/unit/lora/test_experimental_sgl_marlin_policy.py +++ b/test/registered/unit/lora/test_experimental_sgl_marlin_policy.py @@ -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}, ], -) -def test_lora_ep_rejects_nontrivial_placement_features(setting): - 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"), + ids=[ + "init_expert_location", + "ep_num_redundant_experts", + "enable_eplb", + "elastic_ep_backend", + "enable_elastic_expert_backup", + "elastic_ep_rejoin", ], ) -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 - ) +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(**launch, **placement) if __name__ == "__main__": diff --git a/test/registered/unit/lora/test_experimental_sgl_marlin_runtime_unit.py b/test/registered/unit/lora/test_experimental_sgl_marlin_runtime_unit.py index 92fe28e18..117d97c47 100644 --- a/test/registered/unit/lora/test_experimental_sgl_marlin_runtime_unit.py +++ b/test/registered/unit/lora/test_experimental_sgl_marlin_runtime_unit.py @@ -1,8 +1,10 @@ -"""CPU-only runtime-flow tests for experimental_sgl_marlin.""" +"""CPU-only runtime-flow tests for the experimental Marlin MoE-LoRA runner: +recording fakes make its kernel ordering, buffer ownership and expert-parallel +index localization observable without CUDA.""" from __future__ import annotations -import importlib.util +import ast import sys import types from pathlib import Path @@ -11,21 +13,127 @@ from types import SimpleNamespace import pytest import torch -from sglang.test.ci.ci_register import register_cuda_ci +import sglang +from sglang.srt.lora import trtllm_lora_temp +from sglang.srt.lora.marlin_lora_temp import moe_runner +from sglang.srt.lora.trtllm_lora_temp import environ as trtllm_lora_environ +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=2, suite="base-a-test-cpu") + +_SGLANG_ROOT = Path(sglang.__file__).resolve().parent + +# ``torch.Tensor.zero_`` is monkeypatched per flow; hold the pristine method so +# repeated patching inside one session can never chain the recording wrappers. +_ORIGINAL_TENSOR_ZERO = torch.Tensor.zero_ -REPO_ROOT = Path(__file__).resolve().parents[4] -LORA_TEMP_ROOT = REPO_ROOT / "python/sglang/srt/lora" -MARLIN_RUNNER_PATH = LORA_TEMP_ROOT / "marlin_lora_temp/moe_runner.py" -MARLIN_POLICY_PATH = LORA_TEMP_ROOT / "marlin_lora_temp/policy.py" -TWO_STREAM_PATH = LORA_TEMP_ROOT / "trtllm_lora_temp/__init__.py" +# --------------------------------------------------------------------------- +# Import resolution +# --------------------------------------------------------------------------- + + +def _module_source_path(dotted_name: str) -> Path | None: + """Locate an ``sglang.*`` module's source file without importing it.""" + + parts = dotted_name.split(".") + assert parts[0] == "sglang" + base = _SGLANG_ROOT.joinpath(*parts[1:]) + module_file = base.with_suffix(".py") + if module_file.is_file(): + return module_file + package_file = base / "__init__.py" + if package_file.is_file(): + return package_file + return None + + +def _module_level_bindings(path: Path) -> set[str]: + """Names bound at module scope, including inside module-level branches.""" + + names: set[str] = set() + + def walk(body) -> None: + for node in body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + names.add(node.name) + continue + if isinstance(node, (ast.Import, ast.ImportFrom)): + for alias in node.names: + names.add(alias.asname or alias.name.split(".")[0]) + continue + if isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)): + for child in ast.walk(node): + if isinstance(child, ast.Name) and isinstance(child.ctx, ast.Store): + names.add(child.id) + continue + # Module-level control flow (if/try/with/for) can bind names too; + # function and class bodies are deliberately not descended into. + for attribute in ("body", "orelse", "finalbody"): + walk(getattr(node, attribute, ()) or ()) + for handler in getattr(node, "handlers", ()): + walk(handler.body) + + walk(ast.parse(path.read_text()).body) + return names + + +def test_cuda_only_kernel_imports_resolve(): + """Guards resolvability of the module-level ``if _is_cuda:`` kernel imports, + which never run on CPU and are shadowed by this file's fakes; reds when one + goes stale (#32884 repointed a bad ``moe_sum_reduce_triton`` source).""" + + tree = ast.parse(Path(moe_runner.__file__).read_text()) + guarded = [ + node + for node in tree.body + if isinstance(node, ast.If) + and isinstance(node.test, ast.Name) + and node.test.id == "_is_cuda" + ] + assert len(guarded) == 1, ( + "expected exactly one module-level `if _is_cuda:` kernel-import block in " + "moe_runner.py; this test can no longer see the imports it guards" + ) + + imports: list[tuple[str, str]] = [] + for node in ast.walk(guarded[0]): + assert not isinstance(node, ast.Import), ( + "plain `import X` inside the `if _is_cuda:` block is not covered by " + "this resolution check; extend it" + ) + if isinstance(node, ast.ImportFrom): + assert node.level == 0 and node.module, ( + "relative imports inside the `if _is_cuda:` block are not " + "covered by this resolution check; extend it" + ) + imports.extend((node.module, alias.name) for alias in node.names) + assert imports + + third_party = sorted( + {module for module, _ in imports if not module.startswith("sglang.")} + ) + assert third_party == ["sgl_kernel"], ( + "a new non-sglang kernel import appeared in the `if _is_cuda:` block " + f"({third_party}); its symbols cannot be resolved from sglang sources, " + "so extend this test deliberately" + ) + + unresolved: list[str] = [] + for module, symbol in imports: + if not module.startswith("sglang."): + continue + path = _module_source_path(module) + if path is None: + unresolved.append(f"{module} (no such module)") + elif symbol not in _module_level_bindings(path): + unresolved.append(f"{module}.{symbol} (not bound at module scope)") + assert not unresolved, f"unresolvable CUDA-only kernel imports: {unresolved}" + + +# --------------------------------------------------------------------------- +# Runtime-flow harness +# --------------------------------------------------------------------------- def _stub_module(monkeypatch, name: str, **attributes): @@ -46,119 +154,42 @@ def _stub_module(monkeypatch, name: str, **attributes): return module -def _load_file(monkeypatch, name: str, path: Path): - parts = name.split(".") - for end in range(1, len(parts)): - package_name = ".".join(parts[:end]) - if package_name not in sys.modules: - package = types.ModuleType(package_name) - package.__path__ = [] - monkeypatch.setitem(sys.modules, package_name, package) - spec = importlib.util.spec_from_file_location(name, path) - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, name, module) - spec.loader.exec_module(module) - if len(parts) > 1: - parent = sys.modules[".".join(parts[:-1])] - monkeypatch.setattr(parent, parts[-1], module, raising=False) - return module - - -def _load_marlin_runner(monkeypatch, name: str): - _stub_module(monkeypatch, "sglang.srt.utils", is_cuda=lambda: False) - _load_file( - monkeypatch, - "sglang.srt.lora.marlin_lora_temp.policy", - MARLIN_POLICY_PATH, - ) - return _load_file(monkeypatch, name, MARLIN_RUNNER_PATH) - - -@pytest.mark.parametrize(("tokens", "expected"), [(256, True), (257, False)]) -def test_two_stream_token_threshold_is_inclusive(monkeypatch, tokens, expected): - lora_envs = SimpleNamespace( - SGLANG_TWO_STREAM_MAX_TOKENS=SimpleNamespace(get=lambda: 256) - ) - _stub_module(monkeypatch, "sglang.srt.environ", envs=SimpleNamespace()) - _stub_module( - monkeypatch, - "sglang.srt.lora.trtllm_lora_temp.environ", - lora_envs=lora_envs, - ) - module = _load_file(monkeypatch, "_two_stream_under_test", TWO_STREAM_PATH) - assert module.is_two_stream_active(torch.empty(tokens, 1)) is expected - - -@pytest.mark.parametrize( - ("combined_rank", "rank", "expected"), - [(128, 64, True), (128, 128, False), (256, 64, False)], -) -def test_two_stream_dense_lora_rank_falls_back( - monkeypatch, combined_rank, rank, expected -): - lora_envs = SimpleNamespace( - SGLANG_TWO_STREAM_MAX_TOKENS=SimpleNamespace(get=lambda: 256) - ) - _stub_module(monkeypatch, "sglang.srt.environ", envs=SimpleNamespace()) - _stub_module( - monkeypatch, - "sglang.srt.lora.trtllm_lora_temp.environ", - lora_envs=lora_envs, - ) - module = _load_file(monkeypatch, "_two_stream_rank_under_test", TWO_STREAM_PATH) - assert ( - module.supports_two_stream_dense_lora( - torch.empty(1, combined_rank, 1), torch.empty(1, 1, rank) - ) - is expected - ) - - class _CombineInput: + """Stand-in for ``StandardCombineInput``.""" + def __init__(self, hidden_states): self.hidden_states = hidden_states class _DispatchOutput: + """Stand-in for ``StandardDispatchOutput`` (the runner's isinstance gate).""" + def __init__(self, hidden_states, topk_output): self.hidden_states = hidden_states self.topk_output = topk_output -def _run_marlin_policy( +def _run_marlin_flow( monkeypatch, *, tokens: int, - master: bool = True, two_stream: bool = False, capture: bool = False, active_lora: bool = True, base_mapping: bool = False, - direct_decode: bool = False, ep: bool = False, slots: int = 1, rank: int = 1, - shared_outer: bool = True, - base_value: float = 0.0, ): - module = _load_marlin_runner(monkeypatch, "_marlin_runner_under_test") - # The hermetic runner uses tiny CPU tensors. Explicitly emulate the exact - # B200/Inkling eligibility gate so these tests exercise the fused schedule. - module._use_fused_shared_outer_tail = ( - lambda _info, _hidden, num_tokens, _hidden_size, _topk: num_tokens <= 512 - ) - module._use_direct_decode_kernels = lambda *_args, **_kwargs: direct_decode + """Run the real runner against recording fakes and return what it did + (scheduling, stream ownership, buffer identity -- never kernel numerics).""" calls = SimpleNamespace( merged=[], - split_gate_checks=0, - weighted_rank_sums=0, fused_tails=0, - direct_gate=0, - direct_down=0, marlin_is_ep=[], + marlin_output_ids=[], align_num_experts=[], - cache3_was_zero=None, schedule=[], zeroed=[], event_records=[], @@ -210,65 +241,29 @@ def _run_marlin_policy( monkeypatch.setattr(torch.cuda, "stream", _StreamContext) monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capture) - original_zero = torch.Tensor.zero_ - - def tracked_zero(intermediate, *args, **kwargs): - item = ("zero", stream_state["current"].name, id(intermediate)) + def tracked_zero(tensor, *args, **kwargs): + item = ("zero", stream_state["current"].name, id(tensor)) calls.schedule.append(item) calls.zeroed.append(item) - return original_zero(intermediate, *args, **kwargs) + return _ORIGINAL_TENSOR_ZERO(tensor, *args, **kwargs) monkeypatch.setattr(torch.Tensor, "zero_", tracked_zero) def merged_experts_fused_moe_lora_add(**kwargs): intermediate = kwargs.get("intermediate_buffer") stage = kwargs.get("stage") - calls.schedule.append( - ( - "merged", - stage, - stream_state["current"].name, - id(intermediate) if intermediate is not None else None, - ) - ) + buffer_id = id(intermediate) if intermediate is not None else None + stream = stream_state["current"].name + calls.schedule.append(("merged", stage, stream, buffer_id)) calls.merged.append( { "stage": stage, - "stream": stream_state["current"].name, - "intermediate_shape": ( - tuple(kwargs["intermediate_buffer"].shape) - if kwargs.get("intermediate_buffer") is not None - else None - ), - "broadcast": kwargs.get("broadcast_intermediate", False), - "prewarm_a": kwargs.get("prewarm_a_routing", True), - "prewarm_b": kwargs.get("prewarm_b_routing", True), - "topk_shape": tuple(kwargs["topk_ids"].shape), - "cache_id": id(kwargs.get("routing_cache")), - "shared_a": kwargs["experts_shared_outer_loras_a"], + "stream": stream, + "intermediate_id": buffer_id, "shared_b": kwargs["experts_shared_outer_loras_b"], - "fuse_add": kwargs.get("fuse_add_to_output", True), - "direct_expand": kwargs.get("use_direct_expand_add", False), - "mul_routed_weight": kwargs["mul_routed_weight"], "zero_intermediate": kwargs.get("zero_intermediate", False), - "mapping": kwargs["token_lora_mapping"].clone(), - "intermediate_id": ( - id(intermediate) if intermediate is not None else None - ), } ) - if stage == "expand": - if kwargs.get("fuse_add_to_output", True): - active = kwargs["token_lora_mapping"] >= 0 - kwargs["output"][active].add_(1) - else: - kwargs["output"].fill_(0) - if stage == "shrink": - return intermediate - - def is_two_stream_active(_hidden_states): - calls.split_gate_checks += 1 - return two_stream _stub_module( monkeypatch, @@ -276,17 +271,6 @@ def _run_marlin_policy( StandardCombineInput=_CombineInput, StandardDispatchOutput=_DispatchOutput, ) - _stub_module( - monkeypatch, - "sglang.srt.lora.trtllm_lora_temp", - get_lora_side_stream=lambda: side_stream, - is_two_stream_active=is_two_stream_active, - ) - _stub_module( - monkeypatch, - "sglang.srt.lora.trtllm_lora_temp.environ", - experimental_lora_enabled=lambda: master, - ) _stub_module( monkeypatch, "sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts", @@ -297,6 +281,11 @@ def _run_marlin_policy( "sglang.srt.model_executor.runner", get_is_capture_mode=lambda: capture, ) + monkeypatch.setattr(trtllm_lora_temp, "get_lora_side_stream", lambda: side_stream) + monkeypatch.setattr( + trtllm_lora_temp, "is_two_stream_active", lambda _hidden: two_stream + ) + monkeypatch.setattr(trtllm_lora_environ, "experimental_lora_enabled", lambda: True) def fake_align(_topk_ids, _block_size, num_experts, **_kwargs): calls.align_num_experts.append(num_experts) @@ -306,51 +295,25 @@ def _run_marlin_policy( torch.ones(1, dtype=torch.int32), ) - module.moe_align_block_size = fake_align - module.marlin_make_workspace = lambda *_args, **_kwargs: None - module.get_scalar_type = lambda *_args, **_kwargs: None - - def fake_marlin_gemm(_x, output, *_args, **_kwargs): - calls.marlin_is_ep.append(_kwargs["is_ep"]) - if len(calls.marlin_is_ep) == 2: - calls.cache3_was_zero = bool(torch.count_nonzero(output) == 0) + def fake_marlin_gemm(_x, output, *_args, **kwargs): + calls.marlin_is_ep.append(kwargs["is_ep"]) + calls.marlin_output_ids.append(id(output)) calls.schedule.append( - ("marlin", stream_state["current"].name, len(calls.schedule)) + ("marlin", stream_state["current"].name, len(calls.marlin_is_ep)) ) - output.fill_(base_value if len(calls.marlin_is_ep) == 2 else 0) + # ``fill_``, never ``zero_``: the zero-fill contracts are asserted by + # watching ``Tensor.zero_`` calls, so the fakes must not forge one. + output.fill_(0) return output - module.moe_wna16_marlin_gemm = fake_marlin_gemm - def fake_silu_and_mul_add_delta(_x, _delta, output): calls.schedule.append(("activation", stream_state["current"].name)) output.fill_(0) - module.silu_and_mul_add_delta = fake_silu_and_mul_add_delta - module.silu_and_mul = lambda _x, output: output.fill_(0) - - def fake_triton_reduce(_input, output, _scale): - calls.schedule.append(("reduce", stream_state["current"].name)) - output.copy_(_input.sum(dim=1) * _scale) - - module.moe_sum_reduce_triton = fake_triton_reduce - - def fake_weighted_rank_sum(routed_rank, weights, output, scale, *, block_m): - calls.weighted_rank_sums += 1 - calls.schedule.append( - ("weighted", stream_state["current"].name, id(routed_rank)) - ) - output.copy_( - (routed_rank * weights.to(routed_rank.dtype).unsqueeze(-1)).sum(dim=1) - * scale - ) - - module.weighted_topk_rank_sum = fake_weighted_rank_sum - def fake_fused_tail( routed_base, - routed_rank, - weights, + _routed_rank, + _weights, _shared_b, output, scale, @@ -364,26 +327,31 @@ def _run_marlin_policy( ) output.copy_(routed_base.sum(dim=1) * scale) - module.fused_base_shared_lora_reduce = fake_fused_tail - module.fused_base_shared_lora_reduce_config = lambda _tokens: (1, 32) + for name, fake in ( + ("moe_align_block_size", fake_align), + ("marlin_make_workspace", lambda *_args, **_kwargs: None), + ("get_scalar_type", lambda *_args, **_kwargs: None), + ("moe_wna16_marlin_gemm", fake_marlin_gemm), + ("silu_and_mul_add_delta", fake_silu_and_mul_add_delta), + ("fused_base_shared_lora_reduce", fake_fused_tail), + ("fused_base_shared_lora_reduce_config", lambda _tokens: (1, 32)), + # The real gate needs a bf16 SM100 tensor, so CPU can never reach it. + ( + "_use_fused_shared_outer_tail", + lambda _info, _hidden, num_tokens, _hidden_size, _topk: num_tokens <= 512, + ), + ): + monkeypatch.setattr(moe_runner, name, fake, raising=False) - def fake_direct_gate(_shared, _weight, _topk_ids, _mapping, output): - calls.direct_gate += 1 - calls.schedule.append(("direct_gate", stream_state["current"].name)) - output.fill_(0) + # Per-flow capture-event list: the module-level one would accumulate across + # cases now that the module is imported once instead of re-loaded per test. + monkeypatch.setattr(moe_runner, "_MARLIN_LORA_OVERLAP_EVENTS", []) - def fake_direct_down(_activation, _weight, _topk_ids, _mapping, output): - calls.direct_down += 1 - calls.schedule.append(("direct_down", stream_state["current"].name)) - output.fill_(0) - - module.direct_decode_gate_expand = fake_direct_gate - module.direct_decode_down_shrink = fake_direct_down - - hidden_size, num_experts, expert_size, topk, max_rank = 2, 1, 16, 2, rank + hidden_size, num_experts, expert_size, topk = 2, 1, 16, 2 hidden_states = torch.zeros(tokens, hidden_size) topk_ids = torch.zeros(tokens, topk, dtype=torch.int32) if ep: + # Non-local routes arrive already masked out. topk_ids[:, 1] = -1 topk_weights = torch.ones(tokens, topk) dispatch_output = _DispatchOutput( @@ -405,523 +373,75 @@ def _run_marlin_policy( w2_qzeros=None, w2_g_idx=None, w2_g_idx_sort_indices=None, - expert_map=None, - global_num_experts=num_experts, weight_bits=4, is_k_full=True, ) lora_info = SimpleNamespace( lora_use_virtual_experts=True, - max_lora_rank=max_rank, + max_lora_rank=rank, has_active_lora=active_lora, - gate_up_lora_a_weights=torch.zeros( - slots, - 1 if shared_outer else num_experts, - 2 * max_rank, - hidden_size, - ), - gate_up_lora_b_weights=torch.zeros( - slots, num_experts, 2 * expert_size, max_rank - ), - down_lora_a_weights=torch.zeros(slots, num_experts, max_rank, expert_size), - down_lora_b_weights=torch.zeros( - slots, - 1 if shared_outer else num_experts, - hidden_size, - max_rank, - ), + gate_up_lora_a_weights=torch.zeros(slots, 1, 2 * rank, hidden_size), + gate_up_lora_b_weights=torch.zeros(slots, num_experts, 2 * expert_size, rank), + down_lora_a_weights=torch.zeros(slots, num_experts, rank, expert_size), + down_lora_b_weights=torch.zeros(slots, 1, hidden_size, rank), token_lora_mapping=( torch.full((tokens,), -1, dtype=torch.int32) if base_mapping else torch.arange(tokens, dtype=torch.int32).remainder(slots) ), - experts_shared_outer_loras=shared_outer, + experts_shared_outer_loras=True, ) runner_config = SimpleNamespace( activation="silu", routed_scaling_factor=1.0, - num_experts=num_experts, + # EP: this rank owns `num_experts` of `2 * num_experts` global experts. + num_experts=2 * num_experts if ep else num_experts, num_local_experts=num_experts, ) - if ep: - runner_config.num_experts = 2 * num_experts - result = module.fused_experts_experimental_sgl_marlin_lora( + result = moe_runner.fused_experts_experimental_sgl_marlin_lora( dispatch_output, quant_info, runner_config, lora_info ) assert result.hidden_states.shape == hidden_states.shape calls.input_ptr = hidden_states.data_ptr() calls.result_ptr = result.hidden_states.data_ptr() - calls.result = result.hidden_states - calls.capture_event_count = len(module._MARLIN_LORA_OVERLAP_EVENTS) + calls.capture_event_count = len(moe_runner._MARLIN_LORA_OVERLAP_EVENTS) return calls -@pytest.mark.parametrize(("master", "split_gate_checks"), [(False, 0), (True, 1)]) -def test_two_stream_batch_gate_is_master_gated(monkeypatch, master, split_gate_checks): - calls = _run_marlin_policy(monkeypatch, tokens=1, master=master, two_stream=True) - assert calls.split_gate_checks == split_gate_checks - shrinks = [call for call in calls.merged if call["stage"] == "shrink"] - assert shrinks[0]["stream"] == ("side" if master else "main") - - -def test_shared_outer_factorization_runtime_flow(monkeypatch): - tokens = 16 - calls = _run_marlin_policy(monkeypatch, tokens=tokens) - gate_expand = [ - call for call in calls.merged if call["stage"] == "expand" and call["broadcast"] - ] - down_shrink = [call for call in calls.merged if call["stage"] == "shrink"] - routing = [call for call in calls.merged if call["stage"] == "routing"] - - assert len(gate_expand) == 1 - assert len(down_shrink) == 1 - assert calls.weighted_rank_sums == 0 - assert calls.fused_tails == 1 - assert gate_expand[0]["intermediate_shape"] == (tokens, 2) - assert down_shrink[0]["intermediate_shape"] == (tokens, 2, 1) - assert [(call["prewarm_a"], call["prewarm_b"]) for call in routing] == [ - (False, True), - (True, False), - ] - assert down_shrink[0]["prewarm_b"] is False - - -@pytest.mark.parametrize( - ("slots", "rank"), - [(8, 128), (16, 128)], -) -def test_ep_shared_outer_uses_safe_generic_fallback(monkeypatch, slots, rank): - calls = _run_marlin_policy( - monkeypatch, - tokens=1, - ep=True, - slots=slots, - rank=rank, - ) - generic_down = [ - call for call in calls.merged if call["stage"] == "all" and call["shared_b"] - ] - assert len(generic_down) == 1 - assert generic_down[0]["zero_intermediate"] is True - assert generic_down[0]["direct_expand"] is (rank <= 64) - - -def test_ep_shared_outer_low_rank_multi_slot_takes_factored_path(monkeypatch): - # Slot-count gates are lifted: EP shared-outer rank<=64 pools of any size - # collapse through the factored prefill path instead of the zeroed generic - # fallback (routing is by adapter slot, so no unowned regions are read). - calls = _run_marlin_policy( - monkeypatch, - tokens=1, - ep=True, - slots=5, - rank=32, - ) - generic_down = [ - call for call in calls.merged if call["stage"] == "all" and call["shared_b"] - ] - assert not generic_down - - -def test_ep_per_expert_layout_uses_generic_fallback(monkeypatch): - calls = _run_marlin_policy( - monkeypatch, - tokens=1, - ep=True, - slots=8, - rank=128, - shared_outer=False, - ) - generic_down = [ - call - for call in calls.merged - if call["stage"] == "all" and call["mul_routed_weight"] - ] - assert len(generic_down) == 1 - assert generic_down[0]["shared_b"] is False - assert generic_down[0]["zero_intermediate"] is False - assert generic_down[0]["direct_expand"] is False - - -@pytest.mark.parametrize( - ("case", "expected"), - [ - ("supported", True), - ("multi_slot", False), - ("non_shared", False), - ("rank_too_large", False), - ("single_route", False), - ("empty_batch", False), - ("mismatched_experts", False), - ], -) -def test_shared_outer_factorization_eligibility_is_narrow(monkeypatch, case, expected): - module = _load_marlin_runner(monkeypatch, "_marlin_runner_eligibility") - - rank = 65 if case == "rank_too_large" else 32 - slots = 2 if case == "multi_slot" else 1 - experts = 4 - hidden = 8 - intermediate = 3 - info = SimpleNamespace( - max_lora_rank=rank, - experts_shared_outer_loras=case != "non_shared", - gate_up_lora_a_weights=torch.empty(slots, 1, 2 * rank, hidden), - gate_up_lora_b_weights=torch.empty(slots, experts, 2 * intermediate, rank), - down_lora_a_weights=torch.empty( - slots, - experts + (1 if case == "mismatched_experts" else 0), - rank, - intermediate, - ), - down_lora_b_weights=torch.empty(slots, 1, hidden, rank), - ) - tokens = 0 if case == "empty_batch" else 1 - topk = 1 if case == "single_route" else 2 - assert module._use_shared_outer_factorization(info, tokens, topk) is expected - - -@pytest.mark.parametrize( - ("case", "expected"), - [ - ("supported", True), - ("three_slots", True), - ("one_slot", False), - ("five_slots", True), - ("sixteen_slots", True), - ("large_batch", False), - ("ep", False), - ("hopper", False), - ], -) -def test_multi_shared_outer_decode_factorization_is_narrow(monkeypatch, case, expected): - module = _load_marlin_runner(monkeypatch, "_marlin_runner_multi_policy") - monkeypatch.setattr( - torch.cuda, - "get_device_capability", - lambda _device: (9, 0) if case == "hopper" else (10, 0), - ) - slots = ( - 1 - if case == "one_slot" - else ( - 3 - if case == "three_slots" - else 5 if case == "five_slots" else 16 if case == "sixteen_slots" else 4 - ) - ) - info = SimpleNamespace( - max_lora_rank=32, - experts_shared_outer_loras=True, - gate_up_lora_a_weights=SimpleNamespace(shape=(slots, 1, 64, 6144)), - gate_up_lora_b_weights=SimpleNamespace(shape=(slots, 256, 768, 32)), - down_lora_a_weights=SimpleNamespace(shape=(slots, 256, 32, 384)), - down_lora_b_weights=SimpleNamespace(shape=(slots, 1, 6144, 32)), - ) - hidden_states = SimpleNamespace( - is_cuda=True, dtype=torch.bfloat16, device=torch.device("cuda") - ) - assert ( - module._use_multi_shared_outer_decode_factorization( - info, - hidden_states, - num_tokens=33 if case == "large_batch" else 32, - hidden_size=6144, - router_topk=6, - num_experts=256, - intermediate_size=384, - ep_active=case == "ep", - ) - is expected - ) - - -@pytest.mark.parametrize( - ("case", "expected"), - [ - ("supported", True), - ("four_slots", True), - ("decode_boundary", False), - ("one_slot", False), - ("five_slots", True), - ("ep", True), - ("ep_decode", True), - ("non_shared", False), - ("rank_too_large", False), - ("single_route", False), - ("mismatched_shape", False), - ], -) -def test_multi_shared_outer_prefill_factorization_is_narrow( - monkeypatch, case, expected -): - module = _load_marlin_runner(monkeypatch, "_marlin_runner_prefill_policy") - - slots = ( - 1 - if case == "one_slot" - else 5 if case == "five_slots" else 4 if case == "four_slots" else 2 - ) - rank = 65 if case == "rank_too_large" else 32 - experts = 256 - intermediate = 384 - hidden = 6144 - info = SimpleNamespace( - max_lora_rank=rank, - experts_shared_outer_loras=case != "non_shared", - gate_up_lora_a_weights=SimpleNamespace(shape=(slots, 1, 2 * rank, hidden)), - gate_up_lora_b_weights=SimpleNamespace( - shape=(slots, experts, 2 * intermediate, rank) - ), - down_lora_a_weights=SimpleNamespace( - shape=( - slots, - experts + (1 if case == "mismatched_shape" else 0), - rank, - intermediate, - ) - ), - down_lora_b_weights=SimpleNamespace(shape=(slots, 1, hidden, rank)), - ) - assert ( - module._use_multi_shared_outer_prefill_factorization( - info, - num_tokens=32 if case in ("decode_boundary", "ep_decode") else 33, - hidden_size=hidden, - router_topk=1 if case == "single_route" else 6, - num_experts=experts, - intermediate_size=intermediate, - ep_active=case in ("ep", "ep_decode"), - ) - is expected - ) - - -@pytest.mark.parametrize( - ("case", "expected"), - [ - ("supported", True), - ("unfactored", False), - ("unfused", False), - ("ep", False), - ("large_batch", False), - ], -) -def test_direct_decode_selection_is_narrow(monkeypatch, case, expected): - module = _load_marlin_runner(monkeypatch, "_marlin_runner_direct_policy") - info = SimpleNamespace( - gate_up_lora_b_weights=SimpleNamespace(shape=(3, 256, 768, 32)), - down_lora_a_weights=SimpleNamespace(shape=(3, 256, 32, 384)), - ) - assert ( - module._use_direct_decode_kernels( - info, - factored_shared_outer=case != "unfactored", - fused_shared_outer_tail=case != "unfused", - ep_active=case == "ep", - num_tokens=33 if case == "large_batch" else 32, - num_experts=256, - intermediate_size=384, - ) - is expected - ) - - -@pytest.mark.parametrize( - ("case", "expected"), - [ - ("supported", True), - ("boundary_m", True), - ("hopper", False), - ("fp16", False), - ("rank64", False), - ("hidden", False), - ("topk", False), - ("large_m", False), - ], -) -def test_fused_shared_outer_tail_is_b200_inkling_specific(monkeypatch, case, expected): - module = _load_marlin_runner(monkeypatch, "_marlin_runner_tail_policy") - monkeypatch.setattr( - torch.cuda, - "get_device_capability", - lambda _device: (9, 0) if case == "hopper" else (10, 0), - ) - info = SimpleNamespace(max_lora_rank=64 if case == "rank64" else 32) - hidden_states = SimpleNamespace( - is_cuda=True, - dtype=torch.float16 if case == "fp16" else torch.bfloat16, - device=torch.device("cuda"), - ) - assert ( - module._use_fused_shared_outer_tail( - info, - hidden_states, - 513 if case == "large_m" else 512 if case == "boundary_m" else 32, - 4096 if case == "hidden" else 6144, - 8 if case == "topk" else 6, - ) - is expected - ) - - -def test_multi_prefill_collapses_only_shared_factors_and_separates_caches( - monkeypatch, -): - calls = _run_marlin_policy(monkeypatch, tokens=64, slots=3) - - routing = [call for call in calls.merged if call["stage"] == "routing"] - full_routing = [call for call in routing if call["topk_shape"] == (64, 2)] - collapsed_routing = [call for call in routing if call["topk_shape"] == (64, 1)] - assert [(call["prewarm_a"], call["prewarm_b"]) for call in full_routing] == [ - (False, True), # real-route per-expert gate B - (True, False), # real-route per-expert down A - ] - assert [(call["prewarm_a"], call["prewarm_b"]) for call in collapsed_routing] == [ - (True, False), # collapsed selected shared gate A - (False, True), # collapsed selected shared down B - ] - assert len({call["cache_id"] for call in full_routing}) == 1 - assert len({call["cache_id"] for call in collapsed_routing}) == 1 - assert full_routing[0]["cache_id"] != collapsed_routing[0]["cache_id"] - - gate_shrink = next( - call for call in calls.merged if call["stage"] == "shrink" and call["shared_a"] - ) - gate_expand = next( - call for call in calls.merged if call["stage"] == "expand" and call["broadcast"] - ) - down_shrink = next( - call - for call in calls.merged - if call["stage"] == "shrink" and not call["shared_a"] - ) - down_expand = next( - call for call in calls.merged if call["stage"] == "expand" and call["shared_b"] - ) - - assert gate_shrink["topk_shape"] == (64, 1) - assert gate_shrink["intermediate_shape"] == (64, 2) - assert gate_expand["topk_shape"] == (64, 2) - assert gate_expand["intermediate_shape"] == (64, 2) - assert down_shrink["topk_shape"] == (64, 2) - assert down_shrink["intermediate_shape"] == (64, 2, 1) - assert down_expand["topk_shape"] == (64, 1) - assert down_expand["intermediate_shape"] == (64, 1) - assert down_expand["fuse_add"] is True - assert down_expand["direct_expand"] is False - assert down_expand["mul_routed_weight"] is False - assert calls.weighted_rank_sums == 1 - # The mapped one-token-per-CTA tail remains decode-only. - assert calls.fused_tails == 0 - - -def test_multi_prefill_none_rows_preserve_base_reduction(monkeypatch): - calls = _run_marlin_policy( - monkeypatch, - tokens=64, - slots=2, - capture=True, - active_lora=False, - base_mapping=True, - base_value=3.0, - ) - - down_expand = next( - call for call in calls.merged if call["stage"] == "expand" and call["shared_b"] - ) - assert down_expand["topk_shape"] == (64, 1) - assert torch.equal(down_expand["mapping"], torch.full((64,), -1, dtype=torch.int32)) - # The fake Marlin down output is 3 for each of two routes. The collapsed - # shared-B expand masks every None row, so it must leave the base sum at 6. - torch.testing.assert_close(calls.result, torch.full_like(calls.result, 6.0)) - assert calls.fused_tails == 0 - - -def test_direct_decode_skips_virtual_routing_and_zero_fill(monkeypatch): - calls = _run_marlin_policy( - monkeypatch, tokens=16, two_stream=True, direct_decode=True - ) - - assert [call for call in calls.merged if call["stage"] == "routing"] == [] - assert [call for call in calls.merged if call["stage"] == "shrink"] == [] - assert calls.direct_gate == 1 - assert calls.direct_down == 1 - assert calls.zeroed == [] - assert calls.fused_tails == 1 - - def test_ep_uses_local_alignment_and_skips_nonlocal_marlin_blocks(monkeypatch): - calls = _run_marlin_policy(monkeypatch, tokens=16, ep=True) + """Guards expert-parallel index localization and both accumulator clears; + reds when ``moe_align_block_size`` gets the global expert count, ``is_ep`` + is dropped on a Marlin GEMM, or an accumulator is left uninitialized.""" + + calls = _run_marlin_flow(monkeypatch, tokens=16, ep=True, slots=2, rank=128) assert calls.align_num_experts == [1] assert calls.marlin_is_ep == [True, True] - assert calls.cache3_was_zero is True - -def test_factored_decode_two_stream_schedule_and_ownership(monkeypatch): - calls = _run_marlin_policy(monkeypatch, tokens=16, two_stream=True) - shrinks = [call for call in calls.merged if call["stage"] == "shrink"] - - assert len(shrinks) == 1 - assert shrinks[0]["stream"] == "side" - assert shrinks[0]["prewarm_b"] is False - assert len(calls.zeroed) == 1 - assert calls.zeroed[0][1] == "side" - - buffer_id = shrinks[0]["intermediate_id"] - assert calls.zeroed[0][2] == buffer_id - zero_index = calls.schedule.index(calls.zeroed[0]) - shrink_index = next( - index - for index, item in enumerate(calls.schedule) - if item[:3] == ("merged", "shrink", "side") - ) - down_record = calls.event_records[-1] - down_wait = calls.event_waits[-1] - record_index = calls.schedule.index(down_record) - wait_index = calls.schedule.index(down_wait) - fused_index = next( - index for index, item in enumerate(calls.schedule) if item[0] == "fused_tail" - ) - assert down_record[2] == "side" - assert down_wait == ("wait_event", "main", down_record[1]) - assert zero_index < shrink_index < record_index < wait_index < fused_index - assert fused_index == len(calls.schedule) - 1 - - -def test_factored_decode_single_stream_fallback_has_one_main_shrink(monkeypatch): - calls = _run_marlin_policy(monkeypatch, tokens=16, two_stream=False) - shrinks = [call for call in calls.merged if call["stage"] == "shrink"] - - assert len(shrinks) == 1 - assert shrinks[0]["stream"] == "main" - assert shrinks[0]["prewarm_b"] is False - assert len(calls.zeroed) == 1 - assert calls.zeroed[0][1] == "main" - assert calls.event_records == [] - assert calls.event_waits == [] - - buffer_id = shrinks[0]["intermediate_id"] - assert calls.zeroed[0][2] == buffer_id - second_marlin_index = max( + down_gemm_output = calls.marlin_output_ids[1] + assert [item[2] for item in calls.zeroed] == [down_gemm_output] + zero_index = calls.schedule.index(("zero", "main", down_gemm_output)) + down_gemm_index = [ index for index, item in enumerate(calls.schedule) if item[0] == "marlin" - ) - zero_index = calls.schedule.index(calls.zeroed[0]) - shrink_index = next( - index - for index, item in enumerate(calls.schedule) - if item[:3] == ("merged", "shrink", "main") - ) - fused_index = next( - index for index, item in enumerate(calls.schedule) if item[0] == "fused_tail" - ) - assert second_marlin_index < zero_index < shrink_index < fused_index + ][1] + assert zero_index < down_gemm_index + + down_stage = [ + call for call in calls.merged if call["stage"] == "all" and call["shared_b"] + ] + assert len(down_stage) == 1 + assert down_stage[0]["zero_intermediate"] is True -def test_factored_decode_capture_base_rows_keep_main_owned_buffers(monkeypatch): - calls = _run_marlin_policy( +def test_factored_decode_captured_two_stream_schedule_ownership_and_event_lifetime( + monkeypatch, +): + """Guards the captured two-stream decode schedule; reds when LoRA stages gate + on a live adapter rather than capture-or-adapter, the down-shrink clear/event + changes stream or order, an event append drops, or the reduce is in place.""" + + calls = _run_marlin_flow( monkeypatch, tokens=16, two_stream=True, @@ -929,20 +449,42 @@ def test_factored_decode_capture_base_rows_keep_main_owned_buffers(monkeypatch): active_lora=False, base_mapping=True, ) - shrinks = [call for call in calls.merged if call["stage"] == "shrink"] + # Capture records the LoRA stages even with no live adapter. + shrinks = [call for call in calls.merged if call["stage"] == "shrink"] assert len(shrinks) == 1 - assert len(calls.zeroed) == 1 - assert calls.zeroed[0][1] == "side" - assert calls.zeroed[0][2] == shrinks[0]["intermediate_id"] - assert calls.weighted_rank_sums == 0 assert calls.fused_tails == 1 + + # The shrink runs on the side stream and accumulates into a buffer that was + # allocated on main, cleared once, on the stream that writes it. + assert shrinks[0]["stream"] == "side" + assert [item[1:] for item in calls.zeroed] == [ + ("side", shrinks[0]["intermediate_id"]) + ] + + down_record = calls.event_records[-1] + down_wait = calls.event_waits[-1] + assert down_record[2] == "side" + assert down_wait == ("wait_event", "main", down_record[1]) + + zero_index = calls.schedule.index(calls.zeroed[0]) + shrink_index = next( + index + for index, item in enumerate(calls.schedule) + if item[:3] == ("merged", "shrink", "side") + ) + record_index = calls.schedule.index(down_record) + wait_index = calls.schedule.index(down_wait) + fused_index = next( + index for index, item in enumerate(calls.schedule) if item[0] == "fused_tail" + ) + assert zero_index < shrink_index < record_index < wait_index < fused_index + + # gate-up delta done, activation done, down shrink done. assert calls.capture_event_count == 3 + assert calls.result_ptr != calls.input_ptr - torch.testing.assert_close(calls.result, torch.zeros_like(calls.result)) if __name__ == "__main__": - import sys - sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/lora/test_experimental_sgl_marlin_shared_outer_reduce.py b/test/registered/unit/lora/test_experimental_sgl_marlin_shared_outer_reduce.py index afbe56b27..9b8a918e8 100644 --- a/test/registered/unit/lora/test_experimental_sgl_marlin_shared_outer_reduce.py +++ b/test/registered/unit/lora/test_experimental_sgl_marlin_shared_outer_reduce.py @@ -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__": diff --git a/test/registered/unit/lora/test_inkling_linearized_lora_unit.py b/test/registered/unit/lora/test_inkling_linearized_lora_unit.py index 60090176c..c019d36a9 100644 --- a/test/registered/unit/lora/test_inkling_linearized_lora_unit.py +++ b/test/registered/unit/lora/test_inkling_linearized_lora_unit.py @@ -1,1173 +1,509 @@ -"""Regression tests for Inkling's linearized shared-sink LoRA path. - -The production LoRA module has optional GPU/runtime imports that are unavailable -in lightweight unit-test environments. The tests compile selected production -methods directly so every tensor operation remains the real implementation. -""" +"""CPU-only regression tests for Inkling's linearized shared-sink LoRA path: +in-place refresh of the derived decode operands on an adapter swap, and the +shared-sink factor layout from checkpoint to one moe-TP shard's operands.""" from __future__ import annotations -import ast -import logging +import importlib import sys import types -from pathlib import Path from types import SimpleNamespace -from typing import Optional +from unittest import mock import pytest import torch -import torch.nn.functional as F from torch import nn -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.srt.lora import lora_manager as lora_manager_module +from sglang.srt.lora.lora import LoRAAdapter +from sglang.srt.lora.lora_registry import LoRARef +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="refactor-fragile source-parsing unit test; skipped on CI", -) +# Pure layout / bookkeeping math: no CUDA, no distributed groups, no kernels. +register_cpu_ci(est_time=5, suite="base-a-test-cpu") -# Skipped on CI: these hermetic checks AST-extract LoRAManager methods and re-run -# them in a stubbed namespace, so they break whenever the manager's internal +_HIDDEN = 3 +_ADAPTER_RANK = 2 +_NUM_SHARED = 2 +# Per-expert intermediate width before moe-TP sharding, and this rank's shard. +_INTERMEDIATE = 4 +_MOE_TP_SIZE = 2 +_SHARD = _INTERMEDIATE // _MOE_TP_SIZE +_MOE_TP_RANK = 1 -REPO_ROOT = Path(__file__).resolve().parents[4] -LORA_LAYERS_PATH = REPO_ROOT / "python/sglang/srt/lora/layers.py" -LORA_MANAGER_PATH = REPO_ROOT / "python/sglang/srt/lora/lora_manager.py" -INKLING_UTIL_PATH = REPO_ROOT / "python/sglang/srt/models/inkling_common/util.py" -DENSE_MLP_PATH = REPO_ROOT / "python/sglang/srt/models/inkling_common/dense_mlp.py" -INKLING_LAYER_PATH = REPO_ROOT / "python/sglang/srt/models/inkling_common/lora.py" -INKLING_DENSE_PATH = ( - REPO_ROOT / "python/sglang/srt/lora/trtllm_lora_temp/inkling_dense.py" -) +_PREFIX = "model.layers.0.mlp.shared_experts" +_GATE_UP_A = f"{_PREFIX}.gate_up_proj.lora_A.weight" +_GATE_UP_B = f"{_PREFIX}.gate_up_proj.lora_B.weight" +_DOWN_A = f"{_PREFIX}.down_proj.lora_A.weight" +_DOWN_B = f"{_PREFIX}.down_proj.lora_B.weight" -class _Flag: - def __init__(self, value: bool, *, is_set: bool = False): - self.value = value - self.explicitly_set = is_set +class _FakeSharedSink(nn.Module): + """Stand-in for ``InklingBatchDenseMLP``: only the attributes + ``InklingBatchDenseMLPWithLoRA`` reads.""" - def get(self) -> bool: - return self.value - - def is_set(self) -> bool: - return self.explicitly_set + def __init__( + self, + *, + num_shared: int = _NUM_SHARED, + intermediate: int = _SHARD, + moe_tp_size: int = 1, + ): + super().__init__() + self.n_shared_experts = num_shared + self.intermediate_size_per_partition = intermediate + self.moe_tp_size = moe_tp_size + self.moe_tp_rank = 0 + self._linearized_bf16_enabled = True -class _RefreshableSharedSink: - is_shared_fused_moe = True - - def __init__(self, callback): - self._callback = callback - - def on_lora_slots_updated(self, slot_ids): - self._callback(slot_ids) - - -def _load_batch_dense_lora_class(monkeypatch): - """Import the permanent Inkling LoRA layer with lightweight dependencies.""" - impl = _load_inkling_dense_impl() - side_streams: dict[torch.cuda.Stream, torch.cuda.Stream] = {} - - def get_lora_side_stream(): - consumer_stream = torch.cuda.current_stream() - if consumer_stream not in side_streams: - side_streams[consumer_stream] = torch.cuda.Stream() - return side_streams[consumer_stream] - - _stub_module( - monkeypatch, - "sglang.srt.lora.backend.base_backend", - BaseLoRABackend=object, - ) - _stub_module( - monkeypatch, - "sglang.srt.models.inkling_common.dense_mlp", - InklingBatchDenseMLP=_FakeSink, - ) - temp_package = _stub_module( - monkeypatch, - "sglang.srt.lora.trtllm_lora_temp", - get_lora_side_stream=get_lora_side_stream, - ) - temp_package.__path__ = [str(INKLING_DENSE_PATH.parent)] - _stub_module( - monkeypatch, - "sglang.srt.lora.trtllm_lora_temp.inkling_dense", - forward_with_lora=impl.forward_with_lora, - ) - _stub_module(monkeypatch, "sglang.srt.models.inkling_common") - module_name = "sglang.srt.models.inkling_common.lora" - spec = __import__("importlib.util").util.spec_from_file_location( - module_name, INKLING_LAYER_PATH - ) - module = __import__("importlib.util").util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, module_name, module) - monkeypatch.setattr( - sys.modules["sglang.srt.models.inkling_common"], "lora", module, raising=False - ) - spec.loader.exec_module(module) +def _load_sink_lora_class(): + """Import the real Inkling LoRA layer against the fake dense-sink base.""" + stub = types.ModuleType("sglang.srt.models.inkling_common.dense_mlp") + stub.InklingBatchDenseMLP = _FakeSharedSink + patched = dict(sys.modules) + patched["sglang.srt.models.inkling_common.dense_mlp"] = stub + # Drop any cached copy so the real module body re-runs against the stub, + # and let the temporary sys.modules state be discarded on exit. + patched.pop("sglang.srt.models.inkling_common.lora", None) + with mock.patch.dict(sys.modules, patched, clear=True): + module = importlib.import_module("sglang.srt.models.inkling_common.lora") return module.InklingBatchDenseMLPWithLoRA -def _load_bf16_materialization_class(): - return _load_selected_class_methods( - DENSE_MLP_PATH, - "InklingBatchDenseMLP", - { - "weight_loader_fused", - "process_weights_after_loading", - "get_bf16_linearized_weights", - "_refresh_bf16_linearized", - }, - { - "torch": torch, - "FusedMoELoadingMixin": SimpleNamespace( - weight_loader_fused=lambda _self, param, loaded, *_: param.data.copy_( - loaded - ) - ), - "logger": SimpleNamespace( - info=lambda *args: None, info_once=lambda *args: None - ), - "SharedExpertFp4Strategy": SimpleNamespace(FP4=object()), - }, +InklingSharedSinkWithLoRA = _load_sink_lora_class() + + +def _new_sink(*, slots: int = 1, moe_tp_size: int = 1, intermediate: int = _SHARD): + layer = InklingSharedSinkWithLoRA( + moe_tp_size=moe_tp_size, intermediate=intermediate ) - - -def _load_inkling_dense_impl(): - function_names = { - "_apply_per_expert_lora", - "_shared_sink_routing", - "apply_multi_lora", - "forward_with_lora", - } - tree = ast.parse(INKLING_DENSE_PATH.read_text()) - functions = [ - node - for node in tree.body - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) - and node.name in function_names - ] - assert {function.name for function in functions} == function_names - namespace = { - "torch": torch, - "envs": SimpleNamespace( - SGLANG_OPT_USE_INKLING_MULTI_STREAM_OVERLAP=_Flag(True) - ), - "symm_mem_all_reduce": lambda value, _group: value, - } - exec( - compile( - ast.fix_missing_locations(ast.Module(body=functions, type_ignores=[])), - str(INKLING_DENSE_PATH), - "exec", - ), - namespace, - ) - return SimpleNamespace( - **{function_name: namespace[function_name] for function_name in function_names}, - ) - - -def _load_selected_class_methods(path, class_name, method_names, namespace): - """Compile selected production methods into a dependency-free test class.""" - tree = ast.parse(path.read_text()) - source_class = next( - node - for node in tree.body - if isinstance(node, ast.ClassDef) and node.name == class_name - ) - methods = [ - node - for node in source_class.body - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) - and node.name in method_names - ] - assert {method.name for method in methods} == set(method_names) - test_class = ast.ClassDef( - name=f"_{class_name}MethodsUnderTest", - bases=[], - keywords=[], - body=methods, - decorator_list=[], - ) - module_ast = ast.fix_missing_locations( - ast.Module(body=[test_class], type_ignores=[]) - ) - exec(compile(module_ast, str(path), "exec"), namespace) - return namespace[test_class.name] - - -def _load_function(path, function_name, namespace): - tree = ast.parse(path.read_text()) - function = next( - node - for node in tree.body - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) - and node.name == function_name - ) - module_ast = ast.fix_missing_locations(ast.Module(body=[function], type_ignores=[])) - exec(compile(module_ast, str(path), "exec"), namespace) - return namespace[function_name] - - -def _load_manager_methods(method_names, namespace=None): - return _load_selected_class_methods( - LORA_MANAGER_PATH, "LoRAManager", method_names, namespace or {} - ) - - -def _stub_module(monkeypatch, name: str, **attributes): - """Install a small importable module hierarchy for a production-file load.""" - parts = name.split(".") - for end in range(1, len(parts)): - package_name = ".".join(parts[:end]) - if package_name not in sys.modules: - package = types.ModuleType(package_name) - package.__path__ = [] - monkeypatch.setitem(sys.modules, package_name, package) - module = types.ModuleType(name) - for key, value in attributes.items(): - setattr(module, key, value) - monkeypatch.setitem(sys.modules, name, module) - if len(parts) > 1: - parent = sys.modules[".".join(parts[:-1])] - monkeypatch.setattr(parent, parts[-1], module, raising=False) - return module - - -def _load_inkling_util(monkeypatch, state): - class _Dummy: - pass - - for module_name, symbol in ( - ("sglang.srt.layers.moe.fused_moe_triton.layer", "FusedMoE"), - ("sglang.srt.layers.moe.moe_runner.base", "MoeRunnerConfig"), - ("sglang.srt.layers.quantization.base_config", "QuantizationConfig"), - ("sglang.srt.layers.quantization.unquant", "UnquantizedFusedMoEMethod"), - ): - _stub_module(monkeypatch, module_name, **{symbol: _Dummy}) - _stub_module( - monkeypatch, - "sglang.srt.runtime_context", - get_server_args=lambda: state.args, - ) - _stub_module(monkeypatch, "sglang.srt.environ", envs=state.envs) - _stub_module( - monkeypatch, - "sglang.srt.layers.moe", - get_moe_runner_backend=lambda: None, - ) - - module_name = "_inkling_linearized_util_under_test" - spec = __import__("importlib.util").util.spec_from_file_location( - module_name, INKLING_UTIL_PATH - ) - module = __import__("importlib.util").util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, module_name, module) - spec.loader.exec_module(module) - return module - - -@pytest.mark.parametrize( - ( - "enable_lora", - "interleaved", - "serves_fp4", - "expected_fused", - ), - [ - (False, True, False, False), - (False, False, False, True), - (False, True, True, True), - (True, True, False, False), - ], -) -def test_linearized_sink_config_eligibility( - monkeypatch, - enable_lora, - interleaved, - serves_fp4, - expected_fused, -): - state = SimpleNamespace( - args=SimpleNamespace(enable_lora=enable_lora), - envs=SimpleNamespace( - SGLANG_OPT_USE_INKLING_SHARED_FUSED_MOE=_Flag(False), - ), - ) - util = _load_inkling_util(monkeypatch, state) - - assert ( - util.use_inkling_shared_fused_moe( - inference_moe_w13_interleaved=interleaved, - shared_sink_serves_fp4=serves_fp4, - ) - is expected_fused - ) - - -@pytest.mark.parametrize("override", [False, True]) -def test_lora_ignores_fused_shared_expert_override(monkeypatch, override): - state = SimpleNamespace( - args=SimpleNamespace(enable_lora=True), - envs=SimpleNamespace( - SGLANG_OPT_USE_INKLING_SHARED_FUSED_MOE=_Flag(override, is_set=True), - ), - ) - util = _load_inkling_util(monkeypatch, state) - assert not util.use_inkling_shared_fused_moe() - - -class _FakeSink(nn.Module): - def __init__(self, *, hidden_size=3, num_experts=2, expert_size=2, seed=7): - super().__init__() - self.moe_tp_size = 1 - self.moe_tp_rank = 0 - self.intermediate_size_per_partition = expert_size - self.n_shared_experts = num_experts - self.layer_id = 0 - self.inference_moe_w13_interleaved = True - self._linearized_bf16_enabled = True - self._fp4_strategy = SimpleNamespace(serves_fp4=False) - self.tp_group = None - generator = torch.Generator().manual_seed(seed) - self._w13_lin = torch.randn( - num_experts * 2 * expert_size, hidden_size, generator=generator - ) - self._w2_lin = torch.randn( - num_experts * expert_size, hidden_size, generator=generator - ) - self.seen_gammas = [] - - def get_bf16_linearized_weights(self): - return self._w13_lin, self._w2_lin - - def _swiglu(self, gate_up, gammas): - self.seen_gammas.append(gammas.detach().clone()) - gate = gate_up[..., 0::2] - up = gate_up[..., 1::2] - return F.silu(gate) * up * gammas.unsqueeze(-1) - - def _forward_bf16_linearized( - self, x_td, gammas_ts, linearized_weights, use_reduce_scatter - ): - w13_lin, w2_lin = linearized_weights - t = x_td.shape[0] - y = torch.mm(x_td, w13_lin.T).view(t, self.n_shared_experts, -1) - act = self._swiglu(y, gammas_ts) - return torch.mm(act.reshape(t, -1), w2_lin) - - def forward(self, x, gammas, use_reduce_scatter=False): - x_td = x.view(-1, x.size(-1)) if x.ndim != 2 else x - gammas_ts = gammas.view(-1, gammas.size(-1)) if gammas.ndim != 2 else gammas - out_td = self._forward_bf16_linearized( - x_td, - gammas_ts, - self.get_bf16_linearized_weights(), - use_reduce_scatter, - ) - return out_td.view_as(x) if x.ndim == 2 else out_td - - -def test_manager_promotes_dense_sink_in_place(monkeypatch): - lora_cls = _load_batch_dense_lora_class(monkeypatch) - manager_cls = _load_manager_methods( - {"init_lora_modules"}, - { - "BaseLayerWithLoRA": nn.Module, - "Dict": dict, - "FusedMoE": type("_UnusedFusedMoE", (), {}), - "List": list, - "Optional": Optional, - "ParallelLMHead": type("_UnusedParallelLMHead", (), {}), - "VocabParallelEmbedding": type("_UnusedEmbedding", (), {}), - "get_layer_id": lambda _name: 0, - "torch": torch, - }, - ) - layer = _FakeSink() - backend = SimpleNamespace( - max_loras_per_batch=1, - name="torch-test", - is_moe_lora=False, - ) - module_name = "model.layers.0.mlp.shared_experts" - manager = manager_cls() - manager.base_hf_config = SimpleNamespace(num_hidden_layers=1) - manager.base_model = SimpleNamespace(named_modules=lambda: [(module_name, layer)]) - manager.target_modules = {"gate_up_proj", "down_proj"} - manager.lora_backend = backend - - manager.init_lora_modules() - - promoted = manager.lora_modules[0][module_name] - assert promoted is layer - assert type(layer) is lora_cls - assert layer.is_shared_fused_moe is True - assert layer.lora_backend is backend - assert backend.is_moe_lora is True - - -def _make_pool(*, slots: int, max_rank: int, active_rank: int, scale: float): - """Build the real shared-outer memory-pool layouts with zero rank padding.""" - n, f, hidden = 2, 2, 3 - gate_a = torch.zeros(slots, 1, 2 * max_rank, hidden) - gate_b = torch.zeros(slots, n, 2 * f, max_rank) - down_a = torch.zeros(slots, n, max_rank, f) - down_b = torch.zeros(slots, 1, hidden, max_rank) - - base = torch.arange(1, active_rank * hidden + 1, dtype=torch.float32).view( - active_rank, hidden - ) - for slot in range(slots): - slot_scale = scale * (slot + 1) - gate_a[slot, 0, :active_rank] = base * (0.03 * slot_scale) - gate_a[slot, 0, max_rank : max_rank + active_rank] = base * (-0.02 * slot_scale) - gate_b[slot, ..., :active_rank] = 0.05 * slot_scale - down_a[slot, ..., :active_rank, :] = 0.04 * slot_scale - down_b[slot, ..., :active_rank] = -0.06 * slot_scale - return gate_a, gate_b, down_a, down_b - - -def _install_capture_mode(monkeypatch, capture_state): - _stub_module( - monkeypatch, - "sglang.srt.model_executor.runner_utils.capture_mode", - get_is_capture_mode=lambda: capture_state.value, - ) - - -def _make_layer(monkeypatch, *, slots=1, max_rank=2, active_rank=2, scale=1.0): - capture_state = SimpleNamespace(value=False) - _install_capture_mode(monkeypatch, capture_state) - batch_info = SimpleNamespace( - has_active_lora=False, - lora_ranks=[active_rank], - moe_lora_info=SimpleNamespace( - token_lora_mapping=torch.tensor([-1, -1], dtype=torch.int32) - ), - ) - backend = SimpleNamespace( - name="triton" if slots > 1 else "torch-test", - batch_info=batch_info, - max_loras_per_batch=slots, - is_moe_lora=False, - ) - layer = _load_batch_dense_lora_class(monkeypatch)() - layer.initialize_lora(backend) - adapter_values = _make_pool( - slots=slots, max_rank=max_rank, active_rank=active_rank, scale=scale - ) - pool = tuple(torch.zeros_like(tensor) for tensor in adapter_values) - layer.set_lora_info(*pool) - _replace_slot(pool, adapter_values) - layer.on_lora_slots_updated(None) - return layer, pool, batch_info, capture_state - - -def _forward(layer, batch_info, *, active: bool, mapping): - batch_info.has_active_lora = active - moe_lora_info = getattr(batch_info, "moe_lora_info", None) - if moe_lora_info is not None: - mapping_tensor = torch.as_tensor(mapping, dtype=torch.int32).flatten() - if mapping_tensor.numel() == 1: - moe_lora_info.token_lora_mapping.fill_(mapping_tensor.item()) - else: - moe_lora_info.token_lora_mapping.copy_(mapping_tensor) - x = torch.tensor([[0.5, -1.0, 0.25], [1.25, 0.75, -0.5]]) - gammas = torch.tensor([[0.2, 0.8], [0.65, 0.35]]) - output = layer(x, gammas=gammas) - return output, gammas - - -def _replace_slot(pool, replacement): - with torch.no_grad(): - for target, source in zip(pool, replacement): - target.copy_(source) - - -def _clear_slot(pool): - """Mirror production None loading by zeroing both factors.""" - with torch.no_grad(): - for tensor in pool: - tensor.zero_() - - -def test_bf16_materialization_and_w2_reload_refresh_stable_storage(): - layer = _load_bf16_materialization_class()() - layer._linearized_bf16_enabled = True - layer._fp4_strategy = object() - layer._bf16_linearized_ready = False - layer.n_shared_experts = 2 - layer.w13_weight = nn.Parameter(torch.arange(24.0).view(2, 4, 3)) - layer.w2_weight = nn.Parameter(torch.arange(12.0).view(2, 3, 2)) - layer._w2_lin = torch.empty(4, 3) - - layer.process_weights_after_loading() - w13, w2 = layer.get_bf16_linearized_weights() - torch.testing.assert_close(w13, layer.w13_weight.view(8, 3)) - torch.testing.assert_close( - w2, layer.w2_weight.detach().transpose(1, 2).reshape(4, 3) - ) - storage = w2.data_ptr() - - replacement = layer.w2_weight.detach().add(100) - layer.weight_loader_fused(layer.w2_weight, replacement, "w2_weight", "w2") - assert layer._w2_lin.data_ptr() == storage - torch.testing.assert_close(layer._w2_lin, replacement.transpose(1, 2).reshape(4, 3)) - - -def test_capture_like_base_adapter_base_replay_and_direct_gammas(monkeypatch): - layer, pool, batch_info, capture_state = _make_layer(monkeypatch) - capture_state.value = True - - adapter_pool = tuple(tensor.clone() for tensor in pool) - _clear_slot(pool) - layer.on_lora_slots_updated(None) - assert torch.count_nonzero(layer._w1_delta) == 0 - assert torch.count_nonzero(layer._a_cat) == 0 - base_before, gammas = _forward(layer, batch_info, active=False, mapping=-1) - _replace_slot(pool, adapter_pool) - layer.on_lora_slots_updated(None) - adapter, _ = _forward(layer, batch_info, active=True, mapping=0) - _clear_slot(pool) - layer.on_lora_slots_updated(None) - base_after, _ = _forward(layer, batch_info, active=False, mapping=-1) - - torch.testing.assert_close(base_before, base_after, rtol=0, atol=0) - assert not torch.allclose(adapter, base_before) - for seen in layer.seen_gammas: - torch.testing.assert_close(seen, gammas, rtol=0, atol=0) - - -def test_adapter_hot_swap_refreshes_in_place_for_graph_replay(monkeypatch): - layer, pool, batch_info, capture_state = _make_layer(monkeypatch, scale=1.0) - capture_state.value = True - adapter_a, _ = _forward(layer, batch_info, active=True, mapping=0) - pointers_before = (layer._w1_delta.data_ptr(), layer._a_cat.data_ptr()) - contents_before = (layer._w1_delta.clone(), layer._a_cat.clone()) - - adapter_b_pool = _make_pool(slots=1, max_rank=2, active_rank=2, scale=2.5) - _replace_slot(pool, adapter_b_pool) - layer.on_lora_slots_updated(None) - adapter_b, _ = _forward(layer, batch_info, active=True, mapping=0) - - assert (layer._w1_delta.data_ptr(), layer._a_cat.data_ptr()) == pointers_before - assert not torch.equal(layer._w1_delta, contents_before[0]) - assert not torch.equal(layer._a_cat, contents_before[1]) - assert not torch.allclose(adapter_a, adapter_b) - - -def test_slot_update_hook_only_refreshes_changed_slots(monkeypatch): - layer, pool, _, _ = _make_layer(monkeypatch, slots=3) - pointers = (layer._w1_delta.data_ptr(), layer._a_cat.data_ptr()) - running_before = (layer._w1_delta[:2].clone(), layer._a_cat[:2].clone()) - changed_before = (layer._w1_delta[2].clone(), layer._a_cat[2].clone()) - - replacement = _make_pool(slots=3, max_rank=2, active_rank=2, scale=3.0) - with torch.no_grad(): - for target, source in zip(pool, replacement): - target[:2].add_(10) - target[2].copy_(source[2]) - layer.on_lora_slots_updated({2}) - - assert (layer._w1_delta.data_ptr(), layer._a_cat.data_ptr()) == pointers - torch.testing.assert_close(layer._w1_delta[:2], running_before[0], rtol=0, atol=0) - torch.testing.assert_close(layer._a_cat[:2], running_before[1], rtol=0, atol=0) - assert not torch.equal(layer._w1_delta[2], changed_before[0]) - assert not torch.equal(layer._a_cat[2], changed_before[1]) - - -def test_rank_smaller_than_max_rank_matches_compact_rank(monkeypatch): - padded, _, padded_info, padded_capture = _make_layer( - monkeypatch, max_rank=3, active_rank=1, scale=1.3 - ) - padded_capture.value = True - padded_output, _ = _forward(padded, padded_info, active=True, mapping=0) - - compact, _, compact_info, compact_capture = _make_layer( - monkeypatch, max_rank=1, active_rank=1, scale=1.3 - ) - compact_capture.value = True - compact_output, _ = _forward(compact, compact_info, active=True, mapping=0) - - torch.testing.assert_close(padded_output, compact_output, rtol=1e-5, atol=1e-6) - - -def _selected_slot_reference(layer, mapping): - x = layer._w13_lin.new_tensor([[0.5, -1.0, 0.25], [1.25, 0.75, -0.5]]) - gammas = layer._w13_lin.new_tensor([[0.2, 0.8], [0.65, 0.35]]) - t = x.shape[0] - n = layer.n_shared_experts - y = torch.mm(x, layer._w13_lin.T).view(t, n, -1) - for token, slot in enumerate(mapping): - if slot >= 0: - shrink = torch.mm( - x[token : token + 1], layer.gate_up_lora_a_weights[slot, 0].T - ) - y[token : token + 1] += torch.mm(shrink, layer._w1_delta[slot].T).view( - 1, n, -1 - ) - gate = y[..., 0::2] - up = y[..., 1::2] - act = F.silu(gate) * up * gammas.unsqueeze(-1) - out = torch.mm(act.reshape(t, -1), layer._w2_lin) - for token, slot in enumerate(mapping): - if slot >= 0: - shrink = torch.mm( - act[token : token + 1].reshape(1, -1), layer._a_cat[slot].T - ) - out[token : token + 1] += torch.mm( - shrink, layer.down_lora_b_weights[slot, 0].T - ) - return out - - -def test_dense_sink_tp_slices_and_flat_factor_normalization(monkeypatch): - layer, _, _, _ = _make_layer(monkeypatch) - layer.moe_tp_size = 2 - layer.intermediate_size_per_partition = 2 - n, rank, full_intermediate = layer.n_shared_experts, 2, 4 - down_a = torch.arange(n * rank * full_intermediate).view(n, rank, full_intermediate) - gate_up_b = torch.arange(n * 2 * full_intermediate * rank).view( - n, 2 * full_intermediate, rank - ) - expected_b = torch.stack( - [torch.cat([weight[2:4], weight[6:8]], dim=0) for weight in gate_up_b] - ) - - for a, b in ( - (down_a, gate_up_b), - (down_a.transpose(0, 1).reshape(rank, -1), gate_up_b.reshape(-1, rank)), - ): - torch.testing.assert_close( - layer.slice_moe_lora_a_weights(a, 1, "down_proj_moe"), down_a[..., 2:4] - ) - torch.testing.assert_close( - layer.slice_moe_lora_b_weights(b, 1, "gate_up_proj_moe"), expected_b - ) - - hidden_size = layer.gate_up_lora_a_weights.shape[-1] - gate_a = torch.zeros(2 * rank, hidden_size) - down_b = torch.zeros(hidden_size, rank) - assert layer.slice_moe_lora_a_weights(gate_a, 1, "gate_up_proj_moe").shape == ( - 1, - 2 * rank, - hidden_size, - ) - assert layer.slice_moe_lora_b_weights(down_b, 1, "down_proj_moe").shape == ( - 1, - hidden_size, - rank, - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("slots", [1, 2, 4, 5, 8, 16]) -def test_multi_slot_cuda_graph_replay(monkeypatch, slots): - from sglang.srt.lora.backend.triton_backend import TritonLoRABackend - from sglang.srt.lora.utils import LoRABatchInfo, MoELoRABatchInfo - - capture_state = SimpleNamespace(value=True) - _install_capture_mode(monkeypatch, capture_state) - device = torch.device("cuda") - dtype = torch.bfloat16 - rank = 1 - max_rank = 3 - moe_info = MoELoRABatchInfo( - seg_indptr=torch.tensor([0, 1, 2], device=device, dtype=torch.int32), - req_to_lora=torch.tensor([0, slots - 1], device=device, dtype=torch.int32), - adapter_enabled=torch.ones(slots, device=device, dtype=torch.int32), - token_lora_mapping=torch.tensor( - [0, slots - 1], device=device, dtype=torch.int32 - ), - ) - batch_info = LoRABatchInfo( - use_cuda_graph=True, - bs=2, - num_segments=2, - seg_indptr=moe_info.seg_indptr, - weight_indices=moe_info.req_to_lora, - lora_ranks=torch.full((slots,), rank, device=device, dtype=torch.int32), - scalings=torch.full((slots,), 9.0, device=device), - max_len=1, - seg_lens=torch.ones(2, device=device, dtype=torch.int32), - permutation=None, - req_seg_indptr=moe_info.seg_indptr, - req_weight_indices=moe_info.req_to_lora, - moe_lora_info=moe_info, - has_active_lora=True, - ) - backend = TritonLoRABackend(max_loras_per_batch=slots, device=device) - backend.batch_info = batch_info - - layer = _load_batch_dense_lora_class(monkeypatch)() - layer._w13_lin = layer._w13_lin.to(device=device, dtype=dtype) - layer._w2_lin = layer._w2_lin.to(device=device, dtype=dtype) - layer.initialize_lora(backend) - pool = tuple( - tensor.to(device=device, dtype=dtype) - for tensor in _make_pool( - slots=slots, max_rank=max_rank, active_rank=rank, scale=1.0 + # initialize_lora() reads only these two backend fields (and flips + # is_moe_lora on the backend it is handed). + layer.initialize_lora( + SimpleNamespace( + name="triton" if slots > 1 else "torch-test", + max_loras_per_batch=slots, + is_moe_lora=False, ) ) - layer.set_lora_info(*pool) - x = layer._w13_lin.new_tensor([[0.5, -1.0, 0.25], [1.25, 0.75, -0.5]]) - gammas = layer._w13_lin.new_tensor([[0.2, 0.8], [0.65, 0.35]]) - - for _ in range(3): - layer(x, gammas=gammas) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_output = layer(x, gammas=gammas) - graph.replay() - torch.testing.assert_close( - graph_output, - _selected_slot_reference(layer, [0, slots - 1]), - rtol=2e-2, - atol=2e-2, - ) - if slots == 1: - return - - batch_info.weight_indices.copy_( - torch.tensor([slots - 1, 0], device=device, dtype=torch.int32) - ) - batch_info.lora_ranks[0] = 0 - moe_info.token_lora_mapping.copy_( - torch.tensor([slots - 1, -1], device=device, dtype=torch.int32) - ) - graph.replay() - torch.testing.assert_close( - graph_output, - _selected_slot_reference(layer, [slots - 1, -1]), - rtol=2e-2, - atol=2e-2, - ) - - batch_info.lora_ranks[0] = rank - batch_info.weight_indices.copy_( - torch.tensor([0, 1], device=device, dtype=torch.int32) - ) - moe_info.token_lora_mapping.copy_( - torch.tensor([0, 1], device=device, dtype=torch.int32) - ) - graph.replay() - torch.testing.assert_close( - graph_output, - _selected_slot_reference(layer, [0, 1]), - rtol=2e-2, - atol=2e-2, - ) - - -@pytest.mark.skipif( - not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(), - reason="CUDA BF16 is required", -) -def test_split_k_shrink_fp32_feeds_temp_bf16_expand(monkeypatch): - monkeypatch.setenv("SGLANG_EXPERIMENTAL_LORA_OPTI", "1") - monkeypatch.setenv("SGLANG_ENABLE_LORA_SHRINK_SPLIT_K", "1") - monkeypatch.setenv("SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC", "1") - monkeypatch.setenv("SGLANG_OPT_LORA_CUBLAS", "0") - monkeypatch.setenv("SGLANG_OPT_LORA_CUBLAS_A", "0") - monkeypatch.setenv("SGLANG_OPT_LORA_CUBLAS_B", "1") - - from sglang.kernels.ops.gemm.trtllm_lora_temp import sgemm_lora_a as triton_ops - from sglang.srt.lora.trtllm_lora_temp import attention - from sglang.srt.lora.utils import LoRABatchInfo - - device = torch.device("cuda") - dtype = torch.bfloat16 - tokens, input_dim, rank, output_dim = 256, 4096, 64, 128 - batch_info = LoRABatchInfo( - use_cuda_graph=False, - bs=1, - num_segments=1, - seg_indptr=torch.tensor([0, tokens], device=device, dtype=torch.int32), - weight_indices=torch.zeros(1, device=device, dtype=torch.int32), - lora_ranks=torch.full((1,), rank, device=device, dtype=torch.int32), - scalings=torch.full((1,), 0.5, device=device), - max_len=tokens, - seg_lens=torch.full((1,), tokens, device=device, dtype=torch.int32), - permutation=None, - ) - x = torch.ones(tokens, input_dim, device=device, dtype=dtype) - a = torch.full((1, rank, input_dim), 1 / input_dim, device=device, dtype=dtype) - b = torch.full((1, output_dim, rank), 1 / rank, device=device, dtype=dtype) - shrink_dtypes = [] - original_shrink = triton_ops.sgemm_lora_a_fwd - - def record_shrink_dtype(*args, **kwargs): - output = original_shrink(*args, **kwargs) - shrink_dtypes.append(output.dtype) - return output - - def reject_common_expand(**_kwargs): - pytest.fail("two-stream attention must use the temporary expand kernel") - - class QuantMethod: - @staticmethod - def apply(_layer, inputs, bias=None): - return torch.full( - (inputs.shape[0], output_dim), - 0.25, - device=inputs.device, - dtype=inputs.dtype, - ) - - monkeypatch.setattr(triton_ops, "sgemm_lora_a_fwd", record_shrink_dtype) - monkeypatch.setattr(attention, "is_two_stream_active", lambda _inputs: True) - layer = SimpleNamespace( - set_lora=True, - base_layer=SimpleNamespace( - input_is_parallel=True, - tp_rank=0, - tp_size=1, - skip_bias_add=True, - bias=None, - reduce_results=False, - quant_method=QuantMethod(), - ), - lora_backend=SimpleNamespace( - _sgemm_info=lambda: batch_info, - run_lora_b_sgemm=reject_common_expand, - ), - A_buffer=a, - B_buffer=b, - ) - - output, output_bias = attention.row_parallel_lora_forward(layer, x) - torch.cuda.synchronize() - assert shrink_dtypes == [torch.float32] - assert output.dtype == dtype - assert output_bias is None - torch.testing.assert_close(output, torch.full_like(output, 0.75), rtol=0, atol=0) - - -def test_fused_moe_wrapper_reports_local_expert_dimension(monkeypatch): - _stub_module( - monkeypatch, - "sglang.srt.lora.lora_moe_runners", - LoRAInfo=SimpleNamespace, - ) - wrapper_cls = _load_selected_class_methods( - LORA_LAYERS_PATH, - "FusedMoEWithLoRA", - {"_get_lora_info"}, - {}, - ) - moe_lora_info = SimpleNamespace( - seg_indptr=torch.tensor([0, 2], dtype=torch.int32), - req_to_lora=torch.tensor([0], dtype=torch.int32), - adapter_enabled=torch.tensor([1], dtype=torch.int32), - token_lora_mapping=torch.tensor([0, 0], dtype=torch.int32), - ) - wrapper = wrapper_cls() - wrapper._lora_runner_backend = SimpleNamespace( - is_experimental_sgl_trtllm=lambda: True, - is_experimental_sgl_marlin=lambda: False, - ) - wrapper.lora_backend = SimpleNamespace( - batch_info=SimpleNamespace( - lora_ranks=torch.tensor([4], dtype=torch.int32), - moe_lora_info=moe_lora_info, - has_active_lora=True, - ), - moe_cg_buffers={"routing": object()}, - ) - wrapper.base_layer = SimpleNamespace( - num_experts=128, num_local_experts=32, hidden_size=64 - ) - wrapper.gate_up_lora_a_weights = torch.empty(1, 1, 8, 64) - wrapper.gate_up_lora_b_weights = torch.empty(1, 32, 16, 4) - wrapper.down_lora_a_weights = torch.empty(1, 32, 4, 8) - wrapper.down_lora_b_weights = torch.empty(1, 1, 64, 4) - wrapper.experts_shared_outer_loras = True - wrapper.lora_use_virtual_experts = True - wrapper.tp_size = 4 - wrapper.tp_rank = 3 - - info = wrapper._get_lora_info() - - assert info.num_experts == 32 - assert info.num_experts == wrapper.down_lora_a_weights.shape[1] - assert info.max_lora_rank == 4 - assert info.has_active_lora is True - - -def test_manager_refresh_follows_slot_copy_and_only_runs_on_changes(): - manager_cls = _load_manager_methods( - {"fetch_new_loras", "_notify_lora_slots_updated"}, - {"Optional": Optional}, - ) - events = [] - - class _Pool: - def __init__(self): - self.uid_to_buffer_id = {} - - def prepare_lora_batch(self, *, cur_uids, **kwargs): - events.append(("pool", set(cur_uids))) - for uid in cur_uids: - if uid not in self.uid_to_buffer_id: - used = set(self.uid_to_buffer_id.values()) - slot = next((i for i in range(4) if i not in used), 1) - if slot == 1: - self.uid_to_buffer_id = { - resident: resident_slot - for resident, resident_slot in self.uid_to_buffer_id.items() - if resident_slot != slot - } - self.uid_to_buffer_id[uid] = slot - - refreshable = _RefreshableSharedSink( - lambda slots: events.append(("refresh", set(slots))) - ) - manager = manager_cls() - manager.max_loras_per_batch = 4 - manager.memory_pool = _Pool() - manager.loras = {"adapter-a": object(), "adapter-b": object()} - manager.lora_modules = [{"sink": refreshable}] - manager.lora_refs = {} - manager.embed_tokens_module = None - manager.lm_head_module = None - - manager.fetch_new_loras({"adapter-a"}) - assert events == [("pool", {"adapter-a"}), ("refresh", {0})] - - events.clear() - manager.fetch_new_loras({"adapter-a"}) - assert events == [("pool", {"adapter-a"})] - - events.clear() - manager.fetch_new_loras({"adapter-b"}) - assert events == [("pool", {"adapter-b"}), ("refresh", {1})] - - events.clear() - manager.fetch_new_loras({None}) - assert events == [("pool", {None}), ("refresh", {2})] - - events.clear() - manager.loras.update({"adapter-c": object(), "adapter-d": object()}) - manager.fetch_new_loras({"adapter-c", "adapter-d"}, running_loras={"adapter-a"}) - assert events[0] == ("pool", {"adapter-a", "adapter-c", "adapter-d"}) - assert events[1] == ("refresh", {1, 3}) - - -def test_manager_unload_reload_same_uid_refreshes_changed_derived_operands(): - manager_cls = _load_manager_methods( - { - "create_lora_update_result", - "fetch_new_loras", - "_notify_lora_slots_updated", - "unload_lora_adapter", - }, - { - "Dict": dict, - "LoRAAdapter": object, - "LoRARef": object, - "LoRAUpdateOutput": SimpleNamespace, - "Optional": Optional, - "logger": logging.getLogger(__name__), - "get_available_gpu_memory": lambda *args, **kwargs: 0.0, - }, - ) - - class _Pool: - def __init__(self): - self.uid_to_buffer_id = {"same-uid": 0} - self.slot_value = torch.tensor([1.0]) - - def remove_lora(self, uid): - slot = self.uid_to_buffer_id.pop(uid, None) - if slot is not None: - self.slot_value.zero_() - return slot - - def prepare_lora_batch(self, *, cur_uids, lora_adapters, **kwargs): - for uid in cur_uids: - if uid not in self.uid_to_buffer_id: - self.uid_to_buffer_id[uid] = 0 - self.slot_value.fill_(lora_adapters[uid].value) - - pool = _Pool() - derived = torch.tensor([-1.0]) - sink = _RefreshableSharedSink(lambda slots: derived.copy_(pool.slot_value)) - ref = SimpleNamespace( - lora_id="same-uid", lora_name="same", lora_path="old", pinned=False - ) - manager = manager_cls() - manager.device = torch.device("cpu") - manager.max_loras_per_batch = 1 - manager.memory_pool = pool - manager.configs = {"same-uid": object()} - manager.loras = {"same-uid": SimpleNamespace(value=1.0)} - manager.lora_refs = {"same-uid": ref} - manager.num_pinned_loras = 0 - manager.lora_modules = [{"sink": sink}] - manager.embed_tokens_module = None - manager.lm_head_module = None - - result = manager.unload_lora_adapter(ref) - assert result.success - torch.testing.assert_close(derived, torch.zeros_like(derived)) - - manager.configs["same-uid"] = object() - manager.loras["same-uid"] = SimpleNamespace(value=9.0) - manager.lora_refs["same-uid"] = SimpleNamespace( - lora_id="same-uid", lora_name="same", lora_path="new", pinned=False - ) - manager.fetch_new_loras({"same-uid"}) - - torch.testing.assert_close(pool.slot_value, torch.tensor([9.0])) - torch.testing.assert_close(derived, torch.tensor([9.0])) - - -def _make_unconfigured_sink(monkeypatch, *, linearized=True): - layer = _load_batch_dense_lora_class(monkeypatch)() - layer._linearized_bf16_enabled = linearized return layer -@pytest.mark.parametrize( - ("max_loras", "backend_name", "linearized", "expected_error"), - [ - (1, "torch-test", True, None), - (4, "triton", True, None), - (5, "triton", True, None), - (16, "triton", True, None), - (8, "csgmv", True, "requires the Triton backend"), - (1, "torch-test", False, "does not use linearized BF16"), - ], -) -def test_dense_sink_lora_initialization_contract( - monkeypatch, max_loras, backend_name, linearized, expected_error -): - layer = _make_unconfigured_sink(monkeypatch, linearized=linearized) - backend = SimpleNamespace( - name=backend_name, - max_loras_per_batch=max_loras, - is_moe_lora=False, +def _adapter(config) -> LoRAAdapter: + """Bind the real normalizer methods to a bare ``LoRAAdapter``.""" + adapter = LoRAAdapter.__new__(LoRAAdapter) + adapter.base_hf_config = config + return adapter + + +def _hf_config(*, architectures, model_type): + return SimpleNamespace( + architectures=architectures, + model_type=model_type, + n_shared_experts=_NUM_SHARED, ) - if expected_error is None: - layer.initialize_lora(backend) - assert layer.lora_backend is backend - assert layer.is_shared_fused_moe is True - assert backend.is_moe_lora is True - else: - with pytest.raises(ValueError, match=expected_error) as exc_info: - layer.initialize_lora(backend) - assert "InklingBatchDenseMLPWithLoRA is ineligible" in str(exc_info.value) - -@pytest.mark.parametrize( - ("case", "expected_error"), - [ - ("valid", None), - ("ndim", "four 4D MoE buffers"), - ("gate_outer", "same expert layout"), - ("down_outer", "same expert layout"), - ("per_expert", None), - ("expert_count", "expert count does not match"), - ("rank128", None), - ("rank_mismatch", "rank dimensions do not match"), - ], -) -def test_dense_sink_requires_canonical_4d_moe_buffers( - monkeypatch, case, expected_error -): - layer = _make_unconfigured_sink(monkeypatch) - layer.initialize_lora( - SimpleNamespace(name="torch-test", max_loras_per_batch=1, is_moe_lora=False) +def _inkling_config(): + return _hf_config( + architectures=["InklingForConditionalGeneration"], model_type="llama" ) - weights = list(_make_pool(slots=1, max_rank=2, active_rank=2, scale=1.0)) - if case == "ndim": - weights[0] = weights[0][0] - elif case == "gate_outer": - weights[0] = weights[0].expand(-1, 2, -1, -1).clone() - elif case == "down_outer": - weights[3] = weights[3].expand(-1, 2, -1, -1).clone() - elif case == "per_expert": - weights[0] = weights[0].expand(-1, 2, -1, -1).clone() - weights[3] = weights[3].expand(-1, 2, -1, -1).clone() - elif case == "expert_count": - weights[1] = torch.zeros(1, 3, 4, 2) - elif case == "rank128": - weights = [ - torch.zeros(1, 1, 256, 3), - torch.zeros(1, 2, 4, 128), - torch.zeros(1, 2, 128, 2), - torch.zeros(1, 1, 3, 128), + + +def _normalize(config, weights): + """Run the adapter-side shared-sink reshape + gate/up stacking.""" + adapter = _adapter(config) + adapter._normalize_shared_expert_moe(weights) + adapter.normalize_gate_up_proj(list(weights), weights) + return weights + + +def _checkpoint_factors(): + """The four flat 2D factors an Inkling shared-sink adapter ships.""" + + def ramp(*shape): + numel = 1 + for dim in shape: + numel *= dim + # Exact binary fractions keep every downstream matmul exact. + return (torch.arange(numel, dtype=torch.float32) / 8.0).reshape(*shape) + + return { + _GATE_UP_A: ramp(_ADAPTER_RANK, _HIDDEN), + _GATE_UP_B: ramp(_NUM_SHARED * 2 * _INTERMEDIATE, _ADAPTER_RANK), + _DOWN_A: ramp(_ADAPTER_RANK, _NUM_SHARED * _INTERMEDIATE), + _DOWN_B: ramp(_HIDDEN, _ADAPTER_RANK), + } + + +def _shard_factors(layer, normalized): + """Slice the normalized factors down to this moe-TP rank.""" + return ( + layer.slice_moe_lora_a_weights( + normalized[_GATE_UP_A], _MOE_TP_RANK, "gate_up_proj_moe" + ), + layer.slice_moe_lora_b_weights( + normalized[_GATE_UP_B], _MOE_TP_RANK, "gate_up_proj_moe" + ), + layer.slice_moe_lora_a_weights( + normalized[_DOWN_A], _MOE_TP_RANK, "down_proj_moe" + ), + layer.slice_moe_lora_b_weights( + normalized[_DOWN_B], _MOE_TP_RANK, "down_proj_moe" + ), + ) + + +def _slot_shapes(*, slots: int, max_rank: int): + return ( + (slots, 1, 2 * max_rank, _HIDDEN), + (slots, _NUM_SHARED, 2 * _SHARD, max_rank), + (slots, _NUM_SHARED, max_rank, _SHARD), + (slots, 1, _HIDDEN, max_rank), + ) + + +def _empty_pool(*, slots: int, max_rank: int): + return tuple( + torch.zeros(shape) for shape in _slot_shapes(slots=slots, max_rank=max_rank) + ) + + +def _pool_from_shards(shards, *, max_rank: int, slots: int = 1): + """Lay one adapter's shard out as ``LoRAMemoryPool.load_lora_weight_to_buffer`` + does: stacked gate/up LoRA-A halves at ``max_rank`` offsets, rank tails zero.""" + a_gate_up, b_gate_up, a_down, b_down = shards + rank = b_gate_up.shape[-1] + buffers = _empty_pool(slots=slots, max_rank=max_rank) + with torch.no_grad(): + for slot in range(slots): + for half in range(2): + buffers[0][slot, 0, half * max_rank : half * max_rank + rank] = ( + a_gate_up[0, half * rank : (half + 1) * rank] + ) + buffers[1][slot, :, :, :rank] = b_gate_up + buffers[2][slot, :, :rank, :] = a_down + buffers[3][slot, 0, :, :rank] = b_down[0] + return buffers + + +def _sink_with_shard(shards, *, max_rank: int): + layer = _new_sink(moe_tp_size=_MOE_TP_SIZE) + layer.set_lora_info(*_pool_from_shards(shards, max_rank=max_rank)) + return layer + + +def _consumer_delta(layer, x, act): + """The two GEMMs the single-slot shared-outer path runs on the operands, + mirroring ``forward_with_lora`` in ``lora/trtllm_lora_temp/inkling_dense.py``.""" + gate_up_shrink = x @ layer.gate_up_lora_a_weights[0, 0].T + down_shrink = act @ layer._a_cat[0].T + return ( + gate_up_shrink @ layer._w1_delta[0].T, + down_shrink @ layer.down_lora_b_weights[0, 0].T, + ) + + +def _reference_delta(shards, x, act): + """The same delta straight from the sharded adapter factors; the sink's w13 is + gate/up *interleaved*, so gate lands at ``[..., 0::2]`` and up at + ``[..., 1::2]`` of each expert's block.""" + a_gate_up, b_gate_up, a_down, b_down = shards + rank = b_gate_up.shape[-1] + experts, f = b_gate_up.shape[0], b_gate_up.shape[1] // 2 + gate_shrink = x @ a_gate_up[0, :rank].T + up_shrink = x @ a_gate_up[0, rank:].T + y = x.new_zeros(x.shape[0], experts, 2 * f) + for expert in range(experts): + y[:, expert, 0::2] = gate_shrink @ b_gate_up[expert, :f].T + y[:, expert, 1::2] = up_shrink @ b_gate_up[expert, f:].T + down_shrink = torch.einsum("tef,ekf->tk", act.view(-1, experts, f), a_down) + return y.reshape(x.shape[0], -1), down_shrink @ b_down[0].T + + +def _slot_factors(*, gate_up_b: float, down_a: float, max_rank: int): + """One slot's pool contents, each factor a flat constant with a zero rank tail, + so the largest magnitude in a derived operand identifies its source adapter.""" + rank = _ADAPTER_RANK + a_gate_up, b_gate_up, a_down, b_down = _empty_pool(slots=1, max_rank=max_rank) + a_gate_up[0, 0, :rank] = 1.0 + a_gate_up[0, 0, max_rank : max_rank + rank] = 1.0 + b_gate_up[0, ..., :rank] = gate_up_b + a_down[0, :, :rank, :] = down_a + b_down[0, ..., :rank] = 1.0 + return (a_gate_up[0], b_gate_up[0], a_down[0], b_down[0]) + + +def _fake_adapter(*, gate_up_b: float, down_a: float, max_rank: int): + return SimpleNamespace( + factors=_slot_factors(gate_up_b=gate_up_b, down_a=down_a, max_rank=max_rank) + ) + + +def _derived_signature(layer): + """Which adapter the derived operands currently hold.""" + return ( + layer._w1_delta.abs().max().item(), + layer._a_cat.abs().max().item(), + ) + + +class _FakeMemoryPool: + """Minimal ``LoRAMemoryPool``: a new uid takes a free (or evicted) slot and its + factors are copied into the shared buffers there, removing a uid frees its + slot, and a base-model uid (``None``) zeroes the slot.""" + + def __init__(self, buffers, *, max_loras_per_batch: int, events: list): + self.buffers = buffers + self.max_loras_per_batch = max_loras_per_batch + self.uid_to_buffer_id: dict = {} + self.events = events + + def prepare_lora_batch(self, *, cur_uids, lora_adapters, **kwargs): + self.events.append(("pool", set(cur_uids))) + for uid in sorted(cur_uids, key=str): + if uid in self.uid_to_buffer_id: + continue + slot = self._take_slot() + self.uid_to_buffer_id[uid] = slot + self._load(slot, lora_adapters.get(uid)) + + def remove_lora(self, uid): + self.events.append(("remove", uid)) + slot = self.uid_to_buffer_id.pop(uid, None) + if slot is not None: + self._load(slot, None) + return slot + + def _take_slot(self) -> int: + used = set(self.uid_to_buffer_id.values()) + for slot in range(self.max_loras_per_batch): + if slot not in used: + return slot + evicted = next(iter(self.uid_to_buffer_id)) + return self.uid_to_buffer_id.pop(evicted) + + def _load(self, slot: int, adapter) -> None: + with torch.no_grad(): + for index, buffer in enumerate(self.buffers): + if adapter is None: + buffer[slot].zero_() + else: + buffer[slot].copy_(adapter.factors[index]) + + +def _new_manager(pool, layer, *, uids): + """A ``LoRAManager`` carrying only the fields the tested methods touch.""" + manager = lora_manager_module.LoRAManager.__new__(lora_manager_module.LoRAManager) + manager.device = torch.device("cpu") + manager.max_loras_per_batch = pool.max_loras_per_batch + manager.memory_pool = pool + manager.lora_modules = [{_PREFIX: layer}] + manager.embed_tokens_module = None + manager.lm_head_module = None + manager.num_pinned_loras = 0 + manager.loras = {} + manager.configs = {uid: object() for uid in uids} + manager.lora_refs = { + uid: LoRARef(lora_id=uid, lora_name=uid, lora_path=f"/lora/{uid}", pinned=False) + for uid in uids + } + return manager + + +def _record_refreshes(layer, events: list) -> None: + """Log every slot-update notification, forwarding to the real handler.""" + handler = layer.on_lora_slots_updated + + def record(slot_ids): + events.append(("refresh", None if slot_ids is None else set(slot_ids))) + handler(slot_ids) + + layer.on_lora_slots_updated = record + + +def test_derived_operands_refresh_in_place_after_slot_copy(monkeypatch): + """Guards adapter swaps landing in the pre-allocated derived operands; reds when + a swap rebinds instead of copying in place, notifies before the pool copies the + new adapter, or skips the refresh on an unload, a reload, or a base-only batch.""" + max_rank = _ADAPTER_RANK + buffers = _empty_pool(slots=1, max_rank=max_rank) + layer = _new_sink(slots=1) + layer.set_lora_info(*buffers) + assert layer.experts_shared_outer_loras is True + + events: list = [] + _record_refreshes(layer, events) + pool = _FakeMemoryPool(buffers, max_loras_per_batch=1, events=events) + manager = _new_manager(pool, layer, uids=("uid-a", "uid-b")) + manager.loras = { + "uid-a": _fake_adapter(gate_up_b=1.0, down_a=0.5, max_rank=max_rank), + "uid-b": _fake_adapter(gate_up_b=2.0, down_a=0.25, max_rank=max_rank), + } + pointers = (layer._w1_delta.data_ptr(), layer._a_cat.data_ptr()) + + # The pool copies the adapter into its slot first; only then is the slot's + # derived operand rebuilt. Refreshing first would derive it from zeros. + manager.fetch_new_loras({"uid-a"}) + assert events == [("pool", {"uid-a"}), ("refresh", {0})] + assert _derived_signature(layer) == (1.0, 0.5) + + # No uid took a new slot: nothing to re-derive. + events.clear() + manager.fetch_new_loras({"uid-a"}) + assert events == [("pool", {"uid-a"})] + + # Hot swap: uid-b evicts uid-a from slot 0. Refreshing before the pool copy + # would leave uid-a's delta in place for every later replay. + events.clear() + manager.fetch_new_loras({"uid-b"}) + assert events == [("pool", {"uid-b"}), ("refresh", {0})] + assert _derived_signature(layer) == (2.0, 0.25) + + # A base-only batch takes the slot too, and must clear the delta. + events.clear() + manager.fetch_new_loras({None}) + assert events == [("pool", {None}), ("refresh", {0})] + assert _derived_signature(layer) == (0.0, 0.0) + + # Unloading frees the slot, and the freed slot must be re-derived. + events.clear() + manager.fetch_new_loras({"uid-a"}) + assert _derived_signature(layer) == (1.0, 0.5) + # The unload info logs sample GPU memory; there is no GPU on this runner. + monkeypatch.setattr( + lora_manager_module, "get_available_gpu_memory", lambda *args, **kwargs: 0.0 + ) + events.clear() + assert manager.unload_lora_adapter(manager.lora_refs["uid-a"]).success + assert events == [("remove", "uid-a"), ("refresh", {0})] + assert _derived_signature(layer) == (0.0, 0.0) + + # Reloading the same uid with different weights re-derives, even though the + # uid and the slot are unchanged. (Re-seeded the way _load_lora_adapter + # would, without touching a checkpoint.) + events.clear() + manager.configs["uid-a"] = object() + manager.lora_refs["uid-a"] = LoRARef( + lora_id="uid-a", lora_name="uid-a", lora_path="/lora/uid-a-v2", pinned=False + ) + manager.loras["uid-a"] = _fake_adapter( + gate_up_b=3.0, down_a=0.75, max_rank=max_rank + ) + manager.fetch_new_loras({"uid-a"}) + assert events == [("pool", {"uid-a"}), ("refresh", {0})] + assert _derived_signature(layer) == (3.0, 0.75) + + # Every refresh wrote through the buffers allocated at set_lora_info() time. + assert (layer._w1_delta.data_ptr(), layer._a_cat.data_ptr()) == pointers + + +def test_shared_sink_factor_layout_roundtrip(): + """Guards the shared-sink factor layout from checkpoint to one moe-TP shard's + decode operands; reds on a lost ``down_proj`` transpose, gate-major gate/up + LoRA-B, a collapsed gate/up shard, or a rank offset from the wrong rank.""" + rank, hidden = _ADAPTER_RANK, _HIDDEN + experts, full = _NUM_SHARED, _INTERMEDIATE + flat = _checkpoint_factors() + gate_up_a, gate_up_b = flat[_GATE_UP_A], flat[_GATE_UP_B] + down_a, down_b = flat[_DOWN_A], flat[_DOWN_B] + + # -- adapter side ----------------------------------------------------- + # Both gates into the Inkling reshape: the architecture list and the + # model type. + for config in ( + _inkling_config(), + _hf_config(architectures=None, model_type="inkling_text"), + ): + normalized = _normalize(config, dict(flat)) + # gate/up LoRA-A is shared across experts and stacked (gate, up). + torch.testing.assert_close( + normalized[_GATE_UP_A], gate_up_a.unsqueeze(0).repeat(1, 2, 1) + ) + # gate/up LoRA-B is expert-major, gate rows before up rows. + torch.testing.assert_close( + normalized[_GATE_UP_B], gate_up_b.reshape(experts, 2 * full, rank) + ) + # down LoRA-A arrives rank-major and must come out expert-major. + torch.testing.assert_close( + normalized[_DOWN_A], + down_a.reshape(rank, experts, full).transpose(0, 1).contiguous(), + ) + # down LoRA-B is shared across experts. + torch.testing.assert_close(normalized[_DOWN_B], down_b.unsqueeze(0)) + + # A non-Inkling base model keeps the stock 2D shared-expert path. + stock = dict(flat) + _adapter( + _hf_config(architectures=["Qwen3MoeForCausalLM"], model_type="qwen3_moe") + )._normalize_shared_expert_moe(stock) + assert stock[_GATE_UP_B].dim() == 2 + assert stock[_DOWN_A].dim() == 2 + + # A *named* per-expert factor is not a shared-outer factor: leaving it 2D is + # what lets adapter validation reject it later. + named_per_expert = f"{_PREFIX}.1.gate_up_proj.lora_A.weight" + per_expert = {named_per_expert: gate_up_a.clone()} + _adapter(_inkling_config())._normalize_shared_expert_moe(per_expert) + torch.testing.assert_close(per_expert[named_per_expert], gate_up_a) + + # -- layer side: this moe-TP rank's shard ----------------------------- + normalized = _normalize(_inkling_config(), dict(flat)) + layer = _new_sink(moe_tp_size=_MOE_TP_SIZE) + shards = _shard_factors(layer, normalized) + start = _MOE_TP_RANK * _SHARD + # The gate half and the up half are sharded independently and re-paired, so + # this rank's gate rows meet this rank's up rows. + expected_gate_up_b = torch.stack( + [ + torch.cat( + [ + expert_b[start : start + _SHARD], + expert_b[full + start : full + start + _SHARD], + ] + ) + for expert_b in gate_up_b.reshape(experts, 2 * full, rank) ] - elif case == "rank_mismatch": - weights[0] = torch.zeros(1, 1, 3, 3) - - if expected_error is None: - layer.set_lora_info(*weights) - if case == "per_expert": - assert layer.experts_shared_outer_loras is False - assert layer._w1_delta is None - assert layer._a_cat is None - elif case == "rank128": - assert layer.experts_shared_outer_loras is True - assert layer._w1_delta.shape == (1, 8, 256) - assert layer._a_cat.shape == (1, 128, 4) - else: - assert layer._w1_delta.shape == (1, 8, 4) - assert layer._a_cat.shape == (1, 2, 4) - else: - with pytest.raises(ValueError, match=expected_error): - layer.set_lora_info(*weights) - - -def test_outer_factor_detection_bool_and_mixed_rejected(): - manager_cls = _load_manager_methods( - {"_detect_shared_outer_loras"}, - { - "Optional": __import__("typing").Optional, - "re": __import__("re"), - }, ) - routed_shared = "model.layers.0.mlp.experts.gate_up_proj.lora_A.weight" - routed_expert = "model.layers.0.mlp.experts.0.gate_up_proj.lora_A.weight" + torch.testing.assert_close(shards[1], expected_gate_up_b) + # down LoRA-A is sharded on the intermediate axis it contracts over. + expected_down_a = down_a.reshape(rank, experts, full).transpose(0, 1)[ + ..., start : start + _SHARD + ] + torch.testing.assert_close(shards[2], expected_down_a) + # The shrink-side gate/up LoRA-A and the expand-side down LoRA-B are not + # sharded by moe-TP. + torch.testing.assert_close(shards[0], gate_up_a.unsqueeze(0).repeat(1, 2, 1)) + torch.testing.assert_close(shards[3], down_b.unsqueeze(0)) - shared_only = manager_cls() - shared_only.loras = { - "shared": SimpleNamespace( - layers=[SimpleNamespace(weights={routed_shared: torch.empty(1, 8, 4)})] - ), - } - assert shared_only._detect_shared_outer_loras() is True + # The layer also accepts the checkpoint's flat 2D factors, and must produce + # the same shard as the normalized 3D ones. + torch.testing.assert_close( + layer.slice_moe_lora_b_weights(gate_up_b, _MOE_TP_RANK, "gate_up_proj_moe"), + expected_gate_up_b, + ) + torch.testing.assert_close( + layer.slice_moe_lora_a_weights(down_a, _MOE_TP_RANK, "down_proj_moe"), + expected_down_a, + ) - per_expert_only = manager_cls() - per_expert_only.loras = { - "per-expert": SimpleNamespace( - # numbered 2D expert weights must be visible as per-expert layout - layers=[SimpleNamespace(weights={routed_expert: torch.empty(4, 4)})] - ), - } - assert per_expert_only._detect_shared_outer_loras() is False + # -- pool layout + derived operands ----------------------------------- + x = torch.tensor([[0.5, -1.0, 0.25], [1.25, 0.75, -0.5]]) + act = torch.tensor([[0.5, -0.25, 1.0, 0.75], [-1.0, 0.25, 0.5, -0.75]]) + assert x.shape[-1] == hidden and act.shape[-1] == experts * _SHARD + expected_gate_up_delta, expected_down_delta = _reference_delta(shards, x, act) - mixed = manager_cls() - mixed.loras = { - "shared": SimpleNamespace( - layers=[SimpleNamespace(weights={routed_shared: torch.empty(1, 8, 4)})] - ), - "per-expert": SimpleNamespace( - layers=[SimpleNamespace(weights={routed_expert: torch.empty(4, 4)})] - ), - } - with pytest.raises(RuntimeError, match="Mixed shared-outer LoRA formats"): - mixed._detect_shared_outer_loras() + # A pool padded past the adapter's rank must produce the same delta as an + # exactly-sized one: the up half's column offset comes from the padded + # max_rank, matching where the pool stacks the up half of LoRA-A. + for max_rank in (rank, rank + 1): + sink = _sink_with_shard(shards, max_rank=max_rank) + assert sink.experts_shared_outer_loras is True + gate_up_delta, down_delta = _consumer_delta(sink, x, act) + torch.testing.assert_close(gate_up_delta, expected_gate_up_delta) + torch.testing.assert_close(down_delta, expected_down_delta) if __name__ == "__main__": - import sys - sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/lora/test_inkling_lora_normalization_unit.py b/test/registered/unit/lora/test_inkling_lora_normalization_unit.py deleted file mode 100644 index 883f969b3..000000000 --- a/test/registered/unit/lora/test_inkling_lora_normalization_unit.py +++ /dev/null @@ -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"])) diff --git a/test/registered/unit/lora/test_inkling_moe_lora_overlap_unit.py b/test/registered/unit/lora/test_inkling_moe_lora_overlap_unit.py deleted file mode 100644 index c5b29c169..000000000 --- a/test/registered/unit/lora/test_inkling_moe_lora_overlap_unit.py +++ /dev/null @@ -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"]))