DeepInfer/deepinfer/model_executor/models/vision_encoder/vit.py

"""deepseek_vit_opt ViT：结构、权重与推理前向。

state_dict 键经 hugging_face/names.py 映射到 HF ckpt 的 vision.*。
"""

from __future__ import annotations

from typing import Any

import deep_gemm
import torch
import torch.nn.functional as F
from flash_attn.cute.interface import _flash_attn_fwd
from torch import nn

from deep_cluster.hardware import accelerator_kind, is_sm_major_10
from deepinfer.model_executor.layers.tilelang_ops.norm import norm
from deepinfer.timer import cpu_timer
from deepinfer.tk_compat.quant import swiglu_forward
from deepinfer.tk_compat.transform import apply_rotary
from deepinfer.tk_compat.utils import dtype_to_str

# ViT 内全部 RMSNorm 的 eps
NORM_EPS = 1e-6

# V41 唯一在用的 vision encoder（deepseek_vit_opt）结构，不随 cfg.vision 变化
PATCH_SIZE = 14
EMBED_DIM = 1024
DEPTH = 32
NUM_HEADS = 16
MLP_RATIO = 2.7
MULTIPLE_OF = 128
# rope 表按最大网格边长预生成；一张图的 patch 网格宽高都不会超过它
MAX_GRID_SIZE = 4096


def _linear(x: torch.Tensor, w: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor:
    """bf16 GEMM：bias 走 addmm（cublas）；无 bias 走 deep_gemm（输出行 stride 须 16B 对齐）。"""
    if bias is not None:
        return torch.addmm(bias, x, w.T)
    n = w.shape[0]
    if (is_sm_major_10() or accelerator_kind() == "ascend") and x.dtype == torch.bfloat16 and w.dtype == torch.bfloat16 and n > 1:
        alignment = 16 // x.dtype.itemsize
        physical_n = (n + alignment - 1) // alignment * alignment
        out = torch.empty((x.shape[0], physical_n), dtype=x.dtype, device=x.device)
        out.as_strided_((x.shape[0], n), out.stride())
        deep_gemm.bf16_gemm_nt(x, w, out, None)
    else:
        out = torch.empty((x.shape[0], n), dtype=x.dtype, device=x.device)
        deep_gemm.cublaslt_gemm_nt(x, w, out, None)
    return out


def _norm(x: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
    """RMSNorm + 融合残差加；residual=None 时透传 x 作为下一段的 residual。"""
    if residual is None:
        return norm(x, weight, eps=NORM_EPS), x
    return norm(x, weight, residual=residual, eps=NORM_EPS)


class RMSNorm(nn.Module):
    """RMSNorm 参数容器；前向统一走 _norm（需要融合残差加）。"""

    def __init__(self, hidden_size: int) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))


class PatchEmbedFlatten(nn.Module):
    def __init__(self, patch_size: int, in_chans: int, embed_dim: int) -> None:
        super().__init__()
        self.proj = nn.Linear(in_chans * patch_size * patch_size, embed_dim)

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        """images: [n, 3, p, p] -> [n, embed_dim]。"""
        return F.linear(images.view(images.shape[0], 1, -1), self.proj.weight, self.proj.bias).squeeze(-2)


class VisionRotaryEmbedding(nn.Module):
    def __init__(self, dim: int, theta: float = 10000.0) -> None:
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float, device="cuda") / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, seqlen: int) -> torch.Tensor:
        seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        return torch.outer(seq, self.inv_freq)


class Attention(nn.Module):
    def __init__(self, dim: int, num_heads: int) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.qkv_split = [dim, dim, dim]
        self.qkv = nn.Linear(dim, dim * 3)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor, rope: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int) -> torch.Tensor:
        qkv = _linear(x, self.qkv.weight, self.qkv.bias)
        q, k, v = torch.split(qkv, self.qkv_split, dim=1)
        q = q.view(-1, self.num_heads, self.head_dim)
        k = k.view(-1, self.num_heads, self.head_dim)
        v = v.view(-1, self.num_heads, self.head_dim)
        apply_rotary(q, rope, k, interleaved=False)
        # flash-attn 各版本/构建的返回元数不一致（2~3 个或 5 个），此处只取第一个输出
        out, *_ = _flash_attn_fwd(
            q,
            k,
            v,
            cu_seqlens_q=cu_seqlens,
            cu_seqlens_k=cu_seqlens,
            max_seqlen_q=max_seqlen,
            max_seqlen_k=max_seqlen,
            softmax_scale=None,
            causal=False,
            window_size_left=-1,
            window_size_right=-1,
            num_splits=1,
            pack_gqa=None,
            return_lse=False,
        )
        return _linear(out.reshape(x.shape[0], -1), self.proj.weight, self.proj.bias)


class SwiGLU(nn.Module):
    def __init__(self, in_features: int, hidden_features: int, multiple_of: int) -> None:
        super().__init__()
        hidden_features = multiple_of * ((hidden_features + multiple_of - 1) // multiple_of)
        self.linear1 = nn.Linear(in_features, hidden_features * 2, bias=False)
        self.linear2 = nn.Linear(hidden_features, in_features, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = _linear(x, self.linear1.weight, None)
        act = swiglu_forward(gate, dtype_to_str(gate.dtype)).view(gate.shape[0], gate.shape[1] // 2)
        return _linear(act, self.linear2.weight, None)


class Block(nn.Module):
    def __init__(self, dim: int, num_heads: int, mlp_ratio: float, multiple_of: int, return_added_residual: bool) -> None:
        super().__init__()
        self.norm1 = RMSNorm(dim)
        self.attn = Attention(dim, num_heads)
        self.norm2 = RMSNorm(dim)
        self.mlp = SwiGLU(dim, int(dim * mlp_ratio), multiple_of)
        # NOTE: True 时显式 y + residual（bf16 加法），False 时把 residual 交给下一块 norm kernel 融合相加
        self.return_added_residual = return_added_residual

    def forward(
        self,
        x: torch.Tensor,
        residual: torch.Tensor | None,
        rope: torch.Tensor,
        cu_seqlens: torch.Tensor,
        max_seqlen: int,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """返回 (x, residual)；residual 非 None 表示留给下一块的 norm kernel 融合相加。"""
        # 两处 norm 绕开 RMSNorm.forward：融合残差加要返回两个值，签名与 nn.Module 的单值前向不同
        y, residual = _norm(x, self.norm1.weight, residual)
        y = self.attn(y, rope, cu_seqlens, max_seqlen)
        y, residual = _norm(y, self.norm2.weight, residual)
        y = self.mlp(y)
        if self.return_added_residual:
            return y + residual, None
        return y, residual


class VisionTransformer(nn.Module):
    """V4.1 视觉编码器，构造参数取自 model_cfg.vision.params。"""

    def __init__(
        self,
        recomputing_every_layer: int,
        **kwargs: Any,  # cfg.vision 余下的键不参与构造：结构全由本文件常量固定
    ) -> None:
        super().__init__()
        self.patch_size = PATCH_SIZE
        self.patch_embed = PatchEmbedFlatten(PATCH_SIZE, in_chans=3, embed_dim=EMBED_DIM)
        self.rotary_pos_emb = VisionRotaryEmbedding(EMBED_DIM // NUM_HEADS // 2)
        self.blocks = nn.ModuleList(
            [
                Block(
                    dim=EMBED_DIM,
                    num_heads=NUM_HEADS,
                    mlp_ratio=MLP_RATIO,
                    multiple_of=MULTIPLE_OF,
                    # residual 物化点（bf16 加法顺序影响逐 bit 结果）：recompute 段边界与末层物化，
                    # 其余交给下一块的 norm kernel 融合相加
                    return_added_residual=(i + 1) % recomputing_every_layer == 0 or i == DEPTH - 1,
                )
                for i in range(DEPTH)
            ]
        )
        self.fc_norm = RMSNorm(EMBED_DIM)

        freqs = self.rotary_pos_emb(MAX_GRID_SIZE)
        self.rope_table = torch.cat((freqs.cos(), freqs.sin()), dim=-1)
        self.rope_half: int = freqs.shape[1]

    def _rope(self, pos_id: torch.Tensor) -> torch.Tensor:
        """gather 出 [n, 4*half] 的 cos_sin，列序为 (cos_row, cos_col, sin_row, sin_col)。"""
        n = pos_id.shape[0]
        return self.rope_table[pos_id].view(n, 2, 2, self.rope_half).transpose(1, 2).reshape(n, -1)

    @cpu_timer.record("vit_forward")
    def forward(
        self,
        images: torch.Tensor,
        images_global_attn_pos_id: torch.Tensor,
        cu_seqlens_vit: torch.Tensor,
        max_seqlen_vit: int,
        vit_attn_cu_seqlens: list[int] | None = None,
    ) -> torch.Tensor:
        """images: [sum_n_patch, 3, p, p] bf16；pos_id: [sum_n_patch, 2]；返回 [sum_n_patch, dim]。

        vit_attn_cu_seqlens: only temporarily for npu vit attention —— NPU 上 attention 只吃 host list
        （torch_npu 的 actual_seq_qlen），此时用 cu_seqlens_vit[1:] 的 list 形式；GPU 传 None，
        attention 收张量 cu_seqlens_vit。
        """
        assert images.shape[1:] == (3, self.patch_size, self.patch_size), f"unexpected patch shape {tuple(images.shape)}"
        x = self.patch_embed(images)

        rope = self._rope(images_global_attn_pos_id)
        # only temporarily for npu vit attention
        cu_seqlens = cu_seqlens_vit if vit_attn_cu_seqlens is None else vit_attn_cu_seqlens
        residual: torch.Tensor | None = None
        for block in self.blocks:
            x, residual = block(x, residual, rope, cu_seqlens, max_seqlen_vit)

        return norm(x, self.fc_norm.weight, eps=NORM_EPS)