diff --git a/models/world_model/README.md b/models/world_model/README.md new file mode 100644 index 00000000..6952dde8 --- /dev/null +++ b/models/world_model/README.md @@ -0,0 +1,225 @@ +# Fun-Control 1.3B World Model — pypto3.0 Sub-Networks + +本目录包含 Fun-Control 1.3B 世界模型的 5 个 pypto3.0 子网络实现，对应 `infer_fun_control_1_3b_text.py` 推理 pipeline 中 `WanVideoPipeline` 内部的各个计算阶段。 + +## Pipeline 总览 + +``` +┌─────────────────────────────────────────────────────────────────────────────┐ +│ Fun-Control 1.3B Inference Pipeline │ +├─────────────────────────────────────────────────────────────────────────────┤ +│ │ +│ ┌──────────────┐ ┌──────────────┐ │ +│ │ Text Prompt │───→│ T5 Encoder │───→ text_context [1, T5_SEQ, T5_DIM] │ +│ └──────────────┘ └──────────────┘ │ +│ │ +│ ┌──────────────┐ ┌──────────────┐ │ +│ │ Reference │───→│CLIP Encoder │───→ clip_context [1, CLIP_TOKENS, │ +│ │ Image │ └──────────────┘ CLIP_DIM] │ +│ └──────────────┘ │ +│ │ +│ ┌──────────────┐ ┌──────────────┐ │ +│ │ Images │───→│ VAE Encoder │───→ latent [1, 16, LAT_F, LAT_H, │ +│ │ (ref/input/ │ └──────────────┘ LAT_W] │ +│ │ control) │ │ +│ └──────────────┘ │ +│ │ +│ ┌──────────────────────────────────────────────────────┐ │ +│ │ DiT Forward (Denoising) │ │ +│ │ latent + noise + timestep + text_context + │ │ +│ │ clip_context → denoised_latent │ │ +│ └──────────────────────────────────────────────────────┘ │ +│ │ +│ ┌──────────────┐ │ +│ │ VAE Decoder │───→ video_frames [1, 3, T, H, W] │ +│ └──────────────┘ │ +│ │ +└─────────────────────────────────────────────────────────────────────────────┘ +``` + +--- + +## 子网络详解 + +### 1. T5 Encoder (`t5_encoder.py`) + +**职责**: 将文本 prompt 编码为语义特征向量，为 DiT 提供文本条件。 + +**架构**: +``` +RMSNorm → Self-Attention (with relative position bias) → Residual → +RMSNorm → Gated GELU-tanh FFN → Residual +× T5_LAYERS layers +→ Final RMSNorm +``` + +**输入**: +- `x_in`: [T5_SEQ, T5_DIM] — already-embedded text features (FP32) +- `weights`: T5 encoder weights (layer weights, position bias, final norm) + +**输出**: +- `text_context`: [T5_SEQ, T5_DIM] — encoded text features (BF16) + +**Pipeline 位置**: +- 编码 `prompt` 和 `negative_prompt` +- 输出作为 DiT 的 text conditioning input + +**Golden Case 配置** (config.py): +- T5_DIM = 128, T5_SEQ = 32, T5_LAYERS = 1 +- Real case: T5_DIM = 4096, T5_SEQ = 512, T5_LAYERS = 24 + +--- + +### 2. CLIP Encoder (`clip_encoder.py`) + +**职责**: 将参考图像编码为视觉特征向量，为 DiT 提供图像条件。 + +**架构**: +``` +PatchConv2d → CLS + PosEmbed → PreLN → +[LN → MHA → Proj → Res → LN → FFN(QuickGELU) → Res] × CLIP_LAYERS +``` + +**输入**: +- `img_flat`: [1, IMG_FLAT_SIZE] — flattened reference image (IMG_FLAT_SIZE = 3 * CLIP_IMG * CLIP_IMG) +- `weights`: CLIP encoder weights (patch conv, cls, pos embed, layer weights) + +**输出**: +- `clip_context`: [CLIP_TOKENS, CLIP_DIM] — encoded image features (BF16) + +**Pipeline 位置**: +- 编码 `reference_image` (GT first frame) +- 输出作为 DiT 的 image conditioning input + +**Golden Case 配置** (config.py): +- CLIP_DIM = 128, CLIP_TOKENS = 5, CLIP_LAYERS = 2 +- Real case: CLIP_DIM = 1280, CLIP_TOKENS = 257, CLIP_LAYERS = 31 + +--- + +### 3. VAE Encoder (`vae_encoder.py`) + +**职责**: 将视频帧编码到潜在空间，为 DiT 提供 latent 表示。 + +**架构**: +``` +CausalConv3d(3→96) → +Stage 0: ResBlock(96→192) + ResBlock(192) + Downsample → [1, 192, 3, 32, 32] +Stage 1: ResBlock(192→384) + ResBlock(384) + Downsample → [1, 384, 2, 16, 16] +Stage 2: ResBlock(384→768) + ResBlock(768) + Downsample → [1, 768, 1, 8, 8] +Stage 3: ResBlock(768→1536) + ResBlock(1536) → [1, 1536, 1, 8, 8] +Middle: ResBlock + Attention + ResBlock → [1, 1536, 1, 8, 8] +CausalConv3d(1536→32) → chunk → mu [1, 16, 1, 8, 8] +→ scale_norm → output +``` + +**输入**: +- `video`: [1, 3, T, H, W] — video frames (reference/input/control) +- `weights`: VAE encoder weights + +**输出**: +- `latent`: [1, 16, LAT_F, LAT_H, LAT_W] — encoded latent representation + +**Pipeline 位置**: +- 编码三类视频: `reference_image`, `input_image`, `control_video` +- 输出 latent 作为 DiT 的输入 + +**Golden Case 配置** (config.py): +- VAE_Z_DIM = 16, H = W = 64, CHUNK_SIZE = 5 +- LAT_F = 1 (encoder natural output: 5→3→2→1), LAT_H = 8, LAT_W = 8 + +--- + +### 4. DiT Forward (`dit_forward.py`) + +**职责**: 扩散 Transformer 核心，在潜在空间执行去噪过程，生成新的 latent。 + +**架构**: +``` +ref_conv + patch_conv + concat → x_full +text_proj(GELU-tanh) + clip_proj(LN→FC→GELU→FC→LN) → ctx +AdaLN → Self-Attn(RoPE) → Cross-Attn(text+img) → FFN → Head → output +``` + +**输入**: +- `patch_col`: [X_N, PATCH_CONV_COL] — im2col'd noisy latent (BF16) +- `text_raw`: [T5_SEQ, T5_DIM] — text context from T5 encoder (BF16) +- `clip_raw`: [CLIP_PAD, CLIP_DIM] — image context from CLIP encoder (BF16, padded) +- `ref_col`: [REF_N, REF_CONV_COL] — im2col'd reference latent (BF16) +- `timestep`: [1] — current diffusion timestep +- `weights`: DiT weights (time embedding, projections, transformer blocks, head) + +**输出**: +- `noise_prediction`: [X_N, DIT_IN_DIM * 4] — predicted noise (or denoised latent) + +**Pipeline 位置**: +- 核心生成步骤，执行多步去噪 (typically 50 steps) +- 每步: 预测噪声 → 更新 latent → 下一步 +- 支持 CFG (Classifier-Free Guidance): 分别用 positive/negative prompt 推理，加权组合 + +**Golden Case 配置** (config.py): +- DIT_DIM = 128, DIT_HEADS = 4, DIT_LAYERS = 1 + +--- + +### 5. VAE Decoder (`vae_decoder.py`) + +**职责**: 将潜在表示解码回视频帧，生成最终输出。 + +**架构**: +``` +CausalConv3d(16→1536) → +Middle: ResBlock(1536) + Attention(1536) + ResBlock(1536) +Stage 0: ResBlock(1536→768) + ResBlock(768) + Upsample3d(t_up=T) → [1, 768, 2, 16, 16] +Stage 1: ResBlock(768→384) + ResBlock(384) + Upsample3d(t_up=T) → [1, 384, 4, 32, 32] +Stage 2: ResBlock(384→192) + ResBlock(192) + Upsample3d(t_up=F) → [1, 192, 4, 64, 64] +Stage 3: ResBlock(192→96) + ResBlock(96) +CausalConv3d(96→3) → [1, 3, 4, 64, 64] +→ clamp(-1, 1) +``` + +**输入**: +- `latents`: [1, 16, LAT_F, LAT_H, LAT_W] — denoised latent from DiT +- `weights`: VAE decoder weights + +**输出**: +- `video`: [1, 3, T, H, W] — decoded video frames + +**Pipeline 位置**: +- 将 DiT 生成的 denoised latent 解码为视频帧 tensor (值域 [-1, 1]) +- WanVideoPipeline 内部对输出做 denormalize 并转为 PIL Images 后返回 + +**Golden Case 配置** (config.py): +- VAE_Z_DIM = 16, VAE_DEC_CH = [1536, 768, 384, 192, 96] +- LAT_F = 1 (encoder natural output), H = W = 64 + +--- + +## 公共配置 (`config.py`) + +所有子网络共享的模型配置，包含: +- **Golden Case**: 小规模配置，用于快速数值验证 +- **Real Case**: 完整规模配置，匹配真实模型 (注释状态，按需启用) + +主要配置项: +- T5: T5_DIM, T5_SEQ, T5_LAYERS +- CLIP: CLIP_DIM, CLIP_TOKENS, CLIP_LAYERS +- VAE: VAE_Z_DIM, VAE_ENC_CH, VAE_DEC_CH +- DiT: DIT_DIM, DIT_HEADS, DIT_LAYERS +- Video: H, W, CHUNK_SIZE + +--- + +## 依赖关系 + +``` +vae_decoder.py ──→ vae_encoder.py (import vae_encoder as ve) + ├── 复用 NPU kernel runners (_compile, _run) + ├── 复用 building blocks (causal_conv3d_npu_fp32, resblock_npu, etc.) + ├── 复用 host helpers (_pad_cols, _to_2d, etc.) + └── 复用 compile cache (_cache) + +其他子网络相互独立 +``` + +**注意**: `vae_decoder.py` 依赖 `vae_encoder.py`，无法独立运行。如需独立部署，需将共享代码提取为公共模块。 \ No newline at end of file diff --git a/models/world_model/clip_encoder.py b/models/world_model/clip_encoder.py new file mode 100644 index 00000000..151cf2bf --- /dev/null +++ b/models/world_model/clip_encoder.py @@ -0,0 +1,561 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""CLIP image encoder for Fun-Control 1.3B — pypto3.0 implementation. + +Computation logic follows ``test_golden_fun_control_full.py::clip_encode`` +exactly, using **QuickGELU** activation: ``x * sigmoid(1.702 * x)``. + +Architecture:: + + PatchConv2d → CLS+PosEmbed → PreLN → + [LN → MHA → Proj → Res → LN → FFN(QuickGELU) → Res] × L + +Layer count is parameterised via ``CLIP_LAYERS``; weights are concatenated into +2-D tensors and sliced per layer inside a ``pl.unroll`` loop (compile-time +expansion avoids GM buffer dependency issues with ``pl.range``). + +Hardware constraints (Ascend 910B / A2A3): + • Tile rows must be a multiple of ``innerRows = 16`` → ``T = 16`` (pad from 5). + • Row byte-size must be 32-byte aligned → ``COL_PAD = 608`` (from 588). + • Right buffer ≤ 65 536 B → ``N_CHUNK = 32`` for output-N chunking. + • Large-K matmul (K = 608) causes ``tmov acc→acc`` codegen failure → + each N-chunk of the patch matmul gets its own ``pl.at`` block. + • Padding rows [CLIP_TOKENS..T) must be zeroed after LayerNorm (bias would + activate them and corrupt attention scores). + +Usage:: + + python clip_encoder.py -p a2a3 -d 0 +""" + +import argparse + +import pypto.language as pl +import torch +import torch.nn.functional as F + +from config import ( + CLIP_DIM, CLIP_HEADS, CLIP_LAYERS, CLIP_PATCH, CLIP_IMG, CLIP_TOKENS, +) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Derived constants — hardware-aligned padding / tiling. +# ══════════════════════════════════════════════════════════════════════════════ + +CLIP_HEAD_DIM = CLIP_DIM // CLIP_HEADS # 32 +FFN_DIM = CLIP_DIM * 4 # 512 +EPS = 1e-5 +ATTN_SCALE = CLIP_HEAD_DIM ** -0.5 # 1/√32 + +T = ((CLIP_TOKENS + 15) // 16) * 16 # CLIP_TOKENS padded to innerRows=16 +NUM_PATCHES = (CLIP_IMG // CLIP_PATCH) ** 2 +PATCH_H = CLIP_IMG // CLIP_PATCH +PATCH_W = CLIP_IMG // CLIP_PATCH +COL_SIZE = 3 * CLIP_PATCH * CLIP_PATCH # 588 (im2col column width) +COL_PAD = 608 # 32-byte aligned, avoids tmov bug +IMG_FLAT_SIZE = 3 * CLIP_IMG * CLIP_IMG # 2352 +N_CHUNK = 32 # output-N tile (Right buffer safe) + +# ── QuickGELU constant (matches golden reference: x * sigmoid(1.702 * x)) ── +QUICK_GELU_SCALE = 1.702 + +# Geometry assertions — keep at the bottom of the derived-constants block. +assert CLIP_DIM % CLIP_HEADS == 0, "CLIP_HEADS must divide CLIP_DIM" +assert CLIP_HEAD_DIM * CLIP_HEADS == CLIP_DIM +assert T % 16 == 0, "T must be a multiple of innerRows=16" +assert T >= CLIP_TOKENS, "T must accommodate all real tokens" +assert CLIP_TOKENS == NUM_PATCHES + 1, "expected CLS token + NUM_PATCHES patch tokens" +assert COL_PAD >= COL_SIZE and COL_PAD % 4 == 0, "COL_PAD must be ≥ COL_SIZE and 32-byte aligned" +assert FFN_DIM == CLIP_DIM * 4 +assert IMG_FLAT_SIZE == 3 * CLIP_IMG * CLIP_IMG +assert CLIP_DIM % N_CHUNK == 0, "N_CHUNK must divide CLIP_DIM" +assert FFN_DIM % N_CHUNK == 0, "N_CHUNK must divide FFN_DIM" +assert (3 * CLIP_DIM) % N_CHUNK == 0, "N_CHUNK must divide 3*CLIP_DIM (QKV width)" + + +# ══════════════════════════════════════════════════════════════════════════════ +# Reusable inline kernels. +# ══════════════════════════════════════════════════════════════════════════════ + +@pl.jit.inline +def _layernorm( + x: pl.Tensor[[T, CLIP_DIM], pl.FP32], + w: pl.Tensor[[1, CLIP_DIM], pl.FP32], + b: pl.Tensor[[1, CLIP_DIM], pl.FP32], + y: pl.Tensor[[T, CLIP_DIM], pl.FP32], +): + """LayerNorm with padding-row masking. + + Computes ``(x - mean) / std * w + b`` over all T rows, then zeroes rows + ``[CLIP_TOKENS, T)`` so that padding rows stay exactly 0 (the bias would + otherwise activate them and corrupt downstream attention). + """ + scaled = pl.mul(x, 1.0 / CLIP_DIM) + mean = pl.row_sum(scaled) + centred = pl.row_expand_sub(x, mean) + sq = pl.mul(centred, centred) + var = pl.row_sum(sq) + var = pl.mul(var, 1.0 / CLIP_DIM) + var_eps = pl.add(var, EPS) + var_eps = pl.reshape(var_eps, [1, T]) + std = pl.sqrt(var_eps) + std = pl.reshape(std, [T, 1]) + normed = pl.row_expand_div(centred, std) + + scaled_n = pl.col_expand_mul(normed, w) + # A2/A3 does not support pl.full(dtype=pl.FP32). Work around by zeroing + # a same-shape tile then adding 1.0 — produces an all-ones FP32 tile. + ones = pl.sub(x, x) + ones = pl.add(ones, 1.0) + biased = pl.add(scaled_n, pl.col_expand_mul(ones, b)) + + for r in pl.range(CLIP_TOKENS, T): + zero_row = pl.mul(pl.slice(biased, [1, CLIP_DIM], [r, 0]), 0.0) + biased = pl.assemble(biased, zero_row, [r, 0]) + + y = pl.assemble(y, biased, [0, 0]) + return y + + +# ══════════════════════════════════════════════════════════════════════════════ +# Main kernel — full CLIP encoder forward pass. +# ══════════════════════════════════════════════════════════════════════════════ + +@pl.jit +def clip_encoder( + # ── Inputs ── + img_flat: pl.Tensor[[1, IMG_FLAT_SIZE], pl.FP32], + # ── Patch embedding ── + patch_conv_w: pl.Tensor[[CLIP_DIM, COL_PAD], pl.BF16], + cls_token: pl.Tensor[[1, CLIP_DIM], pl.BF16], + pos_embed: pl.Tensor[[T, CLIP_DIM], pl.BF16], + # ── Pre-LayerNorm ── + pre_ln_w: pl.Tensor[[1, CLIP_DIM], pl.FP32], + pre_ln_b: pl.Tensor[[1, CLIP_DIM], pl.FP32], + # ── Stacked transformer weights ── + norm1_w: pl.Tensor[[CLIP_LAYERS, CLIP_DIM], pl.FP32], + norm1_b: pl.Tensor[[CLIP_LAYERS, CLIP_DIM], pl.FP32], + qkv_w: pl.Tensor[[CLIP_LAYERS * 3 * CLIP_DIM, CLIP_DIM], pl.BF16], + proj_w: pl.Tensor[[CLIP_DIM, CLIP_LAYERS * CLIP_DIM], pl.BF16], + norm2_w: pl.Tensor[[CLIP_LAYERS, CLIP_DIM], pl.FP32], + norm2_b: pl.Tensor[[CLIP_LAYERS, CLIP_DIM], pl.FP32], + fc1_w: pl.Tensor[[CLIP_DIM, CLIP_LAYERS * FFN_DIM], pl.BF16], + fc2_w: pl.Tensor[[CLIP_LAYERS * FFN_DIM, CLIP_DIM], pl.BF16], + # ── Output ── + out: pl.Out[pl.Tensor[[CLIP_TOKENS, CLIP_DIM], pl.BF16]], +): + + # ── Scope 1: Patch embedding — im2col (kernel-side, scalar access). ── + # Follows the same pl.read + pl.tensor.write pattern as common/conv2d.py. + col_buf = pl.create_tensor([T, COL_PAD], dtype=pl.FP32) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="im2col"): + for ph in pl.range(PATCH_H): + for pw in pl.range(PATCH_W): + patch_idx = ph * PATCH_W + pw + col_offset = 0 + for c in pl.range(3): + for kh in pl.range(CLIP_PATCH): + for kw in pl.range(CLIP_PATCH): + src_h = ph * CLIP_PATCH + kh + src_w = pw * CLIP_PATCH + kw + src_col = c * CLIP_IMG * CLIP_IMG + src_h * CLIP_IMG + src_w + val = pl.read(img_flat, [0, src_col]) + pl.tensor.write(col_buf, [patch_idx, col_offset], val) + col_offset = col_offset + 1 + + # ── Scope 2: Patch matmul — N-chunked (each chunk in its own pl.at). ── + # Separate pl.at blocks avoid tmov acc→acc codegen failure on K=608. + patch_emb_gm = pl.create_tensor([T, CLIP_DIM], dtype=pl.BF16) + for nc in pl.unroll(CLIP_DIM // N_CHUNK): + n0 = nc * N_CHUNK + with pl.at(level=pl.Level.CORE_GROUP, name_hint="patch_mm"): + col_tile = pl.slice(col_buf, [T, COL_PAD], [0, 0]) + col_bf = pl.cast(col_tile, target_type=pl.BF16) + pw_chunk = pl.slice(patch_conv_w, [N_CHUNK, COL_PAD], [n0, 0]) + patch_out = pl.matmul(col_bf, pw_chunk, b_trans=True, out_dtype=pl.FP32) + patch_emb_gm = pl.assemble( + patch_emb_gm, pl.cast(patch_out, target_type=pl.BF16), [0, n0], + ) + + # ── Scope 3: Prepend CLS token + add positional embedding. ── + x_pos_gm = pl.create_tensor([T, CLIP_DIM], dtype=pl.BF16) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="cls_pos"): + x_pos_gm = pl.assemble(x_pos_gm, pl.slice(cls_token, [1, CLIP_DIM], [0, 0]), [0, 0]) + for i in pl.range(NUM_PATCHES): + x_pos_gm = pl.assemble( + x_pos_gm, pl.slice(patch_emb_gm, [1, CLIP_DIM], [i, 0]), [i + 1, 0], + ) + + x_ln_gm = pl.create_tensor([T, CLIP_DIM], dtype=pl.FP32) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="add_pos"): + x_tile = pl.slice(x_pos_gm, [T, CLIP_DIM], [0, 0]) + p_tile = pl.slice(pos_embed, [T, CLIP_DIM], [0, 0]) + x_ln_gm = pl.assemble( + x_ln_gm, + pl.add(pl.cast(x_tile, target_type=pl.FP32), + pl.cast(p_tile, target_type=pl.FP32)), + [0, 0], + ) + + # ── Scope 4: Pre-LayerNorm. ── + layer_x_gm = pl.create_tensor([T, CLIP_DIM], dtype=pl.FP32) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="pre_ln"): + x_tile = pl.slice(x_ln_gm, [T, CLIP_DIM], [0, 0]) + layer_x_gm = _layernorm(x_tile, pre_ln_w, pre_ln_b, layer_x_gm) + + # ══════════════════════════════════════════════════════════════════════ + # Transformer layers — parameterised loop over CLIP_LAYERS. + # + # GM buffers are pre-allocated once and reused each iteration + # (each scope overwrites before the next scope reads). + # ══════════════════════════════════════════════════════════════════════ + qkv_gm = pl.create_tensor([T, 3 * CLIP_DIM], dtype=pl.BF16) + attn_gm = pl.create_tensor([T, CLIP_DIM], dtype=pl.FP32) + attn_res_gm = pl.create_tensor([T, CLIP_DIM], dtype=pl.FP32) + ln1_gm = pl.create_tensor([T, CLIP_DIM], dtype=pl.FP32) + ln2_gm = pl.create_tensor([T, CLIP_DIM], dtype=pl.FP32) + gelu_gm = pl.create_tensor([T, FFN_DIM], dtype=pl.FP32) + + for layer_idx in pl.unroll(CLIP_LAYERS): + + # ── Scope 5: LN1 + QKV projection. ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ln_qkv"): + ln1_w = pl.slice(norm1_w, [1, CLIP_DIM], [layer_idx, 0]) + ln1_b = pl.slice(norm1_b, [1, CLIP_DIM], [layer_idx, 0]) + x_tile = pl.slice(layer_x_gm, [T, CLIP_DIM], [0, 0]) + ln1_gm = _layernorm(x_tile, ln1_w, ln1_b, ln1_gm) + normed_tile = pl.slice(ln1_gm, [T, CLIP_DIM], [0, 0]) + normed_bf = pl.cast(normed_tile, target_type=pl.BF16) + qkv_offset = layer_idx * (3 * CLIP_DIM) + for n0 in pl.range(0, 3 * CLIP_DIM, N_CHUNK): + qkv_chunk = pl.matmul( + normed_bf, + pl.slice(qkv_w, [N_CHUNK, CLIP_DIM], [qkv_offset + n0, 0]), + b_trans=True, out_dtype=pl.FP32, + ) + qkv_gm = pl.assemble( + qkv_gm, pl.cast(qkv_chunk, target_type=pl.BF16), [0, n0], + ) + for r in pl.range(CLIP_TOKENS, T): + zero_row = pl.full([1, 3 * CLIP_DIM], dtype=pl.BF16, value=0.0) + qkv_gm = pl.assemble(qkv_gm, zero_row, [r, 0]) + + # ── Scope 6: Multi-head self-attention (per-head, with scale). ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="mha"): + for h in pl.range(CLIP_HEADS): + h_col = h * CLIP_HEAD_DIM + q_h = pl.slice(qkv_gm, [T, CLIP_HEAD_DIM], [0, h_col]) + k_h = pl.slice(qkv_gm, [T, CLIP_HEAD_DIM], [0, CLIP_DIM + h_col]) + v_h = pl.slice(qkv_gm, [T, CLIP_HEAD_DIM], [0, 2 * CLIP_DIM + h_col]) + + scores = pl.mul( + pl.matmul(q_h, k_h, b_trans=True, out_dtype=pl.FP32), ATTN_SCALE, + ) + shifted = pl.row_expand_sub(scores, pl.row_max(scores)) + exp_s = pl.exp(shifted) + sm = pl.row_expand_div(exp_s, pl.row_sum(exp_s)) + ctx = pl.matmul(pl.cast(sm, target_type=pl.BF16), v_h, out_dtype=pl.FP32) + attn_gm = pl.assemble(attn_gm, ctx, [0, h_col]) + for r in pl.range(CLIP_TOKENS, T): + zero_attn = pl.full([1, CLIP_DIM], dtype=pl.FP32, value=0.0) + attn_gm = pl.assemble(attn_gm, zero_attn, [r, 0]) + + # ── Scope 7: Output projection + residual. ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="proj_res"): + attn_tile = pl.cast(pl.slice(attn_gm, [T, CLIP_DIM], [0, 0]), target_type=pl.BF16) + proj_offset = layer_idx * CLIP_DIM + proj_w_l = pl.slice(proj_w, [CLIP_DIM, CLIP_DIM], [0, proj_offset]) + proj = pl.matmul(attn_tile, proj_w_l, out_dtype=pl.FP32) + residual = pl.slice(layer_x_gm, [T, CLIP_DIM], [0, 0]) + attn_res_gm = pl.assemble(attn_res_gm, pl.add(residual, proj), [0, 0]) + for r in pl.range(CLIP_TOKENS, T): + zero_res = pl.full([1, CLIP_DIM], dtype=pl.FP32, value=0.0) + attn_res_gm = pl.assemble(attn_res_gm, zero_res, [r, 0]) + + # ── Scope 8: LN2 + FC1 + QuickGELU (x * sigmoid(1.702 * x)). ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ln_fc1"): + ln2_w_l = pl.slice(norm2_w, [1, CLIP_DIM], [layer_idx, 0]) + ln2_b_l = pl.slice(norm2_b, [1, CLIP_DIM], [layer_idx, 0]) + ar_tile = pl.slice(attn_res_gm, [T, CLIP_DIM], [0, 0]) + ln2_gm = _layernorm(ar_tile, ln2_w_l, ln2_b_l, ln2_gm) + ln2_tile = pl.slice(ln2_gm, [T, CLIP_DIM], [0, 0]) + normed_bf = pl.cast(ln2_tile, target_type=pl.BF16) + fc1_offset = layer_idx * FFN_DIM + for n0 in pl.range(0, FFN_DIM, N_CHUNK): + fc1 = pl.matmul( + normed_bf, + pl.slice(fc1_w, [CLIP_DIM, N_CHUNK], [0, fc1_offset + n0]), + out_dtype=pl.FP32, + ) + # QuickGELU: x * sigmoid(1.702 * x) + scaled = pl.mul(fc1, QUICK_GELU_SCALE) + neg_sc = pl.neg(scaled) + exp_neg = pl.exp(neg_sc) + denom = pl.add(exp_neg, 1.0) + sigmoid = pl.recip(denom) + qgelu = pl.mul(fc1, sigmoid) + gelu_gm = pl.assemble(gelu_gm, qgelu, [0, n0]) + + # ── Scope 9: FC2 (K-chunked accumulation) + residual → layer_x_gm. ── + # qgelu [T, FFN_DIM] @ fc2_w [FFN_DIM, CLIP_DIM] → [T, CLIP_DIM] + # K-chunk: [T, N_CHUNK] @ [N_CHUNK, CLIP_DIM] (no b_trans needed) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="fc2_res"): + fc2_offset = layer_idx * FFN_DIM + g_chunk0 = pl.cast(pl.slice(gelu_gm, [T, N_CHUNK], [0, 0]), target_type=pl.BF16) + w_chunk0 = pl.slice(fc2_w, [N_CHUNK, CLIP_DIM], [fc2_offset, 0]) + acc = pl.matmul(g_chunk0, w_chunk0, out_dtype=pl.FP32) + for k0 in pl.range(N_CHUNK, FFN_DIM, N_CHUNK): + g_chunk = pl.cast(pl.slice(gelu_gm, [T, N_CHUNK], [0, k0]), target_type=pl.BF16) + w_chunk = pl.slice(fc2_w, [N_CHUNK, CLIP_DIM], [fc2_offset + k0, 0]) + acc = pl.matmul_acc(acc, g_chunk, w_chunk) + residual = pl.slice(attn_res_gm, [T, CLIP_DIM], [0, 0]) + layer_x_gm = pl.assemble(layer_x_gm, pl.add(residual, acc), [0, 0]) + + # ── Scope 10: Final output — only valid rows (no padding). ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="final"): + final_tile = pl.slice(layer_x_gm, [CLIP_TOKENS, CLIP_DIM], [0, 0]) + out = pl.assemble(out, pl.cast(final_tile, target_type=pl.BF16), [0, 0]) + + return out + + +# ══════════════════════════════════════════════════════════════════════════════ +# Golden reference — mirrors test_golden_fun_control_full.py::clip_encode +# with QuickGELU: x * sigmoid(1.702 * x) (matching golden reference line 742). +# ══════════════════════════════════════════════════════════════════════════════ + +def _bf16(t): + """Round through BF16 then back to FP32 — emulates kernel pl.cast(BF16).""" + if isinstance(t, torch.Tensor): + return t.to(torch.bfloat16).to(torch.float32) + return t + + +def _im2col_clip(img_4d, patch_h, patch_w, clip_patch): + """im2col for CLIP patch conv: matches kernel's scalar read/write order. + img_4d: [B, 3, H, W]. Returns [B, num_patches, col_size]. + """ + B, C, H, W = img_4d.shape + num_patches = patch_h * patch_w + col_size = C * clip_patch * clip_patch + cols = torch.zeros(B, num_patches, col_size) + for b in range(B): + for ph in range(patch_h): + for pw in range(patch_w): + patch_idx = ph * patch_w + pw + col_offset = 0 + for c in range(C): + for kh in range(clip_patch): + for kw in range(clip_patch): + src_h = ph * clip_patch + kh + src_w = pw * clip_patch + kw + cols[b, patch_idx, col_offset] = img_4d[b, c, src_h, src_w] + col_offset += 1 + return cols + + +def golden_clip_encode(img_tensor, w): + """BF16-aware golden matching kernel's exact computation path. + + Implements custom LayerNorm and N-chunked matmul to match kernel exactly. + Output: [CLIP_TOKENS, CLIP_DIM] — only valid rows, no padding. + """ + B = img_tensor.shape[0] + + def custom_layernorm(x, w, b): + """Match kernel's _layernorm exactly.""" + scaled = x * (1.0 / CLIP_DIM) + mean = scaled.sum(dim=-1, keepdim=True) + centred = x - mean + sq = centred * centred + var = sq.sum(dim=-1, keepdim=True) * (1.0 / CLIP_DIM) + var_eps = var + EPS + std = var_eps.sqrt() + normed = centred / std + scaled_n = normed * w + biased = scaled_n + b + # Zero padding rows + if biased.shape[1] > CLIP_TOKENS: + biased[:, CLIP_TOKENS:, :] = 0.0 + return biased + + def n_chunked_matmul(x, w, n_chunk=N_CHUNK): + """Match kernel's N-chunked matmul accumulation order.""" + result = torch.zeros(x.shape[0], w.shape[0], dtype=x.dtype) + for n0 in range(0, w.shape[0], n_chunk): + w_chunk = w[n0:n0+n_chunk] + result[:, n0:n0+n_chunk] = _bf16(x) @ _bf16(w_chunk).T + return result + + col = _im2col_clip(img_tensor, PATCH_H, PATCH_W, CLIP_PATCH) + col_padded = torch.zeros(T, COL_PAD) + col_padded[:NUM_PATCHES, :COL_SIZE] = col.squeeze(0) + w_raw = w["clip_patch_conv"].reshape(CLIP_DIM, -1) + w_padded = torch.zeros(CLIP_DIM, COL_PAD, dtype=w_raw.dtype) + w_padded[:, :w_raw.shape[1]] = w_raw + patch_emb = _bf16(n_chunked_matmul(_bf16(col_padded), _bf16(w_padded))) + x = patch_emb[:NUM_PATCHES, :].unsqueeze(0) + + cls = w["clip_cls"].expand(B, -1, -1) + x = torch.cat([cls, x], dim=1) + x = _bf16(x + w["clip_pos"]) + x = custom_layernorm(x, w["clip_pre_norm_w"], w["clip_pre_norm_b"]) + for i in range(CLIP_LAYERS): + p = f"clip_l{i}" + h = custom_layernorm(x, w[f"{p}_norm1_w"], w[f"{p}_norm1_b"]) + qkv_w = w[f"{p}_qkv"] + qkv = _bf16(n_chunked_matmul(_bf16(h.squeeze(0)), _bf16(qkv_w), n_chunk=3*CLIP_DIM)).unsqueeze(0) + qkv = qkv.view(B, -1, 3, CLIP_HEADS, CLIP_HEAD_DIM) + qkv = qkv.permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + scores = _bf16(q) @ _bf16(k).transpose(-2, -1) * ATTN_SCALE + attn = F.softmax(scores, dim=-1) + attn_out = _bf16(_bf16(attn) @ _bf16(v)).transpose(1, 2).contiguous().view(B, -1, CLIP_DIM) + proj_w = w[f"{p}_proj"] + proj = _bf16(n_chunked_matmul(_bf16(attn_out.squeeze(0)), _bf16(proj_w))).unsqueeze(0) + x = _bf16(x + proj) + # Zero padding after residual + if x.shape[1] > CLIP_TOKENS: + x[:, CLIP_TOKENS:, :] = 0.0 + h = custom_layernorm(x, w[f"{p}_norm2_w"], w[f"{p}_norm2_b"]) + fc1_w = w[f"{p}_fc1"] + ff = _bf16(n_chunked_matmul(_bf16(h.squeeze(0)), _bf16(fc1_w), n_chunk=FFN_DIM)).unsqueeze(0) + ff = ff * torch.sigmoid(1.702 * ff) + fc2_w = w[f"{p}_fc2"] + fc2 = _bf16(n_chunked_matmul(_bf16(ff.squeeze(0)), _bf16(fc2_w))).unsqueeze(0) + x = _bf16(x + fc2) + # Zero padding after residual + if x.shape[1] > CLIP_TOKENS: + x[:, CLIP_TOKENS:, :] = 0.0 + return x.squeeze(0)[:CLIP_TOKENS, :] + + +# ══════════════════════════════════════════════════════════════════════════════ +# Test harness — build_tensor_specs / golden_fn / __main__. +# ══════════════════════════════════════════════════════════════════════════════ + +def build_tensor_specs(): + from golden import TensorSpec + + g = torch.Generator().manual_seed(42) + + patch_conv_w_4d = torch.randn(CLIP_DIM, 3, CLIP_PATCH, CLIP_PATCH, generator=g) * 0.02 + patch_conv_w_padded = torch.zeros(CLIP_DIM, COL_PAD, dtype=torch.bfloat16) + patch_conv_w_padded[:, :COL_SIZE] = patch_conv_w_4d.reshape(CLIP_DIM, COL_SIZE).bfloat16() + + cls_3d = torch.randn(1, 1, CLIP_DIM, generator=g) * 0.02 + pos_3d = torch.randn(1, CLIP_TOKENS, CLIP_DIM, generator=g) * 0.02 + pos_padded = torch.zeros(T, CLIP_DIM, dtype=torch.bfloat16) + pos_padded[:CLIP_TOKENS, :] = pos_3d.squeeze(0).bfloat16() + + specs = [ + TensorSpec("img_flat", [1, IMG_FLAT_SIZE], torch.float32, init_value=torch.randn), + TensorSpec("patch_conv_w", [CLIP_DIM, COL_PAD], torch.bfloat16, init_value=patch_conv_w_padded), + TensorSpec("cls_token", [1, CLIP_DIM], torch.bfloat16, init_value=cls_3d.squeeze(0).bfloat16()), + TensorSpec("pos_embed", [T, CLIP_DIM], torch.bfloat16, init_value=pos_padded), + TensorSpec("pre_ln_w", [1, CLIP_DIM], torch.float32, init_value=torch.ones), + TensorSpec("pre_ln_b", [1, CLIP_DIM], torch.float32, init_value=torch.zeros), + ] + + stacked_norm1_w = torch.ones(CLIP_LAYERS, CLIP_DIM, dtype=torch.float32) + stacked_norm1_b = torch.zeros(CLIP_LAYERS, CLIP_DIM, dtype=torch.float32) + stacked_qkv_w = torch.cat([ + (torch.randn(3 * CLIP_DIM, CLIP_DIM, generator=g) * 0.02).bfloat16() + for _ in range(CLIP_LAYERS) + ], dim=0) + stacked_proj_w = torch.cat([ + (torch.randn(CLIP_DIM, CLIP_DIM, generator=g) * 0.02).bfloat16().T + for _ in range(CLIP_LAYERS) + ], dim=1) + stacked_norm2_w = torch.ones(CLIP_LAYERS, CLIP_DIM, dtype=torch.float32) + stacked_norm2_b = torch.zeros(CLIP_LAYERS, CLIP_DIM, dtype=torch.float32) + stacked_fc1_w = torch.cat([ + (torch.randn(FFN_DIM, CLIP_DIM, generator=g) * 0.02).bfloat16().T + for _ in range(CLIP_LAYERS) + ], dim=1) + stacked_fc2_w = torch.cat([ + (torch.randn(CLIP_DIM, FFN_DIM, generator=g) * 0.02).bfloat16().T + for _ in range(CLIP_LAYERS) + ], dim=0) + + specs.extend([ + TensorSpec("norm1_w", [CLIP_LAYERS, CLIP_DIM], torch.float32, init_value=stacked_norm1_w), + TensorSpec("norm1_b", [CLIP_LAYERS, CLIP_DIM], torch.float32, init_value=stacked_norm1_b), + TensorSpec("qkv_w", [CLIP_LAYERS * 3 * CLIP_DIM, CLIP_DIM], torch.bfloat16, init_value=stacked_qkv_w), + TensorSpec("proj_w", [CLIP_DIM, CLIP_LAYERS * CLIP_DIM], torch.bfloat16, init_value=stacked_proj_w), + TensorSpec("norm2_w", [CLIP_LAYERS, CLIP_DIM], torch.float32, init_value=stacked_norm2_w), + TensorSpec("norm2_b", [CLIP_LAYERS, CLIP_DIM], torch.float32, init_value=stacked_norm2_b), + TensorSpec("fc1_w", [CLIP_DIM, CLIP_LAYERS * FFN_DIM], torch.bfloat16, init_value=stacked_fc1_w), + TensorSpec("fc2_w", [CLIP_LAYERS * FFN_DIM, CLIP_DIM], torch.bfloat16, init_value=stacked_fc2_w), + ]) + + specs.append(TensorSpec("out", [CLIP_TOKENS, CLIP_DIM], torch.bfloat16, is_output=True)) + return specs + + +def golden_clip_encoder(tensors): + """Run golden_clip_encode with QuickGELU, output only valid rows.""" + img_flat = tensors["img_flat"] + img_4d = img_flat.reshape(1, 3, CLIP_IMG, CLIP_IMG).float() + + w = { + "clip_patch_conv": tensors["patch_conv_w"][:, :COL_SIZE].float() + .reshape(CLIP_DIM, 3, CLIP_PATCH, CLIP_PATCH), + "clip_cls": tensors["cls_token"].float().unsqueeze(0), + "clip_pos": tensors["pos_embed"][:CLIP_TOKENS, :].float().unsqueeze(0), + "clip_pre_norm_w": tensors["pre_ln_w"].squeeze(0).float(), + "clip_pre_norm_b": tensors["pre_ln_b"].squeeze(0).float(), + } + for i in range(CLIP_LAYERS): + p = f"clip_l{i}" + w[f"{p}_norm1_w"] = tensors["norm1_w"][i].float() + w[f"{p}_norm1_b"] = tensors["norm1_b"][i].float() + qkv_start = i * 3 * CLIP_DIM + w[f"{p}_qkv"] = tensors["qkv_w"][qkv_start:qkv_start + 3 * CLIP_DIM].float() + proj_start = i * CLIP_DIM + w[f"{p}_proj"] = tensors["proj_w"][:, proj_start:proj_start + CLIP_DIM].float().T + w[f"{p}_norm2_w"] = tensors["norm2_w"][i].float() + w[f"{p}_norm2_b"] = tensors["norm2_b"][i].float() + fc1_start = i * FFN_DIM + w[f"{p}_fc1"] = tensors["fc1_w"][:, fc1_start:fc1_start + FFN_DIM].float().T + fc2_start = i * FFN_DIM + w[f"{p}_fc2"] = tensors["fc2_w"][fc2_start:fc2_start + FFN_DIM].float().T + + result = golden_clip_encode(img_4d, w) + tensors["out"][:] = result.bfloat16() + + +if __name__ == "__main__": + from golden import ratio_allclose, run_jit + + parser = argparse.ArgumentParser() + parser.add_argument("-p", "--platform", type=str, default="a2a3", + choices=["a2a3", "a2a3sim", "a5", "a5sim"]) + parser.add_argument("-d", "--device", type=int, default=0) + parser.add_argument("--enable-l2-swimlane", action="store_true", default=False) + args = parser.parse_args() + + result = run_jit( + fn=clip_encoder, + specs=build_tensor_specs(), + golden_fn=golden_clip_encoder, + runtime_cfg=dict( + platform=args.platform, + device_id=args.device, + enable_l2_swimlane=args.enable_l2_swimlane, + ), + rtol=3e-3, + atol=3e-3, + compare_fn={"out": ratio_allclose(atol=0.1, rtol=0.1, max_error_ratio=0.02)}, + ) + if not result.passed: + if result.error: + print(result.error) + raise SystemExit(1) diff --git a/models/world_model/config.py b/models/world_model/config.py new file mode 100644 index 00000000..95f8ff87 --- /dev/null +++ b/models/world_model/config.py @@ -0,0 +1,174 @@ +# ══════════════════════════════════════════════════════════════════════════════ +# Fun-Control 1.3B — pypto3.0 model configuration. +# +# Two cases are tracked: +# • Golden case — small-scale config for fast numerical validation against +# test_golden_fun_control_full.py. Active by default. +# • Real case — full-size config matching infer_fun_control_1_3b_text.py +# with the real open-clip-xlm-roberta-large-vit-huge-14 checkpoint. +# Commented out; activate when ready. +# ══════════════════════════════════════════════════════════════════════════════ + +# ── Video / image ── +H, W = 64, 64 +NUM_FRAMES = 5 +CHUNK_SIZE = 5 +FALLBACK_PROMPT = "robot arm manipulation" + + +# ══════════════════════════════════════════════════════════════════════════════ +# T5 text encoder +# ══════════════════════════════════════════════════════════════════════════════ + +# ── Golden case (test_golden_fun_control_full.py:65-71) ── +T5_VOCAB = 1000 +T5_DIM = 128 +T5_FFN = 256 +T5_HEADS = 4 +T5_HEAD_DIM = 32 +T5_LAYERS = 1 +T5_SEQ = 32 +T5_NUM_BUCKETS = 32 + +# ── Real case (DiffSynth-Studio t5_umt5-xxl-enc-bf16.pth) ── +# T5_VOCAB = 256384 +# T5_DIM = 4096 +# T5_FFN = 10240 +# T5_HEADS = 64 +# T5_HEAD_DIM = 64 +# T5_LAYERS = 24 +# T5_SEQ = 512 # Reduced from 512 for testing +# T5_NUM_BUCKETS = 32 + + +# ══════════════════════════════════════════════════════════════════════════════ +# CLIP image encoder +# +# Architecture: PatchConv2d → CLS+PosEmbed → PreLN → +# [LN → MHA → Proj → Res → LN → FFN → Res] × L +# +# Source (golden): test_golden_fun_control_full.py:73-79, 127-142, 721-744 +# Source (real): DiffSynth-Studio/diffsynth/models/wan_video_image_encoder.py +# clip_xlm_roberta_vit_h_14() → lines 822-849 +# VisionTransformer.__init__ → lines 386-454 +# VisionTransformer.forward → lines 456-478 +# AttentionBlock.__init__ → lines 289-330 +# XLMRobertaCLIP.__init__ → lines 642-710 +# WanImageEncoder.encode_image → lines 864-878 +# ══════════════════════════════════════════════════════════════════════════════ + +# ── Golden case (test_golden_fun_control_full.py:73-79) ── +# +# Reduced dimensions for fast numerical validation. +# Activation: QuickGELU → ff * sigmoid(1.702 * ff) (golden:742) +# Layers run: all (CLIP_LAYERS = 2) +# Patch conv: Conv2d(3, 128, k=14, s=14, bias=False) (no bias weight in golden:128) +# Output: [1, 5, 128] (full sequence including CLS) +# +CLIP_DIM = 128 +CLIP_HEADS = 4 +CLIP_HEAD_DIM = 32 # CLIP_DIM // CLIP_HEADS +CLIP_LAYERS = 2 +CLIP_PATCH = 14 +CLIP_IMG = 28 # must be divisible by CLIP_PATCH +CLIP_TOKENS = 5 # (CLIP_IMG // CLIP_PATCH) ** 2 + 1 = 4 + 1 +CLIP_FFN = 512 # CLIP_DIM * 4 +CLIP_NORM_EPS = 1e-5 + +# ── Real case (open-clip-xlm-roberta-large-vit-huge-14) ── +# +# Full ViT-H/14 from OpenCLIP, loaded via: +# ModelConfig(path="models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth") +# +# Key differences from golden case: +# 1. Scale: 1280-dim, 16 heads, 32 layers (vs 128/4/2) +# 2. Image: 224×224 → 256 patches → 257 tokens (vs 28×28 → 4 → 5) +# 3. Activation: standard GELU (nn.GELU) — NOT QuickGELU +# Trace: clip_xlm_roberta_vit_h_14 sets activation='gelu' (line 835) +# → XLMRobertaCLIP(activation='gelu') (line 655) +# → VisionTransformer(activation='gelu') (line 702) +# → AttentionBlock(activation='gelu') (line 442) +# → nn.GELU() (line 320: else branch) +# Golden test uses QuickGELU (golden:742) — ACTIVATION MISMATCH +# 4. Blocks: encode_image calls visual(x, use_31_block=True) (line 877) +# → runs transformer[:-1] = first 31 of 32 blocks (line 474) +# → early return, skips block 32 + post_norm + head (line 475) +# 5. Patch conv: Conv2d(3, 1280, k=14, s=14, bias=False) +# bias=False because pre_norm=True (line 429: bias=not pre_norm) +# 6. CLS token: shape [1, 1, 1280], gain-init: 1/sqrt(1280) * randn +# 7. Pos embed: shape [1, 257, 1280], gain-init: 1/sqrt(1280) * randn +# 8. Pre-LN: LayerNorm(1280, eps=1e-5) — exists because pre_norm=True +# 9. Post-LN: LayerNorm(1280, eps=1e-5) — created but SKIPPED (use_31_block) +# 10. Head: nn.Parameter(1280, 1024) — created but SKIPPED (use_31_block) +# 11. Output: [1, 257, 1280] raw hidden states after 31 blocks +# +# Preprocessing (WanImageEncoder.encode_image, lines 864-878): +# PIL Image → resize (width, height) → preprocess_image: [-1, 1] +# → bicubic resize to 224×224 → mul_(0.5).add_(0.5): [0, 1] +# → Normalize(mean, std): +# mean = [0.48145466, 0.4578275, 0.40821073] (line 785) +# std = [0.26862954, 0.26130258, 0.27577711] (line 786) +# +# Downstream usage (wan_video_dit.py WanModel.forward): +# clip_feature [1, 257, 1280] +# → dit.img_emb(clip_feature): MLP(1280 → dim) +# = LayerNorm(1280) → Linear(1280,1280) → GELU → Linear(1280,dim) → LayerNorm(dim) +# → torch.cat([clip_embedding, text_context], dim=1) +# +# CLIP_DIM = 1280 +# CLIP_HEADS = 16 +# CLIP_HEAD_DIM = 80 # 1280 // 16 +# CLIP_LAYERS = 31 # use_31_block=True: first 31 of 32 +# CLIP_PATCH = 14 +# CLIP_IMG = 224 +# CLIP_TOKENS = 257 # (224 // 14) ** 2 + 1 = 256 + 1 +# CLIP_FFN = 5120 # 1280 * 4 +# CLIP_NORM_EPS = 1e-5 +# CLIP_ACTIVATION = "gelu" # standard GELU, NOT QuickGELU +# CLIP_EMBED_DIM = 1024 # joint text-vision embedding (unused by vision path) +# CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073] +# CLIP_STD = [0.26862954, 0.26130258, 0.27577711] + + +# ══════════════════════════════════════════════════════════════════════════════ +# VAE encoder / decoder +# ══════════════════════════════════════════════════════════════════════════════ + +VAE_Z_DIM = 16 +VAE_SPATIAL_FACTOR = 8 +VAE_TEMPORAL_FACTOR = 4 +VAE_ENC_CH = [96, 192, 384, 768, 1536] +VAE_DEC_CH = [1536, 768, 384, 192, 96] + + +# ══════════════════════════════════════════════════════════════════════════════ +# DiT denoising network +# ══════════════════════════════════════════════════════════════════════════════ + +DIT_DIM = 128 +DIT_HEADS = 4 +DIT_HEAD_DIM = 32 +DIT_FFN = 256 +DIT_LAYERS = 1 +DIT_IN_DIM = 16 +DIT_TEXT_DIM = 128 +DIT_FREQ_DIM = 64 +DIT_EPS = 1e-5 + + +# ══════════════════════════════════════════════════════════════════════════════ +# Scheduler +# ══════════════════════════════════════════════════════════════════════════════ + +NUM_STEPS = 1 +SHIFT = 5.0 +CFG_SCALE = 5.0 + + +# ══════════════════════════════════════════════════════════════════════════════ +# Latent space (derived) +# ══════════════════════════════════════════════════════════════════════════════ + +LAT_F = 2 # (CHUNK_SIZE - 1) // VAE_TEMPORAL_FACTOR + 1 +LAT_H = 8 # H // VAE_SPATIAL_FACTOR +LAT_W = 8 # W // VAE_SPATIAL_FACTOR diff --git a/models/world_model/dit_forward.py b/models/world_model/dit_forward.py new file mode 100644 index 00000000..a27a5dc7 --- /dev/null +++ b/models/world_model/dit_forward.py @@ -0,0 +1,1020 @@ +# Copyright (c) PyPTO Contributors. +# Licensed under CANN Open Software License Agreement Version 2.0. +# ----------------------------------------------------------------------------- +"""DiT (Diffusion Transformer) for Fun-Control 1.3B — pypto3.0 implementation. + +Single @pl.jit kernel implementing the full DiT forward pass matching +``test_golden_fun_control_full.py::dit_forward`` (golden:1364-1489). + +Architecture:: + + ref_conv + patch_conv + concat → x_full + text_proj(GELU-tanh) + clip_proj(LN→FC→GELU→FC→LN) → ctx + AdaLN → Self-Attn(RoPE) → Cross-Attn(text+img) → FFN → Head → output + +All arithmetic in NPU kernels. Host prepares im2col (data rearrangement), +RoPE cos/sin tables (lookup), and the static sinusoidal frequency table +(``time_freqs``). The full timestep embedding — ``sin_emb = [cos(t*freqs), +sin(t*freqs)]`` followed by a 3-layer SiLU MLP — runs on NPU. + +Known differences from the golden reference: + • CLIP projection uses GELU-tanh approximation in the kernel (pypto lacks + ``erf`` for exact GELU); the golden uses ``F.gelu()`` (exact). This + difference is intrinsic to the NPU compute model and is tracked by the + permissive K3 tolerance. + +Usage:: + + python dit_forward.py -p a2a3 -d 0 +""" + +import argparse +import math +import sys + +import pypto.language as pl +import torch +import torch.nn.functional as F + +sys.path.insert( + 0, '/data/x00952168/pypto3.0/cann-recipes-embodied-ai/world_model/' + 'agibot-arm-world-model/infer_with_torch') +from test_golden_fun_control_full import ( # noqa: E402 + dit_forward as golden_dit_fn, + sinusoidal_embedding_1d, precompute_freqs_cis_3d, make_weights, +) +from config import ( + DIT_DIM, DIT_HEADS, DIT_HEAD_DIM, DIT_FFN, DIT_IN_DIM, + DIT_TEXT_DIM, DIT_FREQ_DIM, DIT_EPS, + CLIP_DIM, CLIP_TOKENS, T5_SEQ, T5_DIM, LAT_F, LAT_H, LAT_W, +) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Derived constants +# ══════════════════════════════════════════════════════════════════════════════ + +COND_CH = DIT_IN_DIM + 4 + DIT_IN_DIM +FP = LAT_F +HP = LAT_H // 2 +WP = LAT_W // 2 +REF_N = HP * WP +X_N = FP * HP * WP +S = REF_N + X_N +T = 16 +CLIP_PAD = ((CLIP_TOKENS + T - 1) // T) * T + +# 3D RoPE dimension splits +ROPE_F_DIM = DIT_HEAD_DIM - 2 * (DIT_HEAD_DIM // 3) +ROPE_H_DIM = DIT_HEAD_DIM // 3 +ROPE_W_DIM = DIT_HEAD_DIM // 3 +HEAD_HALF = DIT_HEAD_DIM // 2 + +# Numerical constants +HEAD_SCALE = 1.0 / math.sqrt(DIT_HEAD_DIM) +GELU_SQRT_2_OVER_PI = 0.7978845608 +GELU_COEFF = 0.044715 +LN_EPS = 1e-5 +DIM_INV = 1.0 / DIT_DIM + +REF_CONV_COL = T * 4 +PATCH_CONV_COL = T * 13 +TIME_PROJ_N = 6 * DIT_DIM # 768 +TIME_N_CHUNK = 64 # FC3 N-chunk for Right buffer safety (64*128*2=16384 < 65536) + +# Geometry assertions +assert DIT_DIM % DIT_HEADS == 0, "DIT_HEADS must divide DIT_DIM" +assert DIT_HEAD_DIM * DIT_HEADS == DIT_DIM +assert S % T == 0, "S must be a multiple of T" +assert CLIP_PAD % T == 0, "CLIP_PAD must be a multiple of T" +assert T5_SEQ % T == 0, "T5_SEQ must be a multiple of T" +assert X_N % T == 0, "X_N must be a multiple of T" +assert REF_N % T == 0, "REF_N must be a multiple of T" +assert DIT_HEAD_DIM % 2 == 0, "DIT_HEAD_DIM must be even (RoPE half-split)" + + +# ══════════════════════════════════════════════════════════════════════════════ +# Host helpers — data rearrangement only (no arithmetic). +# ══════════════════════════════════════════════════════════════════════════════ + +def _im2col_ref_conv2d(ref): + """im2col for ref Conv2d(kernel=2x2, stride=2). Data rearrangement only.""" + cols = F.unfold(ref.float(), kernel_size=(2, 2), stride=2) + return cols.squeeze(0).transpose(0, 1).contiguous().bfloat16() + + +def _im2col_patch_conv3d(x): + """im2col for patch Conv3d(kernel=(1,2,2), stride=(1,2,2)). Data rearrangement only.""" + B, C, D, H, W = x.shape + col_size = C * 1 * 2 * 2 + xf = x.float() + rows = [] + for n in range(B): + for d in range(D): + for h in range(H // 2): + for w in range(W // 2): + rows.append(xf[n, :, d, h * 2:h * 2 + 2, w * 2:w * 2 + 2].reshape(col_size)) + return torch.stack(rows).contiguous().bfloat16() + + +def _precompute_rope_cos_sin(): + """Precompute 3D RoPE cos/sin tables. Data rearrangement only.""" + freq_f, freq_h, freq_w = precompute_freqs_cis_3d(DIT_HEAD_DIM) + return (freq_f.real.float(), freq_f.imag.float(), + freq_h.real.float(), freq_h.imag.float(), + freq_w.real.float(), freq_w.imag.float()) + + +def _build_rope_map(cos_f, sin_f, cos_h, sin_h, cos_w, sin_w, fl, hl, wl): + """Build per-position RoPE cos/sin lookup tables for complex-pair rotation. + + The kernel applies complex-pair rotation via pl.gather (stride-2 even/odd + split). Shape: ``[S, DIT_DIM//2]`` — half-dim frequencies repeated across + all heads. + """ + total = fl * hl * wl + cos_table = torch.zeros(total, HEAD_HALF) + sin_table = torch.zeros(total, HEAD_HALF) + idx = 0 + for fi in range(fl): + for hi in range(hl): + for wi in range(wl): + for j in range(ROPE_F_DIM // 2): + cos_table[idx, j] = cos_f[fi, j] + sin_table[idx, j] = sin_f[fi, j] + for j in range(ROPE_H_DIM // 2): + col = ROPE_F_DIM // 2 + j + cos_table[idx, col] = cos_h[hi, j] + sin_table[idx, col] = sin_h[hi, j] + for j in range(ROPE_W_DIM // 2): + col = ROPE_F_DIM // 2 + ROPE_H_DIM // 2 + j + cos_table[idx, col] = cos_w[wi, j] + sin_table[idx, col] = sin_w[wi, j] + idx += 1 + return cos_table.repeat(1, DIT_HEADS), sin_table.repeat(1, DIT_HEADS) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Fused kernel — full DiT forward pass. +# ══════════════════════════════════════════════════════════════════════════════ + +@pl.jit +def dit_forward( + # ── Conv inputs (im2col precomputed on host — data rearrangement) ── + ref_col: pl.Tensor[[REF_N, REF_CONV_COL], pl.BF16], + patch_col: pl.Tensor[[X_N, PATCH_CONV_COL], pl.BF16], + ref_conv_w: pl.Tensor[[REF_CONV_COL, DIT_DIM], pl.BF16], + patch_conv_w: pl.Tensor[[PATCH_CONV_COL, DIT_DIM], pl.BF16], + # ── Text + CLIP projection inputs ── + text_raw: pl.Tensor[[T5_SEQ, T5_DIM], pl.BF16], + text_fc1_w: pl.Tensor[[DIT_DIM, DIT_TEXT_DIM], pl.BF16], + text_fc2_w: pl.Tensor[[DIT_DIM, DIT_DIM], pl.BF16], + clip_raw: pl.Tensor[[CLIP_PAD, CLIP_DIM], pl.BF16], + clip_ln1_w: pl.Tensor[[1, CLIP_DIM], pl.FP32], + clip_ln1_b: pl.Tensor[[1, CLIP_DIM], pl.FP32], + clip_fc1_w: pl.Tensor[[CLIP_DIM, CLIP_DIM], pl.BF16], + clip_fc2_w: pl.Tensor[[DIT_DIM, CLIP_DIM], pl.BF16], + clip_ln2_w: pl.Tensor[[1, DIT_DIM], pl.FP32], + clip_ln2_b: pl.Tensor[[1, DIT_DIM], pl.FP32], + # ── Context tables (precomputed on host) ── + clip_mask: pl.Tensor[[1, CLIP_PAD], pl.FP32], + rope_cos: pl.Tensor[[S, DIT_DIM // 2], pl.FP32], + rope_sin: pl.Tensor[[S, DIT_DIM // 2], pl.FP32], + # ── Timestep embedding (computed on NPU) ── + timestep: pl.Tensor[[1], pl.FP32], + time_freqs: pl.Tensor[[T, DIT_FREQ_DIM // 2], pl.FP32], + time_fc1_w: pl.Tensor[[DIT_DIM, DIT_FREQ_DIM], pl.BF16], + time_fc2_w: pl.Tensor[[DIT_DIM, DIT_DIM], pl.BF16], + time_proj_w: pl.Tensor[[6 * DIT_DIM, DIT_DIM], pl.BF16], + # ── Layer weights ── + l_mod: pl.Tensor[[1, 6, DIT_DIM], pl.BF16], + l_norm1_w: pl.Tensor[[1, DIT_DIM], pl.FP32], + l_norm1_b: pl.Tensor[[1, DIT_DIM], pl.FP32], + l_q_w: pl.Tensor[[DIT_DIM, DIT_DIM], pl.BF16], + l_k_w: pl.Tensor[[DIT_DIM, DIT_DIM], pl.BF16], + l_v_w: pl.Tensor[[DIT_DIM, DIT_DIM], pl.BF16], + l_o_w: pl.Tensor[[DIT_DIM, DIT_DIM], pl.BF16], + l_q_rms: pl.Tensor[[1, DIT_DIM], pl.FP32], + l_k_rms: pl.Tensor[[1, DIT_DIM], pl.FP32], + l_norm3_w: pl.Tensor[[1, DIT_DIM], pl.FP32], + l_norm3_b: pl.Tensor[[1, DIT_DIM], pl.FP32], + l_cross_q: pl.Tensor[[DIT_DIM, DIT_DIM], pl.BF16], + l_cross_ktxt: pl.Tensor[[DIT_DIM, DIT_DIM], pl.BF16], + l_cross_vtxt: pl.Tensor[[DIT_DIM, DIT_DIM], pl.BF16], + l_cross_kimg: pl.Tensor[[DIT_DIM, DIT_DIM], pl.BF16], + l_cross_vimg: pl.Tensor[[DIT_DIM, DIT_DIM], pl.BF16], + l_cross_o: pl.Tensor[[DIT_DIM, DIT_DIM], pl.BF16], + l_cross_qrms: pl.Tensor[[1, DIT_DIM], pl.FP32], + l_cross_krms: pl.Tensor[[1, DIT_DIM], pl.FP32], + l_cross_kir: pl.Tensor[[1, DIT_DIM], pl.FP32], + l_norm2_w: pl.Tensor[[1, DIT_DIM], pl.FP32], + l_norm2_b: pl.Tensor[[1, DIT_DIM], pl.FP32], + l_ffn1: pl.Tensor[[DIT_FFN, DIT_DIM], pl.BF16], + l_ffn2: pl.Tensor[[DIT_DIM, DIT_FFN], pl.BF16], + # ── Head weights ── + head_mod_w: pl.Tensor[[1, 2, DIT_DIM], pl.BF16], + head_ln_w: pl.Tensor[[1, DIT_DIM], pl.FP32], + head_ln_b: pl.Tensor[[1, DIT_DIM], pl.FP32], + head_fc: pl.Tensor[[DIT_IN_DIM * 4, DIT_DIM], pl.BF16], + # ── Output ── + out: pl.Out[pl.Tensor[[X_N, DIT_IN_DIM * 4], pl.BF16]], +): + + # ══════════════════════════════════════════════════════════════════════ + # Scope 0: Timestep embedding — sinusoidal + 3-layer SiLU MLP on NPU. + # FC3 is N-chunked to stay within the 65536-byte Right buffer. + # ══════════════════════════════════════════════════════════════════════ + t_mod_gm = pl.create_tensor([6, DIT_DIM], dtype=pl.BF16) + te_fc2_gm = pl.create_tensor([T, DIT_DIM], dtype=pl.FP32) + te_fc3_out = pl.create_tensor([T, TIME_PROJ_N], dtype=pl.FP32) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="time_embed"): + t_scalar = pl.read(timestep, [0]) + freqs_tile = pl.slice(time_freqs, [T, DIT_FREQ_DIM // 2], [0, 0]) + args = pl.mul(freqs_tile, t_scalar) + + # sin_emb = [cos(args) | sin(args)] → FC1 split into two halves + # to avoid assembling cos/sin tiles (tmov valid-shape mismatch). + cos_part = pl.cos(args) + sin_part = pl.sin(args) + cos_bf = pl.cast(cos_part, target_type=pl.BF16) + sin_bf = pl.cast(sin_part, target_type=pl.BF16) + + half_k = DIT_FREQ_DIM // 2 + w_cos = pl.slice(time_fc1_w, [DIT_DIM, half_k], [0, 0]) + w_sin = pl.slice(time_fc1_w, [DIT_DIM, half_k], [0, half_k]) + fc1_cos = pl.matmul(cos_bf, w_cos, b_trans=True, out_dtype=pl.FP32) + fc1_sin = pl.matmul(sin_bf, w_sin, b_trans=True, out_dtype=pl.FP32) + te_fc1 = pl.add(fc1_cos, fc1_sin) + + # SiLU + te_silu_1 = pl.mul(te_fc1, pl.recip(pl.add(pl.exp(pl.neg(te_fc1)), 1.0))) + + # FC2 → store to GM for N-chunked FC3 below + te_fc2 = pl.matmul(pl.cast(te_silu_1, target_type=pl.BF16), time_fc2_w, b_trans=True, out_dtype=pl.FP32) + te_fc2_gm = pl.assemble(te_fc2_gm, te_fc2, [0, 0]) + + # FC3 with N-chunked accumulation — each chunk in its own pl.at. + for n_chunk in pl.unroll(TIME_PROJ_N // TIME_N_CHUNK): + n0 = n_chunk * TIME_N_CHUNK + with pl.at(level=pl.Level.CORE_GROUP, name_hint="time_fc3"): + te_fc2_tile = pl.slice(te_fc2_gm, [T, DIT_DIM], [0, 0]) + # SiLU on the loaded tile + te_silu = pl.mul(te_fc2_tile, pl.recip(pl.add(pl.exp(pl.neg(te_fc2_tile)), 1.0))) + te_silu_bf = pl.cast(te_silu, target_type=pl.BF16) + te_w_chunk = pl.slice(time_proj_w, [TIME_N_CHUNK, DIT_DIM], [n0, 0]) + te_chunk_mm = pl.matmul(te_silu_bf, te_w_chunk, b_trans=True, out_dtype=pl.FP32) + te_fc3_out = pl.assemble(te_fc3_out, te_chunk_mm, [0, n0]) + + with pl.at(level=pl.Level.CORE_GROUP, name_hint="time_final"): + for mod_i in pl.unroll(6): + mod_slice = pl.slice(te_fc3_out, [1, DIT_DIM], [0, mod_i * DIT_DIM]) + t_mod_gm = pl.assemble(t_mod_gm, pl.cast(mod_slice, target_type=pl.BF16), [mod_i, 0]) + + # ══════════════════════════════════════════════════════════════════════ + # Phase A: Conv + Concat (was K1) + # ══════════════════════════════════════════════════════════════════════ + + # ── Scope 1: Reference Conv2d matmul ── + ref_tokens = pl.create_tensor([REF_N, DIT_DIM], dtype=pl.BF16) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ref_conv"): + ref_tokens = pl.cast( + pl.matmul(ref_col, ref_conv_w, out_dtype=pl.FP32), + target_type=pl.BF16, + ) + + # ── Scope 2: Patch Conv3d K-chunked matmul ── + patch_tokens = pl.create_tensor([X_N, DIT_DIM], dtype=pl.BF16) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="patch_conv"): + for row_start in pl.range(0, X_N, T): + col_chunk = pl.slice(patch_col, [T, T], [row_start, 0]) + w_chunk = pl.slice(patch_conv_w, [T, DIT_DIM], [0, 0]) + acc = pl.matmul(col_chunk, w_chunk, out_dtype=pl.FP32) + for k_idx in pl.range(1, 13): + col_chunk = pl.slice(patch_col, [T, T], [row_start, k_idx * T]) + w_chunk = pl.slice(patch_conv_w, [T, DIT_DIM], [k_idx * T, 0]) + acc = pl.matmul_acc(acc, col_chunk, w_chunk) + patch_tokens = pl.assemble( + patch_tokens, pl.cast(acc, target_type=pl.BF16), [row_start, 0], + ) + + # ── Scope 3: Concat ref + patch → x_full ── + x_full = pl.create_tensor([S, DIT_DIM], dtype=pl.BF16) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="concat"): + for row in pl.range(0, REF_N, T): + x_full = pl.assemble( + x_full, pl.slice(ref_tokens, [T, DIT_DIM], [row, 0]), [row, 0], + ) + for row in pl.range(0, X_N, T): + x_full = pl.assemble( + x_full, pl.slice(patch_tokens, [T, DIT_DIM], [row, 0]), + [REF_N + row, 0], + ) + + # ══════════════════════════════════════════════════════════════════════ + # Phase B: Text + CLIP Projection (was K2) + # ══════════════════════════════════════════════════════════════════════ + + text_ctx = pl.create_tensor([T5_SEQ, DIT_DIM], dtype=pl.BF16) + # ── Scope 4: Text projection with GELU-tanh ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="text_proj"): + for row_start in pl.range(0, T5_SEQ, T): + text_tile = pl.slice(text_raw, [T, T5_DIM], [row_start, 0]) + fc1_out = pl.matmul(text_tile, text_fc1_w, b_trans=True, out_dtype=pl.FP32) + # GELU-tanh: 0.5*x*(1+tanh(sqrt(2/pi)*(x+0.044715*x^3))) + x_cubed = pl.mul(pl.mul(fc1_out, fc1_out), fc1_out) + gelu_inner = pl.add(fc1_out, pl.mul(x_cubed, GELU_COEFF)) + gelu_scaled = pl.mul(gelu_inner, GELU_SQRT_2_OVER_PI) + gelu_2z = pl.mul(gelu_scaled, 2.0) + gelu_exp = pl.exp(gelu_2z) + gelu_recip = pl.recip(pl.add(gelu_exp, 1.0)) + tanh_val = pl.sub( + pl.add(pl.sub(fc1_out, fc1_out), 1.0), + pl.add(gelu_recip, gelu_recip), + ) + gelu_out = pl.mul(pl.mul(fc1_out, 0.5), pl.add(tanh_val, 1.0)) + fc2_out = pl.matmul( + pl.cast(gelu_out, target_type=pl.BF16), + text_fc2_w, b_trans=True, out_dtype=pl.FP32, + ) + text_ctx = pl.assemble( + text_ctx, pl.cast(fc2_out, target_type=pl.BF16), [row_start, 0], + ) + + clip_ctx = pl.create_tensor([CLIP_PAD, DIT_DIM], dtype=pl.BF16) + # ── Scope 5: CLIP projection (LN → FC1 → GELU → FC2 → LN) ── + # NOTE: golden reference uses exact F.gelu() for CLIP projection; + # the kernel uses GELU-tanh approximation because pypto lacks erf. + with pl.at(level=pl.Level.CORE_GROUP, name_hint="clip_proj"): + for row_start in pl.range(0, CLIP_PAD, T): + clip_tile = pl.slice(clip_raw, [T, CLIP_DIM], [row_start, 0]) + clip_fp32 = pl.cast(clip_tile, target_type=pl.FP32) + + # LayerNorm 1 + clip_scaled = pl.mul(clip_fp32, 1.0 / CLIP_DIM) + clip_mean = pl.row_sum(clip_scaled) + clip_centered = pl.row_expand_sub(clip_fp32, clip_mean) + clip_var = pl.row_sum(pl.mul(clip_centered, clip_centered)) + clip_var = pl.mul(clip_var, 1.0 / CLIP_DIM) + clip_var_eps = pl.add(clip_var, LN_EPS) + clip_var_eps = pl.reshape(clip_var_eps, [1, T]) + clip_std = pl.sqrt(clip_var_eps) + clip_std = pl.reshape(clip_std, [T, 1]) + clip_normed = pl.col_expand_add( + pl.col_expand_mul(pl.row_expand_div(clip_centered, clip_std), + clip_ln1_w[:, :]), + clip_ln1_b[:, :], + ) + + # FC1 + GELU-tanh + clip_fc1_out = pl.matmul( + pl.cast(clip_normed, target_type=pl.BF16), + clip_fc1_w, b_trans=True, out_dtype=pl.FP32, + ) + clip_x3 = pl.mul(pl.mul(clip_fc1_out, clip_fc1_out), clip_fc1_out) + clip_inner = pl.add(clip_fc1_out, pl.mul(clip_x3, GELU_COEFF)) + clip_sc = pl.mul(clip_inner, GELU_SQRT_2_OVER_PI) + clip_2z = pl.mul(clip_sc, 2.0) + clip_exp = pl.exp(clip_2z) + clip_recip = pl.recip(pl.add(clip_exp, 1.0)) + clip_tanh = pl.sub( + pl.add(pl.sub(clip_fc1_out, clip_fc1_out), 1.0), + pl.add(clip_recip, clip_recip), + ) + clip_gelu = pl.mul(pl.mul(clip_fc1_out, 0.5), pl.add(clip_tanh, 1.0)) + + # FC2 + clip_fc2_out = pl.matmul( + pl.cast(clip_gelu, target_type=pl.BF16), + clip_fc2_w, b_trans=True, out_dtype=pl.FP32, + ) + + # LayerNorm 2 + fc2_scaled = pl.mul(clip_fc2_out, DIM_INV) + fc2_mean = pl.row_sum(fc2_scaled) + fc2_centered = pl.row_expand_sub(clip_fc2_out, fc2_mean) + fc2_var = pl.row_sum(pl.mul(fc2_centered, fc2_centered)) + fc2_var = pl.mul(fc2_var, DIM_INV) + fc2_var_eps = pl.add(fc2_var, DIT_EPS) + fc2_var_eps = pl.reshape(fc2_var_eps, [1, T]) + fc2_std = pl.sqrt(fc2_var_eps) + fc2_std = pl.reshape(fc2_std, [T, 1]) + clip_out = pl.col_expand_add( + pl.col_expand_mul(pl.row_expand_div(fc2_centered, fc2_std), + clip_ln2_w[:, :]), + clip_ln2_b[:, :], + ) + clip_ctx = pl.assemble( + clip_ctx, pl.cast(clip_out, target_type=pl.BF16), [row_start, 0], + ) + + # ══════════════════════════════════════════════════════════════════════ + # Phase C: DiT Blocks + Head (was K3) + # ══════════════════════════════════════════════════════════════════════ + + # ── AdaLN modulation vectors (6 per layer) ── + mod_shift_msa = pl.create_tensor([1, DIT_DIM], dtype=pl.FP32) + mod_scale_msa = pl.create_tensor([1, DIT_DIM], dtype=pl.FP32) + mod_gate_msa = pl.create_tensor([1, DIT_DIM], dtype=pl.FP32) + mod_shift_mlp = pl.create_tensor([1, DIT_DIM], dtype=pl.FP32) + mod_scale_mlp = pl.create_tensor([1, DIT_DIM], dtype=pl.FP32) + mod_gate_mlp = pl.create_tensor([1, DIT_DIM], dtype=pl.FP32) + + # Self-attention buffers + sa_q_gm = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + sa_k_gm = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + sa_v_gm = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + sa_q_rope = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + sa_k_rope = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + sa_attn_out = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + sa_oproj_fp = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + x_after_sa = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + + # Cross-attention buffers + ca_q_gm = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + ca_k_txt_gm = pl.create_tensor([T5_SEQ, DIT_DIM], dtype=pl.FP32) + ca_v_txt_gm = pl.create_tensor([T5_SEQ, DIT_DIM], dtype=pl.FP32) + ca_txt_out = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + ca_k_img_gm = pl.create_tensor([CLIP_PAD, DIT_DIM], dtype=pl.FP32) + ca_v_img_gm = pl.create_tensor([CLIP_PAD, DIT_DIM], dtype=pl.FP32) + ca_img_out = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + + # Post-cross-attn / FFN / Head buffers + x_after_ca = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + x_after_ffn = pl.create_tensor([S, DIT_DIM], dtype=pl.FP32) + head_out_gm = pl.create_tensor([S, DIT_IN_DIM * 4], dtype=pl.BF16) + head_shift = pl.create_tensor([1, DIT_DIM], dtype=pl.FP32) + head_scale = pl.create_tensor([1, DIT_DIM], dtype=pl.FP32) + + # x_full copy for GM safety + x_full_gm = pl.create_tensor([S, DIT_DIM], dtype=pl.BF16) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="copy_input"): + for row_start in pl.range(0, S, T): + x_full_gm = pl.assemble( + x_full_gm, + pl.slice(x_full, [T, DIT_DIM], [row_start, 0]), + [row_start, 0], + ) + + # ── Scope 6: AdaLN Modulation ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="adaln_mod"): + layer_mod_fp = pl.reshape(pl.cast(l_mod, target_type=pl.FP32), [6, DIT_DIM]) + t_mod_fp = pl.cast(t_mod_gm, target_type=pl.FP32) + combined_mod = pl.add(layer_mod_fp, t_mod_fp) + mod_shift_msa = pl.assemble(mod_shift_msa, pl.slice(combined_mod, [1, DIT_DIM], [0, 0]), [0, 0]) + mod_scale_msa = pl.assemble(mod_scale_msa, pl.slice(combined_mod, [1, DIT_DIM], [1, 0]), [0, 0]) + mod_gate_msa = pl.assemble(mod_gate_msa, pl.slice(combined_mod, [1, DIT_DIM], [2, 0]), [0, 0]) + mod_shift_mlp = pl.assemble(mod_shift_mlp, pl.slice(combined_mod, [1, DIT_DIM], [3, 0]), [0, 0]) + mod_scale_mlp = pl.assemble(mod_scale_mlp, pl.slice(combined_mod, [1, DIT_DIM], [4, 0]), [0, 0]) + mod_gate_mlp = pl.assemble(mod_gate_mlp, pl.slice(combined_mod, [1, DIT_DIM], [5, 0]), [0, 0]) + + # ── Scope 7: Self-Attention — LN → AdaLN → QKV projection ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="sa_qkv"): + shift_msa_tile = pl.slice(mod_shift_msa, [1, DIT_DIM], [0, 0]) + scale_msa_tile = pl.slice(mod_scale_msa, [1, DIT_DIM], [0, 0]) + scale_plus_1 = pl.add(scale_msa_tile, 1.0) + + for row_start in pl.range(0, S, T): + x_tile = pl.slice(x_full_gm, [T, DIT_DIM], [row_start, 0]) + x_fp = pl.cast(x_tile, target_type=pl.FP32) + + x_scaled = pl.mul(x_fp, DIM_INV) + x_mean = pl.row_sum(x_scaled) + x_centered = pl.row_expand_sub(x_fp, x_mean) + x_var = pl.row_sum(pl.mul(x_centered, x_centered)) + x_var = pl.mul(x_var, DIM_INV) + x_var_eps = pl.add(x_var, LN_EPS) + x_var_eps = pl.reshape(x_var_eps, [1, T]) + x_std = pl.sqrt(x_var_eps) + x_std = pl.reshape(x_std, [T, 1]) + x_normed = pl.col_expand_add( + pl.col_expand_mul(pl.row_expand_div(x_centered, x_std), + l_norm1_w[:, :]), + l_norm1_b[:, :], + ) + x_modulated = pl.col_expand_add( + pl.col_expand_mul(x_normed, scale_plus_1), + shift_msa_tile, + ) + x_mod_bf = pl.cast(x_modulated, target_type=pl.BF16) + + sa_q_gm = pl.assemble(sa_q_gm, + pl.matmul(x_mod_bf, l_q_w, b_trans=True, out_dtype=pl.FP32), [row_start, 0]) + sa_k_gm = pl.assemble(sa_k_gm, + pl.matmul(x_mod_bf, l_k_w, b_trans=True, out_dtype=pl.FP32), [row_start, 0]) + sa_v_gm = pl.assemble(sa_v_gm, + pl.matmul(x_mod_bf, l_v_w, b_trans=True, out_dtype=pl.FP32), [row_start, 0]) + + # ── Scope 8: Q/K RMSNorm + 3D RoPE ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="sa_rms_rope"): + for row_start in pl.range(0, S, T): + q_tile = pl.slice(sa_q_gm, [T, DIT_DIM], [row_start, 0]) + k_tile = pl.slice(sa_k_gm, [T, DIT_DIM], [row_start, 0]) + cos_tile = pl.slice(rope_cos, [T, DIT_DIM // 2], [row_start, 0]) + sin_tile = pl.slice(rope_sin, [T, DIT_DIM // 2], [row_start, 0]) + + # Q RMSNorm + q_sq = pl.mul(q_tile, q_tile) + q_sq_sum = pl.row_sum(q_sq) + q_sq_sum = pl.reshape(q_sq_sum, [1, T]) + q_rms_var = pl.add(pl.mul(q_sq_sum, DIM_INV), DIT_EPS) + q_inv_rms = pl.recip(pl.sqrt(q_rms_var)) + q_inv_rms = pl.reshape(q_inv_rms, [T, 1]) + q_normed = pl.col_expand_mul(pl.row_expand_mul(q_tile, q_inv_rms), l_q_rms[:, :]) + + # K RMSNorm + k_sq = pl.mul(k_tile, k_tile) + k_sq_sum = pl.row_sum(k_sq) + k_sq_sum = pl.reshape(k_sq_sum, [1, T]) + k_rms_var = pl.add(pl.mul(k_sq_sum, DIM_INV), DIT_EPS) + k_inv_rms = pl.recip(pl.sqrt(k_rms_var)) + k_inv_rms = pl.reshape(k_inv_rms, [T, 1]) + k_normed = pl.col_expand_mul(pl.row_expand_mul(k_tile, k_inv_rms), l_k_rms[:, :]) + + # Complex-pair RoPE via pl.gather stride-2 even/odd split + q_even = pl.gather(q_normed, mask_pattern=pl.tile.MaskPattern.P0101) + q_odd = pl.gather(q_normed, mask_pattern=pl.tile.MaskPattern.P1010) + k_even = pl.gather(k_normed, mask_pattern=pl.tile.MaskPattern.P0101) + k_odd = pl.gather(k_normed, mask_pattern=pl.tile.MaskPattern.P1010) + + q_re = pl.sub(pl.mul(q_even, cos_tile), pl.mul(q_odd, sin_tile)) + q_ro = pl.add(pl.mul(q_odd, cos_tile), pl.mul(q_even, sin_tile)) + k_re = pl.sub(pl.mul(k_even, cos_tile), pl.mul(k_odd, sin_tile)) + k_ro = pl.add(pl.mul(k_odd, cos_tile), pl.mul(k_even, sin_tile)) + + q_buf = pl.full([T, DIT_DIM], dtype=pl.FP32, value=0.0) + q_buf = pl.tensor.scatter(q_re, mask_pattern=pl.tile.MaskPattern.P0101, dst=q_buf) + q_buf = pl.tensor.scatter(q_ro, mask_pattern=pl.tile.MaskPattern.P1010, dst=q_buf) + sa_q_rope = pl.assemble(sa_q_rope, q_buf, [row_start, 0]) + + k_buf = pl.full([T, DIT_DIM], dtype=pl.FP32, value=0.0) + k_buf = pl.tensor.scatter(k_re, mask_pattern=pl.tile.MaskPattern.P0101, dst=k_buf) + k_buf = pl.tensor.scatter(k_ro, mask_pattern=pl.tile.MaskPattern.P1010, dst=k_buf) + sa_k_rope = pl.assemble(sa_k_rope, k_buf, [row_start, 0]) + + # ── Scope 9: Self-Attention MHA (per-head, scaled) ── + attn_scores_tmp = pl.create_tensor([S, S], dtype=pl.FP32) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="sa_mha"): + for head_idx in pl.unroll(DIT_HEADS): + head_col = head_idx * DIT_HEAD_DIM + sa_q_head = pl.slice(sa_q_rope, [S, DIT_HEAD_DIM], [0, head_col]) + sa_k_head = pl.slice(sa_k_rope, [S, DIT_HEAD_DIM], [0, head_col]) + sa_v_head = pl.slice(sa_v_gm, [S, DIT_HEAD_DIM], [0, head_col]) + + sa_scores = pl.matmul(sa_q_head, sa_k_head, b_trans=True, out_dtype=pl.FP32) + sa_scores = pl.mul(sa_scores, HEAD_SCALE) + + for row_start in pl.range(0, S, T): + score_tile = pl.slice(sa_scores, [T, S], [row_start, 0]) + shifted = pl.row_expand_sub(score_tile, pl.row_max(score_tile)) + exp_scores = pl.exp(shifted) + attn_weights = pl.row_expand_div(exp_scores, pl.row_sum(exp_scores)) + attn_scores_tmp = pl.assemble( + attn_scores_tmp, attn_weights, [row_start, 0], + ) + + attn_weights_bf = pl.cast( + pl.slice(attn_scores_tmp, [S, S], [0, 0]), + target_type=pl.BF16, + ) + sa_head_out = pl.matmul(attn_weights_bf, pl.cast(sa_v_head, target_type=pl.BF16), out_dtype=pl.FP32) + sa_attn_out = pl.assemble(sa_attn_out, sa_head_out, [0, head_col]) + + # ── Scope 10: Self-Attention Output Projection ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="sa_oproj"): + oproj_result = pl.matmul(pl.cast(sa_attn_out, target_type=pl.BF16), l_o_w, b_trans=True, out_dtype=pl.FP32) + for row_start in pl.range(0, S, T): + sa_oproj_fp = pl.assemble( + sa_oproj_fp, + pl.slice(oproj_result, [T, DIT_DIM], [row_start, 0]), + [row_start, 0], + ) + + # ── Scope 11: Gated Residual (self-attention) ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="sa_gate"): + gate_msa_tile = pl.slice(mod_gate_msa, [1, DIT_DIM], [0, 0]) + for row_start in pl.range(0, S, T): + x_orig = pl.slice(x_full_gm, [T, DIT_DIM], [row_start, 0]) + x_orig_fp = pl.cast(x_orig, target_type=pl.FP32) + oproj_tile = pl.slice(sa_oproj_fp, [T, DIT_DIM], [row_start, 0]) + gated_oproj = pl.col_expand_mul(oproj_tile, gate_msa_tile) + x_after_sa = pl.assemble( + x_after_sa, pl.add(x_orig_fp, gated_oproj), [row_start, 0], + ) + + # ── Scope 12: Cross-Attention LN + Q projection + RMSNorm ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ca_layernorm"): + for row_start in pl.range(0, S, T): + ca_x_tile = pl.slice(x_after_sa, [T, DIT_DIM], [row_start, 0]) + ca_x_fp = ca_x_tile + ca_scaled = pl.mul(ca_x_fp, DIM_INV) + ca_mean = pl.row_sum(ca_scaled) + ca_centered = pl.row_expand_sub(ca_x_fp, ca_mean) + ca_var = pl.row_sum(pl.mul(ca_centered, ca_centered)) + ca_var = pl.mul(ca_var, DIM_INV) + ca_var_eps = pl.add(ca_var, LN_EPS) + ca_var_eps = pl.reshape(ca_var_eps, [1, T]) + ca_std = pl.sqrt(ca_var_eps) + ca_std = pl.reshape(ca_std, [T, 1]) + ca_normed = pl.col_expand_add( + pl.col_expand_mul(pl.row_expand_div(ca_centered, ca_std), + l_norm3_w[:, :]), + l_norm3_b[:, :], + ) + ca_q_gm = pl.assemble(ca_q_gm, ca_normed, [row_start, 0]) + + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ca_q_proj"): + for row_start in pl.range(0, S, T): + ca_q_tile = pl.slice(ca_q_gm, [T, DIT_DIM], [row_start, 0]) + ca_q_proj = pl.matmul(pl.cast(ca_q_tile, target_type=pl.BF16), l_cross_q, b_trans=True, out_dtype=pl.FP32) + ca_q_sq = pl.mul(ca_q_proj, ca_q_proj) + ca_q_sqsum = pl.row_sum(ca_q_sq) + ca_q_sqsum = pl.reshape(ca_q_sqsum, [1, T]) + ca_q_rms = pl.add(pl.mul(ca_q_sqsum, DIM_INV), DIT_EPS) + ca_q_inv = pl.recip(pl.sqrt(ca_q_rms)) + ca_q_inv = pl.reshape(ca_q_inv, [T, 1]) + ca_q_normed = pl.col_expand_mul( + pl.row_expand_mul(ca_q_proj, ca_q_inv), l_cross_qrms[:, :]) + ca_q_gm = pl.assemble(ca_q_gm, ca_q_normed, [row_start, 0]) + + # ── Scope 13: Cross-Attention Text K/V projection + RMSNorm + MHA ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ca_text_kv"): + for row_start in pl.range(0, T5_SEQ, T): + txt_tile = pl.slice(text_ctx, [T, DIT_DIM], [row_start, 0]) + + txt_k_proj = pl.matmul(txt_tile, l_cross_ktxt, b_trans=True, out_dtype=pl.FP32) + txt_k_sq = pl.mul(txt_k_proj, txt_k_proj) + txt_k_sqsum = pl.row_sum(txt_k_sq) + txt_k_sqsum = pl.reshape(txt_k_sqsum, [1, T]) + txt_k_rms = pl.add(pl.mul(txt_k_sqsum, DIM_INV), DIT_EPS) + txt_k_inv = pl.recip(pl.sqrt(txt_k_rms)) + txt_k_inv = pl.reshape(txt_k_inv, [T, 1]) + txt_k_normed = pl.col_expand_mul( + pl.row_expand_mul(txt_k_proj, txt_k_inv), l_cross_krms[:, :]) + ca_k_txt_gm = pl.assemble(ca_k_txt_gm, txt_k_normed, [row_start, 0]) + + txt_v_proj = pl.matmul(txt_tile, l_cross_vtxt, b_trans=True, out_dtype=pl.FP32) + ca_v_txt_gm = pl.assemble(ca_v_txt_gm, txt_v_proj, [row_start, 0]) + + txt_attn_tmp = pl.create_tensor([S, T5_SEQ], dtype=pl.FP32) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ca_text_mha"): + for head_idx in pl.unroll(DIT_HEADS): + head_col = head_idx * DIT_HEAD_DIM + ca_txt_q_head = pl.slice(ca_q_gm, [S, DIT_HEAD_DIM], [0, head_col]) + txt_k_head = pl.slice(ca_k_txt_gm, [T5_SEQ, DIT_HEAD_DIM], [0, head_col]) + txt_v_head = pl.slice(ca_v_txt_gm, [T5_SEQ, DIT_HEAD_DIM], [0, head_col]) + + ca_txt_scores = pl.matmul(ca_txt_q_head, txt_k_head, b_trans=True, out_dtype=pl.FP32) + ca_txt_scores = pl.mul(ca_txt_scores, HEAD_SCALE) + + for row_start in pl.range(0, S, T): + txt_score_tile = pl.slice(ca_txt_scores, [T, T5_SEQ], [row_start, 0]) + txt_shifted = pl.row_expand_sub(txt_score_tile, pl.row_max(txt_score_tile)) + txt_exp_s = pl.exp(txt_shifted) + txt_attn_w = pl.row_expand_div(txt_exp_s, pl.row_sum(txt_exp_s)) + txt_attn_tmp = pl.assemble(txt_attn_tmp, txt_attn_w, [row_start, 0]) + + txt_attn_bf = pl.cast( + pl.slice(txt_attn_tmp, [S, T5_SEQ], [0, 0]), target_type=pl.BF16, + ) + txt_head_out = pl.matmul(txt_attn_bf, pl.cast(txt_v_head, target_type=pl.BF16), out_dtype=pl.FP32) + ca_txt_out = pl.assemble(ca_txt_out, txt_head_out, [0, head_col]) + + # ── Scope 14: Cross-Attention Image K/V projection + RMSNorm + MHA ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ca_img_kv"): + for row_start in pl.range(0, CLIP_PAD, T): + img_tile = pl.slice(clip_ctx, [T, DIT_DIM], [row_start, 0]) + + img_k_proj = pl.matmul(img_tile, l_cross_kimg, b_trans=True, out_dtype=pl.FP32) + img_k_sq = pl.mul(img_k_proj, img_k_proj) + img_k_sqsum = pl.row_sum(img_k_sq) + img_k_sqsum = pl.reshape(img_k_sqsum, [1, T]) + img_k_rms = pl.add(pl.mul(img_k_sqsum, DIM_INV), DIT_EPS) + img_k_inv = pl.recip(pl.sqrt(img_k_rms)) + img_k_inv = pl.reshape(img_k_inv, [T, 1]) + img_k_normed = pl.col_expand_mul( + pl.row_expand_mul(img_k_proj, img_k_inv), l_cross_kir[:, :]) + ca_k_img_gm = pl.assemble(ca_k_img_gm, img_k_normed, [row_start, 0]) + + img_v_proj = pl.matmul(img_tile, l_cross_vimg, b_trans=True, out_dtype=pl.FP32) + ca_v_img_gm = pl.assemble(ca_v_img_gm, img_v_proj, [row_start, 0]) + + img_attn_tmp = pl.create_tensor([S, CLIP_PAD], dtype=pl.FP32) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ca_img_mha"): + for head_idx in pl.unroll(DIT_HEADS): + head_col = head_idx * DIT_HEAD_DIM + ca_img_q_head = pl.slice(ca_q_gm, [S, DIT_HEAD_DIM], [0, head_col]) + img_k_head = pl.slice(ca_k_img_gm, [CLIP_PAD, DIT_HEAD_DIM], [0, head_col]) + img_v_head = pl.slice(ca_v_img_gm, [CLIP_PAD, DIT_HEAD_DIM], [0, head_col]) + + ca_img_scores = pl.matmul(ca_img_q_head, img_k_head, b_trans=True, out_dtype=pl.FP32) + ca_img_scores = pl.mul(ca_img_scores, HEAD_SCALE) + ca_img_scores = pl.col_expand_add(ca_img_scores, clip_mask) + + for row_start in pl.range(0, S, T): + img_score_tile = pl.slice(ca_img_scores, [T, CLIP_PAD], [row_start, 0]) + img_shifted = pl.row_expand_sub(img_score_tile, pl.row_max(img_score_tile)) + img_exp_s = pl.exp(img_shifted) + img_attn_w = pl.row_expand_div(img_exp_s, pl.row_sum(img_exp_s)) + img_attn_tmp = pl.assemble(img_attn_tmp, img_attn_w, [row_start, 0]) + + img_attn_bf = pl.cast( + pl.slice(img_attn_tmp, [S, CLIP_PAD], [0, 0]), target_type=pl.BF16, + ) + img_head_out = pl.matmul(img_attn_bf, pl.cast(img_v_head, target_type=pl.BF16), out_dtype=pl.FP32) + ca_img_out = pl.assemble(ca_img_out, img_head_out, [0, head_col]) + + # ── Scope 15: Cross-Attention Output Projection + Residual ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ca_oproj"): + combined_out = pl.add(ca_txt_out, ca_img_out) + combined_bf = pl.cast(combined_out, target_type=pl.BF16) + ca_oproj = pl.matmul(combined_bf, l_cross_o, b_trans=True, out_dtype=pl.FP32) + + for row_start in pl.range(0, S, T): + x_sa_tile = pl.slice(x_after_sa, [T, DIT_DIM], [row_start, 0]) + ca_oproj_tile = pl.slice(ca_oproj, [T, DIT_DIM], [row_start, 0]) + x_after_ca = pl.assemble( + x_after_ca, pl.add(x_sa_tile, ca_oproj_tile), [row_start, 0], + ) + + # ── Scope 16: FFN with AdaLN + GELU-tanh + Gated Residual ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ffn"): + shift_mlp_tile = pl.slice(mod_shift_mlp, [1, DIT_DIM], [0, 0]) + scale_mlp_tile = pl.slice(mod_scale_mlp, [1, DIT_DIM], [0, 0]) + gate_mlp_tile = pl.slice(mod_gate_mlp, [1, DIT_DIM], [0, 0]) + scale_mlp_plus1 = pl.add(scale_mlp_tile, 1.0) + + for row_start in pl.range(0, S, T): + ffn_x_tile = pl.slice(x_after_ca, [T, DIT_DIM], [row_start, 0]) + + # LayerNorm + ffn_scaled = pl.mul(ffn_x_tile, DIM_INV) + ffn_mean = pl.row_sum(ffn_scaled) + ffn_centered = pl.row_expand_sub(ffn_x_tile, ffn_mean) + ffn_var = pl.row_sum(pl.mul(ffn_centered, ffn_centered)) + ffn_var = pl.mul(ffn_var, DIM_INV) + ffn_var_eps = pl.add(ffn_var, LN_EPS) + ffn_var_eps = pl.reshape(ffn_var_eps, [1, T]) + ffn_std = pl.sqrt(ffn_var_eps) + ffn_std = pl.reshape(ffn_std, [T, 1]) + ffn_normed = pl.col_expand_add( + pl.col_expand_mul(pl.row_expand_div(ffn_centered, ffn_std), + l_norm2_w[:, :]), + l_norm2_b[:, :], + ) + + # AdaLN modulation + ffn_mod = pl.col_expand_add( + pl.col_expand_mul(ffn_normed, scale_mlp_plus1), + shift_mlp_tile, + ) + ffn_mod_bf = pl.cast(ffn_mod, target_type=pl.BF16) + + # FC1 + GELU-tanh + ffn_fc1 = pl.matmul(ffn_mod_bf, l_ffn1, b_trans=True, out_dtype=pl.FP32) + ffn_x3 = pl.mul(pl.mul(ffn_fc1, ffn_fc1), ffn_fc1) + ffn_inner = pl.add(ffn_fc1, pl.mul(ffn_x3, GELU_COEFF)) + ffn_z = pl.mul(ffn_inner, GELU_SQRT_2_OVER_PI) + ffn_2z = pl.mul(ffn_z, 2.0) + ffn_exp_val = pl.exp(ffn_2z) + ffn_recip = pl.recip(pl.add(ffn_exp_val, 1.0)) + ffn_tanh = pl.sub( + pl.add(pl.sub(ffn_fc1, ffn_fc1), 1.0), + pl.add(ffn_recip, ffn_recip), + ) + ffn_gelu = pl.mul(pl.mul(ffn_fc1, 0.5), pl.add(ffn_tanh, 1.0)) + ffn_gelu_bf = pl.cast(ffn_gelu, target_type=pl.BF16) + + # FC2 + gated residual + ffn_fc2 = pl.matmul(ffn_gelu_bf, l_ffn2, b_trans=True, out_dtype=pl.FP32) + ffn_gated = pl.col_expand_mul(ffn_fc2, gate_mlp_tile) + x_after_ffn = pl.assemble( + x_after_ffn, pl.add(ffn_x_tile, ffn_gated), [row_start, 0], + ) + + # ── Scope 17: Output Head — Modulation ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="head_mod"): + t_mod_fp = pl.cast(t_mod_gm, target_type=pl.FP32) + t_mod_head = pl.slice(t_mod_fp, [2, DIT_DIM], [0, 0]) + head_mod_fp = pl.reshape(pl.cast(head_mod_w, target_type=pl.FP32), [2, DIT_DIM]) + head_combined = pl.add(head_mod_fp, t_mod_head) + head_shift = pl.assemble( + head_shift, pl.slice(head_combined, [1, DIT_DIM], [0, 0]), [0, 0], + ) + head_scale = pl.assemble( + head_scale, pl.slice(head_combined, [1, DIT_DIM], [1, 0]), [0, 0], + ) + + # ── Scope 18: Output Head — LN + AdaLN + FC ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="head_proj"): + head_shift_tile = pl.slice(head_shift, [1, DIT_DIM], [0, 0]) + head_scale_tile = pl.slice(head_scale, [1, DIT_DIM], [0, 0]) + head_scale_plus1 = pl.add(head_scale_tile, 1.0) + + for row_start in pl.range(0, S, T): + head_x_tile = pl.slice(x_after_ffn, [T, DIT_DIM], [row_start, 0]) + + head_scaled = pl.mul(head_x_tile, DIM_INV) + head_mean = pl.row_sum(head_scaled) + head_centered = pl.row_expand_sub(head_x_tile, head_mean) + head_var = pl.row_sum(pl.mul(head_centered, head_centered)) + head_var = pl.mul(head_var, DIM_INV) + head_var_eps = pl.add(head_var, LN_EPS) + head_var_eps = pl.reshape(head_var_eps, [1, T]) + head_std = pl.sqrt(head_var_eps) + head_std = pl.reshape(head_std, [T, 1]) + head_normed = pl.col_expand_add( + pl.col_expand_mul(pl.row_expand_div(head_centered, head_std), + head_ln_w[:, :]), + head_ln_b[:, :], + ) + + head_mod = pl.col_expand_add( + pl.col_expand_mul(head_normed, head_scale_plus1), + head_shift_tile, + ) + head_mod_bf = pl.cast(head_mod, target_type=pl.BF16) + head_fc_out = pl.matmul(head_mod_bf, head_fc, b_trans=True, out_dtype=pl.FP32) + head_out_gm = pl.assemble( + head_out_gm, pl.cast(head_fc_out, target_type=pl.BF16), + [row_start, 0], + ) + + # ── Scope 19: Remove reference tokens → output ── + with pl.at(level=pl.Level.CORE_GROUP, name_hint="remove_ref"): + for row_start in pl.range(0, X_N, T): + out_tile = pl.slice( + head_out_gm, [T, DIT_IN_DIM * 4], [REF_N + row_start, 0], + ) + out = pl.assemble(out, out_tile, [row_start, 0]) + + return out + + +# ══════════════════════════════════════════════════════════════════════════════ +# Test harness +# ══════════════════════════════════════════════════════════════════════════════ + +def build_tensor_specs(): + """Build tensor specs for the fused dit_forward kernel.""" + from golden import TensorSpec + + torch.manual_seed(42) + w = make_weights(seed=42) + + x_input = torch.randn(1, DIT_IN_DIM + COND_CH, FP, LAT_H, LAT_W) + text_raw = torch.randn(1, T5_SEQ, T5_DIM) + clip_raw = torch.randn(1, CLIP_TOKENS, CLIP_DIM) + ref_latents = torch.randn(1, DIT_IN_DIM, LAT_H, LAT_W) + timestep = torch.tensor([500.], dtype=torch.float32) + + # Timestep embedding — frequencies are static; the kernel computes + # sin_emb = [cos(t*freqs), sin(t*freqs)] on NPU using pl.cos/pl.sin. + half_freq = DIT_FREQ_DIM // 2 + time_freqs_1d = torch.exp(-math.log(10000.0) * torch.arange(half_freq).float() / half_freq) + time_freqs = time_freqs_1d.unsqueeze(0).expand(T, half_freq).contiguous() + + # Host im2col (data rearrangement) + ref_conv_w_2d = w["dit_ref_conv"].reshape(DIT_DIM, -1).T.contiguous().bfloat16() + ref_col = _im2col_ref_conv2d(ref_latents.bfloat16()) + patch_conv_w_2d = w["dit_patch"].reshape(DIT_DIM, -1).T.contiguous().bfloat16() + patch_col = _im2col_patch_conv3d(x_input.bfloat16()) + + # RoPE tables (data rearrangement) + cos_f, sin_f, cos_h, sin_h, cos_w, sin_w = _precompute_rope_cos_sin() + rope_cos, rope_sin = _build_rope_map( + cos_f, sin_f, cos_h, sin_h, cos_w, sin_w, FP + 1, HP, WP, + ) + + # CLIP padding + mask + clip_padded = torch.zeros(CLIP_PAD, CLIP_DIM, dtype=torch.bfloat16) + clip_padded[:CLIP_TOKENS, :] = clip_raw.squeeze(0).bfloat16() + clip_mask = torch.zeros(1, CLIP_PAD, dtype=torch.float32) + clip_mask[:, CLIP_TOKENS:] = -65504.0 + + def _identity(t): + return lambda: t + + def spec(name, shape, dtype, value): + return TensorSpec(name, shape, dtype, init_value=_identity(value)) + + specs = [ + # Conv inputs + spec("ref_col", [REF_N, REF_CONV_COL], torch.bfloat16, ref_col), + spec("patch_col", [X_N, PATCH_CONV_COL], torch.bfloat16, patch_col), + spec("ref_conv_w", [REF_CONV_COL, DIT_DIM], torch.bfloat16, ref_conv_w_2d), + spec("patch_conv_w", [PATCH_CONV_COL, DIT_DIM], torch.bfloat16, patch_conv_w_2d), + # Text + CLIP projection + spec("text_raw", [T5_SEQ, T5_DIM], torch.bfloat16, text_raw.squeeze(0).bfloat16()), + spec("text_fc1_w", [DIT_DIM, DIT_TEXT_DIM], torch.bfloat16, w["dit_text_fc1"].bfloat16()), + spec("text_fc2_w", [DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_text_fc2"].bfloat16()), + spec("clip_raw", [CLIP_PAD, CLIP_DIM], torch.bfloat16, clip_padded), + spec("clip_ln1_w", [1, CLIP_DIM], torch.float32, w["dit_clip_ln_w"].unsqueeze(0)), + spec("clip_ln1_b", [1, CLIP_DIM], torch.float32, w["dit_clip_ln_b"].unsqueeze(0)), + spec("clip_fc1_w", [CLIP_DIM, CLIP_DIM], torch.bfloat16, w["dit_clip_fc1"].bfloat16()), + spec("clip_fc2_w", [DIT_DIM, CLIP_DIM], torch.bfloat16, w["dit_clip_fc2"].bfloat16()), + spec("clip_ln2_w", [1, DIT_DIM], torch.float32, w["dit_clip_ln2_w"].unsqueeze(0)), + spec("clip_ln2_b", [1, DIT_DIM], torch.float32, w["dit_clip_ln2_b"].unsqueeze(0)), + # Context tables + spec("clip_mask", [1, CLIP_PAD], torch.float32, clip_mask), + spec("rope_cos", [S, DIT_DIM // 2], torch.float32, rope_cos), + spec("rope_sin", [S, DIT_DIM // 2], torch.float32, rope_sin), + # Timestep embedding (computed on NPU) + spec("timestep", [1], torch.float32, timestep), + spec("time_freqs", [T, half_freq], torch.float32, time_freqs), + spec("time_fc1_w", [DIT_DIM, DIT_FREQ_DIM], torch.bfloat16, w["dit_time_fc1"].bfloat16()), + spec("time_fc2_w", [DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_time_fc2"].bfloat16()), + spec("time_proj_w", [6 * DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_time_proj"].bfloat16()), + # Layer weights + spec("l_mod", [1, 6, DIT_DIM], torch.bfloat16, w["dit_l0_mod"].bfloat16()), + spec("l_norm1_w", [1, DIT_DIM], torch.float32, w["dit_l0_norm1_w"].unsqueeze(0)), + spec("l_norm1_b", [1, DIT_DIM], torch.float32, w["dit_l0_norm1_b"].unsqueeze(0)), + spec("l_q_w", [DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_l0_q"].bfloat16()), + spec("l_k_w", [DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_l0_k"].bfloat16()), + spec("l_v_w", [DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_l0_v"].bfloat16()), + spec("l_o_w", [DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_l0_o"].bfloat16()), + spec("l_q_rms", [1, DIT_DIM], torch.float32, w["dit_l0_q_rms"].unsqueeze(0)), + spec("l_k_rms", [1, DIT_DIM], torch.float32, w["dit_l0_k_rms"].unsqueeze(0)), + spec("l_norm3_w", [1, DIT_DIM], torch.float32, w["dit_l0_norm3_w"].unsqueeze(0)), + spec("l_norm3_b", [1, DIT_DIM], torch.float32, w["dit_l0_norm3_b"].unsqueeze(0)), + spec("l_cross_q", [DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_l0_cross_q"].bfloat16()), + spec("l_cross_ktxt", [DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_l0_cross_k_txt"].bfloat16()), + spec("l_cross_vtxt", [DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_l0_cross_v_txt"].bfloat16()), + spec("l_cross_kimg", [DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_l0_cross_k_img"].bfloat16()), + spec("l_cross_vimg", [DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_l0_cross_v_img"].bfloat16()), + spec("l_cross_o", [DIT_DIM, DIT_DIM], torch.bfloat16, w["dit_l0_cross_o"].bfloat16()), + spec("l_cross_qrms", [1, DIT_DIM], torch.float32, w["dit_l0_cross_q_rms"].unsqueeze(0)), + spec("l_cross_krms", [1, DIT_DIM], torch.float32, w["dit_l0_cross_k_rms"].unsqueeze(0)), + spec("l_cross_kir", [1, DIT_DIM], torch.float32, w["dit_l0_cross_k_img_rms"].unsqueeze(0)), + spec("l_norm2_w", [1, DIT_DIM], torch.float32, w["dit_l0_norm2_w"].unsqueeze(0)), + spec("l_norm2_b", [1, DIT_DIM], torch.float32, w["dit_l0_norm2_b"].unsqueeze(0)), + spec("l_ffn1", [DIT_FFN, DIT_DIM], torch.bfloat16, w["dit_l0_ffn1"].bfloat16()), + spec("l_ffn2", [DIT_DIM, DIT_FFN], torch.bfloat16, w["dit_l0_ffn2"].bfloat16()), + # Head weights + spec("head_mod_w", [1, 2, DIT_DIM], torch.bfloat16, w["dit_head_mod"].bfloat16()), + spec("head_ln_w", [1, DIT_DIM], torch.float32, w["dit_head_ln_w"].unsqueeze(0)), + spec("head_ln_b", [1, DIT_DIM], torch.float32, w["dit_head_ln_b"].unsqueeze(0)), + spec("head_fc", [DIT_IN_DIM * 4, DIT_DIM], torch.bfloat16, w["dit_head_fc"].bfloat16()), + ] + specs.append(TensorSpec("out", [X_N, DIT_IN_DIM * 4], torch.bfloat16, is_output=True)) + + _GOLDEN_DATA["w"] = w + _GOLDEN_DATA["x_input"] = x_input + _GOLDEN_DATA["text_raw"] = text_raw + _GOLDEN_DATA["clip_raw"] = clip_raw + _GOLDEN_DATA["ref_latents"] = ref_latents + _GOLDEN_DATA["timestep"] = timestep + return specs + + +_GOLDEN_DATA = {} + + +def golden_dit_forward_fn(tensors): + """Golden reference — calls test_golden_fun_control_full.py::dit_forward. + + The golden uses exact F.gelu() for CLIP projection (matching the reference + at line 1417). The kernel uses GELU-tanh approximation because pypto lacks + erf — the permissive K3 tolerance accounts for this difference. + """ + w = _GOLDEN_DATA["w"] + x_input = _GOLDEN_DATA["x_input"] + text_raw = _GOLDEN_DATA["text_raw"] + clip_raw = _GOLDEN_DATA["clip_raw"] + ref_latents = _GOLDEN_DATA["ref_latents"] + timestep = _GOLDEN_DATA["timestep"] + context = {"text": text_raw, "clip": clip_raw, "ref_latents": ref_latents} + freqs = precompute_freqs_cis_3d(DIT_HEAD_DIM) + result = golden_dit_fn(x_input, context, freqs, timestep, w) + tensors["out"][:] = result.permute(0, 2, 3, 4, 1).reshape(-1, DIT_IN_DIM * 4).bfloat16() + + +if __name__ == "__main__": + from golden import ratio_allclose, run_jit + + parser = argparse.ArgumentParser(description="DiT forward pass — fused kernel test") + parser.add_argument("-p", "--platform", default="a2a3", + choices=["a2a3", "a2a3sim", "a5", "a5sim"]) + parser.add_argument("-d", "--device", type=int, default=0) + parser.add_argument("--enable-l2-swimlane", action="store_true", default=False) + args = parser.parse_args() + + specs = build_tensor_specs() + + result = run_jit( + fn=dit_forward, + specs=specs, + golden_fn=golden_dit_forward_fn, + runtime_cfg=dict( + platform=args.platform, + device_id=args.device, + enable_l2_swimlane=args.enable_l2_swimlane, + ), + rtol=3e-3, + atol=3e-3, + compare_fn={"out": ratio_allclose(atol=0.6, rtol=0.6, max_error_ratio=0.02)}, + ) + if not result.passed: + if result.error: + print(result.error) + raise SystemExit(1) + print("PASS") diff --git a/models/world_model/t5_encoder.py b/models/world_model/t5_encoder.py new file mode 100644 index 00000000..06d98123 --- /dev/null +++ b/models/world_model/t5_encoder.py @@ -0,0 +1,481 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""T5 Encoder — pypto3.0 kernel for the T5 text encoder sub-network. + +Implements the full T5 encoder stack: RMSNorm → Self-Attention (with relative +position bias) → Residual → RMSNorm → Gated GELU-tanh FFN → Residual, repeated +for T5_LAYERS layers, followed by a final RMSNorm. + +The golden reference uses ``T5LayerNorm``, ``T5RelativeEmbedding``, and +``T5SelfAttention`` from ``test_golden_fun_control_full.py``, which mirrors +the T5 encoder in ``infer_fun_control_1_3b_text.py`` at reduced dimensions. + +Usage:: + + python t5_encoder.py # golden-case precision test (T5_DIM=128) +""" + +import argparse +import sys + +import pypto.language as pl +import torch + +from config import ( + T5_DIM, T5_FFN, T5_HEADS, T5_HEAD_DIM, T5_LAYERS, T5_NUM_BUCKETS, T5_SEQ, +) + +# ── Golden reference (shared classes with infer_fun_control_1_3b_text.py) ── +sys.path.insert(0, '/data/x00952168/pypto3.0/cann-recipes-embodied-ai/world_model/agibot-arm-world-model/infer_with_torch') +from test_golden_fun_control_full import ( # noqa: E402 + T5LayerNorm, T5RelativeEmbedding, T5SelfAttention, +) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Optimization config — tiling parameters per functional stage. +# ══════════════════════════════════════════════════════════════════════════════ + +EPS = 1e-6 # RMSNorm epsilon + +# ── Tiling geometry ── +# SEQ_TILE = row tile for the seq-length dimension (pl.parallel step). +# K_CHUNK = column tile for the hidden-dimension loops (K dimension). +# N_CHUNK = column tile for the output-dimension matmul loops (N dimension). +SEQ_TILE = 16 +K_CHUNK = 512 +N_CHUNK = 64 + +# Dynamic chunking — scale down for the golden case (T5_DIM=128) so loops +# don't exceed the dimension bounds. +if T5_DIM >= 4096: # Real case (t5_umt5-xxl) + pass # defaults above +else: # Golden case (T5_DIM=128) + K_CHUNK = 128 + +# ── Derived block counts ── +HIDDEN_BLOCKS = T5_DIM // K_CHUNK # dim splits for RMSNorm / QKV / out_proj +FFN_BLOCKS = T5_FFN // K_CHUNK # dim splits for the fc2 (down) matmul + +# Geometry assertions — keep at the bottom of the config block. +assert T5_DIM % T5_HEADS == 0, "T5_HEADS must divide T5_DIM" +assert T5_HEAD_DIM * T5_HEADS == T5_DIM +assert T5_DIM % K_CHUNK == 0, "K_CHUNK must divide T5_DIM" +assert T5_DIM % N_CHUNK == 0, "N_CHUNK must divide T5_DIM" +assert T5_FFN % K_CHUNK == 0, "K_CHUNK must divide T5_FFN" +assert T5_FFN % N_CHUNK == 0, "N_CHUNK must divide T5_FFN" +assert T5_SEQ % SEQ_TILE == 0, "SEQ_TILE must divide T5_SEQ" + + +@pl.jit.inline +def _rmsnorm( + x_in: pl.Tensor[[T5_SEQ, T5_DIM], pl.FP32], + weight: pl.Tensor[[1, T5_DIM], pl.FP32], + out: pl.Tensor[[T5_SEQ, T5_DIM], pl.BF16], +) -> pl.Tensor[[T5_SEQ, T5_DIM], pl.BF16]: + """RMSNorm: x * rsqrt(mean(x^2) + eps) * weight. FP32 in → BF16 out.""" + for st in pl.parallel(0, T5_SEQ, SEQ_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="rmsnorm"): + partial_sq = pl.full([1, SEQ_TILE], dtype=pl.FP32, value=0.0) + for kb in pl.range(HIDDEN_BLOCKS): + k0 = kb * K_CHUNK + xc = pl.slice(x_in, [SEQ_TILE, K_CHUNK], [st, k0]) + partial_sq = pl.add(partial_sq, pl.reshape(pl.row_sum(pl.mul(xc, xc)), [1, SEQ_TILE])) + var = pl.reshape(pl.add(pl.mul(partial_sq, 1.0 / T5_DIM), EPS), [SEQ_TILE, 1]) + inv = pl.recip(pl.sqrt(var)) + for kb in pl.range(HIDDEN_BLOCKS): + k0 = kb * K_CHUNK + xc = pl.slice(x_in, [SEQ_TILE, K_CHUNK], [st, k0]) + g = pl.slice(weight, [1, K_CHUNK], [0, k0]) + n = pl.col_expand_mul(pl.row_expand_mul(xc, inv), g) + out = pl.assemble(out, pl.cast(n, target_type=pl.BF16), [st, k0]) + return out + + +@pl.jit.inline +def t5_encoder_layer( + x_in: pl.Tensor[[T5_SEQ, T5_DIM], pl.FP32], + norm1_w: pl.Tensor[[T5_LAYERS, 1, T5_DIM], pl.FP32], + q_w: pl.Tensor[[T5_LAYERS, T5_DIM, T5_DIM], pl.BF16], + k_w: pl.Tensor[[T5_LAYERS, T5_DIM, T5_DIM], pl.BF16], + v_w: pl.Tensor[[T5_LAYERS, T5_DIM, T5_DIM], pl.BF16], + o_w: pl.Tensor[[T5_LAYERS, T5_DIM, T5_DIM], pl.BF16], + pos_bias: pl.Tensor[[T5_LAYERS, T5_HEADS, T5_SEQ, T5_SEQ], pl.FP32], + norm2_w: pl.Tensor[[T5_LAYERS, 1, T5_DIM], pl.FP32], + wi_0: pl.Tensor[[T5_LAYERS, T5_FFN, T5_DIM], pl.BF16], + wi_1: pl.Tensor[[T5_LAYERS, T5_FFN, T5_DIM], pl.BF16], + wo: pl.Tensor[[T5_LAYERS, T5_DIM, T5_FFN], pl.BF16], + out: pl.Tensor[[T5_SEQ, T5_DIM], pl.FP32], + layer_idx: pl.Scalar[pl.INT32], +) -> pl.Tensor[[T5_SEQ, T5_DIM], pl.FP32]: + # ── Per-layer weight slicing ── + norm1_w_layer = pl.slice(norm1_w, [1, 1, T5_DIM], [layer_idx, 0, 0]) + norm1_w_layer = pl.reshape(norm1_w_layer, [1, T5_DIM]) + q_w_layer = pl.slice(q_w, [1, T5_DIM, T5_DIM], [layer_idx, 0, 0]) + q_w_layer = pl.reshape(q_w_layer, [T5_DIM, T5_DIM]) + k_w_layer = pl.slice(k_w, [1, T5_DIM, T5_DIM], [layer_idx, 0, 0]) + k_w_layer = pl.reshape(k_w_layer, [T5_DIM, T5_DIM]) + v_w_layer = pl.slice(v_w, [1, T5_DIM, T5_DIM], [layer_idx, 0, 0]) + v_w_layer = pl.reshape(v_w_layer, [T5_DIM, T5_DIM]) + o_w_layer = pl.slice(o_w, [1, T5_DIM, T5_DIM], [layer_idx, 0, 0]) + o_w_layer = pl.reshape(o_w_layer, [T5_DIM, T5_DIM]) + pos_bias_layer = pl.slice(pos_bias, [1, T5_HEADS, T5_SEQ, T5_SEQ], [layer_idx, 0, 0, 0]) + pos_bias_layer = pl.reshape(pos_bias_layer, [T5_HEADS, T5_SEQ, T5_SEQ]) + norm2_w_layer = pl.slice(norm2_w, [1, 1, T5_DIM], [layer_idx, 0, 0]) + norm2_w_layer = pl.reshape(norm2_w_layer, [1, T5_DIM]) + wi_0_layer = pl.slice(wi_0, [1, T5_FFN, T5_DIM], [layer_idx, 0, 0]) + wi_0_layer = pl.reshape(wi_0_layer, [T5_FFN, T5_DIM]) + wi_1_layer = pl.slice(wi_1, [1, T5_FFN, T5_DIM], [layer_idx, 0, 0]) + wi_1_layer = pl.reshape(wi_1_layer, [T5_FFN, T5_DIM]) + wo_layer = pl.slice(wo, [1, T5_DIM, T5_FFN], [layer_idx, 0, 0]) + wo_layer = pl.reshape(wo_layer, [T5_DIM, T5_FFN]) + + # ── Scope 1 · RMSNorm (pre-attention) ── + normed_gm = pl.create_tensor([T5_SEQ, T5_DIM], dtype=pl.BF16) + normed_gm = _rmsnorm(x_in, norm1_w_layer, normed_gm) + + # ── Scope 2a · Q projection ── + q_gm = pl.create_tensor([T5_SEQ, T5_DIM], dtype=pl.BF16) + for st in pl.parallel(0, T5_SEQ, SEQ_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="q_proj"): + for ob in pl.range(T5_DIM // N_CHUNK): + n0 = ob * N_CHUNK + act_tile = pl.slice(normed_gm, [SEQ_TILE, K_CHUNK], [st, 0]) + w_chunk = pl.slice(q_w_layer, [N_CHUNK, K_CHUNK], [n0, 0]) + q_acc = pl.matmul(act_tile, w_chunk, b_trans=True, out_dtype=pl.FP32) + for kb in pl.range(1, HIDDEN_BLOCKS): + k0 = kb * K_CHUNK + act_chunk = pl.slice(normed_gm, [SEQ_TILE, K_CHUNK], [st, k0]) + w_chunk_i = pl.slice(q_w_layer, [N_CHUNK, K_CHUNK], [n0, k0]) + q_acc = pl.matmul_acc(q_acc, act_chunk, w_chunk_i, b_trans=True) + q_gm = pl.assemble(q_gm, pl.cast(q_acc, target_type=pl.BF16), [st, n0]) + + # ── Scope 2b · K projection ── + k_gm = pl.create_tensor([T5_SEQ, T5_DIM], dtype=pl.BF16) + for st in pl.parallel(0, T5_SEQ, SEQ_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="k_proj"): + for ob in pl.range(T5_DIM // N_CHUNK): + n0 = ob * N_CHUNK + act_tile = pl.slice(normed_gm, [SEQ_TILE, K_CHUNK], [st, 0]) + w_chunk = pl.slice(k_w_layer, [N_CHUNK, K_CHUNK], [n0, 0]) + k_acc = pl.matmul(act_tile, w_chunk, b_trans=True, out_dtype=pl.FP32) + for kb in pl.range(1, HIDDEN_BLOCKS): + k0 = kb * K_CHUNK + act_chunk = pl.slice(normed_gm, [SEQ_TILE, K_CHUNK], [st, k0]) + w_chunk_i = pl.slice(k_w_layer, [N_CHUNK, K_CHUNK], [n0, k0]) + k_acc = pl.matmul_acc(k_acc, act_chunk, w_chunk_i, b_trans=True) + k_gm = pl.assemble(k_gm, pl.cast(k_acc, target_type=pl.BF16), [st, n0]) + + # ── Scope 2c · V projection ── + v_gm = pl.create_tensor([T5_SEQ, T5_DIM], dtype=pl.BF16) + for st in pl.parallel(0, T5_SEQ, SEQ_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="v_proj"): + for ob in pl.range(T5_DIM // N_CHUNK): + n0 = ob * N_CHUNK + act_tile = pl.slice(normed_gm, [SEQ_TILE, K_CHUNK], [st, 0]) + w_chunk = pl.slice(v_w_layer, [N_CHUNK, K_CHUNK], [n0, 0]) + v_acc = pl.matmul(act_tile, w_chunk, b_trans=True, out_dtype=pl.FP32) + for kb in pl.range(1, HIDDEN_BLOCKS): + k0 = kb * K_CHUNK + act_chunk = pl.slice(normed_gm, [SEQ_TILE, K_CHUNK], [st, k0]) + w_chunk_i = pl.slice(v_w_layer, [N_CHUNK, K_CHUNK], [n0, k0]) + v_acc = pl.matmul_acc(v_acc, act_chunk, w_chunk_i, b_trans=True) + v_gm = pl.assemble(v_gm, pl.cast(v_acc, target_type=pl.BF16), [st, n0]) + + # ── Scope 3 · Multi-head self-attention (with relative position bias) ── + ctx_gm = pl.create_tensor([T5_SEQ, T5_DIM], dtype=pl.BF16) + with pl.at(level=pl.Level.CORE_GROUP, name_hint="mha"): + for h in pl.range(T5_HEADS): + head_col = h * T5_HEAD_DIM + q_head = pl.slice(q_gm, [T5_SEQ, T5_HEAD_DIM], [0, head_col]) + k_head = pl.slice(k_gm, [T5_SEQ, T5_HEAD_DIM], [0, head_col]) + v_head = pl.slice(v_gm, [T5_SEQ, T5_HEAD_DIM], [0, head_col]) + raw_scores = pl.matmul(q_head, k_head, b_trans=True, out_dtype=pl.FP32) + + # Add relative position bias for this head + # pos_bias_h: [T5_SEQ, T5_SEQ] - precomputed on host + pos_bias_h = pl.slice(pos_bias_layer, [1, T5_SEQ, T5_SEQ], [h, 0, 0]) + pos_bias_h = pl.reshape(pos_bias_h, [T5_SEQ, T5_SEQ]) + # Add to scores (both are FP32) + raw_scores = pl.add(raw_scores, pos_bias_h) + + row_max_val = pl.row_max(raw_scores) + shifted = pl.row_expand_sub(raw_scores, row_max_val) + exp_scores = pl.exp(shifted) + exp_sum = pl.row_sum(exp_scores) + attn_weights = pl.row_expand_div(exp_scores, exp_sum) + attn_weights_bf = pl.cast(attn_weights, target_type=pl.BF16) + head_out = pl.matmul(attn_weights_bf, v_head, out_dtype=pl.FP32) + ctx_gm = pl.assemble(ctx_gm, pl.cast(head_out, target_type=pl.BF16), [0, head_col]) + + # ── Scope 4 · Output projection ── + attn_out_gm = pl.create_tensor([T5_SEQ, T5_DIM], dtype=pl.BF16) + for st in pl.parallel(0, T5_SEQ, SEQ_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="out_proj"): + for ob in pl.range(T5_DIM // N_CHUNK): + n0 = ob * N_CHUNK + ctx_tile = pl.slice(ctx_gm, [SEQ_TILE, K_CHUNK], [st, 0]) + o_w_chunk = pl.slice(o_w_layer, [N_CHUNK, K_CHUNK], [n0, 0]) + oproj_acc = pl.matmul(ctx_tile, o_w_chunk, b_trans=True, out_dtype=pl.FP32) + for kb in pl.range(1, HIDDEN_BLOCKS): + k0 = kb * K_CHUNK + ctx_chunk = pl.slice(ctx_gm, [SEQ_TILE, K_CHUNK], [st, k0]) + o_w_chunk_i = pl.slice(o_w_layer, [N_CHUNK, K_CHUNK], [n0, k0]) + oproj_acc = pl.matmul_acc(oproj_acc, ctx_chunk, o_w_chunk_i, b_trans=True) + attn_out_gm = pl.assemble(attn_out_gm, pl.cast(oproj_acc, target_type=pl.BF16), [st, n0]) + + # ── Scope 5 · Residual add (attention) ── + x_after_attn_gm = pl.create_tensor([T5_SEQ, T5_DIM], dtype=pl.FP32) + for st in pl.parallel(0, T5_SEQ, SEQ_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="residual"): + for ob in pl.range(T5_DIM // N_CHUNK): + n0 = ob * N_CHUNK + orig_tile = pl.slice(x_in, [SEQ_TILE, N_CHUNK], [st, n0]) + attn_out_fp32 = pl.cast(pl.slice(attn_out_gm, [SEQ_TILE, N_CHUNK], [st, n0]), target_type=pl.FP32) + attention_residual = pl.add(orig_tile, attn_out_fp32) + x_after_attn_gm = pl.assemble(x_after_attn_gm, attention_residual, [st, n0]) + + # ── Scope 6 · RMSNorm (pre-FFN) ── + normed2_gm = pl.create_tensor([T5_SEQ, T5_DIM], dtype=pl.BF16) + normed2_gm = _rmsnorm(x_after_attn_gm, norm2_w_layer, normed2_gm) + + # ── Scope 7 · FFN gate + fc1 (GELU-tanh activation) ── + fc1_gm = pl.create_tensor([T5_SEQ, T5_FFN], dtype=pl.BF16) + for st in pl.parallel(0, T5_SEQ, SEQ_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ffn_gate_fc1"): + for ob in pl.range(T5_FFN // N_CHUNK): + n0 = ob * N_CHUNK + ffn_act_tile = pl.slice(normed2_gm, [SEQ_TILE, K_CHUNK], [st, 0]) + gate_w_chunk = pl.slice(wi_0_layer, [N_CHUNK, K_CHUNK], [n0, 0]) + gate_acc = pl.matmul(ffn_act_tile, gate_w_chunk, b_trans=True, out_dtype=pl.FP32) + fc1_w_chunk = pl.slice(wi_1_layer, [N_CHUNK, K_CHUNK], [n0, 0]) + fc1_acc = pl.matmul(ffn_act_tile, fc1_w_chunk, b_trans=True, out_dtype=pl.FP32) + for kb in pl.range(1, HIDDEN_BLOCKS): + k0 = kb * K_CHUNK + act_chunk = pl.slice(normed2_gm, [SEQ_TILE, K_CHUNK], [st, k0]) + gate_w_chunk_i = pl.slice(wi_0_layer, [N_CHUNK, K_CHUNK], [n0, k0]) + gate_acc = pl.matmul_acc(gate_acc, act_chunk, gate_w_chunk_i, b_trans=True) + fc1_w_chunk_i = pl.slice(wi_1_layer, [N_CHUNK, K_CHUNK], [n0, k0]) + fc1_acc = pl.matmul_acc(fc1_acc, act_chunk, fc1_w_chunk_i, b_trans=True) + # GELU-tanh: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) + x_cubed = pl.mul(pl.mul(gate_acc, gate_acc), gate_acc) + gelu_inner = pl.add(gate_acc, pl.mul(x_cubed, 0.044715)) + gelu_scaled = pl.mul(gelu_inner, 0.7978845608) + gelu_2z = pl.mul(gelu_scaled, 2.0) + gelu_exp = pl.exp(gelu_2z) + tanh_val = pl.div(pl.sub(gelu_exp, 1.0), pl.add(gelu_exp, 1.0)) + gelu_out = pl.mul(pl.mul(gate_acc, 0.5), pl.add(tanh_val, 1.0)) + gated_fc1 = pl.mul(fc1_acc, gelu_out) + fc1_gm = pl.assemble(fc1_gm, pl.cast(gated_fc1, target_type=pl.BF16), [st, n0]) + + # ── Scope 8 · FFN fc2 + residual ── + x_final_gm = pl.create_tensor([T5_SEQ, T5_DIM], dtype=pl.FP32) + for st in pl.parallel(0, T5_SEQ, SEQ_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="ffn_fc2"): + for ob in pl.range(T5_DIM // N_CHUNK): + n0 = ob * N_CHUNK + fc1_out_tile = pl.slice(fc1_gm, [SEQ_TILE, K_CHUNK], [st, 0]) + fc2_w_chunk = pl.slice(wo_layer, [N_CHUNK, K_CHUNK], [n0, 0]) + fc2_acc = pl.matmul(fc1_out_tile, fc2_w_chunk, b_trans=True, out_dtype=pl.FP32) + for kb in pl.range(1, FFN_BLOCKS): + k0 = kb * K_CHUNK + fc1_out_chunk = pl.slice(fc1_gm, [SEQ_TILE, K_CHUNK], [st, k0]) + fc2_w_chunk_i = pl.slice(wo_layer, [N_CHUNK, K_CHUNK], [n0, k0]) + fc2_acc = pl.matmul_acc(fc2_acc, fc1_out_chunk, fc2_w_chunk_i, b_trans=True) + attn_residual = pl.slice(x_after_attn_gm, [SEQ_TILE, N_CHUNK], [st, n0]) + ffn_residual = pl.add(attn_residual, fc2_acc) + x_final_gm = pl.assemble(x_final_gm, ffn_residual, [st, n0]) + + out = pl.assemble(out, x_final_gm, [0, 0]) + return out + + +@pl.jit +def t5_encoder( + x_in: pl.Tensor[[T5_SEQ, T5_DIM], pl.FP32], + norm1_w: pl.Tensor[[T5_LAYERS, 1, T5_DIM], pl.FP32], + q_w: pl.Tensor[[T5_LAYERS, T5_DIM, T5_DIM], pl.BF16], + k_w: pl.Tensor[[T5_LAYERS, T5_DIM, T5_DIM], pl.BF16], + v_w: pl.Tensor[[T5_LAYERS, T5_DIM, T5_DIM], pl.BF16], + o_w: pl.Tensor[[T5_LAYERS, T5_DIM, T5_DIM], pl.BF16], + pos_bias: pl.Tensor[[T5_LAYERS, T5_HEADS, T5_SEQ, T5_SEQ], pl.FP32], + norm2_w: pl.Tensor[[T5_LAYERS, 1, T5_DIM], pl.FP32], + wi_0: pl.Tensor[[T5_LAYERS, T5_FFN, T5_DIM], pl.BF16], + wi_1: pl.Tensor[[T5_LAYERS, T5_FFN, T5_DIM], pl.BF16], + wo: pl.Tensor[[T5_LAYERS, T5_DIM, T5_FFN], pl.BF16], + final_norm_w: pl.Tensor[[1, T5_DIM], pl.FP32], + out: pl.Out[pl.Tensor[[T5_SEQ, T5_DIM], pl.BF16]], +): + # ── Copy input → internal tensor ── + cur = pl.create_tensor([T5_SEQ, T5_DIM], dtype=pl.FP32) + for st in pl.parallel(0, T5_SEQ, SEQ_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="copy_input"): + for kb in pl.range(HIDDEN_BLOCKS): + k0 = kb * K_CHUNK + xc = pl.slice(x_in, [SEQ_TILE, K_CHUNK], [st, k0]) + cur = pl.assemble(cur, xc, [st, k0]) + + # ── Layer loop ── + for layer_idx in pl.range(T5_LAYERS): + next_hidden = pl.create_tensor([T5_SEQ, T5_DIM], dtype=pl.FP32) + cur = t5_encoder_layer( + cur, norm1_w, q_w, k_w, v_w, o_w, pos_bias, norm2_w, + wi_0, wi_1, wo, next_hidden, layer_idx + ) + + # ── Final RMSNorm ── + x_final_normed_gm = pl.create_tensor([T5_SEQ, T5_DIM], dtype=pl.BF16) + x_final_normed_gm = _rmsnorm(cur, final_norm_w, x_final_normed_gm) + + # ── Output copy (BF16) ── + for st in pl.parallel(0, T5_SEQ, SEQ_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="final_output"): + for ob in pl.range(T5_DIM // N_CHUNK): + n0 = ob * N_CHUNK + xf_chunk = pl.slice(x_final_normed_gm, [SEQ_TILE, N_CHUNK], [st, n0]) + out = pl.assemble(out, xf_chunk, [st, n0]) + + return out + + +def _golden_t5_encoder(x, norm1_w, q_w, k_w, v_w, o_w, rel_pos_bias, norm2_w, wi_0, wi_1, wo, final_norm_w): + """Golden reference using the EXACT same classes as test_golden_fun_control_full.py.""" + layer_weights = { + 'norm1_w': norm1_w, + 'norm2_w': norm2_w, + 'attn': {'q': q_w, 'k': k_w, 'v': v_w, 'o': o_w}, + 'ffn': {'wi_0': wi_0, 'wi_1': wi_1, 'wo': wo}, + 'pos_embedding': rel_pos_bias, + } + block = T5SelfAttention( + layer_weights, T5_DIM, T5_DIM, T5_FFN, T5_HEADS, T5_NUM_BUCKETS, + shared_pos=False, dropout=0.1, + ) + x = block(x, mask=None, pos_bias=None) + x = T5LayerNorm(final_norm_w, T5_DIM)(x) + return x + + +# ══════════════════════════════════════════════════════════════════════════════ +# Test harness — build_tensor_specs / golden_fn / __main__. +# ══════════════════════════════════════════════════════════════════════════════ + +_GOLDEN_DATA = None # populated by build_tensor_specs, consumed by golden_t5_encoder_fn + + +def build_tensor_specs(): + global _GOLDEN_DATA + from golden import TensorSpec + + g = torch.Generator().manual_seed(42) + + norm1_w_1d = torch.ones(T5_DIM, dtype=torch.float32) + norm2_w_1d = torch.ones(T5_DIM, dtype=torch.float32) + final_norm_w_1d = torch.ones(T5_DIM, dtype=torch.float32) + + def rn(shape): + return torch.empty(shape).normal_(generator=g) * 0.02 + + q_w = rn([T5_DIM, T5_DIM]) + k_w = rn([T5_DIM, T5_DIM]) + v_w = rn([T5_DIM, T5_DIM]) + o_w = rn([T5_DIM, T5_DIM]) + rel_pos_bias = rn([T5_NUM_BUCKETS, T5_HEADS]) + wi_0 = rn([T5_FFN, T5_DIM]) + wi_1 = rn([T5_FFN, T5_DIM]) + wo = rn([T5_DIM, T5_FFN]) + + x_batch = torch.empty(1, T5_SEQ, T5_DIM).normal_(generator=g) + x_flat = x_batch.squeeze(0).contiguous() + + norm1_w_stacked = norm1_w_1d.unsqueeze(0).unsqueeze(0).expand(T5_LAYERS, 1, T5_DIM).contiguous() + norm2_w_stacked = norm2_w_1d.unsqueeze(0).unsqueeze(0).expand(T5_LAYERS, 1, T5_DIM).contiguous() + final_norm_w_2d = final_norm_w_1d.unsqueeze(0) + + q_w_stacked = q_w.bfloat16().unsqueeze(0).expand(T5_LAYERS, T5_DIM, T5_DIM).contiguous() + k_w_stacked = k_w.bfloat16().unsqueeze(0).expand(T5_LAYERS, T5_DIM, T5_DIM).contiguous() + v_w_stacked = v_w.bfloat16().unsqueeze(0).expand(T5_LAYERS, T5_DIM, T5_DIM).contiguous() + o_w_stacked = o_w.bfloat16().unsqueeze(0).expand(T5_LAYERS, T5_DIM, T5_DIM).contiguous() + wi_0_stacked = wi_0.bfloat16().unsqueeze(0).expand(T5_LAYERS, T5_FFN, T5_DIM).contiguous() + wi_1_stacked = wi_1.bfloat16().unsqueeze(0).expand(T5_LAYERS, T5_FFN, T5_DIM).contiguous() + wo_stacked = wo.bfloat16().unsqueeze(0).expand(T5_LAYERS, T5_DIM, T5_FFN).contiguous() + + pos_emb = T5RelativeEmbedding(rel_pos_bias, T5_NUM_BUCKETS, T5_HEADS, bidirectional=True) + pos_bias = pos_emb(T5_SEQ, T5_SEQ).squeeze(0) + pos_bias_stacked = pos_bias.unsqueeze(0).expand(T5_LAYERS, T5_HEADS, T5_SEQ, T5_SEQ).contiguous() + + def _identity(t): + return lambda: t + + specs = [ + TensorSpec("x_in", [T5_SEQ, T5_DIM], torch.float32, init_value=_identity(x_flat)), + TensorSpec("norm1_w", [T5_LAYERS, 1, T5_DIM], torch.float32, init_value=_identity(norm1_w_stacked)), + TensorSpec("q_w", [T5_LAYERS, T5_DIM, T5_DIM], torch.bfloat16, init_value=_identity(q_w_stacked)), + TensorSpec("k_w", [T5_LAYERS, T5_DIM, T5_DIM], torch.bfloat16, init_value=_identity(k_w_stacked)), + TensorSpec("v_w", [T5_LAYERS, T5_DIM, T5_DIM], torch.bfloat16, init_value=_identity(v_w_stacked)), + TensorSpec("o_w", [T5_LAYERS, T5_DIM, T5_DIM], torch.bfloat16, init_value=_identity(o_w_stacked)), + TensorSpec("pos_bias", [T5_LAYERS, T5_HEADS, T5_SEQ, T5_SEQ], torch.float32, init_value=_identity(pos_bias_stacked)), + TensorSpec("norm2_w", [T5_LAYERS, 1, T5_DIM], torch.float32, init_value=_identity(norm2_w_stacked)), + TensorSpec("wi_0", [T5_LAYERS, T5_FFN, T5_DIM], torch.bfloat16, init_value=_identity(wi_0_stacked)), + TensorSpec("wi_1", [T5_LAYERS, T5_FFN, T5_DIM], torch.bfloat16, init_value=_identity(wi_1_stacked)), + TensorSpec("wo", [T5_LAYERS, T5_DIM, T5_FFN], torch.bfloat16, init_value=_identity(wo_stacked)), + TensorSpec("final_norm_w", [1, T5_DIM], torch.float32, init_value=_identity(final_norm_w_2d)), + ] + specs.append(TensorSpec("out", [T5_SEQ, T5_DIM], torch.bfloat16, is_output=True)) + + _GOLDEN_DATA = { + "x_batch": x_batch, "norm1_w": norm1_w_1d, "q_w": q_w, "k_w": k_w, + "v_w": v_w, "o_w": o_w, "rel_pos_bias": rel_pos_bias, + "norm2_w": norm2_w_1d, "wi_0": wi_0, "wi_1": wi_1, "wo": wo, + "final_norm_w": final_norm_w_1d, + } + return specs + + +def golden_t5_encoder_fn(tensors): + """Run golden T5 encoder and fill tensors['out'].""" + x_expected = _golden_t5_encoder( + _GOLDEN_DATA["x_batch"], _GOLDEN_DATA["norm1_w"], + _GOLDEN_DATA["q_w"], _GOLDEN_DATA["k_w"], _GOLDEN_DATA["v_w"], _GOLDEN_DATA["o_w"], + _GOLDEN_DATA["rel_pos_bias"], _GOLDEN_DATA["norm2_w"], + _GOLDEN_DATA["wi_0"], _GOLDEN_DATA["wi_1"], _GOLDEN_DATA["wo"], + _GOLDEN_DATA["final_norm_w"], + ) + tensors["out"][:] = x_expected.squeeze(0).bfloat16() + + +if __name__ == "__main__": + from golden import ratio_allclose, run_jit + + parser = argparse.ArgumentParser(description="T5 Encoder pypto3.0 kernel test") + parser.add_argument("-p", "--platform", default="a2a3", + choices=["a2a3", "a2a3sim", "a5", "a5sim"]) + parser.add_argument("-d", "--device", type=int, default=0) + parser.add_argument("--enable-l2-swimlane", action="store_true", default=False) + args = parser.parse_args() + + specs = build_tensor_specs() + + result = run_jit( + fn=t5_encoder, + specs=specs, + golden_fn=golden_t5_encoder_fn, + runtime_cfg=dict( + platform=args.platform, + device_id=args.device, + enable_l2_swimlane=args.enable_l2_swimlane, + ), + rtol=3e-3, + atol=3e-3, + compare_fn={"out": ratio_allclose(atol=3e-3, rtol=3e-3, max_error_ratio=0.02)}, + ) + if not result.passed: + if result.error: + print(result.error) + raise SystemExit(1) diff --git a/models/world_model/vae_decoder.py b/models/world_model/vae_decoder.py new file mode 100644 index 00000000..89b3c980 --- /dev/null +++ b/models/world_model/vae_decoder.py @@ -0,0 +1,364 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""VAE decoder for Fun-Control 1.3B — pypto3.0 implementation. + +Implements the VAE decoder sub-network from the pipeline in +``infer_fun_control_1_3b_text.py`` (DiffSynth WanVideoPipeline VAE). +All arithmetic runs in pypto kernels. im2col / reshape / permute / +nearest-neighbour upsample are host-side data rearrangement only. + +The golden reference for precision comparison is +``test_golden_fun_control_full.py::vae_decode``, a pure-PyTorch +reimplementation of the SAME VAE logic used by the DiffSynth pipeline. +The two are logically identical because ``test_golden_fun_control_full.py`` +replicates every VAE computation from the DiffSynth source that +``infer_fun_control_1_3b_text.py`` invokes via ``WanVideoPipeline``. + +Reuses all building blocks from ``vae_encoder`` (ResBlock, AttentionBlock, +CausalConv3d, RMS norm, SiLU, softmax, matmul, etc.). Two new components: + - ``Upsample3d``: temporal CausalConv3d + nearest x2 + spatial Conv2d + - ``clamp_k``: NPU kernel for torch.clamp(-1, 1) + +Legitimate differences from the golden reference: + + 1. **Feature caching not implemented.** The golden reference maintains + ``feat_cache`` / ``feat_idx`` lists for chunked autoregressive + inference where successive chunks share temporal context. The pypto + implementation targets single-chunk inference (CHUNK_SIZE=5 golden + case); multi-chunk caching is a host-side orchestration concern + handled at the pipeline integration level, not inside the NPU kernel. + 2. **LAT_F = 1** (not config.LAT_F = 2). For CHUNK_SIZE=5 the encoder's + natural temporal output is 1 frame (3× downsample: 5→3→2→1). The + golden test uses LAT_F=2 as a DiT-compatible target and interpolates. + The pypto encoder/decoder work at the encoder's natural resolution + (T=1); any DiT-required interpolation is pipeline-level remapping. + 3. **Attention processes per-temporal-slice** (loop over ``batch_temporal``). + The golden uses ``F.scaled_dot_product_attention`` which handles all + heads/slices at once. The pypto version processes each temporal slice + independently because spatial dimensions for the golden case are small + (8×8 → batch_temporal ≤ 2). + +Architecture (cache=None, CHUNK_SIZE=5):: + + Input [1, 16, 1, 8, 8] + → CausalConv3d(16→1536, k=3, s=1) + → Middle: ResBlock(1536)+Attention(1536)+ResBlock(1536) + → Stage 0: ResBlock(1536→768)+ResBlock(768)+Upsample3d(t_up=T) → [1,768,2,16,16] + → Stage 1: ResBlock(768→384)+ResBlock(384)+Upsample3d(t_up=T) → [1,384,4,32,32] + → Stage 2: ResBlock(384→192)+ResBlock(192)+Upsample3d(t_up=F) → [1,192,4,64,64] + → Stage 3: ResBlock(192→96)+ResBlock(96) + → CausalConv3d(96→3, k=3, s=1) → [1,3,4,64,64] + → clamp(-1, 1) + +Precision comparison uses ``ratio_allclose`` with tolerance atol=0.1 / +rtol=0.1 / max_error_ratio=0.02 — same pattern as decode_layer.py but with +looser tolerances justified by the deeper VAE network (25+ RMSNorm ops, +each contributing ~0.012 hardware accuracy loss on Ascend vector unit; +see memory.md Phase 3b for per-op precision analysis). + +Usage:: + + python vae_decoder.py -p a2a3 -d 0 +""" + +import argparse +import sys +import time + +import pypto.language as pl +import torch +import torch.nn.functional as F + +import vae_encoder as ve +from config import VAE_Z_DIM, H, W, VAE_DEC_CH + +# ══════════════════════════════════════════════════════════════════════════════ +# Constants +# ══════════════════════════════════════════════════════════════════════════════ + +# LAT_F = 1 is the encoder's natural temporal output for CHUNK_SIZE=5 +# (3× CausalConv3d stride=(2,1,1) downsample: 5→3→2→1). +# config.LAT_F = 2 is the DiT-compatible target; the golden pipeline +# interpolates to it. The pypto encoder/decoder pair uses the natural +# resolution; DiT interpolation is a pipeline-level remap. +LAT_F = 1 + +# ── Tiling constants for the clamp kernel (reuse encoder's proven tile sizes) ── + +_PAD_COL = ve.COL_CHUNK # 192 — column tile for elementwise clamp (same as encoder) +_R_TILE = ve.R_TILE # 16 — row tile for pl.parallel dispatch (same as encoder) + +# ── Geometry assertions (same pattern as decode_layer.py) ── +assert len(VAE_DEC_CH) == 5, f"VAE_DEC_CH must have 5 entries (middle + 4 stages), got {len(VAE_DEC_CH)}" +assert VAE_DEC_CH[0] == 1536, f"VAE_DEC_CH[0] must be 1536 (middle bottleneck), got {VAE_DEC_CH[0]}" +assert VAE_DEC_CH[4] == 96, f"VAE_DEC_CH[4] must be 96 (last stage out channels), got {VAE_DEC_CH[4]}" +assert LAT_F == 1, f"LAT_F must be 1 (encoder's natural T output for CHUNK_SIZE=5), got {LAT_F}" + + +# ══════════════════════════════════════════════════════════════════════════════ +# NPU Kernel — clamp(-1, 1) +# ══════════════════════════════════════════════════════════════════════════════ + +@pl.jit +def clamp_k( + x: pl.Tensor, + out: pl.Out[pl.Tensor], +): + """clamp(x, -1, 1) on 2D tensor. Pre-padded to R_TILE rows and COL_CHUNK cols.""" + R = x.shape[0] + COLS = x.shape[1] + for r0 in pl.parallel(0, R, _R_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="clamp"): + for c0 in pl.range(0, COLS, _PAD_COL): + row = pl.slice(x, [_R_TILE, _PAD_COL], [r0, c0]) + clipped = pl.minimum(pl.maximum(row, -1.0), 1.0) + out = pl.assemble(out, pl.cast(clipped, target_type=pl.BF16), [r0, c0]) + return out + + +# ══════════════════════════════════════════════════════════════════════════════ +# Host runner — clamp on NPU +# ══════════════════════════════════════════════════════════════════════════════ + +def _npu_clamp(x): + """clamp(x, -1, 1) → BF16. x: 2D FP32 tensor.""" + R, C = x.shape[0], x.shape[1] + xp = ve._pad_cols(x, _PAD_COL) + xp = ve._pad_rows(xp, _R_TILE) + out = torch.zeros(xp.shape[0], xp.shape[1], dtype=torch.bfloat16) + + key = ('clamp', xp.shape[0], xp.shape[1]) + compiled = ve._compile(key, clamp_k, xp, out) + ve._run(compiled, [xp, out]) + return out[:R, :C] + + +# ══════════════════════════════════════════════════════════════════════════════ +# Building block — Upsample3d (host orchestration, NPU arithmetic) +# ══════════════════════════════════════════════════════════════════════════════ + +def upsample3d_npu(x, wd, temporal_up): + """Upsample3d: temporal CausalConv3d + nearest x2 + spatial Conv2d. + + All arithmetic (conv matmul) on NPU. Nearest-neighbour interpolation + and reshape/permute on host (data rearrangement, no arithmetic). + """ + b, c, t, h, w = x.size() + + # ── Temporal upsample ── + if temporal_up: + x = ve.causal_conv3d_npu_fp32(x, wd['time_conv_w'], wd.get('time_conv_b'), + stride=1, padding=(1, 0, 0)) + x = x.reshape(b, 2, c, t, h, w) + x = x.permute(0, 2, 4, 1, 3, 5).contiguous() + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + + # ── Spatial upsample ── + x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w).contiguous() + # Nearest-neighbour x2 — data rearrangement (sampling, no arithmetic) + x = F.interpolate(x, scale_factor=2.0, mode='nearest') + x = ve.conv2d_npu_fp32(x, wd['conv_w'], wd.get('conv_b'), + stride=1, padding=1) + x = x.reshape(b, t, -1, h * 2, w * 2).permute(0, 2, 1, 3, 4).contiguous() + return x + + +# ══════════════════════════════════════════════════════════════════════════════ +# Full VAE decoder — FP32 internal path, all arithmetic on NPU. +# ══════════════════════════════════════════════════════════════════════════════ + +def vae_decode_npu(latents, w): + """Full VAE decoder: latents [1,16,1,8,8] → video [1,3,4,64,64]. + + All arithmetic runs on NPU kernels. im2col, nearest-neighbour upsample, + reshape, and permute are host-side data rearrangement only. + """ + h = latents.float() + + # Input conv: 16 → 1536 + h = ve.causal_conv3d_npu_fp32(h, w["vae_dec_in_w"], w["vae_dec_in_b"], + stride=1, padding=1) + + # Middle: ResBlock(1536) + AttentionBlock(1536) + ResBlock(1536) + mr1w = {k: w[f"vae_dec_mid_r1_{k}"] + for k in ['norm1_w', 'conv1_w', 'conv1_b', + 'norm2_w', 'conv2_w', 'conv2_b']} + h = ve.resblock_npu(h, mr1w, 1536, 1536) + + maw = {k: w[f"vae_dec_mid_attn_{k}"] + for k in ['norm_w', 'to_qkv_w', 'to_qkv_b', 'proj_w', 'proj_b']} + h = ve.attention_block_npu(h, maw, 1536) + + mr2w = {k: w[f"vae_dec_mid_r2_{k}"] + for k in ['norm1_w', 'conv1_w', 'conv1_b', + 'norm2_w', 'conv2_w', 'conv2_b']} + h = ve.resblock_npu(h, mr2w, 1536, 1536) + + # 4 stages: ResBlocks + Upsample + vae_dec_ch = VAE_DEC_CH + for s in range(4): + # ResBlock 1 (may change channels) + r1w = {k: w[f"vae_dec_s{s}_r1_{k}"] + for k in ['norm1_w', 'conv1_w', 'conv1_b', + 'norm2_w', 'conv2_w', 'conv2_b']} + if vae_dec_ch[s] != vae_dec_ch[s + 1]: + r1w['shortcut_w'] = w[f"vae_dec_s{s}_r1_shortcut_w"] + r1w['shortcut_b'] = w.get(f"vae_dec_s{s}_r1_shortcut_b") + h = ve.resblock_npu(h, r1w, vae_dec_ch[s], vae_dec_ch[s + 1]) + + # ResBlock 2 (same channels) + r2w = {k: w[f"vae_dec_s{s}_r2_{k}"] + for k in ['norm1_w', 'conv1_w', 'conv1_b', + 'norm2_w', 'conv2_w', 'conv2_b']} + h = ve.resblock_npu(h, r2w, vae_dec_ch[s + 1], vae_dec_ch[s + 1]) + + # Upsample (stages 0,1,2): temporal_up for s < 2 + if s < 3: + rsw = {k: w[f"vae_dec_s{s}_resample_{k}"] + for k in ['conv_w', 'conv_b', + 'time_conv_w', 'time_conv_b']} + h = upsample3d_npu(h, rsw, temporal_up=(s < 2)) + + # Output conv: 96 → 3 + h = ve.causal_conv3d_npu_fp32(h, w["vae_dec_out_w"], w["vae_dec_out_b"], + stride=1, padding=1) + + # clamp(-1, 1) + h2, info = ve._to_2d(h) + h2 = h2[:ve._real_rows(info)] + r2 = _npu_clamp(h2) + return ve._from_2d(r2, info) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Weight generation — matches golden reference make_weights (decoder part). +# ══════════════════════════════════════════════════════════════════════════════ + +def make_vae_dec_weights(seed=42): + g = torch.Generator().manual_seed(seed) + w = {} + w["vae_dec_in_w"] = torch.randn(1536, VAE_Z_DIM, 3, 3, 3, generator=g) * 0.02 + w["vae_dec_in_b"] = torch.zeros(1536) + + p = "vae_dec_mid_r1" + w[f"{p}_norm1_w"] = torch.ones(1536) + w[f"{p}_conv1_w"] = torch.randn(1536, 1536, 3, 3, 3, generator=g) * 0.02 + w[f"{p}_conv1_b"] = torch.zeros(1536) + w[f"{p}_norm2_w"] = torch.ones(1536) + w[f"{p}_conv2_w"] = torch.randn(1536, 1536, 3, 3, 3, generator=g) * 0.02 + w[f"{p}_conv2_b"] = torch.zeros(1536) + + p = "vae_dec_mid_attn" + w[f"{p}_norm_w"] = torch.ones(1536) + w[f"{p}_to_qkv_w"] = torch.randn(1536 * 3, 1536, 1, 1, generator=g) * 0.02 + w[f"{p}_to_qkv_b"] = torch.zeros(1536 * 3) + w[f"{p}_proj_w"] = torch.randn(1536, 1536, 1, 1, generator=g) * 0.02 + w[f"{p}_proj_b"] = torch.zeros(1536) + + p = "vae_dec_mid_r2" + w[f"{p}_norm1_w"] = torch.ones(1536) + w[f"{p}_conv1_w"] = torch.randn(1536, 1536, 3, 3, 3, generator=g) * 0.02 + w[f"{p}_conv1_b"] = torch.zeros(1536) + w[f"{p}_norm2_w"] = torch.ones(1536) + w[f"{p}_conv2_w"] = torch.randn(1536, 1536, 3, 3, 3, generator=g) * 0.02 + w[f"{p}_conv2_b"] = torch.zeros(1536) + + vae_dec_ch = VAE_DEC_CH + for s in range(4): + p = f"vae_dec_s{s}_r1" + w[f"{p}_norm1_w"] = torch.ones(vae_dec_ch[s]) + w[f"{p}_conv1_w"] = torch.randn(vae_dec_ch[s + 1], vae_dec_ch[s], 3, 3, 3, generator=g) * 0.02 + w[f"{p}_conv1_b"] = torch.zeros(vae_dec_ch[s + 1]) + w[f"{p}_norm2_w"] = torch.ones(vae_dec_ch[s + 1]) + w[f"{p}_conv2_w"] = torch.randn(vae_dec_ch[s + 1], vae_dec_ch[s + 1], 3, 3, 3, generator=g) * 0.02 + w[f"{p}_conv2_b"] = torch.zeros(vae_dec_ch[s + 1]) + if vae_dec_ch[s] != vae_dec_ch[s + 1]: + w[f"{p}_shortcut_w"] = torch.randn(vae_dec_ch[s + 1], vae_dec_ch[s], 1, 1, 1, generator=g) * 0.02 + w[f"{p}_shortcut_b"] = torch.zeros(vae_dec_ch[s + 1]) + + p = f"vae_dec_s{s}_r2" + w[f"{p}_norm1_w"] = torch.ones(vae_dec_ch[s + 1]) + w[f"{p}_conv1_w"] = torch.randn(vae_dec_ch[s + 1], vae_dec_ch[s + 1], 3, 3, 3, generator=g) * 0.02 + w[f"{p}_conv1_b"] = torch.zeros(vae_dec_ch[s + 1]) + w[f"{p}_norm2_w"] = torch.ones(vae_dec_ch[s + 1]) + w[f"{p}_conv2_w"] = torch.randn(vae_dec_ch[s + 1], vae_dec_ch[s + 1], 3, 3, 3, generator=g) * 0.02 + w[f"{p}_conv2_b"] = torch.zeros(vae_dec_ch[s + 1]) + + if s < 3: + p = f"vae_dec_s{s}_resample" + w[f"{p}_conv_w"] = torch.randn(vae_dec_ch[s + 1], vae_dec_ch[s + 1], 3, 3, generator=g) * 0.02 + w[f"{p}_conv_b"] = torch.zeros(vae_dec_ch[s + 1]) + w[f"{p}_time_conv_w"] = torch.randn(vae_dec_ch[s + 1] * 2, vae_dec_ch[s + 1], 3, 1, 1, generator=g) * 0.02 + w[f"{p}_time_conv_b"] = torch.zeros(vae_dec_ch[s + 1] * 2) + + w["vae_dec_out_w"] = torch.randn(3, 96, 3, 3, 3, generator=g) * 0.02 + w["vae_dec_out_b"] = torch.zeros(3) + return w + + +# ══════════════════════════════════════════════════════════════════════════════ +# Golden reference — uses the EXACT vae_decode from test_golden_fun_control_full.py +# ══════════════════════════════════════════════════════════════════════════════ + +def _golden_vae_decode(latents, w): + """Golden reference: imports and runs the SAME code as the test file.""" + _GOLDEN_DIR = '/data/x00952168/pypto3.0/cann-recipes-embodied-ai/world_model/agibot-arm-world-model/infer_with_torch' + if _GOLDEN_DIR not in sys.path: + sys.path.insert(0, _GOLDEN_DIR) + from test_golden_fun_control_full import vae_decode as ref_vae_decode + return ref_vae_decode(latents, w) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Test — uses golden.runner._Stage + golden.validation.validate_golden, +# same harness primitives as run_jit in t5_encoder.py. +# (Multi-kernel architecture: individual @pl.jit kernels are JIT-compiled +# on first use during the runtime stage; no standalone compile pass.) +# ══════════════════════════════════════════════════════════════════════════════ + + +if __name__ == "__main__": + from golden.runner import _Stage + from golden.validation import validate_golden, ratio_allclose + + parser = argparse.ArgumentParser(description="VAE Decoder — pypto3.0") + parser.add_argument( + "-p", "--platform", type=str, default="a2a3", + choices=["a2a3", "a2a3sim", "a5", "a5sim"], + ) + parser.add_argument("-d", "--device", type=int, default=0, + help="NPU device id (soldered chip index)") + args = parser.parse_args() + + ve.set_device(args.device) + t_total = time.time() + + with _Stage("generate inputs"): + torch.manual_seed(42) + w = make_vae_dec_weights(seed=42) + latents = (torch.randn(1, VAE_Z_DIM, LAT_F, H // 8, W // 8) * 0.1).bfloat16() + + with _Stage("compute golden"): + ref = _golden_vae_decode(latents.float(), w) + + with _Stage("runtime"): + npu = vae_decode_npu(latents, w) + + with _Stage("validate"): + validate_golden( + outputs={"out": npu}, + golden={"out": ref}, + rtol=0.1, + atol=0.1, + compare_fn={"out": ratio_allclose(atol=0.1, rtol=0.1, max_error_ratio=0.02)}, + ) + + total = time.time() - t_total + print(f"[RUN] PASS ({total:.2f}s)", flush=True) diff --git a/models/world_model/vae_encoder.py b/models/world_model/vae_encoder.py new file mode 100644 index 00000000..54b6b1e0 --- /dev/null +++ b/models/world_model/vae_encoder.py @@ -0,0 +1,825 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""VAE encoder for Fun-Control 1.3B — pypto3.0 implementation. + +Implements the VAE encoder sub-network from the pipeline in +``infer_fun_control_1_3b_text.py`` (DiffSynth WanVideoPipeline VAE). +All arithmetic (matmul, SiLU, RMS norm, softmax, attention) runs inside +pypto kernels. im2col / reshape / permute are data-rearrangement only. + +The golden reference for precision comparison is +``test_golden_fun_control_full.py::vae_encode``, a pure-PyTorch +reimplementation of the SAME VAE logic used by the DiffSynth pipeline. +The two are logically identical because ``test_golden_fun_control_full.py`` +replicates every VAE computation from the DiffSynth source that +``infer_fun_control_1_3b_text.py`` invokes via ``WanVideoPipeline``. + +Legitimate differences from the golden reference: + + 1. **Feature caching not implemented.** The golden reference maintains + ``feat_cache`` / ``feat_idx`` lists for chunked autoregressive + inference where successive chunks share temporal context. The pypto + implementation targets single-chunk inference (CHUNK_SIZE=5 golden + case); multi-chunk caching is a host-side orchestration concern + handled at the pipeline integration level, not inside the NPU kernel. + 2. **Output temporal dimension = 1** (not 2). For CHUNK_SIZE=5 the + encoder's natural temporal output is 1 (3× CausalConv3d stride=(2,1,1) + downsample: 5→3→2→1). The golden test uses LAT_F=2 as a + DiT-compatible target and interpolates. The pypto encoder/decoder + work at the natural resolution; DiT interpolation is a + pipeline-level remap. + 3. **Attention processes per-temporal-slice** (loop over ``batch_temporal``). + The golden uses ``F.scaled_dot_product_attention`` which handles all + heads/slices at once. The pypto version processes each temporal slice + independently because spatial dimensions are small (8×8 → + batch_temporal ≤ 2). + +Architecture:: + + Input [1,3,5,64,64] + → CausalConv3d(3→96) + → Stage 0: ResBlock(96→192)+ResBlock(192)+Downsample → [1,192,3,32,32] + → Stage 1: ResBlock(192→384)+ResBlock(384)+Downsample → [1,384,2,16,16] + → Stage 2: ResBlock(384→768)+ResBlock(768)+Downsample → [1,768,1,8,8] + → Stage 3: ResBlock(768→1536)+ResBlock(1536) → [1,1536,1,8,8] + → Middle: ResBlock+Attention+ResBlock → [1,1536,1,8,8] + → CausalConv3d(1536→32) → chunk → mu [1,16,1,8,8] + → scale_norm → output [1,16,1,8,8] + +Multi-kernel design rationale: + Convolution im2col (F.pad + .unfold) is host-side data rearrangement. + Each convolution's matmul + each elementwise op (RMS_norm, SiLU, etc.) + runs in its own @pl.jit kernel, dispatched from host. This is the + standard pattern for conv-heavy networks in pypto. + +Precision comparison uses ``ratio_allclose`` with tolerance atol=0.1 / +rtol=0.1 / max_error_ratio=0.02 — same pattern as decode_layer.py but with +looser tolerances justified by the deeper VAE network (25+ RMSNorm ops, +each contributing ~0.012 hardware accuracy loss on Ascend vector unit; +see memory.md Phase 3 for per-op precision analysis). + +Usage:: + + python vae_encoder.py -p a2a3 -d 0 +""" + +import argparse +import sys +import time + +import pypto.language as pl +import torch +import torch.nn.functional as F + +from config import VAE_Z_DIM, H, W, CHUNK_SIZE, VAE_ENC_CH + +# ══════════════════════════════════════════════════════════════════════════════ +# Functional config — model architecture + workload. +# ══════════════════════════════════════════════════════════════════════════════ + +VAE_CH = VAE_ENC_CH # [96, 192, 384, 768, 1536] +VAE_OUT_CH = VAE_Z_DIM * 2 # 32 +T_IN = CHUNK_SIZE # 5 + +# ── Tiling parameters (optimization config) ── +# See decode_layer.py pattern: each kernel type gets its own tile size tuned +# to the NPU's M/N/K alignment requirements. Values verified on Ascend910B. + +R_TILE = 16 # row-tile for pl.parallel dispatch (mm_col_w / softmax / bias / scale / resadd) +K_CHUNK = 1024 # inner K-dim chunk for matmul accumulation loop +N_CHUNK = 64 # outer N-dim chunk for matmul output tiling +SILU_TILE = 64 # vector-row tile for SiLU elementwise +SILU_COL = 192 # vector-col tile for SiLU elementwise +RMS_TILE = 16 # vector-row tile for RMS norm +COL_CHUNK = 192 # vector-col tile for wide-column elementwise ops (add / bias / scale / scale_norm) +RMS_COL = 96 # vector-col chunk for RMS norm inner reduction loop + +RMS_EPS = 1e-12 # matching F.normalize default eps + +# ── Geometry / platform assertions (same pattern as decode_layer.py) ── +assert len(VAE_CH) == 5, f"VAE_CH must have 5 entries (4 stages + middle), got {len(VAE_CH)}" +assert VAE_CH[0] == 96, f"VAE_CH[0] must be 96 (first stage in channels), got {VAE_CH[0]}" +assert VAE_CH[4] == 1536, f"VAE_CH[4] must be 1536 (middle bottleneck), got {VAE_CH[4]}" +assert VAE_OUT_CH == VAE_Z_DIM * 2, f"VAE_OUT_CH must be VAE_Z_DIM*2={VAE_Z_DIM*2}, got {VAE_OUT_CH}" +assert T_IN == CHUNK_SIZE, f"T_IN must equal CHUNK_SIZE={CHUNK_SIZE}, got {T_IN}" + + +# ── Pad helpers (host-side, data rearrangement only) ── + +def _pad_cols(t, align): + c = t.shape[1] + p = ((c + align - 1) // align) * align + if p == c: + return t + o = torch.zeros(t.shape[0], p, dtype=t.dtype) + o[:, :c] = t + return o + + +def _pad_rows(t, align): + r = t.shape[0] + p = ((r + align - 1) // align) * align + if p == r: + return t + o = torch.zeros(p, t.shape[1], dtype=t.dtype) + o[:r] = t + return o + + +# ══════════════════════════════════════════════════════════════════════════════ +# NPU Kernels — each @pl.jit kernel runs one operation on NPU. +# ══════════════════════════════════════════════════════════════════════════════ + +@pl.jit +def mm_col_w( + col: pl.Tensor, + w: pl.Tensor, + out: pl.Out[pl.Tensor], +): + """matmul(col, w^T) with K-chunked accumulation → FP32 output.""" + R = col.shape[0] + K = col.shape[1] + N = w.shape[0] + for r0 in pl.parallel(0, R, R_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="mm"): + for n0 in pl.range(0, N, N_CHUNK): + ct0 = pl.slice(col, [R_TILE, K_CHUNK], [r0, 0]) + wt0 = pl.slice(w, [N_CHUNK, K_CHUNK], [n0, 0]) + acc = pl.matmul(ct0, wt0, b_trans=True, out_dtype=pl.FP32) + for k0 in pl.range(K_CHUNK, K, K_CHUNK): + ct1 = pl.slice(col, [R_TILE, K_CHUNK], [r0, k0]) + wt1 = pl.slice(w, [N_CHUNK, K_CHUNK], [n0, k0]) + acc = pl.matmul_acc(acc, ct1, wt1, b_trans=True) + out = pl.assemble(out, acc, [r0, n0]) + return out + + +@pl.jit +def softmax_k( + x: pl.Tensor, + out: pl.Out[pl.Tensor], +): + """Row-wise softmax.""" + N = x.shape[1] + for r0 in pl.parallel(0, x.shape[0], R_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="softmax"): + row = pl.cast(pl.slice(x, [R_TILE, N], [r0, 0]), target_type=pl.FP32) + shifted = pl.row_expand_sub(row, pl.row_max(row)) + e = pl.exp(shifted) + s = pl.row_expand_div(e, pl.row_sum(e)) + out = pl.assemble(out, pl.cast(s, target_type=pl.BF16), [r0, 0]) + return out + + +@pl.jit +def scale_k( + x: pl.Tensor, + scale: pl.Tensor, + out: pl.Out[pl.Tensor], +): + """Element-wise multiply by scalar. FP32 in → FP32 out.""" + COLS = x.shape[1] + s = pl.read(scale, [0, 0]) + for r0 in pl.parallel(0, x.shape[0], R_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="scale"): + for c0 in pl.range(0, COLS, COL_CHUNK): + row = pl.slice(x, [R_TILE, COL_CHUNK], [r0, c0]) + r = pl.mul(row, s) + out = pl.assemble(out, r, [r0, c0]) + return out + + +@pl.jit +def rms_norm_k( + x: pl.Tensor, + weight: pl.Tensor, + out: pl.Out[pl.Tensor], +): + """RMS_norm(x, weight) FP32 in → FP32 out. Matches F.normalize*√C*w.""" + R = x.shape[0] + COLS = x.shape[1] + cols_fp = pl.cast(COLS, target_type=pl.FP32) + for r0 in pl.parallel(0, R, RMS_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="rms"): + row0 = pl.slice(x, [RMS_TILE, RMS_COL], [r0, 0]) + sq_sum = pl.reshape(pl.row_sum(pl.mul(row0, row0)), [RMS_TILE, 1]) + for c0 in pl.range(RMS_COL, COLS, RMS_COL): + rc = pl.slice(x, [RMS_TILE, RMS_COL], [r0, c0]) + sq_sum = pl.add(sq_sum, pl.reshape(pl.row_sum(pl.mul(rc, rc)), [RMS_TILE, 1])) + variance = pl.add(pl.mul(sq_sum, 1.0 / cols_fp), RMS_EPS) + inv_rms = pl.recip(pl.sqrt(variance)) + for c0 in pl.range(0, COLS, RMS_COL): + rc = pl.slice(x, [RMS_TILE, RMS_COL], [r0, c0]) + g = pl.slice(weight, [1, RMS_COL], [0, c0]) + n = pl.col_expand_mul(pl.row_expand_mul(rc, inv_rms), g) + out = pl.assemble(out, n, [r0, c0]) + return out + + +@pl.jit +def silu_k( + x: pl.Tensor, + out: pl.Out[pl.Tensor], +): + """SiLU(x) FP32 in → FP32 out.""" + R = x.shape[0] + COLS = x.shape[1] + for r0 in pl.parallel(0, R, SILU_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="silu"): + for c0 in pl.range(0, COLS, SILU_COL): + row = pl.slice(x, [SILU_TILE, SILU_COL], [r0, c0]) + sig = pl.recip(pl.add(pl.exp(pl.neg(row)), 1.0)) + y = pl.mul(row, sig) + out = pl.assemble(out, y, [r0, c0]) + return out + + +@pl.jit +def residual_add_k( + x: pl.Tensor, + y: pl.Tensor, + out: pl.Out[pl.Tensor], +): + """x + y, FP32 in → FP32 out. Same shape.""" + COLS = x.shape[1] + for r0 in pl.parallel(0, x.shape[0], R_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="resadd"): + for c0 in pl.range(0, COLS, COL_CHUNK): + xr = pl.slice(x, [R_TILE, COL_CHUNK], [r0, c0]) + yr = pl.slice(y, [R_TILE, COL_CHUNK], [r0, c0]) + z = pl.add(xr, yr) + out = pl.assemble(out, z, [r0, c0]) + return out + + +@pl.jit +def bias_add_k( + x: pl.Tensor, + bias: pl.Tensor, + out: pl.Out[pl.Tensor], +): + """x + bias, FP32 in → FP32 out.""" + COLS = x.shape[1] + for r0 in pl.parallel(0, x.shape[0], R_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="bias"): + for c0 in pl.range(0, COLS, COL_CHUNK): + row = pl.slice(x, [R_TILE, COL_CHUNK], [r0, c0]) + bc = pl.slice(bias, [1, COL_CHUNK], [0, c0]) + z = pl.col_expand_add(row, bc) + out = pl.assemble(out, z, [r0, c0]) + return out + + +@pl.jit +def scale_norm_k( + x: pl.Tensor, + mean: pl.Tensor, + inv_std: pl.Tensor, + out: pl.Out[pl.Tensor], +): + """(x - mean) * inv_std, FP32 in → FP32 out.""" + COLS = x.shape[1] + for r0 in pl.parallel(0, x.shape[0], R_TILE): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="snorm"): + for c0 in pl.range(0, COLS, COL_CHUNK): + row = pl.slice(x, [R_TILE, COL_CHUNK], [r0, c0]) + mc = pl.slice(mean, [1, COL_CHUNK], [0, c0]) + sc = pl.slice(inv_std, [1, COL_CHUNK], [0, c0]) + r = pl.col_expand_add(row, pl.neg(mc)) + r = pl.col_expand_mul(r, sc) + out = pl.assemble(out, r, [r0, c0]) + return out + + +# ══════════════════════════════════════════════════════════════════════════════ +# Host helpers — im2col, reshape, pad (data rearrangement only, no arithmetic). +# ══════════════════════════════════════════════════════════════════════════════ + +def _to_2d(x): + """5D [B,C,T,H,W] or 4D [B,C,H,W] → 2D [B*T*H*W, C].""" + if x.dim() == 5: + B, C, T, Hh, Ww = x.shape + return x.permute(0, 2, 3, 4, 1).reshape(B * T * Hh * Ww, C).contiguous(), (B, T, Hh, Ww, C) + B, C, Hh, Ww = x.shape + return x.permute(0, 2, 3, 1).reshape(B * Hh * Ww, C).contiguous(), (B, Hh, Ww, C) + + +def _from_2d(x2, info): + """Reverse of _to_2d.""" + if len(info) == 5: + B, T, Hh, Ww, C = info + return x2.reshape(B, T, Hh, Ww, C).permute(0, 4, 1, 2, 3) + B, Hh, Ww, C = info + return x2.reshape(B, Hh, Ww, C).permute(0, 3, 1, 2) + + +def _real_rows(info): + if len(info) == 5: + return info[0] * info[1] * info[2] * info[3] + return info[0] * info[1] * info[2] + + +# ══════════════════════════════════════════════════════════════════════════════ +# Kernel runners — compile cache + device dispatch. +# ══════════════════════════════════════════════════════════════════════════════ + +_cache = {} +_DEVICE_ID = 0 + + +def set_device(device_id): + global _DEVICE_ID + _DEVICE_ID = device_id + + +def _compile(key, kernel, *dummies): + from pypto.runtime import RunConfig + if key not in _cache: + _cache[key] = kernel.compile(*dummies, config=RunConfig(platform="a2a3", device_id=_DEVICE_ID)) + return _cache[key] + + +def _run(compiled_info, args): + from pypto.runtime import execute_compiled + execute_compiled(compiled_info.output_dir, args, platform="a2a3", device_id=_DEVICE_ID) + + +# ── Host wrappers (pad inputs → compile/run kernel → slice result) ── + +def _npu_matmul(col, w): + """matmul(col, w^T) → FP32.""" + M_real, K_real, N_real = col.shape[0], col.shape[1], w.shape[0] + col_p = _pad_cols(col, K_CHUNK) + col_p = _pad_rows(col_p, R_TILE) + w_p = _pad_cols(w, K_CHUNK) + w_N_pad = ((N_real + N_CHUNK - 1) // N_CHUNK) * N_CHUNK + if w_N_pad > N_real: + wp2 = torch.zeros(w_N_pad, w_p.shape[1], dtype=w_p.dtype) + wp2[:N_real] = w_p + w_p = wp2 + out = torch.zeros(col_p.shape[0], w_N_pad, dtype=torch.float32) + key = ('mm_col_w', col_p.shape[0], w_p.shape[1], w_N_pad) + compiled = _compile(key, mm_col_w, col_p, w_p, out) + _run(compiled, [col_p, w_p, out]) + return out[:M_real, :N_real] + + +def _npu_softmax(x): + """Row-wise softmax.""" + N = x.shape[1] + xp = _pad_rows(x, R_TILE) + out = torch.zeros(xp.shape[0], N, dtype=torch.bfloat16) + key = ('softmax', xp.shape[0], N) + compiled = _compile(key, softmax_k, xp, out) + _run(compiled, [xp, out]) + return out[:x.shape[0], :] + + +def _npu_scale_fp32(x, scale_val): + """Multiply FP32 tensor by scalar on NPU.""" + R, C = x.shape[0], x.shape[1] + xp = _pad_cols(x, COL_CHUNK) + xp = _pad_rows(xp, R_TILE) + out = torch.zeros(xp.shape[0], xp.shape[1], dtype=torch.float32) + scale_t = torch.tensor([[scale_val]], dtype=torch.float32) + key = ('scale', xp.shape[0], xp.shape[1]) + compiled = _compile(key, scale_k, xp, scale_t, out) + _run(compiled, [xp, scale_t, out]) + return out[:R, :C] + + +def _npu_scale_norm_fp32(x, mean, inv_std): + """(x - mean) * inv_std on NPU.""" + R, C = x.shape[0], x.shape[1] + xp = _pad_cols(x, COL_CHUNK) + xp = _pad_rows(xp, R_TILE) + mp = _pad_cols(mean.reshape(1, -1), COL_CHUNK) + sp = _pad_cols(inv_std.reshape(1, -1), COL_CHUNK) + out = torch.zeros(xp.shape[0], xp.shape[1], dtype=torch.float32) + key = ('snorm', xp.shape[0], xp.shape[1]) + compiled = _compile(key, scale_norm_k, xp, mp, sp, out) + _run(compiled, [xp, mp, sp, out]) + return out[:R, :C] + + +def _npu_rms_norm_fp32(x, weight_1d): + """RMS_norm FP32 in → FP32 out.""" + R, C = x.shape[0], x.shape[1] + xp = _pad_cols(x, RMS_COL) + xp = _pad_rows(xp, RMS_TILE) + wp = _pad_cols(weight_1d.reshape(1, -1), RMS_COL) + out = torch.zeros(xp.shape[0], xp.shape[1], dtype=torch.float32) + key = ('rms', xp.shape[0], xp.shape[1]) + compiled = _compile(key, rms_norm_k, xp, wp, out) + _run(compiled, [xp, wp, out]) + return out[:R, :C] + + +def _npu_silu_fp32(x): + """SiLU FP32 in → FP32 out.""" + R, C = x.shape[0], x.shape[1] + xp = _pad_cols(x, SILU_COL) + xp = _pad_rows(xp, SILU_TILE) + out = torch.zeros(xp.shape[0], xp.shape[1], dtype=torch.float32) + key = ('silu', xp.shape[0], xp.shape[1]) + compiled = _compile(key, silu_k, xp, out) + _run(compiled, [xp, out]) + return out[:R, :C] + + +def _npu_residual_add_fp32(a, b): + """a + b on NPU.""" + R, C = a.shape[0], a.shape[1] + ap = _pad_cols(a, COL_CHUNK) + ap = _pad_rows(ap, R_TILE) + bp = _pad_cols(b, COL_CHUNK) + bp = _pad_rows(bp, R_TILE) + out = torch.zeros(ap.shape[0], ap.shape[1], dtype=torch.float32) + key = ('resadd', ap.shape[0], ap.shape[1]) + compiled = _compile(key, residual_add_k, ap, bp, out) + _run(compiled, [ap, bp, out]) + return out[:R, :C] + + +def _npu_bias_add_fp32(x, bias_1d): + """x + bias on NPU.""" + R, C = x.shape[0], x.shape[1] + xp = _pad_cols(x, COL_CHUNK) + xp = _pad_rows(xp, R_TILE) + bp = _pad_cols(bias_1d.reshape(1, -1), COL_CHUNK) + out = torch.zeros(xp.shape[0], xp.shape[1], dtype=torch.float32) + key = ('bias', xp.shape[0], xp.shape[1]) + compiled = _compile(key, bias_add_k, xp, bp, out) + _run(compiled, [xp, bp, out]) + return out[:R, :C] + + +def _npu_matmul_with_bias(col, w, bias): + """matmul(col, w^T) + bias on NPU.""" + r = _npu_matmul(col, w) + if bias is not None: + r = _npu_bias_add_fp32(r, bias.float()) + return r + + +# ══════════════════════════════════════════════════════════════════════════════ +# Building blocks — FP32 internal path, BF16 at matmul inputs only. +# ALL arithmetic runs on NPU kernels. +# ══════════════════════════════════════════════════════════════════════════════ + +def causal_conv3d_npu_fp32(x, weight, bias=None, stride=1, padding=1): + """Causal 3D conv. im2col on host (data rearrangement), matmul on NPU.""" + out_channels = weight.shape[0] + kd, kh, kw = weight.shape[2], weight.shape[3], weight.shape[4] + if isinstance(stride, int): + stride = (stride, stride, stride) + if isinstance(padding, int): + padding = (padding, padding, padding) + sd, sh, sw = stride + pd, ph, pw = padding + + if sd == 1 and pd == 1: + tpf, tpb = pd, pd + else: + tpf, tpb = pd * 2, 0 + + xp = F.pad(x, [pw, pw, ph, ph, tpf, tpb]) + b, ic, depth, hp, wp = xp.shape + od = (depth - kd) // sd + 1 + oh = (hp - kh) // sh + 1 + ow = (wp - kw) // sw + 1 + xu = xp.unfold(2, kd, sd).unfold(3, kh, sh).unfold(4, kw, sw) + xu = xu.contiguous().view(b, ic, od, oh, ow, kd * kh * kw) + xu = xu.permute(0, 2, 3, 4, 1, 5).contiguous() + col = xu.reshape(b * od * oh * ow, ic * kd * kh * kw).bfloat16() + + w2d = weight.reshape(out_channels, -1).contiguous().bfloat16() + bias_2d = bias.float().reshape(1, -1) if bias is not None else None + result = _npu_matmul_with_bias(col, w2d, bias_2d) + real_rows = od * oh * ow + result_trimmed = result[:real_rows, :] + return result_trimmed.reshape(1, od, oh, ow, out_channels).permute(0, 4, 1, 2, 3).contiguous() + + +def conv2d_npu_fp32(x, weight, bias=None, stride=1, padding=1): + """2D conv. im2col on host, matmul on NPU.""" + out_channels = weight.shape[0] + kh, kw = weight.shape[2], weight.shape[3] + sh = stride if not isinstance(stride, tuple) else stride[0] + sw = stride if not isinstance(stride, tuple) else stride[1] + if isinstance(padding, int): + ph, pw = padding, padding + else: + ph, pw = padding[0], padding[1] + + xp = F.pad(x, [pw, pw, ph, ph]) + b, ic, hp, wp = xp.shape + oh = (hp - kh) // sh + 1 + ow = (wp - kw) // sw + 1 + xu = xp.unfold(2, kh, sh).unfold(3, kw, sw) + xu = xu.contiguous().view(b, ic, oh, ow, kh * kw) + xu = xu.permute(0, 2, 3, 1, 4).contiguous() + col = xu.reshape(b * oh * ow, ic * kh * kw).bfloat16() + + w2d = weight.reshape(out_channels, -1).contiguous().bfloat16() + bias_2d = bias.float().reshape(1, -1) if bias is not None else None + result = _npu_matmul_with_bias(col, w2d, bias_2d) + real_rows = b * oh * ow + return result[:real_rows, :].reshape(b, oh, ow, out_channels).permute(0, 3, 1, 2).contiguous() + + +def rms_norm_npu_fp32(x, weight_1d): + """RMS_norm FP32 in → FP32 out on NPU.""" + x_2d, info = _to_2d(x.float()) + result_2d = _npu_rms_norm_fp32(x_2d, weight_1d.float()) + return _from_2d(result_2d[:_real_rows(info)], info) + + +def silu_npu_fp32(x): + """SiLU FP32 in → FP32 out on NPU.""" + x_2d, info = _to_2d(x.float()) + result_2d = _npu_silu_fp32(x_2d) + return _from_2d(result_2d[:_real_rows(info)], info) + + +def resblock_npu(x, weight_dict, in_dim, out_dim): + """ResBlock: norm→SiLU→conv→norm→SiLU→conv + shortcut. All arithmetic on NPU.""" + x = x.float() + hidden = rms_norm_npu_fp32(x, weight_dict['norm1_w']) + hidden = silu_npu_fp32(hidden) + hidden = causal_conv3d_npu_fp32(hidden, weight_dict['conv1_w'], weight_dict.get('conv1_b'), + stride=1, padding=1) + hidden = rms_norm_npu_fp32(hidden, weight_dict['norm2_w']) + hidden = silu_npu_fp32(hidden) + hidden = causal_conv3d_npu_fp32(hidden, weight_dict['conv2_w'], weight_dict.get('conv2_b'), + stride=1, padding=1) + + if in_dim != out_dim: + shortcut = causal_conv3d_npu_fp32(x, weight_dict['shortcut_w'], weight_dict.get('shortcut_b'), + stride=1, padding=0) + else: + shortcut = x + + hidden_2d, info = _to_2d(hidden) + shortcut_2d, _ = _to_2d(shortcut) + result_2d = _npu_residual_add_fp32(hidden_2d[:_real_rows(info)], shortcut_2d[:_real_rows(info)]) + return _from_2d(result_2d, info) + + +def attention_block_npu(x, weight_dict, dim): + """AttentionBlock: norm→QKV(1x1conv)→SDPA→proj + residual. All arithmetic on NPU.""" + x = x.float() + identity = x + batch, channels, temporal, height, width = x.size() + batch_temporal = batch * temporal + spatial_size = height * width + + x_2d = x.permute(0, 2, 1, 3, 4).reshape(batch_temporal, channels, height, width).contiguous() + x_2d_norm = rms_norm_npu_fp32(x_2d, weight_dict['norm_w']) + + qkv = conv2d_npu_fp32(x_2d_norm, weight_dict['to_qkv_w'], weight_dict.get('to_qkv_b'), + stride=1, padding=0) + qkv_r = qkv.reshape(batch_temporal, channels * 3, spatial_size).permute(0, 2, 1).contiguous() + q_pf = qkv_r[:, :, :channels].contiguous() + k_pf = qkv_r[:, :, channels:2 * channels].contiguous() + v_pf = qkv_r[:, :, 2 * channels:].contiguous() + + scale_factor = 1.0 / (channels ** 0.5) + attn_out_flat = torch.zeros(batch_temporal * spatial_size, channels, dtype=torch.float32) + + for i in range(batch_temporal): + q_i = q_pf[i] + k_i = k_pf[i] + v_i = v_pf[i] + scores = _npu_matmul(q_i.bfloat16(), k_i.bfloat16()) + scores = _npu_scale_fp32(scores, scale_factor) + softmax_out = _npu_softmax(scores.bfloat16()) + head_out = _npu_matmul(softmax_out, v_i.bfloat16().transpose(0, 1)) + attn_out_flat[i * spatial_size:(i + 1) * spatial_size, :] = head_out[:spatial_size, :channels] + + proj_w2d = weight_dict['proj_w'].reshape(channels, channels).bfloat16() + proj_result = _npu_matmul(attn_out_flat.bfloat16(), proj_w2d) + if weight_dict.get('proj_b') is not None: + proj_result = _npu_bias_add_fp32(proj_result, weight_dict['proj_b'].float()) + + x_2d_out = proj_result[:batch_temporal * spatial_size, :channels].reshape( + batch_temporal, spatial_size, channels).permute(0, 2, 1) + x_2d_out = x_2d_out.reshape(batch_temporal, channels, height, width).contiguous() + x_5d_out = x_2d_out.reshape(batch, temporal, channels, height, width).permute(0, 2, 1, 3, 4).contiguous() + + x_5d_2d, info = _to_2d(x_5d_out) + identity_2d, _ = _to_2d(identity) + result_2d = _npu_residual_add_fp32(x_5d_2d[:_real_rows(info)], identity_2d[:_real_rows(info)]) + return _from_2d(result_2d, info) + + +def downsample3d_npu(x, weight_dict, dim): + """Downsample3d: spatial Conv2d(s=2) + temporal CausalConv3d(s=(2,1,1)).""" + x = x.float() + batch, channels, temporal, height, width = x.size() + + x_2d = x.permute(0, 2, 1, 3, 4).reshape(batch * temporal, channels, height, width).contiguous() + x_2d = F.pad(x_2d, (0, 1, 0, 1)) + x_2d_out = conv2d_npu_fp32(x_2d, weight_dict['conv_w'], weight_dict.get('conv_b'), + stride=2, padding=0) + _, _, hd, wd = x_2d_out.shape + x_5d = x_2d_out.reshape(batch, temporal, channels, hd, wd).permute(0, 2, 1, 3, 4).contiguous() + + x_5d = causal_conv3d_npu_fp32(x_5d, weight_dict['time_conv_w'], weight_dict.get('time_conv_b'), + stride=(2, 1, 1), padding=(1, 0, 0)) + return x_5d + + +# ══════════════════════════════════════════════════════════════════════════════ +# Full VAE encoder — FP32 internal path, all arithmetic on NPU. +# ══════════════════════════════════════════════════════════════════════════════ + +def vae_encode_npu(video, w): + """Full VAE encoder: video [1,3,T,64,64] → mu [1,16,1,8,8]. + + All arithmetic runs on NPU kernels. im2col (F.pad + .unfold), reshape, + permute, and dtype cast are host-side data rearrangement only. + """ + hidden = video.float() + hidden = causal_conv3d_npu_fp32(hidden, w["vae_enc_in_w"], w["vae_enc_in_b"], + stride=1, padding=1) + + channels = VAE_CH + for stage_idx in range(4): + r1w = {k: w[f"vae_enc_s{stage_idx}_r1_{k}"] + for k in ['norm1_w', 'conv1_w', 'conv1_b', + 'norm2_w', 'conv2_w', 'conv2_b']} + if channels[stage_idx] != channels[stage_idx + 1]: + r1w['shortcut_w'] = w[f"vae_enc_s{stage_idx}_r1_shortcut_w"] + r1w['shortcut_b'] = w.get(f"vae_enc_s{stage_idx}_r1_shortcut_b") + hidden = resblock_npu(hidden, r1w, channels[stage_idx], channels[stage_idx + 1]) + + r2w = {k: w[f"vae_enc_s{stage_idx}_r2_{k}"] + for k in ['norm1_w', 'conv1_w', 'conv1_b', + 'norm2_w', 'conv2_w', 'conv2_b']} + hidden = resblock_npu(hidden, r2w, channels[stage_idx + 1], channels[stage_idx + 1]) + + if stage_idx < 3: + rsw = {k: w[f"vae_enc_s{stage_idx}_resample_{k}"] + for k in ['conv_w', 'conv_b', 'time_conv_w', 'time_conv_b']} + hidden = downsample3d_npu(hidden, rsw, channels[stage_idx + 1]) + + mr1w = {k: w[f"vae_enc_mid_r1_{k}"] + for k in ['norm1_w', 'conv1_w', 'conv1_b', + 'norm2_w', 'conv2_w', 'conv2_b']} + hidden = resblock_npu(hidden, mr1w, 1536, 1536) + + maw = {k: w[f"vae_enc_mid_attn_{k}"] + for k in ['norm_w', 'to_qkv_w', 'to_qkv_b', 'proj_w', 'proj_b']} + hidden = attention_block_npu(hidden, maw, 1536) + + mr2w = {k: w[f"vae_enc_mid_r2_{k}"] + for k in ['norm1_w', 'conv1_w', 'conv1_b', + 'norm2_w', 'conv2_w', 'conv2_b']} + hidden = resblock_npu(hidden, mr2w, 1536, 1536) + + hidden = causal_conv3d_npu_fp32(hidden, w["vae_enc_out_w"], w["vae_enc_out_b"], + stride=1, padding=1) + mu, _ = hidden.chunk(2, dim=1) + mu_2d, info = _to_2d(mu) + mu_2d = mu_2d[:_real_rows(info)] + result_2d = _npu_scale_norm_fp32(mu_2d, w["vae_scale_mean"].float().flatten(), + w["vae_scale_inv_std"].float().flatten()) + return _from_2d(result_2d, info) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Weight generation — matches golden reference make_weights (encoder part). +# ══════════════════════════════════════════════════════════════════════════════ + +def make_vae_enc_weights(seed=42): + gen = torch.Generator().manual_seed(seed) + weights = {} + weights["vae_enc_in_w"] = torch.randn(96, 3, 3, 3, 3, generator=gen) * 0.02 + weights["vae_enc_in_b"] = torch.zeros(96) + channels = VAE_CH + for stage in range(4): + prefix = f"vae_enc_s{stage}_r1" + weights[f"{prefix}_norm1_w"] = torch.ones(channels[stage]) + weights[f"{prefix}_conv1_w"] = torch.randn(channels[stage+1], channels[stage], 3, 3, 3, generator=gen) * 0.02 + weights[f"{prefix}_conv1_b"] = torch.zeros(channels[stage+1]) + weights[f"{prefix}_norm2_w"] = torch.ones(channels[stage+1]) + weights[f"{prefix}_conv2_w"] = torch.randn(channels[stage+1], channels[stage+1], 3, 3, 3, generator=gen) * 0.02 + weights[f"{prefix}_conv2_b"] = torch.zeros(channels[stage+1]) + if channels[stage] != channels[stage+1]: + weights[f"{prefix}_shortcut_w"] = torch.randn(channels[stage+1], channels[stage], 1, 1, 1, generator=gen) * 0.02 + weights[f"{prefix}_shortcut_b"] = torch.zeros(channels[stage+1]) + prefix = f"vae_enc_s{stage}_r2" + weights[f"{prefix}_norm1_w"] = torch.ones(channels[stage+1]) + weights[f"{prefix}_conv1_w"] = torch.randn(channels[stage+1], channels[stage+1], 3, 3, 3, generator=gen) * 0.02 + weights[f"{prefix}_conv1_b"] = torch.zeros(channels[stage+1]) + weights[f"{prefix}_norm2_w"] = torch.ones(channels[stage+1]) + weights[f"{prefix}_conv2_w"] = torch.randn(channels[stage+1], channels[stage+1], 3, 3, 3, generator=gen) * 0.02 + weights[f"{prefix}_conv2_b"] = torch.zeros(channels[stage+1]) + if stage < 3: + prefix = f"vae_enc_s{stage}_resample" + weights[f"{prefix}_conv_w"] = torch.randn(channels[stage+1], channels[stage+1], 3, 3, generator=gen) * 0.02 + weights[f"{prefix}_conv_b"] = torch.zeros(channels[stage+1]) + weights[f"{prefix}_time_conv_w"] = torch.randn(channels[stage+1], channels[stage+1], 3, 1, 1, generator=gen) * 0.02 + weights[f"{prefix}_time_conv_b"] = torch.zeros(channels[stage+1]) + + prefix = "vae_enc_mid_r1" + for key in ['norm1_w', 'norm2_w']: + weights[f"{prefix}_{key}"] = torch.ones(1536) + for key in ['conv1_w', 'conv2_w']: + weights[f"{prefix}_{key}"] = torch.randn(1536, 1536, 3, 3, 3, generator=gen) * 0.02 + for key in ['conv1_b', 'conv2_b']: + weights[f"{prefix}_{key}"] = torch.zeros(1536) + prefix = "vae_enc_mid_attn" + weights[f"{prefix}_norm_w"] = torch.ones(1536) + weights[f"{prefix}_to_qkv_w"] = torch.randn(1536*3, 1536, 1, 1, generator=gen) * 0.02 + weights[f"{prefix}_to_qkv_b"] = torch.zeros(1536*3) + weights[f"{prefix}_proj_w"] = torch.randn(1536, 1536, 1, 1, generator=gen) * 0.02 + weights[f"{prefix}_proj_b"] = torch.zeros(1536) + prefix = "vae_enc_mid_r2" + for key in ['norm1_w', 'norm2_w']: + weights[f"{prefix}_{key}"] = torch.ones(1536) + for key in ['conv1_w', 'conv2_w']: + weights[f"{prefix}_{key}"] = torch.randn(1536, 1536, 3, 3, 3, generator=gen) * 0.02 + for key in ['conv1_b', 'conv2_b']: + weights[f"{prefix}_{key}"] = torch.zeros(1536) + + weights["vae_enc_out_w"] = torch.randn(VAE_OUT_CH, 1536, 3, 3, 3, generator=gen) * 0.02 + weights["vae_enc_out_b"] = torch.zeros(VAE_OUT_CH) + + vae_mean = torch.tensor([ + -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, + 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921]) + vae_std = torch.tensor([ + 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, + 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160]) + weights["vae_scale_mean"] = vae_mean.view(VAE_Z_DIM, 1, 1, 1) + weights["vae_scale_inv_std"] = (1.0 / vae_std).view(VAE_Z_DIM, 1, 1, 1) + return weights + + +# ══════════════════════════════════════════════════════════════════════════════ +# Golden reference — uses the EXACT vae_encode from test_golden_fun_control_full.py +# ══════════════════════════════════════════════════════════════════════════════ + +def _golden_vae_encode(video, w): + """Golden reference: imports and runs the SAME code as the test file.""" + _GOLDEN_DIR = '/data/x00952168/pypto3.0/cann-recipes-embodied-ai/world_model/agibot-arm-world-model/infer_with_torch' + if _GOLDEN_DIR not in sys.path: + sys.path.insert(0, _GOLDEN_DIR) + from test_golden_fun_control_full import vae_encode as ref_vae_encode + return ref_vae_encode(video, w) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Test — uses golden.runner._Stage + golden.validation.validate_golden, +# same harness primitives as run_jit in t5_encoder.py. +# (Multi-kernel architecture: individual @pl.jit kernels are JIT-compiled +# on first use during the runtime stage; no standalone compile pass.) +# ══════════════════════════════════════════════════════════════════════════════ + + +if __name__ == "__main__": + from golden.runner import _Stage + from golden.validation import validate_golden, ratio_allclose + + parser = argparse.ArgumentParser(description="VAE Encoder — pypto3.0") + parser.add_argument( + "-p", "--platform", type=str, default="a2a3", + choices=["a2a3", "a2a3sim", "a5", "a5sim"], + ) + parser.add_argument("-d", "--device", type=int, default=0, + help="NPU device id (soldered chip index)") + args = parser.parse_args() + + set_device(args.device) + t_total = time.time() + + with _Stage("generate inputs"): + torch.manual_seed(42) + w = make_vae_enc_weights(seed=42) + video = (torch.randn(1, 3, T_IN, H, W) * 0.1).bfloat16() + + with _Stage("compute golden"): + ref = _golden_vae_encode(video.float(), w) + + with _Stage("runtime"): + npu = vae_encode_npu(video, w) + + with _Stage("validate"): + validate_golden( + outputs={"out": npu}, + golden={"out": ref}, + rtol=0.1, + atol=0.1, + compare_fn={"out": ratio_allclose(atol=0.1, rtol=0.1, max_error_ratio=0.02)}, + ) + + total = time.time() - t_total + print(f"[RUN] PASS ({total:.2f}s)", flush=True)