@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c794754062

[ROCm]Remove special-case SiTU support model-specific gating (#50597)

Signed-off-by: Stacy Roberts <sroberts@amd.com> Signed-off-by: Stacy Roberts <robers23@utexas.edu> Co-authored-by: Stacy Roberts <robers23@utexas.edu>

sroberts-amd committed Aug 14, 2026 at 13:38 UTC c794754062d49a8fdb63ab3c5215b488b865030c

3 files changed +97 -146

vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

@@ -484,6 +484,7 @@ class AiterExperts(mk.FusedMoEExpertsModular):

484 return activation in [

485 MoEActivation.SILU,

486 MoEActivation.GELU,

487 + MoEActivation.SITU,

488 MoEActivation.SWIGLUOAI,

489 MoEActivation.SWIGLUOAI_UNINTERLEAVE,

490 ]

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

+45 -19

@@ -644,20 +644,7 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(

644 activation: MoEActivation | None = None,

645 ) -> tuple[int, int]:

646 """Round up hidden_size and intermediate_size based on backend requirements."""

647 - if backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and activation == MoEActivation.SITU:

648 - # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size

649 - # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 (a

650 - # no-op for K3's shapes) rather than the generic ROCm 256 round-up,

651 - # which would inflate weights and OOM.

652 - intermediate_size = round_up(intermediate_size, 128)

653 - hidden_size = round_up(hidden_size, 128)

654 - elif (

655 - backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and activation == MoEActivation.SILU

656 - ):

657 - # AITER's A16W4 SiLU kernel handles native dimensions aligned to 128.

658 - intermediate_size = round_up(intermediate_size, 128)

659 - hidden_size = round_up(hidden_size, 128)

660 - elif backend == Mxfp4MoeBackend.EMULATION:

647 + if backend == Mxfp4MoeBackend.EMULATION:

648 # Emulation has no kernel tile; it only needs OCP MX block alignment so the

649 # per-block scale buffers (`dim // OCP_MX_BLOCK_SIZE`) aren't floor-truncated

650 # by a non-block-aligned TP/DP shard (e.g. 2880 // 4 = 720).

@@ -683,8 +670,18 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(

670 intermediate_size = round_up(intermediate_size, 128)

671 hidden_size = round_up(hidden_size, 128)

672 elif current_platform.is_rocm():

686 - intermediate_size = round_up(intermediate_size, 256)

687 - hidden_size = round_up(hidden_size, 256)

673 + if backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and (

674 + activation == MoEActivation.SITU or activation == MoEActivation.SILU

675 + ):

676 + # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size

677 + # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 (a

678 + # no-op for K3's shapes) rather than the generic ROCm 256 round-up,

679 + # which would inflate weights and OOM.

680 + intermediate_size = round_up(intermediate_size, 128)

681 + hidden_size = round_up(hidden_size, 128)

682 + else:

683 + intermediate_size = round_up(intermediate_size, 256)

684 + hidden_size = round_up(hidden_size, 256)

685 elif backend == Mxfp4MoeBackend.CPU:

686 # CPU AMX kernel uses BLOCK_N=32, align to 32

687 intermediate_size = round_up(intermediate_size, 32)

@@ -1280,6 +1277,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(

1277 w13_bias: torch.Tensor | None = None,

1278 w2_bias: torch.Tensor | None = None,

1279 _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None,

1280 + activation: MoEActivation | None = None,

1281 ) -> tuple[

1282 torch.Tensor,

1283 torch.Tensor,

@@ -1479,13 +1477,41 @@ def convert_weight_to_mxfp4_moe_kernel_format(

1477

1478 import os

1479

1482 - from aiter.ops.shuffle import shuffle_scale as _shuf_s

1483 - from aiter.ops.shuffle import shuffle_weight as _shuf_w

1484 -

1480 # TODO: Remove this once AITER is fixed

1481 # Necessary for AITER side from crashing

1482 os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0"

1483

1484 + if activation == MoEActivation.SITU:

1485 + from aiter.utility.fp4_utils import e8m0_shuffle

1486 +

1487 + from vllm._aiter_ops import rocm_aiter_ops

1488 +

1489 + fp4_dtype = torch.float4_e2m1fn_x2

1490 + e8m0_dtype = torch.float8_e8m0fnu

1491 + # a8w4 uses gate/up-interleaved flydsl kernels;

1492 + # default a16w4 keeps the separated layout.

1493 + guinterleave = rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled()

1494 + w13 = rocm_aiter_ops.shuffle_weight_a16w4(

1495 + w13_weight.data.view(fp4_dtype), 16, guinterleave

1496 + )

1497 + w2 = rocm_aiter_ops.shuffle_weight_a16w4(

1498 + w2_weight.data.view(fp4_dtype), 16, False

1499 + )

1500 + w13_scale_raw = w13_weight_scale.data.view(e8m0_dtype)

1501 + w2_scale_raw = w2_weight_scale.data.view(e8m0_dtype)

1502 + w13_scale = rocm_aiter_ops.shuffle_scale_a16w4(

1503 + w13_scale_raw.view(-1, w13_scale_raw.shape[-1]),

1504 + num_experts,

1505 + guinterleave,

1506 + )

1507 + w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1]))

1508 + w13.is_shuffled = True

1509 + w2.is_shuffled = True

1510 + return (w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)

1511 +

1512 + from aiter.ops.shuffle import shuffle_scale as _shuf_s

1513 + from aiter.ops.shuffle import shuffle_weight as _shuf_w

1514 +

1515 w13_weight = torch.nn.Parameter(

1516 _shuf_w(

1517 w13_weight.data.view(torch.float4_e2m1fn_x2),

vllm/model_executor/layers/quantization/mxfp4.py

+51 -127

@@ -1,8 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -import os

5 -

4 import torch

6 from vllm.logger import init_logger

@@ -22,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.config import (

20 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (

21 TRITON_BACKENDS,

22 Mxfp4MoeBackend,

25 - backend_to_kernel_cls,

23 convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,

24 convert_weight_to_mxfp4_moe_kernel_format,

25 make_mxfp4_moe_kernel,

@@ -472,51 +469,14 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):

469 )

470

471

475 -def _use_k3_situ_aiter(moe: FusedMoEConfig) -> bool:

476 - """Whether Kimi-K3's SiTU MXFP4 MoE should use the AITER A16W4 kernel.

477 -

478 - K3 is weight-only MXFP4 (W4A16) with SiTU activation, which the generic

479 - MXFP4 backend selector does not cover; route it to AITER on gfx950.

480 - """

481 - if not current_platform.is_rocm():

482 - return False

483 - from vllm._aiter_ops import rocm_aiter_ops

484 - from vllm.model_executor.layers.fused_moe.activation import MoEActivation

485 - from vllm.platforms.rocm import on_gfx950

486 -

487 - return (

488 - rocm_aiter_ops.is_fused_moe_enabled()

489 - and on_gfx950()

490 - and moe.activation == MoEActivation.SITU

491 - and moe.activation_situ_linear_beta is not None

492 - and rocm_aiter_ops.get_aiter_activation_type("situ") is not None

493 - )

494 -

495 -

472 class Mxfp4MoEMethod(FusedMoEMethodBase):

473 """MXFP4 MoE quantization method."""

474

475 def __init__(self, moe: FusedMoEConfig):

476 super().__init__(moe)

477 +

478 self.weight_dtype = "mxfp4"

502 - self.is_k3_situ_aiter = _use_k3_situ_aiter(moe)

503 - self.experts_cls: type[mk.FusedMoEExperts] | None

504 - if self.is_k3_situ_aiter:

505 - self.mxfp4_backend = Mxfp4MoeBackend.AITER_MXFP4_BF16

506 - self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0]

507 - logger.info_once("Using AITER_MXFP4_BF16 for Kimi-K3 SiTU MXFP4 MoE.")

508 - from vllm._aiter_ops import rocm_aiter_ops

509 -

510 - if rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled():

511 - # AITER keeps bf16 activations below this token count, which

512 - # would not match the fp8 a8w4 kernels the interleaved SiTU

513 - # path is tuned for. The a16w4 path never reads it.

514 - # TODO: Remove once AITER takes this as a kernel argument.

515 - os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0"

516 - else:

517 - self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(

518 - moe

519 - )

479 + self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(moe)

480

481 self.max_capture_size = moe.max_capture_size

482

@@ -714,62 +674,62 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

674 hidden_size = self.hidden_size

675 sf_block_size = 32

676

717 - # Shape assertions

718 - assert (

719 - w13.dim() == 3

720 - and w13.shape[0] == num_experts

721 - and w13.shape[1] == intermediate_size * self.moe.w13_num_shards

722 - and w13.shape[2] == hidden_size // 2

723 - )

724 - assert (

725 - w13_scale.dim() == 3

726 - and w13_scale.shape[0] == num_experts

727 - and w13_scale.shape[1] == intermediate_size * self.moe.w13_num_shards

728 - and w13_scale.shape[2] == hidden_size // sf_block_size

729 - )

730 - assert (

731 - w2.dim() == 3

732 - and w2.shape[0] == num_experts

733 - and w2.shape[1] == hidden_size

734 - and w2.shape[2] == intermediate_size // 2

735 - )

736 - assert (

737 - w2_scale.dim() == 3

738 - and w2_scale.shape[1] == hidden_size

739 - and w2_scale.shape[2] == intermediate_size // sf_block_size

740 - )

741 - if w13_bias is not None:

677 + # Shape assertions — skipped for SITU since its kernel handles native

678 + # (non-256-aligned) intermediate sizes without prior round-up.

679 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation

680 +

681 + if self.moe.activation != MoEActivation.SITU:

682 assert (

743 - w13_bias.dim() == 2

744 - and w13_bias.shape[0] == num_experts

745 - and w13_bias.shape[1] == intermediate_size * self.moe.w13_num_shards

683 + w13.dim() == 3

684 + and w13.shape[0] == num_experts

685 + and w13.shape[1] == intermediate_size * self.moe.w13_num_shards

686 + and w13.shape[2] == hidden_size // 2

687 )

747 - if w2_bias is not None:

688 assert (

749 - w2_bias.dim() == 2

750 - and w2_bias.shape[0] == num_experts

751 - and w2_bias.shape[1] == hidden_size

689 + w13_scale.dim() == 3

690 + and w13_scale.shape[0] == num_experts

691 + and w13_scale.shape[1] == intermediate_size * self.moe.w13_num_shards

692 + and w13_scale.shape[2] == hidden_size // sf_block_size

693 )

753 -

754 - # Convert weights to kernel format

755 - if self.is_k3_situ_aiter:

756 - w13, w2, w13_scale, w2_scale = (

757 - self._convert_k3_situ_weight_to_kernel_format(layer)

694 + assert (

695 + w2.dim() == 3

696 + and w2.shape[0] == num_experts

697 + and w2.shape[1] == hidden_size

698 + and w2.shape[2] == intermediate_size // 2

699 )

759 - else:

760 - w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (

761 - convert_weight_to_mxfp4_moe_kernel_format(

762 - mxfp4_backend=self.mxfp4_backend,

763 - layer=layer,

764 - w13_weight=w13,

765 - w2_weight=w2,

766 - w13_weight_scale=w13_scale,

767 - w2_weight_scale=w2_scale,

768 - w13_bias=w13_bias,

769 - w2_bias=w2_bias,

770 - _cache_permute_indices=self._cache_permute_indices,

700 + assert (

701 + w2_scale.dim() == 3

702 + and w2_scale.shape[1] == hidden_size

703 + and w2_scale.shape[2] == intermediate_size // sf_block_size

704 + )

705 + if w13_bias is not None:

706 + assert (

707 + w13_bias.dim() == 2

708 + and w13_bias.shape[0] == num_experts

709 + and w13_bias.shape[1] == intermediate_size * self.moe.w13_num_shards

710 )

711 + if w2_bias is not None:

712 + assert (

713 + w2_bias.dim() == 2

714 + and w2_bias.shape[0] == num_experts

715 + and w2_bias.shape[1] == hidden_size

716 + )

717 +

718 + # Convert weights to kernel format

719 + w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (

720 + convert_weight_to_mxfp4_moe_kernel_format(

721 + mxfp4_backend=self.mxfp4_backend,

722 + layer=layer,

723 + w13_weight=w13,

724 + w2_weight=w2,

725 + w13_weight_scale=w13_scale,

726 + w2_weight_scale=w2_scale,

727 + w13_bias=w13_bias,

728 + w2_bias=w2_bias,

729 + _cache_permute_indices=self._cache_permute_indices,

730 + activation=self.moe.activation,

731 )

732 + )

733

734 # For TRITON backends, weights are wrapped tensors from triton_kernels

735 # that don't support .detach(). Manually assign parameters.

@@ -810,42 +770,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

770 routing_tables=layer._expert_routing_tables(),

771 )

772

813 - def _convert_k3_situ_weight_to_kernel_format(

814 - self, layer: RoutedExperts

815 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

816 - # K3's AITER A16W4 kernel wants the separated ([gate_all, up_all])

817 - # stage-1 layout, unlike the interleaved gpt-oss/DeepSeek path in

818 - # convert_weight_to_mxfp4_moe_kernel_format. Preshuffle once here.

819 - from aiter.utility.fp4_utils import e8m0_shuffle

820 -

821 - from vllm._aiter_ops import rocm_aiter_ops

822 -

823 - fp4_dtype = torch.float4_e2m1fn_x2

824 - e8m0_dtype = torch.float8_e8m0fnu

825 - num_experts = layer.w13_weight.shape[0]

826 -

827 - # a8w4 (VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1) uses the gate/up-

828 - # interleaved (_gui_) fp8 flydsl kernels, which need w13 weight+scale

829 - # in interleave layout. Default a16w4 keeps the separated layout.

830 - guinterleave = rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled()

831 - w13 = rocm_aiter_ops.shuffle_weight_a16w4(

832 - layer.w13_weight.data.view(fp4_dtype), 16, guinterleave

833 - )

834 - w2 = rocm_aiter_ops.shuffle_weight_a16w4(

835 - layer.w2_weight.data.view(fp4_dtype), 16, False

836 - )

837 - w13_scale_raw = layer.w13_weight_scale.data.view(e8m0_dtype)

838 - w2_scale_raw = layer.w2_weight_scale.data.view(e8m0_dtype)

839 - w13_scale = rocm_aiter_ops.shuffle_scale_a16w4(

840 - w13_scale_raw.view(-1, w13_scale_raw.shape[-1]), num_experts, guinterleave

841 - )

842 - w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1]))

843 -

844 - w13.is_shuffled = True

845 - w2.is_shuffled = True

846 -

847 - return w13, w2, w13_scale, w2_scale

848 -

773 def process_weights_after_loading(self, layer):

774 w13 = layer.w13_weight

775 w2 = layer.w2_weight