[CI] Fix failing CUDA graph capture in Triton MOE (#46735)
Signed-off-by: Felix Marty <Felix.Marty@amd.com>
fxmarty-amd committed
Jun 26, 2026 at 16:21 UTC
37ce34922f7f5e58241369511130cd99c1c50bfe
2 files changed
+6
-1
vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py
+5
@@ -372,6 +372,11 @@ class Nvfp4QuantizationEmulationTritonExperts(TritonExperts):
372
def expects_unquantized_inputs(self) -> bool:
373
return True
374
375
+ @property
376
+ def a1_scale(self) -> torch.Tensor:
377
+ # Used in experts/triton_moe.py and passed to moe_kernel_quantize_input.
378
+ return self.a1_gscale
379
+
380
@staticmethod
381
def supports_lora() -> bool:
382
return False
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+1
-1
@@ -245,7 +245,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
245
lora_unquantized_hidden_states = hidden_states
246
hidden_states, a1q_scale = moe_kernel_quantize_input(
247
hidden_states,
248
- self.a1_scale or self.a1_gscale,
248
+ self.a1_scale,
249
self.quant_dtype,
250
self.per_act_token_quant,
251
self.block_shape,