[Feature] Optimize TP LMHead with All-to-All (#32313)

test/registered/mock_model/test_e2e_tp.py

@@ -4,7 +4,7 @@ import unittest
 
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.mock_model.utils import run_mock_model_bench_serving
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import CustomTestCase, is_in_amd_ci
 
 register_cuda_ci(est_time=600, stage="extra-a", runner_config="2-gpu-large")
 register_amd_ci(est_time=167, stage="extra-a", runner_config="2-gpu-large-amd")
@@ -16,6 +16,29 @@ class TestE2ETensorParallel(CustomTestCase):
             extra_server_args=["--tp", "2", "--mem-fraction-static", "0.88"],
         )
 
+    @unittest.skipIf(is_in_amd_ci(), "PyNccl CUDA graph smoke test requires CUDA.")
+    def test_tp_lm_head_all_to_all_cuda_graph(self) -> None:
+        """Smoke-test the PyNccl all-to-all path used during graph capture."""
+        run_mock_model_bench_serving(
+            extra_server_args=[
+                "--tp",
+                "2",
+                "--dp",
+                "2",
+                "--enable-dp-attention",
+                "--enable-tp-lm-head-all-to-all",
+                "--cuda-graph-max-bs-decode",
+                "4",
+                "--mem-fraction-static",
+                "0.88",
+                "--attention-backend",
+                "triton",
+            ],
+            num_prompts=2,
+            random_input_len=32,
+            random_output_len=2,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/test_model_overrides.py

@@ -69,6 +69,7 @@ class TestModelOverridableWhitelist(CustomTestCase):
                     "quantization",
                     "enable_dp_attention",
                     "enable_dp_lm_head",
+                    "enable_tp_lm_head_all_to_all",
                     "moe_a2a_backend",
                     "ep_size",
                     "moe_dense_tp_size",