@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7b54690244

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 18:15 UTC

[ROCm][Perf] Expose AITER MoE sorting dispatch policy via env var (#39177)

Signed-off-by: nholmber <nholmber@users.noreply.github.com>

Nico Holmberg committed May 26, 2026 at 22:11 UTC 7b546902447c695c3a555a81352719710d4f1783

3 files changed +32

vllm/_aiter_ops.py

+17

@@ -144,6 +144,7 @@ def _rocm_aiter_fused_moe_impl(

144 intermediate_pad: int = 0,

145 bias1: torch.Tensor | None = None,

146 bias2: torch.Tensor | None = None,

147 + moe_sorting_dispatch_policy: int = 0,

148 ) -> torch.Tensor:

149 from aiter import ActivationType, QuantType

150 from aiter.fused_moe import fused_moe

@@ -171,6 +172,7 @@ def _rocm_aiter_fused_moe_impl(

172 intermediate_pad=intermediate_pad,

173 bias1=bias1,

174 bias2=bias2,

175 + moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,

176 )

177

178

@@ -194,6 +196,7 @@ def _rocm_aiter_fused_moe_fake(

196 intermediate_pad: int = 0,

197 bias1: torch.Tensor | None = None,

198 bias2: torch.Tensor | None = None,

199 + moe_sorting_dispatch_policy: int = 0,

200 ) -> torch.Tensor:

201 if output_dtype is not None:

202 return torch.empty_like(hidden_states, dtype=output_dtype)

@@ -1282,6 +1285,18 @@ class rocm_aiter_ops:

1285 - Triton ops: triton_rotary_embed, triton_fp8_bmm, triton_gemm_a8w8_blockscale

1286 """

1287

1288 + _MOE_DISPATCH_POLICY: int | None = None

1289 +

1290 + @classmethod

1291 + @if_aiter_supported

1292 + def get_moe_dispatch_policy(cls) -> int:

1293 + """Cached MoE sorting dispatch policy."""

1294 + if cls._MOE_DISPATCH_POLICY is None:

1295 + import vllm.envs as envs

1296 +

1297 + cls._MOE_DISPATCH_POLICY = envs.VLLM_ROCM_AITER_MOE_DISPATCH_POLICY

1298 + return cls._MOE_DISPATCH_POLICY

1299 +

1300 # Check if the env variable is set

1301 _AITER_ENABLED = envs.VLLM_ROCM_USE_AITER

1302 _LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR

@@ -1890,6 +1905,7 @@ class rocm_aiter_ops:

1905 intermediate_pad: int = 0,

1906 bias1: torch.Tensor | None = None,

1907 bias2: torch.Tensor | None = None,

1908 + moe_sorting_dispatch_policy: int = 0,

1909 ) -> torch.Tensor:

1910 return torch.ops.vllm.rocm_aiter_fused_moe(

1911 hidden_states,

@@ -1911,6 +1927,7 @@ class rocm_aiter_ops:

1927 intermediate_pad,

1928 bias1,

1929 bias2,

1930 + moe_sorting_dispatch_policy,

1931 )

1932

1933 @staticmethod

vllm/envs.py

+12

@@ -116,6 +116,7 @@ if TYPE_CHECKING:

116 VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False

117 VLLM_ROCM_USE_AITER_LINEAR: bool = True

118 VLLM_ROCM_USE_AITER_MOE: bool = True

119 + VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0

120 VLLM_ROCM_USE_AITER_RMSNORM: bool = True

121 VLLM_ROCM_USE_AITER_MLA: bool = True

122 VLLM_ROCM_USE_AITER_MHA: bool = True

@@ -1103,6 +1104,17 @@ environment_variables: dict[str, Callable[[], Any]] = {

1104 "VLLM_ROCM_USE_AITER_MOE": lambda: (

1105 os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1")

1106 ),

1107 + # MoE sorting dispatch policy for AITER fused MoE kernels.

1108 + # 0 = auto (default): single-pass for small batches, multi-pass

1109 + # for large batches

1110 + # 1 = always single-pass: one kernel launch, no workspace,

1111 + # may be preferred for low-concurrency decode workloads

1112 + # 2 = always multi-pass: can be faster for MoE-heavy models

1113 + # (e.g., +2-5% on Qwen3-Next, +1.5% on DeepSeek-V3 at TP4,

1114 + # see PR #39177 for benchmarks)

1115 + "VLLM_ROCM_AITER_MOE_DISPATCH_POLICY": lambda: int(

1116 + os.getenv("VLLM_ROCM_AITER_MOE_DISPATCH_POLICY", "0")

1117 + ),

1118 # use aiter rms norm op if aiter ops are enabled.

1119 "VLLM_ROCM_USE_AITER_RMSNORM": lambda: (

1120 os.getenv("VLLM_ROCM_USE_AITER_RMSNORM", "True").lower() in ("true", "1")

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

@@ -246,6 +246,7 @@ def rocm_aiter_fused_experts(

246 a1q_scale: torch.Tensor | None = None,

247 num_local_tokens: torch.Tensor | None = None,

248 output_dtype: torch.dtype | None = None,

249 + moe_sorting_dispatch_policy: int = 0,

250 ) -> torch.Tensor:

251 """ROCm AITER fused MoE expert computation."""

252 if quant_config is None:

@@ -361,6 +362,7 @@ def rocm_aiter_fused_experts(

362 intermediate_pad=intermediate_pad // 64 * 64 * 2,

363 bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,

364 bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,

365 + moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,

366 )

367

368

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

506 a1q_scale=a1q_scale,

507 num_local_tokens=num_local_tokens,

508 output_dtype=output.dtype,

509 + moe_sorting_dispatch_policy=rocm_aiter_ops.get_moe_dispatch_policy(),

510 )

511 # avoid redundant copy when output is a view of the result

512 if (