Mixture of Experts và các kiến trúc thay thế: router, expert và Mamba

Bài thứ mười ba kết thúc bằng một lời hứa: với Mixture of Experts, mọi phép tính roofline phải làm lại. Bài này làm phép tính đó, và đi tiếp sang một hướng khác cũng đang thay đổi kiến trúc LLM: bỏ bớt attention.

Câu hỏi của bài: một model dense trả những khoản chi phí nào cho mỗi token, và có cách nào tránh bớt chúng không? Có hai khoản. Khoản thứ nhất là tham số: mỗi token phải nhân với toàn bộ weights, nên model càng lớn thì mỗi token càng đắt. Khoản thứ hai là KV cache: mỗi token mới phải đọc lại K, V của mọi token trước, và cache này lớn dần theo độ dài context. Bài 13 đã đo cả hai khoản: decode ở batch 1 làm khoảng 1 FLOP cho mỗi byte đọc vào, và một sequence 128k token của Llama 3.1 8B tốn 16 GiB KV cache.

Hai hướng trong bài này nhắm vào hai khoản đó:

• Mixture of Experts (MoE) nhắm vào khoản thứ nhất. Model có rất nhiều tham số, nhưng mỗi token chỉ dùng một phần nhỏ. DeepSeek-V3 có 671B tham số, mỗi token chỉ chạy khoảng 37B.
• SSM, Mamba và kiến trúc hybrid nhắm vào khoản thứ hai. Phần lớn layer attention được thay bằng một state có kích thước cố định, nên bộ nhớ không còn tăng theo độ dài context.

Nửa đầu bài đi theo Hình 1 ở đầu trang: một token đi qua layer MoE như thế nào. Nửa sau đi vào SSM và hybrid. Bài 1, mục 7.8 đã giới thiệu MoE ở mức “là gì”: thay FFN bằng nhiều expert, tổng tham số quyết định VRAM, tham số kích hoạt quyết định compute. Bài này không nhắc lại phần định nghĩa đó mà đi vào bên trong.

01Hai khoản chi phí của model dense

Nhớ lại phép đếm của bài 13, mục 2: một lượt forward tốn khoảng 2N FLOPs cho mỗi token và phải đọc toàn bộ 2N byte weights (ở bf16). Với một model dense, N là tất cả tham số. Muốn model biết nhiều hơn thì tăng N, và mỗi token đắt lên đúng theo tỉ lệ đó. Đây là khoản thứ nhất.

Khoản thứ hai nằm ở attention. Token mới phải nhìn lại mọi token trước, nên model giữ K, V của chúng trong KV cache. Bài 13, mục 9 đã chỉ ra context dài đắt chủ yếu ở đây: cache tăng tuyến tính theo độ dài, phải đọc lại ở mọi bước decode, và chiếm chỗ làm giảm batch tối đa.

Hai khoản nằm ở hai nửa khác nhau của một Transformer block:

Con số 81% lấy từ bài 8, mục 6: với GQA và SwiGLU, ba ma trận của FFN chiếm phần lớn tham số mỗi block. Đó là lý do MoE nhắm vào FFN chứ không phải attention.

Hai hướng kết hợp được với nhau: Jamba, Qwen3-Next và Nemotron 3 vừa là MoE vừa là hybrid (mục 9).

02Từ FFN đến expert

Vì sao chia FFN

Bài 8, mục 3 đã giải thích một cách nhìn về FFN: hàng của ma trận thứ nhất là các “key” bắt mẫu trong residual stream, cột của ma trận thứ hai là các “value” đẩy thông tin trở lại. Nếu FFN giống một bộ nhớ key–value, thì một token cụ thể chỉ khớp với một phần nhỏ các key. Phần còn lại nhân ra gần bằng không mà vẫn tốn compute.

MoE biến quan sát đó thành kiến trúc. Thay một FFN lớn bằng E FFN nhỏ hơn gọi là expert, và thêm một router quyết định token nào đi qua expert nào. Mỗi token chỉ qua k expert, với k nhỏ hơn E rất nhiều.

Cấu trúc một layer MoE

Theo thứ tự trong Hình 1:

1. Attention chạy như model dense. MoE không đụng vào attention, và KV cache giữ nguyên.
2. Router là một lớp tuyến tính W_r kích thước d × E. Nó nhân hidden state x của token với W_r để ra E con số, mỗi expert một con số, rồi lấy softmax.
3. Top-k và gate: giữ k expert có xác suất cao nhất. Trọng số của chúng gọi là gate, thường được chuẩn hoá lại để tổng bằng 1.
4. Dispatch: gửi token tới k expert đã chọn. Khi expert nằm trên GPU khác, bước này là một phép truyền dữ liệu giữa các GPU (mục 6).
5. Expert FFN: mỗi expert là một FFN bình thường, ví dụ SwiGLU như bài 8, mục 4.
6. Combine và residual: cộng output của k expert theo trọng số gate, rồi cộng vào residual stream.

Viết gọn thành công thức, theo cách Mixtral trình bày (Jiang et al. 2024, mục 2.1):

y = Σ_{i ∈ TopK} G(x)_i · E_i(x),     G(x) = Softmax(TopK(x · W_r))

TopK đặt logit của các expert ngoài top-k thành −∞, nên sau softmax chúng có trọng số 0 và không cần tính.

Lịch sử ngắn

Phiên bản MoE dùng trong LLM bắt đầu với Shazeer et al. 2017: một layer MoE giữa các layer LSTM, router là noisy top-k gating G(x) = Softmax(KeepTopK(H(x), k)), với H(x) là logit cộng nhiễu Gauss (mục 2.1, phương trình 3–5), và tới 137 tỉ tham số chỉ riêng trong layer MoE (mục 5.2). Ba mốc sau đó:

• GShard (Lepikhin et al. 2020) đưa MoE vào Transformer với top-2 và expert capacity, giới hạn số token mỗi expert nhận (mục 2.2); model lớn nhất hơn 600B tham số (abstract).
• Switch Transformer (Fedus, Zoph & Shazeer 2021) rút xuống top-1 cho router rẻ hơn và truyền dữ liệu đơn giản hơn (mục 2.1); Switch-C có 1,6 nghìn tỉ tham số (mục 5.6).
• Mixtral 8×7B (Jiang et al. 2024) đưa MoE vào model mở phổ biến: 8 expert, top-2, mỗi token “thấy” 47B tham số nhưng chỉ dùng 13B (abstract).

Từ DeepSeekMoE trở đi, xu hướng là nhiều expert nhỏ, k lớn, kèm expert dùng chung (mục 5).

03Bên trong router

Chọn theo từng token, ở từng layer

Router áp lên hidden state của một token ở một layer. Nó không nhìn cả câu, cũng không chọn một lần cho cả request. Hai token liền nhau có thể đi qua hai bộ expert khác nhau, và cùng một token ở layer 3 và layer 30 cũng vậy, vì mỗi layer có router và expert riêng.

Vì vậy “expert số 5” không phải một khối kiến thức chung cho cả model. Mixtral 8×7B có 32 layer × 8 expert = 256 expert, mỗi cái chỉ là một FFN trong một layer cụ thể. Mục 7 xem chúng thực sự học gì.

Top-1, top-2 và vì sao k không phải 1 là phổ biến

Với top-1, nếu chuẩn hoá gate trên một expert thì gate luôn bằng 1 và router không còn nhận gradient qua đường đó, nên Switch giữ gate là xác suất gốc. Với k từ 2 trở lên, router có thêm tín hiệu từ việc so các expert được chọn với nhau. Mixtral, gpt-oss và các model DeepSeek đều dùng k ≥ 2.

Expert choice: đảo chiều chọn

Mọi cách ở trên đều là token chọn expert. Expert choice routing (Zhou et al. 2022) làm ngược lại: mỗi expert chọn k token có điểm cao nhất trong batch, với k = n × c / E, trong đó c là capacity factor (mục 3.2, phương trình 1). Cân bằng tải đạt được theo thiết kế, vì mọi expert nhận đúng bấy nhiêu token. Paper báo thời gian hội tụ khi train nhanh hơn hơn 2 lần so với GShard top-2 (mục 4.2, Hình 2a).

Cái giá nằm ở mục giới hạn của paper: phương pháp “có thể không áp dụng ngay cho sinh văn bản autoregressive”, vì việc chọn top-k cần thấy cả token tương lai trong batch (mục 6). Khi decode chỉ có một token mới mỗi sequence, nên các LLM decoder vẫn dùng token chọn expert.

Code mẫu 2a: một layer MoE bằng PyTorch thuần

Code dưới là một layer MoE đầy đủ các bước của Hình 1, kể cả expert capacity. Phần cuối cho một token của “request của bạn” đi qua cùng một layer trong hai batch khác nhau.

# Python 3.12, torch 2.14.1 — chạy trên CPU, dưới 1 giây
import torch, torch.nn as nn

class MoE(nn.Module):
    """Thay một FFN bằng E expert; router chọn top-k expert cho từng token."""
    def __init__(self, d, d_ff, n_experts=8, k=2, capacity_factor=None):
        super().__init__()
        self.k, self.E, self.cf = k, n_experts, capacity_factor
        self.router = nn.Linear(d, n_experts, bias=False)            # W_r: d → E
        self.experts = nn.ModuleList(
            nn.Sequential(nn.Linear(d, d_ff), nn.SiLU(), nn.Linear(d_ff, d)) for _ in range(n_experts))
        self.register_buffer("bias", torch.zeros(n_experts))        # dùng ở code 2b, không có gradient

    def forward(self, x):                                            # x: (T, d), T token của cả batch
        probs = self.router(x).softmax(-1)                           # (T, E)
        _, idx = (probs + self.bias).topk(self.k, dim=-1)            # bias chỉ ảnh hưởng việc CHỌN expert
        gate = probs.gather(-1, idx)                                 # trọng số lấy từ xác suất gốc
        if self.k > 1:
            gate = gate / gate.sum(-1, keepdim=True)                 # chuẩn hoá lại trên k expert (kiểu Mixtral)
        capacity = None if self.cf is None else int(self.cf * x.shape[0] * self.k / self.E)
        out, dropped = torch.zeros_like(x), 0
        for e, expert in enumerate(self.experts):
            tok, slot = (idx == e).nonzero(as_tuple=True)            # token nào chọn expert e
            if capacity is not None and len(tok) > capacity:         # expert đầy: token xếp sau bị bỏ
                dropped += len(tok) - capacity
                tok, slot = tok[:capacity], slot[:capacity]
            if len(tok):
                out[tok] += gate[tok, slot, None] * expert(x[tok])
        self.last = dict(probs=probs, idx=idx, dropped=dropped)
        return x + out                                               # residual: token bị drop chỉ còn đường này

if __name__ == "__main__":
    torch.manual_seed(0)
    moe = MoE(d=16, d_ff=64, n_experts=8, k=2).eval()
    mine = torch.randn(1, 16)                                        # một token của "request của bạn"
    with torch.no_grad():
        ref = moe(mine)                                              # không giới hạn capacity: kết quả chuẩn
        e1, e2 = moe.last["idx"][0].tolist()
        p1, p2 = moe.last["probs"][0, [e1, e2]].tolist()
        print(f"token của bạn chọn expert {e1} (p={p1:.3f}) và expert {e2} (p={p2:.3f})")
        others_a = torch.randn(15, 16)                               # batch A: 15 token ngẫu nhiên
        others_b = torch.randn(15, 16) + 3 * moe.router.weight[e1]   # batch B: nhiều token cũng "thích" expert e1
        moe.cf = 1.25                                                # capacity = 1,25 × 16·2/8 = 5 slot mỗi expert
        for name, others in [("batch A", others_a), ("batch B", others_b)]:
            y = moe(torch.cat([others, mine]))[-1:]                  # token của bạn đứng cuối hàng
            load = torch.bincount(moe.last["idx"].flatten(), minlength=8).tolist()
            print(f"{name}: tải mỗi expert {load}, token bị drop {moe.last['dropped']:2d}, "
                  f"output của bạn lệch {(y - ref).abs().max():.4f}")

Output:

token của bạn chọn expert 3 (p=0.267) và expert 2 (p=0.139)
batch A: tải mỗi expert [4, 2, 3, 5, 3, 7, 6, 2], token bị drop  3, output của bạn lệch 0.0000
batch B: tải mỗi expert [2, 0, 4, 15, 0, 7, 3, 1], token bị drop 12, output của bạn lệch 0.2768

Ba điều đọc được:

• Router chưa train đã chọn không đều. Ngay ở batch A ngẫu nhiên, expert 5 nhận 7 lượt chọn trong khi expert 1 và 7 chỉ nhận 2. Đó là điểm khởi đầu của bài toán ở mục 4.
• Cùng một token, cùng weights, output khác nhau tuỳ batch. Ở batch A, expert 3 nhận đúng 5 token, vừa đủ capacity, nên token của bạn được xử lý đầy đủ và output khớp kết quả chuẩn. Ở batch B, 14 token khác cũng chọn expert 3. Token của bạn đứng cuối hàng nên bị drop khỏi expert 3, chỉ còn đóng góp của expert 2 và đường residual.
• Capacity tính theo cả batch: capacity_factor × T × k / E, với T là tổng số token của mọi request trong batch. Mục 4 quay lại điểm này.

04Load balancing

Router sụp

Để router tự học, nó có xu hướng dồn token vào vài expert. Shazeer et al. 2017 mô tả: router “có xu hướng hội tụ về trạng thái luôn cho trọng số lớn cho cùng vài expert”, và sự mất cân bằng này tự củng cố, vì expert được ưu tiên được train nhanh hơn nên lại càng được chọn nhiều hơn (mục 4).

Vòng lặp đó nằm ở hàng trên cùng của Hình 2. Mất cân bằng gây hại theo hai cách:

1. Lãng phí tham số. Expert ít được dùng là tham số nằm trong VRAM mà gần như không đóng góp. Model 8 expert sụp về 2 expert thì thực chất chỉ là một model nhỏ hơn.
2. Chậm khi chạy song song. Khi expert nằm trên nhiều GPU (mục 6), GPU giữ expert “nóng” phải xử lý nhiều token hơn, và cả bước phải chờ GPU chậm nhất. Đây là lý do cân bằng tải không chỉ là chuyện “công bằng” giữa các expert.

Lớp 1: auxiliary loss

Cách phổ biến nhất là cộng thêm một loss phạt mất cân bằng. Switch Transformer dùng dạng (mục 2.2, phương trình 4–6):

loss_aux = α · N · Σ_{i=1..N} f_i · P_i

• f_i là tỉ lệ token trong batch được gửi tới expert i. Đây là phép đếm, không có gradient.
• P_i là xác suất router trung bình dành cho expert i trên cả batch. Đây là phần có gradient.
• N là số expert. Khi tải đều tuyệt đối, f_i = P_i = 1/N và loss bằng α. Nhân với N giữ cho giá trị loss không đổi khi đổi số expert.
• Paper dùng α = 10⁻² xuyên suốt, sau khi quét từ 10⁻¹ đến 10⁻⁵.

Vì f_i không có gradient, loss chỉ đẩy P_i xuống ở expert đang nhận nhiều token. GShard dùng dạng gần giống (mục 2.2, Thuật toán 1); Shazeer et al. 2017 dùng importance loss và load loss (mục 4, phụ lục A).

Lớp 2: router z-loss

Bài toán thứ hai không phải cân bằng mà là ổn định. ST-MoE (Zoph et al. 2022) thấy rằng ở quy mô lớn, chỉ tính router bằng float32 như Switch đề xuất là chưa đủ để train ổn định (mục 3.3). Router có hàm mũ trong softmax, và hàm mũ khuếch đại sai số làm tròn: paper lấy ví dụ một sai số 0,5 ở bfloat16 có thể làm output softmax lệch 36% (chú thích 7, mục 3.4). Giải pháp là router z-loss:

L_z = (1/B) · Σ_{i=1..B} ( log Σ_{j=1..N} exp(x_j⁽ⁱ⁾) )²

với B là số token và x là logit vào router (mục 3.3, phương trình 5). Loss này phạt logit lớn, giữ chúng ở vùng mà sai số làm tròn ít ảnh hưởng. Hệ số dùng là c_z = 0,001 (phương trình 6). Trong thử nghiệm của paper, z-loss cho 3/3 lần chạy ổn định so với 4/6 ở baseline, không giảm chất lượng (Bảng 4). Model lớn nhất của ST-MoE có 269B tham số, compute tương đương một model dense 32B (abstract).

Lớp 3: bias, không cần loss phụ

Auxiliary loss có một cái giá: nó là gradient cộng thêm vào gradient của LM loss, và hai mục tiêu có thể kéo nhau. Wang et al. 2024 viết rằng auxiliary loss lớn “đưa vào các gradient nhiễu không thể bỏ qua”, làm giảm chất lượng model (abstract). Loss nhỏ thì nhiễu ít, nhưng cân bằng kém.

Giải pháp của họ, auxiliary-loss-free balancing, tách hẳn việc cân bằng ra khỏi gradient:

• Mỗi expert có một bias b_i. Router chọn top-k theo s_i + b_i, nhưng gate vẫn lấy từ s_i gốc. Paper nhấn mạnh bias “chỉ dùng để điều chỉnh việc chọn top-K”, không cộng vào trọng số output (mục 3, phương trình 3).
• Sau mỗi batch, đếm số token mỗi expert nhận. Expert thừa tải bị giảm bias, expert thiếu tải được tăng: b_i ← b_i + u · sign(ē − c_i), với ē là tải trung bình (Thuật toán 1). Paper chọn u = 0,001 sau khi thử: 0,0001 hội tụ chậm, 0,01 dao động (mục 4.3).

Kết quả trên model 1B và 3B (mục 4.2, Bảng 2), so với baseline dùng auxiliary loss:

MaxVio đo mức lệch của expert tải nặng nhất so với tải trung bình (phương trình 4). Cân bằng tốt hơn hẳn mà perplexity còn giảm nhẹ.

DeepSeek-V3 dùng đúng cơ chế này ở quy mô 671B (mục 2.1.2). Router tính điểm bằng sigmoid thay cho softmax, rồi chuẩn hoá lại trên các expert được chọn. Tốc độ cập nhật bias γ = 0,001 trong 14,3T token đầu, rồi đặt về 0 cho 500B token cuối. Họ vẫn giữ một sequence-wise balance loss với α = 0,0001, rất nhỏ, chỉ để tránh mất cân bằng cực đoan trong một sequence (mục 4.2).

Code mẫu 2b: thí nghiệm cân bằng tải

Code dưới train layer MoE ở code 2a (top-1, 8 expert) trên một bài hồi quy tổng hợp, với bốn cấu hình: không cân bằng, auxiliary loss kiểu Switch ở hai mức α, và bias không cần loss. Mọi seed được cố định. Trước khi chạy, lưu code 2a thành file code2a.py cùng thư mục, vì code 2b import class MoE từ đó.

# Python 3.12, torch 2.14.1 — dùng lại class MoE ở code 2a, chạy trên CPU khoảng một phút
import torch, torch.nn.functional as F
from code2a import MoE

E, d = 8, 16
g0 = torch.Generator().manual_seed(42)
W1, W2 = torch.randn(d, 64, generator=g0) / d**0.5, torch.randn(64, d, generator=g0) / 8
def make_data(n, g):                                   # hồi quy tổng hợp: y = tanh(x·W1)·W2
    x = torch.randn(n, d, generator=g)
    return x, torch.tanh(x @ W1) @ W2

def train(mode, alpha=0.0, u=1e-3, steps=1500):
    torch.manual_seed(0)
    g = torch.Generator().manual_seed(1)
    moe = MoE(d, 64, E, k=1)                           # top-1 như Switch Transformer
    opt = torch.optim.Adam(moe.parameters(), lr=3e-3)
    for _ in range(steps):
        x, y = make_data(256, g)
        loss = F.mse_loss(moe(x) - x, y)
        load = torch.bincount(moe.last["idx"].flatten(), minlength=E).float()
        if alpha:                                      # Switch: α · E · Σ f_i · P_i
            f = load / load.sum()                      # tỉ lệ token gửi tới expert i (đếm, không có gradient)
            P = moe.last["probs"].mean(0)              # xác suất router trung bình (có gradient)
            loss = loss + alpha * E * (f * P).sum()
        opt.zero_grad(); loss.backward(); opt.step()
        if mode == "bias":                             # không loss phụ: chỉnh bias ngoài gradient
            moe.bias += u * torch.sign(load.mean() - load)  # thiếu tải thì tăng, thừa tải thì giảm
    with torch.no_grad():
        x, y = make_data(4096, torch.Generator().manual_seed(7))
        mse = F.mse_loss(moe(x) - x, y).item()
        load = torch.bincount(moe.last["idx"].flatten(), minlength=E).float()
    share = 100 * load / load.sum()
    print(f"{mode:15} tải %: {' '.join(f'{s:4.1f}' for s in share)} | "
          f"MaxVio {load.max() / load.mean() - 1:.2f} | MSE {mse:.4f}")

train("không cân bằng")
train("aux α=0.01", alpha=0.01)
train("aux α=0.05", alpha=0.05)
train("bias", u=1e-3)

Output:

không cân bằng  tải %:  3.8 10.2 15.6 14.1 18.9 18.9 12.8  5.6 | MaxVio 0.52 | MSE 0.0197
aux α=0.01      tải %:  7.3 12.8 14.4 13.5 13.8 15.4 13.2  9.6 | MaxVio 0.23 | MSE 0.0191
aux α=0.05      tải %: 13.3 14.0 13.0 12.0 11.0 13.1 12.7 10.9 | MaxVio 0.12 | MSE 0.0188
bias            tải %: 13.8 13.8 12.6 11.3 10.3 13.6 12.5 12.2 | MaxVio 0.10 | MSE 0.0198

Tải đều tuyệt đối là 12,5% mỗi expert.

• Không cân bằng: expert ít nhất nhận 3,8% token, expert nhiều nhất 18,9%, gấp 1,52 lần mức trung bình (MaxVio 0,52).
• Auxiliary loss: α = 0,01 kéo MaxVio về 0,23, α = 0,05 về 0,12.
• Bias: đạt MaxVio 0,10 mà không có thành phần nào trong loss.
• MSE gần như không đổi giữa bốn cấu hình. Ở bài toán nhỏ này, cân bằng tải không làm giảm chất lượng.

Đây là thí nghiệm đồ chơi một layer, nên mức sụp nhẹ hơn nhiều so với model thật, và quá nhỏ để thấy cái giá “gradient nhiễu” mà Wang et al. đo ở model 1B và 3B. Điều nó cho thấy là cơ chế: bias đạt cân bằng mà không đụng vào loss.

Capacity factor và token dropping

Ngay cả khi cân bằng trung bình tốt, từng batch vẫn dao động. Để tính toán có kích thước cố định (quan trọng với TPU và với kernel GPU), các hệ thống thời Switch đặt một giới hạn cứng:

expert capacity = (số token mỗi batch / số expert) × capacity factor

(Switch, mục 2.2, phương trình 3.) Token vượt mức bị drop: “bỏ qua tính toán, và biểu diễn của token được chuyển thẳng sang layer sau qua residual connection”, nhánh đỏ trong Hình 1. Switch báo tỉ lệ drop thường dưới 1%; ST-MoE khuyến nghị capacity factor 1,25 khi train (mục 5). Như code 2a cho thấy, token có bị drop hay không phụ thuộc vào các token khác cùng batch, mà trong hệ thống phục vụ nhiều người thì đó là token của request khác.

Nối lại giả thuyết ở bài 9

Bài 9, mục 7 có nhắc một giả thuyết: với model MoE có giới hạn sức chứa, việc token được route tới đâu có thể phụ thuộc vào các token khác trong batch, và đó có thể là một nguồn không tất định của các API. Phần trên giải thích cơ chế của giả thuyết đó.

Nhưng nó vẫn chỉ là giả thuyết. Code 2a là layer đồ chơi với capacity đặt chặt có chủ đích, không phải bằng chứng về hệ thống của nhà cung cấp nào, và chưa nhà cung cấp lớn nào xác nhận công khai. Hai lý do để thận trọng:

• Không phải hệ thống nào cũng drop token. DeepSeek-V3 viết rõ model không drop token khi train, và nhờ chiến lược triển khai riêng để cân bằng tải lúc inference, cũng không drop token khi inference (mục 2.1.2).
• Không có drop thì vẫn có những nguồn không tất định khác đã biết, như thứ tự cộng số thực trong kernel phụ thuộc kích thước batch mà bài 9 đã đo.

05Đếm tham số: tổng và kích hoạt

Fine-grained expert và shared expert

Mixtral dùng 8 expert lớn cỡ FFN của một model 7B. DeepSeekMoE (Dai et al. 2024) đề xuất hai thay đổi mà gần như mọi model MoE lớn sau đó đều theo.

Fine-grained expert segmentation (mục 3.1): chia mỗi expert thành m expert nhỏ (chiều trung gian giảm còn 1/m) và chọn nhiều hơn m lần. Compute giữ nguyên nhưng số tổ hợp tăng vọt: 16 expert chọn 2 cho 120 tổ hợp, còn 64 expert chọn 8 cho 4.426.165.368 tổ hợp.

Shared expert isolation (mục 3.2): tách K_s expert làm shared expert, luôn chạy cho mọi token, để kiến thức chung khỏi phải học lặp lại ở các routed expert.

So sánh cấu hình MoE của vài model, đọc từ config.json:

Hướng đi rất rõ: từ 8 expert lớn sang hàng trăm expert nhỏ và k lớn hơn.

Code mẫu 1: đếm tham số từ header safetensors

Cách đáng tin nhất để biết tổng và tham số kích hoạt là đếm thẳng từ checkpoint. File safetensors có header ghi tên, kiểu và kích thước của mọi tensor. huggingface_hub đọc được header này mà không tải weights. Code dưới phân mọi tensor vào ba nhóm: routed expert (tên có .experts.), dense (attention, embedding, norm, router, shared expert), và MTP (module multi-token prediction, không dùng khi sinh văn bản bình thường). Tham số kích hoạt mỗi token bằng phần dense cộng k/E của phần routed expert.

# Python 3.12, huggingface_hub 1.33.0 — chỉ tải config.json và header safetensors, chạy trên CPU
import json, math, os, re
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
from huggingface_hub import hf_hub_download, get_safetensors_metadata

MODELS = {
    "Llama 3.1 8B":     "unsloth/Meta-Llama-3.1-8B",          # dense, mirror không gated
    "Mixtral 8x7B":     "mistralai/Mixtral-8x7B-v0.1",
    "Qwen3-30B-A3B":    "Qwen/Qwen3-30B-A3B-Base",
    "gpt-oss-20b":      "openai/gpt-oss-20b",
    "gpt-oss-120b":     "openai/gpt-oss-120b",
    "Qwen3-Next-80B":   "Qwen/Qwen3-Next-80B-A3B-Instruct",
    "DeepSeek-V3":      "deepseek-ai/DeepSeek-V3",
    "Jamba2 Mini":      "ai21labs/AI21-Jamba2-Mini",          # hybrid Mamba + attention + MoE
    "Nemotron 3 Nano":  "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",  # hybrid Mamba-2 + attention + MoE
}

def n_values(t):
    n = math.prod(t.shape)
    return n * 2 if t.dtype == "U8" else n          # MXFP4 của gpt-oss: 2 giá trị 4-bit mỗi byte

print(f"{'model':15} {'E':>4} {'k':>2} {'shared':>6} {'total B':>7} {'active B':>8} {'act-emb B':>9} {'expert %':>8} {'MTP B':>5}")
for name, repo in MODELS.items():
    cfg = json.load(open(hf_hub_download(repo, "config.json")))
    E = cfg.get("num_local_experts") or cfg.get("num_experts") or cfg.get("n_routed_experts") or 1
    k = cfg.get("num_experts_per_tok", 1)
    L = cfg["num_hidden_layers"]
    buckets = {"routed": 0, "dense": 0, "mtp": 0, "embed": 0}
    for f in get_safetensors_metadata(repo).files_metadata.values():
        for tname, t in f.tensors.items():
            if tname.endswith(("scale_inv", "_scales")):     # hệ số scale FP8/MXFP4, không phải tham số
                continue
            layer = re.search(r"layers\.(\d+)\.", tname)
            if tname.startswith("mtp.") or (layer and int(layer.group(1)) >= L):
                buckets["mtp"] += n_values(t)               # module multi-token prediction
            elif ".experts." in tname:
                buckets["routed"] += n_values(t)            # expert được router chọn
            else:
                buckets["dense"] += n_values(t)             # attention, embedding, norm, router, shared expert
                if "embed" in tname:
                    buckets["embed"] += n_values(t)         # bảng embedding đầu vào: tra bảng, không nhân ma trận
    total = buckets["dense"] + buckets["routed"]
    active = buckets["dense"] + buckets["routed"] * k / E
    shared = cfg.get("n_shared_experts") or (1 if "shared_expert_intermediate_size" in cfg else 0)
    print(f"{name:15} {E:>4} {k:>2} {shared:>6} {total/1e9:>7.2f} {active/1e9:>8.2f} "
          f"{(active - buckets['embed'])/1e9:>9.2f} {buckets['routed']/total:>8.0%} {buckets['mtp']/1e9:>5.2f}")

Output (chạy ngày 09/10/2026):

model              E  k shared total B active B act-emb B expert % MTP B
Llama 3.1 8B       1  1      0    8.03     8.03      7.50       0%  0.00
Mixtral 8x7B       8  2      0   46.70    12.88     12.75      97%  0.00
Qwen3-30B-A3B    128  8      0   30.53     3.35      3.04      95%  0.00
gpt-oss-20b       32  4      0   20.91     4.19      3.61      91%  0.00
gpt-oss-120b     128  4      0  116.83     5.71      5.13      98%  0.00
Qwen3-Next-80B   512 10      1   79.67     3.87      3.56      97%  1.65
DeepSeek-V3      256  8      1  671.03    37.55     36.63      97% 13.46
Jamba2 Mini       16  2      0   51.57    12.11     11.84      87%  0.00
Nemotron 3 Nano  128  6      1   31.58     3.58      3.23      93%  0.00

Cột active B tính cả bảng embedding đầu vào. Cột act-emb B bỏ bảng đó ra, vì tra bảng embedding không tốn phép nhân ma trận. Llama 3.1 8B có repo gốc gated nên code dùng bản mirror unsloth/…, giống bài 13. Mọi repo khác là repo chính thức của hãng, không gated tại ngày chạy.

Đối chiếu với con số công bố

Mỗi hãng đếm theo một quy ước riêng, và đây là chỗ hay gây nhầm:

Vài quan sát:

• Expert chiếm 87–98% tổng tham số. Model MoE thực chất là “một model nhỏ cộng một kho expert rất lớn”. gpt-oss model card cũng ghi phần MoE chiếm hơn 90% tổng số tham số (mục 2.1).
• Tỉ lệ kích hoạt ngày càng nhỏ. Mixtral kích hoạt 28% tham số mỗi token. DeepSeek-V3 khoảng 5,5%. Qwen3-Next khoảng 4,5%, gpt-oss-120b khoảng 4,4%.
• Embedding chiếm phần đáng kể ở model nhỏ có vocab lớn. gpt-oss có vocab 201.088 token, nên bảng embedding 0,58B chiếm hơn 10% tham số kích hoạt của bản 20b. Vì vậy quy ước đếm có tính embedding hay không làm con số lệch tới 14%.

06MoE trên roofline

Bài 13, mục 2 đã dựng mô hình roofline và tính ridge point của RTX 2080 Ti khoảng 87 FLOP/byte. Phần này dùng lại mô hình đó, chỉ thay một giả định: không phải weights nào cũng được đọc ở mỗi bước.

Batch 1: MoE rẻ như model nhỏ

Ở batch 1, mỗi layer chỉ đọc k expert được chọn. Byte phải đọc và FLOPs đều theo tham số kích hoạt. Mixtral decode một token đọc khoảng 12,9B tham số chứ không phải 46,7B, nên với cùng băng thông, tốc độ decode gần với một model dense 13B. Đây là lý do câu “tham số kích hoạt quyết định tốc độ” ở bài 1 đúng ở batch nhỏ.

Bộ nhớ thì vẫn theo tổng tham số: cả 46,7B phải nằm sẵn trong VRAM, vì token tiếp theo có thể chọn bất kỳ expert nào.

Batch B: gần như expert nào cũng bị đọc

Khi batch có B token (ở decode là B sequence, mỗi sequence một token), mỗi token chọn k expert riêng. Giả sử router chọn đều và độc lập, xác suất một expert không được token nào chọn là (1 − k/E)^B. Số expert khác nhau bị đọc trung bình là:

số expert bị đọc ≈ E · (1 − (1 − k/E)^B)

Áp vào vài model (phép tính lý thuyết theo công thức trên):

Mixtral chỉ cần batch 8 là đã đọc 90% expert mỗi bước. Model càng nhiều expert nhỏ thì đường cong càng thoải, nhưng tới batch 128 thì DeepSeek-V3 cũng đọc 98% expert.

Hệ quả: byte đọc mỗi bước tiến dần tới toàn bộ weights, còn FLOPs mỗi token vẫn chỉ theo phần kích hoạt. Một model dense ở batch B có arithmetic intensity khoảng B FLOP/byte. Một model MoE, khi mọi expert đã bị đọc, chỉ có khoảng B × (tham số kích hoạt / tổng tham số).

Nhìn từ phía từng expert thì rõ hơn. Mỗi expert trung bình nhận B × k / E token. Phép nhân ma trận của expert đó có intensity xấp xỉ bằng số token nó nhận. DeepSeek-V3 có k/E = 1/32: ở batch 256, mỗi expert chỉ xử lý khoảng 8 token. Muốn mỗi expert chạm ridge point khoảng 87 FLOP/byte của bài 13, batch phải vào khoảng 87 × 32 ≈ 2.800 token mỗi bước.

Kết luận về phục vụ và chạy local

• MoE hợp với dịch vụ lớn. Nhà cung cấp gom được batch hàng nghìn token, nên mọi expert đều bận, và MoE cho chất lượng của model lớn với FLOPs của model nhỏ.
• MoE không hợp chạy local ở batch 1 trên GPU nhỏ, theo nghĩa nó tốn VRAM của model lớn. Về tốc độ thì nó nhanh như model nhỏ, nếu chứa được toàn bộ weights.
• Expert offloading là cách lách khi VRAM không đủ: để expert ở RAM CPU, chỉ đưa lên GPU expert được chọn; khả thi ở batch 1 nhưng chậm vì băng thông PCIe thấp hơn VRAM nhiều.

Expert parallelism và all-to-all

Bài 10, mục 9 đã hứa nói về expert parallelism: chia các expert lên nhiều GPU, mỗi GPU giữ một phần. Đây là cách tự nhiên để chứa một model 671B, nhưng nó thêm hai bước truyền dữ liệu vào mỗi layer MoE:

1. Dispatch: sau khi router chọn, mỗi token phải được gửi tới GPU đang giữ expert của nó. Mọi GPU gửi cho mọi GPU, nên gọi là all-to-all.
2. Combine: output của expert phải gửi ngược về GPU ban đầu của token để cộng theo gate và đi tiếp.

Đây là chi phí mà model dense không có. DeepSeek-V3 dành nhiều công sức kỹ thuật cho nó:

• Node-limited routing: mỗi token chỉ được gửi tới tối đa M = 4 node (mục 2.1.2, mục 4.2), để giảm truyền dữ liệu qua InfiniBand giữa các node.
• DualPipe chia mỗi phần tính thành attention, all-to-all dispatch, MLP và all-to-all combine, rồi xếp chồng tính toán với truyền dữ liệu để all-to-all “được che hoàn toàn” (mục 3.2.1).
• Kernel all-to-all riêng chỉ cần 20 SM là dùng hết băng thông InfiniBand và NVLink (mục 3.2.2). Trên H800 có 132 SM, tức khoảng 15% chip dành cho truyền dữ liệu (mục 3.5.1).

Bài này dừng ở mức đó. Chi tiết kernel nằm ngoài phạm vi series.

MoE và quantization

Vì expert chiếm hơn 90% weights, quantize expert là cách giảm VRAM hiệu quả nhất. gpt-oss làm đúng như vậy: chỉ phần weights MoE được lượng tử sang MXFP4, khoảng 4,25 bit mỗi tham số, còn attention, router và embedding giữ bf16 (model card, mục 2.1; trường modules_to_not_convert trong config.json). Nhờ vậy bản 120b vừa một GPU 80 GB, bản 20b chạy được với 16 GB bộ nhớ. Model được post-train ở định dạng MXFP4, không phải lượng tử sau khi train xong. Định dạng block-wise và cách lượng tử nói chung đã có ở bài 12, mục 5.

07MoE học được gì, và đánh đổi gì

Cùng compute, MoE train tốt hơn

Switch Transformer báo kết quả rõ nhất: với cùng FLOPs mỗi token và cùng tài nguyên, Switch-Base 64 expert đạt cùng perplexity với T5-Base trong 1/7 thời gian (mục 3.2, Hình 5). Switch-C nhanh hơn 4 lần so với T5-XXL để đạt cùng perplexity (mục 5.6, Bảng 9).

Clark et al. 2022 thấy routing cải thiện model ở mọi kích thước họ thử (mục 1), nhưng lợi ích giảm dần khi model lớn lên (mục 4.2), với ngưỡng ước lượng khoảng 937B cho S-BASE, ngưỡng mà chính paper ghi là phụ thuộc mạnh vào số token train (mục 5.2). Krajewski et al. 2024 thêm số token train và granularity G = d_ff / d_expert (mục 4), và kết luận ngược lại: với cấu hình tối ưu, MoE luôn hiệu quả hơn dense ở mọi ngân sách compute (mục 6.3), còn đặt expert bằng kích thước FFN gốc “không tối ưu ở gần như mọi ngân sách compute” (abstract). Đó là lập luận cho fine-grained expert ở mục 5. Scaling laws gốc có ở bài 10, mục 4.

MoE 47B không ngang dense 47B

Cùng tổng tham số, MoE thường yếu hơn dense; cùng tham số kích hoạt, MoE thường mạnh hơn. Mixtral nằm ở giữa: paper so nó với Llama 2 70B và báo Mixtral vượt hoặc ngang trên gần như mọi benchmark phổ biến, với số tham số kích hoạt ít hơn 5 lần (mục 3, Bảng 2).

Expert không phải chuyên gia theo chủ đề

Cái tên “expert” gợi ý mỗi expert giỏi một lĩnh vực. Phân tích routing của Mixtral không thấy điều đó. Paper viết: “chúng tôi không thấy khuôn mẫu rõ ràng nào trong việc gán expert theo chủ đề”. Phân bố expert ở mọi layer rất giống nhau giữa ArXiv (LaTeX), sinh học (PubMed Abstracts) và triết học (PhilPapers). Chỉ DM Mathematics hơi khác một chút, và paper cho là do dataset này là dữ liệu tổng hợp (mục 5, Hình 7).

Router có hành vi theo cú pháp: từ self trong Python hay từ “Question” trong tiếng Anh thường đi qua cùng một expert, và các token thụt lề trong code luôn được gán cho cùng expert (Hình 8). Các token liền nhau cũng hay chọn cùng expert. Ở layer 15, expert được chọn đầu tiên lặp lại giữa hai token liên tiếp khoảng 28% thời gian trên ArXiv, DM Mathematics và Github, so với 12,5% nếu chọn ngẫu nhiên (bảng “Percentage of expert assignment repetitions”).

ST-MoE thấy kết quả tương tự với model encoder–decoder: expert ở encoder có chuyên môn hoá, nhưng theo loại token (dấu câu, động từ, tên riêng) và không chuyên theo ngôn ngữ (mục 7.3); expert ở decoder gần như không có chuyên môn hoá rõ (mục 7.1, 7.2).

Fine-tune MoE khó hơn

ST-MoE đưa ra bằng chứng cho giả thuyết rằng model sparse dễ overfit khi fine-tune. Trên SuperGLUE, ở tác vụ CB chỉ có 250 mẫu train, model dense tốt hơn model MoE trên tập validation dù MoE học tập train nhanh hơn; ở ReCoRD với 138 nghìn mẫu thì ngược lại, MoE hơn dense vài điểm phần trăm (mục 4.1, Hình 3). Một phát hiện thực tế: chỉ fine-tune phần không phải MoE cho kết quả gần như fine-tune toàn bộ, còn chỉ fine-tune phần MoE (khoảng 80% tham số) thì chất lượng giảm mạnh (mục 4.2, Hình 5). Với LoRA ở bài 12, điều này gợi ý gắn adapter vào attention trước, cân nhắc kỹ khi gắn vào từng expert.

08Thay attention bằng một state cố định

MoE làm mỗi token rẻ hơn, nhưng KV cache vẫn còn nguyên. Nửa sau của bài đi vào khoản chi phí thứ hai.

Attention nhớ mọi thứ, RNN nén mọi thứ

Attention giữ lại K, V của mọi token trước. Bộ nhớ đó tăng tuyến tính theo độ dài, và mỗi bước decode phải đọc lại toàn bộ. Đổi lại, model truy xuất được chính xác bất kỳ token nào trong context.

Họ RNN làm ngược lại: nén toàn bộ quá khứ vào một state có kích thước cố định, cập nhật state sau mỗi token. Decode mỗi token tốn chi phí không đổi, không có KV cache, và bộ nhớ không tăng theo độ dài. Cái giá là nén có mất mát. Mục 9 xem cái giá đó lớn tới đâu.

RNN cổ điển như LSTM còn không train song song theo thời gian được. Các kiến trúc dưới đây giữ dạng hồi quy khi decode nhưng vẫn train song song được.

Linear attention là một RNN

Katharopoulos et al. 2020 chỉ ra rằng nếu thay softmax(q·k) bằng tích của hai feature map φ(q)ᵀφ(k), thì attention có thể viết lại thành một RNN (mục 3.4, phương trình 18–20):

S_i = S_{i−1} + φ(k_i) · v_iᵀ         (state: ma trận cố định)
z_i = z_{i−1} + φ(k_i)                 (bộ chuẩn hoá)
y_i = φ(q_i)ᵀ S_i / φ(q_i)ᵀ z_i

với φ(x) = elu(x) + 1 (phương trình 7). State S có kích thước cố định; decode cập nhật từng bước với chi phí hằng số, và paper báo sinh ảnh autoregressive nhanh hơn tới 4.000 lần trên CIFAR-10 (abstract, Bảng 2). Đổi lại, state chỉ cộng dồn mà không biết quên.

S4: SSM có cấu trúc

Họ state space model (SSM) đi từ một hướng khác: phương trình vi phân tuyến tính của lý thuyết điều khiển. S4 (Gu et al. 2021) bắt đầu từ:

h'(t) = A·h(t) + B·x(t)
y(t)  = C·h(t)

(mục 2.1, phương trình 1). Rời rạc hoá với bước Δ cho dạng hồi quy h_t = Ā·h_{t−1} + B̄·x_t (mục 2.3). Vì A, B, C không đổi theo thời gian, đầu ra cũng là một phép convolution với kernel K̄ = (C̄B̄, C̄ĀB̄, C̄Ā²B̄, …) (mục 2.4): train bằng convolution, decode bằng hồi quy. S4 là model đầu tiên giải được tác vụ Path-X dài 16.384 bước của Long Range Arena (mục 1).

Mamba: SSM có chọn lọc

Mamba (Gu & Dao 2023) gọi thẳng tên điểm yếu của S4: không làm được suy luận dựa trên nội dung (abstract). A, B, C cố định nên mọi token được xử lý như nhau; model không thể quyết định “token này quan trọng, nhớ kỹ” hay “token này là nhiễu”.

Mamba sửa bằng cơ chế chọn lọc: cho Δ, B, C trở thành hàm của input. Cụ thể B = Linear(x), C = Linear(x), và Δ = softplus(tham số + Linear(x)) (mục 3.2, Thuật toán 2). Ý nghĩa trực quan của Δ:

• Δ lớn: Ā = exp(Δ·A) tiến về 0 (vì A âm), state cũ bị quên nhanh, input hiện tại được ghi mạnh. Model “tập trung” vào token này.
• Δ nhỏ: Ā tiến về 1, state gần như giữ nguyên, input hiện tại gần như bị bỏ qua.

Phép rời rạc hoá là zero-order hold: Ā = exp(ΔA), B̄ = (ΔA)⁻¹(exp(ΔA) − I)·ΔB (mục 2, phương trình 4). Trong code 3 bên dưới, B̄ dùng dạng rút gọn Δ·B.

Cái giá: khi A, B, C thay đổi theo từng token, model mất dạng convolution (mục 3.2). Mamba bù lại bằng một parallel scan có hiểu phần cứng: gộp kernel, nạp tham số từ HBM vào SRAM, rời rạc hoá và chạy hồi quy ngay trong SRAM, rồi chỉ ghi output ra HBM; state trung gian không lưu mà tính lại ở backward (mục 3.3.2). Đây là cùng tinh thần với FlashAttention ở bài 13, mục 7: giảm số byte đọc ghi bộ nhớ chậm.

Paper báo throughput inference cao hơn Transformer 5 lần, Mamba-3B ngang Transformer gấp đôi kích thước (abstract), và ở tác vụ induction heads thì tổng quát hoá hoàn hảo tới chuỗi một triệu token (mục 4.1.2). State mỗi kênh có N = 16 (mục 4.5).

Mamba-2: SSM và attention là hai cách tính cùng một thứ

Mamba-2 (Dao & Gu 2024) đưa ra state space duality (SSD). Nếu giới hạn A thành một số vô hướng nhân ma trận đơn vị (mục 2.4), thì phép biến đổi của SSM chính là phép nhân với một ma trận semiseparable (Định lý 3.5). Viết ra, ma trận đó có dạng của masked attention: một ma trận tam giác dưới L chứa hệ số phân rã, nhân từng phần tử với C·Bᵀ, giống Q·Kᵀ trong attention. Paper viết rằng tính trực tiếp SSM vô hướng theo cách này “giống hệt masked kernel attention bậc hai” (mục 5.1).

Dạng hồi quy tuyến tính theo độ dài, hợp với decode; dạng ma trận bậc hai nhưng toàn phép nhân ma trận, chạy tốt trên tensor core. Thuật toán SSD dùng dạng ma trận trong từng khúc và dạng hồi quy giữa các khúc, nhanh hơn selective scan của Mamba 2–8 lần và cho phép state lớn hơn nhiều, N = 64 hay 256 thay cho 16 (mục 9.1, Hình 10).

Code mẫu 3: một SSM, ba cách tính

Code dưới cài một SSM chéo có chọn lọc (Δ, B, C phụ thuộc input như Mamba) theo ba cách: hồi quy từng bước, parallel scan, và dạng ma trận “attention có mask” của Mamba-2. Phần 2 đọc config.json của bốn model thật để so thứ phải giữ lại cho mỗi sequence khi decode.

# Python 3.12, torch 2.14.1, huggingface_hub 1.33.0 — chạy trên CPU, vài giây
import json, math, os, torch
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
from huggingface_hub import hf_hub_download
torch.manual_seed(0)

# Phần 1: một SSM chéo có chọn lọc (kiểu Mamba), tính bằng ba cách
T, D, N = 64, 4, 8                                       # độ dài, số kênh, kích thước state mỗi kênh
x = torch.randn(T, D, dtype=torch.float64)
A = -torch.rand(D, N, dtype=torch.float64) - 0.5         # A âm: state tự phân rã
dt = torch.nn.functional.softplus(torch.randn(T, D, dtype=torch.float64))  # Δ_t phụ thuộc input
B, C = torch.randn(T, N, dtype=torch.float64), torch.randn(T, N, dtype=torch.float64)  # B_t, C_t cũng vậy
a = torch.exp(dt[:, :, None] * A)                        # Ā_t = exp(Δ_t·A), (T, D, N)
b = dt[:, :, None] * B[:, None, :] * x[:, :, None]       # B̄_t·x_t, dạng rút gọn Δ_t·B_t·x_t

def recurrent():                                         # dạng RNN: một state cố định, đi từng bước
    h, ys = torch.zeros(D, N, dtype=torch.float64), []
    for t in range(T):
        h = a[t] * h + b[t]
        ys.append((h * C[t]).sum(-1))                    # y_t = C_t · h_t
    return torch.stack(ys)

def scan():                                              # parallel scan: log2(T) bước, mỗi bước song song
    aa, hh, step = a.clone(), b.clone(), 1               # phép kết hợp (a1,b1)∘(a2,b2) = (a1·a2, a2·b1 + b2)
    while step < T:
        hh = torch.cat([hh[:step], aa[step:] * hh[:-step] + hh[step:]])
        aa = torch.cat([aa[:step], aa[step:] * aa[:-step]])
        step *= 2
    return (hh * C[:, None, :]).sum(-1)

def matrix():                                            # dạng "attention có mask": y = (L ∘ C·Bᵀ) · (Δ·x)
    logcum = torch.cumsum(torch.log(a), dim=0)           # Σ log Ā tới bước t
    diff = logcum[:, None] - logcum[None, :]             # (T, T, D, N): log Π_{s<r≤t} Ā_r
    mask = torch.tril(torch.ones(T, T, dtype=torch.bool))[:, :, None, None]
    L = torch.where(mask, diff, -math.inf).exp()         # ma trận phân rã, tam giác dưới
    M = (L * C[:, None, None, :] * B[None, :, None, :]).sum(-1)  # (T, T, D): "attention score"
    return torch.einsum("tsd,sd->td", M, dt * x)

y_rec, y_scan, y_mat = recurrent(), scan(), matrix()
print(f"scan   vs recurrent: max |Δ| = {(y_scan - y_rec).abs().max():.2e}")
print(f"matrix vs recurrent: max |Δ| = {(y_mat - y_rec).abs().max():.2e}")

# Phần 2: thứ phải giữ lại cho mỗi sequence khi decode, đọc từ config.json
def cache_bytes(cfg, n):                                 # trả về (KV cache tăng theo n, state cố định)
    t = cfg["model_type"]
    if t == "llama":
        attn, kv, hd, fixed = cfg["num_hidden_layers"], cfg["num_key_value_heads"], cfg["head_dim"], 0
    elif t == "nemotron_h":                              # M = Mamba-2, * = attention, E = MoE
        pat = cfg["hybrid_override_pattern"]
        attn, kv, hd = pat.count("*"), cfg["num_key_value_heads"], cfg["head_dim"]
        ssm = cfg["mamba_num_heads"] * cfg["mamba_head_dim"] * cfg["ssm_state_size"] * 4  # state fp32
        conv = (cfg["mamba_num_heads"] * cfg["mamba_head_dim"] + 2 * cfg["n_groups"] * cfg["ssm_state_size"]) \
               * (cfg["conv_kernel"] - 1) * 2
        fixed = pat.count("M") * (ssm + conv)
    elif t == "qwen3_next":                              # 3 lớp Gated DeltaNet rồi 1 lớp attention
        L, every = cfg["num_hidden_layers"], cfg["full_attention_interval"]
        attn, kv, hd = L // every, cfg["num_key_value_heads"], cfg["head_dim"]
        state = cfg["linear_num_value_heads"] * cfg["linear_key_head_dim"] * cfg["linear_value_head_dim"] * 2
        conv = (2 * cfg["linear_num_key_heads"] * cfg["linear_key_head_dim"]
                + cfg["linear_num_value_heads"] * cfg["linear_value_head_dim"]) * (cfg["linear_conv_kernel_dim"] - 1) * 2
        fixed = (L - attn) * (state + conv)
    elif t == "jamba":                                   # 1 lớp attention mỗi 8 lớp, còn lại Mamba
        L = cfg["num_hidden_layers"]
        attn = sum(1 for i in range(L) if i % cfg["attn_layer_period"] == cfg["attn_layer_offset"])
        kv, hd = cfg["num_key_value_heads"], cfg["hidden_size"] // cfg["num_attention_heads"]
        d_in = cfg["mamba_expand"] * cfg["hidden_size"]
        fixed = (L - attn) * d_in * (cfg["mamba_d_state"] + cfg["mamba_d_conv"] - 1) * 2
    return attn * 2 * kv * hd * 2 * n, fixed             # K và V, bf16 = 2 byte

MODELS = {"Llama 3.1 8B": "unsloth/Meta-Llama-3.1-8B",
          "Jamba2 Mini": "ai21labs/AI21-Jamba2-Mini",
          "Nemotron 3 Nano": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
          "Qwen3-Next-80B": "Qwen/Qwen3-Next-80B-A3B-Instruct"}
MiB = 2**20
print(f"\n{'MiB mỗi sequence':16} {'state cố định':>13} " + " ".join(f"{f'n={n//1024}k':>9}" for n in (4096, 32768, 131072)))
for name, repo in MODELS.items():
    cfg = json.load(open(hf_hub_download(repo, "config.json")))
    fixed = cache_bytes(cfg, 0)[1]
    print(f"{name:16} {fixed / MiB:>13.1f} " +
          " ".join(f"{sum(cache_bytes(cfg, n)) / MiB:>9.0f}" for n in (4096, 32768, 131072)))

Output (chạy ngày 09/10/2026):

scan   vs recurrent: max |Δ| = 3.55e-15
matrix vs recurrent: max |Δ| = 1.87e-14

MiB mỗi sequence state cố định      n=4k     n=32k    n=128k
Llama 3.1 8B               0.0       512      4096     16384
Jamba2 Mini                8.3        72       520      2056
Nemotron 3 Nano           46.8        71       239       815
Qwen3-Next-80B            37.7       134       806      3110

Phần 1. Ba cách tính khớp nhau tới sai số làm tròn của float64 (cỡ 10⁻¹⁴). Dạng hồi quy là thứ chạy khi decode, mỗi token chỉ cập nhật một state D × N. Dạng scan cho thấy vì sao train song song được: phép cập nhật h_t = a_t·h_{t−1} + b_t có tính kết hợp, nên tính được trong log₂(T) bước. Dạng ma trận là đối ngẫu của Mamba-2: một ma trận T × T giống attention score, có mask tam giác dưới, cho đúng cùng output.

Phần 2. Đây là số tính từ config.json theo cách cấp phát mặc định của vLLM, không phải số đo bộ nhớ khi chạy thật. Cột “state cố định” là phần không tăng theo độ dài; các cột còn lại là tổng bộ nhớ mỗi sequence (state cộng KV cache) ở ba độ dài context.

• Llama 3.1 8B không có state, chỉ có KV cache: 16 GiB ở 128k token, như bài 13 đã tính.
• Jamba2 Mini chỉ có 4 layer attention trên 32 layer, nên KV cache ở 128k là 2 GiB, nhỏ hơn Llama 8 lần dù model lớn gấp 6 lần. Con số này khớp với Bảng 1 của paper Jamba: 4 GB KV cache ở context 256K cho Jamba, so với 32 GB cho Mixtral và 128 GB cho Llama-2 7B.
• Nemotron 3 Nano có state lớn nhất (46,8 MiB, vì state Mamba-2 lưu fp32 theo trường mamba_ssm_cache_dtype), nhưng chỉ 6 layer attention với 2 KV head, nên ở 128k chỉ tốn 815 MiB, khoảng 1/20 của Llama 3.1 8B.
• Qwen3-Next có 12 layer attention, mỗi layer 2 KV head với head_dim 256, nên KV cache lớn hơn Nemotron, nhưng vẫn chỉ khoảng 1/5 của Llama.

Code giả định state lưu bf16, trừ Nemotron mà config ghi rõ fp32, và conv state chỉ giữ kernel − 1 bước, như mặc định của vLLM. Engine khác có thể tốn gấp đôi: HF transformers giữ state Gated DeltaNet ở fp32 và giữ conv state đủ cả kernel, nên state cố định của Qwen3-Next là khoảng 74 MiB thay vì 37,7, của Jamba2 Mini khoảng 16 MiB thay vì 8,3. Model card Jamba2 khuyên chạy vLLM với --mamba-ssm-cache-dtype float32; theo khuyến nghị đó, state của Jamba2 Mini cũng lên khoảng 15 MiB. Dù vậy, state vẫn nhỏ hơn KV cache hàng chục lần ở context dài, nên các so sánh với Llama ở trên không đổi.

Gated DeltaNet

Linear attention cộng dồn v·kᵀ vào state mà không bao giờ xoá. DeltaNet thay phép cộng đó bằng delta rule: trước khi ghi giá trị mới cho key k, xoá bớt giá trị cũ đang gắn với k, S_t = S_{t−1}(I − β_t·k_t·k_tᵀ) + β_t·v_t·k_tᵀ (Yang et al. 2024, mục 2.2). Mamba-2 thì có cổng quên α_t nhưng quên đều mọi thứ cùng lúc. Gated DeltaNet ghép hai ý: S_t = S_{t−1}(α_t(I − β_t·k_t·k_tᵀ)) + β_t·v_t·k_tᵀ (mục 3.1, phương trình 10). Paper tóm tắt: cổng cho phép xoá bộ nhớ nhanh, delta rule cho phép cập nhật có mục tiêu (abstract). Ở quy mô 1,3B tham số và 100B token, Gated DeltaNet vượt Mamba2 và DeltaNet về độ chính xác trung bình (55,32 so với 54,89 và 52,14), và bản hybrid với sliding-window attention còn cao hơn (Bảng 3). Đây là lớp mà Qwen3-Next dùng cho ba phần tư số layer.

Cùng họ “state cố định” còn có RetNet, Hyena và RWKV; bài này không đi vào chúng.

09Vì sao không bỏ hẳn attention

Giới hạn của state cố định

State cố định nghĩa là nén toàn bộ quá khứ vào một số byte không đổi. Với tác vụ cần nhớ chính xác một chi tiết bất kỳ trong context dài, đó là giới hạn nền tảng, không chỉ là chuyện train chưa đủ.

Jelassi et al. 2024 chứng minh điều đó ở mức lý thuyết: một Transformer hai layer có thể copy chuỗi có độ dài mũ theo số head, trong khi các SSM dạng generalized bị giới hạn cơ bản bởi state cố định (abstract; Định lý 2.3, 2.7). Thực nghiệm cũng khớp:

• Copy: Transformer cần ít hơn 100 lần số mẫu so với SSM tốt nhất để học tác vụ copy (mục 3.2, Hình 1a). Khi tăng độ dài input vượt độ dài lúc train, SSM gần như rơi về 0 ngay, còn Transformer giảm từ từ (mục 3.3).
• Tra danh bạ điện thoại trên model đã pretrain: đưa một danh bạ vào context rồi hỏi số của một người. Transformer nhỏ nhất (Pythia 410M) vượt Mamba lớn nhất (2,8B) khi danh bạ đủ dài, từ 70 mục trở lên (mục 4.3, Hình 1c).

Nghiên cứu của NVIDIA (Waleffe et al. 2024) so ở quy mô lớn hơn: Mamba, Mamba-2 và Transformer 8B tham số, train trên cùng dữ liệu tới 3,5T token. Model thuần SSM ngang hoặc vượt Transformer ở nhiều tác vụ, nhưng tụt lại ở tác vụ cần copy hoặc in-context learning mạnh như MMLU 5-shot, Phonebook, và ở long-context reasoning (abstract):

• MMLU 5-shot ở 1,1T token: Transformer 46,28, Mamba 28,00, Mamba-2 29,19 (Bảng 2).
• Train lên 3,5T token thì khoảng cách thu hẹp: Transformer 50,07, Mamba-2 48,7 (Bảng 3).
• Phonebook thì không thu hẹp: Mamba-2 trả lời sai khi danh bạ dài hơn khoảng 500 token, kể cả ở 3,5T token (mục 3.3.3, Hình 7a).

Paper Jamba có một quan sát cụ thể về in-context learning (mục 6.2, Bảng 6): trên IMDB, nhãn đúng là “Positive” hoặc “Negative”, nhưng model Mamba thuần hay trả lời “Very Good”, “Funny”, “3/10”. Ở model 1,3B, attention đạt 84,1, Mamba 48,8, hybrid 90,9, và các induction head xuất hiện trong layer attention của model hybrid.

Một ít attention là đủ

Cùng nghiên cứu của NVIDIA thử model hybrid Mamba-2-Hybrid 8B: 56 layer, gồm 24 layer Mamba-2 (43%), 4 layer self-attention (7%) và 28 layer MLP (50%) (mục 4.2). Kết quả (abstract, Bảng 7):

• Vượt Transformer 8B trên cả 12 tác vụ chuẩn, trung bình cao hơn 2,65 điểm.
• MMLU 5-shot 53,60 so với 50,07.
• Được dự đoán sinh token nhanh hơn tới 8 lần khi inference, theo ước lượng ở batch 32 với context dài (mục 4.2, Hình 5).
• Bản 128K làm tác vụ Phonebook hoàn hảo với danh bạ hơn 150K token (mục 1).

Ngược lại, trên tác vụ long-context tự nhiên ở 16K/32K, Transformer vẫn nhỉnh hơn khoảng một điểm trung bình (mục 4.3.2, Bảng 9). Mamba-2 paper cũng thấy khoảng 10% layer là attention cho kết quả tốt nhất (mục 9.2.3, Bảng 2).

Hybrid trong các model 2024–2026

Về shared expert của Nemotron 3 Nano: config.json ghi 1 shared expert với moe_shared_expert_intermediate_size 3.712, gấp đôi expert thường (1.856), còn Bảng 1 của tech report ghi 2 shared expert. Hai cách mô tả tương đương về số tham số. Code 1 đọc theo config nên in 1.

Về tốc độ, NVIDIA báo Nemotron-H nhanh hơn tới 3 lần khi inference so với Transformer cùng cỡ (abstract; con số 3 lần là bản 8B so với Llama-3.1-8B, mục 2.6). Với Nemotron 3 Nano, ở cấu hình 8K input và 16K output trên một H200, throughput cao hơn Qwen3-30B-A3B-Thinking-2507 3,3 lần và gpt-oss-20b 2,2 lần (tech report, mục 1, Hình 1). Jamba báo throughput gấp 3 lần Mixtral ở context 128K (mục 3.2, Hình 3b). Đây là số các hãng tự đo trên cấu hình của họ.

Bảng so sánh bốn kiểu

Kết luận của nửa sau: attention không biến mất. Nó trở thành một nguyên liệu đắt, được đặt ở vài layer để lo việc truy xuất chính xác, còn phần lớn layer dùng state cố định rẻ hơn.

10Những chỗ dễ hiểu sai

Chín hiểu lầm dưới đây khớp với hàng trên và hàng dưới của Hình 1.

Tổng kết: hai khoản chi phí, hai cách tránh

Một model dense trả hai khoản cho mỗi token: toàn bộ tham số, và một KV cache tăng theo độ dài. Bài này đi qua hai hướng giảm hai khoản đó.

• MoE thay FFN, không thay attention. Router là một lớp tuyến tính và một phép top-k, chọn theo từng token ở từng layer. Mỗi token chỉ chạy k expert, nên compute theo tham số kích hoạt, còn VRAM theo tổng tham số.
• Load balancing là bài toán trung tâm. Để tự nhiên thì router dồn vào vài expert và vòng đó tự củng cố. Switch dùng auxiliary loss α · N · Σ f_i · P_i; ST-MoE thêm z-loss để ổn định; DeepSeek-V3 dùng bias chỉ để chọn expert, không đi vào gradient. Expert capacity làm token có thể bị drop tuỳ batch, và đó là cơ chế đứng sau giả thuyết về MoE không tất định ở bài 9, một giả thuyết chưa ai xác nhận.
• Đếm tham số phải đọc kỹ quy ước. Code 1 đếm từ header safetensors và khớp con số công bố của cả chín model, khi để ý hãng có tính embedding hay không. Expert chiếm 87–98% tham số; tỉ lệ kích hoạt giảm từ 28% ở Mixtral xuống khoảng 4–5% ở các model 2025.
• Trên roofline, MoE rẻ ở batch 1 nhưng cần batch rất lớn mới dùng hết compute, vì mỗi expert chỉ nhận B·k/E token. Expert parallelism thêm hai lần all-to-all mỗi layer.
• SSM thay attention bằng một state cố định. Linear attention viết được thành RNN; S4 có dạng convolution; Mamba cho Δ, B, C phụ thuộc input để chọn cái cần nhớ, đổi lại phải dùng parallel scan; Mamba-2 chỉ ra SSM và masked attention là hai cách tính cùng một phép biến đổi, như code 3 kiểm chứng.
• State cố định là nén có mất mát. Model thuần SSM yếu ở copy, tra cứu và in-context learning. Một ít attention, khoảng 7–10% số layer theo cách đếm của NVIDIA (tính cả layer FFN/MoE; nếu đếm mỗi layer gồm mixer + FFN như Jamba và Qwen thì là 12,5–25%), là đủ để bù, nên các model mới như Jamba, Nemotron và Qwen3-Next đều là hybrid, thường kèm MoE.

Bài 15 sẽ nói về reasoning model và test-time compute: model sinh hàng nghìn token suy nghĩ trước khi trả lời. MoE làm mỗi token rẻ hơn, và state cố định làm context dài rẻ hơn. Cả hai là lý do việc bắt model “nghĩ lâu hơn” trở nên khả thi về chi phí.

Tài liệu tham khảo

Paper: Mixture of Experts

• Shazeer et al. 2017 — Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• Lepikhin et al. 2020 — GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• Fedus, Zoph & Shazeer 2021 — Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (bản JMLR)
• Zoph et al. 2022 — ST-MoE: Designing Stable and Transferable Sparse Expert Models
• Zhou et al. 2022 — Mixture-of-Experts with Expert Choice Routing
• Clark et al. 2022 — Unified Scaling Laws for Routed Language Models
• Jiang et al. 2024 — Mixtral of Experts
• Dai et al. 2024 — DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models
• Krajewski et al. 2024 — Scaling Laws for Fine-Grained Mixture of Experts
• Wang et al. 2024 — Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
• DeepSeek-AI 2024 — DeepSeek-V3 Technical Report
• OpenAI 2025 — gpt-oss-120b & gpt-oss-20b Model Card

Paper: SSM, linear attention và hybrid

• Katharopoulos et al. 2020 — Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
• Gu, Goel & Ré 2021 — Efficiently Modeling Long Sequences with Structured State Spaces (S4)
• Gu & Dao 2023 — Mamba: Linear-Time Sequence Modeling with Selective State Spaces
• Jelassi et al. 2024 — Repeat After Me: Transformers are Better than State Space Models at Copying
• Lieber et al. 2024 — Jamba: A Hybrid Transformer-Mamba Language Model
• Dao & Gu 2024 — Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (Mamba-2)
• Waleffe et al. 2024 — An Empirical Study of Mamba-based Language Models
• Yang, Kautz & Hatamizadeh 2024 — Gated Delta Networks: Improving Mamba2 with Delta Rule
• NVIDIA 2025 — Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models
• NVIDIA 2025 — Nemotron 3 Nano: Open, Efficient Mixture-of-Experts Hybrid Mamba-Transformer Model for Agentic Reasoning

Model card và config (lấy ngày 09/10/2026)

• mistralai/Mixtral-8x7B-v0.1
• deepseek-ai/DeepSeek-V3, README_WEIGHTS.md
• openai/gpt-oss-120b, openai/gpt-oss-20b
• Qwen/Qwen3-Next-80B-A3B-Instruct, Qwen/Qwen3-30B-A3B-Base
• ai21labs/AI21-Jamba2-Mini
• nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
• unsloth/Meta-Llama-3.1-8B (mirror không gated của Llama 3.1 8B)
• Hugging Face — huggingface_hub: get_safetensors_metadata