Add dedicated FlashInferCuteDslMoE layer for standard-path FP4 MoE (#21339)
This commit is contained in:
Lee Nau
@@ -0,0 +1,353 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.moe.moe_runner.base import (
|
||||
MoeQuantInfo,
|
||||
MoeRunnerConfig,
|
||||
register_fused_func,
|
||||
)
|
||||
from sglang.srt.utils.common import log_info_on_rank0, print_warning_once
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.moe.token_dispatcher import (
|
||||
StandardCombineInput,
|
||||
StandardDispatchOutput,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_FP4_SF_VEC_SIZE = 16
|
||||
_cutedsl_logged_scalarize: set = set()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Weight / scale preparation utilities (called from modelopt_quant.py during
|
||||
# process_weights_after_loading and lazy wrapper init)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def interleave_w13_halves(
|
||||
tensor: torch.Tensor, group_size: int = 64, dim: int = 1
|
||||
) -> torch.Tensor:
|
||||
"""Interleave the two logical W13 halves for CuteDSL's SwiGLU GEMM1 layout.
|
||||
|
||||
The caller is responsible for loading W13 in the expected two-half order.
|
||||
This helper only rewrites the first and second halves into alternating
|
||||
`group_size` chunks along `dim`.
|
||||
"""
|
||||
if tensor.shape[dim] % 2 != 0:
|
||||
raise ValueError(
|
||||
"Expected even size on interleave dimension for W13 half split."
|
||||
)
|
||||
split = tensor.shape[dim] // 2
|
||||
if split % group_size != 0:
|
||||
raise ValueError(
|
||||
f"Expected split dim divisible by group_size={group_size}, got {split}."
|
||||
)
|
||||
first_half = tensor.narrow(dim, 0, split)
|
||||
second_half = tensor.narrow(dim, split, split)
|
||||
first_half_groups = first_half.split(group_size, dim=dim)
|
||||
second_half_groups = second_half.split(group_size, dim=dim)
|
||||
interleaved = [
|
||||
item for pair in zip(first_half_groups, second_half_groups) for item in pair
|
||||
]
|
||||
return torch.cat(interleaved, dim=dim)
|
||||
|
||||
|
||||
def cutedsl_quant_scale_to_scalar(
|
||||
quant_scale: torch.Tensor,
|
||||
*,
|
||||
name: str,
|
||||
) -> torch.Tensor:
|
||||
"""Reduce per-expert quant-domain scale vector to a single scalar.
|
||||
|
||||
The quant domain is the reciprocal of the raw checkpoint scale:
|
||||
quant_scale = 1 / raw_scale
|
||||
|
||||
Returns min(quant_scale) = 1/max(raw_scale), which is the TRTLLM CuteDSL
|
||||
convention for global scalar activation scales (see TRTLLM quantization.py
|
||||
lines 2137-2141: fc2_input_scale = tmp_fc2_input_scale.max().reciprocal()).
|
||||
|
||||
If quant_scale is already scalar (numel==1), returns it unchanged.
|
||||
"""
|
||||
quant_scale = quant_scale.to(torch.float32)
|
||||
if quant_scale.numel() == 0:
|
||||
print_warning_once(
|
||||
f"CuteDSL got empty {name}; using 1.0 fallback.",
|
||||
)
|
||||
return torch.ones(1, device=quant_scale.device, dtype=torch.float32)
|
||||
if quant_scale.numel() == 1:
|
||||
return quant_scale.reshape(1)
|
||||
if name not in _cutedsl_logged_scalarize:
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"CuteDSL: reducing per-expert {name} to scalar via "
|
||||
"min(quant_scale) = 1/max(raw_scale), matching TRTLLM convention.",
|
||||
)
|
||||
_cutedsl_logged_scalarize.add(name)
|
||||
return quant_scale.min().reshape(1)
|
||||
|
||||
|
||||
def resolve_cutedsl_standard_scales(
|
||||
layer: torch.nn.Module,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Resolve standard-path CuteDSL scales (baseline: scalar fc2/w13 input scales).
|
||||
|
||||
Returns (w1_alpha, fc2_input_scale, w2_alpha, used_input_scale).
|
||||
used_input_scale is the scalarized w13 input scale for FP4 quantize and GEMM1.
|
||||
"""
|
||||
|
||||
def _to_fp32_tensor(x: torch.Tensor | float, ref: torch.Tensor) -> torch.Tensor:
|
||||
if not isinstance(x, torch.Tensor):
|
||||
x = torch.tensor(x, device=ref.device)
|
||||
return x.to(device=ref.device, dtype=torch.float32)
|
||||
|
||||
def _align_scale_to_alpha(
|
||||
scale: torch.Tensor, alpha: torch.Tensor, scale_name: str
|
||||
) -> torch.Tensor:
|
||||
scale = scale.to(device=alpha.device, dtype=torch.float32)
|
||||
alpha = alpha.to(torch.float32)
|
||||
if scale.ndim == 0:
|
||||
return scale
|
||||
# Gated weight scales may be (num_experts, 2) with separate gate/up
|
||||
# columns. Collapse to 1D by taking the first column (gate == up for
|
||||
# well-formed checkpoints; mismatch is warned in process_weights_after_loading).
|
||||
if scale.ndim == 2 and scale.shape[1] <= 2:
|
||||
scale = scale[:, 0]
|
||||
if scale.numel() == alpha.numel():
|
||||
return scale
|
||||
if scale.numel() == 1:
|
||||
return scale.reshape(())
|
||||
|
||||
# Some EP setups may carry global-per-expert scale vectors while alphas are
|
||||
# local-per-expert vectors. Slice to this rank's local expert range.
|
||||
num_local_experts = getattr(layer, "num_local_experts", None)
|
||||
num_experts = getattr(layer, "num_experts", None)
|
||||
moe_ep_rank = getattr(layer, "moe_ep_rank", 0)
|
||||
if (
|
||||
num_local_experts is not None
|
||||
and num_experts is not None
|
||||
and scale.numel() == num_experts
|
||||
and alpha.numel() == num_local_experts
|
||||
):
|
||||
start = moe_ep_rank * num_local_experts
|
||||
end = start + num_local_experts
|
||||
return scale[start:end]
|
||||
|
||||
raise ValueError(
|
||||
f"Unable to align {scale_name} shape={tuple(scale.shape)} "
|
||||
f"to alpha shape={tuple(alpha.shape)} for CuteDSL standard scale resolution."
|
||||
)
|
||||
|
||||
def _resolve_w1_alpha_from_scalar_input_scale(
|
||||
used_input_scale: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Resolve GEMM1 alpha consistent with scalarized activation quant scale.
|
||||
|
||||
CuteDSL pre-quantizes x with a single scalar (used_input_scale), but
|
||||
g1_alphas was derived with per-expert activation scales:
|
||||
g1_alphas[e] = (1/w13_isq[e]) * w13_ws2[e]
|
||||
Correct alpha for scalar quantization:
|
||||
w1_alpha[e] = w13_ws2[e] / used_input_scale
|
||||
= g1_alphas[e] * w13_isq[e] / used_input_scale
|
||||
When w13_isq is already scalar, this is a no-op (ratio = 1).
|
||||
"""
|
||||
eps = 1e-12
|
||||
scalar = torch.clamp(used_input_scale.to(torch.float32).reshape(()), min=eps)
|
||||
|
||||
if hasattr(layer, "w13_weight_scale_2"):
|
||||
w13_weight_scale_2 = _align_scale_to_alpha(
|
||||
layer.w13_weight_scale_2, layer.g1_alphas, "w13_weight_scale_2"
|
||||
)
|
||||
return w13_weight_scale_2.to(torch.float32) / scalar
|
||||
|
||||
w13_isq = _align_scale_to_alpha(
|
||||
layer.w13_input_scale_quant, layer.g1_alphas, "w13_input_scale_quant"
|
||||
)
|
||||
w13_isq = torch.clamp(_to_fp32_tensor(w13_isq, layer.g1_alphas), min=eps)
|
||||
return (layer.g1_alphas.to(torch.float32) * w13_isq / scalar).to(torch.float32)
|
||||
|
||||
def _resolve_w2_alpha_from_scalar_fc2_input_scale(
|
||||
fc2_input_scale: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Resolve GEMM2 alpha consistent with scalarized FC2 input scale.
|
||||
|
||||
CuteDSL standard path uses a scalar global scale for GEMM1 FP4 output
|
||||
quantization (`fc2_input_scale`). GEMM2 alpha must use the same scalar
|
||||
convention: alpha2 = w2_weight_scale_2 / fc2_input_scale.
|
||||
"""
|
||||
eps = 1e-12
|
||||
fc2_input_scale = fc2_input_scale.to(torch.float32)
|
||||
fc2_scalar = torch.clamp(fc2_input_scale.reshape(-1)[:1], min=eps).reshape(())
|
||||
|
||||
if hasattr(layer, "w2_weight_scale_2"):
|
||||
w2_weight_scale_2 = _align_scale_to_alpha(
|
||||
layer.w2_weight_scale_2, layer.g2_alphas, "w2_weight_scale_2"
|
||||
)
|
||||
w2_weight_scale_2 = w2_weight_scale_2.to(torch.float32)
|
||||
return w2_weight_scale_2 / fc2_scalar
|
||||
|
||||
w2_q_for_w2 = _align_scale_to_alpha(
|
||||
layer.w2_input_scale_quant, layer.g2_alphas, "w2_input_scale_quant"
|
||||
)
|
||||
w2_q_for_w2 = torch.clamp(
|
||||
_to_fp32_tensor(w2_q_for_w2, layer.g2_alphas), min=eps
|
||||
)
|
||||
w2_weight_scale_2 = layer.g2_alphas.to(torch.float32) * w2_q_for_w2
|
||||
return w2_weight_scale_2 / fc2_scalar
|
||||
|
||||
fc2_input_scale = cutedsl_quant_scale_to_scalar(
|
||||
layer.w2_input_scale_quant,
|
||||
name="w2_input_scale_quant",
|
||||
)
|
||||
w2_alpha = _resolve_w2_alpha_from_scalar_fc2_input_scale(fc2_input_scale)
|
||||
used_input_scale = cutedsl_quant_scale_to_scalar(
|
||||
layer.w13_input_scale_quant,
|
||||
name="w13_input_scale_quant",
|
||||
)
|
||||
w1_alpha = _resolve_w1_alpha_from_scalar_input_scale(used_input_scale)
|
||||
return w1_alpha, fc2_input_scale, w2_alpha, used_input_scale
|
||||
|
||||
|
||||
def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None:
|
||||
"""Lazily create CuteDslMoEWrapper and resolve scales on first forward.
|
||||
|
||||
The wrapper is created lazily (not in __init__ / create_weights) because
|
||||
it depends on final weight shapes and EP configuration. The wrapper's
|
||||
CUDA-graph buffers are allocated inside CuteDslMoEWrapper.__init__, which
|
||||
typically runs during the autotune dummy forward under inference_mode().
|
||||
We wrap the creation in inference_mode(False) so that those pre-allocated
|
||||
buffers are normal tensors -- inference tensors cannot be inplace-updated
|
||||
during later CUDA graph capture, which runs outside inference_mode.
|
||||
"""
|
||||
if getattr(layer, "_cutedsl_wrapper", None) is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
from flashinfer import CuteDslMoEWrapper
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"flashinfer_cutedsl backend requires FlashInfer with CuteDSL support. "
|
||||
"Install with: pip install flashinfer"
|
||||
) from e
|
||||
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
|
||||
assert layer.intermediate_size_per_partition > 0, (
|
||||
f"CuteDSL MoE: intermediate_size_per_partition must be > 0, "
|
||||
f"got {layer.intermediate_size_per_partition}. Check EP/TP configuration."
|
||||
)
|
||||
|
||||
server_args = get_global_server_args()
|
||||
use_cuda_graph = server_args is not None and not server_args.disable_cuda_graph
|
||||
max_num_tokens = max(
|
||||
getattr(server_args, "cuda_graph_max_bs", None) or 512,
|
||||
getattr(server_args, "chunked_prefill_size", None) or 8192,
|
||||
)
|
||||
top_k = layer.top_k if layer.top_k is not None else layer.moe_runner_config.top_k
|
||||
# inference_mode(False) ensures the wrapper's pre-allocated CUDA-graph
|
||||
# buffers are normal tensors. This call typically happens inside
|
||||
# _dummy_run which runs under inference_mode(); inference tensors cannot
|
||||
# be inplace-updated during later CUDA graph capture (which runs outside
|
||||
# inference_mode), so we must opt out here.
|
||||
with torch.inference_mode(False):
|
||||
layer._cutedsl_wrapper = CuteDslMoEWrapper(
|
||||
num_experts=layer.num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=layer.hidden_size,
|
||||
intermediate_size=layer.intermediate_size_per_partition,
|
||||
use_cuda_graph=use_cuda_graph,
|
||||
max_num_tokens=max_num_tokens,
|
||||
num_local_experts=layer.num_local_experts,
|
||||
local_expert_offset=layer.moe_ep_rank * layer.num_local_experts,
|
||||
output_dtype=layer.moe_runner_config.params_dtype,
|
||||
device=str(layer.w13_weight.device),
|
||||
)
|
||||
|
||||
w1_alpha, fc2_input_scale, w2_alpha, used_input_scale = (
|
||||
resolve_cutedsl_standard_scales(layer)
|
||||
)
|
||||
layer._cutedsl_scales = (w1_alpha, fc2_input_scale, w2_alpha)
|
||||
layer._cutedsl_input_scale = used_input_scale
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataclass + fused function for moe_runner dispatch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class CuteDslFp4MoeQuantInfo(MoeQuantInfo):
|
||||
"""Quantization payload consumed by FlashInfer CuteDSL FP4 MoE kernels."""
|
||||
|
||||
# Lazily-created CuteDslMoEWrapper (stashed on layer)
|
||||
wrapper: Any
|
||||
|
||||
# Weights (uint8 FP4 packed)
|
||||
w13_weight: torch.Tensor
|
||||
w2_weight: torch.Tensor
|
||||
|
||||
# Block-scale factors
|
||||
w13_weight_sf: torch.Tensor
|
||||
w2_weight_sf: torch.Tensor
|
||||
|
||||
# Per-expert GEMM scales
|
||||
w1_alpha: torch.Tensor
|
||||
w2_alpha: torch.Tensor
|
||||
|
||||
# Intermediate quantization scale (fc2 input)
|
||||
fc2_input_scale: torch.Tensor
|
||||
|
||||
# Activation quantization scale (scalarized)
|
||||
input_scale: torch.Tensor
|
||||
|
||||
|
||||
@register_fused_func("none", "flashinfer_cutedsl")
|
||||
def fused_experts_none_to_flashinfer_cutedsl_fp4(
|
||||
dispatch_output: StandardDispatchOutput,
|
||||
quant_info: CuteDslFp4MoeQuantInfo,
|
||||
runner_config: MoeRunnerConfig,
|
||||
) -> StandardCombineInput:
|
||||
from flashinfer import fp4_quantize
|
||||
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
|
||||
from sglang.srt.layers.moe.topk import TopKOutputChecker
|
||||
|
||||
assert runner_config.activation == "silu", "Only silu is supported for CuteDSL MoE."
|
||||
|
||||
hidden_states = dispatch_output.hidden_states
|
||||
topk_output = dispatch_output.topk_output
|
||||
assert TopKOutputChecker.format_is_standard(topk_output)
|
||||
|
||||
topk_ids = topk_output.topk_ids
|
||||
topk_weights = topk_output.topk_weights
|
||||
if topk_ids.dtype != torch.int32:
|
||||
topk_ids = topk_ids.to(torch.int32)
|
||||
|
||||
x_fp4, x_sf = fp4_quantize(
|
||||
hidden_states,
|
||||
quant_info.input_scale,
|
||||
sf_vec_size=_FP4_SF_VEC_SIZE,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
|
||||
output = quant_info.wrapper.run(
|
||||
x=x_fp4,
|
||||
x_sf=x_sf,
|
||||
token_selected_experts=topk_ids,
|
||||
token_final_scales=topk_weights,
|
||||
w1_weight=quant_info.w13_weight,
|
||||
w1_weight_sf=quant_info.w13_weight_sf,
|
||||
w1_alpha=quant_info.w1_alpha,
|
||||
fc2_input_scale=quant_info.fc2_input_scale,
|
||||
w2_weight=quant_info.w2_weight,
|
||||
w2_weight_sf=quant_info.w2_weight_sf,
|
||||
w2_alpha=quant_info.w2_alpha,
|
||||
)
|
||||
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
@@ -54,6 +54,8 @@ class MoeRunner:
|
||||
or runner_backend.is_flashinfer_trtllm_routed()
|
||||
):
|
||||
self.runner_core = None # FlashInfer TRT-LLM only supports fused path
|
||||
elif runner_backend.is_flashinfer_cutedsl():
|
||||
self.runner_core = None # FlashInfer CuteDSL only supports fused path
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported runner backend: {runner_backend}")
|
||||
|
||||
|
||||
@@ -88,8 +88,14 @@ class StandardDispatcher(BaseDispatcher):
|
||||
def __init__(self, moe_runner_config: MoeRunnerConfig):
|
||||
super().__init__()
|
||||
self.moe_ep_size = get_moe_expert_parallel_world_size()
|
||||
self.enable_flashinfer_cutlass_moe = (
|
||||
get_moe_runner_backend().is_flashinfer_cutlass()
|
||||
backend = get_moe_runner_backend()
|
||||
self.enable_flashinfer_cutlass_moe = backend.is_flashinfer_cutlass()
|
||||
# FlashInfer CUTLASS and CuteDSL handle EP internally with global expert IDs.
|
||||
# Skip local expert mapping so topk_ids stay in global space.
|
||||
self.skip_local_expert_mapping = (
|
||||
backend.is_flashinfer_cutlass()
|
||||
or backend.is_flashinfer_cutedsl()
|
||||
or backend.is_flashinfer_trtllm_routed()
|
||||
)
|
||||
self.enable_flashinfer_trtllm_routed_moe = (
|
||||
get_moe_runner_backend().is_flashinfer_trtllm_routed()
|
||||
@@ -149,8 +155,7 @@ class StandardDispatcher(BaseDispatcher):
|
||||
|
||||
if (
|
||||
self.moe_ep_size > 1
|
||||
and not self.enable_flashinfer_cutlass_moe
|
||||
and not self.enable_flashinfer_trtllm_routed_moe
|
||||
and not self.skip_local_expert_mapping
|
||||
and TopKOutputChecker.format_is_standard(topk_output)
|
||||
):
|
||||
if self.local_expert_mapping is None:
|
||||
|
||||
@@ -51,7 +51,6 @@ from sglang.srt.layers.quantization.utils import (
|
||||
from sglang.srt.layers.radix_attention import RadixAttention
|
||||
from sglang.srt.layers.utils import copy_or_rebind_param
|
||||
from sglang.srt.utils.common import (
|
||||
get_bool_env_var,
|
||||
is_cuda,
|
||||
is_sm120_supported,
|
||||
next_power_of_2,
|
||||
@@ -167,10 +166,6 @@ if is_cuda() and (not is_sm120_supported()) and (fp4_quantize is not None):
|
||||
return
|
||||
|
||||
|
||||
CUTEDSL_MOE_SCALAR_INPUT_SCALE = get_bool_env_var(
|
||||
"SGLANG_CUTEDSL_MOE_SCALAR_INPUT_SCALE", "true"
|
||||
)
|
||||
|
||||
# FP4 GEMM alignment constant - CUTLASS/FlashInfer kernels require dimensions divisible by 32
|
||||
FP4_GEMM_ALIGNMENT = 32
|
||||
|
||||
@@ -993,9 +988,10 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
|
||||
) -> CombineInput:
|
||||
x = dispatch_output.hidden_states
|
||||
topk_output = dispatch_output.topk_output
|
||||
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
|
||||
from sglang.srt.layers.moe.topk import TopKOutputChecker
|
||||
|
||||
# Fast path: TRT-LLM FP8 per-tensor MoE using BYPASSED TopK routing
|
||||
from sglang.srt.layers.moe.topk import TopKOutputChecker
|
||||
|
||||
if (
|
||||
get_moe_runner_backend().is_flashinfer_trtllm()
|
||||
@@ -1089,8 +1085,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
|
||||
activation_type=activation,
|
||||
)[0]
|
||||
|
||||
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
|
||||
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
quant_info = TritonMoeQuantInfo(
|
||||
@@ -1547,9 +1541,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
|
||||
|
||||
@property
|
||||
def enable_flashinfer_cutedsl_moe(self) -> bool:
|
||||
"""Access the global enable_flashinfer_cutedsl_moe setting."""
|
||||
from sglang.srt.layers.moe import get_moe_runner_backend
|
||||
|
||||
"""Access the global enable_flashinfer_cutedsl_moe setting."""
|
||||
return get_moe_runner_backend().is_flashinfer_cutedsl()
|
||||
|
||||
def create_weights(
|
||||
@@ -1714,19 +1708,12 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
|
||||
w13_input_scale = layer.w13_input_scale.max().to(torch.float32)
|
||||
w2_input_scale = layer.w2_input_scale.max().to(torch.float32)
|
||||
elif self.enable_flashinfer_cutedsl_moe:
|
||||
# All-expert-one-input-scale is mathematically different from default per-expert-input-scale
|
||||
# Thus we allow users to switch the flag to do thorough testing
|
||||
if CUTEDSL_MOE_SCALAR_INPUT_SCALE:
|
||||
w13_input_scale = (
|
||||
layer.w13_input_scale.max()
|
||||
.to(torch.float32)
|
||||
.repeat(layer.w13_input_scale.shape[0])
|
||||
)
|
||||
else:
|
||||
w13_input_scale = layer.w13_input_scale.max(dim=1).values.to(
|
||||
torch.float32
|
||||
)
|
||||
|
||||
# CuteDSL standard path uses a single scalar input scale (all experts).
|
||||
w13_input_scale = (
|
||||
layer.w13_input_scale.max()
|
||||
.to(torch.float32)
|
||||
.repeat(layer.w13_input_scale.shape[0])
|
||||
)
|
||||
w2_input_scale = layer.w2_input_scale
|
||||
|
||||
def _slice_scale(w):
|
||||
@@ -1825,6 +1812,26 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
|
||||
else:
|
||||
# CUTLASS processing - handle w13 and w2 separately
|
||||
|
||||
if self.enable_flashinfer_cutedsl_moe and layer.moe_runner_config.is_gated:
|
||||
# For the CuteDSL FP4 path, interleave the two logical W13 halves
|
||||
# in 64-row chunks before swizzling the block-scales.
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import (
|
||||
interleave_w13_halves,
|
||||
)
|
||||
|
||||
layer.w13_weight = Parameter(
|
||||
interleave_w13_halves(
|
||||
layer.w13_weight.view(torch.uint8), group_size=64, dim=1
|
||||
).contiguous(),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w13_weight_scale = Parameter(
|
||||
interleave_w13_halves(
|
||||
layer.w13_weight_scale, group_size=64, dim=1
|
||||
).contiguous(),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
# Process w13 weights
|
||||
w13_blockscale_swizzled = swizzle_blockscale(layer.w13_weight_scale)
|
||||
copy_or_rebind_param(
|
||||
@@ -1869,6 +1876,45 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
|
||||
layer, "w2_blockscale_swizzled", w2_blockscale_swizzled
|
||||
)
|
||||
|
||||
if self.enable_flashinfer_cutedsl_moe:
|
||||
# CuteDSL expects MMA layout for weight scales. Convert from swizzled bytes.
|
||||
from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout
|
||||
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import (
|
||||
_FP4_SF_VEC_SIZE,
|
||||
)
|
||||
|
||||
sf_vec_size = _FP4_SF_VEC_SIZE
|
||||
num_local_experts = layer.w13_weight.shape[0]
|
||||
w13_m = layer.w13_weight.shape[1]
|
||||
w13_k = layer.w13_weight.shape[2] * 2
|
||||
w2_m = layer.w2_weight.shape[1]
|
||||
w2_k = layer.w2_weight.shape[2] * 2
|
||||
layer.w13_blockscale_mma = Parameter(
|
||||
convert_sf_to_mma_layout(
|
||||
layer.w13_blockscale_swizzled.contiguous()
|
||||
.view(torch.uint8)
|
||||
.reshape(-1),
|
||||
m=w13_m,
|
||||
k=w13_k,
|
||||
num_groups=num_local_experts,
|
||||
sf_vec_size=sf_vec_size,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_blockscale_mma = Parameter(
|
||||
convert_sf_to_mma_layout(
|
||||
layer.w2_blockscale_swizzled.contiguous()
|
||||
.view(torch.uint8)
|
||||
.reshape(-1),
|
||||
m=w2_m,
|
||||
k=w2_k,
|
||||
num_groups=num_local_experts,
|
||||
sf_vec_size=sf_vec_size,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
# Both flashinfer cutlass and regular cutlass use same processing for w2
|
||||
|
||||
# Set up CUTLASS MoE parameters (reuse to keep CUDA graph stable)
|
||||
@@ -1894,27 +1940,34 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
|
||||
|
||||
@property
|
||||
def load_up_proj_weight_first(self) -> bool:
|
||||
# FlashInfer CUTLASS kernel assumes [Up, Gate] Proj as W13
|
||||
return self.enable_flashinfer_cutlass_moe and self.moe_runner_config.is_gated
|
||||
# Load W13 as [Up, Gate] for FlashInfer CUTLASS/CuteDSL kernels.
|
||||
return self.moe_runner_config.is_gated and (
|
||||
self.enable_flashinfer_cutlass_moe or self.enable_flashinfer_cutedsl_moe
|
||||
)
|
||||
|
||||
def create_moe_runner(
|
||||
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
|
||||
):
|
||||
self.moe_runner_config = moe_runner_config
|
||||
if get_moe_runner_backend().is_flashinfer_trtllm():
|
||||
self.runner = MoeRunner(
|
||||
MoeRunnerBackend.FLASHINFER_TRTLLM, moe_runner_config
|
||||
)
|
||||
elif get_moe_runner_backend().is_flashinfer_trtllm_routed():
|
||||
self.runner = MoeRunner(
|
||||
MoeRunnerBackend.FLASHINFER_TRTLLM_ROUTED, moe_runner_config
|
||||
)
|
||||
moe_runner_backend = get_moe_runner_backend()
|
||||
|
||||
if moe_runner_backend.is_auto():
|
||||
# TRTLLM is currently the most performant and tested FP4 MoE
|
||||
# backend, so use it as the default.
|
||||
moe_runner_backend = MoeRunnerBackend.FLASHINFER_TRTLLM
|
||||
|
||||
if moe_runner_backend.is_flashinfer_cutedsl():
|
||||
import sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl # noqa: F401 – triggers @register_fused_func
|
||||
|
||||
if not moe_runner_backend.is_flashinfer_cutlass():
|
||||
self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: FusedMoE,
|
||||
dispatch_output: StandardDispatchOutput,
|
||||
) -> CombineInput:
|
||||
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
|
||||
|
||||
x = dispatch_output.hidden_states
|
||||
x_sf = dispatch_output.hidden_states_scale
|
||||
@@ -1958,6 +2011,33 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
|
||||
|
||||
return self.runner.run(dispatch_output, quant_info)
|
||||
|
||||
if self.enable_flashinfer_cutedsl_moe:
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import (
|
||||
CuteDslFp4MoeQuantInfo,
|
||||
ensure_cutedsl_wrapper,
|
||||
)
|
||||
|
||||
ensure_cutedsl_wrapper(layer)
|
||||
w1_alpha, fc2_input_scale, w2_alpha = layer._cutedsl_scales
|
||||
w1_weight_sf = getattr(
|
||||
layer, "w13_blockscale_mma", layer.w13_blockscale_swizzled
|
||||
)
|
||||
w2_weight_sf = getattr(
|
||||
layer, "w2_blockscale_mma", layer.w2_blockscale_swizzled
|
||||
)
|
||||
quant_info = CuteDslFp4MoeQuantInfo(
|
||||
wrapper=layer._cutedsl_wrapper,
|
||||
w13_weight=layer.w13_weight,
|
||||
w2_weight=layer.w2_weight,
|
||||
w13_weight_sf=w1_weight_sf,
|
||||
w2_weight_sf=w2_weight_sf,
|
||||
w1_alpha=w1_alpha,
|
||||
w2_alpha=w2_alpha,
|
||||
fc2_input_scale=fc2_input_scale,
|
||||
input_scale=layer._cutedsl_input_scale,
|
||||
)
|
||||
return self.runner.run(dispatch_output, quant_info)
|
||||
|
||||
if self.enable_flashinfer_cutlass_moe:
|
||||
from sglang.srt.layers.moe.token_dispatcher import DispatchOutputChecker
|
||||
|
||||
@@ -1997,7 +2077,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
|
||||
fc2_expert_weights=layer.w2_weight.view(torch.long),
|
||||
output_dtype=output_dtype,
|
||||
input_sf=x_sf,
|
||||
# swizzled_input_sf=not get_moe_a2a_backend().is_flashinfer(),
|
||||
# swizzled_input_sf intentionally omitted; not used for this path.
|
||||
quant_scales=[
|
||||
layer.w13_input_scale_quant,
|
||||
layer.w13_blockscale_swizzled.view(torch.int32),
|
||||
@@ -2015,8 +2095,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
|
||||
enable_alltoall=get_moe_a2a_backend().is_flashinfer(),
|
||||
)[0]
|
||||
|
||||
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
|
||||
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4
|
||||
@@ -2038,8 +2116,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
|
||||
apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input,
|
||||
).to(x.dtype)
|
||||
# Scale by routed_scaling_factor is fused into select_experts.
|
||||
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
|
||||
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
def apply_without_routing_weights(
|
||||
|
||||
@@ -2148,6 +2148,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
# TODO: Enable for flashinfer_trtllm_routed once https://github.com/flashinfer-ai/flashinfer/issues/2749 is fixed.
|
||||
# "flashinfer_trtllm_routed",
|
||||
"flashinfer_mxfp4",
|
||||
"flashinfer_cutedsl",
|
||||
# TODO: flashinfer_cutlass will cause some flashinfer compilation errors. To be fixed.
|
||||
# "flashinfer_cutlass",
|
||||
]:
|
||||
|
||||
@@ -2733,6 +2733,26 @@ class ServerArgs:
|
||||
self.tp_size,
|
||||
], "The expert parallel size must be 1 or the same as the tensor parallel size"
|
||||
|
||||
if self.moe_runner_backend == "flashinfer_cutedsl":
|
||||
assert self.quantization in [
|
||||
"modelopt_fp4"
|
||||
], f"Invalid quantization '{self.quantization}'. \nFlashInfer CuteDSL MOE currently supports only: 'modelopt_fp4'."
|
||||
assert self.ep_size in [
|
||||
1,
|
||||
self.tp_size,
|
||||
], "The expert parallel size must be 1 or the same as the tensor parallel size"
|
||||
assert self.moe_a2a_backend in [
|
||||
"none",
|
||||
"deepep",
|
||||
], (
|
||||
f"flashinfer_cutedsl supports moe_a2a_backend='none' (standard path) "
|
||||
f"or 'deepep' (DeepEP low-latency path), got '{self.moe_a2a_backend}'."
|
||||
)
|
||||
self.disable_shared_experts_fusion = True
|
||||
logger.warning(
|
||||
"FlashInfer CuteDSL MoE is enabled. --disable-shared-experts-fusion is automatically set."
|
||||
)
|
||||
|
||||
if self.moe_runner_backend == "flashinfer_trtllm":
|
||||
assert self.quantization in [
|
||||
"modelopt_fp4",
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
"""Backend tests for CuteDSL MoE (FusedMoE + moe_runner, moe_a2a=none).
|
||||
|
||||
Exercises the CuteDSL moe_runner path with ModelOpt FP4 by launching a
|
||||
server with --moe-runner-backend flashinfer_cutedsl.
|
||||
|
||||
Two configurations are tested:
|
||||
- EP=1, TP=4: each GPU holds all experts with TP-sharded intermediate dim
|
||||
- EP=4, TP=4: each GPU holds 1/4 of experts at full intermediate width,
|
||||
partial results combined via all-reduce (no A2A dispatch)
|
||||
|
||||
Requires 4 GPUs. Run from repo root with:
|
||||
python -m pytest test/registered/backends/test_deepseek_v3_fp4_cutedsl_moe.py -v -s
|
||||
Or via the nightly suite:
|
||||
python test/run_suite.py --hw cuda --suite nightly-4-gpu-b200 --nightly
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
write_github_step_summary,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=900, suite="nightly-4-gpu-b200", nightly=True)
|
||||
|
||||
FULL_DEEPSEEK_V3_FP4_MODEL_PATH = "nvidia/DeepSeek-V3-0324-FP4"
|
||||
SERVER_LAUNCH_TIMEOUT = 1000
|
||||
GSM8K_ACCURACY_THRESHOLD = 0.935
|
||||
|
||||
|
||||
class TestDeepseekV3FP4CuteDSLMoE(CustomTestCase):
|
||||
"""CuteDSL standard moe_runner path: flashinfer_cutedsl + modelopt_fp4, EP=1."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = [
|
||||
"--tp",
|
||||
"4",
|
||||
"--ep",
|
||||
"1",
|
||||
"--mem-fraction-static",
|
||||
"0.75",
|
||||
"--attention-backend",
|
||||
"trtllm_mla",
|
||||
"--moe-runner-backend",
|
||||
"flashinfer_cutedsl",
|
||||
"--quantization",
|
||||
"modelopt_fp4",
|
||||
"--model-loader-extra-config",
|
||||
'{"enable_multithread_load": true}',
|
||||
]
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=SERVER_LAUNCH_TIMEOUT,
|
||||
other_args=other_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_a_gsm8k(
|
||||
self,
|
||||
): # Append an "a" to make this test run first (alphabetically) to warm up the server
|
||||
args = SimpleNamespace(
|
||||
num_shots=8,
|
||||
data_path=None,
|
||||
num_questions=1319,
|
||||
parallel=1319,
|
||||
max_new_tokens=512,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
if is_in_ci():
|
||||
write_github_step_summary(
|
||||
f"### test_gsm8k (deepseek-v3-fp4-cutedsl-moe)\n"
|
||||
f'{metrics["accuracy"]=:.3f}\n'
|
||||
)
|
||||
self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD)
|
||||
|
||||
|
||||
class TestDeepseekV3FP4CuteDSLMoEEP4(CustomTestCase):
|
||||
"""CuteDSL standard moe_runner path: flashinfer_cutedsl + modelopt_fp4, EP=TP=4."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = [
|
||||
"--tp",
|
||||
"4",
|
||||
"--ep",
|
||||
"4",
|
||||
"--mem-fraction-static",
|
||||
"0.75",
|
||||
"--attention-backend",
|
||||
"trtllm_mla",
|
||||
"--moe-runner-backend",
|
||||
"flashinfer_cutedsl",
|
||||
"--moe-a2a-backend",
|
||||
"none",
|
||||
"--quantization",
|
||||
"modelopt_fp4",
|
||||
"--model-loader-extra-config",
|
||||
'{"enable_multithread_load": true}',
|
||||
]
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=SERVER_LAUNCH_TIMEOUT,
|
||||
other_args=other_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_a_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=8,
|
||||
data_path=None,
|
||||
num_questions=1319,
|
||||
parallel=1319,
|
||||
max_new_tokens=512,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
if is_in_ci():
|
||||
write_github_step_summary(
|
||||
f"### test_gsm8k (deepseek-v3-fp4-cutedsl-moe-ep4)\n"
|
||||
f'{metrics["accuracy"]=:.3f}\n'
|
||||
)
|
||||
self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,18 +1,22 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import unittest
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
from flashinfer import fp4_quantize, scaled_fp4_grouped_quantize
|
||||
from torch.nn import functional as F
|
||||
|
||||
from sglang.jit_kernel.nvfp4 import scaled_fp4_quant
|
||||
from sglang.srt.layers.activation import SiluAndMul
|
||||
from sglang.srt.layers.moe.flashinfer_cutedsl_moe import flashinfer_cutedsl_moe_masked
|
||||
from sglang.srt.layers.moe.topk import TopKConfig, select_experts
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=20, suite="stage-c-test-4-gpu-b200")
|
||||
try:
|
||||
from flashinfer import CuteDslMoEWrapper
|
||||
from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout
|
||||
except ImportError:
|
||||
CuteDslMoEWrapper = None
|
||||
convert_sf_to_mma_layout = None
|
||||
|
||||
register_cuda_ci(est_time=300, suite="stage-c-test-4-gpu-b200")
|
||||
|
||||
SKIP_TEST = torch.cuda.get_device_capability() < (10, 0)
|
||||
SKIP_REASON = "Nvfp4 Requires compute capability of 10 or above."
|
||||
@@ -78,6 +82,312 @@ def break_fp4_bytes(a, dtype):
|
||||
return values.reshape(m, n * 2).to(dtype=dtype)
|
||||
|
||||
|
||||
def _interleave_w13_halves(
|
||||
x: torch.Tensor, group_size: int = 64, dim: int = -1
|
||||
) -> torch.Tensor:
|
||||
"""Interleave the two logical W13 halves for the CuteDSL wrapper layout."""
|
||||
sizes = x.size()
|
||||
dim = dim % x.dim()
|
||||
assert sizes[dim] % (group_size * 2) == 0
|
||||
prev_sizes = sizes[:dim]
|
||||
post_sizes = sizes[dim + 1 :]
|
||||
x = x.view(*prev_sizes, 2, sizes[dim] // (group_size * 2), group_size, *post_sizes)
|
||||
x = x.transpose(dim, dim + 1).contiguous().view(*sizes)
|
||||
return x
|
||||
|
||||
|
||||
def _create_cutedsl_wrapper_tensors(
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
device: str = "cuda",
|
||||
seed: int = 42,
|
||||
):
|
||||
"""Create quantized tensors for CuteDslMoEWrapper.run() (MMA layout, same as production).
|
||||
|
||||
Returns quantized inputs for the wrapper **and** the original bf16 weights
|
||||
needed to compute a numerical reference. Scale values (w1_alpha, w2_alpha,
|
||||
fc2_input_scale) are derived from weight magnitudes so that scale-contract
|
||||
bugs are caught.
|
||||
"""
|
||||
assert CuteDslMoEWrapper is not None and convert_sf_to_mma_layout is not None
|
||||
torch.manual_seed(seed)
|
||||
sf_vec_size = 16
|
||||
|
||||
x_bf16 = (
|
||||
torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) / 10
|
||||
)
|
||||
a1_gs = torch.tensor([1.0], device=device, dtype=torch.float32)
|
||||
x_quantized, x_sf = fp4_quantize(
|
||||
x_bf16,
|
||||
global_scale=a1_gs,
|
||||
sf_vec_size=sf_vec_size,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
x_sf = x_sf.unsqueeze(-1)
|
||||
|
||||
router_logits = torch.randn(num_tokens, num_experts, device=device)
|
||||
routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
|
||||
routing_weights, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
|
||||
routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
|
||||
routing_weights = routing_weights.float()
|
||||
selected_experts = selected_experts.to(torch.int32)
|
||||
|
||||
# --- GEMM1 weights ---
|
||||
w1_bf16 = (
|
||||
torch.randn(
|
||||
num_experts,
|
||||
2 * intermediate_size,
|
||||
hidden_size,
|
||||
dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
/ 10
|
||||
)
|
||||
w1_bf16_interleaved = _interleave_w13_halves(w1_bf16, group_size=64, dim=1)
|
||||
w1_amax = w1_bf16.abs().amax(dim=(1, 2)).to(torch.float32)
|
||||
w1_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax.mean()
|
||||
w1_gs = w1_gs.unsqueeze(0)
|
||||
w1_flat = w1_bf16_interleaved.view(num_experts * 2 * intermediate_size, hidden_size)
|
||||
w1_q_flat, w1_sf_flat = fp4_quantize(
|
||||
w1_flat,
|
||||
global_scale=w1_gs,
|
||||
sf_vec_size=sf_vec_size,
|
||||
is_sf_swizzled_layout=True,
|
||||
)
|
||||
w1_q = w1_q_flat.view(num_experts, 2 * intermediate_size, hidden_size // 2)
|
||||
w1_weight_sf = convert_sf_to_mma_layout(
|
||||
w1_sf_flat,
|
||||
m=2 * intermediate_size,
|
||||
k=hidden_size,
|
||||
num_groups=num_experts,
|
||||
sf_vec_size=sf_vec_size,
|
||||
)
|
||||
w1_alpha = 1.0 / (a1_gs * w1_gs).expand(num_experts)
|
||||
|
||||
# --- GEMM2 weights ---
|
||||
w2_bf16 = (
|
||||
torch.randn(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size,
|
||||
dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
/ 10
|
||||
)
|
||||
w2_amax = w2_bf16.abs().amax(dim=(1, 2)).to(torch.float32)
|
||||
w2_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax.mean()
|
||||
w2_gs = w2_gs.unsqueeze(0)
|
||||
w2_flat = w2_bf16.view(num_experts * hidden_size, intermediate_size)
|
||||
w2_q_flat, w2_sf_flat = fp4_quantize(
|
||||
w2_flat,
|
||||
global_scale=w2_gs,
|
||||
sf_vec_size=sf_vec_size,
|
||||
is_sf_swizzled_layout=True,
|
||||
)
|
||||
w2_q = w2_q_flat.view(num_experts, hidden_size, intermediate_size // 2)
|
||||
w2_weight_sf = convert_sf_to_mma_layout(
|
||||
w2_sf_flat,
|
||||
m=hidden_size,
|
||||
k=intermediate_size,
|
||||
num_groups=num_experts,
|
||||
sf_vec_size=sf_vec_size,
|
||||
)
|
||||
fc2_input_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax.mean()
|
||||
fc2_input_scale = fc2_input_scale.unsqueeze(0)
|
||||
w2_alpha = 1.0 / (fc2_input_scale * w2_gs).expand(num_experts)
|
||||
|
||||
return {
|
||||
"x": x_quantized,
|
||||
"x_sf": x_sf,
|
||||
"x_bf16": x_bf16,
|
||||
"token_selected_experts": selected_experts,
|
||||
"token_final_scales": routing_weights,
|
||||
"w1_weight": w1_q,
|
||||
"w1_weight_sf": w1_weight_sf,
|
||||
"w1_weight_bf16": w1_bf16,
|
||||
"w1_alpha": w1_alpha,
|
||||
"fc2_input_scale": fc2_input_scale,
|
||||
"w2_weight": w2_q,
|
||||
"w2_weight_sf": w2_weight_sf,
|
||||
"w2_weight_bf16": w2_bf16,
|
||||
"w2_alpha": w2_alpha,
|
||||
# Global scales needed by _quantize_local_expert_weights
|
||||
"a1_gs": a1_gs,
|
||||
"w1_gs": w1_gs,
|
||||
"w2_gs": w2_gs,
|
||||
}
|
||||
|
||||
|
||||
def _quantize_local_expert_weights(
|
||||
w1_bf16_local: torch.Tensor,
|
||||
w2_bf16_local: torch.Tensor,
|
||||
a1_gs: torch.Tensor,
|
||||
w1_gs: torch.Tensor,
|
||||
w2_gs: torch.Tensor,
|
||||
fc2_input_scale: torch.Tensor,
|
||||
):
|
||||
"""Independently quantize and MMA-convert a local expert weight shard.
|
||||
|
||||
Mirrors the per-rank weight preprocessing that happens during model loading
|
||||
in production (each rank holds [num_local_experts, ...] bf16 weights,
|
||||
quantizes them, and calls convert_sf_to_mma_layout with
|
||||
num_groups=num_local_experts).
|
||||
"""
|
||||
sf_vec_size = 16
|
||||
num_local_experts = w1_bf16_local.shape[0]
|
||||
intermediate_size_2x = w1_bf16_local.shape[1]
|
||||
hidden_size = w1_bf16_local.shape[2]
|
||||
intermediate_size = w2_bf16_local.shape[2]
|
||||
|
||||
# GEMM1: interleave -> quantize -> MMA layout
|
||||
w1_interleaved = _interleave_w13_halves(w1_bf16_local, group_size=64, dim=1)
|
||||
w1_flat = w1_interleaved.view(num_local_experts * intermediate_size_2x, hidden_size)
|
||||
w1_q_flat, w1_sf_flat = fp4_quantize(
|
||||
w1_flat,
|
||||
global_scale=w1_gs,
|
||||
sf_vec_size=sf_vec_size,
|
||||
is_sf_swizzled_layout=True,
|
||||
)
|
||||
w1_q = w1_q_flat.view(num_local_experts, intermediate_size_2x, hidden_size // 2)
|
||||
w1_sf = convert_sf_to_mma_layout(
|
||||
w1_sf_flat,
|
||||
m=intermediate_size_2x,
|
||||
k=hidden_size,
|
||||
num_groups=num_local_experts,
|
||||
sf_vec_size=sf_vec_size,
|
||||
)
|
||||
w1_alpha = 1.0 / (a1_gs * w1_gs).expand(num_local_experts)
|
||||
|
||||
# GEMM2: quantize -> MMA layout
|
||||
w2_flat = w2_bf16_local.view(num_local_experts * hidden_size, intermediate_size)
|
||||
w2_q_flat, w2_sf_flat = fp4_quantize(
|
||||
w2_flat,
|
||||
global_scale=w2_gs,
|
||||
sf_vec_size=sf_vec_size,
|
||||
is_sf_swizzled_layout=True,
|
||||
)
|
||||
w2_q = w2_q_flat.view(num_local_experts, hidden_size, intermediate_size // 2)
|
||||
w2_sf = convert_sf_to_mma_layout(
|
||||
w2_sf_flat,
|
||||
m=hidden_size,
|
||||
k=intermediate_size,
|
||||
num_groups=num_local_experts,
|
||||
sf_vec_size=sf_vec_size,
|
||||
)
|
||||
w2_alpha = 1.0 / (fc2_input_scale * w2_gs).expand(num_local_experts)
|
||||
|
||||
return {
|
||||
"w1_weight": w1_q,
|
||||
"w1_weight_sf": w1_sf,
|
||||
"w1_alpha": w1_alpha,
|
||||
"w2_weight": w2_q,
|
||||
"w2_weight_sf": w2_sf,
|
||||
"w2_alpha": w2_alpha,
|
||||
}
|
||||
|
||||
|
||||
def _run_wrapper(wrapper, tensors, **overrides):
|
||||
"""Call wrapper.run() with the standard 11-arg dict from _create_cutedsl_wrapper_tensors."""
|
||||
kwargs = dict(
|
||||
x=tensors["x"],
|
||||
x_sf=tensors["x_sf"],
|
||||
token_selected_experts=tensors["token_selected_experts"],
|
||||
token_final_scales=tensors["token_final_scales"],
|
||||
w1_weight=tensors["w1_weight"],
|
||||
w1_weight_sf=tensors["w1_weight_sf"],
|
||||
w1_alpha=tensors["w1_alpha"],
|
||||
fc2_input_scale=tensors["fc2_input_scale"],
|
||||
w2_weight=tensors["w2_weight"],
|
||||
w2_weight_sf=tensors["w2_weight_sf"],
|
||||
w2_alpha=tensors["w2_alpha"],
|
||||
)
|
||||
kwargs.update(overrides)
|
||||
return wrapper.run(**kwargs)
|
||||
|
||||
|
||||
def _quant_dequant_fp4_reference(
|
||||
tensor: torch.Tensor,
|
||||
global_scale: torch.Tensor,
|
||||
sf_vec_size: int = 16,
|
||||
) -> torch.Tensor:
|
||||
"""Simulate FP4 quant-dequant roundtrip for reference computation."""
|
||||
from flashinfer.fp4_quantization import e2m1_and_ufp8sf_scale_to_float
|
||||
|
||||
tensor_bf16 = tensor.to(torch.bfloat16)
|
||||
fp4_packed, sf = fp4_quantize(
|
||||
tensor_bf16,
|
||||
global_scale=global_scale,
|
||||
sf_vec_size=sf_vec_size,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
sf_uint8 = sf.view(torch.uint8).reshape(-1)
|
||||
dequantized = e2m1_and_ufp8sf_scale_to_float(
|
||||
fp4_packed.cpu(),
|
||||
sf_uint8.cpu(),
|
||||
(1.0 / global_scale).cpu(),
|
||||
sf_vec_size=sf_vec_size,
|
||||
ufp8_type=1,
|
||||
is_sf_swizzled_layout=False,
|
||||
).to(tensor.device)
|
||||
return dequantized.float()
|
||||
|
||||
|
||||
def _compute_reference_moe_fp4(
|
||||
hidden_states: torch.Tensor,
|
||||
gemm1_weights: torch.Tensor,
|
||||
gemm2_weights: torch.Tensor,
|
||||
token_selected_experts: torch.Tensor,
|
||||
token_final_scales: torch.Tensor,
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
fc2_input_scale: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Pure-PyTorch MoE reference using bf16 weights (pre-interleave layout).
|
||||
|
||||
gemm1_weights is [num_experts, 2*intermediate_size, hidden_size] with the
|
||||
*original* (un-interleaved) layout: first half = linear, second half = gate.
|
||||
"""
|
||||
device = hidden_states.device
|
||||
num_tokens = hidden_states.shape[0]
|
||||
hidden_states = hidden_states.float()
|
||||
gemm1_weights = gemm1_weights.float()
|
||||
gemm2_weights = gemm2_weights.float()
|
||||
|
||||
output = torch.zeros(num_tokens, hidden_size, dtype=torch.float32, device=device)
|
||||
|
||||
for token_idx in range(num_tokens):
|
||||
token_input = hidden_states[token_idx : token_idx + 1]
|
||||
for k in range(top_k):
|
||||
expert_idx = token_selected_experts[token_idx, k].item()
|
||||
scale = token_final_scales[token_idx, k].item()
|
||||
if expert_idx < 0 or expert_idx >= num_experts:
|
||||
continue
|
||||
|
||||
w1 = gemm1_weights[expert_idx]
|
||||
gemm1_out = token_input @ w1.T
|
||||
|
||||
linear = gemm1_out[:, :intermediate_size]
|
||||
gate = gemm1_out[:, intermediate_size:]
|
||||
swiglu_out = F.silu(gate) * linear
|
||||
|
||||
if fc2_input_scale is not None:
|
||||
swiglu_out = _quant_dequant_fp4_reference(
|
||||
swiglu_out, fc2_input_scale, sf_vec_size=16
|
||||
)
|
||||
|
||||
w2 = gemm2_weights[expert_idx]
|
||||
gemm2_out = swiglu_out @ w2.T
|
||||
output[token_idx] += scale * gemm2_out.squeeze(0)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def compute_routing(router_logits: torch.Tensor, top_k: int):
|
||||
routing_weights = torch.softmax(router_logits, dim=1, dtype=torch.float)
|
||||
routing_weights, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
|
||||
@@ -109,20 +419,6 @@ def prepare_inputs(
|
||||
return hidden_states_3d, masked_m, topk_idx, routing_weights
|
||||
|
||||
|
||||
MNK_FACTORS = [
|
||||
(2, 1024, 1024),
|
||||
(2, 1024, 1536),
|
||||
(2, 3072, 1024),
|
||||
(2, 3072, 1536),
|
||||
(64, 1024, 1024),
|
||||
(64, 1024, 1536),
|
||||
(64, 3072, 1024),
|
||||
(64, 2048, 1024),
|
||||
(224, 1024, 1024),
|
||||
(224, 1024, 1536),
|
||||
]
|
||||
|
||||
|
||||
# Reference implementation of torch_moe
|
||||
def torch_moe(a, w1, w2, score, topk, expert_map):
|
||||
B, D = a.shape
|
||||
@@ -158,7 +454,7 @@ def torch_moe_nvfp4(a, w1, w2, topk, topk_weight, topk_ids):
|
||||
if mask.sum():
|
||||
m = w1[i].shape[0]
|
||||
assert m % 2 == 0
|
||||
# Note: w1 and w3 are swapped!
|
||||
# The first and second W13 halves feed the two SwiGLU branches.
|
||||
w3_expert, w1_expert = w1[i][m // 2 :, :], w1[i][: m // 2, :]
|
||||
inter = F.silu(a[mask] @ w1_expert.t()) * (a[mask] @ w3_expert.t())
|
||||
inter_gs = torch.tensor(1.0).cuda()
|
||||
@@ -177,128 +473,6 @@ def torch_moe_nvfp4(a, w1, w2, topk, topk_weight, topk_ids):
|
||||
).sum(dim=1)
|
||||
|
||||
|
||||
def check_moe(
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
e: int,
|
||||
topk: int,
|
||||
dtype: torch.dtype,
|
||||
moe_impl: Callable,
|
||||
flip_w13: bool,
|
||||
):
|
||||
torch.manual_seed(7)
|
||||
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
|
||||
w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
|
||||
quant_blocksize = 16
|
||||
round_up = lambda x, y: (x + y - 1) // y * y
|
||||
sf_w1_2n = round_up(2 * n, 128)
|
||||
sf_w1_k = round_up(k // quant_blocksize, 4)
|
||||
w1_blockscale = torch.empty(
|
||||
(e, sf_w1_2n, sf_w1_k), device="cuda", dtype=torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
|
||||
sf_w2_k = round_up(k, 128)
|
||||
sf_w2_n = round_up(n // quant_blocksize, 4)
|
||||
w2_blockscale = torch.empty(
|
||||
(e, sf_w2_k, sf_w2_n), device="cuda", dtype=torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
w1_q = torch.empty((e, 2 * n, k // 2), device="cuda", dtype=torch.uint8)
|
||||
w2_q = torch.empty((e, k, n // 2), device="cuda", dtype=torch.uint8)
|
||||
w1_gs = torch.empty((e,), device="cuda", dtype=torch.float32)
|
||||
w2_gs = torch.empty((e,), device="cuda", dtype=torch.float32)
|
||||
|
||||
for expert in range(e):
|
||||
w1_amax = torch.abs(w1).max().to(torch.float32)
|
||||
w2_amax = torch.abs(w2).max().to(torch.float32)
|
||||
w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax
|
||||
w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax
|
||||
|
||||
w1_q[expert], w1_blockscale[expert] = scaled_fp4_quant(
|
||||
w1[expert], w1_gs[expert]
|
||||
)
|
||||
|
||||
w2_q[expert], w2_blockscale[expert] = scaled_fp4_quant(
|
||||
w2[expert], w2_gs[expert]
|
||||
)
|
||||
|
||||
score = torch.randn((m, e), device="cuda", dtype=dtype)
|
||||
|
||||
topk_output = select_experts(
|
||||
hidden_states=a,
|
||||
router_logits=score,
|
||||
topk_config=TopKConfig(top_k=topk, renormalize=False),
|
||||
)
|
||||
topk_weights, topk_ids, _ = topk_output
|
||||
|
||||
a1_gs = torch.ones((e,), device="cuda", dtype=torch.float32)
|
||||
a2_gs = torch.ones((e,), device="cuda", dtype=torch.float32)
|
||||
test_output = moe_impl(
|
||||
a=a,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
w1_q=w1_q,
|
||||
w2_q=w2_q,
|
||||
a1_gs=a1_gs,
|
||||
w1_blockscale=w1_blockscale,
|
||||
w1_alphas=(1 / w1_gs),
|
||||
a2_gs=a2_gs,
|
||||
w2_blockscale=w2_blockscale,
|
||||
w2_alphas=(1 / w2_gs),
|
||||
)
|
||||
|
||||
# Reference check:
|
||||
a_global_scale = (
|
||||
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
a_fp4, a_scale_interleaved = scaled_fp4_quant(a, a_global_scale)
|
||||
_, m_k = a_fp4.shape
|
||||
a_in_dtype = dequantize_nvfp4_to_dtype(
|
||||
a_fp4,
|
||||
a_scale_interleaved,
|
||||
a_global_scale,
|
||||
dtype=a.dtype,
|
||||
device=a.device,
|
||||
block_size=quant_blocksize,
|
||||
)
|
||||
|
||||
w1_d = torch.empty((e, 2 * n, k), device="cuda", dtype=dtype)
|
||||
w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
|
||||
|
||||
for idx in range(0, e):
|
||||
w1_d[idx] = dequantize_nvfp4_to_dtype(
|
||||
w1_q[idx],
|
||||
w1_blockscale[idx],
|
||||
w1_gs[idx],
|
||||
dtype=w1.dtype,
|
||||
device=w1.device,
|
||||
block_size=quant_blocksize,
|
||||
)
|
||||
w2_d[idx] = dequantize_nvfp4_to_dtype(
|
||||
w2_q[idx],
|
||||
w2_blockscale[idx],
|
||||
w2_gs[idx],
|
||||
dtype=w2.dtype,
|
||||
device=w2.device,
|
||||
block_size=quant_blocksize,
|
||||
)
|
||||
|
||||
if flip_w13:
|
||||
dim = -2
|
||||
size = w1_d.size(dim)
|
||||
assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}"
|
||||
half = size // 2
|
||||
# Reorder weight
|
||||
w1, w3 = w1_d.split(half, dim=dim)
|
||||
w1_d = torch.cat([w3, w1], dim=dim).contiguous()
|
||||
|
||||
torch_output = torch_moe(a_in_dtype, w1_d, w2_d, score, topk, None)
|
||||
|
||||
torch.testing.assert_close(torch_output, test_output, atol=1e-1, rtol=1e-1)
|
||||
|
||||
|
||||
class TestFlashinferCutedslMoe(unittest.TestCase):
|
||||
@unittest.skipIf(SKIP_TEST, SKIP_REASON)
|
||||
def test_flashinfer_cutedsl_moe_masked(self):
|
||||
@@ -316,9 +490,6 @@ class TestFlashinferCutedslMoe(unittest.TestCase):
|
||||
with self.subTest(
|
||||
bs=bs, hidden_dim=hidden_dim, inter_dim=inter_dim, topk=topk
|
||||
):
|
||||
print(
|
||||
f"Testing with bs={bs}, hidden_dim={hidden_dim}, inter_dim={inter_dim}, topk={topk}"
|
||||
)
|
||||
with torch.inference_mode():
|
||||
torch.manual_seed(42)
|
||||
device = "cuda"
|
||||
@@ -476,8 +647,296 @@ class TestFlashinferCutedslMoe(unittest.TestCase):
|
||||
torch.testing.assert_close(
|
||||
out_weighted.cpu(), ref_output.cpu(), atol=5e-2, rtol=5e-2
|
||||
)
|
||||
print(
|
||||
f"Test passed with bs={bs}, hidden_dim={hidden_dim}, inter_dim={inter_dim}, topk={topk}"
|
||||
|
||||
@unittest.skipIf(SKIP_TEST, SKIP_REASON)
|
||||
@unittest.skipIf(
|
||||
CuteDslMoEWrapper is None or convert_sf_to_mma_layout is None,
|
||||
"CuteDslMoEWrapper / convert_sf_to_mma_layout not available",
|
||||
)
|
||||
def test_cutedsl_moe_wrapper_run(self):
|
||||
"""Call CuteDslMoEWrapper.run() with MMA-layout tensors and verify against reference."""
|
||||
test_cases = [
|
||||
# (num_tokens, hidden_size, intermediate_size, num_experts, top_k)
|
||||
# Minimum dimensions match FlashInfer's test_wrapper_accuracy:
|
||||
# num_experts >= 256, hidden_size >= 256, intermediate_size >= 512,
|
||||
# num_tokens >= 128. The CuteDSL GEMM kernels have tile-size
|
||||
# constraints that make smaller dimensions unreliable.
|
||||
(128, 256, 512, 256, 2),
|
||||
(128, 256, 512, 256, 8),
|
||||
(256, 256, 512, 256, 4),
|
||||
]
|
||||
|
||||
for (
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
intermediate_size,
|
||||
num_experts,
|
||||
top_k,
|
||||
) in test_cases:
|
||||
with self.subTest(
|
||||
num_tokens=num_tokens,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
top_k=top_k,
|
||||
):
|
||||
tensors = _create_cutedsl_wrapper_tensors(
|
||||
num_tokens=num_tokens,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
)
|
||||
|
||||
wrapper = CuteDslMoEWrapper(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
use_cuda_graph=False,
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
out = _run_wrapper(wrapper, tensors)
|
||||
|
||||
self.assertEqual(out.shape, (num_tokens, hidden_size))
|
||||
self.assertEqual(out.dtype, torch.bfloat16)
|
||||
self.assertFalse(
|
||||
torch.isnan(out).any().item() or torch.isinf(out).any().item(),
|
||||
"Output contains NaN or Inf",
|
||||
)
|
||||
|
||||
ref_output = _compute_reference_moe_fp4(
|
||||
hidden_states=tensors["x_bf16"].float().cuda(),
|
||||
gemm1_weights=tensors["w1_weight_bf16"].float().cuda(),
|
||||
gemm2_weights=tensors["w2_weight_bf16"].float().cuda(),
|
||||
token_selected_experts=tensors["token_selected_experts"],
|
||||
token_final_scales=tensors["token_final_scales"],
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
fc2_input_scale=tensors["fc2_input_scale"],
|
||||
)
|
||||
|
||||
out_f32 = out.float()
|
||||
ref_f32 = ref_output.float()
|
||||
output_scale = max(ref_f32.std().item(), 0.01)
|
||||
atol = max(0.1, 3.0 * output_scale)
|
||||
rtol = 0.85
|
||||
abs_diff = torch.abs(out_f32 - ref_f32)
|
||||
rel_diff = abs_diff / (torch.abs(ref_f32) + 1e-8)
|
||||
within_tol = (abs_diff < atol) | (rel_diff < rtol)
|
||||
pct_within = within_tol.float().mean().item()
|
||||
self.assertGreaterEqual(
|
||||
pct_within,
|
||||
0.925,
|
||||
f"Only {pct_within * 100:.2f}% of elements within tolerance "
|
||||
f"(atol={atol:.4f})",
|
||||
)
|
||||
|
||||
@unittest.skipIf(SKIP_TEST, SKIP_REASON)
|
||||
@unittest.skipIf(
|
||||
CuteDslMoEWrapper is None or convert_sf_to_mma_layout is None,
|
||||
"CuteDslMoEWrapper / convert_sf_to_mma_layout not available",
|
||||
)
|
||||
def test_cutedsl_cuda_graph_parity(self):
|
||||
"""Verify non-graph and cuda_graph wrappers produce identical results.
|
||||
|
||||
Also checks both match the pure-PyTorch reference, and that a second
|
||||
cuda_graph pass reuses buffers deterministically (subsumes the former
|
||||
cuda_graph_smoke test).
|
||||
"""
|
||||
test_cases = [
|
||||
# (num_tokens, hidden_size, intermediate_size, num_experts, top_k)
|
||||
(128, 256, 512, 256, 2),
|
||||
(256, 256, 512, 256, 4),
|
||||
]
|
||||
|
||||
for (
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
intermediate_size,
|
||||
num_experts,
|
||||
top_k,
|
||||
) in test_cases:
|
||||
with self.subTest(
|
||||
num_tokens=num_tokens,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
top_k=top_k,
|
||||
):
|
||||
tensors = _create_cutedsl_wrapper_tensors(
|
||||
num_tokens=num_tokens,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
)
|
||||
|
||||
wrapper_args = dict(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
)
|
||||
wrapper_no_graph = CuteDslMoEWrapper(
|
||||
**wrapper_args, use_cuda_graph=False
|
||||
)
|
||||
wrapper_graph = CuteDslMoEWrapper(
|
||||
**wrapper_args,
|
||||
use_cuda_graph=True,
|
||||
max_num_tokens=num_tokens,
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
out_no_graph = _run_wrapper(wrapper_no_graph, tensors)
|
||||
out_graph = _run_wrapper(wrapper_graph, tensors)
|
||||
out_graph2 = _run_wrapper(wrapper_graph, tensors)
|
||||
|
||||
torch.testing.assert_close(
|
||||
out_no_graph,
|
||||
out_graph,
|
||||
atol=1e-2,
|
||||
rtol=1e-2,
|
||||
msg="non-graph vs cuda_graph wrapper outputs diverge",
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
out_graph,
|
||||
out_graph2,
|
||||
atol=1e-5,
|
||||
rtol=1e-5,
|
||||
msg="second cuda_graph pass should reuse buffers identically",
|
||||
)
|
||||
|
||||
ref_output = _compute_reference_moe_fp4(
|
||||
hidden_states=tensors["x_bf16"].float().cuda(),
|
||||
gemm1_weights=tensors["w1_weight_bf16"].float().cuda(),
|
||||
gemm2_weights=tensors["w2_weight_bf16"].float().cuda(),
|
||||
token_selected_experts=tensors["token_selected_experts"],
|
||||
token_final_scales=tensors["token_final_scales"],
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
fc2_input_scale=tensors["fc2_input_scale"],
|
||||
)
|
||||
|
||||
out_f32 = out_graph.float()
|
||||
ref_f32 = ref_output.float()
|
||||
output_scale = max(ref_f32.std().item(), 0.01)
|
||||
atol = max(0.1, 3.0 * output_scale)
|
||||
rtol = 0.85
|
||||
abs_diff = torch.abs(out_f32 - ref_f32)
|
||||
rel_diff = abs_diff / (torch.abs(ref_f32) + 1e-8)
|
||||
within_tol = (abs_diff < atol) | (rel_diff < rtol)
|
||||
pct_within = within_tol.float().mean().item()
|
||||
self.assertGreaterEqual(
|
||||
pct_within,
|
||||
0.925,
|
||||
f"graph vs reference: only {pct_within * 100:.2f}% within tol",
|
||||
)
|
||||
|
||||
@unittest.skipIf(SKIP_TEST, SKIP_REASON)
|
||||
@unittest.skipIf(
|
||||
CuteDslMoEWrapper is None or convert_sf_to_mma_layout is None,
|
||||
"CuteDslMoEWrapper / convert_sf_to_mma_layout not available",
|
||||
)
|
||||
def test_cutedsl_ep_sharded_allreduce(self):
|
||||
"""Verify EP-sharded execution: partial outputs from EP ranks sum to full result.
|
||||
|
||||
Simulates the EP=TP all-reduce pattern used by the CuteDSL moe_runner when
|
||||
ep_size > 1 and moe_a2a_backend=none. Each "rank" runs a wrapper with
|
||||
num_local_experts < num_experts and a corresponding local_expert_offset,
|
||||
receiving only the local slice of weights/scales/alphas — matching the
|
||||
real runtime contract where each rank holds only its own expert partition.
|
||||
The partial outputs are summed (simulating tensor_model_parallel_all_reduce)
|
||||
and compared against a single wrapper processing all experts.
|
||||
"""
|
||||
test_cases = [
|
||||
# (num_tokens, hidden_size, intermediate_size, num_experts, top_k, ep_size)
|
||||
# Dimensions match FlashInfer's minimum wrapper requirements.
|
||||
(128, 256, 512, 256, 2, 2),
|
||||
(128, 256, 512, 256, 2, 4),
|
||||
(128, 256, 512, 256, 8, 8),
|
||||
]
|
||||
|
||||
for (
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
intermediate_size,
|
||||
num_experts,
|
||||
top_k,
|
||||
ep_size,
|
||||
) in test_cases:
|
||||
with self.subTest(
|
||||
num_tokens=num_tokens,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
top_k=top_k,
|
||||
ep_size=ep_size,
|
||||
):
|
||||
assert num_experts % ep_size == 0
|
||||
num_local_experts = num_experts // ep_size
|
||||
|
||||
tensors = _create_cutedsl_wrapper_tensors(
|
||||
num_tokens=num_tokens,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
)
|
||||
|
||||
# Full-expert baseline (EP=1): all experts on one "rank"
|
||||
wrapper_full = CuteDslMoEWrapper(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
use_cuda_graph=False,
|
||||
)
|
||||
with torch.no_grad():
|
||||
out_full = _run_wrapper(wrapper_full, tensors)
|
||||
|
||||
# EP-sharded: each rank independently quantizes its local
|
||||
# bf16 weight shard and calls convert_sf_to_mma_layout with
|
||||
# num_groups=num_local_experts — matching the real per-rank
|
||||
# weight preprocessing in the CuteDSL moe_runner path.
|
||||
accumulated = torch.zeros_like(out_full)
|
||||
for rank in range(ep_size):
|
||||
lo = rank * num_local_experts
|
||||
hi = lo + num_local_experts
|
||||
|
||||
local_tensors = _quantize_local_expert_weights(
|
||||
w1_bf16_local=tensors["w1_weight_bf16"][lo:hi],
|
||||
w2_bf16_local=tensors["w2_weight_bf16"][lo:hi],
|
||||
a1_gs=tensors["a1_gs"],
|
||||
w1_gs=tensors["w1_gs"],
|
||||
w2_gs=tensors["w2_gs"],
|
||||
fc2_input_scale=tensors["fc2_input_scale"],
|
||||
)
|
||||
|
||||
wrapper_shard = CuteDslMoEWrapper(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
use_cuda_graph=False,
|
||||
num_local_experts=num_local_experts,
|
||||
local_expert_offset=lo,
|
||||
)
|
||||
with torch.no_grad():
|
||||
partial = _run_wrapper(wrapper_shard, tensors, **local_tensors)
|
||||
accumulated += partial
|
||||
|
||||
torch.testing.assert_close(
|
||||
out_full,
|
||||
accumulated,
|
||||
atol=1e-2,
|
||||
rtol=1e-2,
|
||||
msg=(
|
||||
f"EP-sharded all-reduce mismatch "
|
||||
f"(ep_size={ep_size}, tokens={num_tokens})"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user