move dead sglang.test files to test/manual (#25316)
This commit is contained in:
Liangsheng Yin
@@ -0,0 +1,322 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.distributed import init_distributed_environment
|
||||
from sglang.srt.distributed.parallel_state import (
|
||||
get_tp_group,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from sglang.srt.layers.dp_attention import set_dp_buffer_len
|
||||
from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher
|
||||
from sglang.srt.layers.moe.utils import initialize_moe_config
|
||||
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestFlashinferDispatcher(CustomTestCase):
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.moe_runner_backend = "flashinfer_cutlass"
|
||||
server_args.moe_a2a_backend = "flashinfer"
|
||||
set_global_server_args_for_scheduler(server_args)
|
||||
initialize_moe_config(server_args)
|
||||
|
||||
init_distributed_environment(
|
||||
world_size=-1, # Auto-detect from environment
|
||||
rank=-1, # Auto-detect from environment
|
||||
local_rank=-1, # Auto-detect from environment
|
||||
backend="nccl",
|
||||
)
|
||||
world_size = torch.distributed.get_world_size()
|
||||
rank = torch.distributed.get_rank()
|
||||
device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
|
||||
torch.cuda.set_device(device)
|
||||
initialize_model_parallel(
|
||||
tensor_model_parallel_size=world_size, expert_model_parallel_size=world_size
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
# Clean up distributed environment
|
||||
if torch.distributed.is_initialized():
|
||||
torch.distributed.destroy_process_group()
|
||||
|
||||
def create_dispatcher(
|
||||
self, router_topk=2, num_experts=8, num_local_experts=4, hidden_size=128
|
||||
):
|
||||
"""Helper to create dispatcher instance"""
|
||||
return FlashinferDispatcher(
|
||||
group=get_tp_group().device_group,
|
||||
router_topk=router_topk,
|
||||
num_experts=num_experts,
|
||||
num_local_experts=num_local_experts,
|
||||
hidden_size=hidden_size,
|
||||
params_dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
def test_dispatch_basic(self):
|
||||
"""Test basic dispatch functionality"""
|
||||
num_tokens = 16
|
||||
hidden_size = 128
|
||||
router_topk = 1 # Single expert per token for simplicity
|
||||
world_size = torch.distributed.get_world_size()
|
||||
rank = torch.distributed.get_rank()
|
||||
num_experts = world_size
|
||||
num_local_experts = 1 # One expert per rank
|
||||
|
||||
set_dp_buffer_len(
|
||||
global_dp_buffer_len=num_tokens * world_size,
|
||||
local_dp_buffer_len=num_tokens,
|
||||
dp_max_padding=True,
|
||||
global_num_tokens=None,
|
||||
)
|
||||
|
||||
# Create tokens with rank number
|
||||
hidden_states = torch.full(
|
||||
(num_tokens, hidden_size), 100.0 + rank, dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
|
||||
# Route all tokens from rank i to expert (i+1) % world_size
|
||||
target_rank = (rank + 1) % world_size
|
||||
target_expert = target_rank # Since we have 1 expert per rank
|
||||
|
||||
topk_ids = torch.full(
|
||||
(num_tokens, router_topk), target_expert, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
topk_weights = torch.ones(
|
||||
(num_tokens, router_topk), dtype=torch.float32, device="cuda"
|
||||
)
|
||||
|
||||
from sglang.srt.layers.moe.topk import StandardTopKOutput
|
||||
|
||||
topk_output = StandardTopKOutput(
|
||||
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
|
||||
)
|
||||
|
||||
torch.distributed.barrier()
|
||||
dispatcher = self.create_dispatcher(
|
||||
router_topk=router_topk,
|
||||
num_experts=num_experts,
|
||||
num_local_experts=num_local_experts,
|
||||
hidden_size=hidden_size,
|
||||
)
|
||||
dispatcher.set_quant_config({"input_global_scale": None})
|
||||
|
||||
dispatch_output = dispatcher.dispatch(hidden_states, topk_output)
|
||||
received_hidden_states = dispatch_output.hidden_states
|
||||
self.assertEqual(dispatch_output.hidden_states_scale, None)
|
||||
|
||||
# Expected: we should receive tokens from rank (rank - 1) % world_size
|
||||
expected_source_rank = (rank - 1 + world_size) % world_size
|
||||
|
||||
# Verify we received the right number of tokens
|
||||
self.assertEqual(
|
||||
received_hidden_states.shape[0],
|
||||
num_tokens * world_size,
|
||||
f"Should receive {num_tokens * world_size} tokens",
|
||||
)
|
||||
|
||||
# Verify tokens came from the expected source
|
||||
self.assertTrue(
|
||||
torch.all(
|
||||
received_hidden_states[
|
||||
expected_source_rank
|
||||
* num_tokens : (expected_source_rank + 1)
|
||||
* num_tokens
|
||||
]
|
||||
== 100.0 + expected_source_rank
|
||||
)
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.all(
|
||||
received_hidden_states[: expected_source_rank * num_tokens] == 0.0
|
||||
)
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.all(
|
||||
received_hidden_states[(expected_source_rank + 1) * num_tokens :] == 0.0
|
||||
)
|
||||
)
|
||||
|
||||
def test_dispatch_with_empty_tokens(self):
|
||||
"""Test dispatch when there are no tokens (edge case)"""
|
||||
# This tests the dummy token handling
|
||||
num_tokens = 16
|
||||
hidden_size = 1
|
||||
router_topk = 1 # Single expert per token for simplicity
|
||||
world_size = torch.distributed.get_world_size()
|
||||
rank = torch.distributed.get_rank()
|
||||
num_experts = world_size
|
||||
num_local_experts = 1 # One expert per rank
|
||||
|
||||
set_dp_buffer_len(
|
||||
global_dp_buffer_len=num_tokens * world_size,
|
||||
local_dp_buffer_len=num_tokens,
|
||||
dp_max_padding=False,
|
||||
global_num_tokens=[16, 0, 16, 16],
|
||||
)
|
||||
|
||||
# Route all tokens from rank i to expert (i+1) % world_size
|
||||
target_rank = (rank + 1) % world_size
|
||||
target_expert = target_rank # Since we have 1 expert per rank
|
||||
|
||||
# Create tokens with rank number, rank 1 has no tokens
|
||||
if rank == 1:
|
||||
hidden_states = torch.empty(
|
||||
0, hidden_size, dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
topk_ids = torch.empty(0, router_topk, dtype=torch.int32, device="cuda")
|
||||
topk_weights = torch.empty(
|
||||
0, router_topk, dtype=torch.float32, device="cuda"
|
||||
)
|
||||
else:
|
||||
hidden_states = torch.full(
|
||||
(num_tokens, hidden_size),
|
||||
100.0 + rank,
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
)
|
||||
topk_ids = torch.full(
|
||||
(num_tokens, router_topk),
|
||||
target_expert,
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
topk_weights = torch.ones(
|
||||
(num_tokens, router_topk), dtype=torch.float32, device="cuda"
|
||||
)
|
||||
|
||||
from sglang.srt.layers.moe.topk import StandardTopKOutput
|
||||
|
||||
topk_output = StandardTopKOutput(
|
||||
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
|
||||
)
|
||||
|
||||
dispatcher = self.create_dispatcher(
|
||||
router_topk=router_topk,
|
||||
num_experts=num_experts,
|
||||
num_local_experts=num_local_experts,
|
||||
hidden_size=hidden_size,
|
||||
)
|
||||
dispatcher.set_quant_config({"input_global_scale": None})
|
||||
|
||||
dispatch_output = dispatcher.dispatch(hidden_states, topk_output)
|
||||
received_hidden_states = dispatch_output.hidden_states
|
||||
|
||||
# Expected: we should receive tokens from rank (rank - 1) % world_size
|
||||
expected_source_rank = (rank - 1 + world_size) % world_size
|
||||
|
||||
# Verify we received the right number of tokens
|
||||
self.assertEqual(
|
||||
received_hidden_states.shape[0],
|
||||
num_tokens * world_size,
|
||||
f"Should receive {num_tokens * world_size} tokens",
|
||||
)
|
||||
|
||||
# Verify tokens came from the expected source
|
||||
if rank == 2:
|
||||
# Rank 2 should receive no tokens since rank 1 was empty
|
||||
self.assertTrue(
|
||||
torch.all(received_hidden_states == 0.0),
|
||||
"Rank should receive no tokens",
|
||||
)
|
||||
else:
|
||||
self.assertTrue(
|
||||
torch.all(
|
||||
received_hidden_states[
|
||||
expected_source_rank
|
||||
* num_tokens : (expected_source_rank + 1)
|
||||
* num_tokens
|
||||
]
|
||||
== 100.0 + expected_source_rank
|
||||
),
|
||||
"Rank {rank} should receive tokens from the expected source {expected_source_rank}",
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.all(
|
||||
received_hidden_states[: expected_source_rank * num_tokens] == 0.0
|
||||
),
|
||||
"Rank should receive no tokens from previous ranks",
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.all(
|
||||
received_hidden_states[(expected_source_rank + 1) * num_tokens :]
|
||||
== 0.0
|
||||
),
|
||||
"Rank should receive no tokens from next ranks",
|
||||
)
|
||||
|
||||
def test_dispatch_with_fp4_quantization(self):
|
||||
"""Test dispatch with FP4 quantization enabled"""
|
||||
num_tokens = 128
|
||||
hidden_size = 128
|
||||
router_topk = 1 # Single expert per token for simplicity
|
||||
world_size = torch.distributed.get_world_size()
|
||||
rank = torch.distributed.get_rank()
|
||||
num_experts = world_size
|
||||
num_local_experts = 1 # One expert per rank
|
||||
|
||||
set_dp_buffer_len(
|
||||
global_dp_buffer_len=num_tokens * world_size,
|
||||
local_dp_buffer_len=num_tokens,
|
||||
dp_max_padding=True,
|
||||
global_num_tokens=None,
|
||||
)
|
||||
|
||||
# Create tokens with random values
|
||||
hidden_states = torch.randn(
|
||||
(num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
|
||||
# Route all tokens from rank i to expert (i+1) % world_size
|
||||
target_rank = (rank + 1) % world_size
|
||||
target_expert = target_rank # Since we have 1 expert per rank
|
||||
|
||||
topk_ids = torch.full(
|
||||
(num_tokens, router_topk), target_expert, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
topk_weights = torch.ones(
|
||||
(num_tokens, router_topk), dtype=torch.float32, device="cuda"
|
||||
)
|
||||
|
||||
from sglang.srt.layers.moe.topk import StandardTopKOutput
|
||||
|
||||
topk_output = StandardTopKOutput(
|
||||
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
|
||||
)
|
||||
|
||||
dispatcher = self.create_dispatcher(
|
||||
router_topk=router_topk,
|
||||
num_experts=num_experts,
|
||||
num_local_experts=num_local_experts,
|
||||
hidden_size=hidden_size,
|
||||
)
|
||||
# Set input global scale to enable FP4 quantization
|
||||
input_global_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
|
||||
dispatcher.set_quant_config({"input_global_scale": input_global_scale})
|
||||
|
||||
dispatch_output = dispatcher.dispatch(hidden_states, topk_output)
|
||||
|
||||
self.assertEqual(
|
||||
dispatch_output.hidden_states.shape,
|
||||
(num_tokens * world_size, hidden_size // 2),
|
||||
)
|
||||
self.assertEqual(dispatch_output.hidden_states.dtype, torch.uint8)
|
||||
|
||||
self.assertNotEqual(dispatch_output.hidden_states_scale, None)
|
||||
self.assertEqual(
|
||||
dispatch_output.hidden_states_scale.numel(),
|
||||
num_tokens * world_size * (hidden_size // 16),
|
||||
)
|
||||
self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
"""
|
||||
Usage
|
||||
torchrun --nproc_per_node=4 test_flashinfer_dispatcher.py
|
||||
"""
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user