Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / custom_ops / custom_op_impl.py: 45%
204 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05: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"""MindSpore custom kernel implementations and DFunction wrappers."""
16from dataclasses import dataclass
17import importlib
18import os
19import sys
21import mindspore as ms # pylint: disable=C0415
23from hyper_parallel.core.shard.dfunction import DFunction
26_CC_DIR = os.path.dirname(os.path.abspath(__file__))
27_MS_EXTENSION_NAME = "hyper_parallel_custom_ops_ms"
28_BUILD_LIB = os.path.join(_CC_DIR, "build", "lib")
30if _BUILD_LIB not in sys.path:
31 sys.path.insert(0, _BUILD_LIB)
33_CUSTOM_OP_SOURCES = [
34 os.path.join(_CC_DIR, "module.cc"),
35 os.path.join(_CC_DIR, "dense_lightning_indexer_grad_kl_loss.cc"),
36 os.path.join(_CC_DIR, "dense_lightning_indexer_softmax_lse.cc"),
37 os.path.join(_CC_DIR, "sparse_lightning_indexer_grad_kl_loss.cc"),
38 os.path.join(_CC_DIR, "mhc_post.cc"),
39 os.path.join(_CC_DIR, "mhc_post_backward.cc"),
40 os.path.join(_CC_DIR, "mhc_pre_sinkhorn.cc"),
41 os.path.join(_CC_DIR, "mhc_pre_sinkhorn_backward.cc"),
42 os.path.join(_CC_DIR, "mhc_pre_clamp_sinkhorn.cc"),
43 os.path.join(_CC_DIR, "mhc_pre_clamp_sinkhorn_backward.cc"),
44 os.path.join(_CC_DIR, "lightning_indexer_v2.cc"),
45 os.path.join(_CC_DIR, "sparse_flash_mla.cc"),
46 os.path.join(_CC_DIR, "sparse_flash_mla_grad.cc"),
47 os.path.join(_CC_DIR, "sparse_lightning_indexer_kl_loss_grad.cc"),
48]
49_MHC_PRE_CLAMP_NONE_GRADS = (None,) * 7
52@dataclass(frozen=True)
53class _MhcPreClampArgs:
54 """Bound arguments for npu_mhc_pre_clamp_sinkhorn."""
56 x: ms.Tensor
57 phi: ms.Tensor
58 alpha: ms.Tensor
59 bias: ms.Tensor
60 hc_mult: int
61 num_iters: int
62 hc_eps: float
63 norm_eps: float
64 out_flag: bool
65 clamp_min: float
66 clamp_max: float
69def _bind_mhc_pre_clamp_args(args, kwargs):
70 """Bind npu_mhc_pre_clamp_sinkhorn arguments with Python defaults."""
71 names = (
72 "x", "phi", "alpha", "bias", "hc_mult", "num_iters",
73 "hc_eps", "norm_eps", "out_flag", "clamp_min", "clamp_max",
74 )
75 values = {
76 "hc_mult": 4,
77 "num_iters": 20,
78 "hc_eps": 1e-6,
79 "norm_eps": 1e-6,
80 "out_flag": True,
81 "clamp_min": 0.0,
82 "clamp_max": 0.0,
83 }
84 if len(args) > len(names):
85 raise TypeError(f"npu_mhc_pre_clamp_sinkhorn expected at most {len(names)} arguments")
86 for name, value in zip(names, args):
87 values[name] = value
88 for name, value in kwargs.items():
89 if name in values and name in names[:len(args)]:
90 raise TypeError(f"npu_mhc_pre_clamp_sinkhorn got multiple values for argument '{name}'")
91 if name not in names:
92 raise TypeError(f"npu_mhc_pre_clamp_sinkhorn got an unexpected keyword argument '{name}'")
93 values[name] = value
94 missing = [name for name in names[:4] if name not in values]
95 if missing:
96 raise TypeError(f"npu_mhc_pre_clamp_sinkhorn missing required arguments: {missing}")
97 return _MhcPreClampArgs(*(values[name] for name in names))
100def _build_custom_ops():
101 return ms.ops.CustomOpBuilder(
102 _MS_EXTENSION_NAME,
103 _CUSTOM_OP_SOURCES,
104 backend="Ascend",
105 ).load()
108try:
109 _custom_ops = importlib.import_module(_MS_EXTENSION_NAME)
110except ImportError:
111 # Source-tree development: .so not pre-built; JIT-compile from local .cc files.
112 _custom_ops = _build_custom_ops()
113else:
114 # Rebuild stale source-tree extensions that predate newly added symbols.
115 if not hasattr(_custom_ops, "npu_mhc_pre_clamp_sinkhorn"):
116 _custom_ops = _build_custom_ops()
119def _ensure_contiguous(*tensors):
120 """Ensure all tensors are contiguous (no-op if already contiguous)."""
121 return tuple(t.contiguous() if not t.is_contiguous() else t for t in tensors)
124def _to_list_int64(val):
125 """Convert Tensor(int32) to List[int64] for aclnn kernel consumption."""
126 if isinstance(val, ms.Tensor):
127 return val.asnumpy().astype("int64").tolist()
128 return val
131class NpuDenseLightningIndexerSoftmaxLseDFunction(DFunction): # pylint: disable=W0221
132 """DFunction wrapper for npu_dense_lightning_indexer_softmax_lse on MindSpore.
134 Routes plain-tensor calls directly to the MindSpore custom kernel, and
135 DTensor calls through the distributed dispatch framework using the
136 registered DistributedOp with the same op_name.
138 All 11 forward arguments after ``ctx`` are positional to stay compatible
139 with both MindSpore autograd function conventions.
141 No backward is defined because the operator does not require gradients.
142 """
144 _op_name = "npu_dense_lightning_indexer_softmax_lse"
146 @staticmethod
147 def forward(ctx, query_index, key_index, weights,
148 actual_seq_qlen, actual_seq_klen,
149 layout, sparse_mode, pre_tokens, next_tokens):
150 """Forward pass: delegates to the MindSpore Ascend custom kernel.
152 Args:
153 ctx: Autograd context.
154 query_index: Lightning Indexer query input (Q̃).
155 key_index: Lightning Indexer key input (K̃).
156 weights: Lightning Indexer weight coefficient (W).
157 actual_seq_qlen: Cumulative query sequence lengths; None for BSND.
158 actual_seq_klen: Cumulative key sequence lengths; None for BSND.
159 layout: Data layout format, 'BSND' or 'TND'.
160 sparse_mode: Sparse computation mode (only mode 3 supported).
161 pre_tokens: Number of preceding tokens for sparse attention.
162 next_tokens: Number of following tokens for sparse attention.
164 Returns:
165 tuple[Tensor, Tensor]: (softmax_max_index, softmax_sum_index), both float32.
166 """
167 return _custom_ops.npu_dense_lightning_indexer_softmax_lse(
168 query_index, key_index, weights,
169 _to_list_int64(actual_seq_qlen), _to_list_int64(actual_seq_klen),
170 layout, sparse_mode, pre_tokens, next_tokens,
171 )
173 @staticmethod
174 def backward(ctx, *grad_outputs):
175 """No-op backward — this operator does not require gradients."""
176 return (None,) * 9
179class NpuDenseLightningIndexerGradKlLossDFunction(DFunction): # pylint: disable=W0221
180 """DFunction wrapper for npu_dense_lightning_indexer_grad_kl_loss on MindSpore.
182 Routes plain-tensor calls directly to the MindSpore custom kernel, and
183 DTensor calls through the distributed dispatch framework using the
184 registered DistributedOp with the same op_name.
186 All 18 forward arguments after ``ctx`` are positional to stay compatible
187 with both MindSpore autograd function conventions.
188 """
190 _op_name = "npu_dense_lightning_indexer_grad_kl_loss"
192 @staticmethod
193 def forward(ctx, query, key, query_index, key_index, weights,
194 softmax_max, softmax_sum, softmax_max_index, softmax_sum_index,
195 scale_value, query_rope, key_rope,
196 actual_seq_qlen, actual_seq_klen,
197 layout, sparse_mode, pre_tokens, next_tokens):
198 """Forward pass: delegates to the MindSpore Ascend custom kernel.
200 Args:
201 ctx: Autograd context.
202 query: Main attention query (Q). dtype bfloat16/float16.
203 key: Main attention key (K). dtype bfloat16/float16.
204 query_index: Lightning Indexer query input (Q̃). dtype bfloat16/float16.
205 key_index: Lightning Indexer key input (K̃). dtype bfloat16/float16.
206 weights: Lightning Indexer weight coefficient (W).
207 softmax_max: Attention softmax max values. dtype float32.
208 softmax_sum: Attention softmax sum values. dtype float32.
209 softmax_max_index: Index attention softmax max (from softmax_lse). dtype float32.
210 softmax_sum_index: Index attention softmax sum (from softmax_lse). dtype float32.
211 scale_value: Scaling factor. dtype float32.
212 query_rope: Optional MLA query rope tensor.
213 key_rope: Optional MLA key rope tensor.
214 actual_seq_qlen: Cumulative query sequence lengths; None for BSND.
215 actual_seq_klen: Cumulative key sequence lengths; None for BSND.
216 layout: Data layout format, 'BSND' or 'TND'.
217 sparse_mode: Sparse computation mode (only mode 3 supported).
218 pre_tokens: Number of preceding tokens for sparse attention.
219 next_tokens: Number of following tokens for sparse attention.
221 Returns:
222 tuple[Tensor, Tensor, Tensor, Tensor]:
223 (d_query_index, d_key_index, d_weights, loss).
224 """
225 result = _custom_ops.npu_dense_lightning_indexer_grad_kl_loss(
226 query, key, query_index, key_index, weights,
227 softmax_max, softmax_sum, softmax_max_index, softmax_sum_index,
228 scale_value, query_rope, key_rope,
229 _to_list_int64(actual_seq_qlen), _to_list_int64(actual_seq_klen),
230 layout, sparse_mode, pre_tokens, next_tokens,
231 )
232 ctx.save_for_backward(result[0], result[1], result[2])
233 return result
235 @staticmethod
236 def backward(ctx, *grad_outputs):
237 d_query_index, d_key_index, d_weights = _ensure_contiguous(*ctx.saved_tensors)
238 return (None, None, d_query_index, d_key_index, d_weights,
239 None, None, None, None, None, None, None, None, None, None, None, None, None)
242class NpuSparseLightningIndexerGradKlLossDFunction(DFunction): # pylint: disable=W0221
243 """DFunction wrapper for npu_sparse_lightning_indexer_grad_kl_loss on MindSpore.
245 Routes plain-tensor calls directly to the MindSpore custom kernel, and
246 DTensor calls through the distributed dispatch framework using the
247 registered DistributedOp with the same op_name.
249 All 17 forward arguments after ``ctx`` are positional to stay compatible
250 with both MindSpore autograd function conventions.
251 """
253 _op_name = "npu_sparse_lightning_indexer_grad_kl_loss"
255 @staticmethod
256 def forward(ctx, query, key, query_index, key_index, weights,
257 sparse_indices, softmax_max, softmax_sum, scale_value,
258 query_rope, key_rope,
259 actual_seq_qlen, actual_seq_klen,
260 layout, sparse_mode, pre_tokens, next_tokens):
261 """Forward pass: delegates to the MindSpore Ascend custom kernel.
263 Args:
264 ctx: Autograd context.
265 query: Main attention query (q_t). dtype bfloat16/float16.
266 key: Main attention key (K_t). dtype bfloat16/float16.
267 query_index: Lightning Indexer query input (q̃_t). dtype bfloat16/float16.
268 key_index: Lightning Indexer key input (K̃_t). dtype bfloat16/float16.
269 weights: Lightning Indexer weight coefficient (W_t).
270 sparse_indices: Sorted token indices for key/key_index. dtype bfloat16/float16.
271 softmax_max: Attention softmax max values.
272 softmax_sum: Attention softmax sum values.
273 scale_value: Scaling factor. dtype float.
274 query_rope: Optional MLA query rope tensor.
275 key_rope: Optional MLA key rope tensor.
276 actual_seq_qlen: Cumulative query sequence lengths; None for BSND.
277 actual_seq_klen: Cumulative key sequence lengths; None for BSND.
278 layout: Data layout format, 'BSND' or 'TND'.
279 sparse_mode: Sparse computation mode (only mode 3 supported).
280 pre_tokens: Number of preceding tokens for sparse attention.
281 next_tokens: Number of following tokens for sparse attention.
283 Returns:
284 tuple[Tensor, Tensor, Tensor, Tensor]:
285 (d_query_index, d_key_index, d_weights, loss).
286 """
287 result = _custom_ops.npu_sparse_lightning_indexer_grad_kl_loss(
288 query, key, query_index, key_index, weights,
289 sparse_indices, softmax_max, softmax_sum, scale_value,
290 query_rope, key_rope,
291 _to_list_int64(actual_seq_qlen), _to_list_int64(actual_seq_klen),
292 layout, sparse_mode, pre_tokens, next_tokens,
293 )
294 ctx.save_for_backward(result[0], result[1], result[2])
295 return result
297 @staticmethod
298 def backward(ctx, *grad_outputs):
299 d_query_index, d_key_index, d_weights = _ensure_contiguous(*ctx.saved_tensors)
300 return (None, None, d_query_index, d_key_index, d_weights,
301 None, None, None, None, None, None, None, None, None, None, None, None)
304class NpuMhcPostDFunction(DFunction): # pylint: disable=W0221
305 """DFunction wrapper for npu_mhc_post on MindSpore.
307 Routes plain-tensor calls directly to the MindSpore custom kernel, and
308 DTensor calls through the distributed dispatch framework using the
309 registered DistributedOp with the same op_name.
311 All 4 forward arguments after ``ctx`` are positional to stay compatible
312 with both MindSpore autograd function conventions.
313 """
315 _op_name = "npu_mhc_post"
317 @staticmethod
318 def forward(ctx, x, h_res, h_out, h_post):
319 """Forward pass: delegates to the MindSpore Ascend custom kernel.
321 Args:
322 ctx: Autograd context.
323 x: Input tensor of shape [B,S,N,D] or [T,N,D]. dtype bfloat16/float16.
324 h_res: mHC h_res transformation matrix. dtype float32.
325 h_out: Attention/MLP layer output. dtype bfloat16/float16.
326 h_post: mHC h_post transformation matrix. dtype float32.
328 Returns:
329 Tensor: Output tensor with same shape and dtype as x.
330 """
331 ctx.save_for_backward(x, h_res, h_out, h_post)
332 return _custom_ops.npu_mhc_post(x, h_res, h_out, h_post)
334 @staticmethod
335 def backward(ctx, *grad_outputs):
336 """Backward pass: calls npu_mhc_post_backward kernel.
338 Args:
339 ctx: Autograd context.
340 grad_outputs: Upstream gradients; grad_outputs[0] is grad_y.
342 Returns:
343 tuple: (grad_x, grad_h_res, grad_h_out, grad_h_post).
344 """
345 x, h_res, h_out, h_post = ctx.saved_tensors
346 grad_y, x, h_res, h_out, h_post = _ensure_contiguous(
347 grad_outputs[0], x, h_res, h_out, h_post)
348 grads = _custom_ops.npu_mhc_post_backward(
349 grad_y, x, h_res, h_out, h_post)
350 return grads[0], grads[1], grads[2], grads[3]
353class NpuMhcPreSinkhornDFunction(DFunction): # pylint: disable=W0221
354 """DFunction wrapper for npu_mhc_pre_sinkhorn on MindSpore.
356 Routes plain-tensor calls directly to the MindSpore custom kernel, and
357 DTensor calls through the distributed dispatch framework using the
358 registered DistributedOp with the same op_name.
360 All 9 forward arguments after ``ctx`` are positional to stay compatible
361 with both MindSpore autograd function conventions.
362 """
364 _op_name = "npu_mhc_pre_sinkhorn"
366 @staticmethod
367 def forward(ctx, x, phi, alpha, bias, hc_mult, num_iters, hc_eps, norm_eps, out_flag):
368 """Forward pass: delegates to the MindSpore Ascend custom kernel.
370 Args:
371 ctx: Autograd context.
372 x: Input tensor. dtype bfloat16/float16.
373 phi: mHC parameter matrix. dtype float32.
374 alpha: mHC scaling parameters. dtype float32.
375 bias: mHC bias parameters. dtype float32.
376 hc_mult: HC dimension size (currently only 4 supported).
377 num_iters: Sinkhorn iteration count.
378 hc_eps: H_pre sigmoid eps parameter.
379 norm_eps: RmsNorm eps parameter.
380 out_flag: Whether to output intermediate gradients.
382 Returns:
383 tuple[Tensor, ...]: 8 output tensors
384 (h_in, h_post, h_res, h_pre, hc_before_norm, inv_rms, sum_out, norm_out).
385 """
386 result = _custom_ops.npu_mhc_pre_sinkhorn(
387 x, phi, alpha, bias, hc_mult, num_iters, hc_eps, norm_eps, out_flag
388 )
389 _, _, _, h_pre, hc_before_norm, inv_rms, sum_out, norm_out = result
390 ctx.save_for_backward(x, phi, alpha, bias,
391 h_pre, hc_before_norm, inv_rms, sum_out, norm_out)
392 ctx.hc_eps = hc_eps
393 return result
395 @staticmethod
396 def backward(ctx, *grad_outputs):
397 """Backward pass: calls npu_mhc_pre_sinkhorn_backward kernel.
399 Args:
400 ctx: Autograd context.
401 grad_outputs: Upstream gradients for the 8 forward outputs.
402 grad_outputs[0]=grad_h_in, [1]=grad_h_post, [2]=grad_h_res;
403 [3..7] correspond to saved intermediates and are None.
405 Returns:
406 tuple: (grad_x, grad_phi, grad_alpha, grad_bias, None×5) —
407 gradients for the 9 forward inputs.
408 """
409 x, phi, alpha, bias, h_pre, hc_before_norm, inv_rms, sum_out, norm_out = ctx.saved_tensors
410 (grad_h_in, grad_h_post, grad_h_res,
411 x, phi, alpha, bias,
412 h_pre, hc_before_norm, inv_rms, sum_out, norm_out) = _ensure_contiguous(
413 grad_outputs[0], grad_outputs[1], grad_outputs[2],
414 x, phi, alpha, bias,
415 h_pre, hc_before_norm, inv_rms, sum_out, norm_out)
416 b, s, n = grad_h_post.shape
417 grad_h_res = grad_h_res.reshape(b, s, n, n)
418 grads = _custom_ops.npu_mhc_pre_sinkhorn_backward(
419 grad_h_in, grad_h_post, grad_h_res,
420 x, phi, alpha, bias,
421 h_pre, hc_before_norm, inv_rms, sum_out, norm_out,
422 ctx.hc_eps)
423 return grads[0], grads[1], grads[2], grads[3], None, None, None, None, None
426class NpuMhcPreClampSinkhornDFunction(DFunction): # pylint: disable=W0221
427 """DFunction wrapper for npu_mhc_pre_clamp_sinkhorn on MindSpore.
429 This matches the static-graph aclnnMhcPreClampSinkhorn integration:
430 forward has 11 arguments and returns 9 tensors, and backward consumes
431 h_res_logits plus clamp_min/clamp_max.
432 """
434 _op_name = "npu_mhc_pre_clamp_sinkhorn"
436 @staticmethod
437 def forward(ctx, *args, **kwargs):
438 """Forward pass: delegates to the clamp-enabled Ascend custom kernel."""
439 bound = _bind_mhc_pre_clamp_args(args, kwargs)
440 result = _custom_ops.npu_mhc_pre_clamp_sinkhorn(
441 bound.x, bound.phi, bound.alpha, bound.bias,
442 bound.hc_mult, bound.num_iters, bound.hc_eps, bound.norm_eps,
443 bound.out_flag, bound.clamp_min, bound.clamp_max
444 )
445 _, _, _, h_pre, hc_before_norm, inv_rms, sum_out, norm_out, h_res_logits = result
446 ctx.save_for_backward(bound.x, bound.phi, bound.alpha, bound.bias,
447 h_pre, hc_before_norm, inv_rms, sum_out, norm_out, h_res_logits)
448 ctx.hc_eps = bound.hc_eps
449 ctx.clamp_min = bound.clamp_min
450 ctx.clamp_max = bound.clamp_max
451 return result
453 @staticmethod
454 def backward(ctx, *grad_outputs):
455 """Backward pass: calls npu_mhc_pre_clamp_sinkhorn_backward kernel."""
456 tensors = _ensure_contiguous(
457 grad_outputs[0], grad_outputs[1], grad_outputs[2],
458 *ctx.saved_tensors
459 )
460 n = tensors[1].shape[-1]
461 grad_h_res = ms.ops.reshape(tensors[2], tuple(tensors[2].shape[:-1]) + (n, n))
463 grads = _custom_ops.npu_mhc_pre_clamp_sinkhorn_backward(
464 tensors[0], tensors[1], grad_h_res,
465 tensors[3], tensors[4], tensors[5], tensors[6],
466 tensors[7], tensors[8], tensors[9], tensors[10], tensors[11], tensors[12],
467 ctx.hc_eps, ctx.clamp_min, ctx.clamp_max)
468 return tuple(grads[:4]) + _MHC_PRE_CLAMP_NONE_GRADS
471class NpuLightningIndexerDFunction(DFunction): # pylint: disable=W0221
472 """DFunction wrapper for npu_lightning_indexer.
474 The underlying kernel handles all cmp_ratio values (1 / 4 / 128) directly.
475 Forward-only: indexer gradients are produced by the network's explicit
476 ``sparse_lightning_indexer_kl_loss_grad`` call.
478 Signature mirrors the torch-extension ``lightning_indexer`` benchmark:
479 positional ``(query, key, weights, sparse_count)`` (``sparse_count`` is
480 benchmark ``topk``); the two layouts are merged into a single ``layout``
481 (the kernel is fed identical ``layout_q`` / ``layout_k``).
482 """
484 _op_name = "npu_lightning_indexer"
486 @staticmethod
487 def forward(ctx, query, key, weights, sparse_count,
488 cu_seq_lens_q=None, cu_seq_lens_k=None, cmp_residual_k=None,
489 block_table=None, layout="BSND",
490 sparse_mode=0, cmp_ratio=1, return_value=False):
491 """Forward pass: call the custom kernel for all cmp_ratios.
493 Remaining benchmark kwargs (seqused_q/k, output_idx_offset, metadata,
494 max_seqlen_q) are presently unused by the external API and pinned to
495 ``None`` / ``-1``.
497 Returns:
498 tuple[Tensor, Tensor]: (sparse_indices, sparse_values).
499 """
500 return _custom_ops.npu_lightning_indexer_v2(
501 query, key, weights, sparse_count,
502 cu_seq_lens_q, cu_seq_lens_k,
503 None, None, cmp_residual_k, block_table, None, None, -1,
504 layout, layout, sparse_mode, cmp_ratio, return_value)
506 @staticmethod
507 def backward(ctx, *grad_outputs):
508 """No-op backward — indexer gradients come from kl_loss_grad."""
509 return (None,) * 12
512class NpuSparseFlashMlaDFunction(DFunction): # pylint: disable=W0221
513 """DFunction wrapper for the MLA sparse-attention kernel.
515 Forward runs ``npu_sparse_flash_mla`` (the kernel derives its metadata from
516 the tensor shapes internally); backward runs ``npu_sparse_flash_mla_grad``.
517 """
519 _op_name = "npu_sparse_flash_mla"
521 @staticmethod
522 def forward(ctx, # pylint: disable=too-many-arguments,too-many-locals,too-many-statements
523 query, ori_kv, cmp_kv,
524 cu_seq_lens_q, cu_seq_lens_ori_kv, cu_seq_lens_cmp_kv,
525 ori_sparse_indices, cmp_sparse_indices, sinks,
526 softmax_scale, cmp_ratio, ori_mask_mode, cmp_mask_mode,
527 ori_win_left, ori_win_right,
528 layout_q, layout_kv,
529 cmp_residual_kv=None, seqused_ori_kv=None, seqused_cmp_kv=None,
530 seqused_q=None):
531 """Forward pass: runs MLA sparse attention (metadata computed in-kernel).
533 Args:
534 ctx: Autograd context.
535 query: Query tensor. dtype bfloat16/float16.
536 ori_kv: Original KV tensor; None when absent.
537 cmp_kv: Compressed KV tensor; None when absent.
538 cu_seq_lens_q: Cumulative query seq lengths (TND); None for BSND.
539 cu_seq_lens_ori_kv: Cumulative ori_kv seq lengths; None for PA_ND.
540 cu_seq_lens_cmp_kv: Cumulative cmp_kv seq lengths; None for PA_ND.
541 ori_sparse_indices: Sparse indices for ori_kv; None = band mode.
542 cmp_sparse_indices: Sparse indices for cmp_kv (int32 Tensor).
543 sinks: Attention-sink tensor (float32); None when absent.
544 softmax_scale: Softmax scaling factor (float).
545 cmp_ratio: KV compression ratio (int).
546 ori_mask_mode: Mask mode for q×ori_kv (default 4=band).
547 cmp_mask_mode: Mask mode for q×cmp_kv (default 3=rightDownCausal).
548 ori_win_left: Band-mask left window (default 127).
549 ori_win_right: Band-mask right window (default 0).
550 layout_q: Q data layout — 'BSND' or 'TND'.
551 layout_kv: KV data layout — 'PA_ND' or 'BSND'.
553 Returns:
554 tuple[Tensor, Tensor]: (attention_out, softmax_lse).
555 """
556 if cmp_ratio != 4:
557 cmp_sparse_indices = None
559 # The kernel computes its metadata internally. topk_value_mode=1;
560 # return_softmax_lse is forced True internally so the backward always
561 # receives a valid LSE (a stale/zero LSE makes the grad kernel explode);
562 # the external return value is gated separately by the wrapper's own
563 # return_softmax_lse flag, independent of this.
564 result = _custom_ops.npu_sparse_flash_mla(
565 query, ori_kv, cmp_kv, ori_sparse_indices, cmp_sparse_indices,
566 None, None, # ori_block_table, cmp_block_table
567 cu_seq_lens_q, cu_seq_lens_ori_kv, cu_seq_lens_cmp_kv,
568 seqused_q, seqused_ori_kv, seqused_cmp_kv, # seq_used_q, seq_used_ori_kv, seq_used_cmp_kv
569 cmp_residual_kv, None, None, # cmp_residual_kv, ori_topk_length, cmp_topk_length
570 sinks,
571 softmax_scale, cmp_ratio, ori_mask_mode, cmp_mask_mode,
572 ori_win_left, ori_win_right, layout_q, layout_kv, 1, True,
573 )
574 attention_out, softmax_lse = result[0], result[1]
576 ctx.has_ori_kv = ori_kv is not None
577 ctx.has_cmp_kv = cmp_kv is not None
578 ctx.has_sinks = sinks is not None
579 ctx.has_ori_sparse = ori_sparse_indices is not None
580 ctx.has_cmp_sparse = cmp_sparse_indices is not None
581 ctx.has_cu_q = cu_seq_lens_q is not None
582 ctx.has_cu_ori_kv = cu_seq_lens_ori_kv is not None
583 ctx.has_cu_cmp_kv = cu_seq_lens_cmp_kv is not None
584 ctx.has_cmp_residual = cmp_residual_kv is not None
585 # metadata is NOT saved for backward: the grad kernel asserts metadata
586 # must be nullptr and re-derives its own tiling internally. cmp_residual_kv
587 # IS saved — the grad kernel requires it for CFA/SCFA with cmp_mask_mode=3.
588 ctx.save_for_backward(*[t for t in [
589 query, ori_kv, cmp_kv, sinks, ori_sparse_indices, cmp_sparse_indices,
590 cu_seq_lens_q, cu_seq_lens_ori_kv, cu_seq_lens_cmp_kv,
591 attention_out, softmax_lse, cmp_residual_kv,
592 ] if t is not None])
593 ctx.softmax_scale = softmax_scale
594 ctx.cmp_ratio = cmp_ratio
595 ctx.ori_mask_mode = ori_mask_mode
596 ctx.cmp_mask_mode = cmp_mask_mode
597 ctx.ori_win_left = ori_win_left
598 ctx.ori_win_right = ori_win_right
599 ctx.layout_q = layout_q
600 ctx.layout_kv = layout_kv
601 return attention_out, softmax_lse
603 @staticmethod
604 def backward(ctx, grad_attention_out, grad_softmax_lse): # pylint: disable=unused-argument
605 """Backward pass: calls npu_sparse_flash_mla_grad kernel."""
606 it = iter(ctx.saved_tensors)
607 q = next(it)
608 ori_kv = next(it) if ctx.has_ori_kv else None
609 cmp_kv = next(it) if ctx.has_cmp_kv else None
610 sinks = next(it) if ctx.has_sinks else None
611 ori_sparse_indices = next(it) if ctx.has_ori_sparse else None
612 cmp_sparse_indices = next(it) if ctx.has_cmp_sparse else None
613 cu_seq_lens_q = next(it) if ctx.has_cu_q else None
614 cu_seq_lens_ori_kv = next(it) if ctx.has_cu_ori_kv else None
615 cu_seq_lens_cmp_kv = next(it) if ctx.has_cu_cmp_kv else None
616 attention_out = next(it)
617 softmax_lse = next(it)
618 cmp_residual_kv = next(it) if ctx.has_cmp_residual else None
619 # metadata MUST be None: the grad kernel asserts it is nullptr and
620 # re-derives tiling internally. cmp_residual_kv is passed through —
621 # required for CFA/SCFA (cmp_ratio!=1) with cmp_mask_mode=3.
622 grads = _custom_ops.npu_sparse_flash_mla_grad(
623 q, grad_attention_out, attention_out, softmax_lse,
624 ori_kv, cmp_kv, ori_sparse_indices, cmp_sparse_indices,
625 cu_seq_lens_q, cu_seq_lens_ori_kv, cu_seq_lens_cmp_kv,
626 None, None, None, # seq_used_q, seq_used_ori_kv, seq_used_cmp_kv
627 cmp_residual_kv, None, None, # cmp_residual_kv, ori_topk_length, cmp_topk_length
628 sinks, None, # sinks, metadata(None → grad kernel self-derives)
629 ctx.softmax_scale, ctx.cmp_ratio, ctx.ori_mask_mode, ctx.cmp_mask_mode,
630 ctx.ori_win_left, ctx.ori_win_right, ctx.layout_q, ctx.layout_kv,
631 )
632 d_query = grads[0]
633 d_ori_kv = grads[1] if ori_kv is not None else None
634 d_cmp_kv = grads[2] if cmp_kv is not None else None
635 d_sinks = grads[3] if sinks is not None else None
636 # grads[4], grads[5] = ori/cmp_softmax_l1_norm — discarded here.
637 # 21 positional forward args (ctx excluded):
638 # query, ori_kv, cmp_kv, cu_seq_lens_q, cu_seq_lens_ori_kv, cu_seq_lens_cmp_kv,
639 # ori_sparse_indices, cmp_sparse_indices, sinks,
640 # softmax_scale, cmp_ratio, ori_mask_mode, cmp_mask_mode, ori_win_left, ori_win_right,
641 # layout_q, layout_kv, cmp_residual_kv, seqused_ori_kv, seqused_cmp_kv, seqused_q
642 return (d_query, d_ori_kv, d_cmp_kv,
643 None, None, None,
644 None, None, d_sinks,
645 None, None, None, None, None, None,
646 None, None, None, None, None, None)
649def npu_sparse_flash_mla_grad(*args, **kwargs):
650 """Raw ``sparse_flash_mla_grad`` kernel passthrough (stateless, no autograd).
652 Runs the same backward kernel as ``NpuSparseFlashMlaDFunction.backward``, but
653 returns its full 6-tuple so a network-defined custom backward can also
654 consume ``ori/cmp_softmax_l1_norm`` (the main-attention target distribution
655 ``p`` for the Lightning-Indexer KL loss). Intended to be called from inside
656 another custom function's ``backward`` (autograd already off); it builds no
657 graph. ``metadata`` must be ``None`` — the grad kernel re-derives its own
658 tiling internally.
660 Returns:
661 tuple[Tensor, ...]: ``(d_query, d_ori_kv, d_cmp_kv, d_sinks,
662 ori_softmax_l1_norm, cmp_softmax_l1_norm)``.
663 """
664 return _custom_ops.npu_sparse_flash_mla_grad(*args, **kwargs)
667class NpuSparseLightningIndexerKlLossGradDFunction(DFunction): # pylint: disable=W0221
668 """DFunction wrapper for ``npu_sparse_lightning_indexer_kl_loss_grad``.
670 The kernel takes the pre-computed main-attention target distribution
671 ``attn_softmax_l1_norm`` and produces ``(dq, dk, dw, softmax_out)`` — the
672 gradients w.r.t. ``query``/``key``/``weights`` plus the indexer-branch
673 softmax; it neither recomputes the main attention nor outputs a loss.
674 Metadata is computed inside the kernel from the tensor shapes. Backward
675 propagates ``(dq, dk, dw)`` to those inputs.
676 """
678 _op_name = "npu_sparse_lightning_indexer_kl_loss_grad"
680 @staticmethod
681 def forward(ctx, query, key, weights, sparse_indices, attn_softmax_l1_norm,
682 cu_seq_lens_q, cu_seq_lens_k, seqused_q, seqused_k, cmp_residual_k,
683 layout, mask_mode, cmp_ratio):
684 """Forward pass: runs the KL-loss grad kernel (metadata computed in-kernel).
686 Args:
687 ctx: Autograd context.
688 query: Lightning Indexer query (q̃). dtype bfloat16/float16.
689 key: Lightning Indexer key (k̃). dtype bfloat16/float16.
690 weights: Lightning Indexer weight coefficient (w).
691 sparse_indices: Sorted token indices (int32).
692 attn_softmax_l1_norm: Main-attention target distribution p (float32),
693 pre-computed by the main-attention branch.
694 cu_seq_lens_q: Cumulative query sequence lengths; None for BSND.
695 cu_seq_lens_k: Cumulative key sequence lengths; None for BSND.
696 seqused_q: Used query sequence lengths; None when absent.
697 seqused_k: Used key sequence lengths; None when absent.
698 cmp_residual_k: Optional compressed-KV residual.
699 layout: Data layout format — 'BSND' or 'TND'.
700 mask_mode: Sparse mask mode (only 3 supported).
701 cmp_ratio: KV compression ratio.
703 Returns:
704 tuple[Tensor, Tensor, Tensor, Tensor]:
705 (d_query, d_key, d_weights, softmax_out).
706 """
707 # The kernel computes its metadata internally.
708 result = _custom_ops.npu_sparse_lightning_indexer_kl_loss_grad(
709 query, key, weights, sparse_indices, attn_softmax_l1_norm,
710 cu_seq_lens_q, cu_seq_lens_k, seqused_q, seqused_k, cmp_residual_k,
711 layout, layout, mask_mode, cmp_ratio,
712 )
713 ctx.save_for_backward(result[0], result[1], result[2])
714 return result
716 @staticmethod
717 def backward(ctx, *grad_outputs):
718 """Backward: propagate the fused gradients to query/key/weights inputs."""
719 d_query, d_key, d_weights = _ensure_contiguous(*ctx.saved_tensors)
720 # 13 positional forward args: query, key, weights, sparse_indices,
721 # attn_softmax_l1_norm, cu_seq_lens_q, cu_seq_lens_k, seqused_q, seqused_k,
722 # cmp_residual_k, layout, mask_mode, cmp_ratio.
723 return (d_query, d_key, d_weights,
724 None, None, None, None, None, None, None, None, None, None)