Support custom DeepEP tuning config (#6257)

This commit is contained in:
fzyzcjy
2025-05-18 08:09:42 +08:00
committed by GitHub
parent 26ebb849eb
commit fd08c04821
6 changed files with 79 additions and 5 deletions

View File

@@ -1,3 +1,5 @@
import json
import os
import unittest
from types import SimpleNamespace
@@ -64,8 +66,34 @@ class TestDPAttn(unittest.TestCase):
"2",
"--enable-dp-attention",
"--enable-deepep-moe",
"--deepep-mode",
"normal",
"--disable-cuda-graph",
# Test custom config
"--deepep-config",
json.dumps(
{
"normal_dispatch": {
"num_sms": 20,
"num_max_nvl_chunked_send_tokens": 16,
"num_max_nvl_chunked_recv_tokens": 256,
"num_max_rdma_chunked_send_tokens": 6,
"num_max_rdma_chunked_recv_tokens": 128,
},
"normal_combine": {
"num_sms": 20,
"num_max_nvl_chunked_send_tokens": 6,
"num_max_nvl_chunked_recv_tokens": 256,
"num_max_rdma_chunked_send_tokens": 6,
"num_max_rdma_chunked_recv_tokens": 128,
},
}
),
],
env={
"SGL_ENABLE_JIT_DEEPGEMM": "0",
**os.environ,
},
)
@classmethod