From 121c09e9286fc680bd01cb3620619338a41ea87b Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 2 Oct 2026 03:49:56 +0000 Subject: [PATCH 1/5] perf(proxy): count admission tokens with the Rust fast counter when LITELLM_RUST is on Route.TOKEN_COUNTER rolls out as RUST_OPT_IN, so a process with LITELLM_RUST=1 reaches the native body counter instead of the Python tokenizers, and the counter is built over the byte-level fast path (exact count, no encode). Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/rust_bridge/catalog.py | 2 +- litellm/rust_bridge/token_counter.py | 2 +- tests/unit/rust_bridge/test_catalog.py | 2 +- tests/unit/rust_bridge/test_token_counter.py | 4 ++-- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 40c6456431d..1abd273c0e3 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -88,7 +88,7 @@ RULES: Final[Rules] = ( RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), - RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), + RouteRule(Route.TOKEN_COUNTER, Rollout.RUST_OPT_IN), RouteRule(Route.TOKENIZER, Rollout.PYTHON_ONLY), RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.GOOGLE_KMS.value})), diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index 250ad18d44c..f289357ccd7 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/litellm/rust_bridge/token_counter.py @@ -79,7 +79,7 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: @lru_cache(maxsize=4) def _counter(factory: RustTokenCounterFactory, tokenizer: RustTokenizer) -> RustTokenCounter: - return factory.from_tokenizer(_native_tokenizer(tokenizer)) + return factory.from_tokenizer(_native_tokenizer(tokenizer), fast=True) def _native_tokenizer(tokenizer: RustTokenizer) -> NativeTokenizer: diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index 95e98fe98da..e7c50739e04 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -47,7 +47,7 @@ def test_shipped_decisions( if route is Route.OCR or (route is Route.TRANSCRIPTION and provider == "bedrock"): assert catalog.rollout(context) is Rollout.RUST_REQUIRED assert catalog.decision(context) is Decision.RUST_REQUIRED - elif route is Route.MESSAGES and provider == "anthropic": + elif (route is Route.MESSAGES and provider == "anthropic") or route is Route.TOKEN_COUNTER: assert catalog.rollout(context) is Rollout.RUST_OPT_IN opted_in: Final = environment == "1" or (environment is None and process is True) assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if opted_in else Decision.PYTHON) diff --git a/tests/unit/rust_bridge/test_token_counter.py b/tests/unit/rust_bridge/test_token_counter.py index 3af7127a9b3..d6358bee721 100644 --- a/tests/unit/rust_bridge/test_token_counter.py +++ b/tests/unit/rust_bridge/test_token_counter.py @@ -94,7 +94,7 @@ async def test_native_count_returns_typed_count_and_reuses_one_counter(fake_toke assert second == first assert len(factory.counters) == 1 assert factory.counters[0].bodies == [BODY, BODY] - assert factory.counters[0].fast is False + assert factory.counters[0].fast is True assert factory.counters[0].tokenizer is tokenizer_dispatch.native_anthropic() assert json.loads(factory.counters[0].tokenizer.json or "")["model"]["type"] == "BPE" @@ -113,7 +113,7 @@ async def test_tiktoken_counter_is_built_over_the_shared_encoding_once( assert len(factory.counters) == 1 assert factory.counters[0].tokenizer.name == tokenizer assert factory.counters[0].tokenizer is tokenizer_dispatch.native_encoding(tokenizer) - assert factory.counters[0].fast is False + assert factory.counters[0].fast is True assert factory.counters[0].bodies == [BODY, BODY] From f342dc2249cb2184e58bdd3b09c4d7462681433e Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 2 Oct 2026 04:39:49 +0000 Subject: [PATCH 2/5] perf(router): skip prompt token counting in usage-based-routing-v2 when no deployment declares a tpm limit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_strategy/lowest_tpm_rpm_v2.py | 17 +++++-- .../router_strategy/test_lowest_tpm_rpm.py | 46 +++++++++++++++---- 2 files changed, 50 insertions(+), 13 deletions(-) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 25564a80e0a..d14d1f8b58f 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -70,6 +70,13 @@ class PrefetchedUsage: return _active_prefetched_usage.get() +def _declares_tpm_limit(deployment: Mapping[str, object]) -> bool: + """Whether the deployment carries a ``tpm`` limit that selection has to weigh the prompt against.""" + nested: Final = (deployment.get("litellm_params"), deployment.get("model_info")) + sources: Final = (deployment, *(source for source in nested if isinstance(source, Mapping))) + return any(source.get("tpm") is not None for source in sources) + + class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): """ Updated version of TPM/RPM Logging. @@ -416,10 +423,12 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): for idx, key in enumerate(rpm_keys): rpm_dict[rpm_keys[idx].split(":")[0]] = rpm_values[idx] - try: - input_tokens = token_counter(messages=messages, text=input) - except Exception: - input_tokens = 0 + input_tokens = 0 + if any(_declares_tpm_limit(deployment) for deployment in healthy_deployments): + try: + input_tokens = token_counter(messages=messages, text=input) + except Exception: + input_tokens = 0 verbose_router_logger.debug("input_tokens=%s", input_tokens) # ----------------------- # Find lowest used model diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm.py b/tests/unit/router_strategy/test_lowest_tpm_rpm.py index 625f648bec4..66de439e6bc 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm.py @@ -1,9 +1,10 @@ from datetime import datetime, timedelta from typing import Final -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import pytest +import litellm.router_strategy.lowest_tpm_rpm_v2 as strategy_module from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage @@ -42,14 +43,11 @@ def test_usage_based_routing_v1_selects_the_lowest_recorded_tpm() -> None: } now: Final = datetime.now() cache_keys: Final = tuple( - f"{MODEL_GROUP}:tpm:{(now + timedelta(minutes=offset)).strftime('%H-%M')}" - for offset in range(60) + f"{MODEL_GROUP}:tpm:{(now + timedelta(minutes=offset)).strftime('%H-%M')}" for offset in range(60) ) for cache_key in cache_keys: - router.cache.set_cache( - key=cache_key, value=usage_by_deployment, ttl=float("inf") - ) + router.cache.set_cache(key=cache_key, value=usage_by_deployment, ttl=float("inf")) deployment: Final = router.get_available_deployment( model=MODEL_GROUP, @@ -109,11 +107,41 @@ async def test_v2_subclass_overriding_async_get_available_deployments_with_the_o ) router.lowesttpm_logger_v2 = OldSignatureV2(router_cache=router.cache, routing_args={}) - response: Final = await router.acompletion( - model=MODEL_GROUP, messages=[{"role": "user", "content": "x"}] - ) + response: Final = await router.acompletion(model=MODEL_GROUP, messages=[{"role": "user", "content": "x"}]) assert response.choices[0].message.content in { f"from {HIGH_USAGE_DEPLOYMENT_ID}", f"from {LOW_USAGE_DEPLOYMENT_ID}", } + + +@pytest.mark.asyncio +async def test_v2_counts_prompt_tokens_only_when_a_deployment_declares_a_tpm_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = Mock(return_value=40_000) + monkeypatch.setattr(strategy_module, "token_counter", counter) + router_cache = DualCache() + router_cache.async_batch_get_cache = AsyncMock(return_value=[10, 20, None, None]) # type: ignore[method-assign] + strategy = LowestTPMLoggingHandler_v2(router_cache=router_cache) + messages: Final = [{"role": "user", "content": "a long prompt"}] + unlimited: Final = [ + {"model_name": "g", "litellm_params": {"model": "m"}, "model_info": {"id": "a"}}, + {"model_name": "g", "litellm_params": {"model": "m"}, "model_info": {"id": "b"}}, + ] + + chosen: Final = await strategy.async_get_available_deployments( + model_group="g", healthy_deployments=unlimited, messages=messages + ) + assert chosen["model_info"]["id"] == "a", "the lowest counter still wins without any tpm limit" + assert counter.call_count == 0, "no deployment has a tpm limit, so the prompt is never tokenized for selection" + + limited: Final = [ + {"model_name": "g", "litellm_params": {"model": "m", "tpm": 30_000}, "model_info": {"id": "a"}}, + {"model_name": "g", "litellm_params": {"model": "m"}, "model_info": {"id": "b"}}, + ] + chosen_limited: Final = await strategy.async_get_available_deployments( + model_group="g", healthy_deployments=limited, messages=messages + ) + assert counter.call_count == 1, "a declared tpm limit needs the prompt size to be enforced" + assert chosen_limited["model_info"]["id"] == "b", "a's tpm limit cannot fit the counted prompt, so b is picked" From 7971227e053d3666fe610c7a61e0fe942407bd88 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 2 Oct 2026 04:55:31 +0000 Subject: [PATCH 3/5] fix(router): type the deployment sequence the tpm-limit check iterates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_strategy/lowest_tpm_rpm_v2.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index d14d1f8b58f..4239bdfaa61 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -5,7 +5,7 @@ from collections.abc import Iterator, Mapping, Sequence from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, cast import httpx @@ -424,7 +424,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): rpm_dict[rpm_keys[idx].split(":")[0]] = rpm_values[idx] input_tokens = 0 - if any(_declares_tpm_limit(deployment) for deployment in healthy_deployments): + deployments: Final = cast( + Sequence[Mapping[str, object]], healthy_deployments + ) # cast-ok: router deployments are untyped dicts at this boundary + if any(_declares_tpm_limit(deployment) for deployment in deployments): try: input_tokens = token_counter(messages=messages, text=input) except Exception: From 0d7c17d32eea4dec276def3845ab6b6d2e84a1c8 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 2 Oct 2026 04:59:14 +0000 Subject: [PATCH 4/5] test(router): drop the type ignore from the tpm-limit test and cover a failed prompt count Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../router_strategy/test_lowest_tpm_rpm.py | 23 ++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm.py b/tests/unit/router_strategy/test_lowest_tpm_rpm.py index 66de439e6bc..a3d183ff129 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm.py @@ -121,9 +121,9 @@ async def test_v2_counts_prompt_tokens_only_when_a_deployment_declares_a_tpm_lim ) -> None: counter: Final = Mock(return_value=40_000) monkeypatch.setattr(strategy_module, "token_counter", counter) - router_cache = DualCache() - router_cache.async_batch_get_cache = AsyncMock(return_value=[10, 20, None, None]) # type: ignore[method-assign] - strategy = LowestTPMLoggingHandler_v2(router_cache=router_cache) + router_cache: Final = DualCache() + monkeypatch.setattr(router_cache, "async_batch_get_cache", AsyncMock(return_value=[10, 20, None, None])) + strategy: Final = LowestTPMLoggingHandler_v2(router_cache=router_cache) messages: Final = [{"role": "user", "content": "a long prompt"}] unlimited: Final = [ {"model_name": "g", "litellm_params": {"model": "m"}, "model_info": {"id": "a"}}, @@ -145,3 +145,20 @@ async def test_v2_counts_prompt_tokens_only_when_a_deployment_declares_a_tpm_lim ) assert counter.call_count == 1, "a declared tpm limit needs the prompt size to be enforced" assert chosen_limited["model_info"]["id"] == "b", "a's tpm limit cannot fit the counted prompt, so b is picked" + + +@pytest.mark.asyncio +async def test_v2_treats_a_failed_prompt_count_as_zero_tokens(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(strategy_module, "token_counter", Mock(side_effect=ValueError("no tokenizer"))) + router_cache: Final = DualCache() + monkeypatch.setattr(router_cache, "async_batch_get_cache", AsyncMock(return_value=[10, 20, None, None])) + strategy: Final = LowestTPMLoggingHandler_v2(router_cache=router_cache) + limited: Final = [ + {"model_name": "g", "litellm_params": {"model": "m", "tpm": 30_000}, "model_info": {"id": "a"}}, + {"model_name": "g", "litellm_params": {"model": "m"}, "model_info": {"id": "b"}}, + ] + + chosen: Final = await strategy.async_get_available_deployments( + model_group="g", healthy_deployments=limited, messages=[{"role": "user", "content": "a long prompt"}] + ) + assert chosen["model_info"]["id"] == "a", "a failed count weighs as zero tokens, so the lowest counter still wins" From 19c5790a916fbeca8c1f0b4343facdc52e47d177 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 2 Oct 2026 05:06:53 +0000 Subject: [PATCH 5/5] fix(router): keep the cast-ok suppression on the cast line for the type discipline gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_strategy/lowest_tpm_rpm_v2.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 4239bdfaa61..f630fe373ca 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -423,10 +423,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): for idx, key in enumerate(rpm_keys): rpm_dict[rpm_keys[idx].split(":")[0]] = rpm_values[idx] + deployments: Final = cast(Sequence[Mapping[str, object]], healthy_deployments) # cast-ok: untyped dicts input_tokens = 0 - deployments: Final = cast( - Sequence[Mapping[str, object]], healthy_deployments - ) # cast-ok: router deployments are untyped dicts at this boundary if any(_declares_tpm_limit(deployment) for deployment in deployments): try: input_tokens = token_counter(messages=messages, text=input)