679c6a3ecc · tduncan/Nvidia-Optimization-Stack-VLLM

.buildkite/test-amd.yaml

+1

@@ -2526,6 +2526,7 @@ steps:

2526 - pytest -v -s -x lora/test_llm_with_multi_loras.py

2527 - pytest -v -s -x lora/test_olmoe_tp.py

2528 - pytest -v -s -x lora/test_gptoss_tp.py

2529 + - pytest -v -s -x lora/test_qwen35_densemoel_lora.py

2530

2531

2532 - label: Weight Loading Multiple GPU # 7.5m

tests/lora/test_gptoss_tp.py

+11

@@ -5,6 +5,7 @@ import pytest

5

6 import vllm

7 from vllm.lora.request import LoRARequest

8 +from vllm.platforms import current_platform

9

10 from ..utils import multi_gpu_test

11

@@ -69,6 +70,16 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:

70 assert generated_texts[i].startswith(EXPECTED_LORA_OUTPUT[i])

71

72

73 +@pytest.mark.skipif(

74 + not current_platform.is_cuda(),

75 + reason=(

76 + "Mxfp4 LoRA on ROCm is blocked by a spawn compatibility issue. "

77 + "The fused_moe_lora Triton kernel crashes in spawned subprocesses, "

78 + "and vLLM forces spawn mode when HIP is initialized before "

79 + "multiprocessing. Fixing this requires either making the LoRA "

80 + "Triton kernel spawn-safe or pre-warming the kernel cache."

81 + ),

82 +)

83 @pytest.mark.parametrize("mxfp4_use_marlin", [True, False])

84 @pytest.mark.parametrize("specialize_active_lora", [True, False])

85 def test_gpt_oss_lora(

vllm/lora/layers/fused_moe.py

+2 -2

@@ -109,8 +109,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

109 else: # fall back to the default config

110 get_config_func = functools.partial(

111 try_get_optimal_moe_lora_config,

112 - w1_shape=layer.w13_weight.size(),

113 - w2_shape=layer.w2_weight.size(),

112 + w1_shape=layer.w13_weight.shape,

113 + w2_shape=layer.w2_weight.shape,

114 rank=rank,

115 top_k=top_k,

116 dtype=config_dtype,

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

+5 -1

@@ -379,7 +379,11 @@ def _fused_moe_lora_kernel(

379 )

380 a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak

381

382 - accumulator += tl.dot(a, b)

382 + # Cast operands to matching dtype for tl.dot. On ROCm, Triton's

383 + # compiler may infer different types for a and b when merging

384 + # if/else branches (TMA desc path returns fp32, tl.load returns

385 + # the pointer's element type).

386 + accumulator += tl.dot(a.to(tl.bfloat16), b.to(tl.bfloat16))

387

388 if MUL_ROUTED_WEIGHT:

389 moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0)

vllm/model_executor/layers/fused_moe/config.py

+3

@@ -229,6 +229,9 @@ class FusedMoEQuantConfig:

229 _w1: FusedMoEQuantDesc

230 _w2: FusedMoEQuantDesc

231 is_nvfp4_scale_swizzled: bool = True

232 + # CK MXFP4 (gfx950) padding info for rocm_aiter_ops.fused_moe()

233 + hidden_pad: int = 0

234 + intermediate_pad: int = 0

235

236 def __post_init__(self):

237 assert not self.per_act_token_quant or self.block_shape is None, (

vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py

+5 -5

@@ -257,7 +257,7 @@ def triton_kernel_moe_forward(

257 # sparse_logits.indx contains global expert IDs – remap to local.

258 topk_ids = expert_map[sparse_logits.indx.to(torch.long)]

259 topk_weights = sparse_logits.vals

260 - local_num_experts = w1.size(0)

260 + local_num_experts = w1.shape[0]

261 routing_data, gather_idx, scatter_idx = make_routing_data(

262 topk_ids, topk_weights, local_num_experts

263 )

@@ -604,8 +604,8 @@ class BaseOAITritonExperts(mk.FusedMoEExpertsModular):

604 require a specialized implementation, like MarlinExperts, they are free

605 to override this function.

606 """

607 - assert w1.dim() == 3 and w2.dim() == 3

608 - E, _, N = w1.size()

607 + assert len(w1.shape) == 3 and len(w2.shape) == 3

608 + E, _, N = w1.shape

609 K = a1.size(-1)

610

611 assert a1.dim() == 2

@@ -683,7 +683,7 @@ class OAITritonExperts(BaseOAITritonExperts):

683 if expert_map is not None:

684 topk_ids = expert_map[topk_ids]

685

686 - local_num_experts = w1.size(0)

686 + local_num_experts = w1.shape[0]

687 if global_num_experts == -1:

688 global_num_experts = local_num_experts

689

@@ -781,7 +781,7 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):

781 if expert_map is not None:

782 topk_ids = expert_map[topk_ids]

783

784 - local_num_experts = w1.size(0)

784 + local_num_experts = w1.shape[0]

785 if global_num_experts == -1:

786 global_num_experts = local_num_experts

787

vllm/model_executor/layers/fused_moe/layer.py

+8 -1

@@ -567,6 +567,13 @@ class FusedMoE(CustomOp):

567 # for heuristic purposes, so it must be initialized first.

568 self.quant_method: FusedMoEMethodBase = _get_quant_method()

569

570 + # Quant methods (e.g. Mxfp4MoEMethod) may round up hidden_dim

571 + # and intermediate_size in moe_config during __init__. Sync

572 + # self.hidden_size so downstream consumers (e.g. LoRA) see the

573 + # padded value.

574 + if self.moe_config.hidden_dim != self.hidden_size:

575 + self.hidden_size = self.moe_config.hidden_dim

576 +

577 if not self.moe_config.is_act_and_mul and not current_platform.is_cuda_alike():

578 raise NotImplementedError(

579 "is_act_and_mul=False is supported only for CUDA and ROCm for now"

@@ -586,7 +593,7 @@ class FusedMoE(CustomOp):

593

594 moe_quant_params = {

595 "num_experts": self.local_num_experts,

589 - "hidden_size": hidden_size,

596 + "hidden_size": self.hidden_size,

597 "unpadded_hidden_size": unpadded_hidden_size,

598 "intermediate_size_per_partition": self.intermediate_size_per_partition,

599 "params_dtype": params_dtype,

vllm/model_executor/layers/fused_moe/modular_kernel.py

+3 -3

@@ -768,8 +768,8 @@ class FusedMoEExpertsModular(FusedMoEExperts):

768 require a specialized implementation, like MarlinExperts, they are free

769 to override this function.

770 """

771 - assert w1.dim() == 3 and w2.dim() == 3

772 - E, N, _ = w1.size()

771 + assert len(w1.shape) == 3 and len(w2.shape) == 3

772 + E, N, _ = w1.shape

773 K = a1.size(-1)

774

775 if a1.dim() == 2:

@@ -1349,7 +1349,7 @@ class FusedMoEKernelModularImpl:

1349 else:

1350 output = torch.empty_like(hidden_states)

1351

1352 - local_num_experts = w1.size(0)

1352 + local_num_experts = w1.shape[0]

1353 if global_num_experts == -1:

1354 global_num_experts = local_num_experts

1355

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

+12 -2

@@ -212,7 +212,11 @@ def select_mxfp4_moe_backend(

212 # LoRA: separate experts backend path

213 if config.is_lora_enabled:

214 if not current_platform.is_cuda():

215 - raise NotImplementedError("Mxfp4 LoRA only supported on CUDA Platform.")

215 + # ROCm: Triton mxfp4 LoRA hits GPU memory faults due to

216 + # triton_kernels.tensor.Tensor / HIP read-only page issues

217 + # during weight swizzle and LoRA forward. Needs work from

218 + # the triton_kernels/aiter side.

219 + raise NotImplementedError("Mxfp4 LoRA is currently only supported on CUDA.")

220 if envs.VLLM_MXFP4_USE_MARLIN is False and triton_kernels_supported:

221 logger.info_once("Using Triton backend for mxfp4 lora")

222 return Mxfp4MoeBackend.TRITON_UNFUSED, backend_to_kernel_cls(

@@ -775,6 +779,8 @@ def make_mxfp4_moe_quant_config(

779 w2_scale: Union[torch.Tensor, "PrecisionConfig"],

780 w1_bias: torch.Tensor | None = None,

781 w2_bias: torch.Tensor | None = None,

782 + hidden_pad: int = 0,

783 + intermediate_pad: int = 0,

784 ) -> FusedMoEQuantConfig | None:

785 """Create a FusedMoEQuantConfig for the given MXFP4 backend."""

786 if mxfp4_backend in (

@@ -796,12 +802,16 @@ def make_mxfp4_moe_quant_config(

802 Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,

803 Mxfp4MoeBackend.CK,

804 ):

799 - return mxfp4_w4a16_moe_quant_config(

805 + config = mxfp4_w4a16_moe_quant_config(

806 w1_bias=w1_bias,

807 w2_bias=w2_bias,

808 w1_scale=w1_scale,

809 w2_scale=w2_scale,

810 )

811 + if mxfp4_backend == Mxfp4MoeBackend.CK:

812 + config.hidden_pad = hidden_pad

813 + config.intermediate_pad = intermediate_pad

814 + return config

815 else:

816 return ocp_mx_moe_quant_config(

817 quant_dtype="mxfp4",

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+11 -1

@@ -292,6 +292,8 @@ def rocm_aiter_fused_experts(

292 doweight_stage1=apply_router_weight_on_input,

293 num_local_tokens=num_local_tokens,

294 output_dtype=output_dtype,

295 + hidden_pad=quant_config.hidden_pad,

296 + intermediate_pad=quant_config.intermediate_pad,

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

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

299 )

@@ -332,7 +334,15 @@ class AiterExperts(mk.FusedMoEExpertsModular):

334 (kFp8StaticChannelSym, kFp8DynamicTokenSym),

335 (kMxfp4Static, None),

336 ]

335 - return (weight_key, activation_key) in SUPPORTED_W_A

337 + if (weight_key, activation_key) not in SUPPORTED_W_A:

338 + return False

339 + # CK MXFP4 MoE kernels are only supported on gfx950.

340 + if weight_key == kMxfp4Static:

341 + from vllm.platforms.rocm import on_gfx950

342 +

343 + if not on_gfx950():

344 + return False

345 + return True

346

347 @staticmethod

348 def _supports_activation(activation: MoEActivation) -> bool:

vllm/model_executor/layers/quantization/mxfp4.py

+8

@@ -158,6 +158,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

158 intermediate_size_per_partition_after_pad

159 )

160

161 + # CK (gfx950) padding info for rocm_aiter_ops.fused_moe()

162 + self.hidden_pad = extra_weight_attrs.get("hidden_pad", 0)

163 + self.intermediate_pad = (

164 + intermediate_size_per_partition_after_pad - intermediate_size_per_partition

165 + )

166 +

167 # Fused gate_up_proj (column parallel)

168 w13_weight = torch.nn.Parameter(

169 torch.zeros(

@@ -362,6 +368,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

368 w2_scale=w2_scale,

369 w1_bias=w1_bias,

370 w2_bias=w2_bias,

371 + hidden_pad=self.hidden_pad,

372 + intermediate_pad=self.intermediate_pad,

373 )

374

375 def select_gemm_impl(

[Bugfix][ROCm][MoE] Fix mxfp4 oracle regressions from #37128 (#37787)