Skip to content

Batch detokenization for logprob responses #197

Description

@AgrawalAmey

Problem

When a client requests logprobs > 0, every token gets decoded one-at-a-time in a Python loop at vajra/entrypoints/openai/handlers/base.py:52, 65:

token_str: str = str(self.tokenizer.decode([entry.token]))
...
alt_str: str = str(self.tokenizer.decode([alt.token_id]))

For a 1000-token response with 5 alternatives each, that's 6000 Python → tokenizer → C++ → Python round-trips. Logprob-heavy evaluations take the biggest hit.

Fix

Collect every token_id across the full response into one flat list, call tokenizer.batch_decode(...) once, then scatter the decoded strings back to their positions. Saves 1000× Python overhead for the main tokens plus 5000× for the alternatives.

Even better: move the decode entirely to C++ (Vajra's C++ tokenizer wrapper at csrc/vajra/native/core/tokenizer/Tokenizer.h already has Decode). Surface the decoded strings through TokenLogprob / LogprobData directly.

Acceptance criteria

  • Logprob response construction performs exactly one tokenizer call per request (or zero, if the C++ engine surfaces decoded strings).
  • Output bytes identical to the current path on a fixture.

Parent

Part of VJ-187.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions