Feat/add w4a16 moe support to nemotron (#25655)

2026-06-02 22:42:26 -07:00
parent 512bfbb1e1
commit b8d7351a74
19 changed files with 999 additions and 61 deletions
@@ -1,13 +1,27 @@
 import sys
+from types import SimpleNamespace

 import pytest
 import torch
 from sgl_kernel.scalar_type import scalar_types

 from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm
-from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace
+from sglang.srt.layers.quantization.marlin_utils import (
+    check_marlin_supported,
+    marlin_make_workspace,
+)
+from sglang.srt.layers.quantization.marlin_utils_fp4 import (
+    apply_fp4_marlin_linear,
+    nvfp4_marlin_process_global_scale,
+    prepare_nvfp4_layer_for_marlin,
+)
+from sglang.srt.utils.common import is_sm80_supported, is_sm90_supported
 from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.test_marlin_utils import awq_marlin_quantize, marlin_quantize
+from sglang.test.test_marlin_utils import (
+    awq_marlin_quantize,
+    make_nvfp4_weight_and_ref,
+    marlin_quantize,
+)

 register_cuda_ci(est_time=13, suite="base-b-kernel-unit-1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
@@ -101,5 +115,80 @@ def test_gptq_marlin_gemm(
    assert max_diff < 0.04


+@pytest.mark.skip(reason="Skip, test pass locally but compiling takes too long in CI")
+@pytest.mark.skipif(
+    not (is_sm80_supported() or is_sm90_supported()),
+    reason="NVFP4 Marlin fallback tests require CUDA SM8X/SM9X",
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_nvfp4_marlin_support_and_scale_transforms_sm80_sm90(dtype):
+    major, minor = torch.cuda.get_device_capability()
+    capability = major * 10 + minor
+    assert check_marlin_supported(
+        scalar_types.float4_e2m1f,
+        group_size=16,
+        has_zp=False,
+        device_capability=capability,
+    )
+
+    global_scale = torch.tensor(1.0, dtype=dtype, device="cuda")
+    actual_global_scale = nvfp4_marlin_process_global_scale(global_scale)
+    assert actual_global_scale.is_cuda
+    assert actual_global_scale.ndim == 1
+    assert actual_global_scale.numel() == 1
+    if dtype == torch.float16:
+        assert actual_global_scale.item() == 128.0
+    else:
+        assert actual_global_scale.item() == 2.0**119
+
+
+@pytest.mark.skip(reason="Skip, test pass locally but compiling takes too long in CI")
+@pytest.mark.skipif(
+    not (is_sm80_supported() or is_sm90_supported()),
+    reason="NVFP4 Marlin dense numeric test requires CUDA SM80, SM86, or SM90",
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_nvfp4_marlin_dense_matches_dequant_reference(dtype):
+    torch.manual_seed(0)
+
+    size_m = 17
+    size_k = 256
+    size_n = 192
+    group_size = 16
+
+    a_input = torch.randn((size_m, size_k), dtype=dtype, device="cuda") / 10
+    fp4_weight, scales, global_scale, weight_ref = make_nvfp4_weight_and_ref(
+        size_n, size_k, dtype, group_size=group_size
+    )
+
+    layer = torch.nn.Module()
+    layer.quant_config = SimpleNamespace(group_size=group_size)
+    layer.output_size_per_partition = size_n
+    layer.input_size_per_partition = size_k
+    layer.params_dtype = dtype
+    layer.weight = torch.nn.Parameter(fp4_weight, requires_grad=False)
+    layer.weight_scale = torch.nn.Parameter(scales, requires_grad=False)
+    layer.weight_global_scale = torch.nn.Parameter(
+        global_scale.reshape(1), requires_grad=False
+    )
+    prepare_nvfp4_layer_for_marlin(layer)
+
+    output = apply_fp4_marlin_linear(
+        a_input,
+        layer.weight,
+        layer.weight_scale,
+        layer.weight_global_scale,
+        layer.workspace,
+        size_n,
+        size_k,
+        use_fp32_reduce=True,
+    )
+
+    output_ref = torch.matmul(a_input, weight_ref.T)
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(output, output_ref, rtol=0.04, atol=0.04)
+
+
 if __name__ == "__main__":
    sys.exit(pytest.main([__file__, "-v", "-s"]))
@@ -1,5 +1,6 @@
 import itertools
 import sys
+from types import SimpleNamespace

 import pytest
 import torch
@@ -7,8 +8,17 @@ from sgl_kernel.scalar_type import scalar_types

 from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm
 from sglang.srt.layers.moe.fused_moe_triton import moe_align_block_size
+from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import fused_marlin_moe
+from sglang.srt.layers.quantization.marlin_utils_fp4 import (
+    prepare_moe_nvfp4_layer_for_marlin,
+)
+from sglang.srt.utils.common import is_sm80_supported, is_sm90_supported
 from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.test_marlin_utils import awq_marlin_quantize, marlin_quantize
+from sglang.test.test_marlin_utils import (
+    awq_marlin_quantize,
+    make_nvfp4_weight_and_ref,
+    marlin_quantize,
+)

 register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
@@ -339,5 +349,267 @@ def test_moe_wna16_marlin_gemm(
    torch.testing.assert_close(c_jit, c_aot, rtol=0, atol=0)


+@pytest.mark.skip(reason="Skip, test pass locally but compiling takes too long in CI")
+@pytest.mark.skipif(
+    not (is_sm80_supported() or is_sm90_supported()),
+    reason="Non-gated NVFP4 Marlin fallback test requires CUDA SM8X/SM9X",
+)
+def test_fused_marlin_moe_non_gated_relu2():
+    torch.manual_seed(0)
+
+    m = 17
+    n = 128
+    k = 256
+    e = 4
+    topk = 2
+    dtype = torch.float16
+    group_size = 128
+    quant_type = scalar_types.uint4b8
+
+    hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10
+    w_ref1, qweight1, scales1, zeros1, g_idx1, sort_indices1 = _setup_moe_weights(
+        e, n, k, quant_type, group_size, False, dtype
+    )
+    w_ref2, qweight2, scales2, zeros2, g_idx2, sort_indices2 = _setup_moe_weights(
+        e, k, n, quant_type, group_size, False, dtype
+    )
+
+    router_logits = torch.randn((m, e), device="cuda", dtype=dtype)
+    score_softmax = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
+    topk_weights, topk_ids = torch.topk(score_softmax, topk)
+
+    output = fused_marlin_moe(
+        hidden_states=hidden_states,
+        w1=qweight1,
+        w2=qweight2,
+        w1_scale=scales1,
+        w2_scale=scales2,
+        gating_output=router_logits,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        g_idx1=g_idx1,
+        g_idx2=g_idx2,
+        sort_indices1=sort_indices1,
+        sort_indices2=sort_indices2,
+        w1_zeros=zeros1,
+        w2_zeros=zeros2,
+        num_bits=4,
+        is_k_full=True,
+        routed_scaling_factor=1.0,
+        activation="relu2",
+        is_gated=False,
+    )
+
+    output_ref = torch.zeros_like(hidden_states)
+    for token_idx in range(m):
+        for route_idx in range(topk):
+            expert_id = topk_ids[token_idx, route_idx]
+            intermediate = hidden_states[token_idx] @ w_ref1[expert_id].T
+            intermediate = torch.square(torch.relu(intermediate))
+            routed = intermediate @ w_ref2[expert_id].T
+            output_ref[token_idx] += routed * topk_weights[token_idx, route_idx]
+
+    torch.cuda.synchronize()
+    torch.testing.assert_close(output, output_ref, rtol=0.04, atol=0.04)
+
+
+@pytest.mark.skip(reason="Skip, test pass locally but compiling takes too long in CI")
+@pytest.mark.skipif(
+    not (is_sm80_supported() or is_sm90_supported()),
+    reason="NVFP4 Marlin MoE padding test requires CUDA SM8X/SM9X",
+)
+def test_fused_marlin_moe_nvfp4_non_gated_padded_intermediate_launches():
+    torch.manual_seed(0)
+
+    m = 17
+    intermediate_size = 192
+    hidden_size = 256
+    e = 4
+    topk = 2
+    dtype = torch.bfloat16
+    nvfp4_group_size = 16
+
+    layer = torch.nn.Module()
+    layer.quant_config = SimpleNamespace(group_size=nvfp4_group_size)
+    layer.moe_runner_config = SimpleNamespace(is_gated=False)
+    layer.params_dtype = dtype
+    layer.intermediate_size_per_partition = intermediate_size
+    layer.w13_weight = torch.nn.Parameter(
+        torch.randint(
+            0,
+            256,
+            (e, intermediate_size, hidden_size // 2),
+            device="cuda",
+            dtype=torch.uint8,
+        ),
+        requires_grad=False,
+    )
+    layer.w2_weight = torch.nn.Parameter(
+        torch.randint(
+            0,
+            256,
+            (e, hidden_size, intermediate_size // 2),
+            device="cuda",
+            dtype=torch.uint8,
+        ),
+        requires_grad=False,
+    )
+    layer.w13_weight_scale = torch.nn.Parameter(
+        torch.rand(
+            (e, intermediate_size, hidden_size // nvfp4_group_size),
+            device="cuda",
+            dtype=dtype,
+        ),
+        requires_grad=False,
+    )
+    layer.w2_weight_scale = torch.nn.Parameter(
+        torch.rand(
+            (e, hidden_size, intermediate_size // nvfp4_group_size),
+            device="cuda",
+            dtype=dtype,
+        ),
+        requires_grad=False,
+    )
+    layer.w13_weight_scale_2 = torch.nn.Parameter(
+        torch.ones((e,), device="cuda", dtype=dtype), requires_grad=False
+    )
+    layer.w2_weight_scale_2 = torch.nn.Parameter(
+        torch.ones((e,), device="cuda", dtype=dtype), requires_grad=False
+    )
+    prepare_moe_nvfp4_layer_for_marlin(layer)
+
+    assert layer.w13_weight.shape[1] * 16 == 256
+    assert layer.w2_weight.shape[1] * 16 == 256
+
+    hidden_states = torch.randn((m, hidden_size), device="cuda", dtype=dtype) / 10
+
+    score = torch.randn((m, e), device="cuda", dtype=dtype)
+    score_softmax = torch.softmax(score, dim=-1, dtype=torch.float32)
+    topk_weights, topk_ids = torch.topk(score_softmax, topk)
+
+    out = fused_marlin_moe(
+        hidden_states=hidden_states,
+        w1=layer.w13_weight,
+        w2=layer.w2_weight,
+        w1_scale=layer.w13_weight_scale,
+        w2_scale=layer.w2_weight_scale,
+        gating_output=score,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        w1_global_scale=layer.w13_weight_scale_2,
+        w2_global_scale=layer.w2_weight_scale_2,
+        workspace=layer.workspace,
+        num_bits=4,
+        is_k_full=True,
+        routed_scaling_factor=1.0,
+        activation="relu2",
+        is_gated=False,
+    )
+
+    torch.cuda.synchronize()
+    assert out.shape == (m, hidden_size)
+
+
+@pytest.mark.skip(reason="Skip, test pass locally but compiling takes too long in CI")
+@pytest.mark.skipif(
+    not (is_sm80_supported() or is_sm90_supported()),
+    reason="NVFP4 Marlin MoE numeric test requires CUDA SM80, SM86, or SM90",
+)
+def test_fused_marlin_moe_nvfp4_non_gated_matches_dequant_reference():
+    torch.manual_seed(0)
+
+    m = 17
+    intermediate_size = 192
+    hidden_size = 256
+    e = 4
+    topk = 2
+    dtype = torch.bfloat16
+    group_size = 16
+    routed_scaling_factor = 1.0
+
+    w13_packed_l, w13_scales_l, w13_gscale_l, w13_ref_l = [], [], [], []
+    w2_packed_l, w2_scales_l, w2_gscale_l, w2_ref_l = [], [], [], []
+    for _ in range(e):
+        packed, scales, gscale, ref = make_nvfp4_weight_and_ref(
+            intermediate_size, hidden_size, dtype, group_size=group_size
+        )
+        w13_packed_l.append(packed)
+        w13_scales_l.append(scales)
+        w13_gscale_l.append(gscale)
+        w13_ref_l.append(ref)
+
+        packed, scales, gscale, ref = make_nvfp4_weight_and_ref(
+            hidden_size, intermediate_size, dtype, group_size=group_size
+        )
+        w2_packed_l.append(packed)
+        w2_scales_l.append(scales)
+        w2_gscale_l.append(gscale)
+        w2_ref_l.append(ref)
+
+    layer = torch.nn.Module()
+    layer.quant_config = SimpleNamespace(group_size=group_size)
+    layer.moe_runner_config = SimpleNamespace(is_gated=False)
+    layer.params_dtype = dtype
+    layer.intermediate_size_per_partition = intermediate_size
+    layer.w13_weight = torch.nn.Parameter(
+        torch.stack(w13_packed_l), requires_grad=False
+    )
+    layer.w2_weight = torch.nn.Parameter(torch.stack(w2_packed_l), requires_grad=False)
+    layer.w13_weight_scale = torch.nn.Parameter(
+        torch.stack(w13_scales_l), requires_grad=False
+    )
+    layer.w2_weight_scale = torch.nn.Parameter(
+        torch.stack(w2_scales_l), requires_grad=False
+    )
+    layer.w13_weight_scale_2 = torch.nn.Parameter(
+        torch.stack(w13_gscale_l), requires_grad=False
+    )
+    layer.w2_weight_scale_2 = torch.nn.Parameter(
+        torch.stack(w2_gscale_l), requires_grad=False
+    )
+    prepare_moe_nvfp4_layer_for_marlin(layer)
+
+    # Scale activations down so relu² doesn't blow up intermediate magnitudes;
+    # this keeps output values small so tighter element-wise tolerance is realistic.
+    hidden_states = torch.randn((m, hidden_size), device="cuda", dtype=dtype) / 20
+    router_logits = torch.randn((m, e), device="cuda", dtype=dtype)
+    score_softmax = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
+    topk_weights, topk_ids = torch.topk(score_softmax, topk)
+
+    output = fused_marlin_moe(
+        hidden_states=hidden_states,
+        w1=layer.w13_weight,
+        w2=layer.w2_weight,
+        w1_scale=layer.w13_weight_scale,
+        w2_scale=layer.w2_weight_scale,
+        gating_output=router_logits,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        w1_global_scale=layer.w13_weight_scale_2,
+        w2_global_scale=layer.w2_weight_scale_2,
+        workspace=layer.workspace,
+        num_bits=4,
+        is_k_full=True,
+        routed_scaling_factor=routed_scaling_factor,
+        activation="relu2",
+        is_gated=False,
+    )
+
+    w13_ref = torch.stack(w13_ref_l)
+    w2_ref = torch.stack(w2_ref_l)
+    output_ref = torch.zeros_like(hidden_states)
+    for token_idx in range(m):
+        for route_idx in range(topk):
+            expert_id = topk_ids[token_idx, route_idx]
+            intermediate = hidden_states[token_idx] @ w13_ref[expert_id].T
+            intermediate = torch.square(torch.relu(intermediate))
+            routed = intermediate @ w2_ref[expert_id].T
+            output_ref[token_idx] += routed * topk_weights[token_idx, route_idx]
+    output_ref *= routed_scaling_factor
+
+    torch.cuda.synchronize()
+    torch.testing.assert_close(output, output_ref, rtol=0.05, atol=0.25)
+
+
 if __name__ == "__main__":
    sys.exit(pytest.main([__file__, "-v", "-s"]))