Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / engram.py: 72%
186 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"""Reusable DeepSeek Engram hash, sparse lookup, and fusion modules."""
17from __future__ import annotations
19from collections.abc import Mapping
20from typing import Any
22import torch # pylint: disable=forbidden-backend-import
23import torch.distributed as dist # pylint: disable=forbidden-backend-import
24from torch import nn # pylint: disable=forbidden-backend-import
26from hyper_parallel.distributed.expert_parallel.collectives import ep_all_to_all
27from hyper_parallel.models.materialization import register_rebuildable_buffer
28from hyper_parallel.models.replacement import module_replacement
31def local_token_mask(token_mask: torch.Tensor | None, hidden_states: torch.Tensor,
32 *, tp_rank: int, tp_size: int) -> torch.Tensor | None:
33 """Match a CP-local token mask to the current TP/SP hidden interval.
35 Args:
36 token_mask: CP-local mask replicated before the embedding boundary.
37 hidden_states: Current TP-local hidden tensor.
38 tp_rank: Local TP rank.
39 tp_size: Number of TP ranks.
41 Returns:
42 Mask aligned to the hidden sequence, or ``None``.
43 """
44 if token_mask is None or token_mask.shape[1] == hidden_states.shape[1]:
45 return token_mask
46 local_length = hidden_states.shape[1]
47 if token_mask.shape[1] != local_length * tp_size:
48 raise ValueError("Engram token mask is incompatible with TP/SP hidden")
49 return token_mask[:, tp_rank * local_length:(tp_rank + 1) * local_length]
52class NgramHashMapping(nn.Module):
53 """Map token IDs to one Engram layer's disjoint n-gram hash rows."""
55 DEAD_TOKEN = -1
57 def __init__(self, assets: Mapping[str, Any], layer_id: int) -> None:
58 """Create deterministic, non-persistent hash metadata.
60 Args:
61 assets: Prepared tokenizer map and per-layer hash layout.
62 layer_id: Decoder layer owning this Engram table.
63 """
64 super().__init__()
65 layer_ids = tuple(int(value) for value in assets["layer_ids"])
66 if layer_id not in layer_ids:
67 raise ValueError(f"Engram layer {layer_id} is absent from assets layer_ids={layer_ids}")
68 layer_index = layer_ids.index(layer_id)
69 self.layer_id = layer_id
70 self.max_ngram_size = int(assets["max_ngram_size"])
71 self.num_heads = int(assets["num_heads"])
72 token_map = torch.tensor(assets["token_map"], dtype=torch.long)
73 self.pad_id = int(token_map[int(assets["pad_token_id"])])
74 primes = torch.tensor(assets["primes"][layer_index], dtype=torch.long)
75 if tuple(primes.shape) != (self.max_ngram_size - 1, self.num_heads):
76 raise ValueError(
77 "Engram prime layout must be [max_ngram_size - 1, num_heads], "
78 f"got {tuple(primes.shape)}"
79 )
80 flattened_primes = primes.flatten()
81 offsets = torch.cat(
82 [torch.zeros(1, dtype=torch.long), flattened_primes.cumsum(0)[:-1]]
83 ).view_as(primes)
84 multipliers = torch.tensor(assets["multipliers"][layer_index], dtype=torch.long)
85 if multipliers.numel() != self.max_ngram_size:
86 raise ValueError("Engram multiplier count must equal max_ngram_size")
87 self.logical_num_embeddings = int(flattened_primes.sum().item())
88 initial_buffers = {
89 "token_map": token_map,
90 "primes": primes,
91 "offsets": offsets,
92 "multipliers": multipliers,
93 }
94 for name, value in initial_buffers.items():
95 register_rebuildable_buffer(self, name, value=value, persistent=False)
97 def forward(
98 self,
99 input_ids: torch.Tensor,
100 segment_starts: torch.Tensor | None = None,
101 token_mask: torch.Tensor | None = None,
102 ) -> torch.Tensor:
103 """Return global flat-table rows shaped ``[batch, seq, hash_columns]``.
105 Args:
106 input_ids: Token IDs for hash lookup.
107 segment_starts: Optional packed sequence boundaries.
108 token_mask: Optional mask excluding image tokens.
109 """
110 if input_ids.ndim != 2:
111 raise ValueError(f"Engram expects input_ids [B,S], got {tuple(input_ids.shape)}")
112 compressed = self.token_map[input_ids]
113 if token_mask is not None:
114 if token_mask.shape != input_ids.shape:
115 raise ValueError("Engram token_mask must have the same shape as input_ids")
116 compressed = torch.where(
117 token_mask.to(torch.bool),
118 compressed,
119 compressed.new_full((), self.DEAD_TOKEN),
120 )
121 batch_size, sequence_length = compressed.shape
122 positions = torch.arange(sequence_length, device=input_ids.device).expand(batch_size, -1)
123 segment_ids = None
124 if segment_starts is not None:
125 if segment_starts.shape != input_ids.shape:
126 raise ValueError("Engram segment_starts must have the same shape as input_ids")
127 starts = segment_starts.to(torch.bool).clone()
128 starts[:, 0] = True
129 segment_ids = starts.to(torch.long).cumsum(dim=1)
131 tokens = []
132 blocked = torch.zeros_like(positions, dtype=torch.bool)
133 for shift in range(self.max_ngram_size):
134 source_positions = (positions - shift).clamp_min(0)
135 source = compressed.gather(1, source_positions)
136 blocked = blocked | (positions < shift)
137 if segment_ids is not None:
138 source_segments = segment_ids.gather(1, source_positions)
139 blocked = blocked | (source_segments != segment_ids)
140 blocked = blocked | (source == self.DEAD_TOKEN)
141 tokens.append(torch.where(blocked, self.pad_id, source))
143 products = torch.stack(tokens, dim=-1) * self.multipliers
144 rolling = products[..., 0]
145 hashes = []
146 for ngram_index in range(self.max_ngram_size - 1):
147 rolling = torch.bitwise_xor(rolling, products[..., ngram_index + 1])
148 row = rolling.unsqueeze(-1) % self.primes[ngram_index]
149 hashes.append(row + self.offsets[ngram_index])
150 return torch.cat(hashes, dim=-1)
153def _all_gather_sequence(value: torch.Tensor, group: Any, group_size: int) -> torch.Tensor:
154 """Gather equal-size sequence chunks and concatenate along dimension one."""
155 gathered = [torch.empty_like(value) for _ in range(group_size)]
156 dist.all_gather(gathered, value.contiguous(), group=group)
157 return torch.cat(gathered, dim=1)
160@module_replacement
161class EngramModule(nn.Module):
162 """Structure-preserving Engram replacement with sparse EP row lookup."""
164 def __init__(
165 self,
166 *,
167 module: nn.Module,
168 module_fqn: str = "",
169 context: Mapping[str, Any] | None = None,
170 ) -> None:
171 """Reuse a source Engram module's parameters and hash mapping."""
172 super().__init__()
173 del module_fqn, context
174 required = ("hash_mapping", "embed", "wkv", "q_weight", "k_weight")
175 missing = [name for name in required if not hasattr(module, name)]
176 if missing:
177 raise TypeError(f"Engram source is missing required attributes: {missing}")
178 self.hash_mapping = module.hash_mapping
179 self.embed = module.embed
180 self.wkv = module.wkv
181 self.q_weight = module.q_weight
182 self.k_weight = module.k_weight
183 self.layer_id = int(module.layer_id)
184 self.hidden_size = int(module.hidden_size)
185 self.hc_mult = int(module.hc_mult)
186 self.eps = float(module.eps)
187 self.clamp_value = float(module.clamp_value)
188 self.logical_num_embeddings = int(module.logical_num_embeddings)
189 self.padded_num_embeddings = int(module.padded_num_embeddings)
190 self.train(module.training)
192 def _fuse(self, hidden_states: torch.Tensor, embeddings: torch.Tensor) -> torch.Tensor:
193 """Apply the exact V4.1 WKV projection and signed-root gate."""
194 key_value = self.wkv(embeddings.flatten(start_dim=-2))
195 key, value = key_value.split(
196 [self.hc_mult * self.hidden_size, self.hidden_size], dim=-1
197 )
198 key = key.float().unflatten(-1, (self.hc_mult, self.hidden_size))
199 hidden_fp32 = hidden_states.float()
200 weight = self.q_weight.float() * self.k_weight.float()
201 reciprocal_std = torch.rsqrt(hidden_fp32.square().mean(-1) + self.eps)
202 reciprocal_std = reciprocal_std * torch.rsqrt(key.square().mean(-1) + self.eps)
203 dot = (hidden_fp32 * weight * key).sum(-1)
204 dot = dot * reciprocal_std * self.hidden_size**-0.5
205 root = dot.abs().clamp_min(self.clamp_value).sqrt()
206 signed_root = torch.where(dot < 0, -root, root)
207 gate = torch.sigmoid(signed_root)
208 return (
209 hidden_fp32 + gate.unsqueeze(-1) * value.float().unsqueeze(-2)
210 ).to(hidden_states.dtype)
212 @staticmethod
213 def _sparse_lookup(
214 hash_ids: torch.Tensor,
215 embedding: nn.Embedding,
216 *,
217 ep_group: Any,
218 ep_rank: int,
219 ep_size: int,
220 padded_num_embeddings: int,
221 ) -> torch.Tensor:
222 """Route requested global rows to their EP owners and restore order."""
223 if padded_num_embeddings % ep_size:
224 raise ValueError(
225 f"Engram padded rows {padded_num_embeddings} are not divisible by ep_size={ep_size}"
226 )
227 num_local_rows = padded_num_embeddings // ep_size
228 row_start = ep_rank * num_local_rows
229 flat_ids = hash_ids.reshape(-1)
230 owner = torch.div(flat_ids, num_local_rows, rounding_mode="floor")
231 if owner.numel() and int(owner.max()) >= ep_size:
232 raise IndexError("Engram hash requested a row outside the padded table")
233 local_counts = torch.bincount(owner, minlength=ep_size).to(torch.long)
234 gathered_counts = [torch.empty_like(local_counts) for _ in range(ep_size)]
235 dist.all_gather(gathered_counts, local_counts, group=ep_group)
236 all_counts = torch.stack(gathered_counts)
237 send_counts = local_counts.cpu().tolist()
238 recv_counts = all_counts[:, ep_rank].cpu().tolist()
240 # Ascend ArgSort executes integer inputs on AiCPU. The owner range is
241 # bounded by ep_size, so float32 preserves it exactly and stays on AiCore.
242 sort_indices = torch.argsort(owner.float(), stable=True)
243 sorted_ids = flat_ids.index_select(0, sort_indices)
244 received_ids = ep_all_to_all(sorted_ids, send_counts, recv_counts, ep_group)
245 local_ids = received_ids - row_start
246 if local_ids.numel() and (int(local_ids.min()) < 0 or int(local_ids.max()) >= num_local_rows):
247 raise IndexError("Engram EP routing delivered a row to the wrong owner")
248 # Call the child module instead of reading ``embedding.weight``
249 # directly. This preserves the FSDP pre/post-forward lifecycle when
250 # the EP table is wrapped independently from the dense Engram gates.
251 values = embedding(local_ids)
252 returned = ep_all_to_all(values, recv_counts, send_counts, ep_group)
253 restored = torch.zeros_like(returned)
254 restored.scatter_add_(
255 0,
256 sort_indices.unsqueeze(1).expand(-1, returned.shape[1]),
257 returned,
258 )
259 return restored.view(*hash_ids.shape, embedding.embedding_dim)
261 def _aligned_hash_ids(
262 self,
263 hidden_states: torch.Tensor,
264 input_ids: torch.Tensor,
265 segment_starts: torch.Tensor | None,
266 token_mask: torch.Tensor | None = None,
267 *,
268 cp_group: Any = None,
269 cp_rank: int = 0,
270 cp_size: int = 1,
271 tp_rank: int = 0,
272 tp_size: int = 1,
273 ) -> torch.Tensor:
274 """Hash the local sequence plus the left halo required by n-grams."""
275 cp_local_ids = input_ids
276 if cp_size > 1:
277 input_ids = _all_gather_sequence(input_ids, cp_group, cp_size)
278 if segment_starts is not None:
279 segment_starts = _all_gather_sequence(segment_starts, cp_group, cp_size)
280 if token_mask is not None:
281 token_mask = _all_gather_sequence(token_mask, cp_group, cp_size)
282 local_sequence = hidden_states.shape[1]
283 cp_sequence = cp_local_ids.shape[1]
284 if cp_sequence == local_sequence:
285 sequence_parallel_parts = 1
286 local_tp_rank = 0
287 elif cp_sequence % local_sequence == 0 and cp_sequence // local_sequence == tp_size:
288 sequence_parallel_parts = tp_size
289 local_tp_rank = tp_rank
290 else:
291 raise ValueError(
292 "Engram input_ids length is incompatible with the local hidden sequence: "
293 f"ids={cp_sequence}, hidden={local_sequence}, tp_size={tp_size}"
294 )
295 del sequence_parallel_parts
296 start = cp_rank * cp_sequence + local_tp_rank * local_sequence
297 end = start + local_sequence
298 context_start = max(0, start - self.hash_mapping.max_ngram_size + 1)
299 ids_window = input_ids[:, context_start:end]
300 starts_window = None if segment_starts is None else segment_starts[:, context_start:end]
301 mask_window = None if token_mask is None else token_mask[:, context_start:end]
302 hashes = self.hash_mapping(ids_window, starts_window, mask_window)
303 target_offset = start - context_start
304 return hashes[:, target_offset:target_offset + local_sequence]
306 def forward(
307 self,
308 hidden_states: torch.Tensor,
309 input_ids: torch.Tensor,
310 segment_starts: torch.Tensor | None = None,
311 token_mask: torch.Tensor | None = None,
312 ) -> torch.Tensor:
313 """Run replicated single-card/EP1 Engram.
315 Args:
316 hidden_states: Current decoder hidden rows.
317 input_ids: CP-local token IDs for hash lookup.
318 segment_starts: Optional packed sequence boundaries.
319 token_mask: Optional mask selecting Engram updates.
320 """
321 hash_ids = self._aligned_hash_ids(hidden_states, input_ids, segment_starts, token_mask)
322 fused = self._fuse(hidden_states, self.embed(hash_ids))
323 if token_mask is None:
324 return fused
325 return torch.where(token_mask.to(torch.bool).unsqueeze(-1).unsqueeze(-1), fused, hidden_states)
327 def parallel_forward(
328 self,
329 hidden_states: torch.Tensor,
330 input_ids: torch.Tensor,
331 segment_starts: torch.Tensor | None = None,
332 token_mask: torch.Tensor | None = None,
333 *,
334 ep_group: Any,
335 ep_rank: int,
336 ep_size: int,
337 cp_group: Any = None,
338 cp_rank: int = 0,
339 cp_size: int = 1,
340 tp_rank: int = 0,
341 tp_size: int = 1,
342 ) -> torch.Tensor:
343 """Run CP/TP sequence alignment and sparse EP lookup.
345 Args:
346 hidden_states: Current decoder hidden rows.
347 input_ids: CP-local token IDs for hash construction.
348 segment_starts: Optional packed sequence boundaries.
349 token_mask: CP-local mask selecting Engram updates.
350 ep_group: Expert-parallel process group.
351 ep_rank: Local expert-parallel rank.
352 ep_size: Expert-parallel degree.
353 cp_group: Context-parallel process group.
354 cp_rank: Local context-parallel rank.
355 cp_size: Context-parallel degree.
356 tp_rank: Local tensor-parallel rank.
357 tp_size: Tensor-parallel degree.
358 """
359 hash_ids = self._aligned_hash_ids(
360 hidden_states,
361 input_ids,
362 segment_starts,
363 token_mask,
364 cp_group=cp_group,
365 cp_rank=cp_rank,
366 cp_size=cp_size,
367 tp_rank=tp_rank,
368 tp_size=tp_size,
369 )
370 embeddings = self._sparse_lookup(
371 hash_ids,
372 self.embed,
373 ep_group=ep_group,
374 ep_rank=ep_rank,
375 ep_size=ep_size,
376 padded_num_embeddings=self.padded_num_embeddings,
377 )
378 fused = self._fuse(hidden_states, embeddings)
379 token_mask = local_token_mask(token_mask, hidden_states, tp_rank=tp_rank, tp_size=tp_size)
380 if token_mask is None:
381 return fused
382 return torch.where(token_mask.to(torch.bool).unsqueeze(-1).unsqueeze(-1), fused, hidden_states)
385__all__ = ["EngramModule", "NgramHashMapping", "local_token_mask"]