From 6d4c756115a44ddf758b72997ae7a63f661567fe Mon Sep 17 00:00:00 2001 From: Arbi Admin Date: Fri, 2 Oct 2026 09:46:14 +0100 Subject: [PATCH] fix(args): auto-default --attention-backend to TURBO_ATTN for plugin kv dtypes Forward-port of 0a0fbe182a + 830edc03aa (fix/auto-default-turbo-attn) onto the v0.28.0 seam. A plugin-registered --kv-cache-dtype (tkv) selected TURBO_ATTN only through the platform, leaving attention_config.backend None. gpu_worker._backends_in_use dispatches on_model_loaded / adjust_kv_budget before the attention groups exist, from that field alone, so the plugin's load-time hooks never ran on a `--kv-cache-dtype tkv` boot (no ring reservation, o_proj fold catch-up or pre-flight). An explicit backend is honoured; builtin dtypes are unchanged. The existing arg_utils test, which expected the CUSTOM default the port never carried, now expects TURBO_ATTN. Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/engine/test_arg_utils.py | 11 +++++++---- vllm/config/cache.py | 2 +- vllm/engine/arg_utils.py | 23 +++++++++++++++++++++++ 3 files changed, 31 insertions(+), 5 deletions(-) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index d81f4e1de950..56f86e21a040 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -889,9 +889,12 @@ def test_dp_rank_shards_user_assigned_gpu_ids(self): ) -def test_plugin_kv_cache_dtype_auto_selects_custom_attention_backend(): +def test_plugin_kv_cache_dtype_auto_selects_turbo_attn_backend(): """A plugin-registered kv-cache dtype should auto-default the attention - backend to CUSTOM, but only when the user did not pass an explicit choice. + backend to TURBO_ATTN (the slot such plugins register into), but only when + the user did not pass an explicit choice. The worker's load-time lifecycle + hooks dispatch only to a named backend, so without this a + ``--kv-cache-dtype tkv`` boot never ran the plugin's ``on_model_loaded``. """ import torch @@ -907,7 +910,7 @@ def test_plugin_kv_cache_dtype_auto_selects_custom_attention_backend(): try: parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) - # 1. Plugin dtype + no --attention-backend → CUSTOM auto-selected. + # 1. Plugin dtype + no --attention-backend → TURBO_ATTN auto-selected. args = parser.parse_args( [ "--model", @@ -919,7 +922,7 @@ def test_plugin_kv_cache_dtype_auto_selects_custom_attention_backend(): engine_args = EngineArgs.from_cli_args(args) vllm_config = engine_args.create_engine_config() assert vllm_config.attention_config.backend == ( - AttentionBackendEnum.CUSTOM + AttentionBackendEnum.TURBO_ATTN ) # 2. Plugin dtype + explicit --attention-backend → user choice honoured. diff --git a/vllm/config/cache.py b/vllm/config/cache.py index f2e07693f1db..0e65a1be2f30 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -84,7 +84,7 @@ def is_plugin_cache_dtype(name: str) -> bool: """Return True if ``name`` was registered via ``register_cache_dtype``. Used by arg post-processing to auto-default ``--attention-backend`` to - ``CUSTOM`` when the user selects a plugin-registered KV cache dtype. + ``TURBO_ATTN`` when the user selects a plugin-registered KV cache dtype. """ return name in _PLUGIN_CACHE_DTYPES diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 0fd22fc7771b..d7ecde725270 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -73,6 +73,7 @@ MambaCacheMode, MambaDType, PrefixCachingHashAlgo, + is_plugin_cache_dtype, ) from vllm.config.device import Device from vllm.config.kernel import IrOpPriorityConfig, LinearBackend, MoEBackend @@ -2406,6 +2407,28 @@ def create_engine_config( self.attention_backend ) + # Plugin-registered KV cache dtypes (e.g. tkv) ship their own + # AttentionBackend, registered under AttentionBackendEnum.TURBO_ATTN + # via register_backend(). When the user selects such a dtype without + # --attention-backend, default to TURBO_ATTN: the platform would pick + # it anyway, but only a NAMED backend reaches the lifecycle hooks the + # worker dispatches before its attention groups exist + # (gpu_worker._backends_in_use -> on_model_loaded, adjust_kv_budget; + # kv_cache_manager -> on_kv_manager_created). An explicit choice + # (--attention-backend or --attention-config.backend) is honoured. + if attention_config.backend is None and is_plugin_cache_dtype( + resolved_cache_dtype + ): + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + attention_config.backend = AttentionBackendEnum.TURBO_ATTN + logger.info( + "Auto-selecting attention backend TURBO_ATTN for " + "plugin-registered kv-cache dtype %r. Pass " + "--attention-backend to override.", + resolved_cache_dtype, + ) + # TurboQuant requires FlashAttention 2 — FA3 boundary layers assert # FlashAttentionImpl which fails with TurboQuantAttentionImpl. if resolved_cache_dtype.startswith("turboquant_") and (