diff --git a/python/sglang/kernels/ops/layernorm/__init__.py b/python/sglang/kernels/ops/layernorm/__init__.py index 00f857771..93293d82b 100644 --- a/python/sglang/kernels/ops/layernorm/__init__.py +++ b/python/sglang/kernels/ops/layernorm/__init__.py @@ -514,6 +514,8 @@ _PHASE25_KERNELS = [ ("gemma4_fused_ops", "gemma4_fused_routing", "triton"), ("gemma4_fused_ops", "gemma_qkv_rmsnorm", "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: register_kernel( diff --git a/python/sglang/kernels/ops/layernorm/hy4_ihc.py b/python/sglang/kernels/ops/layernorm/hy4_ihc.py new file mode 100644 index 000000000..b1cea2551 --- /dev/null +++ b/python/sglang/kernels/ops/layernorm/hy4_ihc.py @@ -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 diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 98c6c78be..93246893e 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/python/sglang/kernels/ops/moe/ep_moe_kernels.py @@ -1155,9 +1155,12 @@ def ep_scatter( output_index: torch.Tensor, scale_ue8m0: bool = False, quant_block_size: int = 128, + expert_alignment: int = 128, 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 num_warps = 8 num_experts = num_recv_tokens_per_expert.shape[0] diff --git a/python/sglang/kernels/ops/moe/triton_pad_expert_counts.py b/python/sglang/kernels/ops/moe/triton_pad_expert_counts.py new file mode 100644 index 000000000..b3ff1cfdf --- /dev/null +++ b/python/sglang/kernels/ops/moe/triton_pad_expert_counts.py @@ -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 diff --git a/python/sglang/kernels/ops/moe/triton_sigmoid_gate_mul.py b/python/sglang/kernels/ops/moe/triton_sigmoid_gate_mul.py index 5696a54c7..861ba3ec6 100644 --- a/python/sglang/kernels/ops/moe/triton_sigmoid_gate_mul.py +++ b/python/sglang/kernels/ops/moe/triton_sigmoid_gate_mul.py @@ -37,6 +37,8 @@ def _sigmoid_gate_mul_kernel( def sigmoid_gate_mul(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: """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) n = x.numel() grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),) diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index 6f66b0dec..7d0e109cc 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -44,6 +44,30 @@ from sglang.srt.utils.common import ( 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: """True when ROCm can run the DeepSeek-V4 fp8 wo_a GEMM (gfx950 + aiter).""" try: @@ -159,6 +183,8 @@ def handle_model_specific_adjustments(server_args: Any): "MistralLarge3ForCausalLM", "PixtralForConditionalGeneration", "GlmMoeDsaForCausalLM", + "HYV4ForCausalLM", + "HYV4ForCausalLMNextN", "LongcatFlashForCausalLM", "Dots3NoteForCausalLM", ]: @@ -181,19 +207,7 @@ def handle_model_specific_adjustments(server_args: Any): # The "dsa" attention fill moved to the override registry # (arg_groups/overrides.py: _deepseek_family_overrides). - index_topk_freq = getattr(hf_config, "index_topk_freq", 1) or 1 - 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." - ) + _validate_dsa_tbo_index_sharing(server_args, hf_config) if ( not get_platform().is_npu and not get_platform().is_xpu diff --git a/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py b/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py index 19d48b659..ee192eb88 100644 --- a/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py +++ b/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py @@ -1,6 +1,6 @@ """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 @@ -24,6 +24,8 @@ logger = logging.getLogger(__name__) "MistralLarge3ForCausalLM", "PixtralForConditionalGeneration", "GlmMoeDsaForCausalLM", + "HYV4ForCausalLM", + "HYV4ForCausalLMNextN", "LongcatFlashForCausalLM", "LongcatFlashForCausalLMNextN", "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 the branch.""" 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] = {} + 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 # Set attention backend for DeepSeek if is_attention_backend_not_set(cfg): diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 9ebae0e1c..0415d8f65 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -628,8 +628,16 @@ def _dsa_kv_cache_dtype_default(view: Any) -> dict: ) 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": - 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( 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 ) + 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: backend = "flashinfer_sparse_mla" if not user_set_prefill: @@ -763,6 +792,8 @@ _DEEPSEEK_FAMILY_ARCHS = frozenset( "MistralLarge3ForCausalLM", "PixtralForConditionalGeneration", "GlmMoeDsaForCausalLM", + "HYV4ForCausalLM", + "HYV4ForCausalLMNextN", "LongcatFlashForCausalLM", "LongcatFlashForCausalLMNextN", "Dots3NoteForCausalLM", diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 4dd09822f..9866a4d69 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -862,6 +862,7 @@ def _handle_eagle_family(server_args: ServerArgs) -> None: "MistralLarge3ForCausalLM", "PixtralForConditionalGeneration", "HYV3ForCausalLM", + "HYV4ForCausalLM", ]: if cfg.speculative_draft_model_path is None: declare_resolution( diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index fedbf7708..1c0ccf924 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -17,6 +17,7 @@ from sglang.srt.configs.dots_vlm import DotsVLMConfig from sglang.srt.configs.exaone import ExaoneConfig from sglang.srt.configs.falcon_h1 import FalconH1Config from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig +from sglang.srt.configs.hy_v4 import HYV4Config from sglang.srt.configs.inkling import ( InklingAudioConfig, InklingMMConfig, @@ -118,6 +119,7 @@ __all__ = [ "Dots3Config", "FalconH1Config", "GraniteMoeHybridConfig", + "HYV4Config", "Lfm2Config", "Lfm2MoeConfig", "Lfm2VlConfig", diff --git a/python/sglang/srt/configs/hy_v4.py b/python/sglang/srt/configs/hy_v4.py new file mode 100644 index 000000000..4d9c0c547 --- /dev/null +++ b/python/sglang/srt/configs/hy_v4.py @@ -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"] diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 590d0dd34..abaffcf95 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -66,6 +66,14 @@ def _quant_config_to_dict(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): layout = getattr(hf_config, "attention_projection_layout", None) if layout is None: @@ -134,6 +142,8 @@ def is_deepseek_dsa(config) -> bool: "LongcatFlashForCausalLMNextN", "Dots3NoteForCausalLM", "Dots3NoteForCausalLMNextN", + "HYV4ForCausalLM", + "HYV4ForCausalLMNextN", ) 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: assert is_deepseek_dsa(config) or is_deepseek_v4(config) 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.""" 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. cli_factor = getattr(config, "cli_factor", 1) if cli_factor is None: @@ -794,6 +821,10 @@ class ModelConfig: self.hf_config.architectures[0] = "HYV3ForCausalLMNextN" 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): # Use self.context_len after it has been initialized to prevent using context_len which may be None. self.is_hybrid_swa = ( @@ -865,7 +896,7 @@ class ModelConfig: return getattr( self.hf_text_config, "add_swa_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): is_draft_model = self.is_draft_model @@ -939,6 +970,8 @@ class ModelConfig: or "GlmMoeDsaForCausalLMNextN" in self.hf_config.architectures or "LongcatFlashForCausalLM" 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 "Dots3NoteForCausalLM" in self.hf_config.architectures or "Dots3NoteForCausalLMNextN" in self.hf_config.architectures @@ -969,7 +1002,10 @@ class ModelConfig: else None ) # 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 ( "DeepseekV4ForCausalLM" 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.hidden_size = self.hf_text_config.hidden_size hc_mult = getattr(self.hf_text_config, "hc_mult", 1) - self.spec_hidden_size = ( - self.hidden_size * hc_mult if hc_mult > 1 else self.hidden_size + self.spec_hidden_size, self.hc_hidden_size = resolve_spec_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_attention_layers = self.num_hidden_layers if "LongcatFlashForCausalLM" in self.hf_config.architectures: @@ -1370,11 +1403,8 @@ class ModelConfig: {"quant_algo": "FP8", "quant_method": "modelopt", ...} """ - if "quantization" in quant_config_dict: - json_quant_configs = quant_config_dict["quantization"] - elif "quant_algo" in quant_config_dict: - json_quant_configs = quant_config_dict - else: + json_quant_configs = unwrap_modelopt_quantization_config(quant_config_dict) + if not json_quant_configs: return None quant_algo = json_quant_configs.get("quant_algo", None) diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 43eb943ce..a38f4aa5d 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -89,6 +89,10 @@ from sglang.srt.function_call.utils import ( ) from sglang.srt.managers.io_struct import GenerateReqInput 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 ( MEDIA_URL_PART_TYPES, process_content_for_template_format, @@ -1048,6 +1052,7 @@ class OpenAIServingChat(OpenAIServingBase): request: ChatCompletionRequest, raw_request: Request = None, ) -> tuple[GenerateReqInput, ChatCompletionRequest]: + # 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.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": request.return_output_ids_in_sglext = True - reasoning_effort = ( - request.chat_template_kwargs.pop("reasoning_effort", None) - if request.chat_template_kwargs - else None - ) + reasoning_effort = None + if not uses_hunyuan_reasoning_effort( + self.reasoning_parser, self.template_manager.reasoning_config + ): + 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": raise ValueError( 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: 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 if self.is_gpt_oss or self.is_gemma4: request.skip_special_tokens = False @@ -2575,7 +2589,11 @@ class OpenAIServingChat(OpenAIServingBase): return 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 if self.reasoning_parser == "inkling": @@ -2644,6 +2662,9 @@ class OpenAIServingChat(OpenAIServingBase): ) == "enabled" 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 when reasoning_effort is # "no_think" / "none" / unset; forcing reasoning would route all # output into reasoning_content. diff --git a/python/sglang/srt/function_call/hunyuan_detector.py b/python/sglang/srt/function_call/hunyuan_detector.py index 5e9a794b3..d1aa7e2f8 100644 --- a/python/sglang/srt/function_call/hunyuan_detector.py +++ b/python/sglang/srt/function_call/hunyuan_detector.py @@ -103,6 +103,20 @@ class HunyuanDetector(BaseFormatDetector): def __init__(self, tokenizer=None): super().__init__() + # Hy4 dropped 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 "" for both). + self._has_tool_sep = True + if tokenizer is not None: + try: + self._has_tool_sep = any( + re.fullmatch(r"]+)?>", tok) + for tok in tokenizer.get_vocab() + if isinstance(tok, str) + ) + except Exception: + pass + t = resolve_hunyuan_tokens(tokenizer) tool_calls = t["tool_calls"] tool_call = t["tool_call"] @@ -129,7 +143,13 @@ class HunyuanDetector(BaseFormatDetector): self.tool_call_regex = re.compile( re.escape(tool_call) + r"(.*?)" + + r"(?:" + re.escape(tool_sep) + + r"|(?=" + + re.escape(arg_key) + + r"|" + + re.escape(tc_end) + + r"))" + r"(.*?)" + re.escape(tc_end), re.DOTALL, @@ -383,14 +403,22 @@ class HunyuanDetector(BaseFormatDetector): self._in_tool_calls = False break - sep_pos = self._buffer.find(self.tool_sep_token, tc_start) - if sep_pos == -1: + name_start = tc_start + len(self.tool_call_start_token) + 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:] break + boundary_pos, boundary_token = min(boundaries, key=lambda item: item[0]) - tool_name = self._buffer[ - tc_start + len(self.tool_call_start_token) : sep_pos - ].strip() + tool_name = self._buffer[name_start:boundary_pos].strip() if ( 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. before_name = self._streaming_tool_name @@ -534,6 +564,14 @@ class HunyuanDetector(BaseFormatDetector): return [] def structure_info(self) -> _GetInfoFunc: + # Hy4 has no and no separators between adjacent + # 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( 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}", @@ -542,3 +580,9 @@ class HunyuanDetector(BaseFormatDetector): def supports_structural_tag(self) -> bool: 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 diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 5a5eb5b6b..a7b4e0e31 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -348,6 +348,7 @@ class DeepseekSparseAttnBackend( # Keep original head count if it exceeds current padded variants. self.flashmla_kv_num_q_heads = self.num_q_heads 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) @@ -1885,6 +1886,7 @@ class DeepseekSparseAttnBackend( cos_sin_cache: Optional[torch.Tensor] = None, is_neox: Optional[bool] = False, llama_4_scaling: Optional[torch.Tensor] = None, + attn_sink: Optional[torch.Tensor] = None, ) -> torch.Tensor: causal = not layer.is_cross_attention @@ -1899,6 +1901,10 @@ class DeepseekSparseAttnBackend( ) 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: return self._forward_trtllm( @@ -2124,6 +2130,7 @@ class DeepseekSparseAttnBackend( sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, topk_length=metadata.dsa_cache_seqlens_int32, + attn_sink=attn_sink, ) elif dsa_impl == "flashinfer_sparse_mla": if q_rope is not None: @@ -2195,12 +2202,19 @@ class DeepseekSparseAttnBackend( cos_sin_cache: Optional[torch.Tensor] = None, is_neox: Optional[bool] = False, llama_4_scaling: Optional[torch.Tensor] = None, + attn_sink: Optional[torch.Tensor] = None, ) -> torch.Tensor: causal = not layer.is_cross_attention metadata = self.forward_metadata 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": return self._forward_trtllm( q, @@ -2281,6 +2295,7 @@ class DeepseekSparseAttnBackend( sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, topk_length=metadata.dsa_cache_seqlens_int32, + attn_sink=attn_sink, ) elif self.dsa_decode_impl == "flashinfer_sparse_mla": if q_all is None: @@ -2396,6 +2411,7 @@ class DeepseekSparseAttnBackend( page_table_1: torch.Tensor, sm_scale: float, topk_length: Optional[torch.Tensor] = None, + attn_sink: Optional[torch.Tensor] = None, ) -> torch.Tensor: from sgl_kernel.flash_mla import flash_mla_sparse_fwd @@ -2421,6 +2437,15 @@ class DeepseekSparseAttnBackend( else: 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_input = page_table_1.unsqueeze(1) @@ -2442,6 +2467,7 @@ class DeepseekSparseAttnBackend( indices=indices_input, sm_scale=sm_scale, d_v=v_head_dim, + attn_sink=sink_input, topk_length=topk_length, ) @@ -3437,6 +3463,7 @@ class DeepseekSparseAttnMultiStepBackend: ): self.topk = topk self.speculative_num_steps = speculative_num_steps + self._sink_pad_cache: dict[tuple[int, int], torch.Tensor] = {} self.attn_backends = [] for i in range(self.speculative_num_steps - 1): self.attn_backends.append( diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index 00f756c39..6134374bc 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py @@ -170,6 +170,52 @@ def grouped_gemm_nt_bf16_contig( 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( lhs: Tuple[torch.Tensor, torch.Tensor], rhs: Tuple[torch.Tensor, torch.Tensor], diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 8df1a5c2f..0b7a96a69 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -386,7 +386,9 @@ class LogitsProcessor(nn.Module): self.logit_scale = logit_scale 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_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: self.attn_tp_size = get_parallel().attn_tp_size self.do_tensor_parallel_all_gather = ( diff --git a/python/sglang/srt/layers/moe/cutlass_moe.py b/python/sglang/srt/layers/moe/cutlass_moe.py index 11bcad7dd..6510beda9 100755 --- a/python/sglang/srt/layers/moe/cutlass_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_moe.py @@ -50,6 +50,7 @@ def cutlass_fused_experts_fp8( use_mxfp8: bool = False, output: Optional[torch.Tensor] = None, enable_es: Tuple[bool, bool] = (False, False), + swiglu_limit: Optional[float] = None, ) -> torch.Tensor: """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) - 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: intemediate_q = torch.empty_like(intermediate, dtype=torch.float8_e4m3fn) diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index bc25cbaf2..e43857ae1 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -8,7 +8,11 @@ import torch import triton 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 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 +_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( runner_config: MoeRunnerConfig, quant_info: DeepGemmMoeQuantInfo, @@ -177,6 +207,18 @@ def _should_use_masked_standard_layout( ): 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() if mode not in ("auto", "masked", "compact"): raise ValueError( @@ -287,24 +329,30 @@ class DeepGemmRunnerCore(MoeRunnerCore): hooks: Optional[Any] = None, ) -> DeepGemmRunnerOutput: weight_dtype = quant_info.w13_weight.dtype - 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 - ) + alignment = ( + running_state.get("contiguous_layout_alignment") + if not runner_input.use_masked_gemm + 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: - hidden_states = self._run_contiguous_gemm( - runner_input, quant_info, running_state - ) - else: - if weight_dtype == torch.bfloat16: - hidden_states = self._run_masked_bf16_gemm( - runner_input, quant_info, running_state - ) - else: - hidden_states = self._run_masked_gemm( - runner_input, quant_info, running_state - ) + if weight_dtype == torch.bfloat16: + hidden_states = self._run_masked_bf16_gemm( + 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) def _run_contiguous_gemm( @@ -333,7 +381,7 @@ class DeepGemmRunnerCore(MoeRunnerCore): N = quant_info.w13_weight.size(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 trace_deepep_v2_contig: @@ -344,9 +392,12 @@ class DeepGemmRunnerCore(MoeRunnerCore): (0, K), device=hidden_states_device, dtype=torch.bfloat16 ) - recipe_a, recipe_b = ( - ((1, 128), (1, 32)) if quant_info.is_fp4_experts else (None, None) - ) + if quant_info.use_mxfp8: + 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 = ( quant_info.w13_weight, @@ -469,19 +520,25 @@ class DeepGemmRunnerCore(MoeRunnerCore): 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: down_input = torch.empty( (all_tokens, N // 2), device=gateup_output.device, 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: + 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)) del gateup_output @@ -512,12 +569,17 @@ class DeepGemmRunnerCore(MoeRunnerCore): if deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES: 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( (down_input_fp8, down_input_scale), w2_weight_fp8, down_output, m_indices, - recipe_a=recipe_a, + recipe_a=recipe_a_down, recipe_b=recipe_b, ) 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. # Scatter and post-permute skip non-local experts mapped to -1. - block_e = 128 num_experts = runner_config.num_local_experts 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) 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) valid_tokens_per_expert = tokens_per_expert - tokens_per_expert = (ceil_div(tokens_per_expert, block_e) * block_e).to(torch.int32) - # Keep graph-static shapes by appending padding to the final segment. - # Its m_indices stay -1, so DeepGEMM skips those rows. - tokens_per_expert[-1].add_(all_tokens - tokens_per_expert.sum()) + if _is_cuda: + tokens_per_expert = pad_expert_counts(tokens_per_expert, block_e, all_tokens) + else: + # 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) output_dtype = ( @@ -1023,7 +1092,13 @@ def pre_permute_standard_to_deep_gemm( 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), device=hidden_states_device, dtype=torch.float8_e4m3fn, @@ -1062,6 +1137,7 @@ def pre_permute_standard_to_deep_gemm( src2dst, scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, quant_block_size=(quant_info.block_shape[1] if quant_info.block_shape else 128), + expert_alignment=block_e, expert_start=expert_start, ) 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["src2dst"] = src2dst running_state["all_tokens"] = all_tokens + running_state["contiguous_layout_alignment"] = block_e running_state["mxfp8_act_gran_k"] = ( quant_info.block_shape[1] if quant_info.block_shape else 128 ) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 406ce76ec..b76bd2c3a 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -2000,6 +2000,11 @@ class Fp8MoEMethod(FusedMoEMethodBase): w2_q = layer.w2_weight.data w13_s = layer.w13_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 ( get_moe_runner_backend().is_flashinfer_trtllm() or get_moe_runner_backend().is_flashinfer_trtllm_routed() @@ -2554,6 +2559,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): use_mxfp8=use_mxfp8, output=symm_output, enable_es=(use_mxfp8, use_mxfp8), + swiglu_limit=self.moe_runner_config.swiglu_limit, ) return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index 10ed79457..64bd108e6 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/python/sglang/srt/layers/radix_attention.py @@ -312,6 +312,7 @@ def _unified_attention_with_output_impl( q_rope: Optional[torch.Tensor] = None, k_rope: 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); # they must appear in the schema when --enforce-piecewise-cuda-graph is on. 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] if sinks is not None: kwargs["sinks"] = sinks + if attn_sink is not None: + kwargs["attn_sink"] = attn_sink if cos_sin_cache is not None: kwargs["cos_sin_cache"] = cos_sin_cache if is_neox is not None: @@ -438,6 +441,7 @@ def unified_attention_with_output( q_rope: Optional[torch.Tensor] = None, k_rope: Optional[torch.Tensor] = None, sinks: Optional[torch.Tensor] = None, + attn_sink: Optional[torch.Tensor] = None, cos_sin_cache: Optional[torch.Tensor] = None, is_neox: Optional[bool] = None, llama_4_scaling: Optional[torch.Tensor] = None, @@ -456,6 +460,7 @@ def unified_attention_with_output( q_rope=q_rope, k_rope=k_rope, sinks=sinks, + attn_sink=attn_sink, cos_sin_cache=cos_sin_cache, is_neox=is_neox, llama_4_scaling=llama_4_scaling, @@ -486,6 +491,7 @@ def unified_attention_with_output_and_lse( q_rope: Optional[torch.Tensor] = None, k_rope: Optional[torch.Tensor] = None, sinks: Optional[torch.Tensor] = None, + attn_sink: Optional[torch.Tensor] = None, cos_sin_cache: Optional[torch.Tensor] = None, is_neox: Optional[bool] = None, llama_4_scaling: Optional[torch.Tensor] = None, @@ -504,6 +510,7 @@ def unified_attention_with_output_and_lse( q_rope=q_rope, k_rope=k_rope, sinks=sinks, + attn_sink=attn_sink, cos_sin_cache=cos_sin_cache, is_neox=is_neox, llama_4_scaling=llama_4_scaling, diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 6af4b0c2b..e535ee386 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -118,6 +118,18 @@ def should_defer_dsa_cp_kv_gather( 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: def init_mla_forward(self: DeepseekV2AttentionMLA): self.flashinfer_mla_disable_ragged = ( @@ -161,6 +173,8 @@ class DeepseekMLAForwardMixin: return False if is_kv_b_lora_active(self): 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 # DeepGEMM branches already use different fused paths. 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 # block; also suppresses the duplicate split/rope on that path. 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 = ( self.q_lora_rank is not None @@ -667,6 +686,13 @@ class DeepseekMLAForwardMixin: topk_indices, llama_4_scaling, 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( @@ -681,18 +707,20 @@ class DeepseekMLAForwardMixin: topk_indices, llama_4_scaling, fusion_plan: Optional[MlaBmmFusionPlan] = None, - gate: Optional[torch.Tensor] = None, + attention_output_gate: Optional[torch.Tensor] = None, ): save_kv_cache = True if self.current_attention_backend in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS: 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): - extra_args = { - "cos_sin_cache": self.rotary_emb.cos_sin_cache, - "is_neox": self.rotary_emb.is_neox_style, - "llama_4_scaling": llama_4_scaling, - } + extra_args.update( + cos_sin_cache=self.rotary_emb.cos_sin_cache, + is_neox=self.rotary_emb.is_neox_style, + llama_4_scaling=llama_4_scaling, + ) if fusion_plan is not None: bmm_attention_fn = ( bcg_mla_bmm_then_unified_attention @@ -911,8 +939,10 @@ class DeepseekMLAForwardMixin: attn_bmm_output = apply_kv_b_lora_v_correction( self, attn_output, attn_bmm_output ) - if gate is not None: - attn_bmm_output = self._apply_gated(attn_bmm_output, gate) + if attention_output_gate is not None: + attn_bmm_output = _apply_attention_output_gate( + self, attn_bmm_output, attention_output_gate + ) output, _ = self.o_proj(attn_bmm_output) if self.next_skip_topk is None: diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index 4d17d45e9..cd916cf60 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -147,6 +147,44 @@ def _load_fused_indexer_wk( 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) class NextNEnabledConfig: num_nextn_layers: int @@ -281,6 +319,9 @@ class DeepseekV2WeightLoaderMixin: ): continue + if _load_fused_expert_tensor(name, loaded_weight, params_dict): + continue + if "rotary_emb.inv_freq" in name: continue @@ -594,7 +635,11 @@ class DeepseekV2WeightLoaderMixin: 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. - 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( weight_block_size=( self.quant_config.weight_block_size diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 3470b65c8..9aafb6edd 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -478,7 +478,14 @@ class MoEGate(nn.Module): self.is_nextn = is_nextn self.is_deepseek_v4 = is_deepseek_v4 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: @@ -515,6 +522,9 @@ class MoEGate(nn.Module): gemm_output_zero_allocator: BumpAllocator = 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): return torch.ops.sgl_kernel.weight_packed_linear( hidden_states, diff --git a/python/sglang/srt/models/hunyuan_v4.py b/python/sglang/srt/models/hunyuan_v4.py new file mode 100644 index 000000000..c7971dd89 --- /dev/null +++ b/python/sglang/srt/models/hunyuan_v4.py @@ -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] diff --git a/python/sglang/srt/models/hunyuan_v4_nextn.py b/python/sglang/srt/models/hunyuan_v4_nextn.py new file mode 100644 index 000000000..fce1b9229 --- /dev/null +++ b/python/sglang/srt/models/hunyuan_v4_nextn.py @@ -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] diff --git a/python/sglang/srt/parser/hunyuan_reasoning.py b/python/sglang/srt/parser/hunyuan_reasoning.py new file mode 100644 index 000000000..c6a44d2a4 --- /dev/null +++ b/python/sglang/srt/parser/hunyuan_reasoning.py @@ -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) diff --git a/python/sglang/srt/parser/template_detection.py b/python/sglang/srt/parser/template_detection.py index ad6b3a405..6a82786de 100644 --- a/python/sglang/srt/parser/template_detection.py +++ b/python/sglang/srt/parser/template_detection.py @@ -175,6 +175,15 @@ REASONING_MODE_RULES = ( 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( name="explicit_enable_thinking_default_false", value=ReasoningToggleConfig( @@ -364,8 +373,16 @@ def _is_hunyuan(ctx): sep = ctx.has_text("") or ctx.has_vocab_pattern( r"^]+)?>$" ) - return (tc and sep) or ( - ctx.has_text("reasoning_effort") and ctx.has_text("interleaved_thinking") + return ( + (tc and sep) + or ( + tc + and ctx.reasoning_config + == ReasoningToggleConfig(special_case="hunyuan_effort") + and ctx.has_vocab_pattern(r"^]+)?>$") + and ctx.has_vocab_pattern(r"^]+)?>$") + ) + or (ctx.has_text("reasoning_effort") and ctx.has_text("interleaved_thinking")) ) diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index 26d141ce2..552c6a799 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -38,6 +38,7 @@ from sglang.srt.configs import ( ExaoneConfig, FalconH1Config, GraniteMoeHybridConfig, + HYV4Config, InklingAudioConfig, InklingMMConfig, InklingModelConfig, @@ -122,6 +123,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = { Qwen3NextConfig, FalconH1Config, GraniteMoeHybridConfig, + HYV4Config, DotsVLMConfig, DotsOCRConfig, Dots3Config, diff --git a/test/registered/kernels/ops/attention/test_dsa_sink_pad_cache.py b/test/registered/kernels/ops/attention/test_dsa_sink_pad_cache.py new file mode 100644 index 000000000..2af86eb97 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_dsa_sink_pad_cache.py @@ -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__])) diff --git a/test/registered/kernels/ops/attention/test_hy4_hpc_gated_mla.py b/test/registered/kernels/ops/attention/test_hy4_hpc_gated_mla.py new file mode 100644 index 000000000..15d421761 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_hy4_hpc_gated_mla.py @@ -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__])) diff --git a/test/registered/kernels/ops/layernorm/test_hy4_decode_kernels.py b/test/registered/kernels/ops/layernorm/test_hy4_decode_kernels.py new file mode 100644 index 000000000..dfc0f72d8 --- /dev/null +++ b/test/registered/kernels/ops/layernorm/test_hy4_decode_kernels.py @@ -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() diff --git a/test/registered/kernels/ops/layernorm/test_hy4_hpc_ihc.py b/test/registered/kernels/ops/layernorm/test_hy4_hpc_ihc.py new file mode 100644 index 000000000..203455d95 --- /dev/null +++ b/test/registered/kernels/ops/layernorm/test_hy4_hpc_ihc.py @@ -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__])) diff --git a/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py b/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py index 6a7c709c1..821a1524c 100644 --- a/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py +++ b/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py @@ -243,6 +243,18 @@ def test_standard_masked_runner_matches_compact_end_to_end(monkeypatch, weight_d "use_symmetric_memory", lambda *args, **kwargs: nullcontext(), ) + monkeypatch.setattr( + deep_gemm_runner.deep_gemm_wrapper, + "get_contiguous_layout_alignment", + lambda expected_m, num_groups: 32, + ) + monkeypatch.setattr( + deep_gemm_runner, + "get_exec", + lambda: SimpleNamespace( + deterministic=SimpleNamespace(enable_deterministic_inference=False) + ), + ) # UE8M0 packs four 128-wide scale groups into each int32. Use the smallest # legal K for both the gate/up and down GEMMs. @@ -357,17 +369,23 @@ def test_standard_masked_runner_matches_compact_end_to_end(monkeypatch, weight_d ).hidden_states, ) - compact_is_masked, compact_all_tokens, compact_m_indices, compact_output = ( - run_with_layout("compact") - ) - masked_is_masked, masked_all_tokens, masked_m_indices, masked_output = ( - run_with_layout("masked") - ) + ( + compact_is_masked, + compact_all_tokens, + compact_m_indices, + compact_output, + ) = run_with_layout("compact") + ( + masked_is_masked, + masked_all_tokens, + masked_m_indices, + masked_output, + ) = run_with_layout("masked") torch.cuda.synchronize() assert not compact_is_masked assert masked_is_masked - assert compact_all_tokens == 256 + assert compact_all_tokens == 64 assert masked_all_tokens is None assert masked_m_indices is None valid_assignments = topk_ids[topk_ids >= 0] diff --git a/test/registered/kernels/ops/moe/test_pad_expert_counts.py b/test/registered/kernels/ops/moe/test_pad_expert_counts.py new file mode 100644 index 000000000..3f901f36d --- /dev/null +++ b/test/registered/kernels/ops/moe/test_pad_expert_counts.py @@ -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() diff --git a/test/registered/quant/test_fp8_utils.py b/test/registered/quant/test_fp8_utils.py index c256c5d1d..97d149f0c 100644 --- a/test/registered/quant/test_fp8_utils.py +++ b/test/registered/quant/test_fp8_utils.py @@ -4,17 +4,76 @@ from unittest.mock import patch import torch +from sglang.srt.layers.quantization.fp8 import ( + Fp8MoEMethod, + _is_cuda, + _is_gfx95_supported, + _is_hip, +) from sglang.srt.layers.quantization.fp8_utils import ( inverse_transform_scale_ue8m0, quant_weight_ue8m0, transform_scale_ue8m0, ) +from sglang.srt.runtime_context import get_platform from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=12, stage="base-b", runner_config="1-gpu-large") +class TestMxfp8MoeScaleLayout(CustomTestCase): + @classmethod + def setUpClass(cls): + if not ( + (_is_cuda and get_platform().is_sm100) or (_is_hip and _is_gfx95_supported) + ): + raise unittest.SkipTest( + "MXFP8 MoE quantization requires SM100 or ROCm gfx95" + ) + + def test_cutlass_serialized_scales_remain_expert_first(self): + class CutlassBackend: + def is_cutlass(self): + return True + + def is_flashinfer_trtllm(self): + return False + + def is_flashinfer_trtllm_routed(self): + return False + + def is_deep_gemm(self): + return False + + layer = SimpleNamespace( + w13_weight=torch.nn.Parameter( + torch.zeros((2, 64, 32), dtype=torch.float8_e4m3fn, device="cuda") + ), + w2_weight=torch.nn.Parameter( + torch.zeros((2, 32, 32), dtype=torch.float8_e4m3fn, device="cuda") + ), + w13_weight_scale_inv=torch.nn.Parameter( + torch.zeros((2, 64, 1), dtype=torch.uint8, device="cuda"), + requires_grad=False, + ), + w2_weight_scale_inv=torch.nn.Parameter( + torch.zeros((2, 32, 1), dtype=torch.uint8, device="cuda"), + requires_grad=False, + ), + ) + method = object.__new__(Fp8MoEMethod) + + with patch( + "sglang.srt.layers.quantization.fp8.get_moe_runner_backend", + return_value=CutlassBackend(), + ): + method._process_mxfp8_moe_weights(layer, quantize=False) + + self.assertEqual(tuple(layer.w13_weight_scale_inv.shape), (2, 64, 1)) + self.assertEqual(tuple(layer.w2_weight_scale_inv.shape), (2, 32, 1)) + + class TestInverseTransformScaleUe8m0(CustomTestCase): def test_round_trip(self): for _ in range(100): diff --git a/test/registered/rl/test_fp32_lm_head.py b/test/registered/rl/test_fp32_lm_head.py index 685b0b969..82460fef3 100644 --- a/test/registered/rl/test_fp32_lm_head.py +++ b/test/registered/rl/test_fp32_lm_head.py @@ -10,6 +10,7 @@ from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.runtime_context import get_context from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small") register_amd_ci(est_time=15, suite="stage-b-test-1-gpu-small-amd") @@ -30,7 +31,7 @@ class DummyMeta: def compute_dp_attention_metadata(self): ... -class TestLMHeadFP32(unittest.TestCase): +class TestLMHeadFP32(CustomTestCase): @classmethod def setUpClass(cls): if not torch.cuda.is_available() and not ( @@ -38,13 +39,17 @@ class TestLMHeadFP32(unittest.TestCase): ): raise unittest.SkipTest("needs CUDA GPU or XPU") - def _make_logprocessor(self, vocab_size, enable_fp32): + def _make_logprocessor(self, vocab_size, enable_fp32, config_enable_fp32=False): # LogitsProcessor reads get_exec().features.enable_fp32_lm_head # from the published config. override = get_context().override_server_args(enable_fp32_lm_head=enable_fp32) override.install() self.addCleanup(override.restore) - cfg = SimpleNamespace(vocab_size=vocab_size, final_logit_softcapping=None) + cfg = SimpleNamespace( + vocab_size=vocab_size, + final_logit_softcapping=None, + enable_lm_head_fp32=config_enable_fp32, + ) return LogitsProcessor(cfg, skip_all_gather=True, logit_scale=None) def _run_case( @@ -55,6 +60,7 @@ class TestLMHeadFP32(unittest.TestCase): expected_a_dtype, expected_b_dtype, expected_operation, + config_enable_fp32=False, ): device = get_device() BATCH_SIZE, HIDDEN_SIZE, VOCAB_SIZE = 2, 64, 128 @@ -63,7 +69,9 @@ class TestLMHeadFP32(unittest.TestCase): ) head = LMHeadStub(VOCAB_SIZE, HIDDEN_SIZE, dtype=weights_dtype, device=device) meta = DummyMeta() - logprocessor = self._make_logprocessor(VOCAB_SIZE, enable_fp32) + logprocessor = self._make_logprocessor( + VOCAB_SIZE, enable_fp32, config_enable_fp32 + ) original_matmul = torch.matmul original_mm = torch.mm @@ -173,6 +181,21 @@ class TestLMHeadFP32(unittest.TestCase): "matmul", ) + def test_model_config_enables_fp32_without_server_flag(self): + expected_operation = "mm" if torch.cuda.is_available() else "matmul" + expected_dtype = ( + torch.float32 if expected_operation == "matmul" else torch.bfloat16 + ) + self._run_case( + torch.bfloat16, + False, + torch.bfloat16, + expected_dtype, + expected_dtype, + expected_operation, + config_enable_fp32=True, + ) + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/test/registered/unit/configs/test_model_config_scaling.py b/test/registered/unit/configs/test_model_config_scaling.py index f458b47fd..f2b0a6f76 100644 --- a/test/registered/unit/configs/test_model_config_scaling.py +++ b/test/registered/unit/configs/test_model_config_scaling.py @@ -1,5 +1,6 @@ import math import unittest +from types import SimpleNamespace from sglang.srt.configs.model_config import ModelConfig, compute_mla_mscale_scaling from sglang.test.ci.ci_register import register_cpu_ci @@ -92,6 +93,31 @@ class TestInitMlaScaling(CustomTestCase): base, scaling = _mla_scaling(None) self.assertEqual(scaling, base) + def test_hyv4_shape_derivation_uses_rope_parameters(self): + hf_config = SimpleNamespace( + architectures=["HYV4ForCausalLM"], + model_type="hy_v4", + hidden_size=2816, + num_hidden_layers=34, + num_attention_heads=32, + vocab_size=120832, + head_dim=64, + v_head_dim=256, + kv_lora_rank=512, + qk_nope_head_dim=192, + qk_rope_head_dim=64, + index_topk=2048, + index_head_dim=128, + rope_parameters={"rope_theta": 10_000_000, "rope_type": "default"}, + ) + config = ModelConfig.__new__(ModelConfig) + config.hf_config = hf_config + config.hf_text_config = hf_config + + config._derive_model_shapes() + + self.assertEqual(config.scaling, 1 / math.sqrt(256)) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 204c34ad2..f1123ae57 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -800,6 +800,87 @@ class ServingChatTestCase(unittest.TestCase): self.assertEqual(req.reasoning_effort, "high") + def test_hunyuan_default_reasoning_effort_is_normalized_for_template(self): + self.template_manager.chat_template_name = None + self.template_manager.jinja_template_content_format = "string" + self.template_manager.reasoning_config = ReasoningToggleConfig( + special_case="hunyuan_effort" + ) + self.chat.reasoning_parser = "hunyuan" + self.tm.tokenizer.apply_chat_template.return_value = [1, 2, 3] + + cases = [ + ("none", None, "no_think"), + ("none", "xhigh", "high"), + ] + for default_effort, request_effort, normalized_effort in cases: + with self.subTest( + default_effort=default_effort, request_effort=request_effort + ): + self.chat.default_chat_template_kwargs = { + "reasoning_effort": default_effort + } + req = ChatCompletionRequest( + model="x", + messages=[{"role": "user", "content": "What is 2+2?"}], + reasoning_effort=request_effort, + ) + + self.chat._process_messages(req, is_multimodal=False) + + kwargs = self.tm.tokenizer.apply_chat_template.call_args.kwargs + self.assertEqual(req.reasoning_effort, normalized_effort) + self.assertEqual(kwargs["reasoning_effort"], normalized_effort) + self.assertNotIn("reasoning_effort", req.chat_template_kwargs) + + def test_hunyuan_reasoning_effort_precedence_survives_conversion(self): + self.template_manager.chat_template_name = None + self.template_manager.jinja_template_content_format = "string" + self.template_manager.reasoning_config = ReasoningToggleConfig( + special_case="hunyuan_effort" + ) + self.chat.reasoning_parser = "hunyuan" + self.tm.tokenizer.apply_chat_template.return_value = [1, 2, 3] + + cases = [ + ("xhigh", "none", "high"), + (None, "no_think", "no_think"), + ] + for request_effort, template_effort, normalized_effort in cases: + with self.subTest( + request_effort=request_effort, template_effort=template_effort + ): + req = ChatCompletionRequest( + model="x", + messages=[{"role": "user", "content": "What is 2+2?"}], + reasoning_effort=request_effort, + chat_template_kwargs={"reasoning_effort": template_effort}, + ) + + self.chat._convert_to_internal_request(req) + + kwargs = self.tm.tokenizer.apply_chat_template.call_args.kwargs + self.assertEqual(req.reasoning_effort, normalized_effort) + self.assertEqual(kwargs["reasoning_effort"], normalized_effort) + self.assertNotIn("reasoning_effort", req.chat_template_kwargs) + + def test_non_hunyuan_default_reasoning_effort_is_unchanged(self): + self.template_manager.chat_template_name = None + self.template_manager.jinja_template_content_format = "string" + self.tm.tokenizer.apply_chat_template.return_value = [1, 2, 3] + self.chat.default_chat_template_kwargs = {"reasoning_effort": "medium"} + req = ChatCompletionRequest( + model="x", + messages=[{"role": "user", "content": "What is 2+2?"}], + ) + + self.chat._process_messages(req, is_multimodal=False) + + kwargs = self.tm.tokenizer.apply_chat_template.call_args.kwargs + self.assertEqual(req.reasoning_effort, "medium") + self.assertEqual(kwargs["reasoning_effort"], "medium") + self.assertEqual(req.chat_template_kwargs["reasoning_effort"], "medium") + def test_k2_selected_terminator_reaches_sampling_params(self): self.tm._config_overrides["reasoning_parser"] = "k2_horizon" self.chat = OpenAIServingChat(self.tm, self.template_manager) diff --git a/test/registered/unit/function_call/test_hunyuan_detector.py b/test/registered/unit/function_call/test_hunyuan_detector.py index b1a66b177..80060a5e2 100644 --- a/test/registered/unit/function_call/test_hunyuan_detector.py +++ b/test/registered/unit/function_call/test_hunyuan_detector.py @@ -75,6 +75,17 @@ def _make_tools(): ] +class _Hy4Tokenizer: + def get_vocab(self): + return { + "": 1, + "": 2, + "": 3, + "": 4, + "": 5, + } + + class TestHunyuanDetectorHasToolCall(CustomTestCase): def setUp(self): self.detector = HunyuanDetector() @@ -111,6 +122,21 @@ class TestHunyuanDetectorDetectAndParse(CustomTestCase): self.assertEqual(len(result.calls), 0) self.assertEqual(result.normal_text, text) + def test_hy4_format_without_tool_separator(self): + detector = HunyuanDetector(_Hy4Tokenizer()) + text = ( + "" + "get_weather" + "city" + "Beijing" + "" + ) + + result = detector.detect_and_parse(text, self.tools) + + self.assertEqual(result.calls[0].name, "get_weather") + self.assertEqual(json.loads(result.calls[0].parameters), {"city": "Beijing"}) + def test_zero_arg_inline(self): text = ( "get_current_date" @@ -293,6 +319,43 @@ class TestHunyuanDetectorArgDeserialization(CustomTestCase): args = json.loads(result.calls[0].parameters) self.assertIs(args["verbose"], True) + def test_top_level_composed_schema_args(self): + cases = ( + ("anyOf", {"type": "integer"}, "7", 7), + ("oneOf", {"type": "boolean"}, "true", True), + ("allOf", {"type": "array", "items": {"type": "integer"}}, "[1,2]", [1, 2]), + ) + for keyword, arg_schema, raw_value, expected in cases: + with self.subTest(keyword=keyword): + function_name = f"composed_{keyword}" + tools = [ + Tool( + type="function", + function=Function( + name=function_name, + description="Composed schema", + parameters={ + keyword: [ + { + "type": "object", + "properties": {"value": arg_schema}, + } + ] + }, + ), + ) + ] + text = ( + f"{function_name}" + f"value{raw_value}" + "" + ) + + result = self.detector.detect_and_parse(text, tools) + args = json.loads(result.calls[0].parameters) + + self.assertEqual(args, {"value": expected}) + def test_string_arg_not_deserialized(self): """String-typed args should stay as strings even if they look like JSON.""" text = ( @@ -359,6 +422,23 @@ class TestHunyuanDetectorStreaming(CustomTestCase): self.assertEqual(collected[0]["name"], "get_current_date") self.assertEqual(json.loads(collected[0]["parameters"]), {}) + def test_hy4_format_without_tool_separator_char_by_char(self): + detector = HunyuanDetector(_Hy4Tokenizer()) + text = ( + "" + "get_weather" + "city" + "Tokyo" + "" + ) + all_calls = [] + for char in text: + all_calls.extend(detector.parse_streaming_increment(char, self.tools).calls) + + collected = _collect_streamed_tool_calls(all_calls) + self.assertEqual(collected[0]["name"], "get_weather") + self.assertEqual(json.loads(collected[0]["parameters"]), {"city": "Tokyo"}) + def test_chunked_tool_call(self): detector = self._new_detector() chunks = [ diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index 48a9eddbf..fa607a540 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -552,6 +552,21 @@ class TestParseQuantHfConfig(CustomTestCase): self.assertIn("lm_head", quant_config.ignored_layers) self.assertEqual(quant_config.kv_cache_quant_algo, "FP8") + nested_result = model_config._parse_modelopt_quant_config( + { + "quantization": { + "quantization": { + "quant_algo": "MXFP8", + "group_size": 32, + "exclude_modules": ["lm_head"], + } + } + } + ) + self.assertEqual(nested_result["quant_method"], "mxfp8") + self.assertEqual(nested_result["scale_fmt"], "ue8m0") + self.assertIn("lm_head", nested_result["modules_to_not_convert"]) + def test_modelopt_mxfp8_override(self): """Generic ModelOpt selection must not route MXFP8 to scalar FP8.""" self.assertEqual( diff --git a/test/registered/unit/models/test_hunyuan_v4.py b/test/registered/unit/models/test_hunyuan_v4.py new file mode 100644 index 000000000..ce98b9382 --- /dev/null +++ b/test/registered/unit/models/test_hunyuan_v4.py @@ -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__])) diff --git a/test/registered/unit/models/test_hunyuan_v4_nextn_weight_loading.py b/test/registered/unit/models/test_hunyuan_v4_nextn_weight_loading.py new file mode 100644 index 000000000..f93710d09 --- /dev/null +++ b/test/registered/unit/models/test_hunyuan_v4_nextn_weight_loading.py @@ -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() diff --git a/test/registered/unit/parser/test_hunyuan_reasoning.py b/test/registered/unit/parser/test_hunyuan_reasoning.py new file mode 100644 index 000000000..5013f28f3 --- /dev/null +++ b/test/registered/unit/parser/test_hunyuan_reasoning.py @@ -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__])) diff --git a/test/registered/unit/parser/test_template_manager.py b/test/registered/unit/parser/test_template_manager.py index f1b887e2d..52edc6b83 100644 --- a/test/registered/unit/parser/test_template_manager.py +++ b/test/registered/unit/parser/test_template_manager.py @@ -269,6 +269,42 @@ class TestTemplateManagerReasoningDetection(unittest.TestCase): _, _, parser = self._detect(template, [""]) 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" + "{{ name }}{{ k }}" + ) + + HYV4_VOCAB = [ + "", + "", + "", + "", + ] + + def test_hyv4_effort_template_detected_with_special_case(self): + # Hy4 drops ; 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, [""]) + + self.assertEqual(config, ReasoningToggleConfig(special_case="hunyuan_effort")) + self.assertNotEqual(parser, "hunyuan") + class TestTemplateDetectionRuleMatrix(unittest.TestCase): """Table-driven tests for REASONING_PARSER_RULES and REASONING_MODE_RULES.""" diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index d5f26623e..5aaedbbb1 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1037,6 +1037,7 @@ class TestContextParallelServerArgs(CustomTestCase): moe_dp_size=1, ep_size=1, pp_size=1, + dcp_size=1, enable_aiter_allreduce_fusion=False, ) defaults.update(overrides) @@ -1169,6 +1170,29 @@ class TestContextParallelServerArgs(CustomTestCase): resolution_result(args, "dsa_prefill_cp_mode"), "round-robin-split" ) + def test_canonical_interleave_cp_mirrors_to_dsa_runtime_aliases(self): + server_args = self._new_cp_args( + enable_prefill_cp=True, + cp_strategy="interleave", + attention_backend="dsa", + ) + + handle_legacy_cp_runtime_compatibility(server_args) + handle_context_parallelism(server_args) + + self.assertTrue( + resolution_result(server_args, "enable_dsa_prefill_context_parallel") + ) + self.assertFalse( + resolution_result(server_args, "enable_prefill_context_parallel") + ) + self.assertEqual( + resolution_result(server_args, "dsa_prefill_cp_mode"), "round-robin-split" + ) + self.assertEqual( + resolution_result(server_args, "prefill_cp_mode"), "round-robin-split" + ) + def test_context_parallel_handler_initializes_cp_strategy(self): server_args = self._new_cp_args( enable_prefill_cp=True, diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 3285905f8..dc82825f1 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -1817,8 +1817,8 @@ class TestGoldenModelOverrides(_IsolatedPublish): _dsa_split_backend_resolution, ) - def _view(arch="DeepseekV32ForCausalLM", **kw): - hf = SimpleNamespace(architectures=[arch]) + def _view(arch="DeepseekV32ForCausalLM", learnable_sink=False, **kw): + hf = SimpleNamespace(architectures=[arch], learnable_sink=learnable_sink) defaults = dict( kv_cache_dtype="fp8_e4m3", dsa_prefill_backend=None, @@ -1866,6 +1866,28 @@ class TestGoldenModelOverrides(_IsolatedPublish): "dsa_decode_backend": "flashmla_kv", }, ) + for arch in ("HYV4ForCausalLM", "HYV4ForCausalLMNextN"): + with self.subTest(arch=arch, backends="default"): + self.assertEqual( + _dsa_split_backend_resolution( + _view(arch=arch, learnable_sink=True) + ), + { + "dsa_prefill_backend": "flashmla_sparse", + "dsa_decode_backend": "flashmla_sparse", + }, + ) + for field, value in ( + ("dsa_prefill_backend", "fa3"), + ("dsa_decode_backend", "trtllm"), + ): + with self.subTest(arch=arch, field=field, value=value): + with self.assertRaisesRegex( + ValueError, field.replace("_", "-") + ): + _dsa_split_backend_resolution( + _view(arch=arch, learnable_sink=True, **{field: value}) + ) # non-family arch declares nothing self.assertEqual( _dsa_split_backend_resolution(_view(arch="LlamaForCausalLM")), {} @@ -2842,6 +2864,7 @@ class TestGoldenModelOverrides(_IsolatedPublish): prefill_attention_backend=None, decode_attention_backend=None, enable_prefill_cp=False, + dcp_size=1, ) defaults.update(kw) return SimpleNamespace(**defaults) @@ -2857,6 +2880,22 @@ class TestGoldenModelOverrides(_IsolatedPublish): _deepseek_family_overrides(_args(), None), {"attention_backend": "dsa", "page_size": 64}, ) + for arch in ("HYV4ForCausalLM", "HYV4ForCausalLMNextN"): + hf_config = SimpleNamespace(architectures=[arch]) + with self.subTest(arch=arch, prefill_cp=True): + with self.assertRaisesRegex( + ValueError, "--enable-prefill-cp.*HYV4" + ): + _deepseek_family_overrides( + _args(enable_prefill_cp=True), hf_config + ) + with self.subTest(arch=arch, dcp_size=2): + with self.assertRaisesRegex( + ValueError, "--dcp-size > 1.*HYV4" + ): + _deepseek_family_overrides( + _args(dcp_size=2), hf_config + ) # HIP without the preshuffle path: page 1 with override_platform(is_hip=True): with patch(