fix(token_counter): release the GIL for HuggingFace counts and cap exact counting per string

Both proxy token counting endpoints already count in a worker thread, but the
HuggingFace tokenizer's encode holds the GIL for the whole call, so a 600k-token
count on a Claude model still froze the event loop for up to 0.8 s and every
other request with it. Count through encode_batch_fast, which releases the GIL,
and tokenize at most TOKEN_COUNTER_MAX_EXACT_CHARS characters of any one string
(default 4,000,000), scaling the exact count of that prefix by the string's
length above it so the largest payloads stay bounded.

litellm/constants.py

@@ -391,6 +391,12 @@ TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS: Final = get_env_int_in_range(
     minimum=1,
     maximum=TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS,
 )
+TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range(
+    "TOKEN_COUNTER_MAX_EXACT_CHARS",
+    default=4_000_000,
+    minimum=1,
+    maximum=1_000_000_000,
+)
 MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512))
 MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512))
 OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000))

litellm/litellm_core_utils/token_counter.py

@@ -8,6 +8,7 @@ from typing import Final, Literal, cast
 
 import httpx
 import tiktoken
+from tokenizers import Tokenizer
 
 import litellm
 from litellm import verbose_logger
@@ -21,6 +22,7 @@ from litellm.constants import (
     MAX_TILE_HEIGHT,
     MAX_TILE_WIDTH,
     TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS,
+    TOKEN_COUNTER_MAX_EXACT_CHARS,
 )
 from litellm.litellm_core_utils.default_encoding import encoding as default_encoding
 from litellm.litellm_core_utils.url_utils import safe_get
@@ -538,9 +540,28 @@ def _count_extra(
     return num_tokens
 
 
+def _get_extrapolating_count_function(
+    count_exactly: TokenCounterFunction,
+    max_exact_chars: int = TOKEN_COUNTER_MAX_EXACT_CHARS,
+) -> TokenCounterFunction:
+    def count_tokens(text: str) -> int:
+        if len(text) <= max_exact_chars:
+            return count_exactly(text)
+        return round(count_exactly(text[:max_exact_chars]) * len(text) / max_exact_chars)
+
+    return count_tokens
+
+
 def _get_count_function(
     model: str | None,
     custom_tokenizer: dict | SelectTokenizerResponse | None = None,
+) -> TokenCounterFunction:
+    return _get_extrapolating_count_function(_get_exact_count_function(model, custom_tokenizer))
+
+
+def _get_exact_count_function(
+    model: str | None,
+    custom_tokenizer: dict | SelectTokenizerResponse | None = None,
 ) -> TokenCounterFunction:
     """
     Get the function to count tokens based on the model and custom tokenizer."""
@@ -549,10 +570,10 @@ def _get_count_function(
     if model is not None or custom_tokenizer is not None:
         tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model)
         if tokenizer_json["type"] == "huggingface_tokenizer":
+            tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"]
 
             def count_tokens(text: str) -> int:
-                enc: Final = tokenizer_json["tokenizer"].encode(text)
-                return len(enc.ids)
+                return len(tokenizer.encode_batch_fast([text])[0])
 
             return count_tokens
         elif tokenizer_json["type"] == "openai_tokenizer":

tests/test_litellm/litellm_core_utils/test_token_counter.py

@@ -1,8 +1,10 @@
 #### What this tests ####
 #    This tests litellm.token_counter.token_counter() function
+import asyncio
 import importlib
 import time
 import traceback
+from typing import Final
 from unittest.mock import MagicMock
 
 import pytest
@@ -14,7 +16,12 @@ import litellm
 from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens
 from litellm import token_counter as token_counter_old
 import litellm.constants
-from litellm.litellm_core_utils.token_counter import _get_tiktoken_count_function
+from litellm.litellm_core_utils.asyncify import asyncify
+from litellm.litellm_core_utils.token_counter import (
+    _get_exact_count_function,
+    _get_extrapolating_count_function,
+    _get_tiktoken_count_function,
+)
 from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
 from tests.large_text import text
 from tests.test_litellm.litellm_core_utils.messages_with_counts import (
@@ -120,6 +127,68 @@ def test_valid_chunk_size_config_is_honoured(monkeypatch):
         importlib.reload(litellm.constants)
 
 
+async def _loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]:
+    async def wake_lag() -> float:
+        started: Final = time.perf_counter()
+        await asyncio.sleep(0.001)
+        return time.perf_counter() - started - 0.001
+
+    return tuple([await wake_lag() for _ in iter(until.is_set, True)])
+
+
+async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free():
+    counted: Final = asyncio.Event()
+
+    async def count_off_the_loop() -> tuple[int, float]:
+        started: Final = time.perf_counter()
+        try:
+            tokens: Final = await asyncify(token_counter_new)(model="claude-fable-5", text=text * 100)
+            return tokens, time.perf_counter() - started
+        finally:
+            counted.set()
+
+    (tokens, took), lags = await asyncio.gather(count_off_the_loop(), _loop_wake_lags(counted))
+
+    assert tokens > 0
+    assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count"
+
+
+@pytest.mark.parametrize(
+    ("max_exact_chars", "expected"),
+    [(1_000, 10_000), (5_000, 6_000), (10_000, 6_000)],
+)
+def test_count_above_the_cap_scales_the_exact_count_of_the_prefix(max_exact_chars, expected):
+    def count_exactly(chunk: str) -> int:
+        return chunk.count("a") + len(chunk)
+
+    count_tokens: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars)
+
+    assert count_tokens("a" * 1_000 + "b" * 4_000) == expected
+
+
+def test_token_counter_applies_the_default_cap():
+    max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS
+    prefix: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars]
+    over_the_cap: Final = prefix + "a" * 200_000
+    scaled: Final = round(token_counter_new(model="gpt-5.6", text=prefix) * len(over_the_cap) / max_exact_chars)
+
+    assert token_counter_new(model="gpt-5.6", text=over_the_cap) == scaled
+    assert _get_exact_count_function("gpt-5.6")(over_the_cap) != scaled
+
+
+@pytest.mark.parametrize(
+    ("configured", "expected"),
+    [("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)],
+)
+def test_max_exact_chars_config_is_honoured(monkeypatch, configured, expected):
+    monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured)
+    try:
+        assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected
+    finally:
+        monkeypatch.delenv("TOKEN_COUNTER_MAX_EXACT_CHARS")
+        importlib.reload(litellm.constants)
+
+
 def test_token_counter_with_prefix():
     messages = [
         {"role": "user", "content": "Who won the world cup in 2022?"},

tests/test_litellm/proxy/test_proxy_server.py

@@ -6,6 +6,7 @@ import os
 import re
 import socket
 import subprocess
+import time
 import types
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
@@ -28,7 +29,7 @@ from litellm.caching.caching import RedisCache
 from litellm.caching.redis_cluster_cache import RedisClusterCache
 from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded
 from litellm.caching.dual_cache import DualCache
-from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.proxy_server import app, initialize
 from litellm.utils import _invalidate_model_cost_lowercase_map
@@ -12857,3 +12858,34 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough()
         import litellm.proxy.proxy_server as ps
 
         assert ps.general_settings["enable_openai_websocket_passthrough"] is False
+
+
+async def _loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]:
+    async def wake_lag() -> float:
+        started: Final = time.perf_counter()
+        await asyncio.sleep(0.001)
+        return time.perf_counter() - started - 0.001
+
+    return tuple([await wake_lag() for _ in iter(until.is_set, True)])
+
+
+async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch):
+    from tests.large_text import text
+
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
+    counted: Final = asyncio.Event()
+
+    async def count_off_the_loop() -> tuple[int, float]:
+        started: Final = time.perf_counter()
+        try:
+            response: Final = await proxy_server_module.token_counter(
+                TokenCountRequest(model="claude-fable-5", prompt=text * 100)
+            )
+            return response.total_tokens, time.perf_counter() - started
+        finally:
+            counted.set()
+
+    (tokens, took), lags = await asyncio.gather(count_off_the_loop(), _loop_wake_lags(counted))
+
+    assert tokens > 0
+    assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count"