How DeepEP Controls Streaming Multiprocessors (SM) for Kernel Execution

DeepEP controls SM utilization through a configurable num_sms parameter that propagates from Python's Buffer.set_num_sms() through the C++ runtime to CUDA kernel launch configurations, where it directly determines grid dimensions.

Deep EP (Deep Expert-Parallel) provides fine-grained control over GPU Streaming Multiprocessor (SM) allocation to optimize expert-parallel communication workloads. By adjusting the number of active SMs, developers can balance parallelism against memory bandwidth contention. This article examines the complete flow from Python API to CUDA kernel execution in the deepseek-ai/DeepEP repository, referencing specific source files and line numbers.

Python API Configuration

Setting SM Count with Buffer.set_num_sms()

The entry point for SM control is the static method Buffer.set_num_sms() defined in deep_ep/buffer.py at lines 154-163. This method stores the desired SM count in the class attribute Buffer.num_sms, which defaults to 20.

import deep_ep

# Configure to use 24 SMs

deep_ep.Buffer.set_num_sms(24)

This value persists across the Python session and influences all subsequently created buffers.

Config Object Construction

When instantiating a deep_ep.Config object (lines 202-207 in deep_ep/buffer.py), the current value of Buffer.num_sms is automatically injected into the configuration. The Config object acts as the Python-side container that bridges to the C++ backend.

config = deep_ep.Config(num_sms, nvl_chunk, nvl_buffer)

This Python object is then passed to the C++ runtime (deep_ep_cpp.Buffer) during buffer initialization, transferring the SM requirement across the language boundary.

C++ Runtime Bridge

The Config Struct

The C++ layer defines the configuration structure in csrc/config.hpp (lines 23-31). The Config struct holds the SM count as an integer field that mirrors the Python configuration.

struct Config {
    int num_sms;
    // ... additional fields
};

Cross-Language Data Flow

When the Python deep_ep.Config object crosses into the C++ API, the num_sms value is copied into this struct and stored inside the runtime's deep_ep_cpp::Buffer instance. This ensures the SM preference persists throughout the buffer's lifetime and is available for all subsequent kernel launches.

CUDA Kernel Launch Configuration

SETUP_LAUNCH_CONFIG Macro

DeepEP kernels use the SETUP_LAUNCH_CONFIG macro (defined in csrc/kernels/api.cuh) to translate the stored num_sms value into CUDA grid dimensions. The macro expands to:

dim3 grid(num_sms);
dim3 block(kNumThreads);
cudaLaunchKernel(..., grid, block, …);

This creates a one-dimensional grid where the x-dimension equals the user-specified SM count.

Kernel Implementations

Each high-throughput kernel extracts the SM count from gridDim.x to coordinate parallel execution:

• Intranode kernels (csrc/kernels/intranode.cu, lines 223-227):
  const auto num_sms = static_cast<int>(gridDim.x);

• Internode kernels (csrc/kernels/internode.cu, lines 484-488):
  Same extraction pattern, where each block maps to one SM (or two SMs when num_sms is even).

• Low-latency internode kernels (csrc/kernels/internode_ll.cu, lines 160-164):
  Again uses gridDim.x to determine the active SM count for RDMA operations.

In all cases, the number of SMs equals the grid-dimension X used to launch the kernel, ensuring the GPU executes exactly the requested number of blocks across available SMs.

Performance Tuning and Constraints

Why Control SM Count?

Tuning num_sms allows developers to optimize for specific hardware characteristics:

• Higher parallelism: More SMs increase concurrent execution but may increase contention for NVLink or RDMA bandwidth.
• Memory-bound workloads: Fewer SMs can improve per-SM throughput when memory traffic dominates computation.

Hardware Constraints and Assertions

The codebase enforces that num_sms must be even. In csrc/kernels/intranode.cu at line 226, an assertion checks num_sms % 2 == 0 because many kernels split work into pairs of SMs for channel-wise processing. Odd values will trigger a runtime error.

Automated Testing Patterns

The test suite demonstrates automated SM tuning. In tests/test_intranode.py (lines 205-213) and tests/test_internode.py (lines 251-259), the code iterates over candidate values (e.g., 12, 16, 24) and records best runtimes:

for num_sms in [12, 16, 24]:
    deep_ep.Buffer.set_num_sms(num_sms)
    # ... run benchmark

This pattern allows empirical discovery of the optimal SM count for specific network configurations.

End-to-End Usage Example

The following complete example demonstrates configuring SM count before executing a dispatch operation:

import deep_ep
import torch
import torch.distributed as dist

# Initialize distributed environment

dist.init_process_group(backend='nccl')

# 1. Configure SM count

deep_ep.Buffer.set_num_sms(24)

# 2. Create buffer (automatically builds Config with num_sms=24)

buf = deep_ep.Buffer(
    group=dist.group.WORLD,
    num_nvl_bytes=4 * 1024**2,    # 4 MiB NVLink buffer

    num_rdma_bytes=8 * 1024**2,   # 8 MiB RDMA buffer

)

# 3. Prepare input tensor

x = torch.randn(1024, 4096, dtype=torch.bfloat16, device='cuda')

# 4. Execute dispatch using configured SMs

out = buf.dispatch(x)

When dispatch() executes, it launches CUDA kernels with gridDim.x = 24, restricting execution to 24 Streaming Multiprocessors.

Summary

• DeepEP controls SM execution through the num_sms parameter, configurable via Buffer.set_num_sms().
• Configuration flows from Python (deep_ep/buffer.py) through the C++ struct (csrc/config.hpp) to CUDA kernels.
• Grid dimensions explicitly match the SM count via SETUP_LAUNCH_CONFIG, with gridDim.x representing the number of SMs.
• Hardware constraints require even values for num_sms to support paired SM processing.
• Testing utilities in test_intranode.py and test_internode.py demonstrate automated benchmarking across different SM counts.

Frequently Asked Questions

How do I configure the SM count in DeepEP?

Call deep_ep.Buffer.set_num_sms(n) before creating your buffer, where n is your desired number of SMs (must be even). This value is automatically picked up when constructing deep_ep.Config and passed to the C++ runtime.

Why must num_sms be even?

DeepEP kernels partition work across pairs of SMs for channel-wise processing. The code asserts num_sms % 2 == 0 in csrc/kernels/intranode.cu to ensure proper workload distribution. Odd values will cause a runtime assertion failure.

What files handle the SM configuration?

The configuration travels through three key files: deep_ep/buffer.py (Python API), csrc/config.hpp (C++ struct definition), and the kernel files like csrc/kernels/intranode.cu and csrc/kernels/internode.cu where gridDim.x is evaluated.

How does SM count affect NVLink performance?

Increasing num_sms raises parallelism but can exacerbate contention on NVLink and RDMA bandwidth. Conversely, reducing SM count may improve per-SM throughput for memory-bound operations. The test files demonstrate benchmarking multiple values to find the optimal balance for your specific hardware topology.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}