[TLE][NVIDIA] Add SMEM subslices and multi-writer TMA pipes

[Zhang-kg]

背景

PR837 实现的 TLE pipe lowering 当前默认一个 pipe 只能对应一个 writer task。

MegaMoE 的 warp-specialized pipeline 需要两个独立的 TMA producer 分别加载不同 operand，并共同填充同一个 pipeline stage。同时，math task 需要在不改变 rank 和 shared-memory layout 的前提下，对已经分配的 SMEM tile 创建静态子视图。

因此，本 PR 在 PR837 基础上补充以下两项能力：

1. buffered_tensor 的 SMEM subslice；
2. 同一 pipe 上多个独立的纯 TMA writer。

修改内容

1. 增加 buffered_tensor.subslice

新增接口：

buffer.subslice(start, length, dim)

该接口：

• 仅支持 shared memory；
• start、length 和 dim 必须为编译期常量；
• 保持原 tensor 的 rank；
• 保持原有 shared-memory layout；
• 保留父 buffer 的 alloc_shape；
• 对越界、非法维度和非静态参数进行前端检查。

同时将 buffered_tensor.slot 和 buffered_tensor.subslice 注册到 NVIDIA backend 的 TLE capability 列表中。

2. 支持多个 TMA writer 共享一个 pipe

通过以下环境变量显式启用：

TLE_MULTI_TMA_WRITERS=1

启用后，多个独立 writer task 可以分别向同一个 pipe stage 的不同字段发射 TMA copy。

Lowering 会：

• 收集 pipe 上所有 writer task；
• 根据 writer 数量设置 full barrier 的 arrival count；
• 允许不同 writer 分别 commit pipe 的部分字段；
• 在 NVIDIA token lowering 中保留显式设置的 TMA full_count；
• 将环境变量加入 Triton compilation cache key。

该功能目前仅允许满足以下条件的配置：

• 所有 writer 都必须使用纯 TMA transport；
• 每个 writer task 都必须提供对应的 TMA commit；
• 不同 writer 写入的 pipe 字段必须互不重叠；
• 所有 writer 写入字段的并集必须完整覆盖 pipe；
• 暂不支持 mixed TMA/local-store、mixed cp.async 或 writer close。

不满足约束的配置会在编译期给出明确诊断。

兼容性

多 writer 支持默认关闭。

未设置 TLE_MULTI_TMA_WRITERS=1 时，保持 PR837 原有的单 writer pipe 行为，不改变现有 kernel 的 lowering 结果。

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}