[Feature]: # [Question] Does UCM support Qwen2.5-14B-Instruct-1M at extended context (>32K)?

[baoluo878]

🚀 The feature, motivation and pitch

Environment

• UCM Version: v0.5.0rc2 (develop branch)
• vLLM Version: Tested with 0.11.0 and 0.18.1
• GPU: 1× NVIDIA RTX A6000 (48GB)
• Model: Qwen/Qwen2.5-14B-Instruct-1M
• OS: Ubuntu (Docker container)

Background

We are deploying Qwen2.5-14B-Instruct-1M with UCM for production use. The model ships with a sparse_attention_config.json that enables Qwen's Dual Chunk Attention (DCA) for extended context windows.

Currently, we have the model working at 32K context with UCM prefix caching enabled, which works well. However, we would like to leverage the model's designed 1M context capability.

What We Tried

1. vLLM 0.18.1 + UCM v0.5.0rc2

• UCM patches load successfully (pc_patch only, no sparse_patch for v0180)
• With sparse_attention_config.json present: FlashAttentionImpl.__init__() got an unexpected keyword argument 'layer_idx'
• Without sparse_attention_config.json: works at 32K, but no sparse attention → limited by GPU memory

2. vLLM 0.11.0 + UCM v0.5.0rc2 (ENABLE_SPARSE=1)

• UCM patches load successfully (pc_patch + sparse_patch both applied ✅)
• Same error: FlashAttentionImpl.__init__() got an unexpected keyword argument 'layer_idx'
• The layer_idx parameter is passed by vLLM's Qwen2Attention code regardless of sparse config
• UCM's sparse_patch does not patch FlashAttentionImpl.__init__ to accept layer_idx

3. Qwen's custom vLLM branch

• Supports DCA natively but is incompatible with UCM's KVConnector interface

Questions

1. Does UCM's GSA (Generalized Sparse Attention) work with Qwen2.5-14B-Instruct-1M? The sparse_patch loads but FlashAttentionImpl rejects the layer_idx kwarg.

2. Is there a recommended way to use UCM with long-context models (>32K)? We understand that KV cache memory is the bottleneck on single-GPU setups, but even with multi-GPU, what's the recommended configuration?

3. Will UCM v0.18.0 patches include sparse attention support? Currently v0180/ only has pc_patch, while v0110/ has both pc_patch and sparse_patch.

4. Is there a compatibility matrix for UCM sparse + specific models? We'd like to know which models have been validated with UCM's GSA.

Relevant Code References

UCM apply_patch.py (v0.5.0rc2)

# v0180 only has pc_patch, no sparse_patch
case "0.18.0":
    import ucm.integration.vllm.patch.v0180.vllm.pc_patch
    # No sparse_patch available

# v0110 has both
case "0.11.0":
    import ucm.integration.vllm.patch.v0110.vllm.pc_patch
    if ENABLE_SPARSE:
        import ucm.integration.vllm.patch.v0110.vllm.sparse_patch

Error Stack Trace (both vLLM 0.11.0 and 0.18.1)

File ".../vllm/model_executor/models/qwen2.py", line 165, in __init__
    self.attn = attn_cls(
File ".../vllm/attention/layer.py", line 198, in __init__
    self.impl = impl_cls(num_heads, head_size, scale, num_kv_heads, ...
TypeError: FlashAttentionImpl.__init__() got an unexpected keyword argument 'layer_idx'

Current Workaround

• Remove/rename sparse_attention_config.json from model directory
• Add **kwargs to FlashAttentionImpl.__init__() (monkey-patch)
• Use --max-model-len 32768 (limited by single A6000 48GB)
• UCM prefix caching works well at 32K

Expected Behavior

UCM's sparse attention patches should either:

1. Handle the layer_idx parameter in FlashAttentionImpl, or
2. Provide documentation on compatible models/configurations for extended context

Thank you for your work on UCM! Happy to provide more details or test patches.

Alternatives

No response

Additional context

No response