Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / encoder_parallel / attention.py: 97%

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"""Encoder attention segments and exact execution auditing.""" 

16 

17from __future__ import annotations 

18 

19from contextlib import contextmanager 

20from contextvars import ContextVar 

21from dataclasses import dataclass, field 

22from typing import Iterator, Mapping 

23 

24from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

25from hyper_parallel.distributed.encoder_parallel.contracts import EncoderSegment 

26 

27 

28@dataclass 

29class EncoderAttentionContext: 

30 """One ViT pass: segment boundaries, Ulysses mesh, and call audit.""" 

31 

32 segments: tuple[EncoderSegment, ...] 

33 cp_mesh: DeviceMesh 

34 declared_heads: Mapping[str, int] 

35 expected: frozenset[str] 

36 seen: dict[str, int] = field(default_factory=dict) 

37 

38 def record(self, module_fqn: str) -> None: 

39 """Record one replacement invocation. 

40 

41 Args: 

42 module_fqn: Final module path for the attention block. 

43 """ 

44 self.seen[module_fqn] = self.seen.get(module_fqn, 0) + 1 

45 

46 def assert_exactly_once(self) -> None: 

47 """Reject missing or repeated calls in an active ViT pass.""" 

48 if set(self.seen) != self.expected or any(count != 1 for count in self.seen.values()): 

49 raise RuntimeError( 

50 "vision attention execution mismatch; " 

51 f"expected={sorted(self.expected)}; seen={self.seen}" 

52 ) 

53 

54 def validate_runtime_head_counts(self, *, module_fqn: str, query_heads: int, 

55 key_heads: int, value_heads: int) -> None: 

56 """Check projected head counts before entering a collective.""" 

57 actual = {"query": query_heads, "key": key_heads, "value": value_heads} 

58 if actual != self.declared_heads or any(count % self.cp_mesh.size() for count in actual.values()): 

59 raise ValueError( 

60 "Ulysses runtime head count mismatch; phase=runtime; " 

61 f"module_fqn={module_fqn}; declared_head_counts={self.declared_heads}; " 

62 f"runtime_head_counts={actual}; encoder_cp_size={self.cp_mesh.size()}" 

63 ) 

64 

65 

66_ACTIVE_CONTEXT: ContextVar[EncoderAttentionContext | None] = ContextVar("encoder_attention_context", default=None) 

67 

68 

69def current_encoder_attention_context() -> EncoderAttentionContext | None: 

70 """Return the complete attention state for the current ViT pass.""" 

71 return _ACTIVE_CONTEXT.get() 

72 

73 

74@contextmanager 

75def activate_encoder_attention(context: EncoderAttentionContext) -> Iterator[EncoderAttentionContext]: 

76 """Scope attention work and its exact-call audit to one packed ViT pass. 

77 

78 Args: 

79 context: Segments, executor, and expected module paths for this pass. 

80 """ 

81 token = _ACTIVE_CONTEXT.set(context) 

82 try: 

83 yield context 

84 finally: 

85 _ACTIVE_CONTEXT.reset(token)