DeepEP LL support variable hidden size and tokens num (NVIDIA#6141)

Signed-off-by: Yilin Zhang <18275976+yilin-void@users.noreply.github.com>

Author: yilin-void

cpp/tensorrt_llm/deep_ep/CMakeLists.txt

@@ -1,4 +1,4 @@
-set(DEEP_EP_COMMIT eb3f072664251c05074c3ecc3c3f5dad179c29a9)
+set(DEEP_EP_COMMIT 7b15af835942675df041eca2dcb9930b880287e1)
 set(NVSHMEM_URL_HASH
     SHA256=eb2c8fb3b7084c2db86bd9fd905387909f1dfd483e7b45f7b3c3d5fcf5374b5a)
 

tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py

@@ -100,7 +100,7 @@ class VariableLengthLowLatencyBuffer:
     def __init__(self, mapping: Mapping):
         self.comm = mpi_comm().Split(mapping.pp_rank, mapping.moe_ep_rank)
         self.buffer = None
-        self.num_max_dispatch_tokens_per_rank = None
+        self.num_experts = None
 
     def __del__(self):
         self.comm.Free()
@@ -120,6 +120,7 @@ def reserve(self, num_max_dispatch_tokens_per_rank: int, hidden_size: int,
         allow_nvlink_for_low_latency_mode = (os.environ.get(
             "TRTLLM_DEEP_EP_DISABLE_P2P_FOR_LOW_LATENCY_MODE", "0") == "0")
 
+        assert self.num_experts is None or self.num_experts == num_experts
         # Allocate a buffer if not existed or not enough buffer size
         if self.buffer is None or self.buffer.num_rdma_bytes < num_rdma_bytes:
             # NOTES: for best performance, the QP number **must** be equal to the number of the local experts
@@ -133,17 +134,13 @@ def reserve(self, num_max_dispatch_tokens_per_rank: int, hidden_size: int,
                                  allow_nvlink_for_low_latency_mode=
                                  allow_nvlink_for_low_latency_mode,
                                  comm=self.comm)
+            self.num_experts = num_experts
 
     def low_latency_dispatch(self, hidden_states: torch.Tensor,
                              topk_idx: torch.Tensor,
                              num_max_dispatch_tokens_per_rank: int,
                              num_experts: int):
-        if self.num_max_dispatch_tokens_per_rank is None:
-            self.num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank
-        if num_max_dispatch_tokens_per_rank != self.num_max_dispatch_tokens_per_rank:
-            raise NotImplementedError(
-                "There are issues if `low_latency_dispatch` calls use different `num_max_dispatch_tokens_per_rank` values"
-            )
+        assert num_experts == self.num_experts
 
         # Do MoE dispatch, compatible with CUDA graph (but you may restore some buffer status once you replay)
         recv_hidden_states, recv_expert_count, handle, event, hook = \

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py

@@ -463,15 +463,14 @@ def forward_chunk(
                 if not use_postquant_alltoall:
                     deep_ep_topk_idx = token_selected_slots
                     deep_ep_topk_weights = token_final_scales
+                    assert all_rank_max_num_tokens <= self.deep_ep_max_num_tokens
                     x, recv_expert_count, deep_ep_handle = \
-                        self.deep_ep_buffer.low_latency_dispatch(x, deep_ep_topk_idx, self.deep_ep_max_num_tokens, self.num_slots)
-                    # x shape: [#local experts, EP size * deep_ep_max_num_tokens, hidden_size]
+                        self.deep_ep_buffer.low_latency_dispatch(x, deep_ep_topk_idx, all_rank_max_num_tokens, self.num_slots)
+                    # x shape: [#local experts, EP size * all_rank_max_num_tokens, hidden_size]
                     # recv_expert_count shape: [#local experts]
 
                     # Adapter between `torch.ops.trtllm.fused_moe` and DeepEP
                     # TODO: remove the adapter by changing `torch.ops.trtllm.fused_moe` API
-                    x = x[:, :self.mapping.moe_ep_size *
-                          all_rank_max_num_tokens]
                     mask = torch.arange(
                         x.shape[1], dtype=torch.int32, device=x.device).expand(
                             x.shape[0],
@@ -615,26 +614,14 @@ def forward_chunk(
 
                 deep_ep_topk_idx = token_selected_slots
                 deep_ep_topk_weights = token_final_scales
-                # Each LL combine/dispatch kernel call requires that the `dispatch_rdma_recv_count_buffer` be properly cleaned.
-                # However, the offset of this buffer within the entire RDMA buffer changes according to the hidden size.
-                # Therefore, if the hidden size for the next LL dispatch/combine call is different from the current kernel call, manual cleaning is necessary.
-                if packed_hidden_size != hidden_size:
-                    self.deep_ep_buffer.clean_low_latency_buffer(
-                        self.deep_ep_max_num_tokens, packed_hidden_size,
-                        self.num_slots)
+
+                assert all_rank_max_num_tokens <= self.deep_ep_max_num_tokens
                 fp4_packed_tensor, recv_expert_count, deep_ep_handle = \
-                    self.deep_ep_buffer.low_latency_dispatch(fp4_packed_tensor, deep_ep_topk_idx, self.deep_ep_max_num_tokens, self.num_slots)
-                if packed_hidden_size != hidden_size:
-                    self.deep_ep_buffer.clean_low_latency_buffer(
-                        self.deep_ep_max_num_tokens, hidden_size,
-                        self.num_slots)
+                    self.deep_ep_buffer.low_latency_dispatch(fp4_packed_tensor, deep_ep_topk_idx, all_rank_max_num_tokens, self.num_slots)
                 deep_ep_handle = list(deep_ep_handle)
                 deep_ep_handle[3] = hidden_size
                 deep_ep_handle = tuple(deep_ep_handle)
 
-                fp4_packed_tensor = fp4_packed_tensor[:, :self.mapping.
-                                                      moe_ep_size *
-                                                      all_rank_max_num_tokens]
                 assert fp4_packed_tensor.ndim == 3 and fp4_packed_tensor.shape[
                     2] == packed_hidden_size
                 x_sf = fp4_packed_tensor[:, :, x.shape[1]:x.shape[1] +
@@ -707,23 +694,9 @@ def forward_chunk(
                     final_hidden_states, deep_ep_handle)
             elif self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency:
                 num_tokens_per_expert_for_fused_moe = self.mapping.moe_ep_size * all_rank_max_num_tokens
-                num_tokens_per_expert_for_deep_ep = self.deep_ep_max_num_tokens * self.mapping.moe_ep_size
                 final_hidden_states = final_hidden_states.view(
                     self.expert_size_per_partition,
                     num_tokens_per_expert_for_fused_moe, self.hidden_size)
-                if num_tokens_per_expert_for_deep_ep != num_tokens_per_expert_for_fused_moe:
-                    # Adapter between fused_moe num_tokens and DeepEP num_tokens
-                    # This adapter can be removed if fused_moe accepts DeepEP num_tokens without overhead
-                    final_hidden_states_for_fused_moe = final_hidden_states
-                    final_hidden_states = torch.empty(
-                        self.expert_size_per_partition,
-                        self.deep_ep_max_num_tokens * self.mapping.moe_ep_size,
-                        self.hidden_size,
-                        dtype=final_hidden_states.dtype,
-                        device=final_hidden_states.device)
-                    final_hidden_states[:, :
-                                        num_tokens_per_expert_for_fused_moe] = final_hidden_states_for_fused_moe
-                    del final_hidden_states_for_fused_moe  # Release memory
                 final_hidden_states = self.deep_ep_buffer.low_latency_combine(
                     final_hidden_states, deep_ep_topk_idx, deep_ep_topk_weights,
                     deep_ep_handle)