Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / sparse_flash_mla.py: 90%
171 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"""SparseFlashMla forward/backward adapter for CSA with scalar attention sinks."""
17from __future__ import annotations
19from dataclasses import dataclass
20from functools import lru_cache
21from itertools import accumulate
22from typing import Any, Callable
24import torch # pylint: disable=forbidden-backend-import
27def pack_ranges(tensor: torch.Tensor, ranges: tuple[tuple[int, int], ...]) -> torch.Tensor:
28 """Pack token ranges without copying an already contiguous partition.
30 Args:
31 tensor: Source tensor with tokens on the first dimension.
32 ranges: Nonempty, ordered half-open token intervals.
34 Returns:
35 Contiguous packed tokens; autograd restores gradients to the source.
36 """
37 if all(left[1] == right[0] for left, right in zip(ranges, ranges[1:])):
38 start, end = ranges[0][0], ranges[-1][1]
39 if start == 0 and end == tensor.shape[0]:
40 return tensor.contiguous()
41 return tensor[start:end].contiguous()
42 return torch.cat([tensor[start:end] for start, end in ranges], dim=0).contiguous()
45@dataclass(frozen=True)
46class CSATNDPlan:
47 """Document-local lengths and canonical K offsets; no training activations."""
49 segments: tuple[tuple[int, int, int, int, int], ...]
50 q_ranges: tuple[tuple[int, int], ...]
51 k_ranges: tuple[tuple[int, int], ...]
52 q_lengths: tuple[int, ...]
53 k_lengths: tuple[int, ...]
54 residuals: tuple[int, ...]
56 def query(self, tensor: torch.Tensor) -> torch.Tensor:
57 """Pack a [1,Q,...] tensor in the plan's document order."""
58 return pack_ranges(tensor.squeeze(0), self.q_ranges)
60 def key(self, tensor: torch.Tensor, key_offset: int = 0) -> torch.Tensor:
61 """Pack global document prefixes using the physical origin of a [1,K,...] bank."""
62 ranges = self.k_ranges if not key_offset else tuple((start-key_offset, end-key_offset)
63 for start, end in self.k_ranges)
64 return pack_ranges(tensor.squeeze(0), ranges)
66 def lengths(self, device: torch.device) -> tuple:
67 """Reuse device length tensors, including across inference and recomputation."""
68 return _device_lengths(self, device)
70 def indexer_options(self, device: torch.device, ratio: int) -> dict:
71 """Share the document partition and causal mask between Indexer and KL."""
72 cu_q, cu_k, residuals, _ = self.lengths(device)
73 return {"cu_seqlens_q": cu_q, "cu_seqlens_k": cu_k,
74 "cmp_residual_k": residuals if ratio > 1 else None,
75 "layout_q": "TND", "layout_k": "TND", "mask_mode": 3, "cmp_ratio": ratio}
77 def local_indices(self, indices: torch.Tensor, padded_width: int = 0) -> torch.Tensor:
78 """Translate bank IDs; pad only when the consumer explicitly requires a width."""
79 selected = self.query(indices)
80 offsets = self.lengths(indices.device)[3]
81 selected = torch.where(selected >= 0, selected - offsets, -1).int()
82 if selected.shape[-1] < padded_width:
83 selected = torch.nn.functional.pad(selected, (0, padded_width-selected.shape[-1]), value=-1)
84 return selected.unsqueeze(1).contiguous()
86 def covers_queries(self, length: int) -> bool:
87 """Return whether packed rows already follow the complete original Q order."""
88 return (self.q_ranges[0][0] == 0 and self.q_ranges[-1][1] == length
89 and all(left[1] == right[0] for left, right in zip(self.q_ranges, self.q_ranges[1:])))
91 def restore(self, part: torch.Tensor, output: torch.Tensor) -> None:
92 """Write packed query rows into their original positions, preserving autograd."""
93 cursor = 0
94 for start, end in self.q_ranges:
95 output[0, start:end] = part[cursor:cursor+end-start]
96 cursor += end-start
99@lru_cache(maxsize=128)
100def csa_tnd_plan(segments: tuple, ratio: int, has_keys: bool = True) -> CSATNDPlan | None:
101 """Separate true SWA-only/empty-prefix queries without inventing valid keys."""
102 active = tuple(item for item in segments
103 if item[1] > item[0] and (bool(ratio and item[3] > item[2]) == has_keys))
104 if not active:
105 return None
106 return CSATNDPlan(
107 active, tuple((s, e) for s, e, _, _, _ in active), tuple((s, e) for _, _, s, e, _ in active),
108 tuple(e-s for s, e, _, _, _ in active), tuple(e-s for _, _, s, e, _ in active),
109 tuple(r for _, _, _, _, r in active),
110 )
113@lru_cache(maxsize=128)
114@torch.inference_mode(False)
115def _device_lengths(plan: CSATNDPlan, device: torch.device) -> tuple:
116 """Cache by the complete document partition, not merely by total Q/K sizes."""
117 cu_q = torch.tensor((0, *accumulate(plan.q_lengths)), device=device, dtype=torch.int32)
118 cu_k = torch.tensor((0, *accumulate(plan.k_lengths)), device=device, dtype=torch.int32)
119 residuals = torch.tensor(plan.residuals, device=device, dtype=torch.int32)
120 offsets = torch.tensor([s for s, _ in plan.k_ranges], device=device, dtype=torch.int32)
121 lengths = torch.tensor(plan.q_lengths, device=device, dtype=torch.int64)
122 offsets = torch.repeat_interleave(offsets, lengths, output_size=sum(plan.q_lengths)).unsqueeze(1)
123 return cu_q, cu_k, residuals, offsets
126@lru_cache(maxsize=1)
127def _load_attention_ops() -> tuple[Callable[..., Any], Callable[..., Any], Callable[..., Any]]:
128 """Resolve attention forward, backward and metadata through the official binding."""
129 # Lazy import leaves the eager baseline usable without the NPU extension.
130 try:
131 from cann_ops_transformer import ( # pylint: disable=import-outside-toplevel
132 sparse_flash_mla as sparse_flash_mla_op,
133 sparse_flash_mla_grad, sparse_flash_mla_metadata,
134 )
135 except ImportError as err:
136 raise RuntimeError(
137 "use_fused_ops=True requires cann_ops_transformer.sparse_flash_mla, sparse_flash_mla_grad "
138 "and sparse_flash_mla_metadata. "
139 "See hyper_parallel/components/functional/README.md for ops-transformer installation, "
140 "environment setup and tested versions."
141 ) from err
142 return sparse_flash_mla_op, sparse_flash_mla_grad, sparse_flash_mla_metadata
145def supports_fused_attention_inputs(query: torch.Tensor, ratio: int, window: int, topk: int = 512) -> bool:
146 """Check the validated attention input contract without importing operators.
148 Args:
149 query: Main Q [batch, heads, queries, dimension].
150 ratio: Model ratio, zero for a pure SWA layer.
151 window: Raw attention window, including the current token.
152 topk: Logical compressed selection width, before padding; zero means no complete compressed key.
153 """
154 if query.device.type != "npu" or query.dtype != torch.bfloat16:
155 return False
156 if query.shape[-1] != 512 or window != 128 or ratio not in (0, 1, 2):
157 return False
158 if query.shape[1] not in (1, 2, 4, 8, 16, 32, 64, 128) or query.shape[2] == 0:
159 return False
160 if ratio and not 0 <= topk <= 512:
161 return False
162 return True
165@lru_cache(maxsize=128)
166@torch.inference_mode(False)
167def _attention_metadata(
168 device: torch.device, batch: int, length: int, heads: int, dim: int,
169 raw_length: int, main_length: int, topk: int, ratio: int, residual: int, window: int,
170) -> tuple:
171 """Reuse shape metadata across evaluation and gradient-tracked training.
173 Evaluation may populate the cache under inference_mode. The backward
174 saves the cached residual, so it must be an ordinary tensor regardless
175 of which mode first populated the cache.
176 """
177 # The training backward requires a zero residual tensor even at ratio one;
178 # the forward interface instead omits that optional input at ratio one.
179 residue = torch.full((batch,), residual, dtype=torch.int32, device=device) if main_length else None
180 _, _, metadata_op = _load_attention_ops()
181 metadata = metadata_op(
182 heads, 1, dim, cmp_residual_kv=residue if ratio > 1 else None, batch_size=batch, max_seqlen_q=length,
183 max_seqlen_ori_kv=raw_length, max_seqlen_cmp_kv=main_length, ori_topk=0, cmp_topk=topk,
184 cmp_ratio=ratio, ori_mask_mode=4, cmp_mask_mode=3, ori_win_left=window-1, ori_win_right=0,
185 layout_q="BSND", layout_kv="BSND", has_ori_kv=True, has_cmp_kv=bool(main_length),
186 )
187 return residue, metadata
190class _SparseFlashMla(torch.autograd.Function):
191 """Wrap SparseFlashMla and SparseFlashMlaGrad without substituting the KL teacher."""
193 @staticmethod
194 def forward(
195 ctx: Any, query: torch.Tensor, raw: torch.Tensor, main: torch.Tensor | None,
196 indices: torch.Tensor | None, sinks: torch.Tensor, scale: float,
197 ratio: int, residual: int, window: int,
198 ) -> torch.Tensor:
199 """Save the native output and LSE required by the training backward."""
200 batch, length, heads, dim = query.shape
201 residue, metadata = _attention_metadata(
202 query.device, batch, length, heads, dim, raw.shape[1], 0 if main is None else main.shape[1],
203 0 if indices is None else indices.shape[-1], ratio, residual, window,
204 )
205 options = {"softmax_scale": scale, "cmp_ratio": ratio, "ori_mask_mode": 4, "cmp_mask_mode": 3,
206 "ori_win_left": window-1, "ori_win_right": 0, "layout_q": "BSND", "layout_kv": "BSND"}
207 attention, _, _ = _load_attention_ops()
208 output, lse = attention(
209 query, ori_kv=raw, cmp_kv=main, cmp_sparse_indices=indices, sinks=sinks,
210 cmp_residual_kv=residue if ratio > 1 else None, metadata=metadata, return_softmax_lse=True, **options,
211 )
212 ctx.save_for_backward(query, raw, main, indices, sinks, residue, output, lse)
213 ctx.options = options
214 return output
216 @staticmethod
217 def backward(ctx: Any, grad_output: torch.Tensor) -> tuple:
218 """Return main-branch gradients without substituting the KL teacher."""
219 query, raw, main, indices, sinks, residue, output, lse = ctx.saved_tensors
220 _, attention_grad, _ = _load_attention_ops()
221 grad_query, grad_raw, grad_main, grad_sinks, _, _ = attention_grad(
222 query, grad_output.contiguous(), output, lse, ori_kv=raw, cmp_kv=main,
223 cmp_sparse_indices=indices, sinks=sinks, cmp_residual_kv=residue, **ctx.options,
224 )
225 return grad_query, grad_raw, grad_main if main is not None else None, None, grad_sinks, None, None, None, None
228def _bsnd_sparse_flash_mla(
229 query: torch.Tensor, raw_key: torch.Tensor, main_key: torch.Tensor | None,
230 indices: torch.Tensor | None, sinks: torch.Tensor, scale: float, ratio: int,
231 window: int, query_offset: int,
232 segment: tuple[int, int, int, int, int], raw_key_offset: int = 0, main_key_offset: int = 0,
233 raw_document_starts: tuple = (),
234) -> torch.Tensor:
235 """Run one document segment per batch item with independently routed KV banks.
237 The operator owns Q tiling within this BSND call. Multiple packed
238 documents use the TND entry with cumulative document lengths.
240 Args:
241 query: Local attention Q in [B, H, Q, D].
242 raw_key: Full or halo raw KV in [B, 1, K, D].
243 main_key: Full or prefix compressed KV in [B, Kc, D], or None for SWA only.
244 indices: Global compressed selection [B, Q, topk], invalid slots -1.
245 sinks: One scalar logit per head.
246 scale: Attention dot-product scale.
247 ratio: Compression ratio, zero denotes the SWA-only layer.
248 window: Number of raw tokens visible to each query including itself.
249 query_offset: Global position of the first local query.
250 segment: Local Q start/end, compressed document start/prefix end and residual,
251 shared by all physical batch items.
253 Returns:
254 Attention output in [B, Q, H, D].
255 """
256 if query.dtype != torch.bfloat16 or query.shape[-1] != 512 or window != 128:
257 raise ValueError("V4 attention requires BF16, head_dim=512, and window=128")
258 start, end, key_start, key_end, residual = segment
259 document_start = dict(raw_document_starts).get(start, key_start * max(ratio, 1))
260 raw_left = max(document_start, query_offset + start - window + 1)
261 raw_right = query_offset + end
262 # contiguous() may keep stale singleton strides after slicing a BSND KV bank.
263 batch, _, _, dim = query.shape
264 raw = raw_key[:, :, raw_left-raw_key_offset:raw_right-raw_key_offset]
265 raw = raw.squeeze(1).contiguous().view(batch, raw_right-raw_left, 1, dim)
266 main = None
267 selected = None
268 if ratio and key_end > key_start:
269 main = main_key[:, key_start-main_key_offset:key_end-main_key_offset]
270 main = main.contiguous().view(batch, key_end-key_start, 1, dim)
271 original = indices[:, start:end]
272 selected = torch.where(original >= 0, original-key_start, -1).int()
273 # Backward infers the active count from this width and causality; only KL requires 512 slots.
274 selected = selected.unsqueeze(2).contiguous()
275 return _SparseFlashMla.apply(
276 query[:, :, start:end].transpose(1, 2).contiguous(), raw, main, selected, sinks.float(),
277 scale, max(ratio, 1) if main is not None else 1, residual if main is not None else 0, window,
278 )
281@lru_cache(maxsize=128)
282@torch.inference_mode(False)
283def _tnd_attention_metadata(plan, device, heads, dim, raw_ranges, ratio, window, has_keys, topk):
284 """Describe independent raw and compressed coordinate systems to SMLA."""
285 cu_q, cu_k, residuals, _ = plan.lengths(device)
286 raw_lengths = tuple(end-start for start, end in raw_ranges)
287 cu_raw = torch.tensor((0, *accumulate(raw_lengths)), device=device, dtype=torch.int32)
288 _, _, metadata_op = _load_attention_ops()
289 options = {"layout_q": "TND", "layout_kv": "TND", "cmp_ratio": ratio, "ori_mask_mode": 4,
290 "cmp_mask_mode": 3, "ori_win_left": window-1, "ori_win_right": 0}
291 metadata = metadata_op(
292 heads, 1, dim, cu_seqlens_q=cu_q, cu_seqlens_ori_kv=cu_raw,
293 cu_seqlens_cmp_kv=cu_k if has_keys else None,
294 cmp_residual_kv=residuals if has_keys and ratio > 1 else None,
295 batch_size=len(plan.segments), max_seqlen_q=max(plan.q_lengths),
296 max_seqlen_ori_kv=max(raw_lengths), max_seqlen_cmp_kv=max(plan.k_lengths) if has_keys else 0,
297 ori_topk=0, cmp_topk=topk, has_ori_kv=True, has_cmp_kv=has_keys, **options,
298 )
299 return cu_q, cu_raw, cu_k if has_keys else None, residuals if has_keys else None, metadata
302class _SparseFlashMlaTND(torch.autograd.Function):
303 """Native multi-document forward/backward; save all sequence tensors through autograd."""
305 @staticmethod
306 def forward(
307 ctx: Any, query: torch.Tensor, raw: torch.Tensor, main: torch.Tensor | None,
308 indices: torch.Tensor | None, sinks: torch.Tensor, scale: float, ratio: int, window: int,
309 plan: CSATNDPlan, raw_ranges: tuple,
310 ) -> torch.Tensor:
311 """Save native output/LSE and independent raw/compressed cumulative lengths."""
312 cu_q, cu_raw, cu_k, residuals, metadata = _tnd_attention_metadata(
313 plan, query.device, query.shape[1], query.shape[2], raw_ranges, ratio, window, main is not None,
314 0 if indices is None else indices.shape[-1],
315 )
316 options = {"layout_q": "TND", "layout_kv": "TND", "softmax_scale": scale, "cmp_ratio": ratio,
317 "ori_mask_mode": 4, "cmp_mask_mode": 3, "ori_win_left": window-1, "ori_win_right": 0}
318 forward, _, _ = _load_attention_ops()
319 output, lse = forward(
320 query, ori_kv=raw, cmp_kv=main, cmp_sparse_indices=indices, sinks=sinks,
321 cu_seqlens_q=cu_q, cu_seqlens_ori_kv=cu_raw, cu_seqlens_cmp_kv=cu_k,
322 cmp_residual_kv=residuals if ratio > 1 else None, metadata=metadata,
323 return_softmax_lse=True, **options,
324 )
325 ctx.save_for_backward(query, raw, main, indices, sinks, output, lse, cu_q, cu_raw, cu_k, residuals)
326 ctx.options = options
327 return output
329 @staticmethod
330 def backward(ctx: Any, grad_output: torch.Tensor) -> tuple:
331 """Return each physical bank's gradient to its autograd communication owner."""
332 query, raw, main, indices, sinks, output, lse, cu_q, cu_raw, cu_k, residuals = ctx.saved_tensors
333 _, backward, _ = _load_attention_ops()
334 dq, dr, dm, ds, _, _ = backward(
335 query, grad_output.contiguous(), output, lse, ori_kv=raw, cmp_kv=main,
336 cmp_sparse_indices=indices, sinks=sinks, cu_seqlens_q=cu_q, cu_seqlens_ori_kv=cu_raw,
337 cu_seqlens_cmp_kv=cu_k, cmp_residual_kv=residuals, **ctx.options,
338 )
339 return dq, dr, dm if main is not None else None, None, ds, None, None, None, None, None
342def sparse_flash_mla(
343 query: torch.Tensor, raw_key: torch.Tensor, main_key: torch.Tensor | None,
344 indices: torch.Tensor | None, sinks: torch.Tensor, scale: float, ratio: int,
345 window: int, query_offset: int, segments: tuple, raw_key_offset: int = 0, main_key_offset: int = 0,
346 raw_document_starts: tuple = (),
347) -> torch.Tensor:
348 """Batch packed documents in native TND with independently routed KV banks.
350 Args:
351 query: Local attention Q [batch, heads, queries, 512].
352 raw_key: Full or halo raw KV [batch, 1, keys, 512].
353 main_key: Full or document-prefix compressed KV [batch, compressed keys, 512], or None.
354 indices: Sorted global compressed key IDs, invalid slots -1.
355 sinks: Per-head scalar sink logits.
356 scale: Attention dot-product scale.
357 ratio: Compression ratio; zero selects pure SWA.
358 window: Raw sliding-window size, including the current token.
359 query_offset: Global position of the first local query.
360 segments: Local Q start/end, compressed document start/prefix end and residual.
361 raw_document_starts: Explicit (local Q start, raw document start) pairs for natural lengths.
362 raw_key_offset: Global position of the first row in the compact raw KV bank.
363 main_key_offset: Global compressed ID represented by row zero of main_key.
365 Returns:
366 Attention output [batch, queries, heads, 512]. A single document or
367 ordinary physical batch uses BSND. Empty compressed prefixes form a
368 separate SWA-only TND group without inventing compressed keys.
369 """
370 if len(segments) == 1:
371 return _bsnd_sparse_flash_mla(
372 query, raw_key, main_key, indices, sinks, scale, ratio, window, query_offset, segments[0],
373 raw_key_offset, main_key_offset, raw_document_starts,
374 )
375 if query.shape[0] != 1:
376 raise ValueError("Packed Sparse Flash MLA requires batch size one")
377 if query.dtype != torch.bfloat16 or query.shape[-1] != 512 or window != 128:
378 raise ValueError("V4 attention requires BF16, head_dim=512, and window=128")
379 document_starts = dict(raw_document_starts)
380 output = None
381 for has_keys in (True, False):
382 plan = csa_tnd_plan(segments, ratio, has_keys)
383 if plan is None:
384 continue
385 raw_ranges = tuple(
386 (max(document_starts.get(s, ks * max(ratio, 1)), query_offset+s-window+1)-raw_key_offset,
387 query_offset+e-raw_key_offset)
388 for s, e, ks, _, _ in plan.segments
389 )
390 raw = pack_ranges(raw_key.squeeze(0).squeeze(0), raw_ranges).unsqueeze(1)
391 main = plan.key(main_key, main_key_offset).unsqueeze(1) if has_keys else None
392 selected = plan.local_indices(indices) if has_keys else None
393 part = _SparseFlashMlaTND.apply(
394 plan.query(query.transpose(1, 2)), raw, main, selected, sinks.float(),
395 scale, max(ratio, 1) if has_keys else 1, window, plan, raw_ranges,
396 )
397 # A complete Q partition is already ordered; a canvas would retain another full output.
398 if plan.covers_queries(query.shape[2]):
399 return part.unsqueeze(0)
400 if output is None:
401 output = torch.zeros_like(query.transpose(1, 2))
402 plan.restore(part, output)
403 return torch.zeros_like(query.transpose(1, 2)) if output is None else output