diff --git a/python/sglang/jit_kernel/csrc/elementwise/fused_eh_norm.cuh b/python/sglang/jit_kernel/csrc/elementwise/fused_eh_norm.cuh new file mode 100644 index 000000000..5e3970642 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/elementwise/fused_eh_norm.cuh @@ -0,0 +1,113 @@ +#include +#include + +#include +#include +#include +#include + +#include + +#include + +namespace { + +struct FusedEHNormParams { + const void* __restrict__ embeds; + const void* __restrict__ previous_hidden; + const void* __restrict__ enorm_weight; + const void* __restrict__ hnorm_weight; + void* __restrict__ output; + int64_t embeds_stride; + int64_t previous_hidden_stride; + int64_t output_stride; + float eps; +}; + +template +__global__ void fused_eh_norm_kernel(const __grid_constant__ FusedEHNormParams params) { + using namespace device; + using Storage = norm::StorageType; + + constexpr auto kNumThreads = host::norm::get_cta_threads(); + constexpr auto kNumWarps = kNumThreads / kWarpThreads; + + const auto embeds = static_cast(pointer::offset(params.embeds, blockIdx.x * params.embeds_stride)); + const auto previous_hidden = + static_cast(pointer::offset(params.previous_hidden, blockIdx.x * params.previous_hidden_stride)); + const auto enorm_weight = static_cast(params.enorm_weight); + const auto hnorm_weight = static_cast(params.hnorm_weight); + const auto output = static_cast(pointer::offset(params.output, blockIdx.x * params.output_stride)); + + const auto gmem = tile::Memory::cta(kNumThreads); + __shared__ float smem[norm::kSmemBufferSize]; + + PDLWaitPrimary(); + + const auto embeds_vec = gmem.load(embeds); + const auto enorm_weight_vec = gmem.load(enorm_weight); + const auto embeds_output_vec = + norm::apply_norm_cta(embeds_vec, enorm_weight_vec, params.eps, smem, kNumWarps); + gmem.store(output, embeds_output_vec); + + const auto prev_vec = gmem.load(previous_hidden); + const auto hnorm_weight_vec = gmem.load(hnorm_weight); + const auto prev_output_vec = norm::apply_norm_cta(prev_vec, hnorm_weight_vec, params.eps, smem, kNumWarps); + gmem.store(output + kHidden, prev_output_vec); + + PDLTriggerSecondary(); +} + +template +struct FusedEHNormKernel { + static_assert(host::norm::is_config_supported(), "Unsupported norm configuration"); + static_assert(host::norm::should_use_cta(), "fused_eh_norm requires CTA norm"); + static constexpr auto kernel = fused_eh_norm_kernel; + static constexpr uint32_t kBlockSize = host::norm::get_cta_threads(); + + static void + run(const tvm::ffi::TensorView embeds, + const tvm::ffi::TensorView previous_hidden, + const tvm::ffi::TensorView enorm_weight, + const tvm::ffi::TensorView hnorm_weight, + const tvm::ffi::TensorView output, + float eps) { + using namespace host; + + auto N = SymbolicSize{"num_tokens"}; + auto H = SymbolicSize{"hidden_size"}; + auto H2 = SymbolicSize{"hidden_size_times_2"}; + auto SE = SymbolicSize{"embeds_stride"}; + auto SP = SymbolicSize{"previous_hidden_stride"}; + auto SO = SymbolicSize{"output_stride"}; + auto device_ = SymbolicDevice{}; + H.set_value(kHidden); + H2.set_value(kHidden * 2); + device_.set_options(); + + TensorMatcher({N, H}).with_strides({SE, 1}).with_dtype().with_device(device_).verify(embeds); + TensorMatcher({N, H}).with_strides({SP, 1}).with_dtype().with_device(device_).verify(previous_hidden); + TensorMatcher({H}).with_dtype().with_device(device_).verify(enorm_weight); + TensorMatcher({H}).with_dtype().with_device(device_).verify(hnorm_weight); + TensorMatcher({N, H2}).with_strides({SO, 1}).with_dtype().with_device(device_).verify(output); + + const auto num_tokens = static_cast(N.unwrap()); + const auto params = FusedEHNormParams{ + .embeds = embeds.data_ptr(), + .previous_hidden = previous_hidden.data_ptr(), + .enorm_weight = enorm_weight.data_ptr(), + .hnorm_weight = hnorm_weight.data_ptr(), + .output = output.data_ptr(), + .embeds_stride = SE.unwrap(), + .previous_hidden_stride = SP.unwrap(), + .output_stride = SO.unwrap(), + .eps = eps, + }; + + const auto num_blocks = num_tokens; + LaunchKernel(num_blocks, kBlockSize, device_.unwrap()) // + .enable_pdl(kUsePDL)(kernel, params); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/fused_eh_norm.py b/python/sglang/jit_kernel/fused_eh_norm.py new file mode 100644 index 000000000..d7b7c0817 --- /dev/null +++ b/python/sglang/jit_kernel/fused_eh_norm.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +def is_supported_fused_eh_norm_hidden_size(hidden_size: int) -> bool: + return hidden_size > 256 and hidden_size <= 8192 and hidden_size % 256 == 0 + + +@cache_once +def _jit_fused_eh_norm_module(hidden_size: int, dtype: torch.dtype) -> Module: + args = make_cpp_args(hidden_size, is_arch_support_pdl(), dtype) + return load_jit( + "fused_eh_norm", + *args, + cuda_files=["elementwise/fused_eh_norm.cuh"], + cuda_wrappers=[("fused_eh_norm", f"FusedEHNormKernel<{args}>::run")], + ) + + +def fused_eh_norm( + inputs_embeds: torch.Tensor, + previous_hidden: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Return fused EH norm + cat for contiguous CUDA fp16/bf16 tensors.""" + if inputs_embeds.dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError( + f"fused_eh_norm: unsupported dtype {inputs_embeds.dtype}; " + "expected torch.float16 or torch.bfloat16" + ) + if inputs_embeds.dim() != 2: + raise RuntimeError( + f"fused_eh_norm: inputs_embeds must be 2D, got {inputs_embeds.dim()}D" + ) + hidden_size = inputs_embeds.shape[1] + if not is_supported_fused_eh_norm_hidden_size(hidden_size): + raise RuntimeError( + f"fused_eh_norm: unsupported hidden_size={hidden_size} " + "(must be in (256, 8192] and a multiple of 256)" + ) + output = torch.empty( + (inputs_embeds.shape[0], hidden_size * 2), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + if inputs_embeds.shape[0] == 0: + return output + module = _jit_fused_eh_norm_module(hidden_size, inputs_embeds.dtype) + module.fused_eh_norm( + inputs_embeds, previous_hidden, enorm_weight, hnorm_weight, output, eps + ) + return output diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 74d504458..6847cc180 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -24,6 +24,7 @@ from safetensors.torch import load_file from torch import nn from transformers import PretrainedConfig +from sglang.jit_kernel.fused_eh_norm import fused_eh_norm from sglang.srt.configs.model_config import is_deepseek_dsa from sglang.srt.distributed import get_pp_group from sglang.srt.environ import envs @@ -196,19 +197,27 @@ class DeepseekModelNextN(nn.Module): hidden_states = input_embeds if hidden_states.shape[0] > 0: - eh_input = torch.cat( - ( - self.enorm(hidden_states), - self.hnorm( - forward_batch.spec_info.hidden_states - if self.rot_weight is None - else torch.matmul( - forward_batch.spec_info.hidden_states, self.rot_weight - ) + previous_hidden_states = forward_batch.spec_info.hidden_states + if self.rot_weight is not None: + previous_hidden_states = torch.matmul( + previous_hidden_states, self.rot_weight + ) + if _is_cuda: + eh_input = fused_eh_norm( + hidden_states, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, + ) + else: + eh_input = torch.cat( + ( + self.enorm(hidden_states), + self.hnorm(previous_hidden_states), ), - ), - dim=-1, - ) + dim=-1, + ) if isinstance(self.eh_proj, ReplicatedLinear): hidden_states, _ = self.eh_proj(eh_input) else: diff --git a/test/registered/jit/benchmark/bench_fused_eh_norm.py b/test/registered/jit/benchmark/bench_fused_eh_norm.py new file mode 100644 index 000000000..4069e69d7 --- /dev/null +++ b/test/registered/jit/benchmark/bench_fused_eh_norm.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import torch + +from sglang.jit_kernel.benchmark import marker +from sglang.jit_kernel.fused_eh_norm import fused_eh_norm +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large") + +EPS = 1e-6 + + +def reference( + x: torch.Tensor, + prev: torch.Tensor, + ew: torch.Tensor, + hw: torch.Tensor, + eps: float, +) -> torch.Tensor: + xf = x.float() + pf = prev.float() + x_var = xf.pow(2).mean(dim=-1, keepdim=True) + p_var = pf.pow(2).mean(dim=-1, keepdim=True) + return torch.cat( + ( + (xf * torch.rsqrt(x_var + eps) * ew.float()).to(x.dtype), + (pf * torch.rsqrt(p_var + eps) * hw.float()).to(prev.dtype), + ), + dim=-1, + ) + + +FN_MAP = { + "jit": fused_eh_norm, + "torch": reference, +} + + +@marker.parametrize("dtype", [torch.bfloat16, torch.float16], [torch.bfloat16]) +@marker.parametrize("hidden_size", [6144, 7168], [7168]) +@marker.parametrize("num_tokens", [1, 4, 6, 8, 16, 32, 128, 512], [1, 6]) +@marker.benchmark("impl", ["jit", "torch"]) +def benchmark(num_tokens: int, hidden_size: int, dtype: torch.dtype, impl: str): + x = torch.randn(num_tokens, hidden_size, device="cuda", dtype=dtype) + prev = torch.randn_like(x) + ew = torch.randn(hidden_size, device="cuda", dtype=dtype) + hw = torch.randn(hidden_size, device="cuda", dtype=dtype) + + expected = reference(x, prev, ew, hw, EPS) + actual = fused_eh_norm(x, prev, ew, hw, EPS) + torch.testing.assert_close(actual.float(), expected.float(), rtol=1e-2, atol=1e-2) + + return marker.do_bench( + FN_MAP[impl], + input_args=(x, prev, ew, hw, EPS), + memory_args=(x, prev, ew, hw), + memory_output="out", + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/jit/test_fused_eh_norm.py b/test/registered/jit/test_fused_eh_norm.py new file mode 100644 index 000000000..3f2d1b4af --- /dev/null +++ b/test/registered/jit/test_fused_eh_norm.py @@ -0,0 +1,122 @@ +import sys + +import pytest +import torch + +from sglang.jit_kernel.fused_eh_norm import fused_eh_norm +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=45, suite="base-b-kernel-unit-1-gpu-large") +register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.cuda is None, + reason="fused_eh_norm requires CUDA", +) + + +def _reference( + inputs_embeds: torch.Tensor, + previous_hidden: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + embeds = inputs_embeds.float() + prev = previous_hidden.float() + embeds_var = embeds.pow(2).mean(dim=-1, keepdim=True) + prev_var = prev.pow(2).mean(dim=-1, keepdim=True) + return torch.cat( + ( + (embeds * torch.rsqrt(embeds_var + eps) * enorm_weight.float()).to( + inputs_embeds.dtype + ), + (prev * torch.rsqrt(prev_var + eps) * hnorm_weight.float()).to( + previous_hidden.dtype + ), + ), + dim=-1, + ) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("hidden_size", [6144, 7168]) +@pytest.mark.parametrize("num_tokens", [1, 6, 128]) +def test_fused_eh_norm_matches_reference( + dtype: torch.dtype, hidden_size: int, num_tokens: int +): + torch.manual_seed(0) + eps = 1e-6 + inputs_embeds = torch.randn(num_tokens, hidden_size, device="cuda", dtype=dtype) + previous_hidden = torch.randn_like(inputs_embeds) + enorm_weight = torch.randn(hidden_size, device="cuda", dtype=dtype) + hnorm_weight = torch.randn(hidden_size, device="cuda", dtype=dtype) + + actual = fused_eh_norm( + inputs_embeds, previous_hidden, enorm_weight, hnorm_weight, eps + ) + expected = _reference( + inputs_embeds, previous_hidden, enorm_weight, hnorm_weight, eps + ) + torch.testing.assert_close(actual.float(), expected.float(), rtol=1e-2, atol=1e-2) + + +def test_fused_eh_norm_zero_tokens(): + hidden_size = 7168 + inputs_embeds = torch.empty(0, hidden_size, device="cuda", dtype=torch.bfloat16) + previous_hidden = torch.empty_like(inputs_embeds) + enorm_weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) + hnorm_weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) + + actual = fused_eh_norm( + inputs_embeds, previous_hidden, enorm_weight, hnorm_weight, 1e-6 + ) + assert actual.shape == (0, hidden_size * 2) + assert actual.dtype == inputs_embeds.dtype + assert actual.device == inputs_embeds.device + + +def test_fused_eh_norm_row_strided_inputs(): + torch.manual_seed(1) + hidden_size = 7168 + eps = 1e-6 + base = torch.randn(12, hidden_size, device="cuda", dtype=torch.bfloat16) + prev_base = torch.randn_like(base) + inputs_embeds = base[::2] + previous_hidden = prev_base[::2] + enorm_weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) + hnorm_weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) + + actual = fused_eh_norm( + inputs_embeds, previous_hidden, enorm_weight, hnorm_weight, eps + ) + expected = _reference( + inputs_embeds, previous_hidden, enorm_weight, hnorm_weight, eps + ) + torch.testing.assert_close(actual.float(), expected.float(), rtol=1e-2, atol=1e-2) + + +def test_fused_eh_norm_rejects_unsupported_dtype(): + hidden_size = 7168 + inputs_embeds = torch.randn(1, hidden_size, device="cuda", dtype=torch.float32) + previous_hidden = torch.randn_like(inputs_embeds) + enorm_weight = torch.randn(hidden_size, device="cuda", dtype=torch.float32) + hnorm_weight = torch.randn(hidden_size, device="cuda", dtype=torch.float32) + + with pytest.raises(RuntimeError, match="unsupported dtype"): + fused_eh_norm(inputs_embeds, previous_hidden, enorm_weight, hnorm_weight, 1e-6) + + +def test_fused_eh_norm_rejects_unsupported_hidden_size(): + hidden_size = 5000 + inputs_embeds = torch.randn(1, hidden_size, device="cuda", dtype=torch.bfloat16) + previous_hidden = torch.randn_like(inputs_embeds) + enorm_weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) + hnorm_weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) + + with pytest.raises(RuntimeError, match="unsupported hidden_size"): + fused_eh_norm(inputs_embeds, previous_hidden, enorm_weight, hnorm_weight, 1e-6) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"]))