fused_moe: add VLLM_TRITON_USE_TD tensor-descriptor path (#42436)

Signed-off-by: Artur Fierka <[email protected]>
Signed-off-by: Lena Onyshchenko <[email protected]>
Co-authored-by: Claude <[email protected]>
Co-authored-by: Lena Onyshchenko <[email protected]>

tests/kernels/moe/test_moe.py

@@ -36,6 +36,9 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
     batched_fused_marlin_moe,
     fused_marlin_moe,
 )
+from vllm.model_executor.layers.fused_moe.utils import (
+    moe_use_td_hw_supported,
+)
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     marlin_permute_bias,
 )
@@ -53,9 +56,12 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
 from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights
 from vllm.platforms import current_platform
 from vllm.scalar_type import ScalarType, scalar_types
+from vllm.triton_utils import tl
 from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
+
 
 def iterative_moe(
     hidden_states: torch.Tensor,
@@ -289,6 +295,7 @@ def run_moe_test(
 @pytest.mark.parametrize("ep_size", EP_SIZE)
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("padding", [True, False])
+@pytest.mark.parametrize("use_td", [False, True])
 def test_fused_moe(
     m: int,
     n: int,
@@ -298,9 +305,19 @@ def test_fused_moe(
     ep_size: int,
     dtype: torch.dtype,
     padding: bool,
+    use_td: bool,
     monkeypatch,
     workspace_init,
 ):
+    if use_td and not hasattr(tl, "make_tensor_descriptor"):
+        pytest.skip("Triton < 3.6 lacks tl.make_tensor_descriptor")
+    if use_td and not moe_use_td_hw_supported():
+        pytest.skip(
+            "tensor_descriptor.gather requires XPU or NVIDIA Blackwell "
+            "(sm100+); lowers to tile::gather4 (tcgen05/TMEM), which ptxas "
+            "rejects on Hopper (sm90) and earlier"
+        )
+    monkeypatch.setenv("VLLM_TRITON_USE_TD", "1" if use_td else "0")
     set_random_seed(7)
 
     #
@@ -311,17 +328,17 @@ def test_fused_moe(
     # Setup test data
     #
 
-    a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
-    w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
-    w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
+    a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10
+    w1 = torch.randn((e, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10
+    w2 = torch.randn((e, k, n), device=DEVICE_TYPE, dtype=dtype) / 10
 
-    score = torch.randn((m, e), device="cuda", dtype=dtype)
+    score = torch.randn((m, e), device=DEVICE_TYPE, dtype=dtype)
 
     if ep_size > 1:
         local_e = e // ep_size
-        e_ids = torch.randint(0, e, (local_e,), device="cuda", dtype=torch.int32)
-        e_map = torch.full((e,), -1, device="cuda", dtype=torch.int32)
-        e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)
+        e_ids = torch.randint(0, e, (local_e,), device=DEVICE_TYPE, dtype=torch.int32)
+        e_map = torch.full((e,), -1, device=DEVICE_TYPE, dtype=torch.int32)
+        e_map[e_ids] = torch.arange(local_e, device=DEVICE_TYPE, dtype=torch.int32)
         w1 = w1[e_ids]
         w2 = w2[e_ids]
     else:

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -28,9 +28,12 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
 from vllm.model_executor.layers.fused_moe.utils import (
     enable_swap_ab,
     moe_kernel_quantize_input,
+    resolve_moe_use_td,
+    warn_if_moe_use_td_ineffective,
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.triton_utils.allocation import set_triton_allocator
 from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.platform_utils import get_device_name_as_file_name
 from vllm.utils.torch_utils import direct_register_custom_op
@@ -347,6 +350,8 @@ def fused_moe_kernel(
     per_channel_quant: tl.constexpr,
     HAS_BIAS: tl.constexpr,
     SWAP_AB: tl.constexpr,
+    # Tensor-descriptor path for the A gather and B load in the K-loop.
+    USE_TD: tl.constexpr = False,
 ):
     """
     Implements the fused computation for a Mixture of Experts (MOE) using
@@ -436,7 +441,25 @@ def fused_moe_kernel(
 
     offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
     offs_k = tl.arange(0, BLOCK_SIZE_K)
-    if SWAP_AB:
+    # TD gather and the SWAP_AB accumulator layout are mutually exclusive.
+    tl.static_assert(not (USE_TD and SWAP_AB))
+    if USE_TD:
+        # ``tt.descriptor_gather`` requires block_shape[0] == 1 and i32 idx.
+        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)
+    elif SWAP_AB:
         a_ptrs = a_ptr + (
             offs_k[:, None] * stride_ak + offs_token[None, :] // top_k * stride_am
         )
@@ -454,7 +477,6 @@ def fused_moe_kernel(
             + off_experts * stride_be
             + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
         )
-
     if use_int8_w8a16:
         b_scale_ptrs = (
             b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn
@@ -498,18 +520,21 @@ def fused_moe_kernel(
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         # Load the next block of A and B, generate a mask by checking the
         # K dimension.
-        if SWAP_AB:
+        if USE_TD:
+            a = a_desc.gather(gather_idx, k * BLOCK_SIZE_K)
+            b = b_desc.load([pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]).T
+        elif SWAP_AB:
             a_mask = (offs_k[:, None] < K - k * BLOCK_SIZE_K) & token_mask[None, :]
             b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K
+            a = tl.load(a_ptrs, mask=a_mask, other=0.0)
+            b = tl.load(b_ptrs, mask=b_mask, other=0.0)
         else:
-            a_mask = token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K)
-            b_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K
-        a = tl.load(
-            a_ptrs,
-            mask=a_mask,
-            other=0.0,
-        )
-        b = tl.load(b_ptrs, mask=b_mask, other=0.0)
+            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)
         # We accumulate along the K dimension.
         if use_int8_w8a16:
             accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)
@@ -536,9 +561,10 @@ def fused_moe_kernel(
                     accumulator += tl.dot(a, b)
         else:
             accumulator += tl.dot(a, b)
-        # Advance the ptrs to the next K block.
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
+        if not USE_TD:
+            # Advance the ptrs to the next K block.
+            a_ptrs += BLOCK_SIZE_K * stride_ak
+            b_ptrs += BLOCK_SIZE_K * stride_bk
 
     if SWAP_AB:
         accumulator = tl.trans(accumulator, (1, 0))
@@ -765,6 +791,19 @@ def invoke_fused_moe_triton_kernel(
     else:
         SWAP_AB = False
 
+    # Quantized weights always carry a B_scale (see the asserts below); key off
+    # that rather than enumerating quant flags, which misses w8a16-fp8/nvfp4/etc.
+    is_quantized = B_scale is not None
+    warn_if_moe_use_td_ineffective("TRITON", is_quantized=is_quantized)
+
+    # TD path is unvalidated under quantization; fall back to the pointer path.
+    use_td = resolve_moe_use_td() and not is_quantized
+    if use_td:
+        # The TD path builds a tensor descriptor inside the kernel, which
+        # requires a PyTorch-backed scratch allocator to be registered
+        # (Triton raises "no allocator was set" otherwise on CUDA).
+        set_triton_allocator(A.device)
+
     if use_fp8_w8a8 or use_int8_w8a8:
         assert B_scale is not None
         assert block_shape is None or triton.cdiv(
@@ -805,6 +844,19 @@ def invoke_fused_moe_triton_kernel(
     BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K")
     if block_shape is not None:
         BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1]))
+    if use_td and A.size(1) % BLOCK_SIZE_K != 0:
+        # TD gather/load feeding tl.dot with a non-block-aligned K
+        # miscompiles (~74% of output elements wrong) on real HW;
+        # this is a compiler-codegen issue, not a Python-maskable
+        # boundary gap. Fall back to the pointer-arith path.
+        logger.warning_once(
+            "Disabling VLLM_TRITON_USE_TD for this MoE launch: K=%d is not "
+            "a multiple of BLOCK_SIZE_K=%d, which triggers a known "
+            "Triton tensor-descriptor + tl.dot miscompilation.",
+            A.size(1),
+            BLOCK_SIZE_K,
+        )
+        use_td = False
     fused_moe_kernel[grid](
         A,
         B,
@@ -847,6 +899,7 @@ def invoke_fused_moe_triton_kernel(
         HAS_BIAS=HAS_BIAS,
         BLOCK_SIZE_K=BLOCK_SIZE_K,
         SWAP_AB=SWAP_AB,
+        USE_TD=use_td,
         **config,
     )
 

vllm/model_executor/layers/fused_moe/oracle/unquantized.py

@@ -362,6 +362,13 @@ def make_unquantized_moe_kernel(
     experts_cls: type[mk.FusedMoEExperts],
     routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
 ) -> mk.FusedMoEKernel:
+    from vllm.model_executor.layers.fused_moe.utils import (
+        warn_if_moe_use_td_ineffective,
+    )
+
+    # Warn against the selected backend, not each probed candidate.
+    warn_if_moe_use_td_ineffective(backend.value, is_quantized=False)
+
     # Create Prepare/Finalize
     is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic)
     prepare_finalize = maybe_make_prepare_finalize(

vllm/model_executor/layers/fused_moe/utils.py

@@ -7,7 +7,9 @@ from typing import TYPE_CHECKING
 import torch
 import torch.nn.functional as F
 
+import vllm.envs as envs
 from vllm import _custom_ops as ops
+from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     per_token_group_quant_fp8,
 )
@@ -39,6 +41,8 @@ from vllm.utils.math_utils import cdiv
 if TYPE_CHECKING:
     from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
 
+logger = init_logger(__name__)
+
 
 @triton.jit
 def _count_expert_num_tokens(
@@ -585,3 +589,80 @@ def enable_swap_ab(BLOCK_SIZE_M: int, BLOCK_SIZE_N: int) -> bool:
         and BLOCK_SIZE_M < 64
         and BLOCK_SIZE_N >= 64
     )
+
+
+def moe_use_td_hw_supported() -> bool:
+    """Whether the current device can run the TD (gather) path of
+    ``fused_moe_kernel`` (ignores the ``VLLM_TRITON_USE_TD`` override).
+
+    The A-load uses ``tensor_descriptor.gather``, which lowers to the PTX
+    ``tile::gather4`` instruction. That instruction is part of the
+    ``tcgen05``/Tensor Memory (TMEM) family introduced with Blackwell and has
+    no Hopper (sm90) equivalent -- ptxas rejects it there ("Feature
+    '.tile::gather4 ...' requires .target sm_100 or higher"). Unlike
+    ``scatter4``, ``gather4`` is supported across the whole sm100+ range
+    including consumer Blackwell (sm120/sm121): see triton-lang/triton#8498,
+    which enables ``gather4`` on sm120/sm121 while leaving ``scatter4``
+    unsupported there. So this gates on a blanket ``has_device_capability(100)``
+    rather than the sm100 *family* check used for the scatter store path.
+    """
+    if current_platform.is_xpu():
+        return True
+    if current_platform.is_cuda():
+        return current_platform.has_device_capability(100)
+    return False
+
+
+def resolve_moe_use_td() -> bool:
+    """Tri-state resolver for ``VLLM_TRITON_USE_TD``.
+
+    Unset auto-selects the TD path on XPU only, mirroring the attention
+    dispatcher in ``triton_attn.py``. ``1``/``0`` force it on/off regardless
+    of hardware; forcing ``1`` where it cannot compile (see
+    ``moe_use_td_hw_supported``) fails at ptxas. Blackwell CUDA (sm100+) can
+    compile it but is opt-in only, pending validation.
+    """
+    override = envs.VLLM_TRITON_USE_TD
+    if override is None:
+        return current_platform.is_xpu()
+    return override
+
+
+_warned_moe_use_td_ineffective = False
+
+
+def warn_if_moe_use_td_ineffective(
+    active_backend: str, is_quantized: bool = False
+) -> None:
+    """One-shot warning when ``VLLM_TRITON_USE_TD`` is set but ignored.
+
+    Fires when the user set the env explicitly and either (a) the active
+    MoE backend is not the fused Triton kernel, or (b) the model is
+    quantized (the TD path falls back to the pointer path under any
+    quantization).
+    """
+    global _warned_moe_use_td_ineffective
+    if _warned_moe_use_td_ineffective:
+        return
+    if envs.VLLM_TRITON_USE_TD is None:
+        return
+    is_triton = active_backend.upper() == "TRITON"
+    if is_triton and not is_quantized:
+        return
+    if not is_triton:
+        reason = (
+            f"the active MoE backend is {active_backend!r}; pass "
+            "`--moe-backend triton` to enable the tensor-descriptor path"
+        )
+    else:
+        reason = (
+            "the model uses quantized MoE weights; the TD path is "
+            "currently restricted to non-quantized weights and falls "
+            "back to the pointer path"
+        )
+    logger.warning(
+        "VLLM_TRITON_USE_TD is set to %s but %s.",
+        envs.VLLM_TRITON_USE_TD,
+        reason,
+    )
+    _warned_moe_use_td_ineffective = True