Add fused EH norm for DeepSeek NextN (#29667)
This commit is contained in:
Mohammad Miadh Angkad
@@ -0,0 +1,113 @@
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/runtime.cuh>
|
||||
#include <sgl_kernel/tile.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
|
||||
#include <sgl_kernel/impl/norm.cuh>
|
||||
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
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 <int64_t kHidden, bool kUsePDL, typename T>
|
||||
__global__ void fused_eh_norm_kernel(const __grid_constant__ FusedEHNormParams params) {
|
||||
using namespace device;
|
||||
using Storage = norm::StorageType<T, kHidden>;
|
||||
|
||||
constexpr auto kNumThreads = host::norm::get_cta_threads<T, kHidden>();
|
||||
constexpr auto kNumWarps = kNumThreads / kWarpThreads;
|
||||
|
||||
const auto embeds = static_cast<const T*>(pointer::offset<T>(params.embeds, blockIdx.x * params.embeds_stride));
|
||||
const auto previous_hidden =
|
||||
static_cast<const T*>(pointer::offset<T>(params.previous_hidden, blockIdx.x * params.previous_hidden_stride));
|
||||
const auto enorm_weight = static_cast<const T*>(params.enorm_weight);
|
||||
const auto hnorm_weight = static_cast<const T*>(params.hnorm_weight);
|
||||
const auto output = static_cast<T*>(pointer::offset<T>(params.output, blockIdx.x * params.output_stride));
|
||||
|
||||
const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
|
||||
__shared__ float smem[norm::kSmemBufferSize];
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
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<kHidden>(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<kHidden>(prev_vec, hnorm_weight_vec, params.eps, smem, kNumWarps);
|
||||
gmem.store(output + kHidden, prev_output_vec);
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
template <int64_t kHidden, bool kUsePDL, typename T>
|
||||
struct FusedEHNormKernel {
|
||||
static_assert(host::norm::is_config_supported<T, kHidden>(), "Unsupported norm configuration");
|
||||
static_assert(host::norm::should_use_cta<T, kHidden>(), "fused_eh_norm requires CTA norm");
|
||||
static constexpr auto kernel = fused_eh_norm_kernel<kHidden, kUsePDL, T>;
|
||||
static constexpr uint32_t kBlockSize = host::norm::get_cta_threads<T, kHidden>();
|
||||
|
||||
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<kDLCUDA>();
|
||||
|
||||
TensorMatcher({N, H}).with_strides({SE, 1}).with_dtype<T>().with_device(device_).verify(embeds);
|
||||
TensorMatcher({N, H}).with_strides({SP, 1}).with_dtype<T>().with_device(device_).verify(previous_hidden);
|
||||
TensorMatcher({H}).with_dtype<T>().with_device(device_).verify(enorm_weight);
|
||||
TensorMatcher({H}).with_dtype<T>().with_device(device_).verify(hnorm_weight);
|
||||
TensorMatcher({N, H2}).with_strides({SO, 1}).with_dtype<T>().with_device(device_).verify(output);
|
||||
|
||||
const auto num_tokens = static_cast<uint32_t>(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
|
||||
@@ -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
|
||||
@@ -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,16 +197,24 @@ class DeepseekModelNextN(nn.Module):
|
||||
hidden_states = input_embeds
|
||||
|
||||
if hidden_states.shape[0] > 0:
|
||||
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(
|
||||
forward_batch.spec_info.hidden_states
|
||||
if self.rot_weight is None
|
||||
else torch.matmul(
|
||||
forward_batch.spec_info.hidden_states, self.rot_weight
|
||||
)
|
||||
),
|
||||
self.hnorm(previous_hidden_states),
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
|
||||
@@ -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()
|
||||
@@ -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"]))
|
||||
Reference in New Issue
Block a user