Stage-Separated Profiling with Warm KV Cache Support

[TerrenceZhangX]

Summary

Add a complete stage-separated profiling pipeline (run_stage_profile.py) that captures prefill (EXTEND) and decode traces independently from SGLang inference servers. Includes cross-rank kernel aggregation, kernel shape merging, and comprehensive tests.

Motivation

1. Decouple prefill and decode to reduce profiling cost

The existing run_profile.py captures prefill+decode in one fused trace, so every (prefill_len, decode_len) combination needs its own run. For N prefill lengths × M decode lengths that's N×M profiling passes. Stage-separated profiling captures them independently — N+M runs instead of N×M — and traces can be freely recombined offline.

2. Warm KV cache support for agentic / long-context workloads

In multi-turn chat, tool-use agents, and RAG pipelines, each new request builds on thousands of tokens already in KV cache. The prefill phase only processes new tokens while reading from cached context, producing different kernel shapes and timings than a cold prefill of the same total length. The --existing-ctx flag simulates this by seeding the radix cache before profiling, enabling accurate traces for warm-cache scenarios.

Implementation

This PR implements the full pipeline:

• Separates prefill (EXTEND) from decode traces using SGLang's profile_by_stage API
• Controls exact token counts via input_ids (not approximate text prompts)
• Supports KV cache context simulation (--existing-ctx) for warm prefill scenarios
• Aggregates kernel timings across TP/DP ranks with correct semantics per kernel type

Changes

New files

File	Lines	Purpose
scripts/run_stage_profile.py	925	Main orchestrator: warmup → profile → parse → cross-rank analysis, with --collect {perf,shapes,all} modes
utils/cross_rank_agg.py	491	Cross-rank kernel aggregation (symmetric collectives → per-invocation min, asymmetric → max, compute → mean)
utils/shape_merge.py	365	Merge kernel shape data from no-CUDA-graph pass into timing CSVs with strict 1:1 occurrence matching
utils/net.py	31	Shared wait_for_port helper (extracted from duplicate definitions)
utils/__init__.py	0	Package marker
tests/integration/test_stage_profile_configs.py	518	Integration tests for all 3 --collect modes with shape validation
tests/unit/test_cross_rank_agg.py	248	Unit tests for kernel classification, comm detection, multi-rank aggregation
tests/unit/test_shape_merge.py	322	Unit tests for shape merging (single pair + directory mode)

Modified files

File	What changed
README.md	+109 lines: Stage Profiling section with quick reference, examples, output structure, utilities table
.gitignore	Add __pycache__/
simulator/base_parser.py	None-safe dtype/byte lookups in 3 comm kernel blocks — pytorch_to_nccl_dtype.get() returns None for unknown dtypes, previously caused TypeError
simulator/benchmarks/nccl_benchmarks.py	try/except ValueError around float() parsing in run_nccl_all_reduce_perf and run_nccl_all_gather_perf
scripts/run_simulate.py	Formatting only (black)
tests/unit/test_batch_request.py	Add new fields (text_prompt_len, vision_prompt_len, timestamp), mock requests.get to prevent real HTTP calls, fix base_url
tests/unit/test_kernel_db_coverage.py	Relax missing-op assertion to warnings.warn — unknown kernels need manual kernels.json entries, not test failures
tests/unit/test_defined_len.py	Formatting only (black)
tests/unit/test_llmcompass_backend.py	Remove extra blank line

Key design decisions

Single-point API (not sweep grid)

The script profiles a single (bs, input_len, existing_ctx, decode_tokens) point per invocation. Grid sweeps are handled by the caller (shell loops, CI matrix, etc.). This simplifies the code and makes each run self-contained.

Exact token control via input_ids

Prompts use raw token IDs (input_ids) instead of text to ensure exact prefill token counts. Fixed random seeds guarantee deterministic prompts for radix-cache prefix matching between seed and profile requests.

Fencepost handling

SGLang's profiler stops when batch_count > num_steps (not >=). To capture exactly decode_tokens decode batches, we pass num_steps = max(1, decode_tokens - 1).

Cross-rank aggregation semantics

Kernel type	Aggregation	Rationale
Symmetric collectives (all_reduce, all_gather, reduce_scatter)	Per-invocation min across ranks	All ranks participate equally; min removes straggler noise
Asymmetric collectives (all_to_all)	Max across ranks	Different ranks may send/receive different amounts
Compute kernels	Mean across ranks	Slight per-rank timing variation is noise

Shape merge

CUDA graphs hide kernel input shapes. A second pass without CUDA graphs captures shapes, which are merged into timing CSVs by 1:1 kernel-name occurrence matching. Mismatches raise ValueError to catch configuration bugs early.

Testing

• Unit tests: test_cross_rank_agg.py (13 tests), test_shape_merge.py (10 tests) — run locally with pytest
• Integration tests: test_stage_profile_configs.py (4 tests) — require GPU + Docker container with model configs. Filter with RUN_CONFIGS=P1

# Unit tests (no GPU needed)
cd FlowSim && python -m pytest tests/unit/test_cross_rank_agg.py tests/unit/test_shape_merge.py -v

# Integration tests (inside Docker with GPUs)
RUN_CONFIGS=P1 python -m pytest tests/integration/test_stage_profile_configs.py -v --timeout=1800

Stats

17 files changed, 3173 insertions(+), 84 deletions(-)

[Copilot]

Pull request overview

Adds a stage-separated profiling workflow to FlowSim that profiles SGLang EXTEND (prefill) and DECODE independently (including warm KV-cache scenarios), then post-processes traces into parsed CSVs with cross-rank aggregation and optional kernel-shape merging.

Changes:

• Introduces scripts/run_stage_profile.py to orchestrate stage profiling, parsing, cross-rank analysis, and optional shape collection/merging.
• Adds utilities for cross-rank kernel aggregation (utils/cross_rank_agg.py), shape merging (utils/shape_merge.py), and a shared networking helper (utils/net.py).
• Expands unit/integration test coverage for the new pipeline and adjusts a few existing tests/parsers for robustness.

Reviewed changes

Copilot reviewed 15 out of 17 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
scripts/run_stage_profile.py	New stage-separated profiling orchestrator (collect/parse/analyze/shapes/merge).
utils/cross_rank_agg.py	New cross-rank aggregation logic with comm-vs-compute semantics.
utils/shape_merge.py	New shape-pass → timing-pass CSV merge utility with occurrence matching.
utils/net.py	New shared wait_for_port() helper.
utils/__init__.py	Marks utils as a package.
tests/integration/test_stage_profile_configs.py	Integration coverage for --collect {perf,shapes,all} and workload-shape validation.
tests/unit/test_cross_rank_agg.py	Unit tests for kernel classification + aggregation behavior.
tests/unit/test_shape_merge.py	Unit tests for single-pair and directory shape merge behavior.
simulator/base_parser.py	Makes comm calibration more robust to unknown dtypes/bytes.
simulator/benchmarks/nccl_benchmarks.py	Hardens parsing of benchmark output by skipping non-float values.
tests/unit/test_batch_request.py	Updates request fields and prevents real HTTP calls in unit tests.
tests/unit/test_kernel_db_coverage.py	Downgrades missing-op assertion to a warning.
tests/unit/test_defined_len.py	Formatting-only change.
tests/unit/test_llmcompass_backend.py	Removes an extraneous blank line.
scripts/run_simulate.py	Formatting-only change (black).
README.md	Documents stage profiling usage, modes, outputs, and utilities.
.gitignore	Ignores __pycache__/.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.