Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / context_parallel / dsa_context_parallel.py: 83%

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"""Context parallel style for DeepSeek Sparse Attention (DSA). 

16 

17These PyNative module-level styles prepare DTensor placements at DSA indexer, 

18sparse-attention, and indexer-loss module boundaries. Boundary implementations 

19are responsible for consuming the resulting DTensor or local-tensor inputs. 

20 

21The first implementation intentionally supports only Colossal-style CP: 

22query-side tensors are sharded on sequence, while key-side tensors are gathered 

23to CP-replicated layouts. Ulysses/head sharding is rejected because the current 

24DSA kernels require attention head, index head, head dim and sparse top-k dims to 

25stay replicated. 

26""" 

27from dataclasses import dataclass 

28from typing import Any, Callable, Optional 

29 

30import torch.distributed as dist 

31 

32from hyper_parallel.core.context_parallel.context_parallel import ( 

33 _OUTPUT_NON_CP, 

34 _drop_cp_from_output, 

35 _ensure_1d, 

36 _non_cp_dtensor_layout, 

37 _pop_output_layout, 

38 _push_output_layout, 

39 _to_cp_dtensor, 

40) 

41from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

42from hyper_parallel.core.dtensor.dtensor import DTensor 

43from hyper_parallel.core.dtensor.placement_types import Replicate, Shard 

44from hyper_parallel.core.tensor_parallel.style import ParallelStyle 

45from hyper_parallel.core.context_parallel import utils 

46from hyper_parallel.core.context_parallel.utils import Module 

47 

48 

49 

50_SUPPORTED_LAYOUTS = ("BSND", "TND") 

51_SUPPORTED_LOSS_VARIANTS = ("sparse", "dense") 

52_DEFAULT_ARG_INDEX = object() 

53 

54 

55def _move_dim_to_front(value: Any, dim: int) -> tuple[Any, list[int]]: 

56 """Move ``dim`` to the leading position and return the inverse permutation.""" 

57 order = [dim] + [index for index in range(len(value.shape)) if index != dim] 

58 inverse = [0] * len(order) 

59 for new_index, old_index in enumerate(order): 

60 inverse[old_index] = new_index 

61 return value.permute(order).contiguous(), inverse 

62 

63 

64class _DSASequenceReplicateGradientBridge(utils.Function): 

65 """Reuse one gathered forward value while preserving per-consumer backward.""" 

66 

67 @staticmethod 

68 def forward( 

69 ctx: Any, 

70 local_value: Any, 

71 replicated_value: Any, 

72 group: Any, 

73 world_size: int, 

74 seq_dim: int, 

75 ) -> Any: 

76 """Return the detached shared value and retain only collective metadata.""" 

77 del local_value 

78 ctx.group = group 

79 ctx.world_size = world_size 

80 ctx.seq_dim = seq_dim 

81 return replicated_value 

82 

83 @staticmethod 

84 def backward( # pylint: disable=arguments-differ 

85 ctx: Any, grad_output: Any 

86 ) -> tuple[Any, None, None, None, None]: 

87 """Reduce-scatter this consumer's gradient independently to its local input.""" 

88 if ctx.world_size == 1: 

89 return grad_output, None, None, None, None 

90 grad_front, inverse = _move_dim_to_front(grad_output, ctx.seq_dim) 

91 output_shape = list(grad_front.shape) 

92 if output_shape[0] % ctx.world_size != 0: 

93 raise ValueError( 

94 "DSA shared replicate backward requires a divisible sequence dimension, " 

95 f"got {output_shape[0]} and CP size {ctx.world_size}." 

96 ) 

97 output_shape[0] //= ctx.world_size 

98 local_grad, work = utils.reduce_scatter_single( 

99 grad_front, output_shape, ctx.group, async_op=False 

100 ) 

101 if work is not None: 

102 work.wait() 

103 local_grad = local_grad.permute(inverse).contiguous() 

104 return local_grad, None, None, None, None 

105 

106 

107class DSASequenceReplicateCache: 

108 """Share semantically identical CP AllGather results within one DSA layer. 

109 

110 Callers assign stable semantic slot names only to local tensors that are 

111 mathematically identical views of the same activation. The first consumer 

112 performs the sequence AllGather and later consumers reuse that storage. 

113 Each consumer receives an independent gradient bridge so its reverse 

114 ReduceScatter remains separate, matching the pre-cache accumulation order. 

115 

116 The cache is deliberately scoped to one sparse-attention/indexer-loss 

117 forward interval. :meth:`begin` drops stale entries before sparse attention 

118 and :meth:`clear` releases references after indexer-loss inputs are built. 

119 """ 

120 

121 def __init__(self) -> None: 

122 """Initialize an empty per-layer communication cache.""" 

123 self._values: dict[str, Any] = {} 

124 self._signatures: dict[str, tuple] = {} 

125 

126 @staticmethod 

127 def _signature(value: Any, device_mesh: DeviceMesh, seq_dim: int) -> tuple: 

128 """Build a metadata-only signature without synchronizing the device.""" 

129 if isinstance(value, DTensor): 

130 shape = tuple(value.local_shape) 

131 dtype = value.dtype 

132 else: 

133 shape = tuple(value.shape) 

134 dtype = value.dtype 

135 return shape, dtype, tuple(device_mesh.rank_list), seq_dim 

136 

137 def begin(self) -> None: 

138 """Start a new sparse-attention interval and discard stale references.""" 

139 self.clear() 

140 

141 def replicate(self, slot_name: str, value: Any, device_mesh: DeviceMesh, seq_dim: int) -> Any: 

142 """Return the shared replicated DTensor for one semantic activation.""" 

143 if not _is_tensor_or_dtensor(value): 

144 return value 

145 signature = self._signature(value, device_mesh, seq_dim) 

146 if slot_name in self._values: 

147 if self._signatures[slot_name] != signature: 

148 raise ValueError( 

149 f"DSA shared replicate slot {slot_name!r} received incompatible " 

150 f"metadata: expected {self._signatures[slot_name]}, got {signature}." 

151 ) 

152 replicated = self._values[slot_name] 

153 else: 

154 replicated = _to_sequence_replicate(value, device_mesh, seq_dim) 

155 if isinstance(replicated, DTensor): 

156 replicated = DTensor.from_local_with_layout( 

157 replicated.to_local().detach(), replicated.layout 

158 ) 

159 else: 

160 replicated = replicated.detach() 

161 self._values[slot_name] = replicated 

162 self._signatures[slot_name] = signature 

163 

164 local_value = value.to_local() if isinstance(value, DTensor) else value 

165 replicated_value = replicated.to_local() if isinstance(replicated, DTensor) else replicated 

166 bridged_value = _DSASequenceReplicateGradientBridge.apply( 

167 local_value, 

168 replicated_value, 

169 device_mesh.get_group(), 

170 device_mesh.size(), 

171 seq_dim, 

172 ) 

173 if isinstance(replicated, DTensor): 

174 return DTensor.from_local_with_layout(bridged_value, replicated.layout) 

175 return bridged_value 

176 

177 def clear(self) -> None: 

178 """Release cache-owned references after all forward consumers are wired.""" 

179 self._values.clear() 

180 self._signatures.clear() 

181 

182 

183def _is_tensor_or_dtensor(value: Any) -> bool: 

184 """Return True for framework tensors and HyperParallel DTensors.""" 

185 return isinstance(value, DTensor) or utils.is_tensor(value) 

186 

187 

188def _to_sequence_shard(value: Any, device_mesh: DeviceMesh, seq_dim: int) -> Any: 

189 """Annotate ``value`` as sequence-sharded on ``device_mesh``.""" 

190 if not _is_tensor_or_dtensor(value): 

191 return value 

192 return _to_cp_dtensor(value, device_mesh, (Shard(seq_dim),), (Shard(seq_dim),), seq_dim) 

193 

194 

195def _to_sequence_replicate(value: Any, device_mesh: DeviceMesh, seq_dim: int) -> Any: 

196 """Annotate ``value`` as local sequence shard, then all-gather on CP.""" 

197 if not _is_tensor_or_dtensor(value): 

198 return value 

199 return _to_cp_dtensor(value, device_mesh, (Shard(seq_dim),), (Replicate(),), seq_dim) 

200 

201 

202def _to_query_stats_shard(value: Any, device_mesh: DeviceMesh, stats_seq_dim: int) -> Any: 

203 """Annotate query-side softmax stats as sharded on their sequence dimension.""" 

204 if not _is_tensor_or_dtensor(value): 

205 return value 

206 return _to_cp_dtensor( 

207 value, 

208 device_mesh, 

209 (Shard(stats_seq_dim),), 

210 (Shard(stats_seq_dim),), 

211 stats_seq_dim, 

212 ) 

213 

214 

215def _maybe_replace_arg(args: list, index: Optional[int], fn) -> None: 

216 """Apply ``fn`` to ``args[index]`` when the index points to an existing arg.""" 

217 if index is None or index >= len(args): 

218 return 

219 args[index] = fn(args[index]) 

220 

221 

222def _maybe_replace_kwarg(kwargs: dict, name: Optional[str], fn) -> None: 

223 """Apply ``fn`` to ``kwargs[name]`` when the key exists.""" 

224 if name is None or name not in kwargs: 

225 return 

226 kwargs[name] = fn(kwargs[name]) 

227 

228 

229@dataclass 

230class _ParamSpec: 

231 """Describe one tensor argument location and transform.""" 

232 

233 index: Optional[int] 

234 kwarg_name: Optional[str] 

235 fn: Callable[[Any], Any] 

236 

237 

238def _apply_param_specs(args: list, kwargs: dict, specs: list[_ParamSpec]) -> None: 

239 """Apply each spec's transform to the matching positional or keyword argument.""" 

240 for spec in specs: 

241 _maybe_replace_arg(args, spec.index, spec.fn) 

242 _maybe_replace_kwarg(kwargs, spec.kwarg_name, spec.fn) 

243 

244 

245def _validate_layout_and_mode(style_name: str, layout: str, mode: str) -> tuple[str, int]: 

246 """Return normalized layout and sequence dim for the DSA CP style.""" 

247 layout = layout.upper() 

248 if layout not in _SUPPORTED_LAYOUTS: 

249 raise ValueError(f"layout must be one of {_SUPPORTED_LAYOUTS}, but got {layout!r}.") 

250 if mode != "colossal": 

251 raise ValueError(f"{style_name} currently supports only mode='colossal'.") 

252 return layout, 1 if layout == "BSND" else 0 

253 

254 

255def _validate_loss_variant(loss_variant: str) -> str: 

256 """Return normalized DSA indexer-loss variant.""" 

257 loss_variant = loss_variant.lower() 

258 if loss_variant not in _SUPPORTED_LOSS_VARIANTS: 

259 raise ValueError(f"loss_variant must be one of {_SUPPORTED_LOSS_VARIANTS}, but got {loss_variant!r}.") 

260 return loss_variant 

261 

262 

263def _query_stats_seq_dim(layout: str) -> int: 

264 """Return the sequence dimension used by query-side softmax stats.""" 

265 return 2 if layout == "BSND" else 1 

266 

267 

268def _default_arg_index(index: Any, default: Optional[int]) -> Optional[int]: 

269 """Resolve optional positional index defaults while preserving explicit None.""" 

270 return default if index is _DEFAULT_ARG_INDEX else index 

271 

272 

273def _finalize_output(value: Any, use_local_output: bool, output_layout=None, seq_dim: Optional[int] = None) -> Any: 

274 """Convert direct DTensor outputs, or one-level tuple/list outputs, to local tensors.""" 

275 output_kind, layout = output_layout if output_layout is not None else (None, None) 

276 

277 def finalize_one(item): 

278 if use_local_output: 

279 return item.to_local() if isinstance(item, DTensor) else item 

280 if output_kind == _OUTPUT_NON_CP and seq_dim is not None: 

281 return _drop_cp_from_output(item, layout, (Shard(seq_dim),)) 

282 return item 

283 

284 if isinstance(value, DTensor): 

285 return finalize_one(value) 

286 if isinstance(value, tuple): 

287 return tuple(finalize_one(v) for v in value) 

288 if isinstance(value, list): 

289 return [finalize_one(v) for v in value] 

290 return value 

291 

292 

293def _dtensor_has_partial(value: DTensor) -> bool: 

294 """Return whether ``value`` has any Partial placement.""" 

295 return any(placement.is_partial() for placement in value.placements) 

296 

297 

298def _dtensor_to_local_reducing_partial(value: Any) -> Any: 

299 """Convert a DTensor to local, reducing Partial only when communication is needed.""" 

300 if not isinstance(value, DTensor): 

301 return value 

302 if _dtensor_has_partial(value) and value.device_mesh.mesh.numel() > 1: 

303 value = value.reduce_partial() 

304 return value.to_local() 

305 

306 

307def _register_boundary_hooks(module: Module, pre_hook, use_local_output: bool, seq_dim: int) -> None: 

308 """Register a DSA boundary pre-hook and its public output conversion hook.""" 

309 utils.register_forward_pre_hook(module, pre_hook, with_kwargs=True) 

310 

311 def _finalize_output_hook(hook_module, hook_args, outputs): 

312 del hook_args 

313 return _finalize_output( 

314 outputs, 

315 use_local_output, 

316 _pop_output_layout(hook_module), 

317 seq_dim, 

318 ) 

319 

320 module.register_forward_hook(_finalize_output_hook) 

321 

322 

323def _record_query_output_layout(module: Module, value: Any, cp_mesh: DeviceMesh, seq_dim: int) -> None: 

324 """Remember the non-CP query layout so outputs can drop CP on exit.""" 

325 layout = _non_cp_dtensor_layout(value, cp_mesh, seq_dim) 

326 _push_output_layout(module, (_OUTPUT_NON_CP, layout) if layout is not None else None) 

327 

328 

329def _read_value(args: list, kwargs: dict, index: Optional[int], kwarg_name: Optional[str]) -> Any: 

330 """Read a positional or keyword value from a hook argument pair.""" 

331 if index is not None and index < len(args): 

332 return args[index] 

333 if kwarg_name is not None and kwarg_name in kwargs: 

334 return kwargs[kwarg_name] 

335 return None 

336 

337 

338def _configure_sparse_attention_boundary( # pylint: disable=too-many-arguments 

339 style, 

340 *, 

341 layout: str, 

342 mode: str, 

343 query_index: Optional[int], 

344 key_index: Optional[int], 

345 value_index: Optional[int], 

346 topk_index: Optional[int], 

347 query_kwarg_name: Optional[str], 

348 key_kwarg_name: Optional[str], 

349 value_kwarg_name: Optional[str], 

350 topk_kwarg_name: Optional[str], 

351 query_rope_index: Optional[int], 

352 key_rope_index: Optional[int], 

353 query_rope_kwarg_name: Optional[str], 

354 key_rope_kwarg_name: Optional[str], 

355 use_local_output: bool, 

356 shared_replicate_cache: Optional[DSASequenceReplicateCache], 

357 share_key_value: bool, 

358) -> None: 

359 """Store sparse-attention boundary configuration on ``style``.""" 

360 layout, seq_dim = _validate_layout_and_mode(style.__class__.__name__, layout, mode) 

361 style.layout = layout 

362 style.mode = mode 

363 style.seq_dim = seq_dim 

364 style.query_index = query_index 

365 style.key_index = key_index 

366 style.value_index = value_index 

367 style.topk_index = topk_index 

368 style.query_kwarg_name = query_kwarg_name 

369 style.key_kwarg_name = key_kwarg_name 

370 style.value_kwarg_name = value_kwarg_name 

371 style.topk_kwarg_name = topk_kwarg_name 

372 style.query_rope_index = query_rope_index 

373 style.key_rope_index = key_rope_index 

374 style.query_rope_kwarg_name = query_rope_kwarg_name 

375 style.key_rope_kwarg_name = key_rope_kwarg_name 

376 style.use_local_output = use_local_output 

377 style.shared_replicate_cache = shared_replicate_cache 

378 style.share_key_value = share_key_value 

379 

380 

381def _apply_sparse_attention_boundary( 

382 style, 

383 module: Module, 

384 device_mesh: DeviceMesh, 

385 *, 

386 async_state: Optional[Any] = None, 

387) -> Module: 

388 """Register low-level DSA sparse-attention boundary hooks for ``style``.""" 

389 cp_mesh = _ensure_1d(device_mesh) 

390 

391 def _shard(value: Any) -> Any: 

392 return _to_sequence_shard(value, cp_mesh, style.seq_dim) 

393 

394 def _replicate(slot_name: str): 

395 if async_state is not None: 

396 def _wait_for_replicate(value): 

397 return async_state.wait(slot_name, value) 

398 

399 return _wait_for_replicate 

400 if style.shared_replicate_cache is not None: 

401 def _replicate_from_cache(value): 

402 return style.shared_replicate_cache.replicate( 

403 slot_name, value, cp_mesh, style.seq_dim 

404 ) 

405 

406 return _replicate_from_cache 

407 

408 def _replicate_sequence(value): 

409 return _to_sequence_replicate(value, cp_mesh, style.seq_dim) 

410 

411 return _replicate_sequence 

412 

413 key_slot = "main_kv" if style.share_key_value else "key" 

414 value_slot = "main_kv" if style.share_key_value else "value" 

415 

416 specs = [ 

417 _ParamSpec(style.query_index, style.query_kwarg_name, _shard), 

418 _ParamSpec(style.key_index, style.key_kwarg_name, _replicate(key_slot)), 

419 _ParamSpec(style.value_index, style.value_kwarg_name, _replicate(value_slot)), 

420 _ParamSpec(style.topk_index, style.topk_kwarg_name, _shard), 

421 _ParamSpec(style.query_rope_index, style.query_rope_kwarg_name, _shard), 

422 _ParamSpec(style.key_rope_index, style.key_rope_kwarg_name, _replicate("key_rope")), 

423 ] 

424 

425 def _pre_hook(hook_module, args, kwargs): 

426 new_args = list(args) 

427 new_kwargs = dict(kwargs) 

428 if style.shared_replicate_cache is not None: 

429 style.shared_replicate_cache.begin() 

430 _record_query_output_layout( 

431 hook_module, 

432 _read_value(new_args, new_kwargs, style.query_index, style.query_kwarg_name), 

433 cp_mesh, 

434 style.seq_dim, 

435 ) 

436 _apply_param_specs(new_args, new_kwargs, specs) 

437 return tuple(new_args), new_kwargs 

438 

439 _register_boundary_hooks(module, _pre_hook, style.use_local_output, style.seq_dim) 

440 return module 

441 

442 

443class DSAIndexerContextParallel(ParallelStyle): 

444 """Colossal-style CP hook for a DSA indexer boundary. 

445 

446 This style targets a hookable module/cell whose forward signature is shaped 

447 like ``(query, key, weights, ...)`` and rewrites only that boundary: 

448 

449 - ``query`` and ``weights`` are annotated as ``Shard(seq)``; 

450 - ``key`` is all-gathered to ``Replicate()``. 

451 """ 

452 

453 def __init__( # pylint: disable=too-many-arguments 

454 self, 

455 *, 

456 layout: str = "BSND", 

457 mode: str = "colossal", 

458 query_index: Optional[int] = 0, 

459 key_index: Optional[int] = 1, 

460 weights_index: Optional[int] = 2, 

461 query_kwarg_name: Optional[str] = None, 

462 key_kwarg_name: Optional[str] = None, 

463 weights_kwarg_name: Optional[str] = None, 

464 use_local_output: bool = False, 

465 ) -> None: 

466 super().__init__() 

467 layout, seq_dim = _validate_layout_and_mode(self.__class__.__name__, layout, mode) 

468 self.layout = layout 

469 self.mode = mode 

470 self.seq_dim = seq_dim 

471 self.query_index = query_index 

472 self.key_index = key_index 

473 self.weights_index = weights_index 

474 self.query_kwarg_name = query_kwarg_name 

475 self.key_kwarg_name = key_kwarg_name 

476 self.weights_kwarg_name = weights_kwarg_name 

477 self.use_local_output = use_local_output 

478 

479 def __repr__(self) -> str: 

480 return ( 

481 f"{self.__class__.__name__}(" 

482 f"layout={self.layout!r}, mode={self.mode!r}, " 

483 f"use_local_output={self.use_local_output})" 

484 ) 

485 

486 def _shard_query_side(self, value: Any, device_mesh: DeviceMesh) -> Any: 

487 return _to_sequence_shard(value, device_mesh, self.seq_dim) 

488 

489 def _replicate_key_side(self, value: Any, device_mesh: DeviceMesh) -> Any: 

490 return _to_sequence_replicate(value, device_mesh, self.seq_dim) 

491 

492 def _build_specs(self, cp_mesh: DeviceMesh, key_fn: Callable[[Any], Any]) -> list[_ParamSpec]: 

493 """Build table-driven transforms for the indexer boundary.""" 

494 def shard(value: Any) -> Any: 

495 return self._shard_query_side(value, cp_mesh) 

496 

497 return [ 

498 _ParamSpec(self.query_index, self.query_kwarg_name, shard), 

499 _ParamSpec(self.key_index, self.key_kwarg_name, key_fn), 

500 _ParamSpec(self.weights_index, self.weights_kwarg_name, shard), 

501 ] 

502 

503 def _apply_with_specs(self, module: Module, specs: list[_ParamSpec], cp_mesh: DeviceMesh) -> Module: 

504 """Register indexer hooks driven by parameter specs.""" 

505 def _pre_hook(hook_module, args, kwargs): 

506 new_args = list(args) 

507 new_kwargs = dict(kwargs) 

508 _record_query_output_layout( 

509 hook_module, 

510 _read_value(new_args, new_kwargs, self.query_index, self.query_kwarg_name), 

511 cp_mesh, 

512 self.seq_dim, 

513 ) 

514 _apply_param_specs(new_args, new_kwargs, specs) 

515 return tuple(new_args), new_kwargs 

516 

517 _register_boundary_hooks(module, _pre_hook, self.use_local_output, self.seq_dim) 

518 return module 

519 

520 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module: 

521 """Register DSA indexer CP hooks on ``module`` and return it.""" 

522 cp_mesh = _ensure_1d(device_mesh) 

523 specs = self._build_specs( 

524 cp_mesh, 

525 key_fn=lambda value: self._replicate_key_side(value, cp_mesh), 

526 ) 

527 return self._apply_with_specs(module, specs, cp_mesh) 

528 

529 

530class DSASparseAttentionContextParallel(ParallelStyle): 

531 """Colossal-style CP hook for a DSA sparse-attention boundary. 

532 

533 This style targets a hookable module/cell whose forward signature is shaped 

534 like ``(query, key, value, topk_indices, query_rope, key_rope, ...)`` and 

535 rewrites only that boundary: 

536 

537 - ``query``, ``topk_indices`` and ``query_rope`` are annotated as ``Shard(seq)``; 

538 - ``key``, ``value`` and ``key_rope`` are all-gathered to ``Replicate()``. 

539 """ 

540 

541 def __init__( # pylint: disable=too-many-arguments 

542 self, 

543 *, 

544 layout: str = "BSND", 

545 mode: str = "colossal", 

546 query_index: Optional[int] = 0, 

547 key_index: Optional[int] = 1, 

548 value_index: Optional[int] = 2, 

549 topk_index: Optional[int] = 3, 

550 query_kwarg_name: Optional[str] = None, 

551 key_kwarg_name: Optional[str] = None, 

552 value_kwarg_name: Optional[str] = None, 

553 topk_kwarg_name: Optional[str] = None, 

554 query_rope_index: Optional[int] = 4, 

555 key_rope_index: Optional[int] = 5, 

556 query_rope_kwarg_name: Optional[str] = "query_rope", 

557 key_rope_kwarg_name: Optional[str] = "key_rope", 

558 use_local_output: bool = False, 

559 shared_replicate_cache: Optional[DSASequenceReplicateCache] = None, 

560 share_key_value: bool = False, 

561 ) -> None: 

562 super().__init__() 

563 _configure_sparse_attention_boundary( 

564 self, 

565 layout=layout, 

566 mode=mode, 

567 query_index=query_index, 

568 key_index=key_index, 

569 value_index=value_index, 

570 topk_index=topk_index, 

571 query_kwarg_name=query_kwarg_name, 

572 key_kwarg_name=key_kwarg_name, 

573 value_kwarg_name=value_kwarg_name, 

574 topk_kwarg_name=topk_kwarg_name, 

575 query_rope_index=query_rope_index, 

576 key_rope_index=key_rope_index, 

577 query_rope_kwarg_name=query_rope_kwarg_name, 

578 key_rope_kwarg_name=key_rope_kwarg_name, 

579 use_local_output=use_local_output, 

580 shared_replicate_cache=shared_replicate_cache, 

581 share_key_value=share_key_value, 

582 ) 

583 

584 def __repr__(self) -> str: 

585 return ( 

586 f"{self.__class__.__name__}(" 

587 f"layout={self.layout!r}, mode={self.mode!r}, " 

588 f"use_local_output={self.use_local_output})" 

589 ) 

590 

591 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module: 

592 """Register DSA sparse-attention CP hooks on ``module`` and return it.""" 

593 return _apply_sparse_attention_boundary(self, module, device_mesh) 

594 

595 

596class DSAIndexerLossContextParallel(ParallelStyle): 

597 """Colossal-style CP hook for a DSA indexer-loss kernel boundary. 

598 

599 This style targets a hookable module/cell whose forward signature is shaped 

600 like the sparse indexer-loss variant by default: 

601 ``(query, key, query_index, key_index, weights, topk_indices, 

602 softmax_max, softmax_sum, query_rope, key_rope, ...)``. Set 

603 ``loss_variant="dense"`` for the dense loss signature: 

604 ``(query, key, query_index, key_index, weights, softmax_max, softmax_sum, 

605 softmax_max_index, softmax_sum_index, scale_value, query_rope, key_rope, 

606 ...)``. The boundary is expected to start after MF has already done 

607 local-only bookkeeping such as ``stop_gradient`` and ``split``. 

608 

609 Placements: 

610 - ``query``, ``query_index``, ``weights``, ``topk_indices`` and 

611 ``query_rope`` are annotated as ``Shard(seq)``; 

612 - ``softmax_max``, ``softmax_sum`` and dense index softmax stats are 

613 annotated as query-side stats sharded on their stats sequence dimension; 

614 - ``key``, ``key_index`` and ``key_rope`` are all-gathered to 

615 ``Replicate()``. 

616 """ 

617 

618 def __init__( # pylint: disable=too-many-arguments 

619 self, 

620 *, 

621 layout: str = "BSND", 

622 mode: str = "colossal", 

623 loss_variant: str = "sparse", 

624 query_index: Optional[int] = 0, 

625 key_index: Optional[int] = 1, 

626 query_indexer_index: Optional[int] = 2, 

627 key_indexer_index: Optional[int] = 3, 

628 weights_index: Optional[int] = 4, 

629 topk_index: Optional[int] = _DEFAULT_ARG_INDEX, 

630 softmax_max_index: Optional[int] = _DEFAULT_ARG_INDEX, 

631 softmax_sum_index: Optional[int] = _DEFAULT_ARG_INDEX, 

632 softmax_max_indexer_index: Optional[int] = _DEFAULT_ARG_INDEX, 

633 softmax_sum_indexer_index: Optional[int] = _DEFAULT_ARG_INDEX, 

634 query_rope_index: Optional[int] = _DEFAULT_ARG_INDEX, 

635 key_rope_index: Optional[int] = _DEFAULT_ARG_INDEX, 

636 query_kwarg_name: Optional[str] = None, 

637 key_kwarg_name: Optional[str] = None, 

638 query_indexer_kwarg_name: Optional[str] = None, 

639 key_indexer_kwarg_name: Optional[str] = None, 

640 weights_kwarg_name: Optional[str] = None, 

641 topk_kwarg_name: Optional[str] = None, 

642 softmax_max_kwarg_name: Optional[str] = None, 

643 softmax_sum_kwarg_name: Optional[str] = None, 

644 softmax_max_indexer_kwarg_name: Optional[str] = None, 

645 softmax_sum_indexer_kwarg_name: Optional[str] = None, 

646 query_rope_kwarg_name: Optional[str] = None, 

647 key_rope_kwarg_name: Optional[str] = None, 

648 use_local_output: bool = False, 

649 shared_replicate_cache: Optional[DSASequenceReplicateCache] = None, 

650 ) -> None: 

651 super().__init__() 

652 layout, seq_dim = _validate_layout_and_mode(self.__class__.__name__, layout, mode) 

653 loss_variant = _validate_loss_variant(loss_variant) 

654 is_dense = loss_variant == "dense" 

655 self.layout = layout 

656 self.mode = mode 

657 self.loss_variant = loss_variant 

658 self.seq_dim = seq_dim 

659 self.stats_seq_dim = _query_stats_seq_dim(layout) 

660 self.query_index = query_index 

661 self.key_index = key_index 

662 self.query_indexer_index = query_indexer_index 

663 self.key_indexer_index = key_indexer_index 

664 self.weights_index = weights_index 

665 self.topk_index = _default_arg_index(topk_index, None if is_dense else 5) 

666 self.softmax_max_index = _default_arg_index(softmax_max_index, 5 if is_dense else 6) 

667 self.softmax_sum_index = _default_arg_index(softmax_sum_index, 6 if is_dense else 7) 

668 self.softmax_max_indexer_index = _default_arg_index(softmax_max_indexer_index, 7 if is_dense else None) 

669 self.softmax_sum_indexer_index = _default_arg_index(softmax_sum_indexer_index, 8 if is_dense else None) 

670 self.query_rope_index = _default_arg_index(query_rope_index, 10 if is_dense else 8) 

671 self.key_rope_index = _default_arg_index(key_rope_index, 11 if is_dense else 9) 

672 self.query_kwarg_name = query_kwarg_name 

673 self.key_kwarg_name = key_kwarg_name 

674 self.query_indexer_kwarg_name = query_indexer_kwarg_name 

675 self.key_indexer_kwarg_name = key_indexer_kwarg_name 

676 self.weights_kwarg_name = weights_kwarg_name 

677 self.topk_kwarg_name = topk_kwarg_name 

678 self.softmax_max_kwarg_name = softmax_max_kwarg_name 

679 self.softmax_sum_kwarg_name = softmax_sum_kwarg_name 

680 self.softmax_max_indexer_kwarg_name = softmax_max_indexer_kwarg_name 

681 self.softmax_sum_indexer_kwarg_name = softmax_sum_indexer_kwarg_name 

682 self.query_rope_kwarg_name = query_rope_kwarg_name 

683 self.key_rope_kwarg_name = key_rope_kwarg_name 

684 self.use_local_output = use_local_output 

685 self.shared_replicate_cache = shared_replicate_cache 

686 

687 def __repr__(self) -> str: 

688 return ( 

689 f"{self.__class__.__name__}(" 

690 f"layout={self.layout!r}, mode={self.mode!r}, " 

691 f"loss_variant={self.loss_variant!r}, " 

692 f"use_local_output={self.use_local_output})" 

693 ) 

694 

695 def _shard_query_side(self, value: Any, device_mesh: DeviceMesh) -> Any: 

696 return _to_sequence_shard(value, device_mesh, self.seq_dim) 

697 

698 def _replicate_key_side(self, value: Any, device_mesh: DeviceMesh) -> Any: 

699 return _to_sequence_replicate(value, device_mesh, self.seq_dim) 

700 

701 @staticmethod 

702 def _local_shape(value: Any) -> Optional[tuple]: 

703 if isinstance(value, DTensor): 

704 return value.local_shape 

705 if utils.is_tensor(value): 

706 return value.shape 

707 return None 

708 

709 def _slice_local_key_grad(self, value: Any, module: Module) -> Any: 

710 """Convert d_key_index to the original local key-index shard shape.""" 

711 if not self.use_local_output: 

712 return value 

713 target_shape = getattr(module, "_hp_dsa_loss_key_index_local_shape", None) 

714 if target_shape is None: 

715 return _finalize_output(value, use_local_output=True) 

716 

717 if isinstance(value, DTensor): 

718 value = _dtensor_to_local_reducing_partial(value) 

719 if not utils.is_tensor(value): 

720 return value 

721 

722 target_len = target_shape[self.seq_dim] 

723 if value.shape[self.seq_dim] == target_len: 

724 return value 

725 

726 local_idx = getattr(module, "_hp_dsa_loss_local_idx", 0) 

727 start = local_idx * target_len 

728 return value.narrow(self.seq_dim, start, target_len) 

729 

730 def _process_outputs(self, module: Module, outputs: Any) -> Any: 

731 """Finalize indexer-loss outputs, reducing Partial values when needed.""" 

732 output_layout = _pop_output_layout(module) 

733 if not self.use_local_output: 

734 return _finalize_output(outputs, False, output_layout, self.seq_dim) 

735 if not isinstance(outputs, (tuple, list)) or len(outputs) < 4: 

736 return _finalize_output(outputs, use_local_output=True) 

737 

738 processed = list(outputs) 

739 processed[0] = _dtensor_to_local_reducing_partial(processed[0]) 

740 processed[1] = self._slice_local_key_grad(processed[1], module) 

741 processed[2] = _dtensor_to_local_reducing_partial(processed[2]) 

742 processed[3] = _dtensor_to_local_reducing_partial(processed[3]) 

743 return type(outputs)(processed) 

744 

745 def _build_loss_specs( 

746 self, 

747 cp_mesh: DeviceMesh, 

748 replicate_fn_map: dict[str, Callable[[Any], Any]], 

749 ) -> list[_ParamSpec]: 

750 """Build table-driven transforms for the indexer-loss boundary.""" 

751 def shard(value: Any) -> Any: 

752 return self._shard_query_side(value, cp_mesh) 

753 

754 def stats_shard(value: Any) -> Any: 

755 return _to_query_stats_shard(value, cp_mesh, self.stats_seq_dim) 

756 

757 return [ 

758 _ParamSpec(self.query_index, self.query_kwarg_name, shard), 

759 _ParamSpec(self.key_index, self.key_kwarg_name, replicate_fn_map["key"]), 

760 _ParamSpec(self.query_indexer_index, self.query_indexer_kwarg_name, shard), 

761 _ParamSpec(self.key_indexer_index, self.key_indexer_kwarg_name, replicate_fn_map["key_indexer"]), 

762 _ParamSpec(self.weights_index, self.weights_kwarg_name, shard), 

763 _ParamSpec(self.topk_index, self.topk_kwarg_name, shard), 

764 _ParamSpec(self.softmax_max_index, self.softmax_max_kwarg_name, stats_shard), 

765 _ParamSpec(self.softmax_sum_index, self.softmax_sum_kwarg_name, stats_shard), 

766 _ParamSpec(self.softmax_max_indexer_index, self.softmax_max_indexer_kwarg_name, stats_shard), 

767 _ParamSpec(self.softmax_sum_indexer_index, self.softmax_sum_indexer_kwarg_name, stats_shard),