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
@@ -514,6 +514,8 @@ _PHASE25_KERNELS = [
|
|||||||
("gemma4_fused_ops", "gemma4_fused_routing", "triton"),
|
("gemma4_fused_ops", "gemma4_fused_routing", "triton"),
|
||||||
("gemma4_fused_ops", "gemma_qkv_rmsnorm", "triton"),
|
("gemma4_fused_ops", "gemma_qkv_rmsnorm", "triton"),
|
||||||
("mhc_head", "fused_hc_head", "triton"),
|
("mhc_head", "fused_hc_head", "triton"),
|
||||||
|
("hy4_ihc", "fused_hy4_ihc_pre", "triton"),
|
||||||
|
("hy4_ihc", "fused_hy4_ihc_post", "triton"),
|
||||||
]
|
]
|
||||||
for _mod, _fn, _bk in _PHASE25_KERNELS:
|
for _mod, _fn, _bk in _PHASE25_KERNELS:
|
||||||
register_kernel(
|
register_kernel(
|
||||||
|
|||||||
@@ -0,0 +1,443 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import functools
|
||||||
|
import logging
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import triton
|
||||||
|
import triton.language as tl
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
_HPC_IHC_CAPABILITIES = ((9, 0), (10, 0), (10, 3))
|
||||||
|
_HPC_IHC_HC_MULTS = (4,)
|
||||||
|
_HPC_IHC_HIDDEN_SIZES = (4096, 6144)
|
||||||
|
|
||||||
|
|
||||||
|
@functools.lru_cache(maxsize=None)
|
||||||
|
def _hpc_ihc_op(op_name: str, hc_mult: int, hidden_size: int):
|
||||||
|
try:
|
||||||
|
import hpc
|
||||||
|
except ImportError:
|
||||||
|
return None
|
||||||
|
|
||||||
|
op = getattr(hpc, op_name, None)
|
||||||
|
if op is None:
|
||||||
|
logger.info(
|
||||||
|
"HY4 iHC: the installed hpc build (%s) has no %s; using the "
|
||||||
|
"in-tree Triton kernels.",
|
||||||
|
getattr(hpc, "__version__", "unknown"),
|
||||||
|
op_name,
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
|
||||||
|
from sglang.srt.utils import get_device_capability
|
||||||
|
|
||||||
|
cap = get_device_capability()
|
||||||
|
if cap not in _HPC_IHC_CAPABILITIES:
|
||||||
|
logger.warning(
|
||||||
|
"HY4 iHC: hpc.%s is unavailable on sm%s%s.",
|
||||||
|
op_name,
|
||||||
|
*cap,
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
if hc_mult not in _HPC_IHC_HC_MULTS or hidden_size not in _HPC_IHC_HIDDEN_SIZES:
|
||||||
|
logger.warning(
|
||||||
|
"HY4 iHC: hpc.%s is instantiated for hc_mult in %s and hidden_size "
|
||||||
|
"in %s, got %d / %d.",
|
||||||
|
op_name,
|
||||||
|
_HPC_IHC_HC_MULTS,
|
||||||
|
_HPC_IHC_HIDDEN_SIZES,
|
||||||
|
hc_mult,
|
||||||
|
hidden_size,
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
|
||||||
|
logger.info("HY4 iHC: using hpc.%s.", op_name)
|
||||||
|
return op
|
||||||
|
|
||||||
|
|
||||||
|
@triton.jit
|
||||||
|
def _hy4_ihc_pre_stage1(
|
||||||
|
x_ptr,
|
||||||
|
fn_ptr,
|
||||||
|
part_ptr,
|
||||||
|
K_TOTAL: tl.constexpr,
|
||||||
|
HC_MULT: tl.constexpr,
|
||||||
|
HC_POW2: tl.constexpr,
|
||||||
|
NSPLIT: tl.constexpr,
|
||||||
|
BLOCK_K: tl.constexpr,
|
||||||
|
PART_STRIDE: tl.constexpr,
|
||||||
|
):
|
||||||
|
pid_t = tl.program_id(0).to(tl.int64)
|
||||||
|
pid_s = tl.program_id(1)
|
||||||
|
m_idx = tl.arange(0, HC_POW2)
|
||||||
|
m_mask = m_idx < HC_MULT
|
||||||
|
|
||||||
|
k_offs = pid_s * BLOCK_K + tl.arange(0, BLOCK_K)
|
||||||
|
k_mask = k_offs < K_TOTAL
|
||||||
|
|
||||||
|
x_tile = tl.load(x_ptr + pid_t * K_TOTAL + 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,
|
||||||
|
)
|
||||||
|
|
||||||
|
base = part_ptr + (pid_t * NSPLIT + pid_s) * PART_STRIDE
|
||||||
|
tl.store(base, sumsq)
|
||||||
|
tl.store(base + 1 + m_idx, mix_pre, mask=m_mask)
|
||||||
|
tl.store(base + 1 + HC_POW2 + m_idx, mix_post, mask=m_mask)
|
||||||
|
|
||||||
|
|
||||||
|
@triton.jit
|
||||||
|
def _hy4_ihc_pre_stage2(
|
||||||
|
x_ptr,
|
||||||
|
part_ptr,
|
||||||
|
scale_ptr,
|
||||||
|
base_ptr,
|
||||||
|
y_ptr,
|
||||||
|
post_ptr,
|
||||||
|
hidden_size: tl.constexpr,
|
||||||
|
K_TOTAL: tl.constexpr,
|
||||||
|
HC_MULT: tl.constexpr,
|
||||||
|
HC_POW2: tl.constexpr,
|
||||||
|
NSPLIT: tl.constexpr,
|
||||||
|
PART_STRIDE: tl.constexpr,
|
||||||
|
BLOCK_D: tl.constexpr,
|
||||||
|
magnitude: tl.constexpr,
|
||||||
|
norm_eps: tl.constexpr,
|
||||||
|
hc_eps: tl.constexpr,
|
||||||
|
):
|
||||||
|
pid_t = tl.program_id(0).to(tl.int64)
|
||||||
|
pid_d = tl.program_id(1)
|
||||||
|
m_idx = tl.arange(0, HC_POW2)
|
||||||
|
m_mask = m_idx < HC_MULT
|
||||||
|
|
||||||
|
# The single-CTA kernel folded the BLOCK_K tiles into one accumulator with
|
||||||
|
# '+='; replay that ascending order over the partials so the fp32 sum is
|
||||||
|
# bit-identical.
|
||||||
|
row = part_ptr + pid_t * NSPLIT * PART_STRIDE
|
||||||
|
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 s in tl.static_range(NSPLIT):
|
||||||
|
b = row + s * PART_STRIDE
|
||||||
|
sumsq += tl.load(b)
|
||||||
|
mix_pre += tl.load(b + 1 + m_idx, mask=m_mask, other=0.0)
|
||||||
|
mix_post += tl.load(b + 1 + HC_POW2 + m_idx, mask=m_mask, other=0.0)
|
||||||
|
|
||||||
|
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
|
||||||
|
if pid_d == 0:
|
||||||
|
post = (
|
||||||
|
magnitude * tl.sigmoid(mix_post * rsqrt * scale_post + base_post) + hc_eps
|
||||||
|
)
|
||||||
|
tl.store(post_ptr + pid_t * HC_MULT + m_idx, post, mask=m_mask)
|
||||||
|
|
||||||
|
x_row = x_ptr + pid_t * K_TOTAL
|
||||||
|
d_offs = pid_d * BLOCK_D + 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_ptr + pid_t * hidden_size + d_offs,
|
||||||
|
y_block.to(y_ptr.dtype.element_ty),
|
||||||
|
mask=d_mask,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@triton.jit
|
||||||
|
def _hy4_ihc_post_kernel(
|
||||||
|
out_ptr,
|
||||||
|
res_ptr,
|
||||||
|
post_ptr,
|
||||||
|
y_ptr,
|
||||||
|
hidden_size: tl.constexpr,
|
||||||
|
HC_MULT: tl.constexpr,
|
||||||
|
HC_POW2: tl.constexpr,
|
||||||
|
BLOCK_D: tl.constexpr,
|
||||||
|
):
|
||||||
|
pid_t = tl.program_id(0).to(tl.int64)
|
||||||
|
pid_d = tl.program_id(1)
|
||||||
|
|
||||||
|
m_idx = tl.arange(0, HC_POW2)
|
||||||
|
m_mask = m_idx < HC_MULT
|
||||||
|
post = tl.load(post_ptr + pid_t * HC_MULT + m_idx, mask=m_mask, other=0.0)
|
||||||
|
|
||||||
|
d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||||
|
d_mask = d_offs < hidden_size
|
||||||
|
out_block = tl.load(out_ptr + pid_t * hidden_size + d_offs, mask=d_mask, other=0.0)
|
||||||
|
out_block = out_block.to(tl.float32)
|
||||||
|
|
||||||
|
res_row = res_ptr + pid_t * HC_MULT * hidden_size
|
||||||
|
y_row = y_ptr + pid_t * HC_MULT * hidden_size
|
||||||
|
for m in tl.static_range(HC_MULT):
|
||||||
|
res_block = tl.load(res_row + m * hidden_size + d_offs, mask=d_mask, other=0.0)
|
||||||
|
post_m = tl.sum(tl.where(m_idx == m, post, 0.0), axis=0)
|
||||||
|
y_block = post_m * out_block + res_block.to(tl.float32)
|
||||||
|
tl.store(
|
||||||
|
y_row + m * hidden_size + d_offs,
|
||||||
|
y_block.to(y_ptr.dtype.element_ty),
|
||||||
|
mask=d_mask,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def fused_hy4_ihc_pre(
|
||||||
|
x: torch.Tensor,
|
||||||
|
hc_fn: torch.Tensor,
|
||||||
|
hc_scale: torch.Tensor,
|
||||||
|
hc_base: torch.Tensor,
|
||||||
|
magnitude: float,
|
||||||
|
norm_eps: float,
|
||||||
|
hc_eps: float,
|
||||||
|
rms_weight: torch.Tensor | None = None,
|
||||||
|
rms_eps: float = 0.0,
|
||||||
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||||
|
assert x.dim() == 3, f"x must be 3D (T, hc_mult, hidden_size), got {x.shape}"
|
||||||
|
assert hc_fn.dtype == torch.float32
|
||||||
|
assert hc_scale.dtype == torch.float32 and hc_base.dtype == torch.float32
|
||||||
|
|
||||||
|
x = x.contiguous()
|
||||||
|
hc_fn = hc_fn.contiguous()
|
||||||
|
T, hc_mult, hidden_size = x.shape
|
||||||
|
k_total = hc_mult * hidden_size
|
||||||
|
assert hc_fn.shape == (2 * hc_mult, k_total)
|
||||||
|
assert hc_base.shape == (2 * hc_mult,)
|
||||||
|
assert hc_scale.shape == (2,)
|
||||||
|
|
||||||
|
if T == 0:
|
||||||
|
return (
|
||||||
|
torch.empty((0, hidden_size), dtype=x.dtype, device=x.device),
|
||||||
|
torch.empty((0, hc_mult), dtype=torch.float32, device=x.device),
|
||||||
|
)
|
||||||
|
|
||||||
|
hpc_op = _hpc_ihc_op("fuse_ihc_pre", hc_mult, hidden_size)
|
||||||
|
if hpc_op is not None:
|
||||||
|
return hpc_op(
|
||||||
|
x,
|
||||||
|
hc_fn,
|
||||||
|
hc_scale,
|
||||||
|
hc_base,
|
||||||
|
norm_eps,
|
||||||
|
hc_eps,
|
||||||
|
magnitude,
|
||||||
|
rms_weight,
|
||||||
|
rms_eps,
|
||||||
|
True,
|
||||||
|
)
|
||||||
|
|
||||||
|
y = torch.empty((T, hidden_size), dtype=x.dtype, device=x.device)
|
||||||
|
post = torch.empty((T, hc_mult), dtype=torch.float32, device=x.device)
|
||||||
|
BLOCK_K = 1024
|
||||||
|
BLOCK_D = 1024
|
||||||
|
hc_pow2 = triton.next_power_of_2(hc_mult)
|
||||||
|
# One BLOCK_K tile per CTA: the partials then replay the old kernel's
|
||||||
|
# per-tile accumulation order exactly.
|
||||||
|
nsplit = triton.cdiv(k_total, BLOCK_K)
|
||||||
|
part_stride = 1 + 2 * hc_pow2
|
||||||
|
part = torch.empty((T, nsplit, part_stride), dtype=torch.float32, device=x.device)
|
||||||
|
|
||||||
|
_hy4_ihc_pre_stage1[(T, nsplit)](
|
||||||
|
x,
|
||||||
|
hc_fn,
|
||||||
|
part,
|
||||||
|
K_TOTAL=k_total,
|
||||||
|
HC_MULT=hc_mult,
|
||||||
|
HC_POW2=hc_pow2,
|
||||||
|
NSPLIT=nsplit,
|
||||||
|
BLOCK_K=BLOCK_K,
|
||||||
|
PART_STRIDE=part_stride,
|
||||||
|
num_warps=8,
|
||||||
|
# Disable FMA: eager rounds the fp32 product before summation.
|
||||||
|
enable_fp_fusion=False,
|
||||||
|
)
|
||||||
|
_hy4_ihc_pre_stage2[(T, triton.cdiv(hidden_size, BLOCK_D))](
|
||||||
|
x,
|
||||||
|
part,
|
||||||
|
hc_scale,
|
||||||
|
hc_base,
|
||||||
|
y,
|
||||||
|
post,
|
||||||
|
hidden_size=hidden_size,
|
||||||
|
K_TOTAL=k_total,
|
||||||
|
HC_MULT=hc_mult,
|
||||||
|
HC_POW2=hc_pow2,
|
||||||
|
NSPLIT=nsplit,
|
||||||
|
PART_STRIDE=part_stride,
|
||||||
|
BLOCK_D=BLOCK_D,
|
||||||
|
magnitude=magnitude,
|
||||||
|
norm_eps=norm_eps,
|
||||||
|
hc_eps=hc_eps,
|
||||||
|
num_warps=4,
|
||||||
|
enable_fp_fusion=False,
|
||||||
|
)
|
||||||
|
if rms_weight is not None:
|
||||||
|
y_float = y.float()
|
||||||
|
y = (
|
||||||
|
y_float
|
||||||
|
* torch.rsqrt(y_float.square().mean(dim=-1, keepdim=True) + rms_eps)
|
||||||
|
* rms_weight.float()
|
||||||
|
).to(y.dtype)
|
||||||
|
return y, post
|
||||||
|
|
||||||
|
|
||||||
|
def fused_hy4_ihc_post(
|
||||||
|
output: torch.Tensor,
|
||||||
|
residual: torch.Tensor,
|
||||||
|
post: torch.Tensor,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
assert output.dim() == 2, f"output must be 2D (T, hidden_size), got {output.shape}"
|
||||||
|
assert post.dtype == torch.float32
|
||||||
|
|
||||||
|
output = output.contiguous()
|
||||||
|
residual = residual.contiguous()
|
||||||
|
T, hidden_size = output.shape
|
||||||
|
hc_mult = post.shape[-1]
|
||||||
|
assert residual.shape == (T, hc_mult, hidden_size)
|
||||||
|
assert post.shape == (T, hc_mult)
|
||||||
|
|
||||||
|
if T == 0:
|
||||||
|
return torch.empty(
|
||||||
|
(0, hc_mult, hidden_size), dtype=output.dtype, device=output.device
|
||||||
|
)
|
||||||
|
|
||||||
|
hpc_op = _hpc_ihc_op("fuse_ihc_post", hc_mult, hidden_size)
|
||||||
|
if hpc_op is not None:
|
||||||
|
return hpc_op(output, residual, post)
|
||||||
|
|
||||||
|
y = torch.empty((T, hc_mult, hidden_size), dtype=output.dtype, device=output.device)
|
||||||
|
BLOCK_D = 1024
|
||||||
|
grid = (T, triton.cdiv(hidden_size, BLOCK_D))
|
||||||
|
_hy4_ihc_post_kernel[grid](
|
||||||
|
output,
|
||||||
|
residual,
|
||||||
|
post.contiguous(),
|
||||||
|
y,
|
||||||
|
hidden_size=hidden_size,
|
||||||
|
HC_MULT=hc_mult,
|
||||||
|
HC_POW2=triton.next_power_of_2(hc_mult),
|
||||||
|
BLOCK_D=BLOCK_D,
|
||||||
|
num_warps=4,
|
||||||
|
enable_fp_fusion=False,
|
||||||
|
)
|
||||||
|
return y
|
||||||
|
|
||||||
|
|
||||||
|
def fused_hy4_ihc_post_pre(
|
||||||
|
output: torch.Tensor,
|
||||||
|
residual: torch.Tensor,
|
||||||
|
post: torch.Tensor,
|
||||||
|
hc_fn: torch.Tensor,
|
||||||
|
hc_scale: torch.Tensor,
|
||||||
|
hc_base: torch.Tensor,
|
||||||
|
magnitude: float,
|
||||||
|
norm_eps: float,
|
||||||
|
hc_eps: float,
|
||||||
|
rms_weight: torch.Tensor | None = None,
|
||||||
|
rms_eps: float = 0.0,
|
||||||
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||||
|
num_tokens, hidden_size = output.shape
|
||||||
|
hc_mult = residual.shape[1]
|
||||||
|
if num_tokens == 0:
|
||||||
|
return (
|
||||||
|
torch.empty_like(residual),
|
||||||
|
torch.empty_like(output),
|
||||||
|
torch.empty((0, hc_mult), dtype=torch.float32, device=output.device),
|
||||||
|
)
|
||||||
|
|
||||||
|
hpc_op = _hpc_ihc_op("fuse_ihc_post_pre", hc_mult, hidden_size)
|
||||||
|
if hpc_op is not None:
|
||||||
|
return hpc_op(
|
||||||
|
output.contiguous(),
|
||||||
|
residual.contiguous(),
|
||||||
|
post.contiguous(),
|
||||||
|
hc_fn.contiguous(),
|
||||||
|
hc_scale.contiguous(),
|
||||||
|
hc_base.contiguous(),
|
||||||
|
norm_eps,
|
||||||
|
hc_eps,
|
||||||
|
magnitude,
|
||||||
|
rms_weight,
|
||||||
|
rms_eps,
|
||||||
|
True,
|
||||||
|
)
|
||||||
|
|
||||||
|
next_residual = fused_hy4_ihc_post(output, residual, post)
|
||||||
|
reduced, next_post = fused_hy4_ihc_pre(
|
||||||
|
next_residual,
|
||||||
|
hc_fn,
|
||||||
|
hc_scale,
|
||||||
|
hc_base,
|
||||||
|
magnitude,
|
||||||
|
norm_eps,
|
||||||
|
hc_eps,
|
||||||
|
rms_weight,
|
||||||
|
rms_eps,
|
||||||
|
)
|
||||||
|
return next_residual, reduced, next_post
|
||||||
|
|
||||||
|
|
||||||
|
def fused_hy4_ihc_head(
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
hc_fn: torch.Tensor,
|
||||||
|
hc_scale: torch.Tensor,
|
||||||
|
hc_base: torch.Tensor,
|
||||||
|
norm_eps: float,
|
||||||
|
hc_eps: float,
|
||||||
|
rms_weight: torch.Tensor | None = None,
|
||||||
|
rms_eps: float = 0.0,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
num_tokens, hc_mult, hidden_size = hidden_states.shape
|
||||||
|
if num_tokens == 0:
|
||||||
|
return torch.empty(
|
||||||
|
(0, hidden_size), dtype=hidden_states.dtype, device=hidden_states.device
|
||||||
|
)
|
||||||
|
|
||||||
|
hpc_op = _hpc_ihc_op("fuse_ihc_head", hc_mult, hidden_size)
|
||||||
|
if hpc_op is not None:
|
||||||
|
return hpc_op(
|
||||||
|
hidden_states.contiguous(),
|
||||||
|
hc_fn.contiguous(),
|
||||||
|
hc_scale.contiguous(),
|
||||||
|
hc_base.contiguous(),
|
||||||
|
norm_eps,
|
||||||
|
hc_eps,
|
||||||
|
rms_weight,
|
||||||
|
rms_eps,
|
||||||
|
True,
|
||||||
|
)
|
||||||
|
|
||||||
|
flat = hidden_states.flatten(1).float()
|
||||||
|
scale = torch.rsqrt(flat.square().mean(-1, keepdim=True) + norm_eps)
|
||||||
|
gates = torch.nn.functional.linear(flat, hc_fn) * scale
|
||||||
|
gates = torch.sigmoid(gates * hc_scale + hc_base) + hc_eps
|
||||||
|
output = torch.sum(gates.unsqueeze(-1) * hidden_states.float(), dim=1).to(
|
||||||
|
hidden_states.dtype
|
||||||
|
)
|
||||||
|
if rms_weight is not None:
|
||||||
|
output_float = output.float()
|
||||||
|
output = (
|
||||||
|
output_float
|
||||||
|
* torch.rsqrt(output_float.square().mean(dim=-1, keepdim=True) + rms_eps)
|
||||||
|
* rms_weight.float()
|
||||||
|
).to(output.dtype)
|
||||||
|
return output
|
||||||
@@ -1155,9 +1155,12 @@ def ep_scatter(
|
|||||||
output_index: torch.Tensor,
|
output_index: torch.Tensor,
|
||||||
scale_ue8m0: bool = False,
|
scale_ue8m0: bool = False,
|
||||||
quant_block_size: int = 128,
|
quant_block_size: int = 128,
|
||||||
|
expert_alignment: int = 128,
|
||||||
expert_start: int = 0,
|
expert_start: int = 0,
|
||||||
):
|
):
|
||||||
BLOCK_E = 128 # token num of per expert is aligned to 128
|
# tl.arange needs pow2, and the kernel's unmasked stores need BLOCK_E to
|
||||||
|
# divide the expert_alignment-padded segments; lowbit satisfies both.
|
||||||
|
BLOCK_E = expert_alignment & -expert_alignment
|
||||||
BLOCK_D = quant_block_size # block size of quantization
|
BLOCK_D = quant_block_size # block size of quantization
|
||||||
num_warps = 8
|
num_warps = 8
|
||||||
num_experts = num_recv_tokens_per_expert.shape[0]
|
num_experts = num_recv_tokens_per_expert.shape[0]
|
||||||
|
|||||||
@@ -0,0 +1,41 @@
|
|||||||
|
import torch
|
||||||
|
import triton
|
||||||
|
import triton.language as tl
|
||||||
|
|
||||||
|
|
||||||
|
@triton.jit
|
||||||
|
def _pad_expert_counts_kernel(
|
||||||
|
counts_ptr,
|
||||||
|
out_ptr,
|
||||||
|
num_experts,
|
||||||
|
all_tokens,
|
||||||
|
BLOCK_E: tl.constexpr,
|
||||||
|
NE_POW2: tl.constexpr,
|
||||||
|
):
|
||||||
|
i = tl.arange(0, NE_POW2)
|
||||||
|
m = i < num_experts
|
||||||
|
c = tl.load(counts_ptr + i, mask=m, other=0).to(tl.int32)
|
||||||
|
padded = ((c + BLOCK_E - 1) // BLOCK_E) * BLOCK_E
|
||||||
|
padded = tl.where(m, padded, 0)
|
||||||
|
# The trailing segment absorbs the difference so the total stays
|
||||||
|
# graph-static; its m_indices remain -1 and DeepGEMM skips those rows.
|
||||||
|
slack = all_tokens - tl.sum(padded, axis=0)
|
||||||
|
padded = tl.where(i == num_experts - 1, padded + slack, padded)
|
||||||
|
tl.store(out_ptr + i, padded, mask=m)
|
||||||
|
|
||||||
|
|
||||||
|
def pad_expert_counts(
|
||||||
|
counts: torch.Tensor, block_e: int, all_tokens: int
|
||||||
|
) -> torch.Tensor:
|
||||||
|
ne = counts.numel()
|
||||||
|
out = torch.empty(ne, dtype=torch.int32, device=counts.device)
|
||||||
|
_pad_expert_counts_kernel[(1,)](
|
||||||
|
counts,
|
||||||
|
out,
|
||||||
|
ne,
|
||||||
|
all_tokens,
|
||||||
|
BLOCK_E=block_e,
|
||||||
|
NE_POW2=triton.next_power_of_2(ne),
|
||||||
|
num_warps=4,
|
||||||
|
)
|
||||||
|
return out
|
||||||
@@ -37,6 +37,8 @@ def _sigmoid_gate_mul_kernel(
|
|||||||
|
|
||||||
def sigmoid_gate_mul(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
|
def sigmoid_gate_mul(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
|
||||||
"""Compute ``x * sigmoid(gate)`` in a single fused kernel (same-shape)."""
|
"""Compute ``x * sigmoid(gate)`` in a single fused kernel (same-shape)."""
|
||||||
|
assert x.shape == gate.shape, f"shape mismatch: {x.shape=} {gate.shape=}"
|
||||||
|
assert x.is_contiguous() and gate.is_contiguous(), "inputs must be contiguous"
|
||||||
out = torch.empty_like(x)
|
out = torch.empty_like(x)
|
||||||
n = x.numel()
|
n = x.numel()
|
||||||
grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)
|
grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)
|
||||||
|
|||||||
@@ -44,6 +44,30 @@ from sglang.srt.utils.common import (
|
|||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _validate_dsa_tbo_index_sharing(server_args: Any, hf_config: Any) -> None:
|
||||||
|
cfg = resolving_view(server_args)
|
||||||
|
if not cfg.enable_two_batch_overlap:
|
||||||
|
return
|
||||||
|
|
||||||
|
index_topk_freq = getattr(hf_config, "index_topk_freq", 1) or 1
|
||||||
|
index_topk_pattern = getattr(hf_config, "index_topk_pattern", None)
|
||||||
|
indexer_types = getattr(hf_config, "indexer_types", None)
|
||||||
|
if (
|
||||||
|
index_topk_freq > 1
|
||||||
|
or (index_topk_pattern is not None and "S" in index_topk_pattern)
|
||||||
|
or (indexer_types is not None and "shared" in indexer_types)
|
||||||
|
):
|
||||||
|
raise ValueError(
|
||||||
|
"--enable-two-batch-overlap is not supported with DSA "
|
||||||
|
"index-topk sharing: the TBO op path does not propagate topk "
|
||||||
|
"indices across layers, so shared layers would run sparse "
|
||||||
|
"attention without indices. Got "
|
||||||
|
f"index_topk_freq={index_topk_freq!r}, "
|
||||||
|
f"index_topk_pattern={index_topk_pattern!r}, and "
|
||||||
|
f"indexer_types={indexer_types!r}."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _rocm_fp8_wo_a_supported() -> bool:
|
def _rocm_fp8_wo_a_supported() -> bool:
|
||||||
"""True when ROCm can run the DeepSeek-V4 fp8 wo_a GEMM (gfx950 + aiter)."""
|
"""True when ROCm can run the DeepSeek-V4 fp8 wo_a GEMM (gfx950 + aiter)."""
|
||||||
try:
|
try:
|
||||||
@@ -159,6 +183,8 @@ def handle_model_specific_adjustments(server_args: Any):
|
|||||||
"MistralLarge3ForCausalLM",
|
"MistralLarge3ForCausalLM",
|
||||||
"PixtralForConditionalGeneration",
|
"PixtralForConditionalGeneration",
|
||||||
"GlmMoeDsaForCausalLM",
|
"GlmMoeDsaForCausalLM",
|
||||||
|
"HYV4ForCausalLM",
|
||||||
|
"HYV4ForCausalLMNextN",
|
||||||
"LongcatFlashForCausalLM",
|
"LongcatFlashForCausalLM",
|
||||||
"Dots3NoteForCausalLM",
|
"Dots3NoteForCausalLM",
|
||||||
]:
|
]:
|
||||||
@@ -181,19 +207,7 @@ def handle_model_specific_adjustments(server_args: Any):
|
|||||||
# The "dsa" attention fill moved to the override registry
|
# The "dsa" attention fill moved to the override registry
|
||||||
# (arg_groups/overrides.py: _deepseek_family_overrides).
|
# (arg_groups/overrides.py: _deepseek_family_overrides).
|
||||||
|
|
||||||
index_topk_freq = getattr(hf_config, "index_topk_freq", 1) or 1
|
_validate_dsa_tbo_index_sharing(server_args, hf_config)
|
||||||
index_topk_pattern = getattr(hf_config, "index_topk_pattern", None)
|
|
||||||
if cfg.enable_two_batch_overlap and (
|
|
||||||
index_topk_freq > 1
|
|
||||||
or (index_topk_pattern is not None and "S" in index_topk_pattern)
|
|
||||||
):
|
|
||||||
raise ValueError(
|
|
||||||
"--enable-two-batch-overlap is not supported with DSA "
|
|
||||||
"index-topk sharing (index_topk_freq > 1 or an "
|
|
||||||
"index_topk_pattern containing shared layers): the TBO op "
|
|
||||||
"path does not propagate topk indices across layers, so "
|
|
||||||
"shared layers would run sparse attention without indices."
|
|
||||||
)
|
|
||||||
|
|
||||||
if (
|
if (
|
||||||
not get_platform().is_npu and not get_platform().is_xpu
|
not get_platform().is_npu and not get_platform().is_xpu
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
"""Config-time override declarations for deepseek_v2.
|
"""Config-time override declarations for deepseek_v2.
|
||||||
|
|
||||||
Architectures: DeepseekV32ForCausalLM, DeepseekV3ForCausalLM, Dots3NoteForCausalLM, GlmMoeDsaForCausalLM, KimiK25ForConditionalGeneration, LongcatFlashForCausalLM, LongcatFlashForCausalLMNextN, MistralLarge3ForCausalLM, PixtralForConditionalGeneration.
|
Architectures: DeepseekV32ForCausalLM, DeepseekV3ForCausalLM, Dots3NoteForCausalLM, GlmMoeDsaForCausalLM, HYV4ForCausalLM, HYV4ForCausalLMNextN, KimiK25ForConditionalGeneration, LongcatFlashForCausalLM, LongcatFlashForCausalLMNextN, MistralLarge3ForCausalLM, PixtralForConditionalGeneration.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
@@ -24,6 +24,8 @@ logger = logging.getLogger(__name__)
|
|||||||
"MistralLarge3ForCausalLM",
|
"MistralLarge3ForCausalLM",
|
||||||
"PixtralForConditionalGeneration",
|
"PixtralForConditionalGeneration",
|
||||||
"GlmMoeDsaForCausalLM",
|
"GlmMoeDsaForCausalLM",
|
||||||
|
"HYV4ForCausalLM",
|
||||||
|
"HYV4ForCausalLMNextN",
|
||||||
"LongcatFlashForCausalLM",
|
"LongcatFlashForCausalLM",
|
||||||
"LongcatFlashForCausalLMNextN",
|
"LongcatFlashForCausalLMNextN",
|
||||||
"Dots3NoteForCausalLM",
|
"Dots3NoteForCausalLM",
|
||||||
@@ -35,10 +37,50 @@ def _deepseek_family_overrides(server_args: Any, hf_config: Any) -> dict:
|
|||||||
before it by _set_default_dsa_kv_cache_dtype) and the env writes stay in
|
before it by _set_default_dsa_kv_cache_dtype) and the env writes stay in
|
||||||
the branch."""
|
the branch."""
|
||||||
cfg = resolving_view(server_args)
|
cfg = resolving_view(server_args)
|
||||||
from sglang.srt.configs.model_config import is_deepseek_dsa
|
from sglang.srt.configs.model_config import (
|
||||||
|
is_deepseek_dsa,
|
||||||
|
unwrap_modelopt_quantization_config,
|
||||||
|
)
|
||||||
|
|
||||||
|
model_arch = (getattr(hf_config, "architectures", None) or [None])[0]
|
||||||
|
if model_arch in ("HYV4ForCausalLM", "HYV4ForCausalLMNextN"):
|
||||||
|
if cfg.enable_prefill_cp:
|
||||||
|
raise ValueError(
|
||||||
|
"--enable-prefill-cp is not supported for HYV4 because its "
|
||||||
|
"attention path does not implement DSA context-parallel metadata "
|
||||||
|
f"and sharding. Got architecture={model_arch!r} and "
|
||||||
|
f"enable_prefill_cp={cfg.enable_prefill_cp!r}."
|
||||||
|
)
|
||||||
|
dcp_size = getattr(cfg, "dcp_size", 1)
|
||||||
|
if dcp_size > 1:
|
||||||
|
raise ValueError(
|
||||||
|
"--dcp-size > 1 is not supported for HYV4 because decode context "
|
||||||
|
"parallelism gathers query heads across DCP ranks but does not "
|
||||||
|
"provide single-owner semantics for learnable attention sinks. "
|
||||||
|
f"Got architecture={model_arch!r} and dcp_size={dcp_size!r}."
|
||||||
|
)
|
||||||
|
|
||||||
overrides: Dict[str, Any] = {}
|
overrides: Dict[str, Any] = {}
|
||||||
|
|
||||||
|
if model_arch in ("HYV4ForCausalLM", "HYV4ForCausalLMNextN"):
|
||||||
|
quant_cfg = getattr(hf_config, "quantization_config", None) or {}
|
||||||
|
quant_algo = unwrap_modelopt_quantization_config(quant_cfg).get(
|
||||||
|
"quant_algo", ""
|
||||||
|
)
|
||||||
|
if str(quant_algo).upper() == "MXFP8":
|
||||||
|
from sglang.srt.layers import deep_gemm_wrapper
|
||||||
|
|
||||||
|
# auto would otherwise select an unqualified FP8/MoE path for HYV4 MXFP8.
|
||||||
|
if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
|
||||||
|
if cfg.moe_runner_backend == "auto":
|
||||||
|
overrides["moe_runner_backend"] = "deep_gemm"
|
||||||
|
if cfg.fp8_gemm_runner_backend == "auto":
|
||||||
|
overrides["fp8_gemm_runner_backend"] = "deep_gemm"
|
||||||
|
if overrides:
|
||||||
|
logger.info(
|
||||||
|
"HYV4 MXFP8: defaulting MoE/FP8 GEMM backends to deep_gemm."
|
||||||
|
)
|
||||||
|
|
||||||
if is_deepseek_dsa(hf_config): # DeepSeek 3.2/GLM 5
|
if is_deepseek_dsa(hf_config): # DeepSeek 3.2/GLM 5
|
||||||
# Set attention backend for DeepSeek
|
# Set attention backend for DeepSeek
|
||||||
if is_attention_backend_not_set(cfg):
|
if is_attention_backend_not_set(cfg):
|
||||||
|
|||||||
@@ -628,8 +628,16 @@ def _dsa_kv_cache_dtype_default(view: Any) -> dict:
|
|||||||
)
|
)
|
||||||
|
|
||||||
kv_cache_dtype = view.kv_cache_dtype
|
kv_cache_dtype = view.kv_cache_dtype
|
||||||
|
has_attention_sinks = bool(getattr(hf_config, "learnable_sink", False))
|
||||||
|
if has_attention_sinks and kv_cache_dtype not in ("auto", "bf16", "bfloat16"):
|
||||||
|
raise ValueError(
|
||||||
|
"Learnable DSA attention sinks require a bfloat16 KV cache; "
|
||||||
|
f"got kv_cache_dtype={kv_cache_dtype}."
|
||||||
|
)
|
||||||
if kv_cache_dtype == "auto":
|
if kv_cache_dtype == "auto":
|
||||||
kv_cache_dtype = "fp8_e4m3" if major >= 10 else "bfloat16"
|
kv_cache_dtype = (
|
||||||
|
"fp8_e4m3" if major >= 10 and not has_attention_sinks else "bfloat16"
|
||||||
|
)
|
||||||
logger.warning(
|
logger.warning(
|
||||||
f"Setting KV cache dtype to {kv_cache_dtype} for DeepSeek DSA on SM{major} device."
|
f"Setting KV cache dtype to {kv_cache_dtype} for DeepSeek DSA on SM{major} device."
|
||||||
)
|
)
|
||||||
@@ -697,6 +705,27 @@ def _dsa_split_backend_resolution(view: Any) -> dict:
|
|||||||
and not get_platform().is_hip
|
and not get_platform().is_hip
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if getattr(hf_config, "learnable_sink", False):
|
||||||
|
backend = "flashmla_sparse"
|
||||||
|
for field in ("dsa_prefill_backend", "dsa_decode_backend"):
|
||||||
|
value = getattr(view, field)
|
||||||
|
if value is not None and value != backend:
|
||||||
|
option = "--" + field.replace("_", "-")
|
||||||
|
raise ValueError(
|
||||||
|
f"{model_arch} uses learnable attention sinks and requires "
|
||||||
|
f"{option} {backend!r}; got {value!r}"
|
||||||
|
)
|
||||||
|
if not user_set_prefill:
|
||||||
|
declared["dsa_prefill_backend"] = backend
|
||||||
|
if not user_set_decode:
|
||||||
|
declared["dsa_decode_backend"] = backend
|
||||||
|
logger.warning(
|
||||||
|
"Set DSA backends for learnable attention sinks: "
|
||||||
|
f"prefill={declared.get('dsa_prefill_backend', view.dsa_prefill_backend)}, "
|
||||||
|
f"decode={declared.get('dsa_decode_backend', view.dsa_decode_backend)}."
|
||||||
|
)
|
||||||
|
return declared
|
||||||
|
|
||||||
if is_glm_sm12_fp8:
|
if is_glm_sm12_fp8:
|
||||||
backend = "flashinfer_sparse_mla"
|
backend = "flashinfer_sparse_mla"
|
||||||
if not user_set_prefill:
|
if not user_set_prefill:
|
||||||
@@ -763,6 +792,8 @@ _DEEPSEEK_FAMILY_ARCHS = frozenset(
|
|||||||
"MistralLarge3ForCausalLM",
|
"MistralLarge3ForCausalLM",
|
||||||
"PixtralForConditionalGeneration",
|
"PixtralForConditionalGeneration",
|
||||||
"GlmMoeDsaForCausalLM",
|
"GlmMoeDsaForCausalLM",
|
||||||
|
"HYV4ForCausalLM",
|
||||||
|
"HYV4ForCausalLMNextN",
|
||||||
"LongcatFlashForCausalLM",
|
"LongcatFlashForCausalLM",
|
||||||
"LongcatFlashForCausalLMNextN",
|
"LongcatFlashForCausalLMNextN",
|
||||||
"Dots3NoteForCausalLM",
|
"Dots3NoteForCausalLM",
|
||||||
|
|||||||
@@ -862,6 +862,7 @@ def _handle_eagle_family(server_args: ServerArgs) -> None:
|
|||||||
"MistralLarge3ForCausalLM",
|
"MistralLarge3ForCausalLM",
|
||||||
"PixtralForConditionalGeneration",
|
"PixtralForConditionalGeneration",
|
||||||
"HYV3ForCausalLM",
|
"HYV3ForCausalLM",
|
||||||
|
"HYV4ForCausalLM",
|
||||||
]:
|
]:
|
||||||
if cfg.speculative_draft_model_path is None:
|
if cfg.speculative_draft_model_path is None:
|
||||||
declare_resolution(
|
declare_resolution(
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ from sglang.srt.configs.dots_vlm import DotsVLMConfig
|
|||||||
from sglang.srt.configs.exaone import ExaoneConfig
|
from sglang.srt.configs.exaone import ExaoneConfig
|
||||||
from sglang.srt.configs.falcon_h1 import FalconH1Config
|
from sglang.srt.configs.falcon_h1 import FalconH1Config
|
||||||
from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig
|
from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig
|
||||||
|
from sglang.srt.configs.hy_v4 import HYV4Config
|
||||||
from sglang.srt.configs.inkling import (
|
from sglang.srt.configs.inkling import (
|
||||||
InklingAudioConfig,
|
InklingAudioConfig,
|
||||||
InklingMMConfig,
|
InklingMMConfig,
|
||||||
@@ -118,6 +119,7 @@ __all__ = [
|
|||||||
"Dots3Config",
|
"Dots3Config",
|
||||||
"FalconH1Config",
|
"FalconH1Config",
|
||||||
"GraniteMoeHybridConfig",
|
"GraniteMoeHybridConfig",
|
||||||
|
"HYV4Config",
|
||||||
"Lfm2Config",
|
"Lfm2Config",
|
||||||
"Lfm2MoeConfig",
|
"Lfm2MoeConfig",
|
||||||
"Lfm2VlConfig",
|
"Lfm2VlConfig",
|
||||||
|
|||||||
@@ -0,0 +1,152 @@
|
|||||||
|
from transformers.configuration_utils import PreTrainedConfig
|
||||||
|
|
||||||
|
|
||||||
|
class HYV4Config(PreTrainedConfig):
|
||||||
|
model_type = "hy_v4"
|
||||||
|
keys_to_ignore_at_inference = ["past_key_values"]
|
||||||
|
attribute_map = {"num_local_experts": "n_routed_experts"}
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
vocab_size=120832,
|
||||||
|
hidden_size=2816,
|
||||||
|
intermediate_size=6912,
|
||||||
|
moe_intermediate_size=768,
|
||||||
|
num_hidden_layers=34,
|
||||||
|
num_attention_heads=32,
|
||||||
|
num_key_value_heads=32,
|
||||||
|
hidden_act="silu",
|
||||||
|
max_position_embeddings=262144,
|
||||||
|
rms_norm_eps=1e-5,
|
||||||
|
n_routed_experts=256,
|
||||||
|
n_shared_experts=1,
|
||||||
|
num_experts_per_tok=8,
|
||||||
|
routed_scaling_factor=2.827,
|
||||||
|
norm_topk_prob=True,
|
||||||
|
q_lora_rank=1536,
|
||||||
|
kv_lora_rank=512,
|
||||||
|
qk_nope_head_dim=192,
|
||||||
|
qk_rope_head_dim=64,
|
||||||
|
v_head_dim=256,
|
||||||
|
mlp_layer_types=None,
|
||||||
|
layer_types=None,
|
||||||
|
index_topk=2048,
|
||||||
|
index_head_dim=128,
|
||||||
|
index_n_heads=16,
|
||||||
|
indexer_types=None,
|
||||||
|
enable_lm_head_fp32=True,
|
||||||
|
enable_ihc=True,
|
||||||
|
hc_mult=4,
|
||||||
|
hc_magnitude=2.0,
|
||||||
|
hc_eps=1e-6,
|
||||||
|
gated_mla=True,
|
||||||
|
gating_type="elementwise",
|
||||||
|
learnable_sink=True,
|
||||||
|
learnable_sink_init=0.0,
|
||||||
|
swiglu_limit=10.0,
|
||||||
|
rope_parameters=None,
|
||||||
|
num_nextn_predict_layers=1,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
self.vocab_size = vocab_size
|
||||||
|
self.hidden_size = hidden_size
|
||||||
|
self.intermediate_size = intermediate_size
|
||||||
|
self.moe_intermediate_size = moe_intermediate_size
|
||||||
|
self.num_hidden_layers = num_hidden_layers
|
||||||
|
self.num_attention_heads = num_attention_heads
|
||||||
|
self.num_key_value_heads = num_key_value_heads
|
||||||
|
self.hidden_act = hidden_act
|
||||||
|
self.max_position_embeddings = max_position_embeddings
|
||||||
|
self.rms_norm_eps = rms_norm_eps
|
||||||
|
self.n_routed_experts = n_routed_experts
|
||||||
|
self.n_shared_experts = n_shared_experts
|
||||||
|
self.num_experts_per_tok = num_experts_per_tok
|
||||||
|
self.routed_scaling_factor = routed_scaling_factor
|
||||||
|
self.norm_topk_prob = norm_topk_prob
|
||||||
|
self.q_lora_rank = q_lora_rank
|
||||||
|
self.kv_lora_rank = kv_lora_rank
|
||||||
|
self.qk_nope_head_dim = qk_nope_head_dim
|
||||||
|
self.qk_rope_head_dim = qk_rope_head_dim
|
||||||
|
self.v_head_dim = v_head_dim
|
||||||
|
self.index_topk = index_topk
|
||||||
|
self.index_head_dim = index_head_dim
|
||||||
|
self.index_n_heads = index_n_heads
|
||||||
|
self.enable_lm_head_fp32 = enable_lm_head_fp32
|
||||||
|
self.enable_ihc = enable_ihc
|
||||||
|
self.hc_mult = hc_mult
|
||||||
|
self.hc_magnitude = hc_magnitude
|
||||||
|
self.hc_eps = hc_eps
|
||||||
|
self.gated_mla = gated_mla
|
||||||
|
self.gating_type = gating_type
|
||||||
|
self.learnable_sink = learnable_sink
|
||||||
|
self.learnable_sink_init = learnable_sink_init
|
||||||
|
self.swiglu_limit = swiglu_limit
|
||||||
|
self.rope_parameters = (
|
||||||
|
rope_parameters
|
||||||
|
if rope_parameters is not None
|
||||||
|
else {"rope_theta": 10000000.0, "rope_type": "default"}
|
||||||
|
)
|
||||||
|
self.rope_theta = self.rope_parameters["rope_theta"]
|
||||||
|
self.rope_interleave = True
|
||||||
|
self.indexer_rope_interleave = True
|
||||||
|
self.router_fp32 = True
|
||||||
|
self.num_nextn_predict_layers = num_nextn_predict_layers
|
||||||
|
self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
|
||||||
|
self.head_dim = qk_rope_head_dim
|
||||||
|
self.mlp_layer_types = (
|
||||||
|
mlp_layer_types
|
||||||
|
if mlp_layer_types is not None
|
||||||
|
else ["dense"] + ["sparse"] * (num_hidden_layers - 1)
|
||||||
|
)
|
||||||
|
self.layer_types = (
|
||||||
|
layer_types
|
||||||
|
if layer_types is not None
|
||||||
|
else ["deepseek_sparse_attention"] * num_hidden_layers
|
||||||
|
)
|
||||||
|
self.indexer_types = (
|
||||||
|
indexer_types
|
||||||
|
if indexer_types is not None
|
||||||
|
else [
|
||||||
|
"full" if i == 0 or (i - 1) % 4 == 0 else "shared"
|
||||||
|
for i in range(num_hidden_layers)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
self.first_k_dense_replace = 1
|
||||||
|
self.moe_layer_freq = 1
|
||||||
|
self.scoring_func = "sigmoid"
|
||||||
|
self.topk_method = "noaux_tc"
|
||||||
|
self.n_group = 1
|
||||||
|
self.topk_group = 1
|
||||||
|
self._validate_hy_v4()
|
||||||
|
|
||||||
|
def _validate_hy_v4(self):
|
||||||
|
fields = {
|
||||||
|
"mlp_layer_types": self.mlp_layer_types,
|
||||||
|
"layer_types": self.layer_types,
|
||||||
|
"indexer_types": self.indexer_types,
|
||||||
|
}
|
||||||
|
for name, values in fields.items():
|
||||||
|
if len(values) != self.num_hidden_layers:
|
||||||
|
raise ValueError(
|
||||||
|
f"{name} must contain {self.num_hidden_layers} entries, got {len(values)}"
|
||||||
|
)
|
||||||
|
if set(self.mlp_layer_types) - {"dense", "sparse"}:
|
||||||
|
raise ValueError("mlp_layer_types only supports dense and sparse")
|
||||||
|
if set(self.layer_types) != {"deepseek_sparse_attention"}:
|
||||||
|
raise ValueError("HYV4 only supports deepseek_sparse_attention")
|
||||||
|
if set(self.indexer_types) - {"full", "shared"}:
|
||||||
|
raise ValueError("indexer_types only supports full and shared")
|
||||||
|
if self.indexer_types[0] != "full":
|
||||||
|
raise ValueError("indexer_types must start with a full indexer")
|
||||||
|
if not self.enable_ihc or self.hc_mult <= 0:
|
||||||
|
raise ValueError("HYV4 requires enabled iHC with hc_mult > 0")
|
||||||
|
if not self.gated_mla or self.gating_type != "elementwise":
|
||||||
|
raise ValueError("HYV4 requires elementwise gated MLA")
|
||||||
|
if not self.learnable_sink:
|
||||||
|
raise ValueError("HYV4 requires learnable attention sinks")
|
||||||
|
if self.q_lora_rank is None:
|
||||||
|
raise ValueError("HYV4 sparse attention requires q_lora_rank")
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ["HYV4Config"]
|
||||||
@@ -66,6 +66,14 @@ def _quant_config_to_dict(quant_config):
|
|||||||
return quant_config
|
return quant_config
|
||||||
|
|
||||||
|
|
||||||
|
def unwrap_modelopt_quantization_config(quant_config: dict) -> dict:
|
||||||
|
quantization = quant_config.get("quantization", quant_config)
|
||||||
|
if not isinstance(quantization, dict):
|
||||||
|
return {}
|
||||||
|
nested = quantization.get("quantization")
|
||||||
|
return nested if isinstance(nested, dict) else quantization
|
||||||
|
|
||||||
|
|
||||||
def get_mimo_v2_fused_qkv_expected_tp_size(hf_config):
|
def get_mimo_v2_fused_qkv_expected_tp_size(hf_config):
|
||||||
layout = getattr(hf_config, "attention_projection_layout", None)
|
layout = getattr(hf_config, "attention_projection_layout", None)
|
||||||
if layout is None:
|
if layout is None:
|
||||||
@@ -134,6 +142,8 @@ def is_deepseek_dsa(config) -> bool:
|
|||||||
"LongcatFlashForCausalLMNextN",
|
"LongcatFlashForCausalLMNextN",
|
||||||
"Dots3NoteForCausalLM",
|
"Dots3NoteForCausalLM",
|
||||||
"Dots3NoteForCausalLMNextN",
|
"Dots3NoteForCausalLMNextN",
|
||||||
|
"HYV4ForCausalLM",
|
||||||
|
"HYV4ForCausalLMNextN",
|
||||||
)
|
)
|
||||||
and _hf_attr(config, "index_topk") is not None
|
and _hf_attr(config, "index_topk") is not None
|
||||||
)
|
)
|
||||||
@@ -167,6 +177,17 @@ def is_deepseek_v4(config) -> bool:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_spec_hidden_size(
|
||||||
|
hf_config, hidden_size: int, hc_mult: int
|
||||||
|
) -> tuple[int, Optional[int]]:
|
||||||
|
# Only DSV4 carries the hc-flattened stream across the target→draft
|
||||||
|
# boundary; other hc models (hy_v4) collapse to hidden_size first.
|
||||||
|
if hc_mult <= 1 or not is_deepseek_v4(hf_config):
|
||||||
|
return hidden_size, None
|
||||||
|
hc_hidden_size = hidden_size * hc_mult
|
||||||
|
return hc_hidden_size, hc_hidden_size
|
||||||
|
|
||||||
|
|
||||||
def get_dsa_index_head_dim(config: PretrainedConfig) -> int:
|
def get_dsa_index_head_dim(config: PretrainedConfig) -> int:
|
||||||
assert is_deepseek_dsa(config) or is_deepseek_v4(config)
|
assert is_deepseek_dsa(config) or is_deepseek_v4(config)
|
||||||
return config.index_head_dim
|
return config.index_head_dim
|
||||||
@@ -226,6 +247,12 @@ def dsa_layer_skips_topk(config: PretrainedConfig, layer_id: int) -> bool:
|
|||||||
"""Return whether a DSA layer reuses the previous layer's top-k indices."""
|
"""Return whether a DSA layer reuses the previous layer's top-k indices."""
|
||||||
assert is_deepseek_dsa(config)
|
assert is_deepseek_dsa(config)
|
||||||
|
|
||||||
|
indexer_types = getattr(config, "indexer_types", None)
|
||||||
|
if indexer_types is not None:
|
||||||
|
return (
|
||||||
|
0 <= layer_id < len(indexer_types) and indexer_types[layer_id] == "shared"
|
||||||
|
)
|
||||||
|
|
||||||
# LongCat computes fresh top-k indices every cli_factor layers.
|
# LongCat computes fresh top-k indices every cli_factor layers.
|
||||||
cli_factor = getattr(config, "cli_factor", 1)
|
cli_factor = getattr(config, "cli_factor", 1)
|
||||||
if cli_factor is None:
|
if cli_factor is None:
|
||||||
@@ -794,6 +821,10 @@ class ModelConfig:
|
|||||||
self.hf_config.architectures[0] = "HYV3ForCausalLMNextN"
|
self.hf_config.architectures[0] = "HYV3ForCausalLMNextN"
|
||||||
self.hf_config.num_nextn_predict_layers = 1
|
self.hf_config.num_nextn_predict_layers = 1
|
||||||
|
|
||||||
|
if is_draft_model and self.hf_config.architectures[0] == "HYV4ForCausalLM":
|
||||||
|
self.hf_config.architectures[0] = "HYV4ForCausalLMNextN"
|
||||||
|
self.hf_config.num_nextn_predict_layers = 1
|
||||||
|
|
||||||
def _derive_hybrid_model(self):
|
def _derive_hybrid_model(self):
|
||||||
# Use self.context_len after it has been initialized to prevent using context_len which may be None.
|
# Use self.context_len after it has been initialized to prevent using context_len which may be None.
|
||||||
self.is_hybrid_swa = (
|
self.is_hybrid_swa = (
|
||||||
@@ -865,7 +896,7 @@ class ModelConfig:
|
|||||||
return getattr(
|
return getattr(
|
||||||
self.hf_text_config, "add_swa_attention_sink_bias", False
|
self.hf_text_config, "add_swa_attention_sink_bias", False
|
||||||
) or getattr(self.hf_text_config, "add_full_attention_sink_bias", False)
|
) or getattr(self.hf_text_config, "add_full_attention_sink_bias", False)
|
||||||
return False
|
return bool(getattr(self.hf_text_config, "learnable_sink", False))
|
||||||
|
|
||||||
def _derive_context_length(self, context_length: int):
|
def _derive_context_length(self, context_length: int):
|
||||||
is_draft_model = self.is_draft_model
|
is_draft_model = self.is_draft_model
|
||||||
@@ -939,6 +970,8 @@ class ModelConfig:
|
|||||||
or "GlmMoeDsaForCausalLMNextN" in self.hf_config.architectures
|
or "GlmMoeDsaForCausalLMNextN" in self.hf_config.architectures
|
||||||
or "LongcatFlashForCausalLM" in self.hf_config.architectures
|
or "LongcatFlashForCausalLM" in self.hf_config.architectures
|
||||||
or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures
|
or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures
|
||||||
|
or "HYV4ForCausalLM" in self.hf_config.architectures
|
||||||
|
or "HYV4ForCausalLMNextN" in self.hf_config.architectures
|
||||||
or "DotsVLMForCausalLM" in self.hf_config.architectures
|
or "DotsVLMForCausalLM" in self.hf_config.architectures
|
||||||
or "Dots3NoteForCausalLM" in self.hf_config.architectures
|
or "Dots3NoteForCausalLM" in self.hf_config.architectures
|
||||||
or "Dots3NoteForCausalLMNextN" in self.hf_config.architectures
|
or "Dots3NoteForCausalLMNextN" in self.hf_config.architectures
|
||||||
@@ -969,7 +1002,10 @@ class ModelConfig:
|
|||||||
else None
|
else None
|
||||||
)
|
)
|
||||||
# In transformers v5, rope_scaling is just rope_parameters.
|
# In transformers v5, rope_scaling is just rope_parameters.
|
||||||
self._init_mla_scaling(self.hf_text_config.rope_scaling)
|
rope_scaling = getattr(
|
||||||
|
self.hf_text_config, "rope_parameters", None
|
||||||
|
) or getattr(self.hf_text_config, "rope_scaling", None)
|
||||||
|
self._init_mla_scaling(rope_scaling)
|
||||||
elif (
|
elif (
|
||||||
"DeepseekV4ForCausalLM" in self.hf_config.architectures
|
"DeepseekV4ForCausalLM" in self.hf_config.architectures
|
||||||
or "DeepseekV4ForCausalLMNextN" in self.hf_config.architectures
|
or "DeepseekV4ForCausalLMNextN" in self.hf_config.architectures
|
||||||
@@ -1102,12 +1138,9 @@ class ModelConfig:
|
|||||||
self.num_key_value_heads = self.num_attention_heads
|
self.num_key_value_heads = self.num_attention_heads
|
||||||
self.hidden_size = self.hf_text_config.hidden_size
|
self.hidden_size = self.hf_text_config.hidden_size
|
||||||
hc_mult = getattr(self.hf_text_config, "hc_mult", 1)
|
hc_mult = getattr(self.hf_text_config, "hc_mult", 1)
|
||||||
self.spec_hidden_size = (
|
self.spec_hidden_size, self.hc_hidden_size = resolve_spec_hidden_size(
|
||||||
self.hidden_size * hc_mult if hc_mult > 1 else self.hidden_size
|
self.hf_config, self.hidden_size, hc_mult
|
||||||
)
|
)
|
||||||
# mHC-flattened hidden size; None when not running an mHC model
|
|
||||||
# (e.g. non-DeepSeek-V4 configs without ``hc_mult``).
|
|
||||||
self.hc_hidden_size = self.spec_hidden_size if hc_mult > 1 else None
|
|
||||||
self.num_hidden_layers = self.hf_text_config.num_hidden_layers
|
self.num_hidden_layers = self.hf_text_config.num_hidden_layers
|
||||||
self.num_attention_layers = self.num_hidden_layers
|
self.num_attention_layers = self.num_hidden_layers
|
||||||
if "LongcatFlashForCausalLM" in self.hf_config.architectures:
|
if "LongcatFlashForCausalLM" in self.hf_config.architectures:
|
||||||
@@ -1370,11 +1403,8 @@ class ModelConfig:
|
|||||||
|
|
||||||
{"quant_algo": "FP8", "quant_method": "modelopt", ...}
|
{"quant_algo": "FP8", "quant_method": "modelopt", ...}
|
||||||
"""
|
"""
|
||||||
if "quantization" in quant_config_dict:
|
json_quant_configs = unwrap_modelopt_quantization_config(quant_config_dict)
|
||||||
json_quant_configs = quant_config_dict["quantization"]
|
if not json_quant_configs:
|
||||||
elif "quant_algo" in quant_config_dict:
|
|
||||||
json_quant_configs = quant_config_dict
|
|
||||||
else:
|
|
||||||
return None
|
return None
|
||||||
quant_algo = json_quant_configs.get("quant_algo", None)
|
quant_algo = json_quant_configs.get("quant_algo", None)
|
||||||
|
|
||||||
|
|||||||
@@ -89,6 +89,10 @@ from sglang.srt.function_call.utils import (
|
|||||||
)
|
)
|
||||||
from sglang.srt.managers.io_struct import GenerateReqInput
|
from sglang.srt.managers.io_struct import GenerateReqInput
|
||||||
from sglang.srt.parser.conversation import generate_chat_conv
|
from sglang.srt.parser.conversation import generate_chat_conv
|
||||||
|
from sglang.srt.parser.hunyuan_reasoning import (
|
||||||
|
normalize_hunyuan_reasoning_effort,
|
||||||
|
uses_hunyuan_reasoning_effort,
|
||||||
|
)
|
||||||
from sglang.srt.parser.jinja_template_utils import (
|
from sglang.srt.parser.jinja_template_utils import (
|
||||||
MEDIA_URL_PART_TYPES,
|
MEDIA_URL_PART_TYPES,
|
||||||
process_content_for_template_format,
|
process_content_for_template_format,
|
||||||
@@ -1048,6 +1052,7 @@ class OpenAIServingChat(OpenAIServingBase):
|
|||||||
request: ChatCompletionRequest,
|
request: ChatCompletionRequest,
|
||||||
raw_request: Request = None,
|
raw_request: Request = None,
|
||||||
) -> tuple[GenerateReqInput, ChatCompletionRequest]:
|
) -> tuple[GenerateReqInput, ChatCompletionRequest]:
|
||||||
|
|
||||||
# Header-based opt-in (same rationale as request_headers.py).
|
# Header-based opt-in (same rationale as request_headers.py).
|
||||||
if raw_request is not None and not request.return_input_ids_in_sglext:
|
if raw_request is not None and not request.return_input_ids_in_sglext:
|
||||||
if raw_request.headers.get("x-sglext-return-input-ids") == "1":
|
if raw_request.headers.get("x-sglext-return-input-ids") == "1":
|
||||||
@@ -1057,11 +1062,16 @@ class OpenAIServingChat(OpenAIServingBase):
|
|||||||
if raw_request.headers.get("x-sglext-return-output-ids") == "1":
|
if raw_request.headers.get("x-sglext-return-output-ids") == "1":
|
||||||
request.return_output_ids_in_sglext = True
|
request.return_output_ids_in_sglext = True
|
||||||
|
|
||||||
reasoning_effort = (
|
reasoning_effort = None
|
||||||
request.chat_template_kwargs.pop("reasoning_effort", None)
|
if not uses_hunyuan_reasoning_effort(
|
||||||
if request.chat_template_kwargs
|
self.reasoning_parser, self.template_manager.reasoning_config
|
||||||
else None
|
):
|
||||||
)
|
reasoning_effort = (
|
||||||
|
request.chat_template_kwargs.pop("reasoning_effort", None)
|
||||||
|
if request.chat_template_kwargs
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
|
||||||
if self.is_gpt_oss and reasoning_effort == "none":
|
if self.is_gpt_oss and reasoning_effort == "none":
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"Harmony does not support reasoning effort {reasoning_effort}"
|
f"Harmony does not support reasoning effort {reasoning_effort}"
|
||||||
@@ -1205,6 +1215,10 @@ class OpenAIServingChat(OpenAIServingBase):
|
|||||||
if effort is not None and request.reasoning_effort is None:
|
if effort is not None and request.reasoning_effort is None:
|
||||||
request.reasoning_effort = effort
|
request.reasoning_effort = effort
|
||||||
|
|
||||||
|
normalize_hunyuan_reasoning_effort(
|
||||||
|
request, self.reasoning_parser, self.template_manager.reasoning_config
|
||||||
|
)
|
||||||
|
|
||||||
# GptOss model needs to keep special tokens for harmony parsing
|
# GptOss model needs to keep special tokens for harmony parsing
|
||||||
if self.is_gpt_oss or self.is_gemma4:
|
if self.is_gpt_oss or self.is_gemma4:
|
||||||
request.skip_special_tokens = False
|
request.skip_special_tokens = False
|
||||||
@@ -2575,7 +2589,11 @@ class OpenAIServingChat(OpenAIServingBase):
|
|||||||
return
|
return
|
||||||
|
|
||||||
if self.reasoning_parser == "hunyuan":
|
if self.reasoning_parser == "hunyuan":
|
||||||
request.reasoning_effort = "medium" if enabled else "no_think"
|
config = self.template_manager.reasoning_config
|
||||||
|
if config is not None and config.special_case == "hunyuan_effort":
|
||||||
|
request.reasoning_effort = "high" if enabled else "no_think"
|
||||||
|
else:
|
||||||
|
request.reasoning_effort = "medium" if enabled else "no_think"
|
||||||
return
|
return
|
||||||
|
|
||||||
if self.reasoning_parser == "inkling":
|
if self.reasoning_parser == "inkling":
|
||||||
@@ -2644,6 +2662,9 @@ class OpenAIServingChat(OpenAIServingBase):
|
|||||||
) == "enabled"
|
) == "enabled"
|
||||||
|
|
||||||
if self.reasoning_parser == "hunyuan":
|
if self.reasoning_parser == "hunyuan":
|
||||||
|
config = self.template_manager.reasoning_config
|
||||||
|
if config is not None and config.special_case == "hunyuan_effort":
|
||||||
|
return request.reasoning_effort not in ("none", "no_think")
|
||||||
# Hy3-preview template emits no <think> when reasoning_effort is
|
# Hy3-preview template emits no <think> when reasoning_effort is
|
||||||
# "no_think" / "none" / unset; forcing reasoning would route all
|
# "no_think" / "none" / unset; forcing reasoning would route all
|
||||||
# output into reasoning_content.
|
# output into reasoning_content.
|
||||||
|
|||||||
@@ -103,6 +103,20 @@ class HunyuanDetector(BaseFormatDetector):
|
|||||||
def __init__(self, tokenizer=None):
|
def __init__(self, tokenizer=None):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
|
|
||||||
|
# Hy4 dropped <tool_sep> from its vocab entirely; Hy3 carries it (bare
|
||||||
|
# or suffixed). Detect the dialect from the vocab, not from the
|
||||||
|
# resolved literal (which falls back to "<tool_sep>" for both).
|
||||||
|
self._has_tool_sep = True
|
||||||
|
if tokenizer is not None:
|
||||||
|
try:
|
||||||
|
self._has_tool_sep = any(
|
||||||
|
re.fullmatch(r"<tool_sep(?::[^>]+)?>", tok)
|
||||||
|
for tok in tokenizer.get_vocab()
|
||||||
|
if isinstance(tok, str)
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
t = resolve_hunyuan_tokens(tokenizer)
|
t = resolve_hunyuan_tokens(tokenizer)
|
||||||
tool_calls = t["tool_calls"]
|
tool_calls = t["tool_calls"]
|
||||||
tool_call = t["tool_call"]
|
tool_call = t["tool_call"]
|
||||||
@@ -129,7 +143,13 @@ class HunyuanDetector(BaseFormatDetector):
|
|||||||
self.tool_call_regex = re.compile(
|
self.tool_call_regex = re.compile(
|
||||||
re.escape(tool_call)
|
re.escape(tool_call)
|
||||||
+ r"(.*?)"
|
+ r"(.*?)"
|
||||||
|
+ r"(?:"
|
||||||
+ re.escape(tool_sep)
|
+ re.escape(tool_sep)
|
||||||
|
+ r"|(?="
|
||||||
|
+ re.escape(arg_key)
|
||||||
|
+ r"|"
|
||||||
|
+ re.escape(tc_end)
|
||||||
|
+ r"))"
|
||||||
+ r"(.*?)"
|
+ r"(.*?)"
|
||||||
+ re.escape(tc_end),
|
+ re.escape(tc_end),
|
||||||
re.DOTALL,
|
re.DOTALL,
|
||||||
@@ -383,14 +403,22 @@ class HunyuanDetector(BaseFormatDetector):
|
|||||||
self._in_tool_calls = False
|
self._in_tool_calls = False
|
||||||
break
|
break
|
||||||
|
|
||||||
sep_pos = self._buffer.find(self.tool_sep_token, tc_start)
|
name_start = tc_start + len(self.tool_call_start_token)
|
||||||
if sep_pos == -1:
|
boundaries = [
|
||||||
|
(self._buffer.find(token, name_start), token)
|
||||||
|
for token in (
|
||||||
|
self.tool_sep_token,
|
||||||
|
self.arg_key_start_token,
|
||||||
|
self.tool_call_end_token,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
boundaries = [(pos, token) for pos, token in boundaries if pos != -1]
|
||||||
|
if not boundaries:
|
||||||
self._buffer = self._buffer[tc_start:]
|
self._buffer = self._buffer[tc_start:]
|
||||||
break
|
break
|
||||||
|
boundary_pos, boundary_token = min(boundaries, key=lambda item: item[0])
|
||||||
|
|
||||||
tool_name = self._buffer[
|
tool_name = self._buffer[name_start:boundary_pos].strip()
|
||||||
tc_start + len(self.tool_call_start_token) : sep_pos
|
|
||||||
].strip()
|
|
||||||
|
|
||||||
if (
|
if (
|
||||||
tool_name not in self._tool_indices
|
tool_name not in self._tool_indices
|
||||||
@@ -413,7 +441,9 @@ class HunyuanDetector(BaseFormatDetector):
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
self._buffer = self._buffer[sep_pos + len(self.tool_sep_token) :]
|
if boundary_token == self.tool_sep_token:
|
||||||
|
boundary_pos += len(boundary_token)
|
||||||
|
self._buffer = self._buffer[boundary_pos:]
|
||||||
|
|
||||||
# Phase 2: stream argument JSON of the current tool.
|
# Phase 2: stream argument JSON of the current tool.
|
||||||
before_name = self._streaming_tool_name
|
before_name = self._streaming_tool_name
|
||||||
@@ -534,6 +564,14 @@ class HunyuanDetector(BaseFormatDetector):
|
|||||||
return []
|
return []
|
||||||
|
|
||||||
def structure_info(self) -> _GetInfoFunc:
|
def structure_info(self) -> _GetInfoFunc:
|
||||||
|
# Hy4 has no <tool_sep> and no separators between adjacent
|
||||||
|
# <tool_call> blocks, so the Hy3 begin/end literals never match.
|
||||||
|
if not self._has_tool_sep:
|
||||||
|
return lambda name: StructureInfo(
|
||||||
|
begin=f"{self.bot_token}{self.tool_call_start_token}{name}",
|
||||||
|
end=f"{self.tool_call_end_token}{self.eot_token}",
|
||||||
|
trigger=self.bot_token,
|
||||||
|
)
|
||||||
return lambda name: StructureInfo(
|
return lambda name: StructureInfo(
|
||||||
begin=f"{self.bot_token}\n{self.tool_call_start_token}{name}{self.tool_sep_token}",
|
begin=f"{self.bot_token}\n{self.tool_call_start_token}{name}{self.tool_sep_token}",
|
||||||
end=f"{self.tool_call_end_token}\n{self.eot_token}",
|
end=f"{self.tool_call_end_token}\n{self.eot_token}",
|
||||||
@@ -542,3 +580,9 @@ class HunyuanDetector(BaseFormatDetector):
|
|||||||
|
|
||||||
def supports_structural_tag(self) -> bool:
|
def supports_structural_tag(self) -> bool:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
def parses_required_natively(self) -> bool:
|
||||||
|
# Hy3/Hy4 both emit their structural tool-call format natively; without
|
||||||
|
# a dialect grammar, a json_schema constraint would force plain JSON the
|
||||||
|
# model was not trained to produce here.
|
||||||
|
return True
|
||||||
|
|||||||
@@ -348,6 +348,7 @@ class DeepseekSparseAttnBackend(
|
|||||||
# Keep original head count if it exceeds current padded variants.
|
# Keep original head count if it exceeds current padded variants.
|
||||||
self.flashmla_kv_num_q_heads = self.num_q_heads
|
self.flashmla_kv_num_q_heads = self.num_q_heads
|
||||||
self.enable_auto_select_prefill_impl = self.dsa_prefill_impl == "flashmla_auto"
|
self.enable_auto_select_prefill_impl = self.dsa_prefill_impl == "flashmla_auto"
|
||||||
|
self._sink_pad_cache: dict[tuple[int, int], torch.Tensor] = {}
|
||||||
|
|
||||||
self._arange_buf = torch.arange(16384, device=self.device, dtype=torch.int32)
|
self._arange_buf = torch.arange(16384, device=self.device, dtype=torch.int32)
|
||||||
|
|
||||||
@@ -1885,6 +1886,7 @@ class DeepseekSparseAttnBackend(
|
|||||||
cos_sin_cache: Optional[torch.Tensor] = None,
|
cos_sin_cache: Optional[torch.Tensor] = None,
|
||||||
is_neox: Optional[bool] = False,
|
is_neox: Optional[bool] = False,
|
||||||
llama_4_scaling: Optional[torch.Tensor] = None,
|
llama_4_scaling: Optional[torch.Tensor] = None,
|
||||||
|
attn_sink: Optional[torch.Tensor] = None,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
|
|
||||||
causal = not layer.is_cross_attention
|
causal = not layer.is_cross_attention
|
||||||
@@ -1899,6 +1901,10 @@ class DeepseekSparseAttnBackend(
|
|||||||
)
|
)
|
||||||
else self.dsa_prefill_impl
|
else self.dsa_prefill_impl
|
||||||
)
|
)
|
||||||
|
if attn_sink is not None and dsa_impl != "flashmla_sparse":
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Learnable attention sinks require flashmla_sparse, got {dsa_impl}"
|
||||||
|
)
|
||||||
|
|
||||||
if dsa_impl == "trtllm" and not self.use_mha:
|
if dsa_impl == "trtllm" and not self.use_mha:
|
||||||
return self._forward_trtllm(
|
return self._forward_trtllm(
|
||||||
@@ -2124,6 +2130,7 @@ class DeepseekSparseAttnBackend(
|
|||||||
sm_scale=layer.scaling,
|
sm_scale=layer.scaling,
|
||||||
v_head_dim=layer.v_head_dim,
|
v_head_dim=layer.v_head_dim,
|
||||||
topk_length=metadata.dsa_cache_seqlens_int32,
|
topk_length=metadata.dsa_cache_seqlens_int32,
|
||||||
|
attn_sink=attn_sink,
|
||||||
)
|
)
|
||||||
elif dsa_impl == "flashinfer_sparse_mla":
|
elif dsa_impl == "flashinfer_sparse_mla":
|
||||||
if q_rope is not None:
|
if q_rope is not None:
|
||||||
@@ -2195,12 +2202,19 @@ class DeepseekSparseAttnBackend(
|
|||||||
cos_sin_cache: Optional[torch.Tensor] = None,
|
cos_sin_cache: Optional[torch.Tensor] = None,
|
||||||
is_neox: Optional[bool] = False,
|
is_neox: Optional[bool] = False,
|
||||||
llama_4_scaling: Optional[torch.Tensor] = None,
|
llama_4_scaling: Optional[torch.Tensor] = None,
|
||||||
|
attn_sink: Optional[torch.Tensor] = None,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
|
|
||||||
causal = not layer.is_cross_attention
|
causal = not layer.is_cross_attention
|
||||||
metadata = self.forward_metadata
|
metadata = self.forward_metadata
|
||||||
assert causal, "DSA is causal only"
|
assert causal, "DSA is causal only"
|
||||||
|
|
||||||
|
if attn_sink is not None and self.dsa_decode_impl != "flashmla_sparse":
|
||||||
|
raise RuntimeError(
|
||||||
|
"Learnable attention sinks require flashmla_sparse, got "
|
||||||
|
f"{self.dsa_decode_impl}"
|
||||||
|
)
|
||||||
|
|
||||||
if self.dsa_decode_impl == "trtllm":
|
if self.dsa_decode_impl == "trtllm":
|
||||||
return self._forward_trtllm(
|
return self._forward_trtllm(
|
||||||
q,
|
q,
|
||||||
@@ -2281,6 +2295,7 @@ class DeepseekSparseAttnBackend(
|
|||||||
sm_scale=layer.scaling,
|
sm_scale=layer.scaling,
|
||||||
v_head_dim=layer.v_head_dim,
|
v_head_dim=layer.v_head_dim,
|
||||||
topk_length=metadata.dsa_cache_seqlens_int32,
|
topk_length=metadata.dsa_cache_seqlens_int32,
|
||||||
|
attn_sink=attn_sink,
|
||||||
)
|
)
|
||||||
elif self.dsa_decode_impl == "flashinfer_sparse_mla":
|
elif self.dsa_decode_impl == "flashinfer_sparse_mla":
|
||||||
if q_all is None:
|
if q_all is None:
|
||||||
@@ -2396,6 +2411,7 @@ class DeepseekSparseAttnBackend(
|
|||||||
page_table_1: torch.Tensor,
|
page_table_1: torch.Tensor,
|
||||||
sm_scale: float,
|
sm_scale: float,
|
||||||
topk_length: Optional[torch.Tensor] = None,
|
topk_length: Optional[torch.Tensor] = None,
|
||||||
|
attn_sink: Optional[torch.Tensor] = None,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
from sgl_kernel.flash_mla import flash_mla_sparse_fwd
|
from sgl_kernel.flash_mla import flash_mla_sparse_fwd
|
||||||
|
|
||||||
@@ -2421,6 +2437,15 @@ class DeepseekSparseAttnBackend(
|
|||||||
else:
|
else:
|
||||||
q_input = q_all
|
q_input = q_all
|
||||||
|
|
||||||
|
sink_input = attn_sink
|
||||||
|
if need_padding and attn_sink is not None:
|
||||||
|
key = (attn_sink.data_ptr(), required_padding)
|
||||||
|
sink_input = self._sink_pad_cache.get(key)
|
||||||
|
if sink_input is None:
|
||||||
|
sink_input = attn_sink.new_zeros(required_padding)
|
||||||
|
self._sink_pad_cache[key] = sink_input
|
||||||
|
sink_input[:num_heads].copy_(attn_sink)
|
||||||
|
|
||||||
# indices shape must be (s_q, h_kv=1, topk), keep h_kv=1 unchanged
|
# indices shape must be (s_q, h_kv=1, topk), keep h_kv=1 unchanged
|
||||||
indices_input = page_table_1.unsqueeze(1)
|
indices_input = page_table_1.unsqueeze(1)
|
||||||
|
|
||||||
@@ -2442,6 +2467,7 @@ class DeepseekSparseAttnBackend(
|
|||||||
indices=indices_input,
|
indices=indices_input,
|
||||||
sm_scale=sm_scale,
|
sm_scale=sm_scale,
|
||||||
d_v=v_head_dim,
|
d_v=v_head_dim,
|
||||||
|
attn_sink=sink_input,
|
||||||
topk_length=topk_length,
|
topk_length=topk_length,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -3437,6 +3463,7 @@ class DeepseekSparseAttnMultiStepBackend:
|
|||||||
):
|
):
|
||||||
self.topk = topk
|
self.topk = topk
|
||||||
self.speculative_num_steps = speculative_num_steps
|
self.speculative_num_steps = speculative_num_steps
|
||||||
|
self._sink_pad_cache: dict[tuple[int, int], torch.Tensor] = {}
|
||||||
self.attn_backends = []
|
self.attn_backends = []
|
||||||
for i in range(self.speculative_num_steps - 1):
|
for i in range(self.speculative_num_steps - 1):
|
||||||
self.attn_backends.append(
|
self.attn_backends.append(
|
||||||
|
|||||||
@@ -170,6 +170,52 @@ def grouped_gemm_nt_bf16_contig(
|
|||||||
deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a, b, d, m_indices)
|
deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a, b, d, m_indices)
|
||||||
|
|
||||||
|
|
||||||
|
def get_contiguous_layout_alignment(expected_m: int, num_groups: int) -> int:
|
||||||
|
alignment = deep_gemm.get_mk_alignment_for_contiguous_layout()
|
||||||
|
if isinstance(alignment, tuple):
|
||||||
|
alignment = alignment[0]
|
||||||
|
|
||||||
|
theoretical_alignment = getattr(
|
||||||
|
deep_gemm, "get_theoretical_mk_alignment_for_contiguous_layout", None
|
||||||
|
)
|
||||||
|
if theoretical_alignment is None:
|
||||||
|
return alignment
|
||||||
|
|
||||||
|
per_group_m = (expected_m + num_groups - 1) // num_groups
|
||||||
|
try:
|
||||||
|
try:
|
||||||
|
candidate = theoretical_alignment(per_group_m)
|
||||||
|
except TypeError:
|
||||||
|
candidate = theoretical_alignment(
|
||||||
|
expected_m=expected_m, num_groups=num_groups
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
return alignment
|
||||||
|
return candidate if 0 < candidate <= alignment else alignment
|
||||||
|
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def contiguous_layout_alignment_scope(alignment: Optional[int]):
|
||||||
|
if alignment is None:
|
||||||
|
yield
|
||||||
|
return
|
||||||
|
|
||||||
|
getter = getattr(deep_gemm, "get_mk_alignment_for_contiguous_layout", None)
|
||||||
|
setter = getattr(deep_gemm, "set_mk_alignment_for_contiguous_layout", None)
|
||||||
|
if getter is None or setter is None:
|
||||||
|
yield
|
||||||
|
return
|
||||||
|
|
||||||
|
previous = getter()
|
||||||
|
if isinstance(previous, tuple):
|
||||||
|
previous = previous[0]
|
||||||
|
setter(alignment)
|
||||||
|
try:
|
||||||
|
yield
|
||||||
|
finally:
|
||||||
|
setter(previous)
|
||||||
|
|
||||||
|
|
||||||
def gemm_nt_f8f8bf16(
|
def gemm_nt_f8f8bf16(
|
||||||
lhs: Tuple[torch.Tensor, torch.Tensor],
|
lhs: Tuple[torch.Tensor, torch.Tensor],
|
||||||
rhs: Tuple[torch.Tensor, torch.Tensor],
|
rhs: Tuple[torch.Tensor, torch.Tensor],
|
||||||
|
|||||||
@@ -386,7 +386,9 @@ class LogitsProcessor(nn.Module):
|
|||||||
self.logit_scale = logit_scale
|
self.logit_scale = logit_scale
|
||||||
self.use_attn_tp_group = get_parallel().enable_dp_lm_head
|
self.use_attn_tp_group = get_parallel().enable_dp_lm_head
|
||||||
self.use_tp_lm_head_all_to_all = get_parallel().enable_tp_lm_head_all_to_all
|
self.use_tp_lm_head_all_to_all = get_parallel().enable_tp_lm_head_all_to_all
|
||||||
self.use_fp32_lm_head = get_exec().features.enable_fp32_lm_head
|
self.use_fp32_lm_head = get_exec().features.enable_fp32_lm_head or getattr(
|
||||||
|
config, "enable_lm_head_fp32", False
|
||||||
|
)
|
||||||
if self.use_attn_tp_group:
|
if self.use_attn_tp_group:
|
||||||
self.attn_tp_size = get_parallel().attn_tp_size
|
self.attn_tp_size = get_parallel().attn_tp_size
|
||||||
self.do_tensor_parallel_all_gather = (
|
self.do_tensor_parallel_all_gather = (
|
||||||
|
|||||||
@@ -50,6 +50,7 @@ def cutlass_fused_experts_fp8(
|
|||||||
use_mxfp8: bool = False,
|
use_mxfp8: bool = False,
|
||||||
output: Optional[torch.Tensor] = None,
|
output: Optional[torch.Tensor] = None,
|
||||||
enable_es: Tuple[bool, bool] = (False, False),
|
enable_es: Tuple[bool, bool] = (False, False),
|
||||||
|
swiglu_limit: Optional[float] = None,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
"""Performs Fused MoE computation using CUTLASS-like kernels with FP8 weights and activations.
|
"""Performs Fused MoE computation using CUTLASS-like kernels with FP8 weights and activations.
|
||||||
|
|
||||||
@@ -271,7 +272,12 @@ def cutlass_fused_experts_fp8(
|
|||||||
)
|
)
|
||||||
|
|
||||||
intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype)
|
intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype)
|
||||||
silu_and_mul(c1, intermediate)
|
if swiglu_limit is None:
|
||||||
|
silu_and_mul(c1, intermediate)
|
||||||
|
else:
|
||||||
|
from sglang.kernels.ops.attention.dsv4 import silu_and_mul_clamp
|
||||||
|
|
||||||
|
silu_and_mul_clamp(c1, intermediate, swiglu_limit)
|
||||||
|
|
||||||
if use_mxfp8 and es_down:
|
if use_mxfp8 and es_down:
|
||||||
intemediate_q = torch.empty_like(intermediate, dtype=torch.float8_e4m3fn)
|
intemediate_q = torch.empty_like(intermediate, dtype=torch.float8_e4m3fn)
|
||||||
|
|||||||
@@ -8,7 +8,11 @@ import torch
|
|||||||
import triton
|
import triton
|
||||||
import triton.language as tl
|
import triton.language as tl
|
||||||
|
|
||||||
from sglang.kernels.ops.attention.dsv4 import silu_and_mul_masked_post_quant
|
from sglang.kernels.ops.attention.dsv4 import (
|
||||||
|
silu_and_mul_clamp,
|
||||||
|
silu_and_mul_masked_post_quant,
|
||||||
|
)
|
||||||
|
from sglang.kernels.ops.moe.triton_pad_expert_counts import pad_expert_counts
|
||||||
from sglang.kernels.ops.quantization import per_token_group_quant
|
from sglang.kernels.ops.quantization import per_token_group_quant
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -164,6 +168,32 @@ def _estimate_masked_standard_layout_peak_bytes(
|
|||||||
return runner_config.num_local_experts * padded_m * peak_row_bytes
|
return runner_config.num_local_experts * padded_m * peak_row_bytes
|
||||||
|
|
||||||
|
|
||||||
|
_masked_activation_fallback_logged = False
|
||||||
|
|
||||||
|
|
||||||
|
# Masked clamped/swizzled activation only exists as the DSV4 JIT kernel, which
|
||||||
|
# requires D // 8 >= E and group 128 (silu_and_mul_masked_post_quant.cuh:245).
|
||||||
|
def _masked_activation_unsupported_reason(
|
||||||
|
runner_config: MoeRunnerConfig, quant_info: DeepGemmMoeQuantInfo
|
||||||
|
) -> Optional[str]:
|
||||||
|
if runner_config.swiglu_limit is None and not get_moe_a2a_backend().is_megamoe():
|
||||||
|
return None
|
||||||
|
d = runner_config.intermediate_size_per_partition
|
||||||
|
e = runner_config.num_local_experts
|
||||||
|
if d is None or e is None:
|
||||||
|
return None
|
||||||
|
group_size = quant_info.block_shape[1] if quant_info.block_shape else 128
|
||||||
|
if d // 8 < e:
|
||||||
|
return f"D // 8 ({d // 8}) < num_local_experts ({e})"
|
||||||
|
if group_size != 128:
|
||||||
|
return (
|
||||||
|
f"masked activation group_size {group_size}, DSV4 JIT kernel requires 128"
|
||||||
|
)
|
||||||
|
if d % (group_size * 4) != 0:
|
||||||
|
return f"D ({d}) not divisible by 4 * group_size"
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
def _should_use_masked_standard_layout(
|
def _should_use_masked_standard_layout(
|
||||||
runner_config: MoeRunnerConfig,
|
runner_config: MoeRunnerConfig,
|
||||||
quant_info: DeepGemmMoeQuantInfo,
|
quant_info: DeepGemmMoeQuantInfo,
|
||||||
@@ -177,6 +207,18 @@ def _should_use_masked_standard_layout(
|
|||||||
):
|
):
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
reason = _masked_activation_unsupported_reason(runner_config, quant_info)
|
||||||
|
if reason is not None:
|
||||||
|
global _masked_activation_fallback_logged
|
||||||
|
if not _masked_activation_fallback_logged:
|
||||||
|
_masked_activation_fallback_logged = True
|
||||||
|
logger.info(
|
||||||
|
"DeepGEMM masked standard layout disabled: %s. "
|
||||||
|
"Clamped/swizzled activations on this config must use the "
|
||||||
|
"compact layout.",
|
||||||
|
reason,
|
||||||
|
)
|
||||||
|
return False
|
||||||
mode = envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.get().lower()
|
mode = envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.get().lower()
|
||||||
if mode not in ("auto", "masked", "compact"):
|
if mode not in ("auto", "masked", "compact"):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
@@ -287,24 +329,30 @@ class DeepGemmRunnerCore(MoeRunnerCore):
|
|||||||
hooks: Optional[Any] = None,
|
hooks: Optional[Any] = None,
|
||||||
) -> DeepGemmRunnerOutput:
|
) -> DeepGemmRunnerOutput:
|
||||||
weight_dtype = quant_info.w13_weight.dtype
|
weight_dtype = quant_info.w13_weight.dtype
|
||||||
if not runner_input.use_masked_gemm:
|
alignment = (
|
||||||
if weight_dtype == torch.bfloat16:
|
running_state.get("contiguous_layout_alignment")
|
||||||
hidden_states = self._run_bf16_contiguous_gemm(
|
if not runner_input.use_masked_gemm
|
||||||
runner_input, quant_info, running_state
|
else None
|
||||||
)
|
)
|
||||||
|
with deep_gemm_wrapper.contiguous_layout_alignment_scope(alignment):
|
||||||
|
if not runner_input.use_masked_gemm:
|
||||||
|
if weight_dtype == torch.bfloat16:
|
||||||
|
hidden_states = self._run_bf16_contiguous_gemm(
|
||||||
|
runner_input, quant_info, running_state
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
hidden_states = self._run_contiguous_gemm(
|
||||||
|
runner_input, quant_info, running_state
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
hidden_states = self._run_contiguous_gemm(
|
if weight_dtype == torch.bfloat16:
|
||||||
runner_input, quant_info, running_state
|
hidden_states = self._run_masked_bf16_gemm(
|
||||||
)
|
runner_input, quant_info, running_state
|
||||||
else:
|
)
|
||||||
if weight_dtype == torch.bfloat16:
|
else:
|
||||||
hidden_states = self._run_masked_bf16_gemm(
|
hidden_states = self._run_masked_gemm(
|
||||||
runner_input, quant_info, running_state
|
runner_input, quant_info, running_state
|
||||||
)
|
)
|
||||||
else:
|
|
||||||
hidden_states = self._run_masked_gemm(
|
|
||||||
runner_input, quant_info, running_state
|
|
||||||
)
|
|
||||||
return DeepGemmRunnerOutput(hidden_states=hidden_states)
|
return DeepGemmRunnerOutput(hidden_states=hidden_states)
|
||||||
|
|
||||||
def _run_contiguous_gemm(
|
def _run_contiguous_gemm(
|
||||||
@@ -333,7 +381,7 @@ class DeepGemmRunnerCore(MoeRunnerCore):
|
|||||||
|
|
||||||
N = quant_info.w13_weight.size(1)
|
N = quant_info.w13_weight.size(1)
|
||||||
K = hidden_states_shape[1]
|
K = hidden_states_shape[1]
|
||||||
scale_block_size = 128
|
scale_block_size = quant_info.block_shape[1] if quant_info.use_mxfp8 else 128
|
||||||
|
|
||||||
if all_tokens == 0:
|
if all_tokens == 0:
|
||||||
if trace_deepep_v2_contig:
|
if trace_deepep_v2_contig:
|
||||||
@@ -344,9 +392,12 @@ class DeepGemmRunnerCore(MoeRunnerCore):
|
|||||||
(0, K), device=hidden_states_device, dtype=torch.bfloat16
|
(0, K), device=hidden_states_device, dtype=torch.bfloat16
|
||||||
)
|
)
|
||||||
|
|
||||||
recipe_a, recipe_b = (
|
if quant_info.use_mxfp8:
|
||||||
((1, 128), (1, 32)) if quant_info.is_fp4_experts else (None, None)
|
recipe_a = recipe_b = tuple(quant_info.block_shape)
|
||||||
)
|
elif quant_info.is_fp4_experts:
|
||||||
|
recipe_a, recipe_b = (1, 128), (1, 32)
|
||||||
|
else:
|
||||||
|
recipe_a, recipe_b = None, None
|
||||||
|
|
||||||
w13_weight_fp8 = (
|
w13_weight_fp8 = (
|
||||||
quant_info.w13_weight,
|
quant_info.w13_weight,
|
||||||
@@ -469,19 +520,25 @@ class DeepGemmRunnerCore(MoeRunnerCore):
|
|||||||
sglang_per_token_group_quant_fp8,
|
sglang_per_token_group_quant_fp8,
|
||||||
)
|
)
|
||||||
|
|
||||||
if self.swiglu_limit is not None:
|
|
||||||
gateup_output = _apply_swiglu_limit(
|
|
||||||
gateup_output, swiglu_limit=self.swiglu_limit
|
|
||||||
)
|
|
||||||
|
|
||||||
if not _is_musa:
|
if not _is_musa:
|
||||||
down_input = torch.empty(
|
down_input = torch.empty(
|
||||||
(all_tokens, N // 2),
|
(all_tokens, N // 2),
|
||||||
device=gateup_output.device,
|
device=gateup_output.device,
|
||||||
dtype=torch.bfloat16,
|
dtype=torch.bfloat16,
|
||||||
)
|
)
|
||||||
_legacy_silu_and_mul(gateup_output.view(-1, N), down_input)
|
if self.swiglu_limit is not None:
|
||||||
|
# Fuse the SwiGLU limit with the activation. The quantizing
|
||||||
|
# sibling only supports a group size of 128.
|
||||||
|
silu_and_mul_clamp(
|
||||||
|
gateup_output.view(-1, N), down_input, self.swiglu_limit
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
_legacy_silu_and_mul(gateup_output.view(-1, N), down_input)
|
||||||
else:
|
else:
|
||||||
|
if self.swiglu_limit is not None:
|
||||||
|
gateup_output = _apply_swiglu_limit(
|
||||||
|
gateup_output, swiglu_limit=self.swiglu_limit
|
||||||
|
)
|
||||||
down_input = _silu_and_mul_musa(gateup_output.view(-1, N))
|
down_input = _silu_and_mul_musa(gateup_output.view(-1, N))
|
||||||
del gateup_output
|
del gateup_output
|
||||||
|
|
||||||
@@ -512,12 +569,17 @@ class DeepGemmRunnerCore(MoeRunnerCore):
|
|||||||
if deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES:
|
if deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES:
|
||||||
down_input_scale = tma_align_input_scale(down_input_scale)
|
down_input_scale = tma_align_input_scale(down_input_scale)
|
||||||
|
|
||||||
|
recipe_a_down = (
|
||||||
|
(quant_info.block_shape[0], scale_block_size)
|
||||||
|
if quant_info.use_mxfp8
|
||||||
|
else recipe_a
|
||||||
|
)
|
||||||
deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig(
|
deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig(
|
||||||
(down_input_fp8, down_input_scale),
|
(down_input_fp8, down_input_scale),
|
||||||
w2_weight_fp8,
|
w2_weight_fp8,
|
||||||
down_output,
|
down_output,
|
||||||
m_indices,
|
m_indices,
|
||||||
recipe_a=recipe_a,
|
recipe_a=recipe_a_down,
|
||||||
recipe_b=recipe_b,
|
recipe_b=recipe_b,
|
||||||
)
|
)
|
||||||
if trace_deepep_v2_contig:
|
if trace_deepep_v2_contig:
|
||||||
@@ -976,9 +1038,11 @@ def pre_permute_standard_to_deep_gemm(
|
|||||||
|
|
||||||
# The compact layout avoids scaling masked buffers with the expert count.
|
# The compact layout avoids scaling masked buffers with the expert count.
|
||||||
# Scatter and post-permute skip non-local experts mapped to -1.
|
# Scatter and post-permute skip non-local experts mapped to -1.
|
||||||
block_e = 128
|
|
||||||
num_experts = runner_config.num_local_experts
|
num_experts = runner_config.num_local_experts
|
||||||
num_assignments = topk_ids.numel()
|
num_assignments = topk_ids.numel()
|
||||||
|
block_e = deep_gemm_wrapper.get_contiguous_layout_alignment(
|
||||||
|
num_assignments, num_experts
|
||||||
|
)
|
||||||
all_tokens = _get_compact_all_tokens(num_assignments, num_experts, block_e)
|
all_tokens = _get_compact_all_tokens(num_assignments, num_experts, block_e)
|
||||||
|
|
||||||
tokens_per_expert, unused_masked_dst = fused_moe_dispatch_index(
|
tokens_per_expert, unused_masked_dst = fused_moe_dispatch_index(
|
||||||
@@ -986,10 +1050,15 @@ def pre_permute_standard_to_deep_gemm(
|
|||||||
)
|
)
|
||||||
dispose_tensor(unused_masked_dst)
|
dispose_tensor(unused_masked_dst)
|
||||||
valid_tokens_per_expert = tokens_per_expert
|
valid_tokens_per_expert = tokens_per_expert
|
||||||
tokens_per_expert = (ceil_div(tokens_per_expert, block_e) * block_e).to(torch.int32)
|
if _is_cuda:
|
||||||
# Keep graph-static shapes by appending padding to the final segment.
|
tokens_per_expert = pad_expert_counts(tokens_per_expert, block_e, all_tokens)
|
||||||
# Its m_indices stay -1, so DeepGEMM skips those rows.
|
else:
|
||||||
tokens_per_expert[-1].add_(all_tokens - tokens_per_expert.sum())
|
# The Triton kernel is CUDA-only. Keep the existing MUSA-compatible
|
||||||
|
# tensor implementation for other DeepGEMM backends.
|
||||||
|
tokens_per_expert = (ceil_div(tokens_per_expert, block_e) * block_e).to(
|
||||||
|
torch.int32
|
||||||
|
)
|
||||||
|
tokens_per_expert[-1].add_(all_tokens - tokens_per_expert.sum())
|
||||||
|
|
||||||
k = hidden_states.size(1)
|
k = hidden_states.size(1)
|
||||||
output_dtype = (
|
output_dtype = (
|
||||||
@@ -1023,7 +1092,13 @@ def pre_permute_standard_to_deep_gemm(
|
|||||||
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
|
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
packed_input = torch.zeros(
|
# ep_scatter writes every live row and the grouped GEMM's results for
|
||||||
|
# the alignment padding are dropped by post_reorder, so the zeroing is
|
||||||
|
# dead work -- 174 MB per layer at bs=64. The sibling dispatch path in
|
||||||
|
# this file already allocates its equivalent buffer with torch.empty
|
||||||
|
# unless deterministic inference is on; match it.
|
||||||
|
deterministic = get_exec().deterministic.enable_deterministic_inference
|
||||||
|
packed_input = (torch.zeros if deterministic else torch.empty)(
|
||||||
(all_tokens, k),
|
(all_tokens, k),
|
||||||
device=hidden_states_device,
|
device=hidden_states_device,
|
||||||
dtype=torch.float8_e4m3fn,
|
dtype=torch.float8_e4m3fn,
|
||||||
@@ -1062,6 +1137,7 @@ def pre_permute_standard_to_deep_gemm(
|
|||||||
src2dst,
|
src2dst,
|
||||||
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
|
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
|
||||||
quant_block_size=(quant_info.block_shape[1] if quant_info.block_shape else 128),
|
quant_block_size=(quant_info.block_shape[1] if quant_info.block_shape else 128),
|
||||||
|
expert_alignment=block_e,
|
||||||
expert_start=expert_start,
|
expert_start=expert_start,
|
||||||
)
|
)
|
||||||
if packed_input_source is not hidden_states:
|
if packed_input_source is not hidden_states:
|
||||||
@@ -1080,6 +1156,7 @@ def pre_permute_standard_to_deep_gemm(
|
|||||||
running_state["hidden_states_device"] = hidden_states_device
|
running_state["hidden_states_device"] = hidden_states_device
|
||||||
running_state["src2dst"] = src2dst
|
running_state["src2dst"] = src2dst
|
||||||
running_state["all_tokens"] = all_tokens
|
running_state["all_tokens"] = all_tokens
|
||||||
|
running_state["contiguous_layout_alignment"] = block_e
|
||||||
running_state["mxfp8_act_gran_k"] = (
|
running_state["mxfp8_act_gran_k"] = (
|
||||||
quant_info.block_shape[1] if quant_info.block_shape else 128
|
quant_info.block_shape[1] if quant_info.block_shape else 128
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -2000,6 +2000,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
|||||||
w2_q = layer.w2_weight.data
|
w2_q = layer.w2_weight.data
|
||||||
w13_s = layer.w13_weight_scale_inv.data
|
w13_s = layer.w13_weight_scale_inv.data
|
||||||
w2_s = layer.w2_weight_scale_inv.data
|
w2_s = layer.w2_weight_scale_inv.data
|
||||||
|
elif get_moe_runner_backend().is_cutlass():
|
||||||
|
w13_q = layer.w13_weight.data
|
||||||
|
w2_q = layer.w2_weight.data
|
||||||
|
w13_s = layer.w13_weight_scale_inv.data
|
||||||
|
w2_s = layer.w2_weight_scale_inv.data
|
||||||
elif (
|
elif (
|
||||||
get_moe_runner_backend().is_flashinfer_trtllm()
|
get_moe_runner_backend().is_flashinfer_trtllm()
|
||||||
or get_moe_runner_backend().is_flashinfer_trtllm_routed()
|
or get_moe_runner_backend().is_flashinfer_trtllm_routed()
|
||||||
@@ -2554,6 +2559,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
|||||||
use_mxfp8=use_mxfp8,
|
use_mxfp8=use_mxfp8,
|
||||||
output=symm_output,
|
output=symm_output,
|
||||||
enable_es=(use_mxfp8, use_mxfp8),
|
enable_es=(use_mxfp8, use_mxfp8),
|
||||||
|
swiglu_limit=self.moe_runner_config.swiglu_limit,
|
||||||
)
|
)
|
||||||
return StandardCombineInput(hidden_states=output)
|
return StandardCombineInput(hidden_states=output)
|
||||||
|
|
||||||
|
|||||||
@@ -312,6 +312,7 @@ def _unified_attention_with_output_impl(
|
|||||||
q_rope: Optional[torch.Tensor] = None,
|
q_rope: Optional[torch.Tensor] = None,
|
||||||
k_rope: Optional[torch.Tensor] = None,
|
k_rope: Optional[torch.Tensor] = None,
|
||||||
sinks: Optional[torch.Tensor] = None,
|
sinks: Optional[torch.Tensor] = None,
|
||||||
|
attn_sink: Optional[torch.Tensor] = None,
|
||||||
# MLA / TRT-LLM / NSA paths pass these through RadixAttention.forward(**kwargs);
|
# MLA / TRT-LLM / NSA paths pass these through RadixAttention.forward(**kwargs);
|
||||||
# they must appear in the schema when --enforce-piecewise-cuda-graph is on.
|
# they must appear in the schema when --enforce-piecewise-cuda-graph is on.
|
||||||
cos_sin_cache: Optional[torch.Tensor] = None,
|
cos_sin_cache: Optional[torch.Tensor] = None,
|
||||||
@@ -364,6 +365,8 @@ def _unified_attention_with_output_impl(
|
|||||||
kwargs["k_rope"] = k_rope[:key_value_num_tokens]
|
kwargs["k_rope"] = k_rope[:key_value_num_tokens]
|
||||||
if sinks is not None:
|
if sinks is not None:
|
||||||
kwargs["sinks"] = sinks
|
kwargs["sinks"] = sinks
|
||||||
|
if attn_sink is not None:
|
||||||
|
kwargs["attn_sink"] = attn_sink
|
||||||
if cos_sin_cache is not None:
|
if cos_sin_cache is not None:
|
||||||
kwargs["cos_sin_cache"] = cos_sin_cache
|
kwargs["cos_sin_cache"] = cos_sin_cache
|
||||||
if is_neox is not None:
|
if is_neox is not None:
|
||||||
@@ -438,6 +441,7 @@ def unified_attention_with_output(
|
|||||||
q_rope: Optional[torch.Tensor] = None,
|
q_rope: Optional[torch.Tensor] = None,
|
||||||
k_rope: Optional[torch.Tensor] = None,
|
k_rope: Optional[torch.Tensor] = None,
|
||||||
sinks: Optional[torch.Tensor] = None,
|
sinks: Optional[torch.Tensor] = None,
|
||||||
|
attn_sink: Optional[torch.Tensor] = None,
|
||||||
cos_sin_cache: Optional[torch.Tensor] = None,
|
cos_sin_cache: Optional[torch.Tensor] = None,
|
||||||
is_neox: Optional[bool] = None,
|
is_neox: Optional[bool] = None,
|
||||||
llama_4_scaling: Optional[torch.Tensor] = None,
|
llama_4_scaling: Optional[torch.Tensor] = None,
|
||||||
@@ -456,6 +460,7 @@ def unified_attention_with_output(
|
|||||||
q_rope=q_rope,
|
q_rope=q_rope,
|
||||||
k_rope=k_rope,
|
k_rope=k_rope,
|
||||||
sinks=sinks,
|
sinks=sinks,
|
||||||
|
attn_sink=attn_sink,
|
||||||
cos_sin_cache=cos_sin_cache,
|
cos_sin_cache=cos_sin_cache,
|
||||||
is_neox=is_neox,
|
is_neox=is_neox,
|
||||||
llama_4_scaling=llama_4_scaling,
|
llama_4_scaling=llama_4_scaling,
|
||||||
@@ -486,6 +491,7 @@ def unified_attention_with_output_and_lse(
|
|||||||
q_rope: Optional[torch.Tensor] = None,
|
q_rope: Optional[torch.Tensor] = None,
|
||||||
k_rope: Optional[torch.Tensor] = None,
|
k_rope: Optional[torch.Tensor] = None,
|
||||||
sinks: Optional[torch.Tensor] = None,
|
sinks: Optional[torch.Tensor] = None,
|
||||||
|
attn_sink: Optional[torch.Tensor] = None,
|
||||||
cos_sin_cache: Optional[torch.Tensor] = None,
|
cos_sin_cache: Optional[torch.Tensor] = None,
|
||||||
is_neox: Optional[bool] = None,
|
is_neox: Optional[bool] = None,
|
||||||
llama_4_scaling: Optional[torch.Tensor] = None,
|
llama_4_scaling: Optional[torch.Tensor] = None,
|
||||||
@@ -504,6 +510,7 @@ def unified_attention_with_output_and_lse(
|
|||||||
q_rope=q_rope,
|
q_rope=q_rope,
|
||||||
k_rope=k_rope,
|
k_rope=k_rope,
|
||||||
sinks=sinks,
|
sinks=sinks,
|
||||||
|
attn_sink=attn_sink,
|
||||||
cos_sin_cache=cos_sin_cache,
|
cos_sin_cache=cos_sin_cache,
|
||||||
is_neox=is_neox,
|
is_neox=is_neox,
|
||||||
llama_4_scaling=llama_4_scaling,
|
llama_4_scaling=llama_4_scaling,
|
||||||
|
|||||||
@@ -118,6 +118,18 @@ def should_defer_dsa_cp_kv_gather(
|
|||||||
return dsa_prefill_cp and fuse_rope_for_trtllm_mla
|
return dsa_prefill_cp and fuse_rope_for_trtllm_mla
|
||||||
|
|
||||||
|
|
||||||
|
def _apply_attention_output_gate(module, attn_output, gate):
|
||||||
|
apply_gate = getattr(module, "apply_attention_output_gate", None)
|
||||||
|
if apply_gate is not None:
|
||||||
|
return apply_gate(attn_output, gate)
|
||||||
|
if hasattr(module, "_apply_gated"):
|
||||||
|
return module._apply_gated(attn_output, gate)
|
||||||
|
raise RuntimeError(
|
||||||
|
"Prepared MLA attention gates are unsigmoided and require a "
|
||||||
|
"model-specific application hook"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class DeepseekMLAForwardMixin:
|
class DeepseekMLAForwardMixin:
|
||||||
def init_mla_forward(self: DeepseekV2AttentionMLA):
|
def init_mla_forward(self: DeepseekV2AttentionMLA):
|
||||||
self.flashinfer_mla_disable_ragged = (
|
self.flashinfer_mla_disable_ragged = (
|
||||||
@@ -161,6 +173,8 @@ class DeepseekMLAForwardMixin:
|
|||||||
return False
|
return False
|
||||||
if is_kv_b_lora_active(self):
|
if is_kv_b_lora_active(self):
|
||||||
return False
|
return False
|
||||||
|
if getattr(self, "learnable_sink_param", None) is not None:
|
||||||
|
return False
|
||||||
# The isolated 1-kernel graph is the bf16 fallback BMM. The fp8 and
|
# The isolated 1-kernel graph is the bf16 fallback BMM. The fp8 and
|
||||||
# DeepGEMM branches already use different fused paths.
|
# DeepGEMM branches already use different fused paths.
|
||||||
if self.w_kc.dtype == torch.float8_e4m3fn:
|
if self.w_kc.dtype == torch.float8_e4m3fn:
|
||||||
@@ -291,6 +305,11 @@ class DeepseekMLAForwardMixin:
|
|||||||
# True between the alt-stream fork and its consumption in the born
|
# True between the alt-stream fork and its consumption in the born
|
||||||
# block; also suppresses the duplicate split/rope on that path.
|
# block; also suppresses the duplicate split/rope on that path.
|
||||||
self._q8kv8_qprep_overlap_pending = False
|
self._q8kv8_qprep_overlap_pending = False
|
||||||
|
attention_output_gate = (
|
||||||
|
self.prepare_attention_output_gate(hidden_states)
|
||||||
|
if hasattr(self, "prepare_attention_output_gate")
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
|
||||||
fuse_bmm_attention = (
|
fuse_bmm_attention = (
|
||||||
self.q_lora_rank is not None
|
self.q_lora_rank is not None
|
||||||
@@ -667,6 +686,13 @@ class DeepseekMLAForwardMixin:
|
|||||||
topk_indices,
|
topk_indices,
|
||||||
llama_4_scaling,
|
llama_4_scaling,
|
||||||
fusion_plan,
|
fusion_plan,
|
||||||
|
# Bailing's DsV3MLA appends its own gate to inner_state, so this
|
||||||
|
# slot is emitted only for models owning the gate hook.
|
||||||
|
*(
|
||||||
|
(attention_output_gate,)
|
||||||
|
if hasattr(self, "prepare_attention_output_gate")
|
||||||
|
else ()
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
def forward_absorb_core(
|
def forward_absorb_core(
|
||||||
@@ -681,18 +707,20 @@ class DeepseekMLAForwardMixin:
|
|||||||
topk_indices,
|
topk_indices,
|
||||||
llama_4_scaling,
|
llama_4_scaling,
|
||||||
fusion_plan: Optional[MlaBmmFusionPlan] = None,
|
fusion_plan: Optional[MlaBmmFusionPlan] = None,
|
||||||
gate: Optional[torch.Tensor] = None,
|
attention_output_gate: Optional[torch.Tensor] = None,
|
||||||
):
|
):
|
||||||
save_kv_cache = True
|
save_kv_cache = True
|
||||||
|
|
||||||
if self.current_attention_backend in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS:
|
if self.current_attention_backend in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS:
|
||||||
extra_args = {}
|
extra_args = {}
|
||||||
|
if getattr(self, "learnable_sink_param", None) is not None:
|
||||||
|
extra_args["attn_sink"] = self.learnable_sink_param
|
||||||
if self._fuse_rope_for_trtllm_mla(forward_batch):
|
if self._fuse_rope_for_trtllm_mla(forward_batch):
|
||||||
extra_args = {
|
extra_args.update(
|
||||||
"cos_sin_cache": self.rotary_emb.cos_sin_cache,
|
cos_sin_cache=self.rotary_emb.cos_sin_cache,
|
||||||
"is_neox": self.rotary_emb.is_neox_style,
|
is_neox=self.rotary_emb.is_neox_style,
|
||||||
"llama_4_scaling": llama_4_scaling,
|
llama_4_scaling=llama_4_scaling,
|
||||||
}
|
)
|
||||||
if fusion_plan is not None:
|
if fusion_plan is not None:
|
||||||
bmm_attention_fn = (
|
bmm_attention_fn = (
|
||||||
bcg_mla_bmm_then_unified_attention
|
bcg_mla_bmm_then_unified_attention
|
||||||
@@ -911,8 +939,10 @@ class DeepseekMLAForwardMixin:
|
|||||||
attn_bmm_output = apply_kv_b_lora_v_correction(
|
attn_bmm_output = apply_kv_b_lora_v_correction(
|
||||||
self, attn_output, attn_bmm_output
|
self, attn_output, attn_bmm_output
|
||||||
)
|
)
|
||||||
if gate is not None:
|
if attention_output_gate is not None:
|
||||||
attn_bmm_output = self._apply_gated(attn_bmm_output, gate)
|
attn_bmm_output = _apply_attention_output_gate(
|
||||||
|
self, attn_bmm_output, attention_output_gate
|
||||||
|
)
|
||||||
output, _ = self.o_proj(attn_bmm_output)
|
output, _ = self.o_proj(attn_bmm_output)
|
||||||
|
|
||||||
if self.next_skip_topk is None:
|
if self.next_skip_topk is None:
|
||||||
|
|||||||
@@ -147,6 +147,44 @@ def _load_fused_indexer_wk(
|
|||||||
return True
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def _load_fused_expert_tensor(
|
||||||
|
name: str,
|
||||||
|
loaded_weight: torch.Tensor,
|
||||||
|
params_dict: Dict[str, torch.Tensor],
|
||||||
|
) -> bool:
|
||||||
|
mappings = (
|
||||||
|
(".experts.gate_up_proj_scale", ".experts.w13_weight_scale_inv", "w13"),
|
||||||
|
(".experts.down_proj_scale", ".experts.w2_weight_scale_inv", "w2"),
|
||||||
|
(".experts.gate_up_proj", ".experts.w13_weight", "w13"),
|
||||||
|
(".experts.down_proj", ".experts.w2_weight", "w2"),
|
||||||
|
)
|
||||||
|
for source, target, shard_id in mappings:
|
||||||
|
if not name.endswith(source):
|
||||||
|
continue
|
||||||
|
param_name = name[: -len(source)] + target
|
||||||
|
param = params_dict.get(param_name)
|
||||||
|
if param is None:
|
||||||
|
return False
|
||||||
|
weight_loader = param.weight_loader
|
||||||
|
for expert_id, expert_weight in enumerate(loaded_weight):
|
||||||
|
if shard_id == "w13":
|
||||||
|
gate, up = expert_weight.chunk(2, dim=0)
|
||||||
|
weight_loader(
|
||||||
|
param, gate, param_name, shard_id="w1", expert_id=expert_id
|
||||||
|
)
|
||||||
|
weight_loader(param, up, param_name, shard_id="w3", expert_id=expert_id)
|
||||||
|
else:
|
||||||
|
weight_loader(
|
||||||
|
param,
|
||||||
|
expert_weight,
|
||||||
|
param_name,
|
||||||
|
shard_id="w2",
|
||||||
|
expert_id=expert_id,
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
class NextNEnabledConfig:
|
class NextNEnabledConfig:
|
||||||
num_nextn_layers: int
|
num_nextn_layers: int
|
||||||
@@ -281,6 +319,9 @@ class DeepseekV2WeightLoaderMixin:
|
|||||||
):
|
):
|
||||||
continue
|
continue
|
||||||
|
|
||||||
|
if _load_fused_expert_tensor(name, loaded_weight, params_dict):
|
||||||
|
continue
|
||||||
|
|
||||||
if "rotary_emb.inv_freq" in name:
|
if "rotary_emb.inv_freq" in name:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
@@ -594,7 +635,11 @@ class DeepseekV2WeightLoaderMixin:
|
|||||||
weight = w
|
weight = w
|
||||||
|
|
||||||
# In multiple weight loading scenarios (e.g. RL), we need to inverse the scale of the weights after the requantization happened at the first loading.
|
# In multiple weight loading scenarios (e.g. RL), we need to inverse the scale of the weights after the requantization happened at the first loading.
|
||||||
if (
|
if weight_scale.format_ue8m0 and weight_scale.dtype == torch.uint8:
|
||||||
|
weight_scale = (weight_scale.to(torch.int32) << 23).view(
|
||||||
|
torch.float32
|
||||||
|
)
|
||||||
|
elif (
|
||||||
should_deepgemm_weight_requant_ue8m0(
|
should_deepgemm_weight_requant_ue8m0(
|
||||||
weight_block_size=(
|
weight_block_size=(
|
||||||
self.quant_config.weight_block_size
|
self.quant_config.weight_block_size
|
||||||
|
|||||||
@@ -478,7 +478,14 @@ class MoEGate(nn.Module):
|
|||||||
self.is_nextn = is_nextn
|
self.is_nextn = is_nextn
|
||||||
self.is_deepseek_v4 = is_deepseek_v4
|
self.is_deepseek_v4 = is_deepseek_v4
|
||||||
self.weight = nn.Parameter(
|
self.weight = nn.Parameter(
|
||||||
torch.empty((config.n_routed_experts, config.hidden_size))
|
torch.empty(
|
||||||
|
(config.n_routed_experts, config.hidden_size),
|
||||||
|
dtype=(
|
||||||
|
torch.float32
|
||||||
|
if getattr(config, "router_fp32", False)
|
||||||
|
else torch.get_default_dtype()
|
||||||
|
),
|
||||||
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
if config.topk_method == "noaux_tc" and not is_hash_moe:
|
if config.topk_method == "noaux_tc" and not is_hash_moe:
|
||||||
@@ -515,6 +522,9 @@ class MoEGate(nn.Module):
|
|||||||
gemm_output_zero_allocator: BumpAllocator = None,
|
gemm_output_zero_allocator: BumpAllocator = None,
|
||||||
forward_batch: ForwardBatch = None,
|
forward_batch: ForwardBatch = None,
|
||||||
):
|
):
|
||||||
|
if self.weight.dtype == torch.float32:
|
||||||
|
return F.linear(hidden_states.float(), self.weight)
|
||||||
|
|
||||||
if use_intel_amx_backend(self):
|
if use_intel_amx_backend(self):
|
||||||
return torch.ops.sgl_kernel.weight_packed_linear(
|
return torch.ops.sgl_kernel.weight_packed_linear(
|
||||||
hidden_states,
|
hidden_states,
|
||||||
|
|||||||
@@ -0,0 +1,734 @@
|
|||||||
|
import functools
|
||||||
|
import logging
|
||||||
|
from typing import Iterable, Tuple
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch import nn
|
||||||
|
from transformers import PretrainedConfig
|
||||||
|
|
||||||
|
from sglang.srt.distributed import get_pp_group
|
||||||
|
from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState
|
||||||
|
from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context
|
||||||
|
from sglang.srt.layers.layernorm import RMSNorm
|
||||||
|
from sglang.srt.layers.linear import ColumnParallelLinear, ReplicatedLinear
|
||||||
|
from sglang.srt.layers.logits_processor import LogitsProcessor
|
||||||
|
from sglang.srt.layers.vocab_parallel_embedding import (
|
||||||
|
ParallelLMHead,
|
||||||
|
VocabParallelEmbedding,
|
||||||
|
get_embedding_tp_kwargs,
|
||||||
|
)
|
||||||
|
from sglang.srt.model_executor.forward_context import get_attn_backend
|
||||||
|
from sglang.srt.models.deepseek_common.attention_forward_methods import (
|
||||||
|
AttnForwardMethod,
|
||||||
|
)
|
||||||
|
from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
|
||||||
|
DeepseekV2WeightLoaderMixin,
|
||||||
|
)
|
||||||
|
from sglang.srt.models.deepseek_v2 import (
|
||||||
|
DeepseekV2AttentionMLA,
|
||||||
|
DeepseekV2MLP,
|
||||||
|
DeepseekV2MoE,
|
||||||
|
)
|
||||||
|
from sglang.srt.runtime_context import get_parallel, get_stream
|
||||||
|
from sglang.srt.utils import BumpAllocator, get_device_capability, is_cuda
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _hpc_ihc_available(op_name: str, hc_mult: int, hidden_size: int) -> bool:
|
||||||
|
try:
|
||||||
|
from sglang.kernels.ops.layernorm.hy4_ihc import _hpc_ihc_op
|
||||||
|
except ImportError:
|
||||||
|
return False
|
||||||
|
return _hpc_ihc_op(op_name, hc_mult, hidden_size) is not None
|
||||||
|
|
||||||
|
|
||||||
|
def permute_hyv4_indexer_weight(name, loaded_weight, config):
|
||||||
|
if ".self_attn.indexer.wq_b." in name:
|
||||||
|
group_count = config.index_n_heads
|
||||||
|
elif any(
|
||||||
|
key in name
|
||||||
|
for key in (
|
||||||
|
".self_attn.indexer.wk.",
|
||||||
|
".self_attn.indexer.k_norm.",
|
||||||
|
)
|
||||||
|
):
|
||||||
|
group_count = 1
|
||||||
|
else:
|
||||||
|
return loaded_weight
|
||||||
|
|
||||||
|
shape = loaded_weight.shape
|
||||||
|
loaded_weight = loaded_weight.reshape(
|
||||||
|
group_count,
|
||||||
|
config.index_head_dim,
|
||||||
|
*shape[1:],
|
||||||
|
)
|
||||||
|
rope_dim = config.qk_rope_head_dim
|
||||||
|
return torch.cat(
|
||||||
|
(
|
||||||
|
loaded_weight[:, -rope_dim:],
|
||||||
|
loaded_weight[:, :-rope_dim],
|
||||||
|
),
|
||||||
|
dim=1,
|
||||||
|
).reshape(shape)
|
||||||
|
|
||||||
|
|
||||||
|
class HYV4HCPreLayer(nn.Module):
|
||||||
|
def __init__(self, config: PretrainedConfig, prefix: str):
|
||||||
|
super().__init__()
|
||||||
|
self.hidden_size = config.hidden_size
|
||||||
|
self.hc_mult = config.hc_mult
|
||||||
|
self.magnitude = config.hc_magnitude
|
||||||
|
self.hc_eps = config.hc_eps
|
||||||
|
self.rms_norm_eps = config.rms_norm_eps
|
||||||
|
self.hc_fn = ReplicatedLinear(
|
||||||
|
config.hidden_size * config.hc_mult,
|
||||||
|
2 * config.hc_mult,
|
||||||
|
bias=False,
|
||||||
|
params_dtype=torch.float32,
|
||||||
|
prefix=f"{prefix}.hc_fn",
|
||||||
|
)
|
||||||
|
self.hc_scale = nn.Parameter(torch.empty(2, dtype=torch.float32))
|
||||||
|
self.hc_base = nn.Parameter(
|
||||||
|
torch.empty(2 * config.hc_mult, dtype=torch.float32)
|
||||||
|
)
|
||||||
|
self._fused_ihc_pre_disabled = False
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
rms_weight: torch.Tensor | None = None,
|
||||||
|
rms_eps: float = 0.0,
|
||||||
|
):
|
||||||
|
if hidden_states.is_cuda and not self._fused_ihc_pre_disabled:
|
||||||
|
try:
|
||||||
|
from sglang.kernels.ops.layernorm.hy4_ihc import fused_hy4_ihc_pre
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
return fused_hy4_ihc_pre(
|
||||||
|
hidden_states,
|
||||||
|
self.hc_fn.weight,
|
||||||
|
self.hc_scale,
|
||||||
|
self.hc_base,
|
||||||
|
self.magnitude,
|
||||||
|
self.rms_norm_eps,
|
||||||
|
self.hc_eps,
|
||||||
|
rms_weight,
|
||||||
|
rms_eps,
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
self._fused_ihc_pre_disabled = True
|
||||||
|
logger.warning(
|
||||||
|
"fused_hy4_ihc_pre failed, disabling fused path",
|
||||||
|
exc_info=True,
|
||||||
|
)
|
||||||
|
shape = hidden_states.shape
|
||||||
|
flat = hidden_states.flatten(1).float()
|
||||||
|
scale = torch.rsqrt(flat.square().mean(-1, keepdim=True) + self.rms_norm_eps)
|
||||||
|
gates = self.hc_fn(flat)[0] * scale
|
||||||
|
pre = (
|
||||||
|
torch.sigmoid(
|
||||||
|
gates[..., : self.hc_mult] * self.hc_scale[0]
|
||||||
|
+ self.hc_base[: self.hc_mult]
|
||||||
|
)
|
||||||
|
+ self.hc_eps
|
||||||
|
)
|
||||||
|
post = (
|
||||||
|
self.magnitude
|
||||||
|
* torch.sigmoid(
|
||||||
|
gates[..., self.hc_mult :] * self.hc_scale[1]
|
||||||
|
+ self.hc_base[self.hc_mult :]
|
||||||
|
)
|
||||||
|
+ self.hc_eps
|
||||||
|
)
|
||||||
|
reduced = torch.sum(pre.unsqueeze(-1) * hidden_states.reshape(shape), dim=1)
|
||||||
|
reduced = reduced.to(hidden_states.dtype)
|
||||||
|
if rms_weight is not None:
|
||||||
|
reduced_float = reduced.float()
|
||||||
|
reduced = (
|
||||||
|
reduced_float
|
||||||
|
* torch.rsqrt(
|
||||||
|
reduced_float.square().mean(dim=-1, keepdim=True) + rms_eps
|
||||||
|
)
|
||||||
|
* rms_weight.float()
|
||||||
|
).to(hidden_states.dtype)
|
||||||
|
return reduced, post
|
||||||
|
|
||||||
|
|
||||||
|
class HYV4HCLayer(nn.Module):
|
||||||
|
def __init__(self, config: PretrainedConfig, prefix: str):
|
||||||
|
super().__init__()
|
||||||
|
self.hidden_size = config.hidden_size
|
||||||
|
self.hc_mult = config.hc_mult
|
||||||
|
self.hc_pre = HYV4HCPreLayer(config, f"{prefix}.hc_pre")
|
||||||
|
self._fused_ihc_post_disabled = False
|
||||||
|
self._fused_ihc_post_pre_disabled = False
|
||||||
|
|
||||||
|
def prepare_input(self, hidden_states: torch.Tensor):
|
||||||
|
if hidden_states.ndim == 3:
|
||||||
|
return hidden_states
|
||||||
|
if hidden_states.ndim != 2:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"iHC expects a 2D or 3D tensor, got {hidden_states.shape}"
|
||||||
|
)
|
||||||
|
if hidden_states.shape[-1] == self.hidden_size:
|
||||||
|
return hidden_states.unsqueeze(1).repeat(1, self.hc_mult, 1)
|
||||||
|
if hidden_states.shape[-1] == self.hidden_size * self.hc_mult:
|
||||||
|
return hidden_states.reshape(-1, self.hc_mult, self.hidden_size)
|
||||||
|
raise RuntimeError(
|
||||||
|
"iHC input width must equal hidden_size or hc_mult * hidden_size"
|
||||||
|
)
|
||||||
|
|
||||||
|
def pre(self, hidden_states: torch.Tensor, norm: RMSNorm | None = None):
|
||||||
|
fuse_norm = (
|
||||||
|
norm is not None
|
||||||
|
and hidden_states.is_cuda
|
||||||
|
and _hpc_ihc_available("fuse_ihc_pre", self.hc_mult, self.hidden_size)
|
||||||
|
)
|
||||||
|
reduced, post = self.hc_pre(
|
||||||
|
hidden_states,
|
||||||
|
norm.weight if fuse_norm else None,
|
||||||
|
norm.variance_epsilon if fuse_norm else 0.0,
|
||||||
|
)
|
||||||
|
if norm is not None and not fuse_norm:
|
||||||
|
reduced = norm(reduced)
|
||||||
|
return reduced, post, hidden_states
|
||||||
|
|
||||||
|
def post(self, output, residual, post):
|
||||||
|
if output.is_cuda and not self._fused_ihc_post_disabled:
|
||||||
|
try:
|
||||||
|
from sglang.kernels.ops.layernorm.hy4_ihc import fused_hy4_ihc_post
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
return fused_hy4_ihc_post(output, residual, post)
|
||||||
|
except Exception:
|
||||||
|
self._fused_ihc_post_disabled = True
|
||||||
|
logger.warning(
|
||||||
|
"fused_hy4_ihc_post failed, disabling fused path",
|
||||||
|
exc_info=True,
|
||||||
|
)
|
||||||
|
result = post.float().unsqueeze(-1) * output.float().unsqueeze(1)
|
||||||
|
return (result + residual.float()).to(output.dtype)
|
||||||
|
|
||||||
|
def post_pre(self, output, residual, post, next_layer, norm):
|
||||||
|
if (
|
||||||
|
output.is_cuda
|
||||||
|
and not self._fused_ihc_post_pre_disabled
|
||||||
|
and _hpc_ihc_available("fuse_ihc_post_pre", self.hc_mult, self.hidden_size)
|
||||||
|
):
|
||||||
|
try:
|
||||||
|
from sglang.kernels.ops.layernorm.hy4_ihc import (
|
||||||
|
fused_hy4_ihc_post_pre,
|
||||||
|
)
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
next_residual, reduced, next_post = fused_hy4_ihc_post_pre(
|
||||||
|
output,
|
||||||
|
residual,
|
||||||
|
post,
|
||||||
|
next_layer.hc_pre.hc_fn.weight,
|
||||||
|
next_layer.hc_pre.hc_scale,
|
||||||
|
next_layer.hc_pre.hc_base,
|
||||||
|
next_layer.hc_pre.magnitude,
|
||||||
|
next_layer.hc_pre.rms_norm_eps,
|
||||||
|
next_layer.hc_pre.hc_eps,
|
||||||
|
norm.weight,
|
||||||
|
norm.variance_epsilon,
|
||||||
|
)
|
||||||
|
return reduced, next_post, next_residual
|
||||||
|
except Exception:
|
||||||
|
self._fused_ihc_post_pre_disabled = True
|
||||||
|
logger.warning(
|
||||||
|
"fused_hy4_ihc_post_pre failed, disabling fused path",
|
||||||
|
exc_info=True,
|
||||||
|
)
|
||||||
|
next_residual = self.post(output, residual, post)
|
||||||
|
next_residual = next_layer.prepare_input(next_residual)
|
||||||
|
return next_layer.pre(next_residual, norm)
|
||||||
|
|
||||||
|
|
||||||
|
class HYV4HCHeadLayer(nn.Module):
|
||||||
|
def __init__(self, config: PretrainedConfig, prefix: str):
|
||||||
|
super().__init__()
|
||||||
|
self.config = config
|
||||||
|
self.hc_head_fn = ReplicatedLinear(
|
||||||
|
config.hc_mult * config.hidden_size,
|
||||||
|
config.hc_mult,
|
||||||
|
bias=False,
|
||||||
|
params_dtype=torch.float32,
|
||||||
|
prefix=f"{prefix}.hc_head_fn",
|
||||||
|
)
|
||||||
|
self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32))
|
||||||
|
self.hc_head_base = nn.Parameter(
|
||||||
|
torch.empty(config.hc_mult, dtype=torch.float32)
|
||||||
|
)
|
||||||
|
self._fused_ihc_head_disabled = False
|
||||||
|
|
||||||
|
def forward(self, hidden_states: torch.Tensor, norm: RMSNorm | None = None):
|
||||||
|
if (
|
||||||
|
hidden_states.is_cuda
|
||||||
|
and not self._fused_ihc_head_disabled
|
||||||
|
and _hpc_ihc_available(
|
||||||
|
"fuse_ihc_head", self.config.hc_mult, self.config.hidden_size
|
||||||
|
)
|
||||||
|
):
|
||||||
|
try:
|
||||||
|
from sglang.kernels.ops.layernorm.hy4_ihc import fused_hy4_ihc_head
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
return fused_hy4_ihc_head(
|
||||||
|
hidden_states,
|
||||||
|
self.hc_head_fn.weight,
|
||||||
|
self.hc_head_scale,
|
||||||
|
self.hc_head_base,
|
||||||
|
self.config.rms_norm_eps,
|
||||||
|
self.config.hc_eps,
|
||||||
|
None if norm is None else norm.weight,
|
||||||
|
0.0 if norm is None else norm.variance_epsilon,
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
self._fused_ihc_head_disabled = True
|
||||||
|
logger.warning(
|
||||||
|
"fused_hy4_ihc_head failed, disabling fused path",
|
||||||
|
exc_info=True,
|
||||||
|
)
|
||||||
|
shape = hidden_states.shape
|
||||||
|
flat = hidden_states.flatten(1).float()
|
||||||
|
scale = torch.rsqrt(
|
||||||
|
flat.square().mean(-1, keepdim=True) + self.config.rms_norm_eps
|
||||||
|
)
|
||||||
|
gates = self.hc_head_fn(flat)[0] * scale
|
||||||
|
gates = (
|
||||||
|
torch.sigmoid(gates * self.hc_head_scale + self.hc_head_base)
|
||||||
|
+ self.config.hc_eps
|
||||||
|
)
|
||||||
|
output = torch.sum(gates.unsqueeze(-1) * flat.reshape(shape), dim=1)
|
||||||
|
output = output.to(hidden_states.dtype)
|
||||||
|
return output if norm is None else norm(output)
|
||||||
|
|
||||||
|
|
||||||
|
class HYV4Attention(DeepseekV2AttentionMLA):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config,
|
||||||
|
layer_id,
|
||||||
|
quant_config=None,
|
||||||
|
prefix="",
|
||||||
|
alt_stream=None,
|
||||||
|
is_nextn=False,
|
||||||
|
):
|
||||||
|
rope_parameters = config.rope_parameters
|
||||||
|
super().__init__(
|
||||||
|
config=config,
|
||||||
|
hidden_size=config.hidden_size,
|
||||||
|
num_heads=config.num_attention_heads,
|
||||||
|
qk_nope_head_dim=config.qk_nope_head_dim,
|
||||||
|
qk_rope_head_dim=config.qk_rope_head_dim,
|
||||||
|
v_head_dim=config.v_head_dim,
|
||||||
|
q_lora_rank=config.q_lora_rank,
|
||||||
|
kv_lora_rank=config.kv_lora_rank,
|
||||||
|
rope_theta=rope_parameters["rope_theta"],
|
||||||
|
rope_scaling=(
|
||||||
|
None
|
||||||
|
if rope_parameters.get("rope_type") == "default"
|
||||||
|
else rope_parameters
|
||||||
|
),
|
||||||
|
max_position_embeddings=config.max_position_embeddings,
|
||||||
|
quant_config=quant_config,
|
||||||
|
reduce_results=True,
|
||||||
|
layer_id=layer_id,
|
||||||
|
prefix=prefix,
|
||||||
|
alt_stream=alt_stream,
|
||||||
|
is_nextn=is_nextn,
|
||||||
|
)
|
||||||
|
parallel = get_parallel()
|
||||||
|
self.linear_gate = ColumnParallelLinear(
|
||||||
|
config.hidden_size,
|
||||||
|
config.num_attention_heads * config.v_head_dim,
|
||||||
|
bias=False,
|
||||||
|
quant_config=quant_config,
|
||||||
|
tp_rank=parallel.attn_tp_rank,
|
||||||
|
tp_size=parallel.attn_tp_size,
|
||||||
|
prefix=f"{prefix}.linear_gate",
|
||||||
|
)
|
||||||
|
self.local_gate_width = self.num_local_heads * config.v_head_dim
|
||||||
|
if self.linear_gate.output_size_per_partition != self.local_gate_width:
|
||||||
|
raise ValueError(
|
||||||
|
"HYV4 attention gate shard width must match the local attention "
|
||||||
|
f"output width: {self.linear_gate.output_size_per_partition} != "
|
||||||
|
f"{self.local_gate_width}"
|
||||||
|
)
|
||||||
|
self.learnable_sink_param = nn.Parameter(
|
||||||
|
torch.empty(self.num_local_heads, dtype=torch.float32)
|
||||||
|
)
|
||||||
|
self.learnable_sink_param.weight_loader = self._sink_weight_loader
|
||||||
|
self._gate_fallback_backend = self._resolve_gate_fallback_backend()
|
||||||
|
self._gate_backend = self._resolve_gate_backend(
|
||||||
|
getattr(config, "gating_type", None), self._gate_fallback_backend
|
||||||
|
)
|
||||||
|
if prefix.endswith(".0.self_attn"):
|
||||||
|
logger.info("HYV4 MLA output gate backend: %s", self._gate_backend)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _sink_weight_loader(param, loaded_weight):
|
||||||
|
parallel = get_parallel()
|
||||||
|
heads = loaded_weight.shape[0] // parallel.attn_tp_size
|
||||||
|
start = parallel.attn_tp_rank * heads
|
||||||
|
param.data.copy_(loaded_weight[start : start + heads].float())
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _resolve_gate_fallback_backend() -> str:
|
||||||
|
try:
|
||||||
|
from sglang.kernels.ops.moe.triton_sigmoid_gate_mul import ( # noqa: F401
|
||||||
|
sigmoid_gate_mul,
|
||||||
|
)
|
||||||
|
except ImportError:
|
||||||
|
return "eager"
|
||||||
|
return "triton"
|
||||||
|
|
||||||
|
def _resolve_gate_backend(
|
||||||
|
self, gating_type: str | None, fallback_backend: str
|
||||||
|
) -> str:
|
||||||
|
weight = getattr(self.linear_gate, "weight", None)
|
||||||
|
if self._hpc_gated_mla_supported(
|
||||||
|
gating_type,
|
||||||
|
None if weight is None else weight.dtype,
|
||||||
|
None if weight is None else tuple(weight.shape),
|
||||||
|
self.local_gate_width,
|
||||||
|
self.hidden_size,
|
||||||
|
):
|
||||||
|
return "hpc"
|
||||||
|
return fallback_backend
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
@functools.lru_cache(maxsize=None)
|
||||||
|
def _hpc_gated_mla_supported(
|
||||||
|
gating_type: str | None,
|
||||||
|
weight_dtype: torch.dtype | None,
|
||||||
|
weight_shape: tuple[int, ...] | None,
|
||||||
|
local_gate_width: int,
|
||||||
|
hidden_size: int,
|
||||||
|
) -> bool:
|
||||||
|
try:
|
||||||
|
import hpc
|
||||||
|
except ImportError:
|
||||||
|
return False
|
||||||
|
|
||||||
|
if not hasattr(getattr(hpc, "gemm", None), "gated_mla_gemm"):
|
||||||
|
logger.info(
|
||||||
|
"HY4 gated MLA: the installed hpc build (%s) has no "
|
||||||
|
"gemm.gated_mla_gemm; falling back.",
|
||||||
|
getattr(hpc, "__version__", "unknown"),
|
||||||
|
)
|
||||||
|
return False
|
||||||
|
|
||||||
|
major, minor = get_device_capability()
|
||||||
|
if (major, minor) not in ((10, 0), (10, 3)):
|
||||||
|
logger.warning(
|
||||||
|
"HY4 gated MLA: hpc.gated_mla_gemm is built for sm100/sm103, "
|
||||||
|
"got sm%d%d; falling back.",
|
||||||
|
major,
|
||||||
|
minor,
|
||||||
|
)
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Headwise gating broadcasts one scalar per head over v_head_dim; the
|
||||||
|
# kernel only implements the elementwise product.
|
||||||
|
if gating_type != "elementwise":
|
||||||
|
return False
|
||||||
|
|
||||||
|
if weight_dtype != torch.bfloat16:
|
||||||
|
logger.warning(
|
||||||
|
"HY4 gated MLA: gate weight dtype is %s, the kernel needs "
|
||||||
|
"bfloat16 (keep linear_gate out of quantization).",
|
||||||
|
weight_dtype,
|
||||||
|
)
|
||||||
|
return False
|
||||||
|
expected_shape = (local_gate_width, hidden_size)
|
||||||
|
if weight_shape != expected_shape:
|
||||||
|
logger.warning(
|
||||||
|
"HYV4 gated MLA: local gate weight shape is %s, expected %s.",
|
||||||
|
weight_shape,
|
||||||
|
expected_shape,
|
||||||
|
)
|
||||||
|
return False
|
||||||
|
if local_gate_width % 256 != 0:
|
||||||
|
logger.warning(
|
||||||
|
"HY4 gated MLA: local gate width %s is not a multiple of 256.",
|
||||||
|
local_gate_width,
|
||||||
|
)
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
def _hpc_gated_mla_inputs_supported(self, hidden_states, attn_out) -> bool:
|
||||||
|
weight = self.linear_gate.weight
|
||||||
|
return (
|
||||||
|
hidden_states.dtype == torch.bfloat16
|
||||||
|
and weight.dtype == torch.bfloat16
|
||||||
|
and attn_out.dtype == torch.bfloat16
|
||||||
|
and hidden_states.ndim == 2
|
||||||
|
and weight.ndim == 2
|
||||||
|
and attn_out.ndim == 2
|
||||||
|
and hidden_states.shape[0] == attn_out.shape[0]
|
||||||
|
and hidden_states.shape[1] == weight.shape[1]
|
||||||
|
and weight.shape[0] == self.local_gate_width
|
||||||
|
and attn_out.shape[1] == self.local_gate_width
|
||||||
|
and hidden_states.device == weight.device == attn_out.device
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _apply_attention_output_gate_fallback(attn_out, gate, backend):
|
||||||
|
if gate.shape != attn_out.shape:
|
||||||
|
raise ValueError(
|
||||||
|
"HYV4 projected attention gate shape must match the local attention "
|
||||||
|
f"output shape: {tuple(gate.shape)} != {tuple(attn_out.shape)}"
|
||||||
|
)
|
||||||
|
if backend == "triton" and attn_out.is_cuda:
|
||||||
|
from sglang.kernels.ops.moe.triton_sigmoid_gate_mul import (
|
||||||
|
sigmoid_gate_mul,
|
||||||
|
)
|
||||||
|
|
||||||
|
return sigmoid_gate_mul(attn_out, gate)
|
||||||
|
return attn_out * torch.sigmoid(gate)
|
||||||
|
|
||||||
|
def prepare_attention_output_gate(self, hidden_states):
|
||||||
|
# The hpc kernel fuses the projection in, so hand it the activations and
|
||||||
|
# let apply_attention_output_gate do the GEMM. The other two tiers want
|
||||||
|
# the projected gate.
|
||||||
|
if self._gate_backend == "hpc":
|
||||||
|
return hidden_states
|
||||||
|
return self.linear_gate(hidden_states)[0]
|
||||||
|
|
||||||
|
def apply_attention_output_gate(self, attn_out, gate):
|
||||||
|
if self._gate_backend == "hpc" and self._hpc_gated_mla_inputs_supported(
|
||||||
|
gate, attn_out
|
||||||
|
):
|
||||||
|
from hpc.gemm import gated_mla_gemm
|
||||||
|
|
||||||
|
# linear_gate is column-parallel and unbiased, so the local weight
|
||||||
|
# shard lines up with the local attn_out columns.
|
||||||
|
return gated_mla_gemm(
|
||||||
|
gate.contiguous(), self.linear_gate.weight, attn_out.contiguous()
|
||||||
|
)
|
||||||
|
if self._gate_backend == "hpc":
|
||||||
|
gate = self.linear_gate(gate)[0]
|
||||||
|
backend = self._gate_fallback_backend
|
||||||
|
else:
|
||||||
|
backend = self._gate_backend
|
||||||
|
return self._apply_attention_output_gate_fallback(attn_out, gate, backend)
|
||||||
|
|
||||||
|
def dispatch_attn_forward_method(self, forward_batch):
|
||||||
|
backend = get_attn_backend()
|
||||||
|
backend = getattr(backend, "primary", backend)
|
||||||
|
if hasattr(backend, "use_mha") and backend.use_mha is not False:
|
||||||
|
backend.use_mha = False
|
||||||
|
method = super().dispatch_attn_forward_method(forward_batch)
|
||||||
|
if method != AttnForwardMethod.MLA:
|
||||||
|
raise RuntimeError("HYV4 requires the sparse MLA attention path")
|
||||||
|
return method
|
||||||
|
|
||||||
|
|
||||||
|
class HYV4DecoderLayer(nn.Module):
|
||||||
|
def __init__(self, config, layer_id, quant_config=None, prefix="", alt_stream=None):
|
||||||
|
super().__init__()
|
||||||
|
self.self_attn = HYV4Attention(
|
||||||
|
config, layer_id, quant_config, f"{prefix}.self_attn", alt_stream
|
||||||
|
)
|
||||||
|
self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
|
||||||
|
self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
|
||||||
|
if config.mlp_layer_types[layer_id] == "dense":
|
||||||
|
self.mlp = DeepseekV2MLP(
|
||||||
|
config.hidden_size,
|
||||||
|
config.intermediate_size,
|
||||||
|
config.hidden_act,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=f"{prefix}.mlp",
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.mlp = DeepseekV2MoE(
|
||||||
|
config,
|
||||||
|
layer_id,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=f"{prefix}.mlp",
|
||||||
|
alt_stream=alt_stream,
|
||||||
|
)
|
||||||
|
if hasattr(self.mlp, "shared_experts"):
|
||||||
|
self.mlp.shared_experts.swiglu_limit = None
|
||||||
|
self.hc_attn_layer = HYV4HCLayer(config, f"{prefix}.hc_attn_layer")
|
||||||
|
self.hc_mlp_layer = HYV4HCLayer(config, f"{prefix}.hc_mlp_layer")
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
positions,
|
||||||
|
hidden_states,
|
||||||
|
forward_batch,
|
||||||
|
zero_allocator,
|
||||||
|
prev_topk_indices=None,
|
||||||
|
):
|
||||||
|
hidden_states = self.hc_attn_layer.prepare_input(hidden_states)
|
||||||
|
hidden_states, post, residual = self.hc_attn_layer.pre(
|
||||||
|
hidden_states, self.input_layernorm
|
||||||
|
)
|
||||||
|
get_attn_tp_context().set_attn_inputs(
|
||||||
|
AttentionInputs(
|
||||||
|
hidden_states, forward_batch, self.self_attn.prepare_qkv_latent
|
||||||
|
)
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
hidden_states = self.self_attn(
|
||||||
|
positions,
|
||||||
|
hidden_states,
|
||||||
|
forward_batch,
|
||||||
|
zero_allocator,
|
||||||
|
prev_topk_indices=prev_topk_indices,
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
get_attn_tp_context().clear_attn_inputs()
|
||||||
|
if isinstance(hidden_states, tuple):
|
||||||
|
hidden_states, topk_indices = hidden_states
|
||||||
|
else:
|
||||||
|
topk_indices = None
|
||||||
|
hidden_states, post, residual = self.hc_attn_layer.post_pre(
|
||||||
|
hidden_states,
|
||||||
|
residual,
|
||||||
|
post,
|
||||||
|
self.hc_mlp_layer,
|
||||||
|
self.post_attention_layernorm,
|
||||||
|
)
|
||||||
|
if isinstance(self.mlp, DeepseekV2MoE):
|
||||||
|
hidden_states = self.mlp(hidden_states, forward_batch)
|
||||||
|
else:
|
||||||
|
hidden_states = self.mlp(hidden_states)
|
||||||
|
hidden_states = self.hc_mlp_layer.post(hidden_states, residual, post)
|
||||||
|
return hidden_states, topk_indices
|
||||||
|
|
||||||
|
|
||||||
|
class HYV4Model(nn.Module):
|
||||||
|
def __init__(self, config, quant_config=None, prefix=""):
|
||||||
|
super().__init__()
|
||||||
|
if get_pp_group().world_size != 1:
|
||||||
|
raise ValueError("HYV4 pipeline parallelism is not supported")
|
||||||
|
self.config = config
|
||||||
|
self.start_layer = 0
|
||||||
|
self.end_layer = config.num_hidden_layers
|
||||||
|
self.embed_tokens = VocabParallelEmbedding(
|
||||||
|
config.vocab_size,
|
||||||
|
config.hidden_size,
|
||||||
|
prefix=f"{prefix}.embed_tokens",
|
||||||
|
**get_embedding_tp_kwargs(),
|
||||||
|
)
|
||||||
|
self.alt_stream = get_stream("alt") if is_cuda() else None
|
||||||
|
self.layers = nn.ModuleList(
|
||||||
|
[
|
||||||
|
HYV4DecoderLayer(
|
||||||
|
config,
|
||||||
|
i,
|
||||||
|
quant_config,
|
||||||
|
f"{prefix}.layers.{i}",
|
||||||
|
self.alt_stream,
|
||||||
|
)
|
||||||
|
for i in range(config.num_hidden_layers)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
self.hc_head = HYV4HCHeadLayer(config, f"{prefix}.hc_head")
|
||||||
|
self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
|
||||||
|
|
||||||
|
def forward(self, input_ids, positions, forward_batch, input_embeds=None):
|
||||||
|
hidden_states = (
|
||||||
|
self.embed_tokens(input_ids) if input_embeds is None else input_embeds
|
||||||
|
)
|
||||||
|
zero_allocator = BumpAllocator(
|
||||||
|
buffer_size=2 * len(self.layers),
|
||||||
|
dtype=torch.float32,
|
||||||
|
device=hidden_states.device,
|
||||||
|
)
|
||||||
|
topk_share = IndexTopKShareState(forward_batch, None)
|
||||||
|
for layer in self.layers:
|
||||||
|
hidden_states, topk_indices = layer(
|
||||||
|
positions,
|
||||||
|
hidden_states,
|
||||||
|
forward_batch,
|
||||||
|
zero_allocator,
|
||||||
|
topk_share.topk_indices,
|
||||||
|
)
|
||||||
|
topk_share.update(topk_indices)
|
||||||
|
topk_share.publish()
|
||||||
|
return self.hc_head(hidden_states, self.norm)
|
||||||
|
|
||||||
|
|
||||||
|
class HYV4ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
|
||||||
|
packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
|
||||||
|
return "HYV4 applies the SwiGLU limit only to routed experts."
|
||||||
|
|
||||||
|
def __init__(self, config, quant_config=None, prefix=""):
|
||||||
|
super().__init__()
|
||||||
|
self.config = config
|
||||||
|
self.quant_config = quant_config
|
||||||
|
self.pp_group = get_pp_group()
|
||||||
|
self.model = HYV4Model(config, quant_config, f"{prefix}.model")
|
||||||
|
self.num_fused_shared_experts = max(
|
||||||
|
(
|
||||||
|
layer.mlp.num_fused_shared_experts
|
||||||
|
for layer in self.model.layers
|
||||||
|
if isinstance(layer.mlp, DeepseekV2MoE)
|
||||||
|
),
|
||||||
|
default=0,
|
||||||
|
)
|
||||||
|
self.lm_head = ParallelLMHead(
|
||||||
|
config.vocab_size,
|
||||||
|
config.hidden_size,
|
||||||
|
# Keep checkpoint weights in bf16; LogitsProcessor emits fp32
|
||||||
|
# logits when config.enable_lm_head_fp32 is set.
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=f"{prefix}.lm_head",
|
||||||
|
use_attn_tp_group=get_parallel().enable_dp_lm_head,
|
||||||
|
)
|
||||||
|
self.logits_processor = LogitsProcessor(config)
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def forward(self, input_ids, positions, forward_batch, input_embeds=None):
|
||||||
|
hidden_states = self.model(
|
||||||
|
input_ids, positions, forward_batch, input_embeds=input_embeds
|
||||||
|
)
|
||||||
|
return self.logits_processor(
|
||||||
|
input_ids, hidden_states, self.lm_head, forward_batch
|
||||||
|
)
|
||||||
|
|
||||||
|
def get_embed_and_head(self):
|
||||||
|
return self.model.embed_tokens.weight, self.lm_head.weight
|
||||||
|
|
||||||
|
def set_embed_and_head(self, embed, head):
|
||||||
|
del self.model.embed_tokens.weight
|
||||||
|
del self.lm_head.weight
|
||||||
|
self.model.embed_tokens.weight = embed
|
||||||
|
self.lm_head.weight = head
|
||||||
|
|
||||||
|
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
|
||||||
|
def mapped_weights():
|
||||||
|
for name, loaded_weight in weights:
|
||||||
|
if name.startswith("model.mtp_layers."):
|
||||||
|
continue
|
||||||
|
loaded_weight = permute_hyv4_indexer_weight(
|
||||||
|
name, loaded_weight, self.config
|
||||||
|
)
|
||||||
|
if name.endswith((".hc_fn", ".hc_head_fn")):
|
||||||
|
name += ".weight"
|
||||||
|
if name.endswith(".weight_scale"):
|
||||||
|
name += "_inv"
|
||||||
|
yield name, loaded_weight
|
||||||
|
|
||||||
|
self.do_load_weights(mapped_weights())
|
||||||
|
|
||||||
|
|
||||||
|
EntryClass = [HYV4ForCausalLM]
|
||||||
@@ -0,0 +1,229 @@
|
|||||||
|
import copy
|
||||||
|
from typing import Iterable, Tuple
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch import nn
|
||||||
|
|
||||||
|
from sglang.srt.distributed import get_pp_group
|
||||||
|
from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState
|
||||||
|
from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context
|
||||||
|
from sglang.srt.layers.layernorm import RMSNorm
|
||||||
|
from sglang.srt.layers.logits_processor import LogitsProcessor
|
||||||
|
from sglang.srt.layers.vocab_parallel_embedding import (
|
||||||
|
ParallelLMHead,
|
||||||
|
VocabParallelEmbedding,
|
||||||
|
get_embedding_tp_kwargs,
|
||||||
|
)
|
||||||
|
from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
|
||||||
|
DeepseekV2WeightLoaderMixin,
|
||||||
|
)
|
||||||
|
from sglang.srt.models.deepseek_v2 import DeepseekV2MoE
|
||||||
|
from sglang.srt.models.hunyuan_v4 import (
|
||||||
|
HYV4Attention,
|
||||||
|
HYV4ForCausalLM,
|
||||||
|
permute_hyv4_indexer_weight,
|
||||||
|
)
|
||||||
|
from sglang.srt.runtime_context import get_parallel, get_stream
|
||||||
|
from sglang.srt.utils import BumpAllocator, is_cuda
|
||||||
|
|
||||||
|
|
||||||
|
def _mtp_quant_config(quant_config):
|
||||||
|
if quant_config is None:
|
||||||
|
return None
|
||||||
|
quant_config = copy.deepcopy(quant_config)
|
||||||
|
ignored_layers = getattr(quant_config, "ignored_layers", None)
|
||||||
|
if ignored_layers is not None:
|
||||||
|
quant_config.ignored_layers = [
|
||||||
|
name.replace("model.mtp_layers.0", "model.decoder").replace(
|
||||||
|
"mtp_layers.0", "model.decoder"
|
||||||
|
)
|
||||||
|
for name in ignored_layers
|
||||||
|
]
|
||||||
|
return quant_config
|
||||||
|
|
||||||
|
|
||||||
|
class HYV4MTPDecoderLayer(nn.Module):
|
||||||
|
def __init__(self, config, quant_config=None, prefix="", alt_stream=None):
|
||||||
|
super().__init__()
|
||||||
|
self.self_attn = HYV4Attention(
|
||||||
|
config,
|
||||||
|
0,
|
||||||
|
quant_config,
|
||||||
|
f"{prefix}.self_attn",
|
||||||
|
alt_stream,
|
||||||
|
is_nextn=True,
|
||||||
|
)
|
||||||
|
self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
|
||||||
|
self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
|
||||||
|
self.mlp = DeepseekV2MoE(
|
||||||
|
config,
|
||||||
|
0,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=f"{prefix}.mlp",
|
||||||
|
alt_stream=alt_stream,
|
||||||
|
is_nextn=True,
|
||||||
|
)
|
||||||
|
if hasattr(self.mlp, "shared_experts"):
|
||||||
|
self.mlp.shared_experts.swiglu_limit = None
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
positions,
|
||||||
|
hidden_states,
|
||||||
|
forward_batch,
|
||||||
|
zero_allocator,
|
||||||
|
prev_topk_indices=None,
|
||||||
|
):
|
||||||
|
residual = hidden_states
|
||||||
|
hidden_states = self.input_layernorm(hidden_states)
|
||||||
|
get_attn_tp_context().set_attn_inputs(
|
||||||
|
AttentionInputs(
|
||||||
|
hidden_states, forward_batch, self.self_attn.prepare_qkv_latent
|
||||||
|
)
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
hidden_states = self.self_attn(
|
||||||
|
positions,
|
||||||
|
hidden_states,
|
||||||
|
forward_batch,
|
||||||
|
zero_allocator,
|
||||||
|
prev_topk_indices=prev_topk_indices,
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
get_attn_tp_context().clear_attn_inputs()
|
||||||
|
if isinstance(hidden_states, tuple):
|
||||||
|
hidden_states, topk_indices = hidden_states
|
||||||
|
else:
|
||||||
|
topk_indices = None
|
||||||
|
hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
|
||||||
|
hidden_states = self.mlp(hidden_states, forward_batch)
|
||||||
|
return hidden_states, residual, topk_indices
|
||||||
|
|
||||||
|
|
||||||
|
class HYV4ModelNextN(nn.Module):
|
||||||
|
def __init__(self, config, quant_config=None, prefix=""):
|
||||||
|
super().__init__()
|
||||||
|
self.embed_tokens = VocabParallelEmbedding(
|
||||||
|
config.vocab_size,
|
||||||
|
config.hidden_size,
|
||||||
|
prefix=f"{prefix}.embed_tokens",
|
||||||
|
**get_embedding_tp_kwargs(),
|
||||||
|
)
|
||||||
|
self.enorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
|
||||||
|
self.hnorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
|
||||||
|
self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
|
||||||
|
self.alt_stream = get_stream("alt") if is_cuda() else None
|
||||||
|
self.decoder = HYV4MTPDecoderLayer(
|
||||||
|
config,
|
||||||
|
quant_config,
|
||||||
|
f"{prefix}.decoder",
|
||||||
|
self.alt_stream,
|
||||||
|
)
|
||||||
|
self.shared_head = nn.Module()
|
||||||
|
self.shared_head.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
|
||||||
|
|
||||||
|
def forward(self, input_ids, positions, forward_batch, input_embeds=None):
|
||||||
|
hidden_states = (
|
||||||
|
self.embed_tokens(input_ids) if input_embeds is None else input_embeds
|
||||||
|
)
|
||||||
|
if hidden_states.shape[0] > 0:
|
||||||
|
hidden_states = self.eh_proj(
|
||||||
|
torch.cat(
|
||||||
|
(
|
||||||
|
self.enorm(hidden_states),
|
||||||
|
self.hnorm(forward_batch.spec_info.hidden_states),
|
||||||
|
),
|
||||||
|
dim=-1,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
zero_allocator = BumpAllocator(
|
||||||
|
buffer_size=2,
|
||||||
|
dtype=torch.float32,
|
||||||
|
device=hidden_states.device,
|
||||||
|
)
|
||||||
|
topk_share = IndexTopKShareState.from_mtp_carry(forward_batch)
|
||||||
|
hidden_states, residual, topk_indices = self.decoder(
|
||||||
|
positions,
|
||||||
|
hidden_states,
|
||||||
|
forward_batch,
|
||||||
|
zero_allocator,
|
||||||
|
topk_share.topk_indices,
|
||||||
|
)
|
||||||
|
topk_share.update(topk_indices)
|
||||||
|
topk_share.publish()
|
||||||
|
if forward_batch.forward_mode.is_idle():
|
||||||
|
return hidden_states
|
||||||
|
hidden_states, _ = self.shared_head.norm(hidden_states, residual)
|
||||||
|
return hidden_states
|
||||||
|
|
||||||
|
|
||||||
|
class HYV4ForCausalLMNextN(nn.Module, DeepseekV2WeightLoaderMixin):
|
||||||
|
packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def shared_experts_fusion_disable_reason(hf_config, quant_config):
|
||||||
|
return HYV4ForCausalLM.shared_experts_fusion_disable_reason(
|
||||||
|
hf_config, quant_config
|
||||||
|
)
|
||||||
|
|
||||||
|
def __init__(self, config, quant_config=None, prefix=""):
|
||||||
|
super().__init__()
|
||||||
|
self.config = config
|
||||||
|
self.quant_config = quant_config
|
||||||
|
self.pp_group = get_pp_group()
|
||||||
|
nextn_quant_config = _mtp_quant_config(quant_config)
|
||||||
|
self.model = HYV4ModelNextN(
|
||||||
|
config, nextn_quant_config, prefix=f"{prefix}.model"
|
||||||
|
)
|
||||||
|
self.num_fused_shared_experts = self.model.decoder.mlp.num_fused_shared_experts
|
||||||
|
self.lm_head = ParallelLMHead(
|
||||||
|
config.vocab_size,
|
||||||
|
config.hidden_size,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=f"{prefix}.lm_head",
|
||||||
|
use_attn_tp_group=get_parallel().enable_dp_lm_head,
|
||||||
|
)
|
||||||
|
self.logits_processor = LogitsProcessor(config)
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def forward(self, input_ids, positions, forward_batch):
|
||||||
|
hidden_states = self.model(input_ids, positions, forward_batch)
|
||||||
|
return self.logits_processor(
|
||||||
|
input_ids, hidden_states, self.lm_head, forward_batch
|
||||||
|
)
|
||||||
|
|
||||||
|
def get_embed_and_head(self):
|
||||||
|
return self.model.embed_tokens.weight, self.lm_head.weight
|
||||||
|
|
||||||
|
def set_embed_and_head(self, embed, head):
|
||||||
|
del self.model.embed_tokens.weight
|
||||||
|
del self.lm_head.weight
|
||||||
|
self.model.embed_tokens.weight = embed
|
||||||
|
self.lm_head.weight = head
|
||||||
|
|
||||||
|
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
|
||||||
|
layer_prefix = f"model.layers.{self.config.num_hidden_layers}"
|
||||||
|
|
||||||
|
def mapped_weights():
|
||||||
|
for name, loaded_weight in weights:
|
||||||
|
if not name.startswith("model.mtp_layers.0."):
|
||||||
|
continue
|
||||||
|
name = name.replace("model.mtp_layers.0", layer_prefix)
|
||||||
|
if name.endswith(".final_layernorm.weight"):
|
||||||
|
name = name.replace(
|
||||||
|
".final_layernorm.weight", ".shared_head.norm.weight"
|
||||||
|
)
|
||||||
|
loaded_weight = permute_hyv4_indexer_weight(
|
||||||
|
name, loaded_weight, self.config
|
||||||
|
)
|
||||||
|
if name.endswith(".weight_scale"):
|
||||||
|
name += "_inv"
|
||||||
|
yield name, loaded_weight
|
||||||
|
|
||||||
|
self.do_load_weights(mapped_weights(), is_nextn=True)
|
||||||
|
|
||||||
|
def post_load_weights(self, is_nextn=True, weight_names=None):
|
||||||
|
super().post_load_weights(is_nextn=True, weight_names=weight_names)
|
||||||
|
|
||||||
|
|
||||||
|
EntryClass = [HYV4ForCausalLMNextN]
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
|
||||||
|
from sglang.srt.parser.template_detection import ReasoningToggleConfig
|
||||||
|
|
||||||
|
|
||||||
|
def uses_hunyuan_reasoning_effort(
|
||||||
|
reasoning_parser: Optional[str], reasoning_config: Optional[ReasoningToggleConfig]
|
||||||
|
) -> bool:
|
||||||
|
return (
|
||||||
|
reasoning_parser == "hunyuan"
|
||||||
|
and reasoning_config is not None
|
||||||
|
and reasoning_config.special_case == "hunyuan_effort"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def normalize_hunyuan_reasoning_effort(
|
||||||
|
request: ChatCompletionRequest,
|
||||||
|
reasoning_parser: Optional[str],
|
||||||
|
reasoning_config: Optional[ReasoningToggleConfig],
|
||||||
|
) -> None:
|
||||||
|
if not uses_hunyuan_reasoning_effort(reasoning_parser, reasoning_config):
|
||||||
|
return
|
||||||
|
|
||||||
|
effort = request.reasoning_effort
|
||||||
|
if effort is None and request.chat_template_kwargs is not None:
|
||||||
|
effort = request.chat_template_kwargs.get("reasoning_effort")
|
||||||
|
if effort is None:
|
||||||
|
normalized_effort = "high"
|
||||||
|
elif effort in ("none", "no_think"):
|
||||||
|
normalized_effort = "no_think"
|
||||||
|
elif effort in ("minimal", "low"):
|
||||||
|
normalized_effort = "low"
|
||||||
|
elif effort in ("medium", "high", "xhigh", "max"):
|
||||||
|
normalized_effort = "high"
|
||||||
|
else:
|
||||||
|
raise ValueError(
|
||||||
|
"Hunyuan reasoning_effort must be one of none, minimal, low, "
|
||||||
|
"medium, high, xhigh, or max"
|
||||||
|
)
|
||||||
|
|
||||||
|
request.reasoning_effort = normalized_effort
|
||||||
|
if request.chat_template_kwargs is not None:
|
||||||
|
request.chat_template_kwargs.pop("reasoning_effort", None)
|
||||||
@@ -175,6 +175,15 @@ REASONING_MODE_RULES = (
|
|||||||
ctx.has_text("reasoning_effort") and ctx.has_text("[THINK]")
|
ctx.has_text("reasoning_effort") and ctx.has_text("[THINK]")
|
||||||
),
|
),
|
||||||
),
|
),
|
||||||
|
DetectionRule(
|
||||||
|
name="hunyuan_reasoning_effort",
|
||||||
|
value=ReasoningToggleConfig(special_case="hunyuan_effort"),
|
||||||
|
predicate=lambda ctx: (
|
||||||
|
ctx.has_text("reasoning_effort")
|
||||||
|
and ctx.has_text("reasoning_mode_token")
|
||||||
|
and ctx.has_text("no_think")
|
||||||
|
),
|
||||||
|
),
|
||||||
DetectionRule(
|
DetectionRule(
|
||||||
name="explicit_enable_thinking_default_false",
|
name="explicit_enable_thinking_default_false",
|
||||||
value=ReasoningToggleConfig(
|
value=ReasoningToggleConfig(
|
||||||
@@ -364,8 +373,16 @@ def _is_hunyuan(ctx):
|
|||||||
sep = ctx.has_text("<tool_sep>") or ctx.has_vocab_pattern(
|
sep = ctx.has_text("<tool_sep>") or ctx.has_vocab_pattern(
|
||||||
r"^<tool_sep(?::[^>]+)?>$"
|
r"^<tool_sep(?::[^>]+)?>$"
|
||||||
)
|
)
|
||||||
return (tc and sep) or (
|
return (
|
||||||
ctx.has_text("reasoning_effort") and ctx.has_text("interleaved_thinking")
|
(tc and sep)
|
||||||
|
or (
|
||||||
|
tc
|
||||||
|
and ctx.reasoning_config
|
||||||
|
== ReasoningToggleConfig(special_case="hunyuan_effort")
|
||||||
|
and ctx.has_vocab_pattern(r"^<arg_key(?::[^>]+)?>$")
|
||||||
|
and ctx.has_vocab_pattern(r"^<arg_value(?::[^>]+)?>$")
|
||||||
|
)
|
||||||
|
or (ctx.has_text("reasoning_effort") and ctx.has_text("interleaved_thinking"))
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -38,6 +38,7 @@ from sglang.srt.configs import (
|
|||||||
ExaoneConfig,
|
ExaoneConfig,
|
||||||
FalconH1Config,
|
FalconH1Config,
|
||||||
GraniteMoeHybridConfig,
|
GraniteMoeHybridConfig,
|
||||||
|
HYV4Config,
|
||||||
InklingAudioConfig,
|
InklingAudioConfig,
|
||||||
InklingMMConfig,
|
InklingMMConfig,
|
||||||
InklingModelConfig,
|
InklingModelConfig,
|
||||||
@@ -122,6 +123,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
|
|||||||
Qwen3NextConfig,
|
Qwen3NextConfig,
|
||||||
FalconH1Config,
|
FalconH1Config,
|
||||||
GraniteMoeHybridConfig,
|
GraniteMoeHybridConfig,
|
||||||
|
HYV4Config,
|
||||||
DotsVLMConfig,
|
DotsVLMConfig,
|
||||||
DotsOCRConfig,
|
DotsOCRConfig,
|
||||||
Dots3Config,
|
Dots3Config,
|
||||||
|
|||||||
@@ -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",
|
"use_symmetric_memory",
|
||||||
lambda *args, **kwargs: nullcontext(),
|
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
|
# UE8M0 packs four 128-wide scale groups into each int32. Use the smallest
|
||||||
# legal K for both the gate/up and down GEMMs.
|
# 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,
|
).hidden_states,
|
||||||
)
|
)
|
||||||
|
|
||||||
compact_is_masked, compact_all_tokens, compact_m_indices, compact_output = (
|
(
|
||||||
run_with_layout("compact")
|
compact_is_masked,
|
||||||
)
|
compact_all_tokens,
|
||||||
masked_is_masked, masked_all_tokens, masked_m_indices, masked_output = (
|
compact_m_indices,
|
||||||
run_with_layout("masked")
|
compact_output,
|
||||||
)
|
) = run_with_layout("compact")
|
||||||
|
(
|
||||||
|
masked_is_masked,
|
||||||
|
masked_all_tokens,
|
||||||
|
masked_m_indices,
|
||||||
|
masked_output,
|
||||||
|
) = run_with_layout("masked")
|
||||||
torch.cuda.synchronize()
|
torch.cuda.synchronize()
|
||||||
|
|
||||||
assert not compact_is_masked
|
assert not compact_is_masked
|
||||||
assert masked_is_masked
|
assert masked_is_masked
|
||||||
assert compact_all_tokens == 256
|
assert compact_all_tokens == 64
|
||||||
assert masked_all_tokens is None
|
assert masked_all_tokens is None
|
||||||
assert masked_m_indices is None
|
assert masked_m_indices is None
|
||||||
valid_assignments = topk_ids[topk_ids >= 0]
|
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
|
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 (
|
from sglang.srt.layers.quantization.fp8_utils import (
|
||||||
inverse_transform_scale_ue8m0,
|
inverse_transform_scale_ue8m0,
|
||||||
quant_weight_ue8m0,
|
quant_weight_ue8m0,
|
||||||
transform_scale_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.ci.ci_register import register_cuda_ci
|
||||||
from sglang.test.test_utils import CustomTestCase
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
register_cuda_ci(est_time=12, stage="base-b", runner_config="1-gpu-large")
|
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):
|
class TestInverseTransformScaleUe8m0(CustomTestCase):
|
||||||
def test_round_trip(self):
|
def test_round_trip(self):
|
||||||
for _ in range(100):
|
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.runtime_context import get_context
|
||||||
from sglang.srt.utils import get_device
|
from sglang.srt.utils import get_device
|
||||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
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_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")
|
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): ...
|
def compute_dp_attention_metadata(self): ...
|
||||||
|
|
||||||
|
|
||||||
class TestLMHeadFP32(unittest.TestCase):
|
class TestLMHeadFP32(CustomTestCase):
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls):
|
def setUpClass(cls):
|
||||||
if not torch.cuda.is_available() and not (
|
if not torch.cuda.is_available() and not (
|
||||||
@@ -38,13 +39,17 @@ class TestLMHeadFP32(unittest.TestCase):
|
|||||||
):
|
):
|
||||||
raise unittest.SkipTest("needs CUDA GPU or XPU")
|
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
|
# LogitsProcessor reads get_exec().features.enable_fp32_lm_head
|
||||||
# from the published config.
|
# from the published config.
|
||||||
override = get_context().override_server_args(enable_fp32_lm_head=enable_fp32)
|
override = get_context().override_server_args(enable_fp32_lm_head=enable_fp32)
|
||||||
override.install()
|
override.install()
|
||||||
self.addCleanup(override.restore)
|
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)
|
return LogitsProcessor(cfg, skip_all_gather=True, logit_scale=None)
|
||||||
|
|
||||||
def _run_case(
|
def _run_case(
|
||||||
@@ -55,6 +60,7 @@ class TestLMHeadFP32(unittest.TestCase):
|
|||||||
expected_a_dtype,
|
expected_a_dtype,
|
||||||
expected_b_dtype,
|
expected_b_dtype,
|
||||||
expected_operation,
|
expected_operation,
|
||||||
|
config_enable_fp32=False,
|
||||||
):
|
):
|
||||||
device = get_device()
|
device = get_device()
|
||||||
BATCH_SIZE, HIDDEN_SIZE, VOCAB_SIZE = 2, 64, 128
|
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)
|
head = LMHeadStub(VOCAB_SIZE, HIDDEN_SIZE, dtype=weights_dtype, device=device)
|
||||||
meta = DummyMeta()
|
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_matmul = torch.matmul
|
||||||
original_mm = torch.mm
|
original_mm = torch.mm
|
||||||
@@ -173,6 +181,21 @@ class TestLMHeadFP32(unittest.TestCase):
|
|||||||
"matmul",
|
"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__":
|
if __name__ == "__main__":
|
||||||
unittest.main(verbosity=2)
|
unittest.main(verbosity=2)
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
import math
|
import math
|
||||||
import unittest
|
import unittest
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
from sglang.srt.configs.model_config import ModelConfig, compute_mla_mscale_scaling
|
from sglang.srt.configs.model_config import ModelConfig, compute_mla_mscale_scaling
|
||||||
from sglang.test.ci.ci_register import register_cpu_ci
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
@@ -92,6 +93,31 @@ class TestInitMlaScaling(CustomTestCase):
|
|||||||
base, scaling = _mla_scaling(None)
|
base, scaling = _mla_scaling(None)
|
||||||
self.assertEqual(scaling, base)
|
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__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
@@ -800,6 +800,87 @@ class ServingChatTestCase(unittest.TestCase):
|
|||||||
|
|
||||||
self.assertEqual(req.reasoning_effort, "high")
|
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):
|
def test_k2_selected_terminator_reaches_sampling_params(self):
|
||||||
self.tm._config_overrides["reasoning_parser"] = "k2_horizon"
|
self.tm._config_overrides["reasoning_parser"] = "k2_horizon"
|
||||||
self.chat = OpenAIServingChat(self.tm, self.template_manager)
|
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):
|
class TestHunyuanDetectorHasToolCall(CustomTestCase):
|
||||||
def setUp(self):
|
def setUp(self):
|
||||||
self.detector = HunyuanDetector()
|
self.detector = HunyuanDetector()
|
||||||
@@ -111,6 +122,21 @@ class TestHunyuanDetectorDetectAndParse(CustomTestCase):
|
|||||||
self.assertEqual(len(result.calls), 0)
|
self.assertEqual(len(result.calls), 0)
|
||||||
self.assertEqual(result.normal_text, text)
|
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):
|
def test_zero_arg_inline(self):
|
||||||
text = (
|
text = (
|
||||||
"<tool_calls><tool_call>get_current_date<tool_sep></tool_call></tool_calls>"
|
"<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)
|
args = json.loads(result.calls[0].parameters)
|
||||||
self.assertIs(args["verbose"], True)
|
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):
|
def test_string_arg_not_deserialized(self):
|
||||||
"""String-typed args should stay as strings even if they look like JSON."""
|
"""String-typed args should stay as strings even if they look like JSON."""
|
||||||
text = (
|
text = (
|
||||||
@@ -359,6 +422,23 @@ class TestHunyuanDetectorStreaming(CustomTestCase):
|
|||||||
self.assertEqual(collected[0]["name"], "get_current_date")
|
self.assertEqual(collected[0]["name"], "get_current_date")
|
||||||
self.assertEqual(json.loads(collected[0]["parameters"]), {})
|
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):
|
def test_chunked_tool_call(self):
|
||||||
detector = self._new_detector()
|
detector = self._new_detector()
|
||||||
chunks = [
|
chunks = [
|
||||||
|
|||||||
@@ -552,6 +552,21 @@ class TestParseQuantHfConfig(CustomTestCase):
|
|||||||
self.assertIn("lm_head", quant_config.ignored_layers)
|
self.assertIn("lm_head", quant_config.ignored_layers)
|
||||||
self.assertEqual(quant_config.kv_cache_quant_algo, "FP8")
|
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):
|
def test_modelopt_mxfp8_override(self):
|
||||||
"""Generic ModelOpt selection must not route MXFP8 to scalar FP8."""
|
"""Generic ModelOpt selection must not route MXFP8 to scalar FP8."""
|
||||||
self.assertEqual(
|
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>"])
|
_, _, parser = self._detect(template, ["<minimax:tool_call>"])
|
||||||
self.assertEqual(parser, "minimax")
|
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):
|
class TestTemplateDetectionRuleMatrix(unittest.TestCase):
|
||||||
"""Table-driven tests for REASONING_PARSER_RULES and REASONING_MODE_RULES."""
|
"""Table-driven tests for REASONING_PARSER_RULES and REASONING_MODE_RULES."""
|
||||||
|
|||||||
@@ -1037,6 +1037,7 @@ class TestContextParallelServerArgs(CustomTestCase):
|
|||||||
moe_dp_size=1,
|
moe_dp_size=1,
|
||||||
ep_size=1,
|
ep_size=1,
|
||||||
pp_size=1,
|
pp_size=1,
|
||||||
|
dcp_size=1,
|
||||||
enable_aiter_allreduce_fusion=False,
|
enable_aiter_allreduce_fusion=False,
|
||||||
)
|
)
|
||||||
defaults.update(overrides)
|
defaults.update(overrides)
|
||||||
@@ -1169,6 +1170,29 @@ class TestContextParallelServerArgs(CustomTestCase):
|
|||||||
resolution_result(args, "dsa_prefill_cp_mode"), "round-robin-split"
|
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):
|
def test_context_parallel_handler_initializes_cp_strategy(self):
|
||||||
server_args = self._new_cp_args(
|
server_args = self._new_cp_args(
|
||||||
enable_prefill_cp=True,
|
enable_prefill_cp=True,
|
||||||
|
|||||||
@@ -1817,8 +1817,8 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
|||||||
_dsa_split_backend_resolution,
|
_dsa_split_backend_resolution,
|
||||||
)
|
)
|
||||||
|
|
||||||
def _view(arch="DeepseekV32ForCausalLM", **kw):
|
def _view(arch="DeepseekV32ForCausalLM", learnable_sink=False, **kw):
|
||||||
hf = SimpleNamespace(architectures=[arch])
|
hf = SimpleNamespace(architectures=[arch], learnable_sink=learnable_sink)
|
||||||
defaults = dict(
|
defaults = dict(
|
||||||
kv_cache_dtype="fp8_e4m3",
|
kv_cache_dtype="fp8_e4m3",
|
||||||
dsa_prefill_backend=None,
|
dsa_prefill_backend=None,
|
||||||
@@ -1866,6 +1866,28 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
|||||||
"dsa_decode_backend": "flashmla_kv",
|
"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
|
# non-family arch declares nothing
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
_dsa_split_backend_resolution(_view(arch="LlamaForCausalLM")), {}
|
_dsa_split_backend_resolution(_view(arch="LlamaForCausalLM")), {}
|
||||||
@@ -2842,6 +2864,7 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
|||||||
prefill_attention_backend=None,
|
prefill_attention_backend=None,
|
||||||
decode_attention_backend=None,
|
decode_attention_backend=None,
|
||||||
enable_prefill_cp=False,
|
enable_prefill_cp=False,
|
||||||
|
dcp_size=1,
|
||||||
)
|
)
|
||||||
defaults.update(kw)
|
defaults.update(kw)
|
||||||
return SimpleNamespace(**defaults)
|
return SimpleNamespace(**defaults)
|
||||||
@@ -2857,6 +2880,22 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
|||||||
_deepseek_family_overrides(_args(), None),
|
_deepseek_family_overrides(_args(), None),
|
||||||
{"attention_backend": "dsa", "page_size": 64},
|
{"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
|
# HIP without the preshuffle path: page 1
|
||||||
with override_platform(is_hip=True):
|
with override_platform(is_hip=True):
|
||||||
with patch(
|
with patch(
|
||||||
|
|||||||
Reference in New Issue
Block a user