Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / _npu_contracts.py: 90%
205 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""Shape contracts of the NPU kernels HyperParallel's fused modules call, for the census.
17A train.yaml's ``plan_overrides`` may replace Transformers' modules with
18HyperParallel's fused ones (``replace_module``): its RMSNorm, its grouped-query
19attention, its grouped experts. Their kernels are ``torch_npu``'s, which the
20census, running fake tensors on a host, cannot call. :func:`npu_contracts`
21stands in for ``torch_npu`` with a contract for each kernel those modules
22call: an op returning outputs of the kernel's shapes, which saves for its
23backward what the kernel's backward takes, as ``torch_npu``'s derivatives
24state it.
26- ``npu_rms_norm`` saves its input, its weight and the fp32 reciprocal of
27 each row's root mean square;
28- ``npu_swiglu`` saves its input, the gate and up projections side by side;
29- ``npu_rotary_mul`` saves its input and the two rotary tables;
30- ``npu_fusion_attention`` saves its queries, keys and values, its output
31 and its two fp32 softmax statistics per head and token; the mask it also
32 takes is one every layer shares, which no layer's census counts;
33- ``npu_moe_token_permute`` saves its tokens, their experts and the order it
34 sorts them in, and ``npu_moe_token_unpermute`` the expert outputs, that
35 order and the routing weights;
36- ``npu_grouped_matmul`` runs inside HyperParallel's own autograd functions,
37 which save what their backward takes, so its contract only shapes its
38 output.
40Any other kernel raises when called, naming it.
42A fused module may cache a constant the first time it runs, as the attention
43caches the compressed causal mask its sparse mode takes, one per device.
44Under the census's fake tensors that constant is a fake tensor, which a real
45call would then take from the cache, so leaving the contracts drops what the
46census left (:func:`_drop_fake_caches`, F55).
47"""
48from __future__ import annotations
50import contextlib
51import functools
52import importlib
53import importlib.machinery
54import sys
55import types
56from typing import Any, Iterator, List, Optional, Sequence, Tuple
58import torch # pylint: disable=forbidden-backend-import
59from torch._subclasses.fake_tensor import FakeTensor # pylint: disable=forbidden-backend-import
61# The softmax statistics the fused attention keeps per head and token.
62_STATS = 8
64# What ``sys.modules`` held for ``torch_npu`` when nothing did.
65_ABSENT = object()
68def _rms_norm_outputs(x: torch.Tensor, gamma: torch.Tensor, epsilon: float) -> Tuple[torch.Tensor, torch.Tensor]:
69 """The normalized rows, and each row's fp32 reciprocal root mean square."""
70 del gamma, epsilon
71 return torch.empty_like(x), x.new_empty((*x.shape[:-1], 1), dtype=torch.float32)
74_RMS_NORM = torch.library.custom_op(
75 "nd_census_npu::rms_norm", mutates_args=(),
76 schema="(Tensor x, Tensor gamma, float epsilon) -> (Tensor, Tensor)")(_rms_norm_outputs)
77_RMS_NORM.register_fake(_rms_norm_outputs)
80def _rms_norm_saves(ctx: Any, inputs: Tuple[Any, ...], output: Tuple[torch.Tensor, ...]) -> None:
81 """Keep the input, the weight and the reciprocal root mean squares."""
82 ctx.mark_non_differentiable(output[1])
83 ctx.save_for_backward(inputs[0], inputs[1], output[1])
86def _rms_norm_backward(ctx: Any, *grads: torch.Tensor) -> Tuple[Optional[torch.Tensor], ...]:
87 """Gradients of the input's and the weight's shapes."""
88 del grads
89 x, gamma, _ = ctx.saved_tensors
90 return torch.empty_like(x), torch.empty_like(gamma), None
93_RMS_NORM.register_autograd(_rms_norm_backward, setup_context=_rms_norm_saves)
96def _swiglu_output(x: torch.Tensor, dim: int) -> torch.Tensor:
97 """The gated half of *x*'s width along *dim*."""
98 shape = list(x.shape)
99 shape[dim] //= 2
100 return x.new_empty(shape)
103_SWIGLU = torch.library.custom_op(
104 "nd_census_npu::swiglu", mutates_args=(), schema="(Tensor x, int dim) -> Tensor")(_swiglu_output)
105_SWIGLU.register_fake(_swiglu_output)
106_SWIGLU.register_autograd(
107 lambda ctx, grad: (torch.empty_like(ctx.saved_tensors[0]), None),
108 setup_context=lambda ctx, inputs, output: ctx.save_for_backward(inputs[0]))
111def _rotary_output(x: torch.Tensor, r1: torch.Tensor, r2: torch.Tensor) -> torch.Tensor:
112 """The rotated input, of its shape."""
113 del r1, r2
114 return torch.empty_like(x)
117_ROTARY = torch.library.custom_op(
118 "nd_census_npu::rotary_mul", mutates_args=(), schema="(Tensor x, Tensor r1, Tensor r2) -> Tensor")(_rotary_output)
119_ROTARY.register_fake(_rotary_output)
120_ROTARY.register_autograd(
121 lambda ctx, grad: tuple(torch.empty_like(tensor) for tensor in ctx.saved_tensors),
122 setup_context=lambda ctx, inputs, output: ctx.save_for_backward(*inputs))
125def _attention_outputs(query: torch.Tensor, key: torch.Tensor,
126 value: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
127 """The fused attention's output, as wide as the values, and its two fp32 statistics, BNSD."""
128 del key
129 batch, heads, seq, _ = query.shape
130 stats = [query.new_empty((batch, heads, seq, _STATS), dtype=torch.float32) for _ in range(2)]
131 return query.new_empty((batch, heads, seq, value.shape[-1])), stats[0], stats[1]
134_ATTENTION = torch.library.custom_op(
135 "nd_census_npu::fusion_attention", mutates_args=(),
136 schema="(Tensor query, Tensor key, Tensor value) -> (Tensor, Tensor, Tensor)")(_attention_outputs)
137_ATTENTION.register_fake(_attention_outputs)
140def _attention_saves(ctx: Any, inputs: Tuple[torch.Tensor, ...], output: Tuple[torch.Tensor, ...]) -> None:
141 """Keep the queries, keys and values, the output and the statistics."""
142 ctx.mark_non_differentiable(*output[1:])
143 ctx.save_for_backward(*inputs, *output)
146_ATTENTION.register_autograd(
147 lambda ctx, *grads: tuple(torch.empty_like(tensor) for tensor in ctx.saved_tensors[:3]),
148 setup_context=_attention_saves)
151def _permute_outputs(tokens: torch.Tensor, indices: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
152 """Each token once per expert it routes to, expert by expert, and the order they were sorted in."""
153 routes = indices.numel()
154 return tokens.new_empty((routes, tokens.shape[-1])), indices.new_empty((routes,), dtype=torch.int32)
157_PERMUTE = torch.library.custom_op(
158 "nd_census_npu::moe_token_permute", mutates_args=(),
159 schema="(Tensor tokens, Tensor indices) -> (Tensor, Tensor)")(_permute_outputs)
160_PERMUTE.register_fake(_permute_outputs)
163def _permute_saves(ctx: Any, inputs: Tuple[torch.Tensor, ...], output: Tuple[torch.Tensor, ...]) -> None:
164 """Keep the tokens, their experts and the sorting order."""
165 ctx.mark_non_differentiable(output[1])
166 ctx.save_for_backward(inputs[0], inputs[1], output[1])
169_PERMUTE.register_autograd(
170 lambda ctx, *grads: (torch.empty_like(ctx.saved_tensors[0]), None), setup_context=_permute_saves)
173def _unpermute_output(permuted: torch.Tensor, sorted_indices: torch.Tensor,
174 probs: Optional[torch.Tensor]) -> torch.Tensor:
175 """Each token's routes combined back into one row per token."""
176 tokens = probs.shape[0] if probs is not None else sorted_indices.numel()
177 return permuted.new_empty((tokens, permuted.shape[-1]))
180_UNPERMUTE = torch.library.custom_op(
181 "nd_census_npu::moe_token_unpermute", mutates_args=(),
182 schema="(Tensor permuted, Tensor sorted_indices, Tensor? probs) -> Tensor")(_unpermute_output)
183_UNPERMUTE.register_fake(_unpermute_output)
186def _unpermute_saves(ctx: Any, inputs: Tuple[Any, ...], output: torch.Tensor) -> None:
187 """Keep the expert outputs, the sorting order and the routing weights."""
188 del output
189 ctx.has_probs = inputs[2] is not None
190 ctx.save_for_backward(*(tensor for tensor in inputs if tensor is not None))
193def _unpermute_backward(ctx: Any, grad: torch.Tensor) -> Tuple[Optional[torch.Tensor], ...]:
194 """Gradients of the expert outputs' and the routing weights' shapes."""
195 del grad
196 saved = ctx.saved_tensors
197 return torch.empty_like(saved[0]), None, torch.empty_like(saved[2]) if ctx.has_probs else None
200_UNPERMUTE.register_autograd(_unpermute_backward, setup_context=_unpermute_saves)
203def _v41_sparse_output(query: torch.Tensor, key_value: torch.Tensor, sparse_indices: torch.Tensor,
204 sinks: torch.Tensor, rope_head_dim: int, scale: float) -> torch.Tensor:
205 """The sparse attention's output, a row per token and head, as the kernel's caller reshapes it."""
206 del key_value, sparse_indices, sinks, rope_head_dim, scale
207 batch, heads, seq, head_dim = query.shape
208 return query.new_empty((batch, seq, heads, head_dim))
211_V41_SPARSE = torch.library.custom_op(
212 "nd_census_npu::v41_sparse_attention", mutates_args=(),
213 schema=("(Tensor query, Tensor key_value, Tensor sparse_indices, Tensor sinks, "
214 "int rope_head_dim, float scale) -> Tensor"))(_v41_sparse_output)
215_V41_SPARSE.register_fake(_v41_sparse_output)
218def _v41_sparse_saves(ctx: Any, inputs: Tuple[Any, ...], output: torch.Tensor) -> None:
219 """Keep what the kernel's own caller keeps, in the shapes it keeps them.
221 Copied from the nine tensors
222 ``_NpuSparseAttentionWithScalarSink.forward`` hands
223 ``ctx.save_for_backward``: the queries and the padded keys in the
224 kernel's token-major layout, the two auxiliary rotary coordinates that
225 carry the sink, the selected indices with the sink's column appended,
226 the combined softmax maximum and the corrected sum, the output, and the
227 sink's share of each row's mass. The statistics are one scalar a token
228 and head, not the eight-wide pair the dense fused attention keeps,
229 because that caller views them as (batch, tokens, heads).
230 """
231 # The last input, the softmax scale, shapes nothing the kernel saves.
232 query, key_value, sparse_indices, sinks, rope_head_dim = inputs[:5]
233 batch, heads, seq, head_dim = query.shape
234 keys, selected = key_value.shape[2], sparse_indices.shape[-1]
235 # The gradients go back in the shapes the caller passed, which are not
236 # the token-major ones the kernel is given.
237 ctx.gradients = tuple((tuple(tensor.shape), tensor.dtype) for tensor in (query, key_value, sinks))
238 # The caller pads the keys so every one of them is selected at most
239 # once and the sink still has a row of its own.
240 padded = max(keys + 1, selected + 2)
241 rope = int(rope_head_dim)
242 stats = [query.new_empty((batch * seq, heads, 1), dtype=torch.float32) for _ in range(2)]
243 ctx.save_for_backward(
244 query.new_empty((batch * seq, heads, head_dim)),
245 key_value.new_empty((batch * padded, 1, head_dim)),
246 query.new_empty((batch * seq, heads, rope)),
247 key_value.new_empty((batch * padded, 1, rope)),
248 sparse_indices.new_empty((batch * seq, 1, selected + 1), dtype=torch.int32),
249 stats[0], stats[1], output,
250 query.new_empty((batch, seq, heads), dtype=torch.float32),
251 )
254def _v41_sparse_backward(ctx: Any, grad: torch.Tensor) -> Tuple[Any, ...]:
255 """A gradient for the queries, the keys and the sinks, in the shapes they came in."""
256 gradients = [torch.empty(shape, dtype=dtype, device=grad.device) for shape, dtype in ctx.gradients]
257 return gradients[0], gradients[1], None, gradients[2], None, None
260_V41_SPARSE.register_autograd(_v41_sparse_backward, setup_context=_v41_sparse_saves)
263def npu_sparse_attention_with_scalar_sink(query: torch.Tensor, key_value: torch.Tensor,
264 sparse_indices: torch.Tensor, sinks: torch.Tensor,
265 rope_head_dim: int, scale: float) -> torch.Tensor:
266 """The contract of DeepSeek-V4.1's sparse attention, Omni's kernel and the sink it encodes.
268 Stands in for
269 ``components/modules/shared_compressed_dsa_attention.npu_sparse_attention_with_scalar_sink``,
270 whose kernel is ``omni_training_custom_ops``' and has no torch path a
271 run takes: the reference beside it keeps an fp32 score of one sequence
272 by another, which is not what the run keeps.
273 """
274 return _V41_SPARSE(query, key_value, sparse_indices, sinks, int(rope_head_dim), float(scale))
277def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float = 1e-6) -> Tuple[torch.Tensor, torch.Tensor]:
278 """``torch_npu.npu_rms_norm``'s contract."""
279 return _RMS_NORM(x, gamma, float(epsilon))
282def npu_swiglu(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
283 """``torch_npu.npu_swiglu``'s contract."""
284 return _SWIGLU(x, int(dim) % x.dim())
287def npu_rotary_mul(x: torch.Tensor, r1: torch.Tensor, r2: torch.Tensor, rotary_mode: str = "half") -> torch.Tensor:
288 """``torch_npu.npu_rotary_mul``'s contract."""
289 del rotary_mode
290 return _ROTARY(x, r1, r2)
293def npu_fusion_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, head_num: int,
294 input_layout: str, *args: Any, **kwargs: Any) -> Tuple[Any, ...]:
295 """``torch_npu.npu_fusion_attention``'s contract, for the BNSD and BSND layouts.
297 Returns the kernel's seven outputs: the attention's, the two statistics,
298 an empty softmax output, and the dropout's seed, offset and count.
300 Raises:
301 NotImplementedError: For another layout.
302 """
303 del head_num, args, kwargs
304 if input_layout not in ("BNSD", "BSND"):
305 raise NotImplementedError(f"the census has no contract for npu_fusion_attention's {input_layout} layout")
306 bnsd = input_layout == "BNSD"
307 tensors = [tensor if bnsd else tensor.transpose(1, 2) for tensor in (query, key, value)]
308 output, softmax_max, softmax_sum = _ATTENTION(*tensors)
309 return (output if bnsd else output.transpose(1, 2), softmax_max, softmax_sum,
310 query.new_empty((0,)), 0, 0, 0)
313def npu_moe_token_permute(tokens: torch.Tensor, indices: torch.Tensor, num_out_tokens: Optional[int] = None,
314 padded_mode: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
315 """``torch_npu.npu_moe_token_permute``'s contract, every route kept."""
316 del num_out_tokens, padded_mode
317 return _PERMUTE(tokens, indices)
320def npu_moe_token_unpermute(permuted_tokens: torch.Tensor, sorted_indices: torch.Tensor,
321 probs: Optional[torch.Tensor] = None, padded_mode: bool = False,
322 restore_shape: Optional[Sequence[int]] = None) -> torch.Tensor:
323 """``torch_npu.npu_moe_token_unpermute``'s contract."""
324 del padded_mode, restore_shape
325 return _UNPERMUTE(permuted_tokens, sorted_indices, probs)
328def npu_grouped_matmul(x: Sequence[torch.Tensor], weight: Sequence[torch.Tensor], *args: Any,
329 group_list: Any = None, group_type: int = 0, **kwargs: Any) -> List[torch.Tensor]:
330 """``torch_npu.npu_grouped_matmul``'s output shape, one tensor of a single split.
332 Grouped along the rows (``group_type`` 0), each group's rows times its
333 weight; along the reduced dimension (2, a weight's gradient), one
334 product per group.
335 """
336 del args, kwargs
337 rows, weights = x[0], weight[0]
338 if group_type == 2:
339 groups = group_list.shape[0] if isinstance(group_list, torch.Tensor) else len(group_list)
340 return [rows.new_empty((groups, rows.shape[0], weights.shape[-1]))]
341 return [rows.new_empty((rows.shape[0], weights.shape[-1]))]
344def _missing(name: str) -> Any:
345 """What the stand-in answers for a kernel it has no contract for: one that raises when called."""
346 if name.startswith("_"):
347 raise AttributeError(name)
349 def kernel(*args: Any, **kwargs: Any) -> Any:
350 """Refuse the call, naming the kernel the census has no contract for."""
351 del args, kwargs
352 raise NotImplementedError(f"the census has no contract for torch_npu.{name}")
354 return kernel
357def _holders(bound: Any) -> List[types.ModuleType]:
358 """HyperParallel's modules that hold *bound* as their ``torch_npu``."""
359 return [module for name, module in list(sys.modules.items())
360 if name.startswith("hyper_parallel.") and bound is not _ABSENT
361 and getattr(module, "torch_npu", None) is bound]
364def _stand_in() -> types.ModuleType:
365 """A ``torch_npu`` module of the contracts."""
366 module = types.ModuleType("torch_npu")
367 module.__spec__ = importlib.machinery.ModuleSpec("torch_npu", None)
368 for kernel in (npu_rms_norm, npu_swiglu, npu_rotary_mul, npu_fusion_attention, npu_moe_token_permute,
369 npu_moe_token_unpermute, npu_grouped_matmul):
370 setattr(module, kernel.__name__, kernel)
371 module.__getattr__ = _missing
372 return module
375TORCH_NPU = _stand_in()
378def _drop_fake_caches() -> None:
379 """Drop the fake tensors what ran under the contracts left in the runtime's caches.
381 HyperParallel's fused modules cache a constant per device in a module of
382 their own, such as the compressed causal mask the NPU attention's sparse
383 mode takes (``models/qwen3*/adapter/runtime/attention.py``). A census fills those
384 caches with fake tensors, which a real call in the same process would take
385 from them, so every fake tensor a model module holds in a cache keyed by
386 anything is dropped, leaving the caches as the census found them.
387 """
388 for name, module in list(sys.modules.items()):
389 if not name.startswith("hyper_parallel.models") or module is None:
390 continue
391 for value in list(vars(module).values()):
392 if not isinstance(value, dict):
393 continue
394 for key in [key for key, held in value.items() if isinstance(held, FakeTensor)]:
395 del value[key]
398@contextlib.contextmanager
399def npu_contracts() -> Iterator[types.ModuleType]:
400 """``torch_npu`` as the contracts, for what runs inside.
402 A module that imports ``torch_npu`` inside, as HyperParallel's fused
403 modules do, keeps the contracts; one that had imported the real
404 ``torch_npu`` before calls the contracts until the context ends.
405 HyperParallel's model package, whose build options probe the device when
406 ``torch_npu`` imports, is imported before.
407 """
408 importlib.import_module("hyper_parallel.models")
409 previous = sys.modules.get("torch_npu", _ABSENT)
410 sys.modules["torch_npu"] = TORCH_NPU
411 for module in _holders(previous):
412 module.torch_npu = TORCH_NPU
413 try:
414 yield TORCH_NPU
415 finally:
416 _drop_fake_caches()
417 if previous is _ABSENT:
418 sys.modules.pop("torch_npu", None)
419 else:
420 sys.modules["torch_npu"] = previous
421 # Including a module that imported the stand-in while it stood in.
422 for module in _holders(TORCH_NPU):
423 module.torch_npu = previous
426# Where DeepSeek-V4.1's sparse attention lives, and the two names in it that
427# decide which path a call takes.
428_V41_ATTENTION = "hyper_parallel.components.modules.shared_compressed_dsa_attention"
429_V41_OMNI = "npu_sparse_attention_with_scalar_sink"
430_V41_REFERENCE = "_reference_sparse_attention"
433@contextlib.contextmanager
434def v41_attention_contract(rope_head_dim: int) -> Iterator[bool]:
435 """Legacy Omni sparse attention as its contract, for what runs inside.
437 A V4.1 layer picks its attention by the device it is running on
438 (``self._use_omni_attention and query.device.type == "npu"``), so a
439 census, whose tensors are fake and on no device, always takes the
440 reference beside the kernel. That reference keeps an fp32 score of one
441 sequence by another where the kernel keeps nine tensors that follow the
442 sequence, so measuring it prices a layer no run builds: on the
443 validation crop at 4096 tokens it read 77.9 GiB a rank against a 38.5
444 GiB measured peak.
446 This contract applies only while the module exposes the Omni and legacy
447 reference entry points. Newer adapters have an explicit fusion switch;
448 their eager path must be measured as written, not priced as Omni. Their
449 native kernels need their own saved-tensor contracts.
451 Both legacy names are bound to the contract for as long as the context lasts,
452 so either branch of that choice reaches it. *rope_head_dim* is the
453 width of the auxiliary rotary coordinates the kernel's caller passes and
454 the reference's own signature does not carry, so the census states it
455 from the model's config.
457 Yields:
458 Whether the legacy Omni entry points were present to patch.
459 """
460 try:
461 module = importlib.import_module(_V41_ATTENTION)
462 except ImportError:
463 yield False
464 return
466 if not all(hasattr(module, name) for name in (_V41_OMNI, _V41_REFERENCE)):
467 yield False
468 return
470 def contract(query: torch.Tensor, key_value: torch.Tensor, sparse_indices: torch.Tensor,
471 sinks: torch.Tensor, scale: float) -> torch.Tensor:
472 """The reference's signature, answered by the kernel's contract."""
473 return npu_sparse_attention_with_scalar_sink(
474 query, key_value, sparse_indices, sinks, int(rope_head_dim), float(scale))
476 held = {name: getattr(module, name) for name in (_V41_OMNI, _V41_REFERENCE)}
477 setattr(module, _V41_REFERENCE, contract)
478 setattr(module, _V41_OMNI, npu_sparse_attention_with_scalar_sink)
479 try:
480 yield True
481 finally:
482 for name, value in held.items():
483 setattr(module, name, value)
486# The Lightning Indexer's selection functions, which return the chosen
487# indices and nothing else. One of the four already states that it needs no
488# autograd graph; the other three do not, which is the whole of this contract.
489_V41_SELECTION = (
490 "compressed_causal_topk",
491 "compressed_causal_candidates",
492 "compressed_causal_topk_and_candidates",
493 "compressed_candidate_topk",
494)
497@contextlib.contextmanager
498def v41_indexer_contract() -> Iterator[bool]:
499 """DeepSeek-V4.1's Indexer selection as its fused operator keeps it, which is nothing.
501 The selection returns int32 indices, and the Indexer's own KL loss
502 differentiates itself by hand, under ``no_grad``, from the projections
503 rather than through the selection. So no gradient flows through the
504 scores, and the fused operator keeps nothing for a backward. Running the
505 selection without recording a graph states exactly that, and needs no
506 shape declared, because ``no_grad`` changes whether a graph is kept and
507 never what is computed: the indices are the same indices.
509 It holds on either path. The reference selection records a graph behind
510 its fp32 scores, but that graph ends at the int32 indices, so a run frees
511 it as soon as they are taken and holds the same bytes whichever path it
512 takes. The census once counted that graph, 8.9 GiB a rank on the
513 validation crop at 4096 tokens across the eight layers that carry an
514 Indexer, because its own pack hook held every graph a forward drops
515 alive (D16, fixed in ``_measure``). Measured since on that crop, the
516 reference path holds 26.542 GiB a rank at 4096 tokens with this contract
517 and without it, so it changes no figure the census reports and only
518 spares the census recording a graph no run keeps.
520 Of the four, ``compressed_causal_topk`` already carries
521 ``@torch.no_grad()``. The other three record the graph and drop it,
522 which costs a run nothing it holds once the selection returns.
524 Yields:
525 Whether the selection was there to wrap, which is False on a tree
526 that carries no V4.1 attention.
527 """
528 try:
529 module = importlib.import_module(_V41_ATTENTION)
530 except ImportError:
531 yield False
532 return
534 def without_a_graph(selection: Any) -> Any:
535 """*selection*, computing the same indices without recording a graph."""
536 @functools.wraps(selection)
537 def run(*args: Any, **kwargs: Any) -> Any:
538 """The selection's own indices, with no graph kept behind them."""
539 with torch.no_grad():
540 return selection(*args, **kwargs)
541 return run
543 held = {name: getattr(module, name) for name in _V41_SELECTION if hasattr(module, name)}
544 for name, selection in held.items():
545 setattr(module, name, without_a_graph(selection))
546 try:
547 yield bool(held)
548 finally:
549 for name, selection in held.items():
550 setattr(module, name, selection)