Upload folder using huggingface_hub

model-00016-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1da9dfeda86c4d8bdbe0e0220347260d9a3096462757073d043df97909499b28
+size 6449204804

model-00017-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:709daa789a7fe689885aac68e0318e5d9ff3c9ebe0d8caed11757b161d4de593
+size 6458523600

model-00018-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9dc0cc3eaa551fcc0918ee0e779a9c16e767924a5ef8ee92d3f8b79908fc7b4d
+size 6445125364

model-00019-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21fbeb7933260748dd021c39d6c1806b4369337894ccc3253f3a7af131a1e53f
+size 6458523472

model-00020-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e4699b807ce99d6895ac617d497caa3faffe85128d890a2ce313a01966d8cc7
+size 6449197348

model-00021-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c249772e3c3c657ef301c76ed0f4c0f628145df14f8b8d3e0940cf3a20d3a81e
+size 6458523448

model-00022-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26656dc243c9e21f36520317b5e6f205f23c73be389e967e8558b97f19b3b186
+size 6445125396

model-00023-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:544219a4c26717ff7c48159d4ae784995dc722d3e0c3bf46eada5f07775c36f1
+size 6458523536

model-00024-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9259a0f5c55aa2237c16167caa66db76f6eaecb37f84b56ebfc12d76c00fd043
+size 6449197260

model-00025-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54ea8d09de13508ebce7d9e609cf195ca0dcbf7ea8e16cd2a57752aa1f2194df
+size 6458523608

model-00026-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28d247fe3cbccdd48bb297107824090cb9cd331b286f0f201f5c3dc39783158c
+size 6446436500

model-00027-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d7765703b16007cc52519066a3e37a291115e184c54de4a606d94ce4383e1e72
+size 6457212480

model-00028-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:838ca1c68d974bbd6780bc20bb4873df55dc3898dd2d5e9862d1b8d5b63b3072
+size 6116586562

model-00029-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d453ea5b1139459c277a2ea336ed7065ef67db3f6516b416d1bd0d192cd2aa17
+size 8304

model-00030-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f73d6dcf572c96b1204ed32301104c6f6aaf67a2372cc77add9ed336dab3f966
+size 11012456

model-00031-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71997f843485ea05aaff30b11b84fb3f9e8d4062f1424ab42884c442b82f60d4
+size 4171198532

model-00032-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba08839f9991c0ebd6f648d7df4dcf0c5e6bd9ec5be844a387473ae6ea91c396
+size 4192576620

model-00033-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a50d44978124ecc8b1404e7ef180951690e94224b9414ebdedb76871b5d405a
+size 4171198532

model-00034-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6e426cd19c2ac355de42119e306920f6165bab73a6f197c7bede9e0a99d6e63
+size 4192576620

model-00035-of-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0bfd040848b2f9d1d2100779e7d57a940bd5d3581ea403bb560622c6df619208
+size 4171198532

model-00036-of-00036.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:04c1d3f8536009af1a55739f99146acea6ef08fa0d8bb6f155e156c46464ffed
+size 2127629656

model-parallel-ffn.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:42c8e4de1726b28fedf86d059754093b083b6a2f4cc636fa2c873ba053695e66
+size 1132767752

modeling_agnes.py

@@ -0,0 +1,1784 @@
+# Copyright 2026 The Agnes Model Team. All rights reserved.
+#
+# This modeling code is an independent implementation of the Agnes model
+# that work's Apache-2.0 licensing and attribution accordingly.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers import initialization as init
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache, DynamicSlidingWindowLayer
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_experts_implementation
+from transformers.masking_utils import create_sliding_window_causal_mask
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
+from transformers.utils.generic import maybe_autocast, merge_with_config_defaults
+from transformers.utils.output_capturing import OutputRecorder, capture_outputs
+from transformers.conversion_mapping import register_checkpoint_conversion_mapping
+from transformers.core_model_loading import (
+    Concatenate,
+    MergeModulelist,
+    WeightConverter,
+    WeightRenaming,
+)
+
+from .configuration_agnes import AgnesConfig
+
+
+# ==========================================================================
+# Normalization primitives
+# ==========================================================================
+
+# ---------------------------------------------------------------------------
+# Root-mean-square normalisation.
+#
+# Agnes fuses the RMS statistic, the reciprocal-sqrt rescale and (optionally)
+# the learned gain into a single explicit autograd node rather than letting the
+# framework trace the gradient through the individual pointwise ops. The
+# statistic is always accumulated in float32 for stability; the closed-form
+# backward below is the analytic Jacobian of that statistic, so the saved
+# tensors are just the fp32 input, the per-row reciprocal std and (when
+# present) the gain — no intermediate activation graph is retained.
+#
+# For a row x (length n), with r = 1 / sqrt(mean(x^2) + eps) and x_hat = x * r:
+#     y            = gain * x_hat
+#     dx           = r * (g - x_hat * mean(g * x_hat))     where g = dy * gain
+#     d(gain)      = sum_over_rows(dy * x_hat)
+# ---------------------------------------------------------------------------
+
+
+class _FusedRMSNorm(torch.autograd.Function):
+    """Weighted RMS norm with a hand-derived backward (see module comment)."""
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+        out_dtype = x.dtype
+        xf = x.float()
+        rstd = torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps)
+        x_hat = (xf * rstd).to(out_dtype)
+        ctx.save_for_backward(xf, rstd, weight)
+        ctx.out_dtype = out_dtype
+        return weight * x_hat
+
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor):
+        xf, rstd, weight = ctx.saved_tensors
+        n = xf.shape[-1]
+        x_hat = xf * rstd
+        g = grad_out.float() * weight.float()
+        # projection of g onto the manifold orthogonal to x_hat, then rescale by r
+        row_corr = (g * x_hat).sum(-1, keepdim=True) / n
+        grad_x = (rstd * (g - x_hat * row_corr)).to(ctx.out_dtype)
+        reduce_axes = tuple(range(grad_out.dim() - 1))
+        grad_weight = (grad_out.float() * x_hat).sum(reduce_axes).to(weight.dtype)
+        return grad_x, grad_weight, None
+
+
+class _FusedUnweightedRMSNorm(torch.autograd.Function):
+    """Gain-free RMS norm; backward is the weighted case with gain fixed to 1."""
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, eps: float) -> torch.Tensor:
+        rstd = torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps).to(x.dtype)
+        ctx.save_for_backward(x, rstd)
+        return x * rstd
+
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor):
+        x, rstd = ctx.saved_tensors
+        rf = rstd.float()
+        x_hat = x.float() * rf
+        n = x.shape[-1]
+        g = grad_out.float()
+        row_corr = (g * x_hat).sum(-1, keepdim=True) / n
+        grad_x = (rf * (g - x_hat * row_corr)).to(grad_out.dtype)
+        return grad_x, None
+
+
+class AgnesRMSNorm(nn.Module):
+    """RMS norm with a learned per-channel gain."""
+
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return _FusedRMSNorm.apply(hidden_states, self.weight, self.variance_epsilon)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class AgnesUnweightedRMSNorm(nn.Module):
+    """RMS norm without a learned gain (used inside the attention / mHC blocks)."""
+
+    def __init__(self, eps: float = 1.0e-6):
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return _FusedUnweightedRMSNorm.apply(x, self.eps)
+
+
+
+# ==========================================================================
+# Feed-forward network and mixture-of-experts routing
+# ==========================================================================
+
+class AgnesMLP(nn.Module):
+    def __init__(self, config: AgnesConfig, intermediate_size: int | None = None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+        self.act_fn = ACT2FN[config.hidden_act]
+        self.limit = config.swiglu_limit
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # SwiGLU with a symmetric pre-activation clip: the gate branch is capped
+        # from above, the value branch on both sides, before they are multiplied
+        # and projected back down.
+        activated = self.act_fn(self.gate_proj(x).clamp(max=self.limit))
+        value = self.up_proj(x).clamp(min=-self.limit, max=self.limit)
+        return self.down_proj(activated * value)
+
+
+@use_experts_implementation
+class AgnesExperts(nn.Module):
+    """Routed experts held as stacked 3-D weight tensors (one slab per expert)."""
+
+    def __init__(self, config: AgnesConfig):
+        super().__init__()
+        self.num_experts = config.num_local_experts
+        self.hidden_dim = config.hidden_size
+        self.intermediate_dim = config.intermediate_size
+        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
+        self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
+        self.act_fn = ACT2FN[config.hidden_act]
+        self.limit = config.swiglu_limit
+
+    def forward(
+        self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor
+    ) -> torch.Tensor:
+        # Reference (unfused) dispatch: walk only the experts that at least one
+        # token selected, gather that expert's tokens, run its SwiGLU, and
+        # scatter-add the weighted result back. The heavy fused kernels are
+        # swapped in by `@use_experts_implementation`; this body is the fallback.
+        out = torch.zeros_like(hidden_states)
+        with torch.no_grad():
+            # [experts, slot, token] boolean routing tensor
+            routing = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0)
+            used = (routing.sum(dim=(-1, -2)) > 0).nonzero().flatten()
+        for e in used.tolist():
+            if e == self.num_experts:
+                continue
+            slot, tok = torch.where(routing[e])
+            hidden = self._apply_gate(F.linear(hidden_states[tok], self.gate_up_proj[e]))
+            hidden = F.linear(hidden, self.down_proj[e]) * top_k_weights[tok, slot, None]
+            out.index_add_(0, tok, hidden.to(out.dtype))
+        return out
+
+    def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor:
+        # Method (not inlined) so the fused grouped_mm / batched_mm expert backends
+        # from `@use_experts_implementation` apply the identical clip + SiLU on top
+        # of their packed gate/up output instead of bypassing it.
+        gate, value = gate_up.chunk(2, dim=-1)
+        gated = self.act_fn(gate.clamp(max=self.limit))
+        return gated * value.clamp(min=-self.limit, max=self.limit)
+
+
+class AgnesTopKRouter(nn.Module):
+    def __init__(self, config: AgnesConfig):
+        super().__init__()
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_local_experts
+        self.hidden_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim))
+        self.score_fn = ACT2FN[config.scoring_func]
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.register_buffer("e_score_correction_bias", torch.zeros(self.num_experts), persistent=True)
+
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        tokens = hidden_states.reshape(-1, self.hidden_dim)
+        router_logits = F.linear(tokens, self.weight)
+        affinity = self.score_fn(router_logits)
+        # bias only steers the top-k *selection*; the gathered gate values are the
+        # un-biased affinities, renormalised to sum to one over the chosen experts.
+        chosen = torch.topk(affinity + self.e_score_correction_bias, self.top_k, dim=-1, sorted=False).indices
+        gate = affinity.gather(1, chosen)
+        gate = gate / (gate.sum(dim=-1, keepdim=True) + 1e-20)
+        return router_logits, gate * self.routed_scaling_factor, chosen
+
+
+class AgnesHashRouter(nn.Module):
+    r"""
+    Static hash routing used by the leading `agnes_hash_moe` MoE layers. Which
+    experts a token visits is fixed up front by a frozen `tid2eid[input_ids]`
+    lookup table (token id -> expert ids) rather than by a learned arg-top-k. The
+    learned `weight` is still evaluated to produce the affinities that *weight*
+    those experts' outputs; only the selection itself is frozen.
+    """
+
+    def __init__(self, config: AgnesConfig):
+        super().__init__()
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_local_experts
+        self.hidden_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim))
+        self.score_fn = ACT2FN[config.scoring_func]
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.register_buffer("tid2eid", torch.zeros(config.vocab_size, self.top_k, dtype=torch.long), persistent=True)
+
+    def forward(
+        self, hidden_states: torch.Tensor, input_ids: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        tokens = hidden_states.reshape(-1, self.hidden_dim)
+        router_logits = F.linear(tokens, self.weight)
+        affinity = self.score_fn(router_logits)
+        chosen = self.tid2eid[input_ids.reshape(-1)].long()  # frozen per-token expert ids
+        gate = affinity.gather(1, chosen)
+        gate = gate / (gate.sum(dim=-1, keepdim=True) + 1e-20)
+        return router_logits, gate * self.routed_scaling_factor, chosen
+
+
+class AgnesSparseMoeBlock(nn.Module):
+    def __init__(self, config: AgnesConfig, layer_idx: int):
+        super().__init__()
+        self.is_hash = config.mlp_layer_types[layer_idx] == "agnes_hash_moe"
+        self.gate = AgnesHashRouter(config) if self.is_hash else AgnesTopKRouter(config)
+        self.experts = AgnesExperts(config)
+        self.shared_experts = AgnesMLP(config)
+        pffn_size = getattr(config, "parallel_ffn_intermediate_size", 0) or 0
+        # parallel dense FFN branch; down_proj is zero-initialized at export so the
+        # block is exactly identity-preserving until trained
+        self.parallel_ffn = (
+            AgnesMLP(config, intermediate_size=pffn_size) if pffn_size and not self.is_hash else None
+        )
+
+    def forward(self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None) -> torch.Tensor:
+        b, s, d = hidden_states.shape
+        tokens = hidden_states.reshape(-1, d)
+        # hash layers additionally consume the raw token ids for the frozen lookup
+        if self.is_hash:
+            _, gate_w, gate_idx = self.gate(hidden_states, input_ids)
+        else:
+            _, gate_w, gate_idx = self.gate(hidden_states)
+        result = self.experts(tokens, gate_idx, gate_w).view(b, s, d)
+        result = result + self.shared_experts(hidden_states)
+        if self.parallel_ffn is not None:  # zero-init identity branch until trained
+            result = result + self.parallel_ffn(hidden_states)
+        return result
+
+
+
+# ==========================================================================
+# Residual hyper-connections (mHC)
+# ==========================================================================
+
+class AgnesHyperConnection(nn.Module):
+    r"""
+    Manifold-constrained hyper-connection (mHC). Where a plain block does a
+    single scalar residual add, Agnes keeps `hc_mult` parallel residual streams
+    and learns how to (i) collapse them into the sublayer input, (ii) place the
+    sublayer output back across them, and (iii) re-mix the streams themselves.
+
+    A single small projection (`fn`, biased by `base`, per-branch `scale`) maps
+    the `hc_mult` streams — after an unweighted RMS norm and a flatten — into
+    three groups of logits:
+
+    * `pre`  (length `hc_mult`)         : sigmoid gate that collapses the streams
+                                          into one sequence for the sublayer.
+    * `post` (length `hc_mult`)         : `2·sigmoid` gate in [0, 2] that spreads
+                                          the sublayer output back over streams.
+    * `comb` (`hc_mult × hc_mult`)      : softmaxed then Sinkhorn-projected onto
+                                          the doubly-stochastic manifold; mixes
+                                          the streams across the residual.
+
+    The decoder layer holds two of these, one at the attention site and one at
+    the MLP site.
+    """
+
+    def __init__(self, config: AgnesConfig):
+        super().__init__()
+        self.hc_mult = config.hc_mult
+        self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
+        self.hc_eps = config.hc_eps
+        self.input_norm = AgnesUnweightedRMSNorm(eps=config.rms_norm_eps)
+        n_logits = (2 + self.hc_mult) * self.hc_mult
+        self.fn = nn.Parameter(torch.empty(n_logits, self.hc_mult * config.hidden_size))
+        self.base = nn.Parameter(torch.empty(n_logits))
+        # one learned scale each for the pre / post / comb logit groups
+        self.scale = nn.Parameter(torch.empty(3))
+
+    def forward(self, hidden_streams: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        r"""Return `(post, comb, collapsed)`: the output-placement gate, the
+        doubly-stochastic stream-mixing matrix, and the collapsed sublayer input.
+        `comb` is driven onto the doubly-stochastic manifold by Sinkhorn-Knopp —
+        an initial column normalisation followed by `hc_sinkhorn_iters - 1`
+        alternating row/column passes."""
+        n = self.hc_mult
+        normed = self.input_norm(hidden_streams.flatten(start_dim=2).float())
+        pre_raw, post_raw, comb_raw = F.linear(normed, self.fn.float()).split([n, n, n * n], dim=-1)
+        pre_bias, post_bias, comb_bias = self.base.split([n, n, n * n])
+        pre_scale, post_scale, comb_scale = self.scale.unbind(0)
+
+        pre = torch.sigmoid(pre_raw * pre_scale + pre_bias) + self.hc_eps
+        post = 2 * torch.sigmoid(post_raw * post_scale + post_bias)
+
+        comb = comb_raw.view(*comb_raw.shape[:-1], n, n) * comb_scale + comb_bias.view(n, n)
+        comb = torch.softmax(comb, dim=-1) + self.hc_eps
+        comb = comb / (comb.sum(dim=-2, keepdim=True) + self.hc_eps)
+        for _ in range(self.hc_sinkhorn_iters - 1):
+            comb = comb / (comb.sum(dim=-1, keepdim=True) + self.hc_eps)
+            comb = comb / (comb.sum(dim=-2, keepdim=True) + self.hc_eps)
+
+        # weighted sum over the stream axis -> single sequence fed to the sublayer
+        collapsed = (pre.unsqueeze(-1) * hidden_streams).sum(dim=2).to(hidden_streams.dtype)
+        return post, comb, collapsed
+
+
+class AgnesHyperHead(nn.Module):
+    """Collapses the `hc_mult` residual streams back into one tensor at the top
+    of the stack, just before the final shared RMS norm in `AgnesModel`."""
+
+    def __init__(self, config: AgnesConfig):
+        super().__init__()
+        self.hc_mult = config.hc_mult
+        self.input_norm = AgnesUnweightedRMSNorm(eps=config.rms_norm_eps)
+        self.eps = config.hc_eps
+        self.hc_fn = nn.Parameter(torch.empty(self.hc_mult, self.hc_mult * config.hidden_size))
+        self.hc_base = nn.Parameter(torch.empty(self.hc_mult))
+        self.hc_scale = nn.Parameter(torch.empty(1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # same gate-and-sum as the `pre` branch of AgnesHyperConnection, but with
+        # no sublayer to feed — this is the terminal stream collapse.
+        gate_logits = F.linear(self.input_norm(x.flatten(2).float()), self.hc_fn.float())
+        gate = torch.sigmoid(gate_logits * self.hc_scale.float() + self.hc_base.float()) + self.eps
+        return (gate.unsqueeze(-1) * x).sum(dim=2).to(x.dtype)
+
+
+
+# ==========================================================================
+# Rotary position embedding
+# ==========================================================================
+
+class AgnesRotaryEmbedding(nn.Module):
+    """
+    Rotary embedding that maintains one inverse-frequency buffer per *rope
+    label* and applies interleaved (rather than half-split) rotation with
+    partial coverage of each head.
+
+    Two things are specific to Agnes here:
+
+    * Interleaved layout — a single `θ_i` is shared by the two channels of
+      each consecutive pair, so only `rope_head_dim // 2` frequencies are
+      stored. The doubling back up to the full rope width happens next to the
+      rotation itself (in `apply_agnes_rope`) instead of being baked into the
+      cached cos/sin, which keeps the pairing explicit at the point of use.
+
+    * Rope labels are independent of the attention schedule. The architectural
+      `layer_types` (`agnes_local_attention` / `agnes_sparse_attention` /
+      `agnes_pooled_attention`) are decoupled from the rope labels
+      (`main` / `compress`), which live as keys in `config.rope_parameters`
+      and differ only in their `rope_theta` base. Buffers are therefore built
+      by iterating `rope_parameters` rather than the attention schedule.
+    """
+
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: AgnesConfig):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+        self.config = config
+        # Only the nested per-rope-type sub-dicts are real layer types — the top-level
+        # `rope_type` key that ``convert_rope_params_to_dict`` may leave on
+        # ``config.rope_parameters`` is a flat-shape leftover, not a layer.
+        self.layer_types = [k for k, v in config.rope_parameters.items() if isinstance(v, dict)]
+        self.rope_type = {}
+        for layer_type in self.layer_types:
+            rope_params = config.rope_parameters[layer_type]
+            self.rope_type[layer_type] = rope_params["rope_type"]
+            rope_init_fn = self.compute_default_rope_parameters
+            if self.rope_type[layer_type] != "default":
+                rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type[layer_type]]
+            inv_freq, attention_scaling = rope_init_fn(config, layer_type=layer_type)
+            self.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False)
+            self.register_buffer(f"{layer_type}_original_inv_freq", inv_freq.clone(), persistent=False)
+            setattr(self, f"{layer_type}_attention_scaling", attention_scaling)
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: AgnesConfig | None = None,
+        device: Optional["torch.device"] = None,
+        seq_len: int | None = None,
+        layer_type: str | None = None,
+    ) -> tuple["torch.Tensor", float]:
+        """Plain (non-scaled) inverse frequencies for one rope label.
+
+        Returns `(inv_freq, attention_factor)`. `attention_factor` is always 1.0
+        for this default variant (the scaled variants such as YaRN come from
+        `ROPE_INIT_FUNCTIONS` instead). Only the leading `partial_rotary_factor`
+        fraction of each head receives rotation, so the frequency table has
+        `rope_dim // 2` entries. `seq_len` is accepted for signature parity and
+        unused here.
+        """
+        rope_cfg = config.rope_parameters[layer_type]
+        base = rope_cfg["rope_theta"]
+        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        rope_dim = int(head_dim * rope_cfg.get("partial_rotary_factor", 1.0))
+
+        exponents = torch.arange(0, rope_dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / rope_dim
+        inv_freq = 1.0 / (base**exponents)
+        return inv_freq, 1.0
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids, layer_type=None):
+        # cos/sin carry one entry per interleaved pair; the widening to full rope
+        # width is left to `apply_agnes_rope` so the pairing stays where it is used.
+        inv_freq = getattr(self, f"{layer_type}_inv_freq")
+        scaling = getattr(self, f"{layer_type}_attention_scaling")
+        # outer product positions (B, S) x frequencies (rope_dim/2) -> (B, S, rope_dim/2)
+        freq_col = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        pos_row = position_ids[:, None, :].float()
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with maybe_autocast(device_type=device_type, enabled=False):
+            angles = (freq_col.float() @ pos_row.float()).transpose(1, 2)
+            cos = angles.cos() * scaling
+            sin = angles.sin() * scaling
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def _swap_pairs(x: torch.Tensor) -> torch.Tensor:
+    """Interleaved-pair rotation partner: within every adjacent `(a, b)` pair
+    return `(-b, a)`. Equivalent to the `[-x_odd, x_even]` interleave used by
+    the rotation formula below."""
+    pairs = x.unflatten(-1, (-1, 2))
+    a, b = pairs[..., 0], pairs[..., 1]
+    return torch.stack((-b, a), dim=-1).flatten(-2)
+
+
+def apply_agnes_rope(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, unsqueeze_dim: int = 1
+) -> torch.Tensor:
+    """Apply interleaved rotary position embedding to the trailing rope slice.
+
+    The cached `cos` / `sin` carry one entry per interleaved pair. They are
+    widened to the full rope dimension right here (`repeat_interleave(2)`), the
+    leading "no-position" channels are passed through untouched, and only the
+    final `2 * cos.shape[-1]` channels are rotated via
+    `x * cos + swap(x) * sin`. The rotation math is done in float32 and cast
+    back to `x`'s dtype. Each head is laid out as `[nope | rope]`.
+    """
+    cos = cos.repeat_interleave(2, dim=-1).unsqueeze(unsqueeze_dim)
+    sin = sin.repeat_interleave(2, dim=-1).unsqueeze(unsqueeze_dim)
+    rope_dim = cos.shape[-1]
+    split_at = x.shape[-1] - rope_dim
+    nope, rope = x[..., :split_at], x[..., split_at:]
+    rotated = (rope.float() * cos + _swap_pairs(rope).float() * sin).to(x.dtype)
+    return torch.cat([nope, rotated], dim=-1)
+
+
+
+# ==========================================================================
+# Attention KV caches
+# ==========================================================================
+
+class AgnesHCACache(DynamicSlidingWindowLayer):
+    r"""Cache layer for HCA blocks. Holds the long-range compressor's
+    buffer / running compressed entries / count on top of the sliding-window K=V
+    branch. HCA uses *non-overlapping* windows, so there is *no* overlap state,
+    and HCA has *no* indexer either.
+
+    State is dict-keyed by entry name — HCA only uses `"compressor"`, but
+    :class:`AgnesCSACache` adds `"indexer"` to the same dicts so a single
+    set of methods (`store_compression_weights` / `update_compressor_states`)
+    serves both:
+
+      * `compressed_kv[name]` — the running list of compressed KV entries
+        emitted so far (one every `compress_rate` source tokens; the long-range
+        KVs the attention concatenates onto its sliding-window keys / values).
+      * `buffer_kv[name]` / `buffer_gate[name]` — source tokens that arrived
+        between two full windows; once the buffer hits `compress_rate` tokens
+        the compressor closes a window, emits one entry, and drains the buffer.
+      * `entry_count[name]` — number of compressed entries emitted so far, so
+        `entry_count[name] * compress_rate` is the absolute position of the
+        *next* window's first source token. Tracked separately from
+        `position_ids` so prefill -> decode -> prefill stays consistent.
+    """
+
+    layer_type = "agnes_pooled_attention"
+
+    def __init__(self, config: "AgnesConfig"):
+        super().__init__(config)
+        self.compress_rate = config.compress_rates["agnes_pooled_attention"]
+        self.buffer_kv: dict[str, torch.Tensor | None] = {"compressor": None}
+        self.buffer_gate: dict[str, torch.Tensor | None] = {"compressor": None}
+        self.compressed_kv: dict[str, torch.Tensor | None] = {"compressor": None}
+        self.entry_count: dict[str, int] = {"compressor": 0}
+
+    def update(self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs):
+        """Sliding-window key/value update. Because Agnes is shared-KV MQA, keys
+        and values are one and the same buffer, so both returns are the full
+        (pre-trim) concatenation while the retained cache keeps only the last
+        `sliding_window - 1` positions."""
+        if not self.is_initialized:
+            self.lazy_initialization(key_states, value_states)
+            self.values = self.keys
+        self.cumulative_length += key_states.shape[-2]
+        combined = torch.cat([self.keys, key_states], dim=-2)
+        self.keys = combined[:, :, -self.sliding_window + 1 :, :]
+        self.values = self.keys
+        return combined, combined
+
+    def store_compression_weights(
+        self, name: str, kv: torch.Tensor, gate: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, int]:
+        r"""Prepend the leftover buffered `(kv, gate)` for entry `name`, split
+        off the longest window-aligned prefix (the part that can be compressed
+        now), stash the remainder back in the buffer, and hand back
+        `(chunk_kv, chunk_gate, first_window_position)`. The compressor then
+        softmax-pools each `compress_rate`-token window of that chunk (with its
+        `position_bias`) into one compressed entry.
+        """
+        first_window_position = self.entry_count[name] * self.compress_rate
+        held_kv, held_gate = self.buffer_kv[name], self.buffer_gate[name]
+        if held_kv is not None and held_kv.shape[1]:
+            kv = torch.cat([held_kv, kv], dim=1)
+            gate = torch.cat([held_gate, gate], dim=1)
+        # split at the largest multiple of compress_rate; carry the tail forward
+        cut = (kv.shape[1] // self.compress_rate) * self.compress_rate
+        self.buffer_kv[name], self.buffer_gate[name] = kv[:, cut:], gate[:, cut:]
+        return kv[:, :cut], gate[:, :cut], first_window_position
+
+    def update_compressor_states(self, name: str, compressed: torch.Tensor) -> torch.Tensor:
+        r"""Append the newly emitted compressed entries to the running
+        `compressed_kv[name]`, advance `entry_count[name]`, and return the
+        accumulated tensor."""
+        running = self.compressed_kv[name]
+        if running is None:
+            self.compressed_kv[name] = compressed
+        elif compressed.shape[1] > 0:
+            self.compressed_kv[name] = torch.cat([running, compressed], dim=1)
+        self.entry_count[name] += compressed.shape[1]
+        return self.compressed_kv[name]
+
+
+class AgnesCSACache(AgnesHCACache):
+    r"""CSA cache. On top of :class:`AgnesHCACache` it registers a second entry
+    name, `"indexer"`, in each of the inherited state dicts, and it keeps a small
+    per-name *overlap* buffer required by the two-series windowing.
+
+    Why the overlap buffer exists: for CSA, `kv_proj` / `gate_proj` emit
+    `2 * head_dim` features per token — two interleaved series, Ca in
+    `[..., :head_dim]` and Cb in `[..., head_dim:]`. The pooled entry for window
+    `w` blends window `w-1`'s Ca with window `w`'s Cb (softmax-gated), giving an
+    effective receptive width of `2 * compress_rate_csa` at stride
+    `compress_rate_csa`. The only cross-window dependency is therefore the
+    previous window's Ca slice, so at a forward boundary we persist exactly
+    `chunk[:, -1, :, :head_dim]` (Ca of the last full window) in
+    `overlap_kv[name]` / `overlap_gate[name]`; Cb is never revisited.
+    """
+
+    layer_type = "agnes_sparse_attention"
+
+    def __init__(self, config: "AgnesConfig"):
+        super().__init__(config)
+        self.compress_rate = config.compress_rates["agnes_sparse_attention"]
+        self.buffer_kv["indexer"] = None
+        self.buffer_gate["indexer"] = None
+        self.compressed_kv["indexer"] = None
+        self.entry_count["indexer"] = 0
+        self.overlap_kv: dict[str, torch.Tensor | None] = {"compressor": None, "indexer": None}
+        self.overlap_gate: dict[str, torch.Tensor | None] = {"compressor": None, "indexer": None}
+
+    def update_overlap_state(
+        self, name: str, chunk_kv: torch.Tensor, chunk_gate: torch.Tensor, head_dim: int
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        r"""Swap the overlap buffer: return the Ca slice saved on the previous
+        call (or `None` on the first) and store this call's last-window Ca slice
+        for the next one. Only `[..., :head_dim]` (Ca) is kept, since Cb is
+        already baked into an emitted entry and never read again.
+        """
+        carried_kv, carried_gate = self.overlap_kv[name], self.overlap_gate[name]
+        self.overlap_kv[name] = chunk_kv[:, -1, :, :head_dim].clone()
+        self.overlap_gate[name] = chunk_gate[:, -1, :, :head_dim].clone()
+        return carried_kv, carried_gate
+
+
+
+# ==========================================================================
+# Long-range compressors and lightning indexer
+# ==========================================================================
+
+class AgnesHCACompressor(nn.Module):
+    """
+    Heavily Compressed Attention compressor. Compresses
+    every `compress_rate_hca` (m'=128) source tokens into a single compressed KV
+    entry.
+
+    Each closed window of m' tokens produces one compressed entry:
+    `C^{Comp}_i = Σ_{j∈window} softmax(Z_j + B)_j ⊙ C_j`. RoPE on the trailing
+    `rope_head_dim` slice is applied at the deterministic absolute position
+    `i * compress_rate_hca + first_window_position` so cross-call concatenation
+    stays causality-correct. Returns the running list of *all* compressed
+    entries emitted so far (shape `[B, 1, T, head_dim]` with
+    `T = entry_count["compressor"]`), so the attention can attend over the
+    full long-range history.
+
+    When `past_key_values is None` runs in stateless single-shot mode: compress
+    every complete window from `hidden_states` and discard the remainder
+    (instead of caching it).
+    """
+
+    rope_layer_type = "compress"
+
+    def __init__(self, config: AgnesConfig):
+        super().__init__()
+        self.compress_rate = config.compress_rates["agnes_pooled_attention"]
+        self.head_dim = config.head_dim
+        self.kv_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False)
+        self.gate_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False)
+        self.position_bias = nn.Parameter(torch.empty(self.compress_rate, self.head_dim))
+        self.kv_norm = AgnesRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.rotary_emb = AgnesRotaryEmbedding(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        q_residual: torch.Tensor,
+        position_ids: torch.Tensor,
+        past_key_values: Cache | None,
+        layer_idx: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch = hidden_states.shape[0]
+        cache_layer: AgnesHCACache = past_key_values.layers[layer_idx] if past_key_values is not None else None
+        kv = self.kv_proj(hidden_states)
+        gate = self.gate_proj(hidden_states)
+        # stateless mode drops the ragged tail; cached mode buffers it for later
+        if cache_layer is None:
+            aligned = (kv.shape[1] // self.compress_rate) * self.compress_rate
+            chunk_kv, chunk_gate, first_window_position = kv[:, :aligned], gate[:, :aligned], 0
+        else:
+            chunk_kv, chunk_gate, first_window_position = cache_layer.store_compression_weights("compressor", kv, gate)
+
+        if chunk_kv.shape[1] > 0:  # at least one full window is ready
+            n_windows = chunk_kv.shape[1] // self.compress_rate
+            chunk_kv = chunk_kv.view(batch, n_windows, self.compress_rate, -1)
+            biased_gate = chunk_gate.view(batch, n_windows, self.compress_rate, -1) + self.position_bias
+            # softmax pooling within each window (fp32 for stability), then norm
+            pool_w = biased_gate.softmax(dim=2, dtype=torch.float32).to(chunk_kv.dtype)
+            compressed = self.kv_norm((chunk_kv * pool_w).sum(dim=2))
+            abs_pos = torch.arange(n_windows, device=compressed.device) * self.compress_rate + first_window_position
+            abs_pos = abs_pos.unsqueeze(0).expand(batch, -1)
+            cos, sin = self.rotary_emb(compressed, position_ids=abs_pos, layer_type=self.rope_layer_type)
+            compressed = apply_agnes_rope(compressed.unsqueeze(1), cos, sin).squeeze(1)
+        else:
+            compressed = chunk_kv.new_zeros((batch, 0, self.head_dim))
+
+        if cache_layer is not None:
+            compressed = cache_layer.update_compressor_states("compressor", compressed)
+        compressed_kv = compressed.unsqueeze(1)
+
+        n_entries = compressed_kv.shape[2]
+        seq_len = position_ids.shape[1]
+        if seq_len == 1 or n_entries == 0:
+            return compressed_kv, None
+
+        # a query at position t may only see entry w once t has passed window w,
+        # i.e. w < (t + 1) // compress_rate; everything else is masked to -inf.
+        entry_pos = torch.arange(n_entries, device=compressed_kv.device)
+        ready = (position_ids + 1) // self.compress_rate  # [B, S]
+        block_bias = compressed_kv.new_zeros((batch, 1, seq_len, n_entries))
+        block_bias = block_bias.masked_fill(
+            entry_pos.view(1, 1, 1, -1) >= ready.unsqueeze(1).unsqueeze(-1),
+            float("-inf"),
+        )
+        return compressed_kv, block_bias
+
+
+class AgnesIndexerScorer(nn.Module):
+    r"""Lightning-indexer scoring head: `∑_h w_{t,h} · ReLU(q_{t,h} · K^IComp_s)`."""
+
+    def __init__(self, config: AgnesConfig):
+        super().__init__()
+        self.softmax_scale = config.index_head_dim**-0.5
+        self.weights_scaling = config.index_n_heads**-0.5
+        self.weights_proj = nn.Linear(config.hidden_size, config.index_n_heads, bias=False)
+
+    def forward(self, q: torch.Tensor, compressed_kv: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
+        # per-head query·key, ReLU-rectified and scaled -> [B, S, H, T]
+        qk = torch.matmul(q.float(), compressed_kv.transpose(-1, -2).float().unsqueeze(1))
+        qk = F.relu(qk) * self.softmax_scale
+        head_w = self.weights_proj(hidden_states).float() * self.weights_scaling  # [B, S, H]
+        # collapse the head axis with the learned per-head weights -> [B, S, T]
+        return (qk * head_w.unsqueeze(-1)).sum(dim=2)
+
+
+class AgnesIndexer(nn.Module):
+    r"""Lightning indexer for CSA. For each query it keeps only the top
+    `config.index_topk` of the compressed entries, cutting the effective KV set
+    from `seq_len / compress_rate_csa` down to `index_topk`.
+
+    It builds a miniature compressor of its own at `index_head_dim` over the
+    same windows as the outer CSA compressor, scores each query against those
+    compressed keys as `sum_h w_h * ReLU(q_h . k)`, and returns the winning
+    indices.
+
+    A dedicated rotary lives here because RoPE has to be applied twice, to two
+    tensors whose positions differ: the compressed keys sit at the fixed window
+    positions `i * compress_rate + first_window_position`, whereas the queries
+    sit at the (per-call) `position_ids`. Both use the compressor's theta
+    (`compress_rope_theta`) so that only relative offsets survive the dot
+    product; since the query positions change every call, cos/sin can't be
+    precomputed, so the indexer calls its rotary twice per forward (label
+    always `"compress"`).
+    """
+
+    rope_layer_type = "compress"
+
+    def __init__(self, config: AgnesConfig):
+        super().__init__()
+        self.compress_rate = config.compress_rates["agnes_sparse_attention"]
+        self.num_heads = config.index_n_heads
+        self.head_dim = config.index_head_dim
+        self.index_topk = config.index_topk
+        self.kv_proj = nn.Linear(config.hidden_size, 2 * self.head_dim, bias=False)
+        self.gate_proj = nn.Linear(config.hidden_size, 2 * self.head_dim, bias=False)
+        self.position_bias = nn.Parameter(torch.empty(self.compress_rate, 2 * self.head_dim))
+        self.kv_norm = AgnesRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.head_dim, bias=False)
+        self.rotary_emb = AgnesRotaryEmbedding(config)
+        self.scorer = AgnesIndexerScorer(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        q_residual: torch.Tensor,
+        position_ids: torch.Tensor,
+        past_key_values: Cache | None,
+        layer_idx: int,
+    ) -> torch.LongTensor:
+        batch, seq_len, _ = hidden_states.shape
+        cache_layer: AgnesCSACache = past_key_values.layers[layer_idx] if past_key_values is not None else None
+        kv = self.kv_proj(hidden_states)
+        gate = self.gate_proj(hidden_states)
+
+        if cache_layer is None:
+            aligned = (kv.shape[1] // self.compress_rate) * self.compress_rate
+            chunk_kv, chunk_gate, first_window_position = kv[:, :aligned], gate[:, :aligned], 0
+        else:
+            chunk_kv, chunk_gate, first_window_position = cache_layer.store_compression_weights("indexer", kv, gate)
+
+        if chunk_kv.shape[1] > 0:
+            n_windows = chunk_kv.shape[1] // self.compress_rate
+            m = self.compress_rate
+            chunk_kv = chunk_kv.view(batch, n_windows, m, -1)
+            chunk_gate = chunk_gate.view(batch, n_windows, m, -1) + self.position_bias
+
+            # widen each window to 2m slots: current-window Cb in the top half,
+            # previous-window Ca in the bottom half (same scheme as outer CSA).
+            win_kv = chunk_kv.new_zeros((batch, n_windows, 2 * m, self.head_dim))
+            win_gate = chunk_gate.new_full((batch, n_windows, 2 * m, self.head_dim), float("-inf"))
+            win_kv[:, :, m:] = chunk_kv[..., self.head_dim :]
+            win_gate[:, :, m:] = chunk_gate[..., self.head_dim :]
+            if n_windows > 1:
+                win_kv[:, 1:, :m] = chunk_kv[:, :-1, :, : self.head_dim]
+                win_gate[:, 1:, :m] = chunk_gate[:, :-1, :, : self.head_dim]
+            if cache_layer is not None:
+                carried_kv, carried_gate = cache_layer.update_overlap_state("indexer", chunk_kv, chunk_gate, self.head_dim)
+                if carried_kv is not None:
+                    win_kv[:, 0, :m] = carried_kv.to(win_kv.dtype)
+                    win_gate[:, 0, :m] = carried_gate.to(win_gate.dtype)
+
+            pool_w = win_gate.softmax(dim=2, dtype=torch.float32).to(win_kv.dtype)
+            compressed = self.kv_norm((win_kv * pool_w).sum(dim=2))
+            abs_pos = torch.arange(n_windows, device=compressed.device) * self.compress_rate + first_window_position
+            abs_pos = abs_pos.unsqueeze(0).expand(batch, -1)
+            cos, sin = self.rotary_emb(compressed, position_ids=abs_pos, layer_type=self.rope_layer_type)
+            compressed = apply_agnes_rope(compressed.unsqueeze(1), cos, sin).squeeze(1)
+        else:
+            compressed = chunk_kv.new_zeros((batch, 0, self.head_dim))
+
+        compressed_kv = (
+            compressed if cache_layer is None else cache_layer.update_compressor_states("indexer", compressed)
+        )
+
+        cos_q, sin_q = self.rotary_emb(hidden_states, position_ids=position_ids, layer_type=self.rope_layer_type)
+        q = self.q_b_proj(q_residual).view(batch, seq_len, -1, self.head_dim).transpose(1, 2)
+        q = apply_agnes_rope(q, cos_q, sin_q).transpose(1, 2)
+
+        scores = self.scorer(q, compressed_kv, hidden_states)  # [B, S, T]
+        n_entries = compressed_kv.shape[1]
+        k = min(self.index_topk, n_entries)
+
+        # a query can only see entries whose window has fully closed before it; the
+        # `topk` is taken after masking future entries to -inf, and any pick that
+        # still points past the ready frontier (too few blocks yet) is flagged -1.
+        if n_entries > 0:
+            ready = (position_ids + 1) // self.compress_rate  # [B, S]
+            entry_pos = torch.arange(n_entries, device=scores.device)
+            future = entry_pos.view(1, 1, -1) >= ready.unsqueeze(-1)  # [B, S, T]
+            scores = scores.masked_fill(future, float("-inf"))
+            picks = scores.topk(k, dim=-1).indices  # [B, S, k]
+            stale = picks >= ready.unsqueeze(-1)
+            return torch.where(stale, torch.full_like(picks, -1), picks)
+
+        return scores.topk(k, dim=-1).indices
+
+
+class AgnesCSACompressor(nn.Module):
+    """Compressed Sparse Attention compressor. It pools every `compress_rate_csa`
+    (m=4) source tokens into one entry and pairs the result with a Lightning
+    Indexer that keeps only the top `index_topk` entries per query before core
+    attention runs.
+
+    `kv_proj` / `gate_proj` / `position_bias` all emit `2 * head_dim` features,
+    holding two series per token: Ca in `[..., :head_dim]` (contributes to the
+    *next* window's entry) and Cb in `[..., head_dim:]` (contributes to the
+    *current* one). Entry `w` softmax-blends window `w-1`'s Ca with window `w`'s
+    Cb across `2 * compress_rate_csa` slots — width `2 * compress_rate_csa`,
+    stride `compress_rate_csa`. Window 0 needs the previous forward's last Ca
+    slice, which the cache returns via `overlap_kv`; with no cache (or on the
+    first call) that half is left as zero-kv / `-inf`-gate so it contributes
+    nothing to the softmax.
+    """
+
+    rope_layer_type = "compress"
+
+    def __init__(self, config: AgnesConfig):
+        super().__init__()
+        self.compress_rate = config.compress_rates["agnes_sparse_attention"]
+        self.head_dim = config.head_dim
+        self.kv_proj = nn.Linear(config.hidden_size, 2 * self.head_dim, bias=False)
+        self.gate_proj = nn.Linear(config.hidden_size, 2 * self.head_dim, bias=False)
+        self.position_bias = nn.Parameter(torch.empty(self.compress_rate, 2 * self.head_dim))
+        self.kv_norm = AgnesRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.rotary_emb = AgnesRotaryEmbedding(config)
+        self.indexer = AgnesIndexer(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        q_residual: torch.Tensor,
+        position_ids: torch.Tensor,
+        past_key_values: Cache | None,
+        layer_idx: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch, seq_len, _ = hidden_states.shape
+        cache_layer: AgnesCSACache = past_key_values.layers[layer_idx] if past_key_values is not None else None
+        kv = self.kv_proj(hidden_states)
+        gate = self.gate_proj(hidden_states)
+
+        if cache_layer is None:
+            aligned = (kv.shape[1] // self.compress_rate) * self.compress_rate
+            chunk_kv, chunk_gate, first_window_position = kv[:, :aligned], gate[:, :aligned], 0
+        else:
+            chunk_kv, chunk_gate, first_window_position = cache_layer.store_compression_weights("compressor", kv, gate)
+
+        if chunk_kv.shape[1] > 0:
+            n_windows = chunk_kv.shape[1] // self.compress_rate
+            m = self.compress_rate
+            chunk_kv = chunk_kv.view(batch, n_windows, m, -1)
+            chunk_gate = chunk_gate.view(batch, n_windows, m, -1) + self.position_bias
+
+            # 2m-wide slot layout: top half = current window's Cb, bottom half =
+            # previous window's Ca. For window 0 the bottom half is either the Ca
+            # slice returned by the cache, or (no cache / first call) left as
+            # zero-kv + -inf-gate so it drops out of the softmax.
+            win_kv = chunk_kv.new_zeros((batch, n_windows, 2 * m, self.head_dim))
+            win_gate = chunk_gate.new_full((batch, n_windows, 2 * m, self.head_dim), float("-inf"))
+            win_kv[:, :, m:] = chunk_kv[..., self.head_dim :]
+            win_gate[:, :, m:] = chunk_gate[..., self.head_dim :]
+            if n_windows > 1:
+                win_kv[:, 1:, :m] = chunk_kv[:, :-1, :, : self.head_dim]
+                win_gate[:, 1:, :m] = chunk_gate[:, :-1, :, : self.head_dim]
+            if cache_layer is not None:
+                carried_kv, carried_gate = cache_layer.update_overlap_state(
+                    "compressor", chunk_kv, chunk_gate, self.head_dim
+                )
+                if carried_kv is not None:
+                    win_kv[:, 0, :m] = carried_kv.to(win_kv.dtype)
+                    win_gate[:, 0, :m] = carried_gate.to(win_gate.dtype)
+
+            # fp32 softmax pooling: in bf16/fp16 close logits over wide windows
+            # can collapse together, so accumulate the weights in float.
+            pool_w = win_gate.softmax(dim=2, dtype=torch.float32).to(win_kv.dtype)
+            compressed = self.kv_norm((win_kv * pool_w).sum(dim=2))
+            abs_pos = torch.arange(n_windows, device=compressed.device) * self.compress_rate + first_window_position
+            abs_pos = abs_pos.unsqueeze(0).expand(batch, -1)
+            cos, sin = self.rotary_emb(compressed, position_ids=abs_pos, layer_type=self.rope_layer_type)
+            compressed = apply_agnes_rope(compressed.unsqueeze(1), cos, sin).squeeze(1)
+        else:
+            compressed = chunk_kv.new_zeros((batch, 0, self.head_dim))
+
+        if cache_layer is not None:
+            compressed = cache_layer.update_compressor_states("compressor", compressed)
+        compressed_kv = compressed.unsqueeze(1)
+
+        # ask the indexer which entries each query keeps; it returns -1 for picks
+        # that are not yet attendable. Clamp those to a scratch column, scatter
+        # zeros for the kept ones, then slice the scratch column back off so the
+        # bias is -inf everywhere except the valid picks.
+        picks = self.indexer(hidden_states, q_residual, position_ids, past_key_values, layer_idx)  # [B, S, k]
+        n_entries = compressed_kv.shape[2]
+        kept = picks >= 0  # [B, S, k]
+        gather_idx = torch.where(kept, picks, torch.full_like(picks, n_entries))
+        block_bias = compressed_kv.new_full((batch, 1, seq_len, n_entries + 1), float("-inf"))
+        block_bias.scatter_(-1, gather_idx.unsqueeze(1), 0.0)
+        return compressed_kv, block_bias[..., :n_entries]
+
+
+
+# ==========================================================================
+# Core attention
+# ==========================================================================
+
+class AgnesGroupedLinear(nn.Linear):
+    """Block-diagonal grouped linear used by the grouped output projection
+    The core attention's stacked output is `num_attention_heads* head_dim`-dim,
+    which is *very* large (Agnes-Flash: 32768; Agnes-Pro: 65536). A direct
+    `num_attention_heads*head_dim → hidden_size` projection would dominate the per-token cost.
+
+    Agnes sidesteps that by splitting the heads into `g` groups, projecting
+    each `num_attention_heads * head_dim/g`-dim group independently to a `d_g`-dim intermediate output
+    (with `d_g < num_attention_heads * head_dim/g`), and then mixing the resulting `g·d_g` vector to
+    `hidden_size` through a single follow-up linear (`self_attn.o_b_proj`). This
+    module owns the per-group block (`self_attn.o_a_proj`).
+
+    For Agnes-Flash (num_attention_heads=64, head_dim=512, o_groups=8, o_lora_rank=1024,
+    hidden_size=4096), g=8 groups of 4096-dim each are projected to 1024-dim, then
+    mixed to 4096-dim; for Agnes-Pro (num_attention_heads=128, head_dim=512, o_groups=16,
+    o_lora_rank=1024, hidden_size=7168), g=16 groups of 4096-dim each are projected
+    to 1024-dim, then mixed to 7168-dim.
+    """
+
+    def __init__(self, in_features_per_group: int, out_features: int, n_groups: int, bias: bool = False):
+        super().__init__(in_features_per_group, out_features, bias=bias)
+        self.n_groups = n_groups
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        lead = x.shape[:-2]
+        in_dim = x.shape[-1]
+        # per-group matmul: bring the group axis to the front so a single bmm
+        # applies each group's block independently.
+        blocks = self.weight.view(self.n_groups, -1, in_dim).transpose(1, 2)
+        grouped = x.reshape(-1, self.n_groups, in_dim).transpose(0, 1)
+        out = torch.bmm(grouped, blocks).transpose(0, 1)
+        return out.reshape(*lead, self.n_groups, -1)
+
+
+def _broadcast_kv_heads(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """Broadcast the single shared KV head out to `n_rep` query heads. Same
+    result as `torch.repeat_interleave(x, n_rep, dim=1)`, taking
+    `(B, n_kv, S, D)` to `(B, n_kv * n_rep, S, D)`, but expressed as an
+    `expand` + `reshape` to avoid materialising the intermediate copy.
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def _eager_attention(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float | int = 0.0,
+    **kwargs,
+):
+    # Shared-KV MQA: broadcast the one stored KV head across all query heads.
+    key_states = _broadcast_kv_heads(key, module.num_key_value_groups)
+    value_states = _broadcast_kv_heads(value, module.num_key_value_groups)
+    logits = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        logits = logits + attention_mask
+
+    # Append one learned per-head sink logit as an extra "always-visible" column,
+    # then take the softmax over [keys | sink] and read back only the key
+    # columns. The shift by the row max is a plain numerical-stability
+    # subtraction (matters in bf16/fp16, invariant of the softmax).
+    sink = module.sinks.reshape(1, -1, 1, 1).expand(query.shape[0], -1, query.shape[-2], -1)
+    logits = torch.cat([logits, sink], dim=-1)
+    logits = logits - logits.amax(dim=-1, keepdim=True)
+    probs = F.softmax(logits, dim=-1, dtype=logits.dtype)[..., :-1]
+    probs = nn.functional.dropout(probs, p=dropout, training=module.training).to(value_states.dtype)
+    attn_output = torch.matmul(probs, value_states).transpose(1, 2).contiguous()
+    return attn_output, probs
+
+
+_COMPRESSOR_BY_LAYER_TYPE = {
+    "agnes_local_attention": None,
+    "agnes_sparse_attention": AgnesCSACompressor,
+    "agnes_pooled_attention": AgnesHCACompressor,
+}
+
+
+class AgnesAttention(nn.Module):
+    r"""
+    Agnes attention block. It departs from a textbook multi-head block in five
+    ways:
+
+    * Shared-KV multi-query attention: `num_key_value_heads = 1`; `kv_proj`
+      emits that single KV head and it is read as both key and value.
+    * Partial, interleaved RoPE over the leading `rope_head_dim` of each head.
+      The conjugate rotation (position `-i`) is re-applied to the output rope
+      slice so each entry's contribution depends only on query/key relative
+      distance.
+    * A learned per-head attention-sink logit that absorbs probability mass
+      away from the real keys.
+    * A grouped low-rank output projection, to keep the wide stacked-head
+      output affordable.
+    * Three interchangeable cache regimes: plain sliding window, sliding+CSA,
+      and sliding+HCA.
+    """
+
+    def __init__(self, config: AgnesConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.layer_type = config.layer_types[layer_idx]
+        # Sliding-only layers use the "main" (plain θ=10000) rope; CSA/HCA layers
+        # share the same yarn-scaled "compress" rope as their compressor.
+        self.rope_layer_type = "main" if self.layer_type == "agnes_local_attention" else "compress"
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_groups = config.num_attention_heads  # single KV head, broadcast to all
+        self.head_dim = config.head_dim
+        self.sliding_window = config.sliding_window
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+        self.scaling = self.head_dim**-0.5
+
+        self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False)
+        self.q_a_norm = AgnesRMSNorm(config.q_lora_rank, eps=config.rms_norm_eps)
+        self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.head_dim, bias=False)
+        self.q_b_norm = AgnesUnweightedRMSNorm(eps=config.rms_norm_eps)
+        self.kv_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False)
+        self.kv_norm = AgnesRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.o_a_proj = AgnesGroupedLinear(
+            self.num_heads * self.head_dim // config.o_groups, config.o_groups * config.o_lora_rank, config.o_groups
+        )
+        self.o_b_proj = nn.Linear(config.o_groups * config.o_lora_rank, config.hidden_size, bias=False)
+        self.sinks = nn.Parameter(torch.empty(self.num_heads))
+        self.compressor = (
+            _COMPRESSOR_BY_LAYER_TYPE[self.layer_type](config) if self.layer_type != "agnes_local_attention" else None
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: dict[str, tuple[torch.Tensor, torch.Tensor]] | tuple[torch.Tensor, torch.Tensor],
+        position_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        lead = hidden_states.shape[:-1]
+        head_shape = (*lead, -1, self.head_dim)
+        # the model hands down a {"main", "compress"} cos/sin dict; this layer
+        # takes whichever entry matches its rope label (local -> main, else compress).
+        cos, sin = position_embeddings[self.rope_layer_type]
+
+        # low-rank query path: down-proj + norm, then up-proj into heads, a second
+        # (unweighted) norm, and finally rotary.
+        q_lora = self.q_a_norm(self.q_a_proj(hidden_states))
+        q = self.q_b_norm(self.q_b_proj(q_lora).view(*head_shape).transpose(1, 2))
+        q = apply_agnes_rope(q, cos, sin)
+
+        kv = self.kv_norm(self.kv_proj(hidden_states)).view(*head_shape).transpose(1, 2)
+        kv = apply_agnes_rope(kv, cos, sin)
+
+        if past_key_values is not None:  # shared-KV sliding window (K is V)
+            kv = past_key_values.update(kv, kv, self.layer_idx)[0]
+
+        block_bias = None
+        if self.compressor is not None:  # long-range CSA / HCA entries
+            compressed_kv, block_bias = self.compressor(
+                hidden_states, q_lora, position_ids, past_key_values, self.layer_idx
+            )
+            kv = torch.cat([kv, compressed_kv], dim=2)
+
+        # The compressor path concatenates extra entries onto the KV axis after the
+        # standard sliding-window cache update, so a tensor `attention_mask` (built
+        # for the pre-concat KV length) needs to be extended to cover them. The
+        # compressor returns a `block_bias` carrying per-query causality + indexer
+        # validity over those new slots — cat it in instead of zero-padding (which
+        # would let every query see every compressed slot).
+        if isinstance(attention_mask, torch.Tensor) and kv.shape[2] > attention_mask.shape[-1]:
+            if block_bias is not None:
+                attention_mask = torch.cat([attention_mask, block_bias.to(attention_mask.dtype)], dim=-1)
+            else:
+                attention_mask = F.pad(attention_mask, (0, kv.shape[2] - attention_mask.shape[-1]), value=0.0)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, _eager_attention
+        )
+        attn_output, attn_weights = attention_interface(
+            self,
+            q,
+            kv,
+            kv,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            s_aux=self.sinks,
+            **kwargs,
+        )
+
+        # Because V shares K's storage it also carries K's rotary; undo it on the
+        # output's rope slice with the conjugate rotation (`-sin`) at the query
+        # position, so what reaches the output projection is rope-free. The
+        # transpose pair only re-lays [B, H, S, D] as the [B, S, H, D] that
+        # apply_agnes_rope expects.
+        attn_output = apply_agnes_rope(attn_output.transpose(1, 2), cos, -sin).transpose(1, 2)
+
+        # grouped low-rank output projection: split heads into o_groups, project
+        # each group down, concatenate, then mix up to hidden_size.
+        heads = attn_output.reshape(*lead, self.config.o_groups, -1)
+        mixed = self.o_a_proj(heads).flatten(2)
+        return self.o_b_proj(mixed), attn_weights
+
+
+
+# ==========================================================================
+# Decoder block
+# ==========================================================================
+
+class AgnesDecoderLayer(GradientCheckpointingLayer):
+    r"""Agnes decoder block. Unlike a textbook residual block, the residual
+    here is not a single tensor but a stack of `hc_mult` parallel streams kept
+    in shape `[B, S, hc_mult, D]` for the whole block. Two
+    :class:`AgnesHyperConnection` modules (one at the attention site, one at the
+    MLP site) collapse those streams into the sublayer input and re-expand the
+    sublayer output back across them. Their stream-mixing matrix is projected
+    onto the doubly-stochastic manifold by Sinkhorn-Knopp iteration, which keeps
+    the residual transform non-expansive across a deep stack.
+    """
+
+    def __init__(self, config: AgnesConfig, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.self_attn = AgnesAttention(config, layer_idx)
+        self.mlp = AgnesSparseMoeBlock(config, layer_idx)
+        self.input_layernorm = AgnesRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = AgnesRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attn_hc = AgnesHyperConnection(config)
+        self.ffn_hc = AgnesHyperConnection(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        input_ids: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        # The stream stack stays [B, S, hc_mult, hidden] throughout. Each site
+        # runs the same recipe: collapse streams -> sublayer -> place the output
+        # back across streams (`post`) and carry the old streams through the
+        # doubly-stochastic mixer (`comb`). `post`/`comb` arrive in fp32 (Sinkhorn
+        # runs in float) and are cast down before mixing. `comb` is used
+        # transposed — sum over its FIRST axis, i.e. comb.T @ streams — because
+        # the Sinkhorn matrix is doubly-stochastic but not symmetric.
+        dtype = hidden_states.dtype
+
+        post, comb, collapsed = self.attn_hc(hidden_states)
+        attn_out, _ = self.self_attn(self.input_layernorm(collapsed), **kwargs)
+        placed = post.to(dtype).unsqueeze(-1) * attn_out.unsqueeze(-2)
+        carried = torch.matmul(comb.to(dtype).transpose(-1, -2), hidden_states)
+        hidden_states = placed + carried
+
+        post, comb, collapsed = self.ffn_hc(hidden_states)
+        mlp_out = self.mlp(self.post_attention_layernorm(collapsed), input_ids=input_ids)
+        placed = post.to(dtype).unsqueeze(-1) * mlp_out.unsqueeze(-2)
+        carried = torch.matmul(comb.to(dtype).transpose(-1, -2), hidden_states)
+        return placed + carried
+
+
+
+# ==========================================================================
+# Auxiliary load-balancing loss
+# ==========================================================================
+
+def load_balancing_loss_func(
+    gate_logits: torch.Tensor | tuple[torch.Tensor] | None,
+    num_experts: int | None = None,
+    top_k=2,
+    attention_mask: torch.Tensor | None = None,
+) -> torch.Tensor | int:
+    r"""
+    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
+
+    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
+    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
+    experts is too unbalanced.
+
+    Args:
+        gate_logits:
+            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
+            shape [batch_size X sequence_length, num_experts].
+        num_experts:
+            Number of experts
+        top_k:
+            The number of experts to route per-token, can be also interpreted as the `top-k` routing
+            parameter.
+        attention_mask (`torch.Tensor`, *optional*):
+            The attention_mask used in forward function
+            shape [batch_size X sequence_length] if not None.
+
+    Returns:
+        The auxiliary loss.
+    """
+    if gate_logits is None or not isinstance(gate_logits, tuple):
+        return 0
+
+    if isinstance(gate_logits, tuple):
+        compute_device = gate_logits[0].device
+        concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
+
+    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
+
+    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
+
+    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
+
+    if attention_mask is None:
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.mean(routing_weights, dim=0)
+    else:
+        batch_size, sequence_length = attention_mask.shape
+        num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
+        expert_attention_mask = (
+            attention_mask[None, :, :, None, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
+            .reshape(-1, top_k, num_experts)
+            .to(compute_device)
+        )
+
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
+            expert_attention_mask, dim=0
+        )
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
+        router_per_expert_attention_mask = (
+            attention_mask[None, :, :, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
+            .reshape(-1, num_experts)
+            .to(compute_device)
+        )
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
+            router_per_expert_attention_mask, dim=0
+        )
+
+    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
+    return overall_loss * num_experts
+
+
+
+# ==========================================================================
+# Pretrained base and full models
+# ==========================================================================
+
+@auto_docstring
+class AgnesPreTrainedModel(PreTrainedModel):
+    config: AgnesConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["AgnesDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    # Agnes runs the eager attention path only; the optimized backends are all
+    # unavailable for a concrete structural reason:
+    #
+    #   * FlashAttention 2 / 3 / 4 cap the supported head dim at 256, but Agnes
+    #     uses `head_dim=512`, so those kernels raise
+    #     `FlashAttention forward only supports head dimension at most 256`.
+    #   * SDPA: the fused torch SDPA kernel has no slot for the per-head learned
+    #     attention-sink logit, which is part of Agnes's softmax.
+    #   * FlexAttention: the compressor concatenates a *variable* number of
+    #     entries onto the KV axis inside the block, after the model-level mask
+    #     was already built, so the KV length no longer matches the BlockMask's
+    #     `kv_len`. BlockMask can't be resized at runtime, and the compressor
+    #     already carries its own causal bookkeeping, so wiring that into a
+    #     `mask_mod` is not worthwhile.
+    _supports_flash_attn = False
+    _supports_sdpa = False
+    _supports_flex_attn = False
+    # The compressor's rolling-window buffer / compressed-entries / overlap state
+    # lives on the per-layer cache (:class:`AgnesHCACache` /
+    # :class:`AgnesCSACache`) and isn't compatible with :class:`StaticCache`
+    # — that path would hand the compressor a :class:`StaticSlidingWindowLayer`
+    # with no `store_compression_weights` method. Disabling fullgraph compile
+    # keeps generation tests on the dynamic cache build that does dispatch to
+    # Agnes's own cache layers.
+    _can_compile_fullgraph = False
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "router_logits": OutputRecorder(AgnesTopKRouter, index=0),
+        "hidden_states": AgnesDecoderLayer,
+        "attentions": AgnesAttention,
+    }
+    config_class = AgnesConfig
+    _keep_in_fp32_modules_strict = [
+        "attn_hc",
+        "ffn_hc",
+        "hc_head",
+        "sinks",
+        "position_bias",
+        "e_score_correction_bias",
+        "q_a_norm",
+        "kv_norm",
+        "input_layernorm",
+        "post_attention_layernorm",
+        "norm",
+    ]
+    # Agnes-Flash checkpoints mix FP8 and BF16 in the attention compressor /
+    # indexer branch: these projections ship in BF16 with no companion `scale_inv`.
+    # Listed here (non-strict) so the FP8 quantizer's `get_modules_to_not_convert`
+    # auto-skips them; non-strict has no dtype effect at BF16, so they stay BF16.
+    _keep_in_fp32_modules = [
+        "self_attn.compressor.kv_proj",
+        "self_attn.compressor.gate_proj",
+        "self_attn.compressor.indexer.kv_proj",
+        "self_attn.compressor.indexer.gate_proj",
+        "self_attn.compressor.indexer.scorer.weights_proj",
+    ]
+    _keys_to_ignore_on_load_unexpected = [r"(^|\.)mtp\..*"]
+    # ``_is_stateful`` opts out of generation modes that need to roll the cache
+    # back across drafts (assisted generation, prompt lookup, contrastive search).
+    # The compressor's running-window state isn't rewindable, so `generate`
+    # raises a clear error early instead of failing deep in the compressor with
+    # a missing-method `AttributeError`.
+    _is_stateful = True
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        std = self.config.initializer_range
+        if isinstance(module, (AgnesTopKRouter, AgnesHashRouter)):
+            init.normal_(module.weight, mean=0.0, std=std)
+            if isinstance(module, AgnesTopKRouter):
+                init.zeros_(module.e_score_correction_bias)  # buffer
+            if isinstance(module, AgnesHashRouter):
+                init.zeros_(module.tid2eid)  # buffer; real values come from the checkpoint
+        elif isinstance(module, AgnesExperts):
+            init.normal_(module.gate_up_proj, mean=0.0, std=std)
+            init.normal_(module.down_proj, mean=0.0, std=std)
+        elif isinstance(module, AgnesAttention):
+            init.zeros_(module.sinks)
+        elif isinstance(module, AgnesHyperConnection):
+            init.normal_(module.fn, mean=0.0, std=std)
+            init.zeros_(module.base)
+            init.ones_(module.scale)
+        elif isinstance(module, AgnesHyperHead):
+            init.normal_(module.hc_fn, mean=0.0, std=std)
+            init.zeros_(module.hc_base)
+            init.ones_(module.hc_scale)
+        elif isinstance(module, (AgnesHCACompressor, AgnesCSACompressor, AgnesIndexer)):
+            init.zeros_(module.position_bias)
+        elif isinstance(module, AgnesRotaryEmbedding):
+            for layer_type in module.layer_types:
+                rope_init_fn = module.compute_default_rope_parameters
+                if module.rope_type[layer_type] != "default":
+                    rope_init_fn = ROPE_INIT_FUNCTIONS[module.rope_type[layer_type]]
+                curr_inv_freq, _ = rope_init_fn(module.config, layer_type=layer_type)
+                init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
+                init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)
+
+
+@auto_docstring
+class AgnesModel(AgnesPreTrainedModel):
+    def __init__(self, config: AgnesConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [AgnesDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = AgnesRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = AgnesRotaryEmbedding(config)
+        self.gradient_checkpointing = False
+        self.hc_head = AgnesHyperHead(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+        # what we hand back is the caller's own cache (or None); a cache we spin
+        # up internally is used for this pass only, never returned.
+        return_cache = past_key_values if use_cache else None
+        if past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        if position_ids is None:
+            start = past_key_values.get_seq_length()
+            position_ids = (torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + start).unsqueeze(0)
+
+        # generate() can pass an already-built per-layer-type mask dict; every
+        # Agnes layer shares one sliding-window mask, so reuse it if present.
+        if isinstance(attention_mask, dict):
+            causal_mask = next(iter(attention_mask.values()))
+        else:
+            causal_mask = create_sliding_window_causal_mask(
+                config=self.config,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+                position_ids=position_ids,
+            )
+
+        # fan the embeddings into hc_mult parallel residual streams and precompute
+        # both rope flavours once for the whole stack.
+        hidden_states = inputs_embeds.unsqueeze(2).expand(-1, -1, self.config.hc_mult, -1).contiguous()
+        position_embeddings = {
+            "main": self.rotary_emb(inputs_embeds, position_ids=position_ids, layer_type="main"),
+            "compress": self.rotary_emb(inputs_embeds, position_ids=position_ids, layer_type="compress"),
+        }
+
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                position_embeddings=position_embeddings,
+                position_ids=position_ids,
+                attention_mask=causal_mask,
+                input_ids=input_ids,
+                past_key_values=past_key_values,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(self.hc_head(hidden_states))
+        return MoeModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=return_cache)
+
+
+@auto_docstring
+class AgnesForCausalLM(AgnesPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = AgnesModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.router_aux_loss_coef = config.router_aux_loss_coef
+        self.num_experts = config.num_local_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_router_logits: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeCausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, AutoModelForCausalLM
+
+        >>> model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, trust_remote_code=True)
+        >>> tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+
+        >>> inputs = tokenizer("The capital of France is", return_tensors="pt")
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True)[0]
+        ```"""
+
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs: MoeModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_router_logits=output_router_logits,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+        aux_loss = None
+        if output_router_logits:
+            aux_loss = load_balancing_loss_func(
+                outputs.router_logits,
+                self.num_experts,
+                self.num_experts_per_tok,
+                attention_mask,
+            )
+            if labels is not None:
+                loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )
+
+
+__all__ = ["AgnesPreTrainedModel", "AgnesModel", "AgnesForCausalLM"]
+
+
+
+# ==========================================================================
+# Checkpoint weight conversion
+# ==========================================================================
+
+# Agnes checkpoints ship native-format tensor names (`embed.weight`,
+# `layers.N.attn.*`, `layers.N.ffn.*`, flat `hc_*` hyper-connection params,
+# per-expert `w1`/`w2`/`w3` weights in the native gate/down/up notation). The mapping below
+# renames them onto this implementation's module tree and merges the
+# per-expert weights into the fused `gate_up_proj` / `down_proj` tensors.
+# Registered for `model_type="agnes"`, it is applied automatically by
+# `from_pretrained` and reversed by `save_pretrained`.
+_AGNES_CHECKPOINT_CONVERSION_MAPPING = [
+    WeightRenaming(source_patterns=r"^embed\.weight$", target_patterns="embed_tokens.weight"),
+    WeightRenaming(source_patterns=r"^head\.weight$", target_patterns="lm_head.weight"),
+    WeightRenaming(source_patterns=r"^norm\.weight$", target_patterns="norm.weight"),
+    WeightRenaming(source_patterns=r"^hc_head_fn$", target_patterns="hc_head.hc_fn"),
+    WeightRenaming(source_patterns=r"^hc_head_base$", target_patterns="hc_head.hc_base"),
+    WeightRenaming(source_patterns=r"^hc_head_scale$", target_patterns="hc_head.hc_scale"),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.attn_norm\.",
+        target_patterns=r"layers.\1.input_layernorm.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.ffn_norm\.",
+        target_patterns=r"layers.\1.post_attention_layernorm.",
+    ),
+    WeightRenaming(source_patterns=r"^layers\.(\d+)\.hc_attn_fn$", target_patterns=r"layers.\1.attn_hc.fn"),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.hc_attn_base$", target_patterns=r"layers.\1.attn_hc.base"
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.hc_attn_scale$", target_patterns=r"layers.\1.attn_hc.scale"
+    ),
+    WeightRenaming(source_patterns=r"^layers\.(\d+)\.hc_ffn_fn$", target_patterns=r"layers.\1.ffn_hc.fn"),
+    WeightRenaming(source_patterns=r"^layers\.(\d+)\.hc_ffn_base$", target_patterns=r"layers.\1.ffn_hc.base"),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.hc_ffn_scale$", target_patterns=r"layers.\1.ffn_hc.scale"
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.attn\.",
+        target_patterns=r"layers.\1.self_attn.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.ffn\.",
+        target_patterns=r"layers.\1.mlp.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.attn_sink$",
+        target_patterns=r"layers.\1.self_attn.sinks",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.indexer\.compressor\.norm\.",
+        target_patterns=r"layers.\1.self_attn.compressor.indexer.kv_norm.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.indexer\.compressor\.ape$",
+        target_patterns=r"layers.\1.self_attn.compressor.indexer.position_bias",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.indexer\.compressor\.",
+        target_patterns=r"layers.\1.self_attn.compressor.indexer.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.indexer\.",
+        target_patterns=r"layers.\1.self_attn.compressor.indexer.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.compressor\.indexer\.weights_proj\.",
+        target_patterns=r"layers.\1.self_attn.compressor.indexer.scorer.weights_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.compressor\.norm\.",
+        target_patterns=r"layers.\1.self_attn.compressor.kv_norm.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.compressor\.ape$",
+        target_patterns=r"layers.\1.self_attn.compressor.position_bias",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wq_a\.",
+        target_patterns=r"layers.\1.self_attn.\2.q_a_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wq_b\.",
+        target_patterns=r"layers.\1.self_attn.\2.q_b_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wkv\.",
+        target_patterns=r"layers.\1.self_attn.\2.kv_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wgate\.",
+        target_patterns=r"layers.\1.self_attn.\2.gate_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wo_a\.",
+        target_patterns=r"layers.\1.self_attn.\2.o_a_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.(.*?)\.wo_b\.",
+        target_patterns=r"layers.\1.self_attn.\2.o_b_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.wq_a\.",
+        target_patterns=r"layers.\1.self_attn.q_a_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.wq_b\.",
+        target_patterns=r"layers.\1.self_attn.q_b_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.wkv\.",
+        target_patterns=r"layers.\1.self_attn.kv_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.wo_a\.",
+        target_patterns=r"layers.\1.self_attn.o_a_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.wo_b\.",
+        target_patterns=r"layers.\1.self_attn.o_b_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.self_attn\.q_norm\.",
+        target_patterns=r"layers.\1.self_attn.q_a_norm.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.mlp\.gate\.bias$",
+        target_patterns=r"layers.\1.mlp.gate.e_score_correction_bias",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.mlp\.shared_experts\.w1\.",
+        target_patterns=r"layers.\1.mlp.shared_experts.gate_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.mlp\.shared_experts\.w2\.",
+        target_patterns=r"layers.\1.mlp.shared_experts.down_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.mlp\.shared_experts\.w3\.",
+        target_patterns=r"layers.\1.mlp.shared_experts.up_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.mlp\.parallel_ffn\.w1\.",
+        target_patterns=r"layers.\1.mlp.parallel_ffn.gate_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.mlp\.parallel_ffn\.w2\.",
+        target_patterns=r"layers.\1.mlp.parallel_ffn.down_proj.",
+    ),
+    WeightRenaming(
+        source_patterns=r"^layers\.(\d+)\.mlp\.parallel_ffn\.w3\.",
+        target_patterns=r"layers.\1.mlp.parallel_ffn.up_proj.",
+    ),
+    WeightConverter(
+        source_patterns=[
+            "mlp.experts.*.w1.weight",
+            "mlp.experts.*.w3.weight",
+        ],
+        target_patterns="mlp.experts.gate_up_proj",
+        operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
+    ),
+    WeightConverter(
+        source_patterns="mlp.experts.*.w2.weight",
+        target_patterns="mlp.experts.down_proj",
+        operations=[MergeModulelist(dim=0)],
+    ),
+]
+
+
+try:
+    register_checkpoint_conversion_mapping("agnes", _AGNES_CHECKPOINT_CONVERSION_MAPPING)
+except ValueError:
+    pass  # already registered (module imported more than once)

serve.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+# Agnes PFFN-FP8 serving entrypoint, pinned to lmsysorg/sglang:v0.5.16.
+#
+# Usage:
+#   docker run --gpus all --shm-size 64g -p 30001:30002 \
+#     -v /path/to/this/model:/model \
+#     lmsysorg/sglang:v0.5.16 bash /model/serve.sh
+#
+# Extra sglang flags may be appended after serve.sh, e.g.:
+#   ... bash /model/serve.sh --tp 4
+set -e
+MODEL_DIR="$(cd "$(dirname "$0")" && pwd)"
+SGL_PKG=/sgl-workspace/sglang/python/sglang
+
+# Overlay the Agnes support files onto the stock sglang package
+# (inside this container only; the model directory is never modified).
+cp -r "$MODEL_DIR/sglang_patch/srt" "$MODEL_DIR/sglang_patch/kernels" "$SGL_PKG/"
+
+exec sglang serve --model-path "$MODEL_DIR" --trust-remote-code --tp 8 \
+  --context-length 1048576 --mem-fraction-static 0.90 \
+  --host 0.0.0.0 --port 30002 "$@"

sglang_patch/kernels/ops/layernorm/mhc_head.py

@@ -0,0 +1,151 @@
+"""Fused triton kernel for the DSV4 hc_head LM-head mixer.
+
+Reference torch implementation (deepseek_v4.py AgnesModel.hc_head):
+
+    shape, dtype = x.size(), x.dtype
+    x = x.flatten(1).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=1)
+    return y.to(dtype)
+
+Shapes (DSV4-Pro, hc_mult=4, hidden_size=7168 typical):
+    x      : (T, hc_mult, hidden_size)            bf16
+    hc_fn  : (hc_mult, hc_mult * hidden_size)     fp32
+    scale  : (1,)                                 fp32
+    base   : (hc_mult,)                           fp32
+    out y  : (T, hidden_size)                     bf16
+
+This is a one-shot LM-head op (fires once per forward on the last PP rank), so
+we use a 1-CTA-per-token design that does two passes over x without split-K.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _hc_head_kernel(
+    x_ptr,
+    fn_ptr,
+    scale_ptr,
+    base_ptr,
+    y_ptr,
+    hidden_size: tl.constexpr,
+    HC_MULT: tl.constexpr,
+    K_TOTAL: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    norm_eps: tl.constexpr,
+    hc_eps: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+
+    # ---------- Pass 1: sum_sq over flattened K dim, plus hc_mult inner products ----------
+    sumsq = tl.zeros((), dtype=tl.float32)
+    mix = tl.zeros((HC_MULT,), dtype=tl.float32)
+
+    x_row = x_ptr + pid * K_TOTAL
+    m_idx = tl.arange(0, HC_MULT)
+
+    for k_off in tl.range(0, K_TOTAL, BLOCK_K):
+        k_offs = k_off + tl.arange(0, BLOCK_K)
+        k_mask = k_offs < K_TOTAL
+        x_tile = tl.load(x_row + k_offs, mask=k_mask, other=0.0).to(tl.float32)
+
+        sumsq += tl.sum(x_tile * x_tile, axis=0)
+
+        fn_offs = m_idx[:, None] * K_TOTAL + k_offs[None, :]
+        fn_mask = (m_idx[:, None] < HC_MULT) & k_mask[None, :]
+        fn_tile = tl.load(fn_ptr + fn_offs, mask=fn_mask, other=0.0)
+        mix += tl.sum(fn_tile * x_tile[None, :], axis=1)
+
+    rsqrt = tl.rsqrt(sumsq / K_TOTAL + norm_eps)
+    scale_v = tl.load(scale_ptr).to(tl.float32)
+    base_v = tl.load(base_ptr + m_idx).to(tl.float32)
+
+    # pre[m] = sigmoid(mix[m] * rsqrt * scale + base[m]) + hc_eps
+    pre = tl.sigmoid(mix * rsqrt * scale_v + base_v) + hc_eps
+
+    # ---------- Pass 2: y[d] = sum_m pre[m] * x[m, d]  for d in range(hidden_size) ----------
+    y_row = y_ptr + pid * hidden_size
+
+    for d_off in tl.range(0, hidden_size, BLOCK_D):
+        d_offs = d_off + tl.arange(0, BLOCK_D)
+        d_mask = d_offs < hidden_size
+
+        x_offs = m_idx[:, None] * hidden_size + d_offs[None, :]
+        x_mask = (m_idx[:, None] < HC_MULT) & d_mask[None, :]
+        x_block = tl.load(x_row + x_offs, mask=x_mask, other=0.0).to(tl.float32)
+
+        y_block = tl.sum(pre[:, None] * x_block, axis=0)
+
+        tl.store(y_row + d_offs, y_block.to(y_ptr.dtype.element_ty), mask=d_mask)
+
+
+def fused_hc_head(
+    x: torch.Tensor,
+    hc_fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    norm_eps: float,
+    hc_eps: float,
+) -> torch.Tensor:
+    """Fused (RMSNorm + Linear + Sigmoid-gate + weighted-sum) for the DSV4 hc_head.
+
+    Args:
+        x         : (T, hc_mult, hidden_size) bf16/fp16, must be contiguous
+        hc_fn     : (hc_mult, hc_mult * hidden_size) fp32, contiguous
+        hc_scale  : (1,) fp32 scalar
+        hc_base   : (hc_mult,) fp32
+        norm_eps  : RMS epsilon
+        hc_eps    : additive epsilon after sigmoid
+
+    Returns:
+        y : (T, hidden_size) same dtype as x
+    """
+    assert x.is_contiguous(), "x must be contiguous"
+    assert hc_fn.is_contiguous(), "hc_fn must be contiguous"
+    assert hc_scale.dtype == torch.float32 and hc_base.dtype == torch.float32
+    assert hc_fn.dtype == torch.float32
+    assert x.dim() == 3, f"x must be 3D (T, hc_mult, hidden_size), got {x.shape}"
+
+    T, hc_mult, hidden_size = x.shape
+    assert hc_fn.shape == (hc_mult, hc_mult * hidden_size), (
+        f"hc_fn shape {hc_fn.shape} does not match (hc_mult={hc_mult}, "
+        f"hc_mult*hidden_size={hc_mult * hidden_size})"
+    )
+    assert hc_base.shape == (hc_mult,)
+    assert hc_scale.numel() == 1
+
+    y = torch.empty((T, hidden_size), dtype=x.dtype, device=x.device)
+
+    if T == 0:
+        return y
+
+    BLOCK_K = 512
+    BLOCK_D = 512
+
+    hc_mult_pow2 = max(1, triton.next_power_of_2(hc_mult))
+
+    grid = (T,)
+    _hc_head_kernel[grid](
+        x,
+        hc_fn,
+        hc_scale,
+        hc_base,
+        y,
+        hidden_size=hidden_size,
+        HC_MULT=hc_mult_pow2,
+        K_TOTAL=hc_mult * hidden_size,
+        BLOCK_K=BLOCK_K,
+        BLOCK_D=BLOCK_D,
+        norm_eps=norm_eps,
+        hc_eps=hc_eps,
+        num_warps=4,
+    )
+    return y

sglang_patch/srt/arg_groups/overrides.py

@@ -0,0 +1,2230 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Declarative model-override registry.
+
+Model-identity adjustments to the server configuration are DECLARED here and
+materialized onto ``server_args`` at the end of ``__post_init__`` (gate
+order, last writer wins) — model code never mutates ``ServerArgs`` fields
+imperatively.
+
+Two declaration forms, keyed on ``hf_config.architectures[0]``:
+
+- ``MODEL_OVERRIDES``: pure-constant cases — ``arch -> {field: value}``.
+- ``@register_model_override(arch)``: derived cases — a callable
+  ``fn(server_args, hf_config) -> dict`` that faithfully carries today's
+  conditional logic. ``server_args`` is pristine and must be treated
+  read-only: the callable returns declarations, it never writes.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+
+from sglang.srt.arg_groups.arg_utils import resolvable_fields
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.cuda_graph_config import Backend
+from sglang.srt.utils.common import (
+    cpu_has_amx_support,
+    get_device_capability,
+    get_device_sm,
+    get_nvidia_driver_version,
+    get_quantization_config,
+    is_blackwell_supported,
+    is_cpu,
+    is_cuda,
+    is_flashinfer_available,
+    is_gfx95_supported,
+    is_hip,
+    is_musa,
+    is_npu,
+    is_sm90_supported,
+    is_sm100_supported,
+    is_sm120_supported,
+    is_triton_kernels_available,
+    is_xpu,
+    xpu_has_xmx_support,
+)
+
+logger = logging.getLogger(__name__)
+
+# Constant per-architecture overrides (populated by the migration sweeps).
+MODEL_OVERRIDES: Dict[str, Dict[str, Any]] = {
+    # These models run in bfloat16 regardless of the requested dtype
+    # (faithful port of the legacy unconditional arch branch).
+    "MistralLarge3ForCausalLM": {"dtype": "bfloat16"},
+    "PixtralForConditionalGeneration": {"dtype": "bfloat16"},
+}
+
+# Derived per-architecture override providers, in registration order.
+_MODEL_OVERRIDE_FNS: Dict[str, List[Callable[..., dict]]] = {}
+
+# Predicate-keyed providers, in registration order — for legacy branches
+# matched by substring/predicate on the architecture string rather than an
+# exact name (e.g. '"Step3p5ForCausalLM" in model_arch').
+_PREDICATE_OVERRIDE_FNS: List[Tuple[Callable[[str], bool], Callable[..., dict]]] = []
+
+
+def register_model_override(architecture: str):
+    """Register a derived-override provider for ``architecture``.
+
+    The decorated callable receives ``(server_args, hf_config)``, must not
+    mutate either, and returns a ``{field: resolved_value}`` dict (possibly
+    empty when nothing applies). Providers needing derived model data beyond
+    the HF config go through ``server_args.get_model_config()`` (cached,
+    read-only) — never anything mutating.
+    """
+
+    def decorator(fn: Callable[..., dict]) -> Callable[..., dict]:
+        _MODEL_OVERRIDE_FNS.setdefault(architecture, []).append(fn)
+        return fn
+
+    return decorator
+
+
+def register_model_override_predicate(predicate: Callable[[str], bool]):
+    """Register a derived-override provider keyed by an architecture
+    predicate. Same callable contract as ``register_model_override``."""
+
+    def decorator(fn: Callable[..., dict]) -> Callable[..., dict]:
+        _PREDICATE_OVERRIDE_FNS.append((predicate, fn))
+        return fn
+
+    return decorator
+
+
+def _invoke_provider(
+    fn: Callable[..., dict], server_args: Any, hf_config: Any
+) -> Dict[str, Any]:
+    declared = fn(server_args, hf_config)
+    if not isinstance(declared, dict):
+        raise TypeError(
+            f"model override provider {fn.__qualname__} must return a dict, "
+            f"got {type(declared).__name__}"
+        )
+    return declared
+
+
+class ResolvedView:
+    """Read-only view of the resolving configuration handed to post-process
+    passes: the accumulated declarations overlaid on the pristine
+    ``server_args`` (residual imperative writes of non-resolved fields show
+    through the fallthrough) — exactly the state the legacy handler at the
+    same slot observed. Writes are rejected: passes return declarations.
+    """
+
+    __slots__ = ("_server_args", "_overlay")
+
+    def __init__(self, server_args: Any, overlay: Optional[Dict[str, Any]] = None):
+        object.__setattr__(self, "_server_args", server_args)
+        object.__setattr__(self, "_overlay", overlay or {})
+
+    def __getattr__(self, name: str) -> Any:
+        overlay = object.__getattribute__(self, "_overlay")
+        if name in overlay:
+            return overlay[name]
+        return getattr(object.__getattribute__(self, "_server_args"), name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        raise AttributeError(
+            "ResolvedView is read-only; post-process passes return declarations"
+        )
+
+
+# Ordered post-process passes (the normalization stage). List order is the
+# end-state execution order and mirrors today's handler call sequence in
+# __post_init__; during the transition each pass is invoked from its legacy
+# slot via run_post_process_pass, so ordering is preserved byte-for-byte.
+POST_PROCESS_PASSES: List[Callable[..., dict]] = []
+
+
+def register_post_process(fn: Callable[..., dict]) -> Callable[..., dict]:
+    """Register a post-process pass: ``fn(view) -> {field: resolved_value}``.
+
+    The pass reads a :class:`ResolvedView` (post-model-override state) and
+    must not mutate anything; validations may live in a pass (read + raise).
+    """
+    POST_PROCESS_PASSES.append(fn)
+    return fn
+
+
+def _declaration_overlay(server_args: Any) -> Dict[str, Any]:
+    """Accumulated declared values: declarations never mutate
+    ``server_args``, so mid-resolution readers overlay them from the
+    declaration stash (last writer wins, like the gate)."""
+    overlay: Dict[str, Any] = {}
+    for _source, declared in getattr(server_args, "_resolved_overrides", None) or ():
+        overlay.update(declared)
+    return overlay
+
+
+def run_post_process_pass(server_args: Any, fn: Callable[..., dict]) -> None:
+    """Invoke one pass at its legacy handler slot.
+
+    Evaluates the pass on the resolving state (a read-only view with the
+    accumulated declarations overlaid from the stash) and appends its
+    declaration to the stash. During ``__post_init__`` the fields stay
+    untouched — ``materialize_declarations`` applies the whole stash once at
+    the end of resolution; a pass invoked after materialization (a post-init
+    slot) writes through immediately.
+    """
+    declared = fn(ResolvedView(server_args, overlay=_declaration_overlay(server_args)))
+    if not isinstance(declared, dict):
+        raise TypeError(
+            f"post-process pass {fn.__qualname__} must return a dict, "
+            f"got {type(declared).__name__}"
+        )
+    if declared:
+        entry = (fn.__qualname__, dict(declared))
+        stash = getattr(server_args, "_resolved_overrides", None)
+        if stash is None:
+            # Handlers hosting pass slots may be invoked directly on fixtures
+            # that never ran the monolith dispatch (which owns the stash);
+            # create it lazily. Real publishes always pass through the
+            # dispatch first — the dispatch ASSIGNS the stash, so pass slots
+            # must sit at or after it in __post_init__ order.
+            stash = server_args._resolved_overrides = []
+        stash.append(entry)
+        validate_declarations(server_args, [entry])
+        if getattr(server_args, "_declarations_materialized", False):
+            _apply_fields(server_args, declared)
+
+
+def _apply_fields(server_args: Any, fields: Dict[str, Any]) -> None:
+    """Write fields on behalf of the pipeline (bypasses the strict bare-
+    assignment guard that protects post-resolution mutation)."""
+    object.__setattr__(server_args, "_in_override", True)
+    try:
+        for field, value in fields.items():
+            setattr(server_args, field, value)
+    finally:
+        object.__setattr__(server_args, "_in_override", False)
+
+
+def materialize_declarations(server_args: Any) -> None:
+    """Apply the accumulated declarations onto ``server_args`` once, at the
+    end of ``__post_init__`` (gate order: last writer wins). After this the
+    fields carry the resolved configuration — every post-init reader, in any
+    process, reads them directly; ``resolved_view`` remains an internal
+    helper for mid-resolution code only."""
+    for _source, declared in getattr(server_args, "_resolved_overrides", None) or ():
+        for field, value in declared.items():
+            setattr(server_args, field, value)
+    server_args._declarations_materialized = True
+
+
+def resolved_view(server_args: Any) -> ResolvedView:
+    """Read-only view of the resolving configuration for mid-resolution code
+    that is not a pass (``__post_init__`` handlers and hooks). Internal to
+    the resolution pipeline: after ``materialize_declarations`` runs, the
+    fields themselves carry the resolved values — read them directly."""
+    return ResolvedView(server_args, overlay=_declaration_overlay(server_args))
+
+
+def attention_backends_of(cfg: Any) -> tuple:
+    """(prefill, decode) attention backends of a config-shaped object (a
+    ResolvedView mid-resolution, or pristine server_args at dispatch time):
+    split fields fall back to the base backend."""
+    prefill = (
+        cfg.prefill_attention_backend
+        if cfg.prefill_attention_backend
+        else cfg.attention_backend
+    )
+    decode = (
+        cfg.decode_attention_backend
+        if cfg.decode_attention_backend
+        else cfg.attention_backend
+    )
+    return prefill, decode
+
+
+def mamba_extra_buffer_of(cfg: Any) -> bool:
+    """Mid-resolution equivalent of runtime_context.mamba_extra_buffer_enabled:
+    reads the (possibly overlaid) strategy from a config-shaped object."""
+    return cfg.disable_radix_cache is False and cfg.mamba_radix_cache_strategy in (
+        "extra_buffer",
+        "extra_buffer_lazy",
+    )
+
+
+def declare_load_time_override(source: str, declared: Dict[str, Any]) -> None:
+    """Declare a load-time resolved field (model-file config overrides,
+    weight-resolved dtypes) on the published ``server_args``: resolution has
+    already materialized, so the declaration writes through, joining the
+    declaration stash for provenance and republish consistency."""
+    from sglang.srt.runtime_context import get_context
+
+    server_args = get_context().server_args
+    validate_declarations(server_args, [(source, dict(declared))])
+    override = getattr(server_args, "override", None)
+    if override is not None:
+        override(source, **declared)
+    else:
+        # Config-shaped fixtures without the mutation entry point.
+        _apply_fields(server_args, declared)
+
+
+def collect_model_override_declarations(
+    architecture: str, server_args: Any, hf_config: Any
+) -> List[Tuple[str, Dict[str, Any]]]:
+    """Collect ``(source, declaration)`` pairs for one architecture.
+
+    Application order (last writer wins downstream in the gate): the constant
+    ``MODEL_OVERRIDES`` entry first, then exact-keyed callables in
+    registration order, then matching predicate-keyed callables in
+    registration order. Empty declarations are dropped.
+    """
+    declarations: List[Tuple[str, Dict[str, Any]]] = []
+    const = MODEL_OVERRIDES.get(architecture)
+    if const:
+        declarations.append((f"MODEL_OVERRIDES[{architecture!r}]", dict(const)))
+    for fn in _MODEL_OVERRIDE_FNS.get(architecture, ()):
+        declared = _invoke_provider(fn, server_args, hf_config)
+        if declared:
+            declarations.append((fn.__qualname__, dict(declared)))
+    for predicate, fn in _PREDICATE_OVERRIDE_FNS:
+        if predicate(architecture):
+            declared = _invoke_provider(fn, server_args, hf_config)
+            if declared:
+                declarations.append((fn.__qualname__, dict(declared)))
+    return declarations
+
+
+# ---------------------------------------------------------------------------
+# Derived per-family declarations (faithful ports of legacy arch branches).
+# Callables read the PRISTINE server_args, never write; logging is kept
+# verbatim from the legacy branch for operator-visible fidelity.
+# ---------------------------------------------------------------------------
+
+
+def _register_for(*architectures: str):
+    """Register one provider for several architectures (family lists)."""
+
+    def decorator(fn: Callable[..., dict]) -> Callable[..., dict]:
+        for architecture in architectures:
+            register_model_override(architecture)(fn)
+        return fn
+
+    return decorator
+
+
+@_register_for(
+    "DeepseekV3ForCausalLM",
+    "DeepseekV32ForCausalLM",
+    "KimiK25ForConditionalGeneration",
+    "MistralLarge3ForCausalLM",
+    "PixtralForConditionalGeneration",
+    "GlmMoeDsaForCausalLM",
+    "LongcatFlashForCausalLM",
+    "LongcatFlashForCausalLMNextN",
+)
+def _deepseek_family_overrides(server_args: Any, hf_config: Any) -> dict:
+    """Order-safe declarations of the DeepSeek/DSA branch. The CP parallel
+    writes (enable_dp_attention/ep_size/moe_a2a_backend have post-monolith
+    writers), the kv-cache/split-backend defaults, the quant/moe block (read
+    before it by _set_default_dsa_kv_cache_dtype) and the env writes stay in
+    the branch."""
+    from sglang.srt.configs.model_config import is_deepseek_dsa
+
+    overrides: Dict[str, Any] = {}
+    if is_deepseek_dsa(hf_config):  # DeepSeek 3.2/GLM 5
+        # Set attention backend for DeepSeek
+        if server_args.is_attention_backend_not_set():
+            overrides["attention_backend"] = "dsa"
+            logger.info("Use dsa attention backend for DeepSeek with DSA.")
+        if not is_npu() and not is_xpu():  # CUDA or ROCm GPU
+            if server_args.enable_prefill_cp:
+                logger.warning(
+                    "Context parallel feature is still under experiment. It has only been verified on Hopper platform."
+                )
+                overrides["enable_dp_attention"] = True
+                overrides["moe_dense_tp_size"] = 1
+                if server_args.cp_strategy == "zigzag":
+                    overrides["moe_a2a_backend"] = "deepep"
+                    overrides["ep_size"] = server_args.tp_size
+                    logger.warning(
+                        "zigzag DSA CP requires moe_dense_tp_size=1, "
+                        "moe_a2a_backend=deepep, ep_size=tp_size, batch_size=1."
+                    )
+                else:
+                    assert (
+                        server_args.dp_size == 1
+                    ), "interleave DSA CP does not support DP attention."
+                assert (
+                    server_args.tp_size <= 8
+                ), "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues."
+                # Note(kpham-sgl): Keep attn_tp_size == 1 under DSA CP.
+                # DSACPLayerCommunicator does not all-reduce attention-TP
+                # partial o_proj outputs before replicated dense FFNs.
+                attn_cp_size = server_args.tp_size // server_args.dp_size
+                overrides["attn_cp_size"] = attn_cp_size
+                logger.warning(
+                    "Enabled DSA context parallel: "
+                    f"strategy={server_args.cp_strategy}, dp_size={server_args.dp_size}, "
+                    f"moe_dense_tp_size={overrides['moe_dense_tp_size']}, "
+                    f"ep_size={overrides.get('ep_size', server_args.ep_size)}, tp_size={server_args.tp_size}, "
+                    f"attn_cp_size={attn_cp_size}, "
+                    f"kv_cache_dtype={server_args.kv_cache_dtype}, "
+                    f"moe_a2a_backend={overrides.get('moe_a2a_backend', server_args.moe_a2a_backend)}, "
+                    f"cuda_graph_config[prefill].backend=disabled"
+                )
+
+            # Deferred import to avoid a circular import at module-load
+            # time (dsa.utils imports the runtime-context accessors).
+            from sglang.srt.layers.attention.dsa.utils import (
+                aiter_can_use_preshuffle_paged_mqa,
+            )
+
+            if is_hip() and not aiter_can_use_preshuffle_paged_mqa():
+                # Legacy ROCm DSA path: aiter's gluon paged-MQA kernel is
+                # unavailable (Triton<3.5 and AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS
+                # not set, or SGLANG_DSA_HIP_DISABLE_PRESHUFFLE=1 / SGLANG_USE_AITER=0).
+                overrides["page_size"] = 1
+                logger.warning(
+                    "Setting page size to 1 for DeepSeek DSA on ROCm "
+                    "(aiter preshuffle paged-MQA path unavailable: "
+                    "needs Triton>=3.5.0 or AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1)."
+                )
+            else:
+                overrides["page_size"] = 64
+                logger.warning("Setting page size to 64 for DeepSeek DSA.")
+    else:
+        # DeepSeek V3/R1/V3.1
+        if is_sm100_supported():
+            if (
+                server_args.attention_backend is None
+                and server_args.prefill_attention_backend is None
+                and server_args.decode_attention_backend is None
+            ):
+                overrides["attention_backend"] = "trtllm_mla"
+                logger.info(
+                    "Use trtllm_mla as attention backend on sm100 for DeepseekV3ForCausalLM"
+                )
+        # MLA prefill CP auto-config. Mirrors the NSA CP block above
+        # (minus the in-seq/round-robin mode split, which MLA CP does not support)
+        if server_args.enable_prefill_cp and server_args.use_mla_backend():
+            logger.warning(
+                "MLA prefill context parallel is still experimental. "
+                "Verified on Hopper with the fa3 backend."
+            )
+            overrides["enable_dp_attention"] = True
+            # TODO(kpham-sgl) Supports moe_dense_tp_size != 1.
+            overrides["moe_dense_tp_size"] = 1
+            overrides["moe_a2a_backend"] = "deepep"
+            overrides["ep_size"] = server_args.tp_size
+            logger.warning(
+                "For MLA CP, we have the following restrictions: moe_dense_tp_size == 1, moe_a2a_backend == deepep, ep_size == tp_size, batch_size == 1"
+            )
+            # FIXME(kpham-sgl): Keep attn_tp_size == 1 under MLA CP.
+            # DSACPLayerCommunicator does not all-reduce attention-TP
+            # partial o_proj outputs before replicated dense FFNs.
+            attn_cp_size = server_args.tp_size // server_args.dp_size
+            overrides["attn_cp_size"] = attn_cp_size
+            logger.warning(
+                f"Enable Context Parallel opt for MLA, "
+                f"Setting dp_size == {server_args.dp_size} and "
+                f"attn_cp_size == {attn_cp_size}, "
+                f"moe_dense_tp_size == {overrides['moe_dense_tp_size']}, "
+                f"ep_size == {overrides['ep_size']}, "
+                f"tp_size == {server_args.tp_size}, "
+                f"moe_a2a_backend {overrides['moe_a2a_backend']}, "
+                f"cuda_graph_config[prefill].backend=disabled"
+            )
+    return overrides
+
+
+# Keep in sync with MIMO_V2_MODEL_ARCHS (server_args.py / configs/hf_config.py).
+@_register_for("MiMoV2ForCausalLM", "MiMoV2FlashForCausalLM")
+def _mimo_v2_overrides(server_args: Any, hf_config: Any) -> dict:
+    if server_args.speculative_algorithm == "EAGLE":
+        logger.info("Enable multi-layer EAGLE speculative decoding for MiMoV2 model.")
+        return {"enable_multi_layer_eagle": True}
+    return {}
+
+
+@_register_for("MiniMaxM2ForCausalLM")
+def _minimax_m2_overrides(server_args: Any, hf_config: Any) -> dict:
+    logger.info(
+        "Enable TF32 matmul for MiniMaxM2ForCausalLM model to improve gate gemm performance."
+    )
+    return {"enable_tf32_matmul": True}
+
+
+@_register_for("MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration")
+def _minimax_m3_overrides(server_args: Any, hf_config: Any) -> dict:
+
+    overrides: Dict[str, Any] = {}
+
+    quant_method = get_quantization_config(hf_config)
+    quant_resolved = server_args.quantization
+    if (
+        quant_resolved is None
+        and not server_args._quantization_explicitly_unset
+        and quant_method is not None
+    ):
+        overrides["quantization"] = quant_method
+        quant_resolved = quant_method
+
+    if is_hip():
+        if server_args.is_attention_backend_not_set():
+            overrides["attention_backend"] = "triton"
+        if server_args.moe_runner_backend == "auto" and quant_resolved == "mxfp8":
+            overrides["moe_runner_backend"] = "triton"
+        if not envs.USE_ROCM_AITER_ROPE_BACKEND.is_set():
+            envs.USE_ROCM_AITER_ROPE_BACKEND.set("0")
+        aiter_fusion_resolved = server_args.enable_aiter_allreduce_fusion
+        if (
+            server_args.ep_size > 1
+            and server_args.moe_a2a_backend == "none"
+            and aiter_fusion_resolved
+        ):
+            logger.warning(
+                "Disable --enable-aiter-allreduce-fusion for MiniMax-M3 "
+                "standard EP on ROCm because the deferred fused all-reduce "
+                "corrupts sparse MoE partial outputs."
+            )
+            overrides["enable_aiter_allreduce_fusion"] = False
+            aiter_fusion_resolved = False
+        if not aiter_fusion_resolved:
+            overrides["disable_custom_all_reduce"] = True
+    elif is_sm100_supported():
+        if server_args.is_attention_backend_not_set():
+            overrides["attention_backend"] = "fa4"
+        page_resolved = server_args.page_size
+        if (
+            page_resolved is None
+            and overrides.get("attention_backend", server_args.attention_backend)
+            == "fa4"
+        ):
+            overrides["page_size"] = 128
+            page_resolved = 128
+        if server_args.moe_runner_backend == "auto" and quant_resolved == "mxfp8":
+            overrides["moe_runner_backend"] = "deep_gemm"
+        logger.info(
+            "MiniMax-M3 on SM100: attention_backend="
+            f"{overrides.get('attention_backend', server_args.attention_backend)}, page_size={page_resolved}, "
+            f"moe_runner_backend={overrides.get('moe_runner_backend', server_args.moe_runner_backend)}."
+        )
+    elif is_sm90_supported():
+        if server_args.is_attention_backend_not_set():
+            overrides["attention_backend"] = "fa3"
+        page_resolved = server_args.page_size
+        if (
+            page_resolved is None
+            and overrides.get("attention_backend", server_args.attention_backend)
+            == "fa3"
+        ):
+            overrides["page_size"] = 128
+            page_resolved = 128
+        logger.info(
+            "MiniMax-M3 on Hopper: attention_backend="
+            f"{overrides.get('attention_backend', server_args.attention_backend)}, page_size={page_resolved} "
+            "(MSA is SM100-only; sparse attention runs on the Triton path)."
+        )
+
+    moe_runner_resolved = overrides.get(
+        "moe_runner_backend", server_args.moe_runner_backend
+    )
+    if quant_resolved is None and moe_runner_resolved in ("auto", "deep_gemm"):
+        if moe_runner_resolved == "deep_gemm":
+            logger.warning(
+                "MiniMax-M3: the deep_gemm MoE runner produces corrupted output "
+                "on bf16 full weights; overriding --moe-runner-backend to 'triton'."
+            )
+        overrides["moe_runner_backend"] = "triton"
+
+    return overrides
+
+
+@_register_for(
+    "Gemma2ForCausalLM",
+    "Gemma3ForCausalLM",
+    "Gemma3ForConditionalGeneration",
+    "Gemma3nForCausalLM",
+    "Gemma3nForConditionalGeneration",
+)
+def _gemma2_gemma3_overrides(server_args: Any, hf_config: Any) -> dict:
+    # FIXME: https://github.com/sgl-project/sglang/pull/7367 is not compatible with gemma2 model.
+    # It failed at this test: https://github.com/sgl-project/sglang/actions/runs/16255155597/job/45890331952#step:4:736
+    logger.warning(
+        f"Disable hybrid SWA memory for {hf_config.architectures[0]} as it is not yet supported."
+    )
+    return {"disable_hybrid_swa_memory": True}
+
+
+@_register_for("Exaone4ForCausalLM", "ExaoneMoEForCausalLM")
+def _exaone_overrides(server_args: Any, hf_config: Any) -> dict:
+    if hf_config.sliding_window_pattern is not None:
+        logger.warning(
+            f"Disabling hybrid SWA memory for {hf_config.architectures[0]} as it is not yet supported."
+        )
+        return {"disable_hybrid_swa_memory": True}
+    return {}
+
+
+@_register_for("GptOssForCausalLM")
+def _gpt_oss_overrides(server_args: Any, hf_config: Any) -> dict:
+    overrides: Dict[str, Any] = {}
+    # Set attention backend for GPT-OSS
+    if server_args.is_attention_backend_not_set():
+        if is_sm100_supported():
+            overrides["attention_backend"] = "trtllm_mha"
+        elif is_sm90_supported():
+            overrides["attention_backend"] = "fa3"
+        elif is_cpu() and cpu_has_amx_support():
+            overrides["attention_backend"] = "intel_amx"
+        elif is_xpu():
+            overrides["attention_backend"] = "intel_xpu"
+        elif is_hip():
+            overrides["attention_backend"] = "aiter"
+        else:
+            overrides["attention_backend"] = "triton"
+    if is_xpu():
+        # Check for bf16 dtype on Intel XPU. Reads the pristine dtype request,
+        # which equals the legacy mid-branch read: dtype had no earlier writer
+        # for this arch.
+        if server_args.dtype == "auto":
+            logger.warning(
+                "GptOssForCausalLM on Intel XPU currently supports bfloat16 dtype only"
+            )
+        elif server_args.dtype not in ["bfloat16"]:
+            raise NotImplementedError(
+                f"GptOssForCausalLM on Intel XPU only supports bfloat16 dtype, "
+                f"but got '{server_args.dtype}'. Please use --dtype bfloat16 or remove --dtype to use auto."
+            )
+    quantization_config = getattr(hf_config, "quantization_config", None)
+    is_mxfp4_quant_format = (
+        quantization_config is not None
+        and quantization_config.get("quant_method") == "mxfp4"
+    )
+    if is_mxfp4_quant_format:
+        # use bf16 for mxfp4 triton kernels
+        overrides["dtype"] = "bfloat16"
+    if server_args.moe_runner_backend == "auto":
+
+        if is_sm100_supported() and is_mxfp4_quant_format:
+            overrides["moe_runner_backend"] = "flashinfer_mxfp4"
+            logger.warning(
+                "Detected SM100 and MXFP4 quantization format for GPT-OSS model, enabling FlashInfer MXFP4 MOE kernel."
+            )
+        elif is_sm120_supported() and is_mxfp4_quant_format:
+            # trtllm-gen only supports SM100
+            overrides["moe_runner_backend"] = "marlin"
+            logger.warning(
+                "Detected SM120 and MXFP4 quantization format for GPT-OSS model, enabling Marlin MOE kernel."
+            )
+        elif (is_hip() and envs.SGLANG_USE_AITER.get()) and is_mxfp4_quant_format:
+            overrides["moe_runner_backend"] = "auto"
+            logger.warning(
+                "Detected ROCm and MXFP4 quantization format for GPT-OSS model, enabling aiter MXFP4 MOE kernel."
+            )
+            ## The AITER MXFP4 fused-MoE path for GPT-OSS expects the
+            ## SEPARATED gate/up tile layout (matches the
+            ## `gptoss_fp4_tuned_fmoe.csv` flydsl entries and the
+            ## Mxfp4MoEMethod weight shuffle). Other AITER MXFP4
+            ## callers default to INTERLEAVE; opt this path out
+            ## unless the user explicitly overrode it.
+            # envs.SGLANG_USE_AITER_MOE_GU_ITLV.set(False)
+        elif is_hip() and envs.SGLANG_USE_AITER.get():
+            # For GPT-OSS bf16 on ROCm with aiter, use triton backend
+            # because aiter CK kernel doesn't support all GEMM dimensions
+            overrides["moe_runner_backend"] = "triton"
+            logger.warning(
+                "Detected ROCm with SGLANG_USE_AITER for GPT-OSS bf16 model, using triton MOE kernel."
+            )
+        elif is_musa() and envs.SGLANG_DEEPEP_BF16_DISPATCH.get():
+            overrides["moe_runner_backend"] = "deep_gemm"
+            logger.warning(
+                "Detected MUSA with SGLANG_DEEPEP_BF16_DISPATCH for bf16 model, using deep_gemm kernel."
+            )
+        elif (
+            server_args.ep_size == 1
+            and is_triton_kernels_available()
+            and server_args.quantization is None
+            and not (is_cpu() and cpu_has_amx_support())
+        ):
+            # The triton_kernels package segfaults on Blackwell (B200)
+            # with NVIDIA driver >= 595. Fall back to triton backend.
+            if is_blackwell_supported() and get_nvidia_driver_version() >= (595,):
+                overrides["moe_runner_backend"] = "triton"
+                logger.warning(
+                    "Detected GPT-OSS model on Blackwell with driver >= 595, "
+                    "using triton MOE kernel to avoid triton_kernels SIGSEGV."
+                )
+            else:
+                overrides["moe_runner_backend"] = "triton_kernel"
+                logger.warning(
+                    "Detected GPT-OSS model, enabling triton_kernels MOE kernel."
+                )
+    return overrides
+
+
+# Keep in sync with LLAMA4_MODEL_ARCHS (server_args.py).
+@_register_for("Llama4ForConditionalGeneration", "Llama4ForCausalLM")
+def _llama4_overrides(server_args: Any, hf_config: Any) -> dict:
+    if server_args.device == "cpu":
+        return {}
+    overrides: Dict[str, Any] = {}
+    # Auto-select attention backend for Llama4 if not specified
+    if server_args.attention_backend is None:
+        if is_sm100_supported():
+            backend, platform = "trtllm_mha", "sm100"
+        elif is_sm90_supported():
+            backend, platform = "fa3", "sm90"
+        elif is_hip():
+            backend, platform = "aiter", "hip"
+        elif server_args.device == "xpu":
+            backend, platform = "intel_xpu", "xpu"
+        else:
+            backend, platform = "triton", "other platforms"
+        logger.warning(
+            f"Use {backend} as attention backend on {platform} for Llama4 model"
+        )
+        overrides["attention_backend"] = backend
+    if is_sm100_supported() and server_args.moe_runner_backend == "auto":
+        if server_args.quantization in {"fp8", "modelopt_fp8"}:
+            overrides["moe_runner_backend"] = "flashinfer_trtllm"
+            logger.info(
+                "Use flashinfer_trtllm as MoE runner backend on SM100 for Llama4"
+            )
+    return overrides
+
+
+@_register_for(
+    "Gemma4ForConditionalGeneration",
+    "Gemma4ForCausalLM",
+    "Gemma4UnifiedForConditionalGeneration",
+)
+def _gemma4_overrides(server_args: Any, hf_config: Any) -> dict:
+    overrides: Dict[str, Any] = {}
+    default_attention_backend = "trtllm_mha" if is_sm100_supported() else "triton"
+    if server_args.is_attention_backend_not_set():
+        logger.info(
+            f"Use {default_attention_backend} as default attention backend for Gemma4"
+        )
+        overrides["attention_backend"] = default_attention_backend
+    # If only one split backend is set, keep the other side on a
+    # Gemma4-compatible fallback instead of letting generic backend selection
+    # choose an unsupported backend later.
+    elif server_args.attention_backend is None:
+        overrides["attention_backend"] = default_attention_backend
+    if is_sm100_supported() and server_args.moe_runner_backend == "auto":
+        if server_args.get_model_config().quantization == "modelopt_fp4":
+            overrides["quantization"] = "modelopt_fp4"
+            overrides["moe_runner_backend"] = "flashinfer_trtllm"
+            logger.info(
+                "Use flashinfer_trtllm as MoE runner backend on "
+                "SM100 for Gemma-4 (modelopt_fp4)"
+            )
+    return overrides
+
+
+@_register_for("MossVLForConditionalGeneration")
+def _moss_vl_overrides(server_args: Any, hf_config: Any) -> dict:
+    overrides: Dict[str, Any] = {}
+    if server_args.is_attention_backend_not_set():
+        overrides["prefill_attention_backend"] = "flashinfer"
+        logger.info("Use flashinfer as default prefill attention backend for Moss-VL")
+    prefill_backend = (
+        overrides.get("prefill_attention_backend")
+        or server_args.get_attention_backends()[0]
+    )
+    assert prefill_backend == "flashinfer", (
+        "MossVLForConditionalGeneration requires flashinfer prefill "
+        "attention backend for cross-attention custom mask support."
+    )
+    return overrides
+
+
+@_register_for("MiniCPMV4_6ForConditionalGeneration")
+def _minicpm_v4_6_overrides(server_args: Any, hf_config: Any) -> dict:
+    if is_sm100_supported() and server_args.attention_backend is None:
+        return {"attention_backend": "triton"}
+    return {}
+
+
+@_register_for(
+    "FalconH1ForCausalLM", "JetNemotronForCausalLM", "JetVLMForConditionalGeneration"
+)
+def _falcon_h1_jet_overrides(server_args: Any, hf_config: Any) -> dict:
+    if is_sm100_supported() and server_args.attention_backend is None:
+        return {"attention_backend": "triton"}
+    return {}
+
+
+@_register_for("GraniteMoeHybridForCausalLM")
+def _granite_moe_hybrid_overrides(server_args: Any, hf_config: Any) -> dict:
+    has_mamba = any(
+        layer_type == "mamba" for layer_type in getattr(hf_config, "layer_types", [])
+    )
+    if has_mamba and is_sm100_supported() and server_args.attention_backend is None:
+        return {"attention_backend": "flashinfer"}
+    return {}
+
+
+@_register_for("Lfm2ForCausalLM")
+def _lfm2_overrides(server_args: Any, hf_config: Any) -> dict:
+    if is_sm100_supported() and server_args.attention_backend is None:
+        return {"attention_backend": "flashinfer"}
+    return {}
+
+
+@_register_for("AgnesForCausalLM")
+def _deepseek_v4_overrides(server_args: Any, hf_config: Any) -> dict:
+    """DeepSeek V4 attention/page/window/MoE-runner defaults (from
+    arg_groups/deepseek_v4_hook.py). The kv-cache dtype and NPU split-backend
+    writes, the max_running_requests fill and the validations stay in the
+    hook at its legacy slot."""
+    from sglang.srt.server_args import ServerArgs
+
+    model_arch = hf_config.architectures[0]
+    overrides: Dict[str, Any] = {"attention_backend": "dsv4"}
+
+    page_size = 256
+    if server_args.device == "npu":
+        # NPU keeps the device-aware "dsv4" backend (the registry routes it to
+        # the Ascend V4 subclass); only the pool geometry / dtype differ.
+        # set_default_server_args() pins all three backends to "ascend" for
+        # generic NPU models; override that here so V4 stays consistently on
+        # dsv4.
+        page_size = 128
+        overrides["prefill_attention_backend"] = "dsv4"
+        overrides["decode_attention_backend"] = "dsv4"
+    overrides["page_size"] = page_size
+    logger.info(
+        f"Use dsv4 attention backend for {model_arch}, setting page_size to {page_size}."
+    )
+
+    if server_args.swa_full_tokens_ratio == ServerArgs.swa_full_tokens_ratio:
+        overrides["swa_full_tokens_ratio"] = 0.1
+        logger.info(f"Setting swa_full_tokens_ratio to 0.1 for {model_arch}.")
+
+    # nvidia/DeepSeek-V4-Pro-NVFP4 uses flashinfer_trtllm_routed MoE runner backend.
+    if (
+        server_args.moe_runner_backend == "auto"
+        and server_args.get_model_config().nvfp4_moe_meta is not None
+    ):
+        overrides["moe_runner_backend"] = "flashinfer_trtllm_routed"
+        logger.info(
+            "Use flashinfer_trtllm_routed as MoE runner backend for "
+            f"{model_arch} hybrid FP8+NVFP4 checkpoint."
+        )
+    return overrides
+
+
+@_register_for(
+    "InklingForConditionalGeneration",
+    "InklingForConditionalGenerationMTP",
+)
+def _inkling_overrides(server_args: Any, hf_config: Any) -> dict:
+    """Inkling architecture defaults: SWA / mamba KV-pool ratios tuned for the
+    hybrid-SWA layout, the extra-buffer mamba strategy, and the unified radix
+    tree (which Inkling requires — models/inkling.py asserts it). The full-graph
+    prefill default is set separately (inline, before cuda-graph resolution) —
+    see ServerArgs.__post_init__ / _apply_inkling_prefill_cuda_graph_default. The
+    server-arg defaults each yield to an explicit user value (compared against
+    the ServerArgs class default); the prefill declaration is materialized
+    before _parse_cuda_graph_config folds cuda_graph_backend_prefill into
+    prefill.backend, and an explicit --cuda-graph-backend-prefill /
+    --disable-prefill-cuda-graph still wins. The unified-radix env write follows
+    the MiniMax-M3 handler precedent (env is not a resolvable server-arg)."""
+    from sglang.srt.server_args import ServerArgs
+
+    overrides: Dict[str, Any] = {}
+    # NOTE: the full-graph prefill default is NOT set here. cuda-graph config is
+    # resolved in __post_init__ before declarations are materialized, so a
+    # cuda_graph_backend_prefill declared here lands too late (the breakable
+    # default would already have been auto-disabled for this multimodal arch).
+    # It is set inline before _handle_cuda_graph_config instead.
+    if server_args.swa_full_tokens_ratio == ServerArgs.swa_full_tokens_ratio:
+        overrides["swa_full_tokens_ratio"] = 0.1
+    if server_args.mamba_full_memory_ratio == ServerArgs.mamba_full_memory_ratio:
+        overrides["mamba_full_memory_ratio"] = 0.1
+    # Inkling requires the extra-buffer mamba strategy (inkling.py asserts
+    # enable_mamba_extra_buffer()); the generic "auto" resolution does not cover
+    # Inkling, so pin it here. Yields to an explicit --mamba-scheduler-strategy.
+    if server_args.mamba_radix_cache_strategy == ServerArgs.mamba_radix_cache_strategy:
+        overrides["mamba_radix_cache_strategy"] = "extra_buffer"
+    # Inkling attention runs only on the fa4 (Blackwell) or triton backends --
+    # models/inkling_common/attn.py asserts attention_backend in {fa4, triton}.
+    # The generic resolver would otherwise pick trtllm_mha (SM100) / fa3
+    # (Hopper), so a bare launch fails on the first attention forward. Pin a
+    # supported default when the user left every attention-backend flag unset
+    # (mirrors the MiniMax-M3 SM100 fa4-default above); an explicit
+    # --attention-backend / --prefill/decode-attention-backend still wins.
+    if server_args.is_attention_backend_not_set():
+        inkling_attn_backend = "fa4" if is_sm100_supported() else "triton"
+        overrides["attention_backend"] = inkling_attn_backend
+        logger.info(
+            f"Use {inkling_attn_backend} as the attention backend for Inkling "
+            "(requires fa4 or triton)."
+        )
+    envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.set(True)
+    return overrides
+
+
+@_register_for("NemotronHForCausalLM", "NemotronHPuzzleForCausalLM")
+def _nemotron_h_overrides(server_args: Any, hf_config: Any) -> dict:
+    """NemotronH quantization / MoE runner / attention backend defaults
+    (absorbed from the retired arg_groups/nemotron_h_hook.py; the mamba radix
+    cache handling and the triton-backend assert stay in the arch branch)."""
+    model_arch = hf_config.architectures[0]
+    model_config = server_args.get_model_config()
+    overrides: Dict[str, Any] = {}
+
+    is_modelopt = model_config.quantization in [
+        "modelopt",
+        "modelopt_fp8",
+        "modelopt_fp4",
+        "modelopt_mixed",
+    ]
+    quantization = server_args.quantization
+    if is_modelopt:
+        assert model_config.hf_config.mlp_hidden_act == "relu2"
+        if model_config.quantization == "modelopt":
+            quant_algo = model_config.hf_config.quantization_config["quant_algo"]
+            if quant_algo == "MIXED_PRECISION":
+                quantization = "modelopt_mixed"
+            else:
+                quantization = (
+                    "modelopt_fp4" if quant_algo == "NVFP4" else "modelopt_fp8"
+                )
+        else:
+            quantization = model_config.quantization
+        overrides["quantization"] = quantization
+
+    if (is_modelopt or model_config.quantization is None) and (
+        server_args.moe_runner_backend == "auto"
+    ):
+        if is_sm100_supported() and server_args.moe_a2a_backend == "none":
+            overrides["moe_runner_backend"] = "flashinfer_trtllm"
+            logger.info(
+                f"Use flashinfer_trtllm as MoE runner backend on sm100 for {model_arch}"
+            )
+        elif (
+            (
+                model_config.quantization in ("modelopt_fp4", "modelopt_mixed")
+                or quantization == "modelopt_fp4"
+            )
+            and is_cuda()
+            and (8, 0) <= get_device_capability() < (10, 0)
+        ):
+            overrides["moe_runner_backend"] = "marlin"
+            logger.info(
+                "Use marlin as MoE runner backend on SM80-SM90 for "
+                f"{model_arch} {model_config.quantization}"
+            )
+        else:
+            overrides["moe_runner_backend"] = "flashinfer_cutlass"
+
+    if is_sm100_supported() and server_args.attention_backend is None:
+        overrides["attention_backend"] = "flashinfer"
+    return overrides
+
+
+@_register_for(
+    "Qwen3NextForCausalLM",
+    "Qwen3_5MoeForConditionalGeneration",
+    "InternS2PreviewForConditionalGeneration",
+    "Qwen3_5ForConditionalGeneration",
+)
+def _qwen3_5_hybrid_overrides(server_args: Any, hf_config: Any) -> dict:
+    if not is_sm100_supported() or server_args.attention_backend is not None:
+        return {}
+    sm100_default_attn_backend = "triton"
+    # trtllm_mha requires speculative_eagle_topk == 1 and page_size > 1.
+    # _get_default_attn_backend handles the eagle_topk check.
+    # There is only one case where page_size=1 is required,
+    # which is when radix cache is enabled and both extra_buffer
+    # and spec decoding are disabled.
+    default_attn_backend = server_args._get_default_attn_backend(
+        use_mla_backend=server_args.use_mla_backend(),
+        model_config=server_args.get_model_config(),
+    )
+    # The mamba radix-cache pass runs before this dispatch: read the
+    # declared strategy through the view (the legacy branch observed the
+    # already-written field here).
+    if default_attn_backend == "trtllm_mha" and not (
+        not mamba_extra_buffer_of(resolved_view(server_args))
+        and not server_args.disable_radix_cache
+        and server_args.speculative_algorithm is None
+    ):
+        sm100_default_attn_backend = "trtllm_mha"
+    return {
+        "attention_backend": sm100_default_attn_backend,
+        "page_size": 64 if sm100_default_attn_backend == "trtllm_mha" else 1,
+    }
+
+
+@_register_for("Qwen3VLForConditionalGeneration")
+def _qwen3vl_overrides(server_args: Any, hf_config: Any) -> dict:
+
+    if (
+        is_hip()
+        and envs.SGLANG_USE_AITER_UNIFIED_ATTN.get()
+        and server_args.page_size is None
+    ):
+        logger.info(
+            "Setting page_size=16 for aiter unified attention on Qwen3VLForConditionalGeneration."
+        )
+        return {"page_size": 16}
+    return {}
+
+
+@_register_for(
+    "Qwen3MoeForCausalLM",
+    "Qwen3VLMoeForConditionalGeneration",
+    "Qwen3NextForCausalLM",
+    "Qwen3_5MoeForConditionalGeneration",
+    "InternS2PreviewForConditionalGeneration",
+    "Qwen3_5ForConditionalGeneration",
+)
+def _qwen3_moe_family_overrides(server_args: Any, hf_config: Any) -> dict:
+    overrides: Dict[str, Any] = {}
+    if is_sm100_supported():
+        quant_method = get_quantization_config(hf_config)
+        quantization = server_args.quantization
+        if (
+            quantization is None
+            and not server_args._quantization_explicitly_unset
+            and quant_method is not None
+        ):
+            overrides["quantization"] = quant_method
+            quantization = quant_method
+        if (
+            (quantization in ("fp8", "modelopt_fp4") or quantization is None)
+            and server_args.moe_a2a_backend == "none"
+            and server_args.moe_runner_backend == "auto"
+        ):
+            overrides["moe_runner_backend"] = "flashinfer_trtllm"
+            logger.info(
+                "Use flashinfer_trtllm as MoE runner backend on sm100 for "
+                f"{hf_config.architectures[0]}"
+            )
+    return overrides
+
+
+@_register_for("Glm4MoeForCausalLM")
+def _glm4_moe_overrides(server_args: Any, hf_config: Any) -> dict:
+    overrides: Dict[str, Any] = {}
+    if is_sm100_supported():
+        quantization_config = getattr(hf_config, "quantization_config", None)
+        quant_method = (
+            quantization_config.get("quant_method")
+            if quantization_config is not None
+            else None
+        )
+        quantization = server_args.quantization
+        if (
+            quantization is None
+            and not server_args._quantization_explicitly_unset
+            and quant_method is not None
+        ):
+            overrides["quantization"] = quant_method
+            quantization = quant_method
+        if (
+            quantization in {"modelopt_fp4", None}
+            and server_args.moe_a2a_backend == "none"
+            and server_args.moe_runner_backend == "auto"
+        ):
+            overrides["moe_runner_backend"] = "flashinfer_trtllm"
+            logger.info(
+                "Use flashinfer_trtllm as MoE runner backend on sm100 for Glm4MoeForCausalLM"
+            )
+    logger.info(
+        "Enable TF32 matmul for Glm4MoeForCausalLM model to improve gate gemm performance."
+    )
+    overrides["enable_tf32_matmul"] = True
+    return overrides
+
+
+@_register_for("Olmo2ForCausalLM")
+def _olmo2_overrides(server_args: Any, hf_config: Any) -> dict:
+    overrides: Dict[str, Any] = {}
+    # FIXME: https://github.com/sgl-project/sglang/pull/7367 is not compatible with Olmo3 model.
+    logger.warning(
+        f"Disabling hybrid SWA memory for {hf_config.architectures[0]} as it is not yet supported."
+    )
+    overrides["disable_hybrid_swa_memory"] = True
+    if server_args.attention_backend is None:
+        if is_cuda() and is_sm100_supported():
+            overrides["attention_backend"] = "trtllm_mha"
+        elif is_cuda() and get_device_sm() >= 80:
+            overrides["attention_backend"] = "fa3"
+        else:
+            overrides["attention_backend"] = "triton"
+    return overrides
+
+
+@register_model_override_predicate(
+    lambda arch: "Step3p5ForCausalLM" in arch
+    or "Step3p7ForConditionalGeneration" in arch
+)
+def _step3p_overrides(server_args: Any, hf_config: Any) -> dict:
+    overrides: Dict[str, Any] = {}
+    if server_args.is_attention_backend_not_set():
+        if is_blackwell_supported():
+            logger.info("Auto-select fa4 attention backend for Step3p7 on Blackwell.")
+            overrides["attention_backend"] = "fa4"
+        elif is_sm90_supported():
+            logger.info("Auto-select fa3 attention backend for Step3p7 on Hopper.")
+            overrides["attention_backend"] = "fa3"
+    if server_args.speculative_algorithm == "EAGLE":
+        logger.info(
+            "Enable multi-layer EAGLE speculative decoding for Step3p5ForCausalLM model."
+        )
+        overrides["enable_multi_layer_eagle"] = True
+    if server_args.enable_hierarchical_cache:
+        logger.warning(
+            "Reset swa_full_tokens_ratio to 1.0 for Step3p5ForCausalLM model with hierarchical cache"
+        )
+        overrides["swa_full_tokens_ratio"] = 1.0
+        logger.warning(
+            "Disable hybrid SWA memory for Step3p5ForCausalLM model with hierarchical cache"
+        )
+        overrides["disable_hybrid_swa_memory"] = True
+    return overrides
+
+
+# ---------------------------------------------------------------------------
+# Post-process passes (normalization stage), in end-state execution order.
+# Faithful ports of the legacy __post_init__ handlers; each is invoked from
+# its legacy slot via run_post_process_pass during the transition.
+# ---------------------------------------------------------------------------
+
+
+# Architectures whose monolith branch routes through the mamba radix cache
+# handling (hybrid linear-attention models). Keep in sync with the branch
+# guards in _handle_model_specific_adjustments.
+_MAMBA_RADIX_CACHE_ARCHS = frozenset(
+    {
+        "KimiLinearForCausalLM",
+        "BailingMoeV2_5ForCausalLM",
+        "Qwen3NextForCausalLM",
+        "Qwen3_5MoeForConditionalGeneration",
+        "InternS2PreviewForConditionalGeneration",
+        "Qwen3_5ForConditionalGeneration",
+        "MiniCPMV4_6ForConditionalGeneration",
+        "NemotronHForCausalLM",
+        "NemotronHPuzzleForCausalLM",
+        "FalconH1ForCausalLM",
+        "JetNemotronForCausalLM",
+        "JetVLMForConditionalGeneration",
+        "Lfm2ForCausalLM",
+        "ZayaForCausalLM",
+    }
+)
+
+# Architectures that support the extra_buffer mamba radix cache strategy.
+# Single source of truth: ServerArgs._support_mamba_cache_extra_buffer
+# delegates here.
+_MAMBA_EXTRA_BUFFER_ARCHS = frozenset(
+    {
+        "KimiLinearForCausalLM",
+        "Qwen3_5ForConditionalGeneration",
+        "Qwen3_5MoeForConditionalGeneration",
+        "Qwen3NextForCausalLM",
+        "InternS2PreviewForConditionalGeneration",
+        "MiniCPMV4_6ForConditionalGeneration",
+        "BailingMoeV2_5ForCausalLM",
+        "FalconH1ForCausalLM",
+        "GraniteMoeHybridForCausalLM",
+        "NemotronHForCausalLM",
+        "NemotronHPuzzleForCausalLM",
+    }
+)
+
+
+def supports_mamba_cache_extra_buffer(view: Any, model_arch: str) -> bool:
+    """Whether ``model_arch`` supports the extra_buffer strategy on the
+    configured linear-attention backend (pure read)."""
+    if model_arch in _MAMBA_EXTRA_BUFFER_ARCHS:
+        return view.linear_attn_backend == "triton"
+    return False
+
+
+@register_post_process
+def _mamba_radix_cache_resolution(view: Any) -> dict:
+    """Resolve the hybrid-mamba radix cache fields (pure).
+
+    Slot pass: invoked at each legacy ``_handle_mamba_radix_cache`` slot —
+    the hybrid-spec call at the head of the monolith and the per-arch branch
+    calls — where it reads the mid-resolution ``page_size`` /
+    ``disable_overlap_schedule`` exactly as the legacy helper did. The arch
+    guard replicates the union of the legacy call-site guards so the pass is
+    self-sufficient in the end-state pass list.
+    """
+    from sglang.srt.configs.linear_attn_model_registry import (
+        get_linear_attn_spec_by_arch,
+    )
+
+    hf_config = view.get_model_config().hf_config
+    model_arch = hf_config.architectures[0]
+
+    in_branch = model_arch in _MAMBA_RADIX_CACHE_ARCHS
+    if model_arch == "GraniteMoeHybridForCausalLM":
+        in_branch = any(
+            layer_type == "mamba"
+            for layer_type in getattr(hf_config, "layer_types", [])
+        )
+    spec = get_linear_attn_spec_by_arch(model_arch)
+    if not ((spec is not None and spec.uses_mamba_radix_cache) or in_branch):
+        return {}
+
+    if view.disable_radix_cache:
+        return {}
+
+    declared: Dict[str, Any] = {"uses_mamba_radix_cache": True}
+    if view.mamba_radix_cache_strategy == "auto":
+        wants_overlap = not view.disable_overlap_schedule
+        wants_paging = view.page_size is not None and view.page_size > 1
+        if (wants_overlap or wants_paging) and supports_mamba_cache_extra_buffer(
+            view, model_arch
+        ):
+            declared["mamba_radix_cache_strategy"] = "extra_buffer"
+        else:
+            declared["mamba_radix_cache_strategy"] = "no_buffer"
+            declared["disable_overlap_schedule"] = True
+    return declared
+
+
+@register_post_process
+def _dsa_kv_cache_dtype_default(view: Any) -> dict:
+    """Slot pass in the DSA arm, ordered before the split-backend
+    resolution: default the kv-cache dtype from the device capability
+    (Blackwell FP8, Hopper bf16) and normalize the bf16 alias. Reads the
+    PRISTINE dsa split backends (their resolution runs after this pass)."""
+    from sglang.srt.configs.model_config import is_deepseek_dsa
+
+    hf_config = view.get_model_config().hf_config
+    if hf_config.architectures[0] not in _DEEPSEEK_FAMILY_ARCHS:
+        return {}
+    if not is_deepseek_dsa(hf_config):
+        return {}
+    if is_npu() or is_xpu():
+        return {}
+
+    import torch
+
+    major, _ = torch.cuda.get_device_capability()
+
+    # If user specified a backend but didn't explicitly set kv_cache_dtype,
+    # suggest them to be explicit about kv_cache_dtype to avoid surprises
+    if (
+        view.dsa_prefill_backend is not None or view.dsa_decode_backend is not None
+    ) and view.kv_cache_dtype == "auto":
+        logger.warning(
+            "When specifying --dsa-prefill-backend or --dsa-decode-backend, "
+            "you should also explicitly set --kv-cache-dtype (e.g., 'fp8_e4m3' or 'bfloat16'). "
+            "DeepSeek V3.2 defaults to FP8 KV cache which may not be compatible with all backends."
+        )
+
+    kv_cache_dtype = view.kv_cache_dtype
+    if kv_cache_dtype == "auto":
+        kv_cache_dtype = "fp8_e4m3" if major >= 10 else "bfloat16"
+        logger.warning(
+            f"Setting KV cache dtype to {kv_cache_dtype} for DeepSeek DSA on SM{major} device."
+        )
+    if kv_cache_dtype == "bf16":
+        kv_cache_dtype = "bfloat16"
+    assert kv_cache_dtype in [
+        "bfloat16",
+        "fp8_e4m3",
+    ], "DeepSeek DSA only supports bf16/bfloat16 or fp8_e4m3 kv_cache_dtype"
+    if kv_cache_dtype != view.kv_cache_dtype:
+        return {"kv_cache_dtype": kv_cache_dtype}
+    return {}
+
+
+@register_post_process
+def _dsa_split_backend_resolution(view: Any) -> dict:
+    """Slot pass in the DSA arm: default the DSA prefill/decode split
+    backends from the mid-resolution kv-cache dtype and the device
+    capability. The hisparse arm takes precedence under --enable-hisparse."""
+    from sglang.srt.configs.model_config import is_deepseek_dsa
+
+    hf_config = view.get_model_config().hf_config
+    if hf_config.architectures[0] not in _DEEPSEEK_FAMILY_ARCHS:
+        return {}
+    if not is_deepseek_dsa(hf_config):
+        return {}
+    if is_npu() or is_xpu():
+        return {}
+
+    import torch
+
+    major, _ = torch.cuda.get_device_capability()
+    kv_cache_dtype = view.kv_cache_dtype
+    user_set_prefill = view.dsa_prefill_backend is not None
+    user_set_decode = view.dsa_decode_backend is not None
+    declared: Dict[str, Any] = {}
+
+    if view.enable_hisparse:
+        from sglang.srt.arg_groups.hisparse_hook import _hisparse_default_backend
+
+        backend = _hisparse_default_backend(kv_cache_dtype)
+        if not user_set_prefill:
+            declared["dsa_prefill_backend"] = backend
+        if not user_set_decode:
+            declared["dsa_decode_backend"] = backend
+        prefill = declared.get("dsa_prefill_backend", view.dsa_prefill_backend)
+        decode = declared.get("dsa_decode_backend", view.dsa_decode_backend)
+        logger.warning(
+            f"HiSparse enabled ({kv_cache_dtype}): using DSA backends "
+            f"prefill={prefill}, decode={decode}."
+        )
+        return declared
+
+    if not user_set_prefill and not user_set_decode and is_hip():
+        declared["dsa_prefill_backend"] = "tilelang"
+        declared["dsa_decode_backend"] = "tilelang"
+    elif kv_cache_dtype == "fp8_e4m3":
+        # Blackwell FP8 defaults to trtllm; Hopper FP8 to flashmla_kv.
+        default = "trtllm" if major >= 10 else "flashmla_kv"
+        if not user_set_prefill:
+            declared["dsa_prefill_backend"] = default
+        if not user_set_decode:
+            declared["dsa_decode_backend"] = default
+    else:
+        # Set prefill/decode backends based on hardware architecture.
+        if not user_set_prefill:
+            declared["dsa_prefill_backend"] = "flashmla_sparse"
+        if not user_set_decode:
+            declared["dsa_decode_backend"] = "trtllm" if major >= 10 else "fa3"
+
+    prefill = declared.get("dsa_prefill_backend", view.dsa_prefill_backend)
+    decode = declared.get("dsa_decode_backend", view.dsa_decode_backend)
+    logger.warning(
+        f"Set DSA backends for {kv_cache_dtype} KV Cache: "
+        f"prefill={prefill}, decode={decode}."
+    )
+    return declared
+
+
+# Keep in sync with the DeepSeek family list on _deepseek_family_overrides.
+_DEEPSEEK_FAMILY_ARCHS = frozenset(
+    {
+        "DeepseekV3ForCausalLM",
+        "DeepseekV32ForCausalLM",
+        "KimiK25ForConditionalGeneration",
+        "MistralLarge3ForCausalLM",
+        "PixtralForConditionalGeneration",
+        "GlmMoeDsaForCausalLM",
+        "LongcatFlashForCausalLM",
+        "LongcatFlashForCausalLMNextN",
+    }
+)
+
+
+@register_post_process
+def _deepseek_moe_quant_resolution(view: Any) -> dict:
+    """Slot pass invoked from inside the DeepSeek arch branch ("Set moe
+    backend for DeepSeek"), NOT a dispatch-time declaration: the DSA
+    kv-cache-dtype default earlier in the branch must read the PRISTINE
+    quantization, so this resolution has to stay at its legacy slot."""
+    hf_config = view.get_model_config().hf_config
+    model_arch = hf_config.architectures[0]
+    if model_arch not in _DEEPSEEK_FAMILY_ARCHS:
+        return {}
+    overrides: Dict[str, Any] = {}
+    if is_sm100_supported():
+        quant_method = get_quantization_config(hf_config)
+        quant_cfg = getattr(hf_config, "quantization_config", None) or {}
+        config_groups = quant_cfg.get("config_groups", {})
+        group0 = config_groups.get("group_0", {})
+        weights_cfg = group0.get("weights", {})
+        # this also apply to kimi k2.5
+        # since it follow the compressed tensor int4 recipe
+        # but not kimi k2 instruct or 0905 instruct.
+        is_kimi_k2_k25_thinking_int4 = (
+            quant_method == "compressed-tensors"
+            and weights_cfg.get("num_bits") == 4
+            and weights_cfg.get("group_size") == 32
+            and weights_cfg.get("strategy") == "group"
+            and weights_cfg.get("type") == "int"
+        )
+        quantization = view.quantization
+        if quantization is None and not view._quantization_explicitly_unset:
+            # DeepSeek V3/R1 uses native FP8 MoE experts without
+            # declaring it in quantization_config.  However, other
+            # models that share the same architecture class (e.g.
+            # Moonlight-16B-A3B) are purely BF16.  Check the actual
+            # safetensors header instead of assuming FP8 by arch name.
+            if quant_method is None and model_arch in ["DeepseekV3ForCausalLM"]:
+                from sglang.srt.utils.common import has_fp8_weights_in_checkpoint
+
+                if has_fp8_weights_in_checkpoint(view.model_path):
+                    overrides["quantization"] = quantization = "fp8"
+                    logger.info(
+                        "Detected FP8 expert weights in checkpoint, "
+                        "default to fp8 for DeepSeek on sm100"
+                    )
+                else:
+                    logger.info(
+                        "No FP8 expert weights found in checkpoint, "
+                        "keeping bf16 for DeepSeek-arch model on sm100"
+                    )
+            else:
+                overrides["quantization"] = quantization = quant_method
+        if (
+            view.moe_a2a_backend == "none"
+            and view.moe_runner_backend == "auto"
+            and (
+                quantization
+                in ["fp8", "modelopt_fp8", "modelopt_fp4", "modelopt_mixed"]
+                or is_kimi_k2_k25_thinking_int4
+                or quantization is None
+            )
+        ):
+            overrides["moe_runner_backend"] = "flashinfer_trtllm"
+            if is_kimi_k2_k25_thinking_int4:
+                logger.info(
+                    "Use flashinfer_trtllm as MoE runner backend on Blackwell for Kimi K2 / K2.5 thinking int4"
+                )
+            else:
+                logger.info(
+                    "Use flashinfer_trtllm as MoE runner backend on sm100 for DeepseekV3ForCausalLM"
+                )
+        if (
+            model_arch in ["LongcatFlashForCausalLM", "LongcatFlashForCausalLMNextN"]
+            and view.fp8_gemm_runner_backend == "auto"
+            and quantization in ["fp8", "modelopt_fp8"]
+            and quant_cfg.get("scale_fmt", None) != "ue8m0"
+        ):
+            overrides["fp8_gemm_runner_backend"] = "flashinfer_trtllm"
+            logger.info(
+                "Use flashinfer_trtllm as FP8 GEMM backend on Blackwell for LongCat FP8 "
+                "checkpoint with non-ue8m0 scales"
+            )
+    return overrides
+
+
+@register_post_process
+def _deepseek_spec_moe_resolution(view: Any) -> dict:
+    """Slot pass at the DeepSeek branch's HIP arm: draft (nextn) spec-MoE
+    backends for the DeepSeek fp4 checkpoint. Reads the mid-resolution
+    quantization (after _deepseek_moe_quant_resolution) and the pre-a2a
+    ep_size, exactly like the legacy in-branch writes."""
+
+    hf_config = view.get_model_config().hf_config
+    model_arch = hf_config.architectures[0]
+    if model_arch not in _DEEPSEEK_FAMILY_ARCHS:
+        return {}
+    if not is_hip():
+        return {}
+    if not (
+        view.quantization == "modelopt_fp4"
+        and view.speculative_algorithm == "EAGLE"
+        and (
+            view.speculative_moe_runner_backend is None
+            or view.speculative_moe_a2a_backend is None
+        )
+    ):
+        return {}
+    if envs.SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE.get():
+        logger.info(
+            "Use deep_gemm moe runner and deepep a2a backend for bf16 nextn layer in deepseek fp4 checkpoint."
+        )
+        # Validate usage of ep
+        if view.ep_size == 1:
+            raise ValueError(
+                "Invalid configuration: 'deep_gemm' speculative MoE runner backend with "
+                "'deepep' a2a backend requires expert parallelism (ep_size > 1). "
+                f"Current ep_size is {view.ep_size}. "
+                "Please set --ep-size > 1 (e.g., --ep-size 8) to use this configuration, "
+                "or change --speculative-moe-a2a-backend to 'none' if expert parallelism is not available."
+            )
+        return {
+            "speculative_moe_runner_backend": "deep_gemm",
+            "speculative_moe_a2a_backend": "deepep",
+        }
+    logger.info(
+        "Use triton fused moe by default for bf16 nextn layer in deepseek fp4 checkpoint."
+    )
+    return {
+        "speculative_moe_runner_backend": "triton",
+        "speculative_moe_a2a_backend": "none",
+    }
+
+
+@register_post_process
+def _deepseek_v4_kv_cache_dtype(view: Any) -> dict:
+    """Slot pass in the DeepSeek V4 hook: default the kv-cache dtype to FP8
+    (bfloat16 on NPU, where the pool geometry differs) and validate the
+    result. The NPU split-backend writes stay in the hook."""
+    hf_config = view.get_model_config().hf_config
+    model_arch = hf_config.architectures[0]
+    if model_arch != "AgnesForCausalLM":
+        return {}
+
+    kv_cache_dtype = view.kv_cache_dtype
+    if kv_cache_dtype == "auto":
+        kv_cache_dtype = "fp8_e4m3"
+        logger.warning(f"Setting KV cache dtype to {kv_cache_dtype} for {model_arch}.")
+    if view.device == "npu":
+        kv_cache_dtype = "bfloat16"
+    assert kv_cache_dtype in [
+        "fp8_e4m3",
+        "bfloat16",
+    ], f"{kv_cache_dtype} is not supported for {model_arch}"
+    if kv_cache_dtype != view.kv_cache_dtype:
+        return {"kv_cache_dtype": kv_cache_dtype}
+    return {}
+
+
+@register_post_process
+def _deepseek_v4_sm120_moe(view: Any) -> dict:
+    """Default DeepSeek V4 MXFP4 experts to FlashInfer CUTLASS on SM120."""
+    hf_config = view.get_model_config().hf_config
+    if hf_config.architectures[0] != "AgnesForCausalLM":
+        return {}
+    if is_sm120_supported() and view.moe_runner_backend == "auto":
+        logger.info(
+            "Use flashinfer_mxfp4 as MoE runner backend on SM120 for DeepseekV4"
+        )
+        return {"moe_runner_backend": "flashinfer_mxfp4"}
+    return {}
+
+
+@register_post_process
+def _sparse_head_overlap_disable(view: Any) -> dict:
+
+    if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set():
+        logger.warning(
+            "Overlap scheduler is disabled when using sparse head for embedding model."
+        )
+        return {"disable_overlap_schedule": True}
+    return {}
+
+
+# Architectures with explicit FlashInfer AllReduce Fusion support. Keep in
+# sync with the model-side fusion implementations.
+_FLASHINFER_ALLREDUCE_FUSION_ARCHS = frozenset(
+    {
+        "DeepseekV3ForCausalLM",
+        "DeepseekV32ForCausalLM",
+        "GptOssForCausalLM",
+        "GlmMoeDsaForCausalLM",
+        "Glm4MoeForCausalLM",
+        "Glm4MoeLiteForCausalLM",
+        "MistralLarge3ForCausalLM",
+        "Qwen3MoeForCausalLM",
+        "Qwen3VLMoeForConditionalGeneration",
+        "Qwen3NextForCausalLM",
+        "KimiK25ForConditionalGeneration",
+        "Qwen3_5MoeForConditionalGeneration",
+        "InternS2PreviewForConditionalGeneration",
+        "Qwen3_5ForConditionalGeneration",
+        "NemotronHForCausalLM",
+        "NemotronHPuzzleForCausalLM",
+    }
+)
+
+
+@register_post_process
+def _flashinfer_allreduce_fusion_auto_enable(view: Any) -> dict:
+    """Slot pass at the monolith tail: auto-enable FlashInfer AllReduce
+    Fusion on SM90/SM100 for models with explicit support. auto resolves to
+    mnnvl on Blackwell (single- and multi-node) and trtllm on SM90
+    single-node systems. Reads the mid-resolution enable_dp_attention /
+    moe_a2a_backend (after the DeepSeek CP and a2a declarations), exactly
+    like the legacy tail block."""
+    model_arch = view.get_model_config().hf_config.architectures[0]
+    if (
+        view.flashinfer_allreduce_fusion_backend is None
+        and model_arch in _FLASHINFER_ALLREDUCE_FUSION_ARCHS
+        and (is_sm90_supported() or is_sm100_supported())
+        and view.tp_size > 1
+        and not view.enable_dp_attention
+        and (view.nnodes == 1 or is_sm100_supported())
+        and view.moe_a2a_backend == "none"
+    ):
+        logger.info(
+            f"Auto-enabling FlashInfer AllReduce Fusion on SM90/SM10X for {model_arch}"
+        )
+        return {"flashinfer_allreduce_fusion_backend": "auto"}
+    return {}
+
+
+@register_post_process
+def _enforce_disable_allreduce_fusion(view: Any) -> dict:
+    """Slot pass right after the auto-enable: the user's enforce-disable
+    switch wins over every model-specific adjustment."""
+    if view.enforce_disable_flashinfer_allreduce_fusion:
+        logger.info(
+            "FlashInfer allreduce fusion is forcibly disabled "
+            "via --enforce-disable-flashinfer-allreduce-fusion."
+        )
+        return {"flashinfer_allreduce_fusion_backend": None}
+    return {}
+
+
+@register_post_process
+def _sampling_backend_default(view: Any) -> dict:
+    if view.sampling_backend is None:
+        return {
+            "sampling_backend": (
+                "flashinfer" if is_flashinfer_available() else "pytorch"
+            )
+        }
+    return {}
+
+
+@register_post_process
+def _deterministic_sampling_backend(view: Any) -> dict:
+    if view.enable_deterministic_inference and view.sampling_backend != "ascend":
+        logger.warning(
+            "Sampling backend is set to pytorch for deterministic inference."
+        )
+        return {"sampling_backend": "pytorch"}
+    return {}
+
+
+def _deterministic_is_deepseek_model(view: Any) -> bool:
+    """Faithful copy of the deterministic handler's arch probe (pure read;
+    the handler keeps its own copy for the later deepseek validation)."""
+    from sglang.srt.connector import ConnectorType
+    from sglang.srt.utils.common import parse_connector_type
+
+    if parse_connector_type(view.model_path) == ConnectorType.INSTANCE:
+        return False
+    try:
+        hf_config = view.get_model_config().hf_config
+        return hf_config.architectures[0] in [
+            "DeepseekV2ForCausalLM",
+            "DeepseekV3ForCausalLM",
+            "DeepseekV32ForCausalLM",
+            "MistralLarge3ForCausalLM",
+            "PixtralForConditionalGeneration",
+            "GlmMoeDsaForCausalLM",
+        ]
+    except Exception:
+        return False
+
+
+@register_post_process
+def _deterministic_allreduce_fusion_disable(view: Any) -> dict:
+    if (
+        view.enable_deterministic_inference
+        and view.flashinfer_allreduce_fusion_backend is not None
+    ):
+        logger.warning(
+            "Disable --flashinfer-allreduce-fusion-backend because deterministic inference is enabled."
+        )
+        return {"flashinfer_allreduce_fusion_backend": None}
+    return {}
+
+
+@register_post_process
+def _deterministic_attention_backend(view: Any) -> dict:
+    if not view.enable_deterministic_inference:
+        return {}
+    from sglang.srt.server_args import DETERMINISTIC_ATTENTION_BACKEND_CHOICES
+
+    if view.attention_backend is None:
+        # User didn't specify attention backend, fallback based on GPU architecture
+        if is_sm100_supported() or is_sm120_supported():
+            # Blackwell and newer architectures
+            if _deterministic_is_deepseek_model(view):
+                # fallback to triton for DeepSeek models because flashinfer
+                # doesn't support deterministic inference for DeepSeek models yet
+                backend = "triton"
+            else:
+                # fallback to flashinfer on Blackwell for non-DeepSeek models
+                backend = "flashinfer"
+        else:
+            # Hopper (SM90) and older architectures
+            backend = "fa3"
+        logger.warning(
+            f"Attention backend not specified. Falling back to '{backend}' for deterministic inference. "
+            f"You can explicitly set --attention-backend to one of {DETERMINISTIC_ATTENTION_BACKEND_CHOICES}."
+        )
+        return {"attention_backend": backend}
+    elif view.attention_backend not in DETERMINISTIC_ATTENTION_BACKEND_CHOICES:
+        # User explicitly specified an incompatible attention backend
+        raise ValueError(
+            f"Currently only {DETERMINISTIC_ATTENTION_BACKEND_CHOICES} attention backends are supported for deterministic inference, "
+            f"but you explicitly specified '{view.attention_backend}'."
+        )
+    return {}
+
+
+@register_post_process
+def _attention_backend_default(view: Any) -> dict:
+    if view.prefill_attention_backend is not None and (
+        view.prefill_attention_backend == view.decode_attention_backend
+    ):  # override the default attention backend
+        return {"attention_backend": view.prefill_attention_backend}
+    if view.attention_backend is None:
+        backend = view._get_default_attn_backend(
+            view.use_mla_backend(), view.get_model_config()
+        )
+        logger.info(
+            f"Attention backend not specified. Use {backend} backend by default."
+        )
+        return {"attention_backend": backend}
+    return {}
+
+
+@register_post_process
+def _mla_backend_page_constraints(view: Any) -> dict:
+    """Page-size constraints of the MLA/TRTLLM backend family (the raises and
+    the cutedsl prefill fallback stay in the handler; only the page snaps are
+    declared). The snaps chain on a local value exactly as the legacy blocks
+    chained on self.page_size."""
+    page_size = view.page_size
+    if (
+        view.attention_backend == "flashmla"
+        or view.decode_attention_backend == "flashmla"
+    ):
+        logger.warning(
+            "FlashMLA only supports a page_size of 64, change page_size to 64."
+        )
+        page_size = 64
+    if (
+        view.attention_backend == "cutlass_mla"
+        or view.decode_attention_backend == "cutlass_mla"
+    ):
+        logger.warning(
+            "Cutlass MLA only supports a page_size of 128, change page_size to 128."
+        )
+        page_size = 128
+    if (
+        view.attention_backend == "trtllm_mla"
+        or view.decode_attention_backend == "trtllm_mla"
+    ):
+        if page_size not in [32, 64]:
+            logger.warning(
+                f"TensorRT-LLM MLA only supports page_size of 32 or 64, changing page_size from {page_size} to 64."
+            )
+            page_size = 64
+    if (
+        view.attention_backend == "tokenspeed_mla"
+        or view.decode_attention_backend == "tokenspeed_mla"
+    ):
+        if page_size not in [32, 64]:
+            logger.warning(
+                f"tokenspeed_mla only supports page_size of 32 or 64, changing page_size from {page_size} to 64."
+            )
+            page_size = 64
+    if (
+        view.attention_backend == "cutedsl_mla"
+        or view.decode_attention_backend == "cutedsl_mla"
+        or view.prefill_attention_backend == "cutedsl_mla"
+    ):
+        if page_size not in [32, 64]:
+            logger.warning(
+                f"CuteDSL MLA only supports page_size of 32 or 64, changing page_size from {page_size} to 64."
+            )
+            page_size = 64
+    if (
+        view.attention_backend == "trtllm_mha"
+        or view.decode_attention_backend == "trtllm_mha"
+        or view.prefill_attention_backend == "trtllm_mha"
+    ):
+        if page_size not in [16, 32, 64]:
+            logger.warning(
+                f"TensorRT-LLM MHA only supports page_size of 16, 32 or 64, changing page_size from {page_size} to 64."
+            )
+            page_size = 64
+    if page_size != view.page_size:
+        return {"page_size": page_size}
+    return {}
+
+
+@register_post_process
+def _mla_kv_cache_dtype_checks(view: Any) -> dict:
+    """Read-only validation pass in the attention-backend compatibility
+    handler: the TRT-LLM and tokenspeed MLA backends constrain the resolved
+    kv-cache dtype (declarations never reach the field, so the checks read
+    the view)."""
+    if (
+        view.attention_backend == "trtllm_mla"
+        or view.decode_attention_backend == "trtllm_mla"
+    ):
+        if not is_blackwell_supported():
+            raise ValueError(
+                "TRTLLM MLA backend is only supported on Blackwell GPUs (SM100/SM12x). Please use a different backend."
+            )
+        if view.kv_cache_dtype not in ["fp8_e4m3", "fp4_e2m1", "bf16", "auto"]:
+            raise ValueError(
+                "TensorRT-LLM MLA backend only supports kv-cache-dtype of fp8_e4m3, fp4_e2m1, bf16, or auto."
+            )
+    if (
+        view.attention_backend == "tokenspeed_mla"
+        or view.decode_attention_backend == "tokenspeed_mla"
+    ):
+        if not is_blackwell_supported():
+            raise ValueError(
+                "tokenspeed_mla backend is only supported on Blackwell GPUs (SM100/SM12x)."
+            )
+        if view.kv_cache_dtype not in ["fp8_e4m3"]:
+            raise ValueError(
+                "tokenspeed_mla backend requires kv-cache-dtype=fp8_e4m3, "
+                f"got {view.kv_cache_dtype}."
+            )
+    return {}
+
+
+@register_post_process
+def _hisparse_validation(view: Any) -> dict:
+    """Read-only validation pass: --enable-hisparse constraints (model class,
+    radix cache, kv dtype, DSA backends) read the resolved values through the
+    view."""
+    from sglang.srt.arg_groups.hisparse_hook import validate_hisparse
+
+    validate_hisparse(view)
+    return {}
+
+
+@register_post_process
+def _cutedsl_prefill_backend_fill(view: Any) -> dict:
+    """Slot pass in the attention-backend compatibility handler: CuteDSL MLA
+    is decode-only, so validate the combination and default the prefill side
+    to trtllm_mla. The trtllm_mha check that follows at the legacy slot reads
+    the resolved value through the view."""
+    if not (
+        view.attention_backend == "cutedsl_mla"
+        or view.decode_attention_backend == "cutedsl_mla"
+        or view.prefill_attention_backend == "cutedsl_mla"
+    ):
+        return {}
+    assert (
+        view.prefill_attention_backend != "cutedsl_mla"
+    ), "CuteDSL MLA only supports decoding for now"
+    if not is_sm100_supported():
+        raise ValueError(
+            "CuteDSL MLA backend is only supported on Blackwell GPUs (SM100). Please use a different backend."
+        )
+    if view.kv_cache_dtype not in [
+        "fp8_e4m3",
+        "bf16",
+        "bfloat16",
+        "auto",
+    ]:
+        raise ValueError(
+            "CuteDSL MLA backend only supports kv-cache-dtype of fp8_e4m3, bf16, or auto."
+        )
+    if view.prefill_attention_backend is None:
+        return {"prefill_attention_backend": "trtllm_mla"}
+    return {}
+
+
+@register_post_process
+def _attention_backend_fa3_fp8_fallback(view: Any) -> dict:
+    if view.attention_backend == "fa3" and view.kv_cache_dtype == "fp8_e5m2":
+        logger.warning(
+            "FlashAttention3 only supports fp8_e4m3 if using FP8; "
+            "Setting attention backend to triton."
+        )
+        return {"attention_backend": "triton"}
+    return {}
+
+
+@register_post_process
+def _fa4_page_constraint(view: Any) -> dict:
+    if (
+        (
+            view.attention_backend == "fa4"
+            or view.decode_attention_backend == "fa4"
+            or view.prefill_attention_backend == "fa4"
+        )
+        and not view.use_mla_backend()
+        and is_sm100_supported()
+        # EAGLE topk>1 spec runs the two-pass page-tree cascade, which the FA4
+        # CUTLASS kernel aborts on at page_size>1. That path only works at
+        # page_size==1, so skip the 128 auto-force for it and keep the default.
+        and (view.speculative_eagle_topk or 0) <= 1
+    ):
+        logger.warning(
+            f"FA4 backend only supports page size 128 for non-MLA model architectures, changing page_size from {view.page_size} to 128."
+        )
+        return {"page_size": 128}
+    return {}
+
+
+@register_post_process
+def _attention_backend_platform_fallbacks(view: Any) -> dict:
+    if (
+        view.attention_backend == "intel_amx"
+        and view.device == "cpu"
+        and not cpu_has_amx_support()
+    ):
+        logger.warning(
+            "The current platform does not support Intel AMX, will fallback to torch_native backend."
+        )
+        return {"attention_backend": "torch_native"}
+    if (
+        view.attention_backend == "intel_xpu"
+        and view.device == "xpu"
+        and not xpu_has_xmx_support()
+    ):
+        logger.warning(
+            "The current platform does not support Intel XMX, will fallback to triton backend."
+        )
+        return {"attention_backend": "triton"}
+    return {}
+
+
+@register_post_process
+def _intel_xpu_page_constraint(view: Any) -> dict:
+    _, decode_backend = attention_backends_of(view)
+    if decode_backend == "intel_xpu":
+        if view.use_mla_backend():
+            supported_page_sizes = [16, 32, 64, 128]
+            msg = "Intel XPU attention backend for MLA Decode"
+        else:
+            supported_page_sizes = [64, 128]
+            msg = "Intel XPU attention backend"
+        if view.page_size not in supported_page_sizes:
+            logger.warning(
+                f"{msg} only supports page_sizes of {supported_page_sizes}, changing page_size from {view.page_size} to 128."
+            )
+            return {"page_size": 128}
+    return {}
+
+
+@register_post_process
+def _attention_backend_dual_chunk(view: Any) -> dict:
+    if (
+        getattr(view.get_model_config().hf_config, "dual_chunk_attention_config", None)
+        is not None
+    ):
+        if view.attention_backend is None:
+            logger.info("Dual chunk attention is turned on by default.")
+            return {"attention_backend": "dual_chunk_flash_attn"}
+        elif view.attention_backend != "dual_chunk_flash_attn":
+            raise ValueError(
+                "Dual chunk attention is enabled, but attention backend is set to "
+                f"{view.attention_backend}. Please set it to 'dual_chunk_flash_attn'."
+            )
+    return {}
+
+
+@register_post_process
+def _page_size_default(view: Any) -> dict:
+    if view.page_size is not None:
+        return {}
+
+    # SHUFFLE 5D vectorized KV layout (aiter backend + pa_decode_gluon)
+    # is tuned for and prefers page_size=64 — making it the default
+    # when the layout flag is set avoids users having to pass
+    # --page-size 64 explicitly. The env var is only consumed by the
+    # ROCm AITER backend, so the auto-bump is gated on HIP; on other
+    # platforms the SHUFFLE 5D pool has no consumer kernels and the
+    # env var is silently ignored (see MHATokenToKVPool).
+    if is_hip() and envs.SGLANG_AITER_KV_CACHE_LAYOUT.get().lower() == "vectorized_5d":
+        logger.info(
+            "Setting page_size=64 as default for "
+            "SGLANG_AITER_KV_CACHE_LAYOUT=vectorized_5d."
+        )
+        return {"page_size": 64}
+    if not is_musa():
+        return {"page_size": 1}
+    return {"page_size": 64}
+
+
+@register_post_process
+def _data_parallelism_defaults(view: Any) -> dict:
+    if view.dp_size == 1 and view.ep_join_mode != "scale":
+        return {"enable_dp_attention": False, "enable_dp_lm_head": False}
+    return {}
+
+
+@register_post_process
+def _dp_lm_head_validation(view: Any) -> dict:
+    """Read-only validation pass: dp-attention is a prerequisite for the
+    dp LM head. Reads the mid-resolution values through the view."""
+    if view.enable_dp_lm_head:
+        assert (
+            view.enable_dp_attention
+        ), "Please enable dp attention when setting enable_dp_lm_head. "
+    return {}
+
+
+@register_post_process
+def _moe_runner_backend_quant_constraints(view: Any) -> dict:
+    """The quantization-driven moe_runner_backend resolutions at the head of
+    _handle_moe_kernel_config. The backend-compatibility asserts and the
+    disable_shared_experts_fusion writes (post-publish writers exist for that
+    field) stay in the handler."""
+    moe_runner_backend = view.moe_runner_backend
+    if view.quantization == "nvfp4_online":
+        if not is_sm100_supported():
+            raise ValueError(
+                "--quantization nvfp4_online is supported only on "
+                "NVIDIA Blackwell SM100/SM103 GPUs."
+            )
+        if moe_runner_backend == "auto":
+            moe_runner_backend = "flashinfer_trtllm"
+        elif moe_runner_backend not in [
+            "flashinfer_trtllm",
+            "flashinfer_trtllm_routed",
+        ]:
+            raise ValueError(
+                "--quantization nvfp4_online supports only "
+                "--moe-runner-backend flashinfer_trtllm or "
+                "flashinfer_trtllm_routed."
+            )
+    if view.quantization == "mxfp8":
+        from sglang.srt.server_args import MXFP8_MOE_RUNNER_BACKEND_CHOICES
+
+        is_gfx95_mxfp8 = is_hip() and is_gfx95_supported()
+        allowed = list(MXFP8_MOE_RUNNER_BACKEND_CHOICES)
+        if is_gfx95_mxfp8:
+            allowed.append("triton")
+        mxfp8_default = "triton" if is_gfx95_mxfp8 else "flashinfer_trtllm"
+        if moe_runner_backend == "auto":
+            moe_runner_backend = mxfp8_default
+        elif moe_runner_backend not in allowed:
+            logger.warning(
+                "mxfp8 quantization supports only %s backends. " "Overriding %r.",
+                ", ".join(allowed),
+                moe_runner_backend,
+            )
+            moe_runner_backend = mxfp8_default
+    if (
+        moe_runner_backend == "auto"
+        and view.quantization == "modelopt_fp4"
+        and is_sm120_supported()
+    ):
+        moe_runner_backend = "flashinfer_cutlass"
+        logger.info(
+            "Use flashinfer_cutlass as MoE runner backend on SM120 for "
+            "modelopt_fp4 (trtllm-gen MoE kernels are SM100-only)"
+        )
+    if moe_runner_backend != view.moe_runner_backend:
+        return {"moe_runner_backend": moe_runner_backend}
+    return {}
+
+
+@register_post_process
+def _moe_runner_fusion_disable(view: Any) -> dict:
+    """FlashInfer CuteDSL / TRT-LLM / TRT-LLM-routed MoE runners require the
+    shared-experts fusion disabled; declared at the legacy write slots in
+    _handle_moe_kernel_config (before the deprecated cutlass env override, so
+    the runner value observed is the pre-override one)."""
+    runner = view.moe_runner_backend
+    if runner == "flashinfer_cutedsl":
+        logger.warning(
+            "FlashInfer CuteDSL MoE is enabled. --disable-shared-experts-fusion is automatically set."
+        )
+        return {"disable_shared_experts_fusion": True}
+    if runner in ("flashinfer_trtllm", "experimental_sgl_trtllm"):
+        logger.warning(
+            "FlashInfer TRTLLM MoE is enabled. --disable-shared-experts-fusion is automatically set."
+        )
+        return {"disable_shared_experts_fusion": True}
+    if runner == "flashinfer_trtllm_routed":
+        logger.warning(
+            "FlashInfer TRTLLM routed MoE is enabled. --disable-shared-experts-fusion is automatically set."
+        )
+        return {"disable_shared_experts_fusion": True}
+    return {}
+
+
+def _a2a_fusion_adjustments(view: Any) -> dict:
+    """A2A-backend-driven shared-experts fusion adjustments, declared at the
+    legacy write slots in _handle_a2a_moe: Waterfill requires the
+    fusion enabled; FlashInfer A2A requires it disabled."""
+    if view.moe_a2a_backend in ("deepep", "megamoe") and view.enable_waterfill:
+        if view.disable_shared_experts_fusion:
+            logger.warning(
+                "disable_shared_experts_fusion is overridden to False because Waterfill requires shared expert fusion."
+            )
+            return {"disable_shared_experts_fusion": False}
+        return {}
+    if view.moe_a2a_backend == "flashinfer":
+        logger.warning(
+            "Flashinfer MoE A2A is enabled. --disable-shared-experts-fusion is automatically set."
+        )
+        return {"disable_shared_experts_fusion": True}
+    return {}
+
+
+def _cutlass_moe_env_override(view: Any) -> dict:
+
+    if envs.SGLANG_CUTLASS_MOE.get():
+        logger.warning(
+            "SGLANG_CUTLASS_MOE is deprecated, use --moe-runner-backend=cutlass and/or --speculative-moe-runner-backend=cutlass instead"
+        )
+        assert view.quantization in [
+            "fp8",
+            "mxfp8",
+        ], "cutlass MoE is only supported with fp8/mxfp8 quantization"
+        return {"moe_runner_backend": "cutlass"}
+    return {}
+
+
+# Every A2A backend that forces expert parallelism to span the TP group.
+_A2A_EP_SPANNING_BACKENDS = frozenset(
+    {"megamoe", "deepep", "mooncake", "nixl", "ascend_fuseep", "flashinfer", "mori"}
+)
+
+
+@register_post_process
+def _a2a_backend_overrides(view: Any) -> dict:
+
+    moe_a2a_backend = view.moe_a2a_backend
+    if view.enable_waterfill and moe_a2a_backend not in ("deepep", "megamoe"):
+        logger.warning(
+            "moe_a2a_backend is overridden to 'deepep' because Waterfill "
+            "requires the DeepEP or MegaMOE backend."
+        )
+        moe_a2a_backend = "deepep"
+    if envs.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.get() and moe_a2a_backend != "megamoe":
+        moe_a2a_backend = "megamoe"
+        logger.info(
+            "SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE is set, "
+            "auto-configuring --moe-a2a-backend megamoe."
+        )
+    if moe_a2a_backend != view.moe_a2a_backend:
+        return {"moe_a2a_backend": moe_a2a_backend}
+    return {}
+
+
+@register_post_process
+def _a2a_ep_size(view: Any) -> dict:
+    if view.moe_a2a_backend in _A2A_EP_SPANNING_BACKENDS:
+        return {"ep_size": view.tp_size}
+    return {}
+
+
+@register_post_process
+def _pipeline_parallel_overlap_disable(view: Any) -> dict:
+    if view.pp_size > 1:
+        logger.warning("Pipeline parallelism is incompatible with overlap schedule.")
+        return {"disable_overlap_schedule": True}
+    return {}
+
+
+@register_post_process
+def _speculative_moe_runner_default(view: Any) -> dict:
+    """Default the speculative (draft) MoE runner backend to the resolved
+    target-model backend. Invoked at the head of the speculative-decoding
+    hook, after the MoE kernel chain has resolved."""
+    if view.speculative_moe_runner_backend is None:
+        return {"speculative_moe_runner_backend": view.moe_runner_backend}
+    return {}
+
+
+@register_post_process
+def _gguf_quantization(view: Any) -> dict:
+    from sglang.srt.utils.hf_transformers_utils import check_gguf_file
+
+    if (view.load_format == "auto" or view.load_format == "gguf") and check_gguf_file(
+        view.model_path
+    ):
+        return {"quantization": "gguf"}
+    return {}
+
+
+@register_post_process
+def _dllm_attention_backend(view: Any) -> dict:
+    if view.dllm_algorithm is None:
+        return {}
+    if is_hip():
+        if view.attention_backend not in ["triton", "aiter"]:
+            logger.warning(
+                "Attention backend is set to triton for diffusion LLM inference on AMD GPUs"
+            )
+            return {"attention_backend": "triton"}
+    elif is_npu():
+        if view.attention_backend != "ascend":
+            logger.warning(
+                "Attention backend is overridden to 'ascend' when running on NPU for diffusion LLM inference."
+            )
+            return {"attention_backend": "ascend"}
+    elif view.cuda_graph_config.decode.backend != Backend.DISABLED:
+        if view.attention_backend != "flashinfer":
+            logger.warning(
+                "Attention backend is set to flashinfer because of enabling cuda graph in diffusion LLM inference"
+            )
+            return {"attention_backend": "flashinfer"}
+    return {}
+
+
+@register_post_process
+def _dllm_overlap_disable(view: Any) -> dict:
+    if view.dllm_algorithm is None:
+        return {}
+    if view.disable_overlap_schedule:
+        return {}
+    logger.warning(
+        "Overlap schedule is disabled because of using diffusion LLM inference"
+    )
+    return {"disable_overlap_schedule": True}
+
+
+@register_post_process
+def _dllm_page_size(view: Any) -> dict:
+    if view.dllm_algorithm is None:
+        return {}
+    from sglang.srt.dllm.config import DllmConfig
+
+    config = DllmConfig.from_server_args(view)
+    if not view.disable_radix_cache and view.page_size % config.block_size != 0:
+        logger.warning(
+            f"Setting page size to {config.block_size} for diffusion LLM inference"
+        )
+        return {"page_size": config.block_size}
+    if view.page_size > config.block_size:
+        # Legacy scheduler-init fallback, folded into the pass: the page
+        # size must not exceed the dllm block size.
+        logger.warning(
+            "WARNING: "
+            f"The page size {view.page_size} should not be larger than dllm block size {config.block_size}."
+            f"Page size now falls back to {config.block_size}"
+        )
+        return {"page_size": config.block_size}
+    return {}
+
+
+def validate_declarations(
+    server_args: Any,
+    declarations: Sequence[Tuple[str, Dict[str, Any]]],
+) -> None:
+    """Fail-fast whitelist check at declaration time: a registry typo or a
+    not-yet-resolvable field must be rejected at its slot, not only at
+    publish time. Declarations never mutate ``server_args``.
+    """
+    # Non-dataclass fixtures carry no Arg metadata (mirrors the
+    # resolvable_fields escape); only real ServerArgs is validated.
+    if not dataclasses.is_dataclass(type(server_args)):
+        return
+    whitelist = resolvable_fields(type(server_args))
+    for source, decl in declarations:
+        unknown = set(decl) - whitelist
+        if unknown:
+            raise ValueError(
+                f"{source}: {sorted(unknown)} not model-overridable; "
+                "declarations are limited to the fields the publish gate "
+                "accepts."
+            )
+
+
+def _hrm_text_attention_force(view: Any) -> dict:
+    """HRM-Text's bidirectional prefix attention only works on the Triton
+    backend. Invoked as the last attention declaration of the resolution
+    (mirroring the legacy runner-side force, which ran after the whole
+    pipeline)."""
+    if view.attention_backend not in (None, "triton"):
+        logger.warning(
+            f"Overriding --attention-backend "
+            f"{view.attention_backend!r} -> 'triton': only the "
+            "Triton backend supports HRM-Text's bidirectional prefix "
+            "attention."
+        )
+    return {"attention_backend": "triton"}

sglang_patch/srt/arg_groups/speculative_hook.py

@@ -0,0 +1,792 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+def _disable_overlap_schedule_for_cpu(server_args: ServerArgs) -> None:
+    if server_args.device != "cpu" or server_args.disable_overlap_schedule:
+        return
+
+    server_args.disable_overlap_schedule = True
+    logger.warning(
+        "Overlap schedule is not implemented for speculative decoding on CPU."
+    )
+
+
+def _resolve_speculative_algorithm_alias(
+    speculative_algorithm: Optional[str],
+    speculative_draft_model_path: Optional[str],
+    trust_remote_code: bool = False,
+    kwargs: Optional[dict] = {},
+) -> Optional[str]:
+    """Resolve CLI speculative algorithm; NEXTN/EAGLE may become FROZEN_KV_MTP for Gemma4 assistant drafts."""
+
+    is_gemma4_draft = False
+    if speculative_draft_model_path:
+        from sglang.srt.utils.hf_transformers_utils import get_config
+
+        cfg = get_config(
+            speculative_draft_model_path, trust_remote_code=trust_remote_code, **kwargs
+        )
+        draft_archs = getattr(cfg, "architectures", None) or []
+        is_gemma4_draft = any(
+            arch in ("Gemma4AssistantForCausalLM", "Gemma4UnifiedAssistantForCausalLM")
+            for arch in draft_archs
+        )
+
+    if speculative_algorithm == "EAGLE3" and is_gemma4_draft:
+        raise ValueError(
+            "Gemma4AssistantForCausalLM draft requires "
+            "--speculative-algorithm NEXTN or EAGLE; EAGLE3 is "
+            "not supported for this draft architecture."
+        )
+
+    if speculative_algorithm == "NEXTN" or speculative_algorithm == "EAGLE":
+        if is_gemma4_draft:
+            logger.info(
+                "Detected Gemma4AssistantForCausalLM draft; "
+                f"promoting --speculative-algorithm {speculative_algorithm} to FROZEN_KV_MTP."
+            )
+            return "FROZEN_KV_MTP"
+        return "EAGLE"
+
+    return speculative_algorithm
+
+
+def handle_speculative_decoding(server_args: ServerArgs) -> None:
+    if (
+        server_args.speculative_draft_model_path is not None
+        and server_args.speculative_draft_model_revision is None
+    ):
+        server_args.speculative_draft_model_revision = "main"
+
+    # Moved to the resolution pipeline (arg_groups/overrides.py:
+    # _speculative_moe_runner_default), invoked here at its legacy slot.
+    from sglang.srt.arg_groups.overrides import (
+        _speculative_moe_runner_default,
+        run_post_process_pass,
+    )
+
+    run_post_process_pass(server_args, _speculative_moe_runner_default)
+
+    if server_args.speculative_algorithm is not None:
+        server_args.speculative_algorithm = server_args.speculative_algorithm.upper()
+
+    # Removal notice for the retired env var; raw os.getenv on purpose -- the
+    # Envs descriptor is gone. Drop this check after one release.
+    if os.getenv("SGLANG_ENABLE_SPEC_V2") is not None:
+        logger.warning(
+            "SGLANG_ENABLE_SPEC_V2 has been removed: speculative decoding "
+            "always runs the V2 worker. Use --disable-overlap-schedule to "
+            "select the non-overlap (synchronous) path."
+        )
+
+    kwargs = {}
+
+    override_config_file = server_args.decrypted_draft_config_file
+    if override_config_file and override_config_file.strip():
+        kwargs["_configuration_file"] = override_config_file.strip()
+
+    server_args.speculative_algorithm = _resolve_speculative_algorithm_alias(
+        server_args.speculative_algorithm,
+        server_args.speculative_draft_model_path,
+        trust_remote_code=server_args.trust_remote_code,
+        kwargs=kwargs,
+    )
+
+    # Validate --speculative-draft-window-size once, regardless of algorithm.
+    # Consumed by DFLASH (compact draft KV cache) and Llama EAGLE-3 (drafter attention SWA).
+    if server_args.speculative_draft_window_size is not None:
+        window_size = int(server_args.speculative_draft_window_size)
+        if window_size <= 0:
+            raise ValueError(
+                f"--speculative-draft-window-size must be positive, got {window_size}."
+            )
+        server_args.speculative_draft_window_size = window_size
+        if server_args.speculative_algorithm not in ("EAGLE3", "DFLASH"):
+            logger.warning(
+                "--speculative-draft-window-size has no effect with "
+                "speculative_algorithm=%s (honored by Llama EAGLE-3 and DFLASH only).",
+                server_args.speculative_algorithm,
+            )
+
+    algo = None
+    if server_args.speculative_algorithm is not None:
+        from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+        from sglang.srt.speculative.spec_registry import CustomSpecAlgo
+
+        algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm)
+
+        # TODO: move the per-algorithm validation below into spec module hooks.
+        if isinstance(algo, CustomSpecAlgo) and algo.validate_server_args is not None:
+            algo.validate_server_args(server_args)
+
+    if server_args.speculative_skip_dp_mlp_sync:
+        assert server_args.speculative_algorithm == "EAGLE", (
+            "--speculative-skip-dp-mlp-sync is only supported with "
+            f"speculative_algorithm == EAGLE, got {server_args.speculative_algorithm}."
+        )
+
+    if server_args.speculative_adaptive:
+        _maybe_disable_adaptive(server_args)
+        if server_args.speculative_adaptive:
+            _init_adaptive_speculative_params(server_args)
+
+    if algo is not None:
+        algo.handle_server_args(server_args)
+
+
+def _handle_dflash(server_args: ServerArgs) -> None:
+    from sglang.srt.arg_groups.overrides import resolved_view
+
+    if not server_args.device.startswith("cuda"):
+        raise ValueError("DFLASH speculative decoding only supports CUDA device.")
+
+    if resolved_view(server_args).enable_dp_attention:
+        raise ValueError(
+            "Currently DFLASH speculative decoding does not support dp attention."
+        )
+
+    if server_args.pp_size != 1:
+        raise ValueError(
+            "Currently DFLASH speculative decoding only supports pp_size == 1."
+        )
+
+    if server_args.speculative_draft_model_path is None:
+        raise ValueError(
+            "DFLASH speculative decoding requires setting --speculative-draft-model-path."
+        )
+
+    # DFLASH does not use EAGLE-style `num_steps`/`topk`, but those fields still
+    # affect generic scheduler/KV-cache accounting (buffer sizing, KV freeing,
+    # RoPE reservation). Force them to 1 to avoid surprising memory behavior.
+    #
+    # For DFlash, the natural unit is `block_size` (verify window length).
+    if server_args.speculative_num_steps is None:
+        server_args.speculative_num_steps = 1
+    elif int(server_args.speculative_num_steps) != 1:
+        logger.warning(
+            "DFLASH only supports speculative_num_steps == 1; overriding speculative_num_steps=%s to 1.",
+            server_args.speculative_num_steps,
+        )
+        server_args.speculative_num_steps = 1
+
+    if server_args.speculative_eagle_topk is None:
+        server_args.speculative_eagle_topk = 1
+    elif int(server_args.speculative_eagle_topk) != 1:
+        logger.warning(
+            "DFLASH only supports speculative_eagle_topk == 1; overriding speculative_eagle_topk=%s to 1.",
+            server_args.speculative_eagle_topk,
+        )
+        server_args.speculative_eagle_topk = 1
+
+    if server_args.speculative_dflash_block_size is not None:
+        if int(server_args.speculative_dflash_block_size) <= 0:
+            raise ValueError(
+                "DFLASH requires --speculative-dflash-block-size to be positive, "
+                f"got {server_args.speculative_dflash_block_size}."
+            )
+        if server_args.speculative_num_draft_tokens is not None and int(
+            server_args.speculative_num_draft_tokens
+        ) != int(server_args.speculative_dflash_block_size):
+            raise ValueError(
+                "Both --speculative-num-draft-tokens and --speculative-dflash-block-size are set "
+                "but they differ. For DFLASH they must match. "
+                f"speculative_num_draft_tokens={server_args.speculative_num_draft_tokens}, "
+                f"speculative_dflash_block_size={server_args.speculative_dflash_block_size}."
+            )
+        server_args.speculative_num_draft_tokens = int(
+            server_args.speculative_dflash_block_size
+        )
+
+    if server_args.speculative_num_draft_tokens is None:
+        from sglang.srt.speculative.dflash_utils import (
+            parse_dflash_draft_config,
+        )
+
+        model_override_args = json.loads(server_args.json_model_override_args)
+        inferred_block_size = None
+        try:
+            from sglang.srt.utils.hf_transformers_utils import get_config
+
+            draft_hf_config = get_config(
+                server_args.speculative_draft_model_path,
+                trust_remote_code=server_args.trust_remote_code,
+                revision=server_args.speculative_draft_model_revision,
+                model_override_args=model_override_args,
+            )
+            inferred_block_size = parse_dflash_draft_config(
+                draft_hf_config=draft_hf_config
+            ).resolve_block_size(default=None)
+        except Exception as e:
+            logger.warning(
+                "Failed to infer DFLASH block_size from draft model config; "
+                "defaulting speculative_num_draft_tokens to 16. Error: %s",
+                e,
+            )
+
+        if inferred_block_size is None:
+            inferred_block_size = 16
+            logger.warning(
+                "speculative_num_draft_tokens is not set; defaulting to %d for DFLASH.",
+                inferred_block_size,
+            )
+        server_args.speculative_num_draft_tokens = inferred_block_size
+
+    if server_args.speculative_draft_window_size is not None:
+        draft_tokens = int(server_args.speculative_num_draft_tokens)
+        if server_args.speculative_draft_window_size < draft_tokens:
+            raise ValueError(
+                "--speculative-draft-window-size must be >= "
+                "--speculative-num-draft-tokens (block_size). "
+                f"window_size={server_args.speculative_draft_window_size}, block_size={draft_tokens}."
+            )
+
+    _resolve_dflash_draft_attention_backend(server_args)
+
+    if server_args.max_running_requests is None:
+        server_args.max_running_requests = 48
+        logger.warning(
+            "Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
+        )
+
+    if server_args.enable_mixed_chunk:
+        server_args.enable_mixed_chunk = False
+        logger.warning(
+            "Mixed chunked prefill is disabled because of using dflash speculative decoding."
+        )
+
+
+def _target_checkpoint_bundles_dspark_draft(server_args: ServerArgs) -> bool:
+    from sglang.srt.speculative.dspark_components.dspark_config import (
+        checkpoint_bundles_dspark_draft,
+    )
+
+    return checkpoint_bundles_dspark_draft(server_args.get_model_config().hf_config)
+
+
+def _handle_dspark(server_args: ServerArgs) -> None:
+    if not server_args.device.startswith("cuda"):
+        raise ValueError("DSpark speculative decoding only supports CUDA device.")
+
+    if server_args.enable_dp_attention:
+        if not server_args.enable_dp_lm_head:
+            raise ValueError("DSpark with dp attention requires --enable-dp-lm-head.")
+        if server_args.moe_a2a_backend != "none":
+            raise ValueError(
+                "DSpark with dp attention only supports the built-in TP MoE "
+                f"(moe_a2a_backend='none'), got {server_args.moe_a2a_backend!r}."
+            )
+        if server_args.attn_cp_size > 1:
+            raise ValueError(
+                "DSpark with dp attention does not support context parallel "
+                f"(attn_cp_size={server_args.attn_cp_size})."
+            )
+        if (
+            server_args.speculative_moe_a2a_backend is not None
+            and server_args.speculative_moe_a2a_backend != server_args.moe_a2a_backend
+        ):
+            raise ValueError(
+                "DSpark ignores --speculative-moe-a2a-backend; with dp attention it "
+                f"must match the target moe_a2a_backend={server_args.moe_a2a_backend!r} "
+                f"(got {server_args.speculative_moe_a2a_backend!r})."
+            )
+
+    if server_args.pp_size != 1:
+        raise ValueError(
+            "Currently DSpark speculative decoding only supports pp_size == 1."
+        )
+
+    if server_args.speculative_draft_model_path is None:
+        if _target_checkpoint_bundles_dspark_draft(server_args):
+            server_args.speculative_draft_model_path = server_args.model_path
+            server_args.speculative_draft_model_revision = server_args.revision
+            logger.info(
+                "DSpark draft weights are bundled in the target checkpoint; "
+                "defaulting --speculative-draft-model-path to --model-path (%s).",
+                server_args.model_path,
+            )
+        else:
+            raise ValueError(
+                "DSpark dense speculative decoding requires setting "
+                "--speculative-draft-model-path."
+            )
+
+    if server_args.speculative_num_steps is None:
+        server_args.speculative_num_steps = 1
+    elif int(server_args.speculative_num_steps) != 1:
+        logger.warning(
+            "DSpark only supports speculative_num_steps == 1; overriding speculative_num_steps=%s to 1.",
+            server_args.speculative_num_steps,
+        )
+        server_args.speculative_num_steps = 1
+
+    if server_args.speculative_eagle_topk is None:
+        server_args.speculative_eagle_topk = 1
+    elif int(server_args.speculative_eagle_topk) != 1:
+        logger.warning(
+            "DSpark only supports speculative_eagle_topk == 1; overriding speculative_eagle_topk=%s to 1.",
+            server_args.speculative_eagle_topk,
+        )
+        server_args.speculative_eagle_topk = 1
+
+    gamma: Optional[int] = None
+    if server_args.speculative_dspark_block_size is not None:
+        if int(server_args.speculative_dspark_block_size) <= 0:
+            raise ValueError(
+                "DSpark requires --speculative-dspark-block-size to be positive, "
+                f"got {server_args.speculative_dspark_block_size}."
+            )
+        gamma = int(server_args.speculative_dspark_block_size)
+    else:
+        from sglang.srt.speculative.dspark_components.dspark_config import (
+            DEFAULT_DSPARK_GAMMA,
+            read_draft_checkpoint_gamma,
+        )
+
+        try:
+            gamma = read_draft_checkpoint_gamma(server_args=server_args)
+        except Exception as e:
+            logger.warning(
+                "Failed to read DSpark gamma from draft model config; "
+                "cannot cross-check --speculative-num-draft-tokens. Error: %s",
+                e,
+            )
+        if gamma is None and server_args.speculative_num_draft_tokens is None:
+            gamma = DEFAULT_DSPARK_GAMMA
+            logger.warning(
+                "DSpark gamma is not set; defaulting to %d.",
+                gamma,
+            )
+
+    if gamma is not None:
+        verify_window = int(gamma) + 1
+        if (
+            server_args.speculative_num_draft_tokens is not None
+            and int(server_args.speculative_num_draft_tokens) != verify_window
+        ):
+            raise ValueError(
+                "DSpark speculative_num_draft_tokens must equal gamma + 1 "
+                f"(= {verify_window} for gamma={gamma}), but got "
+                f"speculative_num_draft_tokens={server_args.speculative_num_draft_tokens}."
+            )
+        server_args.speculative_num_draft_tokens = verify_window
+
+    if server_args.speculative_num_draft_tokens is None:
+        raise ValueError(
+            "DSpark could not resolve speculative_num_draft_tokens; set "
+            "--speculative-dspark-block-size (= gamma)."
+        )
+    if int(server_args.speculative_num_draft_tokens) < 2:
+        raise ValueError(
+            "DSpark speculative_num_draft_tokens must be >= 2 (= gamma + 1), "
+            f"got {server_args.speculative_num_draft_tokens}."
+        )
+
+    if server_args.max_running_requests is None:
+        server_args.max_running_requests = 48
+        logger.warning(
+            "Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
+        )
+
+    if server_args.enable_mixed_chunk:
+        server_args.enable_mixed_chunk = False
+        logger.warning(
+            "Mixed chunked prefill is disabled because of using dspark speculative decoding."
+        )
+
+    from sglang.srt.speculative.ragged_verify import (
+        RaggedVerifyMode,
+        read_ragged_verify_mode,
+    )
+
+    ragged_mode = read_ragged_verify_mode()
+    if (
+        server_args.speculative_dspark_align_verify_tokens_to_graph_tier
+        and ragged_mode is not RaggedVerifyMode.COMPACT
+    ):
+        logger.warning(
+            "--speculative-dspark-align-verify-tokens-to-graph-tier only takes "
+            "effect with SGLANG_RAGGED_VERIFY_MODE=compact (got %r); it will be "
+            "a no-op.",
+            ragged_mode.value,
+        )
+    if (
+        server_args.speculative_dspark_sps_table_path
+        and ragged_mode is RaggedVerifyMode.STATIC
+    ):
+        logger.warning(
+            "--speculative-dspark-sps-table-path feeds the ragged-verify budget "
+            "scheduler, which is off under SGLANG_RAGGED_VERIFY_MODE=static; it "
+            "will be a no-op."
+        )
+
+
+def _resolve_dflash_draft_attention_backend(server_args: ServerArgs) -> None:
+    """Resolve `speculative_draft_attention_backend` to a final, supported value.
+
+    Consumed by ModelRunner's `is_draft_worker` override (one backend for all
+    draft modes).
+    """
+    from sglang.srt.utils import is_hip
+
+    supported_draft_backends = ("flashinfer", "fa3", "fa4", "triton", "ascend")
+    # Use triton on ROCm (no FlashInfer), flashinfer on CUDA.
+    fallback_backend = "triton" if is_hip() else "flashinfer"
+
+    draft_backend = server_args.speculative_draft_attention_backend
+    if draft_backend is None:
+        from sglang.srt.arg_groups.overrides import (
+            attention_backends_of,
+            resolved_view,
+        )
+
+        draft_backend, _ = attention_backends_of(resolved_view(server_args))
+    if draft_backend is None:
+        draft_backend = fallback_backend
+    elif draft_backend == "trtllm_mha":
+        logger.warning(
+            "DFLASH draft worker does not support 'trtllm_mha' because the "
+            "draft path requires per-layer DFlash attention. Falling back to "
+            "'%s'.",
+            fallback_backend,
+        )
+        draft_backend = fallback_backend
+    elif draft_backend not in supported_draft_backends:
+        logger.warning(
+            "DFLASH draft worker only supports attention_backend in %s for now, "
+            "but got %r. Falling back to '%s'.",
+            supported_draft_backends,
+            draft_backend,
+            fallback_backend,
+        )
+        draft_backend = fallback_backend
+    # FIXME: avoid overriding server args directly; pass the resolved draft
+    # backend to the draft worker explicitly instead.
+    server_args.speculative_draft_attention_backend = draft_backend
+
+
+def _handle_frozen_kv_mtp(server_args: ServerArgs) -> None:
+    if server_args.max_running_requests is None:
+        server_args.max_running_requests = 48
+        logger.warning(
+            "Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
+        )
+
+    if server_args.enable_mixed_chunk:
+        server_args.enable_mixed_chunk = False
+        logger.warning(
+            "Mixed chunked prefill is disabled because of using "
+            "Frozen-KV MTP speculative decoding."
+        )
+
+
+def _handle_eagle_family(server_args: ServerArgs) -> None:
+    from sglang.srt.arg_groups.overrides import (
+        attention_backends_of,
+        resolved_view,
+    )
+
+    if (
+        server_args.speculative_algorithm == "STANDALONE"
+        and resolved_view(server_args).enable_dp_attention
+    ):
+        # TODO: support dp attention for standalone speculative decoding
+        raise ValueError(
+            "Currently standalone speculative decoding does not support dp attention."
+        )
+
+    if server_args.max_running_requests is None:
+        server_args.max_running_requests = 48
+        logger.warning(
+            "Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
+        )
+
+    _disable_overlap_schedule_for_cpu(server_args)
+
+    if resolved_view(server_args).disable_overlap_schedule:
+        logger.warning(
+            "Non-overlap (synchronous) spec v2 is used for eagle/eagle3/standalone "
+            "speculative decoding."
+        )
+
+    if server_args.enable_mixed_chunk:
+        server_args.enable_mixed_chunk = False
+        logger.warning(
+            "Mixed chunked prefill is disabled because of using "
+            "eagle speculative decoding."
+        )
+
+    model_arch = server_args.get_model_config().hf_config.architectures[0]
+    if model_arch in [
+        "DeepseekV32ForCausalLM",
+        "DeepseekV3ForCausalLM",
+        "AgnesForCausalLM",
+        "Glm4MoeForCausalLM",
+        "Glm4MoeLiteForCausalLM",
+        "GlmMoeDsaForCausalLM",
+        "BailingMoeForCausalLM",
+        "BailingMoeV2ForCausalLM",
+        "BailingMoeV2_5ForCausalLM",
+        "MistralLarge3ForCausalLM",
+        "PixtralForConditionalGeneration",
+        "HYV3ForCausalLM",
+    ]:
+        if server_args.speculative_draft_model_path is None:
+            server_args.speculative_draft_model_path = server_args.model_path
+            server_args.speculative_draft_model_revision = server_args.revision
+        else:
+            if model_arch not in [
+                "MistralLarge3ForCausalLM",
+                "PixtralForConditionalGeneration",
+            ]:
+                logger.warning(
+                    "DeepSeek MTP does not require setting speculative_draft_model_path."
+                )
+
+    if (
+        not server_args.speculative_adaptive
+        and server_args.speculative_num_steps is None
+    ):
+        assert (
+            server_args.speculative_eagle_topk is None
+            and server_args.speculative_num_draft_tokens is None
+        )
+
+        (
+            server_args.speculative_num_steps,
+            server_args.speculative_eagle_topk,
+            server_args.speculative_num_draft_tokens,
+        ) = _auto_choose_speculative_params(server_args, model_arch)
+
+    if "trtllm_mha" in attention_backends_of(resolved_view(server_args)):
+        if server_args.speculative_eagle_topk > 1:
+            raise ValueError(
+                "trtllm_mha backend only supports topk = 1 for speculative decoding."
+            )
+
+    if server_args.speculative_use_rejection_sampling:
+        # Resolved alias by now: NEXTN -> EAGLE, Gemma4 draft -> FROZEN_KV_MTP.
+        # Only the EAGLE/EAGLE3 draft workers emit a target-vocab proposal that
+        # the rejection-sampling kernel consumes; everything else (STANDALONE,
+        # FROZEN_KV_MTP, NGRAM, DFLASH) is unsupported.
+        if server_args.speculative_algorithm not in ("EAGLE", "EAGLE3"):
+            raise NotImplementedError(
+                "--speculative-use-rejection-sampling is only supported for "
+                "EAGLE / EAGLE3 / NEXTN, not "
+                f"speculative_algorithm={server_args.speculative_algorithm}."
+            )
+        if server_args.speculative_eagle_topk != 1:
+            raise ValueError(
+                "--speculative-use-rejection-sampling requires --speculative-eagle-topk=1."
+            )
+        if (
+            server_args.speculative_accept_threshold_single != 1.0
+            or server_args.speculative_accept_threshold_acc != 1.0
+        ):
+            raise ValueError(
+                "--speculative-use-rejection-sampling is incompatible with "
+                "--speculative-accept-threshold-single / "
+                "--speculative-accept-threshold-acc; rejection sampling ignores "
+                "the accept thresholds."
+            )
+        if server_args.enable_deterministic_inference:
+            raise ValueError(
+                "--speculative-use-rejection-sampling is incompatible with "
+                "--enable-deterministic-inference; the sampling kernel draws "
+                "coins from the global RNG and is not batch-invariant."
+            )
+
+        from sglang.srt.arg_groups.overrides import resolved_view
+
+        if (
+            resolved_view(server_args).enable_multi_layer_eagle
+            and server_args.speculative_eagle_topk != 1
+        ):
+            raise ValueError(
+                "--speculative-use-rejection-sampling with multi-layer EAGLE "
+                "(--enable-multi-layer-eagle) requires --speculative-eagle-topk 1; "
+                "rejection sampling is only implemented for the linear (topk=1) chain."
+            )
+        logger.info(
+            "Rejection sampling is enabled for speculative decoding "
+            "(speculative_use_rejection_sampling=True)."
+        )
+
+    if (
+        server_args.speculative_eagle_topk == 1
+        and server_args.speculative_num_draft_tokens
+        != server_args.speculative_num_steps + 1
+    ):
+        logger.warning(
+            "speculative_num_draft_tokens is adjusted to speculative_num_steps + 1 when speculative_eagle_topk == 1"
+        )
+        server_args.speculative_num_draft_tokens = server_args.speculative_num_steps + 1
+
+    # topk > 1 + page_size > 1 needs the two-pass cascade draft-decode (shared prefix
+    # pass + per-branch expand pass with prefix-tail dup). Only these backends implement
+    # it; flashmla / trtllm_mla / cutlass_mla can't express the per-branch tree, so reject.
+    _PAGE_TREE_SPEC_BACKENDS = ("flashinfer", "fa3", "triton")
+    view = resolved_view(server_args)
+    if (
+        server_args.speculative_eagle_topk > 1
+        and view.page_size > 1
+        and view.attention_backend not in _PAGE_TREE_SPEC_BACKENDS
+    ):
+        raise ValueError(
+            f"speculative_eagle_topk > 1 with page_size > 1 is only supported on "
+            f"{_PAGE_TREE_SPEC_BACKENDS}; got attention_backend="
+            f"{view.attention_backend!r}. Use page_size == 1 or one of those backends."
+        )
+
+
+def _handle_ngram(server_args: ServerArgs) -> None:
+    if server_args.device not in ("cuda", "cpu"):
+        raise ValueError(
+            "Ngram speculative decoding only supports CUDA or CPU devices."
+        )
+
+    _disable_overlap_schedule_for_cpu(server_args)
+
+    if server_args.max_running_requests is None:
+        server_args.max_running_requests = 48
+        logger.warning(
+            "Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
+        )
+
+    server_args.enable_mixed_chunk = False
+    server_args.speculative_eagle_topk = server_args.speculative_ngram_max_bfs_breadth
+    if server_args.speculative_num_draft_tokens is None:
+        server_args.speculative_num_draft_tokens = 12
+        logger.warning(
+            "speculative_num_draft_tokens is set to 12 by default for ngram speculative decoding. "
+            "You can override this by explicitly setting --speculative-num-draft-tokens."
+        )
+    if server_args.speculative_num_steps is None:
+        server_args.speculative_num_steps = (
+            server_args.speculative_num_draft_tokens
+            // server_args.speculative_eagle_topk
+        )
+    if server_args.speculative_ngram_external_corpus_path is not None:
+        if server_args.speculative_ngram_external_sam_budget <= 0:
+            raise ValueError(
+                "--speculative-ngram-external-sam-budget must be positive when "
+                "--speculative-ngram-external-corpus-path is set."
+            )
+        if server_args.speculative_ngram_external_corpus_max_tokens <= 0:
+            raise ValueError(
+                "--speculative-ngram-external-corpus-max-tokens must be positive when "
+                "--speculative-ngram-external-corpus-path is set."
+            )
+        if (
+            server_args.speculative_ngram_external_sam_budget
+            > server_args.speculative_num_draft_tokens - 1
+        ):
+            raise ValueError(
+                "speculative_ngram_external_sam_budget must be less than or equal to "
+                f"speculative_num_draft_tokens - 1 ({server_args.speculative_num_draft_tokens - 1})."
+            )
+    logger.warning(
+        "The mixed chunked prefill are disabled because of "
+        "using ngram speculative decoding."
+    )
+
+    from sglang.srt.arg_groups.overrides import resolved_view
+
+    view = resolved_view(server_args)
+    if (
+        server_args.speculative_eagle_topk > 1
+        and view.page_size > 1
+        and view.attention_backend != "flashinfer"
+    ):
+        raise ValueError(
+            f"speculative_eagle_topk({server_args.speculative_eagle_topk}) > 1 "
+            f"with page_size({view.page_size}) > 1 is unstable "
+            "and produces incorrect results for paged attention backends. "
+            "This combination is only supported for the 'flashinfer' backend."
+        )
+    if view.enable_dp_attention:
+        # TODO: support dp attention for ngram speculative decoding
+        raise ValueError(
+            "Currently ngram speculative decoding does not support dp attention."
+        )
+
+
+def _maybe_disable_adaptive(server_args: ServerArgs) -> None:
+    from sglang.srt.speculative.adaptive_spec_params import (
+        adaptive_unsupported_reason,
+    )
+
+    reason = adaptive_unsupported_reason(server_args)
+    if reason is not None:
+        logger.warning(
+            f"speculative_adaptive disabled: {reason}. "
+            "Falling back to static speculative params."
+        )
+        server_args.speculative_adaptive = False
+
+
+def _init_adaptive_speculative_params(server_args: ServerArgs) -> None:
+    from sglang.srt.speculative.adaptive_spec_params import (
+        resolve_candidate_steps_from_config,
+    )
+
+    candidate_steps = resolve_candidate_steps_from_config(
+        cfg_path=server_args.speculative_adaptive_config,
+    )
+
+    if server_args.speculative_eagle_topk is None:
+        server_args.speculative_eagle_topk = 1
+
+    if server_args.speculative_num_steps is None:
+        server_args.speculative_num_steps = candidate_steps[len(candidate_steps) // 2]
+
+    if server_args.speculative_num_steps not in candidate_steps:
+        raise ValueError(
+            f"--speculative-num-steps={server_args.speculative_num_steps} "
+            f"is not in the adaptive config candidate_steps {candidate_steps}. "
+            "Pass one of those values."
+        )
+
+    server_args.speculative_num_draft_tokens = server_args.speculative_num_steps + 1
+
+
+def _auto_choose_speculative_params(server_args: ServerArgs, model_arch: str) -> tuple:
+    """
+    Automatically choose the parameters for speculative decoding.
+
+    You can tune them on your own models and prompts with scripts/playground/bench_speculative.py
+    """
+    if server_args.speculative_algorithm == "STANDALONE":
+        return (3, 1, 4)
+    if model_arch in ["LlamaForCausalLM"]:
+        return (5, 4, 8)
+    elif model_arch in [
+        "DeepseekV32ForCausalLM",
+        "DeepseekV3ForCausalLM",
+        "DeepseekV2ForCausalLM",
+        "GptOssForCausalLM",
+        "Glm4MoeForCausalLM",
+        "Glm4MoeLiteForCausalLM",
+        "GlmMoeDsaForCausalLM",
+        "BailingMoeForCausalLM",
+        "BailingMoeV2ForCausalLM",
+        "BailingMoeV2_5ForCausalLM",
+        "MistralLarge3ForCausalLM",
+        "PixtralForConditionalGeneration",
+        "MiMoV2ForCausalLM",
+        "MiMoV2FlashForCausalLM",
+    ]:
+        return (3, 1, 4)
+    elif model_arch in ["Grok1ForCausalLM", "Grok1VForCausalLM"]:
+        return (5, 4, 8)
+    else:
+        return (3, 1, 4)

sglang_patch/srt/batch_overlap/operations_strategy.py

@@ -0,0 +1,378 @@
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+
+from sglang.srt.batch_overlap import operations
+from sglang.srt.batch_overlap.operations import Operation
+from sglang.srt.layers.moe.token_dispatcher import DeepEPConfig
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
+
+
+@dataclass
+class OperationsStrategy:
+    operations: List[Operation]
+    deep_gemm_num_sms: Optional[int] = None
+    tbo_delta_stages: Optional[int] = None
+
+    @classmethod
+    def concat(cls, items: List["OperationsStrategy"]) -> "OperationsStrategy":
+        return OperationsStrategy(
+            operations=[x for item in items for x in item.operations],
+            deep_gemm_num_sms=_assert_all_same(
+                [item.deep_gemm_num_sms for item in items]
+            ),
+            tbo_delta_stages=_assert_all_same(
+                [item.tbo_delta_stages for item in items]
+            ),
+        )
+
+    @staticmethod
+    def init_new_tbo(
+        layers: torch.nn.ModuleList,
+        forward_mode: ForwardMode,
+    ) -> "OperationsStrategy":
+        layer_name = layers[0].__class__.__name__
+        if layer_name == "DeepseekV2DecoderLayer":
+            return OperationsStrategy.concat(
+                [
+                    _compute_moe_deepseek_layer_operations_strategy_tbo(
+                        layer, forward_mode
+                    )
+                    for layer in layers
+                ]
+            )
+        elif layer_name == "Qwen3MoeDecoderLayer":
+            return OperationsStrategy.concat(
+                [
+                    _compute_moe_qwen3_layer_operations_strategy_tbo(
+                        layer, forward_mode
+                    )
+                    for layer in layers
+                ]
+            )
+        elif layer_name == "MiMoV2DecoderLayer":
+            return OperationsStrategy.concat(
+                [
+                    _compute_moe_mimov2_layer_operations_strategy_tbo(
+                        layer, forward_mode
+                    )
+                    for layer in layers
+                ]
+            )
+        elif layer_name == "AgnesDecoderLayer":
+            return OperationsStrategy.concat(
+                [
+                    _compute_moe_deepseek_v4_layer_operations_strategy_tbo(
+                        layer, forward_mode
+                    )
+                    for layer in layers
+                ]
+            )
+        else:
+            raise NotImplementedError
+
+
+def _assert_all_same(items: List):
+    assert all(item == items[0] for item in items)
+    return items[0]
+
+
+# -------------------------------- Strategy for DeepSeek ---------------------------------------
+
+
+# TODO can refactor to make it more fancy if we have more complex strategies
+def _compute_moe_deepseek_layer_operations_strategy_tbo(
+    layer: torch.nn.Module,
+    forward_mode: ForwardMode,
+) -> OperationsStrategy:
+    assert layer.is_layer_sparse, "dense layer TBO not yet implemented"
+    if forward_mode == ForwardMode.EXTEND:
+        return _compute_moe_deepseek_blog_prefill(layer)
+    elif (
+        forward_mode == ForwardMode.DECODE or forward_mode == ForwardMode.TARGET_VERIFY
+    ):
+        return _compute_moe_deepseek_blog_decode(layer)
+    else:
+        raise NotImplementedError(f"Unsupported {forward_mode=}")
+
+
+def _compute_moe_deepseek_blog_prefill(layer):
+    device_properties = torch.cuda.get_device_properties(device="cuda")
+    total_num_sms = device_properties.multi_processor_count
+    deep_gemm_num_sms = None
+    if not _is_hip:
+        deep_gemm_num_sms = total_num_sms - DeepEPConfig.get_instance().num_sms
+
+    return OperationsStrategy(
+        deep_gemm_num_sms=deep_gemm_num_sms,
+        tbo_delta_stages=0,
+        operations=[
+            layer.op_comm_prepare_attn,
+            layer.self_attn.op_prepare,
+            layer.self_attn.op_core,
+            layer.op_comm_prepare_mlp,
+            layer.mlp.op_gate,
+            layer.mlp.op_select_experts,
+            layer.mlp.op_dispatch_a,
+            operations.YieldOperation(),
+            layer.mlp.op_dispatch_b,
+            layer.mlp.op_experts,
+            layer.mlp.op_combine_a,
+            operations.YieldOperation(),
+            layer.mlp.op_shared_experts,
+            layer.mlp.op_combine_b,
+            layer.mlp.op_output,
+            layer.op_comm_postprocess_layer,
+        ],
+    )
+
+
+def _compute_moe_deepseek_blog_decode(layer):
+    return OperationsStrategy(
+        deep_gemm_num_sms=None,
+        tbo_delta_stages=2,
+        operations=[
+            layer.op_comm_prepare_attn,
+            layer.self_attn.op_prepare,
+            operations.YieldOperation(),
+            layer.self_attn.op_core,
+            layer.op_comm_prepare_mlp,
+            layer.mlp.op_gate,
+            layer.mlp.op_select_experts,
+            operations.YieldOperation(),
+            layer.mlp.op_dispatch_a,
+            layer.mlp.op_shared_experts,
+            operations.YieldOperation(),
+            layer.mlp.op_dispatch_b,
+            layer.mlp.op_experts,
+            layer.mlp.op_combine_a,
+            operations.YieldOperation(),
+            layer.mlp.op_combine_b,
+            operations.YieldOperation(),
+            layer.mlp.op_output,
+            layer.op_comm_postprocess_layer,
+        ],
+    )
+
+
+# -------------------------------- Strategy for DeepSeek V4 ---------------------------------------
+
+
+# DSV4 prefill TBO (EP / mori path). Cross-layer mHC fusion is disabled under
+# TBO, so each layer is self-contained: attn-side mHC pre+norm -> attn ->
+# ffn-side mHC pre+norm -> MoE (a2a dispatch/combine overlapped) -> mHC post.
+# The MoE ops are reused from self.mlp (DeepseekV2MoE) and decompose
+# forward_deepep; the layer-level op_mhc_* wrap DSV4's hc_pre / hc_post.
+def _compute_moe_deepseek_v4_layer_operations_strategy_tbo(
+    layer: torch.nn.Module,
+    forward_mode: ForwardMode,
+) -> OperationsStrategy:
+    if forward_mode == ForwardMode.EXTEND:
+        return _compute_moe_deepseek_v4_prefill(layer)
+    else:
+        # Decode TBO for DSV4 is not implemented yet (ATOM data: decode TBO
+        # regresses; needs cuda-graph capture work). Prefill-only for now.
+        raise NotImplementedError(
+            f"DeepseekV4 TBO only supports prefill (EXTEND), got {forward_mode=}"
+        )
+
+
+def _compute_moe_deepseek_v4_prefill(layer):
+    from sglang.srt.layers.moe import get_moe_a2a_backend
+
+    if get_moe_a2a_backend().is_none():
+        # Non-EP DP TP-MoE: overlap the DP all_gatherv (gather) + reduce_scatterv
+        # (combine) with the other ubatch's attn+MoE compute (ATOM's DSV4 path).
+        ops = [
+            layer.op_mhc_prepare_attn,
+            layer.self_attn.op_attn,
+            layer.op_mhc_post_attn_pre_mlp,
+            layer.op_gather_a,
+            operations.YieldOperation(),
+            layer.op_gather_b,
+            layer.op_moe,
+            layer.op_combine_a,
+            operations.YieldOperation(),
+            layer.op_combine_b,
+            layer.op_mhc_postprocess,
+        ]
+    else:
+        # EP / mori a2a: reuse DeepseekV2MoE's deepep dispatch/combine ops.
+        ops = [
+            layer.op_mhc_prepare_attn,
+            layer.self_attn.op_attn,
+            layer.op_mhc_post_attn_pre_mlp,
+            layer.mlp.op_gate,
+            layer.mlp.op_select_experts,
+            layer.mlp.op_dispatch_a,
+            operations.YieldOperation(),
+            layer.mlp.op_dispatch_b,
+            layer.mlp.op_experts,
+            layer.mlp.op_combine_a,
+            operations.YieldOperation(),
+            layer.mlp.op_shared_experts,
+            layer.mlp.op_combine_b,
+            layer.mlp.op_output,
+            layer.op_mhc_postprocess,
+        ]
+    return OperationsStrategy(
+        deep_gemm_num_sms=None,
+        tbo_delta_stages=0,
+        operations=ops,
+    )
+
+
+# -------------------------------- Strategy for Qwen3 ---------------------------------------
+
+
+# TODO: unstable, current strategy is almost the same as DeepSeek, keep redundant code here for
+# convenience to adjust strategy
+def _compute_moe_qwen3_layer_operations_strategy_tbo(
+    layer: torch.nn.Module,
+    forward_mode: ForwardMode,
+) -> OperationsStrategy:
+    assert layer.is_layer_sparse, "qwen3 moe only support sparse layers"
+    if forward_mode == ForwardMode.EXTEND:
+        return _compute_moe_qwen3_prefill(layer)
+    elif (
+        forward_mode == ForwardMode.DECODE or forward_mode == ForwardMode.TARGET_VERIFY
+    ):
+        return _compute_moe_qwen3_decode(layer)
+    else:
+        raise NotImplementedError(f"Unsupported {forward_mode=}")
+
+
+def _compute_moe_qwen3_prefill(layer):
+    device_properties = torch.cuda.get_device_properties(device="cuda")
+    total_num_sms = device_properties.multi_processor_count
+    deep_gemm_num_sms = None
+    if not _is_hip:
+        deep_gemm_num_sms = total_num_sms - DeepEPConfig.get_instance().num_sms
+
+    return OperationsStrategy(
+        deep_gemm_num_sms=deep_gemm_num_sms,
+        tbo_delta_stages=0,
+        operations=[
+            layer.op_comm_prepare_attn,
+            layer.self_attn.op_prepare,
+            layer.self_attn.op_core,
+            layer.op_comm_prepare_mlp,
+            layer.mlp.op_gate,
+            layer.mlp.op_select_experts,
+            layer.mlp.op_dispatch_a,
+            operations.YieldOperation(),
+            layer.mlp.op_dispatch_b,
+            layer.mlp.op_experts,
+            layer.mlp.op_combine_a,
+            operations.YieldOperation(),
+            layer.mlp.op_combine_b,
+            layer.mlp.op_output,
+            layer.op_comm_postprocess_layer,
+        ],
+    )
+
+
+def _compute_moe_qwen3_decode(layer):
+    return OperationsStrategy(
+        deep_gemm_num_sms=None,
+        tbo_delta_stages=2,
+        operations=[
+            layer.op_comm_prepare_attn,
+            layer.self_attn.op_prepare,
+            operations.YieldOperation(),
+            layer.self_attn.op_core,
+            layer.op_comm_prepare_mlp,
+            layer.mlp.op_gate,
+            layer.mlp.op_select_experts,
+            operations.YieldOperation(),
+            layer.mlp.op_dispatch_a,
+            operations.YieldOperation(),
+            layer.mlp.op_dispatch_b,
+            layer.mlp.op_experts,
+            layer.mlp.op_combine_a,
+            operations.YieldOperation(),
+            layer.mlp.op_combine_b,
+            layer.mlp.op_output,
+            layer.op_comm_postprocess_layer,
+            operations.YieldOperation(),
+        ],
+    )
+
+
+# -------------------------------- Strategy for MiMoV2DecoderLayer ---------------------------------------
+
+
+# TODO: unstable; current strategy matches DeepSeek for the common operations (MiMoV2 has no op_shared_experts),
+# so we keep this redundant code here for convenience when adjusting the strategy
+def _compute_moe_mimov2_layer_operations_strategy_tbo(
+    layer: torch.nn.Module,
+    forward_mode: ForwardMode,
+) -> OperationsStrategy:
+    assert layer.is_layer_sparse, "MiMoV2DecoderLayer moe only support sparse layers"
+    if forward_mode == ForwardMode.EXTEND:
+        return _compute_moe_mimov2_prefill(layer)
+    elif (
+        forward_mode == ForwardMode.DECODE or forward_mode == ForwardMode.TARGET_VERIFY
+    ):
+        return _compute_moe_mimov2_decode(layer)
+    else:
+        raise NotImplementedError(f"Unsupported {forward_mode=}")
+
+
+def _compute_moe_mimov2_prefill(layer):
+    device_properties = torch.cuda.get_device_properties(device="cuda")
+    total_num_sms = device_properties.multi_processor_count
+    deep_gemm_num_sms = total_num_sms - DeepEPConfig.get_instance().num_sms
+
+    return OperationsStrategy(
+        deep_gemm_num_sms=deep_gemm_num_sms,
+        tbo_delta_stages=0,
+        operations=[
+            layer.op_comm_prepare_attn,
+            layer.self_attn.op_prepare,
+            layer.self_attn.op_core,
+            layer.op_comm_prepare_mlp,
+            layer.mlp.op_gate,
+            layer.mlp.op_select_experts,
+            layer.mlp.op_dispatch_a,
+            operations.YieldOperation(),
+            layer.mlp.op_dispatch_b,
+            layer.mlp.op_experts,
+            layer.mlp.op_combine_a,
+            operations.YieldOperation(),
+            layer.mlp.op_combine_b,
+            layer.mlp.op_output,
+            layer.op_comm_postprocess_layer,
+        ],
+    )
+
+
+def _compute_moe_mimov2_decode(layer):
+    return OperationsStrategy(
+        deep_gemm_num_sms=None,
+        tbo_delta_stages=2,
+        operations=[
+            layer.op_comm_prepare_attn,
+            layer.self_attn.op_prepare,
+            operations.YieldOperation(),
+            layer.self_attn.op_core,
+            layer.op_comm_prepare_mlp,
+            layer.mlp.op_gate,
+            layer.mlp.op_select_experts,
+            operations.YieldOperation(),
+            layer.mlp.op_dispatch_a,
+            operations.YieldOperation(),
+            layer.mlp.op_dispatch_b,
+            layer.mlp.op_experts,
+            layer.mlp.op_combine_a,
+            operations.YieldOperation(),
+            layer.mlp.op_combine_b,
+            layer.mlp.op_output,
+            layer.op_comm_postprocess_layer,
+            operations.YieldOperation(),
+        ],
+    )

sglang_patch/srt/configs/agnes.py

@@ -0,0 +1,110 @@
+import logging
+import os
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+from transformers import PretrainedConfig
+
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+logger = logging.getLogger(__name__)
+
+
+def try_detect_fp4_experts(model_path: str) -> Optional[bool]:
+    """True = mxfp4-packed (U8/I8/F4), False = converted FP8 (F8_E4M3),
+    None when the header isn't readable (HF slug not cached yet, etc.).
+    Caller falls back to user default. Pure read; never mutates env.
+    """
+    from sglang.srt.model_loader.weight_utils import (
+        probe_routed_expert_weight_dtype,
+    )
+    from sglang.srt.utils import find_local_repo_dir
+
+    if os.path.isdir(model_path):
+        local_path = model_path
+    else:
+        local_path = find_local_repo_dir(model_path)
+        if not local_path or not os.path.isdir(local_path):
+            return None
+
+    try:
+        dtype = probe_routed_expert_weight_dtype(local_path)
+    except Exception as e:
+        logger.warning("Failed to probe routed-expert dtype for %s: %s", model_path, e)
+        return None
+    if dtype is None:
+        return None
+    if dtype in ("U8", "I8", "F4"):
+        return True
+    if dtype == "F8_E4M3":
+        return False
+    logger.warning(
+        "Unexpected routed-expert safetensors dtype=%s for DeepSeek V4", dtype
+    )
+    return None
+
+
+@dataclass(kw_only=True)
+class AgnesConfig(PretrainedConfig):
+    architectures: List[str]
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    bos_token_id: int = 0
+    eos_token_id: int = 1
+    ep_size: int = 1
+    first_k_dense_replace: int = 0
+    hidden_act: str = "silu"
+    hidden_size: int = 4096
+    index_head_dim: int = 128
+    index_n_heads: int = 64
+    index_topk: int = 512
+    initializer_range: float = 0.02
+    intermediate_size: int = 2048
+    kv_lora_rank: int = 512
+    max_position_embeddings: int = 65536
+    model_type: str = "agnes"
+    moe_intermediate_size: int = 2048
+    moe_layer_freq: int = 1
+    n_group: int = 8
+    n_routed_experts: int = 256
+    n_shared_experts: int = 1
+    norm_topk_prob: bool = True
+
+    num_attention_heads: int = 64
+    num_experts_per_tok: int = 6
+    num_hidden_layers: int = 43
+    num_key_value_heads: int = 1
+
+    q_lora_rank: int = 1024
+    qk_nope_head_dim: int = 448
+    qk_rope_head_dim: int = 64
+
+    quantization_config: QuantizationConfig = field(default_factory=QuantizationConfig)
+
+    rms_norm_eps: float = 1e-6
+
+    rope_scaling: Dict[str, float] = field(default_factory=dict)
+    rope_theta: int = 10000
+
+    routed_scaling_factor: float = 1.5
+    scoring_func: str = "sqrtsoftplus"
+
+    tie_word_embeddings: bool = False
+
+    topk_group: int = 8
+    topk_method: str = "noaux_tc"
+
+    use_cache: bool = True
+    v_head_dim: int = 512
+    vocab_size: int = 129280
+    o_lora_rank: int = 1024
+    o_groups: int = 8
+    window_size: int = 128
+
+    compress_rope_theta: int = 40000
+    compress_ratios: List[int] = field(default_factory=list)
+
+    n_hash_layers: int = 3
+    hc_mult: int = 4
+    hc_sinkhorn_iters: int = 20
+    hc_eps: float = 1e-6

sglang_patch/srt/configs/model_config.py

@@ -0,0 +1,2055 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import copy
+import json
+import logging
+import math
+import os
+from enum import Enum, IntEnum, auto
+from functools import cached_property
+from pathlib import Path
+from typing import Any, List, Optional, Set, Union
+
+import torch
+from transformers import PretrainedConfig
+
+from sglang.srt.configs.linear_attn_model_registry import get_linear_attn_config
+from sglang.srt.environ import envs
+from sglang.srt.layers.quantization import QUANTIZATION_METHODS
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import is_hip, is_sm100_supported, retry
+from sglang.srt.utils.hf_transformers_utils import (
+    get_config,
+    get_context_length,
+    get_generation_config,
+    get_hf_text_config,
+    get_sparse_attention_config,
+)
+from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
+from sglang.utils import is_in_ci
+
+logger = logging.getLogger(__name__)
+
+MIMO_V2_MODEL_ARCHS = (
+    "MiMoV2ForCausalLM",
+    "MiMoV2FlashForCausalLM",
+)
+MIMO_V2_MULTIMODAL_ARCHS = ("MiMoV2ForCausalLM",)
+
+
+def get_mimo_v2_fused_qkv_expected_tp_size(hf_config):
+    layout = getattr(hf_config, "attention_projection_layout", None)
+    if layout is None:
+        return None
+    if layout != "fused_qkv":
+        raise ValueError(
+            "MiMoV2 hf_config has unsupported "
+            f"attention_projection_layout={layout!r}; expected 'fused_qkv' "
+            "or unset."
+        )
+
+    num_key_value_heads = getattr(hf_config, "num_key_value_heads", None)
+    text_config = getattr(hf_config, "text_config", None)
+    if num_key_value_heads is None and text_config is not None:
+        num_key_value_heads = getattr(text_config, "num_key_value_heads", None)
+    if num_key_value_heads is None:
+        raise ValueError(
+            "MiMoV2 hf_config has attention_projection_layout='fused_qkv' "
+            "but num_key_value_heads is missing; this value is required to "
+            "derive the fused qkv_proj TP size."
+        )
+    return num_key_value_heads
+
+
+class AttentionArch(IntEnum):
+    MLA = auto()
+    MHA = auto()
+
+
+class ModelImpl(str, Enum):
+    AUTO = "auto"
+    SGLANG = "sglang"
+    TRANSFORMERS = "transformers"
+    MINDSPORE = "mindspore"
+
+
+def _hf_arch(config) -> Optional[str]:
+    """First architecture from a HF config dict or PretrainedConfig (or None)."""
+    archs = (
+        config.get("architectures")
+        if isinstance(config, dict)
+        else getattr(config, "architectures", None)
+    )
+    return archs[0] if archs else None
+
+
+def _hf_attr(config, name):
+    """Read an arbitrary field from a HF config dict or PretrainedConfig."""
+    if isinstance(config, dict):
+        return config.get(name)
+    return getattr(config, name, None)
+
+
+def is_deepseek_dsa(config) -> bool:
+    return (
+        _hf_arch(config)
+        in (
+            "DeepseekV3ForCausalLM",
+            "DeepseekV32ForCausalLM",
+            "DeepseekV3ForCausalLMNextN",
+            "MistralLarge3ForCausalLM",
+            "PixtralForConditionalGeneration",
+            "GlmMoeDsaForCausalLM",
+            "GlmMoeDsaForCausalLMNextN",
+            "LongcatFlashForCausalLM",
+            "LongcatFlashForCausalLMNextN",
+        )
+        and _hf_attr(config, "index_topk") is not None
+    )
+
+
+def is_deepseek_v4(config) -> bool:
+    return _hf_arch(config) in (
+        "AgnesForCausalLM",
+        "AgnesForCausalLMNextN",
+        "AgnesForCausalLMDSpark",
+    )
+
+
+def get_dsa_index_head_dim(config: PretrainedConfig) -> int:
+    assert is_deepseek_dsa(config) or is_deepseek_v4(config)
+    return config.index_head_dim
+
+
+def is_minimax_sparse(config: PretrainedConfig) -> bool:
+    arch = (config.architectures or [None])[0]
+    return arch in (
+        "MiniMaxM3SparseForCausalLM",
+        "MiniMaxM3SparseForConditionalGeneration",
+    )
+
+
+def get_minimax_sparse_attention_config(config: PretrainedConfig) -> dict:
+    text_cfg = getattr(config, "text_config", None)
+    cfg = (
+        getattr(text_cfg, "sparse_attention_config", None)
+        if text_cfg is not None
+        else None
+    )
+    if cfg is None:
+        cfg = getattr(config, "sparse_attention_config", None)
+    if cfg is None:
+        raise ValueError("Could not find sparse config. Is it MiniMax M3 Sparse model?")
+    return cfg
+
+
+def get_minimax_sparse_layer_ids(sparse_cfg: dict) -> tuple[list[int], list[int]]:
+    sparse_freq = sparse_cfg["sparse_attention_freq"]
+    dense_layer_ids = [i for i, f in enumerate(sparse_freq) if f == 0]
+    sparse_layer_ids = [i for i, f in enumerate(sparse_freq) if f != 0]
+    return dense_layer_ids, sparse_layer_ids
+
+
+def get_minimax_sparse_disable_value_layer_ids(sparse_cfg: dict) -> list[int]:
+    flags = sparse_cfg.get("sparse_disable_index_value")
+    if flags is None:
+        return []
+    return [i for i, f in enumerate(flags) if f != 0]
+
+
+def get_minimax_sparse_score_type(sparse_cfg: dict) -> str:
+    score_type = sparse_cfg.get("sparse_score_type", "max")
+    assert score_type in (
+        "max",
+        "lse",
+    ), f"sparse_score_type must be 'max' or 'lse', got {score_type!r}"
+    return score_type
+
+
+def get_dsa_index_topk(config: PretrainedConfig) -> int:
+    assert is_deepseek_dsa(config)
+    return config.index_topk
+
+
+def dsa_layer_skips_topk(config: PretrainedConfig, layer_id: int) -> bool:
+    """Return whether a DSA layer reuses the previous layer's top-k indices."""
+    assert is_deepseek_dsa(config)
+
+    pattern = getattr(config, "index_topk_pattern", None)
+    if pattern is not None:
+        return layer_id < len(pattern) and pattern[layer_id] == "S"
+
+    freq = getattr(config, "index_topk_freq", 1)
+    if freq is None:
+        freq = 1
+    assert freq > 0, f"index_topk_freq must be positive, got {freq}"
+    offset = getattr(config, "index_skip_topk_offset", None)
+    if offset is not None:
+        assert offset > 0, (
+            "index_skip_topk_offset must be positive; offset <= 0 "
+            "marks layer 0 as skip_topk with no prior topk to reuse"
+        )
+        return max(layer_id - offset + 1, 0) % freq != 0
+
+    return max(layer_id - 1, 0) % freq != 0
+
+
+def get_dsa_index_n_heads(config: PretrainedConfig) -> int:
+    assert is_deepseek_dsa(config)
+    return config.index_n_heads
+
+
+def get_num_indexer_layers(config) -> int:
+    """Layer count for the global indexer-topk capturer's host buffer.
+
+    DSA models (V3.2) instantiate an Indexer on every transformer layer.
+    With index_topk_freq > 1 some layers reuse prev layer's topk; those still
+    get a slot (mirrored at the MLA call site). DSv4 has C4 indexers only on
+    layers whose compress_ratio == 4. Other architectures: set
+    num_indexer_layers on hf_text_config; 0 disables the capturer.
+    """
+    if is_deepseek_dsa(config):
+        return config.num_hidden_layers
+    if is_deepseek_v4(config):
+        compress_ratios = getattr(config, "compress_ratios", None) or []
+        return sum(1 for r in compress_ratios if r == 4)
+    return getattr(config, "num_indexer_layers", 0)
+
+
+class ModelConfig:
+    def __init__(
+        self,
+        model_path: str,
+        trust_remote_code: bool = True,
+        revision: Optional[str] = None,
+        context_length: Optional[int] = None,
+        model_override_args: str = "{}",
+        is_embedding: Optional[bool] = None,
+        enable_multimodal: Optional[bool] = None,
+        dtype: str = "auto",
+        quantization: Optional[str] = None,
+        override_config_file: Optional[str] = None,
+        is_draft_model: bool = False,
+        model_impl: Union[str, ModelImpl] = ModelImpl.AUTO,
+        sampling_defaults: str = "openai",
+        quantize_and_serve: bool = False,
+        is_multi_layer_eagle: bool = False,
+        encoder_only: bool = False,
+        language_only: bool = False,
+        disable_hybrid_swa_memory: bool = False,
+        model_config_parser: str = "auto",
+        speculative_algorithm: Optional[str] = None,
+    ) -> None:
+        # Parse args
+        self.model_path = model_path
+        self.revision = revision
+        self.quantization = quantization
+        self.is_draft_model = is_draft_model
+        self.speculative_algorithm = speculative_algorithm
+        self.model_impl = model_impl
+        self.sampling_defaults = sampling_defaults
+        self.quantize_and_serve = quantize_and_serve
+        self.is_multi_layer_eagle = is_multi_layer_eagle
+        self.disable_hybrid_swa_memory = disable_hybrid_swa_memory
+        self.model_config_parser = model_config_parser
+
+        # Validate quantize_and_serve configuration
+        self._validate_quantize_and_serve_config()
+
+        # Get hf config
+        self._maybe_pull_model_for_runai(self.model_path)
+        self._maybe_pull_model_tokenizer_from_remote()
+        self.model_override_args = json.loads(model_override_args)
+        kwargs = {}
+        if override_config_file and override_config_file.strip():
+            kwargs["_configuration_file"] = override_config_file.strip()
+        # get_config() is cached. ModelConfig mutates hf_config for draft-model
+        # remapping and architecture-specific normalization, so each instance
+        # must own an isolated copy.
+        self.hf_config = copy.deepcopy(
+            get_config(
+                self.model_path,
+                trust_remote_code=trust_remote_code,
+                revision=revision,
+                model_override_args=self.model_override_args,
+                model_config_parser=model_config_parser,
+                **kwargs,
+            )
+        )
+        self.hf_text_config = get_hf_text_config(self.hf_config)
+
+        rope_scaling = getattr(self.hf_text_config, "rope_parameters", None) or getattr(
+            self.hf_text_config, "rope_scaling", {}
+        )
+        self.model_is_mrope = (
+            rope_scaling is not None and "mrope_section" in rope_scaling
+        )
+
+        self.hf_generation_config = get_generation_config(
+            self.model_path,
+            trust_remote_code=trust_remote_code,
+            revision=revision,
+            **kwargs,
+        )
+
+        # Set enable_multimodal
+        if enable_multimodal is None:
+            mm_disabled_models = [
+                "Gemma3ForConditionalGeneration",
+                "Llama4ForConditionalGeneration",
+                "Step3VLForConditionalGeneration",
+                "InklingForConditionalGeneration",
+            ]
+            if (
+                self.hf_config.architectures[0] in mm_disabled_models
+                and self.model_impl != ModelImpl.TRANSFORMERS
+            ):
+                enable_multimodal = False
+                logger.info(
+                    f"Multimodal is disabled for {self.hf_config.model_type}. To enable it, set --enable-multimodal."
+                )
+            elif self.hf_config.architectures[0] in MIMO_V2_MULTIMODAL_ARCHS and not (
+                hasattr(self.hf_config, "vision_config")
+                and hasattr(self.hf_config, "audio_config")
+            ):
+                enable_multimodal = False
+                logger.info(
+                    "Multimodal is disabled for this MiMoV2 checkpoint: "
+                    "vision_config/audio_config not found in the model config "
+                    "(likely a text-only MiMoV2 variant)."
+                )
+            else:
+                enable_multimodal = True
+
+        # Config draft model
+        self._config_draft_model()
+
+        # DSV4 expert layout: env (default True = mxfp4) applies only to V4.
+        # Other FP8 MoE models (for example DeepSeek V3.2) must keep the normal
+        # FP8 expert tensor layout.
+        self.is_fp4_experts: bool = False
+        if is_deepseek_v4(self.hf_config):
+            self.is_fp4_experts = envs.SGLANG_DSV4_FP4_EXPERTS.get()
+            if (
+                not envs.SGLANG_DSV4_FP4_EXPERTS.is_set()
+                or envs.SGLANG_DSV4_FP4_DEQUANT.is_set()
+            ):
+                from sglang.srt.configs.agnes import try_detect_fp4_experts
+
+                detected = try_detect_fp4_experts(self.model_path)
+                if detected is not None:
+                    self.is_fp4_experts = detected
+                    logger.info(
+                        "Auto-detected DSV4 routed-expert layout: is_fp4_experts=%s",
+                        self.is_fp4_experts,
+                    )
+            if envs.SGLANG_DSV4_FP4_DEQUANT.get():
+                envs.SGLANG_DSV4_FP4_DEQUANT.set(self.is_fp4_experts is not None)
+
+            # HF config.json inherits topk_group=4 from the V3 template, but
+            # DSV4 trains with no group limiting (sqrtsoftplus + full-expert
+            # top-k). Force topk_group == n_group so deepseek_v2.py:531's
+            # `n_group > topk_group` evaluates False and routes to the
+            # ungrouped sqrtsoftplus path. The grouped impl only supports
+            # sigmoid scoring (topk.py:722) and would silently corrupt expert
+            # weights if hit.
+            n_group = getattr(self.hf_config, "n_group", None)
+            if n_group is not None:
+                self.hf_config.topk_group = n_group
+
+        # Handle hybrid NVFP4 moe (nvidia/DeepSeek-V4-Pro-NVFP4)
+        self.nvfp4_moe_meta: Optional[dict] = None
+        hybrid_quant_cfg = getattr(self.hf_config, "quantization_config", None)
+        if hybrid_quant_cfg is not None and not isinstance(hybrid_quant_cfg, dict):
+            hybrid_quant_cfg = hybrid_quant_cfg.to_dict()
+        if (
+            hybrid_quant_cfg is not None
+            and str(hybrid_quant_cfg.get("quant_algo", "")).upper() == "MIXED_PRECISION"
+            and str(hybrid_quant_cfg.get("moe_quant_algo", "")).upper() == "NVFP4"
+            and hybrid_quant_cfg.get("group_size") is not None
+        ):
+            self.nvfp4_moe_meta = {
+                "group_size": int(hybrid_quant_cfg["group_size"]),
+                "exclude_modules": list(hybrid_quant_cfg.get("ignore") or []),
+            }
+            logger.info(
+                "Auto-detected hybrid FP8+NVFP4 checkpoint "
+                "(NVFP4 MoE group_size=%d, %d exclude_modules)",
+                self.nvfp4_moe_meta["group_size"],
+                len(self.nvfp4_moe_meta["exclude_modules"]),
+            )
+
+        # Check model type
+        self.attention_chunk_size = getattr(
+            self.hf_text_config, "attention_chunk_size", None
+        )
+        self.sliding_window_size = self._get_sliding_window_size()
+        self.is_generation = is_generation_model(
+            self.hf_config.architectures, is_embedding
+        )
+        # The vision_config/audio_config attribute heuristic is only applied when
+        # the transformers backend is explicitly requested. Some text-only models
+        # (e.g. xai-org/grok-2 with model_type="git") would otherwise be
+        # false-positively detected because their HF config auto-populates a
+        # `vision_config` in __post_init__.
+        has_multimodal_subconfig = self.hf_config is not self.hf_text_config or (
+            self.model_impl == ModelImpl.TRANSFORMERS
+            and (
+                hasattr(self.hf_config, "vision_config")
+                or hasattr(self.hf_config, "audio_config")
+            )
+        )
+        self.is_multimodal = enable_multimodal and (
+            is_multimodal_model(self.hf_config.architectures)
+            or has_multimodal_subconfig
+        )
+        self.is_audio_model = enable_multimodal and is_audio_model(
+            self.hf_config.architectures
+        )
+        # TODO: requires further polishing
+        self.is_image_understandable_model = enable_multimodal and hasattr(
+            self.hf_config, "vision_config"
+        )
+
+        # Models expose audio_config at different nesting levels:
+        #   - top-level audio_config: e.g. Qwen2Audio
+        #   - thinker_config.audio_config: Qwen3-Omni, Qwen3-ASR (nested thinker arch)
+        #   - sound_config: Nemotron AVLM with Parakeet audio encoder
+        #   - is_audio_model(): Whisper, Qwen3-ASR (architecture-based fallback)
+        # TODO: Handle this more robustly by standardizing the config structure in the future
+        self.is_audio_understandable_model = enable_multimodal and (
+            hasattr(self.hf_config, "audio_config")
+            or hasattr(getattr(self.hf_config, "thinker_config", None), "audio_config")
+            or getattr(self.hf_config, "sound_config", None) is not None
+            or is_audio_model(self.hf_config.architectures)
+        )
+
+        self.is_multimodal_chunked_prefill_supported = (
+            enable_multimodal
+            and is_multimodal_chunked_prefill_supported(self.hf_config.architectures)
+        )
+        self.is_encoder_decoder = is_encoder_decoder_model(self.hf_config.architectures)
+        self.is_local_attention_model = is_local_attention_model(
+            self.hf_config.architectures
+        )
+        self.use_ngram_embedding = getattr(self.hf_config, "use_ngram_embedding", False)
+        # A multimodal arch is piecewise-incompatible until its LM prefill is validated.
+        self.is_piecewise_cuda_graph_disabled_model = (
+            is_piecewise_cuda_graph_disabled_model(self.hf_config.architectures)
+            or (
+                self.is_multimodal
+                and not is_multimodal_piecewise_cuda_graph_supported(
+                    self.hf_config.architectures
+                )
+            )
+        )
+        # Multimodal archs whose language-model prefill is verified safe to capture
+        # under piecewise CUDA graph. ServerArgs otherwise disables prefill piecewise
+        # CG for every multimodal model; this opt-in re-enables it for listed archs
+        # (the vision encoder still runs eagerly via general_mm_embed_routine, only the
+        # LM forward is captured).
+        self.is_multimodal_piecewise_cuda_graph_supported = enable_multimodal and (
+            is_multimodal_piecewise_cuda_graph_supported(self.hf_config.architectures)
+        )
+        self.is_multimodal_breakable_cuda_graph_supported = enable_multimodal and (
+            is_multimodal_breakable_cuda_graph_supported(self.hf_config.architectures)
+        )
+        self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype)
+
+        # Derive context length and model shapes
+        self._derive_context_length(context_length)
+        self._derive_model_shapes()
+
+        # Update hybrid model
+        self._derive_hybrid_model()
+
+        # Verify quantization
+        self._verify_quantization()
+
+        # Verify dual-chunk attention config
+        self._verify_dual_chunk_attention_config()
+
+        # Cache attributes
+        self.hf_eos_token_id = self._get_hf_eos_token_id()
+        # Set by scheduler when reasoning_parser is enabled
+        self.think_end_id: Optional[int] = None
+
+        # multimodal
+        self.image_token_id = getattr(
+            self.hf_config, "image_token_id", None
+        ) or getattr(self.hf_config, "image_token_index", None)
+
+        self.hf_config.encoder_only = encoder_only
+        self.hf_config.language_only = language_only
+
+        # matryoshka embeddings
+        self.matryoshka_dimensions = getattr(
+            self.hf_config, "matryoshka_dimensions", None
+        )
+        self.is_matryoshka = self.matryoshka_dimensions or getattr(
+            self.hf_config, "is_matryoshka", False
+        )
+
+    @staticmethod
+    def from_server_args(
+        server_args: ServerArgs,
+        model_path: str = None,
+        model_revision: str = None,
+        is_draft_model: bool = False,
+        context_length: Optional[int] = None,
+        **kwargs,
+    ):
+        quantization = (
+            server_args.speculative_draft_model_quantization
+            if is_draft_model
+            else server_args.quantization
+        )
+        override_config_file = (
+            server_args.decrypted_draft_config_file
+            if is_draft_model
+            else server_args.decrypted_config_file
+        )
+        return ModelConfig(
+            model_path=model_path or server_args.model_path,
+            trust_remote_code=server_args.trust_remote_code,
+            revision=model_revision or server_args.revision,
+            context_length=(
+                context_length
+                if context_length is not None
+                else server_args.context_length
+            ),
+            model_override_args=server_args.json_model_override_args,
+            is_embedding=server_args.is_embedding,
+            enable_multimodal=server_args.enable_multimodal,
+            dtype=server_args.dtype,
+            quantization=quantization,
+            model_impl=server_args.model_impl,
+            sampling_defaults=server_args.sampling_defaults,
+            quantize_and_serve=server_args.quantize_and_serve,
+            override_config_file=override_config_file,
+            is_multi_layer_eagle=server_args.enable_multi_layer_eagle,
+            language_only=server_args.language_only,
+            encoder_only=server_args.encoder_only,
+            is_draft_model=is_draft_model,
+            disable_hybrid_swa_memory=server_args.disable_hybrid_swa_memory,
+            model_config_parser=server_args.model_config_parser,
+            speculative_algorithm=server_args.speculative_algorithm,
+            **kwargs,
+        )
+
+    def _config_draft_model(self):
+        is_draft_model = self.is_draft_model
+
+        if is_draft_model and self.hf_config.architectures[0] in [
+            "DeepseekV3ForCausalLM",
+            "DeepseekV32ForCausalLM",
+        ]:
+            self.hf_config.architectures[0] = "DeepseekV3ForCausalLMNextN"
+
+        if is_draft_model and self.hf_config.architectures[0] == "GlmMoeDsaForCausalLM":
+            self.hf_config.architectures[0] = "GlmMoeDsaForCausalLMNextN"
+
+        if (
+            is_draft_model
+            and self.hf_config.architectures[0] == "AgnesForCausalLM"
+        ):
+            from sglang.srt.speculative.dspark_components.dspark_config import (
+                checkpoint_bundles_dspark_draft,
+            )
+
+            # A dspark-bundled checkpoint may also carry MTP layers; the
+            # selected algorithm decides which draft arch to load.
+            if checkpoint_bundles_dspark_draft(self.hf_config) and (
+                self.speculative_algorithm in (None, "DSPARK")
+            ):
+                self.hf_config.architectures[0] = "AgnesForCausalLMDSpark"
+                logger.info(
+                    "Draft checkpoint bundles a DSpark head; loading draft arch "
+                    "AgnesForCausalLMDSpark."
+                )
+            else:
+                self.hf_config.architectures[0] = "AgnesForCausalLMNextN"
+                self.hf_config.num_nextn_predict_layers = 1
+
+        if is_draft_model and self.hf_config.architectures[0] == "Glm4MoeForCausalLM":
+            self.hf_config.architectures[0] = "Glm4MoeForCausalLMNextN"
+        if (
+            is_draft_model
+            and self.hf_config.architectures[0] == "Glm4MoeLiteForCausalLM"
+        ):
+            self.hf_config.architectures[0] = "Glm4MoeLiteForCausalLMNextN"
+
+        if is_draft_model and self.hf_config.architectures[0] in [
+            "GlmOcrForConditionalGeneration",
+        ]:
+            self.hf_config.architectures[0] = "GlmOcrForConditionalGenerationNextN"
+
+        if (
+            is_draft_model
+            and self.hf_config.architectures[0] == "LongcatFlashForCausalLM"
+        ):
+            self.hf_config.architectures[0] = "LongcatFlashForCausalLMNextN"
+            self.hf_config.num_hidden_layers = self.hf_config.num_nextn_predict_layers
+
+        if is_draft_model and self.hf_config.architectures[0] == "MiMoForCausalLM":
+            self.hf_config.architectures[0] = "MiMoMTP"
+        if is_draft_model and self.hf_config.architectures[0] in MIMO_V2_MODEL_ARCHS:
+            self.hf_config.architectures[0] = "MiMoV2MTP"
+        if is_draft_model and self.hf_config.architectures[0] == "Step3p5ForCausalLM":
+            self.hf_config.architectures[0] = "Step3p5MTP"
+        if (
+            is_draft_model
+            and self.hf_config.architectures[0] == "InklingForConditionalGeneration"
+        ):
+            self.hf_config.architectures[0] = "InklingForConditionalGenerationMTP"
+        if (
+            is_draft_model
+            and self.hf_config.architectures[0] == "Step3p7ForConditionalGeneration"
+        ):
+            self.hf_config = self.hf_text_config
+            self.hf_config.architectures = ["Step3p5MTP"]
+        if is_draft_model and self.hf_config.architectures[0] in [
+            "BailingMoeV2ForCausalLM",
+            "BailingMoeForCausalLM",
+            "BailingMoeV2_5ForCausalLM",
+        ]:
+            self.hf_config.architectures[0] = "BailingMoeForCausalLMNextN"
+        if (
+            is_draft_model
+            and self.hf_config.architectures[0] == "Ernie4_5_MoeForCausalLM"
+        ):
+            self.hf_config.architectures[0] = "Ernie4_5_MoeForCausalLMMTP"
+
+        if is_draft_model and self.hf_config.architectures[0] == "Qwen3NextForCausalLM":
+            self.hf_config.architectures[0] = "Qwen3NextForCausalLMMTP"
+            self.hf_config.num_nextn_predict_layers = 1
+
+        if is_draft_model and self.hf_config.architectures[0] == "Qwen3MoeForCausalLM":
+            self.hf_config.architectures[0] = "Qwen3MoeForCausalLMMTP"
+            self.hf_config.num_nextn_predict_layers = 1
+
+        if is_draft_model and self.hf_config.architectures[0] in [
+            "Qwen3_5ForConditionalGeneration",
+            "Qwen3_5MoeForConditionalGeneration",
+            "InternS2PreviewForConditionalGeneration",
+        ]:
+            self.hf_config.architectures[0] = "Qwen3_5ForCausalLMMTP"
+            self.hf_config.num_nextn_predict_layers = 1
+
+        if is_draft_model and self.hf_config.architectures[0] == "ExaoneMoEForCausalLM":
+            self.hf_config.architectures[0] = "ExaoneMoEForCausalLMMTP"
+            self.hf_config.num_nextn_predict_layers = 1
+
+        if is_draft_model and self.hf_config.architectures[0] in [
+            "NemotronHForCausalLM",
+            "NemotronHPuzzleForCausalLM",
+        ]:
+            self.hf_config.architectures[0] = "NemotronHForCausalLMMTP"
+            self.hf_config.num_nextn_predict_layers = 1
+
+        if is_draft_model and self.hf_config.architectures[0] == "HYV3ForCausalLM":
+            self.hf_config.architectures[0] = "HYV3ForCausalLMNextN"
+            self.hf_config.num_nextn_predict_layers = 1
+
+    def _derive_hybrid_model(self):
+        # Use self.context_len after it has been initialized to prevent using context_len which may be None.
+        self.is_hybrid_swa = (
+            is_hybrid_swa_model(self.hf_config.architectures, self.hf_text_config)
+            and not self.disable_hybrid_swa_memory
+        )
+
+        if self.is_hybrid_swa:
+            logger.info(f"Hybrid swa model: {self.hf_config.architectures=}")
+
+            self.is_deepseek_v4_arch = any(
+                arch
+                in [
+                    "AgnesForCausalLM",
+                    "AgnesForCausalLMNextN",
+                    "AgnesForCausalLMDSpark",
+                ]
+                for arch in self.hf_config.architectures
+            )
+
+            if not self.is_deepseek_v4_arch:
+                self.swa_attention_layer_ids, self.full_attention_layer_ids = (
+                    get_hybrid_layer_ids(
+                        self.hf_config.architectures,
+                        self.hf_text_config,
+                    )
+                )
+
+        self.has_attention_sinks = self._detect_attention_sinks()
+
+        self.is_hybrid_swa_compress = self.hf_config.architectures[0] in [
+            *MIMO_V2_MODEL_ARCHS,
+            "MiMoV2MTP",
+            "Gemma4ForCausalLM",
+            "Gemma4ForConditionalGeneration",
+            "InklingForConditionalGeneration",
+            "InklingForConditionalGenerationMTP",
+            "Gemma4UnifiedForConditionalGeneration",
+        ]
+
+    @cached_property
+    def linear_attn_registry_result(self) -> Any:
+        return get_linear_attn_config(self.hf_config)
+
+    def _detect_attention_sinks(self) -> bool:
+        """Check whether the model uses learned attention sinks.
+
+        Attention sinks are per-head scalars added to the softmax denominator
+        to compensate for evicted KV-cache entries under sliding-window
+        attention.  Not every hybrid-SWA model uses them.
+        """
+        archs = self.hf_config.architectures or []
+        # GptOss always creates sinks unconditionally.
+        if "GptOssForCausalLM" in archs:
+            return True
+
+        # MiMoV2 creates sinks only when the config flags are set.
+        if any(a in archs for a in (*MIMO_V2_MODEL_ARCHS, "MiMoV2MTP")):
+            return getattr(
+                self.hf_text_config, "add_swa_attention_sink_bias", False
+            ) or getattr(self.hf_text_config, "add_full_attention_sink_bias", False)
+        return False
+
+    def _derive_context_length(self, context_length: int):
+        is_draft_model = self.is_draft_model
+        derived_context_len = get_context_length(self.hf_text_config)
+
+        if context_length is not None:
+            if context_length > derived_context_len:
+                reason = "Target model's" if is_draft_model else "User-specified"
+                msg = (
+                    f"Warning: {reason} context_length ({context_length}) is greater than the derived context_length ({derived_context_len}). "
+                    f"This may lead to incorrect model outputs or CUDA errors. Note that the derived context_length may differ from max_position_embeddings in the model's config."
+                )
+                if (
+                    envs.SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.get()
+                    or is_in_ci()  # FIXME: fix this special case
+                ):
+                    logger.warning(msg)
+                    self.context_len = context_length
+                    if is_draft_model:
+                        self.hf_text_config.max_position_embeddings = context_length
+                        logger.warning(
+                            f"Overriding the draft model's max_position_embeddings to {context_length}."
+                        )
+                else:
+                    raise ValueError(
+                        f"{msg} To allow overriding this maximum, set the env var SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1"
+                    )
+            else:
+                self.context_len = context_length
+        else:
+            self.context_len = derived_context_len
+
+        # Transfer context_len to HuggingFace config so models can access it
+        self.hf_config.context_len = self.context_len
+
+    def _derive_model_shapes(self):
+        # Unify the config keys for hf_text_config
+        self.head_dim = getattr(self.hf_text_config, "head_dim", None)
+        if self.head_dim is None:
+            self.head_dim = (
+                self.hf_text_config.hidden_size
+                // self.hf_text_config.num_attention_heads
+            )
+            setattr(self.hf_text_config, "head_dim", self.head_dim)
+
+        self.v_head_dim = getattr(self.hf_text_config, "v_head_dim", None)
+        if self.v_head_dim is None:
+            self.v_head_dim = self.head_dim
+            setattr(self.hf_text_config, "v_head_dim", self.v_head_dim)
+
+        self.swa_head_dim = getattr(self.hf_text_config, "swa_head_dim", None)
+        if self.swa_head_dim is None:
+            self.swa_head_dim = self.head_dim
+            setattr(self.hf_text_config, "swa_head_dim", self.swa_head_dim)
+
+        self.swa_v_head_dim = getattr(self.hf_text_config, "swa_v_head_dim", None)
+        if self.swa_v_head_dim is None:
+            self.swa_v_head_dim = self.swa_head_dim
+            setattr(self.hf_text_config, "swa_v_head_dim", self.swa_v_head_dim)
+        # FIXME: temporary special judge for MLA architecture
+        if (
+            "DeepseekV2ForCausalLM" in self.hf_config.architectures
+            or "DeepseekV32ForCausalLM" in self.hf_config.architectures
+            or "DeepseekV3ForCausalLM" in self.hf_config.architectures
+            or "DeepseekV3ForCausalLMNextN" in self.hf_config.architectures
+            or "Glm4MoeLiteForCausalLM" in self.hf_config.architectures
+            or "Glm4MoeLiteForCausalLMNextN" in self.hf_config.architectures
+            or "GlmMoeDsaForCausalLM" in self.hf_config.architectures
+            or "GlmMoeDsaForCausalLMNextN" in self.hf_config.architectures
+            or "LongcatFlashForCausalLM" in self.hf_config.architectures
+            or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures
+            or "DotsVLMForCausalLM" in self.hf_config.architectures
+            or "MistralLarge3ForCausalLM" in self.hf_config.architectures
+            or (
+                "PixtralForConditionalGeneration" in self.hf_config.architectures
+                and getattr(self.hf_text_config, "kv_lora_rank", None) is not None
+            )
+            or "MistralLarge3ForCausalLMEagle" in self.hf_config.architectures
+            or "KimiK25ForConditionalGeneration" in self.hf_config.architectures
+            or "Eagle3DeepseekV2ForCausalLM" in self.hf_config.architectures
+        ):
+            self.head_dim = 256
+            self.attention_arch = AttentionArch.MLA
+            self.kv_lora_rank = self.hf_text_config.kv_lora_rank
+            self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
+            self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
+            self.v_head_dim = self.hf_text_config.v_head_dim
+            self.index_head_dim = (
+                get_dsa_index_head_dim(self.hf_text_config)
+                if is_deepseek_dsa(self.hf_text_config)
+                else None
+            )
+            # Handle rope scaling
+            self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim)
+            # in transformers v5, rope_scaling is just rope_parameters for backward compatibility
+            rope_scaling = self.hf_text_config.rope_scaling
+            if rope_scaling:
+                # v5 uses "rope_type", v4 uses "type"
+                rope_type = (
+                    rope_scaling.get("rope_type")
+                    or rope_scaling.get("type")
+                    or "default"
+                )
+                if rope_type != "default":
+                    self.scaling = compute_mla_mscale_scaling(
+                        rope_scaling, self.scaling
+                    )
+        elif (
+            "AgnesForCausalLM" in self.hf_config.architectures
+            or "AgnesForCausalLMNextN" in self.hf_config.architectures
+            or "AgnesForCausalLMDSpark" in self.hf_config.architectures
+        ):
+            self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim
+            self.qk_nope_head_dim = self.hf_config.head_dim - self.qk_rope_head_dim
+            self.window_size = self.hf_config.sliding_window
+            self.head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+            self.v_head_dim = self.head_dim
+            self.index_head_dim = self.hf_config.index_head_dim
+            self.compress_ratios = self.hf_config.compress_ratios
+            self.attention_arch = AttentionArch.MHA
+            self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim)
+            if self.hf_config.rope_scaling:
+                self.scaling = compute_mla_mscale_scaling(
+                    self.hf_config.rope_scaling, self.scaling
+                )
+        elif "Glm4MoeForCausalLMNextN" in self.hf_config.architectures:
+            if self.head_dim is None:
+                self.head_dim = (
+                    self.hf_text_config.hidden_size
+                    // self.hf_text_config.num_attention_heads
+                )
+            if self.swa_head_dim is None:
+                self.swa_head_dim = self.head_dim
+            self.v_head_dim = self.head_dim
+            self.swa_v_head_dim = self.swa_head_dim
+            self.attention_arch = AttentionArch.MHA
+        elif "MiniCPM3ForCausalLM" in self.hf_config.architectures:
+            self.head_dim = 128
+            self.attention_arch = AttentionArch.MLA
+            self.kv_lora_rank = self.hf_config.kv_lora_rank
+            self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim
+        elif "DeepseekVL2ForCausalLM" in self.hf_config.architectures and getattr(
+            self.hf_text_config, "use_mla", True
+        ):
+            self.head_dim = 256
+            self.attention_arch = AttentionArch.MLA
+            self.kv_lora_rank = self.hf_text_config.kv_lora_rank
+            self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
+            self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
+        elif "KimiVLForConditionalGeneration" in self.hf_config.architectures:
+            self.head_dim = 256
+            self.attention_arch = AttentionArch.MLA
+            self.kv_lora_rank = self.hf_text_config.kv_lora_rank
+            self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
+            self.v_head_dim = self.hf_text_config.v_head_dim
+            self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
+        elif "KimiLinearForCausalLM" in self.hf_config.architectures:
+            self.head_dim = 72
+            self.attention_arch = AttentionArch.MLA
+            self.kv_lora_rank = self.hf_config.kv_lora_rank
+            self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim
+            self.v_head_dim = self.hf_config.v_head_dim
+            self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim
+            self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim)
+            if self.hf_config.rope_scaling:
+                self.scaling = compute_mla_mscale_scaling(
+                    self.hf_config.rope_scaling, self.scaling
+                )
+        elif (
+            "BailingMoeV2_5ForCausalLM" in self.hf_config.architectures
+            or "BailingMoeForCausalLMNextN" in self.hf_config.architectures
+        ):
+            self.head_dim = self.hf_text_config.head_dim
+            self.attention_arch = AttentionArch.MLA
+            self.kv_lora_rank = self.hf_text_config.kv_lora_rank
+            self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
+            self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
+            self.v_head_dim = self.hf_config.v_head_dim
+            # Handle rope scaling with yarn
+            self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim)
+            if self.hf_config.rope_scaling:
+                self.scaling = compute_mla_mscale_scaling(
+                    self.hf_config.rope_scaling, self.scaling
+                )
+        elif "SarvamMLAForCausalLM" in self.hf_config.architectures:
+            self.head_dim = (
+                self.hf_config.qk_nope_head_dim + self.hf_config.qk_rope_head_dim
+            )
+            self.attention_arch = AttentionArch.MLA
+            self.kv_lora_rank = self.hf_config.kv_lora_rank
+            self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim
+            self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim
+            self.v_head_dim = self.hf_config.v_head_dim
+            self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim)
+            if self.hf_config.rope_scaling:
+                self.scaling = compute_mla_mscale_scaling(
+                    self.hf_config.rope_scaling, self.scaling
+                )
+        else:
+            if (
+                "MistralModel" in self.hf_config.architectures
+                or "MixtralForCausalLM" in self.hf_config.architectures
+                or "MistralForCausalLM" in self.hf_config.architectures
+            ):
+                if getattr(self, "head_dim", None) is None:
+                    self.head_dim = (
+                        self.hf_config.hidden_size // self.hf_config.num_attention_heads
+                    )
+                    # In transformers==4.52.3, the head_dim is null in MistralConfig
+                    if (
+                        not hasattr(self.hf_text_config, "head_dim")
+                        or self.hf_text_config.head_dim is None
+                    ):
+                        setattr(self.hf_text_config, "head_dim", self.head_dim)
+
+            elif "BaichuanForCausalLM" in self.hf_config.architectures:
+                self.use_alibi = self.hf_config.hidden_size != 4096
+
+            self.attention_arch = AttentionArch.MHA
+
+        self.num_attention_heads = self.hf_text_config.num_attention_heads
+        self.num_key_value_heads = getattr(
+            self.hf_text_config, "num_key_value_heads", None
+        )
+        self.first_k_dense_replace = getattr(
+            self.hf_text_config, "first_k_dense_replace", None
+        )
+        self.full_attention_interval = getattr(
+            self.hf_text_config, "full_attention_interval", None
+        )
+
+        # for Dbrx and MPT models
+        if self.hf_config.model_type in ["dbrx", "mpt"]:
+            self.num_key_value_heads = getattr(
+                self.hf_config.attn_config, "kv_n_heads", None
+            )
+
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+        self.hidden_size = self.hf_text_config.hidden_size
+        hc_mult = getattr(self.hf_text_config, "hc_mult", 1)
+        self.spec_hidden_size = (
+            self.hidden_size * hc_mult if hc_mult > 1 else self.hidden_size
+        )
+        # mHC-flattened hidden size; None when not running an mHC model
+        # (e.g. non-DeepSeek-V4 configs without ``hc_mult``).
+        self.hc_hidden_size = self.spec_hidden_size if hc_mult > 1 else None
+        self.num_hidden_layers = self.hf_text_config.num_hidden_layers
+        self.num_attention_layers = self.num_hidden_layers
+        if "LongcatFlashForCausalLM" in self.hf_config.architectures:
+            self.num_attention_layers = self.num_hidden_layers * 2
+        if "IQuestLoopCoderForCausalLM" in self.hf_config.architectures:
+            loop_num = getattr(self.hf_text_config, "loop_num", 1)
+            self.num_attention_layers = int(self.num_hidden_layers * int(loop_num))
+        if "WhisperForConditionalGeneration" in self.hf_config.architectures:
+            # Whisper has unique layer ID scheme:
+            # - Encoder self-attention: 0 to encoder_layers-1 (no KV cache)
+            # - Decoder self-attention: encoder_layers to encoder_layers+decoder_layers-1 (uses KV cache)
+            # - Decoder cross-attention: encoder_layers+decoder_layers to encoder_layers+2*decoder_layers-1
+            # Even though cross-attention doesn't save KV cache, attention backend needs buffer to exist
+            encoder_layers = getattr(self.hf_text_config, "encoder_layers", 0)
+            decoder_layers = getattr(
+                self.hf_text_config, "decoder_layers", self.num_hidden_layers
+            )
+            self.num_attention_layers = encoder_layers + 2 * decoder_layers
+        self.num_nextn_predict_layers = getattr(
+            self.hf_text_config, "num_nextn_predict_layers", None
+        )
+        self.vocab_size = self.hf_text_config.vocab_size
+        # GLM-Image is the only model here whose output head predicts vision tokens.
+        # Use vision_vocab_size for lm_head, LogitsProcessor, and graph-mode logits buffers.
+        if _hf_arch(self.hf_config) == "GlmImageForConditionalGeneration":
+            self.vocab_size = self.hf_text_config.vision_vocab_size
+
+    def get_total_num_attention_heads(self) -> int:
+        return self.num_attention_heads
+
+    def get_num_attention_heads(self, tensor_parallel_size) -> int:
+        total_num_attention_heads = self.num_attention_heads
+        return max(1, total_num_attention_heads // tensor_parallel_size)
+
+    # adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py#L289
+    def get_total_num_kv_heads(self) -> int:
+        """Returns the total number of KV heads."""
+        # For GPTBigCode & Falcon:
+        # NOTE: for falcon, when new_decoder_architecture is True, the
+        # multi_query flag is ignored and we use n_head_kv for the number of
+        # KV heads.
+        falcon_model_types = ["falcon", "RefinedWeb", "RefinedWebModel"]
+        new_decoder_arch_falcon = (
+            self.hf_config.model_type in falcon_model_types
+            and getattr(self.hf_config, "new_decoder_architecture", False)
+        )
+        if not new_decoder_arch_falcon and getattr(
+            self.hf_text_config, "multi_query", False
+        ):
+            # Multi-query attention, only one KV head.
+            # Currently, tensor parallelism is not supported in this case.
+            return 1
+
+        # For DBRX and MPT
+        if self.hf_config.model_type in ["mpt"]:
+            if "kv_n_heads" in self.hf_config.attn_config:
+                return self.hf_config.attn_config["kv_n_heads"]
+            return self.hf_config.num_attention_heads
+        if self.hf_config.model_type in ["dbrx"]:
+            return getattr(
+                self.hf_config.attn_config,
+                "kv_n_heads",
+                self.hf_config.num_attention_heads,
+            )
+        if self.hf_config.model_type in ["nemotron-nas"]:
+            nkvh = {
+                self.hf_config.num_attention_heads // block.attention.n_heads_in_group
+                for block in self.hf_config.block_configs
+                if not block.attention.no_op
+            }
+            if len(nkvh) == 0:
+                raise RuntimeError("Couldn't determine number of kv heads")
+            if len(nkvh) > 1:
+                raise ValueError(
+                    "Variable GQA (VGQA) is not yet supported for nemotron-nas in sglang"
+                )
+            return next(iter(nkvh))
+
+        attributes = [
+            # For Falcon:
+            "n_head_kv",
+            "num_kv_heads",
+            # For LLaMA-2:
+            "num_key_value_heads",
+            # For ChatGLM:
+            "multi_query_group_num",
+            # For Step3
+            "num_attention_groups",
+        ]
+        for attr in attributes:
+            num_kv_heads = getattr(self.hf_text_config, attr, None)
+            if num_kv_heads is not None:
+                return num_kv_heads
+
+        # For non-grouped-query attention models, the number of KV heads is
+        # equal to the number of attention heads.
+        return self.hf_text_config.num_attention_heads
+
+    def get_num_kv_heads(self, tensor_parallel_size) -> int:
+        """Returns the number of KV heads per GPU."""
+        total_num_kv_heads = self.get_total_num_kv_heads()
+        # If tensor parallelism is used, we divide the number of KV heads by
+        # the tensor parallel size. We will replicate the KV heads in the
+        # case where the number of KV heads is smaller than the tensor
+        # parallel size so each GPU has at least one KV head.
+        return max(1, total_num_kv_heads // tensor_parallel_size)
+
+    def get_swa_num_kv_heads(self, tensor_parallel_size) -> int:
+        """Similar to get_num_kv_heads(), but for SWA."""
+        if hasattr(self.hf_text_config, "swa_num_key_value_heads"):
+            total_num_kv_heads = self.hf_text_config.swa_num_key_value_heads
+            return max(1, total_num_kv_heads // tensor_parallel_size)
+        elif hasattr(self.hf_text_config, "attention_other_setting"):  # For step3p5
+            total_num_kv_heads = self.hf_text_config.attention_other_setting.get(
+                "num_attention_groups"
+            )
+            return max(1, total_num_kv_heads // tensor_parallel_size)
+        else:
+            return self.get_num_kv_heads(tensor_parallel_size)
+
+    # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py
+    def _parse_quant_hf_config(self):
+        quant_cfg = getattr(self.hf_config, "quantization_config", None)
+        if quant_cfg is not None and not isinstance(quant_cfg, dict):
+            quant_cfg = quant_cfg.to_dict()
+        if quant_cfg is not None:
+            # Identify modelopt quantization
+            if (
+                "quant_method" not in quant_cfg
+                or quant_cfg["quant_method"] == "modelopt"
+            ):
+                parsed_cfg = self._parse_modelopt_quant_config(
+                    {"quantization": quant_cfg}
+                )
+                if parsed_cfg:
+                    quant_cfg.update(parsed_cfg)
+
+        if quant_cfg is None:
+            # compressed-tensors uses a "compression_config" key
+            quant_cfg = getattr(self.hf_config, "compression_config", None)
+        if quant_cfg is None:
+            # check if is modelopt or mixed-precision model -- Both of them don't have corresponding field
+            # in hf `config.json` but has a standalone `hf_quant_config.json` in the root directory
+            # example: https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-FP8/tree/main
+            # example: https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/tree/main
+            is_local = os.path.exists(self.model_path)
+            if not is_local:
+                # Conditional import based on SGLANG_USE_MODELSCOPE environment variable
+                if envs.SGLANG_USE_MODELSCOPE.get():
+
+                    from modelscope import HubApi, model_file_download
+
+                    hf_api = HubApi()
+                else:
+                    import huggingface_hub
+                    from huggingface_hub import HfApi, hf_hub_download
+
+                    hf_api = HfApi()
+                try:
+                    # In offline mode, skip file_exists check to avoid OfflineModeIsEnabled error
+                    # Instead, directly try to download/read from cache with local_files_only
+                    file_exists = False  # Initialize to avoid UnboundLocalError
+                    if not huggingface_hub.constants.HF_HUB_OFFLINE:
+                        # Online mode: check if file exists before attempting download (optimization)
+                        file_exists = retry(
+                            lambda: hf_api.file_exists(
+                                self.model_path, "hf_quant_config.json"
+                            ),
+                            max_retry=2,
+                            initial_delay=1.0,
+                            max_delay=5.0,
+                        )
+                        if not file_exists:
+                            # File doesn't exist on hub, no need to try downloading
+                            return quant_cfg  # None
+
+                    # Download (online mode) or read from cache (offline mode)
+                    if envs.SGLANG_USE_MODELSCOPE.get():
+                        quant_config_file = model_file_download(
+                            model_id=self.model_path,
+                            file_path="hf_quant_config.json",
+                            revision=self.revision,
+                        )
+                    else:
+                        quant_config_file = hf_hub_download(
+                            repo_id=self.model_path,
+                            filename="hf_quant_config.json",
+                            revision=self.revision,
+                            local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+                        )
+                    with open(quant_config_file) as f:
+                        quant_config_dict = json.load(f)
+                    quant_cfg = self._parse_modelopt_quant_config(quant_config_dict)
+                except huggingface_hub.errors.LocalEntryNotFoundError:
+                    # Offline mode and file not in cache - this is normal for non-quantized models
+                    logger.debug(
+                        f"hf_quant_config.json not found in cache for {self.model_path} "
+                        "(offline mode, normal for non-quantized models)"
+                    )
+                except huggingface_hub.errors.OfflineModeIsEnabled:
+                    # Should not reach here after our changes, but keep for safety
+                    logger.warning(
+                        "Offline mode is enabled, skipping hf_quant_config.json check"
+                    )
+                except Exception as e:
+                    logger.warning(
+                        "Failed to load hf_quant_config.json for model %s: %s",
+                        self.model_path,
+                        e,
+                    )
+            elif os.path.exists(os.path.join(self.model_path, "hf_quant_config.json")):
+                quant_config_file = os.path.join(
+                    self.model_path, "hf_quant_config.json"
+                )
+                with open(quant_config_file) as f:
+                    quant_config_dict = json.load(f)
+                quant_cfg = self._parse_modelopt_quant_config(quant_config_dict)
+        return quant_cfg
+
+    def _find_quant_modelslim_config(self):
+        if self.is_draft_model:
+            return None
+        quant_config_file = Path(self.model_path, "quant_model_description.json")
+        quant_cfg = None
+        if quant_config_file.is_file():
+            with open(quant_config_file) as f:
+                quant_cfg = json.load(f)
+            # This field is required for flagless model loading but is not present in
+            # modelslim model description, so we're adding it here manually.
+            quant_cfg["quant_method"] = "modelslim"
+
+        return quant_cfg
+
+    def _parse_modelopt_quant_config(self, quant_config_dict: dict) -> Optional[dict]:
+        """Parse ModelOpt quantization config and return the appropriate quant_method."""
+        json_quant_configs = quant_config_dict["quantization"]
+        quant_algo = json_quant_configs.get("quant_algo", None)
+
+        if quant_algo == "MIXED_PRECISION":
+            quantized_layers = json_quant_configs.get("quantized_layers") or {}
+            has_modelopt_nvfp4_layers = any(
+                str(layer_info.get("quant_algo", "")).upper()
+                in ("NVFP4", "W4A16_NVFP4")
+                for layer_info in quantized_layers.values()
+                if isinstance(layer_info, dict)
+            )
+            if has_modelopt_nvfp4_layers:
+                return {"quant_method": "modelopt_mixed", "quant_algo": quant_algo}
+            return {"quant_method": "w4afp8", "quant_algo": quant_algo}
+        elif quant_algo and ("FP4" in quant_algo or "NVFP4" in quant_algo):
+            return {"quant_method": "modelopt_fp4", "quant_algo": quant_algo}
+        elif quant_algo and "FP8" in quant_algo:
+            return {"quant_method": "modelopt_fp8", "quant_algo": quant_algo}
+        else:
+            return None
+
+    def get_quantization_config_log_str(self) -> Optional[str]:
+        """
+        Get a concise string representation of the quantization config for logging.
+        Returns something like "quant=fp8, fmt=e4m3" or "quant=gptq, bits=4".
+        """
+        try:
+            quant_cfg = self._parse_quant_hf_config()
+            if not quant_cfg:
+                return None
+
+            quant_method = quant_cfg.get("quant_method", "quantized")
+            log_str = f"quant={quant_method}"
+
+            # Append interesting fields if they exist
+            for field in ["bits", "quant_algo", "fmt"]:
+                if field in quant_cfg:
+                    log_str += f", {field}={quant_cfg[field]}"
+
+            return log_str
+        except Exception:
+            return None
+
+    def _is_already_quantized(self) -> bool:
+        """Check if the model is already quantized based on config files."""
+        # Check for quantization in hf_config (config.json)
+        if getattr(self.hf_config, "quantization_config", None) or getattr(
+            self.hf_config, "compression_config", None
+        ):
+            return True
+
+        # Check for HuggingFace quantization config
+        quant_cfg = getattr(self.hf_config, "quantization_config", None)
+        if quant_cfg is None:
+            from sglang.srt.utils import has_hf_quant_config
+
+            return has_hf_quant_config(self.model_path)
+        return True
+
+    def _get_modelopt_quant_type(self) -> str:
+        """Extract ModelOpt quantization type from unified quantization flag."""
+        if self.quantization == "modelopt_fp8":
+            return "fp8"
+        elif self.quantization == "modelopt_fp4":
+            return "nvfp4"
+        elif self.quantization == "modelopt_mixed":
+            raise ValueError(
+                "modelopt_mixed is only supported for pre-quantized checkpoints."
+            )
+        elif self.quantization == "modelopt":
+            # Auto-detect from model config
+            quant_cfg = self._parse_quant_hf_config()
+            if quant_cfg:
+                quant_method = quant_cfg.get("quant_method", "").lower()
+                if "fp4" in quant_method:
+                    return "fp4"
+                elif "fp8" in quant_method:
+                    return "fp8"
+            # Default to fp8 if can't detect
+            return "fp8"
+        else:
+            return "fp8"  # Default fallback
+
+    def _get_sliding_window_size(self) -> Optional[int]:
+        for key in ("sliding_window_size", "sliding_window", "window_size"):
+            value = getattr(self.hf_text_config, key, None)
+            if value is not None:
+                return value
+        return None
+
+    def _validate_quantize_and_serve_config(self):
+        """Validate quantize_and_serve configuration."""
+        if not self.quantize_and_serve:
+            return
+
+        # Check if ModelOpt quantization is specified
+        _MODELOPT_QUANTIZATION_METHODS = [
+            "modelopt",
+            "modelopt_fp8",
+            "modelopt_fp4",
+            "nvfp4_online",
+            "modelopt_mixed",
+        ]
+        modelopt_quantization_specified = (
+            self.quantization in _MODELOPT_QUANTIZATION_METHODS
+        )
+
+        if not modelopt_quantization_specified:
+            raise ValueError(
+                "quantize_and_serve requires ModelOpt quantization (set with --quantization "
+                f"{{{', '.join(sorted(_MODELOPT_QUANTIZATION_METHODS))}}})"
+            )
+
+        # quantize_and_serve is disabled due to compatibility issues
+        raise NotImplementedError(
+            "quantize_and_serve functionality is currently disabled due to compatibility issues. "
+            "Please use the separate quantize-then-deploy workflow instead. "
+            "Step 1: Quantize and export model. "
+            "Step 2: Deploy the exported model."
+        )
+
+    # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py
+    def _verify_quantization(self) -> None:
+        supported_quantization = [*QUANTIZATION_METHODS]
+        rocm_supported_quantization = [
+            "awq",
+            "gptq",
+            "fp8",
+            "compressed_tensors",
+            "compressed-tensors",
+            "w8a8_fp8",
+            "petit_nvfp4",
+            "quark",
+            "mxfp4",
+            "mxfp8",
+            "auto-round",
+            "auto-round-int8",
+            "quark_int4fp8_moe",
+            "quark_mxfp4",
+        ]
+        optimized_quantization_methods = [
+            "fp8",
+            "marlin",
+            "modelopt_fp8",
+            "modelopt_fp4",
+            "modelopt_mixed",
+            "nvfp4_online",
+            "gptq_marlin_24",
+            "gptq_marlin",
+            "awq_marlin",
+            "compressed_tensors",
+            "compressed-tensors",
+            "experts_int8",
+            "w8a8_int8",
+            "w8a8_fp8",
+            "moe_wna16",
+            "w4afp8",
+            "petit_nvfp4",
+            "quark",
+            "modelslim",
+            "humming",
+            "quark_mxfp4",
+        ]
+        compatible_quantization_methods = {
+            "modelopt_fp8": ["modelopt"],
+            "modelopt_fp4": ["modelopt"],
+            "modelopt_mixed": ["modelopt"],
+            "nvfp4_online": ["fp8"],
+            "petit_nvfp4": ["modelopt"],
+            "w8a8_int8": ["compressed-tensors", "compressed_tensors"],
+            "w8a8_fp8": ["compressed-tensors", "compressed_tensors"],
+            "auto-round-int8": ["compressed-tensors", "compressed_tensors"],
+        }
+        if self.quantization is not None:
+            self.quantization = self.quantization.lower()
+
+        # Parse quantization method from the HF and ModelSlim model config, if available.
+        # Only one function should return config, other should return None.
+        cfg_list = []
+        hf_config = self._parse_quant_hf_config()
+        modelslim_config = self._find_quant_modelslim_config()
+        quant_config = modelslim_config or hf_config
+        if quant_config is not None:
+            cfg_list.append(quant_config)
+
+        # Filter out None values
+        cfg_list = [item for item in cfg_list if item is not None]
+        if len(cfg_list) > 1:
+            raise ValueError(
+                "Config list contains configs from 2 methods, must be only 1"
+            )
+        quant_cfg = cfg_list[0] if cfg_list else None
+
+        if quant_cfg is not None:
+            quant_method = quant_cfg.get(
+                "quant_method", "" if not self.quantization else self.quantization
+            ).lower()
+
+            # Detect which checkpoint is it
+            for _, method in QUANTIZATION_METHODS.items():
+                quantization_override = method.override_quantization_method(
+                    quant_cfg, self.quantization
+                )
+                if quantization_override:
+                    quant_method = quantization_override
+                    self.quantization = quantization_override
+                    break
+
+            # Verify quantization configurations.
+            if self.quantization is None:
+                self.quantization = quant_method
+            elif self.quantization != quant_method:
+                # Check if the CLI-specified quantization is compatible with HF config's quant_method
+                is_compatible = (
+                    self.quantization in compatible_quantization_methods
+                    and quant_method
+                    in compatible_quantization_methods[self.quantization]
+                )
+                if is_compatible:
+                    # Keep the CLI-specified quantization (e.g., modelopt_fp4) even if
+                    # HF config says "modelopt" - they are compatible
+                    logger.info(
+                        f"Using CLI-specified quantization ({self.quantization}) which is "
+                        f"compatible with HF config quant_method ({quant_method})."
+                    )
+                elif self.is_draft_model:
+                    # Allow auto-detection of quantization from checkpoint for draft model
+                    # only if the CLI quantization is not compatible
+                    logger.info(
+                        f"Draft model quantization ({quant_method}) differs from "
+                        f"main model quantization ({self.quantization}). "
+                        f"Using draft model's detected quantization: {quant_method}"
+                    )
+                    self.quantization = quant_method
+                else:
+                    raise ValueError(
+                        "Quantization method specified in the model config "
+                        f"({quant_method}) does not match the quantization "
+                        f"method specified in the `quantization` argument "
+                        f"({self.quantization})."
+                    )
+
+            # Warn if DeepGemm is enabled for a non-ue8m0 checkpoint on Blackwell.
+            # MXFP8 stores E8M0 block scales that DeepGemm consumes losslessly, so skip the warning there.
+            self.use_scale_ue8m0 = quant_cfg.get("scale_fmt", None) == "ue8m0"
+            from sglang.srt.layers import deep_gemm_wrapper
+
+            if (
+                not self.use_scale_ue8m0
+                and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
+                and self.quantization != "mxfp8"
+            ):
+                logger.warning(
+                    "DeepGemm is enabled but the scale_fmt of checkpoint is not ue8m0. This might cause accuracy degradation on Blackwell."
+                )
+
+        if self.quantization is not None:
+            if self.quantization not in supported_quantization:
+                raise ValueError(
+                    f"Unknown quantization method: {self.quantization}. Must "
+                    f"be one of {supported_quantization}."
+                )
+            if is_hip() and self.quantization not in rocm_supported_quantization:
+                raise ValueError(
+                    f"{self.quantization} quantization is currently not "
+                    f"supported in ROCm."
+                )
+            if self.quantization not in optimized_quantization_methods:
+                # Don't warn for MXFP4/MXFP8 on SM100 since they have optimized kernels
+                if not (
+                    self.quantization in ["mxfp4", "mxfp8"] and is_sm100_supported()
+                ):
+                    logger.warning(
+                        "%s quantization is not fully "
+                        "optimized yet. The speed can be slower than "
+                        "non-quantized models.",
+                        self.quantization,
+                    )
+
+    def _verify_dual_chunk_attention_config(self) -> None:
+        if hasattr(self.hf_config, "dual_chunk_attention_config"):
+            # Try loading the sparse attention config
+            sparse_attn_config = get_sparse_attention_config(self.model_path)
+            if not sparse_attn_config:
+                return
+            self.hf_config.dual_chunk_attention_config["sparse_attention_config"] = (
+                sparse_attn_config
+            )
+            if (
+                "sparse_attention_enabled"
+                not in self.hf_config.dual_chunk_attention_config
+            ):
+                self.hf_config.dual_chunk_attention_config[
+                    "sparse_attention_enabled"
+                ] = True
+
+    def _get_hf_eos_token_id(self) -> Optional[Set[int]]:
+        eos_ids = getattr(self.hf_config, "eos_token_id", None)
+        if eos_ids is not None:
+            # it can be either int or list of int
+            eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids)
+        if eos_ids is None:
+            eos_ids = set()
+        if self.hf_generation_config:
+            generation_eos_ids = getattr(
+                self.hf_generation_config, "eos_token_id", None
+            )
+            if generation_eos_ids:
+                generation_eos_ids = (
+                    {generation_eos_ids}
+                    if isinstance(generation_eos_ids, int)
+                    else set(generation_eos_ids)
+                )
+                eos_ids = eos_ids | generation_eos_ids
+        return eos_ids
+
+    def get_default_sampling_params(self) -> dict[str, Any]:
+        """
+        Get default sampling parameters from the model's generation config.
+
+        This method returns non-default sampling parameters from the model's
+        generation_config.json when sampling_defaults is set to "model".
+
+        Returns:
+            A dictionary containing the non-default sampling parameters.
+        """
+        if self.sampling_defaults != "model":
+            return {}
+
+        if self.hf_generation_config is None:
+            return {}
+
+        config = self.hf_generation_config.to_dict()
+
+        available_params = [
+            "repetition_penalty",
+            "temperature",
+            "top_k",
+            "top_p",
+            "min_p",
+        ]
+
+        default_sampling_params = {
+            p: config.get(p) for p in available_params if config.get(p) is not None
+        }
+
+        return default_sampling_params
+
+    def _maybe_pull_model_for_runai(self, model: str) -> None:
+        if is_runai_obj_uri(model):
+            # local path for loading the config
+            self.model_path = ObjectStorageModel.get_path(model)
+            # remote path for loading the weights
+            self.model_weights = model
+
+    def _maybe_pull_model_tokenizer_from_remote(self) -> None:
+        """
+        Pull the model config files to a temporary
+        directory in case of remote.
+
+        Args:
+            model: The model name or path.
+
+        """
+        from sglang.srt.connector import create_remote_connector
+        from sglang.srt.utils import is_remote_url
+
+        if is_remote_url(self.model_path):
+            logger.info("Pulling model configs from remote...")
+            # BaseConnector implements __del__() to clean up the local dir.
+            # Since config files need to exist all the time, so we DO NOT use
+            # with statement to avoid closing the client.
+            client = create_remote_connector(self.model_path)
+            if is_remote_url(self.model_path):
+                client.pull_files(allow_pattern=["*config.json"])
+                self.model_weights = self.model_path
+                self.model_path = client.get_local_dir()
+
+
+# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py
+_STR_DTYPE_TO_TORCH_DTYPE = {
+    "half": torch.float16,
+    "float16": torch.float16,
+    "float": torch.float32,
+    "float32": torch.float32,
+    "bfloat16": torch.bfloat16,
+}
+
+
+# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py
+def _get_and_verify_dtype(
+    config: PretrainedConfig,
+    dtype: Union[str, torch.dtype],
+) -> torch.dtype:
+    # NOTE: getattr(config, "torch_dtype", torch.float32) is not correct
+    # because config.torch_dtype can be None.
+    if isinstance(config, dict):
+        config_dtype = config.get("dtype", None) or config.get("torch_dtype", None)
+        model_type = config.get("model_type", "")
+    else:
+        config_dtype = getattr(config, "dtype", None)
+        model_type = getattr(config, "model_type", "")
+    if isinstance(config_dtype, str):
+        config_dtype = _STR_DTYPE_TO_TORCH_DTYPE.get(config_dtype, None)
+    if config_dtype is None:
+        config_dtype = torch.float32
+
+    if isinstance(dtype, str):
+        dtype = dtype.lower()
+        if dtype == "auto":
+            if config_dtype == torch.float32:
+                if model_type.startswith("gemma"):
+                    if model_type == "gemma":
+                        gemma_version = ""
+                    else:
+                        gemma_version = model_type[5]
+                    logger.info(
+                        f"For Gemma {gemma_version}, we downcast float32 to bfloat16 instead "
+                        "of float16 by default. Please specify `dtype` if you "
+                        "want to use float16."
+                    )
+                    torch_dtype = torch.bfloat16
+                else:
+                    # Following the common practice, we use float16 for float32
+                    # models.
+                    torch_dtype = torch.float16
+            else:
+                torch_dtype = config_dtype
+        else:
+            if dtype not in _STR_DTYPE_TO_TORCH_DTYPE:
+                raise ValueError(f"Unknown dtype: {dtype}")
+            torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype]
+    elif isinstance(dtype, torch.dtype):
+        torch_dtype = dtype
+    else:
+        raise ValueError(f"Unknown dtype: {dtype}")
+
+    # Verify the dtype.
+    if torch_dtype != config_dtype:
+        if torch_dtype == torch.float32:
+            # Upcasting to float32 is allowed.
+            logger.debug("Upcasting %s to %s.", config_dtype, torch_dtype)
+            pass
+        elif config_dtype == torch.float32:
+            # Downcasting from float32 to float16 or bfloat16 is allowed.
+            logger.debug("Downcasting %s to %s.", config_dtype, torch_dtype)
+            pass
+        else:
+            # Casting between float16 and bfloat16 is allowed with a warning.
+            logger.debug("Casting %s to %s.", config_dtype, torch_dtype)
+
+    return torch_dtype
+
+
+def is_generation_model(model_architectures: List[str], is_embedding: bool = False):
+    # We have two ways to determine whether a model is a generative model.
+    # 1. Check the model architecture
+    # 2. check the `is_embedding` server args
+
+    if (
+        "LlamaEmbeddingModel" in model_architectures
+        or "MistralModel" in model_architectures
+        or "LlamaForSequenceClassification" in model_architectures
+        or "LlamaForSequenceClassificationWithNormal_Weights" in model_architectures
+        or "InternLM2ForRewardModel" in model_architectures
+        or "Qwen2ForRewardModel" in model_architectures
+        or "Qwen3ForRewardModel" in model_architectures
+        or "Qwen2ForSequenceClassification" in model_architectures
+        or "Qwen3ForSequenceClassification" in model_architectures
+        or "CLIPModel" in model_architectures
+        or "BertModel" in model_architectures
+        or "Contriever" in model_architectures
+        or "BertForSequenceClassification" in model_architectures
+        or "XLMRobertaModel" in model_architectures
+        or "XLMRobertaForSequenceClassification" in model_architectures
+        or "Gemma2ForSequenceClassification" in model_architectures
+    ):
+        return False
+    else:
+        return not is_embedding
+
+
+multimodal_model_archs = [
+    "CLIPModel",
+    "Cohere2VisionForConditionalGeneration",
+    "DeepseekVL2ForCausalLM",
+    "Ernie4_5_VLMoeForConditionalGeneration",
+    "MiniMaxM3SparseForConditionalGeneration",
+    "Gemma3ForConditionalGeneration",
+    "Gemma3nForConditionalGeneration",
+    "Gemma4ForConditionalGeneration",
+    "Gemma4UnifiedForConditionalGeneration",
+    "Glm4vForConditionalGeneration",
+    "Glm4vMoeForConditionalGeneration",
+    "GlmOcrForConditionalGeneration",
+    "GlmAsrForConditionalGeneration",
+    "GlmImageForConditionalGeneration",
+    "Grok1VForCausalLM",
+    "Grok1AForCausalLM",
+    "LlavaLlamaForCausalLM",
+    "Llama4ForConditionalGeneration",
+    "LlavaMistralForCausalLM",
+    "LlavaQwenForCausalLM",
+    "LlavaForConditionalGeneration",
+    "LlavaVidForCausalLM",
+    "Lfm2VlForConditionalGeneration",
+    "LightOnOCRForConditionalGeneration",
+    *MIMO_V2_MULTIMODAL_ARCHS,
+    "MiMoV2ASRForCausalLM",
+    "MiniCPMO",
+    "MiniCPMV",
+    "Mistral3ForConditionalGeneration",
+    "MultiModalityCausalLM",
+    "MllamaForConditionalGeneration",
+    "MossVLForConditionalGeneration",
+    "NemotronH_Nano_VL_V2",
+    "NemotronH_Nano_Omni_Reasoning_V3",
+    "PixtralForConditionalGeneration",
+    "Qwen2AudioForConditionalGeneration",
+    "Qwen2VLForConditionalGeneration",
+    "Qwen2_5_VLForConditionalGeneration",
+    "Qwen3VLForConditionalGeneration",
+    "Qwen3VLMoeForConditionalGeneration",
+    "Qwen3_5ForConditionalGeneration",
+    "Qwen3_5MoeForConditionalGeneration",
+    "InternS2PreviewForConditionalGeneration",
+    "Qwen3ASRForConditionalGeneration",
+    "Qwen3OmniMoeForConditionalGeneration",
+    "KimiVLForConditionalGeneration",
+    "LocateAnythingForConditionalGeneration",
+    "InternVLChatModel",
+    "InternS1ForConditionalGeneration",
+    "InternS1ProForConditionalGeneration",
+    "Phi4MMForCausalLM",
+    "VoxtralForConditionalGeneration",
+    "WhisperForConditionalGeneration",
+    "Step3VLForConditionalGeneration",
+    "POINTSV15ChatModel",
+    "DotsVLMForCausalLM",
+    "DotsOCRForCausalLM",
+    "Sarashina2VisionForCausalLM",
+    "NVILAForConditionalGeneration",
+    "NVILALiteForConditionalGeneration",
+    "DeepseekOCRForCausalLM",
+    "UnlimitedOCRForCausalLM",
+    "JetVLMForConditionalGeneration",
+    "PaddleOCRVLForConditionalGeneration",
+    "MiDashengLMModel",
+    "StepVLForConditionalGeneration",
+    "Step3p7ForConditionalGeneration",
+    "KimiK25ForConditionalGeneration",
+]
+
+piecewise_cuda_graph_disabled_model_archs = [
+    "AgnesForCausalLM",
+    "AgnesForCausalLMNextN",
+    "AgnesForCausalLMDSpark",
+    "Qwen3NextForCausalLM",
+    "BailingMoeV2_5ForCausalLM",
+    "LLaDAModelLM",
+]
+
+# Multimodal archs allowed to keep prefill piecewise CUDA graph enabled. The
+# generic "multimodal model" rule in ServerArgs disables prefill piecewise CG for
+# all multimodal models; archs here opt back in because their LM prefill captures
+# cleanly (vision encoder runs eagerly outside the graph via general_mm_embed_routine).
+multimodal_piecewise_cuda_graph_supported_model_archs = [
+    "Cohere2VisionForConditionalGeneration",
+    "KimiK25ForConditionalGeneration",
+    "MiniMaxM3SparseForCausalLM",
+    "MiniMaxM3SparseForConditionalGeneration",
+]
+
+# Multimodal archs whose LM prefill is validated under breakable CUDA graph;
+# embed-carrying batches are rejected at replay (can_run_graph) and run eager.
+multimodal_breakable_cuda_graph_supported_model_archs = [
+    "Qwen3_5ForConditionalGeneration",
+    "Qwen3_5MoeForConditionalGeneration",
+]
+
+if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get():
+    multimodal_model_archs.append(external_mm_model_arch)
+
+
+def is_multimodal_model(model_architectures: List[str]):
+    if any(
+        multi_model_arch in model_architectures
+        for multi_model_arch in multimodal_model_archs
+    ):
+        return True
+    else:
+        return False
+
+
+def is_audio_model(model_architectures: List[str]):
+    models = [
+        "WhisperForConditionalGeneration",
+        "Qwen3ASRForConditionalGeneration",
+    ]
+    return any(model in model_architectures for model in models)
+
+
+def is_encoder_decoder_model(model_architectures: List[str]):
+    models = [
+        "WhisperForConditionalGeneration",
+        "MllamaForConditionalGeneration",
+        "MossVLForConditionalGeneration",
+    ]
+    return any(model in model_architectures for model in models)
+
+
+def is_local_attention_model(model_architectures: List[str]):
+    return "Llama4ForConditionalGeneration" in model_architectures
+
+
+def is_multimodal_chunked_prefill_supported(model_architectures: List[str]):
+    """Check if chunked prefill is supported for a MultiModal model."""
+    unsupported = [
+        "Grok1VForCausalLM",
+        "Grok1AForCausalLM",
+        "LlavaLlamaForCausalLM",
+        "MllamaForConditionalGeneration",
+        "MossVLForConditionalGeneration",
+        "CLIPModel",
+    ]
+    if any(multi_model_arch in unsupported for multi_model_arch in model_architectures):
+        return False
+    else:
+        return True
+
+
+def is_piecewise_cuda_graph_disabled_model(model_architectures: List[str]):
+    return any(
+        arch in piecewise_cuda_graph_disabled_model_archs
+        for arch in model_architectures
+    )
+
+
+def is_multimodal_piecewise_cuda_graph_supported(model_architectures: List[str]):
+    """Whether a multimodal arch may keep prefill piecewise CUDA graph enabled."""
+    return any(
+        arch in multimodal_piecewise_cuda_graph_supported_model_archs
+        for arch in model_architectures
+    )
+
+
+def is_multimodal_breakable_cuda_graph_supported(model_architectures: List[str]):
+    """Whether a multimodal arch may keep prefill breakable CUDA graph enabled."""
+    return any(
+        arch in multimodal_breakable_cuda_graph_supported_model_archs
+        for arch in model_architectures
+    )
+
+
+# SequenceClassification models that use CrossEncodingPooler
+_cross_encoding_pooler_archs = [
+    "BertForSequenceClassification",
+    "XLMRobertaForSequenceClassification",
+]
+
+
+def is_cross_encoding_pooler_model(model_architectures: List[str]) -> bool:
+    return any(arch in _cross_encoding_pooler_archs for arch in model_architectures)
+
+
+def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
+    if scale <= 1:
+        return 1.0
+    return 0.1 * mscale * math.log(scale) + 1.0
+
+
+def compute_mla_mscale_scaling(rope_scaling: dict, base_scaling: float) -> float:
+    """Compute MLA attention scaling factor from rope_scaling with mscale.
+
+    Used by DeepSeek, BailingMoe, SarvamMLA and similar MLA models.
+    Warns if 'factor' is missing from rope_scaling (common in v5 configs).
+    """
+    if not rope_scaling.get("apply_yarn_scaling", True) or not rope_scaling.get(
+        "apply_scale", True
+    ):
+        return base_scaling
+
+    mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
+    if "factor" not in rope_scaling:
+        logger.warning(
+            "rope_scaling missing 'factor', defaulting to 1.0. "
+            "Check model accuracy.",
+        )
+    scaling_factor = rope_scaling.get("factor", 1.0)
+    mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
+    return base_scaling * mscale * mscale
+
+
+def is_hybrid_swa_model(
+    model_architectures: List[str],
+    hf_text_config: Optional[PretrainedConfig] = None,
+):
+
+    hybrid_swa_archs = {
+        "Llama4ForConditionalGeneration",
+        "AgnesForCausalLM",
+        "AgnesForCausalLMNextN",
+        "AgnesForCausalLMDSpark",
+        "GptOssForCausalLM",
+        *MIMO_V2_MODEL_ARCHS,
+        "MiMoV2MTP",
+        "Step3p5ForCausalLM",
+        "Step3p5MTP",
+        "Step3p7ForConditionalGeneration",
+        "Gemma4ForCausalLM",
+        "Gemma4ForConditionalGeneration",
+        "Gemma4UnifiedForConditionalGeneration",
+        "LagunaForCausalLM",
+        "MellumForCausalLM",
+        "InklingForConditionalGeneration",
+        "InklingForConditionalGenerationMTP",
+        "UnlimitedOCRForCausalLM",
+    }
+    if any(arch in hybrid_swa_archs for arch in model_architectures):
+        # Only treat Laguna as hybrid SWA when it actually has a sliding window.
+        if (
+            "LagunaForCausalLM" in model_architectures
+            and hf_text_config is not None
+            and not getattr(hf_text_config, "sliding_window", 0)
+        ):
+            return False
+        return True
+    # Also recognize models that explicitly opt-in via their HF text config,
+    # so custom hybrid-SWA architectures don't need to be added to the allowlist.
+    if hf_text_config is not None and getattr(hf_text_config, "is_hybrid_swa", False):
+        return True
+    return False
+
+
+def get_hybrid_layer_ids(
+    model_architectures: List[str],
+    hf_text_config: PretrainedConfig,
+):
+    num_hidden_layers = hf_text_config.num_hidden_layers
+    if "Llama4ForConditionalGeneration" in model_architectures:
+        swa_attention_layer_ids = [
+            i for i in range(num_hidden_layers) if (i + 1) % 4 != 0
+        ]
+        full_attention_layer_ids = [
+            i for i in range(num_hidden_layers) if (i + 1) % 4 == 0
+        ]
+    elif "GptOssForCausalLM" in model_architectures:
+        layer_types = getattr(hf_text_config, "layer_types", [])
+        swa_attention_layer_ids = [
+            i for i, x in enumerate(layer_types) if x == "sliding_attention"
+        ]
+        full_attention_layer_ids = [
+            i for i, x in enumerate(layer_types) if x == "full_attention"
+        ]
+    elif any(arch in MIMO_V2_MODEL_ARCHS for arch in model_architectures):
+        hybrid_layer_pattern = getattr(hf_text_config, "hybrid_layer_pattern", None)
+        swa_attention_layer_ids = [
+            i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 1
+        ]
+        full_attention_layer_ids = [
+            i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 0
+        ]
+    elif "MiMoV2MTP" in model_architectures:
+        swa_attention_layer_ids = [0]
+        full_attention_layer_ids = []
+    elif (
+        "Step3p5ForCausalLM" in model_architectures
+        or "Step3p7ForConditionalGeneration" in model_architectures
+    ):
+        layer_types = hf_text_config.layer_types
+        swa_attention_layer_ids = [
+            i
+            for i, x in enumerate(layer_types)
+            if x == "sliding_attention" and i < num_hidden_layers
+        ]
+        full_attention_layer_ids = [
+            i
+            for i, x in enumerate(layer_types)
+            if x == "full_attention" and i < num_hidden_layers
+        ]
+    elif "Step3p5MTP" in model_architectures:
+        swa_attention_layer_ids = [0]
+        full_attention_layer_ids = []
+    elif (
+        "Gemma4ForCausalLM" in model_architectures
+        or "Gemma4ForConditionalGeneration" in model_architectures
+        or "Gemma4UnifiedForConditionalGeneration" in model_architectures
+        or "LagunaForCausalLM" in model_architectures
+        or "MellumForCausalLM" in model_architectures
+    ):
+        layer_types = getattr(hf_text_config, "layer_types", [])
+        swa_attention_layer_ids = [
+            i for i, x in enumerate(layer_types) if x == "sliding_attention"
+        ]
+        full_attention_layer_ids = [
+            i for i, x in enumerate(layer_types) if x == "full_attention"
+        ]
+    elif "InklingForConditionalGenerationMTP" in model_architectures:
+        # One block per MTP depth; a banded head marks its sliding-window depths
+        # in mtp_local_layer_ids. The per-depth pool routing in the KV-cache
+        # mixin is authoritative; this keeps model_config's swa/full lists
+        # self-consistent for other consumers.
+        mtp_local_layer_ids = hf_text_config.mtp_local_layer_ids
+        if mtp_local_layer_ids:
+            num_depths = hf_text_config.num_nextn_predict_layers
+            local_set = set(mtp_local_layer_ids)
+            swa_attention_layer_ids = sorted(local_set)
+            full_attention_layer_ids = [
+                i for i in range(num_depths) if i not in local_set
+            ]
+        else:
+            swa_attention_layer_ids = []
+            full_attention_layer_ids = [0]
+    elif "InklingForConditionalGeneration" in model_architectures:
+        local_layer_ids = hf_text_config.local_layer_ids
+        local_layer_id_set = set(local_layer_ids)
+        assert len(local_layer_id_set) == len(
+            local_layer_ids
+        ), f"Inkling local_layer_ids must be unique: {local_layer_ids}"
+        assert all(
+            0 <= layer_id < num_hidden_layers for layer_id in local_layer_id_set
+        ), f"Inkling local_layer_ids must be in [0, {num_hidden_layers}): {local_layer_ids}"
+        swa_attention_layer_ids = [
+            i for i in range(num_hidden_layers) if i in local_layer_id_set
+        ]
+        full_attention_layer_ids = [
+            i for i in range(num_hidden_layers) if i not in local_layer_id_set
+        ]
+    elif "UnlimitedOCRForCausalLM" in model_architectures:
+        swa_attention_layer_ids = list(range(num_hidden_layers))
+        full_attention_layer_ids = []
+    elif getattr(hf_text_config, "hybrid_layer_pattern", None) is not None:
+        # Generic fallback for custom hybrid SWA models that opt in via
+        # hf_text_config.is_hybrid_swa and expose a hybrid_layer_pattern
+        # (1 = SWA, 0 = full) without needing to be added to the allowlist.
+        hybrid_layer_pattern = hf_text_config.hybrid_layer_pattern
+        swa_attention_layer_ids = [
+            i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 1
+        ]
+        full_attention_layer_ids = [
+            i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 0
+        ]
+    else:
+        swa_attention_layer_ids = None
+        full_attention_layer_ids = None
+    return swa_attention_layer_ids, full_attention_layer_ids

sglang_patch/srt/entrypoints/openai/chat_encoding.py

@@ -0,0 +1,97 @@
+"""Single home for the chat-encoding dispatch.
+
+Which encoder turns chat messages into prompt tokens is a property of the
+model, so the serving path and offline tools (benchmarks, evals) must resolve
+it here instead of re-deriving it from model architectures themselves.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+
+def resolve_chat_encoding_spec(
+    *,
+    hf_config: Any,
+    tokenizer: Any,
+    tool_call_parser: Optional[str] = None,
+) -> Optional[str]:
+    """Return the chat encoding spec for a model: "dsv4", "dsv32", "inkling", or None.
+
+    None means the default path (HF chat template).
+    """
+    if tool_call_parser == "deepseekv4":
+        return "dsv4"
+    if tool_call_parser == "deepseekv32":
+        return "dsv32"
+
+    architectures = hf_config.architectures
+    arch = architectures[0] if architectures else ""
+
+    if "Agnes" in arch:
+        return "dsv4"
+
+    # Inkling has no Jinja chat_template and uses a tiktoken base + a special-token
+    # overlay + negative MM placeholders, so it can't go through apply_chat_template;
+    # render input_ids directly via the Inkling renderer (serving_chat._encode_messages).
+    if "InklingForConditionalGeneration" in arch:
+        return "inkling"
+
+    has_chat_template = tokenizer is not None and tokenizer.chat_template is not None
+    if "DeepseekV3" in arch and not has_chat_template:
+        return "dsv32"
+    return None
+
+
+def encode_simple_chat(
+    *,
+    tokenizer: Any,
+    spec: Optional[str],
+    messages: List[Dict[str, Any]],
+    thinking_mode: str = "chat",
+) -> List[int]:
+    """Encode a plain-text chat conversation into prompt token ids.
+
+    Minimal encode for offline tools: no tools, no multimodal content, no
+    continue_final_message; the serving path keeps its full request-level
+    pipeline in ``serving_chat``. Like
+    ``serving_chat``, an empty system message is prepended when the
+    conversation does not start with one (for the dsv4/dsv32 encoders this
+    currently renders to zero tokens, but keeping the insertion explicit ties
+    this helper to the serving semantics rather than to that coincidence).
+    """
+    if spec == "inkling":
+        from sglang.srt.parser.inkling_renderer import render_inkling_messages
+        from sglang.srt.parser.inkling_tokenizer import InklingTokenizer
+
+        return render_inkling_messages(
+            messages,
+            InklingTokenizer(tokenizer=tokenizer),
+            add_generation_prompt=False,
+        )
+
+    if spec in ("dsv4", "dsv32"):
+        if messages and messages[0]["role"] != "system":
+            messages = [{"role": "system", "content": ""}] + list(messages)
+        if spec == "dsv4":
+            from sglang.srt.entrypoints.openai import encoding_dsv4
+
+            real_input = encoding_dsv4.encode_messages(
+                messages, thinking_mode=thinking_mode
+            )
+        else:
+            from sglang.srt.entrypoints.openai import encoding_dsv32
+
+            real_input = encoding_dsv32.encode_messages(
+                messages, thinking_mode=thinking_mode
+            )
+        return tokenizer.encode(real_input)
+
+    if getattr(tokenizer, "chat_template", None) is None:
+        raise ValueError(
+            "This model has no HF chat template and no custom chat encoder; "
+            f"cannot encode chat messages with {getattr(tokenizer, 'name_or_path', tokenizer)!r}."
+        )
+    return tokenizer.apply_chat_template(
+        messages, add_generation_prompt=True, tokenize=True
+    )

sglang_patch/srt/entrypoints/openai/serving_chat.py

@@ -0,0 +1,2264 @@
+from __future__ import annotations
+
+import os
+import copy
+import json
+import logging
+import math
+import time
+import uuid
+from enum import Enum
+from http import HTTPStatus
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional, Union
+
+
+class ThinkingMode(str, Enum):
+    """Mode for message encoding - chat vs thinking/reasoning."""
+
+    CHAT = "chat"
+    THINKING = "thinking"
+
+
+import jinja2
+import orjson
+from fastapi import Request
+from fastapi.responses import ORJSONResponse, StreamingResponse
+from jsonschema import Draft202012Validator, SchemaError
+
+from sglang.srt.entrypoints.openai import encoding_dsv4, encoding_dsv32
+from sglang.srt.entrypoints.openai.protocol import (
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionResponseChoice,
+    ChatCompletionResponseStreamChoice,
+    ChatCompletionStreamResponse,
+    ChatCompletionTokenLogprob,
+    ChatMessage,
+    ChoiceLogprobs,
+    DeltaMessage,
+    ErrorResponse,
+    FunctionResponse,
+    LogProbs,
+    MessageProcessingResult,
+    PromptTokensDetails,
+    ResponseParserProtocol,
+    SglExt,
+    ToolCall,
+    ToolCallProcessingResult,
+    ToolChoice,
+    TopLogprob,
+)
+from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
+from sglang.srt.entrypoints.openai.sse_utils import build_sse_content
+from sglang.srt.entrypoints.openai.usage_processor import UsageProcessor
+from sglang.srt.entrypoints.openai.utils import (
+    cached_tokens_details_from_dict,
+    process_cached_tokens_details_from_ret,
+    process_hidden_states_from_ret,
+    process_routed_experts_from_ret,
+    should_include_usage,
+    to_openai_style_logprobs,
+)
+from sglang.srt.environ import envs
+from sglang.srt.function_call.core_types import ToolCallItem
+from sglang.srt.function_call.function_call_parser import FunctionCallParser
+from sglang.srt.function_call.json_array_parser import JsonArrayParser
+from sglang.srt.function_call.utils import (
+    get_json_schema_constraint,
+    normalize_json_schema_types,
+)
+from sglang.srt.managers.io_struct import GenerateReqInput
+from sglang.srt.parser.conversation import generate_chat_conv
+from sglang.srt.parser.jinja_template_utils import process_content_for_template_format
+from sglang.srt.parser.reasoning_parser import ReasoningParser
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.tokenizer_manager import TokenizerManager
+    from sglang.srt.parser.template_manager import TemplateManager
+
+logger = logging.getLogger(__name__)
+
+
+def normalize_tool_content(role: str, content):
+    """Normalize tool message content from OpenAI array format to plain string.
+
+    OpenAI clients may send tool content as a list of content parts
+    (e.g. [{"type":"text","text":"..."}]) but most chat templates expect
+    a plain string for tool messages. Only flatten when ALL items are
+    pure OpenAI text parts; preserve lists containing non-text-type items
+    that some templates intentionally iterate over.
+    """
+    if role != "tool" or not isinstance(content, list):
+        return content
+    parts = content
+    is_openai_text_parts = all(
+        (isinstance(p, dict) and p.get("type") == "text") or isinstance(p, str)
+        for p in parts
+    )
+    if is_openai_text_parts:
+        text_parts = [p.get("text", "") if isinstance(p, dict) else p for p in parts]
+        return " ".join(text_parts)
+    return content
+
+
+def parse_tool_call_arguments(arguments: str) -> Dict[str, Any]:
+    """Parse OpenAI tool call arguments for chat templates."""
+    try:
+        parsed_arguments = orjson.loads(arguments)
+    except orjson.JSONDecodeError as exc:
+        raise ValueError(
+            "Assistant tool call function.arguments must be valid JSON."
+        ) from exc
+
+    if not isinstance(parsed_arguments, dict):
+        raise ValueError(
+            "Assistant tool call function.arguments must be a JSON object."
+        )
+
+    return parsed_arguments
+
+
+def normalize_assistant_tool_call_arguments(message: Dict[str, Any]) -> None:
+    """Normalize assistant history tool call arguments in-place."""
+    if message.get("role") != "assistant" or not isinstance(
+        message.get("tool_calls"), list
+    ):
+        return
+
+    for item in message["tool_calls"]:
+        function = item.get("function") if isinstance(item, dict) else None
+        if not isinstance(function, dict):
+            continue
+        if "arguments" in function and isinstance(function["arguments"], str):
+            function["arguments"] = parse_tool_call_arguments(function["arguments"])
+
+
+def _extract_max_dynamic_patch(request: ChatCompletionRequest):
+    img_vals = []
+    vid_vals = []
+    for msg in request.messages or []:
+        content = getattr(msg, "content", None)
+        if not isinstance(content, list):
+            continue
+        for part in content:
+            # pydantic object or dict type
+            if getattr(part, "type", None) == "image_url":
+                iu = getattr(part, "image_url", None)
+                mdp = getattr(iu, "max_dynamic_patch", None) if iu else None
+                if mdp is not None:
+                    img_vals.append(int(mdp))
+            elif getattr(part, "type", None) == "video_url":
+                vu = getattr(part, "video_url", None)
+                mdp = getattr(vu, "max_dynamic_patch", None) if vu else None
+                if mdp is not None:
+                    vid_vals.append(int(mdp))
+
+    # TODO(yuan-luo): per-item max_dynamic_patch for both image and video
+    img_max_dynamic_patch = min(img_vals) if img_vals else None
+    vid_max_dynamic_patch = min(vid_vals) if vid_vals else None
+    return img_max_dynamic_patch, vid_max_dynamic_patch
+
+
+class OpenAIServingChat(OpenAIServingBase):
+    """Handler for /v1/chat/completions requests"""
+
+    _default_sampling_params_logged = False
+
+    def __init__(
+        self,
+        tokenizer_manager: TokenizerManager,
+        template_manager: TemplateManager,
+    ):
+        super().__init__(tokenizer_manager)
+        self.template_manager = template_manager
+        self.tool_call_parser = self.tokenizer_manager.server_args.tool_call_parser
+        self.reasoning_parser = self.tokenizer_manager.server_args.reasoning_parser
+        self.default_chat_template_kwargs = (
+            self.tokenizer_manager.server_args.default_chat_template_kwargs or {}
+        )
+        self._reasoning_detector = None
+        if self.reasoning_parser:
+            try:
+                rp = ReasoningParser(
+                    model_type=self.reasoning_parser,
+                    stream_reasoning=True,
+                    tokenizer=self.tokenizer_manager.tokenizer,
+                )
+                self._reasoning_detector = rp.detector
+            except ValueError as e:
+                logger.warning(
+                    "Failed to initialize reasoning detector for parser '%s': %s",
+                    self.reasoning_parser,
+                    e,
+                )
+
+        # Get default sampling parameters from model's generation config
+        self.default_sampling_params = (
+            self.tokenizer_manager.model_config.get_default_sampling_params()
+        )
+        if (
+            self.default_sampling_params
+            and not OpenAIServingChat._default_sampling_params_logged
+        ):
+            logger.info(
+                f"Using default chat sampling params from model generation config: {self.default_sampling_params}",
+            )
+            OpenAIServingChat._default_sampling_params_logged = True
+
+        # Check if the model is a GPT-OSS model
+        self.is_gpt_oss = (
+            hasattr(self.tokenizer_manager.model_config, "hf_config")
+            and hasattr(self.tokenizer_manager.model_config.hf_config, "model_type")
+            and self.tokenizer_manager.model_config.hf_config.model_type == "gpt_oss"
+        )
+        self.is_gemma4 = (
+            hasattr(self.tokenizer_manager.model_config, "hf_config")
+            and hasattr(self.tokenizer_manager.model_config.hf_config, "model_type")
+            and self.tokenizer_manager.model_config.hf_config.model_type
+            in ("gemma4", "gemma4_unified")
+        )
+
+        # Which Python-based chat encoder (if any) bypasses apply_chat_template.
+        # Values: "dsv32", "dsv4", or custom values set by subclass. None for default.
+        self.chat_encoding_spec = self._resolve_chat_encoding_spec()
+
+        # Resolve the env-configured Inkling effort default once: the env var is
+        # frozen for the server's lifetime, and a misconfigured value should
+        # fail at boot, not 400 every request.
+        self._inkling_default_reasoning_effort: Optional[float] = (
+            self._get_inkling_default_reasoning_effort()
+            if self.chat_encoding_spec == "inkling"
+            else None
+        )
+
+        # Per-request response parser for custom decoding (set by _encode_messages)
+        self._response_parser: Optional[ResponseParserProtocol] = None
+
+        # Probe whether ``encode("")`` returns specials. If it does, we must
+        # keep ``add_special_tokens=False`` at the chat-template encode site
+        # to avoid double BOS; otherwise the kwarg is a no-op and dropping it
+        # lets slow tokenizers (e.g. Kimi's TikTokenTokenizer) stay on the
+        # fast internal path.
+        try:
+            self._tokenizer_auto_adds_specials = (
+                len(self.tokenizer_manager.tokenizer.encode("")) > 0
+            )
+        except Exception:
+            self._tokenizer_auto_adds_specials = True
+
+    def _handle_last_assistant_message(
+        self,
+        messages: List[Dict[str, Any]],
+        request: ChatCompletionRequest,
+    ) -> tuple[List[Dict[str, Any]], Optional[str]]:
+        """
+        Handle continue_final_message feature: separate final assistant message.
+
+        If continue_final_message is enabled and the last message is from assistant,
+        extract its content and remove it from the message list.
+        If continue_final_message is False and the last message is from assistant,
+        convert it to a user message to ensure the last message is always from user.
+
+        Only processes text-based content (strings), ignoring multimodal content (lists).
+
+        Args:
+            messages: List of message dictionaries
+            request: ChatCompletionRequest with continue_final_message flag
+
+        Returns:
+            Tuple of (processed_messages, assistant_prefix)
+            - processed_messages: Messages with last assistant message handled appropriately
+            - assistant_prefix: Content of the last assistant message (string only), or None
+        """
+        assistant_prefix = None
+        if messages and messages[-1].get("role") == "assistant":
+            last_content = messages[-1].get("content")
+            # Only process string content, ignore multimodal content (lists)
+            if isinstance(last_content, str):
+                if request.continue_final_message:
+                    # Extract content and remove the assistant message
+                    assistant_prefix = last_content
+                    messages = messages[:-1]
+                else:
+                    # Convert the last assistant message to user message
+                    messages[-1] = {"role": "user", "content": last_content}
+        return messages, assistant_prefix
+
+    def _append_assistant_prefix_to_prompt_ids(
+        self, prompt_ids: List[int], assistant_prefix: str
+    ) -> List[int]:
+        """
+        Append assistant prefix to prompt_ids.
+
+        Args:
+            prompt_ids: Current prompt token IDs
+            assistant_prefix: Assistant message content to append
+
+        Returns:
+            Updated prompt_ids with assistant prefix appended
+        """
+        encoded = self.tokenizer_manager.tokenizer.encode(assistant_prefix)
+        if encoded and encoded[0] == self.tokenizer_manager.tokenizer.bos_token_id:
+            encoded = encoded[1:]
+        return prompt_ids + encoded
+
+    def _resolve_chat_encoding_spec(self) -> Optional[str]:
+        """Determine which chat encoding spec to use.
+
+        Override in subclass to add custom encoding specs.
+        """
+        from sglang.srt.entrypoints.openai.chat_encoding import (
+            resolve_chat_encoding_spec,
+        )
+
+        return resolve_chat_encoding_spec(
+            hf_config=self.tokenizer_manager.model_config.hf_config,
+            tokenizer=self.tokenizer_manager.tokenizer,
+            tool_call_parser=self.tool_call_parser,
+        )
+
+    def _request_id_prefix(self) -> str:
+        return "chatcmpl-"
+
+    def _encode_messages(
+        self,
+        messages: List[Dict[str, Any]],
+        request: ChatCompletionRequest,
+        thinking_mode: ThinkingMode,
+        tools: Optional[List[Dict]] = None,
+    ) -> Optional[List[int]]:
+        """Encode messages for custom chat_encoding_spec values.
+
+        Returns prompt_ids if handled, None to use default encoding.
+        """
+        if self.chat_encoding_spec == "inkling":
+            # Inkling: render messages -> input_ids with framing tokens + ONE placeholder per
+            # media (encoding/expansion happens later in InklingMultimodalProcessor). The
+            # server's tokenizer is the base tiktoken backend; wrap it so encode_special
+            # supplies the framing-token overlay.
+            from sglang.srt.parser.inkling_renderer import render_inkling_messages
+            from sglang.srt.parser.inkling_tokenizer import (
+                CONTENT_TEXT,
+                MESSAGE_MODEL,
+                InklingTokenizer,
+            )
+
+            inkling_tokenizer = InklingTokenizer(
+                tokenizer=self.tokenizer_manager.tokenizer
+            )
+            reasoning_effort = self._parse_inkling_reasoning_effort(
+                request.reasoning_effort
+            )
+            if reasoning_effort is None:
+                reasoning_effort = self._inkling_default_reasoning_effort
+            assistant_prefix = self._pop_inkling_assistant_prefix(messages, request)
+            prompt_ids = render_inkling_messages(
+                messages,
+                inkling_tokenizer,
+                add_generation_prompt=False,
+                tools=tools,
+                reasoning_effort=reasoning_effort,
+            )
+            if assistant_prefix is not None:
+                # Continue the final assistant message inside an OPEN model text
+                # block: header + payload, no <|end_message|> and no
+                # <|content_model_end_sampling|>, so the model resumes the turn.
+                prompt_ids += [
+                    inkling_tokenizer.encode_special(MESSAGE_MODEL),
+                    inkling_tokenizer.encode_special(CONTENT_TEXT),
+                    *inkling_tokenizer.encode_text(assistant_prefix),
+                ]
+            return prompt_ids
+        return None
+
+    @staticmethod
+    def _pop_inkling_assistant_prefix(
+        messages: List[Dict[str, Any]],
+        request: ChatCompletionRequest,
+    ) -> Optional[str]:
+        """Extract the trailing assistant text for ``continue_final_message``.
+
+        Only a plain-string assistant message with no tool calls and no
+        reasoning content can be continued; anything else renders as a closed
+        historical turn. Mutates ``messages`` in place (callers pass a copy).
+        """
+        if not request.continue_final_message or not messages:
+            return None
+        last = messages[-1]
+        if (
+            last.get("role") != "assistant"
+            or not isinstance(last.get("content"), str)
+            or last.get("tool_calls")
+            or last.get("reasoning_content")
+        ):
+            return None
+        messages.pop()
+        return last["content"]
+
+    @staticmethod
+    def _parse_inkling_reasoning_effort(
+        value: Optional[Union[str, float]],
+    ) -> Optional[float]:
+        """Convert an OpenAI-style reasoning_effort to an Inkling float."""
+        if value is None:
+            return None
+        if isinstance(value, bool):
+            raise ValueError("Inkling reasoning_effort must not be a boolean")
+        if isinstance(value, (int, float)):
+            parsed = float(value)
+            if not math.isfinite(parsed) or not 0.0 <= parsed <= 0.99:
+                raise ValueError("Inkling reasoning_effort must be in [0.0, 0.99]")
+            return parsed
+        _EFFORT_MAP = {
+            "none": 0.0,
+            "minimal": 0.1,
+            "low": 0.2,
+            "medium": 0.7,
+            "high": 0.9,
+            "xhigh": 0.99,
+            "max": 0.99,
+        }
+        if value in _EFFORT_MAP:
+            return _EFFORT_MAP[value]
+        try:
+            parsed = float(value)
+        except (ValueError, TypeError) as exc:
+            raise ValueError(f"invalid Inkling reasoning_effort: {value!r}") from exc
+        if not math.isfinite(parsed) or not 0.0 <= parsed <= 0.99:
+            raise ValueError("Inkling reasoning_effort must be in [0.0, 0.99]")
+        return parsed
+
+    @staticmethod
+    def _get_inkling_default_reasoning_effort() -> float:
+        """Read the default Inkling reasoning effort from the environment."""
+        from sglang.srt.environ import envs
+
+        val = envs.SGLANG_INKLING_DEFAULT_REASONING_EFFORT.get()
+        if not val:
+            return 0.9
+        try:
+            parsed = float(val)
+        except (ValueError, TypeError) as exc:
+            raise ValueError(
+                "SGLANG_INKLING_DEFAULT_REASONING_EFFORT must be numeric"
+            ) from exc
+        if not math.isfinite(parsed) or not 0.0 <= parsed <= 0.99:
+            raise ValueError(
+                "SGLANG_INKLING_DEFAULT_REASONING_EFFORT must be in [0.0, 0.99]"
+            )
+        return parsed
+
+    def _decode_response(self, ret_item: Dict[str, Any]) -> Union[str, ErrorResponse]:
+        """Extract text from response."""
+        return ret_item["text"]
+
+    def _get_parsed_response_fields(
+        self,
+        reasoning_text: Optional[str],
+        tool_calls: Optional[List[Dict]],
+    ) -> tuple[Optional[str], Optional[List[Dict]]]:
+        """Post-process reasoning and tool_calls before building response."""
+        return reasoning_text, tool_calls
+
+    def _continuous_usage_cached_details(
+        self, content: Dict[str, Any]
+    ) -> Optional[PromptTokensDetails]:
+        if not self.tokenizer_manager.server_args.enable_cache_report:
+            return None
+        return UsageProcessor._details_if_cached(
+            content["meta_info"].get("cached_tokens", 0)
+        )
+
+    async def _generate_stream_content(
+        self,
+        content: Dict[str, Any],
+        index: int,
+        request: ChatCompletionRequest,
+        stream_offsets: Dict[int, int],
+        reasoning_parser_dict: Dict,
+        parser_dict: Dict,
+        has_tool_calls: Dict[int, bool],
+        choice_logprobs: Optional[Dict],
+        finish_reason_type: Optional[str],
+        continuous_usage_stats: bool,
+        prompt_tokens: Dict[int, int],
+        reasoning_tokens: Dict[int, int],
+        completion_tokens: Dict[int, int],
+    ) -> AsyncGenerator[str, None]:
+        """Generate SSE chunks for streaming content."""
+        offset = stream_offsets.get(index, 0)
+        if self.tokenizer_manager.server_args.incremental_streaming_output:
+            delta = content["text"]
+        else:
+            delta = content["text"][offset:]
+            stream_offsets[index] = len(content["text"])
+
+        # Attach logprobs to the first chunk emitted this step (reasoning,
+        # tool-call, or content) so they aren't dropped when a parser is active
+        # nor duplicated across chunks; flush any leftover at the end.
+        remaining_logprobs = choice_logprobs
+
+        # Handle reasoning content
+        if self.reasoning_parser and request.separate_reasoning:
+            reasoning_text, delta = self._process_reasoning_stream(
+                index,
+                delta,
+                reasoning_parser_dict,
+                content,
+                request,
+                finish_reason_type,
+            )
+            if reasoning_text:
+                usage = None
+                if continuous_usage_stats:
+                    usage = UsageProcessor.calculate_token_usage(
+                        prompt_tokens=prompt_tokens.get(index, 0),
+                        reasoning_tokens=reasoning_tokens.get(index, 0),
+                        completion_tokens=completion_tokens.get(index, 0),
+                        cached_tokens=self._continuous_usage_cached_details(content),
+                    ).model_dump()
+
+                yield build_sse_content(
+                    chunk_id=content["meta_info"]["id"],
+                    created=int(time.time()),
+                    model=request.model,
+                    index=index,
+                    reasoning_content=reasoning_text,
+                    logprobs=remaining_logprobs,
+                    usage=usage,
+                )
+                remaining_logprobs = None
+
+        # Handle tool calls
+        if request.tool_choice != "none" and request.tools and self.tool_call_parser:
+            async for chunk in self._process_tool_call_stream(
+                index,
+                delta,
+                parser_dict,
+                content,
+                request,
+                has_tool_calls,
+                continuous_usage_stats,
+            ):
+                if chunk:
+                    yield chunk
+
+            # Send any remaining tool call arguments when generation finishes
+            if finish_reason_type is not None and index in parser_dict:
+                parser = parser_dict[index]
+                remaining_chunk = self._check_for_unstreamed_tool_args(
+                    parser, content, request, index
+                )
+                if remaining_chunk:
+                    yield remaining_chunk
+
+        else:
+            # Regular content
+            if delta:
+                usage = None
+                if continuous_usage_stats:
+                    usage = UsageProcessor.calculate_token_usage(
+                        prompt_tokens=prompt_tokens.get(index, 0),
+                        reasoning_tokens=reasoning_tokens.get(index, 0),
+                        completion_tokens=completion_tokens.get(index, 0),
+                        cached_tokens=self._continuous_usage_cached_details(content),
+                    ).model_dump()
+
+                yield build_sse_content(
+                    chunk_id=content["meta_info"]["id"],
+                    created=int(time.time()),
+                    model=request.model,
+                    index=index,
+                    content=delta,
+                    logprobs=remaining_logprobs,
+                    usage=usage,
+                )
+                remaining_logprobs = None
+
+        # Flush logprobs still unattached this step — only when a parser is
+        # active, since _process_tool_call_stream may consume the delta and emit
+        # no content chunk. On the plain path an empty-delta step has no chunk
+        # to attach to either way, and a standalone empty-delta logprobs chunk
+        # is not a shape clients expect.
+        if remaining_logprobs is not None and (
+            self.reasoning_parser or self.tool_call_parser
+        ):
+            usage = None
+            if continuous_usage_stats:
+                usage = UsageProcessor.calculate_token_usage(
+                    prompt_tokens=prompt_tokens.get(index, 0),
+                    reasoning_tokens=reasoning_tokens.get(index, 0),
+                    completion_tokens=completion_tokens.get(index, 0),
+                    cached_tokens=self._continuous_usage_cached_details(content),
+                ).model_dump()
+
+            yield build_sse_content(
+                chunk_id=content["meta_info"]["id"],
+                created=int(time.time()),
+                model=request.model,
+                index=index,
+                logprobs=remaining_logprobs,
+                usage=usage,
+            )
+
+    def _validate_request(self, request: ChatCompletionRequest) -> Optional[str]:
+        """Validate that the input is valid."""
+        if not request.messages:
+            return "Messages cannot be empty."
+
+        if (
+            isinstance(request.tool_choice, str)
+            and request.tool_choice.lower() == "required"
+            and not request.tools
+        ):
+            return "Tools cannot be empty if tool choice is set to required."
+
+        if request.tool_choice is not None and not isinstance(request.tool_choice, str):
+            if not request.tools:
+                return "Tools cannot be empty if tool choice is set to a specific tool."
+            tool_name = request.tool_choice.function.name
+            tool_exists = any(tool.function.name == tool_name for tool in request.tools)
+            if not tool_exists:
+                return f"Tool '{tool_name}' not found in tools list."
+
+        # Validate tool definitions
+        for i, tool in enumerate(request.tools or []):
+            if tool.function.parameters is None:
+                continue
+            try:
+                # Rewrite DB/ORM-style aliases (e.g. "varchar", "enum", "int")
+                # to standard JSON Schema types before validation. RecursionError
+                # guards against hand-crafted cyclic schemas so the request gets
+                # a 400 instead of crashing into a 500.
+                normalize_json_schema_types(tool.function.parameters)
+                Draft202012Validator.check_schema(tool.function.parameters)
+            except SchemaError as e:
+                return f"Tool {i} function has invalid 'parameters' schema: {str(e)}"
+            except RecursionError:
+                return (
+                    f"Tool {i} function 'parameters' schema is too deeply nested "
+                    "or contains a cycle."
+                )
+
+        max_output_tokens = request.max_completion_tokens or request.max_tokens
+        server_context_length = self.tokenizer_manager.server_args.context_length
+        if (
+            max_output_tokens
+            and server_context_length
+            and max_output_tokens > server_context_length
+        ) and not self.tokenizer_manager.server_args.allow_auto_truncate:
+            return (
+                f"max_completion_tokens is too large: {max_output_tokens}."
+                f"This model supports at most {server_context_length} completion tokens."
+            )
+
+        if request.response_format and request.response_format.type == "json_schema":
+            schema = getattr(request.response_format.json_schema, "schema_", None)
+            if schema is None:
+                return "schema_ is required for json_schema response format request."
+
+        return None
+
+    def _convert_to_internal_request(
+        self,
+        request: ChatCompletionRequest,
+        raw_request: Request = None,
+    ) -> tuple[GenerateReqInput, ChatCompletionRequest]:
+        reasoning_effort = (
+            request.chat_template_kwargs.pop("reasoning_effort", None)
+            if request.chat_template_kwargs
+            else None
+        )
+        if self.is_gpt_oss and reasoning_effort == "none":
+            raise ValueError(
+                f"Harmony does not support reasoning effort {reasoning_effort}"
+            )
+
+        if reasoning_effort is not None:
+            request.reasoning_effort = reasoning_effort
+
+        if request.stream:
+            if request.return_prompt_token_ids:
+                raise ValueError(
+                    "return_prompt_token_ids is not supported with streaming. "
+                    "Please set stream=false when using return_prompt_token_ids=true."
+                )
+            if request.return_meta_info:
+                raise ValueError(
+                    "return_meta_info is not supported with streaming. "
+                    "Please set stream=false when using return_meta_info=true."
+                )
+
+        is_multimodal = self.tokenizer_manager.model_config.is_multimodal
+
+        # Process messages and apply chat template
+        processed_messages = self._process_messages(request, is_multimodal)
+        # Build sampling parameters
+        sampling_params = request.to_sampling_params(
+            stop=processed_messages.stop,
+            model_generation_config=self.default_sampling_params,
+            tool_call_constraint=processed_messages.tool_call_constraint,
+        )
+
+        # Handle single vs multiple requests
+        if request.input_ids is not None:
+            prompt_kwargs = {"input_ids": processed_messages.prompt_ids}
+        elif is_multimodal:
+            # Standard VLMs render a text prompt (with placeholder strings) for the MM
+            # processor to tokenize. Inkling's custom encoder instead produces pre-rendered
+            # input_ids with single placeholders; pass those through so the MM processor
+            # expands them rather than re-tokenizing an empty prompt. Gated on the Inkling
+            # encoding spec so every other model keeps the standard text path.
+            if (
+                self.chat_encoding_spec == "inkling"
+                and isinstance(processed_messages.prompt_ids, list)
+                and processed_messages.prompt_ids
+            ):
+                prompt_kwargs = {"input_ids": processed_messages.prompt_ids}
+            else:
+                prompt_kwargs = {"text": processed_messages.prompt}
+        else:
+            if isinstance(processed_messages.prompt_ids, str):
+                prompt_kwargs = {"text": processed_messages.prompt_ids}
+            else:
+                prompt_kwargs = {"input_ids": processed_messages.prompt_ids}
+
+        # Extract custom labels from raw request headers
+        custom_labels = self.extract_custom_labels(raw_request)
+
+        # Extract routed_dp_rank from header (has higher priority than body)
+        effective_routed_dp_rank = self.extract_routed_dp_rank_from_header(
+            raw_request, request.routed_dp_rank
+        )
+
+        # Resolve LoRA adapter from model parameter or explicit lora_path
+        lora_path = self._resolve_lora_path(request.model, request.lora_path)
+        img_max_dynamic_patch, vid_max_dynamic_patch = _extract_max_dynamic_patch(
+            request
+        )
+        require_reasoning = self._get_reasoning_from_request(request)
+
+        adapted_request = GenerateReqInput(
+            **prompt_kwargs,
+            image_data=processed_messages.image_data,
+            video_data=processed_messages.video_data,
+            audio_data=processed_messages.audio_data,
+            sampling_params=sampling_params,
+            return_logprob=request.logprobs,
+            logprob_start_len=-1,
+            top_logprobs_num=request.top_logprobs or 0,
+            stream=request.stream,
+            return_text_in_logprobs=True,
+            modalities=processed_messages.modalities,
+            lora_path=lora_path,
+            bootstrap_host=request.bootstrap_host,
+            bootstrap_port=request.bootstrap_port,
+            bootstrap_room=request.bootstrap_room,
+            routed_dp_rank=effective_routed_dp_rank,
+            disagg_prefill_dp_rank=request.disagg_prefill_dp_rank,
+            return_hidden_states=request.return_hidden_states,
+            return_routed_experts=request.return_routed_experts,
+            routed_experts_start_len=request.routed_experts_start_len,
+            rid=request.rid,
+            session_id=request.session_id,
+            extra_key=self._compute_extra_key(request),
+            require_reasoning=require_reasoning,
+            priority=request.priority,
+            routing_key=self.extract_routing_key(raw_request),
+            custom_labels=custom_labels,
+            custom_logit_processor=request.custom_logit_processor,
+            images_config=getattr(request, "images_config", None),
+            image_max_dynamic_patch=img_max_dynamic_patch,
+            video_max_dynamic_patch=vid_max_dynamic_patch,
+            max_dynamic_patch=getattr(request, "max_dynamic_patch", None),
+            use_audio_in_video=getattr(request, "use_audio_in_video", False),
+            return_prompt_token_ids=request.return_prompt_token_ids,
+        )
+
+        return adapted_request, request
+
+    def _process_messages(
+        self, request: ChatCompletionRequest, is_multimodal: bool
+    ) -> MessageProcessingResult:
+        """Process chat messages and apply chat template"""
+        if self.default_chat_template_kwargs:
+            ctk = dict(request.chat_template_kwargs or {})
+            for k, v in self.default_chat_template_kwargs.items():
+                ctk.setdefault(k, v)
+            request.chat_template_kwargs = ctk
+            effort = ctk.get("reasoning_effort")
+            if effort is not None and request.reasoning_effort is None:
+                request.reasoning_effort = effort
+
+        # GptOss model needs to keep special tokens for harmony parsing
+        if self.is_gpt_oss or self.is_gemma4:
+            request.skip_special_tokens = False
+
+        self._patch_reasoning_skip_special_tokens(request)
+
+        thinking_mode = self._get_reasoning_from_request(request)
+        # SGLang's ReasonerGrammarBackend owns the reasoning prefix
+        # when --reasoning-parser is configured, so builtin xgrammar
+        # tags must describe only the post-reasoning tool-call suffix.
+        xgrammar_reasoning = thinking_mode and (
+            self.tokenizer_manager.server_args.reasoning_parser is None
+        )
+        tool_call_constraint = None
+
+        # Apply chat template and its stop strings
+        tools = None
+        if request.tools and request.tool_choice != "none":
+            request.skip_special_tokens = False
+            if not isinstance(request.tool_choice, str):
+                tools = [
+                    item.model_dump()
+                    for item in request.tools
+                    if item.function.name == request.tool_choice.function.name
+                ]
+            else:
+                tools = [item.model_dump() for item in request.tools]
+            if self.tool_call_parser:
+                parser = FunctionCallParser(
+                    request.tools,
+                    self.tool_call_parser,
+                    tokenizer=self.tokenizer_manager.tokenizer,
+                )
+                tool_call_constraint = parser.get_structure_constraint(
+                    request.tool_choice,
+                    parallel_tool_calls=request.parallel_tool_calls,
+                    thinking_mode=xgrammar_reasoning,
+                )
+            # Fallback: use generic JSON schema for required/named tool choice
+            # only when no parser-specific constraint was set
+            if tool_call_constraint is None and (
+                request.tool_choice == "required"
+                or isinstance(request.tool_choice, ToolChoice)
+            ):
+                json_schema = get_json_schema_constraint(
+                    request.tools,
+                    request.tool_choice,
+                    parallel_tool_calls=request.parallel_tool_calls,
+                )
+                tool_call_constraint = ("json_schema", json_schema)
+
+        # When input_ids are provided, skip template tokenization entirely;
+        # only stop tokens and tool_call_constraint are needed.
+        if request.input_ids is not None:
+            result = MessageProcessingResult(
+                prompt="",
+                prompt_ids=request.input_ids,
+                image_data=None,
+                audio_data=None,
+                video_data=None,
+                modalities=[],
+                stop=request.stop or [],
+            )
+        elif self.template_manager.chat_template_name is None:
+            result = self._apply_jinja_template(request, tools, is_multimodal)
+        else:
+            result = self._apply_conversation_template(request, is_multimodal)
+
+        result.tool_call_constraint = tool_call_constraint
+        return result
+
+    def _apply_jinja_template(
+        self,
+        request: ChatCompletionRequest,
+        tools: Optional[List[Dict]],
+        is_multimodal: bool,
+    ) -> MessageProcessingResult:
+        """Apply Jinja chat template"""
+        prompt = ""
+        prompt_ids = []
+        openai_compatible_messages = []
+        image_data = []
+        video_data = []
+        audio_data = []
+        modalities = []
+
+        template_content_format = self.template_manager.jinja_template_content_format
+
+        # Try custom encoding first (override in subclass for custom renderers)
+        thinking_requested = (request.chat_template_kwargs or {}).get(
+            "thinking", envs.SGLANG_DEFAULT_THINKING.get()
+        )
+        thinking_mode = (
+            ThinkingMode.THINKING if thinking_requested else ThinkingMode.CHAT
+        )
+        messages = [msg.model_dump() for msg in request.messages]
+        for message in messages:
+            normalize_assistant_tool_call_arguments(message)
+
+        prompt_ids = self._encode_messages(
+            copy.deepcopy(messages),
+            request,
+            thinking_mode,
+            tools=tools,
+        )
+
+        if prompt_ids is not None:
+            # Custom encoding produced prompt_ids. Text-only encoders (dsv4/dsv32) need
+            # nothing more; Inkling is the only multimodal custom encoder and still needs the
+            # image/audio media harvested from the messages for the MM processor.
+            if self.chat_encoding_spec == "inkling":
+                for message in request.messages:
+                    msg_dict = message.model_dump()
+                    if msg_dict.get("content") is None:
+                        msg_dict["content"] = ""
+                    process_content_for_template_format(
+                        msg_dict,
+                        "openai",
+                        image_data,
+                        video_data,
+                        audio_data,
+                        modalities,
+                    )
+        elif self.chat_encoding_spec is not None:
+            # dsv4/dsv32 encoding path
+            messages = copy.deepcopy(messages)
+
+            # dsv4/dsv32 are text-only and consume string content; flatten
+            # OpenAI parts-list content here so the encoder sees a plain string.
+            for i, msg in enumerate(messages):
+                if isinstance(msg.get("content"), list):
+                    messages[i] = process_content_for_template_format(
+                        msg, "string", [], [], [], []
+                    )
+
+            for msg in messages:
+                if msg.get("content") is None:
+                    msg["content"] = ""
+                processed_msg = process_content_for_template_format(
+                    msg,
+                    template_content_format,
+                    image_data,
+                    video_data,
+                    audio_data,
+                    modalities,
+                    use_dpsk_v32_encoding=self.chat_encoding_spec == "dsv32",
+                )
+                msg.update(processed_msg)
+
+            # Handle continue_final_message: separate final assistant message
+            messages, assistant_prefix = self._handle_last_assistant_message(
+                messages, request
+            )
+
+            if messages[0]["role"] != "system":
+                # Insert a placeholder system message so tool schemas have
+                # somewhere to render. $AGNES_DEFAULT_SYSTEM_PROMPT gives it
+                # default content; unset (or empty) restores upstream behaviour
+                # of an empty system message. A system message supplied by the
+                # request is never overridden -- this branch only runs when the
+                # conversation does not open with one.
+                messages.insert(
+                    0,
+                    {
+                        "role": "system",
+                        "content": os.environ.get("AGNES_DEFAULT_SYSTEM_PROMPT", ""),
+                    },
+                )
+            if request.tools:
+                messages[0]["tools"] = [tool.model_dump() for tool in request.tools]
+
+            # Default encoding (dsv4/dsv32)
+            if self.chat_encoding_spec == "dsv4":
+                # V4 encoder only accepts "max" / "high" / None.
+                # OpenAI protocol defaults to "medium" which V4 rejects; drop it.
+                # Fallback: if request didn't set it, try env SGLANG_DSV4_REASONING_EFFORT.
+                effort_source = request.reasoning_effort
+                if effort_source is None:
+                    env_val = envs.SGLANG_DSV4_REASONING_EFFORT.get()
+                    if env_val:
+                        effort_source = env_val
+                v4_reasoning_effort = (
+                    effort_source if effort_source in ("max", "high") else None
+                )
+                if request.task is not None:
+                    encoding_dsv4.attach_task_to_last_user_message(
+                        messages, request.task
+                    )
+                real_input = encoding_dsv4.encode_messages(
+                    messages,
+                    thinking_mode=thinking_mode,
+                    reasoning_effort=v4_reasoning_effort,
+                )
+                prompt_ids = self.tokenizer_manager.tokenizer.encode(real_input)
+            else:
+                real_input = encoding_dsv32.encode_messages(
+                    messages, thinking_mode=thinking_mode
+                )
+                prompt_ids = self.tokenizer_manager.tokenizer.encode(real_input)
+
+            # Append assistant prefix if continue_final_message is enabled
+            if assistant_prefix:
+                prompt_ids = self._append_assistant_prefix_to_prompt_ids(
+                    prompt_ids, assistant_prefix
+                )
+        else:
+            for msg_dict in copy.deepcopy(messages):
+                if msg_dict.get("content") is None:
+                    msg_dict["content"] = ""
+
+                # Process content based on detected template format
+                processed_msg = process_content_for_template_format(
+                    msg_dict,
+                    template_content_format,
+                    image_data,
+                    video_data,
+                    audio_data,
+                    modalities,
+                )
+
+                processed_msg["content"] = normalize_tool_content(
+                    processed_msg["role"], processed_msg.get("content")
+                )
+
+                openai_compatible_messages.append(processed_msg)
+
+            # Handle continue_final_message: separate final assistant message
+            openai_compatible_messages, assistant_prefix = (
+                self._handle_last_assistant_message(openai_compatible_messages, request)
+            )
+
+            extra_template_kwargs = {}
+            if request.reasoning_effort is not None:
+                extra_template_kwargs["reasoning_effort"] = request.reasoning_effort
+            if request.chat_template_kwargs:
+                extra_template_kwargs.update(request.chat_template_kwargs)
+
+            rc = self.template_manager.reasoning_config
+            if rc is not None and rc.effort_kwarg is not None:
+                if request.reasoning_effort == "low":
+                    extra_template_kwargs.setdefault(rc.effort_kwarg, True)
+                elif request.reasoning_effort in ("medium", "high", "max"):
+                    logger.warning(
+                        "Model '%s' supports only 'low' reasoning effort; "
+                        "requested '%s' treated as default thinking",
+                        self.tokenizer_manager.server_args.served_model_name,
+                        request.reasoning_effort,
+                    )
+
+            # Split apply_chat_template(tokenize=True) into render + encode so we
+            # can skip add_special_tokens=False on tokenizers that don't auto-add
+            # specials (Kimi-like, OpenAI-chat analogue of #25265). Chat
+            # templates already include role/special tokens, so the encode must
+            # avoid double BOS on tokenizers that would add it.
+            encode_kwargs = (
+                {"add_special_tokens": False}
+                if self._tokenizer_auto_adds_specials
+                else {}
+            )
+            try:
+                rendered_prompt = self.tokenizer_manager.tokenizer.apply_chat_template(
+                    openai_compatible_messages,
+                    tokenize=False,
+                    add_generation_prompt=True,
+                    tools=tools,
+                    return_dict=False,
+                    **extra_template_kwargs,
+                )
+                prompt_ids = self.tokenizer_manager.tokenizer.encode(
+                    rendered_prompt, **encode_kwargs
+                )
+            except Exception:
+                # If the first attempt fails, try with flat function-only format.
+                # Some templates (e.g. Mistral) expect tools without the OpenAI wrapper.
+                tools = (
+                    [t["function"] if "function" in t else t for t in tools]
+                    if tools
+                    else None
+                )
+                try:
+                    rendered_prompt = (
+                        self.tokenizer_manager.tokenizer.apply_chat_template(
+                            openai_compatible_messages,
+                            tokenize=False,
+                            add_generation_prompt=True,
+                            tools=tools,
+                            return_dict=False,
+                            **extra_template_kwargs,
+                        )
+                    )
+                    prompt_ids = self.tokenizer_manager.tokenizer.encode(
+                        rendered_prompt, **encode_kwargs
+                    )
+                except (jinja2.TemplateError, TypeError) as template_error:
+                    # Template errors (e.g., from raise_exception in Jinja templates)
+                    # and TypeError (e.g., tojson filter on Jinja2 Undefined variables)
+                    # should be treated as client errors (400 BadRequest)
+                    raise ValueError(str(template_error)) from template_error
+
+            # Append assistant prefix if continue_final_message is enabled
+            if assistant_prefix:
+                prompt_ids = self._append_assistant_prefix_to_prompt_ids(
+                    prompt_ids, assistant_prefix
+                )
+
+            if is_multimodal:
+                prompt = self.tokenizer_manager.tokenizer.decode(prompt_ids)
+
+        stop = request.stop
+        image_data = image_data if image_data else None
+        audio_data = audio_data if audio_data else None
+        video_data = video_data if video_data else None
+        modalities = modalities if modalities else []
+        return MessageProcessingResult(
+            prompt=prompt,
+            prompt_ids=prompt_ids,
+            image_data=image_data,
+            video_data=video_data,
+            audio_data=audio_data,
+            modalities=modalities,
+            stop=stop,
+        )
+
+    def _apply_conversation_template(
+        self,
+        request: ChatCompletionRequest,
+        is_multimodal: bool,
+    ) -> MessageProcessingResult:
+        """Apply conversation template"""
+        prompt = ""
+        prompt_ids = []
+        conv = generate_chat_conv(request, self.template_manager.chat_template_name)
+
+        # If we should continue the final assistant message, adjust the conversation.
+        if (
+            request.continue_final_message
+            and request.messages
+            and request.messages[-1].role == "assistant"
+        ):
+            # Remove the auto-added blank assistant turn, if present.
+            if conv.messages and conv.messages[-1][1] is None:
+                conv.messages.pop()
+            # Rebuild the prompt from the conversation.
+            prompt = conv.get_prompt()
+            # Strip trailing stop tokens or separators that indicate end-of-assistant.
+            if isinstance(conv.stop_str, list):
+                for stop_token in conv.stop_str:
+                    if prompt.endswith(stop_token):
+                        prompt = prompt[: -len(stop_token)]
+            elif isinstance(conv.stop_str, str) and prompt.endswith(conv.stop_str):
+                prompt = prompt[: -len(conv.stop_str)]
+            if conv.sep and prompt.endswith(conv.sep):
+                prompt = prompt[: -len(conv.sep)]
+            if getattr(conv, "sep2", None) and prompt.endswith(conv.sep2):
+                prompt = prompt[: -len(conv.sep2)]
+        else:
+            prompt = conv.get_prompt()
+            if self._get_reasoning_from_request(request) and (
+                self._reasoning_detector is None
+                or not self._reasoning_detector.thinks_internally
+            ):
+                # Models with thinks_internally=True think without a leading <think> token
+                prompt += "<think>"  # Note(Xinyuan): hard code thinking token
+
+        image_data = conv.image_data if conv.image_data else None
+        video_data = conv.video_data if conv.video_data else None
+        audio_data = conv.audio_data if conv.audio_data else None
+        modalities = conv.modalities if conv.modalities else []
+        stop = copy.copy(conv.stop_str or [] if not request.ignore_eos else [])
+
+        if request.stop:
+            if isinstance(request.stop, str):
+                stop.append(request.stop)
+            else:
+                stop.extend(request.stop)
+
+        if not is_multimodal:
+            prompt_ids = self.tokenizer_manager.tokenizer.encode(prompt)
+
+        return MessageProcessingResult(
+            prompt=prompt,
+            prompt_ids=prompt_ids,
+            image_data=image_data,
+            video_data=video_data,
+            audio_data=audio_data,
+            modalities=modalities,
+            stop=stop,
+        )
+
+    async def _handle_streaming_request(
+        self,
+        adapted_request: GenerateReqInput,
+        request: ChatCompletionRequest,
+        raw_request: Request,
+    ) -> Union[StreamingResponse, ErrorResponse]:
+        """Handle streaming chat completion request"""
+        generator = self._generate_chat_stream(adapted_request, request, raw_request)
+
+        # Kick-start the generator to trigger validation before HTTP 200 is sent.
+        # If validation fails (e.g., context length exceeded), we can still return
+        # a proper HTTP 400 error response instead of streaming it as SSE payload.
+        try:
+            first_chunk = await generator.__anext__()
+        except ValueError as e:
+            return self.create_error_response(str(e))
+
+        async def prepend_first_chunk():
+            yield first_chunk
+            async for chunk in generator:
+                yield chunk
+
+        return StreamingResponse(
+            prepend_first_chunk(),
+            media_type="text/event-stream",
+            background=self.tokenizer_manager.create_abort_task(adapted_request),
+        )
+
+    async def _generate_chat_stream(
+        self,
+        adapted_request: GenerateReqInput,
+        request: ChatCompletionRequest,
+        raw_request: Request,
+    ) -> AsyncGenerator[str, None]:
+        """Generate streaming chat completion response"""
+        # Parsers for tool calls and reasoning
+        parser_dict = {}
+        reasoning_parser_dict = {}
+
+        # State tracking for streaming
+        is_firsts = {}
+        stream_offsets = {}
+        n_prev_tokens = {}
+        has_tool_calls = {}
+        finish_reasons = {}
+
+        # Usage tracking
+        prompt_tokens = {}
+        reasoning_tokens = {}
+        completion_tokens = {}
+        cached_tokens = {}
+        hidden_states = {}
+        routed_experts = {}
+        cached_tokens_details = {}
+        image_tokens = {}
+        audio_tokens = {}
+        video_tokens = {}
+
+        stream_started = False
+        try:
+            include_usage, continuous_usage_stats = should_include_usage(
+                request.stream_options,
+                self.tokenizer_manager.server_args.stream_response_default_include_usage,
+            )
+
+            async for content in self.tokenizer_manager.generate_request(
+                adapted_request, raw_request
+            ):
+                index = content.get("index", 0)
+
+                prompt_tokens[index] = content["meta_info"].get("prompt_tokens", 0)
+                completion_tokens[index] = content["meta_info"].get(
+                    "completion_tokens", 0
+                )
+                reasoning_tokens[index] = content["meta_info"].get(
+                    "reasoning_tokens", 0
+                )
+                cached_tokens[index] = content["meta_info"].get("cached_tokens", 0)
+                hidden_states[index] = content["meta_info"].get("hidden_states", None)
+                routed_experts[index] = content["meta_info"].get("routed_experts", None)
+                cached_tokens_details[index] = content["meta_info"].get(
+                    "cached_tokens_details", None
+                )
+                image_tokens[index] = content["meta_info"].get("image_tokens", 0)
+                audio_tokens[index] = content["meta_info"].get("audio_tokens", 0)
+                video_tokens[index] = content["meta_info"].get("video_tokens", 0)
+
+                # Handle logprobs
+                choice_logprobs = None
+                if request.logprobs:
+                    n_prev_token = n_prev_tokens.get(index, 0)
+                    total_output_logprobs = content["meta_info"][
+                        "output_token_logprobs_length"
+                    ]
+                    if n_prev_token < total_output_logprobs:
+                        choice_logprobs = self._process_streaming_logprobs(
+                            content, n_prev_token, total_output_logprobs
+                        ).model_dump()
+                    n_prev_tokens[index] = total_output_logprobs
+
+                finish_reason = content["meta_info"].get("finish_reason", None)
+                finish_reason_type = finish_reason["type"] if finish_reason else None
+
+                # Track finish_reason for each index
+                if finish_reason_type:
+                    # Abort with an explicit error status_code is a system error
+                    # (timeout, OOM, validation): emit a streaming error chunk.
+                    # A graceful abort (no status_code, e.g. user-initiated via
+                    # /abort_request or session lifecycle cleanup) falls through
+                    # to the normal chunk path, matching the non-stream behavior
+                    # in tokenizer_manager._handle_abort_finish_reason.
+                    if finish_reason_type == "abort" and isinstance(
+                        finish_reason.get("status_code"), HTTPStatus
+                    ):
+                        code = finish_reason["status_code"]
+                        error = self.create_streaming_error_response(
+                            finish_reason.get("message", "Generation aborted."),
+                            code.name,
+                            code.value,
+                        )
+                        yield f"data: {error}\n\n"
+                        break
+                    finish_reasons[index] = finish_reason
+
+                # First chunk with role
+                if is_firsts.get(index, True):
+                    is_firsts[index] = False
+                    yield build_sse_content(
+                        chunk_id=content["meta_info"]["id"],
+                        created=int(time.time()),
+                        model=request.model,
+                        index=index,
+                        role="assistant",
+                        content="",
+                    )
+                    stream_started = True
+
+                # Generate streaming content (override in subclass for custom behavior)
+                async for chunk in self._generate_stream_content(
+                    content=content,
+                    index=index,
+                    request=request,
+                    stream_offsets=stream_offsets,
+                    reasoning_parser_dict=reasoning_parser_dict,
+                    parser_dict=parser_dict,
+                    has_tool_calls=has_tool_calls,
+                    choice_logprobs=choice_logprobs,
+                    finish_reason_type=finish_reason_type,
+                    continuous_usage_stats=continuous_usage_stats,
+                    prompt_tokens=prompt_tokens,
+                    reasoning_tokens=reasoning_tokens,
+                    completion_tokens=completion_tokens,
+                ):
+                    yield chunk
+
+            # Send finish_reason chunks for each index that completed
+            for idx, finish_reason_data in finish_reasons.items():
+                finish_reason_type = finish_reason_data["type"]
+
+                # Change finish_reason to "tool_calls" if we had tool calls and stopped naturally
+                final_finish_reason = finish_reason_type
+                if has_tool_calls.get(idx, False) and finish_reason_type == "stop":
+                    final_finish_reason = "tool_calls"
+
+                matched_stop = finish_reason_data.get("matched")
+                yield build_sse_content(
+                    chunk_id=content["meta_info"]["id"],
+                    created=int(time.time()),
+                    model=request.model,
+                    index=idx,
+                    finish_reason=final_finish_reason,
+                    matched_stop=matched_stop,
+                )
+
+            # Send hidden states if requested
+            if request.return_hidden_states and hidden_states:
+                for index, choice_hidden_states in hidden_states.items():
+                    if choice_hidden_states:
+                        last_token_hidden_states = (
+                            choice_hidden_states[-1]
+                            if len(choice_hidden_states) > 1
+                            else []
+                        )
+                        hidden_states_chunk = ChatCompletionStreamResponse(
+                            id=content["meta_info"]["id"],
+                            created=int(time.time()),
+                            choices=[
+                                ChatCompletionResponseStreamChoice(
+                                    index=index,
+                                    delta=DeltaMessage(
+                                        hidden_states=last_token_hidden_states
+                                    ),
+                                    finish_reason=None,  # Hidden states don't need finish_reason
+                                )
+                            ],
+                            model=request.model,
+                        )
+                        yield f"data: {hidden_states_chunk.model_dump_json()}\n\n"
+
+            sglext_routed = None
+            if request.return_routed_experts and routed_experts:
+                sglext_routed = next(
+                    (v for v in routed_experts.values() if v is not None), None
+                )
+
+            sglext_details = None
+            if request.return_cached_tokens_details and cached_tokens_details:
+                first_details = next(
+                    (v for v in cached_tokens_details.values() if v is not None), None
+                )
+                if first_details is not None:
+                    sglext_details = cached_tokens_details_from_dict(first_details)
+
+            if sglext_routed is not None or sglext_details is not None:
+                sglext_chunk = ChatCompletionStreamResponse(
+                    id=content["meta_info"]["id"],
+                    created=int(time.time()),
+                    choices=[],  # sglext is at response level
+                    model=request.model,
+                    sglext=SglExt(
+                        routed_experts=sglext_routed,
+                        cached_tokens_details=sglext_details,
+                    ),
+                )
+                yield f"data: {sglext_chunk.model_dump_json()}\n\n"
+
+            # Additional usage chunk
+            if include_usage:
+                # Multimodal tokens are per-prompt (input side), so aggregate
+                # once per prompt (first choice), matching prompt/cached semantics.
+                total_image_tokens = sum(
+                    tok for idx, tok in image_tokens.items() if idx % request.n == 0
+                )
+                total_audio_tokens = sum(
+                    tok for idx, tok in audio_tokens.items() if idx % request.n == 0
+                )
+                total_video_tokens = sum(
+                    tok for idx, tok in video_tokens.items() if idx % request.n == 0
+                )
+                usage = UsageProcessor.calculate_streaming_usage(
+                    prompt_tokens,
+                    reasoning_tokens,
+                    completion_tokens,
+                    cached_tokens=cached_tokens,
+                    n_choices=request.n,
+                    enable_cache_report=self.tokenizer_manager.server_args.enable_cache_report,
+                    image_tokens=total_image_tokens,
+                    audio_tokens=total_audio_tokens,
+                    video_tokens=total_video_tokens,
+                )
+                usage_chunk = ChatCompletionStreamResponse(
+                    id=content["meta_info"]["id"],
+                    created=int(time.time()),
+                    choices=[],  # Empty choices array as per OpenAI spec
+                    model=request.model,
+                    usage=usage,
+                )
+                yield f"data: {usage_chunk.model_dump_json()}\n\n"
+
+        except ValueError as e:
+            if not stream_started:
+                raise
+            error = self.create_streaming_error_response(str(e))
+            yield f"data: {error}\n\n"
+
+        yield "data: [DONE]\n\n"
+
+    async def _handle_non_streaming_request(
+        self,
+        adapted_request: GenerateReqInput,
+        request: ChatCompletionRequest,
+        raw_request: Request,
+    ) -> Union[ChatCompletionResponse, ErrorResponse, ORJSONResponse]:
+        """Handle non-streaming chat completion request"""
+        try:
+            ret = await self.tokenizer_manager.generate_request(
+                adapted_request, raw_request
+            ).__anext__()
+        except ValueError as e:
+            return self.create_error_response(str(e))
+
+        if not isinstance(ret, list):
+            ret = [ret]
+
+        response = self._build_chat_response(
+            request,
+            ret,
+            int(time.time()),
+        )
+
+        return response
+
+    def _build_chat_response(
+        self,
+        request: ChatCompletionRequest,
+        ret: List[Dict[str, Any]],
+        created: int,
+    ) -> Union[ChatCompletionResponse, ORJSONResponse]:
+        """Build chat completion response from generation results"""
+        choices = []
+
+        # Build sglext at response level (from first ret_item, as these are per-request)
+        first_ret = ret[0]
+        routed_experts = process_routed_experts_from_ret(first_ret, request)
+        cached_tokens_details = process_cached_tokens_details_from_ret(
+            first_ret, request
+        )
+        response_sglext = None
+        if routed_experts or cached_tokens_details:
+            response_sglext = SglExt(
+                routed_experts=routed_experts,
+                cached_tokens_details=cached_tokens_details,
+            )
+
+        for idx, ret_item in enumerate(ret):
+            # Process logprobs
+            choice_logprobs = None
+            if request.logprobs:
+                choice_logprobs = self._process_response_logprobs(ret_item)
+
+            # Handle hidden states
+            hidden_states = process_hidden_states_from_ret(ret_item, request)
+
+            finish_reason = ret_item["meta_info"]["finish_reason"]
+
+            text = self._decode_response(ret_item)
+            if isinstance(text, ErrorResponse):
+                return ORJSONResponse(content=text.model_dump(), status_code=text.code)
+
+            # Handle reasoning content
+            reasoning_text = None
+            if self.reasoning_parser and request.separate_reasoning:
+                force_reasoning = (
+                    self.template_manager.force_reasoning
+                    or self._get_reasoning_from_request(request)
+                )
+                try:
+                    parser = ReasoningParser(
+                        model_type=self.reasoning_parser,
+                        stream_reasoning=False,
+                        force_reasoning=force_reasoning,
+                        request=request,
+                        tokenizer=self.tokenizer_manager.tokenizer,
+                    )
+                    reasoning_text, text = parser.parse_non_stream(text)
+                except Exception as e:
+                    logger.error(f"Reasoning parsing error: {e}")
+                    return self.create_error_response(
+                        "Failed to parse reasoning content",
+                        err_type="InternalServerError",
+                        status_code=500,
+                    )
+
+            # Handle tool calls
+            tool_calls = None
+            if (
+                request.tool_choice != "none"
+                and request.tools
+                and self.tool_call_parser
+            ):
+                history_tool_calls_cnt = self._get_history_tool_calls_cnt(request)
+                tool_calls, text, finish_reason = self._process_tool_calls(
+                    text,
+                    request.tools,
+                    finish_reason,
+                    request.tool_choice,
+                    history_tool_calls_cnt,
+                )
+
+            # Extract prompt_token_ids if requested
+            choice_prompt_token_ids = (
+                ret_item.get("prompt_token_ids")
+                if request.return_prompt_token_ids
+                else None
+            )
+
+            choice_meta_info = (
+                ret_item["meta_info"] if request.return_meta_info else None
+            )
+            # NOTE: content should not be None but empty string to make sure retokenize consistency.
+            reasoning_text, tool_calls = self._get_parsed_response_fields(
+                reasoning_text, tool_calls
+            )
+
+            choice_data = ChatCompletionResponseChoice(
+                index=idx,
+                message=ChatMessage(
+                    role="assistant",
+                    content=text if text else "",
+                    tool_calls=tool_calls,
+                    reasoning_content=reasoning_text if reasoning_text else None,
+                ),
+                logprobs=choice_logprobs,
+                finish_reason=finish_reason["type"] if finish_reason else None,
+                matched_stop=(
+                    finish_reason["matched"]
+                    if finish_reason and "matched" in finish_reason
+                    else None
+                ),
+                hidden_states=hidden_states,
+                prompt_token_ids=choice_prompt_token_ids,
+                meta_info=choice_meta_info,
+            )
+            choices.append(choice_data)
+
+        # Calculate usage. Multimodal tokens are per-prompt (input side), so
+        # aggregate once per prompt (stride by n), matching prompt/cached semantics.
+        image_tokens = sum(
+            ret[i]["meta_info"].get("image_tokens", 0)
+            for i in range(0, len(ret), request.n)
+        )
+        audio_tokens = sum(
+            ret[i]["meta_info"].get("audio_tokens", 0)
+            for i in range(0, len(ret), request.n)
+        )
+        video_tokens = sum(
+            ret[i]["meta_info"].get("video_tokens", 0)
+            for i in range(0, len(ret), request.n)
+        )
+        usage = UsageProcessor.calculate_response_usage(
+            ret,
+            n_choices=request.n,
+            enable_cache_report=self.tokenizer_manager.server_args.enable_cache_report,
+            image_tokens=image_tokens,
+            audio_tokens=audio_tokens,
+            video_tokens=video_tokens,
+        )
+
+        return ChatCompletionResponse(
+            id=ret[0]["meta_info"]["id"],
+            created=created,
+            model=request.model,
+            choices=choices,
+            usage=usage,
+            metadata={"weight_version": ret[0]["meta_info"]["weight_version"]},
+            sglext=response_sglext,
+        )
+
+    def _process_logprobs_tokens(
+        self, logprobs: LogProbs, use_token_index: bool = False
+    ) -> List[ChatCompletionTokenLogprob]:
+        """Common helper to process logprobs tokens for both streaming and non-streaming
+
+        Args:
+            logprobs: LogProbs data from model
+            use_token_index: True for non-streaming (use token_idx), False for streaming (use index 0)
+        """
+        token_logprobs = []
+
+        for token_idx, (token, logprob) in enumerate(
+            zip(logprobs.tokens, logprobs.token_logprobs)
+        ):
+            token_bytes = list(token.encode("utf-8"))
+            top_logprobs = []
+            if logprobs.top_logprobs:
+                # - Non-streaming (use_token_index=True): uses token_idx for full data
+                # - Streaming (use_token_index=False): uses index 0 for pre-sliced data
+                top_logprobs_idx = token_idx if use_token_index else 0
+                for top_token, top_logprob in logprobs.top_logprobs[
+                    top_logprobs_idx
+                ].items():
+                    top_token_bytes = list(top_token.encode("utf-8"))
+                    top_logprobs.append(
+                        TopLogprob(
+                            token=top_token,
+                            bytes=top_token_bytes,
+                            logprob=top_logprob,
+                        )
+                    )
+            token_logprobs.append(
+                ChatCompletionTokenLogprob(
+                    token=token,
+                    bytes=token_bytes,
+                    logprob=logprob,
+                    top_logprobs=top_logprobs,
+                )
+            )
+
+        return token_logprobs
+
+    def _process_response_logprobs(self, ret_item: Dict[str, Any]) -> ChoiceLogprobs:
+        """Process logprobs for non-streaming response"""
+        logprobs = to_openai_style_logprobs(
+            output_token_logprobs=ret_item["meta_info"]["output_token_logprobs"],
+            output_top_logprobs=ret_item["meta_info"].get("output_top_logprobs", None),
+        )
+
+        token_logprobs = self._process_logprobs_tokens(logprobs, use_token_index=True)
+        return ChoiceLogprobs(content=token_logprobs)
+
+    def _process_tool_call_id(
+        self,
+        call_item: ToolCallItem,
+        history_tool_calls_cnt: int,
+    ) -> str:
+        """Process for generating a new and unique `tool_call_id`"""
+        if self.tool_call_parser != "kimi_k2":
+            # A simple uuid is sufficient for all models except for Kimi-K2.
+            tool_call_id = f"call_{uuid.uuid4().hex[:24]}"
+            return tool_call_id
+        else:
+            # Align with Kimi-K2 format: functions.{name}:{index}
+            # Kimi-K2 allows multiple tool_calls in one message; SGLang sets call_item.tool_index to the *local* position inside that message.
+            # Therefore, the index must be corrected by using `history_tool_calls_cnt + call_item.tool_index` to ensure globally unique and properly ordered.
+            tool_call_id = f"functions.{call_item.name}:{history_tool_calls_cnt+call_item.tool_index}"
+            logger.debug(
+                f"Process tool call idx, parser: {self.tool_call_parser}, tool_call_id: {tool_call_id}, history_cnt: {history_tool_calls_cnt}"
+            )
+            return tool_call_id
+
+    def _process_tool_calls(
+        self,
+        text: str,
+        tools: List[Any],
+        finish_reason: Dict[str, Any],
+        tool_choice: Optional[Union[str, ToolChoice]] = None,
+        history_tool_calls_cnt: int = 0,
+    ) -> ToolCallProcessingResult:
+        """Process tool calls in the response"""
+
+        is_required = tool_choice == "required" or isinstance(tool_choice, ToolChoice)
+
+        # Try model-specific parser when output is in native format.
+        # For required/named: only use parser when structural_tag was used
+        # as constraint (mirrors the streaming path). For auto: always try.
+        if self.tool_call_parser:
+            parser = FunctionCallParser(
+                tools, self.tool_call_parser, tokenizer=self.tokenizer_manager.tokenizer
+            )
+            should_try_parser = (
+                not is_required or parser.detector.supports_structural_tag()
+            )
+            if should_try_parser and parser.has_tool_call(text):
+                try:
+                    text, call_info_list = parser.parse_non_stream(text)
+                    if not call_info_list:
+                        return ToolCallProcessingResult(None, text, finish_reason)
+
+                    tool_calls = []
+                    for call_info in call_info_list:
+                        tool_id = self._process_tool_call_id(
+                            call_info, history_tool_calls_cnt
+                        )
+                        tool_calls.append(
+                            ToolCall(
+                                id=tool_id,
+                                index=getattr(call_info, "tool_index", None),
+                                function=FunctionResponse(
+                                    name=call_info.name,
+                                    arguments=call_info.parameters,
+                                ),
+                            )
+                        )
+                    if finish_reason["type"] == "stop":
+                        finish_reason["type"] = "tool_calls"
+                        finish_reason["matched"] = None
+                    return ToolCallProcessingResult(tool_calls, text, finish_reason)
+                except Exception as e:
+                    logger.error(f"Tool call parsing error: {e}")
+                    return ToolCallProcessingResult(None, text, finish_reason)
+
+        # json_schema constraint → JSON array output for required/named
+        if is_required:
+            original_finish_type = finish_reason["type"]
+            if finish_reason["type"] == "stop":
+                finish_reason["type"] = "tool_calls"
+                finish_reason["matched"] = None
+            try:
+                tool_call_data = orjson.loads(text)
+                tool_calls = []
+                for i, tool in enumerate(tool_call_data):
+                    call_info = ToolCallItem(
+                        tool_index=i,
+                        name=tool["name"],
+                        parameters=json.dumps(tool["parameters"], ensure_ascii=False),
+                    )
+                    tool_id = self._process_tool_call_id(
+                        call_info, history_tool_calls_cnt
+                    )
+                    tool_calls.append(
+                        ToolCall(
+                            id=tool_id,
+                            index=i,
+                            function=FunctionResponse(
+                                name=tool["name"],
+                                arguments=json.dumps(
+                                    tool["parameters"], ensure_ascii=False
+                                ),
+                            ),
+                        )
+                    )
+                return ToolCallProcessingResult(tool_calls, "", finish_reason)
+            except Exception as e:
+                logger.error(f"Tool call parsing error: {e}")
+                finish_reason["type"] = original_finish_type
+                return ToolCallProcessingResult(None, text, finish_reason)
+
+        return ToolCallProcessingResult(None, text, finish_reason)
+
+    def _process_streaming_logprobs(
+        self,
+        content: Dict[str, Any],
+        n_prev_token: int,
+        total_output_logprobs: int,
+    ) -> ChoiceLogprobs:
+        """Process logprobs for streaming response"""
+        output_token_logprobs = content["meta_info"]["output_token_logprobs"]
+        output_top_logprobs = content["meta_info"].get("output_top_logprobs", [])
+        if not self.tokenizer_manager.server_args.incremental_streaming_output:
+            output_token_logprobs = output_token_logprobs[
+                n_prev_token:total_output_logprobs
+            ]
+            output_top_logprobs = output_top_logprobs[
+                n_prev_token:total_output_logprobs
+            ]
+        logprobs = to_openai_style_logprobs(
+            output_token_logprobs=output_token_logprobs,
+            output_top_logprobs=output_top_logprobs,
+        )
+
+        token_logprobs = self._process_logprobs_tokens(logprobs, use_token_index=False)
+        return ChoiceLogprobs(content=token_logprobs)
+
+    def _process_reasoning_stream(
+        self,
+        index: int,
+        delta: str,
+        reasoning_parser_dict: Dict[int, ReasoningParser],
+        content: Dict[str, Any],
+        request: ChatCompletionRequest,
+        finish_reason_type: Optional[str] = None,
+    ) -> tuple[Optional[str], str]:
+        """Process reasoning content in streaming response"""
+        if index not in reasoning_parser_dict:
+            is_force_reasoning = (
+                self.template_manager.force_reasoning
+                or self._get_reasoning_from_request(request)
+            )
+            reasoning_parser_dict[index] = ReasoningParser(
+                self.reasoning_parser,
+                request.stream_reasoning,
+                is_force_reasoning,
+                request,
+                tokenizer=self.tokenizer_manager.tokenizer,
+            )
+        reasoning_parser = reasoning_parser_dict[index]
+        reasoning_text, normal_text = reasoning_parser.parse_stream_chunk(delta)
+        if finish_reason_type is not None and finish_reason_type != "abort":
+            end_reasoning_text, end_normal_text = reasoning_parser.parse_stream_end()
+            if end_reasoning_text:
+                reasoning_text = (reasoning_text or "") + end_reasoning_text
+            if end_normal_text:
+                normal_text = (normal_text or "") + end_normal_text
+        return reasoning_text, normal_text
+
+    def _get_history_tool_calls_cnt(self, request: ChatCompletionRequest) -> int:
+        """Counts the number of tool calls in the request's message history.
+
+        NOTE: This method is only useful for models that include self-increasing
+        history tool call idx in tool calls id, such as kimi-k2
+
+        Args:
+            request: The chat completion request object.
+
+        Returns:
+            The total number of tool calls in the history, or 0 if not applicable.
+        """
+        messages = getattr(request, "messages", [])
+        idx = 0
+        for msg in messages:
+            if msg.role == "assistant":
+                tool_calls = getattr(msg, "tool_calls", None)
+                idx += len(list(tool_calls)) if tool_calls is not None else 0  # noqa
+        return idx
+
+    def _patch_reasoning_skip_special_tokens(
+        self, request: ChatCompletionRequest
+    ) -> None:
+        """Keep parser-specific reasoning markers in the decoded text.
+
+        Some reasoning parsers rely on special-token delimiters that would be
+        removed during detokenization when ``skip_special_tokens=True``.
+        """
+        if self.reasoning_parser == "apertus2509":
+            request.skip_special_tokens = False
+
+        if (
+            self.reasoning_parser in ["mistral"]
+            and request.reasoning_effort is not None
+            and request.reasoning_effort != "none"
+        ):
+            request.skip_special_tokens = False
+        elif self.reasoning_parser == "inkling":
+            request.skip_special_tokens = False
+
+    def wrap_reasoning_history(self, reasoning_text: str) -> str:
+        """Wrap prior-turn reasoning in the detector's own start/end tokens.
+
+        Pulling the delimiters from the detector keeps adapters in lockstep
+        with any future parser that ships non-``<think>`` markers — Mistral's
+        ``[THINK]``, Gemma4's ``think_start_self_label = "thought\\n"``, etc.
+        Falling back to a plain string is unsafe: it would let prior
+        thinking text reach a non-reasoning model as ordinary assistant
+        content, so the caller must surface this state, not paper over it.
+        """
+        if self._reasoning_detector is None:
+            raise ValueError(
+                "Cannot rewrap thinking history: no reasoning detector is "
+                "configured for this model"
+            )
+        d = self._reasoning_detector
+        return (
+            f"{d.think_start_token}{d.think_start_self_label}"
+            f"{reasoning_text}\n{d.think_end_token}"
+        )
+
+    def _reasoning_default_mode(self) -> Optional[str]:
+        if self._reasoning_detector is None:
+            return None
+        return self._reasoning_detector.reasoning_default
+
+    def _get_reasoning_toggle_param(self) -> Optional[str]:
+        """Resolve the chat-template kwarg that toggles reasoning, if any."""
+        config = self.template_manager.reasoning_config
+        if config is not None:
+            return config.toggle_param
+
+        mode = self._reasoning_default_mode()
+        if mode in ("thinking", "enable_thinking"):
+            return mode
+        if mode in ("explicit_thinking", "explicit_enable_thinking"):
+            return mode.replace("explicit_", "")
+        return None
+
+    def apply_reasoning_enabled(
+        self, request: ChatCompletionRequest, enabled: bool
+    ) -> None:
+        """Force the request into the requested reasoning-on/off mode.
+
+        Mirrors the read-side logic in ``_get_reasoning_from_request``;
+        the two must stay in sync. Always-on models cannot be disabled,
+        so explicit ``enabled=False`` raises rather than silently leaving
+        reasoning on.
+        """
+        if not self.reasoning_parser:
+            if enabled:
+                raise ValueError(
+                    "Anthropic thinking is not supported for models without "
+                    "a reasoning parser"
+                )
+            return
+
+        if self.reasoning_parser == "hunyuan":
+            request.reasoning_effort = "medium" if enabled else "no_think"
+            return
+
+        config = self.template_manager.reasoning_config
+        is_mistral = (config is not None and config.special_case == "mistral") or (
+            config is None and self._reasoning_default_mode() == "mistral"
+        )
+        if is_mistral:
+            request.reasoning_effort = "medium" if enabled else "none"
+            return
+
+        is_always_on = (config is not None and config.special_case == "always") or (
+            config is None and self._reasoning_default_mode() == "always"
+        )
+        if is_always_on:
+            if not enabled:
+                raise ValueError(
+                    f"Reasoning parser '{self.reasoning_parser}' is always-on "
+                    f"and cannot be disabled via Anthropic thinking"
+                )
+            return
+
+        toggle_param = self._get_reasoning_toggle_param()
+        # The read side (``_get_reasoning_from_request``) returns False
+        # whenever ``config.toggle_param is None`` OR
+        # ``config.default_enabled is None``. The write side must mirror
+        # both conditions: if ``default_enabled`` is unset we cannot
+        # actually honor an ``enabled=True`` request even when the toggle
+        # name itself is resolvable, so writing the kwarg would set up the
+        # template to emit reasoning tokens while the parser ignores them
+        # (literal ``<think>`` markers leak into the assistant text).
+        config = self.template_manager.reasoning_config
+        read_side_supported = toggle_param is not None and (
+            config is None or config.default_enabled is not None
+        )
+        if not read_side_supported:
+            if not enabled:
+                return
+            raise ValueError(
+                f"Anthropic thinking is not supported for reasoning parser "
+                f"'{self.reasoning_parser}'"
+            )
+
+        chat_template_kwargs = dict(request.chat_template_kwargs or {})
+        chat_template_kwargs[toggle_param] = enabled
+        request.chat_template_kwargs = chat_template_kwargs
+
+    def _get_reasoning_from_request(self, request: ChatCompletionRequest) -> bool:
+        """Determine whether reasoning mode should be enabled for this request.
+
+        NOTE: This is predefined based on model's chat template
+        """
+        if not self.reasoning_parser:
+            return False
+
+        if self.reasoning_parser == "minimax-m3":
+            # M3 template prefills <mm:think> for thinking_mode=enabled, so it never
+            # appears in output and reasoning must be forced. Mirrors reasoning_parser.py.
+            return (request.chat_template_kwargs or {}).get(
+                "thinking_mode"
+            ) == "enabled"
+
+        if self.reasoning_parser == "hunyuan":
+            # Hy3-preview template emits no <think> when reasoning_effort is
+            # "no_think" / "none" / unset; forcing reasoning would route all
+            # output into reasoning_content.
+            return request.reasoning_effort not in (None, "none", "no_think")
+
+        config = self.template_manager.reasoning_config
+        if config is None:
+            # Fallback to parser-level defaults when template toggle config
+            # cannot be inferred (e.g., parser-only <think> templates).
+            mode = (
+                self._reasoning_detector.reasoning_default
+                if self._reasoning_detector is not None
+                else None
+            )
+            if mode is None:
+                return False
+            if mode == "always":
+                return True
+            if mode == "mistral":
+                return (
+                    request.reasoning_effort is not None
+                    and request.reasoning_effort != "none"
+                )
+            if mode in ("thinking", "enable_thinking"):
+                return (
+                    not request.chat_template_kwargs
+                    or request.chat_template_kwargs.get(mode) is not False
+                )
+            if mode in ("explicit_thinking", "explicit_enable_thinking"):
+                toggle = mode.replace("explicit_", "")
+                return (
+                    request.chat_template_kwargs is not None
+                    and request.chat_template_kwargs.get(toggle) is True
+                )
+            logger.warning(
+                "Unknown reasoning_default mode '%s', defaulting to reasoning disabled",
+                mode,
+            )
+            return False
+
+        if config.special_case == "always":
+            return True
+
+        if config.special_case == "mistral":
+            return (
+                request.reasoning_effort is not None
+                and request.reasoning_effort != "none"
+            )
+
+        if config.toggle_param is None or config.default_enabled is None:
+            return False
+
+        if config.default_enabled:
+            return (
+                not request.chat_template_kwargs
+                or request.chat_template_kwargs.get(config.toggle_param) is not False
+            )
+        return (
+            request.chat_template_kwargs is not None
+            and request.chat_template_kwargs.get(config.toggle_param) is True
+        )
+
+    async def _process_tool_call_stream(
+        self,
+        index: int,
+        delta: str,
+        parser_dict: Dict[int, FunctionCallParser],
+        content: Dict[str, Any],
+        request: ChatCompletionRequest,
+        has_tool_calls: Dict[int, bool],
+        continuous_usage_stats: bool = False,
+    ):
+        """Process tool calls in streaming response"""
+        if index not in parser_dict:
+            is_required = request.tool_choice == "required" or isinstance(
+                request.tool_choice, ToolChoice
+            )
+            # For required/named tool choice: use JsonArrayParser when the
+            # constrained output is plain JSON (detector doesn't support
+            # structural_tag or no parser configured). Use FunctionCallParser
+            # only when the detector supports structural_tag and will produce
+            # native format output.
+            if is_required:
+                use_native_parser = False
+                if self.tool_call_parser:
+                    probe = FunctionCallParser(
+                        tools=request.tools,
+                        tool_call_parser=self.tool_call_parser,
+                        tokenizer=self.tokenizer_manager.tokenizer,
+                    )
+                    use_native_parser = probe.detector.supports_structural_tag()
+                if use_native_parser:
+                    parser_dict[index] = probe
+                else:
+                    parser_dict[index] = JsonArrayParser()
+            else:
+                parser_dict[index] = FunctionCallParser(
+                    tools=request.tools,
+                    tool_call_parser=self.tool_call_parser,
+                    tokenizer=self.tokenizer_manager.tokenizer,
+                )
+
+        parser = parser_dict[index]
+
+        # Handle both FunctionCallParser and JsonArrayParser
+        if isinstance(parser, JsonArrayParser):
+            result = parser.parse_streaming_increment(delta, request.tools)
+            normal_text, calls = result.normal_text, result.calls
+        else:
+            normal_text, calls = parser.parse_stream_chunk(delta)
+
+        # Yield normal text
+        if normal_text:
+            choice_data = ChatCompletionResponseStreamChoice(
+                index=index,
+                delta=DeltaMessage(content=normal_text),
+                finish_reason=None,
+            )
+            chunk = ChatCompletionStreamResponse(
+                id=content["meta_info"]["id"],
+                created=int(time.time()),
+                choices=[choice_data],
+                model=request.model,
+            )
+
+            # Add usage stats if continuous_usage_stats is enabled
+            if continuous_usage_stats:
+                prompt_tokens = content["meta_info"].get("prompt_tokens", 0)
+                completion_tokens = content["meta_info"].get("completion_tokens", 0)
+                reasoning_tokens = content["meta_info"].get("reasoning_tokens", 0)
+                chunk.usage = UsageProcessor.calculate_token_usage(
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                    reasoning_tokens=reasoning_tokens,
+                    cached_tokens=self._continuous_usage_cached_details(content),
+                )
+
+            yield f"data: {chunk.model_dump_json()}\n\n"
+
+        # Yield tool calls
+        history_tool_calls_cnt = self._get_history_tool_calls_cnt(request)
+        for call_item in calls:
+            # Mark that this choice has tool calls
+            has_tool_calls[index] = True
+
+            # Tool call ID should be generated only once per tool call
+            if call_item.name:
+                # First chunk: include ID and function name
+                tool_call_id = self._process_tool_call_id(
+                    call_item, history_tool_calls_cnt
+                )
+                function_name = call_item.name
+            else:
+                # Subsequent chunks: null ID and name for argument deltas
+                tool_call_id = None
+                function_name = None
+
+            tool_call = ToolCall(
+                id=tool_call_id,
+                index=call_item.tool_index,
+                function=FunctionResponse(
+                    name=function_name,
+                    arguments=call_item.parameters,
+                ),
+            )
+
+            choice_data = ChatCompletionResponseStreamChoice(
+                index=index,
+                delta=DeltaMessage(tool_calls=[tool_call]),
+                finish_reason=None,
+            )
+            chunk = ChatCompletionStreamResponse(
+                id=content["meta_info"]["id"],
+                created=int(time.time()),
+                choices=[choice_data],
+                model=request.model,
+            )
+
+            # Add usage stats if continuous_usage_stats is enabled
+            if continuous_usage_stats:
+                prompt_tokens = content["meta_info"].get("prompt_tokens", 0)
+                completion_tokens = content["meta_info"].get("completion_tokens", 0)
+                reasoning_tokens = content["meta_info"].get("reasoning_tokens", 0)
+                chunk.usage = UsageProcessor.calculate_token_usage(
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                    reasoning_tokens=reasoning_tokens,
+                    cached_tokens=self._continuous_usage_cached_details(content),
+                )
+
+            yield f"data: {chunk.model_dump_json()}\n\n"
+
+    def _check_for_unstreamed_tool_args(
+        self,
+        parser: Union[FunctionCallParser, JsonArrayParser],
+        content: Dict[str, Any],
+        request: ChatCompletionRequest,
+        index: int,
+    ) -> Optional[str]:
+        """
+        Check for any remaining tool call arguments that need to be streamed
+        when generation finishes. This ensures tool calls are properly completed
+        even if the model generates the final arguments in the last chunk.
+        """
+        # Get the detector - either from FunctionCallParser or directly if json detector
+        detector = parser.detector if hasattr(parser, "detector") else parser
+
+        # Only check if we have tool calls and the detector has tracked data
+        if (
+            not hasattr(detector, "prev_tool_call_arr")
+            or not detector.prev_tool_call_arr
+        ):
+            return None
+
+        if (
+            not hasattr(detector, "streamed_args_for_tool")
+            or not detector.streamed_args_for_tool
+        ):
+            return None
+
+        # Get the last tool call that was being processed
+        tool_index = len(detector.prev_tool_call_arr) - 1
+        if tool_index < 0 or tool_index >= len(detector.streamed_args_for_tool):
+            return None
+
+        # Get expected vs actual arguments
+        expected_args = detector.prev_tool_call_arr[tool_index].get("arguments", {})
+        if isinstance(expected_args, str):
+            expected_call = expected_args
+        else:
+            expected_call = json.dumps(expected_args, ensure_ascii=False)
+        actual_call = detector.streamed_args_for_tool[tool_index]
+
+        # Check if there are remaining arguments to send
+        remaining_call = (
+            expected_call[len(actual_call) :]
+            if expected_call.startswith(actual_call)
+            else ""
+        )
+
+        if remaining_call:
+            # Create tool call chunk with remaining arguments
+            tool_call = ToolCall(
+                id=None,  # No ID for argument deltas
+                index=tool_index,
+                function=FunctionResponse(
+                    name=None,  # No name for argument deltas
+                    arguments=remaining_call,
+                ),
+            )
+
+            choice_data = ChatCompletionResponseStreamChoice(
+                index=index,
+                delta=DeltaMessage(tool_calls=[tool_call]),
+                finish_reason=None,  # Don't send finish_reason with this chunk
+            )
+
+            chunk = ChatCompletionStreamResponse(
+                id=content["meta_info"]["id"],
+                created=int(time.time()),
+                choices=[choice_data],
+                model=request.model,
+            )
+
+            return f"data: {chunk.model_dump_json()}\n\n"
+
+        return None

sglang_patch/srt/function_call/function_call_parser.py

@@ -0,0 +1,286 @@
+import inspect
+import logging
+from typing import Dict, List, Literal, Optional, Set, Tuple, Type, Union
+
+from sglang.srt.entrypoints.openai.protocol import (
+    LegacyStructuralTagResponseFormat,
+    StructuralTagResponseFormat,
+    StructuresResponseFormat,
+    Tool,
+    ToolCallConstraint,
+    ToolChoice,
+)
+from sglang.srt.environ import ToolStrictLevel, envs
+from sglang.srt.function_call.apertus2509_detector import Apertus2509Detector
+from sglang.srt.function_call.base_format_detector import BaseFormatDetector
+from sglang.srt.function_call.cohere_command4_detector import CohereCommand4Detector
+from sglang.srt.function_call.core_types import ToolCallItem
+from sglang.srt.function_call.deepseekv3_detector import DeepSeekV3Detector
+from sglang.srt.function_call.deepseekv4_detector import DeepSeekV4Detector
+from sglang.srt.function_call.deepseekv31_detector import DeepSeekV31Detector
+from sglang.srt.function_call.deepseekv32_detector import DeepSeekV32Detector
+from sglang.srt.function_call.gemma4_detector import Gemma4Detector
+from sglang.srt.function_call.gigachat3_detector import GigaChat3Detector
+from sglang.srt.function_call.glm4_moe_detector import Glm4MoeDetector
+from sglang.srt.function_call.glm47_moe_detector import Glm47MoeDetector
+from sglang.srt.function_call.gpt_oss_detector import GptOssDetector
+from sglang.srt.function_call.hermes_detector import HermesDetector
+from sglang.srt.function_call.hunyuan_detector import HunyuanDetector
+from sglang.srt.function_call.inkling_detector import InklingDetector
+from sglang.srt.function_call.internlm_detector import InternlmDetector
+from sglang.srt.function_call.kimik2_detector import KimiK2Detector
+from sglang.srt.function_call.lfm2_detector import Lfm2Detector
+from sglang.srt.function_call.llama32_detector import Llama32Detector
+from sglang.srt.function_call.mimo_detector import MiMoDetector
+from sglang.srt.function_call.minicpm5_detector import MiniCPM5Detector
+from sglang.srt.function_call.minimax_m2 import MinimaxM2Detector
+from sglang.srt.function_call.minimax_m3 import MinimaxM3Detector
+from sglang.srt.function_call.mistral_detector import MistralDetector
+from sglang.srt.function_call.poolside_v1_detector import PoolsideV1Detector
+from sglang.srt.function_call.pythonic_detector import PythonicDetector
+from sglang.srt.function_call.qwen3_coder_detector import Qwen3CoderDetector
+from sglang.srt.function_call.qwen25_detector import Qwen25Detector
+from sglang.srt.function_call.step3_detector import Step3Detector
+from sglang.srt.function_call.trinity_detector import TrinityDetector
+from sglang.srt.function_call.utils import (
+    _get_tool_schema_defs,
+    get_json_schema_constraint,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class FunctionCallParser:
+    """
+    Parser for function/tool calls in model outputs.
+
+    This class handles both streaming and non-streaming parsing of function calls using a detector.
+    In streaming scenarios, each time new_text is received, it calls detector.parse_streaming_increment
+    and returns the resulting normal_text and calls to the upper layer (or SSE).
+    """
+
+    ToolCallParserEnum: Dict[str, Type[BaseFormatDetector]] = {
+        "apertus2509": Apertus2509Detector,
+        "cohere_command4": CohereCommand4Detector,
+        "deepseekv3": DeepSeekV3Detector,
+        "deepseekv31": DeepSeekV31Detector,
+        "deepseekv32": DeepSeekV32Detector,
+        "agnes": DeepSeekV4Detector,  # alias: Agnes uses the deepseekv4 format
+        "deepseekv4": DeepSeekV4Detector,
+        "glm": Glm4MoeDetector,
+        "glm45": Glm4MoeDetector,
+        "glm47": Glm47MoeDetector,
+        "gpt-oss": GptOssDetector,
+        "kimi_k2": KimiK2Detector,
+        "lfm2": Lfm2Detector,
+        "llama3": Llama32Detector,
+        "mimo": MiMoDetector,
+        "minicpm5": MiniCPM5Detector,
+        "mistral": MistralDetector,
+        "poolside_v1": PoolsideV1Detector,
+        "pythonic": PythonicDetector,
+        "qwen": Qwen25Detector,
+        "qwen25": Qwen25Detector,
+        "qwen3_coder": Qwen3CoderDetector,
+        "step3": Step3Detector,
+        "step3p5": Qwen3CoderDetector,
+        "minimax-m2": MinimaxM2Detector,
+        "minimax-m3": MinimaxM3Detector,
+        "trinity": TrinityDetector,
+        "interns1": InternlmDetector,
+        "hermes": HermesDetector,
+        "hunyuan": HunyuanDetector,
+        "gigachat3": GigaChat3Detector,
+        "gemma4": Gemma4Detector,
+        "inkling": InklingDetector,
+    }
+
+    def __init__(self, tools: List[Tool], tool_call_parser: str, tokenizer=None):
+        detector_class = self.ToolCallParserEnum.get(tool_call_parser)
+        if detector_class:
+            kwargs = {}
+            if tokenizer is not None:
+                sig = inspect.signature(detector_class)
+                if "tokenizer" in sig.parameters:
+                    kwargs["tokenizer"] = tokenizer
+            detector = detector_class(**kwargs)
+        else:
+            raise ValueError(f"Unsupported tool_call_parser: {tool_call_parser}")
+
+        self.detector = detector
+        self.tools = tools
+        self.tool_strict_level = envs.SGLANG_TOOL_STRICT_LEVEL.get()
+
+    def has_tool_call(self, text: str) -> bool:
+        """
+        Check if the given text contains a tool call in the format supported by this parser.
+        This delegates to the detector's implementation.
+
+        Args:
+            text: The text to check for tool calls
+
+        Returns:
+            True if the text contains a tool call, False otherwise
+        """
+        if not self.tools:
+            return False
+        return self.detector.has_tool_call(text)
+
+    def parse_non_stream(self, full_text: str) -> Tuple[str, list[ToolCallItem]]:
+        """
+        One-time parsing of the full text to extract tool calls.
+
+        Args:
+            full_text: The complete text to parse
+
+        Returns:
+            A tuple containing:
+            - The remaining text after parsing that was not consumed by the detector (can be treated as normal text)
+            - A list of tool calls parsed from the text
+        """
+        if not self.tools:
+            return full_text, []
+        has_tool_call = self.detector.has_tool_call(full_text)
+        parsed_result = self.detector.detect_and_parse(full_text, self.tools)
+        tool_call_list = parsed_result.calls
+        if tool_call_list or has_tool_call:
+            return parsed_result.normal_text, tool_call_list
+        else:
+            return full_text, []
+
+    def parse_stream_chunk(self, chunk_text: str) -> Tuple[str, list[ToolCallItem]]:
+        """
+        Streaming incremental parsing of chunks of text as they arrive.
+
+        Args:
+            chunk_text: The new chunk of text to parse
+
+        Returns:
+            A tuple containing:
+            - The normal text that should be displayed to the user
+            - A list of tool calls parsed from the chunk
+        """
+        if not self.tools:
+            return chunk_text, []
+        final_normal_text = ""
+        final_calls = []
+
+        sp_result = self.detector.parse_streaming_increment(chunk_text, self.tools)
+        if sp_result.normal_text:
+            final_normal_text = sp_result.normal_text
+        if sp_result.calls:
+            final_calls.extend(sp_result.calls)
+            final_normal_text = sp_result.normal_text
+
+        return final_normal_text, final_calls
+
+    def get_legacy_structural_tag(
+        self, at_least_one: bool = False
+    ) -> StructuralTagResponseFormat:
+        """
+        Generate a structural tag response format for all available tools.
+
+        This creates the necessary structural tags that guide the model's output format.
+
+        Args:
+            at_least_one: If True, the grammar forces at least one tool call
+                (no free text allowed). Used for required/named tool_choice.
+
+        Raises:
+            ValueError: If tools have conflicting $defs schemas.
+        """
+        # Validate $defs consistency before building structural tags
+        _get_tool_schema_defs(self.tools)
+
+        tool_structures: List[StructuresResponseFormat] = list()
+        tool_trigger_set: Set[str] = set()
+
+        get_structure_info = self.detector.structure_info()
+        for tool in self.tools:
+            function = tool.function
+            name = function.name
+            assert name is not None
+            info = get_structure_info(name)
+
+            # accept all if not strict, otherwise only accept the schema
+            is_strict = (
+                function.strict or self.tool_strict_level >= ToolStrictLevel.PARAMETER
+            )
+            schema = function.parameters if is_strict else {}
+
+            tool_structures.append(
+                StructuresResponseFormat(
+                    begin=info.begin,
+                    schema=schema or {},  # type: ignore
+                    end=info.end,
+                )
+            )
+            tool_trigger_set.add(info.trigger)
+
+        # TODO(dark): move this into new structural tag format
+        # This requires all grammar backend support the new format
+        return LegacyStructuralTagResponseFormat(
+            type="structural_tag",
+            structures=tool_structures,
+            triggers=list(tool_trigger_set),
+            at_least_one=at_least_one,
+        )
+
+    def get_structure_constraint(
+        self,
+        tool_choice: Union[ToolChoice, Literal["auto", "required"]],
+        parallel_tool_calls: bool = True,
+        thinking_mode: bool = False,
+    ) -> Optional[ToolCallConstraint]:
+        """
+        Returns the appropriate structure constraint for tool calls based on the tool_choice.
+        The constraint is used to guide the model's output format.
+
+        Args:
+            tool_choice: The tool choice setting from the request
+
+        Returns:
+            A tuple of (constraint_type, constraint_value) to be added to sampling parameters,
+            or None if no constraint applies.
+        """
+        is_required = tool_choice == "required" or isinstance(tool_choice, ToolChoice)
+        should_constrain_auto = tool_choice == "auto" and (
+            any(tool.function.strict for tool in self.tools)
+            or self.tool_strict_level >= ToolStrictLevel.FUNCTION
+        )
+
+        # Highest priority: model-native structural_tag when available.
+        try:
+            if tool_choice == "auto" and not should_constrain_auto:
+                structural_tag = self.detector.get_auto_tool_call_structural_tag(
+                    tools=self.tools
+                )
+                if structural_tag is not None:
+                    return ("structural_tag", structural_tag)
+
+            if is_required or should_constrain_auto:
+                structural_tag = self.detector.get_structural_tag(
+                    tools=self.tools,
+                    thinking_mode=thinking_mode,
+                    tool_choice=tool_choice,
+                )
+                if structural_tag is not None:
+                    return ("structural_tag", structural_tag)
+
+                # Fallback to legacy structural tag if model-native tag is not supported.
+                if self.detector.supports_structural_tag():
+                    # For "required"/named: always use structural_tag to preserve the
+                    # model's native tool call format. Schema is only included when
+                    # strict=True, per OpenAI protocol semantics.
+                    # For "auto": only constrain when strict is enabled.
+                    tag = self.get_legacy_structural_tag(at_least_one=is_required)
+                    return ("structural_tag", tag)
+
+            if tool_choice == "required" or isinstance(tool_choice, ToolChoice):
+                json_schema = get_json_schema_constraint(
+                    self.tools, tool_choice, parallel_tool_calls=parallel_tool_calls
+                )
+                return ("json_schema", json_schema)
+        except Exception as e:
+            logger.error(f"Error getting structure constraint: {e}")
+            return None

sglang_patch/srt/layers/attention/dsv4/compressor.py

@@ -0,0 +1,480 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Literal, NamedTuple, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from sglang.jit_kernel.dsv4 import linear_bf16_fp32, triton_create_paged_compress_data
+from sglang.jit_kernel.dsv4.compress_old import (
+    CompressorDecodePlan,
+    CompressorPrefillPlan,
+    compress_forward,
+    compress_fused_norm_rope_inplace,
+)
+from sglang.kernels.ops.attention.dsa.triton_kernel import act_quant
+from sglang.kernels.ops.attention.dsv4.quant_k_cache import (
+    quant_to_nope_fp8_rope_bf16_pack_triton,
+)
+from sglang.srt.configs.agnes import AgnesConfig
+from sglang.srt.environ import envs
+from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
+from sglang.srt.layers.layernorm import RMSNorm
+from sglang.srt.layers.linear import ReplicatedLinear
+from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output
+from sglang.srt.layers.utils.multi_platform import MultiPlatformOp
+from sglang.srt.mem_cache.deepseek_v4_compress_state import (
+    CompressStatePool,
+)
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+from sglang.srt.model_executor.forward_context import get_attn_backend
+from sglang.srt.models.deepseek_v2 import _is_hip
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils import add_prefix, is_npu, set_weight_attrs
+
+_is_npu = is_npu()
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+    from sglang.srt.layers.attention.deepseek_v4_backend import DeepseekV4AttnBackend
+    from sglang.srt.layers.rotary_embedding import RotaryEmbedding
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class FusedCompressMetadata(NamedTuple):
+    write_loc: torch.Tensor
+    extra_data: Optional[torch.Tensor]
+    plan: Union[CompressorDecodePlan, CompressorPrefillPlan]
+
+    def copy_(self, other: FusedCompressMetadata) -> None:
+        from .metadata import maybe_copy_inplace
+
+        self.write_loc.copy_(other.write_loc)
+        maybe_copy_inplace(self.extra_data, src=other.extra_data)
+        self.plan.copy_(other.plan)
+
+
+class CompressorBackendMixin:
+    def get_paged_compress_metadata(self, compress_ratio: int) -> FusedCompressMetadata:
+        attr_name = f"c{compress_ratio}_compress_metadata"
+        metadata = getattr(self.forward_metadata, attr_name)
+        assert isinstance(metadata, FusedCompressMetadata)
+        return metadata
+
+    def forward_compress(
+        self,
+        *,
+        kv_score_buffer: torch.Tensor,
+        kv_score_input: torch.Tensor,
+        ape: torch.Tensor,
+        head_dim: int,
+        norm: RMSNorm,
+        freqs_cis_cache: torch.Tensor,
+        rotate: bool,
+        forward_batch: ForwardBatch,
+        compress_ratio: int,
+        is_paged: bool = False,
+    ) -> torch.Tensor:
+        from sglang.srt.layers.attention.dsa.dsa_indexer import rotate_activation
+
+        assert compress_ratio in (
+            4,
+            128,
+        ), f"DSV4 supports CSA(4x) and HCA(128x) only, got {compress_ratio=}"
+        if is_paged:
+            metadata = self.get_paged_compress_metadata(compress_ratio)
+            coff = 2 if is_overlap_compress(compress_ratio) else 1
+            if compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
+                kv_score_buffer = kv_score_buffer.view(-1, 1, head_dim * 3)
+            else:
+                last_dim = 2 * head_dim * coff
+                assert kv_score_buffer.shape[-1] == last_dim
+                kv_score_buffer = kv_score_buffer.view(-1, compress_ratio, last_dim)
+        else:
+            plan = make_compressor_plan(compress_ratio, forward_batch)
+            metadata = (forward_batch.req_pool_indices.to(torch.int32), None, plan)
+        indices, extra_data, plan = metadata
+
+        if _is_hip:
+            if not is_paged:
+                raise NotImplementedError("HIP fused compressor expects paged metadata")
+
+            from sglang.kernels.ops.attention.dsv4.fused_compress_triton import (
+                hip_compress_forward,
+                hip_compress_fused_norm_rope_hadamard_inplace,
+                hip_compress_fused_norm_rope_inplace,
+            )
+
+            kv_compressed = hip_compress_forward(
+                kv_score_buffer=kv_score_buffer,
+                kv_score_input=kv_score_input,
+                ape=ape,
+                indices=indices,
+                plan=plan,
+                compress_ratio=compress_ratio,
+                head_dim=head_dim,
+                extra_data=extra_data,
+            )
+            norm_eps = (
+                norm.variance_epsilon if hasattr(norm, "variance_epsilon") else norm.eps
+            )
+            if rotate:
+                hip_compress_fused_norm_rope_hadamard_inplace(
+                    kv_compressed,
+                    norm.weight,
+                    norm_eps,
+                    freqs_cis_cache,
+                    plan,
+                    head_dim,
+                )
+            else:
+                hip_compress_fused_norm_rope_inplace(
+                    kv_compressed,
+                    norm.weight,
+                    norm_eps,
+                    freqs_cis_cache,
+                    plan,
+                )
+            return kv_compressed
+
+        kv_compressed = compress_forward(
+            kv_score_buffer=kv_score_buffer,
+            kv_score_input=kv_score_input,
+            ape=ape,
+            indices=indices,
+            plan=plan,
+            compress_ratio=compress_ratio,
+            head_dim=head_dim,
+            extra_data=extra_data,
+        )
+        compress_fused_norm_rope_inplace(
+            kv_compressed,
+            norm.weight,
+            norm.variance_epsilon,
+            freqs_cis_cache,
+            plan,
+        )
+        return rotate_activation(kv_compressed) if rotate else kv_compressed
+
+    def forward_core_compressor(
+        self,
+        x: torch.Tensor,
+        forward_batch: ForwardBatch,
+        layer_id: int,
+        compressor: Compressor,
+    ) -> None:
+        if forward_batch.forward_mode.is_idle():
+            return
+        token_to_kv_pool = self.token_to_kv_pool
+        if TYPE_CHECKING:
+            assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
+
+        new_compressed_kv = compressor(x, forward_batch, attn_backend=self)
+        core_metadata = self.forward_metadata.core_metadata
+        out_loc = (
+            core_metadata.c4_out_loc
+            if compressor.ratio == 4
+            else core_metadata.c128_out_loc
+        )
+        if out_loc.shape[0] > new_compressed_kv.shape[0]:
+            out_loc = out_loc[: new_compressed_kv.shape[0]]
+        if envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get():
+            token_to_kv_pool.set_extra_key_buffer_fused(
+                layer_id=layer_id,
+                loc=out_loc,
+                cache_k=new_compressed_kv,
+            )
+        else:
+            pack = quant_to_nope_fp8_rope_bf16_pack_triton(new_compressed_kv.bfloat16())
+            token_to_kv_pool.set_extra_key_buffer(layer_id, out_loc, pack)
+
+    def forward_indexer_compressor(
+        self,
+        x: torch.Tensor,
+        forward_batch: ForwardBatch,
+        layer_id: int,
+        compressor: Compressor,
+    ) -> None:
+        assert is_overlap_compress(compressor.ratio)
+        token_to_kv_pool = self.token_to_kv_pool
+        if TYPE_CHECKING:
+            assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
+
+        new_compressed_kv = compressor(x, forward_batch, attn_backend=self)
+        out_loc = self.forward_metadata.core_metadata.c4_out_loc
+        if out_loc.shape[0] > new_compressed_kv.shape[0]:
+            out_loc = out_loc[: new_compressed_kv.shape[0]]
+        if self.enable_deepseek_v4_fp4_indexer:
+            token_to_kv_pool.set_index_k_fp4(
+                layer_id=layer_id,
+                loc=out_loc,
+                cache_k=new_compressed_kv,
+            )
+        elif envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get():
+            token_to_kv_pool.set_index_k_fused(
+                layer_id=layer_id,
+                loc=out_loc,
+                cache_k=new_compressed_kv,
+            )
+        else:
+            new_compressed_kv_fp8, new_compressed_kv_scale = act_quant(
+                new_compressed_kv
+            )
+            token_to_kv_pool.set_index_k_scale_buffer(
+                layer_id=layer_id,
+                loc=out_loc,
+                index_k=new_compressed_kv_fp8,
+                index_k_scale=new_compressed_kv_scale,
+            )
+
+
+def is_overlap_compress(compress_ratio: int) -> bool:
+    return compress_ratio == 4
+
+
+def make_compressor_plan(
+    compress_ratio: Literal[4, 128],
+    forward_batch: ForwardBatch,
+) -> Union[CompressorDecodePlan, CompressorPrefillPlan]:
+    if forward_batch.forward_mode.is_decode():
+        seq_lens_32 = forward_batch.seq_lens.to(torch.int32)
+        return CompressorDecodePlan(compress_ratio, seq_lens_32)
+    if forward_batch.forward_mode.is_prefill():
+        assert not forward_batch.forward_mode.is_target_verify()
+        extend_lens_list = forward_batch.extend_seq_lens_cpu
+        seq_lens_cpu = forward_batch.seq_lens_cpu
+        assert extend_lens_list is not None and seq_lens_cpu is not None
+        return CompressorPrefillPlan.generate(
+            compress_ratio=compress_ratio,
+            num_q_tokens=sum(extend_lens_list),
+            seq_lens=seq_lens_cpu,
+            extend_lens=torch.tensor(extend_lens_list),
+            device=forward_batch.seq_lens.device,
+        )
+    elif forward_batch.forward_mode.is_target_verify():
+        raise NotImplementedError("target verify mode to be implemented")
+    else:
+        raise NotImplementedError(f"unsupported mode {forward_batch.forward_mode=}")
+
+
+def create_paged_compressor_data(
+    compress_ratio: Literal[4, 128],
+    *,
+    is_prefill: bool,
+    token_to_kv_pool: DeepSeekV4TokenToKVPool,
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    seq_lens: torch.Tensor,
+    extend_lens: Optional[torch.Tensor] = None,
+    seq_lens_cpu: Optional[List[int]] = None,
+    extend_lens_cpu: Optional[List[int]] = None,
+    use_prefill_cuda_graph: bool = False,
+    num_q_tokens: Optional[int] = None,
+) -> FusedCompressMetadata:
+    swa_page_size = token_to_kv_pool.swa_page_size
+    ring_size = token_to_kv_pool.get_ring_size(compress_ratio=compress_ratio)
+    # assert ring_size % compress_ratio == 0
+
+    def clip_down(positions: torch.Tensor) -> torch.Tensor:
+        return positions // compress_ratio * compress_ratio
+
+    def get_raw_loc(positions: torch.Tensor) -> torch.Tensor:
+        positions = positions.masked_fill(positions < 0, 0)
+        if compress_ratio == 128:
+            state_loc = req_pool_indices * ring_size + positions % ring_size
+        else:
+            loc = req_to_token[req_pool_indices, positions]
+            swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(loc)
+            swa_pages = swa_loc // swa_page_size
+            state_loc = swa_pages * ring_size + swa_loc % ring_size
+        return (state_loc // compress_ratio).to(torch.int32)
+
+    is_overlap = is_overlap_compress(compress_ratio)
+
+    if is_prefill:
+        assert extend_lens is not None
+        write_loc, extra_data = triton_create_paged_compress_data(
+            compress_ratio=compress_ratio,
+            is_overlap=is_overlap,
+            swa_page_size=swa_page_size,
+            ring_size=ring_size,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            extend_seq_lens=extend_lens,
+            req_to_token=req_to_token,
+            full_to_swa_index_mapping=token_to_kv_pool.full_to_swa_index_mapping,
+        )
+
+        plan_kwargs: dict
+        if seq_lens_cpu is None:
+            assert num_q_tokens is not None
+            plan_kwargs = dict(
+                num_q_tokens=num_q_tokens,
+                seq_lens=seq_lens,
+                extend_lens=extend_lens,
+            )
+        else:
+            assert extend_lens_cpu is not None
+            plan_kwargs = dict(
+                num_q_tokens=sum(extend_lens_cpu),
+                seq_lens=torch.tensor(seq_lens_cpu),
+                extend_lens=torch.tensor(extend_lens_cpu),
+            )
+        plan = CompressorPrefillPlan.generate(
+            compress_ratio=compress_ratio,
+            device=seq_lens.device,
+            use_cuda_graph=use_prefill_cuda_graph,
+            **plan_kwargs,
+        )
+    else:
+        write_positions = clip_down(seq_lens - 1)
+        write_loc = get_raw_loc(write_positions)
+        if is_overlap:
+            write_overlap_loc = get_raw_loc(write_positions - compress_ratio)
+            extra_data = write_overlap_loc.view(-1, 1)
+        elif _is_hip:
+            extra_data = get_raw_loc(write_positions - compress_ratio)
+        else:
+            extra_data = None
+        plan = CompressorDecodePlan(compress_ratio, seq_lens.to(torch.int32))
+
+    return FusedCompressMetadata(write_loc=write_loc, extra_data=extra_data, plan=plan)
+
+
+class Compressor(MultiPlatformOp):
+    def __init__(
+        self,
+        config: AgnesConfig,
+        layer_id: int,
+        is_in_indexer: bool,
+        freqs_cis: torch.Tensor,
+        compress_ratio: Literal[0, 4, 128],
+        head_dim: int,
+        rotate: bool = False,
+        prefix: str = "",
+        rotary_emb: Optional[RotaryEmbedding] = None,
+    ) -> None:
+        super().__init__()
+        self.layer_id = layer_id
+        self.is_in_indexer = is_in_indexer
+        self.dim = config.hidden_size
+        self.head_dim = head_dim
+        self.rope_head_dim = getattr(config, "qk_rope_head_dim", 64)
+        assert compress_ratio != 0, "compress_ratio should not be 0"
+        self.ratio = compress_ratio
+        self.overlap = self.ratio == 4
+        self.rotate = rotate
+        self.coff = coff = 1 + self.overlap
+
+        self.ape = nn.Parameter(
+            torch.empty(self.ratio, coff * self.head_dim, dtype=torch.float32)
+        )
+        set_weight_attrs(self.ape, {"weight_loader": self.load_ape_weight})
+        wkv_gate_dtype = torch.bfloat16
+        self.wkv_gate = ReplicatedLinear(
+            self.dim,
+            2 * coff * self.head_dim,
+            bias=False,
+            quant_config=None,
+            prefix=add_prefix("wkv_gate", prefix),
+            params_dtype=wkv_gate_dtype,
+        )
+        self.norm = RMSNorm(
+            self.head_dim, eps=config.rms_norm_eps, weight_dtype=torch.float32
+        )
+        self.rotary_emb = rotary_emb
+        self.freqs_cis = freqs_cis
+
+        self.ape_converted = False
+
+    def _apply_ape_hotfix(self):
+        self.ape_converted = True
+
+        if _is_npu:
+            return
+
+        if self.overlap:
+            ape = torch.chunk(self.ape.data, 2, dim=-1)
+            ape = torch.cat([ape[0], ape[1]], dim=0)
+            self.ape.data.copy_(ape.view(self.ratio, -1))
+
+    def apply_ape_hotfix(self):
+        assert not self.ape_converted
+        self._apply_ape_hotfix()
+
+    def load_ape_weight(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        assert param is self.ape
+        assert loaded_weight.shape == param.shape
+        param.data.copy_(loaded_weight)
+        self._apply_ape_hotfix()
+
+    def get_state_pool(self, attn_backend: AttentionBackend) -> CompressStatePool:
+        token_to_kv_pool = attn_backend.token_to_kv_pool
+        assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
+        if self.is_in_indexer:
+            ret = token_to_kv_pool.get_indexer_compress_states(self.layer_id)
+        else:
+            ret = token_to_kv_pool.get_attention_compress_states(self.layer_id)
+        assert isinstance(ret, CompressStatePool)
+        return ret
+
+    def compute_kv_score(self, x: torch.Tensor, forward_batch: ForwardBatch):
+        kv_score = linear_bf16_fp32(x, self.wkv_gate.weight)
+
+        # CUDA path: delegate to backend
+        if dsa_use_prefill_cp(forward_batch):
+            kv_score = cp_all_gather_rerange_output(
+                kv_score,
+                get_parallel().attn_cp_size,
+                forward_batch,
+                torch.cuda.current_stream(),
+            )
+        return kv_score
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+        forward_batch: ForwardBatch,
+        attn_backend: Optional[AttentionBackend] = None,
+    ) -> torch.Tensor:
+        if forward_batch.forward_mode.is_idle():
+            assert x.shape[0] == 0
+            return x.new_empty(0, self.head_dim)
+
+        kv_score = self.compute_kv_score(x, forward_batch)
+
+        if TYPE_CHECKING:
+            assert isinstance(attn_backend, DeepseekV4AttnBackend)
+        kv_score_buffer = self.get_state_pool(attn_backend).kv_score_buffer.kv_score
+        return attn_backend.forward_compress(
+            kv_score_buffer=kv_score_buffer,
+            kv_score_input=kv_score,
+            ape=self.ape.view(-1, self.head_dim),
+            head_dim=self.head_dim,
+            norm=self.norm,
+            freqs_cis_cache=self.freqs_cis,
+            rotate=self.rotate,
+            compress_ratio=self.ratio,
+            forward_batch=forward_batch,
+            is_paged=True,
+        )
+
+    def forward_npu(
+        self,
+        x: torch.Tensor,
+        forward_batch: ForwardBatch,
+        attn_backend: Optional[AttentionBackend] = None,
+    ) -> torch.Tensor:
+        if forward_batch.forward_mode.is_idle():
+            assert x.shape[0] == 0
+            return x.new_empty(0, self.head_dim)
+
+        if dsa_use_prefill_cp(forward_batch):
+            x = cp_all_gather_rerange_output(
+                x,
+                get_parallel().attn_cp_size,
+                forward_batch,
+                torch.cuda.current_stream(),
+            )
+
+        return get_attn_backend().forward_compress(self, x, forward_batch)

sglang_patch/srt/layers/attention/dsv4/indexer.py

@@ -0,0 +1,887 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple, TypeAlias, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.jit_kernel.dsv4 import (
+    fused_q_indexer_rope_hadamard_fp4_quant,
+    fused_q_indexer_rope_hadamard_quant,
+    topk_transform_512,
+    topk_transform_512_v2,
+)
+from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
+from sglang.srt.configs.agnes import AgnesConfig
+from sglang.srt.environ import envs
+from sglang.srt.layers.attention.dsv4.compressor import Compressor
+from sglang.srt.layers.attention.dsv4.metadata import (
+    NonPagedIndexerPlan,
+    PagedIndexerMetadata,
+)
+from sglang.srt.layers.linear import ReplicatedLinear
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+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 (
+    is_in_tc_piecewise_cuda_graph,
+)
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
+from sglang.srt.utils import add_prefix, is_cuda, is_hip, is_xpu
+from sglang.srt.utils.common import is_sm120_supported
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+    from sglang.srt.layers.attention.dsv4.compressor import (
+        CompressorBackendMixin,
+    )
+    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
+
+
+FP8_DTYPE = torch.float8_e4m3fnuz if is_fp8_fnuz() else torch.float8_e4m3fn
+
+
+IndexerQuery: TypeAlias = Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
+
+
+_arange_cache = {}
+
+
+def fp8_paged_mqa_logits_torch(
+    q_fp8: torch.Tensor,
+    kvcache_fp8: torch.Tensor,
+    weight: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_table: torch.Tensor,
+    deep_gemm_metadata: Any,
+    max_seq_len: int,
+    clean_logits: bool = True,
+) -> torch.Tensor:
+    """Vectorized implementation compatible with CUDA graph capture."""
+    _ = deep_gemm_metadata
+    batch_size, _, num_heads, head_dim = q_fp8.shape
+    block_size = kvcache_fp8.shape[1]
+
+    assert head_dim == 128
+    assert block_size == 64
+    assert q_fp8.shape == (batch_size, 1, num_heads, head_dim)
+    assert kvcache_fp8.shape[1:] == (block_size, 1, head_dim + 4)
+    assert weight.shape == (batch_size, num_heads)
+    assert seq_lens.shape == (batch_size,)
+    assert page_table.shape[0] == batch_size
+    assert clean_logits == False
+
+    max_num_pages = page_table.shape[1]
+    SCALE_OFFSET = block_size * head_dim
+    total_dim = block_size * (head_dim + 4)
+
+    kvcache_flat = kvcache_fp8.view(-1, total_dim)
+
+    pages_clamped = page_table.clamp(min=0)
+    kvcache_gathered = kvcache_flat[pages_clamped]
+
+    kv_values_raw = kvcache_gathered[..., :SCALE_OFFSET].contiguous()
+    kv_values_fp8 = kv_values_raw.view(dtype=FP8_DTYPE)
+    kv_values = kv_values_fp8.to(torch.bfloat16)
+    kv_values = kv_values.reshape(batch_size, max_num_pages * block_size, head_dim)
+
+    kv_scales_raw = kvcache_gathered[..., SCALE_OFFSET:].contiguous()
+    kv_scales = kv_scales_raw.view(dtype=torch.float32)
+    kv_scales = kv_scales.reshape(batch_size, max_num_pages * block_size)
+
+    q_float = q_fp8[:, 0].to(torch.bfloat16)
+    scores = torch.bmm(kv_values, q_float.transpose(1, 2))
+    scores = F.relu(scores)
+    scores = scores * weight.unsqueeze(1)
+    scores = scores.sum(dim=2)
+    scores = scores * kv_scales
+
+    padded_seq_len = max_num_pages * block_size
+    cache = _arange_cache
+    arange_key = f"arange_{padded_seq_len}_{scores.device}"
+    if arange_key not in cache:
+        cache[arange_key] = torch.arange(padded_seq_len, device=scores.device)
+    positions = cache[arange_key].unsqueeze(0)
+    valid_mask = positions < seq_lens.unsqueeze(1)
+    scores = scores.masked_fill(~valid_mask, 0.0)
+
+    if padded_seq_len < max_seq_len:
+        scores = F.pad(scores, (0, max_seq_len - padded_seq_len), value=0.0)
+    else:
+        scores = scores[:, :max_seq_len]
+
+    return scores
+
+
+def _aiter_fp8_paged_mqa_logits(
+    q_fp8: torch.Tensor,
+    kvcache_fp8: torch.Tensor,
+    weight: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_table: torch.Tensor,
+    deep_gemm_metadata: Any,
+    max_seq_len: int,
+    clean_logits: bool = False,
+) -> torch.Tensor:
+    """Wrapper adapting aiter's deepgemm_fp8_paged_mqa_logits to SGLang's interface."""
+    from aiter.ops.triton.attention.pa_mqa_logits import (
+        deepgemm_fp8_paged_mqa_logits,
+    )
+
+    batch_size = q_fp8.shape[0]
+    next_n = q_fp8.shape[1]
+    total_tokens = batch_size * next_n
+    _sl = seq_lens.squeeze(-1) if seq_lens.dim() == 2 else seq_lens
+    kv_block_size = kvcache_fp8.shape[1]
+    logits = torch.empty(
+        total_tokens,
+        max_seq_len,
+        dtype=torch.float32,
+        device=q_fp8.device,
+    )
+    deepgemm_fp8_paged_mqa_logits(
+        q_fp8,
+        kvcache_fp8,
+        weight,
+        logits,
+        _sl.to(torch.int32),
+        page_table.to(torch.int32),
+        max_seq_len,
+        KVBlockSize=kv_block_size,
+        Preshuffle=True,
+    )
+    return logits
+
+
+def fp8_paged_mqa_logits_torch_sm120(
+    q_fp8: torch.Tensor,
+    kvcache_fp8: torch.Tensor,
+    weight: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_table: torch.Tensor,
+    deep_gemm_metadata: Any,
+    max_seq_len: int,
+    clean_logits: bool = True,
+) -> torch.Tensor:
+    """CUDA-graph-compatible FP8 paged MQA logits for SM120 (vectorized, no .item())."""
+    _ = deep_gemm_metadata
+    batch_size, _, num_heads, head_dim = q_fp8.shape
+    block_size = kvcache_fp8.shape[1]
+    device = q_fp8.device
+
+    _QUERY_CHUNK = 1024
+    if batch_size > _QUERY_CHUNK:
+        return torch.cat(
+            [
+                fp8_paged_mqa_logits_torch_sm120(
+                    q_fp8[start : start + _QUERY_CHUNK],
+                    kvcache_fp8,
+                    weight[start : start + _QUERY_CHUNK],
+                    seq_lens[start : start + _QUERY_CHUNK],
+                    page_table[start : start + _QUERY_CHUNK],
+                    deep_gemm_metadata,
+                    max_seq_len,
+                    clean_logits=clean_logits,
+                )
+                for start in range(0, batch_size, _QUERY_CHUNK)
+            ],
+            dim=0,
+        )
+
+    assert head_dim == 128, "Vectorized torch impl hardcodes DSV4 indexer head_dim=128"
+    assert (
+        block_size == 64
+    ), "Vectorized torch impl hardcodes block_size=64 cache layout"
+    assert q_fp8.shape == (batch_size, 1, num_heads, head_dim)
+    assert kvcache_fp8.shape[1:] == (block_size, 1, head_dim + 4)
+    assert weight.shape == (batch_size, num_heads)
+    if seq_lens.dim() > 1:
+        seq_lens = seq_lens.squeeze(-1)
+    assert seq_lens.shape == (batch_size,)
+    assert page_table.shape[0] == batch_size
+    assert clean_logits == False
+
+    max_pages = (max_seq_len + block_size - 1) // block_size
+    max_padded_seq = max_pages * block_size
+
+    kvcache_flat = kvcache_fp8.view(-1, block_size * (head_dim + 4))
+    SCALE_OFFSET = block_size * head_dim
+
+    page_ids = page_table[:, :max_pages]
+    kvcache_gathered = kvcache_flat[page_ids]
+
+    kv_value_raw = kvcache_gathered[..., :SCALE_OFFSET]
+    kv_scale_raw = kvcache_gathered[..., SCALE_OFFSET:]
+
+    kv_value = kv_value_raw.contiguous().view(dtype=FP8_DTYPE).to(torch.bfloat16)
+    kv_value = kv_value.view(batch_size, max_padded_seq, head_dim)
+
+    kv_scale = kv_scale_raw.contiguous().view(dtype=torch.float32)
+    kv_scale = kv_scale.view(batch_size, max_padded_seq)
+
+    q = q_fp8[:, 0].to(torch.bfloat16)
+
+    score = torch.bmm(kv_value, q.transpose(1, 2))
+
+    score = F.relu(score)
+    score = score * weight.unsqueeze(1)
+    score = score.sum(dim=2)
+
+    score = score * kv_scale
+
+    out_width = min(max_padded_seq, max_seq_len)
+    logits = score.new_full((batch_size, max_seq_len), float("-inf"))
+    logits[:, :out_width] = score[:, :out_width]
+
+    positions = torch.arange(max_seq_len, device=device)
+    invalid_mask = positions.unsqueeze(0) >= seq_lens.unsqueeze(1)
+    logits.masked_fill_(invalid_mask, float("-inf"))
+
+    return logits
+
+
+def topk_transform_512_pytorch_vectorized(
+    scores: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_tables: torch.Tensor,
+    out_page_indices: torch.Tensor,
+    page_size: int,
+    out_raw_indices: Optional[torch.Tensor] = None,
+) -> None:
+    """Vectorized PyTorch fallback for topk_transform_512.
+    All helper tensors (arange, zeros) are cached to avoid device-tensor
+    creation during HIP/CUDA graph capture."""
+
+    TOPK = out_page_indices.shape[1]
+    batch_size = scores.shape[0]
+    max_seq_len = scores.shape[1]
+    device = scores.device
+
+    page_bits = (page_size - 1).bit_length() if page_size > 1 else 0
+    page_mask = page_size - 1
+
+    cache = _arange_cache
+    key_seq = f"arange_{max_seq_len}_{device}"
+    key_topk = f"arange_{TOPK}_{device}"
+    key_bs = f"arange_{batch_size}_{device}"
+    if key_seq not in cache:
+        cache[key_seq] = torch.arange(max_seq_len, device=device)
+    if key_topk not in cache:
+        cache[key_topk] = torch.arange(TOPK, device=device, dtype=torch.int32)
+    if key_bs not in cache:
+        cache[key_bs] = torch.arange(batch_size, device=device)
+
+    positions = cache[key_seq].unsqueeze(0).expand(batch_size, -1)
+    valid_mask = positions < seq_lens.unsqueeze(1)
+
+    masked_scores = scores.clone()
+    masked_scores.masked_fill_(~valid_mask, float("-inf"))
+
+    actual_k = min(TOPK, max_seq_len)
+    _, raw_indices = torch.topk(
+        masked_scores, k=actual_k, dim=1, largest=True, sorted=False
+    )
+    raw_indices = raw_indices.to(torch.int32)
+
+    if actual_k < TOPK:
+        raw_indices = F.pad(raw_indices, (0, TOPK - actual_k), value=0)
+
+    batch_indices = cache[key_bs].unsqueeze(1).expand(-1, TOPK)
+    gathered_scores = scores[
+        batch_indices.flatten(), raw_indices.clamp(min=0).flatten()
+    ].view(batch_size, TOPK)
+
+    valid_topk = gathered_scores != float("-inf")
+    if actual_k < TOPK:
+        pad_mask = cache[key_topk].unsqueeze(0) >= actual_k
+        valid_topk = valid_topk & ~pad_mask
+
+    needs_sequential = seq_lens <= TOPK
+    sequential_indices = cache[key_topk].unsqueeze(0).expand(batch_size, -1)
+    sequential_valid = sequential_indices < seq_lens.unsqueeze(1)
+
+    seq_indices_or_neg1 = sequential_indices.clone()
+    seq_indices_or_neg1.masked_fill_(~sequential_valid, -1)
+
+    needs_seq_mask = needs_sequential.unsqueeze(1).expand(-1, TOPK)
+    raw_indices = torch.where(needs_seq_mask, seq_indices_or_neg1, raw_indices)
+    valid_topk = torch.where(needs_seq_mask, sequential_valid, valid_topk)
+
+    page_idx = raw_indices >> page_bits
+    offset_in_page = raw_indices & page_mask
+
+    page_idx_clamped = torch.clamp(page_idx, min=0)
+    physical_pages = torch.gather(page_tables, dim=1, index=page_idx_clamped.long())
+
+    page_indices = (physical_pages << page_bits) | offset_in_page
+    page_indices = page_indices.to(torch.int32)
+    page_indices.masked_fill_(~valid_topk, -1)
+
+    out_page_indices.copy_(page_indices)
+
+    if out_raw_indices is not None:
+        raw_indices = raw_indices.clone()
+        raw_indices.masked_fill_(~valid_topk, -1)
+        out_raw_indices.copy_(raw_indices)
+
+
+class C4IndexerBackendMixin:
+    def __init__(self):
+        super().__init__()
+        self.debug_use_external_c4_sparse_indices: bool = False
+
+    def _forward_prepare_multi_stream(
+        self,
+        x: torch.Tensor,
+        q_lora: torch.Tensor,
+        c4_indexer: C4Indexer,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        alt_streams: Optional[List[torch.cuda.Stream]] = None,
+        q_lora_ready: Optional[torch.cuda.Event] = None,
+    ) -> Tuple[IndexerQuery, torch.Tensor]:
+        if TYPE_CHECKING:
+            assert isinstance(self, CompressorBackendMixin)
+
+        assert alt_streams is not None
+        assert len(alt_streams) >= 2
+        current_stream = torch.cuda.current_stream()
+        stream_q = alt_streams[0]
+        stream_weights = alt_streams[1]
+
+        stream_q.wait_stream(current_stream)
+        stream_weights.wait_stream(current_stream)
+
+        self.forward_indexer_compressor(
+            x=x,
+            forward_batch=forward_batch,
+            layer_id=c4_indexer.layer_id,
+            compressor=c4_indexer.compressor,
+        )
+
+        # The weight projection is small and fast; compute it on its own
+        # stream, then have the Q stream wait on it before launching the big
+        # fused Q kernel (which folds rope, hadamard, quantization, and
+        # weight scaling into one pass).
+        with torch.cuda.stream(stream_weights):
+            weights = c4_indexer.compute_weights(x, skip_scale=True)
+            weights_ready = stream_weights.record_event()
+
+        with torch.cuda.stream(stream_q):
+            if q_lora_ready is not None:
+                stream_q.wait_event(q_lora_ready)
+            stream_q.wait_event(weights_ready)
+            q, weights = c4_indexer.compute_q(q_lora, positions, weights)
+
+        current_stream.wait_stream(stream_q)
+        return q, weights
+
+    def _forward_prepare_normal(
+        self,
+        x: torch.Tensor,
+        q_lora: torch.Tensor,
+        c4_indexer: C4Indexer,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        skip_compressor: bool = False,
+    ) -> Tuple[IndexerQuery, torch.Tensor]:
+        if TYPE_CHECKING:
+            assert isinstance(self, CompressorBackendMixin)
+
+        weights = c4_indexer.compute_weights(x, skip_scale=True)
+        q, weights = c4_indexer.compute_q(q_lora, positions, weights)
+        if not skip_compressor:
+            self.forward_indexer_compressor(
+                x=x,
+                forward_batch=forward_batch,
+                layer_id=c4_indexer.layer_id,
+                compressor=c4_indexer.compressor,
+            )
+        return q, weights
+
+    def _can_use_nonpaged_indexer(
+        self,
+        *,
+        c4_indexer: C4Indexer,
+        forward_batch: ForwardBatch,
+        indexer_metadata: PagedIndexerMetadata,
+    ) -> bool:
+        if not envs.SGLANG_OPT_DSV4_NONPAGED_INDEXER.get():
+            return False
+        # This path calls CUDA DeepGEMM and assumes the CUDA FP8+FP32 packed
+        # indexer cache layout. Explicitly reject HIP, NPU, and other devices.
+        if not is_cuda() or is_hip():
+            return False
+        # The gather plan is built from eager, child-local ForwardBatch metadata.
+        # Rewritten, TBO-split, and graph-backed batches must use the paged path.
+        if (
+            forward_batch.forward_mode != ForwardMode.EXTEND
+            or forward_batch._original_forward_mode is not None
+            or forward_batch.tbo_parent_token_range is not None
+            or forward_batch.batch_size != 1
+            or indexer_metadata.use_prefill_cuda_graph
+        ):
+            return False
+        if (
+            c4_indexer.use_fp4_indexer
+            or envs.SGLANG_OPT_USE_TILELANG_INDEXER.get()
+            or envs.SGLANG_OPT_USE_AITER_INDEXER.get()
+            or envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get()
+        ):
+            return False
+        if (
+            get_parallel().attn_cp_size != 1
+            or self.hisparse_coordinator is not None
+            or is_in_tc_piecewise_cuda_graph()
+            or is_in_breakable_cuda_graph()
+        ):
+            return False
+        return not torch.cuda.is_current_stream_capturing()
+
+    def _get_nonpaged_indexer_plan(
+        self,
+        *,
+        c4_indexer: C4Indexer,
+        forward_batch: ForwardBatch,
+        indexer_metadata: PagedIndexerMetadata,
+        page_table: torch.Tensor,
+        c4_seq_lens: torch.Tensor,
+        query_rows: int,
+    ) -> Optional[NonPagedIndexerPlan]:
+        if query_rows < envs.SGLANG_OPT_DSV4_NONPAGED_INDEXER_MIN_QUERY_TOKENS.get():
+            return None
+        if not self._can_use_nonpaged_indexer(
+            c4_indexer=c4_indexer,
+            forward_batch=forward_batch,
+            indexer_metadata=indexer_metadata,
+        ):
+            return None
+        if indexer_metadata.nonpaged_plan is not None:
+            return indexer_metadata.nonpaged_plan
+
+        if (
+            forward_batch.seq_lens is None
+            or forward_batch.seq_lens_cpu is None
+            or forward_batch.extend_seq_lens_cpu is None
+            or forward_batch.extend_seq_lens is None
+            or forward_batch.extend_start_loc is None
+            or forward_batch.extend_num_tokens is None
+        ):
+            return None
+
+        def to_cpu_int_list(values) -> Optional[List[int]]:
+            if isinstance(values, torch.Tensor):
+                if values.device.type != "cpu":
+                    return None
+                values = values.tolist()
+            return [int(value) for value in values]
+
+        extend_lens_cpu = to_cpu_int_list(forward_batch.extend_seq_lens_cpu)
+        seq_lens_cpu = to_cpu_int_list(forward_batch.seq_lens_cpu)
+        if (
+            extend_lens_cpu is None
+            or seq_lens_cpu is None
+            or len(extend_lens_cpu) != 1
+            or len(seq_lens_cpu) != 1
+            or extend_lens_cpu[0] <= 0
+        ):
+            return None
+
+        actual_queries = extend_lens_cpu[0]
+        if (
+            actual_queries != query_rows
+            or int(forward_batch.extend_num_tokens) != query_rows
+            or forward_batch.seq_lens.numel() != 1
+            or forward_batch.extend_seq_lens.numel() != 1
+            or forward_batch.extend_start_loc.numel() != 1
+            or page_table.dim() != 2
+            or page_table.shape[0] < query_rows
+            or c4_seq_lens.numel() < query_rows
+        ):
+            return None
+
+        final_c4_len = seq_lens_cpu[0] // 4
+        if final_c4_len <= 0:
+            return None
+
+        request_page_table = page_table[:1].contiguous()
+        ke = c4_seq_lens[:query_rows].reshape(-1).to(torch.int32).contiguous()
+        gather_seq_lens = ke[-1:]
+        ks = torch.zeros_like(ke)
+        c4_page_size = indexer_metadata.c4_page_size
+        max_seqlen_k = (final_c4_len + c4_page_size - 1) // c4_page_size * c4_page_size
+        plan = NonPagedIndexerPlan(
+            page_table=request_page_table,
+            gather_seq_lens=gather_seq_lens,
+            ks=ks,
+            ke=ke,
+            seq_len_sum=final_c4_len,
+            max_seq_len=final_c4_len,
+            max_seqlen_k=max_seqlen_k,
+            query_rows=query_rows,
+        )
+        indexer_metadata.nonpaged_plan = plan
+        return plan
+
+    @staticmethod
+    def _forward_nonpaged_indexer(
+        *,
+        q_indexer: torch.Tensor,
+        weights: torch.Tensor,
+        c4_indexer: C4Indexer,
+        token_to_kv_pool: DeepSeekV4TokenToKVPool,
+        plan: NonPagedIndexerPlan,
+    ) -> torch.Tensor:
+        import deep_gemm
+
+        k_u8, scale_u8 = token_to_kv_pool.get_index_k_scale_buffer(
+            layer_id=c4_indexer.layer_id,
+            seq_len_tensor=plan.gather_seq_lens,
+            page_indices=plan.page_table,
+            seq_len_sum=plan.seq_len_sum,
+            max_seq_len=plan.max_seq_len,
+        )
+        k_fp8 = k_u8.view(FP8_DTYPE)
+        k_scale = scale_u8.view(torch.float32).squeeze(-1)
+        return deep_gemm.fp8_mqa_logits(
+            q_indexer[: plan.query_rows],
+            (k_fp8, k_scale),
+            weights[: plan.query_rows],
+            plan.ks,
+            plan.ke,
+            clean_logits=False,
+            max_seqlen_k=plan.max_seqlen_k,
+        )
+
+    def forward_c4_indexer(
+        self,
+        x: torch.Tensor,
+        q_lora: torch.Tensor,
+        c4_indexer: C4Indexer,
+        forward_batch: ForwardBatch,
+        alt_streams: Optional[List[torch.cuda.Stream]] = None,
+        enable_multi_stream: bool = False,
+        q_lora_ready: Optional[torch.cuda.Event] = None,
+        skip_compressor: bool = False,
+    ) -> None:
+        if forward_batch.forward_mode.is_idle():
+            return
+        token_to_kv_pool = self.token_to_kv_pool
+
+        if TYPE_CHECKING:
+            assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
+            assert isinstance(self, CompressorBackendMixin)
+
+        metadata = self.forward_metadata
+        indexer_metadata = metadata.indexer_metadata
+        core_metadata = metadata.core_metadata
+
+        assert isinstance(indexer_metadata, PagedIndexerMetadata)
+
+        positions = core_metadata.positions
+        num_queries = min(x.shape[0], q_lora.shape[0], positions.shape[0])
+        if x.shape[0] != num_queries:
+            x = x[:num_queries]
+        if q_lora.shape[0] != num_queries:
+            q_lora = q_lora[:num_queries]
+        if positions.shape[0] != num_queries:
+            positions = positions[:num_queries]
+
+        if enable_multi_stream:
+            q_indexer, weights = self._forward_prepare_multi_stream(
+                x=x,
+                q_lora=q_lora,
+                c4_indexer=c4_indexer,
+                positions=positions,
+                forward_batch=forward_batch,
+                alt_streams=alt_streams,
+                q_lora_ready=q_lora_ready,
+            )
+        else:
+            assert q_lora_ready is None
+            q_indexer, weights = self._forward_prepare_normal(
+                x=x,
+                q_lora=q_lora,
+                c4_indexer=c4_indexer,
+                positions=positions,
+                forward_batch=forward_batch,
+                skip_compressor=skip_compressor,
+            )
+
+        use_fp4_indexer = c4_indexer.use_fp4_indexer
+
+        if use_fp4_indexer:
+            q_fp4, q_sf = q_indexer
+            assert len(q_fp4.shape) == 3
+            assert len(q_sf.shape) == 2
+            q = (q_fp4.unsqueeze(1), q_sf.unsqueeze(1))
+        else:
+            assert len(q_indexer.shape) == 3
+            q = q_indexer.unsqueeze(1)
+
+        assert len(weights.shape) == 3
+        weights = weights.squeeze(2)
+        if use_fp4_indexer:
+            weights = weights.float()
+            if envs.SGLANG_OPT_USE_TILELANG_INDEXER.get():
+                raise RuntimeError("DeepSeek V4 FP4 indexer requires DeepGEMM indexer.")
+            from deep_gemm import fp8_fp4_paged_mqa_logits as fn
+        elif envs.SGLANG_OPT_USE_TILELANG_INDEXER.get():
+            from sglang.kernels.ops.attention.dsa.tilelang_kernel import (
+                tilelang_fp8_paged_mqa_logits as fn,
+            )
+        elif envs.SGLANG_OPT_USE_AITER_INDEXER.get():
+            fn = _aiter_fp8_paged_mqa_logits
+        elif envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get():
+            if is_sm120_supported():
+                fn = fp8_paged_mqa_logits_torch_sm120
+            else:
+                fn = fp8_paged_mqa_logits_torch
+        elif is_xpu():
+            from sgl_kernel import fp8_paged_mqa_logits_triton
+
+            # TODO: switch from triton to SYCL when OOM is resolved
+
+            fn = fp8_paged_mqa_logits_triton
+        else:
+            from deep_gemm import fp8_paged_mqa_logits as fn
+
+        query_rows = q_indexer[0].shape[0] if use_fp4_indexer else q_indexer.shape[0]
+
+        def match_num_queries(tensor: torch.Tensor, value: int) -> torch.Tensor:
+            if tensor.shape[0] == query_rows:
+                return tensor
+            if tensor.shape[0] > query_rows:
+                return tensor[:query_rows]
+            pad = (0, 0) * (tensor.dim() - 1) + (0, query_rows - tensor.shape[0])
+            return F.pad(tensor, pad, value=value)
+
+        c4_seq_lens = match_num_queries(indexer_metadata.c4_seq_lens, value=1)
+        _c4sl = c4_seq_lens
+        page_table = match_num_queries(indexer_metadata.page_table, value=0)
+        c4_sparse_page_indices = match_num_queries(
+            core_metadata.c4_sparse_page_indices, value=-1
+        )
+        _use_tilelang = (
+            envs.SGLANG_OPT_USE_TILELANG_INDEXER.get() and not use_fp4_indexer
+        )
+        _use_aiter = envs.SGLANG_OPT_USE_AITER_INDEXER.get() and not use_fp4_indexer
+        if _c4sl.dim() == 1 and not _use_tilelang and not _use_aiter:
+            _c4sl = _c4sl.unsqueeze(-1)
+        nonpaged_plan = self._get_nonpaged_indexer_plan(
+            c4_indexer=c4_indexer,
+            forward_batch=forward_batch,
+            indexer_metadata=indexer_metadata,
+            page_table=page_table,
+            c4_seq_lens=c4_seq_lens,
+            query_rows=query_rows,
+        )
+        if nonpaged_plan is not None:
+            assert isinstance(q_indexer, torch.Tensor)
+            logits = self._forward_nonpaged_indexer(
+                q_indexer=q_indexer,
+                weights=weights,
+                c4_indexer=c4_indexer,
+                token_to_kv_pool=token_to_kv_pool,
+                plan=nonpaged_plan,
+            )
+        else:
+            c4_indexer_kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(
+                layer_id=c4_indexer.layer_id,
+            )
+            assert c4_indexer_kv_cache.dim() == 2
+            head_dim_with_sf = 68 if use_fp4_indexer else 132
+            c4_indexer_kv_cache = c4_indexer_kv_cache.view(
+                c4_indexer_kv_cache.shape[0], 64, 1, head_dim_with_sf
+            )
+            logits = fn(
+                q,
+                c4_indexer_kv_cache,
+                weights,
+                _c4sl,
+                page_table,
+                indexer_metadata.deep_gemm_metadata,
+                indexer_metadata.max_c4_seq_len,
+                False,
+            )
+
+        assert indexer_metadata.page_table is core_metadata.page_table
+        if self.debug_use_external_c4_sparse_indices:
+            return
+
+        indexer_capturer = get_global_indexer_capturer()
+        capture_enabled = indexer_capturer is not None
+
+        hisparse_coordinator = self.hisparse_coordinator
+        hisparse_decode = (
+            hisparse_coordinator is not None and forward_batch.forward_mode.is_decode()
+        )
+
+        raw_indices = None
+        if capture_enabled:
+            raw_indices = torch.empty_like(c4_sparse_page_indices)
+        elif hisparse_decode:
+            raw_indices = hisparse_coordinator.raw_indices_buffer[
+                : c4_sparse_page_indices.size(0)
+            ]
+        elif core_metadata.c4_sparse_raw_indices is not None:
+            raw_indices = core_metadata.c4_sparse_raw_indices
+
+        if envs.SGLANG_TOPK_TRANSFORM_512_TORCH.get():
+            topk_transform_512_pytorch_vectorized(
+                logits,
+                c4_seq_lens,
+                page_table,
+                c4_sparse_page_indices,
+                indexer_metadata.c4_page_size,
+                raw_indices,
+            )
+        elif envs.SGLANG_OPT_USE_TOPK_V2.get() and raw_indices is None:
+            topk_transform_512_v2(
+                logits,
+                c4_seq_lens,
+                page_table,
+                c4_sparse_page_indices,
+                indexer_metadata.c4_page_size,
+                indexer_metadata.topk_metadata,
+            )
+        else:
+            topk_transform_512(
+                logits,
+                c4_seq_lens,
+                page_table,
+                c4_sparse_page_indices,
+                indexer_metadata.c4_page_size,
+                raw_indices,
+            )
+        if hisparse_coordinator is not None:
+            if hisparse_decode:
+                compress_layer_id = token_to_kv_pool.layer_mapping[
+                    c4_indexer.layer_id
+                ].compress_layer_id
+                core_metadata.c4_sparse_page_indices = (
+                    hisparse_coordinator.swap_in_selected_pages(
+                        req_pool_indices=forward_batch.req_pool_indices,
+                        compressed_seq_lens=indexer_metadata.c4_seq_lens,
+                        top_k_result=raw_indices,
+                        layer_id=compress_layer_id,
+                    )
+                )
+            else:
+                # flash_mla C4 attention requires int32 page indices.
+                core_metadata.c4_sparse_page_indices = (
+                    token_to_kv_pool.c4_kv_pool.translate_loc_to_hisparse_device(
+                        core_metadata.c4_sparse_page_indices
+                    ).to(torch.int32)
+                )
+
+        if capture_enabled:
+            compress_layer_id = token_to_kv_pool.layer_mapping[
+                c4_indexer.layer_id
+            ].compress_layer_id
+            indexer_capturer.capture(compress_layer_id, raw_indices)
+
+
+class C4Indexer(nn.Module):
+    def __init__(
+        self,
+        config: AgnesConfig,
+        layer_id: int,
+        freqs_cis: torch.Tensor,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+        alt_streams: Optional[List[torch.cuda.Stream]] = None,
+        rotary_emb=None,
+    ):
+        super().__init__()
+        self.layer_id = layer_id
+        self.dim = config.hidden_size
+        self.n_heads = config.index_n_heads
+        self.head_dim = config.index_head_dim
+        self.rope_head_dim = config.qk_rope_head_dim
+        self.index_topk = config.index_topk
+        self.q_lora_rank = config.q_lora_rank
+        self.softmax_scale = self.head_dim**-0.5
+        self.n_local_heads = self.n_heads
+        self.wq_b = ReplicatedLinear(
+            self.q_lora_rank,
+            self.n_heads * self.head_dim,
+            bias=False,
+            quant_config=quant_config,
+            params_dtype=torch.bfloat16,
+            prefix=add_prefix("wq_b", prefix),
+        )
+        self.weights_proj = ReplicatedLinear(
+            self.dim,
+            self.n_heads,
+            bias=False,
+            quant_config=None,
+            params_dtype=torch.bfloat16,
+            prefix=add_prefix("weights_proj", prefix),
+        )
+        self.compressor = Compressor(
+            config,
+            self.layer_id,
+            True,
+            freqs_cis,
+            compress_ratio=4,
+            head_dim=self.head_dim,
+            rotate=True,
+            prefix=add_prefix("compressor", prefix),
+            rotary_emb=rotary_emb,
+        )
+        self.rotary_emb = rotary_emb
+        self.freqs_cis = freqs_cis
+        self.weight_scale: float = self.softmax_scale * self.n_heads**-0.5
+        from sglang.srt.runtime_context import get_server_args
+
+        self.use_fp4_indexer = get_server_args().enable_deepseek_v4_fp4_indexer
+        self.alt_streams = alt_streams
+
+    def compute_q(
+        self,
+        q_lora: torch.Tensor,
+        positions: torch.Tensor,
+        weight: torch.Tensor,
+    ) -> Tuple[IndexerQuery, torch.Tensor]:
+        q, _ = self.wq_b(q_lora)
+        q = q.view(-1, self.n_local_heads, self.head_dim)
+        if self.use_fp4_indexer:
+            return fused_q_indexer_rope_hadamard_fp4_quant(
+                q.contiguous(), weight, self.weight_scale, self.freqs_cis, positions
+            )
+        return fused_q_indexer_rope_hadamard_quant(
+            q, weight, self.weight_scale, self.freqs_cis, positions
+        )
+
+    def compute_weights(self, x: torch.Tensor, skip_scale=False) -> torch.Tensor:
+        out, _ = self.weights_proj(x)
+        if not skip_scale:
+            out = out * self.weight_scale
+        return out
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        q_lora: torch.Tensor,
+        forward_batch: ForwardBatch,
+        attn_backend: AttentionBackend,
+        enable_multi_stream: bool = False,
+        q_lora_ready: Optional[torch.cuda.Event] = None,
+        skip_compressor: bool = False,
+    ) -> None:
+        return attn_backend.forward_c4_indexer(
+            x=x,
+            q_lora=q_lora,
+            forward_batch=forward_batch,
+            c4_indexer=self,
+            alt_streams=self.alt_streams,
+            enable_multi_stream=enable_multi_stream,
+            q_lora_ready=q_lora_ready,
+            skip_compressor=skip_compressor,
+        )

sglang_patch/srt/layers/quantization/modelslim/modelslim.py

@@ -0,0 +1,485 @@
+from __future__ import annotations
+
+import logging
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Tuple, Union, cast
+
+import torch
+
+from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
+    _NPULinearMethodBase,
+)
+from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig
+from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend
+from sglang.srt.layers.quantization.base_config import (
+    FusedMoEMethodBase,
+    QuantizationConfig,
+)
+from sglang.srt.layers.quantization.modelslim.schemes import (
+    ModelSlimMXFP4Scheme,
+    ModelSlimMXFP4W4A8Scheme,
+    ModelSlimMXFP8Scheme,
+    ModelSlimW4A4Int4,
+    ModelSlimW4A4Int4MoE,
+    ModelSlimW4A8Int8MoE,
+    ModelSlimW8A8Int8,
+    ModelSlimW8A8Int8MoE,
+)
+from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
+from sglang.srt.utils import apply_module_patch
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe import MoeRunnerConfig
+    from sglang.srt.layers.moe.token_dispatcher import (
+        CombineInput,
+        StandardDispatchOutput,
+    )
+    from sglang.srt.layers.quantization.base_config import QuantizeMethodBase
+    from sglang.srt.layers.quantization.modelslim.schemes import (
+        ModelSlimLinearScheme,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+# func refers to RMSNorm.__init__
+def npu_wrapper_rmsnorm_init(func):
+    def init(self, hidden_size: int, **extra_args) -> None:
+        func(self, hidden_size, **extra_args)
+        self.ignore_anti = True
+        # The Ascend w8a8_int8 quantization requires adding a bias in rmsnorm
+        self.bias = torch.nn.Parameter(torch.zeros(hidden_size), requires_grad=False)
+
+    return init
+
+
+# func refers to RMSNorm.forward_oot
+def npu_wrapper_rmsnorm_forward(func):
+    def _rmsnorm_forward_oot(
+        self,
+        x: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+        post_residual_addition: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if not x.is_contiguous():
+            x = x.contiguous()
+        if residual is not None:
+            if post_residual_addition is not None:
+                residual = residual + post_residual_addition
+            from sgl_kernel_npu.norm.add_rmsnorm_bias import add_rmsnorm_bias
+
+            out, residual_out = add_rmsnorm_bias(
+                x,
+                residual,
+                self.weight.data,
+                self.bias,
+                self.variance_epsilon,
+            )
+            return out.to(x.dtype), residual_out
+
+        out = torch.ops.npu.npu_rms_norm(x, self.weight.data, self.variance_epsilon)[0]
+        out = out + self.bias
+        return out.to(x.dtype)
+
+    return _rmsnorm_forward_oot
+
+
+class ModelSlimConfig(QuantizationConfig):
+    """
+    Config class for ModelSlim Quantization, a NPU-specific quantization type.
+    """
+
+    def __init__(self, quant_config: Dict[str, Any] = {}):
+        super().__init__()
+        keys = [k for k in quant_config if isinstance(k, str)]
+        is_dsv4 = any(k.startswith("hc_head_") for k in keys)
+        if is_dsv4:
+            from sglang.srt.models.agnes import AgnesForCausalLM
+
+            remap = AgnesForCausalLM.remap_weight_name_to_dpsk_hf_format
+            quant_config = {
+                (remap(k) if isinstance(k, str) else k): v
+                for k, v in quant_config.items()
+            }
+
+        self.quant_description = quant_config
+        ignore = cast(List[str], quant_config.get("ignore", []))
+        self.ignore = ignore if ignore is not None else []
+        packed_modules_mapping = quant_config.get("packed_modules_mapping", {})
+        self.packed_modules_mapping = (
+            packed_modules_mapping if packed_modules_mapping is not None else {}
+        )
+
+        for name in self.quant_description.keys():
+            if "norm.bias" in name:
+                apply_module_patch(
+                    "sglang.srt.layers.layernorm.RMSNorm",
+                    "__init__",
+                    [npu_wrapper_rmsnorm_init],
+                )
+                apply_module_patch(
+                    "sglang.srt.layers.layernorm.RMSNorm",
+                    "forward_npu",
+                    [npu_wrapper_rmsnorm_forward],
+                )
+
+    def update_packed_modules_mapping(self, mapping: Dict[str, List[str]]) -> None:
+        self.packed_modules_mapping.update(mapping)
+
+    def get_linear_method(self) -> ModelSlimLinearMethod:
+        return ModelSlimLinearMethod(self)
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.int8, torch.float16, torch.bfloat16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 0
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "modelslim"
+
+    @classmethod
+    def get_config_filenames(cls) -> List[str]:
+        filenames = ["quant_model_description.json"]
+        return filenames
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> ModelSlimConfig:
+        return cls(config)
+
+    def get_quant_method(
+        self,
+        layer: torch.nn.Module,
+        prefix: str,
+    ) -> Optional[QuantizeMethodBase]:
+        from sglang.srt.layers.linear import LinearBase
+        from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
+
+        if isinstance(layer, LinearBase):
+            # TODO: we should remove this code and switch to the packed_modules_mapping declared inside the modeling files
+            key = "model"
+            if "vision_model" in prefix:
+                key = "vision_model"
+            elif "visual" in prefix:
+                key = "visual"
+            if "vision_tower" in prefix or "mm_projector" in prefix:
+                prefix = prefix.replace(r"attn.qkv_proj", r"wqkv")
+                prefix = prefix.replace(r"attn.proj", r"wo")
+            packed_modules_mapping_subset = self.packed_modules_mapping.get(key, {})
+            prefix_in_quant_config = prefix
+            proj_name = prefix.split(".")[-1]
+            if proj_name in packed_modules_mapping_subset:
+                prefix_in_quant_config = prefix.replace(
+                    proj_name, packed_modules_mapping_subset[proj_name][0]
+                )
+            if self.is_layer_skipped(
+                prefix, packed_modules_mapping_subset
+            ) or self.is_layer_skipped(prefix, self.packed_modules_mapping):
+                return UnquantizedLinearMethod()
+            layer.scheme = self.get_linear_scheme(layer, prefix_in_quant_config)
+            if layer.scheme is None:
+                return UnquantizedLinearMethod()
+            return ModelSlimLinearMethod(self)
+        elif isinstance(layer, FusedMoE):
+            moe_schemes = self.get_moe_scheme(layer, prefix)
+            if moe_schemes is None:
+                raise ValueError(f"No ModelSlim MoE scheme found for layer {prefix}")
+            layer.w13_scheme, layer.w2_scheme = moe_schemes
+            layer.w13_kernel, layer.w2_kernel = (
+                layer.w13_scheme.kernel,
+                layer.w2_scheme.kernel,
+            )
+            return ModelSlimFusedMoEMethod(self)
+        return None
+
+    def get_linear_scheme(
+        self, layer: torch.nn.Module, prefix: Optional[str] = None
+    ) -> Optional[ModelSlimLinearScheme]:
+        """
+        get_scheme method adjusted for modelslim, taken from
+        python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
+        """
+
+        linear_quant_schemes = [
+            ("W4A4_DYNAMIC", ModelSlimW4A4Int4),
+            ("W8A8", ModelSlimW8A8Int8),
+            ("W8A8_DYNAMIC", ModelSlimW8A8Int8),
+            ("W8A8_MXFP8", ModelSlimMXFP8Scheme),
+            ("W4A8_MXFP", ModelSlimMXFP4W4A8Scheme),
+            ("W4A4_MXFP4", ModelSlimMXFP4Scheme),
+        ]
+
+        quant_schemes = [self.quant_description.get(prefix + ".weight", "")]
+
+        for scheme_name, scheme_class in linear_quant_schemes:
+            if any(s == scheme_name for s in quant_schemes):
+                logger.info_once(f"Using {scheme_class.__name__}")
+                return scheme_class(quant_config=self.quant_description, prefix=prefix)
+
+        logger.warning(
+            f"Unsupported Linear modelslim scheme: "
+            f"{quant_schemes} in layer: {prefix}"
+        )
+        return None
+
+    def get_moe_scheme(
+        self,
+        layer: torch.nn.Module,
+        prefix: str,
+    ):
+        moe_quant_schemes = [
+            ("W4A4_DYNAMIC", ModelSlimW4A4Int4MoE),
+            ("W4A8_DYNAMIC", ModelSlimW4A8Int8MoE),
+            ("W8A8_DYNAMIC", ModelSlimW8A8Int8MoE),
+        ]
+
+        # Try multiple naming conventions:
+        #   (gate_proj, up_proj, down_proj) – standard compressed-tensors format
+        #   (w1, w3, w2)                      – MiniMax-M2.5 / some other models
+        naming_conventions = [
+            ("gate_proj", "up_proj", "down_proj"),
+            ("w1", "w3", "w2"),
+        ]
+
+        w13_scheme_name = None
+        w2_scheme_name = None
+        for gate_name, up_name, down_name in naming_conventions:
+            w13_keys = [
+                f"{prefix}.0.{gate_name}.weight",
+                f"{prefix}.0.{up_name}.weight",
+            ]
+            w2_key = f"{prefix}.0.{down_name}.weight"
+            w13_entries = {
+                key: self.quant_description[key]
+                for key in w13_keys
+                if key in self.quant_description
+            }
+            if w13_entries and w2_key in self.quant_description:
+                w13_names = list(w13_entries.values())
+                # For w13, both projections must agree on the scheme
+                unique_w13 = set(w13_names)
+                if len(unique_w13) > 1:
+                    raise ValueError(
+                        f"Mismatched ModelSlim quantization for W13 in layer {prefix}: "
+                        f"{w13_entries}"
+                    )
+                w13_scheme_name = w13_names[0]
+                w2_scheme_name = self.quant_description[w2_key]
+                break
+
+        if w13_scheme_name is None:
+            # Build a helpful error message listing all attempted key patterns
+            all_attempted = []
+            for gate_name, up_name, down_name in naming_conventions:
+                w13_keys = [
+                    f"{prefix}.0.{gate_name}.weight",
+                    f"{prefix}.0.{up_name}.weight",
+                ]
+                w2_key = f"{prefix}.0.{down_name}.weight"
+                w13_found = any(k in self.quant_description for k in w13_keys)
+                w2_found = w2_key in self.quant_description
+                status = (
+                    f"({gate_name}/{up_name}={'found' if w13_found else 'missing'}, "
+                    f"{down_name}={'found' if w2_found else 'missing'})"
+                )
+                all_attempted.append(status)
+            raise ValueError(
+                f"Missing ModelSlim MoE quantization description for layer {prefix}: "
+                + "; ".join(all_attempted)
+            )
+
+        # Map scheme names to classes
+        scheme_map = dict(
+            moe_quant_schemes
+        )  # dict: "W4A4_DYNAMIC" -> ModelSlimW4A4Int4MoE, etc.
+
+        # Instantiate the schemes
+        def instantiate(name, weight_group):
+            cls = scheme_map.get(name)
+            if cls is None:
+                logger.warning(f"Unsupported scheme '{name}' for layer {prefix}")
+                return None
+            return cls(self, weight_group)
+
+        w13_scheme = instantiate(w13_scheme_name, weight_group="w13")
+        w2_scheme = instantiate(w2_scheme_name, weight_group="w2")
+        if w13_scheme is None or w2_scheme is None:
+            raise ValueError(
+                f"Unsupported ModelSlim MoE schemes for layer {prefix}: "
+                f"W13='{w13_scheme_name}', W2='{w2_scheme_name}'"
+            )
+        logger.info_once(f"Using {type(w13_scheme).__name__} for W13")
+        logger.info_once(f"Using {type(w2_scheme).__name__} for W2")
+
+        return w13_scheme, w2_scheme
+
+    def is_layer_skipped(
+        self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({})
+    ):
+        # adapted from vllm.model_executor.layers.quantization.utils.quant_utils.is_layer_skipped
+        proj_name = prefix.split(".")[-1]
+        if proj_name in fused_mapping:
+            shard_prefixes = [
+                prefix.replace(proj_name, shard_proj_name)
+                for shard_proj_name in fused_mapping[proj_name]
+            ]
+
+            is_skipped = None
+            for shard_prefix in shard_prefixes:
+                is_shard_skipped = (
+                    self.quant_description.get(shard_prefix + ".weight", "") == "FLOAT"
+                )
+
+                if is_skipped is None:
+                    is_skipped = is_shard_skipped
+                elif is_shard_skipped != is_skipped:
+                    raise ValueError(
+                        f"Detected some but not all shards of {prefix} "
+                        "are quantized. All shards of fused layers "
+                        "to have the same precision."
+                    )
+        else:
+            is_skipped = self.quant_description.get(prefix + ".weight", "") == "FLOAT"
+
+        assert is_skipped is not None
+        return is_skipped
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+
+class ModelSlimLinearMethod(_NPULinearMethodBase):
+
+    def __init__(self, quantization_config: ModelSlimConfig):
+        self.quantization_config = quantization_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.scheme.process_weights_after_loading(layer)
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        """
+        Use the ModelSlimLinearScheme associated with the layer to create
+        the necessary parameters for the layer. See LinearMethodBase for param
+        details
+        """
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        layer.scheme.create_weights(
+            layer=layer,
+            input_size=input_size,
+            input_size_per_partition=input_size_per_partition,
+            output_partition_sizes=output_partition_sizes,
+            output_size=output_size,
+            params_dtype=params_dtype,
+            weight_loader=weight_loader,
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ):
+        """
+        Use the output of create_weights and the ModelSlimLinearScheme
+        associated with the layer to apply the forward pass with the
+        layer input.  See LinearMethodBase for param details
+
+        """
+
+        scheme = layer.scheme
+        if scheme is None:
+            raise ValueError("A scheme must be defined for each layer")
+        return scheme.apply_weights(layer, x, bias=bias)
+
+
+class ModelSlimFusedMoEMethod(FusedMoEMethodBase):
+    """
+    Fused MoE method for ModelSlim quantization on Ascend NPU.
+
+    Delegates routing, activation, and finalization to the modular NPU MoE
+    components introduced in the hardware backend refactoring.
+    """
+
+    def __init__(self, quantization_config: ModelSlimConfig):
+        self.quantization_config = quantization_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.w13_scheme.process_weights_after_loading(layer)
+        layer.w2_scheme.process_weights_after_loading(layer)
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        """
+        Use the ModelSlimMoEScheme associated with the layer to create
+        the necessary parameters for the layer. See FusedMoEMethodBase for param
+        details
+        """
+        layer.w13_scheme.create_weights(
+            layer=layer,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size_per_partition=intermediate_size_per_partition,
+            weight_prefix="w13",
+            **extra_weight_attrs,
+        )
+        layer.w2_scheme.create_weights(
+            layer=layer,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size_per_partition=intermediate_size_per_partition,
+            weight_prefix="w2",
+            **extra_weight_attrs,
+        )
+
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        moe_runner_config.layer = layer
+        self.moe_runner_config = moe_runner_config
+        backend = get_moe_runner_backend()
+        if backend.is_auto():
+            backend = MoeRunnerBackend.ASCEND
+        self.runner = MoeRunner(backend, moe_runner_config)
+
+    # ------------------------------------------------------------------
+    # Main apply()
+    # ------------------------------------------------------------------
+    def apply(
+        self,
+        layer,
+        dispatch_output: StandardDispatchOutput,
+    ) -> CombineInput:
+        from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo
+
+        quant_info = AscendQuantInfo(
+            w13_weight=layer.w13_weight,
+            w2_weight=layer.w2_weight,
+            w13_weight_scale=layer.w13_weight_scale,
+            w2_weight_scale=layer.w2_weight_scale,
+            w13_weight_offset=layer.w13_weight_offset,
+            w2_weight_offset=layer.w2_weight_offset,
+            w13_scale_bias=getattr(layer, "w13_scale_bias", None),
+            w2_scale_bias=getattr(layer, "w2_scale_bias", None),
+            w13_weight_bias=getattr(layer, "w13_weight_bias", None),
+            w2_weight_bias=getattr(layer, "w2_weight_bias", None),
+        )
+        return self.runner.run(dispatch_output, quant_info)

sglang_patch/srt/mem_cache/kv_cache_configurator.py

@@ -0,0 +1,1867 @@
+from __future__ import annotations
+
+import logging
+import math
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Optional
+
+import msgspec
+import torch
+
+from sglang.srt.configs.hybrid_arch import hybrid_gdn_config, mambaish_config
+from sglang.srt.configs.model_config import (
+    ModelConfig,
+    get_dsa_index_head_dim,
+    get_minimax_sparse_attention_config,
+    get_minimax_sparse_disable_value_layer_ids,
+    get_minimax_sparse_layer_ids,
+    is_deepseek_dsa,
+    is_deepseek_v4,
+    is_minimax_sparse,
+)
+from sglang.srt.distributed.parallel_state import get_world_group
+from sglang.srt.environ import envs
+from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
+    get_kv_cache_quant_method,
+    resolve_kv_cache_quant,
+)
+from sglang.srt.mem_cache.allocation_sizing import get_req_to_token_extra_context_len
+from sglang.srt.mem_cache.allocator import (
+    BaseTokenToKVPoolAllocator,
+    PagedTokenToKVPoolAllocator,
+    TokenToKVPoolAllocator,
+)
+from sglang.srt.mem_cache.allocator.hisparse import (
+    DeepSeekV4HiSparseTokenToKVPoolAllocator,
+    HiSparseTokenToKVPoolAllocator,
+)
+from sglang.srt.mem_cache.allocator.swa import (
+    PureSWATokenToKVPoolAllocator,
+    SWATokenToKVPoolAllocator,
+)
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool
+from sglang.srt.mem_cache.memory_pool import (
+    DSATokenToKVPool,
+    HybridLinearKVPool,
+    HybridReqToTokenPool,
+    KVCache,
+    MHATokenToKVPool,
+    MHATokenToKVPoolFP4,
+    MHATokenToKVPoolMXFP8,
+    MiniMaxSparseKVPool,
+    MLATokenToKVPool,
+    MLATokenToKVPoolFP4,
+    NoOpMHATokenToKVPool,
+    PageMajorMHATokenToKVPool,
+    ReqToTokenPool,
+)
+from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
+from sglang.srt.platforms import current_platform
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.utils.common import (
+    get_available_gpu_memory,
+    get_device_memory_capacity,
+    is_float4_e2m1fn_x2,
+    is_hip,
+    is_npu,
+)
+
+logger = logging.getLogger(__name__)
+
+_is_hip = is_hip()
+
+
+def _get_dsv4_compress_state_dtypes() -> tuple[torch.dtype, torch.dtype]:
+    dtype_name = envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.get().strip().lower()
+    if dtype_name in ("float32", "fp32"):
+        return torch.float32, torch.float32
+    if dtype_name in ("bfloat16", "bf16"):
+        return torch.bfloat16, torch.bfloat16
+    raise ValueError(
+        "Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE="
+        f"{dtype_name!r}. Expected one of: float32, fp32, bfloat16, bf16."
+    )
+
+
+_is_npu = is_npu()
+
+
+def _should_enable_lazy_compaction() -> bool:
+    """Lazy compaction default — ON unless
+    `SGLANG_DISABLE_LAZY_COMPACTION=1` (escape hatch for A/B / rollback).
+    Centralized here so both unified-memory-pool factory call sites stay in sync.
+    """
+    return not envs.SGLANG_DISABLE_LAZY_COMPACTION.get()
+
+
+# the ratio of mamba cache pool size to max_running_requests
+MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO = 3
+MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP = 2
+MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP_LAZY = 1
+MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP = 1
+
+if TYPE_CHECKING:
+    from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+    from sglang.srt.mem_cache.unified_memory_pool import (
+        UnifiedKVPool,
+        UnifiedPoolBundle,
+    )
+    from sglang.srt.model_executor.model_runner_components.layer_setup import (
+        ModelLayerInfo,
+    )
+    from sglang.srt.model_executor.model_runner_components.spec_aux_hidden_state import (
+        SpecAuxHiddenStateConfig,
+    )
+    from sglang.srt.model_executor.pool_configurator import (
+        MemoryPoolConfig,
+    )
+
+
+class KVCacheConfigResult(msgspec.Struct, frozen=True, kw_only=True):
+    max_total_num_tokens: int
+    max_running_requests: int
+    full_max_total_num_tokens: Optional[int]
+    swa_max_total_num_tokens: Optional[int]
+    req_to_token_pool: ReqToTokenPool
+    token_to_kv_pool: KVCache
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
+    memory_pool_config: MemoryPoolConfig
+    unified_memory_pool: Optional[UnifiedKVPool] = None
+
+
+class _InitializedPools(msgspec.Struct, frozen=True, kw_only=True):
+    req_to_token_pool: ReqToTokenPool
+    token_to_kv_pool: KVCache
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
+    unified_memory_pool: Optional[UnifiedKVPool] = None
+
+
+class _PoolSizes(msgspec.Struct, frozen=True, kw_only=True):
+    max_total_num_tokens: int
+    max_running_requests: int
+    full_max_total_num_tokens: Optional[int]
+    swa_max_total_num_tokens: Optional[int]
+    c4_max_total_num_tokens: int
+    c128_max_total_num_tokens: int
+    c4_state_pool_size: int
+    c128_state_pool_size: int
+    c4_state_dtype: Optional[torch.dtype]
+    c128_state_dtype: Optional[torch.dtype]
+
+
+@dataclass(slots=True, kw_only=True)
+class KVCacheConfigurator:
+    device: str
+    gpu_id: int
+    ps: ParallelState
+    pp_group: Any
+    model: Any
+    model_config: ModelConfig
+    server_args: ServerArgs
+    kv_cache_dtype: torch.dtype
+    model_dtype: torch.dtype
+    page_size: int
+    sliding_window_size: Optional[int]
+    spec_algorithm: SpeculativeAlgorithm
+    is_draft_worker: bool
+    post_capture_kv_active: bool
+    spec_aux_config: SpecAuxHiddenStateConfig
+    is_hybrid_swa: bool
+    is_hybrid_swa_compress: bool
+    use_mla_backend: bool
+    layer_info: ModelLayerInfo
+    forward_stream: Any
+    req_to_token_pool: Optional[ReqToTokenPool]
+    token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator]
+    memory_pool_config: Optional[MemoryPoolConfig]
+    draft_model_idx: Optional[int] = None
+    mambaish_config: Optional[Any] = field(init=False)
+    hybrid_gdn_config: Optional[Any] = field(init=False)
+    is_inkling_mtp_draft: bool = field(init=False)
+    draft_swa_full_capacity: bool = field(init=False)
+
+    def __post_init__(self) -> None:
+        self.mambaish_config = mambaish_config(self.model_config)
+        self.hybrid_gdn_config = hybrid_gdn_config(self.model_config)
+        # Each multi-layer EAGLE MTP head owns one transformer block at
+        # layer_id=draft_model_idx; heads at a banded 's' depth route that layer
+        # into the SWA ring sub-pool (draft_swa_full_capacity) so the SWA
+        # store/read path activates for this depth, exactly like a trunk local
+        # layer.
+        self.is_inkling_mtp_draft = (
+            self.is_draft_worker
+            and self.draft_model_idx is not None
+            and self.model_config.hf_config.architectures[0]
+            == "InklingForConditionalGenerationMTP"
+        )
+        self.draft_swa_full_capacity = self.is_inkling_mtp_draft and (
+            self.draft_model_idx
+            in set(self.model_config.hf_text_config.mtp_local_layer_ids)
+        )
+
+    def _build_fp4_quant_method(self, *, num_layers: int):
+        if not is_float4_e2m1fn_x2(self.kv_cache_dtype):
+            return None
+        quant_name = resolve_kv_cache_quant(self.server_args.kv_cache_dtype)
+        if quant_name is None:
+            return None
+        quant_method = get_kv_cache_quant_method(
+            quant_name,
+            num_layers=num_layers,
+            device=self.device,
+        )
+        quant_method.load_scales_from_model(self.model)
+        return quant_method
+
+    def configure(self, *, pre_model_load_memory: int) -> KVCacheConfigResult:
+        """Apply a resolved MemoryPoolConfig and initialize pools."""
+        if not self.spec_algorithm.is_none() and self.is_draft_worker:
+            assert (
+                self.memory_pool_config is not None
+            ), "Draft worker requires memory_pool_config"
+            config = self.memory_pool_config
+        else:
+            config = self._resolve_memory_pool_config(pre_model_load_memory)
+
+        sizes = self._derive_pool_sizes(config=config)
+
+        pools = self._init_pools(
+            sizes=sizes,
+            req_to_token_pool=self.req_to_token_pool,
+            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+        )
+
+        logger.info(
+            f"Memory pool end. "
+            f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
+        )
+
+        return KVCacheConfigResult(
+            max_total_num_tokens=sizes.max_total_num_tokens,
+            max_running_requests=sizes.max_running_requests,
+            full_max_total_num_tokens=sizes.full_max_total_num_tokens,
+            swa_max_total_num_tokens=sizes.swa_max_total_num_tokens,
+            req_to_token_pool=pools.req_to_token_pool,
+            token_to_kv_pool=pools.token_to_kv_pool,
+            token_to_kv_pool_allocator=pools.token_to_kv_pool_allocator,
+            memory_pool_config=config,
+            unified_memory_pool=pools.unified_memory_pool,
+        )
+
+    def _derive_pool_sizes(self, *, config: MemoryPoolConfig) -> _PoolSizes:
+        max_total_num_tokens = config.max_total_num_tokens
+        max_running_requests = config.max_running_requests
+        full_max_total_num_tokens = None
+        swa_max_total_num_tokens = None
+        if self.is_hybrid_swa:
+            full_max_total_num_tokens = config.full_max_total_num_tokens
+            swa_max_total_num_tokens = config.swa_max_total_num_tokens
+
+        # DSV4 compressed-attention pool sizes. Draft worker reuses target's
+        # full/swa sizes but does NOT own c4/c128/state pools (those live on
+        # the target rank only); zero them out regardless of what config holds.
+        if self.is_draft_worker:
+            c4_max_total_num_tokens = 0
+            c128_max_total_num_tokens = 0
+            c4_state_pool_size = 0
+            c128_state_pool_size = 0
+        else:
+            c4_max_total_num_tokens = config.c4_max_total_num_tokens
+            c128_max_total_num_tokens = config.c128_max_total_num_tokens
+            c4_state_pool_size = config.c4_state_pool_size
+            c128_state_pool_size = config.c128_state_pool_size
+
+        # Draft worker does not own the compression-state pools, but keep the
+        # dtype attributes initialized so _init_pools can share one code path.
+        c4_state_dtype: Optional[torch.dtype] = None
+        c128_state_dtype: Optional[torch.dtype] = None
+        if is_deepseek_v4(self.model_config.hf_config):
+            c4_state_dtype, c128_state_dtype = _get_dsv4_compress_state_dtypes()
+
+        return _PoolSizes(
+            max_total_num_tokens=max_total_num_tokens,
+            max_running_requests=max_running_requests,
+            full_max_total_num_tokens=full_max_total_num_tokens,
+            swa_max_total_num_tokens=swa_max_total_num_tokens,
+            c4_max_total_num_tokens=c4_max_total_num_tokens,
+            c128_max_total_num_tokens=c128_max_total_num_tokens,
+            c4_state_pool_size=c4_state_pool_size,
+            c128_state_pool_size=c128_state_pool_size,
+            c4_state_dtype=c4_state_dtype,
+            c128_state_dtype=c128_state_dtype,
+        )
+
+    def _init_pools(
+        self,
+        *,
+        sizes: _PoolSizes,
+        req_to_token_pool: Optional[ReqToTokenPool],
+        token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator],
+    ) -> _InitializedPools:
+        """Initialize the memory pools."""
+        token_to_kv_pool = None
+
+        # Unified-pool fast path: build req_to_token + token_to_kv pool + allocator
+        # from one byte buffer, then return. Gated to the target worker
+        # (req_to_token_pool is None); supports hybrid Mamba and hybrid SWA (not DSV4).
+        if (
+            self.server_args.enable_unified_memory
+            and self.server_args.disaggregation_mode == "null"
+            and req_to_token_pool is None
+        ):
+            if self.mambaish_config is not None:
+                bundle = self._init_unified_mamba_pools(
+                    max_num_reqs=sizes.max_running_requests,
+                    max_total_num_tokens=sizes.max_total_num_tokens,
+                )
+            elif self.is_hybrid_swa and not is_deepseek_v4(self.model_config.hf_config):
+                bundle = self._init_unified_swa_pools(
+                    max_num_reqs=sizes.max_running_requests,
+                    full_max_total_num_tokens=sizes.full_max_total_num_tokens,
+                    swa_max_total_num_tokens=sizes.swa_max_total_num_tokens,
+                )
+            else:
+                # Fail loud, not silently fall through to the normal pools (which would
+                # leave the flag a no-op). The feature replaces the HYBRID pools only.
+                raise ValueError(
+                    "--enable-unified-memory only supports hybrid Mamba and "
+                    "hybrid sliding-window-attention models (DeepSeek-V4 excluded); "
+                    f"the current model ({self.model_config.hf_config.architectures}) "
+                    "is neither, so the unified memory pool cannot be built. Drop "
+                    "--enable-unified-memory for this model."
+                )
+            return _InitializedPools(
+                req_to_token_pool=bundle.req_to_token_pool,
+                token_to_kv_pool=bundle.token_to_kv_pool,
+                token_to_kv_pool_allocator=bundle.token_to_kv_pool_allocator,
+                unified_memory_pool=bundle.unified_memory_pool,
+            )
+
+        # Initialize req_to_token_pool
+        if req_to_token_pool is None:
+            req_to_token_pool = self._build_req_to_token_pool(
+                max_num_reqs=sizes.max_running_requests
+            )
+        else:
+            # Draft worker shares req_to_token_pool with the target worker.
+            assert self.is_draft_worker
+            # Each multi-layer EAGLE MTP head owns one transformer block at
+            # layer_id=draft_model_idx and needs its own sconv/mamba cache while
+            # sharing the target's request-to-token mapping.
+            if self.is_inkling_mtp_draft and isinstance(
+                req_to_token_pool, HybridReqToTokenPool
+            ):
+                # speculative_num_draft_tokens=None: draft heads never run
+                # TARGET_VERIFY, so their pools skip the per-step intermediate
+                # (SpeculativeState) buffers only the target pool consumes.
+                req_to_token_pool = req_to_token_pool.clone_with_new_mamba(
+                    mamba_size=self.server_args.max_mamba_cache_size,
+                    mamba_spec_state_size=sizes.max_running_requests,
+                    cache_params=self.mambaish_config.mamba2_cache_params,
+                    device=self.device,
+                    enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
+                    draft_model_idx=self.draft_model_idx,
+                    speculative_eagle_topk=self.server_args.speculative_eagle_topk,
+                )
+
+        # Initialize token_to_kv_pool
+        is_dsa_model = is_deepseek_dsa(self.model_config.hf_config)
+        is_dsv4_model = is_deepseek_v4(self.model_config.hf_config)
+
+        self._validate_prefill_only_disable_kv_cache_pool_family(
+            is_dsa_model, is_dsv4_model, current_platform
+        )
+
+        token_to_kv_pool = self._build_token_to_kv_pool(
+            sizes=sizes,
+            is_dsa_model=is_dsa_model,
+            is_dsv4_model=is_dsv4_model,
+            req_to_token_pool=req_to_token_pool,
+        )
+
+        token_to_kv_pool_allocator = self._build_token_to_kv_pool_allocator(
+            sizes=sizes,
+            token_to_kv_pool=token_to_kv_pool,
+            is_dsv4_model=is_dsv4_model,
+            req_to_token_pool=req_to_token_pool,
+            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+        )
+
+        # Defensive check: the explicit validation above should reject known
+        # unsupported pool families before allocation. Keep this guard here so
+        # future pool-selection refactors fail at boot instead of on first use.
+        if (
+            self.server_args.prefill_only_disable_kv_cache
+            and not self.is_draft_worker
+            and not isinstance(token_to_kv_pool, NoOpMHATokenToKVPool)
+        ):
+            raise RuntimeError(
+                "--prefill-only-disable-kv-cache expected NoOpMHATokenToKVPool but the "
+                f"runtime pool is {type(token_to_kv_pool).__name__}. This pool "
+                "family is not yet supported by --prefill-only-disable-kv-cache. "
+                "Supported configurations today: plain MHA models on CUDA with the FA "
+                "(fa3/fa4) prefill backend, --is-embedding, --chunked-prefill-size=-1, "
+                "--disable-radix-cache, no context-parallel attention, no HiSparse, "
+                "and --kv-cache-dtype not in {nvfp4, fp4_mx_block16}."
+            )
+        return _InitializedPools(
+            req_to_token_pool=req_to_token_pool,
+            token_to_kv_pool=token_to_kv_pool,
+            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+        )
+
+    def _init_unified_mamba_pools(
+        self, *, max_num_reqs: int, max_total_num_tokens: int
+    ) -> UnifiedPoolBundle:
+        """Build the shared-KV-pool stack for a hybrid-Mamba model:
+        one byte buffer split between the full-attn MHA KV pool and the
+        per-request Mamba state pool, with virtual slot ids above the
+        allocator."""
+        from sglang.srt.mem_cache.unified_memory_pool import init_unified_mamba_pools
+
+        config = self.mambaish_config
+        assert config is not None
+        assert (
+            not self.use_mla_backend
+        ), "unified memory pool does not support MLA-hybrid-Mamba yet"
+        # The full sub-pool is page-aware (via `MultiEndedAllocator(page_size=...)`);
+        # the mamba sub-pool stays page=1.
+        assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}"
+        # Mirror the non-shared path's extra_max_context_len computation.
+        extra_max_context_len = 4
+        if self.server_args.speculative_num_draft_tokens is not None:
+            extra_max_context_len += self.server_args.speculative_num_draft_tokens
+
+        mamba_layer_ids = [
+            i
+            for i in config.mamba2_cache_params.layers
+            if self.layer_info.start_layer <= i < self.layer_info.end_layer
+        ]
+        full_attention_layer_ids = [
+            i
+            for i in config.full_attention_layer_ids
+            if self.layer_info.start_layer <= i < self.layer_info.end_layer
+        ]
+
+        bundle = init_unified_mamba_pools(
+            device=self.device,
+            kv_cache_dtype=self.kv_cache_dtype,
+            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_dim=self.model_config.head_dim,
+            page_size=self.page_size,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+            is_draft_worker=self.is_draft_worker,
+            use_mla_backend=self.use_mla_backend,
+            mamba_layer_ids=mamba_layer_ids,
+            full_attention_layer_ids=full_attention_layer_ids,
+            mamba2_cache_params=config.mamba2_cache_params,
+            model_context_len=self.model_config.context_len,
+            extra_max_context_len=extra_max_context_len,
+            max_total_num_tokens=max_total_num_tokens,
+            max_mamba_cache_size=self.server_args.max_mamba_cache_size,
+            max_num_reqs=max_num_reqs,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
+            speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens,
+            disable_overlap_schedule=self.server_args.disable_overlap_schedule,
+            need_sort=self.server_args.disaggregation_mode in ("decode", "prefill"),
+            mamba_full_memory_ratio=self.server_args.mamba_full_memory_ratio,
+            # Overlap mode: the allocator's `free` drops a wait_stream(forward_stream)
+            # barrier so eager compaction serializes after the in-flight forward's
+            # v2p/KV reads. Near-no-op in normal mode.
+            forward_stream=self.forward_stream,
+            # Lazy compaction: default ON, env-var escape hatch for rollback / A/B.
+            lazy_compaction=_should_enable_lazy_compaction(),
+        )
+        return bundle
+
+    def _init_unified_swa_pools(
+        self,
+        *,
+        max_num_reqs: int,
+        full_max_total_num_tokens: Optional[int],
+        swa_max_total_num_tokens: Optional[int],
+    ) -> UnifiedPoolBundle:
+        """Build the unified-pool stack for a hybrid-SWA model (Triton): one byte
+        buffer split between the full-attention and SWA KV pools."""
+        from sglang.srt.mem_cache.unified_memory_pool import (
+            UnifiedPoolBundle,
+            init_unified_swa_pools,
+        )
+
+        assert self.is_hybrid_swa, "_init_unified_swa_pools called on a non-SWA model"
+        # Both sub-pools are page-aware; the SWA composite runs alloc_extend_kernel
+        # once in virtual space and binds the new pages on both sub-allocators.
+        assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}"
+        assert (
+            not self.use_mla_backend
+        ), "unified memory pool does not support MLA-SWA hybrid yet"
+        # Mirror the non-shared path's extra_max_context_len computation.
+        extra_max_context_len = 4
+        if self.server_args.speculative_num_draft_tokens is not None:
+            extra_max_context_len += self.server_args.speculative_num_draft_tokens
+        req_to_token_pool = ReqToTokenPool(
+            size=max_num_reqs,
+            max_context_len=self.model_config.context_len + extra_max_context_len,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+        )
+
+        head_num = self.model_config.get_num_kv_heads(get_parallel().attn_tp_size)
+        head_dim = self.model_config.head_dim
+        if self.is_hybrid_swa_compress:
+            # Asymmetric head dims between full and SWA (NPU compress path):
+            # pull SWA-specific dims from the hf text config.
+            v_head_dim = self.model_config.hf_text_config.v_head_dim
+            swa_head_num = max(
+                1,
+                self.model_config.hf_text_config.swa_num_key_value_heads
+                // get_parallel().attn_tp_size,
+            )
+            swa_head_dim = self.model_config.hf_text_config.swa_head_dim
+            swa_v_head_dim = self.model_config.hf_text_config.swa_v_head_dim
+        else:
+            v_head_dim = head_dim
+            swa_head_num = head_num
+            swa_head_dim = head_dim
+            swa_v_head_dim = head_dim
+
+        # Filter layer ids to this worker's [start_layer, end_layer) range.
+        swa_attention_layer_ids = [
+            i
+            for i in self.model_config.swa_attention_layer_ids
+            if self.layer_info.start_layer <= i < self.layer_info.end_layer
+        ]
+        full_attention_layer_ids = [
+            i
+            for i in self.model_config.full_attention_layer_ids
+            if self.layer_info.start_layer <= i < self.layer_info.end_layer
+        ]
+
+        bundle = init_unified_swa_pools(
+            device=self.device,
+            kv_cache_dtype=self.kv_cache_dtype,
+            head_num=head_num,
+            head_dim=head_dim,
+            v_head_dim=v_head_dim,
+            swa_head_num=swa_head_num,
+            swa_head_dim=swa_head_dim,
+            swa_v_head_dim=swa_v_head_dim,
+            page_size=self.page_size,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+            swa_attention_layer_ids=swa_attention_layer_ids,
+            full_attention_layer_ids=full_attention_layer_ids,
+            full_max_total_num_tokens=full_max_total_num_tokens,
+            swa_max_total_num_tokens=swa_max_total_num_tokens,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            need_sort=self.server_args.disaggregation_mode in ("decode", "prefill"),
+            # Overlap mode: same wait_stream(forward_stream) rationale as
+            # `_init_unified_mamba_pools`.
+            forward_stream=self.forward_stream,
+            # Lazy compaction: default ON, with env var escape hatch for rollback / A/B.
+            lazy_compaction=_should_enable_lazy_compaction(),
+        )
+        return UnifiedPoolBundle(
+            unified_memory_pool=bundle.unified_memory_pool,
+            token_to_kv_pool=bundle.token_to_kv_pool,
+            token_to_kv_pool_allocator=bundle.token_to_kv_pool_allocator,
+            req_to_token_pool=req_to_token_pool,
+        )
+
+    def _validate_prefill_only_disable_kv_cache_pool_family(
+        self,
+        is_dsa_model: bool,
+        is_dsv4_model: bool,
+        current_platform,
+    ):
+        if not self.server_args.prefill_only_disable_kv_cache or self.is_draft_worker:
+            return
+
+        unsupported_pool_family = None
+        if is_dsv4_model:
+            unsupported_pool_family = "DeepSeekV4TokenToKVPool"
+        elif current_platform.is_out_of_tree() and not self.mambaish_config:
+            unsupported_pool_family = "out-of-tree platform KV pool"
+        elif (
+            self.server_args.attention_backend == "ascend" and not self.mambaish_config
+        ):
+            unsupported_pool_family = "NPU/Ascend KV pool"
+        elif self.use_mla_backend and is_dsa_model:
+            unsupported_pool_family = "DSA/MLA KV pool"
+        elif self.use_mla_backend and not self.mambaish_config:
+            unsupported_pool_family = "MLA KV pool"
+        elif self.is_hybrid_swa:
+            unsupported_pool_family = "SWA KV pool"
+        elif self.mambaish_config:
+            unsupported_pool_family = "hybrid linear/Mamba KV pool"
+        elif is_float4_e2m1fn_x2(self.kv_cache_dtype):
+            unsupported_pool_family = "FP4 MHA KV pool"
+
+        if unsupported_pool_family is not None:
+            raise RuntimeError(
+                "--prefill-only-disable-kv-cache is not supported for "
+                f"{unsupported_pool_family}. Supported configurations today: plain MHA "
+                "models on CUDA with the FA (fa3/fa4) prefill backend, --is-embedding, "
+                "--chunked-prefill-size=-1, --disable-radix-cache, no context-parallel "
+                "attention, no HiSparse, and --kv-cache-dtype not in {nvfp4, fp4_mx_block16}."
+            )
+
+    def _build_req_to_token_pool(self, *, max_num_reqs: int) -> ReqToTokenPool:
+        extra_max_context_len = get_req_to_token_extra_context_len(self.server_args)
+
+        if self.server_args.disaggregation_mode == "decode":
+            # Extra slots for pre-allocated requests
+            pre_alloc_size = self.server_args.disaggregation_decode_extra_slots
+            if self.mambaish_config:
+                req_to_token_pool = self._build_hybrid_mamba_decode_req_pool(
+                    max_num_reqs=max_num_reqs,
+                    extra_max_context_len=extra_max_context_len,
+                    pre_alloc_size=pre_alloc_size,
+                )
+            else:
+                req_to_token_pool = self._build_decode_req_pool(
+                    max_num_reqs=max_num_reqs,
+                    extra_max_context_len=extra_max_context_len,
+                    pre_alloc_size=pre_alloc_size,
+                )
+        elif self.mambaish_config:
+            req_to_token_pool = self._build_hybrid_req_pool(
+                max_num_reqs=max_num_reqs,
+                extra_max_context_len=extra_max_context_len,
+            )
+        else:
+            req_to_token_pool = self._build_default_req_pool(
+                max_num_reqs=max_num_reqs,
+                extra_max_context_len=extra_max_context_len,
+            )
+        return req_to_token_pool
+
+    def _build_hybrid_mamba_decode_req_pool(
+        self,
+        *,
+        max_num_reqs: int,
+        extra_max_context_len: int,
+        pre_alloc_size: int,
+    ) -> ReqToTokenPool:
+        from sglang.srt.disaggregation.decode import (
+            HybridMambaDecodeReqToTokenPool,
+        )
+
+        req_to_token_pool = HybridMambaDecodeReqToTokenPool(
+            size=max_num_reqs,
+            max_context_len=self.model_config.context_len + extra_max_context_len,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            cache_params=self.mambaish_config.mamba2_cache_params,
+            mamba_layer_ids=(
+                [
+                    i
+                    for i in self.mambaish_config.mamba2_cache_params.layers
+                    if self.layer_info.start_layer <= i < self.layer_info.end_layer
+                ]
+            ),
+            speculative_num_draft_tokens=self.server_args.max_speculative_num_draft_tokens,
+            speculative_eagle_topk=self.server_args.speculative_eagle_topk,
+            enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
+            pre_alloc_size=pre_alloc_size,
+            enable_overlap_schedule=not self.server_args.disable_overlap_schedule,
+            mamba_size=self.server_args.max_mamba_cache_size,
+            start_layer=self.layer_info.start_layer,
+        )
+        return req_to_token_pool
+
+    def _build_decode_req_pool(
+        self,
+        *,
+        max_num_reqs: int,
+        extra_max_context_len: int,
+        pre_alloc_size: int,
+    ) -> ReqToTokenPool:
+        from sglang.srt.disaggregation.decode import DecodeReqToTokenPool
+
+        req_to_token_pool = DecodeReqToTokenPool(
+            size=max_num_reqs,
+            max_context_len=self.model_config.context_len + extra_max_context_len,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            pre_alloc_size=pre_alloc_size,
+        )
+        return req_to_token_pool
+
+    def _build_hybrid_req_pool(
+        self,
+        *,
+        max_num_reqs: int,
+        extra_max_context_len: int,
+    ) -> ReqToTokenPool:
+        req_to_token_pool = HybridReqToTokenPool(
+            size=max_num_reqs,
+            mamba_size=self.server_args.max_mamba_cache_size,
+            mamba_spec_state_size=max_num_reqs,
+            max_context_len=self.model_config.context_len + extra_max_context_len,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            cache_params=self.mambaish_config.mamba2_cache_params,
+            mamba_layer_ids=(
+                [
+                    i
+                    for i in self.mambaish_config.mamba2_cache_params.layers
+                    if self.layer_info.start_layer <= i < self.layer_info.end_layer
+                ]
+            ),
+            enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
+            enable_mamba_extra_buffer_lazy=self.server_args.enable_mamba_extra_buffer_lazy(),
+            speculative_num_draft_tokens=self.server_args.max_speculative_num_draft_tokens,
+            speculative_eagle_topk=self.server_args.speculative_eagle_topk,
+            enable_overlap_schedule=not self.server_args.disable_overlap_schedule,
+            start_layer=self.layer_info.start_layer,
+            enable_linear_replayssm=self.server_args.enable_linear_replayssm,
+            linear_replayssm_cache_len=self.server_args.linear_replayssm_cache_len,
+            mamba_envelope_layout=self.server_args.enable_page_major_kv_layout,
+            # ReplaySSM spec-verify is GDN-only: activate the pool machinery
+            # (rings + cursors + the intermediate_ssm gate) only for GDN-hybrid
+            # models, so any other mamba-ish model (Mamba2/Nemotron, lightning,
+            # ...) run with the flag set stays byte-identical to flag-off.
+            enable_gdn_replayssm_spec=(
+                self.server_args.enable_gdn_replayssm_spec
+                and self.hybrid_gdn_config is not None
+            ),
+        )
+        return req_to_token_pool
+
+    def _build_default_req_pool(
+        self,
+        *,
+        max_num_reqs: int,
+        extra_max_context_len: int,
+    ) -> ReqToTokenPool:
+        # DSV4 on NPU needs an extended ReqToTokenPool holding per-req
+        # swa/c4/c128/c{4,128}_state tables; others stay on the stock one.
+        req_to_token_pool_cls = ReqToTokenPool
+        if _is_npu and is_deepseek_v4(self.model_config.hf_config):
+            from sglang.srt.hardware_backend.npu.dsv4.dsv4_req_to_token_pool import (
+                DSV4NPUReqToTokenPool,
+            )
+
+            req_to_token_pool_cls = DSV4NPUReqToTokenPool
+
+        req_to_token_pool = req_to_token_pool_cls(
+            size=max_num_reqs,
+            max_context_len=self.model_config.context_len + extra_max_context_len,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+        )
+        return req_to_token_pool
+
+    def _build_token_to_kv_pool(
+        self,
+        *,
+        sizes: _PoolSizes,
+        is_dsa_model: bool,
+        is_dsv4_model: bool,
+        req_to_token_pool: ReqToTokenPool,
+    ) -> KVCache:
+        # Page-granularity envelope layout for the MHA-shaped (full / SWA) pools,
+        # selected by swapping in the PageMajorMHATokenToKVPool subclass. The
+        # default keeps upstream's per-layer layout. The Mamba state pool is routed
+        # separately via `mamba_envelope_layout` on the req-to-token pool above.
+        enable_page_major = self.server_args.enable_page_major_kv_layout
+        mha_pool_class = (
+            PageMajorMHATokenToKVPool if enable_page_major else MHATokenToKVPool
+        )
+
+        if is_dsv4_model:
+            token_to_kv_pool = self._build_dsv4_kv_pool(
+                max_running_requests=sizes.max_running_requests,
+                swa_max_total_num_tokens=sizes.swa_max_total_num_tokens,
+                c4_max_total_num_tokens=sizes.c4_max_total_num_tokens,
+                c128_max_total_num_tokens=sizes.c128_max_total_num_tokens,
+                c4_state_pool_size=sizes.c4_state_pool_size,
+                c128_state_pool_size=sizes.c128_state_pool_size,
+                c4_state_dtype=sizes.c4_state_dtype,
+                c128_state_dtype=sizes.c128_state_dtype,
+                req_to_token_pool=req_to_token_pool,
+            )
+        elif current_platform.is_out_of_tree() and not self.mambaish_config:
+            if self.use_mla_backend and is_dsa_model:
+                token_to_kv_pool = self._build_oot_dsa_kv_pool(
+                    max_total_num_tokens=sizes.max_total_num_tokens,
+                )
+            elif self.use_mla_backend:
+                token_to_kv_pool = self._build_oot_mla_kv_pool(
+                    max_total_num_tokens=sizes.max_total_num_tokens,
+                    is_dsa_model=is_dsa_model,
+                )
+            else:
+                token_to_kv_pool = self._build_oot_mha_kv_pool(
+                    max_total_num_tokens=sizes.max_total_num_tokens,
+                )
+        elif (
+            self.server_args.attention_backend == "ascend" and not self.mambaish_config
+        ):
+            if self.is_hybrid_swa:
+                token_to_kv_pool = self._build_ascend_swa_kv_pool(
+                    full_max_total_num_tokens=sizes.full_max_total_num_tokens,
+                    swa_max_total_num_tokens=sizes.swa_max_total_num_tokens,
+                )
+            elif self.use_mla_backend:
+                token_to_kv_pool = self._build_ascend_mla_kv_pool(
+                    max_total_num_tokens=sizes.max_total_num_tokens,
+                    is_dsa_model=is_dsa_model,
+                )
+            else:
+                token_to_kv_pool = self._build_ascend_mha_kv_pool(
+                    max_total_num_tokens=sizes.max_total_num_tokens,
+                )
+        elif self.use_mla_backend and is_dsa_model:
+            token_to_kv_pool = self._build_dsa_kv_pool(
+                max_total_num_tokens=sizes.max_total_num_tokens,
+            )
+        elif self.use_mla_backend and not self.mambaish_config:
+            assert not is_dsa_model
+            if is_float4_e2m1fn_x2(self.kv_cache_dtype):
+                token_to_kv_pool = self._build_mla_fp4_kv_pool(
+                    max_total_num_tokens=sizes.max_total_num_tokens,
+                )
+            else:
+                token_to_kv_pool = self._build_mla_kv_pool(
+                    max_total_num_tokens=sizes.max_total_num_tokens,
+                )
+        else:
+            if self.is_hybrid_swa:
+                token_to_kv_pool = self._build_hybrid_swa_kv_pool(
+                    full_max_total_num_tokens=sizes.full_max_total_num_tokens,
+                    swa_max_total_num_tokens=sizes.swa_max_total_num_tokens,
+                    mha_pool_class=mha_pool_class,
+                )
+            elif is_minimax_sparse(self.model_config.hf_config):
+                token_to_kv_pool = self._build_minimax_sparse_kv_pool(
+                    max_total_num_tokens=sizes.max_total_num_tokens,
+                )
+            elif self.mambaish_config:
+                token_to_kv_pool = self._build_hybrid_linear_kv_pool(
+                    max_total_num_tokens=sizes.max_total_num_tokens,
+                    req_to_token_pool=req_to_token_pool,
+                    mha_pool_class=mha_pool_class,
+                )
+            else:
+                quant_method = None
+                if is_float4_e2m1fn_x2(self.kv_cache_dtype):
+                    assert (
+                        not enable_page_major
+                    ), "page-major KV layout is not supported with fp4 KV cache"
+                    quant_method = self._build_fp4_quant_method(
+                        num_layers=self.layer_info.num_effective_layers
+                    )
+                token_to_kv_pool = self._build_mha_kv_pool(
+                    max_total_num_tokens=sizes.max_total_num_tokens,
+                    mha_pool_class=mha_pool_class,
+                    quant_method=quant_method,
+                )
+        return token_to_kv_pool
+
+    def _build_dsv4_kv_pool(
+        self,
+        *,
+        max_running_requests: int,
+        swa_max_total_num_tokens: Optional[int],
+        c4_max_total_num_tokens: int,
+        c128_max_total_num_tokens: int,
+        c4_state_pool_size: int,
+        c128_state_pool_size: int,
+        c4_state_dtype: Optional[torch.dtype],
+        c128_state_dtype: Optional[torch.dtype],
+        req_to_token_pool: ReqToTokenPool,
+    ) -> KVCache:
+        swa_page_size = self.server_args.page_size
+        if not _is_npu:
+            assert swa_page_size == 256, "In paged swa mode, page_size must be 256."
+
+        if self.is_draft_worker:
+            from sglang.srt.models.agnes_nextn import (
+                COMPRESS_RATIO_NEXTN_LAYER,
+            )
+
+            compression_ratios = [
+                COMPRESS_RATIO_NEXTN_LAYER
+            ] * self.layer_info.num_effective_layers
+        else:
+            compression_ratios = self.model_config.compress_ratios
+
+        # NPU + DSV4 → paged-state subclass: the fused compressor kernel
+        # needs cache_mode=1 (paged); Atlas A3 rejects cache_mode=2 (ring),
+        # so the CUDA ring-buffer state path can't be shared. CUDA keeps
+        # DeepSeekV4TokenToKVPool unchanged; NPU recomputes state sizes below.
+        if _is_npu:
+            from sglang.srt.hardware_backend.npu.dsv4.dsv4_memory_pool import (
+                DSV4NPUTokenToKVPool,
+                npu_state_pool_size,
+            )
+
+            pool_cls = DSV4NPUTokenToKVPool
+            # Recompute state pool sizes for the NPU paged formula (CUDA's
+            # ring sizes are dropped here). Tail-only allocation keeps the
+            # per-req-budget formula sufficient at any prefill length: long
+            # prompts allocate only ``tail+128`` (c4) / ``tail`` (c128)
+            # slots (tail = seq_len % 128), and decode is drained by
+            # sliding eviction in ``ScheduleBatch._evict_swa``.
+            c4_state_pool_size = npu_state_pool_size(
+                ratio=4,
+                page_size=self.server_args.page_size,
+                max_num_reqs=max_running_requests,
+            )
+            c128_state_pool_size = npu_state_pool_size(
+                ratio=128,
+                page_size=self.server_args.page_size,
+                max_num_reqs=max_running_requests,
+            )
+        else:
+            pool_cls = DeepSeekV4TokenToKVPool
+            c4_state_pool_size = c4_state_pool_size
+            c128_state_pool_size = c128_state_pool_size
+
+        token_to_kv_pool = pool_cls(
+            max_num_reqs=max_running_requests,
+            # SWA ring is indexed by req_pool_idx; PD decode inflates req_to_token
+            # past max_running_requests (pre-alloc), so size to the real capacity.
+            num_req_slots=req_to_token_pool.req_to_token.shape[0],
+            swa_size=swa_max_total_num_tokens,
+            c4_size=c4_max_total_num_tokens,
+            c128_size=c128_max_total_num_tokens,
+            c4_state_pool_size=c4_state_pool_size,
+            c128_state_pool_size=c128_state_pool_size,
+            page_size=self.server_args.page_size,
+            swa_page_size=swa_page_size,
+            sliding_window=self.model_config.window_size,
+            dtype=self.kv_cache_dtype,
+            c4_state_dtype=c4_state_dtype,
+            c128_state_dtype=c128_state_dtype,
+            qk_nope_head_dim=self.model_config.qk_nope_head_dim,
+            qk_rope_head_dim=self.model_config.qk_rope_head_dim,
+            indexer_head_dim=self.model_config.index_head_dim,
+            layer_num=self.layer_info.num_effective_layers,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            compression_ratios=compression_ratios,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+            enable_hisparse=self.server_args.enable_hisparse,
+            online_mtp_max_draft_tokens=(
+                self.server_args.max_speculative_num_draft_tokens or 0
+            ),
+        )
+        return token_to_kv_pool
+
+    def _build_oot_dsa_kv_pool(self, *, max_total_num_tokens: int) -> KVCache:
+        PoolCls = current_platform.get_dsa_kv_pool_cls()
+        token_to_kv_pool = PoolCls(
+            max_total_num_tokens,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            kv_lora_rank=self.model_config.kv_lora_rank,
+            qk_rope_head_dim=self.model_config.qk_rope_head_dim,
+            layer_num=self.layer_info.num_effective_layers,
+            device=self.device,
+            kv_cache_dim=calculate_mla_kv_cache_dim(
+                model_config=self.model_config,
+                kv_cache_dtype=self.kv_cache_dtype,
+                server_args=self.server_args,
+            ),
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+            index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config),
+        )
+        return token_to_kv_pool
+
+    def _build_oot_mla_kv_pool(
+        self, *, max_total_num_tokens: int, is_dsa_model: bool
+    ) -> KVCache:
+        PoolCls = current_platform.get_mla_kv_pool_cls()
+        token_to_kv_pool = PoolCls(
+            max_total_num_tokens,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            kv_lora_rank=self.model_config.kv_lora_rank,
+            qk_rope_head_dim=self.model_config.qk_rope_head_dim,
+            index_head_dim=(self.model_config.index_head_dim if is_dsa_model else None),
+            layer_num=self.layer_info.num_effective_layers,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+        )
+        return token_to_kv_pool
+
+    def _build_oot_mha_kv_pool(self, *, max_total_num_tokens: int) -> KVCache:
+        PoolCls = current_platform.get_mha_kv_pool_cls()
+        token_to_kv_pool = PoolCls(
+            max_total_num_tokens,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_dim=self.model_config.head_dim,
+            layer_num=self.layer_info.num_effective_layers,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+        )
+        return token_to_kv_pool
+
+    def _build_ascend_swa_kv_pool(
+        self,
+        *,
+        full_max_total_num_tokens: Optional[int],
+        swa_max_total_num_tokens: Optional[int],
+    ) -> KVCache:
+        from sglang.srt.hardware_backend.npu.memory_pool_npu import (
+            NPUMHATokenToKVPool,
+        )
+
+        kwargs = {}
+        if self.is_hybrid_swa_compress:
+            kwargs = {
+                "swa_head_num": max(
+                    1,
+                    self.model_config.hf_text_config.swa_num_key_value_heads
+                    // get_parallel().attn_tp_size,
+                ),
+                "swa_head_dim": self.model_config.swa_head_dim,
+                "swa_v_head_dim": self.model_config.swa_v_head_dim,
+                "v_head_dim": self.model_config.v_head_dim,
+            }
+        token_to_kv_pool = SWAKVPool(
+            size=full_max_total_num_tokens,
+            size_swa=swa_max_total_num_tokens,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            post_capture_active=self.post_capture_kv_active,
+            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_dim=self.model_config.head_dim,
+            swa_attention_layer_ids=self.model_config.swa_attention_layer_ids,
+            full_attention_layer_ids=self.model_config.full_attention_layer_ids,
+            device=self.device,
+            token_to_kv_pool_class=NPUMHATokenToKVPool,
+            **kwargs,
+        )
+        return token_to_kv_pool
+
+    def _build_ascend_mla_kv_pool(
+        self, *, max_total_num_tokens: int, is_dsa_model: bool
+    ) -> KVCache:
+        from sglang.srt.hardware_backend.npu.memory_pool_npu import (
+            NPUMLATokenToKVPool,
+        )
+
+        token_to_kv_pool = NPUMLATokenToKVPool(
+            max_total_num_tokens,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            kv_lora_rank=self.model_config.kv_lora_rank,
+            qk_rope_head_dim=self.model_config.qk_rope_head_dim,
+            index_head_dim=(self.model_config.index_head_dim if is_dsa_model else None),
+            layer_num=self.layer_info.num_effective_layers,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+        )
+        return token_to_kv_pool
+
+    def _build_ascend_mha_kv_pool(self, *, max_total_num_tokens: int) -> KVCache:
+        from sglang.srt.hardware_backend.npu.memory_pool_npu import (
+            NPUMHATokenToKVPool,
+        )
+
+        token_to_kv_pool = NPUMHATokenToKVPool(
+            max_total_num_tokens,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_dim=self.model_config.head_dim,
+            layer_num=self.layer_info.num_effective_layers,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+        )
+        return token_to_kv_pool
+
+    def _build_dsa_kv_pool(self, *, max_total_num_tokens: int) -> KVCache:
+        from sglang.srt.layers.cp.utils import get_glm_dsa_cp_layer_shard_info
+
+        (
+            dsa_cp_layer_shard_rank,
+            dsa_cp_layer_shard_size,
+        ) = get_glm_dsa_cp_layer_shard_info(self)
+        pool_kwargs = {}
+        if self.server_args.enable_hisparse:
+            PoolCls = HiSparseDSATokenToKVPool
+            from sglang.srt.mem_cache.sparsity import parse_hisparse_config
+
+            pool_kwargs["host_to_device_ratio"] = parse_hisparse_config(
+                self.server_args
+            ).host_to_device_ratio
+        elif dsa_cp_layer_shard_rank is not None:
+            # DSA cache layer split: shard KV/indexer layers across CP ranks.
+            from sglang.srt.mem_cache.dsa_cache_layer_split import (
+                LayerSplitDSATokenToKVPool,
+            )
+
+            PoolCls = LayerSplitDSATokenToKVPool
+            pool_kwargs["layer_shard_rank"] = dsa_cp_layer_shard_rank
+            pool_kwargs["layer_shard_size"] = dsa_cp_layer_shard_size
+        else:
+            PoolCls = DSATokenToKVPool
+        token_to_kv_pool = PoolCls(
+            max_total_num_tokens,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            kv_lora_rank=self.model_config.kv_lora_rank,
+            qk_rope_head_dim=self.model_config.qk_rope_head_dim,
+            layer_num=self.layer_info.num_effective_layers,
+            device=self.device,
+            kv_cache_dim=calculate_mla_kv_cache_dim(
+                model_config=self.model_config,
+                kv_cache_dtype=self.kv_cache_dtype,
+                server_args=self.server_args,
+            ),
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+            index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config),
+            **pool_kwargs,
+        )
+        return token_to_kv_pool
+
+    def _build_mla_fp4_kv_pool(self, *, max_total_num_tokens: int) -> KVCache:
+        token_to_kv_pool = MLATokenToKVPoolFP4(
+            max_total_num_tokens,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            kv_lora_rank=self.model_config.kv_lora_rank,
+            qk_rope_head_dim=self.model_config.qk_rope_head_dim,
+            layer_num=self.layer_info.num_effective_layers,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+        )
+        return token_to_kv_pool
+
+    def _build_mla_kv_pool(self, *, max_total_num_tokens: int) -> KVCache:
+        token_to_kv_pool = MLATokenToKVPool(
+            max_total_num_tokens,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            kv_lora_rank=self.model_config.kv_lora_rank,
+            qk_rope_head_dim=self.model_config.qk_rope_head_dim,
+            layer_num=self.layer_info.num_effective_layers,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+        )
+        return token_to_kv_pool
+
+    def _build_hybrid_swa_kv_pool(
+        self,
+        *,
+        full_max_total_num_tokens: Optional[int],
+        swa_max_total_num_tokens: Optional[int],
+        mha_pool_class: type,
+    ) -> KVCache:
+        kwargs = {}
+        if self.is_hybrid_swa_compress:
+            kwargs = {
+                "swa_head_num": max(
+                    1,
+                    self.model_config.hf_text_config.swa_num_key_value_heads
+                    // get_parallel().attn_tp_size,
+                ),
+                "swa_head_dim": self.model_config.swa_head_dim,
+                "swa_v_head_dim": self.model_config.swa_v_head_dim,
+                "v_head_dim": self.model_config.v_head_dim,
+            }
+        swa_pool_class = (
+            MHATokenToKVPoolMXFP8
+            if self.server_args.kv_cache_dtype == "mxfp8"
+            else mha_pool_class
+        )
+        swa_attention_layer_ids = self.model_config.swa_attention_layer_ids
+        full_attention_layer_ids = self.model_config.full_attention_layer_ids
+        if self.is_inkling_mtp_draft:
+            if self.draft_swa_full_capacity:
+                # Banded 's' depth: route the draft's single layer into the SWA
+                # ring sub-pool so use_sliding_window_kv_pool activates the SWA
+                # store/read path for this depth, exactly like a trunk local
+                # layer.
+                swa_attention_layer_ids = [self.draft_model_idx]
+                full_attention_layer_ids = []
+            else:
+                swa_attention_layer_ids = []
+                full_attention_layer_ids = [self.draft_model_idx]
+        # Size the banded draft's SWA ring to FULL draft capacity (not the
+        # trunk-window-derived swa_max): with the identity full->swa mapping
+        # registered in _build_token_to_kv_pool_allocator, every logical slot
+        # the shared target allocator hands out (up to full_max) must be
+        # addressable in the ring, whatever the head-vs-trunk window
+        # relationship.
+        size_swa = (
+            full_max_total_num_tokens
+            if self.draft_swa_full_capacity
+            else swa_max_total_num_tokens
+        )
+        token_to_kv_pool = SWAKVPool(
+            size=full_max_total_num_tokens,
+            size_swa=size_swa,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            post_capture_active=self.post_capture_kv_active,
+            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_dim=self.model_config.head_dim,
+            swa_attention_layer_ids=swa_attention_layer_ids,
+            full_attention_layer_ids=full_attention_layer_ids,
+            device=self.device,
+            enable_kv_cache_copy=(self.server_args.speculative_algorithm is not None),
+            token_to_kv_pool_class=swa_pool_class,
+            **kwargs,
+        )
+        return token_to_kv_pool
+
+    def _build_minimax_sparse_kv_pool(self, *, max_total_num_tokens: int) -> KVCache:
+        _hf_config = self.model_config.hf_config
+        sparse_cfg = get_minimax_sparse_attention_config(_hf_config)
+        dense_layer_ids, sparse_layer_ids = get_minimax_sparse_layer_ids(sparse_cfg)
+        disable_value_sparse_layer_ids = get_minimax_sparse_disable_value_layer_ids(
+            sparse_cfg
+        )
+        token_to_kv_pool = MiniMaxSparseKVPool(
+            size=max_total_num_tokens,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            index_dtype=self.model_dtype,
+            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_dim=self.model_config.head_dim,
+            idx_head_dim=sparse_cfg["sparse_index_dim"],
+            dense_layer_ids=dense_layer_ids,
+            sparse_layer_ids=sparse_layer_ids,
+            disable_value_sparse_layer_ids=disable_value_sparse_layer_ids,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+        )
+        return token_to_kv_pool
+
+    def _build_hybrid_linear_kv_pool(
+        self,
+        *,
+        max_total_num_tokens: int,
+        req_to_token_pool: ReqToTokenPool,
+        mha_pool_class: type,
+    ) -> KVCache:
+        extra_args = {}
+        if self.use_mla_backend:
+            extra_args = {
+                "kv_lora_rank": self.model_config.kv_lora_rank,
+                "qk_rope_head_dim": self.model_config.qk_rope_head_dim,
+            }
+        full_attention_layer_ids = (
+            [0]
+            if self.is_draft_worker
+            else [
+                i
+                for i in self.mambaish_config.full_attention_layer_ids
+                if self.layer_info.start_layer <= i < self.layer_info.end_layer
+            ]
+        )
+        quant_method = self._build_fp4_quant_method(
+            num_layers=len(full_attention_layer_ids)
+        )
+        # MXFP8 KV cache needs the block-scaled pool (data + UE8M0 scale
+        # buffers) for the full-attention layers, same as the SWA branch.
+        full_pool_class = (
+            MHATokenToKVPoolMXFP8
+            if self.server_args.kv_cache_dtype == "mxfp8" and not self.use_mla_backend
+            else mha_pool_class
+        )
+        token_to_kv_pool = HybridLinearKVPool(
+            page_size=self.server_args.page_size,
+            size=max_total_num_tokens,
+            dtype=self.kv_cache_dtype,
+            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_dim=self.model_config.head_dim,
+            # if draft worker, we only need 1 attention layer's kv pool
+            full_attention_layer_ids=full_attention_layer_ids,
+            device=self.device,
+            mamba_pool=req_to_token_pool.mamba_pool,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            enable_kv_cache_copy=(self.server_args.speculative_algorithm is not None),
+            use_mla=self.use_mla_backend,
+            start_layer=self.layer_info.start_layer,
+            full_kv_pool_class=full_pool_class,
+            quant_method=quant_method,
+            post_capture_active=self.post_capture_kv_active and quant_method is None,
+            **extra_args,
+        )
+        return token_to_kv_pool
+
+    def _build_mha_fp4_kv_pool(self, *, max_total_num_tokens: int) -> KVCache:
+        token_to_kv_pool = MHATokenToKVPoolFP4(
+            max_total_num_tokens,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_dim=self.model_config.head_dim,
+            v_head_dim=self.model_config.v_head_dim,
+            layer_num=self.layer_info.num_effective_layers,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+            enable_alt_stream=not self.server_args.enable_pdmux,
+            enable_kv_cache_copy=(self.server_args.speculative_algorithm is not None),
+        )
+        return token_to_kv_pool
+
+    def _build_mha_kv_pool(
+        self, *, max_total_num_tokens: int, mha_pool_class: type, quant_method=None
+    ) -> KVCache:
+        if self.server_args.kv_cache_dtype == "mxfp8":
+            pool_cls = MHATokenToKVPoolMXFP8
+        else:
+            pool_cls = (
+                NoOpMHATokenToKVPool
+                if self.server_args.prefill_only_disable_kv_cache
+                else mha_pool_class
+            )
+        pool_kwargs = {}
+        if quant_method is not None:
+            pool_kwargs["quant_method"] = quant_method
+        else:
+            pool_kwargs["post_capture_active"] = self.post_capture_kv_active
+        token_to_kv_pool = pool_cls(
+            max_total_num_tokens,
+            page_size=self.server_args.page_size,
+            dtype=self.kv_cache_dtype,
+            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_dim=self.model_config.head_dim,
+            v_head_dim=self.model_config.v_head_dim,
+            layer_num=self.layer_info.num_effective_layers,
+            device=self.device,
+            enable_memory_saver=self.server_args.enable_memory_saver,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+            enable_alt_stream=not self.server_args.enable_pdmux,
+            enable_kv_cache_copy=(self.server_args.speculative_algorithm is not None),
+            **pool_kwargs,
+        )
+        return token_to_kv_pool
+
+    def _build_token_to_kv_pool_allocator(
+        self,
+        *,
+        sizes: _PoolSizes,
+        token_to_kv_pool: KVCache,
+        is_dsv4_model: bool,
+        req_to_token_pool: ReqToTokenPool,
+        token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator],
+    ) -> BaseTokenToKVPoolAllocator:
+        # Initialize token_to_kv_pool_allocator
+        need_sort = self.server_args.disaggregation_mode in ("decode", "prefill")
+        if token_to_kv_pool_allocator is None:
+            if current_platform.is_out_of_tree():
+                AllocatorCls = current_platform.get_paged_allocator_cls()
+                token_to_kv_pool_allocator = AllocatorCls(
+                    sizes.max_total_num_tokens,
+                    page_size=self.server_args.page_size,
+                    dtype=self.kv_cache_dtype,
+                    device=self.device,
+                    kvcache=token_to_kv_pool,
+                    need_sort=need_sort,
+                )
+            elif _is_npu and (
+                self.server_args.attention_backend == "ascend"
+                or is_dsv4_model
+                or self.hybrid_gdn_config is not None
+            ):
+                if self.is_hybrid_swa:
+                    # DSV4 on NPU: SWA allocator subclass that also drives the
+                    # c4/c128 allocators, producing a DSV4OutCacheLoc per alloc.
+                    if is_dsv4_model:
+                        from sglang.srt.hardware_backend.npu.dsv4.dsv4_allocator import (
+                            DSV4NPUTokenToKVPoolAllocator,
+                        )
+
+                        swa_allocator_cls = DSV4NPUTokenToKVPoolAllocator
+                    else:
+                        swa_allocator_cls = SWATokenToKVPoolAllocator
+                    token_to_kv_pool_allocator = swa_allocator_cls(
+                        sizes.full_max_total_num_tokens,
+                        sizes.swa_max_total_num_tokens,
+                        page_size=self.server_args.page_size,
+                        dtype=self.kv_cache_dtype,
+                        device=self.device,
+                        kvcache=token_to_kv_pool,
+                        need_sort=need_sort,
+                    )
+                else:
+                    from sglang.srt.hardware_backend.npu.allocator_npu import (
+                        NPUPagedTokenToKVPoolAllocator,
+                    )
+
+                    token_to_kv_pool_allocator = NPUPagedTokenToKVPoolAllocator(
+                        sizes.max_total_num_tokens,
+                        page_size=self.server_args.page_size,
+                        dtype=self.kv_cache_dtype,
+                        device=self.device,
+                        kvcache=token_to_kv_pool,
+                        need_sort=need_sort,
+                    )
+            else:
+                if self.is_hybrid_swa and sizes.full_max_total_num_tokens == 0:
+                    token_to_kv_pool_allocator = PureSWATokenToKVPoolAllocator(
+                        sizes.swa_max_total_num_tokens,
+                        page_size=self.server_args.page_size,
+                        dtype=self.kv_cache_dtype,
+                        device=self.device,
+                        kvcache=token_to_kv_pool,
+                        need_sort=need_sort,
+                    )
+                elif self.is_hybrid_swa:
+                    token_to_kv_pool_allocator = SWATokenToKVPoolAllocator(
+                        sizes.full_max_total_num_tokens,
+                        sizes.swa_max_total_num_tokens,
+                        page_size=self.server_args.page_size,
+                        dtype=self.kv_cache_dtype,
+                        device=self.device,
+                        kvcache=token_to_kv_pool,
+                        need_sort=need_sort,
+                    )
+                else:
+                    if self.server_args.enable_hisparse:
+                        from sglang.srt.mem_cache.sparsity import (
+                            parse_hisparse_config,
+                        )
+
+                        hisparse_cfg = parse_hisparse_config(self.server_args)
+                        token_to_kv_pool_allocator = HiSparseTokenToKVPoolAllocator(
+                            sizes.max_total_num_tokens,
+                            page_size=self.server_args.page_size,
+                            dtype=self.kv_cache_dtype,
+                            device=self.device,
+                            kvcache=token_to_kv_pool,
+                            need_sort=need_sort,
+                            host_to_device_ratio=hisparse_cfg.host_to_device_ratio,
+                        )
+                    elif (
+                        self.server_args.page_size == 1
+                        and self.server_args.dcp_size == 1
+                    ):
+                        token_to_kv_pool_allocator = TokenToKVPoolAllocator(
+                            sizes.max_total_num_tokens,
+                            dtype=self.kv_cache_dtype,
+                            device=self.device,
+                            kvcache=token_to_kv_pool,
+                            need_sort=need_sort,
+                        )
+                    else:
+                        token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator(
+                            sizes.max_total_num_tokens * self.server_args.dcp_size,
+                            page_size=self.server_args.page_size
+                            * self.server_args.dcp_size,
+                            dtype=self.kv_cache_dtype,
+                            device=self.device,
+                            kvcache=token_to_kv_pool,
+                            need_sort=need_sort,
+                        )
+
+            if self.server_args.enable_hisparse and is_dsv4_model:
+                assert self.is_hybrid_swa, "DeepSeek V4 HiSparse requires SWA mode."
+                token_to_kv_pool_allocator = DeepSeekV4HiSparseTokenToKVPoolAllocator(
+                    token_to_kv_pool_allocator
+                )
+
+            # DSV4-NPU: wire allocator back-ref into req_to_token_pool so its
+            # free(req) can release c4/c128 pool pages alongside the slot.
+            if hasattr(req_to_token_pool, "register_dsv4_allocator"):
+                req_to_token_pool.register_dsv4_allocator(token_to_kv_pool_allocator)
+
+        else:
+            assert self.is_draft_worker
+            if self.is_hybrid_swa:
+                if self.draft_swa_full_capacity:
+                    # Banded depth: the SWA ring is full draft capacity, so use
+                    # an IDENTITY full->swa mapping — store and read locs both
+                    # equal out_cache_loc, and a slot is never evicted before
+                    # the request frees it. The window itself is enforced by the
+                    # FA sliding-window kernel, not by the ring. Layout mirrors
+                    # SWATokenToKVPoolAllocator's mapping (size + page_size
+                    # entries + trailing -1 sentinel so a -1 last_loc maps
+                    # to -1).
+                    n = sizes.full_max_total_num_tokens + self.page_size
+                    identity_mapping = torch.arange(
+                        n + 1, dtype=torch.int64, device=self.device
+                    )
+                    identity_mapping[-1] = -1
+                    token_to_kv_pool.register_mapping(identity_mapping)
+                else:
+                    swa_allocator = getattr(
+                        token_to_kv_pool_allocator,
+                        "logical_attn_allocator",
+                        token_to_kv_pool_allocator,
+                    )
+                    assert isinstance(swa_allocator, SWATokenToKVPoolAllocator)
+                    token_to_kv_pool.register_mapping(
+                        swa_allocator.full_to_swa_index_mapping
+                    )
+        return token_to_kv_pool_allocator
+
+    def _profile_available_bytes(self, pre_model_load_memory: int) -> int:
+        # KV pool budget = currently-free GPU memory minus the non-static runtime
+        # slack (pre_model_load_memory * (1 - mem_fraction_static)). Whatever is
+        # already resident (model weights, etc.) is thus charged against it.
+        available_gpu_memory = get_available_gpu_memory(
+            self.device,
+            self.gpu_id,
+            distributed=get_world_group().world_size > 1,
+            cpu_group=get_world_group().cpu_group,
+        )
+
+        slack_gb = pre_model_load_memory * (1 - self.server_args.mem_fraction_static)
+        if self.mambaish_config is not None and self.post_capture_kv_active:
+            # Mamba state is a fixed pre-capture allocation, so it can't ride the ~0 post-capture slack.
+            slack_gb = max(
+                slack_gb,
+                self.server_args.mamba_pre_capture_reserve_mb(
+                    get_device_memory_capacity(self.device)
+                )
+                / 1024,
+            )
+        rest_memory = available_gpu_memory - slack_gb
+        if self.mambaish_config is not None:
+            rest_memory = self._handle_max_mamba_cache(rest_memory)
+
+        # Loaded weights (target + draft) can exceed the static budget
+        if rest_memory <= 0:
+            minimum_mem_fraction_static = (
+                1 - available_gpu_memory / pre_model_load_memory
+            )
+            suggested_mem_fraction_static = (
+                math.ceil(minimum_mem_fraction_static * 1000) / 1000
+            )
+            raise ValueError(
+                f"Loaded weights leave no GPU memory for the KV cache under "
+                f"--mem-fraction-static={self.server_args.mem_fraction_static}. "
+                f"Raise --mem-fraction-static above "
+                f"{suggested_mem_fraction_static:.3f} "
+                f"(minimum viable = 1 - available/pre = "
+                f"{minimum_mem_fraction_static:.4f}). If using speculative "
+                f"decoding, draft weights are now counted."
+            )
+
+        return int(rest_memory * (1 << 30))  # return in bytes
+
+    def _calculate_mamba_ratio(self) -> int:
+        if self.server_args.disable_radix_cache:
+            return 1
+
+        additional_ratio = 0
+        if self.server_args.enable_mamba_extra_buffer():
+            # ping-pong buffer size is 2 when overlap schedule is on, 1 otherwise.
+            # Lazy mode saves 1 slot (2 → 1) for overlap; non-overlap already uses 1.
+            if not self.server_args.disable_overlap_schedule:
+                if self.server_args.enable_mamba_extra_buffer_lazy():
+                    additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP_LAZY
+                else:
+                    additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP
+            else:
+                assert (
+                    not self.server_args.enable_mamba_extra_buffer_lazy()
+                ), "Lazy extra buffer requires overlap schedule (--disable-overlap-schedule is incompatible)"
+                additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP
+
+        return MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO + additional_ratio
+
+    def _apply_token_constraints(self, token_capacity: int) -> int:
+        """Apply external constraints to token capacity: user cap, PP sync.
+
+        Page alignment is handled by the configurator, not here.
+        If constraints change the value, the configurator re-runs and re-aligns.
+        """
+        user_limit = self.server_args.max_total_tokens
+
+        # Apply user-specified upper bound
+        if user_limit is not None:
+            if user_limit > token_capacity:
+                logging.warning(
+                    f"max_total_tokens={user_limit} is larger than the profiled value "
+                    f"{token_capacity}. Use the profiled value instead."
+                )
+            token_capacity = min(token_capacity, user_limit)
+
+        # Sync across PP ranks (each may have different layer counts)
+        if self.server_args.pp_size > 1:
+            tensor = torch.tensor(token_capacity, dtype=torch.int64)
+            torch.distributed.all_reduce(
+                tensor,
+                op=torch.distributed.ReduceOp.MIN,
+                group=get_world_group().cpu_group,
+            )
+            token_capacity = tensor.item()
+
+        return token_capacity
+
+    def resolve_max_num_reqs(self, token_capacity: int) -> int:
+        """Compute max concurrent requests (per dp worker) from the finalized
+        token capacity."""
+        # Estimate pool size (used as upper bound when user specifies max_running_requests)
+        estimated = int(token_capacity / self.model_config.context_len * 512)
+        estimated = max(min(estimated, 4096), 2048)
+
+        max_num_reqs = self.server_args.max_running_requests
+        if max_num_reqs is not None:
+            requested_per_worker = max_num_reqs // self.ps.attn_dp_size
+            max_num_reqs = min(requested_per_worker, token_capacity // 2)
+        else:
+            requested_per_worker = None
+            max_num_reqs = min(estimated, token_capacity // 2)
+
+        if self.mambaish_config is not None:
+            ratio = self._calculate_mamba_ratio()
+            max_num_reqs = min(
+                max_num_reqs, self.server_args.max_mamba_cache_size // ratio
+            )
+
+            if max_num_reqs <= 0:
+                raise RuntimeError(
+                    f"Hybrid (mamba/linear-attention) state cache is too small to serve "
+                    f"any requests. max_mamba_cache_size={self.server_args.max_mamba_cache_size}, "
+                    f"mamba_ratio={ratio}, resulting max_num_reqs={max_num_reqs}. "
+                    f"Try: (1) reduce --max-running-requests, "
+                    f"(2) increase --mem-fraction-static, or "
+                    f"(3) use GPUs with more memory."
+                )
+        if requested_per_worker is not None and max_num_reqs < requested_per_worker:
+            logger.warning(
+                "max_running_requests was reduced from the requested %d to %d "
+                "(per dp worker) due to the available KV cache capacity.",
+                requested_per_worker,
+                max_num_reqs,
+            )
+        return max_num_reqs
+
+    def _resolve_memory_pool_config(
+        self, pre_model_load_memory: int
+    ) -> MemoryPoolConfig:
+        """Profile GPU memory and resolve all pool parameters into a config."""
+        from sglang.srt.model_executor.pool_configurator import (
+            create_memory_pool_configurator,
+        )
+
+        available_bytes = self._profile_available_bytes(pre_model_load_memory)
+        config = self.config_from_budget(available_bytes)
+        config.max_running_requests = self.resolve_max_num_reqs(
+            config.max_total_num_tokens
+        )
+        configurator = create_memory_pool_configurator(self)
+        config = configurator.finalize_with_max_running_requests(config)
+        config.mem_fraction_static = self.server_args.mem_fraction_static
+        return config
+
+    def config_from_budget(
+        self, budget_bytes: int, *, cap_tokens: Optional[int] = None
+    ) -> MemoryPoolConfig:
+        """Turn a KV byte budget into a pool config via the configurator, re-applying
+        the external token constraints (user cap, page alignment, PP sync) and the
+        optional ``cap_tokens`` clamp."""
+        # Local import avoids a pool_configurator import cycle.
+        from sglang.srt.model_executor.pool_configurator import (
+            create_memory_pool_configurator,
+        )
+
+        configurator = create_memory_pool_configurator(self)
+        config = configurator.calculate_pool_sizes(
+            budget_bytes, self.server_args.page_size
+        )
+        max_tokens = self._apply_token_constraints(config.max_total_num_tokens)
+        if cap_tokens is not None:
+            max_tokens = min(max_tokens, cap_tokens)
+        if max_tokens != config.max_total_num_tokens:
+            config = configurator.calculate_pool_sizes_from_max_tokens(
+                max_tokens, self.server_args.page_size
+            )
+        return config
+
+    def _handle_max_mamba_cache(self, total_rest_memory):
+        config = self.mambaish_config
+        server_args = self.server_args
+        assert config is not None
+
+        has_spec_dec = not self.spec_algorithm.is_none()
+        if has_spec_dec:
+            assert server_args.speculative_num_draft_tokens is not None
+            assert server_args.max_running_requests is not None
+
+        if server_args.max_mamba_cache_size is not None:
+            # Use explicitly set max_mamba_cache_size
+            server_args.override(
+                "mamba_pool.per_dp_shard",
+                max_mamba_cache_size=server_args.max_mamba_cache_size
+                // self.ps.attn_dp_size,
+            )
+            # Reserve intermediate memory based on capped max_num_reqs
+            if has_spec_dec:
+                ratio = self._calculate_mamba_ratio()
+                capped_reqs = min(
+                    server_args.max_running_requests // self.ps.attn_dp_size,
+                    server_args.max_mamba_cache_size // ratio,
+                )
+                intermediate_size = (
+                    config.mamba2_cache_params.mamba_cache_per_req
+                    * capped_reqs
+                    * server_args.speculative_num_draft_tokens
+                )
+                total_rest_memory = total_rest_memory - (intermediate_size / (1 << 30))
+        elif (
+            server_args.disable_radix_cache
+            and server_args.max_running_requests is not None
+        ):
+            # Use explicitly set max_running_requests when radix cache is disabled
+            server_args.override(
+                "mamba_pool.from_max_running_requests",
+                max_mamba_cache_size=server_args.max_running_requests
+                // self.ps.attn_dp_size,
+            )
+            # Reserve intermediate memory based on capped max_num_reqs
+            if has_spec_dec:
+                intermediate_size = (
+                    config.mamba2_cache_params.mamba_cache_per_req
+                    * server_args.max_mamba_cache_size
+                    * server_args.speculative_num_draft_tokens
+                )
+                total_rest_memory = total_rest_memory - (intermediate_size / (1 << 30))
+        else:
+            # Use ratio-based calculation to auto-fit available memory
+            assert config.mamba2_cache_params.mamba_cache_per_req > 0
+            per_req = config.mamba2_cache_params.mamba_cache_per_req
+
+            # Solve jointly for max_mamba_cache_size accounting for intermediate memory.
+            # The mamba budget (from the ratio split) must cover both:
+            #   1. main mamba state: max_mamba_cache_size * per_req
+            #   2. intermediate states: (max_mamba_cache_size / ratio) * D * per_req
+            # So: max_mamba_cache_size * per_req * (1 + D/ratio) = mamba_budget_bytes
+            mamba_budget = (
+                total_rest_memory
+                * server_args.mamba_full_memory_ratio
+                / (1 + server_args.mamba_full_memory_ratio)
+            )
+            mamba_budget_bytes = mamba_budget * (1 << 30)
+
+            if has_spec_dec:
+                ratio = self._calculate_mamba_ratio()
+                D = server_args.speculative_num_draft_tokens
+                # Joint solve: main_state + intermediate = mamba_budget
+                server_args.override(
+                    "mamba_pool.memory_budget_spec",
+                    max_mamba_cache_size=int(
+                        mamba_budget_bytes // (per_req * (1 + D / ratio))
+                    ),
+                )
+                # Intermediate memory is included in mamba_budget, subtract it
+                # so the return value only has main_state subtracted from total
+                capped_reqs = min(
+                    server_args.max_running_requests // self.ps.attn_dp_size,
+                    server_args.max_mamba_cache_size // ratio,
+                )
+                intermediate_size = per_req * capped_reqs * D
+                total_rest_memory = total_rest_memory - (intermediate_size / (1 << 30))
+            else:
+                server_args.override(
+                    "mamba_pool.memory_budget",
+                    max_mamba_cache_size=int(mamba_budget_bytes // per_req),
+                )
+
+        # Validate: max_mamba_cache_size must be positive after memory allocation.
+        # A non-positive value means GPU memory is insufficient for the requested
+        # configuration. Fail fast with actionable advice instead of silently
+        # producing garbled output at runtime.
+        if server_args.max_mamba_cache_size <= 0:
+            raise RuntimeError(
+                f"Not enough GPU memory for hybrid (mamba/linear-attention) state cache. "
+                f"Computed max_mamba_cache_size={server_args.max_mamba_cache_size} "
+                f"(total_rest_memory={total_rest_memory:.2f} GB, "
+                f"mamba_cache_per_req={config.mamba2_cache_params.mamba_cache_per_req / (1 << 20):.2f} MB). "
+                f"Try: (1) reduce --max-running-requests, "
+                f"(2) increase --mem-fraction-static, "
+                f"(3) reduce --speculative-num-draft-tokens, or "
+                f"(4) use GPUs with more memory."
+            )
+
+        mamba_state_memory = (
+            server_args.max_mamba_cache_size
+            * config.mamba2_cache_params.mamba_cache_per_req
+            / (1 << 30)
+        )
+        return total_rest_memory - mamba_state_memory
+
+
+def calculate_mla_kv_cache_dim(
+    *,
+    model_config: ModelConfig,
+    kv_cache_dtype: torch.dtype,
+    server_args: ServerArgs,
+) -> int:
+    is_dsa_model = is_deepseek_dsa(model_config.hf_config)
+    kv_cache_dtype = kv_cache_dtype
+    kv_lora_rank = model_config.kv_lora_rank
+    qk_rope_head_dim = model_config.qk_rope_head_dim
+    kv_cache_dim = kv_lora_rank + qk_rope_head_dim  # default mla kv cache dim
+
+    # For non-DSA models, MLA kv cache dim is simply kv_lora_rank + qk_rope_head_dim
+    if not is_dsa_model:
+        return kv_cache_dim
+
+    # TRTLLM backend does not override kv_cache_dim for MLA kv cache
+    # Assuming dsa prefill and decode backends are the same when using trtllm MLA backend,
+    # since it is not compatible for trtllm and other mla attn backend due to the different
+    # kv cache layout.
+    if (
+        server_args.dsa_prefill_backend == "trtllm"
+        or server_args.dsa_decode_backend == "trtllm"
+    ):
+        return kv_cache_dim
+
+    # On HIP, TileLang and AITER DSA kernels consume the raw MLA KV layout:
+    # nope(512 fp8) + rope(64 fp8), without extra per-block scales.
+    if _is_hip and (
+        server_args.dsa_prefill_backend in ("tilelang", "aiter")
+        or server_args.dsa_decode_backend in ("tilelang", "aiter")
+    ):
+        return kv_cache_dim
+
+    quant_block_size = DSATokenToKVPool.quant_block_size
+    rope_storage_dtype = DSATokenToKVPool.rope_storage_dtype
+    # Calculate override_kv_cache_dim for FP8 storage in backends that use scaled KV layout
+    # (excluding TRTLLM and HIP raw-layout kernels).
+    # kv_lora_rank + scale storage (kv_lora_rank // quant_block_size * 4 bytes) + rope dimension storage
+    # Note: rope dimension is stored in original dtype (bf16), not quantized to fp8
+    if kv_cache_dtype == torch.float8_e4m3fn:
+        assert (
+            kv_lora_rank % quant_block_size == 0
+        ), f"kv_lora_rank {kv_lora_rank} must be multiple of quant_block_size {quant_block_size}"
+
+        return (
+            kv_lora_rank
+            + kv_lora_rank // quant_block_size * 4
+            + qk_rope_head_dim * rope_storage_dtype.itemsize
+        )
+
+    return kv_cache_dim

sglang_patch/srt/models/agnes_dspark.py

@@ -0,0 +1,892 @@
+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.agnes import AgnesConfig
+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.agnes import (
+    DEEPSEEK_V4_STACKED_PARAMS_MAPPING,
+    AgnesDecoderLayer,
+    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: AgnesConfig,
+        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: AgnesConfig, 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(AgnesDecoderLayer):
+
+    def __init__(
+        self,
+        config: AgnesConfig,
+        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: AgnesConfig,
+        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 AgnesForCausalLMDSpark(nn.Module):
+
+    def __init__(
+        self,
+        config: AgnesConfig,
+        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(
+                "AgnesForCausalLMDSpark 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(
+                "AgnesForCausalLMDSpark 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 = [AgnesForCausalLMDSpark]

sglang_patch/srt/models/agnes_nextn.py

@@ -0,0 +1,283 @@
+import logging
+from typing import Iterable, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import PretrainedConfig
+
+from sglang.srt.distributed import get_pp_group
+from sglang.srt.layers.attention.dsa.utils import (
+    can_dsa_cp_split,
+    dsa_use_prefill_cp,
+    is_dsa_enable_prefill_cp,
+    is_dsa_prefill_cp_round_robin_split,
+)
+from sglang.srt.layers.dp_attention import (
+    dp_gather_partial,
+    get_global_dp_buffer_len,
+    is_dp_attention_enabled,
+)
+from sglang.srt.layers.layernorm import RMSNorm
+from sglang.srt.layers.linear import ReplicatedLinear
+from sglang.srt.layers.logits_processor import LogitsProcessor
+from sglang.srt.layers.moe.utils import get_moe_a2a_backend
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
+from sglang.srt.layers.utils.cp_utils import (
+    cp_all_gather_rerange_output,
+    cp_round_robin_input_ids,
+    cp_split_and_rebuild_data,
+    cp_split_and_rebuild_position,
+    prepare_context_parallel_metadata,
+)
+from sglang.srt.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.forward_context import get_attn_backend
+from sglang.srt.models.agnes import AgnesDecoderLayer, AgnesForCausalLM
+from sglang.srt.runtime_context import get_parallel, get_server_args
+from sglang.srt.utils import add_prefix
+
+logger = logging.getLogger(__name__)
+
+COMPRESS_RATIO_NEXTN_LAYER = 0
+
+
+class AgnesModelNextN(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            enable_tp=not is_dp_attention_enabled(),
+            prefix=add_prefix("embed_tokens", prefix),
+        )
+
+        self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rms_norm_eps = config.rms_norm_eps
+
+        self.hc_eps = config.hc_eps
+        self.hc_mult = hc_mult = config.hc_mult
+        hc_dim = hc_mult * config.hidden_size
+        self.hc_head_fn = nn.Parameter(
+            torch.empty(hc_mult, hc_dim, dtype=torch.float32)
+        )
+        self.hc_head_base = nn.Parameter(torch.empty(hc_mult, dtype=torch.float32))
+        self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32))
+
+        self.e_proj = ReplicatedLinear(
+            config.hidden_size,
+            config.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("e_proj", prefix),
+        )
+        self.h_proj = ReplicatedLinear(
+            config.hidden_size,
+            config.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("h_proj", prefix),
+        )
+        if isinstance(quant_config, ModelSlimConfig):
+            prefix = "mtp.0"
+        else:
+            prefix = add_prefix("decoder", prefix)
+
+        self.decoder = AgnesDecoderLayer(
+            config,
+            layer_id=0,
+            quant_config=quant_config,
+            is_nextn=True,
+            prefix=prefix,
+            alt_streams=None,
+            compress_ratio_override=COMPRESS_RATIO_NEXTN_LAYER,
+        )
+
+        self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
+        if self.dsa_enable_prefill_cp:
+            self.cp_size = get_parallel().attn_cp_size
+        else:
+            self.cp_size = None
+
+        self.shared_head = nn.Module()
+        self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def hc_head(
+        self,
+        x: torch.Tensor,
+        hc_fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+    ):
+        shape, dtype = x.size(), x.dtype
+        x = x.flatten(1).float()
+        rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.rms_norm_eps)
+        mixes = F.linear(x, hc_fn) * rsqrt
+        pre = torch.sigmoid(mixes * hc_scale + hc_base) + self.hc_eps
+        y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=1)
+        return y.to(dtype)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        input_embeds: torch.Tensor = None,
+    ) -> torch.Tensor:
+        if input_embeds is None:
+            hidden_states = self.embed_tokens(input_ids)
+        else:
+            hidden_states = input_embeds
+
+        if hidden_states.shape[0] > 0:
+            n_tokens = hidden_states.shape[0]
+            d = self.config.hidden_size
+            hc_flat = forward_batch.spec_info.hidden_states.view(
+                n_tokens * self.hc_mult, d
+            )
+            h_proj_out, _ = self.h_proj(self.hnorm(hc_flat))
+            h_proj_hidden_states = h_proj_out.view(n_tokens, self.hc_mult, d)
+
+            e_proj_hidden_states, _ = self.e_proj(self.enorm(hidden_states))
+            hidden_states = e_proj_hidden_states[:, None, :] + h_proj_hidden_states
+        else:
+            hidden_states = hidden_states.unsqueeze(1).repeat(1, self.hc_mult, 1)
+
+        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,
+            )
+            dp_gather_partial(input_ids_global, input_ids[:, None], forward_batch)
+            input_ids_global = input_ids_global.squeeze(-1)
+        else:
+            input_ids_global = input_ids
+
+        if dsa_use_prefill_cp(forward_batch):
+            hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
+            positions = cp_split_and_rebuild_position(forward_batch, positions)
+            input_ids = cp_round_robin_input_ids(input_ids)
+            input_ids_global = input_ids
+
+        hidden_states, residual, post, comb = self.decoder(
+            positions=positions,
+            hidden_states=hidden_states,
+            forward_batch=forward_batch,
+            input_ids=input_ids,
+            input_ids_global=input_ids_global,
+        )
+        if residual is not None:
+            # NextN has a single decoder layer, so no later layer can consume a
+            # deferred fused hc_post state.
+            hidden_states = self.decoder.hc_post(hidden_states, residual, post, comb)
+
+        if dsa_use_prefill_cp(forward_batch):
+            hidden_states = cp_all_gather_rerange_output(
+                hidden_states,
+                self.cp_size,
+                forward_batch,
+                torch.cuda.current_stream(),
+            )
+
+        pre_hc_head = hidden_states.flatten(1)
+
+        hidden_states = self.hc_head(
+            hidden_states, self.hc_head_fn, self.hc_head_scale, self.hc_head_base
+        )
+        hidden_states = self.shared_head.norm(hidden_states)
+
+        return hidden_states, pre_hc_head
+
+
+class AgnesForCausalLMNextN(AgnesForCausalLM):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        nn.Module.__init__(self)
+        self.config = config
+        self.tp_size = get_parallel().tp_size
+        self.pp_group = get_pp_group()
+        self.quant_config = quant_config
+        self.determine_num_fused_shared_experts()
+        self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
+        if self.dsa_enable_prefill_cp:
+            self.cp_rank = get_parallel().attn_cp_rank
+            self.cp_size = get_parallel().attn_cp_size
+        else:
+            self.cp_rank = None
+            self.cp_size = None
+
+        self.model = AgnesModelNextN(
+            config, quant_config, prefix=add_prefix("model", prefix)
+        )
+        self.lm_head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=add_prefix("model.shared_head.head", prefix),
+            use_attn_tp_group=get_server_args().enable_dp_lm_head,
+        )
+        self.logits_processor = LogitsProcessor(config)
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ) -> torch.Tensor:
+        if self.dsa_enable_prefill_cp:
+            if can_dsa_cp_split(len(input_ids), self.cp_size, True, forward_batch):
+                forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
+                    len(input_ids),
+                    self.cp_rank,
+                    self.cp_size,
+                    forward_batch.seq_lens_cpu.tolist(),
+                    extend_seqs_len=forward_batch.extend_seq_lens_cpu,
+                )
+                if is_dsa_prefill_cp_round_robin_split():
+                    attn_backend = get_attn_backend()
+                    metadata = attn_backend.forward_metadata
+                    core_meta = metadata.core_attn_metadata
+                    core_meta.apply_cp_reindex()
+                    core_meta.init_flashmla_related(is_prefill=True)
+                    if metadata.indexer_metadata is not None:
+                        metadata.indexer_metadata = (
+                            attn_backend.init_forward_metadata_indexer(core_meta)
+                        )
+
+        hidden_states, pre_hc_head = self.model(input_ids, positions, forward_batch)
+        return self.logits_processor(
+            input_ids,
+            hidden_states,
+            self.lm_head,
+            forward_batch,
+            hidden_states_before_norm=pre_hc_head,
+        )
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        super().load_weights(weights, is_nextn=True)
+
+    def post_load_weights(self, is_nextn=False, weight_names=None):
+        super().post_load_weights(is_nextn=True, weight_names=weight_names)
+
+
+EntryClass = [AgnesForCausalLMNextN]

sglang_patch/srt/parser/reasoning_parser.py

@@ -0,0 +1,1495 @@
+import inspect
+import re
+from typing import Dict, List, Optional, Tuple, Type
+
+from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
+from sglang.srt.function_call.hunyuan_detector import resolve_hunyuan_tokens
+from sglang.srt.parser.harmony_parser import HarmonyParser
+from sglang.srt.parser.inkling_tokenizer import (
+    CONTENT_INVOKE_TOOL_JSON,
+    CONTENT_MODEL_END_SAMPLING,
+    CONTENT_TEXT,
+    CONTENT_THINKING,
+    END_MESSAGE,
+    INKLING_CONTROL_TOKENS,
+    MESSAGE_MODEL,
+)
+
+
+class StreamingParseResult:
+    """Result of streaming incremental parsing."""
+
+    def __init__(
+        self,
+        normal_text: Optional[str] = None,
+        reasoning_text: Optional[str] = None,
+    ):
+        self.normal_text = normal_text or ""
+        self.reasoning_text = reasoning_text or ""
+
+
+class BaseReasoningFormatDetector:
+    """Base class providing two sets of interfaces: one-time and streaming incremental."""
+
+    def __init__(
+        self,
+        think_start_token: str,
+        think_end_token: str,
+        think_excluded_tokens: Optional[List[str]] = None,
+        force_reasoning: bool = False,
+        stream_reasoning: bool = True,
+        tool_start_token: Optional[str] = None,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        thinks_internally: bool = False,
+        reasoning_default: str = "always",
+        force_nonempty_content: bool = False,
+    ):
+        self.think_start_token = think_start_token
+        self.think_end_token = think_end_token
+        self.think_excluded_tokens = think_excluded_tokens
+        self.tool_start_token = tool_start_token
+        self.force_reasoning = force_reasoning
+        self._in_reasoning = force_reasoning
+        self.stream_reasoning = stream_reasoning
+        self.thinks_internally = thinks_internally
+        self.reasoning_default = reasoning_default
+
+        self._buffer = ""
+        self.stripped_think_start = False
+        self.think_start_self_label = ""
+
+        self._force_nonempty_content = force_nonempty_content
+        self._accumulated_reasoning = ""
+
+        self.continue_final_message = continue_final_message
+        if self.continue_final_message:
+            self.previous_content = previous_content
+            self.previous_count = len(previous_content)
+        else:
+            self.previous_content = ""
+            self.previous_count = 0
+
+        if self.think_start_token in self.previous_content:
+            self._in_reasoning = True
+        if self.think_end_token in self.previous_content:
+            self._in_reasoning = False
+
+    def _maybe_apply_force_nonempty_content(
+        self, ret: StreamingParseResult
+    ) -> StreamingParseResult:
+        if self._force_nonempty_content and not ret.normal_text:
+            ret.normal_text, ret.reasoning_text = ret.reasoning_text, ret.normal_text
+        return ret
+
+    def detect_and_parse(self, text: str) -> StreamingParseResult:
+        """
+        One-time parsing: Detects and parses reasoning sections in the provided text.
+        Returns both reasoning content and normal text separately.
+        """
+        return self._maybe_apply_force_nonempty_content(
+            self._detect_and_parse_impl(text)
+        )
+
+    def _detect_and_parse_impl(self, text: str) -> StreamingParseResult:
+        in_reasoning = self._in_reasoning or self.think_start_token in text
+
+        if not in_reasoning:
+            return StreamingParseResult(normal_text=text)
+
+        # The text is considered to be in a reasoning block.
+        think_start_text = self.think_start_token + self.think_start_self_label
+        processed_text = text
+        while processed_text.startswith(think_start_text):
+            processed_text = processed_text[len(think_start_text) :]
+
+        if (
+            self.think_end_token not in processed_text
+            and self.think_end_token not in self.previous_content
+        ):
+            # Check for tool_start_token interruption
+            if (
+                in_reasoning
+                and self.tool_start_token is not None
+                and self.tool_start_token in processed_text
+            ):
+                # Find the first occurrence of tool_start_token and split there
+                tool_idx = processed_text.find(self.tool_start_token)
+                reasoning_text = processed_text[:tool_idx]
+                # Preserve tool_start_token in normal text
+                normal_text = processed_text[tool_idx:]
+                return StreamingParseResult(
+                    normal_text=normal_text, reasoning_text=reasoning_text
+                )
+            # Assume reasoning was truncated before end token
+            return StreamingParseResult(reasoning_text=processed_text)
+
+        # Extract reasoning content
+        if self.think_end_token in processed_text:
+            splits = processed_text.split(self.think_end_token, maxsplit=1)
+            reasoning_text = splits[0]
+            normal_text = splits[1]
+
+            return StreamingParseResult(
+                normal_text=normal_text, reasoning_text=reasoning_text
+            )
+        else:
+            # think_end_token is in self.previous_content for continue_final_message=True case
+            return StreamingParseResult(normal_text=processed_text)
+
+    def parse_streaming_increment(self, new_text: str) -> StreamingParseResult:
+        """
+        Streaming incremental parsing for reasoning content.
+        Handles partial reasoning tags and content.
+
+        If stream_reasoning is False:
+            Accumulates reasoning content until the end tag is found
+        If stream_reasoning is True:
+            Streams reasoning content as it arrives
+        """
+        ret = self._parse_streaming_increment_impl(new_text)
+        if self._force_nonempty_content:
+            if self._in_reasoning:
+                self._accumulated_reasoning += ret.reasoning_text
+            else:
+                self._accumulated_reasoning = ""
+        return ret
+
+    def _parse_streaming_increment_impl(self, new_text: str) -> StreamingParseResult:
+        self._buffer += new_text
+        current_text = self._buffer
+
+        think_start_text = self.think_start_token + self.think_start_self_label
+
+        # If the current text is a prefix of the think token, keep buffering
+        tokens_to_check = [think_start_text, self.think_end_token]
+        if self.tool_start_token:
+            tokens_to_check.append(self.tool_start_token)
+        if any(
+            token.startswith(current_text) and token != current_text
+            for token in tokens_to_check
+        ):
+            return StreamingParseResult()
+
+        # Strip `<think>` token if present
+        if not self.stripped_think_start and think_start_text in current_text:
+            current_text = current_text.replace(think_start_text, "", 1)
+            self.stripped_think_start = True
+            self._in_reasoning = True
+
+        # Handle end of reasoning block
+        if self._in_reasoning and self.think_end_token in current_text:
+            end_idx = current_text.find(self.think_end_token)
+
+            reasoning_text = current_text[:end_idx]
+
+            self._buffer = ""
+            self._in_reasoning = False
+            normal_text = current_text[end_idx + len(self.think_end_token) :]
+
+            return StreamingParseResult(
+                normal_text=normal_text, reasoning_text=reasoning_text
+            )
+
+        # Continue with reasoning content
+        if self._in_reasoning:
+            # Check for tool_start_token interruption
+            if self.tool_start_token and self.tool_start_token in current_text:
+                tool_idx = current_text.find(self.tool_start_token)
+                reasoning_text = current_text[:tool_idx]
+                # Preserve tool_start_token in normal text
+                normal_text = current_text[tool_idx:]
+                self._buffer = ""
+                self._in_reasoning = False
+                return StreamingParseResult(
+                    normal_text=normal_text, reasoning_text=reasoning_text
+                )
+            if self.stream_reasoning:
+                # Stream the content immediately
+                self._buffer = ""
+                return StreamingParseResult(reasoning_text=current_text)
+            else:
+                return StreamingParseResult()
+
+        # If we're not in a reasoning block return as normal text
+        if not self._in_reasoning:
+            self._buffer = ""
+            return StreamingParseResult(normal_text=current_text)
+
+        return StreamingParseResult()
+
+    def finish(self) -> StreamingParseResult:
+        """
+        Called once when the stream ends. If force_nonempty_content is set
+        and the stream ended mid-reasoning, reclassifies the accumulated
+        reasoning (plus any partial token still buffered) as normal text.
+        """
+        if self._force_nonempty_content and self._in_reasoning:
+            # stream_reasoning=False never clears _buffer, so the opening think
+            # token (stripped only from the base class's local view) survives here.
+            buffer = self._buffer
+            think_start_text = self.think_start_token + self.think_start_self_label
+            if buffer.startswith(think_start_text):
+                buffer = buffer[len(think_start_text) :]
+            normal_text = self._accumulated_reasoning + buffer
+            self._accumulated_reasoning = ""
+            self._buffer = ""
+            if normal_text:
+                return StreamingParseResult(normal_text=normal_text)
+        return StreamingParseResult()
+
+
+class DeepSeekR1Detector(BaseReasoningFormatDetector):
+    """
+    Detector for DeepSeek-R1 model.
+    Assumes reasoning format:
+      (<think>)*(.*)</think>
+    Returns all the text before the </think> tag as `reasoning_text`
+    and the rest of the text as `normal_text`.
+
+    Supported models:
+      - DeepSeek-R1: Always generates thinking content without <think> start tag
+      - DeepSeek-R1-0528: Generates thinking content with <think> start tag
+
+    Format patterns:
+      - DeepSeek-R1: "I need to think about this...</think>The answer is 42."
+      - DeepSeek-R1-0528: "<think>I need to think about this...</think>The answer is 42."
+
+    Args:
+        stream_reasoning (bool): If False, accumulates reasoning content until the end tag.
+            If True, streams reasoning content as it arrives.
+    """
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = True,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        # DeepSeek-R1 is assumed to be reasoning until `</think>` token
+        super().__init__(
+            "<think>",
+            "</think>",
+            force_reasoning=True,
+            stream_reasoning=stream_reasoning,
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            force_nonempty_content=force_nonempty_content,
+        )
+        # https://github.com/sgl-project/sglang/pull/3202#discussion_r1950153599
+
+
+class Qwen3Detector(BaseReasoningFormatDetector):
+    """
+    Detector for Qwen3 models (e.g., Qwen/Qwen3-235B-A22B).
+    Assumes reasoning format:
+      (<think>)*(.*)</think>
+
+    Qwen3 models released before 07/2025 supports switching between thinking mode and normal
+    mode using `enable_thinking` parameter in the request parameter.
+      - enable_thinking=True: "<think>reasoning content</think>The answer is 42."
+      - enable_thinking=False: "The answer is 42." (no thinking tokens)
+
+    Args:
+        stream_reasoning (bool): If False, accumulates reasoning content until the end tag.
+            If True, streams reasoning content as it arrives.
+    """
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = False,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        think_excluded_tokens = [
+            "<tool_call>",
+            "</tool_call>",
+            "<|im_end|>",
+            "<|endoftext|>",
+        ]
+        super().__init__(
+            "<think>",
+            "</think>",
+            think_excluded_tokens=think_excluded_tokens,
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+            # Qwen3.5 sometimes opens ``<tool_call>`` without closing
+            # ``</think>``; treat it as an implicit reasoning close.
+            tool_start_token="<tool_call>",
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            thinks_internally=True,
+            reasoning_default="enable_thinking",
+            force_nonempty_content=force_nonempty_content,
+        )
+
+
+class KimiDetector(BaseReasoningFormatDetector):
+    """
+    Detector for Kimi Thinking model.
+    Assumes reasoning format:
+      ◁think▷*(.*)◁/think▷
+    Returns all the text before the ◁/think▷ tag as `reasoning_text`
+    and the rest of the text as `normal_text`.
+    """
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = False,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        super().__init__(
+            "◁think▷",
+            "◁/think▷",
+            force_reasoning=False,
+            stream_reasoning=stream_reasoning,
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            force_nonempty_content=force_nonempty_content,
+        )
+
+
+class KimiK2Detector(BaseReasoningFormatDetector):
+    """
+    Detector for Kimi K2 models.
+    Assumes reasoning format:
+      (<think>)*(.*)</think>
+
+    Kimi K2 can switch from reasoning to tool-call section with
+    `<|tool_calls_section_begin|>` before emitting `</think>`.
+    """
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = False,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        think_excluded_tokens = [
+            "<think>",
+            "<|tool_calls_section_begin|>",
+            "<|tool_call_begin|>",
+            "<|tool_call_argument_begin|>",
+            "<|tool_call_section_end|>",
+            "<|tool_call_end|>",
+            "[EOS]",
+            "<|im_end|>",
+            "<|end_header_id|>",
+            "[EOT]",
+        ]
+        super().__init__(
+            "<think>",
+            "</think>",
+            think_excluded_tokens=think_excluded_tokens,
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+            tool_start_token="<|tool_calls_section_begin|>",
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            reasoning_default="thinking",
+            force_nonempty_content=force_nonempty_content,
+        )
+
+
+class Glm45Detector(BaseReasoningFormatDetector):
+    """
+    Detector for GLM-4.5 models.
+    Assumes reasoning format:
+      (<think>)*(.*)</think>
+
+    GLM-4.5 uses `<tool_call>` as the tool start token to switch from reasoning mode to normal mode.
+
+    Args:
+        stream_reasoning (bool): If False, accumulates reasoning content until the end tag.
+            If True, streams reasoning content as it arrives.
+    """
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = False,
+        force_nonempty_content: bool = False,
+    ):
+        think_excluded_tokens = [
+            "<tool_call>",
+            "</tool_call>",
+            "<eop>",
+            "<|user|>",
+            "<|endoftext|>",
+        ]
+        super().__init__(
+            "<think>",
+            "</think>",
+            think_excluded_tokens=think_excluded_tokens,
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+            tool_start_token="<tool_call>",
+            thinks_internally=True,
+            reasoning_default="enable_thinking",
+            force_nonempty_content=force_nonempty_content,
+        )
+
+
+class GptOssDetector(BaseReasoningFormatDetector):
+    """
+    Detector for T4-style reasoning format (GPT-OSS), using the HarmonyParser.
+    """
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = True,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        super().__init__(
+            "<|channel|>analysis<|message|>",
+            "<|end|>",
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            force_nonempty_content=force_nonempty_content,
+        )
+        self.parser = HarmonyParser()
+
+    def detect_and_parse(self, text: str) -> StreamingParseResult:
+        events = self.parser.parse(text)
+        # Flush the buffer for one-shot parsing
+        events += self.parser.parse("")
+
+        reasoning_text = "".join(
+            [e.content for e in events if e.event_type == "reasoning"]
+        )
+        normal_parts = []
+        for e in events:
+            if e.event_type == "normal":
+                normal_parts.append(e.content)
+            elif e.event_type == "tool_call":
+                # Use raw_text to preserve structural markers for function call detector
+                normal_parts.append(e.raw_text if e.raw_text else e.content)
+        normal_text = "".join(normal_parts)
+        # Tool call events preserve raw text with structural markers
+
+        return self._maybe_apply_force_nonempty_content(
+            StreamingParseResult(
+                normal_text=normal_text,
+                reasoning_text=reasoning_text,
+            )
+        )
+
+    def parse_streaming_increment(self, new_text: str) -> StreamingParseResult:
+        events = self.parser.parse(new_text)
+
+        reasoning_text = "".join(
+            [e.content for e in events if e.event_type == "reasoning"]
+        )
+        normal_parts = []
+        for e in events:
+            if e.event_type == "normal":
+                normal_parts.append(e.content)
+            elif e.event_type == "tool_call":
+                # Use raw_text to preserve structural markers for function call detector
+                normal_parts.append(e.raw_text if e.raw_text else e.content)
+        normal_text = "".join(normal_parts)
+
+        return StreamingParseResult(
+            normal_text=normal_text,
+            reasoning_text=reasoning_text,
+        )
+
+
+class MiniMaxAppendThinkDetector(BaseReasoningFormatDetector):
+    """
+    Append `<think>` token to the beginning of the text.
+    """
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = False,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        # scheduler.py need `reasoning_parser.detector.think_end_token`
+        super().__init__(
+            "<think>",
+            "</think>",
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            force_nonempty_content=force_nonempty_content,
+        )
+        self.is_first_chunk = False
+
+    def parse_streaming_increment(self, new_text: str) -> StreamingParseResult:
+        if not self.is_first_chunk:
+            self.is_first_chunk = True
+            new_text = self.think_start_token + new_text
+        return StreamingParseResult(normal_text=new_text)
+
+    def detect_and_parse(self, text: str) -> StreamingParseResult:
+        return StreamingParseResult(normal_text=self.think_start_token + text)
+
+
+class Nemotron3Detector(BaseReasoningFormatDetector):
+    """
+    Detector for Nemotron3 model.
+    Uses the same reasoning format as DeepSeek-R1: (<think>)*(.*)</think>
+
+    """
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = False,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        super().__init__(
+            "<think>",
+            "</think>",
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+            tool_start_token="<tool_call>",
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            reasoning_default="enable_thinking",
+            force_nonempty_content=force_nonempty_content,
+        )
+
+
+class MiniMaxM3Detector(BaseReasoningFormatDetector):
+    """MiniMax-M3 detector. Format: (<mm:think>)*(.*)</mm:think>.
+
+    In multi-turn chats M3 prefixes earlier non-thinking turns with a bare
+    ``</mm:think>``, so a non-thinking reply may open with one stray closer; drop it unless thinking.
+    """
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = False,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        super().__init__(
+            "<mm:think>",
+            "</mm:think>",
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+        )
+        self._lead_buffer = ""
+        self._checked_leading_close = False
+        self._force_nonempty_content = force_nonempty_content
+
+    def detect_and_parse(self, text: str) -> StreamingParseResult:
+        if not self._in_reasoning and text.lstrip().startswith(self.think_end_token):
+            text = text.lstrip()[len(self.think_end_token) :]
+        ret = super().detect_and_parse(text)
+        if self._force_nonempty_content and not ret.normal_text:
+            ret.normal_text, ret.reasoning_text = ret.reasoning_text, ret.normal_text
+        return ret
+
+    def parse_streaming_increment(self, new_text: str) -> StreamingParseResult:
+        # ``</mm:think>`` is a single token, so a stray leading closer arrives whole.
+        if not self._checked_leading_close and not self._in_reasoning:
+            self._lead_buffer += new_text
+            stripped = self._lead_buffer.lstrip()
+            if not stripped:
+                return StreamingParseResult()
+            self._checked_leading_close = True
+            if stripped.startswith(self.think_end_token):
+                new_text = stripped[len(self.think_end_token) :]
+            else:
+                new_text = self._lead_buffer
+            self._lead_buffer = ""
+            if not new_text:
+                return StreamingParseResult()
+        return super().parse_streaming_increment(new_text)
+
+
+class MistralDetector(BaseReasoningFormatDetector):
+    """
+    Detector for Mistral models with reasoning (e.g., Mistral-Small-4-119B-2603).
+    Assumes reasoning format:
+      [THINK]reasoning content[/THINK]answer
+
+    Reasoning is optional — it only appears when reasoning_effort="high" is set.
+    When reasoning_effort="none", the model outputs directly without thinking tokens.
+    """
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = False,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        super().__init__(
+            "[THINK]",
+            "[/THINK]",
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            reasoning_default="mistral",
+            force_nonempty_content=force_nonempty_content,
+        )
+
+
+class HunyuanDetector(BaseReasoningFormatDetector):
+    """
+    Detector for Hunyuan models (e.g., tencent/Hunyuan-A13B-Instruct).
+
+    Like Glm45Detector but uses ``<tool_calls>`` (plural) as the tool start token.
+    """
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = False,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        tokenizer=None,
+        force_nonempty_content: bool = False,
+    ):
+        t = resolve_hunyuan_tokens(tokenizer)
+        think_open = t["think"]
+        think_close = (
+            "</" + think_open[1:] if think_open.startswith("<") else think_open
+        )
+        super().__init__(
+            think_open,
+            think_close,
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+            tool_start_token=t["tool_calls"],
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            force_nonempty_content=force_nonempty_content,
+        )
+
+
+class Gemma4Detector(BaseReasoningFormatDetector):
+    """Gemma4 reasoning detector."""
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = False,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        super().__init__(
+            "<|channel>",
+            "<channel|>",
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            reasoning_default="explicit_enable_thinking",
+            force_nonempty_content=force_nonempty_content,
+        )
+        self.think_start_self_label = "thought\n"
+
+
+_INKLING_CONTENT_KINDS = {
+    CONTENT_THINKING: "reasoning",
+    CONTENT_TEXT: "content",
+}
+_INKLING_END_TOKENS = {
+    CONTENT_MODEL_END_SAMPLING,
+    END_MESSAGE,
+}
+_INKLING_CONTROL_TOKENS = INKLING_CONTROL_TOKENS
+_INKLING_CONTROL_RE = re.compile(
+    "|".join(re.escape(t) for t in sorted(_INKLING_CONTROL_TOKENS))
+)
+
+
+class InklingDetector(BaseReasoningFormatDetector):
+    """Detector for Inkling typed content blocks."""
+
+    # Parse the model's sequence of typed content blocks, for example:
+    #   <|message_model|><|content_thinking|>reasoning<|end_message|>
+    #   <|message_model|><|content_text|>visible answer<|end_message|>
+    #   <|content_model_end_sampling|>
+    # Special tokens must decode literally so thinking and visible text can be
+    # routed to their respective response fields.
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = False,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        del force_nonempty_content
+        super().__init__(
+            CONTENT_THINKING,
+            END_MESSAGE,
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            thinks_internally=False,
+            reasoning_default="always",
+        )
+
+        self._kind: str | None = None
+        self._pending_header = ""
+        self._pending_reasoning = ""
+
+    def detect_and_parse(self, text: str) -> StreamingParseResult:
+        self._buffer = ""
+        self._kind = None
+        self._pending_header = ""
+        self._pending_reasoning = ""
+        ret = self._parse_blocks(text)
+        if self._kind == "reasoning" and not self.stream_reasoning:
+            ret.reasoning_text += self._pending_reasoning
+        self._kind = None
+        self._pending_header = ""
+        self._pending_reasoning = ""
+        return ret
+
+    def parse_streaming_increment(self, new_text: str) -> StreamingParseResult:
+        text = self._buffer + new_text
+        partial_len = self._partial_control_length(text)
+        if partial_len:
+            self._buffer = text[-partial_len:]
+            text = text[:-partial_len]
+        else:
+            self._buffer = ""
+        return self._parse_blocks(text)
+
+    def finish(self) -> StreamingParseResult:
+        # Flush reasoning buffered under stream_reasoning=False when the stream
+        # ends before a control/end token closes the block (e.g. max_tokens cut
+        # a thinking block short). Mirrors the non-streaming flush in
+        # detect_and_parse; without it the trailing reasoning trace is dropped.
+        reasoning_text = ""
+        if self._kind == "reasoning" and not self.stream_reasoning:
+            reasoning_text = self._pending_reasoning
+        self._buffer = ""
+        self._pending_reasoning = ""
+        self._pending_header = ""
+        self._kind = None
+        return StreamingParseResult(reasoning_text=reasoning_text)
+
+    @staticmethod
+    def _partial_control_length(text: str) -> int:
+        max_token_len = max(map(len, _INKLING_CONTROL_TOKENS))
+        for length in range(min(len(text), max_token_len - 1), 0, -1):
+            suffix = text[-length:]
+            if any(
+                len(suffix) < len(token) and token.startswith(suffix)
+                for token in _INKLING_CONTROL_TOKENS
+            ):
+                return length
+        return 0
+
+    def _parse_blocks(self, text: str) -> StreamingParseResult:
+        reasoning: list[str] = []
+        content: list[str] = []
+        saw_control = False
+        pos = 0
+
+        def emit(text: str) -> None:
+            if self._kind == "reasoning":
+                if self.stream_reasoning:
+                    reasoning.append(text)
+                else:
+                    self._pending_reasoning += text
+            elif self._kind == "content":
+                content.append(text)
+            elif self._kind == "tool":
+                content.append(text)
+            elif self._kind == "header":
+                self._pending_header += text
+            elif text:
+                # No open block — e.g. a continue_final_message stream resuming
+                # mid text block. Route to visible content, matching the
+                # no-control-token path below.
+                content.append(text)
+
+        def flush_reasoning() -> None:
+            if self._kind == "reasoning" and not self.stream_reasoning:
+                reasoning.append(self._pending_reasoning)
+                self._pending_reasoning = ""
+
+        for match in _INKLING_CONTROL_RE.finditer(text):
+            saw_control = True
+            emit(text[pos : match.start()])
+
+            token = match.group(0)
+            pos = match.end()
+            if token == MESSAGE_MODEL:
+                if self._kind in (None, "header"):
+                    flush_reasoning()
+                    self._pending_header = ""
+                    self._kind = "header"
+                else:
+                    # Inside an open block a decoded <|message_model|> string
+                    # is payload the model wrote (e.g. quoting the protocol) —
+                    # a real header can only follow an end token. Preserve it
+                    # instead of rerouting the rest of the block into a header.
+                    emit(token)
+            elif token == CONTENT_INVOKE_TOOL_JSON:
+                flush_reasoning()
+                if self._kind == "header":
+                    content.extend(
+                        (MESSAGE_MODEL, self._pending_header, CONTENT_INVOKE_TOOL_JSON)
+                    )
+                    self._pending_header = ""
+                else:
+                    content.append(token)
+                self._kind = "tool"
+            elif self._kind == "tool":
+                content.append(token)
+                if token in _INKLING_END_TOKENS:
+                    self._kind = None
+            elif token in _INKLING_CONTENT_KINDS:
+                flush_reasoning()
+                self._pending_header = ""
+                self._kind = _INKLING_CONTENT_KINDS[token]
+            elif token in _INKLING_END_TOKENS:
+                flush_reasoning()
+                self._pending_header = ""
+                self._kind = None
+
+        tail = text[pos:]
+        if saw_control or self._kind is not None:
+            emit(tail)
+        else:
+            content.append(text)
+
+        return StreamingParseResult(
+            normal_text="".join(content),
+            reasoning_text="".join(reasoning),
+        )
+
+
+class _DeepSeekV3Detector(Qwen3Detector):
+    """DeepSeek-V3 reuses Qwen3 tokens but requires explicit thinking=True to enable."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.reasoning_default = "explicit_thinking"
+
+
+class _MimoDetector(Qwen3Detector):
+    """MIMO reuses Qwen3 tokens but requires explicit enable_thinking=True to enable."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.reasoning_default = "explicit_enable_thinking"
+
+
+class _PoolsideV1Detector(Qwen3Detector):
+    """Poolside v1 (Laguna-XS.2) reuses Qwen3 <think> tokens but the HF chat template
+    defaults `enable_thinking=False`; reasoning is opt-in via `enable_thinking=True`."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.reasoning_default = "explicit_enable_thinking"
+
+
+class Apertus2509Detector(BaseReasoningFormatDetector):
+    """
+    Detector for Apertus 2509 models
+
+    Reasoning blocks are delimited by:
+        <|inner_prefix|> ... <|inner_suffix|>
+    """
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = False,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        super().__init__(
+            "<|inner_prefix|>",
+            "<|inner_suffix|>",
+            force_reasoning=False,
+            stream_reasoning=stream_reasoning,
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            force_nonempty_content=force_nonempty_content,
+        )
+        self._force_reasoning = force_reasoning
+        self._tool_start_token = "<|tools_prefix|>["
+        self._tool_end_token = "<|tools_suffix|>"
+        self._reasoning_acc: str = ""
+        self._in_inner_tool: bool = False
+
+    @staticmethod
+    def _ends_with_partial_token(buffer: str, token: str) -> int:
+        for i in range(1, min(len(buffer) + 1, len(token))):
+            if token.startswith(buffer[-i:]):
+                return i
+        return 0
+
+    def detect_and_parse(self, text: str) -> StreamingParseResult:
+        blocks = self.detect_and_parse_block_sequence(text)
+        reasoning_parts = [t for k, t in blocks if k == "reasoning"]
+        text_parts = [t for k, t in blocks if k == "text"]
+        ret = StreamingParseResult(
+            normal_text="".join(text_parts),
+            reasoning_text="".join(reasoning_parts),
+        )
+        return self._maybe_apply_force_nonempty_content(ret)
+
+    def detect_and_parse_block_sequence(self, text: str) -> list[tuple[str, str]]:
+        """Return an ordered sequence of blocks: [("reasoning"|"text", content), ...]"""
+        start_tok = self.think_start_token
+        end_tok = self.think_end_token
+        blocks: list[tuple[str, str]] = []
+        cursor = 0
+
+        # continue_final_message can resume inside an existing inner
+        if self._in_reasoning:
+            if (e := text.find(end_tok, cursor)) == -1:
+                blocks.extend(self._split_inner_reasoning(text[cursor:]))
+                blocks.append(("text", ""))
+                return blocks
+            blocks.extend(self._split_inner_reasoning(text[cursor:e]))
+            cursor = e + len(end_tok)
+
+        while True:
+            if (s := text.find(start_tok, cursor)) == -1:
+                # Always include the trailing text block (may be empty)
+                blocks.append(("text", text[cursor:]))
+                break
+            if s > cursor:
+                blocks.append(("text", text[cursor:s]))
+
+            cursor = s + len(start_tok)
+            if (e := text.find(end_tok, cursor)) == -1:
+                blocks.extend(self._split_inner_reasoning(text[cursor:]))
+                blocks.append(("text", ""))
+                break
+            blocks.extend(self._split_inner_reasoning(text[cursor:e]))
+            cursor = e + len(end_tok)
+
+        last_idx = len(blocks) - 1
+        blocks = [
+            (k, t)
+            for i, (k, t) in enumerate(blocks)
+            if not (k == "text" and t == "" and i != last_idx)
+        ]
+
+        return blocks
+
+    def _split_inner_reasoning(self, inner_text: str) -> list[tuple[str, str]]:
+        """
+        Split content inside <|inner_prefix|>...<|inner_suffix|> into:
+          - ("reasoning", <thoughts text>)
+          - ("text", <|tools_prefix|>[...]<|tools_suffix|>) for any tool calls inside reasoning
+        """
+        tool_start = self._tool_start_token
+        tool_end = self._tool_end_token
+        out: list[tuple[str, str]] = []
+        cursor = 0
+
+        while True:
+            if (s := inner_text.find(tool_start, cursor)) == -1:
+                if (tail := inner_text[cursor:]) != "":
+                    out.append(("reasoning", tail))
+                break
+            if s > cursor:
+                out.append(("reasoning", inner_text[cursor:s]))
+
+            if (e := inner_text.find(tool_end, s)) == -1:
+                out.append(("text", inner_text[s:]))
+                break
+
+            out.append(("text", inner_text[s : e + len(tool_end)]))
+            cursor = e + len(tool_end)
+
+        return out
+
+    def parse_streaming_increment(self, new_text: str) -> StreamingParseResult:
+        self._buffer += new_text
+
+        out_reasoning = ""
+        out_normal = ""
+
+        start_tok = self.think_start_token
+        end_tok = self.think_end_token
+        tool_start = self._tool_start_token
+        tool_end = self._tool_end_token
+
+        while True:
+            if not self._in_reasoning:
+                if (s := self._buffer.find(start_tok)) == -1:
+                    if partial := self._ends_with_partial_token(
+                        self._buffer, start_tok
+                    ):
+                        out_normal += self._buffer[:-partial]
+                        self._buffer = self._buffer[-partial:]
+                    else:
+                        out_normal += self._buffer
+                        self._buffer = ""
+                    return StreamingParseResult(
+                        normal_text=out_normal, reasoning_text=out_reasoning
+                    )
+
+                out_normal += self._buffer[:s]
+                self._buffer = self._buffer[s + len(start_tok) :]
+                self._in_reasoning = True
+                self._reasoning_acc = ""
+                self._in_inner_tool = False
+                continue
+
+            if self._in_inner_tool:
+                if (end_pos := self._buffer.find(tool_end)) == -1:
+                    if (
+                        hold := self._ends_with_partial_token(self._buffer, tool_end)
+                    ) != 0:
+                        out_normal += self._buffer[:-hold]
+                        self._buffer = self._buffer[-hold:]
+                    else:
+                        out_normal += self._buffer
+                        self._buffer = ""
+                    return StreamingParseResult(
+                        normal_text=out_normal, reasoning_text=out_reasoning
+                    )
+
+                out_normal += self._buffer[: end_pos + len(tool_end)]
+                self._buffer = self._buffer[end_pos + len(tool_end) :]
+                self._in_inner_tool = False
+                continue
+
+            pos_tool = self._buffer.find(tool_start)
+            pos_end = self._buffer.find(end_tok)
+
+            if pos_tool == -1 and pos_end == -1:
+                if self.stream_reasoning:
+                    if (
+                        hold := max(
+                            self._ends_with_partial_token(self._buffer, end_tok),
+                            self._ends_with_partial_token(self._buffer, tool_start),
+                        )
+                    ) != 0:
+                        out_reasoning += self._buffer[:-hold]
+                        self._buffer = self._buffer[-hold:]
+                    else:
+                        out_reasoning += self._buffer
+                        self._buffer = ""
+                return StreamingParseResult(
+                    normal_text=out_normal, reasoning_text=out_reasoning
+                )
+
+            next_pos = min(p for p in [pos_tool, pos_end] if p != -1)
+
+            if pos_end != -1 and pos_end == next_pos:
+                reasoning_chunk = self._buffer[:pos_end]
+                if self.stream_reasoning:
+                    out_reasoning += reasoning_chunk
+                else:
+                    self._reasoning_acc += reasoning_chunk
+                    out_reasoning += self._reasoning_acc
+                    self._reasoning_acc = ""
+                self._buffer = self._buffer[pos_end + len(end_tok) :]
+                self._in_reasoning = False
+                continue
+
+            reasoning_chunk = self._buffer[:pos_tool]
+            if self.stream_reasoning:
+                out_reasoning += reasoning_chunk
+            else:
+                self._reasoning_acc += reasoning_chunk
+            self._buffer = self._buffer[pos_tool:]
+            self._in_inner_tool = True
+            continue
+
+
+class CohereCommand4Detector(BaseReasoningFormatDetector):
+    """Detector for Cohere Command4 / Command-A family (incl. cohere2_moe and
+    cohere2_vision Command-A-Plus).
+
+    Generated format (the assistant prefix in the chat template already emits
+    ``<|START_THINKING|>`` when ``reasoning=True``, so the *generated* text
+    typically begins inside the thinking block):
+
+        thinking_content<|END_THINKING|><|START_TEXT|>final_answer<|END_TEXT|>
+
+    When ``reasoning=False`` the chat template emits both START/END_THINKING
+    in the prefix and the generated text is just::
+
+        <|START_TEXT|>final_answer<|END_TEXT|>
+
+    This detector returns:
+      - ``reasoning_text`` = the thinking block (between START_THINKING and
+        END_THINKING, with the START tag stripped if the model echoed it).
+      - ``normal_text`` = the content between ``<|START_TEXT|>`` and
+        ``<|END_TEXT|>``, with both markers stripped. If no ``<|START_TEXT|>``
+        appears (the model exhausted max_new_tokens still inside thinking),
+        ``normal_text`` is the empty string.
+
+    Matches the public token names from the model's
+    ``special_tokens_map.json`` (``<|START_THINKING|>`` etc.).
+    """
+
+    TEXT_START_TOKEN = "<|START_TEXT|>"
+    TEXT_END_TOKEN = "<|END_TEXT|>"
+    # When the model decides to call tools instead of producing a final text
+    # block, it emits an action block instead of a text block. The reasoning
+    # parser must leave that block intact so the downstream tool-call parser
+    # can pick it up.
+    ACTION_START_TOKEN = "<|START_ACTION|>"
+
+    def __init__(
+        self,
+        stream_reasoning: bool = True,
+        force_reasoning: bool = True,
+        continue_final_message: bool = False,
+        previous_content: str = "",
+        force_nonempty_content: bool = False,
+    ):
+        # The chat template puts <|START_THINKING|> in the assistant prefix
+        # when reasoning is enabled, so the *generated* text usually starts
+        # already inside thinking. ``force_reasoning=True`` makes the base
+        # detector treat the leading bytes as reasoning even though the
+        # generated stream typically does not echo <|START_THINKING|>.
+        super().__init__(
+            think_start_token="<|START_THINKING|>",
+            think_end_token="<|END_THINKING|>",
+            force_reasoning=force_reasoning,
+            stream_reasoning=stream_reasoning,
+            continue_final_message=continue_final_message,
+            previous_content=previous_content,
+            force_nonempty_content=force_nonempty_content,
+        )
+        # Streaming state machine. The model emits, in order:
+        #   1. reasoning  (between START_THINKING [in prefix] and END_THINKING)
+        #   2. either ``<|START_TEXT|>...<|END_TEXT|>`` (final answer) or
+        #      ``<|START_ACTION|>...<|END_ACTION|>`` (tool calls) -- never both.
+        # When ``reasoning=False`` the chat template emits both START/END
+        # thinking in the prefix and step 1 is empty; the generated stream
+        # then starts directly with the text or action block.
+        self._reasoning_done = False
+        self._saw_text_start = False
+        self._saw_text_end = False
+        self._in_action_mode = False
+
+    @classmethod
+    def _strip_text_markers(cls, raw: str) -> str:
+        """Extract the substring between ``<|START_TEXT|>`` and
+        ``<|END_TEXT|>``. If ``<|START_TEXT|>`` is absent but a
+        ``<|START_ACTION|>`` block is present, the model produced a tool
+        call instead of a text answer -- return the raw text untouched so
+        the downstream tool-call parser can pick up the action block. If
+        neither marker is present (ran out of tokens still inside
+        thinking) return ``""``. If ``<|END_TEXT|>`` is absent (stop token
+        or max_new_tokens cut the stream off inside the text block) return
+        everything after ``<|START_TEXT|>``.
+        """
+        if not raw:
+            return ""
+        s = raw.find(cls.TEXT_START_TOKEN)
+        if s == -1:
+            if cls.ACTION_START_TOKEN in raw:
+                return raw
+            return ""
+        s += len(cls.TEXT_START_TOKEN)
+        tail = raw[s:]
+        e = tail.find(cls.TEXT_END_TOKEN)
+        if e == -1:
+            return tail
+        return tail[:e]
+
+    def detect_and_parse(self, text: str) -> StreamingParseResult:
+        # Direct parse: split on the (single) ``<|END_THINKING|>`` token if
+        # present. Anything before is reasoning, anything after is the
+        # final-text block. If no END_THINKING but a START_TEXT exists,
+        # we're in the reasoning=False case (chat template emitted both
+        # START/END thinking in the prefix; the model only generated the
+        # text block). Otherwise the model exhausted tokens still thinking
+        # and ``normal_text`` ends up empty -- matching the convention of
+        # the other detectors in this module (DeepSeekR1, Qwen3, ...). The
+        # empty content is propagated as ``message.content = None`` by
+        # serving_chat, and downstream code is expected to treat that as
+        # "no answer" rather than falling back to ``reasoning_content``.
+        end_think_idx = text.find(self.think_end_token)
+        text_start_idx = text.find(self.TEXT_START_TOKEN)
+        action_start_idx = text.find(self.ACTION_START_TOKEN)
+        if end_think_idx != -1:
+            reasoning = text[:end_think_idx]
+            rest = text[end_think_idx + len(self.think_end_token) :]
+        elif text_start_idx != -1:
+            reasoning = text[:text_start_idx]
+            rest = text[text_start_idx:]
+        elif action_start_idx != -1:
+            # reasoning=False + tool call: chat template emitted both
+            # START/END thinking in the prefix, the model only generated
+            # an action block. Treat the prefix before the action block as
+            # (probably empty) reasoning so the action block reaches the
+            # tool-call parser intact.
+            reasoning = text[:action_start_idx]
+            rest = text[action_start_idx:]
+        else:
+            reasoning = text
+            rest = ""
+
+        # Some checkpoints echo the START_THINKING token even though the
+        # chat template put it in the prefix; drop it if so.
+        think_start_text = self.think_start_token + self.think_start_self_label
+        if reasoning.startswith(think_start_text):
+            reasoning = reasoning[len(think_start_text) :]
+
+        return self._maybe_apply_force_nonempty_content(
+            StreamingParseResult(
+                normal_text=self._strip_text_markers(rest),
+                reasoning_text=reasoning,
+            )
+        )
+
+    def parse_streaming_increment(self, new_text: str) -> StreamingParseResult:
+        """Streaming parse. Custom state machine -- we don't reuse the base
+        class because Cohere's "reasoning=False" path (the model emits no
+        ``<|END_THINKING|>``, just goes straight to a text or action block)
+        is fundamentally incompatible with the base detector's
+        ``force_reasoning`` semantics."""
+        self._buffer += new_text
+        buf = self._buffer
+
+        if not self._reasoning_done:
+            # Look for any marker that ends reasoning: an explicit
+            # END_THINKING, or an implicit transition via the start of the
+            # final-text or action block (reasoning=False case).
+            markers = (
+                (self.think_end_token, "think_end"),
+                (self.TEXT_START_TOKEN, "text"),
+                (self.ACTION_START_TOKEN, "action"),
+            )
+            first_pos = None
+            first_marker = None
+            first_kind = None
+            for marker_text, kind in markers:
+                p = buf.find(marker_text)
+                if p != -1 and (first_pos is None or p < first_pos):
+                    first_pos, first_marker, first_kind = p, marker_text, kind
+            if first_pos is None:
+                # No marker seen yet. Stream the reasoning prefix, but keep
+                # enough tail in the buffer to recognise a marker split
+                # across chunk boundaries.
+                if not self.stream_reasoning:
+                    return StreamingParseResult()
+                max_keep = max(len(m) for m, _ in markers) - 1
+                if len(buf) > max_keep:
+                    head = buf[:-max_keep]
+                    self._buffer = buf[-max_keep:]
+                    return StreamingParseResult(reasoning_text=head)
+                return StreamingParseResult()
+
+            reasoning_chunk = buf[:first_pos]
+            if first_kind == "think_end":
+                self._buffer = buf[first_pos + len(first_marker) :]
+            else:
+                # Implicit reasoning-end: leave the start-of-block marker in
+                # the buffer for the post-thinking branch below to consume.
+                self._buffer = buf[first_pos:]
+            self._reasoning_done = True
+            if reasoning_chunk:
+                return StreamingParseResult(reasoning_text=reasoning_chunk)
+            buf = self._buffer
+
+        # Reasoning is closed. Decide between text-stripping and
+        # action-passthrough on first sight of a marker.
+        if self._in_action_mode:
+            if not buf:
+                return StreamingParseResult()
+            self._buffer = ""
+            return StreamingParseResult(normal_text=buf)
+
+        if not self._saw_text_start:
+            s_text = buf.find(self.TEXT_START_TOKEN)
+            s_action = buf.find(self.ACTION_START_TOKEN)
+            picks = [
+                (p, k) for p, k in ((s_text, "text"), (s_action, "action")) if p != -1
+            ]
+            if not picks:
+                max_keep = (
+                    max(len(self.TEXT_START_TOKEN), len(self.ACTION_START_TOKEN)) - 1
+                )
+                if len(buf) > max_keep:
+                    self._buffer = buf[-max_keep:]
+                return StreamingParseResult()
+            picks.sort()
+            first_pos, first_kind = picks[0]
+            if first_kind == "action":
+                self._in_action_mode = True
+                out_normal = buf[first_pos:]
+                self._buffer = ""
+                return StreamingParseResult(normal_text=out_normal)
+            # Found <|START_TEXT|>. Drop everything up to and including the
+            # marker -- text content streams next.
+            self._buffer = buf[first_pos + len(self.TEXT_START_TOKEN) :]
+            self._saw_text_start = True
+            buf = self._buffer
+
+        if self._saw_text_start and not self._saw_text_end:
+            e = buf.find(self.TEXT_END_TOKEN)
+            if e == -1:
+                # Emit everything except a possible partial END_TEXT tail.
+                keep = len(self.TEXT_END_TOKEN) - 1
+                if len(buf) > keep:
+                    out_normal = buf[:-keep]
+                    self._buffer = buf[-keep:]
+                    return StreamingParseResult(normal_text=out_normal)
+                return StreamingParseResult()
+            out_normal = buf[:e]
+            self._buffer = buf[e + len(self.TEXT_END_TOKEN) :]
+            self._saw_text_end = True
+            return StreamingParseResult(normal_text=out_normal)
+
+        return StreamingParseResult()
+
+
+class ReasoningParser:
+    """
+    Parser that handles both streaming and non-streaming scenarios for extracting
+    reasoning content from model outputs.
+
+    Args:
+        model_type (str): Type of model to parse reasoning from
+        stream_reasoning (bool): If False, accumulates reasoning content until complete.
+            If True, streams reasoning content as it arrives.
+    """
+
+    DetectorMap: Dict[str, Type[BaseReasoningFormatDetector]] = {
+        "apertus2509": Apertus2509Detector,
+        "deepseek-r1": DeepSeekR1Detector,
+        "deepseek-v3": _DeepSeekV3Detector,
+        "agnes": _DeepSeekV3Detector,  # alias: Agnes uses the deepseek-v4 format
+        "deepseek-v4": _DeepSeekV3Detector,
+        "glm45": Glm45Detector,
+        "hunyuan": HunyuanDetector,
+        "gpt-oss": GptOssDetector,
+        "kimi": KimiDetector,
+        "kimi_k2": KimiK2Detector,
+        "mimo": _MimoDetector,
+        "poolside_v1": _PoolsideV1Detector,
+        "qwen3": Qwen3Detector,
+        "qwen3-thinking": Qwen3Detector,
+        "minimax": Qwen3Detector,
+        "minimax-append-think": MiniMaxAppendThinkDetector,
+        "minimax-m3": MiniMaxM3Detector,
+        "step3": DeepSeekR1Detector,
+        "step3p5": DeepSeekR1Detector,
+        "mistral": MistralDetector,
+        "nemotron_3": Nemotron3Detector,
+        "interns1": Qwen3Detector,
+        "gemma4": Gemma4Detector,
+        "inkling": InklingDetector,
+        "cohere_command4": CohereCommand4Detector,
+    }
+
+    def __init__(
+        self,
+        model_type: Optional[str] = None,
+        stream_reasoning: bool = True,
+        force_reasoning: Optional[bool] = None,
+        request: ChatCompletionRequest = None,
+        tokenizer=None,
+    ):
+        if not model_type:
+            raise ValueError("Model type must be specified")
+
+        detector_class = self.DetectorMap.get(model_type.lower())
+        if not detector_class:
+            raise ValueError(f"Unsupported model type: {model_type}")
+
+        chat_template_kwargs = getattr(request, "chat_template_kwargs", None) or {}
+
+        # Special cases where we override force_reasoning
+        if model_type.lower() in {
+            "qwen3-thinking",
+            "gpt-oss",
+            "minimax",
+        }:
+            force_reasoning = True
+
+        # M3 consumes the <mm:think> start tag only for thinking_mode=enabled
+        # (absent from output → must force); mirror serving_chat's M3 branch.
+        if model_type.lower() == "minimax-m3" and force_reasoning is None:
+            force_reasoning = chat_template_kwargs.get("thinking_mode") == "enabled"
+
+        # Only pass force_reasoning if explicitly set, let detectors use their defaults
+        kwargs = {"stream_reasoning": stream_reasoning}
+        if force_reasoning is not None:
+            kwargs["force_reasoning"] = force_reasoning
+
+        if (
+            request is not None
+            and isinstance(request, ChatCompletionRequest)
+            and request.continue_final_message
+            and request.messages[-1].role == "assistant"
+        ):
+            kwargs["continue_final_message"] = True
+            kwargs["previous_content"] = request.messages[-1].content
+
+        if chat_template_kwargs.get("force_nonempty_content") is True:
+            kwargs["force_nonempty_content"] = True
+
+        if tokenizer is not None:
+            sig = inspect.signature(detector_class)
+            if "tokenizer" in sig.parameters:
+                kwargs["tokenizer"] = tokenizer
+
+        self.detector = detector_class(**kwargs)
+
+    def parse_non_stream(self, full_text: str) -> Tuple[Optional[str], Optional[str]]:
+        """Non-streaming call: one-time parsing"""
+        ret = self.detector.detect_and_parse(full_text)
+        return ret.reasoning_text, ret.normal_text
+
+    def parse_non_stream_blocks(self, full_text: str) -> list[dict]:
+        """Non-streaming call: return an ordered sequence of reasoning/text blocks"""
+        if hasattr(self.detector, "detect_and_parse_block_sequence"):
+            seq = self.detector.detect_and_parse_block_sequence(full_text)
+            return [{"type": k, "text": t} for k, t in seq]
+
+        ret = self.detector.detect_and_parse(full_text)
+        blocks: list[dict] = []
+        if ret.reasoning_text:
+            blocks.append({"type": "reasoning", "text": ret.reasoning_text})
+        blocks.append({"type": "text", "text": ret.normal_text or ""})
+        return blocks
+
+    def parse_stream_chunk(
+        self, chunk_text: str
+    ) -> Tuple[Optional[str], Optional[str]]:
+        """Streaming call: incremental parsing"""
+        ret = self.detector.parse_streaming_increment(chunk_text)
+        return ret.reasoning_text, ret.normal_text
+
+    def parse_stream_end(self) -> Tuple[Optional[str], Optional[str]]:
+        """Streaming call: flush any detector-specific buffered state once
+        the stream ends."""
+        ret = self.detector.finish()
+        return ret.reasoning_text, ret.normal_text

sglang_patch/srt/parser/template_detection.py

@@ -0,0 +1,707 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""
+Template detection utilities for auto-detecting reasoning and tool-call parsers.
+
+Provides rule-based detection of reasoning mode, reasoning parser, and tool-call
+parser from chat templates and tokenizer vocabularies.
+"""
+
+import logging
+import os
+import re
+from dataclasses import dataclass
+from typing import Callable, Optional, Tuple
+
+import jinja2
+import jinja2.ext
+import jinja2.sandbox
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class TemplateDetectionContext:
+    template: str
+    reasoning_config: Optional["ReasoningToggleConfig"]
+    force_reasoning: bool
+    vocab: set[str]
+
+    def has_text(self, needle: str) -> bool:
+        return needle in self.template
+
+    def has_vocab(self, token: str) -> bool:
+        return token in self.vocab
+
+    def has_pattern(self, pattern: str, flags: int = 0) -> bool:
+        return re.search(pattern, self.template, flags) is not None
+
+    def has_vocab_pattern(self, pattern: str) -> bool:
+        compiled = re.compile(pattern)
+        return any(isinstance(tok, str) and compiled.search(tok) for tok in self.vocab)
+
+
+@dataclass(frozen=True)
+class DetectionRule:
+    name: str
+    value: object
+    predicate: Callable[[TemplateDetectionContext], bool]
+
+
+@dataclass(frozen=True)
+class ReasoningToggleConfig:
+    toggle_param: Optional[str] = None
+    default_enabled: Optional[bool] = None
+    special_case: Optional[str] = None
+    effort_kwarg: Optional[str] = None
+
+    @property
+    def always_on(self) -> bool:
+        return self.special_case == "always"
+
+
+# ---------------------------------------------------------------------------
+# Reasoning mode rules (detect toggle config from template)
+# ---------------------------------------------------------------------------
+
+REASONING_MODE_RULES = (
+    DetectionRule(
+        name="gpt_oss_channel_markers",
+        value=ReasoningToggleConfig(special_case="always"),
+        predicate=lambda ctx: ctx.has_text("<|channel|>"),
+    ),
+    DetectionRule(
+        name="force_reasoning_pattern",
+        value=ReasoningToggleConfig(special_case="always"),
+        predicate=lambda ctx: ctx.has_pattern(r"<\|im_start\|>assistant\\n<think>\\n")
+        and not ctx.has_text("enable_thinking")
+        and not ctx.has_text("thinking"),
+    ),
+    DetectionRule(
+        name="mistral_reasoning_effort",
+        value=ReasoningToggleConfig(special_case="mistral"),
+        predicate=lambda ctx: ctx.has_text("reasoning_effort")
+        and ctx.has_text("[THINK]"),
+    ),
+    DetectionRule(
+        name="explicit_enable_thinking_default_false",
+        value=ReasoningToggleConfig(
+            toggle_param="enable_thinking", default_enabled=False
+        ),
+        predicate=lambda ctx: ctx.has_pattern(
+            r"{%\s*if\s+not\s+enable_thinking\s+is\s+defined\s*%}.*?"
+            r"{%\s*set\s+enable_thinking\s*=\s*(?:false|False)\s*%}",
+            re.DOTALL,
+        ),
+    ),
+    DetectionRule(
+        name="nemotron_3_super_low_effort",
+        value=ReasoningToggleConfig(
+            toggle_param="enable_thinking",
+            default_enabled=True,
+            effort_kwarg="low_effort",
+        ),
+        predicate=lambda ctx: ctx.has_text("low_effort")
+        and ctx.has_text("truncate_history_thinking"),
+    ),
+    DetectionRule(
+        name="enable_thinking_default_true",
+        value=ReasoningToggleConfig(
+            toggle_param="enable_thinking", default_enabled=True
+        ),
+        predicate=lambda ctx: ctx.has_pattern(
+            r"{%\s*if\s+not\s+enable_thinking\s+is\s+defined\s*%}.*?"
+            r"{%\s*set\s+enable_thinking\s*=\s*(?:true|True)\s*%}",
+            re.DOTALL,
+        )
+        or ctx.has_pattern(
+            r"set\s+enable_thinking\s*=\s*enable_thinking\s+if\s+enable_thinking\s+is\s+defined\s+else\s+(?:true|True)"
+        )
+        or ctx.has_pattern(
+            r"enable_thinking\s+is\s+defined\s+and\s+(?:enable_thinking\s+is\s+false|not\s+enable_thinking)"
+        )
+        or ctx.has_pattern(
+            r"enable_thinking\s+is\s+not\s+defined\s+or\s+enable_thinking"
+        )
+        or ctx.has_pattern(r"namespace\([^)]*enable_thinking\s*=\s*true"),
+    ),
+    DetectionRule(
+        name="explicit_thinking_default_false",
+        value=ReasoningToggleConfig(toggle_param="thinking", default_enabled=False),
+        predicate=lambda ctx: ctx.has_pattern(
+            r"{%\s*if\s+not\s+thinking\s+is\s+defined\s*%}.*?"
+            r"{%\s*set\s+thinking\s*=\s*(?:false|False)\s*%}",
+            re.DOTALL,
+        ),
+    ),
+    DetectionRule(
+        name="thinking_default_true",
+        value=ReasoningToggleConfig(toggle_param="thinking", default_enabled=True),
+        predicate=lambda ctx: ctx.has_pattern(
+            r"{%\s*if\s+not\s+thinking\s+is\s+defined\s*%}.*?"
+            r"{%\s*set\s+thinking\s*=\s*(?:true|True)\s*%}",
+            re.DOTALL,
+        )
+        or ctx.has_pattern(
+            r"set\s+thinking\s*=\s*thinking\s+if\s+thinking\s+is\s+defined\s+else\s+(?:true|True)"
+        )
+        or ctx.has_pattern(
+            r"thinking\s+is\s+defined\s+and\s+(?:thinking\s+is\s+false|not\s+thinking)"
+        )
+        or ctx.has_pattern(r"thinking\s+is\s+not\s+defined\s+or\s+thinking")
+        or ctx.has_pattern(r"namespace\([^)]*thinking\s*=\s*true"),
+    ),
+)
+
+
+# ---------------------------------------------------------------------------
+# Shared predicates for model-family detection
+# ---------------------------------------------------------------------------
+
+
+def _is_apertus2509(ctx):
+    return ctx.has_vocab("<|inner_prefix|>")
+
+
+def _is_gemma4(ctx):
+    return ctx.has_text("<|channel>")
+
+
+def _is_kimi(ctx):
+    return ctx.has_text("◁think▷")
+
+
+def _is_interns1(ctx):
+    return ctx.has_text("default_thinking_sys") and ctx.reasoning_config == (
+        ReasoningToggleConfig(toggle_param="enable_thinking", default_enabled=True)
+    )
+
+
+def _is_mistral(ctx):
+    return (
+        ctx.reasoning_config is not None
+        and ctx.reasoning_config.special_case == "mistral"
+    )
+
+
+def _is_gpt_oss(ctx):
+    return ctx.has_text("<|channel|>")
+
+
+def _is_kimi_k2(ctx):
+    return ctx.has_vocab("<|tool_calls_section_begin|>")
+
+
+def _is_nemotron_3(ctx):
+    return ctx.has_text("truncate_history_thinking") and (
+        ctx.reasoning_config is not None
+        and ctx.reasoning_config.toggle_param == "enable_thinking"
+        and ctx.reasoning_config.default_enabled is True
+    )
+
+
+def _is_glm45(ctx):
+    return (
+        (
+            ctx.has_text("[gMASK]<sop>")
+            or ctx.has_pattern(r"(?<!<)/nothink")
+            or ctx.has_pattern(r"(?<!<)/think")
+        )
+        and ctx.has_vocab("<tool_call>")
+        and ctx.reasoning_config
+        == ReasoningToggleConfig(toggle_param="enable_thinking", default_enabled=True)
+        and (ctx.has_vocab("<|user|>") or ctx.has_vocab("<|endoftext|>"))
+    )
+
+
+def _is_glm47(ctx):
+    return _is_glm45(ctx) and ctx.has_pattern(
+        r"\{\{[-\s]*['\"]<tool_call>['\"]\s*\+\s*tc\.name"
+    )
+
+
+def _is_xml_kv_tool_call(ctx):
+    # Structural signature for the GLM-4.5 / GLM-4.6 style tool-call format
+    # (`<tool_call>name<arg_key>k</arg_key>\n<arg_value>v</arg_value>...</tool_call>`).
+    # Matches any model whose tokenizer carries `<arg_key>` and `<arg_value>` as
+    # added tokens — e.g., inclusionAI/Ring-2.6, which borrows GLM's tool-call
+    # format but doesn't share the `[gMASK]<sop>` / `enable_thinking` family
+    # signature checked by `_is_glm45`.
+    return ctx.has_vocab("<arg_key>") and ctx.has_vocab("<arg_value>")
+
+
+def _is_deepseek_v31(ctx):
+    return ctx.has_text("<｜tool▁calls▁begin｜>") and ctx.has_text("<｜tool▁sep｜>")
+
+
+def _is_deepseek_v32(ctx):
+    return ctx.has_text("<｜DSML｜function_calls>")
+
+
+def _is_deepseek_v4(ctx):
+    return ctx.has_text("<｜DSML｜tool_calls>")
+
+
+def _is_hunyuan(ctx):
+    # The shipping Hy3 tokenizer appends a shared suffix to each special token
+    # (e.g. ``<tool_calls:opensource>``), so match the bare or suffixed form.
+    tc = ctx.has_text("<tool_calls>") or ctx.has_vocab_pattern(
+        r"^<tool_calls(?::[^>]+)?>$"
+    )
+    sep = ctx.has_text("<tool_sep>") or ctx.has_vocab_pattern(
+        r"^<tool_sep(?::[^>]+)?>$"
+    )
+    return (tc and sep) or (
+        ctx.has_text("reasoning_effort") and ctx.has_text("interleaved_thinking")
+    )
+
+
+def _is_poolside_v1(ctx):
+    has_poolside_tool_format = (
+        ctx.has_text("unescaped XML-like object")
+        and ctx.has_text("<tool_call>function-name")
+        and ctx.has_text("<arg_key>")
+        and ctx.has_text("<arg_value>")
+    )
+    return has_poolside_tool_format or (
+        ctx.reasoning_config
+        == ReasoningToggleConfig(toggle_param="enable_thinking", default_enabled=False)
+        and not _is_hunyuan(ctx)
+        and (ctx.has_text("<arg_key>") or ctx.has_vocab("<arg_key>"))
+        and (ctx.has_text("<arg_value>") or ctx.has_vocab("<arg_value>"))
+    )
+
+
+def _is_mimo(ctx):
+    return ctx.reasoning_config == ReasoningToggleConfig(
+        toggle_param="enable_thinking", default_enabled=False
+    )
+
+
+def _is_minimax(ctx):
+    return ctx.has_text("<minimax:tool_call>")
+
+
+def _is_minimax_m3(ctx):
+    return ctx.has_text("<mm:think>") or ctx.has_text("]<]minimax[>[")
+
+
+def _is_minicpm5(ctx):
+    if ctx.has_vocab("<function") and ctx.has_vocab("<param"):
+        return True
+    return ctx.has_pattern(r"<function\s+name=") and ctx.has_pattern(r"<param\s+name=")
+
+
+def _is_lfm2(ctx):
+    return (
+        ctx.has_text("<|tool_call_start|>") or ctx.has_vocab("<|tool_call_start|>")
+    ) and (ctx.has_text("<|tool_call_end|>") or ctx.has_vocab("<|tool_call_end|>"))
+
+
+def _is_step3p5(ctx):
+    return ctx.has_pattern(r"Step-?3(?:\.|p)?[57]", re.IGNORECASE) or (
+        ctx.has_text("reasoning_effort")
+        and ctx.has_text("Reasoning: ")
+        and _is_qwen3_coder(ctx)
+    )
+
+
+def _is_step3(ctx):
+    return ctx.has_text("<steptml:invoke") or (
+        ctx.has_text("<｜tool_calls_begin｜>") and ctx.has_text("<｜tool_sep｜>")
+    )
+
+
+def _is_qwen3_coder(ctx):
+    return ctx.has_text("<function=") and ctx.has_text("<parameter=")
+
+
+def _is_qwen3(ctx):
+    return ctx.reasoning_config == ReasoningToggleConfig(
+        toggle_param="enable_thinking", default_enabled=True
+    )
+
+
+def _is_deepseek_v3(ctx):
+    return ctx.reasoning_config == ReasoningToggleConfig(
+        toggle_param="thinking", default_enabled=False
+    )
+
+
+def _is_deepseek_r1(ctx):
+    return ctx.force_reasoning
+
+
+def _is_deepseek_r1_think_tags(ctx):
+    return not _is_lfm2(ctx) and (ctx.has_text("<think>") or ctx.has_text("</think>"))
+
+
+# ---------------------------------------------------------------------------
+# Reasoning parser rules
+# ---------------------------------------------------------------------------
+
+REASONING_PARSER_RULES = (
+    DetectionRule(name="apertus2509", value="apertus2509", predicate=_is_apertus2509),
+    DetectionRule(name="gemma4", value="gemma4", predicate=_is_gemma4),
+    DetectionRule(name="kimi", value="kimi", predicate=_is_kimi),
+    DetectionRule(name="interns1", value="interns1", predicate=_is_interns1),
+    DetectionRule(name="mistral", value="mistral", predicate=_is_mistral),
+    DetectionRule(name="gpt_oss", value="gpt-oss", predicate=_is_gpt_oss),
+    DetectionRule(name="kimi_k2", value="kimi_k2", predicate=_is_kimi_k2),
+    DetectionRule(name="nemotron_3", value="nemotron_3", predicate=_is_nemotron_3),
+    DetectionRule(name="glm45", value="glm45", predicate=_is_glm45),
+    DetectionRule(name="hunyuan", value="hunyuan", predicate=_is_hunyuan),
+    DetectionRule(name="poolside_v1", value="poolside_v1", predicate=_is_poolside_v1),
+    DetectionRule(name="mimo", value="mimo", predicate=_is_mimo),
+    DetectionRule(name="minimax_m3", value="minimax-m3", predicate=_is_minimax_m3),
+    DetectionRule(name="minimax", value="minimax", predicate=_is_minimax),
+    DetectionRule(name="step3p5", value="step3p5", predicate=_is_step3p5),
+    DetectionRule(name="step3", value="step3", predicate=_is_step3),
+    DetectionRule(name="qwen3", value="qwen3", predicate=_is_qwen3),
+    DetectionRule(name="deepseek_v4", value="deepseek-v4", predicate=_is_deepseek_v4),
+    DetectionRule(name="deepseek_v3", value="deepseek-v3", predicate=_is_deepseek_v3),
+    DetectionRule(
+        name="deepseek_r1_force", value="deepseek-r1", predicate=_is_deepseek_r1
+    ),
+    DetectionRule(
+        name="deepseek_r1_think_tags",
+        value="deepseek-r1",
+        predicate=_is_deepseek_r1_think_tags,
+    ),
+)
+
+# ---------------------------------------------------------------------------
+# Tool-call parser rules (reuse shared predicates, different values)
+# ---------------------------------------------------------------------------
+
+TOOL_CALL_PARSER_RULES = (
+    DetectionRule(name="apertus2509", value="apertus2509", predicate=_is_apertus2509),
+    DetectionRule(name="gemma4", value="gemma4", predicate=_is_gemma4),
+    DetectionRule(name="gpt_oss", value="gpt-oss", predicate=_is_gpt_oss),
+    DetectionRule(name="kimi_k2", value="kimi_k2", predicate=_is_kimi_k2),
+    DetectionRule(name="minimax_m3", value="minimax-m3", predicate=_is_minimax_m3),
+    DetectionRule(name="minimax", value="minimax-m2", predicate=_is_minimax),
+    DetectionRule(name="interns1", value="interns1", predicate=_is_interns1),
+    DetectionRule(name="mistral", value="mistral", predicate=_is_mistral),
+    DetectionRule(name="deepseek_v4", value="deepseekv4", predicate=_is_deepseek_v4),
+    DetectionRule(name="deepseek_v32", value="deepseekv32", predicate=_is_deepseek_v32),
+    DetectionRule(name="deepseek_v31", value="deepseekv31", predicate=_is_deepseek_v31),
+    DetectionRule(name="lfm2", value="lfm2", predicate=_is_lfm2),
+    DetectionRule(name="glm47", value="glm47", predicate=_is_glm47),
+    DetectionRule(name="glm45", value="glm45", predicate=_is_glm45),
+    DetectionRule(name="minicpm5", value="minicpm5", predicate=_is_minicpm5),
+    DetectionRule(name="hunyuan", value="hunyuan", predicate=_is_hunyuan),
+    DetectionRule(name="poolside_v1", value="poolside_v1", predicate=_is_poolside_v1),
+    DetectionRule(name="step3p5", value="step3p5", predicate=_is_step3p5),
+    DetectionRule(name="step3", value="step3", predicate=_is_step3),
+    DetectionRule(
+        name="xml_kv_tool_call", value="glm45", predicate=_is_xml_kv_tool_call
+    ),
+    DetectionRule(name="mimo", value="mimo", predicate=_is_mimo),
+    DetectionRule(name="qwen3_coder", value="qwen3_coder", predicate=_is_qwen3_coder),
+    DetectionRule(name="qwen", value="qwen", predicate=_is_qwen3),
+    DetectionRule(name="deepseek_v3", value="deepseekv3", predicate=_is_deepseek_v3),
+    DetectionRule(name="deepseek_r1", value="deepseekv3", predicate=_is_deepseek_r1),
+)
+
+
+# ---------------------------------------------------------------------------
+# Detection functions
+# ---------------------------------------------------------------------------
+
+
+def build_detection_context(
+    template: Optional[str],
+    tokenizer,
+    reasoning_config: Optional[ReasoningToggleConfig] = None,
+    force_reasoning: bool = False,
+) -> Optional[TemplateDetectionContext]:
+    if template is None:
+        return None
+    vocab = set()
+    if tokenizer is not None:
+        try:
+            vocab = set(tokenizer.get_vocab().keys())
+        except Exception as e:
+            logger.warning(
+                "Failed to load tokenizer vocab for template detection: %s. "
+                "Vocab-dependent detection rules will be skipped.",
+                e,
+            )
+    return TemplateDetectionContext(
+        template=template,
+        reasoning_config=reasoning_config,
+        force_reasoning=force_reasoning,
+        vocab=vocab,
+    )
+
+
+def match_rules(
+    ctx: TemplateDetectionContext,
+    rules: Tuple[DetectionRule, ...],
+    label: str,
+) -> Optional[str]:
+    for rule in rules:
+        try:
+            if rule.predicate(ctx):
+                return rule.value
+        except Exception as e:
+            logger.warning(
+                "Detection rule '%s' for %s raised an exception: %s. Skipping.",
+                rule.name,
+                label,
+                e,
+                exc_info=True,
+            )
+    return None
+
+
+def detect_reasoning_pattern(
+    template: Optional[str],
+) -> Tuple[bool, Optional[ReasoningToggleConfig]]:
+    """Detect if the chat template contains reasoning/thinking patterns."""
+    if template is None:
+        return False, None
+
+    ctx = TemplateDetectionContext(
+        template=template,
+        reasoning_config=None,
+        force_reasoning=False,
+        vocab=set(),
+    )
+    for rule in REASONING_MODE_RULES:
+        if rule.predicate(ctx):
+            return rule.value.always_on, rule.value
+
+    return False, None
+
+
+def detect_reasoning_parser(
+    template: Optional[str],
+    tokenizer,
+    reasoning_config: Optional[ReasoningToggleConfig] = None,
+    force_reasoning: bool = False,
+) -> Optional[str]:
+    """Auto-detect which reasoning parser to use from the chat template."""
+    ctx = build_detection_context(
+        template, tokenizer, reasoning_config, force_reasoning
+    )
+    if ctx is None:
+        return None
+    return match_rules(ctx, REASONING_PARSER_RULES, "reasoning parser")
+
+
+def detect_tool_call_parser(
+    template: Optional[str],
+    tokenizer,
+    reasoning_config: Optional[ReasoningToggleConfig] = None,
+    force_reasoning: bool = False,
+) -> Optional[str]:
+    """Auto-detect which tool-call parser to use from the chat template."""
+    ctx = build_detection_context(
+        template, tokenizer, reasoning_config, force_reasoning
+    )
+    if ctx is None:
+        return None
+    return match_rules(ctx, TOOL_CALL_PARSER_RULES, "tool-call parser")
+
+
+def detect_inline_system_support(chat_template: Optional[str]) -> bool:
+    """True if mid-conversation ``role: "system"`` renders inline; False if the
+    template raises or silently drops it (then merge into the leading block).
+
+    The probe requires the second system's sentinel to appear in the output —
+    not raising isn't enough, since some templates ignore non-leading system."""
+    if not chat_template:
+        return False
+    sentinel = "__sglang_inline_system_sentinel__"
+    try:
+        env = jinja2.sandbox.ImmutableSandboxedEnvironment(
+            trim_blocks=True,
+            lstrip_blocks=True,
+            extensions=[jinja2.ext.loopcontrols],
+        )
+        rendered = env.from_string(chat_template).render(
+            messages=[
+                {"role": "system", "content": "t"},
+                {"role": "user", "content": "t"},
+                {"role": "system", "content": sentinel},
+                {"role": "user", "content": "t"},
+            ],
+            add_generation_prompt=False,
+        )
+        return sentinel in rendered
+    except jinja2.TemplateError:
+        return False
+    except Exception:
+        return False
+
+
+def _resolve_auto_parser(
+    server_args,
+    attr: str,
+    ctx: TemplateDetectionContext,
+    rules: Tuple[DetectionRule, ...],
+    label: str,
+) -> None:
+    """Resolve a single auto parser, updating server_args in place."""
+    detected = match_rules(ctx, rules, label)
+    if detected:
+        server_args.override(source="template-detection", **{attr: detected})
+        logger.info(
+            f"Auto-detected --{attr.replace('_', '-')} as '{detected}' from chat template"
+        )
+    else:
+        logger.warning(
+            f"--{attr.replace('_', '-')}=auto specified but could not detect "
+            f"{label} from chat template. Disabling {label}."
+        )
+        server_args.override(source="template-detection", **{attr: None})
+
+
+def _load_explicit_jinja_template(chat_template_arg: Optional[str]) -> Optional[str]:
+    if not chat_template_arg or not isinstance(chat_template_arg, str):
+        return None
+    if not chat_template_arg.endswith(".jinja") or not os.path.exists(
+        chat_template_arg
+    ):
+        return None
+    with open(chat_template_arg, encoding="utf-8") as f:
+        return f.read().replace("\\n", "\n")
+
+
+def _disable_auto_parser(server_args, attr: str, label: str) -> None:
+    logger.warning(
+        f"--{attr.replace('_', '-')}=auto specified but could not detect "
+        f"{label} from chat template. Disabling {label}."
+    )
+    server_args.override(source="template-detection", **{attr: None})
+
+
+def _resolve_architecture_auto_parsers(server_args) -> None:
+    from sglang.srt.utils.hf_transformers_utils import get_config
+
+    config = get_config(
+        server_args.model_path,
+        trust_remote_code=server_args.trust_remote_code,
+        revision=getattr(server_args, "revision", None),
+        model_config_parser=getattr(server_args, "model_config_parser", "auto"),
+    )
+    architectures = getattr(config, "architectures", None) or []
+    arch = architectures[0] if architectures else ""
+
+    if "Agnes" in arch:
+        reasoning_parser, tool_call_parser = "deepseek-v4", "deepseekv4"
+    elif "DeepseekV3" in arch:
+        reasoning_parser, tool_call_parser = "deepseek-v3", "deepseekv32"
+    else:
+        return
+
+    for attr, detected in (
+        ("reasoning_parser", reasoning_parser),
+        ("tool_call_parser", tool_call_parser),
+    ):
+        if getattr(server_args, attr) == "auto":
+            server_args.override(source="template-detection", **{attr: detected})
+            logger.info(
+                f"Auto-detected --{attr.replace('_', '-')} as '{detected}' "
+                f"from model architecture '{arch}'"
+            )
+
+
+def resolve_auto_parsers(server_args) -> None:
+    """Resolve --reasoning-parser=auto and --tool-call-parser=auto before scheduler.
+
+    This performs a lightweight tokenizer load to detect parsers from the chat
+    template. Called early in engine init before scheduler subprocesses are spawned.
+    """
+    needs_reasoning = server_args.reasoning_parser == "auto"
+    needs_tool_call = server_args.tool_call_parser == "auto"
+
+    if not needs_reasoning and not needs_tool_call:
+        return
+
+    from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+
+    chat_template_arg = getattr(server_args, "chat_template", None)
+    try:
+        explicit_jinja_template = _load_explicit_jinja_template(chat_template_arg)
+    except Exception as e:
+        logger.warning("Failed to load explicit Jinja chat template: %s", e)
+        explicit_jinja_template = None
+    has_explicit_template_without_detection = (
+        chat_template_arg is not None and explicit_jinja_template is None
+    )
+
+    tokenizer = None
+    try:
+        tokenizer = get_tokenizer(
+            server_args.model_path,
+            trust_remote_code=server_args.trust_remote_code,
+        )
+    except Exception as e:
+        logger.warning(f"Failed to load tokenizer for auto-detection: {e}")
+
+    template = explicit_jinja_template
+    if template is None and tokenizer is not None:
+        template = getattr(tokenizer, "chat_template", None)
+
+    force_reasoning, reasoning_config = detect_reasoning_pattern(template)
+    ctx = build_detection_context(
+        template, tokenizer, reasoning_config, force_reasoning
+    )
+    if ctx is None:
+        if has_explicit_template_without_detection:
+            logger.warning(
+                "--chat-template=%s is explicit but is not a readable Jinja template, so "
+                "parser auto-detection from chat template is not available.",
+                chat_template_arg,
+            )
+        else:
+            try:
+                _resolve_architecture_auto_parsers(server_args)
+            except Exception as e:
+                logger.warning(
+                    "Failed to load model config for architecture-based auto-detection: %s",
+                    e,
+                )
+        if needs_reasoning:
+            if server_args.reasoning_parser == "auto":
+                _disable_auto_parser(
+                    server_args, "reasoning_parser", "reasoning parser"
+                )
+        if needs_tool_call:
+            if server_args.tool_call_parser == "auto":
+                _disable_auto_parser(
+                    server_args, "tool_call_parser", "tool-call parser"
+                )
+        return
+
+    if needs_reasoning:
+        _resolve_auto_parser(
+            server_args,
+            "reasoning_parser",
+            ctx,
+            REASONING_PARSER_RULES,
+            "reasoning parser",
+        )
+
+    if needs_tool_call:
+        _resolve_auto_parser(
+            server_args,
+            "tool_call_parser",
+            ctx,
+            TOOL_CALL_PARSER_RULES,
+            "tool-call parser",
+        )

sglang_patch/srt/utils/hf_transformers/common.py

@@ -0,0 +1,541 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Shared helpers used by config, tokenizer, and processor modules."""
+
+import json
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional, Type, Union
+
+import torch
+from huggingface_hub import snapshot_download
+
+from sglang.srt.configs import (
+    AfmoeConfig,
+    BailingHybridConfig,
+    ChatGLMConfig,
+    DbrxConfig,
+    DeepseekVL2Config,
+    DotsOCRConfig,
+    DotsVLMConfig,
+    ExaoneConfig,
+    FalconH1Config,
+    GraniteMoeHybridConfig,
+    InklingAudioConfig,
+    InklingMMConfig,
+    InklingModelConfig,
+    InklingVisionConfig,
+    InternS2PreviewConfig,
+    JetNemotronConfig,
+    JetVLMConfig,
+    KimiK25Config,
+    KimiLinearConfig,
+    KimiVLConfig,
+    LagunaConfig,
+    LocateAnythingConfig,
+    LongcatFlashConfig,
+    MiniCPMV4_6Config,
+    MiniCPMV4_6VisionConfig,
+    MiniMaxM3VLConfig,
+    MultiModalityConfig,
+    NemotronH_Nano_Omni_Reasoning_V3_Config,
+    NemotronH_Nano_VL_V2_Config,
+    NemotronHConfig,
+    NemotronHPuzzleConfig,
+    Olmo3Config,
+    Qwen3_5Config,
+    Qwen3_5MoeConfig,
+    Qwen3NextConfig,
+    Step3p5Config,
+    Step3p7Config,
+    Step3VLConfig,
+)
+from sglang.srt.configs.deepseek_ocr import DeepseekVLV2Config
+from sglang.srt.configs.internvl import InternVLChatConfig
+from sglang.srt.utils import get_bool_env_var, logger, lru_cache_frozenset
+from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
+
+from ..hf_transformers_patches import normalize_rope_scaling_compat
+
+if get_bool_env_var("SGLANG_USE_MODELSCOPE"):
+    from modelscope import AutoConfig, GenerationConfig
+else:
+    from transformers import AutoConfig, GenerationConfig
+
+from transformers import PretrainedConfig
+
+# ---------------------------------------------------------------------------
+# Config registry
+# ---------------------------------------------------------------------------
+
+_CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
+    cls.model_type: cls
+    for cls in [
+        AfmoeConfig,
+        BailingHybridConfig,
+        ChatGLMConfig,
+        DbrxConfig,
+        ExaoneConfig,
+        DeepseekVL2Config,
+        MultiModalityConfig,
+        KimiVLConfig,
+        LocateAnythingConfig,
+        InternVLChatConfig,
+        LagunaConfig,
+        Step3VLConfig,
+        LongcatFlashConfig,
+        Olmo3Config,
+        KimiLinearConfig,
+        Qwen3NextConfig,
+        FalconH1Config,
+        GraniteMoeHybridConfig,
+        DotsVLMConfig,
+        DotsOCRConfig,
+        NemotronH_Nano_VL_V2_Config,
+        NemotronH_Nano_Omni_Reasoning_V3_Config,
+        NemotronHConfig,
+        NemotronHPuzzleConfig,
+        DeepseekVLV2Config,
+        Qwen3_5Config,
+        Qwen3_5MoeConfig,
+        InternS2PreviewConfig,
+        JetNemotronConfig,
+        JetVLMConfig,
+        KimiK25Config,
+        Step3p5Config,
+        Step3p7Config,
+        MiniCPMV4_6Config,
+        MiniCPMV4_6VisionConfig,
+        InklingModelConfig,
+        InklingAudioConfig,
+        InklingVisionConfig,
+        InklingMMConfig,
+        MiniMaxM3VLConfig,
+    ]
+}
+
+# DeepSeek V3.2 / V4 reuse the V3 config schema. Subclass the upstream
+# transformers class with each model_type so AutoConfig.register passes its
+# consistency check (which requires class.model_type == registered key).
+# Default-value divergences (e.g. V4's topk_group) are handled in
+# model_config.py post-load.
+try:
+    from transformers import DeepseekV3Config as _HFDeepseekV3Config
+
+    class _DeepseekV32ConfigAlias(_HFDeepseekV3Config):
+        model_type = "deepseek_v32"
+
+    class _AgnesConfigAlias(_HFDeepseekV3Config):
+        model_type = "agnes"
+
+    _CONFIG_REGISTRY["deepseek_v32"] = _DeepseekV32ConfigAlias
+    _CONFIG_REGISTRY["agnes"] = _AgnesConfigAlias
+
+    # For kimi_k25_eagle3
+    class _KimiK2ConfigAlias(_HFDeepseekV3Config):
+        model_type = "kimi_k2"
+
+    _CONFIG_REGISTRY["kimi_k2"] = _KimiK2ConfigAlias
+except ImportError:
+    pass
+
+# Newer transformers versions (>=5.10.2) expose MellumConfig directly,
+# but fallback to Qwen3MoeConfig for older versions.
+try:
+    import transformers as _hf_transformers
+
+    _HFMellumConfig = getattr(_hf_transformers, "MellumConfig", None)
+
+    if _HFMellumConfig is not None:
+        _CONFIG_REGISTRY["mellum"] = _HFMellumConfig
+    else:
+        from transformers import Qwen3MoeConfig as _HFQwen3MoeConfig
+
+        class _MellumConfigAlias(_HFQwen3MoeConfig):
+            model_type = "mellum"
+
+            def __post_init__(self, **kwargs):
+                # Qwen3MoeConfig.__post_init__ wipes sliding_window unless
+                # use_sliding_window=True. Mellum gates sliding attention
+                # per-layer via layer_types, so preserve sliding_window
+                # regardless of the legacy use_sliding_window flag.
+                sliding_window = getattr(self, "sliding_window", None)
+                super().__post_init__(**kwargs)
+                self.sliding_window = sliding_window
+
+        _CONFIG_REGISTRY["mellum"] = _MellumConfigAlias
+
+except ImportError:
+    pass
+
+
+try:
+    from transformers import Gemma4Config as _HFGemma4Config
+
+    class _Gemma4UnifiedConfigAlias(_HFGemma4Config):
+        model_type = "gemma4_unified"
+
+    _CONFIG_REGISTRY["gemma4_unified"] = _Gemma4UnifiedConfigAlias
+except ImportError:
+    pass
+
+for name, cls in _CONFIG_REGISTRY.items():
+    try:
+        AutoConfig.register(name, cls)
+    except ValueError as e:
+        err = str(e).lower()
+        if "already registered" not in err and "already used" not in err:
+            logger.warning("Failed to register config %s: %s", name, e)
+
+
+# ---------------------------------------------------------------------------
+# Download / path helpers
+# ---------------------------------------------------------------------------
+
+
+def download_from_hf(
+    model_path: str,
+    allow_patterns: Optional[Union[str, list]] = None,
+):
+    if os.path.exists(model_path):
+        return model_path
+
+    if not allow_patterns:
+        allow_patterns = ["*.json", "*.bin", "*.model"]
+
+    return snapshot_download(model_path, allow_patterns=allow_patterns)
+
+
+def resolve_runai_obj_uri(model_name_or_path: str) -> str:
+    if is_runai_obj_uri(model_name_or_path):
+        return ObjectStorageModel.get_path(model_name_or_path)
+    return model_name_or_path
+
+
+def _resolve_local_or_cached_file(model_name_or_path, filename, revision=None):
+    """Resolve a file from a local directory or HF hub cache (no network)."""
+    local_path = Path(model_name_or_path) / filename
+    if local_path.is_file():
+        return str(local_path)
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(
+        model_name_or_path, filename, revision=revision, local_files_only=True
+    )
+
+
+def _cached_file_exists(model_name_or_path, filename, revision=None) -> bool:
+    """Whether *filename* is available locally or in the HF cache (no network)."""
+    try:
+        _resolve_local_or_cached_file(model_name_or_path, filename, revision)
+        return True
+    except Exception:
+        return False
+
+
+def _remote_file_exists(repo_id, filename, revision=None) -> bool:
+    """Whether *filename* exists on the HF hub (HEAD request only, no download).
+
+    Returns False on any error (offline, gated, network, invalid id) so callers
+    fall back to their default path instead of crashing.
+    """
+    from huggingface_hub.constants import HF_HUB_OFFLINE
+
+    if HF_HUB_OFFLINE:
+        return False
+    try:
+        from huggingface_hub import HfApi
+
+        return HfApi().file_exists(repo_id, filename, revision=revision)
+    except Exception:
+        return False
+
+
+def check_gguf_file(model: Union[str, os.PathLike]) -> bool:
+    model = Path(model)
+    if not model.is_file():
+        return False
+    elif model.suffix == ".gguf":
+        return True
+
+    with open(model, "rb") as f:
+        header = f.read(4)
+    return header == b"GGUF"
+
+
+# ---------------------------------------------------------------------------
+# Rope / text config helpers
+# ---------------------------------------------------------------------------
+
+
+def get_rope_config(config):
+    """Get (rope_theta, rope_params) from config, supporting both v4 and v5.
+
+    Trust-remote-code configs or parent configs passed to sub-models may not
+    have the v5 ``rope_parameters`` property, so we fall back to the v4-style
+    ``config.rope_theta`` / ``config.rope_scaling`` attributes.
+
+    Returns:
+        (rope_theta, rope_params): In v5, rope_params is the full
+        rope_parameters dict (which subsumes rope_scaling and includes
+        rope_theta). In v4, rope_params is the rope_scaling dict or None.
+    """
+    rope_params = getattr(config, "rope_parameters", None)
+    if rope_params is not None:
+        return rope_params["rope_theta"], rope_params
+    return getattr(config, "rope_theta", 10000), getattr(config, "rope_scaling", None)
+
+
+def _patch_text_config(parent_config: PretrainedConfig, text_config):
+    """Synchronize standard attributes between parent config and text sub-config.
+
+    In transformers v5, the "untangle config" refactor removed automatic
+    inheritance of top-level PretrainedConfig attributes (pad_token_id,
+    tie_word_embeddings, etc.) from sub-configs. Downstream code expects
+    these attributes to be present on both configs (some models pass the
+    parent directly to the language model, others pass the text sub-config),
+    so we propagate in both directions when an attribute is missing.
+    (See https://github.com/huggingface/transformers/pull/41541)
+    """
+    _ATTRS_TO_PROPAGATE = [
+        "pad_token_id",
+        "bos_token_id",
+        "eos_token_id",
+        "tie_word_embeddings",
+    ]
+    for attr in _ATTRS_TO_PROPAGATE:
+        parent_has = hasattr(parent_config, attr)
+        text_has = hasattr(text_config, attr)
+        if parent_has and not text_has:
+            setattr(text_config, attr, getattr(parent_config, attr))
+        elif text_has and not parent_has:
+            setattr(parent_config, attr, getattr(text_config, attr))
+    return text_config
+
+
+def get_hf_text_config(config: PretrainedConfig):
+    """Get the "sub" config relevant to llm for multi modal models.
+    No op for pure text models.
+    """
+    if config.architectures is not None:
+        class_name = config.architectures[0]
+        if class_name.startswith("Llava") and class_name.endswith("ForCausalLM"):
+            # We support non-hf version of llava models, so we do not want to
+            # read the wrong values from the unused default text_config.
+            # NOTE(HandH1998): We set `torch_dtype` of config to `torch.float16` for the weights, as
+            # `torch.float16` is default used for image features in `python/sglang/srt/models/llava.py`.
+            setattr(config, "dtype", torch.float16)
+            return config
+
+    text_config = None
+
+    # Some models (e.g. DeepSeek-OCR) store sub-configs as plain dicts.
+    # Convert to PretrainedConfig early so hasattr() checks and asserts work.
+    parent_dtype = getattr(config, "dtype", None)
+    for _attr in ("text_config", "llm_config", "language_config", "thinker_config"):
+        _sub = getattr(config, _attr, None)
+        if isinstance(_sub, dict):
+            _converted = PretrainedConfig(**_sub)
+            if getattr(_converted, "dtype", None) is None and parent_dtype is not None:
+                _converted.dtype = parent_dtype
+            setattr(config, _attr, _converted)
+        elif _sub is not None and parent_dtype is not None:
+            # transformers v5 multimodal configs (e.g. Mistral3Config) carry
+            # `dtype` only on the top-level config, leaving the sub-configs at
+            # None. Without this, _get_and_verify_dtype falls back to float32
+            # and then "auto" downcasts to float16, which overflows the Pixtral
+            # vision tower on real images and produces NaN features.
+            if getattr(_sub, "dtype", None) is None:
+                _sub.dtype = parent_dtype
+
+    # Priority: thinker_config > llm_config > language_config > text_config
+    if hasattr(config, "thinker_config"):
+        # qwen2.5 omni
+        thinker_config = config.thinker_config
+        if hasattr(thinker_config, "text_config"):
+            setattr(
+                thinker_config.text_config,
+                "dtype",
+                getattr(thinker_config, "dtype", None),
+            )
+            text_config = thinker_config.text_config
+        else:
+            text_config = thinker_config
+    elif hasattr(config, "llm_config"):
+        # PointsV1.5 Chat Model
+        assert hasattr(config.llm_config, "num_attention_heads")
+        text_config = config.llm_config
+    elif hasattr(config, "language_config"):
+        text_config = config.language_config
+    elif hasattr(config, "text_config"):
+        # The code operates under the assumption that text_config should have
+        # `num_attention_heads` (among others). Assert here to fail early
+        # if transformers config doesn't align with this assumption.
+        assert hasattr(config.text_config, "num_attention_heads")
+        text_config = config.text_config
+
+    # Ensure rope_scaling dicts have "type" for remote-code compat (v5).
+    normalize_rope_scaling_compat(config)
+
+    if text_config is not None:
+        return _patch_text_config(config, text_config)
+    return config
+
+
+# ---------------------------------------------------------------------------
+# Model-specific helpers
+# ---------------------------------------------------------------------------
+
+
+def _ensure_sub_configs(config: PretrainedConfig, *attr_names: str) -> None:
+    """Convert dict-valued sub-configs to proper AutoConfig objects in-place."""
+    for attr in attr_names:
+        sub = getattr(config, attr, None)
+        if sub is not None and isinstance(sub, dict):
+            setattr(config, attr, AutoConfig.for_model(**sub))
+
+
+def _is_deepseek_ocr_model(config: PretrainedConfig) -> bool:
+    # TODO: Remove this workaround once AutoConfig correctly identifies deepseek-ocr.
+    # Hugging Face's AutoConfig currently misidentifies it as deepseekvl2.
+    auto_map = getattr(config, "auto_map", None) or {}
+    return auto_map.get("AutoModel") == "modeling_deepseekocr.DeepseekOCRForCausalLM"
+
+
+def _is_deepseek_ocr2_model(config: PretrainedConfig) -> bool:
+    auto_map = getattr(config, "auto_map", None) or {}
+    return auto_map.get("AutoModel") == "modeling_deepseekocr2.DeepseekOCR2ForCausalLM"
+
+
+def _override_v_head_dim_if_zero(config: PretrainedConfig, patch: int = 128) -> None:
+    patched = False
+    for attr in ("text_config", "language_config"):
+        sub = getattr(config, attr, None)
+        if sub is None:
+            continue
+        if isinstance(sub, dict):
+            if sub.get("v_head_dim") == 0:
+                sub["v_head_dim"] = patch
+                patched = True
+        elif getattr(sub, "v_head_dim", None) == 0:
+            sub.v_head_dim = patch
+            patched = True
+    if patched:
+        logger.warning(
+            f"Overriding v_head_dim from 0 to {patch} to avoid potential issues."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Context length / generation config / sparse attention
+# ---------------------------------------------------------------------------
+
+# Models don't use the same configuration key for determining the maximum
+# context length.  Store them here so we can sanely check them.
+# NOTE: The ordering here is important. Some models have two of these and we
+# have a preference for which value gets used.
+CONTEXT_LENGTH_KEYS = [
+    "max_sequence_length",
+    "seq_length",
+    "max_seq_len",
+    "model_max_length",
+    "max_position_embeddings",
+]
+
+
+def get_context_length(config):
+    """Get the context length of a model from a huggingface model configs."""
+    text_config = config
+    rope_scaling = getattr(text_config, "rope_scaling", None)
+    if rope_scaling:
+        rope_scaling_factor = rope_scaling.get("factor", 1)
+        if "original_max_position_embeddings" in rope_scaling:
+            rope_scaling_factor = 1
+        if rope_scaling.get("rope_type", None) == "llama3":
+            rope_scaling_factor = 1
+    else:
+        rope_scaling_factor = 1
+
+    for key in CONTEXT_LENGTH_KEYS:
+        val = getattr(text_config, key, None)
+        if val is not None:
+            return int(rope_scaling_factor * val)
+    return 2048
+
+
+@lru_cache_frozenset(maxsize=32)
+def get_generation_config(
+    model: str,
+    trust_remote_code: bool,
+    revision: Optional[str] = None,
+    **kwargs,
+):
+    try:
+        return GenerationConfig.from_pretrained(
+            model, trust_remote_code=trust_remote_code, revision=revision, **kwargs
+        )
+    except FileNotFoundError:
+        return None
+    except OSError as e:
+        logger.warning(
+            "Failed to load generation config for %s: %s. "
+            "Proceeding without generation config.",
+            model,
+            e,
+        )
+        return None
+
+
+# Qwen-1M related
+def get_sparse_attention_config(
+    model: str,
+    sparse_attention_config_filename: str = "sparse_attention_config.json",
+) -> Dict[str, Any]:
+    is_local = os.path.isdir(model)
+    if not is_local:
+        model = download_from_hf(model, allow_patterns=["*.json"])
+
+    config_file = os.path.join(model, sparse_attention_config_filename)
+    if not os.path.exists(config_file):
+        return {}
+
+    with open(config_file) as f:
+        config = json.load(f)
+    return config
+
+
+# ---------------------------------------------------------------------------
+# Tokenizer / processor helpers
+# ---------------------------------------------------------------------------
+
+
+# Some models don't have an available processor, e.g.: InternVL
+def get_tokenizer_from_processor(processor):
+    from transformers import PreTrainedTokenizerBase
+
+    if isinstance(processor, PreTrainedTokenizerBase):
+        return processor
+    return processor.tokenizer
+
+
+# Turn-final markers that some checkpoints ship without EOS metadata:
+# <|eom_id|> (Llama-3 tool use) and <|content_model_end_sampling|> (Inkling,
+# whose bundled tokenizer config leaves eos_token unset).
+_ADDITIONAL_STOP_TOKEN_TEXTS = ("<|eom_id|>", "<|content_model_end_sampling|>")
+
+
+def attach_additional_stop_token_ids(tokenizer):
+    added = tokenizer.get_added_vocab()
+    stop_ids = {added[text] for text in _ADDITIONAL_STOP_TOKEN_TEXTS if text in added}
+    tokenizer.additional_stop_token_ids = stop_ids or None

tokenizer_config.json

@@ -0,0 +1,34 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<｜begin▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "legacy": true,
+  "model_max_length": 1048576,
+  "pad_token": {
+    "__type": "AddedToken",
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sp_model_kwargs": {},
+  "unk_token": null,
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}