Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / shared_compressed_dsa_attention.py: 68%
1040 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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 functools import partial
32from typing import Any, Protocol
34import torch # pylint: disable=forbidden-backend-import
35from torch import nn # pylint: disable=forbidden-backend-import
36from torch.utils._pytree import register_pytree_node # pylint: disable=forbidden-backend-import
38from hyper_parallel.components.functional.aux_loss import aux_loss_auto_scale
39from hyper_parallel.components.functional.lightning_indexer import (
40 lightning_indexer, sort_key_indices, validate_indexer_tensors,
41)
42from hyper_parallel.components.functional.sparse_flash_mla import (
43 supports_fused_attention_inputs, sparse_flash_mla,
44)
45from hyper_parallel.components.functional.sparse_lightning_indexer_kl_loss_grad import (
46 gather_selected_keys, sparse_lightning_indexer_kl_loss, validate_fused_kl_inputs,
47)
48from hyper_parallel.core.activation_memory.recompute_state import is_recomputing
49from hyper_parallel.core.activation_memory.wrapper import checkpoint_exclude_wrapper
50from hyper_parallel.distributed.context_parallel.compressed_document import (
51 CompressionBoundaryRoute, document_compression_plan,
52)
53from hyper_parallel.distributed.context_parallel.indexer_balance import (
54 resolve_indexer_balance_block_size, validate_indexer_hybrid_threshold,
55)
56from hyper_parallel.distributed.context_parallel.sequence_dispatch import (
57 SequenceDispatchRoute, document_prefix_route, sliding_window_route,
58)
59from hyper_parallel.models.replacement import module_replacement
62class SequenceGatherHandle(Protocol):
63 """Handle returned by asynchronous differentiable bank communication."""
65 def wait(self) -> torch.Tensor:
66 """Wait for communication and return the differentiable result."""
69@dataclass
70class _DeferredSequenceGather:
71 """Synchronous fallback with the same interface as an async gather."""
73 tensor: torch.Tensor
74 sequence_dim: int
75 gather_sequence: Callable[[torch.Tensor, int], torch.Tensor]
77 def wait(self) -> torch.Tensor:
78 """Run the fallback gather at the first consumer boundary."""
79 return self.gather_sequence(self.tensor, self.sequence_dim)
82@dataclass
83class _DeferredBankLaunch:
84 """Keep main K behind Index K in the shared communication queue."""
86 launcher: Callable[[], SequenceGatherHandle] | None
87 handle: SequenceGatherHandle | None = None
89 def start(self) -> SequenceGatherHandle:
90 """Launch once in the differentiable forward, then release captured inputs."""
91 if self.handle is None:
92 self.handle = self.launcher()
93 self.launcher = None
94 return self.handle
96 def wait(self) -> torch.Tensor:
97 """Resolve the same differentiable handle at its first consumer."""
98 return self.start().wait()
101@dataclass
102class SharedCompressedAttentionState:
103 """Autograd-bearing CSA2 state scoped to one training forward.
105 State is addressed by the layer that produced it instead of using one
106 mutable "latest value" slot. This matters for the released V4.1 topology,
107 which has several Full/Reindex groups, and makes a consumer's dependency
108 explicit when activation recomputation revisits modules out of forward
109 order.
111 Nothing drops an entry early. The state travels into every layer as one of
112 the checkpoint region's inputs, so the graph keeps it alive from the forward
113 pass to the end of backward, and a layer whose forward is recomputed reads
114 its entries a second time there. A published entry therefore has to stay
115 readable until backward is done with every consumer, which is also why the
116 publishers skip their write while recomputing.
117 """
119 compressed_kv_by_source: dict[int, torch.Tensor] = field(default_factory=dict)
120 index_key_by_source: dict[int, torch.Tensor] = field(default_factory=dict)
121 topk_indices_by_source: dict[int, torch.Tensor] = field(default_factory=dict)
122 candidate_blocks_by_source: dict[int, torch.Tensor] = field(default_factory=dict)
123 compressed_kv_offsets: dict[int, int] = field(default_factory=dict)
125 @staticmethod
126 def _require(
127 values: dict[int, torch.Tensor],
128 source_layer: int | None,
129 value_name: str,
130 consumer_layer: int,
131 ) -> torch.Tensor:
132 """Return one published tensor or report the broken layer dependency."""
133 if source_layer is None or source_layer not in values:
134 raise RuntimeError(
135 f"layer {consumer_layer} requires {value_name} from source "
136 f"layer {source_layer}, but that source has not run in this forward"
137 )
138 return values[source_layer]
140 def publish_compressed_kv(self, source_layer: int, value: torch.Tensor, key_offset: int = 0) -> None:
141 """Publish compressed K=V and its global origin, preserving the autograd graph."""
142 self.compressed_kv_by_source[source_layer] = value
143 self.compressed_kv_offsets[source_layer] = key_offset
145 def require_compressed_kv(self, source_layer: int | None, consumer_layer: int) -> torch.Tensor:
146 """Read compressed K=V from the consumer's configured Full layer."""
147 return self._require(
148 self.compressed_kv_by_source,
149 source_layer,
150 "compressed K=V",
151 consumer_layer,
152 )
154 def publish_index_key(self, source_layer: int, value: torch.Tensor) -> None:
155 """Publish the Full Indexer's shared key tensor."""
156 self.index_key_by_source[source_layer] = value
158 def require_index_key(self, source_layer: int | None, consumer_layer: int) -> torch.Tensor:
159 """Read the Indexer key associated with a compressed-KV source."""
160 return self._require(
161 self.index_key_by_source,
162 source_layer,
163 "Indexer key",
164 consumer_layer,
165 )
167 def publish_topk_indices(self, source_layer: int, value: torch.Tensor) -> None:
168 """Publish one Full/Reindex layer's token selection."""
169 self.topk_indices_by_source[source_layer] = value
171 def require_topk_indices(self, source_layer: int | None, consumer_layer: int) -> torch.Tensor:
172 """Read the latest configured Full/Reindex selection."""
173 return self._require(
174 self.topk_indices_by_source,
175 source_layer,
176 "Top-K indices",
177 consumer_layer,
178 )
180 def publish_candidate_blocks(self, source_layer: int, value: torch.Tensor) -> None:
181 """Publish hierarchical candidate blocks from the configured source."""
182 self.candidate_blocks_by_source[source_layer] = value
184 def require_candidate_blocks(
185 self,
186 source_layer: int | None,
187 consumer_layer: int,
188 ) -> torch.Tensor:
189 """Read hierarchical candidates for a later Reindex layer."""
190 return self._require(
191 self.candidate_blocks_by_source,
192 source_layer,
193 "candidate blocks",
194 consumer_layer,
195 )
198@dataclass
199class SharedCompressedIndexerOutput:
200 """Tensors produced by one Full or Reindex Indexer invocation."""
202 topk_indices: torch.Tensor
203 index_key: torch.Tensor
204 candidate_blocks: torch.Tensor | None
205 index_query: torch.Tensor
206 merge_weight: torch.Tensor
209@dataclass(frozen=True)
210class SharedCompressedAttentionCPContext:
211 """Communication for local queries and shared compressed banks.
213 ``launch_raw`` exchanges the same-document SWA halo when supplied.
214 Optional ``launch_main`` exchanges main KV document prefixes.
215 Index K retains differentiable all-gather.
216 Contexts without ``launch_raw`` provide the full raw-bank reference.
217 """
219 size: int
220 rank: int
221 gather_sequence: Callable[[torch.Tensor, int], torch.Tensor]
222 launch_sequence: Callable[[torch.Tensor, int], SequenceGatherHandle] | None = None
223 exchange_boundary: Callable[[torch.Tensor, CompressionBoundaryRoute], torch.Tensor] | None = None
224 launch_raw: Callable[[torch.Tensor, int, SequenceDispatchRoute], SequenceGatherHandle] | None = None
225 balance_indexer: Callable[..., torch.Tensor] | None = None
226 launch_main: Callable[[torch.Tensor, int, SequenceDispatchRoute], SequenceGatherHandle] | None = None
227 launch_balance: Callable[..., Any] | None = None
228 indexer_priority: bool = False
230 def launch(self, tensor: torch.Tensor, sequence_dim: int) -> SequenceGatherHandle:
231 """Launch an async gather or create a deferred synchronous gather."""
232 if self.launch_sequence is not None:
233 return self.launch_sequence(tensor, sequence_dim)
234 return _DeferredSequenceGather(tensor, sequence_dim, self.gather_sequence)
237@dataclass(frozen=True)
238class SharedCompressedAttentionTPContext:
239 """Tensor-parallel score reduction context for a head-sharded Indexer."""
241 size: int
242 rank: int
243 reduce_sum: Callable[[torch.Tensor], torch.Tensor]
246@dataclass(frozen=True)
247class SharedCompressedPackedSequence:
248 """Compact sample boundaries for one contiguous CP query shard."""
250 cu_seq_lens: torch.Tensor
251 local_query_start: int
252 local_query_length: int
253 global_sequence_length: int
254 _prepared_starts: dict[torch.device, torch.Tensor] = field(
255 default_factory=dict, init=False, repr=False, compare=False
256 )
257 _prepared_minima: dict[tuple[torch.device, int], torch.Tensor] = field(
258 default_factory=dict, init=False, repr=False, compare=False
259 )
260 _prepared_bounds: dict[tuple[torch.device, int], torch.Tensor] = field(
261 default_factory=dict, init=False, repr=False, compare=False
262 )
263 _packing_alignment: int = field(default=0, init=False, repr=False, compare=False)
264 _compression_rope: dict[int, tuple] = field(default_factory=dict, init=False, repr=False, compare=False)
265 _snapshot: dict[str, Any] = field(default_factory=dict, init=False, repr=False, compare=False)
266 _prepared_positions: dict[torch.device, torch.Tensor] = field(
267 default_factory=dict, init=False, repr=False, compare=False
268 )
269 _prepared_masks: dict[torch.device, torch.Tensor] = field(
270 default_factory=dict, init=False, repr=False, compare=False
271 )
272 _indexer_segments: dict[int, tuple[tuple[int, int, int, int, int], ...]] = field(
273 default_factory=dict, init=False, repr=False, compare=False
274 )
276 _prepared_sequences: dict[torch.device, torch.Tensor] = field(
277 default_factory=dict, init=False, repr=False, compare=False
278 )
279 _use_prepared_packing_inputs: bool = field(default=False, init=False, repr=False, compare=False)
280 _restored_from_packing_tensors: bool = field(default=False, init=False, repr=False, compare=False)
282 def prepare_packing_geometry(self, device: torch.device) -> None:
283 """Prepare CP1 graph metadata independently of compression or indexing.
285 The runtime calls this before capture. Eager retains its real boundaries;
286 graph inputs carry token-shaped metadata, not a variable document list.
287 """
288 if self.local_query_start != 0 or self.local_query_length != self.global_sequence_length:
289 return
290 self.segment_start_mask(device)
291 object.__setattr__(self, "_use_prepared_packing_inputs", True)
293 def router_metadata(self, device: torch.device) -> tuple[torch.Tensor, int]:
294 """Return document IDs and storage capacity; empty rows do not affect the mean."""
295 count = self.local_query_length if self._restored_from_packing_tensors else self.cu_seq_lens.numel() - 1
296 return self._prepared_sequences[torch.device(device)], count
298 def prepare(self, device: torch.device, compress_ratios: tuple[int, ...]) -> SharedCompressedPackedSequence:
299 """Prepare an owned boundary snapshot, reused until the input is mutated.
301 Normal in-place tensor updates invalidate the cached snapshot. Existing
302 consumers retain their owned geometry. Prepared tensors are read-only;
303 callers must not mutate boundaries through ``.data`` or external storage.
304 Inference tensors without a version counter are always snapshotted anew.
306 Args:
307 device: Device on which tensors and prepared geometry are consumed.
308 compress_ratios: Compression ratios used by the layers in this forward.
309 """
310 device = torch.device(device)
311 # Torch exposes tensor mutation tracking only through its version counter.
312 version = None if self.cu_seq_lens.is_inference() else self.cu_seq_lens._version # pylint: disable=protected-access
313 if version is not None and self._snapshot.get("version") == version:
314 prepared: SharedCompressedPackedSequence = self if self._snapshot.get("owned") else self._snapshot["value"]
315 else:
316 prepared = type(self)(
317 self.cu_seq_lens.detach().clone(), self.local_query_start,
318 self.local_query_length, self.global_sequence_length,
319 )
320 SharedCompressedPackedSequence._mark_owned(prepared)
321 self._snapshot.clear()
322 self._snapshot.update(version=version, value=prepared)
323 SharedCompressedPackedSequence._prepare_owned_metadata(prepared, device, compress_ratios)
324 return prepared
326 def _mark_owned(self) -> None:
327 """Record the owned tensor's mutation counter for subsequent layer reuse."""
328 if not self.cu_seq_lens.is_inference():
329 # Torch has no public read API for its tensor mutation counter.
330 self._snapshot.update(owned=True, version=self.cu_seq_lens._version) # pylint: disable=protected-access
332 def _prepare_owned_metadata(self, device: torch.device, compress_ratios: tuple[int, ...]) -> None:
333 """Populate caches on the owned snapshot, never on caller-owned boundaries."""
334 starts = self.local_segment_starts(device)
335 self._prepared_starts[device] = starts
336 if device not in self._prepared_sequences:
337 self._prepared_sequences[device] = torch.bucketize(
338 self.local_positions(device).squeeze(0), self.cu_seq_lens.to(device=device)[1:], right=True,
339 ).unsqueeze(0)
340 for ratio in set(compress_ratios):
341 self.validate_compression_ratio(ratio)
342 if ratio > 0 and (device, ratio) not in self._prepared_minima:
343 if self._restored_from_packing_tensors:
344 raise ValueError(
345 "graph packing inputs carry no compressed-key metadata for ratio "
346 f"{ratio}; configure every active model compress ratio on the "
347 "runtime input adapter (compress_ratios) so the graph boundary "
348 "can prepare it"
349 )
350 self._prepared_minima[(device, ratio)] = self.minimum_key_indices(device, ratio)
352 def local_positions(self, device: torch.device) -> torch.Tensor:
353 """Return the immutable local interval's global token positions."""
354 device = torch.device(device)
355 if device not in self._prepared_positions:
356 self._prepared_positions[device] = torch.arange(
357 self.local_query_start, self.local_query_start + self.local_query_length, device=device,
358 ).unsqueeze(0)
359 return self._prepared_positions[device]
361 def segment_start_mask(self, device: torch.device) -> torch.Tensor:
362 """Reuse the per-batch sample-start mask across model forwards."""
363 device = torch.device(device)
364 if device not in self._prepared_masks:
365 self._prepared_masks[device] = self.local_positions(device) == self.local_segment_starts(device)
366 return self._prepared_masks[device]
368 def document_boundaries(self) -> tuple[int, ...]:
369 """Reuse the owned batch's host boundaries for routing across layers."""
370 if "boundaries" not in self._snapshot:
371 self._snapshot["boundaries"] = tuple(self.cu_seq_lens.detach().cpu().tolist())
372 return self._snapshot["boundaries"]
374 def indexer_segments(self, ratio: int) -> tuple[tuple[int, int, int, int, int], ...]:
375 """Describe packed Q intervals and completed K prefixes without per-layer synchronization.
377 Args:
378 ratio: Compression ratio used to express each document's key prefix.
379 """
380 if self._restored_from_packing_tensors:
381 raise NotImplementedError(
382 "graph packing inputs do not carry Python document segments; the "
383 "reference selection consumes per-token tensor bounds instead"
384 )
385 if ratio not in self._indexer_segments:
386 self.validate_compression_ratio(ratio)
387 boundaries = self.document_boundaries()
388 local_start = self.local_query_start
389 local_end = local_start + self.local_query_length
390 segments = []
391 key_base = 0
392 for sample_start, sample_end in zip(boundaries, boundaries[1:]):
393 start, end = max(local_start, sample_start), min(local_end, sample_end)
394 if start < end:
395 segments.append((start-local_start, end-local_start, key_base,
396 key_base+(end-sample_start)//ratio, (end-sample_start) % ratio))
397 key_base += (sample_end-sample_start)//ratio
398 self._indexer_segments[ratio] = tuple(segments)
399 return self._indexer_segments[ratio]
401 def minimum_key_indices(
402 self, device: torch.device, compress_ratio: int, segment_starts: torch.Tensor | None = None,
403 ) -> torch.Tensor:
404 """Return the first compressed key belonging to each query's sample.
406 Args:
407 device: Device on which tensors and prepared geometry are consumed.
408 compress_ratio: Compression ratio of the consuming layer.
409 segment_starts: Optional precomputed global sample starts for local queries.
410 """
411 if compress_ratio <= 0:
412 raise ValueError("minimum_key_indices requires a positive compression ratio")
413 cached = self._prepared_minima.get((torch.device(device), compress_ratio))
414 if cached is not None:
415 return cached
416 self.validate_compression_ratio(compress_ratio)
417 del segment_starts
418 device = torch.device(device)
419 # Pure tensor arithmetic from cu_seq_lens, identical for eager and the
420 # graph boundary: count each document's complete blocks, then take the
421 # exclusive cumsum as the document's global compressed-key base.
422 boundaries = self.cu_seq_lens.to(device=device, dtype=torch.long)
423 documents = self._prepared_sequences.get(device)
424 if documents is None:
425 documents = torch.bucketize(
426 self.local_positions(device).squeeze(0), boundaries[1:], right=True,
427 ).unsqueeze(0)
428 counts = (boundaries[1:] - boundaries[:-1]) // compress_ratio
429 offsets = torch.cat((counts.new_zeros(1), counts.cumsum(dim=0)[:-1]))
430 return offsets.index_select(0, documents.squeeze(0)).unsqueeze(0)
432 def causal_key_upper_bounds(
433 self, device: torch.device, compress_ratio: int,
434 ) -> torch.Tensor:
435 """Return the exclusive upper bound of compressed keys visible to each query.
437 Both execution modes share one derivation: a query at document-local
438 position ``p`` sees compressed keys ``[document base, base + (p+1)//ratio)``.
439 The result is a ``[1, local queries]`` tensor, so graph capture replays it
440 across batches whose document count and lengths differ.
442 Args:
443 device: Device on which tensors and prepared geometry are consumed.
444 compress_ratio: Compression ratio of the consuming layer.
445 """
446 if compress_ratio <= 0:
447 raise ValueError("causal_key_upper_bounds requires a positive compression ratio")
448 key = (torch.device(device), compress_ratio)
449 cached = self._prepared_bounds.get(key)
450 if cached is None:
451 device = torch.device(device)
452 minima = self.minimum_key_indices(device, compress_ratio)
453 cached = minima + (
454 (self.local_positions(device) - self.local_segment_starts(device) + 1) // compress_ratio
455 )
456 self._prepared_bounds[key] = cached
457 return cached
459 def local_segment_starts(self, device: torch.device) -> torch.Tensor:
460 """Return each local query's global packed-sample start position.
462 Args:
463 device: Device on which tensors and prepared geometry are consumed.
464 """
465 cached = self._prepared_starts.get(torch.device(device))
466 if cached is not None:
467 return cached
468 if self.local_query_start < 0 or self.local_query_length < 0 or (
469 self.local_query_start + self.local_query_length > self.global_sequence_length
470 ):
471 raise ValueError("packed local query interval must lie within the global sequence")
472 boundaries = self.cu_seq_lens.to(dtype=torch.long)
473 if boundaries.ndim != 1 or boundaries.numel() < 2 or int(boundaries[0]) != 0:
474 raise ValueError("packed cu_seq_lens must be one-dimensional and start with zero")
475 if int(boundaries[-1]) != self.global_sequence_length:
476 raise ValueError(
477 "packed cu_seq_lens must cover global_sequence_length: "
478 f"{int(boundaries[-1])} versus {self.global_sequence_length}"
479 )
480 if torch.any(boundaries[1:] <= boundaries[:-1]):
481 raise ValueError("packed cu_seq_lens must be strictly increasing")
482 boundaries = boundaries.to(device=device)
483 query_positions = self.local_positions(device).squeeze(0)
484 segment_ids = torch.bucketize(query_positions, boundaries[1:], right=True)
485 return boundaries[:-1].index_select(0, segment_ids).unsqueeze(0)
487 @staticmethod
488 def validate_compression_ratio(compress_ratio: int) -> None:
489 """Validate the ratio; complete block planning handles natural document lengths."""
490 if compress_ratio < 0:
491 raise ValueError("compression ratio must be nonnegative")
493 def is_compression_aligned(self, compress_ratio: int) -> bool:
494 """Report whether every document boundary lies on a complete-block edge.
496 Eager reads the real boundaries; graph packing inputs read the alignment
497 the graph boundary validated before capture. Both answer the same
498 question for the same batch.
500 Args:
501 compress_ratio: Compression ratio of the consuming layer.
502 """
503 if compress_ratio <= 1:
504 return True
505 if self._restored_from_packing_tensors:
506 return 0 < self._packing_alignment and self._packing_alignment % compress_ratio == 0
507 return not any(boundary % compress_ratio for boundary in self.document_boundaries())
509 def compressed_key_length(self, compress_ratio: int) -> int:
510 """Return the total number of complete compressed keys in the batch.
512 For alignment-validated graph packing inputs the contiguous count equals
513 the per-document sum; eager keeps exact arithmetic for natural lengths.
515 Args:
516 compress_ratio: Compression ratio of the consuming layer.
517 """
518 ratio = max(1, compress_ratio)
519 if self._restored_from_packing_tensors:
520 return self.local_query_length // ratio
521 boundaries = self.document_boundaries()
522 return sum((end - start) // ratio for start, end in zip(boundaries, boundaries[1:]))
524 def block_plan(self, ratio: int, size: int, rank: int) -> Any:
525 """Use the original contiguous compressor fast path when every group is aligned.
527 Args:
528 ratio: Compression ratio of the consuming layer.
529 size: Number of contiguous CP shards.
530 rank: Local rank within the CP group.
532 Returns:
533 Document-local block plan, or None when contiguous compression suffices.
534 """
535 if self._restored_from_packing_tensors:
536 # The graph boundary rejects misaligned documents before capture, so
537 # contiguous compression is exact and needs no Python-side plan.
538 return None
539 if ratio <= 1 or (self.local_query_length % ratio == 0
540 and all(boundary % ratio == 0 for boundary in self.document_boundaries())):
541 return None
542 return document_compression_plan(self.document_boundaries(), ratio, self.local_query_length, size, rank)
544 def compressed_positions(self, plan: Any, embeddings: tuple, context: Any) -> tuple:
545 """Reuse boundary RoPE once per batch/source embedding, retaining its original precision.
547 Args:
548 plan: Document-local compression plan for the current CP rank.
549 embeddings: Cosine and sine coefficients for the local raw-token shard.
550 context: CP boundary-dispatch context, or None when no exchange is needed.
552 Returns:
553 Cosine and sine coefficients at the first token of each owned block.
554 """
555 cos, sin = embeddings
556 # Preserve cache correctness when caller-owned RoPE tensors are updated in place.
557 versions = tuple(None if item.is_inference() else item._version for item in embeddings) # pylint: disable=protected-access
558 cached = self._compression_rope.get(plan.ratio)
559 if (cached is not None and cached[0] is cos and cached[1] is sin and cached[2] == versions
560 and None not in versions and not cos.requires_grad and not sin.requires_grad):
561 return cached[3]
562 result = plan.positions(cos, sin, context)
563 if not cos.requires_grad and not sin.requires_grad:
564 self._compression_rope[plan.ratio] = (cos, sin, versions, result)
565 return result
567 def raw_document_starts(self) -> tuple:
568 """Keep raw token coordinates separate from compact compressed-key IDs."""
569 if self._restored_from_packing_tensors:
570 # Only the fused attention kernel consumes these Python
571 # coordinates; fused execution is gated out of graph mode.
572 return ()
573 first, last = self.local_query_start, self.local_query_start + self.local_query_length
574 return tuple((max(first, start)-first, start)
575 for start, end in zip(self.document_boundaries(), self.document_boundaries()[1:])
576 if max(first, start) < min(last, end))
579 def pytree_flatten(self) -> tuple[list[Any], tuple]:
580 """Expose packing tensors with a schema independent of document count and S."""
581 if self._use_prepared_packing_inputs:
582 return [self._prepared_starts, self._prepared_positions,
583 self._prepared_masks, self._prepared_sequences,
584 self._prepared_minima], ("packing", self._packing_alignment)
585 return [self.cu_seq_lens], (self.local_query_start, self.local_query_length, self.global_sequence_length)
587 @classmethod
588 def pytree_unflatten(cls, children: Any, context: tuple) -> SharedCompressedPackedSequence:
589 """Restore prepared metadata without reading boundary values during capture."""
590 if context[0] != "packing":
591 return cls(children[0], *context)
592 starts, positions, masks, sequences, minima = children
593 reference = next(iter(starts.values()))
594 length = reference.shape[-1]
595 packed = cls(reference.new_empty(0), 0, length, length)
596 packed._prepared_starts.update(starts)
597 packed._prepared_positions.update(positions)
598 packed._prepared_masks.update(masks)
599 packed._prepared_sequences.update(sequences)
600 packed._prepared_minima.update(minima)
601 object.__setattr__(packed, "_packing_alignment", context[1])
602 object.__setattr__(packed, "_use_prepared_packing_inputs", True)
603 object.__setattr__(packed, "_restored_from_packing_tensors", True)
604 packed._mark_owned()
605 return packed
608class GraphPackedSequence(SharedCompressedPackedSequence):
609 """Tensor-only packing inputs created exclusively at the graph boundary.
611 The subclass is a registration marker: only graph-boundary inputs are
612 flattened as pytree nodes, so eager FSDP keeps treating packed inputs as
613 opaque leaves that retain real boundaries and caches. Every semantic lives
614 in the shared base, including the compressed-key metadata that lets the
615 graph carry indexer and compressor execution with the same code path.
616 """
619register_pytree_node(GraphPackedSequence, GraphPackedSequence.pytree_flatten,
620 GraphPackedSequence.pytree_unflatten)
623def _rotate_half(tensor: torch.Tensor) -> torch.Tensor:
624 """Rotate adjacent real/imaginary pairs for V4 interleaved RoPE."""
625 first, second = tensor.unflatten(-1, (-1, 2)).unbind(-1)
626 return torch.stack((-second, first), dim=-1).flatten(-2)
629def _apply_v41_rope(
630 tensor: torch.Tensor,
631 cos: torch.Tensor,
632 sin: torch.Tensor,
633 unsqueeze_dim: int = 1,
634) -> torch.Tensor:
635 """Apply V4 interleaved RoPE to the trailing rotary channels."""
636 cos = cos.repeat_interleave(2, dim=-1).unsqueeze(unsqueeze_dim)
637 sin = sin.repeat_interleave(2, dim=-1).unsqueeze(unsqueeze_dim)
638 rope_dim = cos.shape[-1]
639 pass_through, rotary = tensor.split((tensor.shape[-1] - rope_dim, rope_dim), dim=-1)
640 rotary = ((rotary.float() * cos) + (_rotate_half(rotary).float() * sin)).to(tensor.dtype)
641 return torch.cat((pass_through, rotary), dim=-1) if pass_through.shape[-1] else rotary
644def build_sliding_window_indices(
645 batch_size: int,
646 sequence_length: int,
647 window_size: int,
648 device: torch.device,
649 *,
650 query_offset: int = 0,
651 key_length: int | None = None,
652) -> torch.Tensor:
653 """Build fixed-width causal sliding-window indices."""
654 if key_length is None:
655 key_length = sequence_length
656 width = min(key_length, window_size)
657 query = torch.arange(
658 query_offset,
659 query_offset + sequence_length,
660 device=device,
661 ).unsqueeze(-1)
662 offsets = torch.arange(width - 1, -1, -1, device=device)
663 indices = query - offsets
664 indices = indices.masked_fill(indices < 0, -1)
665 return indices.unsqueeze(0).expand(batch_size, -1, -1)
668def _run_lightning_indexer(
669 query: torch.Tensor, key: torch.Tensor, merge_weight: torch.Tensor,
670 top_k: int, compress_ratio: int, query_offset: int,
671 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None,
672 minimum_key_indices: torch.Tensor | None, query_segments: tuple | None,
673) -> torch.Tensor:
674 """Run LI V2 with explicit adapter errors and native capability validation."""
675 if reduce_sum is not None:
676 raise NotImplementedError("Fused LightningIndexer does not support TP score reduction before TopK")
677 if minimum_key_indices is not None and query_segments is None:
678 raise ValueError("Fused LightningIndexer requires query_segments when minimum_key_indices is set")
679 validate_indexer_tensors(query, key, merge_weight)
680 if query.device.type != "npu":
681 raise RuntimeError(f"Fused LightningIndexer requires NPU inputs; got device={query.device}")
682 if query_segments is None:
683 length = query.shape[1]
684 raw_end = query_offset + length
685 query_segments = ((0, length, 0, raw_end // compress_ratio, raw_end % compress_ratio),)
686 return lightning_indexer(query, key, merge_weight, compress_ratio, top_k, query_segments)
689def _causal_key_bounds(
690 query: torch.Tensor, ratio: int, query_offset: int,
691 minimum_key_indices: torch.Tensor | None, segments: tuple | None,
692 causal_upper: torch.Tensor | None = None,
693) -> tuple[torch.Tensor, torch.Tensor | None]:
694 """Derive visibility from original query positions within each independently compressed document.
696 The returned upper bound always has shape ``[1, queries]`` (shared across
697 the batch), so chunked consumers slice it along the sequence dimension.
698 ``causal_upper`` is the packed per-query exclusive bound shared by eager and
699 graph execution; Python ``segments`` remain available for eager callers and
700 the fused native kernels.
701 """
702 batch_size, length = query.shape[:2]
703 if causal_upper is not None:
704 return causal_upper, minimum_key_indices
705 if segments is None:
706 positions = torch.arange(query_offset, query_offset + length, device=query.device)
707 return ((positions + 1) // ratio).unsqueeze(0), minimum_key_indices
708 upper = torch.empty(length, device=query.device, dtype=torch.long)
709 lower = torch.empty_like(upper)
710 for begin, end, key_start, key_end, residual in segments:
711 first = ratio * key_end + residual - (end - begin)
712 positions = torch.arange(first, first + end - begin, device=query.device)
713 upper[begin:end] = (positions + 1) // ratio
714 lower[begin:end] = key_start
715 return upper.unsqueeze(0), lower.unsqueeze(0).expand(batch_size, -1)
718@torch.no_grad()
719def compressed_causal_topk(
720 query: torch.Tensor,
721 key: torch.Tensor,
722 merge_weight: torch.Tensor,
723 *,
724 compress_ratio: int,
725 sparse_count: int,
726 query_offset: int = 0,
727 query_chunk_size: int = 256,
728 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None = None,
729 minimum_key_indices: torch.Tensor | None = None,
730 query_segments: tuple[tuple[int, int, int, int, int], ...] | None = None,
731 causal_upper: torch.Tensor | None = None,
732 use_fused: bool = False,
733) -> torch.Tensor:
734 """Select compressed positions with V4.1's ratio-aware causal rule.
736 CSA2 keys represent closed groups of ``compress_ratio`` source tokens.
737 Both LI V2 and the batched reference preserve the source rule
738 ``key < floor((query + 1) / ratio)``. The optional native path requires
739 document-local K prefixes; the reference retains TP score reductions.
741 Args:
742 query_segments: Local Q intersections and global compressed K prefixes.
743 causal_upper: Packed per-query exclusive causal bound shared by eager
744 and graph execution; supersedes ``query_segments`` for reference
745 selection when provided.
746 use_fused: Select LI V2 and raise if a nonempty call is unsupported.
747 False explicitly selects the numerical reference.
748 """
749 if compress_ratio <= 0:
750 raise ValueError(f"compress_ratio must be positive, got {compress_ratio}")
751 if query_chunk_size <= 0:
752 raise ValueError(f"query_chunk_size must be positive, got {query_chunk_size}")
753 if query.ndim != 4:
754 raise ValueError(f"Indexer query must have shape [B, Q, H, D], got {tuple(query.shape)}")
755 batch_size, sequence_length, _, head_dim = query.shape
756 if key.ndim != 3 or key.shape[::2] != (batch_size, head_dim):
757 raise ValueError(
758 "compressed index key must have shape [batch, compressed_sequence, head_dim]"
759 )
760 if merge_weight.shape != query.shape[:3]:
761 raise ValueError(
762 f"merge_weight must have shape {tuple(query.shape[:3])}, got {tuple(merge_weight.shape)}"
763 )
764 compressed_length = key.shape[1]
765 if minimum_key_indices is not None and minimum_key_indices.shape != query.shape[:2]:
766 raise ValueError("minimum_key_indices must have shape [batch, query]")
767 top_k = min(sparse_count, compressed_length)
768 if top_k <= 0 or sequence_length == 0:
769 return torch.empty(
770 batch_size,
771 sequence_length,
772 max(top_k, 0),
773 dtype=torch.int32,
774 device=query.device,
775 )
777 if use_fused:
778 return _run_lightning_indexer(
779 query, key, merge_weight, top_k, compress_ratio, query_offset,
780 reduce_sum, minimum_key_indices, query_segments,
781 )
783 visible_keys, minimum_key_indices = _causal_key_bounds(
784 query, compress_ratio, query_offset, minimum_key_indices, query_segments, causal_upper,
785 )
786 key_fp32 = key.float().transpose(1, 2).unsqueeze(1)
787 key_positions = torch.arange(compressed_length, device=query.device).view(1, 1, -1)
788 selected_chunks = []
789 for start in range(0, sequence_length, query_chunk_size):
790 end = min(start + query_chunk_size, sequence_length)
791 # SAC may retain the matmul result; do not mutate its cached tensor.
792 scores = torch.matmul(query[:, start:end].float(), key_fp32).relu()
793 scores = (scores * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
794 if reduce_sum is not None:
795 scores = reduce_sum(scores)
796 # Bounds are [1, S]; slice the query axis so the mask covers this chunk.
797 visible = visible_keys[:, start:end].unsqueeze(-1)
798 scores.masked_fill_(key_positions >= visible, float("-inf"))
799 if minimum_key_indices is not None:
800 minimum = minimum_key_indices[:, start:end].unsqueeze(-1)
801 scores.masked_fill_(key_positions < minimum, float("-inf"))
802 top = scores.topk(top_k, dim=-1, sorted=False)
803 indices = top.indices.masked_fill(~torch.isfinite(top.values), compressed_length)
804 indices = sort_key_indices(indices, compressed_length)
805 indices = indices.masked_fill(indices == compressed_length, -1)
806 selected_chunks.append(indices.to(torch.int32))
807 return torch.cat(selected_chunks, dim=1)
810def select_candidate_blocks(
811 logits: torch.Tensor,
812 compress_lens: torch.Tensor | int,
813 topk_blocks: int,
814 block_size: int,
815) -> torch.Tensor:
816 """Select the official block-max hierarchical candidate pool.
818 The newest reachable partial block is pinned into the candidate set, then
819 the remaining blocks compete by their maximum position score. ``logits``
820 must already contain ``-inf`` at causally unreachable positions.
821 """
822 if topk_blocks <= 0:
823 raise ValueError(f"topk_blocks must be positive, got {topk_blocks}")
824 if block_size <= 0:
825 raise ValueError(f"block_size must be positive, got {block_size}")
826 width = logits.shape[-1]
827 if width == 0:
828 return torch.zeros_like(logits, dtype=torch.bool)
829 padding = -width % block_size
830 scores = torch.nn.functional.pad(logits, (0, padding), value=float("-inf"))
831 scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1)
832 num_blocks = scores.shape[-1]
833 if isinstance(compress_lens, torch.Tensor):
834 last = (compress_lens.to(device=logits.device) - 1) // block_size
835 else:
836 last = (compress_lens - 1) // block_size
837 block_positions = torch.arange(num_blocks, device=logits.device)
838 scores = scores.masked_fill(block_positions == last, float("inf"))
839 top = scores.topk(min(topk_blocks, num_blocks), dim=-1)
840 keep = torch.zeros_like(scores, dtype=torch.bool)
841 keep.scatter_(-1, top.indices, top.values > float("-inf"))
842 return keep.repeat_interleave(block_size, dim=-1)[..., :width]
845def select_candidate_block_indices(
846 logits: torch.Tensor,
847 compress_lens: torch.Tensor | int,
848 topk_blocks: int,
849 block_size: int,
850) -> torch.Tensor:
851 """Return compact block ids for the official hierarchical candidates.
853 Storing block ids uses ``block_size`` times less cross-layer state than
854 expanded position ids and avoids the full ``[batch, query, key]`` boolean
855 mask used by the released inference reference. Invalid slots are ``-1``.
856 """
857 if topk_blocks <= 0:
858 raise ValueError(f"topk_blocks must be positive, got {topk_blocks}")
859 if block_size <= 0:
860 raise ValueError(f"block_size must be positive, got {block_size}")
861 width = logits.shape[-1]
862 if width == 0:
863 return torch.empty(*logits.shape[:-1], 0, dtype=torch.int32, device=logits.device)
864 padding = -width % block_size
865 scores = torch.nn.functional.pad(logits, (0, padding), value=float("-inf"))
866 scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1)
867 num_blocks = scores.shape[-1]
868 if isinstance(compress_lens, torch.Tensor):
869 last = (compress_lens.to(device=logits.device) - 1) // block_size
870 else:
871 last = (compress_lens - 1) // block_size
872 block_positions = torch.arange(num_blocks, device=logits.device)
873 scores = scores.masked_fill(block_positions == last, float("inf"))
874 top = scores.topk(min(topk_blocks, num_blocks), dim=-1, sorted=False)
875 return top.indices.masked_fill(top.values == float("-inf"), -1).to(torch.int32)
878def compressed_causal_candidates(
879 query: torch.Tensor,
880 key: torch.Tensor,
881 merge_weight: torch.Tensor,
882 *,
883 compress_ratio: int,
884 topk_blocks: int,
885 block_size: int,
886 query_offset: int = 0,
887 query_chunk_size: int = 256,
888 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None = None,
889 minimum_key_indices: torch.Tensor | None = None,
890) -> torch.Tensor:
891 """Build compact blockwise candidate ids in bounded query chunks."""
892 _, sequence_length, _, _ = query.shape
893 compressed_length = key.shape[1]
894 if minimum_key_indices is not None and minimum_key_indices.shape != query.shape[:2]:
895 raise ValueError("minimum_key_indices must have shape [batch, query]")
896 key_fp32 = key.float().transpose(1, 2).unsqueeze(1)
897 key_positions = torch.arange(compressed_length, device=query.device).view(1, 1, -1)
898 candidate_chunks = []
899 for start in range(0, sequence_length, query_chunk_size):
900 end = min(start + query_chunk_size, sequence_length)
901 scores = torch.matmul(query[:, start:end].float(), key_fp32).relu()
902 scores = (scores * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
903 if reduce_sum is not None:
904 scores = reduce_sum(scores)
905 visible = (
906 torch.arange(query_offset + start, query_offset + end, device=query.device) + 1
907 ) // compress_ratio
908 scores.masked_fill_(key_positions >= visible.view(1, -1, 1), float("-inf"))
909 if minimum_key_indices is not None:
910 minimum = minimum_key_indices[:, start:end].unsqueeze(-1)
911 scores.masked_fill_(key_positions < minimum, float("-inf"))
912 candidate_chunks.append(
913 select_candidate_block_indices(
914 scores,
915 visible.view(1, -1, 1),
916 topk_blocks,
917 block_size,
918 )
919 )
920 return torch.cat(candidate_chunks, dim=1)
923def compressed_causal_topk_and_candidates(
924 query: torch.Tensor,
925 key: torch.Tensor,
926 merge_weight: torch.Tensor,
927 *,
928 compress_ratio: int,
929 sparse_count: int,
930 topk_blocks: int,
931 block_size: int,
932 query_offset: int = 0,
933 query_chunk_size: int = 256,
934 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None = None,
935 minimum_key_indices: torch.Tensor | None = None,
936 causal_upper: torch.Tensor | None = None,
937) -> tuple[torch.Tensor, torch.Tensor]:
938 """Select Full-layer Top-K and candidate blocks from one score pass.
940 ``causal_upper`` supplies the packed per-query causal bound for eager and
941 graph execution; it must be compression-aligned packed input when given.
942 """
943 if compress_ratio <= 0:
944 raise ValueError(f"compress_ratio must be positive, got {compress_ratio}")
945 if query_chunk_size <= 0:
946 raise ValueError(f"query_chunk_size must be positive, got {query_chunk_size}")
947 batch_size, sequence_length, _, head_dim = query.shape
948 if key.ndim != 3 or key.shape[::2] != (batch_size, head_dim):
949 raise ValueError(
950 "compressed index key must have shape [batch, compressed_sequence, head_dim]"
951 )
952 if merge_weight.shape != query.shape[:3]:
953 raise ValueError(
954 f"merge_weight must have shape {tuple(query.shape[:3])}, got {tuple(merge_weight.shape)}"
955 )
956 compressed_length = key.shape[1]
957 if minimum_key_indices is not None and minimum_key_indices.shape != query.shape[:2]:
958 raise ValueError("minimum_key_indices must have shape [batch, query]")
959 top_k = min(sparse_count, compressed_length)
960 if top_k <= 0:
961 empty_topk = torch.empty(
962 batch_size,
963 sequence_length,
964 0,
965 dtype=torch.int32,
966 device=query.device,
967 )
968 return empty_topk, empty_topk
970 key_fp32 = key.float().transpose(1, 2).unsqueeze(1)
971 key_positions = torch.arange(compressed_length, device=query.device).view(1, 1, -1)
972 selected_chunks = []
973 candidate_chunks = []
974 for start in range(0, sequence_length, query_chunk_size):
975 end = min(start + query_chunk_size, sequence_length)
976 scores = torch.matmul(query[:, start:end].float(), key_fp32).relu()
977 scores = (scores * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
978 if reduce_sum is not None:
979 scores = reduce_sum(scores)
980 visible = (
981 causal_upper[:, start:end] if causal_upper is not None else
982 (torch.arange(query_offset + start, query_offset + end, device=query.device) + 1) // compress_ratio
983 )
984 scores.masked_fill_(key_positions >= visible.view(1, -1, 1), float("-inf"))
985 if minimum_key_indices is not None:
986 minimum = minimum_key_indices[:, start:end].unsqueeze(-1)
987 scores.masked_fill_(key_positions < minimum, float("-inf"))
988 top = scores.topk(top_k, dim=-1, sorted=False)
989 indices = top.indices.masked_fill(~torch.isfinite(top.values), compressed_length)
990 indices = indices.sort(dim=-1).values
991 selected_chunks.append(indices.masked_fill(indices == compressed_length, -1).to(torch.int32))
992 candidate_chunks.append(
993 select_candidate_block_indices(
994 scores,
995 visible.view(1, -1, 1),
996 topk_blocks,
997 block_size,
998 )
999 )
1000 return torch.cat(selected_chunks, dim=1), torch.cat(candidate_chunks, dim=1)
1003def compressed_candidate_topk(
1004 query: torch.Tensor,
1005 key: torch.Tensor,
1006 merge_weight: torch.Tensor,
1007 candidate_blocks: torch.Tensor,
1008 *,
1009 compress_ratio: int,
1010 sparse_count: int,
1011 block_size: int,
1012 query_offset: int = 0,
1013 query_chunk_size: int = 256,
1014 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None = None,
1015 minimum_key_indices: torch.Tensor | None = None,
1016 causal_upper: torch.Tensor | None = None,
1017) -> torch.Tensor:
1018 """Score only hierarchical candidate blocks and return global key ids.
1020 ``causal_upper`` supplies the packed per-query causal bound for eager and
1021 graph execution; it must be compression-aligned packed input when given.
1022 """
1023 if compress_ratio <= 0:
1024 raise ValueError(f"compress_ratio must be positive, got {compress_ratio}")
1025 if block_size <= 0:
1026 raise ValueError(f"block_size must be positive, got {block_size}")
1027 if query_chunk_size <= 0:
1028 raise ValueError(f"query_chunk_size must be positive, got {query_chunk_size}")
1029 batch_size, sequence_length, _, head_dim = query.shape
1030 if key.ndim != 3 or key.shape[::2] != (batch_size, head_dim):
1031 raise ValueError(
1032 "compressed index key must have shape [batch, compressed_sequence, head_dim]"
1033 )
1034 if merge_weight.shape != query.shape[:3]:
1035 raise ValueError(
1036 f"merge_weight must have shape {tuple(query.shape[:3])}, got {tuple(merge_weight.shape)}"
1037 )
1038 if candidate_blocks.ndim != 3 or candidate_blocks.shape[:2] != query.shape[:2]:
1039 raise ValueError(
1040 "candidate_blocks must have shape [batch, query, candidate_blocks], "
1041 f"got {tuple(candidate_blocks.shape)}"
1042 )
1043 if minimum_key_indices is not None and minimum_key_indices.shape != query.shape[:2]:
1044 raise ValueError("minimum_key_indices must have shape [batch, query]")
1046 compressed_length = key.shape[1]
1047 candidate_width = candidate_blocks.shape[-1] * block_size
1048 top_k = min(sparse_count, candidate_width)
1049 if top_k <= 0:
1050 return torch.empty(
1051 batch_size,
1052 sequence_length,
1053 0,
1054 dtype=torch.int32,
1055 device=query.device,
1056 )
1058 block_offsets = torch.arange(block_size, device=query.device)
1059 batch_indices = torch.arange(batch_size, device=query.device).view(-1, 1, 1)
1060 selected_chunks = []
1061 for start in range(0, sequence_length, query_chunk_size):
1062 end = min(start + query_chunk_size, sequence_length)
1063 blocks = candidate_blocks[:, start:end].long()
1064 block_positions = blocks.unsqueeze(-1) * block_size + block_offsets
1065 positions = block_positions.flatten(-2)
1066 visible = (
1067 causal_upper[:, start:end] if causal_upper is not None else
1068 (torch.arange(query_offset + start, query_offset + end, device=query.device) + 1) // compress_ratio
1069 )
1070 valid = (blocks.unsqueeze(-1) >= 0).expand_as(block_positions).flatten(-2)
1071 valid = valid & (positions < compressed_length) & (positions < visible.view(1, -1, 1))
1072 if minimum_key_indices is not None:
1073 valid = valid & (positions >= minimum_key_indices[:, start:end].unsqueeze(-1))
1074 safe_positions = positions.clamp(min=0, max=max(compressed_length - 1, 0))
1075 selected_key = key[batch_indices, safe_positions].float()
1076 dots = torch.einsum(
1077 "bchd,bckd->bchk",
1078 query[:, start:end].float(),
1079 selected_key,
1080 ).relu()
1081 scores = (dots * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
1082 if reduce_sum is not None:
1083 scores = reduce_sum(scores)
1084 scores.masked_fill_(~valid, float("-inf"))
1085 top = scores.topk(top_k, dim=-1, sorted=False)
1086 selected = positions.gather(-1, top.indices)
1087 selected.masked_fill_(~torch.isfinite(top.values), compressed_length)
1088 selected = selected.sort(dim=-1).values
1089 selected.masked_fill_(selected == compressed_length, -1)
1090 selected_chunks.append(selected.to(torch.int32))
1091 return torch.cat(selected_chunks, dim=1)
1094class _SharedCompressedIndexerKLLoss(torch.autograd.Function):
1095 """PanGu-style precomputed-gradient KL loss for compressed index keys."""
1097 @staticmethod
1098 def forward(
1099 ctx: Any,
1100 index_query: torch.Tensor,
1101 index_key: torch.Tensor,
1102 merge_weight: torch.Tensor,
1103 attention_query: torch.Tensor,
1104 compressed_key: torch.Tensor,
1105 topk_indices: torch.Tensor,
1106 sinks: torch.Tensor,
1107 attention_scale: float,
1108 loss_coeff: float,
1109 query_chunk_size: int,
1110 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None,
1111 ) -> torch.Tensor:
1112 """Compute sparse KL and save only the three Indexer gradients."""
1113 batch_size, sequence_length, _, _ = index_query.shape
1114 denominator = batch_size * sequence_length
1115 total_loss = index_query.new_zeros((), dtype=torch.float32)
1116 grad_index_query = torch.zeros_like(index_query, dtype=torch.float32)
1117 grad_index_key = torch.zeros_like(index_key, dtype=torch.float32)
1118 grad_merge_weight = torch.zeros_like(merge_weight, dtype=torch.float32)
1120 with torch.no_grad(), torch.autocast(device_type=index_query.device.type, enabled=False):
1121 for start in range(0, sequence_length, query_chunk_size):
1122 end = min(start + query_chunk_size, sequence_length)
1123 selected = topk_indices[:, start:end]
1124 valid = selected >= 0
1125 # Graph capture forbids data-dependent control flow, so empty
1126 # rows contribute through masks instead of a `torch.any` skip:
1127 # their targets are all-zero, their gradients are all-zero, and
1128 # eager evaluates exactly the same masked expressions.
1129 safe_indices = selected.clamp_min(0).long()
1131 query_chunk = index_query[:, start:end].float()
1132 weight_chunk = merge_weight[:, start:end].float()
1133 selected_index_key = gather_selected_keys(index_key, safe_indices).float()
1134 index_dots = torch.einsum(
1135 "bcid,bckd->bcik", query_chunk, selected_index_key
1136 )
1137 index_relu = index_dots.relu()
1138 index_scores = (index_relu * weight_chunk.unsqueeze(-1)).sum(dim=2)
1139 if reduce_sum is not None:
1140 index_scores = reduce_sum(index_scores)
1141 index_scores.masked_fill_(~valid, -1.0e9)
1143 selected_attention_key = gather_selected_keys(compressed_key, safe_indices).float()
1144 attention_scores = torch.einsum(
1145 "bhcd,bckd->bhck",
1146 attention_query[:, :, start:end].float(),
1147 selected_attention_key,
1148 ) * attention_scale
1149 attention_scores.masked_fill_(~valid.unsqueeze(1), -1.0e9)
1150 sink_logits = sinks.float().view(1, -1, 1, 1).expand(
1151 batch_size, -1, end - start, -1
1152 )
1153 target = torch.cat((attention_scores, sink_logits), dim=-1).softmax(dim=-1)
1154 target = target[..., :-1].masked_fill(~valid.unsqueeze(1), 0.0).sum(dim=1)
1155 if reduce_sum is not None:
1156 target = reduce_sum(target)
1157 target = target / target.sum(dim=-1, keepdim=True).clamp_min(
1158 torch.finfo(torch.float32).tiny
1159 )
1161 log_prediction = index_scores.log_softmax(dim=-1)
1162 target_log = target.clamp_min(torch.finfo(torch.float32).tiny).log()
1163 row_loss = (target * (target_log - log_prediction)).sum(dim=-1)
1164 # Masked reduction instead of boolean indexing: identical value
1165 # in eager and traceable across runtime-dependent row validity.
1166 row_mask = valid.any(dim=-1).to(row_loss.dtype)
1167 total_loss.add_((row_loss * row_mask).sum() * (loss_coeff / denominator))
1169 grad_scores = log_prediction.exp() * target.sum(-1, keepdim=True) - target
1170 grad_scores.mul_(loss_coeff / denominator)
1171 grad_scores.masked_fill_(~valid, 0.0)
1172 grad_dots = (
1173 grad_scores.unsqueeze(2)
1174 * weight_chunk.unsqueeze(-1)
1175 * (index_dots > 0).to(index_dots.dtype)
1176 )
1177 grad_index_query[:, start:end] = torch.einsum(
1178 "bcik,bckd->bcid", grad_dots, selected_index_key
1179 )
1180 grad_merge_weight[:, start:end] = (
1181 grad_scores.unsqueeze(2) * index_relu
1182 ).sum(dim=-1)
1183 selected_key_grad = torch.einsum(
1184 "bcik,bcid->bckd", grad_dots, query_chunk
1185 )
1186 # Preserve the einsum result if SAC cached its storage.
1187 selected_key_grad = selected_key_grad.masked_fill(~valid.unsqueeze(-1), 0.0)
1188 scatter_indices = safe_indices.unsqueeze(-1).expand_as(selected_key_grad)
1189 grad_index_key.scatter_add_(
1190 1,
1191 scatter_indices.reshape(batch_size, -1, index_key.shape[-1]),
1192 selected_key_grad.reshape(batch_size, -1, index_key.shape[-1]),
1193 )
1195 ctx.save_for_backward(grad_index_query, grad_index_key, grad_merge_weight)
1196 return total_loss
1198 @staticmethod
1199 def backward(ctx: Any, grad_loss: torch.Tensor) -> tuple:
1200 """Return the precomputed Indexer gradients and detach the teacher."""
1201 grad_index_query, grad_index_key, grad_merge_weight = ctx.saved_tensors
1202 return (
1203 grad_index_query * grad_loss,
1204 grad_index_key * grad_loss,
1205 grad_merge_weight * grad_loss,
1206 *((None,) * 8),
1207 )
1210def shared_compressed_indexer_kl_loss(
1211 index_query: torch.Tensor,
1212 index_key: torch.Tensor,
1213 merge_weight: torch.Tensor,
1214 attention_query: torch.Tensor,
1215 compressed_key: torch.Tensor,
1216 topk_indices: torch.Tensor,
1217 sinks: torch.Tensor,
1218 *,
1219 attention_scale: float,
1220 loss_coeff: float,
1221 query_chunk_size: int = 256,
1222 tp_context: SharedCompressedAttentionTPContext | None = None,
1223 compress_ratio: int = 1,
1224 query_segments: tuple[tuple[int, int, int, int, int], ...] | None = None,
1225 use_fused: bool = False,
1226) -> torch.Tensor:
1227 """Compute the sparse-stage DSA KL objective on V4.1 compressed keys.
1229 The target first softmaxes selected main-attention scores per attention
1230 head, then averages and L1-normalizes across heads. Indexer inputs are
1231 detached by the caller, so this objective updates only Indexer parameters.
1233 ``use_fused`` is valid only for ordinary causal TopK selections, without
1234 candidate-pool holes. ``query_segments`` supplies each document-local K
1235 prefix; its ratio and residual encode the CP query offset for the kernel.
1236 ``use_fused=True`` raises for unsupported calls when ``loss_coeff != 0``;
1237 False explicitly selects the numerical reference. Native errors propagate.
1238 """
1239 if query_chunk_size <= 0:
1240 raise ValueError(f"query_chunk_size must be positive, got {query_chunk_size}")
1241 if not loss_coeff:
1242 return index_query.sum(dtype=torch.float32) * 0.0
1243 if index_key.shape[1] == 0:
1244 # Keep zero gradients connected even when no document contains a complete block.
1245 return (index_query[..., :0].sum(dtype=torch.float32) + index_key.sum(dtype=torch.float32)
1246 + merge_weight[..., :0].sum(dtype=torch.float32)) * 0.0
1247 if topk_indices.shape[:2] != index_query.shape[:2]:
1248 raise ValueError("topk_indices and index_query must share batch/query dimensions")
1249 if use_fused:
1250 if tp_context is not None:
1251 raise NotImplementedError("Fused KL does not support TP score/teacher reduction")
1252 if query_segments is None:
1253 raise ValueError("Fused KL requires document-local causal query_segments")
1254 validate_fused_kl_inputs(index_query, index_key, merge_weight, topk_indices)
1255 return sparse_lightning_indexer_kl_loss(
1256 index_query, index_key, merge_weight, attention_query.detach(), compressed_key.detach(),
1257 topk_indices, sinks.detach(), attention_scale, loss_coeff, query_chunk_size,
1258 compress_ratio, query_segments,
1259 )
1260 return _SharedCompressedIndexerKLLoss.apply(
1261 index_query,
1262 index_key,
1263 merge_weight,
1264 attention_query.detach(),
1265 compressed_key.detach(),
1266 topk_indices,
1267 sinks.detach(),
1268 attention_scale,
1269 loss_coeff,
1270 query_chunk_size,
1271 None if tp_context is None else tp_context.reduce_sum,
1272 )
1275class SharedCompressedDSAIndexer(nn.Module):
1276 """Reusable V4.1 compressed Lightning Indexer semantics."""
1278 # Installed by :meth:`enable_selection_exclude` for a layer that recomputes.
1279 # It is written through ``object.__setattr__`` so the exclusion region stays
1280 # out of ``modules()`` and the state dict.
1281 _selection_region: Callable[..., Any] | None = None
1282 _input_region: Callable[..., Any] | None = None
1283 _prepare_region: Callable[..., Any] | None = None
1284 _result_region: Callable[..., Any] | None = None
1285 _wait_region: Callable[..., Any] | None = None
1287 def __init__(self, module: nn.Module) -> None:
1288 """Transfer the source projections without changing checkpoint names."""
1289 super().__init__()
1290 required = ("q_b_proj", "weights_proj")
1291 missing = [name for name in required if not hasattr(module, name)]
1292 if missing:
1293 raise TypeError(f"shared compressed DSA indexer is missing: {missing}")
1294 for name, child in module._modules.items(): # pylint: disable=protected-access
1295 self.add_module(name, child)
1296 for name, parameter in module._parameters.items(): # pylint: disable=protected-access
1297 self.register_parameter(name, parameter)
1298 self.use_fused = False
1299 self.compress_ratio = module.compress_ratio
1300 self.global_num_heads = module.num_heads
1301 self.head_dim = module.head_dim
1302 self.index_topk = module.index_topk
1303 self.owns_key = module.owns_key
1304 self.is_candidate_source = module.is_candidate_source
1305 self.uses_candidates = module.uses_candidates
1306 self.candidate_topk_blocks = module.candidate_topk_blocks
1307 self.candidate_block_size = module.candidate_block_size
1308 self.loss_coeff = module.loss_coeff
1309 self.query_chunk_size = int(getattr(module, "query_chunk_size", 256))
1310 self.train(module.training)
1312 def _project_index_query(
1313 self, query_residual: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
1314 *, sequence_major: bool = False,
1315 ) -> torch.Tensor:
1316 """Project queries with their original rotary positions, including moved tasks."""
1317 query = self.q_b_proj(query_residual.detach())
1318 if query.shape[-1] % self.head_dim:
1319 raise ValueError(
1320 f"Indexer query width {query.shape[-1]} is not divisible by head_dim {self.head_dim}"
1321 )
1322 query = query.view(*query_residual.shape[:2], -1, self.head_dim)
1323 if sequence_major:
1324 return _apply_v41_rope(query, cos, sin, unsqueeze_dim=2)
1325 return _apply_v41_rope(query.transpose(1, 2), cos, sin).transpose(1, 2)
1327 @torch.no_grad()
1328 def _project_worker_index_query(
1329 self, query_residual: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
1330 owner_query_length: int,
1331 ) -> torch.Tensor:
1332 """Preserve the owner's projection row shape for discrete worker selection.
1334 BF16 GEMM rounding can depend on its row count. Match the original
1335 owner's GEMM shape so task redistribution does not perturb TopK inputs.
1336 Padding is confined to the no-grad projection and never reaches LI.
1338 Args:
1339 query_residual: Low-rank query rows assigned to this worker.
1340 cos: Original-position cosine rows for the worker queries.
1341 sin: Original-position sine rows for the worker queries.
1342 owner_query_length: Query count in the original owner projection.
1344 Returns:
1345 Worker Index Q in sequence-major layout, excluding padded rows.
1346 """
1347 if owner_query_length <= 0:
1348 raise ValueError("Indexer owner query length must be positive")
1349 batch, length = query_residual.shape[:2]
1350 if length == owner_query_length:
1351 return self._project_index_query(query_residual, cos, sin, sequence_major=True)
1352 if not length:
1353 query = self.q_b_proj(query_residual.detach())
1354 return query.view(batch, 0, query.shape[-1] // self.head_dim, self.head_dim)
1356 query = None
1357 for start in range(0, length, owner_query_length):
1358 end = min(start + owner_query_length, length)
1359 chunk = query_residual[:, start:end].detach()
1360 if end - start < owner_query_length:
1361 chunk = torch.nn.functional.pad(chunk, (0, 0, 0, owner_query_length - (end - start)))
1362 projected = self.q_b_proj(chunk)
1363 if query is None:
1364 query = projected.new_empty((batch, length, projected.shape[-1]))
1365 query[:, start:end].copy_(projected[:, :end - start])
1366 del projected, chunk
1367 query = query.view(batch, length, query.shape[-1] // self.head_dim, self.head_dim)
1368 return _apply_v41_rope(query, cos, sin, unsqueeze_dim=2)
1370 def _project_index_key(
1371 self,
1372 latent: torch.Tensor | None,
1373 cos: torch.Tensor,
1374 sin: torch.Tensor,
1375 index_key: torch.Tensor | None,
1376 cp_context: SharedCompressedAttentionCPContext | None,
1377 block_plan: Any = None,
1378 key_position_embeddings: tuple | None = None,
1379 ) -> tuple[torch.Tensor, "SequenceGatherHandle | None"]:
1380 """Project or adopt the index key and launch its CP gather."""
1381 if self.owns_key:
1382 if latent is None:
1383 raise RuntimeError("a Full Indexer requires the compressor latent")
1384 key = self.k_norm(self.wk(latent.detach()))
1385 compressed_length = key.shape[1]
1386 if key_position_embeddings is None:
1387 key_cos = cos[:, :compressed_length * self.compress_ratio:self.compress_ratio]
1388 key_sin = sin[:, :compressed_length * self.compress_ratio:self.compress_ratio]
1389 else:
1390 key_cos, key_sin = key_position_embeddings
1391 key = _apply_v41_rope(key.unsqueeze(1), key_cos, key_sin).squeeze(1)
1392 key_handle = None
1393 if cp_context is not None:
1394 key_handle = (cp_context.launch(key, 1) if block_plan is None
1395 else block_plan.launch_bank(key, cp_context))
1396 return key, key_handle
1397 if index_key is None:
1398 raise RuntimeError("a Reindex layer requires index K from the preceding Full layer")
1399 return index_key, None
1401 def forward(
1402 self,
1403 hidden_states: torch.Tensor,
1404 query_residual: torch.Tensor,
1405 latent: torch.Tensor | None,
1406 compress_position_embeddings: tuple[torch.Tensor, torch.Tensor],
1407 *,
1408 index_key: torch.Tensor | None = None,
1409 candidate_blocks: torch.Tensor | None = None,
1410 cp_context: SharedCompressedAttentionCPContext | None = None,
1411 tp_context: SharedCompressedAttentionTPContext | None = None,
1412 query_offset: int = 0,
1413 minimum_key_indices: torch.Tensor | None = None,
1414 query_segments: tuple[tuple[int, int, int, int, int], ...] | None = None,
1415 document_lengths: tuple[int, ...] = (),
1416 prepare_attention: Callable[[], None] | None = None,
1417 block_plan: Any = None,
1418 key_position_embeddings: tuple | None = None,
1419 causal_upper: torch.Tensor | None = None,
1420 ) -> SharedCompressedIndexerOutput:
1421 """Project local queries and score them against global compressed keys.
1423 When the enclosing layer is recomputed, :meth:`enable_selection_exclude`
1424 makes the discrete selection a checkpoint exclusion region: it runs in
1425 the original forward and its cached result is reused during the
1426 recomputation, so the chosen key set can never drift and the checkpoint
1427 keeps saving the same tensors in both passes.
1428 """
1429 cos, sin = compress_position_embeddings
1430 if cp_context is not None and cp_context.launch_balance is not None:
1431 return self._forward_dispatched(
1432 hidden_states, query_residual, latent, (cos, sin), index_key, cp_context,
1433 tp_context, document_lengths, prepare_attention, block_plan, key_position_embeddings,
1434 )
1435 key, key_handle = self._project_index_key(
1436 latent, cos, sin, index_key, cp_context, block_plan, key_position_embeddings,
1437 )
1438 if prepare_attention is not None:
1439 prepare_attention()
1441 query = self._project_index_query(query_residual, cos, sin)
1442 local_num_heads = query.shape[2]
1443 merge_weight = self.weights_proj(hidden_states.detach())
1444 if merge_weight.shape[-1] != local_num_heads:
1445 raise ValueError(
1446 "Indexer q_b_proj and weights_proj must use the same TP head shard: "
1447 f"{local_num_heads} query heads versus {merge_weight.shape[-1]} weights"
1448 )
1449 merge_weight = merge_weight * (self.head_dim**-0.5 * self.global_num_heads**-0.5)
1450 reduce_sum = None if tp_context is None else tp_context.reduce_sum
1451 if key_handle is not None:
1452 key = key_handle.wait()
1453 topk_indices, candidate_blocks = self._run_selection(
1454 query, key, merge_weight, query_residual, (cos, sin), cp_context, tp_context,
1455 candidate_blocks, query_offset, reduce_sum, minimum_key_indices, query_segments,
1456 document_lengths, causal_upper,
1457 )
1458 return SharedCompressedIndexerOutput(
1459 topk_indices=topk_indices,
1460 index_key=key,
1461 candidate_blocks=candidate_blocks,
1462 index_query=query,
1463 merge_weight=merge_weight,
1464 )
1466 def _forward_dispatched(
1467 self, hidden_states: torch.Tensor, query_residual: torch.Tensor, latent: torch.Tensor | None,
1468 embeddings: tuple, index_key: torch.Tensor | None, context: SharedCompressedAttentionCPContext,
1469 tp_context: SharedCompressedAttentionTPContext | None, document_lengths: tuple,
1470 prepare_attention: Callable[[], None] | None,
1471 block_plan: Any = None, key_position_embeddings: tuple | None = None,
1472 ) -> SharedCompressedIndexerOutput:
1473 """Overlap worker inputs and key gathering while retaining owner-side KL edges."""
1474 cos, sin = embeddings
1475 if tp_context is not None or self.is_candidate_source or self.uses_candidates:
1476 raise NotImplementedError("Indexer overlap requires ordinary causal selection and TP=1")
1477 weights = self.weights_proj(hidden_states.detach())
1478 weights = weights * (self.head_dim**-0.5 * self.global_num_heads**-0.5)
1479 launch = self._launch_selection if self._input_region is None else self._input_region
1480 wave = launch(query_residual, weights, (cos, sin), document_lengths, context)
1481 key, handle = self._project_index_key(
1482 latent, cos, sin, index_key, context, block_plan, key_position_embeddings,
1483 )
1484 if prepare_attention is not None:
1485 prepare_attention()
1486 prepare = self._prepare_dispatched if self._prepare_region is None else self._prepare_region
1487 prepared = prepare(wave) if context.indexer_priority else None
1488 # KL uses this autograd edge; only selection transport runs under no_grad.
1489 if handle is not None:
1490 key = handle.wait()
1491 if not context.indexer_priority:
1492 prepared = prepare(wave)
1493 select = self._select_dispatched if self._result_region is None else self._result_region
1494 result = select(prepared, key)
1495 del prepared
1496 query = self._project_index_query(query_residual, cos, sin)
1497 if weights.shape[-1] != query.shape[2]:
1498 raise ValueError("Indexer query and merge weights must have matching head counts")
1499 finish = self._finish_selection if self._wait_region is None else self._wait_region
1500 return SharedCompressedIndexerOutput(finish(result), key, None, query, weights)
1502 def enable_selection_exclude(self) -> None:
1503 """Keep discrete selection fixed when the enclosing layer recomputes."""
1504 object.__setattr__(
1505 self, "_selection_region", checkpoint_exclude_wrapper(self._select_tokens)
1506 )
1507 object.__setattr__(self, "_input_region", checkpoint_exclude_wrapper(self._launch_selection))
1508 # This output goes straight into an exclusion region; retain IDs, not expanded worker Q.
1509 object.__setattr__(self, "_prepare_region", checkpoint_exclude_wrapper( # pylint: disable=E1123
1510 self._prepare_dispatched, save_output=False,
1511 ))
1512 object.__setattr__(self, "_result_region", checkpoint_exclude_wrapper(self._select_dispatched))
1513 object.__setattr__(self, "_wait_region", checkpoint_exclude_wrapper(self._finish_selection))
1515 @torch.no_grad()
1516 def _launch_selection(self, qr: torch.Tensor, weights: torch.Tensor, embeddings: tuple,
1517 lengths: tuple, context: SharedCompressedAttentionCPContext) -> Any:
1518 """Cache the input launch so checkpoint replay does not repeat communication."""
1519 return context.launch_balance(qr, weights, embeddings, lengths, self.compress_ratio)
1521 @torch.no_grad()
1522 def _prepare_dispatched(self, wave: Any) -> tuple:
1523 """Expand received qr while Index K may still be in flight."""
1524 qr, weights, cos, sin = wave.wait()
1525 query = self._project_worker_index_query(qr, cos, sin, wave.owner_queries)
1526 return query, weights, wave
1528 @torch.no_grad()
1529 def _select_dispatched(self, prepared: tuple, key: torch.Tensor) -> Any:
1530 """Consume worker Q once and launch the inverse ID transfer."""
1531 query, weights, wave = prepared
1532 selected = self._select_indices(query, key, weights, None, 0, None, None, wave.route.segments)[0]
1533 return wave.launch_result(selected[0])
1535 @staticmethod
1536 @torch.no_grad()
1537 def _finish_selection(result: Any) -> torch.Tensor:
1538 """Cache completed IDs; replay never consumes an already finished handle."""
1539 return result.wait().unsqueeze(0)
1541 def _run_selection(self, *args: Any, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor | None]:
1542 """Select directly, or reuse the result retained by the checkpoint region."""
1543 selection = self._select_tokens if self._selection_region is None else self._selection_region
1544 return selection(*args, **kwargs)
1546 def _select_tokens(
1547 self, query: torch.Tensor, key: torch.Tensor, merge_weight: torch.Tensor,
1548 query_residual: torch.Tensor, embeddings: tuple,
1549 cp_context: SharedCompressedAttentionCPContext | None,
1550 tp_context: SharedCompressedAttentionTPContext | None,
1551 candidate_blocks: torch.Tensor | None, query_offset: int,
1552 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None,
1553 minimum_key_indices: torch.Tensor | None, query_segments: tuple | None,
1554 document_lengths: tuple, causal_upper: torch.Tensor | None = None,
1555 ) -> tuple[torch.Tensor, torch.Tensor | None]:
1556 """Cache dispatch, discrete selection and TopK return as one recompute region."""
1557 cos, sin = embeddings
1558 if cp_context is not None and cp_context.balance_indexer is not None:
1559 if tp_context is not None or self.is_candidate_source or self.uses_candidates:
1560 raise NotImplementedError("Indexer load balance currently requires ordinary causal selection and TP=1")
1562 def select_tasks(task_qr: torch.Tensor, task_weights: torch.Tensor,
1563 task_embeddings: tuple, segments: tuple) -> torch.Tensor:
1564 """Project on the receiver while this layer's parameters are resident."""
1565 task_query = self._project_worker_index_query(
1566 task_qr, *task_embeddings, query_residual.shape[1],
1567 )
1568 return self._select_indices(
1569 task_query, key, task_weights, None, 0, None, None, segments,
1570 )[0]
1572 topk_indices = cp_context.balance_indexer(
1573 query_residual, merge_weight, (cos, sin), select_tasks, document_lengths, self.compress_ratio,
1574 )
1575 candidate_blocks = None
1576 else:
1577 topk_indices, candidate_blocks = self._select_indices(
1578 query, key, merge_weight, candidate_blocks, query_offset,
1579 reduce_sum, minimum_key_indices, query_segments, causal_upper,
1580 )
1581 return topk_indices, candidate_blocks
1583 def _select_indices(
1584 self, query: torch.Tensor, key: torch.Tensor, merge_weight: torch.Tensor,
1585 candidate_blocks: torch.Tensor | None, query_offset: int,
1586 reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None,
1587 minimum_key_indices: torch.Tensor | None, query_segments: tuple | None,
1588 causal_upper: torch.Tensor | None = None,
1589 ) -> tuple[torch.Tensor, torch.Tensor | None]:
1590 """Select the requested implementation without changing candidate policies."""
1591 if self.use_fused and (self.is_candidate_source or self.uses_candidates):
1592 raise NotImplementedError("V4.1 fused Indexer does not support candidate-pool selection")
1593 if self.is_candidate_source:
1594 topk_indices, candidate_blocks = compressed_causal_topk_and_candidates(
1595 query,
1596 key,
1597 merge_weight,
1598 compress_ratio=self.compress_ratio,
1599 sparse_count=self.index_topk,
1600 topk_blocks=self.candidate_topk_blocks,
1601 block_size=self.candidate_block_size,
1602 query_offset=query_offset,
1603 query_chunk_size=self.query_chunk_size,
1604 reduce_sum=reduce_sum,
1605 minimum_key_indices=minimum_key_indices,
1606 causal_upper=causal_upper,
1607 )
1608 elif self.uses_candidates:
1609 if candidate_blocks is None:
1610 raise RuntimeError("a hierarchical Reindex layer requires candidate block ids")
1611 topk_indices = compressed_candidate_topk(
1612 query,
1613 key,
1614 merge_weight,
1615 candidate_blocks,
1616 compress_ratio=self.compress_ratio,
1617 sparse_count=self.index_topk,
1618 block_size=self.candidate_block_size,
1619 query_offset=query_offset,
1620 query_chunk_size=self.query_chunk_size,
1621 reduce_sum=reduce_sum,
1622 minimum_key_indices=minimum_key_indices,
1623 causal_upper=causal_upper,
1624 )
1625 else:
1626 candidate_blocks = None
1627 topk_indices = compressed_causal_topk(
1628 query,
1629 key,
1630 merge_weight,
1631 compress_ratio=self.compress_ratio,
1632 sparse_count=self.index_topk,
1633 query_offset=query_offset,
1634 query_chunk_size=self.query_chunk_size,
1635 reduce_sum=reduce_sum,
1636 minimum_key_indices=minimum_key_indices,
1637 query_segments=query_segments,
1638 causal_upper=causal_upper,
1639 use_fused=self.use_fused,
1640 )
1641 return topk_indices, candidate_blocks
1644def reference_sparse_attention(
1645 query: torch.Tensor,
1646 key_value: torch.Tensor,
1647 sparse_indices: torch.Tensor,
1648 sinks: torch.Tensor,
1649 scale: float,
1650) -> torch.Tensor:
1651 """Compute sink attention with eager Torch operations on any device.
1653 Args:
1654 query: Main queries [batch, heads, local queries, head dimension].
1655 key_value: Shared K=V bank [batch, 1, keys, head dimension].
1656 sparse_indices: Visible bank IDs [batch, local queries, selections];
1657 negative entries are padding and repeated IDs count only once.
1658 sinks: Learnable scalar sink logits [heads], with zero-valued output.
1659 scale: Main attention logit scale.
1661 Returns:
1662 Attention values [batch, local queries, heads, head dimension].
1663 Softmax uses FP32; matrix multiplications retain the input precision.
1665 Note:
1666 This reference materializes a dense query/key mask and scores. It is
1667 intended for explicit unfused comparisons, not long-sequence efficiency.
1668 """
1669 batch_size, num_heads, sequence_length, _ = query.shape
1670 key_length = key_value.shape[2]
1671 valid = sparse_indices >= 0
1672 safe_indices = sparse_indices.clamp_min(0).long()
1673 allowed_counts = torch.zeros(
1674 batch_size,
1675 sequence_length,
1676 key_length,
1677 dtype=torch.int32,
1678 device=query.device,
1679 )
1680 allowed_counts.scatter_add_(2, safe_indices, valid.to(torch.int32))
1681 allowed = allowed_counts > 0
1682 logits = torch.matmul(query, key_value.transpose(-1, -2)) * scale
1683 logits = logits.masked_fill(~allowed.unsqueeze(1), float("-inf"))
1684 sink_logits = sinks.view(1, num_heads, 1, 1).expand(batch_size, -1, sequence_length, -1)
1685 probabilities = torch.cat((logits, sink_logits), dim=-1).softmax(dim=-1, dtype=torch.float32)
1686 return torch.matmul(probabilities[..., :-1].to(key_value.dtype), key_value).transpose(1, 2)
1689class SharedCompressedDSAAttentionBase(nn.Module):
1690 """Shared V4.1 attention semantics with selectable reference execution."""
1692 def __init__(
1693 self,
1694 module: nn.Module,
1695 *,
1696 replace_indexer: bool,
1697 use_fused_ops: bool,
1698 ) -> None:
1699 """Transfer source state and configure the execution implementation."""
1700 super().__init__()
1701 for name, child in module._modules.items(): # pylint: disable=protected-access
1702 self.add_module(name, child)
1703 for name, parameter in module._parameters.items(): # pylint: disable=protected-access
1704 self.register_parameter(name, parameter)
1705 indexer = getattr(self, "indexer", None)
1706 if (
1707 replace_indexer
1708 and indexer is not None
1709 and not isinstance(indexer, SharedCompressedDSAIndexer)
1710 ):
1711 self.indexer = SharedCompressedDSAIndexer(indexer)
1712 self.config = module.config
1713 self.layer_idx = module.layer_idx
1714 self.num_heads = module.num_heads
1715 self.num_key_value_groups = module.num_key_value_groups
1716 self.num_groups = module.config.o_groups
1717 self.compress_ratio = module.compress_ratio
1718 self.is_kv_source = module.is_kv_source
1719 self.is_index_source = module.is_index_source
1720 self.kv_source_layer_idx = module.kv_source_layer_idx
1721 self.index_source_layer_idx = module.index_source_layer_idx
1722 self.candidate_source_layer_idx = module.candidate_source_layer_idx
1723 self.head_dim = module.head_dim
1724 self.rope_head_dim = module.config.qk_rope_head_dim
1725 self.sliding_window = module.sliding_window
1726 self.scaling = module.scaling
1727 self.use_fused_ops = use_fused_ops
1728 if self.is_index_source:
1729 self.indexer.use_fused = use_fused_ops
1730 self.train(module.training)
1732 @staticmethod
1733 def _validate_cp_inputs(
1734 cp_context: SharedCompressedAttentionCPContext,
1735 position_ids: torch.Tensor,
1736 sequence_length: int,
1737 ) -> tuple[int, int]:
1738 """Validate contiguous CP shards and return global query geometry."""
1739 if cp_context.size <= 1 or not 0 <= cp_context.rank < cp_context.size:
1740 raise ValueError(
1741 "shared compressed attention CP requires size > 1 and "
1742 f"0 <= rank < size, got rank={cp_context.rank}, size={cp_context.size}"
1743 )
1744 query_offset = cp_context.rank * sequence_length
1745 if not torch.all(position_ids[..., 0] == query_offset):
1746 raise ValueError(
1747 "shared compressed attention CP requires contiguous global position_ids; "
1748 f"rank {cp_context.rank} expected first position {query_offset}"
1749 )
1750 return query_offset, sequence_length * cp_context.size
1752 def _project_raw_kv(
1753 self,
1754 hidden_states: torch.Tensor,
1755 cos: torch.Tensor,
1756 sin: torch.Tensor,
1757 cp_context: SharedCompressedAttentionCPContext | None,
1758 raw_route: SequenceDispatchRoute | None = None,
1759 ) -> tuple[torch.Tensor | None, SequenceGatherHandle | None]:
1760 """Project raw K=V and launch CP communication before query GEMMs."""
1761 key_value = self.kv_norm(self.kv_proj(hidden_states)).unsqueeze(1)
1762 key_value = _apply_v41_rope(key_value, cos, sin)
1763 if cp_context is None:
1764 return key_value, None
1765 if raw_route is not None:
1766 return None, cp_context.launch_raw(key_value, 2, raw_route)
1767 return None, cp_context.launch(key_value, 2)
1769 def _project_query(
1770 self,
1771 hidden_states: torch.Tensor,
1772 cos: torch.Tensor,
1773 sin: torch.Tensor,
1774 ) -> tuple[torch.Tensor, torch.Tensor]:
1775 """Project and normalize the local query heads."""
1776 query_residual = self.q_a_norm(self.q_a_proj(hidden_states))
1777 return query_residual, self._expand_query(query_residual, cos, sin)
1779 def _expand_query(self, query_residual: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
1780 """Expand main heads after low-rank query inputs have been dispatched."""
1781 query = self.q_b_proj(query_residual).view(*query_residual.shape[:2], -1, self.head_dim)
1782 # Preserve sequence-major storage for sparse attention while keeping the head-major API.
1783 return _apply_v41_rope(query, cos, sin, unsqueeze_dim=2).transpose(1, 2)
1785 @staticmethod
1786 def _validate_forward_inputs(
1787 hidden_states: torch.Tensor,
1788 attention_mask: torch.Tensor | None,
1789 past_key_values: Any | None,
1790 kwargs: dict[str, Any],
1791 ) -> tuple[
1792 "SharedCompressedAttentionState",
1793 "SharedCompressedPackedSequence | None",
1794 "SharedCompressedAttentionCPContext | None",
1795 "SharedCompressedAttentionTPContext | None",
1796 int,
1797 int,
1798 ]:
1799 """Pop and validate the shared-attention keyword arguments."""
1800 if past_key_values is not None:
1801 raise NotImplementedError("shared compressed DSA training does not support KV cache")
1802 shared_state = kwargs.pop("shared_attention_state", None)
1803 if not isinstance(shared_state, SharedCompressedAttentionState):
1804 raise ValueError("shared_attention_state is required for shared compressed DSA")
1805 if attention_mask is not None:
1806 raise NotImplementedError(
1807 "shared compressed DSA consumes compact packed_seq_params instead of a dense mask"
1808 )
1809 batch_size, sequence_length, _ = hidden_states.shape
1810 packed_sequence = kwargs.pop("packed_seq_params", None)
1811 if packed_sequence is not None and not isinstance(
1812 packed_sequence, SharedCompressedPackedSequence
1813 ):
1814 raise TypeError(
1815 "packed_seq_params must be SharedCompressedPackedSequence, "
1816 f"got {type(packed_sequence).__name__}"
1817 )
1818 cp_context = kwargs.pop("shared_attention_cp_context", None)
1819 if cp_context is not None and not isinstance(cp_context, SharedCompressedAttentionCPContext):
1820 raise TypeError(
1821 "shared_attention_cp_context must be SharedCompressedAttentionCPContext, "
1822 f"got {type(cp_context).__name__}"
1823 )
1824 tp_context = kwargs.pop("shared_attention_tp_context", None)
1825 if tp_context is not None and not isinstance(tp_context, SharedCompressedAttentionTPContext):
1826 raise TypeError(
1827 "shared_attention_tp_context must be SharedCompressedAttentionTPContext, "
1828 f"got {type(tp_context).__name__}"
1829 )
1830 return shared_state, packed_sequence, cp_context, tp_context, batch_size, sequence_length
1832 def _resolve_packed_segments(
1833 self,
1834 packed_sequence: "SharedCompressedPackedSequence | None",
1835 batch_size: int,
1836 sequence_length: int,
1837 query_offset: int,
1838 global_sequence_length: int,
1839 device: torch.device,
1840 ) -> tuple[torch.Tensor | None, torch.Tensor | None, tuple, torch.Tensor | None]:
1841 """Validate packed document boundaries and derive local Q/K geometry.
1843 Both execution modes consume the same per-token tensors: document
1844 starts, per-query compressed-key minima and the packed causal upper
1845 bound. Python ``query_segments`` stay available for the eager fused
1846 kernels only; graph packing inputs never carry them.
1847 """
1848 segment_starts = None
1849 minimum_key_indices = None
1850 causal_upper = None
1851 ratio = max(1, self.compress_ratio)
1852 raw_end = query_offset + sequence_length
1853 query_segments = ((0, sequence_length, 0, raw_end // ratio, raw_end % ratio),)
1854 if packed_sequence is not None:
1855 if batch_size != 1:
1856 raise ValueError("V4.1 compact packed attention currently requires micro_batch_size=1")
1857 if (
1858 packed_sequence.local_query_start != query_offset
1859 or packed_sequence.local_query_length != sequence_length
1860 or packed_sequence.global_sequence_length != global_sequence_length
1861 ):
1862 packed_shard = (
1863 packed_sequence.local_query_start,
1864 packed_sequence.local_query_length,
1865 packed_sequence.global_sequence_length,
1866 )
1867 raise ValueError(
1868 "packed sequence geometry does not match the local CP shard: "
1869 f"packed={packed_shard}, "
1870 f"attention={(query_offset, sequence_length, global_sequence_length)}"
1871 )
1872 prepared = packed_sequence.prepare(device, (self.compress_ratio,))
1873 segment_starts = prepared.local_segment_starts(device)
1874 if self.compress_ratio:
1875 minimum_key_indices = prepared.minimum_key_indices(device, self.compress_ratio)
1876 causal_upper = prepared.causal_key_upper_bounds(device, self.compress_ratio)
1877 if (self.candidate_source_layer_idx is not None
1878 and not prepared.is_compression_aligned(ratio)):
1879 raise ValueError("candidate-pool packed attention still requires compression-aligned documents")
1880 query_segments = (
1881 None if prepared._restored_from_packing_tensors # pylint: disable=protected-access
1882 else prepared.indexer_segments(ratio)
1883 )
1884 return segment_starts, minimum_key_indices, query_segments, causal_upper
1886 def _prepare_compression(
1887 self, packed: SharedCompressedPackedSequence | None, context: SharedCompressedAttentionCPContext | None,
1888 embeddings: dict, global_length: int, device: torch.device,
1889 ) -> tuple:
1890 """Prepare block-start positions only for a source that needs document-local grouping."""
1891 if not self.is_kv_source or self.compress_ratio <= 1:
1892 return None, None
1893 size, rank = (1, 0) if context is None else (context.size, context.rank)
1894 if packed is None:
1895 local_length = global_length // size
1896 if local_length % self.compress_ratio == 0:
1897 return None, None
1898 packed = SharedCompressedPackedSequence(
1899 torch.tensor((0, global_length)), rank * local_length, local_length, global_length,
1900 )
1901 packed = packed.prepare(device, (self.compress_ratio,))
1902 plan = packed.block_plan(self.compress_ratio, size, rank)
1903 positions = None if plan is None else packed.compressed_positions(plan, embeddings["compress"], context)
1904 return plan, positions
1906 def _plan_bank_dispatch(
1907 self, packed_sequence: SharedCompressedPackedSequence | None,
1908 cp_context: SharedCompressedAttentionCPContext | None,
1909 sequence_length: int, global_sequence_length: int, device: torch.device,
1910 ) -> tuple[SequenceDispatchRoute | None, SequenceDispatchRoute | None, tuple[int, ...]]:
1911 """Reuse document boundaries to route each bank in its own physical coordinates.
1913 Boundary reads happen only for CP communication planning; without a CP
1914 context the packed metadata carries no Python document list and the
1915 per-document token accounting lives entirely in tensors.
1916 """
1917 raw_route, main_route = None, None
1918 if cp_context is None:
1919 return raw_route, main_route, ()
1920 boundaries = (0, global_sequence_length)
1921 if packed_sequence is not None:
1922 boundaries = packed_sequence.prepare(device, (self.compress_ratio,)).document_boundaries()
1923 if cp_context.launch_raw is not None:
1924 raw_route = sliding_window_route(
1925 boundaries, sequence_length, cp_context.size, cp_context.rank, self.sliding_window,
1926 )
1927 if self.compress_ratio and cp_context.launch_main is not None:
1928 main_route = document_prefix_route(
1929 boundaries, sequence_length, cp_context.size, cp_context.rank, self.compress_ratio,
1930 )
1931 return raw_route, main_route, tuple(end - start for start, end in zip(boundaries, boundaries[1:]))
1933 def _compute_compressed_kv(
1934 self,
1935 hidden_states: torch.Tensor,
1936 position_embeddings: dict[str, tuple[torch.Tensor, torch.Tensor]],
1937 shared_state: "SharedCompressedAttentionState",
1938 cp_context: "SharedCompressedAttentionCPContext | None",
1939 main_route: SequenceDispatchRoute | None = None,
1940 block_plan: Any = None,
1941 key_position_embeddings: tuple | None = None,
1942 defer_communication: bool = False,
1943 ) -> tuple[torch.Tensor | None, "SequenceGatherHandle | None"]:
1944 """Compress local hidden states and publish or launch the compressed KV.
1946 The compressed K=V reaches its consumers - the producing Full layer
1947 included - through the shared state, so only the publish (or the CP
1948 gather handle that feeds it) leaves this call.
1949 """
1950 latent = None
1951 compressed_handle = None
1952 if self.is_kv_source:
1953 if block_plan is None:
1954 latent, compressed = self.compressor(hidden_states, position_embeddings["compress"])
1955 else:
1956 latent, compressed = self.compressor(hidden_states, key_position_embeddings, block_plan, cp_context)
1957 if cp_context is None:
1958 if not is_recomputing():
1959 shared_state.publish_compressed_kv(self.layer_idx, compressed)
1960 else:
1961 launch = partial(self._launch_compressed_bank, compressed, cp_context, main_route, block_plan)
1962 compressed_handle = _DeferredBankLaunch(launch) if defer_communication else launch()
1963 return latent, compressed_handle
1965 @staticmethod
1966 def _launch_compressed_bank(compressed: torch.Tensor, context: SharedCompressedAttentionCPContext,
1967 route: SequenceDispatchRoute | None, block_plan: Any) -> SequenceGatherHandle:
1968 """Use the configured physical bank layout at the scheduled launch point."""
1969 if route is not None:
1970 return context.launch_main(compressed, 1, route)
1971 if block_plan is not None:
1972 return block_plan.launch_bank(compressed, context)
1973 return context.launch(compressed, 1)
1975 def _run_indexer(
1976 self,
1977 hidden_states: torch.Tensor,
1978 query_residual: torch.Tensor,
1979 latent: torch.Tensor | None,
1980 position_embeddings: dict[str, tuple[torch.Tensor, torch.Tensor]],
1981 shared_state: "SharedCompressedAttentionState",
1982 cp_context: "SharedCompressedAttentionCPContext | None",
1983 tp_context: "SharedCompressedAttentionTPContext | None",
1984 query_offset: int,
1985 minimum_key_indices: torch.Tensor | None,
1986 query_segments: tuple,
1987 document_lengths: tuple[int, ...],
1988 block_plan: Any = None,
1989 key_position_embeddings: tuple | None = None,
1990 prepare_attention: Callable[[], None] | None = None,
1991 causal_upper: torch.Tensor | None = None,
1992 ) -> Any | None:
1993 """Run the indexer on index-source layers and publish its forward outputs."""
1994 if not self.is_index_source:
1995 return None
1996 index_key = (
1997 None
1998 if self.indexer.owns_key
1999 else shared_state.require_index_key(self.kv_source_layer_idx, self.layer_idx)
2000 )
2001 candidate_blocks = (
2002 shared_state.require_candidate_blocks(
2003 self.candidate_source_layer_idx,
2004 self.layer_idx,
2005 )
2006 if self.indexer.uses_candidates else None
2007 )
2008 indexer_output = self.indexer(
2009 hidden_states,
2010 query_residual,
2011 latent,
2012 position_embeddings["compress"],
2013 prepare_attention=prepare_attention,
2014 index_key=index_key,
2015 candidate_blocks=candidate_blocks,
2016 cp_context=cp_context,
2017 tp_context=tp_context,
2018 query_offset=query_offset,
2019 minimum_key_indices=minimum_key_indices,
2020 query_segments=query_segments,
2021 block_plan=block_plan,
2022 key_position_embeddings=key_position_embeddings,
2023 document_lengths=document_lengths,
2024 causal_upper=causal_upper,
2025 )
2026 if is_recomputing():
2027 # A recompute must not publish again: its consumers read the entries
2028 # the forward pass wrote, and a second write would break the
2029 # single-publisher contract.
2030 return indexer_output
2031 if self.indexer.owns_key:
2032 shared_state.publish_index_key(self.layer_idx, indexer_output.index_key)
2033 shared_state.publish_topk_indices(self.layer_idx, indexer_output.topk_indices)
2034 if indexer_output.candidate_blocks is not None:
2035 shared_state.publish_candidate_blocks(
2036 self.layer_idx,
2037 indexer_output.candidate_blocks,
2038 )
2039 return indexer_output
2041 def _finish_bank_communication(
2042 self, key_value: torch.Tensor | None, raw_kv_handle: SequenceGatherHandle | None,
2043 compressed_handle: SequenceGatherHandle | None, shared_state: SharedCompressedAttentionState,
2044 main_route: SequenceDispatchRoute | None,
2045 ) -> torch.Tensor:
2046 """Resolve differentiable banks only when attention and KL can consume them."""
2047 if raw_kv_handle is not None:
2048 key_value = raw_kv_handle.wait()
2049 if compressed_handle is not None:
2050 compressed = compressed_handle.wait()
2051 if not is_recomputing():
2052 shared_state.publish_compressed_kv(
2053 self.layer_idx, compressed, 0 if main_route is None else main_route.global_start,
2054 )
2055 if key_value is None:
2056 raise RuntimeError("raw KV communication did not produce a tensor")
2058 return key_value
2060 def _build_sparse_indices(
2061 self,
2062 key_value: torch.Tensor,
2063 compressed_kv: torch.Tensor | None,
2064 topk_indices: torch.Tensor | None,
2065 sequence_length: int,
2066 query_offset: int,
2067 global_sequence_length: int,
2068 segment_starts: torch.Tensor | None,
2069 raw_key_offset: int = 0,
2070 main_key_offset: int = 0,
2071 ) -> tuple[torch.Tensor, torch.Tensor]:
2072 """Assemble dense-reference banks only after reference dispatch is selected."""
2073 window = build_sliding_window_indices(
2074 key_value.shape[0], sequence_length, self.sliding_window, key_value.device,
2075 query_offset=query_offset, key_length=global_sequence_length,
2076 )
2077 if segment_starts is not None:
2078 window.masked_fill_(window < segment_starts.unsqueeze(-1), -1)
2079 window = torch.where(window >= 0, window - raw_key_offset, -1)
2080 if compressed_kv is None:
2081 return key_value, window
2082 combined_key_value = torch.cat((key_value, compressed_kv.unsqueeze(1)), dim=2)
2083 compressed_indices = torch.where(
2084 topk_indices >= 0, topk_indices.long() - main_key_offset + key_value.shape[2], -1,
2085 )
2086 return combined_key_value, torch.cat((window, compressed_indices), dim=-1)
2088 def _apply_indexer_loss(
2089 self,
2090 indexer_output: Any | None,
2091 query: torch.Tensor,
2092 compressed_kv: torch.Tensor | None,
2093 tp_context: "SharedCompressedAttentionTPContext | None",
2094 query_segments: tuple,
2095 main_key_offset: int = 0,
2096 ) -> torch.Tensor:
2097 """Auto-scale the query by the indexer KL loss when configured."""
2098 if (
2099 indexer_output is not None
2100 and self.training
2101 and self.indexer.loss_coeff
2102 and compressed_kv is not None
2103 ):
2104 # KL uses the original owner's queries. A view aligns its student K
2105 # with the compact teacher bank while preserving global Index K gradients.
2106 index_key = indexer_output.index_key
2107 if main_key_offset or compressed_kv.shape[1] != index_key.shape[1]:
2108 index_key = index_key[:, main_key_offset:main_key_offset+compressed_kv.shape[1]]
2109 selected = indexer_output.topk_indices
2110 if main_key_offset:
2111 selected = torch.where(selected >= 0, selected - main_key_offset, -1)
2112 query_segments = tuple((s, e, ks-main_key_offset, ke-main_key_offset, residual)
2113 for s, e, ks, ke, residual in query_segments)
2114 indexer_loss = shared_compressed_indexer_kl_loss(
2115 indexer_output.index_query,
2116 index_key,
2117 indexer_output.merge_weight,
2118 query,
2119 compressed_kv,
2120 selected,
2121 self.sinks,
2122 attention_scale=self.scaling,
2123 loss_coeff=self.indexer.loss_coeff,
2124 query_chunk_size=self.indexer.query_chunk_size,
2125 tp_context=tp_context,
2126 compress_ratio=self.compress_ratio,
2127 query_segments=query_segments,
2128 use_fused=self.use_fused_ops,
2129 )
2130 query = aux_loss_auto_scale(query, indexer_loss)
2131 return query
2133 def _run_sparse_attention(
2134 self, query: torch.Tensor, key_value: torch.Tensor, compressed_kv: torch.Tensor | None,
2135 topk_indices: torch.Tensor | None, query_offset: int, global_sequence_length: int,
2136 segment_starts: torch.Tensor | None, query_segments: tuple, raw_key_offset: int = 0,
2137 main_key_offset: int = 0, raw_document_starts: tuple = (), global_key_length: int | None = None,
2138 ) -> torch.Tensor:
2139 """Run the chosen implementation; unsupported explicit fusion always raises."""
2140 if self.use_fused_ops:
2141 expected_topk = min(512, global_key_length if global_key_length is not None
2142 else global_sequence_length // max(1, self.compress_ratio))
2143 if not (self.candidate_source_layer_idx is None
2144 and (topk_indices is None or topk_indices.shape[-1] == expected_topk)
2145 and supports_fused_attention_inputs(
2146 query, self.compress_ratio, self.sliding_window,
2147 0 if topk_indices is None else topk_indices.shape[-1])):
2148 raise RuntimeError(
2149 "V4.1 fused attention is required but this call is unsupported: "
2150 f"layer={self.layer_idx}, device={query.device}, dtype={query.dtype}, Q={tuple(query.shape)}, "
2151 f"ratio={self.compress_ratio}, window={self.sliding_window}. "
2152 "Check the SMLA training extension and native input contract; no reference fallback was run."
2153 )
2154 output = sparse_flash_mla(
2155 query, key_value, compressed_kv, topk_indices, self.sinks, self.scaling,
2156 self.compress_ratio, self.sliding_window, query_offset, query_segments, raw_key_offset, main_key_offset,
2157 raw_document_starts,
2158 )
2159 if compressed_kv is not None and all(end == start for _, _, start, end, _ in query_segments):
2160 # Empty-prefix ranks must still participate in the shared bank's collective backward.
2161 output = output + compressed_kv[:, :0].sum().to(output.dtype)
2162 return output
2163 combined_key_value, sparse_indices = self._build_sparse_indices(
2164 key_value, compressed_kv, topk_indices, query.shape[2], query_offset,
2165 global_sequence_length, segment_starts, raw_key_offset, main_key_offset,
2166 )
2167 return reference_sparse_attention(query, combined_key_value, sparse_indices, self.sinks, self.scaling)
2169 def _project_output(
2170 self,
2171 attention_output: torch.Tensor,
2172 cos: torch.Tensor,
2173 sin: torch.Tensor,
2174 batch_size: int,
2175 sequence_length: int,
2176 ) -> torch.Tensor:
2177 """Undo RoPE and apply the grouped low-rank output projection."""
2178 attention_output = _apply_v41_rope(attention_output, cos, -sin, unsqueeze_dim=2)
2179 grouped = attention_output.reshape(batch_size, sequence_length, self.num_groups, -1)
2180 hidden_per_group = grouped.shape[-1]
2181 output_per_group = self.o_a_proj.weight.shape[0] // self.num_groups
2182 projection = self.o_a_proj.weight.view(
2183 self.num_groups, output_per_group, hidden_per_group
2184 ).transpose(1, 2)
2185 flattened = grouped.reshape(-1, self.num_groups, hidden_per_group).transpose(0, 1)
2186 projected = torch.bmm(flattened, projection).transpose(0, 1)
2187 projected = projected.reshape(batch_size, sequence_length, -1)
2188 return self.o_b_proj(projected)
2190 def forward(
2191 self,
2192 hidden_states: torch.Tensor,
2193 position_embeddings: dict[str, tuple[torch.Tensor, torch.Tensor]],
2194 position_ids: torch.Tensor,
2195 attention_mask: torch.Tensor | None,
2196 past_key_values: Any | None = None,
2197 **kwargs: Any,
2198 ) -> tuple[torch.Tensor, torch.Tensor | None]:
2199 """Publish or consume compressed KV and execute indexed K=V attention."""
2200 (
2201 shared_state,
2202 packed_sequence,
2203 cp_context,
2204 tp_context,
2205 batch_size,
2206 sequence_length,
2207 ) = self._validate_forward_inputs(hidden_states, attention_mask, past_key_values, kwargs)
2208 query_offset = 0
2209 global_sequence_length = sequence_length
2210 if cp_context is not None:
2211 query_offset, global_sequence_length = self._validate_cp_inputs(
2212 cp_context,
2213 position_ids,
2214 sequence_length,
2215 )
2216 if packed_sequence is not None:
2217 packed_sequence = packed_sequence.prepare(hidden_states.device, (self.compress_ratio,))
2218 if (self.use_fused_ops and packed_sequence._restored_from_packing_tensors): # pylint: disable=protected-access
2219 raise NotImplementedError(
2220 "fused DSA operators require Python document segments and are not "
2221 "supported with graph packing inputs; use the reference implementation"
2222 )
2223 segment_starts, minimum_key_indices, query_segments, causal_upper = self._resolve_packed_segments(
2224 packed_sequence,
2225 batch_size,
2226 sequence_length,
2227 query_offset,
2228 global_sequence_length,
2229 hidden_states.device,
2230 )
2232 block_plan, key_positions = self._prepare_compression(
2233 packed_sequence, cp_context, position_embeddings, global_sequence_length, hidden_states.device,
2234 )
2235 raw_document_starts = () if packed_sequence is None else packed_sequence.raw_document_starts()
2236 raw_route, main_route, document_lengths = self._plan_bank_dispatch(
2237 packed_sequence, cp_context, sequence_length, global_sequence_length, hidden_states.device,
2238 )
2239 if packed_sequence is not None:
2240 global_key_length = packed_sequence.compressed_key_length(self.compress_ratio)
2241 else:
2242 global_key_length = global_sequence_length // max(1, self.compress_ratio)
2243 raw_key_offset = 0 if raw_route is None else raw_route.global_start
2244 rope_type = "compress" if self.compress_ratio else "main"
2245 cos, sin = position_embeddings[rope_type]
2246 key_value, raw_kv_handle = self._project_raw_kv(hidden_states, cos, sin, cp_context, raw_route)
2247 overlap = (self.is_index_source and cp_context is not None
2248 and (cp_context.indexer_priority or cp_context.launch_balance is not None))
2249 prepare_attention = None
2250 if overlap:
2251 query_residual = self.q_a_norm(self.q_a_proj(hidden_states))
2252 query = None
2254 def _prepare_attention() -> None:
2255 """Keep differentiable main Q work outside selection checkpoint regions."""
2256 nonlocal query
2257 if isinstance(compressed_handle, _DeferredBankLaunch):
2258 compressed_handle.start()
2259 query = self._expand_query(query_residual, cos, sin)
2261 prepare_attention = _prepare_attention
2262 else:
2263 query_residual, query = self._project_query(hidden_states, cos, sin)
2264 latent, compressed_handle = self._compute_compressed_kv(
2265 hidden_states, position_embeddings, shared_state, cp_context, main_route, block_plan, key_positions,
2266 defer_communication=overlap and cp_context.indexer_priority,
2267 )
2268 indexer_output = self._run_indexer(
2269 hidden_states, query_residual, latent, position_embeddings, shared_state, cp_context, tp_context,
2270 query_offset, minimum_key_indices, query_segments, document_lengths, block_plan, key_positions,
2271 prepare_attention=prepare_attention, causal_upper=causal_upper,
2272 )
2274 key_value = self._finish_bank_communication(
2275 key_value, raw_kv_handle, compressed_handle, shared_state, main_route,
2276 )
2278 compressed_kv = None
2279 main_key_offset = 0
2280 topk_indices = None
2281 if self.compress_ratio:
2282 compressed_kv = shared_state.require_compressed_kv(self.kv_source_layer_idx, self.layer_idx)
2283 topk_indices = shared_state.require_topk_indices(self.index_source_layer_idx, self.layer_idx)
2284 main_key_offset = shared_state.compressed_kv_offsets.get(self.kv_source_layer_idx, 0)
2285 query = self._apply_indexer_loss(
2286 indexer_output, query, compressed_kv, tp_context, query_segments, main_key_offset,
2287 )
2288 attention_output = self._run_sparse_attention(
2289 query, key_value, compressed_kv, topk_indices, query_offset, global_sequence_length,
2290 segment_starts, query_segments, raw_key_offset, main_key_offset, raw_document_starts,
2291 global_key_length,
2292 )
2293 output = self._project_output(attention_output, cos, sin, batch_size, sequence_length)
2294 return output, None
2297@module_replacement
2298class SharedCompressedDSAAttention(SharedCompressedDSAAttentionBase):
2299 """High-performance V4.1 shared compressed attention replacement."""
2301 def __init__(
2302 self,
2303 module: nn.Module,
2304 module_fqn: str = "",
2305 context: Mapping[str, Any] | None = None,
2306 *,
2307 use_fused_ops: bool = False,
2308 indexer_load_balance: str = "none",
2309 main_kv_comm: str = "all_gather",
2310 indexer_balance_block_size: int | None = None,
2311 indexer_zigzag_min_length: int | None = None,
2312 ) -> None:
2313 """Transfer source state and select the implementation used for comparisons.
2315 Args:
2316 module: Source attention module exposing the shared-compressed DSA
2317 structural contract.
2318 module_fqn: Fully qualified source name supplied by replacement.
2319 context: Read-only replacement context supplied by Trainer.
2320 use_fused_ops: Use LI V2, SMLA and fused KL when True,
2321 raising on unsupported calls without automatic fallback.
2322 False selects eager PyTorch attention, Indexer and KL on every
2323 device, without loading optional operator extensions.
2324 indexer_load_balance: 'none', 'zigzag', 'greedy' or 'hybrid' query-task allocation.
2325 Selection exchanges low-rank qr, merge weights and rotary embeddings;
2326 attention and KL retain the original owner and autograd graph.
2327 indexer_balance_block_size: Maximum original query rows per greedy task,
2328 before adjacent tasks on one worker are merged. Greedy/hybrid default to 2048.
2329 Explicit values require greedy/hybrid and must be positive and no larger
2330 than local Q. Omit this option for none/zigzag.
2331 indexer_zigzag_min_length: Explicit hybrid-only minimum document length in raw tokens.
2332 Eligible head/tail fragments stay intact; leading remainders and other documents
2333 share the greedy pool and row budget. Documents shorter than 2*CP have no complete
2334 zigzag fragments and enter the greedy pool.
2335 main_kv_comm: 'all_to_all_v' dispatches document-causal main KV prefixes;
2336 'all_gather' retains the full-bank comparison.
2338 Note:
2339 All implementations retain the same parameter and CP/TP contracts.
2340 Selecting fusion preserves the existing external teacher and scalar
2341 KL. Reuse and disabled KL do not introduce extra operator calls.
2342 Fusion supports ordinary causal r1/r2 training on NPU with BF16
2343 main attention, D=512, SWA=128, Indexer D=128, TopK=512 and TP=1.
2344 Active fused KL requires 8, 16, 32 or 64 Indexer heads. Candidate
2345 pools are unsupported. Packed inputs compress each document independently;
2346 multiple local document segments use native TND batching.
2347 The ops-transformer training extension is loaded on the first
2348 supported fused call; a missing extension raises ImportError.
2349 """
2350 if not isinstance(use_fused_ops, bool):
2351 raise ValueError(f"use_fused_ops must be a bool, got {use_fused_ops!r}")
2352 if main_kv_comm not in ("all_gather", "all_to_all_v"):
2353 raise ValueError("main_kv_comm must be 'all_gather' or 'all_to_all_v'")
2354 indexer_balance_block_size = resolve_indexer_balance_block_size(
2355 indexer_load_balance, indexer_balance_block_size,
2356 )
2357 validate_indexer_hybrid_threshold(
2358 indexer_load_balance, indexer_zigzag_min_length,
2359 )
2360 del module_fqn, context
2361 super().__init__(
2362 module,
2363 replace_indexer=True,
2364 use_fused_ops=use_fused_ops,
2365 )
2366 self.indexer_load_balance = indexer_load_balance
2367 self.indexer_balance_block_size = indexer_balance_block_size
2368 self.indexer_zigzag_min_length = indexer_zigzag_min_length
2369 self.main_kv_comm = main_kv_comm
2372__all__ = [
2373 "SharedCompressedAttentionCPContext",
2374 "SharedCompressedAttentionState",
2375 "SharedCompressedPackedSequence",
2376 "SharedCompressedAttentionTPContext",
2377 "SharedCompressedDSAAttention",
2378 "SharedCompressedDSAAttentionBase",
2379 "SharedCompressedDSAIndexer",
2380 "build_sliding_window_indices",
2381 "compressed_candidate_topk",
2382 "compressed_causal_topk_and_candidates",
2383 "compressed_causal_topk",
2384 "select_candidate_block_indices",
2385 "select_candidate_blocks",
2386 "shared_compressed_indexer_kl_loss",
2387 "reference_sparse_attention",
2388]