FlashMLA 的核心优化点之一 序列维 Stream-K 调度策略，是一种针对可变长度序列的高效并行计算方法。其核心思想是通过沿序列维度动态划分计算任务至 GPU 的流多处理器（SM），最大化硬件利用率并减少调度开销。以下是该优化的技术细节与实现逻辑：
🔧 1. 序列维任务划分原理
动态分块与 SM 分配
将输入序列（batch × seq）沿序列维度（seq）拆分为多个子块（tiles），每个子块的计算任务分配给一个 SM。具体实现中：
短序列合并：多个短序列的子块可能被分配到同一个 SM 上，避免 SM 空闲。
长序列拆分：长序列被拆分为多个子块，分布到不同 SM 并行处理。
资源对齐：通过设置 CTA（线程块）数量等于 SM 数量，确保每个 SM 至少处理一个任务块，实现硬件资源与计算负载的严格匹配。
元数据调度机制
使用 get_mla_metadata() 函数生成动态调度计划（tile_scheduler_metadata），根据序列实际长度（cache_seqlens）和注意力头数（h_kv）计算任务分块策略（num_splits），实现负载均衡。
⚡ 2. 计算碎片消除与调度优化
全局 Wave 数目为 1
传统 GEMM 计算可能因任务大小不均产生多个调度波次（Waves），导致 SM 间等待和资源闲置。Stream-K 通过 单波次全局调度（仅一个 Wave）实现：
消除波次间同步开销：避免多波次调度导致的上下文切换延迟。
计算碎片最小化：任务划分粒度与 SM 数量对齐，确保所有 SM 同时完成计算。
双模式执行引擎适配
根据序列长度动态切换计算策略：
短序列 → 计算优先模式：利用 Tensor Core 密集计算（达 580 TFLOPS）。
长序列 → 内存优先模式：优化内存访问（带宽 3000 GB/s），减少显存瓶颈。
🧩 3. 与硬件架构的协同优化
Hopper GPU 特性适配
Tensor Core 利用率：通过循环步长调整（如 i+=16）匹配 Tensor Core 的 16×16 矩阵计算粒度，提升计算密度。
异步内存拷贝：利用 TMA（Tensor Memory Accelerator）指令实现 HBM → SRAM 零拷贝传输，预取下一子块的 KV 数据，隐藏内存延迟。
共享内存与寄存器优化
Shared Memory 分块缓存：将子块 KV 数据缓存在共享内存（48 KB 用于状态缓存，16 KB 用于局部 KV），供线程块内复用。
寄存器累加中间结果：减少全局内存写入次数，提升计算效率。
💎 4. 性能优势总结
优化维度 传统方法 FlashMLA Stream-K 提升效果
调度开销 多波次调度（Waves > 1） 单波次全局调度（Wave = 1） 调度延迟降低 40%+
负载均衡 长序列对齐导致 SM 闲置 动态分块匹配 SM 数量 SM 利用率达 95%+
内存访问 连续大块内存易碎片化 分页式 KV 缓存（块大小=64） 显存碎片减少 93.3%
计算密度 固定粒度并行 序列维细粒度分块 + Tensor Core 优化 计算性能 580 TFLOPS
💡 5. 实际应用场景
长文本推理：通过分块调度支持万 token 级序列，显存占用仅传统方法的 6.7%。
高并发服务：动态任务分配允许多个短序列（如用户对话）高效共享 SM 资源，提升 QPS。
低延迟场景：单波次调度减少内核启动延迟，适用于实时翻译、搜索增强等任务。
🔍 技术对比参考
FlashMLA 的 Stream-K 与相近优化技术（如 Marlin Kernel、FlashAttention-3）的区别在于：
维度选择：沿序列维度（而非头维度或批维度）划分任务，更适配变长序列。
动态性：通过元数据实时调整分块策略，避免静态划分的负载不均问题。
📌 核心价值：Stream-K 通过 “硬件资源-任务粒度” 严格匹配 + 单波次全局调度，在保持低秩压缩（如 KV 缓存减少 93.3%）的同时，彻底释放 Hopper GPU 的并行潜力，成为大模型高吞吐推理的关键突破。
