[Bugfix] Widen the QSA raw-key ring instead of asserting divisibility

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)

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,