执行摘要

• 一句话：fused_moe 新增 TD 路径，XPU 默认启用
• 推荐动作：值得精读。重点看三处设计：resolve_moe_use_td 的三态解析与硬件门控分离、B_scale 驱动的量化检测、以及 K 非对齐时对编译器 bug 的回退处理。这些模式对后续 Triton kernel 的 TD 化很有参考价值。

功能与动机

PR body 说明：为 Triton MoE 内核做性能优化，参照 VLLM_TRITON_ATTN_USE_TD（PR #40327）与 TD 采纳策略 RFC #42545 的单开关设计，统一 VLLM_TRITON_USE_TD 作为跨 Triton kernel 的 gate。XPU 的 Tensor Memory/descriptor 路径有明显收益，Blackwell 可编译但默认 opt-in；同时修复 E2E 验证中发现的 Blackwell 上 K 非对齐导致的编译器误编译正确性问题。

实现拆解

1. 在 vllm/model_executor/layers/fused_moe/utils.py 中新增 3 个函数：moe_use_td_hw_supported() 判断硬件能否编译 TD 路径（XPU 恒真，CUDA 要求 has_device_capability(100)，因为 gather 降级为 PTX tile::gather4，属于 Blackwell 的 tcgen05/TMEM 指令族）；resolve_moe_use_td() 三态解析 VLLM_TRITON_USE_TD，未设置时仅 XPU 自动开启；warn_if_moe_use_td_ineffective() 在用户显式设置但路径未生效时做一次性告警（非 Triton 后端或量化权重）。
2. 在 fused_moe.py 的 fused_moe_kernel 中加入 USE_TD: tl.constexpr = False 参数。启用时用 tl.make_tensor_descriptor 为 A 构造 gather 描述符（block_shape[0] == 1、i32 索引）、为 B 构造 load 描述符，K 循环内以 a_desc.gather(...) 与 b_desc.load(...).T 替代指针运算；tl.static_assert 禁止 USE_TD 与 SWAP_AB 同时开启。
3. 在 invoke_fused_moe_triton_kernel 启动点：以 B_scale is not None 判定量化（比枚举 quant flag 更全面），量化时强制 use_td = False；use_td 时通过 set_triton_allocator(A.device) 注册 Triton scratch allocator；当 A.size(1) % BLOCK_SIZE_K != 0 时回退指针路径并 logger.warning_once（规避 Triton 对 TD + tl.dot 在非对齐 K 下的误编译，见 triton-lang/triton#10927）。
4. 在 oracle/unquantized.py 的 make_unquantized_moe_kernel 中调用 warn_if_moe_use_td_ineffective(backend.value, is_quantized=False)，针对最终选定的后端而非每个探测候选做告警。
5. 测试配套：tests/kernels/moe/test_moe.py 参数化 use_td，Triton < 3.6 或硬件不支持时 skip；将测试数据创建从 device="cuda" 泛化为 DEVICE_TYPE 以支持 XPU。量化与非量化场景均有现有参数化覆盖。

关键文件：

• vllm/model_executor/layers/fused_moe/utils.py（模块 融合 MoE；类别 source；类型 data-contract；符号 moe_use_td_hw_supported, resolve_moe_use_td, warn_if_moe_use_td_ineffective）: 新增硬件能力检测、三态解析和一次性告警，构成 TD 路径的开关与风险控制核心。
• vllm/model_executor/layers/fused_moe/fused_moe.py（模块 融合 MoE；类别 source；类型 data-contract；符号 fused_moe_kernel, invoke_fused_moe_triton_kernel）: 内核与启动点改动：USE_TD 分支实现 TD gather/load，以及量化回退、K 对齐回退和 scratch allocator 注册。
• vllm/model_executor/layers/fused_moe/oracle/unquantized.py（模块 融合 MoE；类别 source；类型 data-contract；符号 make_unquantized_moe_kernel）: 在非量化后端选择入口加入一次性告警，确保用户显式设置但后端不对时得到提示。
• tests/kernels/moe/test_moe.py（模块 MoE 测试；类别 test；类型 test-coverage；符号 test_fused_moe）: 测试覆盖 TD 路径，增加 Triton 版本与硬件 skip 保护，并将设备泛化到 XPU。

关键符号：moe_use_td_hw_supported, resolve_moe_use_td, warn_if_moe_use_td_ineffective, fused_moe_kernel, invoke_fused_moe_triton_kernel

关键源码片段

vllm/model_executor/layers/fused_moe/fused_moe.py

内核与启动点改动：USE_TD 分支实现 TD gather/load，以及量化回退、K 对齐回退和 scratch allocator 注册。

# vllm/model_executor/layers/fused_moe/fused_moe.py（head 版本整理）
# 内核 K 循环中的加载分支：TD 路径与指针路径
if USE_TD:
    # tt.descriptor_gather 要求 block_shape[0] == 1 且索引为 i32
    m_td = num_valid_tokens // top_k
    a_desc = tl.make_tensor_descriptor(
        base=a_ptr,
        shape=(m_td, K),
        strides=(stride_am, stride_ak),
        block_shape=(1, BLOCK_SIZE_K),
    )
    b_desc = tl.make_tensor_descriptor(
        base=b_ptr + off_experts * stride_be,
        shape=(N, K),
        strides=(stride_bn, stride_bk),
        block_shape=(BLOCK_SIZE_N, BLOCK_SIZE_K),
    )
    gather_idx = (offs_token // top_k).to(tl.int32)
​
# TD 与 SWAP_AB 的累加器布局不兼容，编译期强制互斥
tl.static_assert(not (USE_TD and SWAP_AB))
​
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
    if USE_TD:
        # 描述符按形状自动填零，无需显式 mask
        a = a_desc.gather(gather_idx, k * BLOCK_SIZE_K)
        b = b_desc.load([pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]).T
    else:
        # 指针路径：显式 mask 处理边界
        a = tl.load(
            a_ptrs,
            mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
            other=0.0,
        )
        b = tl.load(
            b_ptrs,
            mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
            other=0.0,
        )
    accumulator += tl.dot(a, b)
    if not USE_TD:
        # TD 路径无需手工推进指针，描述符按 k 偏移加载
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk

评论区精华

Review 中核心交锋集中在三点：

• mayuyuace 质疑硬编码 is_quantized = use_fp8_w8a8 or use_int8_w8a8 or use_int8_w8a16 or use_int4_w4a16 不安全；oonyshch 随后在提交 c7ca8ff 中改为基于 B_scale is not None 判断，避免遗漏 w8a16-fp8/nvfp4 等量化形式。
• quinnlp 询问是否也为输出 store 增加 TD 路径，该问题留待后续（本 PR 仅覆盖 A gather 与 B load）。

• gemini-code-assist[bot] 曾指出早期迭代中 test_batched_moe.py 缺少 Triton 版本 skip；作者回复已在 HEAD 应用相同保护（但最终 PR 仅保留 test_moe.py）。

• 量化检测方式不完整 (design): oonyshch 在提交 c7ca8ff 中改为基于 B_scale is not None 检测，量化权重必然携带 B_scale，覆盖更全面。

• K 非对齐时的编译器误编译 (correctness): 在启动点增加 A.size(1) % BLOCK_SIZE_K != 0 回退到指针路径，并 warning_once 提示。
• TD 路径是否扩展输出 store (design): 本 PR 仅覆盖 A gather 与 B load，输出 store 留作后续扩展（可能复用同一 VLLM_TRITON_USE_TD）。
• Triton 版本兼容性测试 (testing): afierka-intel 回复已在 HEAD 应用 hasattr(tl, 'make_tensor_descriptor') skip guard（最终 PR 的 test_moe.py 同样包含）。

风险与影响

• 风险：
  1. 量化回退依赖 B_scale is not None 判定，若未来出现无 B_scale 的量化格式可能误入 TD 路径，需关注后续量化扩展。
  2. Blackwell 上 K 非对齐已通过回退规避，但 Triton 编译器对 TD + tl.dot 的误编译仍可能在其他边界 shape 出现，上游修复前需保持该回退。
  3. TD 路径仅在 XPU/Blackwell 实测，Hopper/Ampere 显式开启会在 ptxas 阶段直接失败，告警只能提示无法编译。
  4. 性能无普适提升：B200 上三个数据集均略降，B70 上仅 sharegpt 有增益；默认关闭避免了负优化风险。
     - 影响：影响范围：所有使用 Triton fused MoE kernel 的用户。XPU 用户默认启用 TD 路径，行为变化最大（部分场景吞吐 +6%，部分持平）；CUDA/ROCm 用户默认不变，仅显式设置后才受影响。量化模型自动回退，保证与 main 字节级一致。对团队而言，该 PR 确立了 VLLM_TRITON_USE_TD 单开关在 MoE 内核的落地范式，后续 batched MoE TD 路径（#46340）将复用同一 flag。
     - 风险标记：XPU 默认启用行为变化, 量化回退依赖 B_scale 检测, Blackwell K 非对齐编译器 bug 回退, 仅非量化 bf16 验证, 性能无普适提升

关联脉络

• PR #50874 [Bugfix][R3] Size monolithic routing replay buffer for DP: 同属 fused_moe 内核演进：修复 DP/EP 下 MoE 路由回放缓冲容量与分片捕获，与本 PR 的 MoE 内核优化同模块相连。
• PR #52114 [Model] [Quantization] Add Ling hybrid MXFP4 routed experts support: 增加 MXFP4 路由专家量化支持，而本 PR 的 TD 路径在量化权重下会回退，两者在量化 MoE 路线上互补。