Load tuned fused_moe_lora shrink and expand kernel configs separately

[yugong333]

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.

Purpose

1. Loading the kernel configs of fused_moe_lora kernel
2. Separate shrink and expand kernel configs in the fused_moe_lora function
3. Adding num_stages and num_warps parameters in the configs

Note: Based on PR #21229 and #26319

Test Plan

To benchmark the shrink and expand kernel in gate/up and down in the benchmark_lora.py:

python3 benchmarks/kernels/benchmark_lora.py list_bench --arg-pool-size 32 \ --batch-sizes 1 \ --dtype torch.bfloat16 \ --hidden-sizes 2048 \ --num-loras 16 \ --lora-ranks 16 \ --num-loras 4 \ --op-types fused_moe_lora_gate_up_shrink \ --seq-lengths 1 \ --sort-by-lora-id 1 \ --cuda-graph-nops 32 \ --top-k-nums 4 \ --num-experts 128

Test Result

Together with #26319 we can improve the OTPS 80% - 90% in GPT-OSS-120B when concurrency is 1 or 2.

(Optional) Documentation Update

[gemini-code-assist]

Code Review

This pull request introduces separate loading for shrink and expand kernel configurations for fused_moe_lora, which is a good performance optimization. My review focuses on improving code maintainability by addressing significant code duplication in two key areas. By refactoring the duplicated logic into helper functions, the code will be cleaner, easier to read, and less prone to errors in future modifications.

[chatgpt-codex-connector]

💡 Codex Review

vllm/vllm/lora/ops/triton_ops/fused_moe_lora_op.py

Lines 339 to 355 in 6d47ff8

	def _fused_moe_lora_fake(
	output: torch.Tensor,
	qcurr_hidden_states: torch.Tensor,
	lora_a_stacked: list[torch.Tensor],
	lora_b_stacked: list[torch.Tensor],
	topk_weights: torch.Tensor,
	sorted_token_ids: torch.Tensor,
	expert_ids: torch.Tensor,
	num_tokens_post_padded: torch.Tensor,
	max_lora_rank: int,
	top_k_num: int,
	block_size_m: int,
	block_size_n: int,
	block_size_k: int,
	group_size_m: int,
	mul_routed_weight: bool = False,
	) -> None:

Update fake fused_moe_lora signature for new kernel params

The real _fused_moe_lora now takes twelve tuning scalars (separate shrink/expand block sizes, warps, stages), but _fused_moe_lora_fake still exposes the old four-parameter signature. Because direct_register_custom_op registers this fake as the tracing/compile-time implementation, any torch.compile/fake tensor invocation will raise a TypeError when the fake is called with the expanded argument list. The fake implementation should mirror the updated signature so graph capture continues to work.

ℹ️ 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".

[mergify]

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

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

[mergify]

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

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

[varun-sundar-rabindranath]

looks like this change is a no-op ?

[yugong333]

This line determines the config file name based on the op_type. fused_moe_lora kernel will use this line.

[varun-sundar-rabindranath]

Yes. But, on main we have,

        if op_type == "shrink":
            config_fname = f"{gpu_name}_{op_type.upper()}.json"
        else:
            assert op_type == "expand"
            config_fname = (
                f"{gpu_name}_{op_type.upper()}_{str(add_inputs).upper()}.json"
            )

it looks like this block of change does the same, except that the if and else part are switched. Is that right ? am I missing something ?

[yugong333]

Yes I switch the if and else as only expand has add_inputs check and others do not.

[varun-sundar-rabindranath]

Thank you for working on the @yugong333 🙌

I left a few minor comments. But it generally looks good to me.

About the performance numbers reported, are those for the default config ? Also can you share some numbers for higher concurrency ? (IIUC concurrency refers to number of requests running concurrently - if that is the case, 8 seems too low, can you report numbers for concurrency in multiples of 32 upto 512) - Mostly worried about regressions for decode. Thanks.

[jeejeelee]

Some init NIT comments, mainly suggesting that W13/W2 naming should be used

[mergify]

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

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

[varun-sundar-rabindranath]

I think it is better to lower-case all keys of the configs here so we dont have to do the
config.get("UPPER_CASE", None) or config.get("lower_case") check everywhere ?

[yugong333]

Thank you @varun-sundar-rabindranath I add an function to normalize the key name into upper case to be consistent with the moe layer config name, which also is able to remove the name check everywhere.

[varun-sundar-rabindranath]

LGTM ! Thank you for working on this @yugong333 . I left some minor nit that'd be nice to have before landing this PR.

[dcmaddix]

LGTM ! Thank you for working on this @yugong333 . I left some minor nit that'd be nice to have before landing this PR.

Great thanks a lot @varun-sundar-rabindranath!