Use TMA to optimize internode combine. (#287)

* Let forwarders use a dedicated SM

* Shuffle rdma idx

* Sender use TMA.

* Adjust the tuning chunk size.

* Modify NVL chunk layout.

* Update some combine config.

* Small lint

* Minor fix

* Overlap TMA store

---------
Co-authored-by: Chenggang Zhao <chenggangz@deepseek.com>

deep_ep/buffer.py

@@ -231,9 +231,9 @@ class Buffer:
             2: Config(Buffer.num_sms, 10, 256, 6, 128),
             4: Config(Buffer.num_sms, 9, 256, 6, 128),
             8: Config(Buffer.num_sms, 4, 256, 6, 128),
-            16: Config(Buffer.num_sms, 2, 288, 28, 128),
-            24: Config(Buffer.num_sms, 1, 288, 20, 128),
-            32: Config(Buffer.num_sms, 1, 288, 20, 128),
+            16: Config(Buffer.num_sms, 4, 288, 16, 128),
+            24: Config(Buffer.num_sms, 1, 288, 8, 128),
+            32: Config(Buffer.num_sms, 1, 288, 8, 128),
             64: Config(Buffer.num_sms, 1, 288, 20, 128),
             128: Config(Buffer.num_sms, 1, 560, 12, 128),
             144: Config(Buffer.num_sms, 2, 720, 8, 128),

tests/test_internode.py

@@ -209,7 +209,7 @@ def test_main(args: argparse.Namespace, num_sms: int,
 
     # Tune combine performance
     best_time, best_results = 1e10, None
-    for nvl_chunk_size in range(1, 5, 1):
+    for nvl_chunk_size in range(1, 13, 1):
         for rdma_chunk_size in range(8, 33, 4):
             config = deep_ep.Config(num_sms, nvl_chunk_size, nvl_buffer_size, rdma_chunk_size, rdma_buffer_size)
             tune_args = {'x': recv_x, 'handle': handle, 'config': config}