Torch Profiling with DeepEP PD Disaggregation

[liz-badada]

Motivation

Support torch profiling when enable SGLang PD Disaggregation. Will Generate prefill & decode profile separately.

Test based on #5435

Modifications

Add start_profile / stop_profile request to mini_lb.py

Examples (TP16 + DP16 Attn + DeepEP on H100)

Prepare ./pd_node*.json first.

Server

SGLANG_TORCH_PROFILER_DIR=./profile_log MOONCAKE_CONFIG_PATH=./pd_node0.json python -m sglang.launch_server --model-path ~/.cache/huggingface/hub/models--chwan--DeepSeek-V3-5layer/snapshots/38a0c0ee55158e7d2ac9a6af1de94c4dfe084872/ --disaggregation-mode prefill --host 10.10.38.7 --port 30000 --trust-remote-code --dist-init-addr 10.10.38.7:5000 --nnodes 2 --node-rank 0 --tp-size 16 --dp-size 16 --enable-dp-attention --enable-deepep-moe --deepep-mode normal --mem-fraction-static 0.8

SGLANG_TORCH_PROFILER_DIR=./profile_log MOONCAKE_CONFIG_PATH=./pd_node1.json python -m sglang.launch_server --model-path ~/.cache/huggingface/hub/models--chwan--DeepSeek-V3-5layer/snapshots/38a0c0ee55158e7d2ac9a6af1de94c4dfe084872/ --disaggregation-mode prefill --host 10.10.38.8 --port 30000 --trust-remote-code --dist-init-addr 10.10.38.7:5000 --nnodes 2 --node-rank 1 --tp-size 16 --dp-size 16 --enable-dp-attention --enable-deepep-moe --deepep-mode normal --mem-fraction-static 0.8

SGLANG_TORCH_PROFILER_DIR=./profile_log MOONCAKE_CONFIG_PATH=./pd_node2.json python -m sglang.launch_server --model-path ~/.cache/huggingface/hub/models--chwan--DeepSeek-V3-5layer/snapshots/38a0c0ee55158e7d2ac9a6af1de94c4dfe084872/ --disaggregation-mode decode --host 10.10.38.10 --port 30001 --trust-remote-code --dist-init-addr 10.10.38.10:5000 --nnodes 2 --node-rank 0 --tp-size 16 --dp-size 16 --enable-dp-attention --enable-deepep-moe --deepep-mode low_latency --mem-fraction-static 0.8 --cuda-graph-max-bs 128 --max-running-requests 128

SGLANG_TORCH_PROFILER_DIR=./profile_log MOONCAKE_CONFIG_PATH=./pd_node3.json python -m sglang.launch_server --model-path ~/.cache/huggingface/hub/models--chwan--DeepSeek-V3-5layer/snapshots/38a0c0ee55158e7d2ac9a6af1de94c4dfe084872/ --disaggregation-mode decode --host 10.10.38.4 --port 30001 --trust-remote-code --dist-init-addr 10.10.38.10:5000 --nnodes 2 --node-rank 1 --tp-size 16 --dp-size 16 --enable-dp-attention --enable-deepep-moe --deepep-mode low_latency --mem-fraction-static 0.8 --cuda-graph-max-bs 128 --max-running-requests 128

Proxy

SGLANG_TORCH_PROFILER_DIR=./profile_log python3 -m sglang.srt.disaggregation.mini_lb --prefill http://10.10.38.7:30000/ --decode http://10.10.38.10:30001/ --host 0.0.0.0 --port 8000

Bench with Profiling

SGLANG_TORCH_PROFILER_DIR=./profile_log python -m sglang.bench_serving --backend sglang --model ~/.cache/huggingface/hub/models--chwan--DeepSeek-V3-5layer/snapshots/38a0c0ee55158e7d2ac9a6af1de94c4dfe084872/ --host 127.0.0.1 --port 8000 --num-prompts 20 --dataset-name random --random-input-len 1000 --random-output-len 4 --profile

event_loop_normal_disagg_prefill

event_loop_normal_disagg_decode

Note

Known bug from PyTroch Profiler UTF-8 decode: pytorch/pytorch#64345

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[whybeyoung]

LGTM

[merrymercy]

please fix the lint

[ch-wan]

Is changing the output directory necessary?

[ch-wan]

Creating profile_dir is not the job of lb.