Configuration for sparse attention.
| 185 | |
| 186 | |
| 187 | class BaseSparseAttentionConfig(StrictBaseModel): |
| 188 | """ |
| 189 | Configuration for sparse attention. |
| 190 | """ |
| 191 | seq_len_threshold: Optional[int] = Field( |
| 192 | default=None, |
| 193 | description= |
| 194 | "The sequence length threshold for separating short and long sequences." |
| 195 | ) |
| 196 | |
| 197 | @property |
| 198 | def algorithm(self) -> str: |
| 199 | raise NotImplementedError("Algorithm must be implemented in subclasses") |
| 200 | |
| 201 | @classmethod |
| 202 | def from_dict(cls, data: dict): |
| 203 | # dispatch to the correct sparse attention config |
| 204 | config_classes = { |
| 205 | "rocket": RocketSparseAttentionConfig, |
| 206 | "dsa": DeepSeekSparseAttentionConfig, |
| 207 | "skip_softmax": SkipSoftmaxAttentionConfig, |
| 208 | } |
| 209 | |
| 210 | algorithm = data.get("algorithm", None) |
| 211 | if algorithm is None: |
| 212 | raise ValueError(f"Sparse attention algorithm is required") |
| 213 | |
| 214 | config_class = config_classes.get(algorithm.lower()) |
| 215 | if config_class is None: |
| 216 | raise ValueError(f"Invalid algorithm: {algorithm}") |
| 217 | |
| 218 | # Remove 'algorithm' before passing to subclass constructor |
| 219 | # It's a ClassVar in subclasses, and used for dispatching to the correct subclass |
| 220 | data = {k: v for k, v in data.items() if k != 'algorithm'} |
| 221 | return config_class(**data) |
| 222 | |
| 223 | def _check_fields(self): |
| 224 | pass |
| 225 | |
| 226 | def supports_backend(self, backend: str) -> bool: |
| 227 | """ |
| 228 | Override if the speculation algorithm does not support |
| 229 | a subset of the possible backends. |
| 230 | """ |
| 231 | return True |
| 232 | |
| 233 | def get_indices_block_size(self) -> int: |
| 234 | return 1 |
| 235 | |
| 236 | def needs_separate_short_long_cuda_graphs(self) -> bool: |
| 237 | """ |
| 238 | Determines whether to capture a dedicated CUDA graph for batches consisting entirely of short sequences. |
| 239 | If True, capture distinct graphs for short-only batches and general cases (e.g., long or mixed batches). |
| 240 | If False, capture a single unified CUDA graph for all sequences regardless of length. |
| 241 | The seq_len_threshold parameter defines the cutoff boundary between short and long sequences. |
| 242 | """ |
| 243 | return False |
| 244 |