AMD ATOM Inference Engine: วิธีใช้ KV Cache แบบแบ่งหน้าและการคอมไพล์แบบแบ่งส่วนเพื่อเพิ่มประสิทธิภาพการอนุมานโมเดลใหญ่เป็นสองเท่า

2026年5月4日 pm8:16 • การอนุมานโมเดลขนาดใหญ่ • 46 views

เมื่อพูดถึงการปรับปรุงประสิทธิภาพการอนุมานของโมเดลขนาดใหญ่ สิ่งที่ดึงดูดความสนใจมากที่สุดมักจะเป็นประสิทธิภาพที่เพิ่มขึ้นของโอเปอเรเตอร์เดี่ยวๆ เช่น GEMM เร่งความเร็วได้เท่าไหร่ Attention เร็วขึ้นกี่เปอร์เซ็นต์ หรือ MoE routing ถูกผสานรวมแล้วหรือยัง อย่างไรก็ตาม สิ่งที่กำหนดประสบการณ์การให้บริการออนไลน์อย่างแท้จริงนั้น ไม่ใช่ตัวชี้วัดสูงสุดของเคอร์เนลเดี่ยวๆ แต่เป็นวงจรชีวิตที่สมบูรณ์ของคำขอตั้งแต่เข้าสู่ระบบผ่าน HTTP ผ่านกระบวนการ分词 การจัดตาราง การเขียนลงใน KV cache การดำเนินการตามกราฟ GPU การสุ่มตัวอย่าง ไปจนถึงการส่งคืนแบบสตรีม

แผนภูมิด้านล่างแสดงความสัมพันธ์แบบแลกเปลี่ยนระหว่างปริมาณงานต่อการ์ด (Token throughput) และเวลาแฝงแบบ end-to-end ของโมเดล DeepSeek R1 0528 ที่ความแม่นยำ FP8 บน GPU และเฟรมเวิร์กการอนุมานที่แตกต่างกัน แกนนอนแสดงถึงเวลาแฝง (วินาที) แกนตั้งแสดงถึงปริมาณงาน (tok/s/gpu) แนวโน้มโดยรวมแสดงให้เห็นว่าปริมาณงานเพิ่มขึ้นตามเวลาแฝงที่สูงขึ้น ซึ่งเน้นย้ำถึงการแลกเปลี่ยนประสิทธิภาพระหว่าง “ความเร็ว-เวลาแฝง” ทั่วไปในการอนุมานโมเดลขนาดใหญ่ โดยเฉพาะอย่างยิ่ง AMD MI355X ที่จับคู่กับเฟรมเวิร์ก ATOM มีประสิทธิภาพโดดเด่นที่สุด โดยมีปริมาณงานทะลุ 1100 tok/s/gpu ที่เวลาแฝงประมาณ 26 วินาที SGLang เวอร์ชันในซีรีส์เดียวกันอยู่ในอันดับที่สอง ในขณะที่ NVIDIA B200 (TRT) มีปริมาณงานต่ำกว่าอย่างเห็นได้ชัดที่เวลาแฝงใกล้เคียงกัน ซึ่งแสดงให้เห็นถึงความแตกต่างอย่างชัดเจนของประสิทธิภาพการอนุมานโมเดลนี้บนฮาร์ดแวร์และสแต็กการปรับให้เหมาะสมที่แตกต่างกัน ดูข้อมูลเพิ่มเติมได้ที่: https://inferencex.semianalysis.com/

ATOM (AiTer Optimized Model) เป็นการใช้งานสไตล์ vLLM ที่มีน้ำหนักเบา ซึ่งมุ่งเน้นไปที่การบูรณาการและการปรับให้เหมาะสมตาม AITER
ที่เก็บโค้ด: https://github.com/ROCm/ATOM
ที่อยู่เอกสาร: https://rocm.github.io/ATOM/docs
บทความประมาณ 5,000 คำ ใช้เวลาอ่าน 22 นาที ความยาวพอดแคสต์ 27 นาที

ATOM (AiTer Optimized Model) เป็นโปรเจกต์ดังกล่าวอย่างแท้จริง: มันไม่ใช่เฟรมเวิร์กการฝึก/อนุมานที่ครอบคลุมทุกอย่าง แต่เป็นเอนจิ้นการอนุมานสไตล์ vLLM ที่มีน้ำหนักเบา มันใช้เคอร์เนล AITER บนแพลตฟอร์ม AMD ROCm เพื่อเชื่อมต่อการดำเนินการโมเดล, KV cache แบบแบ่งหน้า, การจับภาพกราฟที่คอมไพล์, การทำงานแบบขนานหลาย GPU และอินเทอร์เฟซที่เข้ากันได้กับ OpenAI เข้าด้วยกันเป็นระบบที่สมบูรณ์ซึ่งสามารถรัน, ทดสอบโหลด และวิเคราะห์ได้ คุณค่าของมันไม่ได้อยู่ที่ “รันได้” เท่านั้น แต่ยังรวมถึงการทำให้เส้นทางการปรับให้เหมาะสมของการอนุมานโมเดลขนาดใหญ่บน GPU ของ AMD เป็นรูปธรรมในรูปแบบโครงสร้างทางวิศวกรรม

ตารางด้านล่างแสดงตระกูลโมเดลหลักที่ ATOM รองรับในปัจจุบัน รวมถึงสถาปัตยกรรม HuggingFace, ประเภท Dense และ Mixture of Experts (MoE) และหมายเหตุสำคัญ เพื่อให้เข้าใจขอบเขตความเข้ากันได้และข้อควรระวังได้อย่างรวดเร็ว

ตระกูลโมเดล	สถาปัตยกรรม HF	Dense/MoE	หมายเหตุ
Llama[1]	`LlamaForCausalLM`	Dense	Llama 2, Llama 3, Llama 3.1
Qwen3[2]	`Qwen3ForCausalLM`	Dense
Qwen3-MoE[3]	`Qwen3MoeForCausalLM`	MoE	128 ผู้เชี่ยวชาญ, top-8 routing
Qwen3-Next[4]	`Qwen3NextForCausalLM`	MoE	ผสม Full Attention + Gated DeltaNet
DeepSeek V2/V3[5]	`DeepseekV3ForCausalLM`	MoE	MLA Attention, MTP Speculative Decoding
Mixtral[6]	`MixtralForCausalLM`	MoE	8 ผู้เชี่ยวชาญ, top-2 routing
GLM-4-MoE[7]	`Glm4MoeForCausalLM`	MoE
GLM-5[8]	`GlmMoeDsaForCausalLM`	MoE	MLA Attention, คล้าย DeepSeek V3.2 ดูรายละเอียดเพิ่มเติมที่ recipe[9]
GPT-OSS[10]	`GptOssForCausalLM`	MoE	Sliding Window + Attention Aggregation
Kimi-K2[11]	ผ่าน `--trust-remote-code`	MoE	ดูรายละเอียดเพิ่มเติมที่ recipe[12]
MiMo-V2-Flash[13]	`MiMoV2FlashForCausalLM`	MoE	ผสม Full Attention + SWA, MTP 3 ชั้น ดูรายละเอียดเพิ่มเติมที่ recipe[14]

สารบัญบทความนี้

หนึ่ง. เริ่มต้นใช้งานอย่างรวดเร็ว: ทำให้ ATOM ทำงานก่อน
สอง. ตำแหน่งของโปรเจกต์: ไม่ใช่ “ไลบรารีโมเดล” แต่เป็นห้องตัวอย่างห่วงโซ่การอนุมาน ROCm
2.1 เป้าหมายหลักของ ATOM
2.2 ความสัมพันธ์ของส่วนประกอบโดยรวม
สาม. วงจรชีวิตของคำขอ: พรอมต์หนึ่งตัวเดินทางผ่าน ATOM ได้อย่างไร
3.1 จาก API สู่ Sequence
3.2 Sequence คือบัตรประจำตัวของคำขอในระบบ
สี่. ตัวจัดตาราง: การแลกเปลี่ยนระหว่าง prefill-first และ continuous batching
4.1 ทำไมต้อง prefill ก่อน
4.2 ขั้นตอน Decode: แย่งชิงเมื่อพื้นที่ KV ไม่เพียงพอ
ห้า. Paged KV Cache: การตัดหน่วยความจำ GPU ให้เป็นบล็อกที่นำกลับมาใช้ใหม่ได้
5.1 สาระสำคัญของ BlockManager
5.2 Prefix caching: การใช้ xxhash64 เพื่อระบุคำนำหน้าที่เหมือนกัน
หก. ชั้นโอเปอเรเตอร์: AITER คือรากฐานประสิทธิภาพของ ATOM
6.1 Linear: จ่ายงาน GEMM ตามประเภท quantization
6.2 Attention: การห่อหุ้มตรรกะไดนามิกที่ซับซ้อนเป็น custom op ที่แบ่งส่วนได้
เจ็ด. การคอมไพล์และ CUDA Graph: การลดโอเวอร์เฮด CPU ของ decode ให้เหลือน้อยที่สุด
7.1 ระดับการคอมไพล์สี่ระดับ
7.2 ทำไม prefill ใช้ eager, decode ใช้ graph replay
แปด. MoE, Quantization และหลายการ์ด: พื้นที่ความกดดันที่แท้จริงของ ATOM สำหรับโมเดลขนาดใหญ่
8.1 เส้นทาง MoE: Routing, Expert Parallelism และการสื่อสาร MORI
8.2 Quantization ไม่ใช่รูปแบบเดียว แต่เป็นชุดสัญญาการดำเนินการ
เก้า. การแยก P/D: การแยก prefill และ decode ไปยังโหนด GPU ที่แตกต่างกัน
สิบ. จะประเมิน ATOM อย่างไร: ความหนาแน่นของวิศวกรรมระบบภายใต้เปลือกที่มีน้ำหนักเบา
บทสรุป: ความสำคัญของ ATOM คือการผลักดันจาก “ROCm สามารถรันโมเดลขนาดใหญ่” ไปสู่ “ROCm สามารถให้บริการโมเดลขนาดใหญ่”

หนึ่ง. เริ่มต้นใช้งานอย่างรวดเร็ว: ทำให้ ATOM ทำงานก่อน

README อย่างเป็นทางการของ ATOM มีเส้นทางที่ง่ายที่สุด: ใช้ Docker image nightly โดยตรง

การตั้งค่าสภาพแวดล้อมและการเริ่มต้นใช้งานอย่างรวดเร็ว

ข้อกำหนดหลักในการรันเอนจิ้นการอนุมาน ATOM คือการมี AMD GPU, ติดตั้งซอฟต์แวร์ ROCm และเตรียมสภาพแวดล้อม Docker วิธีการปรับใช้ที่官方แนะนำมากที่สุดคือการดึง Docker image ที่สร้างไว้ล่วงหน้าโดยตรง ซึ่ง image นี้ได้รวม AITER และ dependencies ทั้งหมดของ ATOM ไว้แล้ว ตามคำอธิบายใน README.md ของโปรเจกต์ วิธีการเริ่มต้นมีดังนี้:

docker pull rocm/atom-dev:latest  

docker run -it --network=host   
  --device=/dev/kfd   
  --device=/dev/dri   
  --group-add video   
  --cap-add=SYS_PTRACE   
  --security-opt seccomp=unconfined   
  -v $HOME:/home/$USER   
  -v /mnt:/mnt   
  -v /data:/data   
  --shm-size=16G   
  --ulimit memlock=-1   
  --ulimit stack=67108864   
rocm/atom-dev:latest

หากนักพัฒนาต้องการสร้างด้วยตนเองจาก ROCm PyTorch image พื้นฐาน README ก็มีอีกเส้นทางหนึ่ง: ขั้นแรกให้เข้าไปใน image พื้นฐาน rocm/pytorch:rocm7.0.2_ubuntu24.04_py3.12_pytorch_release_2.8.0 จากนั้นติดตั้ง AITER และ ATOM

pip install amd-aiter  
git clone https://github.com/ROCm/ATOM.git && pip install ./ATOM

การรันตัวอย่างการอนุมานขั้นต่ำสามารถทำได้โดยตรงผ่านโมดูล Python:

python -m atom.examples.simple_inference --model meta-llama/Meta-Llama-3-8B --kv_cache_dtype fp8

สำหรับผู้ใช้ที่สนใจการปรับใช้ในรูปแบบบริการมากกว่า ATOM มีเซิร์ฟเวอร์ที่เข้ากันได้กับอินเทอร์เฟซ OpenAI ซึ่งสามารถเปิดเผยอินเทอร์เฟซคล้าย /v1/chat/completions และ /v1/completions ได้โดยตรง:

# สถานการณ์ GPU เดียว  
python -m atom.entrypoints.openai_server --model Qwen/Qwen3-0.6B --kv_cache_dtype fp8  

# สถานการณ์หลาย GPU, เปิดใช้งาน Tensor Parallelism  
python -m atom.entrypoints.openai_server --model deepseek-ai/DeepSeek-R1 --kv_cache_dtype fp8 -tp 8  

# เปิดใช้งาน MTP Speculative Decoding  
python -m atom.entrypoints.openai_server --model deepseek-ai/DeepSeek-R1 --kv_cache_dtype fp8 -tp 8   
  --method mtp --num-speculative-tokens 3

สำหรับรายละเอียดเพิ่มเติมเกี่ยวกับการออกแบบสถาปัตยกรรมที่สมบูรณ์ยิ่งขึ้น, พารามิเตอร์การกำหนดค่า, รายการโมเดลที่รองรับ, กลยุทธ์การจัดตาราง, กระบวนการคอมไพล์, แผนการปรับใช้หลายการ์ด และการทดสอบประสิทธิภาพ โปรดดูที่入口เอกสารโปรเจกต์อย่างเป็นทางการ: rocm.github.io/ATOM/docs[15]

สอง. ตำแหน่งของโปรเจกต์: ไม่ใช่ “ไลบรารีโมเดล” แต่เป็นห้องตัวอย่างห่วงโซ่การอนุมาน ROCm

2.1 เป้าหมายหลักของ ATOM

README ให้คำจำกัดความของ ATOM ไว้อย่างชัดเจน: มันคือการใช้งานสไตล์ vLLM ที่มีน้ำหนักเบา ซึ่งแกนหลักคือการบูรณาการและการปรับให้เหมาะสมตาม AITER[16] กล่าวอีกนัยหนึ่ง จุดเน้นของ ATOM ไม่ใช่การคิดค้นสถาปัตยกรรม Transformer ขึ้นมาใหม่ แต่เป็นการตอบคำถามเชิงระบบมากกว่า:

หาก AMD มีชุดเคอร์เนล ROCm ประสิทธิภาพสูงอยู่แล้ว จะจัดระเบียบพวกมันให้เป็นเอนจิ้นการอนุมานโมเดลขนาดใหญ่ที่ให้บริการได้จริงได้อย่างไร?

เป้าหมายนี้กำหนดจุดเน้นทางสถาปัตยกรรมของ ATOM ให้เป็นไปในทางวิศวกรรมอย่างมาก:

ชั้นโอเปอเรเตอร์: จัดการเส้นทางการคำนวณหนัก เช่น Linear, Attention, MoE, Norm, Sampling โดย AITER
ชั้นโมเดล: รองรับสถาปัตยกรรมหลักที่หลากหลาย เช่น Llama, Qwen3, DeepSeek, Mixtral, GLM, GPT-OSS, Kimi, MiMo
ชั้นการดำเนินการ: บรรลุการประมวลผลแบบกลุ่มต่อเนื่องและการจัดการ KV cache ผ่านส่วนประกอบ Scheduler, BlockManager, ModelRunner
ชั้นการคอมไพล์: มีระดับการคอมไพล์สี่ระดับตั้งแต่ 0 ถึง 3 สภาพแวดล้อมการผลิตใช้ torch.compile แบบแบ่งส่วนร่วมกับ CUDA graph เป็นค่าเริ่มต้น
ชั้นบริการ: เปิดเผย API ที่เข้ากันได้กับ OpenAI และมีเครื่องมือทดสอบประสิทธิภาพ, การวิเคราะห์ประสิทธิภาพ และการตรวจสอบความแม่นยำ
ชั้นกระจาย: รองรับ Tensor Parallelism (TP), Data Parallelism (DP), Expert Parallelism (EP), การสื่อสาร MoE all-to-all ตาม MORI และการส่ง KV แบบแยก P/D

ดังนั้น ATOM จึงเปรียบเสมือน “ห้องตัวอย่างระบบอนุมาน ROCm”: มันรวมเคอร์เนลประสิทธิภาพสูง, กลยุทธ์การคอมไพล์ และกลไกการอนุมานออนไลน์บน GPU ของ AMD เข้าด้วยกันในเฟรมเวิร์กที่ตรวจสอบได้

2.2 ความสัมพันธ์ของส่วนประกอบโดยรวม

เอกสารสถาปัตยกรรมอย่างเป็นทางการระบุห่วงโซ่หลักดังนี้: LLMEngine หันหน้าเข้าหาผู้ใช้, CoreManager รับผิดชอบการจัดระเบียบหลายกระบวนการ, EngineCore ถือตัวจัดตารางและตัวดำเนินการโมเดล, ModelRunner ขับเคลื่อนการคำนวณ forward ของโมเดลบน GPU แต่ละตัวจริงๆ

เราสามารถเปรียบเทียบสิ่งนี้กับขั้นตอนการทำงานของร้านอาหาร:

LLMEngine คือแผนกต้อนรับ รับผิดชอบรับคำขอจากผู้ใช้
InputOutputProcessor คือพนักงานรับออเดอร์ รับผิดชอบการ分词, การถอด分词 และงานสถิติ
Scheduler คือผู้จัดตารางในครัว ตัดสินใจว่างานใดควรดำเนินการก่อน
BlockManager คือผู้จัดการคลังสินค้า จัดการชั้นวางจัดเก็บ KV cache
ModelRunner คือพ่อครัว เรียกใช้เคอร์เนล AITER เพื่อคำนวณจริง
CoreManager และ ZMQ เปรียบเสมือนระบบส่งอาหาร กระจายคำขอไปยังกระบวนการต่างๆ

สาม. วงจรชีวิตของคำขอ: พรอมต์หนึ่งตัวเดินทางผ่าน ATOM ได้อย่างไร

3.1 จาก API สู่ Sequence

入口ผู้ใช้ของ ATOM ถูกกำหนดไว้ในไฟล์ atom/model_engine/llm_engine.py

เอาล่ะ นี่คือผลลัพธ์ของการเขียนใหม่เชิงลึกและการลดความซ้ำซ้อนของส่วนบทความที่ระบุตามความต้องการของคุณ

ขั้นตอนการดำเนินการของ LLMEngine.generate() เริ่มต้นจากการเรียก add_request() เมธอดนี้จะส่งพรอมต์ที่ผู้ใช้ป้อนและพารามิเตอร์การสุ่มตัวอย่างไปยัง InputOutputProcessor.preprocess() ซึ่งจะสร้างออบเจ็กต์ Sequence ภายใน จากนั้นออบเจ็กต์นี้จะถูกส่งต่อไปยัง CoreManager เพื่อดำเนินการต่อไป

# atom/model_engine/llm_engine.py  
def generate(  
    self,  
    prompts: list[str],  
    sampling_params: SamplingParams | list[SamplingParams],  
) -> list[str]:  
    # รีเซ็ตตัวนับ round-robin เพื่อให้แน่ใจว่า DP core dump ไม่สอดคล้องกัน  
    self.core_mgr._rr_counter = 0  

    self.add_request(prompts, sampling_params)  
    outputs = {}  
    while not self.is_finished() and (  
        self.core_mgr.is_alive() or self.core_mgr.is_rest()  
    ):  
        seqs = self.step()  
        outs = self.io_processor.postprocess(seqs)  
        outputs.update(outs)  

    outputs = [outputs[seq_id] for seq_id in sorted(outputs)]  
    return outputs

แม้ว่าโค้ดนี้จะดูเรียบง่าย แต่ก็สะท้อนรูปแบบการควบคุมหลักของ ATOM อย่างแม่นยำ: เธรดผู้ใช้ไม่ได้ขับเคลื่อนการอนุมานโมเดลโดยตรง แต่ผ่านลูปที่ทำงานอย่างต่อเนื่อง ดึงออบเจ็กต์ Sequence ที่ประมวลผลเสร็จแล้วจาก engine core จากนั้นจึงดำเนินการถอดรหัส (detokenize) และงานสถิติในภายหลัง

3.2 Sequence: ตัวระบุเฉพาะของคำขอในระบบ

ออบเจ็กต์ Sequence มีข้อมูลวงจรชีวิตที่สมบูรณ์ของคำขอ รวมถึง prompt token, token ที่สร้าง, KV block table, สถานะ, พารามิเตอร์การสุ่มตัวอย่าง รวมถึงตัวชี้วัดประสิทธิภาพหลัก เช่น เวลาที่มาถึง, เวลาที่สร้าง token แรก, เวลาที่ออก มันเป็นทั้งหน่วยปฏิบัติการพื้นฐานของระบบจัดตารางและหน่วยทางสถิติสำหรับการรวบรวมข้อมูลประสิทธิภาพ

การเปลี่ยนแปลงสถานะของคำขอใน ATOM โดยประมาณจะเป็นไปตามเส้นทางนี้:

WAITING -> RUNNING(PREFILL) -> RUNNING(DECODE) -> FINISHED

เมื่อพื้นที่ KV cache ไม่เพียงพอ คำขอที่กำลังทำงานอยู่อาจถูก抢占 (preempt) และถูกบังคับให้กลับไปยังคิวรอ เพื่อผ่านขั้นตอน prefill อีกครั้ง กลไกนี้ดูเหมือนจะทำให้เกิดการสิ้นเปลืองการคำนวณ แต่ในสถานการณ์ที่มีการทำงานพร้อมกันสูง มันทำให้ระบบมีความสามารถในการรักษาปริมาณงานโดยรวมที่สูงภายใต้ทรัพยากร KV cache ที่จำกัด

สี่. ตัวจัดตาราง: การแลกเปลี่ยนระหว่าง prefill-first และ continuous batching

4.1 ทำไมต้อง优先执行 prefill

กระบวนการอนุมานโมเดลขนาดใหญ่แบ่งออกเป็นสองขั้นตอน: ขั้นตอน prefill จะประมวลผล prompt อินพุตที่สมบูรณ์ มีปริมาณการคำนวณสูงและรูปร่างอินพุตที่แปรผัน; ขั้นตอน decode โดยปกติจะสร้าง token เพียงหนึ่งหรือไม่กี่ token ในแต่ละครั้ง การดำเนินการบ่อยครั้ง ไวต่อเวลาแฝง และเหมาะสมอย่างยิ่งสำหรับการปรับให้เหมาะสมผ่านการเล่นซ้ำด้วย CUDA graph

เมธอด Scheduler.schedule() ของ ATOM ใช้กลยุทธ์ prefill-first: ตราบใดที่มีคำขอที่รออยู่ซึ่งสามารถจัดตารางได้ ระบบจะ优先将它们组成一个 prefill batch; เฉพาะเมื่อไม่มีคำขอ prefill ที่ต้องดำเนินการเท่านั้น จึงจะจัดตารางงาน decode

ตรรกะโค้ดหลักมาจาก atom/model_engine/scheduler.py:

# atom/model_engine/scheduler.py  
def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]:  
    """เลือก batch ถัดไปของ sequences สำหรับ forward pass  

    พยายาม prefill ก่อน; หากไม่มี prefill ใหม่ที่พร้อม ให้ถอยไปที่  
    การถอดรหัส sequences ที่กำลังทำงานอยู่  
    """  
    scheduled_seqs = {}  
    num_seqs_prefill = 0  
    num_batched_tokens = 0  
    skipped_waiting_requests: deque[Sequence] = deque()  
    num_scheduled_tokens: list[int] = []  
    scheduled_spec_decode_tokens: dict[int, np.ndarray] = {}  

    if not self.running and not self.waiting:  
        return None  

    # --- การจัดตาราง Prefill ---  
    while self.waiting and num_seqs_prefill < self.max_num_seqs:  
        seq = self.waiting.popleft()  

        num_new_tokens = seq.num_tokens - seq.num_cached_tokens  
        if (  
            num_batched_tokens + num_new_tokens > self.max_num_batched_tokens  
            or not self.block_manager.can_allocate(seq)  
        ):  
            self.waiting.appendleft(seq)  
            break  

        self.block_manager.allocate(seq)  

        num_seqs_prefill += 1  
        num_new_tokens = seq.num_tokens - seq.num_cached_tokens  
        if self.cache_stats:  
            self.cache_stats.update(seq.num_cached_tokens, seq.num_tokens)  
        num_batched_tokens += num_new_tokens  
        seq.status = SequenceStatus.RUNNING  
        seq.type = SequenceType.PREFILL  
        self.running.append(seq)  
        scheduled_seqs[seq.id] = seq  
        num_scheduled_tokens.append(num_new_tokens)

แกนหลักของตรรกะการจัดตารางนี้อยู่ที่การควบคุมงบประมาณสำคัญสองประการ:

4.2 ขั้นตอน Decode: แย่งชิงเมื่อพื้นที่ KV ไม่เพียงพอ

เมื่อตัวจัดตารางไม่พบงาน prefill ใดๆ มันจะเปลี่ยนไปยังขั้นตอน decode ในขั้นตอนนี้ มันจะตรวจสอบแต่ละ sequence ที่กำลังทำงาน (running sequence) ทีละรายการว่ายังมีพื้นที่สำหรับ追加 KV block ใหม่หรือไม่ หากพื้นที่ไม่เพียงพอ มันจะ抢占คำขอที่อยู่ท้ายคิว

# atom/model_engine/scheduler.py  
# --- การจัดตาราง Decode ---  
num_seqs_decode = 0  
num_new_tokens = self.mtp_k + 1  
while self.running and num_seqs_decode < self.max_num_seqs:  
seq = self.running.popleft()  
while not self.block_manager.can_append(seq, num_new_tokens):  
if self.running:  
self.preempt(self.running.pop())  
else:  
self.preempt(seq)  
break  
else:  
if seq.spec_token_ids.size > 0:  
scheduled_spec_decode_tokens[seq.id] = seq.spec_token_ids  
num_seqs_decode += 1  
if not getattr(seq, "is_first_decode", False):  
self.block_manager.may_append(seq, num_new_tokens)  
scheduled_seqs[seq.id] = seq  
seq.type = SequenceType.DECODE  
num_scheduled_tokens.append(num_new_tokens)  
seq.is_first_decode = False

ในโค้ดนี้ การออกแบบ num_new_tokens = self.mtp_k + 1 นั้นชาญฉลาดมาก เมื่อเปิดใช้งาน MTP speculative decoding แล้ว แต่ละขั้นตอน decode ไม่เพียงแต่ต้อง预留พื้นที่ KV สำหรับ target token สุดท้ายเท่านั้น แต่ยังต้องจัดสรรตำแหน่ง缓存ให้กับ draft token ล่วงหน้าด้วย ตัวจัดตาราง本身ไม่สนใจว่าโมเดลจะตรวจสอบ draft token ภายในอย่างไร แต่ต้องแน่ใจว่า token เหล่านี้ที่อาจถูกเขียนมีตำแหน่ง缓存ที่พร้อมใช้งาน

ห้า. Paged KV Cache: การตัดหน่วยความจำ GPU ให้เป็นบล็อกที่นำกลับมาใช้ใหม่ได้

5.1 สาระสำคัญของ BlockManager

ในสถานการณ์การอนุมานโมเดลขนาดใหญ่ คอขวดของหน่วยความจำ GPU มักไม่ใช่น้ำหนักโมเดล แต่เป็น KV cache ATOM จัดการ KV block ขนาดคงที่ผ่าน BlockManager โดยค่าเริ่มต้นแต่ละ block มีขนาด 16 token แต่ละ Sequence จะรักษา block_table ไว้ ซึ่งบันทึก block ทั้งหมดที่มันครอบครอง

โครงสร้างข้อมูลของ Block นั้น精简มาก:

# atom/model_engine/block_manager.py  
class Block:  

def __init__(self, block_id):  
self.block_id = block_id  
self.ref_count = 0  
self.hash = -1  
self.token_ids = []  

def update(self, hash: int, token_ids: list[int]):  
self.hash = hash  
self.token_ids = token_ids  

def reset(self):  
self.ref_count = 1  
self.hash = -1  
self.token_ids = []

เราสามารถจินตนาการว่า KV cache เป็นโรงแรม: เมื่อคำขอมาถึงก็จัดสรรห้อง เมื่อคำขอสิ้นสุดก็เช็คเอาท์; หากแขกสองคนมี prompt นำหน้าที่เหมือนกันทุกประการ พวกเขาสามารถแชร์ห้องสองสามห้องแรกได้ด้วยซ้ำ ที่นี่ ref_count คือจำนวนผู้เข้าพักที่บันทึกห้องที่ใช้ร่วมกัน

5.2 Prefix caching: การใช้ xxhash64 เพื่อระบุคำนำหน้าที่เหมือนกัน

กลไก prefix caching ของ ATOM ใช้ xxhash64 และใช้กลยุทธ์链式哈希 (chained hash): ค่าแฮชของ block ปัจจุบันจะรวมค่าแฮชของ block ก่อนหน้า ด้วยวิธีนี้ แม้ว่าเนื้อหา token ของสอง block จะเหมือนกันทุกประการ แต่ตราบใดที่บริบท (context) ที่它们อยู่ต่างกัน (เช่น ข้อความก่อนหน้าต่างกัน) ค่าแฮชสุดท้ายก็จะต่างกัน จึงหลีกเลี่ยงการแชร์บริบทที่ผิดพลาด

# atom/model_engine/block_manager.py  
@classmethod  
def compute_hash(cls, token_ids: list[int], prefix: int = -1):  
h = xxhash.xxh64()  
if prefix != -1:  
h.update(prefix.to_bytes(8, "little"))  
h.update(np.array(token_ids).tobytes())  
return h.intdigest()

นี่ไม่ใช่แค่ “การทำแฮชกับรายการ token” ธรรมดา 链式哈希ช่วยให้แน่ใจว่าความหมายของแต่ละ KV block ผูกพันอย่างแน่นหนากับบริบทที่มันอยู่ สำหรับสถานการณ์ต่างๆ เช่น การสนทนาหลายรอบ, พรอมต์ที่คล้ายกันจำนวนมาก และการ复用เทมเพลต RAG กลไกนี้สามารถลดการคำนวณ prefill ที่ซ้ำซ้อนได้อย่างมาก ซึ่งช่วยเพิ่มปริมาณงานได้อย่างมาก

หก. ชั้นโอเปอเรเตอร์: AITER คือรากฐานประสิทธิภาพของ ATOM

6.1 Linear: จ่ายงาน GEMM ตามประเภท quantization

ไฟล์ atom/model_ops/linear.py ของ ATOM เป็นไฟล์หลักในการทำความเข้าใจกลยุทธ์โอเปอเรเตอร์ของมัน ไฟล์นี้จะจ่ายงานการคำนวณไปยังการใช้งาน GEMM ที่สอดคล้องกันใน AITER ตามรูปแบบ quantization ที่แตกต่างกัน:

ไม่มี quantization: tgemm.mm
per-tensor FP8: tgemm.mm พร้อม scale
per-token INT8: gemm_a8w8
per-token FP8: gemm_a8w8_bpreshuffle
per-1×128 FP8: gemm_a8w8_blockscale_bpreshuffle
per-1×32 MXFP4: gemm_a4w4

ตรรกะการจัดตารางหลักมีดังนี้:

# atom/model_ops/linear.py  
@mark_trace  
def forward(  
self, x: torch.Tensor, x_scale: Optional[torch.Tensor] = None, otype=dtypes.bf16  
) -> torch.Tensor:  
if self.quant_type.value == QuantType.No.value:  
y = tgemm.mm(  
x,  
self.weight,  
self.bias,  
otype=otype,  
)  
else:  
if x_scale is None:  
quant_func = self.quant_func  
if self.quant_type.value == QuantType.per_1x128.value:  
quant_func = functools_partial(  
self.quant_func, transpose_scale=True  
)  
if self.quant_type.value != QuantType.per_1x32.value:  
x, x_scale = quant_func(  
x,  
quant_dtype=self.params_dtype,  
scale=getattr(self, "input_scale", None),  
)  
if self.quant_type.value == QuantType.per_Tensor.value:  
y = tgemm.mm(  
x,  
self.weight,  
self.bias,  
otype=otype,  
scale_a=x_scale,  
scale_b=self.weight_scale,  
)  
elif self.quant_type.value == QuantType.per_Token.value:  
if self.params_dtype == dtypes.i8:  
y = gemm_a8w8(  
x,  
self.weight,  
x_scale,  
self.weight_scale,  
self.bias,  
dtype=otype,  
)  
else:  
y = gemm_a8w8_bpreshuffle(  
x,  
self.weight,  
x_scale,  
self.weight_scale,  
dtype=otype,  
)  
elif self.quant_type.value == QuantType.per_1x128.value:  
y = gemm_a8w8_blockscale_preshuffle_impl(  
x,  
self.weight,  
x_scale,  
self.weight_scale,  
dtype=otype,  
prefix=self.prefix,  
)  
elif self.quant_type.value == QuantType.per_1x32.value:  
y = gemm_a4w4_quant(  
x,  
x_scale,  
self.weight,  
otype,  
self.weight_scale.data,  
self.params_dtype,  
getattr(self, "input_scale", None),  
self.output_size,  
)  
if self.tp_dim == 1 and self.tp_size > 1 and self.reduce_results:  
y = get_tp_group().all_reduce(y, ca_fp8_quant=False)  
return y

มีแนวคิดการออกแบบหลักสองประการซ่อนอยู่ที่นี่:

ประการแรก ATOM ไม่ได้มองว่า quantization เป็น “การดำเนินการแปลงที่ทำครั้งเดียวเมื่อโหลดโมเดล” แต่ในระหว่างการ forward ของแต่ละชั้น (layer forward) จะคงเส้นทางที่ชัดเจนของ scale, weight layout และ kernel dispatch ไว้เสมอ

ประการสอง ตรรกะที่เกี่ยวข้องกับ Tensor Parallel ก็ถูกรวมเข้าภายใน Linear wrapper เพื่อจัดการเช่นกัน ตัวอย่างเช่น การดำเนินการ all-reduce ที่จำเป็นสำหรับเอาต์พุต row parallel จะเสร็จสมบูรณ์ที่นี่ ทำให้ไฟล์โมเดลสามารถรักษาโครงสร้างที่ค่อนข้าง简洁ได้