sitecustomize.py is imported automatically by Python at startup
in every process, including the SGLang scheduler/TP worker
processes spawned for tensor parallelism. Adding the p800_timing
import here ensures the monkey-patches are applied to all 8 TP
workers, not just the launch_server parent process (which never
runs the model forward and would record empty timing data).
The timing module self-disables unless SGLANG_TIMING_ENABLED=1,
so this import is a no-op when timing is not in use.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
PyTorch Profiler crashes on P800 XPU due to a CUPTI bug in
cuptiActivityDisable() that SIGKILLs the server process. This
module provides an alternative by monkey-patching SGLang's
DeepSeek-V4 decoder layer and its sub-operators to record
per-op wall-clock timing.
What it does:
- Patches DeepseekV4DecoderLayer.forward and sub-operators
(hc_pre/hc_post, RMSNorm, MQALayer attention, QKV projection,
DeepseekV2MoE and its gate/shared/routed experts) with timed
wrappers that call torch.cuda.synchronize() before and after.
- Distinguishes prefill vs decode via forward_batch flags
(is_prefill_only / is_extend_in_batch / extend_num_tokens)
and propagates the phase to sub-operators via threading.local.
- Writes per-PID JSON summaries (count/avg/p50/p95/p99) every
5000 records plus on atexit/SIGTERM, so data survives when
the server is stopped.
Usage: set SGLANG_TIMING_ENABLED=1 in the container environment.
Output: /tmp/p800_timing_results_{pid}.json
Known limitations:
- torch.cuda.synchronize() adds overhead, inflating small-op times.
- moe_routed_ms only covers forward_normal (65% of calls); the
dual_stream path taken during CUDA graph capture is not patched.
- All-reduce patch fails (Communicator import path mismatch).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>