mul / models /subdivider.py
Upload folder using huggingface_hub
90ca518 verified 16 days ago
53.3 kB
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from einops import rearrange
from modules.attention import MLA
from .point_conditioner import PointConditioner

try:
    from flash_attn import flash_attn_varlen_func
except Exception:
    try:
        from flash_attn.flash_attn_interface import flash_attn_varlen_func
    except Exception:
        flash_attn_varlen_func = None

class SinusoidalPosEmb(nn.Module):
    def __init__(self, dim, scale=100.0):
        super().__init__()
        self.dim = dim
        self.scale = scale

    def forward(self, x):
        device = x.device
        dtype = x.dtype
        x = x.float() * self.scale
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = x[:, None] * emb[None, :]
        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
        return emb.to(dtype=dtype)

class RMSNorm(nn.Module):
    """现代 LLM 标配：去均值的 RMSNorm"""
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        # 只需要学习一个缩放参数，去掉 bias
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        # 计算均方根 (Root Mean Square)
        norm_x = torch.mean(x.pow(2), dim=-1, keepdim=True)
        x_normed = x * torch.rsqrt(norm_x + self.eps)
        return self.weight * x_normed

class SwiGLU(nn.Module):
    """现代 LLM 标配：SwiGLU 激活的门控 FFN"""
    def __init__(self, d_model, hidden_dim=None):
        super().__init__()
        # 将原先 4 * d_model 的维度缩小到 8/3，保持参数量对等; hidden_dim 可显式覆盖
        # (如 1024 宽度时 8/3 给 2730 非 256 倍数, 显式取 2816 对齐 Tensor Core)
        hidden_dim = int((8 / 3) * d_model) if hidden_dim is None else int(hidden_dim)
        # 提示：在实际工程中，通常会把 hidden_dim 向上取整到 128 或 256 的倍数，以最大化 Tensor Core 效率
        
        self.w1 = nn.Linear(d_model, hidden_dim, bias=False)
        self.w2 = nn.Linear(d_model, hidden_dim, bias=False)
        self.w3 = nn.Linear(hidden_dim, d_model, bias=False)

    def forward(self, x):
        # 公式: Swish(xW1) ⊗ (xW2) -> W3
        # PyTorch 中 F.silu 就是 Swish
        return self.w3(F.silu(self.w1(x)) * self.w2(x))

class ModernAttention(nn.Module):
    def __init__(self, d_model, nhead, dropout=0., qk_norm=False):
        super().__init__()
        assert d_model % nhead == 0, "d_model 必须能被 nhead 整除"
        self.nhead = nhead
        self.d_head = d_model // nhead

        # 现代架构通常在 Attention 的投影层中也丢弃 bias
        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, d_model, bias=False)
        self.v_proj = nn.Linear(d_model, d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)
        self.dropout = dropout
        # QK-norm: 逐 head 对 Q/K 在 head_dim 上做 RMSNorm(在 RoPE 之前),稳定 attention logits、
        # 防幅度爆炸。默认关(不改变既有模型);flow 等新模型可开。新增参数 -> 旧 ckpt load_strict=False。
        self.qk_norm = qk_norm
        if qk_norm:
            self.q_norm = RMSNorm(self.d_head)
            self.k_norm = RMSNorm(self.d_head)

    def forward(self, q, k, v, mask=None, is_causal=False, rope_apply_fn=None):
        B, L_q, _ = q.size()
        _, L_k, _ = k.size()

        # 投影并调整形状为 [Batch, Head, SeqLen, HeadDim]
        Q = self.q_proj(q).view(B, L_q, self.nhead, self.d_head).transpose(1, 2)
        K = self.k_proj(k).view(B, L_k, self.nhead, self.d_head).transpose(1, 2)
        V = self.v_proj(v).view(B, L_k, self.nhead, self.d_head).transpose(1, 2)

        if self.qk_norm:                                 # QK-norm(RoPE 前),保持 dtype 供 SDPA/flash
            Q = self.q_norm(Q).to(Q.dtype); K = self.k_norm(K).to(K.dtype)

        # 注入旋转位置编码 (RoPE)
        if rope_apply_fn is not None:
            Q, K = rope_apply_fn(Q, K)

        # 使用 PyTorch 2.0+ 内置的 SDPA，它会自动调用 Flash Attention 节省巨量显存
        out = F.scaled_dot_product_attention(
            Q, K, V, 
            attn_mask=mask, 
            dropout_p=self.dropout if self.training else 0.0, 
            is_causal=is_causal
        )
        
        # 恢复形状并输出
        out = out.transpose(1, 2).contiguous().view(B, L_q, -1)
        return self.out_proj(out)

    def forward_varlen(self, q, k, v, lengths_q, lengths_k=None, is_causal=False, rope_apply_fn=None,
                       max_len_q=None, max_len_k=None):
        """max_len_q / max_len_k: flash 需要的 max_seqlen。**传 python int 可避免 GPU->CPU 同步**。
        不传时回退到 `int(lengths.max().item())` —— 那是一次 cudaStreamSynchronize: CPU 必须等
        GPU 把已排队的活干完才能读到值, 于是无法提前排队 kernel。本模型一次前向有 36 次 attention
        调用(18 self + 18 cross), 即 72 次同步, 是「kernel 又多又小、GPU 有空隙」的直接成因,
        也让 torch.compile 在这里 graph break(数据依赖的 unbacked symint)。
        调用方本来就知道这些长度(lengths 是从 [c.numel() for c in coord_list] 建的), 应当传进来。
        """
        if lengths_k is None:
            lengths_k = lengths_q
            if max_len_k is None:
                max_len_k = max_len_q

        total_q = q.size(0)
        total_k = k.size(0)
        Q = self.q_proj(q).view(total_q, self.nhead, self.d_head)
        K = self.k_proj(k).view(total_k, self.nhead, self.d_head)
        V = self.v_proj(v).view(total_k, self.nhead, self.d_head)

        if self.qk_norm:                                 # QK-norm(RoPE 前),.to 保持 dtype 供 flash-attn
            Q = self.q_norm(Q).to(Q.dtype); K = self.k_norm(K).to(K.dtype)

        if rope_apply_fn is not None:
            Q, K = rope_apply_fn(Q, K)

        dropout_p = self.dropout if self.training else 0.0
        if can_use_flash_attn(Q):
            out = flash_attn_varlen_func(
                Q.contiguous(),
                K.contiguous(),
                V.contiguous(),
                lengths_to_cu_seqlens(lengths_q),
                lengths_to_cu_seqlens(lengths_k),
                int(lengths_q.max().item()) if max_len_q is None else int(max_len_q),
                int(lengths_k.max().item()) if max_len_k is None else int(max_len_k),
                dropout_p=dropout_p,
                causal=is_causal,
            )
        else:
            out_chunks = []
            q_start = 0
            k_start = 0
            for q_len, k_len in zip(lengths_q.tolist(), lengths_k.tolist()):
                q_end = q_start + q_len
                k_end = k_start + k_len
                q_i = Q[q_start:q_end].transpose(0, 1).unsqueeze(0)
                k_i = K[k_start:k_end].transpose(0, 1).unsqueeze(0)
                v_i = V[k_start:k_end].transpose(0, 1).unsqueeze(0)
                out_i = F.scaled_dot_product_attention(
                    q_i,
                    k_i,
                    v_i,
                    dropout_p=dropout_p,
                    is_causal=is_causal,
                )
                out_chunks.append(out_i.squeeze(0).transpose(0, 1))
                q_start = q_end
                k_start = k_end
            out = torch.cat(out_chunks, dim=0)

        return self.out_proj(out.contiguous().view(total_q, -1))

# class CrossAttnTransformerBlock(nn.Module):
#     def __init__(self, d_model, nhead, dropout=0.):
#         super().__init__()
#         self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
#         self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)

#         self.norm1 = nn.LayerNorm(d_model)
#         self.norm2 = nn.LayerNorm(d_model)
#         self.norm3 = nn.LayerNorm(d_model)

#         self.ffn = nn.Sequential(
#             nn.Linear(d_model, d_model * 4),
#             nn.GELU(),
#             nn.Dropout(dropout),
#             nn.Linear(d_model * 4, d_model),
#         )
#         self.dropout = nn.Dropout(dropout)

#     def forward(self, x, cond):
#         # 1) token self-attention
#         h = self.norm1(x)
#         sa_out, _ = self.self_attn(h, h, h, need_weights=False)
#         x = x + self.dropout(sa_out)

#         # 2) cross-attention: query from current tokens, key/value from condition tokens
#         h = self.norm2(x)
#         ca_out, _ = self.cross_attn(h, cond, cond, need_weights=False)
#         x = x + self.dropout(ca_out)

#         # 3) feed-forward
#         h = self.norm3(x)
#         ff_out = self.ffn(h)
#         x = x + self.dropout(ff_out)
#         return x
class ModernCrossAttnTransformerBlock(nn.Module):
    def __init__(self, d_model, nhead, dropout=0.):
        super().__init__()
        
        # 替换为支持 Flash Attention 的现代 Attention
        self.self_attn = ModernAttention(d_model, nhead, dropout=dropout)
        self.cross_attn = ModernAttention(d_model, nhead, dropout=dropout)

        # 替换 LayerNorm 为 RMSNorm
        self.norm1 = RMSNorm(d_model)
        self.norm2 = RMSNorm(d_model)
        self.norm3 = RMSNorm(d_model)

        # 替换传统 FFN 为 SwiGLU
        self.ffn = SwiGLU(d_model)
        
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, cond, rope_apply_fn=None, causal_mask=False):
        # 1) token self-attention
        # 你的 Pre-Norm 逻辑保持不变，这是正确的
        h = self.norm1(x)
        # 注意：在这里把 RoPE 传进去处理 Query 和 Key
        sa_out = self.self_attn(h, h, h, is_causal=causal_mask, rope_apply_fn=rope_apply_fn)
        x = x + self.dropout(sa_out)

        # 2) cross-attention: query from current tokens, key/value from condition tokens
        if cond is not None:
            h = self.norm2(x)
            # Cross Attention 一般不使用 RoPE，因为输入源不同
            ca_out = self.cross_attn(h, cond, cond)
            x = x + self.dropout(ca_out)

        # 3) feed-forward (SwiGLU)
        h = self.norm3(x)
        ff_out = self.ffn(h)
        x = x + self.dropout(ff_out)
        
        return x

    def forward_varlen(self, x, lengths, cond=None, cond_lengths=None, rope_apply_fn=None, causal_mask=False):
        h = self.norm1(x)
        sa_out = self.self_attn.forward_varlen(
            h,
            h,
            h,
            lengths,
            is_causal=causal_mask,
            rope_apply_fn=rope_apply_fn,
        )
        x = x + self.dropout(sa_out)

        if cond is not None:
            h = self.norm2(x)
            ca_out = self.cross_attn.forward_varlen(
                h,
                cond,
                cond,
                lengths,
                lengths_k=cond_lengths,
            )
            x = x + self.dropout(ca_out)

        h = self.norm3(x)
        ff_out = self.ffn(h)
        x = x + self.dropout(ff_out)
        return x

def rotate_half(x):
    """
    旋转特征的一半维度。
    对于向量 [x1, x2, ..., x_d/2, x_d/2+1, ..., x_d]，
    将其转换为 [-x_d/2+1, ..., -x_d, x1, x2, ..., x_d/2]。
    这在数学上等价于复数平面上的乘以 i。
    """
    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 和 sin 注入到 Query 和 Key 中。
    q, k 形状期望为: (Batch, Num_Heads, Seq_Len, Head_Dim)
    cos, sin 形状期望为: (1, 1, Seq_Len, Head_Dim)
    """
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

def apply_rotary_pos_emb_varlen(q, k, cos, sin):
    """RoPE。**cos/sin 先转成 q/k 的 dtype**，保证输出 dtype 不被 cos/sin 带跑。
    为什么必须这么写：cos/sin 由 `rotary_emb.forward_positions(..., dtype=x.dtype)` 生成，而
    9 尺度的 rope_fn 是在 `x = raw_emb + refine_emb` 之后立刻构造的 —— `nn.Embedding` 不在 autocast
    的转型名单里，那时 x 还是 **fp32**，于是 cos/sin 是 fp32。bf16 的 Q/K 乘上 fp32 的 cos/sin 会
    类型提升成 fp32，`can_use_flash_attn` 判否，直接掉进 `ModernAttention.forward_varlen` 里
    **SDPA 的逐序列 python 循环**。命中的恰好是 enc_blocks_9 / dec_blocks_9 这 4 层,
    也就是跑在满 token 长度(其余层已被 hourglass 降到 1/3、1/9)上最贵的 4 个注意力。
    3 尺度与 bottleneck 的 rope_fn 因为构造时 x 已是 bf16，本来就走 flash —— 即同一网络里
    4 层 fp32 慢路径 + 14 层 bf16 快路径，不一致。统一到 bf16。
    """
    cos_q, sin_q = cos.to(q.dtype), sin.to(q.dtype)
    q_embed = (q * cos_q) + (rotate_half(q) * sin_q)
    cos_k, sin_k = (cos_q, sin_q) if k.dtype == q.dtype else (cos.to(k.dtype), sin.to(k.dtype))
    k_embed = (k * cos_k) + (rotate_half(k) * sin_k)
    return q_embed, k_embed

def lengths_to_cu_seqlens(lengths):
    lengths = lengths.to(dtype=torch.int32)
    prefix = torch.empty(lengths.numel() + 1, device=lengths.device, dtype=torch.int32)
    prefix[0] = 0
    prefix[1:] = torch.cumsum(lengths, dim=0, dtype=torch.int32)
    return prefix.contiguous()

def lengths_to_position_ids(lengths):
    total = int(lengths.sum().item())
    if total == 0:
        return torch.empty(0, device=lengths.device, dtype=torch.long)
    starts = torch.cumsum(lengths, dim=0) - lengths
    return torch.arange(total, device=lengths.device, dtype=torch.long) - torch.repeat_interleave(starts, lengths)

def can_use_flash_attn(x):
    return (
        flash_attn_varlen_func is not None
        and x.is_cuda
        and x.dtype in (torch.float16, torch.bfloat16)
    )

class RotaryEmbedding(nn.Module):
    """
    预计算和缓存 RoPE 的 cos 和 sin 矩阵。
    """
    def __init__(self, dim, max_position_embeddings=90000, base=10000.0, device=None):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        
        # 计算旋转频率: theta_i = 10000^(-2(i-1)/d)
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # 预先构建缓存
        self._set_cos_sin_cache(seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype())

    def _set_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len
        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)

        # freqs 形状: (seq_len, dim/2)
        freqs = torch.outer(t, self.inv_freq)
        
        # 将频率在特征维度上复制一份，形状变为: (seq_len, dim)
        # 前半部分和后半部分的频率一致，为了和 rotate_half 的逻辑对齐
        emb = torch.cat((freqs, freqs), dim=-1)
        
        # 增加 Batch 和 Head 的广播维度: (1, 1, Seq_Len, Head_Dim)
        self.register_buffer("cos_cached", emb.cos().to(dtype).unsqueeze(0).unsqueeze(0), persistent=False)
        self.register_buffer("sin_cached", emb.sin().to(dtype).unsqueeze(0).unsqueeze(0), persistent=False)

    def forward(self, x, seq_len=None):
        """
        x 形状: (Batch, Num_Heads, Seq_Len, Head_Dim)
        """
        # 如果当前序列长度超过了预计算的缓存，则动态扩展缓存
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

        # 截取当前需要的序列长度返回
        return (
            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
        )

    def forward_positions(self, position_ids, dtype=None):
        if dtype is None:
            dtype = self.cos_cached.dtype
        if position_ids.numel() == 0:
            empty = torch.empty(0, 1, self.dim, device=position_ids.device, dtype=dtype)
            return empty, empty
        seq_len = int(position_ids.max().item()) + 1
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len=seq_len, device=position_ids.device, dtype=dtype)
        cos = self.cos_cached[0, 0, position_ids, :].to(device=position_ids.device, dtype=dtype)
        sin = self.sin_cached[0, 0, position_ids, :].to(device=position_ids.device, dtype=dtype)
        return cos.unsqueeze(1), sin.unsqueeze(1)

class LinearDownsample(nn.Module):
    def __init__(self, dim, shorten_factor):
        super().__init__()
        self.proj = nn.Linear(dim * shorten_factor, dim)
        self.shorten_factor = shorten_factor

    def forward(self, x):
        x = rearrange(x, 'b (n s) d -> b n (s d)', s=self.shorten_factor)
        return self.proj(x)

    def forward_varlen(self, x, lengths):
        if torch.any(lengths % self.shorten_factor != 0):
            raise ValueError(f"All sequence lengths must be divisible by {self.shorten_factor} for hourglass downsample.")
        x = rearrange(x, '(n s) d -> n (s d)', s=self.shorten_factor)
        return self.proj(x), lengths // self.shorten_factor

class LinearUpsample(nn.Module):
    def __init__(self, dim, shorten_factor):
        super().__init__()
        self.proj = nn.Linear(dim, dim * shorten_factor)
        self.shorten_factor = shorten_factor

    def forward(self, x):
        x = self.proj(x)
        return rearrange(x, 'b n (s d) -> b (n s) d', s=self.shorten_factor)

    def forward_varlen(self, x, lengths):
        x = self.proj(x)
        x = rearrange(x, 'n (s d) -> (n s) d', s=self.shorten_factor)
        return x, lengths * self.shorten_factor

class DiscreteDiffusionTransformer(nn.Module):
    def __init__(self, 
                 vocab_A, 
                 vocab_B, 
                 max_seq_len, 
                 d_model=768, 
                 nhead=12, 
                 depth = [4,4,8],
                 dropout=0.,
                 encoder_freeze=True,
                 encoder_name='miche-256-feature',
                 use_pc=True,
                 time_embed_scale=100.0):
        super().__init__()
        self.vocab_A = vocab_A
        self.vocab_B = vocab_B
        self.d_model = d_model
        
        # MASK Token 的 ID，我们将其设为 B（因为真实的 refine_token 是 0 到 B-1）
        self.mask_token_id = vocab_B 
        
        # 1. 词嵌入 (Embeddings)
        self.raw_embedding = nn.Embedding(vocab_A, d_model)
        # refine_token 的词表大小是 B + 1 (包含了 MASK token)
        self.refine_embedding = nn.Embedding(vocab_B + 1, d_model) 
        
        # 2. 位置编码 (Positional Encoding)
        # self.pos_embedding = nn.Embedding(max_seq_len, d_model)
        
        # 3. 时间步编码 (Timestep Embedding)
        self.time_mlp = nn.Sequential(
            SinusoidalPosEmb(d_model, scale=time_embed_scale),
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Linear(d_model * 4, d_model)
        )
        
        # 4. 主干网络 (Denoising Backbone)
        # U-Net/Hourglass 主干:
        # enc(9) -> down -> enc(3) -> down -> bottleneck(1) -> up+skip -> dec(3) -> up+skip -> dec(9)
        self.depth = depth

        self.enc_blocks_9 = nn.ModuleList(
            [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[0])]
        )
        self.enc_blocks_3 = nn.ModuleList(
            [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[1])]
        )
        self.bottleneck_blocks = nn.ModuleList(
            [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[2])]
        )
        self.dec_blocks_3 = nn.ModuleList(
            [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[1])]
        )
        self.dec_blocks_9 = nn.ModuleList(
            [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[0])]
        )

        self.down_9_to_3 = LinearDownsample(d_model, shorten_factor=3)
        self.down_3_to_1 = LinearDownsample(d_model, shorten_factor=3)
        self.up_1_to_3 = LinearUpsample(d_model, shorten_factor=3)
        self.up_3_to_9 = LinearUpsample(d_model, shorten_factor=3)

        self.rotary_emb = RotaryEmbedding(dim=d_model//nhead, max_position_embeddings=max_seq_len)
        # 5. 输出预测头 (只预测真实的 B 个 Token)
        self.lm_head = nn.Linear(d_model, vocab_B)
        self.use_pc = use_pc
        if use_pc:
            print(f'Point cloud encoder: {encoder_name} | freeze: {encoder_freeze}')
            self.conditioner = PointConditioner(model_name=encoder_name, freeze=encoder_freeze, feature_dim=d_model)
            # Learnable null condition for classifier-free guidance.
            # Replaces self.conditioner(pc) when pc is None, so the model can learn
            # an "unconditional" representation during cond-drop fine-tuning.
            self.null_cond_len = 257  # matches PointConditioner.cond_length
            self.null_cond = nn.Parameter(torch.randn(self.null_cond_len, d_model) * 0.02)

    def forward(self, raw_tokens, noisy_refine_tokens, t, pc=None):
        """
        raw_tokens: [b, n], 取值 0~A-1
        noisy_refine_tokens: [b, n], 取值 0~B (B是MASK)
        t: [b], 当前的时间步
        pc: [b, n, 6], 点云数据
        """
        if isinstance(raw_tokens, (list, tuple)):
            return self.forward_varlen(raw_tokens, noisy_refine_tokens, t, pc)

        b, n = raw_tokens.shape
        device = raw_tokens.device
        
        # 获取 Embedding
        raw_emb = self.raw_embedding(raw_tokens)               # [b, n, d_model]
        refine_emb = self.refine_embedding(noisy_refine_tokens) # [b, n, d_model]
        
        # 获取位置和时间 Embedding
        time_emb = self.time_mlp(t).unsqueeze(1)               # [b, 1, d_model]
        
        # x 作为被去噪序列，cond 作为每层 cross-attn 的条件序列
        x = raw_emb + refine_emb + time_emb

        if self.use_pc:
            if pc is None:
                # Classifier-free guidance / cond-drop: use learnable null condition.
                cond = self.null_cond.unsqueeze(0).expand(b, -1, -1)
            else:
                cond = self.conditioner(pc) # [b, 257, d_model]
        else:
            cond = None

        if n % 9 != 0:
            raise ValueError(f"Sequence length {n} is not compatible with 9->3->1 hourglass compression.")

        dummy_tensor = torch.empty(1, device=x.device, dtype=torch.float32)

        def rope_fn_for(seq_len):
            cos, sin = self.rotary_emb(dummy_tensor, seq_len=seq_len)
            return lambda q, k: apply_rotary_pos_emb(q, k, cos, sin)

        # Encoder @ length 9
        rope_fn_9 = rope_fn_for(x.shape[1])
        for block in self.enc_blocks_9:
            x = block(x, cond, rope_apply_fn=rope_fn_9)
        skip_9 = x

        # Downsample 9 -> 3
        x = self.down_9_to_3(x)

        # Encoder @ length 3
        rope_fn_3 = rope_fn_for(x.shape[1])
        for block in self.enc_blocks_3:
            x = block(x, cond, rope_apply_fn=rope_fn_3)
        skip_3 = x

        # Downsample 3 -> 1 (or in general: n/3 -> n/9)
        x = self.down_3_to_1(x)

        # Bottleneck
        rope_fn_1 = rope_fn_for(x.shape[1])
        for block in self.bottleneck_blocks:
            x = block(x, cond, rope_apply_fn=rope_fn_1)

        # Upsample 1 -> 3 and long skip connection
        x = self.up_1_to_3(x)
        x = x + skip_3

        # Decoder @ length 3
        for block in self.dec_blocks_3:
            x = block(x, cond, rope_apply_fn=rope_fn_3)

        # Upsample 3 -> 9 and long skip connection
        x = self.up_3_to_9(x)
        x = x + skip_9

        # Decoder @ length 9
        for block in self.dec_blocks_9:
            x = block(x, cond, rope_apply_fn=rope_fn_9)
        
        # 预测 Logits
        logits = self.lm_head(x) # [b, n, vocab_B]
        return logits

    def forward_varlen(self, raw_tokens, noisy_refine_tokens, t, pc=None):
        device = t.device
        raw_tokens = [seq.to(device=device, dtype=torch.long) for seq in raw_tokens]
        noisy_refine_tokens = [seq.to(device=device, dtype=torch.long) for seq in noisy_refine_tokens]
        lengths = torch.tensor([seq.numel() for seq in raw_tokens], device=device, dtype=torch.long)
        if len(raw_tokens) == 0 or int(lengths.sum().item()) == 0:
            raise ValueError("forward_varlen received an empty batch.")
        if any(a.numel() != b.numel() for a, b in zip(raw_tokens, noisy_refine_tokens)):
            raise ValueError("raw token lengths and refine token lengths must match per sample.")
        if torch.any(lengths % 9 != 0):
            bad = lengths[lengths % 9 != 0].detach().cpu().tolist()
            raise ValueError(f"Varlen hourglass requires each sequence length to be divisible by 9, got {bad}.")

        raw_packed = torch.cat(raw_tokens, dim=0)
        refine_packed = torch.cat(noisy_refine_tokens, dim=0)
        raw_emb = self.raw_embedding(raw_packed)
        refine_emb = self.refine_embedding(refine_packed)
        time_emb = self.time_mlp(t).repeat_interleave(lengths, dim=0)
        x = raw_emb + refine_emb + time_emb

        if self.use_pc:
            if pc is None:
                # Null condition path for CFG / cond-drop. Replicate the learned
                # null_cond per sample and flatten to match the varlen layout.
                bsz = len(raw_tokens)
                null = self.null_cond  # [257, d_model]
                cond = null.unsqueeze(0).expand(bsz, -1, -1).reshape(-1, null.size(-1))
                cond_lengths = torch.full((bsz,), self.null_cond_len, device=device, dtype=torch.long)
            else:
                pc = pc.to(device=device)
                cond = self.conditioner(pc)
                cond_lengths = torch.full((cond.size(0),), cond.size(1), device=device, dtype=torch.long)
                cond = cond.reshape(-1, cond.size(-1))
        else:
            cond = None
            cond_lengths = None

        def rope_fn_for(cur_lengths):
            position_ids = lengths_to_position_ids(cur_lengths)
            cos, sin = self.rotary_emb.forward_positions(position_ids, dtype=x.dtype)
            return lambda q, k: apply_rotary_pos_emb_varlen(q, k, cos, sin)

        rope_fn_9 = rope_fn_for(lengths)
        for block in self.enc_blocks_9:
            x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_9)
        skip_9 = x
        lengths_9 = lengths

        x, lengths = self.down_9_to_3.forward_varlen(x, lengths)
        rope_fn_3 = rope_fn_for(lengths)
        for block in self.enc_blocks_3:
            x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_3)
        skip_3 = x
        lengths_3 = lengths

        x, lengths = self.down_3_to_1.forward_varlen(x, lengths)
        rope_fn_1 = rope_fn_for(lengths)
        for block in self.bottleneck_blocks:
            x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_1)

        x, lengths = self.up_1_to_3.forward_varlen(x, lengths)
        if not torch.equal(lengths, lengths_3):
            raise RuntimeError("Internal length mismatch after 1->3 upsample.")
        x = x + skip_3
        for block in self.dec_blocks_3:
            x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_3)

        x, lengths = self.up_3_to_9.forward_varlen(x, lengths)
        if not torch.equal(lengths, lengths_9):
            raise RuntimeError("Internal length mismatch after 3->9 upsample.")
        x = x + skip_9
        for block in self.dec_blocks_9:
            x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_9)

        return self.lm_head(x)

def cosine_schedule(t, T):
    """
    余弦调度函数。
    t: 当前进度步数 (0 表示起始/全噪，T 表示结束/真实数据)
    返回: 需要 mask 的比例 (t=0时返回1.0，t=T时返回0.0)
    """
    # 确保 t 是浮点数以进行安全除法
    return torch.cos((t / T) * (math.pi / 2))

    return torch.cos((t / T) * (math.pi / 2))


def precompute_topology_tensors(raw_faces, tokens_per_face=9):
    """
    预处理：完全向量化的拓扑索引构建（在去噪循环外执行一次即可）
    返回的是一维的张量，方便在 GPU 上做极速的 Scatter/Gather 计算
    """
    B, num_faces, _, _ = raw_faces.shape
    device = raw_faces.device

    # 转为 Python list 加速哈希构建
    raw_faces_list = raw_faces.cpu().tolist()

    token_indices_list = []
    v_indices_list = []
    v_shared_count = 0

    for b in range(B):
        vertex_dict = {}
        for f in range(num_faces):
            for v_idx in range(3):
                coord = tuple(raw_faces_list[b][f][v_idx])

                # 计算展平后的【全局绝对 Token 索引】
                global_token_idx = b * (num_faces * tokens_per_face) + f * tokens_per_face + v_idx

                if coord not in vertex_dict:
                    vertex_dict[coord] = []
                vertex_dict[coord].append(global_token_idx)

        # 仅保留被共享的顶点
        for indices in vertex_dict.values():
            if len(indices) > 1:
                token_indices_list.extend(indices)
                v_indices_list.extend([v_shared_count] * len(indices)) # 记录属于哪个全局共享顶点 ID
                v_shared_count += 1

    # 如果没有任何共享顶点，直接返回空
    if len(token_indices_list) == 0:
        return None, None, None, 0

    # [M] 个共享 Token 的绝对索引
    token_idx = torch.tensor(token_indices_list, dtype=torch.long, device=device)
    # [M] 每个 Token 对应的全局顶点 ID
    v_idx = torch.tensor(v_indices_list, dtype=torch.long, device=device)
    # [M] 每个 Token 对应的绝对面片(Face)索引
    face_idx = token_idx // tokens_per_face

    return token_idx, v_idx, face_idx, v_shared_count

def sample_refine_tokens_with_topology(
    model,
    mask_token_id,
    raw_tokens,
    pc,
    T_infer=10,
    T_train=100,
    tokens_per_face=9,
    use_gumbel=False,
    gumbel_temp_init=1.0,
    use_soft_consensus=False,
    consensus_start_step_ratio=0.0,
    cfg_scale=0.0,
    null_pc=None,
    diag_stats=None,
    forbidden_token_ids=(0,),
    forbid_mask_token=True,
    empty_token_ids=(),
    empty_token_bias=0.0,
):
    """
    极致优化的拓扑一致性采样代码（零 Python Loop）
    empty_token_ids / empty_token_bias: 对"空/背景"占用 token(实测 258,770 各~33%)的 logits
      减去 empty_token_bias(负偏置)。>0 -> 空 token 更不易被选 -> 更密的细分(更多面)。默认 no-op。
    优化 flag:
      use_gumbel               : 重 mask 时给 confidence 加温度退火的 Gumbel 噪声 (MaskGIT 标配)
      gumbel_temp_init         : 初始温度，随 step 线性退火到 0
      use_soft_consensus       : 共享顶点用概率分布聚合代替硬多数表决
      consensus_start_step_ratio: 在 step / T_infer 达到该值前不做 topology consensus
      cfg_scale                : Classifier-free guidance 强度. 0 = 不开启 (单 forward).
                                 >0 时每步多一次 forward (null_pc 作 uncond).
                                 logits = logits_c + cfg_scale * (logits_c - logits_u)
      null_pc                  : (B, N, C) 作为 uncond 输入. 若为 None, 用 zeros_like(pc).
                                 注: 当前模型未训练 cond-drop, 用 null pc 是 OOD 行为, 效果不保证.
      diag_stats               : 可选 dict, 把每步的诊断信息追加进去 (供测试用)
      forbidden_token_ids      : 禁止模型"决定"采样到的 token id 集合. 这些 id 的 logits
                                 在 softmax 前被设为 -inf, 后续 argmax / consensus 全部
                                 自动避开. 默认 (0,) 即 pad/EOS——堵住"模型中段吐出 0
                                 -> save_stage_prediction 截断 -> mesh 只剩前半段"这条 bug.
      forbid_mask_token        : 同时也禁止预测 mask_token_id (默认 True). mask 是占位符,
                                 模型不应把它当最终答案.
    """
    model.eval()
    b, n = raw_tokens.shape
    device = raw_tokens.device
    num_faces = n // tokens_per_face

    # 收集要屏蔽的 token id; 真正裁剪并搬到 GPU 留到拿到 logits 后再做
    # (mask_token_id 通常等于 vocab_B, 不在输出 logits 的维度内, 直接索引会越界)
    _forbidden_raw = list(forbidden_token_ids) if forbidden_token_ids is not None else []
    if forbid_mask_token:
        _forbidden_raw.append(int(mask_token_id))
    _forbidden_raw = sorted(set(int(x) for x in _forbidden_raw))
    forbidden_t = None  # 在首次 forward 后用 vocab_size 过滤构建

    # --- 1. 预处理 (仅执行一次，耗时极短) ---
    raw_faces = raw_tokens.view(b, num_faces, 3, 3)
    token_idx, v_idx, face_idx, v_shared_count = precompute_topology_tensors(raw_faces, tokens_per_face)

    # 初始化
    current_tokens = torch.full((b, n), mask_token_id, dtype=torch.long, device=device)
    unmasked_flags = torch.zeros((b, n), dtype=torch.bool, device=device)
    vocab_size = None # 将在第一次 forward 后获取

    def cosine_schedule(t, max_val):
        return torch.cos(t * torch.pi * 0.5) * max_val

    # 预备 null_pc (CFG)
    # 三种模式:
    #   - null_pc 是 tensor : 用该 tensor 作为 uncond 输入(老模型可用 zeros_like(pc))
    #   - null_pc == "learned" : 调用 model(..., pc=None), 走模型的 self.null_cond 路径
    #                            (只有 cond-drop fine-tune 过的模型才会有意义)
    #   - null_pc is None (默认): 用 zeros_like(pc), 行为同老版
    use_learned_null = (isinstance(null_pc, str) and null_pc.lower() == "learned")
    if cfg_scale > 0.0 and not use_learned_null:
        if null_pc is None:
            null_pc = torch.zeros_like(pc)
        null_pc = null_pc.to(device=device, dtype=pc.dtype)

    for step in range(T_infer):
        current_progress_ratio = step / T_infer
        t_model = torch.full((b,), current_progress_ratio, device=device, dtype=torch.float)

        # [1] 前向预测 (可选 CFG: 双 forward 线性外推)
        logits_c = model(raw_tokens, current_tokens, t_model, pc) # [b, n, vocab_size]
        if cfg_scale > 0.0:
            if use_learned_null:
                logits_u = model(raw_tokens, current_tokens, t_model, None)
            else:
                logits_u = model(raw_tokens, current_tokens, t_model, null_pc)
            logits = logits_c + cfg_scale * (logits_c - logits_u)
        else:
            logits = logits_c
        if vocab_size is None:
            vocab_size = logits.shape[-1]
            # 现在知道真实输出 vocab 大小, 把越界的 id (例如 mask_token_id=vocab_B)
            # 剔除掉后再构建 forbidden_t. 否则 logits[..., id] 会触发 CUDA assert.
            _valid = [x for x in _forbidden_raw if 0 <= x < vocab_size]
            if len(_valid) > 0:
                forbidden_t = torch.tensor(_valid, dtype=torch.long, device=device)

        # 禁止采样到 pad / EOS. softmax 前置 -inf, 后续 argmax 与
        # consensus 聚合 (软/硬) 都因 probs[forbidden]=0 自动避开.
        if forbidden_t is not None:
            logits[..., forbidden_t] = float('-inf')

        # 空行/空占用 token 负偏置: 减小其 logits -> 更不易被选 -> 更密的细分
        if empty_token_bias != 0.0 and len(empty_token_ids) > 0:
            _et = torch.tensor([int(i) for i in empty_token_ids if 0 <= int(i) < vocab_size],
                               dtype=torch.long, device=device)
            if _et.numel() > 0:
                logits[..., _et] = logits[..., _et] - float(empty_token_bias)

        probs = F.softmax(logits, dim=-1)
        max_probs, pred_tokens = probs.max(dim=-1) # [b, n]

        force_mask = torch.zeros((b, n), dtype=torch.bool, device=device)
        target_vals = None  # 最后一步若开启 consensus 会用到

        # ---------------------------------------------------------
        # [2] 拓扑共识 (硬多数表决 / 软概率聚合)
        # 通过 consensus_start_step_ratio 控制何时开启
        # ---------------------------------------------------------
        apply_consensus = (v_shared_count > 0) and (current_progress_ratio >= consensus_start_step_ratio)
        if apply_consensus:
            if use_soft_consensus:
                # 软共识: 直接 scatter_add 共享 token 的概率分布
                shared_probs = probs.view(-1, vocab_size)[token_idx]  # [M, V]
                agg_probs = torch.zeros((v_shared_count, vocab_size), dtype=shared_probs.dtype, device=device)
                v_idx_expanded = v_idx.unsqueeze(1).expand(-1, vocab_size)
                agg_probs.scatter_add_(0, v_idx_expanded, shared_probs)

                mode_vals = agg_probs.argmax(dim=-1)
                target_vals = mode_vals[v_idx]

                preds_flat = pred_tokens.view(-1)
                shared_vals = preds_flat[token_idx]
                is_disagree = (shared_vals != target_vals)
            else:
                # 硬多数表决
                preds_flat = pred_tokens.view(-1)
                shared_vals = preds_flat[token_idx]

                one_hot_vals = F.one_hot(shared_vals, num_classes=vocab_size) # [M, V]
                votes = torch.zeros((v_shared_count, vocab_size), dtype=torch.long, device=device)
                v_idx_expanded = v_idx.unsqueeze(1).expand(-1, vocab_size)
                votes.scatter_add_(0, v_idx_expanded, one_hot_vals)

                mode_vals = votes.argmax(dim=-1)
                target_vals = mode_vals[v_idx]
                is_disagree = (shared_vals != target_vals)

            if is_disagree.any():
                disagree_faces = face_idx[is_disagree]
                force_mask_faces = torch.zeros(b * num_faces, dtype=torch.bool, device=device)
                force_mask_faces[disagree_faces] = True
                force_mask_flat = force_mask_faces.repeat_interleave(tokens_per_face)
                force_mask = force_mask_flat.view(b, n)

        # ---------------------------------------------------------
        # [3] 置信度干预与 Mask 更新
        # ---------------------------------------------------------
        # 注意: 这里先克隆，避免 Gumbel 噪声破坏掉之后我们对 unmasked/force 的 ±inf 标记
        confidence = max_probs.clone()
        if use_gumbel:
            # 温度随进度线性退火到 0
            temperature = gumbel_temp_init * (1.0 - current_progress_ratio)
            if temperature > 0:
                uniform = torch.rand_like(confidence).clamp_(1e-9, 1.0 - 1e-9)
                gumbel = -torch.log(-torch.log(uniform))
                confidence = confidence + temperature * gumbel

        # 已确认的位置不应被再 mask
        confidence[unmasked_flags] = float('inf')

        if apply_consensus:
            confidence[force_mask] = -float('inf')
            unmasked_flags = unmasked_flags & (~force_mask)

        # 诊断信息
        if diag_stats is not None:
            diag_stats.setdefault('force_count', []).append(int(force_mask.sum().item()))
            diag_stats.setdefault('avg_conf', []).append(float(max_probs.mean().item()))

        if step == T_infer - 1:
            # 最后一步：用共识算出的最终值硬覆盖所有可能残余的不一致
            if apply_consensus and target_vals is not None:
                preds_flat = pred_tokens.view(-1)
                preds_flat[token_idx] = target_vals
                pred_tokens = preds_flat.view(b, n)

            current_tokens = torch.where(unmasked_flags, current_tokens, pred_tokens)

        else:
            next_progress_ratio = (step + 1) / T_infer
            ratio = cosine_schedule(torch.tensor(next_progress_ratio), 1.0).item()
            num_to_mask = int(ratio * n)

            max_forced_in_batch = int(force_mask.sum(dim=-1).max().item()) if apply_consensus else 0
            actual_k = min(max(num_to_mask, max_forced_in_batch), n)

            _, mask_indices = torch.topk(-confidence, k=actual_k, dim=-1)

            new_mask = torch.zeros_like(unmasked_flags)
            new_mask.scatter_(1, mask_indices, True)

            # 兜底: force_mask 必入 new_mask (Gumbel 噪声不会让 -inf 漏选, 但保留以防边界条件)
            new_mask = new_mask | force_mask

            current_tokens = torch.where(new_mask, mask_token_id, pred_tokens)
            unmasked_flags = ~new_mask

    return current_tokens


def precompute_topology_tensors_fm(raw_faces, tokens_per_face=9):
    """precompute_topology_tensors 的扩展版: 除了 F(refine 位置 0/1/2 = 3 个顶点/父的占用,
    按顶点 coord 分组), 额外把 intra(M_inner, refine 位置 3/4/5, 分别对应父 0/1/2 的内部连接
    矩阵)也按其父 coord 分组。F 与 intra 用不同"通道标签"分开归组, 绝不互相表决。
    token 布局(与 fix3 同构, 见 tokenizer decode 的 chunk[9..17]):
      refine[0,1,2] = F0,F1,F2(父 0/1/2 占用)   -> 按父 coord 分组(每父冗余, 应一致)
      refine[3,4,5] = M00,M11,M22(父 0/1/2 intra)-> 按父 coord 分组(每父冗余, 应一致)
      refine[6,7,8] = M01,M02,M12(inter, 每面本地)-> 不分组、不表决
    返回 (token_idx, v_idx, face_idx, v_shared_count), 结构与原版完全一致, 可直接喂给
    采样循环(F 组的众数落在 F 值域, intra 组的众数落在 intra 值域, 天然不串)。
    """
    B, num_faces, _, _ = raw_faces.shape
    device = raw_faces.device
    raw_faces_list = raw_faces.cpu().tolist()

    # 每个"冗余通道" = (refine 位置, 该位置对应的父编号, 通道标签)。
    # F: 位置 p -> 父 p; intra: 位置 3+p -> 父 p。inter(6/7/8)不在其中。
    channels = [(0, 0, "F"), (1, 1, "F"), (2, 2, "F"),
                (3, 0, "I"), (4, 1, "I"), (5, 2, "I")]

    token_indices_list = []
    v_indices_list = []
    v_shared_count = 0
    for b in range(B):
        group_dict = {}
        for f in range(num_faces):
            for pos, parent, tag in channels:
                coord = tuple(raw_faces_list[b][f][parent])   # 该父的坐标
                key = (tag, coord)                            # F / intra 分开归组
                global_token_idx = b * (num_faces * tokens_per_face) + f * tokens_per_face + pos
                group_dict.setdefault(key, []).append(global_token_idx)
        for indices in group_dict.values():
            if len(indices) > 1:                              # 仅保留被共享(>1 面引用)的
                token_indices_list.extend(indices)
                v_indices_list.extend([v_shared_count] * len(indices))
                v_shared_count += 1

    if len(token_indices_list) == 0:
        return None, None, None, 0
    token_idx = torch.tensor(token_indices_list, dtype=torch.long, device=device)
    v_idx = torch.tensor(v_indices_list, dtype=torch.long, device=device)
    face_idx = token_idx // tokens_per_face
    return token_idx, v_idx, face_idx, v_shared_count


def sample_refine_tokens_with_topology_fm(
    model,
    mask_token_id,
    raw_tokens,
    pc,
    T_infer=10,
    T_train=100,
    tokens_per_face=9,
    use_gumbel=False,
    gumbel_temp_init=1.0,
    use_soft_consensus=False,
    consensus_start_step_ratio=0.0,
    cfg_scale=0.0,
    null_pc=None,
    diag_stats=None,
    forbidden_token_ids=(0,),
    forbid_mask_token=True,
):
    """sample_refine_tokens_with_topology 的"F + M_inner(intra)双对齐"版。
    与原版唯一区别: 拓扑分组用 precompute_topology_tensors_fm ——
      原版只把 refine 位置 0/1/2(F, 顶点占用)按共享顶点分组表决;
      本版额外把 3/4/5(intra, 每父内部连接矩阵)也按共享父分组表决。
    因此每步的共识检测/末步硬覆盖会同时作用在 F 与 intra 上, 任一不一致的面整块 9 token
    重掩重预测。inter(6/7/8)每面本地、无冗余, 仍不表决(但会随所在面被连坐重掩)。
    其余采样逻辑(CFG 双前向、forbidden、Gumbel、cosine 揭示)与原版逐行一致。
    参数含义见 sample_refine_tokens_with_topology。
    """
    model.eval()
    b, n = raw_tokens.shape
    device = raw_tokens.device
    num_faces = n // tokens_per_face

    _forbidden_raw = list(forbidden_token_ids) if forbidden_token_ids is not None else []
    if forbid_mask_token:
        _forbidden_raw.append(int(mask_token_id))
    _forbidden_raw = sorted(set(int(x) for x in _forbidden_raw))
    forbidden_t = None

    # --- 1. 预处理: F + intra 双通道拓扑分组(唯一改动) ---
    raw_faces = raw_tokens.view(b, num_faces, 3, 3)
    token_idx, v_idx, face_idx, v_shared_count = precompute_topology_tensors_fm(raw_faces, tokens_per_face)

    current_tokens = torch.full((b, n), mask_token_id, dtype=torch.long, device=device)
    unmasked_flags = torch.zeros((b, n), dtype=torch.bool, device=device)
    vocab_size = None

    def cosine_schedule(t, max_val):
        return torch.cos(t * torch.pi * 0.5) * max_val

    use_learned_null = (isinstance(null_pc, str) and null_pc.lower() == "learned")
    if cfg_scale > 0.0 and not use_learned_null:
        if null_pc is None:
            null_pc = torch.zeros_like(pc)
        null_pc = null_pc.to(device=device, dtype=pc.dtype)

    for step in range(T_infer):
        current_progress_ratio = step / T_infer
        t_model = torch.full((b,), current_progress_ratio, device=device, dtype=torch.float)

        # [1] 前向预测 (可选 CFG)
        logits_c = model(raw_tokens, current_tokens, t_model, pc)
        if cfg_scale > 0.0:
            if use_learned_null:
                logits_u = model(raw_tokens, current_tokens, t_model, None)
            else:
                logits_u = model(raw_tokens, current_tokens, t_model, null_pc)
            logits = logits_c + cfg_scale * (logits_c - logits_u)
        else:
            logits = logits_c
        if vocab_size is None:
            vocab_size = logits.shape[-1]
            _valid = [x for x in _forbidden_raw if 0 <= x < vocab_size]
            if len(_valid) > 0:
                forbidden_t = torch.tensor(_valid, dtype=torch.long, device=device)

        if forbidden_t is not None:
            logits[..., forbidden_t] = float('-inf')

        probs = F.softmax(logits, dim=-1)
        max_probs, pred_tokens = probs.max(dim=-1)

        force_mask = torch.zeros((b, n), dtype=torch.bool, device=device)
        target_vals = None

        # [2] 拓扑共识 (F + intra 同时, 因两者在 v_idx 里是不同组, 自动分别表决)
        apply_consensus = (v_shared_count > 0) and (current_progress_ratio >= consensus_start_step_ratio)
        if apply_consensus:
            if use_soft_consensus:
                shared_probs = probs.view(-1, vocab_size)[token_idx]
                agg_probs = torch.zeros((v_shared_count, vocab_size), dtype=shared_probs.dtype, device=device)
                v_idx_expanded = v_idx.unsqueeze(1).expand(-1, vocab_size)
                agg_probs.scatter_add_(0, v_idx_expanded, shared_probs)
                mode_vals = agg_probs.argmax(dim=-1)
                target_vals = mode_vals[v_idx]
                preds_flat = pred_tokens.view(-1)
                shared_vals = preds_flat[token_idx]
                is_disagree = (shared_vals != target_vals)
            else:
                preds_flat = pred_tokens.view(-1)
                shared_vals = preds_flat[token_idx]
                one_hot_vals = F.one_hot(shared_vals, num_classes=vocab_size)
                votes = torch.zeros((v_shared_count, vocab_size), dtype=torch.long, device=device)
                v_idx_expanded = v_idx.unsqueeze(1).expand(-1, vocab_size)
                votes.scatter_add_(0, v_idx_expanded, one_hot_vals)
                mode_vals = votes.argmax(dim=-1)
                target_vals = mode_vals[v_idx]
                is_disagree = (shared_vals != target_vals)

            if is_disagree.any():
                disagree_faces = face_idx[is_disagree]
                force_mask_faces = torch.zeros(b * num_faces, dtype=torch.bool, device=device)
                force_mask_faces[disagree_faces] = True
                force_mask_flat = force_mask_faces.repeat_interleave(tokens_per_face)
                force_mask = force_mask_flat.view(b, n)

        # [3] 置信度干预与 Mask 更新
        confidence = max_probs.clone()
        if use_gumbel:
            temperature = gumbel_temp_init * (1.0 - current_progress_ratio)
            if temperature > 0:
                uniform = torch.rand_like(confidence).clamp_(1e-9, 1.0 - 1e-9)
                gumbel = -torch.log(-torch.log(uniform))
                confidence = confidence + temperature * gumbel

        confidence[unmasked_flags] = float('inf')
        if apply_consensus:
            confidence[force_mask] = -float('inf')
            unmasked_flags = unmasked_flags & (~force_mask)

        if diag_stats is not None:
            diag_stats.setdefault('force_count', []).append(int(force_mask.sum().item()))
            diag_stats.setdefault('avg_conf', []).append(float(max_probs.mean().item()))

        if step == T_infer - 1:
            if apply_consensus and target_vals is not None:
                preds_flat = pred_tokens.view(-1)
                preds_flat[token_idx] = target_vals
                pred_tokens = preds_flat.view(b, n)
            current_tokens = torch.where(unmasked_flags, current_tokens, pred_tokens)
        else:
            next_progress_ratio = (step + 1) / T_infer
            ratio = cosine_schedule(torch.tensor(next_progress_ratio), 1.0).item()
            num_to_mask = int(ratio * n)
            max_forced_in_batch = int(force_mask.sum(dim=-1).max().item()) if apply_consensus else 0
            actual_k = min(max(num_to_mask, max_forced_in_batch), n)
            _, mask_indices = torch.topk(-confidence, k=actual_k, dim=-1)
            new_mask = torch.zeros_like(unmasked_flags)
            new_mask.scatter_(1, mask_indices, True)
            new_mask = new_mask | force_mask
            current_tokens = torch.where(new_mask, mask_token_id, pred_tokens)
            unmasked_flags = ~new_mask

    return current_tokens




def precompute_topology_tensors_fmi(raw_faces, tokens_per_face=9):
    """F + intra + inter 三通道拓扑分组(逐步共识用)。
      F    (refine 0/1/2): 按单父 coord 分组(每父冗余, 引用同父的面应一致)
      intra(refine 3/4/5): 按单父 coord 分组(每父冗余)
      inter(refine 6/7/8): 按"无序父对 coord"分组 —— 两个面若恰好共享 2 顶点(=共享一条粗边)
                           就共享同一父对, 其 inter 矩阵应一致; tokenizer 每面 3 父排序+按
                           (小,大)取 inter, 同一父对在任何面里 token 值相同(朝向天然规范)。
    三通道用不同 key 前缀("F"/"I"/"O")互不表决。返回结构与 precompute_topology_tensors 完全一致,
    可直接喂采样循环(各通道众数落各自值域, 不串)。
    """
    B, num_faces, _, _ = raw_faces.shape
    device = raw_faces.device
    raw_faces_list = raw_faces.cpu().tolist()
    single_ch = [(0, 0, "F"), (1, 1, "F"), (2, 2, "F"), (3, 0, "I"), (4, 1, "I"), (5, 2, "I")]
    pair_ch = [(6, 0, 1), (7, 0, 2), (8, 1, 2)]   # inter 位置 -> 父对
    token_indices_list, v_indices_list, v_shared_count = [], [], 0
    for b in range(B):
        gd = {}
        for f in range(num_faces):
            fc = raw_faces_list[b][f]
            base = b * (num_faces * tokens_per_face) + f * tokens_per_face
            for pos, parent, tag in single_ch:
                gd.setdefault((tag, tuple(fc[parent])), []).append(base + pos)
            for pos, a, bp in pair_ch:
                ca, cb = tuple(fc[a]), tuple(fc[bp])
                pair = (ca, cb) if ca <= cb else (cb, ca)
                gd.setdefault(("O", pair), []).append(base + pos)
        for indices in gd.values():
            if len(indices) > 1:
                token_indices_list.extend(indices)
                v_indices_list.extend([v_shared_count] * len(indices))
                v_shared_count += 1
    if len(token_indices_list) == 0:
        return None, None, None, 0
    token_idx = torch.tensor(token_indices_list, dtype=torch.long, device=device)
    v_idx = torch.tensor(v_indices_list, dtype=torch.long, device=device)
    face_idx = token_idx // tokens_per_face
    return token_idx, v_idx, face_idx, v_shared_count