feat(profile): merge Proton and VizTracer traces onto one timeline#662
Draft
antiagainst wants to merge 5 commits into
Draft
feat(profile): merge Proton and VizTracer traces onto one timeline#662antiagainst wants to merge 5 commits into
antiagainst wants to merge 5 commits into
Conversation
Proton could not profile serving workloads: GPU kernels run in scheduler subprocesses (mp.Process spawned in _launch_subprocesses), and those are torn down with SIGKILL on both crash and clean engine shutdown. The import-time TOKENSPEED_KERNEL_PROFILE bootstrap finalizes via atexit, which never runs under SIGKILL, so no profile was ever written; all ranks also shared one default output path. Fix by running Proton where the kernels launch, controlled through the existing profile control plane: * request_handler: new PROTON activity for /start_profile and /stop_profile. Each scheduler rank starts its own Proton session writing <output_dir>/<profile_id>-TP-<rank>[-STAGE].proton.<fmt> and finalizes it on /stop_profile while the process is alive, so nothing depends on atexit or graceful exit. The existing start_step / num_steps / profile_by_stage step windows apply to Proton unchanged. * init_profile validation: reject PROTON combined with GPU or CUDA_PROFILER (CUPTI/roctracer allows one GPU profiling client per process; host-side CPU/MEM/VIZTRACER still compose), and fail cleanly when Proton is unavailable or an env-bootstrap session already owns the process. * tokenspeed_kernel.profiling: expose proton_available() and profile_config_from_env(output=...) so the control plane reuses the TOKENSPEED_KERNEL_PROFILE_* env config (data/backend/mode/ hook) with a per-rank output override. Env bootstrap behavior is unchanged and remains the right tool for single-process scripts. * bench serve: new --profile-activities flag (e.g. PROTON) sent in the /start_profile body. The profile request now uses a dedicated body dict, so profiler fields like num_steps no longer leak into every generation request payload. * Document the serving flow in the tokenspeed-kernel README; add unit tests for the new helpers and the scheduler-side PROTON dispatch, activity-conflict rejection, and step-window finalize. Signed-off-by: Lei Zhang <antiagainst@gmail.com>
Signed-off-by: Lei Zhang <antiagainst@gmail.com>
attn_tp_rank is computed modulo the attention TP size, and the DP controller launches scheduler processes as dp_rank * attn_tp_size + attn_tp_rank, so with attention DP/CP > 1 several processes share the same TP rank. Their profile outputs under a shared output_dir used the same <profile_id>-TP-<tp_rank> path and would race and overwrite each other. This affected the new Proton output as well as the pre-existing torch trace, memory snapshot, and viztracer file names. Build one process-unique rank tag per scheduler at init — TP<rank>, prefixed with DP<rank> and/or CP<rank> when those dimensions exist (e.g. DP1-CP0-TP2) — and use it for all four profiler outputs. Add unit tests for the tag shape and for two DP peers with the same TP rank producing distinct Proton output paths. Signed-off-by: Lei Zhang <antiagainst@gmail.com>
init_profile set self.profile_by_stage before the PROTON activity validation, so a rejected request (e.g. activities=["PROTON", "GPU"] with profile_by_stage=true) returned failure while leaving profile_by_stage=True behind with the stage counters still None. The event loop runs _profile_batch_predicate on every batch, so the next scheduled batch entered the stage-profiling branch and crashed the scheduler with a TypeError on `None += 1`. Hoist the validation block above the first state mutation so a rejected control request leaves the handler exactly as it was, and add a regression test that replays the rejected request and drives the next-batch predicate. Signed-off-by: Lei Zhang <antiagainst@gmail.com>
Profiling a scheduler rank with both VIZTRACER and PROTON activities yields two chrome-trace files whose event timestamps are relative to each profiler's own start, so they could not be viewed on a shared time axis. Both files carry an absolute wall-clock anchor for ts=0: VizTracer reports store viztracer_metadata.baseTimeNanoseconds, and Proton chrome traces now store a top-level baseTimeNanoseconds (added to the tokenspeed-triton fork, mirroring the torch>=2.4 chrome trace contract). That makes a fully offline merge possible. * New tokenspeed.trace_merge module and `tokenspeed merge-traces` CLI: shifts every Proton event by the delta between the two anchors and appends them to the VizTracer report, producing one Perfetto-loadable JSON where Python frames and Proton's kernel / scope lanes (including launch->kernel flow arrows) share a time axis. One file pair per scheduler rank. * Proton's synthetic pid-0 process is renamed "Trace" -> "Proton" in the merged view, and its metadata events stay timestamp-free so they cannot clobber the report's process/thread names. * Inputs missing their anchor fail with a pointed error (old tokenspeed-triton or viztracer versions). Alignment accuracy is us-to-ms grade (each profiler reconciles GPU and CPU clocks independently): good for correlating host stalls with GPU gaps, not for sub-microsecond attribution. Tested with synthetic trace pairs encoding both formats' contracts, plus an end-to-end merge of a real VizTracer report. Signed-off-by: Lei Zhang <antiagainst@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Profiling a scheduler rank with both VIZTRACER and PROTON activities yields two chrome-trace files whose event timestamps are relative to each profiler's own start, so they could not be viewed on a shared time axis.
Both files carry an absolute wall-clock anchor for ts=0: VizTracer reports store viztracer_metadata.baseTimeNanoseconds, and Proton chrome traces now store a top-level baseTimeNanoseconds (added to the tokenspeed-triton fork, mirroring the torch>=2.4 chrome trace contract). That makes a fully offline merge possible.
New tokenspeed.trace_merge module and
tokenspeed merge-tracesCLI: shifts every Proton event by the delta between the two anchors and appends them to the VizTracer report, producing one Perfetto-loadable JSON where Python frames and Proton's kernel / scope lanes (including launch->kernel flow arrows) share a time axis. One file pair per scheduler rank.Proton's synthetic pid-0 process is renamed "Trace" -> "Proton" in the merged view, and its metadata events stay timestamp-free so they cannot clobber the report's process/thread names.
Inputs missing their anchor fail with a pointed error (old tokenspeed-triton or viztracer versions).
Alignment accuracy is us-to-ms grade (each profiler reconciles GPU and CPU clocks independently): good for correlating host stalls with GPU gaps, not for sub-microsecond attribution.
Tested with synthetic trace pairs encoding both formats' contracts, plus an end-to-end merge of a real VizTracer report.