Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / mla_context_parallel.py: 65%

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"""Dense MLA Ulysses with expanded QKV or shared latent-KV communication.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Callable, Sequence 

20from dataclasses import dataclass 

21from functools import lru_cache, wraps 

22import math 

23from typing import Any, Literal 

24 

25# Native Torch is required, as in the existing CP adapters; the legacy lint rule predates that migration. 

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

27from torch import nn # pylint: disable=forbidden-backend-import 

28from torch.nn import functional as F # pylint: disable=forbidden-backend-import 

29 

30from hyper_parallel.core.dtensor.dtensor import DTensor 

31from hyper_parallel.distributed._builder.forward_rewriter import _ForwardRewriteRequest 

32from hyper_parallel.distributed.context_parallel.mla_backward_overlap import ( 

33 MLAExpandedInputA2AState, MLAInputA2AState, 

34 _InputA2ALaunch, _LatentInputA2AWait, _LatentKVLinear, 

35 finish_expanded_inputs, launch_input_a2a, 

36) 

37from hyper_parallel.distributed.context_parallel.gated_mla_backward_overlap import ( 

38 GatedMLAInputState, finish_gated_expanded_inputs, 

39) 

40from hyper_parallel.components import functional 

41from hyper_parallel.components.modules.mla_attention import prepare_mla_attention_mask 

42from hyper_parallel.distributed.context_parallel.collectives import ( 

43 _build_hybrid_cp_submeshes, async_cp_allgather_launch, async_ulysses_seq_to_head_launch, 

44 ulysses_head_to_seq, ulysses_seq_to_head, 

45) 

46from hyper_parallel.components.functional.npu_fusion_attention import ( 

47 PACKED_SEQUENCE_ARGUMENTS, npu_fusion_attention_forward, resolve_packed_sequence_lengths, 

48) 

49from hyper_parallel.core.utils.communication import ( 

50 differentiable_all_to_all_single, 

51 differentiable_all_to_all_single_async, wait_async_tensor, 

52) 

53 

54 

55MLACPStrategy = Literal["expanded_ulysses", "latent_kv_head"] 

56MLA_CP_STRATEGIES = ("expanded_ulysses", "latent_kv_head") 

57 

58 

59@dataclass(frozen=True) 

60class MLADimensions: 

61 """Global MLA dimensions; the latent ranks are not partitioned across heads.""" 

62 

63 heads: int 

64 q_rank: int 

65 kv_rank: int 

66 nope_dim: int 

67 rope_dim: int 

68 value_dim: int 

69 

70 def __post_init__(self) -> None: 

71 for name, value in vars(self).items(): 

72 if not isinstance(value, int) or isinstance(value, bool) or value <= 0: 

73 raise ValueError(f"{name} must be a positive integer, got {value}") 

74 if self.rope_dim % 2: 

75 raise ValueError("rope_dim must be even") 

76 

77 @property 

78 def qk_dim(self) -> int: 

79 """Return the score dimension, including the decoupled RoPE band.""" 

80 return self.nope_dim + self.rope_dim 

81 

82 

83@dataclass(frozen=True) 

84class MLAState: 

85 """Normalized local-token latent tensors in BSR layout, before up projection.""" 

86 

87 query_latent: torch.Tensor 

88 kv_latent: torch.Tensor 

89 key_rope: torch.Tensor 

90 

91 

92@dataclass(frozen=True) 

93class MLACPPlan: 

94 """Validated explicit execution strategy; no uncalibrated automatic selection.""" 

95 

96 dimensions: MLADimensions 

97 cp_degree: int 

98 strategy: MLACPStrategy = "expanded_ulysses" 

99 tp_degree: int = 1 

100 overlap: bool = False 

101 ulysses_degree: int | None = None 

102 kv_exchange: Literal["all_gather", "all_to_all_v"] = "all_gather" 

103 

104 def __post_init__(self) -> None: 

105 if not isinstance(self.overlap, bool): 

106 raise ValueError("overlap must be bool") 

107 if self.ulysses_degree is None: 

108 object.__setattr__(self, "ulysses_degree", self.cp_degree) 

109 for name in ("cp_degree", "tp_degree", "ulysses_degree"): 

110 value = getattr(self, name) 

111 if not isinstance(value, int) or isinstance(value, bool) or value < 1: 

112 raise ValueError(f"{name} must be a positive integer") 

113 if self.strategy not in MLA_CP_STRATEGIES: 

114 raise ValueError(f"Unknown MLA CP strategy: {self.strategy}") 

115 if self.cp_degree % self.ulysses_degree: 

116 raise ValueError("ulysses_degree must divide cp_degree") 

117 if self.dimensions.heads % (self.tp_degree * self.ulysses_degree): 

118 raise ValueError("MLA heads must be divisible by TP degree times Ulysses degree") 

119 if self.kv_exchange not in ("all_gather", "all_to_all_v"): 

120 raise ValueError("kv_exchange must be all_gather or all_to_all_v") 

121 if self.kv_exchange == "all_to_all_v" and self.all_gather_degree == 1: 

122 raise ValueError("selective KV exchange requires query-partitioned MLA CP") 

123 if self.tp_degree > 1 and self.all_gather_degree > 1: 

124 raise ValueError("MLA KV AllGather and hybrid CP currently require TP1") 

125 

126 @property 

127 def local_heads(self) -> int: 

128 """Heads in the persistent TP parameter shard, before CP redistribution.""" 

129 return self.dimensions.heads // self.tp_degree 

130 

131 @property 

132 def compute_heads(self) -> int: 

133 """Heads used by this rank's query-partition attention kernel.""" 

134 return self.local_heads // self.ulysses_degree 

135 

136 @property 

137 def all_gather_degree(self) -> int: 

138 """Number of query owners for each Ulysses head partition.""" 

139 return self.cp_degree // self.ulysses_degree 

140 

141 def validate_scale(self, scale: float) -> None: 

142 """Reject invalid scales without replacing an explicit YaRN scale. 

143 

144 Args: 

145 scale: The model's attention score scale, including any YaRN factor. 

146 

147 Raises: 

148 ValueError: If the scale is nonfinite or nonpositive. 

149 """ 

150 if not math.isfinite(scale) or scale <= 0: 

151 raise ValueError("MLA attention scaling must be finite and positive") 

152 

153 

154def mla_all_gather(tensor: torch.Tensor, sequence_dim: int, cp_mesh: Any) -> torch.Tensor: 

155 """Gather source tokens and sum all consumer gradients back to their owner. 

156 

157 Args: 

158 tensor: Equal-sized local sequence shard. 

159 sequence_dim: Dimension along which source shards are concatenated. 

160 cp_mesh: CP mesh, or None for communication-free execution. 

161 

162 Returns: 

163 Full-sequence tensor in CP group-rank order. 

164 """ 

165 if cp_mesh is None or cp_mesh.size() == 1: 

166 return tensor 

167 # Share the overlap path's SUM ReduceScatter adjoint, but wait immediately. 

168 # PyTorch's generic autograd AllGather falls back to AllToAll + local sum on 

169 # HCCL, changing BF16 reduction order when only the overlap flag changes. 

170 return async_cp_allgather_launch(tensor, sequence_dim, cp_mesh).wait() 

171 

172 

173def apply_mla_rope( 

174 tensor: torch.Tensor, 

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

176 *, 

177 interleaved: bool = False, 

178) -> torch.Tensor: 

179 """Rotate one local BSND tensor so Q and shared K can be scheduled independently.""" 

180 if position_embeddings is None: 

181 return tensor 

182 cos, sin = (part.to(tensor.dtype) for part in position_embeddings) 

183 if cos.ndim == 2: 

184 cos, sin = cos.unsqueeze(0), sin.unsqueeze(0) 

185 if interleaved: 

186 cos, sin = (part[..., :tensor.shape[-1] // 2].repeat_interleave(2, dim=-1) 

187 for part in (cos, sin)) 

188 cos, sin = (part.expand(tensor.shape[0], -1, -1).transpose(0, 1) for part in (cos, sin)) 

189 rotated = functional.apply_rotary_pos_emb_single( 

190 tensor.transpose(0, 1), cos, sin, 

191 rotary_interleaved=interleaved, use_fused_rotary_pos_emb=True, 

192 ) 

193 if interleaved: 

194 # V3 consumes real/imaginary half-bands; the kernel returns adjacent pairs. 

195 rotated = rotated.permute(1, 2, 0, 3) 

196 return rotated.unflatten(-1, (-1, 2)).transpose(-1, -2).flatten(-2).transpose(1, 2) 

197 return rotated.transpose(0, 1) 

198 

199 

200def _validate_query_partition(causal, attention_mask): 

201 """Reject query-partition modes whose FA metadata contract is not implemented.""" 

202 if not causal or attention_mask is not None: 

203 raise ValueError("MLA KV AllGather and hybrid CP require causal attention without an explicit mask") 

204 

205 

206@dataclass(frozen=True) 

207class _CausalPartition: 

208 """One continuous query partition and its document-local causal KV prefixes.""" 

209 

210 key_begin: int 

211 key_end: int 

212 query_lengths: tuple[int, ...] 

213 key_lengths: tuple[int, ...] 

214 documents: tuple[int, ...] 

215 

216 def select(self, tensor: torch.Tensor, sequence_dim: int) -> torch.Tensor: 

217 """Keep a view of the needed prefix before either projection or attention. 

218 

219 Args: 

220 tensor: Gathered KV features in global sequence order. 

221 sequence_dim: Axis containing the gathered tokens. 

222 

223 Returns: 

224 A view of the KV interval needed by this query partition. 

225 """ 

226 return tensor.narrow(sequence_dim, self.key_begin, self.key_end - self.key_begin) 

227 

228 

229def _causal_partition(lengths, query_length, query_rank, batch) -> _CausalPartition: 

230 """Coalesce adjacent document prefixes without copying their KV tensors.""" 

231 query_begin = query_rank * query_length 

232 query_end = query_begin + query_length 

233 query_ends, key_ends = [], [] 

234 query_count = key_count = document_begin = 0 

235 key_begin = None 

236 for document_end in lengths: 

237 begin, end = max(document_begin, query_begin), min(document_end, query_end) 

238 if begin < end: 

239 if key_begin is None: 

240 key_begin = document_begin 

241 query_count += end - begin 

242 key_count += end - document_begin 

243 query_ends.append(query_count) 

244 key_ends.append(key_count) 

245 document_begin = document_end 

246 if query_count != query_length: 

247 raise ValueError("MLA document boundaries do not cover the query partition") 

248 # Continuous Q ownership makes all selected document prefixes adjacent. 

249 return _CausalPartition( 

250 key_begin, query_end, 

251 tuple(end + index * query_count for index in range(batch) for end in query_ends), 

252 tuple(end + index * key_count for index in range(batch) for end in key_ends), 

253 tuple(lengths), 

254 ) 

255 

256 

257@lru_cache(maxsize=128) 

258def _kv_exchange_ranges(documents, source_length, query_length, head_copies, degree, rank): 

259 """Cache only host intervals shared by layers; never retain activation tensors.""" 

260 destinations = tuple( 

261 _causal_partition(documents, query_length, peer // head_copies, 1) 

262 for peer in range(degree) 

263 ) 

264 start = rank * source_length 

265 sends, receives = [], [] 

266 local = destinations[rank] 

267 for peer, destination in enumerate(destinations): 

268 begin, end = max(start, destination.key_begin), min(start + source_length, destination.key_end) 

269 sends.append((begin - start, end - begin) if begin < end else (0, 0)) 

270 begin, end = max(peer * source_length, local.key_begin), min((peer + 1) * source_length, local.key_end) 

271 receives.append(max(0, end - begin)) 

272 local_only = all(peer * source_length <= item.key_begin and item.key_end <= (peer + 1) * source_length 

273 for peer, item in enumerate(destinations)) 

274 return tuple(sends), tuple(receives), local_only 

275 

276 

277@dataclass 

278class _SelectiveKVExchange: 

279 """A shared differentiable A2Av and its batch-major receive layout.""" 

280 

281 output: torch.Tensor 

282 receive_tokens: tuple[int, ...] 

283 batch: int 

284 

285 def wait(self) -> torch.Tensor: 

286 """Finish communication before rebuilding batch-major token order.""" 

287 if not self.receive_tokens: 

288 return self.output 

289 output = wait_async_tensor(self.output) 

290 parts, start = [], 0 

291 for count in self.receive_tokens: 

292 rows = self.batch * count 

293 if count: 

294 parts.append(output.narrow(0, start, rows).reshape(self.batch, count, *output.shape[1:])) 

295 start += rows 

296 return parts[0] if len(parts) == 1 else torch.cat(parts, dim=1) 

297 

298 

299def _selective_kv_launch(payload, partition, mesh, head_copies, overlap): 

300 """Send only each query owner's causal document interval, with SUM gradients through repeated slices.""" 

301 batch, source_length = payload.shape[:2] 

302 rank, degree = mesh.get_local_rank(), mesh.size() 

303 sends, receives, local_only = _kv_exchange_ranges( 

304 partition.documents, source_length, partition.query_lengths[-1] // batch, head_copies, degree, rank, 

305 ) 

306 if local_only: 

307 begin, count = sends[rank] 

308 return _SelectiveKVExchange(payload.narrow(1, begin, count), (), batch) 

309 send = torch.cat([payload.narrow(1, begin, count).flatten(0, 1) for begin, count in sends], dim=0) 

310 input_splits = [batch * count for _, count in sends] 

311 output_splits = [batch * count for count in receives] 

312 exchange = differentiable_all_to_all_single_async if overlap else differentiable_all_to_all_single 

313 output = exchange(send, input_splits, output_splits, mesh.get_group()) 

314 return _SelectiveKVExchange(output, receives, batch) 

315 

316 

317def _positions( 

318 embeddings: tuple[torch.Tensor, torch.Tensor] | None, 

319 state: MLAState, 

320 rope_dim: int, 

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

322 """Validate local-token RoPE frequencies and broadcast only the batch axis.""" 

323 if embeddings is None: 

324 return None 

325 if not isinstance(embeddings, (tuple, list)) or len(embeddings) != 2: 

326 raise ValueError("position_embeddings must be a local (cos, sin) pair") 

327 batch, sequence = state.query_latent.shape[:2] 

328 parts = [] 

329 for tensor in embeddings: 

330 if tensor.ndim == 2: 

331 tensor = tensor.unsqueeze(0) 

332 if tensor.shape not in ((1, sequence, rope_dim), (batch, sequence, rope_dim)): 

333 raise ValueError("RoPE cos/sin must describe local tokens with shape [B, S_local, rope_dim]") 

334 if tensor.requires_grad: 

335 raise ValueError("MLA CP currently requires non-trainable RoPE frequencies") 

336 if tensor.device != state.query_latent.device: 

337 raise ValueError("RoPE and MLA activations must be on the same device") 

338 parts.append(tensor.expand(batch, -1, -1)) 

339 return tuple(parts) 

340 

341 

342def _validate_mask( 

343 mask: torch.Tensor | None, 

344 lengths: Sequence[int] | None, 

345 batch: int, 

346 sequence: int, 

347 device: torch.device, 

348) -> None: 

349 """Require global boolean masks and reject unsupported packed-mask combinations.""" 

350 if mask is None: 

351 return 

352 if lengths is not None: 

353 raise ValueError("packed MLA CP currently accepts document causal/noncausal masks only") 

354 if mask.dtype != torch.bool: 

355 raise ValueError("MLA CP masks must be boolean (True=allowed); additive bias is not supported") 

356 valid = ( 

357 mask.ndim == 2 and mask.shape == (sequence, sequence) 

358 or mask.ndim == 4 and mask.shape[0] in (1, batch) 

359 and mask.shape[1] == 1 and mask.shape[-2:] == (sequence, sequence) 

360 ) 

361 if not valid or mask.device != device: 

362 raise ValueError("attention_mask must cover global Q/K: [S, S] or [B|1, 1, S, S], on the input device") 

363 

364 

365def _query( 

366 module: nn.Module, 

367 latent: torch.Tensor, 

368 positions: tuple[torch.Tensor, torch.Tensor] | None, 

369 heads: int, 

370 rotary_fn: Callable[..., torch.Tensor], 

371) -> torch.Tensor: 

372 """Expand queries through the model's rotary callback and managed projection.""" 

373 projected = module.q_b_proj(latent) 

374 projected = projected.view(*latent.shape[:2], heads, module.qk_head_dim) 

375 content, rotary = projected.split((module.qk_nope_head_dim, module.qk_rope_head_dim), dim=-1) 

376 rotated = rotary_fn(rotary, positions, interleaved=module.rotary_interleaved) 

377 # An identity callback needs no split/concatenate copy of the complete Q. 

378 if rotated is rotary: 

379 return projected.transpose(1, 2) 

380 return torch.cat((content, rotated), dim=-1).transpose(1, 2) 

381 

382 

383def _key_value( 

384 module: nn.Module, 

385 latent: torch.Tensor, 

386 key_rope: torch.Tensor, 

387 heads: int, 

388 head_range: tuple[int, int] | None, 

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

390 """Expand the requested KV heads through the managed projection module.""" 

391 projected = (module.kv_b_proj(latent) if head_range is None 

392 else module.kv_b_proj(latent, mla_head_range=head_range)) 

393 return _assemble_key_value(module, projected, key_rope, heads) 

394 

395 

396def _assemble_key_value(module, projected, key_rope, heads): 

397 """Build per-head keys and values from shared projection results.""" 

398 projected = projected.view(*projected.shape[:2], heads, module.qk_nope_head_dim + module.v_head_dim) 

399 content, value = projected.split((module.qk_nope_head_dim, module.v_head_dim), dim=-1) 

400 key = torch.cat((content, key_rope.unsqueeze(2).expand(-1, -1, heads, -1)), dim=-1) 

401 return key.transpose(1, 2), value.transpose(1, 2) 

402 

403 

404def _validate_forward_kwargs(kwargs, position_embeddings): 

405 """Validate forwarded model options before routing any latent.""" 

406 position_ids = kwargs.pop("position_ids", None) 

407 if position_ids is not None and position_embeddings is None: 

408 raise ValueError("position_ids require precomputed local position_embeddings in MLA CP") 

409 kwargs.pop("cache_position", None) 

410 if kwargs.pop("use_cache", False) or kwargs.pop("output_attentions", False): 

411 raise ValueError("MLA CP does not return KV caches or attention weights") 

412 if kwargs: 

413 raise ValueError(f"Unsupported MLA CP forward arguments: {sorted(kwargs)}") 

414 

415 

416def _validate_state(module: nn.Module, plan: MLACPPlan, state: MLAState) -> None: 

417 """Check latent shape and NPU kernel requirements before collectives.""" 

418 dimensions = plan.dimensions 

419 query_latent, kv_latent, key_rope = state.query_latent, state.kv_latent, state.key_rope 

420 if query_latent.ndim != 3: 

421 raise ValueError("MLA CP latents must have BSR layout") 

422 batch, sequence = query_latent.shape[:2] 

423 if batch < 1 or sequence < 1: 

424 raise ValueError("MLA CP requires nonempty batches and sequence shards") 

425 for tensor, width in ((query_latent, dimensions.q_rank), (kv_latent, dimensions.kv_rank), 

426 (key_rope, dimensions.rope_dim)): 

427 if (tensor.shape != (batch, sequence, width) or tensor.device != query_latent.device 

428 or tensor.dtype != query_latent.dtype): 

429 raise ValueError("MLA CP latent shapes/devices/dtypes do not match the plan") 

430 if module.training and module.attention_dropout != 0: 

431 raise ValueError("MLA CP requires attention dropout=0") 

432 plan.validate_scale(module.scaling) 

433 if query_latent.device.type != "npu" or query_latent.dtype not in (torch.float16, torch.bfloat16): 

434 raise ValueError("NPU MLA CP requires NPU BF16/FP16 activations") 

435 if dimensions.value_dim > dimensions.qk_dim: 

436 raise ValueError("NPU MLA CP requires value_dim <= qk_dim") 

437 

438 

439class MLAContextParallel: 

440 """Run CP over normalized latents; parameter-gradient reduction belongs to the trainer.""" 

441 

442 def __init__( 

443 self, 

444 plan: MLACPPlan, 

445 cp_mesh: Any, 

446 *, 

447 rotary_fn: Callable[..., torch.Tensor] | None = None, 

448 attention_interface: Callable[..., tuple[torch.Tensor, Any]] | None = None, 

449 ) -> None: 

450 """Bind communication to optional local rotary and attention implementations. 

451 

452 Args: 

453 plan: Validated MLA communication and head plan. 

454 cp_mesh: Framework-owned CP mesh, or None for degree one. 

455 rotary_fn: Local RoPE callback receiving tensor, positions and interleaved. 

456 It preserves the tensor's shape, dtype and device. None uses generic RoPE. 

457 attention_interface: Local attention callback with the generic attention contract. 

458 It returns BSND output and optional weights. Extra keywords mla_head_offset 

459 and mla_local_heads describe its head slice within the TP-local heads. 

460 None uses the existing NPU attention implementation. 

461 

462 Note: 

463 Callbacks can run again during recomputation; their side effects belong to the caller. 

464 Communication, projections and output restoration remain owned by this runtime. 

465 """ 

466 if (1 if cp_mesh is None else cp_mesh.size()) != plan.cp_degree: 

467 raise ValueError("MLACPPlan degree must match the CP mesh") 

468 if rotary_fn is not None and not callable(rotary_fn): 

469 raise TypeError("rotary_fn must be callable") 

470 if attention_interface is not None and not callable(attention_interface): 

471 raise TypeError("attention_interface must be callable") 

472 self.plan = plan 

473 self.rotary_fn = rotary_fn 

474 self.attention_interface = attention_interface 

475 self.cp_mesh = cp_mesh 

476 rank = 0 if cp_mesh is None else cp_mesh.get_local_rank() 

477 self.query_rank, self.head_rank = divmod(rank, plan.ulysses_degree) 

478 self.ulysses_mesh = self.all_gather_mesh = cp_mesh 

479 if 1 < plan.ulysses_degree < plan.cp_degree: 

480 self.ulysses_mesh, self.all_gather_mesh = _build_hybrid_cp_submeshes(cp_mesh, plan.ulysses_degree) 

481 

482 def __call__( 

483 self, 

484 module: nn.Module, 

485 query_latent: torch.Tensor, 

486 kv_latent: torch.Tensor, 

487 key_rope: torch.Tensor, 

488 *, 

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

490 attention_mask: torch.Tensor | None = None, 

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

492 is_causal: bool | None = None, 

493 restore_output: bool = True, 

494 overlap_input_backward: bool = False, 

495 input_gate: GatedMLAInputState | None = None, 

496 **kwargs: Any, 

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

498 """Run attention and optionally restore its output to token-local BSND. 

499 

500 Args: 

501 overlap_input_backward: Pair input autograd nodes only for pure Ulysses. 

502 AG/hybrid input gradients remain native; output backward is selected 

503 independently by the model adapter. 

504 input_gate: Model-owned gate to cover expanded pure-Ulysses input exchanges. 

505 restore_output: False exposes the head-sharded output so model-owned 

506 independent work can run before its return exchange completes. 

507 """ 

508 packed_kwargs = {name: kwargs.pop(name) for name in PACKED_SEQUENCE_ARGUMENTS if name in kwargs} 

509 packed_kwargs["actual_seq_len"] = actual_seq_len 

510 _validate_forward_kwargs(kwargs, position_embeddings) 

511 state = MLAState(query_latent, kv_latent, key_rope) 

512 dimensions, degree = self.plan.dimensions, self.plan.cp_degree 

513 batch, sequence = query_latent.shape[:2] 

514 lengths, key_lengths = resolve_packed_sequence_lengths( 

515 packed_kwargs, batch * sequence * degree, batch * sequence * degree, 

516 ) 

517 if lengths is not None and batch != 1: 

518 raise ValueError("packed MLA CP currently requires batch_size=1") 

519 if lengths != key_lengths: 

520 raise ValueError("MLA CP requires matching global query and key document boundaries") 

521 causal = getattr(module, "is_causal", True) if is_causal is None else is_causal 

522 if not isinstance(causal, bool): 

523 raise ValueError("is_causal must be bool") 

524 _validate_mask(attention_mask, lengths, batch, sequence * degree, query_latent.device) 

525 attention_mask = prepare_mla_attention_mask(attention_mask, sequence * degree, sequence * degree) 

526 partition = None 

527 if self.plan.all_gather_degree > 1: 

528 _validate_query_partition(causal, attention_mask) 

529 partition = _causal_partition( 

530 lengths or (sequence * degree,), sequence * self.plan.ulysses_degree, self.query_rank, batch, 

531 ) 

532 lengths, key_lengths = partition.query_lengths, partition.key_lengths 

533 positions = _positions(position_embeddings, state, dimensions.rope_dim) 

534 _validate_state(module, self.plan, state) 

535 rotary_fn = apply_mla_rope if self.rotary_fn is None else self.rotary_fn 

536 rotated_key = rotary_fn( 

537 state.key_rope.unsqueeze(2), positions, interleaved=module.rotary_interleaved, 

538 ).squeeze(2) 

539 query, key, value = self._prepare_attention_inputs( 

540 module, state, positions, rotated_key, partition, overlap_input_backward, input_gate, rotary_fn, 

541 ) 

542 attention_interface = ( 

543 npu_fusion_attention_forward if self.attention_interface is None else self.attention_interface 

544 ) 

545 output = attention_interface( 

546 module, query, key, value, attention_mask, scaling=module.scaling, 

547 is_causal=causal, actual_seq_qlen=lengths, actual_seq_kvlen=key_lengths, 

548 pre_tokens=2147483647, next_tokens=0 if causal else 2147483647, 

549 mla_head_offset=self.head_rank * self.plan.compute_heads, 

550 mla_local_heads=self.plan.local_heads, 

551 )[0] 

552 if restore_output: 

553 output = self.restore_output(output) 

554 return output, None 

555 

556 def _prepare_attention_inputs(self, module, state, positions, key_rope, partition, 

557 overlap_input_backward, input_gate, rotary_fn): 

558 """Select input preparation; paired input backward requires pure Ulysses. 

559 

560 Forward overlap also supports KV AllGather and hybrid CP. Those paths 

561 retain native input collective adjoints; their output dO overlap is 

562 selected independently by the model adapter. 

563 """ 

564 use_forward_overlap = self.plan.overlap and self.plan.cp_degree > 1 

565 pure_ulysses = self.plan.all_gather_degree == 1 

566 expanded = self.plan.strategy == "expanded_ulysses" 

567 supports_input_gate = use_forward_overlap and pure_ulysses and expanded 

568 if input_gate is not None and not supports_input_gate: 

569 raise ValueError("Input gate overlap requires expanded pure Ulysses with overlap enabled") 

570 if overlap_input_backward and not pure_ulysses: 

571 raise ValueError("Paired MLA input backward currently requires pure Ulysses") 

572 if not use_forward_overlap: 

573 query = _query(module, state.query_latent, positions, self.plan.local_heads, rotary_fn) 

574 prepare = self._expanded if expanded or self.plan.cp_degree == 1 else self._latent 

575 return prepare(module, state, query, key_rope, partition) 

576 # Pure Ulysses has no outer KV AllGather. Expanded Q/K/V use the same 

577 # V/K/Q forward order with or without a gate; overlap_input_backward 

578 # selects paired input gradients versus the native collective adjoints. 

579 if pure_ulysses and expanded: 

580 return self._prepare_expanded_ulysses_inputs( 

581 module, state, positions, key_rope, input_gate, overlap_input_backward, rotary_fn, 

582 ) 

583 if overlap_input_backward: 

584 # Expanded has returned above and AG/hybrid were rejected earlier: 

585 # this is latent pure Ulysses, pairing Q A2A with KV projection/RS. 

586 return self._prepare_latent_ulysses_inputs(module, state, positions, key_rope, rotary_fn) 

587 # AG/hybrid inputs and forward-only latent Ulysses use native adjoints. 

588 prepare = self._expanded_overlap if expanded else self._latent_overlap 

589 return prepare(module, state, positions, key_rope, partition, rotary_fn=rotary_fn) 

590 

591 def restore_output(self, output: torch.Tensor) -> torch.Tensor: 

592 """Restore token ownership after attention, before the model output projection. 

593 

594 Args: 

595 output: Attention result in BSND head-sharded layout. 

596 

597 Returns: 

598 Original local-token, all-local-head BSND tensor. 

599 """ 

600 batch, sequence = output.shape[:2] 

601 sequence //= self.plan.ulysses_degree 

602 if self.plan.ulysses_degree > 1: 

603 output = ulysses_head_to_seq(output, 1, 2, self.ulysses_mesh) 

604 expected = (batch, sequence, self.plan.local_heads, self.plan.dimensions.value_dim) 

605 if tuple(output.shape) != expected: 

606 raise RuntimeError(f"MLA CP output has shape {tuple(output.shape)}, expected {expected}") 

607 return output 

608 

609 

610 def _expanded(self, module, state, query, key_rope, partition=None): 

611 """Pack unequal-width QKV into one sequence-to-head exchange.""" 

612 key, value = _key_value(module, state.kv_latent, key_rope, self.plan.local_heads, None) 

613 if self.plan.ulysses_degree > 1: 

614 widths = (query.shape[-1], key.shape[-1], value.shape[-1]) 

615 payload = torch.cat(tuple(tensor.transpose(1, 2) for tensor in (query, key, value)), dim=-1) 

616 payload = ulysses_seq_to_head(payload, 1, 2, self.ulysses_mesh) 

617 query, key, value = (tensor.transpose(1, 2) for tensor in payload.split(widths, dim=-1)) 

618 if self.plan.all_gather_degree > 1: 

619 widths = (key.shape[-1], value.shape[-1]) 

620 payload = torch.cat((key.transpose(1, 2), value.transpose(1, 2)), dim=-1) 

621 if self.plan.kv_exchange == "all_to_all_v": 

622 payload = _selective_kv_launch(payload, partition, self.all_gather_mesh, 1, False).wait() 

623 else: 

624 payload = partition.select(mla_all_gather(payload, 1, self.all_gather_mesh), 1) 

625 key, value = (tensor.transpose(1, 2) for tensor in payload.split(widths, dim=-1)) 

626 return query, key, value 

627 

628 def _latent(self, module, state, query, key_rope, partition=None): 

629 """Exchange Q heads, gather shared KV latents, then expand only owned heads.""" 

630 if self.plan.ulysses_degree > 1: 

631 query = ulysses_seq_to_head(query, 2, 1, self.ulysses_mesh) 

632 payload = torch.cat((state.kv_latent, key_rope), dim=-1) 

633 if self.plan.kv_exchange == "all_to_all_v": 

634 payload = _selective_kv_launch(payload, partition, self.cp_mesh, self.plan.ulysses_degree, False).wait() 

635 partition = None 

636 else: 

637 payload = mla_all_gather(payload, 1, self.cp_mesh) 

638 return query, *self._expand_latent(module, payload, partition) 

639 

640 def _expand_latent(self, module, payload, partition=None): 

641 """Keep the same managed head-selective projection for both communication schedules.""" 

642 if partition is not None: 

643 payload = partition.select(payload, 1) 

644 dimensions = self.plan.dimensions 

645 heads = self.plan.compute_heads 

646 start = self.head_rank * heads 

647 kv_latent, key_rope = payload.split((dimensions.kv_rank, dimensions.rope_dim), dim=-1) 

648 # A contiguous latent preserves the native Linear bias-fusion path after payload splitting. 

649 return _key_value(module, kv_latent.contiguous(), key_rope, heads, (start, start + heads)) 

650 

651 def _expanded_overlap(self, module, state, positions, key_rope, partition=None, *, rotary_fn): 

652 """Overlap KV transfer with Q preparation, respecting the two-stage hybrid dependency.""" 

653 key, value = _key_value(module, state.kv_latent, key_rope, self.plan.local_heads, None) 

654 payload = torch.cat((key.transpose(1, 2), value.transpose(1, 2)), dim=-1) 

655 use_ulysses = self.plan.ulysses_degree > 1 

656 if use_ulysses: 

657 kv_transfer_handle = async_ulysses_seq_to_head_launch(payload, 1, 2, self.ulysses_mesh) 

658 else: 

659 kv_transfer_handle = self._expanded_kv_launch(payload, partition) 

660 query = _query(module, state.query_latent, positions, self.plan.local_heads, rotary_fn) 

661 query_a2a_handle = (async_ulysses_seq_to_head_launch(query.transpose(1, 2), 1, 2, self.ulysses_mesh) 

662 if use_ulysses else None) 

663 payload = kv_transfer_handle.wait() 

664 if use_ulysses and self.plan.all_gather_degree > 1: 

665 payload = self._expanded_kv_launch(payload, partition).wait() 

666 if query_a2a_handle is not None: 

667 query = query_a2a_handle.wait().transpose(1, 2) 

668 if partition is not None and self.plan.kv_exchange == "all_gather": 

669 payload = partition.select(payload, 1) 

670 key, value = payload.split((module.qk_head_dim, module.v_head_dim), dim=-1)