[GH-ISSUE #760] [Performance] DeepEP V2 combine epilogue low batch performance #417

New issue

Open

opened 2026-10-26 10:00:20 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 10:00:20 +00:00

Owner

Originally created by @elvircrn on GitHub (Sep 22, 2026).
Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/760

While doing low batch (less than 8 concurrent requests) DP4/TP8 P/D deployments with Kimi K3 on H200, I noticed that decode traces have straggler ranks, spending a long time in combine_reduce_epilogue_impl:

Because this kernel assigns one warp per output token, it is underperforming for cases where we have a handful of output tokens per dp rank.

Profiling this kernel further points to this being the case:

Please let me know if my analysis is sound here and if I can provide more profiling data to help diagnose this further. I have two thoughts on this, provided findings are sound:

specialize the combine_reduce_epilogue_impl so that it assigns multiple warps to a single output token
expose a skip_epilogue flag and let the library consumer take over the epilogue work. This would have a benefit of allowing fusions with downstream ops

I have both versions of this running on my DeepEP fork and the corresponding vllm fork.

Reproduction

What follows are some reproduction steps:

repro_dp4_tp8_combine_only.py

#!/usr/bin/env python3
"""DP4 x TP8 DeepEP combine-only repro.

One real dispatch establishes the production EPHandle and receive layout.  The
measured loop then repeatedly calls only DeepEP v2 combine with the same
received metadata and fake expert output.  This removes dispatch, CUDA graph,
PyTorch profiler, and routing variability from the combine epilogue.
"""

from __future__ import annotations

import os

import torch
import torch.distributed as dist

import deep_ep


TP_SIZE = 8
DP_SIZE = 4
WORLD_SIZE = 32
HIDDEN = 3584
NUM_EXPERTS = 896
TOPK = 16
NUM_MAX_TOKENS_PER_RANK = 1
NUM_SMS = 132
NUM_QPS = 0
WARMUP_COMBINES = 10
MEASURED_COMBINES = 20


def main() -> None:
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    local_rank = int(os.environ.get("LOCAL_RANK", rank % TP_SIZE))
    if world_size != WORLD_SIZE:
        raise RuntimeError(f"expected WORLD_SIZE={WORLD_SIZE}, got {world_size}")
    torch.cuda.set_device(local_rank)
    dist.init_process_group("nccl")
    device = torch.device(f"cuda:{local_rank}")

    tp_rank = rank % TP_SIZE
    dp_rank = rank // TP_SIZE
    local_experts = NUM_EXPERTS // WORLD_SIZE

    # Exact vLLM TP8/SP one-token layout: one physical row per TP rank, with
    # only TP rank 0 real and TP ranks 1..7 carrying -1 padding routes.
    x = torch.zeros((1, HIDDEN), dtype=torch.bfloat16, device=device)
    topk_idx = torch.full((1, TOPK), -1, dtype=deep_ep.topk_idx_t, device=device)
    topk_weights = torch.zeros((1, TOPK), dtype=torch.float32, device=device)
    if tp_rank == 0:
        topk_idx[0] = (
            torch.arange(TOPK, device=device, dtype=torch.int64) * local_experts
        ).to(deep_ep.topk_idx_t)
        topk_weights[0].fill_(1.0 / TOPK)
        x.fill_(float(dp_rank + 1))

    buffer = deep_ep.ElasticBuffer(
        group=dist.group.WORLD,
        num_max_tokens_per_rank=NUM_MAX_TOKENS_PER_RANK,
        hidden=HIDDEN,
        num_topk=TOPK,
        allow_hybrid_mode=False,
        allow_multiple_reduction=False,
        sl_idx=3,
        explicitly_destroy=True,
    )

    # One production dispatch only: this supplies the real handle, receive
    # metadata, communication slots, and token layout consumed by combine.
    recv_x, recv_topk_idx, recv_topk_weights, handle, event = buffer.dispatch(
        x=x,
        topk_idx=topk_idx,
        topk_weights=topk_weights,
        num_experts=NUM_EXPERTS,
        num_max_tokens_per_rank=NUM_MAX_TOKENS_PER_RANK,
        num_sms=NUM_SMS,
        num_qps=NUM_QPS,
        do_expand=False,
        do_cpu_sync=False,
        async_with_compute_stream=False,
    )
    if event.event is not None:
        event.current_stream_wait()
    if isinstance(recv_x, tuple) or recv_topk_weights is None:
        raise RuntimeError("expected BF16 dispatch with received top-k weights")

    # Fake expert/GEMM output exactly as in the full repro: each valid received
    # row contains its valid local gate-weight sum; padded rows are zero.
    valid_weight = recv_topk_weights.masked_fill(
        recv_topk_idx < 0, 0
    ).sum(dim=1, keepdim=True)
    recv_x.copy_(valid_weight.to(recv_x.dtype))
    valid_recv_rows = int((recv_topk_idx >= 0).any(dim=1).sum().item())
    valid_recv_routes = int((recv_topk_idx >= 0).sum().item())
    torch.cuda.synchronize()
    dist.barrier()

    combined_x = None
    for _ in range(WARMUP_COMBINES):
        combined_x, _, event = buffer.combine(
            x=recv_x,
            handle=handle,
            topk_weights=None,
            num_sms=NUM_SMS,
            num_qps=NUM_QPS,
            async_with_compute_stream=False,
        )
        if event.event is not None:
            event.current_stream_wait()
        torch.cuda.synchronize()
    dist.barrier()

    comm_stream = buffer.get_comm_stream()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    for _ in range(MEASURED_COMBINES):
        with torch.cuda.stream(comm_stream):
            start.record()
        combined_x, _, event = buffer.combine(
            x=recv_x,
            handle=handle,
            topk_weights=None,
            num_sms=NUM_SMS,
            num_qps=NUM_QPS,
            async_with_compute_stream=False,
        )
        if event.event is not None:
            event.current_stream_wait()
        with torch.cuda.stream(comm_stream):
            end.record()
        end.synchronize()

    expected = 1 if tp_rank == 0 else 0
    error = (combined_x - expected).abs().max().to(torch.float32)
    dist.all_reduce(error, op=dist.ReduceOp.MAX)
    elapsed_us = start.elapsed_time(end) * 1000.0
    print(
        f"rank={rank} dp={dp_rank} tp={tp_rank} "
        f"received_rows={recv_x.shape[0]} valid_recv_rows={valid_recv_rows} "
        f"valid_recv_routes={valid_recv_routes} combined_rows={combined_x.shape[0]} "
        f"last_combine_us={elapsed_us:.2f} max_error={error.item():.6f}",
        flush=True,
    )

    buffer.destroy()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()

Repro script for ncu profiler:

#!/usr/bin/env bash
set -euo pipefail

if [[ "${RANK}" == "8" ]]; then
    exec /opt/nvidia/nsight-compute/2025.4.1/ncu \
        --metrics "gpu__time_duration.sum,sm__cycles_elapsed.avg,sm__cycles_active.avg,sm__warps_active.avg.pct_of_peak_sustained_active,launch__occupancy_limit_shared_mem,launch__occupancy_limit_registers,smsp__warp_issue_stalled_barrier_per_warp_active,smsp__warp_issue_stalled_branch_resolving_per_warp_active,smsp__warp_issue_stalled_dispatch_stall_per_warp_active,smsp__warp_issue_stalled_drain_per_warp_active,smsp__warp_issue_stalled_gmma_per_warp_active,smsp__warp_issue_stalled_imc_miss_per_warp_active,smsp__warp_issue_stalled_lg_throttle_per_warp_active,smsp__warp_issue_stalled_long_scoreboard_per_warp_active,smsp__warp_issue_stalled_math_pipe_throttle_per_warp_active,smsp__warp_issue_stalled_membar_per_warp_active,smsp__warp_issue_stalled_mio_throttle_per_warp_active,smsp__warp_issue_stalled_misc_per_warp_active,smsp__warp_issue_stalled_no_instruction_per_warp_active,smsp__warp_issue_stalled_not_selected_per_warp_active,smsp__warp_issue_stalled_selected_per_warp_active,smsp__warp_issue_stalled_short_scoreboard_per_warp_active,smsp__warp_issue_stalled_sleeping_per_warp_active,smsp__warp_issue_stalled_tex_throttle_per_warp_active,smsp__warp_issue_stalled_wait_per_warp_active" \
        --kernel-name 'regex:.*combine_reduce_epilogue_impl.*' \
        --launch-skip 5 \
        --launch-count 1 \
        --replay-mode kernel \
        --cache-control none \
        --clock-control none \
        --force-overwrite \
        --export /tmp/deepep-ncu/combine-rank8-stalls \
        /usr/bin/python3 /tmp/repro_dp4_tp8_combine_only.py
fi

exec /usr/bin/python3 /tmp/repro_dp4_tp8_combine_only.py

Originally created by @elvircrn on GitHub (Sep 22, 2026). Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/760 While doing low batch (less than 8 concurrent requests) DP4/TP8 P/D deployments with Kimi K3 on H200, I noticed that decode traces have straggler ranks, spending a long time in `combine_reduce_epilogue_impl`: <img width="3248" height="2122" alt="Image" src="https://github.com/user-attachments/assets/b0660025-01ea-4dc1-be21-0c23544ea445" /> Because this kernel assigns one warp per output token, it is underperforming for cases where we have a handful of output tokens per dp rank. Profiling this kernel further points to this being the case: <img width="1512" height="839" alt="Image" src="https://github.com/user-attachments/assets/a48aebc0-1ad1-46d3-a9d6-9b46a8d87e27" /> Please let me know if my analysis is sound here and if I can provide more profiling data to help diagnose this further. I have two thoughts on this, provided findings are sound: 1. specialize the `combine_reduce_epilogue_impl` so that it assigns multiple warps to a single output token 2. expose a `skip_epilogue` flag and let the library consumer take over the epilogue work. This would have a benefit of allowing fusions with downstream ops I have both versions of this running on my DeepEP [fork](https://github.com/elvircrn/DeepEP/tree/combine_epilogue_small_batch) and the corresponding vllm [fork](https://github.com/elvircrn/vllm/tree/deepep-triton-epilogue). ## Reproduction What follows are some reproduction steps: `repro_dp4_tp8_combine_only.py` ``` #!/usr/bin/env python3 """DP4 x TP8 DeepEP combine-only repro. One real dispatch establishes the production EPHandle and receive layout. The measured loop then repeatedly calls only DeepEP v2 combine with the same received metadata and fake expert output. This removes dispatch, CUDA graph, PyTorch profiler, and routing variability from the combine epilogue. """ from __future__ import annotations import os import torch import torch.distributed as dist import deep_ep TP_SIZE = 8 DP_SIZE = 4 WORLD_SIZE = 32 HIDDEN = 3584 NUM_EXPERTS = 896 TOPK = 16 NUM_MAX_TOKENS_PER_RANK = 1 NUM_SMS = 132 NUM_QPS = 0 WARMUP_COMBINES = 10 MEASURED_COMBINES = 20 def main() -> None: rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) local_rank = int(os.environ.get("LOCAL_RANK", rank % TP_SIZE)) if world_size != WORLD_SIZE: raise RuntimeError(f"expected WORLD_SIZE={WORLD_SIZE}, got {world_size}") torch.cuda.set_device(local_rank) dist.init_process_group("nccl") device = torch.device(f"cuda:{local_rank}") tp_rank = rank % TP_SIZE dp_rank = rank // TP_SIZE local_experts = NUM_EXPERTS // WORLD_SIZE # Exact vLLM TP8/SP one-token layout: one physical row per TP rank, with # only TP rank 0 real and TP ranks 1..7 carrying -1 padding routes. x = torch.zeros((1, HIDDEN), dtype=torch.bfloat16, device=device) topk_idx = torch.full((1, TOPK), -1, dtype=deep_ep.topk_idx_t, device=device) topk_weights = torch.zeros((1, TOPK), dtype=torch.float32, device=device) if tp_rank == 0: topk_idx[0] = ( torch.arange(TOPK, device=device, dtype=torch.int64) * local_experts ).to(deep_ep.topk_idx_t) topk_weights[0].fill_(1.0 / TOPK) x.fill_(float(dp_rank + 1)) buffer = deep_ep.ElasticBuffer( group=dist.group.WORLD, num_max_tokens_per_rank=NUM_MAX_TOKENS_PER_RANK, hidden=HIDDEN, num_topk=TOPK, allow_hybrid_mode=False, allow_multiple_reduction=False, sl_idx=3, explicitly_destroy=True, ) # One production dispatch only: this supplies the real handle, receive # metadata, communication slots, and token layout consumed by combine. recv_x, recv_topk_idx, recv_topk_weights, handle, event = buffer.dispatch( x=x, topk_idx=topk_idx, topk_weights=topk_weights, num_experts=NUM_EXPERTS, num_max_tokens_per_rank=NUM_MAX_TOKENS_PER_RANK, num_sms=NUM_SMS, num_qps=NUM_QPS, do_expand=False, do_cpu_sync=False, async_with_compute_stream=False, ) if event.event is not None: event.current_stream_wait() if isinstance(recv_x, tuple) or recv_topk_weights is None: raise RuntimeError("expected BF16 dispatch with received top-k weights") # Fake expert/GEMM output exactly as in the full repro: each valid received # row contains its valid local gate-weight sum; padded rows are zero. valid_weight = recv_topk_weights.masked_fill( recv_topk_idx < 0, 0 ).sum(dim=1, keepdim=True) recv_x.copy_(valid_weight.to(recv_x.dtype)) valid_recv_rows = int((recv_topk_idx >= 0).any(dim=1).sum().item()) valid_recv_routes = int((recv_topk_idx >= 0).sum().item()) torch.cuda.synchronize() dist.barrier() combined_x = None for _ in range(WARMUP_COMBINES): combined_x, _, event = buffer.combine( x=recv_x, handle=handle, topk_weights=None, num_sms=NUM_SMS, num_qps=NUM_QPS, async_with_compute_stream=False, ) if event.event is not None: event.current_stream_wait() torch.cuda.synchronize() dist.barrier() comm_stream = buffer.get_comm_stream() start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) for _ in range(MEASURED_COMBINES): with torch.cuda.stream(comm_stream): start.record() combined_x, _, event = buffer.combine( x=recv_x, handle=handle, topk_weights=None, num_sms=NUM_SMS, num_qps=NUM_QPS, async_with_compute_stream=False, ) if event.event is not None: event.current_stream_wait() with torch.cuda.stream(comm_stream): end.record() end.synchronize() expected = 1 if tp_rank == 0 else 0 error = (combined_x - expected).abs().max().to(torch.float32) dist.all_reduce(error, op=dist.ReduceOp.MAX) elapsed_us = start.elapsed_time(end) * 1000.0 print( f"rank={rank} dp={dp_rank} tp={tp_rank} " f"received_rows={recv_x.shape[0]} valid_recv_rows={valid_recv_rows} " f"valid_recv_routes={valid_recv_routes} combined_rows={combined_x.shape[0]} " f"last_combine_us={elapsed_us:.2f} max_error={error.item():.6f}", flush=True, ) buffer.destroy() dist.destroy_process_group() if __name__ == "__main__": main() ``` Repro script for `ncu` profiler: ``` #!/usr/bin/env bash set -euo pipefail if [[ "${RANK}" == "8" ]]; then exec /opt/nvidia/nsight-compute/2025.4.1/ncu \ --metrics "gpu__time_duration.sum,sm__cycles_elapsed.avg,sm__cycles_active.avg,sm__warps_active.avg.pct_of_peak_sustained_active,launch__occupancy_limit_shared_mem,launch__occupancy_limit_registers,smsp__warp_issue_stalled_barrier_per_warp_active,smsp__warp_issue_stalled_branch_resolving_per_warp_active,smsp__warp_issue_stalled_dispatch_stall_per_warp_active,smsp__warp_issue_stalled_drain_per_warp_active,smsp__warp_issue_stalled_gmma_per_warp_active,smsp__warp_issue_stalled_imc_miss_per_warp_active,smsp__warp_issue_stalled_lg_throttle_per_warp_active,smsp__warp_issue_stalled_long_scoreboard_per_warp_active,smsp__warp_issue_stalled_math_pipe_throttle_per_warp_active,smsp__warp_issue_stalled_membar_per_warp_active,smsp__warp_issue_stalled_mio_throttle_per_warp_active,smsp__warp_issue_stalled_misc_per_warp_active,smsp__warp_issue_stalled_no_instruction_per_warp_active,smsp__warp_issue_stalled_not_selected_per_warp_active,smsp__warp_issue_stalled_selected_per_warp_active,smsp__warp_issue_stalled_short_scoreboard_per_warp_active,smsp__warp_issue_stalled_sleeping_per_warp_active,smsp__warp_issue_stalled_tex_throttle_per_warp_active,smsp__warp_issue_stalled_wait_per_warp_active" \ --kernel-name 'regex:.*combine_reduce_epilogue_impl.*' \ --launch-skip 5 \ --launch-count 1 \ --replay-mode kernel \ --cache-control none \ --clock-control none \ --force-overwrite \ --export /tmp/deepep-ncu/combine-rank8-stalls \ /usr/bin/python3 /tmp/repro_dp4_tp8_combine_only.py fi exec /usr/bin/python3 /tmp/repro_dp4_tp8_combine_only.py ```

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

deepseek-ai/DeepEP#417

No description provided.

Rows
Columns