Skip to content

Inference split: keep host ops between step annotations - #1022

Open
ajassani wants to merge 3 commits into
mainfrom
ajassani/split-keep-interstep-cpu
Open

ajassani wants to merge 3 commits into
mainfrom
ajassani/split-keep-interstep-cpu

Conversation

@ajassani

@ajassani ajassani commented Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Split windows previously kept CPU ops only while they were nested inside each iteration annotation (ts .. ts+dur). Serving stacks launch work after that envelope and before the next step — vLLM compute_logits / lm_head GEMMs are the usual example — so those ops never made it into the extracted trace or GEMM tables.
  • Each root's CPU window now runs until the next same-pattern iteration annotation on that thread. GPU kernels still attach by correlation, so they can sit outside the host window. Neighboring steps do not leak in.
  • Not every model: this bites dense TP=1 vLLM serving where sample sits outside execute_* (Qwen3-0.6B/8B/14B, Llama-3.1-8B). Other vLLM configs can still keep vocab GEMMs in the split report. SGLang in-graph lm_head is a separate --capture_folder issue.

Fixes #1023

Test plan

  • pytest tests/test_split_inference_trace_annotation.py
  • Re-split a vLLM serving trace and confirm compute_logits / vocab-sized GEMMs appear in the extracted window and GEMM sheet
  • Spot-check an SGLang serving split still isolates step[DECODE|EXTEND] windows

vLLM compute_logits and similar post-step work sit after the execute_* envelope; bounding the CPU window by the next matching annotation keeps lm_head GEMMs without leaking the following step.

Co-authored-by: Cursor <cursoragent@cursor.com>
@codecov-commenter

codecov-commenter commented Sep 10, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

Comment thread TraceLens/TraceUtils/split_inference/trace_extraction.py Outdated
Comment thread TraceLens/TraceUtils/split_inference/trace_extraction.py Outdated
Comment thread TraceLens/TraceUtils/split_inference/trace_extraction.py Outdated
ajassani and others added 2 commits September 10, 2026 22:57
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
@gabeweisz
gabeweisz self-requested a review September 11, 2026 17:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Inference split drops vLLM compute_logits / lm_head after execute_* envelope

3 participants