[Gelato1337]

Overview

UnfusedOAITritonExperts now accepts SiTU, giving Kimi-K3 MXFP4 MoE layers a native backend on GPUs without an AITER SiTU kernel, e.g. MI250X (gfx90a).

Claims

• Kimi-K3 (moonshotai/Kimi-K3, MXFP4) serves on 4 nodes x 8 MI250X GCDs with --moe-backend triton_unfused; without this change no native backend accepts SiTU there and EMULATION is used.
• GSM8K (full test set, 1319 questions, greedy): 97.50% (see Validation).
• One-line change; the activation already exists in apply_moe_activation.

Validation

New test_unfused_oai_triton_moe_situ in tests/kernels/moe/test_modular_oai_triton_moe.py compares the unfused Triton MoE with SiTU (Kimi-K3's beta=4.0, linear_beta=25.0) against a torch reference, for the existing MNK cases.

On 1x MI250X GCD, vLLM v0.31.0 ROCm image with this change bind-mounted:

pytest tests/kernels/moe/test_modular_oai_triton_moe.py -k situ
# without this change:
E       AssertionError: assert False
E        +  where False = <function UnfusedOAITritonExperts._supports_activation ...>(<MoEActivation.SITU: 'situ'>)
5 failed
# with this change, whole file:
pytest tests/kernels/moe/test_modular_oai_triton_moe.py
26 passed

Model evaluation, Kimi-K3, 4 nodes x 8 MI250X (LUMI), TP=8 PP=4, --enable-expert-parallel --moe-backend triton_unfused, CUDA graphs, max-model-len 8192:

Eval	Result
GSM8K, 1319 questions, greedy, chat, answer = last number, max 2048 tokens, --gpu-memory-utilization 0.88	97.50% (1286 correct; 31 wrong; 2 client timeouts counted as wrong)
Throughput, 128-token outputs	6.8 tok/s at 1 request; 110 / 178 / 193 tok/s at 64 / 128 / 256 concurrent

Without this change (same setup, automatic selection) the engine selects EMULATION.

Details

• Root cause: UnfusedOAITritonExperts runs GEMM1, then the activation through the base activation() → apply_moe_activation (which implements SiTU with activation_situ_beta/activation_situ_linear_beta from the MoE config), then GEMM2. Its _supports_activation list omitted SiTU.
• SiTU takes the [gate | up] layout (as SiLU does), not the interleaved layout of SWIGLUOAI; the MXFP4 weight conversion for TRITON_UNFUSED already keeps that layout. This matches the contiguous [w1; w3] layout listed for SiTU on TRITON_UNFUSED in [Quantization] Clarify MXFP4 MoE W13 layout requirement from backends #52240. The test helper make_weights() gains an interleave_gate_up switch and pads each half separately so the split point stays in the middle; existing callers are unchanged.
• Limitation (existing code, not changed here): the unfused path accumulates expert outputs in a temporary float32 tensor (`acc.index_add_(..., down.to(torch.float32))`) whose size follows how many tokens reach a rank's experts. With expert parallelism and 128 concurrent prefills, real routing exceeded what memory profiling reserved and one rank ran out of memory at `--gpu-memory-utilization 0.95`; the evaluation uses 0.88.
• Not a duplicate: no open PR adds SiTU to the OAI Triton experts.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀