Apply suggestion from @ProExpertProg

Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
input not output
2026-04-29 16:03:22 -04:00 · 2026-04-29 16:02:56 -04:00 · 2026-04-29 18:15:03 +00:00
3 changed files with 7 additions and 9 deletions
@@ -97,13 +97,6 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
                f"attention backend '{attn_backend.backend.name}'"
            )

-        # TODO: remove this after finishing migration from envs to model kwargs
-        if model_name == "openai/gpt-oss-20b":
-            from .common import is_blackwell
-
-            if is_blackwell():
-                monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1")
-
        # Disable, compile cache to make sure custom passes run.
        # Otherwise, we can't verify fusion happened through the logs.
        monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
@@ -214,4 +214,5 @@ gpt_oss_20b = ModelFusionInfo(
        sequence_parallel=n_layers * 2 + 1,
        async_tp=n_layers * 2,
    ),
+    model_kwargs={"moe_backend": "flashinfer_trtllm_afp8"} if is_blackwell() else {},
 )
@@ -112,7 +112,9 @@ MoEBackend = Literal[
    "deep_gemm_mega_moe",
    "cutlass",
    "flashinfer_trtllm",
+    "flashinfer_trtllm_afp8",
    "flashinfer_cutlass",
+    "flashinfer_cutlass_afp8",
    "flashinfer_cutedsl",
    "marlin",
    "aiter",
@@ -141,8 +143,10 @@ class KernelConfig:
    - "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only)
    - "deep_gemm_mega_moe": Use DeepGEMM mega MoE kernels
    - "cutlass": Use vLLM CUTLASS kernels
-    - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels
-    - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
+    - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels (BF16 input)
+    - "flashinfer_trtllm_afp8": Use FlashInfer with TRTLLM-GEN kernels (FP8 input)
+    - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels (BF16 input)
+    - "flashinfer_cutlass_afp8": Use FlashInfer with CUTLASS kernels (FP8 input)
    - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)
    - "marlin": Use Marlin kernels (weight-only quantization)
    - "aiter": Use AMD AITer kernels (ROCm only)