Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / shared_compressed_dsa_attention.py: 69%
725 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""High-performance shared compressed DeepSeek Sparse Attention.
17This module implements the DeepSeek-V4.1 CSA2 training contract: Full layers
18publish compressed K=V and index K, Reindex layers produce fresh Top-K indices,
19and Reuse layers consume the latest selection. The decoder's first Full layer
20can additionally publish the blockwise candidate pool used by later Reindex
21layers. The attention replacement selects eager PyTorch references or fused
22Indexer, sparse attention and external-teacher KL operators while preserving
23the same sharing and parallel contexts. Explicit fusion rejects unsupported
24calls instead of falling back to another implementation.
25"""
27from __future__ import annotations
29from collections.abc import Callable, Mapping
30from dataclasses import dataclass, field
31from typing import Any, Protocol
33import torch # pylint: disable=forbidden-backend-import
34from torch import nn # pylint: disable=forbidden-backend-import
36from hyper_parallel.components.functional.aux_loss import aux_loss_auto_scale
37from hyper_parallel.components.functional.lightning_indexer import (
38 lightning_indexer, sort_key_indices, validate_indexer_tensors,
39)
40from hyper_parallel.components.functional.sparse_flash_mla import (
41 supports_fused_attention_inputs, sparse_flash_mla,
42)
43from hyper_parallel.components.functional.sparse_lightning_indexer_kl_loss_grad import (
44 gather_selected_keys, sparse_lightning_indexer_kl_loss, validate_fused_kl_inputs,
45)
46from hyper_parallel.core.activation_memory.wrapper import checkpoint_exclude_wrapper
47from hyper_parallel.core.activation_memory.recompute_state import is_recomputing
48from hyper_parallel.models.replacement import module_replacement
51class SequenceGatherHandle(Protocol):
52 """Handle returned by an asynchronous sequence all-gather."""
54 def wait(self) -> torch.Tensor:
55 """Wait for communication and return the differentiable result."""
58@dataclass
59class _DeferredSequenceGather:
60 """Synchronous fallback with the same interface as an async gather."""
62 tensor: torch.Tensor
63 sequence_dim: int
64 gather_sequence: Callable[[torch.Tensor, int], torch.Tensor]
66 def wait(self) -> torch.Tensor:
67 """Run the fallback gather at the first consumer boundary."""
68 return self.gather_sequence(self.tensor, self.sequence_dim)
71@dataclass
72class SharedCompressedAttentionState:
73 """Autograd-bearing CSA2 state scoped to one training forward.
75 State is addressed by the layer that produced it instead of using one
76 mutable "latest value" slot. This matters for the released V4.1 topology,
77 which has several Full/Reindex groups, and makes a consumer's dependency
78 explicit when activation recomputation revisits modules out of forward
79 order.
81 Nothing drops an entry early. The state travels into every layer as one of
82 the checkpoint region's inputs, so the graph keeps it alive from the forward
83 pass to the end of backward, and a layer whose forward is recomputed reads
84 its entries a second time there. A published entry therefore has to stay
85 readable until backward is done with every consumer, which is also why the
86 publishers skip their write while recomputing.
87 """
89 compressed_kv_by_source: dict[int, torch.Tensor] = field(default_factory=dict)
90 index_key_by_source: dict[int, torch.Tensor] = field(default_factory=dict)
91 topk_indices_by_source: dict[int, torch.Tensor] = field(default_factory=dict)
92 candidate_blocks_by_source: dict[int, torch.Tensor] = field(default_factory=dict)
94 @staticmethod
95 def _require(
96 values: dict[int, torch.Tensor],
97 source_layer: int | None,
98 value_name: str,
99 consumer_layer: int,
100 ) -> torch.Tensor:
101 """Return one published tensor or report the broken layer dependency."""
102 if source_layer is None or source_layer not in values:
103 raise RuntimeError(
104 f"layer {consumer_layer} requires {value_name} from source "
105 f"layer {source_layer}, but that source has not run in this forward"
106 )
107 return values[source_layer]
109 def publish_compressed_kv(self, source_layer: int, value: torch.Tensor) -> None:
110 """Publish compressed K=V while preserving its autograd graph."""
111 self.compressed_kv_by_source[source_layer] = value
113 def require_compressed_kv(self, source_layer: int | None, consumer_layer: int) -> torch.Tensor:
114 """Read compressed K=V from the consumer's configured Full layer."""
115 return self._require(
116 self.compressed_kv_by_source,
117 source_layer,
118 "compressed K=V",
119 consumer_layer,
120 )
122 def publish_index_key(self, source_layer: int, value: torch.Tensor) -> None:
123 """Publish the Full Indexer's shared key tensor."""
124 self.index_key_by_source[source_layer] = value
126 def require_index_key(self, source_layer: int | None, consumer_layer: int) -> torch.Tensor:
127 """Read the Indexer key associated with a compressed-KV source."""
128 return self._require(
129 self.index_key_by_source,
130 source_layer,
131 "Indexer key",
132 consumer_layer,
133 )
135 def publish_topk_indices(self, source_layer: int, value: torch.Tensor) -> None:
136 """Publish one Full/Reindex layer's token selection."""
137 self.topk_indices_by_source[source_layer] = value
139 def require_topk_indices(self, source_layer: int | None, consumer_layer: int) -> torch.Tensor:
140 """Read the latest configured Full/Reindex selection."""
141 return self._require(
142 self.topk_indices_by_source,
143 source_layer,
144 "Top-K indices",
145 consumer_layer,
146 )
148 def publish_candidate_blocks(self, source_layer: int, value: torch.Tensor) -> None:
149 """Publish hierarchical candidate blocks from the configured source."""
150 self.candidate_blocks_by_source[source_layer] = value
152 def require_candidate_blocks(
153 self,
154 source_layer: int | None,
155 consumer_layer: int,
156 ) -> torch.Tensor:
157 """Read hierarchical candidates for a later Reindex layer."""
158 return self._require(
159 self.candidate_blocks_by_source,
160 source_layer,
161 "candidate blocks",
162 consumer_layer,
163 )
166@dataclass
167class SharedCompressedIndexerOutput:
168 """Tensors produced by one Full or Reindex Indexer invocation."""
170 topk_indices: torch.Tensor
171 index_key: torch.Tensor
172 candidate_blocks: torch.Tensor | None
173 index_query: torch.Tensor
174 merge_weight: torch.Tensor
177@dataclass(frozen=True)
178class SharedCompressedAttentionCPContext:
179 """KV-all-gather context for shared compressed DSA.
181 Query-side tensors remain local sequence shards. Raw KV, compressed KV,
182 and index keys are gathered in CP-rank order. ``launch_sequence`` is
183 optional so numerical tests and non-overlapped backends can supply only a
184 differentiable synchronous gather.
185 """
187 size: int
188 rank: int
189 gather_sequence: Callable[[torch.Tensor, int], torch.Tensor]
190 launch_sequence: Callable[[torch.Tensor, int], SequenceGatherHandle] | None = None
192 def launch(self, tensor: torch.Tensor, sequence_dim: int) -> SequenceGatherHandle:
193 """Launch an async gather or create a deferred synchronous gather."""
194 if self.launch_sequence is not None:
195 return self.launch_sequence(tensor, sequence_dim)
196 return _DeferredSequenceGather(tensor, sequence_dim, self.gather_sequence)
199@dataclass(frozen=True)
200class SharedCompressedAttentionTPContext:
201 """Tensor-parallel score reduction context for a head-sharded Indexer."""
203 size: int
204 rank: int
205 reduce_sum: Callable[[torch.Tensor], torch.Tensor]
208@dataclass(frozen=True)
209class SharedCompressedPackedSequence:
210 """Compact sample boundaries for one contiguous CP query shard."""
212 cu_seq_lens: torch.Tensor
213 local_query_start: int
214 local_query_length: int
215 global_sequence_length: int
216 _prepared_starts: dict[torch.device, torch.Tensor] = field(
217 default_factory=dict, init=False, repr=False, compare=False
218 )
219 _prepared_minima: dict[tuple[torch.device, int], torch.Tensor] = field(
220 default_factory=dict, init=False, repr=False, compare=False
221 )
222 _validated_ratios: set[int] = field(default_factory=set, init=False, repr=False, compare=False)
223 _snapshot: dict[str, Any] = field(default_factory=dict, init=False, repr=False, compare=False)
224 _prepared_positions: dict[torch.device, torch.Tensor] = field(
225 default_factory=dict, init=False, repr=False, compare=False
226 )
227 _prepared_masks: dict[torch.device, torch.Tensor] = field(
228 default_factory=dict, init=False, repr=False, compare=False
229 )
230 _indexer_segments: dict[int, tuple[tuple[int, int, int, int, int], ...]] = field(
231 default_factory=dict, init=False, repr=False, compare=False
232 )
234 def prepare(self, device: torch.device, compress_ratios: tuple[int, ...]) -> SharedCompressedPackedSequence:
235 """Prepare an owned boundary snapshot, reused until the input is mutated.
237 Normal in-place tensor updates invalidate the cached snapshot. Existing
238 consumers retain their owned geometry. Prepared tensors are read-only;
239 callers must not mutate boundaries through ``.data`` or external storage.
240 Inference tensors without a version counter are always snapshotted anew.
242 Args:
243 device: Device on which tensors and prepared geometry are consumed.
244 compress_ratios: Compression ratios used by the layers in this forward.
245 """
246 device = torch.device(device)
247 # Torch exposes tensor mutation tracking only through its version counter.
248 version = None if self.cu_seq_lens.is_inference() else self.cu_seq_lens._version # pylint: disable=protected-access
249 if version is not None and self._snapshot.get("version") == version:
250 prepared: SharedCompressedPackedSequence = self if self._snapshot.get("owned") else self._snapshot["value"]
251 else:
252 prepared = SharedCompressedPackedSequence(
253 self.cu_seq_lens.detach().clone(), self.local_query_start,
254 self.local_query_length, self.global_sequence_length,
255 )
256 SharedCompressedPackedSequence._mark_owned(prepared)
257 self._snapshot.clear()
258 self._snapshot.update(version=version, value=prepared)
259 SharedCompressedPackedSequence._prepare_owned_metadata(prepared, device, compress_ratios)
260 return prepared
262 def _mark_owned(self) -> None:
263 """Record the owned tensor's mutation counter for subsequent layer reuse."""
264 if not self.cu_seq_lens.is_inference():
265 # Torch has no public read API for its tensor mutation counter.
266 self._snapshot.update(owned=True, version=self.cu_seq_lens._version) # pylint: disable=protected-access
268 def _prepare_owned_metadata(self, device: torch.device, compress_ratios: tuple[int, ...]) -> None:
269 """Populate caches on the owned snapshot, never on caller-owned boundaries."""
270 starts = self.local_segment_starts(device)
271 self._prepared_starts[device] = starts
272 for ratio in set(compress_ratios):
273 self.validate_compression_alignment(ratio)
274 self._validated_ratios.add(ratio)
275 if ratio > 0 and (device, ratio) not in self._prepared_minima:
276 self._prepared_minima[(device, ratio)] = starts // ratio
278 def local_positions(self, device: torch.device) -> torch.Tensor:
279 """Return the immutable local interval's global token positions."""
280 device = torch.device(device)
281 if device not in self._prepared_positions:
282 self._prepared_positions[device] = torch.arange(
283 self.local_query_start, self.local_query_start + self.local_query_length, device=device,
284 ).unsqueeze(0)
285 return self._prepared_positions[device]
287 def segment_start_mask(self, device: torch.device) -> torch.Tensor:
288 """Reuse the per-batch sample-start mask across model forwards."""
289 device = torch.device(device)
290 if device not in self._prepared_masks:
291 self._prepared_masks[device] = self.local_positions(device) == self.local_segment_starts(device)
292 return self._prepared_masks[device]
294 def indexer_segments(self, ratio: int) -> tuple[tuple[int, int, int, int, int], ...]:
295 """Describe packed Q intervals and completed K prefixes without per-layer synchronization."""
296 if ratio not in self._indexer_segments:
297 self.validate_compression_alignment(ratio)
298 boundaries = self.cu_seq_lens.detach().cpu().tolist()
299 local_start = self.local_query_start
300 local_end = local_start + self.local_query_length
301 segments = []
302 for sample_start, sample_end in zip(boundaries, boundaries[1:]):
303 start, end = max(local_start, sample_start), min(local_end, sample_end)
304 if start < end:
305 segments.append((start-local_start, end-local_start, sample_start//ratio, end//ratio,
306 (end-sample_start) % ratio))
307 self._indexer_segments[ratio] = tuple(segments)
308 return self._indexer_segments[ratio]
310 def minimum_key_indices(
311 self, device: torch.device, compress_ratio: int, segment_starts: torch.Tensor | None = None,
312 ) -> torch.Tensor:
313 """Return the first compressed key belonging to each query's sample.
315 Args:
316 device: Device on which tensors and prepared geometry are consumed.
317 compress_ratio: Compression ratio of the consuming layer.
318 segment_starts: Optional precomputed global sample starts for local queries.
319 """
320 if compress_ratio <= 0:
321 raise ValueError("minimum_key_indices requires a positive compression ratio")
322 cached = self._prepared_minima.get((torch.device(device), compress_ratio))
323 if cached is not None:
324 return cached
325 self.validate_compression_alignment(compress_ratio)
326 if segment_starts is None:
327 segment_starts = self.local_segment_starts(device)
328 return segment_starts // compress_ratio
330 def local_segment_starts(self, device: torch.device) -> torch.Tensor:
331 """Return each local query's global packed-sample start position.
333 Args:
334 device: Device on which tensors and prepared geometry are consumed.
335 """
336 cached = self._prepared_starts.get(torch.device(device))
337 if cached is not None:
338 return cached
339 if self.local_query_start < 0 or self.local_query_length < 0 or (
340 self.local_query_start + self.local_query_length > self.global_sequence_length
341 ):
342 raise ValueError("packed local query interval must lie within the global sequence")
343 boundaries = self.cu_seq_lens.to(dtype=torch.long)
344 if boundaries.ndim != 1 or boundaries.numel() < 2 or int(boundaries[0]) != 0:
345 raise ValueError("packed cu_seq_lens must be one-dimensional and start with zero")
346 if int(boundaries[-1]) != self.global_sequence_length:
347 raise ValueError(
348 "packed cu_seq_lens must cover global_sequence_length: "
349 f"{int(boundaries[-1])} versus {self.global_sequence_length}"
350 )
351 if torch.any(boundaries[1:] <= boundaries[:-1]):
352 raise ValueError("packed cu_seq_lens must be strictly increasing")
353 boundaries = boundaries.to(device=device)
354 query_positions = self.local_positions(device).squeeze(0)
355 segment_ids = torch.bucketize(query_positions, boundaries[1:], right=True)
356 return boundaries[:-1].index_select(0, segment_ids).unsqueeze(0)
358 def validate_compression_alignment(self, compress_ratio: int) -> None:
359 """Reject packed samples whose compressor groups would cross boundaries.
361 Args:
362 compress_ratio: Compression ratio of the consuming layer.
363 """
364 if compress_ratio <= 1 or compress_ratio in self._validated_ratios:
365 return
366 boundaries = self.cu_seq_lens.to(dtype=torch.long)
367 misaligned = boundaries[boundaries.remainder(compress_ratio) != 0]
368 if misaligned.numel():
369 raise ValueError(
370 "packed sample boundaries must align with the CSA2 compression ratio "
371 f"{compress_ratio}; misaligned boundaries={misaligned.tolist()}"
372 )
375def _rotate_half(tensor: torch.Tensor) -> torch.Tensor:
376 """Rotate adjacent real/imaginary pairs for V4 interleaved RoPE."""
377 first = tensor[..., 0::2]
378 second = tensor[..., 1::2]
379 return torch.stack((-second, first), dim=-1).flatten(-2)
382def _apply_v41_rope(
383 tensor: torch.Tensor,
384 cos: torch.Tensor,
385 sin: torch.Tensor,
386 unsqueeze_dim: int = 1,
387) -> torch.Tensor:
388 """Apply V4 interleaved RoPE to the trailing rotary channels."""
389 cos = cos.repeat_interleave(2, dim=-1).unsqueeze(unsqueeze_dim)
390 sin = sin.repeat_interleave(2, dim=-1).unsqueeze(unsqueeze_dim)
391 rope_dim = cos.shape[-1]
392 pass_through, rotary = tensor[..., :-rope_dim], tensor[..., -rope_dim:]
393 rotary = ((rotary.float() * cos) + (_rotate_half(rotary).float() * sin)).to(tensor.dtype)
394 return torch.cat((pass_through, rotary), dim=-1)
397def build_sliding_window_indices(
398 batch_size: int,
399 sequence_length: int,
400 window_size: int,
401 device: torch.device,
402 *,
403 query_offset: int = 0,
404 key_length: int | None = None,
405) -> torch.Tensor:
406 """Build fixed-width causal sliding-window indices."""
407 if key_length is None:
408 key_length = sequence_length
409 width = min(key_length, window_size)
410 query = torch.arange(
411 query_offset,
412 query_offset + sequence_length,
413 device=device,
414 ).unsqueeze(-1)
415 offsets = torch.arange(width - 1, -1, -1, device=device)
416 indices = query - offsets
417 indices = indices.masked_fill(indices < 0, -1)
418 return indices.unsqueeze(0).expand(batch_size, -1, -1)
421def _run_lightning_indexer(
422 query: torch.Tensor, key: torch.Tensor, merge_weight: torch.Tensor,
423 top_k: int, compress_ratio: int, query_offset: int,
424 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None,
425 minimum_key_indices: torch.Tensor | None, query_segments: tuple | None,
426) -> torch.Tensor:
427 """Run LI V2 with explicit adapter errors and native capability validation."""
428 if reduce_sum is not None:
429 raise NotImplementedError("Fused LightningIndexer does not support TP score reduction before TopK")
430 if minimum_key_indices is not None and query_segments is None:
431 raise ValueError("Fused LightningIndexer requires query_segments when minimum_key_indices is set")
432 validate_indexer_tensors(query, key, merge_weight)
433 if query.device.type != "npu":
434 raise RuntimeError(f"Fused LightningIndexer requires NPU inputs; got device={query.device}")
435 if query_segments is None:
436 length = query.shape[1]
437 raw_end = query_offset + length
438 query_segments = ((0, length, 0, raw_end // compress_ratio, raw_end % compress_ratio),)
439 return lightning_indexer(query, key, merge_weight, compress_ratio, top_k, query_segments)
442@torch.no_grad()
443def compressed_causal_topk(
444 query: torch.Tensor,
445 key: torch.Tensor,
446 merge_weight: torch.Tensor,
447 *,
448 compress_ratio: int,
449 sparse_count: int,
450 query_offset: int = 0,
451 query_chunk_size: int = 256,
452 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None = None,
453 minimum_key_indices: torch.Tensor | None = None,
454 query_segments: tuple[tuple[int, int, int, int, int], ...] | None = None,
455 use_fused: bool = False,
456) -> torch.Tensor:
457 """Select compressed positions with V4.1's ratio-aware causal rule.
459 CSA2 keys represent closed groups of ``compress_ratio`` source tokens.
460 Both LI V2 and the batched reference preserve the source rule
461 ``key < floor((query + 1) / ratio)``. The optional native path requires
462 document-local K prefixes; the reference retains TP score reductions.
464 Args:
465 query_segments: Local Q intersections and global compressed K prefixes.
466 use_fused: Select LI V2 and raise if a nonempty call is unsupported.
467 False explicitly selects the numerical reference.
468 """
469 if compress_ratio <= 0:
470 raise ValueError(f"compress_ratio must be positive, got {compress_ratio}")
471 if query_chunk_size <= 0:
472 raise ValueError(f"query_chunk_size must be positive, got {query_chunk_size}")
473 if query.ndim != 4:
474 raise ValueError(f"Indexer query must have shape [B, Q, H, D], got {tuple(query.shape)}")
475 batch_size, sequence_length, _, head_dim = query.shape
476 if key.ndim != 3 or key.shape[::2] != (batch_size, head_dim):
477 raise ValueError(
478 "compressed index key must have shape [batch, compressed_sequence, head_dim]"
479 )
480 if merge_weight.shape != query.shape[:3]:
481 raise ValueError(
482 f"merge_weight must have shape {tuple(query.shape[:3])}, got {tuple(merge_weight.shape)}"
483 )
484 compressed_length = key.shape[1]
485 if minimum_key_indices is not None and minimum_key_indices.shape != query.shape[:2]:
486 raise ValueError("minimum_key_indices must have shape [batch, query]")
487 top_k = min(sparse_count, compressed_length)
488 if top_k <= 0 or sequence_length == 0:
489 return torch.empty(
490 batch_size,
491 sequence_length,
492 max(top_k, 0),
493 dtype=torch.int32,
494 device=query.device,
495 )
497 if use_fused:
498 return _run_lightning_indexer(
499 query, key, merge_weight, top_k, compress_ratio, query_offset,
500 reduce_sum, minimum_key_indices, query_segments,
501 )
503 key_fp32 = key.float().transpose(1, 2).unsqueeze(1)
504 key_positions = torch.arange(compressed_length, device=query.device).view(1, 1, -1)
505 selected_chunks = []
506 for start in range(0, sequence_length, query_chunk_size):
507 end = min(start + query_chunk_size, sequence_length)
508 # SAC may retain the matmul result; do not mutate its cached tensor.
509 scores = torch.matmul(query[:, start:end].float(), key_fp32).relu()
510 scores = (scores * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
511 if reduce_sum is not None:
512 scores = reduce_sum(scores)
513 visible = (
514 torch.arange(query_offset + start, query_offset + end, device=query.device) + 1
515 ) // compress_ratio
516 scores.masked_fill_(key_positions >= visible.view(1, -1, 1), float("-inf"))
517 if minimum_key_indices is not None:
518 minimum = minimum_key_indices[:, start:end].unsqueeze(-1)
519 scores.masked_fill_(key_positions < minimum, float("-inf"))
520 top = scores.topk(top_k, dim=-1, sorted=False)
521 indices = top.indices.masked_fill(~torch.isfinite(top.values), compressed_length)
522 indices = sort_key_indices(indices, compressed_length)
523 indices = indices.masked_fill(indices == compressed_length, -1)
524 selected_chunks.append(indices.to(torch.int32))
525 return torch.cat(selected_chunks, dim=1)
528def select_candidate_blocks(
529 logits: torch.Tensor,
530 compress_lens: torch.Tensor | int,
531 topk_blocks: int,
532 block_size: int,
533) -> torch.Tensor:
534 """Select the official block-max hierarchical candidate pool.
536 The newest reachable partial block is pinned into the candidate set, then
537 the remaining blocks compete by their maximum position score. ``logits``
538 must already contain ``-inf`` at causally unreachable positions.
539 """
540 if topk_blocks <= 0:
541 raise ValueError(f"topk_blocks must be positive, got {topk_blocks}")
542 if block_size <= 0:
543 raise ValueError(f"block_size must be positive, got {block_size}")
544 width = logits.shape[-1]
545 if width == 0:
546 return torch.zeros_like(logits, dtype=torch.bool)
547 padding = -width % block_size
548 scores = torch.nn.functional.pad(logits, (0, padding), value=float("-inf"))
549 scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1)
550 num_blocks = scores.shape[-1]
551 if isinstance(compress_lens, torch.Tensor):
552 last = (compress_lens.to(device=logits.device) - 1) // block_size
553 else:
554 last = (compress_lens - 1) // block_size
555 block_positions = torch.arange(num_blocks, device=logits.device)
556 scores = scores.masked_fill(block_positions == last, float("inf"))
557 top = scores.topk(min(topk_blocks, num_blocks), dim=-1)
558 keep = torch.zeros_like(scores, dtype=torch.bool)
559 keep.scatter_(-1, top.indices, top.values > float("-inf"))
560 return keep.repeat_interleave(block_size, dim=-1)[..., :width]
563def select_candidate_block_indices(
564 logits: torch.Tensor,
565 compress_lens: torch.Tensor | int,
566 topk_blocks: int,
567 block_size: int,
568) -> torch.Tensor:
569 """Return compact block ids for the official hierarchical candidates.
571 Storing block ids uses ``block_size`` times less cross-layer state than
572 expanded position ids and avoids the full ``[batch, query, key]`` boolean
573 mask used by the released inference reference. Invalid slots are ``-1``.
574 """
575 if topk_blocks <= 0:
576 raise ValueError(f"topk_blocks must be positive, got {topk_blocks}")
577 if block_size <= 0:
578 raise ValueError(f"block_size must be positive, got {block_size}")
579 width = logits.shape[-1]
580 if width == 0:
581 return torch.empty(*logits.shape[:-1], 0, dtype=torch.int32, device=logits.device)
582 padding = -width % block_size
583 scores = torch.nn.functional.pad(logits, (0, padding), value=float("-inf"))
584 scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1)
585 num_blocks = scores.shape[-1]
586 if isinstance(compress_lens, torch.Tensor):
587 last = (compress_lens.to(device=logits.device) - 1) // block_size
588 else:
589 last = (compress_lens - 1) // block_size
590 block_positions = torch.arange(num_blocks, device=logits.device)
591 scores = scores.masked_fill(block_positions == last, float("inf"))
592 top = scores.topk(min(topk_blocks, num_blocks), dim=-1, sorted=False)
593 return top.indices.masked_fill(top.values == float("-inf"), -1).to(torch.int32)
596def compressed_causal_candidates(
597 query: torch.Tensor,
598 key: torch.Tensor,
599 merge_weight: torch.Tensor,
600 *,
601 compress_ratio: int,
602 topk_blocks: int,
603 block_size: int,
604 query_offset: int = 0,
605 query_chunk_size: int = 256,
606 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None = None,
607 minimum_key_indices: torch.Tensor | None = None,
608) -> torch.Tensor:
609 """Build compact blockwise candidate ids in bounded query chunks."""
610 _, sequence_length, _, _ = query.shape
611 compressed_length = key.shape[1]
612 if minimum_key_indices is not None and minimum_key_indices.shape != query.shape[:2]:
613 raise ValueError("minimum_key_indices must have shape [batch, query]")
614 key_fp32 = key.float().transpose(1, 2).unsqueeze(1)
615 key_positions = torch.arange(compressed_length, device=query.device).view(1, 1, -1)
616 candidate_chunks = []
617 for start in range(0, sequence_length, query_chunk_size):
618 end = min(start + query_chunk_size, sequence_length)
619 scores = torch.matmul(query[:, start:end].float(), key_fp32).relu()
620 scores = (scores * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
621 if reduce_sum is not None:
622 scores = reduce_sum(scores)
623 visible = (
624 torch.arange(query_offset + start, query_offset + end, device=query.device) + 1
625 ) // compress_ratio
626 scores.masked_fill_(key_positions >= visible.view(1, -1, 1), float("-inf"))
627 if minimum_key_indices is not None:
628 minimum = minimum_key_indices[:, start:end].unsqueeze(-1)
629 scores.masked_fill_(key_positions < minimum, float("-inf"))
630 candidate_chunks.append(
631 select_candidate_block_indices(
632 scores,
633 visible.view(1, -1, 1),
634 topk_blocks,
635 block_size,
636 )
637 )
638 return torch.cat(candidate_chunks, dim=1)
641def compressed_causal_topk_and_candidates(
642 query: torch.Tensor,
643 key: torch.Tensor,
644 merge_weight: torch.Tensor,
645 *,
646 compress_ratio: int,
647 sparse_count: int,
648 topk_blocks: int,
649 block_size: int,
650 query_offset: int = 0,
651 query_chunk_size: int = 256,
652 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None = None,
653 minimum_key_indices: torch.Tensor | None = None,
654) -> tuple[torch.Tensor, torch.Tensor]:
655 """Select Full-layer Top-K and candidate blocks from one score pass."""
656 if compress_ratio <= 0:
657 raise ValueError(f"compress_ratio must be positive, got {compress_ratio}")
658 if query_chunk_size <= 0:
659 raise ValueError(f"query_chunk_size must be positive, got {query_chunk_size}")
660 batch_size, sequence_length, _, head_dim = query.shape
661 if key.ndim != 3 or key.shape[::2] != (batch_size, head_dim):
662 raise ValueError(
663 "compressed index key must have shape [batch, compressed_sequence, head_dim]"
664 )
665 if merge_weight.shape != query.shape[:3]:
666 raise ValueError(
667 f"merge_weight must have shape {tuple(query.shape[:3])}, got {tuple(merge_weight.shape)}"
668 )
669 compressed_length = key.shape[1]
670 if minimum_key_indices is not None and minimum_key_indices.shape != query.shape[:2]:
671 raise ValueError("minimum_key_indices must have shape [batch, query]")
672 top_k = min(sparse_count, compressed_length)
673 if top_k <= 0:
674 empty_topk = torch.empty(
675 batch_size,
676 sequence_length,
677 0,
678 dtype=torch.int32,
679 device=query.device,
680 )
681 return empty_topk, empty_topk
683 key_fp32 = key.float().transpose(1, 2).unsqueeze(1)
684 key_positions = torch.arange(compressed_length, device=query.device).view(1, 1, -1)
685 selected_chunks = []
686 candidate_chunks = []
687 for start in range(0, sequence_length, query_chunk_size):
688 end = min(start + query_chunk_size, sequence_length)
689 scores = torch.matmul(query[:, start:end].float(), key_fp32).relu()
690 scores = (scores * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
691 if reduce_sum is not None:
692 scores = reduce_sum(scores)
693 visible = (
694 torch.arange(query_offset + start, query_offset + end, device=query.device) + 1
695 ) // compress_ratio
696 scores.masked_fill_(key_positions >= visible.view(1, -1, 1), float("-inf"))
697 if minimum_key_indices is not None:
698 minimum = minimum_key_indices[:, start:end].unsqueeze(-1)
699 scores.masked_fill_(key_positions < minimum, float("-inf"))
700 top = scores.topk(top_k, dim=-1, sorted=False)
701 indices = top.indices.masked_fill(~torch.isfinite(top.values), compressed_length)
702 indices = indices.sort(dim=-1).values
703 selected_chunks.append(indices.masked_fill(indices == compressed_length, -1).to(torch.int32))
704 candidate_chunks.append(
705 select_candidate_block_indices(
706 scores,
707 visible.view(1, -1, 1),
708 topk_blocks,
709 block_size,
710 )
711 )
712 return torch.cat(selected_chunks, dim=1), torch.cat(candidate_chunks, dim=1)
715def compressed_candidate_topk(
716 query: torch.Tensor,
717 key: torch.Tensor,
718 merge_weight: torch.Tensor,
719 candidate_blocks: torch.Tensor,
720 *,
721 compress_ratio: int,
722 sparse_count: int,
723 block_size: int,
724 query_offset: int = 0,
725 query_chunk_size: int = 256,
726 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None = None,
727 minimum_key_indices: torch.Tensor | None = None,
728) -> torch.Tensor:
729 """Score only hierarchical candidate blocks and return global key ids."""
730 if compress_ratio <= 0:
731 raise ValueError(f"compress_ratio must be positive, got {compress_ratio}")
732 if block_size <= 0:
733 raise ValueError(f"block_size must be positive, got {block_size}")
734 if query_chunk_size <= 0:
735 raise ValueError(f"query_chunk_size must be positive, got {query_chunk_size}")
736 batch_size, sequence_length, _, head_dim = query.shape
737 if key.ndim != 3 or key.shape[::2] != (batch_size, head_dim):
738 raise ValueError(
739 "compressed index key must have shape [batch, compressed_sequence, head_dim]"
740 )
741 if merge_weight.shape != query.shape[:3]:
742 raise ValueError(
743 f"merge_weight must have shape {tuple(query.shape[:3])}, got {tuple(merge_weight.shape)}"
744 )
745 if candidate_blocks.ndim != 3 or candidate_blocks.shape[:2] != query.shape[:2]:
746 raise ValueError(
747 "candidate_blocks must have shape [batch, query, candidate_blocks], "
748 f"got {tuple(candidate_blocks.shape)}"
749 )
750 if minimum_key_indices is not None and minimum_key_indices.shape != query.shape[:2]:
751 raise ValueError("minimum_key_indices must have shape [batch, query]")
753 compressed_length = key.shape[1]
754 candidate_width = candidate_blocks.shape[-1] * block_size
755 top_k = min(sparse_count, candidate_width)
756 if top_k <= 0:
757 return torch.empty(
758 batch_size,
759 sequence_length,
760 0,
761 dtype=torch.int32,
762 device=query.device,
763 )
765 block_offsets = torch.arange(block_size, device=query.device)
766 batch_indices = torch.arange(batch_size, device=query.device).view(-1, 1, 1)
767 selected_chunks = []
768 for start in range(0, sequence_length, query_chunk_size):
769 end = min(start + query_chunk_size, sequence_length)
770 blocks = candidate_blocks[:, start:end].long()
771 block_positions = blocks.unsqueeze(-1) * block_size + block_offsets
772 positions = block_positions.flatten(-2)
773 visible = (
774 torch.arange(query_offset + start, query_offset + end, device=query.device) + 1
775 ) // compress_ratio
776 valid = (blocks.unsqueeze(-1) >= 0).expand_as(block_positions).flatten(-2)
777 valid = valid & (positions < compressed_length) & (positions < visible.view(1, -1, 1))
778 if minimum_key_indices is not None:
779 valid = valid & (positions >= minimum_key_indices[:, start:end].unsqueeze(-1))
780 safe_positions = positions.clamp(min=0, max=max(compressed_length - 1, 0))
781 selected_key = key[batch_indices, safe_positions].float()
782 dots = torch.einsum(
783 "bchd,bckd->bchk",
784 query[:, start:end].float(),
785 selected_key,
786 ).relu()
787 scores = (dots * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
788 if reduce_sum is not None:
789 scores = reduce_sum(scores)
790 scores.masked_fill_(~valid, float("-inf"))
791 top = scores.topk(top_k, dim=-1, sorted=False)
792 selected = positions.gather(-1, top.indices)
793 selected.masked_fill_(~torch.isfinite(top.values), compressed_length)
794 selected = selected.sort(dim=-1).values
795 selected.masked_fill_(selected == compressed_length, -1)
796 selected_chunks.append(selected.to(torch.int32))
797 return torch.cat(selected_chunks, dim=1)
800class _SharedCompressedIndexerKLLoss(torch.autograd.Function):
801 """PanGu-style precomputed-gradient KL loss for compressed index keys."""
803 @staticmethod
804 def forward(
805 ctx: Any,
806 index_query: torch.Tensor,
807 index_key: torch.Tensor,
808 merge_weight: torch.Tensor,
809 attention_query: torch.Tensor,
810 compressed_key: torch.Tensor,
811 topk_indices: torch.Tensor,
812 sinks: torch.Tensor,
813 attention_scale: float,
814 loss_coeff: float,
815 query_chunk_size: int,
816 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None,
817 ) -> torch.Tensor:
818 """Compute sparse KL and save only the three Indexer gradients."""
819 batch_size, sequence_length, _, _ = index_query.shape
820 denominator = batch_size * sequence_length
821 total_loss = index_query.new_zeros((), dtype=torch.float32)
822 grad_index_query = torch.zeros_like(index_query, dtype=torch.float32)
823 grad_index_key = torch.zeros_like(index_key, dtype=torch.float32)
824 grad_merge_weight = torch.zeros_like(merge_weight, dtype=torch.float32)
826 with torch.no_grad(), torch.autocast(device_type=index_query.device.type, enabled=False):
827 for start in range(0, sequence_length, query_chunk_size):
828 end = min(start + query_chunk_size, sequence_length)
829 selected = topk_indices[:, start:end]
830 valid = selected >= 0
831 valid_rows = valid.any(dim=-1)
832 if not torch.any(valid_rows):
833 continue
834 safe_indices = selected.clamp_min(0).long()
836 query_chunk = index_query[:, start:end].float()
837 weight_chunk = merge_weight[:, start:end].float()
838 selected_index_key = gather_selected_keys(index_key, safe_indices).float()
839 index_dots = torch.einsum(
840 "bcid,bckd->bcik", query_chunk, selected_index_key
841 )
842 index_relu = index_dots.relu()
843 index_scores = (index_relu * weight_chunk.unsqueeze(-1)).sum(dim=2)
844 if reduce_sum is not None:
845 index_scores = reduce_sum(index_scores)
846 index_scores.masked_fill_(~valid, -1.0e9)
848 selected_attention_key = gather_selected_keys(compressed_key, safe_indices).float()
849 attention_scores = torch.einsum(
850 "bhcd,bckd->bhck",
851 attention_query[:, :, start:end].float(),
852 selected_attention_key,
853 ) * attention_scale
854 attention_scores.masked_fill_(~valid.unsqueeze(1), -1.0e9)
855 sink_logits = sinks.float().view(1, -1, 1, 1).expand(
856 batch_size, -1, end - start, -1
857 )
858 target = torch.cat((attention_scores, sink_logits), dim=-1).softmax(dim=-1)
859 target = target[..., :-1].masked_fill(~valid.unsqueeze(1), 0.0).sum(dim=1)
860 if reduce_sum is not None:
861 target = reduce_sum(target)
862 target = target / target.sum(dim=-1, keepdim=True).clamp_min(
863 torch.finfo(torch.float32).tiny
864 )
866 log_prediction = index_scores.log_softmax(dim=-1)
867 target_log = target.clamp_min(torch.finfo(torch.float32).tiny).log()
868 row_loss = (target * (target_log - log_prediction)).sum(dim=-1)
869 total_loss.add_(row_loss[valid_rows].sum() * (loss_coeff / denominator))
871 grad_scores = log_prediction.exp() * target.sum(-1, keepdim=True) - target
872 grad_scores.mul_(loss_coeff / denominator)
873 grad_scores.masked_fill_(~valid, 0.0)
874 grad_dots = (
875 grad_scores.unsqueeze(2)
876 * weight_chunk.unsqueeze(-1)
877 * (index_dots > 0).to(index_dots.dtype)
878 )
879 grad_index_query[:, start:end] = torch.einsum(
880 "bcik,bckd->bcid", grad_dots, selected_index_key
881 )
882 grad_merge_weight[:, start:end] = (
883 grad_scores.unsqueeze(2) * index_relu
884 ).sum(dim=-1)
885 selected_key_grad = torch.einsum(
886 "bcik,bcid->bckd", grad_dots, query_chunk
887 )
888 # Preserve the einsum result if SAC cached its storage.
889 selected_key_grad = selected_key_grad.masked_fill(~valid.unsqueeze(-1), 0.0)
890 scatter_indices = safe_indices.unsqueeze(-1).expand_as(selected_key_grad)
891 grad_index_key.scatter_add_(
892 1,
893 scatter_indices.reshape(batch_size, -1, index_key.shape[-1]),
894 selected_key_grad.reshape(batch_size, -1, index_key.shape[-1]),
895 )
897 ctx.save_for_backward(grad_index_query, grad_index_key, grad_merge_weight)
898 return total_loss
900 @staticmethod
901 def backward(ctx: Any, grad_loss: torch.Tensor) -> tuple:
902 """Return the precomputed Indexer gradients and detach the teacher."""
903 grad_index_query, grad_index_key, grad_merge_weight = ctx.saved_tensors
904 return (
905 grad_index_query * grad_loss,
906 grad_index_key * grad_loss,
907 grad_merge_weight * grad_loss,
908 *((None,) * 8),
909 )
912def shared_compressed_indexer_kl_loss(
913 index_query: torch.Tensor,
914 index_key: torch.Tensor,
915 merge_weight: torch.Tensor,
916 attention_query: torch.Tensor,
917 compressed_key: torch.Tensor,
918 topk_indices: torch.Tensor,
919 sinks: torch.Tensor,
920 *,
921 attention_scale: float,
922 loss_coeff: float,
923 query_chunk_size: int = 256,
924 tp_context: SharedCompressedAttentionTPContext | None = None,
925 compress_ratio: int = 1,
926 query_segments: tuple[tuple[int, int, int, int, int], ...] | None = None,
927 use_fused: bool = False,
928) -> torch.Tensor:
929 """Compute the sparse-stage DSA KL objective on V4.1 compressed keys.
931 The target first softmaxes selected main-attention scores per attention
932 head, then averages and L1-normalizes across heads. Indexer inputs are
933 detached by the caller, so this objective updates only Indexer parameters.
935 ``use_fused`` is valid only for ordinary causal TopK selections, without
936 candidate-pool holes. ``query_segments`` supplies each document-local K
937 prefix; its ratio and residual encode the CP query offset for the kernel.
938 ``use_fused=True`` raises for unsupported calls when ``loss_coeff != 0``;
939 False explicitly selects the numerical reference. Native errors propagate.
940 """
941 if query_chunk_size <= 0:
942 raise ValueError(f"query_chunk_size must be positive, got {query_chunk_size}")
943 if not loss_coeff:
944 return index_query.sum(dtype=torch.float32) * 0.0
945 if topk_indices.shape[:2] != index_query.shape[:2]:
946 raise ValueError("topk_indices and index_query must share batch/query dimensions")
947 if use_fused:
948 if tp_context is not None:
949 raise NotImplementedError("Fused KL does not support TP score/teacher reduction")
950 if query_segments is None:
951 raise ValueError("Fused KL requires document-local causal query_segments")
952 validate_fused_kl_inputs(index_query, index_key, merge_weight, topk_indices)
953 return sparse_lightning_indexer_kl_loss(
954 index_query, index_key, merge_weight, attention_query.detach(), compressed_key.detach(),
955 topk_indices, sinks.detach(), attention_scale, loss_coeff, query_chunk_size,
956 compress_ratio, query_segments,
957 )
958 return _SharedCompressedIndexerKLLoss.apply(
959 index_query,
960 index_key,
961 merge_weight,
962 attention_query.detach(),
963 compressed_key.detach(),
964 topk_indices,
965 sinks.detach(),
966 attention_scale,
967 loss_coeff,
968 query_chunk_size,
969 None if tp_context is None else tp_context.reduce_sum,
970 )
973class SharedCompressedDSAIndexer(nn.Module):
974 """Reusable V4.1 compressed Lightning Indexer semantics."""
976 # Installed by :meth:`enable_selection_exclude` for a layer that recomputes.
977 # It is written through ``object.__setattr__`` so the exclusion region stays
978 # out of ``modules()`` and the state dict.
979 _selection_region: Callable[..., Any] | None = None
981 def __init__(self, module: nn.Module) -> None:
982 """Transfer the source projections without changing checkpoint names."""
983 super().__init__()
984 required = ("q_b_proj", "weights_proj")
985 missing = [name for name in required if not hasattr(module, name)]
986 if missing:
987 raise TypeError(f"shared compressed DSA indexer is missing: {missing}")
988 for name, child in module._modules.items(): # pylint: disable=protected-access
989 self.add_module(name, child)
990 for name, parameter in module._parameters.items(): # pylint: disable=protected-access
991 self.register_parameter(name, parameter)
992 self.use_fused = False
993 self.compress_ratio = module.compress_ratio
994 self.global_num_heads = module.num_heads
995 self.head_dim = module.head_dim
996 self.index_topk = module.index_topk
997 self.owns_key = module.owns_key
998 self.is_candidate_source = module.is_candidate_source
999 self.uses_candidates = module.uses_candidates
1000 self.candidate_topk_blocks = module.candidate_topk_blocks
1001 self.candidate_block_size = module.candidate_block_size
1002 self.loss_coeff = module.loss_coeff
1003 self.query_chunk_size = int(getattr(module, "query_chunk_size", 256))
1004 self.train(module.training)
1006 def _project_index_key(
1007 self,
1008 latent: torch.Tensor | None,
1009 cos: torch.Tensor,
1010 sin: torch.Tensor,
1011 index_key: torch.Tensor | None,
1012 cp_context: SharedCompressedAttentionCPContext | None,
1013 ) -> tuple[torch.Tensor, "SequenceGatherHandle | None"]:
1014 """Project or adopt the index key and launch its CP gather."""
1015 if self.owns_key:
1016 if latent is None:
1017 raise RuntimeError("a Full Indexer requires the compressor latent")
1018 key = self.k_norm(self.wk(latent.detach()))
1019 compressed_length = key.shape[1]
1020 key_cos = cos[:, :compressed_length * self.compress_ratio:self.compress_ratio]
1021 key_sin = sin[:, :compressed_length * self.compress_ratio:self.compress_ratio]
1022 key = _apply_v41_rope(key.unsqueeze(1), key_cos, key_sin).squeeze(1)
1023 key_handle = cp_context.launch(key, 1) if cp_context is not None else None
1024 return key, key_handle
1025 if index_key is None:
1026 raise RuntimeError("a Reindex layer requires index K from the preceding Full layer")
1027 return index_key, None
1029 def forward(
1030 self,
1031 hidden_states: torch.Tensor,
1032 query_residual: torch.Tensor,
1033 latent: torch.Tensor | None,
1034 compress_position_embeddings: tuple[torch.Tensor, torch.Tensor],
1035 *,
1036 index_key: torch.Tensor | None = None,
1037 candidate_blocks: torch.Tensor | None = None,
1038 cp_context: SharedCompressedAttentionCPContext | None = None,
1039 tp_context: SharedCompressedAttentionTPContext | None = None,
1040 query_offset: int = 0,
1041 minimum_key_indices: torch.Tensor | None = None,
1042 query_segments: tuple[tuple[int, int, int, int, int], ...] | None = None,
1043 ) -> SharedCompressedIndexerOutput:
1044 """Project local queries and score them against global compressed keys.
1046 When the enclosing layer is recomputed, :meth:`enable_selection_exclude`
1047 makes the discrete selection a checkpoint exclusion region: it runs in
1048 the original forward and its cached result is reused during the
1049 recomputation, so the chosen key set can never drift and the checkpoint
1050 keeps saving the same tensors in both passes.
1051 """
1052 batch_size, sequence_length, _ = hidden_states.shape
1053 cos, sin = compress_position_embeddings
1055 key, key_handle = self._project_index_key(latent, cos, sin, index_key, cp_context)
1057 query = self.q_b_proj(query_residual.detach())
1058 if query.shape[-1] % self.head_dim:
1059 raise ValueError(
1060 f"Indexer query width {query.shape[-1]} is not divisible by head_dim {self.head_dim}"
1061 )
1062 local_num_heads = query.shape[-1] // self.head_dim
1063 query = query.view(batch_size, sequence_length, local_num_heads, self.head_dim)
1064 query = _apply_v41_rope(query.transpose(1, 2), cos, sin).transpose(1, 2)
1065 merge_weight = self.weights_proj(hidden_states.detach())
1066 if merge_weight.shape[-1] != local_num_heads:
1067 raise ValueError(
1068 "Indexer q_b_proj and weights_proj must use the same TP head shard: "
1069 f"{local_num_heads} query heads versus {merge_weight.shape[-1]} weights"
1070 )
1071 merge_weight = merge_weight * (self.head_dim**-0.5 * self.global_num_heads**-0.5)
1072 reduce_sum = None if tp_context is None else tp_context.reduce_sum
1073 if key_handle is not None:
1074 key = key_handle.wait()
1075 topk_indices, candidate_blocks = self._run_selection(
1076 query, key, merge_weight, candidate_blocks, query_offset,
1077 reduce_sum, minimum_key_indices, query_segments,
1078 )
1079 return SharedCompressedIndexerOutput(
1080 topk_indices=topk_indices,
1081 index_key=key,
1082 candidate_blocks=candidate_blocks,
1083 index_query=query,
1084 merge_weight=merge_weight,
1085 )
1087 def enable_selection_exclude(self) -> None:
1088 """Keep discrete selection fixed when the enclosing layer recomputes."""
1089 object.__setattr__(
1090 self, "_selection_region", checkpoint_exclude_wrapper(self._select_indices)
1091 )
1093 def _run_selection(self, *args: Any, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor | None]:
1094 """Select directly, or reuse the result retained by the checkpoint region."""
1095 selection = self._select_indices if self._selection_region is None else self._selection_region
1096 return selection(*args, **kwargs)
1098 def _select_indices(
1099 self, query: torch.Tensor, key: torch.Tensor, merge_weight: torch.Tensor,
1100 candidate_blocks: torch.Tensor | None, query_offset: int,
1101 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None,
1102 minimum_key_indices: torch.Tensor | None, query_segments: tuple | None,
1103 ) -> tuple[torch.Tensor, torch.Tensor | None]:
1104 """Select the requested implementation without changing candidate policies."""
1105 if self.use_fused and (self.is_candidate_source or self.uses_candidates):
1106 raise NotImplementedError("V4.1 fused Indexer does not support candidate-pool selection")
1107 if self.is_candidate_source:
1108 topk_indices, candidate_blocks = compressed_causal_topk_and_candidates(
1109 query,
1110 key,
1111 merge_weight,
1112 compress_ratio=self.compress_ratio,
1113 sparse_count=self.index_topk,
1114 topk_blocks=self.candidate_topk_blocks,
1115 block_size=self.candidate_block_size,
1116 query_offset=query_offset,
1117 query_chunk_size=self.query_chunk_size,
1118 reduce_sum=reduce_sum,
1119 minimum_key_indices=minimum_key_indices,
1120 )
1121 elif self.uses_candidates:
1122 if candidate_blocks is None:
1123 raise RuntimeError("a hierarchical Reindex layer requires candidate block ids")
1124 topk_indices = compressed_candidate_topk(
1125 query,
1126 key,
1127 merge_weight,
1128 candidate_blocks,
1129 compress_ratio=self.compress_ratio,
1130 sparse_count=self.index_topk,
1131 block_size=self.candidate_block_size,
1132 query_offset=query_offset,
1133 query_chunk_size=self.query_chunk_size,
1134 reduce_sum=reduce_sum,
1135 minimum_key_indices=minimum_key_indices,
1136 )
1137 else:
1138 candidate_blocks = None
1139 topk_indices = compressed_causal_topk(
1140 query,
1141 key,
1142 merge_weight,
1143 compress_ratio=self.compress_ratio,
1144 sparse_count=self.index_topk,
1145 query_offset=query_offset,
1146 query_chunk_size=self.query_chunk_size,
1147 reduce_sum=reduce_sum,
1148 minimum_key_indices=minimum_key_indices,
1149 query_segments=query_segments,
1150 use_fused=self.use_fused,
1151 )
1152 return topk_indices, candidate_blocks
1155def reference_sparse_attention(
1156 query: torch.Tensor,
1157 key_value: torch.Tensor,
1158 sparse_indices: torch.Tensor,
1159 sinks: torch.Tensor,
1160 scale: float,
1161) -> torch.Tensor:
1162 """Compute sink attention with eager Torch operations on any device.
1164 Args:
1165 query: Main queries [batch, heads, local queries, head dimension].
1166 key_value: Shared K=V bank [batch, 1, keys, head dimension].
1167 sparse_indices: Visible bank IDs [batch, local queries, selections];
1168 negative entries are padding and repeated IDs count only once.
1169 sinks: Learnable scalar sink logits [heads], with zero-valued output.
1170 scale: Main attention logit scale.
1172 Returns:
1173 Attention values [batch, local queries, heads, head dimension].
1174 Softmax uses FP32; matrix multiplications retain the input precision.
1176 Note:
1177 This reference materializes a dense query/key mask and scores. It is
1178 intended for explicit unfused comparisons, not long-sequence efficiency.
1179 """
1180 batch_size, num_heads, sequence_length, _ = query.shape
1181 key_length = key_value.shape[2]
1182 valid = sparse_indices >= 0
1183 safe_indices = sparse_indices.clamp_min(0).long()
1184 allowed_counts = torch.zeros(
1185 batch_size,
1186 sequence_length,
1187 key_length,
1188 dtype=torch.int32,
1189 device=query.device,
1190 )
1191 allowed_counts.scatter_add_(2, safe_indices, valid.to(torch.int32))
1192 allowed = allowed_counts > 0
1193 logits = torch.matmul(query, key_value.transpose(-1, -2)) * scale
1194 logits = logits.masked_fill(~allowed.unsqueeze(1), float("-inf"))
1195 sink_logits = sinks.view(1, num_heads, 1, 1).expand(batch_size, -1, sequence_length, -1)
1196 probabilities = torch.cat((logits, sink_logits), dim=-1).softmax(dim=-1, dtype=torch.float32)
1197 return torch.matmul(probabilities[..., :-1].to(key_value.dtype), key_value).transpose(1, 2)
1200class SharedCompressedDSAAttentionBase(nn.Module):
1201 """Shared V4.1 attention semantics with selectable reference execution."""
1203 def __init__(
1204 self,
1205 module: nn.Module,
1206 *,
1207 replace_indexer: bool,
1208 use_fused_ops: bool,
1209 ) -> None:
1210 """Transfer source state and configure the execution implementation."""
1211 super().__init__()
1212 for name, child in module._modules.items(): # pylint: disable=protected-access
1213 self.add_module(name, child)
1214 for name, parameter in module._parameters.items(): # pylint: disable=protected-access
1215 self.register_parameter(name, parameter)
1216 indexer = getattr(self, "indexer", None)
1217 if (
1218 replace_indexer
1219 and indexer is not None
1220 and not isinstance(indexer, SharedCompressedDSAIndexer)
1221 ):
1222 self.indexer = SharedCompressedDSAIndexer(indexer)
1223 self.config = module.config
1224 self.layer_idx = module.layer_idx
1225 self.num_heads = module.num_heads
1226 self.num_key_value_groups = module.num_key_value_groups
1227 self.num_groups = module.config.o_groups
1228 self.compress_ratio = module.compress_ratio
1229 self.is_kv_source = module.is_kv_source
1230 self.is_index_source = module.is_index_source
1231 self.kv_source_layer_idx = module.kv_source_layer_idx
1232 self.index_source_layer_idx = module.index_source_layer_idx
1233 self.candidate_source_layer_idx = module.candidate_source_layer_idx
1234 self.head_dim = module.head_dim
1235 self.rope_head_dim = module.config.qk_rope_head_dim
1236 self.sliding_window = module.sliding_window
1237 self.scaling = module.scaling
1238 self.use_fused_ops = use_fused_ops
1239 if self.is_index_source:
1240 self.indexer.use_fused = use_fused_ops
1241 self.train(module.training)
1243 @staticmethod
1244 def _validate_cp_inputs(
1245 cp_context: SharedCompressedAttentionCPContext,
1246 position_ids: torch.Tensor,
1247 sequence_length: int,
1248 compress_ratio: int,
1249 is_kv_source: bool,
1250 ) -> tuple[int, int]:
1251 """Validate contiguous CP shards and return global query geometry."""
1252 if cp_context.size <= 1 or not 0 <= cp_context.rank < cp_context.size:
1253 raise ValueError(
1254 "shared compressed attention CP requires size > 1 and "
1255 f"0 <= rank < size, got rank={cp_context.rank}, size={cp_context.size}"
1256 )
1257 if is_kv_source and sequence_length % compress_ratio:
1258 raise ValueError(
1259 "compressed-KV CP requires every local sequence shard to divide "
1260 f"compress_ratio={compress_ratio}, got {sequence_length}"
1261 )
1262 query_offset = cp_context.rank * sequence_length
1263 if not torch.all(position_ids[..., 0] == query_offset):
1264 raise ValueError(
1265 "shared compressed attention CP requires contiguous global position_ids; "
1266 f"rank {cp_context.rank} expected first position {query_offset}"
1267 )
1268 return query_offset, sequence_length * cp_context.size
1270 def _project_raw_kv(
1271 self,
1272 hidden_states: torch.Tensor,
1273 cos: torch.Tensor,
1274 sin: torch.Tensor,
1275 cp_context: SharedCompressedAttentionCPContext | None,
1276 ) -> tuple[torch.Tensor | None, SequenceGatherHandle | None]:
1277 """Project raw K=V and launch CP communication before query GEMMs."""
1278 key_value = self.kv_norm(self.kv_proj(hidden_states)).unsqueeze(1)
1279 key_value = _apply_v41_rope(key_value, cos, sin)
1280 if cp_context is None:
1281 return key_value, None
1282 return None, cp_context.launch(key_value, 2)
1284 def _project_query(
1285 self,
1286 hidden_states: torch.Tensor,
1287 cos: torch.Tensor,
1288 sin: torch.Tensor,
1289 ) -> tuple[torch.Tensor, torch.Tensor]:
1290 """Project and normalize the local query heads."""
1291 batch_size, sequence_length, _ = hidden_states.shape
1292 query_residual = self.q_a_norm(self.q_a_proj(hidden_states))
1293 query = self.q_b_proj(query_residual).view(batch_size, sequence_length, -1, self.head_dim)
1294 query = query.transpose(1, 2)
1295 return query_residual, _apply_v41_rope(query, cos, sin)
1297 @staticmethod
1298 def _validate_forward_inputs(
1299 hidden_states: torch.Tensor,
1300 attention_mask: torch.Tensor | None,
1301 past_key_values: Any | None,
1302 kwargs: dict[str, Any],
1303 ) -> tuple[
1304 "SharedCompressedAttentionState",
1305 "SharedCompressedPackedSequence | None",
1306 "SharedCompressedAttentionCPContext | None",
1307 "SharedCompressedAttentionTPContext | None",
1308 int,
1309 int,
1310 ]:
1311 """Pop and validate the shared-attention keyword arguments."""
1312 if past_key_values is not None:
1313 raise NotImplementedError("shared compressed DSA training does not support KV cache")
1314 shared_state = kwargs.pop("shared_attention_state", None)
1315 if not isinstance(shared_state, SharedCompressedAttentionState):
1316 raise ValueError("shared_attention_state is required for shared compressed DSA")
1317 if attention_mask is not None:
1318 raise NotImplementedError(
1319 "shared compressed DSA consumes compact packed_seq_params instead of a dense mask"
1320 )
1321 batch_size, sequence_length, _ = hidden_states.shape
1322 packed_sequence = kwargs.pop("packed_seq_params", None)
1323 if packed_sequence is not None and not isinstance(
1324 packed_sequence, SharedCompressedPackedSequence
1325 ):
1326 raise TypeError(
1327 "packed_seq_params must be SharedCompressedPackedSequence, "
1328 f"got {type(packed_sequence).__name__}"
1329 )
1330 cp_context = kwargs.pop("shared_attention_cp_context", None)
1331 if cp_context is not None and not isinstance(cp_context, SharedCompressedAttentionCPContext):
1332 raise TypeError(
1333 "shared_attention_cp_context must be SharedCompressedAttentionCPContext, "
1334 f"got {type(cp_context).__name__}"
1335 )
1336 tp_context = kwargs.pop("shared_attention_tp_context", None)
1337 if tp_context is not None and not isinstance(tp_context, SharedCompressedAttentionTPContext):
1338 raise TypeError(
1339 "shared_attention_tp_context must be SharedCompressedAttentionTPContext, "
1340 f"got {type(tp_context).__name__}"
1341 )
1342 return shared_state, packed_sequence, cp_context, tp_context, batch_size, sequence_length
1344 def _resolve_packed_segments(
1345 self,
1346 packed_sequence: "SharedCompressedPackedSequence | None",
1347 batch_size: int,
1348 sequence_length: int,
1349 query_offset: int,
1350 global_sequence_length: int,
1351 device: torch.device,
1352 ) -> tuple[torch.Tensor | None, torch.Tensor | None, tuple]:
1353 """Validate packed document boundaries and derive local Q/K segments."""
1354 segment_starts = None
1355 minimum_key_indices = None
1356 ratio = max(1, self.compress_ratio)
1357 raw_end = query_offset + sequence_length
1358 query_segments = ((0, sequence_length, 0, raw_end // ratio, raw_end % ratio),)
1359 if packed_sequence is not None:
1360 if batch_size != 1:
1361 raise ValueError("V4.1 compact packed attention currently requires micro_batch_size=1")
1362 if (
1363 packed_sequence.local_query_start != query_offset
1364 or packed_sequence.local_query_length != sequence_length
1365 or packed_sequence.global_sequence_length != global_sequence_length
1366 ):
1367 packed_shard = (
1368 packed_sequence.local_query_start,
1369 packed_sequence.local_query_length,
1370 packed_sequence.global_sequence_length,
1371 )
1372 raise ValueError(
1373 "packed sequence geometry does not match the local CP shard: "
1374 f"packed={packed_shard}, "
1375 f"attention={(query_offset, sequence_length, global_sequence_length)}"
1376 )
1377 packed_sequence = packed_sequence.prepare(device, (self.compress_ratio,))
1378 segment_starts = packed_sequence.local_segment_starts(device)
1379 if self.compress_ratio:
1380 minimum_key_indices = packed_sequence.minimum_key_indices(device, self.compress_ratio)
1381 query_segments = packed_sequence.indexer_segments(ratio)
1382 return segment_starts, minimum_key_indices, query_segments
1384 def _compute_compressed_kv(
1385 self,
1386 hidden_states: torch.Tensor,
1387 position_embeddings: dict[str, tuple[torch.Tensor, torch.Tensor]],
1388 shared_state: "SharedCompressedAttentionState",
1389 cp_context: "SharedCompressedAttentionCPContext | None",
1390 ) -> tuple[torch.Tensor | None, "SequenceGatherHandle | None"]:
1391 """Compress local hidden states and publish or launch the compressed KV.
1393 The compressed K=V reaches its consumers - the producing Full layer
1394 included - through the shared state, so only the publish (or the CP
1395 gather handle that feeds it) leaves this call.
1396 """
1397 latent = None
1398 compressed_handle = None
1399 if self.is_kv_source:
1400 latent, compressed = self.compressor(hidden_states, position_embeddings["compress"])
1401 if cp_context is None:
1402 if not is_recomputing():
1403 shared_state.publish_compressed_kv(self.layer_idx, compressed)
1404 else:
1405 compressed_handle = cp_context.launch(compressed, 1)
1406 return latent, compressed_handle
1408 def _run_indexer(
1409 self,
1410 hidden_states: torch.Tensor,
1411 query_residual: torch.Tensor,
1412 latent: torch.Tensor | None,
1413 position_embeddings: dict[str, tuple[torch.Tensor, torch.Tensor]],
1414 shared_state: "SharedCompressedAttentionState",
1415 cp_context: "SharedCompressedAttentionCPContext | None",
1416 tp_context: "SharedCompressedAttentionTPContext | None",
1417 query_offset: int,
1418 minimum_key_indices: torch.Tensor | None,
1419 query_segments: tuple,
1420 ) -> Any | None:
1421 """Run the indexer on index-source layers and publish its forward outputs."""
1422 if not self.is_index_source:
1423 return None
1424 index_key = (
1425 None
1426 if self.indexer.owns_key
1427 else shared_state.require_index_key(self.kv_source_layer_idx, self.layer_idx)
1428 )
1429 candidate_blocks = (
1430 shared_state.require_candidate_blocks(
1431 self.candidate_source_layer_idx,
1432 self.layer_idx,
1433 )
1434 if self.indexer.uses_candidates else None
1435 )
1436 indexer_output = self.indexer(
1437 hidden_states,
1438 query_residual,
1439 latent,
1440 position_embeddings["compress"],
1441 index_key=index_key,
1442 candidate_blocks=candidate_blocks,
1443 cp_context=cp_context,
1444 tp_context=tp_context,
1445 query_offset=query_offset,
1446 minimum_key_indices=minimum_key_indices,
1447 query_segments=query_segments,
1448 )
1449 if is_recomputing():
1450 # A recompute must not publish again: its consumers read the entries
1451 # the forward pass wrote, and a second write would break the
1452 # single-publisher contract.
1453 return indexer_output
1454 if self.indexer.owns_key:
1455 shared_state.publish_index_key(self.layer_idx, indexer_output.index_key)
1456 shared_state.publish_topk_indices(self.layer_idx, indexer_output.topk_indices)
1457 if indexer_output.candidate_blocks is not None:
1458 shared_state.publish_candidate_blocks(
1459 self.layer_idx,
1460 indexer_output.candidate_blocks,
1461 )
1462 return indexer_output
1464 def _build_sparse_indices(
1465 self,
1466 key_value: torch.Tensor,
1467 compressed_kv: torch.Tensor | None,
1468 topk_indices: torch.Tensor | None,
1469 sequence_length: int,
1470 query_offset: int,
1471 global_sequence_length: int,
1472 segment_starts: torch.Tensor | None,
1473 ) -> tuple[torch.Tensor, torch.Tensor]:
1474 """Assemble dense-reference banks only after reference dispatch is selected."""
1475 window = build_sliding_window_indices(
1476 key_value.shape[0], sequence_length, self.sliding_window, key_value.device,
1477 query_offset=query_offset, key_length=global_sequence_length,
1478 )
1479 if segment_starts is not None:
1480 window.masked_fill_(window < segment_starts.unsqueeze(-1), -1)
1481 if compressed_kv is None:
1482 return key_value, window
1483 combined_key_value = torch.cat((key_value, compressed_kv.unsqueeze(1)), dim=2)
1484 compressed_indices = torch.where(topk_indices >= 0, topk_indices.long() + global_sequence_length, -1)
1485 return combined_key_value, torch.cat((window, compressed_indices), dim=-1)
1487 def _apply_indexer_loss(
1488 self,
1489 indexer_output: Any | None,
1490 query: torch.Tensor,
1491 compressed_kv: torch.Tensor | None,
1492 tp_context: "SharedCompressedAttentionTPContext | None",
1493 query_segments: tuple,
1494 ) -> torch.Tensor:
1495 """Auto-scale the query by the indexer KL loss when configured."""
1496 if (
1497 indexer_output is not None
1498 and self.training
1499 and self.indexer.loss_coeff
1500 and compressed_kv is not None
1501 ):
1502 indexer_loss = shared_compressed_indexer_kl_loss(
1503 indexer_output.index_query,
1504 indexer_output.index_key,
1505 indexer_output.merge_weight,
1506 query,
1507 compressed_kv,
1508 indexer_output.topk_indices,
1509 self.sinks,
1510 attention_scale=self.scaling,
1511 loss_coeff=self.indexer.loss_coeff,
1512 query_chunk_size=self.indexer.query_chunk_size,
1513 tp_context=tp_context,
1514 compress_ratio=self.compress_ratio,
1515 query_segments=query_segments,
1516 use_fused=self.use_fused_ops,
1517 )
1518 query = aux_loss_auto_scale(query, indexer_loss)
1519 return query
1521 def _run_sparse_attention(
1522 self, query: torch.Tensor, key_value: torch.Tensor, compressed_kv: torch.Tensor | None,
1523 topk_indices: torch.Tensor | None, query_offset: int, global_sequence_length: int,
1524 segment_starts: torch.Tensor | None, query_segments: tuple,
1525 ) -> torch.Tensor:
1526 """Run the chosen implementation; unsupported explicit fusion always raises."""
1527 if self.use_fused_ops:
1528 if not (self.candidate_source_layer_idx is None
1529 and (topk_indices is None or topk_indices.shape[-1] == min(512, compressed_kv.shape[1]))
1530 and supports_fused_attention_inputs(query, self.compress_ratio, self.sliding_window,
1531 0 if topk_indices is None else topk_indices.shape[-1])):
1532 raise RuntimeError(
1533 "V4.1 fused attention is required but this call is unsupported: "
1534 f"layer={self.layer_idx}, device={query.device}, dtype={query.dtype}, Q={tuple(query.shape)}, "
1535 f"ratio={self.compress_ratio}, window={self.sliding_window}. "
1536 "Check the SMLA training extension and native input contract; no reference fallback was run."
1537 )
1538 return sparse_flash_mla(
1539 query, key_value, compressed_kv, topk_indices, self.sinks, self.scaling,
1540 self.compress_ratio, self.sliding_window, query_offset, query_segments,
1541 )
1542 combined_key_value, sparse_indices = self._build_sparse_indices(
1543 key_value, compressed_kv, topk_indices, query.shape[2], query_offset,
1544 global_sequence_length, segment_starts,
1545 )
1546 return reference_sparse_attention(query, combined_key_value, sparse_indices, self.sinks, self.scaling)
1548 def _project_output(
1549 self,
1550 attention_output: torch.Tensor,
1551 cos: torch.Tensor,
1552 sin: torch.Tensor,
1553 batch_size: int,
1554 sequence_length: int,
1555 ) -> torch.Tensor:
1556 """Undo RoPE and apply the grouped low-rank output projection."""
1557 attention_output = _apply_v41_rope(attention_output.transpose(1, 2), cos, -sin).transpose(1, 2)
1558 grouped = attention_output.reshape(batch_size, sequence_length, self.num_groups, -1)
1559 hidden_per_group = grouped.shape[-1]
1560 output_per_group = self.o_a_proj.weight.shape[0] // self.num_groups
1561 projection = self.o_a_proj.weight.view(
1562 self.num_groups, output_per_group, hidden_per_group
1563 ).transpose(1, 2)
1564 flattened = grouped.reshape(-1, self.num_groups, hidden_per_group).transpose(0, 1)
1565 projected = torch.bmm(flattened, projection).transpose(0, 1)
1566 projected = projected.reshape(batch_size, sequence_length, -1)
1567 return self.o_b_proj(projected)
1569 def forward(
1570 self,
1571 hidden_states: torch.Tensor,
1572 position_embeddings: dict[str, tuple[torch.Tensor, torch.Tensor]],
1573 position_ids: torch.Tensor,
1574 attention_mask: torch.Tensor | None,
1575 past_key_values: Any | None = None,
1576 **kwargs: Any,
1577 ) -> tuple[torch.Tensor, torch.Tensor | None]:
1578 """Publish or consume compressed KV and execute indexed K=V attention."""
1579 (
1580 shared_state,
1581 packed_sequence,
1582 cp_context,
1583 tp_context,
1584 batch_size,
1585 sequence_length,
1586 ) = self._validate_forward_inputs(hidden_states, attention_mask, past_key_values, kwargs)
1587 query_offset = 0
1588 global_sequence_length = sequence_length
1589 if cp_context is not None:
1590 query_offset, global_sequence_length = self._validate_cp_inputs(
1591 cp_context,
1592 position_ids,
1593 sequence_length,
1594 self.compress_ratio,
1595 self.is_kv_source,
1596 )
1597 segment_starts, minimum_key_indices, query_segments = self._resolve_packed_segments(
1598 packed_sequence,
1599 batch_size,
1600 sequence_length,
1601 query_offset,
1602 global_sequence_length,
1603 hidden_states.device,
1604 )
1606 rope_type = "compress" if self.compress_ratio else "main"
1607 cos, sin = position_embeddings[rope_type]
1608 key_value, raw_kv_handle = self._project_raw_kv(hidden_states, cos, sin, cp_context)
1609 query_residual, query = self._project_query(hidden_states, cos, sin)
1610 latent, compressed_handle = self._compute_compressed_kv(
1611 hidden_states, position_embeddings, shared_state, cp_context
1612 )
1613 indexer_output = self._run_indexer(
1614 hidden_states,
1615 query_residual,
1616 latent,
1617 position_embeddings,
1618 shared_state,
1619 cp_context,
1620 tp_context,
1621 query_offset,
1622 minimum_key_indices,
1623 query_segments,
1624 )
1626 if raw_kv_handle is not None:
1627 key_value = raw_kv_handle.wait()
1628 if compressed_handle is not None:
1629 compressed = compressed_handle.wait()
1630 if not is_recomputing():
1631 shared_state.publish_compressed_kv(self.layer_idx, compressed)
1632 if key_value is None:
1633 raise RuntimeError("raw KV all-gather did not produce a tensor")
1635 compressed_kv = None
1636 topk_indices = None
1637 if self.compress_ratio:
1638 compressed_kv = shared_state.require_compressed_kv(self.kv_source_layer_idx, self.layer_idx)
1639 topk_indices = shared_state.require_topk_indices(self.index_source_layer_idx, self.layer_idx)
1640 query = self._apply_indexer_loss(indexer_output, query, compressed_kv, tp_context, query_segments)
1641 attention_output = self._run_sparse_attention(
1642 query, key_value, compressed_kv, topk_indices, query_offset, global_sequence_length,
1643 segment_starts, query_segments,
1644 )
1645 output = self._project_output(attention_output, cos, sin, batch_size, sequence_length)
1646 return output, None
1649@module_replacement
1650class SharedCompressedDSAAttention(SharedCompressedDSAAttentionBase):
1651 """High-performance V4.1 shared compressed attention replacement."""
1653 def __init__(
1654 self,
1655 module: nn.Module,
1656 module_fqn: str = "",
1657 context: Mapping[str, Any] | None = None,
1658 *,
1659 use_fused_ops: bool = False,
1660 ) -> None:
1661 """Transfer source state and select the implementation used for comparisons.
1663 Args:
1664 module: Source attention module exposing the shared-compressed DSA
1665 structural contract.
1666 module_fqn: Fully qualified source name supplied by replacement.
1667 context: Read-only replacement context supplied by Trainer.
1668 use_fused_ops: Use LI V2, SMLA and fused KL when True,
1669 raising on unsupported calls without automatic fallback.
1670 False selects eager PyTorch attention, Indexer and KL on every
1671 device, without loading optional operator extensions.
1673 Note:
1674 All implementations retain the same parameter and CP/TP contracts.
1675 Selecting native preserves the existing external teacher and scalar
1676 KL. Reuse and disabled KL do not introduce extra operator calls.
1677 Fusion supports ordinary causal r1/r2 training on NPU with BF16
1678 main attention, D=512, SWA=128, Indexer D=128, TopK=512 and TP=1.
1679 Active fused KL requires 8, 16, 32 or 64 Indexer heads. Candidate
1680 pools are unsupported. Packed inputs retain the existing alignment
1681 constraints and are dispatched as separate BSND document segments.
1682 The ops-transformer training extension is loaded on the first
1683 supported fused call; a missing extension raises ImportError.
1684 """
1685 if not isinstance(use_fused_ops, bool):
1686 raise ValueError(f"use_fused_ops must be a bool, got {use_fused_ops!r}")
1687 del module_fqn, context
1688 super().__init__(
1689 module,
1690 replace_indexer=True,
1691 use_fused_ops=use_fused_ops,
1692 )
1695__all__ = [
1696 "SharedCompressedAttentionCPContext",
1697 "SharedCompressedAttentionState",
1698 "SharedCompressedPackedSequence",
1699 "SharedCompressedAttentionTPContext",
1700 "SharedCompressedDSAAttention",
1701 "SharedCompressedDSAAttentionBase",
1702 "SharedCompressedDSAIndexer",
1703 "build_sliding_window_indices",
1704 "compressed_candidate_topk",
1705 "compressed_causal_topk_and_candidates",
1706 "compressed_causal_topk",
1707 "select_candidate_block_indices",
1708 "select_candidate_blocks",
1709 "shared_compressed_indexer_kl_loss",
1710 "reference_sparse_attention",
1711]