Support Hy4-preview (#36805)
Co-authored-by: BBuf <1182563586@qq.com> Co-authored-by: alphabetc1 <2508695655@qq.com>
This commit is contained in:
co-authored by
BBuf
alphabetc1
parent
85da5457de
commit
55bf3380e0
@@ -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__]))
|
||||
@@ -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__]))
|
||||
@@ -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()
|
||||
@@ -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__]))
|
||||
@@ -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]
|
||||
|
||||
@@ -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()
|
||||
@@ -4,17 +4,76 @@ from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.quantization.fp8 import (
|
||||
Fp8MoEMethod,
|
||||
_is_cuda,
|
||||
_is_gfx95_supported,
|
||||
_is_hip,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
inverse_transform_scale_ue8m0,
|
||||
quant_weight_ue8m0,
|
||||
transform_scale_ue8m0,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_platform
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=12, stage="base-b", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
class TestMxfp8MoeScaleLayout(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if not (
|
||||
(_is_cuda and get_platform().is_sm100) or (_is_hip and _is_gfx95_supported)
|
||||
):
|
||||
raise unittest.SkipTest(
|
||||
"MXFP8 MoE quantization requires SM100 or ROCm gfx95"
|
||||
)
|
||||
|
||||
def test_cutlass_serialized_scales_remain_expert_first(self):
|
||||
class CutlassBackend:
|
||||
def is_cutlass(self):
|
||||
return True
|
||||
|
||||
def is_flashinfer_trtllm(self):
|
||||
return False
|
||||
|
||||
def is_flashinfer_trtllm_routed(self):
|
||||
return False
|
||||
|
||||
def is_deep_gemm(self):
|
||||
return False
|
||||
|
||||
layer = SimpleNamespace(
|
||||
w13_weight=torch.nn.Parameter(
|
||||
torch.zeros((2, 64, 32), dtype=torch.float8_e4m3fn, device="cuda")
|
||||
),
|
||||
w2_weight=torch.nn.Parameter(
|
||||
torch.zeros((2, 32, 32), dtype=torch.float8_e4m3fn, device="cuda")
|
||||
),
|
||||
w13_weight_scale_inv=torch.nn.Parameter(
|
||||
torch.zeros((2, 64, 1), dtype=torch.uint8, device="cuda"),
|
||||
requires_grad=False,
|
||||
),
|
||||
w2_weight_scale_inv=torch.nn.Parameter(
|
||||
torch.zeros((2, 32, 1), dtype=torch.uint8, device="cuda"),
|
||||
requires_grad=False,
|
||||
),
|
||||
)
|
||||
method = object.__new__(Fp8MoEMethod)
|
||||
|
||||
with patch(
|
||||
"sglang.srt.layers.quantization.fp8.get_moe_runner_backend",
|
||||
return_value=CutlassBackend(),
|
||||
):
|
||||
method._process_mxfp8_moe_weights(layer, quantize=False)
|
||||
|
||||
self.assertEqual(tuple(layer.w13_weight_scale_inv.shape), (2, 64, 1))
|
||||
self.assertEqual(tuple(layer.w2_weight_scale_inv.shape), (2, 32, 1))
|
||||
|
||||
|
||||
class TestInverseTransformScaleUe8m0(CustomTestCase):
|
||||
def test_round_trip(self):
|
||||
for _ in range(100):
|
||||
|
||||
@@ -10,6 +10,7 @@ from sglang.srt.layers.logits_processor import LogitsProcessor
|
||||
from sglang.srt.runtime_context import get_context
|
||||
from sglang.srt.utils import get_device
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small")
|
||||
register_amd_ci(est_time=15, suite="stage-b-test-1-gpu-small-amd")
|
||||
@@ -30,7 +31,7 @@ class DummyMeta:
|
||||
def compute_dp_attention_metadata(self): ...
|
||||
|
||||
|
||||
class TestLMHeadFP32(unittest.TestCase):
|
||||
class TestLMHeadFP32(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if not torch.cuda.is_available() and not (
|
||||
@@ -38,13 +39,17 @@ class TestLMHeadFP32(unittest.TestCase):
|
||||
):
|
||||
raise unittest.SkipTest("needs CUDA GPU or XPU")
|
||||
|
||||
def _make_logprocessor(self, vocab_size, enable_fp32):
|
||||
def _make_logprocessor(self, vocab_size, enable_fp32, config_enable_fp32=False):
|
||||
# LogitsProcessor reads get_exec().features.enable_fp32_lm_head
|
||||
# from the published config.
|
||||
override = get_context().override_server_args(enable_fp32_lm_head=enable_fp32)
|
||||
override.install()
|
||||
self.addCleanup(override.restore)
|
||||
cfg = SimpleNamespace(vocab_size=vocab_size, final_logit_softcapping=None)
|
||||
cfg = SimpleNamespace(
|
||||
vocab_size=vocab_size,
|
||||
final_logit_softcapping=None,
|
||||
enable_lm_head_fp32=config_enable_fp32,
|
||||
)
|
||||
return LogitsProcessor(cfg, skip_all_gather=True, logit_scale=None)
|
||||
|
||||
def _run_case(
|
||||
@@ -55,6 +60,7 @@ class TestLMHeadFP32(unittest.TestCase):
|
||||
expected_a_dtype,
|
||||
expected_b_dtype,
|
||||
expected_operation,
|
||||
config_enable_fp32=False,
|
||||
):
|
||||
device = get_device()
|
||||
BATCH_SIZE, HIDDEN_SIZE, VOCAB_SIZE = 2, 64, 128
|
||||
@@ -63,7 +69,9 @@ class TestLMHeadFP32(unittest.TestCase):
|
||||
)
|
||||
head = LMHeadStub(VOCAB_SIZE, HIDDEN_SIZE, dtype=weights_dtype, device=device)
|
||||
meta = DummyMeta()
|
||||
logprocessor = self._make_logprocessor(VOCAB_SIZE, enable_fp32)
|
||||
logprocessor = self._make_logprocessor(
|
||||
VOCAB_SIZE, enable_fp32, config_enable_fp32
|
||||
)
|
||||
|
||||
original_matmul = torch.matmul
|
||||
original_mm = torch.mm
|
||||
@@ -173,6 +181,21 @@ class TestLMHeadFP32(unittest.TestCase):
|
||||
"matmul",
|
||||
)
|
||||
|
||||
def test_model_config_enables_fp32_without_server_flag(self):
|
||||
expected_operation = "mm" if torch.cuda.is_available() else "matmul"
|
||||
expected_dtype = (
|
||||
torch.float32 if expected_operation == "matmul" else torch.bfloat16
|
||||
)
|
||||
self._run_case(
|
||||
torch.bfloat16,
|
||||
False,
|
||||
torch.bfloat16,
|
||||
expected_dtype,
|
||||
expected_dtype,
|
||||
expected_operation,
|
||||
config_enable_fp32=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import math
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.configs.model_config import ModelConfig, compute_mla_mscale_scaling
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
@@ -92,6 +93,31 @@ class TestInitMlaScaling(CustomTestCase):
|
||||
base, scaling = _mla_scaling(None)
|
||||
self.assertEqual(scaling, base)
|
||||
|
||||
def test_hyv4_shape_derivation_uses_rope_parameters(self):
|
||||
hf_config = SimpleNamespace(
|
||||
architectures=["HYV4ForCausalLM"],
|
||||
model_type="hy_v4",
|
||||
hidden_size=2816,
|
||||
num_hidden_layers=34,
|
||||
num_attention_heads=32,
|
||||
vocab_size=120832,
|
||||
head_dim=64,
|
||||
v_head_dim=256,
|
||||
kv_lora_rank=512,
|
||||
qk_nope_head_dim=192,
|
||||
qk_rope_head_dim=64,
|
||||
index_topk=2048,
|
||||
index_head_dim=128,
|
||||
rope_parameters={"rope_theta": 10_000_000, "rope_type": "default"},
|
||||
)
|
||||
config = ModelConfig.__new__(ModelConfig)
|
||||
config.hf_config = hf_config
|
||||
config.hf_text_config = hf_config
|
||||
|
||||
config._derive_model_shapes()
|
||||
|
||||
self.assertEqual(config.scaling, 1 / math.sqrt(256))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -800,6 +800,87 @@ class ServingChatTestCase(unittest.TestCase):
|
||||
|
||||
self.assertEqual(req.reasoning_effort, "high")
|
||||
|
||||
def test_hunyuan_default_reasoning_effort_is_normalized_for_template(self):
|
||||
self.template_manager.chat_template_name = None
|
||||
self.template_manager.jinja_template_content_format = "string"
|
||||
self.template_manager.reasoning_config = ReasoningToggleConfig(
|
||||
special_case="hunyuan_effort"
|
||||
)
|
||||
self.chat.reasoning_parser = "hunyuan"
|
||||
self.tm.tokenizer.apply_chat_template.return_value = [1, 2, 3]
|
||||
|
||||
cases = [
|
||||
("none", None, "no_think"),
|
||||
("none", "xhigh", "high"),
|
||||
]
|
||||
for default_effort, request_effort, normalized_effort in cases:
|
||||
with self.subTest(
|
||||
default_effort=default_effort, request_effort=request_effort
|
||||
):
|
||||
self.chat.default_chat_template_kwargs = {
|
||||
"reasoning_effort": default_effort
|
||||
}
|
||||
req = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[{"role": "user", "content": "What is 2+2?"}],
|
||||
reasoning_effort=request_effort,
|
||||
)
|
||||
|
||||
self.chat._process_messages(req, is_multimodal=False)
|
||||
|
||||
kwargs = self.tm.tokenizer.apply_chat_template.call_args.kwargs
|
||||
self.assertEqual(req.reasoning_effort, normalized_effort)
|
||||
self.assertEqual(kwargs["reasoning_effort"], normalized_effort)
|
||||
self.assertNotIn("reasoning_effort", req.chat_template_kwargs)
|
||||
|
||||
def test_hunyuan_reasoning_effort_precedence_survives_conversion(self):
|
||||
self.template_manager.chat_template_name = None
|
||||
self.template_manager.jinja_template_content_format = "string"
|
||||
self.template_manager.reasoning_config = ReasoningToggleConfig(
|
||||
special_case="hunyuan_effort"
|
||||
)
|
||||
self.chat.reasoning_parser = "hunyuan"
|
||||
self.tm.tokenizer.apply_chat_template.return_value = [1, 2, 3]
|
||||
|
||||
cases = [
|
||||
("xhigh", "none", "high"),
|
||||
(None, "no_think", "no_think"),
|
||||
]
|
||||
for request_effort, template_effort, normalized_effort in cases:
|
||||
with self.subTest(
|
||||
request_effort=request_effort, template_effort=template_effort
|
||||
):
|
||||
req = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[{"role": "user", "content": "What is 2+2?"}],
|
||||
reasoning_effort=request_effort,
|
||||
chat_template_kwargs={"reasoning_effort": template_effort},
|
||||
)
|
||||
|
||||
self.chat._convert_to_internal_request(req)
|
||||
|
||||
kwargs = self.tm.tokenizer.apply_chat_template.call_args.kwargs
|
||||
self.assertEqual(req.reasoning_effort, normalized_effort)
|
||||
self.assertEqual(kwargs["reasoning_effort"], normalized_effort)
|
||||
self.assertNotIn("reasoning_effort", req.chat_template_kwargs)
|
||||
|
||||
def test_non_hunyuan_default_reasoning_effort_is_unchanged(self):
|
||||
self.template_manager.chat_template_name = None
|
||||
self.template_manager.jinja_template_content_format = "string"
|
||||
self.tm.tokenizer.apply_chat_template.return_value = [1, 2, 3]
|
||||
self.chat.default_chat_template_kwargs = {"reasoning_effort": "medium"}
|
||||
req = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[{"role": "user", "content": "What is 2+2?"}],
|
||||
)
|
||||
|
||||
self.chat._process_messages(req, is_multimodal=False)
|
||||
|
||||
kwargs = self.tm.tokenizer.apply_chat_template.call_args.kwargs
|
||||
self.assertEqual(req.reasoning_effort, "medium")
|
||||
self.assertEqual(kwargs["reasoning_effort"], "medium")
|
||||
self.assertEqual(req.chat_template_kwargs["reasoning_effort"], "medium")
|
||||
|
||||
def test_k2_selected_terminator_reaches_sampling_params(self):
|
||||
self.tm._config_overrides["reasoning_parser"] = "k2_horizon"
|
||||
self.chat = OpenAIServingChat(self.tm, self.template_manager)
|
||||
|
||||
@@ -75,6 +75,17 @@ def _make_tools():
|
||||
]
|
||||
|
||||
|
||||
class _Hy4Tokenizer:
|
||||
def get_vocab(self):
|
||||
return {
|
||||
"<tool_calls:opensource>": 1,
|
||||
"<tool_call:opensource>": 2,
|
||||
"<arg_key:opensource>": 3,
|
||||
"<arg_value:opensource>": 4,
|
||||
"<think:opensource>": 5,
|
||||
}
|
||||
|
||||
|
||||
class TestHunyuanDetectorHasToolCall(CustomTestCase):
|
||||
def setUp(self):
|
||||
self.detector = HunyuanDetector()
|
||||
@@ -111,6 +122,21 @@ class TestHunyuanDetectorDetectAndParse(CustomTestCase):
|
||||
self.assertEqual(len(result.calls), 0)
|
||||
self.assertEqual(result.normal_text, text)
|
||||
|
||||
def test_hy4_format_without_tool_separator(self):
|
||||
detector = HunyuanDetector(_Hy4Tokenizer())
|
||||
text = (
|
||||
"<tool_calls:opensource>"
|
||||
"<tool_call:opensource>get_weather"
|
||||
"<arg_key:opensource>city</arg_key:opensource>"
|
||||
"<arg_value:opensource>Beijing</arg_value:opensource>"
|
||||
"</tool_call:opensource></tool_calls:opensource>"
|
||||
)
|
||||
|
||||
result = detector.detect_and_parse(text, self.tools)
|
||||
|
||||
self.assertEqual(result.calls[0].name, "get_weather")
|
||||
self.assertEqual(json.loads(result.calls[0].parameters), {"city": "Beijing"})
|
||||
|
||||
def test_zero_arg_inline(self):
|
||||
text = (
|
||||
"<tool_calls><tool_call>get_current_date<tool_sep></tool_call></tool_calls>"
|
||||
@@ -293,6 +319,43 @@ class TestHunyuanDetectorArgDeserialization(CustomTestCase):
|
||||
args = json.loads(result.calls[0].parameters)
|
||||
self.assertIs(args["verbose"], True)
|
||||
|
||||
def test_top_level_composed_schema_args(self):
|
||||
cases = (
|
||||
("anyOf", {"type": "integer"}, "7", 7),
|
||||
("oneOf", {"type": "boolean"}, "true", True),
|
||||
("allOf", {"type": "array", "items": {"type": "integer"}}, "[1,2]", [1, 2]),
|
||||
)
|
||||
for keyword, arg_schema, raw_value, expected in cases:
|
||||
with self.subTest(keyword=keyword):
|
||||
function_name = f"composed_{keyword}"
|
||||
tools = [
|
||||
Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name=function_name,
|
||||
description="Composed schema",
|
||||
parameters={
|
||||
keyword: [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {"value": arg_schema},
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
)
|
||||
]
|
||||
text = (
|
||||
f"<tool_calls><tool_call>{function_name}<tool_sep>"
|
||||
f"<arg_key>value</arg_key><arg_value>{raw_value}</arg_value>"
|
||||
"</tool_call></tool_calls>"
|
||||
)
|
||||
|
||||
result = self.detector.detect_and_parse(text, tools)
|
||||
args = json.loads(result.calls[0].parameters)
|
||||
|
||||
self.assertEqual(args, {"value": expected})
|
||||
|
||||
def test_string_arg_not_deserialized(self):
|
||||
"""String-typed args should stay as strings even if they look like JSON."""
|
||||
text = (
|
||||
@@ -359,6 +422,23 @@ class TestHunyuanDetectorStreaming(CustomTestCase):
|
||||
self.assertEqual(collected[0]["name"], "get_current_date")
|
||||
self.assertEqual(json.loads(collected[0]["parameters"]), {})
|
||||
|
||||
def test_hy4_format_without_tool_separator_char_by_char(self):
|
||||
detector = HunyuanDetector(_Hy4Tokenizer())
|
||||
text = (
|
||||
"<tool_calls:opensource>"
|
||||
"<tool_call:opensource>get_weather"
|
||||
"<arg_key:opensource>city</arg_key:opensource>"
|
||||
"<arg_value:opensource>Tokyo</arg_value:opensource>"
|
||||
"</tool_call:opensource></tool_calls:opensource>"
|
||||
)
|
||||
all_calls = []
|
||||
for char in text:
|
||||
all_calls.extend(detector.parse_streaming_increment(char, self.tools).calls)
|
||||
|
||||
collected = _collect_streamed_tool_calls(all_calls)
|
||||
self.assertEqual(collected[0]["name"], "get_weather")
|
||||
self.assertEqual(json.loads(collected[0]["parameters"]), {"city": "Tokyo"})
|
||||
|
||||
def test_chunked_tool_call(self):
|
||||
detector = self._new_detector()
|
||||
chunks = [
|
||||
|
||||
@@ -552,6 +552,21 @@ class TestParseQuantHfConfig(CustomTestCase):
|
||||
self.assertIn("lm_head", quant_config.ignored_layers)
|
||||
self.assertEqual(quant_config.kv_cache_quant_algo, "FP8")
|
||||
|
||||
nested_result = model_config._parse_modelopt_quant_config(
|
||||
{
|
||||
"quantization": {
|
||||
"quantization": {
|
||||
"quant_algo": "MXFP8",
|
||||
"group_size": 32,
|
||||
"exclude_modules": ["lm_head"],
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
self.assertEqual(nested_result["quant_method"], "mxfp8")
|
||||
self.assertEqual(nested_result["scale_fmt"], "ue8m0")
|
||||
self.assertIn("lm_head", nested_result["modules_to_not_convert"])
|
||||
|
||||
def test_modelopt_mxfp8_override(self):
|
||||
"""Generic ModelOpt selection must not route MXFP8 to scalar FP8."""
|
||||
self.assertEqual(
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from sglang.srt.models import hunyuan_v4
|
||||
from sglang.srt.models.deepseek_common.attention_forward_methods import forward_mla
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=4, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def test_attention_gate_uses_attention_tp(monkeypatch):
|
||||
attn_tp_size = 2
|
||||
parallel = SimpleNamespace(
|
||||
attn_tp_rank=min(1, attn_tp_size - 1), attn_tp_size=attn_tp_size
|
||||
)
|
||||
captured = {}
|
||||
|
||||
def fake_attention_init(module, **kwargs):
|
||||
nn.Module.__init__(module)
|
||||
module.hidden_size = kwargs["hidden_size"]
|
||||
module.num_local_heads = kwargs["num_heads"] // parallel.attn_tp_size
|
||||
|
||||
class FakeColumnParallelLinear(nn.Module):
|
||||
def __init__(self, input_size, output_size, **kwargs):
|
||||
super().__init__()
|
||||
captured.update(kwargs)
|
||||
self.output_size_per_partition = output_size // kwargs["tp_size"]
|
||||
|
||||
monkeypatch.setattr(
|
||||
hunyuan_v4.DeepseekV2AttentionMLA, "__init__", fake_attention_init
|
||||
)
|
||||
monkeypatch.setattr(hunyuan_v4, "ColumnParallelLinear", FakeColumnParallelLinear)
|
||||
monkeypatch.setattr(hunyuan_v4, "get_parallel", lambda: parallel)
|
||||
monkeypatch.setattr(
|
||||
hunyuan_v4.HYV4Attention,
|
||||
"_hpc_gated_mla_supported",
|
||||
staticmethod(lambda *args: False),
|
||||
)
|
||||
|
||||
config = SimpleNamespace(
|
||||
rope_parameters={"rope_theta": 10_000, "rope_type": "default"},
|
||||
hidden_size=6144,
|
||||
num_attention_heads=64,
|
||||
qk_nope_head_dim=8,
|
||||
qk_rope_head_dim=4,
|
||||
v_head_dim=256,
|
||||
q_lora_rank=32,
|
||||
kv_lora_rank=16,
|
||||
max_position_embeddings=1024,
|
||||
gating_type="elementwise",
|
||||
)
|
||||
|
||||
attention = hunyuan_v4.HYV4Attention(config, layer_id=0)
|
||||
|
||||
assert captured["tp_rank"] == parallel.attn_tp_rank
|
||||
assert captured["tp_size"] == parallel.attn_tp_size
|
||||
assert attention.local_gate_width == (64 // attn_tp_size) * 256
|
||||
assert attention.linear_gate.output_size_per_partition == attention.local_gate_width
|
||||
|
||||
|
||||
class TupleLinear(nn.Module):
|
||||
def __init__(self, input_size, output_size, dtype):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(
|
||||
torch.randn(output_size, input_size, dtype=dtype), requires_grad=False
|
||||
)
|
||||
|
||||
def forward(self, inputs):
|
||||
return nn.functional.linear(inputs, self.weight), None
|
||||
|
||||
|
||||
def test_attention_gate_non_bf16_model_fallback_parity():
|
||||
torch.manual_seed(0)
|
||||
attention = hunyuan_v4.HYV4Attention.__new__(hunyuan_v4.HYV4Attention)
|
||||
nn.Module.__init__(attention)
|
||||
attention.linear_gate = TupleLinear(8, 256, torch.float32)
|
||||
attention.local_gate_width = 256
|
||||
attention._gate_backend = "eager"
|
||||
attention._gate_fallback_backend = "eager"
|
||||
hidden_states = torch.randn(3, 8)
|
||||
attn_out = torch.randn(3, 256)
|
||||
|
||||
gate = attention.prepare_attention_output_gate(hidden_states)
|
||||
actual = attention.apply_attention_output_gate(attn_out, gate)
|
||||
expected = attn_out * torch.sigmoid(
|
||||
nn.functional.linear(hidden_states, attention.linear_gate.weight)
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual, expected)
|
||||
|
||||
|
||||
def test_prepared_attention_gate_requires_model_application_hook():
|
||||
with pytest.raises(RuntimeError, match="unsigmoided"):
|
||||
forward_mla._apply_attention_output_gate(
|
||||
SimpleNamespace(), torch.ones(1), torch.ones(1)
|
||||
)
|
||||
|
||||
|
||||
def test_hpc_attention_gate_is_bf16_only(monkeypatch):
|
||||
fake_hpc = SimpleNamespace(
|
||||
gemm=SimpleNamespace(gated_mla_gemm=object()), __version__="test"
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "hpc", fake_hpc)
|
||||
monkeypatch.setattr(hunyuan_v4, "get_device_capability", lambda: (10, 0))
|
||||
supported = hunyuan_v4.HYV4Attention._hpc_gated_mla_supported
|
||||
supported.cache_clear()
|
||||
try:
|
||||
assert supported("elementwise", torch.bfloat16, (256, 6144), 256, 6144)
|
||||
assert not supported("elementwise", torch.float32, (256, 6144), 256, 6144)
|
||||
finally:
|
||||
supported.cache_clear()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,42 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.models.hunyuan_v4_nextn import HYV4ForCausalLMNextN
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=4, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestHunyuanV4NextNWeightLoading(unittest.TestCase):
|
||||
def test_indexer_checkpoint_layout_is_permuted(self):
|
||||
model = object.__new__(HYV4ForCausalLMNextN)
|
||||
model.config = SimpleNamespace(
|
||||
num_hidden_layers=80,
|
||||
index_n_heads=2,
|
||||
index_head_dim=4,
|
||||
qk_rope_head_dim=2,
|
||||
)
|
||||
captured = []
|
||||
object.__setattr__(
|
||||
model,
|
||||
"do_load_weights",
|
||||
lambda weights, **kwargs: captured.extend(weights),
|
||||
)
|
||||
loaded_weight = torch.arange(8).reshape(8, 1)
|
||||
|
||||
model.load_weights(
|
||||
[("model.mtp_layers.0.self_attn.indexer.wq_b.weight", loaded_weight)]
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
captured[0][0], "model.layers.80.self_attn.indexer.wq_b.weight"
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
captured[0][1].flatten(), torch.tensor([2, 3, 0, 1, 6, 7, 4, 5])
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,43 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
|
||||
from sglang.srt.parser.hunyuan_reasoning import normalize_hunyuan_reasoning_effort
|
||||
from sglang.srt.parser.template_detection import ReasoningToggleConfig
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("effort", "normalized"),
|
||||
[
|
||||
(None, "high"),
|
||||
("none", "no_think"),
|
||||
("minimal", "low"),
|
||||
("low", "low"),
|
||||
("medium", "high"),
|
||||
("high", "high"),
|
||||
("xhigh", "high"),
|
||||
("max", "high"),
|
||||
],
|
||||
)
|
||||
def test_hunyuan_reasoning_effort_normalization(effort, normalized):
|
||||
request = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
reasoning_effort=effort,
|
||||
)
|
||||
|
||||
normalize_hunyuan_reasoning_effort(
|
||||
request,
|
||||
reasoning_parser="hunyuan",
|
||||
reasoning_config=ReasoningToggleConfig(special_case="hunyuan_effort"),
|
||||
)
|
||||
|
||||
assert request.reasoning_effort == normalized
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -269,6 +269,42 @@ class TestTemplateManagerReasoningDetection(unittest.TestCase):
|
||||
_, _, parser = self._detect(template, ["<minimax:tool_call>"])
|
||||
self.assertEqual(parser, "minimax")
|
||||
|
||||
HYV4_TEMPLATE = (
|
||||
"{%- set reasoning_mode_token = '<|reasoning_mode:opensource|>' %}\n"
|
||||
"{%- if not reasoning_effort is defined %}\n"
|
||||
" {%- set reasoning_effort = 'high' %}\n"
|
||||
"{%- elif reasoning_effort not in ['high', 'low', 'no_think'] %}\n"
|
||||
"{%- endif %}\n"
|
||||
"<tool_call:opensource>{{ name }}<arg_key:opensource>{{ k }}</arg_key:opensource>"
|
||||
)
|
||||
|
||||
HYV4_VOCAB = [
|
||||
"<tool_calls:opensource>",
|
||||
"<tool_call:opensource>",
|
||||
"<arg_key:opensource>",
|
||||
"<arg_value:opensource>",
|
||||
]
|
||||
|
||||
def test_hyv4_effort_template_detected_with_special_case(self):
|
||||
# Hy4 drops <tool_sep>; detection must key on the effort-mode template
|
||||
# signature plus the suffixed arg tokens instead.
|
||||
force, config, parser = self._detect(self.HYV4_TEMPLATE, self.HYV4_VOCAB)
|
||||
|
||||
self.assertEqual(config, ReasoningToggleConfig(special_case="hunyuan_effort"))
|
||||
self.assertEqual(parser, "hunyuan")
|
||||
self.assertEqual(
|
||||
detect_tool_call_parser(
|
||||
self.HYV4_TEMPLATE, _DummyTokenizer(self.HYV4_VOCAB), config, force
|
||||
),
|
||||
"hunyuan",
|
||||
)
|
||||
|
||||
def test_hyv4_template_without_arg_tokens_not_hunyuan(self):
|
||||
_, config, parser = self._detect(self.HYV4_TEMPLATE, ["<tool_call:opensource>"])
|
||||
|
||||
self.assertEqual(config, ReasoningToggleConfig(special_case="hunyuan_effort"))
|
||||
self.assertNotEqual(parser, "hunyuan")
|
||||
|
||||
|
||||
class TestTemplateDetectionRuleMatrix(unittest.TestCase):
|
||||
"""Table-driven tests for REASONING_PARSER_RULES and REASONING_MODE_RULES."""
|
||||
|
||||
@@ -1037,6 +1037,7 @@ class TestContextParallelServerArgs(CustomTestCase):
|
||||
moe_dp_size=1,
|
||||
ep_size=1,
|
||||
pp_size=1,
|
||||
dcp_size=1,
|
||||
enable_aiter_allreduce_fusion=False,
|
||||
)
|
||||
defaults.update(overrides)
|
||||
@@ -1169,6 +1170,29 @@ class TestContextParallelServerArgs(CustomTestCase):
|
||||
resolution_result(args, "dsa_prefill_cp_mode"), "round-robin-split"
|
||||
)
|
||||
|
||||
def test_canonical_interleave_cp_mirrors_to_dsa_runtime_aliases(self):
|
||||
server_args = self._new_cp_args(
|
||||
enable_prefill_cp=True,
|
||||
cp_strategy="interleave",
|
||||
attention_backend="dsa",
|
||||
)
|
||||
|
||||
handle_legacy_cp_runtime_compatibility(server_args)
|
||||
handle_context_parallelism(server_args)
|
||||
|
||||
self.assertTrue(
|
||||
resolution_result(server_args, "enable_dsa_prefill_context_parallel")
|
||||
)
|
||||
self.assertFalse(
|
||||
resolution_result(server_args, "enable_prefill_context_parallel")
|
||||
)
|
||||
self.assertEqual(
|
||||
resolution_result(server_args, "dsa_prefill_cp_mode"), "round-robin-split"
|
||||
)
|
||||
self.assertEqual(
|
||||
resolution_result(server_args, "prefill_cp_mode"), "round-robin-split"
|
||||
)
|
||||
|
||||
def test_context_parallel_handler_initializes_cp_strategy(self):
|
||||
server_args = self._new_cp_args(
|
||||
enable_prefill_cp=True,
|
||||
|
||||
@@ -1817,8 +1817,8 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
||||
_dsa_split_backend_resolution,
|
||||
)
|
||||
|
||||
def _view(arch="DeepseekV32ForCausalLM", **kw):
|
||||
hf = SimpleNamespace(architectures=[arch])
|
||||
def _view(arch="DeepseekV32ForCausalLM", learnable_sink=False, **kw):
|
||||
hf = SimpleNamespace(architectures=[arch], learnable_sink=learnable_sink)
|
||||
defaults = dict(
|
||||
kv_cache_dtype="fp8_e4m3",
|
||||
dsa_prefill_backend=None,
|
||||
@@ -1866,6 +1866,28 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
||||
"dsa_decode_backend": "flashmla_kv",
|
||||
},
|
||||
)
|
||||
for arch in ("HYV4ForCausalLM", "HYV4ForCausalLMNextN"):
|
||||
with self.subTest(arch=arch, backends="default"):
|
||||
self.assertEqual(
|
||||
_dsa_split_backend_resolution(
|
||||
_view(arch=arch, learnable_sink=True)
|
||||
),
|
||||
{
|
||||
"dsa_prefill_backend": "flashmla_sparse",
|
||||
"dsa_decode_backend": "flashmla_sparse",
|
||||
},
|
||||
)
|
||||
for field, value in (
|
||||
("dsa_prefill_backend", "fa3"),
|
||||
("dsa_decode_backend", "trtllm"),
|
||||
):
|
||||
with self.subTest(arch=arch, field=field, value=value):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, field.replace("_", "-")
|
||||
):
|
||||
_dsa_split_backend_resolution(
|
||||
_view(arch=arch, learnable_sink=True, **{field: value})
|
||||
)
|
||||
# non-family arch declares nothing
|
||||
self.assertEqual(
|
||||
_dsa_split_backend_resolution(_view(arch="LlamaForCausalLM")), {}
|
||||
@@ -2842,6 +2864,7 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
||||
prefill_attention_backend=None,
|
||||
decode_attention_backend=None,
|
||||
enable_prefill_cp=False,
|
||||
dcp_size=1,
|
||||
)
|
||||
defaults.update(kw)
|
||||
return SimpleNamespace(**defaults)
|
||||
@@ -2857,6 +2880,22 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
||||
_deepseek_family_overrides(_args(), None),
|
||||
{"attention_backend": "dsa", "page_size": 64},
|
||||
)
|
||||
for arch in ("HYV4ForCausalLM", "HYV4ForCausalLMNextN"):
|
||||
hf_config = SimpleNamespace(architectures=[arch])
|
||||
with self.subTest(arch=arch, prefill_cp=True):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "--enable-prefill-cp.*HYV4"
|
||||
):
|
||||
_deepseek_family_overrides(
|
||||
_args(enable_prefill_cp=True), hf_config
|
||||
)
|
||||
with self.subTest(arch=arch, dcp_size=2):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "--dcp-size > 1.*HYV4"
|
||||
):
|
||||
_deepseek_family_overrides(
|
||||
_args(dcp_size=2), hf_config
|
||||
)
|
||||
# HIP without the preshuffle path: page 1
|
||||
with override_platform(is_hip=True):
|
||||
with patch(
|
||||
|
||||
Reference in New Issue
Block a user