🚀 Fast-FlashTalk: Efficient FlashTalk 14B Inference on RTX 4090 (2× Speedup, Reduced VRAM)

[di-osc]

Hi, thanks for the great work on FlashTalk!

I’d like to share a project I’ve been working on: Fast-FlashTalk, an optimized inference implementation designed to improve the efficiency of FlashTalk, especially on consumer GPUs like RTX 4090.

The focus is to make FlashTalk 14B (DiT-based) more practical to run on a single GPU, with better speed and lower memory usage.

---

🎯 Motivation

Running FlashTalk efficiently on a single GPU can be challenging due to:

• High VRAM requirements for large DiT models
• Compute overhead in attention and positional encoding
• Redundant computation in encoder components

This project aims to improve practical usability under these constraints.

---

⚡ Key Results (RTX 4090)

• 🚀 ~2× inference speedup compared to baseline
• 💾 Reduced VRAM usage, enabling 14B DiT on a single GPU

---

🔧 Core Optimizations

1. DiT Dynamic Parameter Offloading

• Dynamically schedules parameters between CPU and GPU
• Controlled via num_persistent_param_in_dit
• Reduces peak VRAM usage while keeping performance reasonable

---

2. GemLite A8W8 Quantization

• INT8 dynamic quantization for linear layers
• Reduces both memory footprint and compute cost
• Keeps key modules (e.g., time_embedding, head) in higher precision

---

3. RoPE Optimization

• Uses flash_attention rotary embedding kernel
• Replaces element-wise complex operations
• Improves efficiency via CUDA kernel fusion

---

4. SageAttention Integration

• Integrates SageAttention (sageattn / sageattn_varlen)
• Falls back to flash_attn for short sequences
• Improves overall attention efficiency

---

🔗 Project Link

👉 https://github.com/di-osc/fast-flashtalk

---

Thanks again for the great project!

[jokerz0624]

Thanks for your PR! I'll check it as soon as possible.