[ROCm][Bugfix] Fix Triton "out of resource: shared memory" Error In One-Shot LoRA MoE (#47209)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Micah Williamson committed
Jul 1, 2026 at 01:24 UTC
b446792306c2dd1d751b5e42a0d8ce72c1c63608
1 file changed
+9
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+9
@@ -10,6 +10,7 @@ from vllm.distributed import (
10
from vllm.platforms import current_platform
11
from vllm.triton_utils import tl, triton
12
from vllm.triton_utils.allocation import set_triton_allocator
13
+from vllm.utils.mem_utils import get_max_shared_memory_bytes
14
from vllm.utils.torch_utils import direct_register_custom_op
15
16
from .utils import supports_pdl, supports_tma
@@ -432,6 +433,14 @@ def _run_fused_moe_lora_one_shot(
433
block_n, nw, ns = 128, 8, 3
434
else:
435
block_n, nw, ns = 128, 4, 3
436
+
437
+ # Devices with max shmem size less than 68KB can't support 3-stage
438
+ # pipeline. Fall back to a 2-stage on such devices
439
+ if current_platform.is_cuda_alike():
440
+ max_shmem_bytes = 68 * 1024
441
+ if get_max_shared_memory_bytes(device.index) < max_shmem_bytes:
442
+ ns = min(ns, 2)
443
+
444
# BLOCK_K choice: for hidden-sized K (≥256, i.e. the K=hidden_size
445
# shrink input on w13) force BLOCK_K=128 -- the wider tile halves the
446
# K-loop trip count and removes the scoreboard stalls that dominated