Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / expert_parallel / ep_context.py: 89%

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"""Routed-slot transport, independent of routing and expert computation.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass 

20from typing import Any 

21 

22import torch 

23from torch import Tensor 

24 

25from hyper_parallel.distributed.expert_parallel.collectives import ( 

26 ep_all_to_all, ep_exchange_counts, ep_split_sizes, 

27) 

28 

29 

30@dataclass 

31class DispatchReceipt: 

32 """Per-call inverse permutations and split sizes; never cached on a module.""" 

33 

34 owner: object 

35 send_counts: list[int] 

36 recv_counts: list[int] 

37 send_order: Tensor 

38 expert_order: Tensor 

39 tokens: int 

40 top_k: int 

41 consumed: bool = False 

42 

43 

44@dataclass(frozen=True) 

45class DispatchedExperts: 

46 """Expert-major slots; probability application belongs to the expert algorithm.""" 

47 

48 hidden_states: Tensor 

49 routing_weights: Tensor 

50 expert_indices: Tensor 

51 tokens_per_expert: Tensor 

52 receipt: DispatchReceipt 

53 

54 

55def _validate_slots(hidden: Tensor, indices: Tensor, weights: Tensor, experts: int) -> None: 

56 """Validate routed token slots, expert indices and optional routing probabilities.""" 

57 if hidden.ndim != 2 or indices.ndim != 2 or weights.shape != indices.shape: 

58 raise ValueError("expected hidden[T,H], expert_indices[T,K], routing_weights[T,K]") 

59 if hidden.shape[0] != indices.shape[0] or indices.shape[1] < 1: 

60 raise ValueError("token dimensions must agree and top-k must be positive") 

61 if indices.dtype not in (torch.int32, torch.int64) or not weights.is_floating_point(): 

62 raise ValueError("expert IDs must be integers and routing weights floating point") 

63 if hidden.device != indices.device or hidden.device != weights.device: 

64 raise ValueError("all routing inputs must be on the same device") 

65 if bool(((indices < 0) | (indices >= experts)).any()): 

66 raise ValueError("every supplied slot must have an expert ID in [0, num_experts); drop is unsupported") 

67 

68 

69class EPContext: 

70 """Balanced contiguous expert ownership from HP mesh and Shard(0) layout.""" 

71 

72 def __init__(self, mesh: Any, *, num_experts: int) -> None: 

73 """Borrow the EP mesh; never create a group or reduce parameter gradients.""" 

74 if mesh is None or mesh.ndim != 1: 

75 raise ValueError("EPContext requires an existing one-dimensional HP mesh") 

76 if num_experts < 1 or num_experts % mesh.size(): 

77 raise ValueError("num_experts must be positive and divisible by EP size") 

78 self.mesh = mesh 

79 self.num_experts = num_experts 

80 self.local_experts = num_experts // mesh.size() 

81 self.rank = mesh.get_local_rank() 

82 self._identity = object() 

83 

84 def _exchange(self, tensor, send, receive): 

85 if self.mesh.size() == 1: 

86 return tensor 

87 return ep_all_to_all(tensor, send, receive, self.mesh.get_group()) 

88 

89 def dispatch(self, hidden_states: Tensor, expert_indices: Tensor, routing_weights: Tensor) -> DispatchedExperts: 

90 """Transport every already-selected slot, including zero-probability slots. 

91 

92 Args: 

93 hidden_states: Source tokens [T,H]. 

94 expert_indices: Global selected expert IDs [T,K]; no capacity/drop mode. 

95 routing_weights: Already normalized/scaled probabilities [T,K], with gradients. 

96 """ 

97 _validate_slots(hidden_states, expert_indices, routing_weights, self.num_experts) 

98 tokens, top_k = expert_indices.shape 

99 indices = expert_indices.reshape(-1).to(torch.int64) 

100 destinations = indices // self.local_experts 

101 send_order = torch.argsort(destinations, stable=True) 

102 counts = torch.bincount(destinations, minlength=self.mesh.size()) 

103 received_counts = (ep_exchange_counts(counts, self.mesh.get_group()) 

104 if self.mesh.size() > 1 else counts) 

105 send, receive = ep_split_sizes(counts), ep_split_sizes(received_counts) 

106 expanded = hidden_states[:, None, :].expand(-1, top_k, -1).reshape(-1, hidden_states.shape[-1]) 

107 values = self._exchange(expanded[send_order], send, receive) 

108 weights = self._exchange(routing_weights.reshape(-1, 1)[send_order], send, receive).reshape(-1) 

109 ids = self._exchange(indices[send_order, None], send, receive).reshape(-1) - self.rank * self.local_experts 

110 expert_order = torch.argsort(ids, stable=True) 

111 receipt = DispatchReceipt(self._identity, send, receive, send_order, expert_order, tokens, top_k) 

112 return DispatchedExperts(values[expert_order], weights[expert_order], ids[expert_order], 

113 torch.bincount(ids, minlength=self.local_experts), receipt) 

114 

115 def combine(self, values: Tensor, receipt: DispatchReceipt) -> Tensor: 

116 """Return computed slots to their source tokens and sum top-k contributions. 

117 

118 Args: 

119 values: Expert-major weighted slot outputs [received_slots,H]. 

120 receipt: Receipt from exactly one dispatch on this context. 

121 """ 

122 if receipt.owner is not self._identity or receipt.consumed: 

123 raise ValueError("dispatch receipt belongs to another context or was already consumed") 

124 if values.ndim != 2 or values.shape[0] != sum(receipt.recv_counts): 

125 raise ValueError("expert output must have one row per received slot") 

126 peer_major = values[torch.argsort(receipt.expert_order)] 

127 source_slots = self._exchange(peer_major, receipt.recv_counts, receipt.send_counts) 

128 source_slots = source_slots[torch.argsort(receipt.send_order)] 

129 output = source_slots.reshape(receipt.tokens, receipt.top_k, values.shape[-1]).sum(dim=1) 

130 receipt.consumed = True 

131 return output