[Bugfix][LoRA] Fix garbage/NaN lora_expand output on SM90 with fp32 i… (#61906)

Signed-off-by: BeichenHuang <polarishuang0.0@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>

vllm/lora/ops/triton_ops/lora_expand_op.py

@@ -15,6 +15,7 @@ from vllm.lora.ops.triton_ops.utils import (
     get_lora_op_configs,
     supports_pdl,
 )
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.torch_utils import direct_register_custom_op
 
@@ -231,6 +232,10 @@ def _lora_expand(
         torch.bfloat16,
     ]:
         CAST_TYPE = True
+        # Pipelined wgmma with a register A operand is miscompiled by
+        # ptxas < 13.3 on SM90 (triton-lang/triton#9433).
+        if current_platform.is_cuda() and current_platform.is_device_capability(90):
+            NUM_STAGES = 1
 
     # TODO (varun): This grid formulation maximizes parallelization at the
     # cost of wasteful thread block launch when only a few input tokens require