将因果注意力掩码与扩散变换器（DiT）结合的生成范式。 核心约束：每个位置只能关注它之前的序列，杜绝“看到未来”。
这个文件是 Matrix-Game 4 推理链路中注意力计算的“统一调度层”，源自 Alibaba Wan 团队代码，核心作用是根据硬件和配置，将 Q/K/V 分发到 FlashAttention 3、FlashAttention 2 或 PyTorch SDPA 上执行。它本身不定义网络结构，而是决定“因果性如何被物理执行”。
它暴露了 flash_attention 和 attention 两个函数。flash_attention 负责变长序列的底层内核调用，而 attention 则是上层封装，供 DiT 的各个 Transformer Block 调用。整个文件的设计目标很明确：在保持因果约束的前提下，尽可能压榨推理速度，这与 Matrix-Game 4 “单卡实时生成”的定位直接挂钩。
1. causal 是显式开关，而非隐式掩码
flash_attention 函数签名里直接有 causal=False 参数，注释写着“Whether to apply causal attention mask”。这意味着因果性不是靠外部构造下三角矩阵实现的，而是直接传给 FlashAttention 内核的布尔标志。可以顺着这个 causal 参数，去追踪它在 Matrix-Game 4 的 Transformer Block 里被设置成 True 的调用链，是因果 DiT 的“开关点”。
2. window_size 暴露了“滑动窗口”的流式逻辑
参数 window_size=(-1, -1) 的注释是“If not (-1, -1), apply sliding window local attention”。这直接对应 Matrix-Game 4 的流式生成机制——每个新 chunk 只关注最近若干块组成的滑动窗口，而不是全部历史。(-1, -1) 表示全注意力，非负值则激活局部窗口。这是一个非常具体的可调参数，值得检查推理配置里它实际被设成了什么。
3. 变长序列（q_lens / k_lens）是为因果缓存服务的
函数专门处理 q_lens 和 k_lens 两个列表，用于拼接不同长度的序列。在因果自回归推理中，历史 KV 缓存长度和当前 query 长度天然不等，这个变长支持就是为流式 KV 缓存准备的。你可以看 attention 封装层如何传入这两个长度，来理解 Matrix-Game 4 是如何把“历史块”和“当前块”拼在一起送进注意力的。
4. 版本降级与硬件判断是“为实时性兜底”
代码顶部有 is_hopper_gpu() 检查，只在 H100/Hopper 上启用 FlashAttention 3，否则回退到 FlashAttention 2。同时，当请求 FA3 但不可用时，会显式发出警告并降级。这说明 Matrix-Game 4 的推理不是绑死在某一种注意力内核上，而是根据显卡能力动态选择最快路径，这是单卡 60 FPS 的工程基础之一。
文件没有自己构造任何 attn_mask 张量。因果性完全由底层内核的 causal=True 实现，滑动窗口由 window_size 实现。这意味着 Matrix-Game 4 的“因果 DiT”在注意力层面是零额外掩码开销的——因果约束被下沉到了高度优化的 CUDA 内核里，而不是在 Python 层用巨大的布尔矩阵去乘。从这里分析性能瓶颈，是一个很好的切入点。