Overview

HiSparse crashes when a batch has exactly max_num_seqs * max_decode_query_len + 1 tokens. With max_num_seqs=1 the engine fails at startup, and with max_num_seqs=2 a 3-token prompt crashes at runtime:

ValueError: HiSparse swap rows exceed the configured speculative decode capacity: stop=2, capacity=1.

Part of #60072.

Claims

• HiSparse now routes a batch to swap_in only when the batch fits the swap rows that swap_in consumes. Larger batches take the existing all-resident fallback.
• No performance impact: the change is a Python-side routing comparison against a fixed size. Batches up to the swap capacity take the same path as before; only the single size that crashed changes path.

Validation

1×H200, upstream main 242e421. Repro (tiny DeepSeek-V3.2, dummy weights):

import sys
from vllm import LLM, SamplingParams
from vllm.config import AttentionConfig, HiSparseConfig, KVTransferConfig

llm = LLM(
    "deepseek-ai/DeepSeek-V3.2", load_format="dummy",
    hf_overrides={"num_hidden_layers": 8, "hidden_size": 256, "intermediate_size": 512,
                  "num_attention_heads": 8, "num_key_value_heads": 1, "n_routed_experts": 8,
                  "num_experts_per_tok": 2, "index_topk": 128},
    max_model_len=320, block_size=64, max_num_seqs=int(sys.argv[1]),
    enable_prefix_caching=False, gpu_memory_utilization=0.3,
    kv_transfer_config=KVTransferConfig(kv_connector="HiSparseConnector", kv_role="kv_both",
                                        kv_connector_extra_config={"host_pool_gib": 1}),
    attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
)
for n in (2, 3, 200):
    llm.generate([{"prompt_token_ids": list(range(100, 100 + n))}],
                 SamplingParams(max_tokens=8, ignore_eos=True))

Tests on this PR (each command run separately):

• pytest tests/v1/attention/test_sparse_mla_backends.py -k "hisparse or index_group": 40 passed. The regression case test_hisparse_resident_prefill_uses_attention_block_stride fails on main and passes with this PR.
• pytest tests/v1/kv_connector/unit/test_hisparse_connector.py: 9 passed
• pytest tests/v1/core -k hisparse: 72 passed
• pytest tests/v1/e2e/general/test_hisparse.py: 4 passed

Details

Root cause. HiSparseMLAIndexGroup.convert_logical_to_physical_topk sends a batch to swap_in whenever the batch fits physical_topk_indices. That index-group workspace is allocated with one spare row (max_decode_rows + 1). swap_in, however, writes into the HiSparse runtime's shared swap rows, which hold exactly min(max_num_batched_tokens, max_num_seqs * max_decode_query_len) rows. A batch one token larger than the swap capacity therefore passes the routing check and then overflows in _step_rows.

With max_num_seqs=1, the CUDA-graph capture sizes are [1, 2]. The 2-token piecewise capture follows the all-resident prefill path into this conversion, so the engine fails at startup. With max_num_seqs=2, capture sizes [1, 2, 4] skip 3 tokens, so the failure only appears when a real 3-token batch arrives.

Fix. Route on HiSparseRuntime.max_swap_rows, the capacity that swap_in actually consumes. The non-HiSparse SparseMLAIndexGroup does not change.

🤖 Generated with Claude Code