write-triton-sampling-kernel

Skill: Write a Triton Sampling Kernel

Purpose

Guide the agent through implementing a Triton kernel for LLM decode-time token sampling: take a [batch, vocab] logits tensor, apply per-request temperature, top-k, and top-p (nucleus) filtering, renormalize, and draw one token per request. This is the last hot kernel on every decode step — it runs once per generated token, so latency directly translates into tokens/second.

---

Use this when

• You need a sampling strategy that vLLM, SGLang, TGI, or FlashInfer do not expose (typical-p, mirostat, classifier-free guidance, fused repetition penalty, structured-generation logit bias, contrastive decoding).
• You need heterogeneous per-request sampling — each request has its own T, k, p, seed, and possibly its own logit-bias mask — and you want one fused kernel rather than N samplers.
• You are willing to special-case the greedy path (T == 0 or top-k == 1) to skip softmax and sort.
• Decode batch size is large enough (B >= 8) that one-program-per-request is worthwhile. For B == 1, a CPU-side argmax/multinomial is usually fine.

---

Do not use this when

• A vendor sampler covers your case. vLLM's Sampler and FlashInfer's top_k_top_p_sampling_from_probs are heavily tuned and handle edge cases (extremely peaked distributions, ties, deterministic argmax fallback). Re-implementing without a concrete reason is a likely source of subtle bias.
• You need only argmax. logits.argmax(-1) from PyTorch is competitive and avoids every numerical pitfall in this skill.
• The strategy requires global communication across the batch (beam search, speculative decoding verification). Those are not multinomial-per-request.
• You need provably uniform reproducibility across hardware. RNG semantics, cumsum reduction order, and sort tie-breaking are all platform-dependent.

---

Inputs the agent should gather first

Before writing any code, confirm:

1. Vocab size V. Typical: 32k (Llama-2), 128k (Llama-3), 256k (Gemma). Determines whether the row fits in one BLOCK or needs multi-block streaming.
2. Batch size B. Number of concurrent requests in the decode step. Each request maps to one program.
3. Per-request sampling params. Are T, k, p scalars (uniform) or tensors of shape [B] (heterogeneous)? Heterogeneous is the realistic case in continuous-batching servers.
4. Logits dtype. Almost always fp16 or bf16 from the LM head. Sampling internally promotes to fp32.
5. RNG source. Stateful Philox seed/offset (advanced once per decode step) or a precomputed [B] tensor of uniforms. Stateful is more flexible; precomputed is simpler and easier to test.
6. Greedy fallback policy. Is T == 0 legal? Is top_k == 1 legal? Both must short-circuit to argmax.
7. Logit bias / mask. Per-request additive bias (e.g., grammar-constrained decoding) is added to logits before temperature scaling.
8. Maximum top_k. A hard upper bound (e.g., K_MAX = 1024) lets you pick a sort strategy at compile time. Without a bound, you cannot size a fixed on-chip sort buffer.