feat: add beam search torch support to mlu. by phantomlei3 · Pull Request #1835 · xLLM-AI/xllm

phantomlei3 · 2026-06-27T08:35:54Z

Description

Related Issues

Change Type

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 support for MLU device-side beam search, implementing the beam_search kernel, integrating it into the BeamSearcher framework, and updating SamplingParameters to propagate and serialize beam_width. Feedback on the changes highlights a critical issue in the MLU kernel where integer division (idx / topk) could cause a runtime crash in PyTorch 1.6+ due to true division, recommending idx.div(topk, "trunc") instead. Additionally, several style guide violations were identified in the test suite, specifically regarding the use of PascalCase for function names and using push_back instead of emplace_back.

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 · 2026-06-27T08:37:54Z

+  // MLU torch::topk returns int64 indices.
+  torch::Tensor idx = std::get<1>(topk_result);  // [num_group, beam_width]
+
+  torch::Tensor src_local = idx / topk;  // source beam within the group


In PyTorch 1.6+, the / operator on integer tensors performs true division and returns a float tensor. Since idx is an integer tensor, idx / topk will produce a float tensor for src_local. This will cause a runtime crash when src_local is used as an index in gather (via row_index), which strictly expects integer types (kInt64 or kInt32).

To prevent this and ensure truncating integer division, use idx.div(topk, "trunc") or torch::div(idx, topk, "trunc").

Suggested change

torch::Tensor src_local = idx / topk; // source beam within the group

torch::Tensor src_local = idx.div(topk, "trunc"); // source beam within the group

gemini-code-assist · 2026-06-27T08:37:54Z

+  int64_t topk;
+};
+
+std::string PrintBeamCase(const testing::TestParamInfo<BeamCase>& info) {


According to the repository style guide (Section 1, Naming Conventions), function names must use snake_case. Please rename PrintBeamCase to print_beam_case.

Suggested change

std::string PrintBeamCase(const testing::TestParamInfo<BeamCase>& info) {

std::string print_beam_case(const testing::TestParamInfo<BeamCase>& info) {

References

Function names must use snake_case. ^(link)

gemini-code-assist · 2026-06-27T08:37:54Z

+  }
+}
+
+std::vector<BeamCase> GenerateCases() {


According to the repository style guide (Section 1, Naming Conventions), function names must use snake_case. Please rename GenerateCases to generate_cases.

Suggested change

std::vector<BeamCase> GenerateCases() {

std::vector<BeamCase> generate_cases() {

References

Function names must use snake_case. ^(link)

gemini-code-assist · 2026-06-27T08:37:54Z

+      int64_t seq = g * beam_width + b;
+      for (int64_t c = 0; c < topk; ++c) {
+        float value = lp[seq] + tlp[seq][c];
+        candidates.push_back({value, b, c});


According to the repository style guide (Section 9, STL Best Practices), we should prefer emplace_back over push_back to construct elements in-place and avoid unnecessary copies.

Suggested change

candidates.push_back({value, b, c});

candidates.emplace_back(Candidate{value, b, c});

References

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

gemini-code-assist · 2026-06-27T08:37:54Z

+      ref.src_seq_idxes.push_back(static_cast<int32_t>(src_seq));
+      ref.out_tokens.push_back(static_cast<int32_t>(tt[src_seq][cand.col]));
+      ref.out_logprobs.push_back(cand.value);


According to the repository style guide (Section 9, STL Best Practices), we should prefer emplace_back over push_back to construct elements in-place and avoid unnecessary copies.

Suggested change

ref.src_seq_idxes.push_back(static_cast<int32_t>(src_seq));

ref.out_tokens.push_back(static_cast<int32_t>(tt[src_seq][cand.col]));

ref.out_logprobs.push_back(cand.value);

ref.src_seq_idxes.emplace_back(static_cast<int32_t>(src_seq));

ref.out_tokens.emplace_back(static_cast<int32_t>(tt[src_seq][cand.col]));

ref.out_logprobs.emplace_back(cand.value);

References

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

gemini-code-assist · 2026-06-27T08:37:54Z

+INSTANTIATE_TEST_SUITE_P(BeamSearchMlu,
+                         BeamSearchMluTest,
+                         testing::ValuesIn(GenerateCases()),
+                         PrintBeamCase);


Please update the instantiation to use the renamed generate_cases and print_beam_case functions to adhere to the naming conventions.

Suggested change

INSTANTIATE_TEST_SUITE_P(BeamSearchMlu,

BeamSearchMluTest,

testing::ValuesIn(GenerateCases()),

PrintBeamCase);

INSTANTIATE_TEST_SUITE_P(BeamSearchMlu,

BeamSearchMluTest,

testing::ValuesIn(generate_cases()),

print_beam_case);

References

Function names must use snake_case. ^(link)

Add an int32_t beam_width field to the device-side SamplingParameters and plumb it through init aggregation, to(device) transfer, microbatch concat (with a consistency CHECK), and the shared-memory write/read round-trip so the MLU beam-search device kernel can recover the per-group beam width. NPU/CUDA/CPU paths are unaffected; non-beam requests keep beam_width=0. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Implement the MLU device path for general LLM beam search so beam_width>1 requests run through the same engine pipeline as NPU instead of hitting NOT_IMPLEMENTED(). Add kernels/mlu/beam_search.cpp (pure torch ops: combined logprobs -> per-group topk -> reconstruct src_seq_idxes/out_tokens/ out_logprobs), wire a USE_MLU branch in BeamSearcher::forward with a beam_width default param (NPU/rec callers unchanged) and divide-by-zero guards, pass sampling_params.beam_width from llm_worker_impl, and add an MLU cc_test that strict-compares against a torch-CPU reference. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Integer tensor operator/ performs true division in torch2.10 and yields a float tensor, which is invalid as a gather index and crashes the MLU beam_search device path. Use torch::div with floor rounding to keep the source beam index int64. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

apply_kv_block_swaps() was a no-op on MLU: its platform guards only covered USE_NPU/USE_CUDA/USE_DCU, so MLU fell through to `#else return` and the torch swap path was compiled out. Beam-search copy-on-write (block_manager_pool.cpp) allocates a fresh KV block for each diverging beam and relies on this swap to populate it with the source block's KV. With the swap dropped, every diverging beam read stale/uninitialized pool memory, whose physical-block identity is nondeterministic, making beam search (beam_width>=2) nondeterministic and degenerate even for a single sequential request at temperature 0. greedy/beam_width=1 never trigger copy-on-write and were unaffected. Add USE_MLU to both guards so MLU runs the existing device-agnostic torch swap path (KVCache::swap_blocks -> index_select/index_copy_). Verified on Qwen3-0.6B (MLU): beam_width 2/4 go from 7/8 distinct degenerate outputs to fully deterministic, coherent results across repeated single-request runs; beam-search correctness checks pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

The device beam-search kernel (enable_beam_search_kernel=true) assumes the batch is num_groups blocks of beam_width contiguous, fully-expanded, decoding rows sharing one width. Continuous batching violates this (mid-flight beam expansion, per-beam early finish, mixed widths, prefill/decode mix, non-beam sequences), so group*beam_width+i indexed into a foreign request. Generalize the existing partial-finished guard into beam_kernel_batch_eligible() and apply it in both prepare_forward_input overloads (the MLU executor path used the unguarded one): ineligible batches drop acc_logprob and fall back to the verified host beam path. Also fix an off-by-one CHECK_LE->CHECK_LT on src_seq_idx. Verified on Qwen3-0.6B/MLU: kernel-server output matches host-server output for both isolated and concurrent mixed-width workloads; test_beam_search 48/0. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

phantomlei3 changed the title ~~Feat/beam search torch~~ feat: add beam search torch support to mlu. Jun 27, 2026

gemini-code-assist Bot reviewed Jun 27, 2026

View reviewed changes

phantomlei3 and others added 2 commits July 7, 2026 15:41

phantomlei3 force-pushed the feat/beam-search-torch branch from d8334cf to b6cac38 Compare July 7, 2026 07:42

phantomlei3 and others added 2 commits July 7, 2026 16:24

phantomlei3 force-pushed the feat/beam-search-torch branch from a744be2 to d885288 Compare July 8, 2026 06:39

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: add beam search torch support to mlu.#1835

feat: add beam search torch support to mlu.#1835
phantomlei3 wants to merge 5 commits into
xLLM-AI:mainfrom
phantomlei3:feat/beam-search-torch

phantomlei3 commented Jun 27, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

	torch::Tensor src_local = idx / topk; // source beam within the group
	torch::Tensor src_local = idx.div(topk, "trunc"); // source beam within the group

	std::string PrintBeamCase(const testing::TestParamInfo<BeamCase>& info) {
	std::string print_beam_case(const testing::TestParamInfo<BeamCase>& info) {

	std::vector<BeamCase> GenerateCases() {
	std::vector<BeamCase> generate_cases() {

	candidates.push_back({value, b, c});
	candidates.emplace_back(Candidate{value, b, c});

Uh oh!

Conversation

phantomlei3 commented Jun 27, 2026

Description

Related Issues

Change Type

Pull Request Checklist

PR Title and Commit Messages

Pre-commit Checks

Self Review

Build and Test Coverage

Reviewer Notes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant