Fix Mega MoE task info slot release ordering

[zheanxu]

No description provided.

[ds-review-bot]

🔵 suggestion: 非阻塞建议：非 leader CTA 的消费者与位于 CTA 0 的调度 warp 是跨 CTA 同步，严格按 PTX 内存模型应使用 .cluster scope（fence.acq_rel.cluster 或 mbarrier.arrive.release.cluster.shared::cluster）。实际硬件上 fence.acq_rel.cta 已足以让本线程未完成的 shared 加载在 arrive 前完成，因此行为安全且开销更低；建议在注释中补充说明这一取舍（为何选 .cta 而非 .cluster），或在 ld_st.cuh 中同时提供 fence_acq_rel_cluster() 变体以备后续需要。

_{🤖 v5}

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

⚠️ 未完成评审（upstream_error：模型上游服务不可用）

v5

该提交为 Mega MoE 调度器的 task info 槽位释放路径补上了内存序保证。原实现中，消费者线程在 get_next_task() 里用普通 ld.shared 拷贝 task_infos[stage]，随后在 release_task_info() 中通过 CUTLASS 的 ClusterBarrier::arrive(cta_id) 远程 arrive 到 CTA 0 的 empty barrier；该内联汇编（third-party/cutlass/include/cutlass/arch/barrier.h:485-492）没有 "memory" clobber，且 task_info 的字段（valid_m / pool_block_idx / n_cluster_idx / block_phase）都在 arrive 之后才被使用（如 sm100_bf16_mega_moe.cuh:863-872），编译器可以把字段加载下沉到 arrive 之后。一旦调度 warp 观察到 empty 并用 st_async_cluster 覆盖该槽，消费者就会读到下一个任务的元数据，导致错误的 m_idx / n_idx / block_phase（静默错误结果或越界）。新增的 ptx::fence_acq_rel_cta()（fence.acq_rel.cta + "memory" clobber）同时在编译器层阻止重排、在硬件层保证先前的 shared 加载在 arrive 前完成，能够封住这个窗口。验证结论：变更最小且正确；两个 impl（bf16 / fp8_fp4）都通过同一个 release_task_info() 受益；sched_stage_idx ^ 1 语义未变；新封装放在 /// Memory ordering 分节下与现有 ld_acq / red_xor_rel 风格一致。唯一可商榷之处是 fence scope（见评论），属非阻塞建议。总体建议合并。

v4

The MR fixes a real producer/consumer race on the double-buffered task_infos schedule slots. In MegaMoEScheduler::release_task_info(), the epilogue threads signal task_info_empty_barriers so the scheduler can reuse the slot. Previously the plain (non-atomic) reads of task_infos[sched_stage_idx] performed in get_next_task() were not explicitly ordered before the barrier arrival, so the compiler/hardware could let the scheduler overwrite the slot while metadata reads were still in flight. Adding ptx::fence_acq_rel_cta() immediately before the arrive makes the release ordering explicit, and the "memory" clobber also blocks compiler reordering. The fence scope matches the access pattern: the slot lives in the local CTA's shared memory and the barrier is the CTA/cluster mbarrier used to gate reuse. The change is minimal, well-placed in the shared scheduler header, and therefore applies to both sm100_bf16_mega_moe and sm100_fp8_fp4_mega_moe. The new helper is added next to the other memory-ordering/PTX wrappers in ld_st.cuh and is correctly annotated CUTLASS_DEVICE. No correctness issues found.

Files reviewed: 2
Issues found: 🔵 1 suggestion
Inline comments posted: 1

⚠️ Parse warning: [v6] upstream_error：模型上游服务不可用