sgl-project--sglang/python/sglang/srt/models/deepseek_v4_dspark.py

from __future__ import annotations

import logging
from typing import Iterable, List, Optional, Tuple

import msgspec
import torch
import torch.nn.functional as F
from torch import nn

from sglang.jit_kernel.dsv4 import fused_q_norm_rope, fused_rope_inplace
from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
from sglang.srt.environ import envs
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.dbrx import ReplicatedLinear
from sglang.srt.models.deepseek_v4 import (
    DEEPSEEK_V4_STACKED_PARAMS_MAPPING,
    DeepseekV4DecoderLayer,
    MqaAttentionBase,
    _dequant_fp8_wo_a,
    hc_head_torch,
    make_hc_head_params,
)
from sglang.srt.models.dspark import (
    DSparkConfidenceHead,
    StepSampler,
    gather_and_crop_vocab,
    run_markov_block,
)
from sglang.srt.runtime_context import get_parallel
from sglang.srt.speculative.dspark_components.dspark_config import (
    parse_dspark_draft_config,
)
from sglang.srt.speculative.dspark_components.kernels.dspark_draft_model import (
    BuildStepLocal,
    CommitKvProj,
)
from sglang.srt.speculative.ragged_verify import (
    RaggedVerifyMode,
    read_ragged_verify_mode,
)
from sglang.srt.utils import add_prefix, is_blackwell_supported
from sglang.srt.utils.async_probe import maybe_detect_in_closed_range

logger = logging.getLogger(__name__)

_PAD_NUM_HEADS = 64


def apply_rotary_emb(
    x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False
) -> torch.Tensor:
    y = x
    x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2)))
    if inverse:
        freqs_cis = freqs_cis.conj()
    if x.ndim == 3:
        freqs_cis = freqs_cis.view(x.size(0), 1, x.size(-1))
    else:
        freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))
    x = torch.view_as_real(x * freqs_cis).flatten(-2)
    y.copy_(x)
    return y


class DSparkAttention(MqaAttentionBase):

    def __init__(
        self,
        config: DeepSeekV4Config,
        layer_id: int,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        alt_streams: Optional[List[torch.cuda.Stream]] = None,
    ) -> None:
        super().__init__(
            config,
            layer_id,
            quant_config,
            prefix,
            attn_tp_rank=get_parallel().attn_tp_rank,
            attn_tp_size=get_parallel().attn_tp_size,
            compress_ratio=0,
            fuse_wqa_wkv=False,
            wo_a_fp8=False,
            wo_a_keeps_quant_config=False,
            wo_b_reduce_results=True,
            rope_original_seq_len=0,
        )
        assert (
            self.compress_ratio == 0
        ), "DSpark draft attention requires compress_ratio == 0."
        self.window_size = int(
            getattr(config, "sliding_window", None) or config.window_size
        )

        self.attn = 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", prefix),
        )

        self._use_fast_kernel = envs.SGLANG_DSPARK_FAST_KERNEL.get()
        self.alt_streams = alt_streams
        self._multi_stream_bs_limit = 128 if is_blackwell_supported() else 64

    def kv_proj_only(self, x: torch.Tensor) -> torch.Tensor:
        kv, _ = self.wkv(x)
        return kv

    def _local_attn_sink(self) -> torch.Tensor:
        if self.attn_tp_size == 1:
            return self.attn_sink
        if self._attn_sink_local is None:
            rank = self.attn_tp_rank
            num_heads = self.n_local_heads
            sink = self.attn_sink.new_zeros(max(num_heads, _PAD_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

    def _store_block_kv(
        self,
        *,
        kv: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        attn_backend,
        pool: DeepSeekV4TokenToKVPool,
    ) -> None:
        pool.set_swa_key_buffer_radix_fused_norm_rope(
            layer_id=self.layer_id,
            swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch),
            kv=kv,
            kv_weight=self.kv_norm.weight.data,
            eps=self.eps,
            freqs_cis=self.freqs_cis,
            positions=positions,
        )

    def _compute_q(
        self,
        x: torch.Tensor,
        positions: torch.Tensor,
        q_out: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        q, _ = self.wq_a(x)
        q = self.q_norm(q)
        q, _ = self.wq_b(q)
        q = q.view(-1, self.n_local_heads, self.head_dim)
        if self._use_fast_kernel:
            if q_out is None:
                q_out = torch.empty_like(q)
            fused_q_norm_rope(q, q_out, self.eps, self.freqs_cis, positions)
            return q_out
        else:
            q = q * torch.rsqrt(
                q.float().square().mean(-1, keepdim=True) + self.eps
            ).to(q.dtype)
            apply_rotary_emb(q[..., -self.rope_head_dim :], self.freqs_cis[positions])
            if q_out is not None:
                q_out.copy_(q)
                return q_out
            return q

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
    ) -> torch.Tensor:
        from sglang.srt.model_executor.forward_context import get_attn_backend

        pool = _resolve_dspark_pool()
        attn_backend = get_attn_backend()
        rd = self.rope_head_dim

        enable_multi_stream = (
            self.alt_streams is not None
            and get_is_capture_mode()
            and hidden_states.shape[0] <= self._multi_stream_bs_limit
        )

        q_padded: Optional[torch.Tensor] = None
        q_out: Optional[torch.Tensor] = None
        if self.n_local_heads < _PAD_NUM_HEADS:
            q_padded = hidden_states.new_empty(
                hidden_states.shape[0], _PAD_NUM_HEADS, self.head_dim
            )
            q_out = q_padded[:, : self.n_local_heads, :]

        if enable_multi_stream:
            current_stream = torch.cuda.current_stream()
            stream_kv = self.alt_streams[0]
            stream_kv.wait_stream(current_stream)
            with torch.cuda.stream(stream_kv):
                kv = self.kv_proj_only(hidden_states)
                self._store_block_kv(
                    kv=kv,
                    positions=positions,
                    forward_batch=forward_batch,
                    attn_backend=attn_backend,
                    pool=pool,
                )
            q = self._compute_q(hidden_states, positions, q_out=q_out)
            current_stream.wait_stream(stream_kv)
        else:
            kv = self.kv_proj_only(hidden_states)
            self._store_block_kv(
                kv=kv,
                positions=positions,
                forward_batch=forward_batch,
                attn_backend=attn_backend,
                pool=pool,
            )
            q = self._compute_q(hidden_states, positions, q_out=q_out)

        if q_padded is not None:
            q = q_padded
        attn_sink = self._local_attn_sink()

        o = attn_backend.forward(
            q=q,
            k=kv,
            v=kv,
            layer=self.attn,
            forward_batch=forward_batch,
            compress_ratio=0,
            attn_sink=attn_sink,
            save_kv_cache=False,
        )
        if o.shape[1] != self.n_local_heads:
            o = o[:, : self.n_local_heads, :]

        if self._use_fast_kernel:
            fused_rope_inplace(
                o[..., -rd:], None, self.freqs_cis, positions=positions, inverse=True
            )
        else:
            apply_rotary_emb(o[..., -rd:], self.freqs_cis[positions], inverse=True)

        o = o.view(
            o.shape[0],
            self.n_local_groups,
            o.shape[1] * o.shape[2] // self.n_local_groups,
        )
        wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
        if self._use_fast_kernel:
            o = torch.einsum("bgd,grd->bgr", o, wo_a)
        else:
            o = torch.einsum("bgd,grd->bgr", o.float(), wo_a.float()).to(q.dtype)
        out, _ = self.wo_b(o.reshape(o.shape[0], o.shape[1] * o.shape[2]))
        return out


def _resolve_dspark_pool() -> DeepSeekV4TokenToKVPool:
    pool = get_token_to_kv_pool()
    assert isinstance(pool, DeepSeekV4TokenToKVPool), (
        "DSpark draft attention requires a DeepSeekV4TokenToKVPool, "
        f"got {type(pool).__name__}."
    )
    return pool


class MarkovW2ShardGeometry(msgspec.Struct, frozen=True):

    tp_size: int
    org_vocab_start: int
    org_vocab_end: int
    num_embeddings_per_partition: int
    num_embeddings_padded: int


class DSparkV4MarkovHead(nn.Module):

    markov_head_type = "vanilla"

    def __init__(self, *, vocab_size: int, markov_rank: int) -> None:
        super().__init__()
        self.vocab_size = int(vocab_size)
        self.markov_rank = int(markov_rank)
        if self.markov_rank <= 0:
            raise ValueError(
                f"DSparkV4MarkovHead requires markov_rank > 0, got {self.markov_rank}."
            )
        self.markov_w1 = VocabParallelEmbedding(
            self.vocab_size, self.markov_rank, enable_tp=False
        )
        self._opt_markov_w2_bf16 = envs.SGLANG_DSPARK_OPT_MARKOV_W2_BF16.get()
        self._opt_markov_w2_tp_shard = envs.SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD.get()
        markov_w2_dtype = torch.bfloat16 if self._opt_markov_w2_bf16 else torch.float32
        self.markov_w2 = nn.Linear(
            self.markov_rank, self.vocab_size, bias=False, dtype=markov_w2_dtype
        )
        self._tp_shard: Optional[MarkovW2ShardGeometry] = None

    def configure_tp_shard(self, *, lm_head: nn.Module) -> None:
        if not self._opt_markov_w2_tp_shard:
            return
        if int(lm_head.org_vocab_size) != self.vocab_size:
            raise ValueError(
                "DSpark markov_w2 TP-shard requires lm_head.org_vocab_size == "
                f"markov vocab_size, got {int(lm_head.org_vocab_size)} vs "
                f"{self.vocab_size}."
            )
        tp_size = int(lm_head.tp_size)
        per_partition = int(lm_head.num_embeddings_per_partition)
        num_padded = int(lm_head.num_embeddings_padded)
        if per_partition * tp_size != num_padded:
            raise ValueError(
                "DSpark markov_w2 TP-shard could not align to the lm_head partition: "
                f"num_embeddings_per_partition({per_partition}) * tp_size({tp_size}) != "
                f"num_embeddings_padded({num_padded})."
            )
        attn_tp_size = get_parallel().attn_tp_group.world_size
        if attn_tp_size != tp_size:
            raise ValueError(
                "DSpark markov_w2 TP-shard needs the attn-TP group (used for the per-step "
                f"all-gather) to equal the lm_head shard group, got attn_tp_size="
                f"{attn_tp_size} vs lm_head tp_size={tp_size}. This config (e.g. DP "
                "attention without --enable-dp-lm-head, where lm_head shards over the "
                "global TP group) is unsupported; disable "
                "SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD."
            )
        self._tp_shard = MarkovW2ShardGeometry(
            tp_size=tp_size,
            org_vocab_start=int(lm_head.shard_indices.org_vocab_start_index),
            org_vocab_end=int(lm_head.shard_indices.org_vocab_end_index),
            num_embeddings_per_partition=per_partition,
            num_embeddings_padded=num_padded,
        )

    def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor:
        return self.markov_w1(token_ids.long())

    def project_bias(
        self, latent_states: torch.Tensor, *, weight: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        weight = self.markov_w2.weight if weight is None else weight
        if self._opt_markov_w2_bf16:
            return F.linear(latent_states.to(weight.dtype), weight).float()
        return F.linear(latent_states.float(), weight)

    def compute_step_bias(
        self, token_ids: torch.Tensor, hidden_states: Optional[torch.Tensor]
    ) -> torch.Tensor:
        del hidden_states
        return self.project_bias(self.get_prev_embeddings(token_ids))

    def apply_step_logits(
        self,
        logits: torch.Tensor,
        *,
        token_ids: torch.Tensor,
        hidden_states: Optional[torch.Tensor],
    ) -> torch.Tensor:
        if self._tp_shard is not None:
            return self._apply_step_logits_sharded(
                base_local=logits, token_ids=token_ids
            )
        return logits + self.compute_step_bias(token_ids, hidden_states)

    def _apply_step_logits_sharded(
        self, *, base_local: torch.Tensor, token_ids: torch.Tensor
    ) -> torch.Tensor:
        shard = self._tp_shard
        latent = self.get_prev_embeddings(token_ids)
        weight_local = self.markov_w2.weight[
            shard.org_vocab_start : shard.org_vocab_end
        ]
        if self._opt_markov_w2_bf16:
            bias = F.linear(latent.to(weight_local.dtype), weight_local)
        else:
            bias = F.linear(latent.float(), weight_local)
        step_local = BuildStepLocal.execute(bias=bias, base_local=base_local)
        if shard.tp_size > 1:
            full = get_parallel().attn_tp_group.all_gather(step_local, dim=-1)
        else:
            full = step_local
        return full[..., : self.vocab_size]

    def forward(self, token_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        embed = self.get_prev_embeddings(token_ids)
        logits = self.project_bias(embed)
        return logits, embed

    def sample_block(
        self,
        base_logits: torch.Tensor,
        *,
        first_prev_tokens: torch.Tensor,
        hidden_states: Optional[torch.Tensor],
        sampler: StepSampler,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        return run_markov_block(
            self,
            base_logits,
            first_prev_tokens=first_prev_tokens,
            hidden_states=hidden_states,
            sampler=sampler,
        )


def build_dspark_v4_confidence_head(
    *, config: DeepSeekV4Config, markov_rank: int
) -> Optional[DSparkConfidenceHead]:
    if read_ragged_verify_mode() is RaggedVerifyMode.STATIC:
        return None
    if not hasattr(config, "enable_confidence_head"):
        logger.warning(
            "DSpark draft config has no enable_confidence_head field; treating the "
            "confidence head as enabled."
        )
    with_markov_cfg = getattr(config, "confidence_head_with_markov", None)
    with_markov = (
        (markov_rank > 0) if with_markov_cfg is None else bool(with_markov_cfg)
    )
    if with_markov and markov_rank <= 0:
        raise ValueError(
            "DSpark V4 confidence_head_with_markov requires markov_rank > 0, "
            f"got markov_rank={markov_rank}."
        )
    return DSparkConfidenceHead(
        hidden_size=int(config.hidden_size),
        markov_rank=int(markov_rank),
        with_markov=with_markov,
        bias=False,
    )


class DSparkV4Stage(DeepseekV4DecoderLayer):

    def __init__(
        self,
        config: DeepSeekV4Config,
        layer_id: int,
        stage_id: int,
        num_stages: int,
        num_target_layers: int,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        alt_streams: Optional[List[torch.cuda.Stream]] = None,
    ) -> None:
        super().__init__(
            config=config,
            layer_id=layer_id,
            quant_config=quant_config,
            prefix=prefix,
            is_nextn=True,
            alt_streams=alt_streams,
        )
        self.stage_id = stage_id
        self.dim = config.hidden_size

        if stage_id == 0:
            if num_target_layers <= 0:
                raise ValueError(
                    "DSpark needs target layers for the target-hidden projection."
                )
            self.main_proj = ReplicatedLinear(
                config.hidden_size * num_target_layers,
                config.hidden_size,
                bias=False,
                quant_config=quant_config,
                prefix=add_prefix("main_proj", prefix),
            )
            self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        if stage_id == num_stages - 1:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
            (
                self.hc_head_fn,
                self.hc_head_base,
                self.hc_head_scale,
            ) = make_hc_head_params(config.hc_mult, config.hidden_size)

    def _build_self_attn(
        self,
        *,
        config: DeepSeekV4Config,
        layer_id: int,
        quant_config: Optional[QuantizationConfig],
        prefix: str,
        alt_streams: Optional[List[torch.cuda.Stream]],
        compress_ratio_override: Optional[int],
    ) -> nn.Module:
        del compress_ratio_override
        return DSparkAttention(
            config=config,
            layer_id=layer_id,
            quant_config=quant_config,
            prefix=prefix,
            alt_streams=alt_streams,
        )

    def _hc_pre_block(
        self,
        x: torch.Tensor,
        hc_fn: torch.Tensor,
        hc_scale: torch.Tensor,
        hc_base: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        y, post, comb, _ = self.hc_pre(x, hc_fn, hc_scale, hc_base)
        return y, post, comb

    def _hc_post_block(
        self,
        x: torch.Tensor,
        residual: torch.Tensor,
        post: torch.Tensor,
        comb: torch.Tensor,
    ) -> torch.Tensor:
        return self.hc_post(x, residual, post, comb)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
    ) -> torch.Tensor:
        residual = hidden_states
        x, post, comb = self._hc_pre_block(
            hidden_states, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
        )
        x = self.input_layernorm(x)
        x = self.self_attn(positions, x, forward_batch)
        x = self._hc_post_block(x, residual, post, comb)

        residual = x
        x, post, comb = self._hc_pre_block(
            x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
        )
        x = self.post_attention_layernorm(x)
        x = self._run_ffn(x, forward_batch)
        x = self._hc_post_block(x, residual, post, comb)
        return x

    def _run_ffn(self, x: torch.Tensor, forward_batch: ForwardBatch) -> torch.Tensor:
        shape = x.shape
        x = x.reshape(-1, self.dim)
        y = self._run_moe_ffn_dp_sync(
            x, forward_batch, input_ids=None, input_ids_global=None
        )
        return y.view(shape)


class DeepseekV4ForCausalLMDSpark(nn.Module):

    def __init__(
        self,
        config: DeepSeekV4Config,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config
        self.quant_config = quant_config

        dspark_config = parse_dspark_draft_config(draft_hf_config=config)
        if not dspark_config.require_markov():
            raise ValueError(
                "DSpark V4 draft requires markov_rank > 0, "
                f"got markov_rank={dspark_config.markov_rank}."
            )
        self.gamma = int(
            dspark_config.resolve_gamma(default=int(config.num_hidden_layers))
        )
        self.block_size = self.gamma
        if dspark_config.target_layer_ids is not None:
            self.num_stages = len(dspark_config.target_layer_ids)
        else:
            self.num_stages = int(getattr(config, "num_nextn_predict_layers", 1) or 1)

        target_num_layers = (
            int(dspark_config.num_target_layers)
            if dspark_config.num_target_layers is not None
            else int(getattr(config, "num_hidden_layers", 1))
        )
        if dspark_config.target_layer_ids is not None:
            self.num_target_features = len(dspark_config.target_layer_ids)
        else:
            self.num_target_features = target_num_layers

        self.start_layer = 0
        self.end_layer = self.num_stages
        use_multi_stream = (
            envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get()
            and envs.SGLANG_DSPARK_ENABLE_MULTI_STREAM.get()
            and torch.cuda.is_available()
        )
        self.alt_streams: Optional[List[torch.cuda.Stream]] = (
            [torch.cuda.Stream()] if use_multi_stream else None
        )
        self.stages = nn.ModuleList(
            [
                DSparkV4Stage(
                    config=config,
                    layer_id=stage_id,
                    stage_id=stage_id,
                    num_stages=self.num_stages,
                    num_target_layers=self.num_target_features,
                    quant_config=quant_config,
                    prefix=add_prefix(f"stages.{stage_id}", prefix),
                    alt_streams=self.alt_streams,
                )
                for stage_id in range(self.num_stages)
            ]
        )
        self.markov_head = DSparkV4MarkovHead(
            vocab_size=int(config.vocab_size),
            markov_rank=int(dspark_config.markov_rank),
        )
        self.confidence_head = build_dspark_v4_confidence_head(
            config=config, markov_rank=int(dspark_config.markov_rank)
        )
        self.hc_mult = int(config.hc_mult)
        self.norm_eps = float(config.rms_norm_eps)
        self.hc_eps = float(config.hc_eps)

        self.embed_tokens: Optional[nn.Module] = None
        self.lm_head: Optional[nn.Module] = None
        self._use_fp32_lm_head = envs.SGLANG_DSPARK_FP32_LM_HEAD.get()
        self._opt_markov_w2_tp_shard = envs.SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD.get()

    @property
    def enable_confidence_head(self) -> bool:
        return self.confidence_head is not None

    def attach_shared_modules(
        self, *, embed_tokens: nn.Module, lm_head: nn.Module
    ) -> None:
        self.embed_tokens = embed_tokens
        self.lm_head = lm_head
        self.markov_head.configure_tp_shard(lm_head=lm_head)

    def project_target_hidden(self, main_hidden: torch.Tensor) -> torch.Tensor:
        stage0 = self.stages[0]
        projected, _ = stage0.main_proj(main_hidden)
        return stage0.main_norm(projected)

    def write_target_hidden_kv(
        self,
        *,
        main_hidden: torch.Tensor,
        swa_loc: torch.Tensor,
        positions: torch.Tensor,
        pool: DeepSeekV4TokenToKVPool,
    ) -> None:
        main_x = self.project_target_hidden(main_hidden)
        swa_loc = swa_loc.to(torch.int32)
        kvs = CommitKvProj.execute(
            main_x=main_x,
            wkv_linears=[stage.self_attn.wkv for stage in self.stages],
        )
        for stage, kv in zip(self.stages, kvs):
            attn = stage.self_attn
            pool.set_swa_key_buffer_radix_fused_norm_rope(
                layer_id=attn.layer_id,
                swa_loc=swa_loc,
                kv=kv,
                kv_weight=attn.kv_norm.weight.data,
                eps=attn.eps,
                freqs_cis=attn.freqs_cis,
                positions=positions,
            )

    def forward_embed(self, input_ids: torch.Tensor) -> torch.Tensor:
        if self.embed_tokens is None:
            raise ValueError(
                "DeepseekV4ForCausalLMDSpark requires the target embed_tokens "
                "(call attach_shared_modules first)."
            )
        x = self.embed_tokens(input_ids)
        x = x.unsqueeze(1).repeat(1, self.hc_mult, 1)
        return x

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        input_embeds: Optional[torch.Tensor] = None,
        get_embedding: bool = False,
        pp_proxy_tensors=None,
    ) -> LogitsProcessorOutput:
        del get_embedding, pp_proxy_tensors
        if input_embeds is None:
            input_embeds = self.forward_embed(input_ids)
        x = input_embeds
        for stage in self.stages:
            x = stage(positions, x, forward_batch)

        return LogitsProcessorOutput(next_token_logits=None, hidden_states=x)

    def collapse_hc_head(self, x: torch.Tensor) -> torch.Tensor:
        last = self.stages[-1]
        return hc_head_torch(
            x,
            last.hc_head_fn,
            last.hc_head_scale,
            last.hc_head_base,
            norm_eps=self.norm_eps,
            hc_eps=self.hc_eps,
        )

    def compute_base_logits(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

        x_post_hc = self.collapse_hc_head(x)
        return self._logits_from_x_post_hc(x_post_hc), x_post_hc

    def _logits_from_x_post_hc(self, x_post_hc: torch.Tensor) -> torch.Tensor:
        if self.lm_head is None:
            raise ValueError(
                "DeepseekV4ForCausalLMDSpark requires the target lm_head "
                "(call attach_shared_modules first)."
            )
        last = self.stages[-1]
        x = last.norm(x_post_hc)
        weight = self.lm_head.weight
        if self._use_fp32_lm_head:
            local_logits = F.linear(x.float(), weight.float())
        else:
            local_logits = torch.matmul(x.to(weight.dtype), weight.T)
        if self._opt_markov_w2_tp_shard:
            return local_logits
        return gather_and_crop_vocab(local_logits, self.lm_head)

    def compute_confidence(
        self,
        *,
        anchor_tokens: torch.Tensor,
        sampled_tokens: torch.Tensor,
        x_post_hc: torch.Tensor,
    ) -> Optional[torch.Tensor]:
        confidence_head = self.confidence_head
        if confidence_head is None:
            return None
        bs = int(anchor_tokens.shape[0])
        x_post_hc = x_post_hc.view(bs, self.gamma, -1)
        if confidence_head.with_markov:
            prev_seq = torch.cat(
                [anchor_tokens.view(-1, 1), sampled_tokens[:, : self.gamma - 1]], dim=1
            )
            markov_embed_stack = self.markov_head.get_prev_embeddings(prev_seq)
        else:
            markov_embed_stack = None
        confidence_raw = confidence_head(x_post_hc, markov_embed_stack)
        confidence = confidence_head.apply_sts(confidence_raw)
        maybe_detect_in_closed_range(
            confidence, 0.0, 1.0, "DSpark confidence must lie in [0, 1]."
        )
        return confidence

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> None:
        params_dict = dict(self.named_parameters())
        loaded_params = set()

        weights = list(weights)
        if any(name.endswith(".wo_a.scale") for name, _ in weights):
            weights = list(_dequant_fp8_wo_a(weights))

        stacked_params_mapping = DEEPSEEK_V4_STACKED_PARAMS_MAPPING
        from sglang.srt.layers.moe.fused_moe_triton import FusedMoE

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

        for name, loaded_weight in weights:
            mapped = self._remap_dspark_weight_name(name)
            if mapped is None:
                continue

            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in mapped:
                    continue
                candidate = mapped.replace(weight_name, param_name)
                if candidate not in params_dict:
                    continue
                param = params_dict[candidate]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                loaded_params.add(candidate)
                break
            else:
                for (
                    param_name,
                    weight_name,
                    expert_id,
                    shard_id,
                ) in expert_params_mapping:
                    if weight_name not in mapped:
                        continue
                    candidate = mapped.replace(weight_name, param_name)
                    if candidate not in params_dict:
                        continue
                    param = params_dict[candidate]
                    weight_loader = param.weight_loader
                    weight_loader(
                        param,
                        loaded_weight,
                        candidate,
                        shard_id=shard_id,
                        expert_id=expert_id,
                    )
                    loaded_params.add(candidate)
                    break
                else:
                    if mapped not in params_dict:
                        logger.warning(
                            "DSpark V4 draft: unexpected weight %r -> %r", name, mapped
                        )
                        continue
                    param = params_dict[mapped]
                    weight_loader = getattr(
                        param, "weight_loader", default_weight_loader
                    )
                    weight_loader(param, loaded_weight)
                    loaded_params.add(mapped)

        self._assert_confidence_head_loaded(
            params_dict=params_dict, loaded_params=loaded_params
        )

    def _assert_confidence_head_loaded(
        self, *, params_dict: dict, loaded_params: set
    ) -> None:
        if self.confidence_head is None:
            return
        confidence_param_names = {
            name for name in params_dict if name.startswith("confidence_head.")
        }
        missing = confidence_param_names - loaded_params
        if missing:
            raise ValueError(
                f"DSpark V4 confidence head is enabled but the checkpoint is missing "
                f"{sorted(missing)}. Provide a checkpoint with trained confidence weights, "
                f"or disable the confidence head (enable_confidence_head=False)."
            )

    def _remap_dspark_weight_name(self, name: str) -> Optional[str]:
        if name.startswith(("embed.", "embed_tokens.", "head.", "lm_head.")):
            return None
        if "rotary_emb.inv_freq" in name:
            return None

        if not name.startswith("mtp."):
            return None
        parts = name.split(".", 2)
        if len(parts) < 3:
            return None
        stage_id, rest = parts[1], parts[2]

        if rest.startswith("markov_head."):
            return f"markov_head.{rest[len('markov_head.'):]}"

        if rest.startswith("confidence_head."):
            if self.confidence_head is None:
                return None
            return f"confidence_head.{rest[len('confidence_head.'):]}"

        mapped_rest = rest
        mapped_rest = mapped_rest.replace("attn.", "self_attn.", 1)
        mapped_rest = mapped_rest.replace("ffn.", "mlp.", 1)
        mapped_rest = mapped_rest.replace("attn_norm.", "input_layernorm.", 1)
        mapped_rest = mapped_rest.replace("ffn_norm.", "post_attention_layernorm.", 1)
        mapped_rest = mapped_rest.replace(".w1.", ".gate_proj.")
        mapped_rest = mapped_rest.replace(".w2.", ".down_proj.")
        mapped_rest = mapped_rest.replace(".w3.", ".up_proj.")
        mapped_rest = mapped_rest.replace(".gate.tid2eid", ".topk.tid2eid")
        mapped_rest = mapped_rest.replace(".gate.bias", ".gate.e_score_correction_bias")
        mapped_rest = mapped_rest.replace(".scale", ".weight_scale_inv")
        return f"stages.{stage_id}.{mapped_rest}"


EntryClass = [DeepseekV4ForCausalLMDSpark]