55 lines
2.6 KiB
Diff
55 lines
2.6 KiB
Diff
# UNDF: UNDF-2026-000001175
|
||
# vllm-0002: Grok2Tokenizer decode/convert_ids_to_tokens O(N×S) dict.values() scan
|
||
# CWE-407 — Algorithmic Complexity
|
||
#
|
||
# In Grok2Tokenizer.decode() and convert_ids_to_tokens(), filtering special tokens
|
||
# uses `token_id not in self._special_tokens.values()` where `self._special_tokens`
|
||
# is a dict[str, int]. dict.values() is a view — membership test is O(V) linear scan
|
||
# (V = number of special tokens). This check runs per output token in the decode loop,
|
||
# making the total cost O(N × V) where N = sequence length.
|
||
#
|
||
# The Mistral tokenizer in the same codebase already fixes this correctly by caching
|
||
# `_special_token_ids_set: frozenset[int]` at init time. Grok2 missed that pattern.
|
||
#
|
||
# Fix: add `self._special_token_ids: frozenset[int]` at __init__ time and replace
|
||
# all `.values()` membership checks with the frozenset lookup.
|
||
#
|
||
# Severity: HIGH — decode() is called per output token during streaming LLM inference.
|
||
# At N=2048 tokens and V=200 special tokens: 409,600 unnecessary comparisons per request.
|
||
# With 100 concurrent requests: 40M comparisons per batch step.
|
||
# Speedup: ~200x at V=200 for the membership check portion.
|
||
#
|
||
# File: vllm/tokenizers/grok2.py
|
||
# Class: Grok2Tokenizer
|
||
# Methods: __init__, decode, convert_ids_to_tokens
|
||
--- a/vllm/tokenizers/grok2.py
|
||
+++ b/vllm/tokenizers/grok2.py
|
||
@@ -278,6 +278,9 @@ class Grok2Tokenizer:
|
||
self._eos_token_id = self._special_tokens.get(EOS, self._bos_token_id)
|
||
self._pad_token_id = self._special_tokens.get(PAD, self._eos_token_id)
|
||
self._unk_token_id = self._pad_token_id
|
||
+
|
||
+ # Cache special token IDs as a frozenset for O(1) membership checks in decode.
|
||
+ self._special_token_ids: frozenset[int] = frozenset(self._special_tokens.values())
|
||
|
||
self._max_chars_per_token = max(len(tok) for tok in self._token_to_id)
|
||
|
||
@@ -354,7 +357,7 @@ class Grok2Tokenizer:
|
||
if skip_special_tokens:
|
||
ids = [
|
||
token_id
|
||
- for token_id in ids
|
||
- if token_id not in self._special_tokens.values()
|
||
+ for token_id in ids
|
||
+ if token_id not in self._special_token_ids
|
||
]
|
||
return self._tokenizer.decode(ids)
|
||
|
||
@@ -376,7 +379,7 @@ class Grok2Tokenizer:
|
||
tokens = []
|
||
for token_id in ids:
|
||
- if skip_special_tokens and token_id in self._special_tokens.values():
|
||
+ if skip_special_tokens and token_id in self._special_token_ids:
|
||
continue
|
||
tokens.append(self._id_to_token.get(token_id, "<|unk|>"))
|
||
return tokens
|