Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / context_parallel / dsa_context_parallel.py: 85%
291 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Context parallel style for DeepSeek Sparse Attention (DSA).
17These PyNative module-level styles are companions to the DSA distributed
18operators. The distributed operators define the per-op layout rules for
19``lightning_indexer``, ``npu_sparse_flash_attention`` and indexer-loss custom
20ops; these styles prepare module inputs so those rules can be selected by the
21DTensor dispatcher.
23The first implementation intentionally supports only Colossal-style CP:
24query-side tensors are sharded on sequence, while key-side tensors are gathered
25to CP-replicated layouts. Ulysses/head sharding is rejected because the current
26DSA kernels require attention head, index head, head dim and sparse top-k dims to
27stay replicated.
28"""
29from dataclasses import dataclass
30from typing import Any, Callable, Optional
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.platform import get_platform
47platform = get_platform()
48Module = platform.Module
51_SUPPORTED_LAYOUTS = ("BSND", "TND")
52_SUPPORTED_LOSS_VARIANTS = ("sparse", "dense")
53_DEFAULT_ARG_INDEX = object()
56def _is_tensor_or_dtensor(value: Any) -> bool:
57 """Return True for framework tensors and HyperParallel DTensors."""
58 return isinstance(value, DTensor) or platform.is_tensor(value)
61def _to_sequence_shard(value: Any, device_mesh: DeviceMesh, seq_dim: int) -> Any:
62 """Annotate ``value`` as sequence-sharded on ``device_mesh``."""
63 if not _is_tensor_or_dtensor(value):
64 return value
65 return _to_cp_dtensor(value, device_mesh, (Shard(seq_dim),), (Shard(seq_dim),), seq_dim)
68def _to_sequence_replicate(value: Any, device_mesh: DeviceMesh, seq_dim: int) -> Any:
69 """Annotate ``value`` as local sequence shard, then all-gather on CP."""
70 if not _is_tensor_or_dtensor(value):
71 return value
72 return _to_cp_dtensor(value, device_mesh, (Shard(seq_dim),), (Replicate(),), seq_dim)
75def _to_query_stats_shard(value: Any, device_mesh: DeviceMesh, stats_seq_dim: int) -> Any:
76 """Annotate query-side softmax stats as sharded on their sequence dimension."""
77 if not _is_tensor_or_dtensor(value):
78 return value
79 return _to_cp_dtensor(
80 value,
81 device_mesh,
82 (Shard(stats_seq_dim),),
83 (Shard(stats_seq_dim),),
84 stats_seq_dim,
85 )
88def _maybe_replace_arg(args: list, index: Optional[int], fn) -> None:
89 """Apply ``fn`` to ``args[index]`` when the index points to an existing arg."""
90 if index is None or index >= len(args):
91 return
92 args[index] = fn(args[index])
95def _maybe_replace_kwarg(kwargs: dict, name: Optional[str], fn) -> None:
96 """Apply ``fn`` to ``kwargs[name]`` when the key exists."""
97 if name is None or name not in kwargs:
98 return
99 kwargs[name] = fn(kwargs[name])
102@dataclass
103class _ParamSpec:
104 """Describe one tensor argument location and transform."""
106 index: Optional[int]
107 kwarg_name: Optional[str]
108 fn: Callable[[Any], Any]
111def _apply_param_specs(args: list, kwargs: dict, specs: list[_ParamSpec]) -> None:
112 """Apply each spec's transform to the matching positional or keyword argument."""
113 for spec in specs:
114 _maybe_replace_arg(args, spec.index, spec.fn)
115 _maybe_replace_kwarg(kwargs, spec.kwarg_name, spec.fn)
118def _validate_layout_and_mode(style_name: str, layout: str, mode: str) -> tuple[str, int]:
119 """Return normalized layout and sequence dim for the DSA CP style."""
120 layout = layout.upper()
121 if layout not in _SUPPORTED_LAYOUTS:
122 raise ValueError(f"layout must be one of {_SUPPORTED_LAYOUTS}, but got {layout!r}.")
123 if mode != "colossal":
124 raise ValueError(f"{style_name} currently supports only mode='colossal'.")
125 return layout, 1 if layout == "BSND" else 0
128def _validate_loss_variant(loss_variant: str) -> str:
129 """Return normalized DSA indexer-loss variant."""
130 loss_variant = loss_variant.lower()
131 if loss_variant not in _SUPPORTED_LOSS_VARIANTS:
132 raise ValueError(f"loss_variant must be one of {_SUPPORTED_LOSS_VARIANTS}, but got {loss_variant!r}.")
133 return loss_variant
136def _query_stats_seq_dim(layout: str) -> int:
137 """Return the sequence dimension used by query-side softmax stats."""
138 return 2 if layout == "BSND" else 1
141def _default_arg_index(index: Any, default: Optional[int]) -> Optional[int]:
142 """Resolve optional positional index defaults while preserving explicit None."""
143 return default if index is _DEFAULT_ARG_INDEX else index
146def _finalize_output(value: Any, use_local_output: bool, output_layout=None, seq_dim: Optional[int] = None) -> Any:
147 """Convert direct DTensor outputs, or one-level tuple/list outputs, to local tensors."""
148 output_kind, layout = output_layout if output_layout is not None else (None, None)
150 def finalize_one(item):
151 if use_local_output:
152 return item.to_local() if isinstance(item, DTensor) else item
153 if output_kind == _OUTPUT_NON_CP and seq_dim is not None:
154 return _drop_cp_from_output(item, layout, (Shard(seq_dim),))
155 return item
157 if isinstance(value, DTensor):
158 return finalize_one(value)
159 if isinstance(value, tuple):
160 return tuple(finalize_one(v) for v in value)
161 if isinstance(value, list):
162 return [finalize_one(v) for v in value]
163 return value
166def _dtensor_has_partial(value: DTensor) -> bool:
167 """Return whether ``value`` has any Partial placement."""
168 return any(placement.is_partial() for placement in value.placements)
171def _dtensor_to_local_reducing_partial(value: Any) -> Any:
172 """Convert a DTensor to local, reducing Partial only when communication is needed."""
173 if not isinstance(value, DTensor):
174 return value
175 if _dtensor_has_partial(value) and value.device_mesh.mesh.numel() > 1:
176 value = value.reduce_partial()
177 return value.to_local()
180def _register_boundary_hooks(module: Module, pre_hook, use_local_output: bool, seq_dim: int) -> None:
181 """Register a DSA boundary pre-hook and its public output conversion hook."""
182 platform.register_forward_pre_hook(module, pre_hook, with_kwargs=True)
183 def _finalize_output_hook(hook_module, hook_args, outputs):
184 del hook_args
185 return _finalize_output(
186 outputs,
187 use_local_output,
188 _pop_output_layout(hook_module),
189 seq_dim,
190 )
192 module.register_forward_hook(_finalize_output_hook)
195def _record_query_output_layout(module: Module, value: Any, cp_mesh: DeviceMesh, seq_dim: int) -> None:
196 """Remember the non-CP query layout so outputs can drop CP on exit."""
197 layout = _non_cp_dtensor_layout(value, cp_mesh, seq_dim)
198 _push_output_layout(module, (_OUTPUT_NON_CP, layout) if layout is not None else None)
201def _read_value(args: list, kwargs: dict, index: Optional[int], kwarg_name: Optional[str]) -> Any:
202 """Read a positional or keyword value from a hook argument pair."""
203 if index is not None and index < len(args):
204 return args[index]
205 if kwarg_name is not None and kwarg_name in kwargs:
206 return kwargs[kwarg_name]
207 return None
210def _configure_sparse_attention_boundary( # pylint: disable=too-many-arguments
211 style,
212 *,
213 layout: str,
214 mode: str,
215 query_index: Optional[int],
216 key_index: Optional[int],
217 value_index: Optional[int],
218 topk_index: Optional[int],
219 query_kwarg_name: Optional[str],
220 key_kwarg_name: Optional[str],
221 value_kwarg_name: Optional[str],
222 topk_kwarg_name: Optional[str],
223 query_rope_index: Optional[int],
224 key_rope_index: Optional[int],
225 query_rope_kwarg_name: Optional[str],
226 key_rope_kwarg_name: Optional[str],
227 use_local_output: bool,
228) -> None:
229 """Store sparse-attention boundary configuration on ``style``."""
230 layout, seq_dim = _validate_layout_and_mode(style.__class__.__name__, layout, mode)
231 style.layout = layout
232 style.mode = mode
233 style.seq_dim = seq_dim
234 style.query_index = query_index
235 style.key_index = key_index
236 style.value_index = value_index
237 style.topk_index = topk_index
238 style.query_kwarg_name = query_kwarg_name
239 style.key_kwarg_name = key_kwarg_name
240 style.value_kwarg_name = value_kwarg_name
241 style.topk_kwarg_name = topk_kwarg_name
242 style.query_rope_index = query_rope_index
243 style.key_rope_index = key_rope_index
244 style.query_rope_kwarg_name = query_rope_kwarg_name
245 style.key_rope_kwarg_name = key_rope_kwarg_name
246 style.use_local_output = use_local_output
249def _apply_sparse_attention_boundary(
250 style,
251 module: Module,
252 device_mesh: DeviceMesh,
253 *,
254 async_state: Optional[Any] = None,
255) -> Module:
256 """Register low-level DSA sparse-attention boundary hooks for ``style``."""
257 cp_mesh = _ensure_1d(device_mesh)
259 def _shard(value: Any) -> Any:
260 return _to_sequence_shard(value, cp_mesh, style.seq_dim)
262 def _replicate(slot_name: str):
263 if async_state is not None:
264 return lambda value: async_state.wait(slot_name, value)
265 return lambda value: _to_sequence_replicate(value, cp_mesh, style.seq_dim)
267 specs = [
268 _ParamSpec(style.query_index, style.query_kwarg_name, _shard),
269 _ParamSpec(style.key_index, style.key_kwarg_name, _replicate("key")),
270 _ParamSpec(style.value_index, style.value_kwarg_name, _replicate("value")),
271 _ParamSpec(style.topk_index, style.topk_kwarg_name, _shard),
272 _ParamSpec(style.query_rope_index, style.query_rope_kwarg_name, _shard),
273 _ParamSpec(style.key_rope_index, style.key_rope_kwarg_name, _replicate("key_rope")),
274 ]
276 def _pre_hook(hook_module, args, kwargs):
277 new_args = list(args)
278 new_kwargs = dict(kwargs)
279 _record_query_output_layout(
280 hook_module,
281 _read_value(new_args, new_kwargs, style.query_index, style.query_kwarg_name),
282 cp_mesh,
283 style.seq_dim,
284 )
285 _apply_param_specs(new_args, new_kwargs, specs)
286 return tuple(new_args), new_kwargs
288 _register_boundary_hooks(module, _pre_hook, style.use_local_output, style.seq_dim)
289 return module
292class DSAIndexerContextParallel(ParallelStyle):
293 """Colossal-style CP hook for a DSA indexer boundary.
295 This style targets a hookable module/cell whose forward signature is shaped
296 like ``(query, key, weights, ...)`` and rewrites only that boundary:
298 - ``query`` and ``weights`` are annotated as ``Shard(seq)``;
299 - ``key`` is all-gathered to ``Replicate()``.
300 """
302 def __init__( # pylint: disable=too-many-arguments
303 self,
304 *,
305 layout: str = "BSND",
306 mode: str = "colossal",
307 query_index: Optional[int] = 0,
308 key_index: Optional[int] = 1,
309 weights_index: Optional[int] = 2,
310 query_kwarg_name: Optional[str] = None,
311 key_kwarg_name: Optional[str] = None,
312 weights_kwarg_name: Optional[str] = None,
313 use_local_output: bool = False,
314 ) -> None:
315 super().__init__()
316 layout, seq_dim = _validate_layout_and_mode(self.__class__.__name__, layout, mode)
317 self.layout = layout
318 self.mode = mode
319 self.seq_dim = seq_dim
320 self.query_index = query_index
321 self.key_index = key_index
322 self.weights_index = weights_index
323 self.query_kwarg_name = query_kwarg_name
324 self.key_kwarg_name = key_kwarg_name
325 self.weights_kwarg_name = weights_kwarg_name
326 self.use_local_output = use_local_output
328 def __repr__(self) -> str:
329 return (
330 f"{self.__class__.__name__}("
331 f"layout={self.layout!r}, mode={self.mode!r}, "
332 f"use_local_output={self.use_local_output})"
333 )
335 def _shard_query_side(self, value: Any, device_mesh: DeviceMesh) -> Any:
336 return _to_sequence_shard(value, device_mesh, self.seq_dim)
338 def _replicate_key_side(self, value: Any, device_mesh: DeviceMesh) -> Any:
339 return _to_sequence_replicate(value, device_mesh, self.seq_dim)
341 def _build_specs(self, cp_mesh: DeviceMesh, key_fn: Callable[[Any], Any]) -> list[_ParamSpec]:
342 """Build table-driven transforms for the indexer boundary."""
343 def shard(value: Any) -> Any:
344 return self._shard_query_side(value, cp_mesh)
346 return [
347 _ParamSpec(self.query_index, self.query_kwarg_name, shard),
348 _ParamSpec(self.key_index, self.key_kwarg_name, key_fn),
349 _ParamSpec(self.weights_index, self.weights_kwarg_name, shard),
350 ]
352 def _apply_with_specs(self, module: Module, specs: list[_ParamSpec], cp_mesh: DeviceMesh) -> Module:
353 """Register indexer hooks driven by parameter specs."""
354 def _pre_hook(hook_module, args, kwargs):
355 new_args = list(args)
356 new_kwargs = dict(kwargs)
357 _record_query_output_layout(
358 hook_module,
359 _read_value(new_args, new_kwargs, self.query_index, self.query_kwarg_name),
360 cp_mesh,
361 self.seq_dim,
362 )
363 _apply_param_specs(new_args, new_kwargs, specs)
364 return tuple(new_args), new_kwargs
366 _register_boundary_hooks(module, _pre_hook, self.use_local_output, self.seq_dim)
367 return module
369 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module:
370 """Register DSA indexer CP hooks on ``module`` and return it."""
371 cp_mesh = _ensure_1d(device_mesh)
372 specs = self._build_specs(
373 cp_mesh,
374 key_fn=lambda value: self._replicate_key_side(value, cp_mesh),
375 )
376 return self._apply_with_specs(module, specs, cp_mesh)
379class DSASparseAttentionContextParallel(ParallelStyle):
380 """Colossal-style CP hook for a DSA sparse-attention boundary.
382 This style targets a hookable module/cell whose forward signature is shaped
383 like ``(query, key, value, topk_indices, query_rope, key_rope, ...)`` and
384 rewrites only that boundary:
386 - ``query``, ``topk_indices`` and ``query_rope`` are annotated as ``Shard(seq)``;
387 - ``key``, ``value`` and ``key_rope`` are all-gathered to ``Replicate()``.
388 """
390 def __init__( # pylint: disable=too-many-arguments
391 self,
392 *,
393 layout: str = "BSND",
394 mode: str = "colossal",
395 query_index: Optional[int] = 0,
396 key_index: Optional[int] = 1,
397 value_index: Optional[int] = 2,
398 topk_index: Optional[int] = 3,
399 query_kwarg_name: Optional[str] = None,
400 key_kwarg_name: Optional[str] = None,
401 value_kwarg_name: Optional[str] = None,
402 topk_kwarg_name: Optional[str] = None,
403 query_rope_index: Optional[int] = 4,
404 key_rope_index: Optional[int] = 5,
405 query_rope_kwarg_name: Optional[str] = "query_rope",
406 key_rope_kwarg_name: Optional[str] = "key_rope",
407 use_local_output: bool = False,
408 ) -> None:
409 super().__init__()
410 _configure_sparse_attention_boundary(
411 self,
412 layout=layout,
413 mode=mode,
414 query_index=query_index,
415 key_index=key_index,
416 value_index=value_index,
417 topk_index=topk_index,
418 query_kwarg_name=query_kwarg_name,
419 key_kwarg_name=key_kwarg_name,
420 value_kwarg_name=value_kwarg_name,
421 topk_kwarg_name=topk_kwarg_name,
422 query_rope_index=query_rope_index,
423 key_rope_index=key_rope_index,
424 query_rope_kwarg_name=query_rope_kwarg_name,
425 key_rope_kwarg_name=key_rope_kwarg_name,
426 use_local_output=use_local_output,
427 )
429 def __repr__(self) -> str:
430 return (
431 f"{self.__class__.__name__}("
432 f"layout={self.layout!r}, mode={self.mode!r}, "
433 f"use_local_output={self.use_local_output})"
434 )
436 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module:
437 """Register DSA sparse-attention CP hooks on ``module`` and return it."""
438 return _apply_sparse_attention_boundary(self, module, device_mesh)
441class DSAIndexerLossContextParallel(ParallelStyle):
442 """Colossal-style CP hook for a DSA indexer-loss kernel boundary.
444 This style targets a hookable module/cell whose forward signature is shaped
445 like the sparse indexer-loss variant by default:
446 ``(query, key, query_index, key_index, weights, topk_indices,
447 softmax_max, softmax_sum, query_rope, key_rope, ...)``. Set
448 ``loss_variant="dense"`` for the dense loss signature:
449 ``(query, key, query_index, key_index, weights, softmax_max, softmax_sum,
450 softmax_max_index, softmax_sum_index, scale_value, query_rope, key_rope,
451 ...)``. The boundary is expected to start after MF has already done
452 local-only bookkeeping such as ``stop_gradient`` and ``split``.
454 Placements:
455 - ``query``, ``query_index``, ``weights``, ``topk_indices`` and
456 ``query_rope`` are annotated as ``Shard(seq)``;
457 - ``softmax_max``, ``softmax_sum`` and dense index softmax stats are
458 annotated as query-side stats sharded on their stats sequence dimension;
459 - ``key``, ``key_index`` and ``key_rope`` are all-gathered to
460 ``Replicate()``.
461 """
463 def __init__( # pylint: disable=too-many-arguments
464 self,
465 *,
466 layout: str = "BSND",
467 mode: str = "colossal",
468 loss_variant: str = "sparse",
469 query_index: Optional[int] = 0,
470 key_index: Optional[int] = 1,
471 query_indexer_index: Optional[int] = 2,
472 key_indexer_index: Optional[int] = 3,
473 weights_index: Optional[int] = 4,
474 topk_index: Optional[int] = _DEFAULT_ARG_INDEX,
475 softmax_max_index: Optional[int] = _DEFAULT_ARG_INDEX,
476 softmax_sum_index: Optional[int] = _DEFAULT_ARG_INDEX,
477 softmax_max_indexer_index: Optional[int] = _DEFAULT_ARG_INDEX,
478 softmax_sum_indexer_index: Optional[int] = _DEFAULT_ARG_INDEX,
479 query_rope_index: Optional[int] = _DEFAULT_ARG_INDEX,
480 key_rope_index: Optional[int] = _DEFAULT_ARG_INDEX,
481 query_kwarg_name: Optional[str] = None,
482 key_kwarg_name: Optional[str] = None,
483 query_indexer_kwarg_name: Optional[str] = None,
484 key_indexer_kwarg_name: Optional[str] = None,
485 weights_kwarg_name: Optional[str] = None,
486 topk_kwarg_name: Optional[str] = None,
487 softmax_max_kwarg_name: Optional[str] = None,
488 softmax_sum_kwarg_name: Optional[str] = None,
489 softmax_max_indexer_kwarg_name: Optional[str] = None,
490 softmax_sum_indexer_kwarg_name: Optional[str] = None,
491 query_rope_kwarg_name: Optional[str] = None,
492 key_rope_kwarg_name: Optional[str] = None,
493 use_local_output: bool = False,
494 ) -> None:
495 super().__init__()
496 layout, seq_dim = _validate_layout_and_mode(self.__class__.__name__, layout, mode)
497 loss_variant = _validate_loss_variant(loss_variant)
498 is_dense = loss_variant == "dense"
499 self.layout = layout
500 self.mode = mode
501 self.loss_variant = loss_variant
502 self.seq_dim = seq_dim
503 self.stats_seq_dim = _query_stats_seq_dim(layout)
504 self.query_index = query_index
505 self.key_index = key_index
506 self.query_indexer_index = query_indexer_index
507 self.key_indexer_index = key_indexer_index
508 self.weights_index = weights_index
509 self.topk_index = _default_arg_index(topk_index, None if is_dense else 5)
510 self.softmax_max_index = _default_arg_index(softmax_max_index, 5 if is_dense else 6)
511 self.softmax_sum_index = _default_arg_index(softmax_sum_index, 6 if is_dense else 7)
512 self.softmax_max_indexer_index = _default_arg_index(softmax_max_indexer_index, 7 if is_dense else None)
513 self.softmax_sum_indexer_index = _default_arg_index(softmax_sum_indexer_index, 8 if is_dense else None)
514 self.query_rope_index = _default_arg_index(query_rope_index, 10 if is_dense else 8)
515 self.key_rope_index = _default_arg_index(key_rope_index, 11 if is_dense else 9)
516 self.query_kwarg_name = query_kwarg_name
517 self.key_kwarg_name = key_kwarg_name
518 self.query_indexer_kwarg_name = query_indexer_kwarg_name
519 self.key_indexer_kwarg_name = key_indexer_kwarg_name
520 self.weights_kwarg_name = weights_kwarg_name
521 self.topk_kwarg_name = topk_kwarg_name
522 self.softmax_max_kwarg_name = softmax_max_kwarg_name
523 self.softmax_sum_kwarg_name = softmax_sum_kwarg_name
524 self.softmax_max_indexer_kwarg_name = softmax_max_indexer_kwarg_name
525 self.softmax_sum_indexer_kwarg_name = softmax_sum_indexer_kwarg_name
526 self.query_rope_kwarg_name = query_rope_kwarg_name
527 self.key_rope_kwarg_name = key_rope_kwarg_name
528 self.use_local_output = use_local_output
530 def __repr__(self) -> str:
531 return (
532 f"{self.__class__.__name__}("
533 f"layout={self.layout!r}, mode={self.mode!r}, "
534 f"loss_variant={self.loss_variant!r}, "
535 f"use_local_output={self.use_local_output})"
536 )
538 def _shard_query_side(self, value: Any, device_mesh: DeviceMesh) -> Any:
539 return _to_sequence_shard(value, device_mesh, self.seq_dim)
541 def _replicate_key_side(self, value: Any, device_mesh: DeviceMesh) -> Any:
542 return _to_sequence_replicate(value, device_mesh, self.seq_dim)
544 @staticmethod
545 def _local_shape(value: Any) -> Optional[tuple]:
546 if isinstance(value, DTensor):
547 return value.local_shape
548 if platform.is_tensor(value):
549 return value.shape
550 return None
552 def _slice_local_key_grad(self, value: Any, module: Module) -> Any:
553 """Convert d_key_index to the original local key-index shard shape."""
554 if not self.use_local_output:
555 return value
556 target_shape = getattr(module, "_hp_dsa_loss_key_index_local_shape", None)
557 if target_shape is None:
558 return _finalize_output(value, use_local_output=True)
560 if isinstance(value, DTensor):
561 value = _dtensor_to_local_reducing_partial(value)
562 if not platform.is_tensor(value):
563 return value
565 target_len = target_shape[self.seq_dim]
566 if value.shape[self.seq_dim] == target_len:
567 return value
569 local_idx = getattr(module, "_hp_dsa_loss_local_idx", 0)
570 start = local_idx * target_len
571 return value.narrow(self.seq_dim, start, target_len)
573 def _process_outputs(self, module: Module, outputs: Any) -> Any:
574 """Finalize indexer-loss outputs, reducing Partial values when needed."""
575 output_layout = _pop_output_layout(module)
576 if not self.use_local_output:
577 return _finalize_output(outputs, False, output_layout, self.seq_dim)
578 if not isinstance(outputs, (tuple, list)) or len(outputs) < 4:
579 return _finalize_output(outputs, use_local_output=True)
581 processed = list(outputs)
582 processed[0] = _dtensor_to_local_reducing_partial(processed[0])
583 processed[1] = self._slice_local_key_grad(processed[1], module)
584 processed[2] = _dtensor_to_local_reducing_partial(processed[2])
585 processed[3] = _dtensor_to_local_reducing_partial(processed[3])
586 return type(outputs)(processed)
588 def _build_loss_specs(
589 self,
590 cp_mesh: DeviceMesh,
591 replicate_fn_map: dict[str, Callable[[Any], Any]],
592 ) -> list[_ParamSpec]:
593 """Build table-driven transforms for the indexer-loss boundary."""
594 def shard(value: Any) -> Any:
595 return self._shard_query_side(value, cp_mesh)
597 def stats_shard(value: Any) -> Any:
598 return _to_query_stats_shard(value, cp_mesh, self.stats_seq_dim)
600 return [
601 _ParamSpec(self.query_index, self.query_kwarg_name, shard),
602 _ParamSpec(self.key_index, self.key_kwarg_name, replicate_fn_map["key"]),
603 _ParamSpec(self.query_indexer_index, self.query_indexer_kwarg_name, shard),
604 _ParamSpec(self.key_indexer_index, self.key_indexer_kwarg_name, replicate_fn_map["key_indexer"]),
605 _ParamSpec(self.weights_index, self.weights_kwarg_name, shard),
606 _ParamSpec(self.topk_index, self.topk_kwarg_name, shard),
607 _ParamSpec(self.softmax_max_index, self.softmax_max_kwarg_name, stats_shard),
608 _ParamSpec(self.softmax_sum_index, self.softmax_sum_kwarg_name, stats_shard),
609 _ParamSpec(self.softmax_max_indexer_index, self.softmax_max_indexer_kwarg_name, stats_shard),
610 _ParamSpec(self.softmax_sum_indexer_index, self.softmax_sum_indexer_kwarg_name, stats_shard),
611 _ParamSpec(self.query_rope_index, self.query_rope_kwarg_name, shard),
612 _ParamSpec(self.key_rope_index, self.key_rope_kwarg_name, replicate_fn_map["key_rope"]),
613 ]
615 def _read_key_indexer_shape(self, args: list, kwargs: dict) -> Optional[tuple]:
616 """Read the original local key-indexer shape before hook transforms."""
617 if self.key_indexer_index is not None and self.key_indexer_index < len(args):
618 return self._local_shape(args[self.key_indexer_index])
619 if self.key_indexer_kwarg_name and self.key_indexer_kwarg_name in kwargs:
620 return self._local_shape(kwargs[self.key_indexer_kwarg_name])
621 return None
623 @staticmethod
624 def _get_local_idx(cp_mesh: DeviceMesh) -> int:
625 """Return current rank's index in the CP mesh rank list."""
626 rank_list = list(cp_mesh.rank_list)
627 rank = platform.get_rank()
628 return rank_list.index(rank) if rank in rank_list else 0
630 def _apply_with_loss_specs(
631 self,
632 module: Module,
633 specs: list[_ParamSpec],
634 local_idx: int,
635 cp_mesh: DeviceMesh,
636 ) -> Module:
637 """Register indexer-loss hooks driven by parameter specs."""
638 def _pre_hook(hook_module, args, kwargs):
639 new_args = list(args)
640 new_kwargs = dict(kwargs)
641 _record_query_output_layout(
642 hook_module,
643 _read_value(new_args, new_kwargs, self.query_index, self.query_kwarg_name),
644 cp_mesh,
645 self.seq_dim,
646 )
647 key_shape = self._read_key_indexer_shape(new_args, new_kwargs)
648 setattr(hook_module, "_hp_dsa_loss_key_index_local_shape", key_shape)
649 setattr(hook_module, "_hp_dsa_loss_local_idx", local_idx)
650 _apply_param_specs(new_args, new_kwargs, specs)
651 return tuple(new_args), new_kwargs
653 platform.register_forward_pre_hook(module, _pre_hook, with_kwargs=True)
654 module.register_forward_hook(lambda _module, _args, outputs: self._process_outputs(_module, outputs))
655 return module
657 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module:
658 """Register DSA indexer-loss CP hooks on ``module`` and return it."""
659 cp_mesh = _ensure_1d(device_mesh)
661 def replicate(value: Any) -> Any:
662 return self._replicate_key_side(value, cp_mesh)
664 specs = self._build_loss_specs(
665 cp_mesh,
666 replicate_fn_map={
667 "key": replicate,
668 "key_indexer": replicate,
669 "key_rope": replicate,
670 },
671 )
672 return self._apply_with_loss_specs(module, specs, self._get_local_idx(cp_mesh), cp_mesh)