[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)