503 lines
18 KiB
Python
503 lines
18 KiB
Python
import unittest
|
|
|
|
import torch
|
|
|
|
from sglang.srt.distributed import init_distributed_environment
|
|
from sglang.srt.distributed.parallel_state import (
|
|
destroy_distributed_environment,
|
|
destroy_model_parallel,
|
|
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.runtime_context import get_context, publish
|
|
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
|
from sglang.test.test_utils import CustomTestCase, publish_build_topology
|
|
|
|
|
|
class TestFlashinferDispatcher(CustomTestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
# Dist-init first: world_size (and so the tp/ep width ServerArgs must
|
|
# carry) is only known after it, and init_distributed_environment
|
|
# itself reads no published config.
|
|
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)
|
|
|
|
server_args = ServerArgs(
|
|
model_path="dummy", tp_size=world_size, ep_size=world_size
|
|
)
|
|
server_args.moe_runner_backend = "flashinfer_cutlass"
|
|
server_args.moe_a2a_backend = "flashinfer"
|
|
cls.server_args = server_args
|
|
set_global_server_args_for_scheduler(server_args)
|
|
publish(server_args, role="scheduler")
|
|
initialize_moe_config()
|
|
|
|
publish_build_topology(tp_size=world_size, ep_size=world_size, world_rank=rank)
|
|
initialize_model_parallel()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
try:
|
|
from flashinfer.comm.trtllm_moe_alltoall import MoeAlltoAll
|
|
|
|
for workspace in MoeAlltoAll._WORKSPACE_CACHE.values():
|
|
mnnvl_mem = workspace.get("mnnvl_mem")
|
|
if mnnvl_mem is not None and "ptr" in vars(mnnvl_mem):
|
|
del mnnvl_mem.ptr
|
|
MoeAlltoAll._WORKSPACE_CACHE.clear()
|
|
except ImportError:
|
|
pass
|
|
destroy_model_parallel()
|
|
destroy_distributed_environment()
|
|
|
|
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 set_dispatch_type(self, dispatch_type):
|
|
get_context().override(
|
|
"test_flashinfer_dispatcher",
|
|
flashinfer_a2a_dispatch_type=dispatch_type,
|
|
)
|
|
|
|
def _zero_moe_a2a_dispatch_payloads(self):
|
|
# Shared MoeAlltoAll workspaces keep stale recv payloads across tests.
|
|
# Zero only the payload region so unused-source == 0 asserts stay valid.
|
|
try:
|
|
from flashinfer.comm.trtllm_moe_alltoall import (
|
|
MoeAlltoAll,
|
|
get_moe_alltoall_module,
|
|
)
|
|
except ImportError:
|
|
return
|
|
|
|
module = get_moe_alltoall_module()
|
|
for ws in MoeAlltoAll._WORKSPACE_CACHE.values():
|
|
workspace = ws["workspace"]
|
|
aux = int(
|
|
module.moe_a2a_get_aux_data_size(
|
|
ws["ep_size"],
|
|
ws["max_num_tokens"],
|
|
ws["eplb_stats_num_experts"],
|
|
)
|
|
)
|
|
aux = ((aux + 127) // 128) * 128
|
|
if aux < workspace.shape[1]:
|
|
workspace[:, aux:].zero_()
|
|
|
|
def test_dispatch_basic(self):
|
|
"""Test basic dispatch functionality"""
|
|
self.set_dispatch_type("bf16")
|
|
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)"""
|
|
self.set_dispatch_type("bf16")
|
|
# This tests the dummy token handling
|
|
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=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
|
|
)
|
|
|
|
self._zero_moe_a2a_dispatch_payloads()
|
|
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
|
|
|
|
# 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"""
|
|
self.set_dispatch_type("nvfp4")
|
|
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)
|
|
|
|
def test_dispatch_with_mxfp8_quantization(self):
|
|
"""Test dispatch with MXFP8 quantization enabled"""
|
|
self.set_dispatch_type("mxfp8")
|
|
num_tokens = 128
|
|
hidden_size = 128
|
|
router_topk = 1
|
|
world_size = torch.distributed.get_world_size()
|
|
rank = torch.distributed.get_rank()
|
|
num_experts = world_size
|
|
num_local_experts = 1
|
|
|
|
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,
|
|
)
|
|
|
|
hidden_states = torch.randn(
|
|
(num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
|
|
)
|
|
|
|
target_rank = (rank + 1) % world_size
|
|
target_expert = target_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,
|
|
)
|
|
dispatcher.set_quant_config({"input_global_scale": None, "use_mxfp8": True})
|
|
|
|
dispatch_output = dispatcher.dispatch(hidden_states, topk_output)
|
|
|
|
self.assertEqual(
|
|
dispatch_output.hidden_states.shape,
|
|
(num_tokens * world_size, hidden_size),
|
|
)
|
|
self.assertEqual(dispatch_output.hidden_states.dtype, torch.float8_e4m3fn)
|
|
self.assertEqual(dispatch_output.output_dtype, torch.bfloat16)
|
|
|
|
self.assertIsNotNone(dispatch_output.hidden_states_scale)
|
|
self.assertEqual(
|
|
dispatch_output.hidden_states_scale.shape,
|
|
(num_tokens * world_size, hidden_size // 32),
|
|
)
|
|
self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)
|
|
self.assertEqual(
|
|
dispatch_output.topk_output.topk_ids.shape,
|
|
(num_tokens * world_size, router_topk),
|
|
)
|
|
self.assertEqual(dispatch_output.topk_output.topk_ids.dtype, torch.int32)
|
|
|
|
def test_dispatch_with_mxfp8_quantization_and_empty_rank(self):
|
|
"""All ranks must contribute the same payload dtypes, including empty ranks."""
|
|
self.set_dispatch_type("mxfp8")
|
|
num_tokens = 16
|
|
hidden_size = 128
|
|
router_topk = 1
|
|
world_size = torch.distributed.get_world_size()
|
|
rank = torch.distributed.get_rank()
|
|
empty_rank = 1
|
|
|
|
global_num_tokens = [num_tokens] * world_size
|
|
global_num_tokens[empty_rank] = 0
|
|
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=global_num_tokens,
|
|
)
|
|
|
|
local_tokens = 0 if rank == empty_rank else num_tokens
|
|
hidden_states = torch.randn(
|
|
(local_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
|
|
)
|
|
target_expert = (rank + 1) % world_size
|
|
topk_ids = torch.full(
|
|
(local_tokens, router_topk),
|
|
target_expert,
|
|
dtype=torch.int32,
|
|
device="cuda",
|
|
)
|
|
topk_weights = torch.ones(
|
|
(local_tokens, router_topk), dtype=torch.float32, device="cuda"
|
|
)
|
|
|
|
from sglang.srt.layers.moe.topk import StandardTopKOutput
|
|
|
|
dispatcher = self.create_dispatcher(
|
|
router_topk=router_topk,
|
|
num_experts=world_size,
|
|
num_local_experts=1,
|
|
hidden_size=hidden_size,
|
|
)
|
|
dispatcher.set_quant_config({"input_global_scale": None, "use_mxfp8": True})
|
|
self._zero_moe_a2a_dispatch_payloads()
|
|
torch.distributed.barrier()
|
|
dispatch_output = dispatcher.dispatch(
|
|
hidden_states,
|
|
StandardTopKOutput(
|
|
topk_weights=topk_weights,
|
|
topk_ids=topk_ids,
|
|
router_logits=None,
|
|
),
|
|
)
|
|
|
|
self.assertEqual(dispatch_output.hidden_states.dtype, torch.float8_e4m3fn)
|
|
self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)
|
|
self.assertEqual(
|
|
dispatch_output.hidden_states.shape,
|
|
(num_tokens * world_size, hidden_size),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
"""
|
|
Usage
|
|
torchrun --nproc_per_node=4 test_flashinfer_dispatcher.py
|
|
"""
|
|
unittest.main()
|