feat(examples): exercise and audit top-k/top-p sampling end to end

[aleozlx]

📌 Description

Closes a structural hole in the e2e LLM example harness (#4083, which this
stacks on — the diff includes #4083's commits until it lands).

Every stage of the harness ran greedy: --temperature defaults to 0.0 and
sample_tokens() short-circuits to logits.argmax(), so
flashinfer.sampling.top_k_top_p_sampling_from_logits was never invoked
and the sampling and topk JIT modules were never even compiled.
Measured on B200: the greedy path builds 4 modules, the sampling path builds
2 more. Nothing in the suite had ever touched them.

Running the kernel is not enough — a wrong answer has to be caught. So
--check-sampling audits every sampled token against a plain-torch reference
support (top-k ∩ nucleus, from the same temperature-scaled logits). Both
boundaries close over ties and the eps slack is applied on the permissive
side only, so the mask is a superset of what any correct implementation may
keep and a violation is unambiguous. This mirrors how
tests/utils/test_sampling.py asserts membership.

Membership alone has two blind spots, both covered:

• a vacuous check — with top_k=0/top_p=1.0 the admissible set is the
  whole vocabulary and can never fail. So the mean admissible set size is
  reported and bounded. Measured: 4.3 of 151936, i.e. genuinely selective.
• a sampler degenerated to argmax — always inside the support, so
  membership is blind to it. So divergences from greedy are counted against
  how many the filtered distribution predicts. Measured: 15 observed vs
  16.58 expected.

Also audited: ids in range, seeded replay on identical logits, and the same
membership check with per-request tensor top_k/top_p — a different
kernel path, since a scalar top_k on a large vocab only ever exercises the
top_k_first fast path.

smoke_test.py gains a sampling run pair alongside the greedy pair, asserting
all of the above plus seeded cross-process determinism and sampling/topk
JIT cache reuse. Greedy remains the default everywhere and its existing
assertions are untouched; --skip-sampling restores the old behavior.

🔍 Related Issues

Stacks on #4083. Relevant to #4189 (__launch_bounds__ gated to SM107): that
change makes B200 sampling codegen match what main already ships (main has
zero launch_bounds in sampling.cuh), so it is not a new Blackwell path —
but this is the first e2e coverage of sampling either way.

🧪 Tests

Validated on B200 (sm100a) against flashinfer 0.6.16rc3 (the tree that
carries #4189's gate — verified present in the tested checkout):

Check	Result
sample_violations / 72 draws	0
sample_out_of_range	0
sample_allowed_mean	4.3 of 151936 (bar ≤ 2·top_k)
divergences vs expected	15 vs 16.58
sample_replay_match	1
sample_perreq_violations	0
smoke_test.py (4 runs)	PASS — run 3 built topk+sampling, run 4 built 0
reference_check.py --arch dense	unregressed

Note on the cross-process sampled-determinism bar: it is a hard assertion
rather than gated on a logits fingerprint, because the harness already
asserts greedy determinism as a hard bar. If model numerics ever stopped being
bit-reproducible, greedy determinism would fail first and more legibly.

✅ Checklist

• Passes pre-commit run -a
• Validated on real hardware (B200)
• Greedy default path unchanged
• Sampling under TEP/torchrun (non-zero ranks send stdout to devnull, so a
  violation there surfaces only as a nonzero exit code)

AI-assisted (Claude Code).

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 1e6c06c8-bf08-4c63-ab2a-776ea9bd6313

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}