End-to-end DeepEP v2 non-expand + SonicMoE integration #68

Hi, thanks for your great work!

I am working on an end-to-end MoE layer integration of DeepEP V2 and SonicMoE for EP8 training. The goal is to keep DeepEP v2 non-expand mode, because it avoids sending duplicated activation rows when one token hits multiple local experts on the same destination rank.

In non-expand mode, after dispatch we get:

 recv_x            [N_recv, H]
 recv_topk_idx     [N_recv, K]   # local expert id, -1 for invalid slot
 recv_topk_weights [N_recv, K]
 handle            # used later by combine

The integration is functionally straightforward, but there is one layout mismatch.

SonicMoE’s current general routing API expects compact pair-list inputs:

token_indices  [P]
  expert_indices [P]
  router_scores  [P]

where each entry is one valid (token, expert) pair.

Therefore, with DeepEP non-expand we currently need an extra pair extraction step.
This can be made sync-free by using the pair count from the DeepEP handle, but it still creates extra [P] temporary tensors and partially duplicates SonicMoE’s own routing metadata construction.
What I would like to ask is whether SonicMoE should support DeepEP-style non-expand inputs directly, for example:

  moe_deepep_nonexpand_inputs(
      recv_x,             # [N_recv, H]
      recv_topk_idx,      # [N_recv, K], -1 means invalid
      recv_topk_weights,  # [N_recv, K]
      total_pairs,        # P, known from DeepEP handle
      w1,
      b1,
      w2,
      b2,
      E,
      stream_id,
      activation_type,
      ...
  )

Internally this API would directly build the metadata currently needed by SonicMoE grouped GEMM:

expert_frequency
expert_frequency_offset
x_gather_idx
s_scatter_idx
s_reverse_scatter_idx
num_activated_expert_per_token_offset
compact router_scores [P]

from [N_recv, K], instead of requiring the caller to materialize:

token_indices/expert_indices/router_scores [P]

The grouped GEMM path itself could stay unchanged.

A simplified target dataflow would be:

DeepEP dispatch(non-expand)
  recv_x / recv_topk_idx / recv_topk_weights / handle
        │
        ▼
SonicMoE native non-expand metadata + local expert compute
        │
        ▼
local_out [N_recv, H]
        │
        ▼
DeepEP combine(handle)
        │
        ▼
routed_out [T, H]

Questions:

1. Does this look like the right SonicMoE-side abstraction for end-to-end DeepEP v2 non-expand integration?
2. Is there any correctness or performance assumption in moe_general_routing_inputs or general_routing_router_metadata_triton that would make direct [N_recv, K] input problematic?
3. token_indices is currently assumed to be sorted ascending. Is this only needed for building num_activated_expert_per_token_offset, or does it affect other parts of the forward/backward path?
4. For backward, would you recommend returning/scattering gradients for routing scores back to [N_recv, K], or is there a better metadata representation to keep the full forward+backward path efficient?

Any advice on the preferred API/kernel design would be very helpful. My preference is to avoid changing the QuACK grouped GEMM kernels and only add a SonicMoE metadata path that consumes DeepEP non-expand layout directly.ttering it back to [N_recv, K]?

Any guidance on the preferred design would be very helpful. Thank you very much!