Inference: Cudagraph-aware admission gating in prefill scheduler - #4870
Conversation
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
The code changes LGTM but can we add some tests which exercise the code path that defers a request due to lack of a CUDA graph match? Basically we should confirm that even if a request is deferred: |
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| def _find_cg_chunk_size(self, max_chunk_tokens: int) -> Optional[int]: | ||
| """Return the largest chunk size <= max_chunk_tokens where batch matches a captured graph, | ||
| or None if no graph covers any chunk in the budget. | ||
|
|
||
| Walks the captured-CG list (sorted descending by token_count) and returns the first chunk | ||
| that falls within budget and produces an applicable batch_dim under the engine's matching | ||
| mode (strict for hybrid models). Callers must explicitly handle the None case by deferring | ||
| the admission rather than scheduling eagerly. | ||
| """ | ||
| active_tok = self.context.active_token_count | ||
| active_p = self.context.num_prefill_requests | ||
| active_d = self.context.num_decode_requests | ||
| strict = self.context.is_hybrid_model | ||
|
|
||
| for cg in self.context.cuda_graph_batch_dimensions_list: | ||
| chunk = cg.token_count - active_tok | ||
| if chunk < 1: | ||
| continue | ||
| if chunk > max_chunk_tokens: | ||
| continue | ||
| candidate = InferenceBatchDimensions( | ||
| token_count=cg.token_count, | ||
| prefill_req_count=active_p + 1, | ||
| decode_req_count=active_d, | ||
| ) | ||
| if cg.is_applicable_for_batch_dim(candidate, strict=strict): | ||
| return chunk | ||
|
|
||
| return None |
There was a problem hiding this comment.
[SUGGESTION Simplification]: _find_cg_chunk_size inlines CG matching by calling cg.is_applicable_for_batch_dim(candidate) on each CG individually, while _cg_admission_check delegates to CUDAGraphBatchDimensionBuilder.match_graph_config() which wraps the same is_applicable_for_batch_dim in a filter-then-min-select loop.
Both are correct today, but they form two independent matching code paths against the same CG list. If match_graph_config gains new filtering constraints in the future (beyond the EP sync it already does), this method won't inherit them.
This is fine as-is — _find_cg_chunk_size fundamentally can't reuse match_graph_config without calling it O(n) times (once per candidate chunk), since the token count is the search variable. Just flagging the coupling for future awareness.
Minor readability note: candidate.token_count is always set to cg.token_count, so the token-dimension check inside is_applicable_for_batch_dim (self.token_count >= real.token_count) is tautologically True. The method is effectively filtering on P/D compatibility only. A one-line comment would help a future reader who wonders why the candidate token_count mirrors the CG's.
| # Flash-attn guard: if this chunk would leave exactly 1 token for the | ||
| # final chunk, reduce by 1 (or defer if we only have 1 token of budget). | ||
| # See https://github.com/Dao-AILab/flash-attention/issues/1537 | ||
| if remaining_len - prefill_chunk_length == 1: | ||
| if prefill_chunk_length > 1: | ||
| prefill_chunk_length -= 1 | ||
| else: | ||
| can_schedule = False | ||
| break |
There was a problem hiding this comment.
[SUGGESTION Clarity]: The flash-attn guard runs after CG snapping, so prefill_chunk_length -= 1 can shift the batch 1 token below the snapped CG boundary. This is safe because is_applicable_for_batch_dim uses >= matching (a CG with token_count=T covers any batch with token_count <= T), so the same CG still applies.
A brief inline comment noting this non-obvious invariant ("-1 stays within the matched CG's capacity because matching is >=, not exact") would help a future reader who sees the adjustment and worries about CG misalignment.
There was a problem hiding this comment.
Review Summary
CRITICAL: 0 | IMPORTANT: 0 | SUGGESTION: 2
Thorough review of the CG-aware admission gating logic, the schedule_chunked_prefill refactoring, and the new helper methods.
What I verified
- Refactoring correctness: The
schedule_chunked_prefillmerge oftoken_fully_can_be_added/token_partially_can_be_addedinto a unifiedmax_chunk = min(remaining_len, token_budget)path is behavior-preserving. All original admission/deferral/break semantics are maintained. - CG gating logic:
_find_cg_chunk_sizecorrectly searches descending-by-token CGs, checking P/D compatibility viais_applicable_for_batch_dim, returning the largest CG-aligned chunk within budget._cg_admission_checkcorrectly delegates tomatch_graph_configwithmatch_ep_token_counts=Falseto avoid per-attempt NCCL all-reduce. - Continuing chunked prefill bypass: Correctly skips CG gating for mid-flight chunked prefills (
is_continuing_chunked_prefill), preventing deadlock. - Flash-attn guard interaction: The
-1adjustment after CG snapping is safe because CG matching uses>=(not exact), so the matched CG still covers the adjusted batch. - Starvation warning:
cg_wait_iterstracks CG-specific deferrals, resets on CG-matched admission, and correctly does not count flash-attn deferrals. - Backward compatibility: Gating is strictly opt-in via
cuda_graph_all_prefills && use_cuda_graphs_for_non_decode_steps && non-empty CG list. Existing tests/configs are unaffected. The_StubEnginefix intest_dynamic_prefix_caching.pycorrectly adds the new attribute. - Non-chunked prefill path: CG check correctly constructs the candidate with
active_tokens + remaining_prompt_tokensandprefill_count + 1, and thebreakon miss is consistent with the existing FIFO scheduling invariant. - New
cg_wait_itersfield: Default-valued dataclass field onDynamicInferenceRequest, no serialization or compatibility concerns. - Test coverage: Comprehensive — activation conditions, chunk-size snapping, strict/non-strict matching, warning thresholds, deferral/resume flows, and the "no eager fallback" invariant.
Suggestions posted
- Note the dual CG-matching code paths (
_find_cg_chunk_sizeinline vs_cg_admission_checkviamatch_graph_config) for future maintenance awareness. - A clarity comment on the flash-attn guard's safety after CG snapping, documenting the
>=matching invariant.
Clean implementation. LGTM.
|
/ok to test 06a9508 |
|
/ok to test d940d1a |
|
/ok to test efa61b3 |
|
/ok to test 4cd2c33 |
Signed-off-by: Helen Ngo <helenn@nvidia.com>
Signed-off-by: Helen Ngo <helenn@nvidia.com>
Signed-off-by: Helen Ngo <helenn@nvidia.com>
Signed-off-by: Helen Ngo <helenn@nvidia.com>
Signed-off-by: Helen Ngo <helenn@nvidia.com>
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/27453419722 |
6 participants
What does this PR do ?
Presently,
schedule_chunked_prefillandschedule_non_chunked_prefilladmit requests based on the available token/request budget only without consideration for whether the resulting batch shape will match any captured CG. When a request pushes the batch into a shape the captured set doesn't cover, the engine silently falls back to eager mode. While Transformer models can absorb extra decodes into prefill slots, for hybrid models, the matcher requires captured_decode_req_count >= real_decode_req_count (strict mode), so we fallback to eager more often.Now, when
cuda_graph_all_prefillsis on,_find_cg_chunk_size(max_chunk_tokens)traverses the list of cudagraphs to find the best-fit cudagraph for(active_token_count + chunk_size, num_prefill_requests + 1, num_decode_requests).schedule_chunked_prefillsnapsprefill_chunk_lengthto the largest CG-aligned boundary in the token budget, and if no CG covers the resulting shape, the request is deferred.This is designed without an eager fallback: if there is no cudagraph, the deferred request waits until the present decode finishes, when the next prefill will be automatically admitted. The worst case of this is
max_sequence_length. This avoids the one-off version of forcing an eager step which then continues to trigger eager due to the imbalance, and ensures that we run every step cudagraphed.#3509 needs to be merged first to avoid unnecessary starvation for hybrid inference. These two PRs together ensure that no steps will run eager.
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.
For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.