From fb9dea85cb73c5e87c101d393aa75882f1ce8437 Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 23 Oct 2026 12:05:16 +0800 Subject: [PATCH] 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 --- .../tilelang_ops/neuron/fused_sample.py | 32 +++++++++++++++---- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py index a8c6a87..741c401 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py +++ b/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)