[V1] EP/TP MoE + DP Attention (#13931)

2025-03-05 00:27:26 -05:00
parent 0a995d5434
commit 72c62eae5f
17 changed files with 250 additions and 75 deletions
--- a/tests/kernels/test_moe.py
+++ b/tests/kernels/test_moe.py
@@ -217,6 +217,7 @@ def test_mixtral_moe(dtype: torch.dtype):
        intermediate_size=config.intermediate_size,
        params_dtype=dtype,
        tp_size=1,
+        dp_size=1,
    ).cuda()

    # Load the weights