@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e7b68f4d6c
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 10, 10:15 UTC

[Bugfix] Fix Triton FusedMoE LoRA (#30585)

Signed-off-by: Xin Yang <xyangx@amazon.com>

Xin Yang committed Jan 9, 2026 at 03:46 UTC e7b68f4d6c7e4b71820994566b7317f24e613b2e
3 files changed +53 -37
pyproject.toml
+1
@@ -167,6 +167,7 @@ depthwise_seperable_CNN = "depthwise_seperable_CNN"
167 [tool.typos.default.extend-words]
168 iy = "iy"
169 tendencias = "tendencias"
170 +indx = "indx"
171 # intel cpu features
172 tme = "tme"
173 dout = "dout"
tests/lora/test_gptoss_tp.py
+47 -34
@@ -69,41 +69,54 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:
69 assert generated_texts[i].startswith(EXPECTED_LORA_OUTPUT[i])
70
71
72 -def test_gpt_oss_lora(gptoss20b_lora_files):
73 - llm = vllm.LLM(
74 - MODEL_PATH,
75 - max_model_len=1024,
76 - enable_lora=True,
77 - max_loras=4,
78 - max_lora_rank=8,
79 - max_num_seqs=2,
80 - max_num_batched_tokens=2048,
81 - compilation_config=vllm.config.CompilationConfig( # Avoid OOM
82 - cudagraph_specialize_lora=False,
83 - ),
84 - )
85 -
86 - generate_and_test(llm, gptoss20b_lora_files, lora_id=1)
87 - generate_and_test(llm, gptoss20b_lora_files, lora_id=2)
72 +@pytest.mark.parametrize("mxfp4_use_marlin", [True, False])
73 +def test_gpt_oss_lora(
74 + monkeypatch: pytest.MonkeyPatch, gptoss20b_lora_files, mxfp4_use_marlin
75 +):
76 + with monkeypatch.context() as m:
77 + m.setenv("VLLM_MXFP4_USE_MARLIN", "1" if mxfp4_use_marlin else "0")
78 + llm = vllm.LLM(
79 + MODEL_PATH,
80 + max_model_len=1024,
81 + enable_lora=True,
82 + max_loras=4,
83 + max_lora_rank=8,
84 + max_num_seqs=2,
85 + max_num_batched_tokens=2048,
86 + compilation_config=vllm.config.CompilationConfig( # Avoid OOM
87 + cudagraph_specialize_lora=False,
88 + ),
89 + )
90 +
91 + generate_and_test(llm, gptoss20b_lora_files, lora_id=1)
92 + generate_and_test(llm, gptoss20b_lora_files, lora_id=2)
93
94
95 @multi_gpu_test(num_gpus=2)
96 @pytest.mark.parametrize("fully_sharded_loras", [False, True])
92 -def test_gpt_oss_lora_tp2(gptoss20b_lora_files, fully_sharded_loras):
93 - llm = vllm.LLM(
94 - MODEL_PATH,
95 - max_model_len=1024,
96 - enable_lora=True,
97 - max_loras=2,
98 - max_num_seqs=2,
99 - max_num_batched_tokens=2048,
100 - tensor_parallel_size=2,
101 - gpu_memory_utilization=0.8,
102 - fully_sharded_loras=fully_sharded_loras,
103 - compilation_config=vllm.config.CompilationConfig( # Avoid OOM
104 - cudagraph_specialize_lora=False,
105 - ),
106 - )
107 -
108 - generate_and_test(llm, gptoss20b_lora_files, lora_id=1)
109 - generate_and_test(llm, gptoss20b_lora_files, lora_id=2)
97 +@pytest.mark.parametrize("mxfp4_use_marlin", [True, False])
98 +def test_gpt_oss_lora_tp2(
99 + monkeypatch: pytest.MonkeyPatch,
100 + gptoss20b_lora_files,
101 + fully_sharded_loras,
102 + mxfp4_use_marlin,
103 +):
104 + with monkeypatch.context() as m:
105 + m.setenv("VLLM_MXFP4_USE_MARLIN", "1" if mxfp4_use_marlin else "0")
106 + llm = vllm.LLM(
107 + MODEL_PATH,
108 + max_model_len=1024,
109 + enable_lora=True,
110 + max_loras=2,
111 + max_num_seqs=2,
112 + max_num_batched_tokens=2048,
113 + tensor_parallel_size=2,
114 + gpu_memory_utilization=0.8,
115 + fully_sharded_loras=fully_sharded_loras,
116 + compilation_config=vllm.config.CompilationConfig( # Avoid OOM
117 + cudagraph_specialize_lora=False,
118 + ),
119 + )
120 +
121 + generate_and_test(llm, gptoss20b_lora_files, lora_id=1)
122 + generate_and_test(llm, gptoss20b_lora_files, lora_id=2)
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+5 -3
@@ -502,16 +502,18 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
502 )
503
504 self.activation(
505 - activation, intermediate_cache2, intermediate_cache1.view(-1, N)
505 + activation,
506 + intermediate_cache2,
507 + intermediate_cache1.view(-1, N)[gather_indx.dst_indx],
508 )
509
510 # matmul_ogs grouped reduction fuse sum across multiple experts:
509 - # y[dst_ind // n_expts_act, :] += x[src_ind, :]
511 + # y[dst_indx // n_expts_act, :] += x
512 # Need to set n_expts_act to 1 to unfuse moe_sum
513 routing_data.n_expts_act = 1
514
515 matmul_ogs(
514 - intermediate_cache2,
516 + intermediate_cache2[gather_indx.src_indx],
517 w2,
518 self.quant_config.w2_bias,
519 routing_data,