Fuse Nemotron latent MoE projection and shared add (#30430)
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com> Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai> Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
This commit is contained in:
Samuel Nordmann
co-authored by
Po-Han Huang
Mohammad Angkad
Mohammad Miadh Angkad
Po-Han Huang
Mohammad Angkad
Mohammad Miadh Angkad
parent
1d5d85260c
commit
214313ee79
@@ -0,0 +1,205 @@
|
||||
"""
|
||||
Unit tests for UnquantizedLinearMethod.apply_with_addend.
|
||||
|
||||
The cuBLAS route folds the addend into the GEMM beta input,
|
||||
writing back into that buffer; every other route must add separately
|
||||
and leave the caller's buffer intact.
|
||||
"""
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=25, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
import unittest
|
||||
from contextlib import ExitStack
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
import sglang.srt.layers.quantization.unquant as unquant
|
||||
from sglang.srt.layers.linear import ReplicatedLinear
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
# addmm rounds once in the accumulator where the separate add rounds twice;
|
||||
# with 8 BF16 mantissa bits the paths agree only to this precision.
|
||||
_BF16_RTOL = 1e-2
|
||||
_BF16_ATOL = 3.125e-2
|
||||
|
||||
|
||||
def _cublas_only_backend():
|
||||
# The TORCH backend leaves every custom-kernel global unpopulated.
|
||||
return patch.object(unquant, "_BF16_GEMM_BACKEND", unquant.Bf16GemmBackend.TORCH)
|
||||
|
||||
|
||||
def _fake_kernel(calls, name):
|
||||
def kernel(x, weight, bias=None, *args):
|
||||
calls.append(name)
|
||||
return F.linear(x, weight, bias)
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
|
||||
class TestApplyWithAddend(CustomTestCase):
|
||||
def setUp(self):
|
||||
torch.manual_seed(0)
|
||||
self.projection = ReplicatedLinear(
|
||||
64, 128, bias=False, params_dtype=torch.bfloat16
|
||||
).cuda()
|
||||
self.projection.weight.copy_(
|
||||
torch.randn_like(self.projection.weight) / 8.0 # 1 / sqrt(64)
|
||||
)
|
||||
self.method = self.projection.quant_method
|
||||
|
||||
def _reference(self, x, addend, bias=None):
|
||||
return F.linear(x, self.projection.weight, bias) + addend
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_cublas_route_accumulates_into_addend(self):
|
||||
with _cublas_only_backend():
|
||||
x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
|
||||
addend = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16)
|
||||
reference = self._reference(x, addend)
|
||||
|
||||
output = self.method.apply_with_addend(self.projection, x, addend)
|
||||
|
||||
self.assertEqual(output.data_ptr(), addend.data_ptr())
|
||||
torch.testing.assert_close(
|
||||
output, reference, rtol=_BF16_RTOL, atol=_BF16_ATOL
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_cuda_graph_replay_reads_the_replayed_addend(self):
|
||||
"""A graph replay must consume the addend written on that replay;
|
||||
addmm(out=addend) reads and writes the one buffer."""
|
||||
with _cublas_only_backend():
|
||||
x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
|
||||
produced = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
# Initialize the cuBLAS workspace before capture.
|
||||
self.method.apply_with_addend(self.projection, x, produced.clone())
|
||||
torch.cuda.synchronize()
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
# clone() stands in for the shared expert,
|
||||
# which rewrites its output buffer on every replay.
|
||||
captured = self.method.apply_with_addend(
|
||||
self.projection, x, produced.clone()
|
||||
)
|
||||
|
||||
x.copy_(torch.randn_like(x))
|
||||
produced.copy_(torch.randn_like(produced))
|
||||
reference = self._reference(x, produced)
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
captured, reference, rtol=_BF16_RTOL, atol=_BF16_ATOL
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_batch_invariant_mode_keeps_separate_add(self):
|
||||
"""Deterministic inference must not reach the fused route;
|
||||
batch-invariant mode does not override aten::addmm.out."""
|
||||
with (
|
||||
_cublas_only_backend(),
|
||||
patch.object(unquant, "is_batch_invariant_mode_enabled", return_value=True),
|
||||
):
|
||||
x = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16)
|
||||
addend = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16)
|
||||
before = addend.clone()
|
||||
|
||||
output = self.method.apply_with_addend(self.projection, x, addend)
|
||||
|
||||
self.assertNotEqual(output.data_ptr(), addend.data_ptr())
|
||||
torch.testing.assert_close(addend, before, rtol=0, atol=0)
|
||||
torch.testing.assert_close(
|
||||
output, self._reference(x, before), rtol=_BF16_RTOL, atol=_BF16_ATOL
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_unfused_routes_leave_the_addend_intact(self):
|
||||
"""Every route that cannot use the GEMM beta input adds separately;
|
||||
consuming the addend there corrupts the caller's buffer."""
|
||||
for route in (
|
||||
"non_nvidia",
|
||||
"compiling",
|
||||
"cutedsl",
|
||||
"splitk",
|
||||
"hopper_gemv",
|
||||
"noncontiguous_addend",
|
||||
"bias",
|
||||
):
|
||||
with self.subTest(route=route):
|
||||
self._assert_route_is_unfused(route)
|
||||
|
||||
def _assert_route_is_unfused(self, route: str):
|
||||
x = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16)
|
||||
addend = (
|
||||
torch.randn(128, 4, device="cuda", dtype=torch.bfloat16).t()
|
||||
if route == "noncontiguous_addend"
|
||||
else torch.randn(4, 128, device="cuda", dtype=torch.bfloat16)
|
||||
)
|
||||
bias = (
|
||||
torch.randn(128, device="cuda", dtype=torch.bfloat16)
|
||||
if route == "bias"
|
||||
else None
|
||||
)
|
||||
before = addend.clone()
|
||||
kernel_calls = []
|
||||
|
||||
with ExitStack() as stack:
|
||||
enter = stack.enter_context
|
||||
enter(_cublas_only_backend())
|
||||
if route == "non_nvidia":
|
||||
enter(patch.object(unquant, "_is_cuda", False))
|
||||
elif route == "compiling":
|
||||
enter(patch.object(torch.compiler, "is_compiling", return_value=True))
|
||||
elif route == "cutedsl":
|
||||
enter(patch.object(unquant, "_use_cutedsl_bf16_gemm", lambda *a: True))
|
||||
enter(
|
||||
patch.object(
|
||||
unquant,
|
||||
"_cutedsl_bf16_gemm",
|
||||
_fake_kernel(kernel_calls, route),
|
||||
)
|
||||
)
|
||||
elif route == "splitk":
|
||||
enter(patch.object(unquant, "_enable_bf16_splitk_gemm", True))
|
||||
enter(
|
||||
patch.object(
|
||||
unquant, "use_flashinfer_pr4266_bf16_gemm", lambda *a: True
|
||||
)
|
||||
)
|
||||
enter(
|
||||
patch.object(
|
||||
unquant,
|
||||
"_flashinfer_pr4266_bf16_gemm",
|
||||
_fake_kernel(kernel_calls, route),
|
||||
)
|
||||
)
|
||||
elif route == "hopper_gemv":
|
||||
enter(patch.object(unquant, "_use_hopper_bf16_gemv", lambda *a: True))
|
||||
enter(
|
||||
patch.object(
|
||||
unquant,
|
||||
"_hopper_bf16_gemv",
|
||||
_fake_kernel(kernel_calls, route),
|
||||
)
|
||||
)
|
||||
|
||||
output = self.method.apply_with_addend(self.projection, x, addend, bias)
|
||||
|
||||
self.assertNotEqual(output.data_ptr(), addend.data_ptr())
|
||||
torch.testing.assert_close(addend, before, rtol=0, atol=0)
|
||||
torch.testing.assert_close(
|
||||
output, self._reference(x, before, bias), rtol=_BF16_RTOL, atol=_BF16_ATOL
|
||||
)
|
||||
if route in ("cutedsl", "splitk", "hopper_gemv"):
|
||||
self.assertEqual(kernel_calls, [route])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,127 @@
|
||||
"""
|
||||
Unit tests for the NemotronHMoE latent-projection shared-expert add.
|
||||
|
||||
The fused path calls the projection's quant method directly,
|
||||
not ``ReplicatedLinear.forward``.
|
||||
These cases pin the gate that decides when the substitution is safe.
|
||||
"""
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=20, suite="base-a-test-cpu")
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.srt.layers.linear import ReplicatedLinear
|
||||
from sglang.srt.lora.layers import ReplicatedLinearWithLoRA
|
||||
from sglang.srt.models.nemotron_h import (
|
||||
NemotronHMoE,
|
||||
_latent_proj_fuses_shared_add,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
# Stands in for a quantized linear method, which has no addend entry point.
|
||||
class _DoubleMethod:
|
||||
def apply(self, layer, x, bias):
|
||||
return 2 * x
|
||||
|
||||
|
||||
class _FakeLoRABackend:
|
||||
batch_info = object()
|
||||
skip_inactive_lora_batches = False
|
||||
|
||||
def run_lora_a_sgemm(self, x, weights):
|
||||
return 2 * x
|
||||
|
||||
def run_lora_b_sgemm(self, *, x, weights, output_offset, base_output):
|
||||
return base_output + x
|
||||
|
||||
|
||||
def _apply(projection, routed, shared):
|
||||
moe = SimpleNamespace(fc2_latent_proj=projection)
|
||||
return NemotronHMoE._apply_latent_projection(moe, routed, shared)
|
||||
|
||||
|
||||
class TestNemotronHSharedAdd(CustomTestCase):
|
||||
def setUp(self):
|
||||
torch.manual_seed(0)
|
||||
self.routed = torch.randn(4, 8)
|
||||
self.shared = torch.randn(4, 8)
|
||||
|
||||
def test_gate_accepts_plain_bias_free_projection(self):
|
||||
"""A bias-free unquantized ReplicatedLinear must stay eligible;
|
||||
a gate that degrades to always-false silently drops the fusion."""
|
||||
projection = ReplicatedLinear(8, 8, bias=False)
|
||||
|
||||
self.assertTrue(_latent_proj_fuses_shared_add(projection))
|
||||
|
||||
def test_gate_rejects_lora_wrapper(self):
|
||||
"""LoRA swaps a wrapper module over fc2_latent_proj after model init.
|
||||
Calling the base quant method there would drop the adapter update."""
|
||||
projection = ReplicatedLinear(8, 8, bias=False)
|
||||
with torch.no_grad():
|
||||
projection.weight.zero_()
|
||||
wrapped = ReplicatedLinearWithLoRA(projection, _FakeLoRABackend())
|
||||
wrapped.set_lora_info(torch.empty(1, 8), torch.empty(8, 1))
|
||||
|
||||
self.assertFalse(_latent_proj_fuses_shared_add(wrapped))
|
||||
torch.testing.assert_close(
|
||||
_apply(wrapped, self.routed, self.shared),
|
||||
2 * self.routed + self.shared,
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
|
||||
def test_gate_rejects_quantized_projection(self):
|
||||
"""Only UnquantizedLinearMethod implements apply_with_addend."""
|
||||
projection = ReplicatedLinear(8, 8, bias=False)
|
||||
projection.quant_method = _DoubleMethod()
|
||||
|
||||
self.assertFalse(_latent_proj_fuses_shared_add(projection))
|
||||
torch.testing.assert_close(
|
||||
_apply(projection, self.routed, self.shared),
|
||||
2 * self.routed + self.shared,
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
|
||||
def test_gate_rejects_projection_with_bias(self):
|
||||
"""forward defers the bias under skip_bias_add and runs module hooks;
|
||||
both are lost if a bias-bearing projection takes the fused call."""
|
||||
projection = ReplicatedLinear(8, 8, bias=True, skip_bias_add=True)
|
||||
with torch.no_grad():
|
||||
projection.weight.copy_(torch.randn_like(projection.weight))
|
||||
projection.bias.copy_(torch.randn_like(projection.bias))
|
||||
hook_calls = []
|
||||
projection.register_forward_hook(lambda *args: hook_calls.append(True))
|
||||
|
||||
self.assertFalse(_latent_proj_fuses_shared_add(projection))
|
||||
torch.testing.assert_close(
|
||||
_apply(projection, self.routed, self.shared),
|
||||
F.linear(self.routed, projection.weight) + self.shared,
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
self.assertEqual(hook_calls, [True])
|
||||
|
||||
def test_no_shared_output_keeps_plain_projection(self):
|
||||
"""Layers without shared experts pass shared_output=None."""
|
||||
projection = ReplicatedLinear(8, 8, bias=False)
|
||||
with torch.no_grad():
|
||||
projection.weight.copy_(torch.randn_like(projection.weight))
|
||||
|
||||
torch.testing.assert_close(
|
||||
_apply(projection, self.routed, None),
|
||||
F.linear(self.routed, projection.weight),
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user