Support true on-policy (#12058)

test/srt/cpu/test_activation.py

@@ -6,6 +6,7 @@ import torch
 import torch.nn.functional as F
 from utils import GeluAndMul, SiluAndMul, precision
 
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.test_utils import CustomTestCase
 
 torch.manual_seed(1234)
@@ -17,6 +18,8 @@ class TestActivation(CustomTestCase):
     dtype = [torch.float16, torch.bfloat16]
 
     def _silu_and_mul_test(self, m, n, dtype):
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
         x = torch.randn([m, n], dtype=dtype)
 
         out = torch.ops.sgl_kernel.silu_and_mul_cpu(x)

test/srt/cpu/test_shared_expert.py

@@ -20,6 +20,7 @@ from utils import (
     torch_w8a8_per_column_moe,
 )
 
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.test_utils import CustomTestCase
 
 torch.manual_seed(1234)
@@ -149,6 +150,8 @@ class TestSharedExpert(CustomTestCase):
                 self._int8_shared_expert(*params)
 
     def _fp8_shared_expert(self, M, N, K, routed_scaling_factor):
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
         dtype = torch.bfloat16
         prepack = True
 

test/srt/quant/test_block_int8.py

@@ -6,6 +6,7 @@ import torch
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_moe
 from sglang.srt.layers.moe.topk import TopKConfig, select_experts
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -96,6 +97,8 @@ def native_w8a8_block_int8_matmul(A, B, As, Bs, block_size, output_dtype=torch.f
 def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
     """This function performs fused moe with block-wise quantization using native torch."""
 
+    set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
     B, D = a.shape
     a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
     out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)

test/srt/quant/test_int8_kernel.py

@@ -7,6 +7,7 @@ from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_moe
 from sglang.srt.layers.moe.topk import TopKConfig, select_experts
 from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -35,6 +36,8 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
 def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk):
     """This function performs fused moe with per-column int8 quantization using native torch."""
 
+    set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
     B, D = a.shape
     # Perform per-token quantization
     a_q, a_s = per_token_quant_int8(a)

test/srt/test_fused_moe.py

@@ -9,6 +9,7 @@ from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_moe
 from sglang.srt.layers.moe.topk import TopKConfig, select_experts
 from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.layers.quantization.fp8_utils import normalize_e4m3fn_to_e4m3fnuz
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.srt.utils import is_hip
 from sglang.test.test_utils import CustomTestCase
 
@@ -63,6 +64,8 @@ class TestFusedMOE(CustomTestCase):
         a1_scale=None,
         a2_scale=None,
     ):
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
         B, D = a.shape
         a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
         out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)

test/srt/test_triton_fused_moe.py

@@ -9,6 +9,7 @@ from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
 from sglang.srt.layers.moe.moe_runner.triton_kernels import TritonKernelsQuantInfo
 from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
 from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -56,6 +57,8 @@ class TestFusedMOE(CustomTestCase):
         topk,
         return_per_expert: bool = False,
     ):
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
         B, D = a.shape
         a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
         out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)

test/srt/test_triton_moe_channel_fp8_kernel.py

@@ -7,6 +7,7 @@ from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_moe
 from sglang.srt.layers.moe.topk import TopKConfig, select_experts
 from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -40,6 +41,8 @@ def fp8_mask(a, mask):
 def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk):
     """This function performs fused moe with per-column int8 quantization using native torch."""
 
+    set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
     B, D = a.shape
     # Perform per-token quantization
     a_q, a_s = scaled_fp8_quant(a, use_per_token_if_dynamic=True)

test/srt/test_triton_moe_wna16.py

@@ -6,6 +6,7 @@ import torch
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_moe
 from sglang.srt.layers.moe.topk import TopKConfig, select_experts
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 
 NUM_EXPERTS = [8, 64]
 TOP_KS = [2, 6]
@@ -116,6 +117,8 @@ def quantize_weights(
 
 
 def torch_moe(a, w1, w2, score, topk):
+    set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
     B, D = a.shape
     a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
     out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)