6.2 Suy luận tensor parallel

Mục 6.1 đã nói: khác PP, TP để nhiều card cùng lúc tính cùng một Token, nên nó là lựa chọn đầu để hạ độ trễ một request. Cái giá là mỗi lớp đều phải đồng bộ. Mục này trước hết làm rõ cách cắt một lớp và hai AllReduce, rồi tính sổ giao tiếp, cuối cùng bàn mức TP nên lấy bao nhiêu và vì sao TP hầu như không xuyên node.

Chương 3 mục 3.3 đã đọc ColumnParallelLinear và RowParallelLinear trong nano-vllm. Ở đây nhìn từ góc người triển khai: những cách cắt đó nghĩa là gì trên mô hình và phần cứng thật.

⏰ Thời hạn thông tin: Phần nguyên lý không phụ thuộc framework; mô tả vLLM theo v0.32.0 (phát hành 2026-10-18). Băng thông phần cứng lấy từ trang sản phẩm NVIDIA.

📑 Mục lục

• 1. Một lớp cắt thế nào: cặp cột và hàng
• 2. Attention cắt theo head: nhân bản KV head dưới GQA
• 3. Sổ giao tiếp: hai AllReduce tốn bao lâu
• 4. Vì sao TP không ra khỏi node
• 5. Mức TP: tách độ trễ khỏi thông lượng
• 6. Chi tiết triển khai TP trong vLLM
• 7. Tính tay: mô hình thô mức TP và thời gian mỗi Token
• Tóm tắt
• Danh sách tự kiểm
• Tài liệu tham khảo

---

1. Một lớp cắt thế nào: cặp cột và hàng

Một lớp Transformer có hai lớp con: attention và FFN. Cách cắt TP đến từ paper Megatron-LM; cốt lõi là để cắt cột và cắt hàng xuất hiện theo cặp.

Lấy FFN. Đầu vào $X\in\mathbb R^{n\times h}$, FFN là $Y=\text{act}(XA)\,B$, với $A\in\mathbb R^{h\times f}$, $B\in\mathbb R^{f\times h}$.

• Ma trận thứ nhất cắt theo cột: $A=[A_1,\ldots,A_t]$; card $i$ tính $\text{act}(XA_i)$. Hàm kích hoạt theo từng phần tử; mỗi card tự tính cột của mình, kết quả đúng, không cần giao tiếp.
• Ma trận thứ hai cắt theo hàng: $B$ cắt theo chiều $f$ thành $B_1,\ldots,B_t$; card $i$ tính $\text{act}(XA_i)B_i$, được tổng cục bộ hình dạng đầy đủ $n\times h$.

$$Y=\sum_{i=1}^{t}\text{act}(XA_i)\,B_i$$

Cộng $t$ tổng cục bộ lại chính là một AllReduce. Activation trên chiều trung gian $f$ luôn ở lại từng card, không bao giờ ghép đủ.

Attention cùng mẫu: chiếu QKV cắt cột theo head; mỗi card tính attention đầy đủ trên vài head của mình; chiếu ra o_proj cắt hàng, cuối cùng AllReduce một lần nữa. Vậy mỗi lớp đúng 2 AllReduce.

Hình dưới là luồng dữ liệu một lớp Llama-3.1-70B ở TP=4; bốn màu là phân mảnh của bốn GPU; thanh tím là AllReduce cần mọi card:

Các số lấy từ cấu hình mô hình: 64 Q head chia 4 được 16, 8 KV head chia 4 được 2, chiều trung gian FFN 28672 chia 4 được 7168.

📌 Điểm then chốt: Tính đúng của TP chỉ dựa vào một việc: sau cắt hàng, các tổng cục bộ cộng lại bằng kết quả gốc. Mọi tổ hợp “cột + hàng” cuối phải có một AllReduce; thiếu một lần là sai kết quả.

---

2. Attention cắt theo head: nhân bản KV head dưới GQA

2.1 Số head quyết định TP lấy được những giá trị nào

Attention cắt theo head, nên số Q head phải chia hết cho kích thước TP. Llama-3.1-70B có 64 Q head, TP có thể 1, 2, 4, 8, 16…; Qwen3-32B cũng 64 Q head, tương tự. Thực tế TP còn bị số GPU mỗi node hạn chế; giá trị hay gặp là 2, 4, 8.

2.2 Khi KV head ít hơn TP thì sao

GQA để nhiều Q head chia sẻ một nhóm K, V. Llama-3.1-70B chỉ có 8 KV head; TP=16 thì mỗi card không đủ một KV head đầy đủ. Cách của vLLM là nhân bản: khi TP lớn hơn số KV head, mỗi card lưu 1 KV head, mỗi KV head được sao sang TP / số KV head card (khi đó yêu cầu TP chia hết cho số KV head).

Hình so sánh TP=8 và TP=16; ô xanh là số Q head mỗi card, ô xanh lục là số hiệu KV head card đó lưu:

Hệ quả của nhân bản:

• KV Cache mỗi card không còn nhỏ theo TP. TP từ 8 lên 16, trọng số mỗi card giảm nửa, nhưng KV mỗi Token giữ nguyên. Trong ước lượng mục 6.1, dung lượng KV ở TP=16 chỉ nhiều hơn TP=8 khoảng 18%.
• Tổng bộ nhớ KV nhân đôi. Cùng một KV lưu trên hai card; trên 16 card tổng KV gấp đôi trường hợp không nhân bản.

Vậy với mô hình GQA, số KV head là trần tự nhiên của TP. Vượt qua rồi muốn scale thêm thì đổi sang PP hoặc DP. vLLM còn có Decode Context Parallel (DCP), cắt KV Cache theo chiều Token để bỏ phần trùng này; mục 6.6 sẽ nói.

⚠️ Lưu ý: MLA (như DeepSeek-V3) chỉ cache một latent đã nén, không có “KV head” để cắt; dưới TP mọi card đều phải lưu đủ latent KV. Đó cũng là một lý do mô hình kiểu DeepSeek-V3 thường để attention đi DP chứ không TP lớn (mục 6.4).

---

3. Sổ giao tiếp: hai AllReduce tốn bao lâu

3.1 Chi phí một AllReduce

Tài liệu hiệu năng nccl-tests suy ra thời gian truyền AllReduce: dữ liệu kích thước $S$ trên $t$ card, mỗi phần tử cần $2(t-1)$ lần truyền, chia cho $t$ liên kết băng thông $B$, thời gian lý tưởng là $\frac{2(t-1)}{t}\cdot\frac{S}{B}$. Cộng overhead cố định, tổng thời gian khoảng:

$$T_{\text{AR}}\approx\alpha+\frac{2(t-1)}{t}\cdot\frac{S}{B}$$

$B$ là băng thông liên kết mỗi card; $\alpha$ là chi phí cố định (khởi động Kernel, căn hàng các card, độ trễ liên kết). Hạng đầu không phụ thuộc khối lượng; hạng sau tỉ lệ với khối lượng.

3.2 Tin Decode rất nhỏ

Decode thì $S=n\times h\times b$, $n$ là số Token bước này. Llama-3.1-70B ($h=8192$, BF16):

Decode hạng băng thông chỉ vài micro giây—cùng bậc với overhead cố định hoặc nhỏ hơn—nên giao tiếp TP giai đoạn Decode bị giới hạn bởi độ trễ. Mỗi bước 160 AllReduce, mỗi lần thêm 10 µs thì TPOT thêm 1.6 ms.

Prefill một lần xử lý hàng nghìn Token, tin lớn tới hàng trăm MiB, trở thành giới hạn băng thông; khi đó khoảng cách NVLink và PCIe bị phóng đại vài lần.

💡 Gợi ý: Dùng all_reduce_perf của nccl-tests quét kích thước tin từ 8 B đến hàng trăm MiB trên máy mình để thấy ranh giới vùng độ trễ và vùng băng thông. busbw nó báo đã nhân $2(t-1)/t$, so trực tiếp với băng thông liên kết được.

---

4. Vì sao TP không ra khỏi node

Trong một máy 8 card H100 SXM, GPU nối full qua NVLink và NVSwitch; NVIDIA công bố băng thông NVLink mỗi card 900 GB/s (cộng hai chiều). Xuyên node phải qua NIC: lấy DGX B200, 8 GPU kèm 8 ConnectX-7 một cổng, mỗi cái 400 Gb/s, trung bình mỗi GPU khoảng 50 GB/s.

Trải TP=16 sang hai máy 8 card, mỗi AllReduce đều phải có dữ liệu qua NIC. Thuật toán Ring bị đoạn chậm nhất kéo; băng thông từ 450 GB/s xuống 50 GB/s, độ trễ xuyên máy cũng cao hơn trong máy. Mỗi lớp hai lần, mỗi bước hàng trăm lần—chi phí cộng lại rất rõ.

Hình trái: TP=16 ngang hai máy. Phải: TP=8 giữ trong máy, hai máy nối bằng PP=2:

Đó là gợi ý của tài liệu vLLM: kích thước TP = số GPU mỗi node, kích thước PP = số node. PP mỗi bước chỉ truyền activation một lần ở biên đoạn; số lần từ hàng trăm xuống hàng đơn vị.

Ngoại lệ là hệ NVLink cấp rack, như GB200 NVL72, 72 GPU trong cùng miền NVLink. Trên hệ đó biên “node” đổi; TP có thể xuyên máy vật lý, tiêu chí đổi thành “có cùng miền NVLink hay không”.

---

5. Mức TP: tách độ trễ khỏi thông lượng

Tăng mức TP thì vài việc xảy ra cùng lúc:

• Mỗi card đọc ít trọng số hơn: thời gian truy cập bộ nhớ mỗi bước Decode giảm khoảng $1/t$—nguồn chính của lợi độ trễ.
• Số lần giao tiếp không đổi, mỗi lần nhiều card hơn: mỗi bước vẫn $2L$ AllReduce; overhead cố định không giảm theo $t$, tỉ trọng trong tổng thời gian ngày càng lớn.
• Ma trận mỗi card nhỏ hơn: chiều N hoặc K của GEMM bị cắt còn $1/t$; ma trận nhỏ khó đầy Tensor Core, Prefill càng rõ.
• Mỗi card trống thêm bộ nhớ: trọng số và KV (trong phạm vi số KV head) đều nhỏ theo $1/t$, chứa được nhiều request đồng thời hơn.

Ba điểm đầu nói TP đôi không làm độ trễ giảm một nửa, thông lượng mỗi card thường còn giảm. Điểm bốn nói khi hệ bị KV kẹt, tăng TP để mỗi card phục vụ thêm vài request, thông lượng mỗi card có thể tăng.

Nghiên cứu suy luận PaLM của Pope et al. cũng kết luận tương tự: không có bố cục cắt nào tối ưu ở mọi kích thước Batch; độ trễ thấp và thông lượng cao cần bố cục khác nhau. Với triển khai vLLM:

⚠️ Lưu ý: So thông lượng các mức TP khác nhau phải chuẩn hóa theo số card, và so dưới cùng ràng buộc độ trễ. Thông lượng một bản sao TP=8 cao hơn TP=4 không có nghĩa 8 card mở một TP=8 hơn hai TP=4.

---

6. Chi tiết triển khai TP trong vLLM

6.1 Cách khởi động

Một node chỉ cần một tham số:

vllm serve meta-llama/Llama-3.1-70B-Instruct --tensor-parallel-size 4

vLLM tạo một tiến trình Worker cho mỗi card. Một node mặc định dùng multiprocessing Python (mp) quản lý Worker; nhiều node có thể dùng Ray (mục 6.5). Worker giao tiếp tập hợp qua NCCL.

6.2 Những lớp nào bị cắt

Vocab của Embedding và LM Head thường rất lớn (Llama-3.1 là 128256); cắt theo vocab tiết kiệm đáng kể bộ nhớ. Mục 3.3 đã đọc VocabParallelEmbedding và ParallelLMHead tương ứng trong nano-vllm.

6.3 Custom AllReduce

Thuật toán tổng quát của NCCL nhắm nhiều topology; với tin nhỏ trong node không phải nhanh nhất. vLLM có AllReduce Kernel riêng: khi GPU trong node đọc trực tiếp được bộ nhớ nhau, các card đọc buffer của nhau rồi cộng, bớt vài vòng truyền vòng—thân thiện hơn với tin nhỏ Decode bị giới hạn độ trễ. Điều kiện không đủ hoặc tin lớn thì quay lại NCCL.

Gặp lỗi liên quan (ví dụ môi trường ảo hóa P2P không dùng được) có thể --disable-custom-all-reduce tắt nó, mọi AllReduce đi NCCL để khoanh vùng vấn đề.

💡 Gợi ý: Custom AllReduce không bật thì log khởi động in lý do—ví dụ P2P fail, hoặc hơn hai GPU chỉ nối PCIe. TP hiệu năng bất thường thì đọc đoạn log đó trước.

---

7. Tính tay: mô hình thô mức TP và thời gian mỗi Token

Dưới đây là mô hình chi phí giả định: chỉ tính thời gian đọc trọng số, đọc KV và giao tiếp TP của một bước Decode; bỏ tính toán, khởi động Kernel và lập lịch. $\alpha=10$ µs là giả định, không khớp đo thật nào. TP=1 và TP=2 thì mô hình 70B thực tế không vừa; ở đây chỉ để nhìn xu hướng.

# Mô hình chi phí giả định: chỉ ước lượng thời gian truy cập bộ nhớ và giao tiếp TP
# của một bước Decode; bỏ tính toán và khởi động Kernel
HBM = 3.35e12          # băng thông bộ nhớ H100 SXM, B/s
NVLINK = 450e9         # NVLink 4 một chiều, B/s (900 GB/s là tổng hai chiều)
ALPHA = 10e-6          # overhead cố định mỗi AllReduce, giả định
LAYERS, HIDDEN = 80, 8192
WEIGHT_BYTES = 70.6e9 * 2
KV_PER_TOKEN = 320 * 1024        # KV cả mô hình mỗi Token, B
CTX = 4096                        # độ dài ngữ cảnh trung bình mỗi request


def step_ms(tp, batch):
    weight = WEIGHT_BYTES / tp / HBM
    kv = batch * CTX * KV_PER_TOKEN / min(tp, 8) / HBM   # TP>8 thì KV không nhỏ thêm
    msg = batch * HIDDEN * 2
    comm = 0 if tp == 1 else 2 * LAYERS * (ALPHA + 2 * (tp - 1) / tp * msg / NVLINK)
    return (weight + kv + comm) * 1e3, comm * 1e3


print(f"{'TP':>3} {'batch':>5} {'step ms':>8} {'comm':>8} {'Token/s':>9} {'Token/s/GPU':>10}")
for batch in [1, 64]:
    for tp in [1, 2, 4, 8]:
        t, c = step_ms(tp, batch)
        print(f"{tp:>3} {batch:>5} {t:>8.2f} {c:>8.2f} {batch / t * 1e3:>9.0f} {batch / t * 1e3 / tp:>10.0f}")

Output:

 TP batch    step ms      comm   Token/s  Token/s/GPU
  1     1    42.55     0.00        24         24
  2     1    22.88     1.61        44         22
  4     1    12.25     1.61        82         20
  8     1     6.93     1.61       144         18
  1    64    67.79     0.00       944        944
  2    64    35.87     1.97      1784        892
  4    64    19.11     2.16      3350        837
  8    64    10.73     2.25      5967        746

Với giả định này:

• TP từ 1 đến 8, thời gian mỗi bước một request từ 42.6 ms xuống 6.9 ms, khoảng 6.1 lần, không đủ 8 lần. Chênh nằm ở 1.6 ms giao tiếp không nhỏ theo TP.
• Batch=64, thông lượng mỗi card ở TP=8 thấp hơn TP=2 khoảng 16%. Nếu dung lượng KV không phải nút thắt, cùng 8 card mở 4 bản sao TP=2 thì tổng thông lượng cao hơn.
• Đổi ALPHA thành 30 µs rồi xem lại TP=8 sẽ cảm nhận vì sao xuyên node hoặc không có NVLink thì TP không đáng.

Hệ thật còn thời gian tính toán, khác biệt hiệu quả Attention Kernel và overhead lập lịch. Mô hình chỉ để hiểu xu hướng; quyết định triển khai lấy theo đo vllm bench serve (mục 6.6).

---

📝 Tóm tắt

• TP dùng cặp “cắt cột + cắt hàng”; activation trung gian ở lại từng card; mỗi lớp 2 AllReduce, mô hình 80 lớp mỗi bước 160 lần.
• Attention cắt theo head; số Q head phải chia hết cho TP; mô hình GQA khi TP vượt số KV head thì vLLM nhân bản KV head, KV mỗi card không nhỏ thêm.
• Tin AllReduce của Decode chỉ mức KiB đến MiB, bị giới hạn độ trễ; Prefill tin lớn, bị giới hạn băng thông.
• NVLink 4 một chiều 450 GB/s; NIC 400 Gb/s xuyên node khoảng 50 GB/s, nên TP giữ trong node, xuyên node dùng PP.
• Tăng TP hạ độ trễ nhưng thông lượng mỗi card thường giảm; khi KV là nút thắt thì ngược lại—so sau khi chuẩn hóa theo số card.
• Custom AllReduce của vLLM tăng tốc tin nhỏ trong node; gặp sự cố dùng --disable-custom-all-reduce để khoanh vùng.

🎯 Danh sách tự kiểm

• Suy ra được vì sao FFN cắt cột + cắt hàng bằng kết quả gốc, và AllReduce đặt ở đâu
• Theo số Q head và KV head phán được mức TP dùng được, và từ mức TP nào bắt đầu nhân bản KV head
• Dùng $\alpha+\frac{2(t-1)}{t}\frac{S}{B}$ ước lượng một AllReduce và phán bị giới hạn độ trễ hay băng thông
• So được băng thông mỗi card của NVLink, PCIe, InfiniBand và giải thích vì sao TP không xuyên node
• Giải thích được vì sao TP đôi thì độ trễ giảm không đủ một nửa, thông lượng mỗi card có thể giảm
• Với hai mục tiêu “độ trễ thấp nhất” và “thông lượng cao nhất” chọn được mức TP và số bản sao
• Nói được lớp nào trong vLLM cắt theo vocab, và --disable-custom-all-reduce dùng để làm gì

📚 Tài liệu tham khảo

• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Efficiently Scaling Transformer Inference
• NVIDIA H100 Tensor Core GPU
• NVIDIA NVLink and NVLink Switch
• NVIDIA DGX B200 User Guide
• nccl-tests Performance (mô hình thời gian AllReduce và busbw)
• vLLM docs: Parallelism and Scaling
• vLLM source: vllm/model_executor/layers/linear.py (v0.32.0)
• vLLM source: vllm/distributed/device_communicators/custom_all_reduce.py (v0.32.0)