Remap to max supported topk instead of assert

[tchedaTT]

Ticket

#35661

Problem description

Currently this code crashes in vllm v0 which does not have a workaround for it, and v1 has a workaround that should be removed (re-mapping should be done here, so that when/if capability is added down the line, the code only needs to be changed in one place

This is covered by another problem on main, but shows up here https://github.com/tenstorrent/tt-metal/actions/runs/21754091798/job/62760237173

What's changed

map to 32 instead of asserting

Checklist

• 
• 
• 
• New/Existing tests provide coverage for changes

Model tests

If your changes cover model-related code, you should run tests corresponding to affected models and platforms (Single card, T3K, Galaxy). "Choose your pipeline" workflows facilitate running multiple kinds of tests in a single run. Each offers models-mandatory and models-extended presets.
The former includes a minimal set of tests, to be run always. The latter extends that with additional ones - use your best judgement in deciding which is the most appropriate for your PR.

• 

  • models-mandatory preset (runs: Device perf regressions and Frequent model and ttnn tests)
  • models-extended preset (runs: the mandatory tests, plus Demo and Model perf tests)
  • other selection - specify runs

• 

  • models-mandatory preset (runs: Unit tests)
  • models-extended preset (runs: the mandatory tests, plus Demo and Model perf tests)
  • other selection - specify runs

• 

  • models-mandatory preset (runs: Quick tests)
  • models-extended preset (runs: the mandatory tests, plus Demo and Model perf tests)
  • other selection - specify runs

[tchedaTT]

/codeowners ping

[tenstorrent-github-bot]

CodeOwners Group Analysis

This PR requires approval from one member of each of the following groups:

Summary: 1 pending groups, 0 approved groups

Group Information:

• ⏳ models/common (Group) - Members: Gongyu Wang, Mark O'Connor, Miguel Tairum, Utku Aydonat | Pending approval

  📁 Files owned by this group (1 files)

  • models/common/sampling/generator.py

Note: At least one approval from each group is sufficient.

[tenstorrent-github-bot]

Hi Gongyu Wang (@gwangTT), Miguel Tairum (@mtairum), this PR Remap to max supported topk instead of assert by Tomasz Cheda (@tchedaTT) needs your approval/review to merge this.

[tchedaTT]

cc @skhorasganiTT

[Copilot]

Pull request overview

This PR prevents crashes/undefined behavior in on-device sampling by clamping top_k to the currently supported maximum (32) rather than asserting, aligning with the expectation that callers (e.g., vLLM v0) may pass top_k > 32.

Changes:

• Replace the assert top_k <= 32 check with logic that re-maps top_k > 32 to 32.
• Update the inline comment to reflect the new top_k handling contract.

[Copilot]

Clamping top_k values >32 is now silent. Since this used to assert, callers may not realize their request is being changed, which can make sampling behavior hard to debug. Consider emitting a warning when top_k is re-mapped (e.g., only when the value is modified) so the change is discoverable without crashing.

[tchedaTT]

other vllm backends don't emit warnings when remapping, and based on discussions with @skhorasganiTT we don't want to either

[Copilot]

This block duplicates clamping logic and repeats the magic constant 32 twice, even though the file already has a clamp() helper and uses named constants for top-p. Consider introducing a TOP_K_MAX = 32 constant and using clamp() (or a single min/max) so the limit is defined once and easier to change later.

[gwangTT]

The code change looks fine

Please run some CI pipelines to check for regressions. I suggest T3K demo and Galaxy demo pipelines

[tchedaTT]

vllm ci: https://github.com/tenstorrent/tt-metal/actions/runs/21759446790

[tchedaTT]

skipping galaxy to save resources as this path is already tested in the vllm ci linked above

[tchedaTT]

T3k failures look unrelated
main: T3K demo https://github.com/tenstorrent/tt-metal/actions/runs/21759598642
This branch: https://github.com/tenstorrent/tt-metal/actions/runs/21848999695