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 5, 11:00 UTC
[XPU] Quick fix for TritonMLA to remove cuda hardcode (#39088)
Signed-off-by: Chendi Xue <chendi.xue@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Chendi.Xue committed
Apr 7, 2026 at 11:17 UTC
92b9afeecde83bb0e0fef8ec8e30e7dae6b43bdc
2 files changed
+3
-2
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+1
-1
@@ -222,7 +222,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
222
self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
223
else:
224
self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
225
- elif current_platform.is_xpu():
225
+ elif self.unquantized_backend == UnquantizedMoeBackend.XPU:
226
w13 = layer.w13_weight
227
w2 = layer.w2_weight
228
vllm/v1/attention/backends/mla/triton_mla.py
+2
-1
@@ -13,6 +13,7 @@ from vllm.model_executor.layers.attention.mla_attention import (
13
MLACommonImpl,
14
MLACommonMetadata,
15
)
16
+from vllm.platforms import current_platform
17
from vllm.platforms.interface import DeviceCapability
18
from vllm.triton_utils import triton
19
from vllm.utils.torch_utils import is_quantized_kv_cache
@@ -116,7 +117,7 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
117
if is_quantized_kv_cache(self.kv_cache_dtype):
118
self.supports_quant_query_input = False
119
119
- self._sm_count = torch.cuda.get_device_properties(0).multi_processor_count
120
+ self._sm_count = current_platform.num_compute_units()
121
122
def _flash_attn_varlen_diff_headdims(
123
self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs