Skip to content

[Refactor][TCPCG] Bind live forward batches in BCG eager calls (2/9) - #42286

Merged
Oasis-Git merged 8 commits into
sgl-project:mainfrom
Oasis-Git:refactor/tcpcg-02-eager-batch
Oct 8, 2026
Merged

Oasis-Git merged 8 commits into
sgl-project:mainfrom
Oasis-Git:refactor/tcpcg-02-eager-batch

Conversation

@Oasis-Git

@Oasis-Git Oasis-Git commented Oct 2, 2026 •

Copy link
Copy Markdown
Collaborator

This is step 2/9, split from the original TCPCG removal PR #41634.

Step 1 (#42285) is merged. Upstream main was merged into this branch with a normal merge commit (no history rewrite). The Files changed tab now contains only step 2: three files. Pre-commit and the eight CPU tests described below passed again on the merged tree; upstream CI must validate the updated head.

BCG eager calls currently retain their capture-time arguments. To let later migrations pass forward_batch directly, replay must supply the prepared serving batch instead of retaining the capture batch. This PR adds that binding while preserving the existing decorator forms for callers migrated in later steps.

Series

  1. Relocate shared graph tensor and DSA head-gate helpers — [Refactor][TCPCG] Relocate shared graph tensor and DSA head-gate helpers (1/9) #42285.
  2. Add explicit batch support to BCG eager regions — this PR.
  3. Retire TCPCG backend selection and configuration; remove TCPCG-only documentation and tests while retaining internal machinery needed by unmigrated callers.
  4. Consolidate Radix attention eager regions and isolate Inkling-specific behavior, including the dependent MLA wrapper migration.
  5. Migrate DSA and DeepSeek eager regions to explicit methods and batch arguments.
  6. Migrate MoE and Mamba eager regions.
  7. Migrate diffusion eager wrappers and standardize their names. This step depends only on step 2.
  8. Remove obsolete runtime contexts, layer registries, split-op plumbing, and temporary compatibility after caller migration.
  9. Simplify quantization paths separately because of their torch.compile implications.

Only steps 1 and 2 are submitted upstream so far. Later steps will be submitted after review. The disputed relocation of inactive compiler/reference files is excluded pending a separate decision.

Changes

  • Support bare @eager_on_graph on functions and methods, including methods with a named forward_batch argument.
  • Bind the prepared serving batch at replay, supporting both positional and keyword arguments. Clear the binding afterward, including on failure, so the graph does not retain capture or serving batches.
  • Preserve graph-owned tensor/output storage and capture-stub behavior.
  • Pass static_forward_batch from the BCG backend into replay.
  • Temporarily accept eager_on_graph(True), enable=True/False, and the existing capture-stub form so unmigrated callers continue to work.

No backend is retired in this step. Callers without a named forward_batch retain their existing argument behavior.

Validation

  • Pre-commit passed on the changed files.
  • Eight CPU tests passed for positional/keyword batch binding, stable output addresses, capture stubs, batch lifetime, eager passthrough, and legacy decorator forms.
  • Those tests use a macOS import shim while exercising the real BCG implementation. They do not exercise CUDA graph capture; GPU execution and upstream CI remain to be validated.

CI States

Latest PR Test (Base): ❌ Run #37548288530
Latest PR Test (Extra): ❌ Run #37548288241
Latest PR Test (AMD ROCm 10): ❌ Run #37548288542

@Oasis-Git Oasis-Git added run-ci CI: run the baseline test suite on this PR bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) labels Oct 6, 2026
@Oasis-Git

Copy link
Copy Markdown
Collaborator Author

@Oasis-Git

Copy link
Copy Markdown
Collaborator Author

Fix by #42768

@Oasis-Git
Oasis-Git merged commit 54b551d into sgl-project:main Oct 8, 2026
256 of 280 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) piecewise-cuda-graph run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant