diff --git a/.github/workflows/benchmark-tmpl.yml b/.github/workflows/benchmark-tmpl.yml index 35548863fe..f45c32865e 100644 --- a/.github/workflows/benchmark-tmpl.yml +++ b/.github/workflows/benchmark-tmpl.yml @@ -94,6 +94,11 @@ on: required: false type: string default: "" + profile: + description: "Op-attribution profiling settings (JSON; vLLM, SGLang and TRT-LLM single-node srt-slurm recipes); empty disables" + required: false + type: string + default: "" env: PYTHONPATH: ${{ github.workspace }}/inferencex-e2e INFERENCEX_E2E_ROOT: ${{ github.workspace }} @@ -150,6 +155,7 @@ env: AIPERF_EXPERIMENTAL_FAST: ${{ inputs.agentx-fast && '1' || '0' }} EVAL_LIMIT: ${{ inputs.eval-limit }} SWEBENCH_GEN_MODE: ${{ inputs.swebench-gen-mode || 'agentic' }} + INFX_PROFILE: ${{ inputs.profile }} AIPERF_FAILED_REQUEST_THRESHOLD: '0.10' RESULT_DIR: /workspace/results PYTHONDONTWRITEBYTECODE: '1' @@ -353,7 +359,8 @@ jobs: fi - name: Validate agentic result - if: ${{ !inputs.eval-only && inputs.scenario-type == 'agentic-coding' }} + # A synthetic profile (INFX_PROFILE mode=synthetic) replays no AgentX traces. + if: ${{ !inputs.eval-only && inputs.scenario-type == 'agentic-coding' && !contains(inputs.profile, 'synthetic') }} env: PYTHONPATH: ${{ github.workspace }}/.result-tooling/inferencex-e2e run: | @@ -404,6 +411,48 @@ jobs: !${{ env.INFERENCEX_E2E_ROOT }}/results/aiperf_artifacts/profile_export_raw.jsonl if-no-files-found: ignore + - name: Upload op-attribution profile + if: ${{ always() && inputs.profile != '' }} + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: profile_${{ env.RESULT_FILENAME }} + path: ${{ env.INFERENCEX_E2E_ROOT }}/infx-profile.tar + compression-level: 1 + if-no-files-found: warn + + # Per-step breakdown of the profile above: one file per step and rank, + # with index.json and report.json. Best effort: the raw profile is the record. + - name: Extract op-attribution profile by step + if: ${{ always() && inputs.profile != '' }} + continue-on-error: true + run: | + set -euo pipefail + cd "$INFERENCEX_E2E_ROOT" + if [[ ! -f infx-profile.tar ]]; then + echo "No profile to extract" + exit 0 + fi + rm -rf infx-profile-raw infx-profile-steps + mkdir -p infx-profile-raw + tar xf infx-profile.tar -C infx-profile-raw + if ! "${INFERENCEX_RESULTS_PYTHON:-python3}" benchmarks/profiling/extract.py \ + infx-profile-raw/infx_profile infx-profile-steps > infx-profile-extract.log 2>&1; then + tail -50 infx-profile-extract.log + exit 1 + fi + cp infx-profile-extract.log infx-profile-steps/extract.log + python3 -c "import json; r = json.load(open('infx-profile-steps/report.json')); print(json.dumps({k: v for k, v in r.items() if k != 'traces'}, indent=1)); [print(json.dumps(t)) for t in r['traces']]" + rm -rf infx-profile-raw + + - name: Upload op-attribution profile by step + if: ${{ always() && inputs.profile != '' }} + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: profile_steps_${{ env.RESULT_FILENAME }} + path: ${{ env.INFERENCEX_E2E_ROOT }}/infx-profile-steps/ + compression-level: 0 + if-no-files-found: warn + - name: Upload server logs if: always() uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 diff --git a/.github/workflows/e2e-tests.yml b/.github/workflows/e2e-tests.yml index 3310fee119..95129eb3ae 100644 --- a/.github/workflows/e2e-tests.yml +++ b/.github/workflows/e2e-tests.yml @@ -101,6 +101,11 @@ on: # zizmor: ignore[concurrency-limits] required: false type: string default: "[]" + profile: + description: "Op-attribution profiling settings (JSON) for vLLM AgentX points; empty disables" + required: false + type: string + default: "" workflow_call: secrets: INFERENCEX_OFFICIAL_RO_HF_TOKEN: @@ -203,6 +208,11 @@ on: # zizmor: ignore[concurrency-limits] required: false type: string default: "[]" + profile: + description: "Op-attribution profiling settings (JSON) for vLLM AgentX points; empty disables" + required: false + type: string + default: "" env: PYTHONPATH: ${{ github.workspace }}/inferencex-e2e @@ -439,6 +449,7 @@ jobs: run-eval: false require-power: ${{ inputs.require-power }} scenario-type: agentic-coding + profile: ${{ inputs.profile }} ref: ${{ needs.get-jobs.outputs.benchmark-ref }} test-sweep-agentic-evals: diff --git a/inferencex-e2e/benchmarks/benchmark_lib.sh b/inferencex-e2e/benchmarks/benchmark_lib.sh index 39a7e22df3..5f9f09ddd4 100644 --- a/inferencex-e2e/benchmarks/benchmark_lib.sh +++ b/inferencex-e2e/benchmarks/benchmark_lib.sh @@ -3280,6 +3280,10 @@ build_replay_cmd() { duration=1200 warmup_requests_per_lane=1 fi + # Op-attribution profiling replays only as long as its windows need. + if [[ -n "${INFX_PROFILE_DURATION:-}" ]]; then + duration="$INFX_PROFILE_DURATION" + fi export AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES="${AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES}" # Dataset configuration takes 4-5 min on fast /tmp (B300) but reached 14 min diff --git a/inferencex-e2e/benchmarks/profiling/README.md b/inferencex-e2e/benchmarks/profiling/README.md new file mode 100644 index 0000000000..a81c896566 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/README.md @@ -0,0 +1,248 @@ +# Op-attributed engine profiles for AgentX points + +Profiles a single-node vLLM, SGLang or TRT-LLM srt-slurm AgentX point (for +example `dsv4-fp4-b200-vllm-agentic-mtp`) during its real replay and +attributes every device kernel to the CPU-side call stack that launched it: the +module stack, the torch op, and for kernels launched straight from Python, the +launcher and its engine callers. Kernels replayed from CUDA graphs are +attributed through the graph's capture. + +## Layout + +- `sitecustomize.py`: loads through the workers' `PYTHONPATH` and boots + `infx_prof` when `INFX_PROF_DIR` is set. +- `infx_prof/`: the in-engine hooks. The core is engine-neutral: + - `imports.py`: patch a module when it is imported (or now, if it already is). + - `steps.py`: the per-step record and `infx_step#k` marker. + - `windows.py`: profiler windows, and the graph-capture profile. + - `graphs.py`: CUDA graph ordinals and stream markers. + - `modules.py`: module markers, eager and Inductor-compiled. + - `launchers.py`: Python launchers of kernels without a torch op. + - `routing.py`: per-step MoE top-k ids. + - `envinfo.py`: versions and device identity per rank. + + `infx_prof/engines/.py` adapts one engine (`INFX_PROF_ENGINE`): + where its steps, ranks, graph captures, models and MoE routers are, and + which kernel packages launch from Python. +- `profile_windows.py`: the window client, run beside the benchmark client; + `clock_sampler.py` polls GPU clocks for it. +- `synthetic_load.py`: the fixed-shape load of synthetic mode. +- `extract.py`: per-step attribution from a profile artifact. + +| | vLLM | SGLang | TRT-LLM | +|---|---|---|---| +| Step | `execute_model` (dummy: `execute_dummy_batch`) | `Scheduler.run_batch` (dummy: an IDLE batch) | `PyExecutor._forward_step` (dummy: only padding requests) | +| Windows | vLLM's profiler, `/start_profile` | `control/windows.jsonl` | `control/windows.jsonl` | +| Graph capture | `capture_model` | every `*CudaGraphRunner.capture` | `PyTorchModelEngine.warmup` | +| Graph mode | the dispatcher's `cg_mode` | the target graph runner that replayed (decode FULL; prefill by its backend) | `CUDAGraphRunner.replay` (FULL), other graph replays (PIECEWISE) | +| MoE routing | the routed-experts `capture_fn` | the ids `FusedMoE.forward_impl` is handed, else recomputed | `routing_method.apply`, else recomputed | + +Inside torch.compile (SGLang's piecewise prefill traces the MoE layers) the copy into +the routing buffer is the custom op `infx::capture_routing`, which Dynamo +leaves opaque, so compiled graphs still record routing. + +When a router fuses top-k into the MoE kernel (TRT-LLM's TRTLLM-gen backend, +FlashInfer TRT-LLM-gen MoE in SGLang), no ids exist outside the kernel. SGLang +then hands the MoE layer a `BypassedTopKOutput`, and `FusedMoE.forward_impl` +recomputes its ids eagerly with SGLang's own `select_experts` +(`to_standard`). In TRT-LLM, routing methods with a closed form (top-k of the logits, or of +sigmoid plus the correction bias, as MiniMax-M3's) are then recomputed from the +layer's router logits, the same top-k selection on the same inputs, under an +`infx_py#infx_prof.routing.recompute` marker (SGLang's recompute too): those +kernels are the profiler's, not the engine's. Otherwise +`routing//meta.json` says routing is unavailable, as it does for SGLang +DP attention with a gathered (not all-to-all) MoE, whose router sees every DP +rank's tokens in a padded layout. + +## Running + +Dispatch `e2e-tests.yml` on a branch carrying this directory with a `profile` +input (a JSON object; `{}` takes every default): + +```sh +gh workflow run e2e-tests.yml -R SemiAnalysisAI/InferenceX --ref --json <<'EOF' +{ + "ref": "", + "agentx-fast": "true", + "test-name": "profile test", + "profile": "{}", + "generate-cli-command": "test-config --config-files configs/nvidia-master.yaml --scenario-type agentic-coding --no-evals --conc 32 --config-keys dsv4-fp4-b200-vllm-agentic-mtp" +} +EOF +``` + +One concurrency selects one recipe variant (for DSV4 on B200: `tp8_c*` or +`dep8_c*`). Settings (`infx/srt_slurm/single_node.py`, `PROFILE_DEFAULTS`): + +| Setting | Default | Meaning | +|---|---|---| +| `windows` | `[["warmup", 0, 32], ["decode", 0, 32]]` | `[anchor, delay_seconds, iterations]` per torch profiler window. `warmup` and `profiling` anchor on aiperf logging that phase's start; AgentX warmup is the lanes' long first turns, all sent at its start (prefill-heavy; at TP8 c4 they have drained by 60 s in). `decode` anchors on steady decode: from the warmup start, the first 20 s in which at least 90% of the engines' logged steps replayed FULL CUDA graphs (falling back to 150 s before the measured replay's cap). Each measured-phase turn re-prefills first, so under data parallelism, where one prefilling rank keeps every rank eager, steady decode arrives late and at a time that depends on concurrency | +| `capture_ranks` | engine default | Workers whose CUDA graph capture is profiled (`"all"` for every rank). By default every rank under SGLang or TRT-LLM attention DP, whose ranks capture different graphs (per-rank buffers change copy kernels), else `"dp0_tp0"`; each rank joins its own capture when profiled | +| `host_headroom_gib` | `128` | Host memory taken from a CPU KV tier for the profiler's buffers: vLLM's SimpleCPUOffload pool or embedded Mooncake store (`global_segment_size`), SGLang's HiCache `hicache-size`, TRT-LLM's `kv_cache_config.host_cache_size`. A HiCache pool sized by `hicache-ratio` (a multiple of the GPU KV pool) is left as is | +| `duration` | last `profiling` delay + 600 s | Cap on the measured replay. Once its last window closes (and the measured phase has run 60 s), the window client SIGINTs aiperf, which exports what it measured and exits zero, so the run does not replay past its windows | +| `mode` | `"agentic"` | `"agentic"` replays AgentX; `"synthetic"` drives fixed shapes (`synthetic_load.py`): CONC random-token prompts of `isl` tokens with one-token outputs while the prefill window is open, then `osl`-token generations until the decode window closes. No trace dataset or AgentX warmup: a point takes minutes after engine start-up | +| `isl`, `osl` | `8192`, `4096` | Synthetic prompt and output lengths. `isl` stands in for the KV context: decode steps attend over `isl` plus the tokens generated so far | + +A profiled point holds its node for 30 to 60 minutes on B200 (engine start-up +and graph capture, then an AgentX warmup that grows with concurrency). Its +throughput is not a result: windows pause the engine while they export, and +profiled runs tolerate failed requests. + +At DSV4 DEP8 c192 on B200, a warmup window has preceded an engine fault in two +of three runs (`CUBLAS_STATUS_EXECUTION_FAILED` in the attention compressor's +`torch.mm`, when warmup's cache-pressure requests arrive, minutes after the +window closed); the same point without it (`{"windows": [["decode", 0, 32]]}`) +and unprofiled production runs complete. Profile such a point's prefill +(`{"windows": [["warmup", 0, 32]]}`) and decode in separate runs. + +Deviations from the recipe, all recorded in the run's config: the hooks' +environment, compile caches off (compiled code gets its module markers when +Inductor generates it), longer engine timeouts (a window's export pauses its +rank), CPU KV tiers smaller by `host_headroom_gib`, and per engine: + +- vLLM: its torch profiler config. +- SGLang: `watchdog-timeout` and `dist-timeout` of at least 1800 s. +- TRT-LLM: none beyond the above. + +## Example + +DSV4-Pro FP4 on B200, TP8 c4. A decode step: one FULL CUDA-graph replay, +every kernel joined to the op, launcher and module that launched it: + +![Decode step timeline with per-kernel host attribution](docs/decode_step.png) + +A prefill step with each kernel's SM clock; the SW power cap engages and the +clock steps on NVML's ~100 ms grid: + +![Prefill step timeline with per-kernel SM clock](docs/prefill_step.png) + +## Artifacts + +- `profile_`: `infx-profile.tar`, the full record. + - `torch/`: the engine's per-rank traces, one per window. + - `capture/`: the CUDA graph capture traces, one per capture session, with + module and launcher markers (no Python stacks: over every captured size + their export outlasts start-up). + - `steps/.jsonl`: every step's batch composition on that rank. + - `copies/.jsonl`: every CPU KV-offload block copy the rank queued + (direction, blocks, bytes, step, vLLM callers; vLLM's CPU offload only). + - `routing//step.npz`, `routing//meta.json`: each profiled + step's MoE routing: `counts` (tokens per expert per layer, int32 + `[layers, experts]`), `tokens`, `t0_ns`, and `ids` (each token's top-k + expert ids per layer, int16 `[tokens, layers, topk]`, -1 for none) when + the step has at most 1024 tokens. `meta.json` maps layer ids to module + names and lists layers the capture could not bind. + - `clocks/window.csv`, `clocks/gpus.json`: every GPU's graphics, SM, + memory and video clocks and clock event reasons, polled through NVML every + 250 us per GPU for the length of each window, each poll stamped when it + began and returned (`clock_sampler.py`). + - `graphs/`, `env/`, `profiler/`, `control/`, `windows_conc.jsonl`: + graph ordinals, versions and config per rank, each rank's window + start/stop, window requests, and the window client's log. +- `profile_steps_`: the same profile broken down by step + (`extract.py`). + - `index.json`: every step file with its kernel count, device span, + tokens, request count, phase and CUDA graph mode. + - `report.json`: per-trace attribution coverage, graph joins, and the + kernel-name check. + - `window//step.json.gz` (`dummy` for an idle DP rank's + forward, `unstepped` for kernels outside any step): + + ``` + {"rank", "window", "kind", "step", + "summary": {"kernels", "busy_us", "t0_us", "t1_us", "span_us", "sources"}, + "batch": {"phase", "graph", "total_tokens", + "reqs": [{"req", "scheduled", "computed", "spec", "new", "prompt_len"}], ...}, + "routing": {"source_rank", "tokens", "topk", "num_experts", "has_ids", "file", + "expert_tokens": {"": [[expert, tokens], ...]}}, + "kernels": [{"kernel", "cat", "stream", "device", "ts_us", "dur_us", + "source": "eager" | "graph" | "offload_copy" | "unattributed", + "module_stack": [[qualified name, [[shape, dtype], ...]], ...], + "op", "op_chain", "input_dims", "input_types", "concrete_inputs", + "launcher", "launcher_callers", "annotations", + "graph", "node_pos", "graph_node_id", "grid", "block", + "copy": {"store", "blocks", "bytes", "issue_lag_us"}, + "clocks": {"graphics_mhz": [min, max], "sm_mhz": [min, max], + "mem_mhz": [min, max], "video_mhz": [min, max], + "event_reasons", "samples", "prior_us"}, ...}]} + ``` + +## How attribution works + +`sitecustomize.py` loads through `PYTHONPATH` and does nothing unless +`INFX_PROF_DIR` is set. Its markers are `record_function` ranges emitted only +while a profiler runs: + +- `infx_step#k` / `infx_dummy#k`: a scheduled step and an idle DP rank's + dummy forward. Each step's batch composition is logged under the same `k`. +- `infx_mod##:;...` (e.g. `7x7168:bfloat16`): every + module call and its tensor inputs, through global module hooks registered + only while a window or capture is profiled (they do nothing while Dynamo + traces). Inductor-compiled code marks the module of each scheduler node it + generated, and engine ops that take a `layer_name` mark that layer. +- `infx_py##|...`: the Python entry points that launch + kernels without a torch op (Triton, TileLang, CuTe DSL, and the engine's + kernel packages: DeepGEMM, FlashInfer, sgl-kernel, ...). +- `infx_graph_capture#n` / `infx_graph_replay#n`: CUDA graph ordinals. + +Marker names carry no quotes or backslashes: Kineto writes event names into the +trace JSON unescaped, and the extractor fails on a trace it cannot decode. + +An eager kernel joins its launch through the CUDA correlation id. A kernel +replayed from graph `n` carries a `graph node id`: ordered by node id, graph +`n`'s nodes pair with the launches recorded while graph `n` was captured, +which carry that launch's full context. `report.json` counts a graph whose +node count differs from its captured launches as `count_mismatch`, and lists +graph kernels whose op never launched that kernel eagerly. + +HIP graph replays carry no node id. Each replay stream runs its nodes in +capture order, so capture order is an interleaving of the per-stream sequences. +A beam search over the streams' heads picks the interleaving that best matches +each launch: a kernel family (name without template arguments) its op or +launcher ran eagerly, its Inductor or Triton kernel, and a memset or memcpy only +where one was captured; ties go to device start order. On MI355X DSV4 it finds +the exact optimum on every two-stream graph and orders the three-stream decode +graph with no kernel/copy mismatches (device order alone has 112). +`report.json` counts the positions that had name evidence as +`stream_order_evidence` of `stream_order_positions` (~70% on MI355X DSV4 +decode); positions without evidence follow device order. + +CPU KV offload (SimpleCPUOffloadConnector) copies blocks with +`cuMemcpyBatchAsync` from the connector's copy thread. Kineto records neither +that driver call nor `record_function` ranges on that thread, so those memcpys +have no CPU launch. Each queued copy is one memcpy on its direction's stream, +run in queue order, and `copies/` logs each one as it is queued. Per direction, +the extractor pairs the memcpys in order with the logged copies issued before +them (with equal bytes), taking the pairing with the least total issue lag. +The step markers put the log's wall clock on the trace clock. + +On AMD the sampler reads `amdsmi`'s gpu_metrics instead of NVML: the graphics +clock fills both `graphics_mhz` and `sm_mhz`, and `event_reasons` holds the +gpu_metrics throttle status (AMD's bit meanings, not NVML's). `gpus.json` +records which metrics fields were read; ranks join to GPUs by UUID or PCI +address. + +A kernel's `clocks` are its GPU's polls over its lifetime: the last poll that +returned before it started (`prior_us` earlier) and every poll overlapping it +(`samples`), as each clock's min and max and the OR of the event-reason +bitmasks (NVML `nvmlClocksEventReason*`). The window client, not the engine, +polls NVML, one thread per GPU, from just before each window opens until every +engine that started profiling has stopped (`profiler/.jsonl`, logged by +the patch) plus twice the window's longest step, as the GPU trails the engine; +the env record's GPU UUID ties a rank to its polls. On B200 an NVML poll occasionally blocks for 10 to 55 ms, more often +under prefill load; a reading is from somewhere inside its poll, and a kernel +within such a poll shows it as a large `prior_us`. +`report.json` gives each trace's clock coverage and `prior_us` percentiles. + +MoE routing is dynamic, so each step records what its router chose. The +engine adapter hands every MoE layer's top-k expert ids (target model, not the +draft) to the core capturer, before CUDA graph capture; in vLLM through the +routed-experts `capture_fn`. It copies the ids into a device buffer, a GPU copy +that graphs capture and replay. +Inside a window, every step's rows are copied to pinned host memory +asynchronously (`infx_routing_copy`) and written at the window's stop. The ids +are logical expert ids before EPLB remapping, for this DP rank's tokens; TP +ranks route the same tokens, so TP rank 0 writes for its group. A step file's +`routing.expert_tokens` is keyed by the MoE module name, which joins to the +`module_stack` of that layer's MoE kernels. The capture adds one small copy +kernel per MoE layer per step to every profiled run. diff --git a/inferencex-e2e/benchmarks/profiling/clock_sampler.py b/inferencex-e2e/benchmarks/profiling/clock_sampler.py new file mode 100644 index 0000000000..28cdf4b0f0 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/clock_sampler.py @@ -0,0 +1,200 @@ +"""Sample every GPU's clocks, for one profile window: NVML on NVIDIA, amdsmi on AMD. + +Runs in the window client's process, not the engine's. One thread per GPU +polls, every POLL_INTERVAL_S, the graphics, SM, memory and video clocks and +the clock event (throttle) reasons. On NVIDIA the reasons are NVML's +nvmlClocksEventReason* bits; on AMD, gpu_metrics' throttle status, and the +graphics clock fills both graphics_mhz and sm_mhz (AMD has no separate SM +clock). A poll is stamped with the wall clock (the one the engines' step log +uses) before and after its calls: those occasionally block for tens of +milliseconds, and the reading is from somewhere inside that interval. Samples +are kept in memory and written when the window's sampling stops. Output, +under OUT_DIR: + + gpus.json {"": {"uuid", "bdf", "backend", ...}} + window.csv t0_ns,t1_ns,gpu,graphics_mhz,sm_mhz,mem_mhz,video_mhz,event_reasons + +The extractor puts each kernel's lifetime on this timeline, matching a rank to +its GPU by UUID or PCI address. +""" + +import array +import ctypes +import json +import os +import threading +import time + +CLOCKS = (("graphics_mhz", 0), ("sm_mhz", 1), ("mem_mhz", 2), ("video_mhz", 3)) +FIELDS = ("t0_ns", "t1_ns", "gpu") + tuple(name for name, _ in CLOCKS) + ("event_reasons",) +# Clock changes seen on B200 are milliseconds apart; faster polling only burns a core. +POLL_INTERVAL_S = 250e-6 + + +class Nvml: + def __init__(self): + self.lib = ctypes.CDLL("libnvidia-ml.so.1") + self._check(self.lib.nvmlInit_v2()) + count = ctypes.c_uint() + self._check(self.lib.nvmlDeviceGetCount_v2(ctypes.byref(count))) + self.handles = [] + for i in range(count.value): + handle = ctypes.c_void_p() + self._check(self.lib.nvmlDeviceGetHandleByIndex_v2(i, ctypes.byref(handle))) + self.handles.append(handle) + # Renamed from ClocksThrottleReasons in recent drivers; same bitmask. + self._reasons = getattr(self.lib, "nvmlDeviceGetCurrentClocksEventReasons", None) or \ + self.lib.nvmlDeviceGetCurrentClocksThrottleReasons + + @staticmethod + def _check(status): + if status != 0: + raise RuntimeError(f"NVML error {status}") + + def gpus(self): + out = {} + buf = ctypes.create_string_buffer(96) + for i, handle in enumerate(self.handles): + self._check(self.lib.nvmlDeviceGetUUID(handle, buf, len(buf))) + out[str(i)] = {"uuid": buf.value.decode(), "bdf": None, "backend": "nvml"} + return out + + def poller(self, i): + """A function returning GPU i's (graphics, sm, mem, video MHz, event reasons); -1 on failure.""" + handle, lib, reasons = self.handles[i], self.lib, self._reasons + mhz, mask = ctypes.c_uint(), ctypes.c_ulonglong() + mhz_ref, mask_ref = ctypes.byref(mhz), ctypes.byref(mask) + + def poll(): + values = [] + for _, clock_type in CLOCKS: + values.append(mhz.value if lib.nvmlDeviceGetClockInfo(handle, clock_type, mhz_ref) == 0 + else -1) + values.append(mask.value if reasons(handle, mask_ref) == 0 else -1) + return values + return poll + + +class AmdSmi: + """amdsmi (ROCm's Python bindings): clocks and throttle status from gpu_metrics. + + gpu_metrics field names vary with the ASIC and amdsmi version (MI300-class + parts report per-XCC lists), so the first read picks the fields present. + """ + + CLOCK_FIELDS = { # column -> candidate gpu_metrics keys, first present wins + "graphics_mhz": ("current_gfxclk", "current_gfxclks", "average_gfxclk_frequency"), + "mem_mhz": ("current_uclk", "average_uclk_frequency"), + "video_mhz": ("current_vclk0", "current_vclk0s", "average_vclk0_frequency"), + "event_reasons": ("throttle_status", "indep_throttle_status"), + } + + def __init__(self): + import amdsmi + + self.lib = amdsmi + amdsmi.amdsmi_init() + self.handles = list(amdsmi.amdsmi_get_processor_handles()) + if not self.handles: + raise RuntimeError("amdsmi found no GPUs") + sample = amdsmi.amdsmi_get_gpu_metrics_info(self.handles[0]) + self.fields = {column: next((k for k in keys if self._value(sample.get(k)) is not None), None) + for column, keys in self.CLOCK_FIELDS.items()} + + @staticmethod + def _value(raw): + """A gpu_metrics value as an int: the first valid entry of a per-XCC list.""" + if isinstance(raw, (list, tuple)): + raw = next((v for v in raw if isinstance(v, int) and 0 < v < 0xFFFF), None) + return raw if isinstance(raw, int) and raw != 0xFFFF and raw != 0xFFFFFFFF else None + + def gpus(self): + out = {} + for i, handle in enumerate(self.handles): + info = {"backend": "amdsmi", "fields": self.fields} + for key, call in (("uuid", "amdsmi_get_gpu_device_uuid"), ("bdf", "amdsmi_get_gpu_device_bdf")): + try: + info[key] = str(getattr(self.lib, call)(handle)) + except Exception: + info[key] = None + out[str(i)] = info + return out + + def poller(self, i): + handle, lib, fields, value = self.handles[i], self.lib, self.fields, self._value + + def poll(): + try: + metrics = lib.amdsmi_get_gpu_metrics_info(handle) + except Exception: + return [-1] * 5 + got = {c: (value(metrics.get(k)) if k else None) for c, k in fields.items()} + gfx = got["graphics_mhz"] + return [v if v is not None else -1 + for v in (gfx, gfx, got["mem_mhz"], got["video_mhz"], got["event_reasons"])] + return poll + + +class ClockSampler: + """Polls every GPU, one thread each, between start(window) and stop().""" + + def __init__(self, out_dir): + self.out_dir = out_dir + self.nvml = None + self.error = None + errors = [] + for backend in (Nvml, AmdSmi): + try: + self.nvml = backend() + break + except Exception as e: + errors.append(f"{backend.__name__}: {e}") + if self.nvml is None: # no GPU telemetry here: windows still profile, without clocks + self.error = "; ".join(errors) + else: + try: + os.makedirs(out_dir, exist_ok=True) + with open(os.path.join(out_dir, "gpus.json"), "w") as f: + json.dump(self.nvml.gpus(), f) + except Exception as e: + self.error, self.nvml = str(e), None + self._threads = [] + self._samples = [] + self._stop = threading.Event() + self._window = None + self.polls = 0 + + def start(self, window): + if self.nvml is None: + return + self._stop.clear() + self._window = window + self._samples = [array.array("q") for _ in self.nvml.handles] + self._threads = [threading.Thread(target=self._run, args=(i,), daemon=True) + for i in range(len(self.nvml.handles))] + for thread in self._threads: + thread.start() + + def stop(self): + if not self._threads: + return + self._stop.set() + for thread in self._threads: + thread.join() + self._threads = [] + width = len(FIELDS) + self.polls = sum(len(s) // width for s in self._samples) + with open(os.path.join(self.out_dir, f"window{self._window}.csv"), "w") as f: + f.write(",".join(FIELDS) + "\n") + for samples in self._samples: + for k in range(0, len(samples), width): + f.write(",".join(map(str, samples[k:k + width])) + "\n") + self._samples = [] + + def _run(self, i): + poll, out, stop = self.nvml.poller(i), self._samples[i], self._stop + while not stop.is_set(): + t0 = time.time_ns() + values = poll() + out.extend((t0, time.time_ns(), i, *values)) + time.sleep(POLL_INTERVAL_S) diff --git a/inferencex-e2e/benchmarks/profiling/docs/decode_step.png b/inferencex-e2e/benchmarks/profiling/docs/decode_step.png new file mode 100644 index 0000000000..44cb442917 Binary files /dev/null and b/inferencex-e2e/benchmarks/profiling/docs/decode_step.png differ diff --git a/inferencex-e2e/benchmarks/profiling/docs/prefill_step.png b/inferencex-e2e/benchmarks/profiling/docs/prefill_step.png new file mode 100644 index 0000000000..872ed084c7 Binary files /dev/null and b/inferencex-e2e/benchmarks/profiling/docs/prefill_step.png differ diff --git a/inferencex-e2e/benchmarks/profiling/extract.py b/inferencex-e2e/benchmarks/profiling/extract.py new file mode 100644 index 0000000000..7c045978b8 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/extract.py @@ -0,0 +1,925 @@ +"""Attribute every device kernel of a profiled engine step to the op that launched it. + +Usage: extract.py PROFILE_DIR OUT_DIR + +PROFILE_DIR is an unpacked profile artifact (infx_profile/): torch/ holds the +engine's per-rank torch traces, capture/ the CUDA graph capture traces, +steps/ the per-step batch log, copies/ the CPU KV-offload copy log, clocks/ +the window client's NVML clock samples, routing/ the MoE routing of each +profiled step and env/ the per-rank environment. + +A kernel launched eagerly joins its CPU launch through the CUDA correlation +id. A kernel replayed from CUDA graph n carries a ``graph node id``; ordered +by node id, graph n's nodes pair with the launches recorded inside +``infx_graph_capture#n``, which carry the op, shapes and stack of the +capture. Capture ordinals follow capture order, which every rank shares; an +engine that captures in several sessions (TRT-LLM's estimation pass, SGLang's +graph runners) writes one capture trace per session, and their ordinals union. +HIP graph replays carry no node id: each stream runs its nodes in capture +order, so the capture order is the interleaving of the per-stream sequences +that best agrees with the kernel families each op launched eagerly. +A CPU KV-offload memcpy has no CPU launch (a driver call Kineto does not +record, from the connector's copy thread); per direction, the memcpys pair in +order with the logged copies issued before them, least total issue lag first. +Each kernel's clocks are the NVML polls of its rank's GPU over its lifetime: +the last poll that returned before it started, and every poll overlapping it. + +Outputs, per replay trace: kernels.jsonl.gz (one row per device activity) +and steps.jsonl (per step: batch composition and device time), plus +report.json with the join checks. +""" + +import ast +import bisect +import collections +import glob +import gzip +import json +import os +import re +import shutil +import struct +import sys +import zipfile + +MARK = re.compile(r"^infx_(step|dummy|graph_replay|graph_capture|piece|stream)#(\d+)$") +MODULE_MARK = "infx_mod#" +LAUNCHER_MARK = "infx_py#" +DEVICE_CATS = {"kernel", "gpu_memcpy", "gpu_memset"} +LAUNCH_CATS = {"cuda_runtime", "cuda_driver"} +CONTEXT_CATS = {"cpu_op", "user_annotation", "python_function"} +NODE_ID_MASK = (1 << 32) - 1 + + +KEPT_ARGS = { + "correlation", "External id", "Input Dims", "Input type", "Concrete Inputs", + "kernel_file", "graph node id", "stream", "device", "grid", "block", "bytes", +} + + +class Event: + """The fields of one complete ("X") trace event the joins use.""" + + __slots__ = ("cat", "name", "ts", "dur", "pid", "tid", "args") + + def __init__(self, raw, names): + self.cat = raw.get("cat") + self.name = names.setdefault(raw.get("name", ""), raw.get("name", "")) + self.ts = raw["ts"] + self.dur = raw.get("dur", 0) + self.pid = raw.get("pid") + self.tid = raw.get("tid") + args = raw.get("args") + # Python frames carry nothing the joins read; everything else keeps its ids. + self.args = ({k: v for k, v in args.items() if k in KEPT_ARGS} + if args and self.cat != "python_function" else {}) + + def __getitem__(self, key): + return getattr(self, key) + + def get(self, key, default=None): + value = getattr(self, key, None) + return default if value is None else value + + +def iter_trace_events(path, chunk_size=1 << 22): + """Stream the traceEvents array one event at a time; traces reach several GB.""" + decoder = json.JSONDecoder() + opener = gzip.open if path.endswith(".gz") else open + with opener(path, "rt") as f: + buf, pos = "", 0 + while True: # find the array + chunk = f.read(chunk_size) + if not chunk: + return + buf += chunk + start = buf.find('"traceEvents"') + if start >= 0: + bracket = buf.find("[", start) + if bracket >= 0: + pos = bracket + 1 + break + while True: + while pos < len(buf) and buf[pos] in " \t\r\n,": + pos += 1 + if pos < len(buf) and buf[pos] == "]": + return + try: + event, end = decoder.raw_decode(buf, pos) + except ValueError as error: + # An event never spans more than a chunk or two; a growing + # undecodable buffer is corrupt input, not a partial read. + if len(buf) - pos > 8 * chunk_size: + raise ValueError(f"{path}: undecodable trace near offset {pos}: {error}") from error + chunk = f.read(chunk_size) + if not chunk: + raise ValueError(f"{path}: trace ends inside an event: {error}") from error + buf, pos = buf[pos:] + chunk, 0 + continue + yield event + pos = end + if pos > chunk_size: + buf, pos = buf[pos:], 0 + + +def load_events(path, keep=None): + names = {} + return [Event(raw, names) for raw in iter_trace_events(path) + if raw.get("ph") == "X" and (keep is None or keep(raw))] + + +def replayed_graphs(paths): + """Graph ordinals the replay traces replay (their infx_graph_replay#n markers).""" + graphs = set() + for path in paths: + for raw in iter_trace_events(path): + m = MARK.match(raw.get("name", "")) if raw.get("ph") == "X" else None + if m and m.group(1) == "graph_replay": + graphs.add(int(m.group(2))) + return graphs + + +def capture_ranges(path, graphs): + """Sorted, merged (start, end) of the host-side infx_graph_capture#n ranges for the given ordinals. + + ROCm traces also carry each marker's GPU-timeline copy (gpu_user_annotation), + which can fall inside another graph's host range; only host markers count. + """ + ranges = [] + for raw in iter_trace_events(path): + m = MARK.match(raw.get("name", "")) if raw.get("ph") == "X" else None + if (m and m.group(1) == "graph_capture" and int(m.group(2)) in graphs + and raw.get("cat") != "gpu_user_annotation"): + ranges.append((raw["ts"], raw["ts"] + raw.get("dur", 0))) + merged = [] + for start, end in sorted(ranges): + if merged and start <= merged[-1][1]: + merged[-1] = (merged[-1][0], max(merged[-1][1], end)) + else: + merged.append((start, end)) + return merged + + +def overlapping(ranges): + """Predicate: a raw event overlaps one of the sorted, disjoint ranges.""" + starts = [a for a, _ in ranges] + + def keep(raw): + ts = raw.get("ts", 0) + end = ts + raw.get("dur", 0) + k = bisect.bisect_right(starts, end) - 1 # last range starting at or before the event's end + return k >= 0 and ranges[k][1] >= ts + return keep + + +def parse_signature(sig): + """[[shape, dtype], ...] from '7x7168:bfloat16;1x2:int64'.""" + if not sig: + return [] + out = [] + for item in sig.split(";"): + dims, _, dtype = item.rpartition(":") + out.append([[int(d) for d in dims.split("x")] if dims else [], dtype]) + return out + + +def parse_launcher(text): + """[launcher, [vLLM caller frames]] from '#|'.""" + label, _, callers = text.partition("#") + return [label, [c for c in callers.split("|") if c]] + + +def launch_kind(name): + """kernel / memcpy / memset for APIs that put work on a stream, else None.""" + if "Graph" in name: + return None + if "LaunchKernel" in name or "LaunchCooperativeKernel" in name: + return "kernel" + if "Memcpy" in name: + return "memcpy" + if "Memset" in name: + return "memset" + return None + + +class Trace: + """One Kineto trace with each CPU launch resolved to its enclosing context.""" + + def __init__(self, path, keep=None): + self.path = path + self.events = load_events(path, keep) + self._index() + + def _index(self): + by_thread = collections.defaultdict(list) + self.launches = {} # correlation -> launch event index + self.device = [] + for i, e in enumerate(self.events): + cat = e.get("cat") + if cat in CONTEXT_CATS: + by_thread[(e["pid"], e["tid"])].append(i) + elif cat in LAUNCH_CATS: + by_thread[(e["pid"], e["tid"])].append(i) + corr = e.get("args", {}).get("correlation") + if corr is not None: + self.launches[corr] = i + elif cat in DEVICE_CATS: + self.device.append(i) + # Sweep each thread's intervals; every launch gets its enclosing stack. + self.stack_of = {} + for items in by_thread.values(): + items.sort(key=lambda i: (self.events[i]["ts"], -self.events[i].get("dur", 0))) + stack = [] + for i in items: + e = self.events[i] + ts = e["ts"] + while stack and self.events[stack[-1]]["ts"] + self.events[stack[-1]].get("dur", 0) <= ts: + stack.pop() + if e["cat"] in LAUNCH_CATS: + self.stack_of[i] = tuple(stack) + else: + stack.append(i) + + def context(self, launch): + """Markers, op chain, innermost op args and module path of a launch.""" + marks = {} + ops = [] + annotations = [] + modules = [] # qualified names from infx_mod markers + stack_modules = [] # class-instance names from Python stacks (capture only) + frame = None + module_stack = [] # [qualified name, input signature], outermost first + launcher = None # innermost [launcher, vLLM caller frames] + for i in self.stack_of.get(launch, ()): + e = self.events[i] + name = e["name"] + m = MARK.match(name) + if m: + marks[m.group(1)] = int(m.group(2)) + elif name.startswith(MODULE_MARK): + qualname, _, sig = name[len(MODULE_MARK):].partition("#") + modules.append(qualname) + module_stack.append([qualname, parse_signature(sig)]) + elif name.startswith(LAUNCHER_MARK): + launcher = parse_launcher(name[len(LAUNCHER_MARK):]) + elif e["cat"] == "cpu_op": + ops.append(i) + elif e["cat"] == "user_annotation": + annotations.append(name) # e.g. vLLM's per-step execute_context_* scope + elif e["cat"] == "python_function": + if name.startswith("nn.Module: "): + stack_modules.append(name[len("nn.Module: "):]) + else: + frame = name + inner = self.events[ops[-1]] if ops else None + args = inner.get("args", {}) if inner else {} + return { + "marks": marks, + "op": inner["name"] if inner else None, + "op_chain": [self.events[i]["name"] for i in ops], + "annotations": annotations, + "input_dims": args.get("Input Dims"), + "input_types": args.get("Input type"), + "concrete_inputs": args.get("Concrete Inputs"), + "kernel_file": args.get("kernel_file"), + "module_path": modules or stack_modules, + "module_stack": module_stack, + "launcher": launcher[0] if launcher else None, + "launcher_callers": launcher[1] if launcher else None, + "py_frame": frame, + "launch_api": self.events[launch]["name"], + } + + def process_ids(self): + return {e["pid"] for e in self.events if e.get("cat") in CONTEXT_CATS} + + +def capture_launches(trace): + """graph ordinal -> ordered contexts of every stream launch captured into it. + + A launch inside a capture range that ran on the device then and there (a + stream not being captured: HIP runs those eagerly) is not one of the graph's + nodes; captured launches produce no device activity until replay. + """ + executed = {trace.events[i].get("args", {}).get("correlation") for i in trace.device} + graphs = collections.defaultdict(list) + order = sorted( + (i for i in trace.stack_of if launch_kind(trace.events[i]["name"])), + key=lambda i: trace.events[i]["ts"], + ) + for i in order: + if trace.events[i].get("args", {}).get("correlation") in executed: + continue + ctx = trace.context(i) + gid = ctx["marks"].get("graph_capture") + if gid is not None: + ctx["kind"] = launch_kind(trace.events[i]["name"]) + graphs[gid].append(ctx) + return graphs + + +def rank_of_pid(profile_dir): + ranks = {} + for path in glob.glob(os.path.join(profile_dir, "env", "*.json")): + with open(path) as f: + info = json.load(f) + ranks[info["pid"]] = info["rank"] + return ranks + + +def step_key(record): + """("step", k) for scheduled steps, ("dummy", k) for idle-rank dummy forwards.""" + if "step" in record: + return ("step", record["step"]) + if "dummy" in record: + return ("dummy", record["dummy"]) + return None + + +def load_steps(profile_dir): + steps = {} + for path in glob.glob(os.path.join(profile_dir, "steps", "*.jsonl")): + rank = os.path.basename(path)[: -len(".jsonl")] + with open(path) as f: + steps[rank] = {step_key(r): r for r in map(json.loads, f)} + return steps + + +def read_npz(path): + """name -> (shape, flat values) for an .npz of little-endian int arrays (stdlib only).""" + out = {} + with zipfile.ZipFile(path) as z: + for member in z.namelist(): + data = z.read(member) + assert data[:6] == b"\x93NUMPY", member + major = data[6] + header_len = struct.unpack(" (routing dir, meta) for ranks whose MoE routing was captured. + + Tensor-parallel ranks route the same tokens, so only TP rank 0 writes; the + others take their DP group's TP rank 0. + """ + found = {} + for meta_path in glob.glob(os.path.join(profile_dir, "routing", "*", "meta.json")): + with open(meta_path) as f: + meta = json.load(f) + found[meta["rank"]] = (os.path.dirname(meta_path), meta) + routing = {} + for rank, (path, meta) in found.items(): + source = rank if meta.get("write") else re.sub(r"_tp\d+$", "_tp0", rank) + if source in found and found[source][1].get("write"): + routing[rank] = (source, *found[source]) + return routing + + +def step_routing(routing, step, out_dir): + """A step's per-layer expert token counts; copies its routing file into out_dir.""" + if routing is None or step is None: + return None + source, path, meta = routing + npz = os.path.join(path, f"step{step:06d}.npz") + if not os.path.exists(npz): + return None + arrays = read_npz(npz) + (layers, experts), counts = arrays["counts"] + names = meta.get("bound", {}) + per_layer = {} + for layer in range(layers): + row = counts[layer * experts:(layer + 1) * experts] + if any(row): + per_layer[names.get(str(layer), str(layer))] = [[e, n] for e, n in enumerate(row) if n] + dest = os.path.join(out_dir, "routing", source, os.path.basename(npz)) + os.makedirs(os.path.dirname(dest), exist_ok=True) + if not os.path.exists(dest): + shutil.copyfile(npz, dest) + return {"source_rank": source, "tokens": arrays["tokens"][1][0], "topk": meta.get("topk"), + "num_experts": meta.get("num_experts"), "has_ids": "ids" in arrays, + "file": os.path.relpath(dest, out_dir), "expert_tokens": per_layer} + + +def load_copies(profile_dir, ranks): + """rank -> logged CPU KV-offload copies (one memcpy each), in issue order. + + A log written before the rank's process groups existed is named pid. + """ + copies = collections.defaultdict(list) + for path in glob.glob(os.path.join(profile_dir, "copies", "*.jsonl")): + tag = os.path.basename(path)[: -len(".jsonl")] + if tag.startswith("pid") and tag[3:].isdigit(): + tag = ranks.get(int(tag[3:]), tag) + with open(path) as f: + copies[tag] += (r for r in map(json.loads, f) if r.get("blocks")) + for logged in copies.values(): + logged.sort(key=lambda r: r["t_ns"]) + return copies + + +def unix_to_trace_us(trace, step_log): + """Offset from unix microseconds to the trace clock, from the step markers.""" + deltas = [] + for e in trace.events: + m = MARK.match(e["name"]) if e.get("cat") in CONTEXT_CATS else None + key = (m.group(1), int(m.group(2))) if m else None + if key in step_log and "t0_ns" in step_log[key]: + deltas.append(e["ts"] - step_log[key]["t0_ns"] / 1e3) + deltas.sort() + return deltas[len(deltas) // 2] if deltas else None + + +CLOCK_NAMES = ("graphics_mhz", "sm_mhz", "mem_mhz", "video_mhz") + + +def normalize_uuid(uuid): + text = str(uuid).lower() + for prefix in ("gpu-", "mig-"): + text = text.removeprefix(prefix) + return text + + +def load_clocks(profile_dir): + """rank -> (poll start ns, poll end ns, [(graphics, sm, mem, video MHz, event reasons)]). + + Polls are in time order; a reading was taken between its poll's start and end. + """ + clock_dir = os.path.join(profile_dir, "clocks") + try: + with open(os.path.join(clock_dir, "gpus.json")) as f: + gpus = json.load(f) + except (OSError, ValueError): + return {} + gpu_of = {} # normalized UUID or PCI address -> sampler GPU index + for i, info in gpus.items(): + info = info if isinstance(info, dict) else {"uuid": info} # older files: index -> UUID + for key in (info.get("uuid"), info.get("bdf")): + if key: + gpu_of[normalize_uuid(key)] = int(i) + rank_of_gpu = {} + for path in glob.glob(os.path.join(profile_dir, "env", "*.json")): + with open(path) as f: + info = json.load(f) + gpu = next((gpu_of[k] for k in (normalize_uuid(info.get("device_uuid", "")), + normalize_uuid(info.get("device_bdf", ""))) if k in gpu_of), None) + if gpu is not None: + rank_of_gpu[gpu] = info["rank"] + samples = collections.defaultdict(list) + for path in glob.glob(os.path.join(clock_dir, "window*.csv")): + with open(path) as f: + header = next(f, "").rstrip("\n").split(",") + stamps = 2 if header[:2] == ["t0_ns", "t1_ns"] else 1 # one stamp: an instant + for line in f: + parts = line.rstrip("\n").split(",") + if len(parts) != len(header): + continue # a row cut short when the client stopped + rank = rank_of_gpu.get(int(parts[stamps])) + if rank is not None: + samples[rank].append((int(parts[0]), int(parts[stamps - 1]), + tuple(int(v) for v in parts[stamps + 1:]))) + tracks = {} + for rank, rows in samples.items(): + rows.sort() + tracks[rank] = ([r[0] for r in rows], [r[1] for r in rows], [r[2] for r in rows]) + return tracks + + +def kernel_clocks(track, offset_us, ts_us, dur_us): + """Clock ranges over a kernel's lifetime from its GPU's sample track, or None.""" + if track is None or offset_us is None: + return None + starts, ends, values = track + start = (ts_us - offset_us) * 1e3 + first = bisect.bisect_right(ends, start) - 1 # last poll done before the kernel began + last = bisect.bisect_left(starts, start + dur_us * 1e3) # polls begun before it ended + if first < 0: + return None + span = values[first:last] + out = {name: [min(v[k] for v in span), max(v[k] for v in span)] + for k, name in enumerate(CLOCK_NAMES)} + reasons = 0 + for v in span: + if v[4] >= 0: + reasons |= v[4] + out.update(event_reasons=reasons, samples=last - first - 1, + prior_us=round((start - ends[first]) / 1e3, 1)) + return out + + +COPY_LAUNCHER = "vllm.v1.simple_kv_offload.copy_backend.DmaCopyBackend.launch_copy" +COPY_MAX_LAG_US = 60e6 +COPY_MAX_SPLIT = 8 # memcpys one logged copy may arrive as + + +def match_copies(trace, orphans, copies, offset_us): + """device index -> logged copy for orphan memcpys, per direction. + + Each logged copy is one memcpy on that direction's stream, or a run of + consecutive ones whose bytes add up to it (the backend splits large + copies: a 4.6 GB load arrived as 3.2 GB + 1.5 GB), executed in issue order + after its issue. Of the order-preserving pairings with every copy issued + before its first memcpy starts, take the one with least total + issue-to-start lag: copies logged before the window whose memcpys ran + before it, and copies still queued at its end, stay unpaired. + """ + matched = {} + if offset_us is None: + return matched + for store, marker in ((True, "DtoH"), (False, "HtoD")): + memcpys = sorted((i for i in orphans if marker in trace.events[i]["name"]), + key=lambda i: trace.events[i]["ts"]) + if not memcpys: + continue + # The log spans the run; a copy's memcpy starts within a minute of its issue. + lo = trace.events[memcpys[0]]["ts"] - offset_us - COPY_MAX_LAG_US + hi = trace.events[memcpys[-1]]["ts"] - offset_us + logged = [c for c in copies if c["store"] == store and lo <= c["t_ns"] / 1e3 <= hi] + n, m = len(memcpys), len(logged) + if not n or not m: + continue + inf = float("inf") + sizes = [trace.events[i].get("args", {}).get("bytes") for i in memcpys] + cost = [[0.0] * (m + 1)] + [[inf] * (m + 1) for _ in range(n)] + take = [[0] * (m + 1) for _ in range(n + 1)] # memcpys the copy j covers, ending at a + for a in range(1, n + 1): + for j in range(1, m + 1): + cost[a][j] = cost[a][j - 1] + c = logged[j - 1] + for k in range(1, min(a, COPY_MAX_SPLIT) + 1): + first = trace.events[memcpys[a - k]] + lag = first["ts"] - (c["t_ns"] / 1e3 + offset_us) + if lag < 0: + break + part = sizes[a - k:a] + if k > 1 and (c.get("bytes") is None or None in part): + break + if c.get("bytes") is not None and None not in part and sum(part) != c["bytes"]: + if sum(part) > c["bytes"]: + break + continue + if cost[a - k][j - 1] + lag < cost[a][j]: + cost[a][j] = cost[a - k][j - 1] + lag + take[a][j] = k + if cost[n][m] == inf: + continue # more memcpys than pairable copies: leave them unattributed + a, j = n, m + while a: + if take[a][j]: + for i in memcpys[a - take[a][j]:a]: + matched[i] = logged[j - 1] + a -= take[a][j] + j -= 1 + return matched + + +def kernel_family(name): + """A kernel name without template arguments, tile suffixes or C++ mangling.""" + if name.startswith("_Z"): + i, parts = (3 if name.startswith("_ZN") else 2), [] + while i < len(name) and (name[i].isdigit() or name[i] == "L" and name[i + 1:i + 2].isdigit()): + i += name[i] == "L" # internal linkage + j = i + while j < len(name) and name[j].isdigit(): + j += 1 + parts.append(name[j:j + int(name[i:j])]) + i = j + int(name[i:j]) + if parts: + return "::".join(parts) + return re.split(r"[<(]|_MT\d|_GROUP_|_BLOCK_|\.kd$", name.removeprefix("void "))[0] + + +def activity_kind(event): + """kernel / memcpy / memset, counting HIP's blit kernels as the copies they are.""" + if event["cat"] != "kernel": + return event["cat"].removeprefix("gpu_") + if event["name"].startswith("__amd_rocclr_fill"): + return "memset" + return "memcpy" if event["name"].startswith("__amd_rocclr_copy") else "kernel" + + +def join_by_stream_marks(trace, items, launches): + """Exact order for a node-id-less replay when the capture marked each launch's stream. + + Capture launches group by their infx_stream mark (None: the capture + stream); each group pairs with the one replay stream whose activity count + and kernel/memcpy/memset sequence equal it, and joins in order. Returns + the activities in capture order, or None when the groups do not pair up + one to one. + """ + groups = collections.defaultdict(list) + for p, launch in enumerate(launches): + groups[launch["marks"].get("stream")].append(p) + streams = collections.defaultdict(list) + for i in sorted(items, key=lambda i: trace.events[i]["ts"]): + streams[trace.events[i]["args"].get("stream")].append(i) + if len(groups) != len(streams): + return None + order, used = [None] * len(launches), set() + for positions in groups.values(): + kinds = [launches[p]["kind"] for p in positions] + matches = [s for s, seq in streams.items() if s not in used and len(seq) == len(positions) + and [activity_kind(trace.events[i]) for i in seq] == kinds] + if len(matches) != 1: + return None + used.add(matches[0]) + for p, i in zip(positions, streams[matches[0]]): + order[p] = i + return order + + +def align_streams(trace, items, launches, families, beam=64): + """Order a node-id-less graph replay's activities as its captured launches. + + Each stream keeps capture order, so capture order interleaves the + per-stream sequences. A beam search over the streams' heads picks the + interleaving that best matches each launch: +2 when the activity's kernel + family is one its op or launcher ran eagerly, is its Inductor kernel or its + Triton kernel, -2 when a memset or memcpy meets a kernel launch or the + reverse; ties go to device start order. Returns the ordered activities and + the number of evidence-matched positions. + """ + streams = collections.defaultdict(list) + for i in sorted(items, key=lambda i: trace.events[i]["ts"]): + streams[trace.events[i]["args"].get("stream")].append(i) + seqs = list(streams.values()) + + def score(p, i): + cap, e = launches[p], trace.events[i] + if (cap["kind"] == "kernel") != (activity_kind(e) == "kernel"): + return -2 + key, family = cap["op"] or cap["launcher"], kernel_family(e["name"]) + launcher = cap["launcher"] or "" + if (family in families.get(key, ()) or family == key + or launcher.startswith("triton:") and launcher[len("triton:"):] in e["name"]): + return 2 + return 0 + + # Each state: (score, -start-order penalty, heads, picks as a (stream, parent) chain). + states = [(0, 0, (0,) * len(seqs), None)] + for p in range(len(launches)): + best = {} + for total, penalty, heads, chain in states: + ready = sorted((trace.events[seqs[s][h]]["ts"], s) for s, h in enumerate(heads) + if h < len(seqs[s])) + for rank, (_, s) in enumerate(ready): + nxt = heads[:s] + (heads[s] + 1,) + heads[s + 1:] + state = (total + score(p, seqs[s][heads[s]]), penalty - rank, nxt, (s, chain)) + if nxt not in best or state[:2] > best[nxt][:2]: + best[nxt] = state + states = sorted(best.values(), key=lambda st: st[:2], reverse=True)[:beam] + picks, chain = [], states[0][3] + while chain: + picks.append(chain[0]) + chain = chain[1] + heads, order = [0] * len(seqs), [] + for s in reversed(picks): + order.append(seqs[s][heads[s]]) + heads[s] += 1 + return order, sum(score(p, i) > 0 for p, i in enumerate(order)) + + +def extract_replay(trace, captured, rank, window, step_log, copies, clocks, routing, out_dir, + pairs): + """Attribute every device activity of one replay trace; write one file per step. + + Returns the trace's report entry and its index entries. `pairs` accumulates + (op, kernel) pairs by source for the cross-trace kernel-name check, and + under "families" the kernel families each op or launcher ran eagerly. + """ + rows = [] + graph_replays = collections.defaultdict(list) # (graph, launch) -> device indices + orphans = [] # device activities with no CPU launch + for i in trace.device: + e = trace.events[i] + args = e.get("args", {}) + launch = trace.launches.get(args.get("correlation")) + if launch is None: + orphans.append(i) + continue + ctx = trace.context(launch) + if "Graph" in trace.events[launch]["name"]: + graph_replays[(ctx["marks"].get("graph_replay"), launch)].append(i) + continue + rows.append({"source": "eager", "device_index": i, **ctx}) + pairs["eager"][(ctx["op"], e["name"])] += 1 + pairs["families"][ctx["op"] or ctx["launcher"]].add(kernel_family(e["name"])) + + offset_us = unix_to_trace_us(trace, step_log) + offload = match_copies(trace, orphans, copies, offset_us) + # Work launched before the profiler started (the GPU still finishing the step + # before the window) runs ahead of everything the window launched. + orphan_set = set(orphans) + first_launched = min((trace.events[i]["ts"] for i in trace.device if i not in orphan_set), default=None) + unattributed = 0 + for i in orphans: + copy = offload.get(i) + if copy is None: + e = trace.events[i] + if first_launched is not None and e["ts"] + e.get("dur", 0) <= first_launched: + rows.append({"source": "pre_window", "device_index": i}) + continue + unattributed += 1 + rows.append({"source": "unattributed", "device_index": i}) + continue + step = copy.get("step") + rows.append({ + "source": "offload_copy", "device_index": i, + "marks": {"step": step} if step is not None else {}, + "launch_api": "cuMemcpyBatchAsync", "launcher": COPY_LAUNCHER, + "launcher_callers": copy.get("callers"), "op": None, "module_path": [], + "copy": {"store": copy["store"], "blocks": copy["blocks"], "bytes": copy.get("bytes"), + "issue_lag_us": trace.events[i]["ts"] - (copy["t_ns"] / 1e3 + offset_us)}, + }) + + graph_checks = collections.Counter() + interleavings = {} # (graph, per-stream kernel names) -> stream of each capture position + for (gid, launch), items in graph_replays.items(): + replay_ctx = trace.context(launch) + launches = captured.get(gid) if gid is not None else None + join = "count_mismatch" if launches is None or len(launches) != len(items) else None + if join is None and all("graph node id" in trace.events[i]["args"] for i in items): + items.sort(key=lambda i: trace.events[i]["args"]["graph node id"] & NODE_ID_MASK) + elif join is None: + streams = collections.defaultdict(list) + for i in sorted(items, key=lambda i: trace.events[i]["ts"]): + streams[trace.events[i]["args"].get("stream")].append(i) + shape = (gid, tuple(sorted((str(s), tuple(trace.events[i]["name"] for i in seq)) + for s, seq in streams.items()))) + if shape not in interleavings: + order, exact = join_by_stream_marks(trace, items, launches), True + if order is None: + order, evidence = align_streams(trace, items, launches, pairs["families"]) + exact = False + graph_checks["stream_order_aligned"] += 1 + graph_checks["stream_order_evidence"] += evidence + graph_checks["stream_order_positions"] += len(items) + interleavings[shape] = ([trace.events[i]["args"].get("stream") for i in order], exact) + queues = {s: iter(seq) for s, seq in streams.items()} + items = [next(queues[s]) for s in interleavings[shape][0]] + graph_checks["joined_by_stream_marks" if interleavings[shape][1] else "joined_by_alignment"] += 1 + if join: + graph_checks[join] += 1 + for pos, i in enumerate(items): + rows.append({"source": "graph", "graph": gid, "node_pos": pos, "device_index": i, + "marks": replay_ctx["marks"], "capture_join": join}) + continue + graph_checks["joined"] += 1 + for pos, (i, cap) in enumerate(zip(items, launches)): + rows.append({"source": "graph", "graph": gid, "node_pos": pos, "device_index": i, + **cap, "marks": replay_ctx["marks"]}) + pairs["graph"][(cap["op"], trace.events[i]["name"])] += 1 + + by_step = collections.defaultdict(list) + for row in rows: + e = trace.events[row.pop("device_index")] + args = e.get("args", {}) + row.update( + rank=rank, window=window, kernel=e["name"], cat=e["cat"], + stream=args.get("stream", e.get("tid")), device=args.get("device", e.get("pid")), + ts_us=e["ts"], dur_us=e.get("dur", 0), graph_node_id=args.get("graph node id"), + grid=args.get("grid"), block=args.get("block"), + clocks=kernel_clocks(clocks, offset_us, e["ts"], e.get("dur", 0)), + ) + key = step_key(row.get("marks") or {}) + row["step"] = list(key) if key else None + by_step[key].append(row) + + rank_dir = os.path.join(out_dir, f"window{window}", str(rank)) + os.makedirs(rank_dir, exist_ok=True) + index = [] + for key in sorted(by_step, key=lambda k: (k is None, k or ("", 0))): + kernels = sorted(by_step[key], key=lambda r: r["ts_us"]) + t0 = kernels[0]["ts_us"] + t1 = max(r["ts_us"] + r["dur_us"] for r in kernels) + summary = { + "kernels": len(kernels), "busy_us": sum(r["dur_us"] for r in kernels), + "t0_us": t0, "t1_us": t1, "span_us": t1 - t0, + "sources": dict(collections.Counter(r["source"] for r in kernels)), + } + name = f"{key[0]}{key[1]:06d}.json.gz" if key else "unstepped.json.gz" + batch = step_log.get(key) if key else None + moe = step_routing(routing, key[1], out_dir) if key and key[0] == "step" else None + with gzip.open(os.path.join(rank_dir, name), "wt") as f: + json.dump({"rank": rank, "window": window, "kind": key[0] if key else None, + "step": key[1] if key else None, "summary": summary, "batch": batch, + "routing": moe, "kernels": kernels}, f, separators=(",", ":")) + reqs = (batch or {}).get("reqs") or [] + index.append({ + "window": window, "rank": rank, "kind": key[0] if key else None, + "step": key[1] if key else None, + "file": os.path.relpath(os.path.join(rank_dir, name), out_dir), + **summary, "total_tokens": (batch or {}).get("total_tokens"), "num_reqs": len(reqs), + "graph": (batch or {}).get("graph"), "phase": (batch or {}).get("phase"), + "cudagraph": [d.get("cg_mode") for d in (batch or {}).get("dispatch") or []], + "routing": bool(moe), + }) + entry = { + "trace": os.path.basename(trace.path), "rank": rank, "window": window, + "device_activities": len(rows), "unattributed": unattributed, + "graph_replays": dict(graph_checks), + "sources": dict(collections.Counter(r["source"] for r in rows)), + "steps": sum(1 for k in by_step if k is not None), + "clocks": clock_coverage(rows), + "steps_with_routing": sum(1 for e in index if e["routing"]), + } + return entry, index + + +def clock_coverage(rows): + """How many activities carry clocks, and how stale their start state is.""" + ages = sorted(r["clocks"]["prior_us"] for r in rows if r.get("clocks")) + if not ages: + return {"activities_with_clocks": 0} + return {"activities_with_clocks": len(ages), + "prior_us_p50": ages[len(ages) // 2], "prior_us_p99": ages[int(len(ages) * 0.99)], + "prior_us_max": ages[-1]} + + +def trace_time(path): + """The export timestamp in a window trace's file name (..pt.trace...).""" + parts = os.path.basename(path).split(".") + return int(parts[1]) if len(parts) > 1 and parts[1].isdigit() else 0 + + +def main(): + profile_dir, out_dir = sys.argv[1], sys.argv[2] + os.makedirs(out_dir, exist_ok=True) + report = {"traces": []} + pairs = {"eager": collections.Counter(), "graph": collections.Counter(), + "families": collections.defaultdict(set)} + traces = sorted(glob.glob(os.path.join(profile_dir, "torch", "**", "*.pt.trace.json*"), + recursive=True), key=trace_time) + # Capture traces record every CUDA graph size the engine captured (GBs with + # many sizes); load only the graphs the windows replay. A rank's replays join + # its own capture when that rank's was profiled (DP-attention ranks' graphs can + # differ: per-rank buffers change copy kernels), else the first profiled rank's. + needed = replayed_graphs(traces) + sessions = collections.defaultdict(list) # rank -> its capture traces, in capture order + for path in glob.glob(os.path.join(profile_dir, "capture", "*.json*")): + name = os.path.basename(path).split(".") + session = int(name[1]) if len(name) > 1 and name[1].isdigit() else 0 + sessions[name[0]].append((session, path)) + captures = {} # rank -> graph ordinal -> captured launches + for capture_rank in sorted(sessions): + captured = {} + paths = [path for _, path in sorted(sessions[capture_rank])] + for path in paths: + ranges = capture_ranges(path, needed) + trace = Trace(path, overlapping(ranges)) if ranges else None + if trace: + captured.update(capture_launches(trace)) + del trace + captures[capture_rank] = captured + report.setdefault("capture", {})[capture_rank] = { + "paths": [os.path.relpath(p, profile_dir) for p in paths], "graphs": len(captured), + "graphs_replayed": len(needed), "launches": sum(map(len, captured.values()))} + default_capture = captures[min(captures)] if captures else {} + ranks = rank_of_pid(profile_dir) + steps = load_steps(profile_dir) + copies = load_copies(profile_dir, ranks) + clocks = load_clocks(profile_dir) + routing = load_routing(profile_dir) + report["routing"] = {rank: {"bound_layers": len(meta.get("bound", {})), "failed": meta.get("failed")} + for rank, (_, _, meta) in sorted(routing.items())} + windows_seen = collections.Counter() + index = [] + for path in traces: + trace = Trace(path) + rank = next((ranks[p] for p in trace.process_ids() if p in ranks), None) + window = windows_seen[rank] + windows_seen[rank] += 1 + entry, entries = extract_replay(trace, captures.get(rank, default_capture), rank, window, + steps.get(rank, {}), + copies.get(rank, []), clocks.get(rank), + routing.get(rank), out_dir, pairs) + report["traces"].append(entry) + index += entries + del trace + # A graph-replayed kernel whose op never launched that kernel eagerly anywhere. + eager_ops = {op for op, _ in pairs["eager"]} + report["graph_kernels_unseen_eagerly"] = sorted( + ([op, kernel, n] for (op, kernel), n in pairs["graph"].items() + if op is not None and op in eager_ops and (op, kernel) not in pairs["eager"]), + key=lambda item: -item[2]) + with open(os.path.join(out_dir, "index.json"), "w") as f: + json.dump({"capture": report.get("capture"), "steps": index}, f) + with open(os.path.join(out_dir, "report.json"), "w") as f: + json.dump(report, f, indent=1) + print(json.dumps(report, indent=1)) + if traces and not sum(t["device_activities"] for t in report["traces"]): + sys.exit(f"no device activity attributed across {len(traces)} traces") + + +if __name__ == "__main__": + main() diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/__init__.py b/inferencex-e2e/benchmarks/profiling/infx_prof/__init__.py new file mode 100644 index 0000000000..e78afdda37 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/__init__.py @@ -0,0 +1,81 @@ +"""Op-attribution profiling for inference engines, loaded through PYTHONPATH (sitecustomize). + +Inactive unless INFX_PROF_DIR is set. INFX_PROF_ENGINE (vllm, sglang, trtllm) +selects the adapter. The core (this package) records, in every engine process: + +- graphs: CUDA graph capture/replay ordinals and stream markers; +- modules: module markers for eager code and Inductor-compiled code; +- launchers: Python kernel launchers (Triton, TileLang, CuTe, vendored packages); +- steps: per-step batch records and markers, the file window controller; +- routing: per-step MoE top-k ids (device buffer, graph-safe); +- windows: profiler windows and the graph-capture profile. + +The adapter (engines/.py) wires these to the engine's step loop, graph +capture, models and MoE layers. Outputs go under $INFX_PROF_DIR: env/, +capture/, graphs/, steps/, profiler/, torch/, routing/, copies/, errors/, and +stacks/.log: every thread's Python stack every STACK_DUMP_S seconds (300), to +see where a hung engine is. +Nothing here may break serving: every hook degrades to the original call. +""" + +import atexit +import importlib +import os + +from . import graphs, imports, launchers, modules, runtime, steps, windows + +ENGINES = {"vllm": "vllm", "sglang": "sglang", "trtllm": "trtllm", "trt": "trtllm"} +STACK_DUMP_S = int(os.environ.get("INFX_PROF_STACK_DUMP_S") or 300) + + +def dump_stacks_periodically(): + """A daemon thread writes every thread's stack to stacks/.log every STACK_DUMP_S seconds. + + Not faulthandler.dump_traceback_later: it holds a single timer per process, + which an engine re-arming its own (SGLang) silently replaces. The thread runs + whenever the hung code releases the GIL (GPU and collective waits do). + """ + import faulthandler + import threading + import time + + if STACK_DUMP_S <= 0: + return + + def run(): + path = os.path.join(runtime.DIR, "stacks", f"{os.getpid()}.log") + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "a", buffering=1) as stream: + while True: + time.sleep(STACK_DUMP_S) + stream.write(f"=== {time.strftime('%Y-%m-%dT%H:%M:%S')} pid {os.getpid()}\n") + faulthandler.dump_traceback(file=stream, all_threads=True) + + def start(): + threading.Thread(target=run, name="infx_prof_stacks", daemon=True).start() + + start() + os.register_at_fork(after_in_child=start) # fork does not copy threads + + +def boot(): + if not runtime.DIR: + return + try: + name = ENGINES.get(os.environ.get("INFX_PROF_ENGINE", "vllm"), "vllm") + engine = importlib.import_module(f".engines.{name}", __name__).ENGINE + registry = imports.HookRegistry() + graphs.install(registry) + modules.install(registry) + launchers.install(registry, engine.caller_roots, engine.launcher_modules, engine.launcher_packages) + if engine.window_control == "file": + steps.controller = windows.FileWindows() + engine.install(registry) + registry.install() + atexit.register(runtime.flush_all) + except Exception: + runtime.write_error("boot") + try: + dump_stacks_periodically() + except Exception: + runtime.write_error("stack dumps") diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/engines/__init__.py b/inferencex-e2e/benchmarks/profiling/infx_prof/engines/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/engines/base.py b/inferencex-e2e/benchmarks/profiling/infx_prof/engines/base.py new file mode 100644 index 0000000000..397b71b1a0 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/engines/base.py @@ -0,0 +1,28 @@ +"""What an engine adapter declares. + +An adapter is a module-level ``Engine`` subclass instance. The core installs +its own hooks (graph ordinals, stream and module markers, Inductor lines, +DSL launchers) and then the adapter's: + +- ``caller_roots``: source path segments of the engine's frames, quoted in + launcher markers; +- ``launcher_modules`` / ``launcher_packages``: where Python kernel launchers + live (packages: every module under them, private functions included); +- ``window_control``: "engine" when the engine's own profiler opens windows + (the adapter logs its start/stop), "file" for the core's FileWindows; +- ``install(registry)``: the engine-specific patches (step boundary and + record, graph capture session, module roots, routing, rank naming). +""" + + +class Engine: + name = None + caller_roots = () + launcher_modules = () + launcher_packages = () + version_packages = () + env_prefixes = () + window_control = "file" + + def install(self, registry): + raise NotImplementedError diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/engines/sglang.py b/inferencex-e2e/benchmarks/profiling/infx_prof/engines/sglang.py new file mode 100644 index 0000000000..a0ca98e152 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/engines/sglang.py @@ -0,0 +1,396 @@ +"""SGLang adapter (one scheduler process per TP rank; v0.5.12 through the runner/ refactor and main). + +Hook points every supported release has: + +- Rank: run_scheduler_process's tp_rank / dp_rank arguments (SGLANG_DP_RANK + when the launcher sets it instead), patched before the spawned scheduler + process calls it. +- Step: Scheduler.run_batch (every rank, every step; an IDLE batch on an idle + DP-attention rank is a dummy step). The forward, speculative draft and verify + passes included, runs synchronously inside it on the scheduler thread. The + record's requests come from the ScheduleBatch; its token count from the + target ModelRunner.forward's batch (a verify step's bs x draft tokens). +- Graph mode: which target CUDA graph runner replayed (replay() through + v0.5.12, execute() from the runner/ refactor): the decode runner FULL; the + prefill runner by its backend (BREAKABLE, PIECEWISE, FULL); NONE otherwise. +- Windows: the core FileWindows, at run_batch boundaries. +- Graph capture: every *CudaGraphRunner.capture (and the breakable runner's + _capture_all) runs inside the core capture session. +- Modules: ModelRunner.load_model registers its model, as "draft" for draft + workers (whose MoE layers record no routing). +- Routing: TopK / HashTopK outputs, eagerly (Mega-MoE routes through TopK + without FusedMoE); and every MoE layer's routing reaches FusedMoE.forward_impl (or + forward_deferred_finalize), whoever produced it (TopK, DeepSeek-V4's HashTopK, + DeepSeek-V4.1's vision_topk) and however the layer was called (eager, the + piecewise custom op, the dual-stream graph path). Materialised ids go to the + core capturer, compile-safe (piecewise prefill traces the call). When the MoE + kernel routes itself (a BypassedTopKOutput: FlashInfer TRT-LLM-gen MoE), the + ids are recomputed eagerly with SGLang's own select_experts + (BypassedTopKOutput.to_standard) under an infx_prof.routing.recompute + launcher marker. DP attention with a gathered (non all-to-all) MoE routes + every DP rank's tokens in a padded layout; its routing is not recorded. +""" + +import functools +import inspect +import itertools +import os +import sys +import threading + +from .. import envinfo, graphs, launchers, modules, routing, runtime, steps, windows +from .base import Engine + +_rank = {"tag": None, "dp": 0} +_tls = threading.local() + + +def rank_tag(): + return _rank["tag"] or f"pid{os.getpid()}" + + +def patch_scheduler_module(module): + run_process = getattr(module, "run_scheduler_process", None) + if run_process is not None and not getattr(run_process, "_infx_patched", False): + signature = inspect.signature(run_process) + + @functools.wraps(run_process) + def run_scheduler_process(*args, **kwargs): + try: + bound = signature.bind_partial(*args, **kwargs).arguments + dp = bound.get("dp_rank") + if dp is None: + dp = os.environ.get("SGLANG_DP_RANK") + _rank["dp"] = int(dp or 0) + _rank["tag"] = f"dp{_rank['dp']}_tp{bound.get('tp_rank') or 0}" + except Exception: + runtime.write_error("sglang rank") + return run_process(*args, **kwargs) + + run_scheduler_process._infx_patched = True + module.run_scheduler_process = run_scheduler_process + + cls = getattr(module, "Scheduler", None) + if cls is None or getattr(cls, "_infx_patched", False) or not hasattr(cls, "run_batch"): + return + orig = cls.run_batch + counter = itertools.count() + + def run_batch(self, batch, *args, **kwargs): + try: + record = batch_record(batch) + except Exception: + runtime.write_error("sglang step record") + record = {} + kind = "dummy" if record.get("phase") == "idle" else "step" + _tls.graph, _tls.forward = "NONE", None + with steps.Step(rank_tag(), next(counter), record, kind=kind) as rec: + try: + return orig(self, batch, *args, **kwargs) + finally: + rec["graph"] = _tls.graph + rec["forward_ct"] = getattr(self, "forward_ct", None) + try: + apply_forward(rec, _tls.forward) + except Exception: + runtime.write_error("sglang step forward") + + cls.run_batch = run_batch + cls._infx_patched = True + + +def batch_record(batch): + mode = getattr(batch, "forward_mode", None) + mode_name = getattr(mode, "name", str(mode)) + phase = {"EXTEND": "prefill", "DECODE": "decode", "MIXED": "mixed", "IDLE": "idle", + "TARGET_VERIFY": "decode"}.get(mode_name, mode_name.lower() if mode_name else None) + reqs = list(getattr(batch, "reqs", None) or ()) + extend_lens = list(getattr(batch, "extend_lens", None) or ()) + prefix_lens = list(getattr(batch, "prefix_lens", None) or ()) + seq_lens = getattr(batch, "seq_lens_cpu", None) + seq_lens = seq_lens.tolist() if hasattr(seq_lens, "tolist") else list(seq_lens or ()) + rows, total = [], 0 + for i, req in enumerate(reqs): + prompt = getattr(req, "origin_input_ids", None) + if phase in ("prefill", "mixed") and i < len(extend_lens): + scheduled, computed = int(extend_lens[i]), int(prefix_lens[i]) if i < len(prefix_lens) else None + else: + scheduled = 1 + # seq_lens: the KV the step attends over (output_ids lags a step under overlap) + computed = int(seq_lens[i]) if i < len(seq_lens) else None + total += scheduled + rows.append({"req": getattr(req, "rid", None), "scheduled": scheduled, "computed": computed, + "prompt_len": len(prompt) if prompt is not None else None, + "new": phase in ("prefill", "mixed") and not computed, "spec": 0}) + extend_tokens = getattr(batch, "extend_num_tokens", None) + if phase in ("prefill", "mixed") and extend_tokens: + total = int(extend_tokens) + return {"phase": phase, "mode": mode_name, "total_tokens": total, "reqs": rows} + + +def apply_forward(rec, forward): + """Token counts from the target model's forward batch (the verify pass under speculation).""" + if not forward: + return + rec["total_tokens"] = forward["tokens"] + rec["forward_mode"] = forward["mode"] + draft = forward.get("draft_tokens") + if forward["mode"] == "TARGET_VERIFY" and draft: + for row in rec.get("reqs") or (): + row["scheduled"], row["spec"] = int(draft), int(draft) - 1 + + +def patch_model_runner(module): + cls = getattr(module, "ModelRunner", None) + if cls is None or getattr(cls, "_infx_patched", False): + return + orig_load, orig_forward = cls.load_model, cls.forward + + def load_model(self, *args, **kwargs): + result = orig_load(self, *args, **kwargs) + draft = bool(getattr(self, "is_draft_worker", False)) + try: + model = getattr(self, "model", None) + modules.register_root("draft" if draft else "model", model) + mark_draft(model, draft) + if not draft: + envinfo.write(rank_tag(), ("sglang", "sgl_kernel"), ("SGLANG_", "SGL_"), + {"runner_class": f"{type(self).__module__}.{type(self).__qualname__}", + "server_args": str(getattr(self, "server_args", None))}) + setup_routing(self) + except Exception: + runtime.write_error("sglang load_model hooks") + return result + + def forward(self, forward_batch, *args, **kwargs): + if not getattr(self, "is_draft_worker", False): + try: + mode = getattr(getattr(forward_batch, "forward_mode", None), "name", "") + ids = getattr(forward_batch, "input_ids", None) + spec = getattr(forward_batch, "spec_info", None) + _tls.forward = {"mode": mode, "tokens": int(ids.shape[0]) if ids is not None else None, + "draft_tokens": getattr(spec, "draft_token_num", None)} + except Exception: + pass + return orig_forward(self, forward_batch, *args, **kwargs) + + cls.load_model = load_model + cls.forward = forward + cls._infx_patched = True + + +def mark_draft(model, draft): + """Draft-model modules: their MoE layers reuse the target's layer numbers, so they record no routing.""" + if not draft or model is None or not hasattr(model, "named_modules"): + return + for _, mod in model.named_modules(): + try: + mod._infx_draft = True + except Exception: + pass + + +def _graph_label(runner): + name = type(runner).__name__ + if not any(k in name for k in ("Prefill", "Piecewise", "Breakable")): + return "FULL" # decode runners + backend = type(getattr(runner, "backend", None)).__name__ if getattr(runner, "backend", None) is not None else name + for key, label in (("Breakable", "BREAKABLE"), ("Piecewise", "PIECEWISE"), ("Full", "FULL")): + if key in backend or key in name: + return label + return "PIECEWISE" + + +def _is_target(runner): + model_runner = getattr(runner, "model_runner", None) + return model_runner is not None and not getattr(model_runner, "is_draft_worker", False) + + +def patch_graph_runners(module): + """Graph runner classes: capture inside the core capture session; replays label the step.""" + for name, cls in list(vars(module).items()): + if not isinstance(cls, type) or cls.__module__ != module.__name__ or not name.endswith("CudaGraphRunner"): + continue + for method in ("capture", "_capture_all"): + orig = vars(cls).get(method) + if orig is None or getattr(orig, "_infx_patched", False): + continue + + def capture(self, *args, _orig=orig, **kwargs): + with windows.capture_session(rank_tag()): + return _orig(self, *args, **kwargs) + + capture._infx_patched = True + setattr(cls, method, capture) + if method == "capture": + break # a runner with capture() calls its own _capture_all inside it + for method in ("replay", "execute"): + orig = vars(cls).get(method) + if orig is None or getattr(orig, "_infx_patched", False): + continue + + def run(self, *args, _orig=orig, **kwargs): + before = graphs.replays() + result = _orig(self, *args, **kwargs) + if graphs.replays() > before and _is_target(self): + _tls.graph = _graph_label(self) + return result + + run._infx_patched = True + setattr(cls, method, run) + + +def setup_routing(runner): + config = getattr(getattr(runner, "model_config", None), "hf_config", None) + text = getattr(config, "text_config", None) or config + num_layers = getattr(text, "num_hidden_layers", None) + topk = getattr(text, "num_experts_per_tok", None) + experts = (getattr(text, "n_routed_experts", None) or getattr(text, "num_experts", None) + or getattr(text, "num_local_experts", None)) + if not (num_layers and topk and experts): + return + args = getattr(runner, "server_args", None) + max_tokens = max(int(getattr(args, "chunked_prefill_size", 0) or 0), int(getattr(args, "max_prefill_tokens", 0) or 0), + 8192) + try: + from sglang.srt.distributed.parallel_state import get_attn_tp_group + + attn_tp_rank = get_attn_tp_group().rank_in_group + except Exception: + runtime.write_error("sglang attention TP rank") + attn_tp_rank = int(getattr(runner, "tp_rank", 0) or 0) + extra = {"source": "sglang TopK outputs; bypassed routing recomputed with select_experts"} + write = attn_tp_rank == 0 + gathered = bool(getattr(args, "enable_dp_attention", False)) and getattr(args, "moe_a2a_backend", "none") in ( + None, "none") + if gathered: + write = False + extra["unavailable"] = ("DP attention with a gathered MoE: routing rows are every DP rank's tokens " + "in a padded layout, not recorded") + routing.setup(rank_tag(), max_tokens, int(num_layers), int(topk), int(experts), write=write, extra=extra) + + +def _ids(result): + ids = getattr(result, "topk_ids", None) + if ids is None: + packed = getattr(result, "packed_topk_ids", None) + if packed is not None: + ids = packed >> 16 # PackedTopKOutput: expert id in the high half + return ids # BypassedTopKOutput / TritonKernelTopKOutput carry none + + +def _topk_output(args, kwargs): + return args[1] if len(args) > 1 else kwargs.get("topk_output") + + +def capture_moe_routing(moe, topk_output): + """Record the top-k ids an experts layer is handed; recompute them when routing is left to its kernel.""" + import torch + + layer = getattr(moe, "layer_id", None) + if not isinstance(layer, int) or getattr(moe, "_infx_draft", False) or topk_output is None: + return + if torch.compiler.is_compiling(): # piecewise prefill traces the experts call: copy only, compile-safe + routing.capture(layer, _ids(topk_output)) + return + if not routing.active() or getattr(_tls, "routed", None) is topk_output: # captured at TopK, or nested call + return + _tls.routed = topk_output + try: + ids = _ids(topk_output) + if ids is None and type(topk_output).__name__ == "BypassedTopKOutput": + with launchers.marker("infx_prof.routing.recompute", skip=2): + ids = _ids(topk_output.to_standard(layer)) + routing.capture(layer, ids) + except Exception: + runtime.write_error("sglang moe routing capture") + + +def capture_topk(topk, result): + """TopK / HashTopK outputs, for MoE paths that do not go through FusedMoE.forward_impl + (Mega-MoE). Eager only: compiled paths are captured at the experts layer, which skips + an output already captured here.""" + import torch + + if torch.compiler.is_compiling() or not routing.active(): + return + layer = getattr(topk, "layer_id", None) + ids = _ids(result) + if not isinstance(layer, int) or getattr(topk, "_infx_draft", False) or ids is None: + return # bypassed outputs carry no ids: the experts layer recomputes them + _tls.routed = result + try: + routing.capture(layer, ids) + except Exception: + runtime.write_error("sglang topk capture") + + +def patch_topk(module): + """TopK's leaf forward_* methods (its forward dispatches to one) and HashTopK.forward.""" + for cls_name, methods in (("TopK", ("forward_native", "forward_cuda", "forward_cpu", "forward_npu", + "forward_xpu", "forward_musa")), + ("HashTopK", ("forward",))): + cls = getattr(module, cls_name, None) + if cls is None or getattr(cls, "_infx_patched", False): + continue + for method in methods: + orig = vars(cls).get(method) + if orig is None: + continue + + def wrapped(self, *args, _orig=orig, **kwargs): + result = _orig(self, *args, **kwargs) + capture_topk(self, result) + return result + + setattr(cls, method, wrapped) + cls._infx_patched = True + + +def patch_fused_moe(module): + """FusedMoE and its subclasses: every MoE layer's routing arrives at forward_impl (or + forward_deferred_finalize), whichever producer built it (TopK, HashTopK, a model's own + routing such as DeepSeek-V4.1's vision_topk) and whichever path called it (eager, + piecewise custom op, dual-stream graph).""" + layer_module = sys.modules.get("sglang.srt.layers.moe.fused_moe_triton.layer") + base = getattr(layer_module, "FusedMoE", None) + if base is None: + return + for cls in list(vars(module).values()): + if not isinstance(cls, type) or cls.__module__ != module.__name__ or not issubclass(cls, base): + continue + for method in ("forward_impl", "forward_deferred_finalize"): + orig = vars(cls).get(method) + if orig is None or getattr(orig, "_infx_patched", False): + continue + + def wrapped(self, *args, _orig=orig, **kwargs): + capture_moe_routing(self, _topk_output(args, kwargs)) + return _orig(self, *args, **kwargs) + + wrapped._infx_patched = True + setattr(cls, method, wrapped) + + +class SglangEngine(Engine): + name = "sglang" + caller_roots = ("sglang", "sgl_kernel") + # device_communicators: PyNccl and the custom all-reduces, which CUDA graphs capture + # (NCCL through ctypes, no torch op). aiter: AMD's kernels, e.g. its fused all-reduce. + launcher_packages = ("flashinfer", "sgl_kernel", "sglang.jit_kernel", "sglang.kernels", "deep_gemm", + "flash_mla", "flash_attn", "sglang.srt.distributed.device_communicators", "aiter") + version_packages = ("sglang", "sgl_kernel") + env_prefixes = ("SGLANG_", "SGL_") + window_control = "file" + + def install(self, registry): + registry.on("sglang.srt.managers.scheduler", patch_scheduler_module) + registry.on("sglang.srt.model_executor.model_runner", patch_model_runner) + registry.on_package("sglang.srt.layers.moe", patch_fused_moe) + registry.on("sglang.srt.layers.moe.topk", patch_topk) + registry.on("sglang.srt.layers.moe.hash_topk", patch_topk) + registry.on_package("sglang.srt.model_executor", patch_graph_runners) + registry.on_package("sglang.srt.speculative", patch_graph_runners) + + +ENGINE = SglangEngine() diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/engines/trtllm.py b/inferencex-e2e/benchmarks/profiling/infx_prof/engines/trtllm.py new file mode 100644 index 0000000000..079ec0a256 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/engines/trtllm.py @@ -0,0 +1,314 @@ +"""TensorRT-LLM PyTorch-backend adapter (trtllm-serve; one MPI process per rank). + +- Rank, env, module roots: PyExecutor.__init__ (model_engine.mapping; the target + model, which carries a one-model Eagle3 draft as a submodule, and a + two-model draft engine's). +- Step: PyExecutor._forward_step, once per iteration per rank, on the executor + thread, skipping the KV-estimation executor's serving pass (is_warmup). The + record comes from ScheduledRequests (attention-DP and CUDA-graph padding + dummies excluded; a step of only dummies is a dummy step). Graph mode: FULL + when CUDAGraphRunner.replay ran, PIECEWISE when only torch.compile's + piecewise graphs replayed (prefill-bearing steps up to the largest + capture_num_tokens), else NONE. +- Windows: the core FileWindows, at _forward_step boundaries (TRT-LLM has no + on-demand profiler). +- Graph capture: PyTorchModelEngine.warmup runs inside the core capture + session. It runs twice (the KV-estimation executor, then the final one); + graphs are released in between, so the second capture holds the live graphs. +- Routing: MoE backends' forward_impl names the layer (layer_idx); it runs + inside TRT-LLM's MoE custom op, opaque to torch.compile, so eagerly and + during graph capture. A routing method's apply hands its ids to the core + capturer when the backend calls it (attention DP, load balancer, separated + routing). The TRTLLM-gen backend otherwise routes inside its kernel; for + routing methods with a known closed form (softmax / sigmoid + correction bias + top-k) the ids are then recomputed from the layer's router logits under an + ``infx_py#infx_prof.routing.recompute`` marker; other methods leave routing + unavailable (meta.json says so). +""" + +import itertools +import os +import sys +import threading + +from .. import envinfo, graphs, launchers, modules, routing, runtime, steps, windows +from .base import Engine + +_rank = {"tag": None} +_tls = threading.local() + + +def rank_tag(): + return _rank["tag"] or f"pid{os.getpid()}" + + +def _tag(mapping): + tp = getattr(mapping, "tp_rank", 0) or 0 + pp = getattr(mapping, "pp_rank", 0) or 0 + return f"dp0_tp{tp}" if not pp else f"dp0_pp{pp}_tp{tp}" + + +def setup_routing(engine): + """Allocate the routing buffer once, before the first warmup captures graphs that write it.""" + if routing.active(): + return + config = getattr(getattr(engine, "model", None), "config", None) + config = getattr(config, "text_config", None) or config + num_layers = getattr(config, "num_hidden_layers", None) + topk = getattr(config, "num_experts_per_tok", None) or getattr(config, "moe_topk", None) + experts = (getattr(config, "n_routed_experts", None) or getattr(config, "num_experts", None) + or getattr(config, "num_local_experts", None)) + if not (num_layers and topk and experts): + return + mapping = getattr(engine, "mapping", None) + adp = bool(getattr(mapping, "enable_attention_dp", False)) + write = adp or (getattr(mapping, "tp_rank", 0) or 0) == 0 # attention DP: every rank routes its own tokens + max_tokens = int(getattr(engine, "max_num_tokens", 0) or 8192) + routing.setup(rank_tag(), max_tokens, int(num_layers), int(topk), int(experts), write=write, + extra={"source": "TRT-LLM routing_method.apply, else recomputed from router logits"}) + + +def patch_py_executor(module): + cls = getattr(module, "PyExecutor", None) + if cls is None or getattr(cls, "_infx_patched", False): + return + orig_init, orig_step = cls.__init__, cls._forward_step + counter = itertools.count() + + def __init__(self, *args, **kwargs): + orig_init(self, *args, **kwargs) + try: + engine = self.model_engine + _rank["tag"] = _tag(getattr(engine, "mapping", None)) + modules.register_root("model", getattr(engine, "model", None)) + draft = getattr(self, "draft_model_engine", None) + if draft is not None: + modules.register_root("draft", getattr(draft, "model", None)) + envinfo.write(rank_tag(), ("tensorrt_llm",), ("TLLM_", "TRTLLM_"), + {"model_engine_class": f"{type(engine).__module__}.{type(engine).__qualname__}", + "llm_args": str(getattr(engine, "llm_args", None))}) + except Exception: + runtime.write_error("trtllm executor init hooks") + + def _forward_step(self, scheduled_requests, *args, **kwargs): + if getattr(self, "is_warmup", False): + return orig_step(self, scheduled_requests, *args, **kwargs) + try: + record = requests_record(scheduled_requests) + except Exception: + runtime.write_error("trtllm step record") + record = {} + kind = "dummy" if record.get("phase") == "idle" else "step" + _tls.full = False + replays = graphs.replays() + with steps.Step(rank_tag(), next(counter), record, kind=kind) as rec: + try: + return orig_step(self, scheduled_requests, *args, **kwargs) + finally: + rec["graph"] = ("FULL" if _tls.full else "PIECEWISE" if graphs.replays() > replays else "NONE") + rec["iter"] = getattr(self, "iter_counter", None) + + cls.__init__ = __init__ + cls._forward_step = _forward_step + cls._infx_patched = True + + +def _draft_len(req): + # The C++ draft_tokens are sized for this step before scheduling; py_draft_tokens + # arrive from the sampler one step late under the overlap scheduler. + for attr in ("draft_tokens", "py_draft_tokens"): + try: + tokens = getattr(req, attr, None) + if tokens is not None: + return len(tokens) + except Exception: + continue + return 0 + + +def requests_record(scheduled): + rows, total = [], 0 + context = list(getattr(scheduled, "context_requests", None) or ()) + generation = list(getattr(scheduled, "generation_requests", None) or ()) + dummies = 0 + for req, is_context in [(r, True) for r in context] + [(r, False) for r in generation]: + if getattr(req, "is_dummy", False) or getattr(req, "is_attention_dp_dummy", False) \ + or getattr(req, "is_cuda_graph_dummy", False): + dummies += 1 + continue + if is_context: + draft = 0 + scheduled_tokens = int(getattr(req, "context_chunk_size", 0) or 0) + computed = int(getattr(req, "context_current_position", 0) or 0) + else: + draft = _draft_len(req) + scheduled_tokens = 1 + draft + try: + computed = int(req.get_num_tokens(0)) # host-side count; lags a step under overlap + except Exception: + computed = None + total += scheduled_tokens + rows.append({"req": getattr(req, "py_request_id", None), "scheduled": scheduled_tokens, "computed": computed, + "prompt_len": getattr(req, "py_prompt_len", None), "new": is_context and not computed, + "spec": draft}) + phase = ("idle" if not rows and dummies else "mixed" if context and generation + else "prefill" if context else "decode") + return {"phase": phase, "total_tokens": total, "reqs": rows, "dummies": dummies} + + +def patch_model_engine(module): + cls = getattr(module, "PyTorchModelEngine", None) + if cls is None or getattr(cls, "_infx_patched", False) or not hasattr(cls, "warmup"): + return + orig = cls.warmup + + def warmup(self, *args, **kwargs): + if getattr(self, "is_draft_model", False): + return orig(self, *args, **kwargs) + try: + _rank["tag"] = _tag(getattr(self, "mapping", None)) + setup_routing(self) + except Exception: + runtime.write_error("trtllm routing setup") + with windows.capture_session(rank_tag()): + return orig(self, *args, **kwargs) + + cls.warmup = warmup + cls._infx_patched = True + + +def patch_cuda_graph_runner(module): + cls = getattr(module, "CUDAGraphRunner", None) + if cls is None or getattr(cls, "_infx_patched", False) or not hasattr(cls, "replay"): + return + orig = cls.replay + + def replay(self, *args, **kwargs): + _tls.full = True + return orig(self, *args, **kwargs) + + cls.replay = replay + cls._infx_patched = True + + +def _recompute_ids(method, logits): + """Top-k ids of routing methods with a closed form, else None. Selection only: weights are not needed.""" + import torch + + name = type(method).__name__ + k = getattr(method, "top_k", None) + if not k: + return None + bias = getattr(method, "e_score_correction_bias", None) + if bias is not None and getattr(method, "n_group", 1) in (None, 1): + scores = torch.sigmoid(logits) + bias # sigmoid + correction bias (MiniMax get_scores, ungrouped DSv3) + elif name in ("DefaultMoeRoutingMethod", "RenormalizeMoeRoutingMethod", "RenormalizeNaiveMoeRoutingMethod"): + scores = logits.float() # softmax is monotonic: top-k of the logits + else: + return None + return torch.topk(scores, int(k), dim=-1).indices.int() + + +def patch_moe_backends(module): + """Every MoE class's forward_impl in this module names its layer and fills in fused routing.""" + interface = __import__("sys").modules.get("tensorrt_llm._torch.modules.fused_moe.interface") + base = getattr(interface, "MoE", None) + for cls in list(vars(module).values()): + if not isinstance(cls, type) or cls.__module__ != module.__name__ or "forward_impl" not in vars(cls): + continue + if base is not None and not issubclass(cls, base) and cls.__name__ != "ConfigurableMoE": + continue + orig = vars(cls)["forward_impl"] + if getattr(orig, "_infx_patched", False): + continue + + def forward_impl(self, x, router_logits, *args, _orig=orig, **kwargs): + layer = getattr(self, "layer_idx", None) + if layer is None or getattr(_tls, "layer", None) is not None: # inner backend of a ConfigurableMoE + return _orig(self, x, router_logits, *args, **kwargs) + _tls.layer, _tls.routed = layer, False + try: + result = _orig(self, x, router_logits, *args, **kwargs) + if not _tls.routed and routing.active() and not runtime.compiling(): + try: + with launchers.marker("infx_prof.routing.recompute", skip=2): + ids = _recompute_ids(getattr(self, "routing_method", None), router_logits) + if ids is not None: + routing.capture(int(layer), ids) + except Exception: + runtime.write_error("trtllm routing recompute") + return result + finally: + _tls.layer = None + + forward_impl._infx_patched = True + cls.forward_impl = forward_impl + + +def patch_routing_methods(module): + base = getattr(module, "BaseMoeRoutingMethod", None) + if base is None: + return + for cls in [base] + [c for c in vars(module).values() if isinstance(c, type) and issubclass(c, base)]: + orig = vars(cls).get("apply") + if orig is None or getattr(orig, "_infx_patched", False): + continue + + def apply(self, *args, _orig=orig, **kwargs): + result = _orig(self, *args, **kwargs) + layer = getattr(_tls, "layer", None) + if layer is not None and routing.active(): + try: + routing.capture(int(layer), result[0]) + _tls.routed = True + except Exception: + pass + return result + + apply._infx_patched = True + cls.apply = apply + + +def patch_thop(module): + """The C++ attention entry points TRT-LLM's FMHA libraries call as nanobind functions, not torch ops.""" + internal = getattr(getattr(module, "bindings", None), "internal", None) + thop = getattr(internal, "thop", None) + if thop is not None: + launchers.wrap_builtins(thop, "tensorrt_llm.bindings.internal.thop") + # Wrapped at the source, before modules from-import them (also patched there: patch_kv_copies). + kv_utils = getattr(getattr(internal, "batch_manager", None), "kv_cache_manager_v2_utils", None) + if kv_utils is not None: + launchers.wrap_builtins(kv_utils, "tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils") + + +def patch_kv_copies(module): + """KV cache manager v2's block-offset and host-tier copies (nanobind functions bound by name).""" + launchers.wrap_builtins(module, module.__name__, names={ + "copy_batch_block_offsets_to_device", "copy_host_to_device", "copy_device_to_host", + "copy_device_to_device"}) + + +class TrtllmEngine(Engine): + name = "trtllm" + caller_roots = ("tensorrt_llm",) + # fmha_sm100: MiniMax-M3's sparse attention (MSA); attention_backend.fmha: the FMHA + # libraries' private launchers of FlashInfer's TRT-LLM-gen kernels. + launcher_packages = ("flashinfer", "fmha_sm100", "tensorrt_llm.deep_gemm", "tensorrt_llm._torch.cute_dsl_kernels", + "tensorrt_llm._torch.cuda_tile_kernels", "tensorrt_llm._torch.attention_backend.fmha") + version_packages = ("tensorrt_llm",) + env_prefixes = ("TLLM_", "TRTLLM_") + window_control = "file" + + def install(self, registry): + registry.on("tensorrt_llm._torch.pyexecutor.py_executor", patch_py_executor) + registry.on("tensorrt_llm._torch.pyexecutor.model_engine", patch_model_engine) + registry.on("tensorrt_llm._torch.pyexecutor.cuda_graph_runner", patch_cuda_graph_runner) + registry.on_package("tensorrt_llm._torch.modules.fused_moe", patch_moe_backends) + registry.on("tensorrt_llm._torch.modules.fused_moe.routing", patch_routing_methods) + registry.on("tensorrt_llm", patch_thop) + registry.on("tensorrt_llm._torch.attention_backend.trtllm", lambda m: patch_thop(sys.modules["tensorrt_llm"])) + registry.on("tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2", patch_kv_copies) + registry.on("tensorrt_llm.runtime.kv_cache_manager_v2._copy_engine", patch_kv_copies) + + +ENGINE = TrtllmEngine() diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/engines/vllm.py b/inferencex-e2e/benchmarks/profiling/infx_prof/engines/vllm.py new file mode 100644 index 0000000000..8ff6efcc3e --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/engines/vllm.py @@ -0,0 +1,386 @@ +"""vLLM adapter (V1 and V2 GPU model runners). + +- Step: GPUModelRunner.execute_model (scheduled steps) and Worker.execute_dummy_batch + (idle DP ranks); the step's cudagraph mode comes from the dispatch log + (V2: dp_utils.dispatch_cg_and_sync_dp, V1: CudagraphDispatcher.dispatch, last call wins). +- Windows: vLLM's own profiler (--profiler-config, opened by the window client's + /start_profile); its WorkerProfiler start/stop are logged and toggle the markers. +- Graph capture: GPUModelRunner.capture_model runs inside the core capture session. +- Modules: the target model and the speculator's; compiled pieces set their root. + Attention and other ops vLLM registers with a ``layer_name`` argument run their + layer without calling it (from compiled code too), so they mark that layer. +- Routing: the core capturer bound to every MoE layer's capture_fn hook. +- CPU KV offload (SimpleCPUOffload): each launch_copy is logged so the extractor + can pair its driver memcpys, which Kineto records without a CPU launch. +""" + +import functools +import inspect +import itertools +import os +import threading +import time + +from .. import envinfo, graphs, launchers, modules, routing, runtime, steps, windows +from .base import Engine + +_dispatch_log = threading.local() +_current_step = None + + +def rank_tag(runner=None): + """dp_tp for the calling worker; pid when ranks are not set up.""" + try: + from vllm.distributed import parallel_state as ps + + tp = ps.get_tensor_model_parallel_rank() + dp = runner.parallel_config.data_parallel_rank if runner is not None else ps.get_dp_group().rank_in_group + return f"dp{dp}_tp{tp}" + except Exception: + return f"pid{os.getpid()}" + + +def _forward_context(): + from vllm.forward_context import get_forward_context, is_forward_context_available + + if not is_forward_context_available(): + return None + ctx = get_forward_context() + return {"cudagraph_runtime_mode": str(getattr(ctx, "cudagraph_runtime_mode", None)), + "batch_descriptor": repr(getattr(ctx, "batch_descriptor", None))} + + +def _batch_desc_record(result): + try: + batch_desc, dp_sync = result + record = {"batch_desc": repr(batch_desc)} + for field in ("cg_mode", "num_tokens", "num_reqs", "uniform_token_count"): + if hasattr(batch_desc, field): + record[field] = str(getattr(batch_desc, field)) + if dp_sync is not None: + record["dp_sync"] = repr(dp_sync) + return record + except Exception: + return {"batch_desc": "unparsed"} + + +def patch_dp_utils(module): + orig = module.dispatch_cg_and_sync_dp + if getattr(orig, "_infx_patched", False): + return + + def dispatch_cg_and_sync_dp(*args, **kwargs): + result = orig(*args, **kwargs) + calls = getattr(_dispatch_log, "calls", None) + if calls is not None: + calls.append(_batch_desc_record(result)) + return result + + dispatch_cg_and_sync_dp._infx_patched = True + module.dispatch_cg_and_sync_dp = dispatch_cg_and_sync_dp + + +def patch_cudagraph_dispatcher(module): + """V1 dispatches before and again after the DP batch sync; the last call is the step's mode.""" + cls = getattr(module, "CudagraphDispatcher", None) + if cls is None or getattr(cls, "_infx_patched", False): + return + orig = cls.dispatch + + def dispatch(self, *args, **kwargs): + result = orig(self, *args, **kwargs) + try: + calls = getattr(_dispatch_log, "calls", None) + if calls is not None and getattr(_dispatch_log, "v1", False): + mode, batch_desc = result + record = {"batch_desc": repr(batch_desc), "cg_mode": str(getattr(mode, "name", mode))} + for field in ("num_tokens", "num_reqs", "uniform"): + if hasattr(batch_desc, field): + record[field] = str(getattr(batch_desc, field)) + calls.append(record) + except Exception: + pass + return result + + cls.dispatch = dispatch + cls._infx_patched = True + + +def step_record(scheduler_output): + reqs, computed, prompt_len = [], {}, {} + for new in scheduler_output.scheduled_new_reqs: + computed[new.req_id] = new.num_computed_tokens + ids = getattr(new, "prompt_token_ids", None) + prompt_len[new.req_id] = len(ids) if ids is not None else None + cached = scheduler_output.scheduled_cached_reqs + for req_id, n in zip(cached.req_ids, cached.num_computed_tokens): + computed[req_id] = n + spec = scheduler_output.scheduled_spec_decode_tokens or {} + for req_id, n in scheduler_output.num_scheduled_tokens.items(): + reqs.append({"req": req_id, "scheduled": n, "computed": computed.get(req_id), + "spec": len(spec.get(req_id, ())), "new": req_id in prompt_len, + "prompt_len": prompt_len.get(req_id)}) + return {"total_tokens": scheduler_output.total_num_scheduled_tokens, "reqs": reqs} + + +def register_models(runner): + modules.register_root("model", getattr(runner, "model", None)) + speculator = getattr(runner, "speculator", None) or getattr(runner, "drafter", None) # V2 / V1 + if speculator is not None: + modules.register_root("draft", getattr(speculator, "model", None)) + + +def bind_routing(runner, tag): + """Route every MoE layer's top-k ids of the target model to the core capturer.""" + model_config = runner.vllm_config.model_config + try: + from vllm.model_executor.layers.fused_moe.layer import MoERunner + except Exception: + MoERunner = () + bound, failed, sources = {}, {}, [] + for name, module in runner.model.named_modules(prefix="model"): + try: + fn = functools.partial(routing.capture, module.layer_id) if hasattr(module, "layer_id") else None + if MoERunner and isinstance(module, MoERunner): + quant_method = module._quant_method + if quant_method.is_monolithic: + impl = getattr(getattr(quant_method, "moe_kernel", None), "impl", None) + sources.append((getattr(impl, "fused_experts").set_capture_fn, fn)) + else: + sources.append((module.router.set_capture_fn, fn)) + elif hasattr(module, "capture_fn") and fn is not None: + # A capture source (e.g. DSV4's mega-MoE experts) calls it with its ids. + sources.append((functools.partial(setattr, module, "capture_fn"), fn)) + else: + continue + bound[module.layer_id] = name + except Exception as e: + failed[name] = repr(e) + try: + from vllm.distributed import parallel_state as ps + + write = ps.get_tensor_model_parallel_rank() == 0 # TP ranks route the same tokens + except Exception: + write = True + routing.setup(tag, runner.scheduler_config.max_num_batched_tokens, model_config.get_total_num_hidden_layers(), + model_config.get_num_experts_per_tok(), model_config.get_num_experts(), write=write, + bound=bound, failed=failed) + for set_fn, fn in sources: + set_fn(fn) + + +def patch_model_runner(module, v1=False): + import torch + + cls = module.GPUModelRunner + if vars(cls).get("_infx_patched", False): # its own flag: one runner may subclass the other + return + orig_capture, orig_execute = cls.capture_model, cls.execute_model + orig_sample = getattr(cls, "sample_tokens", None) + step_counter = itertools.count() + + def capture_model(self, *args, **kwargs): + tag = rank_tag(self) + envinfo.write(tag, ("vllm",), ("VLLM_",), + {"runner_class": f"{type(self).__module__}.{type(self).__qualname__}", + "vllm_config": str(getattr(self, "vllm_config", None))}) + for what, fn in (("register module names", register_models), ("bind routing", lambda r: bind_routing(r, tag))): + try: + fn(self) + except Exception: + runtime.write_error(what) + with windows.capture_session(tag): + return orig_capture(self, *args, **kwargs) + + def execute_model(self, scheduler_output, *args, **kwargs): + global _current_step + if kwargs.get("dummy_run") or (args and args[1:2] == (True,)): + return orig_execute(self, scheduler_output, *args, **kwargs) + k = next(step_counter) + _current_step = k + try: + record = step_record(scheduler_output) + except Exception: + runtime.write_error("step record") + record = {} + _dispatch_log.calls, _dispatch_log.v1 = [], v1 + try: + with steps.Step(rank_tag(self), k, record) as rec: + try: + return orig_execute(self, scheduler_output, *args, **kwargs) + finally: # inside the step, so its record carries the dispatch + calls = _dispatch_log.calls or [] + rec["dispatch"] = calls + if calls: + rec["graph"] = calls[-1].get("cg_mode") + finally: + _dispatch_log.calls = None + + def sample_tokens(self, *args, **kwargs): + if not runtime.profiling(): + return orig_sample(self, *args, **kwargs) + with torch.autograd.profiler.record_function("infx_sample"): + return orig_sample(self, *args, **kwargs) + + cls.capture_model = capture_model + cls.execute_model = execute_model + if orig_sample is not None: + cls.sample_tokens = sample_tokens + cls._infx_patched = True + + +def patch_gpu_worker(module): + """Idle DP ranks run a dummy forward each step to keep EP collectives in step.""" + cls = module.Worker + if getattr(cls, "_infx_patched", False) or not hasattr(cls, "execute_dummy_batch"): + return + orig = cls.execute_dummy_batch + counter = itertools.count() + + def execute_dummy_batch(self, *args, **kwargs): + # vLLM advances its profiler only for scheduled steps, so an idle DP rank's + # window never reaches max_iterations: its dummy forwards are its steps. + profiler = getattr(self, "profiler", None) + if profiler is not None and hasattr(profiler, "step"): + try: + profiler.step() + except Exception: + runtime.write_error("dummy profiler step") + with steps.Step(rank_tag(getattr(self, "model_runner", None)), next(counter), kind="dummy"): + return orig(self, *args, **kwargs) + + cls.execute_dummy_batch = execute_dummy_batch + cls._infx_patched = True + + +def patch_profiler_wrapper(module): + cls = module.WorkerProfiler + if getattr(cls, "_infx_patched", False): + return + orig_start, orig_stop = cls._call_start, cls._call_stop + + def _call_start(self, *args, **kwargs): + result = orig_start(self, *args, **kwargs) + windows.window_opened(rank_tag()) + return result + + def _call_stop(self, *args, **kwargs): + windows.window_closing(rank_tag()) # logged before the export, which runs inside the stop + return orig_stop(self, *args, **kwargs) + + cls._call_start = _call_start + cls._call_stop = _call_stop + cls._infx_patched = True + + +def patch_piecewise_backend(module): + import torch + + cls = module.PiecewiseBackend + if getattr(cls, "_infx_patched", False): + return + orig_call = cls.__call__ + + def __call__(self, *args): + if not runtime.profiling(): + return orig_call(self, *args) + tag = str(getattr(getattr(self, "vllm_backend", None), "prefix", "") or "") + modules.set_compiled_root("draft" if any(k in tag for k in ("eagle", "draft", "mtp", "spec", "dspark")) + else "model") + try: + with torch.autograd.profiler.record_function(f"infx_piece#{getattr(self, 'piecewise_compile_index', '?')}"): + return orig_call(self, *args) + finally: + modules.close_compiled_module() + + cls.__call__ = __call__ + cls._infx_patched = True + + +def _mark_layer_op(op_func): + """An op called with layer_name runs that layer without calling it: mark the launches as its.""" + try: + params = inspect.signature(op_func).parameters + except (TypeError, ValueError): + return op_func + if "layer_name" not in params or getattr(op_func, "_infx_patched", False): + return op_func + position = list(params).index("layer_name") + + @functools.wraps(op_func) + def op(*args, **kwargs): + if not runtime.profiling() or runtime.compiling(): + return op_func(*args, **kwargs) + name = kwargs.get("layer_name", args[position] if len(args) > position else None) + with modules.marked(modules.resolve(str(name)) if name else op_func.__name__, args): + return op_func(*args, **kwargs) + + op._infx_patched = True + return runtime.keep_signature(op, op_func) # registered as a custom op next + + +def patch_custom_op_registry(module): + orig = getattr(module, "direct_register_custom_op", None) + if orig is None or getattr(orig, "_infx_patched", False): + return + + @functools.wraps(orig) + def direct_register_custom_op(op_name, op_func, *args, **kwargs): + return orig(op_name, _mark_layer_op(op_func), *args, **kwargs) + + direct_register_custom_op._infx_patched = True + module.direct_register_custom_op = direct_register_custom_op + + +def patch_copy_backend(module): + """SimpleCPUOffload copies are driver memcpys from its copy thread: log each queued copy.""" + cls = getattr(module, "DmaCopyBackend", None) + if cls is None or getattr(cls, "_infx_patched", False) or not hasattr(cls, "launch_copy"): + return + orig = cls.launch_copy + block_bytes = {} # id(params) -> bytes of one block across all layers + + def launch_copy(self, src_blocks, dst_blocks, is_store, *args, **kwargs): + try: + params = self._store_params if is_store else self._load_params + if id(params) not in block_bytes: + block_bytes[id(params)] = int(sum(int(b) for b in params.bpb)) + runtime.sink("copies", rank_tag()).write({ + "t_ns": time.time_ns(), "store": bool(is_store), "blocks": len(src_blocks), + "bytes": block_bytes[id(params)] * len(src_blocks), "step": _current_step, + "callers": launchers.callers(2)}) + except Exception: + runtime.write_error("copy log") + return orig(self, src_blocks, dst_blocks, is_store, *args, **kwargs) + + cls.launch_copy = launch_copy + cls._infx_patched = True + + +class VllmEngine(Engine): + name = "vllm" + caller_roots = ("vllm",) + launcher_modules = ("vllm.utils.deep_gemm", "vllm.utils.flashinfer") + launcher_packages = ("vllm.third_party", "flashinfer") + version_packages = ("vllm",) + env_prefixes = ("VLLM_",) + window_control = "engine" + + def install(self, registry): + graphs.capture_context = _forward_context + for name, fn in { + "vllm.v1.worker.gpu.dp_utils": patch_dp_utils, + "vllm.v1.worker.gpu.model_runner": patch_model_runner, + "vllm.v1.worker.gpu_model_runner": functools.partial(patch_model_runner, v1=True), + "vllm.v1.cudagraph_dispatcher": patch_cudagraph_dispatcher, + "vllm.profiler.wrapper": patch_profiler_wrapper, + "vllm.compilation.piecewise_backend": patch_piecewise_backend, + "vllm.v1.worker.gpu_worker": patch_gpu_worker, + "vllm.v1.simple_kv_offload.copy_backend": patch_copy_backend, + "vllm.utils.torch_utils": patch_custom_op_registry, + "vllm.utils": patch_custom_op_registry, + }.items(): + registry.on(name, fn) + + +ENGINE = VllmEngine() diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/envinfo.py b/inferencex-e2e/benchmarks/profiling/infx_prof/envinfo.py new file mode 100644 index 0000000000..9b422d1131 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/envinfo.py @@ -0,0 +1,34 @@ +"""env/.json: versions, device identity (to join the clock samples) and engine config.""" + +import os +import time + +from . import runtime + + +def write(rank, packages, env_prefixes, extra=None): + try: + import torch + + info = {"pid": os.getpid(), "rank": rank, "t_ns": time.time_ns()} + for name in tuple(packages) + ("torch", "triton", "flashinfer", "deep_gemm"): + try: + module = __import__(name) + info[f"{name}_version"] = getattr(module, "__version__", None) + except Exception as e: + info[f"{name}_version"] = f"unavailable: {e}" + info["cuda_runtime"] = torch.version.cuda + info["hip_runtime"] = getattr(torch.version, "hip", None) + try: + info["device_name"] = torch.cuda.get_device_name() + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + info["device_uuid"] = str(props.uuid) + info["device_bdf"] = f"{props.pci_domain_id:04x}:{props.pci_bus_id:02x}:{props.pci_device_id:02x}.0" + except Exception as e: + info.setdefault("device_uuid", f"unavailable: {e}") + info["env"] = {k: v for k, v in os.environ.items() + if k.startswith(tuple(env_prefixes) + ("INFX_", "TORCH", "CUDA", "NCCL", "INDUCTOR"))} + info.update(extra or {}) + runtime.write_json(os.path.join("env", f"{rank}.json"), info) + except Exception: + runtime.write_error("write env") diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/graphs.py b/inferencex-e2e/benchmarks/profiling/infx_prof/graphs.py new file mode 100644 index 0000000000..accaf96ad6 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/graphs.py @@ -0,0 +1,128 @@ +"""CUDA graph ordinals and stream markers. + +Each CUDA graph a process captures gets an ordinal n: its capture runs inside +an ``infx_graph_capture#`` record_function and each replay inside +``infx_graph_replay#``, so the extractor can pair a replayed kernel with the +launch recorded when graph n was captured. ``with torch.cuda.stream(s)`` blocks +are marked ``infx_stream#``. +""" + +import itertools +import os +import threading +import time + +from . import runtime + +_graph_ids = itertools.count() +capture_context = None # engine adapters may set a callable returning a JSON-able dict +after_capture = [] # callables run after each graph's capture ends (the capture session rotates its profiler) +_replays = threading.local() + + +def replays(): + """CUDA graph replays launched from this thread so far (adapters diff it around a step).""" + return getattr(_replays, "n", 0) + + +def patch_cuda_graph(graphs_module): + import torch + + cls = graphs_module.CUDAGraph + if getattr(cls, "_infx_patched", False): + return + orig_begin, orig_end, orig_replay = cls.capture_begin, cls.capture_end, cls.replay + + def capture_begin(self, *args, **kwargs): + try: + gid = next(_graph_ids) + self._infx_id = gid + rf = torch.autograd.profiler.record_function(f"infx_graph_capture#{gid}") + rf.__enter__() + self._infx_rf = rf + context = None + if capture_context is not None: + try: + context = capture_context() + except Exception: + context = None + runtime.sink("graphs", f"pid{os.getpid()}").write({ + "graph": gid, "event": "capture_begin", "t_ns": time.time_ns(), + "device": torch.cuda.current_device(), "context": context, + }) + except Exception: + runtime.write_error("capture_begin") + return orig_begin(self, *args, **kwargs) + + def capture_end(self, *args, **kwargs): + try: + return orig_end(self, *args, **kwargs) + finally: + try: + rf = getattr(self, "_infx_rf", None) + if rf is not None: + rf.__exit__(None, None, None) + self._infx_rf = None + runtime.sink("graphs", f"pid{os.getpid()}").write({ + "graph": getattr(self, "_infx_id", None), "event": "capture_end", "t_ns": time.time_ns(), + }) + except Exception: + runtime.write_error("capture_end") + for fn in list(after_capture): + try: + fn() + except Exception: + runtime.write_error("after capture") + + def replay(self, *args, **kwargs): + _replays.n = getattr(_replays, "n", 0) + 1 + gid = getattr(self, "_infx_id", None) + if gid is None or not runtime.profiling(): + return orig_replay(self, *args, **kwargs) + with torch.autograd.profiler.record_function(f"infx_graph_replay#{gid}"): + return orig_replay(self, *args, **kwargs) + + cls.capture_begin = capture_begin + cls.capture_end = capture_end + cls.replay = replay + cls._infx_patched = True + + +def patch_stream_context(cuda_module): + """Mark `with torch.cuda.stream(s)` blocks: infx_stream#.""" + import torch + + cls = getattr(cuda_module, "StreamContext", None) + if cls is None or getattr(cls, "_infx_patched", False): + return + orig_enter, orig_exit = cls.__enter__, cls.__exit__ + + def __enter__(self): + out = orig_enter(self) + stream = getattr(self, "stream", None) + if stream is not None and runtime.profiling(): + try: + rf = torch.autograd.profiler.record_function(f"infx_stream#{int(stream.cuda_stream)}") + rf.__enter__() + self._infx_rf = rf + except Exception: + pass + return out + + def __exit__(self, *exc): + rf = getattr(self, "_infx_rf", None) + if rf is not None: + self._infx_rf = None + try: + rf.__exit__(None, None, None) + except Exception: + pass + return orig_exit(self, *exc) + + cls.__enter__, cls.__exit__ = __enter__, __exit__ + cls._infx_patched = True + + +def install(registry): + registry.on("torch.cuda.graphs", patch_cuda_graph) + registry.on("torch.cuda", patch_stream_context) diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/imports.py b/inferencex-e2e/benchmarks/profiling/infx_prof/imports.py new file mode 100644 index 0000000000..32e62bd035 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/imports.py @@ -0,0 +1,70 @@ +"""Patch modules as they are imported. + +Hooks are keyed by exact module name or by package prefix (every module under +it). Modules already imported when the hooks are installed are patched at once. +""" + +import importlib.abc +import sys + +from . import runtime + + +class HookRegistry: + def __init__(self): + self.exact = {} # module name -> [fn(module)] + self.packages = [] # (package, fn(module)) + + def on(self, name, fn): + self.exact.setdefault(name, []).append(fn) + + def on_package(self, package, fn): + self.packages.append((package, fn)) + + def hooks_for(self, name): + hooks = list(self.exact.get(name, ())) + hooks += [fn for package, fn in self.packages if name == package or name.startswith(package + ".")] + return hooks + + def run(self, name, module): + for hook in self.hooks_for(name): + try: + hook(module) + except Exception: + runtime.write_error(f"patch {name}") + + def install(self): + for name in list(sys.modules): + if self.hooks_for(name): + self.run(name, sys.modules[name]) + sys.meta_path.insert(0, _PostImportFinder(self)) + + +class _PostImportFinder(importlib.abc.MetaPathFinder): + def __init__(self, registry): + self.registry = registry + + def find_spec(self, name, path, target=None): + if not self.registry.hooks_for(name): + return None + for finder in sys.meta_path: + if finder is self: + continue + find = getattr(finder, "find_spec", None) + if find is None: + continue + spec = find(name, path, target) + if spec is not None and spec.loader is not None: + break + else: + return None + loader = spec.loader + exec_module = loader.exec_module + registry = self.registry + + def exec_and_patch(module, _exec=exec_module, _name=name): + _exec(module) + registry.run(_name, module) + + loader.exec_module = exec_and_patch + return spec diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/launchers.py b/inferencex-e2e/benchmarks/profiling/infx_prof/launchers.py new file mode 100644 index 0000000000..d55de13d14 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/launchers.py @@ -0,0 +1,176 @@ +"""Python launcher markers. + +Kernels launched straight from Python (Triton, TileLang, CuTe DSL, DeepGEMM, +FlashInfer and vendored kernel packages) have no torch op. Their launch entry +points mark themselves while profiling: ``infx_py##|``, +the frames being the engine's own callers above the launch, outermost first. + +Engines declare which modules and packages hold launchers (``Launchers``); the +DSL entry points (Triton's JITFunction.run, TileLang and CuTe DSL kernel calls) +are hooked for every engine. +""" + +import functools +import os +import sys +import types + +from . import runtime + +CALLER_FRAMES = 6 +_WRAP_SKIP_PREFIXES = ("_", "is_", "has_", "should_", "supports") +_OWN_DIR = os.path.dirname(os.path.abspath(__file__)) +_caller_roots = () # e.g. ("vllm",): source path segments of the engine's own frames + + +def callers(skip): + """The engine's frames above the caller, outermost first (path:line:function).""" + frame = sys._getframe(skip) + chain = [] + while frame is not None and len(chain) < CALLER_FRAMES: + path = frame.f_code.co_filename + if not os.path.abspath(path).startswith(_OWN_DIR): + for root in _caller_roots: + cut = path.rfind(f"{os.sep}{root}{os.sep}") + if cut >= 0: + chain.append(f"{path[cut + 1:]}:{frame.f_lineno}:{frame.f_code.co_name}") + break + frame = frame.f_back + return chain[::-1] + + +def marker(label, skip=3): + import torch + + return torch.autograd.profiler.record_function(f"infx_py#{label}#{'|'.join(callers(skip))}") + + +def wrap(fn, label): + @functools.wraps(fn) + def launcher(*args, **kwargs): + if runtime.compiling() or not runtime.profiling(): + return fn(*args, **kwargs) + with marker(label): + return fn(*args, **kwargs) + + launcher._infx_launcher = True + return runtime.keep_signature(launcher, fn) # launchers get registered as custom ops + + +# Functions of compiled extensions: CPython/pybind11 builtins, nanobind and Cython functions. +_COMPILED_FUNCTION_TYPES = {"builtin_function_or_method", "nb_func", "nb_method", "cython_function_or_method"} + + +def compiled_function(obj): + return type(obj).__name__ in _COMPILED_FUNCTION_TYPES + + +def patch_module(module, package=None): + """Mark every public Python function `module` defines, and its classes' public methods. + + With `package` (a vendored kernel package the module belongs to), also its + single-underscore functions (vendored packages keep real launchers private) + and code of the package reached outside the module namespace: functions + whose source lies under the package, builtins of its compiled extensions. + """ + private = package is not None + root = os.path.dirname(getattr(sys.modules.get(package), "__file__", None) or "") if private else "" + + def launchable(name, obj): + skipped = name.startswith("__") or ( + name.startswith(_WRAP_SKIP_PREFIXES) and not (private and name.startswith("_"))) + return (isinstance(obj, types.FunctionType) and obj.__module__ == module.__name__ + and not skipped and not getattr(obj, "_infx_launcher", False)) + + def package_code(name, obj): + if not root or name.startswith("__") or getattr(obj, "_infx_launcher", False): + return False + if isinstance(obj, types.FunctionType): + return obj.__module__ != module.__name__ and obj.__code__.co_filename.startswith(root + os.sep) + owner_file = getattr(getattr(obj, "__self__", None), "__file__", None) or "" + return compiled_function(obj) and ( + (getattr(obj, "__module__", None) or "").startswith(package + ".") + or owner_file.startswith(root + os.sep)) + + for name, obj in list(vars(module).items()): + if launchable(name, obj) or package_code(name, obj): + setattr(module, name, wrap(obj, f"{module.__name__}.{name}")) + elif (private and callable(obj) and not isinstance(obj, (type, types.ModuleType)) + and not name.startswith("__") and not getattr(obj, "_infx_launcher", False)): + kind = type(obj) + runtime.note_once("launchers", "unwrapped", + f"{module.__name__}.{name}: {kind.__module__}.{kind.__qualname__}" + f" from {getattr(obj, '__module__', None)}") + elif isinstance(obj, type) and obj.__module__ == module.__name__ and not name.startswith("_"): + for attr, method in list(vars(obj).items()): + label = f"{module.__name__}.{name}.{attr}" + if launchable(attr, method): + setattr(obj, attr, wrap(method, label)) + elif isinstance(method, (staticmethod, classmethod)) and launchable(attr, method.__func__): + setattr(obj, attr, type(method)(wrap(method.__func__, label))) # e.g. a plan's generate() + + +def wrap_builtins(module, label_prefix, names=None): + """Mark compiled-extension functions (nanobind / pybind builtins) bound as `module` attributes. + + For engines that call kernel bindings directly rather than through torch ops, + looking them up on `module` at call time. `names` limits it to those attributes. + """ + for name, obj in list(vars(module).items()): + if names is not None and name not in names: + continue + if (compiled_function(obj) and not name.startswith("_") + and not getattr(obj, "_infx_launcher", False)): + setattr(module, name, wrap(obj, f"{label_prefix}.{name}")) + + +def patch_calls(module, attr, label_of): + """Mark calls of every class in `module` that defines `attr` itself.""" + for cls in list(vars(module).values()): + if not isinstance(cls, type) or cls.__module__ != module.__name__ or attr not in vars(cls): + continue + orig = vars(cls)[attr] + if getattr(orig, "_infx_launcher", False): + continue + + def call(self, *args, _orig=orig, **kwargs): + if runtime.compiling() or not runtime.profiling(): + return _orig(self, *args, **kwargs) + try: + label = label_of(self) + except Exception: + label = type(self).__qualname__ + with marker(label): + return _orig(self, *args, **kwargs) + + call._infx_launcher = True + setattr(cls, attr, call) + + +def triton_label(jit_function): + fn = getattr(jit_function, "fn", None) + return f"triton:{getattr(fn, '__qualname__', None) or type(jit_function).__qualname__}" + + +def named_label(prefix): + def label(obj): + for attr in ("kernel_name", "name", "__name__", "func_name"): + value = getattr(obj, attr, None) + if isinstance(value, str) and value: + return f"{prefix}:{value}" + return f"{prefix}:{type(obj).__qualname__}" + + return label + + +def install(registry, caller_roots, modules=(), packages=()): + """Hook the DSL launchers, the engine's launcher `modules` and every module under `packages`.""" + global _caller_roots + _caller_roots = tuple(caller_roots) + registry.on("triton.runtime.jit", lambda m: patch_calls(m, "run", triton_label)) + registry.on("tilelang.jit.kernel", lambda m: patch_calls(m, "__call__", named_label("tilelang"))) + registry.on("cutlass.cutlass_dsl.tvm_ffi_provider", lambda m: patch_calls(m, "__call__", named_label("cute"))) + for name in modules: + registry.on(name, patch_module) + for package in packages: + registry.on_package(package, functools.partial(patch_module, package=package)) diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/modules.py b/inferencex-e2e/benchmarks/profiling/infx_prof/modules.py new file mode 100644 index 0000000000..daefdf8a6f --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/modules.py @@ -0,0 +1,394 @@ +"""Module markers: every launch carries the module that issued it. + +Profiled windows run without Python stacks (their export would stall the +engine), so modules mark themselves: ``infx_mod##``. + +- Eagerly running modules: global forward hooks, enabled while a profiler is + recording (windows, graph capture). They return at once while Dynamo traces, + so compiled graphs are unchanged. +- Inductor-compiled code: one host line per scheduler node is written into the + wrapper code, naming the deepest module common to the node's sources + (nn_module_stack, else the weights it reads); at run time it opens a marker. + Host-only: kernels and fusion are unchanged. + +Engine adapters register model roots with ``register_root(prefix, module)`` +("model" for the target, "draft" for a speculator) and may set the current +compiled root with ``set_compiled_root``. +""" + +import re +import threading +import weakref + +from . import runtime + +_names = weakref.WeakKeyDictionary() # module -> qualified name +_hooks = [] +_tls = threading.local() +_compiled_names = {} # (root, relative path) -> registered qualified name + + +def register_root(prefix, root): + if root is None or not hasattr(root, "named_modules"): + return + for name, mod in root.named_modules(prefix=prefix): + # torch.compile's OptimizedModule nests the module as _orig_mod; name the eager tree. + _names.setdefault(mod, name.replace("._orig_mod", "")) + _mark_direct_calls(root, methods=True) + for mod in root.children(): + _mark_direct_calls(mod) + _compiled_names.clear() + _resolved.clear() + + +_resolved = {} + + +def _marked_call(self, call, args): + """Run call() as module `self`'s, unless `self` is already the innermost marked module.""" + import torch + + if torch.compiler.is_compiling() or not _hooks: + return call() + stack = getattr(_tls, "stack", None) + if stack is None: + stack = _tls.stack = [] + if stack and stack[-1][0] is self: # called through __call__ (the pre-hook marked it) or nested + return call() + rf = torch.autograd.profiler.record_function( + f"infx_mod#{_names.get(self) or type(self).__name__}#{tensor_signature(args)}") + rf.__enter__() + entry = (self, rf) + stack.append(entry) + try: + return call() + finally: + if entry in stack: # the post-hook pops through entries it passes + while stack: + top = stack.pop() + top[1].__exit__(None, None, None) + if top is entry: + break + + +def _marking(fn): + """`fn` (a model class's own forward or method) marked as its module when called directly. + + The original function stays unbound and gets `self` explicitly, so Dynamo still + names the module `self` (compiled code's module paths are unchanged), and every + instance, copies included, marks itself. + """ + import functools + + @functools.wraps(fn) + def method(self, *args, **kwargs): + return _marked_call(self, lambda: fn(self, *args, **kwargs), args) + + method._infx_marking = True + return method + + +def _own_methods(cls): + """Public functions a model class defines itself (not nn.Module's), e.g. compute_logits.""" + import types + + import torch + + base = set(dir(torch.nn.Module)) + return {name for name, obj in vars(cls).items() + if isinstance(obj, types.FunctionType) and not name.startswith("_") and name not in base} + + +def _mark_direct_calls(module, methods=False): + """Engines call a model's forward directly (SGLang: model.forward, then model.model.forward), + and its other methods (vLLM: compute_logits; a speculative head's logits step), which skips + the module hooks; mark such calls. The module's classes are patched, not the instance: + engines copy models (SGLang's breakable prefill graphs), and an instance attribute would + carry the original module into the copy.""" + import types + + import torch + + if isinstance(module, (torch.nn.ModuleList, torch.nn.ModuleDict)) or "forward" in vars(module): + return + for cls in type(module).__mro__: + if cls is torch.nn.Module or not issubclass(cls, torch.nn.Module) or cls.__module__.startswith("torch."): + continue + names = (_own_methods(cls) if methods else set()) | {"forward"} + for name in names: + fn = vars(cls).get(name) + if not isinstance(fn, types.FunctionType) or getattr(fn, "_infx_marking", False): + continue + try: + setattr(cls, name, _marking(fn)) + except Exception: + pass + + +def resolve(name, root="model"): + """The registered qualified name ending with `name` (an engine's own layer prefix), shortest first.""" + key = (root, name) + if key not in _resolved: + suffix = f".{name}" + matches = sorted((n for n in list(_names.values()) + if (n.startswith(root + ".") or not root) and (n.endswith(suffix) or n == name)), key=len) + _resolved[key] = matches[0] if matches else name + return _resolved[key] + + +class marked: + """Mark the enclosed launches as module `name`'s (for code that runs a module without calling it).""" + + def __init__(self, name, args=()): + self.name, self.args, self.rf = name, args, None + + def __enter__(self): + if runtime.profiling() and not runtime.compiling(): + try: + import torch + + self.rf = torch.autograd.profiler.record_function( + f"infx_mod#{self.name}#{tensor_signature(self.args)}") + self.rf.__enter__() + except Exception: + self.rf = None + return self + + def __exit__(self, *exc): + if self.rf is not None: + try: + self.rf.__exit__(None, None, None) + except Exception: + pass + return False + + +def tensor_signature(args): + """'7x7168:bfloat16;1x2:int64' for the tensors among a call's (one-level nested) args. + + Marker names must stay free of quotes and backslashes: Kineto writes event + names into the trace JSON without escaping them. + """ + import torch + + sig = [] + for arg in args: + items = arg if isinstance(arg, (list, tuple)) else (arg,) + for item in items: + if isinstance(item, torch.Tensor): + dims = "x".join(str(d) for d in item.shape) + sig.append(f"{dims}:{str(item.dtype).removeprefix('torch.')}") + return ";".join(sig) + + +def _pre_hook(module, args): + if runtime.compiling(): + return + try: + import torch + + name = _names.get(module) or type(module).__name__ + # Python-launched kernels have no op shapes of their own; the module's inputs stand in. + rf = torch.autograd.profiler.record_function(f"infx_mod#{name}#{tensor_signature(args)}") + rf.__enter__() + stack = getattr(_tls, "stack", None) + if stack is None: + stack = _tls.stack = [] + stack.append((module, rf)) + except Exception: + pass + + +def _post_hook(module, args, output): + if runtime.compiling(): + return + stack = getattr(_tls, "stack", None) + while stack: + mod, rf = stack.pop() + try: + rf.__exit__(None, None, None) + except Exception: + pass + if mod is module: + break + + +def enable(): + if _hooks: + return + from torch.nn.modules import module as nn_module + + _hooks.append(nn_module.register_module_forward_pre_hook(_pre_hook)) + _hooks.append(nn_module.register_module_forward_hook(_post_hook, always_call=True)) + + +def disable(): + while _hooks: + _hooks.pop().remove() + stack = getattr(_tls, "stack", None) + while stack: + try: + stack.pop()[1].__exit__(None, None, None) + except Exception: + pass + + +# --- Inductor-compiled code ------------------------------------------------------ + +_PLACEHOLDER_MODULE = re.compile(r"^l_self_(?:modules_)?(.*?)_(?:parameters|buffers)_") +_compile_inputs = threading.local() # stack of Dynamo placeholder names, one list per compile_fx call + + +def _relative_module_path(fqn): + """"L['self'].layers[3].self_attn" -> "layers.3.self_attn".""" + fqn = re.sub(r"^L\['self'\]\.?", "", fqn) + fqn = re.sub(r"_modules\['([^']+)'\]", r"\1", fqn) + fqn = re.sub(r"\[(\d+)\]", r".\1", fqn) + fqn = re.sub(r"\['([^']+)'\]", r".\1", fqn) + return fqn.strip(".").replace("..", ".") + + +def _stack_modules(node): + paths = [] + for snode in node.get_nodes(): + if snode.node is None: + continue + for origin in snode.node.get_origins(): + stack = origin.meta.get("nn_module_stack") + if stack: + paths.append(_relative_module_path(list(stack.values())[-1][0]).split(".")) + return paths + + +def _weight_modules(node, graph): + """Module paths of the weights (graph inputs) the node reads. + + Inductor's own placeholders (arg0_1, ...) keep the order of the graph + handed to compile_fx, whose Dynamo names spell the owning module + (l_self_modules_layers_modules_3_modules_mlp_..._parameters_weight_). + """ + stack = getattr(_compile_inputs, "stack", None) + if not stack: + return [] + placeholders = [n.name for n in graph.graph.find_nodes(op="placeholder")] + if len(placeholders) != len(stack[-1]): + return [] + dynamo_name = dict(zip(placeholders, stack[-1])) + paths = [] + for dep in node.read_writes.reads: + match = _PLACEHOLDER_MODULE.match(dynamo_name.get(dep.name, "")) + if match: + paths.append(match.group(1).split("_modules_")) + return paths + + +def _common_module(paths): + if not paths: + return None + common = paths[0] + for parts in paths[1:]: + k = 0 + while k < min(len(common), len(parts)) and common[k] == parts[k]: + k += 1 + common = common[:k] + return ".".join(common) + + +def patch_compile_fx(module): + orig = getattr(module, "compile_fx", None) + if orig is None or getattr(orig, "_infx_patched", False): + return + + def compile_fx(model_, example_inputs_, *args, **kwargs): + graph = getattr(model_, "graph", None) + names = [n.name for n in graph.find_nodes(op="placeholder")] if graph is not None else [] + stack = getattr(_compile_inputs, "stack", None) + if stack is None: + stack = _compile_inputs.stack = [] + stack.append(names) + try: + return orig(model_, example_inputs_, *args, **kwargs) + finally: + stack.pop() + + compile_fx._infx_patched = True + module.compile_fx = compile_fx + + +def patch_inductor_scheduler(module): + cls = module.Scheduler + if getattr(cls, "_infx_patched", False) or not hasattr(cls, "enter_context"): + return + orig = cls.enter_context + + def enter_context(self, node): + orig(self, node) + try: + # Imported here, not at patch time: the scheduler can be first imported while + # torch._inductor.debug is still initializing, and the wrapper module imports from it. + from torch._inductor.codegen.wrapper import PythonWrapperCodegen + from torch._inductor.virtualized import V + + wrapper = V.graph.wrapper_code + # Subgraph wrappers are Python too; C++ wrappers (a subclass) are not. + if getattr(V.graph, "cpp_wrapper", False) or not isinstance(wrapper, PythonWrapperCodegen): + runtime.note_once("compile", "module_line_skipped", type(wrapper).__name__) + return + stacked = _stack_modules(node) + source = "nn_module_stack" if stacked else "weights" + path = _common_module(stacked or _weight_modules(node, V.graph)) + if path is None: + path, source = "", "none" # no module: close the previous one rather than extend it + if getattr(wrapper, "_infx_module", None) != path: + wrapper.writeline(f"__import__('infx_prof.modules').modules.compiled_module({path!r})") + wrapper._infx_module = path + runtime.note_once("compile", "module_source", f"{type(wrapper).__name__}:{source}") + except Exception: + runtime.write_error("inductor module line") + + cls.enter_context = enter_context + cls._infx_patched = True + + +def set_compiled_root(root): + """'model' or 'draft': whose modules the compiled code about to run belongs to.""" + _tls.compiled_root = root + + +def compiled_module(path): + """Called from Inductor wrapper code: mark the following kernels as `path`'s.""" + if not runtime.profiling(): + return + import torch + + current = getattr(_tls, "compiled", None) + if current is not None and current[0] == path: + return + close_compiled_module() + if not path: + return + root = getattr(_tls, "compiled_root", "model") + name = _compiled_names.get((root, path)) + if name is None: + suffix = f".{path}" + matches = sorted((n for n in list(_names.values()) if n.startswith(root + ".") and n.endswith(suffix)), + key=len) + name = _compiled_names[(root, path)] = matches[0] if matches else f"{root}:{path}" + rf = torch.autograd.profiler.record_function(f"infx_mod#{name}#") + rf.__enter__() + _tls.compiled = (path, rf) + + +def close_compiled_module(): + current = getattr(_tls, "compiled", None) + if current is not None: + _tls.compiled = None + try: + current[1].__exit__(None, None, None) + except Exception: + pass + + +def install(registry): + registry.on("torch._inductor.scheduler", patch_inductor_scheduler) + registry.on("torch._inductor.compile_fx", patch_compile_fx) diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/routing.py b/inferencex-e2e/benchmarks/profiling/infx_prof/routing.py new file mode 100644 index 0000000000..46bdf7e227 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/routing.py @@ -0,0 +1,152 @@ +"""MoE routing: each step's per-layer top-k expert ids. + +An engine adapter calls ``setup`` once per rank with the model's shape, then +routes every MoE layer's top-k ids to ``capture(layer_id, topk_ids)``. The ids +are copied into a device buffer, a GPU copy that CUDA graphs capture, so every +replay records its routing. Inside torch.compile the copy is the custom op +``infx::capture_routing``, opaque to Dynamo, so compiled graphs keep it. While a window is open, ``snapshot`` queues each +step's rows to pinned host memory asynchronously; ``flush`` (at the window's +stop) writes routing//step.npz: per-layer expert token counts, and +per-token ids when the step has at most IDS_MAX_TOKENS tokens. +""" + +import os + +from . import runtime + +IDS_MAX_TOKENS = 1024 + + +class _State: + buffer = None + rank = None + write = False + num_experts = None + captured = False + pending = [] + + +_s = _State() + + +def _define_op(): + import torch + + if hasattr(torch.ops.infx, "capture_routing"): + return + + @torch.library.custom_op("infx::capture_routing", mutates_args=("buf",)) + def capture_routing(buf: torch.Tensor, ids: torch.Tensor, layer: int) -> None: + _copy(buf, ids, layer) + + @capture_routing.register_fake + def _(buf, ids, layer): + return None + + +def _copy(buf, ids, layer): + n = min(ids.shape[0], buf.shape[0]) + buf[:n, layer, :ids.shape[1]] = ids[:n] + _s.captured = True + + +def setup(rank, max_tokens, num_layers, topk, num_experts, write=True, bound=None, failed=None, extra=None): + """Allocate the buffer and write routing//meta.json. `write`: this rank records (one per TP group). + + Once per process: CUDA graphs captured after it write this buffer. + """ + import torch + + if _s.buffer is not None: + return + try: + _define_op() + except Exception: + runtime.write_error("routing custom op") + + device = "cuda" if torch.cuda.is_available() else "cpu" + _s.buffer = torch.full((max_tokens, num_layers, topk), -1, dtype=torch.int32, device=device) + _s.rank, _s.write, _s.num_experts = rank, write, num_experts + meta = {"rank": rank, "write": write, "layers": num_layers, "topk": topk, "num_experts": num_experts, + "max_tokens": max_tokens, "bound": {str(k): v for k, v in sorted((bound or {}).items())}, + "failed": failed or {}} + meta.update(extra or {}) + runtime.write_json(os.path.join("routing", rank, "meta.json"), meta) + + +def update_meta(**fields): + import json + + path = os.path.join(runtime.DIR, "routing", _s.rank, "meta.json") + try: + with open(path) as f: + meta = json.load(f) + except Exception: + meta = {} + meta.update(fields) + runtime.write_json(os.path.join("routing", _s.rank, "meta.json"), meta) + + +def active(): + return _s.buffer is not None + + +def capture(layer_id, topk_ids): + """Record layer `layer_id`'s top-k ids [tokens, k]; traceable by Dynamo (fullgraph).""" + import torch + + buf = _s.buffer + if buf is None or topk_ids is None or not isinstance(topk_ids, torch.Tensor) or topk_ids.dim() != 2: + return + if not 0 <= layer_id < buf.shape[1]: + return + if topk_ids.shape[1] > buf.shape[2]: + # Fused shared experts ride along as trailing columns (SGLang shared-expert + # fusion appends expert id num_experts); the routed top-k come first. + topk_ids = topk_ids[:, :buf.shape[2]] + if torch.compiler.is_compiling(): + torch.ops.infx.capture_routing(buf, topk_ids, layer_id) + else: + _copy(buf, topk_ids, layer_id) + + +def snapshot(step, num_tokens, t0_ns): + """Queue this step's routing rows for the host (asynchronous).""" + if _s.buffer is None or not _s.write or not num_tokens: + return + import torch + + with torch.autograd.profiler.record_function("infx_routing_copy"): + rows = _s.buffer[:num_tokens].to(torch.int16) + host = torch.empty(rows.shape, dtype=torch.int16, pin_memory=rows.is_cuda) + host.copy_(rows, non_blocking=True) + _s.pending.append((step, num_tokens, t0_ns, host, rows)) + + +def flush(): + """Write the window's queued routing: per-layer expert counts, and ids for small steps.""" + pending, _s.pending = _s.pending, [] + if not pending: + return + if not _s.captured: # no layer handed over ids, during capture or eagerly + update_meta(unavailable="no MoE layer exposed its top-k ids (fused routing inside the MoE kernel)") + return + import numpy as np + import torch + + if torch.cuda.is_available(): + torch.cuda.current_stream().synchronize() # the queued host copies + out_dir = os.path.join(runtime.DIR, "routing", _s.rank) + os.makedirs(out_dir, exist_ok=True) + for step, num_tokens, t0_ns, host, _ in pending: + ids = host.numpy().astype(np.int32) # [tokens, layers, topk]; -1 marks no expert + # EPLB hands over physical expert ids, which can exceed the logical count. + experts = max(_s.num_experts or 0, int(ids.max()) + 1) + counts = np.zeros((ids.shape[1], experts), dtype=np.int32) + for layer in range(ids.shape[1]): + valid = ids[:, layer][ids[:, layer] >= 0] + counts[layer] = np.bincount(valid, minlength=experts)[:experts] + arrays = {"counts": counts, "t0_ns": np.array(t0_ns), "tokens": np.array(num_tokens)} + if num_tokens <= IDS_MAX_TOKENS: + arrays["ids"] = ids.astype(np.int16) + np.savez_compressed(os.path.join(out_dir, f"step{step:06d}.npz"), **arrays) diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/runtime.py b/inferencex-e2e/benchmarks/profiling/infx_prof/runtime.py new file mode 100644 index 0000000000..ab3704953c --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/runtime.py @@ -0,0 +1,126 @@ +"""Process-wide state shared by the core and the engine adapters. + +Everything goes under $INFX_PROF_DIR; nothing here may raise into the engine. +""" + +import json +import os +import threading +import time +import traceback + +DIR = os.environ.get("INFX_PROF_DIR") + + +def write_error(where): + """Append the current exception to errors/.log.""" + try: + path = os.path.join(DIR, "errors", f"{os.getpid()}.log") + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "a") as f: + f.write(f"--- {time.time():.3f} {where}\n{traceback.format_exc()}\n") + except Exception: + pass + + +class JsonlSink: + """Line-buffered: engines are killed at teardown, so nothing may sit in a buffer.""" + + def __init__(self, path): + os.makedirs(os.path.dirname(path), exist_ok=True) + self._f = open(path, "a", buffering=1) + self._lock = threading.Lock() + + def write(self, record): + line = json.dumps(record, separators=(",", ":"), default=str) + with self._lock: + self._f.write(line + "\n") + + def flush(self): + with self._lock: + self._f.flush() + + +_sinks = {} + + +def sink(kind, tag): + """The JSONL file /.jsonl.""" + key = (kind, tag) + if key not in _sinks: + _sinks[key] = JsonlSink(os.path.join(DIR, kind, f"{tag}.jsonl")) + return _sinks[key] + + +def flush_all(): + for s in list(_sinks.values()): + try: + s.flush() + except Exception: + pass + + +_notes = set() + + +def note_once(kind, event, detail): + """One record per (event, detail) in /pid.jsonl, to see which paths ran.""" + if (kind, event, detail) in _notes: + return + _notes.add((kind, event, detail)) + try: + sink(kind, f"pid{os.getpid()}").write({"event": event, "detail": detail, "t_ns": time.time_ns()}) + except Exception: + pass + + +def profiling(): + """Whether a torch profiler is recording in this process.""" + try: + import torch + + return torch._C._autograd._profiler_enabled() + except Exception: + return True + + +def compiling(): + """Whether Dynamo is tracing (hooks must then be no-ops so compiled graphs are unchanged).""" + try: + import torch + + return torch.compiler.is_compiling() + except Exception: + return False + + +def capture_ranks(): + return os.environ.get("INFX_PROF_CAPTURE_RANKS", "dp0_tp0").split(",") + + +def is_capture_rank(tag): + ranks = capture_ranks() + return tag in ranks or "all" in ranks + + +def keep_signature(wrapper, fn): + """Give `wrapper` the signature of `fn` with string annotations resolved in fn's globals. + + Engines register wrapped callables as torch custom ops; torch.library.infer_schema + resolves `from __future__ import annotations` strings in the callable's own globals, + which for a wrapper are not fn's. + """ + import inspect + + try: + wrapper.__signature__ = inspect.signature(fn, eval_str=True) + except Exception: + pass + return wrapper + + +def write_json(relpath, data): + path = os.path.join(DIR, relpath) + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w") as f: + json.dump(data, f, indent=1, default=str) diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/steps.py b/inferencex-e2e/benchmarks/profiling/infx_prof/steps.py new file mode 100644 index 0000000000..c0b98762ee --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/steps.py @@ -0,0 +1,65 @@ +"""Per-step record and marker. + +``Step`` wraps one engine step on one rank: + + with steps.Step(rank, k, record) as rec: + rec["graph"] = ... # FULL / PIECEWISE / BREAKABLE / NONE, when known + result = run_step() + +At the boundary it ticks the file window controller (if the engine uses one), +runs the step inside ``infx_#k`` while profiling, queues the step's +routing, and appends the record to steps/.jsonl with t0_ns/t1_ns. + +Record fields (engine-neutral; adapters fill what they know): + step | dummy step index (dummy: idle-rank forward) + phase prefill / decode / mixed / idle + graph FULL / PIECEWISE / BREAKABLE / NONE + total_tokens tokens this rank's forward processes + reqs [{req, scheduled, computed, prompt_len, new, spec}] + route_rows routing rows the MoE layers saw, when not total_tokens + (MoE over tokens gathered from every DP rank) +""" + +import time + +from . import routing, runtime + +controller = None # windows.FileWindows for engines without an on-demand profiler + + +class Step: + def __init__(self, rank, k, record=None, kind="step"): + self.rank, self.k, self.kind = rank, k, kind + self.record = record if record is not None else {} + self.rf = None + + def __enter__(self): + if controller is not None: + controller.on_step_begin(self.rank, self.k) + self.t0 = time.time_ns() + if runtime.profiling(): + try: + import torch + + self.rf = torch.autograd.profiler.record_function(f"infx_{self.kind}#{self.k}") + self.rf.__enter__() + except Exception: + self.rf = None + return self.record + + def __exit__(self, *exc): + if self.rf is not None: + try: + self.rf.__exit__(None, None, None) + if self.kind == "step": + rows = self.record.get("route_rows") or self.record.get("total_tokens") + routing.snapshot(self.k, rows, self.t0) + except Exception: + runtime.write_error("step close") + try: + self.record.update({"dummy" if self.kind == "dummy" else "step": self.k, + "t0_ns": self.t0, "t1_ns": time.time_ns()}) + runtime.sink("steps", self.rank).write(self.record) + except Exception: + runtime.write_error("step write") + return False diff --git a/inferencex-e2e/benchmarks/profiling/infx_prof/windows.py b/inferencex-e2e/benchmarks/profiling/infx_prof/windows.py new file mode 100644 index 0000000000..68d8fc6deb --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/infx_prof/windows.py @@ -0,0 +1,221 @@ +"""Profiler windows and the graph-capture profile. + +Engines whose own profiler can be opened on demand (vLLM's /start_profile) +only log their profiler's start/stop with ``log_profiler_event``. For the +others, ``FileWindows`` runs torch profiler windows itself, at step +boundaries: the window client appends requests to control/windows.jsonl, + + {"window": 0, "steps": 32, "start_at": {"dp0_tp0": 1210, ...}} + {"stop": 0} + +and each rank starts the window at the first step k >= its start_at (the +client sets it a couple of steps past the rank's latest logged step, so TP +ranks in lockstep start together), records `steps` steps, then stops and +exports torch/..pt.trace.json.gz. A stop request closes an open +window early. Both paths log start/stop to profiler/.jsonl, which the +client waits on. +""" + +import json +import os +import time + +from . import graphs, modules, routing, runtime + +# Graphs per capture trace: one profiler session over many large captures (SGLang's +# breakable prefill graphs) overflows its activity buffer and silently drops launches. +CAPTURE_GRAPHS_PER_FILE = int(os.environ.get("INFX_PROF_CAPTURE_GRAPHS_PER_FILE") or 0) # 0: no rotation + + +def log_profiler_event(rank, event): + try: + runtime.sink("profiler", rank).write({"event": event, "t_ns": time.time_ns()}) + except Exception: + runtime.write_error(f"profiler {event} log") + + +def window_opened(rank): + """Module markers on for an open window.""" + log_profiler_event(rank, "start") + try: + modules.enable() + except Exception: + runtime.write_error("enable module markers") + + +def window_closing(rank): + """Before the profiler stops: log the stop, write the window's routing, markers off.""" + log_profiler_event(rank, "stop") + try: + routing.flush() + except Exception: + runtime.write_error("routing flush") + try: + modules.disable() + except Exception: + runtime.write_error("disable module markers") + + +def _quiesce(): + """Wait for the device's queued work before the profiler starts or stops. + + On ROCm an SGLang rank hung in hipGraphLaunch right after a window's profiler + started with the previous (overlapped) step's work still in flight, which + deadlocked every rank's collectives. Tracing now starts and ends idle. + """ + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.synchronize() + except Exception: + runtime.write_error("synchronize") + + +def _new_profiler(): + from torch.profiler import ProfilerActivity, profile + + return profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True, + # Markers carry the context; Python stacks make exports outlast the engine's patience. + with_stack=False) + + +class FileWindows: + """Torch profiler windows driven by control/windows.jsonl, opened and closed at step boundaries.""" + + def __init__(self): + self.path = os.path.join(runtime.DIR, "control", "windows.jsonl") + self.offset = 0 + self.pending = [] # requests not yet started + self.stops = set() + self.active = None # (window, profiler, start step, steps) + self.mtime = None + + def _poll(self): + try: + st = os.stat(self.path) + # Appends within one coarse mtime tick change only the size. + mtime = (st.st_mtime_ns, st.st_size) + except OSError: + return + if mtime == self.mtime: + return + self.mtime = mtime + with open(self.path) as f: + f.seek(self.offset) + while line := f.readline(): + if not line.endswith("\n"): + break + self.offset = f.tell() + request = json.loads(line) + if "stop" in request: + self.stops.add(request["stop"]) + else: + self.pending.append(request) + + def on_step_begin(self, rank, step): + try: + self._poll() + if self.active is not None: + window, _, start, steps = self.active + if step >= start + steps or window in self.stops: + self._stop(rank) + if self.active is None and self.pending: + request = self.pending[0] + if step >= request.get("start_at", {}).get(rank, 0): + self.pending.pop(0) + if request["window"] not in self.stops: + self._start(rank, request, step) + except Exception: + runtime.write_error("window control") + + def _start(self, rank, request, step): + prof = _new_profiler() + _quiesce() + prof.start() + self.active = (request["window"], prof, step, int(request["steps"])) + window_opened(rank) + + def _stop(self, rank): + window, prof, _, _ = self.active + self.active = None + window_closing(rank) + _quiesce() + prof.stop() + path = os.path.join(runtime.DIR, "torch", f"{rank}.{time.time_ns()}.pt.trace.json.gz") + os.makedirs(os.path.dirname(path), exist_ok=True) + prof.export_chrome_trace(path) + + def close(self, rank): + if self.active is not None: + try: + self._stop(rank) + except Exception: + runtime.write_error("window close") + + +# Profilers a capture session rotated out stay referenced for the process's life: +# freeing a stopped ROCm profiler as SGLang collected garbage at the end of its +# capture (freeze_gc) crashed the process (SIGSEGV). Capture runs once, at start-up. +_retired = [] + + +class capture_session: + """Profile CUDA graph capture on capture ranks: capture/[.].pt.trace.json.gz. + + Every launch captured into a graph is recorded with its op, input shapes and + markers; the extractor pairs replayed kernels with these launches. + """ + + def __init__(self, rank): + self.rank = rank + self.prof = None + self.graphs = 0 + + def __enter__(self): + # A re-capture during serving would nest in an open window's profiler. + if not runtime.is_capture_rank(self.rank) or runtime.profiling(): + return self + try: + self.prof = _new_profiler() + self.prof.__enter__() + modules.enable() + graphs.after_capture.append(self._graph_captured) + except Exception: + runtime.write_error("capture profiler start") + self.prof = None + return self + + def _graph_captured(self): + """Every CAPTURE_GRAPHS_PER_FILE graphs, export this trace and continue in a new one.""" + self.graphs += 1 + if self.prof is None or CAPTURE_GRAPHS_PER_FILE <= 0 or self.graphs % CAPTURE_GRAPHS_PER_FILE: + return + self._export() + _retired.append(self.prof) + self.prof = _new_profiler() + self.prof.__enter__() + + def _export(self): + try: + self.prof.__exit__(None, None, None) + base = os.path.join(runtime.DIR, "capture", self.rank) + path, n = f"{base}.pt.trace.json.gz", 1 + while os.path.exists(path): # several sessions or files per rank: capture order is file order + path, n = f"{base}.{n}.pt.trace.json.gz", n + 1 + os.makedirs(os.path.dirname(path), exist_ok=True) + self.prof.export_chrome_trace(path) + except Exception: + runtime.write_error("capture profiler export") + + def __exit__(self, *exc): + if self.prof is not None: + try: + graphs.after_capture.remove(self._graph_captured) + except ValueError: + pass + modules.disable() + self._export() + self.prof = None + runtime.flush_all() + return False diff --git a/inferencex-e2e/benchmarks/profiling/profile_windows.py b/inferencex-e2e/benchmarks/profiling/profile_windows.py new file mode 100644 index 0000000000..f242800206 --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/profile_windows.py @@ -0,0 +1,365 @@ +"""Open torch profiler windows on every engine during an AgentX replay. + +Usage: profile_windows.py WINDOWS_JSON LOG_PATH AIPERF_LOG STEPS_DIR [SERVER_URL...] + +WINDOWS_JSON lists [anchor, delay_seconds, iterations] windows. An anchor +"warmup" or "profiling" is an aiperf phase: the window opens delay_seconds +after AIPERF_LOG records that phase's start. AgentX warmup is the lanes' long +first turns (prefill-heavy). The anchor "decode" is steady decode: from the +warmup start, the window opens delay_seconds after the engines' step logs in +STEPS_DIR show at least DECODE_SHARE of recent steps replaying FULL CUDA +graphs. Each new turn re-prefills, so how long that takes depends on +concurrency; if it has not happened DECODE_DEADLINE_S into the measured phase, +the window opens then. If aiperf never logs a phase, windows fall back to this +script's start. + +How a window opens depends on the engine (INFX_PROF_CONTROL): + +- http (vLLM): POST /start_profile to every server; vLLM stops on its own + after its configured max_iterations, and a /stop_profile after a grace + period closes a window that did not fill. +- file (SGLang, TRT-LLM): append a request to control/windows.jsonl for the + profiling hook in every rank (infx_prof.windows.FileWindows). Every rank + starts at the same step, a little past the latest one logged, and stops on + its own after `iterations` steps; a stop request after the grace period + closes a window that did not fill. + +Each window's GPU clocks are sampled through NVML (clock_sampler.py) from +just before it opens until every engine that started profiling has stopped +(logged by the engine patch under profiler/), plus the GPU's lag behind it. + +The measured phase is capped at INFX_PROFILE_DURATION seconds, long enough for +steady decode at any concurrency. Once the last window has closed and the +measured phase has results to export, this script ends the replay early with +SIGINT, which aiperf treats as a user cancel: it stops issuing requests, +exports what it measured and exits zero. +""" + +import glob +import json +import os +import re +import signal +import sys +import time +import urllib.request + +from clock_sampler import ClockSampler + +PHASE_START = re.compile(r"Phase (\S+) \((\S+)\) started") +POLL_S = 5.0 +NO_LOG_FALLBACK_S = 900.0 # aiperf logs its first phase within minutes of starting +PHASE_TIMEOUT_S = 7200.0 +DECODE_SHARE = 0.9 +DECODE_LOOKBACK_S = 20.0 +DECODE_MIN_STEPS = 50 +MEASURED_S = float(os.environ.get("INFX_PROFILE_DURATION") or 600) +# Latest a decode window opens in the measured phase, leaving it room to close. +DECODE_DEADLINE_S = max(0.0, MEASURED_S - 150.0) +# Measured replay before an early stop: aiperf fails a run with no results. +MIN_MEASURED_S = 60.0 +CLOCK_WINDOW_MAX_S = 120.0 # a window's steps take seconds; stop sampling regardless +# The GPU finishes a window's last step after its engine stops profiling +# (launches are asynchronous): keep sampling for twice its longest step past that. +CLOCK_TAIL_S = 0.5 +CLOCK_TAIL_MAX_S = 10.0 +PROFILER_START_WAIT_S = 60.0 # engines that never log a start: fall back to step counts +# File control: a window starts this far past the latest logged step, so every +# rank sees the request before reaching it. +START_LEAD_S = 1.0 +START_LEAD_MIN_STEPS = 4 + + +def post(url: str) -> str: + request = urllib.request.Request(url, data=b"", method="POST") + try: + with urllib.request.urlopen(request, timeout=600) as response: + return f"{response.status}" + except Exception as error: # a failed window must not fail the benchmark + return f"error: {error}" + + +_seen: dict[str, dict[str, float]] = {} + + +def phase_starts(aiperf_log: str) -> dict[str, float]: + """aiperf phase name -> wall time this script first saw its start logged.""" + seen = _seen.setdefault(aiperf_log, {}) + try: + with open(aiperf_log) as f: + for line in f: + m = PHASE_START.search(line) + if m: + seen.setdefault(m.group(2), time.time()) + except OSError: + pass + return seen + + +def wait_for_phase(aiperf_log: str, phase: str, started: float) -> tuple[float, str]: + """Wall time of `phase`'s start (to within one poll), else a fallback and why.""" + while True: + starts = phase_starts(aiperf_log) + if phase in starts: + return starts[phase], "aiperf" + elapsed = time.time() - started + if not starts and elapsed > NO_LOG_FALLBACK_S: + return started, "no aiperf phases logged" + if elapsed > PHASE_TIMEOUT_S: + return time.time(), "phase never started" + time.sleep(POLL_S) + + +_offsets: dict[str, int] = {} + + +def step_graph(record: dict) -> str | None: + """The CUDA graph mode a logged step ran in.""" + if record.get("graph"): + return record["graph"] + dispatch = record.get("dispatch") # step logs from before the record had `graph` + return dispatch[-1].get("cg_mode") if dispatch else None + +_steps: list[tuple[float, bool]] = [] + + +def read_steps(steps_dir: str) -> list[tuple[float, bool]]: + """(step start unix time, replayed a FULL graph) for every engine step logged so far.""" + for path in glob.glob(os.path.join(steps_dir, "*.jsonl")): + try: + with open(path) as f: + f.seek(_offsets.get(path, 0)) + while line := f.readline(): + if not line.endswith("\n"): + break # the engine is mid-write; reread it next poll + _offsets[path] = f.tell() + record = json.loads(line) + if "step" in record: + _steps.append((record["t0_ns"] / 1e9, step_graph(record) == "FULL")) + except (OSError, ValueError): + continue + return _steps + + +def steady_decode(steps_dir: str) -> bool: + now = time.time() + recent = [full for t, full in read_steps(steps_dir) if t >= now - DECODE_LOOKBACK_S] + return len(recent) >= DECODE_MIN_STEPS and sum(recent) >= DECODE_SHARE * len(recent) + + +def wait_for_decode(aiperf_log: str, steps_dir: str, started: float) -> tuple[float, str]: + """Wall time steady decode was first seen after warmup began, else the measured-phase deadline.""" + wait_for_phase(aiperf_log, "warmup", started) + while True: + if steady_decode(steps_dir): + return time.time(), "steady decode" + measured = phase_starts(aiperf_log).get("profiling") + if measured is not None and time.time() > measured + DECODE_DEADLINE_S: + return time.time(), "measured-phase deadline" + if time.time() - started > PHASE_TIMEOUT_S: + return time.time(), "decode never steady" + time.sleep(POLL_S) + + +class StepCounter: + """Engine steps (scheduled or dummy) each rank's step log records after a time.""" + + def __init__(self, steps_dir: str): + self.steps_dir = steps_dir + self.offsets: dict[str, int] = {} + self.times: dict[str, list[tuple[int, int]]] = {} # (t0_ns, t1_ns) per step + self.latest: dict[str, int] = {} # rank -> latest step index (scheduled or dummy) + + def counts_since(self, t_ns: int) -> dict[str, int]: + for path in glob.glob(os.path.join(self.steps_dir, "*.jsonl")): + try: + with open(path) as f: + f.seek(self.offsets.get(path, 0)) + while line := f.readline(): + if not line.endswith("\n"): + break + self.offsets[path] = f.tell() + record = json.loads(line) + k = record.get("step", record.get("dummy")) + if k is not None: + rank = os.path.basename(path)[:-len(".jsonl")] + self.latest[rank] = max(self.latest.get(rank, -1), int(k)) + if "t0_ns" in record: + self.times.setdefault(path, []).append( + (record["t0_ns"], record.get("t1_ns", record["t0_ns"]))) + except (OSError, ValueError): + continue + return {path: sum(t0 >= t_ns for t0, _ in times) for path, times in self.times.items()} + + def start_steps(self) -> dict[str, int]: + """rank -> the step a window requested now starts at (the same on every rank).""" + self.counts_since(time.time_ns()) + if not self.latest: + return {} + recent = sorted(b - a for times in self.times.values() for a, b in times[-64:]) + step_s = max(recent[len(recent) // 2] / 1e9, 1e-3) if recent else 1.0 + start = max(self.latest.values()) + max(START_LEAD_MIN_STEPS, int(START_LEAD_S / step_s) + 1) + return {rank: start for rank in self.latest} + + def longest_between(self, t0_ns: int, t1_ns: int) -> float: + """Seconds of the longest step any rank logged starting between t0_ns and t1_ns.""" + self.counts_since(t0_ns) + return max((b - a for times in self.times.values() for a, b in times + if t0_ns <= a <= t1_ns), default=0) / 1e9 + + +def profiler_events(prof_dir: str, t_ns: int) -> dict[str, list[tuple[str, int]]]: + """rank -> the engine's profiler start/stop events after t_ns.""" + events = {} + for path in glob.glob(os.path.join(prof_dir, "profiler", "*.jsonl")): + try: + with open(path) as f: + rows = [json.loads(line) for line in f if line.endswith("\n")] + except (OSError, ValueError): + continue + events[os.path.basename(path)] = [(r["event"], r["t_ns"]) for r in rows if r["t_ns"] >= t_ns] + return events + + +def wait_for_profilers(prof_dir: str, t_ns: int) -> int | None: + """Wall time (ns) the last engine that started profiling after t_ns stopped, else None.""" + deadline = time.time() + CLOCK_WINDOW_MAX_S + while time.time() < deadline: + events = profiler_events(prof_dir, t_ns) + started = {rank: ev for rank, ev in events.items() if any(e == "start" for e, _ in ev)} + if not started and time.time() - t_ns / 1e9 > PROFILER_START_WAIT_S: + return None + stops = [max(t for e, t in ev if e == "stop") for ev in started.values() + if any(e == "stop" for e, _ in ev)] + if started and len(stops) == len(started): + return max(stops) + time.sleep(0.2) + return None + + +def wait_for_window_steps(counter: StepCounter, t_ns: int, iterations: int) -> None: + """Until every rank that stepped since t_ns has stepped `iterations` times.""" + deadline = time.time() + CLOCK_WINDOW_MAX_S + while time.time() < deadline: + active = [n for n in counter.counts_since(t_ns).values() if n] + if active and min(active) >= iterations: + return + time.sleep(0.2) + + +class HttpControl: + """Engines with an on-demand profiler endpoint (vLLM).""" + + def __init__(self, servers: list[str], note): + self.servers, self.note = servers, note + + def open(self, index: int, iterations: int) -> None: + for server in self.servers: + self.note(event="start", window=index, server=server, iterations=iterations, + status=post(f"{server}/start_profile")) + + def close(self, index: int) -> None: + for server in self.servers: + self.note(event="stop", window=index, server=server, status=post(f"{server}/stop_profile")) + + +class FileControl: + """Engines whose profiling hook reads window requests from control/windows.jsonl (SGLang, TRT-LLM).""" + + def __init__(self, prof_dir: str, counter: StepCounter, note): + self.path = os.path.join(prof_dir, "control", "windows.jsonl") + self.counter, self.note = counter, note + os.makedirs(os.path.dirname(self.path), exist_ok=True) + + def _append(self, request: dict) -> None: + with open(self.path, "a") as f: + f.write(json.dumps(request) + "\n") + + def open(self, index: int, iterations: int) -> None: + start_at = self.counter.start_steps() + self._append({"window": index, "steps": iterations, "start_at": start_at}) + self.note(event="start", window=index, iterations=iterations, start_at=start_at) + + def close(self, index: int) -> None: + self._append({"stop": index}) + self.note(event="stop", window=index) + + +def aiperf_pids() -> list[int]: + """The aiperf SystemController, which titles itself "aiperf system_controller".""" + pids = [] + for entry in os.listdir("/proc"): + if not entry.isdigit(): + continue + try: + with open(f"/proc/{entry}/cmdline", "rb") as f: + title = f.read().replace(b"\0", b" ").decode(errors="replace").split() + except OSError: + continue + if title[:2] == ["aiperf", "system_controller"]: + pids.append(int(entry)) + return pids + + +def stop_replay(aiperf_log: str, started: float, note) -> None: + """End the measured phase once it has results; aiperf exports them and exits zero.""" + measured, basis = wait_for_phase(aiperf_log, "profiling", started) + if basis != "aiperf": + note(event="replay_stop", status=f"skipped: {basis}") + return + time.sleep(max(0.0, measured + MIN_MEASURED_S - time.time())) + pids = aiperf_pids() + for pid in pids: + try: + os.kill(pid, signal.SIGINT) + except OSError: + pass + note(event="replay_stop", pids=pids) + + +def main() -> None: + windows = json.loads(sys.argv[1]) + log_path, aiperf_log, steps_dir = sys.argv[2], sys.argv[3], sys.argv[4] + servers = [url.rstrip("/") for url in sys.argv[5:]] + started = time.time() + clocks = ClockSampler(os.path.join(os.path.dirname(steps_dir), "clocks")) + counter = StepCounter(steps_dir) + with open(log_path, "a", buffering=1) as log: + def note(**record): + record["t_unix"] = time.time() + log.write(json.dumps(record) + "\n") + + mode = os.environ.get("INFX_PROF_CONTROL", "http") + control = (FileControl(os.path.dirname(steps_dir), counter, note) if mode == "file" + else HttpControl(servers, note)) + note(event="servers", servers=servers, windows=windows, aiperf_log=aiperf_log, control=mode, + clocks_error=clocks.error) + for index, (anchor, delay, iterations) in enumerate(windows): + if anchor == "decode": + anchor_t, basis = wait_for_decode(aiperf_log, steps_dir, started) + else: + anchor_t, basis = wait_for_phase(aiperf_log, anchor, started) + note(event="anchor", window=index, anchor=anchor, basis=basis, anchor_unix=anchor_t) + time.sleep(max(0.0, anchor_t + float(delay) - time.time())) + window_t0 = time.time() + clocks.start(index) + control.open(index, int(iterations)) + window_ns = int(window_t0 * 1e9) + stopped_ns = wait_for_profilers(os.path.dirname(steps_dir), window_ns) + if stopped_ns is None: # no profiler log: the window's steps, as logged + wait_for_window_steps(counter, window_ns, int(iterations)) + stopped_ns = time.time_ns() + tail = 2 * counter.longest_between(window_ns, stopped_ns) + time.sleep(max(0.0, stopped_ns / 1e9 + CLOCK_TAIL_S + min(tail, CLOCK_TAIL_MAX_S) + - time.time())) + clocks.stop() + note(event="clocks", window=index, polls=clocks.polls, + seconds=round(time.time() - window_t0, 3)) + # Generous grace: a step is well under a second at every agentic point. + time.sleep(max(0.0, window_t0 + max(120.0, 4.0 * float(iterations)) - time.time())) + control.close(index) + stop_replay(aiperf_log, started, note) + + +if __name__ == "__main__": + main() diff --git a/inferencex-e2e/benchmarks/profiling/sitecustomize.py b/inferencex-e2e/benchmarks/profiling/sitecustomize.py new file mode 100644 index 0000000000..40b6fcdaab --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/sitecustomize.py @@ -0,0 +1,36 @@ +"""Load the op-attribution profiler (infx_prof) in every engine process; see infx_prof/__init__.py. + +The engine's launch puts this directory on PYTHONPATH, which shadows the +image's own sitecustomize, if any: run that one too. +""" + +import importlib.machinery +import importlib.util +import os +import sys + + +def _chain_next_sitecustomize(): + here = os.path.dirname(os.path.abspath(__file__)) + for entry in sys.path: + if not entry or os.path.abspath(entry) == here: + continue + spec = importlib.machinery.PathFinder.find_spec("sitecustomize", [entry]) + if spec is not None and spec.origin and spec.loader is not None: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return + + +if os.environ.get("INFX_PROF_DIR"): + import infx_prof + + infx_prof.boot() + +try: + _chain_next_sitecustomize() +except Exception: + if os.environ.get("INFX_PROF_DIR"): + from infx_prof import runtime + + runtime.write_error("chain sitecustomize") diff --git a/inferencex-e2e/benchmarks/profiling/synthetic_load.py b/inferencex-e2e/benchmarks/profiling/synthetic_load.py new file mode 100644 index 0000000000..5d2cdb7b4c --- /dev/null +++ b/inferencex-e2e/benchmarks/profiling/synthetic_load.py @@ -0,0 +1,104 @@ +"""Fixed-shape load for synthetic profiling: a prefill phase, then a decode phase. + +Usage: synthetic_load.py SERVER_URL PHASE_LOG WINDOWS_LOG RESULT_JSON + +Requests are random token-id prompts of exactly INFX_SYNTH_ISL tokens sent to +/v1/completions, CONC at a time, so no tokenizer or dataset is involved. The +prefill phase asks for one token per request until the window client closes +its first window; the decode phase then generates INFX_SYNTH_OSL tokens per +request (ignore_eos) until STOP_FILE appears. Each phase start is written to +PHASE_LOG in aiperf's "Phase () started" form, so the window +client anchors its windows on them unchanged: warmup = prefill, profiling = +decode. The prompt length stands in for KV-cache context: decode steps attend +over ISL + generated tokens per request. +""" + +import json +import os +import random +import sys +import threading +import time +import urllib.request + +PHASE_MAX_S = float(os.environ.get("INFX_SYNTH_PHASE_MAX_S") or 900) +TOKEN_RANGE = (1000, 30000) # ordinary vocabulary ids, clear of special tokens + + +def post(url, body, timeout=1800): + request = urllib.request.Request(url, data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(request, timeout=timeout) as response: + return json.loads(response.read()) + + +def model_name(server): + with urllib.request.urlopen(f"{server}/v1/models", timeout=60) as response: + return json.loads(response.read())["data"][0]["id"] + + +def first_window_done(windows_log): + """The window client has closed window 0 (its clock note follows the profiler stop).""" + try: + with open(windows_log) as f: + return any(json.loads(line).get("event") == "clocks" and json.loads(line).get("window") == 0 + for line in f if line.strip()) + except (OSError, ValueError): + return False + + +def run_phase(name, server, model, conc, isl, osl, done, stats, phase_log): + with open(phase_log, "a", buffering=1) as f: + f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - synthetic_load - NOTICE - " + f"Phase {name} ({name}) started | conc={conc} isl={isl} osl={osl}\n") + deadline = time.time() + PHASE_MAX_S + rng = random.Random(name) + lock = threading.Lock() + + def lane(): + while not done() and time.time() < deadline: + with lock: + prompt = [rng.randrange(*TOKEN_RANGE) for _ in range(isl)] + body = {"model": model, "prompt": prompt, "max_tokens": osl, "min_tokens": osl, + "ignore_eos": True, "temperature": 0.0} + try: + post(f"{server}/v1/completions", body) + key = "completed" + except Exception as error: # a request lost to a window's export stall is not fatal + key = "errors" + stats.setdefault("last_error", repr(error)[:300]) + with lock: + stats[key] = stats.get(key, 0) + 1 + + started = time.time() + for _ in range(conc): + threading.Thread(target=lane, daemon=True).start() + # Return when the phase is done; lanes still mid-request finish or die with the process. + while not done() and time.time() < deadline: + time.sleep(0.5) + stats["seconds"] = round(time.time() - started, 1) + stats["stopped_by"] = "deadline" if time.time() >= deadline else "window client" + + +def main(): + server, phase_log, windows_log, result_json = sys.argv[1:5] + conc = int(os.environ["CONC"]) + isl = int(os.environ.get("INFX_SYNTH_ISL") or 8192) + osl = int(os.environ.get("INFX_SYNTH_OSL") or 4096) + stop_file = os.environ["INFX_SYNTH_STOP_FILE"] + model = model_name(server) + result = {"mode": "synthetic", "model": model, "conc": conc, "isl": isl, "osl": osl, + "prefill": {}, "decode": {}} + run_phase("warmup", server, model, conc, isl, 1, lambda: first_window_done(windows_log), + result["prefill"], phase_log) + run_phase("profiling", server, model, conc, isl, osl, lambda: os.path.exists(stop_file), + result["decode"], phase_log) + with open(result_json, "w") as f: + json.dump(result, f, indent=1) + print(json.dumps(result)) + # Decode requests run OSL tokens; the window may close before any finishes. + return 0 if result["prefill"].get("completed") else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/inferencex-e2e/benchmarks/srt_agentic.sh b/inferencex-e2e/benchmarks/srt_agentic.sh index e4298e64d6..795cb75098 100644 --- a/inferencex-e2e/benchmarks/srt_agentic.sh +++ b/inferencex-e2e/benchmarks/srt_agentic.sh @@ -56,6 +56,39 @@ if [[ -z "${AIPERF_SERVER_URL:-}" ]]; then fi echo "Using srt-slurm frontend endpoint: $AIPERF_SERVER_URL" +# srt-slurm's logical worker endpoints: each vLLM server's control port. +profile_servers=() +if [[ -n "${INFX_PROFILE_WINDOWS:-}" ]]; then + IFS=',' read -r -a profile_endpoints <<< "${SRT_AGG_ENDPOINTS:-}" + for endpoint in "${profile_endpoints[@]}"; do + [[ -n "$endpoint" ]] && profile_servers+=("http://$endpoint") + done + (( ${#profile_servers[@]} )) || profile_servers=("$AIPERF_SERVER_URL") +fi + +# Synthetic op-attribution profiling (INFX_PROFILE mode=synthetic): fixed-shape +# prefill then decode load in place of the AgentX replay; no trace dataset or aiperf. +if [[ "${INFX_PROFILE_MODE:-}" == "synthetic" ]]; then + mkdir -p "$INFX_PROF_DIR" + synth_phases="$INFX_PROF_DIR/synthetic_phases.log" + synth_windows="$INFX_PROF_DIR/windows_conc${CONC}.jsonl" + export INFX_SYNTH_STOP_FILE="$INFX_PROF_DIR/synthetic.stop" + : > "$synth_phases" + rm -f "$INFX_SYNTH_STOP_FILE" + python3 "$INFMAX_CONTAINER_WORKSPACE/benchmarks/profiling/profile_windows.py" \ + "$INFX_PROFILE_WINDOWS" "$synth_windows" "$synth_phases" "$INFX_PROF_DIR/steps" \ + "${profile_servers[@]}" & + synth_windows_pid=$! + python3 "$INFMAX_CONTAINER_WORKSPACE/benchmarks/profiling/synthetic_load.py" \ + "$AIPERF_SERVER_URL" "$synth_phases" "$synth_windows" \ + "${AGENTIC_OUTPUT_DIR:-/logs}/${RESULT_FILENAME}.json" & + synth_load_pid=$! + wait "$synth_windows_pid" || true + touch "$INFX_SYNTH_STOP_FILE" + wait "$synth_load_pid" + exit $? +fi + # A router frontend does not re-export engine metrics; read them from each worker. if [[ -z "${AIPERF_SERVER_METRICS_URLS:-}" && "${SRTCTL_FRONTEND_TYPE:-}" != dynamo ]]; then endpoints="${SRT_AGG_ENDPOINTS:-${SRT_PREFILL_ENDPOINTS:+$SRT_PREFILL_ENDPOINTS,}${SRT_DECODE_ENDPOINTS:-}}" @@ -88,4 +121,18 @@ fi if [[ -n "${AIPERF_BENCHMARK_GRACE_PERIOD:-}" ]]; then REPLAY_CMD+=" --benchmark-grace-period $AIPERF_BENCHMARK_GRACE_PERIOD" fi +# Op-attribution profiling (INFX_PROFILE): torch windows on every worker. +profile_windows_pid="" +if [[ -n "${INFX_PROFILE_WINDOWS:-}" ]]; then + mkdir -p "$INFX_PROF_DIR" + python3 "$INFMAX_CONTAINER_WORKSPACE/benchmarks/profiling/profile_windows.py" \ + "$INFX_PROFILE_WINDOWS" "$INFX_PROF_DIR/windows_conc${CONC}.jsonl" \ + "$RESULT_DIR/aiperf_artifacts/logs/aiperf.log" "$INFX_PROF_DIR/steps" \ + "${profile_servers[@]}" & + profile_windows_pid=$! +fi run_agentic_replay_and_write_outputs "$RESULT_DIR" +if [[ -n "$profile_windows_pid" ]]; then + kill "$profile_windows_pid" 2>/dev/null || true + wait "$profile_windows_pid" 2>/dev/null || true +fi diff --git a/inferencex-e2e/infx/launch/drivers/srt/collect.py b/inferencex-e2e/infx/launch/drivers/srt/collect.py index 0324f4e775..5681537a29 100644 --- a/inferencex-e2e/infx/launch/drivers/srt/collect.py +++ b/inferencex-e2e/infx/launch/drivers/srt/collect.py @@ -9,6 +9,7 @@ import os import shutil import sys +import tarfile import time from collections.abc import Callable from pathlib import Path @@ -39,6 +40,8 @@ SINGLE_NODE_LOGS = "srt-single-node-logs.tar.gz" MULTINODE_LOGS = "multinode_server_logs.tar.gz" +PROFILE_DIR = "infx_profile" +PROFILE_ARCHIVE = "infx-profile.tar" def _copy_tree_into(source: Path, destination: Path) -> None: @@ -58,8 +61,18 @@ def finish_single_node(run: SrtRun, submitted: Submitted, fetched: Path) -> int: if not output.is_dir(): return 0 rc = 0 - bundle_server_logs(output, run.workspace / SINGLE_NODE_LOGS) logs = output / "logs" + # Op-attribution profiles (INFX_PROFILE) upload as their own artifact. + profile = logs / PROFILE_DIR + if profile.is_dir(): + try: + with tarfile.open(run.workspace / PROFILE_ARCHIVE, "w") as archive: + archive.add(profile, arcname=PROFILE_DIR) + shutil.rmtree(profile) + except OSError as error: + print(f"ERROR: failed to stage the op-attribution profile: {error}", file=sys.stderr) + rc = 1 + bundle_server_logs(output, run.workspace / SINGLE_NODE_LOGS) result = logs / f"{run.request.result_filename}.json" for artifact in [result, *sorted(logs.glob("gpu_metrics*"))]: if artifact.is_file(): diff --git a/inferencex-e2e/infx/srt_slurm/single_node.py b/inferencex-e2e/infx/srt_slurm/single_node.py index 136b1afe51..38529b4a71 100644 --- a/inferencex-e2e/infx/srt_slurm/single_node.py +++ b/inferencex-e2e/infx/srt_slurm/single_node.py @@ -3,10 +3,12 @@ from __future__ import annotations import argparse +import copy import json import os import re from collections.abc import Mapping +from dataclasses import dataclass from pathlib import Path from typing import Any @@ -136,6 +138,334 @@ def validate_recipe(recipe: dict[str, Any], environment: Mapping[str, str]) -> N raise ValueError(f"Single-node SRT {name}: recipe/matrix {actual!r} != {wanted!r}") +PROFILE_DIR = "/logs/infx_profile" +PROFILE_DEFAULTS: dict[str, Any] = { + # [anchor, seconds after it, engine iterations] per torch window. Anchors are + # aiperf phases ("warmup", "profiling") or "decode", the first steady decode + # (FULL CUDA graph replays) after warmup starts. AgentX warmup is the lanes' + # long first turns, all sent at its start (prefill-heavy; at low concurrency + # they have drained within a minute); each measured-phase turn re-prefills + # first, so steady decode's arrival depends on concurrency. + "windows": [["warmup", 0, 32], ["decode", 0, 32]], + # workers whose CUDA graph capture is profiled; "all" profiles every rank. + # Unset: the engine's default (every rank under SGLang / TRT-LLM attention DP, + # whose ranks' graphs differ; else dp0_tp0) + "capture_ranks": None, + # Graphs per capture trace (a long capture overflows one profiler session's + # buffer); 0 keeps one trace per capture session. Off until rotation is shown not + # to cause the SIGSEGV seen in SGLang's end-of-capture garbage collection on MI355X. + "capture_graphs_per_file": 0, + # Host memory kept free of CPU KV offload for the profiler's trace buffers; + # offload recipes otherwise size the pool to nearly the whole host. + "host_headroom_gib": 128, + # "agentic" replays AgentX; "synthetic" drives fixed shapes instead: CONC + # random-token prompts of `isl` tokens, one-token outputs for the prefill + # window, then `osl`-token generations for the decode window. No trace + # dataset or AgentX warmup, so a point takes minutes after engine start-up; + # `isl` stands in for the KV context decode steps attend over. + "mode": "agentic", + "isl": 8192, + "osl": 4096, +} +PROFILE_MODES = ("agentic", "synthetic") +PROFILE_PHASES = ("warmup", "decode", "profiling") +# Cap on the measured replay past its last phase-anchored window's start. Steady +# decode can take minutes to arrive under data parallelism; the window client +# ends the replay once its windows close. A profiled run's throughput is not a result. +PROFILE_MEASURED_CAP_SECONDS = 600 + + +def offload_headroom_arguments(role_args: Mapping[str, Any], headroom_gib: float) -> list[str]: + """Shrink a CPU KV-offload pool so profiling leaves the host headroom_gib free.""" + raw = role_args.get("kv-transfer-config") + if not raw or headroom_gib <= 0: + return [] + transfer = json.loads(raw) if isinstance(raw, str) else dict(raw) + extra = transfer.get("kv_connector_extra_config") or {} + per_rank = extra.get("cpu_bytes_to_use_per_rank") + if per_rank is None: + return [] + ranks = int(role_args.get("data-parallel-size", 1)) * int( + role_args.get("tensor-parallel-size", 1) + ) + shrunk = int(per_rank) - int(headroom_gib * 2**30 / ranks) + if shrunk <= 0: + raise ValueError("INFX_PROFILE host_headroom_gib exceeds the recipe's CPU offload pool") + transfer["kv_connector_extra_config"] = {**extra, "cpu_bytes_to_use_per_rank": shrunk} + return ["--set", f"roles.agg.args.kv-transfer-config={json.dumps(transfer)}"] + + +def mooncake_headroom_arguments( + recipe: Mapping[str, Any], role_args: Mapping[str, Any], headroom_gib: float +) -> list[str]: + """Shrink an embedded Mooncake store's per-rank segment so profiling leaves the host headroom_gib. + + Each TP rank registers global_segment_size of host memory for RDMA; with the + profiler's buffers on top, B300 Kimi K3 ran out registering it. + """ + if headroom_gib <= 0: + return [] + ranks = int(role_args.get("data-parallel-size", 1)) * int( + role_args.get("tensor-parallel-size", 1) + ) + services = copy.deepcopy(list(recipe.get("services") or [])) + changed = False + for service in services: + if not isinstance(service, Mapping) or service.get("type") != "mooncake-master": + continue + store = (service.get("options") or {}).get("store_config") or {} + match = re.fullmatch( + r"(\d+(?:\.\d+)?)\s*([KMGT]?B)", + str(store.get("global_segment_size") or ""), + re.IGNORECASE, + ) + if not match: + continue + scale = {"KB": 1 / 2**20, "MB": 1 / 2**10, "GB": 1, "TB": 2**10}[match.group(2).upper()] + shrunk = float(match.group(1)) * scale - headroom_gib / ranks + if shrunk <= 0: + raise ValueError("INFX_PROFILE host_headroom_gib exceeds the recipe's Mooncake segment") + store["global_segment_size"] = f"{int(shrunk)}GB" + changed = True + # srtctl applies --set to base and every override_* section, so an indexed + # services[i] path fails on variants without services; replace the whole list. + return ["--set", f"services={json.dumps(services)}"] if changed else [] + + +def _recipe_env_int(recipe: Mapping[str, Any] | None, name: str) -> int: + """An integer the recipe's agg role (or the whole recipe) sets in its environment, else 0.""" + for env in ( + ((recipe or {}).get("roles") or {}).get("agg", {}).get("env") or {}, + (recipe or {}).get("environment") or {}, + ): + try: + return int(env[name]) + except (KeyError, TypeError, ValueError): + continue + return 0 + + +def _set(path: str, value: Any) -> list[str]: + return ["--set", f"{path}={json.dumps(value)}"] + + +@dataclass(frozen=True) +class ProfileRequest: + """One profiled point: INFX_PROFILE's settings and the selected recipe variant.""" + + settings: Mapping[str, Any] + role_args: Mapping[str, Any] + recipe: Mapping[str, Any] + + @property + def headroom_gib(self) -> float: + return float(self.settings["host_headroom_gib"]) + + +class ProfiledEngine: + """How one engine is profiled: its hooks, how windows open, and the recipe changes that needs. + + Every engine loads the infx_prof hooks through PYTHONPATH (sitecustomize); + INFX_PROF_ENGINE picks its adapter. `control` is how the window client + opens windows: "http" (the engine's own profiler endpoint) or "file" + (control/windows.jsonl, read by the hooks at step boundaries). + """ + + name = "" + control = "file" + attention_dp_arg = "" # the role arg enabling attention DP, when its ranks' graphs differ + + def worker_env(self, profile: ProfileRequest) -> dict[str, str]: + return { + "INFX_PROF_DIR": PROFILE_DIR, + "INFX_PROF_ENGINE": self.name, + "INFX_PROF_CAPTURE_GRAPHS_PER_FILE": str( + int(profile.settings["capture_graphs_per_file"]) + ), + "INFX_PROF_CAPTURE_RANKS": profile.settings["capture_ranks"] + or self.capture_ranks(profile), + "PYTHONPATH": "/infmax-workspace/benchmarks/profiling", + # Compiled code gets its module markers when Inductor generates its + # wrapper, so compile from scratch rather than load cached wrappers. + "TORCHINDUCTOR_FX_GRAPH_CACHE": "0", + "TORCHINDUCTOR_AUTOGRAD_CACHE": "0", + } + + def role_arguments(self, profile: ProfileRequest) -> list[str]: + """--set overrides of the engine's arguments (and services) for a profiled run.""" + raise NotImplementedError + + def capture_ranks(self, profile: ProfileRequest) -> str: + """Workers whose graph capture is profiled when INFX_PROFILE does not say: every rank + under attention DP (SGLang / TRT-LLM ranks capture different graphs: per-rank + buffers change copy kernels), else dp0_tp0.""" + dp = self.attention_dp_arg and profile.role_args.get(self.attention_dp_arg) + return "all" if dp else "dp0_tp0" + + +class VllmProfiling(ProfiledEngine): + name = "vllm" + control = "http" + + def worker_env(self, profile: ProfileRequest) -> dict[str, str]: + timeout = max(1800, _recipe_env_int(profile.recipe, "VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS")) + ready = max(3600, _recipe_env_int(profile.recipe, "VLLM_ENGINE_READY_TIMEOUT_S")) + return { + **super().worker_env(profile), + # A window's trace export blocks its worker; keep peers from timing out + # (VLLM_RPC_TIMEOUT before vLLM 0.31, VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS from it). + "VLLM_RPC_TIMEOUT": "1800000", + "VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS": str(timeout), + # Uncached compiles lengthen start-up past the 600 s default. + "VLLM_ENGINE_READY_TIMEOUT_S": str(ready), + "VLLM_DISABLE_COMPILE_CACHE": "1", + } + + def role_arguments(self, profile: ProfileRequest) -> list[str]: + # vLLM's profiler window length is fixed per engine; every window uses the first's. + profiler_config = { + "profiler": "torch", + "torch_profiler_dir": f"{PROFILE_DIR}/torch", + "torch_profiler_record_shapes": True, + # Python stacks make exports of eager steps outlast the RPC timeout; + # modules mark themselves instead (benchmarks/profiling/infx_prof). + "torch_profiler_with_stack": False, + # The summary table walks every event in Python on all ranks at once; + # on offload-sized hosts that exhausted memory. The raw trace suffices. + "torch_profiler_dump_cuda_time_total": False, + "ignore_frontend": True, + "max_iterations": int(profile.settings["windows"][0][2]), + } + return [ + "--set", + f"roles.agg.args.profiler-config={json.dumps(profiler_config)}", + *offload_headroom_arguments(profile.role_args, profile.headroom_gib), + *mooncake_headroom_arguments(profile.recipe, profile.role_args, profile.headroom_gib), + ] + + +class SglangProfiling(ProfiledEngine): + name = "sglang" + + attention_dp_arg = "enable-dp-attention" + + # A window's export blocks its scheduler; neither the watchdog nor peers' + # collectives may give up on it. + TIMEOUT_SECONDS = 1800 + + def role_arguments(self, profile: ProfileRequest) -> list[str]: + args = profile.role_args + overrides = [] + for flag in ("watchdog-timeout", "dist-timeout"): + overrides += _set( + f"roles.agg.args.{flag}", max(self.TIMEOUT_SECONDS, int(args.get(flag) or 0)) + ) + # HiCache's host pool is hicache-size GB per rank when set (it overrides + # hicache-ratio); a ratio-sized pool depends on the GPU pool and is left as is. + size = args.get("hicache-size") + if args.get("enable-hierarchical-cache") and size and profile.headroom_gib > 0: + shrunk = int( + float(size) - profile.headroom_gib / int(args.get("tensor-parallel-size", 1)) + ) + if shrunk <= 0: + raise ValueError("INFX_PROFILE host_headroom_gib exceeds the recipe's HiCache pool") + overrides += _set("roles.agg.args.hicache-size", shrunk) + return overrides + + +class TrtllmProfiling(ProfiledEngine): + name = "trtllm" + + attention_dp_arg = "enable_attention_dp" + + def role_arguments(self, profile: ProfileRequest) -> list[str]: + # The KV host tier is kv_cache_config.host_cache_size bytes per rank. + args = profile.role_args + size = (args.get("kv_cache_config") or {}).get("host_cache_size") + if not size or profile.headroom_gib <= 0: + return [] + ranks = int(args.get("tensor_parallel_size", 1)) * int( + args.get("pipeline_parallel_size", 1) + ) + shrunk = int(size) - int(profile.headroom_gib * 2**30 / ranks) + if shrunk <= 0: + raise ValueError("INFX_PROFILE host_headroom_gib exceeds the recipe's KV host cache") + return _set("roles.agg.args.kv_cache_config.host_cache_size", shrunk) + + +PROFILED_ENGINES = { + engine.name: engine for engine in (VllmProfiling(), SglangProfiling(), TrtllmProfiling()) +} + + +def profiling_arguments( + environment: Mapping[str, str], + role_args: Mapping[str, Any] | None = None, + recipe: Mapping[str, Any] | None = None, +) -> list[str]: + """Op-attribution profiling: capture hooks, step log and torch windows. + + INFX_PROFILE is a JSON object overriding PROFILE_DEFAULTS; empty disables. + Everything lands in PROFILE_DIR, which the launcher uploads as its own artifact. + """ + raw = environment.get("INFX_PROFILE", "") + if not raw: + return [] + engine = PROFILED_ENGINES.get(SINGLE_NODE_ENGINES.get(environment["FRAMEWORK"], "")) + if engine is None: + raise ValueError(f"INFX_PROFILE supports {sorted(PROFILED_ENGINES)} recipes") + settings = {**PROFILE_DEFAULTS, **json.loads(raw)} + windows = settings["windows"] + if not windows or any( + len(window) != 3 + or window[0] not in PROFILE_PHASES + or int(window[1]) < 0 + or int(window[2]) <= 0 + for window in windows + ): + raise ValueError( + f"INFX_PROFILE windows must be [phase in {PROFILE_PHASES}, delay_seconds, iterations]" + ) + if [PROFILE_PHASES.index(w[0]) for w in windows] != sorted( + PROFILE_PHASES.index(w[0]) for w in windows + ): + raise ValueError("INFX_PROFILE windows must be ordered by phase") + profile = ProfileRequest(settings, role_args or {}, recipe or {}) + overrides = engine.role_arguments(profile) + for name, value in engine.worker_env(profile).items(): + overrides += _set(f"roles.agg.env.{name}", value) + measured = [int(w[1]) for w in windows if w[0] == "profiling"] + duration = int( + settings.get("duration") or max(measured, default=0) + PROFILE_MEASURED_CAP_SECONDS + ) + for name, value in ( + ("INFX_PROFILE_DURATION", str(duration)), + ("INFX_PROFILE_WINDOWS", json.dumps(windows)), + ("INFX_PROF_DIR", PROFILE_DIR), + ("INFX_PROF_CONTROL", engine.control), + ): + overrides += _set(f"benchmark.env.{name}", value) + if settings["mode"] not in PROFILE_MODES: + raise ValueError(f"INFX_PROFILE mode must be one of {PROFILE_MODES}") + if settings["mode"] == "synthetic": + for name in ("isl", "osl"): + if not isinstance(settings[name], int) or settings[name] <= 0: + raise ValueError(f"INFX_PROFILE {name} must be a positive integer") + for name, value in ( + ("INFX_PROFILE_MODE", "synthetic"), + ("INFX_SYNTH_ISL", settings["isl"]), + ("INFX_SYNTH_OSL", settings["osl"]), + ): + overrides += _set(f"benchmark.env.{name}", str(value)) + # A window's export pauses the engine; its requests must not abort the replay. + for name in ("AIPERF_FAILED_REQUEST_THRESHOLD", "AIPERF_LIVE_FAILED_REQUEST_THRESHOLD"): + overrides += _set(f"benchmark.env.{name}", "1.0") + # The window client ends the replay with a user cancel, after which aiperf + # exports results but not server metrics; a profiled run charts none. + overrides += _set("benchmark.env.AIPERF_REQUIRED_SERVER_METRIC_PREFIX", "") + return overrides + + def runtime_arguments(config: str, environment: Mapping[str, str]) -> list[str]: """Bind only runtime-owned values after validating the selected recipe.""" _, recipe = select_recipe(config, environment) @@ -145,6 +475,7 @@ def runtime_arguments(config: str, environment: Mapping[str, str]) -> list[str]: # Exclusive nodes include idle GPUs. Restrict each server/client step to # the serving GPU count so client-side power collection sees the same set. overrides = ["--set", f"srun_options.gpus-per-node={json.dumps(environment['GPU_COUNT'])}"] + overrides += profiling_arguments(environment, recipe["roles"]["agg"]["args"], recipe) if environment.get("SRT_SRUN_OPTIONS"): options = json.loads(environment["SRT_SRUN_OPTIONS"]) if not isinstance(options, dict) or any( diff --git a/inferencex-e2e/infx/tests/launch/fake_slurm.py b/inferencex-e2e/infx/tests/launch/fake_slurm.py index 851a62b06d..648e5256bf 100644 --- a/inferencex-e2e/infx/tests/launch/fake_slurm.py +++ b/inferencex-e2e/infx/tests/launch/fake_slurm.py @@ -127,6 +127,9 @@ for conc in os.environ["CONC_LIST"].split(): (workspace / f"{result}_conc{conc}.json").write_text(json.dumps({"conc": int(conc)})) (logs / "agentic" / f"conc_{conc}").mkdir(parents=True) +if os.environ.get("FAKE_PROFILE"): + (logs / "infx_profile" / "steps").mkdir(parents=True) + (logs / "infx_profile" / "steps" / "dp0_tp0.jsonl").write_text('{"step": 0}\n') if os.environ.get("RUN_EVAL") == "true" or os.environ.get("EVAL_ONLY") == "true": (logs / "eval_results").mkdir() (logs / "eval_results" / "results_gsm8k.json").write_text("{}") diff --git a/inferencex-e2e/infx/tests/launch/test_srt_driver.py b/inferencex-e2e/infx/tests/launch/test_srt_driver.py index 176f8cb71a..52a03f9d2b 100644 --- a/inferencex-e2e/infx/tests/launch/test_srt_driver.py +++ b/inferencex-e2e/infx/tests/launch/test_srt_driver.py @@ -153,6 +153,17 @@ def test_single_node_point_stages_workflow_artifacts(harness): assert lines(harness.logs, "scancel") == [] +def test_single_node_profile_stages_its_own_archive(harness): + workspace = harness.workspace + env = {**single_node_env(harness, "h200-cw"), "FAKE_PROFILE": "1"} + assert_ok(launch(env, harness.config, workspace)) + + with tarfile.open(workspace / "infx-profile.tar") as archive: + assert "infx_profile/steps/dp0_tp0.jsonl" in archive.getnames() + with tarfile.open(workspace / "srt-single-node-logs.tar.gz") as logs: + assert not any("infx_profile" in name for name in logs.getnames()) + + def test_single_node_eval_requires_a_successful_eval(harness): env = single_node_env(harness, "h100-cw", RUN_EVAL="true", MAX_MODEL_LEN="1024") assert_ok(launch(env, harness.config, harness.workspace))