Lab tổng hợp Qwen-Image-2.1 — chạy toàn trình 1A → 6B bằng NumPy

1Định nghĩa

Bài học này là lab thực hành tổng hợp duy nhất của khóa Qwen-Image-2.1 Architecture, kết nối toàn bộ mười khối kiến trúc 1A đến 6B trên sơ đồ 4 cột thành một vòng lặp khử nhiễu Flow-Matching chạy trực tiếp bằng NumPy trong trình duyệt. Thay vì học rời rạc từng công thức, bạn sẽ theo dõi cùng một chuỗi 11 token ghép (7 token Prefix gồm chữ và ảnh điều kiện cộng 4 token ảnh đích) đi từ bước chuẩn hóa RGBA 64-Channel VAE và mở rộng 4× slot ảnh cho tới lúc giải mã ra 4 kênh R, G, B, Alpha.

Trọng tâm kỹ thuật của bài lab là chứng minh bằng số thật vì sao một mô hình Diffusion Transformer có thể tái sử dụng bộ nhớ đệm giống như LLM: nhờ điều biến nhân quả giữ Prefix ở t = 0 kết hợp với mặt nạ Block-Causal 11×11 chặn ảnh đích rò rỉ ngược về phía trước, toàn bộ Key và Value của 7 token Prefix được trích xuất ở bước đầu tiên trong Prefix KV Cache và dùng lại nguyên vẹn ở các bước Euler tiếp theo với sai lệch chỉ ở mức 2.22e-16.

Ví dụ đời thường: Quy trình của lab giống như một xưởng vẽ tranh minh họa nhiều lớp trên giấy kính trong suốt: ở bước đầu tiên (Step 1 extract), họa sĩ đọc kỹ đề bài và bức ảnh mẫu rồi ghim cố định tấm phác thảo tham chiếu lên bảng sáng (Prefix KV Cache tại t = 0); trong suốt các lượt tô màu và tinh chỉnh chi tiết tiếp theo (Steps 2..3 cached), họa sĩ chỉ việc vẽ đè lên tờ giấy kính của bức tranh đích mà không bao giờ phải chép lại đề bài hay vẽ lại tấm ảnh mẫu nữa. Phép so sánh này dừng lại ở chỗ mọi bước trong mô hình đều là phép nhân ma trận tất định trên số dấu phẩy động chứ không có người quan sát bằng mắt.

2Luồng

Cách đọc hình 4 cột: bắt đầu ở 1A dưới đáy Cột 1 đi lên 1B → 2A; sang Cột 2 đi từ 3A → 3B → 3C để cấp tọa độ 3D RoPE, hệ số điều biến t=0 và mặt nạ Block-Causal cho 4A ở nửa trên Cột 3 (nối hai chiều Extract/Cached với trụ bộ nhớ Prefix KV Cache 6A ở nửa dưới Cột 3); sau đó sang Cột 4 từ mạng SwiGLU 5A ở dưới đi lên đầu ra 6B (Scale-Only AdaLN + Euler → RGBA). Mọi số liệu bên dưới lấy nguyên văn từ qwen-image-21.py (seed 21).

Sơ đồ luồng dữ liệu 4 cột · 10 thành phần (1A → 6B)

Kiến trúc Qwen-Image-2.1 & Ánh xạ 1:1 với Bài Lab tổng hợp

Bắt đầu từ 1A ở góc dưới Cột 1 đi lên 1B → 2A, sang Cột 2 (3A → 3B → 3C), qua Cột 3 (4A ↔ 6A) và kết thúc ở Cột 4 (5A → 6B). Di chuột hoặc bấm vào từng khối để xem pop-over và đối chiếu với Lab.

Tất cả (10)1. RGBA VAE, VLM & Packing2. 3D RoPE, t=0 & Block Mask3. Single-Stream & KV Cache4. SwiGLU & 8-Step Euler

Ẩn khung overlay

1A1B2A3A3B3C4A6A5A6B

1A1 · RGBA VAE, VLM & Joint Packing

RGBA [B, 4, H, W] → AutoencoderKLQwenImage21 → ẑ = (z − mean)/std ∈ [B, N_img, 64] (patch_size=1) → img_in(64 → 4096)

1A · RGBA 64-Ch VAE & Unpatched Latents

Biến ảnh thành token mà Transformer đọc được: VAE nén ảnh RGBA thành một lưới latent nhỏ, mỗi ô có 64 kênh, rồi mỗi ô được chiếu thành một token 4096 chiều.

Mục đích

Transformer không thể xử lý trực tiếp hàng triệu điểm ảnh. Bước này nén ảnh xuống một lưới gọn nhưng vẫn giữ đủ thông tin, kể cả kênh trong suốt Alpha, để cuối pipeline (6B) giải nén lại thành ảnh.

Đầu vào

Ảnh tham chiếu 4 kênh R, G, B, Alpha có dạng [B, 4, H, W]. Ảnh đích không đi qua encoder: nó bắt đầu là nhiễu ngẫu nhiên tạo thẳng trên lưới latent cùng kích thước.

Đầu ra

Mỗi ảnh thành N_img token 4096 chiều (1 ô latent = 1 token vì patch_size = 1), sau khi chuẩn hóa từng kênh (z − mean) / std và chiếu img_in 64 → 4096. Các token này được lắp vào chuỗi chung ở 2A.

Bắt đầu đọc sơ đồ từ góc dưới bên trái (1A): AutoencoderKLQwenImage21 nhận ảnh 4 kênh R, G, B, Alpha và nén thành các vector latent 64 kênh. Trước khi đi vào mạng DiT, mỗi kênh được chuẩn hóa theo thống kê riêng ẑ = (z − latents_mean) / latents_std rồi chiếu thẳng qua img_in (64 → 4096) với patch_size = 1 (không gộp ô 2×2 tại đây) để bảo toàn chi tiết biên trong suốt.

Bài 1A · RGBA 64-Channel VAE & Unpatched Latents↗Xem bước [1A] trong Bài Lab tổng hợp↗

flowchart TD
  step1["1A–1B · Chuẩn hóa RGBA Latents #amp; Ngữ cảnh Qwen3-VL"]
  step2["2A · Mở rộng 4× Slot ảnh #amp; Đóng gói chuỗi 11 token"]
  step3["3A–3C · 3D RoPE tâm không, Điều biến t=0 #amp; Mask 11×11"]
  step4["4A · Single-Stream DiT Attention #amp; QK-RMSNorm"]
  step5["5A · Modulated SwiGLU FFN (3× Ratio) #amp; Cổng Tanh"]
  step6["6A–6B · Prefix KV Cache #amp; Vòng lặp Euler → RGBA"]
  step1 --> step2
  step2 --> step3
  step3 --> step4
  step4 --> step5
  step5 --> step6

3Giải pháp

Mục tiêu thực hành: sau khi đọc 10 bài giải phẫu ở Tab 01, bài lab này giúp bạn tự tay kiểm chứng toàn bộ dây chuyền số học của Qwen-Image-2.1 trên một cấu hình thu nhỏ tất định (seed 21, DIM = 8, 2 lớp, 2 head × 4 chiều, lưới ảnh 2×2 × 4 kênh RGBA). Không cần GPU hay thư viện ngoài: mỗi khối in ra bảng số trung gian để bạn đối chiếu trực tiếp với công thức trên sơ đồ.

Ở Cột 1 (1A → 1B → 2A), script chuẩn hóa 4 pixel ảnh điều kiện theo từng kênh z = (z − mean) / std rồi chiếu qua img_in (PATCH_SIZE = 1, giữ nguyên 4 vị trí). Chuỗi ngữ cảnh VLM gồm 4 slot [text, text, COND_IMG, text] đi qua Zero-Center RMSNorm nhân với (w + 1); sau đó bước 2A dùng repeats = [1, 1, 4, 1] bung slot ảnh thành 4 token và nối thêm 4 token nhiễu của ảnh đích, tạo thành chuỗi ghép S_joint = 11 token với image_ids = [-1, -1, 0, 0, 0, 0, -1, 1, 1, 1, 1].

Ở Cột 2 (3A → 3B → 3C), bạn sẽ thấy rõ hai thiết kế giữ cho Prefix bất biến theo thời gian: tại 3A, 3 token chữ tăng dần tọa độ (0,0,0) → (1,1,1) → (2,2,2) trong khi cả hai khối ảnh 2×2 đều ghim frame = 0 và dùng lưới tâm không [(-1,-1), (-1,0), (0,-1), (0,0)]; tại 3B, bảng điều biến có 2 hàng [t_step, 0.0] nên 7 token Prefix luôn đọc đúng hàng t = 0.0 (lệch 0.00e+00 giữa Step 1 ở t=1000 và Step 2 ở t=850); và tại 3C, ma trận mặt nạ 11×11 cho phép khối ảnh 2..5 nhìn nhau song hướng (6 khóa hợp lệ mỗi hàng) nhưng chặn tuyệt đối không cho đọc 4 token ảnh đích ở 7..10.

Ở Cột 3 và Cột 4 (4A → 5A → 6A → 6B), Step 1 (sigma = 1.0, mode = extract) chạy cả 11 token qua 2 lớp Single-Stream DiT (QK-RMSNorm, 3D RoPE, cổng tanh(gate1), SwiGLU 8 → 24 → 8, cổng tanh(gate2)) và lưu cached_k, cached_v có shape (7, 2, 4). Sang Step 2 (sigma = 0.85) và Step 3 (sigma = 0.45) ở chế độ cached, mạng chỉ chiếu Q, K, V cho đúng 4 token ảnh đích (shape 4, 2, 4) rồi nối với cache (7, 2, 4) thành (11, 2, 4). Cuối mỗi bước, Scale-Only Final AdaLN nhân (1 + scale_final) không có shift, chiếu ra vận tốc v_theta để cập nhật Euler x_{k+1} = x_k + (sigma_{k+1} − sigma_k) · v_theta trước khi giải mã về ảnh RGBA 2×2×4.

Giới hạn của kết luận: ba phép kiểm chứng cuối bài xác nhận tính tương đương số học giữa đường dùng Prefix KV Cache và đường tính lại toàn chuỗi (sai lệch 2.22e-16 đến 4.44e-16 < 1e-9), chứng minh nếu tắt t = 0 trên Prefix thì cache bị trôi ngay 9.61e-01 đến 1.68e+00, và xác nhận tính cô lập một chiều từ Prefix sang Target. Tuy nhiên, vì trọng số được khởi tạo ngẫu nhiên với seed 21 và chưa qua huấn luyện, bức ảnh RGBA 2×2 đầu ra chỉ dùng để kiểm tra đường đi của tensor chứ không mang nội dung thị giác thật hay số đo tốc độ GPU.

4Profiling

• Kiểm tra tính đúng đắn của Prefix KV Cache trước khi bàn về tốc độ

  Luôn chạy đối chiếu hai đường trên cùng hạt giống ngẫu nhiên (seed 21): đường 1 dùng kv_cache_mode = extract ở bước đầu và cached ở các bước sau; đường 2 tính lại toàn bộ 11 token ở mọi bước Euler. Sai lệch tuyệt đối tối đa trên trường vận tốc v_theta (2.22e-16 ở bước 2, 4.44e-16 ở bước 3) và trên ảnh RGBA cuối (1.11e-16) phải nằm dưới ngưỡng 1e-9.

• Kiểm tra điều kiện bất biến thời gian (t = 0) trên Prefix

  Thử nghiệm đối chứng bằng cách đặt causal_condition = False (tức bắt 7 token Prefix đọc hàng t_step thay vì hàng t = 0.0): đo độ lệch giữa cached_k ở Step 1 (t = 1000) và K của Prefix tính lại ở Step 2 (t = 850). Độ lệch vọt lên 9.61e-01 ở lớp 0 và 1.68e+00 ở lớp 1 chứng minh Prefix KV Cache chỉ hợp lệ khi Prefix được điều biến cố định ở t = 0.

• Tỷ lệ tiết kiệm số token phải tính ở các bước Cached

  Đếm số lượng token phải đi qua phép chiếu Q/K/V và mạng SwiGLU FFN ở mỗi bước: ở Step 1 (extract) là S_joint = prefix_len + target_len (11 token trong lab); ở mỗi bước Steps 2..N (cached) chỉ còn đúng target_len (4 token trong lab, giảm 7/11 ≈ 63.6% số token phải tính mỗi bước). Khi có tới 10 ảnh tham chiếu trên thực tế, tỷ lệ prefix_len / S_joint càng lớn thì phần tính toán tiết kiệm được ở các bước 2..N càng cao.

Mẹo đo lường: ghim số liệu trên workload thật, ghi cả p50/p95/p99 thay vì chỉ trung bình, và luôn ghi lại điều kiện chạy (batch, cache, đồng thời, phần cứng) trước và sau khi tối ưu. Số đo thiếu điều kiện chạy thì không so sánh được.

5Ví dụ

6Thuật ngữ

Chuỗi ghép S_joint = 11 (Prefix 7 + Target 4)

Gồm 2 token chữ đầu, 4 token của ảnh điều kiện 2×2 (bung 4× từ 1 slot VLM), 1 token chữ cuối (tổng 7 token Prefix) và 4 token của ảnh đích 2×2 đang khử nhiễu.

image_ids = [-1, -1, 0, 0, 0, 0, -1, 1, 1, 1, 1]

Nhãn phân biệt phương thức trên chuỗi ghép: -1 là văn bản (chỉ nhìn nhân quả), 0 là ảnh điều kiện thứ nhất, 1 là ảnh đích (các token cùng mã >= 0 nhìn nhau song hướng).

Bảng điều biến hai hàng [t_step, 0.0]

Khi bật causal_condition, hàng 0 mang hệ số của bước nhiễu hiện tại t_step cho 4 token ảnh đích, còn hàng 1 giữ cố định mốc t = 0.0 cho 7 token Prefix.

kv_cache_mode: extract vs cached

Ở bước Euler đầu tiên (extract), mạng chạy đủ 11 token và lưu bản sao K, V của 7 token Prefix; ở các bước sau (cached), mạng chỉ chạy 4 token ảnh đích và nối K, V đã lưu.

7Bài liên quan

8Nguồn

9Bài lab