diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index f8d03952478c..dc1ef3c4712b 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/python/sglang/srt/models/deepseek_v4_nextn.py @@ -2,7 +2,6 @@ from typing import Iterable, Optional, Tuple import torch -import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig @@ -37,12 +36,19 @@ ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_context import get_attn_backend -from sglang.srt.models.deepseek_v4 import DeepseekV4DecoderLayer, DeepseekV4ForCausalLM +from sglang.srt.model_executor.runner import compile_in_capture_mode +from sglang.srt.models.deepseek_v4 import ( + DeepseekV4DecoderLayer, + DeepseekV4ForCausalLM, + hc_head_torch, +) from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import add_prefix +from sglang.srt.utils import add_prefix, is_cuda logger = logging.getLogger(__name__) +_is_cuda = is_cuda() + COMPRESS_RATIO_NEXTN_LAYER = 0 @@ -122,13 +128,16 @@ def hc_head( hc_scale: torch.Tensor, hc_base: torch.Tensor, ): - shape, dtype = x.size(), x.dtype - x = x.flatten(1).float() - rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.rms_norm_eps) - mixes = F.linear(x, hc_fn) * rsqrt - pre = torch.sigmoid(mixes * hc_scale + hc_base) + self.hc_eps - y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=1) - return y.to(dtype) + kwargs = { + "norm_eps": self.rms_norm_eps, + "hc_eps": self.hc_eps, + } + + if x.shape[0] == 0 or not _is_cuda: + return hc_head_torch(x, hc_fn, hc_scale, hc_base, **kwargs) + + hc_head_impl = compile_in_capture_mode(hc_head_torch) + return hc_head_impl(x, hc_fn, hc_scale, hc_base, **kwargs) def forward( self,