Add vLLM acceleration plugin (nas_child_vl_vllm package)

nas_child_vl_vllm/config.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+"""NAS-Child-VL model config for vLLM."""
+
+from transformers import AutoConfig, Qwen3VLConfig
+
+
+class NasChildVLConfig(Qwen3VLConfig):
+    """Config for NAS-Child-VL model, extending Qwen3VLConfig
+    with per-layer NAS configuration."""
+
+    model_type = "nas-child-vl"
+
+    # FLA library defaults for hidden_size=2048
+    # These match the training configuration exactly.
+    DEFAULT_GDN_CONFIG = {
+        "num_heads": 6,
+        "head_k_dim": 256,
+        "head_v_dim": 512,  # expand_v=2 -> head_k_dim * 2
+        "conv_size": 4,
+    }
+    DEFAULT_KDA_CONFIG = {
+        "num_heads": 16,
+        "head_dim": 128,
+        "conv_size": 4,
+    }
+    DEFAULT_MLA_CONFIG = {
+        "num_heads": 16,
+        "q_lora_rank": 64,
+        "qk_rope_head_dim": 64,
+        "kv_lora_rank": 512,
+        "v_head_dim": 128,
+        "qk_nope_head_dim": 128,
+    }
+
+    def __init__(self, nas_layer_config=None, **kwargs):
+        super().__init__(**kwargs)
+        self.nas_layer_config = nas_layer_config or []
+
+    def get_layer_attention_type(self, layer_idx: int) -> str:
+        if layer_idx < len(self.nas_layer_config):
+            return self.nas_layer_config[layer_idx].get(
+                "attention_type", "full_attention"
+            )
+        return "full_attention"
+
+    def get_layer_config(self, layer_idx: int) -> dict:
+        if layer_idx < len(self.nas_layer_config):
+            return self.nas_layer_config[layer_idx]
+        return {"attention_type": "full_attention", "ffn_type": "ffn"}
+
+    def get_gdn_config(self, layer_idx: int) -> dict:
+        """Get GDN config for a layer, falling back to FLA defaults."""
+        lc = self.get_layer_config(layer_idx)
+        explicit = lc.get("gdn_config", {})
+        result = dict(self.DEFAULT_GDN_CONFIG)
+        result.update(explicit)
+        return result
+
+    def get_kda_config(self, layer_idx: int) -> dict:
+        """Get KDA config for a layer, falling back to FLA defaults."""
+        lc = self.get_layer_config(layer_idx)
+        explicit = lc.get("kda_config", {})
+        result = dict(self.DEFAULT_KDA_CONFIG)
+        result.update(explicit)
+        return result
+
+    def get_mla_config(self, layer_idx: int) -> dict:
+        """Get MLA config for a layer, falling back to FLA defaults."""
+        lc = self.get_layer_config(layer_idx)
+        explicit = lc.get("mla_config", {})
+        result = dict(self.DEFAULT_MLA_CONFIG)
+        result.update(explicit)
+        return result
+
+    @property
+    def layers_block_type(self) -> list[str]:
+        """Return block type per layer for get_num_layers_of_type.
+
+        'attention' for FULL/SWA/MLA layers, 'mamba' for GDN/KDA layers.
+        """
+        result = []
+        num_layers = getattr(self, "num_hidden_layers", 0)
+        if hasattr(self, "text_config"):
+            num_layers = getattr(self.text_config, "num_hidden_layers", num_layers)
+        for i in range(num_layers):
+            attn_type = self.get_layer_attention_type(i)
+            if attn_type in ("gdn", "kda"):
+                result.append("mamba")
+            else:
+                result.append("attention")
+        return result
+
+
+AutoConfig.register("nas-child-vl", NasChildVLConfig)

nas_child_vl_vllm/model.py

@@ -0,0 +1,2022 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Inference-only NAS-Child-VL model for vLLM.
+
+NAS-Child-VL is based on Qwen3-VL with per-layer mixed attention:
+  FULL / SWA : Standard GQA via vLLM PagedAttention
+  GDN        : Native GatedDeltaNet (MambaBase, causal_conv1d + chunk ops)
+  KDA        : Native KimiDeltaAttention (MambaBase, causal_conv1d + chunk ops)
+  MLA        : fla wrapper with dummy MambaBase state
+
+GDN/KDA use MambaBase for proper state management via the V1 engine.
+MLA uses fla wrapper with use_cache=True and a dummy MambaBase state.
+FULL/SWA use standard PagedAttention KV cache.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+from itertools import islice
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
+from vllm.distributed import (
+    divide,
+    get_pp_group,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.fla.ops import (
+    RMSNormGated,
+    chunk_gated_delta_rule,
+    fused_recurrent_gated_delta_rule,
+)
+from vllm.model_executor.layers.fla.ops.kda import (
+    FusedRMSNormGated,
+    chunk_kda,
+    fused_kda_gate,
+    fused_recurrent_kda,
+)
+from vllm.model_executor.layers.fla.ops.l2norm import l2norm_fwd
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.mamba.abstract import MambaBase
+from vllm.model_executor.layers.mamba.mamba_utils import (
+    MambaStateCopyFunc,
+    MambaStateCopyFuncCalculator,
+    MambaStateDtypeCalculator,
+    MambaStateShapeCalculator,
+)
+from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
+    causal_conv1d_fn,
+    causal_conv1d_update,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    sharded_weight_loader,
+)
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.sequence import IntermediateTensors
+from nas_child_vl_vllm.config import NasChildVLConfig
+from vllm.utils.torch_utils import direct_register_custom_op
+from vllm.v1.attention.backend import AttentionMetadata
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
+
+from vllm.model_executor.models.interfaces import HasInnerState, SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.qwen3_vl import (
+    Qwen3VLDummyInputsBuilder,
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLMultiModalProcessor,
+    Qwen3VLProcessingInfo,
+    Qwen3_VisionTransformer,
+)
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    WeightsMapper,
+    extract_layer_index,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+)
+
+logger = init_logger(__name__)
+
+# Page size padded for ALL mamba-like layers (GDN, KDA).
+# Must be >= each layer's actual state size AND a common multiple of every
+# attention layer's real_page_size_bytes so that
+# `unify_kv_cache_spec_page_size` can adjust block_size cleanly.
+#   FULL/SWA real_page_size = 16 * 8 * (128+128) * 2 = 65536
+#   MLA      real_page_size = 16 * 16 * (192+192) * 2 = 196608
+# Choose smallest k * 196608 that is >= 3211264 (original mamba sizing) and
+# also divisible by 65536. 196608 = 3*65536 so any multiple of 196608 already
+# divides 65536-multiples evenly.
+#   k = ceil(3211264 / 196608) = 17 -> 17 * 196608 = 3342336
+# 3342336 % 65536  == 0 (51) ✓
+# 3342336 % 196608 == 0 (17) ✓
+_NAS_PAGE_SIZE_PADDED = 3342336
+
+
+def _get_unified_state_shape(tp_size: int):
+    # 🌟 核心修改：将 (dim, 3) 改为 (3, dim)，匹配 vLLM 底层 Cache 的约定
+    return (
+        (3, 6144 // tp_size),   # State 0: GDN conv (6144) / KDA conv_q (2048)
+        (3, 2048 // tp_size),   # State 1: KDA conv_k
+        (3, 2048 // tp_size),   # State 2: KDA conv_v
+        (786432 // tp_size,),   # State 3: GDN ssm (6*256*512=786432) / KDA ssm (16*128*128=262144)
+    )
+
+def _get_unified_state_dtype(model_dtype: torch.dtype):
+    return (model_dtype, model_dtype, model_dtype, torch.float32)
+
+
+
+
+
+
+
+
+# ============================================================================
+# Custom ops for GDN and KDA
+# ============================================================================
+
+def nas_gdn_attention_core(
+    mixed_qkv: torch.Tensor,
+    b: torch.Tensor,
+    a: torch.Tensor,
+    core_attn_out: torch.Tensor,
+    layer_name: str,
+) -> None:
+    forward_context: ForwardContext = get_forward_context()
+    self = forward_context.no_compile_layers[layer_name]
+    self._forward_core(
+        mixed_qkv=mixed_qkv,
+        b=b,
+        a=a,
+        core_attn_out=core_attn_out,
+    )
+
+
+def nas_gdn_attention_core_fake(
+    mixed_qkv: torch.Tensor,
+    b: torch.Tensor,
+    a: torch.Tensor,
+    core_attn_out: torch.Tensor,
+    layer_name: str,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="nas_gdn_attention_core",
+    op_func=nas_gdn_attention_core,
+    mutates_args=["core_attn_out"],
+    fake_impl=nas_gdn_attention_core_fake,
+)
+
+
+def nas_kda_attention_core(
+    q_proj_states: torch.Tensor,
+    k_proj_states: torch.Tensor,
+    v_proj_states: torch.Tensor,
+    g1: torch.Tensor,
+    beta: torch.Tensor,
+    core_attn_out: torch.Tensor,
+    layer_name: str,
+) -> None:
+    forward_context: ForwardContext = get_forward_context()
+    self = forward_context.no_compile_layers[layer_name]
+    self._forward_core(
+        q_proj_states=q_proj_states,
+        k_proj_states=k_proj_states,
+        v_proj_states=v_proj_states,
+        g1=g1,
+        beta=beta,
+        core_attn_out=core_attn_out,
+    )
+
+
+def nas_kda_attention_core_fake(
+    q_proj_states: torch.Tensor,
+    k_proj_states: torch.Tensor,
+    v_proj_states: torch.Tensor,
+    g1: torch.Tensor,
+    beta: torch.Tensor,
+    core_attn_out: torch.Tensor,
+    layer_name: str,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="nas_kda_attention_core",
+    op_func=nas_kda_attention_core,
+    mutates_args=["core_attn_out"],
+    fake_impl=nas_kda_attention_core_fake,
+)
+
+
+# ============================================================================
+# FULL / SWA attention - vLLM PagedAttention
+# ============================================================================
+
+class NasAttention(nn.Module):
+    """Standard Qwen3-style GQA attention for FULL and SWA layers."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        rope_parameters: dict,
+        max_position: int,
+        rms_norm_eps: float,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        sliding_window: int | None = None,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = head_dim
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim ** -0.5
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=max_position,
+            rope_parameters=rope_parameters,
+        )
+
+        from vllm.model_executor.layers.attention import Attention
+        from vllm.v1.attention.backend import AttentionType
+        attn_kwargs: dict = dict(
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+            attn_type=AttentionType.DECODER,
+        )
+        if sliding_window is not None:
+            attn_kwargs["per_layer_sliding_window"] = sliding_window
+        self.attn = Attention(
+            self.num_heads, self.head_dim, self.scaling, **attn_kwargs
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        output: torch.Tensor,
+    ) -> None:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q = self.q_norm(q.reshape(-1, self.num_heads, self.head_dim))
+        k = self.k_norm(k.reshape(-1, self.num_kv_heads, self.head_dim))
+        q = q.reshape(-1, self.num_heads * self.head_dim)
+        k = k.reshape(-1, self.num_kv_heads * self.head_dim)
+        q, k = self.rotary_emb(positions, q, k)
+        attn_output = self.attn(q, k, v)
+        output[:], _ = self.o_proj(attn_output)
+
+
+# ============================================================================
+# GDN attention - Native implementation with MambaBase
+# ============================================================================
+
+class NasGDNAttention(nn.Module, MambaBase):
+    """GatedDeltaNet with native causal_conv1d + chunk ops and MambaBase state."""
+
+    @property
+    def mamba_type(self) -> str:
+        return "gdn_attention"  # 🌟 统一命名
+
+    def get_state_dtype(self) -> tuple[torch.dtype, ...]:
+        return _get_unified_state_dtype(self.model_config.dtype)
+
+    def get_state_shape(self) -> tuple[tuple[int, ...], ...]:
+        return _get_unified_state_shape(self.tp_size)
+
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig):
+        spec = super().get_kv_cache_spec(vllm_config)
+        from dataclasses import replace
+        
+        # 🌟 核心修改：如果 MambaBase 返回的 block_size 是 None，则赋予默认的 block_size
+        block_size = spec.block_size if spec.block_size is not None else vllm_config.cache_config.block_size
+        
+        return replace(spec, page_size_padded=_NAS_PAGE_SIZE_PADDED, block_size=block_size)
+
+    def __init__(
+        self,
+        hidden_size: int,
+        layer_idx: int,
+        cache_config: CacheConfig | None = None,
+        model_config: ModelConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        **kwargs,
+    ):
+        super().__init__()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_rank = get_tensor_model_parallel_rank()
+        self.hidden_size = hidden_size
+        self.model_config = model_config
+        self.cache_config = cache_config
+        self.layer_idx = layer_idx
+        self.prefix = prefix
+
+        # GDN config: num_heads=6, head_k_dim=256, head_v_dim=512, conv_size=4
+        self.num_heads = 6
+        self.head_k_dim = 256
+        self.head_v_dim = 512  # expand_v=2
+        self.conv_kernel_size = 4
+
+        self.key_dim = self.head_k_dim * self.num_heads
+        self.value_dim = self.head_v_dim * self.num_heads
+        self.conv_dim = self.key_dim * 2 + self.value_dim
+
+        # Merged conv1d for Q+K+V
+        self.conv1d = ColumnParallelLinear(
+            input_size=self.conv_kernel_size,
+            output_size=self.conv_dim,
+            bias=False,
+            prefix=f"{prefix}.conv1d",
+        )
+        self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
+
+        # Input projection: Q, K, V, Z (gate)
+        self.in_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[self.key_dim, self.key_dim, self.value_dim, self.value_dim],
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.in_proj",
+        )
+
+        # Beta and alpha projections
+        self.in_proj_ba = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[self.num_heads, self.num_heads],
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.in_proj_ba",
+        )
+
+        # A_log and dt_bias
+        self.A_log = nn.Parameter(
+            torch.empty(divide(self.num_heads, self.tp_size), dtype=torch.float32)
+        )
+        set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
+
+        self.dt_bias = nn.Parameter(
+            torch.empty(divide(self.num_heads, self.tp_size), dtype=torch.float32)
+        )
+        set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
+
+        # Output norm and projection
+        # Match fla GatedDeltaNet: FusedRMSNormGated(head_v_dim, activation="swish")
+        # which computes rms_norm(x) * weight * silu(z) (norm_before_gate=True)
+        self.norm = RMSNormGated(
+            self.head_v_dim,
+            eps=1e-6,
+            norm_before_gate=True,
+            activation="swish",
+        )
+        self.out_proj = RowParallelLinear(
+            self.value_dim,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.out_proj",
+        )
+
+        compilation_config = get_current_vllm_config().compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+
+    def forward(self, positions, hidden_states, output):
+        num_tokens = hidden_states.size(0)
+
+        # Input projection
+        projected, _ = self.in_proj(hidden_states)
+        q_proj, k_proj, v_proj, z = projected.split(
+            [self.key_dim // self.tp_size,
+             self.key_dim // self.tp_size,
+             self.value_dim // self.tp_size,
+             self.value_dim // self.tp_size],
+            dim=-1,
+        )
+        mixed_qkv = torch.cat([q_proj, k_proj, v_proj], dim=-1)
+
+        ba, _ = self.in_proj_ba(hidden_states)
+        b, a = ba.split([self.num_heads // self.tp_size,
+                         self.num_heads // self.tp_size], dim=-1)
+
+        # Core attention via custom op
+        core_attn_out = torch.zeros(
+            (num_tokens, self.num_heads // self.tp_size, self.head_v_dim),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+
+        torch.ops.vllm.nas_gdn_attention_core(
+            mixed_qkv,
+            b,
+            a,
+            core_attn_out,
+            self.prefix,
+        )
+
+        # Output: norm(out, gate) -> o_proj
+        z_reshaped = z.reshape(-1, self.head_v_dim)
+        core_attn_out = core_attn_out.reshape(-1, self.head_v_dim)
+        core_attn_out = self.norm(core_attn_out, z_reshaped)
+        core_attn_out = core_attn_out.reshape(num_tokens, -1)
+        output[:], _ = self.out_proj(core_attn_out)
+
+    def _forward_core(
+        self,
+        mixed_qkv: torch.Tensor,
+        b: torch.Tensor,
+        a: torch.Tensor,
+        core_attn_out: torch.Tensor,
+    ):
+        forward_context = get_forward_context()
+        attn_metadata: AttentionMetadata = forward_context.attn_metadata
+
+        if attn_metadata is None:
+            # V1 profile run
+            return
+
+        assert isinstance(attn_metadata, dict)
+        attn_metadata = attn_metadata[self.prefix]
+        assert isinstance(attn_metadata, GDNAttentionMetadata)
+        has_initial_state = attn_metadata.has_initial_state
+        non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
+        non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor
+        num_actual_tokens = attn_metadata.num_actual_tokens
+        self_kv_cache = self.kv_cache[forward_context.virtual_engine]
+        conv_state = self_kv_cache[0].transpose(-1, -2)
+        # ssm_state = self_kv_cache[1]
+        ssm_state = self_kv_cache[3].view(
+            -1, 
+            self.num_heads // self.tp_size, 
+            self.head_v_dim,   # 先 V (512)
+            self.head_k_dim    # 后 K (256)
+        )
+
+
+
+        mixed_qkv = mixed_qkv[:num_actual_tokens]
+        b = b[:num_actual_tokens]
+        a = a[:num_actual_tokens]
+
+        # 1. Convolution
+        conv_weights = self.conv1d.weight.view(
+            self.conv1d.weight.size(0), self.conv1d.weight.size(2)
+        )
+
+        if attn_metadata.num_prefills > 0:
+            mixed_qkv_T = mixed_qkv.transpose(0, 1)
+            mixed_qkv = causal_conv1d_fn(
+                mixed_qkv_T,
+                conv_weights,
+                self.conv1d.bias,
+                activation="silu",
+                conv_states=conv_state,
+                has_initial_state=has_initial_state,
+                cache_indices=non_spec_state_indices_tensor,
+                query_start_loc=non_spec_query_start_loc,
+                metadata=attn_metadata,
+            ).transpose(0, 1)
+        else:
+            decode_conv_indices = non_spec_state_indices_tensor[:num_actual_tokens]
+            mixed_qkv = causal_conv1d_update(
+                mixed_qkv,
+                conv_state,
+                conv_weights,
+                self.conv1d.bias,
+                activation="silu",
+                conv_state_indices=decode_conv_indices,
+                validate_data=True,
+            )
+
+        # Split Q, K, V
+        query, key, value = torch.split(
+            mixed_qkv,
+            [self.key_dim // self.tp_size,
+             self.key_dim // self.tp_size,
+             self.value_dim // self.tp_size],
+            dim=-1,
+        )
+        query = rearrange(query, "l (h d) -> 1 l h d", d=self.head_k_dim)
+        key = rearrange(key, "l (h d) -> 1 l h d", d=self.head_k_dim)
+        value = rearrange(value, "l (h d) -> 1 l h d", d=self.head_v_dim)
+
+        # 2. Compute gating: g = -exp(A_log) * softplus(a + dt_bias), beta = sigmoid(b)
+        num_heads_local = self.num_heads // self.tp_size
+        # g: (1, L, H) -> expand to head_k_dim
+        a_float = a.float()
+        dt = a_float + self.dt_bias.unsqueeze(0)
+        softplus_dt = torch.where(
+            dt > 20.0, dt, torch.log1p(torch.exp(dt))
+        )
+        g = -torch.exp(self.A_log.float()).unsqueeze(0) * softplus_dt
+        g = g.unsqueeze(0)  # (1, L, H)
+
+        beta = b.float().sigmoid()
+        beta = beta.unsqueeze(0)  # (1, L, H)
+
+        # 3. Recurrent attention
+        if attn_metadata.num_prefills > 0:
+            zero_idx = non_spec_state_indices_tensor[~has_initial_state]
+            ssm_state[zero_idx] = 0
+            initial_state = ssm_state[non_spec_state_indices_tensor].contiguous()
+            (
+                core_attn_out_non_spec,
+                last_recurrent_state,
+            ) = chunk_gated_delta_rule(
+                q=query,
+                k=key,
+                v=value,
+                g=g,
+                beta=beta,
+                initial_state=initial_state,
+                output_final_state=True,
+                cu_seqlens=non_spec_query_start_loc,
+                use_qk_l2norm_in_kernel=True,
+            )
+            ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to(
+                ssm_state.dtype
+            )
+        else:
+            (
+                core_attn_out_non_spec,
+                last_recurrent_state,
+            ) = fused_recurrent_gated_delta_rule(
+                q=query,
+                k=key,
+                v=value,
+                g=g,
+                beta=beta,
+                initial_state=ssm_state,
+                inplace_final_state=True,
+                cu_seqlens=non_spec_query_start_loc[:attn_metadata.num_decodes + 1],
+                ssm_state_indices=non_spec_state_indices_tensor,
+                use_qk_l2norm_in_kernel=True,
+            )
+
+        core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)[:num_actual_tokens]
+
+
+# ============================================================================
+# KDA attention - Native implementation with MambaBase
+# ============================================================================
+
+class NasKDAAttention(nn.Module, MambaBase):
+    """KimiDeltaAttention with native causal_conv1d + chunk ops and MambaBase state."""
+
+    @property
+    def mamba_type(self) -> str:
+        return "gdn_attention"  # 🌟 统一命名
+
+    def get_state_dtype(self) -> tuple[torch.dtype, ...]:
+        return _get_unified_state_dtype(self.model_config.dtype)
+
+    def get_state_shape(self) -> tuple[tuple[int, ...], ...]:
+        return _get_unified_state_shape(self.tp_size)
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig):
+        spec = super().get_kv_cache_spec(vllm_config)
+        from dataclasses import replace
+        
+        # 🌟 核心修改：如果 MambaBase 返回的 block_size 是 None，则赋予默认的 block_size
+        block_size = spec.block_size if spec.block_size is not None else vllm_config.cache_config.block_size
+        
+        return replace(spec, page_size_padded=_NAS_PAGE_SIZE_PADDED, block_size=block_size)
+
+    def __init__(
+        self,
+        hidden_size: int,
+        layer_idx: int,
+        cache_config: CacheConfig | None = None,
+        model_config: ModelConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        **kwargs,
+    ):
+        super().__init__()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_rank = get_tensor_model_parallel_rank()
+        self.hidden_size = hidden_size
+        self.model_config = model_config
+        self.cache_config = cache_config
+        self.layer_idx = layer_idx
+        self.prefix = prefix
+
+        # KDA config: num_heads=16, head_dim=128, conv_size=4
+        self.num_heads = 16
+        self.head_dim = 128
+        self.conv_size = 4
+        assert self.num_heads % self.tp_size == 0
+        self.local_num_heads = divide(self.num_heads, self.tp_size)
+
+        projection_size = self.head_dim * self.num_heads
+
+        # Q, K, V projections
+        self.q_proj = ColumnParallelLinear(
+            self.hidden_size,
+            projection_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.q_proj",
+        )
+        self.k_proj = ColumnParallelLinear(
+            self.hidden_size,
+            projection_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.k_proj",
+        )
+        self.v_proj = ColumnParallelLinear(
+            self.hidden_size,
+            projection_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.v_proj",
+        )
+
+        # Convolutions for Q, K, V
+        self.q_conv1d = ColumnParallelLinear(
+            input_size=self.conv_size,
+            output_size=projection_size,
+            bias=False,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.q_conv1d",
+        )
+        self.k_conv1d = ColumnParallelLinear(
+            input_size=self.conv_size,
+            output_size=projection_size,
+            bias=False,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.k_conv1d",
+        )
+        self.v_conv1d = ColumnParallelLinear(
+            input_size=self.conv_size,
+            output_size=projection_size,
+            bias=False,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.v_conv1d",
+        )
+        # unsqueeze to fit conv1d weights shape
+        self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1)
+        self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1)
+        self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1)
+
+        # f_gate: f_proj.0 (ReplicatedLinear) + f_proj.1 (ColumnParallelLinear)
+        self.f_proj_0 = ReplicatedLinear(
+            self.hidden_size,
+            self.head_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.f_proj.0",
+        )
+        self.f_proj_1 = ColumnParallelLinear(
+            self.head_dim,
+            projection_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.f_proj.1",
+        )
+
+        # dt_bias
+        self.dt_bias = nn.Parameter(
+            torch.empty(divide(projection_size, self.tp_size), dtype=torch.float32)
+        )
+        set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
+
+        # A_log
+        self.A_log = nn.Parameter(
+            torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32)
+        )
+        set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(2)})
+
+        # beta projection
+        self.b_proj = ColumnParallelLinear(
+            self.hidden_size,
+            self.num_heads,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.b_proj",
+        )
+
+        # output_gate: g_proj.0 (ReplicatedLinear) + g_proj.1 (ColumnParallelLinear, bias=True)
+        self.g_proj_0 = ReplicatedLinear(
+            self.hidden_size,
+            self.head_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.g_proj.0",
+        )
+        self.g_proj_1 = ColumnParallelLinear(
+            self.head_dim,
+            projection_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.g_proj.1",
+        )
+
+        # Output norm and projection
+        self.o_norm = FusedRMSNormGated(
+            self.head_dim, eps=1e-6, activation="sigmoid"
+        )
+        self.o_proj = RowParallelLinear(
+            projection_size,
+            self.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        compilation_config = get_current_vllm_config().compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+
+    def forward(self, positions, hidden_states, output):
+        num_tokens = hidden_states.size(0)
+
+        # Project Q, K, V
+        q = self.q_proj(hidden_states)[0]
+        k = self.k_proj(hidden_states)[0]
+        v = self.v_proj(hidden_states)[0]
+
+        # Beta
+        beta = self.b_proj(hidden_states)[0].float().sigmoid()
+
+        # f_gate
+        g1 = self.f_proj_1(self.f_proj_0(hidden_states)[0])[0]
+        g1 = fused_kda_gate(g1, self.A_log, self.head_dim, g_bias=self.dt_bias)
+        beta = beta.unsqueeze(0)
+        g1 = g1.unsqueeze(0)
+
+        # Output gate
+        g_proj_states = self.g_proj_1(self.g_proj_0(hidden_states)[0])[0]
+        g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
+
+        # Core attention via custom op
+        core_attn_out = torch.zeros(
+            (1, num_tokens, self.local_num_heads, self.head_dim),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+        torch.ops.vllm.nas_kda_attention_core(
+            q,
+            k,
+            v,
+            g1,
+            beta,
+            core_attn_out,
+            self.prefix,
+        )
+
+        # Output: o_norm(core_out, output_gate) -> o_proj
+        core_attn_out = self.o_norm(core_attn_out, g2)
+        core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
+        output[:] = self.o_proj(core_attn_out)[0]
+
+    def _forward_core(
+        self,
+        q_proj_states: torch.Tensor,
+        k_proj_states: torch.Tensor,
+        v_proj_states: torch.Tensor,
+        g1: torch.Tensor,
+        beta: torch.Tensor,
+        core_attn_out: torch.Tensor,
+    ) -> None:
+        forward_context = get_forward_context()
+        attn_metadata: AttentionMetadata = forward_context.attn_metadata
+
+        if attn_metadata is None:
+            # V1 profile run
+            return
+
+        assert isinstance(attn_metadata, dict)
+        attn_metadata = attn_metadata[self.prefix]
+        assert isinstance(attn_metadata, GDNAttentionMetadata)
+        has_initial_state = attn_metadata.has_initial_state
+        non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
+        non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor
+        num_actual_tokens = attn_metadata.num_actual_tokens
+
+        constant_caches = self.kv_cache[forward_context.virtual_engine]
+
+        # 🌟 核心修改：因为 6144 现在在最后一维，所以切片要放在最后一维 [..., :, :2048]
+        conv_state_q = constant_caches[0][..., :, :2048 // self.tp_size].transpose(-1, -2)
+        conv_state_k = constant_caches[1].transpose(-1, -2)
+        conv_state_v = constant_caches[2].transpose(-1, -2)
+        
+        # recurrent_state 只需要 262144，切片后 reshape
+        recurrent_state = constant_caches[3][..., :262144 // self.tp_size].view(
+            -1, self.local_num_heads, self.head_dim, self.head_dim
+        )
+
+        q_proj_states = q_proj_states[:num_actual_tokens]
+        k_proj_states = k_proj_states[:num_actual_tokens]
+        v_proj_states = v_proj_states[:num_actual_tokens]
+        g1 = g1[:, :num_actual_tokens]
+        beta = beta[:, :num_actual_tokens]
+
+        # (conv_state_q, conv_state_k, conv_state_v, recurrent_state) = constant_caches
+
+
+        q_conv_weights = self.q_conv1d.weight.view(
+            self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2)
+        )
+        k_conv_weights = self.k_conv1d.weight.view(
+            self.k_conv1d.weight.size(0), self.k_conv1d.weight.size(2)
+        )
+        v_conv_weights = self.v_conv1d.weight.view(
+            self.v_conv1d.weight.size(0), self.v_conv1d.weight.size(2)
+        )
+
+        if attn_metadata.num_prefills > 0:
+            q_proj_states = q_proj_states.transpose(0, 1)
+            k_proj_states = k_proj_states.transpose(0, 1)
+            v_proj_states = v_proj_states.transpose(0, 1)
+            q = causal_conv1d_fn(
+                q_proj_states,
+                q_conv_weights,
+                self.q_conv1d.bias,
+                activation="silu",
+                conv_states=conv_state_q,
+                has_initial_state=has_initial_state,
+                cache_indices=non_spec_state_indices_tensor,
+                query_start_loc=non_spec_query_start_loc,
+                metadata=attn_metadata,
+            ).transpose(0, 1)
+            k = causal_conv1d_fn(
+                k_proj_states,
+                k_conv_weights,
+                self.k_conv1d.bias,
+                activation="silu",
+                conv_states=conv_state_k,
+                has_initial_state=has_initial_state,
+                cache_indices=non_spec_state_indices_tensor,
+                query_start_loc=non_spec_query_start_loc,
+                metadata=attn_metadata,
+            ).transpose(0, 1)
+            v = causal_conv1d_fn(
+                v_proj_states,
+                v_conv_weights,
+                self.v_conv1d.bias,
+                activation="silu",
+                conv_states=conv_state_v,
+                has_initial_state=has_initial_state,
+                cache_indices=non_spec_state_indices_tensor,
+                query_start_loc=non_spec_query_start_loc,
+                metadata=attn_metadata,
+            ).transpose(0, 1)
+        else:
+            decode_conv_indices = non_spec_state_indices_tensor[:num_actual_tokens]
+            q = causal_conv1d_update(
+                q_proj_states,
+                conv_state_q,
+                q_conv_weights,
+                self.q_conv1d.bias,
+                activation="silu",
+                conv_state_indices=decode_conv_indices,
+                validate_data=True,
+            )
+            k = causal_conv1d_update(
+                k_proj_states,
+                conv_state_k,
+                k_conv_weights,
+                self.k_conv1d.bias,
+                activation="silu",
+                conv_state_indices=decode_conv_indices,
+                validate_data=True,
+            )
+            v = causal_conv1d_update(
+                v_proj_states,
+                conv_state_v,
+                v_conv_weights,
+                self.v_conv1d.bias,
+                activation="silu",
+                conv_state_indices=decode_conv_indices,
+                validate_data=True,
+            )
+
+        q, k, v = map(
+            lambda x: rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim), (q, k, v)
+        )
+
+        if attn_metadata.num_prefills > 0:
+            zero_idx = non_spec_state_indices_tensor[~has_initial_state]
+            recurrent_state[zero_idx] = 0
+            initial_state = recurrent_state[non_spec_state_indices_tensor].contiguous()
+            (
+                core_attn_out_non_spec,
+                last_recurrent_state,
+            ) = chunk_kda(
+                q=q,
+                k=k,
+                v=v,
+                g=g1,
+                beta=beta,
+                initial_state=initial_state,
+                output_final_state=True,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=non_spec_query_start_loc,
+            )
+            recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state
+        else:
+            (
+                core_attn_out_non_spec,
+                last_recurrent_state,
+            ) = fused_recurrent_kda(
+                q=q,
+                k=k,
+                v=v,
+                g=g1,
+                beta=beta,
+                initial_state=recurrent_state,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=non_spec_query_start_loc[:attn_metadata.num_decodes + 1],
+                ssm_state_indices=non_spec_state_indices_tensor,
+            )
+        core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[0, :num_actual_tokens]
+
+
+# ============================================================================
+# MLA attention - Native vLLM PagedAttention (DeepSeek V2 pattern)
+# ============================================================================
+
+# class NasMLAAttention(nn.Module):
+#     """MultiheadLatentAttention with LoRA-compressed Q/KV projections + PagedAttention.
+
+#     Architecture (matching fla MLA exactly):
+#     - Q: hidden → q_lora_rank → RMSNorm → num_heads * qk_head_dim
+#     - K_rope: hidden → qk_rope_head_dim (shared across heads)
+#     - KV: hidden → kv_lora_rank → RMSNorm → num_heads * (qk_nope_head_dim + v_head_dim)
+#     - Attention: standard causal attention via vLLM PagedAttention
+#     - Output: num_heads * v_head_dim → hidden
+#     """
+
+#     def __init__(
+#         self,
+#         hidden_size: int,
+#         layer_idx: int,
+#         cache_config: CacheConfig | None = None,
+#         model_config: ModelConfig | None = None,
+#         quant_config: QuantizationConfig | None = None,
+#         prefix: str = "",
+#         **kwargs,
+#     ):
+#         super().__init__()
+#         self.hidden_size = hidden_size
+#         tp_size = get_tensor_model_parallel_world_size()
+
+#         # MLA dimensions
+#         self.num_heads = 16
+#         self.q_lora_rank = 64
+#         self.qk_rope_head_dim = 64
+#         self.kv_lora_rank = 512
+#         self.v_head_dim = 128
+#         self.qk_nope_head_dim = 128
+#         self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim  # 192
+
+#         assert self.num_heads % tp_size == 0
+#         self.num_local_heads = self.num_heads // tp_size
+#         self.scaling = self.qk_head_dim ** -0.5
+
+#         # Q projection: hidden → q_lora_rank → RMSNorm → num_heads * qk_head_dim
+#         self.q_a_proj = ReplicatedLinear(
+#             self.hidden_size,
+#             self.q_lora_rank,
+#             bias=False,
+#             quant_config=quant_config,
+#             prefix=f"{prefix}.q_a_proj",
+#         )
+#         self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=1e-6)
+#         self.q_b_proj = ColumnParallelLinear(
+#             self.q_lora_rank,
+#             self.num_heads * self.qk_head_dim,
+#             bias=False,
+#             quant_config=quant_config,
+#             prefix=f"{prefix}.q_b_proj",
+#         )
+
+#         # KV projection: hidden → kv_lora_rank → RMSNorm → num_heads * (nope + v)
+#         self.kv_a_proj = ReplicatedLinear(
+#             self.hidden_size,
+#             self.kv_lora_rank,
+#             bias=False,
+#             quant_config=quant_config,
+#             prefix=f"{prefix}.kv_a_proj",
+#         )
+#         self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=1e-6)
+#         self.kv_b_proj = ColumnParallelLinear(
+#             self.kv_lora_rank,
+#             self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
+#             bias=False,
+#             quant_config=quant_config,
+#             prefix=f"{prefix}.kv_b_proj",
+#         )
+
+#         # K rope projection (shared across heads, single "head")
+#         self.k_rope_proj = ReplicatedLinear(
+#             self.hidden_size,
+#             self.qk_rope_head_dim,
+#             bias=False,
+#             quant_config=quant_config,
+#             prefix=f"{prefix}.k_rope_proj",
+#         )
+
+#         # Output projection
+#         self.o_proj = RowParallelLinear(
+#             self.num_heads * self.v_head_dim,
+#             self.hidden_size,
+#             bias=False,
+#             quant_config=quant_config,
+#             prefix=f"{prefix}.o_proj",
+#         )
+
+#         # Standard 1D RoPE (not mrope) for the rope portion
+#         self.rotary_emb = get_rope(
+#             self.qk_rope_head_dim,
+#             max_position=262144,
+#             rope_parameters={"rope_type": "default"},
+#             is_neox_style=False,
+#         )
+
+#         # PagedAttention
+#         from vllm.model_executor.layers.attention import Attention
+#         from vllm.v1.attention.backend import AttentionType
+#         self.attn = Attention(
+#             self.num_local_heads,
+#             self.qk_head_dim,
+#             self.scaling,
+#             num_kv_heads=self.num_local_heads,
+#             cache_config=cache_config,
+#             quant_config=quant_config,
+#             prefix=f"{prefix}.attn",
+#             attn_type=AttentionType.DECODER,
+#         )
+
+#     def forward(self, positions, hidden_states, output):
+#         # Q: hidden → compress → norm → expand
+#         q = self.q_a_proj(hidden_states)[0]
+#         q = self.q_a_layernorm(q)
+#         q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+#         q_nope, q_pe = q.split(
+#             [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
+#         )
+
+#         # KV: hidden → compress → norm → expand → split(k_nope, v)
+#         kv = self.kv_a_proj(hidden_states)[0]
+#         kv = self.kv_a_layernorm(kv)
+#         kv = self.kv_b_proj(kv)[0].view(
+#             -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
+#         )
+#         k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+#         # K rope (shared across heads)
+#         k_pe = self.k_rope_proj(hidden_states)[0]  # [N, 64]
+#         k_pe = k_pe.unsqueeze(1)  # [N, 1, 64]
+
+#         # Apply RoPE (need 1D positions; Qwen3-VL passes [3, N] for mrope)
+#         pos_1d = positions[0] if positions.ndim == 2 else positions
+#         q_pe, k_pe = self.rotary_emb(pos_1d, q_pe, k_pe)
+
+#         # Assemble full Q and K
+#         q[..., self.qk_nope_head_dim:] = q_pe
+#         k = torch.empty_like(q)
+#         k[..., :self.qk_nope_head_dim] = k_nope
+#         k[..., self.qk_nope_head_dim:] = k_pe  # broadcasts [N,1,64] → [N,H,64]
+
+#         # Pad V to qk_head_dim for PagedAttention KV cache alignment
+#         v = F.pad(v, [0, self.qk_head_dim - self.v_head_dim], value=0)
+#         v = v.view(-1, self.num_local_heads * self.qk_head_dim)
+
+#         # PagedAttention (q is [N, H, D], k is [N, H, D], v is [N, H*D])
+#         attn_output = self.attn(q, k, v)
+
+#         # Strip V padding and project output
+#         attn_output = attn_output.view(
+#             -1, self.num_local_heads, self.qk_head_dim
+#         )[..., :self.v_head_dim].reshape(-1, self.num_local_heads * self.v_head_dim)
+#         output[:], _ = self.o_proj(attn_output)
+
+class NasMLAAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        layer_idx: int,
+        cache_config: CacheConfig | None = None,
+        model_config: ModelConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        rope_parameters: dict | None = None,  
+        **kwargs,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+
+        # 🌟 从 model_config.hf_config 动态获取参数，提供默认值以防 config 中没有
+        hf_config = model_config.hf_config
+        self.num_heads = num_heads
+        self.q_lora_rank = getattr(hf_config, "q_lora_rank", 64)
+        self.qk_rope_head_dim = getattr(hf_config, "qk_rope_head_dim", 64)
+        self.kv_lora_rank = getattr(hf_config, "kv_lora_rank", 512)
+        self.v_head_dim = getattr(hf_config, "v_head_dim", 128)
+        self.qk_nope_head_dim = getattr(hf_config, "qk_nope_head_dim", 128)
+        
+        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+
+        assert self.num_heads % tp_size == 0
+        self.num_local_heads = self.num_heads // tp_size
+        self.scaling = self.qk_head_dim ** -0.5
+
+        # 1. 严格对齐 q_proj (Sequential: Linear -> RMSNorm -> Linear)
+        if self.q_lora_rank is not None:
+            self.q_proj = nn.ModuleList([
+                ReplicatedLinear(self.hidden_size, self.q_lora_rank, bias=False, quant_config=quant_config, prefix=f"{prefix}.q_proj.0"),
+                RMSNorm(self.q_lora_rank, eps=1e-6),
+                ColumnParallelLinear(self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False, quant_config=quant_config, prefix=f"{prefix}.q_proj.2")
+            ])
+        else:
+            self.q_proj = ColumnParallelLinear(self.hidden_size, self.num_heads * self.qk_head_dim, bias=False, quant_config=quant_config, prefix=f"{prefix}.q_proj")
+
+        # 2. 🌟 严格对齐 kv_proj (Sequential: Linear -> RMSNorm -> Linear)
+        # 这将完美匹配 Checkpoint 中的 kv_proj.0, kv_proj.1, kv_proj.2
+        self.kv_proj = nn.ModuleList([
+            ReplicatedLinear(self.hidden_size, self.kv_lora_rank, bias=False, quant_config=quant_config, prefix=f"{prefix}.kv_proj.0"),
+            RMSNorm(self.kv_lora_rank, eps=1e-6),
+            ColumnParallelLinear(self.kv_lora_rank, self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False, quant_config=quant_config, prefix=f"{prefix}.kv_proj.2")
+        ])
+
+        # 3. 严格对齐 k_rope
+        self.k_rope = ReplicatedLinear(
+            self.hidden_size,
+            self.qk_rope_head_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.k_rope"
+        )
+
+        # 4. 严格对齐 o_proj
+        self.o_proj = RowParallelLinear(
+            self.num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj"
+        )
+        import copy
+        mla_rope_params = copy.deepcopy(rope_parameters) if rope_parameters else {}
+        
+        # 🌟 核心修改：只做纯粹的数学缩放，绝不修改 type 或 rope_type 字段
+        if "mrope_section" in mla_rope_params:
+            orig_sections = mla_rope_params["mrope_section"]
+            target_sum = self.qk_rope_head_dim // 2  # 对于 64 维，这里是 32
+            orig_sum = sum(orig_sections)
+            
+            if orig_sum > 0 and orig_sum != target_sum:
+                # 按比例缩放 (例如 64 -> 32，比例为 0.5)
+                scale = target_sum / orig_sum
+                new_sections = [int(s * scale) for s in orig_sections]
+                # 修正由于向下取整可能导致的微小误差，把差值补到最后一个维度
+                if sum(new_sections) != target_sum:
+                    new_sections[-1] += target_sum - sum(new_sections)
+                
+                mla_rope_params["mrope_section"] = new_sections
+
+        self.rotary_emb = get_rope(
+            self.qk_rope_head_dim,
+            max_position=262144,
+            rope_parameters=mla_rope_params,  # 🌟 使用缩放后的参数
+            is_neox_style=False,
+        )
+
+        from vllm.model_executor.layers.attention import Attention
+        from vllm.v1.attention.backend import AttentionType
+        self.attn = Attention(
+            self.num_local_heads,
+            self.qk_head_dim,
+            self.scaling,
+            num_kv_heads=self.num_local_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+            attn_type=AttentionType.DECODER,
+        )
+
+    def forward(self, positions, hidden_states, output):
+        # Q
+        if isinstance(self.q_proj, nn.ModuleList):
+            q = self.q_proj[0](hidden_states)[0]
+            q = self.q_proj[1](q)
+            q = self.q_proj[2](q)[0]
+        else:
+            q = self.q_proj(hidden_states)[0]
+        q = q.view(-1, self.num_local_heads, self.qk_head_dim)
+        q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+
+        # KV (依次通过 kv_proj 的 3 层)
+        kv = self.kv_proj[0](hidden_states)[0]
+        kv = self.kv_proj[1](kv)
+        kv = self.kv_proj[2](kv)[0]
+        kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
+        k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+        # K rope
+        k_pe = self.k_rope(hidden_states)[0]
+        k_pe = k_pe.unsqueeze(1)
+
+        # Apply RoPE
+        pos_1d = positions[0] if positions.ndim == 2 else positions
+        # q_pe, k_pe = self.rotary_emb(pos_1d, q_pe, k_pe)
+        q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
+
+        # Assemble full Q and K
+        q[..., self.qk_nope_head_dim:] = q_pe
+        k = torch.empty_like(q)
+        k[..., :self.qk_nope_head_dim] = k_nope
+        k[..., self.qk_nope_head_dim:] = k_pe
+
+        # Pad V to qk_head_dim for PagedAttention KV cache alignment
+        v = F.pad(v, [0, self.qk_head_dim - self.v_head_dim], value=0)
+        v = v.view(-1, self.num_local_heads * self.qk_head_dim)
+
+        # PagedAttention
+        attn_output = self.attn(q, k, v)
+
+        # Strip V padding and project output
+        attn_output = attn_output.view(-1, self.num_local_heads, self.qk_head_dim)[..., :self.v_head_dim].reshape(-1, self.num_local_heads * self.v_head_dim)
+        output[:], _ = self.o_proj(attn_output)
+
+
+# ============================================================================
+# MLP
+# ============================================================================
+
+class NasMLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, quant_config=None, prefix=""):
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            hidden_size, [intermediate_size] * 2, bias=False,
+            quant_config=quant_config, prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size, hidden_size, bias=False,
+            quant_config=quant_config, prefix=f"{prefix}.down_proj",
+        )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+# ============================================================================
+# Decoder layer
+# ============================================================================
+
+class NasDecoderLayer(nn.Module):
+    """NAS decoder layer with per-layer attention type dispatch."""
+
+    def __init__(
+        self,
+        config: NasChildVLConfig,
+        text_config,
+        layer_idx: int,
+        cache_config: CacheConfig | None = None,
+        model_config: ModelConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.hidden_size = text_config.hidden_size
+
+        layer_config = config.get_layer_config(layer_idx)
+        self.attention_type = layer_config.get("attention_type", "full_attention")
+
+        # 🌟 核心修改：把 rope_parameters 的解析提前到所有 Attention 初始化之前
+        rope_parameters = getattr(text_config, "rope_parameters", None)
+        if rope_parameters is None:
+            rope_scaling = getattr(text_config, "rope_scaling", None)
+            rope_parameters = {"rope_type": "default", "factor": 1.0}
+            if rope_scaling:
+                rope_parameters.update(rope_scaling)
+
+        if self.attention_type in ("full_attention", "swa"):
+            sliding_window = None
+            if self.attention_type == "swa":
+                sliding_window = layer_config.get(
+                    "sliding_window",
+                    getattr(text_config, "sliding_window", 1024),
+                )
+            num_kv_heads = layer_config.get("gqa_num_kv_heads", 0) or text_config.num_key_value_heads
+            
+            self.self_attn = NasAttention(
+                hidden_size=text_config.hidden_size,
+                num_heads=text_config.num_attention_heads,
+                num_kv_heads=num_kv_heads,
+                head_dim=getattr(text_config, "head_dim",
+                                 text_config.hidden_size // text_config.num_attention_heads),
+                rope_parameters=rope_parameters, # 已经提前解析
+                max_position=text_config.max_position_embeddings,
+                rms_norm_eps=text_config.rms_norm_eps,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.self_attn",
+                sliding_window=sliding_window,
+            )
+
+        elif self.attention_type == "gdn":
+            self.self_attn = NasGDNAttention(
+                hidden_size=text_config.hidden_size,
+                layer_idx=layer_idx,
+                cache_config=cache_config,
+                model_config=model_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.self_attn",
+            )
+
+        elif self.attention_type == "kda":
+            self.self_attn = NasKDAAttention(
+                hidden_size=text_config.hidden_size,
+                layer_idx=layer_idx,
+                cache_config=cache_config,
+                model_config=model_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.self_attn",
+            )
+
+        elif self.attention_type == "mla":
+            self.self_attn = NasMLAAttention(
+                hidden_size=text_config.hidden_size,
+                layer_idx=layer_idx,
+                cache_config=cache_config,
+                model_config=model_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.self_attn",
+                num_heads=16,
+                rope_parameters=rope_parameters,  # 🌟 新增：把 rope_parameters 传给 MLA
+            )
+
+        else:
+            raise ValueError(f"Unsupported attention type: {self.attention_type}")
+
+        # Determine whether this layer has MLP based on ffn_type from config.
+        # Training code: ffn_type="ffn" -> has MLP + post_attention_layernorm;
+        # ffn_type="no-op" -> no MLP. This applies to ALL attention types
+        # including GDN and KDA (only layer 14 is no-op in current config).
+        ffn_type = layer_config.get("ffn_type", "ffn")
+        self.has_mlp = ffn_type != "no-op"
+
+        self.input_layernorm = RMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps)
+
+        if self.has_mlp:
+            intermediate_size = layer_config.get("child_intermediate_size", text_config.intermediate_size)
+            self.mlp = NasMLP(
+                hidden_size=text_config.hidden_size,
+                intermediate_size=intermediate_size,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+            self.post_attention_layernorm = RMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps)
+        else:
+            self.mlp = None
+            self.post_attention_layernorm = None
+
+        logger.info("NAS layer %d: %s (%s), has_mlp: %s", layer_idx, self.attention_type, type(self.self_attn).__name__, self.has_mlp)
+
+    def forward(self, positions, hidden_states, residual):
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        attn_output = torch.empty_like(hidden_states)
+
+        # All attention types use same interface
+        self.self_attn(positions=positions, hidden_states=hidden_states, output=attn_output)
+
+        hidden_states = attn_output
+        
+        # Only execute post_norm and mlp for layers that have them
+        # (based on ffn_type from config, not attention_type)
+        if self.has_mlp:
+            hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+            hidden_states = self.mlp(hidden_states)
+
+        return hidden_states, residual
+
+
+# ============================================================================
+# Text model
+# ============================================================================
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        # positions is (3, seq_len) for mrope, or (seq_len,) otherwise
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+        # IntermediateTensors-like dict, same leading dim as inputs_embeds
+        "deepstack_input_embeds": 0,
+    }
+)
+class NasChildVLTextModel(nn.Module):
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config: NasChildVLConfig = vllm_config.model_config.hf_config
+        text_config = config.text_config
+        cache_config = vllm_config.cache_config
+        model_config = vllm_config.model_config
+        quant_config = vllm_config.quant_config
+
+        self.config = text_config
+        self.nas_config = config
+        self.vocab_size = text_config.vocab_size
+
+        if get_pp_group().is_first_rank:
+            self.embed_tokens = VocabParallelEmbedding(
+                text_config.vocab_size, text_config.hidden_size,
+                prefix=f"{prefix}.embed_tokens",
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+
+        def get_layer(prefix: str):
+            layer_idx = extract_layer_index(prefix)
+            return NasDecoderLayer(
+                config=config,
+                text_config=text_config,
+                layer_idx=layer_idx,
+                cache_config=cache_config,
+                model_config=model_config,
+                quant_config=quant_config,
+                prefix=prefix,
+            )
+
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            text_config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers",
+        )
+
+        if get_pp_group().is_last_rank:
+            self.norm = RMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], text_config.hidden_size
+        )
+
+    def embed_input_ids(self, input_ids):
+        return self.embed_tokens(input_ids)
+
+    def forward(self, input_ids, positions, intermediate_tensors=None,
+                inputs_embeds=None, deepstack_input_embeds=None):
+        if get_pp_group().is_first_rank:
+            hidden_states = inputs_embeds if inputs_embeds is not None else self.embed_input_ids(input_ids)
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+
+        for layer_idx, layer in islice(enumerate(self.layers), self.start_layer, self.end_layer):
+            hidden_states, residual = layer(positions=positions, hidden_states=hidden_states, residual=residual)
+            if deepstack_input_embeds is not None and layer_idx in range(len(deepstack_input_embeds)):
+                hidden_states = hidden_states + deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"]
+
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors({"hidden_states": hidden_states, "residual": residual})
+        hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+
+# ============================================================================
+# CausalLM wrapper - HasInnerState
+# (NOT IsHybrid: would conflict with the manual `_NAS_PAGE_SIZE_PADDED` override
+#  in get_kv_cache_spec; HybridAttentionMambaModelConfig auto-pads to a different
+#  value and clashes with MLA's larger page size.)
+# ============================================================================
+
+class NasChildVLForCausalLM(nn.Module, HasInnerState, SupportsPP):
+    """NAS-Child-VL causal language model with HasInnerState for mamba state management."""
+
+    has_inner_state = True
+
+    packed_modules_mapping = {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config: NasChildVLConfig = vllm_config.model_config.hf_config
+        text_config = config.text_config
+        quant_config = vllm_config.quant_config
+
+        self.config = text_config
+        self.nas_config = config
+        self.model = NasChildVLTextModel(
+            vllm_config=vllm_config.with_hf_config(config),
+            prefix=maybe_prefix(prefix, "model"),
+        )
+
+        if get_pp_group().is_last_rank:
+            if getattr(config, "tie_word_embeddings", True):
+                self.lm_head = self.model.embed_tokens
+            else:
+                self.lm_head = ParallelLMHead(
+                    text_config.vocab_size, text_config.hidden_size,
+                    quant_config=quant_config,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )
+        else:
+            self.lm_head = PPMissingLayer()
+
+        self.logits_processor = LogitsProcessor(text_config.vocab_size)
+        self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
+
+    @property
+    def start_layer(self):
+        return self.model.start_layer
+
+    def embed_input_ids(self, input_ids):
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(self, input_ids, positions, intermediate_tensors=None,
+                inputs_embeds=None, deepstack_input_embeds=None, **kwargs):
+        return self.model(input_ids, positions, intermediate_tensors,
+                          inputs_embeds, deepstack_input_embeds)
+
+    def compute_logits(self, hidden_states):
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(
+        cls,
+        vllm_config: VllmConfig,
+    ) -> tuple[torch.dtype, ...]:
+        """Return state dtypes for the largest mamba layer type (GDN 2-state)."""
+        return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
+            vllm_config.model_config.dtype,
+            vllm_config.cache_config.mamba_cache_dtype,
+            vllm_config.cache_config.mamba_ssm_cache_dtype,
+        )
+
+    @classmethod
+    def get_mamba_state_shape_from_config(
+        cls, vllm_config: VllmConfig,
+    ) -> tuple[tuple[int, ...], ...]:
+        """Return state shapes for GDN layers (the reference mamba type)."""
+        tp_size = vllm_config.parallel_config.tensor_parallel_size
+        # GDN: num_heads=6, head_k_dim=256, head_v_dim=512, conv_kernel_size=4
+        return MambaStateShapeCalculator.gated_delta_net_state_shape(
+            tp_size,
+            6,   # num_k_heads
+            6,   # num_v_heads
+            256, # head_k_dim
+            512, # head_v_dim
+            4,   # conv_kernel_size
+        )
+
+    @classmethod
+    def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, ...]:
+        return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights with name remapping for all layer types.
+
+        Checkpoint format:
+          - FULL/SWA: q_proj, k_proj, v_proj (stacked to qkv_proj), q_norm, k_norm, o_proj
+          - GDN: native weights loaded into self_attn.XXX directly
+          - KDA: native weights loaded into self_attn.XXX with name remapping
+          - MLA: kv_down_proj.0->kv_proj.0, kv_down_proj.1->kv_proj.1, kv_up_proj->kv_proj.2
+                 loaded into self_attn.mla.XXX
+        """
+        stacked_params_mapping = [
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+
+        # Build attention type map
+        layer_attn_types: dict[int, str] = {}
+        for i, lc in enumerate(self.nas_config.nas_layer_config):
+            layer_attn_types[i] = lc.get("attention_type", "full_attention")
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        def _load_param(param_name: str, loaded_weight: torch.Tensor) -> bool:
+            if param_name not in params_dict:
+                logger.debug("Param not found: %s", param_name)
+                return False
+            param = params_dict[param_name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            try:
+                weight_loader(param, loaded_weight)
+            except (RuntimeError, ValueError, AssertionError) as e:
+                logger.warning("Skip loading %s: %s", param_name, e)
+                return False
+            loaded_params.add(param_name)
+            return True
+
+        sa_pattern = re.compile(r"(model\.layers\.(\d+)\.self_attn)\.(.*)")
+
+        for name, loaded_weight in weights:
+            m = sa_pattern.match(name)
+            if m:
+                sa_prefix = m.group(1)
+                layer_idx = int(m.group(2))
+                sa_suffix = m.group(3)
+                attn_type = layer_attn_types.get(layer_idx, "full_attention")
+
+                # ---- GDN ----
+                if attn_type == "gdn":
+                    # GDN native uses merged projections:
+                    # in_proj = [q_proj, k_proj, v_proj, g_proj] merged
+                    # in_proj_ba = [b_proj, a_proj] merged
+                    # Map checkpoint separate weights to merged params
+                    gdn_stacked = {
+                        "q_proj": ("in_proj", 0),
+                        "k_proj": ("in_proj", 1),
+                        "v_proj": ("in_proj", 2),
+                        "g_proj": ("in_proj", 3),
+                        "b_proj": ("in_proj_ba", 0),
+                        "a_proj": ("in_proj_ba", 1),
+                    }
+                    # Extract base name (without .weight/.bias)
+                    base = sa_suffix
+                    leaf = ""
+                    for lf in (".weight", ".bias"):
+                        if sa_suffix.endswith(lf):
+                            base = sa_suffix[:-len(lf)]
+                            leaf = lf
+                            break
+
+                    if base in gdn_stacked:
+                        merged_name, shard_id = gdn_stacked[base]
+                        target = f"{sa_prefix}.{merged_name}{leaf}"
+                        if target in params_dict:
+                            param = params_dict[target]
+                            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                            try:
+                                weight_loader(param, loaded_weight, shard_id)
+                            except (RuntimeError, ValueError, AssertionError) as e:
+                                logger.warning("Skip loading %s: %s", target, e)
+                                continue
+                            loaded_params.add(target)
+                        continue
+
+                    # o_norm.weight -> norm.weight
+                    if base == "o_norm":
+                        target = f"{sa_prefix}.norm{leaf}"
+                        _load_param(target, loaded_weight)
+                        continue
+
+                    # o_proj -> out_proj
+                    if base == "o_proj":
+                        target = f"{sa_prefix}.out_proj{leaf}"
+                        _load_param(target, loaded_weight)
+                        continue
+
+                    # q_conv1d/k_conv1d/v_conv1d -> merged conv1d
+                    # checkpoint: q_conv1d(1536,1,4), k_conv1d(1536,1,4), v_conv1d(3072,1,4)
+                    # model: conv1d(6144,1,4) = [q(1536) + k(1536) + v(3072)]
+                    if base in ("q_conv1d", "k_conv1d", "v_conv1d"):
+                        target = f"{sa_prefix}.conv1d{leaf}"
+                        if target in params_dict:
+                            param = params_dict[target]
+                            # Manual slice loading
+                            key_dim_local = 1536 // self.tp_size if hasattr(self, 'tp_size') else 1536
+                            # For TP=1: q=[0:1536], k=[1536:3072], v=[3072:6144]
+                            tp_size = get_tensor_model_parallel_world_size()
+                            kd = 1536 // tp_size
+                            vd = 3072 // tp_size
+                            offsets = {"q_conv1d": 0, "k_conv1d": kd, "v_conv1d": kd * 2}
+                            sizes = {"q_conv1d": kd, "k_conv1d": kd, "v_conv1d": vd}
+                            offset = offsets[base]
+                            size = sizes[base]
+                            # Shard the loaded weight for TP
+                            tp_rank = get_tensor_model_parallel_rank()
+                            if loaded_weight.size(0) > size:
+                                loaded_weight = loaded_weight[tp_rank * size:(tp_rank + 1) * size]
+                            param.data[offset:offset + size] = loaded_weight
+                            loaded_params.add(target)
+                        continue
+
+                    # Direct: A_log, dt_bias
+                    target = f"{sa_prefix}.{sa_suffix}"
+                    _load_param(target, loaded_weight)
+                    continue
+
+                # ---- KDA ----
+                if attn_type == "kda":
+                    # KDA native weight mapping:
+                    # checkpoint f_proj.0.* -> self_attn.f_proj_0.*
+                    # checkpoint f_proj.1.* -> self_attn.f_proj_1.*
+                    # checkpoint g_proj.0.* -> self_attn.g_proj_0.*
+                    # checkpoint g_proj.1.* -> self_attn.g_proj_1.*
+                    # checkpoint *_conv1d.weight (D, 1, K) -> load directly
+                    # checkpoint A_log (16,) -> reshape to (1, 1, local_num_heads, 1)
+                    kda_suffix = sa_suffix
+                    kda_suffix = kda_suffix.replace("f_proj.0.", "f_proj_0.")
+                    kda_suffix = kda_suffix.replace("f_proj.1.", "f_proj_1.")
+                    kda_suffix = kda_suffix.replace("g_proj.0.", "g_proj_0.")
+                    kda_suffix = kda_suffix.replace("g_proj.1.", "g_proj_1.")
+                    target = f"{sa_prefix}.{kda_suffix}"
+
+                    if "A_log" in kda_suffix and loaded_weight.dim() == 1:
+                        # Reshape from (num_heads,) to (1, 1, num_heads, 1)
+                        loaded_weight = loaded_weight.reshape(1, 1, -1, 1)
+
+                    if "conv1d.weight" in kda_suffix and loaded_weight.dim() == 3:
+                        # Already 3D (D, 1, K) - load directly
+                        _load_param(target, loaded_weight)
+                    else:
+                        _load_param(target, loaded_weight)
+                    continue
+
+                # ---- MLA ----
+                if attn_type == "mla":
+                    # Remap checkpoint names to native MLA submodules
+                    # Checkpoint: q_proj.{0,1,2} -> q_a_proj, q_a_layernorm, q_b_proj
+                    # Checkpoint: kv_down_proj.{0,1} -> kv_a_proj, kv_a_layernorm
+                    # Checkpoint: kv_up_proj -> kv_b_proj
+                    # Checkpoint: k_rope -> k_rope_proj
+                    # Checkpoint: o_proj -> o_proj
+                    # 数学等价映射：将 Checkpoint 的 Sequential 映射到 vLLM 的拆分模块
+                    target = f"{sa_prefix}.{sa_suffix}"
+                    _load_param(target, loaded_weight)
+                    continue
+                
+                    # mla_remap = {
+                    #     "q_proj.0": "q_a_proj",
+                    #     "q_proj.1": "q_a_layernorm",
+                    #     "q_proj.2": "q_b_proj",
+                    #     "k_rope": "k_rope_proj",
+                    #     "kv_down_proj.0": "kv_a_proj",
+                    #     "kv_down_proj.1": "kv_a_layernorm",
+                    #     "kv_up_proj": "kv_b_proj",
+                    #     "o_proj": "o_proj",
+                    # }
+                    # # Extract base and leaf (.weight/.bias)
+                    # mla_base = sa_suffix
+                    # mla_leaf = ""
+                    # for lf in (".weight", ".bias"):
+                    #     if sa_suffix.endswith(lf):
+                    #         mla_base = sa_suffix[:-len(lf)]
+                    #         mla_leaf = lf
+                    #         break
+                    # mapped = mla_remap.get(mla_base)
+                    # if mapped:
+                    #     target = f"{sa_prefix}.{mapped}{mla_leaf}"
+                    # else:
+                    #     target = f"{sa_prefix}.{sa_suffix}"
+                    # _load_param(target, loaded_weight)
+                    # continue
+
+                # ---- FULL / SWA: handle stacked params ----
+                shard_matched = False
+                for param_name, weight_name, shard_id in stacked_params_mapping:
+                    if sa_suffix == f"{weight_name}.weight" or sa_suffix == f"{weight_name}.bias":
+                        leaf = ".weight" if sa_suffix.endswith(".weight") else ".bias"
+                        candidate = f"{sa_prefix}.{param_name}{leaf}"
+                        if candidate in params_dict:
+                            param = params_dict[candidate]
+                            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                            try:
+                                weight_loader(param, loaded_weight, shard_id)
+                            except (RuntimeError, ValueError) as e:
+                                logger.warning("Skip loading %s: %s", candidate, e)
+                                break
+                            loaded_params.add(candidate)
+                            shard_matched = True
+                        break
+
+                if not shard_matched:
+                    # Direct load (o_proj, q_norm, k_norm)
+                    _load_param(f"{sa_prefix}.{sa_suffix}", loaded_weight)
+                continue
+
+            # ---- Non self_attn weights (MLP, layernorm, embed, norm, lm_head) ----
+            # Handle gate_proj/up_proj stacking for MLP
+            shard_matched = False
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                candidate = name.replace(weight_name, param_name)
+                if candidate in params_dict:
+                    param = params_dict[candidate]
+                    weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                    try:
+                        weight_loader(param, loaded_weight, shard_id)
+                    except (RuntimeError, ValueError, AssertionError) as e:
+                        logger.warning("Skip loading %s: %s", candidate, e)
+                        break
+                    loaded_params.add(candidate)
+                    shard_matched = True
+                break
+
+            if not shard_matched:
+                _load_param(name, loaded_weight)
+
+        return loaded_params
+
+
+# ============================================================================
+# Top-level multimodal model
+# ============================================================================
+
+from vllm.tokenizers.registry import cached_tokenizer_from_config
+from vllm.multimodal import MULTIMODAL_REGISTRY
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen3VLMultiModalProcessor,
+    info=Qwen3VLProcessingInfo,
+    dummy_inputs=Qwen3VLDummyInputsBuilder,
+)
+class NasChildVLForConditionalGeneration(Qwen3VLForConditionalGeneration):
+    """NAS-Child-VL multimodal model."""
+
+    has_inner_state = True
+
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "visual.": "visual.",
+            "lm_head.": "language_model.lm_head.",
+            "embed_tokens.": "language_model.model.embed_tokens.",
+            "norm.": "language_model.model.norm.",
+            "layers.": "language_model.model.layers.",
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
+        nn.Module.__init__(self)
+        config: NasChildVLConfig = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+
+        # Set triton allocator for chunk_gated_delta_rule etc. (runs in subprocess)
+        import triton
+        triton.set_allocator(
+            lambda size, align, stream: torch.cuda.caching_allocator_alloc(
+                size, stream=stream))
+
+        self.config = config
+        self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
+        self.multimodal_config = multimodal_config
+        self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+        self.video_pruning_rate = multimodal_config.video_pruning_rate
+        self.is_multimodal_pruning_enabled = multimodal_config.is_multimodal_pruning_enabled()
+
+        self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes")
+        self.deepstack_num_level = (
+            len(config.vision_config.deepstack_visual_indexes) if self.use_deepstack else 0
+        )
+        self.visual_dim = config.vision_config.out_hidden_size
+        self.multiscale_dim = self.visual_dim * self.deepstack_num_level
+
+        with self._mark_tower_model(vllm_config, {"image", "video"}):
+            self.visual = Qwen3_VisionTransformer(
+                config.vision_config,
+                norm_eps=getattr(config, "rms_norm_eps", 1e-6),
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "visual"),
+            )
+            if self.use_deepstack:
+                self.deepstack_input_embeds = [
+                    torch.zeros(
+                        vllm_config.scheduler_config.max_num_batched_tokens,
+                        config.text_config.hidden_size,
+                    )
+                    for _ in range(self.deepstack_num_level)
+                ]
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = NasChildVLForCausalLM(
+                vllm_config=vllm_config.with_hf_config(config),
+                prefix=maybe_prefix(prefix, "language_model"),
+            )
+
+        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
+
+    # def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    #     loader = AutoWeightsLoader(self)
+    #     return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        vision_weights = []
+        lm_weights = []
+        
+        for name, data in weights:
+            if name.startswith("visual."):
+                # 视觉侧保留原始 name
+                vision_weights.append((name, data))
+            else:
+                # 手动将 HuggingFace 的权重名映射为 NasChildVLForCausalLM 期望的内部名
+                if name.startswith("layers."):
+                    lm_name = "model." + name
+                elif name.startswith("embed_tokens."):
+                    lm_name = "model." + name
+                elif name.startswith("norm."):
+                    lm_name = "model." + name
+                elif name.startswith("lm_head."):
+                    lm_name = name  # lm_head 在 NasChildVLForCausalLM 中没有 model. 前缀
+                else:
+                    lm_name = name  # fallback
+                    
+                lm_weights.append((lm_name, data))
+                
+        loaded_params = set()
+        
+        # 1. 加载视觉侧权重 (AutoWeightsLoader 知道如何使用 mapper)
+        if vision_weights:
+            loader = AutoWeightsLoader(self)
+            vision_loaded = loader.load_weights(vision_weights, mapper=self.hf_to_vllm_mapper)
+            loaded_params.update(vision_loaded)
+            
+        # 2. 加载语言模型侧权重
+        if lm_weights:
+            # 此时传入的 lm_name 已经是 "model.layers.0..." 格式
+            # 完美匹配 NasChildVLForCausalLM 内部的正则和参数名
+            lm_loaded = self.language_model.load_weights(lm_weights)
+            
+            for name in lm_loaded:
+                # 补齐前缀以匹配顶层多模态模型 (NasChildVLForConditionalGeneration) 的 named_parameters
+                loaded_params.add(f"language_model.{name}")
+                
+        return loaded_params
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(
+        cls,
+        vllm_config: VllmConfig,
+    ) -> tuple[torch.dtype, ...]:
+        return _get_unified_state_dtype(vllm_config.model_config.dtype)
+
+    @classmethod
+    def get_mamba_state_shape_from_config(
+        cls, vllm_config: VllmConfig,
+    ) -> tuple[tuple[int, ...], ...]:
+        return _get_unified_state_shape(vllm_config.parallel_config.tensor_parallel_size)
+
+    @classmethod
+    def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, ...]:
+        # 🌟 核心修改：从 mamba_utils 导入正确的拷贝函数
+        from vllm.model_executor.layers.mamba.mamba_utils import (
+            get_conv_copy_spec,
+            get_temporal_copy_spec,
+        )
+        # 统一状态有 4 个：前 3 个是 conv，第 4 个是 temporal (ssm)
+        return (
+            get_conv_copy_spec,
+            get_conv_copy_spec,
+            get_conv_copy_spec,
+            get_temporal_copy_spec,
+        )

nas_child_vl_vllm/register.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Register NAS-Child-VL model with vLLM via monkey patch.
+
+Call `register()` once before using vLLM with NAS-Child-VL models.
+This is automatically done when `import nas_child_vl_vllm`.
+"""
+
+_registered = False
+
+
+def register():
+    global _registered
+    if _registered:
+        return
+    _registered = True
+
+    # 1. Register Config class in vLLM config discovery system
+    from vllm.transformers_utils.configs import _CLASS_TO_MODULE, __all__
+    from nas_child_vl_vllm.config import NasChildVLConfig
+
+    _CLASS_TO_MODULE["NasChildVLConfig"] = "nas_child_vl_vllm.config"
+    __all__.append("NasChildVLConfig")
+
+    # 2. Register model_type -> Config mapping
+    from vllm.transformers_utils.config import _CONFIG_REGISTRY
+    _CONFIG_REGISTRY["nas-child-vl"] = NasChildVLConfig
+
+    # 3. Register HuggingFace architecture -> model class mapping (lazy)
+    from vllm.model_executor.models.registry import ModelRegistry
+    ModelRegistry.register_model(
+        "NasChildVLModelForCausalLM",
+        "nas_child_vl_vllm.model:NasChildVLForConditionalGeneration",
+    )
+
+    # 4. Register Models Config (MambaModelConfig for hybrid state management)
+    from vllm.model_executor.models.config import MODELS_CONFIG_MAP, MambaModelConfig
+    MODELS_CONFIG_MAP["NasChildVLModelForCausalLM"] = MambaModelConfig
+
+    # 5. Monkey patch rmsnorm_fn (add sigmoid activation support)
+    import vllm.model_executor.layers.fla.ops.layernorm_guard as ln_mod
+    from nas_child_vl_vllm.rmsnorm_patch import rmsnorm_fn
+    ln_mod.rmsnorm_fn = rmsnorm_fn

nas_child_vl_vllm/rmsnorm_patch.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pure PyTorch rmsnorm_fn with sigmoid activation support.
+
+Replaces vLLM's Triton-based rmsnorm_fn which does not support
+sigmoid activation and can cause issues with fla GDN/KDA layers.
+"""
+
+import torch
+import torch.nn.functional as F
+
+
+def rmsnorm_fn(
+    x,
+    weight,
+    bias,
+    eps,
+    z=None,
+    group_size=None,
+    norm_before_gate=True,
+    activation="swish",
+    **kwargs,
+):
+    # 1. Native RMSNorm computation
+    variance = x.pow(2).mean(-1, keepdim=True)
+    out = x * torch.rsqrt(variance + eps)
+
+    if weight is not None:
+        out = out * weight
+    if bias is not None:
+        out = out + bias
+
+    # 2. If z is provided, this is a Fused RMSNorm + Gating
+    # Match fla FusedRMSNormGated behavior:
+    #   norm_before_gate=True  -> rms_norm(x) * weight * activation(z)
+    #   norm_before_gate=False -> rms_norm(x * activation(z)) * weight
+    # fla always does norm_before_gate=True (norm first, then multiply by gate).
+    if z is not None:
+        if activation == "swish" or activation == "silu":
+            gate = F.silu(z)
+        elif activation == "sigmoid":
+            gate = torch.sigmoid(z)
+        else:
+            gate = z
+        out = out * gate
+
+    return out