modeling_qwen2_hybrid.py

from __future__ import annotations
from .configuration_qwen2_hybrid import Qwen2HybridConfig
from typing import Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint

from transformers.cache_utils import Cache
from transformers.generation.utils import GenerationMixin
from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.utils import add_start_docstrings, logging

from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
from transformers.models.qwen2.modeling_qwen2 import (
    Qwen2Attention,
    Qwen2MLP,
    Qwen2PreTrainedModel,
    Qwen2RMSNorm,
    Qwen2RotaryEmbedding,
    apply_rotary_pos_emb,
    repeat_kv,
)
import transformers.models.qwen2.modeling_qwen2 as qwen2_modeling

logger = logging.get_logger(__name__)

_GQA_LAYERS       = set(range(0, 7))
_SHARED_LAYER     = 7
_SOFT_MID_LAYERS  = set(range(8,  23))
_SOFT_DEEP_LAYERS = set(range(23, 28))

_GQA_SLIDING_WINDOW  = 32768  # 前几层的SW为什么这么大
# _SOFT_SLIDING_WINDOW = 4096
_SOFT_SLIDING_WINDOW = 8192

_SHARED_RANK    = 320  # hidenstage是1536
_SOFT_RANK_MID  = 192
_SOFT_RANK_DEEP = 128

def _layer_role(layer_idx: int) -> str:
    if layer_idx in _GQA_LAYERS:   return "gqa"
    if layer_idx == _SHARED_LAYER: return "shared_mla"
    return "soft_mla"

def _mla_rank(layer_idx: int) -> int:
    if layer_idx == _SHARED_LAYER:    return _SHARED_RANK
    if layer_idx in _SOFT_MID_LAYERS: return _SOFT_RANK_MID
    return _SOFT_RANK_DEEP

def _mla_sliding_window(layer_idx: int) -> Optional[int]:
    return None if layer_idx == _SHARED_LAYER else _SOFT_SLIDING_WINDOW

def _mla_zone(layer_idx: int) -> str:
    if layer_idx in _GQA_LAYERS:      return "gqa"
    if layer_idx == _SHARED_LAYER:    return "shared"
    if layer_idx in _SOFT_MID_LAYERS: return "mid"
    return "deep"

# HybridCache:支持"Attention Sinks"的双模缓存
# 这部分的两个关键：混合缓存管理（HybridCache） 与 跨层特征共享（SharedLatentGate）
# HybirdModle主干文件中有实例化HybridCache的代码
class HybridCache(Cache):  # 这里继承了hf的Cache类
    def __init__(self, config: Qwen2Config):
        try:
            super().__init__(layers=config.num_hidden_layers)  # 新版本需要传入模型的层数
        except TypeError:
            super().__init__()

        self.config = config
        n = config.num_hidden_layers
        self._gqa_k:  List[Optional[torch.Tensor]] = [None] * n  # 维度：通常为 [batch, num_kv_heads, seq_len, head_dim]
        self._gqa_v:  List[Optional[torch.Tensor]] = [None] * n
        self._latent: List[Optional[torch.Tensor]] =[None] * n  # 第 7 层的 _latent 还会被 SharedLatentGate 调用，实现跨层特征传递
        self._seen_tokens: int = 0  # 记录模型迄今为止已经处理过的Token总数，计算CachePosition和RoPE的关键

    # 感觉好多此一举，为什么不直接调用update_gqa函数
    def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
        return self.update_gqa(key_states, value_states, layer_idx)

    # 返回现在已经处理了多长的序列了
    def get_seq_length(self, layer_idx: int = 0) -> int:
        return self._seen_tokens

    # 这啥意思？
    def get_max_cache_shape(self) -> Optional[int]:
        return None

    def update_gqa(self, key, value, layer_idx, sliding_window=_GQA_SLIDING_WINDOW):
        if self._gqa_k[layer_idx] is None:
            self._gqa_k[layer_idx] = key
            self._gqa_v[layer_idx] = value
        else:
            self._gqa_k[layer_idx] = torch.cat([self._gqa_k[layer_idx], key], dim=2)
            self._gqa_v[layer_idx] = torch.cat([self._gqa_v[layer_idx], value], dim=2)
        T = self._gqa_k[layer_idx].shape[2]  # seq_len当前历史信息长度
        
        # update_gqa的话只保留最后的sliding_window大小
        if T > sliding_window:
            self._gqa_k[layer_idx] = self._gqa_k[layer_idx][:, :, -sliding_window:, :]
            self._gqa_v[layer_idx] = self._gqa_v[layer_idx][:, :, -sliding_window:, :]
        if layer_idx == 0:
            self._seen_tokens += key.shape[2]  # 我对一次输入一个token还能理解，一会儿一次输入一个一会儿一次输出多个这件事不是特别理解
        return self._gqa_k[layer_idx], self._gqa_v[layer_idx]  # 返回加上了历史信息的KVCache

    # 我要修改一下这个方法，变成StreamingLLM的思路
    # def update_latent(self, c_kv, layer_idx, sliding_window=None):
    #     if self._latent[layer_idx] is None:
    #         self._latent[layer_idx] = c_kv
    #     else:
    #         self._latent[layer_idx] = torch.cat([self._latent[layer_idx], c_kv], dim=1)
    #     if sliding_window is not None:
    #         T = self._latent[layer_idx].shape[1]
    #         if T > sliding_window:
    #             self._latent[layer_idx] = self._latent[layer_idx][:, -sliding_window:, :]
    #     return self._latent[layer_idx]
    
    # 更新隐藏状态
    def update_latent(self, c_kv, layer_idx, sliding_window=None, sink_size=64):  # MLA因为SW比GQA小很多，所以需要sink
        if self._latent[layer_idx] is None:
            self._latent[layer_idx] = c_kv
        else:
            self._latent[layer_idx] = torch.cat([self._latent[layer_idx], c_kv], dim=1)  # latent这里的dim和上面gqa不太一样...

        if sliding_window is not None:
            T = self._latent[layer_idx].shape[1]
            if T > sliding_window:
                # 🚀 Attention Sinks: 保留头部 sink_size 个 Token，和尾部最新 Token！
                sink_tokens = self._latent[layer_idx][:, :sink_size, :]  # 保留前sink_size个记忆，这段记忆会一直保留，因为每次超出size，获取sink_size获取的都是sink_tokens
                recent_tokens = self._latent[layer_idx][:, -(sliding_window - sink_size):, :]  # 因为加入了sink_tokens所以SW要适当减小
                self._latent[layer_idx] = torch.cat([sink_tokens, recent_tokens], dim=1)  # 不过感觉这部分有些荣誉计算
        return self._latent[layer_idx]  # 返回新缓存
    
    # 返回SHARED_LAYER的Cache
    def get_shared_latent(self) -> Optional[torch.Tensor]:
        return self._latent[_SHARED_LAYER]

    # 好像是个移动都某个设备不是特别理解
    def to(self, device):
        # 模型参数一般调用model.to('cuda')还是device就可以移动到显卡了
        # 但是Cache类里的张量列表需要手动移动到GPU中确保可以顺利进行计算
        for i in range(len(self._gqa_k)):
            if self._gqa_k[i] is not None:
                self._gqa_k[i] = self._gqa_k[i].to(device)
                self._gqa_v[i] = self._gqa_v[i].to(device)
            if self._latent[i] is not None:
                self._latent[i] = self._latent[i].to(device)
        return self

# 为了把HybridCache伪装成一个Cache,从而兼容之前的代码逻辑
# 大概理解它的用途，但是不清楚调用和使用时机
class _GQASlotAdapter:
    def __init__(self, cache: HybridCache, sliding_window: int = _GQA_SLIDING_WINDOW):
        self._cache  = cache
        self._window = sliding_window

    def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
        return self._cache.update_gqa(key_states, value_states, layer_idx, self._window)

    def get_seq_length(self, layer_idx: int = 0) -> int:
        return self._cache.get_seq_length(layer_idx)

    def get_max_cache_shape(self) -> Optional[int]:
        return None

# 主要实现跨层特征通信和平滑微调
# 本质是一个带门控的残差投影器
# 让深层网络能够站在巨人的肩膀上，直接利用已经提取好的特征
class SharedLatentGate(nn.Module):
    def __init__(self, config: Qwen2Config):
        super().__init__()
        H = config.hidden_size
        self.cross_proj   = nn.Linear(_SHARED_RANK, H, bias=False)  # 从SHARED_RANK投影会H维度
        self.gate         = nn.Parameter(torch.full((H,), -4.0))  # H是标量，(H，)是一维向量，每个维度一个独立的门控机制
        self.warmup_alpha = nn.Parameter(torch.tensor(0.0))  # warmup_alpha是控制整体的一个加入比列，总阀门
        self.norm         = Qwen2RMSNorm(H, eps=config.rms_norm_eps)

    def forward(self, hidden_states, cache=None, explicit_shared=None):
        # 为了兼容训练/预填充模式和推理生成模式
        # 训练或首次输入时会使用explicit_shared
        if cache is not None and cache.get_shared_latent() is not None:  # 这里get_shared_latent是什么意思？
            shared = cache.get_shared_latent()  # 返回第七层截止目前的Cache
        elif explicit_shared is not None:  # 训练时选择显示传参，可以减少频繁读写Cache带来的不必要的开销
            shared = explicit_shared
        else:  # else主要是处理
            return hidden_states

        B, T, _ = hidden_states.shape  # 这不是当前输入长度吗
        T_full = shared.shape[1]  # 获取shared info的序列长度
        
        # 🚀 降维打击修复：只提取当前需要的 Token 进行投影，防止历史污染
        # 保证长度一致，就是每个ids的token只能获得相同ids token的浅层抽象信息
        # 这里其实让我有些疑惑，这样的机制是否真的有用，把浅层的东西往深层直接传递的意义是什么？
        if T_full != T:
            shared = shared[:, -T:, :]

        # 对我们把符合要求的C_kv找出来，然后要把维度从rank扩张会H,因为这个要加到当前输入的token的H上。
        proj = self.cross_proj(shared)
        proj = self.norm(proj)

        # 制作gate
        gate_weight = torch.sigmoid(self.gate) * self.warmup_alpha
        # hidden_states应该是[batchsize,seqlen,dim]
        return hidden_states + gate_weight.unsqueeze(0).unsqueeze(0) * proj  # unsqueeze是解压缩，也有增加维度的意思


class Qwen2MLASoftAttention(nn.Module):
    def __init__(self, config, layer_idx, kv_lora_rank, sliding_window):
        super().__init__()
        self.config         = config
        self.layer_idx      = layer_idx
        self.kv_lora_rank   = kv_lora_rank
        self.sliding_window = sliding_window

        H   = config.hidden_size
        nh  = config.num_attention_heads  # config中是12吧，能求出head_dim是128
        nkv = config.num_key_value_heads  # config中是2，用的也是GQA
        self.head_dim      = getattr(config, "head_dim", H // nh)
        self.num_heads     = nh
        self.num_kv_heads  = nkv
        self.num_kv_groups = nh // nkv  # repeatKV的时候需要这个group的参数，12heads 2kvheads，kv_group就是6（每6个heads一组）
        self.scaling       = self.head_dim ** -0.5  # 缩放系数，通过把方差拉回1来避免,softmax前数据分布太大，导致梯度消失，参数不更新

        self.q_proj = nn.Linear(H, nh * self.head_dim, bias=True)
        self.kv_down_proj = nn.Linear(H, kv_lora_rank, bias=False)  # 原本是2 x self.num_kv_heads x self.head_dim = 512 , 这里直接压成了kv_lora_rank{7:320,8~22:192,23~27:128}，最后实测表明这里压得有些多了
        self.k_up_proj    = nn.Linear(kv_lora_rank, nkv * self.head_dim, bias=True)  # 把低秩投会全注意力做计算这种合适吗，信息不是还是低秩的吗？
        self.v_up_proj    = nn.Linear(kv_lora_rank, nkv * self.head_dim, bias=True)  # 低秩投影回全注意力和GQA复制回全注意力，哪种更好？
        
        self.o_proj = nn.Linear(nh * self.head_dim, H, bias=False)  
        # 下面这两个norm是哪里做的？
        self.k_norm = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.v_norm = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        # 旋转emb层
        self.rotary_emb = Qwen2RotaryEmbedding(config=config)
        self.output_alpha = nn.Parameter(torch.tensor(0.0))

    # 这份代码中主要是一个是KVCache,一个是Mask，一个是postion的问题，不容易想明白
    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: Tuple[torch.Tensor, torch.Tensor],  # 这个还有些疑惑，position_embeddings是如何工作的？
        attention_mask: Optional[torch.Tensor],  # 这里传入的mask是4D的形式吗？
        past_key_values: Optional[HybridCache] = None,  # 这个是怎么用？
        cache_position: Optional[torch.LongTensor] = None,  # cache_Position怎么用？
        full_position_ids: Optional[torch.LongTensor] = None,  # 这里还有个position如何用？
        **kwargs,  # 这里有什么参数？
    ) -> Tuple[torch.Tensor, None]:
        B, T, H = hidden_states.shape
        cos, sin = position_embeddings  # 还没看内部

        # 这里q投影前后都没有进行norm，难道是上一层对输入x进行的norm吗
        q = self.q_proj(hidden_states)
        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        q, _ = apply_rotary_pos_emb(q, q, cos, sin)  # 这个要看一下

        # [batch_size,seq_len,kv_latent_dim]
        c_kv = self.kv_down_proj(hidden_states)
        
        # 🚀 终极防切片崩溃修复：独立拼接与缓存
        # 这里涉及kvcache的使用，是推理部分的核心，需要我去好好看一下，等下回来我先去看kvcache
        if past_key_values is not None:
            past_latent = past_key_values._latent[self.layer_idx]  # 这是什么意思，为什么这里获取past还有这个奇怪逻辑
            if past_latent is not None:
                full_c_kv = torch.cat([past_latent, c_kv], dim=1)
            else:
                full_c_kv = c_kv
            past_key_values.update_latent(c_kv, self.layer_idx, sliding_window=self.sliding_window)
        else:
            full_c_kv = c_kv
            
        T_kv = full_c_kv.shape[1]

        k = self.k_up_proj(full_c_kv).view(B, T_kv, self.num_kv_heads, self.head_dim)
        v = self.v_up_proj(full_c_kv).view(B, T_kv, self.num_kv_heads, self.head_dim)
        # 这里这个norm我不是很理解，为什么要获取kv后进行一次norm，为什么是先norm再transpose
        k = self.k_norm(k).transpose(1, 2)
        v = self.v_norm(v).transpose(1, 2)

        # # 🚀 绝对时空锁定修复：完美支持 bs>1 的 Left-Padding
        # if full_position_ids is not None:
        #     full_pos_ids = full_position_ids[:, -T_kv:]
        
        # 🚀 绝对时空锁定修复：完美支持 bs>1 的 Left-Padding
        # 下面这三行我也要替换掉
        # if full_position_ids is not None:
        #     full_pos_ids = full_position_ids[:, -T_kv:].contiguous()
        # elif cache_position is not None:
        
        # 🚀 绝对时空锁定修复：支持 Attention Sinks 与 Left-Padding
        S = 64  # Sink 大小，必须与 Cache 中保持一致
        # 这个full_position_ids还有些不清楚
        if full_position_ids is not None:
            total_seq_len = full_position_ids.shape[1]
            # 如果没超过滑动窗口，或者处于 Prefill 阶段 (T_kv == total_seq_len)，则直接取尾部
            if self.sliding_window is None or total_seq_len <= self.sliding_window or T_kv == total_seq_len:
                full_pos_ids = full_position_ids[:, -T_kv:].contiguous()
            else:
                # 触发 Sink 拼接逻辑：提取头部的 S 个位置，和尾部的残余位置
                sink_pos = full_position_ids[:, :S]
                recent_pos = full_position_ids[:, -(T_kv - S):]
                full_pos_ids = torch.cat([sink_pos, recent_pos], dim=1).contiguous()
        elif cache_position is not None:
            last_abs_pos_t = cache_position[-1]
            full_pos_ids = (torch.arange(T_kv, device=hidden_states.device, dtype=torch.long) + (last_abs_pos_t + 1 - T_kv)).unsqueeze(0)
        else:
            full_pos_ids = torch.arange(T_kv, device=hidden_states.device, dtype=torch.long).unsqueeze(0)

        # 生成rotary的逻辑也需要好好看一下
        cos_k, sin_k = self.rotary_emb(k, full_pos_ids)
        k, _ = apply_rotary_pos_emb(k, k, cos_k, sin_k)

        k = repeat_kv(k, self.num_kv_groups)
        v = repeat_kv(v, self.num_kv_groups)

        # 这里切换成连续是什么意思？
        q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
        
        # kv_seq_len = k.shape[2]
        # if attention_mask is not None and attention_mask.shape[-1] > kv_seq_len:
        #     attention_mask = attention_mask[..., :, -kv_seq_len:]
        # 修改后逻辑，加contiguous
        # kv_seq_len = k.shape[2]
        # if attention_mask is not None and attention_mask.shape[-1] > kv_seq_len:
        #     attention_mask = attention_mask[..., :, -kv_seq_len:].contiguous()
        # 下面这里也是我新修改的，稍微有些难理解，和sink有关系
        kv_seq_len = k.shape[2]
        if attention_mask is not None and attention_mask.shape[-1] > kv_seq_len:
            total_mask_len = attention_mask.shape[-1]
            if self.sliding_window is None or total_mask_len <= self.sliding_window or kv_seq_len == total_mask_len:
                attention_mask = attention_mask[..., :, -kv_seq_len:].contiguous()
            else:
                # 🚀 掩码也要同步拼接 Sink
                sink_mask = attention_mask[..., :, :S]
                recent_mask = attention_mask[..., :, -(kv_seq_len - S):]
                attention_mask = torch.cat([sink_mask, recent_mask], dim=-1).contiguous()
        
        
        is_causal = True if (attention_mask is None and T > 1) else False
        
        out = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=attention_mask,
            dropout_p=0.0,
            is_causal=is_causal,
            scale=self.scaling
        )

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        out = self.o_proj(out) * self.output_alpha
        return out, c_kv
    

# 上一个层self.layers就是堆叠了一堆decoder
class Qwen2HybridDecoderLayer(nn.Module):
    def __init__(self, config: Qwen2Config, layer_idx: int):
        super().__init__()
        self.layer_idx  = layer_idx
        self.layer_role = _layer_role(layer_idx)

        if self.layer_role == "gqa":
            attn_impl = getattr(config, "_attn_implementation", "sdpa")
            attn_class = getattr(qwen2_modeling, "QWEN2_ATTENTION_CLASSES", {}).get(attn_impl, Qwen2Attention)
            self.self_attn = attn_class(config=config, layer_idx=layer_idx)
        else:
            self.self_attn = Qwen2MLASoftAttention(
                config=config, layer_idx=layer_idx,
                kv_lora_rank=_mla_rank(layer_idx), sliding_window=_mla_sliding_window(layer_idx)
            )

        self.shared_gate = SharedLatentGate(config) if self.layer_role == "soft_mla" else None
        self.mlp = Qwen2MLP(config)
        self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self, hidden_states, attention_mask=None, position_ids=None, past_key_values=None,
        use_cache=False, cache_position=None, position_embeddings=None, output_attentions=False,
        shared_latent=None, full_position_ids=None, **kwargs,
    ):
        if self.shared_gate is not None:
            # 在模型的前 6 层，为了兼容 GQA，传入的是 _GQASlotAdapter
            # 不是很理解这里的适配，前六层既然是适配器了，为什么还需要调用sharedgate
            real_cache = past_key_values._cache if isinstance(past_key_values, _GQASlotAdapter) else past_key_values
            # 这里的real_cache是一个HybridCache对象
            hidden_states = self.shared_gate(hidden_states, cache=real_cache, explicit_shared=shared_latent)

        # Decoder的前半部分mid_output = x + Atten(Norm(x))
        residual = hidden_states  # 一个decoder要进行残差链接的
        normed_input = self.input_layernorm(hidden_states)  # Attention前做了input_norm了

        # 为什么gqa传的position_ids，mla传的是full_position_ids
        if self.layer_role == "gqa":
            attn_outputs = self.self_attn(
                hidden_states=normed_input, attention_mask=attention_mask, position_ids=position_ids,
                past_key_value=past_key_values, output_attentions=output_attentions, use_cache=use_cache,
                cache_position=cache_position, position_embeddings=position_embeddings,  # gqa的位置信息已经被处理过一部分了，是增量处理
            )
            if len(attn_outputs) == 3:
                attn_out, _, past_key_values = attn_outputs
            elif len(attn_outputs) == 2:
                attn_out, past_key_values = attn_outputs
            else:
                attn_out = attn_outputs[0]; past_key_values = None
            hidden_states = attn_out
        else:
            attn_out, c_kv = self.self_attn(
                hidden_states=normed_input, position_embeddings=position_embeddings, attention_mask=attention_mask,
                past_key_values=past_key_values, cache_position=cache_position, full_position_ids=full_position_ids,  # mla需要全量处理所有位置信息（），是全量处理
            )
            hidden_states = attn_out
            if self.layer_role == "shared_mla":
                shared_latent = c_kv

        hidden_states = residual + hidden_states

        # 下面是标准Decoder的后半块，output = x + MLP(Norm(x))
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states, shared_latent  # 返回残差块输出hidden_states可以理解，但shared_latent是什么意思，是训练时的显示串联吗？

@add_start_docstrings("Qwen2.5-Coder 非对称混合架构主干，v9。")
class Qwen2HybridModel(Qwen2PreTrainedModel):
    config_class = Qwen2HybridConfig  # <--- 就是缺了这一行！
    def __init__(self, config: Qwen2HybridConfig):
        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([Qwen2HybridDecoderLayer(config, i) for i in range(config.num_hidden_layers)])
        self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = Qwen2RotaryEmbedding(config=config)
        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self): return self.embed_tokens
    def set_input_embeddings(self, value): self.embed_tokens = value

    def forward(
        self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None,
        inputs_embeds=None, use_cache=None, cache_position=None, output_attentions=False,
        output_hidden_states=False, return_dict=True, **kwargs,
    ):
        # 输入处理
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("必须且只能指定 input_ids 或 inputs_embeds 之一")
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        B, T, _ = inputs_embeds.shape

        # 判断是否使用Cache，如果使用且没有创建合适类型就在这里创建
        if use_cache:
            if not isinstance(past_key_values, HybridCache):
                past_key_values = HybridCache(config=self.config)

        # 生成当前输入token在整个序列中的"绝对位置索引流水号"
        # cache_postion是给新来的每个Token分配的唯一门牌号，有些迷惑
        if cache_position is None:
            past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(past_seen, past_seen + T, device=inputs_embeds.device)

        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)
        # # 🚀 绝对时空锁定：提取真实的 Position IDs，完美解决 Left-Padding 导致的 RoPE 错位！
        # if getattr(self.config, "_attn_implementation", "sdpa") == "sdpa" and not output_attentions and attention_mask is None:
        #     causal_4d = None
        #     full_position_ids = None
        # else:
        #     past_kv_len = int(cache_position[0].item()) if T > 0 else 0
        #     causal_4d   = _prepare_4d_causal_attention_mask(
        #         attention_mask, (B, T), inputs_embeds, past_kv_len, sliding_window=None
        #     )
        #     if attention_mask is not None and attention_mask.dim() == 2:
        #         full_position_ids = attention_mask.long().cumsum(-1) - 1
        #         full_position_ids = full_position_ids.masked_fill(attention_mask == 0, 1)
        #     else:
        #         full_position_ids = None
        
        # 🚀 绝对时空锁定：提取真实的 Position IDs，完美解决 Left-Padding 导致的 RoPE 错位！
        # 解决Left-Padding导致的位移偏差
        # 下面这部分代码只有预填充阶段进行，会根据attention_mask的情况计算每个token在序列中的绝对位置，同时能够处理好Left-Padding
        # 训练阶段是不是也一直走这部分逻辑，但是我传入的bin文件，是如何产生attention_mask的？
        if attention_mask is not None and attention_mask.dim() == 2:  # ！只有预填充时mask才是2d，推理Decoder到之后传递的就变成4d的mask了
            full_position_ids = attention_mask.long().cumsum(-1) - 1  # 前缀和累加+索引对齐
            full_position_ids = full_position_ids.masked_fill(attention_mask == 0, 1)  # 
        else:
            full_position_ids = None

        # 🌟 新增拦截器：如果 mask 存在但全是 1（无 padding），强行设为 None，保住 Flash Attention！
        # attention_mask是一个2d的提示器，主要适用于识别padding的，全1说明没有Padding
        is_all_ones = (attention_mask is None) or (attention_mask.min() == 1) 
        # output_attentions是布尔开关，是否需要每层计算出注意力权重（应该是用来调试的，观察每层的状态）
        if getattr(self.config, "_attn_implementation", "sdpa") == "sdpa" and not output_attentions and is_all_ones:
            causal_4d = None  # 没有padding直接用None，启用sdpa内部的causal mask逻辑
        else:  # 这里的意思是，如果没有加速，或者说就是需要使用自定义mask，走下面的逻辑
            past_kv_len = int(cache_position[0].item()) if T > 0 else 0
            # 将2d的attention_mask转成4d的mask张量
            causal_4d   = _prepare_4d_causal_attention_mask(
                attention_mask, (B, T), inputs_embeds, past_kv_len, sliding_window=None
            )

        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids)
        # 构建一个gqa适配器，给前六层用，7层以后的模型直接用past_key_values就行
        # 主要是因为前六层调用的是Transformers库里的Attention所以得把HybridCache封装的和之前的DynamicCache一样
        gqa_adapter = _GQASlotAdapter(past_key_values) if past_key_values is not None else None
        all_hidden_states = () if output_hidden_states else None
        shared_latent = None
        
        # 这里是按层遍历的逻辑
        for layer in self.layers:
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)
            effective_cache = gqa_adapter if layer.layer_role == "gqa" else past_key_values

            if self.gradient_checkpointing and self.training:
                if cache_position is not None:
                    assert cache_position.device == inputs_embeds.device
                outputs = torch.utils.checkpoint.checkpoint(
                    layer, hidden_states, causal_4d, position_ids, None, False, cache_position, 
                    position_embeddings, output_attentions, shared_latent, full_position_ids,
                    use_reentrant=False,
                )
                hidden_states, shared_latent = outputs[0], outputs[1]
            else:
                outputs = layer(
                    hidden_states, attention_mask=causal_4d, position_ids=position_ids,
                    past_key_values=effective_cache, use_cache=use_cache, cache_position=cache_position,
                    position_embeddings=position_embeddings, output_attentions=output_attentions,
                    shared_latent=shared_latent, full_position_ids=full_position_ids,
                )
                hidden_states, shared_latent = outputs[0], outputs[1]

        # 遍历完要进行一下norm这里是RMSnorm
        hidden_states = self.norm(hidden_states)
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(v for v in[hidden_states, past_key_values if use_cache else None, all_hidden_states, None] if v is not None)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None,
            hidden_states=all_hidden_states, attentions=None,
        )

class Qwen2HybridForCausalLM(Qwen2PreTrainedModel, GenerationMixin):
    _tied_weights_keys = ["lm_head.weight"]
    config_class = Qwen2HybridConfig  # <--- 就是缺了这一行！
    def __init__(self, config: Qwen2HybridConfig):
        super().__init__(config)
        self.model      = Qwen2HybridModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head    = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def _init_weights(self, module: nn.Module):
        super()._init_weights(module)
        if isinstance(module, Qwen2MLASoftAttention):
            nn.init.zeros_(module.output_alpha)
        elif isinstance(module, SharedLatentGate):
            nn.init.zeros_(module.warmup_alpha)
            nn.init.constant_(module.gate, -4.0)

    def get_input_embeddings(self): return self.model.embed_tokens
    def set_input_embeddings(self, value): self.model.embed_tokens = value
    def get_output_embeddings(self): return self.lm_head
    def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings
    def set_decoder(self, decoder): self.model = decoder
    def get_decoder(self): return self.model

    def forward(
        self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None,
        inputs_embeds=None, labels=None, use_cache=None, cache_position=None, output_attentions=False,
        output_hidden_states=False, return_dict=True, **kwargs,
    ) -> Union[CausalLMOutputWithPast, Tuple]:

        outputs = 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,
            cache_position=cache_position, output_attentions=output_attentions,
            output_hidden_states=output_hidden_states, return_dict=True,
        )
        hidden_states = outputs.last_hidden_state
        logits = self.lm_head(hidden_states).float()

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1), ignore_index=-100)

        if not return_dict:
            out = (logits,)
            if use_cache: out = out + (outputs.past_key_values,)
            if output_hidden_states: out = out + (outputs.hidden_states,)
            return ((loss,) + out) if loss is not None else out

        return CausalLMOutputWithPast(
            loss=loss, logits=logits, past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states, attentions=outputs.attentions,
        )

    def prepare_inputs_for_generation(
        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, **kwargs,
    ) -> dict:
        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0

        if past_len > 0:
            if inputs_embeds is not None:
                inputs_embeds = inputs_embeds[:, -1:]
            else:
                input_ids = input_ids[:, -1:]

        position_ids = kwargs.get("position_ids", None)
        if attention_mask is not None and position_ids is None:
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids = position_ids.masked_fill(attention_mask == 0, 1)
            if past_len > 0:
                position_ids = position_ids[:, -input_ids.shape[1]:]

        # 好像是decode的生成阶段执行的
        if cache_position is None:
            cache_position = torch.arange(past_len, past_len + input_ids.shape[1], device=input_ids.device)

        model_inputs = {}
        if inputs_embeds is not None and past_len == 0:
            model_inputs["inputs_embeds"] = inputs_embeds
        else:
            model_inputs["input_ids"] = input_ids

        model_inputs.update({
            "past_key_values": past_key_values,
            "use_cache": kwargs.get("use_cache", True),
            "attention_mask": attention_mask,
            "position_ids": position_ids,
            "cache_position": cache_position,
        })
        return model_inputs

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        for i in range(len(past_key_values._gqa_k)):
            if past_key_values._gqa_k[i] is not None:
                past_key_values._gqa_k[i] = past_key_values._gqa_k[i].index_select(0, beam_idx)
                past_key_values._gqa_v[i] = past_key_values._gqa_v[i].index_select(0, beam_idx)
            if past_key_values._latent[i] is not None:
                past_key_values._latent[i] = past_key_values._latent[i].index_select(0, beam_idx)
        return past_key_values

def _svd_project_kv(k_weight, v_weight, kv_rank, k_bias=None, v_bias=None):
    nkv_d = k_weight.shape[0]
    orig_dtype = k_weight.dtype
    M = torch.cat([k_weight, v_weight], dim=0).float()
    U, S, Vh = torch.linalg.svd(M, full_matrices=False)
    r = min(kv_rank, S.shape[0])
    S_sqrt = S[:r].sqrt().unsqueeze(0)
    down_w = Vh[:r, :].to(orig_dtype)
    k_up_w = (U[:nkv_d, :r] * S_sqrt).to(orig_dtype)
    v_up_w = (U[nkv_d:, :r] * S_sqrt).to(orig_dtype)
    k_up_bias = k_bias.to(orig_dtype) if k_bias is not None else None
    v_up_bias = v_bias.to(orig_dtype) if v_bias is not None else None
    return down_w, k_up_w, v_up_w, k_up_bias, v_up_bias

def migrate_weights_from_qwen2(hybrid_model, original_state_dict, svd_verbose=True):
    hybrid_sd = hybrid_model.state_dict()
    new_sd, unmapped = {},[]
    layer_kv = {}
    for orig_key, orig_val in original_state_dict.items():
        if not orig_key.startswith("model.layers."): continue
        parts = orig_key.split(".")
        layer_idx = int(parts[2])
        suffix = ".".join(parts[3:])
        if _layer_role(layer_idx) == "gqa": continue
        if suffix == "self_attn.k_proj.weight": layer_kv.setdefault(layer_idx, {})["k_w"] = orig_val
        elif suffix == "self_attn.v_proj.weight": layer_kv.setdefault(layer_idx, {})["v_w"] = orig_val
        elif suffix == "self_attn.k_proj.bias": layer_kv.setdefault(layer_idx, {})["k_b"] = orig_val
        elif suffix == "self_attn.v_proj.bias": layer_kv.setdefault(layer_idx, {})["v_b"] = orig_val

    for orig_key, orig_val in original_state_dict.items():
        if not orig_key.startswith("model.layers."):
            if orig_key in hybrid_sd: new_sd[orig_key] = orig_val
            else: unmapped.append(orig_key)
            continue
        parts = orig_key.split(".")
        layer_idx = int(parts[2])
        suffix = ".".join(parts[3:])
        role = _layer_role(layer_idx)
        tgt = f"model.layers.{layer_idx}.{suffix}"
        if role == "gqa":
            if tgt in hybrid_sd: new_sd[tgt] = orig_val
            else: unmapped.append(orig_key)
            continue
        if suffix in ("self_attn.q_proj.weight", "self_attn.q_proj.bias"):
            if tgt in hybrid_sd: new_sd[tgt] = orig_val
        elif suffix in ("self_attn.k_proj.weight", "self_attn.v_proj.weight", "self_attn.k_proj.bias", "self_attn.v_proj.bias"):
            pass 
        elif suffix == "self_attn.o_proj.weight":
            if tgt in hybrid_sd and hybrid_sd[tgt].shape == orig_val.shape: new_sd[tgt] = orig_val
            else: unmapped.append(f"{orig_key} [shape mismatch or missing]")
        elif "mlp." in suffix or "layernorm" in suffix:
            if tgt in hybrid_sd: new_sd[tgt] = orig_val
        else:
            unmapped.append(orig_key)

    svd_done, svd_errors = 0,[]
    for layer_idx in sorted(layer_kv.keys()):
        kv = layer_kv[layer_idx]
        k_w, v_w = kv.get("k_w"), kv.get("v_w")
        if k_w is None or v_w is None:
            svd_errors.append(f"Layer {layer_idx}: 缺少 k_w 或 v_w")
            continue
        rank = _mla_rank(layer_idx)
        zone = _mla_zone(layer_idx)
        k_b, v_b = kv.get("k_b"), kv.get("v_b")
        if svd_verbose:
            bias_info = "w/ bias" if k_b is not None else "no bias"
            print(f"  [SVD] Layer {layer_idx:2d} [{zone:6s}] k{list(k_w.shape)} + v{list(v_w.shape)} → rank={rank:3d} ({bias_info})")
        try:
            down_w, k_up_w, v_up_w, k_up_b, v_up_b = _svd_project_kv(k_w, v_w, rank, k_bias=k_b, v_bias=v_b)
        except Exception as exc:
            svd_errors.append(f"Layer {layer_idx}: SVD failed — {exc}")
            continue
        pfx = f"model.layers.{layer_idx}.self_attn"
        for key, weight in[(f"{pfx}.kv_down_proj.weight", down_w), (f"{pfx}.k_up_proj.weight", k_up_w), (f"{pfx}.v_up_proj.weight", v_up_w)]:
            if key in hybrid_sd and hybrid_sd[key].shape == weight.shape: new_sd[key] = weight
            else: svd_errors.append(f"{key}: shape mismatch")
        for key, bias_val in[(f"{pfx}.k_up_proj.bias", k_up_b), (f"{pfx}.v_up_proj.bias", v_up_b)]:
            if bias_val is not None and key in hybrid_sd:
                if hybrid_sd[key].shape == bias_val.shape: new_sd[key] = bias_val
        svd_done += 1

    custom_written = 0
    for key in hybrid_sd:
        if key.endswith(".self_attn.output_alpha"):
            new_sd[key] = torch.tensor(0.0)
            custom_written += 1
        elif key.endswith(".shared_gate.warmup_alpha"):
            new_sd[key] = torch.tensor(0.0)
            custom_written += 1
        elif key.endswith(".shared_gate.gate"):
            new_sd[key] = torch.full(hybrid_sd[key].shape, -4.0)
            custom_written += 1

    missing, unexpected = hybrid_model.load_state_dict(new_sd, strict=False)
    if svd_verbose:
        sep = "=" * 65
        print(f"\n{sep}\n[migrate_weights_v9] Qwen2 → Hybrid v9  迁移完成\n{sep}")
        print(f"  Rank: shared(L7)={_SHARED_RANK} | mid(L8-22)={_SOFT_RANK_MID} | deep(L23-27)={_SOFT_RANK_DEEP}")
        print(f"  SVD 热启动      : {svd_done} 层\n  自定义参数写入  : {custom_written} 个\n  总写入 keys     : {len(new_sd)}")
        print(f"  缺失(新增模块)  : {len(missing):3d}\n  意外(多余)      : {len(unexpected):3d}\n  未映射原始 keys : {len(unmapped):3d}")
        if svd_errors:
            for e in svd_errors: print(f"  ⚠ {e}")
        print(f"{sep}\n")
    return unmapped

def get_alpha_param_groups(model, base_lr, alpha_lr_scale=10.0):
    alpha_params, base_params, alpha_names = [], [],[]
    for name, param in model.named_parameters():
        if not param.requires_grad: continue
        if name.endswith(".self_attn.output_alpha") or name.endswith(".shared_gate.warmup_alpha"):
            alpha_params.append(param)
            alpha_names.append(name)
        else: base_params.append(param)
    print(f"[get_alpha_param_groups]\n  Base params  : {len(base_params):4d}  lr={base_lr:.2e}\n  Alpha params : {len(alpha_params):4d}  lr={base_lr * alpha_lr_scale:.2e}")
    return[{"params": base_params, "lr": base_lr, "name": "base"}, {"params": alpha_params, "lr": base_lr * alpha_lr_scale, "name": "alpha_gate"}]

def verify_no_nan(model):
    nan_params =[f"  ✗ NaN in {n}  shape={list(p.shape)}" for n, p in model.named_parameters() if p.data.isnan().any()]
    if nan_params:
        print("[verify_no_nan] 发现 NaN 参数：\n" + "\n".join(nan_params))
        return False
    print(f"[verify_no_nan] ✓ 所有 {sum(1 for _ in model.parameters())} 个参数均无 NaN")
    return True

def verify_alpha_zero(model):
    problems =[]
    for name, param in model.named_parameters():
        if name.endswith(".self_attn.output_alpha") or name.endswith(".shared_gate.warmup_alpha"):
            if abs(param.item()) > 1e-6: problems.append(f"  ✗ {name} = {param.item():.6f}（应为 0.0）")
    if problems:
        print("[verify_alpha_zero] Alpha 初始化异常：\n" + "\n".join(problems))
        return False
    print("[verify_alpha_zero] ✓ 所有 output_alpha / warmup_alpha = 0.0")
    return True

__all__ =[
    "_SHARED_RANK", "_SOFT_RANK_MID", "_SOFT_RANK_DEEP", "_layer_role", "_mla_rank", "_mla_zone", "_mla_sliding_window",
    "_svd_project_kv", "HybridCache", "SharedLatentGate", "Qwen2MLASoftAttention", "Qwen2HybridDecoderLayer", 
    "Qwen2HybridModel", "Qwen2HybridForCausalLM", "migrate_weights_from_qwen2", "get_alpha_param_groups", 
    "verify_no_nan", "verify_alpha_zero",
]