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, 18:15 UTC
[Kernel] Enable TritonW4A16LinearKernel as CUDA fallback for non-Marlin-aligned W4A16 shapes (#43731)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Luciano Martins committed
May 27, 2026 at 07:36 UTC
22720624714d3299d2df87e0c2c67639202c4b8a
2 files changed
+3
-2
vllm/model_executor/kernels/linear/__init__.py
+1
@@ -336,6 +336,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
336
MarlinLinearKernel,
337
ConchLinearKernel,
338
ExllamaLinearKernel,
339
+ TritonW4A16LinearKernel,
340
],
341
PlatformEnum.ROCM: [
342
TritonW4A16LinearKernel,
vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py
+2
-2
@@ -288,8 +288,8 @@ class TritonW4A16LinearKernel(MPLinearKernel):
288
289
@classmethod
290
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
291
- if not current_platform.is_rocm():
292
- return False, "TritonW4A16LinearKernel only targets ROCm"
291
+ if not (current_platform.is_rocm() or current_platform.is_cuda()):
292
+ return False, "TritonW4A16LinearKernel requires CUDA or ROCm"
293
294
if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
295
return (