model.py

import inspect
import torch
import math
import torch.nn as nn
from torch.nn import functional as F
from config import GPTConfig





# =============================================================================
# ========================= 基础组件: RMSNorm ==================================
# =============================================================================

class RMSNorm(nn.Module):
    """
    RMSNorm: 去掉 mean centering 和 bias，
    只保留 RMS 归一化 + 可学习缩放因子 γ。
    参考: https://arxiv.org/abs/1910.07467
    """
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        return output * self.weight


# =============================================================================
# ========================= 基础组件: RoPE ====================================
# =============================================================================

class RotaryPositionalEmbedding(nn.Module):
    """
    RoPE (Rotary Positional Embedding):
    将相对位置信息通过旋转变换注入 Q/K，支持长度外推。
    参考: https://arxiv.org/abs/2104.09864
    """
    def __init__(self, dim: int, max_seq_len: int = 8192, base: float = 10000.0):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = base
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._build_cache(max_seq_len)

    def _build_cache(self, seq_len: int):
        t = torch.arange(seq_len, dtype=self.inv_freq.dtype, device=self.inv_freq.device)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos(), persistent=False)
        self.register_buffer("sin_cached", emb.sin(), persistent=False)

    def forward(self, x, seq_len: int):
        if seq_len > self.max_seq_len:
            self._build_cache(seq_len)
            self.max_seq_len = seq_len
        return (
            self.cos_cached[:seq_len].to(x.dtype),
            self.sin_cached[:seq_len].to(x.dtype),
        )


def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin):
    cos = cos.unsqueeze(0).unsqueeze(0)
    sin = sin.unsqueeze(0).unsqueeze(0)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


# =============================================================================
# ========================= GQA 注意力 (支持 KV Cache) =========================
# =============================================================================

class GroupedQueryAttention(nn.Module):
    """
    GQA (Grouped-Query Attention):
    Q heads 分组共享 K/V heads，推理时大幅减少 KV Cache 显存。  # 训练技巧：降显存，提token吞吐量
    训练时正常计算，推理时支持 KV Cache。
    参考: https://arxiv.org/abs/2305.13245
    """
    def __init__(self, config, layer_idx: int = 0):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        assert config.n_head % config.n_kv_head == 0

        self.layer_idx = layer_idx
        self.n_head = config.n_head
        self.n_kv_head = config.n_kv_head
        self.n_rep = self.n_head // self.n_kv_head # 每个 KV head 需要被重复的次数
        self.head_dim = config.n_embd // config.n_head
        self.n_embd = config.n_embd

        self.wq = nn.Linear(config.n_embd, config.n_head * self.head_dim, bias=False)
        self.wk = nn.Linear(config.n_embd, config.n_kv_head * self.head_dim, bias=False)
        self.wv = nn.Linear(config.n_embd, config.n_kv_head * self.head_dim, bias=False)
        self.wo = nn.Linear(config.n_head * self.head_dim, config.n_embd, bias=False)
        self.wo.NANOGPT_SCALE_INIT = 1

        self.rotary_emb = RotaryPositionalEmbedding(
            self.head_dim, max_seq_len=config.block_size,
        )

    def _repeat_kv(self, x: torch.Tensor) -> torch.Tensor:
        if self.n_rep == 1:
            return x
        B, n_kv_head, T, head_dim = x.shape
        x = x[:, :, None, :, :].expand(B, n_kv_head, self.n_rep, T, head_dim)
        return x.reshape(B, n_kv_head * self.n_rep, T, head_dim)

    def forward(self, x, kv_cache=None, start_pos: int = 0):
        """
        Args:
            x: (B, T, C)
            kv_cache: KVCache 对象 (推理时使用, 训练时为 None)
            start_pos: 当前 token 在完整序列中的起始位置 (推理用)
        """
        B, T, C = x.size()

        q = self.wq(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2)

        # RoPE (位置偏移处理)
        if kv_cache is not None:
            cos, sin = self.rotary_emb(q, start_pos + T)
            cos = cos[start_pos:start_pos + T]
            sin = sin[start_pos:start_pos + T]
        else:
            cos, sin = self.rotary_emb(q, T)

        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # KV Cache
        if kv_cache is not None:
            k, v = kv_cache.update(self.layer_idx, k, v)

        k = self._repeat_kv(k)
        v = self._repeat_kv(v)


        # # 标准注意力计算: att = softmax((q @ k^T) / sqrt(d_k)) @ v
        # att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        # k_len = k.size(-2)
        # mask = torch.tril(
        #     torch.ones(T, k_len, device=x.device, dtype=torch.bool)
        # ).view(1, 1, T, k_len)
        # att = att.masked_fill(~mask, torch.finfo(att.dtype).min)
        # att = F.softmax(att, dim=-1)
        # y = att @ v  # 计算加权和，得到注意力输出（B, nh, T, T）x (B, nh, T, hs) -> (B, nh, T, hs)


        # Flash Attention (is_causal: 推理decode阶段T=1时不需要causal mask)  # 训练技巧：降显存，提token吞吐量
        is_causal = True if kv_cache is None else (T > 1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)


        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.wo(y)
        return y


# =============================================================================
# ========================= KV Cache ==========================================
# =============================================================================

class KVCache:
    """
    KV Cache: 自回归生成时缓存已计算的 K/V。
    预分配显存，避免反复 cat 造成碎片化。

    GQA 优势: 缓存大小 = n_kv_head * head_dim (而非 n_head * head_dim)
    """
    def __init__(self, config, batch_size: int, max_seq_len: int,
                 device: torch.device, dtype=torch.bfloat16):
        self.n_layers = config.n_layer
        self.n_kv_head = config.n_kv_head
        self.head_dim = config.n_embd // config.n_head
        self.max_seq_len = max_seq_len

        # 预分配: (n_layers, B, n_kv_head, max_seq_len, head_dim)
        self.k_cache = torch.zeros(
            self.n_layers, batch_size, self.n_kv_head, max_seq_len, self.head_dim,
            device=device, dtype=dtype
        )
        self.v_cache = torch.zeros(
            self.n_layers, batch_size, self.n_kv_head, max_seq_len, self.head_dim,
            device=device, dtype=dtype
        )
        self.seq_len = 0

    def update(self, layer_idx: int, k_new: torch.Tensor, v_new: torch.Tensor):
        """写入新 K/V 并返回完整缓存"""
        new_len = k_new.size(2)
        end = self.seq_len + new_len
        self.k_cache[layer_idx, :, :, self.seq_len:end, :] = k_new
        self.v_cache[layer_idx, :, :, self.seq_len:end, :] = v_new
        return (
            self.k_cache[layer_idx, :, :, :end, :],
            self.v_cache[layer_idx, :, :, :end, :],
        )

    def advance(self, n: int = 1):
        self.seq_len += n

    def reset(self):
        self.seq_len = 0
        self.k_cache.zero_()
        self.v_cache.zero_()

    @staticmethod
    def memory_footprint(config, batch_size=1, seq_len=2048, dtype=torch.bfloat16):
        """分析 KV Cache 显存占用"""
        bpe = 2 if dtype in (torch.bfloat16, torch.float16) else 4
        head_dim = config.n_embd // config.n_head
        gqa = 2 * config.n_layer * batch_size * config.n_kv_head * seq_len * head_dim * bpe
        mha = 2 * config.n_layer * batch_size * config.n_head * seq_len * head_dim * bpe
        return gqa, mha


# =============================================================================
# ========================= SwiGLU FFN ========================================
# =============================================================================

class SwiGLUFFN(nn.Module):
    """
    SwiGLU FFN: SiLU 门控 + 3 个权重矩阵
    SwiGLU(x) = (SiLU(xW1) ⊙ xW3) W2
    参考: https://arxiv.org/abs/2002.05202
    """
    def __init__(self, config):
        super().__init__()
        hidden_dim = int(4 * config.n_embd * 2 / 3)
        if hasattr(config, 'multiple_of'):
            hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)
        self.w1 = nn.Linear(config.n_embd, hidden_dim, bias=False)
        self.w3 = nn.Linear(config.n_embd, hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim, config.n_embd, bias=False)
        self.w2.NANOGPT_SCALE_INIT = 1

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


# =============================================================================
# ========================= MoE 层 ============================================
# =============================================================================

class MoEGate(nn.Module):
    """MoE 路由门控: Top-K 路由 + 辅助负载均衡损失"""
    def __init__(self, config):
        super().__init__()
        self.n_experts = config.n_experts
        self.top_k = config.n_experts_per_tok
        self.gate = nn.Linear(config.n_embd, config.n_experts, bias=False)
        self.aux_loss_coeff = config.aux_loss_coeff

    def forward(self, x):
        logits = self.gate(x)
        weights, indices = torch.topk(logits, self.top_k, dim=-1)
        weights = F.softmax(weights, dim=-1)

        if self.training:
            probs = F.softmax(logits, dim=-1)
            mask = F.one_hot(indices, num_classes=self.n_experts).sum(dim=1)
            f = mask.float().mean(dim=0)
            P = probs.mean(dim=0)
            aux_loss = self.aux_loss_coeff * self.n_experts * (f * P).sum()
        else:
            aux_loss = torch.tensor(0.0, device=x.device)

        return weights, indices, aux_loss


class MoEFFN(nn.Module):
    """
    MoE FFN: 共享专家 + 路由专家
    参考: DeepSeekMoE (https://arxiv.org/abs/2401.06066)
    """
    def __init__(self, config):
        super().__init__()
        self.n_experts = config.n_experts
        self.top_k = config.n_experts_per_tok
        self.n_embd = config.n_embd

        self.gate = MoEGate(config)
        self.experts = nn.ModuleList([SwiGLUFFN(config) for _ in range(config.n_experts)])

        self.n_shared_experts = getattr(config, 'n_shared_experts', 1)
        if self.n_shared_experts > 0:
            self.shared_experts = nn.ModuleList(
                [SwiGLUFFN(config) for _ in range(self.n_shared_experts)]
            )

    def forward(self, x):
        B, T, C = x.shape
        x_flat = x.view(-1, C)

        # 共享专家
        shared_output = torch.zeros_like(x_flat)
        if self.n_shared_experts > 0:
            for se in self.shared_experts:
                shared_output = shared_output + se(x_flat)
            if self.n_shared_experts > 1:
                shared_output = shared_output / self.n_shared_experts

        # 路由专家
        weights, indices, aux_loss = self.gate(x_flat)
        routed_output = torch.zeros_like(x_flat)

        for i in range(self.n_experts):
            token_idx, slot_idx = torch.where(indices == i)
            if token_idx.numel() == 0:
                continue
            expert_input = x_flat[token_idx]
            expert_output = self.experts[i](expert_input)
            expert_weights = weights[token_idx, slot_idx].unsqueeze(-1)
            # routed_output.index_add_(0, token_idx, expert_output * expert_weights)
            contrib = (expert_output * expert_weights).to(routed_output.dtype)
            routed_output.index_add_(0, token_idx, contrib)

        output = (shared_output + routed_output).view(B, T, C)
        return output, aux_loss


# =============================================================================
# ========================= Transformer Block ==================================
# =============================================================================

class Block(nn.Module):
    """
    Transformer Block:
    Pre-RMSNorm → GQA → 残差
    Pre-RMSNorm → SwiGLU (MoE) → 残差
    支持 Gradient Checkpointing
    """
    def __init__(self, config, layer_idx: int = 0):
        super().__init__()
        self.ln_1 = RMSNorm(config.n_embd, eps=config.norm_eps)
        self.attn = GroupedQueryAttention(config, layer_idx=layer_idx)
        self.ln_2 = RMSNorm(config.n_embd, eps=config.norm_eps)

        self.use_moe = config.use_moe
        if self.use_moe:
            self.ffn = MoEFFN(config)
        else:
            self.ffn = SwiGLUFFN(config)

        self.use_checkpoint = config.use_gradient_checkpointing

    def _attn_block(self, x):
        """注意力子块 (可被 checkpoint 包裹)"""
        return self.attn(self.ln_1(x))

    def _ffn_block(self, x):
        """FFN 子块 (可被 checkpoint 包裹)"""
        if self.use_moe:
            return self.ffn(self.ln_2(x))
        else:
            return self.ffn(self.ln_2(x)), torch.tensor(0.0, device=x.device)

    def forward(self, x, kv_cache=None, start_pos: int = 0):
        """
        前向传播，支持三种模式:
        1. 训练 + checkpoint: 使用 gradient checkpointing 节省显存
        2. 训练 - checkpoint: 标准前向
        3. 推理 (kv_cache is not None): 使用 KV Cache
        """
        # --- 注意力 ---
        if self.use_checkpoint and self.training and kv_cache is None:
            attn_out = torch.utils.checkpoint.checkpoint(
                self._attn_block, x,
                use_reentrant=False,
                preserve_rng_state=True,
            )
        else:
            attn_out = self.attn(self.ln_1(x), kv_cache=kv_cache, start_pos=start_pos)
        x = x + attn_out

        # --- FFN ---
        if self.use_checkpoint and self.training and kv_cache is None:
            if self.use_moe:
                ffn_out, aux_loss = torch.utils.checkpoint.checkpoint(
                    self._ffn_block, x,
                    use_reentrant=False,
                    preserve_rng_state=True,
                )
            else:
                def _dense_ffn(x_in):
                    out = self.ffn(self.ln_2(x_in))
                    return out, torch.tensor(0.0, device=x_in.device)
                ffn_out, aux_loss = torch.utils.checkpoint.checkpoint(
                    _dense_ffn, x,
                    use_reentrant=False,
                    preserve_rng_state=True,
                )
        else:
            if self.use_moe:
                ffn_out, aux_loss = self.ffn(self.ln_2(x))
            else:
                ffn_out = self.ffn(self.ln_2(x))
                aux_loss = torch.tensor(0.0, device=x.device)

        x = x + ffn_out
        return x, aux_loss


# =============================================================================
# ========================= GPT 模型 ==========================================
# =============================================================================

class GPT(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.config = config

        self.transformer = nn.ModuleDict(dict(
            wte=nn.Embedding(config.vocab_size, config.n_embd),
            h=nn.ModuleList([Block(config, layer_idx=i) for i in range(config.n_layer)]),
            ln_f=RMSNorm(config.n_embd, eps=config.norm_eps),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.transformer.wte.weight = self.lm_head.weight  # 权重共享

        self.apply(self._init_weights)
        self._print_model_info()

    # -------------------------------------------------------------------------
    def _init_weights(self, module):  # 权重初始化: Linear 正态分布, Embedding 正态分布, bias 置零
        if isinstance(module, nn.Linear):
            std = 0.02
            if hasattr(module, 'NANOGPT_SCALE_INIT'):
                std *= (2 * self.config.n_layer) ** -0.5
            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    # -------------------------------------------------------------------------
    def _print_model_info(self):  # 模型参数统计和配置信息展示
        cfg = self.config
        total = sum(p.numel() for p in self.parameters())
        print(f"\n{'='*60}")
        print(f"  Layers={cfg.n_layer}, Heads={cfg.n_head}, KV Heads={cfg.n_kv_head}")
        print(f"  Embed={cfg.n_embd}, BlockSize={cfg.block_size}")
        if cfg.use_moe:
            print(f"  MoE: {cfg.n_experts} experts, top-{cfg.n_experts_per_tok}, "
                  f"{cfg.n_shared_experts} shared")
        print(f"  GradCheckpoint={cfg.use_gradient_checkpointing}")
        print(f"  Total Parameters: {total:,} ({total/1e6:.1f}M)")
        if cfg.use_moe:
            active = self._estimate_active_params()
            print(f"  Active Parameters/token: ~{active:,} ({active/1e6:.1f}M)")
        print(f"{'='*60}\n")

    # -------------------------------------------------------------------------
    def _estimate_active_params(self):  # 估算每 token 活跃参数量 (MoE 模型)
        cfg = self.config
        emb = cfg.vocab_size * cfg.n_embd
        hd = cfg.n_embd // cfg.n_head
        attn = cfg.n_embd * cfg.n_head * hd + cfg.n_embd * cfg.n_kv_head * hd * 2 + cfg.n_head * hd * cfg.n_embd
        norms = cfg.n_embd * 2
        hidden = int(4 * cfg.n_embd * 2 / 3)
        hidden = cfg.multiple_of * ((hidden + cfg.multiple_of - 1) // cfg.multiple_of)
        ffn_single = 3 * cfg.n_embd * hidden
        if cfg.use_moe:
            active_ffn = cfg.n_experts_per_tok * ffn_single + cfg.n_shared_experts * ffn_single
            active_ffn += cfg.n_embd * cfg.n_experts  # gate
        else:
            active_ffn = ffn_single
        per_layer = attn + norms + active_ffn
        return emb + per_layer * cfg.n_layer + cfg.n_embd

    # -------------------------------------------------------------------------
    def forward(self, idx, targets=None, kv_cache=None, start_pos: int = 0):
        B, T = idx.size()
        assert T <= self.config.block_size, \
            f"Sequence length {T} exceeds block_size {self.config.block_size}"

        x = self.transformer.wte(idx)

        total_aux_loss = torch.tensor(0.0, device=idx.device)
        for block in self.transformer.h:
            x, aux_loss = block(x, kv_cache=kv_cache, start_pos=start_pos)
            total_aux_loss = total_aux_loss + aux_loss

        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
            if self.config.use_moe:
                loss = loss + total_aux_loss

        return logits, loss

    # -------------------------------------------------------------------------
    def configure_optimizers(self, weight_decay, learning_rate, device_type, master_process=True):
        # 只对需要梯度更新的参数进行优化器配置，区分权重衰减和非权重衰减参数

        param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad}
        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]

        optim_groups = [
            {'params': decay_params, 'weight_decay': weight_decay},
            {'params': nodecay_params, 'weight_decay': 0.0},
        ]
        if master_process:
            nd = sum(p.numel() for p in decay_params)
            nn_ = sum(p.numel() for p in nodecay_params)
            print(f"  Decayed tensors: {len(decay_params)}, params: {nd:,}")
            print(f"  Non-decayed tensors: {len(nodecay_params)}, params: {nn_:,}")

        # Fused AdamW 的原理是将多个小的内核调用合并为一个大内核，减少内核启动和内存访问开销，提升训练速度。
        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
        use_fused = fused_available and device_type == "cuda"
        if master_process:
            print(f"  Fused AdamW: {use_fused}")

        optimizer = torch.optim.AdamW(
            optim_groups, lr=learning_rate,
            betas=(0.9, 0.95), eps=1e-8, fused=use_fused  # 训练技巧：显存不变，提token吞吐量
        )
        return optimizer

    # -------------------------------------------------------------------------
    def set_gradient_checkpointing(self, enabled: bool):
        """训练/推理切换时动态开关 checkpointing，推理时关闭以节省计算开销"""
        for block in self.transformer.h:
            block.use_checkpoint = enabled

    # -------------------------------------------------------------------------
    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=50, top_p=0.9):
        """
        使用 KV Cache 的高效自回归生成
        """
        B, prompt_len = idx.shape
        total_len = prompt_len + max_new_tokens
        assert total_len <= self.config.block_size

        # 推理模式: 关闭 checkpoint
        self.set_gradient_checkpointing(False)
        self.eval()

        kv_cache = KVCache(
            self.config, B, total_len, idx.device, dtype=torch.bfloat16
        )

        # --- Prefill: 处理整个 prompt ---
        x = self.transformer.wte(idx)
        for block in self.transformer.h:
            x, _ = block(x, kv_cache=kv_cache, start_pos=0)
        kv_cache.advance(prompt_len)

        x = self.transformer.ln_f(x)
        logits = self.lm_head(x[:, -1, :])
        next_token = self._sample(logits, temperature, top_k, top_p)
        generated = [next_token]

        # --- Decode: 逐 token 生成 ---
        for step in range(1, max_new_tokens):
            x = self.transformer.wte(next_token.unsqueeze(1))
            pos = prompt_len + step - 1
            for block in self.transformer.h:
                x, _ = block(x, kv_cache=kv_cache, start_pos=pos)
            kv_cache.advance(1)

            x = self.transformer.ln_f(x)
            logits = self.lm_head(x[:, -1, :])
            next_token = self._sample(logits, temperature, top_k, top_p)
            generated.append(next_token)

        generated = torch.stack(generated, dim=1)
        return torch.cat([idx, generated], dim=1)

    # -------------------------------------------------------------------------
    def _sample1(self, logits, temperature, top_k, top_p):

        # 从 logits 中采样下一个 token，支持 temperature、top-k 和 top-p 策略

        if temperature == 0:
            return logits.argmax(dim=-1)
        logits = logits / temperature
        if top_k > 0:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = -float('Inf')
        if top_p < 1.0:
            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
            cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
            remove = cum_probs > top_p
            remove[..., 1:] = remove[..., :-1].clone()
            remove[..., 0] = False
            indices_to_remove = remove.scatter(1, sorted_idx, remove)
            logits[indices_to_remove] = -float('Inf')
        probs = F.softmax(logits, dim=-1)
        return torch.multinomial(probs, num_samples=1).squeeze(-1)



    def _sample(self, logits, temperature, top_k, top_p):
        # 转为 float32 避免 bfloat16 精度问题
        logits = logits.float()

        if temperature == 0:
            return logits.argmax(dim=-1)
        logits = logits / temperature
        if top_k > 0:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = -float('Inf')
        if top_p < 1.0:
            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
            cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
            remove = cum_probs > top_p
            remove[..., 1:] = remove[..., :-1].clone()
            remove[..., 0] = False
            indices_to_remove = remove.scatter(1, sorted_idx, remove)
            logits[indices_to_remove] = -float('Inf')
        probs = F.softmax(logits, dim=-1)
        # 安全检查：将可能的 nan/inf 替换为 0，避免 multinomial 崩溃
        probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)
        # 如果某行全为 0（极端情况），给均匀分布
        zero_rows = (probs.sum(dim=-1) == 0)
        if zero_rows.any():
            probs[zero_rows] = 1.0 / probs.size(-1)
        return torch.multinomial(probs, num_samples=1).squeeze(-1)