L2 · Operator graph

Mô hình transformer biến token thành dự đoán như thế nào

Theo dõi token qua một lớp Transformer để hiểu cách mô hình lấy thông tin từ ngữ cảnh, tạo điểm dự đoán và dùng lại KV cache khi sinh tiếp.

LESSON ID transformer GROUP Kiến trúc Inference

1Định nghĩa

Transformer là mạng nơ-ron biến chuỗi vector biểu diễn token thành các vector đã kết hợp thông tin ngữ cảnh. Trong mô hình sinh văn bản của bài này, mã token từ tokenizer được đổi thành vector, đi qua các lớp Transformer, rồi được LM head chấm điểm trên bảng từ vựng. Hãy phân biệt hai việc: mạng tính điểm cho các token ứng viên; bước chọn token quyết định lấy token nào từ những điểm đó.

Một lớp có hai công việc chính: attention lấy thông tin từ các vị trí được phép nhìn, còn FFN biến đổi riêng vector tại từng vị trí. Chuẩn hóa và residual bao quanh hai nhánh để điều chỉnh biên độ và cộng phần cập nhật vào thông tin cũ. Với decoder causal, mặt nạ nhân quả ngăn vị trí trước đọc vị trí sau. Sau bài học, bạn cần theo được đường đi này trên hình và giải thích vì sao cùng weights có thể xử lý cả prompt hoặc chỉ token mới với KV cache.

Ví dụ đời thường: Hãy hình dung một nhóm biên tập sửa bản thảo qua nhiều vòng. Ở mỗi vòng, một biên tập viên đối chiếu đoạn mình phụ trách với những đoạn được phép đọc, rồi tự sửa đoạn đó; attention tương tự bước đối chiếu, FFN tương tự bước xử lý riêng. Residual giữ bản hiện có để cộng thêm phần sửa. Ví dụ này giúp nhớ hai công việc khác nhau; mạng thực tế chỉ tính trên vector số, không có người đọc hay chuyên gia hiểu câu chữ bên trong từng head.

2Luồng

Cách đọc hình: bắt đầu từ Hidden State ở góc dưới cột 1, đi lên Q/K/V, theo các mũi tên sang attention ở cột 2 và đi lên đầu ra ở cột 4. Cột 3 là nhánh cấp K/V cho attention, không phải công đoạn chỉ bắt đầu sau cột 2. Các mục MHA/GQA/MQA/MLA và Dense/MoE là lựa chọn kiến trúc, không phải các bước phải chạy nối tiếp. Lab bên dưới chọn MHA và Dense SwiGLU.

Sơ đồ kiến trúc tương tác · 12 thành phần

Kiến trúc Transformer Layer & Cơ chế Attention (Overlay giải thích chi tiết)

Di chuột hoặc bấm vào từng vùng đánh số (1A → 4C) trên hình kiến trúc bên dưới để xem giải thích cơ chế, kích thước tensor và mở bài học chi tiết của thành phần đó.

Sơ đồ kiến trúc Transformer Decoder Layer và cơ chế Attention gồm 4 chặng: Input & QKV Projection, Attention Core & Causal Mask, KV Cache & Head Sharing (MHA/GQA/MQA/MLA), và Output, Residual & FFN/MoE

1A 1 · Input & QKV Projection

x_t: [1, d_model] → x̂_t = RMSNorm(x_t)

Hidden State x_t & RMSNorm (Pre-Norm)

Chuẩn hóa biên độ vector ẩn đầu vào trước khi tách nhánh chiếu Q, K, V.

Bắt đầu từ vector của một vị trí: ở lớp đầu nó đến từ embedding; ở lớp sau nó đến từ lớp trước. RMSNorm điều chỉnh biên độ trước khi tạo Q/K/V. Hãy giữ riêng vector gốc x_t: bước 4A còn cần nó để cộng residual. Trong lab, so hàng đầu trước và sau RMSNorm; số chiều không đổi, chỉ giá trị đổi. Đây là chuẩn bị dữ liệu, chưa có bước trộn thông tin giữa token.

Normalization (RMSNorm / Pre-Norm) Chuẩn hóa & Residual Kiến trúc Transformer

3Giải pháp

Bắt đầu từ câu hỏi: một token cần biết những gì để dự đoán token tiếp theo? Mã số tự nó chưa đủ. Bảng embedding đổi mã thành vector; ở mỗi lớp, vector đó được chuẩn hóa rồi chiếu thành Q, K và V. Q/K dùng để đối chiếu độ liên quan, V mang nội dung cần gom. Trong kiến trúc minh họa, RoPE xoay Q/K theo vị trí, giúp phép đối chiếu phân biệt trật tự token.

Tiếp theo, hãy theo đường từ 2B đến 2E trên hình. Tích Q với K tạo điểm liên quan, nhưng điểm chưa cho biết nên lấy bao nhiêu thông tin. Causal mask chặn các vị trí tương lai; softmax đổi điểm hợp lệ thành trọng số có tổng bằng 1; nhân trọng số với V mới tạo ra thông tin ngữ cảnh. Đừng nhầm softmax này với phân phối chọn token đầu ra: nó phân bổ trọng số trên các vị trí đầu vào.

Ở cột 3, đặt hai câu hỏi riêng. MHA/GQA/MQA trả lời mỗi Q head dùng bộ K/V nào; KV cache trả lời có thể giữ lại K/V của token nào để khỏi tính lại. Lab chọn MHA, tức mỗi Q head có K/V riêng. Khi sinh tiếp, cache giữ K/V token cũ, còn token mới vẫn phải đi qua attention và FFN. Cột này cấp dữ liệu cho attention, không phải một lớp xử lý đứng sau attention.

Cuối lớp, kết quả các head được ghép và chiếu qua W_O, rồi cộng vào vector đầu vào. Một nhánh chuẩn hóa–FFN–residual tiếp tục cập nhật từng vector. Nếu còn lớp, kết quả đi vào lớp kế; chỉ sau lớp cuối mới chuẩn hóa lần cuối và đưa qua LM head để tạo logits. Sơ đồ có nhánh Dense hoặc MoE; hãy theo nhánh Dense SwiGLU trong lab, không đọc hai lựa chọn này như hai bước nối tiếp.

Khi đã hiểu đường tính, mới xét cách chạy nhanh. Kernel có thể gộp các phép toán và giảm đọc/ghi bộ nhớ mà vẫn giữ kết quả trong sai số cho phép. Ngược lại, thay MHA bằng GQA hoặc Dense bằng MoE là thay cấu trúc model; không thể coi đó chỉ là bật một tùy chọn runtime cho checkpoint bất kỳ.

Để so trọng số attention trên các vị trí với logits trên toàn bộ từ vựng bằng số thật, hãy chạy khối 5 (Transformer × N layers) trong Bài Lab tổng hợp.

4Profiling

Kiểm tra dữ liệu có đi đúng đường không
Trước khi đo nhanh/chậm, hãy đọc kích thước: ba token, vector bốn số phải tạo đầu vào (3, 4); LM head với sáu token ứng viên phải tạo logits (3, 6). Nếu kích thước sai, quay lại chỗ ghép heads hoặc phép chiếu, không cố giải thích ý nghĩa bảng số sai.
Kiểm tra token có nhìn trộm tương lai không
Quan sát hàng đầu của trọng số attention: nó chỉ được có trọng số tại chính vị trí đầu. Sau đó thay token cuối và so logits của các vị trí trước. Giữ nguyên weights và các token còn lại; nếu logits trước thay đổi ngoài sai số, kiểm tra mask và vị trí trước khi kết luận model hoạt động đúng.
Kiểm tra dùng cache có giữ kết quả không
Chạy cùng chuỗi bằng hai đường: tính lại toàn chuỗi và chỉ đưa token mới qua mạng với cache. So hàng logits cuối bằng độ lệch tuyệt đối lớn nhất; lab yêu cầu không vượt 1e-9. Đồng thời kiểm tra prefix cache không bị sửa. Một kết quả khớp xác nhận trường hợp đã thử, không thay cho mọi kiểm thử của model production.
Chỉ đo tốc độ khi đã giữ cố định điều kiện chạy
Nếu chuyển sang checkpoint thật, đo riêng prefill và decode, ghi độ dài prompt, số token sinh, batch, kiểu số và thiết bị. Lab NumPy này không có benchmark GPU và không đưa ra mức tăng tốc. Cache giảm phần tính lại K/V cũ nhưng attention của token mới vẫn phải đọc thông tin quá khứ; chi phí không biến mất.

5Ví dụ

Cùng một token, ngữ cảnh khác nhau

Giả sử hai chuỗi có cùng token ở vị trí cuối nhưng phần trước khác nhau. Embedding của token cuối giống nhau, nhưng attention được phép đọc các token trước khác nhau nên trạng thái cuối và logits có thể khác. Điều cần nhớ là biểu diễn có ngữ cảnh không chỉ phụ thuộc mã token tại chính vị trí đó.

Xử lý prompt rồi sinh tiếp

Prefill xử lý ba mã [1, 2, 3], tạo ba hàng logits và K/V cho ba vị trí. Hàng logits cuối được dùng để chọn token tiếp theo. Token vừa chọn được đưa vào lượt decode với cache ba vị trí; mạng nối K/V của nó, tạo một hàng logits mới để chọn token sau nữa. Tách rõ chọn token và xử lý token sẽ giúp bạn tránh đếm nhầm số vị trí trong cache.

Huấn luyện khác thực hành này ở đâu?

Khi huấn luyện dự đoán token kế tiếp, logits tại mỗi vị trí được so với nhãn của token sau vị trí đó để tính loss; lan truyền ngược và optimizer mới cập nhật weights. Lab ở cuối trang chỉ chạy forward và kiểm tra tính nhất quán. Việc in PASS không có nghĩa là weights đã học hay mô hình dự đoán đúng từ tiếp theo trong văn bản thật.

6Thuật ngữ

Hidden state — trạng thái ẩn: Vector biểu diễn một vị trí ở một thời điểm trong mạng. Đầu lớp nó nhận thông tin từ lớp trước; cuối lớp nó đã được cập nhật bằng attention và FFN.
Attention head — đầu chú ý: Một nhánh Q/K/V tạo trọng số và gom thông tin từ các vị trí được phép nhìn. Nhiều head có tham số riêng; không mặc định gán mỗi head cho một nhiệm vụ ngôn ngữ cố định.
MHA — Multi-Head Attention: Cách tổ chức attention mà mỗi Q head có K/V riêng. Trong lab có hai head; Q0 dùng K0/V0 và Q1 dùng K1/V1, không dùng chung K/V giữa hai head.
Residual connection — đường cộng tắt: Giữ đầu vào của một nhánh rồi cộng với kết quả xử lý. Hãy đọc x + Attention(x) như giữ thông tin cũ và bổ sung phần cập nhật, không phải thay toàn bộ x.
LM head — đầu ra bảng từ vựng: Phép chiếu trạng thái ẩn cuối thành một điểm cho mỗi token ứng viên. Nó tạo logits; chọn token từ các điểm này là bước tiếp theo, không phải chính phép chiếu.

Glossary (10)

Token ID và embedding: Token ID là số dùng để tra bảng; embedding là vector lấy từ hàng tương ứng trong bảng đó. Hai token có ID gần nhau không vì thế mà có nghĩa gần nhau.
Query / Key / Value — Q/K/V: Ba phép chiếu từ vector đã chuẩn hóa. Q và K được so sánh để tạo trọng số; V là nội dung đem gom theo trọng số. Chúng là các vector số, không phải ba câu văn.
RoPE: Phép xoay Q và K theo vị trí token, giúp điểm Q·K phản ánh khoảng cách tương đối. V không xoay trong lab; vị trí token mới phải tiếp tục từ độ dài cache, không quay về 0.
RMSNorm: Chuẩn hóa biên độ một vector theo căn trung bình bình phương các phần tử. Lab dùng bản không có hệ số học được; mục đích là quan sát luồng tính, không tái tạo đầy đủ checkpoint thật.
Causal mask — mặt nạ nhân quả: Quy tắc vị trí i chỉ nhìn vị trí j không lớn hơn i. Lab đặt điểm vị trí tương lai thành âm vô cực trước softmax, khiến trọng số ở đó bằng 0.
Logits: Điểm số chưa chuẩn hóa trên bảng từ vựng, không phải xác suất. Đọc hàng cuối để biết điểm cho token đi sau toàn bộ chuỗi đầu vào.
KV cache: K/V của các token đã xử lý được giữ lại để dùng ở lượt sau. MHA quyết định head nào có bộ K/V nào; cache quyết định dữ liệu token cũ nào được tái sử dụng.
Prefill và decode: Prefill xử lý chuỗi prompt ban đầu. Trong mỗi lượt decode của lab, token vừa chọn được đưa lại vào mạng để tính điểm cho token tiếp theo, đồng thời nối K/V mới vào cache.
Dense SwiGLU: FFN dùng một bộ trọng số chung tại mọi vị trí; một nhánh có cổng điều tiết nhánh còn lại. Nó xử lý từng vector riêng, không tự trao đổi thông tin giữa các vị trí như attention.
Forward pass và weights: Forward tính đầu ra bằng bộ tham số hiện có, không tự cập nhật tham số. Trong lab weights là số ngẫu nhiên cố định; cần huấn luyện mới học được khả năng ngôn ngữ.

7Bài liên quan

Chuyển văn bản thành đơn vị mà mô hình hiểu Vị trí · RoPE, partial RoPE, NoPE và YaRN MHA · Multi-Head Attention, độc lập đầu Q/K/V và chi phí KV cache Query, key và value kết hợp context như thế nào Chuẩn hóa vector và đường cộng residual trong transformer Mạng truyền thẳng và cổng kích hoạt sau attention Bộ nhớ đệm key và value của attention Chọn token từ điểm số và xuất văn bản Hàm mất mát: model gán bao nhiêu xác suất cho đáp án đúng

8Nguồn

Đọc cơ chế attention và multi-head trong Attention Is All You Need; lưu ý bài gốc dùng encoder–decoder, không trùng toàn bộ decoder trong lab Đọc cách RoPE đưa vị trí vào Q/K trong RoFormer Đọc biến thể FFN có cổng, nền tảng của SwiGLU trong GLU Variants Improve Transformer Đối chiếu nhiệm vụ dự đoán token kế tiếp và cách tạo nhãn trong hướng dẫn causal language modeling

9Bài lab

Theo dõi ba token đi qua một lớp Transformer

Thực hành NumPy · quan sát cơ chế, không huấn luyện mô hình

Ở section 2, bạn đã thấy các khối của một lớp Transformer. Bây giờ hãy theo dõi dữ liệu thật đi qua chúng: từ ba mã token đến bảng điểm dự đoán, rồi xử lý thêm một token bằng KV cache. Bạn không cần nhớ hết công thức trước khi chạy; mỗi bước sẽ chỉ ra nên quan sát gì và kết quả có ý nghĩa gì.

Bạn sẽ học được gì?

Sau bài này, bạn cần tự giải thích được ba việc: token lấy thông tin từ đâu; vector sau một lớp được đổi thành điểm dự đoán như thế nào; vì sao có thể dùng lại K/V của token cũ khi xử lý token mới.

Bạn sẽ nối được mỗi kết quả in ra với một vùng 1A–4C trên hình, thay vì chỉ nhớ tên các khối. Hãy đọc lời hướng dẫn trước mỗi bảng số rồi đối chiếu vùng tương ứng.

Đầu vào của bài lab

Đầu vào là DEMO_PROMPT = [1, 2, 3]. Đây là ba mã token giả lập, coi như đã qua tokenizer; chúng không gắn với ba từ tiếng Việt cụ thể. Bảng từ vựng có sáu mã, từ 0 đến 5. Không cần tải dữ liệu hay nhập API key.

Mỗi mã được tra thành vector bốn số. Vì có ba vị trí, bảng vector đầu vào có ba hàng và bốn cột, viết là (3, 4). Khi gặp shape trong output, hãy đọc các con số như số hàng/cột hoặc số head, đừng coi chúng là công thức mới.

Mô hình bạn đang sử dụng

Lab chọn MHA: mỗi Q head có K/V riêng. Có hai head, Q0 dùng K0/V0 và Q1 dùng K1/V1. Mỗi vector head có hai số (d_head = 2); ghép hai head lại được bốn số. MHA trả lời ai dùng bộ K/V nào; tái sử dụng cache trả lời dữ liệu token cũ nào không cần tính lại.

Q/K nhận thông tin vị trí bằng RoPE. Sau attention, lab dùng FFN Dense SwiGLU, không dùng MoE. GQA/MQA/MLA là các lựa chọn khác trên hình, không chạy trong lab này. Mô hình chỉ có một lớp và chưa được huấn luyện: bạn đang kiểm tra cơ chế tính toán, không đánh giá chất lượng sinh văn bản.

KV Cache theo MHA ở đây là mảng NumPy lưu K đã xoay RoPE và V chưa xoay. Hình vẽ minh họa cả HBM/phân trang của hệ thống lớn; lab không có HBM thật, bộ cấp phát phân trang hay GPU.

Cách thực hiện và đọc output

1. Giữ nguyên code và bấm Chạy cell. Lượt đầu xử lý cả ba token được gọi là prefill. Ở 1A–1B, so vector trước/sau chuẩn hóa và trước/sau RoPE để thấy mã token đã trở thành dữ liệu attention như thế nào.

2. Ở 2A–2E, đọc hàng đầu của bảng trọng số head 0. Bạn sẽ thấy [1, 0, 0]: vị trí đầu chỉ nhìn chính nó. Đọc tiếp hàng cuối để thấy thông tin được lấy từ các vị trí hợp lệ theo những trọng số nào. Đây là trọng số trên vị trí đầu vào, không phải xác suất chọn từ trong bảng từ vựng.

3. Ở 3A–3B, xác định bộ K/V riêng của từng head. Cache có shape (2, 3, 2): hai head, ba token, mỗi vector hai số. Nhánh này cung cấp K/V cho attention, nên output in 3A/3B trước 2A–2E dù cột 3 nằm bên phải trong hình.

4. Ở 4A–4C, theo dõi vector được bổ sung qua hai đường residual, rồi đọc hàng cuối bảng logits. Logits là điểm số, không phải xác suất. Lab chọn mã có điểm cao nhất chỉ để tạo token đầu vào cho lượt tiếp.

5. Đọc lượt decode: token vừa chọn đi lại qua một lớp, lần này dùng cache đã có. Cache thành (2, 4, 2); ba vị trí cũ giữ nguyên, chỉ nối thêm một vị trí. Bảng logits mới có một hàng, dùng để dự đoán token sau token vừa xử lý.

Kết quả này chứng tỏ điều gì?

Kiểm tra causal thay token cuối rồi so hai hàng logits trước nó. Hai hàng không đổi cho thấy dữ liệu tương lai không lọt vào các vị trí trước trong demo này. Tổng mỗi hàng trọng số bằng 1 và các ô tương lai bằng 0 xác nhận bước chuẩn hóa/che mặt nạ làm đúng điều ta vừa quan sát.

Kiểm tra cache chạy hai cách với cùng weights và cùng chuỗi: chỉ xử lý token mới bằng cache, hoặc tính lại cả chuỗi. Sai khác logits không vượt 1e-9 cho thấy cách tái sử dụng cache giữ kết quả trong ngưỡng kiểm tra. Đây không phải phép đo tốc độ, cũng không chứng tỏ mô hình hiểu ngôn ngữ.

Kết quả mong đợiBạn sẽ thấy bảng logits kích thước (3, 6): ba hàng ứng với ba vị trí đầu vào, sáu cột ứng với sáu token ứng viên. Đọc hàng cuối để xem điểm cho token đi tiếp sau toàn bộ prompt. Sau đó lab xử lý token vừa chọn, so kết quả dùng cache với tính lại toàn chuỗi và in kết luận kiểm tra.

Tải transformer-decoder.py ↓

transformer-decoder.pyPYTHON

"""Lab Transformer: một lớp decoder-only chạy NumPy thuần, đồng bộ thuật ngữ
với khung kiến trúc tương tác AttentionArchitectureOverlay (mã 1A..4C).

Con đường kiến trúc lab đã chọn (phải nói rõ, không nói chung chung):
  - MHA (Multi-Head Attention): n_q = n_kv = 2, mỗi Q head có K/V riêng.
  - RoPE chuẩn cho từng head, chỉ xoay Q và K, theo VỊ TRÍ TUYỆT ĐỐI
    (= độ dài cache trước đó + chỉ số token trong pass hiện tại).
  - Pre-Norm RMSNorm, causal scaled dot-product attention (chia √d_head),
    output projection W_O, và đủ hai phép cộng residual.
  - FFN Dense SwiGLU (không MoE), cuối cùng là RMSNorm + LM head.

KV cache bắt chước cache thật: lưu K đã xoay RoPE và V chưa xoay, theo từng
head; decode nối 1 token mới vào cache thay vì tính lại token cũ. Đây là mảng
NumPy giáo dục: KHÔNG có HBM thật, block phân trang (PagedAttention) hay GPU.
GQA/MQA/MLA/MoE chỉ được nêu tên ở 3B/4B như các lựa chọn kiến trúc, KHÔNG
được chạy trong lab.
"""
import numpy as np

# ── Cấu hình demo: số nhỏ để giá trị có thể in thẳng ra màn hình ─────────────
VOCAB = 6           # bảng từ vựng demo
D_MODEL = 4         # chiều vector ẩn
N_HEADS = 2         # MHA: n_q = n_kv = 2
D_HEAD = 2          # chiều mỗi head (n_heads * d_head == d_model)
MAX_CONTEXT = 12
FFN_HIDDEN = 8      # chiều ẩn mạng Dense SwiGLU
ROPE_BASE = 10000.0

# weights chưa học: ngẫu nhiên có seed cố định, chạy lại cho ra cùng số.
rng = np.random.default_rng(4)
E = rng.normal(0, 0.25, (VOCAB, D_MODEL))                # bảng nhúng token
# MHA: MỖI Q head có W_Q / W_K / W_V riêng — không chia sẻ như GQA/MQA/MLA.
Wq = rng.normal(0, 0.30, (N_HEADS, D_MODEL, D_HEAD))
Wk = rng.normal(0, 0.30, (N_HEADS, D_MODEL, D_HEAD))
Wv = rng.normal(0, 0.30, (N_HEADS, D_MODEL, D_HEAD))
Wo = rng.normal(0, 0.30, (D_MODEL, D_MODEL))             # output projection
Wg = rng.normal(0, 0.30, (D_MODEL, FFN_HIDDEN))          # nhánh gate SwiGLU
Wz = rng.normal(0, 0.30, (D_MODEL, FFN_HIDDEN))          # nhánh up SwiGLU
Wf = rng.normal(0, 0.20, (FFN_HIDDEN, D_MODEL))          # chiếu về d_model
U = rng.normal(0, 0.20, (D_MODEL, VOCAB))                # LM head

DEMO_PROMPT = [1, 2, 3]        # Giữ làm mốc so sánh của bài học.
TRY_PROMPT = [1, 2, 4]         # Bài tập: chỉ đổi dòng này, dự đoán trước khi chạy.
DEMO_CONTINUATION = [4, 0, 2]  # token demo decode thêm (chỉ minh họa số học)

def rmsnorm(x):
    """RMSNorm không tham số (như cũ): giữ hướng, chỉ chuẩn hóa biên độ."""
    return x / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + 1e-6)

def rope_angles(positions):
    """Góc RoPE cho các vị trí tuyệt đối: [half, T]."""
    half = D_HEAD // 2
    freqs = ROPE_BASE ** (-2.0 * np.arange(half) / D_HEAD)
    return np.asarray(positions, dtype=float)[None, :] * freqs[:, None]

def rope_rotate(x, positions):
    """Xoay RoPE từng cặp nửa chiều của x: áp dụng cho Q và K, KHÔNG cho V.

x: [n_heads, T, d_head]; positions là vị trí tuyệt đối (đã cộng offset
    của cache), cho nên token decode thứ 4 xoay đúng góc của vị trí 3.
    """
    half = D_HEAD // 2
    ang = rope_angles(positions)                 # [half, T]
    cos = np.cos(ang).T[None, :, :]              # [1, T, half]
    sin = np.sin(ang).T[None, :, :]              # [1, T, half]
    x1, x2 = x[..., :half], x[..., half:]
    r1 = x1 * cos - x2 * sin
    r2 = x1 * sin + x2 * cos
    return np.concatenate([r1, r2], axis=-1)

def empty_cache():
    """KV cache trống theo MHA: K và V riêng cho từng head, chưa có token nào."""
    return {
        "k": np.zeros((N_HEADS, 0, D_HEAD)),
        "v": np.zeros((N_HEADS, 0, D_HEAD)),
    }

def attention_traced(ids, cache=None):
    """Tính một pass forward và trả về MỌI tensor trung gian (B=1 bị ẩn).

Thứ tự tính theo phụ thuộc số học (không theo cột vẽ trong sơ đồ):
    1A → 1B → 3A/3B → 2A → 2B → 2C → 2D → 2E → 4A → 4B → 4C.
    Cột 3 của sơ đồ là nhánh hỗ trợ attention: cache phải sẵn sàng TRƯỚC
    khi 2B tạo scores và 2E nhân với V, không phải bước chạy sau attention.
    """
    ids = np.asarray(ids, dtype=int)
    assert ids.ndim == 1 and len(ids) > 0
    assert np.all((ids >= 0) & (ids < VOCAB))
    if cache is None:
        cache = empty_cache()
    Kc, Vc = cache["k"], cache["v"]
    Tc = int(Kc.shape[1])
    assert Vc.shape[1] == Tc, "K cache và V cache lệch độ dài token"
    assert Kc.shape == (N_HEADS, Tc, D_HEAD) and Vc.shape == (N_HEADS, Tc, D_HEAD), (
        "cache phải có shape (n_heads, token, d_head) theo MHA"
    )
    T = len(ids)
    assert Tc + T <= MAX_CONTEXT, "Vượt demo ngữ cảnh tối đa"

# Vị trí tuyệt đối: offset theo độ dài cache, tức là 3A cung cấp "đã có bao nhiêu token".
    q_pos = np.arange(Tc, Tc + T)          # ví dụ decode: [3] thay vì [0]
    key_index = np.arange(Tc + T)

# 1A — Hidden State & RMSNorm (Pre-Norm)
    X = E[ids].astype(float)               # [T, d_model] (chiều B=1 đã ẩn)
    N = rmsnorm(X)

# 1B — Chiếu Q/K/V theo head (MHA: W riêng từng head)
    Q = np.einsum("td,hde->hte", N, Wq)    # [n_heads, T, d_head]
    Knew = np.einsum("td,hde->hte", N, Wk)
    Vnew = np.einsum("td,hde->hte", N, Wv)
    Qr = rope_rotate(Q, q_pos)             # RoPE trên Q
    Kr = rope_rotate(Knew, q_pos)          # RoPE trên K; V KHÔNG xoay

# 3A — KV Cache theo MHA: nối K/V mới vào cache (K đã xoay, V chưa xoay)
    K = np.concatenate([Kc, Kr], axis=1)   # [n_heads, Tc+T, d_head]
    V = np.concatenate([Vc, Vnew], axis=1)

# 2A/2B — mặt nạ nhân quả & scores chia √d_head (KHÔNG phải √d_model)
    scores = np.einsum("htd,hkd->htk", Qr, K) / np.sqrt(D_HEAD)   # [n_heads, Tq, Tkey]
    allow = key_index[None, :] <= q_pos[:, None]                  # j ≤ i (kể cả cache)
    masked = np.where(allow[None, :, :], scores, -np.inf)         # 2C — cộng mask

# 2D — softmax từng hàng (trọng số chú ý trên các token)
    expd = np.exp(masked - masked.max(axis=-1, keepdims=True))
    prob = expd / expd.sum(axis=-1, keepdims=True)                # [n_heads, Tq, Tkey]

# 2E — tổng hợp ngữ cảnh P·V (V giữ nguyên, không xoay)
    ctx = np.einsum("htk,hkd->htd", prob, V)                      # [n_heads, Tq, d_head]
    o = ctx.transpose(1, 0, 2).reshape(T, D_MODEL)                # ghép 2 head về d_model

# 4A — Output projection W_O + residual thứ nhất
    attn_out = o @ Wo
    H = X + attn_out

# 4B — RMSNorm + FFN Dense SwiGLU (không MoE) + residual thứ hai
    M2 = rmsnorm(H)
    gate = M2 @ Wg
    up = M2 @ Wz
    ffn = (gate / (1.0 + np.exp(-gate))) * up                     # silu(gate) · up
    F = H + ffn @ Wf

# 4C — RMSNorm cuối + LM head → logits (điểm số, KHÔNG phải xác suất)
    logits = rmsnorm(F) @ U                                       # [T, VOCAB]

return {
        "ids": ids, "positions": q_pos, "cache_past": Tc,
        "x": X, "n": N, "q": Q, "k_new": Knew, "v_new": Vnew,
        "q_rot": Qr, "k_rot": Kr,
        "k": K, "v": V, "scores": scores, "allow": allow,
        "masked": masked, "prob": prob, "ctx": ctx, "attn_out": attn_out,
        "h1": H, "gate": gate, "up": up, "ffn": ffn, "f2": F,
        "logits": logits, "cache": {"k": K, "v": V},
    }

def forward(ids, cache=None, trace=False):
    """Forward một lớp; im lặng mặc định, in chi tiết 1A..4C khi trace=True.

Trả về (logits, cache, shapes) — cache là dict {"k", "v"} theo head.
    """
    st = attention_traced(ids, cache)
    if trace:
        _print_pass(st, "Prefill" if st["cache_past"] == 0 else "Decode")
    return (
        st["logits"],
        st["cache"],
        {
            "embedding": st["x"].shape,
            "attention weights": st["prob"].shape,
            "MLP hidden": st["ffn"].shape,
            "logits": st["logits"].shape,
        },
    )

# ── Bộ in câu chuyện 1A..4C (formatter dùng chung, giữ output dễ đọc) ───────

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

def fmt_rows(a, nd=3, max_rows=None):
    a = np.asarray(a, dtype=float)
    rows = a if a.ndim > 1 else a[None, :]
    if max_rows:
        rows = rows[:max_rows]
    return ["  " + fmt_vec(r, nd) for r in rows]

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

def _print_pass(st, phase):
    """In lần lượt từng mã overlay cho một pass; giá trị đều là giá trị thật."""
    T = len(st["ids"])
    Tkey = st["k"].shape[1]
    print(f"--- {phase}: xử lý mã token {st['ids'].tolist()} tại vị trí {st['positions'].tolist()} ---")
    say("1A", "Hidden State x_t & RMSNorm (Pre-Norm) — Chuẩn bị vector", [
        "Quan sát: mã token chỉ dùng để tra bảng E. Vector mới là dữ liệu đi qua lớp.",
        f"shape x=({T}, {D_MODEL}) (nghĩa là {T} vị trí × d_model={D_MODEL}; batch B=1 bị ẩn).",
        "hàng 0 của X: " + fmt_vec(st["x"][0]),
        "hàng 0 sau RMSNorm: " + fmt_vec(st["n"][0]),
        "Kết luận: số chiều giữ nguyên; RMSNorm điều chỉnh biên độ, chưa trộn thông tin giữa token.",
        "Giữ X gốc để cộng residual ở 4A. Không dùng vector chuẩn hóa thay cho nhánh gốc này.",
    ])
    rope_idx = 1 if T > 1 else 0
    rope_pos = int(st["positions"][rope_idx])
    say("1B", "Chiếu Query/Key/Value (W_Q, W_K, W_V) & RoPE", [
        "Quan sát: một vector được chiếu thành Q để hỏi, K để đối chiếu, V làm nội dung cần lấy.",
        f"shape Q, K, V=({N_HEADS}, {T}, {D_HEAD}) (n_heads · n_tokens · d_head).",
        f"Q head 0, hàng 0 (vị trí {int(st['positions'][0])}): " + fmt_vec(st["q_rot"][0][0]),
        f"Q head 0, hàng {rope_idx} (vị trí {rope_pos} — đã xoay RoPE): " + fmt_vec(st["q_rot"][0][rope_idx])
        + " | chưa xoay: " + fmt_vec(st["q"][0][rope_idx]),
        "K head 0 sau RoPE, hàng đang xem: " + fmt_vec(st["k_rot"][0][rope_idx]),
        "V head 0, hàng 0: " + fmt_vec(st["v_new"][0][0]),
        f"Kết luận: RoPE xoay Q/K theo vị trí {st['positions'].tolist()}, V không xoay.",
        "Ở vị trí 0 góc xoay bằng 0; ở vị trí khác, hãy so số trước/sau để nhận ra thông tin vị trí.",
    ])
    say("3A", "KV Cache theo MHA", [
        f"shape K cache, V cache=({N_HEADS}, {Tkey}, {D_HEAD}) — head × token × d_head.",
        f"MHA: mỗi Q head có K/V riêng — Q0 → (K0, V0), Q1 → (K1, V1).",
        f"Độ dài token cache của pass này: {st['cache_past']} → {Tkey}.",
        "K head 0, vị trí cuối cache: " + fmt_vec(st["k"][0][-1]),
        "V head 0, vị trí cuối cache: " + fmt_vec(st["v"][0][-1]),
        "Quan sát: cache lưu K đã xoay RoPE và V chưa xoay của từng head.",
        ("Kết luận: prefill vừa lưu K/V của cả prompt; lượt decode sau có thể dùng lại chúng."
         if st["cache_past"] == 0 else
         f"Kết luận: giữ {st['cache_past']} vị trí cũ, chỉ tính rồi nối K/V của {T} token mới."),
        "Đây là mảng NumPy, không phải HBM hay bộ nhớ phân trang thật trên hình.",
    ])
    say("3B", "Tổ chức head: Lab Transformer chọn MHA", [
        f"n_q = n_kv = {N_HEADS} (tỉ lệ 1:1:1). GQA (g:1:1), MQA (n:1:1), MLA (latent) KHÔNG chạy trong lab.",
        "Quan sát hai head tại cùng vị trí 0:",
        "K0: " + fmt_vec(st["k"][0][0]) + " | K1: " + fmt_vec(st["k"][1][0]),
        "Kết luận: MHA là cách TỔ CHỨC HEAD; Q0 chỉ đọc K0/V0, Q1 chỉ đọc K1/V1.",
        "Cache ở 3A tái sử dụng THEO TOKEN qua các lượt; giữ token cũ không có nghĩa là chia sẻ giữa head.",
    ])
    say("2A", "Causal Mask (0 tại j≤i, -inf tại j>i)", [
        f"shape mask=({T}, {Tkey}); cột j là vị trí nguồn, hàng là token đang hỏi.",
        f"Hàng đầu đang ở vị trí tuyệt đối {int(st['positions'][0])}; chỉ được đọc j không lớn hơn vị trí này.",
        "Mask hàng đầu: " + ", ".join("0" if a else "-inf" for a in st["allow"][0]),
        "Kết luận: 0 nghĩa là cho phép, -inf nghĩa là chặn. Khi decode ở vị trí 3, các cột 0..3 đều hợp lệ.",
    ])
    say("2B", "Scaled Dot-Product Attention: Q · K^T / √d_head", [
        f"shape scores=({N_HEADS}, {T}, {Tkey}).",
        "scores head 0:"
        ] + fmt_rows(st["scores"][0]) + [
        "Cách đọc: ô hàng i, cột j là điểm đối chiếu Q của token i với K ở vị trí j, cùng head.",
        f"Chia √d_head = √{D_HEAD} ≈ {np.sqrt(D_HEAD):.3f} để kiểm soát độ lớn điểm.",
        "Kết luận: đây mới là điểm liên quan thô, có thể âm; chưa phải trọng số để gom nội dung.",
    ])
    say("2C", "Add Causal Mask (cộng mask vào scores)", [
        f"shape masked=({N_HEADS}, {T}, {Tkey}); head 0 sau khi cộng mask:"
        ] + fmt_rows(st["masked"][0]) + [
        "Ô hợp lệ (j≤i) cộng 0 nên giữ nguyên điểm; ô tương lai (j>i) cộng -∞ và sau softmax sẽ có xác suất 0.",
    ])
    say("2D", "Softmax Attention Weights", [
        f"shape P=({N_HEADS}, {T}, {Tkey}) — xác suất TRỌNG SỐ trên CÁC TOKEN, mỗi hàng tổng 1, không âm.",
        "P head 0:"
        ] + fmt_rows(st["prob"][0]) + [
        "Quan sát hàng đầu: khi đang ở vị trí 0, [1, 0, 0] nghĩa là chỉ lấy thông tin từ chính nó.",
        "Tổng hàng cuối head 0: " + f"{st['prob'][0][-1].sum():.6f}",
        "Kết luận: mỗi hàng phân bổ trọng số trên vị trí đầu vào, không phải trên sáu token ứng viên.",
        "Đây là trọng số chọn nguồn thông tin; điểm chọn token kế tiếp chỉ có ở 4C.",
    ])
    say("2E", "Context Aggregation: O = P · V", [
        f"shape mỗi head=({T}, {D_HEAD}); ghép 2 head thành ({T}, {D_MODEL}).",
        "O head 0, hàng 0: " + fmt_vec(st["ctx"][0][0]),
        "V head 0, vị trí nguồn đầu: " + fmt_vec(st["v"][0][0]),
        "Cách tính: nhân mỗi trọng số trong hàng P với V cùng cột, rồi cộng các vector lại.",
        "Kết luận: ở vị trí 0, trọng số [1, 0, 0] trả về đúng V đầu; các vị trí sau gom nhiều nguồn hợp lệ.",
    ])
    say("4A", "Output Projection W_O + Residual", [
        f"shape AttnOut=({T}, {D_MODEL}); H = X + AttnOut=({T}, {D_MODEL}).",
        "AttnOut hàng 0: " + fmt_vec(st["attn_out"][0]),
        "X gốc hàng 0: " + fmt_vec(st["x"][0]),
        "H sau cộng hàng 0: " + fmt_vec(st["h1"][0]),
        "Kết luận: W_O trộn kết quả hai head; residual cộng phần cập nhật vào X, không bỏ thông tin gốc.",
    ])
    say("4B", "RMSNorm + FFN (Dense SwiGLU, không MoE)", [
        f"shape gate/up=({T}, {FFN_HIDDEN}), đầu SwiGLU=({T}, {FFN_HIDDEN}); F = H + FFN=({T}, {D_MODEL}).",
        "gate hàng 0: " + fmt_vec(st["gate"][0]),
        "F sau cộng residual lần hai, hàng 0: " + fmt_vec(st["f2"][0]),
        "Quan sát: mỗi vị trí được chuẩn hóa rồi xử lý riêng bằng Dense SwiGLU, sau đó cộng vào H.",
        "Kết luận: attention đọc giữa các vị trí; FFN xử lý từng vector với cùng weights. Nhánh MoE không chạy.",
    ])
    say("4C", "Next Layer Hidden State / LM Head Logits", [
        f"shape logits=({T}, {VOCAB}) — MỖI VỊ TRÍ có một hàng điểm số cho {VOCAB} token.",
        "logits (3 chữ số, mỗi hàng một vị trí):"
        ] + fmt_rows(st["logits"]) + [
        "Cột theo thứ tự mã token: " + ", ".join(str(i) for i in range(VOCAB)),
        "Cách đọc: hàng cuối cho điểm của token đi sau toàn bộ chuỗi đang có; mỗi cột ứng với một mã ứng viên.",
        "Kết luận: logits là điểm số chưa chuẩn hóa (không phải xác suất). Không cần cộng chúng thành 1.",
        "Lab chỉ có một lớp nên chiếu LM head ngay ở đây; mô hình nhiều lớp phải đi hết stack trước.",
    ])

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

print("=== THEO DÕI BA TOKEN ĐI QUA MỘT LỚP TRANSFORMER ===")
print("MỤC TIÊU: giải thích token lấy thông tin ở đâu, tạo điểm dự đoán ra sao và dùng lại K/V thế nào.")
print(f"ĐẦU VÀO: {DEMO_PROMPT}, ba mã giả lập trong bảng từ vựng 0..5, không phải ba từ tiếng Việt.")
print("CÁCH HỌC: đọc câu Quan sát, xem số liệu rồi thử tự nói lại Kết luận của từng bước 1A..4C.")
print("Lab chọn MHA: mỗi Q head có K/V riêng, cùng FFN Dense SwiGLU và RoPE trên Q/K.")
print("Chỉ một chuỗi (B=1, ẩn trong shape), một lớp NumPy; weights chưa huấn luyện.")
print("Các số kiểm chứng cơ chế, không chứng minh khả năng ngôn ngữ hay tốc độ GPU.")
print("THỨ TỰ THỰC THI theo phụ thuộc số học (không theo cột vẽ):")
print("  1A → 1B → 3A/3B → 2A → 2B → 2C → 2D → 2E → 4A → 4B → 4C")
print("  Cột 3 là nhánh HỖ TRỢ attention: cache phải sẵn sàng TRƯỚC khi 2B/2E dùng K,V.")
print("")

# Prefill: in toàn bộ câu chuyện 1A..4C với giá trị thật.
prefill_logits, cache, prefill_shapes = forward(DEMO_PROMPT, trace=True)
assert prefill_shapes["logits"] == (len(DEMO_PROMPT), VOCAB)
print("")
print(f"argmax hàng cuối prefill → token demo {int(np.argmax(prefill_logits[-1]))} "
      "(chọn bằng số, không phải chất lượng ngôn ngữ).")
demo_token = int(np.argmax(prefill_logits[-1]))
print("")

# Decode đúng 1 token, mượn cache vừa ghi ở prefill.
past_before = int(cache["k"].shape[1])
assert cache["k"].shape == (N_HEADS, past_before, D_HEAD)
history_prompt_prefix = {key: value.copy() for key, value in cache.items()}
st_dec = attention_traced([demo_token], cache)
print(f"--- DECODE: đưa token vừa chọn ({demo_token}) trở lại mạng với cache {past_before} token ---")
print("Token thứ tư đã được chọn từ hàng cuối prefill. Lượt này xử lý nó để dự đoán token thứ năm.")
_print_pass(st_dec, "Decode")

# Logits theo từng vị trí của prompt + điểm next-token từ hàng cuối.
print(f"shape logits prefill: {prefill_logits.shape} — hàng cuối là next-token scores:")
print("  " + fmt_vec(prefill_logits[-1]))

# ── Kiểm tra số học (các forward nội bộ im lặng) ────────────────────────────
full_next_logits, _, _ = forward(list(DEMO_PROMPT) + [demo_token])
error = float(np.max(np.abs(st_dec["logits"][0] - full_next_logits[-1])))
assert error <= 1e-9
assert int(st_dec["cache"]["k"].shape[1]) == past_before + 1
# Prefix cache bất biến trước khi decode.
assert np.array_equal(st_dec["cache"]["k"][:, :past_before], history_prompt_prefix["k"])
assert np.array_equal(st_dec["cache"]["v"][:, :past_before], history_prompt_prefix["v"])
# Deterministic: chạy lại cùng prompt cho cùng logits/cache.
again_logits, again_cache, _ = forward(DEMO_PROMPT)
assert np.array_equal(again_logits, prefill_logits)
assert np.array_equal(again_cache["k"], cache["k"]) and np.array_equal(again_cache["v"], cache["v"])
# Causal mask: xác suất tương lai đúng bằng 0, mỗi hàng tổng 1 và không âm.
st3 = attention_traced(DEMO_PROMPT)
P = st3["prob"]
future = ~st3["allow"]
assert np.all(P[:, future] == 0.0)
assert np.all(P >= 0.0)
assert np.allclose(P.sum(axis=-1), 1.0, atol=1e-12)
# Hàng đầu mỗi head chỉ được nhìn chính nó.
assert np.allclose(P[:, 0, 0], 1.0) and np.all(P[:, 0, 1:] == 0.0)
# Thay token tương lai không đổi logits trước đó.
base, _, _ = forward(DEMO_PROMPT)
changed, _, _ = forward([1, 2, 4])
assert np.allclose(base[:2], changed[:2], atol=1e-10)
# Decode nhiều token (vị trí khác 0) qua cache khớp tính lại toàn bộ — im lặng.
full_ids = list(DEMO_PROMPT) + [demo_token]
cache = st_dec["cache"]   # bắt đầu từ cache sau pass decode demo (4 token)
for token in DEMO_CONTINUATION:
    full_ids.append(token)
    full_logits, _, _ = forward(full_ids)
    cached_logits, cache, _ = forward([token], cache)
    step_error = float(np.max(np.abs(cached_logits[-1] - full_logits[-1])))
    assert step_error <= 1e-9, f"decode lệch full recompute tại token {token}: {step_error}"
    assert int(cache["k"].shape[1]) == len(full_ids)

print("\nTHỬ NGHIỆM: Token tương lai có làm đổi quá khứ không?")
print(f"Mốc so sánh: {DEMO_PROMPT}; prompt bạn thử: {TRY_PROMPT}")
trial_logits, _, _ = forward(TRY_PROMPT)
if len(TRY_PROMPT) == len(DEMO_PROMPT):
    changed_positions = np.flatnonzero(np.asarray(TRY_PROMPT) != np.asarray(DEMO_PROMPT))
    first_changed = int(changed_positions[0]) if len(changed_positions) else len(DEMO_PROMPT)
    row_errors = np.max(np.abs(trial_logits - prefill_logits), axis=1)
    for i, delta in enumerate(row_errors):
        expectation = "phải giữ nguyên" if i < first_changed else "được phép đổi"
        print(f"  Hàng {i}: lệch lớn nhất = {delta:.3e}; dự đoán: {expectation}.")
    assert np.all(row_errors[:first_changed] <= 1e-9)
    print("Kết luận: các vị trí trước token đầu tiên bị sửa không đọc được thay đổi đó.")
    print("Được phép đổi không có nghĩa là bắt buộc đổi: hãy phân biệt quy tắc mask với giá trị cụ thể.")
else:
    print("Để so từng hàng, hãy dùng TRY_PROMPT có cùng số token với DEMO_PROMPT.")

print("\nTỔNG KẾT: Bạn vừa quan sát được gì?")
print("1. Causal mask cho phép tính cả prompt mà token trước không nhìn token sau.")
print("2. Attention weights chọn nguồn thông tin; logits ở LM head chấm điểm token ứng viên.")
print(f"3. Dùng cache và tính lại cả chuỗi có độ lệch logits lớn nhất {error:.3e} (ngưỡng 1e-9).")
print("   Các vị trí cũ trong cache giữ nguyên; cùng weights và input cho cùng kết quả.")
print("Các kiểm tra này xác nhận trường hợp demo, không chứng minh model hiểu ngôn ngữ hoặc tăng tốc bao nhiêu.")
print("PASS: decoder một lớp MHA + RoPE + SwiGLU chạy đúng và khớp 12 mã overlay kiến trúc.")

Python outputSẵn sàng.

Bài tập 1 · Token tương lai có làm đổi quá khứ không?

Trước khi chạy, dự đoán hàng logits nào được phép đổi khi thay token cuối của [1, 2, 3] thành 4. Trong code, tìm TRY_PROMPT và đặt [1, 2, 4], rồi chạy lại. Đọc phần THỬ NGHIỆM ở cuối output.

Đối chiếu sau khi tự trả lời

Hai hàng đầu phải giữ nguyên trong sai số; hàng cuối được phép đổi. Vì causal mask chặn nhìn tương lai, vị trí 0 và 1 không đọc token ở vị trí 2. Nếu đổi token đầu, cả ba hàng đều có thể đổi vì cả ba vị trí được phép đọc nó.

Bài tập 2 · Đọc đúng bảng điểm dự đoán

Ở bước 4C của lượt prefill, hãy tự tìm cột có điểm lớn nhất ở hàng cuối. Cột đó tương ứng mã token nào? Đối chiếu với token demo mà lab chọn và giải thích vì sao không đọc hàng đầu để sinh tiếp toàn bộ prompt.

Đối chiếu sau khi tự trả lời

Chỉ số cột lớn nhất là mã token được chọn bằng argmax trong lab. Hàng đầu chỉ dùng thông tin tới vị trí đầu; hàng cuối dùng cả ba token hợp lệ nên cho điểm của token đi sau toàn bộ prompt. Điểm lớn nhất không có nghĩa là từ đó đúng về ngôn ngữ vì weights chưa học.

Bài tập 3 · MHA và KV cache giải quyết hai việc gì?

Không chạy thêm code, hãy dự đoán shape cache khi đã xử lý tổng cộng năm token trong cấu hình này. Giải thích vì sao vẫn cần hai bộ K/V dù cache giúp tránh tính lại token cũ.

Đối chiếu sau khi tự trả lời

Mỗi tensor K hoặc V có shape (2, 5, 2): hai head, năm token, mỗi vector hai số. MHA giữ K/V riêng cho từng Q head; cache chỉ giữ lại kết quả của các token đã xử lý, không biến MHA thành GQA hay MQA.