修复 fused RoPE 丢弃 mrope 高度宽度坐标
值得精读,尤其是 kernel 层与 rotary 层配合的写法。关注 4 个设计点: 1. 用 `MROPE: tl.constexpr` 区分 1-D 与 mrope 特化,既修复 bug 又不损失 1-D 性能(两个 benchmark 数据佐证)。 2. 行距取 `positions.stride(0)` 而不是 3,兼容 CUDA-graph decode 从宽 buffer 切出的子行。 3. `_build_axis_map` 将三种布局收敛为统一契约,并用 `_legacy_axis_map` 隔离新旧 kernel 的输入差异。 4. 测试特意使用三轴互不相同的位置(t != h != w)来捕捉“只取 row 0”类 bug,并复现 CUDA-graph 的 buffer 切片场景。
参与讨论