323 lines
12 KiB
Python
323 lines
12 KiB
Python
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()
|