wehub-resource-sync
773 lines
24 KiB
Python
773 lines
24 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from typing import cast
|
|
|
|
import huggingface_hub
|
|
import pytest
|
|
import torch
|
|
from safetensors import safe_open
|
|
|
|
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
|
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
|
from vllm.model_executor.layers.fused_moe.config import (
|
|
FusedMoEConfig,
|
|
FusedMoEParallelConfig,
|
|
FusedMoEQuantConfig,
|
|
RoutingMethodType,
|
|
nvfp4_moe_quant_config,
|
|
)
|
|
from vllm.model_executor.layers.fused_moe.experts.nvfp4_emulation_moe import (
|
|
Nvfp4QuantizationEmulationTritonExperts,
|
|
)
|
|
from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
|
|
from vllm.model_executor.layers.quantization.utils import (
|
|
nvfp4_emulation_utils,
|
|
)
|
|
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
|
|
dequantize_to_dtype,
|
|
ref_nvfp4_quant_dequant,
|
|
)
|
|
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
|
QuantKey,
|
|
kNvfp4Dynamic,
|
|
kNvfp4Static,
|
|
)
|
|
from vllm.platforms import current_platform
|
|
from vllm.triton_utils import triton
|
|
|
|
if current_platform.is_rocm():
|
|
from vllm.platforms.rocm import on_gfx950
|
|
else:
|
|
|
|
def on_gfx950() -> bool:
|
|
return False
|
|
|
|
|
|
MOE_MODEL_CONFIGS = {
|
|
"nvidia/Qwen3-30B-A3B-NVFP4": {
|
|
"shards": ["model-00001-of-00004.safetensors"],
|
|
"expert_prefix": "model.layers.9.mlp.experts.",
|
|
# Position of the expert index in the dot-split key.
|
|
"expert_idx_pos": 5,
|
|
}
|
|
}
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def loaded_model_files():
|
|
return {
|
|
model_id: huggingface_hub.snapshot_download(
|
|
repo_id=model_id, allow_patterns=config["shards"]
|
|
)
|
|
for model_id, config in MOE_MODEL_CONFIGS.items()
|
|
}
|
|
|
|
|
|
class Nvfp4QuantizationEmulationTritonExpertsReference(TritonExperts):
|
|
"""
|
|
Extension of TritonExperts to support emulated NVFP4 MoE experts.
|
|
|
|
It may be used for NVFP4 models when the device does not have
|
|
native support for this dtype.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
moe_config: FusedMoEConfig,
|
|
quant_config: FusedMoEQuantConfig,
|
|
):
|
|
super().__init__(moe_config, quant_config)
|
|
|
|
# `TritonExperts.apply` expects pre-dequantized weights,
|
|
# which we handle in `apply` below.
|
|
self.w1_scale_val = self.quant_config.w1_scale
|
|
self.w2_scale_val = self.quant_config.w2_scale
|
|
|
|
self.quant_config._w1.scale = None
|
|
self.quant_config._w2.scale = None
|
|
|
|
self.quantization_emulation = True
|
|
|
|
@property
|
|
def quant_dtype(self) -> torch.dtype | str | None:
|
|
return "nvfp4"
|
|
|
|
@property
|
|
def a1_scale(self) -> torch.Tensor | None:
|
|
return self.quant_config.a1_gscale
|
|
|
|
@property
|
|
def expects_unquantized_inputs(self) -> bool:
|
|
return True
|
|
|
|
@staticmethod
|
|
def _supports_quant_scheme(
|
|
weight_key: QuantKey | None,
|
|
activation_key: QuantKey | None,
|
|
) -> bool:
|
|
return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic)
|
|
|
|
def apply(
|
|
self,
|
|
output: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
w1: torch.Tensor,
|
|
w2: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
activation: MoEActivation,
|
|
global_num_experts: int,
|
|
expert_map: torch.Tensor | None,
|
|
a1q_scale: torch.Tensor | None,
|
|
a2_scale: torch.Tensor | None,
|
|
workspace13: torch.Tensor,
|
|
workspace2: torch.Tensor,
|
|
expert_tokens_meta: mk.ExpertTokensMetadata | None,
|
|
apply_router_weight_on_input: bool,
|
|
):
|
|
assert w1.dtype == torch.uint8
|
|
assert w2.dtype == torch.uint8
|
|
|
|
# Dequantize w1 from packed NVFP4 to fp16/bf16
|
|
w13_global_scale = self.quant_config.g1_alphas
|
|
|
|
w1_dequant = dequantize_to_dtype(
|
|
tensor_fp4=w1,
|
|
tensor_sf=self.w1_scale_val,
|
|
global_scale=w13_global_scale,
|
|
dtype=hidden_states.dtype,
|
|
block_size=16,
|
|
swizzle=False,
|
|
)
|
|
|
|
# Dequantize w2 from packed NVFP4 to fp16/bf16
|
|
w2_global_scale = self.quant_config.g2_alphas
|
|
|
|
w2_dequant = dequantize_to_dtype(
|
|
tensor_fp4=w2,
|
|
tensor_sf=self.w2_scale_val,
|
|
global_scale=w2_global_scale,
|
|
dtype=hidden_states.dtype,
|
|
block_size=16,
|
|
swizzle=False,
|
|
)
|
|
|
|
super().apply(
|
|
output=output,
|
|
hidden_states=hidden_states,
|
|
w1=w1_dequant,
|
|
w2=w2_dequant,
|
|
topk_weights=topk_weights,
|
|
topk_ids=topk_ids,
|
|
activation=activation,
|
|
global_num_experts=global_num_experts,
|
|
expert_map=expert_map,
|
|
a1q_scale=None,
|
|
a2_scale=self.quant_config.a2_gscale,
|
|
workspace13=workspace13,
|
|
workspace2=workspace2,
|
|
expert_tokens_meta=expert_tokens_meta,
|
|
apply_router_weight_on_input=apply_router_weight_on_input,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("config_kwargs", "expected_reason"),
|
|
[
|
|
({"has_bias": True}, "kernel does not support bias"),
|
|
({"is_lora_enabled": True}, "kernel does not support LoRA"),
|
|
],
|
|
)
|
|
def test_nvfp4_emulation_support_check_rejects_bias_and_lora(
|
|
config_kwargs: dict[str, bool],
|
|
expected_reason: str,
|
|
) -> None:
|
|
moe_config = FusedMoEConfig(
|
|
num_experts=2,
|
|
experts_per_token=1,
|
|
hidden_dim=16,
|
|
intermediate_size=16,
|
|
num_local_experts=2,
|
|
num_logical_experts=2,
|
|
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
|
|
activation=MoEActivation.SILU,
|
|
in_dtype=torch.bfloat16,
|
|
device="cuda",
|
|
routing_method=RoutingMethodType.TopK,
|
|
**config_kwargs,
|
|
)
|
|
|
|
supported, reason = Nvfp4QuantizationEmulationTritonExperts.is_supported_config(
|
|
Nvfp4QuantizationEmulationTritonExperts,
|
|
moe_config,
|
|
kNvfp4Static,
|
|
kNvfp4Dynamic,
|
|
mk.FusedMoEActivationFormat.Standard,
|
|
)
|
|
|
|
assert not supported
|
|
assert reason == expected_reason
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda_alike(),
|
|
reason="Triton NVFP4 kernel requires CUDA.",
|
|
)
|
|
def test_triton_dequantize_nvfp4(monkeypatch, loaded_model_files) -> None:
|
|
"""Test the Triton dequantization kernel against the CPU reference
|
|
using real NVFP4 weights from a checkpoint.
|
|
|
|
Tests both 2D (attention projection) and 3D (stacked MoE experts).
|
|
"""
|
|
checkpoint_path = loaded_model_files["nvidia/Qwen3-30B-A3B-NVFP4"]
|
|
shards = cast(list[str], MOE_MODEL_CONFIGS["nvidia/Qwen3-30B-A3B-NVFP4"]["shards"])
|
|
shard_path = f"{checkpoint_path}/{shards[0]}"
|
|
block_size = 16
|
|
|
|
with safe_open(shard_path, framework="pt", device="cpu") as f:
|
|
all_keys = list(f.keys())
|
|
|
|
# 2D case: attention projection
|
|
tensor_fp4_2d = f.get_tensor("model.layers.9.self_attn.k_proj.weight")
|
|
tensor_sf_2d = f.get_tensor("model.layers.9.self_attn.k_proj.weight_scale")
|
|
global_scale_2d = f.get_tensor("model.layers.9.self_attn.k_proj.weight_scale_2")
|
|
|
|
# 3D case: stack ALL experts for layer 9 up_proj
|
|
expert_prefix = "model.layers.9.mlp.experts."
|
|
expert_indices = sorted(
|
|
int(key.split(".")[5])
|
|
for key in all_keys
|
|
if key.startswith(expert_prefix) and key.endswith(".up_proj.weight")
|
|
)
|
|
assert len(expert_indices) > 0
|
|
|
|
all_fp4 = []
|
|
all_sf = []
|
|
all_global_scale = []
|
|
for index in expert_indices:
|
|
name = f"{expert_prefix}{index}.up_proj"
|
|
all_fp4.append(f.get_tensor(f"{name}.weight"))
|
|
all_sf.append(f.get_tensor(f"{name}.weight_scale"))
|
|
all_global_scale.append(f.get_tensor(f"{name}.weight_scale_2"))
|
|
|
|
tensor_fp4_3d = torch.stack(all_fp4)
|
|
tensor_sf_3d = torch.stack(all_sf)
|
|
global_scale_3d = torch.stack(all_global_scale)
|
|
|
|
test_cases = [
|
|
("2D base", tensor_fp4_2d, tensor_sf_2d, global_scale_2d),
|
|
(
|
|
"2D 2x rows",
|
|
tensor_fp4_2d.repeat(2, 1),
|
|
tensor_sf_2d.repeat(2, 1),
|
|
global_scale_2d,
|
|
),
|
|
(
|
|
"2D 4x rows",
|
|
tensor_fp4_2d.repeat(4, 1),
|
|
tensor_sf_2d.repeat(4, 1),
|
|
global_scale_2d,
|
|
),
|
|
(
|
|
"2D 2x cols",
|
|
tensor_fp4_2d.repeat(1, 2),
|
|
tensor_sf_2d.repeat(1, 2),
|
|
global_scale_2d,
|
|
),
|
|
("3D base", tensor_fp4_3d, tensor_sf_3d, global_scale_3d),
|
|
(
|
|
"3D 2x experts",
|
|
tensor_fp4_3d.repeat(2, 1, 1),
|
|
tensor_sf_3d.repeat(2, 1, 1),
|
|
global_scale_3d.repeat(2),
|
|
),
|
|
(
|
|
"3D 2x rows",
|
|
tensor_fp4_3d.repeat(1, 2, 1),
|
|
tensor_sf_3d.repeat(1, 2, 1),
|
|
global_scale_3d,
|
|
),
|
|
(
|
|
"3D 2x cols",
|
|
tensor_fp4_3d.repeat(1, 1, 2),
|
|
tensor_sf_3d.repeat(1, 1, 2),
|
|
global_scale_3d,
|
|
),
|
|
]
|
|
|
|
quantiles = [0.5, 0.001, 0.999]
|
|
|
|
# Move the E2M1 lookup table to CUDA ahead of time, as would normally
|
|
# happen during model loading (process_weights_after_loading). Both the
|
|
# Triton and PyTorch reference paths run on CUDA.
|
|
nvfp4_emulation_utils.kE2M1ToFloat_handle.val = (
|
|
nvfp4_emulation_utils.kE2M1ToFloat_handle.val.cuda()
|
|
)
|
|
|
|
for label, tensor_fp4, tensor_sf, global_scale in test_cases:
|
|
fp4_cuda = tensor_fp4.cuda()
|
|
sf_cuda = tensor_sf.cuda()
|
|
gs_cuda = global_scale.cuda()
|
|
|
|
# Triton path
|
|
triton_result = dequantize_to_dtype(
|
|
fp4_cuda,
|
|
sf_cuda,
|
|
gs_cuda,
|
|
torch.bfloat16,
|
|
block_size,
|
|
swizzle=False,
|
|
)
|
|
|
|
# Reference path (PyTorch ops on CUDA, Triton dispatch disabled)
|
|
with monkeypatch.context() as m:
|
|
m.setattr(
|
|
nvfp4_emulation_utils.current_platform,
|
|
"is_cuda_alike",
|
|
lambda: False,
|
|
)
|
|
reference = dequantize_to_dtype(
|
|
fp4_cuda,
|
|
sf_cuda,
|
|
gs_cuda,
|
|
torch.bfloat16,
|
|
block_size,
|
|
swizzle=False,
|
|
)
|
|
|
|
torch.testing.assert_close(triton_result, reference, atol=0, rtol=0)
|
|
|
|
# Benchmark
|
|
shape = list(tensor_fp4.shape)
|
|
|
|
def _triton_bench(
|
|
fp4_cuda=fp4_cuda,
|
|
scale_cuda=sf_cuda,
|
|
global_scale_cuda=gs_cuda,
|
|
block_size=block_size,
|
|
):
|
|
return dequantize_to_dtype(
|
|
fp4_cuda,
|
|
scale_cuda,
|
|
global_scale_cuda,
|
|
torch.bfloat16,
|
|
block_size,
|
|
swizzle=False,
|
|
)
|
|
|
|
triton_ms, triton_min, triton_max = triton.testing.do_bench(
|
|
_triton_bench, quantiles=quantiles
|
|
)
|
|
|
|
def _reference_bench(
|
|
fp4_cuda=fp4_cuda,
|
|
scale_cuda=sf_cuda,
|
|
global_scale_cuda=gs_cuda,
|
|
block_size=block_size,
|
|
):
|
|
with monkeypatch.context() as m2:
|
|
m2.setattr(
|
|
nvfp4_emulation_utils.current_platform,
|
|
"is_cuda_alike",
|
|
lambda: False,
|
|
)
|
|
dequantize_to_dtype(
|
|
fp4_cuda,
|
|
scale_cuda,
|
|
global_scale_cuda,
|
|
torch.bfloat16,
|
|
block_size,
|
|
swizzle=False,
|
|
)
|
|
|
|
ref_ms, ref_min, ref_max = triton.testing.do_bench(
|
|
_reference_bench, quantiles=quantiles
|
|
)
|
|
|
|
speedup = ref_ms / triton_ms if triton_ms > 0 else float("inf")
|
|
print(f" dequantize {label} {shape}:")
|
|
print(
|
|
f" triton: median={triton_ms:.3f}ms, "
|
|
f"min={triton_min:.3f}ms, max={triton_max:.3f}ms"
|
|
)
|
|
print(
|
|
f" reference: median={ref_ms:.3f}ms, "
|
|
f"min={ref_min:.3f}ms, max={ref_max:.3f}ms"
|
|
)
|
|
print(f" speedup: {speedup:.2f}x")
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda_alike(),
|
|
reason="Triton NVFP4 kernel requires CUDA.",
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"m, k",
|
|
[
|
|
(1, 16),
|
|
(1, 4096),
|
|
(2, 4096),
|
|
(4, 4096),
|
|
(8, 4096),
|
|
(16, 4096),
|
|
(24, 4096),
|
|
(32, 4096),
|
|
(1, 8192),
|
|
(2, 8192),
|
|
(4, 8192),
|
|
(8, 8192),
|
|
(16, 8192),
|
|
(24, 8192),
|
|
(32, 8192),
|
|
(1, 32),
|
|
(2, 48),
|
|
(7, 64),
|
|
(16, 128),
|
|
(33, 160),
|
|
(128, 256),
|
|
(256, 512),
|
|
(1024, 1024),
|
|
(5120, 2048),
|
|
(2048, 4096),
|
|
(4096, 7168),
|
|
(8192, 8192),
|
|
(128, 16384),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("global_scale_value", [0.5, 1.0, 0.001])
|
|
def test_triton_nvfp4_quant_dequant(
|
|
monkeypatch, m: int, k: int, global_scale_value: float
|
|
) -> None:
|
|
"""Test the Triton quant-dequant kernel against the CPU reference."""
|
|
block_size = 16
|
|
x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
|
|
global_scale = torch.tensor(global_scale_value, dtype=torch.float32, device="cuda")
|
|
|
|
# Triton path
|
|
triton_result = ref_nvfp4_quant_dequant(x, global_scale, block_size)
|
|
|
|
# CPU reference path
|
|
with monkeypatch.context() as mp:
|
|
mp.setattr(
|
|
nvfp4_emulation_utils.current_platform,
|
|
"is_cuda_alike",
|
|
lambda: False,
|
|
)
|
|
reference = ref_nvfp4_quant_dequant(x.cpu(), global_scale.cpu(), block_size)
|
|
|
|
torch.testing.assert_close(triton_result.cpu(), reference, atol=0, rtol=0)
|
|
|
|
# Benchmark (both paths on CUDA tensors for fair comparison)
|
|
quantiles = [0.5, 0.001, 0.999]
|
|
|
|
def _triton_bench(
|
|
input_tensor=x, input_global_scale=global_scale, input_block_size=block_size
|
|
):
|
|
return ref_nvfp4_quant_dequant(
|
|
input_tensor, input_global_scale, input_block_size
|
|
)
|
|
|
|
triton_ms, triton_min, triton_max = triton.testing.do_bench(
|
|
_triton_bench, quantiles=quantiles
|
|
)
|
|
|
|
def _reference_bench(
|
|
input_tensor=x, input_global_scale=global_scale, input_block_size=block_size
|
|
):
|
|
with monkeypatch.context() as mp2:
|
|
mp2.setattr(
|
|
nvfp4_emulation_utils.current_platform,
|
|
"is_cuda_alike",
|
|
lambda: False,
|
|
)
|
|
ref_nvfp4_quant_dequant(input_tensor, input_global_scale, input_block_size)
|
|
|
|
ref_ms, ref_min, ref_max = triton.testing.do_bench(
|
|
_reference_bench, quantiles=quantiles
|
|
)
|
|
|
|
speedup = ref_ms / triton_ms if triton_ms > 0 else float("inf")
|
|
print(f" quant_dequant [{m}x{k}] gs={global_scale_value}:")
|
|
print(
|
|
f" triton: median={triton_ms:.3f}ms, "
|
|
f"min={triton_min:.3f}ms, max={triton_max:.3f}ms"
|
|
)
|
|
print(
|
|
f" reference: median={ref_ms:.3f}ms, "
|
|
f"min={ref_min:.3f}ms, max={ref_max:.3f}ms"
|
|
)
|
|
print(f" speedup: {speedup:.2f}x")
|
|
|
|
|
|
def _load_nvfp4_moe_weights(
|
|
model_files: dict[str, str],
|
|
model_id: str,
|
|
tensor_parallel_size: int,
|
|
max_experts: int | None = None,
|
|
):
|
|
"""Load and stack NVFP4 MoE weights from checkpoint shards.
|
|
|
|
Returns (w1, w1_scale, w1_gscale, w2, w2_scale, w2_gscale,
|
|
a1_gscale, a2_gscale, num_experts, hidden_dim,
|
|
intermediate_size).
|
|
|
|
When max_experts is set, only the first max_experts experts are loaded.
|
|
|
|
When tensor_parallel_size > 1, the N dimension of w1 and the K
|
|
dimension of w2 are narrowed to the first TP shard (simulating
|
|
column-parallel on w1 / row-parallel on w2).
|
|
"""
|
|
cfg = MOE_MODEL_CONFIGS[model_id]
|
|
shards = cast(list[str], cfg["shards"])
|
|
checkpoint_path = model_files[model_id]
|
|
|
|
expert_prefix = cfg["expert_prefix"]
|
|
idx_pos = cast(int, cfg["expert_idx_pos"])
|
|
|
|
# Collect all tensors across shards into a flat dict — an expert's
|
|
# tensors may be split across multiple shard files.
|
|
all_tensors: dict[str, torch.Tensor] = {}
|
|
for shard_name in shards:
|
|
shard_path = f"{checkpoint_path}/{shard_name}"
|
|
with safe_open(shard_path, framework="pt", device="cpu") as f:
|
|
for key in f.keys(): # noqa: SIM118
|
|
if key.startswith(expert_prefix):
|
|
all_tensors[key] = f.get_tensor(key)
|
|
|
|
expert_indices = sorted(
|
|
{
|
|
int(key.split(".")[idx_pos])
|
|
for key in all_tensors
|
|
if key.endswith(".gate_proj.weight")
|
|
}
|
|
)
|
|
if max_experts is not None:
|
|
expert_indices = expert_indices[:max_experts]
|
|
num_experts = len(expert_indices)
|
|
|
|
gate_weights, up_weights, down_weights = [], [], []
|
|
gate_scales, up_scales, down_scales = [], [], []
|
|
gate_gscales, up_gscales, down_gscales = [], [], []
|
|
a1_scales, a2_scales = [], []
|
|
|
|
for idx in expert_indices:
|
|
prefix = f"{expert_prefix}{idx}"
|
|
gate_weights.append(all_tensors[f"{prefix}.gate_proj.weight"])
|
|
gate_scales.append(all_tensors[f"{prefix}.gate_proj.weight_scale"])
|
|
gate_gscales.append(all_tensors[f"{prefix}.gate_proj.weight_scale_2"])
|
|
up_weights.append(all_tensors[f"{prefix}.up_proj.weight"])
|
|
up_scales.append(all_tensors[f"{prefix}.up_proj.weight_scale"])
|
|
up_gscales.append(all_tensors[f"{prefix}.up_proj.weight_scale_2"])
|
|
down_weights.append(all_tensors[f"{prefix}.down_proj.weight"])
|
|
down_scales.append(all_tensors[f"{prefix}.down_proj.weight_scale"])
|
|
down_gscales.append(all_tensors[f"{prefix}.down_proj.weight_scale_2"])
|
|
a1_scales.append(all_tensors[f"{prefix}.gate_proj.input_scale"])
|
|
a2_scales.append(all_tensors[f"{prefix}.down_proj.input_scale"])
|
|
|
|
# Stack into MoE format.
|
|
# w1 = [E, 2*intermediate, hidden//2] (gate + up concatenated)
|
|
w1 = torch.stack(
|
|
[torch.cat([g, u], dim=0) for g, u in zip(gate_weights, up_weights)]
|
|
).cuda()
|
|
w1_scale = torch.stack(
|
|
[torch.cat([g, u], dim=0) for g, u in zip(gate_scales, up_scales)]
|
|
).cuda()
|
|
w1_gscale = torch.stack(gate_gscales).cuda()
|
|
|
|
# w2 = [E, hidden, intermediate//2]
|
|
w2 = torch.stack(down_weights).cuda()
|
|
w2_scale = torch.stack(down_scales).cuda()
|
|
w2_gscale = torch.stack(down_gscales).cuda()
|
|
|
|
a13_scale_raw = torch.stack(a1_scales).cuda()
|
|
a2_scale_raw = torch.stack(a2_scales).cuda()
|
|
|
|
# Apply EMULATION transforms (matches oracle/nvfp4.py).
|
|
nvfp4_emulation_utils.kE2M1ToFloat_handle.val = (
|
|
nvfp4_emulation_utils.kE2M1ToFloat_handle.val.cuda()
|
|
)
|
|
a1_gscale = 1.0 / a13_scale_raw.max().to(torch.float32)
|
|
a2_gscale = 1.0 / a2_scale_raw.max().to(torch.float32)
|
|
|
|
# ── Simulate TP sharding ──
|
|
# w1 (gate_up): column-parallel → shard the N dimension (dim 1).
|
|
# w2 (down): row-parallel → shard the K dimension (dim 2,
|
|
# which is the packed K//2 dim).
|
|
# Scales follow the same sharding on the corresponding dimension.
|
|
tp = tensor_parallel_size
|
|
if tp > 1:
|
|
n1 = w1.size(1) // tp
|
|
w1 = w1[:, :n1, :].contiguous()
|
|
w1_scale = w1_scale[:, :n1, :].contiguous()
|
|
|
|
k2_packed = w2.size(2) // tp
|
|
k2_scale = w2_scale.size(2) // tp
|
|
w2 = w2[:, :, :k2_packed].contiguous()
|
|
w2_scale = w2_scale[:, :, :k2_scale].contiguous()
|
|
|
|
hidden_dim = w1.size(2) * 2
|
|
intermediate_size = w1.size(1) // 2
|
|
|
|
return (
|
|
w1,
|
|
w1_scale,
|
|
w1_gscale,
|
|
w2,
|
|
w2_scale,
|
|
w2_gscale,
|
|
a1_gscale,
|
|
a2_gscale,
|
|
num_experts,
|
|
hidden_dim,
|
|
intermediate_size,
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda_alike(),
|
|
reason="Triton NVFP4 kernel requires CUDA.",
|
|
)
|
|
@pytest.mark.parametrize("num_tokens", [1, 2, 4, 1024])
|
|
@pytest.mark.parametrize("top_k", [4])
|
|
@pytest.mark.parametrize("model_id", list(MOE_MODEL_CONFIGS.keys()))
|
|
@pytest.mark.parametrize(
|
|
"tensor_parallel_size",
|
|
[pytest.param(val, id=f"tensor_parallel_size:{val}") for val in [1, 2, 4, 8]],
|
|
)
|
|
def test_nvfp4_moe_correctness(
|
|
loaded_model_files,
|
|
num_tokens: int,
|
|
top_k: int,
|
|
model_id: str,
|
|
tensor_parallel_size: int,
|
|
) -> None:
|
|
"""Compare Nvfp4QuantizationEmulationTritonExperts (fused weight dequant + compute)
|
|
against the unfused reference Nvfp4QuantizationEmulationTritonExpertsReference.
|
|
|
|
Both must produce bit-identical results.
|
|
"""
|
|
num_test_experts = max(8, top_k)
|
|
(
|
|
w1,
|
|
w1_scale,
|
|
w1_gscale,
|
|
w2,
|
|
w2_scale,
|
|
w2_gscale,
|
|
a1_gscale,
|
|
a2_gscale,
|
|
num_experts,
|
|
hidden_dim,
|
|
intermediate_size,
|
|
) = _load_nvfp4_moe_weights(
|
|
loaded_model_files,
|
|
model_id,
|
|
tensor_parallel_size,
|
|
max_experts=num_test_experts,
|
|
)
|
|
|
|
moe_config = FusedMoEConfig(
|
|
num_experts=num_experts,
|
|
experts_per_token=top_k,
|
|
hidden_dim=hidden_dim,
|
|
intermediate_size=intermediate_size,
|
|
num_local_experts=num_experts,
|
|
num_logical_experts=num_experts,
|
|
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
|
|
activation=MoEActivation.SILU,
|
|
in_dtype=torch.bfloat16,
|
|
device="cuda",
|
|
routing_method=RoutingMethodType.TopK,
|
|
max_num_tokens=512,
|
|
)
|
|
|
|
def _make_quant_config():
|
|
return nvfp4_moe_quant_config(
|
|
g1_alphas=w1_gscale.clone(),
|
|
g2_alphas=w2_gscale.clone(),
|
|
a1_gscale=a1_gscale.clone(),
|
|
a2_gscale=a2_gscale.clone(),
|
|
w1_scale=w1_scale.clone(),
|
|
w2_scale=w2_scale.clone(),
|
|
)
|
|
|
|
ref_experts = Nvfp4QuantizationEmulationTritonExpertsReference(
|
|
moe_config=moe_config,
|
|
quant_config=_make_quant_config(),
|
|
)
|
|
fused_experts = Nvfp4QuantizationEmulationTritonExperts(
|
|
moe_config=moe_config,
|
|
quant_config=_make_quant_config(),
|
|
)
|
|
|
|
torch.manual_seed(42)
|
|
hidden_states = torch.randn(
|
|
num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda"
|
|
)
|
|
|
|
topk_weights = torch.randn(
|
|
num_tokens, top_k, dtype=torch.float32, device="cuda"
|
|
).softmax(dim=-1)
|
|
topk_ids = torch.stack(
|
|
[torch.randperm(num_experts, device="cuda")[:top_k] for _ in range(num_tokens)]
|
|
).to(torch.int32)
|
|
|
|
N = w1.size(1) # 2 * intermediate
|
|
K = hidden_dim
|
|
|
|
ws13_size = num_tokens * top_k * max(intermediate_size, K)
|
|
ws2_size = num_tokens * top_k * max(N, K)
|
|
|
|
workspace13_ref = torch.zeros(ws13_size, dtype=torch.bfloat16, device="cuda")
|
|
workspace2_ref = torch.zeros(ws2_size, dtype=torch.bfloat16, device="cuda")
|
|
output_ref = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device="cuda")
|
|
|
|
workspace13_fused = torch.zeros_like(workspace13_ref)
|
|
workspace2_fused = torch.zeros_like(workspace2_ref)
|
|
output_fused = torch.zeros_like(output_ref)
|
|
|
|
apply_kwargs = dict(
|
|
hidden_states=hidden_states,
|
|
w1=w1,
|
|
w2=w2,
|
|
topk_weights=topk_weights,
|
|
topk_ids=topk_ids,
|
|
activation=MoEActivation.SILU,
|
|
global_num_experts=num_experts,
|
|
expert_map=None,
|
|
a1q_scale=None,
|
|
a2_scale=None,
|
|
expert_tokens_meta=None,
|
|
apply_router_weight_on_input=False,
|
|
)
|
|
|
|
# Unfused reference.
|
|
ref_experts.apply(
|
|
output=output_ref,
|
|
workspace13=workspace13_ref,
|
|
workspace2=workspace2_ref,
|
|
**apply_kwargs,
|
|
)
|
|
|
|
# Fused implementation.
|
|
fused_experts.apply(
|
|
output=output_fused,
|
|
workspace13=workspace13_fused,
|
|
workspace2=workspace2_fused,
|
|
**apply_kwargs,
|
|
)
|
|
|
|
# Not strict equality on H100, MI325, MI300 (< 0.1% elements).
|
|
# The fused on-the-fly dequant path can lower to a slightly
|
|
# different Triton/MMA tiling than the pre-dequantized
|
|
# reference; experiments with reference-like tiling/masking
|
|
# reduced some diffs were not kept because they regress
|
|
# the fused kernel speed.
|
|
# Strict equality validated on MI355.
|
|
torch.testing.assert_close(
|
|
output_fused,
|
|
output_ref,
|
|
atol=0.0 if on_gfx950() else 0.02,
|
|
rtol=0,
|
|
)
|