Bài Lab tổng hợp — DeepSeek-V4.1-Flash Architecture

Giới thiệu bài lab

Lab DeepSeek-V4.1-Flash: chạy thử một lượt suy luận thu nhỏ từ 1A đến 6B bằng NumPy

Bài lab là một script Python dùng NumPy (thư viện tính toán trên mảng số) chạy ngay trong trình duyệt. Nó dựng lại cả 10 khối 1A đến 6B của sơ đồ kiến trúc ở kích thước thu nhỏ: vector 8 chiều thay vì 5120, 5 lớp thay vì 40. Rồi nó cho một prompt nhỏ [1, 2, 3, 4] chạy qua, trong đó vị trí 1 là một token ảnh, và in từng bảng số.

Đọc mô tả thì dễ gật đầu mà khó biết mình đã hiểu chưa. Khi thấy số thật, ví dụ tổng mỗi hàng và mỗi cột của ma trận Sinkhorn (phép chuẩn hóa lặp để ma trận 4×4 có tổng hàng và cột đều bằng 1) đúng 1.000000, hay sáu khoản bộ nhớ cộng ra 890 byte mỗi token, bạn tự kiểm chứng được cơ chế trong bài kiến trúc hoạt động thế nào.

Bài lab gồm

Sơ đồ 4 cột có khung tương tác ở đầu trang, để biết mỗi bảng số thuộc khối nào.
Bảng ánh xạ 1:1 cho từng mã 1A đến 6B: lab dựng gì và nhìn số nào để kiểm chứng.
Trình chạy trực tiếp với mã deepseek-v41-flash.py sửa được, và các nút Chạy cell, Dừng, Đặt lại.
Output khoảng 180 dòng chia thành các khối [1A] đến [6B], mỗi khối có lời giải thích, bảng số, Quan sát và Kết luận, rồi mục [KIỂM TRA & THỬ NGHIỆM] với ngưỡng sai số in rõ.
Trang bài đọc Lab đầy đủ (liên kết ở thanh công cụ) có hướng dẫn chi tiết và 4 bài tập kèm đáp án.

Làm xong, bạn có thể

Giải thích vì sao vector tại vị trí ảnh bị thay hẳn bằng vector gộp 3×3 (phủ đè, không cộng) và vì sao Engram bỏ qua cặp token chạm ảnh.
Giải thích vì sao KV chỉ nằm ở lớp nguồn (lớp 1 và 3 trong lab) mà lớp 2 và 4 vẫn attention được.
Đọc nhật ký của Compressor: prefill 4 token phát 2 latent, decode token lẻ trả None cho tới khi đủ cặp.
Tự cộng 256 + 64 + 480 + 32 + 10 + 48 ra 890 byte mỗi token.
Nói được lab chứng minh gì (cơ chế tính đúng ở cỡ nhỏ) và không chứng minh gì (chất lượng ngôn ngữ, tốc độ GPU).

Phạm vi: lab chỉ chạy suy luận (forward) của một mô hình thu nhỏ và không có bước huấn luyện; trọng số là số ngẫu nhiên cố định bằng seed 41, nên chạy lại cho ra đúng các số cũ, nhưng token đoán ra và điểm confidence không mang nghĩa ngôn ngữ. Muốn học huấn luyện, hãy xem khóa LLM Training (Text), vốn chỉ dạy mô hình văn bản; huấn luyện mô hình đa phương thức chưa có bài nào trên site.

Cách dùng bài lab

Xem sơ đồ 4 cột ở đầu trang: rê chuột hoặc bấm vào khung để mở pop-over, các nút lọc theo cột làm mờ khung của những cột còn lại, và nút Ẩn khung overlay cho bạn đọc ảnh gốc.
Đọc bảng ánh xạ bên dưới: cột thứ ba cho biết lab dựng gì ở mỗi mã, cột thứ tư cho biết nhìn số nào để kiểm chứng.
Đọc cấu hình thu nhỏ ở đầu mã và đầu output: DEMO_PROMPT = [1, 2, 3, 4] là 4 mã token (mỗi mã chỉ định danh một mục trong bảng từ vựng, không tự mang nghĩa), IMAGE_POS = [1] đánh dấu token ảnh, hai bước decode thêm token 5 rồi 2.
Với mỗi khối, dự đoán trước rồi bấm Chạy cell. Đọc theo thứ tự prefill (xử lý cả 4 token một lượt), decode (thêm từng token mới và dùng lại cache), rồi [KIỂM TRA & THỬ NGHIỆM].
Đọc từng khối từ [1A] đến [6B] và đối chiếu số in ra với cột Cách đọc số liệu của bảng ánh xạ; đừng chỉ nhìn chữ PASS mà bỏ qua bảng số.
Sửa một hằng số trong ô mã rồi chạy lại, ví dụ đổi HC_SINKHORN_ITERS từ 20 xuống 1 và xem sai lệch của ma trận Sinkhorn vượt ngưỡng 1e-6; nút Đặt lại khôi phục mã gốc.

Đọc kết quả đúng cách: PASS chỉ cho biết phép kiểm chứng cụ thể đó đạt ở cỡ nhỏ. Với trọng số chưa huấn luyện, argmax (token có điểm cao nhất) và confidence không phải dự đoán có nghĩa, đừng đọc chúng như vậy.

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

Kiến trúc DeepSeek-V4.1-Flash & Á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.

Sơ đồ luồng dữ liệu 4 cột của kiến trúc DeepSeek-V4.1-Flash gồm 10 thành phần từ 1A đến 6B

1A1 · Input, Engram & 4× mHC Stream

E[ids]: [B, T, 5120] · Image Patches 3×3 → DeepSeek-ViT (2D-RoPE) + Aligner → [B, N_img, 5120]

1A · Multimodal Input & Aligner 3×3

Điểm xuất phát ở đáy Cột 1: nhúng token chữ E[ids] và gộp mỗi ô 3×3 patch ảnh từ DeepSeek-ViT (2D-RoPE) phủ đè vào vị trí ảnh.

Bắt đầu đọc sơ đồ từ góc dưới bên trái (1A): các token văn bản được tra bảng E[ids] thành khối vector [B, T, 5120] (ở lab là [4, 8]). Nếu đầu vào có ảnh, DeepSeek-ViT (32 lớp, xoay vị trí 2D-RoPE) mã hóa các patch rồi Aligner gộp mỗi cụm 3×3 patch thành 1 vector (giảm 9 lần số token ảnh) và ghi đè trực tiếp lên span <｜deepseek_image｜> trước khi đẩy luồng ẩn đi lên 1B.

Bài 1A · Đầu vào đa phương thức & Aligner 3×3 Xem bước [1A] trong Bài Lab tổng hợp

BẢNG ÁNH XẠ 1:1 · MỖI KHỐI KIẾN TRÚC Ở TAB 01 ĐƯỢC DỰNG VÀ KIỂM CHỨNG TRONG LAB NHƯ THẾ NÀO?

Mã	Thành phần ở Tab 01 (bấm để mở bài)	Lab NumPy dựng gì (`deepseek-v41-flash.py`)	Cách đọc số liệu & phép kiểm chứng
1A	Đầu vào Văn bản & Hình ảnh (ViT + Aligner 3×3) ↗	Tra bảng E[ids] ra vector (4, 8) và gộp ô 3×3 patch ảnh phủ đè trực tiếp lên vị trí ảnh (IMAGE_POS = [1]).	Đối chiếu vector tại vị trí 1 đã bị thay hoàn toàn bởi vector gộp 3×3 (phủ đè, không cộng).
1B	Bộ nhớ có điều kiện Engram (N-gram XOR-Prime) ↗	Đánh dấu vị trí ảnh DEAD = -1, băm cặp 2-gram qua số nguyên tố tại lớp 1 và nhân cổng σ(signed_sqrt(dot)).	Cặp (3, 4) rơi vào bucket 9 và cộng vào hc_pre của lớp 1; mọi cặp chạm vị trí ảnh 1 đều bị DEAD chặn.
2A	Single-Pass mHC (4 làn Residual & Sinkhorn 4×4) ↗	Mở rộng trạng thái thành 4 làn (4, 4, 8), gộp hc_pre (4, 8) trước lớp con và trộn hc_post bằng ma trận Sinkhorn 4×4.	Kiểm chứng (1): sau 20 vòng lặp Sinkhorn, tổng mỗi hàng và mỗi cột của comb bằng 1.000000 (sai lệch 2.22e-16 < 1e-6).
3A	Causal Encoder–Decoder (CED: L0..2 → L3..4) ↗	5 lớp thu nhỏ với COMPRESS_RATIOS = (0, 2, 2, 1, 1); L1 giữ KV nén của Encoder, L3 giữ Global KV dùng chung cho Decoder L4.	Kiểm chứng (3): L2 đọc nguyên mảng KV của L1 và L4 đọc nguyên mảng Global KV của L3 mà không tính lại.
3B	CSA2 Compressor (Gộp r=2 & 3 chế độ Full/Reindex/Reuse) ↗	Gộp mỗi cặp r = 2 token bằng softmax(wgate(x)); bước decode lẻ giữ tạm trong kv_state và trả None cho đến khi đủ cặp.	Prefill 4 token phát 2 latent (end 1, 3); decode bước 1 (pos 4 lẻ) trả None, bước 2 (pos 5 chẵn) phát latent thứ 3 (end 5).
3C	Ring SWA (2 slot) & Hierarchical Sparse Indexer ↗	Bộ đệm vòng window_kv_cache giữ 2 token gần nhất kết hợp Indexer chấm điểm relu(q·k)/√d × w chọn index_topk = 1 vị trí nén.	Các lớp không phải nguồn tái sử dụng trực tiếp chỉ mục top-k do lớp nguồn (L1, L3) phát hành.
4A	Multi-Head Latent Attention (MLA, Inverse RoPE & W_O nhóm) ↗	Q hạ chiều (8 → 4 → 2×4), KV 1 head (1×4) chỉ xoay RoPE 2 chiều cuối, thêm cột attn_sink, xoay ngược RoPE ở đầu ra và chiếu W_O 2 nhóm.	Kiểm chứng (2) & (4): giữ nguyên tính nhân quả khi đổi token cuối và decode trên cache khớp tính lại toàn chuỗi (0.00e+00 < 1e-9).
5A	Fine-Grained MoE (Cổng sqrtsoftplus, tách bias & SwiGLU Clamp) ↗	Điểm thô gate = √softplus(scores); chọn top-2/4 chuyên gia theo scores + bias (có bias_vl ở vị trí ảnh), nhân theo gate thô và kẹp SwiGLU [-10, 10].	Đổi bias_vl làm đổi chỉ số chuyên gia được chọn tại vị trí ảnh nhưng công thức tính trọng số nhân vẫn lấy từ gate thô.
6A	Ngân sách KV 890 bytes/token (FP4 E2M1 & SWA Bounded Replay) ↗	Tính đơn giá bộ nhớ tại 4 lớp nguồn [2, 8, 14, 20] với tỷ lệ nén (2, 2, 2, 1), FP4 E2M1 + scale E4M3/16 cho KV và E8M0/32 cho Indexer K.	Cộng chính xác 6 hạng mục đơn giá (256 + 64 + 480 + 32 + 10 + 48) ra đúng 890 B/token và đối chiếu tỷ lệ giảm với GQA/MHA.
6B	Đầu ra LM Head & DSpark Speculative Decoding ↗	Gộp 4 làn mHC ra logits (4, 8), nối trung bình lớp đích (3, 4) vào DSpark kết hợp MarkovHead (hạng 4) và ConfidenceHead vẽ 3 token nháp.	Quan sát chuỗi 3 token nháp được sinh bán tự hồi quy kèm điểm confidence của từng vị trí.

Trình chạy Lab trực tiếp · `deepseek-v41-flash.py`

deepseek-v41-flash.pyPYTHON

"""Lab DeepSeek-V4.1-Flash: chạy toàn trình 1A→6B bằng NumPy thuần.

Bài lab giải phẫu MỘT lượt forward đầy đủ của kiến trúc DeepSeek-V4.1-Flash
ở kích thước thu nhỏ, đồng bộ mã khối với sơ đồ kiến trúc Tab 1 (1A..6B):

[1A] Đầu vào đa phương thức (Embedding + Aligner ảnh 3×3)
  [1B] Bộ nhớ có điều kiện Engram (n-gram băm XOR-prime + cổng sigmoid
       căn bậc hai có dấu; vị trí ảnh bị che thành DEAD = -1)
  [2A] Single-pass mHC: 4 làn residual + ma trận Sinkhorn 4×4
  [3A] Causal Encoder-Decoder (CED): lớp 0..2 Encoder, lớp 3..4 Decoder
  [3B] CSA2: Compressor gộp r token → 1 latent; chia sẻ KV giữa các lớp
  [3C] Ring SWA + Hierarchical Sparse Indexer
  [4A] MLA: Q low-rank, K/V một head, RoPE nửa cuối, attn_sink,
       inverse RoPE và W_O nhóm (Group-LoRA)
  [5A] MoE fine-grained: cổng sqrt(softplus), tách bias chọn chuyên gia
       khỏi trọng số nhân, SwiGLU có clamp, cộng shared expert
  [6A] Ngân sách KV: bảng FP4 E2M1 (890 B/token) so với GQA/MHA chuẩn
  [6B] LM head + DSpark speculative decoding (MarkovHead có bias)

Tiến trình thực thi: PREFILL xử lý DEMO_PROMPT = [1, 2, 3, 4] (vị trí 1 là
một token ảnh ghép ACL 3×3, các vị trí 0/2/3 là token chữ), sau đó 2 bước
DECODE_TOKENS = [5, 2] tại start_pos = 4 và 5, rồi [KIỂM TRA & THỬ NGHIỆM].

Mọi số in ra đều là giá trị TÍNH THẬT của pass đang chạy (không hardcode).
Weights chưa huấn luyện, sinh một lần bằng np.random.default_rng(41);
chạy lại cho ra cùng số. Lab mô phỏng cơ chế trung thực ở vùng số nhỏ;
KHÔNG tái tạo checkpoint thật, KHÔNG đo tốc độ GPU, không chứng minh chất
lượng ngôn ngữ.
"""
import numpy as np

# ── Cấu hình thu nhỏ (sửa các số này để thí nghiệm tiếp) ────────────────────
VOCAB = 8                # bảng từ vựng demo
DIM = 8                  # chiều vector ẩn
HC_MULT = 4              # số làn residual của mHC
HC_SINKHORN_ITERS = 20   # số vòng chuẩn hóa ma trận trộn Sinkhorn
N_LAYERS = 5             # L0..L4
# r mỗi lớp: L0 SWA thuần (r=0); L1..L2 Causal Encoder r=2; L3..L4 Decoder r=1
COMPRESS_RATIOS = (0, 2, 2, 1, 1)
KV_SOURCE_LAYERS = (1, 3)    # L1 giữ Encoder KV nén; L3 giữ Decoder Global KV
INDEX_SOURCE_LAYERS = (1, 3) # nguồn của Indexer phân cấp (L2/L4 dùng lại)
WINDOW_SIZE = 2              # số slot ring buffer của SWA
INDEX_TOPK = 1               # số latent Indexer chọn mỗi head
N_HEADS = 2
HEAD_DIM = 4
ROPE_HEAD_DIM = 2        # chỉ nửa cuối (2 chiều cuối) được xoay RoPE
Q_LORA_RANK = 4          # Q low-rank: 8 -> 4 -> 2×4
O_GROUPS = 2             # W_O nhóm: 2 group, mỗi group LoRA hạng 3
O_LORA_RANK = 3
N_ROUTED_EXPERTS = 4
N_ACTIVATED_EXPERTS = 2
MOE_INTER_DIM = 6
SWIGLU_LIMIT = 10.0
ENGRAM_LAYERS = (1,)     # Engram chỉ chạy ở L1
ENGRAM_BUCKETS = 16
ENGRAM_PRIME = 31
DSPARK_BLOCK_SIZE = 3    # 3 token nháp mỗi khối
DSPARK_TARGET_LAYERS = (3, 4)
DSPARK_MARKOV_RANK = 4
DSPARK_MARKOV_ALPHA = 0.8
ROPE_BASE = 10000.0
IMAGE_POS = (1,)         # vị trí 1 của prompt là token ảnh
DEAD = -1                # mã "chết" trên bản đồ token nén
DEMO_PROMPT = [1, 2, 3, 4]
DECODE_TOKENS = [5, 2]
TRY_PROMPT = [1, 2, 3, 5]

assert N_HEADS * HEAD_DIM == DIM
assert O_GROUPS * (N_HEADS // O_GROUPS) == N_HEADS

# ── Weights (chưa học): sinh MỘT lần, mọi pass dùng chung ───────────────────
rng = np.random.default_rng(41)
def _n(shape, s):
    return rng.normal(0.0, s, shape)

E = _n((VOCAB, DIM), 0.35)                        # bảng nhúng token chữ
W_ALIGN = _n((9, DIM), 0.30)                      # Aligner: 3×3=9 -> DIM
B_ALIGN = _n((DIM,), 0.20)
IMG_PATCH = (np.arange(9).reshape(3, 3) / 9.0) * 0.8 + 0.1  # ảnh demo cố định

ENG_TABLE = _n((ENGRAM_BUCKETS, DIM), 0.40)       # bộ nhớ n-gram

W_HC_UP = _n((HC_MULT, DIM), 0.30)                # mở rộng thành 4 làn
W_HC_DOWN = _n((HC_MULT, DIM), 0.30)              # gộp 4 làn về DIM
MIX_LOGITS = _n((HC_MULT, HC_MULT), 0.50)         # hạt giống ma trận trộn

W_QA = _n((DIM, Q_LORA_RANK), 0.30)               # Q low-rank phần a
W_QB = _n((Q_LORA_RANK, N_HEADS * HEAD_DIM), 0.30)
W_KV = _n((DIM, 2 * HEAD_DIM), 0.30)              # K/V MỘT head (MLA)
SINK = _n((N_LAYERS, N_HEADS), 0.50)              # attention sink mỗi lớp/head
WO_A = _n((O_GROUPS, HEAD_DIM, O_LORA_RANK), 0.30)
WO_B = _n((O_GROUPS * O_LORA_RANK, DIM), 0.30)

W_ROUTER = _n((DIM, N_ROUTED_EXPERTS), 0.60)      # điểm router
BIAS_TEXT = _n((N_ROUTED_EXPERTS,), 0.10)         # bias chọn chuyên gia (chữ)
BIAS_VL = _n((N_ROUTED_EXPERTS,), 0.40)           # bias cộng thêm tại token ảnh
EXP_GATE = _n((N_ROUTED_EXPERTS, DIM, MOE_INTER_DIM), 0.30)
EXP_UP = _n((N_ROUTED_EXPERTS, DIM, MOE_INTER_DIM), 0.30)
EXP_DOWN = _n((N_ROUTED_EXPERTS, MOE_INTER_DIM, DIM), 0.25)
SH_GATE = _n((DIM, MOE_INTER_DIM), 0.30)
SH_UP = _n((DIM, MOE_INTER_DIM), 0.30)
SH_DOWN = _n((MOE_INTER_DIM, DIM), 0.25)

W_LM = _n((DIM, VOCAB), 0.30)                     # LM head cho mô hình lớn
W_DRAFT = _n((DIM, VOCAB), 0.30)                  # LM head nháp của DSpark
MK_A = _n((VOCAB, DSPARK_MARKOV_RANK), 0.50)      # MarkovHead hạng 4
MK_B = _n((DSPARK_MARKOV_RANK, VOCAB), 0.50)
INDEX_W = _n((64,), 0.60)                         # trọng số indexer gán theo slot

# ── Hàm số học dùng chung ───────────────────────────────────────────────────
def rmsnorm(x):
    """Pre-Norm không tham số: giữ hướng, chuẩn hóa biên độ vector."""
    return x / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + 1e-6)

def silu(x):
    return x / (1.0 + np.exp(-x))

def softplus(x):
    return np.logaddexp(0.0, x)

def signed_sqrt(x):
    return np.sign(x) * np.sqrt(np.abs(x))

def softmax(v):
    v = np.asarray(v, dtype=float)
    e = np.exp(v - v.max())
    return e / e.sum()

def sinkhorn(logits):
    """Đưa ma trận vuông về doubly stochastic bằng vòng lặp hàng–cột."""
    P = np.exp(logits - logits.max())
    for _ in range(HC_SINKHORN_ITERS):
        P = P / P.sum(axis=1, keepdims=True)
        P = P / P.sum(axis=0, keepdims=True)
    return P

MIX_COMB = sinkhorn(MIX_LOGITS)

ROPE_FREQ = ROPE_BASE ** (-2.0 * np.arange(ROPE_HEAD_DIM // 2) / ROPE_HEAD_DIM)

def rope_rotate(x, pos, inverse=False):
    """Xoay NỬA CUỐI (ROPE_HEAD_DIM chiều cuối) của x theo vị trí pos.

x: shape (..., HEAD_DIM). Chiều 0..1 giữ nguyên (nửa "không xoay"),
    chiều 2..3 (cặp cuối) xoay góc pos. inverse=True xoay ngược dấu —
    dùng để đưa đầu ra attention về không gian không vị trí (inverse RoPE).
    """
    sign = -1.0 if inverse else 1.0
    ang = np.asarray(pos, dtype=float) * ROPE_FREQ * sign
    cos, sin = np.cos(ang), np.sin(ang)
    rest, rot = x[..., :-ROPE_HEAD_DIM], x[..., -ROPE_HEAD_DIM:]
    x1, x2 = rot[..., : ROPE_HEAD_DIM // 2], rot[..., ROPE_HEAD_DIM // 2:]
    r1 = x1 * cos - x2 * sin
    r2 = x1 * sin + x2 * cos
    return np.concatenate([rest, r1, r2], axis=-1)

def comp_token(tok, pos):
    """Bản đồ token nén cho băm n-gram: token ảnh bị che thành DEAD = -1."""
    return DEAD if pos in IMAGE_POS else int(tok)

def engram_bucket(prev_c, cur_c):
    """Băm 2-gram XOR-prime: (prev × 31) ^ cur, chia dư số bucket."""
    return ((int(prev_c) * ENGRAM_PRIME) ^ int(cur_c)) % ENGRAM_BUCKETS

def engram_contribution(bucket):
    """Cộng gộp có cổng: sigmoid theo chiều × căn bậc hai có dấu."""
    tbl = ENG_TABLE[bucket]
    gate = 1.0 / (1.0 + np.exp(-tbl))
    return gate * signed_sqrt(tbl), gate

# ── Trạng thái luồng (streaming) ────────────────────────────────────────────
CACHE_OWNERS = (0,) + tuple(KV_SOURCE_LAYERS)   # lớp tự giữ cache riêng

def owner_of(layer):
    """L1 là nguồn KV cho L2; L3 là nguồn Global KV cho L4."""
    if layer in CACHE_OWNERS:
        return layer
    return 1 if layer == 2 else 3

def new_state():
    caches = {l: {"latents": [], "kv_state": [], "ring": []} for l in CACHE_OWNERS}
    return {"caches": caches, "S": None, "ids": [], "last_y": {}}

def attention_block(layer, pos, h_norm, state, trace):
    """MLA + chọn nguồn ngữ cảnh (ring / latent qua Indexer) cho MỘT token.

Trả về (attn_out CHƯA cộng residual (8,), chi tiết chọn nguồn để in).
    Truy cập ngữ cảnh: các latent nén đã phát (đủ điều kiện end <= pos)
    được Indexer lựa chọn per head; thêm ring SWA và token hiện tại.
    """
    qa = h_norm @ W_QA
    q_raw = (qa @ W_QB).reshape(N_HEADS, HEAD_DIM)
    q_rot = np.stack([rope_rotate(q_raw[h], pos) for h in range(N_HEADS)])
    kv = W_KV @ h_norm
    k_raw, v_new = kv[:HEAD_DIM], kv[HEAD_DIM:]
    k_rot = rope_rotate(k_raw, pos)
    cache = state["caches"][owner_of(layer)]
    ring = [e for e in cache["ring"] if e["pos"] < pos]

# --- Indexer phân cấp: relu(q·k) × trọng số, chọn INDEX_TOPK latent ---
    indexer = []
    per_head_picks = []
    if layer != 0 and cache["latents"]:
        eligible = [(i, lt) for i, lt in enumerate(cache["latents"]) if lt["end"] <= pos]
        for h in range(N_HEADS):
            sc = np.array([max(0.0, float(q_rot[h] @ lt["k"])) / np.sqrt(HEAD_DIM)
                           * lt["w"] for _, lt in eligible])
            order = np.argsort(-sc)[:INDEX_TOPK]
            picks = [(eligible[int(o)][0], eligible[int(o)][1]) for o in order]
            per_head_picks.append(picks)
            indexer.append({
                "eligible_ends": [lt["end"] for _, lt in eligible],
                "scores": sc.tolist(),
                "picked_idx": [i for i, _ in picks],
                "picked_end": [lt["end"] for _, lt in picks],
            })
    else:
        per_head_picks = [[] for _ in range(N_HEADS)]

ctx_heads = []
    head_report = []
    for h in range(N_HEADS):
        pieces = [["latent", lt["k"], lt["v"], lt["end"], lt["w"]]
                  for _, lt in per_head_picks[h]]
        pieces += [["ring", e["k"], e["v"], e["pos"], None] for e in ring]
        pieces.append(["cur", k_rot, v_new, pos, None])
        keys = np.stack([p[1] for p in pieces])
        vals = np.stack([p[2] for p in pieces])
        scores = np.array([float(q_rot[h] @ p[1]) / np.sqrt(HEAD_DIM) for p in pieces])
        w_ext = softmax(np.concatenate([scores, [SINK[layer][h]]]))
        weights, sink_w = w_ext[:-1], float(w_ext[-1])
        ctx_h = weights @ vals
        ctx_heads.append(ctx_h)
        head_report.append({
            "sources": [p[0] for p in pieces],
            "src_pos": [p[3] for p in pieces],
            "scores": scores.tolist(),
            "weights": weights.tolist(),
            "sink_w": sink_w,
        })

ctx_pre = np.stack(ctx_heads)                          # (N_HEADS, HEAD_DIM)
    ctx = np.stack([rope_rotate(ctx_pre[h], pos, inverse=True)  # inverse RoPE
                    for h in range(N_HEADS)])
    ctx_flat = ctx.reshape(N_HEADS * HEAD_DIM)
    ctx_flat = ctx_flat.reshape(O_GROUPS, HEAD_DIM)        # gom head theo group
    grouped = [ctx_flat[g] @ WO_A[g] for g in range(O_GROUPS)]
    attn_out = np.concatenate(grouped) @ WO_B              # (2×3)=6 -> DIM

report = {
        "k_rot": k_rot, "k_raw": k_raw, "v_new": v_new,
        "q_raw": q_raw, "q_rot": q_rot, "sink": SINK[layer],
        "ctx_pre": ctx_pre, "ctx": ctx,
        "ring_len": len(ring), "ring_pos": [e["pos"] for e in ring],
        "latent_n": len(cache["latents"]),
        "latent_end_list": [lt["end"] for lt in cache["latents"]],
        "latent_k0": cache["latents"][0]["k"].copy() if cache["latents"] else None,
        "heads": head_report, "indexer": indexer,
    }
    if trace is not None:
        trace["attn"].append({"layer": layer, "pos": pos, **report})
    return attn_out, report

def moe_block(h2, pos, layer, trace):
    """MoE fine-grained: chọn bằng scores+bias, nhân bằng sqrt(softplus)."""
    r_norm = rmsnorm(h2)
    scores = r_norm @ W_ROUTER
    gate = np.sqrt(softplus(scores))                 # cổng sqrt(softplus)
    bias = BIAS_TEXT + (BIAS_VL if pos in IMAGE_POS else 0.0)
    sel = np.argsort(-(scores + bias))[:N_ACTIVATED_EXPERTS]
    w_raw = gate[sel]
    w_route = w_raw / w_raw.sum()                    # trọng số từ gate THÔ
    moe_out = np.zeros(DIM)
    clamped = 0
    for j, e in enumerate(sel):
        hidden = silu(r_norm @ EXP_GATE[e]) * (r_norm @ EXP_UP[e])
        if np.any(np.abs(hidden) > SWIGLU_LIMIT):
            clamped += 1
        hidden = np.clip(hidden, -SWIGLU_LIMIT, SWIGLU_LIMIT)
        moe_out += w_route[j] * (hidden @ EXP_DOWN[e])
    sh_hidden = silu(r_norm @ SH_GATE) * (r_norm @ SH_UP)
    sh_hidden = np.clip(sh_hidden, -SWIGLU_LIMIT, SWIGLU_LIMIT)
    shared_out = sh_hidden @ SH_DOWN
    y = h2 + moe_out + shared_out
    rep = {"scores": scores, "gate": gate, "bias": bias, "sel": sel,
           "w_route": w_route, "clamped": clamped, "shared_norm": float(np.linalg.norm(shared_out))}
    if trace is not None:
        trace["moe"].append({"layer": layer, "pos": pos, **{k: v for k, v in rep.items()}})
    return y, rep

def emit_latents(cache, r, trace_events):
    """Gộp các cặp KV đang chờ thành latent nén theo tỉ lệ r."""
    while r > 0 and len(cache["kv_state"]) >= r:
        block = cache["kv_state"][:r]
        cache["kv_state"] = cache["kv_state"][r:]
        latent = {
            "k": np.mean([b["k"] for b in block], axis=0),
            "v": np.mean([b["v"] for b in block], axis=0),
            "end": block[-1]["pos"],
            "w": float(INDEX_W[len(cache["latents"])]),
        }
        cache["latents"].append(latent)
        trace_events.append({"emitted": True, "end": latent["end"],
                             "members": [b["pos"] for b in block]})
    if not cache["kv_state"]:
        return
    trace_events.append({"emitted": False,
                         "pending": len(cache["kv_state"]),
                         "pending_pos": [b["pos"] for b in cache["kv_state"]]})

def run_token(tok, pos, state, trace=None):
    """Chạy MỘT token qua cả 5 lớp theo thứ tự streaming; trả về logits."""
    is_image = pos in IMAGE_POS
    x0 = (IMG_PATCH.reshape(3, 3).reshape(9) @ W_ALIGN + B_ALIGN if is_image
          else E[int(tok)])
    if state["S"] is None:
        state["S"] = np.einsum("d,md->md", x0, W_HC_UP)   # mở rộng 4 làn
    S = state["S"]
    if trace is not None:
        trace["x0"].append(("img" if is_image else "text", pos, x0.copy()))

prev_c = None
    if state["ids"]:
        prev_pos = pos - 1
        prev_c = comp_token(state["ids"][-1], prev_pos)
    cur_c = comp_token(tok, pos)
    engram_done = None
    if trace is not None:
        trace["engram_pair"].append({"pos": pos, "prev_c": prev_c, "cur_c": cur_c})

for layer in range(N_LAYERS):
        if layer == 4:
            trace["lanes_l4"].append(S.copy()) if trace is not None else None
        hc_pre = np.einsum("md,md->d", S, W_HC_DOWN)      # gộp 4 làn -> (8,)
        if trace is not None and layer == 4:
            trace["hc_pre_l4"].append(hc_pre.copy())
        if layer in ENGRAM_LAYERS:
            if cur_c != DEAD and prev_c is not None and prev_c != DEAD:
                bucket = engram_bucket(prev_c, cur_c)
                contrib, gate = engram_contribution(bucket)
                hc_pre = hc_pre + contrib
                engram_done = (pos, bucket, contrib, gate)
        h_norm = rmsnorm(hc_pre)
        attn_out, attn_rep = attention_block(layer, pos, h_norm, state, trace)
        if trace is not None and layer == 1 and engram_done:
            trace["engram_applied"].append({
                "pos": pos, "bucket": engram_done[1],
                "gate": engram_done[3], "hc_pre": hc_pre.copy()})
        post = hc_pre + attn_out
        y, moe_rep = moe_block(post, pos, layer, trace)
        state["last_y"][layer] = y.copy()
        # mHC: nâng y thành 4 làn, trộn qua ma trận Sinkhorn, cộng vào S.
        U = np.einsum("d,md->md", y, W_HC_UP)
        S = np.einsum("mn,md->nd", MIX_COMB, S) + U
        r_r = COMPRESS_RATIOS[layer]
        if layer in CACHE_OWNERS:
            cache = state["caches"][layer]
            if layer in KV_SOURCE_LAYERS and r_r > 0:
                cache["kv_state"].append({"k": attn_rep["k_rot"],
                                          "v": attn_rep["v_new"], "pos": pos})
                events_local = []
                emit_latents(cache, r_r, events_local)
                if trace is not None:
                    for ev in events_local:
                        trace["compressor"].append({"layer": layer, **ev})
            cache["ring"].append({"k": attn_rep["k_rot"], "v": attn_rep["v_new"],
                                  "pos": pos})
            cache["ring"] = cache["ring"][-WINDOW_SIZE:]
    state["S"] = S
    h_final = np.einsum("md,md->d", S, W_HC_DOWN)
    logits = h_final @ W_LM
    state["ids"].append(int(tok))
    if trace is not None:
        trace["logits"].append(logits.copy())
    return logits, attn_rep, engram_done

def new_trace():
    return {"x0": [], "engram_pair": [], "engram_applied": [], "lanes_l4": [],
            "hc_pre_l4": [], "attn": [], "moe": [], "compressor": [],
            "logits": []}

def run_sequence(ids, state=None, trace=None, start_pos=0):
    """Chuỗi token; state tra về để decode tiếp. Trả về state."""
    if state is None:
        state = new_state()
    for i, tok in enumerate(ids):
        run_token(tok, start_pos + i, state, trace)
    return state

def markov_head():
    """MarkovHead hạng 4: ma trận chuyển (VOCAB × VOCAB) chưa chuẩn hóa."""
    return MK_A @ MK_B

def dspark_draft(state):
    """DSpark: gộp hidden của lớp (3, 4), vẽ 3 token nháp bằng MarkovHead."""
    ys = [state["last_y"][l] for l in DSPARK_TARGET_LAYERS]
    fused = np.mean(ys, axis=0)
    base = fused @ W_DRAFT
    M = markov_head()
    def mrow(prev):
        return softmax(M[prev])          # xác suất chuyển > 0 cho burst log an toàn
    bias_row = mrow(int(state["ids"][-1]))
    dist1 = softmax(base + DSPARK_MARKOV_ALPHA * np.log(bias_row + 1e-9))
    drafts = [{"token": int(np.argmax(dist1)), "conf": float(dist1[int(np.argmax(dist1))]),
               "source": "LM nháp gộp (3, 4) + bias Markov"}]
    prev = drafts[0]["token"]
    for _ in range(DSPARK_BLOCK_SIZE - 1):
        p = mrow(prev)
        t = int(np.argmax(p))
        drafts.append({"token": t, "conf": float(p[t]),
                       "source": "chuỗi Markov thuần"})
        prev = t
    return {"fused": fused, "base": base, "M_rank": DSPARK_MARKOV_RANK,
            "drafts": drafts, "bias_last": DSPARK_MARKOV_ALPHA * np.log(bias_row + 1e-9)}

def kv_budget_rows():
    """Bảng ngân sách KV sản xuất khai báo đơn giá; tổng phải đúng 890 B/token."""
    rows = [
        ("Global KV của Decoder (1 latent dùng chung)",
         "512 chiều × 0.5 B (FP4 E2M1)", 512 * 0.5),
        ("Rope tail của Global KV",
         "64 chiều × 1 B (FP8 E4M3)", 64 * 1.0),
        ("Encoder KV nén CSA2 (r=2, gộp 2 token)",
         "3 lớp nguồn × (256 + 64) / 2", 3 * (512 * 0.5 + 64 * 1.0) / 2),
        ("Index KV phân cấp (Indexer)",
         "2 lớp nguồn × 32 × 0.5 B (FP4)", 2 * 32 * 0.5),
        ("attn_sink + metadata",
         "10 phần tử × 1 B (FP8)", 10 * 1.0),
        ("Ring SWA ghi trung bình mỗi token",
         "K+V 2 × 24 chiều × 1 B (FP8)", 24 * 2 * 1.0),
    ]
    return rows

def print_header():
    print("=== LAB DEEPSEEK-V4.1-FLASH: TOÀN TRÌNH 1A → 6B BẰNG NUMPY ===")
    print("MỤC TIÊU: chạy một forward đa phương thức từ prompt ảnh/chữ đến logits,")
    print("đọc được từng khối 1A..6B và kiểm chứng tính nhân quả + tái sử dụng KV.")
    print(f"ĐẦU VÀO: PREFILL {DEMO_PROMPT} (vị trí 1 là token ảnh), "
          f"DECODE {DECODE_TOKENS} tại start_pos 4/5, THỬ {TRY_PROMPT}.")
    print(f"CẤU HÌNH THU NHỎ: vocab={VOCAB}, dim={DIM}, lớp={N_LAYERS}, r={COMPRESS_RATIOS},")
    print(f"  nguồn KV {KV_SOURCE_LAYERS}, cửa sổ {WINDOW_SIZE}, indexer top-{INDEX_TOPK},")
    print(f"  heads={N_HEADS}, d_head={HEAD_DIM}, MoE {N_ACTIVATED_EXPERTS}/{N_ROUTED_EXPERTS},")
    print(f"  Engram ở lớp {ENGRAM_LAYERS}, DSpark lớp {DSPARK_TARGET_LAYERS}.")
    print("Weights chưa huấn luyện, seed 41 — các số chỉ kiểm chứng CƠ CHẾ, không chứng")
    print("minh chất lượng ngôn ngữ và không đo tốc độ production.")
    print()

def fmt_vec(v, nd=3):
    return "[" + ", ".join(f"{x:.{nd}f}" for x in np.asarray(v, dtype=float)) + "]"

def say(code, title, lines):
    print(f"[{code}] {title}")
    for ln in lines:
        print("   " + ln)
    print()

def fmt_ints(seq):
    return ", ".join(str(int(x)) for x in seq)

def print_sections(prefill_trace, prefill_logits, snap, decode_steps,
                   dspark, state):
    # ── [1A] ────────────────────────────────────────────────────────────────
    x0 = prefill_trace["x0"]
    img_vec = x0[1][2]
    say("1A", "ĐẦU VÀO ĐA PHƯƠNG THỨC (Embedding + Aligner 3×3)", [
        "Quan sát: token chữ chỉ tra bảng E; token ảnh bị PHỦ ĐÈ bằng vector gộp.",
        f"shape X0=({len(DEMO_PROMPT)}, {DIM}) — bốn vị trí × dim {DIM}; "
        f"vị trí ảnh = {fmt_ints(IMAGE_POS)}.",
        f"Mã token đầu vào: {fmt_ints(DEMO_PROMPT)} (giả lập qua tokenizer).",
        "Ảnh 3×3 dùng cho gộp (giá trị cố định 0.1..0.9):",
        "   hàng 0: " + fmt_vec(IMG_PATCH[0]),
        "   hàng 1: " + fmt_vec(IMG_PATCH[1]),
        "   hàng 2: " + fmt_vec(IMG_PATCH[2]),
        "Aligner 9 → dim phủ đè vocab[1] đã tra: " + fmt_vec(img_vec),
        "E[1] trước khi bị thay thế: " + fmt_vec(E[1]),
        "X0[0] (token 1 chữ, không bị đụng): " + fmt_vec(x0[0][2]),
        "Kết luận: đa phương thức chỉ khác ở bước nhúng; từ 1A trở đi mọi token",
        "là vector (8,) như nhau — các lớp sau không biết token nào là ảnh.",
    ])

# ── [1B] ────────────────────────────────────────────────────────────────
    pair_lines = []
    for e in prefill_trace["engram_pair"]:
        if e["cur_c"] == DEAD:
            pair_lines.append(f"   pos {e['pos']}: token ảnh → cur=DEAD, bỏ lookup")
        elif e["prev_c"] is None:
            pair_lines.append(f"   pos {e['pos']}: chưa có token trước, bỏ lookup")
        elif e["prev_c"] == DEAD:
            pair_lines.append(f"   pos {e['pos']}: prev=DEAD (bắt qua ảnh) → "
                              "băm bị chặn cắt ngang ảnh")
        else:
            pair_lines.append(f"   pos {e['pos']}: cặp ({e['prev_c']}, {e['cur_c']}) → lookup")
    b_line = []
    for ap in prefill_trace["engram_applied"]:
        bucket = ap["bucket"]
        b_line.append(f"   pos {ap['pos']}: bucket {bucket} → cộng gộp "
                      f"‖Δ‖={np.linalg.norm(ENG_TABLE[bucket]):.3f}")
    if not b_line:
        b_line = ["   (prefill không có cặp hợp lệ — chỉ decode mới có)"]
    lines = [
        f"Map nén token của prompt: "
        f"[{', '.join(str(comp_token(t, p)) for p, t in enumerate(DEMO_PROMPT))}] — "
        f"vị trí ảnh thành {DEAD} (DEAD).",
    ] + pair_lines + b_line

for ap in prefill_trace["engram_applied"]:
        lines.append("   gate của bucket này: " + fmt_vec(ap["gate"]))
        lines.append("   hc_pre của L1 SAU khi cộng engram tại pos: " + fmt_vec(ap["hc_pre"]))
    for ds in decode_steps:
        for ap in ds["engram"]:
            lines.append(f"Decode ▸ pos {ap['pos']}: bucket {ap['bucket']} — "
                         f"pair hợp lệ nên lookup chạy cả khi decode.")
    lines += [
        "Kết luận: Engram chỉ thêm vector khi cặp n-gram không chạm ảnh; token ảnh",
        "được che bằng DEAD nên bộ nhớ n-gram không nhảy qua vùng đa phương thức.",
    ]
    say("1B", "BỘ NHỚ CÓ ĐIỀU KIỆN ENGRAM (N-gram XOR-Prime Hash & Gate)", lines)

# ── [2A] ────────────────────────────────────────────────────────────────
    lanes = np.stack(prefill_trace["lanes_l4"], axis=1)      # (HC_MULT, T, DIM)
    hc_pre = np.stack(prefill_trace["hc_pre_l4"])            # (T, DIM)
    say("2A", "SINGLE-PASS mHC (4 làn Residual & Ma trận Sinkhorn 4×4)", [
        f"shape mở rộng (lanes): {lanes.shape} = {HC_MULT} làn × {len(DEMO_PROMPT)} token "
        f"× {DIM} chiều (chụp trước L4 của prefill).",
        "Lanes của token 3 (làn 0..3):",
    ] + ["   làn " + str(m) + ": " + fmt_vec(lanes[m][-1])
         for m in range(HC_MULT)] + [
        f"hc_pre gộp: shape {hc_pre.shape} — hàng 3 (đầu vào L4 của token 3): "
        + fmt_vec(hc_pre[-1]),
        "Ma trận trộn comb sau Sinkhorn (đầu ra thực của forward):",
    ] + ["   hàng " + str(m) + ": " + fmt_vec(MIX_COMB[m], 6)
         for m in range(HC_MULT)] + [
        "Tổng mỗi hàng: " + fmt_vec(MIX_COMB.sum(axis=1), 6),
        "Tổng mỗi cột: " + fmt_vec(MIX_COMB.sum(axis=0), 6),
        "Kết luận: cập nhật lớp đi qua comb trước khi về 4 làn; mỗi lần trộn giữ",
        "tổng hàng và cột = 1 nên không làn nào được chiếm quan trọng dần.",
    ])

def print_sections_cba(prefill_trace, snap, decode_steps, state):
    # ── [3A] ────────────────────────────────────────────────────────────────
    role_lines = []
    for l in range(N_LAYERS):
        r = COMPRESS_RATIOS[l]
        if r == 0:
            role, kv_owner = "SWA thuần", "tự giữ ring riêng (không ai dùng lại)"
        elif l in KV_SOURCE_LAYERS:
            kind = "Encoder nén" if r == 2 else "Decoder Global KV"
            role = f"{kind} (r={r})"
            kv_owner = f"TỰ LÀ NGUỒN KV — ghi cache cho L{l + 1}"
        else:
            role = f"thuần attention (r={r})"
            kv_owner = f"DÙNG LẠI KV do L{owner_of(l)} ghi — không tính lại K/V"
        role_lines.append(f"   L{l}: {role} · {kv_owner}")
    l1, l3 = state["caches"][1], state["caches"][3]
    role_lines += [
        f"Sau prefill: L1 có {snap['lat_n'][1]} latent (end "
        f"{fmt_ints(snap['ends'][1])}) và L3 có {snap['lat_n'][3]} Global latent "
        f"(end {fmt_ints(snap['ends'][3])}).",
        f"Sau 2 decode: L1 thêm latent end {fmt_ints([e['end'] for e in l1['latents'][snap['lat_n'][1]:]])} "
        f"và L3 có tổng {len(l3['latents'])} latent (end {fmt_ints([e['end'] for e in l3['latents']])}).",
        "Kết luận: ranh giới Encoder (L0..L2) và Decoder (L3..L4) là ranh giới",
        "NGUỒN KV, không phải ranh giới tính toán: cả 5 lớp cùng một forward.",
    ]
    say("3A", "CAUSAL ENCODER-DECODER (CED: Lớp 0..2 Encoder → Lớp 3..4 Decoder)", role_lines)

# ── [3B] ────────────────────────────────────────────────────────────────
    comp_lines = [
        f"Compressor gộp {COMPRESS_RATIOS[1]} token → 1 latent (pool lấy trung bình K và V):",
    ]
    for e in prefill_trace["compressor"]:
        if e["layer"] != 1:
            continue
        if e["emitted"]:
            comp_lines.append(f"   prefill pos {e['end']}: đủ {COMPRESS_RATIOS[1]} token "
                              f"({fmt_ints(e['members'])}) → PHÁT latent end={e['end']}")
        else:
            comp_lines.append(f"   prefill pos {e['pending_pos'][0]}: lẻ ({e['pending']}) "
                              "→ buffer kv_state, trả None")
    for ds in decode_steps:
        for e in ds["compressor_l1"]:
            if e["emitted"]:
                comp_lines.append(f"   decode pos {e['end']}: đủ cặp ({fmt_ints(e['members'])}) "
                                  f"→ PHÁT latent end={e['end']}")
            else:
                comp_lines.append(f"   decode pos {e['pending_pos'][0]}: lẻ ({e['pending']}) "
                                  "→ buffer kv_state, trả None")
    comp_lines += [
        "Giải thích số lẻ: khối đo bằng số token CHƯA gộp trong kv_state; prefill",
        "trả 4 token → 2 latent; decode bước 1 đẩy pending từ 0 lên 1 (lẻ, None);",
        "decode bước 2 hoàn tất cặp → phát latent thứ 3.",
        "Chia sẻ: L2 đọc đúng mảng latent của L1; L4 đọc đúng Global KV của L3",
        "(kiểm chứng array_equal nằm ở [KIỂM TRA & THỬ NGHIỆM]).",
    ]
    say("3B", "COMPRESSED SPARSE ATTENTION 2 (Compressor gộp r=2 & Chia sẻ KV)", comp_lines)

# ── [3C] ────────────────────────────────────────────────────────────────
    idx_lines = [
        f"Ring SWA {WINDOW_SIZE} slot: chỉ giữ kv của {WINDOW_SIZE} token trước; "
        "nguồn: latent nén qua Indexer + ring + token hiện tại.",
        f"Ring L0 sau prefill: pos {fmt_ints(snap['ring'][0])} (riêng, không ai dùng lại)",
        f"Ring L1 sau prefill: pos {fmt_ints(snap['ring'][1])} (L2 dùng lại)",
        f"Ring L3 sau prefill: pos {fmt_ints(snap['ring'][3])} (L4 dùng lại)",
        indexer_example_line(prefill_trace, 3),
    ]
    for ds in decode_steps:
        for a in ds["attn_l1"]:
            for h, ix in enumerate(a["indexer"]):
                if not ix["scores"]:
                    continue
                idx_lines.append(
                    f"Decode ▸ pos {a['pos']} head {h}: eligible end {fmt_ints(ix['eligible_ends'])} "
                    f"→ chọn end {fmt_ints(ix['picked_end'])}")
    idx_lines += [
        "Lưu ý đọc số: một vị trí có thể xuất hiện 2 lần (qua latent nén và qua ring)",
        "— attention phân cấp trộn cả hai kênh, đúng như thiết kế.",
        "Kết luận: Indexer chọn mức duy nhất (top-1) giữa các latent đã phát;",
        "các slot ring cung cấp cửa sổ gần chính xác theo vị trí tuyệt đối.",
    ]
    say("3C", "RING SWA & HIERARCHICAL SPARSE INDEXER", idx_lines)

def indexer_example_line(prefill_trace, pos):
    rep = [a for a in prefill_trace["attn"] if a["layer"] == 1 and a["pos"] == pos][0]
    msgs = []
    for h, ix in enumerate(rep["indexer"]):
        if not ix["scores"]:
            continue
        msgs.append("head " + str(h) + ": eligible end " + fmt_ints(ix["eligible_ends"])
                    + " → chọn end " + fmt_ints(ix["picked_end"])
                    + " (score ∝ " + fmt_vec(ix["scores"], 2) + ")")
    return ("Indexer tại prefill pos " + str(pos) + ": "
            + (" / ".join(msgs) if msgs else "không có latent hợp lệ"))

def print_sections_dsparks(prefill_trace, decode_steps, dspark, prefill_logits):
    # ── [4A] ────────────────────────────────────────────────────────────────
    rep = [a for a in prefill_trace["attn"] if a["layer"] == 1 and a["pos"] == 3][0]
    hr = rep["heads"][0]
    say("4A", "MULTI-HEAD LATENT ATTENTION (MLA, Inverse RoPE & Group-LoRA W_O)", [
        f"Q low-rank: ({DIM}) → ({Q_LORA_RANK}) → {rep['q_raw'].shape} "
        f"= {N_HEADS} head × {HEAD_DIM} (hạng chỉ {Q_LORA_RANK}).",
        "Q head 0 chưa xoay: " + fmt_vec(rep["q_raw"][0]),
        "Q head 0 đã xoay RoPE (pos 3): " + fmt_vec(rep["q_rot"][0]),
        f"K/V MỘT head dùng chung cho cả {N_HEADS} Q head; K shape ({HEAD_DIM},).",
        "K chưa xoay: " + fmt_vec(rep["k_raw"]) + " | K đã xoay: " + fmt_vec(rep["k_rot"]),
        f"RoPE chỉ áp NỬA CUỐI ({ROPE_HEAD_DIM} chiều cuối); {HEAD_DIM - ROPE_HEAD_DIM} chiều đầu giữ nguyên.",
        "V token hiện tại: " + fmt_vec(rep["v_new"]),
        "attn_sink (thêm một cột điểm trước softmax, mỗi head một sink): "
        + fmt_vec(rep["sink"]),
        f"Weights của head 0 trên các nguồn {fmt_ints(hr['src_pos'])}: "
        + fmt_vec(hr["weights"]) + f" (phần rơi vào sink = {hr['sink_w']:.3f})",
        "ctx head 0 TRƯỚC inverse-RoPE: " + fmt_vec(rep["ctx_pre"][0]),
        "ctx head 0 SAU inverse-RoPE (xoay ngược góc pos 3): " + fmt_vec(rep["ctx"][0]),
        f"W_O nhóm: ctx đầu ra ({N_HEADS}, {HEAD_DIM}) gom thành {O_GROUPS} group, "
        f"mỗi group {HEAD_DIM} → {O_LORA_RANK} (wo_a), ghép "
        f"{O_GROUPS * O_LORA_RANK} → {DIM} (wo_b).",
        "Kết luận: MLA giảm KV bằng K/V latent một head; inverse RoPE đưa đầu ra",
        "về không gian không vị trí trước W_O nhóm, giữ cấu trúc low-rank.",
    ])

# ── [5A] ────────────────────────────────────────────────────────────────
    moe_lines = []
    for m in prefill_trace["moe"]:
        if m["layer"] != 1 or m["pos"] not in (1, 3):
            continue
        tag = " (token ảnh: +bias_vl)" if m["pos"] in IMAGE_POS else ""
        moe_lines += [
            f"pos {m['pos']}{tag}: scores=" + fmt_vec(m["scores"]),
            "   gate = sqrt(softplus(scores)) = " + fmt_vec(m["gate"]),
            f"   bias chọn chuyên gia = " + fmt_vec(m["bias"]),
            f"   chọn theo scores+bias: chuyên gia {fmt_ints(m['sel'])} — "
            f"nhưng TRỌNG SỐ lấy từ gate thô chuẩn hóa: " + fmt_vec(m["w_route"]),
            f"   SwiGLU clamp ±{SWIGLU_LIMIT:g}: {m['clamped']} expert bị kẹp; "
            f"shared expert cộng đầy đủ (‖shared‖={m['shared_norm']:.3f})",
        ]
    for ds in decode_steps:
        for m in ds["moe_l1"]:
            moe_lines.append(f"Decode ▸ pos {m['pos']}: chuyên gia {fmt_ints(m['sel'])}"
                             f", trọng số " + fmt_vec(m["w_route"]))
    moe_lines += [
        "Kết luận: softmax KHÔNG được dùng để chọn; chọn là top-k của scores+bias",
        "(bias_vl chỉ cộng tại vị trí ảnh), nhân lại bằng gate thô — tách vai trò.",
    ]
    say("5A", "FINE-GRAINED MoE (Cổng sqrtsoftplus, tách bias & SwiGLU clamp)", moe_lines)

# ── [6A] ────────────────────────────────────────────────────────────────
    rows = kv_budget_rows()
    total = sum(r[2] for r in rows)
    assert total == 890.0, f"ngân sách KV phải đúng 890 B/token, đang {total}"
    kv_lines = ["Bảng ngân sách KV của sản phẩm (đơn giá khai báo trong lab):"]
    for name, detail, b in rows:
        kv_lines.append(f"   {b:6.0f} B  ← {name} ({detail})")
    mha = 2 * 128 * 128 * 2 * 61
    gqa = 2 * 8 * 128 * 2 * 61
    kv_lines += [
        f"TOTAL: {total:.0f} B/token cho DeepSeek-V4.1-Flash.",
        "So sánh cùng lớp 61, head 128×128, BF16=2 B:",
        f"   MHA đầy đủ: {mha:,} B/token ({mha / total:.0f}× lớn hơn)",
        f"   GQA 8 KV head: {gqa:,} B/token ({gqa / total:.0f}× lớn hơn)",
        "Cache hình học của lab này (số nhỏ, mảng NumPy):",
        f"   L1 latent mỗi cái k/v ({HEAD_DIM},) — đang có {len(state_latents(state, 1))} latent",
        f"   L3 Global latent: {len(state_latents(state, 3))} cái; pending kv_state L1: "
        f"{len(state['caches'][1]['kv_state'])}",
        f"   Ring {WINDOW_SIZE} slot × ({HEAD_DIM},) mỗi lớp nguồn {CACHE_OWNERS}.",
        "Kết luận: giá rẻ nhờ (1) latent FP4, (2) nén r, (3) chia sẻ KV giữa các lớp.",
    ]
    say("6A", "NGÂN SÁCH BỘ NHỚ KV (FP4 E2M1 & Công thức 890 bytes/token)", kv_lines)

# ── [6B] ────────────────────────────────────────────────────────────────
    arg_rows = [int(np.argmax(r)) for r in prefill_logits]
    lm_lines = [
        f"Logits prefill gộp từ {HC_MULT} làn sau L4: shape {prefill_logits.shape} "
        f"= {len(arg_rows)} vị trí × {VOCAB} token ứng viên.",
        "argmax theo hàng: " + fmt_ints(arg_rows)
        + " (chọn bằng số, weights chưa học).",
        "DSpark: gộp hidden lớp " + fmt_ints(DSPARK_TARGET_LAYERS) + " (0.5/0.5) → "
        + fmt_vec(dspark["fused"]),
        f"MarkovHead: ma trận (8×8) hạng {DSPARK_MARKOV_RANK} (= A(8,4)·B(4,8)); "
        f"block {DSPARK_BLOCK_SIZE} token nháp.",
        "bias Markov cho token cuối (alpha "
        f"{DSPARK_MARKOV_ALPHA}): " + fmt_vec(dspark["bias_last"]),
    ]
    for i, d in enumerate(dspark["drafts"], 1):
        lm_lines.append(f"   nháp {i}: token {d['token']} — confidence "
                        f"{d['conf']:.3f} ({d['source']})")
    lm_lines += [
        "Kết luận: draft 1 dùng logits thật + bias Markov; nháp sau kéo theo chuỗi",
        "Markov; confidence chỉ là xác suất của chính con số nháp đó, chưa verify.",
    ]
    say("6B", "ĐẦU RA LM HEAD & DSPARK SPECULATIVE DECODING", lm_lines)

def state_latents(state, l):
    return state["caches"][l]["latents"]

def print_checks(prefill_trace, prefill_logits, decode_steps, state):
    # (1) Sinkhorn phải doubly stochastic.
    err = float(max(np.max(np.abs(MIX_COMB.sum(axis=1) - 1.0)),
                    np.max(np.abs(MIX_COMB.sum(axis=0) - 1.0))))
    print("[KIỂM TRA & THỬ NGHIỆM]")
    print(f"  (1) Sinkhorn 4×4 doubly-stochastic: lệch tối đa hàng/cột = "
          f"{err:.2e} (ngưỡng 1e-6).")
    assert err < 1e-6
    print("      → PASS: ma trận trộn mHC giữ tổng mỗi hàng và mỗi cột bằng 1.")

# (2) Tính nhân quả: đổi token cuối, các vị trí trước phải đứng yên.
    try_state = new_state()
    try_trace = new_trace()
    run_sequence(TRY_PROMPT, try_state, try_trace)
    try_logits = np.stack(try_trace["logits"])
    diffs = np.max(np.abs(try_logits[:3] - prefill_logits[:3]), axis=1)
    for i, d in enumerate(diffs):
        print(f"  (2) THỬ {TRY_PROMPT} vs {DEMO_PROMPT}: hàng logits {i} lệch "
              f"{d:.2e} (ngưỡng 1e-9).")
    assert diffs.max() <= 1e-9
    diff_last = float(np.max(np.abs(try_logits[3] - prefill_logits[3])))
    print(f"      hàng 3 (vị trí token bị đổi): lệch {diff_last:.2e} — được phép đổi.")
    print("      → PASS: dữ liệu sau vị trí i không rỉ vào logits của vị trí ≤ i.")

# (3) Tái sử dụng KV/Index giữa lớp nguồn và lớp dùng lại.
    for l_re in (2, 4):
        rep = [a for a in prefill_trace["attn"] if a["layer"] == l_re
               and a["pos"] == 3][0]
        src = owner_of(l_re)
        store = state["caches"][src]
        assert np.array_equal(rep["latent_k0"], store["latents"][0]["k"]), (
            f"L{l_re} phải đọc đúng latent của L{src}")
        assert rep["latent_end_list"] == [lt["end"] for lt in store["latents"]][:len(
            rep["latent_end_list"])]
    print("  (3) L2 đọc đúng mảng latent K của L1; L4 đọc đúng Global KV của L3")
    print("      (array_equal trên latent đầu + danh sách end khớp snapshot).")
    assert state["caches"][owner_of(2)] is state["caches"][1]
    assert state["caches"][owner_of(4)] is state["caches"][3]
    print("      → PASS: lớp dùng lại KHÔNG mang cache riêng; nguồn chiếm đúng 1 bản sao.")

# (4) Decode streaming khớp tính lại toàn chuỗi.
    full_state = new_state()
    full_trace = new_trace()
    run_sequence(list(DEMO_PROMPT) + list(DECODE_TOKENS), full_state, full_trace)
    full_logits = np.stack(full_trace["logits"])
    for j, ds in enumerate(decode_steps):
        row = len(DEMO_PROMPT) + j
        d_err = float(np.max(np.abs(np.asarray(ds["logits"]) - full_logits[row])))
        print(f"  (4) Decode start_pos={row} khớp tính lại toàn chuỗi: lệch "
              f"{d_err:.2e} (ngưỡng 1e-9).")
        assert d_err <= 1e-9
    ev_odd = [e for e in decode_steps[0]["compressor_l1"] if not e["emitted"]]
    ev_even = [e for e in decode_steps[1]["compressor_l1"] if e["emitted"]]
    assert ev_odd and ev_odd[0]["pending"] == 1, "decode 1 phải chờ cặp lẻ"
    assert ev_even and ev_even[0]["end"] == 5, "decode 2 phải phát latent end=5"
    print("  (5) Decode bước lẻ chờ kv_state (None), bước pair phát latent end=5.")
    print("PASS: toàn trình 1A→6B chạy đúng; các số trên mới kiểm chứng CƠ CHẾ")
    print("đã khai báo, không chứng minh chất lượng ngôn ngữ hay tốc độ production.")

# ── Demo chạy khi thực thi (trình duyệt/CLI/test đều dùng chung) ────────────
print_header()

# 1) PREFILL DEMO_PROMPT và chụp cache sau prefill.
state = new_state()
prefill_trace = new_trace()
run_sequence(DEMO_PROMPT, state, prefill_trace, start_pos=0)
prefill_logits = np.stack(prefill_trace["logits"])
assert prefill_logits.shape == (len(DEMO_PROMPT), VOCAB)
snap = {
    "ring": {l: [e["pos"] for e in state["caches"][l]["ring"]] for l in CACHE_OWNERS},
    "lat_n": {l: len(state["caches"][l]["latents"]) for l in CACHE_OWNERS},
    "ends": {l: [lt["end"] for lt in state["caches"][l]["latents"]] for l in CACHE_OWNERS},
    "pend": {l: len(state["caches"][l]["kv_state"]) for l in CACHE_OWNERS},
}
assert snap["lat_n"][1] == 2 and snap["ends"][1] == [1, 3]
assert snap["lat_n"][3] == len(DEMO_PROMPT)

# 2) Hai bước decode nối tiếp trên cùng state (streaming).
decode_steps = []
pos = 4
for tok in DECODE_TOKENS:
    step_trace = new_trace()
    logits, _, _ = run_token(tok, pos, state, step_trace)
    decode_steps.append({
        "pos": pos, "tok": tok, "logits": logits,
        "compressor_l1": [e for e in step_trace["compressor"] if e["layer"] == 1],
        "engram": step_trace["engram_applied"],
        "attn_l1": [a for a in step_trace["attn"] if a["layer"] == 1],
        "moe_l1": [m for m in step_trace["moe"] if m["layer"] == 1],
    })
    pos += 1
assert len(state["caches"][1]["latents"]) == 3
assert [e["end"] for e in state["caches"][3]["latents"]] == [0, 1, 2, 3, 4, 5]

# 3) DSpark vẽ token nháp từ trạng thái hiện tại.
dspark = dspark_draft(state)
assert len(dspark["drafts"]) == DSPARK_BLOCK_SIZE
assert all(0 <= d["token"] < VOCAB and 0.0 < d["conf"] <= 1.0
           for d in dspark["drafts"])

# 4) In toàn bộ câu chuyện 1A..6B bằng số thật.
print_sections(prefill_trace, prefill_logits, snap, decode_steps, dspark, state)
print_sections_cba(prefill_trace, snap, decode_steps, state)
print_sections_dsparks(prefill_trace, decode_steps, dspark, prefill_logits)
print()

# 5) Kiểm chứng số học.
print_checks(prefill_trace, prefill_logits, decode_steps, state)
print()
print("TỔNG KẾT:")
print(f"1. Token ảnh được phủ đè ở 1A và DEAD=-1 chặn Engram/CSA2 qua vùng ảnh (1B/3B).")
print(f"2. mHC trộn 4 làn bằng Sinkhorn {HC_MULT}×{HC_MULT} chuẩn hóa kép; logits gộp (4, {VOCAB}).")
print(f"3. L1 nén r=2, phát latent theo cặp; L3 giữ Global KV; L2/L4 dùng lại thành công.")
print(f"4. MLA giảm KV về một head; MoE chọn theo scores+bias nhưng nhân bằng gate thô.")
print(f"5. Ngân sách KV 890 B/token là khai báo đơn giá; lab không đo tốc độ thật.")

Python outputSẵn sàng.

Nguồn công bố & mã nguồn tham chiếu

deepseek-ai/DeepSeek-V4.1-Flash (Hugging Face) ↗ Model card, config.json, DeepSeek_V41_Tech_Report.pdf & thư mục inference/ (MIT) Engram · Conditional Memory via Scalable Lookup ↗ Bảng băm N-gram 196B tham số tại lớp 1 và 14 DeepSeek-V3 Technical Report ↗ Nền tảng Multi-head Latent Attention (MLA) và cân bằng tải MoE không aux-loss

Danh sách bài học (11)