vllm - 💡(How to fix) Fix [RFC] Add Top-nσ logit truncation example via custom logits processor [2 pull requests]

RFC: Add Top-nσ Logit Truncation Example via Custom Logits Processor

Reference: Tang et al., "Top-nσ: Not All Logits Are You Need", ACL 2025. arXiv:2411.07641

Motivation

Top-nσ is a logit-space dynamic truncation method proposed by Tang et al. (ACL 2025). Unlike probability-space filters such as top_p or min_p, it applies before softmax and keeps the candidate set less sensitive to temperature changes.

This issue proposes adding a research-oriented example implementation using vLLM's existing custom logits processor API. It is not a request to add a new built-in sampling parameter.

Scope

This proposal only adds an example file demonstrating how to implement top-nσ as a custom logits processor. It:

• Does not add a new SamplingParams field
• Does not change any existing sampling behavior
• Does not modify the core Sampler class

Algorithm

For each request row:

1. max_logit = logits.max()
2. std_logit = logits.std()
3. threshold = max_logit - n * std_logit
4. logits[logits < threshold] = -inf   (mask unlikely tokens)
5. Let existing vLLM sampler handle temperature, softmax, and sampling

The standard deviation directly measures the "peakiness" of the logit distribution:

• Sharp distribution (model is certain): small std → narrow threshold → fewer candidates
• Flat distribution (model is uncertain): large std → wide threshold → more candidates

This is based on the empirical observation (verified in Tang et al., 2025) that logit tails are approximately Gaussian, making the nσ threshold a principled percentile-based cutoff.

Usage (via custom logits processor)

from vllm import SamplingParams

sampling_params = SamplingParams(
    temperature=0.8,
    extra_args={"top_n_sigma": 2.0}
)

The custom logit processor reads top_n_sigma from extra_args and applies truncation before the standard sampling pipeline.

Safety / Edge Cases

Why This Is Useful

Existing methods don't use the spread of the logit distribution as a signal:

Top-nσ adapts to distribution spread, making it useful for:

• Models with highly variable logit distributions across prompts
• Scenarios where top_k/top_p over-filter or under-filter
• Research comparing logit-space vs probability-space filtering

A key property verified in our experiments: Top-nσ candidate count is invariant to temperature (0.1–10.0), while top_p varies by 19,665× and min_p by 20,000×. See experiment report for details.

Comparison Example

# Logit distribution: max=10.0, std=2.0

top_k=5:           Keeps exactly 5 tokens
top_p=0.9:         Keeps tokens summing to 90% cumulative probability
min_p=0.05:        Keeps tokens with prob >= 5% of max
top_n_sigma=2.0:   Keeps tokens with logit >= 10.0 - 2*2.0 = 6.0

Deliverables (if maintainers are interested)

1. Example file: examples/features/logits_processor/top_n_sigma.py
2. Unit tests: Verify masking, edge cases, and compatibility with other samplers
3. Benchmark sketch: Compare output quality vs top_p/min_p on diverse prompts
4. Documentation: One-paragraph note in sampling docs

Minimal Prototype (for initial feedback)

Before investing in a full PR, I'll share a standalone prototype script that demonstrates the algorithm. If maintainers find it interesting, I'll follow up with the full example, tests, and docs.

References: Tang et al., "Top-nσ: Not All Logits Are You Need", ACL 2025. arXiv:2411.07641