from __future__ import annotations

import concurrent.futures
import functools
import logging
import time
from array import array
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
from sglang.kernels.ops.attention.flash_mla_sm120 import SM120_DECODE_MAX_TOKENS
from sglang.kernels.ops.gemm.dsv4_wo_a import (
    wo_a_bf16_gemv,
    wo_a_bf16_small_batch,
    wo_a_bf16_small_batch_mxfp8,
)
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 (
    use_npu_arch35_mxfp8_wo_a,
)
from sglang.srt.layers.attention.dsa.utils import (
    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.cp.cp_decode_attn_tp import get_cp_decode_attn_tp_ctx
from sglang.srt.layers.cp.interleave import attn_cp_interleave_gather
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_gather_into_tensor,
    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.layer_boundary import get_attn_tp_context
from sglang.srt.layers.layer_boundary.ops import attn_cp_interleave_reduce_scatter
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,
    should_skip_post_experts_all_reduce,
    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.mxfp4_flashinfer_trtllm_moe import (
    can_fuse_all_reduce,
)
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 import deepseek_v4_mhc as mhc
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.decoder_swa_bounded_replay import (
    enter_late_layer_dp,
    exit_late_layer_dp,
    late_layer_dp_counts,
)
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.models.utils import WeightsMapper
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_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_pp_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
_hip = None
if _is_hip:
    from sglang.srt.models.deepseek_common.amd import deepseek_v4_hip as _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,
    fp8_grid: bool = False,
    emit_fp8: bool = False,
) -> torch.Tensor | Mxfp8SwizzledInput | Fp8GridActivation | Mxfp8Activation:
    # 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
    # the gfx950 routes, like the CUDA ones, are V4.1's (fast_path); DSv4 keeps the
    # aiter batched GEMM / einsum below
    hip_fast_path = fast_path and _is_hip and _is_gfx95_supported
    hip_decode_verify = (
        hip_fast_path
        and (
            (is_decode and o.shape[0] == 1 and o.is_contiguous())
            or (is_target_verify and 2 <= o.shape[0] <= 384)
        )
        and _hip.wo_a_split_k_allowed()
    )
    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
                )
            )
        )
        or hip_decode_verify
        or (hip_fast_path and is_prefill and 4096 <= o.shape[0] <= 65536)
    ) 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
    # the 16-row decode kernels also serve V4.1's target-verify rows
    if (is_decode or is_target_verify) and hip_fast_path:
        y = _hip.wo_a_fp8_grid_matmul(o, wo_a, fp8_grid, emit_fp8=emit_fp8)
        if y is not None:
            return y
    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)


def _materialize_cp_unified_fp8_kv(
    k_nope: torch.Tensor,
    k_rope: torch.Tensor,
    forward_batch: ForwardBatch,
    stream: Optional[Any],
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Gather a two-pool unified-KV pair in global logical-token order.

    E8M0 scales live inline in the noPE bytes, so both halves travel as one raw
    byte row and the CP permutation cannot separate scales, noPE and RoPE.
    """
    nope_bytes = k_nope.view(torch.uint8)
    rope_bytes = k_rope.view(torch.uint8)
    packed = torch.cat((nope_bytes, rope_bytes), dim=-1)
    packed = cp_materialize_global_token_order(packed, forward_batch, stream)
    nope, rope = packed.split((nope_bytes.shape[-1], rope_bytes.shape[-1]), dim=-1)
    return nope.contiguous().view(k_nope.dtype), rope.contiguous().view(k_rope.dtype)


if TYPE_CHECKING:
    from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import (
        Fp8GridActivation,
        Mxfp8Activation,
    )
    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=(
            _hip.live_rows(q_lora, real_num_tokens)
            if _is_hip
            else 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)


def deepseek_v4_enter_late_layer_tail(sentinel: torch.Tensor) -> torch.Tensor:
    # Decoder SWA bounded replay under the prefill CUDA graph: swap the attention
    # backend onto the tail metadata of the *live* batch (the captured args hold
    # the capture-time batch). The sentinel keeps the break in the replay sequence.
    forward_batch = get_tc_piecewise_forward_context().forward_batch
    get_attn_backend().prefill_graph_enter_late_layer_tail(forward_batch)
    return sentinel


def deepseek_v4_exit_late_layer_tail(sentinel: torch.Tensor) -> torch.Tensor:
    forward_batch = get_tc_piecewise_forward_context().forward_batch
    get_attn_backend().prefill_graph_exit_late_layer_tail(forward_batch)
    return sentinel


bcg_deepseek_v4_enter_late_layer_tail = eager_on_graph(True)(
    deepseek_v4_enter_late_layer_tail
)
bcg_deepseek_v4_exit_late_layer_tail = eager_on_graph(True)(
    deepseek_v4_exit_late_layer_tail
)


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,
        *,
        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.attn_tp_rank: int = get_parallel().attn_tp_rank
        self.attn_tp_size: int = get_parallel().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),
            parallel_group="attn_tp",
        )
        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),
            parallel_group="attn_tp",
            **({} 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),
            parallel_group="attn_tp",
        )

        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 (
            (
                # the gfx950 overlap is validated on V4.1 only; DSv4 on ROCm keeps one stream
                (_is_cuda or (_is_gfx95_supported and self.is_dsv41))
                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 _hip.use_fused_qk_norm_rope(self)
        # 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)
        )
        if _is_hip:
            _hip.init_mqa_layer(self, quant_config)

        # 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.
        if _is_hip:
            return _hip.q_norm_for_wq_b(self, q)
        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 q_out is not None and q_out.dtype == torch.float8_e4m3fn:
                fused_q_norm_rope(q, q_out, None, self.freqs_cis, positions)
                return q_out
            # TODO: enable zero-copy path to skip the extra no-rope copy overhead
            if (_is_cuda or _is_gfx95_supported) and q_out is not None:
                fused_q_norm_rope(
                    q,
                    q_out,
                    None,  # eps = None means no norm
                    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,
        q_rope: 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).
        q_rope ([T, H, head_dim]) has its query heads roped by the same launch.
        """
        if envs.SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE.get():
            assert q_rope is None
            # 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,
            # the query RoPE fused into this store is HIP only
            **({} if q_rope is None else {"q": q_rope}),
        )

    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, q_for_wqb = self._normalize_q_lora(q_lora)
            q = self._compute_q_b(q_for_wqb, 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 = is_cp_active(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,
                )
                if use_cp:
                    # Store globally ordered KV after the CP gather.
                    swa_cache, swa_loc = None, None

            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 use_cp:
                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 not (unified and fuse_verify):
                kv = None
        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)
            if _is_hip:
                q, q_lora, fuse_q_rope = _hip.compute_q_b(
                    self, q_lora, q_for_wqb, positions, q_out, unified, use_cp
                )
            else:
                q = self._compute_q_b(q_for_wqb, positions, q_out)
                fuse_q_rope = False
            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)
                if use_cp:
                    # 2-source attention reads the whole logical chunk.
                    kv = cp_materialize_global_token_order(
                        kv.contiguous(),
                        forward_batch,
                        torch.cuda.current_stream(),
                    )
            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,
                    q_rope=q if fuse_q_rope else None,
                )
                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 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: Optional[torch.Tensor] = None,
        defer_all_reduce: bool = False,
    ) -> Union[torch.Tensor, mhc.AttnOutput]:
        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 is_cp_active(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) or _is_gfx95_supported)
            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 (
                        _is_gfx95_supported
                        or 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."
            )

        tp_slice, q_padded, q_out, q_rope = slice(None), None, None, None
        k_nope, k_rope = None, None
        if (
            self.is_dsv41
            and not self.q_head_norm
            and get_token_to_kv_pool().uniform_fp8
            and self.n_local_heads in (8, 16, 32, 64, 128)
        ):
            kernel_num_heads = self.n_local_heads
            q_out = torch.empty(
                (x.shape[0], self.n_local_heads, self.head_dim),
                dtype=torch.float8_e4m3fn,
                device=x.device,
            )
        elif 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.n_local_heads
                if _is_hip and not unified and _hip.skip_head_pad(self)
                else 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,
            )

        if unified_fp8_prefill and is_cp_active(forward_batch):
            # kv is the packed noPE half the fused prefill store left on the kv slot.
            kv, k_rope = _materialize_cp_unified_fp8_kv(
                kv, k_rope, forward_batch, torch.cuda.current_stream()
            )

        # 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

        # HIP: the attention kernel may apply the inverse RoPE itself; the else-branch below keeps it
        inv_rope = (
            _hip.attention_inv_rope(
                self,
                positions,
                forward_batch,
                unified,
                wo_a_applies_inv_rope=self.wo_a_fp8
                and _wo_a_fp8_mxscale_fused_invrope is not None,
            )
            if _is_hip
            else None
        )
        if unified:
            # only the HIP radix backend takes these; 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
            if inv_rope is not None:
                rope_kwargs["inv_rope"] = inv_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,
                    # HIP only, for the same reason as rope_kwargs above
                    **({} if inv_rope is None else {"inv_rope": inv_rope}),
                )
            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 inv_rope is None and 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,
                            fp8_grid=_is_hip and _hip.wo_a_emits_fp8_grid(self),
                            emit_fp8=_is_hip
                            and _hip.wo_b_emits_mxfp8(self, o.shape[0]),
                        )
                else:
                    o = _apply_gguf_grouped_wo_a(
                        o,
                        self.wo_a.qweight,
                        self.wo_a.qweight_type.weight_type,
                        self.o_lora_rank,
                    )

        return self._project_wo_b(o, defer_all_reduce)

    def _project_wo_b(self, o, defer_all_reduce: bool = False):
        from sglang.srt.layers.moe.utils import should_skip_mlp_all_reduce

        defer_all_reduce = defer_all_reduce and (
            self.attn_tp_size > 1
            and self.attn_tp_size == get_parallel().tp_size
            and self.wo_b.reduce_results
            and not self.wo_b.use_decode_attn_tp
            and not should_skip_mlp_all_reduce()
        )
        o, _ = self.wo_b(
            o.flatten(1) if isinstance(o, torch.Tensor) else o,
            skip_all_reduce=defer_all_reduce,
        )
        if defer_all_reduce:
            return mhc.AttnOutput(o)
        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


def _attn_tp_all_gather_chunks(
    local: torch.Tensor, chunk_rows: List[int]
) -> torch.Tensor:
    # Tensor all_gather stays on aiter AG / pynccl under graph capture; the list
    # all_gather falls back to ProcessGroupNCCL, which invalidates HIP capture.
    max_rows = max(chunk_rows)
    padded = local
    if local.shape[0] != max_rows:
        padded = local.new_zeros((max_rows, *local.shape[1:]))
        padded[: local.shape[0]].copy_(local)
    gathered = local.new_empty((len(chunk_rows) * max_rows, *local.shape[1:]))
    attn_tp_all_gather_into_tensor(gathered, padded.contiguous())
    if all(rows == max_rows for rows in chunk_rows):
        return gathered
    return torch.cat(
        [
            gathered[i * max_rows : i * max_rows + rows]
            for i, rows in enumerate(chunk_rows)
        ]
    )


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
        if _is_hip:
            _hip.init_decoder_layer(self, quant_config)
        self.hc_cfg = mhc.HcConfig(
            mult=hc_mult,
            sinkhorn_iters=config.hc_sinkhorn_iters,
            eps=config.hc_eps,
            rms_eps=config.rms_norm_eps,
            hidden=config.hidden_size,
            pre_from_prev=self.hc_pre_from_prev_sublayer,
            cp_prefill=self.dsa_enable_prefill_cp,
        )
        self.attn_hc: Optional[mhc.HcSubLayer] = None
        self.ffn_hc: Optional[mhc.HcSubLayer] = None
        self._hc_attn_tf32_parts = self._hc_ffn_tf32_parts = None
        self._hc_attn_bf16_parts = self._hc_ffn_bf16_parts = None
        self._init_hyper_connections()
        self._next_layer: Optional[DeepseekV4DecoderLayer] = None
        self.local_boundary: Optional[mhc.HcNextBoundary] = None
        self.next_boundary: Optional[mhc.HcNextBoundary] = None
        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],
    ) -> MQALayer:
        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 bind_next(self, next_layer: Optional[nn.Module]) -> None:
        """Save the next layer for mhc. Must be called before `_init_boundaries`"""
        if not self.hc_pre_from_prev_sublayer:
            return
        self._next_layer = None
        if isinstance(next_layer, DeepseekV4DecoderLayer) and next_layer.engram is None:
            self._next_layer = next_layer

    def _init_hyper_connections(self) -> None:
        """(Re)bind the two hyper-connections to the current mixing weights."""
        if not self.hc_pre_from_prev_sublayer:
            return
        self.attn_hc = mhc.HcSubLayer(
            self.hc_cfg,
            self.hc_attn_fn,
            self.hc_attn_scale,
            self.hc_attn_base,
            self.input_layernorm,
            self._hc_attn_tf32_parts,
            self._hc_attn_bf16_parts,
        )
        self.ffn_hc = mhc.HcSubLayer(
            self.hc_cfg,
            self.hc_ffn_fn,
            self.hc_ffn_scale,
            self.hc_ffn_base,
            self.post_attention_layernorm,
            self._hc_ffn_tf32_parts,
            self._hc_ffn_bf16_parts,
        )

    def _init_boundaries(self) -> None:
        """Resolved on first forward: `accepts_mxfp8_swizzled_input` is not settled
        until `process_weights_after_loading`."""
        if self.local_boundary is not None:
            return
        # ATTN[i] -> MOE[i]
        self.local_boundary = mhc.make_boundary(
            self.post_attention_layernorm,
            accepts_mxfp8=False,
            hc=self.ffn_hc,
        )
        # MOE[i] -> ATTN[i+1]
        if self._next_layer is not None:
            self.next_boundary = mhc.make_boundary(
                self._next_layer.input_layernorm,
                accepts_mxfp8=self._next_layer.self_attn.accepts_mxfp8_swizzled_input(),
                hc=self._next_layer.attn_hc,
            )

    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 or get_platform().is_sm90)
            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_bf16_hc_weight

            if get_platform().is_sm90:
                # Hopper's compensated projection does not require DeepGEMM.
                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)
            else:
                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 callable(getattr(deep_gemm, "tf32_hc_prenorm_gemm", None)):
                        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
                        )
                        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
                        )
        if self.hc_pre_from_prev_sublayer:
            self._init_hyper_connections()
        # The fuse gates and boundaries snapshot load-time facts; weight updates
        # re-run this, so drop them and re-resolve on the next forward.
        self.__dict__.pop("_can_fuse_attn_mhc", None)
        self.__dict__.pop("_can_fuse_ffn_mhc", None)
        self.local_boundary = None
        self.next_boundary = None

    def _input_norm(
        self,
        hidden_states: torch.Tensor,
        allow_aiter_quant: bool = True,
        coefficients=None,
    ) -> Tuple[torch.Tensor, Optional[Tuple]]:
        """input_layernorm(hidden_states) as (the bf16 norm attention reads, the pre-quantized
        operand of its dense projections or None); coefficients is a ROCm boundary's pending
        reduce + sinkhorn."""
        if _is_hip:
            return _hip.input_norm(
                self,
                hidden_states,
                allow_aiter_quant,
                coefficients,
                _fused_rmsnorm_fp8_quant,
            )
        return self.input_layernorm(hidden_states), None

    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,
    ) -> torch.Tensor:
        return mhc.post(self.hc_cfg, 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.
                    hidden_states, x_quant = self._input_norm(
                        hidden_states, allow_aiter_quant=_is_gfx95_supported
                    )
                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:
                    # gfx1250 takes plain layernorm here
                    hidden_states, x_quant = self._input_norm(
                        hidden_states, allow_aiter_quant=_is_gfx95_supported
                    )
                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:
                # gfx1250 takes plain layernorm here
                hidden_states, x_quant = self._input_norm(
                    hidden_states, allow_aiter_quant=_is_gfx95_supported
                )
            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

    @functools.cached_property
    def _can_fuse_attn_mhc(self) -> bool:
        """The static half of the attention fuse gate (dsv4.1 only); load-time
        facts, dropped by `refresh_mhc_norm_weight_cache`. Row capacity is the
        dynamic half, checked per forward via `can_fuse_all_reduce`."""
        return (
            get_parallel().tp_size == self.self_attn.attn_tp_size == 4
            and self.self_attn.wo_b.reduce_results
            and self.local_boundary.norm_fusable
            and mhc.can_fuse_post(self.hc_cfg)
        )

    @functools.cached_property
    def _can_fuse_ffn_mhc(self) -> bool:
        """The MoE counterpart of `_can_fuse_attn_mhc`."""
        return (
            self.mlp.tp_size == 4
            and get_moe_a2a_backend().is_none()
            and not self.mlp._shared_expert_tp1
            and not should_skip_post_experts_all_reduce(is_tp_path=True)
            and mhc.can_fuse_post(self.hc_cfg)
        )

    def forward_hc_pre_from_prev(
        self,
        positions: torch.Tensor,
        state: mhc.HcState,
        input_ids: torch.Tensor,
        forward_batch: ForwardBatch,
        input_ids_global: torch.Tensor,
    ) -> mhc.HcState:
        """The layer's two hyper-connections, each collapsing with the previous one's
        pre-mix."""
        self._init_boundaries()
        stats_stream = None
        if mhc.use_stats_stream(self.hc_cfg, forward_batch, state.residual):
            stats_stream = self.hc_stats_stream

        def run_attn_hc(state: mhc.HcState) -> mhc.HcState:
            assert self.attn_hc is not None
            residual = state.residual
            quantized = [] if self.self_attn.accepts_mxfp8_swizzled_input() else None
            mhc.fork_stats_stream(stats_stream)
            x = mhc.combine(self.attn_hc, state, quantized)
            precomputed = state.stats
            state.release()
            del state
            world_size = self.self_attn.attn_tp_size
            fuse_all_reduce_mhc = self._can_fuse_attn_mhc and can_fuse_all_reduce(
                x.shape[0], self.hc_cfg.hidden
            )
            with self.self_attn.maybe_use_decode_attn_tp(forward_batch):
                y = self.self_attn(
                    x=x,
                    positions=positions,
                    forward_batch=forward_batch,
                    x_quant=quantized[0] if quantized else None,
                    defer_all_reduce=fuse_all_reduce_mhc,
                )
            del x
            return mhc.run_attn_post(
                self.attn_hc,
                y,
                residual,
                stats_stream=stats_stream,
                next=self.local_boundary,
                world_size=world_size,
                precomputed=precomputed,
            )

        def run_ffn_hc(state: mhc.HcState) -> mhc.HcState:
            assert self.ffn_hc is not None
            residual = state.residual
            mhc.fork_stats_stream(stats_stream)
            x = mhc.combine(self.ffn_hc, state)
            precomputed = state.stats
            state.release()
            del state
            fuse_all_reduce_mhc = self._can_fuse_ffn_mhc and can_fuse_all_reduce(
                x.shape[0], self.hc_cfg.hidden
            )
            y = self._run_moe_ffn_dp_sync(
                x,
                forward_batch,
                input_ids=input_ids,
                input_ids_global=input_ids_global,
                return_moe_output=fuse_all_reduce_mhc,
            )
            del x
            return mhc.run_moe_post(
                self.ffn_hc,
                y,
                residual,
                stats_stream=stats_stream,
                next=self.next_boundary,
                world_size=self.mlp.tp_size,
                precomputed=precomputed,
            )

        return run_ffn_hc(run_attn_hc(state))

    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],
        return_moe_output: bool = False,
    ) -> Union[torch.Tensor, deepseek_v2.MoEOutput]:
        _use_cp = is_cp_active(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
            or (_use_tp_moe_gather and should_use_dp_reduce_scatterv())
        )
        # 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/attn_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 ~attn_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 = attn_cp_interleave_gather(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)
            # The gather ships the whole padded chunk and the MoE runs on it, but
            # graph replay never refreshes pad rows, so they hold stale pool
            # memory (bf16-max scale) that overflows the MoE and the combine.
            if _is_hip:
                _mask_dp_pad_rows(local_hidden_states, forward_batch)
            # 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()
        )
        with (
            get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter),
            gathered_rows,
        ):
            return_moe_output = (
                return_moe_output
                and not _use_cp
                and not _use_tp_moe_gather
                and not _use_tp_attn_a2a_scatter
            )
            hidden_states = self.mlp(
                hidden_states,
                forward_batch,
                input_ids=input_ids,
                input_ids_global=input_ids_global,
                skip_shared_experts=_do_shared_local,
                return_moe_output=return_moe_output,
            )

        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)
                hidden_states = hidden_states.tensor_split(parallel.attn_cp_size)[
                    parallel.attn_cp_rank
                ].contiguous()
            else:
                hidden_states = attn_cp_interleave_reduce_scatter(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
            if _is_hip:
                hidden_states = _attn_tp_all_gather_chunks(
                    hidden_states, chunk_rows=[t.shape[0] for t in _a2a_scatter_chunks]
                )
            else:
                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:
            hidden_states, x_quant = self._input_norm(hidden_states)
        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).
            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 _mask_dp_pad_rows(local_hidden_states, forward_batch) -> None:
    n_real = forward_batch.num_token_non_padded
    num_rows = local_hidden_states.shape[0]
    if n_real is None or num_rows == 0:
        return
    pad = torch.arange(num_rows, device=local_hidden_states.device) >= n_real
    local_hidden_states.masked_fill_(
        pad.view(-1, *([1] * (local_hidden_states.dim() - 1))), 0
    )


def _gather_late_layer_ids(
    input_ids: torch.Tensor, counts: List[int], forward_batch: ForwardBatch
) -> torch.Tensor:
    # Re-gather the compact tail IDs in the late-layer MoE layout.
    gathered = input_ids.new_empty((sum(counts), 1))
    dp_gather_replicate(gathered, input_ids[:, None].clone(), forward_batch)
    return gathered.squeeze(-1)


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).
    if tail.scatter_indices is not None:
        # Prefill CUDA graph: fixed-capacity tail block; padding rows land on the
        # dummy row past the extend and are dropped by the slice.
        full = rows.new_zeros((num_tokens + 1, rows.shape[1]))
        full.index_copy_(0, tail.scatter_indices, rows)
        return full[:num_tokens]
    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_pp_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,
            ),
            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}"
            assert getattr(config, "num_hash_layers", 0) <= self.late_layer_start, (
                "decoder SWA bounded replay needs hash routing before the late layers"
            )
        # Only vision routing reads token IDs in the late layers.
        self.late_layer_vl_routing = (
            config.model_type == "deepseek_v41"
            and config.vision_n_layers > 0
            and not getattr(config, "language_model_only", False)
        )

    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: hash the whole prompt
                total = int(forward_batch.attn_cp_metadata.total_seq_lens)
                hash_ids = self.engram_hasher(
                    forward_batch.input_ids[:total], 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
        late_dp_counts = None
        saved_dp = None
        if (
            self.late_layer_start is not None
            and get_parallel().attn_dp_size > 1
            and (forward_batch.is_extend_in_batch or _is_hip)
        ):
            late_dp_counts = late_layer_dp_counts(
                forward_batch,
                (
                    tail.token_indices.shape[0]
                    if tail is not None
                    else hidden_states.shape[0]
                ),
            )
        tail_in_graph = (
            tail is not None
            and tail.scatter_indices is not None
            and is_in_breakable_cuda_graph()
        )
        tail_sentinel = positions.new_zeros((1,)) if tail_in_graph else None
        saved_full = None
        # A pending post never meets a residual reader: the HIP boundary's defer_post
        # gate excludes Engram/DSpark-capture layers and the model end.
        # mHC assumes full token rows per rank; LayerNorm SP needs its own path.
        assert not get_forward().sp_active
        state = mhc.HcState(hidden_states)
        for i in range(self.start_layer, self.end_layer):
            if tail is not None and i == self.late_layer_start:
                # Decode reaches back at most SWA_WINDOW positions.
                if tail_in_graph:
                    # Eager break: re-run at replay so the backend swaps onto the
                    # live batch's tail metadata; the gathers below are captured
                    # against the persistent index buffers.
                    tail_sentinel = bcg_deepseek_v4_enter_late_layer_tail(tail_sentinel)
                    saved_full = attn_backend.prefill_graph_late_layer_saved
                else:
                    saved_full = attn_backend.enter_late_layer_tail(forward_batch)
                state = state.take_rows(tail.rows)
                input_ids = tail.rows(input_ids)
                if not (
                    get_parallel().attn_dp_size > 1 and get_moe_a2a_backend().is_none()
                ):
                    input_ids_global = tail.rows(input_ids_global)
                positions = tail.positions
                if hash_ids is not None:
                    hash_ids = tail.rows(hash_ids)
            if late_dp_counts is not None and i == self.late_layer_start:
                saved_dp = enter_late_layer_dp(forward_batch, late_dp_counts)
                if get_moe_a2a_backend().is_none():
                    # Full-prefill IDs are not indexed by this rank's tail rows.
                    input_ids_global = (
                        _gather_late_layer_ids(input_ids, late_dp_counts, forward_batch)
                        if self.late_layer_vl_routing
                        else None
                    )
            engram = self.layers[i].engram
            if engram is not None:
                before_engram = state.residual
                hidden_states = engram(
                    state.residual,
                    hash_ids[:, engram.layer_hash_index],
                    forward_batch,
                    cp_all_tokens=cp_extend,
                )
                # Only multimodal placeholders become image_token_id, so a text-only
                # batch skips this full-residual where.
                if (
                    self.config.model_type == "deepseek_v41"
                    and self.config.vision_n_layers > 0
                    and forward_batch.contains_mm_inputs()
                ):
                    hidden_states = torch.where(
                        (input_ids == self.config.image_token_id)[:, None, None],
                        before_engram,
                        hidden_states,
                    )
                state = state.with_residual(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 = state.residual
                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)
            )
            with ctx:
                if _is_hip and self.layers[i].hc_boundary_fused:
                    state = _hip.forward_layer_fused_boundary(
                        self,
                        i,
                        state,
                        positions=positions,
                        input_ids=input_ids,
                        forward_batch=forward_batch,
                        input_ids_global=input_ids_global,
                        capture_dspark=capture_dspark,
                    )
                    continue
                state = self.layers[i].forward_hc_pre_from_prev(
                    positions=positions,
                    state=state,
                    input_ids=input_ids,
                    forward_batch=forward_batch,
                    input_ids_global=input_ids_global,
                )
        state = state.materialized(self.layers[self.end_layer - 1].hc_cfg)
        if saved_dp is not None:
            exit_late_layer_dp(forward_batch, saved_dp)
        if saved_full is not None:
            if tail_in_graph:
                tail_sentinel = bcg_deepseek_v4_exit_late_layer_tail(tail_sentinel)
            else:
                attn_backend.exit_late_layer_tail(saved_full, forward_batch)
            return state.residual, state.pre, tail
        return state.residual, state.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 is_cp_active(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_key,
        )

        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],
        )

        return _model_forward_tbo_merge_key(
            outputs_arr[0], outputs_arr[1], "hidden_states", hidden_states.shape[0]
        )

    @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
            from sglang.kernels.ops.layernorm.mhc import hc_broadcast

            hidden_states = hc_broadcast(hidden_states, self.hc_mult)
        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


# Quark keys its per-layer specs by checkpoint module names (layers.N.attn.wq_a,
# mtp.0.ffn.experts.N.w1, ...); this maps them onto the module prefixes built here.
_QUARK_CKPT_TO_SGLANG_MAPPER = WeightsMapper(
    orig_to_new_prefix={"layers.": "model.layers.", "mtp.": "model.mtp."},
    orig_to_new_substr={".attn.": ".self_attn.", ".ffn.": ".mlp."},
    orig_to_new_suffix={".w1": ".gate_proj", ".w2": ".down_proj", ".w3": ".up_proj"},
)


def _is_quark_checkpoint(config: DeepSeekV4Config) -> bool:
    quantization_config = getattr(config, "quantization_config", None) or {}
    return quantization_config.get("quant_method") == "quark"


class DeepseekV4ForCausalLM(nn.Module):
    supports_cuda_vmm_feature_transport = True

    @classmethod
    def get_hf_to_sglang_mapper(cls, config) -> Optional[WeightsMapper]:
        # The loader hands this mapper to every quant config; other formats were
        # written against different names, so only Quark checkpoints get it.
        return _QUARK_CKPT_TO_SGLANG_MAPPER if _is_quark_checkpoint(config) else None

    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_parallel().pp_group.world_size != 1
                or not (
                    get_moe_a2a_backend().is_none()
                    or (_is_cuda and get_moe_a2a_backend().is_megamoe())
                )
            ):
                raise ValueError(
                    "V4.1 vision supports TP/EP/DP without CP or PP; "
                    "MoE A2A requires CUDA MegaMoE"
                )

            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: array, mm_inputs) -> array:
        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)
            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)
            if tail.real_rows_count is not None:
                # Fixed-capacity tail block from the prefill CUDA graph: drop padding.
                n = tail.real_rows_count
                input_ids = input_ids[:n]
                hidden_states = hidden_states[:n]
                pre_hc_head = pre_hc_head[:n]
                if aux_hidden_states is not None:
                    aux_hidden_states = [a[:n] for a in aux_hidden_states]
            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[: tail.real_rows_count]
                if tail.real_rows_count is not None
                else 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()
        layers = self.model.layers
        for i, layer in enumerate(layers):
            if isinstance(layer, DeepseekV4DecoderLayer):
                layer.bind_next(layers[i + 1] if i + 1 < len(layers) else None)

    def precompile_kernels_after_loading(self) -> None:
        from sglang.srt.layers.moe.mega_moe import (
            build_mega_moe_shared_weights,
            should_fuse_mega_moe_shared_experts,
        )

        # The only deployments whose fused-collective gates can fire; registration
        # must happen eagerly here, before graph capture.
        if get_platform().is_blackwell and get_parallel().tp_size == 4:
            from sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe import (
                register_fused_all_reduce_comm,
            )

            register_fused_all_reduce_comm()

        for module in self.modules():
            if isinstance(
                module, deepseek_v2.DeepseekV2MoE
            ) and should_fuse_mega_moe_shared_experts(module):
                module.mega_shared_l1_weights, module.mega_shared_l2_weights = (
                    build_mega_moe_shared_weights(module.shared_experts)
                )

    @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
    if quant_config is not None and quant_config.get_name() == "quark":
        weights = _rename_quark_fp8_block_scales(weights)
    return _dequant_fp8_wo_a_streaming(weights)


def _rename_quark_fp8_block_scales(
    weights: Iterable[Tuple[str, torch.Tensor]],
) -> Iterable[Tuple[str, torch.Tensor]]:
    # Quark spells the FP8-block scale ".weight_scale" where native V4 spells it
    # ".scale" (both e8m0); MXFP4 scales are uint8 and keep their name.
    for name, tensor in weights:
        if name.endswith(".weight_scale") and tensor.dtype == torch.float8_e8m0fnu:
            name = name.removesuffix(".weight_scale") + ".scale"
        yield name, tensor


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)