Skip to content

refactor(pytorch): clarify scheduler and input-maker control flow#4727

Draft
grimoire wants to merge 6 commits into
InternLM:mainfrom
grimoire:readable-refactor-review
Draft

refactor(pytorch): clarify scheduler and input-maker control flow#4727
grimoire wants to merge 6 commits into
InternLM:mainfrom
grimoire:readable-refactor-review

Conversation

@grimoire

@grimoire grimoire commented Jul 2, 2026

Copy link
Copy Markdown
Collaborator

Description
This PR is a readability-focused refactor of the PyTorch scheduler/input-maker path. It keeps scheduling behavior UNCHANGED while making long-context chunk planning, forward input selection, prefill admission, and BlockTrie ownership easier to reason about.

What Changed

  • Added characterization tests for long-context chunk reservation contracts.
  • Shared long-context chunk planning between scheduler and input maker.
  • Extracted _ForwardInputsTask from InputsMakerAsync._make_forward_inputs() to separate selection policy from payload construction.
  • Split scheduler prefill reorder/admission flow into _PrefillReorderer and _PrefillAdmissionAttempt.
  • Clarified BlockTrie ownership, SSM checkpoint lifecycle, eviction, rollback, and leaf bookkeeping with docstrings/comments.
  • Added warning-only diagnostics for impossible internal states.

Non-Goals

  • No new scheduling feature.
  • No intended behavior change.

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.

1 participant