Hướng dẫn sử dụng SGLang trên Linux: Tăng tốc độ Inference LLM và tối ưu RadixAttention cho AI API

Artificial Intelligence tutorial - IT technology blog
Artificial Intelligence tutorial - IT technology blog

Bối cảnh: Khi TTFT trở thành nút thắt cổ chai

Nửa năm trước, team mình nâng cấp backend cho một AI agent chuyên đọc tài liệu pháp lý và chat đa lượt. Vấn đề lập tức xuất hiện: Time to First Token (TTFT) vọt lên hơn 1.8 giây mỗi khi có từ 15 người dùng đồng thời. Mỗi request đều cõng theo một system prompt dài 2.000 token kèm ngữ cảnh hợp đồng cố định. GPU liên tục phải tính toán lại KV cache cho cùng một chuỗi dữ liệu đầu vào. Lãng phí tài nguyên vô cùng.

Team đã thử vLLM lẫn TGI. Cả hai đều tốt, nhưng chưa giải quyết triệt để bài toán prefix dài bị lặp lại liên tục. Chuyển sang SGLang (engine từ nhóm LMSYS và UC Berkeley), kết quả thay đổi rõ rệt.

Nhờ cơ chế RadixAttention, SGLang quản lý KV cache dưới dạng cây Radix Tree. Thay vì vứt bỏ cache sau mỗi lượt sinh token, nó giữ lại và tái sử dụng cho các request sau có chung phần mở đầu. Với kịch bản RAG hay chatbot nhiều vòng, latency phản hồi token đầu tiên giảm một mạch từ 1.800ms xuống còn chưa đầy 220ms.

Nếu chi phí thuê GPU hàng tháng của bạn đang tăng nhanh theo lượng truy vấn, SGLang là phương án đáng để đưa lên bàn cân.

Chuẩn bị môi trường máy chủ Linux

Bạn cần một máy chủ Ubuntu (tốt nhất là 22.04 LTS) gắn GPU NVIDIA kiến trúc Ampere trở lên. Cụ thể, các dòng như RTX 3090, 4090, A10G, A100 hoặc H100 đều chạy rất mượt. Hãy chắc chắn máy đã cài sẵn NVIDIA Driver bản 535+ và CUDA 12.1 trở lên.

Khởi tạo một virtual environment riêng biệt với Python 3.10 hoặc 3.11:

# Tạo và kích hoạt môi trường ảo
python3 -m venv venv-sglang
source venv-sglang/bin/activate

# Cập nhật công cụ đóng gói
pip install --upgrade pip setuptools wheel

Cài đặt gói SGLang bản mới nhất kèm theo FlashInfer để tăng tốc nhân tính toán attention:

# Cài đặt sglang engine cùng flashinfer
pip install "sglang[all]>=0.4.0" --find-links https://flashinfer.ai/whl/cu121/torch2.4/flashinfer/

Chạy thử một dòng lệnh kiểm tra GPU:

python3 -c "import torch; print('CUDA Ready:', torch.cuda.is_available(), '| Device:', torch.cuda.get_device_name(0))"

Khởi chạy và Cấu hình Production

Khởi động OpenAI-compatible Server

Điểm tiện lợi lớn nhất của SGLang là hỗ trợ sẵn chuẩn API của OpenAI. Bạn có thể trỏ thẳng service backend vào endpoint mới mà không cần sửa logic ứng dụng. Dưới đây là lệnh host model Meta-Llama-3.1-8B-Instruct trên port 30000:

python3 -m sglang.launch_server \
    --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \
    --host 0.0.0.0 \
    --port 30000 \
    --mem-fraction-static 0.85 \
    --context-length 8192 \
    --log-level info

Tinh chỉnh các tham số ăn tiền

  • --mem-fraction-static: Dung lượng VRAM dành riêng cho KV cache sau khi nạp model weight. Mặc định là 0.88. Nếu trên máy chủ còn chạy kèm embedding model hoặc Celery worker, hãy hạ xuống 0.80 – 0.83 để tránh lỗi CUDA OOM đột ngột.
  • --tp (Tensor Parallelism): Dành cho mô hình cỡ lớn như Qwen-2.5-72B. Ví dụ với box 4x A100 (80GB), hãy gắn cờ --tp 4 để phân chia trọng số qua giao tiếp NVLink/NCCL tốc độ cao.
  • --schedule-policy lpm: Thuật toán Longest Prefix Match. Đây là linh hồn của RadixAttention: request nào có khả năng khớp KV cache cao nhất sẽ được xếp ưu tiên vào batch xử lý trước.

Test nhanh từ client backend

Chỉ cần dùng SDK openai quen thuộc:

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:30000/v1",
    api_key="EMPTY"  # Server nội bộ không cần pass key
)

# System prompt dài đóng vai trò là shared prefix cacheable
system_prompt = "Bạn là kỹ sư backend cao cấp. Luôn phản hồi ngắn gọn, ưu tiên code mẫu production-ready."

response = client.chat.completions.create(
    model="meta-llama/Meta-Llama-3.1-8B-Instruct",
    messages=[
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": "Cho ví dụ cấu hình connection pool trong asyncpg."}
    ],
    temperature=0.1,
    max_tokens=256
)

print(response.choices[0].message.content)

Kiểm thử Tải & Giám sát thực tế

Stress test với công cụ tích hợp

Đừng vội đưa lên production nếu chưa đo tải. SGLang có sẵn công cụ benchmark giả lập traffic thực tế cực kỳ tiện lợi:

python3 -m sglang.bench_serving \
    --backend sglang \
    --num-prompts 500 \
    --request-rate 25 \
    --dataset-name sharegpt \
    --random-input-len 1024 \
    --random-output-len 256

Hãy chú ý kỹ hai con số: Median TTFT và Output Throughput (tok/s). Trên cùng một con GPU A10G, khi các prompt chia sẻ khoảng 35% nội dung prefix, throughput của SGLang đạt ~720 token/s, cao hơn khoảng 2.2 lần so với khi tắt tính năng reuse cache.

Đẩy Metrics lên Grafana qua Prometheus

SGLang tự động mở endpoint tại http://0.0.0.0:30000/metrics. Hãy nối endpoint này vào hệ thống monitoring của bạn và dựng cảnh báo với 3 chỉ số then chốt:

  • sglang:radix_cache_hit_rate: Tỷ lệ trúng cache. Với ứng dụng chatbot hoặc RAG, tỷ lệ này lý tưởng nên duy trì ở mức 60% – 85%. Nếu tụt dưới 30%, bạn nên xem lại cách sắp xếp cấu trúc prompt đầu vào.
  • sglang:num_waiting_reqs: Hàng đợi request chờ xử lý. Con số này vượt quá 10 trong thời gian dài nghĩa là phần cứng đang quá tải. Lúc này bạn cần scale out thêm GPU instance hoặc tinh chỉnh lại batch size.
  • sglang:gpu_cache_usage_sys: Bộ nhớ KV cache đang chiếm dụng. Đồ thị chạm ngưỡng 95% liên tục là tín hiệu cần giảm --context-length hoặc nới thêm VRAM.

Nắm chắc các thông số này, bạn sẽ làm chủ hoàn toàn cụm serving LLM mà không lo hóa đơn đám mây nhảy vọt mất kiểm soát.

Share: