Support GPT-OSS zigzag CP with TRTLLM-MHA (#31732)

2026-07-20 00:48:04 -07:00
parent 2f14d6c6f2
commit 9668d9ea72
5 changed files with 143 additions and 33 deletions
@@ -0,0 +1,32 @@
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.gpt_oss_common import BaseTestGptOss
+
+register_cuda_ci(est_time=220, stage="extra-b", runner_config="4-gpu-b200")
+
+
+class TestGptOss4GpuMxfp4CP(BaseTestGptOss):
+    def test_mxfp4_120b(self):
+        self.run_test(
+            model_variant="120b",
+            quantization="mxfp4",
+            expected_score_of_reasoning_effort={
+                "low": 0.58,
+            },
+            other_args=[
+                "--tp",
+                "4",
+                "--enable-prefill-cp",
+                "--attn-cp-size",
+                "4",
+                "--cp-strategy",
+                "zigzag",
+                "--cuda-graph-max-bs-decode",
+                "200",
+            ],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()