diff --git a/docker/Dockerfile b/docker/Dockerfile index 5af24ff30f13..ad488626310f 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -18,7 +18,7 @@ ARG PIP_DEFAULT_INDEX ARG UBUNTU_MIRROR ARG GITHUB_ARTIFACTORY=github.com ARG INSTALL_FLASHINFER_JIT_CACHE=0 -ARG FLASHINFER_VERSION=0.7.2 +ARG FLASHINFER_VERSION=0.7.3 ARG MOONCAKE_VERSION=0.3.14 ARG MSCCLPP_VERSION=sglang-v0.10.1 diff --git a/docker/Dockerfile.cu134 b/docker/Dockerfile.cu134 index 64ef7a17e01b..6f40a749e408 100644 --- a/docker/Dockerfile.cu134 +++ b/docker/Dockerfile.cu134 @@ -34,7 +34,7 @@ ARG PIP_DEFAULT_INDEX ARG UBUNTU_MIRROR ARG GITHUB_ARTIFACTORY=github.com ARG INSTALL_FLASHINFER_JIT_CACHE=1 -ARG FLASHINFER_VERSION=0.7.2 +ARG FLASHINFER_VERSION=0.7.3 ARG MOONCAKE_VERSION=0.3.13.post1 ARG MSCCLPP_VERSION=sglang-v0.10.1 diff --git a/python/pyproject.toml b/python/pyproject.toml index cde9e6303183..dd4433e28d51 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -38,7 +38,7 @@ dependencies = [ "easydict", # Required by remote model code (e.g. DeepSeek-OCR) loaded via trust_remote_code; validated by transformers 5.4+ check_imports "einops", "fastapi", - "flashinfer_python[cu13]==0.7.2", # keep it aligned with jit-cache version in Dockerfile + "flashinfer_python[cu13]==0.7.3", # keep it aligned with jit-cache version in Dockerfile "gguf", "humming-kernels[cu13]==0.1.18", "interegular", diff --git a/python/sglang/kernels/jit/csrc/minimax/nvfp4_store_kv_index.cuh b/python/sglang/kernels/jit/csrc/minimax/nvfp4_store_kv_index.cuh index 21c35647290c..1bb006e3f086 100644 --- a/python/sglang/kernels/jit/csrc/minimax/nvfp4_store_kv_index.cuh +++ b/python/sglang/kernels/jit/csrc/minimax/nvfp4_store_kv_index.cuh @@ -63,7 +63,8 @@ SGL_DEVICE uint32_t f32x8_to_e2m1x8(const float (&x)[8]) { /// /// Reproduces the reference operation order exactly: /// `sf = rcp(gs) * (amax * rcp(6))`, rounded to E4M3 and read back, -/// `osc = amax != 0 ? rcp(sf * gs) : 0`, +/// `osc = sf != 0 ? rcp(sf * gs) : 0`, so a block whose scale underflows +/// to zero stores an all-zero payload, /// each value scaled by `osc` then `cvt.rn.satfinite.e2m1x2.f32`. /// `amax` is exact regardless of reduction order because `fabsf` of a BF16 /// value is representable in BF16. @@ -86,7 +87,7 @@ SGL_DEVICE uint64_t quantize_block(const bf16_t* __restrict__ src, float global_ const __nv_fp8_e4m3 sf_e4m3 = __nv_fp8_e4m3(sf); *scale_out = sf_e4m3.__x; sf = static_cast(sf_e4m3); - const float osc = amax != 0.0f ? rcp_approx(sf * global_scale) : 0.0f; + const float osc = sf != 0.0f ? rcp_approx(sf * global_scale) : 0.0f; float lo[8], hi[8]; #pragma unroll diff --git a/python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh b/python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh index ac46f636bcb5..a5999485550a 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh @@ -134,6 +134,8 @@ __global__ void TreeSpeculativeSamplingTargetOnly( DType u = coin * sum_relu_q_minus_p; DType aggregate_relu_q_minus_p(0); + // Each thread's fallback candidate; FlashInfer reduces it only if the scan selects nothing. + int last_valid_index = -1; for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { q_vec.fill(DType(0)); p_vec.fill(DType(0)); @@ -152,7 +154,14 @@ __global__ void TreeSpeculativeSamplingTargetOnly( } DeviceSamplingFromProb( - i, d, [&](DType x) { return x > 0; }, u, relu_q_minus_p_vec, aggregate_relu_q_minus_p, &temp_storage); + i, + d, + [&](DType x) { return x > 0; }, + u, + relu_q_minus_p_vec, + aggregate_relu_q_minus_p, + last_valid_index, + &temp_storage); if (aggregate_relu_q_minus_p > u) { break; } diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 46b645de77be..88143bbd6910 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -1994,7 +1994,7 @@ def _set_envs_and_config(server_args: ServerArgs): ): assert_pkg_version( "flashinfer_python", - "0.7.2", + "0.7.3", "Please uninstall the old version and " "reinstall the latest version by following the instructions " "at https://docs.flashinfer.ai/installation.html.", diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 888b0421bdea..9ab4de936aff 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -307,6 +307,10 @@ def fast_prefill_plan( ) self._cached_o_data_type = o_data_type self._block_tables = None + # FlashInfer's plan() clears this on every call and may set it again for its + # SM90 fp8-KV verify kernel. This plan never takes that route, so clear it + # here too rather than let run() launch it with a previous plan's state. + self._eagle_verify_fp8kv_sm90_state = None args = [ self._float_workspace_buffer, @@ -329,6 +333,7 @@ def fast_prefill_plan( False, # disable_split_kv 0, # num_colocated_ctas 0, # uniform_q_len + False, # fixed_uniform_q_len ] self._plan_info = self._cached_module.plan(*args) diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 3a02a8af107a..aabf20e47cd0 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -34,7 +34,12 @@ register_fused_func, ) from sglang.srt.layers.utils import copy_or_rebind_param -from sglang.srt.runtime_context import get_forward, get_parallel +from sglang.srt.runtime_context import ( + get_context, + get_exec, + get_forward, + get_parallel, +) from sglang.srt.utils.device import is_flashinfer_available from sglang.srt.utils.triton_utils import next_power_of_2 @@ -130,6 +135,24 @@ def _make_deferred_finalize_output( ) +def _finalize_in_order(result, *, top_k: int, out: torch.Tensor) -> torch.Tensor: + """Sum FlashInfer's ``do_finalize=False`` outputs into ``out``, adding the + top-k terms in the same order whatever the batch size.""" + from sglang.kernels.ops.moe.moe_finalize_fuse_shared import moe_finalize_fuse_shared + + deferred = _make_deferred_finalize_output(result, top_k=top_k) + num_tokens = out.shape[0] + return moe_finalize_fuse_shared( + deferred.gemm2_out.view(-1, deferred.gemm2_out.shape[-1]), + deferred.expanded_idx_to_permuted_idx.view(-1)[: num_tokens * top_k], + deferred.expert_weights.view(-1, top_k)[:num_tokens], + None, + top_k, + enable_pdl=trtllm_moe_enable_pdl(num_tokens), + out=out, + ) + + def round_up_to_multiple(x: int, m: int) -> int: """Round up *x* to the nearest multiple of *m*.""" return (x + m - 1) // m * m @@ -1559,6 +1582,22 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( symmetric=hidden_pad == 0, ) + # FlashInfer's small-batch finalize sums the top-k outputs in a different order + # from its large-batch finalize (flashinfer-ai/flashinfer#8831), so a token's + # output would depend on its batch size. In deterministic mode, take the + # unsummed outputs and sum them in one fixed order at every batch size. Only + # the pre-routed call: its routing weights already carry any routed scaling. + finalize_in_order = ( + use_routed_topk + and not defer_finalize + and hidden_pad == 0 + and symm_output.dtype == torch.bfloat16 + # Kernel tests call this runner without a published server config. + and get_context().is_config_namespace_published("exec") + and get_exec().deterministic.enable_deterministic_inference + ) + flashinfer_finalize = not (defer_finalize or finalize_in_order) + if use_routed_topk: routing = _get_routing_for_flashinfer_routed(topk_output) moe_kwargs = dict( @@ -1588,7 +1627,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( local_num_experts=quant_info.local_num_experts, routed_scaling_factor=None, routing_method_type=1, # Unused, but must be 1 to pass validation. - do_finalize=not defer_finalize, + do_finalize=flashinfer_finalize, activation_type=activation_type, tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]), enable_pdl=trtllm_moe_enable_pdl(hs_fp4.shape[0]), @@ -1601,6 +1640,10 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( result, top_k=_routing_top_k(routing), ) + elif finalize_in_order: + result = _finalize_in_order( + result, top_k=_routing_top_k(routing), out=symm_output + ) else: result = result[0] else: diff --git a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py index 261e9fc2bd61..174ba5b2e8a2 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py @@ -41,16 +41,18 @@ get_flags, get_forward, get_parallel, - get_spec, max_prefill_buffer_tokens, ) -from sglang.srt.speculative.spec_info import SpeculativeAlgorithm try: from flashinfer import nvfp4_block_scale_interleave - from flashinfer.comm import MoeAlltoAll, moe_a2a_get_workspace_size_per_rank - from flashinfer.comm.mapping import Mapping - from flashinfer.comm.mnnvl import MnnvlConfig + from flashinfer.moe_ep import ( + BootstrapConfig, + MoEEpCommParams, + NVLinkOneSidedAlltoAll, + NVLinkOneSidedConfig, + create_communication, + ) from sglang.srt.layers.quantization.fp4_utils import fp4_quantize @@ -60,11 +62,11 @@ logger = logging.getLogger(__name__) -# FlashInfer keys MNNVL allocations by workspace size; aligned tail padding gives -# concurrently live paths distinct persistent workspaces without extra token work. -_WORKSPACE_NAMESPACE_ALIGNMENT = 128 - +# FlashInfer shares one workspace among communication instances with the same +# geometry, and those instances must not interleave dispatch/combine rounds. +# Padding the token capacity by the namespace gives concurrently live paths +# distinct persistent workspaces without extra token work. class _WorkspaceNamespace(enum.IntEnum): DECODE = 0 SPECULATIVE = 1 @@ -92,8 +94,8 @@ def _scattered_source_token_counts( return counts -def _workspace_size_for_namespace( - workspace_size: int, +def _max_tokens_for_namespace( + max_num_tokens: int, *, namespace: Optional[_WorkspaceNamespace] = None, speculative: Optional[bool] = None, @@ -106,7 +108,7 @@ def _workspace_size_for_namespace( if speculative else _WorkspaceNamespace.DECODE ) - return workspace_size + int(namespace) * _WORKSPACE_NAMESPACE_ALIGNMENT + return max_num_tokens + int(namespace) class FlashinferDispatchOutput(NamedTuple): @@ -195,90 +197,45 @@ def __init__( moe_runner_config = MoeRunnerConfig( num_experts=num_experts, num_local_experts=num_local_experts, + num_fused_shared_experts=0, hidden_size=hidden_size, top_k=router_topk, ) self.prefill_dispatcher = StandardDispatcher(moe_runner_config) # FlashInfer sizes the workspace from the maximum dispatched tokens per - # EP rank. See FlashInfer's moe_a2a_get_workspace_size_per_rank(), - # which reserves ep_size * max_num_tokens * payload bytes, and the C++ - # dispatch op's epSize * runtimeMaxTokensPerRank payload buffer. + # EP rank: its dispatch receive buffer is ep_size * max_tokens_per_rank rows. # # Cover both the largest eager prefill chunk (including PP dynamic # chunking) and the largest decode batch using upstream's shared bound. self.max_num_tokens = flashinfer_a2a_max_dispatch_tokens_per_rank( max_prefill_buffer_tokens() ) - speculative_algo = SpeculativeAlgorithm.from_string( - get_spec().speculative_algorithm + bootstrap = BootstrapConfig( + world_size=self.ep_size, rank=self.ep_rank, process_group=group ) - can_use_quantized_dispatch = not speculative_algo.is_eagle() - topk_id_and_weight_bytes = self.router_topk * 4 + self.router_topk * 4 - bf16_dispatch_payload_size_per_token = ( - hidden_size * 2 + topk_id_and_weight_bytes # bf16 hidden states + a2a_config = NVLinkOneSidedConfig( + comm_backend=TorchDistributedCommBackend(group) ) - if ( - self.dispatch_type == FlashinferA2ADispatchType.NVFP4 - and can_use_quantized_dispatch - ): - quantized_dispatch_payload_size_per_token = ( - hidden_size // 2 # nvfp4 hidden states - + hidden_size // 16 # uint8 scaling factors - + topk_id_and_weight_bytes - ) - total_dispatch_payload_size_per_token = max( - quantized_dispatch_payload_size_per_token, - bf16_dispatch_payload_size_per_token, - ) - elif ( - self.dispatch_type == FlashinferA2ADispatchType.MXFP8 - and can_use_quantized_dispatch - ): - quantized_dispatch_payload_size_per_token = ( - hidden_size # fp8 hidden states - + hidden_size // 32 # ue8m0 scaling factors - + topk_id_and_weight_bytes - ) - total_dispatch_payload_size_per_token = max( - quantized_dispatch_payload_size_per_token, - bf16_dispatch_payload_size_per_token, - ) - else: - total_dispatch_payload_size_per_token = bf16_dispatch_payload_size_per_token - combine_payload_size_per_token = hidden_size * 2 # bf16 hidden states - self.workspace_size = moe_a2a_get_workspace_size_per_rank( - ep_size=self.ep_size, - max_num_tokens=self.max_num_tokens, - total_dispatch_payload_size_per_token=total_dispatch_payload_size_per_token, - combine_payload_size_per_token=combine_payload_size_per_token, - ) - - self.mapping = Mapping( - rank=self.ep_rank, - tp_size=self.ep_size, - moe_ep_size=self.ep_size, - world_size=self.ep_size, - gpus_per_node=torch.cuda.device_count(), - pp_size=1, - cp_size=1, - ) - mnnvl_config = MnnvlConfig(comm_backend=TorchDistributedCommBackend(group)) is_speculative_model = get_flags().moe.speculative_context - def make_moe_a2a(namespace: _WorkspaceNamespace) -> MoeAlltoAll: - workspace_size = _workspace_size_for_namespace( - self.workspace_size, - namespace=namespace, - ) - return MoeAlltoAll( - mapping=self.mapping, - max_num_tokens=self.max_num_tokens, - top_k=self.router_topk, + def make_moe_a2a(namespace: _WorkspaceNamespace) -> NVLinkOneSidedAlltoAll: + params = MoEEpCommParams( num_experts=self.num_experts, - workspace_size_per_rank=workspace_size, - mnnvl_config=mnnvl_config, + top_k=self.router_topk, + max_tokens_per_rank=_max_tokens_for_namespace( + self.max_num_tokens, namespace=namespace + ), + hidden_size=self.hidden_size, + # No dispatch_format: BF16 rows bound every quantized format and + # cover layers whose dispatch falls back to BF16 (draft, EAGLE). + dtype=torch.bfloat16, + invalid_expert_id=self.invalid_token_expert_id, ) + # FlashInfer's host-side workspace metadata must stay on the CPU even + # when layers are built under a CUDA default-device context. + with torch.device("cpu"): + return create_communication(bootstrap, params, backend=a2a_config) decode_namespace = ( _WorkspaceNamespace.SPECULATIVE @@ -309,11 +266,7 @@ def _start_a2a(self, *, sp_prefill: bool) -> None: ) def _clear_active_a2a(self) -> None: - for attr in ( - "runtime_max_tokens_per_rank", - "active_moe_a2a", - "active_a2a_namespace", - ): + for attr in ("active_moe_a2a", "active_a2a_namespace"): if hasattr(self, attr): delattr(self, attr) @@ -424,8 +377,8 @@ def dispatch( # FlashInfer dispatch requires materialized top-k IDs and weights. if TopKOutputChecker.format_is_bypassed(topk_output): topk_output = topk_output.to_standard() - # FlashInfer MoeAlltoAll's expert-ID ABI is int32. This dispatcher is - # only selected for moe_a2a_backend="flashinfer". + # FlashInfer NVLinkOneSidedAlltoAll's expert-ID ABI is int32. This + # dispatcher is only selected for moe_a2a_backend="flashinfer". topk_ids = topk_output.topk_ids.to(torch.int32) topk_weights = topk_output.topk_weights @@ -460,19 +413,9 @@ def dispatch( 0, self.hidden_size // 32, dtype=torch.uint8, device=x.device ) - payloads = [] - payloads.append(x) - if x_sf is not None: - payloads.append(x_sf) - expert_id_payload_index = 2 - else: - expert_id_payload_index = 1 - payloads.append(topk_ids) - payloads.append(topk_weights) - # runtime_max_tokens_per_rank selection # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # MoeAlltoAll uses fixed-geometry buffers shaped + # NVLinkOneSidedAlltoAll uses fixed-geometry buffers shaped # [ep_size, runtime_max_tokens_per_rank, ...], so every EP rank must pass # the SAME value. This code (Python) runs during eager forwards and during # CUDA-graph *capture*; on *replay* dispatch() is not re-executed and the @@ -507,8 +450,8 @@ def dispatch( ) runtime_max_tokens_per_rank = x.shape[0] - # MoeAlltoAll does not resize its max_num_tokens workspace; reject larger - # runtime geometry here before it becomes an illegal memory access. + # The FlashInfer workspace is sized once for max_num_tokens; reject larger + # runtime geometry here with a sizing hint rather than FlashInfer's error. assert runtime_max_tokens_per_rank <= self.max_num_tokens, ( "FlashInfer A2A runtime token geometry exceeds its fixed workspace: " f"runtime_max_tokens_per_rank={runtime_max_tokens_per_rank} > " @@ -525,54 +468,42 @@ def dispatch( # count) rather than a benign case. assert runtime_max_tokens_per_rank >= x.shape[0], ( f"runtime_max_tokens_per_rank={runtime_max_tokens_per_rank} < " - f"x.shape[0]={x.shape[0]}: MoeAlltoAll recv buffer would overflow." + f"x.shape[0]={x.shape[0]}: A2A recv buffer would overflow." ) - # Passing topk_ids + invalid_token_expert_id triggers the sanitize step - # inside moe_a2a. The recv buffer has shape - # [ep_size, max_tokens_per_rank, ...], so any rank below max leaves - # padding slots whose expert_id would otherwise route to a real expert - # and waste downstream MoE compute. Sanitizing the padding to a - # sentinel id is structural, not optional. - self.runtime_max_tokens_per_rank = runtime_max_tokens_per_rank + # The recv buffer has shape [ep_size * runtime_max_tokens_per_rank, ...], + # so any rank below max leaves padding slots; FlashInfer sanitizes their + # expert ids to MoEEpCommParams.invalid_expert_id so that downstream MoE + # compute skips them. self._start_a2a(sp_prefill=is_sp_prefill) try: - recv_tensors = self.active_moe_a2a.dispatch( + received = self.active_moe_a2a.dispatch( + x, topk_ids, - payloads, - runtime_max_tokens_per_rank, - invalid_token_expert_id=self.invalid_token_expert_id, - expert_id_payload_index=expert_id_payload_index, + topk_weights, + hidden_states_scale=x_sf, + runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, ) - if x_sf is not None: - x_recv, x_sf_recv, topk_ids_recv, topk_weights_recv = recv_tensors - x_sf = x_sf_recv.view(-1, x_sf_recv.shape[-1]) - # TODO: fuse interleave into cutlass moe - if dispatch_type == FlashinferA2ADispatchType.NVFP4 and ( - self.runner_backend is not None - and self.runner_backend.is_flashinfer_cutlass() - ): - x_sf = nvfp4_block_scale_interleave(x_sf) - else: - x_recv, topk_ids_recv, topk_weights_recv = recv_tensors - x = x_recv.view(-1, x_recv.shape[-1]) - if dispatch_type == FlashinferA2ADispatchType.MXFP8: - x = x.view(torch.float8_e4m3fn) - topk_ids = topk_ids_recv.view(-1, topk_ids_recv.shape[-1]) - topk_weights = topk_weights_recv.view(-1, topk_weights_recv.shape[-1]) + x = received.hidden_states + x_sf = received.hidden_states_scale + # TODO: fuse interleave into cutlass moe + if ( + dispatch_type == FlashinferA2ADispatchType.NVFP4 + and self.runner_backend is not None + and self.runner_backend.is_flashinfer_cutlass() + ): + x_sf = nvfp4_block_scale_interleave(x_sf) # Provide an output tensor to fused_moe so it writes directly to our buffer moe_output = None if self.payload_in_workspace: - moe_output = ( - self.active_moe_a2a.get_combine_payload_tensor_in_workspace( - runtime_max_tokens_per_rank, self.hidden_size, output_dtype - ).view(-1, self.hidden_size) - ) + moe_output = self.active_moe_a2a.get_combine_input_buffer(output_dtype) return FlashinferDispatchOutput( x, x_sf, - StandardTopKOutput(topk_weights, topk_ids, topk_output.router_logits), + StandardTopKOutput( + received.topk_weights, received.topk_ids, topk_output.router_logits + ), moe_output, output_dtype, ) @@ -621,16 +552,8 @@ def combine( "FlashInfer A2A FP8 MoE combine payload must be BF16, but " f"received {hidden_states.dtype}." ) - output_hidden_size = hidden_states.shape[-1] try: - return self.active_moe_a2a.combine( - hidden_states.view( - self.ep_size, - self.runtime_max_tokens_per_rank, - output_hidden_size, - ), - self.runtime_max_tokens_per_rank, - payload_in_workspace=self.payload_in_workspace, - ) + # FlashInfer detects a payload already in its combine input buffer. + return self.active_moe_a2a.combine(hidden_states) finally: self._clear_active_a2a() diff --git a/test/registered/kernels/ops/kvcache/test_minimax_m31_nvfp4.py b/test/registered/kernels/ops/kvcache/test_minimax_m31_nvfp4.py index 4b646f6eb8a4..839ac9d2ed5b 100644 --- a/test/registered/kernels/ops/kvcache/test_minimax_m31_nvfp4.py +++ b/test/registered/kernels/ops/kvcache/test_minimax_m31_nvfp4.py @@ -48,7 +48,7 @@ def _zeros(shape): def _tiny(shape): # Below the E4M3 scale underflow point, where the reference writes a zero - # scale and a saturated payload. + # scale and a zero payload. return torch.full(shape, 1e-4, dtype=torch.bfloat16, device="cuda") diff --git a/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py b/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py index 82f9ac39f249..7c1468f253ac 100644 --- a/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py +++ b/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py @@ -8,8 +8,8 @@ PermuteMethodPool, ) from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( + _max_tokens_for_namespace, _max_tokens_per_scattered_source, - _workspace_size_for_namespace, ) from sglang.srt.layers.quantization import fp8 # noqa: F401 from sglang.test.ci.ci_register import register_cuda_ci @@ -38,11 +38,11 @@ def test_dp4_tp4_uses_physical_source_rank_geometry(self): self.assertEqual(_max_tokens_per_scattered_source([512] * 16, 1), 512) def test_target_and_draft_decode_use_distinct_workspaces(self): - sizes = { - _workspace_size_for_namespace(4096, speculative=speculative) + capacities = { + _max_tokens_for_namespace(4096, speculative=speculative) for speculative in (False, True) } - self.assertEqual(sizes, {4096, 4224}) + self.assertEqual(capacities, {4096, 4097}) def test_deepgemm_dispatch_marks_empty_expert_lanes_invalid(self): topk_ids = torch.tensor([-1, 0, -1, 1], dtype=torch.int32, device="cuda") diff --git a/test/registered/unit/layers/moe/test_flashinfer_dispatcher.py b/test/registered/unit/layers/moe/test_flashinfer_dispatcher.py index fee0392a99be..e2d885be10d8 100644 --- a/test/registered/unit/layers/moe/test_flashinfer_dispatcher.py +++ b/test/registered/unit/layers/moe/test_flashinfer_dispatcher.py @@ -21,12 +21,26 @@ register_cpu_ci(est_time=11, suite="base-a-test-cpu") -class _RecordingMoeAlltoAll: - def dispatch(self, token_selected_experts, payloads, *_args, **kwargs): - self.token_selected_experts = token_selected_experts - self.payloads = payloads - self.expert_id_payload_index = kwargs["expert_id_payload_index"] - return payloads +class _RecordingCommunication: + def dispatch( + self, + hidden_states, + topk_ids, + topk_weights=None, + *, + hidden_states_scale=None, + **_kwargs, + ): + self.hidden_states = hidden_states + self.hidden_states_scale = hidden_states_scale + self.topk_ids = topk_ids + self.topk_weights = topk_weights + return SimpleNamespace( + hidden_states=hidden_states, + hidden_states_scale=hidden_states_scale, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) def _make_dispatcher(*, dispatch_type, max_num_tokens, quant_config): @@ -39,7 +53,7 @@ def _make_dispatcher(*, dispatch_type, max_num_tokens, quant_config): dispatcher.invalid_token_expert_id = 8 dispatcher.payload_in_workspace = False dispatcher.quant_config = quant_config - dispatcher.moe_a2a = _RecordingMoeAlltoAll() + dispatcher.moe_a2a = _RecordingCommunication() dispatcher.decode_a2a_namespace = _WorkspaceNamespace.DECODE dispatcher.sp_prefill_moe_a2a = None return dispatcher @@ -78,16 +92,22 @@ def test_empty_mxfp8_dispatch_uses_same_payload_dtype_as_nonempty_rank(self): topk_output=topk_output, ) + moe_a2a = dispatcher.moe_a2a self.assertEqual( - [payload.dtype for payload in dispatcher.moe_a2a.payloads], + [ + moe_a2a.hidden_states.dtype, + moe_a2a.hidden_states_scale.dtype, + moe_a2a.topk_ids.dtype, + moe_a2a.topk_weights.dtype, + ], [torch.float8_e4m3fn, torch.uint8, torch.int32, torch.float32], ) self.assertEqual(output.hidden_states.dtype, torch.float8_e4m3fn) self.assertEqual(output.hidden_states_scale.dtype, torch.uint8) def test_int64_expert_ids_reach_moe_a2a_as_int32(self): - """Expert IDs reach MoeAlltoAll as int32 even when routing emits int64, - as static EP dispatch does; MoeAlltoAll rejects any other dtype.""" + """Expert IDs reach the A2A as int32 even when routing emits int64, + as static EP dispatch does; FlashInfer's kernel rejects any other dtype.""" dispatcher = _make_dispatcher( dispatch_type=FlashinferA2ADispatchType.BF16, max_num_tokens=4, @@ -107,16 +127,11 @@ def test_int64_expert_ids_reach_moe_a2a_as_int32(self): ) moe_a2a = dispatcher.moe_a2a - self.assertEqual(moe_a2a.token_selected_experts.dtype, torch.int32) - self.assertEqual( - moe_a2a.payloads[moe_a2a.expert_id_payload_index].dtype, torch.int32 - ) - self.assertTrue( - torch.equal(moe_a2a.token_selected_experts, topk_ids.to(torch.int32)) - ) + self.assertEqual(moe_a2a.topk_ids.dtype, torch.int32) + self.assertTrue(torch.equal(moe_a2a.topk_ids, topk_ids.to(torch.int32))) def test_bypassed_topk_is_routed_before_moe_a2a(self): - """Runners that fuse routing hand over bypassed top-k; MoeAlltoAll still + """Runners that fuse routing hand over bypassed top-k; the A2A still needs explicit expert IDs and weights, so routing must run first.""" dispatcher = _make_dispatcher( dispatch_type=FlashinferA2ADispatchType.BF16, @@ -150,13 +165,12 @@ def test_bypassed_topk_is_routed_before_moe_a2a(self): ) moe_a2a = dispatcher.moe_a2a - expert_ids = moe_a2a.token_selected_experts + expert_ids = moe_a2a.topk_ids self.assertEqual(expert_ids.dtype, torch.int32) self.assertEqual( expert_ids.sort(dim=-1).values.tolist(), [[0, 1], [2, 3], [1, 3]] ) - topk_weights = moe_a2a.payloads[moe_a2a.expert_id_payload_index + 1] - self.assertEqual(topk_weights.shape, (3, 2)) + self.assertEqual(moe_a2a.topk_weights.shape, (3, 2)) class TestFlashinferRunnerDispatchFormats(CustomTestCase): diff --git a/test/registered/unit/layers/muse_spark/test_muse_spark_runner_dispatch.py b/test/registered/unit/layers/muse_spark/test_muse_spark_runner_dispatch.py index c76859f05386..3cc59ecd3027 100644 --- a/test/registered/unit/layers/muse_spark/test_muse_spark_runner_dispatch.py +++ b/test/registered/unit/layers/muse_spark/test_muse_spark_runner_dispatch.py @@ -22,13 +22,17 @@ register_cpu_ci(est_time=2, suite="base-a-test-cpu") -class _CpuMoeAlltoAll: - def __init__(self, **_kwargs): - self.invalid_token_expert_id = None +class _CpuCommunication: + def __init__(self, _bootstrap, params, **_kwargs): + self.params = params - def dispatch(self, _topk_ids, payloads, _max_tokens, **kwargs): - self.invalid_token_expert_id = kwargs["invalid_token_expert_id"] - return payloads + def dispatch(self, hidden_states, topk_ids, topk_weights, **_kwargs): + return SimpleNamespace( + hidden_states=hidden_states, + hidden_states_scale=None, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) class TestMuseSparkRunnerDispatch(CustomTestCase): @@ -83,11 +87,11 @@ def test_custom_expert_config_selects_trtllm_receive_sentinel(self): with ExitStack() as stack: for name, value in { "use_flashinfer": True, - "moe_a2a_get_workspace_size_per_rank": lambda **_: 128, - "Mapping": lambda **_: None, - "MnnvlConfig": lambda **_: None, + "BootstrapConfig": lambda **_: None, + "MoEEpCommParams": SimpleNamespace, + "NVLinkOneSidedConfig": lambda **_: None, "TorchDistributedCommBackend": lambda _: None, - "MoeAlltoAll": _CpuMoeAlltoAll, + "create_communication": _CpuCommunication, "get_flashinfer_a2a_dispatch_type": lambda: ( FlashinferA2ADispatchType.BF16 ), @@ -118,7 +122,7 @@ def test_custom_expert_config_selects_trtllm_receive_sentinel(self): torch.zeros((2, 128), dtype=torch.bfloat16), topk ) self.assertEqual( - dispatcher.sp_prefill_moe_a2a.invalid_token_expert_id, -1 + dispatcher.sp_prefill_moe_a2a.params.invalid_expert_id, -1 ) torch.testing.assert_close( output.topk_output.topk_ids, topk.topk_ids