Skip to content

fix(engine): preserve greedy sampling in mixed batches - #471

Open
taking-lying-flat wants to merge 1 commit into
FlashML-org:mainfrom
taking-lying-flat:fix/greedy-mixed-batch
Open

taking-lying-flat wants to merge 1 commit into
FlashML-org:mainfrom
taking-lying-flat:fix/greedy-mixed-batch

Conversation

@taking-lying-flat

@taking-lying-flat taking-lying-flat commented Sep 14, 2026

Copy link
Copy Markdown
Contributor

A greedy request with fixed logits [1.0, 1.0, 0.0] returns token 0 alone, but can return token 1 when batched with a random request: the mixed path approximates greedy with temperature 1e-6, leaving the tied maxima equally likely. Also, temperature=0, top_p=0.9 unnecessarily enters probability sampling.

Keep a per-request greedy_mask and explicitly select argmax for greedy rows. Treat temperature <= 0 or top_k == 1 as greedy regardless of top_p, and use neutral backend parameters for those rows. The all-greedy fast path is preserved. Mixed batches still run probability sampling across all rows; skipping that computation is outside this correctness fix.

Related: #307. This fixes sampling-policy differences for identical logits; it does not establish or resolve the cause of the serial model-level divergence reported there.

Validation:

  • Ran local deterministic CPU checks covering tied maxima, standalone versus mixed batches, FP32/FP16/BF16, and all four top-k/top-p modes. Before the fix, the mixed regression returned [1, 2, 2] instead of [0, 2, 1], and zero-temperature requests with top_p < 1 entered probability sampling. The validation file is retained locally and is not included in this PR.
  • On upstream base e0886cc plus this patch: 73 passed, 4 skipped, including real Triton CUDA sampling. The four FlashInfer cases were skipped because FlashInfer is not installed. git diff --check passed.
  • Hardware: NVIDIA RTX A1000 Laptop GPU; Intel Core i7-12800HX; NVIDIA driver 595.84; PyTorch 2.11.0+cu130. Checkpoint: N/A, tests use fixed logits without loading a model.

Exact local validation command, from the repository root, using the untracked test file:

env PYTHONPATH=python /tmp/freetoken-glm-shared-venv/bin/python -m pytest tests/engine/test_sample.py -q -rs

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.

1 participant