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, 01:00 UTC
[Bugfix] Add Marlin kernel in block scaled mm kernel selection. (#40105)
Signed-off-by: maral <maralbahari.98@gmail.com>
Maral committed
Apr 17, 2026 at 18:20 UTC
c0c98b8b9a392c7e8b36b68cf477e245dda48d80
1 file changed
+16
-2
vllm/model_executor/kernels/linear/__init__.py
+16
-2
@@ -186,12 +186,13 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =
186
187
# in priority/performance order (when available)
188
_POSSIBLE_FP8_BLOCK_KERNELS: dict[
189
- PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel]]
189
+ PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel | FP8ScaledMMLinearKernel]]
190
] = {
191
PlatformEnum.CUDA: [
192
FlashInferFp8DeepGEMMDynamicBlockScaledKernel,
193
DeepGemmFp8BlockScaledMMKernel,
194
CutlassFp8BlockScaledMMKernel,
195
+ MarlinFP8ScaledMMLinearKernel,
196
TritonFp8BlockScaledMMKernel,
197
],
198
PlatformEnum.ROCM: [
@@ -392,6 +393,19 @@ def init_fp8_linear_kernel(
393
scope="global",
394
)
395
396
+ # TODO make scaled_mm kernels inherit from MMLinearKernel
397
+ # only MarlinFP8ScaledMMLinearKernel is a type of FP8ScaledMMLinearKernel.
398
+ if issubclass(kernel_type, FP8ScaledMMLinearKernel):
399
+ return kernel_type(
400
+ scaled_mm_linear_kernel_config,
401
+ layer_param_names=[
402
+ "weight",
403
+ "weight_scale",
404
+ "input_scale",
405
+ "input_scale_ub",
406
+ ],
407
+ )
408
+
409
return kernel_type(
410
scaled_mm_linear_kernel_config,
411
)
@@ -399,7 +413,7 @@ def init_fp8_linear_kernel(
413
else:
414
kernel_type = choose_scaled_mm_linear_kernel(
415
config=scaled_mm_linear_kernel_config,
402
- possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[misc]
416
+ possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[arg-type]
417
force_kernel=force_kernel,
418
)
419
if module_name: