Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / sparse_flash_mla.py: 87%
166 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""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) -> torch.Tensor:
61 """Pack complete visible document prefixes from a [1,K,...] bank."""
62 return pack_ranges(tensor.squeeze(0), self.k_ranges)
64 def lengths(self, device: torch.device) -> tuple:
65 """Reuse device length tensors, including across inference and recomputation."""
66 return _device_lengths(self, device)
68 def indexer_options(self, device: torch.device, ratio: int) -> dict:
69 """Share the document partition and causal mask between Indexer and KL."""
70 cu_q, cu_k, residuals, _ = self.lengths(device)
71 return {"cu_seqlens_q": cu_q, "cu_seqlens_k": cu_k,
72 "cmp_residual_k": residuals if ratio > 1 else None,
73 "layout_q": "TND", "layout_k": "TND", "mask_mode": 3, "cmp_ratio": ratio}
75 def local_indices(self, indices: torch.Tensor, padded_width: int = 0) -> torch.Tensor:
76 """Translate bank IDs; pad only when the consumer explicitly requires a width."""
77 selected = self.query(indices)
78 offsets = self.lengths(indices.device)[3]
79 selected = torch.where(selected >= 0, selected - offsets, -1).int()
80 if selected.shape[-1] < padded_width:
81 selected = torch.nn.functional.pad(selected, (0, padded_width-selected.shape[-1]), value=-1)
82 return selected.unsqueeze(1).contiguous()
84 def covers_queries(self, length: int) -> bool:
85 """Return whether packed rows already follow the complete original Q order."""
86 return (self.q_ranges[0][0] == 0 and self.q_ranges[-1][1] == length
87 and all(left[1] == right[0] for left, right in zip(self.q_ranges, self.q_ranges[1:])))
89 def restore(self, part: torch.Tensor, output: torch.Tensor) -> None:
90 """Write packed query rows into their original positions, preserving autograd."""
91 cursor = 0
92 for start, end in self.q_ranges:
93 output[0, start:end] = part[cursor:cursor+end-start]
94 cursor += end-start
97@lru_cache(maxsize=128)
98def csa_tnd_plan(segments: tuple, ratio: int, has_keys: bool = True) -> CSATNDPlan | None:
99 """Separate true SWA-only/empty-prefix queries without inventing valid keys."""
100 active = tuple(item for item in segments
101 if item[1] > item[0] and (bool(ratio and item[3] > item[2]) == has_keys))
102 if not active:
103 return None
104 return CSATNDPlan(
105 active, tuple((s, e) for s, e, _, _, _ in active), tuple((s, e) for _, _, s, e, _ in active),
106 tuple(e-s for s, e, _, _, _ in active), tuple(e-s for _, _, s, e, _ in active),
107 tuple(r for _, _, _, _, r in active),
108 )
111@lru_cache(maxsize=128)
112@torch.inference_mode(False)
113def _device_lengths(plan: CSATNDPlan, device: torch.device) -> tuple:
114 """Cache by the complete document partition, not merely by total Q/K sizes."""
115 cu_q = torch.tensor((0, *accumulate(plan.q_lengths)), device=device, dtype=torch.int32)
116 cu_k = torch.tensor((0, *accumulate(plan.k_lengths)), device=device, dtype=torch.int32)
117 residuals = torch.tensor(plan.residuals, device=device, dtype=torch.int32)
118 offsets = torch.tensor([s for s, _ in plan.k_ranges], device=device, dtype=torch.int32)
119 lengths = torch.tensor(plan.q_lengths, device=device, dtype=torch.int64)
120 offsets = torch.repeat_interleave(offsets, lengths, output_size=sum(plan.q_lengths)).unsqueeze(1)
121 return cu_q, cu_k, residuals, offsets
124@lru_cache(maxsize=1)
125def _load_attention_ops() -> tuple[Callable[..., Any], Callable[..., Any], Callable[..., Any]]:
126 """Resolve attention forward, backward and metadata through the official binding."""
127 # Lazy import leaves the eager baseline usable without the NPU extension.
128 try:
129 from cann_ops_transformer import ( # pylint: disable=import-outside-toplevel
130 sparse_flash_mla as sparse_flash_mla_op,
131 sparse_flash_mla_grad, sparse_flash_mla_metadata,
132 )
133 except ImportError as err:
134 raise RuntimeError(
135 "use_fused_ops=True requires cann_ops_transformer.sparse_flash_mla, sparse_flash_mla_grad "
136 "and sparse_flash_mla_metadata. "
137 "See hyper_parallel/components/functional/README.md for ops-transformer installation, "
138 "environment setup and tested versions."
139 ) from err
140 return sparse_flash_mla_op, sparse_flash_mla_grad, sparse_flash_mla_metadata
143def supports_fused_attention_inputs(query: torch.Tensor, ratio: int, window: int, topk: int = 512) -> bool:
144 """Check the validated attention input contract without importing operators.
146 Args:
147 query: Main Q [batch, heads, queries, dimension].
148 ratio: Model ratio, zero for a pure SWA layer.
149 window: Raw attention window, including the current token.
150 topk: Logical compressed selection width, before padding.
151 """
152 if query.device.type != "npu" or query.dtype != torch.bfloat16:
153 return False
154 if query.shape[-1] != 512 or window != 128 or ratio not in (0, 1, 2):
155 return False
156 if query.shape[1] not in (1, 2, 4, 8, 16, 32, 64, 128) or query.shape[2] == 0:
157 return False
158 if ratio and not 1 <= topk <= 512:
159 return False
160 return True
163@lru_cache(maxsize=128)
164@torch.inference_mode(False)
165def _attention_metadata(
166 device: torch.device, batch: int, length: int, heads: int, dim: int,
167 raw_length: int, main_length: int, topk: int, ratio: int, residual: int, window: int,
168) -> tuple:
169 """Reuse shape metadata across evaluation and gradient-tracked training.
171 Evaluation may populate the cache under inference_mode. The backward
172 saves the cached residual, so it must be an ordinary tensor regardless
173 of which mode first populated the cache.
174 """
175 # The training backward requires a zero residual tensor even at ratio one;
176 # the forward interface instead omits that optional input at ratio one.
177 residue = torch.full((batch,), residual, dtype=torch.int32, device=device) if main_length else None
178 _, _, metadata_op = _load_attention_ops()
179 metadata = metadata_op(
180 heads, 1, dim, cmp_residual_kv=residue if ratio > 1 else None, batch_size=batch, max_seqlen_q=length,
181 max_seqlen_ori_kv=raw_length, max_seqlen_cmp_kv=main_length, ori_topk=0, cmp_topk=topk,
182 cmp_ratio=ratio, ori_mask_mode=4, cmp_mask_mode=3, ori_win_left=window-1, ori_win_right=0,
183 layout_q="BSND", layout_kv="BSND", has_ori_kv=True, has_cmp_kv=bool(main_length),
184 )
185 return residue, metadata
188class _SparseFlashMla(torch.autograd.Function):
189 """Wrap SparseFlashMla and SparseFlashMlaGrad without substituting the KL teacher."""
191 @staticmethod
192 def forward(
193 ctx: Any, query: torch.Tensor, raw: torch.Tensor, main: torch.Tensor | None,
194 indices: torch.Tensor | None, sinks: torch.Tensor, scale: float,
195 ratio: int, residual: int, window: int,
196 ) -> torch.Tensor:
197 """Save the native output and LSE required by the training backward."""
198 batch, length, heads, dim = query.shape
199 residue, metadata = _attention_metadata(
200 query.device, batch, length, heads, dim, raw.shape[1], 0 if main is None else main.shape[1],
201 0 if indices is None else indices.shape[-1], ratio, residual, window,
202 )
203 options = {"softmax_scale": scale, "cmp_ratio": ratio, "ori_mask_mode": 4, "cmp_mask_mode": 3,
204 "ori_win_left": window-1, "ori_win_right": 0, "layout_q": "BSND", "layout_kv": "BSND"}
205 attention, _, _ = _load_attention_ops()
206 output, lse = attention(
207 query, ori_kv=raw, cmp_kv=main, cmp_sparse_indices=indices, sinks=sinks,
208 cmp_residual_kv=residue if ratio > 1 else None, metadata=metadata, return_softmax_lse=True, **options,
209 )
210 ctx.save_for_backward(query, raw, main, indices, sinks, residue, output, lse)
211 ctx.options = options
212 return output
214 @staticmethod
215 def backward(ctx: Any, grad_output: torch.Tensor) -> tuple:
216 """Return main-branch gradients without substituting the KL teacher."""
217 query, raw, main, indices, sinks, residue, output, lse = ctx.saved_tensors
218 _, attention_grad, _ = _load_attention_ops()
219 grad_query, grad_raw, grad_main, grad_sinks, _, _ = attention_grad(
220 query, grad_output.contiguous(), output, lse, ori_kv=raw, cmp_kv=main,
221 cmp_sparse_indices=indices, sinks=sinks, cmp_residual_kv=residue, **ctx.options,
222 )
223 return grad_query, grad_raw, grad_main if main is not None else None, None, grad_sinks, None, None, None, None
226def _bsnd_sparse_flash_mla(
227 query: torch.Tensor, raw_key: torch.Tensor, main_key: torch.Tensor | None,
228 indices: torch.Tensor | None, sinks: torch.Tensor, scale: float, ratio: int,
229 window: int, query_offset: int,
230 segment: tuple[int, int, int, int, int],
231) -> torch.Tensor:
232 """Run one document segment per batch item using the original all-gather banks.
234 The operator owns Q tiling within this BSND call. Multiple packed
235 documents use the TND entry with cumulative document lengths.
237 Args:
238 query: Local attention Q in [B, H, Q, D].
239 raw_key: Global raw KV in [B, 1, K, D].
240 main_key: Global compressed KV in [B, Kc, D], or None for SWA only.
241 indices: Global compressed selection [B, Q, topk], invalid slots -1.
242 sinks: One scalar logit per head.
243 scale: Attention dot-product scale.
244 ratio: Compression ratio, zero denotes the SWA-only layer.
245 window: Number of raw tokens visible to each query including itself.
246 query_offset: Global position of the first local query.
247 segment: Local Q start/end, compressed document start/prefix end and residual,
248 shared by all physical batch items.
250 Returns:
251 Attention output in [B, Q, H, D].
252 """
253 if query.dtype != torch.bfloat16 or query.shape[-1] != 512 or window != 128:
254 raise ValueError("V4 attention requires BF16, head_dim=512, and window=128")
255 start, end, key_start, key_end, residual = segment
256 document_start = key_start * max(ratio, 1)
257 raw_left = max(document_start, query_offset + start - window + 1)
258 raw_right = query_offset + end
259 raw = raw_key[:, :, raw_left:raw_right].transpose(1, 2).contiguous()
260 main = None
261 selected = None
262 if ratio and key_end > key_start:
263 main = main_key[:, key_start:key_end].unsqueeze(2).contiguous()
264 original = indices[:, start:end]
265 selected = torch.where(original >= 0, original-key_start, -1).int()
266 # Backward infers the active count from this width and causality; only KL requires 512 slots.
267 selected = selected.unsqueeze(2).contiguous()
268 return _SparseFlashMla.apply(
269 query[:, :, start:end].transpose(1, 2).contiguous(), raw, main, selected, sinks.float(),
270 scale, max(ratio, 1) if main is not None else 1, residual if main is not None else 0, window,
271 )
274@lru_cache(maxsize=128)
275@torch.inference_mode(False)
276def _tnd_attention_metadata(plan, device, heads, dim, raw_ranges, ratio, window, has_keys, topk):
277 """Describe independent raw and compressed coordinate systems to SMLA."""
278 cu_q, cu_k, residuals, _ = plan.lengths(device)
279 raw_lengths = tuple(end-start for start, end in raw_ranges)
280 cu_raw = torch.tensor((0, *accumulate(raw_lengths)), device=device, dtype=torch.int32)
281 _, _, metadata_op = _load_attention_ops()
282 options = {"layout_q": "TND", "layout_kv": "TND", "cmp_ratio": ratio, "ori_mask_mode": 4,
283 "cmp_mask_mode": 3, "ori_win_left": window-1, "ori_win_right": 0}
284 metadata = metadata_op(
285 heads, 1, dim, cu_seqlens_q=cu_q, cu_seqlens_ori_kv=cu_raw,
286 cu_seqlens_cmp_kv=cu_k if has_keys else None,
287 cmp_residual_kv=residuals if has_keys and ratio > 1 else None,
288 batch_size=len(plan.segments), max_seqlen_q=max(plan.q_lengths),
289 max_seqlen_ori_kv=max(raw_lengths), max_seqlen_cmp_kv=max(plan.k_lengths) if has_keys else 0,
290 ori_topk=0, cmp_topk=topk, has_ori_kv=True, has_cmp_kv=has_keys, **options,
291 )
292 return cu_q, cu_raw, cu_k if has_keys else None, residuals if has_keys else None, metadata
295class _SparseFlashMlaTND(torch.autograd.Function):
296 """Native multi-document forward/backward; save all sequence tensors through autograd."""
298 @staticmethod
299 def forward(
300 ctx: Any, query: torch.Tensor, raw: torch.Tensor, main: torch.Tensor | None,
301 indices: torch.Tensor | None, sinks: torch.Tensor, scale: float, ratio: int, window: int,
302 plan: CSATNDPlan, raw_ranges: tuple,
303 ) -> torch.Tensor:
304 """Save native output/LSE and independent raw/compressed cumulative lengths."""
305 cu_q, cu_raw, cu_k, residuals, metadata = _tnd_attention_metadata(
306 plan, query.device, query.shape[1], query.shape[2], raw_ranges, ratio, window, main is not None,
307 0 if indices is None else indices.shape[-1],
308 )
309 options = {"layout_q": "TND", "layout_kv": "TND", "softmax_scale": scale, "cmp_ratio": ratio,
310 "ori_mask_mode": 4, "cmp_mask_mode": 3, "ori_win_left": window-1, "ori_win_right": 0}
311 forward, _, _ = _load_attention_ops()
312 output, lse = forward(
313 query, ori_kv=raw, cmp_kv=main, cmp_sparse_indices=indices, sinks=sinks,
314 cu_seqlens_q=cu_q, cu_seqlens_ori_kv=cu_raw, cu_seqlens_cmp_kv=cu_k,
315 cmp_residual_kv=residuals if ratio > 1 else None, metadata=metadata,
316 return_softmax_lse=True, **options,
317 )
318 ctx.save_for_backward(query, raw, main, indices, sinks, output, lse, cu_q, cu_raw, cu_k, residuals)
319 ctx.options = options
320 return output
322 @staticmethod
323 def backward(ctx: Any, grad_output: torch.Tensor) -> tuple:
324 """Return each physical bank's gradient to its autograd communication owner."""
325 query, raw, main, indices, sinks, output, lse, cu_q, cu_raw, cu_k, residuals = ctx.saved_tensors
326 _, backward, _ = _load_attention_ops()
327 dq, dr, dm, ds, _, _ = backward(
328 query, grad_output.contiguous(), output, lse, ori_kv=raw, cmp_kv=main,
329 cmp_sparse_indices=indices, sinks=sinks, cu_seqlens_q=cu_q, cu_seqlens_ori_kv=cu_raw,
330 cu_seqlens_cmp_kv=cu_k, cmp_residual_kv=residuals, **ctx.options,
331 )
332 return dq, dr, dm if main is not None else None, None, ds, None, None, None, None, None
335def sparse_flash_mla(
336 query: torch.Tensor, raw_key: torch.Tensor, main_key: torch.Tensor | None,
337 indices: torch.Tensor | None, sinks: torch.Tensor, scale: float, ratio: int,
338 window: int, query_offset: int, segments: tuple,
339) -> torch.Tensor:
340 """Batch packed documents in native TND while preserving the all-gather banks.
342 Args:
343 query: Local attention Q [batch, heads, queries, 512].
344 raw_key: Global raw KV [batch, 1, keys, 512].
345 main_key: Global compressed KV [batch, compressed keys, 512], or None.
346 indices: Sorted global compressed key IDs, invalid slots -1.
347 sinks: Per-head scalar sink logits.
348 scale: Attention dot-product scale.
349 ratio: Compression ratio; zero selects pure SWA.
350 window: Raw sliding-window size, including the current token.
351 query_offset: Global position of the first local query.
352 segments: Local Q start/end, compressed document start/prefix end and residual.
354 Returns:
355 Attention output [batch, queries, heads, 512]. A single document or
356 ordinary physical batch uses BSND. Empty compressed prefixes form a
357 separate SWA-only TND group without inventing compressed keys.
358 """
359 if len(segments) == 1:
360 return _bsnd_sparse_flash_mla(
361 query, raw_key, main_key, indices, sinks, scale, ratio, window, query_offset, segments[0],
362 )
363 if query.shape[0] != 1:
364 raise ValueError("Packed Sparse Flash MLA requires batch size one")
365 if query.dtype != torch.bfloat16 or query.shape[-1] != 512 or window != 128:
366 raise ValueError("V4 attention requires BF16, head_dim=512, and window=128")
367 output = None
368 for has_keys in (True, False):
369 plan = csa_tnd_plan(segments, ratio, has_keys)
370 if plan is None:
371 continue
372 raw_ranges = tuple(
373 (max(ks * max(ratio, 1), query_offset+s-window+1), query_offset+e)
374 for s, e, ks, _, _ in plan.segments
375 )
376 raw = pack_ranges(raw_key.squeeze(0).squeeze(0), raw_ranges).unsqueeze(1)
377 main = plan.key(main_key).unsqueeze(1) if has_keys else None
378 selected = plan.local_indices(indices) if has_keys else None
379 part = _SparseFlashMlaTND.apply(
380 plan.query(query.transpose(1, 2)), raw, main, selected, sinks.float(),
381 scale, max(ratio, 1) if has_keys else 1, window, plan, raw_ranges,
382 )
383 # A complete Q partition is already ordered; a canvas would retain another full output.
384 if plan.covers_queries(query.shape[2]):
385 return part.unsqueeze(0)
386 if output is None:
387 output = torch.zeros_like(query.transpose(1, 2))
388 plan.restore(part, output)
389 return torch.zeros_like(query.transpose(1, 2)) if output is None else output