เร่งสปีด Video Diffusion บน TPU ด้วยเทคนิค Sparse Attention

30 ก.ย. 2026 โดยปกติแล้วโมเดล Video Diffusion มักประสบปัญหาการสร้างวิดีโอที่ล่าช้าจากสองปัจจัยหลัก คือจำนวนขั้นตอนการ Denoising ที่มาก และต้นทุนการคำนวณที่สูงในแต่ละขั้นตอน เมื่อความยาวของลำดับวิดีโอ (Sequence length) เพิ่มขึ้น Self-attention จะกลายเป็นปัจจัยสำคัญที่ทำให้เกิดความหน่วง ตัวอย่างเช่น วิดีโอความละเอียด 720p เพียง 81 เฟรม อาจมี Sequence length สูงถึง 50K ถึง 400K ในโมเดลยอดนิยมปัจจุบัน

ตารางที่ 1: ความหน่วงของ Self-attention และสัดส่วนของความหน่วงใน Transformer block ชั้นเดียวในแต่ละความละเอียด (วิดีโอ 81 เฟรม บนชิป TPU v6e แปดตัว)

การขยายขนาดจาก 720p เป็น 1440p (2K) ส่งผลให้ Sequence length เพิ่มขึ้นถึงสี่เท่า และเนื่องจาก Full attention ขยายตัวแบบยกกำลังสองตามความยาวลำดับ สัดส่วนความหน่วงต่อชั้นจึงอาจพุ่งสูงขึ้นจาก 55.5% เป็น 88.2% เนื่องจาก Attention เป็นสาเหตุหลักของความหน่วง การเพิ่มประสิทธิภาพการอนุมาน (Inference optimization) จึงต้องมุ่งเน้นไปที่จุดนี้เป็นอันดับแรก ซึ่งใน Video Diffusion นั้น Attention มีโครงสร้างที่ชัดเจน โดยการโต้ตอบส่วนใหญ่มีน้ำหนักน้อยและสามารถข้ามได้ ในขณะที่ Dense attention จะคำนวณทุกการโต้ตอบ แต่ Sparse attention จะใช้ Mask เพื่อคงไว้เฉพาะส่วนที่สำคัญเท่านั้น

รูปที่ 1: น้ำหนัก attention สำหรับ spatial head (Head 12) ในการจัดเรียงแบบเน้นเฟรม (ซ้าย) attention จะรวมตัวกันอย่างหนาแน่นตามเฟรมท้องถิ่น ในการจัดเรียงแบบลำดับเวลา (ขวา) น้ำหนักเดียวกันนั้นจะดูเหมือนกระจายไปตาม tokens ซึ่งแสดงให้เห็นว่าทำไมการจัดเรียง token ต้องสอดคล้องกับประเภทของ head

อย่างไรก็ตาม รูปแบบของ Attention จะแตกต่างกันไปตาม Head และชั้นการทำงาน งานวิจัย Sparse VideoGen (SVG) ได้จำแนก Head ออกเป็น Spatial heads ที่สนใจข้อมูลภายในเฟรมเดียวกันหรือใกล้เคียง และ Temporal heads ที่เน้นการติดตามความเคลื่อนไหวผ่านเฟรมจำนวนมากในพื้นที่ขนาดเล็ก

รูปที่ 2: ความเชี่ยวชาญเฉพาะด้านของ attention head (Step 6, Layer 32) กล่องสีฟ้าทำเครื่องหมายตำแหน่งเชิงพื้นที่ (y,x) ของ query token บนเฟรม F06 ใน spatial head (บน) 94.8% ของน้ำหนัก attention จะกระจายกว้างไปทั่วเฟรมของ query สำหรับการสังเคราะห์แบบ 2D ใน temporal head (ล่าง) attention จะรวมตัวตามท่อเชิงพื้นที่ (y,x) เดียวกันนั้นผ่านเฟรมก่อนหน้า (F04 - F05) เพื่อติดตามการเคลื่อนไหวเฉพาะจุด

SVG ใช้วิธีการ Profiling attention heads แบบไดนามิกเพื่อเลือก Sparse mask ที่เหมาะสมที่สุด ซึ่งช่วยรักษาคุณภาพของวิดีโอไว้ได้ในขณะที่มีความเบาบาง (Sparsity) สูง โดยทั้ง Spatial และ Temporal mask จะยังคงรักษา Full attention ไปยังเฟรมแรก (F00) เพื่อทำหน้าที่เป็น Attention sink ในการยึดเหนี่ยวรูปลักษณ์ของฉาก

รูปที่ 3: Sparse attention masks ในการจัดเรียงแบบเน้นเฟรม ซ้าย: Spatial mask (แถบท้องถิ่นที่ต่อเนื่องกัน + sink เฟรมแรกที่ F00) ขวา: Temporal mask ซึ่งปรากฏเป็นเส้นทแยงมุมแบบละเอียดในการจัดเรียงแบบเน้นเฟรม ซึ่งเป็นแรงจูงใจในการสลับเปลี่ยน token ของเราในภายหลัง

การนำเทคนิคนี้มาใช้งานจริงบนฮาร์ดแวร์ต้องอาศัยการปรับปรุง Kernel โดยเราได้เปรียบเทียบการทำงานของ JAX และ Pallas Splash Attention บน TPU v6e ซึ่งผลลัพธ์เบื้องต้นพบว่าการใช้ Sparse mask แบบทั่วไป (B1) กลับมีความหน่วงสูงกว่า Dense attention ถึง 22% เนื่องจากการประเมิน Mask ในทุก Tile กลายเป็นคอขวดของระบบ

ตารางที่ 2: ความก้าวหน้าของการเพิ่มประสิทธิภาพ sparse Splash Attention kernel บนชิป TPU v6e ตัวเดียว (75,600 tokens, 10 heads, มิติ head 128)

รูปที่ 4: กลยุทธ์การท่องไปใน tile ที่สอดคล้องกับตารางที่ 2 (จากซ้ายไปขวา: B0 Dense, B1 Exact sparse, B2 Classified tiles, และ B3 Balanced tile rounding)

รูปที่ 5: การจำแนก hardware tile สำหรับ sparse local-window attention. Tile จะถูกแบ่งออกเป็น full tiles (ดำเนินการผ่าน fast path ที่ไม่มี mask), boundary tiles (ต้องใช้ elementwise coordinate masking) และ skipped tiles (ถูกละเว้นทั้งหมดเพื่อประหยัดการคำนวณและ bandwidth) เพื่อแก้ปัญหานี้ เราได้พัฒนาแนวทางใหม่โดยการแยกประเภทของ Tile เป็น Full tiles, Boundary tiles และ Empty tiles โดยมอบเส้นทาง Fast path ให้กับ Full tiles ซึ่งช่วยลดความหน่วงลงได้ 44% (B2) และเมื่อปรับจูนการจัดแนว Mask ให้ตรงกับขนาด Tile (B3) ความหน่วงก็ลดลงเหลือเพียง 32.76 ms ซึ่งเร็วกว่า Dense Splash ถึง 2.40 เท่า

รูปที่ 6: ขอบเขต tile ขนาดใหญ่เทียบกับขนาดเล็ก เส้นทึบสีดำลากตามขอบเขตของ SVG mask (รวมถึง sink เฟรมแรกในแนวตั้งด้านซ้าย) tile ขนาดเล็กจะตามขอบเขตได้ใกล้เคียงกว่า ช่วยลด boundary tiles สีเหลืองโดยแลกกับประสิทธิภาพของ compute tile ที่เล็กลง

รูปที่ 7: กราฟเส้นแสดงความหน่วงของ kernel ในหน่วยมิลลิวินาทีเทียบกับเปอร์เซ็นต์ของงานที่ต้องการ intra-tile masking ในขนาด query tile สี่ขนาดบน TPU v6e แสดงให้เห็นถึงความหน่วงต่ำสุดที่ BQ เท่ากับ 3328

สำหรับการใช้งานในการอนุมานแบบกระจาย (Distributed inference) เราได้จัดวาง Token layout แบบไดนามิกเพื่อให้สอดคล้องกับประเภทของ Head โดยการสลับการจัดเรียงจากเน้นเฟรม (F, H, W) เป็นเน้นเวลา (H, W, F) ซึ่งช่วยให้ Kernel สามารถประมวลผล Temporal mask ได้อย่างมีประสิทธิภาพมากขึ้น โดยไม่เพิ่มภาระในการสื่อสารระหว่างอุปกรณ์

รูปที่ 8: end-to-end execution pipeline สำหรับ dynamic spatio-temporal attention การทำ dynamic routing จะประเมิน query ที่สุ่มตัวอย่างเพื่อสร้าง boolean flags ต่อ head ซึ่งจะเลือก token layouts บน tensors ที่มีรูปทรงคงที่ ก่อนจะส่งไปยัง sparse Splash Attention kernel และกู้คืนลำดับของผลลัพธ์

รูปที่ 9: การแปลง token memory layout การจัดเรียงแบบเน้นเฟรม (F,H,W) จะเก็บ spatial tokens ของเฟรมที่ 1 ไว้ต่อเนื่องกัน การสลับเปลี่ยนเป็นลำดับเน้นเวลา (H,W,F) จะจัดกลุ่ม token จากตำแหน่งเชิงพื้นที่เดียวกันในเฟรมที่ต่อเนื่องกัน (B1 - B4) ลงในหน่วยความจำที่ติดกันเพื่อให้การท่องไปของ temporal kernel เป็นไปอย่างมีประสิทธิภาพ

เมื่อทดสอบแบบ End-to-end บนวิดีโอ 720p พบว่าการตั้งค่าแบบ Aggressive ช่วยลดเวลา Denoising ลงได้ 22% และเมื่อเพิ่มความละเอียดเป็น 1440p ประสิทธิภาพความเร็วจะเพิ่มขึ้นอย่างเห็นได้ชัดถึง 1.69 เท่า ช่วยประหยัดเวลาในการสร้างวิดีโอได้มากกว่า 16 นาทีต่อคลิป

ตารางที่ 3: ความหน่วง end-to-end ของการทำ 720p denoising (81 เฟรม, 40 ขั้นตอน บนชิป TPU v6e แปดตัว) และคุณภาพของผลลัพธ์ในแต่ละ SVG sparsity schedules

ตารางที่ 4: ความหน่วง end-to-end ของ denoising และการขยายคุณภาพในความละเอียดต่างๆ ภายใต้ Aggressive SVG sparsity schedule บนชิป TPU v6e แปดตัว

รูปที่ 10: การเปรียบเทียบแบบเคียงข้างกันของการสร้างวิดีโอ 720p 81 เฟรม ภายใต้ dense attention (ซ้าย) เทียบกับ Aggressive SVG sparsity schedule (ขวา เร็วขึ้น 1.28 เท่า, 24.80 dB PSNR)

สรุปได้ว่าการเพิ่มประสิทธิภาพ Sparse Attention บน TPU ไม่ได้ขึ้นอยู่กับจำนวนคู่ที่ถูก Mask ออกไปเพียงอย่างเดียว แต่ขึ้นอยู่กับความสามารถของฮาร์ดแวร์ในการข้ามงานที่ไม่จำเป็นและการจัดเรียงข้อมูลให้เข้าถึงได้รวดเร็ว ซึ่งเป็นกุญแจสำคัญในการเร่งความเร็วการสร้างวิดีโอในอนาคต