Adjust token major layout (#844)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-18 01:37:24 +00:00 · 2026-07-16 18:56:10 -07:00
parent 39813bfff1
commit dd9fa728ad
15 changed files with 190 additions and 60 deletions
--- a/python/mscclpp/ep/README.md
+++ b/python/mscclpp/ep/README.md
@@ -60,11 +60,11 @@ class MoECommunicatorConfig:
    hidden_size: int = 0
    topk: int = 0
    max_tokens_per_rank: int = 0
-    max_recv_tokens_per_rank: Optional[int] = None

    # Runtime mode and output layout
    mode: MoEMode = MoEMode.LOW_LATENCY
    output_layout: Optional[DispatchLayout] = None  # default is derived from mode
+    token_major_init_padding: bool = False

    # Quantization defaults
    quant: Optional[QuantConfig] = None
@@ -137,8 +137,8 @@ a later version can add an explicit `expert_map` for arbitrary placement.
 |---|---|
 | `mode` | Backend selection (`MoEMode.LOW_LATENCY` or `MoEMode.HIGH_THROUGHPUT`) |
 | `output_layout` | MLP input layout returned by dispatch |
+| `token_major_init_padding` | Initialize token-major padding metadata for fixed-capacity Triton kernels |
 | `max_tokens_per_rank` | dispatch capacity |
-| `max_recv_tokens_per_rank` | recv buffer capacity |
 | scratch buffers | internally sized from mode, capacity, topology, and shape |
 | `num_sms` | backend launch/resource tuning |
 | `dispatch_config`, `combine_config` | backend-specific tuning configs |
@@ -310,6 +310,7 @@ class TokenMajorCombineContext:
    hidden_size: int
    source_token_ids: torch.Tensor
    num_tokens_per_rank: torch.Tensor
+    rank_offsets: torch.Tensor
    num_max_dispatch_tokens_per_rank: int


@@ -514,9 +515,24 @@ For token-major LL layout:
 output_buffer: [world_size * max_tokens_per_rank, hidden]
 ```

-The token-major rows are grouped into fixed source-rank regions. For source rank
-`r`, only the first `dispatch_out.layout.num_tokens_per_rank[r]` rows in region
-`[r * max_tokens_per_rank : (r + 1) * max_tokens_per_rank]` are valid.
+The caller chooses `max_tokens_per_rank`, allocates that fixed capacity, and may
+pass the resulting row capacity directly to a Triton kernel. This avoids a CPU
+synchronization while keeping a static CUDA Graph shape. All valid rows are
+compacted into one contiguous prefix. For source rank `r`, its rows are:
+
+```python
+begin = dispatch_out.layout.offsets[r]
+end = dispatch_out.layout.offsets[r + 1]
+```
+
+`offsets[-1]` is the total number of valid rows.
+
+Rows after `offsets[-1]` are padding. Set
+`token_major_init_padding=True` when a fixed-capacity Triton kernel should
+process the entire allocation: padding `topk_ids` are then initialized to `-1`
+and weights to `0`, so the kernel can skip a row when all expert IDs are `-1`.
+The option defaults to `False` to avoid initialization overhead when the MLP
+uses the compact valid length.

 The dtype must match the dispatch output dtype. For BF16 dispatch it is BF16.
 For FP8 dispatch it is FP8 and the returned `DispatchOutput.quant` carries the
@@ -560,8 +576,10 @@ dispatch_out.topk_ids          # [world_size * max_tokens_per_rank, K], int32 lo
 dispatch_out.weights           # [world_size * max_tokens_per_rank, K], float32
 ```

-Non-local top-k entries use expert ID `-1` and weight `0`. The valid row count in each
-source-rank region is returned in `dispatch_out.layout.num_tokens_per_rank`.
+Only the prefix ending at `dispatch_out.layout.offsets[-1]` contains valid
+tokens. Non-local entries always use expert ID `-1` and weight `0`; padding uses
+the same sentinel values when `token_major_init_padding=True`. Per-source-rank
+counts are returned in `dispatch_out.layout.num_tokens_per_rank`.
 For expert-major output, only the first
 `dispatch_out.layout.num_tokens_per_expert[i]` slots are valid:

--- a/python/mscclpp/ep/low_latency.py
+++ b/python/mscclpp/ep/low_latency.py
@@ -54,11 +54,19 @@ class LowLatencyRuntime:
        hidden: int,
        num_experts: int,
        num_topk: int,
+        initialize_token_major_padding: bool,
    ) -> None:
        self.rank: int = comm.my_rank
        self.group_size: int = comm.nranks
        self.comm = comm
-        self.cpp_runtime = MoERuntime(comm.communicator, max_tokens_per_rank, hidden, num_experts, num_topk)
+        self.cpp_runtime = MoERuntime(
+            comm.communicator,
+            max_tokens_per_rank,
+            hidden,
+            num_experts,
+            num_topk,
+            initialize_token_major_padding,
+        )

    def is_available(self) -> bool:
        return self.cpp_runtime.is_available()
@@ -90,6 +98,7 @@ class LowLatencyBackend:
        self.num_blocks = config.low_latency_num_blocks
        self.num_sms = self.num_blocks - 2
        self.combine_mode = config.low_latency_combine_mode
+        self.token_major_init_padding = config.token_major_init_padding
        self.enable_overlap = config.enable_overlap

        if self.output_layout not in (DispatchLayout.EXPERT_MAJOR, DispatchLayout.TOKEN_MAJOR):
@@ -100,6 +109,10 @@ class LowLatencyBackend:
            raise ValueError("low_latency_num_blocks must be between world_size + 2 and 130")
        if not isinstance(self.combine_mode, CombineMode):
            raise TypeError("low_latency_combine_mode must be a CombineMode")
+        if not isinstance(self.token_major_init_padding, bool):
+            raise TypeError("token_major_init_padding must be a bool")
+        if self.token_major_init_padding and self.output_layout != DispatchLayout.TOKEN_MAJOR:
+            raise ValueError("token_major_init_padding requires TOKEN_MAJOR output")
        if self.output_layout == DispatchLayout.TOKEN_MAJOR and self.combine_mode != CombineMode.RANK_LOCAL_REDUCE:
            raise ValueError("TOKEN_MAJOR output requires RANK_LOCAL_REDUCE combine")

@@ -111,8 +124,6 @@ class LowLatencyBackend:
            local_expert_start=config.local_expert_start,
        )

-        if config.max_recv_tokens_per_rank not in (None, self.max_tokens_per_rank):
-            raise NotImplementedError("low-latency mode currently uses max_tokens_per_rank as recv capacity")
        self.dispatch_data_type = _resolve_dispatch_data_type(config.quant)

        self._dispatch_scales: Optional[torch.Tensor] = None
@@ -128,6 +139,7 @@ class LowLatencyBackend:
            hidden=self.hidden_size,
            num_experts=self.num_experts,
            num_topk=self.topk,
+            initialize_token_major_padding=self.token_major_init_padding,
        )
        self._is_internode = self._runtime.is_internode_available()

@@ -188,8 +200,15 @@ class LowLatencyBackend:
        )
        if self.output_layout == DispatchLayout.EXPERT_MAJOR:
            layout_info = DispatchLayoutInfo(kind=self.output_layout, num_tokens_per_expert=count)
+        elif self.output_layout == DispatchLayout.TOKEN_MAJOR:
+            assert layout_range is not None
+            layout_info = DispatchLayoutInfo(
+                kind=self.output_layout,
+                num_tokens_per_rank=count,
+                offsets=layout_range,
+            )
        else:
-            layout_info = DispatchLayoutInfo(kind=self.output_layout, num_tokens_per_rank=count)
+            raise ValueError(f"unsupported low-latency output layout: {self.output_layout}")
        output_info = DispatchOutputInfo(layout=layout_info, quant=output_quant)
        dispatch_out = DispatchOutput(
            tokens=out_buf,
@@ -213,7 +232,8 @@ class LowLatencyBackend:
                    num_max_dispatch_tokens_per_rank=self.max_tokens_per_rank,
                ),
            )
-        else:
+        elif self.output_layout == DispatchLayout.TOKEN_MAJOR:
+            assert layout_range is not None
            handle = TokenMajorDispatchHandle(
                output_info=output_info,
                combine_context=TokenMajorCombineContext(
@@ -223,9 +243,12 @@ class LowLatencyBackend:
                    hidden_size=self.hidden_size,
                    source_token_ids=src_info,
                    num_tokens_per_rank=count,
+                    rank_offsets=layout_range,
                    num_max_dispatch_tokens_per_rank=self.max_tokens_per_rank,
                ),
            )
+        else:
+            raise ValueError(f"unsupported low-latency output layout: {self.output_layout}")
        return dispatch_out, handle

    def combine(
@@ -246,7 +269,7 @@ class LowLatencyBackend:
            context = handle.combine_context
            topk_weights = None
            src_info = context.source_token_ids
-            layout_range = None
+            layout_range = context.rank_offsets
        else:
            raise ValueError("DispatchHandle does not contain low-latency combine context")
        if out is None:
@@ -292,17 +315,19 @@ class LowLatencyBackend:
                        (self.num_local_experts, num_scales, slots_per_expert), dtype=torch.float32, device=device
                    )
                    self._dispatch_scales = scale_storage.transpose(1, 2)
-            else:
+            elif self.output_layout == DispatchLayout.TOKEN_MAJOR:
                token_capacity = self.world_size * self.max_tokens_per_rank
                self._dispatch_src_info = torch.empty((token_capacity,), dtype=torch.int32, device=device)
                self._dispatch_topk_ids = torch.empty((token_capacity, self.topk), dtype=torch.int32, device=device)
                self._dispatch_weights = torch.empty((token_capacity, self.topk), dtype=torch.float32, device=device)
-                self._dispatch_layout_range = None
+                self._dispatch_layout_range = torch.empty((self.world_size + 1,), dtype=torch.int64, device=device)
                self._dispatch_count = torch.empty((self.world_size,), dtype=torch.int32, device=device)
                if self.dispatch_data_type == DispatchDataType.FP8_E4M3:
                    self._dispatch_scales = torch.empty(
                        (token_capacity, self.hidden_size // 128), dtype=torch.float32, device=device
                    )
+            else:
+                raise ValueError(f"unsupported low-latency output layout: {self.output_layout}")
        assert self._dispatch_src_info is not None
        assert self._dispatch_count is not None
        return (
@@ -346,8 +371,10 @@ class LowLatencyBackend:
        slots_per_expert = self.world_size * self.max_tokens_per_rank
        if self.output_layout == DispatchLayout.EXPERT_MAJOR:
            expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size)
-        else:
+        elif self.output_layout == DispatchLayout.TOKEN_MAJOR:
            expected_shape = (self.world_size * self.max_tokens_per_rank, self.hidden_size)
+        else:
+            raise ValueError(f"unsupported low-latency output layout: {self.output_layout}")
        if output_buffer.dim() != len(expected_shape) or not output_buffer.is_contiguous():
            raise ValueError(f"output_buffer must be a contiguous {self.output_layout} tensor")
        expected_dtype = torch.float8_e4m3fn if self.dispatch_data_type == DispatchDataType.FP8_E4M3 else torch.bfloat16
@@ -371,8 +398,12 @@ class LowLatencyBackend:
        slots_per_expert = self.world_size * self.max_tokens_per_rank
        if handle.output_info.layout.kind == DispatchLayout.EXPERT_MAJOR:
            expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size)
-        else:
+        elif handle.output_info.layout.kind == DispatchLayout.TOKEN_MAJOR:
            expected_shape = (self.world_size * self.max_tokens_per_rank, self.hidden_size)
+            if handle.output_info.layout.offsets is None:
+                raise ValueError("TOKEN_MAJOR DispatchHandle is missing compact rank offsets")
+        else:
+            raise ValueError(f"unsupported low-latency output layout: {handle.output_info.layout.kind}")
        if expert_output.dim() != len(expected_shape) or not expert_output.is_contiguous():
            raise ValueError("expert_output must keep dispatch output's contiguous layout")
        if tuple(expert_output.shape) != expected_shape:
--- a/python/mscclpp/ep/types.py
+++ b/python/mscclpp/ep/types.py
@@ -47,11 +47,11 @@ class MoECommunicatorConfig:
    hidden_size: int = 0
    topk: int = 0
    max_tokens_per_rank: int = 0
-    max_recv_tokens_per_rank: Optional[int] = None

    # Runtime mode and output layout
    mode: MoEMode = MoEMode.LOW_LATENCY
    output_layout: Optional[DispatchLayout] = None
+    token_major_init_padding: bool = False

    # Quantization defaults
    quant: Optional[QuantConfig] = None
@@ -127,6 +127,7 @@ class TokenMajorCombineContext:
    hidden_size: int
    source_token_ids: torch.Tensor
    num_tokens_per_rank: torch.Tensor
+    rank_offsets: torch.Tensor
    num_max_dispatch_tokens_per_rank: int


--- a/src/core/include/unix_socket.hpp
+++ b/src/core/include/unix_socket.hpp
@@ -24,6 +24,7 @@ class UnixSocketServer {
  int registerFd(int fd);
  void unregisterFd(int fd);
  std::string getSocketPath() const;
+  ~UnixSocketServer();

 private:
  int listenUnixSockFd_ = -1;
@@ -36,6 +37,7 @@ class UnixSocketServer {

  UnixSocketServer();
  void mainLoop(int listenUnixSockFd);
+  void shutdown();
 };

 class UnixSocketClient {
--- a/src/core/unix_socket.cc
+++ b/src/core/unix_socket.cc
@@ -140,9 +140,13 @@ void UnixSocketServer::start() {
 }

 void UnixSocketServer::stop() {
+  if (mainThread_.joinable()) INFO(MSCCLPP_INIT, "Stopping unix socket server");
+  shutdown();
+}
+
+void UnixSocketServer::shutdown() {
  *abortFlag_ = 1;
  if (mainThread_.joinable()) {
-    INFO(MSCCLPP_INIT, "Stopping unix socket server");
    mainThread_.join();
  }
  ::close(listenUnixSockFd_);
@@ -252,6 +256,8 @@ void UnixSocketServer::mainLoop(int listenUnixSockFd) {

 UnixSocketServer::UnixSocketServer() : abortFlagStorage_(new uint32_t(0)), abortFlag_(abortFlagStorage_.get()) {}

+UnixSocketServer::~UnixSocketServer() { shutdown(); }
+
 std::string UnixSocketServer::getSocketPath() const { return listenUnixSockPath_; }

 UnixSocketClient& UnixSocketClient::instance() {
--- a/src/ext/ep/README.md
+++ b/src/ext/ep/README.md
@@ -34,8 +34,12 @@ LL dispatch supports two user-visible layouts:

 - `EXPERT_MAJOR`: one row per `(token, local expert)`.
 - `TOKEN_MAJOR`: one row per `(token, destination rank)`, plus local top-k expert
-  IDs, routing weights, source-token IDs, and per-source-rank counts. The caller
-  must produce one pre-weighted local partial per row before combine.
+  IDs, routing weights, source-token IDs, per-source-rank counts, and exclusive
+  offsets. Valid rows occupy a compact prefix of the caller-provided capacity
+  buffer. With `token_major_init_padding=True`, padding rows have top-k IDs
+  `-1`, allowing fixed-capacity Triton kernels to skip them without a CPU count
+  synchronization. The option is disabled by default. The caller must produce
+  one pre-weighted local partial per valid row before combine.

 ### High throughput

--- a/src/ext/ep/bindings.cpp
+++ b/src/ext/ep/bindings.cpp
@@ -57,8 +57,9 @@ NB_MODULE(mscclpp_ep_cpp, m) {
      .value("MXFP8_E4M3", mscclpp::ep::low_latency::DispatchDataType::MXFP8_E4M3);

  nb::class_<mscclpp::ep::MoERuntime>(m, "MoERuntime")
-      .def(nb::init<mscclpp::Communicator&, int, int, int, int>(), nb::arg("comm"), nb::arg("max_tokens_per_rank"),
-           nb::arg("hidden"), nb::arg("num_experts"), nb::arg("num_topk"))
+      .def(nb::init<mscclpp::Communicator&, int, int, int, int, bool>(), nb::arg("comm"),
+           nb::arg("max_tokens_per_rank"), nb::arg("hidden"), nb::arg("num_experts"), nb::arg("num_topk"),
+           nb::arg("initialize_token_major_padding"))
      .def("is_available", &mscclpp::ep::MoERuntime::isAvailable)
      .def("is_internode_available", &mscclpp::ep::MoERuntime::isInternodeAvailable)
      .def(
--- a/src/ext/ep/include/api.cuh
+++ b/src/ext/ep/include/api.cuh
@@ -144,6 +144,8 @@ struct Workload {
  int maxTokensPerRank_;
  /// User-visible dispatch output layout.
  DispatchLayout outputLayout_;
+  /// Whether token-major padding metadata is initialized to sentinel values.
+  bool initializeTokenMajorPadding_;
  /// Dispatch payload data format.
  DispatchDataType dispatchDataType_;
 };
@@ -183,7 +185,7 @@ size_t workspaceSize(int numRanks, int numExperts);
 /// @param[out] outputTopkWeights Token-major routing weights
 /// [num_ranks * max_tokens_per_rank, num_topk], or nullptr.
 /// @param[out] outputLayout Per-[local expert, source rank] packed count and offset for expert-major output, or
-/// nullptr.
+/// token-major exclusive source-rank offsets [num_ranks + 1].
 /// @param[out] outputCount Per-local-expert counts for expert-major output or per-source-rank counts for token-major.
 /// @param[in] input Local input tokens [num_tokens, hidden].
 /// @param[in] topkIdx Global expert indices [num_tokens, num_topk].
@@ -206,7 +208,8 @@ void dispatch(void* output, float* outputScales, int* outputSrcInfo, int* output
 /// @param[in] topkIdx Global expert indices [num_tokens, num_topk].
 /// @param[in] topkWeights Routing weights [num_tokens, num_topk], or nullptr for unit weights.
 /// @param[in] srcInfo Original source-token index for every packed expert row.
-/// @param[in] layoutRange Per-[local expert, source rank] packed count and offset.
+/// @param[in] layoutRange Per-[local expert, source rank] packed count and offset for expert-major input, or
+/// token-major exclusive source-rank offsets [num_ranks + 1].
 /// @param[in] workload Per-call workload dimensions.
 /// @param[in,out] recvBuffer Current symmetric ping-pong buffer receiving partials or expert rows.
 /// @param[in] dispatchRecvBuffer Previous dispatch buffer containing rewritten routing metadata.
--- a/src/ext/ep/low_latency/combine.cu
+++ b/src/ext/ep/low_latency/combine.cu
@@ -191,7 +191,8 @@ MSCCLPP_DEVICE_INLINE void sendRankReducedPartials(const void* expertOutput, int
 }

 template <int Hidden>
-MSCCLPP_DEVICE_INLINE void sendTokenMajorPartials(const void* expertOutput, const int* srcInfo, int maxTokensPerRank,
+MSCCLPP_DEVICE_INLINE void sendTokenMajorPartials(const void* expertOutput, const int* srcInfo,
+                                                  const int64_t* rankOffsets, int maxTokensPerRank,
                                                  void* combineRecvBuffer, const TransportView& transport,
                                                  WorkspaceView& workspaceView, uint8_t* sharedMemory) {
  const int threadId = static_cast<int>(threadIdx.x);
@@ -205,12 +206,13 @@ MSCCLPP_DEVICE_INLINE void sendTokenMajorPartials(const void* expertOutput, cons
       taskIdx += static_cast<int>(gridDim.x)) {
    const RecvTask recvTask = loadRecvTask(workspaceView.dispatchRecvTasks_, taskIdx);
    const int sourceRank = recvTask.sourceRank_;
+    const int rankOffset = warpBroadcast(get_lane_id() == 0 ? static_cast<int>(rankOffsets[sourceRank]) : 0, 0);

    for (int sourceTokenSlot = recvTask.tokenBegin_; sourceTokenSlot < recvTask.tokenEnd_;
         ++sourceTokenSlot, ++tokenIteration) {
      const int stage = tokenIteration % CombineNStages;
      auto* outputTile = reinterpret_cast<int4*>(outputTiles + static_cast<size_t>(stage) * HiddenBytes);
-      const int inputRowOffset = sourceRank * maxTokensPerRank + sourceTokenSlot;
+      const int inputRowOffset = rankOffset + sourceTokenSlot;
      const auto* inputRow =
          reinterpret_cast<const int4*>(expertOutput) + static_cast<size_t>(inputRowOffset) * HiddenInt4;
      int4 values[ChunksPerThread];
@@ -438,8 +440,8 @@ __global__ __launch_bounds__(CombineNThreads, 1) void combineKernel(

  if constexpr (Layout == DispatchLayout::TOKEN_MAJOR) {
    static_assert(Mode == low_latency::CombineMode::RANK_LOCAL_REDUCE);
-    sendTokenMajorPartials<Hidden>(expertOutput, srcInfo, maxTokensPerRank, combineRecvBuffer, transport, workspaceView,
-                                   sharedMemory);
+    sendTokenMajorPartials<Hidden>(expertOutput, srcInfo, layoutRange, maxTokensPerRank, combineRecvBuffer, transport,
+                                   workspaceView, sharedMemory);
  } else if constexpr (Mode == low_latency::CombineMode::RANK_LOCAL_REDUCE) {
    sendRankReducedPartials<Hidden, DispatchType, ScaleBlockSize>(
        expertOutput, nExperts, nRanks, nTopk, maxTokensPerRank, combineRecvBuffer, dispatchRecvBuffer, transport,
@@ -552,9 +554,9 @@ inline void combine(void* output, const void* expertOutput, const int64_t* topkI
  const int rank = comm.rank_;
  const int nRanks = comm.numRanks_;

-  EP_HOST_ASSERT(output != nullptr);
+  EP_HOST_ASSERT(workload.numTokens_ == 0 || output != nullptr);
  EP_HOST_ASSERT(expertOutput != nullptr);
-  EP_HOST_ASSERT(topkIndices != nullptr);
+  EP_HOST_ASSERT(workload.numTokens_ == 0 || topkIndices != nullptr);
  EP_HOST_ASSERT(recvBuffer != nullptr);
  EP_HOST_ASSERT(dispatchRecvBuffer != nullptr);
  EP_HOST_ASSERT(comm.symmetricBufferBase_ != nullptr);
@@ -574,6 +576,7 @@ inline void combine(void* output, const void* expertOutput, const int64_t* topkI
  if (workload.outputLayout_ == DispatchLayout::TOKEN_MAJOR) {
    EP_HOST_ASSERT(mode == low_latency::CombineMode::RANK_LOCAL_REDUCE);
    EP_HOST_ASSERT(srcInfo != nullptr);
+    EP_HOST_ASSERT(layoutRange != nullptr);
  } else if (mode == low_latency::CombineMode::DIRECT_SEND) {
    EP_HOST_ASSERT(srcInfo != nullptr);
    EP_HOST_ASSERT(layoutRange != nullptr);
--- a/src/ext/ep/low_latency/dispatch.cu
+++ b/src/ext/ep/low_latency/dispatch.cu
@@ -331,6 +331,12 @@ MSCCLPP_DEVICE_INLINE void dispatchRecvScheduler(int64_t* outputLayout, int* out
    activeRankPrefix += sharedMem[nRankWarps + warpId];
    const int nTotalTokens = sharedMem[2 * nRankWarps];
    const int nActiveRanks = sharedMem[2 * nRankWarps + 1];
+    if constexpr (Layout == DispatchLayout::TOKEN_MAJOR) {
+      if (sourceRank < nRanks) {
+        outputLayout[sourceRank] = rankTokenPrefix - nRankTokens;
+        if (sourceRank == nRanks - 1) outputLayout[nRanks] = rankTokenPrefix;
+      }
+    }
    const int nTasks = nTotalTokens < nWorkerBlocks ? nTotalTokens : nWorkerBlocks;

    // Reserve one task for every active rank. Distribute the remaining tasks
@@ -473,15 +479,13 @@ template <int Hidden, DispatchDataType DataType, int ScaleBlockSize>
 MSCCLPP_DEVICE_INLINE bool dispatchRecvTokenMajorOutput(void* output, float* outputScales, int* outputSrcInfo,
                                                        int* outputTopkIdx, float* outputTopkWeights,
                                                        const DispatchPayloadView<DataType>& payloadView,
-                                                        const void* sourcePayload, int localExpertIdx, int sourceRank,
-                                                        int sourceTokenSlot, int sourceTokenIdx, int nTopk,
-                                                        int maxTokensPerRank, uint8_t* sharedTile, uint64_t* tmaBarrier,
-                                                        uint32_t& recvTmaPhase) {
+                                                        const void* sourcePayload, int localExpertIdx, int outputRow,
+                                                        int sourceTokenIdx, int nTopk, uint8_t* sharedTile,
+                                                        uint64_t* tmaBarrier, uint32_t& recvTmaPhase) {
  using OutputType = DispatchElementType<DataType>;
  constexpr size_t OutputBytes = static_cast<size_t>(Hidden) * sizeof(OutputType);
  constexpr int NumScales = DataType == DispatchDataType::BF16 ? 0 : Hidden / ScaleBlockSize;
  const int laneId = get_lane_id();
-  const int outputRow = sourceRank * maxTokensPerRank + sourceTokenSlot;

  if (laneId == 0) outputSrcInfo[outputRow] = sourceTokenIdx;
  if (laneId < nTopk) {
@@ -567,17 +571,19 @@ MSCCLPP_DEVICE_INLINE void dispatchRecvWorker(void* output, float* outputScales,
          sourceTokenIdx, nLocalExperts, nRanks, nTopk, maxTokensPerRank, workspaceView, sharedTile, tmaBarrier,
          recvTmaPhase);
    } else {
+      const int outputRow =
+          warpBroadcast(laneId == 0 ? static_cast<int>(outputLayout[sourceRank]) : 0, 0) + sourceTokenSlot;
      hasPendingStore = dispatchRecvTokenMajorOutput<Hidden, DataType, ScaleBlockSize>(
          output, outputScales, outputSrcInfo, outputTopkIdx, outputTopkWeights, payloadView, sourcePayload,
-          localExpertIdx, sourceRank, sourceTokenSlot, sourceTokenIdx, nTopk, maxTokensPerRank, sharedTile, tmaBarrier,
-          recvTmaPhase);
+          localExpertIdx, outputRow, sourceTokenIdx, nTopk, sharedTile, tmaBarrier, recvTmaPhase);
    }
  }

  if (hasPendingStore) waitBulkGroup();
 }

-template <int Hidden, DispatchDataType DataType, int ScaleBlockSize, DispatchLayout Layout>
+template <int Hidden, DispatchDataType DataType, int ScaleBlockSize, DispatchLayout Layout,
+          bool InitializeTokenMajorPadding>
 __global__ __launch_bounds__(DispatchNThreads,
                             1) void dispatchKernel(void* output, float* outputScales, int* outputSrcInfo,
                                                    int* outputTopkIdx, float* outputTopkWeights, int64_t* outputLayout,
@@ -596,12 +602,21 @@ __global__ __launch_bounds__(DispatchNThreads,
  const TransportView transport(comm);
  WorkspaceView workspaceView(workspace, nRanks, nExperts);
  const uint32_t dispatchEpoch = *workspaceView.dispatchEpoch_ + 1;
+  static_assert(!InitializeTokenMajorPadding || Layout == DispatchLayout::TOKEN_MAJOR);

  dispatchSend<Hidden, DataType, ScaleBlockSize>(inputTokens, transport, nExperts, nRanks, topkIndices, topkWeights,
                                                 nTokens, nTopk, maxTokensPerRank, recvBuffer, workspace, dispatchEpoch,
                                                 sharedMem);

  if (static_cast<int>(blockIdx.x) == 0) {
+    if constexpr (InitializeTokenMajorPadding) {
+      const int nMetadataEntries = nRanks * maxTokensPerRank * nTopk;
+      for (int idx = static_cast<int>(threadIdx.x); idx < nMetadataEntries; idx += static_cast<int>(blockDim.x)) {
+        outputTopkIdx[idx] = -1;
+        outputTopkWeights[idx] = 0.0f;
+      }
+      __syncthreads();
+    }
    dispatchRecvScheduler<Layout>(outputLayout, outputCount, transport, nExperts, nRanks, recvBuffer, workspace,
                                  dispatchEpoch, sharedMem);
  } else if (static_cast<int>(blockIdx.x) <= nWorkerBlocks) {
@@ -614,7 +629,8 @@ __global__ __launch_bounds__(DispatchNThreads,
  }
 }

-template <int Hidden, DispatchDataType DataType, int ScaleBlockSize, DispatchLayout Layout>
+template <int Hidden, DispatchDataType DataType, int ScaleBlockSize, DispatchLayout Layout,
+          bool InitializeTokenMajorPadding>
 inline void dispatchHiddenMode(void* output, float* outputScales, int* outputSrcInfo, int* outputTopkIdx,
                               float* outputTopkWeights, int64_t* outputLayout, int* outputCount, const void* input,
                               const int64_t* topkIdx, const float* topkWeights, const low_latency::Workload& workload,
@@ -630,17 +646,18 @@ inline void dispatchHiddenMode(void* output, float* outputScales, int* outputSrc

  const size_t dynamicSharedBytes = dispatchSharedBytes<Hidden, DataType, ScaleBlockSize>(nRanks, nExperts, nTopk);
  static thread_local KernelConfigCache kernelConfig;
-  const int residentBlocks = configureKernel(dispatchKernel<Hidden, DataType, ScaleBlockSize, Layout>, DispatchNThreads,
-                                             dynamicSharedBytes, comm, kernelConfig);
+  const int residentBlocks =
+      configureKernel(dispatchKernel<Hidden, DataType, ScaleBlockSize, Layout, InitializeTokenMajorPadding>,
+                      DispatchNThreads, dynamicSharedBytes, comm, kernelConfig);
  EP_HOST_ASSERT(residentBlocks >= numBlocks);
-  dispatchKernel<Hidden, DataType, ScaleBlockSize, Layout>
+  dispatchKernel<Hidden, DataType, ScaleBlockSize, Layout, InitializeTokenMajorPadding>
      <<<dim3(numBlocks), dim3(DispatchNThreads), dynamicSharedBytes, stream>>>(
          output, outputScales, outputSrcInfo, outputTopkIdx, outputTopkWeights, outputLayout, outputCount, topkIdx,
          topkWeights, input, workload, recvBuffer, comm, workspace);
  CUDA_CHECK(cudaGetLastError());
 }

-template <int Hidden, DispatchLayout Layout>
+template <int Hidden, DispatchLayout Layout, bool InitializeTokenMajorPadding>
 inline void dispatchHidden(void* output, float* outputScales, int* outputSrcInfo, int* outputTopkIdx,
                           float* outputTopkWeights, int64_t* outputLayout, int* outputCount, const void* input,
                           const int64_t* topkIdx, const float* topkWeights, const low_latency::Workload& workload,
@@ -648,11 +665,11 @@ inline void dispatchHidden(void* output, float* outputScales, int* outputSrcInfo
                           cudaStream_t stream) {
  switch (workload.dispatchDataType_) {
    case DispatchDataType::BF16:
-      return dispatchHiddenMode<Hidden, DispatchDataType::BF16, 0, Layout>(
+      return dispatchHiddenMode<Hidden, DispatchDataType::BF16, 0, Layout, InitializeTokenMajorPadding>(
          output, outputScales, outputSrcInfo, outputTopkIdx, outputTopkWeights, outputLayout, outputCount, input,
          topkIdx, topkWeights, workload, recvBuffer, comm, workspace, numBlocks, stream);
    case DispatchDataType::FP8_E4M3:
-      return dispatchHiddenMode<Hidden, DispatchDataType::FP8_E4M3, 128, Layout>(
+      return dispatchHiddenMode<Hidden, DispatchDataType::FP8_E4M3, 128, Layout, InitializeTokenMajorPadding>(
          output, outputScales, outputSrcInfo, outputTopkIdx, outputTopkWeights, outputLayout, outputCount, input,
          topkIdx, topkWeights, workload, recvBuffer, comm, workspace, numBlocks, stream);
    case DispatchDataType::MXFP8_E4M3:
@@ -668,11 +685,16 @@ inline void dispatchLayout(void* output, float* outputScales, int* outputSrcInfo
                           void* recvBuffer, const low_latency::CommContext& comm, void* workspace, int numBlocks,
                           cudaStream_t stream) {
  if (workload.outputLayout_ == DispatchLayout::EXPERT_MAJOR) {
-    return dispatchHidden<Hidden, DispatchLayout::EXPERT_MAJOR>(
+    return dispatchHidden<Hidden, DispatchLayout::EXPERT_MAJOR, false>(
        output, outputScales, outputSrcInfo, outputTopkIdx, outputTopkWeights, outputLayout, outputCount, input,
        topkIdx, topkWeights, workload, recvBuffer, comm, workspace, numBlocks, stream);
  }
-  return dispatchHidden<Hidden, DispatchLayout::TOKEN_MAJOR>(
+  if (workload.initializeTokenMajorPadding_) {
+    return dispatchHidden<Hidden, DispatchLayout::TOKEN_MAJOR, true>(
+        output, outputScales, outputSrcInfo, outputTopkIdx, outputTopkWeights, outputLayout, outputCount, input,
+        topkIdx, topkWeights, workload, recvBuffer, comm, workspace, numBlocks, stream);
+  }
+  return dispatchHidden<Hidden, DispatchLayout::TOKEN_MAJOR, false>(
      output, outputScales, outputSrcInfo, outputTopkIdx, outputTopkWeights, outputLayout, outputCount, input, topkIdx,
      topkWeights, workload, recvBuffer, comm, workspace, numBlocks, stream);
 }
@@ -699,18 +721,18 @@ inline void dispatch(void* output, float* outputScales, int* outputSrcInfo, int*
  EP_HOST_ASSERT(output != nullptr);
  EP_HOST_ASSERT(workload.outputLayout_ == DispatchLayout::EXPERT_MAJOR ||
                 workload.outputLayout_ == DispatchLayout::TOKEN_MAJOR);
+  EP_HOST_ASSERT(!workload.initializeTokenMajorPadding_ || workload.outputLayout_ == DispatchLayout::TOKEN_MAJOR);
  EP_HOST_ASSERT(isSupportedDispatchDataType(workload.dispatchDataType_));
  EP_HOST_ASSERT(workload.dispatchDataType_ == DispatchDataType::BF16 || outputScales != nullptr);
  EP_HOST_ASSERT(outputSrcInfo != nullptr);
  EP_HOST_ASSERT(outputCount != nullptr);
-  if (workload.outputLayout_ == DispatchLayout::EXPERT_MAJOR) {
-    EP_HOST_ASSERT(outputLayout != nullptr);
-  } else {
+  EP_HOST_ASSERT(outputLayout != nullptr);
+  if (workload.outputLayout_ == DispatchLayout::TOKEN_MAJOR) {
    EP_HOST_ASSERT(outputTopkIdx != nullptr);
    EP_HOST_ASSERT(outputTopkWeights != nullptr);
  }
-  EP_HOST_ASSERT(input != nullptr);
-  EP_HOST_ASSERT(topkIdx != nullptr);
+  EP_HOST_ASSERT(workload.numTokens_ == 0 || input != nullptr);
+  EP_HOST_ASSERT(workload.numTokens_ == 0 || topkIdx != nullptr);
  EP_HOST_ASSERT(recvBuffer != nullptr);
  EP_HOST_ASSERT(comm.symmetricBufferBase_ != nullptr);
  EP_HOST_ASSERT(comm.peerMappedBufferBases_ != nullptr);
--- a/src/ext/ep/moe_runtime.cc
+++ b/src/ext/ep/moe_runtime.cc
@@ -17,11 +17,12 @@ namespace mscclpp {
 namespace ep {

 MoERuntime::MoERuntime(mscclpp::Communicator& communicator, int maxTokensPerRank, int hidden, int numExperts,
-                       int numTopk)
+                       int numTopk, bool initializeTokenMajorPadding)
    : rank_(communicator.bootstrap()->getRank()),
      numRanks_(communicator.bootstrap()->getNranks()),
      symmetricBufferBytes_(static_cast<int64_t>(
          low_latency::symmetricBufferSize(maxTokensPerRank, hidden, numRanks_, numExperts, numTopk))),
+      initializeTokenMajorPadding_(initializeTokenMajorPadding),
      communicator_(&communicator) {
  EP_HOST_ASSERT(communicator_ != nullptr);
  EP_HOST_ASSERT(symmetricBufferBytes_ % NUM_BUFFER_ALIGNMENT_BYTES == 0);
@@ -135,6 +136,7 @@ void MoERuntime::dispatch(void* output, float* outputScales, int* outputSrcInfo,
                                       .numExperts_ = numExperts,
                                       .maxTokensPerRank_ = maxTokensPerRank,
                                       .outputLayout_ = dispatchLayout,
+                                       .initializeTokenMajorPadding_ = initializeTokenMajorPadding_,
                                       .dispatchDataType_ = dispatchDataType};
  const size_t workspaceBytes = low_latency::workspaceSize(numRanks_, numExperts);
  EP_HOST_ASSERT(workspaceBytes <= NUM_WORKSPACE_BYTES);
@@ -163,6 +165,7 @@ void MoERuntime::combine(void* output, const void* input, const int64_t* topkIdx
                                       .numExperts_ = numExperts,
                                       .maxTokensPerRank_ = maxTokensPerRank,
                                       .outputLayout_ = dispatchLayout,
+                                       .initializeTokenMajorPadding_ = false,
                                       .dispatchDataType_ = dispatchDataType};
  low_latency::combine(output, input, topkIdx, topkWeights, srcInfo, layoutRange, workload, combineRecvBuffer,
                       dispatchRecvBuffer, commContext_, workspace_, numBlocks, mode, stream);
--- a/src/ext/ep/moe_runtime.hpp
+++ b/src/ext/ep/moe_runtime.hpp
@@ -20,7 +20,8 @@ namespace ep {

 class MoERuntime {
 public:
-  MoERuntime(mscclpp::Communicator& communicator, int maxTokensPerRank, int hidden, int numExperts, int numTopk);
+  MoERuntime(mscclpp::Communicator& communicator, int maxTokensPerRank, int hidden, int numExperts, int numTopk,
+             bool initializeTokenMajorPadding);
  ~MoERuntime() noexcept(false);

  bool isAvailable() const;
@@ -44,6 +45,7 @@ class MoERuntime {
  int numRanksPerIpcDomain_;
  int deviceId_;
  int64_t symmetricBufferBytes_;
+  bool initializeTokenMajorPadding_;
  bool available_ = false;
  void* symmetricBuffer_ = nullptr;
  void* workspace_ = nullptr;
--- a/test/python/ep/ep_bench_ll.py
+++ b/test/python/ep/ep_bench_ll.py
@@ -142,6 +142,11 @@ def parse_args() -> argparse.Namespace:
        default="expert_major",
        help="low-latency dispatch output layout",
    )
+    p.add_argument(
+        "--token-major-init-padding",
+        action="store_true",
+        help="initialize unused token-major top-k IDs and weights for fixed-capacity kernels",
+    )
    p.add_argument("--num-blocks", type=int, default=130, help="total low-latency dispatch blocks")
    p.add_argument(
        "--no-kernel-timing",
@@ -396,6 +401,7 @@ def main() -> None:
        low_latency_num_blocks=args.num_blocks,
        low_latency_combine_mode=combine_mode,
        output_layout=output_layout,
+        token_major_init_padding=args.token_major_init_padding,
        quant=dispatch_quant,
    )
    assert moe_comm.is_available()
--- a/test/python/ep/run_ep_bench.py
+++ b/test/python/ep/run_ep_bench.py
@@ -134,6 +134,11 @@ def parse_args() -> argparse.Namespace:
        default="expert_major",
        help="MSCCL++ Python low-latency output layout",
    )
+    p.add_argument(
+        "--token-major-init-padding",
+        action="store_true",
+        help="initialize token-major padding metadata for fixed-capacity kernels",
+    )
    p.add_argument("--num-blocks", type=int, default=130, help="MSCCL++ low-latency dispatch blocks")

    # Launch / fabric.
@@ -333,6 +338,8 @@ def build_mscclpp_cmd(args: argparse.Namespace) -> str:
        f"--dispatch-dtype {args.dispatch_dtype} --combine-mode {args.combine_mode} "
        f"--output-layout {args.output_layout} --num-blocks {args.num_blocks}"
    )
+    if args.token_major_init_padding:
+        bench_flags += " --token-major-init-padding"
    cupti_build = ""
    extra_exports = ""
    if args.cupti_inproc or args.kernel_only:
--- a/test/python/ep/test_low_latency_multirank.py
+++ b/test/python/ep/test_low_latency_multirank.py
@@ -86,6 +86,11 @@ def parse_args():
        default="expert_major",
        help="Low-latency dispatch output layout",
    )
+    parser.add_argument(
+        "--token-major-init-padding",
+        action="store_true",
+        help="Initialize unused token-major top-k IDs to -1 and weights to zero",
+    )
    parser.add_argument("--bench", action="store_true", help="Run dispatch/combine benchmark after correctness")
    parser.add_argument(
        "--cuda-graph",
@@ -249,6 +254,7 @@ def validate_token_major_dispatch(
    all_topk_weights,
    all_x,
    expected_scales,
+    initialize_padding,
 ):
    assert all_x is not None
    assert dispatch_out.topk_ids is not None
@@ -258,11 +264,23 @@ def validate_token_major_dispatch(
    assert dispatch_out.weights.shape == (num_ranks * num_tokens, num_topk)
    source_token_ids = handle.combine_context.source_token_ids
    assert source_token_ids.shape == (num_ranks * num_tokens,)
+    rank_offsets = handle.combine_context.rank_offsets
+    assert rank_offsets.shape == (num_ranks + 1,)
+    assert dispatch_out.layout.offsets is rank_offsets
+    assert int(rank_offsets[0].item()) == 0
+    total_recv_tokens = int(rank_offsets[-1].item())
+    assert total_recv_tokens == int(packed_recv_count.sum().item())
+    if initialize_padding:
+        assert torch.all(dispatch_out.topk_ids[total_recv_tokens:] == -1)
+        assert torch.all(dispatch_out.weights[total_recv_tokens:] == 0)
    local_expert_begin = rank * num_local_experts
    local_expert_end = local_expert_begin + num_local_experts

    for source_rank in range(num_ranks):
        recv_count = int(packed_recv_count[source_rank].item())
+        row_begin = int(rank_offsets[source_rank].item())
+        row_end = int(rank_offsets[source_rank + 1].item())
+        assert row_end - row_begin == recv_count
        source_routing = all_topk_idx[source_rank]
        expected_source_tokens = (
            ((source_routing >= local_expert_begin) & (source_routing < local_expert_end))
@@ -274,8 +292,6 @@ def validate_token_major_dispatch(
        if recv_count == 0:
            continue

-        row_begin = source_rank * num_tokens
-        row_end = row_begin + recv_count
        source_tokens = source_token_ids[row_begin:row_end].long()
        assert torch.equal(torch.sort(source_tokens).values, expected_source_tokens)

@@ -361,6 +377,7 @@ def reconstruct_token_major_reference(
    group,
 ):
    source_token_ids = handle.combine_context.source_token_ids
+    rank_offsets = handle.combine_context.rank_offsets
    destination_ranks = torch.where(
        all_topk_idx >= 0,
        all_topk_idx // num_local_experts,
@@ -372,8 +389,9 @@ def reconstruct_token_major_reference(
        source_count = int(packed_recv_count[source_rank].item())
        if source_count == 0:
            continue
-        row_begin = source_rank * num_tokens
-        row_end = row_begin + source_count
+        row_begin = int(rank_offsets[source_rank].item())
+        row_end = int(rank_offsets[source_rank + 1].item())
+        assert row_end - row_begin == source_count
        source_tokens = source_token_ids[row_begin:row_end].long()
        selected = first_destination_rank[source_rank, source_tokens] == rank
        dispatched_reference_x[source_rank, source_tokens[selected]] = dequantized_x[row_begin:row_end][selected]
@@ -473,6 +491,7 @@ def main():
        low_latency_num_blocks=args.num_blocks,
        low_latency_combine_mode=combine_mode,
        output_layout=output_layout,
+        token_major_init_padding=args.token_major_init_padding,
        quant=dispatch_quant,
    )
    if rank == 0:
@@ -519,7 +538,8 @@ def main():
    torch.cuda.synchronize()
    print(f"[rank {rank}] post-dispatch", flush=True)
    # expert-major: packed_recv_x [num_local_experts, num_ranks * max_tokens, hidden]
-    # token-major:  packed_recv_x [num_ranks * max_tokens, hidden], rank-grouped
+    # token-major: packed_recv_x has worst-case capacity, with valid rows compacted
+    # into [0 : layout.offsets[-1]).

    # Reference: gather source tokens, routing IDs, and weights from all ranks.
    all_topk_idx = torch.empty((num_ranks, num_tokens, num_topk), dtype=topk_idx.dtype, device="cuda")
@@ -579,6 +599,7 @@ def main():
            all_topk_weights=all_topk_weights,
            all_x=all_x,
            expected_scales=expected_scales,
+            initialize_padding=args.token_major_init_padding,
        )

    if rank == 0: