Batch detokenization across positions in detokenize_top_logprobs_tokens

[Aphoh]

Motivation

TokenizerManager.detokenize_top_logprobs_tokens previously iterated over positions and called tokenizer.batch_decode once per position, leading to N tokenizer round-trips for N positions. The original code carried a TODO to batch all top-k tokens across all positions; this PR addresses it.

Modifications

• Rewrote detokenize_top_logprobs_tokens in python/sglang/srt/managers/tokenizer_manager.py to flatten every position's top-k token ids into a single decode call over one-token sequences, then slice the decoded texts back to per-position lists. When decode_to_text=False, the tokenizer is skipped entirely.
• Added _batch_decode_token_ids so both regular logprob detokenization and top-logprob detokenization use backend_tokenizer.decode_batch(..., skip_special_tokens=False) when available, with a fallback to tokenizer.batch_decode.
• Added TestDetokenizeTopLogprobsTokens coverage for empty inputs, mixed empty/non-empty positions, the single decode-call path, backend decode preference, and equivalence against a per-position reference implementation.

Performance / Repro

Helper-level benchmark only, no server e2e path. This compares the old per-position helper against the PR implementation on a long output sequence:

• Machine: macOS 26.3 arm64
• Python: 3.12.11
• transformers: 5.5.4
• tokenizers: 0.22.2
• Tokenizer: gpt2
• Input shape: 100,000 output positions, top_k=5, so 500,000 top-logprob token ids decoded with decode_to_text=True
• Median of 7 measured runs after warmup

implementation	tokenizer decode calls	median time	output positions/s	top-logprob token ids/s
old per-position loop	100,000	0.669s	149,468	747,342
PR batched backend decode	1	0.268s	373,415	1,867,076

Speedup: 2.50x.

Repro script, run from the repo root in an environment with SGLang installed, for example after pip install -e "python":

import gc
import os
import statistics
import time

from transformers import AutoTokenizer

from sglang.srt.managers.tokenizer_manager import TokenizerManager

MODEL = os.environ.get("MODEL", "gpt2")
POSITIONS = int(os.environ.get("POSITIONS", "100000"))
TOP_K = int(os.environ.get("TOP_K", "5"))
REPS = int(os.environ.get("REPS", "7"))
LOCAL_FILES_ONLY = os.environ.get("LOCAL_FILES_ONLY", "0") == "1"


class BenchManager:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer

    def detokenize_logprob_tokens(self, vals, idxs, decode_to_text):
        return TokenizerManager.detokenize_logprob_tokens(
            self, vals, idxs, decode_to_text
        )

    def _batch_decode_token_ids(self, token_ids):
        return TokenizerManager._batch_decode_token_ids(self, token_ids)


def old_detokenize_logprob_tokens(tokenizer, vals, idxs):
    token_texts = tokenizer.batch_decode([[idx] for idx in idxs])
    return list(zip(vals, idxs, token_texts))


def old_detokenize_top_logprobs_tokens(tokenizer, vals, idxs):
    ret = []
    for i in range(len(vals)):
        if vals[i]:
            ret.append(old_detokenize_logprob_tokens(tokenizer, vals[i], idxs[i]))
        else:
            ret.append(None)
    return ret


def new_detokenize_top_logprobs_tokens(mgr, vals, idxs):
    return TokenizerManager.detokenize_top_logprobs_tokens(
        mgr, vals, idxs, decode_to_text=True
    )


def make_case(positions, top_k, vocab_size):
    vals = [[-(j + 1) / 1000 for j in range(top_k)] for _ in range(positions)]
    idxs = [
        [(i * top_k + j) % vocab_size for j in range(top_k)]
        for i in range(positions)
    ]
    return vals, idxs


def run_once(fn, *args):
    gc.collect()
    gc.disable()
    start = time.perf_counter()
    ret = fn(*args)
    elapsed = time.perf_counter() - start
    gc.enable()
    assert len(ret) == POSITIONS and ret[-1][-1][2]
    return elapsed


tokenizer = AutoTokenizer.from_pretrained(MODEL, local_files_only=LOCAL_FILES_ONLY)
mgr = BenchManager(tokenizer)
vals, idxs = make_case(POSITIONS, TOP_K, tokenizer.vocab_size)
assert old_detokenize_top_logprobs_tokens(
    tokenizer, vals[:4], idxs[:4]
) == new_detokenize_top_logprobs_tokens(mgr, vals[:4], idxs[:4])

old_detokenize_top_logprobs_tokens(tokenizer, vals[:1000], idxs[:1000])
new_detokenize_top_logprobs_tokens(mgr, vals[:1000], idxs[:1000])

old_samples = [
    run_once(old_detokenize_top_logprobs_tokens, tokenizer, vals, idxs)
    for _ in range(REPS)
]
new_samples = [
    run_once(new_detokenize_top_logprobs_tokens, mgr, vals, idxs)
    for _ in range(REPS)
]

old_s = statistics.median(old_samples)
new_s = statistics.median(new_samples)
top_logprob_tokens = POSITIONS * TOP_K

print(f"Input: {POSITIONS:,} output positions, top_k={TOP_K}")
print(f"Top-logprob token ids: {top_logprob_tokens:,}")
print(
    f"old: {old_s:.3f}s, {POSITIONS / old_s:,.0f} positions/s, "
    f"{top_logprob_tokens / old_s:,.0f} top-logprob token ids/s"
)
print(
    f"new: {new_s:.3f}s, {POSITIONS / new_s:,.0f} positions/s, "
    f"{top_logprob_tokens / new_s:,.0f} top-logprob token ids/s"
)
print(f"speedup: {old_s / new_s:.2f}x")

Checklist

• Add unit tests according to the Run and add unit tests.
• Follow the SGLang code style guidance.

[gemini-code-assist]

Code Review

This pull request optimizes the detokenize_top_logprobs_tokens method in tokenizer_manager.py by batching the detokenization of top-k tokens across all positions, replacing the previous per-position approach to reduce overhead. A new test suite was also added to validate the batched implementation. Feedback suggests reusing the detokenize_logprob_tokens method when text decoding is not required to minimize code duplication.