minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Support Hy4-preview (#36805)

Browse Source 
Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: alphabetc1 <2508695655@qq.com>
This commit is contained in:
Xinyuan Tong
2026-09-04 18:03:49 -07:00
committed by GitHub
co-authored by BBuf alphabetc1
parent 85da5457de
commit 55bf3380e0
47 changed files with 3387 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/ops/attention/test_dsa_sink_pad_cache.py
+47
View File
@@ -0,0 +1,47 @@
import sys
import pytest
import sgl_kernel.flash_mla as flash_mla
import torch
from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_sparse_flashmla_sink_padding_refreshes_reused_buffer(monkeypatch):
backend = object.__new__(DeepseekSparseAttnBackend)
backend.device_sm_major = 10
backend.dsa_index_topk = 2
backend._sink_pad_cache = {}
captured_sinks = []
def capture_flash_mla_sparse_fwd(**kwargs):
captured_sinks.append(kwargs["attn_sink"].clone())
q = kwargs["q"]
return q.new_zeros((*q.shape[:2], kwargs["d_v"])), None, None
monkeypatch.setattr(flash_mla, "flash_mla_sparse_fwd", capture_flash_mla_sparse_fwd)
q = torch.zeros((1, 64, 8), device="cuda")
kv_cache = torch.zeros((1, 1, 8), device="cuda")
page_table = torch.zeros((1, 2), dtype=torch.int32, device="cuda")
sink = torch.arange(64, dtype=torch.float32, device="cuda")
backend._forward_flashmla_sparse(q, kv_cache, 8, page_table, 1.0, attn_sink=sink)
cached_sink = next(iter(backend._sink_pad_cache.values()))
cached_ptr = cached_sink.data_ptr()
sink.add_(100)
backend._forward_flashmla_sparse(q, kv_cache, 8, page_table, 1.0, attn_sink=sink)
assert next(iter(backend._sink_pad_cache.values())).data_ptr() == cached_ptr
torch.testing.assert_close(captured_sinks[0][:64], sink - 100)
torch.testing.assert_close(captured_sinks[1][:64], sink)
torch.testing.assert_close(captured_sinks[1][64:], torch.zeros_like(sink))
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/attention/test_hy4_hpc_gated_mla.py
+145
View File
@@ -0,0 +1,145 @@
import sys
import pytest
import torch
from torch import nn
from sglang.kernels.jit.utils import get_ci_test_range
from sglang.srt.models.hunyuan_v4 import HYV4Attention
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(
est_time=60,
nightly=False,
disabled=None,
stage="base-b-kernel-unit",
runner_config="4-gpu-b200",
)
register_cuda_ci(
est_time=120,
nightly=True,
disabled=None,
stage="nightly",
runner_config="4-gpu-b200",
)
# Guards attention-TP shards from silently reverting to the global N=16384 shape.
LOCAL_GATE_WIDTHS = get_ci_test_range(
full_range=[256, 512, 1024, 2048, 4096, 8192, 16384],
ci_range=[256, 2048, 16384],
)
def _hpc_gated_mla_available():
try:
from hpc.gemm import gated_mla_gemm # noqa: F401
except (AttributeError, ImportError):
return False
return torch.cuda.is_available() and torch.cuda.get_device_capability() in (
(10, 0),
(10, 3),
)
class TupleLinear(nn.Module):
def __init__(self, weight):
super().__init__()
self.weight = nn.Parameter(weight, requires_grad=False)
def forward(self, inputs):
return nn.functional.linear(inputs, self.weight), None
def _make_attention(weight):
attention = HYV4Attention.__new__(HYV4Attention)
nn.Module.__init__(attention)
attention.linear_gate = TupleLinear(weight)
attention.local_gate_width = weight.shape[0]
attention._gate_backend = "hpc"
attention._gate_fallback_backend = "eager"
return attention
@pytest.mark.skipif(
not _hpc_gated_mla_available(),
reason="requires HPC-Ops gated MLA on SM100 or SM103",
)
@pytest.mark.parametrize("local_gate_width", LOCAL_GATE_WIDTHS)
def test_hy4_gated_mla_attention_tp_eager_graph_parity(local_gate_width):
torch.manual_seed(local_gate_width)
hidden_size = 6144
batch_size = 7
weight = torch.randn(
local_gate_width, hidden_size, dtype=torch.bfloat16, device="cuda"
)
hidden_states = torch.randn(
batch_size, hidden_size, dtype=torch.bfloat16, device="cuda"
)
attn_out = torch.randn(
batch_size, local_gate_width, dtype=torch.bfloat16, device="cuda"
)
attention = _make_attention(weight)
def run():
gate = attention.prepare_attention_output_gate(hidden_states)
return attention.apply_attention_output_gate(attn_out, gate)
expected = attn_out * torch.sigmoid(nn.functional.linear(hidden_states, weight))
eager = run()
torch.testing.assert_close(eager, expected, rtol=0.08, atol=0.01)
stream = torch.cuda.Stream()
stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(stream):
run()
torch.cuda.current_stream().wait_stream(stream)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
graph_out = run()
graph.replay()
torch.testing.assert_close(graph_out, eager, rtol=0, atol=0)
@pytest.mark.skipif(
not _hpc_gated_mla_available(),
reason="requires HPC-Ops gated MLA on SM100 or SM103",
)
@pytest.mark.parametrize("batch_size", [128, 129])
def test_hy4_gated_mla_dispatch_boundary_graph_parity(batch_size):
torch.manual_seed(batch_size)
hidden_size = 6144
local_gate_width = 256
weight = torch.randn(
local_gate_width, hidden_size, dtype=torch.bfloat16, device="cuda"
)
hidden_states = torch.randn(
batch_size, hidden_size, dtype=torch.bfloat16, device="cuda"
)
attn_out = torch.randn(
batch_size, local_gate_width, dtype=torch.bfloat16, device="cuda"
)
attention = _make_attention(weight)
def run():
gate = attention.prepare_attention_output_gate(hidden_states)
return attention.apply_attention_output_gate(attn_out, gate)
eager = run()
expected = attn_out * torch.sigmoid(nn.functional.linear(hidden_states, weight))
torch.testing.assert_close(eager, expected, rtol=0.08, atol=0.01)
stream = torch.cuda.Stream()
stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(stream):
run()
torch.cuda.current_stream().wait_stream(stream)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
graph_out = run()
graph.replay()
torch.testing.assert_close(graph_out, eager, rtol=0, atol=0)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/layernorm/test_hy4_decode_kernels.py
+229
View File
@@ -0,0 +1,229 @@
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import torch
import triton
import triton.language as tl
from torch import nn
import sglang.kernels.ops.layernorm.hy4_ihc as hy4_ihc
from sglang.srt.models import hunyuan_v4
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=35, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@triton.jit
def _reference_hy4_ihc_pre_kernel(
x_ptr,
fn_ptr,
scale_ptr,
base_ptr,
y_ptr,
post_ptr,
hidden_size: tl.constexpr,
HC_MULT: tl.constexpr,
HC_POW2: tl.constexpr,
K_TOTAL: tl.constexpr,
BLOCK_K: tl.constexpr,
BLOCK_D: tl.constexpr,
magnitude: tl.constexpr,
norm_eps: tl.constexpr,
hc_eps: tl.constexpr,
):
pid = tl.program_id(0).to(tl.int64)
x_row = x_ptr + pid * K_TOTAL
m_idx = tl.arange(0, HC_POW2)
m_mask = m_idx < HC_MULT
sumsq = tl.zeros((), dtype=tl.float32)
mix_pre = tl.zeros((HC_POW2,), dtype=tl.float32)
mix_post = tl.zeros((HC_POW2,), dtype=tl.float32)
for k_off in tl.range(0, K_TOTAL, BLOCK_K):
k_offs = k_off + tl.arange(0, BLOCK_K)
k_mask = k_offs < K_TOTAL
x_tile = tl.load(x_row + k_offs, mask=k_mask, other=0.0).to(tl.float32)
sumsq += tl.sum(x_tile * x_tile, axis=0)
fn_offs = m_idx[:, None] * K_TOTAL + k_offs[None, :]
fn_mask = m_mask[:, None] & k_mask[None, :]
mix_pre += tl.sum(
tl.load(fn_ptr + fn_offs, mask=fn_mask, other=0.0) * x_tile[None, :],
axis=1,
)
mix_post += tl.sum(
tl.load(
fn_ptr + HC_MULT * K_TOTAL + fn_offs,
mask=fn_mask,
other=0.0,
)
* x_tile[None, :],
axis=1,
)
rsqrt = tl.rsqrt(sumsq / K_TOTAL + norm_eps)
scale_pre = tl.load(scale_ptr)
scale_post = tl.load(scale_ptr + 1)
base_pre = tl.load(base_ptr + m_idx, mask=m_mask, other=0.0)
base_post = tl.load(base_ptr + HC_MULT + m_idx, mask=m_mask, other=0.0)
pre = tl.sigmoid(mix_pre * rsqrt * scale_pre + base_pre) + hc_eps
post = magnitude * tl.sigmoid(mix_post * rsqrt * scale_post + base_post) + hc_eps
tl.store(post_ptr + pid * HC_MULT + m_idx, post, mask=m_mask)
y_row = y_ptr + pid * hidden_size
for d_off in tl.range(0, hidden_size, BLOCK_D):
d_offs = d_off + tl.arange(0, BLOCK_D)
d_mask = d_offs < hidden_size
y_block = tl.zeros((BLOCK_D,), dtype=tl.float32)
for m in tl.static_range(HC_MULT):
x_m = tl.load(x_row + m * hidden_size + d_offs, mask=d_mask, other=0.0)
pre_m = tl.sum(tl.where(m_idx == m, pre, 0.0), axis=0)
y_block += pre_m * x_m.to(tl.float32)
tl.store(
y_row + d_offs,
y_block.to(y_ptr.dtype.element_ty),
mask=d_mask,
)
def _reference_hy4_ihc_pre(x, hc_fn, hc_scale, hc_base):
num_tokens, hc_mult, hidden_size = x.shape
k_total = hc_mult * hidden_size
y = torch.empty((num_tokens, hidden_size), dtype=x.dtype, device=x.device)
post = torch.empty((num_tokens, hc_mult), dtype=torch.float32, device=x.device)
if num_tokens == 0:
return y, post
_reference_hy4_ihc_pre_kernel[(num_tokens,)](
x,
hc_fn,
hc_scale,
hc_base,
y,
post,
hidden_size=hidden_size,
HC_MULT=hc_mult,
HC_POW2=triton.next_power_of_2(hc_mult),
K_TOTAL=k_total,
BLOCK_K=1024,
BLOCK_D=1024,
magnitude=2.0,
norm_eps=1e-6,
hc_eps=1e-6,
num_warps=8,
enable_fp_fusion=False,
)
return y, post
@unittest.skipUnless(torch.cuda.is_available(), "HYV4 Triton kernels need CUDA")
class TestHy4DecodeKernels(CustomTestCase):
def test_split_k_matches_single_cta(self):
torch.manual_seed(0)
for num_tokens, hidden_size in (
(0, 6144),
(1, 4096),
(3, 4100),
(31, 6144),
(64, 6144),
):
with self.subTest(num_tokens=num_tokens, hidden_size=hidden_size):
hc_mult = 4
k_total = hc_mult * hidden_size
x = torch.randn(
num_tokens,
hc_mult,
hidden_size,
device="cuda",
dtype=torch.bfloat16,
)
hc_fn = (
torch.randn(
2 * hc_mult,
k_total,
device="cuda",
dtype=torch.float32,
)
* 0.02
)
hc_scale = torch.tensor([0.7, 1.3], device="cuda", dtype=torch.float32)
hc_base = (
torch.randn(2 * hc_mult, device="cuda", dtype=torch.float32) * 0.5
)
expected = _reference_hy4_ihc_pre(x, hc_fn, hc_scale, hc_base)
with patch.object(hy4_ihc, "_hpc_ihc_op", return_value=None):
actual = hy4_ihc.fused_hy4_ihc_pre(
x, hc_fn, hc_scale, hc_base, 2.0, 1e-6, 1e-6
)
self.assertTrue(torch.equal(actual[0], expected[0]))
self.assertTrue(torch.equal(actual[1], expected[1]))
def test_fused_ihc_failure_disables_each_path(self):
class TupleLinear(nn.Module):
def __init__(self, input_size, output_size):
super().__init__()
self.weight = nn.Parameter(torch.randn(output_size, input_size))
def forward(self, inputs):
return nn.functional.linear(inputs, self.weight), None
config = SimpleNamespace(
hidden_size=8,
hc_mult=2,
hc_magnitude=2.0,
hc_eps=1e-6,
rms_norm_eps=1e-5,
)
counts = {"pre": 0, "post": 0, "post_pre": 0, "head": 0}
def fail(name):
def raise_error(*args, **kwargs):
counts[name] += 1
raise RuntimeError(name)
return raise_error
def make_linear(input_size, output_size, **kwargs):
return TupleLinear(input_size, output_size)
with patch.object(hunyuan_v4, "ReplicatedLinear", make_linear):
pre_layer = hunyuan_v4.HYV4HCPreLayer(config, "pre").cuda()
layer = hunyuan_v4.HYV4HCLayer(config, "layer").cuda()
next_layer = hunyuan_v4.HYV4HCLayer(config, "next").cuda()
head_layer = hunyuan_v4.HYV4HCHeadLayer(config, "head").cuda()
next_layer.hc_pre._fused_ihc_pre_disabled = True
norm = hunyuan_v4.RMSNorm(
config.hidden_size, eps=config.rms_norm_eps, force_native=True
).cuda()
hidden_states = torch.randn(3, 2, 8, device="cuda")
output = torch.randn(3, 8, device="cuda")
residual = torch.randn(3, 2, 8, device="cuda")
post = torch.randn(3, 2, device="cuda")
with (
patch.object(hy4_ihc, "fused_hy4_ihc_pre", fail("pre")),
patch.object(hy4_ihc, "fused_hy4_ihc_post", fail("post")),
patch.object(hy4_ihc, "fused_hy4_ihc_post_pre", fail("post_pre")),
patch.object(hy4_ihc, "fused_hy4_ihc_head", fail("head")),
patch.object(hunyuan_v4, "_hpc_ihc_available", return_value=True),
self.assertLogs(hunyuan_v4.logger, level="WARNING") as logs,
):
for _ in range(2):
pre_layer(hidden_states)
layer.post(output, residual, post)
layer.post_pre(output, residual, post, next_layer, norm)
head_layer(hidden_states)
self.assertEqual(counts, {"pre": 1, "post": 1, "post_pre": 1, "head": 1})
self.assertEqual(len(logs.records), 4)
self.assertTrue(all(record.exc_info is not None for record in logs.records))
if __name__ == "__main__":
unittest.main()
test/registered/kernels/ops/layernorm/test_hy4_hpc_ihc.py
+111
View File
@@ -0,0 +1,111 @@
import sys
import pytest
import torch
from sglang.kernels.ops.layernorm.hy4_ihc import (
_hpc_ihc_op,
fused_hy4_ihc_head,
fused_hy4_ihc_post_pre,
fused_hy4_ihc_pre,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@pytest.mark.parametrize("hidden_size", [4096, 6144])
def test_hpc_ihc_eager_graph_parity(hidden_size):
if not torch.cuda.is_available():
pytest.skip("requires CUDA")
if _hpc_ihc_op("fuse_ihc_post_pre", 4, hidden_size) is None:
pytest.skip("requires a compatible HPC-Ops iHC build")
torch.manual_seed(13)
num_tokens = 7
hc_mult = 4
norm_eps, hc_eps, magnitude = 1e-5, 1e-6, 2.0
x = torch.rand(
(num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16, device="cuda"
)
output = torch.rand((num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda")
pre_weight = (
torch.rand(
(2 * hc_mult, hc_mult * hidden_size),
dtype=torch.float32,
device="cuda",
)
* 6e-3
)
next_weight = torch.rand_like(pre_weight) * 6e-3
head_weight = (
torch.rand(
(hc_mult, hc_mult * hidden_size),
dtype=torch.float32,
device="cuda",
)
* 6e-3
)
pre_scale = torch.rand((2,), dtype=torch.float32, device="cuda")
next_scale = torch.rand((2,), dtype=torch.float32, device="cuda")
head_scale = torch.rand((1,), dtype=torch.float32, device="cuda")
pre_base = torch.rand((2 * hc_mult,), dtype=torch.float32, device="cuda")
next_base = torch.rand((2 * hc_mult,), dtype=torch.float32, device="cuda")
head_base = torch.rand((hc_mult,), dtype=torch.float32, device="cuda")
rms_weight = torch.rand((hidden_size,), dtype=torch.bfloat16, device="cuda")
def run():
_, post = fused_hy4_ihc_pre(
x,
pre_weight,
pre_scale,
pre_base,
magnitude,
norm_eps,
hc_eps,
rms_weight,
norm_eps,
)
residual, reduced, next_post = fused_hy4_ihc_post_pre(
output,
x,
post,
next_weight,
next_scale,
next_base,
magnitude,
norm_eps,
hc_eps,
rms_weight,
norm_eps,
)
head = fused_hy4_ihc_head(
residual,
head_weight,
head_scale,
head_base,
norm_eps,
hc_eps,
rms_weight,
norm_eps,
)
return residual, reduced, next_post, head
warmup_stream = torch.cuda.Stream()
warmup_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(warmup_stream):
run()
torch.cuda.current_stream().wait_stream(warmup_stream)
eager = tuple(value.clone() for value in run())
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
graph_outputs = run()
graph.replay()
for eager_output, graph_output in zip(eager, graph_outputs):
assert torch.equal(eager_output, graph_output)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/kernels/ops/moe/test_minimax_quant_scatter.py
+25 -7
View File
@@ -243,6 +243,18 @@ def test_standard_masked_runner_matches_compact_end_to_end(monkeypatch, weight_d
"use_symmetric_memory",
lambda *args, **kwargs: nullcontext(),
)
monkeypatch.setattr(
deep_gemm_runner.deep_gemm_wrapper,
"get_contiguous_layout_alignment",
lambda expected_m, num_groups: 32,
)
monkeypatch.setattr(
deep_gemm_runner,
"get_exec",
lambda: SimpleNamespace(
deterministic=SimpleNamespace(enable_deterministic_inference=False)
),
)
# UE8M0 packs four 128-wide scale groups into each int32. Use the smallest
# legal K for both the gate/up and down GEMMs.
@@ -357,17 +369,23 @@ def test_standard_masked_runner_matches_compact_end_to_end(monkeypatch, weight_d
).hidden_states,
)
compact_is_masked, compact_all_tokens, compact_m_indices, compact_output = (
run_with_layout("compact")
)
masked_is_masked, masked_all_tokens, masked_m_indices, masked_output = (
run_with_layout("masked")
)
(
compact_is_masked,
compact_all_tokens,
compact_m_indices,
compact_output,
) = run_with_layout("compact")
(
masked_is_masked,
masked_all_tokens,
masked_m_indices,
masked_output,
) = run_with_layout("masked")
torch.cuda.synchronize()
assert not compact_is_masked
assert masked_is_masked
assert compact_all_tokens == 256
assert compact_all_tokens == 64
assert masked_all_tokens is None
assert masked_m_indices is None
valid_assignments = topk_ids[topk_ids >= 0]
test/registered/kernels/ops/moe/test_pad_expert_counts.py
+37
View File
@@ -0,0 +1,37 @@
import unittest
import torch
from sglang.kernels.ops.moe.triton_pad_expert_counts import pad_expert_counts
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@unittest.skipUnless(torch.cuda.is_available(), "pad_expert_counts needs CUDA")
class TestPadExpertCounts(CustomTestCase):
def test_matches_eager(self):
cases = (
([0], 16, 32),
([1, 8, 9, 0], 8, 48),
([0, 1, 7, 8, 9, 31, 32], 16, 160),
)
for dtype in (torch.int32, torch.int64):
for values, block_e, all_tokens in cases:
with self.subTest(dtype=dtype, values=values):
counts = torch.tensor(values, device="cuda", dtype=dtype)
expected = (((counts + block_e - 1) // block_e) * block_e).to(
torch.int32
)
expected[-1].add_(all_tokens - expected.sum())
actual = pad_expert_counts(counts, block_e, all_tokens)
self.assertTrue(torch.equal(actual, expected))
self.assertEqual(actual.dtype, torch.int32)
self.assertEqual(actual.sum().item(), all_tokens)
if __name__ == "__main__":
unittest.main()