@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6f00a1ae3b

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

Signed-off-by: Artur Fierka <artur.fierka@intel.com> Signed-off-by: Lena Onyshchenko <162571002+oonyshch@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Lena Onyshchenko <162571002+oonyshch@users.noreply.github.com>

Artur Fierka committed Jul 29, 2026 at 07:15 UTC 6f00a1ae3bd4b86168667bce673998218f461c0f

4 files changed +179 -21

tests/kernels/moe/test_moe.py

+24 -7

@@ -36,6 +36,9 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (

36 batched_fused_marlin_moe,

37 fused_marlin_moe,

38 )

39 +from vllm.model_executor.layers.fused_moe.utils import (

40 + moe_use_td_hw_supported,

41 +)

42 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

43 marlin_permute_bias,

44 )

@@ -53,9 +56,12 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (

56 from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights

57 from vllm.platforms import current_platform

58 from vllm.scalar_type import ScalarType, scalar_types

59 +from vllm.triton_utils import tl

60 from vllm.utils.math_utils import next_power_of_2

61 from vllm.utils.torch_utils import set_random_seed

63 +DEVICE_TYPE = current_platform.device_type

64 +

66 def iterative_moe(

67 hidden_states: torch.Tensor,

@@ -289,6 +295,7 @@ def run_moe_test(

295 @pytest.mark.parametrize("ep_size", EP_SIZE)

296 @pytest.mark.parametrize("dtype", [torch.bfloat16])

297 @pytest.mark.parametrize("padding", [True, False])

298 +@pytest.mark.parametrize("use_td", [False, True])

299 def test_fused_moe(

300 m: int,

301 n: int,

@@ -298,9 +305,19 @@ def test_fused_moe(

305 ep_size: int,

306 dtype: torch.dtype,

307 padding: bool,

308 + use_td: bool,

309 monkeypatch,

310 workspace_init,

311 ):

312 + if use_td and not hasattr(tl, "make_tensor_descriptor"):

313 + pytest.skip("Triton < 3.6 lacks tl.make_tensor_descriptor")

314 + if use_td and not moe_use_td_hw_supported():

315 + pytest.skip(

316 + "tensor_descriptor.gather requires XPU or NVIDIA Blackwell "

317 + "(sm100+); lowers to tile::gather4 (tcgen05/TMEM), which ptxas "

318 + "rejects on Hopper (sm90) and earlier"

319 + )

320 + monkeypatch.setenv("VLLM_TRITON_USE_TD", "1" if use_td else "0")

321 set_random_seed(7)

322

323 #

@@ -311,17 +328,17 @@ def test_fused_moe(

328 # Setup test data

329 #

330

314 - a = torch.randn((m, k), device="cuda", dtype=dtype) / 10

315 - w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10

316 - w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10

331 + a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10

332 + w1 = torch.randn((e, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10

333 + w2 = torch.randn((e, k, n), device=DEVICE_TYPE, dtype=dtype) / 10

334

318 - score = torch.randn((m, e), device="cuda", dtype=dtype)

335 + score = torch.randn((m, e), device=DEVICE_TYPE, dtype=dtype)

336

337 if ep_size > 1:

338 local_e = e // ep_size

322 - e_ids = torch.randint(0, e, (local_e,), device="cuda", dtype=torch.int32)

323 - e_map = torch.full((e,), -1, device="cuda", dtype=torch.int32)

324 - e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)

339 + e_ids = torch.randint(0, e, (local_e,), device=DEVICE_TYPE, dtype=torch.int32)

340 + e_map = torch.full((e,), -1, device=DEVICE_TYPE, dtype=torch.int32)

341 + e_map[e_ids] = torch.arange(local_e, device=DEVICE_TYPE, dtype=torch.int32)

342 w1 = w1[e_ids]

343 w2 = w2[e_ids]

344 else:

vllm/model_executor/layers/fused_moe/fused_moe.py

+67 -14

@@ -28,9 +28,12 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (

28 from vllm.model_executor.layers.fused_moe.utils import (

29 enable_swap_ab,

30 moe_kernel_quantize_input,

31 + resolve_moe_use_td,

32 + warn_if_moe_use_td_ineffective,

33 )

34 from vllm.platforms import current_platform

35 from vllm.triton_utils import tl, triton

36 +from vllm.triton_utils.allocation import set_triton_allocator

37 from vllm.utils.math_utils import next_power_of_2

38 from vllm.utils.platform_utils import get_device_name_as_file_name

39 from vllm.utils.torch_utils import direct_register_custom_op

@@ -347,6 +350,8 @@ def fused_moe_kernel(

350 per_channel_quant: tl.constexpr,

351 HAS_BIAS: tl.constexpr,

352 SWAP_AB: tl.constexpr,

353 + # Tensor-descriptor path for the A gather and B load in the K-loop.

354 + USE_TD: tl.constexpr = False,

355 ):

356 """

357 Implements the fused computation for a Mixture of Experts (MOE) using

@@ -436,7 +441,25 @@ def fused_moe_kernel(

441

442 offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N

443 offs_k = tl.arange(0, BLOCK_SIZE_K)

439 - if SWAP_AB:

444 + # TD gather and the SWAP_AB accumulator layout are mutually exclusive.

445 + tl.static_assert(not (USE_TD and SWAP_AB))

446 + if USE_TD:

447 + # ``tt.descriptor_gather`` requires block_shape[0] == 1 and i32 idx.

448 + m_td = num_valid_tokens // top_k

449 + a_desc = tl.make_tensor_descriptor(

450 + base=a_ptr,

451 + shape=(m_td, K),

452 + strides=(stride_am, stride_ak),

453 + block_shape=(1, BLOCK_SIZE_K),

454 + )

455 + b_desc = tl.make_tensor_descriptor(

456 + base=b_ptr + off_experts * stride_be,

457 + shape=(N, K),

458 + strides=(stride_bn, stride_bk),

459 + block_shape=(BLOCK_SIZE_N, BLOCK_SIZE_K),

460 + )

461 + gather_idx = (offs_token // top_k).to(tl.int32)

462 + elif SWAP_AB:

463 a_ptrs = a_ptr + (

464 offs_k[:, None] * stride_ak + offs_token[None, :] // top_k * stride_am

465 )

@@ -454,7 +477,6 @@ def fused_moe_kernel(

477 + off_experts * stride_be

478 + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

479 )

457 -

480 if use_int8_w8a16:

481 b_scale_ptrs = (

482 b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn

@@ -498,18 +520,21 @@ def fused_moe_kernel(

520 for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):

521 # Load the next block of A and B, generate a mask by checking the

522 # K dimension.

501 - if SWAP_AB:

523 + if USE_TD:

524 + a = a_desc.gather(gather_idx, k * BLOCK_SIZE_K)

525 + b = b_desc.load([pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]).T

526 + elif SWAP_AB:

527 a_mask = (offs_k[:, None] < K - k * BLOCK_SIZE_K) & token_mask[None, :]

528 b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K

529 + a = tl.load(a_ptrs, mask=a_mask, other=0.0)

530 + b = tl.load(b_ptrs, mask=b_mask, other=0.0)

531 else:

505 - a_mask = token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K)

506 - b_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K

507 - a = tl.load(

508 - a_ptrs,

509 - mask=a_mask,

510 - other=0.0,

511 - )

512 - b = tl.load(b_ptrs, mask=b_mask, other=0.0)

532 + a = tl.load(

533 + a_ptrs,

534 + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),

535 + other=0.0,

536 + )

537 + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)

538 # We accumulate along the K dimension.

539 if use_int8_w8a16:

540 accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)

@@ -536,9 +561,10 @@ def fused_moe_kernel(

561 accumulator += tl.dot(a, b)

562 else:

563 accumulator += tl.dot(a, b)

539 - # Advance the ptrs to the next K block.

540 - a_ptrs += BLOCK_SIZE_K * stride_ak

541 - b_ptrs += BLOCK_SIZE_K * stride_bk

564 + if not USE_TD:

565 + # Advance the ptrs to the next K block.

566 + a_ptrs += BLOCK_SIZE_K * stride_ak

567 + b_ptrs += BLOCK_SIZE_K * stride_bk

568

569 if SWAP_AB:

570 accumulator = tl.trans(accumulator, (1, 0))

@@ -765,6 +791,19 @@ def invoke_fused_moe_triton_kernel(

791 else:

792 SWAP_AB = False

793

794 + # Quantized weights always carry a B_scale (see the asserts below); key off

795 + # that rather than enumerating quant flags, which misses w8a16-fp8/nvfp4/etc.

796 + is_quantized = B_scale is not None

797 + warn_if_moe_use_td_ineffective("TRITON", is_quantized=is_quantized)

798 +

799 + # TD path is unvalidated under quantization; fall back to the pointer path.

800 + use_td = resolve_moe_use_td() and not is_quantized

801 + if use_td:

802 + # The TD path builds a tensor descriptor inside the kernel, which

803 + # requires a PyTorch-backed scratch allocator to be registered

804 + # (Triton raises "no allocator was set" otherwise on CUDA).

805 + set_triton_allocator(A.device)

806 +

807 if use_fp8_w8a8 or use_int8_w8a8:

808 assert B_scale is not None

809 assert block_shape is None or triton.cdiv(

@@ -805,6 +844,19 @@ def invoke_fused_moe_triton_kernel(

844 BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K")

845 if block_shape is not None:

846 BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1]))

847 + if use_td and A.size(1) % BLOCK_SIZE_K != 0:

848 + # TD gather/load feeding tl.dot with a non-block-aligned K

849 + # miscompiles (~74% of output elements wrong) on real HW;

850 + # this is a compiler-codegen issue, not a Python-maskable

851 + # boundary gap. Fall back to the pointer-arith path.

852 + logger.warning_once(

853 + "Disabling VLLM_TRITON_USE_TD for this MoE launch: K=%d is not "

854 + "a multiple of BLOCK_SIZE_K=%d, which triggers a known "

855 + "Triton tensor-descriptor + tl.dot miscompilation.",

856 + A.size(1),

857 + BLOCK_SIZE_K,

858 + )

859 + use_td = False

860 fused_moe_kernel[grid](

861 A,

862 B,

@@ -847,6 +899,7 @@ def invoke_fused_moe_triton_kernel(

899 HAS_BIAS=HAS_BIAS,

900 BLOCK_SIZE_K=BLOCK_SIZE_K,

901 SWAP_AB=SWAP_AB,

902 + USE_TD=use_td,

903 **config,

904 )

905

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

@@ -362,6 +362,13 @@ def make_unquantized_moe_kernel(

362 experts_cls: type[mk.FusedMoEExperts],

363 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

364 ) -> mk.FusedMoEKernel:

365 + from vllm.model_executor.layers.fused_moe.utils import (

366 + warn_if_moe_use_td_ineffective,

367 + )

368 +

369 + # Warn against the selected backend, not each probed candidate.

370 + warn_if_moe_use_td_ineffective(backend.value, is_quantized=False)

371 +

372 # Create Prepare/Finalize

373 is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic)

374 prepare_finalize = maybe_make_prepare_finalize(

vllm/model_executor/layers/fused_moe/utils.py

+81

@@ -7,7 +7,9 @@ from typing import TYPE_CHECKING

7 import torch

8 import torch.nn.functional as F

10 +import vllm.envs as envs

11 from vllm import _custom_ops as ops

12 +from vllm.logger import init_logger

13 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

14 per_token_group_quant_fp8,

15 )

@@ -39,6 +41,8 @@ from vllm.utils.math_utils import cdiv

41 if TYPE_CHECKING:

42 from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig

44 +logger = init_logger(__name__)

45 +

47 @triton.jit

48 def _count_expert_num_tokens(

@@ -585,3 +589,80 @@ def enable_swap_ab(BLOCK_SIZE_M: int, BLOCK_SIZE_N: int) -> bool:

589 and BLOCK_SIZE_M < 64

590 and BLOCK_SIZE_N >= 64

591 )

592 +

593 +

594 +def moe_use_td_hw_supported() -> bool:

595 + """Whether the current device can run the TD (gather) path of

596 + ``fused_moe_kernel`` (ignores the ``VLLM_TRITON_USE_TD`` override).

597 +

598 + The A-load uses ``tensor_descriptor.gather``, which lowers to the PTX

599 + ``tile::gather4`` instruction. That instruction is part of the

600 + ``tcgen05``/Tensor Memory (TMEM) family introduced with Blackwell and has

601 + no Hopper (sm90) equivalent -- ptxas rejects it there ("Feature

602 + '.tile::gather4 ...' requires .target sm_100 or higher"). Unlike

603 + ``scatter4``, ``gather4`` is supported across the whole sm100+ range

604 + including consumer Blackwell (sm120/sm121): see triton-lang/triton#8498,

605 + which enables ``gather4`` on sm120/sm121 while leaving ``scatter4``

606 + unsupported there. So this gates on a blanket ``has_device_capability(100)``

607 + rather than the sm100 *family* check used for the scatter store path.

608 + """

609 + if current_platform.is_xpu():

610 + return True

611 + if current_platform.is_cuda():

612 + return current_platform.has_device_capability(100)

613 + return False

614 +

615 +

616 +def resolve_moe_use_td() -> bool:

617 + """Tri-state resolver for ``VLLM_TRITON_USE_TD``.

618 +

619 + Unset auto-selects the TD path on XPU only, mirroring the attention

620 + dispatcher in ``triton_attn.py``. ``1``/``0`` force it on/off regardless

621 + of hardware; forcing ``1`` where it cannot compile (see

622 + ``moe_use_td_hw_supported``) fails at ptxas. Blackwell CUDA (sm100+) can

623 + compile it but is opt-in only, pending validation.

624 + """

625 + override = envs.VLLM_TRITON_USE_TD

626 + if override is None:

627 + return current_platform.is_xpu()

628 + return override

629 +

630 +

631 +_warned_moe_use_td_ineffective = False

632 +

633 +

634 +def warn_if_moe_use_td_ineffective(

635 + active_backend: str, is_quantized: bool = False

636 +) -> None:

637 + """One-shot warning when ``VLLM_TRITON_USE_TD`` is set but ignored.

638 +

639 + Fires when the user set the env explicitly and either (a) the active

640 + MoE backend is not the fused Triton kernel, or (b) the model is

641 + quantized (the TD path falls back to the pointer path under any

642 + quantization).

643 + """

644 + global _warned_moe_use_td_ineffective

645 + if _warned_moe_use_td_ineffective:

646 + return

647 + if envs.VLLM_TRITON_USE_TD is None:

648 + return

649 + is_triton = active_backend.upper() == "TRITON"

650 + if is_triton and not is_quantized:

651 + return

652 + if not is_triton:

653 + reason = (

654 + f"the active MoE backend is {active_backend!r}; pass "

655 + "`--moe-backend triton` to enable the tensor-descriptor path"

656 + )

657 + else:

658 + reason = (

659 + "the model uses quantized MoE weights; the TD path is "

660 + "currently restricted to non-quantized weights and falls "

661 + "back to the pointer path"

662 + )

663 + logger.warning(

664 + "VLLM_TRITON_USE_TD is set to %s but %s.",

665 + envs.VLLM_TRITON_USE_TD,

666 + reason,

667 + )

668 + _warned_moe_use_td_ineffective = True