diff --git a/python/sglang/srt/arg_groups/model_overrides/__init__.py b/python/sglang/srt/arg_groups/model_overrides/__init__.py index a9e536147..b9431a2aa 100644 --- a/python/sglang/srt/arg_groups/model_overrides/__init__.py +++ b/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 diff --git a/python/sglang/srt/arg_groups/model_overrides/cohere2_moe.py b/python/sglang/srt/arg_groups/model_overrides/cohere2_moe.py new file mode 100644 index 000000000..5b9125779 --- /dev/null +++ b/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"} diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index e58734bf7..1682879d3 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/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", diff --git a/python/sglang/srt/models/cohere2_moe.py b/python/sglang/srt/models/cohere2_moe.py index aa3e39a29..10c428243 100644 --- a/python/sglang/srt/models/cohere2_moe.py +++ b/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) diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 3659a2444..27e18bc76 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/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.