Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 63 additions & 0 deletions tests/unit_tests/dispatch/test_config_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
# Copyright (c) 2026 BAAI. All rights reserved.

from types import SimpleNamespace
from unittest.mock import patch

from vllm_fl.dispatch.config.utils import get_config_path, load_platform_config


def _platform_with_capability(major: int, minor: int = 0):
return SimpleNamespace(
get_device_capability=lambda: SimpleNamespace(major=major, minor=minor)
)


def test_hopper_optimization_is_disabled_by_default(monkeypatch):
monkeypatch.delenv("VLLM_FL_HOPPER_LONG_CONTEXT_OPT", raising=False)
platform = _platform_with_capability(9)
with patch("vllm.platforms.current_platform", platform):
path = get_config_path("nvidia")
config = load_platform_config("nvidia")

assert path is not None
assert path.name == "nvidia.yaml"
assert config["op_backends"]["attention_backend"][0] == "flagos"
assert "mm" not in config["flagos_blacklist"]


def test_hopper_uses_architecture_specific_config_when_enabled(monkeypatch):
monkeypatch.setenv("VLLM_FL_HOPPER_LONG_CONTEXT_OPT", "1")
platform = _platform_with_capability(9)
with patch("vllm.platforms.current_platform", platform):
path = get_config_path("nvidia")
config = load_platform_config("nvidia")

assert path is not None
assert path.name == "nvidia_hopper.yaml"
assert config["op_backends"]["attention_backend"][0] == "vendor"
assert {"mm", "mm_out"}.issubset(config["flagos_blacklist"])


def test_non_hopper_nvidia_keeps_vendor_wide_config(monkeypatch):
monkeypatch.setenv("VLLM_FL_HOPPER_LONG_CONTEXT_OPT", "1")
platform = _platform_with_capability(8)
with patch("vllm.platforms.current_platform", platform):
path = get_config_path("nvidia")
config = load_platform_config("nvidia")

assert path is not None
assert path.name == "nvidia.yaml"
assert config["op_backends"]["attention_backend"][0] == "flagos"
assert "mm" not in config["flagos_blacklist"]


def test_capability_probe_failure_falls_back_to_vendor_config(monkeypatch):
monkeypatch.setenv("VLLM_FL_HOPPER_LONG_CONTEXT_OPT", "true")
platform = SimpleNamespace(
get_device_capability=lambda: (_ for _ in ()).throw(RuntimeError("no GPU"))
)
with patch("vllm.platforms.current_platform", platform):
path = get_config_path("nvidia")

assert path is not None
assert path.name == "nvidia.yaml"
4 changes: 4 additions & 0 deletions vllm_fl/dispatch/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -323,6 +323,7 @@ Environment variables can override specific items from platform config. If not s
| `VLLM_FL_PREFER` | `flagos` | Preferred backend: `flagos`, `vendor`, `reference` |
| `VLLM_FL_STRICT` | `0` | Strict mode: `1` = fail on error, `0` = try fallback |
| `VLLM_FL_PER_OP` | (none) | Per-operator order: `op1=a\|b\|c;op2=x\|y` |
| `VLLM_FL_HOPPER_LONG_CONTEXT_OPT` | `0` | On NVIDIA Hopper, route attention to FA3 and `mm`/`mm_out` to native CUDA |
| `VLLM_FL_ALLOW_VENDORS` | (none) | Vendor whitelist, comma-separated |
| `VLLM_FL_DENY_VENDORS` | (none) | Vendor blacklist, comma-separated |

Expand Down Expand Up @@ -385,6 +386,9 @@ export VLLM_FL_FLAGOS_WHITELIST="silu_and_mul,rms_norm"
# Specify per-operator order
export VLLM_FL_PER_OP="rms_norm=vendor|flagos|reference"

# Opt in to the NVIDIA Hopper long-context routing for an A/B run
export VLLM_FL_HOPPER_LONG_CONTEXT_OPT=1

# Use completely custom config file
export VLLM_FL_CONFIG=/path/to/my_config.yaml

Expand Down
43 changes: 43 additions & 0 deletions vllm_fl/dispatch/config/nvidia_hopper.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
# vLLM-FL Dispatch Configuration for NVIDIA Hopper GPUs

# Keep FlagOS as the default implementation and override only the operators
# whose Hopper paths have a measured regression against the CUDA vendor stack.
# This config is opt-in via VLLM_FL_HOPPER_LONG_CONTEXT_OPT=1.
prefer: flagos
strict: false

allow_vendors:
- cuda

op_backends:
# FlashAttention-3 is substantially faster than Triton unified attention
# for long-context full-attention prefill on Hopper.
attention_backend:
- vendor
- flagos
- reference
rms_norm:
- flagos
- vendor
- reference
silu_and_mul:
- flagos
- vendor
- reference
rotary_embedding:
- flagos
- vendor
- reference

flagos_blacklist:
- index_put_
- index_put
- _index_put_impl_
- nonzero
- copy_
- to_copy
- index
# FlagGems' Hopper host-TMA mm.out path is slower than the native dense
# GEMM for the long-context projection shapes used by vLLM.
- mm
- mm_out
30 changes: 30 additions & 0 deletions vllm_fl/dispatch/config/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,30 @@
_CONFIG_DIR = Path(__file__).parent


def _get_arch_config_name(platform: str) -> Optional[str]:
"""Return an architecture-specific config name when one is available."""
if platform != "nvidia":
return None
if os.environ.get("VLLM_FL_HOPPER_LONG_CONTEXT_OPT", "0").lower() not in (
"1",
"true",
):
return None

try:
from vllm.platforms import current_platform

capability = current_platform.get_device_capability()
except Exception:
# Device discovery is not guaranteed to be available in API-only or
# model-inspection processes. Fall back to the vendor-wide defaults.
return None

if capability is not None and capability.major == 9:
return "nvidia_hopper"
return None


def get_platform_name() -> str:
"""
Detect the current hardware platform.
Expand Down Expand Up @@ -75,6 +99,12 @@ def get_config_path(platform: Optional[str] = None) -> Optional[Path]:
if platform is None:
platform = get_platform_name()

arch_config = _get_arch_config_name(platform)
if arch_config is not None:
config_file = _CONFIG_DIR / f"{arch_config}.yaml"
if config_file.exists():
return config_file

# Try platform-specific config
config_file = _CONFIG_DIR / f"{platform}.yaml"
if config_file.exists():
Expand Down
Loading