Support custom DeepEP tuning config (#6257)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user