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

2026-08-31 19:37:47 -07:00
parent 1591dcd91a
commit e6f21cdadc
5 changed files with 196 additions and 17 deletions
@@ -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.