Skip to content

fix(sampler): bound top_k to the vocabulary before the int32 cast - #848

Closed
techsaurabh-394 wants to merge 1 commit into
FlashML-org:mainfrom
techsaurabh-394:fix/topk-int32-overflow
Closed

techsaurabh-394 wants to merge 1 commit into
FlashML-org:mainfrom
techsaurabh-394:fix/topk-int32-overflow

Conversation

@techsaurabh-394

Copy link
Copy Markdown

Fixes #789

Sampler.prepare() failed with an OverflowError when a request's top_k did not fit in int32 (for example 2**31 or 2**80), and one such row broke preparation of a whole mixed batch. A top_k at or above the vocabulary size is the same as no top-k, so it is now capped at vocab_size before the int32 tensor is built.

Tested: ran the issue's repro on CPU with make_device_tensor replaced by a plain CPU tensor and CUDA-only modules mocked (no GPU on my machine). top_k=2**31 and 2**80 no longer raise and prepare() returns top_k=None; top_k=5 still gives tensor([5]). Not run on a GPU; the repo's pytest suite was not run. No test added.

KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Nov 2, 2026
FlashML-org#850 adopted with fixups, FlashML-org#848/FlashML-org#849 superseded (FlashML-org#790/FlashML-org#788), FlashML-org#845 and FlashML-org#844
deferred (no NVFP4 qwen4exp GGUF; GPU driver down). FlashML-org#316 was closed upstream
without merge and is kept for the ROCm branches.

Assisted-by: Claude Opus 5.5
@jason-fxz

Copy link
Copy Markdown
Collaborator

Close as #790 fix it.

@jason-fxz jason-fxz closed this Nov 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Large top-k values overflow int32 in sampler preparation

2 participants