[diffusion] optimization: INT8 Linear + pluggable DiT attention backends for MiniMax-H3 on consumer-level GPUs (#34581)

python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import inspect
 import re
 
 import torch
@@ -19,6 +20,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)
 
 _SOL_ATTN_HEAD_DIM = 128
+_DENSE_BACKENDS = {"fa", "sage_attn"}
 
 
 def _parse_layer_ranges(spec: str | int | None) -> frozenset[int]:
@@ -57,7 +59,16 @@ def _resolve_kv_splits(q: torch.Tensor, kv_splits: int | str | None) -> int:
 def _get_sol_attn_runtime_config() -> dict:
     server_args = get_global_server_args()
     cfg = getattr(server_args, "attention_backend_config", None) or {}
-    dense_layers = cfg.get("dense_layers", "0,1")
+    dense_backend = (
+        str(cfg.get("dense_backend", "fa")).strip().lower().replace("-", "_")
+    )
+    if dense_backend in {"sage", "sageattention"}:
+        dense_backend = "sage_attn"
+    if dense_backend not in _DENSE_BACKENDS:
+        raise ValueError(
+            f"Unsupported sol_attn dense_backend={dense_backend!r}; "
+            f"expected one of {sorted(_DENSE_BACKENDS)}"
+        )
     sink_start = cfg.get("sink_start", 0)
     return {
         "tau": float(cfg.get("tau", 1.0)),
@@ -66,7 +77,8 @@ def _get_sol_attn_runtime_config() -> dict:
         "sink_tokens": int(cfg.get("sink_tokens", 0)),
         "sink_start": None if sink_start is None else int(sink_start),
         "dense_steps": int(cfg.get("dense_steps", 10)),
-        "dense_layers": _parse_layer_ranges(dense_layers),
+        "dense_layers": _parse_layer_ranges(cfg.get("dense_layers", "0,1")),
+        "dense_backend": dense_backend,
     }
 
 
@@ -107,13 +119,12 @@ class SolAttnImpl(AttentionImpl):
         self.softmax_scale = softmax_scale
         self.prefix = prefix
         self.layer_idx = self._parse_layer_idx(prefix)
+        self._sol_params: frozenset[str] | None = None
 
     @staticmethod
     def _parse_layer_idx(prefix: str) -> int | None:
         match = re.search(r"blocks\.(\d+)", prefix)
-        if match is None:
-            return None
-        return int(match.group(1))
+        return int(match.group(1)) if match else None
 
     def _should_use_dense(self) -> bool:
         cfg = _get_sol_attn_runtime_config()
@@ -125,13 +136,11 @@ class SolAttnImpl(AttentionImpl):
             step = int(get_forward_context().current_timestep)
         except AssertionError:
             step = 0
-        if step < cfg["dense_steps"]:
-            return True
-        if self.layer_idx is not None and self.layer_idx in cfg["dense_layers"]:
-            return True
-        return False
+        return step < cfg["dense_steps"] or (
+            self.layer_idx is not None and self.layer_idx in cfg["dense_layers"]
+        )
 
-    def _dense_varlen(
+    def _dense_fa(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
@@ -153,6 +162,46 @@ class SolAttnImpl(AttentionImpl):
         )
         return output[0] if isinstance(output, tuple) else output
 
+    def _dense_sage(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        *,
+        cu_seqlens: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """SageAttention dense path.
+
+        - batched NHD ``[B, T, H, D]`` when ``cu_seqlens is None``
+        - packed ``[total, H, D]`` when ``cu_seqlens`` is provided
+        """
+        from sageattention import sageattn
+
+        if cu_seqlens is None:
+            return sageattn(
+                query.contiguous(),
+                key.contiguous(),
+                value.contiguous(),
+                tensor_layout="NHD",
+                is_causal=self.causal,
+                sm_scale=self.softmax_scale,
+            )
+
+        bounds = [int(x) for x in cu_seqlens.tolist()]
+        output = torch.empty_like(query)
+        for start, stop in zip(bounds[:-1], bounds[1:]):
+            if start == stop:
+                continue
+            output[start:stop] = sageattn(
+                query[start:stop].unsqueeze(0).contiguous(),
+                key[start:stop].unsqueeze(0).contiguous(),
+                value[start:stop].unsqueeze(0).contiguous(),
+                tensor_layout="NHD",
+                is_causal=self.causal,
+                sm_scale=self.softmax_scale,
+            )[0]
+        return output
+
     def _run_sol_attn_thd(
         self,
         query: torch.Tensor,
@@ -167,17 +216,24 @@ class SolAttnImpl(AttentionImpl):
         v = value.unsqueeze(0).contiguous()
         if q.dtype != torch.bfloat16:
             raise TypeError(f"Sol-Attn requires bfloat16 activations, got {q.dtype}")
-        out = sol_attn(
-            q,
-            k,
-            v,
-            tau=cfg["tau"],
-            thresh_type=cfg["thresh_type"],
-            kv_splits=_resolve_kv_splits(q, cfg["kv_splits"]),
-            sink_start=cfg["sink_start"],
-            sink_tokens=cfg["sink_tokens"],
-        )
-        return out.squeeze(0)
+
+        if self._sol_params is None:
+            self._sol_params = frozenset(inspect.signature(sol_attn).parameters)
+
+        kwargs = {
+            "tau": cfg["tau"],
+            "thresh_type": cfg["thresh_type"],
+            "kv_splits": _resolve_kv_splits(q, cfg["kv_splits"]),
+            "sink_start": cfg["sink_start"],
+            "sink_tokens": cfg["sink_tokens"],
+        }
+        # Wan2GP Ada port: INT8-QK Triton; official NVlabs API has no int8_qk.
+        if "int8_qk" in self._sol_params and tuple(
+            torch.cuda.get_device_capability(q.device)
+        ) >= (8, 9):
+            kwargs["int8_qk"] = True
+        kwargs = {k: v for k, v in kwargs.items() if k in self._sol_params}
+        return sol_attn(q, k, v, **kwargs).squeeze(0)
 
     def forward(
         self,
@@ -187,14 +243,17 @@ class SolAttnImpl(AttentionImpl):
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
         del attn_metadata
+        # ``query`` is NHD: [B, T, H, D]
         if self._should_use_dense():
-            q = query.transpose(1, 2).reshape(
+            if _get_sol_attn_runtime_config()["dense_backend"] == "sage_attn":
+                return self._dense_sage(query, key, value)
+            # NHD [B, T, H, D] → packed THD [B*T, H, D] (plain reshape; do not
+            # transpose — that would scramble token order for flash_attn_varlen).
+            q = query.reshape(
                 query.shape[0] * query.shape[1], query.shape[2], query.shape[3]
             )
-            k = key.transpose(1, 2).reshape(
-                key.shape[0] * key.shape[1], key.shape[2], key.shape[3]
-            )
-            v = value.transpose(1, 2).reshape(
+            k = key.reshape(key.shape[0] * key.shape[1], key.shape[2], key.shape[3])
+            v = value.reshape(
                 value.shape[0] * value.shape[1], value.shape[2], value.shape[3]
             )
             cu_seqlens = torch.arange(
@@ -204,14 +263,11 @@ class SolAttnImpl(AttentionImpl):
                 device=query.device,
                 dtype=torch.int32,
             )
-            out = self._dense_varlen(
-                q,
-                k,
-                v,
-                cu_seqlens=cu_seqlens,
-                max_seqlen=query.shape[1],
+            out = self._dense_fa(
+                q, k, v, cu_seqlens=cu_seqlens, max_seqlen=query.shape[1]
             )
             return out.reshape(query.shape[0], query.shape[1], query.shape[2], -1)
+
         q = query.reshape(query.shape[0] * query.shape[1], query.shape[2], -1)
         k = key.reshape(key.shape[0] * key.shape[1], key.shape[2], -1)
         v = value.reshape(value.shape[0] * value.shape[1], value.shape[2], -1)
@@ -230,11 +286,9 @@ class SolAttnImpl(AttentionImpl):
     ) -> torch.Tensor:
         del cu_seqlens_host
         if self._should_use_dense():
-            return self._dense_varlen(
-                query,
-                key,
-                value,
-                cu_seqlens=cu_seqlens,
-                max_seqlen=max_seqlen,
+            if _get_sol_attn_runtime_config()["dense_backend"] == "sage_attn":
+                return self._dense_sage(query, key, value, cu_seqlens=cu_seqlens)
+            return self._dense_fa(
+                query, key, value, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen
             )
         return self._run_sol_attn_thd(query, key, value)

python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py

@@ -8,6 +8,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
 )
+from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import (
+    KitchenInt8Config,
+)
 from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config
 from sglang.multimodal_gen.runtime.layers.quantization.modelopt_fp8 import (
     ModelOptFp8Config as ModelOptFp8DiffusionConfig,
@@ -33,6 +36,7 @@ QuantizationMethods = Literal[
     "mxfp8",
     "mxfp4",
     "mxfp4_npu",
+    "kitchen_int8",
 ]
 
 QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
@@ -48,6 +52,7 @@ _CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {
     "mxfp4": Mxfp4Config,
     "mxfp8": MXFP8Config,
     "mxfp4_npu": NPUMXFP4Config,
+    "kitchen_int8": KitchenInt8Config,
 }
 
 

python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_int8_config.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Config for online INT8 ConvRot quantization via comfy_kitchen.
+
+A no-arg ``KitchenInt8Config()`` is the only supported form: weights load in
+their source dtype and are quantized in ``process_weights_after_loading``.
+
+Registered CLI name: ``kitchen_int8``.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.layers.quantization.utils import is_layer_skipped
+
+logger = init_logger(__name__)
+
+_SUPPORTED_GROUP_SIZES = (16, 64, 256)
+
+
+class KitchenInt8Config(QuantizationConfig):
+    """Config for online INT8 ConvRot quantization via comfy_kitchen.
+
+    A no-arg ``KitchenInt8Config()`` is the only supported form: weights load in
+    their source dtype and are quantized in ``process_weights_after_loading``.
+    """
+
+    def __init__(
+        self,
+        group_size: int = 256,
+        ignored_layers: list[str] | None = None,
+        packed_modules_mapping: dict[str, list[str]] | None = None,
+    ) -> None:
+        super().__init__()
+        if group_size not in _SUPPORTED_GROUP_SIZES:
+            raise ValueError(
+                f"kitchen_int8 group_size must be one of {_SUPPORTED_GROUP_SIZES}, "
+                f"got {group_size}"
+            )
+        self.group_size = group_size
+        self.ignored_layers = ignored_layers or []
+        self.packed_modules_mapping = packed_modules_mapping or {}
+        # Which layers actually got quantized is worth stating plainly in the
+        # log: a silent fallback to BF16 looks exactly like a slow kernel.
+        self.selected: list[str] = []
+        self.skipped: list[str] = []
+        self._processed = 0
+        self._quantized_bytes = 0
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "kitchen_int8"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        # INT8 tensor cores land on Turing.
+        return 75
+
+    @classmethod
+    def get_config_filenames(cls) -> list[str]:
+        return []
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> KitchenInt8Config:
+        return cls(
+            group_size=cls.get_from_keys_or(config, ["group_size"], 256),
+            ignored_layers=cls.get_from_keys_or(config, ["ignored_layers"], None),
+        )
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> QuantizeMethodBase | None:
+        from sglang.multimodal_gen.runtime.layers.linear import LinearBase
+        from sglang.multimodal_gen.runtime.layers.quantization.kitchen_int8 import (
+            KitchenInt8LinearMethod,
+        )
+
+        if not isinstance(layer, LinearBase):
+            return None
+        if is_layer_skipped(
+            prefix, self.ignored_layers, fused_mapping=self.packed_modules_mapping
+        ):
+            self.skipped.append(prefix)
+            return UnquantizedLinearMethod()
+        # The rotation partitions the input dim into fixed-size groups, so a
+        # layer whose input does not divide evenly simply stays in BF16 rather
+        # than failing the whole model. H3's adaln projections (in=2688) are
+        # the case this exists for, and they cost 0.2% of a step anyway.
+        if layer.input_size % self.group_size:
+            self.skipped.append(f"{prefix}(in={layer.input_size})")
+            return UnquantizedLinearMethod()
+        self.selected.append(prefix)
+        return KitchenInt8LinearMethod(self)
+
+    def note_quantized(self, saved_bytes: int) -> None:
+        self._processed += 1
+        self._quantized_bytes += saved_bytes
+        if self._processed == len(self.selected):
+            logger.info(
+                "kitchen_int8: quantized %d linear layers (%.2f GiB of BF16 weights "
+                "-> %.2f GiB INT8), left %d in BF16",
+                self._processed,
+                self._quantized_bytes / 1024**3,
+                self._quantized_bytes / 2 / 1024**3,
+                len(self.skipped),
+            )
+            logger.debug("kitchen_int8: layers left in BF16: %s", self.skipped)
+
+    def get_scaled_act_names(self) -> list[str]:
+        return []

python/sglang/multimodal_gen/runtime/layers/quantization/kitchen_int8.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+"""INT8 weight-only-storage linear backed by comfy_kitchen's fused ConvRot kernel.
+
+On Ada (RTX 4090) INT8 is only worth doing with the right kernel: on MiniMax H3
+shapes `torch._int_mm` measures 0.46-0.90x of BF16 (i.e. slower) and a Triton
+INT8 GEMM roughly ties BF16, while `comfy_kitchen.int8_linear` reaches 2.49x.
+The difference is that it is a single fused op -- it takes a BF16 activation and
+does the Hadamard rotation, dynamic per-row activation quantization, IMMA GEMM,
+dequantization and bias add without ever materializing the intermediates.
+
+Quantization is data-free (group-wise Hadamard rotation + per-output-channel
+absmax), so weights are quantized here after loading rather than read from a
+pre-quantized checkpoint. That keeps this usable with the stock BF16 checkpoint
+and avoids depending on any external file layout.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+from torch.nn.parameter import Parameter
+
+from sglang.multimodal_gen.runtime.layers.linear import LinearMethodBase
+from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import (
+    KitchenInt8Config,
+)
+from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
+
+__all__ = ["KitchenInt8Config", "KitchenInt8LinearMethod"]
+
+# comfy_kitchen's dtype codes for the fused op's output.
+_OUT_DTYPE_CODE = {torch.float32: 0, torch.float16: 1, torch.bfloat16: 2}
+
+# comfy_kitchen picks its CUTLASS tile configuration from a threshold tree
+# (select_fused_int8_config in cutlass_gemm_int8.cu). Shapes whose N falls under
+# its 24832 cutoff but whose M is large get a Stream-K schedule, which exists to
+# balance load when there are too few tiles to fill the GPU. At H3's 32700 tokens
+# qkv_proj already launches ~21k CTAs over the 4090's 128 SMs, so Stream-K's
+# workspace and fixup reduction are pure overhead: 26.5 ms against 17.9 ms for
+# the identical tile without it. Capping rows per call keeps the plain
+# data-parallel config, and is bit-exact because splitting rows does not change
+# any single row's arithmetic.
+_MAX_ROWS_PER_CALL = int(os.environ.get("SGLANG_KITCHEN_INT8_MAX_ROWS", "8192"))
+# Narrow outputs do not recover the cost of writing results back through a
+# preallocated buffer; H3's out_proj and fc2 (N=5376) both measure slower split.
+_MIN_SPLIT_OUTPUT = int(os.environ.get("SGLANG_KITCHEN_INT8_MIN_SPLIT_N", "8192"))
+
+
+def _row_split(rows: int, out_features: int) -> int | None:
+    """Rows per `int8_linear` call, or None to issue one call for everything."""
+    if _MAX_ROWS_PER_CALL <= 0 or rows <= _MAX_ROWS_PER_CALL:
+        return None
+    if out_features < _MIN_SPLIT_OUTPUT:
+        return None
+    return _MAX_ROWS_PER_CALL
+
+
+def _load_comfy_kitchen():
+    try:
+        import comfy_kitchen  # noqa: F401
+    except ImportError as exc:  # pragma: no cover - depends on optional dep
+        raise ImportError(
+            "kitchen_int8 quantization requires the `comfy-kitchen` package "
+            "(pip install comfy-kitchen). It is a self-contained abi3 extension "
+            "and does not link against libtorch, so any torch version works."
+        ) from exc
+    if not hasattr(torch.ops.comfy_kitchen, "int8_linear"):
+        raise RuntimeError(
+            "comfy_kitchen is installed but did not register "
+            "torch.ops.comfy_kitchen.int8_linear"
+        )
+
+
+class KitchenInt8LinearMethod(LinearMethodBase):
+    """Quantizes BF16 weights to INT8 after load and runs the fused kernel."""
+
+    def __init__(self, quant_config: KitchenInt8Config) -> None:
+        self.quant_config = quant_config
+        _load_comfy_kitchen()
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ) -> None:
+        # get_quant_method already screened the unsharded input size, so this
+        # only fires under TP > 1, where a row-parallel layer splits the very
+        # dimension the rotation groups over.
+        if input_size_per_partition % self.quant_config.group_size:
+            raise ValueError(
+                f"kitchen_int8 needs input_size_per_partition "
+                f"({input_size_per_partition}) divisible by group_size "
+                f"{self.quant_config.group_size}"
+            )
+
+        # Deliberately identical to UnquantizedLinearMethod: weights load as
+        # BF16 through the model's existing loaders (H3 for instance installs a
+        # custom qkv loader that reorders the grouped checkpoint layout), and
+        # only then get replaced by their quantized form.
+        weight = Parameter(
+            torch.empty(
+                sum(output_partition_sizes),
+                input_size_per_partition,
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
+        layer.register_parameter("weight", weight)
+        set_weight_attrs(weight, extra_weight_attrs)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        from comfy_kitchen.tensor.int8 import TensorWiseINT8Layout
+
+        weight = layer.weight.data
+        if weight.dtype == torch.int8:  # already processed
+            return
+
+        # Quantization runs on CUDA, but the model may still be staged on CPU
+        # for offload. Round-trip one layer at a time rather than relying on
+        # the loader's whole-model device move, which would not fit in VRAM.
+        home = weight.device
+        qdata, params = TensorWiseINT8Layout.quantize(
+            weight.to("cuda", non_blocking=True),
+            is_weight=True,
+            per_channel=True,
+            convrot=True,
+            convrot_groupsize=self.quant_config.group_size,
+            stochastic_rounding=0,
+        )
+        layer.weight = Parameter(qdata.to(home), requires_grad=False)
+        layer.register_parameter(
+            "weight_scale",
+            Parameter(
+                params.scale.to(device=home, dtype=torch.float32), requires_grad=False
+            ),
+        )
+        self.quant_config.note_quantized(weight.numel() * weight.element_size())
+        del qdata, params
+        torch.cuda.empty_cache()
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        out_code = _OUT_DTYPE_CODE.get(x.dtype)
+        if out_code is None:
+            raise ValueError(
+                f"kitchen_int8 does not support activation dtype {x.dtype}"
+            )
+
+        # The kernel takes 2D activations; callers may pass [..., K].
+        orig_shape = x.shape
+        if x.dim() != 2:
+            x = x.reshape(-1, orig_shape[-1])
+        x = x.contiguous()
+
+        def run(rows: torch.Tensor) -> torch.Tensor:
+            return torch.ops.comfy_kitchen.int8_linear(
+                rows,
+                layer.weight,
+                layer.weight_scale,
+                bias,
+                out_code,
+                True,  # convrot
+                self.quant_config.group_size,
+            )
+
+        n_rows, n_out = x.shape[0], layer.weight.shape[0]
+        split = _row_split(n_rows, n_out)
+        if split is None:
+            out = run(x)
+        else:
+            # Row slices of a contiguous 2D tensor are themselves contiguous, so
+            # this splits without copying the activation.
+            out = torch.empty(n_rows, n_out, dtype=x.dtype, device=x.device)
+            for start in range(0, n_rows, split):
+                out[start : start + split] = run(x[start : start + split])
+
+        if len(orig_shape) != 2:
+            out = out.reshape(*orig_shape[:-1], out.shape[-1])
+        return out

python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py

@@ -742,13 +742,13 @@ def _resolve_quant_config(
         if server_args.quantization == "modelslim":
             return get_quant_config(hf_config, component_model_path)
 
-        # Online-quant convention: for `fp8` and `mxfp4`, a no-arg
-        # QuantizationConfig() selects the post-load path -- weights load
-        # in source dtype and are quantized in
+        # Online-quant convention: for `fp8`, `mxfp4` and `kitchen_int8`, a
+        # no-arg QuantizationConfig() selects the post-load path -- weights
+        # load in source dtype and are quantized in
         # process_weights_after_loading.
         quant_cls = get_quantization_config(server_args.quantization)
         quant_kwargs = {}
-        if server_args.quantization in {"fp8", "mxfp4"}:
+        if server_args.quantization in {"fp8", "mxfp4", "kitchen_int8"}:
             quant_kwargs["ignored_layers"] = getattr(
                 server_args, "quantization_ignored_layers", None
             )

python/sglang/multimodal_gen/test/unit/test_sol_attn_backend.py

@@ -7,6 +7,7 @@ import torch
 from sglang.multimodal_gen.runtime.layers.attention.backends.sol_attn import (
     SolAttnBackend,
     SolAttnImpl,
+    _get_sol_attn_runtime_config,
     _parse_layer_ranges,
 )
 from sglang.multimodal_gen.runtime.platforms.cuda import CudaPlatformBase
@@ -43,6 +44,32 @@ class TestSolAttnBackend(unittest.TestCase):
     def test_parse_layer_ranges(self):
         self.assertEqual(_parse_layer_ranges("0,1,3-5"), frozenset({0, 1, 3, 4, 5}))
 
+    def test_dense_backend_aliases(self):
+        for raw, expected in (
+            ("fa", "fa"),
+            ("sage", "sage_attn"),
+            ("sage_attn", "sage_attn"),
+        ):
+            server_args = MagicMock()
+            server_args.attention_backend_config = {"dense_backend": raw}
+            with patch(
+                "sglang.multimodal_gen.runtime.layers.attention.backends.sol_attn.get_global_server_args",
+                return_value=server_args,
+            ):
+                self.assertEqual(
+                    _get_sol_attn_runtime_config()["dense_backend"], expected
+                )
+        server_args = MagicMock()
+        server_args.attention_backend_config = {"dense_backend": "torch_sdpa"}
+        with (
+            patch(
+                "sglang.multimodal_gen.runtime.layers.attention.backends.sol_attn.get_global_server_args",
+                return_value=server_args,
+            ),
+            self.assertRaises(ValueError),
+        ):
+            _get_sol_attn_runtime_config()
+
     def test_backend_head_size(self):
         self.assertEqual(SolAttnBackend.get_supported_head_sizes(), [128])