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 apply_mome( 

43 hidden_states: torch.Tensor, 

44 mome_mask: torch.Tensor | None, 

45 convolution: nn.Conv1d, 

46 *, 

47 fused: bool, 

48) -> torch.Tensor: 

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

50 if mome_mask is None: 

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

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

53 raise ValueError( 

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

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

56 ) 

57 

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

59 padding = convolution.kernel_size[0] - 1 

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

61 if fused: 

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

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

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

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

66 weight, 

67 padded_mask, 

68 ) 

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

70 else: 

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

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

73 return hidden_states + mixed_states 

74 

75 

76def _apply_attention_rope( 

77 query: torch.Tensor, 

78 key: torch.Tensor, 

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

80 *, 

81 interleaved: bool, 

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

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

84 if position_embeddings is None: 

85 return query, key 

86 cos, sin = position_embeddings 

87 rotary_op = apply_rotary_pos_emb_interleave if interleaved else apply_rotary_pos_emb 

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

89 

90 

91def _apply_index_rope( 

92 query: torch.Tensor, 

93 key: torch.Tensor, 

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

95 *, 

96 rope_head_dim: int, 

97 index_head_dim: int, 

98 interleaved: bool, 

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

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

101 if position_embeddings is None: 

102 return query, key 

103 cos, sin = position_embeddings 

104 split_sizes = (rope_head_dim, index_head_dim - rope_head_dim) 

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

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

107 rotary_op = apply_rotary_pos_emb_interleave if interleaved else apply_rotary_pos_emb 

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

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

110 

111 

112def _absorb_query_projection( 

113 query: torch.Tensor, 

114 key_up_weight: torch.Tensor, 

115 *, 

116 num_heads: int, 

117 batch_size: int, 

118 seq_length: int, 

119 input_dim: int, 

120 output_dim: int, 

121) -> torch.Tensor: 

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

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

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

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

126 

127 

128def _restore_attention_projection( 

129 output: torch.Tensor, 

130 value_up_weight: torch.Tensor, 

131 *, 

132 num_heads: int, 

133 batch_size: int, 

134 seq_length: int, 

135 kv_lora_rank: int, 

136 value_head_dim: int, 

137) -> torch.Tensor: 

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

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

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

141 output = torch.bmm(output, value_up_weight) 

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

143 batch_size, seq_length, -1 

144 ) 

145 

146 

147@module_replacement 

148class DeepseekV32DSAAttention(nn.Module): 

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

150 

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

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

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

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

155 """ 

156 

157 @staticmethod 

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

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

160 required = ( 

161 "q_a_proj", "q_a_layernorm", "q_b_proj", "kv_a_proj_with_mqa", 

162 "kv_a_layernorm", "kv_b_proj", "o_proj", "indexer", 

163 ) 

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

165 if missing: 

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

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

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

169 if indexer_missing: 

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

171 

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

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

174 config = self.config 

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

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

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

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

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

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

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

182 self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim 

183 if self.q_lora_rank is None: 

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

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

186 self.num_index_heads = getattr( 

187 module.indexer, 

188 "n_heads", 

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

190 ) 

191 if self.num_index_heads is None: 

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

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

194 

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

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

197 config = self.config 

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

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

200 if isinstance(attention_dropout, nn.Dropout): 

201 attention_dropout = attention_dropout.p 

202 if attention_dropout != 0.0: 

203 raise ValueError( 

204 "DeepseekV32DSAAttention requires attention_dropout=0 because " 

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

206 ) 

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

208 if source_rotary_interleaved is None: 

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

210 if source_rotary_interleaved is None: 

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

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

213 self.index_rotary_interleaved = bool( 

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

215 ) 

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

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

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

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

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

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

222 

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

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

225 q_a = module.q_a_proj 

226 kv_a = module.kv_a_proj_with_mqa 

227 self._q_latent_output_size = q_a.out_features 

228 self._kv_latent_output_size = kv_a.out_features 

229 can_fuse = ( 

230 q_a.in_features == kv_a.in_features 

231 and q_a.weight.requires_grad == kv_a.weight.requires_grad 

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

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

234 ) 

235 if not can_fuse: 

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

237 self.linear_qkv = nn.Linear( 

238 q_a.in_features, 

239 q_a.out_features + kv_a.out_features, 

240 bias=q_a.bias is not None, 

241 device=q_a.weight.device, 

242 dtype=q_a.weight.dtype, 

243 ) 

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

245 if q_a.bias is not None: 

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

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

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

249 

250 def __init__( 

251 self, 

252 *, 

253 module: nn.Module, 

254 module_fqn: str = "", 

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

256 ) -> None: 

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

258 super().__init__() 

259 del module_fqn, context 

260 self._validate_source(module) 

261 config = module.config 

262 self.config = config 

263 if hasattr(config, "indexer_types"): 

264 raise TypeError( 

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

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

267 "a dedicated replacement" 

268 ) 

269 self._initialize_dimensions(module) 

270 self._initialize_options(module) 

271 self._initialize_source_layers(module) 

272 self.train(module.training) 

273 

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

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

276 transforms: list[WeightConverter] = [ 

277 WeightConverter( 

278 source_patterns=[ 

279 "q_a_proj.weight", 

280 "kv_a_proj_with_mqa.weight", 

281 ], 

282 target_patterns="linear_qkv.weight", 

283 operations=[ 

284 ConcatenateWithSections( 

285 sections=( 

286 self._q_latent_output_size, 

287 self._kv_latent_output_size, 

288 ), 

289 dim=0, 

290 ) 

291 ], 

292 ), 

293 ] 

294 if self.linear_qkv.bias is not None: 

295 transforms.insert( 

296 1, 

297 WeightConverter( 

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

299 target_patterns="linear_qkv.bias", 

300 operations=[ 

301 ConcatenateWithSections( 

302 sections=( 

303 self._q_latent_output_size, 

304 self._kv_latent_output_size, 

305 ), 

306 dim=0, 

307 ) 

308 ], 

309 ), 

310 ) 

311 return transforms 

312 

313 @staticmethod 

314 def _get_actual_seq_len( 

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

316 batch_size: int, 

317 seq_length: int, 

318 device: torch.device, 

319 ) -> torch.Tensor: 

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

321 if isinstance(actual_seq_len, torch.Tensor): 

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

323 if actual_seq_len is not None: 

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

325 return torch.arange( 

326 seq_length, 

327 (batch_size + 1) * seq_length, 

328 seq_length, 

329 dtype=torch.int32, 

330 device=device, 

331 ) 

332 

333 @staticmethod 

334 def _invalid_position_ids( 

335 position_ids: torch.Tensor | None, 

336 batch_size: int, 

337 seq_length: int, 

338 ) -> bool: 

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

340 if position_ids is None: 

341 return False 

342 if position_ids.ndim != 2: 

343 return True 

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

345 

346 @staticmethod 

347 def _validate_forward_inputs( 

348 hidden_states: torch.Tensor, 

349 attention_mask: torch.Tensor | None, 

350 past_key_values: Any | None, 

351 position_ids: torch.Tensor | None, 

352 kwargs: dict[str, Any], 

353 ) -> None: 

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

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

356 if past_key_values is not None: 

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

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

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

360 if attention_mask is not None: 

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

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

363 raise ValueError( 

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

365 "padding and custom masks are not supported" 

366 ) 

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

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

369 

370 def _project_attention_states( 

371 self, 

372 hidden_states: torch.Tensor, 

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

374 """Project hidden states into absorbed-query and latent KV states.""" 

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

376 latent_states = self.linear_qkv(hidden_states) 

377 q_latent, kv_nope, k_rot = torch.split( 

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

379 ) 

380 q_resid = self.q_a_layernorm(q_latent) 

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

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

383 kv_weight = self.kv_b_proj.weight.view( 

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

385 ) 

386 absorbed_query = _absorb_query_projection( 

387 q_pass, 

388 kv_weight[:, : self.qk_nope_head_dim], 

389 num_heads=self.num_heads, 

390 batch_size=batch_size, 

391 seq_length=seq_length, 

392 input_dim=self.qk_nope_head_dim, 

393 output_dim=self.kv_lora_rank, 

394 ) 

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

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

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

398 

399 def _project_index_states( 

400 self, 

401 hidden_states: torch.Tensor, 

402 q_resid: torch.Tensor, 

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

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

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

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

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

408 batch_size, seq_length, self.num_index_heads, self.index_head_dim 

409 ) 

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

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

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

413 index_query, index_key = _apply_index_rope( 

414 index_query, 

415 index_key, 

416 position_embeddings, 

417 rope_head_dim=self.qk_rope_head_dim, 

418 index_head_dim=self.index_head_dim, 

419 interleaved=self.index_rotary_interleaved, 

420 ) 

421 return index_query, index_key, merge_weight 

422 

423 def _apply_auxiliary_loss( 

424 self, 

425 attn_output: torch.Tensor, 

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

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

428 topk_indices: torch.Tensor, 

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

430 actual_q_len: torch.Tensor, 

431 actual_kv_len: torch.Tensor, 

432 ) -> torch.Tensor: 

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

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

435 return attn_output 

436 query_tnd, key_tnd, q_rot_tnd, k_rot_tnd = ( 

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

438 for tensor in attention_states 

439 ) 

440 aux_loss = dsa_kl_loss( # pylint: disable=not-callable 

441 *index_states, query_tnd, key_tnd, 

442 topk_indices, *softmax_stats, q_rot_tnd, k_rot_tnd, 

443 actual_q_len, actual_kv_len, self.scaling, self.dsa_loss_coeff, 

444 ) 

445 return aux_loss_auto_scale(attn_output, aux_loss) 

446 

447 def _prepare_attention_states(self, hidden_states, position_embeddings): 

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

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

450 hidden_states 

451 ) 

452 q_rot, k_rot = _apply_attention_rope( 

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

454 ) 

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

456 

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

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

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

460 packed_kwargs = dict(kwargs) 

461 packed_kwargs["actual_seq_len"] = actual_seq_len 

462 actual_q_len, actual_kv_len = resolve_packed_sequence_lengths( 

463 packed_kwargs, 

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

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

466 ) 

467 actual_q_len = self._get_actual_seq_len( 

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

469 ) 

470 actual_kv_len = self._get_actual_seq_len( 

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

472 ) 

473 indexed = dsa_indexer( # pylint: disable=not-callable 

474 *index_states, actual_q_len, actual_kv_len, self.index_topk 

475 ) 

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

477 

478 def _compute_sparse_attention( 

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

480 ): 

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

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

483 hidden_states, q_resid, position_embeddings, actual_seq_len, kwargs, 

484 ) 

485 attn_output, softmax_max, softmax_sum = dsa_sparse_attention( # pylint: disable=not-callable 

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

487 ) 

488 return self._apply_auxiliary_loss( 

489 attn_output, index_states, attention_states, topk_indices, 

490 (softmax_max, softmax_sum), actual_q_len, actual_kv_len, 

491 ) 

492 

493 def _compute_attention_output( 

494 self, hidden_states, position_embeddings, actual_seq_len, kwargs, 

495 ): 

496 """Compute sparse attention and restore its value projection.""" 

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

498 q_resid, attention_states, kv_weight = self._prepare_attention_states( 

499 hidden_states, position_embeddings, 

500 ) 

501 attn_output = self._compute_sparse_attention( 

502 hidden_states, q_resid, attention_states, position_embeddings, actual_seq_len, kwargs, 

503 ) 

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

505 return _restore_attention_projection( 

506 attn_output, 

507 value_up_weight, 

508 num_heads=self.num_heads, 

509 batch_size=batch_size, 

510 seq_length=seq_length, 

511 kv_lora_rank=self.kv_lora_rank, 

512 value_head_dim=self.v_head_dim, 

513 ) 

514 

515 def forward( 

516 self, 

517 hidden_states: torch.Tensor, 

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

519 attention_mask: torch.Tensor | None = None, 

520 past_key_values: Any | None = None, 

521 position_ids: torch.Tensor | None = None, 

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

523 **kwargs: Any, 

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

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

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

527 attn_output = self._compute_attention_output( 

528 hidden_states, position_embeddings, actual_seq_len, kwargs, 

529 ) 

530 attn_output = self.o_proj(attn_output) 

531 return attn_output, None 

532 

533 

534@module_replacement 

535class DSAAttention(nn.Module): 

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

537 

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

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

540 preserved, so no checkpoint conversion is required. 

541 """ 

542 

543 def __init__( 

544 self, 

545 *, 

546 module: nn.Module, 

547 module_fqn: str = "", 

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

549 ) -> None: 

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

551 super().__init__() 

552 del module_fqn, context 

553 required = ( 

554 "linear_qkv", 

555 "q_layernorm", 

556 "k_layernorm", 

557 "linear_qb", 

558 "linear_kvb", 

559 "linear_proj", 

560 "index_linear_qb", 

561 "index_linear_k", 

562 "index_k_layernorm", 

563 "linear_merge_weight", 

564 ) 

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

566 if missing: 

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

568 

569 unsupported = { 

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

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

572 } 

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

574 if enabled: 

575 raise NotImplementedError( 

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

577 ) 

578 

579 self.config = module.config 

580 self.layer_number = module.layer_number 

581 self.num_heads = module.num_heads 

582 self.q_lora_rank = module.q_lora_rank 

583 self.kv_lora_rank = module.kv_lora_rank 

584 self.qk_rope_head_dim = module.qk_rope_head_dim 

585 self.qk_nope_head_dim = module.qk_nope_head_dim 

586 self.qk_head_dim = module.qk_head_dim 

587 self.v_head_dim = module.v_head_dim 

588 self.index_head_dim = module.index_head_dim 

589 self.num_index_heads = module.num_index_heads 

590 self.index_topk = module.index_topk 

591 self.rotary_interleaved = module.rotary_interleaved 

592 self.dsa_loss_coeff = module.dsa_loss_coeff 

593 self.freeze_dsa = module.freeze_dsa 

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

595 self.use_mome = bool(getattr(module, "use_mome", False)) 

596 self.use_fused_mome = bool(getattr(module, "use_fused_mome", False)) 

597 self.param_sink_number = int(getattr(module, "param_sink_number", 0)) 

598 self.param_sink_scalar = getattr(module, "param_sink_scalar", None) 

599 self.apply_fa_rescale = bool( 

600 getattr(module, "apply_fa_rescale", False) 

601 ) 

602 self.attention_dropout = getattr(module, "attention_dropout", nn.Dropout(0.0)) 

603 if self.param_sink_number > 0: 

604 if self.param_sink_scalar: 

605 raise NotImplementedError( 

606 "DSAAttention does not support scalar parameter sink" 

607 ) 

608 if not self.apply_fa_rescale: 

609 raise NotImplementedError( 

610 "DSAAttention supports parameter sink through " 

611 "apply_fa_rescale only" 

612 ) 

613 

614 self.linear_qkv = module.linear_qkv 

615 self.q_layernorm = module.q_layernorm 

616 self.k_layernorm = module.k_layernorm 

617 self.linear_qb = module.linear_qb 

618 self.linear_kvb = module.linear_kvb 

619 self.linear_proj = module.linear_proj 

620 self.index_linear_qb = module.index_linear_qb 

621 self.index_linear_k = module.index_linear_k 

622 self.index_k_layernorm = module.index_k_layernorm 

623 self.linear_merge_weight = module.linear_merge_weight 

624 if self.use_mome: 

625 self.qa_conv = module.qa_conv 

626 self.compresskv_conv = module.compresskv_conv 

627 self.o_conv = module.o_conv 

628 if self.param_sink_number > 0: 

629 self.param_sink_k_pe = module.param_sink_k_pe 

630 self.param_sink_compressed_kv = module.param_sink_compressed_kv 

631 self.train(module.training) 

632 

633 @staticmethod 

634 def _get_actual_seq_len( 

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

636 batch_size: int, 

637 seq_length: int, 

638 device: torch.device, 

639 ) -> torch.Tensor: 

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

641 if actual_seq_len is None: 

642 return torch.arange( 

643 seq_length, 

644 (batch_size + 1) * seq_length, 

645 seq_length, 

646 dtype=torch.int32, 

647 device=device, 

648 ) 

649 return torch.as_tensor(actual_seq_len, dtype=torch.int32, device=device) 

650 

651 def _prepare_param_sink( 

652 self, batch_size: int 

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

654 """Construct the DSA sink key and value in BSND layout.""" 

655 sink_rotary_key = self.param_sink_k_pe.unsqueeze(0).unsqueeze(2).expand( 

656 batch_size, -1, -1, -1 

657 ) 

658 sink_latent = self.param_sink_compressed_kv.unsqueeze(0).unsqueeze(2).expand( 

659 batch_size, -1, -1, -1 

660 ) 

661 sink_latent = self.k_layernorm(sink_latent) 

662 sink_value = ( 

663 F.pad(sink_latent, (0, self.qk_rope_head_dim)) 

664 if self.use_flash_attn and self.qk_rope_head_dim > 0 

665 else sink_latent 

666 ) 

667 sink_key = torch.cat((sink_latent, sink_rotary_key), dim=-1) 

668 return sink_key, sink_value 

669 

670 @staticmethod 

671 def _validate_forward_inputs( 

672 attention_mask: torch.Tensor | None, 

673 kv_reuse_states: Any | None, 

674 past_key_values: Any | None, 

675 cache_position: torch.Tensor | None, 

676 output_attentions: bool, 

677 ) -> None: 

678 """Validate inputs supported by the optimized DSA path.""" 

679 if attention_mask is not None: 

680 raise ValueError( 

681 "DSAAttention does not consume attention_mask; pass packed sequence " 

682 "boundaries through actual_seq_len" 

683 ) 

684 if kv_reuse_states is not None: 

685 raise NotImplementedError("DSAAttention does not support KV reuse") 

686 if past_key_values is not None or cache_position is not None: 

687 raise NotImplementedError("DSAAttention does not support KV cache") 

688 if output_attentions: 

689 raise NotImplementedError("DSAAttention does not expose attention weights") 

690 

691 def _project_attention_states( 

692 self, 

693 hidden_states: torch.Tensor, 

694 mome_mask: torch.Tensor | None, 

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

696 """Project hidden states into absorbed-query and latent KV states.""" 

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

698 latent_states, _ = self.linear_qkv(hidden_states) 

699 q_latent, kv_nope, k_rot = torch.split( 

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

701 ) 

702 if self.use_mome: 

703 q_latent = apply_mome(q_latent, mome_mask, self.qa_conv, fused=False) 

704 kv_nope = apply_mome(kv_nope, mome_mask, self.compresskv_conv, fused=False) 

705 q_resid = self.q_layernorm(q_latent) 

706 q_states, _ = self.linear_qb(q_resid.contiguous()) 

707 q_states = q_states.view(batch_size, seq_length, self.num_heads, self.qk_head_dim) 

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

709 kv_weight = self.linear_kvb.weight.view( 

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

711 ) 

712 absorbed_query = _absorb_query_projection( 

713 q_pass, 

714 kv_weight[:, : self.qk_nope_head_dim],