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 2, 19:30 UTC
[XPU] Fix fused MoE LoRA kernel crash on XPU by using platform-agnos num_compute_units (#43646)
Signed-off-by: Chaojun,Zhang <chaojun.zhang@intel.com>
Chaojun Zhang committed
May 26, 2026 at 18:40 UTC
861b97765d6cd2587ba74e7787a89d9b8aa1c9d4
1 file changed
+3
-2
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+3
-2
@@ -7,6 +7,7 @@ from vllm.distributed import (
7
tensor_model_parallel_all_gather,
8
tensor_model_parallel_all_reduce,
9
)
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.torch_utils import direct_register_custom_op
@@ -406,7 +407,7 @@ def _run_fused_moe_lora_one_shot(
407
408
# NPID_FACTOR heuristic: scale N-axis parallelism when base CTA count is
409
# short of saturating the SM array. Cap by the cost of redundant shrink.
409
- sm_count = torch.cuda.get_device_properties(device).multi_processor_count
410
+ sm_count = current_platform.num_compute_units(device.index)
411
base_programs = max(M_blocks * num_slices * grid_lora_dim, 1)
412
shrink_ratio = K / max(K + N_per_slice, 1)
413
max_npid_by_budget = max(1, int(1.5 / max(shrink_ratio, 1e-3)) + 1)
@@ -786,7 +787,7 @@ def _run_fused_moe_lora_small_batch(
787
N_tiles = triton.cdiv(N_per_slice, BLOCK_N)
788
pair_slices = M_grid * num_slices
789
789
- sm_count = torch.cuda.get_device_properties(device).multi_processor_count
790
+ sm_count = current_platform.num_compute_units(device.index)
791
n_tiles_per_program = _pick_small_batch_chunk(pair_slices, N_tiles, sm_count)
792
n_chunks = triton.cdiv(N_tiles, n_tiles_per_program)
793
work_total = pair_slices * n_chunks