Skip to content

Conversation

@xyang16
Copy link
Contributor

@xyang16 xyang16 commented Nov 18, 2025

Purpose

This PR is to support FusedMoE LoRA Triton kernel for mxfp4 model.

  • Add UnfusedOAITritonExperts
  • Inject lora module in activation: Since matmul_ogs can fuse activation, set fused_activation to None to unfuse activation in the first matmul_ogs.
  • Inject lora module in moe_sum: This need to unfuse sum in second matmul_ogs. Grouped reduction does scatter + accumulate, it is essentially equal to: y[dst_indx // n_expts_act, :] += x[src_indx, :], so that scatter sum across multiple experts, and collapse M * topk to M rows. Therefore, we need to set routing_data.n_expts_act to 1, so it doesn't sum across multiple experts, in order unfuse moe_sum in the second matmul_ogs.
  • Added test_modular_oai_triton_moe.py

Test Plan

pytest -s -v tests/kernels/moe/test_modular_oai_triton_moe.py

Test Result

Tests passed.

Benchmark

Baseline (marlin):

VLLM_MXFP4_USE_MARLIN=1 vllm serve openai/gpt-oss-20b \
  --tensor-parallel-size 1 \
  --max-num-seqs 16 \
  --compilation-config '{"max_cudagraph_capture_size": 128, "compile_sizes": [1, 2, 4, 8, 16]}' \
  --enable-lora \
  --max-loras 6 \
  --lora-modules lora1=/opt/dlami/nvme/models/gpt-oss-20b-lora-gsm8k \
  --max-lora-rank 64
vllm bench serve \
  --model openai/gpt-oss-20b \
  --lora-modules lora1 \
  --dataset-name sharegpt \
  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
  --max-concurrency 16 \
  --num-prompts 1000 \
  --num-warmups 60 \
  --ignore-eos
============ Serving Benchmark Result ============
Successful requests:                     1000      
Failed requests:                         0         
Maximum request concurrency:             16        
Benchmark duration (s):                  142.34    
Total input tokens:                      215312    
Total generated tokens:                  199033    
Request throughput (req/s):              7.03      
Output token throughput (tok/s):         1398.29   
Peak output token throughput (tok/s):    1677.00   
Peak concurrent requests:                29.00     
Total Token throughput (tok/s):          2910.95   
---------------Time to First Token----------------
Mean TTFT (ms):                          43.24     
Median TTFT (ms):                        26.77     
P99 TTFT (ms):                           140.09    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          11.19     
Median TPOT (ms):                        10.93     
P99 TPOT (ms):                           17.44     
---------------Inter-token Latency----------------
Mean ITL (ms):                           11.12     
Median ITL (ms):                         9.73      
P99 ITL (ms):                            56.55     
==================================================

PR (triton):

Install triton_kernels

pip install "git+https://github.com/triton-lang/triton.git@0a2e3a391cbb9e13d29bf12a2a0005e358102d74#subdirectory=python/triton_kernels"
vllm serve openai/gpt-oss-20b \
  --tensor-parallel-size 1 \
  --max-num-seqs 16 \
  --compilation-config '{"max_cudagraph_capture_size": 128, "compile_sizes": [1, 2, 4, 8, 16]}' \
  --enable-lora \
  --max-loras 6 \
  --lora-modules lora1=/opt/dlami/nvme/models/gpt-oss-20b-lora-gsm8k \
  --max-lora-rank 64
vllm bench serve \
  --model openai/gpt-oss-20b \
  --lora-modules lora1 \
  --dataset-name sharegpt \
  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
  --max-concurrency 16 \
  --num-prompts 1000 \
  --num-warmups 60 \
  --ignore-eos
============ Serving Benchmark Result ============
Successful requests:                     1000      
Failed requests:                         0         
Maximum request concurrency:             16        
Benchmark duration (s):                  125.59    
Total input tokens:                      215312    
Total generated tokens:                  199033    
Request throughput (req/s):              7.96      
Output token throughput (tok/s):         1584.77   
Peak output token throughput (tok/s):    2171.00   
Peak concurrent requests:                29.00     
Total Token throughput (tok/s):          3299.15   
---------------Time to First Token----------------
Mean TTFT (ms):                          52.14     
Median TTFT (ms):                        19.03     
P99 TTFT (ms):                           129.60    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          9.82      
Median TPOT (ms):                        9.28      
P99 TPOT (ms):                           21.31     
---------------Inter-token Latency----------------
Mean ITL (ms):                           9.72      
Median ITL (ms):                         7.32      
P99 ITL (ms):                            90.10     
==================================================

Accuracy Testing

  • 20b deepep triton
VLLM_ALL2ALL_BACKEND="deepep_high_throughput" vllm serve openai/gpt-oss-20b \
  --tensor-parallel-size 1 \
  --data-parallel-size 2 \
  --enable-expert-parallel \
  --no-enable-prefix-caching
OPENAI_API_KEY=EMPTY python3 -m gpt_oss.evals --model openai/gpt-oss-20b --eval gpqa --n-threads 200 --reasoning-effort low
Writing report to /tmp/gpqa_openai__gpt-oss-20b-low_temp1.0_20251125_122239.html
{'chars': np.float64(55.219065656565654), 'chars:std': np.float64(220.291583893894), 'score': np.float64(0.5707070707070707), 'score:std': np.float64(0.4949752621616814)}
Writing results to /tmp/gpqa_openai__gpt-oss-20b-low_temp1.0_20251125_122239.json
Writing all results to /tmp/gpqa_openai__gpt-oss-20b-low_temp1.0_20251125_122239_allresults.json
[{'eval_name': 'gpqa', 'model_name': 'openai__gpt-oss-20b-low_temp1.0_20251125_122239', 'metric': 0.5707070707070707}]
  • 20b lora triton
vllm serve openai/gpt-oss-20b \
  --tensor-parallel-size 1 \
  --compilation-config '{"cudagraph_mode": "PIECEWISE"}' \
  --max-num-seqs 16 \
  --enable-lora \
  --max-loras 6 \
  --lora-modules lora1=/opt/dlami/nvme/models/gpt-oss-20b-lora/checkpoint-13 \
  --max-lora-rank 64
OPENAI_API_KEY=EMPTY python3 -m gpt_oss.evals --model lora1 --eval gpqa --n-threads 200 --reasoning-effort low
Writing report to /tmp/gpqa_openai__gpt-oss-20b-low_temp1.0_20251126_124509.html
{'chars': np.float64(62.859848484848484), 'chars:std': np.float64(230.21663448129942), 'score': np.float64(0.577020202020202), 'score:std': np.float64(0.4940322747359399)}
Writing results to /tmp/gpqa_openai__gpt-oss-20b-low_temp1.0_20251126_124509.json
Writing all results to /tmp/gpqa_openai__gpt-oss-20b-low_temp1.0_20251126_124509_allresults.json
[{'eval_name': 'gpqa', 'model_name': 'openai__gpt-oss-20b-low_temp1.0_20251126_124509', 'metric': 0.577020202020202}]
  • 20b lora marlin
VLLM_MXFP4_USE_MARLIN=1 vllm serve openai/gpt-oss-20b \
  --tensor-parallel-size 1 \
  --compilation-config '{"cudagraph_mode": "PIECEWISE"}' \
  --max-num-seqs 16 \
  --enable-lora \
  --max-loras 6 \
  --lora-modules lora1=/opt/dlami/nvme/models/gpt-oss-20b-lora/checkpoint-13 \
  --max-lora-rank 64
OPENAI_API_KEY=EMPTY python3 -m gpt_oss.evals --model lora1 --eval gpqa --n-threads 200 --reasoning-effort low
Writing report to /tmp/gpqa_lora1-low_temp1.0_20251126_125725.html
{'chars': np.float64(25.56691919191919), 'chars:std': np.float64(165.94325370775), 'score': np.float64(0.5669191919191919), 'score:std': np.float64(0.4955015860245882)}
Writing results to /tmp/gpqa_lora1-low_temp1.0_20251126_125725.json
Writing all results to /tmp/gpqa_lora1-low_temp1.0_20251126_125725_allresults.json
[{'eval_name': 'gpqa', 'model_name': 'lora1-low_temp1.0_20251126_125725', 'metric': 0.5669191919191919}]

Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
  • (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces support for FusedMoE LoRA with Triton kernels for mxfp4 models, which provides a significant performance improvement as shown in the benchmarks. The changes are well-structured, adding the necessary logic to select the Triton backend and adapting the kernels for this new path. However, I've identified a critical issue where attributes in OAITritonExperts are used without being initialized, which could lead to a runtime error in non-LoRA use cases. Please address this to ensure the stability of the implementation.

Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

@xyang16 xyang16 force-pushed the fused_moe_lora_triton branch 3 times, most recently from 8acf2f1 to 8f24eec Compare November 19, 2025 00:59
@robertgshaw2-redhat
Copy link
Collaborator

nice speedup!

@xyang16 xyang16 force-pushed the fused_moe_lora_triton branch 4 times, most recently from b8fd020 to 168d8cd Compare November 19, 2025 02:00
@xyang16 xyang16 changed the title Support FusedMoE LoRA Triton kernel for mxfp4 model [LoRA] Support FusedMoE LoRA Triton kernel for mxfp4 model Nov 19, 2025
modular_triton_fused_moe,
try_get_optimal_moe_config,
)
from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@xyang16 can you add a unit test for gpt-oss lora + triton_kernels. The test can be predicated on has_triton_kernels like in https://github.com/vllm-project/vllm/blob/main/tests/kernels/moe/test_gpt_oss_triton_kernels.py

Copy link
Contributor Author

@xyang16 xyang16 Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added test_modular_oai_triton_moe.py. Thanks!

Copy link
Contributor

@varun-sundar-rabindranath varun-sundar-rabindranath left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The changes to triton_kernel_fused_experts are invasive and it is a bit confusing reason about the fused_act=True, fuse_sum=True and fuse_act=False,fused_sum=False cases as the assumptions and expectations from matmul_ogs is different in both cases.

The main difference between the non-LoRA and the LoRA case,

  • For the non-LoRA case, no assumptions about the sizes of matmul_ogs output tensors are made. The only requirement here is that the second matmul_ogs must return a tensor of size [M, K]. For the LoRA case, we expect the outputs to be of a specific shape - This pattern is similar to TritonExperts
  • For the non-LoRA case, there are no requirements on the gather_indx and scatter_indx sizes. The LoRA case requires the tensors in these objects to be a specific shape.

For these reasons, I think it will be better to create a separate implementation of BaseOAITritonExperts class for the LoRA case, naming it something like UnfusedOAITritonExperts. Apart of being easier to assert for expectations, with this we can create correct and adequate workspace shapes for both workspace13 and workspace2 and reuse them properly in the implementation. Please refer to TritonExperts I think the implementation here would be very similar and all the logic could be contained within the apply function, thus not disturbing the existing triton_kernel_fused_experts function. something like,

def apply():
        routing_data, gather_indx, scatter_indx = self._make_routing_data(
            topk_ids, topk_weights, local_num_experts
        )
        matmul_ogs(..., 
                             y = intermediate_cache1)
        activation(intermediate_cache2, intermediate_cache1)
       matmul_ogs(...,
                            y = intermediate_cache3)


if with_lora_support:
return get_mxfp4_backend_with_lora()

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should maintain get_mxfp4_backend_with_lora() and return the appropriate backend from within that function. This is because, there is no guarantee that the logic below will choose a LoRA compatible backend.

Copy link
Contributor Author

@xyang16 xyang16 Nov 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added back get_mxfp4_backend_with_lora. Thanks for review!

precision_config=quant_config.w1_precision,
gammas=gammas if apply_router_weight_on_input else None,
fused_activation=act,
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@xyang16 To guarantee that matmul_ogs will return an output shaped [M * topk, N], I think it is better to pass in the output tensor ourselves using the argument y. Note that matmul_ogs actually checks if the output is of expected size here https://github.com/triton-lang/triton/blob/c3c476f357f1e9768ea4e45aa5c17528449ab9ef/python/triton_kernels/triton_kernels/matmul_ogs.py#L180 . That way it is guaranteed that matmul_ogs will respect the contract.

Same for the second matmul_ogs also.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Raised a separate PR #29219 for this

apply_router_weight_on_input: bool,
weight_and_reduce_impl: mk.TopKWeightAndReduce,
) -> None:
) -> torch.Tensor:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This violates the base class contract at

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changed back. Thanks.

@github-project-automation github-project-automation bot moved this from To Triage to In progress in gpt-oss Issues & Enhancements Nov 19, 2025
@xyang16 xyang16 force-pushed the fused_moe_lora_triton branch from 168d8cd to 1569681 Compare November 19, 2025 17:18
@mergify
Copy link

mergify bot commented Nov 19, 2025

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @xyang16.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Nov 19, 2025
@xyang16 xyang16 force-pushed the fused_moe_lora_triton branch from 1569681 to ae8d73f Compare November 23, 2025 06:23
@mergify mergify bot removed the needs-rebase label Nov 23, 2025
@xyang16 xyang16 force-pushed the fused_moe_lora_triton branch 2 times, most recently from 7d727f8 to b7ca4ae Compare November 23, 2025 06:32
Signed-off-by: Xin Yang <[email protected]>
Signed-off-by: Xin Yang <[email protected]>
Signed-off-by: Xin Yang <[email protected]>
@xyang16 xyang16 force-pushed the fused_moe_lora_triton branch from b24893c to 867d973 Compare November 26, 2025 07:08
@xyang16
Copy link
Contributor Author

xyang16 commented Nov 26, 2025

@xyang16 Could you please resolve conflicts?

@jeejeelee I have resolved the conflicts. Thanks a lot!

@xyang16
Copy link
Contributor Author

xyang16 commented Nov 26, 2025

Thank you for the work @xyang16. It generally looks good to me - just left some minor comments. Can you add some gpt-oss eval numbers for OAITritonExperts and UnfusedOAITritonExperts please.

I have added gpt-oss eval numbers. Please take a look. Thanks!

I see the vllm serve commands have --compilation-config '{"max_cudagraph_capture_size": 16, "compile_sizes": [1, 2, 4, 8, 16]}' - can you confirm that it passes without the --compilation-config arguments ? Thanks.

@varun-sundar-rabindranath I have rerun the gpt-oss eval numbers, but I have to set --compilation-config '{"cudagraph_mode": "PIECEWISE"}'. Otherwise I will get !!! in output. See the details in #29539.

Fixing the !!! issue is beyond the scope of this PR. So I have to set PIECEWISE mode. Thanks!

  • 20b lora triton
vllm serve openai/gpt-oss-20b \
  --tensor-parallel-size 1 \
  --compilation-config '{"cudagraph_mode": "PIECEWISE"}' \
  --max-num-seqs 16 \
  --enable-lora \
  --max-loras 6 \
  --lora-modules lora1=/opt/dlami/nvme/models/gpt-oss-20b-lora/checkpoint-13 \
  --max-lora-rank 64
Writing report to /tmp/gpqa_openai__gpt-oss-20b-low_temp1.0_20251126_124509.html
{'chars': np.float64(62.859848484848484), 'chars:std': np.float64(230.21663448129942), 'score': np.float64(0.577020202020202), 'score:std': np.float64(0.4940322747359399)}
Writing results to /tmp/gpqa_openai__gpt-oss-20b-low_temp1.0_20251126_124509.json
Writing all results to /tmp/gpqa_openai__gpt-oss-20b-low_temp1.0_20251126_124509_allresults.json
[{'eval_name': 'gpqa', 'model_name': 'openai__gpt-oss-20b-low_temp1.0_20251126_124509', 'metric': 0.577020202020202}]
  • 20b lora marlin
Writing report to /tmp/gpqa_lora1-low_temp1.0_20251126_125725.html
{'chars': np.float64(25.56691919191919), 'chars:std': np.float64(165.94325370775), 'score': np.float64(0.5669191919191919), 'score:std': np.float64(0.4955015860245882)}
Writing results to /tmp/gpqa_lora1-low_temp1.0_20251126_125725.json
Writing all results to /tmp/gpqa_lora1-low_temp1.0_20251126_125725_allresults.json
[{'eval_name': 'gpqa', 'model_name': 'lora1-low_temp1.0_20251126_125725', 'metric': 0.5669191919191919}]

@jeejeelee
Copy link
Collaborator

@xyang16 All tests pass correctly, but I realized that the GPT-OSS LoRA test does not select the Triton backend. Could you verify whether the Triton backend can pass correctly?

@xyang16
Copy link
Contributor Author

xyang16 commented Nov 28, 2025

@xyang16 All tests pass correctly, but I realized that the GPT-OSS LoRA test does not select the Triton backend. Could you verify whether the Triton backend can pass correctly?

@jeejeelee Thanks a lot! Triton backend can pass with slight modification of test_gpt_oss_lora_tp2 like below. Because FULL_AND_PIECEWISE cudagraph mode is causing some !!! in output, you can see the details in #29539. But fixing that problem is beyond the scope of this PR.

        compilation_config=vllm.config.CompilationConfig(  # Avoid OOM
            cudagraph_mode=CUDAGraphMode.PIECEWISE,
            cudagraph_specialize_lora=False,
        ),

@jeejeelee
Copy link
Collaborator

@xyang16 Any plan to fix this issue, if we have ,I think we can land this PR now

@xyang16
Copy link
Contributor Author

xyang16 commented Nov 28, 2025

@xyang16 Any plan to fix this issue, if we have ,I think we can land this PR now

@jeejeelee Yes, I think the issue should be fixed soon. I talked about this issue with @varun-sundar-rabindranath and he suggested me to create an issue for this, and see what other people think how to fix it. Maybe quickest is to not make FULL_AND_PIECEWISE default (FULL_AND_PIECEWISE was made default recently starting from 0.11.0).

Copy link
Collaborator

@jeejeelee jeejeelee left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for contribution

@github-project-automation github-project-automation bot moved this from In progress to Ready in gpt-oss Issues & Enhancements Nov 28, 2025
@jeejeelee jeejeelee merged commit 745a3ba into vllm-project:main Nov 28, 2025
60 checks passed
@xyang16 xyang16 deleted the fused_moe_lora_triton branch November 28, 2025 03:11
hl475 added a commit to hl475/vllm that referenced this pull request Nov 28, 2025
@DarkLight1337
Copy link
Member

Reverted by #29697, please redo

vllm-bot pushed a commit that referenced this pull request Nov 28, 2025
kitaekatt pushed a commit to kitaekatt/vllm that referenced this pull request Dec 1, 2025
kitaekatt pushed a commit to kitaekatt/vllm that referenced this pull request Dec 1, 2025
amd-hhashemi pushed a commit to amd-hhashemi/vllm that referenced this pull request Dec 2, 2025
)

Signed-off-by: Xin Yang <[email protected]>
Co-authored-by: Jee Jee Li <[email protected]>
Signed-off-by: Hashem Hashemi <[email protected]>
amd-hhashemi pushed a commit to amd-hhashemi/vllm that referenced this pull request Dec 2, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

gpt-oss Related to GPT-OSS models ready ONLY add when PR is ready to merge/full CI is needed

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

5 participants