diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index 9d30365b8151..266096547a29 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -13,11 +13,16 @@ Qwen3_5ForConditionalGenerationConfig, Qwen4ExpForConditionalGenerationConfig, ) +from vllm.models.qwen4_exp.common.qsa_cache import ( + QSA_RING_MAX_WIDENING, + qsa_ring_capacity, +) from vllm.models.qwen4_exp.config import ( Qwen4ExpConfig, Qwen4ExpTextConfig, ) from vllm.models.qwen4_exp.nvidia.model_state import Qwen4ExpModelState +from vllm.utils.math_utils import cdiv from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState from ...utils import spawn_new_process_for_each_test @@ -246,3 +251,39 @@ def test_qwen4_exp_model_state_prepares_stable_dummy_ngram_inputs() -> None: ) assert second["query_start_loc"].data_ptr() == query_start_loc_ptr assert second["ngram_context"].data_ptr() == ngram_context_ptr + + +@pytest.mark.parametrize("block_size", [848, 1616]) +def test_qsa_ring_capacity_divides_block_size(block_size: int) -> None: + compress_ratio = 4 + # Depths 0..12 are servable on both hybrid block sizes (16 * 53, 16 * 101). + for num_spec in range(13): + span = compress_ratio + num_spec + minimal = compress_ratio * cdiv(span, compress_ratio) + capacity = qsa_ring_capacity(compress_ratio, num_spec, block_size) + assert capacity >= span + assert capacity % compress_ratio == 0 + assert block_size % capacity == 0 + assert capacity <= QSA_RING_MAX_WIDENING * minimal + if block_size % minimal == 0: + # Every previously legal depth keeps its ring size. + assert capacity == minimal + + +def test_qsa_ring_capacity_widens_within_bound() -> None: + # num_speculative_tokens 5..8 need 12 rows; 848 and 1616 have no factor 3, + # so the ring widens to 16, the next multiple of 4 that divides them. + assert qsa_ring_capacity(4, 5, 848) == 16 + assert qsa_ring_capacity(4, 8, 1616) == 16 + assert qsa_ring_capacity(4, 5, 16) == 16 + + +@pytest.mark.parametrize("block_size", [16, 848, 1616]) +def test_qsa_ring_capacity_refuses_disproportionate_widening( + block_size: int, +) -> None: + # num_speculative_tokens 13..16 need 20 rows. The next divisors of 848 and + # 1616 that are multiples of 4 are 212 and 404 (16 * 53, 16 * 101); 16 has + # none. All three must fail loudly instead of allocating a 10x ring. + with pytest.raises(ValueError, match="QSA ring"): + qsa_ring_capacity(4, 13, block_size) diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py index e806dc0375de..552b65fdba4e 100644 --- a/vllm/models/qwen4_exp/common/qsa_cache.py +++ b/vllm/models/qwen4_exp/common/qsa_cache.py @@ -21,6 +21,7 @@ from vllm.config import CacheConfig, VllmConfig from vllm.config.cache import CacheDType +from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, tl, triton @@ -44,6 +45,8 @@ MLAAttentionSpec, ) +logger = init_logger(__name__) + def canonical_qsa_rope_positions(positions: torch.Tensor) -> torch.Tensor: """Return exact per-token positions as ``[tokens, 1, 3]`` int64 rows.""" @@ -808,6 +811,46 @@ def get_attn_backend(self) -> type[AttentionBackend]: return QSAStateBackend +# The ring may be widened past the minimal whole-group size by at most this +# factor. Widening keeps the scheduler LCM unchanged, but every request holds +# one ring block for its lifetime, so a large divisor must not be picked +# silently (e.g. 20 rows -> 212 on a block size of 848 = 16 * 53). +QSA_RING_MAX_WIDENING = 2 + + +def qsa_ring_capacity( + compress_ratio: int, num_speculative_tokens: int, block_size: int +) -> int: + """Rows of the QSA raw-key ring for a given attention block size. + + The ring holds the open group's committed keys plus every row a + speculative step stores before acceptance is known (``span``). Anything + narrower lets a rejected draft row overwrite a committed key the next + step needs to close the group, so ``span`` is a lower bound; a wider ring + is slack. The ring is rounded up to whole groups, and a whole-group size + that divides ``block_size`` is preferred so that adding the circular-buffer + group does not increase the scheduler block size (the LCM over all + groups). When the minimal size does not divide ``block_size``, the next + multiple of ``compress_ratio`` that does is used, up to + ``QSA_RING_MAX_WIDENING`` times the minimal size; beyond that a + ``ValueError`` is raised rather than allocating a disproportionate ring. + """ + span = compress_ratio + num_speculative_tokens + minimal = compress_ratio * cdiv(span, compress_ratio) + limit = min(block_size, QSA_RING_MAX_WIDENING * minimal) + capacity = minimal + while capacity <= limit: + if block_size % capacity == 0: + return capacity + capacity += compress_ratio + raise ValueError( + f"QSA ring needs {minimal} rows for num_speculative_tokens=" + f"{num_speculative_tokens} (compress ratio {compress_ratio}), and no " + f"multiple of {compress_ratio} up to {limit} divides the attention " + f"block size {block_size}; choose a depth whose ring divides it" + ) + + class QSAKeyStateCache(_QSAStateCache): """Raw BF16 key, optionally followed by exact int64 MRoPE positions.""" @@ -839,17 +882,20 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: self.rope_position_cache = None def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: - # Hold the open group's committed keys plus every row a speculative - # step stores before acceptance is known, rounded up to whole groups so - # the ring divides the attention block size (it joins the LCM that sets - # the scheduler block size). Anything narrower lets a rejected draft row - # overwrite a committed key the next step needs to close the group. - span = self.compress_ratio + vllm_config.num_speculative_tokens - capacity = self.compress_ratio * cdiv(span, self.compress_ratio) - assert self.cache_config.block_size % capacity == 0, ( - f"QSA ring capacity {capacity} must divide the attention block " - f"size {self.cache_config.block_size}" - ) + num_spec = vllm_config.num_speculative_tokens + block_size = self.cache_config.block_size + span = self.compress_ratio + num_spec + minimal = self.compress_ratio * cdiv(span, self.compress_ratio) + capacity = qsa_ring_capacity(self.compress_ratio, num_spec, block_size) + if capacity != minimal: + logger.info_once( + "QSA ring widened from %d to %d rows so that it divides the " + "attention block size %d (num_speculative_tokens=%d).", + minimal, + capacity, + block_size, + num_spec, + ) return CircularBufferSpec( block_size=capacity, num_kv_heads=1,