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 8, 01:15 UTC
[Model Bash]: Improve FP8 Oracle for Config Specific Kernel Selection (#34260)
Signed-off-by: Elizabeth Thomas <email2eliza@gmail.com> Signed-off-by: Robert Shaw <robertgshaw2-redhat@h100-02.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Robert Shaw <robertgshaw2-redhat@h100-02.nemg-001.lab.rdu2.dc.redhat.com> Co-authored-by: Robert Shaw <robertgshaw2@gmail.com>
Elizabeth Thomas committed
Feb 19, 2026 at 23:29 UTC
81bfc21a6ad0cb498dbe5466ccf2987624efbba5
1 file changed
+46
-13
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+46
-13
@@ -34,6 +34,8 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
34
)
35
from vllm.model_executor.layers.quantization.utils.quant_utils import (
36
QuantKey,
37
+ kFp8Dynamic128Sym,
38
+ kFp8Static128BlockSym,
39
)
40
from vllm.platforms import current_platform
41
@@ -55,6 +57,49 @@ class Fp8MoeBackend(Enum):
57
XPU = "XPU"
58
59
60
+def _get_priority_backends(
61
+ moe_config: FusedMoEConfig,
62
+ weight_key: QuantKey | None,
63
+ activation_key: QuantKey | None,
64
+) -> list[Fp8MoeBackend]:
65
+ """
66
+ Get available backends in priority order based on platform and config.
67
+
68
+ This function can be extended to become more complex as needed.
69
+ """
70
+
71
+ _AVAILABLE_BACKENDS = [
72
+ Fp8MoeBackend.AITER,
73
+ Fp8MoeBackend.FLASHINFER_TRTLLM,
74
+ Fp8MoeBackend.FLASHINFER_CUTLASS,
75
+ Fp8MoeBackend.DEEPGEMM,
76
+ Fp8MoeBackend.VLLM_CUTLASS,
77
+ Fp8MoeBackend.TRITON,
78
+ Fp8MoeBackend.MARLIN,
79
+ Fp8MoeBackend.BATCHED_DEEPGEMM,
80
+ Fp8MoeBackend.BATCHED_VLLM_CUTLASS,
81
+ Fp8MoeBackend.BATCHED_TRITON,
82
+ Fp8MoeBackend.XPU,
83
+ ]
84
+
85
+ def _move_to_front(backends: list[Fp8MoeBackend], backend: Fp8MoeBackend) -> None:
86
+ backends.insert(0, backends.pop(backends.index(backend)))
87
+
88
+ # On Hopper for Block Fp8, prefer Triton for TP and FI CUTLASS for EP.
89
+ if (
90
+ current_platform.is_cuda()
91
+ and current_platform.is_device_capability(90)
92
+ and activation_key == kFp8Dynamic128Sym
93
+ and weight_key == kFp8Static128BlockSym
94
+ ):
95
+ if moe_config.moe_parallel_config.ep_size > 1:
96
+ _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.FLASHINFER_CUTLASS)
97
+ else:
98
+ _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.TRITON)
99
+
100
+ return _AVAILABLE_BACKENDS
101
+
102
+
103
def backend_to_kernel_cls(
104
backend: Fp8MoeBackend,
105
) -> type[mk.FusedMoEPermuteExpertsUnpermute]:
@@ -151,19 +196,7 @@ def select_fp8_moe_backend(
196
return Fp8MoeBackend.TRITON, backend_to_kernel_cls(Fp8MoeBackend.TRITON)
197
198
# NOTE: the kernels are selected in the following order.
154
- AVAILABLE_BACKENDS = [
155
- Fp8MoeBackend.AITER,
156
- Fp8MoeBackend.FLASHINFER_TRTLLM,
157
- Fp8MoeBackend.FLASHINFER_CUTLASS,
158
- Fp8MoeBackend.DEEPGEMM,
159
- Fp8MoeBackend.BATCHED_DEEPGEMM,
160
- Fp8MoeBackend.VLLM_CUTLASS,
161
- Fp8MoeBackend.BATCHED_VLLM_CUTLASS,
162
- Fp8MoeBackend.TRITON,
163
- Fp8MoeBackend.BATCHED_TRITON,
164
- Fp8MoeBackend.MARLIN,
165
- Fp8MoeBackend.XPU,
166
- ]
199
+ AVAILABLE_BACKENDS = _get_priority_backends(config, weight_key, activation_key)
200
201
# NOTE(rob): We need to peak into the P/F selection to determine
202
# if we are using the batched or standard expert format, which