Neuron sampler: skip the top-p/top-k bisection when no row filters (#27)

The bisection is 32 rounds over the vocabulary: ~0.5 s a decode step at 129280
even compiled, and the default parameters (top_p 1, top_k -1) filter nothing.
Such batches take the plain temperature log-softmax: 529.6 -> 1.3 ms on one
core, same tokens, logprobs within 1e-9.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py

@@ -116,7 +116,16 @@ def _logprobs_and_greedy(logits, temperatures, top_p, top_k):  # noqa: ANN001, A
     return logprobs, logprobs.argmax(dim=1).to(torch.int32)
 
 
-_COMPILED = None
+def _plain_logprobs_and_greedy(logits, temperatures):  # noqa: ANN001, ANN202
+    """``compute_logprobs`` for rows with no top-p / top-k filter: the temperature log-softmax alone."""
+    x = logits.float()
+    inv_t = (1.0 / temperatures.float()).unsqueeze(1)
+    scaled = (x - x.amax(dim=1, keepdim=True)) * inv_t
+    logprobs = scaled - torch.log(torch.exp(scaled).sum(dim=1, keepdim=True))
+    return logprobs, logprobs.argmax(dim=1).to(torch.int32)
+
+
+_COMPILED: dict = {}
 
 
 def _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size: int):  # noqa: ANN001, ANN202
@@ -135,13 +144,24 @@ def _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_si
     from deepinfer.neuron_ops import neuron_compile  # noqa: F401  (registers the "deepinfer_neuron" backend)
     from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda
 
-    if _COMPILED is None:
-        _COMPILED = torch.compile(_logprobs_and_greedy, backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, dynamic=False)
-    with hide_cuda():
-        logprobs, greedy = _COMPILED(logits, temperatures, top_p, top_k)
     from deepinfer.sampling_params import SAMPLING_EPS
 
-    if bool((temperatures.cpu() <= SAMPLING_EPS).all()):  # host sync on the [rows] temperatures
+    def compiled(fn):  # noqa: ANN001, ANN202
+        if fn not in _COMPILED:
+            _COMPILED[fn] = torch.compile(fn, backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, dynamic=False)
+        return _COMPILED[fn]
+
+    # Host copies of the [rows] parameters (they come from the host; the sync waits for the logits anyway).
+    t_h, p_h, k_h = temperatures.cpu(), top_p.cpu(), top_k.cpu()
+    # The top-p / top-k bisection is 32 rounds over the vocabulary (~0.5 s at 129280 even compiled): only rows that
+    # filter need it, and the default parameters filter nothing.
+    filtering = bool(((p_h < 1.0) | ((k_h >= 0) & (k_h < logits.shape[1]))).any())
+    with hide_cuda():
+        if filtering:
+            logprobs, greedy = compiled(_logprobs_and_greedy)(logits, temperatures, top_p, top_k)
+        else:
+            logprobs, greedy = compiled(_plain_logprobs_and_greedy)(logits, temperatures)
+    if bool((t_h <= SAMPLING_EPS).all()):
         return logprobs, greedy
     return logprobs, two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size)