Skip to content

fix(sampling): apply presence and frequency penalties - #527

Open
taking-lying-flat wants to merge 2 commits into
FlashML-org:mainfrom
taking-lying-flat:fix/sampling-penalties
Open

taking-lying-flat wants to merge 2 commits into
FlashML-org:mainfrom
taking-lying-flat:fix/sampling-penalties

Conversation

@taking-lying-flat

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

Copy link
Copy Markdown
Contributor

Chat and text completion requests accept presence_penalty and frequency_penalty, but currently drop both values before sampling. This change carries them through SamplingParams and applies logit -= frequency * count + presence * (count > 0) before greedy selection or temperature/top-k/top-p sampling.

Counts belong to each request and include generated tokens only. They are updated on the sampling stream, so overlapping execution sees the latest sampled token before its CPU acknowledgement arrives. Intermediate prefill chunks do not contribute counts; requests with zero penalties allocate no counter. The HTTP models reject non-finite penalty values to prevent NaN logits.

Validation: 181 checks passed: 73 existing API, wire-format, accounting and scheduler checks, plus 108 local HTTP/CUDA checks for parameter forwarding, positive/negative penalties, final-token handling, mixed greedy batches, changing batch membership, stream ordering, non-finite inputs and zero-penalty parity with the previous sampler. The post-submission checks also cover vocabulary sizes 128,256 and 248,320, batch sizes 1/8/32, and FP16/BF16/FP32 logits. This branch contains production code only; the additional checks were run from files outside the repository.

Environment: Intel Core i7-12800HX, NVIDIA RTX A1000 Laptop GPU (4 GiB), driver 595.84, PyTorch 2.13.0+cu130, Triton 3.7.1. GPU checks used synthetic logits; no model checkpoint or full-model/multi-GPU inference was run. The installed Torch/Triton versions differ from the repository's pinned ranges.

Exact validation commands, run from the repository root (163 checks, then 18 additional checks):

/tmp/freetoken-497-venv/bin/python -m pytest -q -o pythonpath=python \
  /tmp/freetoken-penalty-artifacts/test_penalty_review.py \
  /tmp/freetoken-penalty-artifacts/test_sampling_penalties.py \
  tests/server/test_openai_api.py \
  tests/server/test_effort_dialect.py \
  tests/server/test_message_wire.py \
  tests/server/test_generation_accounting.py \
  tests/scheduler/test_req_append_host.py \
  tests/scheduler/test_scheduler_chunked_prefill.py \
  --maxfail=2 --tb=short -o faulthandler_timeout=30

/tmp/freetoken-497-venv/bin/python -m pytest -q \
  /tmp/freetoken-penalty-artifacts/test_post_pr_review.py \
  --maxfail=1 --tb=short -o faulthandler_timeout=30

JUNQINGV587 added a commit to JUNQINGV587/FreeToken that referenced this pull request Sep 23, 2026
…ty application already present)

Upstream FlashML-org#527 carries two commits: the substantive one applies presence/frequency
penalties end to end, the other constrains both to finite values at the API boundary.

This fork already implements the penalty path (Sampler.apply_penalties + the per-row
output_token_counts accumulation in engine/sample.py, wired through core.py and both
OpenAI entrypoints), so the substantive commit is a duplicate implementation. Kept ours
in every conflict hunk -- it is a strict superset: upstream's hunks carry no identifier
ours lacks, and ours additionally threads the logprobs plumbing (PR FlashML-org#224 lineage).
What this merge actually adds is the API validation: presence_penalty /
frequency_penalty now fail closed on inf/nan (4 declarations).

Verified: 1639 passed, 34 skipped across tests/{attention,checkpoint,daemon,kvcache,
layers,scheduler,server,tokenizer,utils}.

This branch has not been deployed

No deployments
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