diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index b1e0f1ebb..95355f730 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -234,7 +234,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--lora-backend` | Choose the kernel backend for multi-LoRA serving. | `triton` | `triton`, `csgmv` | | `--max-lora-chunk-size` | Maximum chunk size for the ChunkedSGMV LoRA backend. Only used when `--lora-backend` is `csgmv`. Larger values may improve performance. | `16` | `16`, `32`, `64`, `128` | -## Kernel backend +## Kernel Backends (Attention, Sampling, Grammar, GEMM) | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--attention-backend` | Choose the kernels for attention layers. | `None` | `triton`, `torch_native`, `flex_attention`, `nsa`, `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` | @@ -245,6 +245,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--mm-attention-backend` | Set multimodal attention backend. | `None` | `sdpa`, `fa3`, `triton_attn`, `ascend_attn`, `aiter_attn` | | `--nsa-prefill` | Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention). | `flashmla_sparse` | `flashmla_sparse`, `flashmla_decode`, `fa3`, `tilelang`, `aiter` | | `--nsa-decode` | Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding. | `flashmla_kv` | `flashmla_prefill`, `flashmla_kv`, `fa3`, `tilelang`, `aiter` | +| `--fp8-gemm-backend` | Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (optimal for Blackwell and low-latency), 'cutlass' (optimal for Hopper/Blackwell GPUs and high-throughput), 'triton' (fallback, widely compatible), 'aiter' (ROCm only). **NOTE**: This replaces the deprecated environment variables SGLANG_ENABLE_FLASHINFER_FP8_GEMM and SGLANG_SUPPORT_CUTLASS_BLOCK_FP8. | `auto` | `auto`, `deep_gemm`, `flashinfer_trtllm`, `cutlass`, `triton`, `aiter` | ## Speculative decoding | Argument | Description | Defaults | Options | diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 67e3f9295..8d450a48f 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -42,7 +42,7 @@ SGLang supports various environment variables that can be used to configure its | Environment Variable | Description | Default Value | | --- | --- | --- | -| `SGLANG_ENABLE_JIT_DEEPGEMM` | Enable Just-In-Time compilation of DeepGEMM kernels | `"true"` | +| `SGLANG_ENABLE_JIT_DEEPGEMM` | Enable Just-In-Time compilation of DeepGEMM kernels (enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) GPUs when the DeepGEMM package is installed; set to `"0"` to disable) | `"true"` | | `SGLANG_JIT_DEEPGEMM_PRECOMPILE` | Enable precompilation of DeepGEMM kernels | `"true"` | | `SGLANG_JIT_DEEPGEMM_COMPILE_WORKERS` | Number of workers for parallel DeepGEMM kernel compilation | `4` | | `SGL_IN_DEEPGEMM_PRECOMPILE_STAGE` | Indicator flag used during the DeepGEMM precompile script | `"false"` | @@ -78,9 +78,9 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_INT4_WEIGHT` | Enable INT4 weight quantization | `false` | | `SGLANG_MOE_PADDING` | Enable MoE padding (sets padding size to 128 if value is `1`, often set to `1` in Docker builds) | `0` | | `SGLANG_FORCE_FP8_MARLIN` | Force using FP8 MARLIN kernels even if other FP8 kernels are available | `false` | -| `SGLANG_ENABLE_FLASHINFER_FP8_GEMM` | Use flashinfer kernels when running blockwise fp8 GEMM on Blackwell GPUs | `false` | +| `SGLANG_ENABLE_FLASHINFER_FP8_GEMM` (deprecated) | Use flashinfer kernels when running blockwise fp8 GEMM on Blackwell GPUs. **DEPRECATED**: Please use `--fp8-gemm-backend=flashinfer_trtllm` instead. | `false` | | `SGLANG_FLASHINFER_FP4_GEMM_BACKEND` | Select backend for `mm_fp4` on Blackwell GPUS | `` | -| `SGLANG_SUPPORT_CUTLASS_BLOCK_FP8` | Use Cutlass kernels when running blockwise fp8 GEMM on Hopper or Blackwell GPUs | `false` | +| `SGLANG_SUPPORT_CUTLASS_BLOCK_FP8` (deprecated) | Use Cutlass kernels when running blockwise fp8 GEMM on Hopper or Blackwell GPUs. **DEPRECATED**: Please use `--fp8-gemm-backend=cutlass` instead. | `false` | | `SGLANG_CUTLASS_MOE` (deprecated) | Use Cutlass FP8 MoE kernel on Blackwell GPUs (deprecated, use --moe-runner-backend=cutlass) | `false` | diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index 25b16d310..2a5164299 100644 --- a/python/sglang/bench_one_batch.py +++ b/python/sglang/bench_one_batch.py @@ -67,6 +67,7 @@ from sglang.srt.configs.model_config import ModelConfig from sglang.srt.distributed.parallel_state import destroy_distributed_environment from sglang.srt.entrypoints.engine import _set_envs_and_config from sglang.srt.layers.moe import initialize_moe_config +from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config from sglang.srt.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.managers.scheduler_dp_attn_mixin import prepare_mlp_sync_batch_raw from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -631,6 +632,7 @@ def latency_test( tp_rank, ): initialize_moe_config(server_args) + initialize_fp8_gemm_config(server_args) # Set CPU affinity if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"): diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e161926d5..3afc41ca5 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -368,6 +368,15 @@ def _print_deprecated_env(new_name: str, old_name: str): os.environ[new_name] = os.environ[old_name] +def _warn_deprecated_env_to_cli_flag(env_name: str, suggestion: str): + """Warn when a deprecated environment variable is used. + + This is for env vars that are deprecated in favor of CLI flags. + """ + if env_name in os.environ: + warnings.warn(f"Environment variable {env_name} is deprecated. {suggestion}") + + def _convert_SGL_to_SGLANG(): _print_deprecated_env("SGLANG_LOG_GC", "SGLANG_GC_LOG") _print_deprecated_env( @@ -388,6 +397,19 @@ def _convert_SGL_to_SGLANG(): _convert_SGL_to_SGLANG() +_warn_deprecated_env_to_cli_flag( + "SGLANG_ENABLE_FLASHINFER_FP8_GEMM", + "It will be completely removed in 0.5.7. Please use '--fp8-gemm-backend=flashinfer_trtllm' instead.", +) +_warn_deprecated_env_to_cli_flag( + "SGLANG_ENABLE_FLASHINFER_GEMM", + "It will be completely removed in 0.5.7. Please use '--fp8-gemm-backend=flashinfer_trtllm' instead.", +) +_warn_deprecated_env_to_cli_flag( + "SGLANG_SUPPORT_CUTLASS_BLOCK_FP8", + "It will be completely removed in 0.5.7. Please use '--fp8-gemm-backend=cutlass' instead.", +) + def example_with_exit_stack(): # Use this style of context manager in unit test diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 70914c9d3..73d72a104 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1,4 +1,8 @@ -from typing import Callable, List, Optional, Tuple +from __future__ import annotations + +import logging +from enum import Enum +from typing import TYPE_CHECKING, Callable, List, Optional, Tuple import torch @@ -6,7 +10,9 @@ from sglang.srt.environ import envs from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.quantization.fp8_kernel import sglang_per_token_group_quant_fp8 from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil -from sglang.srt.utils import ceil_div, is_blackwell_supported, offloader + +if TYPE_CHECKING: + from sglang.srt.server_args import ServerArgs try: from vllm import _custom_ops as ops @@ -29,14 +35,20 @@ from sglang.srt.layers.quantization.fp8_kernel import ( ) from sglang.srt.utils import ( ceil_align, + ceil_div, get_bool_env_var, get_cuda_version, get_device_capability, + is_blackwell_supported, is_cuda, is_flashinfer_available, is_hip, + is_sm90_supported, + offloader, ) +logger = logging.getLogger(__name__) + _is_hip = is_hip() _is_cuda = is_cuda() _is_fp8_fnuz = is_fp8_fnuz() @@ -125,42 +137,166 @@ def normalize_e4m3fn_to_e4m3fnuz( return weight, weight_scale, input_scale -# TODO(ch-wan): define these backends in --moe-runner-backend -def cutlass_block_fp8_supported() -> bool: - if not get_bool_env_var("SGLANG_SUPPORT_CUTLASS_BLOCK_FP8"): - return False - if _is_cuda: - major, minor = torch.cuda.get_device_capability() - sm_version = major * 10 + minor - cuda_version = tuple(map(int, torch.version.cuda.split("."))) - if cuda_version >= (12, 0) and sm_version >= 90: - return True - return False +class Fp8GemmRunnerBackend(Enum): + """Enum for FP8 GEMM runner backend selection.""" + + AUTO = "auto" + FLASHINFER = "flashinfer_trtllm" + CUTLASS = "cutlass" + DEEP_GEMM = "deep_gemm" + TRITON = "triton" + AITER = "aiter" + + def is_auto(self) -> bool: + return self == Fp8GemmRunnerBackend.AUTO + + def is_flashinfer(self) -> bool: + return self == Fp8GemmRunnerBackend.FLASHINFER + + def is_cutlass(self) -> bool: + return self == Fp8GemmRunnerBackend.CUTLASS + + def is_deep_gemm(self) -> bool: + return self == Fp8GemmRunnerBackend.DEEP_GEMM + + def is_triton(self) -> bool: + return self == Fp8GemmRunnerBackend.TRITON + + def is_aiter(self) -> bool: + return self == Fp8GemmRunnerBackend.AITER -CUTLASS_BLOCK_FP8_SUPPORTED = cutlass_block_fp8_supported() -ENABLE_FLASHINFER_FP8_GEMM = ( - envs.SGLANG_ENABLE_FLASHINFER_FP8_GEMM.get() - and is_blackwell_supported() - and is_flashinfer_available() -) -if ENABLE_FLASHINFER_FP8_GEMM: +FP8_GEMM_RUNNER_BACKEND: Fp8GemmRunnerBackend | None = None + + +def _check_cutlass_block_fp8_hardware_support() -> bool: + """Return True if CUTLASS block FP8 is supported (Hopper or newer with CUDA 12.0+).""" + return is_sm90_supported() or is_blackwell_supported() + + +if is_blackwell_supported() and is_flashinfer_available(): from flashinfer.gemm import gemm_fp8_nt_groupwise def dispatch_w8a8_block_fp8_linear() -> Callable: - if ENABLE_FLASHINFER_FP8_GEMM: + """ + Dispatch to the appropriate FP8 block linear implementation. + + This function selects the backend based on: + 1. The --fp8-gemm-backend server argument (preferred) + 2. Auto-detection based on hardware capabilities + """ + backend = get_fp8_gemm_runner_backend() + + # Handle explicit backend selection via --fp8-gemm-backend + if not backend.is_auto(): + return _dispatch_explicit_backend(backend) + + # Auto mode: Select based purely on hardware/backend availability + return _dispatch_auto_backend() + + +def _dispatch_explicit_backend(backend: Fp8GemmRunnerBackend) -> Callable: + """Dispatch based on explicitly selected backend.""" + if backend.is_flashinfer(): + if not (is_blackwell_supported() and is_flashinfer_available()): + raise RuntimeError( + "FlashInfer FP8 GEMM requested via --fp8-gemm-backend=flashinfer_trtllm, " + "but FlashInfer is not available or not supported on this hardware. " + "FlashInfer FP8 GEMM requires Blackwell GPUs and FlashInfer to be installed." + ) return flashinfer_gemm_w8a8_block_fp8_linear - elif CUTLASS_BLOCK_FP8_SUPPORTED: + + elif backend.is_cutlass(): + if not _check_cutlass_block_fp8_hardware_support(): + raise RuntimeError( + "CUTLASS block FP8 requested via --fp8-gemm-backend=cutlass, " + "but hardware does not support it. CUTLASS block FP8 requires " + "Hopper (SM90+) GPUs with CUDA 12.0+." + ) + return cutlass_w8a8_block_fp8_linear_with_fallback + + elif backend.is_aiter(): + if not _use_aiter: + raise RuntimeError( + "AITER backend requested via --fp8-gemm-backend=aiter, " + "but AITER is not available. AITER requires AMD GPUs with " + "SGLANG_USE_AITER=1 environment variable set." + ) + return aiter_w8a8_block_fp8_linear + + elif backend.is_deep_gemm(): + if not deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: + raise RuntimeError( + "DeepGEMM backend requested via --fp8-gemm-backend=deep_gemm, " + "but DeepGEMM is not available. This usually means the deep_gemm package " + "is not installed or has been disabled via SGLANG_ENABLE_JIT_DEEPGEMM=0." + ) + return deepgemm_w8a8_block_fp8_linear_with_fallback + + elif backend.is_triton(): + return triton_w8a8_block_fp8_linear + + else: + raise ValueError(f"Unknown FP8 GEMM backend: {backend}") + + +def _dispatch_auto_backend() -> Callable: + """Auto-select the best backend based on hardware capabilities.""" + # Priority order for auto selection: + # 1. DeepGEMM (if enabled and available) + # 2. FlashInfer TRTLLM (if Blackwell GPU and FlashInfer available) + # 3. CUTLASS (if Hopper+ GPU and CUDA 12.0+) + # 4. AITER (if AMD GPU with AITER enabled) + # 5. Triton (fallback) + + if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: + return deepgemm_w8a8_block_fp8_linear_with_fallback + elif is_blackwell_supported() and is_flashinfer_available(): + return flashinfer_gemm_w8a8_block_fp8_linear + elif _check_cutlass_block_fp8_hardware_support(): return cutlass_w8a8_block_fp8_linear_with_fallback elif _use_aiter: return aiter_w8a8_block_fp8_linear - elif deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: - return deepgemm_w8a8_block_fp8_linear_with_fallback else: return triton_w8a8_block_fp8_linear +def initialize_fp8_gemm_config(server_args: ServerArgs) -> None: + """Initialize FP8 GEMM configuration.""" + global FP8_GEMM_RUNNER_BACKEND + + backend = server_args.fp8_gemm_runner_backend + + # TODO(brayden): Remove env-based overrides in v0.5.7, they will be fully removed in v0.5.7. + # Only check environment variables when the server args is not set, server args should take priority. + if backend == "auto": + if envs.SGLANG_ENABLE_FLASHINFER_FP8_GEMM.get(): + backend = "flashinfer_trtllm" + elif envs.SGLANG_SUPPORT_CUTLASS_BLOCK_FP8.get(): + backend = "cutlass" + else: + if ( + envs.SGLANG_ENABLE_FLASHINFER_FP8_GEMM.get() + or envs.SGLANG_SUPPORT_CUTLASS_BLOCK_FP8.get() + ): + logger.warning( + f"FP8 GEMM backend set to '{backend}' via --fp8-gemm-backend overrides " + "environment variables SGLANG_ENABLE_FLASHINFER_FP8_GEMM and " + "SGLANG_SUPPORT_CUTLASS_BLOCK_FP8. Using server argument value." + ) + + FP8_GEMM_RUNNER_BACKEND = Fp8GemmRunnerBackend(backend) + + +def get_fp8_gemm_runner_backend() -> Fp8GemmRunnerBackend: + """Get the current FP8 GEMM runner backend.""" + global FP8_GEMM_RUNNER_BACKEND + if FP8_GEMM_RUNNER_BACKEND is None: + FP8_GEMM_RUNNER_BACKEND = Fp8GemmRunnerBackend.AUTO + return FP8_GEMM_RUNNER_BACKEND + + def flashinfer_gemm_w8a8_block_fp8_linear( input: torch.Tensor, weight: torch.Tensor, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c47d5ab65..baaf63e97 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -65,6 +65,7 @@ from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers.dp_attention import compute_dp_attention_world_info from sglang.srt.layers.moe import initialize_moe_config +from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config from sglang.srt.managers.io_struct import ( AbortReq, BaseBatchReq, @@ -305,6 +306,9 @@ class Scheduler( # Init moe config self.init_moe_config() + # Init GEMM config (FP8 GEMM, etc.) + self.init_gemm_config() + # Check whether overlap can be enabled if not self.is_generation: self.enable_overlap = False @@ -966,6 +970,12 @@ class Scheduler( if hasattr(self.model_config.hf_config, "num_experts_per_tok"): initialize_moe_config(self.server_args) + def init_gemm_config(self): + # Initialize GEMM-related configuration (currently FP8 Blockwise GEMM backend). + # Other GEMM backends (e.g. FP4, BF16, etc.) can be added here in the future. + # This is needed for FP8 quantization. + initialize_fp8_gemm_config(self.server_args) + @DynamicGradMode() def event_loop_normal(self): """A normal scheduler loop.""" diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 68bd6709c..97b5200c8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -170,6 +170,15 @@ MOE_RUNNER_BACKEND_CHOICES = [ MOE_A2A_BACKEND_CHOICES = ["none", "deepep", "mooncake", "ascend_fuseep"] +FP8_GEMM_RUNNER_BACKEND_CHOICES = [ + "auto", + "deep_gemm", + "flashinfer_trtllm", + "cutlass", + "triton", + "aiter", +] + MAMBA_SSM_DTYPE_CHOICES = ["float32", "bfloat16"] @@ -198,6 +207,10 @@ def add_moe_runner_backend_choices(choices): MOE_RUNNER_BACKEND_CHOICES.extend(choices) +def add_fp8_gemm_runner_backend_choices(choices): + FP8_GEMM_RUNNER_BACKEND_CHOICES.extend(choices) + + def add_deterministic_attention_backend_choices(choices): DETERMINISTIC_ATTENTION_BACKEND_CHOICES.extend(choices) @@ -380,6 +393,7 @@ class ServerArgs: sampling_backend: Optional[str] = None grammar_backend: Optional[str] = None mm_attention_backend: Optional[str] = None + fp8_gemm_runner_backend: str = "auto" nsa_prefill_backend: str = "flashmla_sparse" nsa_decode_backend: str = "fa3" enable_flashinfer_autotune: bool = False @@ -2952,6 +2966,22 @@ class ServerArgs: type=str, choices=NSA_CHOICES, ) + parser.add_argument( + "--fp8-gemm-backend", + type=str, + choices=FP8_GEMM_RUNNER_BACKEND_CHOICES, + default=ServerArgs.fp8_gemm_runner_backend, + dest="fp8_gemm_runner_backend", + help="Choose the runner backend for Blockwise FP8 GEMM operations. " + "Options: 'auto' (default, auto-selects based on hardware), " + "'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), " + "'flashinfer_trtllm' (optimal for Blackwell and low-latency), " + "'cutlass' (optimal for Hopper/Blackwell GPUs and high-throughput), " + "'triton' (fallback, widely compatible), " + "'aiter' (ROCm only). " + "NOTE: This replaces the deprecated environment variables " + "SGLANG_ENABLE_FLASHINFER_FP8_GEMM and SGLANG_SUPPORT_CUTLASS_BLOCK_FP8.", + ) parser.add_argument( "--enable-flashinfer-autotune", default=ServerArgs.enable_flashinfer_autotune, diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index cf7264225..fd9a61ef2 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -227,9 +227,7 @@ def is_blackwell(): def is_blackwell_supported(device=None) -> bool: if not is_cuda_alike(): return False - return (torch.cuda.get_device_capability(device)[0] in [10, 12]) and ( - torch.version.cuda >= "12.8" - ) + return is_sm100_supported(device) or is_sm120_supported(device) @lru_cache(maxsize=1) @@ -243,7 +241,7 @@ def is_sm120_supported(device=None) -> bool: @lru_cache(maxsize=1) def is_sm100_supported(device=None) -> bool: - if not is_cuda_alike(): + if not is_cuda(): return False return (torch.cuda.get_device_capability(device)[0] == 10) and ( torch.version.cuda >= "12.8" @@ -252,7 +250,7 @@ def is_sm100_supported(device=None) -> bool: @lru_cache(maxsize=1) def is_sm90_supported(device=None) -> bool: - if not is_cuda_alike(): + if not is_cuda(): return False return (torch.cuda.get_device_capability(device)[0] == 9) and ( torch.version.cuda >= "12.3" diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index b7b38bf2b..278135c7c 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -170,6 +170,7 @@ suites = { "per-commit-4-gpu-b200": [ TestFile("test_deepseek_v3_fp4_4gpu.py", 1800), TestFile("test_flash_attention_4.py", 300), + TestFile("test_fp8_blockwise_gemm.py", 280), TestFile("test_gpt_oss_4gpu.py", 600), TestFile("test_llama31_fp4.py", 300), TestFile("test_eagle_infer_beta_dp_attention.py", 300), diff --git a/test/srt/test_fp8_blockwise_gemm.py b/test/srt/test_fp8_blockwise_gemm.py new file mode 100644 index 000000000..cea9b098f --- /dev/null +++ b/test/srt/test_fp8_blockwise_gemm.py @@ -0,0 +1,73 @@ +import unittest +from types import SimpleNamespace +from urllib.parse import urlparse + +from sglang.srt.utils import get_device_sm, kill_process_tree +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + popen_launch_server, + try_cached_model, +) + +MODEL_PATH = "Qwen/Qwen3-4B-Instruct-2507-FP8" + + +class FP8BlockwiseGemmBase: + backend = None + + @classmethod + def setUpClass(cls): + if cls.backend is None: + raise NotImplementedError("Subclass must set 'backend' attribute") + cls.model = try_cached_model(MODEL_PATH) + cls.base_url = DEFAULT_URL_FOR_TEST + other_args = [ + "--trust-remote-code", + "--fp8-gemm-backend", + cls.backend, + ] + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + parsed_url = urlparse(self.base_url) + args = SimpleNamespace( + num_shots=8, + data_path=None, + num_questions=1319, + max_new_tokens=512, + parallel=200, + host=f"{parsed_url.scheme}://{parsed_url.hostname}", + port=parsed_url.port, + ) + metrics = run_eval_few_shot_gsm8k(args) + print(metrics) + + self.assertGreaterEqual(metrics["accuracy"], 0.41) + + +class TestFP8BlockwiseGemmTriton(FP8BlockwiseGemmBase, unittest.TestCase): + backend = "triton" + + +class TestFP8BlockwiseGemmDeepGemm(FP8BlockwiseGemmBase, unittest.TestCase): + backend = "deep_gemm" + + +@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher") +class TestFP8BlockwiseGemmFlashinferTrtllm(FP8BlockwiseGemmBase, unittest.TestCase): + backend = "flashinfer_trtllm" + + +if __name__ == "__main__": + unittest.main()