Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / dsa_kl_loss.py: 83%
30 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"""DSA sparse-attention indexer KL loss."""
17from typing import Any, Tuple
19import torch # pylint: disable=forbidden-backend-import
21from hyper_parallel.tools.logging import get_logger
23logger = get_logger()
26def _sparse_indexer_kl_loss_grad(
27 query: torch.Tensor,
28 key: torch.Tensor,
29 query_index: torch.Tensor,
30 key_index: torch.Tensor,
31 weights: torch.Tensor,
32 sparse_indices: torch.Tensor,
33 softmax_max: torch.Tensor,
34 softmax_sum: torch.Tensor,
35 scale: float,
36 query_rope: torch.Tensor | None,
37 key_rope: torch.Tensor | None,
38 actual_seq_qlen: torch.Tensor | None,
39 actual_seq_klen: torch.Tensor | None,
40 layout: str,
41) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
42 """Compute teacher, indexer gradients and summed KL inside CANN."""
43 # Keep the optional NPU dependency out of CPU and Omni-only imports.
44 from torch_npu import ( # pylint: disable=import-outside-toplevel,forbidden-backend-import
45 npu_sparse_lightning_indexer_grad_kl_loss,
46 )
48 return npu_sparse_lightning_indexer_grad_kl_loss(
49 query.contiguous(), key.contiguous(), query_index.contiguous(), key_index.contiguous(),
50 weights.contiguous(), sparse_indices.contiguous(), softmax_max.contiguous(),
51 softmax_sum.contiguous(), scale,
52 query_rope=None if query_rope is None else query_rope.contiguous(),
53 key_rope=None if key_rope is None else key_rope.contiguous(),
54 actual_seq_qlen=None if actual_seq_qlen is None else actual_seq_qlen.tolist(),
55 actual_seq_klen=None if actual_seq_klen is None else actual_seq_klen.tolist(),
56 layout=layout, sparse_mode=3,
57 )
60class _DSAKLLoss(torch.autograd.Function):
61 """Autograd bridge for the DSA indexer KL-loss custom operator."""
63 @staticmethod
64 def forward(
65 ctx: Any,
66 index_query: torch.Tensor,
67 index_key: torch.Tensor,
68 merge_weight: torch.Tensor,
69 query: torch.Tensor,
70 key: torch.Tensor,
71 topk_indices: torch.Tensor,
72 softmax_max: torch.Tensor,
73 softmax_sum: torch.Tensor,
74 query_rope: torch.Tensor,
75 key_rope: torch.Tensor,
76 actual_seq_qlen: torch.Tensor,
77 actual_seq_kvlen: torch.Tensor,
78 scale: float,
79 loss_coeff: float,
80 backend: str,
81 ) -> torch.Tensor:
82 """Compute the KL loss and save its precomputed indexer gradients."""
83 divisor = query.size(0)
84 if divisor == 0:
85 raise ValueError("DSA KL requires at least one valid query token")
86 # Load only the selected optional backend so either can be used independently.
87 # pylint: disable=import-outside-toplevel
88 if backend == "cann":
89 result = _sparse_indexer_kl_loss_grad(
90 query,
91 key,
92 index_query,
93 index_key,
94 merge_weight,
95 topk_indices,
96 softmax_max,
97 softmax_sum,
98 scale,
99 query_rope,
100 key_rope,
101 actual_seq_qlen,
102 actual_seq_kvlen,
103 "TND",
104 )
105 else:
106 import omni_training_custom_ops # noqa: F401 # pylint: disable=unused-import
108 result = torch.ops.custom.npu_sparse_lightning_indexer_grad_kl_loss_enhance(
109 query=query,
110 key=key,
111 query_index=index_query,
112 key_index=index_key,
113 weights=merge_weight,
114 sparse_indices=topk_indices,
115 softmax_max=softmax_max,
116 softmax_sum=softmax_sum,
117 scale_value=scale,
118 query_rope=query_rope,
119 key_rope=key_rope,
120 actual_seq_qlen=actual_seq_qlen,
121 actual_seq_klen=actual_seq_kvlen,
122 layout="TND",
123 sparse_mode=3,
124 deterministic=torch.are_deterministic_algorithms_enabled(),
125 sparse_block_size=1,
126 )
127 # pylint: enable=import-outside-toplevel
128 grad_index_query, grad_index_key, grad_merge_weight, inner_loss = result
129 ctx.save_for_backward(
130 grad_index_query / divisor * loss_coeff,
131 grad_index_key / divisor * loss_coeff,
132 grad_merge_weight / divisor * loss_coeff,
133 )
134 return inner_loss.squeeze() / divisor * loss_coeff
136 @staticmethod
137 def backward(ctx: Any, grad_loss: torch.Tensor) -> Tuple:
138 """Scale the precomputed indexer gradients by the upstream gradient."""
139 grad_index_query, grad_index_key, grad_merge_weight = ctx.saved_tensors
140 return (
141 grad_index_query * grad_loss,
142 grad_index_key * grad_loss,
143 grad_merge_weight * grad_loss,
144 *((None,) * 12),
145 )
148def dsa_kl_loss(
149 index_query: torch.Tensor,
150 index_key: torch.Tensor,
151 merge_weight: torch.Tensor,
152 query: torch.Tensor,
153 key: torch.Tensor,
154 topk_indices: torch.Tensor,
155 softmax_max: torch.Tensor,
156 softmax_sum: torch.Tensor,
157 query_rope: torch.Tensor,
158 key_rope: torch.Tensor,
159 actual_seq_qlen: torch.Tensor,
160 actual_seq_kvlen: torch.Tensor,
161 scale: float,
162 loss_coeff: float,
163 *,
164 backend: str = "cann",
165) -> torch.Tensor:
166 """Compute the DSA indexer KL loss.
168 ``index_query``, ``index_key``, and ``merge_weight`` are the flattened
169 tensors returned by :func:`dsa_indexer`. ``topk_indices``, ``softmax_max``,
170 and ``softmax_sum`` come from :func:`dsa_sparse_attention`. The MLA query,
171 key, and rotary states use TND layout; sequence lengths are cumulative
172 int32 tensors, following the custom-operator call contract.
173 The result is a scalar auxiliary loss. The keyword-only ``backend`` selects
174 ``"cann"`` (default) or ``"omni"`` for the legacy implementation.
175 """
176 if backend not in ("cann", "omni"):
177 raise ValueError(f"backend must be 'cann' or 'omni', got {backend!r}")
178 if backend == "omni":
179 logger.warning(
180 "dsa_kl_loss backend='omni' is deprecated and will be removed in a future release; "
181 "use backend='cann'."
182 )
183 return _DSAKLLoss.apply(
184 index_query,
185 index_key,
186 merge_weight,
187 query,
188 key,
189 topk_indices,
190 softmax_max,
191 softmax_sum,
192 query_rope,
193 key_rope,
194 actual_seq_qlen,
195 actual_seq_kvlen,
196 scale,
197 loss_coeff,
198 backend,
199 )