Vendor integration#

Chip vendors integrate by adding a backend directory under dispatch/backends/vendor/:

cp -r sglang_fl/dispatch/backends/vendor/template/ \
      sglang_fl/dispatch/backends/vendor/my_chip/

The following table lists the existing vendors:

Vendor

Directory

Hardware Detection

NVIDIA CUDA

vendor/cuda/

sgl_kernel importable

Huawei Ascend

vendor/ascend/

torch_npu importable

Template

vendor/template/

Always False (reference only)

To integrate with a new vendor, you need to implement two files:

1. Backend class (my_chip.py)#

from sglang_fl.dispatch.backends import Backend


class MyChipBackend(Backend):
    _available = None

    @property
    def name(self) -> str:
        return "my_chip"

    def is_available(self) -> bool:
        if MyChipBackend._available is None:
            try:
                import my_chip_sdk
                MyChipBackend._available = my_chip_sdk.device_count() > 0
            except ImportError:
                MyChipBackend._available = False
        return MyChipBackend._available

    def silu_and_mul(self, obj, x):
        from .impl.activation import silu_and_mul_my_chip
        return silu_and_mul_my_chip(obj, x)

    def rms_norm(self, obj, x, residual=None):
        from .impl.normalization import rms_norm_my_chip
        return rms_norm_my_chip(obj, x, residual)

    def rotary_embedding(self, obj, query, key, cos, sin, position_ids,
                         rotary_interleaved=False, inplace=True):
        from .impl.rotary import rotary_embedding_my_chip
        return rotary_embedding_my_chip(
            obj, query, key, cos, sin, position_ids, rotary_interleaved, inplace
        )

2. Registration (register_ops.py)#

import functools
from sglang_fl.dispatch.types import OpImpl, BackendImplKind, BackendPriority


def _bind_is_available(fn, is_available_fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        return fn(*args, **kwargs)
    wrapper._is_available = is_available_fn
    return wrapper


def register_builtins(registry) -> None:
    from .my_chip import MyChipBackend

    backend = MyChipBackend()
    is_avail = backend.is_available

    impls = [
        OpImpl(
            op_name="silu_and_mul",
            impl_id="vendor.my_chip",
            kind=BackendImplKind.VENDOR,
            fn=_bind_is_available(backend.silu_and_mul, is_avail),
            vendor="my_chip",
            priority=BackendPriority.VENDOR,
        ),
        OpImpl(
            op_name="rms_norm",
            impl_id="vendor.my_chip",
            kind=BackendImplKind.VENDOR,
            fn=_bind_is_available(backend.rms_norm, is_avail),
            vendor="my_chip",
            priority=BackendPriority.VENDOR,
        ),
        OpImpl(
            op_name="rotary_embedding",
            impl_id="vendor.my_chip",
            kind=BackendImplKind.VENDOR,
            fn=_bind_is_available(backend.rotary_embedding, is_avail),
            vendor="my_chip",
            priority=BackendPriority.VENDOR,
        ),
    ]
    registry.register_many(impls)

3. Operator implementations (impl/)#

Each op function receives standardized arguments (same as vllm-plugin-FL):

Op

Signature

silu_and_mul

fn(obj, x: Tensor) -> Tensor

rms_norm

fn(obj, x: Tensor, residual: Optional[Tensor] = None) -> Tensor | tuple[Tensor, Tensor]

rotary_embedding

fn(obj, query, key, cos, sin, position_ids, rotary_interleaved=False, inplace=True) -> tuple[Tensor, Tensor]

The obj parameter provides access to layer attributes (obj.weight, obj.variance_epsilon, etc.). These attribute names are identical between SGLang and vLLM, so the same impl works for both frameworks.

Vendor’s backend auto-discovery#

The plugin scans dispatch/backends/vendor/*/register_ops.py at startup. If is_available() returns True, the vendor’s ops are registered. No other files need modification.