Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / dsa_attention.py: 0%

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"""Model-compatible DSA modules using the Hyper high-performance operators.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Mapping, Sequence 

20from typing import Any 

21 

22# This package provides PyTorch-specific high-performance modules. 

23# pylint: disable=forbidden-backend-import 

24import torch # pylint: disable=forbidden-backend-import 

25from torch import nn 

26from torch.nn import functional as F 

27from hyper_parallel.components.checkpoint.weight_conversion import WeightConverter 

28 

29from hyper_parallel.components.checkpoint import ConcatenateWithSections 

30from hyper_parallel.models.replacement import module_replacement 

31from hyper_parallel.components.functional import aggregate_hidden, aux_loss_auto_scale 

32from hyper_parallel.components.functional import apply_rotary_pos_emb, apply_rotary_pos_emb_interleave 

33from hyper_parallel.components.functional import ( 

34 dsa_indexer, 

35 dsa_kl_loss, 

36 dsa_sparse_attention, 

37 dsa_sparse_attention_rescale, 

38) 

39from hyper_parallel.components.functional.npu_fusion_attention import resolve_packed_sequence_lengths 

40 

41 

42def _dense_flash_attention( 

43 module: nn.Module, 

44 query: torch.Tensor, 

45 key: torch.Tensor, 

46 value: torch.Tensor, 

47 sequence_lengths: tuple[torch.Tensor, torch.Tensor], 

48 *, 

49 need_stats: bool, 

50) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: 

51 """Run one TND FA and return dense-indexer statistics.""" 

52 query_ends, key_ends = (lengths.tolist() for lengths in sequence_lengths) 

53 q_flat, k_flat, v_flat = (tensor.flatten(0, 1) for tensor in (query, key, value)) 

54 num_heads = q_flat.shape[1] 

55 causal_mask = torch.ones((2048, 2048), dtype=torch.bool, device=query.device).triu(diagonal=1) 

56 import torch_npu # pylint: disable=import-outside-toplevel,forbidden-backend-import 

57 

58 result = torch_npu.npu_fusion_attention( 

59 q_flat.contiguous(), k_flat.contiguous(), v_flat.contiguous(), num_heads, "TND", 

60 pse=None, padding_mask=None, atten_mask=causal_mask, scale=module.scaling, 

61 pre_tockens=getattr(module, "pre_tockens", 1048576), 

62 next_tockens=getattr(module, "next_tockens", 0), keep_prob=1.0, inner_precise=0, 

63 actual_seq_qlen=query_ends, actual_seq_kvlen=key_ends, sparse_mode=3, 

64 ) 

65 output = result[0].reshape(*query.shape[:3], value.shape[-1]) 

66 if not need_stats: 

67 return output, None, None 

68 

69 # TND FA stores each packed segment as [head, token, 8], while exposing a 

70 # physical [token, head, 8] shape. Dense KL consumes [head, token, 1]. 

71 statistics = [] 

72 starts = [0, *query_ends[:-1]] 

73 for statistic in result[1:3]: 

74 rows = statistic.reshape(-1, statistic.shape[-1]) 

75 segments = [ 

76 rows[start * num_heads : end * num_heads].reshape( 

77 num_heads, end - start, statistic.shape[-1] 

78 )[..., :1] 

79 for start, end in zip(starts, query_ends) 

80 ] 

81 statistics.append(torch.cat(segments, dim=1).detach()) 

82 return output, statistics[0], statistics[1] 

83 

84 

85def _deepseek_v32_dense_attention( 

86 module: nn.Module, 

87 index_states: tuple[torch.Tensor, ...], 

88 attention_states: tuple[torch.Tensor, ...], 

89 sequence_lengths: tuple[torch.Tensor, torch.Tensor], 

90) -> torch.Tensor: 

91 """Run unabsorbed MLA FA and the fused dense indexer KL on local states.""" 

92 query, compressed_kv, query_rope, key_rope = attention_states 

93 projected_kv = module.kv_b_proj(compressed_kv.squeeze(2)).view( 

94 *compressed_kv.shape[:2], module.num_heads, module.qk_nope_head_dim + module.v_head_dim, 

95 ) 

96 key, value = projected_kv.split((module.qk_nope_head_dim, module.v_head_dim), dim=-1) 

97 key_rope = key_rope.expand(-1, -1, module.num_heads, -1) 

98 need_stats = bool(module.training and not module.freeze_dsa and module.dsa_loss_coeff) 

99 output, maximum, total = _dense_flash_attention( 

100 module, torch.cat((query, query_rope), dim=-1), torch.cat((key, key_rope), dim=-1), 

101 value, sequence_lengths, need_stats=need_stats, 

102 ) 

103 if need_stats: 

104 # Dense CANN operators are optional for sparse and Omni execution. 

105 # pylint: disable=import-outside-toplevel 

106 from hyper_parallel.components.functional.dsa_dense import ( 

107 dsa_dense_indexer_kl_loss, 

108 dsa_dense_indexer_lse, 

109 ) 

110 # pylint: enable=import-outside-toplevel 

111 

112 query_lengths, key_lengths = sequence_lengths 

113 index_query, index_key, weights = (tensor.flatten(0, 1) for tensor in index_states) 

114 indexer_stats = dsa_dense_indexer_lse( 

115 index_query, index_key, weights, query_lengths=query_lengths, key_lengths=key_lengths, 

116 ) 

117 loss = dsa_dense_indexer_kl_loss( 

118 index_query, index_key, weights, query.flatten(0, 1), key.flatten(0, 1), 

119 maximum, total, indexer_stats, scale=module.scaling, 

120 query_rope=query_rope.flatten(0, 1), key_rope=key_rope.flatten(0, 1), 

121 query_lengths=query_lengths, key_lengths=key_lengths, loss_coefficient=module.dsa_loss_coeff, 

122 ) 

123 output = aux_loss_auto_scale(output, loss) 

124 return output 

125 

126 

127def apply_mome( 

128 hidden_states: torch.Tensor, 

129 mome_mask: torch.Tensor | None, 

130 convolution: nn.Conv1d, 

131 *, 

132 fused: bool, 

133) -> torch.Tensor: 

134 """Apply masked causal depthwise convolution and residual.""" 

135 if mome_mask is None: 

136 raise ValueError("mome_mask is required when MOME is enabled") 

137 if mome_mask.shape != hidden_states.shape[:2]: 

138 raise ValueError( 

139 f"mome_mask must have shape {hidden_states.shape[:2]}, " 

140 f"but got {tuple(mome_mask.shape)}" 

141 ) 

142 

143 mome_mask = mome_mask.to(device=hidden_states.device, dtype=torch.bool) 

144 padding = convolution.kernel_size[0] - 1 

145 padded_states = F.pad(hidden_states, (0, 0, padding, 0)) 

146 if fused: 

147 padded_mask = F.pad(mome_mask, (padding, 0), value=False) 

148 weight = convolution.weight.squeeze(1).transpose(0, 1) 

149 mixed_states = aggregate_hidden( # pylint: disable=not-callable 

150 padded_states.transpose(0, 1).contiguous(), 

151 weight, 

152 padded_mask, 

153 ) 

154 mixed_states = mixed_states[padding:].transpose(0, 1).contiguous() 

155 else: 

156 mixed_states = convolution(padded_states.transpose(1, 2)).transpose(1, 2) 

157 mixed_states = mixed_states * mome_mask.unsqueeze(-1).to(mixed_states.dtype) 

158 return hidden_states + mixed_states 

159 

160 

161def _apply_attention_rope( 

162 query: torch.Tensor, 

163 key: torch.Tensor, 

164 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None, 

165 *, 

166 interleaved: bool, 

167) -> tuple[torch.Tensor, torch.Tensor]: 

168 """Apply the configured RoPE convention to an attention query/key pair.""" 

169 if position_embeddings is None: 

170 return query, key 

171 cos, sin = position_embeddings 

172 rotary_op = apply_rotary_pos_emb_interleave if interleaved else apply_rotary_pos_emb 

173 return rotary_op(query, key, cos, sin, unsqueeze_dim=2) 

174 

175 

176def _apply_index_rope( 

177 query: torch.Tensor, 

178 key: torch.Tensor, 

179 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None, 

180 *, 

181 rope_head_dim: int, 

182 index_head_dim: int, 

183 interleaved: bool, 

184) -> tuple[torch.Tensor, torch.Tensor]: 

185 """Apply RoPE to the rotary portion of an indexer query/key pair.""" 

186 if position_embeddings is None: 

187 return query, key 

188 cos, sin = position_embeddings 

189 split_sizes = (rope_head_dim, index_head_dim - rope_head_dim) 

190 query_rot, query_pass = torch.split(query, split_sizes, dim=-1) 

191 key_rot, key_pass = torch.split(key, split_sizes, dim=-1) 

192 rotary_op = apply_rotary_pos_emb_interleave if interleaved else apply_rotary_pos_emb 

193 query_rot, key_rot = rotary_op(query_rot, key_rot, cos, sin, unsqueeze_dim=2) 

194 return torch.cat((query_rot, query_pass), dim=-1), torch.cat((key_rot, key_pass), dim=-1) 

195 

196 

197def _absorb_query_projection( 

198 query: torch.Tensor, 

199 key_up_weight: torch.Tensor, 

200 *, 

201 num_heads: int, 

202 batch_size: int, 

203 seq_length: int, 

204 input_dim: int, 

205 output_dim: int, 

206) -> torch.Tensor: 

207 """Absorb the key up-projection weight into query states.""" 

208 query = query.permute(2, 0, 1, 3).reshape(num_heads, batch_size * seq_length, input_dim) 

209 query = torch.bmm(query, key_up_weight).view(num_heads, batch_size, seq_length, output_dim) 

210 return query.permute(1, 2, 0, 3) 

211 

212 

213def _restore_attention_projection( 

214 output: torch.Tensor, 

215 value_up_weight: torch.Tensor, 

216 *, 

217 num_heads: int, 

218 batch_size: int, 

219 seq_length: int, 

220 kv_lora_rank: int, 

221 value_head_dim: int, 

222) -> torch.Tensor: 

223 """Apply the absorbed value up-projection and restore the hidden layout.""" 

224 output = output[..., :kv_lora_rank] 

225 output = output.permute(2, 0, 1, 3).reshape(num_heads, batch_size * seq_length, kv_lora_rank) 

226 output = torch.bmm(output, value_up_weight) 

227 return output.view(num_heads, batch_size, seq_length, value_head_dim).permute(1, 2, 0, 3).reshape( 

228 batch_size, seq_length, -1 

229 ) 

230 

231 

232@module_replacement 

233class DeepseekV32DSAAttention(nn.Module): 

234 """NPU DSA replacement for Transformers ``DeepseekV32Attention``. 

235 

236 It supports causal training and packed sequences. Packed boundaries must 

237 be supplied through ``actual_seq_len``. The optional ``attention_mask`` is 

238 assumed to be the standard causal mask created by Transformers; padding, 

239 custom attention masks, dropout, and KV cache are not supported. 

240 """ 

241 

242 @staticmethod 

243 def _validate_source(module: nn.Module) -> None: 

244 """Validate the required DSA module and indexer surface.""" 

245 required = ( 

246 "q_a_proj", "q_a_layernorm", "q_b_proj", "kv_a_proj_with_mqa", 

247 "kv_a_layernorm", "kv_b_proj", "o_proj", "indexer", 

248 ) 

249 missing = [name for name in required if not hasattr(module, name)] 

250 if missing: 

251 raise TypeError(f"DeepseekV32DSAAttention source module is missing: {missing}") 

252 indexer_required = ("wq_b", "wk", "k_norm", "weights_proj") 

253 indexer_missing = [name for name in indexer_required if not hasattr(module.indexer, name)] 

254 if indexer_missing: 

255 raise TypeError(f"DeepseekV32DSAAttention source indexer is missing: {indexer_missing}") 

256 

257 def _initialize_dimensions(self, module: nn.Module) -> None: 

258 """Copy main-attention and indexer dimensions from the source module.""" 

259 config = self.config 

260 self.layer_idx = getattr(module, "layer_idx", getattr(module, "layer_number", None)) 

261 self.num_heads = getattr(module, "num_heads", config.num_attention_heads) 

262 self.q_lora_rank = getattr(module, "q_lora_rank", config.q_lora_rank) 

263 self.kv_lora_rank = getattr(module, "kv_lora_rank", config.kv_lora_rank) 

264 self.qk_rope_head_dim = getattr(module, "qk_rope_head_dim", config.qk_rope_head_dim) 

265 self.qk_nope_head_dim = getattr(module, "qk_nope_head_dim", config.qk_nope_head_dim) 

266 self.v_head_dim = getattr(module, "v_head_dim", config.v_head_dim) 

267 self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim 

268 if self.q_lora_rank is None: 

269 raise ValueError("DeepseekV32DSAAttention requires the Q-LoRA projection path") 

270 self.index_head_dim = getattr(module.indexer, "head_dim", config.index_head_dim) 

271 self.num_index_heads = getattr( 

272 module.indexer, 

273 "n_heads", 

274 getattr(config, "index_n_heads", getattr(config, "index_num_attention_heads", None)), 

275 ) 

276 if self.num_index_heads is None: 

277 raise ValueError("DeepseekV32DSAAttention requires the number of indexer heads") 

278 self.index_topk = getattr(module.indexer, "index_topk", config.index_topk) 

279 

280 def _initialize_options(self, module: nn.Module) -> None: 

281 """Copy supported DSA options and reject unsupported source modes.""" 

282 config = self.config 

283 self.scaling = getattr(module, "scaling", self.qk_head_dim**-0.5) 

284 attention_dropout = getattr(module, "attention_dropout", getattr(config, "attention_dropout", 0.0)) 

285 if isinstance(attention_dropout, nn.Dropout): 

286 attention_dropout = attention_dropout.p 

287 if attention_dropout != 0.0: 

288 raise ValueError( 

289 "DeepseekV32DSAAttention requires attention_dropout=0 because " 

290 "the sparse-attention operator does not apply dropout" 

291 ) 

292 source_rotary_interleaved = getattr(module, "rotary_interleaved", None) 

293 if source_rotary_interleaved is None: 

294 source_rotary_interleaved = getattr(config, "rope_interleave", None) 

295 if source_rotary_interleaved is None: 

296 source_rotary_interleaved = getattr(config, "rope_interleaved", None) 

297 self.rotary_interleaved = bool(True if source_rotary_interleaved is None else source_rotary_interleaved) 

298 self.index_rotary_interleaved = bool( 

299 getattr(module.indexer, "rotary_interleaved", getattr(config, "index_rope_interleaved", False)) 

300 ) 

301 self.dsa_loss_coeff = getattr(config, "dsa_loss_coeff", 0.0) 

302 self.freeze_dsa = bool(getattr(config, "freeze_DSA", getattr(config, "freeze_dsa", False))) 

303 self.dsa_dense_warm_up = bool( 

304 getattr(module, "dsa_dense_warm_up", getattr(config, "dsa_dense_warm_up", False)) 

305 ) 

306 if getattr(module, "param_sink_number", 0) > 0: 

307 raise ValueError("DeepseekV32DSAAttention does not support the parameter-sink path") 

308 if getattr(module, "use_mome", False): 

309 raise ValueError("DeepseekV32DSAAttention does not support the MOME path") 

310 

311 def _initialize_source_layers(self, module: nn.Module) -> None: 

312 """Fuse latent projections and attach the remaining source layers.""" 

313 q_a = module.q_a_proj 

314 kv_a = module.kv_a_proj_with_mqa 

315 self._q_latent_output_size = q_a.out_features 

316 self._kv_latent_output_size = kv_a.out_features 

317 can_fuse = ( 

318 q_a.in_features == kv_a.in_features 

319 and q_a.weight.requires_grad == kv_a.weight.requires_grad 

320 and (q_a.bias is None) == (kv_a.bias is None) 

321 and (q_a.bias is None or q_a.bias.requires_grad == kv_a.bias.requires_grad) 

322 ) 

323 if not can_fuse: 

324 raise ValueError("Q and KV latent projections cannot be represented by one fused projection") 

325 self.linear_qkv = nn.Linear( 

326 q_a.in_features, 

327 q_a.out_features + kv_a.out_features, 

328 bias=q_a.bias is not None, 

329 device=q_a.weight.device, 

330 dtype=q_a.weight.dtype, 

331 ) 

332 self.linear_qkv.weight.requires_grad_(q_a.weight.requires_grad) 

333 if q_a.bias is not None: 

334 self.linear_qkv.bias.requires_grad_(q_a.bias.requires_grad) 

335 for name in ("q_a_layernorm", "kv_a_layernorm", "q_b_proj", "kv_b_proj", "o_proj", "indexer"): 

336 setattr(self, name, getattr(module, name)) 

337 

338 def __init__( 

339 self, 

340 *, 

341 module: nn.Module, 

342 module_fqn: str = "", 

343 context: Mapping[str, Any] | None = None, 

344 ) -> None: 

345 """Build the high-performance module from a Transformers DSA module.""" 

346 super().__init__() 

347 del module_fqn, context 

348 self._validate_source(module) 

349 config = module.config 

350 self.config = config 

351 if hasattr(config, "indexer_types"): 

352 raise TypeError( 

353 "DeepseekV32DSAAttention supports the DeepSeek-V3.2 DSA contract; " 

354 "GLM-MOE-DSA shared indexers and three-value forward result require " 

355 "a dedicated replacement" 

356 ) 

357 self._initialize_dimensions(module) 

358 self._initialize_options(module) 

359 self._initialize_source_layers(module) 

360 self.train(module.training) 

361 

362 def make_transforms(self) -> list[WeightConverter]: 

363 """Describe reversible source-checkpoint to high-performance conversion.""" 

364 transforms: list[WeightConverter] = [ 

365 WeightConverter( 

366 source_patterns=[ 

367 "q_a_proj.weight", 

368 "kv_a_proj_with_mqa.weight", 

369 ], 

370 target_patterns="linear_qkv.weight", 

371 operations=[ 

372 ConcatenateWithSections( 

373 sections=( 

374 self._q_latent_output_size, 

375 self._kv_latent_output_size, 

376 ), 

377 dim=0, 

378 ) 

379 ], 

380 ), 

381 ] 

382 if self.linear_qkv.bias is not None: 

383 transforms.insert( 

384 1, 

385 WeightConverter( 

386 source_patterns=["q_a_proj.bias", "kv_a_proj_with_mqa.bias"], 

387 target_patterns="linear_qkv.bias", 

388 operations=[ 

389 ConcatenateWithSections( 

390 sections=( 

391 self._q_latent_output_size, 

392 self._kv_latent_output_size, 

393 ), 

394 dim=0, 

395 ) 

396 ], 

397 ), 

398 ) 

399 return transforms 

400 

401 @staticmethod 

402 def _get_actual_seq_len( 

403 actual_seq_len: torch.Tensor | Sequence[int] | None, 

404 batch_size: int, 

405 seq_length: int, 

406 device: torch.device, 

407 ) -> torch.Tensor: 

408 """Normalize cumulative sequence lengths for the TND custom operators.""" 

409 if isinstance(actual_seq_len, torch.Tensor): 

410 return actual_seq_len.to(device=device, dtype=torch.int32) 

411 if actual_seq_len is not None: 

412 return torch.tensor(actual_seq_len, device=device, dtype=torch.int32) 

413 return torch.arange( 

414 seq_length, 

415 (batch_size + 1) * seq_length, 

416 seq_length, 

417 dtype=torch.int32, 

418 device=device, 

419 ) 

420 

421 @staticmethod 

422 def _invalid_position_ids( 

423 position_ids: torch.Tensor | None, 

424 batch_size: int, 

425 seq_length: int, 

426 ) -> bool: 

427 """Return whether position IDs violate the supported dense layout.""" 

428 if position_ids is None: 

429 return False 

430 if position_ids.ndim != 2: 

431 return True 

432 return position_ids.shape[0] not in (1, batch_size) or position_ids.shape[1] != seq_length 

433 

434 @staticmethod 

435 def _validate_forward_inputs( 

436 hidden_states: torch.Tensor, 

437 attention_mask: torch.Tensor | None, 

438 past_key_values: Any | None, 

439 position_ids: torch.Tensor | None, 

440 kwargs: dict[str, Any], 

441 ) -> None: 

442 """Validate inputs supported by the DeepSeek sparse-attention path.""" 

443 batch_size, seq_length = hidden_states.shape[:-1] 

444 if past_key_values is not None: 

445 raise NotImplementedError("DeepseekV32DSAAttention does not support KV cache") 

446 if DeepseekV32DSAAttention._invalid_position_ids(position_ids, batch_size, seq_length): 

447 raise ValueError("position_ids must have shape [1 or batch_size, sequence_length]") 

448 if attention_mask is not None: 

449 expected_shape = (batch_size, 1, seq_length, seq_length) 

450 if tuple(attention_mask.shape) != expected_shape: 

451 raise ValueError( 

452 "DeepseekV32DSAAttention supports only a Transformers 4D causal mask; " 

453 "padding and custom masks are not supported" 

454 ) 

455 if kwargs.pop("output_attentions", False): 

456 raise NotImplementedError("DeepseekV32DSAAttention does not expose attention weights") 

457 

458 def _project_attention_states( 

459 self, 

460 hidden_states: torch.Tensor, 

461 *, 

462 use_dense: bool = False, 

463 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: 

464 """Project MLA Q/KV, retaining Q before absorption during dense warm-up.""" 

465 batch_size, seq_length = hidden_states.shape[:-1] 

466 latent_states = self.linear_qkv(hidden_states) 

467 q_latent, kv_nope, k_rot = torch.split( 

468 latent_states, (self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim), dim=-1 

469 ) 

470 q_resid = self.q_a_layernorm(q_latent) 

471 q_states = self.q_b_proj(q_resid).view(batch_size, seq_length, self.num_heads, self.qk_head_dim) 

472 q_pass, q_rot = torch.split(q_states, (self.qk_nope_head_dim, self.qk_rope_head_dim), dim=-1) 

473 kv_weight = self.kv_b_proj.weight.view( 

474 self.num_heads, self.qk_nope_head_dim + self.v_head_dim, self.kv_lora_rank 

475 ) 

476 absorbed_query = q_pass if use_dense else _absorb_query_projection( 

477 q_pass, 

478 kv_weight[:, : self.qk_nope_head_dim], 

479 num_heads=self.num_heads, 

480 batch_size=batch_size, 

481 seq_length=seq_length, 

482 input_dim=self.qk_nope_head_dim, 

483 output_dim=self.kv_lora_rank, 

484 ) 

485 kv_nope = self.kv_a_layernorm(kv_nope).view(batch_size, seq_length, 1, self.kv_lora_rank) 

486 k_rot = k_rot.view(batch_size, seq_length, 1, self.qk_rope_head_dim) 

487 return q_resid, absorbed_query, kv_nope, q_rot, k_rot, kv_weight 

488 

489 def _project_index_states( 

490 self, 

491 hidden_states: torch.Tensor, 

492 q_resid: torch.Tensor, 

493 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None, 

494 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 

495 """Build indexer query, key, and merge-weight states.""" 

496 batch_size, seq_length = hidden_states.shape[:-1] 

497 index_query = self.indexer.wq_b(q_resid.detach()).view( 

498 batch_size, seq_length, self.num_index_heads, self.index_head_dim 

499 ) 

500 index_key = self.indexer.k_norm(self.indexer.wk(hidden_states.detach())).unsqueeze(2) 

501 merge_weight = self.indexer.weights_proj(hidden_states.detach()) 

502 merge_weight = merge_weight * self.num_index_heads**-0.5 * self.index_head_dim**-0.5 

503 index_query, index_key = _apply_index_rope( 

504 index_query, 

505 index_key, 

506 position_embeddings, 

507 rope_head_dim=self.qk_rope_head_dim, 

508 index_head_dim=self.index_head_dim, 

509 interleaved=self.index_rotary_interleaved, 

510 ) 

511 return index_query, index_key, merge_weight 

512 

513 def _apply_auxiliary_loss( 

514 self, 

515 attn_output: torch.Tensor, 

516 index_states: tuple[torch.Tensor, torch.Tensor, torch.Tensor], 

517 attention_states: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], 

518 topk_indices: torch.Tensor, 

519 softmax_stats: tuple[torch.Tensor, torch.Tensor], 

520 actual_q_len: torch.Tensor, 

521 actual_kv_len: torch.Tensor, 

522 ) -> torch.Tensor: 

523 """Attach the DSA KL auxiliary loss when training enables it.""" 

524 if not (self.training and not self.freeze_dsa and self.dsa_loss_coeff): 

525 return attn_output 

526 query_tnd, key_tnd, q_rot_tnd, k_rot_tnd = ( 

527 tensor.reshape(-1, tensor.shape[2], tensor.shape[3]) 

528 for tensor in attention_states 

529 ) 

530 aux_loss = dsa_kl_loss( 

531 *index_states, query_tnd, key_tnd, 

532 topk_indices, *softmax_stats, q_rot_tnd, k_rot_tnd, 

533 actual_q_len, actual_kv_len, self.scaling, self.dsa_loss_coeff, 

534 backend="cann", 

535 ) 

536 return aux_loss_auto_scale(attn_output, aux_loss) 

537 

538 def _prepare_attention_states(self, hidden_states, position_embeddings): 

539 """Project attention states and apply their rotary embeddings.""" 

540 q_resid, absorbed_query_states, kv_nope, q_rot, k_rot, kv_weight = self._project_attention_states( 

541 hidden_states, use_dense=self.dsa_dense_warm_up, 

542 ) 

543 q_rot, k_rot = _apply_attention_rope( 

544 q_rot, k_rot, position_embeddings, interleaved=self.rotary_interleaved 

545 ) 

546 return q_resid, (absorbed_query_states, kv_nope, q_rot, k_rot), kv_weight 

547 

548 def _run_indexer(self, hidden_states, q_resid, position_embeddings, actual_seq_len, kwargs): 

549 """Resolve packed boundaries and select sparse attention indices.""" 

550 index_states = self._project_index_states(hidden_states, q_resid, position_embeddings) 

551 packed_kwargs = dict(kwargs) 

552 packed_kwargs["actual_seq_len"] = actual_seq_len 

553 actual_q_len, actual_kv_len = resolve_packed_sequence_lengths( 

554 packed_kwargs, 

555 hidden_states.shape[0] * hidden_states.shape[1], 

556 hidden_states.shape[0] * hidden_states.shape[1], 

557 ) 

558 actual_q_len = self._get_actual_seq_len( 

559 actual_q_len, *hidden_states.shape[:-1], hidden_states.device, 

560 ) 

561 actual_kv_len = self._get_actual_seq_len( 

562 actual_kv_len, *hidden_states.shape[:-1], hidden_states.device, 

563 ) 

564 indexed = dsa_indexer( 

565 *index_states, actual_q_len, actual_kv_len, self.index_topk, backend="cann", 

566 ) 

567 return indexed[0], indexed[1:], actual_q_len, actual_kv_len 

568 

569 def _compute_sparse_attention( 

570 self, hidden_states, q_resid, attention_states, position_embeddings, actual_seq_len, kwargs, 

571 ): 

572 """Run sparse attention and attach its optional auxiliary loss.""" 

573 topk_indices, index_states, actual_q_len, actual_kv_len = self._run_indexer( 

574 hidden_states, q_resid, position_embeddings, actual_seq_len, kwargs, 

575 ) 

576 attn_output, softmax_max, softmax_sum = dsa_sparse_attention( 

577 *attention_states, topk_indices, self.scaling, actual_q_len, actual_kv_len, 

578 backend="cann", 

579 ) 

580 return self._apply_auxiliary_loss( 

581 attn_output, index_states, attention_states, topk_indices, 

582 (softmax_max, softmax_sum), actual_q_len, actual_kv_len, 

583 ) 

584 

585 def _compute_attention_output( 

586 self, hidden_states, position_embeddings, actual_seq_len, kwargs, 

587 ): 

588 """Run the selected dense/sparse stage and restore the value projection.""" 

589 batch_size, seq_length = hidden_states.shape[:-1] 

590 q_resid, attention_states, kv_weight = self._prepare_attention_states( 

591 hidden_states, position_embeddings, 

592 ) 

593 if self.dsa_dense_warm_up: 

594 index_states = self._project_index_states(hidden_states, q_resid, position_embeddings) 

595 packed_kwargs = {**kwargs, "actual_seq_len": actual_seq_len} 

596 lengths = resolve_packed_sequence_lengths( 

597 packed_kwargs, batch_size * seq_length, batch_size * seq_length, 

598 ) 

599 lengths = tuple(self._get_actual_seq_len( 

600 value, batch_size, seq_length, hidden_states.device, 

601 ) for value in lengths) 

602 output = _deepseek_v32_dense_attention(self, index_states, attention_states, lengths) 

603 return output.flatten(2) 

604 attn_output = self._compute_sparse_attention( 

605 hidden_states, q_resid, attention_states, position_embeddings, actual_seq_len, kwargs, 

606 ) 

607 value_up_weight = kv_weight[:, self.qk_nope_head_dim :].transpose(1, 2) 

608 return _restore_attention_projection( 

609 attn_output, 

610 value_up_weight, 

611 num_heads=self.num_heads, 

612 batch_size=batch_size, 

613 seq_length=seq_length, 

614 kv_lora_rank=self.kv_lora_rank, 

615 value_head_dim=self.v_head_dim, 

616 ) 

617 

618 def forward( 

619 self, 

620 hidden_states: torch.Tensor, 

621 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, 

622 attention_mask: torch.Tensor | None = None, 

623 past_key_values: Any | None = None, 

624 position_ids: torch.Tensor | None = None, 

625 actual_seq_len: torch.Tensor | Sequence[int] | None = None, 

626 **kwargs: Any, 

627 ) -> tuple[torch.Tensor, None]: 

628 """Run causal or packed DSA with the NPU sparse-attention kernels.""" 

629 self._validate_forward_inputs(hidden_states, attention_mask, past_key_values, position_ids, kwargs) 

630 attn_output = self._compute_attention_output( 

631 hidden_states, position_embeddings, actual_seq_len, kwargs, 

632 ) 

633 attn_output = self.o_proj(attn_output) 

634 return attn_output, None 

635 

636 

637@module_replacement 

638class DSAAttention(nn.Module): 

639 """NPU-optimized replacement for a DSA attention module. 

640 

641 Packed boundaries are supplied through ``actual_seq_len``. Attention masks, 

642 KV reuse, and KV cache are not supported. The source parameter layout is 

643 preserved, so no checkpoint conversion is required. 

644 """ 

645 

646 def __init__( 

647 self, 

648 *, 

649 module: nn.Module, 

650 module_fqn: str = "", 

651 context: Mapping[str, Any] | None = None, 

652 ) -> None: 

653 """Build the replacement from a compatible DSA attention module.""" 

654 super().__init__() 

655 del module_fqn, context 

656 required = ( 

657 "linear_qkv", 

658 "q_layernorm", 

659 "k_layernorm", 

660 "linear_qb", 

661 "linear_kvb", 

662 "linear_proj", 

663 "index_linear_qb", 

664 "index_linear_k", 

665 "index_k_layernorm", 

666 "linear_merge_weight", 

667 ) 

668 missing = [name for name in required if not hasattr(module, name)] 

669 if missing: 

670 raise TypeError(f"DSAAttention source module is missing: {missing}") 

671 

672 unsupported = { 

673 "mla_mm_split": bool(getattr(module, "mla_mm_split", False)), 

674 "dense warm-up": bool(getattr(module, "dsa_dense_warm_up", False)), 

675 } 

676 enabled = [name for name, value in unsupported.items() if value] 

677 if enabled: 

678 raise NotImplementedError( 

679 f"DSAAttention does not yet support: {', '.join(enabled)}" 

680 ) 

681 

682 self.config = module.config 

683 self.layer_number = module.layer_number 

684 self.num_heads = module.num_heads 

685 self.q_lora_rank = module.q_lora_rank 

686 self.kv_lora_rank = module.kv_lora_rank 

687 self.qk_rope_head_dim = module.qk_rope_head_dim 

688 self.qk_nope_head_dim = module.qk_nope_head_dim 

689 self.qk_head_dim = module.qk_head_dim 

690 self.v_head_dim = module.v_head_dim 

691 self.index_head_dim = module.index_head_dim 

692 self.num_index_heads = module.num_index_heads 

693 self.index_topk = module.index_topk 

694 self.rotary_interleaved = module.rotary_interleaved 

695 self.dsa_loss_coeff = module.dsa_loss_coeff 

696 self.freeze_dsa = module.freeze_dsa 

697 self.use_flash_attn = bool(getattr(module, "use_flash_attn", True))