07 DeepEP Metadata 系统深度分析：Blog 概念 vs 源码实现

Layout Metadata (Where?)          Identity Metadata (Who?)
┌─────────────────────┐          ┌────────────────────────────┐
│ count per expert    │          │ token id (which token?)    │
│ prefix sum / offset │          │ expert id (which expert?)  │
│                     │          │ gate weight (how much?)    │
│ → dst = prefix[e]++ │          │ top-k slot (which slot?)   │
└─────────────────────┘          └────────────────────────────┘
         ↓                                ↓
  写入位置计算                      Combine 语义恢复

EPHandle 字段分类
├── Layout Metadata (Where?)
│   ├── psum_num_recv_tokens_per_scaleup_rank  ← per-rank 前缀和
│   ├── psum_num_recv_tokens_per_expert        ← per-expert 前缀和
│   ├── num_unaligned_recv_tokens_per_expert   ← 实际未对齐计数
│   ├── num_recv_tokens_per_expert_list        ← CPU 侧镜像
│   ├── dst_buffer_slot_idx                    ← 写入位置（核心）
│   ├── token_metadata_at_forward              ← hybrid 前传位置
│   └── channel_linked_list                    ← hybrid 链表
│
├── Identity Metadata (Who?)
│   ├── recv_src_metadata                      ← 源 token 索引 + slot 索引
│   ├── topk_idx                               ← expert 选择
│   └── cached_recv_src_metadata_before_sort   ← 原始顺序备份
│
└── 控制字段
    ├── do_expand, num_experts, expert_alignment
    ├── num_max_tokens_per_rank, num_sms
    ├── num_recv_tokens, num_expanded_tokens
    └── deterministic_sort() 方法

// Notify warps: 每个 warp 处理一部分 token，原子累加到 smem
const auto global_warp_idx = warp_idx * kNumSMs + sm_idx;
for (int i = global_warp_idx; i < num_tokens; i += kNumNotifyWarps * kNumSMs) {
    const auto dst_expert_idx = lane_idx < kNumTopk ?
        static_cast<int>(__ldg(topk_idx + i * kNumTopk + lane_idx)) : -1;
    if (dst_expert_idx >= 0)
        atomicAdd_block(expert_count + dst_expert_idx, 1);

    // Rank 维度需要去重（一个 token 可能多个 topk 落在同一 rank）
    const auto dst_rank_idx = dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerRank : -1;
    if (ptx::deduplicate(dst_rank_idx, lane_idx) and dst_rank_idx >= 0)
        atomicAdd_block(rank_count + dst_rank_idx, 1);
}

// 每个 SM 将局部计数写入 workspace，等待全部 SM 完成后归约
for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) {
    const int64_t counter = (1ll << 32ll) | rank_expert_count[i];
    ptx::red_add(workspace_layout.get_notify_reduction_workspace_ptr() + i, counter);
}
// 编码：高 32 位 = 已到达 SM 计数，低 32 位 = token 计数

// Expert count 对齐到 expert_alignment
for (int i = thread_idx; i < kNumExpertsPerRank; i += kNumNotifyThreads) {
    int sum = 0;
    for (int j = 0; j < kNumRanks; ++ j)
        sum += expert_count[j * kNumExpertsPerRank + i];
    // 先保存未对齐值
    if (num_unaligned_recv_tokens_per_expert != nullptr)
        num_unaligned_recv_tokens_per_expert[i] = sum;
    expert_count[i] = math::align(sum, kExpertAlignment);
}

// Inclusive prefix sum (warp 0 处理 rank)
do_psum(rank_count, psum_num_recv_tokens_per_scaleup_rank, kNumRanks, 0);
// Exclusive prefix sum (warp 1 处理 expert，用于 expand)
do_psum(expert_count, psum_num_recv_tokens_per_expert, kNumExpertsPerRank, 1);

if (not cached_mode) {
    if (do_expand) {
        // expand 模式：slice 出 exclusive 部分，后续 atomicAdd 转 inclusive
        psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert.slice(0, 0, num_local_experts);
    } else {
        // non-expand 模式：slice 出 inclusive 部分（后续不再被 epilogue 使用）
        psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert.slice(0, 1, num_local_experts + 1);
    }
}

flowchart TD
    A[topk_idx: num_tokens x num_topk] --> B[Notify Warps]
    B --> C[per-SM rank_count + expert_count]
    C --> D[Grid Reduction → Workspace]
    D --> E[SM 0 归约 + Expert Alignment]
    E --> F[psum_num_recv_tokens_per_scaleup_rank]
    E --> G[psum_num_recv_tokens_per_expert]
    E --> H[num_unaligned_recv_tokens_per_expert]
    F --> I[Dispatch Warps: 确定接收 token 边界]
    G --> J[Copy Epilogue: expand 模式 atomicAdd 定位]
    I --> K[dst_buffer_slot_idx 生成]
    K --> L[目标 buffer 写入位置]

recv_src_metadata 布局 (kMetadataStride = 2 + topk)
┌────────────────────────────────────────────────────────────────┐
│ [i, 0] │ src_token_global_idx                                 │  ← 全局唯一 token ID (rank * max_tokens + local_idx)
│ [i, 1] │ src_rank_topk_idx (编码 rank_idx * topk + topk_lane) │  ← 联合索引
│ [i, 2] │ topk slot 0 (expand mode 的 dst_tensor_idx)          │  ← expert slot
│ [i, 3] │ topk slot 1                                          │
│ ...    │ ...                                                  │
│ [i, 2+topk-1] │ topk slot (topk-1)                            │
└────────────────────────────────────────────────────────────────┘

// 在 dispatch warps 中，每个 token 写入全局唯一 ID
if (ptx::elect_one_sync())
    *tma_buffer.get_src_token_global_idx_ptr() = rank_idx * kNumMaxTokensPerRank + token_idx;
ptx::tma_store_fence();

constexpr int kMetadataStride = 2 + kNumTopk;
if constexpr (not kCachedMode) {
    if (ptx::elect_one_sync()) {
        // Field 0: 全局源 token 索引
        recv_src_metadata[i * kMetadataStride + 0] = *tma_buffer.get_src_token_global_idx_ptr();
        // Field 1: 联合编码 (rank_idx * topk + master_topk_lane)
        if constexpr (kNumScaleoutRanks == 1) {
            recv_src_metadata[i * kMetadataStride + 1] = current_rank_idx * kNumTopk + master_src_topk_idx;
        } else {
            recv_src_metadata[i * kMetadataStride + 1] = (i - current_rank_start) * kNumTopk + master_src_topk_idx;
        }
    }
    // Fields 2..2+topk-1: expand 模式下的目标 tensor slot
    if (kDoExpand and lane_idx < kNumTopk)
        recv_src_metadata[i * kMetadataStride + 2 + lane_idx] = dst_tensor_idx;
}

constexpr int kMetadataStride = 2 + kNumTopk;
// 解码源 token 索引
const int src_token_idx = __ldg(src_metadata + i * kMetadataStride) % kNumMaxTokensPerRank;
// 解码联合索引 → rank_idx + topk_lane
const int src_rank_topk_idx = __ldg(src_metadata + i * kMetadataStride + 1);
const int src_rank_idx = src_rank_topk_idx / kNumTopk;
const int src_topk_idx = src_rank_topk_idx % kNumTopk;

// 使用 src_rank_idx 定位写入目标
layout::TokenLayout master_token_buffer = [=]() {
    if (nvlink_bypass) {
        auto token_buffer = recv_buffer.get_rank_buffer(...).get_token_buffer(src_token_idx);
        token_buffer.set_base_ptr(gin.get_sym_ptr<team_t>(token_buffer.get_base_ptr(), src_rank_idx));
        return token_buffer;
    }
    return send_buffer.get_rank_buffer(src_rank_idx).get_token_buffer(src_token_idx);
}();

int stored_dst_slot_idx = -1;
if constexpr (kReuseSlotIndices) {
    // Cached dispatch: 从 handle 中复用
    if (lane_idx < kNumTopk)
        stored_dst_slot_idx = __ldg(dst_buffer_slot_idx + token_idx * kNumTopk + lane_idx);
    stored_dst_slot_idx = stored_dst_slot_idx >= 0 ?
        (stored_dst_slot_idx - rank_idx * kNumMaxTokensPerRank) : -1;
} else {
    // 首次 dispatch: atomic 分配 slot
    if (ptx::deduplicate(stored_dst_rank_idx, lane_idx) and stored_dst_rank_idx >= 0)
        stored_dst_slot_idx = atomicAdd(workspace_layout.get_scaleup_atomic_sender_counter() + stored_dst_rank_idx, 1);
    if (lane_idx < kNumTopk) {
        const auto value = stored_dst_slot_idx >= 0 ?
            rank_idx * kNumMaxTokensPerRank + stored_dst_slot_idx : -1;
        dst_buffer_slot_idx[token_idx * kNumTopk + lane_idx] = value;
    }
}

dst_buffer_slot_idx
├── Dispatch 阶段 (作为输出)
│   ├── 写入目标 buffer 位置: recv_buffer.get_token_buffer(stored_dst_slot_idx)
│   └── TMA store 目标: gin.get_sym_ptr(recv_buffer.get_token_buffer(stored_dst_slot_idx), dst_rank)
│
├── Cached dispatch 阶段 (作为输入)
│   └── 复用已有 slot 映射，跳过 atomic 分配（kReuseSlotIndices=true）
│
└── 语义本质
    └── 回答了 Blog 的 "Where?" → 这个 token 的这份 copy 应该写到目标 rank 的哪个 slot

def deterministic_sort(self, do_cpu_sync, is_cached_dispatch,
                       recv_x, recv_sf, recv_topk_idx, recv_topk_weights,
                       channel_linked_list):
    # 关键：只在首次 dispatch 时缓存，后续 cached dispatch 复用
    if not is_cached_dispatch:
        self.cached_recv_src_metadata_before_sort = self.recv_src_metadata.clone()
    sort_keys = self.cached_recv_src_metadata_before_sort[:, 0]  # src_token_global_idx

    # 越界 token 设为 max（排到最后）
    if not do_cpu_sync:
        oob_tokens_mask = torch.arange(0, self.recv_src_metadata.shape[0], ...) >= num_recv_tokens
        sort_keys = sort_keys.clone()
        sort_keys[oob_tokens_mask] = torch.iinfo(sort_keys.dtype).max
    orig_indices = torch.sort(sort_keys).indices

    if not self.do_expand:
        # Non-expand: 按 src_token_global_idx 排序
        permute(recv_x, orig_indices)
        permute(recv_sf, orig_indices)
        permute(recv_topk_weights, orig_indices)
        permute(recv_topk_idx, orig_indices)
        if not is_cached_dispatch:
            permute(self.recv_src_metadata, orig_indices)
        # 更新 channel_linked_list 中的索引
        ...
    elif not is_cached_dispatch:
        # Expand: 按 (expert_idx, src_token_global_idx) 两键排序
        # 排序 key = expert_idx * 1e10 + (-5e9 + src_token_global_idx)
        # 保证 valid token 在 padding 之前，且 valid 按 src_token_global_idx 排序
        expert_token_idx_start = self.psum_num_recv_tokens_per_expert - self.num_unaligned_recv_tokens_per_expert
        token_idx2expert_idx = torch.bucketize(...)
        sort_keys_for_expanded_tensors = token_idx2expert_idx * src_token_global_index_max_x2
        # scatter_add 将 src_token_global_idx 编码进 slot 位置
        ...

排序前（到达顺序不确定）          排序后（确定顺序）
┌──────────────────────┐       ┌──────────────────────┐
│ recv token 0: src=42 │       │ recv token 0: src=7  │  ← 按 src_token_global_idx 升序
│ recv token 1: src=7  │  →    │ recv token 1: src=17 │
│ recv token 2: src=103│       │ recv token 2: src=42 │
│ recv token 3: src=17 │       │ recv token 3: src=103│
└──────────────────────┘       └──────────────────────┘

flowchart TB
    subgraph Dispatch["Dispatch Phase"]
        TOPK[topk_idx<br/>num_tokens × num_topk] -->|per token per k| NOTIFY[Notify Warps<br/>count + dedup]
        NOTIFY -->|rank_count, expert_count| REDUCE[Grid Reduction]
        REDUCE -->|aligned count| PSUM[Prefix Sum<br/>warp 0: rank<br/>warp 1: expert]
        PSUM --> PSRANK[psum_num_recv_tokens<br/>_per_scaleup_rank]
        PSUM --> PSEXPERT[psum_num_recv_tokens<br/>_per_expert]
        PSEXPERT --> UNALIGN[num_unaligned_recv<br/>_tokens_per_expert]

        TOPK -->|per token per k| DISP[Dispatch Warps]
        DISP -->|atomicAdd| SLOT[dst_buffer_slot_idx]
        DISP -->|TMA store| TMA[TMA Buffer<br/>+ src_token_global_idx]
    end

    subgraph CopyEpilogue["Copy Epilogue Phase"]
        TMA -->|TMA load| BUF[Scaleup Buffer]
        PSRANK -->|rank boundary| SRC[rank_idx per token]
        BUF -->|metadata read| META[recv_src_metadata<br/>stride=2+topk]
        PSEXPERT -->|atomicAdd| DSTTENSOR[dst_tensor_idx<br/>for expand]
        META --> OUT[recv_src_metadata]
        DSTTENSOR -->|write| OUT
    end

    subgraph EPHandle["EPHandle (持久化)"]
        PSRANK_H[psum_per_scaleup_rank]
        PSEXPERT_H[psum_per_expert]
        UNALIGN_H[num_unaligned_per_expert]
        SLOT_H[dst_buffer_slot_idx]
        META_H[recv_src_metadata]
        TOPK_H[topk_idx clone]
    end

    subgraph Combine["Combine Phase"]
        META_H -->|decode src_rank_idx<br/>src_token_idx, src_topk_idx| COMB[Combine Warps]
        PSRANK_H -->|num_reduced_tokens| COMB
        TOPK_H -->|combined_topk_idx| COMB
        X[x: expert output] --> COMB
        COMB -->|NVLink/RDMA| RET[Return to source ranks]
    end

    PSUM --> PSRANK_H
    PSEXPERT --> PSEXPERT_H
    UNALIGN --> UNALIGN_H
    SLOT --> SLOT_H
    OUT --> META_H
    TOPK --> TOPK_H

    style Dispatch fill:#e1f5fe
    style CopyEpilogue fill:#fff3e0
    style EPHandle fill:#f3e5f5
    style Combine fill:#e8f5e9

*tma_buffer.get_src_token_global_idx_ptr() = rank_idx * kNumMaxTokensPerRank + token_idx;

const int src_token_idx = __ldg(src_metadata + i * kMetadataStride) % kNumMaxTokensPerRank;

// Non-hybrid: 编码 scaleup rank + master topk lane
recv_src_metadata[i * kMetadataStride + 1] = current_rank_idx * kNumTopk + master_src_topk_idx;
// Hybrid: 编码 slot 偏移 + master topk lane
recv_src_metadata[i * kMetadataStride + 1] = (i - current_rank_start) * kNumTopk + master_src_topk_idx;

const int src_rank_topk_idx = __ldg(src_metadata + i * kMetadataStride + 1);
const int src_rank_idx = src_rank_topk_idx / kNumTopk;
const int src_topk_idx = src_rank_topk_idx % kNumTopk;

struct TokenLayout {
    // Metadata 字节数: topk indices + weights + optional (src_token + linked_list)
    int num_metadata_bytes = num_topk * (sizeof(int) + sizeof(float)) +
                             (with_metadata ? (1 + num_topk) * sizeof(int) : 0);

    int* get_topk_idx_ptr() const { return get_metadata_ptr(); }
    float* get_topk_weights_ptr() const { return advance_ptr<float>(get_metadata_ptr(), num_topk * sizeof(int)); }
    int* get_src_token_global_idx_ptr() const { return advance_ptr<int>(get_topk_weights_ptr(), num_topk * sizeof(float)); }
    int* get_linked_list_idx_ptr() const { return get_src_token_global_idx_ptr() + 1; }
};

return {recv_x, recv_sf,
        recv_topk_idx, recv_topk_weights,
        copied_topk_idx,
        num_recv_tokens, num_expanded_tokens,
        num_recv_tokens_per_expert_list,
        psum_num_recv_tokens_per_scaleup_rank,   // → EPHandle.psum_num_recv_tokens_per_scaleup_rank
        psum_num_recv_tokens_per_expert,         // → EPHandle.psum_num_recv_tokens_per_expert
        num_unaligned_recv_tokens_per_expert,    // → EPHandle.num_unaligned_recv_tokens_per_expert
        recv_src_metadata,                       // → EPHandle.recv_src_metadata
        dst_buffer_slot_idx,                     // → EPHandle.dst_buffer_slot_idx
        token_metadata_at_forward,               // → EPHandle.token_metadata_at_forward
        channel_linked_list,                     // → EPHandle.channel_linked_list
        event};

combined_x, combined_topk_weights, event = \
    self.runtime.combine(x, topk_weights,
                         bias_0, bias_1,
                         handle.recv_src_metadata,                         # Identity
                         handle.topk_idx,                                   # Identity
                         handle.psum_num_recv_tokens_per_scaleup_rank,      # Layout
                         handle.token_metadata_at_forward,                   # Layout (hybrid)
                         handle.channel_linked_list,                         # Layout (hybrid)
                         handle.num_experts,
                         handle.num_max_tokens_per_rank,
                         ...)

Blog 概念模型准确性评估
├── 概念层面: ★★★★★ (5/5)
│   └── Layout/Identity 二分法准确抓住本质
├── 字段覆盖: ★★★☆☆ (3/5)
│   └── 遗漏 dst_buffer_slot_idx、deterministic_sort 等重要字段
├── 流程描述: ★★★☆☆ (3/5)
│   └── "count → prefix sum → scatter" 正确但过度简化
├── 实现细节: ★★☆☆☆ (2/5)
│   └── 编码方式、联合索引、去重机制未涉及
└── 总体: 概念模型正确，可作为理解入口，但不足以指导实现

Blog 描述:                    DeepEP 实际:
prefix[expert]++              atomicAdd(atomic_counter[rank], 1)
                              ↓
                              dst_buffer_slot_idx[token][k] = rank * max_tokens + slot
                              ↓
                              recv_buffer[slot] = data  (间接寻址)

Blog: Identity 直接映射       Expert ID 在 topk_idx 中
                              Token ID 在 recv_src_metadata[:, 0]
                              Gate weight 在 recv_topk_weights 中
                              三者分离存储，combine 时联合解码

Layout Metadata 回答: "数据应该去哪？"
  → 所有 prefix sum / slot / offset 字段服务于此
  → 与 token 身份无关，只与数量和位置有关

Identity Metadata 回答: "这份数据是谁？"
  → src_token_global_idx / topk slot 服务于此
  → 与路由无关，只与 token 身份有关

Blog 概念                     DeepEP 实现
─────────                    ──────────
Layout (count/psum/offset) → 4 张量 + hybrid 扩展 + alignment 处理
Identity (token/expert/weight/slot) → 联合编码 3 段 + topk_idx 分离
Counting Sort              → notify warps + grid reduction + psum
Combine 语义恢复            → + deterministic_sort + cached dispatch

假设: T=num_tokens, R=num_ranks, RU=num_scaleup_ranks, E=num_experts,
      EL=num_local_experts=E/R, K=num_topk, M=num_max_tokens_per_rank

生成阶段:
  rank_count:             [RU]          (smmem)
  expert_count:           [E]           (smmem)
  psum_per_scaleup_rank:  [RU]          ← EPHandle 字段
  psum_per_expert:        [EL+1] → [EL] ← EPHandle 字段 (slice 后)
  num_unaligned_per_expert: [EL]        ← EPHandle 字段
  dst_buffer_slot_idx:    [T, K]        ← EPHandle 字段
  src_token_global_idx:   scalar/TMA    ← 写入 per-token metadata
  recv_src_metadata:      [num_recv, 2+K] ← EPHandle 字段

消费阶段:
  combine 读取:
    recv_src_metadata[i, 0]     → src_token_global_idx
    recv_src_metadata[i, 1]     → src_rank_topk_idx
    recv_src_metadata[i, 2+k]   → expand slot (或 -1)
    psum_per_scaleup_rank[-1]   → num_reduced_tokens