Cp chunkprefill upstream 20260610 latest

[cchh05]

Description

Related Issues

Change Type

• Bug fix
• New feature
• Performance improvement
• Refactor
• Documentation
• Test
• Build or CI

Pull Request Checklist

Thank you for contributing to xLLM. Before requesting review, please make sure the following items are complete.

PR Title and Commit Messages

• The PR title and each commit message follow the xLLM commit format: <type>: <subject>.

Allowed types: feat, bugfix, docs, test, refactor, chore, style, revert, perf, model, build, release.
The subject should use clear English, start with a verb, include at least 4 words, and end with ..

Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit or an equivalent command.
• I have installed the hooks with pre-commit install.
• I have run pre-commit run --all-files and fixed any reported issues.

If you are unsure how to set up pre-commit, see the pre-commit documentation.

Self Review

• I have self-reviewed the code according to .agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.
• I have rebased this PR onto the latest main branch.

Build and Test Coverage

• Tests have been added or updated as needed.
• CUDA: python setup.py build test has passed on a CUDA machine.
• NPU: python setup.py build test has passed on an NPU machine.
• MLU: python setup.py build test has passed on an MLU machine.

Reviewer Notes

[gemini-code-assist]

Code Review

This pull request introduces a proper fix for prefix cache asymmetry in disaggregated PD mode using a sender-driven transfer window, alongside adding decode-first scheduling, step isolation, and prefill chunk capping to the MixScheduler. It also refactors cluster linking to support batch operations and updates profiling logic. The review feedback highlights a logic error in MixScheduler when restoring the sequence budget under prefill chunk limits, and identifies several style guide violations, including the use of at::ScalarType instead of torch::ScalarType, missing braces in loops, using push_back instead of emplace_back, and using auto for primitive types.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Use torch::ScalarType instead of at::ScalarType to adhere to Section 7 of the repository style guide, which specifies using the torch:: namespace instead of at:: or c10:: wherever possible.

Suggested change

	::std::optional<at::ScalarType>{resolved_output_dtype});
	::std::optional<torch::ScalarType>{resolved_output_dtype});

References

1. Use torch:: namespace instead of at:: or c10:: wherever possible. Prefer the highest-level PyTorch C++ API. ^(link)

[gemini-code-assist]

There is a logic error in how consumed_in_phase2 and the restored remaining_seq_budget are calculated when max_prefill_chunks is enabled. Because remaining_seq_budget was capped to max_prefill_chunks before phase 2, subtracting the post-phase-2 remaining_seq_budget from the uncapped saved_seq_budget results in an incorrect (inflated) count of consumed sequences. This causes the restored remaining_seq_budget to be much lower than it should be (effectively leaving it at the capped remaining value). To fix this, calculate the actual consumed sequences by subtracting the post-phase-2 remaining_seq_budget from the capped budget value, and then subtract that actual consumed count from saved_seq_budget.

      if (max_prefill_chunks > 0) {
        const size_t phase2_initial_budget = std::min(saved_seq_budget, static_cast<size_t>(max_prefill_chunks));
        const size_t consumed_in_phase2 = phase2_initial_budget > remaining_seq_budget
                                              ? phase2_initial_budget - remaining_seq_budget
                                              : 0;
        remaining_seq_budget = saved_seq_budget > consumed_in_phase2
                                   ? saved_seq_budget - consumed_in_phase2
                                   : 0;
      }

[gemini-code-assist]

Adhere to the repository style guide by:

1. Always using braces {} for for loops, even for single-line bodies (Section 8).
2. Preferring emplace_back over push_back to avoid unnecessary copies (Section 9).

Suggested change

	for (auto& r : decode_queue) running_queue_.push_back(r);
	for (auto& r : prefill_queue) running_queue_.push_back(r);
	for (auto& r : decode_queue) {
	running_queue_.emplace_back(r);
	}
	for (auto& r : prefill_queue) {
	running_queue_.emplace_back(r);
	}

References

1. Always use braces {} with if, while, for, even for single-line bodies. Prefer emplace_back over push_back to construct elements in-place and avoid unnecessary copies. ^(link)

[gemini-code-assist]

Prefer emplace_back over push_back to construct elements in-place and avoid unnecessary copies, in accordance with Section 9 of the repository style guide.

Suggested change

	for (auto& req : running_queue_) {
	if (req->sequences()[0]->kv_state().kv_cache_tokens_num() > 0) {
	decode_queue.push_back(req);
	} else {
	prefill_queue.push_back(req);
	}
	}
	for (auto& req : running_queue_) {
	if (req->sequences()[0]->kv_state().kv_cache_tokens_num() > 0) {
	decode_queue.emplace_back(req);
	} else {
	prefill_queue.emplace_back(req);
	}
	}

References

1. Prefer emplace_back over push_back to construct elements in-place and avoid unnecessary copies. ^(link)

[gemini-code-assist]

Do not use auto for simple/primitive types like uint64_t or status codes, in accordance with Section 3 of the repository style guide. Please use the explicit type instead.

Suggested change

	auto ret =
	llm_data_dist_->UnlinkLlmClusters(clusters, rets, 1000, force_flag);
	uint64_t ret =
	llm_data_dist_->UnlinkLlmClusters(clusters, rets, 1000, force_flag);

References

1. Do not use auto for simple/primitive types. auto is acceptable for complex types (iterators, lambdas, template-deduced types) but not for int32_t, float, bool, std::string, etc. ^(link)