Skip to content

[Transformers backend] Support batch-sharded sampling - #66545

Merged
DarkLight1337 merged 2 commits into
vllm-project:mainfrom
hmellor:transformers-backend-local-argmax
Nov 9, 2026
Merged

DarkLight1337 merged 2 commits into
vllm-project:mainfrom
hmellor:transformers-backend-local-argmax

Conversation

@hmellor

@hmellor hmellor commented Nov 9, 2026 •

Copy link
Copy Markdown
Member

Overview

Add compute_logits_local to the Transformers modeling backend's causal models, so batch-sharded sampling (enable_batch_sharded_sampling) works for them as it does for native Llama. Part of serving Llama with the backend (#66547).

Claims

  • Backend causal models no longer log "Disabling batch-sharded sampling: ... does not implement compute_logits_local" and fall back to replicated sampling.
  • Sharded and unsharded backend outputs, logprobs and prompt logprobs match.

Validation

B200, Llama-3.2-1B-Instruct, TP=2:

pytest -v "tests/v1/e2e/general/test_sharded_sampling.py::test_sharded_sampling_outputs_match[llama-transformers]" "tests/v1/e2e/general/test_sharded_sampling.py::test_sharded_sampling_outputs_match[llama]"

2 passed. [llama-transformers] fails on main (AssertionError: [False, False]: the sharded sampler never ran).


Pull Request Checklist
  • I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

  • AI assistance was used during the creation of this PR.

  • Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

  • Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

  • Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

  • Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

Implement `compute_logits_local` in the backend's `CausalMixin` so that
`enable_batch_sharded_sampling` is no longer silently disabled for
Transformers modeling backend models, matching native Llama.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@DarkLight1337 DarkLight1337 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tbh we should have defined this method in the model interface but that's outside the scope of this PR. cc @WoosukKwon

@DarkLight1337

Copy link
Copy Markdown
Member

/ci run

@DarkLight1337
DarkLight1337 enabled auto-merge (squash) November 9, 2026 11:27
@github-actions github-actions Bot added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 9, 2026
@github-actions

github-actions Bot commented Nov 9, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #101090 for commit 07e6ed566325.

ℹ️ This PR is 48 commits behind upstream main. Merged upstream main into this branch as 07e6ed566325 before starting CI. Pull the branch before pushing more changes.

@DarkLight1337
DarkLight1337 merged commit 7c35455 into vllm-project:main Nov 9, 2026
125 checks passed
@github-project-automation github-project-automation Bot moved this from In review to Shipped in Transformers modeling backend Nov 9, 2026
@hmellor
hmellor deleted the transformers-backend-local-argmax branch November 9, 2026 12:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready ONLY add when PR is ready to merge/full CI is needed

Projects

Development

Successfully merging this pull request may close these issues.

3 participants