MetaX TileLang Inference Learning Framework

这个仓库是一个面向学习的推理引擎实验框架：用真实 vLLM、真实 Qwen 模型和真实 MetaX C500 环境，演示如何把 TileLang decode attention kernel 接进 vLLM attention backend。

目标不是做一个完整生产 serving 系统，而是让新人能沿着一条清晰路径理解：

• vLLM 如何选择和调用 attention backend。
• Prefill 和 decode 在推理链路里如何分流。
• TileLang PageAttention decode kernel 如何读取 paged KV cache。
• Qwen GQA head padding/reorder 为什么会影响 correctness 和性能。
• 怎么用 cold/warm、TTFT、decode tok/s、batch e2e tok/s 观察性能。

当前主线很简单：

Stage	Implementation	Why
Prefill	vllm_metax Maca FlashAttention fallback	保证真实模型稳定可用
Decode	TileLang PageAttention	学习自定义 decode kernel 接入
Backend	vLLM CUSTOM attention backend	和 FLASH_ATTN baseline 对比

Quick Start

在当前机器上，1.5B、3B、7B 模型都已经放在 /data/huggingface_home/models。第一次上手建议只跑单卡 1.5B。

# 1. 进入项目
cd /data/pd-sep-tutorial-metax

# 2. 验证 TileLang kernel correctness
python tests/test_compile_tilelang.py

# 3. 用 CUSTOM backend 跑单卡 1.5B
python serve.py \
  --profile single \
  --attention-backend CUSTOM \
  --max-tokens 32 \
  --temperature 0.0 \
  --prompt '请用一句话介绍人工智能。' \
  --enforce-eager

# 4. 用 FLASH_ATTN baseline 对比
python serve.py \
  --profile single \
  --attention-backend FLASH_ATTN \
  --max-tokens 32 \
  --temperature 0.0 \
  --prompt '请用一句话介绍人工智能。' \
  --enforce-eager

看到类似下面的中文输出，就说明端到端链路已经跑通：

人工智能是模拟人类智能的技术，使计算机能够执行需要人类智能的任务...

Common Runs

# 内置学习 prompt，适合快速 smoke test
python serve.py --profile single --auto --max-tokens 64 --enforce-eager

# 交互式聊天
python serve.py --profile single --interactive --max-tokens 128 --enforce-eager

# 双卡 3B
python serve.py --profile dual --attention-backend CUSTOM --max-tokens 32 --enforce-eager

# 四卡 7B
python serve.py --profile quad --attention-backend CUSTOM --max-tokens 32 --enforce-eager

# batch decode 基础路径
python serve.py --profile single --attention-backend CUSTOM --max-tokens 8 --batch-size 2 --repeat 2 --enforce-eager

Model Profiles

serve.py --profile ... 会自动设置模型路径、tensor parallel 和默认 max_model_len。

Profile	Cards / TP	Model	Default path	Status
single	1	Qwen2.5-1.5B-Instruct	/data/huggingface_home/models/Qwen2.5-1.5B-Instruct	verified
dual	2	Qwen2.5-3B-Instruct	/data/huggingface_home/models/Qwen2.5-3B-Instruct	verified
quad	4	Qwen2.5-7B-Instruct	/data/huggingface_home/models/Qwen2.5-7B-Instruct	verified

也可以显式覆盖：

python serve.py --model /path/to/model --tp 1 --max-model-len 512 --attention-backend CUSTOM

Model Setup

当前机器可以直接验证模型目录：

du -sh /data/huggingface_home/models/Qwen2.5-1.5B-Instruct \
  /data/huggingface_home/models/Qwen2.5-3B-Instruct \
  /data/huggingface_home/models/Qwen2.5-7B-Instruct

换新机器时，推荐下载到固定目录：

mkdir -p /data/huggingface_home/models

huggingface-cli download Qwen/Qwen2.5-1.5B-Instruct --local-dir /data/huggingface_home/models/Qwen2.5-1.5B-Instruct
huggingface-cli download Qwen/Qwen2.5-3B-Instruct --local-dir /data/huggingface_home/models/Qwen2.5-3B-Instruct
huggingface-cli download Qwen/Qwen2.5-7B-Instruct --local-dir /data/huggingface_home/models/Qwen2.5-7B-Instruct

下载后检查：

python -c 'from pathlib import Path; root=Path("/data/huggingface_home/models"); names=["Qwen2.5-1.5B-Instruct","Qwen2.5-3B-Instruct","Qwen2.5-7B-Instruct"]; [print(name, "OK" if (root/name/"config.json").exists() else "MISSING") for name in names]'

Benchmarks

serve.py 会输出这些学习指标：

Metric	Meaning
latency_ms	一次 llm.generate() 端到端耗时
tokens	本次生成 token 数
e2e_tok/s	端到端 token/s，包含 prefill 和调度开销
ttft_ms	Time To First Token
decode_tok/s	基于首末 token 时间估算的 decode token/s

推荐用 bench.py 做可复现对比：

# 1. 跑 1.5B / 3B / 7B 的 CUSTOM vs FLASH_ATTN
python bench.py \
  --profiles single dual quad \
  --backends CUSTOM FLASH_ATTN \
  --max-tokens 8 \
  --repeat 3 \
  --timeout 900 \
  --json benchmark_results.json

# 2. 扫 batch size
python bench.py \
  --profiles single \
  --backends CUSTOM FLASH_ATTN \
  --max-tokens 8 \
  --repeat 5 \
  --warmup-runs 2 \
  --batch-sizes 1 2 4 \
  --timeout 900 \
  --json benchmark_batch_results.json

# 3. 从已有 JSON 生成 Markdown 表格
python bench.py --report-from benchmark_results.json --markdown report.md

--repeat 会在同一个 vLLM engine 内连续生成，用来区分 cold first run 和 warm steady state。--warmup-runs 控制 warm_summary 丢弃多少个初始 repeat。

2026-06-11 的短输出 benchmark 摘要：

Profile	CUSTOM warm decode P50	FLASH_ATTN warm decode P50
single / 1.5B / TP=1	97.26 tok/s	110.97 tok/s
dual / 3B / TP=2	51.40 tok/s	52.40 tok/s
quad / 7B / TP=4	60.12 tok/s	64.94 tok/s

batch decode 摘要：

Batch	CUSTOM warm batch e2e P50	FLASH_ATTN warm batch e2e P50
2	149.63 tok/s	153.79 tok/s
4	279.57 tok/s	313.50 tok/s

完整 benchmark 说明见 docs/benchmark.md。

Project Map

.
├── serve.py                         # 主入口：generate / auto / interactive / batch
├── bench.py                         # benchmark runner + JSON/Markdown report
├── profiles.py                      # 1.5B / 3B / 7B profile 与模型路径
├── vllm_integration/
│   └── tilelang_backend.py           # vLLM CUSTOM backend：prefill fallback + TileLang decode
├── tilelang_kernels/
│   └── attention/
│       ├── flash_attn_tilelang.py    # TileLang FlashAttention，对照与 prefill 实验
│       └── page_attn_tilelang.py     # TileLang PageAttention decode kernel
├── tests/
│   └── test_compile_tilelang.py      # attention kernel correctness test
└── docs/
    ├── architecture.md               # 架构和 decode contract
    └── benchmark.md                  # benchmark 方法和记录

Learning Path

建议按这个顺序读代码：

1. serve.py：看 vLLM 如何加载模型、构造 prompt、选择 backend。
2. profiles.py：看 single/dual/quad profile 如何映射模型和 TP。
3. vllm_integration/tilelang_backend.py：看 TileLangAttentionImpl.forward() 如何分 prefill/decode。
4. tilelang_kernels/attention/page_attn_tilelang.py：看 paged KV cache、GQA、softmax 和 V 聚合。
5. tests/test_compile_tilelang.py：看如何用 PyTorch reference 验证 kernel correctness。
6. docs/architecture.md：看整体 runtime flow 和 decode kernel contract。
7. docs/benchmark.md：看 cold/warm、batch 和性能解释。

Current Status

已经完成：

• 1.5B / 3B / 7B 都能用 CUSTOM backend 端到端生成。
• Decode 阶段接入 TileLang PageAttention。
• Prefill 默认使用 Maca FlashAttention fallback，保证稳定。
• Qwen GQA head padding/reorder 已修正。
• batch decode 基础路径已接入，batch=2 接近 baseline，batch=4 仍有约 11% 差距。
• bench.py 支持 cold/warm JSON 记录、batch sweep 和 Markdown report。

还在继续优化：

• 更长输出、多 prompt 的正式 benchmark。
• batch=4+ decode 吞吐，重点看 kernel batch 维并行度和 vLLM 调度同步成本。
• TileLang kernel cache warning 的日志降噪。
• 实验性 TileLang prefill 路径稳定化。
• 更多 metadata、head padding、batch decode 测试。

Debug Flags

# 打印 decode block table、seq len 等信息
TILELANG_DEBUG_DECODE=1 python serve.py --profile single --attention-backend CUSTOM --max-tokens 8 --enforce-eager

# 实验性启用 TileLang prefill；默认关闭
TILELANG_ENABLE_PREFILL=1 python serve.py --profile single --attention-backend CUSTOM --max-tokens 8 --enforce-eager

常见判断：

• 日志出现 Using AttentionBackendEnum.CUSTOM backend. 表示 vLLM 已选中自定义 backend。
• 第一次运行慢通常是 cold path，包括模型 worker、TileLang materialize 和 cache 建立。
• 判断性能时优先看 bench.py --repeat 的 warm summary。
• 生成重复或语义异常时，优先检查 GQA head padding/reorder 和 slice-back 逻辑。