minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity
Files
92632a60bad2148eff27ebe22e9d92d5a81c1a40
sglang/python/sglang/srt/models/deepseek_v4.py
T
SYChen123 3810f531a8 Support V4.1 decode vision MegaMoE 
(cherry picked from commit 98aa79972d6d31f30a3c5bc83b9d7e3009ae638a)
2026-09-20 22:12:43 +08:00

5952 lines
238 KiB
Python
Raw Blame History

from __future__ import annotations
import concurrent.futures
import functools
import logging
import time
from contextlib import contextmanager, nullcontext
from types import SimpleNamespace
from typing import (
TYPE_CHECKING,
Any,
Callable,
Iterable,
List,
NamedTuple,
Optional,
Set,
Tuple,
Union,
)
import torch
import torch.nn as nn
import torch.nn.functional as F
import sglang.srt.models.deepseek_v2 as deepseek_v2
from sglang.kernels.ops.attention.dsv4 import (
fused_norm_rope_inplace,
fused_q_norm_rope,
fused_rope_inplace,
sglang_per_token_group_quant_fp8_dsv4_wo_a,
)
from sglang.kernels.ops.attention.dsv4.wo_a import MAX_M as _FUSED_WO_A_MAX_TOKENS
from sglang.kernels.ops.attention.dsv4.wo_a import (
fused_rope_wo_a_bf16,
wo_a_bf16_gemv,
wo_a_bf16_small_batch,
wo_a_bf16_small_batch_mxfp8,
)
from sglang.kernels.ops.attention.flash_mla_sm120 import SM120_DECODE_MAX_TOKENS
from sglang.kernels.ops.layernorm.mhc_post_split_h import mhc_post_split_h
from sglang.kernels.ops.quantization.fp8_kernel import (
sglang_per_token_group_quant_fp8,
)
from sglang.srt.compilation.compilation_config import register_split_op
from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import (
Dsv4NpuRoPE,
prime_rope_cos_sin,
rope_cos_sin,
)
from sglang.srt.hardware_backend.npu.utils import (
is_npu_arch35,
use_npu_arch35_mxfp8_wo_a,
)
from sglang.srt.layers.attention.dsa.utils import (
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
)
from sglang.srt.layers.attention.dsv4.compressor import Compressor
from sglang.srt.layers.attention.dsv4.dsv41_sparse import (
DeepseekV41Compressor,
DeepseekV41Indexer,
)
from sglang.srt.layers.attention.dsv4.indexer import C4Indexer
from sglang.srt.layers.communicator import get_attn_tp_context
from sglang.srt.layers.communicator_dsa_cp import (
dsa_cp_gather_hidden_states,
dsa_cp_reduce_scatter_hidden_states,
)
from sglang.srt.layers.cp.cp_decode_attn_tp import get_cp_decode_attn_tp_ctx
from sglang.srt.layers.cp.utils import (
cp_gather_full_sequence_states,
cp_materialize_global_token_order,
is_cp_active,
)
from sglang.srt.layers.deep_gemm_wrapper.configurer import DEEPGEMM_SCALE_UE8M0
from sglang.srt.layers.dp_attention import (
_tbo_event,
attn_tp_all_gather,
attn_tp_all_reduce,
dp_gather_partial,
dp_gather_replicate,
dp_reduce_scatter_tensor,
dp_reduce_scatterv_async,
dp_scatter,
get_dp_global_num_tokens,
get_dp_tbo_comm_stream,
get_global_dp_buffer,
get_global_dp_buffer_len,
get_local_dp_buffer,
get_local_dp_buffer_len,
get_tbo_persistent_buffer,
is_allocation_symmetric,
is_dp_attention_enabled,
is_dp_gatherv_active,
)
from sglang.srt.layers.engram import Engram, EngramHasher, EngramLayout
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor
from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
from sglang.srt.layers.moe.utils import (
is_shared_experts_fusion_disabled,
uses_per_rank_fused_shared_slots,
)
from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod
from sglang.srt.layers.quantization.fp8_utils import (
Mxfp8DenseGemmBackend,
view_aiter_fused_rms_transposed_fp8_scale,
)
from sglang.srt.layers.quantization.mxfp8_input import Mxfp8SwizzledInput
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
from sglang.srt.managers.mm_utils import (
MultiModalityDataPaddingPatternMultimodalTokens,
embed_mm_inputs,
)
from sglang.srt.managers.schedule_batch import MM_PAD_SHIFT_VALUE, MultimodalInputs
from sglang.srt.mem_cache.memory_pool import RadixAttention
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
PPProxyTensors,
)
from sglang.srt.model_executor.forward_context import (
get_attn_backend,
get_token_to_kv_pool,
)
from sglang.srt.model_executor.runner import (
compile_in_capture_mode,
get_is_capture_mode,
)
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
eager_on_graph,
)
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
is_in_breakable_cuda_graph,
)
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
get_tc_piecewise_forward_context,
)
from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load
from sglang.srt.model_loader.weight_utils import (
RUNAI_STREAMER_TENSOR_ATTR,
default_weight_loader,
)
from sglang.srt.models.dbrx import ReplicatedLinear
from sglang.srt.models.deepseek_common.amd.deepseek_v4_fused_mhc import (
apply_mhc_post_pre_boundary,
is_cross_layer_mhc_fusion_enabled,
)
from sglang.srt.models.deepseek_common.utils import (
_use_aiter_bpreshuffle_gfx95,
is_wint4afp8_or_wint4a16_config,
quant_blocks_shared_experts_fusion,
)
from sglang.srt.models.deepseek_v2 import (
ParallelLMHead,
_is_cuda,
_is_hip,
_is_npu,
_is_xpu,
)
from sglang.srt.models.deepseek_v41_vit import Aligner, ViT
from sglang.srt.multimodal.deepseek_v41_image_processing import (
GPU_PLAN_KEY,
image_token_types,
materialize_image_gpu,
)
from sglang.srt.runtime_context import (
get_device,
get_disagg,
get_exec,
get_forward,
get_parallel,
get_platform,
)
from sglang.srt.utils import (
LazyValue,
add_prefix,
get_bool_env_var,
is_gfx95_supported,
is_gfx942_supported,
is_gfx1250_supported,
log_info_on_rank0,
make_layers,
)
from sglang.srt.utils.custom_op import register_custom_op
from sglang.srt.utils.hf_transformers_utils import get_rope_config
# NPU-only: bind torch_npu here so _compute_q_b / _forward_prepare can call
# torch_npu.npu_rms_norm directly (imports elsewhere aren't visible in this module).
if _is_npu:
import torch_npu
class MhcOps(NamedTuple):
hc_split_sinkhorn: Callable[..., Any]
mhc_fused_post_pre: Optional[Callable[..., Any]]
npu_hc_pre: Optional[Callable[..., Any]]
mhc_pre: Optional[Callable[..., Any]]
mhc_post: Optional[Callable[..., Any]]
fused_hc_head: Optional[Callable[..., Any]]
@functools.cache
def _get_mhc_ops() -> MhcOps:
"""Load MHC kernels only when a DeepSeek-V4 layer needs them.
Model modules are imported eagerly by the registry. Importing
``sglang.kernels.ops.layernorm.mhc`` owns TileLang-backed MHC kernels.
Import it only when a DeepSeek-V4 layer executes so registry discovery
cannot initialize an optional CUDA runtime before unrelated models set up
their communication workspaces. DeepSeek-V4 is the sole consumer here.
"""
if _is_xpu:
from sgl_kernel import (
fused_hc_head,
hc_post,
hc_split_sinkhorn,
mhc_fused_post_pre,
mhc_pre,
)
return MhcOps(
hc_split_sinkhorn=hc_split_sinkhorn,
mhc_fused_post_pre=mhc_fused_post_pre,
npu_hc_pre=None,
mhc_pre=mhc_pre,
mhc_post=hc_post,
fused_hc_head=fused_hc_head,
)
from sglang.kernels.ops.layernorm.mhc import (
hc_split_sinkhorn,
mhc_fused_post_pre,
npu_hc_pre,
)
return MhcOps(
hc_split_sinkhorn=hc_split_sinkhorn,
mhc_fused_post_pre=mhc_fused_post_pre,
npu_hc_pre=npu_hc_pre,
mhc_pre=None,
mhc_post=None,
fused_hc_head=None,
)
logger = logging.getLogger(__name__)
_FP8_WO_A_GEMM = envs.SGLANG_OPT_FP8_WO_A_GEMM.get()
_FP8_WO_A_UE8M0 = _FP8_WO_A_GEMM and DEEPGEMM_SCALE_UE8M0
def wo_a_fp8_gemm_enabled(quant_config: Optional[QuantizationConfig]) -> bool:
"""The fp8 wo_a absorb GEMM (DeepGEMM fp8_einsum, aiter mxscale) takes 128x128
block scales only; any other layout dequantizes wo_a to bf16 at load."""
return (
_FP8_WO_A_GEMM
and isinstance(quant_config, Fp8Config)
and quant_config.weight_block_size == [128, 128]
)
_NPU_BF16_WO_A_GEMM = _is_npu and envs.SGLANG_OPT_NPU_BF16_WO_A_GEMM.get()
_MHC_POST_MULT_VALUE = 2.0
_HC_PRENORM_DEEPGEMM_MIN_TOKENS = 1024
DEEPSEEK_V4_STACKED_PARAMS_MAPPING: List[Tuple[str, str, int]] = [
("gate_up_proj", "gate_proj", 0),
("gate_up_proj", "up_proj", 1),
]
def _is_fused_mhc_post_pre_enabled_xpu() -> bool:
if _is_xpu:
return envs.SGLANG_OPT_FUSE_MHC_POST_PRE.get()
return False
# FlashInfer's mhc_pre_big_fuse only accepts these split-K counts.
_FLASHINFER_MHC_PRE_SPLITS = (1, 2, 4, 8, 16)
@functools.cache
def _cuda_sm_count() -> int:
return torch.cuda.get_device_properties(0).multi_processor_count
def _flashinfer_mhc_pre_num_splits(num_tokens: int, hc_hidden_size: int) -> int:
block_m = block_k = 64
grid_m = (num_tokens + block_m - 1) // block_m
num_block_k = (hc_hidden_size + block_k - 1) // block_k
raw = max(1, min(_cuda_sm_count() // max(grid_m, 1), num_block_k // 4))
best = 1
for split in _FLASHINFER_MHC_PRE_SPLITS:
if split <= raw:
best = split
return best
def _flashinfer_hc_pre(
x: torch.Tensor,
hc_fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
*,
rms_eps: float,
hc_eps: float,
sinkhorn_iters: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
from flashinfer.mhc import mhc_pre_big_fuse
from sglang.srt.layers.deep_gemm_wrapper.entrypoint import tf32_hc_prenorm_gemm
num_tokens, hc_mult, hidden_size = x.shape
hc_hidden_size = hc_mult * hidden_size
mix_dim = hc_fn.shape[0] # hc_mult * (2 + hc_mult) == 24
n_splits = _flashinfer_mhc_pre_num_splits(num_tokens, hc_hidden_size)
dot_mix = torch.empty(
(n_splits, num_tokens, mix_dim), dtype=torch.float32, device=x.device
)
sqrsum = torch.empty((n_splits, num_tokens), dtype=torch.float32, device=x.device)
tf32_hc_prenorm_gemm(
x.reshape(num_tokens, hc_hidden_size), hc_fn, dot_mix, sqrsum, n_splits
)
if n_splits == 1:
dot_mix = dot_mix.squeeze(0)
sqrsum = sqrsum.squeeze(0)
post, comb, layer_input = mhc_pre_big_fuse(
dot_mix,
sqrsum,
x,
hc_scale,
hc_base,
hc_hidden_size,
rms_eps=rms_eps,
mhc_pre_eps=hc_eps,
mhc_sinkhorn_eps=hc_eps,
mhc_post_mult_value=_MHC_POST_MULT_VALUE,
sinkhorn_repeat=sinkhorn_iters,
num_splits=n_splits,
)
return layer_input, post.squeeze(-1), comb
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
# PoC: compute the (replicated TP1) shared expert on LOCAL hidden before the dp
# gather instead of on the gathered global buffer. Requires
# SGLANG_SHARED_EXPERT_TP1=1 (replicated shared expert). Default OFF.
_SHARED_EXPERT_LOCAL = get_bool_env_var("SGLANG_DP_SHARED_EXPERT_LOCAL")
_is_gfx95_supported = is_gfx95_supported()
_is_gfx942_supported = is_gfx942_supported()
_is_gfx1250_supported = is_gfx1250_supported()
if _use_aiter:
if _is_gfx95_supported or _is_gfx1250_supported:
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
def _wo_a_aiter_gemm_eligible(
flag: bool, use_aiter: bool, is_hip: bool, is_gfx95: bool
) -> bool:
"""Static eligibility for the aiter ``wo_a`` reroute.
Folds the opt-in flag, the global ``SGLANG_USE_AITER`` switch, and the
HIP/gfx95 platform gates into one predicate. Evaluated once at import (see
``_wo_a_aiter_batched_gemm_enabled``) so none of it runs on the per-token
decode critical path.
"""
return bool(flag and use_aiter and is_hip and is_gfx95)
# Read the opt-in flag and import the aiter kernel ONCE at module import: the
# decode ``wo_a`` matmul runs per layer/token on the critical path, so it must
# not pay an ``EnvBool.get()`` plus a function-local import on every call. If the
# path is eligible but the kernel import fails, disable it here and fall back to
# the einsum for the process (logged once) instead of retrying every step.
_wo_a_aiter_batched_gemm_enabled = _wo_a_aiter_gemm_eligible(
envs.SGLANG_OPT_USE_AITER_BATCHED_GEMM.get(),
_use_aiter,
_is_hip,
_is_gfx95_supported,
)
_wo_a_batched_gemm_bf16 = None
if _wo_a_aiter_batched_gemm_enabled:
try:
from aiter.ops.triton.gemm.batched.batched_gemm_bf16 import (
batched_gemm_bf16 as _wo_a_batched_gemm_bf16,
)
except Exception as err: # pragma: no cover - env-dependent
_wo_a_aiter_batched_gemm_enabled = False
logger.warning(
"aiter wo_a batched_gemm_bf16 import failed; using einsum for wo_a "
"for the rest of this process: %s",
err,
)
# Flipped once if the (already-imported) aiter kernel raises at runtime, so a
# per-call kernel failure falls back to the einsum for the rest of the process
# instead of re-raising (and re-logging) on every layer/token.
_wo_a_aiter_batched_gemm_disabled = False
# ROCm fp8 wo_a. The CUDA fp8 path below is built on DeepGEMM's fp8_einsum, so
# gfx950 runs the equivalent aiter e8m0 block-scale batched GEMM instead. Both
# the kernel availability and the weight-scale converter resolve once at import;
# ``None`` here means the platform keeps the bf16 absorb GEMM.
_wo_a_fp8_mxscale = None
_wo_a_fp8_mxscale_fused_invrope = None
_wo_a_weight_scale_to_e8m0 = None
if _is_hip:
from sglang.srt.models.deepseek_common.amd.deepseek_v4_wo_a_fp8 import (
apply_wo_a_fp8_mxscale,
apply_wo_a_fp8_mxscale_fused_invrope,
is_wo_a_fp8_fused_invrope_supported,
is_wo_a_fp8_mxscale_supported,
wo_a_weight_scale_to_e8m0,
)
if is_wo_a_fp8_mxscale_supported():
_wo_a_fp8_mxscale = apply_wo_a_fp8_mxscale
_wo_a_weight_scale_to_e8m0 = wo_a_weight_scale_to_e8m0
# Opt-in fused inverse-RoPE + quant front end (env-gated for A/B). Only
# bind it when both the flatmm and the fused aiter op are available.
if (
envs.SGLANG_OPT_FP8_WO_A_FUSED_INVROPE.get()
and is_wo_a_fp8_fused_invrope_supported()
):
_wo_a_fp8_mxscale_fused_invrope = apply_wo_a_fp8_mxscale_fused_invrope
@functools.lru_cache(maxsize=1)
def _fused_wo_a_arch_supported() -> bool:
return _is_cuda and torch.cuda.get_device_capability()[0] == 10
def _apply_wo_a_bf16_matmul(
o: torch.Tensor,
wo_a: torch.Tensor,
is_decode: bool,
is_target_verify: bool = False,
fuse_mxfp8_quant: bool = False,
is_prefill: bool = False,
fast_path: bool = False,
) -> torch.Tensor | Mxfp8SwizzledInput:
# o [T, G, D] @ wo_a [G, R, D] -> [T, G, R]; the fast paths below are gated
# on the exact validated TP4 shapes and write token-major output directly.
global _wo_a_aiter_batched_gemm_disabled
if (
fast_path
and _is_cuda
and (
(
is_decode
and o.shape[0] == 1
and (get_platform().is_blackwell or get_platform().is_sm90)
)
or (
is_target_verify
and 0 < o.shape[0] <= 384
and get_platform().is_blackwell
)
or (
is_prefill
and 4096 <= o.shape[0] <= 65536
and get_platform().is_blackwell
)
)
and o.shape[1:] == (2, 4096)
and wo_a.shape == (2, 1024, 4096)
and o.dtype == wo_a.dtype == torch.bfloat16
and o.stride(2) == 1
and o.stride(1) == 4096
and o.stride(0) >= 8192
and wo_a.is_contiguous()
):
if is_decode and o.shape[0] == 1:
return wo_a_bf16_gemv(o, wo_a)
if 2 <= o.shape[0] <= 8:
if fuse_mxfp8_quant:
return Mxfp8SwizzledInput(*wo_a_bf16_small_batch_mxfp8(o, wo_a))
return wo_a_bf16_small_batch(o, wo_a)
result = torch.empty(
(o.shape[0], wo_a.shape[0], wo_a.shape[1]), dtype=o.dtype, device=o.device
)
# The strided destination keeps the einsum reduction while producing the
# layout wo_b consumes; draft warmup/capture can enter with grad enabled.
with torch.no_grad():
torch.bmm(
o.transpose(0, 1), wo_a.transpose(1, 2), out=result.transpose(0, 1)
)
return result
if (
is_decode
and _wo_a_aiter_batched_gemm_enabled
and not _wo_a_aiter_batched_gemm_disabled
):
try:
# aiter batched_gemm_bf16: XQ[B,M,K] @ WQ[B,N,K]^T -> [B,M,N].
# Here batch = group G: XQ = o.transpose(0,1) [G,T,D], WQ = wo_a
# [G,R,D] -> [G,T,R] -> transpose back to [T,G,R].
xq = o.transpose(0, 1).contiguous()
y = _wo_a_batched_gemm_bf16(xq, wo_a, dtype=torch.bfloat16)
return y.transpose(0, 1).contiguous()
except Exception as err:
_wo_a_aiter_batched_gemm_disabled = True
logger.warning(
"aiter wo_a batched_gemm_bf16 failed; disabling the reroute and "
"falling back to einsum for the rest of this process: %s",
err,
)
return torch.einsum("tgd,grd->tgr", o, wo_a)
def _fused_rmsnorm_fp8_quant(hidden_states, weight, eps):
x_quant, x_bf16, _, _ = fused_rms_fp8_group_quant(
hidden_states,
weight,
eps,
inp2=None,
inp2_weight=None,
inp2_epsilon=None,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=True,
transpose_scale=_use_aiter_bpreshuffle_gfx95,
)
if _use_aiter_bpreshuffle_gfx95:
x_quant = (
x_quant[0],
view_aiter_fused_rms_transposed_fp8_scale(x_quant[1]),
)
return x_quant, x_bf16
def make_hc_mixing_params(
hc_mult: int, hidden_size: int
) -> Tuple[
nn.Parameter, nn.Parameter, nn.Parameter, nn.Parameter, nn.Parameter, nn.Parameter
]:
mix_hc = (2 + hc_mult) * hc_mult
hc_dim = hc_mult * hidden_size
return (
nn.Parameter(torch.empty(mix_hc, hc_dim, dtype=torch.float32)),
nn.Parameter(torch.empty(mix_hc, hc_dim, dtype=torch.float32)),
nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)),
nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)),
nn.Parameter(torch.empty(3, dtype=torch.float32)),
nn.Parameter(torch.empty(3, dtype=torch.float32)),
)
def make_hc_head_params(
hc_mult: int, hidden_size: int
) -> Tuple[nn.Parameter, nn.Parameter, nn.Parameter]:
hc_dim = hc_mult * hidden_size
return (
nn.Parameter(torch.empty(hc_mult, hc_dim, dtype=torch.float32)),
nn.Parameter(torch.empty(hc_mult, dtype=torch.float32)),
nn.Parameter(torch.empty(1, dtype=torch.float32)),
)
def hc_head_torch(
x: torch.Tensor,
hc_fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
*,
norm_eps: float,
hc_eps: float,
) -> torch.Tensor:
shape, dtype = x.size(), x.dtype
x = x.flatten(-2).float()
rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + norm_eps)
mixes = F.linear(x, hc_fn) * rsqrt
pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps
y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=-2)
return y.to(dtype)
_FREQS_CIS_TO_COS_SIN: dict[
Tuple[int, torch.dtype, torch.device], Tuple[torch.Tensor, torch.Tensor]
] = {}
def _freqs_cis_to_cos_sin(
freqs_cis: torch.Tensor, dtype: torch.dtype, device: torch.device
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Derive (cos, sin) bf16 contiguous tables from a complex64 `freqs_cis`,
cached by `(id(freqs_cis), dtype, device)` so that all layers sharing the
same `freqs_cis` (via `precompute_freqs_cis`'s lru_cache) reuse one pair."""
key = (id(freqs_cis), dtype, device)
cached = _FREQS_CIS_TO_COS_SIN.get(key)
if cached is not None:
return cached
fr = torch.view_as_real(freqs_cis)
cos = fr[..., 0].to(device=device, dtype=dtype).contiguous()
sin = fr[..., 1].to(device=device, dtype=dtype).contiguous()
_FREQS_CIS_TO_COS_SIN[key] = (cos, sin)
return cos, sin
def _apply_gguf_grouped_wo_a(
o: torch.Tensor,
qweight: torch.Tensor,
qweight_type: int,
o_lora_rank: int,
matmul_fn: Optional[Callable] = None,
) -> torch.Tensor:
if matmul_fn is None:
from sglang.srt.layers.quantization.gguf import fused_mul_mat_gguf
matmul_fn = fused_mul_mat_gguf
group_outputs = []
for group_id in range(o.shape[1]):
start = group_id * o_lora_rank
group_outputs.append(
matmul_fn(
o[:, group_id, :].contiguous(),
qweight[start : start + o_lora_rank],
qweight_type,
)
)
return torch.stack(group_outputs, dim=1)
if TYPE_CHECKING:
from sglang.srt.layers.attention.deepseek_v4_backend import (
DeepseekV4AttnBackend,
LateLayerTail,
)
from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import (
DeepseekV4HipRadixBackend,
)
from sglang.srt.layers.quantization import QuantizationConfig
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@register_custom_op(mutates_args=["output"])
@register_split_op()
def deepseek_v4_attention_with_output(
query: torch.Tensor,
key_value: torch.Tensor,
output: torch.Tensor,
layer_id: int,
compress_ratio: int,
attn_sink: torch.Tensor,
save_kv_cache: bool,
) -> None:
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
attention_layers = context.attention_layers
attention_layer = attention_layers[layer_id]
real_num_tokens = forward_batch.global_num_token_non_padded_cpu
if real_num_tokens == 0:
output.zero_()
return
query = query[:real_num_tokens]
key_value = key_value[:real_num_tokens]
original_out_cache_loc = forward_batch.out_cache_loc
forward_batch.out_cache_loc = original_out_cache_loc[:real_num_tokens]
attn_backend = get_attn_backend()
try:
ret = attn_backend.forward(
q=query,
k=key_value,
v=key_value,
layer=attention_layer,
forward_batch=forward_batch,
compress_ratio=compress_ratio,
attn_sink=attn_sink,
save_kv_cache=save_kv_cache,
)
finally:
forward_batch.out_cache_loc = original_out_cache_loc
assert output[:real_num_tokens].numel() == ret.numel(), (
f"Output tensor element mismatch: {output[:real_num_tokens].numel()} != {ret.numel()}"
)
output[:real_num_tokens].view(ret.shape).copy_(ret)
output[real_num_tokens:].zero_()
return
bcg_deepseek_v4_attention_with_output = eager_on_graph(True)(
deepseek_v4_attention_with_output
)
def deepseek_v4_low_ratio_sources(layer, x, q_lora, positions) -> None:
# The compressor and prefill indexer sync with the host, like the attention.
forward_batch = get_tc_piecewise_forward_context().forward_batch
real_num_tokens = forward_batch.global_num_token_non_padded_cpu
if real_num_tokens == 0:
return
get_attn_backend().forward_low_ratio_sources(
layer=layer,
x=x[:real_num_tokens],
q_lora=q_lora[:real_num_tokens],
positions=positions[:real_num_tokens],
forward_batch=forward_batch,
)
bcg_deepseek_v4_low_ratio_sources = eager_on_graph(True)(deepseek_v4_low_ratio_sources)
def deepseek_v4_engram_hash_ids(hasher, input_ids: torch.Tensor) -> torch.Tensor:
# The hasher reads per-request rows, so it cannot run inside the CUDA graph.
forward_batch = get_tc_piecewise_forward_context().forward_batch
return hasher(input_ids, forward_batch)
bcg_deepseek_v4_engram_hash_ids = eager_on_graph(True)(deepseek_v4_engram_hash_ids)
class MqaAttentionBase(nn.Module):
# Class-level default for subclasses that read it without running __init__.
wo_a_fp8: bool = False
def __init__(
self,
config: DeepSeekV4Config,
layer_id: int,
quant_config: Optional[QuantizationConfig],
prefix: str,
*,
attn_tp_rank: Optional[int] = None,
attn_tp_size: Optional[int] = None,
compress_ratio: Optional[int] = None,
fuse_wqa_wkv: Optional[bool] = None,
wo_a_fp8: Optional[bool] = None,
wo_a_keeps_quant_config: Optional[bool] = None,
wo_b_reduce_results: Optional[bool] = None,
rope_original_seq_len: Optional[int] = None,
) -> None:
super().__init__()
self.is_dsv41 = getattr(config, "model_type", None) == "deepseek_v41"
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if attn_tp_rank is None or attn_tp_size is None:
attn_tp_rank = get_parallel().attn_tp_rank
attn_tp_size = get_parallel().attn_tp_size
self.attn_tp_rank: int = attn_tp_rank
self.attn_tp_size: int = attn_tp_size
self.layer_id = layer_id
self.dim = config.hidden_size
self.hidden_size = config.hidden_size
self.qk_rope_head_dim = config.qk_rope_head_dim
self.qk_nope_head_dim = config.head_dim - config.qk_rope_head_dim
self.head_dim = self.qk_rope_head_dim + self.qk_nope_head_dim
self.rope_head_dim = config.qk_rope_head_dim
self.n_heads = config.num_attention_heads
self.n_local_heads = self.n_heads // self.attn_tp_size
self.n_groups = config.o_groups
self.n_local_groups = self.n_groups // self.attn_tp_size
self.q_lora_rank = config.q_lora_rank
self.o_lora_rank = config.o_lora_rank
self.eps = config.rms_norm_eps
self.softmax_scale = self.head_dim**-0.5
self.q_head_norm = config.q_head_norm
self.compress_ratio: int = (
compress_ratio
if compress_ratio is not None
else config.compress_ratios[layer_id]
)
assert self.compress_ratio in (
0,
1,
2,
4,
128,
), (
f"compress_ratio: expected one of (0, 1, 2, 4, 128), got {self.compress_ratio}"
)
assert self.head_dim == config.head_dim
assert config.num_key_value_heads == 1
fuse: bool = (
envs.SGLANG_OPT_FUSE_WQA_WKV.get() if fuse_wqa_wkv is None else fuse_wqa_wkv
)
fp8: bool = (
wo_a_fp8_gemm_enabled(quant_config) if wo_a_fp8 is None else wo_a_fp8
)
reduce_results: bool = (
(self.attn_tp_size == get_parallel().tp_size and self.attn_tp_size > 1)
if wo_b_reduce_results is None
else wo_b_reduce_results
)
# NPU arch35 runs wo_a as a batched MXFP8 GEMM instead of deep_gemm's FP8 one,
# but it needs the same quantized weights.
self.use_npu_arch35_mxfp8_wo_a = use_npu_arch35_mxfp8_wo_a(quant_config)
quantize_wo_a = fp8 or self.use_npu_arch35_mxfp8_wo_a
if wo_a_keeps_quant_config is None:
keep_source_quant = (
quant_config is not None and quant_config.get_name() == "expert_pack"
)
wo_a_quant_config: Optional[QuantizationConfig] = (
quant_config if quantize_wo_a or keep_source_quant else None
)
elif wo_a_keeps_quant_config:
wo_a_quant_config = quant_config
else:
wo_a_quant_config = None
self.fuse_wqa_wkv = fuse
self.wo_a_fp8 = fp8
self.attn_sink = nn.Parameter(torch.empty(self.n_heads, dtype=torch.float32))
self._attn_sink_local: Optional[torch.Tensor] = None
if fuse:
self.wqkv_a = ReplicatedLinear(
self.hidden_size,
self.q_lora_rank + self.head_dim,
bias=False,
quant_config=quant_config,
prefix=add_prefix("wqkv_a", prefix),
)
else:
self.wq_a = ReplicatedLinear(
self.hidden_size,
self.q_lora_rank,
bias=False,
quant_config=quant_config,
prefix=add_prefix("wq_a", prefix),
)
self.wkv = ReplicatedLinear(
self.hidden_size,
self.head_dim,
bias=False,
quant_config=quant_config,
prefix=add_prefix("wkv", prefix),
)
self.q_norm = RMSNorm(self.q_lora_rank, eps=self.eps)
self.wq_b = ColumnParallelLinear(
self.q_lora_rank,
self.n_heads * self.head_dim,
bias=False,
quant_config=quant_config,
prefix=add_prefix("wq_b", prefix),
tp_rank=self.attn_tp_rank,
tp_size=self.attn_tp_size,
)
self.kv_norm = RMSNorm(self.head_dim, eps=self.eps)
self.wo_a = ColumnParallelLinear(
self.n_heads * self.head_dim // self.n_groups,
self.n_groups * self.o_lora_rank,
bias=False,
quant_config=wo_a_quant_config,
prefix=add_prefix("wo_a", prefix),
tp_rank=self.attn_tp_rank,
tp_size=self.attn_tp_size,
**({} if quantize_wo_a else {"params_dtype": torch.bfloat16}),
)
if quantize_wo_a:
assert hasattr(self.wo_a, "weight_scale_inv"), (
"FP8 quant_config must create weight_scale_inv"
)
if self.use_npu_arch35_mxfp8_wo_a:
# Read by the NPU arch35 MXFP8 weight processor to batch the
# weight/scale per attention group for npu_transpose_quant_batchmatmul.
self.wo_a._dsv4_npu_arch35_mxfp8_wo_a = True
self.wo_a._dsv4_num_groups = self.n_local_groups
self.wo_a._dsv4_o_lora_rank = self.o_lora_rank
elif fp8:
self.wo_a.weight_scale_inv.format_ue8m0 = _FP8_WO_A_UE8M0
# wo_a is quantized but never *applied* through its quant method:
# the absorb GEMM in forward() reads .weight / .weight_scale_inv and
# runs its own batched kernel (DeepGEMM fp8_einsum on CUDA, aiter
# mxscale BMM on gfx950), both of which want the plain row-major
# [G, R, D] weight. Opt out of any backend-private weight layout the
# linear method would otherwise install for its own GEMM -- on ROCm
# that is aiter's B-preshuffle, which silently permutes the weight
# in place (same shape, dtype and strides) and makes this GEMM
# return noise.
self.wo_a.keep_plain_weight_layout = True
self.wo_b = RowParallelLinear(
self.n_groups * self.o_lora_rank,
self.hidden_size,
bias=False,
quant_config=quant_config,
reduce_results=reduce_results,
prefix=add_prefix("wo_b", prefix),
tp_rank=self.attn_tp_rank,
tp_size=self.attn_tp_size,
)
from sglang.kernels.ops.attention.deepseek_v4_rope import precompute_freqs_cis
rope_theta, rope_scaling = get_rope_config(config)
self.rope_scaling = dict(rope_scaling) if rope_scaling else None
scaling = self.rope_scaling or {}
# RoPE is selected at layer granularity in the reference model. Pure
# SWA layers use the main unscaled RoPE, while C4/C128 layers use the
# compressed YaRN RoPE for Q, their SWA branch, and compressed KV.
self.rope_base = (
config.compress_rope_theta if self.compress_ratio else rope_theta
)
original_seq_len: int = (
rope_original_seq_len
if rope_original_seq_len is not None
else (
scaling["original_max_position_embeddings"]
if self.compress_ratio
else 0
)
)
freqs_cis = precompute_freqs_cis(
dim=self.qk_rope_head_dim,
seqlen=config.max_position_embeddings,
original_seq_len=original_seq_len,
base=self.rope_base,
factor=scaling.get("factor", 1.0),
beta_fast=scaling.get("beta_fast", 32),
beta_slow=scaling.get("beta_slow", 1),
)
self.register_buffer("freqs_cis", freqs_cis, persistent=False)
self.freqs_cis: torch.Tensor
@functools.cached_property
def use_flashinfer_mxfp8_wo_b(self) -> bool:
"""Whether wo_b consumes FlashInfer-swizzled MXFP8, so wo_a can fuse the
quantization into its epilogue. Not known until wo_b's weights load."""
quant_method = getattr(self.wo_b, "quant_method", None)
return getattr(
quant_method, "mxfp8_dense_backend", None
) == Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL and (
getattr(quant_method, "use_mxfp8", False)
or getattr(self.wo_b, "block_fp8_mxfp8_ready", False)
)
def _kernel_num_heads(self, num_tokens: int) -> int:
if self.attn_tp_size == 1:
return self.n_local_heads
if get_platform().is_sm120:
# Prefill already accepts the native per-rank query width.
if num_tokens > SM120_DECODE_MAX_TOKENS:
return self.n_local_heads
if envs.SGLANG_SM120_FLASHMLA_BACKEND.get() == "flashinfer":
from sglang.kernels.ops.attention.flash_mla_sm120 import (
flashinfer_dsv4_decode_supports_num_heads,
)
if flashinfer_dsv4_decode_supports_num_heads(
self.n_local_heads, num_tokens
):
return self.n_local_heads
# Other FlashMLA implementations retain their existing padded shape.
return 64 if self.n_local_heads <= 64 else self.n_heads
def _local_attn_sink(self, kernel_num_heads: Optional[int] = None) -> torch.Tensor:
if self.attn_tp_size == 1:
return self.attn_sink
rank = self.attn_tp_rank
num_heads = self.n_local_heads
padded_num_heads = 64 if num_heads <= 64 else self.n_heads
if kernel_num_heads is None:
# Preserve the legacy contract for subclasses such as DSpark that
# always pad their attention query independently of this helper.
kernel_num_heads = padded_num_heads
assert kernel_num_heads >= num_heads
# Keep one fallback-width allocation and return a view matching Q.
# Prefill and decode can alternate, and CUDA graphs can retain the
# view, so replacing this tensor when the path changes would
# both reallocate every transition and risk invalidating a captured
# pointer.
sink_num_heads = max(kernel_num_heads, padded_num_heads)
if self._attn_sink_local is None:
sink = self.attn_sink.new_zeros(sink_num_heads)
sink[:num_heads] = self.attn_sink[rank * num_heads : (rank + 1) * num_heads]
self._attn_sink_local = sink
return self._attn_sink_local[:kernel_num_heads]
@contextmanager
def maybe_use_decode_attn_tp(self, forward_batch: ForwardBatch):
ctx = get_cp_decode_attn_tp_ctx()
attn = self.attn_mqa if isinstance(self, MQALayer) else self.attn
with ctx.maybe_use_decode_attn_tp(
forward_batch,
[self.wq_b, self.wo_a, self.wo_b],
radix_attn=attn,
):
if ctx.use_decode_attn_tp:
orig = (
self.n_local_heads,
self.n_local_groups,
self.attn_tp_rank,
self.attn_tp_size,
)
decode_tp_size = ctx.decode_tp_size
self.n_local_heads = self.n_heads // decode_tp_size
self.n_local_groups = self.n_groups // decode_tp_size
self.attn_tp_rank = ctx.decode_tp_rank
self.attn_tp_size = decode_tp_size
try:
yield
finally:
(
self.n_local_heads,
self.n_local_groups,
self.attn_tp_rank,
self.attn_tp_size,
) = orig
else:
yield
class MQALayer(MqaAttentionBase):
is_dsv41: bool = False
def __init__(
self,
config: DeepSeekV4Config,
layer_id: int,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
alt_streams: Optional[List[torch.cuda.Stream]] = None,
compress_ratio_override: Optional[int] = None,
) -> None:
super().__init__(
config,
layer_id,
quant_config,
prefix,
compress_ratio=compress_ratio_override,
)
active_rope_scaling = None
if self.compress_ratio:
active_rope_scaling = dict(self.rope_scaling or {})
active_rope_scaling["rope_type"] = "deepseek_yarn"
self.rotary_emb = get_rope_wrapper(
head_size=self.rope_head_dim,
rotary_dim=self.rope_head_dim,
max_position=config.max_position_embeddings,
base=self.rope_base,
rope_scaling=active_rope_scaling,
is_neox_style=False,
device=get_device().device,
)
if _is_npu:
rope = Dsv4NpuRoPE.for_freqs(
self.freqs_cis, getattr(self, "rotary_emb", None)
)
# fp32 tables feed the compressor gather; bf16 tables make the
# activation-dtype gathers cast-free. Bit-identical values:
# rounding the table once equals rounding each gathered element.
rope.ensure_tables(torch.float32)
rope.ensure_tables(torch.bfloat16)
# npu_rms_norm has no weight-free overload; the per-head q norm
# reads this cached ones vector instead of paying a per-call
# alloc + fill.
self.register_buffer(
"q_rms_norm_ones",
torch.ones(self.head_dim, dtype=torch.bfloat16),
persistent=False,
)
if _is_hip:
cos_cache = (
self.freqs_cis.real.to(torch.bfloat16).unsqueeze(-2).unsqueeze(-2)
)
sin_cache = (
self.freqs_cis.imag.to(torch.bfloat16).unsqueeze(-2).unsqueeze(-2)
)
self.register_buffer("cos_cache", cos_cache, persistent=False)
self.register_buffer("sin_cache", sin_cache, persistent=False)
if alt_streams is not None and (
(_is_cuda and envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get())
or (_is_npu and envs.SGLANG_NPU_USE_MULTI_STREAM.get())
):
self.alt_streams = alt_streams[:3]
self.alt_streams_indexer = alt_streams[-2:]
else:
self.alt_streams = None
self.alt_streams_indexer = None
self._multi_stream_bs_limit = 128 if get_platform().is_blackwell else 64
self.compressor = None
self.indexer = None
if self.compress_ratio in (4, 128):
expert_pack_quant_config = (
quant_config
if quant_config is not None and quant_config.get_name() == "expert_pack"
else None
)
self.compressor = Compressor(
config,
layer_id=self.layer_id,
is_in_indexer=False,
freqs_cis=self.freqs_cis,
compress_ratio=self.compress_ratio,
head_dim=self.head_dim,
rotate=False,
prefix=add_prefix("compressor", prefix),
quant_config=expert_pack_quant_config,
rotary_emb=self.rotary_emb,
)
if self.compress_ratio == 4:
self.indexer = C4Indexer(
config,
freqs_cis=self.freqs_cis,
layer_id=layer_id,
quant_config=quant_config,
prefix=add_prefix("indexer", prefix),
alt_streams=self.alt_streams_indexer,
rotary_emb=self.rotary_emb,
fp4_cos=(self.cos_cache[:, 0, 0, :] if _is_hip else None),
fp4_sin=(self.sin_cache[:, 0, 0, :] if _is_hip else None),
)
elif self.compress_ratio in (1, 2):
# The layers in between read both through the attention backend.
if self.layer_id in config.kv_source_layer_ids:
self.compressor = DeepseekV41Compressor(
hidden_size=config.hidden_size,
head_dim=self.head_dim,
compress_ratio=self.compress_ratio,
eps=config.rms_norm_eps,
)
if self.layer_id in config.index_source_layer_ids:
self.indexer = DeepseekV41Indexer(
config,
layer_id=self.layer_id,
head_dim=self.head_dim,
quant_config=quant_config,
prefix=add_prefix("indexer", prefix),
)
self.attn_mqa = RadixAttention(
self.n_local_heads,
self.head_dim,
self.softmax_scale,
num_kv_heads=1,
layer_id=layer_id,
quant_config=quant_config,
prefix=add_prefix("attn_mqa", prefix),
)
self.use_fused_qk_norm_rope = (
_is_hip and envs.SGLANG_OPT_USE_FUSED_QK_NORM_ROPE.get()
)
# Static eligibility; token count and wo_b's output format are checked
# in forward, after weights have loaded.
self.use_fused_wo_a = (
self.is_dsv41
and envs.SGLANG_DSV41_FUSED_WO_A.get()
and _fused_wo_a_arch_supported()
and not self.wo_a_fp8
and not self.use_npu_arch35_mxfp8_wo_a
and self.wo_a.weight.dtype == torch.bfloat16
and self.wo_a.weight.shape == (self.n_local_groups * self.o_lora_rank, 4096)
and (self.n_local_groups, self.o_lora_rank) == (2, 1024)
)
# KV cache write is always fused into the K kernel
# (`_compute_kv_to_cache`), so the legacy "overlap store cache" flag
# has no effect here -- the fused path is on by default.
def _apply(self, fn, recurse=True):
result = super()._apply(fn, recurse=recurse)
if self.indexer is not None and hasattr(self.indexer.compressor, "fp4_cos"):
self.indexer.compressor.fp4_cos = self.cos_cache[:, 0, 0, :]
self.indexer.compressor.fp4_sin = self.sin_cache[:, 0, 0, :]
return result
def _get_npu_rope_position_cache(
self,
forward_batch: ForwardBatch,
positions: torch.Tensor,
dtype: torch.dtype,
inverse: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor]:
# ``rotary_emb`` is shared by layers with the same RoPE configuration and
# can also be shared by the target and NextN models. Only the immutable
# full table is cached on it; position-gathered tensors are memoized per
# forward (prime_rope_cos_sin / rope_cos_sin), never across forwards --
# reusing them based on shape alone gives MTP decode the previous step's
# RoPE values when positions change but batch size does not.
return rope_cos_sin(
self.freqs_cis,
getattr(self, "rotary_emb", None),
forward_batch,
positions,
dtype,
inverse=inverse,
)
def accepts_mxfp8_swizzled_input(self) -> bool:
"""Whether the first projection consumes a 128x4 MXFP8 activation tuple."""
cached = getattr(self, "_accepts_mxfp8_swizzled_input", None)
if cached is not None:
return cached
if self.fuse_wqa_wkv:
linears = [getattr(self, "wqkv_a", None)]
else:
# Both projections read the same activation on this path.
linears = [getattr(self, "wq_a", None), getattr(self, "wkv", None)]
def _takes_swizzled(linear) -> bool:
method = getattr(linear, "quant_method", None)
return bool(
linear is not None
and getattr(method, "mxfp8_dense_backend", None)
in (
Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL,
Mxfp8DenseGemmBackend.FLASHINFER_CUTLASS,
)
and (
getattr(method, "use_mxfp8", False)
or getattr(linear, "block_fp8_mxfp8_ready", False)
)
)
ok = all(_takes_swizzled(linear) for linear in linears)
self._accepts_mxfp8_swizzled_input = ok
return ok
def _normalize_q_lora(
self, q: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor | Mxfp8SwizzledInput]:
# The indexer needs the BF16 normalized row; wq_b needs the quantized one.
method = self.wq_b.quant_method
if (
_is_cuda
and self.is_dsv41
and get_platform().is_blackwell
and q.dtype == self.q_norm.weight.dtype == torch.bfloat16
and q.ndim == 2
and 0 < q.shape[0] <= 8
and q.shape[1] == 1280
and q.stride(1) == 1
and getattr(method, "mxfp8_dense_backend", None)
== Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL
and (
getattr(method, "use_mxfp8", False)
or getattr(self.wq_b, "block_fp8_mxfp8_ready", False)
)
):
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
from sglang.srt.runtime_context import get_exec
if not (
is_batch_invariant_mode_enabled()
or get_exec().deterministic.enable_deterministic_inference
):
from sglang.kernels.ops.layernorm.mxfp8_epilogue import rmsnorm_mxfp8
y, quant, scale = rmsnorm_mxfp8(
q, self.q_norm.weight, self.q_norm.variance_epsilon
)
return y, Mxfp8SwizzledInput(quant, scale)
q = self.q_norm(q)
return q, q
def _compute_q_a(
self,
x: torch.Tensor,
qkv_a: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor | Mxfp8SwizzledInput]:
if qkv_a is not None:
q = qkv_a[..., : self.q_lora_rank]
else:
q, _ = self.wq_a(x)
return self._normalize_q_lora(q)
def _compute_q_b(
self,
q: torch.Tensor,
positions: torch.Tensor,
q_out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
q, _ = self.wq_b(q)
q = q.view(-1, self.n_local_heads, self.head_dim)
if not self.q_head_norm:
if (
_is_cuda
and q_out is not None
and (
0 < q.shape[0] <= 8
or (
self.is_dsv41
and get_platform().is_blackwell
and self.n_local_heads == 16
and 4096 <= q.shape[0] <= 65536
)
)
and self.head_dim == 512
and self.qk_rope_head_dim == 64
and q.dtype == q_out.dtype == torch.bfloat16
and q.stride(1) == q_out.stride(1) == 512
and q.stride(2) == q_out.stride(2) == 1
):
from sglang.kernels.ops.attention.dsv4.q_rope_store import q_rope_store
q_rope_store(q, q_out, self.freqs_cis, positions)
return q_out
fused_rope_inplace(
q[..., -self.qk_rope_head_dim :],
None,
self.freqs_cis,
positions=positions,
)
if q_out is None:
return q
q_out.copy_(q)
return q_out
if q_out is None:
q_out = torch.empty_like(q)
# Fused warp-per-(token, head) rmsnorm-self + RoPE + write to q_out.
fused_q_norm_rope(q, q_out, self.eps, self.freqs_cis, positions)
return q_out
def _compute_kv_to_cache(
self,
x: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
attn_backend,
qkv_a: Optional[torch.Tensor] = None,
) -> None:
"""Fused: rmsnorm + RoPE + write directly to FlashMLA paged cache.
Replaces the bf16-kv-intermediate path. Used everywhere except the DSA
prefill-CP case (which needs bf16 kv for the cross-rank all-gather).
"""
if envs.SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE.get():
# Quantize the nope payload from bf16-rounded values (the fused
# kernel quantizes from fp32 registers; the bf16 rounding moves
# values across fp8 bins relative to bf16-sourced consumers).
kv = self._compute_kv_bf16(x, positions, qkv_a=qkv_a)
attn_backend.store_cache(
layer_id=self.layer_id, swa_k=kv, forward_batch=forward_batch
)
return
if qkv_a is not None:
kv = qkv_a[..., self.q_lora_rank :]
else:
kv, _ = self.wkv(x)
token_to_kv_pool = get_token_to_kv_pool()
if TYPE_CHECKING:
assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
token_to_kv_pool.set_swa_key_buffer_radix_fused_norm_rope(
layer_id=self.layer_id,
swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch),
kv=kv,
kv_weight=self.kv_norm.weight.data,
eps=self.eps,
freqs_cis=self.freqs_cis,
positions=positions,
)
def _compute_kv_bf16(
self,
x: torch.Tensor,
positions: torch.Tensor,
qkv_a: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Bf16-kv path used by the DSA prefill-CP case (needs all-gather)."""
if qkv_a is not None:
kv = qkv_a[..., self.q_lora_rank :]
else:
kv, _ = self.wkv(x)
kv = kv.contiguous()
fused_norm_rope_inplace(
kv,
self.kv_norm.weight.data,
self.eps,
self.freqs_cis,
positions,
)
return kv
def _forward_prepare_multi_stream(
self,
x: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
attn_backend,
q_out: Optional[torch.Tensor] = None,
x_quant=None,
) -> torch.Tensor:
assert self.alt_streams is not None
assert len(self.alt_streams) >= 3
current_stream = torch.cuda.current_stream()
stream_kv = self.alt_streams[0]
stream_compressor = self.alt_streams[1]
stream_indexer = self.alt_streams[2]
stream_kv.wait_stream(current_stream)
stream_compressor.wait_stream(current_stream)
stream_indexer.wait_stream(current_stream)
x_linear = x_quant if x_quant is not None else x
qkv_a: Optional[torch.Tensor] = None
qkv_a_ready: Optional[torch.cuda.Event] = None
if self.fuse_wqa_wkv:
qkv_a, _ = self.wqkv_a(x_linear)
qkv_a_ready = current_stream.record_event()
q_lora, q_for_wqb = self._compute_q_a(x_linear, qkv_a=qkv_a)
q_lora_ready = current_stream.record_event()
if self.indexer is not None:
with torch.cuda.stream(stream_indexer):
self.indexer(
x=x,
q_lora=q_lora,
forward_batch=forward_batch,
attn_backend=attn_backend,
enable_multi_stream=True,
q_lora_ready=q_lora_ready,
)
with torch.cuda.stream(stream_kv):
if qkv_a_ready is not None:
stream_kv.wait_event(qkv_a_ready)
# Fused norm + rope + cache write -- no bf16 KV intermediate.
self._compute_kv_to_cache(
x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
)
if self.compressor is not None:
with torch.cuda.stream(stream_compressor):
attn_backend.forward_core_compressor(
x, forward_batch, self.layer_id, self.compressor
)
q = self._compute_q_b(q_for_wqb, positions, q_out)
current_stream.wait_stream(stream_kv)
current_stream.wait_stream(stream_compressor)
current_stream.wait_stream(stream_indexer)
del qkv_a
return q
def _forward_prepare_low_ratio_multi_stream(
self,
x: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
attn_backend,
q_out: Optional[torch.Tensor] = None,
x_quant=None,
) -> torch.Tensor:
# Both side streams are joined before returning, and nothing they read is
# released before the join.
assert self.alt_streams is not None
current_stream = torch.cuda.current_stream()
stream_kv = self.alt_streams[0]
stream_sources = self.alt_streams[-1]
x_linear = x_quant if x_quant is not None else x
# NOTE: wait for x ready
if self.compressor is not None:
stream_sources.wait_stream(current_stream)
qkv_a: Optional[torch.Tensor] = None
if self.fuse_wqa_wkv:
qkv_a, _ = self.wqkv_a(x_linear)
if self.compressor is not None:
with torch.cuda.stream(stream_sources):
attn_backend.forward_low_ratio_sources(
layer=self,
x=x,
q_lora=None,
positions=positions,
forward_batch=forward_batch,
run_indexer=False,
)
stream_kv.wait_stream(current_stream)
q_lora, q_for_wqb = self._compute_q_a(x_linear, qkv_a=qkv_a)
# NOTE: wait for the q_lora ready
if self.indexer is not None:
stream_sources.wait_stream(current_stream)
q = self._compute_q_b(q_for_wqb, positions, q_out)
if self.indexer is not None:
# Forked above, right after q_lora; recorded here, after the Q chain.
with torch.cuda.stream(stream_sources):
attn_backend.forward_low_ratio_sources(
layer=self,
x=x,
q_lora=q_lora,
positions=positions,
forward_batch=forward_batch,
run_compressor=False,
)
with torch.cuda.stream(stream_kv):
self._compute_kv_to_cache(
x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
)
current_stream.wait_stream(stream_kv)
if self.compressor is not None or self.indexer is not None:
current_stream.wait_stream(stream_sources)
return q
def _forward_prepare_multi_stream_npu(
self,
x: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
attn_backend,
q_out: Optional[torch.Tensor] = None,
x_quant=None,
) -> torch.Tensor:
# NPU multi-stream: KV on stream_kv, Q on stream_q, overlapped with
# indexer/compressor on current. rope is split; the kv-only call passes
# kv.unsqueeze(1) as q_rope so the op sees [T,1,1,head_dim] like the
# fused path.
assert self.alt_streams is not None
current_stream = torch.npu.current_stream()
stream_kv = self.alt_streams[0]
stream_q = self.alt_streams[1]
stream_kv.wait_stream(current_stream)
stream_q.wait_stream(current_stream)
x_linear = x_quant if x_quant is not None else x
qkv_a: Optional[torch.Tensor] = None
qkv_a_ready = None
if self.fuse_wqa_wkv:
qkv_a, _ = self.wqkv_a(x_linear)
qkv_a_ready = current_stream.record_event()
if qkv_a is not None:
q_lora = qkv_a[..., : self.q_lora_rank]
else:
q_lora, _ = self.wq_a(x_linear)
q_lora = self.q_norm(q_lora)
q_lora_ready = current_stream.record_event()
# KV block on stream_kv.
with torch.npu.stream(stream_kv):
if qkv_a_ready is not None:
stream_kv.wait_event(qkv_a_ready)
if qkv_a is not None:
kv = qkv_a[..., self.q_lora_rank :]
else:
kv, _ = self.wkv(x)
kv = self.kv_norm(kv)
cos4_k, sin4_k = self._get_npu_rope_position_cache(
forward_batch, positions, kv.dtype, inverse=False
)
Dsv4NpuRoPE.apply_rotary_mul_inplace(
kv.unsqueeze(1),
None,
cos4_k,
sin4_k,
qk_nope_dim=self.qk_nope_head_dim,
)
attn_backend.store_cache(
layer_id=self.layer_id,
swa_k=kv,
forward_batch=forward_batch,
)
# Q block on stream_q (needs only q_lora).
with torch.npu.stream(stream_q):
stream_q.wait_event(q_lora_ready)
q, _ = self.wq_b(q_lora)
q = q.view(-1, self.n_local_heads, self.head_dim)
q = torch_npu.npu_rms_norm(q, self.q_rms_norm_ones, self.eps)[0]
cos4_q, sin4_q = self._get_npu_rope_position_cache(
forward_batch, positions, q.dtype, inverse=False
)
Dsv4NpuRoPE.apply_rotary_mul_inplace(
q,
None,
cos4_q,
sin4_q,
qk_nope_dim=self.qk_nope_head_dim,
)
if q_out is not None:
q_out.copy_(q)
q.record_stream(stream_q)
# Indexer + compressor: serial on current.
if self.indexer is not None:
self.indexer(
x=x,
q_lora=q_lora,
forward_batch=forward_batch,
attn_backend=attn_backend,
)
if self.compressor is not None:
attn_backend.forward_core_compressor(
x,
forward_batch,
self.layer_id,
self.compressor,
)
# Join stream_kv + stream_q before downstream attention.
current_stream.wait_stream(stream_kv)
current_stream.wait_stream(stream_q)
del qkv_a
return q
def _forward_prepare_multi_stream_hip(
self,
x: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
attn_backend,
q_out: Optional[torch.Tensor] = None,
x_quant=None,
) -> torch.Tensor:
"""ATOM-style ROCm path: overlap compressors, keep Q/KV on main stream."""
assert self.alt_streams is not None
assert len(self.alt_streams) >= 1
current_stream = torch.cuda.current_stream()
stream_compressor = self.alt_streams[0]
stream_indexer_compressor = (
self.alt_streams[1] if len(self.alt_streams) > 1 else None
)
if self.compressor is not None:
stream_compressor.wait_stream(current_stream)
with torch.cuda.stream(stream_compressor):
attn_backend.forward_core_compressor(
x, forward_batch, self.layer_id, self.compressor
)
if self.indexer is not None and stream_indexer_compressor is not None:
stream_indexer_compressor.wait_stream(current_stream)
with torch.cuda.stream(stream_indexer_compressor):
attn_backend.forward_indexer_compressor(
x=x,
forward_batch=forward_batch,
layer_id=self.indexer.layer_id,
compressor=self.indexer.compressor,
)
x_linear = x_quant if x_quant is not None else x
if self.fuse_wqa_wkv:
qkv_a, _ = self.wqkv_a(x_linear)
q_lora = qkv_a[..., : self.q_lora_rank]
else:
q_lora, _ = self.wq_a(x_linear)
qkv_a = None
if self.use_fused_qk_norm_rope:
if _is_gfx95_supported or _is_gfx1250_supported:
q_for_wqb, q_lora = _fused_rmsnorm_fp8_quant(
q_lora,
self.q_norm.weight,
self.q_norm.variance_epsilon,
)
q, _ = self.wq_b(q_for_wqb)
else:
q_lora = self.q_norm(q_lora)
q, _ = self.wq_b(q_lora)
kv = (
qkv_a[..., self.q_lora_rank :]
if qkv_a is not None
else self.wkv(x_linear)[0]
)
from sglang.kernels.ops.attention.fused_qk_norm_rope_store import (
fused_qk_norm_rope_swa_store,
)
token_to_kv_pool = get_token_to_kv_pool()
swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch)
swa_cache = token_to_kv_pool.get_swa_raw_buffer(self.layer_id)
swa_page_size = token_to_kv_pool.swa_page_size
q = fused_qk_norm_rope_swa_store(
q=q,
kv=kv,
q_norm_weight=None,
kv_norm_weight=self.kv_norm.weight,
q_rms_eps=self.eps,
kv_rms_eps=self.eps,
rope_head_dim=self.qk_rope_head_dim,
cos_cache=self.cos_cache,
sin_cache=self.sin_cache,
positions=positions,
swa_cache=swa_cache,
swa_loc=swa_loc,
swa_page_size=swa_page_size,
q_out=q_out,
dtype=x.dtype,
)
else:
q_lora = self.q_norm(q_lora)
q = self._compute_q_b(q_lora, positions, q_out)
self._compute_kv_to_cache(
x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
)
del qkv_a
if self.indexer is not None:
current_stream.wait_stream(stream_compressor)
if stream_indexer_compressor is not None:
current_stream.wait_stream(stream_indexer_compressor)
self.indexer(
x=x,
q_lora=q_lora,
forward_batch=forward_batch,
attn_backend=attn_backend,
skip_compressor=True,
)
elif self.compressor is not None:
current_stream.wait_stream(stream_compressor)
return q
def _forward_prepare(
self,
x: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
attn_backend,
q_out: Optional[torch.Tensor] = None,
x_quant=None,
q_rope_out: Optional[torch.Tensor] = None,
k_nope_out: Optional[torch.Tensor] = None,
k_rope_out: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
x_linear = x_quant if x_quant is not None else x
if self.fuse_wqa_wkv:
qkv_a, _ = self.wqkv_a(x_linear)
q_lora = qkv_a[..., : self.q_lora_rank]
else:
q_lora, _ = self.wq_a(x_linear)
qkv_a = None
use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)
kv: Optional[torch.Tensor]
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
is_unified_kv_fp8,
is_unified_kv_triton,
)
unified = is_unified_kv_triton()
fp8_2buff = is_unified_kv_fp8()
is_decode = forward_batch.forward_mode.is_decode_or_idle()
# The kernel is token-indexed (q, kv and positions are all length M), so
# a verify batch carrying several draft tokens per request is a shape it
# already handles. Only the cache store differs between decode and
# verify, and under fp8 that store takes the packed pair instead of bf16.
fuse_verify = (
envs.SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY.get()
and forward_batch.forward_mode.is_target_verify()
)
# fp8 verify packs like prefill but keeps verify's store timing: the pair
# lands in the caller's buffers and the backend writes the ring off the
# per-token slot map before attention. Keyed off those buffers the same
# way fuse_prefill is, so the two arms cannot disagree about the layout.
fuse_verify_fp8 = (
fuse_verify
and unified
and fp8_2buff
and k_nope_out is not None
and k_rope_out is not None
)
# Prefill under fp8 goes through the same fused store: the 2-source
# kernel reads this chunk as its extend region in the pool's packed form,
# and the ring write after attention reuses those same rows, so they are
# materialised once here rather than quantized on both sides. Keyed off
# the caller's buffers the way q_rope_out keys the packed Q, so the two
# cannot disagree about the layout; both halves are required because the
# nope one leaves on the kv slot and a missing one would read as "the
# fused store did not run". Verify packs the same way but is its own arm
# above: it stores before attention, not after.
fuse_prefill = (
unified
and fp8_2buff
and k_nope_out is not None
and k_rope_out is not None
and not is_decode
and not forward_batch.forward_mode.is_target_verify()
)
do_fused_qk_norm_rope = (
unified and (is_decode or fuse_verify or fuse_prefill)
) or (not unified and self.use_fused_qk_norm_rope)
if do_fused_qk_norm_rope:
if _is_gfx95_supported or _is_gfx1250_supported:
q_for_wqb, q_lora = _fused_rmsnorm_fp8_quant(
q_lora,
self.q_norm.weight,
self.q_norm.variance_epsilon,
)
q, _ = self.wq_b(q_for_wqb)
else:
q_lora, q_for_wqb = self._normalize_q_lora(q_lora)
q, _ = self.wq_b(q_for_wqb)
kv = (
qkv_a[..., self.q_lora_rank :]
if qkv_a is not None
else self.wkv(x_linear)[0]
)
token_to_kv_pool = get_token_to_kv_pool()
swa_rope_cache = None
if unified and fuse_verify:
# Target-verify runs through the unified_kv decode path. The
# backend writes the current chunk's KV into the ring *before*
# attention (save_kv_cache=True -> store_swa_into_unified ahead
# of runtime.decode), and per-token causal index streams -- built
# once per step in the backend metadata -- keep each draft query
# attending only to positions up to itself. Causal masking among
# the draft tokens comes from those index streams, not from store
# timing. So this path skips only the fused kernel's *own* store
# and returns kv, letting that existing causally-indexed backend
# store run unchanged; we fuse just the norm+RoPE. swa_loc is not
# computed -- it only addresses the kernel store this path drops.
#
# kv is a strided slice of qkv_a and the ring store requires a
# contiguous buffer, so materialise it before the kernel norms
# it in place. The unfused path pays the same copy inside
# _compute_kv_bf16.
#
# Under fp8 the kernel writes the packed pair to the caller's
# buffers rather than norming kv in place, and the same backend
# store takes that pair -- only the row format changes.
kv = kv.contiguous()
swa_cache, swa_loc = None, None
swa_page_size, bf16_store = 1, not fuse_verify_fp8
elif unified and fuse_prefill:
# No pools, so the kernel norms + RoPEs + packs and writes no
# ring row. It must not: those rows are this fwd's extend region
# and the prefix pool has to stay as attention expects to find
# it. The backend stores them after attention from the pair.
swa_cache, swa_loc = None, None
swa_page_size, bf16_store = 1, False
# kv stays the strided slice of qkv_a. Under fp8 the kernel only
# reads it -- the packed pair goes to k_nope_out/k_rope_out, it
# does not norm in place -- and it takes the row stride as an
# argument, so materialising it was a copy on every fp8 layer.
elif unified:
swa_cache = token_to_kv_pool.get_unified_kv(self.layer_id)
# swa_loc is layer-independent; computed once per forward by the
# backend and cached on the metadata (read here by every layer).
swa_loc = attn_backend.get_unified_swa_loc(forward_batch)
swa_page_size, bf16_store = 1, not fp8_2buff
if fp8_2buff:
swa_rope_cache = token_to_kv_pool.get_unified_kv_rope(self.layer_id)
# kv stays the strided slice of qkv_a -- the group-quant
# kernel takes the row stride as an argument.
else:
swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch)
swa_cache = token_to_kv_pool.get_swa_raw_buffer(self.layer_id)
swa_page_size, bf16_store = (
token_to_kv_pool.swa_page_size,
False,
)
from sglang.kernels.ops.attention.fused_qk_norm_rope_store import (
fused_qk_norm_rope_swa_store,
)
q = fused_qk_norm_rope_swa_store(
q=q,
kv=kv,
q_norm_weight=None,
kv_norm_weight=self.kv_norm.weight,
q_rms_eps=self.eps,
kv_rms_eps=self.eps,
rope_head_dim=self.qk_rope_head_dim,
cos_cache=self.cos_cache,
sin_cache=self.sin_cache,
positions=positions,
swa_cache=swa_cache,
swa_loc=swa_loc,
swa_page_size=swa_page_size,
q_out=q_out,
dtype=x.dtype,
bf16_store=bf16_store,
fp8_2buff=fp8_2buff,
swa_rope_cache=swa_rope_cache,
k_nope_out=k_nope_out if (fuse_prefill or fuse_verify_fp8) else None,
k_rope_out=k_rope_out if (fuse_prefill or fuse_verify_fp8) else None,
q_rope_out=q_rope_out,
)
# On the verify path the kernel normed + RoPE'd kv in place and wrote
# nothing, so hand it back: the caller feeds it to attention as the
# current chunk (attn_k = kv) and save_kv_cache = kv is not None lets
# the backend do its normal causally-indexed store into the ring
# before the decode kernel runs -- exactly as the unfused path did.
if unified and (fuse_prefill or fuse_verify_fp8):
# The packed nope half rides out on the kv slot -- attention
# takes it as attn_k and save_kv_cache stays on so the backend
# does the ring write. Its rope half went to the caller's buffer,
# which has no second return slot here. Prefill's write lands
# after attention, verify's before it; both read this pair.
kv = k_nope_out
elif not (unified and fuse_verify):
kv = None
if not unified and use_cp:
# DSA CP: keep bf16 kv around for the cross-rank all-gather, then
# write to the FlashMLA cache after gather.
kv = self._compute_kv_bf16(x, positions, qkv_a=qkv_a)
kv = cp_materialize_global_token_order(
kv.contiguous(),
forward_batch,
torch.cuda.current_stream(),
)
elif _is_npu:
q_lora = self.q_norm(q_lora)
q, _ = self.wq_b(q_lora)
q = q.view(-1, self.n_local_heads, self.head_dim)
q = torch_npu.npu_rms_norm(q, self.q_rms_norm_ones, self.eps)[0]
if qkv_a is not None:
kv = qkv_a[..., self.q_lora_rank :]
else:
kv, _ = self.wkv(x)
kv = self.kv_norm(kv)
cos4, sin4 = self._get_npu_rope_position_cache(
forward_batch, positions, q.dtype, inverse=False
)
Dsv4NpuRoPE.apply_rotary_mul_inplace(
q,
kv.unsqueeze(1),
cos4,
sin4,
qk_nope_dim=self.qk_nope_head_dim,
)
kv_for_cache = kv
if use_cp:
kv_for_cache = cp_gather_full_sequence_states(
kv.contiguous(),
forward_batch,
torch.cuda.current_stream(),
)
attn_backend.store_cache(
layer_id=self.layer_id,
swa_k=kv_for_cache,
forward_batch=forward_batch,
)
kv = None
if q_out is not None:
q_out.copy_(q)
else:
q_lora, q_for_wqb = self._normalize_q_lora(q_lora)
q = self._compute_q_b(q_for_wqb, positions, q_out)
if unified:
# unified_kv prefill: keep bf16 kv; the backend writes
# the ring AFTER attention (2-source path).
kv = self._compute_kv_bf16(x_linear, positions, qkv_a=qkv_a)
elif use_cp and not self.is_dsv41:
kv = self._compute_kv_bf16(x_linear, positions, qkv_a=qkv_a)
kv = cp_materialize_global_token_order(
kv.contiguous(),
forward_batch,
torch.cuda.current_stream(),
)
attn_backend.store_cache(
layer_id=self.layer_id,
swa_k=kv,
forward_batch=forward_batch,
)
elif use_cp:
# every rank writes the whole chunk's window KV with the fused fp32 store
if qkv_a is not None:
kv = qkv_a[..., self.q_lora_rank :]
else:
kv, _ = self.wkv(x_linear)
kv = cp_materialize_global_token_order(
kv.contiguous(),
forward_batch,
torch.cuda.current_stream(),
)
tail = attn_backend.forward_metadata.late_layer_tail
global_positions = (
tail.pos_global
if tail is not None
else forward_batch.positions[: kv.shape[0]]
)
get_token_to_kv_pool().set_swa_key_buffer_radix_fused_norm_rope(
layer_id=self.layer_id,
swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch),
kv=kv,
kv_weight=self.kv_norm.weight.data,
eps=self.eps,
freqs_cis=self.freqs_cis,
positions=global_positions,
)
kv = None
else:
self._compute_kv_to_cache(
x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
)
kv = None
del qkv_a
if self.compress_ratio in (1, 2) and (
self.compressor is not None or self.indexer is not None
):
if (
forward_batch.forward_mode.is_extend()
and is_in_breakable_cuda_graph()
and (
dsa_use_prefill_cp(forward_batch)
or not getattr(attn_backend, "low_ratio_prefill_graph", False)
)
):
bcg_deepseek_v4_low_ratio_sources(self, x, q_lora, positions)
else:
attn_backend.forward_low_ratio_sources(
layer=self,
x=x,
q_lora=q_lora,
positions=positions,
forward_batch=forward_batch,
)
else:
use_npu_cp_full_metadata = use_cp and _is_npu
if self.indexer is not None:
if use_npu_cp_full_metadata:
with attn_backend.use_dsv4_cp_full_metadata(forward_batch):
attn_backend.forward_indexer_compressor(
x,
forward_batch,
self.indexer.layer_id,
self.indexer.compressor,
)
self.indexer(
x=x,
q_lora=q_lora,
forward_batch=forward_batch,
attn_backend=attn_backend,
skip_compressor=True,
)
else:
self.indexer(
x=x,
q_lora=q_lora,
forward_batch=forward_batch,
attn_backend=attn_backend,
)
if self.compressor is not None:
if use_npu_cp_full_metadata:
with attn_backend.use_dsv4_cp_full_metadata(forward_batch):
attn_backend.forward_core_compressor(
x,
forward_batch,
self.layer_id,
self.compressor,
)
else:
attn_backend.forward_core_compressor(
x,
forward_batch,
self.layer_id,
self.compressor,
)
return q, kv
def forward(
self,
x: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
x_quant=None,
) -> torch.Tensor:
if not get_attn_tp_context().input_scattered and x.shape[0] == 0:
return x
attn_backend = get_attn_backend()
if TYPE_CHECKING:
assert isinstance(
attn_backend,
(DeepseekV4AttnBackend, DeepseekV4HipRadixBackend),
)
enable_multi_stream = (
envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get()
and self.alt_streams is not None
and get_is_capture_mode()
and (
is_in_breakable_cuda_graph()
or x.shape[0] <= self._multi_stream_bs_limit
)
and not (self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch))
and not (_is_hip and self.compressor is None)
and self.compress_ratio not in (1, 2)
) or (
_is_npu
and envs.SGLANG_NPU_USE_MULTI_STREAM.get()
and self.alt_streams is not None
and x.shape[0] <= self._multi_stream_bs_limit
and not forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed()
)
low_ratio_multi_stream = (
_is_cuda
and get_platform().is_blackwell
and self.compress_ratio in (1, 2)
and self.alt_streams is not None
and (
forward_batch.forward_mode.is_decode()
or (
forward_batch.forward_mode.is_target_verify()
# Other MXFP8 backends may share mutable GEMM workspace.
and getattr(self.wq_b.quant_method, "mxfp8_dense_backend", None)
== Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL
)
)
)
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
is_unified_kv_fp8,
is_unified_kv_triton,
)
unified = is_unified_kv_triton()
unified_fp8_verify = (
unified
and is_unified_kv_fp8()
and forward_batch.forward_mode.is_target_verify()
)
# The v4 nm asm reader takes Q in the pool's own packed form, so fp8
# decode wants a contiguous fp8 buffer of exactly the local heads --
# q_padded below is a FlashMLA layout and buys nothing here. Verify runs
# that same reader over the ring, so it takes the same Q.
unified_fp8_decode = (
unified
and is_unified_kv_fp8()
and (forward_batch.forward_mode.is_decode_or_idle() or unified_fp8_verify)
)
# The 2-source prefill kernel wants the same packed Q plus this chunk's
# K in the pool's layout. Verify is not prefill here even though it takes
# the same branch below -- it reads rows the ring already holds, so it
# goes with decode above. Multi-stream picks a different prepare that has
# no unified arm at all, so it keeps the bf16 buffers it always had.
unified_fp8_prefill = (
unified
and is_unified_kv_fp8()
and not enable_multi_stream
and not forward_batch.forward_mode.is_decode_or_idle()
and not forward_batch.forward_mode.is_target_verify()
)
if unified_fp8_verify and not envs.SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY.get():
# The packed pair is produced by the fused norm+RoPE store; with that
# off the unfused arm hands the backend bf16 kv and the ring scatter
# dies on a dtype assert that says nothing about MTP.
raise NotImplementedError(
"fp8 two-pool unified_kv needs the fused verify store for "
"speculative decoding: set "
"SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY=1, or run with "
"SGLANG_DSV4_UNIFIED_KV_FP8=0."
)
if (
unified
and is_unified_kv_fp8()
and self.dsa_enable_prefill_cp
and dsa_use_prefill_cp(forward_batch)
and not forward_batch.forward_mode.is_decode_or_idle()
):
# The gather hands back bf16 kv in global token order *after*
# norm+RoPE, so packing would have to move ahead of it and re-derive
# RoPE from global-order positions. Whether the CP path has those
# ready is unverified, so refuse instead of packing the wrong order.
raise NotImplementedError(
"fp8 two-pool unified_kv does not support DSA prefill CP "
"(SGLANG_DSV4_UNIFIED_KV_FP8=1 with cp_size > 1)."
)
tp_slice, q_padded, q_out, q_rope = slice(None), None, None, None
k_nope, k_rope = None, None
if unified_fp8_decode or unified_fp8_prefill:
# width and dtype come off the pools themselves; the kernel reads Q
# with the kv row stride, so the two must not drift
kv_pool = get_token_to_kv_pool()
nope_pool = kv_pool.get_unified_kv(self.layer_id)
rope_pool = kv_pool.get_unified_kv_rope(self.layer_id)
q_out = nope_pool.new_empty(
(x.shape[0], self.n_local_heads, nope_pool.shape[-1])
)
q_rope = rope_pool.new_empty(
(x.shape[0], self.n_local_heads, rope_pool.shape[-1])
)
if unified_fp8_prefill or unified_fp8_verify:
k_nope = nope_pool.new_empty((x.shape[0], nope_pool.shape[-1]))
k_rope = rope_pool.new_empty((x.shape[0], rope_pool.shape[-1]))
kernel_num_heads = self.n_local_heads
else:
kernel_num_heads = self._kernel_num_heads(x.shape[0])
if kernel_num_heads != self.n_local_heads:
# Backends without an exact-head specialization retain the existing
# padded shape. attn_sink is sliced to this rank and padded to match.
if self.is_dsv41:
# V4.1 kernels read all padded heads, so the padding must be
# zero. The buffer is reused per layer; no consumer may keep it.
want = (x.shape[0], kernel_num_heads, self.head_dim)
meta = getattr(attn_backend, "forward_metadata", None)
q_padded = getattr(meta, "q_pad_buffer", None)
if (
q_padded is None
or tuple(q_padded.shape) != want
or q_padded.dtype != x.dtype
):
q_padded = x.new_zeros(*want)
if meta is not None:
try:
meta.q_pad_buffer = q_padded
except (AttributeError, TypeError):
pass
elif _is_gfx942_supported:
# Uninitialized padded TP heads inject NaN into attention on gfx942
# (fnuz), so zero-init there; other archs tolerate new_empty and skip
# the per-forward memset.
q_padded = x.new_zeros(x.shape[0], kernel_num_heads, self.head_dim)
else:
q_padded = x.new_empty(x.shape[0], kernel_num_heads, self.head_dim)
tp_slice = slice(0, self.n_local_heads)
q_out = q_padded[:, tp_slice, :]
attn_sink = self._local_attn_sink(kernel_num_heads)
if enable_multi_stream:
# Multi-stream path always fuses cache write into the K kernel,
# so the bf16 KV intermediate is gone.
if _is_hip:
q = self._forward_prepare_multi_stream_hip(
x,
positions,
forward_batch,
attn_backend,
q_out,
x_quant=x_quant,
)
elif _is_npu:
q = self._forward_prepare_multi_stream_npu(
x,
positions,
forward_batch,
attn_backend,
q_out,
x_quant=x_quant,
)
else:
q = self._forward_prepare_multi_stream(
x,
positions,
forward_batch,
attn_backend,
q_out,
x_quant=x_quant,
)
kv = None
elif low_ratio_multi_stream:
q = self._forward_prepare_low_ratio_multi_stream(
x,
positions,
forward_batch,
attn_backend,
q_out,
x_quant=x_quant,
)
kv = None
else:
q, kv = self._forward_prepare(
x,
positions,
forward_batch,
attn_backend,
q_out,
x_quant=x_quant,
q_rope_out=q_rope,
k_nope_out=k_nope,
k_rope_out=k_rope,
)
# save_kv_cache = kv is not None selects who writes the ring. When kv is
# None the store was already fused into _forward_prepare* (decode) or
# done inline, so the backend skips its own store_cache; pass `q` as a
# sentinel for the `k is v` assert (attention won't read it once
# save_kv_cache=False). When kv is not None (target-verify, or DSA-CP),
# _forward_prepare* deliberately left the store off and the backend does
# its normal causally-indexed store from attn_k = kv.
attn_k = kv if kv is not None else q
if unified:
# only the HIP radix backend takes these two; passing them always would
# leave non-ROCm depending on the **_ in its forward() to drop them, and
# no test on that side would notice if the **_ went away
rope_kwargs = {}
if q_rope is not None:
rope_kwargs["q_rope"] = q_rope
if k_rope is not None:
rope_kwargs["k_rope"] = k_rope
o = attn_backend.forward(
q=q_out if q_out is not None else q,
k=attn_k,
v=attn_k,
layer=self.attn_mqa,
forward_batch=forward_batch,
compress_ratio=self.compress_ratio,
attn_sink=attn_sink[: self.n_local_heads],
save_kv_cache=kv is not None,
**rope_kwargs,
)
else:
attn_q = q_padded if q_padded is not None else q
save_kv_cache = False
if forward_batch.forward_mode.is_extend() and is_in_breakable_cuda_graph():
o = attn_q.new_empty(
(*attn_q.shape[:-1], self.attn_mqa.v_head_dim),
)
bcg_deepseek_v4_attention_with_output(
attn_q,
attn_k,
o,
self.attn_mqa.layer_id,
self.compress_ratio,
attn_sink,
save_kv_cache,
)
else:
o = attn_backend.forward(
q=attn_q,
k=attn_k,
v=attn_k,
layer=self.attn_mqa,
forward_batch=forward_batch,
compress_ratio=self.compress_ratio,
attn_sink=attn_sink,
save_kv_cache=save_kv_cache,
)
o = o[:, tp_slice, :]
if (
self.wo_a_fp8
and _wo_a_fp8_mxscale_fused_invrope is not None
and not _is_npu
):
# ROCm gfx950 fused path: inverse-RoPE + per-token-group mxfp8 quant
# in one aiter kernel on the pre-view [T,H,Dh] output, then the a8w8
# mxscale absorb GEMM. Replaces the standalone inverse RoPE, the
# [T,G,D] view, and the quant inside the two-kernel fp8 path below.
G = self.n_local_groups
cos_c, sin_c = _freqs_cis_to_cos_sin(self.freqs_cis, o.dtype, o.device)
o = _wo_a_fp8_mxscale_fused_invrope(
o,
positions,
cos_c,
sin_c,
G,
self.wo_a.weight.view(G, self.o_lora_rank, -1),
self.wo_a.weight_scale_inv.data,
)
else:
fuse_mxfp8_quant = (
self.use_flashinfer_mxfp8_wo_b and not get_forward().sp_active
)
fuse_rope_wo_a = (
self.use_fused_wo_a and 0 < o.shape[0] <= _FUSED_WO_A_MAX_TOKENS
)
if _is_npu:
cos4, sin4 = self._get_npu_rope_position_cache(
forward_batch, positions, o.dtype, inverse=True
)
Dsv4NpuRoPE.apply_rotary_mul_inplace(
o,
None,
cos4,
sin4,
qk_nope_dim=self.qk_nope_head_dim,
)
elif not fuse_rope_wo_a:
# The fused path folds this in; it must not run twice.
fused_rope_inplace(
o[..., -self.qk_rope_head_dim :],
None,
self.freqs_cis,
positions=positions,
inverse=True,
)
o = o.view(o.shape[0], self.n_local_groups, -1)
if fuse_rope_wo_a:
wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
out = fused_rope_wo_a_bf16(
o,
wo_a,
torch.view_as_real(self.freqs_cis).flatten(1),
positions,
out_mxfp8=fuse_mxfp8_quant,
)
o = Mxfp8SwizzledInput(*out) if fuse_mxfp8_quant else out[0]
elif self.use_npu_arch35_mxfp8_wo_a:
o, o_scale = torch_npu.npu_dynamic_mx_quant(
o, dst_type=torch.float8_e4m3fn
)
o = torch_npu.npu_transpose_quant_batchmatmul(
o,
self.wo_a.weight,
dtype=torch.bfloat16,
bias=None,
group_sizes=(0, 0, 32),
x1_scale=o_scale.view(torch.float8_e8m0fnu),
x2_scale=self.wo_a.weight_scale_inv.view(torch.float8_e8m0fnu),
perm_x1=(1, 0, 2),
perm_x2=(0, 1, 2),
perm_y=(1, 0, 2),
)
elif self.wo_a_fp8 and _wo_a_fp8_mxscale is not None:
# ROCm gfx950: same fp8 absorb GEMM as the DeepGEMM path below,
# but through aiter's e8m0 block-scale batched GEMM. The
# activation is quantized per token-group inside the helper.
T, G, D = o.shape
o = _wo_a_fp8_mxscale(
o,
self.wo_a.weight.view(G, self.o_lora_rank, D),
self.wo_a.weight_scale_inv.data,
)
elif self.wo_a_fp8:
import deep_gemm
T, G, D = o.shape
R = self.o_lora_rank
if _FP8_WO_A_UE8M0:
# Blackwell (including SM120): UE8M0 scales via the dedicated
# JIT kernel.
o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o)
recipe = (1, 1, 128)
else:
# sm90 (Hopper): fp32 scales.
o_fp8, o_s = sglang_per_token_group_quant_fp8(
o.reshape(T * G, D).contiguous(),
group_size=128,
scale_ue8m0=False,
)
o_fp8 = o_fp8.view(T, G, D)
o_s = o_s.view(T, G, -1)
recipe = (1, 128, 128)
output = torch.empty(T, G, R, device=o.device, dtype=torch.bfloat16)
deep_gemm.fp8_einsum(
"bhr,hdr->bhd",
(o_fp8, o_s),
(
self.wo_a.weight.view(G, R, D),
self.wo_a.weight_scale_inv.data,
),
output,
recipe=recipe,
)
o = output
else:
wo_a_weight = getattr(self.wo_a, "weight", None)
if wo_a_weight is not None:
if (
_NPU_BF16_WO_A_GEMM
and forward_batch.forward_mode.is_decode()
and self.n_local_groups == 1
and o.dtype == wo_a_weight.dtype == torch.bfloat16
and wo_a_weight.is_contiguous()
):
# One local group needs no grouped contraction; linear
# avoids materializing a transpose of the BF16 weight.
o = F.linear(o, wo_a_weight)
else:
wo_a = wo_a_weight.view(
self.n_local_groups, self.o_lora_rank, -1
)
o = _apply_wo_a_bf16_matmul(
o,
wo_a,
is_decode=forward_batch.forward_mode.is_decode(),
is_target_verify=forward_batch.forward_mode.is_target_verify(),
is_prefill=forward_batch.forward_mode.is_extend_without_speculative(),
fast_path=self.is_dsv41,
fuse_mxfp8_quant=fuse_mxfp8_quant,
)
else:
o = _apply_gguf_grouped_wo_a(
o,
self.wo_a.qweight,
self.wo_a.qweight_type.weight_type,
self.o_lora_rank,
)
from sglang.srt.layers.moe.mhc_post_fusion import current_mhc_post_fusion
mhc = current_mhc_post_fusion()
o, _ = self.wo_b(
o if isinstance(o, Mxfp8SwizzledInput) else o.flatten(1),
skip_all_reduce=mhc is not None,
)
if mhc is not None and mhc.overlap_only:
mhc.start_stats_before_all_reduce()
o = attn_tp_all_reduce(o)
elif mhc is not None:
from sglang.kernels.ops.communication.all_reduce_mhc import (
all_reduce_mhc_norm,
)
mhc.materialize_stats()
if mhc.stats_stream is not None:
torch.cuda.current_stream().wait_stream(mhc.stats_stream)
if mhc.combine_only:
from sglang.kernels.ops.communication.all_reduce_mhc_combine import (
all_reduce_mhc_combine,
)
o, mhc.output, mhc.combined = all_reduce_mhc_combine(
o,
mhc.residual,
mhc.post,
mhc.comb,
mhc.pre,
world_size=self.attn_tp_size,
)
else:
o, mhc.output, mhc.normalized = all_reduce_mhc_norm(
o,
mhc.residual,
mhc.post,
mhc.comb,
mhc.pre,
mhc.norm_weight,
mhc.norm_eps,
world_size=self.attn_tp_size,
)
if self.attn_tp_size > 1 and self.attn_tp_size < get_parallel().tp_size:
o = attn_tp_all_reduce(o)
return o
# ---- TBO op decomposition (prefill two-batch-overlap) ----
def op_attn(self, state):
"""Run the attention forward as a single TBO op.
Consumes the post-input-norm hidden states produced by
``DeepseekV4DecoderLayer.op_mhc_prepare_attn`` and stores the attention
output for ``op_mhc_post_attn_pre_mlp``.
"""
state.hidden_states_after_attn = self.forward(
x=state.pop("hidden_states_after_input_norm"),
positions=state.positions,
forward_batch=state.forward_batch,
x_quant=state.pop("attn_x_quant"),
)
@contextmanager
def _every_row_routed(forward_batch: ForwardBatch, num_rows: int):
# Under CP the real rows are not a prefix, so every row must be routed.
saved = (
forward_batch.num_token_non_padded,
forward_batch.global_num_token_non_padded_cpu,
)
if saved[0] is not None:
forward_batch.num_token_non_padded = torch.full_like(saved[0], num_rows)
forward_batch.global_num_token_non_padded_cpu = num_rows
try:
yield
finally:
(
forward_batch.num_token_non_padded,
forward_batch.global_num_token_non_padded_cpu,
) = saved
class DeepseekV4DecoderLayer(nn.Module):
def __init__(
self,
config: DeepSeekV4Config,
layer_id: int,
quant_config: Optional[QuantizationConfig] = None,
moe_quant_config_override: Optional[QuantizationConfig] = None,
is_nextn: bool = False,
prefix: str = "",
alt_streams: Optional[List[torch.cuda.Stream]] = None,
compress_ratio_override: Optional[int] = None,
engram_layout: Optional[EngramLayout] = None,
hc_stats_stream: Optional[torch.cuda.Stream] = None,
moe_routed_quant_stream: Optional[torch.cuda.Stream] = None,
) -> None:
super().__init__()
self.hc_stats_stream = hc_stats_stream
self.config = config
self.hidden_size = config.hidden_size
self.layer_id = layer_id
self.self_attn = self._build_self_attn(
config=config,
layer_id=layer_id,
quant_config=quant_config,
prefix=add_prefix("self_attn", prefix),
alt_streams=alt_streams,
compress_ratio_override=compress_ratio_override,
)
moe_alt_stream = (
alt_streams[0]
if (
alt_streams is not None
and (
_is_cuda
or envs.SGLANG_ROCM_USE_MULTI_STREAM.get()
or envs.SGLANG_NPU_USE_MULTI_STREAM.get()
)
)
else None
)
self.mlp = deepseek_v2.DeepseekV2MoE(
config=config,
quant_config=moe_quant_config_override or quant_config,
prefix=add_prefix("mlp", prefix),
layer_id=self.layer_id,
alt_stream=moe_alt_stream,
routed_quant_stream=moe_routed_quant_stream,
is_nextn=is_nextn,
is_deepseek_v4=True,
vl_correction_bias=config.model_type == "deepseek_v41"
and config.vision_n_layers > 0
and not getattr(config, "language_model_only", False),
)
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
self.hc_mult = hc_mult = config.hc_mult
self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
self.hc_eps = config.hc_eps
(
self.hc_attn_fn,
self.hc_ffn_fn,
self.hc_attn_base,
self.hc_ffn_base,
self.hc_attn_scale,
self.hc_ffn_scale,
) = make_hc_mixing_params(hc_mult, config.hidden_size)
self.rms_norm_eps = config.rms_norm_eps
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
self.use_fused_mhc_post_pre = (
is_cross_layer_mhc_fusion_enabled() or _is_fused_mhc_post_pre_enabled_xpu()
)
# The fused post+pre boundary bakes in the same-sublayer pre-mix.
self.hc_pre_from_prev_sublayer = config.hc_pre_from_prev_sublayer
if self.hc_pre_from_prev_sublayer:
self.use_fused_mhc_post_pre = False
self.engram = None
if engram_layout is not None and layer_id in engram_layout.layer_ids:
self.engram = Engram(
config,
layer_id,
engram_layout,
quant_config=quant_config,
prefix=add_prefix("engram", prefix),
)
self._input_layernorm_weight_bf16 = None
self._post_attention_layernorm_weight_bf16 = None
def _build_self_attn(
self,
*,
config: DeepSeekV4Config,
layer_id: int,
quant_config: Optional[QuantizationConfig],
prefix: str,
alt_streams: Optional[List[torch.cuda.Stream]],
compress_ratio_override: Optional[int],
) -> nn.Module:
return MQALayer(
config=config,
layer_id=layer_id,
quant_config=quant_config,
prefix=prefix,
alt_streams=alt_streams,
compress_ratio_override=compress_ratio_override,
)
def refresh_mhc_norm_weight_cache(self):
# Cache bf16 norm weights so the fused path does not allocate/cast per forward.
self._input_layernorm_weight_bf16 = (
self.input_layernorm.weight.data.bfloat16().contiguous()
)
self._post_attention_layernorm_weight_bf16 = (
self.post_attention_layernorm.weight.data.bfloat16().contiguous()
)
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
# The original FP32 parameters stay intact for small rows and invariant mode.
self._hc_attn_tf32_parts = self._hc_ffn_tf32_parts = None
self._hc_attn_bf16_parts = self._hc_ffn_bf16_parts = None
if (
self.hc_pre_from_prev_sublayer
and get_platform().is_sm100
and self.hc_attn_fn.shape == (24, 20480)
and envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get()
and getattr(self.config, "model_type", None) == "deepseek_v41"
and not is_batch_invariant_mode_enabled()
):
from sglang.kernels.ops.layernorm.mhc import (
split_tf32_hc_weight,
)
from sglang.srt.layers.deep_gemm_wrapper.configurer import (
ENABLE_JIT_DEEPGEMM,
)
if ENABLE_JIT_DEEPGEMM:
import deep_gemm
if not callable(getattr(deep_gemm, "tf32_hc_prenorm_gemm", None)):
return
self._hc_attn_tf32_parts = split_tf32_hc_weight(self.hc_attn_fn.data)
self._hc_ffn_tf32_parts = split_tf32_hc_weight(self.hc_ffn_fn.data)
if (
getattr(getattr(self, "config", None), "model_type", None)
== "deepseek_v41"
):
from sglang.kernels.ops.layernorm.mhc import (
split_bf16_hc_weight,
)
self._hc_attn_bf16_parts = split_bf16_hc_weight(
self.hc_attn_fn.data
)
self._hc_ffn_bf16_parts = split_bf16_hc_weight(self.hc_ffn_fn.data)
def hc_pre(
self,
x: torch.Tensor,
hc_fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
norm: Optional[nn.Module] = None,
forward_batch: Optional[ForwardBatch] = None,
):
"""If *norm* is given and the TileLang path is active, the returned
hidden_states are already post-norm (the norm is fused into the kernel)."""
@compile_in_capture_mode
def hc_pre_torch_impl(x, hc_fn):
x_flat = x.flatten(1).float()
rsqrt = torch.rsqrt(
x_flat.square().mean(-1, keepdim=True) + self.rms_norm_eps
)
mixes = (F.linear(x_flat, hc_fn) * rsqrt).unsqueeze(1)
return x_flat, mixes
shape, dtype = x.size(), x.dtype
if _is_npu:
return _get_mhc_ops().npu_hc_pre(
x,
hc_fn,
hc_scale,
hc_base,
hc_mult=self.hc_mult,
hc_sinkhorn_iters=self.hc_sinkhorn_iters,
rms_norm_eps=self.rms_norm_eps,
hc_eps=self.hc_eps,
forward_batch=forward_batch,
)
if x.shape[0] == 0:
y = torch.empty((0, shape[-1]), dtype=dtype, device=x.device)
post = torch.empty((0, self.hc_mult), dtype=torch.float32, device=x.device)
comb = torch.empty(
(0, self.hc_mult, self.hc_mult), dtype=torch.float32, device=x.device
)
return y, post, comb, False
if _is_xpu:
norm_kwargs = {}
if norm is not None:
norm_kwargs["norm_weight"] = norm.weight.data
norm_kwargs["norm_eps"] = norm.variance_epsilon
post, comb, y = _get_mhc_ops().mhc_pre(
residual=x,
fn=hc_fn,
hc_scale=hc_scale,
hc_base=hc_base,
rms_eps=self.rms_norm_eps,
hc_pre_eps=self.hc_eps,
hc_sinkhorn_eps=self.hc_eps,
hc_post_mult_value=_MHC_POST_MULT_VALUE,
sinkhorn_repeat=self.hc_sinkhorn_iters,
**norm_kwargs,
)
return y, post, comb, norm is not None
if envs.SGLANG_OPT_USE_FLASHINFER_MHC.get():
y, post, comb = _flashinfer_hc_pre(
x,
hc_fn,
hc_scale,
hc_base,
rms_eps=self.rms_norm_eps,
hc_eps=self.hc_eps,
sinkhorn_iters=self.hc_sinkhorn_iters,
)
return y, post, comb, False
if envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get():
from sglang.kernels.ops.layernorm.mhc import mhc_pre
norm_kwargs = {}
if norm is not None:
norm_kwargs["norm_weight"] = norm.weight.data
norm_kwargs["norm_eps"] = norm.variance_epsilon
post, comb, y = mhc_pre(
residual=x,
fn=hc_fn,
hc_scale=hc_scale,
hc_base=hc_base,
rms_eps=self.rms_norm_eps,
hc_pre_eps=self.hc_eps,
hc_sinkhorn_eps=self.hc_eps,
hc_post_mult_value=_MHC_POST_MULT_VALUE,
sinkhorn_repeat=self.hc_sinkhorn_iters,
**norm_kwargs,
)
return y, post.squeeze(-1), comb, norm is not None
if _is_hip:
from aiter.ops.mhc import mhc_pre
post, comb, y = mhc_pre(
residual=x,
fn=hc_fn,
hc_scale=hc_scale,
hc_base=hc_base,
rms_eps=self.rms_norm_eps,
hc_pre_eps=self.hc_eps,
hc_sinkhorn_eps=self.hc_eps,
hc_post_mult_value=_MHC_POST_MULT_VALUE,
sinkhorn_repeat=self.hc_sinkhorn_iters,
)
return y, post.squeeze(-1), comb, False
# The deepgemm tf32 gemm wins at large M (prefill) but its fixed
# dispatch cost dominates at small M (decode): dispatch by token count.
if (
envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get()
and x.shape[0] >= _HC_PRENORM_DEEPGEMM_MIN_TOKENS
):
from sglang.srt.layers.deep_gemm_wrapper.entrypoint import (
tf32_hc_prenorm_gemm,
)
x_flat = x.flatten(1).bfloat16()
m, k = x_flat.shape
mix_hc = hc_fn.size(0)
d_out = torch.empty((m, mix_hc), dtype=torch.float, device=x.device)
s_out = torch.empty((m,), dtype=torch.float, device=x.device)
tf32_hc_prenorm_gemm(
x_flat, hc_fn.float().contiguous(), d_out, s_out, num_splits=None
)
rsqrt = torch.rsqrt(s_out / k + self.rms_norm_eps)
mixes = (d_out * rsqrt.unsqueeze(1)).unsqueeze(1)
else:
x_flat, mixes = hc_pre_torch_impl(x, hc_fn)
pre, post, comb = _get_mhc_ops().hc_split_sinkhorn(
mixes,
hc_scale,
hc_base,
self.hc_mult,
self.hc_sinkhorn_iters,
self.hc_eps,
)
from sglang.kernels.ops.layernorm.mhc import hc_combine
# y is the post-norm activation fed into the MoE. Allocate it in the
# symmetric memory pool so the downstream all-reduce uses the low-latency
# NCCL symmetric path: the Triton inplace MoE runner writes the expert
# output back into this buffer, so a symmetric input yields a symmetric
# all-reduce input. Gated by is_allocation_symmetric() (mirrors the
# TileLang path in _mhc_pre_impl / mhc_fused_post_pre).
with use_symmetric_memory(
get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
y = hc_combine(x_flat, pre.squeeze(1), self.hc_mult, dtype)
return y, post.squeeze(1), comb.squeeze(1), False
def hc_post(
self,
x: torch.Tensor,
residual: torch.Tensor,
post: torch.Tensor,
comb: torch.Tensor,
):
if x.shape[0] == 0:
return torch.empty(
(0, self.hc_mult, x.shape[-1]), dtype=x.dtype, device=x.device
)
if _is_npu:
if not is_npu_arch35():
return torch.ops.custom.npu_hc_post(x, residual, post, comb)
# The A5 build of npu_hc_post is batched — it requires a leading
# batch axis on every operand.
return torch.ops.custom.npu_hc_post(
x.unsqueeze(0),
residual.unsqueeze(0),
post.unsqueeze(0),
comb.unsqueeze(0),
).squeeze(0)
if _is_xpu:
return _get_mhc_ops().mhc_post(x, residual, post, comb)
if (
_is_cuda
and get_platform().is_blackwell
and self.hc_pre_from_prev_sublayer
and self.hc_mult == 4
and x.shape[1] == 5120
and x.shape[0] <= 384
and x.dtype == residual.dtype == torch.bfloat16
and post.dtype == comb.dtype == torch.float32
and all(t.is_contiguous() for t in (x, residual, post, comb))
):
return mhc_post_split_h(x, residual, post, comb)
if envs.SGLANG_OPT_USE_FLASHINFER_MHC.get():
from flashinfer.mhc import mhc_post
return mhc_post(x, residual, post, comb)
if envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get():
if (
self.hc_pre_from_prev_sublayer
and get_platform().is_sm90
and x.is_cuda
and 1 <= x.shape[0] <= 64
and x.shape[1] == 5120
and residual.shape == (x.shape[0], 4, 5120)
and x.dtype == residual.dtype == torch.bfloat16
and post.dtype == comb.dtype == torch.float32
and post.numel() == x.shape[0] * 4
and comb.shape == (x.shape[0], 4, 4)
and all(t.is_contiguous() for t in (x, residual, post, comb))
):
return mhc_post_split_h(x, residual, post, comb)
from sglang.kernels.ops.layernorm.mhc import mhc_post
return mhc_post(x, residual, post, comb)
elif _is_hip:
from aiter.ops.mhc import mhc_post
result = torch.empty_like(residual)
mhc_post(result, x, residual, post, comb)
return result
assert residual.shape == (x.shape[0], self.hc_mult, x.shape[-1])
assert post.shape == (x.shape[0], self.hc_mult)
assert comb.shape == (x.shape[0], self.hc_mult, self.hc_mult)
@compile_in_capture_mode
def hc_post_torch_impl(x, residual, post, comb):
return (
post.unsqueeze(-1) * x.unsqueeze(1)
+ (comb.unsqueeze(-1) * residual.unsqueeze(2)).sum(dim=1)
).type_as(x)
return hc_post_torch_impl(x, residual, post, comb)
def forward(
self,
positions: torch.tensor,
hidden_states: torch.Tensor,
input_ids: torch.Tensor,
forward_batch: ForwardBatch,
input_ids_global: torch.Tensor,
prev_residual: Optional[torch.Tensor] = None,
prev_post: Optional[torch.Tensor] = None,
prev_comb: Optional[torch.Tensor] = None,
) -> Tuple[
torch.Tensor,
Optional[torch.Tensor],
Optional[torch.Tensor],
Optional[torch.Tensor],
]:
use_fused = self.use_fused_mhc_post_pre
if prev_residual is not None and use_fused:
# Dispatch cascade: aiter HIP (gfx95) -> Triton (gfx95 small-batch
# <=64 tokens, or gfx1250 all sizes) -> TileLang -> None.
input_norm_weight = (
self._input_layernorm_weight_bf16
if self._input_layernorm_weight_bf16 is not None
else self.input_layernorm.weight.data
)
fused = apply_mhc_post_pre_boundary(
hidden_states,
prev_residual,
prev_post,
prev_comb,
self.hc_attn_fn,
self.hc_attn_scale,
self.hc_attn_base,
self.hc_mult,
self.rms_norm_eps,
self.hc_eps,
_MHC_POST_MULT_VALUE,
self.hc_sinkhorn_iters,
input_norm_weight,
self.input_layernorm.variance_epsilon,
fn_transpose=False,
)
if fused is not None:
residual, hidden_states, post, comb, norm_fused = fused
if not norm_fused:
# Triton fused post+pre (gfx95 small-batch or gfx1250) returns
# norm_fused=False — the input layernorm is NOT folded.
# gfx95 takes the fp8-quant path; gfx1250 takes plain layernorm.
if _use_aiter and _is_gfx95_supported:
x_quant, hidden_states = _fused_rmsnorm_fp8_quant(
hidden_states,
self.input_layernorm.weight,
self.rms_norm_eps,
)
else:
hidden_states = self.input_layernorm(hidden_states)
x_quant = None
else:
x_quant = None
else:
hidden_states = self.hc_post(
hidden_states, prev_residual, prev_post, prev_comb
)
residual = hidden_states
hidden_states, post, comb, norm_fused = self.hc_pre(
hidden_states,
self.hc_attn_fn,
self.hc_attn_scale,
self.hc_attn_base,
norm=self.input_layernorm,
forward_batch=forward_batch,
)
if not norm_fused:
if _use_aiter and _is_gfx95_supported:
x_quant, hidden_states = _fused_rmsnorm_fp8_quant(
hidden_states,
self.input_layernorm.weight,
self.rms_norm_eps,
)
else:
hidden_states = self.input_layernorm(hidden_states)
x_quant = None
else:
x_quant = None
else:
residual = hidden_states
hidden_states, post, comb, norm_fused = self.hc_pre(
hidden_states,
self.hc_attn_fn,
self.hc_attn_scale,
self.hc_attn_base,
norm=self.input_layernorm,
forward_batch=forward_batch,
)
if not norm_fused:
if _use_aiter and _is_gfx95_supported:
x_quant, hidden_states = _fused_rmsnorm_fp8_quant(
hidden_states,
self.input_layernorm.weight,
self.rms_norm_eps,
)
else:
hidden_states = self.input_layernorm(hidden_states)
x_quant = None
else:
x_quant = None
with self.self_attn.maybe_use_decode_attn_tp(forward_batch):
hidden_states = self.self_attn(
x=hidden_states,
positions=positions,
forward_batch=forward_batch,
x_quant=x_quant,
)
if use_fused:
post_attn_norm_weight = (
self._post_attention_layernorm_weight_bf16
if self._post_attention_layernorm_weight_bf16 is not None
else self.post_attention_layernorm.weight.data
)
fused = apply_mhc_post_pre_boundary(
hidden_states,
residual,
post,
comb,
self.hc_ffn_fn,
self.hc_ffn_scale,
self.hc_ffn_base,
self.hc_mult,
self.rms_norm_eps,
self.hc_eps,
_MHC_POST_MULT_VALUE,
self.hc_sinkhorn_iters,
post_attn_norm_weight,
self.post_attention_layernorm.variance_epsilon,
fn_transpose=True,
)
if fused is not None:
residual, hidden_states, post, comb, norm_fused = fused
if not norm_fused:
hidden_states = self.post_attention_layernorm(hidden_states)
else:
hidden_states = self.hc_post(hidden_states, residual, post, comb)
residual = hidden_states
hidden_states, post, comb, norm_fused = self.hc_pre(
hidden_states,
self.hc_ffn_fn,
self.hc_ffn_scale,
self.hc_ffn_base,
norm=self.post_attention_layernorm,
forward_batch=forward_batch,
)
if not norm_fused:
hidden_states = self.post_attention_layernorm(hidden_states)
else:
hidden_states = self.hc_post(hidden_states, residual, post, comb)
residual = hidden_states
hidden_states, post, comb, norm_fused = self.hc_pre(
hidden_states,
self.hc_ffn_fn,
self.hc_ffn_scale,
self.hc_ffn_base,
norm=self.post_attention_layernorm,
forward_batch=forward_batch,
)
if not norm_fused:
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self._run_moe_ffn_dp_sync(
hidden_states,
forward_batch,
input_ids=input_ids,
input_ids_global=input_ids_global,
)
if not use_fused:
hidden_states = self.hc_post(hidden_states, residual, post, comb)
return hidden_states, None, None, None
# Return the deferred FFN hc_post state; the next layer consumes it with
# cross-layer fusion, and the final layer is completed in DeepseekV4Model.
return hidden_states, residual, post, comb
def _hc_combine(
self,
x: torch.Tensor,
apply_pre: Optional[torch.Tensor],
norm: RMSNorm,
stats_stream: Optional[torch.cuda.Stream] = None,
quantized: Optional[list] = None,
normalized: Optional[torch.Tensor] = None,
precomputed: Optional[tuple] = None,
combined: Optional[torch.Tensor] = None,
) -> torch.Tensor:
from sglang.kernels.ops.layernorm.mhc import hc_combine
quantize = quantized is not None
x_flat = x.flatten(1)
tiny = 0 < x.shape[0] <= 8
if stats_stream is not None and not tiny:
stats_stream.wait_stream(torch.cuda.current_stream())
def combine_and_norm():
if precomputed is not None:
assert quantized is not None
quantized.append(precomputed[1])
return precomputed[0]
if normalized is not None:
# Prefill projections still quantize the BF16 input themselves;
# the optional fused-quantization list stays empty for this case.
assert not quantize or 4096 <= x.shape[0] <= 65536
return normalized
if combined is not None:
if (
4096 <= combined.shape[0] <= 65536
and norm.weight.dtype == torch.bfloat16
and not norm.cast_x_before_out_mul
and norm.variance_size_override is None
):
from sglang.kernels.ops.layernorm.mhc_post_combine import (
hc_norm_prefill,
)
return hc_norm_prefill(combined, norm.weight, norm.variance_epsilon)
return norm(combined)
if apply_pre is None:
return norm(x[:, 0, :].contiguous())
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
if (
x.is_cuda
and get_platform().is_blackwell
and (
0 < x.shape[0] <= 96
or (
self.config.model_type == "deepseek_v41"
and 4096 <= x.shape[0] <= 65536
)
)
and self.hc_mult == 4
and x_flat.shape[1] == 20480
and x.dtype == norm.weight.dtype == torch.bfloat16
and apply_pre.stride(1) == 1
and not norm.cast_x_before_out_mul
and norm.variance_size_override is None
and not is_batch_invariant_mode_enabled()
):
# The fused scale writer supports the small decode/verify tile only.
if quantize and x.shape[0] <= 8:
from sglang.kernels.ops.layernorm.hc_combine_norm import (
hc_combine_norm_mxfp8,
)
y, y_q, y_sf = hc_combine_norm_mxfp8(
x_flat, apply_pre, norm.weight, norm.variance_epsilon
)
quantized.append(Mxfp8SwizzledInput(y_q, y_sf))
return y
from sglang.kernels.ops.layernorm.hc_combine_norm import hc_combine_norm
return hc_combine_norm(
x_flat, apply_pre, norm.weight, norm.variance_epsilon
)
return norm(hc_combine(x_flat, apply_pre, self.hc_mult, x.dtype))
y = combine_and_norm()
if stats_stream is not None and tiny:
stats_stream.wait_stream(torch.cuda.current_stream())
return y
def _hc_mix_stats(
self,
x: torch.Tensor,
hc_fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
stats_stream: Optional[torch.cuda.Stream] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
from sglang.kernels.ops.layernorm.mhc import hc_mix_stats, hc_mix_stats_sinkhorn
x_flat = x.flatten(1)
if (
x.is_cuda
and torch.version.cuda is not None
and (
get_platform().is_blackwell
or (get_platform().is_sm90 and x.shape[0] == 1)
)
and x.dtype == torch.bfloat16
):
# Fusing the split-K reduction with sinkhorn keeps it batch-invariant.
main_stream = torch.cuda.current_stream()
if stats_stream is not None:
x.record_stream(stats_stream)
with (
torch.cuda.stream(stats_stream)
if stats_stream is not None
else nullcontext()
):
from sglang.srt.batch_invariant_ops import (
is_batch_invariant_mode_enabled,
)
parts = bf16_parts = None
if (
x_flat.shape[0] >= 128
and x_flat.is_contiguous()
and get_platform().is_sm100
and envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get()
and not is_batch_invariant_mode_enabled()
):
if hc_fn is self.hc_attn_fn:
parts = getattr(self, "_hc_attn_tf32_parts", None)
bf16_parts = getattr(self, "_hc_attn_bf16_parts", None)
elif hc_fn is self.hc_ffn_fn:
parts = getattr(self, "_hc_ffn_tf32_parts", None)
bf16_parts = getattr(self, "_hc_ffn_bf16_parts", None)
if bf16_parts is not None and 4096 <= x_flat.shape[0] <= 65536:
from sglang.kernels.ops.layernorm.mhc import (
hc_mix_stats_sinkhorn_bf16x3,
)
pre, post, comb = hc_mix_stats_sinkhorn_bf16x3(
x_flat,
bf16_parts,
hc_scale,
hc_base,
self.hc_sinkhorn_iters,
self.rms_norm_eps,
self.hc_eps,
)
elif parts is not None:
from sglang.kernels.ops.layernorm.mhc import (
hc_mix_stats_sinkhorn_deepgemm,
)
pre, post, comb = hc_mix_stats_sinkhorn_deepgemm(
x_flat,
parts,
hc_scale,
hc_base,
self.hc_sinkhorn_iters,
self.rms_norm_eps,
self.hc_eps,
)
else:
pre, post, comb = hc_mix_stats_sinkhorn(
x_flat,
hc_fn,
hc_scale,
hc_base,
self.hc_mult,
self.hc_sinkhorn_iters,
self.rms_norm_eps,
self.hc_eps,
)
if stats_stream is not None:
# Allocated on the side stream, read on the main stream after the join.
for coefficient in (pre, post, comb):
coefficient.record_stream(main_stream)
return pre, post, comb
if x.is_cuda and torch.version.cuda is not None:
# cuBLAS/torch reductions can change order with num_tokens; this kernel
# keeps the mixing and RMS reductions batch-invariant.
mixes = hc_mix_stats(x_flat, hc_fn, self.rms_norm_eps).unsqueeze(1)
else:
x_flat = x_flat.float()
rsqrt = torch.rsqrt(
x_flat.square().mean(-1, keepdim=True) + self.rms_norm_eps
)
mixes = (F.linear(x_flat, hc_fn) * rsqrt).unsqueeze(1)
pre, post, comb = _get_mhc_ops().hc_split_sinkhorn(
mixes,
hc_scale,
hc_base,
self.hc_mult,
self.hc_sinkhorn_iters,
self.hc_eps,
)
return pre.squeeze(1), post.squeeze(1), comb.squeeze(1)
def _hc_mix_and_combine(
self,
x,
hc_fn,
hc_scale,
hc_base,
apply_pre,
norm,
stats_stream=None,
quantized=None,
normalized=None,
precomputed=None,
):
y = DeepseekV4DecoderLayer._hc_combine(
self, x, apply_pre, norm, stats_stream, quantized, normalized, precomputed
)
return (
y,
*DeepseekV4DecoderLayer._hc_mix_stats(
self, x, hc_fn, hc_scale, hc_base, stats_stream
),
)
def _get_hc_stats_stream(self, hidden_states, forward_batch):
# Prefill stats share one model-wide stream. Start them immediately
# before the sublayer's all-reduce, after its compute has completed.
if (
self.config.model_type == "deepseek_v41"
and hidden_states.is_cuda
and get_platform().is_blackwell
and forward_batch.forward_mode.is_extend_without_speculative()
and 4096 <= hidden_states.shape[0] <= 65536
and get_parallel().attn_dp_size == 1
and not get_forward().sp_active
and not self.dsa_enable_prefill_cp
):
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
if not is_batch_invariant_mode_enabled():
return self.hc_stats_stream
# Verify batches can also compute coefficients beside the
# sublayer; each branch joins before hc_post reads those coefficients.
return (
self.hc_stats_stream
if (
forward_batch.forward_mode.is_decode()
or (
forward_batch.forward_mode.is_target_verify()
and hidden_states.shape[0] > 0
)
)
and (not get_platform().is_sm90 or hidden_states.shape[0] == 1)
else None
)
def _hc_post_with_combine(
self, x, residual, post, comb, pre, forward_batch, norm=None
):
"""Return updated HC streams and optional combined/normalized inputs."""
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
if (
self.config.model_type == "deepseek_v41"
and x.is_cuda
and get_platform().is_blackwell
and (
(
128 <= x.shape[0] <= 384
and (
forward_batch.forward_mode.is_decode()
or forward_batch.forward_mode.is_target_verify()
)
)
or (
4096 <= x.shape[0] <= 65536
and forward_batch.forward_mode.is_extend_without_speculative()
and envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get()
and not envs.SGLANG_OPT_USE_FLASHINFER_MHC.get()
)
)
and x.shape[1] == 5120
and self.hc_mult == 4
and x.dtype == residual.dtype == torch.bfloat16
and post.dtype == comb.dtype == pre.dtype == torch.float32
and all(t.is_contiguous() for t in (x, residual, post, comb, pre))
and get_parallel().attn_dp_size == 1
and not get_forward().sp_active
and not self.dsa_enable_prefill_cp
and not is_batch_invariant_mode_enabled()
):
if (
x.shape[0] >= 4096
and norm is not None
and not norm.cast_x_before_out_mul
and norm.variance_size_override is None
and norm.weight.dtype == torch.bfloat16
and norm.weight.shape == (5120,)
and norm.weight.is_contiguous()
and all(t.data_ptr() % 16 == 0 for t in (x, residual, norm.weight))
):
from sglang.kernels.ops.layernorm.mhc_post_combine_norm_prefill import (
mhc_post_combine_norm_prefill,
)
updated, normalized = mhc_post_combine_norm_prefill(
x, residual, post, comb, pre, norm.weight, norm.variance_epsilon
)
return updated, None, normalized
from sglang.kernels.ops.layernorm.mhc_post_combine import mhc_post_combine
updated, combined = mhc_post_combine(x, residual, post, comb, pre)
return updated, combined, None
return self.hc_post(x, residual, post, comb), None, None
def forward_hc_pre_from_prev(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
input_ids: torch.Tensor,
forward_batch: ForwardBatch,
input_ids_global: torch.Tensor,
prev_pre: Optional[torch.Tensor],
precomputed_attn: Optional[tuple] = None,
next_norm: Optional[RMSNorm] = None,
next_input: Optional[list] = None,
combined_attn: Optional[torch.Tensor] = None,
normalized_attn: Optional[torch.Tensor] = None,
next_combined: Optional[list] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Layer forward where each sublayer consumes the previous sublayer's
pre-mix. Returns (hidden_states, ffn_pre)."""
from functools import partial
from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
from sglang.srt.layers.moe.mhc_post_fusion import (
MhcPostFusion,
use_mhc_post_fusion,
)
stats_stream = self._get_hc_stats_stream(hidden_states, forward_batch)
residual = hidden_states
attn_quantized: Optional[list] = (
[] if self.self_attn.accepts_mxfp8_swizzled_input() else None
)
attn_stats = partial(
self._hc_mix_stats,
hidden_states,
self.hc_attn_fn,
self.hc_attn_scale,
self.hc_attn_base,
stats_stream,
)
x = self._hc_combine(
hidden_states,
apply_pre=prev_pre,
norm=self.input_layernorm,
stats_stream=stats_stream,
quantized=attn_quantized,
precomputed=precomputed_attn,
combined=combined_attn,
normalized=normalized_attn,
)
prefill_overlap = (
stats_stream is not None
and forward_batch.forward_mode.is_extend_without_speculative()
)
medium_verify = 128 <= x.shape[0] <= 384 and (
forward_batch.forward_mode.is_decode()
or forward_batch.forward_mode.is_target_verify()
)
attn_mhc = None
if (
self.config.model_type == "deepseek_v41"
and x.is_cuda
and get_platform().is_blackwell
and (0 < x.shape[0] <= 8 or medium_verify)
and x.shape[1] == 5120
and self.hc_mult == 4
and x.dtype == residual.dtype == torch.bfloat16
and residual.is_contiguous()
and get_parallel().attn_dp_size == 1
and get_parallel().tp_size == self.self_attn.attn_tp_size == 4
and self.self_attn.wo_b.reduce_results
and not get_forward().sp_active
and not self.dsa_enable_prefill_cp
and not self.post_attention_layernorm.cast_x_before_out_mul
and self.post_attention_layernorm.variance_size_override is None
and self.post_attention_layernorm.weight.dtype == torch.bfloat16
):
from sglang.kernels.ops.communication.all_reduce_fusion import (
get_registered_comm,
)
comm_ready = get_registered_comm(self.self_attn.attn_tp_size) is not None
if medium_verify and not is_batch_invariant_mode_enabled():
from sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe import (
_fused_finalize_all_reduce_comm_world_size,
)
comm_ready = (
_fused_finalize_all_reduce_comm_world_size()
== self.self_attn.attn_tp_size
)
if not is_batch_invariant_mode_enabled() and comm_ready:
attn_mhc = MhcPostFusion(
residual,
None,
None,
stats_stream,
record_stats=attn_stats,
norm_weight=self.post_attention_layernorm.weight,
norm_eps=self.post_attention_layernorm.variance_epsilon,
combine_only=medium_verify,
)
if (
prefill_overlap
and get_parallel().tp_size == self.self_attn.attn_tp_size == 4
and self.self_attn.wo_b.reduce_results
):
attn_mhc = MhcPostFusion(
residual,
None,
None,
stats_stream,
overlap_only=True,
record_stats=attn_stats,
)
context = (
use_mhc_post_fusion(attn_mhc) if attn_mhc is not None else nullcontext()
)
with context, self.self_attn.maybe_use_decode_attn_tp(forward_batch):
x = self.self_attn(
x=x,
positions=positions,
forward_batch=forward_batch,
x_quant=attn_quantized[0] if attn_quantized else None,
)
ffn_combined = None
ffn_normalized = None
if attn_mhc is not None:
attn_mhc.materialize_stats()
if attn_mhc is not None and attn_mhc.output is not None:
attn_pre = attn_mhc.pre
hidden_states = attn_mhc.output
ffn_combined = attn_mhc.combined
ffn_normalized = attn_mhc.normalized
else:
attn_pre, attn_post, attn_comb = (
(attn_mhc.pre, attn_mhc.post, attn_mhc.comb)
if attn_mhc is not None
else attn_stats()
)
if stats_stream is not None:
torch.cuda.current_stream().wait_stream(stats_stream)
hidden_states, ffn_combined, ffn_normalized = self._hc_post_with_combine(
x,
residual,
attn_post,
attn_comb,
attn_pre,
forward_batch,
norm=self.post_attention_layernorm,
)
residual = hidden_states
ffn_stats = partial(
self._hc_mix_stats,
hidden_states,
self.hc_ffn_fn,
self.hc_ffn_scale,
self.hc_ffn_base,
stats_stream,
)
x = self._hc_combine(
hidden_states,
apply_pre=attn_pre,
norm=self.post_attention_layernorm,
stats_stream=stats_stream,
normalized=ffn_normalized,
combined=ffn_combined,
)
mhc = None
if (
self.config.model_type == "deepseek_v41"
and x.is_cuda
and get_platform().is_blackwell
and (0 < x.shape[0] <= 8 or (medium_verify and next_combined is not None))
and x.shape[1] == 5120
and self.hc_mult == 4
and x.dtype == residual.dtype == torch.bfloat16
and residual.is_contiguous()
and get_parallel().attn_dp_size == 1
and get_moe_a2a_backend().is_none()
and not get_forward().sp_active
and not self.dsa_enable_prefill_cp
and not self.mlp._shared_expert_tp1
and self.mlp.tp_size == 4
and (not medium_verify or not is_batch_invariant_mode_enabled())
):
mhc = MhcPostFusion(
residual,
None,
None,
stats_stream,
record_stats=ffn_stats,
combine_only=medium_verify,
)
if next_norm is not None:
mhc.norm_weight = next_norm.weight
mhc.norm_eps = next_norm.variance_epsilon
if (
prefill_overlap
and self.mlp.tp_size == 4
and get_moe_a2a_backend().is_none()
):
mhc = MhcPostFusion(
residual,
None,
None,
stats_stream,
overlap_only=True,
record_stats=ffn_stats,
)
context = use_mhc_post_fusion(mhc) if mhc is not None else nullcontext()
with context:
x = self._run_moe_ffn_dp_sync(
x, forward_batch, input_ids=input_ids, input_ids_global=input_ids_global
)
if mhc is not None:
mhc.materialize_stats()
ffn_pre, ffn_post, ffn_comb = mhc.pre, mhc.post, mhc.comb
else:
ffn_pre, ffn_post, ffn_comb = ffn_stats()
if mhc is not None and mhc.output is not None:
hidden_states = mhc.output
if next_input is not None and mhc.quantized is not None:
next_input.append((mhc.normalized, Mxfp8SwizzledInput(*mhc.quantized)))
if next_combined is not None and mhc.combined is not None:
next_combined.append((mhc.combined, None))
else:
if stats_stream is not None:
torch.cuda.current_stream().wait_stream(stats_stream)
if next_combined is not None:
hidden_states, combined, normalized = self._hc_post_with_combine(
x,
residual,
ffn_post,
ffn_comb,
ffn_pre,
forward_batch,
norm=next_norm,
)
if combined is not None or normalized is not None:
next_combined.append((combined, normalized))
else:
hidden_states = self.hc_post(x, residual, ffn_post, ffn_comb)
return hidden_states, ffn_pre
def _run_moe_ffn_dp_sync(
self,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
*,
input_ids: Optional[torch.Tensor],
input_ids_global: Optional[torch.Tensor],
) -> torch.Tensor:
_use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)
_use_tp_moe_gather = (
not _use_cp
and get_parallel().attn_dp_size > 1
and get_moe_a2a_backend().is_none()
)
_use_tp_attn_a2a_scatter = (
not _use_cp
and get_parallel().attn_tp_size > 1
and not get_moe_a2a_backend().is_none()
)
# symmetric gather+scatter for the no-EP TP-MoE dp-attn path:
# all_gatherv gather (in self.mlp's dp_gather) + reduce_scatterv combine.
# The experts ARE TP-sharded by intermediate (moe_tp_size==tp_size), so
# the post-experts reduce is a SUM. reduce_scatterv does that sum+scatter
# in ONE op, REPLACING the MoE-internal post-experts all_reduce — so we
# MUST tell the MoE to skip it (mlp_reduce_scatter=True) or it
# double-reduces. Env-gated via SGLANG_DP_USE_GATHERV, default OFF.
_use_reduce_scatterv = (
_use_tp_moe_gather
and is_dp_gatherv_active()
and forward_batch.dp_padding_mode is not None
and not forward_batch.dp_padding_mode.is_max_len()
)
# SGLANG_DP_USE_REDUCE_SCATTER: in the MAX_LEN decode path (equal per-rank
# padding, gatherv inactive, no EP), replace the MoE-internal post-experts
# all_reduce + dp_scatter with an equal-chunk reduce_scatter. On ROCm this
# uses the aiter custom kernel (so BOTH gather and combine are aiter custom),
# elsewhere RCCL reduce_scatter; either way it cuts combine traffic ~2x vs
# all_reduce. tp_size==attn_dp_size required so the global buffer splits
# evenly into per-rank chunks.
_use_reduce_scatter = (
envs.SGLANG_DP_USE_REDUCE_SCATTER.get()
and _use_tp_moe_gather
and not _use_reduce_scatterv
and not should_use_dp_reduce_scatterv()
and forward_batch.dp_padding_mode is not None
and forward_batch.dp_padding_mode.is_max_len()
and get_parallel().tp_size == get_parallel().attn_dp_size
)
mlp_reduce_scatter = _use_cp or _use_reduce_scatterv or _use_reduce_scatter
# PoC (SGLANG_DP_SHARED_EXPERT_LOCAL): compute the replicated shared expert
# on LOCAL hidden before the gather and add it back after the combine
# (reduce_scatterv OR dp_scatter), instead of on the gathered global buffer.
# Applies to BOTH prefill and decode: the shared expert is a per-token MLP,
# so computing it on this rank's local tokens (M_local rows) is identical to
# computing it on the gathered global buffer (M_global rows) and keeping the
# local slice -- but costs 1/dp_size the rows. With a replicated (TP1) shared
# expert this cancels the TP1 "full-dim" cost in decode (M_local * dim ==
# M_global * dim/tp), so decode no longer pays the ~dp_size x penalty.
_shared_local = None
_do_shared_local = (
_SHARED_EXPERT_LOCAL
and _use_tp_moe_gather
and getattr(self.mlp, "shared_experts", None) is not None
and getattr(self.mlp, "_shared_expert_tp1", False)
)
if _use_cp:
moe_a2a_backend = get_moe_a2a_backend()
if moe_a2a_backend.is_none():
hidden_states = dsa_cp_gather_hidden_states(hidden_states)
else:
assert (
moe_a2a_backend.is_deepep()
or moe_a2a_backend.is_megamoe()
or moe_a2a_backend.is_mori()
), (
"CP requires moe_a2a_backend in ('deepep', 'megamoe', 'mori'), "
f"got {moe_a2a_backend.value!r}."
)
elif _use_tp_moe_gather:
hidden_states, local_hidden_states = (
get_global_dp_buffer(get_parallel().tp_group),
hidden_states,
)
if _do_shared_local and local_hidden_states.shape[0] > 0:
_shared_local = self.mlp._forward_shared_experts(local_hidden_states)
# self_attn has already reduced across attention TP, so these hidden
# states are replicated and must not be summed by a partial gather.
dp_gather_replicate(hidden_states, local_hidden_states, forward_batch)
_a2a_scatter_chunks: Optional[List[torch.Tensor]] = None
if _use_tp_attn_a2a_scatter:
s, r = get_parallel().attn_tp_size, get_parallel().attn_tp_rank
_a2a_scatter_chunks = list(hidden_states.tensor_split(s))
hidden_states = _a2a_scatter_chunks[r].contiguous()
# DSpark next-token layers are not hash-routed and intentionally do not
# carry token IDs. Only split IDs for callers that actually provide them.
if input_ids is not None:
input_ids = input_ids.tensor_split(s)[r].contiguous()
if input_ids_global is not None:
input_ids_global = input_ids_global.tensor_split(s)[r].contiguous()
# Skip the MoE-internal post-experts all_reduce when we will do the
# reduce via reduce_scatterv/reduce_scatter at the combine below
# (else double-reduce).
gathered_rows = (
_every_row_routed(forward_batch, hidden_states.shape[0])
if _use_cp and get_moe_a2a_backend().is_none()
else nullcontext()
)
# The MoE sees DP-gathered rows, so this rank's local count cannot mask them.
# The standard dispatcher masks padding in the gathered buffer.
saved_num_token_non_padded = forward_batch.num_token_non_padded
if _use_tp_moe_gather:
forward_batch.num_token_non_padded = None
try:
with (
get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter),
gathered_rows,
):
hidden_states = self.mlp(
hidden_states,
forward_batch,
input_ids=input_ids,
input_ids_global=input_ids_global,
skip_shared_experts=_do_shared_local,
)
finally:
forward_batch.num_token_non_padded = saved_num_token_non_padded
if _use_cp and get_moe_a2a_backend().is_none():
if self.config.model_type == "deepseek_v41":
parallel = get_parallel()
hidden_states = parallel.tp_group.all_reduce(hidden_states)
parallel = get_parallel()
hidden_states = hidden_states.tensor_split(parallel.attn_cp_size)[
parallel.attn_cp_rank
].contiguous()
else:
hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states)
elif _use_tp_moe_gather:
hidden_states, global_hidden_states = (
get_local_dp_buffer(get_parallel().tp_group),
hidden_states,
)
if should_use_dp_reduce_scatterv() or _use_reduce_scatterv:
# SUM the TP-sharded per-rank partial expert outputs AND scatter
# each rank its own token slice, in one op. Correct because the
# MoE-internal all_reduce was skipped (mlp_reduce_scatter above).
# This is the symmetric inverse of the all_gatherv gather.
get_parallel().tp_group.reduce_scatterv(
global_hidden_states,
output=hidden_states,
sizes=get_dp_global_num_tokens(),
)
elif _use_reduce_scatter:
# Equal-chunk reduce_scatter: SUM the TP-sharded per-rank partial
# expert outputs AND scatter each rank its own (MAX_LEN-padded)
# token chunk in one op (symmetric inverse of the MAX_LEN
# all_gather). Correct because the MoE-internal all_reduce was
# skipped (mlp_reduce_scatter above). dp_reduce_scatter_tensor
# routes to the equal-chunk reduce_scatter_tensor here (its
# variable-length reduce_scatterv branch is gated by
# is_dp_gatherv_active(), which is False under MAX_LEN), which in
# turn uses the aiter custom kernel when it fits (else RCCL).
dp_reduce_scatter_tensor(hidden_states, global_hidden_states)
else:
dp_scatter(hidden_states, global_hidden_states, forward_batch)
# PoC: add the locally-computed shared-expert output to this rank's
# reduce-scattered / dp-scattered local slice (skipped inside self.mlp
# above). Covers both prefill (gatherv) and decode (dp_scatter).
if _shared_local is not None:
n = hidden_states.shape[0]
hidden_states = hidden_states + _shared_local[:n]
if _use_tp_attn_a2a_scatter:
assert _a2a_scatter_chunks is not None
gathered = [torch.empty_like(t) for t in _a2a_scatter_chunks]
attn_tp_all_gather(gathered, hidden_states.contiguous())
hidden_states = torch.cat(gathered)
return hidden_states
# ------------------------------------------------------------------
# TBO op decomposition (prefill two-batch-overlap, EP / mori path)
#
# These mirror the NON-fused branch of ``forward`` (cross-layer mHC
# fusion is disabled under TBO, so every layer is self-contained), split
# into ops so the operations engine can overlap one ubatch's MoE a2a
# dispatch/combine with the other ubatch's attention + expert GEMM.
# The MoE ops themselves (op_gate / op_select_experts / op_dispatch_a/b /
# op_experts / op_combine_a/b / op_shared_experts / op_output) are reused
# as-is from ``self.mlp`` (DeepseekV2MoE) — they decompose ``forward_deepep``.
# ------------------------------------------------------------------
def op_mhc_prepare_attn(
self,
state,
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
residual: Optional[torch.Tensor] = None,
tbo_subbatch_index: Optional[int] = None,
**kwargs,
):
# Non-fused attention-side mHC pre + input layernorm.
attn_residual = hidden_states
hidden_states, post, comb, norm_fused = self.hc_pre(
hidden_states,
self.hc_attn_fn,
self.hc_attn_scale,
self.hc_attn_base,
norm=self.input_layernorm,
forward_batch=forward_batch,
)
if not norm_fused:
if _use_aiter and (_is_gfx95_supported or _is_gfx1250_supported):
x_quant, hidden_states = _fused_rmsnorm_fp8_quant(
hidden_states,
self.input_layernorm.weight,
self.rms_norm_eps,
)
else:
hidden_states = self.input_layernorm(hidden_states)
x_quant = None
else:
x_quant = None
state.attn_residual = attn_residual
state.attn_post = post
state.attn_comb = comb
state.hidden_states_after_input_norm = hidden_states
state.attn_x_quant = x_quant
# mori's op_output slices final_hidden_states[:num_tokens].
if get_moe_a2a_backend().is_mori():
state.num_tokens = attn_residual.shape[0]
state.update(
dict(
forward_batch=forward_batch,
positions=positions,
tbo_subbatch_index=tbo_subbatch_index,
)
)
def op_mhc_post_attn_pre_mlp(self, state):
# Close the attention mHC (hc_post), then open the FFN-side mHC pre +
# post-attention layernorm. Produces the 2D MoE input.
#
# Pop each boundary tensor from the state EXACTLY ONCE, up front, and
# reuse the locals for both the fused attempt and the non-fused
# fallback. apply_mhc_post_pre_boundary() returns None when it declines
# to fuse -- most importantly for the 0-token DP two-batch-overlap idle
# ubatch -- in which case control must fall through to the unfused
# hc_post. Popping in the fused call's arguments and again in the
# fallback would double-pop -> KeyError: 'hidden_states_after_attn' on
# every idle DP rank. use_fused_mhc_post_pre is on whenever the aiter
# gfx95 mHC path is available (is_cross_layer_mhc_fusion_enabled), so
# this fallback is reached under DP regardless of the TileLang env.
hidden_states_after_attn = state.pop("hidden_states_after_attn")
attn_residual = state.pop("attn_residual")
attn_post = state.pop("attn_post")
attn_comb = state.pop("attn_comb")
if self.use_fused_mhc_post_pre:
post_attn_norm_weight = (
self._post_attention_layernorm_weight_bf16
if self._post_attention_layernorm_weight_bf16 is not None
else self.post_attention_layernorm.weight.data
)
fused = apply_mhc_post_pre_boundary(
hidden_states_after_attn,
attn_residual,
attn_post,
attn_comb,
self.hc_ffn_fn,
self.hc_ffn_scale,
self.hc_ffn_base,
self.hc_mult,
self.rms_norm_eps,
self.hc_eps,
_MHC_POST_MULT_VALUE,
self.hc_sinkhorn_iters,
post_attn_norm_weight,
self.post_attention_layernorm.variance_epsilon,
fn_transpose=True,
)
if fused is not None:
ffn_residual, hidden_states, post, comb, norm_fused = fused
if not norm_fused:
# The Triton fused post+pre skips the post-attention
# layernorm (norm_fused=False); apply it before the MoE,
# matching the unfused hc_pre path below.
hidden_states = self.post_attention_layernorm(hidden_states)
state.ffn_residual = ffn_residual
state.ffn_post = post
state.ffn_comb = comb
state.hidden_states_mlp_input = hidden_states
return
hidden_states = self.hc_post(
hidden_states_after_attn,
attn_residual,
attn_post,
attn_comb,
)
ffn_residual = hidden_states
hidden_states, post, comb, norm_fused = self.hc_pre(
hidden_states,
self.hc_ffn_fn,
self.hc_ffn_scale,
self.hc_ffn_base,
norm=self.post_attention_layernorm,
forward_batch=state.forward_batch,
)
if not norm_fused:
hidden_states = self.post_attention_layernorm(hidden_states)
state.ffn_residual = ffn_residual
state.ffn_post = post
state.ffn_comb = comb
state.hidden_states_mlp_input = hidden_states
def op_mhc_postprocess(self, state):
# Close the FFN mHC (hc_post) and emit the next layer's input dict.
hidden_states = self.hc_post(
state.pop("hidden_states_mlp_output"),
state.pop("ffn_residual"),
state.pop("ffn_post"),
state.pop("ffn_comb"),
)
output = dict(
positions=state.positions,
hidden_states=hidden_states,
# DSV4 non-fused layers carry no residual across layers; the key is
# required by the next layer's op_mhc_prepare_attn (ignored) and by
# _model_forward_tbo_merge_outputs (None -> None).
residual=None,
forward_batch=state.forward_batch,
tbo_subbatch_index=state.tbo_subbatch_index,
)
state.clear(
expect_keys={
"positions",
"forward_batch",
"tbo_subbatch_index",
}
)
return output
# ------------------------------------------------------------------
# Non-EP (DP TP-MoE) TBO ops. Overlap the DP all_gatherv (pre-MoE gather)
# + reduce_scatterv (post-MoE combine) with the OTHER ubatch's attn+MoE
# compute. Used when moe_a2a_backend is "none" (DP-attention, TP-MoE) —
# the path ATOM uses for DSV4 (+~7.7% prefill). Replaces the EP mori
# op_dispatch/op_combine. op_mhc_* and op_attn are reused (local hidden).
# ------------------------------------------------------------------
def op_gather_a(self, state):
# Launch the all_gatherv (local hidden -> global buffer) + the input_ids
# replicate-gather on the shared comm stream; record an event.
fb = state.forward_batch
local = state.pop("hidden_states_mlp_input") # LOCAL [M_local, hidden]
# Shared-expert-local: compute on LOCAL hidden before the gather; added
# back after the combine (same as the non-fused forward). Skipped in the
# global MoE via skip_shared_experts.
do_shared_local = (
_SHARED_EXPERT_LOCAL
and getattr(self.mlp, "shared_experts", None) is not None
and getattr(self.mlp, "_shared_expert_tp1", False)
)
state.do_shared_local = do_shared_local
state.shared_local = (
self.mlp._forward_shared_experts(local)
if (do_shared_local and local.shape[0] > 0)
else None
)
# Persistent grow-only scratch (keyed per ubatch) instead of a fresh
# torch.empty each layer -> stops the allocator's `reserved` from
# ballooning at large prefill chunks. input_ids_global is gathered ONCE
# per ubatch in _forward_layers_tbo (cached on fb), not here.
sub = state.tbo_subbatch_index
global_rows = get_global_dp_buffer_len()
global_hidden = get_tbo_persistent_buffer(
("gh", sub), global_rows, local.shape[1], local.dtype, local.device
)
comm = get_dp_tbo_comm_stream()
compute = torch.cuda.current_stream()
with torch.cuda.stream(comm):
comm.wait_stream(compute)
dp_gather_partial(global_hidden, local, fb)
state.gather_event = _tbo_event(("gather", sub))
state.gather_event.record(comm)
state.gather_keepalive = local
state.global_hidden = global_hidden
def op_gather_b(self, state):
torch.cuda.current_stream().wait_event(state.pop("gather_event"))
# Compute now ordered after the gather -> the gather input is safe to
# release (freed on the compute stream, no record_stream deferral).
state.pop("gather_keepalive")
def op_moe(self, state):
# MoE (gate/topk/experts) on the GLOBAL gathered buffer. mlp_reduce_scatter
# skips the MoE-internal all_reduce (we reduce_scatterv in op_combine).
fb = state.forward_batch
global_hidden = state.pop("global_hidden")
global_ids = fb._tbo_global_input_ids
with get_forward().scoped(mlp_reduce_scatter=True):
state.global_expert_out = self.mlp(
global_hidden,
fb,
input_ids=global_ids,
input_ids_global=global_ids,
skip_shared_experts=state.do_shared_local,
)
def op_combine_a(self, state):
# Launch reduce_scatterv (global partial expert sums -> per-rank local) on
# the comm stream; record an event. Symmetric inverse of the all_gatherv.
global_out = state.pop("global_expert_out")
local_out = get_tbo_persistent_buffer(
("lo", state.tbo_subbatch_index),
get_local_dp_buffer_len(),
global_out.shape[1],
global_out.dtype,
global_out.device,
)
state.combine_event = dp_reduce_scatterv_async(
local_out,
global_out,
get_dp_global_num_tokens(),
event_key=("combine", state.tbo_subbatch_index),
)
state.local_out = local_out
# Keep the (variable-size) MoE output alive until op_combine_b waits on
# the combine event (replaces record_stream; avoids reserved churn).
state.combine_keepalive = global_out
def op_combine_b(self, state):
torch.cuda.current_stream().wait_event(state.pop("combine_event"))
state.pop("combine_keepalive")
hidden = state.pop("local_out")
shared_local = state.pop("shared_local")
state.pop("do_shared_local")
if shared_local is not None:
n = hidden.shape[0]
hidden = hidden + shared_local[:n]
state.hidden_states_mlp_output = hidden
def _scatter_tail_rows(
tail: LateLayerTail, rows: torch.Tensor, num_tokens: int
) -> torch.Tensor:
# Rows outside the tail are never read (see _check_late_layer_tail_readers).
full = rows.new_empty((num_tokens, rows.shape[1]))
if tail.contiguous_start is not None:
full[tail.contiguous_start :].copy_(rows)
else:
full[tail.token_indices] = rows[: tail.token_indices.shape[0]]
return full
class DeepseekV4Model(nn.Module):
fall_back_to_pt_during_load = False
def __init__(
self,
config: DeepSeekV4Config,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__()
self.config = config
self.pp_group = get_parallel().pp_group
self.hidden_size = config.hidden_size
if self.pp_group.is_first_rank:
embedding_quant_config = (
quant_config
if quant_config is not None and quant_config.get_name() == "expert_pack"
else None
)
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
config.hidden_size,
enable_tp=not is_dp_attention_enabled(),
quant_config=embedding_quant_config,
prefix=add_prefix("embed_tokens", prefix),
)
else:
self.embed_tokens = PPMissingLayer()
self.rms_norm_eps = config.rms_norm_eps
use_stream_pool = (
_is_cuda
or (
_is_hip
and (
envs.SGLANG_ROCM_USE_MULTI_STREAM.get()
or envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get()
)
)
or (_is_npu and envs.SGLANG_NPU_USE_MULTI_STREAM.get())
)
device_module = torch.get_device_module()
num_alt_streams = 5 if (_is_cuda or _is_npu) else 2
self.alt_streams = (
[device_module.Stream() for _ in range(num_alt_streams)]
if use_stream_pool
else None
)
# Routed-MoE input pre-quant, separate from the attention/indexer and
# shared-expert streams.
self.moe_routed_quant_stream = (
device_module.Stream()
if _is_cuda
and config.hc_pre_from_prev_sublayer
and envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get()
else None
)
# One shared stream for all layers; every sublayer joins it before
# reusing its residual.
self.hc_stats_stream = (
device_module.Stream()
if _is_cuda
and (get_platform().is_blackwell or get_platform().is_sm90)
and envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get()
and config.hc_pre_from_prev_sublayer
else None
)
self.engram_layout = EngramLayout.from_config(config)
self.layers, self.start_layer, self.end_layer = make_layers(
config.num_hidden_layers,
lambda idx, prefix: DeepseekV4DecoderLayer(
config=config,
layer_id=idx,
quant_config=quant_config,
prefix=prefix,
alt_streams=self.alt_streams,
engram_layout=self.engram_layout,
hc_stats_stream=self.hc_stats_stream,
moe_routed_quant_stream=self.moe_routed_quant_stream,
),
pp_rank=self.pp_group.rank_in_group,
pp_size=self.pp_group.world_size,
prefix=add_prefix("layers", prefix),
)
if self.pp_group.is_last_rank:
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
else:
self.norm = PPMissingLayer()
self.gemm_output_zero_allocator_size = 0
self.hc_eps = config.hc_eps
self.hc_mult = hc_mult = config.hc_mult
self.norm_eps = config.rms_norm_eps
self.hc_pre_from_prev_sublayer = config.hc_pre_from_prev_sublayer
self.hc_head_fn = self.hc_head_base = self.hc_head_scale = None
if self.pp_group.is_last_rank and not self.hc_pre_from_prev_sublayer:
(
self.hc_head_fn,
self.hc_head_base,
self.hc_head_scale,
) = make_hc_head_params(hc_mult, config.hidden_size)
self.engram_hasher = None
if self.engram_layout is not None:
self.engram_hasher = EngramHasher.from_config(
config,
self.engram_layout,
image_token_id=(
config.image_token_id
if config.model_type == "deepseek_v41"
and config.vision_n_layers > 0
else None
),
)
self.use_fused_mhc_post_pre = (
is_cross_layer_mhc_fusion_enabled() or _is_fused_mhc_post_pre_enabled_xpu()
)
self.dspark_layers_to_capture: Optional[List[int]] = None
# Decoder SWA bounded replay: layers past the last kv_source layer run over
# each request's last SWA_WINDOW extend tokens only.
self.late_layer_start: Optional[int] = None
if get_exec().features.enable_decoder_swa_bounded_replay:
assert config.kv_source_layer_ids, (
"decoder SWA bounded replay needs kv_source_layer_ids"
)
self.late_layer_start = max(config.kv_source_layer_ids) + 1
late_ratios = set(
config.compress_ratios[self.late_layer_start : config.num_hidden_layers]
)
assert late_ratios <= {
0,
1,
}, f"late layers must not compress on their own, got ratios {late_ratios}"
def get_input_embeddings(self) -> nn.Module:
return self.embed_tokens
def hc_head(
self,
x: torch.Tensor,
hc_fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
):
if x.numel() > 0:
if _is_xpu:
return _get_mhc_ops().fused_hc_head(
x.contiguous(),
hc_fn,
hc_scale,
hc_base,
norm_eps=self.norm_eps,
hc_eps=self.hc_eps,
)
from sglang.kernels.ops.layernorm.mhc_head import fused_hc_head
return fused_hc_head(
x.contiguous(),
hc_fn,
hc_scale,
hc_base,
norm_eps=self.norm_eps,
hc_eps=self.hc_eps,
)
return hc_head_torch(
x,
hc_fn,
hc_scale,
hc_base,
norm_eps=self.norm_eps,
hc_eps=self.hc_eps,
)
def _check_late_layer_tail_readers(self, forward_batch: ForwardBatch) -> None:
# Rows outside the tail are never computed past the last kv_source layer.
if (
forward_batch.capture_hidden_mode == CaptureHiddenMode.FULL
and self.dspark_layers_to_capture is None
):
raise ValueError(
"decoder SWA bounded replay cannot capture hidden states of all "
"prompt tokens"
)
if forward_batch.return_logprob and any(
start < n
for start, n in zip(
forward_batch.extend_logprob_start_lens_cpu,
forward_batch.extend_seq_lens_cpu,
)
):
raise ValueError(
"decoder SWA bounded replay cannot return logprobs of prompt tokens; "
"set logprob_start_len to the prompt length"
)
def _forward_layers_hc_pre_from_prev(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
input_ids: torch.Tensor,
input_ids_global: torch.Tensor,
capture_dspark: bool,
dspark_aux_hidden_states: List[torch.Tensor],
) -> Tuple[torch.Tensor, torch.Tensor, Optional[LateLayerTail]]:
assert self.pp_group.world_size == 1, "pre-mix hand-off across PP is not wired"
hash_ids = None
cp_extend = (
is_cp_active(forward_batch) and forward_batch.forward_mode.is_extend()
)
if self.engram_hasher is not None:
if cp_extend:
# N-gram hashing needs each token's predecessors, so hash the
# whole prompt before selecting this CP rank's interleaved rows.
# The hasher builds request-to-token indices dynamically; keep
# that work at an eager break during breakable graph capture.
total = int(forward_batch.attn_cp_metadata.total_seq_lens)
global_input_ids = forward_batch.input_ids[:total]
if is_in_breakable_cuda_graph():
hash_ids = bcg_deepseek_v4_engram_hash_ids(
self.engram_hasher, global_input_ids
)
else:
hash_ids = self.engram_hasher(global_input_ids, forward_batch)
parallel = get_parallel()
hash_ids = hash_ids[parallel.attn_cp_rank :: parallel.attn_cp_size]
pad_rows = hidden_states.shape[0] - hash_ids.shape[0]
if pad_rows > 0:
hash_ids = torch.cat(
[hash_ids, hash_ids.new_zeros(pad_rows, *hash_ids.shape[1:])]
)
elif (
forward_batch.forward_mode.is_extend() and is_in_breakable_cuda_graph()
):
hash_ids = bcg_deepseek_v4_engram_hash_ids(
self.engram_hasher, input_ids
)
else:
hash_ids = self.engram_hasher(input_ids, forward_batch)
tail = None
if (
self.late_layer_start is not None
and forward_batch.forward_mode.is_extend_without_speculative()
):
self._check_late_layer_tail_readers(forward_batch)
attn_backend = get_attn_backend()
tail = attn_backend.tail_forward_metadata.late_layer_tail
saved_full = None
prev_pre = None
precomputed_attn = None
combined_attn = None
normalized_attn = None
for i in range(self.start_layer, self.end_layer):
if tail is not None and i == self.late_layer_start:
combined_attn = None
normalized_attn = None
# Decode reaches back at most SWA_WINDOW positions.
saved_full = attn_backend.enter_late_layer_tail(forward_batch)
hidden_states, prev_pre, input_ids, input_ids_global = (
tail.rows(hidden_states),
tail.rows(prev_pre),
tail.rows(input_ids),
tail.rows(input_ids_global),
)
positions = tail.positions
if hash_ids is not None:
hash_ids = tail.rows(hash_ids)
engram = self.layers[i].engram
if engram is not None:
precomputed_attn = None
combined_attn = None
normalized_attn = None
before_engram = hidden_states
hidden_states = engram(
hidden_states,
hash_ids[:, engram.layer_hash_index],
forward_batch,
cp_all_tokens=cp_extend,
)
if (
self.config.model_type == "deepseek_v41"
and self.config.vision_n_layers > 0
):
hidden_states = torch.where(
(input_ids == self.config.image_token_id)[:, None, None],
before_engram,
hidden_states,
)
if capture_dspark and i in self.dspark_layers_to_capture:
# The draft head reads the attention input of its target layers.
aux = hidden_states
if tail is not None and i < self.late_layer_start:
aux = tail.rows(aux)
dspark_aux_hidden_states.append(aux.mean(dim=1))
ctx = (
nullcontext()
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
else get_global_expert_distribution_recorder().with_current_layer(i)
)
next_norm = None
next_input = []
# The next layer can consume a collapsed input only if no Engram
# or row selection changes the residual between the two layers.
next_combined = (
[]
if (
self.config.model_type == "deepseek_v41"
and (
128 <= hidden_states.shape[0] <= 384
or (
4096 <= hidden_states.shape[0] <= 65536
and forward_batch.forward_mode.is_extend_without_speculative()
)
)
and i + 1 < self.end_layer
and tail is None
and self.layers[i + 1].engram is None
)
else None
)
if next_combined is not None and hidden_states.shape[0] >= 4096:
next_norm = self.layers[i + 1].input_layernorm
if (
self.config.model_type == "deepseek_v41"
and i + 1 < self.end_layer
and tail is None
and hidden_states.is_cuda
and get_platform().is_blackwell
and 0 < hidden_states.shape[0] <= 8
and (
forward_batch.forward_mode.is_decode()
or forward_batch.forward_mode.is_target_verify()
)
and not get_forward().sp_active
and self.layers[i + 1].engram is None
and self.layers[i + 1].self_attn.accepts_mxfp8_swizzled_input()
):
from sglang.srt.batch_invariant_ops import (
is_batch_invariant_mode_enabled,
)
norm = self.layers[i + 1].input_layernorm
if (
not norm.cast_x_before_out_mul
and norm.variance_size_override is None
and norm.weight.dtype == torch.bfloat16
and norm.weight.shape == (5120,)
and norm.weight.is_contiguous()
and not is_batch_invariant_mode_enabled()
):
next_norm = norm
with ctx:
hidden_states, prev_pre = self.layers[i].forward_hc_pre_from_prev(
positions=positions,
hidden_states=hidden_states,
input_ids=input_ids,
forward_batch=forward_batch,
input_ids_global=input_ids_global,
prev_pre=prev_pre,
precomputed_attn=precomputed_attn,
next_norm=next_norm,
next_input=next_input,
combined_attn=combined_attn,
normalized_attn=normalized_attn,
next_combined=next_combined,
)
precomputed_attn = next_input[0] if next_input else None
combined_attn, normalized_attn = (
next_combined[0] if next_combined else (None, None)
)
if saved_full is not None:
attn_backend.exit_late_layer_tail(saved_full, forward_batch)
return hidden_states, prev_pre, tail
return hidden_states, prev_pre, None
def _can_run_tbo(self, forward_batch: ForwardBatch) -> bool:
"""DSV4 prefill-only two-batch-overlap gate.
TBO batch prep (tbo_split_seq_index / tbo_children) is populated
model-agnostically when --enable-two-batch-overlap is set and the
DP-attention preparer allows it (mori `normal` mode permits prefill
TBO). We additionally restrict to prefill (EXTEND), single PP, and
non-CP paths supported by the DSV4 op strategy.
"""
from sglang.srt.layers.moe import is_tbo_enabled
path_ok = not dsa_use_prefill_cp(forward_batch) and (
not _is_hip
or not get_moe_a2a_backend().is_none()
or get_parallel().attn_dp_size > 1
)
return (
is_tbo_enabled()
and forward_batch.can_run_tbo
and forward_batch.tbo_children is not None
and forward_batch.global_forward_mode is not None
# MTP target-verify also reports is_extend(); only real prefill
# should enter the prefill TBO strategy.
and forward_batch.global_forward_mode.is_extend_without_speculative()
and path_ok
and self.pp_group.world_size == 1
)
def _forward_layers_tbo(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
) -> torch.Tensor:
from sglang.srt.batch_overlap.operations import execute_overlapped_operations
from sglang.srt.batch_overlap.operations_strategy import OperationsStrategy
from sglang.srt.batch_overlap.two_batch_overlap import (
_model_forward_filter_inputs,
_model_forward_tbo_merge_outputs,
)
layers = [self.layers[i] for i in range(self.start_layer, self.end_layer)]
operations_strategy = OperationsStrategy.init_new_tbo(
layers, forward_batch.global_forward_mode
)
# Split the per-rank batch into the 2 ubatches (token-range slice + pad
# to tbo_padded_len). residual is unused by the DSV4 non-fused layer ops.
inputs_arr = [
_model_forward_filter_inputs(
hidden_states=hidden_states,
residual=None,
positions=positions,
output_forward_batch=child,
tbo_subbatch_index=idx,
)
for idx, child in enumerate(forward_batch.tbo_children)
]
# Non-EP DP TP-MoE: the per-ubatch DP gather/combine (op_gather/op_combine)
# needs each ubatch's per-rank token counts, but tbo_padded_len is computed
# per-rank locally (not synced). All-gather both ubatches' padded lengths
# once across DP ranks, then populate each child's global_num_tokens +
# global_dp_buffer_len so the gatherv/reduce_scatterv buffers size correctly.
if get_moe_a2a_backend().is_none() and get_parallel().attn_dp_size > 1:
tp_group = get_parallel().tp_group
world = tp_group.world_size
children = forward_batch.tbo_children
local_lens = torch.tensor(
[int(c.tbo_padded_len) for c in children],
dtype=torch.int64,
device=hidden_states.device,
)
gathered = torch.empty(
(world, local_lens.shape[0]),
dtype=torch.int64,
device=hidden_states.device,
)
tp_group.all_gather_into_tensor(gathered, local_lens)
gathered_cpu = gathered.tolist()
rank = tp_group.rank_in_group
for idx, child in enumerate(children):
sizes = [gathered_cpu[r][idx] for r in range(world)]
child.global_num_tokens_cpu = sizes
child.global_num_tokens_gpu = gathered[:, idx].contiguous()
child.global_dp_buffer_len = sum(sizes)
# Gather the ubatch's input_ids -> global ONCE here (cached on the
# child) instead of per-layer in op_gather_a. The hash MoE reads
# the SAME global ids every layer, so 61x2 per-layer all_gatherv of
# VARYING size (-> RCCL registers a new internal buffer per size ->
# HSA_STATUS_ERROR_OUT_OF_RESOURCES) collapses to 1 per ubatch.
local_ids = child.input_ids
rows = sizes[rank]
if local_ids.shape[0] < rows:
padded_ids = local_ids.new_zeros((rows,))
padded_ids[: local_ids.shape[0]] = local_ids
elif local_ids.shape[0] > rows:
padded_ids = local_ids[:rows]
else:
padded_ids = local_ids
gids = torch.empty(
(sum(sizes),), dtype=local_ids.dtype, device=local_ids.device
)
tp_group.all_gatherv(padded_ids, sizes=sizes, output=gids)
child._tbo_global_input_ids = gids
outputs_arr = execute_overlapped_operations(
inputs_arr=inputs_arr,
operations_arr=[operations_strategy.operations] * 2,
delta_stages=[0, operations_strategy.tbo_delta_stages],
)
hidden_states, _ = _model_forward_tbo_merge_outputs(
outputs_arr[0], outputs_arr[1], hidden_states.shape[0]
)
return hidden_states
@torch.no_grad()
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds: Optional[torch.Tensor],
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> Union[torch.Tensor, PPProxyTensors]:
if self.pp_group.is_first_rank:
if input_embeds is None:
hidden_states = self.embed_tokens(input_ids)
else:
hidden_states = input_embeds
hidden_states = hidden_states.unsqueeze(1).repeat(1, self.hc_mult, 1)
else:
assert pp_proxy_tensors is not None
hidden_states = pp_proxy_tensors["hidden_states"]
# Unflatten 2D PP IPC tensor back to 3D mHC shape.
if hidden_states.ndim == 2:
hidden_states = hidden_states.view(
hidden_states.shape[0], self.hc_mult, self.hidden_size
)
if get_parallel().attn_dp_size > 1 and get_moe_a2a_backend().is_none():
input_ids_global = torch.empty(
(get_global_dp_buffer_len(), 1),
dtype=input_ids.dtype,
device=input_ids.device,
)
# Token ids are replicated within an attention-TP group. Use replicate
# gather here to avoid summing duplicated ids when attention_tp_size > 1.
# Clone because the MAX_LEN gather may zero its local input in place.
dp_gather_replicate(
input_ids_global, input_ids[:, None].clone(), forward_batch
)
input_ids_global = input_ids_global.squeeze(-1)
else:
input_ids_global = getattr(forward_batch, "input_ids_global", input_ids)
capture_dspark = self.dspark_layers_to_capture is not None
dspark_aux_hidden_states: List[torch.Tensor] = []
attn_backend = get_attn_backend()
if _is_npu and forward_batch.attn_cp_metadata is not None:
attn_backend.prepare_dsv4_cp_metadata(forward_batch)
local_positions = getattr(forward_batch, "dsv4_cp_local_positions", None)
if (
local_positions is not None
and positions.shape[0] == local_positions.shape[0]
):
forward_batch.positions = positions
# Reset Compressor's per-step freqs_cis cache from any previous step.
for _attr in ("freqs_cis_c4", "freqs_cis_c128"):
if hasattr(forward_batch, _attr):
delattr(forward_batch, _attr)
run_tbo = self._can_run_tbo(forward_batch) and not capture_dspark
if _is_npu and not run_tbo:
# Rope cos/sin for the whole forward: one bf16 gather per rope
# config on the current stream, before the layer loop forks the
# KV/Q side streams. TBO children carry their own positions and
# recompute per layer.
prime_rope_cos_sin(
(
self.layers[i].self_attn
for i in range(self.start_layer, self.end_layer)
),
forward_batch,
positions,
)
last_pre = None
tail = None
if self.hc_pre_from_prev_sublayer:
assert not run_tbo, "two-batch overlap is not wired for this hc scheme"
hidden_states, last_pre, tail = self._forward_layers_hc_pre_from_prev(
positions,
hidden_states,
forward_batch,
input_ids,
input_ids_global,
capture_dspark,
dspark_aux_hidden_states,
)
elif run_tbo:
# Two-batch-overlap prefill (EP / mori). Cross-layer mHC fusion is
# disabled here (each layer self-contained), so no trailing hc_post.
hidden_states = self._forward_layers_tbo(
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
)
else:
use_fused = self.use_fused_mhc_post_pre
prev_residual, prev_post, prev_comb = None, None, None
last_layer = None
for i in range(self.start_layer, self.end_layer):
layer = self.layers[i]
last_layer = layer
ctx = (
nullcontext()
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
else get_global_expert_distribution_recorder().with_current_layer(i)
)
with ctx:
hidden_states, prev_residual, prev_post, prev_comb = layer(
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
input_ids=input_ids,
input_ids_global=input_ids_global,
prev_residual=prev_residual,
prev_post=prev_post,
prev_comb=prev_comb,
)
if capture_dspark and i in self.dspark_layers_to_capture:
if use_fused:
completed = layer.hc_post(
hidden_states, prev_residual, prev_post, prev_comb
)
else:
completed = hidden_states
dspark_aux_hidden_states.append(completed.mean(dim=1))
if use_fused and last_layer is not None:
hidden_states = last_layer.hc_post(
hidden_states, prev_residual, prev_post, prev_comb
)
if not self.pp_group.is_last_rank:
# Flatten 3D mHC tensor for PP IPC.
return PPProxyTensors({"hidden_states": hidden_states.flatten(1)})
pre_hc_head = hidden_states.flatten(1)
if self.hc_pre_from_prev_sublayer:
from sglang.kernels.ops.layernorm.mhc import hc_combine
hidden_states = hc_combine(
pre_hc_head.float(), last_pre, self.hc_mult, hidden_states.dtype
)
else:
hidden_states = self.hc_head(
hidden_states, self.hc_head_fn, self.hc_head_scale, self.hc_head_base
)
hidden_states = self.norm(hidden_states)
if tail is not None and not capture_dspark:
# The logits processor indexes rows by the full extend layout.
num_tokens = input_ids.shape[0]
hidden_states = _scatter_tail_rows(
tail=tail, rows=hidden_states, num_tokens=num_tokens
)
pre_hc_head = _scatter_tail_rows(
tail=tail, rows=pre_hc_head, num_tokens=num_tokens
)
if capture_dspark:
return (hidden_states, pre_hc_head), dspark_aux_hidden_states
return hidden_states, pre_hc_head
def _v41_vision_a2a_supported() -> bool:
backend = get_moe_a2a_backend()
return backend.is_none() or (
backend.is_megamoe() and get_disagg().disaggregation_mode == "decode"
)
class DeepseekV4ForCausalLM(nn.Module):
supports_cuda_vmm_feature_transport = True
def __init__(
self,
config: DeepSeekV4Config,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__()
# DeepseekV4 enables, by default, the CK w8a8-block GEMM (MLA proj) and the
# batched/contiguous-load rope kernels (faster on gfx95; .
# Module-level toggles default OFF; flipped True here for DSV4
if _is_hip:
from sglang.kernels.ops.attention.deepseek_v4_rope import set_batched_rope
from sglang.srt.layers.quantization.fp8_utils import set_force_ck_w8a8
set_force_ck_w8a8(True)
set_batched_rope(True)
self.config = config
self.tp_size = get_parallel().tp_size
self.quant_config = quant_config
self.wo_a_fp8 = wo_a_fp8_gemm_enabled(quant_config)
self.determine_num_fused_shared_experts()
self.vision = None
if (
config.model_type == "deepseek_v41"
and config.vision_n_layers > 0
and not getattr(config, "language_model_only", False)
):
if (
get_parallel().attn_cp_size != 1
or get_pp_group().world_size != 1
or not _v41_vision_a2a_supported()
):
raise ValueError(
"V4.1 vision supports TP/EP/DP without CP or PP; "
"MoE A2A is supported only with MegaMoE on a PD decode node"
)
args = SimpleNamespace(**vars(config), dim=config.hidden_size)
self.vision = ViT(args)
self.aligner = Aligner(args)
self.image_start = nn.Parameter(torch.empty(config.hidden_size))
self.image_end = nn.Parameter(torch.empty(config.hidden_size))
self.image_newline = nn.Parameter(torch.empty(config.hidden_size))
self.model = DeepseekV4Model(
config, quant_config, prefix=add_prefix("model", prefix)
)
self.pp_group = get_parallel().pp_group
if self.pp_group.is_last_rank:
if self.pp_group.world_size == 1 and config.tie_word_embeddings:
self.lm_head = self.model.embed_tokens
else:
self.lm_head = ParallelLMHead(
config.vocab_size,
config.hidden_size,
quant_config=quant_config,
prefix=add_prefix("lm_head", prefix),
use_attn_tp_group=get_parallel().enable_dp_lm_head,
)
else:
self.lm_head = PPMissingLayer()
self.logits_processor = LogitsProcessor(config)
self.capture_aux_hidden_states = False
get_attn_tp_context().init_context(config.q_lora_rank, is_dsa=True)
self._routed_experts_weights_of_layer = LazyValue(
lambda: {
layer_id: self.model.layers[layer_id].mlp.get_moe_weights()
for layer_id in range(self.model.start_layer, self.model.end_layer)
if isinstance(
self.model.layers[layer_id].mlp, deepseek_v2.DeepseekV2MoE
)
}
)
# Expose start_layer/end_layer for model_runner PP support
self.start_layer = self.model.start_layer
self.end_layer = self.model.end_layer
# update_weights_from_disk/_tensor/_distributed re-enter load_weights
# mid-serving (RL refit sends many partial batches); the prewarm and
# its barrier must only run on the first (startup) load.
self._mhc_prewarmed_at_load = False
@torch.inference_mode()
def wants_prefill_autotune(self) -> bool:
return getattr(self.config, "model_type", None) == "deepseek_v41"
def autotune_prefill_kernels(self, num_tokens: int, *, dtype: torch.dtype) -> int:
"""Tune resident MXFP8 linears for every M bucket up to ``num_tokens``.
The quant method is called directly, so no TP collectives run and no
request/KV/draft state is touched; the runner owns the autotune context."""
if getattr(self.config, "model_type", None) != "deepseek_v41":
return 0
seen = set()
# The backbone excludes vision and lm_head, whose prefill shapes differ.
for layer in self.model.modules():
method = getattr(layer, "quant_method", None)
if not isinstance(method, Fp8LinearMethod):
continue
if not (method.use_mxfp8 or method.block_fp8_as_mxfp8):
continue
if method.block_fp8_as_mxfp8 and not getattr(
layer, "block_fp8_mxfp8_ready", False
):
# No swizzled MXFP8 scale buffer: these kept the block-FP8 fallback.
continue
backend = method.mxfp8_dense_backend
if backend is None or not backend.is_flashinfer_cutedsl():
continue
if method.block_fp8_as_mxfp8:
# Small shapes and deterministic execution keep their pinned tactic.
method.mxfp8_prefill_autotune_min_tokens = 4096
weight = layer.weight
scale = layer.weight_scale_inv_swizzled
key = (
weight.shape,
weight.stride(),
weight.dtype,
scale.shape,
scale.stride(),
scale.dtype,
)
if key in seen:
continue
seen.add(key)
x = torch.zeros(
(num_tokens, weight.shape[1]),
dtype=dtype,
device=weight.device,
)
method.apply(layer, x)
del x
if seen:
logger.info(
"FlashInfer prefill autotune: %d MXFP8 weight layouts at M=%d.",
len(seen),
num_tokens,
)
return len(seen)
@property
def routed_experts_weights_of_layer(self):
return self._routed_experts_weights_of_layer.value
def pad_input_ids(self, input_ids, mm_inputs):
return MultiModalityDataPaddingPatternMultimodalTokens().pad_input_tokens(
input_ids, mm_inputs
)
def get_image_feature(self, items):
"""Return complete spans for the shared MM cache and chunk scheduler."""
spans = []
device, dtype = self.image_start.device, self.image_start.dtype
for item in items:
item.reconstruct(device.index, ipc_consumer_count=self.tp_size)
h, w = int(item.n_vit_h), int(item.n_vit_w)
pixels = torch.as_tensor(item.feature, device=device)
plan = item.model_specific_data.get(GPU_PLAN_KEY)
patches = (
materialize_image_gpu(pixels, plan).to(dtype)
if plan is not None
else pixels.to(dtype)
)
features = self.aligner(self.vision(patches, h, w), h, w)
r = self.config.vision_downsample_ratio
types = image_token_types((h + r - 1) // r, (w + r - 1) // r).to(device)
span = torch.empty(
(len(types), self.config.hidden_size), device=device, dtype=dtype
)
span[types == 0] = self.image_start
span[types == 1] = features.to(dtype)
span[types == 2] = self.image_newline
span[types == 3] = self.image_end
spans.append(span)
return spans
def _prepare_mm_embeddings(self, input_ids, forward_batch):
# Keep scheduler hash IDs intact: the shared embedder clamps its input in place.
input_embeds, _ = embed_mm_inputs(
mm_inputs_list=[
item if item is not None else MultimodalInputs(mm_items=[])
for item in forward_batch.mm_inputs
],
extend_prefix_lens=forward_batch.extend_prefix_lens_cpu,
extend_seq_lens=forward_batch.extend_seq_lens_cpu,
input_ids=input_ids.clone(),
input_embedding=self.get_input_embeddings(),
multimodal_model=self,
)
forward_batch.mm_input_embeds = input_embeds
return input_embeds
def get_input_embeddings(self) -> nn.Module:
return self.model.get_input_embeddings()
def set_dspark_layers_to_capture(self, layer_ids: List[int]) -> None:
if not self.pp_group.is_last_rank:
return
if layer_ids is None:
raise ValueError(
"DSPARK requires explicit layer_ids for aux hidden capture."
)
self.capture_aux_hidden_states = True
self.model.dspark_layers_to_capture = list(layer_ids)
@classmethod
def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
"""V4 only fuses when explicitly asked to, and then the checkpoint must
carry exactly one shared expert. Asked by the loader before any layer is
built."""
# Need to disable if quant precision mismatch, even if
# --enforce-shared-experts-fusion is specified
if quant_blocks_shared_experts_fusion(quant_config):
return (
"Quantization keeps shared experts at a higher precision than the "
"routed experts, so they cannot be fused into the quantized "
"routed-expert path."
)
if get_parallel().moe_ep_size > 1 and not uses_per_rank_fused_shared_slots():
return (
"Expert parallelism keeps only a slice of the routed experts on "
"each rank, so the fused shared expert cannot be appended to the "
"routed weight tensor (only DeepEP/MegaMOE per-rank shared slots "
"support fusion under EP)."
)
if not get_exec().moe.enforce_shared_experts_fusion:
return "Config does not support fused shared expert(s)."
if hf_config.n_shared_experts != 1:
raise ValueError(
"DeepSeek V4 shared-experts fusion expects exactly one shared "
f"expert, but got n_shared_experts={hf_config.n_shared_experts}."
)
return None
def determine_num_fused_shared_experts(self):
# The decision was installed by the loader; this only reads it.
self.num_fused_shared_experts = (
0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
)
def prepare_language_model_inputs(
self,
input_ids: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
"""Prepare full-sequence image embeddings and model IDs before CP splits.
Scheduler hash IDs stay intact for multimodal cache keys; the language
model uses image_token_id for Engram masking and visual MoE routing.
"""
if (
getattr(self, "vision", None) is not None
and not forward_batch.forward_mode.is_decode()
and not forward_batch.forward_mode.is_target_verify()
and forward_batch.mm_inputs is not None
and any(x is not None for x in forward_batch.mm_inputs)
):
if input_embeds is not None:
raise ValueError("Cannot combine input_embeds and image inputs")
input_embeds = self._prepare_mm_embeddings(input_ids, forward_batch)
if getattr(self, "vision", None) is not None and not (
forward_batch.forward_mode.is_decode_or_idle()
or forward_batch.forward_mode.is_target_verify()
):
# Decode/verify IDs are already vocabulary IDs; remap prompt image
# hashes for Engram and routing.
input_ids = input_ids.masked_fill(
input_ids >= MM_PAD_SHIFT_VALUE, self.config.image_token_id
)
return input_ids, input_embeds
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds: Optional[torch.Tensor] = None,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> torch.Tensor:
input_ids, input_embeds = self.prepare_language_model_inputs(
input_ids, forward_batch, input_embeds
)
with get_attn_tp_context().maybe_input_scattered(forward_batch):
hidden_states = self.model.forward(
input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
)
if not self.pp_group.is_last_rank:
return hidden_states
aux_hidden_states = None
if self.capture_aux_hidden_states:
hidden_states, aux_hidden_states = hidden_states
hidden_states, pre_hc_head = hidden_states
logits_metadata = forward_batch
tail = None
if (
self.capture_aux_hidden_states
and self.model.late_layer_start is not None
and forward_batch.forward_mode.is_extend_without_speculative()
):
tail = get_attn_backend().tail_forward_metadata.late_layer_tail
input_ids = tail.rows(input_ids)
logits_metadata = LogitsMetadata.from_forward_batch(forward_batch)
logits_metadata.extend_seq_lens = tail.extend_seq_lens
logits_metadata.extend_seq_lens_cpu = tail.extend_seq_lens_cpu
logits_metadata.extend_logprob_start_lens_cpu = tail.extend_seq_lens_cpu
output = self.logits_processor(
input_ids,
hidden_states,
self.lm_head,
logits_metadata,
aux_hidden_states,
hidden_states_before_norm=(
None if aux_hidden_states is not None else pre_hc_head
),
)
if tail is not None:
output.hidden_states_token_indices = tail.token_indices
return output
def _setup_fp8_wo_a_scales(self, is_nextn: bool) -> None:
if _FP8_WO_A_UE8M0:
from deep_gemm import transform_sf_into_required_layout
if is_nextn:
layers = [self.model.decoder]
else:
layers = [
self.model.layers[layer_id]
for layer_id in range(self.model.start_layer, self.model.end_layer)
]
for layer in layers:
attn = layer.self_attn
G = attn.n_local_groups
R = attn.o_lora_rank
D = attn.wo_a.weight.shape[1]
if _wo_a_weight_scale_to_e8m0 is not None:
# ROCm: aiter's mxscale GEMM reads uint8 e8m0 block scales, and
# requantizes the weight when the checkpoint's scales are not
# already powers of two. It also needs the weight row-major, so
# check the linear method honoured keep_plain_weight_layout: a
# preshuffled weight has the same shape, dtype and strides and
# would only show up as garbage output.
assert not getattr(attn.wo_a, "aiter_bpreshuffled", False), (
"DSV4 wo_a was B-preshuffled by the fp8 linear method; the "
"aiter mxscale absorb GEMM needs the row-major weight"
)
weight, scale = _wo_a_weight_scale_to_e8m0(
attn.wo_a.weight.data,
attn.wo_a.weight_scale_inv.data,
G,
R,
)
attn.wo_a.weight.data = weight.view(G * R, D)
attn.wo_a.weight_scale_inv.data = scale
attn.wo_a.weight_scale_inv.format_ue8m0 = True
continue
raw_scale = attn.wo_a.weight_scale_inv.data.view(G, R // 128, D // 128)
if _FP8_WO_A_UE8M0:
attn.wo_a.weight_scale_inv.data = transform_sf_into_required_layout(
raw_scale,
mn=R,
k=D,
recipe=(1, 128, 128),
num_groups=G,
is_sfa=False,
)
attn.wo_a.weight_scale_inv.format_ue8m0 = True
else:
attn.wo_a.weight_scale_inv.data = raw_scale.contiguous()
attn.wo_a.weight_scale_inv.format_ue8m0 = False
def post_load_weights(self, is_nextn=False, weight_names=None):
if self.wo_a_fp8:
self._setup_fp8_wo_a_scales(is_nextn)
if is_nextn:
return
for layer_id in range(self.model.start_layer, self.model.end_layer):
layer = self.model.layers[layer_id]
self_attn = layer.self_attn
if (
self_attn.compress_ratio in (4, 128)
and not self_attn.compressor.ape_converted
):
self_attn.compressor.apply_ape_hotfix()
if (
self_attn.compress_ratio == 4
and not self_attn.indexer.compressor.ape_converted
):
self_attn.indexer.compressor.apply_ape_hotfix()
layer.refresh_mhc_norm_weight_cache()
@staticmethod
def remap_weight_name_to_dpsk_hf_format(
name: str,
is_nextn: bool = False,
num_hidden_layers: Optional[int] = None,
) -> str:
if name.startswith("vision."):
return name.replace(".attn.wqkv.", ".attn.qkv_proj.").replace(
".attn.wo.", ".attn.proj."
)
if name.startswith(("aligner.", "image_")):
return name
if name.startswith("embed."):
return "model.embed_tokens." + name.removeprefix("embed.")
if name.startswith("head."):
return "lm_head." + name.removeprefix("head.")
if name == "norm.weight":
return "model.norm.weight"
if name.startswith("hc_head_"):
return "model." + name
if is_nextn and name.startswith("mtp."):
parts = name.split(".", 2)
if len(parts) >= 3:
rest = parts[2]
nextn_spec_prefixes = [
"e_proj",
"h_proj",
"emb",
"enorm",
"hnorm",
"norm",
"head",
"hc_head",
]
is_nextn_spec = any(rest.startswith(p) for p in nextn_spec_prefixes)
if is_nextn_spec:
if rest.startswith("emb.tok_emb"):
rest = rest.replace("emb.tok_emb", "embed_tokens")
elif rest == "norm.weight":
rest = "shared_head.norm.weight"
elif rest.startswith("head."):
rest = "shared_head.head.weight"
elif rest == "e_proj.scale":
rest = "e_proj.weight_scale_inv"
elif rest == "h_proj.scale":
rest = "h_proj.weight_scale_inv"
name = f"model.layers.{num_hidden_layers}." + rest
if name.startswith("layers."):
name = "model." + name
name = name.replace(".attn.", ".self_attn.")
name = name.replace(".ffn.", ".mlp.")
name = name.replace(".attn_norm.", ".input_layernorm.")
name = name.replace(".ffn_norm.", ".post_attention_layernorm.")
if "self_attn" in name and name.endswith(".scale"):
name = name.removesuffix(".scale") + ".weight_scale_inv"
if ".engram.wkv." in name and name.endswith(".scale"):
name = name.removesuffix(".scale") + ".weight_scale_inv"
name = name.replace(".gate.tid2eid", ".topk.tid2eid")
name = name.replace(".gate.bias", ".gate.e_score_correction_bias")
name = name.replace(".w1.", ".gate_proj.")
name = name.replace(".w2.", ".down_proj.")
name = name.replace(".w3.", ".up_proj.")
if "mlp" in name and name.endswith(".scale"):
name = name.removesuffix(".scale") + ".weight_scale_inv"
return name
def _prewarm_mhc_kernels(self) -> None:
"""One-shot MHC JIT prewarm at load time, synced across ranks.
Runs before any forward so the compile burst stays off the serving
path; the barrier keeps ranks from proceeding while a peer is still
compiling. The early returns below must stay rank-uniform.
"""
if self._mhc_prewarmed_at_load:
return
self._mhc_prewarmed_at_load = True
if _is_npu or _is_xpu or not envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get():
return
layer = next(
(m for m in self.model.layers if isinstance(m, DeepseekV4DecoderLayer)),
None,
)
if layer is None:
return
from sglang.kernels.ops.layernorm.mhc import mhc_post, prewarm_mhc_pre
tic = time.perf_counter()
residual = torch.zeros(
(1, layer.hc_mult, layer.hidden_size),
dtype=torch.bfloat16,
device=layer.hc_attn_fn.device,
)
prewarm_mhc_pre(
# Template carrying dtype/device; buckets allocate their own sizes.
residual=residual,
fn=layer.hc_attn_fn,
hc_scale=layer.hc_attn_scale,
hc_base=layer.hc_attn_base,
rms_eps=layer.rms_norm_eps,
hc_pre_eps=layer.hc_eps,
hc_sinkhorn_eps=layer.hc_eps,
hc_post_mult_value=_MHC_POST_MULT_VALUE,
sinkhorn_repeat=layer.hc_sinkhorn_iters,
n_splits=1,
n_splits_pre=32,
norm_weight=layer.input_layernorm.weight.data,
norm_eps=layer.input_layernorm.variance_epsilon,
)
mhc_post(
x=residual.new_zeros((1, layer.hidden_size)),
residual=residual,
post_layer_mix=torch.zeros(
(1, layer.hc_mult, 1),
dtype=torch.float32,
device=residual.device,
),
comb_res_mix=torch.zeros(
(1, layer.hc_mult, layer.hc_mult),
dtype=torch.float32,
device=residual.device,
),
)
torch.cuda.synchronize()
compile_secs = time.perf_counter() - tic
# Runs before init_memory_pool(); don't let transients skew pool sizing.
torch.cuda.empty_cache()
get_parallel().tp_group.barrier()
logger.info(
"DeepSeek V4 MHC prewarm at load: compile %.1fs, rank sync +%.1fs",
compile_secs,
time.perf_counter() - tic - compile_secs,
)
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False):
params_dict = dict(self.named_parameters())
loaded_params: Set[str] = set()
if is_nextn:
if hasattr(self.config, "num_nextn_predict_layers"):
num_nextn_layers = self.config.num_nextn_predict_layers
assert num_nextn_layers == 1, "Only 1 nextn layer is supported"
nextn_layer_id = (
0
if self.config.num_hidden_layers == 1
else self.config.num_hidden_layers
)
else:
raise ValueError("num_nextn_predict_layers is not in the config")
# Must mirror MQALayer.__init__'s `quantize_wo_a`: dequantizing wo_a here
# while the layer allocated an FP8 parameter (or vice versa) fails the
# weight loader's dtype check.
if not (self.wo_a_fp8 or use_npu_arch35_mxfp8_wo_a(self.quant_config)):
weights = _prepare_deepseek_v4_weights(weights, self.quant_config)
stacked_params_mapping = DEEPSEEK_V4_STACKED_PARAMS_MAPPING
expert_params_mapping = FusedMoE.make_expert_params_mapping(
ckpt_gate_proj_name="gate_proj",
ckpt_down_proj_name="down_proj",
ckpt_up_proj_name="up_proj",
num_experts=self.config.n_routed_experts + self.num_fused_shared_experts,
)
if is_wint4afp8_or_wint4a16_config(self.quant_config):
expert_params_mapping += FusedMoE.make_expert_input_scale_params_mapping(
num_experts=self.config.n_routed_experts
)
cache_compressor_weight = {}
COMPRESSOR_PART = ".compressor.w"
fuse_wqa_wkv = envs.SGLANG_OPT_FUSE_WQA_WKV.get()
cache_wqkv_a_weight: dict[str, dict[str, torch.Tensor]] = {}
skipped_by_group: dict[str, int] = {}
# V4 checkpoints must load every compressor / indexer tensor.
is_dsv41 = getattr(self.config, "model_type", None) == "deepseek_v41"
def auto_weight_loader(module):
return getattr(module, "weight_loader", default_weight_loader)
if is_nextn:
nextn_layer_prefix = f"model.layers.{nextn_layer_id}"
nextn_spec_weight_names_out_of_layer = [
"shared_head.norm",
"shared_head.head",
"embed_tokens",
".e_proj",
"h_proj",
"enorm",
"hnorm",
"hc_head_base",
"hc_head_fn",
"hc_head_scale",
]
if self.num_fused_shared_experts > 0:
assert self.num_fused_shared_experts == 1
log_info_on_rank0(logger, "Shared experts fusion optimization enabled.")
with concurrent.futures.ThreadPoolExecutor() as executor:
futures = []
weight_names = []
for name, loaded_weight in weights:
if (
self.wo_a_fp8
and name.endswith(".wo_a.weight")
and loaded_weight.dtype != torch.float8_e4m3fn
):
raise ValueError(
f"SGLANG_OPT_FP8_WO_A_GEMM is enabled but {name} has "
f"dtype {loaded_weight.dtype}, expected "
"torch.float8_e4m3fn. This checkpoint does not provide "
"a supported fp8-quantized wo_a; rerun with "
"SGLANG_OPT_FP8_WO_A_GEMM=0."
)
try:
use_async_loading = should_async_load(loaded_weight)
name = self.remap_weight_name_to_dpsk_hf_format(
name,
is_nextn=is_nextn,
num_hidden_layers=self.config.num_hidden_layers,
)
# V4.1 checkpoint tensors with no module in the text model yet.
skip_group = None
if not is_dsv41:
pass
elif self.vision is None and name.startswith(
("vision.", "aligner.", "image_")
):
skip_group = "vision"
elif self.vision is None and name.endswith(
".gate.e_score_correction_bias_vl"
):
skip_group = "gate.bias_vl"
if skip_group is not None:
skipped_by_group[skip_group] = (
skipped_by_group.get(skip_group, 0) + 1
)
continue
layer_id = get_layer_id(name)
if (
layer_id is not None
and hasattr(self.model, "start_layer")
and (
layer_id < self.model.start_layer
or layer_id >= self.model.end_layer
)
):
continue
if (
self.num_fused_shared_experts > 0
and "mlp.shared_experts" in name
):
name = name.replace(
"mlp.shared_experts",
f"mlp.experts.{self.config.n_routed_experts}",
)
weight_names.append(name)
if not is_nextn:
if hasattr(self.config, "num_nextn_predict_layers"):
num_nextn_layers = self.config.num_nextn_predict_layers
if num_nextn_layers > 0 and name.startswith("model.layers"):
name_list = name.split(".")
if (
len(name_list) >= 3
and int(name_list[2])
>= self.config.num_hidden_layers
):
continue
if name.startswith("mtp"):
continue
else:
if "shared_head.head" in name or "embed_tokens" in name:
continue
if not name.startswith(nextn_layer_prefix):
continue
in_decoder = True
for weight_name in nextn_spec_weight_names_out_of_layer:
if weight_name in name:
in_decoder = False
name = name.replace(nextn_layer_prefix, "model")
break
if in_decoder:
name = name.replace(nextn_layer_prefix, "model.decoder")
if "rotary_emb.inv_freq" in name:
continue
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
if _is_npu:
name = name.replace("weight_packed", "weight")
if ("mlp.experts." in name) and name not in params_dict:
continue
name = name.replace(weight_name, param_name)
if name.endswith(".bias") and name not in params_dict:
continue
if name not in params_dict and name.startswith("mtp"):
break
param = params_dict[name]
weight_loader = param.weight_loader
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(param, loaded_weight, shard_id),
)
loaded_params.add(name)
break
else:
skip_unmaterialized_expert_param = False
for mapping in expert_params_mapping:
param_name, weight_name, expert_id, shard_id = mapping
if weight_name not in name:
continue
if _is_npu:
name = name.replace("weight_packed", "weight")
resolved_name = name.replace(weight_name, param_name)
if resolved_name not in params_dict:
skip_unmaterialized_expert_param = True
continue
param = params_dict[resolved_name]
weight_loader = param.weight_loader
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(
param,
loaded_weight,
resolved_name,
),
func_kwargs={
"shard_id": shard_id,
"expert_id": expert_id,
},
)
loaded_params.add(resolved_name)
break
else:
if skip_unmaterialized_expert_param:
continue
if name.endswith(".bias") and name not in params_dict:
continue
if (
".embed_tokens." in name
and not self.pp_group.is_first_rank
):
continue
if (
name == "model.norm.weight"
and not self.pp_group.is_last_rank
):
continue
if (
name.startswith("model.hc_head_")
or name == "lm_head.weight"
) and not self.pp_group.is_last_rank:
continue
elif (
COMPRESSOR_PART in name
and ".wkv_gate." not in name
and (name.rsplit(".", 2)[0] + ".wkv_gate.weight")
in params_dict
):
# Split-projection modules load per parameter instead.
is_kv = name.endswith(".wkv.weight")
is_wgate = name.endswith(".wgate.weight")
assert is_kv != is_wgate
key = name.rsplit(".", 2)[0]
assert key.endswith(".compressor")
if key not in cache_compressor_weight:
cache_compressor_weight[key] = (
is_kv,
_clone_if_runai_streamed_tensor(loaded_weight),
)
else:
assert key in cache_compressor_weight
cached_is_kv, cached_weight = (
cache_compressor_weight[key]
)
assert cached_is_kv != is_kv
kv = loaded_weight if is_kv else cached_weight
wgate = loaded_weight if is_wgate else cached_weight
fused_weight = torch.cat([kv, wgate], dim=0)
param_name = key + ".wkv_gate.weight"
param = params_dict[param_name]
weight_loader = auto_weight_loader(param)
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(param, fused_weight),
)
loaded_params.add(param_name)
cache_compressor_weight.pop(key)
elif (
fuse_wqa_wkv
and ".compressor." not in name
and ".engram." not in name
and (
name.endswith(".wq_a.weight")
or name.endswith(".wq_a.weight_scale_inv")
or name.endswith(".wkv.weight")
or name.endswith(".wkv.weight_scale_inv")
or name.endswith(".wq_a.qweight")
or name.endswith(".wkv.qweight")
or name.endswith(".wq_a.qweight_type")
or name.endswith(".wkv.qweight_type")
)
):
is_q = ".wq_a." in name
param_name = name.replace(
".wq_a." if is_q else ".wkv.", ".wqkv_a."
)
bucket = cache_wqkv_a_weight.setdefault(param_name, {})
shard_key = "q" if is_q else "kv"
assert shard_key not in bucket, (
f"duplicate shard {shard_key} for {param_name}"
)
bucket[shard_key] = _clone_if_runai_streamed_tensor(
loaded_weight
)
if len(bucket) == 2:
fused_weight = _fuse_deepseek_v4_wqkv_a_pair(
param_name, bucket
)
param = params_dict[param_name]
weight_loader = auto_weight_loader(param)
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(param, fused_weight),
)
loaded_params.add(param_name)
cache_wqkv_a_weight.pop(param_name)
else:
if (
"k_scale" in name or "v_scale" in name
) and name not in params_dict:
for scale in ["k_scale", "v_scale"]:
if scale in name:
name = name.replace(
f"{scale[0]}_proj", "attn_mqa"
)
break
if name not in params_dict:
if not name.startswith("mtp"):
logger.warning(
f"{name} not found in params_dict."
)
continue
param = params_dict[name]
weight_loader = auto_weight_loader(param)
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(param, loaded_weight),
)
loaded_params.add(name)
except Exception as e:
e.add_note(f"{name=} {loaded_weight.shape=}")
raise
for future in concurrent.futures.as_completed(futures):
future.result()
assert len(cache_compressor_weight) == 0
assert len(cache_wqkv_a_weight) == 0, cache_wqkv_a_weight.keys()
if skipped_by_group:
log_info_on_rank0(
logger,
"Skipped checkpoint tensors not wired yet: "
+ ", ".join(f"{k}={v}" for k, v in sorted(skipped_by_group.items())),
)
unloaded_params = params_dict.keys() - loaded_params
skipped_checking_patterns = [
"attn_mqa.k_scale",
"attn_mqa.v_scale",
"blockscale_swizzled",
]
if not self.pp_group.is_first_rank:
skipped_checking_patterns.append("embed_tokens")
if not self.pp_group.is_last_rank:
skipped_checking_patterns.append("model.norm.")
skipped_checking_patterns.extend(["lm_head", "hc_head_"])
if is_nextn:
skipped_checking_patterns.extend(["lm_head", "embed_tokens"])
unloaded_params = {
p
for p in unloaded_params
if all(
skipped_checking_pattern not in p
for skipped_checking_pattern in skipped_checking_patterns
)
}
if unloaded_params:
logger.warning(
f"Some weights are not initialized from checkpoints: {unloaded_params}"
)
self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names)
if not is_nextn:
for i, layer in enumerate(self.model.layers):
if getattr(layer, "engram", None) is not None:
layer.engram.embed.finish_load(label=f"layer {i}")
self._prewarm_mhc_kernels()
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
# Hot weight reload (RL workflows). Use the device-agnostic module
# accessor so this works on both CUDA/HIP and NPU.
torch.get_device_module().empty_cache()
torch.get_device_module().synchronize()
@classmethod
def get_model_config_for_expert_location(cls, config):
return ModelConfigForExpertLocation(
num_layers=config.num_hidden_layers,
num_logical_experts=config.n_routed_experts,
num_groups=None,
)
EntryClass = [DeepseekV4ForCausalLM]
def _dequant_fp8(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
from einops import rearrange
assert weight.dtype == torch.float8_e4m3fn, (
f"expected fp8_e4m3fn, got {weight.dtype}"
)
assert scale.dtype in (
torch.float8_e8m0fnu,
torch.float32,
), f"expected fp8_e8m0fnu or float32, got {scale.dtype}"
# Block size is per-checkpoint: V4 128x128, V4.1 32x32.
bn = weight.shape[0] // scale.shape[0]
bk = weight.shape[1] // scale.shape[1]
weight_f32 = rearrange(
weight.float(), "(sn bn) (sk bk) -> sn bn sk bk", bn=bn, bk=bk
)
result = rearrange(
weight_f32 * scale.float()[:, None, :, None], "sn bn sk bk -> (sn bn) (sk bk)"
)
return result.to(torch.bfloat16)
def _clone_if_runai_streamed_tensor(tensor: torch.Tensor) -> torch.Tensor:
if getattr(tensor, RUNAI_STREAMER_TENSOR_ATTR, False):
return tensor.clone().detach()
return tensor
def _dequant_fp8_wo_a_streaming(
weights: Iterable[Tuple[str, torch.Tensor]],
) -> Iterable[Tuple[str, torch.Tensor]]:
pending: dict[str, dict[str, torch.Tensor]] = {}
saw_wo_a_scale = False
emitted = False
for name, tensor in weights:
if name.endswith(".wo_a.weight"):
prefix = name[: -len(".weight")]
bucket = pending.setdefault(prefix, {})
scale = bucket.pop("scale", None)
if scale is not None:
pending.pop(prefix, None)
emitted = True
yield name, _dequant_fp8(tensor, scale)
else:
bucket["weight"] = _clone_if_runai_streamed_tensor(tensor)
continue
if name.endswith(".wo_a.scale"):
saw_wo_a_scale = True
prefix = name[: -len(".scale")]
bucket = pending.setdefault(prefix, {})
weight = bucket.pop("weight", None)
if weight is not None:
pending.pop(prefix, None)
emitted = True
yield prefix + ".weight", _dequant_fp8(weight, tensor)
else:
bucket["scale"] = _clone_if_runai_streamed_tensor(tensor)
continue
yield name, tensor
if emitted:
logger.info("Finished streaming dequant fp8 wo_a")
for prefix, bucket in pending.items():
if "weight" in bucket:
assert not saw_wo_a_scale, f"{prefix}.scale is missing"
yield prefix + ".weight", bucket["weight"]
if "scale" in bucket:
yield prefix + ".scale", bucket["scale"]
def _dequant_fp8_wo_a(
weights: Iterable[Tuple[str, torch.Tensor]],
) -> Iterable[Tuple[str, torch.Tensor]]:
weights_dict = dict(weights)
for name in list(weights_dict.keys()):
if name not in weights_dict:
continue
if not name.endswith(".wo_a.weight"):
continue
scale_name = name.replace(".wo_a.weight", ".wo_a.scale")
assert scale_name in weights_dict
weight = weights_dict.pop(name)
scale = weights_dict.pop(scale_name)
yield name, _dequant_fp8(weight, scale)
yield from weights_dict.items()
def _prepare_deepseek_v4_weights(
weights: Iterable[Tuple[str, torch.Tensor]],
quant_config: Optional[QuantizationConfig],
) -> Iterable[Tuple[str, torch.Tensor]]:
"""Keep Expert Pack GGUF weights on the streaming load path."""
if quant_config is not None and quant_config.get_name() == "expert_pack":
logger.info("Keep Expert Pack GGUF weights on the streaming load path")
return weights
return _dequant_fp8_wo_a_streaming(weights)
def _fuse_deepseek_v4_wqkv_a_pair(
param_name: str, bucket: dict[str, torch.Tensor]
) -> torch.Tensor:
"""Fuse Q/KV rows while preserving their common GGUF type scalar."""
q = bucket["q"]
kv = bucket["kv"]
if param_name.endswith(".qweight_type"):
if q.numel() != 1 or kv.numel() != 1 or q.item() != kv.item():
raise ValueError(
f"cannot fuse different GGUF qweight types for {param_name}: "
f"q={q.tolist()} kv={kv.tolist()}"
)
return q
return torch.cat([q, kv], dim=0)
Reference in New Issue View Git Blame Copy Permalink