more (#275)

tests/test_low_latency.py

@@ -164,7 +164,8 @@ def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
         print(f'Allocating buffer size: {num_rdma_bytes / 1e6} MB ...', flush=True)
     buffer = deep_ep.Buffer(group, num_rdma_bytes=num_rdma_bytes, low_latency_mode=True,
                             num_qps_per_rank=num_experts // num_ranks,
-                            allow_nvlink_for_low_latency_mode=not args.disable_nvlink, explicitly_destroy=True)
+                            allow_nvlink_for_low_latency_mode=not args.disable_nvlink, explicitly_destroy=True,
+                            allow_mnnvl=args.allow_mnnvl)
     test_main(num_tokens, hidden, num_experts, num_topk, rank, num_ranks, group, buffer,
               use_logfmt=args.use_logfmt, seed=1)
 
@@ -198,6 +199,8 @@ if __name__ == '__main__':
                        help='Number of top-k experts (default: 8)')
     parser.add_argument('--num-experts', type=int, default=288,
                        help='Number of experts (default: 288)')
+    parser.add_argument('--allow-mnnvl', action="store_true",
+                        help='Allow MNNVL for communication')
     parser.add_argument('--disable-nvlink', action='store_true',
                         help='Whether to disable NVLink for testing')
     parser.add_argument('--use-logfmt', action='store_true',