ac1ccbceaa
Integrate COMPASS XAttention algorithm into nano-vllm's CPU offload execution path. Uses FlashAttention with native GQA support for offload mode. New files: - nanovllm/kvcache/sparse/utils.py: find_blocks_chunked() utility - nanovllm/kvcache/sparse/kernels.py: Triton kernels for XAttention - nanovllm/kvcache/sparse/xattn.py: XAttentionPolicy implementation Modified: - nanovllm/config.py: Add XATTN configuration parameters - nanovllm/engine/model_runner.py: Support XATTN policy - nanovllm/kvcache/sparse/__init__.py: Register XAttentionPolicy - tests/test_ruler.py: Add --sparse-policy parameter Test results (32k ruler): - NIAH tasks: 12/12 (100%) - QA/Recall tasks: 11/15 (73%) - Overall: 23/27 (85%) Co-Authored-By: Claude <noreply@anthropic.com>
96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
"""
|
|
Sparse Attention Policy module.
|
|
|
|
Provides pluggable policies for selecting which KV blocks to load
|
|
during chunked attention with CPU offload.
|
|
|
|
Usage:
|
|
from nanovllm.kvcache.sparse import create_sparse_policy, SparsePolicyType
|
|
|
|
# Create policy using factory function
|
|
policy = create_sparse_policy(SparsePolicyType.QUEST, topk_blocks=8)
|
|
|
|
# Or create custom policy
|
|
class MyPolicy(SparsePolicy):
|
|
supports_prefill = True
|
|
supports_decode = True
|
|
|
|
def select_blocks(self, available_blocks, ctx):
|
|
return available_blocks[:5] # Just first 5 blocks
|
|
"""
|
|
|
|
from nanovllm.config import SparsePolicyType
|
|
from nanovllm.kvcache.sparse.policy import SparsePolicy, PolicyContext
|
|
from nanovllm.kvcache.sparse.full_policy import FullAttentionPolicy
|
|
from nanovllm.kvcache.sparse.quest import QuestPolicy, QuestConfig, BlockMetadataManager
|
|
from nanovllm.kvcache.sparse.minference import MInferencePolicy
|
|
from nanovllm.kvcache.sparse.xattn import XAttentionPolicy
|
|
|
|
|
|
def create_sparse_policy(policy_type: SparsePolicyType, **kwargs) -> SparsePolicy:
|
|
"""
|
|
Create a sparse policy instance from an enum type.
|
|
|
|
The returned policy is not yet initialized. Call policy.initialize()
|
|
or let the framework call it during KV cache allocation.
|
|
|
|
Args:
|
|
policy_type: SparsePolicyType enum value
|
|
**kwargs: Policy-specific configuration options
|
|
|
|
Returns:
|
|
SparsePolicy instance (not initialized)
|
|
|
|
Example:
|
|
policy = create_sparse_policy(SparsePolicyType.QUEST, topk_blocks=4)
|
|
policy.initialize(num_layers=28, num_kv_heads=8, ...)
|
|
"""
|
|
if policy_type == SparsePolicyType.FULL:
|
|
return FullAttentionPolicy()
|
|
|
|
elif policy_type == SparsePolicyType.QUEST:
|
|
config = QuestConfig(
|
|
topk_blocks=kwargs.get("topk_blocks", 8),
|
|
threshold_blocks=kwargs.get("threshold_blocks", 4),
|
|
include_sink_blocks=kwargs.get("include_sink_blocks", 0),
|
|
include_recent_blocks=kwargs.get("include_recent_blocks", 0),
|
|
)
|
|
return QuestPolicy(config)
|
|
|
|
elif policy_type == SparsePolicyType.MINFERENCE:
|
|
return MInferencePolicy(
|
|
vertical_size=kwargs.get("vertical_size", 1000),
|
|
slash_size=kwargs.get("slash_size", 6096),
|
|
adaptive_budget=kwargs.get("adaptive_budget", 0.3),
|
|
num_sink_tokens=kwargs.get("num_sink_tokens", 30),
|
|
num_recent_diags=kwargs.get("num_recent_diags", 100),
|
|
)
|
|
|
|
elif policy_type == SparsePolicyType.XATTN:
|
|
return XAttentionPolicy(
|
|
stride=kwargs.get("stride", 8),
|
|
threshold=kwargs.get("threshold", 0.9),
|
|
chunk_size=kwargs.get("chunk_size", 16384),
|
|
use_triton=kwargs.get("use_triton", True),
|
|
keep_sink=kwargs.get("keep_sink", False),
|
|
keep_recent=kwargs.get("keep_recent", False),
|
|
norm=kwargs.get("norm", 1.0),
|
|
)
|
|
|
|
else:
|
|
raise ValueError(f"Unknown policy type: {policy_type}")
|
|
|
|
|
|
__all__ = [
|
|
"SparsePolicy",
|
|
"PolicyContext",
|
|
"SparsePolicyType",
|
|
"FullAttentionPolicy",
|
|
"QuestPolicy",
|
|
"QuestConfig",
|
|
"BlockMetadataManager",
|
|
"MInferencePolicy",
|
|
"XAttentionPolicy",
|
|
"create_sparse_policy",
|
|
]
|