"""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 tile_kernels.quant import swiglu_forward from tile_kernels.transform import apply_rotary from tile_kernels.utils import dtype_to_str 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 # 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)