[Cohere Command-A-Plus] Optimize decode and BCG capture on SM10X (#36624)

python/sglang/srt/arg_groups/model_overrides/__init__.py

@@ -9,6 +9,7 @@ the order of the imports below. ``test_model_override_split.py`` forbids the
 overlap, which is why this list needs no particular order.
 """
 
+from sglang.srt.arg_groups.model_overrides import cohere2_moe  # noqa: F401
 from sglang.srt.arg_groups.model_overrides import deepseek_v2  # noqa: F401
 from sglang.srt.arg_groups.model_overrides import deepseek_v4  # noqa: F401
 from sglang.srt.arg_groups.model_overrides import exaone  # noqa: F401

python/sglang/srt/arg_groups/model_overrides/cohere2_moe.py

@@ -0,0 +1,48 @@
+"""Config-time override declarations for cohere2_moe.
+
+Architectures: Cohere2MoeForCausalLM, Cohere2VisionForConditionalGeneration.
+"""
+
+import logging
+from typing import Any
+
+from sglang.srt.arg_groups.model_override_base import (
+    _register_for,
+    model_config_of,
+    resolving_view,
+)
+from sglang.srt.runtime_context import get_platform
+
+logger = logging.getLogger(__name__)
+
+
+def _is_nvfp4_pack_quantized(hf_config: Any) -> bool:
+    # Note(mmangkad): nvfp4-pack-quantized is llm-compressor's output format.
+    qc = getattr(hf_config, "quantization_config", None)
+    if not isinstance(qc, dict):
+        return False
+    groups = qc.get("config_groups") or {}
+    formats = [qc.get("format", "")] + [
+        g.get("format", "") for g in groups.values() if isinstance(g, dict)
+    ]
+    return any("nvfp4" in str(fmt) for fmt in formats)
+
+
+@_register_for(
+    "Cohere2VisionForConditionalGeneration",
+    "Cohere2MoeForCausalLM",
+)
+def _cohere2_moe_runner_overrides(server_args: Any, hf_config: Any) -> dict:
+    cfg = resolving_view(server_args)
+    if cfg.moe_runner_backend != "auto":
+        return {}
+    if not get_platform().is_sm100:
+        return {}
+    if model_config_of(server_args).quantization is not None:
+        if not _is_nvfp4_pack_quantized(hf_config):
+            return {}
+    logger.info(
+        "Command-A-Plus on SM10X: moe_runner_backend=flashinfer_trtllm "
+        "(trtllm-gen fused MoE)."
+    )
+    return {"moe_runner_backend": "flashinfer_trtllm"}

python/sglang/srt/configs/model_config.py

@@ -1971,7 +1971,6 @@ piecewise_cuda_graph_disabled_model_archs = [
 # all multimodal models; archs here opt back in because their LM prefill captures
 # cleanly (vision encoder runs eagerly outside the graph via general_mm_embed_routine).
 multimodal_piecewise_cuda_graph_supported_model_archs = [
-    "Cohere2VisionForConditionalGeneration",
     "KimiK25ForConditionalGeneration",
     "MiniMaxM3SparseForCausalLM",
     "MiniMaxM3SparseForConditionalGeneration",
@@ -1984,6 +1983,7 @@ multimodal_piecewise_cuda_graph_supported_model_archs = [
 # generic multimodal rule disabled prefill CG for them despite the LM prefill
 # capturing cleanly.
 multimodal_breakable_cuda_graph_supported_model_archs = [
+    "Cohere2VisionForConditionalGeneration",
     "InternS2MobiusForConditionalGeneration",
     "PaddleOCRVLForConditionalGeneration",
     "Qwen3_5ForConditionalGeneration",

python/sglang/srt/models/cohere2_moe.py

@@ -31,7 +31,7 @@ from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
-from sglang.srt.runtime_context import get_parallel
+from sglang.srt.runtime_context import get_parallel, get_stream
 from sglang.srt.utils import add_prefix, get_compiler_backend, is_cuda, make_layers
 
 
@@ -290,6 +290,8 @@ class Cohere2MoeSparseMoeBlock(nn.Module):
             layer_id=layer_id,
             prefix=add_prefix("experts", prefix),
             routing_method_type=routing_method_type,
+            # The decoder reads this buffer on another stream, so in-place races it.
+            inplace=False,
         )
 
         num_shared_experts = getattr(config, "num_shared_experts", 0)
@@ -310,13 +312,10 @@ class Cohere2MoeSparseMoeBlock(nn.Module):
         )
         assert self.shared_expert_combination_strategy in ("average", "sum")
 
-        # Auxiliary CUDA stream so shared_experts can overlap with the
-        # gate + routed-experts path inside a captured CUDA graph. Only used
-        # during capture/replay; outside capture the sync overhead outweighs it.
+        # Leased by role, not per layer. A stream per layer gives the caching
+        # allocator a segregated pool each, and prefill graph capture then OOMs.
         self.alt_stream = (
-            torch.cuda.Stream()
-            if is_cuda() and self.shared_experts is not None
-            else None
+            get_stream("alt") if is_cuda() and self.shared_experts is not None else None
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -329,9 +328,7 @@ class Cohere2MoeSparseMoeBlock(nn.Module):
             final_hidden_states = self.experts(hidden_states, topk_output)
             return final_hidden_states.view(orig_shape)
 
-        # FusedMoE.experts can write back into its input buffer (observed for
-        # the unquantized triton BF16 path). Snapshot the post-norm input so
-        # the shared-expert branch sees the original layernorm output.
+        # Unnecessary now that experts is inplace=False, but sharing measured slower.
         shared_input = hidden_states.clone()
 
         if self.alt_stream is not None and get_is_capture_mode():
@@ -405,6 +402,10 @@ class Cohere2MoeDecoderLayer(nn.Module):
         self.input_layernorm = Cohere2MoeLayerNorm(config.hidden_size, eps=norm_eps)
         self.tp_size = get_parallel().tp_size
 
+        # The parallel block makes attn(norm(x)) and mlp(norm(x)) independent, so
+        # the two chains are safe to run concurrently on a second side stream.
+        self.mlp_stream = get_stream("alt_mlp") if is_cuda() else None
+
     def forward(
         self,
         positions: torch.Tensor,
@@ -416,12 +417,31 @@ class Cohere2MoeDecoderLayer(nn.Module):
         # sum-then-allreduce, halving per-layer all-reduces.
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
-        attn_out = self.self_attn(
-            positions=positions,
-            hidden_states=hidden_states,
-            forward_batch=forward_batch,
-        )
-        mlp_out = self.mlp(hidden_states)
+        if (
+            self.mlp_stream is not None
+            and get_is_capture_mode()
+            and forward_batch.forward_mode.is_decode()
+        ):
+            # Decode only. A prefill chunk already saturates the GPU, so there the
+            # split buys nothing and costs event overhead.
+            current_stream = torch.cuda.current_stream()
+            self.mlp_stream.wait_stream(current_stream)
+            with torch.cuda.stream(self.mlp_stream):
+                mlp_out = self.mlp(hidden_states)
+            attn_out = self.self_attn(
+                positions=positions,
+                hidden_states=hidden_states,
+                forward_batch=forward_batch,
+            )
+            # Neither tensor needs record_stream, both stay referenced until here.
+            current_stream.wait_stream(self.mlp_stream)
+        else:
+            attn_out = self.self_attn(
+                positions=positions,
+                hidden_states=hidden_states,
+                forward_batch=forward_batch,
+            )
+            mlp_out = self.mlp(hidden_states)
         combined = attn_out + mlp_out
         if self.tp_size > 1:
             combined = tensor_model_parallel_all_reduce(combined)

test/registered/unit/test_model_overrides.py

@@ -675,6 +675,116 @@ class TestGoldenModelOverrides(_IsolatedPublish):
         self._publish(nvfp4)
         self.assertEqual(get_exec().moe.moe_runner_backend, "flashinfer_trtllm_routed")
 
+    # ---- Command-A-Plus (Cohere2Moe) MoE runner gate ----
+
+    _NVFP4_QUANT = {
+        "quant_method": "compressed-tensors",
+        "format": "nvfp4-pack-quantized",
+        "config_groups": {"group_0": {"targets": ["Linear"]}},
+    }
+    _FP8_QUANT = {
+        "quant_method": "compressed-tensors",
+        "format": "float-quantized",
+        "config_groups": {"group_0": {"format": "float-quantized"}},
+    }
+
+    def test_cohere2_moe_runner_gate(self):
+        """FP8 must stay on auto. flashinfer_trtllm rejects its quant info at
+        the first forward, so a wrong answer here crashes mid-serving."""
+        with override_platform(is_sm100=True):
+            nvfp4 = self._construct(
+                "Cohere2MoeForCausalLM",
+                "cohere2_moe",
+                config_extra={"quantization_config": self._NVFP4_QUANT},
+            )
+            fp8 = self._construct(
+                "Cohere2MoeForCausalLM",
+                "cohere2_moe",
+                config_extra={"quantization_config": self._FP8_QUANT},
+            )
+            bf16 = self._construct("Cohere2MoeForCausalLM", "cohere2_moe")
+            explicit = self._construct(
+                "Cohere2MoeForCausalLM",
+                "cohere2_moe",
+                config_extra={"quantization_config": self._NVFP4_QUANT},
+                moe_runner_backend="triton",
+            )
+
+        b = "moe_runner_backend"
+        self.assertEqual(self._resolved(nvfp4, b), "flashinfer_trtllm")
+        self.assertEqual(self._resolved(bf16, b), "flashinfer_trtllm")
+        self.assertEqual(self._resolved(fp8, b), "auto")
+        self.assertEqual(self._resolved(explicit, b), "triton")
+        self.assertIn(
+            (
+                "_cohere2_moe_runner_overrides",
+                {"moe_runner_backend": "flashinfer_trtllm"},
+            ),
+            nvfp4._resolved_overrides,
+        )
+
+        # Non-SM10X keeps the existing auto behavior.
+        with override_platform(is_sm100=False):
+            non_sm10x = self._construct(
+                "Cohere2MoeForCausalLM",
+                "cohere2_moe",
+                config_extra={"quantization_config": self._NVFP4_QUANT},
+            )
+        self.assertEqual(self._resolved(non_sm10x, b), "auto")
+
+        self._publish(nvfp4)
+        self.assertEqual(get_exec().moe.moe_runner_backend, "flashinfer_trtllm")
+
+        # The vision wrapper carries a text_config that holds no
+        # quantization_config of its own, so the format is read at the top level.
+        from sglang.srt.arg_groups.model_overrides.cohere2_moe import (
+            _is_nvfp4_pack_quantized,
+        )
+
+        self.assertTrue(
+            _is_nvfp4_pack_quantized(
+                SimpleNamespace(
+                    text_config=SimpleNamespace(),
+                    quantization_config=self._NVFP4_QUANT,
+                )
+            )
+        )
+
+    def test_cohere2_moe_runner_gate_fails_closed(self):
+        """A quantized checkpoint we cannot positively read as NVFP4 stays on
+        auto. Treating an unreadable config as BF16 would force the runner."""
+        from sglang.srt.arg_groups.model_overrides.cohere2_moe import (
+            _cohere2_moe_runner_overrides,
+        )
+
+        def _args(quantization):
+            # A fixture-supplied `_model_config` is what `model_config_of`
+            # hands back, so this needs no checkpoint on disk.
+            return SimpleNamespace(
+                moe_runner_backend="auto",
+                _model_config=SimpleNamespace(quantization=quantization),
+            )
+
+        with override_platform(is_sm100=True):
+            # quantization_config the walk cannot read (an object, or absent
+            # because it lives in a standalone hf_quant_config.json).
+            for hf_config in (
+                SimpleNamespace(quantization_config=object()),
+                SimpleNamespace(),
+            ):
+                self.assertEqual(
+                    _cohere2_moe_runner_overrides(
+                        _args("compressed-tensors"), hf_config
+                    ),
+                    {},
+                )
+
+            # Genuinely unquantized -> trtllm-gen.
+            self.assertEqual(
+                _cohere2_moe_runner_overrides(_args(None), SimpleNamespace()),
+                {"moe_runner_backend": "flashinfer_trtllm"},
+            )
+
     def test_mimo_v2_declarations(self):
         # Callable-level golden: MiMoV2 archs are hybrid (config-shape heavy),
         # so the declaration is pinned directly for both provider inputs.