Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / collectives.py: 31%
269 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-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# ============================================================================
16"""context_parallel.collectives: CP collective primitives (05 §4.4.2 / §6.3.4).
18- ``flex_cp_allgather``: K/V all-gather along the CP dim (reuses
19 cp_mesh.get_group(); new_group is forbidden);
20- ``ulysses_seq_to_head`` / ``ulysses_head_to_seq``: differentiable Pure
21 Ulysses all-to-all layout transforms;
22- async collective launches (``AsyncCPCollective`` /
23 ``async_cp_allgather_launch`` / ``async_ulysses_seq_to_head_launch``);
24- hybrid (Ulysses + all-gather) CP sub-mesh construction and attention
25 orchestration (``_build_hybrid_cp_submeshes`` / ``hybrid_cp_attention``).
27Note (G5): the seq_len % (2*cp) padding constraint originates from the
28zigzag/ring load-balancing scheme; this design uses all-gather K/V +
29contiguous chunk (D-01'' rejected ring), so each rank's Q chunk is
30equal-length and FLOPs are naturally balanced -- the constraint is redundant
31but harmless: the implementation is kept and documented here.
33Split out of components/distributed/cp_utils.py in stage 4e; the attention
34adaptations live in context_parallel/attention.py and the batch sharding in
35hyper_parallel/data/parallel/batch_parallel.py.
36"""
38from dataclasses import dataclass
39from typing import Any, Callable, Optional
40import torch
41import torch.distributed as dist
42from torch import Tensor
43from torch.distributed.nn.functional import all_gather as differentiable_all_gather
44from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
45from hyper_parallel.core.utils import communication
48_ULYSSES_WRAPPED_FLAG = "_hyper_ulysses_wrapped"
51_HYBRID_MESH_CACHE = {}
54class _SeqAllToAll(torch.autograd.Function):
55 """Autograd-aware exchange from one tensor dimension to another."""
57 @staticmethod
58 def forward(
59 ctx: Any,
60 group: Any,
61 tensor: Tensor,
62 scatter_dim: int,
63 gather_dim: int,
64 ) -> Tensor:
65 """Exchange tensor shards between the sequence and head dimensions."""
66 ctx.group = group
67 ctx.scatter_dim = scatter_dim
68 ctx.gather_dim = gather_dim
69 world_size = dist.get_world_size(group)
70 if tensor.size(scatter_dim) % world_size:
71 raise ValueError(
72 f"dimension {scatter_dim} size {tensor.size(scatter_dim)} "
73 f"is not divisible by CP size {world_size}")
74 inputs = [
75 part.contiguous()
76 for part in tensor.chunk(world_size, dim=scatter_dim)
77 ]
78 outputs = [torch.empty_like(inputs[0]) for _ in range(world_size)]
79 dist.all_to_all(outputs, inputs, group=group)
80 return torch.cat(outputs, dim=gather_dim).contiguous()
82 @staticmethod
83 def backward(
84 ctx: Any,
85 grad_output: Tensor,
86 ) -> tuple[None, Tensor, None, None]:
87 """Apply the inverse all-to-all to the output gradient."""
88 grad_input = _SeqAllToAll.apply(
89 ctx.group, grad_output, ctx.gather_dim, ctx.scatter_dim)
90 return None, grad_input, None, None
93@dataclass(frozen=True)
94class _UlyssesContext:
95 """Communication context derived from the framework-owned CP mesh."""
97 cp_mesh: Any
99 @property
100 def group(self) -> Any:
101 """Process group of the CP mesh."""
102 return self.cp_mesh.get_group()
104 @property
105 def size(self) -> int:
106 """Number of ranks in the CP mesh."""
107 return self.cp_mesh.size()
109 @property
110 def rank(self) -> int:
111 """Local rank of this process within the CP mesh."""
112 return self.cp_mesh.get_local_rank()
115def _sequence_to_head(tensor, context):
116 """[B, S/CP, H, D] -> [B, S, H/CP, D]."""
117 return _SeqAllToAll.apply(context.group, tensor, 2, 1)
120def _head_to_sequence(tensor, context):
121 """[B, S, H/CP, D] -> [B, S/CP, H, D]."""
122 return _SeqAllToAll.apply(context.group, tensor, 1, 2)
125def _gather_sequence(tensor, context):
126 """Gather sequence shards and sum their gradients during backward."""
127 return torch.cat(
128 differentiable_all_gather(
129 tensor.contiguous(), group=context.group), dim=1)
132def _slice_sequence(tensor, dim, context):
133 if tensor.size(dim) % context.size:
134 raise ValueError(
135 f"sequence length {tensor.size(dim)} is not divisible by "
136 f"CP size {context.size}")
137 length = tensor.size(dim) // context.size
138 return tensor.narrow(
139 dim, context.rank * length, length).contiguous()
142def _global_seq_len(actual, length, device):
143 if actual is None:
144 return [length]
145 if isinstance(actual, Tensor):
146 return actual.to(device=device, dtype=torch.int32)
147 return actual
150def _slice_sink(tensor, context):
151 if tensor is None or tensor.size(2) == 1:
152 return tensor
153 if tensor.size(2) % context.size:
154 raise ValueError(
155 f"parameter-sink heads {tensor.size(2)} are not divisible by "
156 f"CP size {context.size}")
157 count = tensor.size(2) // context.size
158 start = context.rank * count
159 return tensor[:, :, start:start + count].contiguous()
162def _normalize_dim(dim: int, ndim: int) -> int:
163 """Normalize and validate a tensor dimension."""
164 normalized = dim + ndim if dim < 0 else dim
165 if normalized < 0 or normalized >= ndim:
166 raise ValueError(f"dimension {dim} is out of range for a {ndim}D tensor")
167 return normalized
170def _reconstruct_all_to_all(output: Tensor, concat_dim: int) -> Tensor:
171 """Move the source-rank axis next to ``concat_dim`` and merge it."""
172 rank_axis_position = concat_dim + 1
173 permutation = (
174 list(range(1, rank_axis_position))
175 + [0]
176 + list(range(rank_axis_position, output.dim()))
177 )
178 reconstructed = output.permute(permutation).contiguous()
179 shape = list(reconstructed.shape)
180 merged = shape[concat_dim] * shape[concat_dim + 1]
181 return reconstructed.reshape(shape[:concat_dim] + [merged] + shape[concat_dim + 2:])
184def _move_dim_to_front(tensor: Tensor, dim: int) -> Tensor:
185 """Move one dimension to the leading communication dimension."""
186 dim = _normalize_dim(dim, tensor.dim())
187 if dim == 0:
188 return tensor.contiguous()
189 return tensor.movedim(dim, 0).contiguous()
192def _move_dim_from_front(tensor: Tensor, dim: int) -> Tensor:
193 """Restore a leading communication dimension to its original position."""
194 dim = _normalize_dim(dim, tensor.dim())
195 if dim == 0:
196 return tensor.contiguous()
197 return tensor.movedim(0, dim).contiguous()
200def _prepare_ulysses_send(
201 tensor: Tensor,
202 scatter_dim: int,
203 world_size: int,
204) -> Tensor:
205 """Split one tensor dimension and move the destination-rank axis first."""
206 ndim = tensor.dim()
207 scatter_dim = _normalize_dim(scatter_dim, ndim)
208 shape = list(tensor.shape)
209 scatter_size = shape[scatter_dim]
210 if scatter_size % world_size:
211 raise ValueError(
212 f"tensor dimension {scatter_dim} size ({scatter_size}) must be "
213 f"divisible by Ulysses degree ({world_size})"
214 )
215 split_shape = (
216 shape[:scatter_dim]
217 + [world_size, scatter_size // world_size]
218 + shape[scatter_dim + 1:]
219 )
220 permutation = (
221 [scatter_dim]
222 + list(range(scatter_dim))
223 + list(range(scatter_dim + 1, ndim + 1))
224 )
225 return tensor.contiguous().reshape(split_shape).permute(permutation).contiguous()
228class _AsyncAllGatherWait(torch.autograd.Function):
229 """Attach a pre-launched all-gather to autograd at its wait point."""
231 @staticmethod
232 def forward(
233 ctx: Any,
234 tensor: Tensor,
235 work: Any,
236 output: Tensor,
237 group: Any,
238 world_size: int,
239 gather_dim: int,
240 ) -> Tensor:
241 """Wait for a launched all-gather and expose its autograd edge."""
242 del tensor
243 ctx.group = group
244 ctx.world_size = world_size
245 ctx.gather_dim = gather_dim
246 work.wait()
247 return _move_dim_from_front(output, gather_dim)
249 @staticmethod
250 def backward(
251 ctx: Any,
252 grad_output: Tensor,
253 ) -> tuple[Tensor, None, None, None, None, None]:
254 """Reduce and select this rank's input-gradient sequence shard."""
255 grad = grad_output.contiguous().clone()
256 work = dist.all_reduce(grad, group=ctx.group, async_op=True)
257 work.wait()
258 rank = dist.get_rank(ctx.group)
259 local = torch.chunk(
260 grad, ctx.world_size, dim=ctx.gather_dim
261 )[rank].contiguous()
262 return (
263 local,
264 None,
265 None,
266 None,
267 None,
268 None,
269 )
272class _AsyncUlyssesWait(torch.autograd.Function):
273 """Attach a pre-launched sequence-to-head A2A to autograd."""
275 @staticmethod
276 def forward(
277 ctx: Any,
278 tensor: Tensor,
279 work: Any,
280 output: Tensor,
281 group: Any,
282 world_size: int,
283 seq_dim: int,
284 head_dim: int,
285 ) -> Tensor:
286 """Wait for a launched A2A and expose its autograd edge."""
287 del tensor
288 ctx.group = group
289 ctx.world_size = world_size
290 ctx.seq_dim = seq_dim
291 ctx.head_dim = head_dim
292 work.wait()
293 return _reconstruct_all_to_all(output, seq_dim)
295 @staticmethod
296 def backward(
297 ctx: Any,
298 grad_output: Tensor,
299 ) -> tuple[Tensor, None, None, None, None, None, None]:
300 """Apply the inverse A2A to restore the input-gradient layout."""
301 send = _prepare_ulysses_send(
302 grad_output,
303 scatter_dim=ctx.seq_dim,
304 world_size=ctx.world_size,
305 )
306 output = torch.empty(
307 list(send.shape), device=send.device, dtype=send.dtype
308 )
309 work = dist.all_to_all_single(
310 output, send, group=ctx.group, async_op=True
311 )
312 work.wait()
313 grad_input = _reconstruct_all_to_all(output, ctx.head_dim)
314 return grad_input, None, None, None, None, None, None
317@dataclass
318class AsyncCPCollective:
319 """A local-tensor CP collective launched before its consumer."""
321 tensor: Tensor
322 work: Any
323 output: Tensor
324 group: Any
325 world_size: int
326 kind: str
327 seq_dim: int
328 head_dim: Optional[int] = None
329 waited: bool = False
331 def wait(self) -> Tensor:
332 """Materialize the collective result exactly once."""
333 if self.waited:
334 raise RuntimeError("an async CP collective handle cannot be waited twice")
335 self.waited = True
336 if self.work is None:
337 return self.tensor
338 if self.kind == "allgather":
339 return _AsyncAllGatherWait.apply(
340 self.tensor,
341 self.work,
342 self.output,
343 self.group,
344 self.world_size,
345 self.seq_dim,
346 )
347 if self.kind == "ulysses":
348 return _AsyncUlyssesWait.apply(
349 self.tensor,
350 self.work,
351 self.output,
352 self.group,
353 self.world_size,
354 self.seq_dim,
355 self.head_dim,
356 )
357 raise RuntimeError(f"unknown async CP collective kind {self.kind!r}")
360def async_cp_allgather_launch(
361 tensor: Tensor,
362 gather_dim: int,
363 cp_mesh: DeviceMesh,
364) -> AsyncCPCollective:
365 """Launch an asynchronous all-gather along gather_dim."""
366 world_size = cp_mesh.size()
367 if world_size <= 1:
368 return AsyncCPCollective(
369 tensor, None, tensor, None, world_size, "allgather", gather_dim
370 )
371 gather_dim = _normalize_dim(gather_dim, tensor.dim())
372 send = _move_dim_to_front(tensor.detach(), gather_dim)
373 output_shape = list(send.shape)
374 output_shape[0] *= world_size
375 group = cp_mesh.get_group()
376 output = torch.empty(
377 output_shape, device=send.device, dtype=send.dtype
378 )
379 work = dist.all_gather_into_tensor(
380 output, send, group=group, async_op=True
381 )
382 return AsyncCPCollective(
383 tensor, work, output, group, world_size, "allgather", gather_dim
384 )
387def async_ulysses_seq_to_head_launch(
388 tensor: Tensor,
389 seq_dim: int,
390 head_dim: int,
391 cp_mesh: DeviceMesh,
392) -> AsyncCPCollective:
393 """Launch an asynchronous Ulysses sequence-to-head A2A."""
394 world_size = cp_mesh.size()
395 if world_size <= 1:
396 return AsyncCPCollective(
397 tensor,
398 None,
399 tensor,
400 None,
401 world_size,
402 "ulysses",
403 seq_dim,
404 head_dim,
405 )
406 seq_dim = _normalize_dim(seq_dim, tensor.dim())
407 head_dim = _normalize_dim(head_dim, tensor.dim())
408 if seq_dim == head_dim:
409 raise ValueError("Ulysses seq_dim and head_dim must be different")
410 send = _prepare_ulysses_send(
411 tensor.detach(), scatter_dim=head_dim, world_size=world_size
412 )
413 group = cp_mesh.get_group()
414 output = torch.empty(
415 list(send.shape), device=send.device, dtype=send.dtype
416 )
417 work = dist.all_to_all_single(
418 output, send, group=group, async_op=True
419 )
420 return AsyncCPCollective(
421 tensor,
422 work,
423 output,
424 group,
425 world_size,
426 "ulysses",
427 seq_dim,
428 head_dim,
429 )
432def _ulysses_all_to_all(
433 tensor: Tensor,
434 *,
435 scatter_dim: int,
436 gather_dim: int,
437 cp_mesh: DeviceMesh,
438) -> Tensor:
439 """Split one tensor dimension and gather another through all-to-all."""
440 cp_size = cp_mesh.size()
441 if cp_size <= 1:
442 return tensor
444 ndim = tensor.dim()
445 scatter_dim = _normalize_dim(scatter_dim, ndim)
446 gather_dim = _normalize_dim(gather_dim, ndim)
447 if scatter_dim == gather_dim:
448 raise ValueError("Ulysses scatter_dim and gather_dim must be different")
450 shape = list(tensor.shape)
451 scatter_size = shape[scatter_dim]
452 if scatter_size % cp_size != 0:
453 raise ValueError(
454 f"tensor dimension {scatter_dim} size ({scatter_size}) must be "
455 f"divisible by Ulysses degree ({cp_size})"
456 )
458 split_shape = (
459 shape[:scatter_dim]
460 + [cp_size, scatter_size // cp_size]
461 + shape[scatter_dim + 1:]
462 )
463 split_ndim = ndim + 1
464 permutation = (
465 [scatter_dim]
466 + list(range(scatter_dim))
467 + list(range(scatter_dim + 1, split_ndim))
468 )
469 send = tensor.contiguous().reshape(split_shape).permute(permutation).contiguous()
470 received = communication.differentiable_all_to_all(
471 send, list(send.shape), cp_mesh.get_group())
472 return _reconstruct_all_to_all(received, gather_dim)
475def ulysses_seq_to_head(
476 tensor: Tensor,
477 seq_dim: int,
478 head_dim: int,
479 cp_mesh: DeviceMesh,
480) -> Tensor:
481 """Convert a sequence shard into a full-sequence head shard.
483 Args:
484 tensor: Local Q, K, or V tensor.
485 seq_dim: Sequence dimension to gather across the CP group.
486 head_dim: Head dimension to split across the CP group.
487 cp_mesh: One-dimensional CP DeviceMesh.
489 Returns:
490 Tensor with the sequence dimension gathered and head dimension sharded.
491 """
492 return _ulysses_all_to_all(
493 tensor,
494 scatter_dim=head_dim,
495 gather_dim=seq_dim,
496 cp_mesh=cp_mesh,
497 )
500def ulysses_head_to_seq(
501 tensor: Tensor,
502 seq_dim: int,
503 head_dim: int,
504 cp_mesh: DeviceMesh,
505) -> Tensor:
506 """Convert a full-sequence head shard back into a sequence shard.
508 Args:
509 tensor: Local attention output in Ulysses head-sharded layout.
510 seq_dim: Full sequence dimension to split across the CP group.
511 head_dim: Head dimension to gather across the CP group.
512 cp_mesh: One-dimensional CP DeviceMesh.
514 Returns:
515 Tensor restored to the original local sequence-sharded layout.
516 """
517 return _ulysses_all_to_all(
518 tensor,
519 scatter_dim=seq_dim,
520 gather_dim=head_dim,
521 cp_mesh=cp_mesh,
522 )
525def _validate_ulysses_head_counts(
526 query: Tensor, key: Tensor, value: Tensor,
527 ulysses_degree: int, context: str) -> None:
528 """Require every Q/K/V head dimension to be shardable by Ulysses."""
529 head_counts = {
530 "query": query.shape[1],
531 "key": key.shape[1],
532 "value": value.shape[1],
533 }
534 invalid = {
535 name: count
536 for name, count in head_counts.items()
537 if count % ulysses_degree
538 }
539 if invalid:
540 details = ", ".join(
541 f"{name}_heads={count}" for name, count in invalid.items()
542 )
543 raise ValueError(
544 f"{context} requires every Q/K/V head count to be divisible by "
545 f"ulysses_degree ({ulysses_degree}), got {details}"
546 )
549class _AllGatherAlongDim(torch.autograd.Function):
550 """all-gather along cp_dim + backward reduce-scatter semantics (sum across
551 ranks, then take this rank's chunk)."""
553 @staticmethod
554 def forward(
555 ctx: Any,
556 t: Tensor,
557 cp_dim: int,
558 group: Any,
559 cp_size: int,
560 ) -> Tensor:
561 """All-gather the tensor along cp_dim in CP-rank order."""
562 ctx.cp_dim = cp_dim
563 ctx.group = group
564 ctx.cp_size = cp_size
565 world_t = [torch.empty_like(t) for _ in range(cp_size)]
566 dist.all_gather(world_t, t.contiguous(), group=group)
567 # cat in cp_rank order: [chunk_rank0, chunk_rank1, ...]
568 return torch.cat(world_t, dim=cp_dim)
570 @staticmethod
571 def backward(
572 ctx: Any,
573 grad_output: Tensor,
574 ) -> tuple[Tensor, None, None, None]:
575 """Sum the gradient across ranks and take this rank's chunk."""
576 # reduce-scatter: sum the gradient across ranks, take this rank's chunk
577 grad = grad_output.contiguous().clone()
578 dist.all_reduce(grad, group=ctx.group)
579 rank = dist.get_rank(ctx.group)
580 local = torch.chunk(grad, ctx.cp_size, dim=ctx.cp_dim)[rank]
581 return local.contiguous(), None, None, None
584def flex_cp_allgather(
585 k: Tensor,
586 v: Tensor,
587 cp_dim: int,
588 cp_mesh: DeviceMesh,
589) -> tuple[Tensor, Tensor]:
590 """All-gather K/V along CP dimension for context parallel attention.
592 Forward: all-gather K/V along cp_dim (each rank ends up holding the full K/V).
593 Backward: reduce-scatter semantics (gradients summed across ranks, then this
594 rank's chunk is taken -- implemented explicitly by the _AllGatherAlongDim
595 autograd.Function, since plain ``dist.all_gather`` has no autograd kernel).
597 Args:
598 k, v: [B, N, S_local, H] (cp_dim=2 is the sequence dim).
599 cp_dim: the gather dimension.
600 cp_mesh: DeviceMesh of the CP dim. The communication group is taken from
601 ``cp_mesh.get_group()`` -- already created and cached at DeviceMesh
602 construction; **dist.new_group must NOT be called here** (otherwise
603 every forward leaks one process group, and the semantics would be
604 misaligned).
605 """
606 cp_size = cp_mesh.size()
607 if cp_size <= 1:
608 return k, v
609 group = cp_mesh.get_group()
610 return (_AllGatherAlongDim.apply(k, cp_dim, group, cp_size),
611 _AllGatherAlongDim.apply(v, cp_dim, group, cp_size))
614def _build_hybrid_cp_submeshes(cp_mesh, ulysses_degree: int):
615 """Build cached communication submeshes for local-tensor Hybrid CP."""
616 if cp_mesh is None or cp_mesh.size() <= 1:
617 raise ValueError("Hybrid CP requires an active CP mesh")
618 if isinstance(ulysses_degree, bool) or not isinstance(ulysses_degree, int):
619 raise TypeError(
620 "Hybrid CP requires integer ulysses_degree, got "
621 f"{type(ulysses_degree).__name__}"
622 )
623 cp_size = cp_mesh.size()
624 if not 1 < ulysses_degree < cp_size:
625 raise ValueError(
626 "Hybrid CP requires 1 < ulysses_degree < cp_size, got "
627 f"ulysses_degree={ulysses_degree}, cp_size={cp_size}"
628 )
629 if cp_size % ulysses_degree:
630 raise ValueError(
631 f"cp_size ({cp_size}) must be divisible by ulysses_degree "
632 f"({ulysses_degree})"
633 )
635 cache_key = (cp_mesh.to_hash(), ulysses_degree)
636 cached = _HYBRID_MESH_CACHE.get(cache_key)
637 if cached is not None:
638 return cached
640 colossal_degree = cp_size // ulysses_degree
641 if cp_mesh.ndim == 2:
642 hybrid_mesh = cp_mesh
643 elif (cp_mesh.root_mesh is not None
644 and cp_mesh.mesh_dim_names == ("cp",)):
645 hybrid_mesh = cp_mesh._unflatten( # pylint: disable=protected-access
646 "cp",
647 (colossal_degree, ulysses_degree),
648 ("colossal", "ulysses"),
649 )
650 else:
651 ranks = list(cp_mesh.rank_list)
652 hybrid_mesh = DeviceMesh(
653 cp_mesh.device_type,
654 [
655 ranks[index * ulysses_degree:(index + 1) * ulysses_degree]
656 for index in range(colossal_degree)
657 ],
658 mesh_dim_names=("colossal", "ulysses"),
659 )
660 dim_names = hybrid_mesh.mesh_dim_names
661 if not dim_names or len(dim_names) != 2:
662 raise ValueError("Hybrid CP requires a named two-dimensional submesh")
663 colossal_mesh = hybrid_mesh[dim_names[0]]
664 ulysses_mesh = hybrid_mesh[dim_names[1]]
665 result = (ulysses_mesh, colossal_mesh)
666 _HYBRID_MESH_CACHE[cache_key] = result
667 return result
670def hybrid_cp_attention(
671 attention_fn: Callable[[Tensor, Tensor, Tensor, dict[str, Any]], Any],
672 query: Tensor, key: Tensor, value: Tensor,
673 attention_kwargs: dict[str, Any], cp_mesh: Any,
674 ulysses_degree: int) -> Any:
675 """Run local-tensor Hybrid communication around an attention callable.
677 ``attention_kwargs`` are forwarded unchanged. Mask construction and other
678 model-input semantics belong to the caller's input-preparation contract.
679 """
680 _validate_ulysses_head_counts(
681 query, key, value, ulysses_degree, "Hybrid CP"
682 )
683 ulysses_mesh, colossal_mesh = (
684 _build_hybrid_cp_submeshes(cp_mesh, ulysses_degree)
685 )
686 local_shape = tuple(query.shape)
687 query, key, value = (
688 ulysses_seq_to_head(tensor, 2, 1, ulysses_mesh)
689 for tensor in (query, key, value)
690 )
691 key, value = flex_cp_allgather(key, value, 2, colossal_mesh)
692 output = attention_fn(query, key, value, attention_kwargs)
693 output = ulysses_head_to_seq(output, 2, 1, ulysses_mesh)
694 if tuple(output.shape) != local_shape:
695 raise RuntimeError(
696 f"Hybrid CP output shape {tuple(output.shape)} does not restore "
697 f"the local query shape {local_shape}"
698 )
699 return output