以内核为中心的路径：在 Trainium 上实现实时视频生成

团队用 NKI 直接编写运行在 NeuronCore 上的计算内核，替换通用编译器难以优化的瓶颈算子：3D-RoPE 内核从五秒降至 1.8 毫秒，缓存拷贝从每层 23 毫秒降至 1.9 毫秒，注意力转置通过融合进注意力内核被完全消除。 相对依赖通用编译器的路径，这里把实时视频生成中反复出现的动态形状、特殊访存模式与重缓存管理三类问题下沉到硬件级内核解决，并给出可量化的算子级收益。 文中给出具体算子级延迟数字（五秒到 1.8 毫秒、23 毫秒到 1.9 毫秒）与内存结果（流水线占用 11 GB 高带宽内存，而标准 eager 模式路径内存不足），属于工程实测报告。

针对自注意力（23,400 个查询 token 关注 32,760 个上下文 token，约占 70% 计算时间），团队采用序列并行与张量并行的混合分片：把注意力头分到 4 个核、序列分到 2 个核，兼顾 12 个注意力头与每芯片 8 个核的整除问题以及 3D 视频 token 的帧粒度约束。 单独使用张量并行需要填充而浪费计算，单独使用序列并行可能破坏 3D 结构；混合方案同时保持数学正确与数据布局正确，并报告 VAE 解码器采用空间 W 轴分片取得 8.25 倍超线性加速。 文中给出注意力规模、计算占比、分片维度与 8.25 倍加速等具体数字，属于针对该模型结构的工程测量。

团队改写模型代码，把扩散阶段与缓存更新阶段共有的组件批处理在一起，仅在组件略有差异时再拆分处理，以缓解缓存更新阶段计算量小、单独执行时硬件利用率低的问题。 相对把两阶段作为两次独立运行的做法，这种批处理针对 Trainium 每实例 16 芯片、每芯片 8 核的拓扑，避免把计算切分到过多核上导致每次调用计算量过小。 文中以工程师对两阶段计算量与硬件利用率的描述为依据，未给出该改写单独的加速数字。

综合上述优化后，团队在 Trainium 上首次端到端运行即成功生成正确视频，并强调这些技术面向所有具有实时流式需求的自回归扩散模型，而非只优化单一模型。 相对针对单个模型的调优，团队把 Rolling Forcing 作为代理，主张其包含编码器、DiT、VAE、解码器等完整系统组件，因此所得方法可迁移到更复杂的系统。 这是团队自述的首次端到端成功与可推广性判断，文中未提供与其他硬件或基线的对照数据。

这项工作面向需要在 Trainium 上部署实时流式自回归扩散视频模型的开发者与平台方，适用场景是生成必须持续领先播放时间线、帧率需高于 16 fps 的交互式视频生成；团队以 Rolling Forcing 作为代理模型验证路径，并主张其方法可迁移到包含编码器、DiT、VAE、解码器的更复杂系统。

文中未给出与其他硬件或基线的对照实验，也未说明这些优化在更大或不同架构的自回归扩散模型上的表现；算子级延迟与 8.25 倍加速的具体测量条件、以及首次端到端成功所对应的视频质量与稳定性，读者仍需结合后续公开材料进一步确认。