Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / lightning_indexer.py: 72%

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"""LightningIndexer V2 adapter for document-local causal TopK selection.""" 

16 

17from __future__ import annotations 

18 

19from functools import lru_cache 

20from typing import Any, Callable 

21 

22import torch # pylint: disable=forbidden-backend-import 

23 

24from hyper_parallel.components.functional.sparse_flash_mla import csa_tnd_plan 

25 

26 

27# Float32 represents every integer up to this inclusive bound exactly. 

28_FP32_EXACT_INT_LIMIT = 1 << 24 

29 

30 

31def validate_indexer_tensors(query: torch.Tensor, key: torch.Tensor, weights: torch.Tensor) -> None: 

32 """Check logical tensor relationships before packing Indexer or KL inputs. 

33 

34 Native operators validate their hardware-specific dtype and shape limits. 

35 

36 Args: 

37 query: Indexer query [batch, queries, heads, dimension]. 

38 key: Indexer key bank [batch, keys, dimension]. 

39 weights: Per-query head weights [batch, queries, heads]. 

40 

41 Raises: 

42 ValueError: Tensor layouts, devices or query/key dtypes disagree. 

43 """ 

44 if query.ndim != 4: 

45 raise ValueError(f"Indexer query must have shape [B, Q, H, D], got {tuple(query.shape)}") 

46 if key.ndim != 3 or key.shape[::2] != (query.shape[0], query.shape[-1]): 

47 raise ValueError(f"Indexer key must have shape [B, K, D] matching Q, got Q={tuple(query.shape)}, " 

48 f"K={tuple(key.shape)}") 

49 if weights.shape != query.shape[:3]: 

50 raise ValueError(f"Indexer weights must have shape {tuple(query.shape[:3])}, got {tuple(weights.shape)}") 

51 if key.device != query.device or weights.device != query.device: 

52 raise ValueError(f"Indexer Q/K/weights must share a device, got {query.device}, {key.device}, {weights.device}") 

53 if key.dtype != query.dtype: 

54 raise ValueError(f"Indexer Q/K must share a dtype, got {query.dtype} and {key.dtype}") 

55 

56 

57def sort_key_indices(indices: torch.Tensor, maximum: int) -> torch.Tensor: 

58 """Sort key IDs without losing integer precision or changing their dtype. 

59 

60 Args: 

61 indices: Integer key IDs in [-1, maximum], including any sentinel. 

62 maximum: A shape-derived upper bound; does not synchronize the device. 

63 """ 

64 if maximum <= _FP32_EXACT_INT_LIMIT: 

65 return indices.float().sort(dim=-1).values.to(indices.dtype) 

66 return indices.sort(dim=-1).values 

67 

68 

69@lru_cache(maxsize=1) 

70def _load_indexer_op() -> Callable[..., Any]: 

71 """Import the LI V2 entry only when fused selection is requested.""" 

72 # Keep the optional NPU extension out of CPU and eager-only imports. 

73 try: 

74 from cann_ops_transformer import ( # pylint: disable=import-outside-toplevel 

75 lightning_indexer as lightning_indexer_op, 

76 ) 

77 except ImportError as err: 

78 raise RuntimeError( 

79 "use_fused_ops=True requires cann_ops_transformer.lightning_indexer (LI V2). " 

80 "See hyper_parallel/components/functional/README.md for ops-transformer installation, " 

81 "environment setup and tested versions." 

82 ) from err 

83 return lightning_indexer_op 

84 

85 

86@lru_cache(maxsize=1) 

87def _load_indexer_metadata_op() -> Callable[..., Any]: 

88 """Import the optional LI scheduling entry only when it will be used.""" 

89 # The verified 910B kernel computes its own split and does not read metadata. 

90 try: 

91 from cann_ops_transformer import lightning_indexer_metadata # pylint: disable=import-outside-toplevel 

92 except ImportError as err: 

93 raise RuntimeError( 

94 "use_fused_ops=True requires cann_ops_transformer.lightning_indexer_metadata on this device. " 

95 "See hyper_parallel/components/functional/README.md for ops-transformer installation, " 

96 "environment setup and tested versions." 

97 ) from err 

98 return lightning_indexer_metadata 

99 

100 

101@lru_cache(maxsize=16) 

102def _uses_indexer_metadata(device: torch.device) -> bool: 

103 """Omit unused LI scheduling metadata on the verified 910B path. 

104 

105 LI V2's arch22 kernel computes its split internally; arch35 requires 

106 metadata. Keep the standard metadata call on other devices rather than 

107 assuming they share the verified 910B behavior. TND still needs cu_seqlens. 

108 """ 

109 return "910B" not in torch.npu.get_device_name(device) 

110 

111 

112@lru_cache(maxsize=128) 

113@torch.inference_mode(False) 

114def _bsnd_indexer_metadata(device, batch, length, heads, dim, key_length, topk, ratio, residual): 

115 """Reuse shape-only BSND metadata without caching inference-mode tensors.""" 

116 residue = torch.full((batch,), residual, device=device, dtype=torch.int32) 

117 metadata = None 

118 if _uses_indexer_metadata(device): 

119 metadata = _load_indexer_metadata_op()( 

120 heads, 1, dim, topk, cmp_residual_k=residue, 

121 batch_size=batch, max_seqlen_q=length, max_seqlen_k=key_length, 

122 layout_q="BSND", layout_k="BSND", mask_mode=3, cmp_ratio=ratio, 

123 ) 

124 return residue, metadata 

125 

126 

127@torch.no_grad() 

128def _bsnd_lightning_indexer( 

129 query: torch.Tensor, key: torch.Tensor, weights: torch.Tensor, ratio: int, topk: int, 

130 segment: tuple[int, int, int, int, int], 

131) -> torch.Tensor: 

132 """Select one document segment per batch item in one BSND call. 

133 

134 Query tiling is owned by the operator, independently of the external 

135 teacher chunk size. Multiple packed documents use the TND entry. 

136 

137 Args: 

138 query: Local Indexer queries [batch, queries, heads, 128]. 

139 key: Global Indexer key bank [batch, compressed keys, 128]. 

140 weights: Local per-head merge weights, including model scaling. 

141 ratio: Model compression ratio. 

142 topk: Output selection width, including invalid tail slots. 

143 segment: Local Q start/end, global compressed K start/end, and raw 

144 prefix length modulo ratio, shared by all physical batch items. 

145 """ 

146 batch, length, heads, dim = query.shape 

147 result = torch.full((batch, length, topk), -1, device=query.device, dtype=torch.int32) 

148 start, end, key_start, key_end, residual = segment 

149 if key_end <= key_start or start == end: 

150 return result 

151 indexer = _load_indexer_op() 

152 residue, metadata = _bsnd_indexer_metadata( 

153 query.device, batch, end-start, heads, dim, key_end-key_start, topk, ratio, residual, 

154 ) 

155 with torch.autocast(device_type=query.device.type, enabled=False): 

156 indices, _ = indexer( 

157 query[:, start:end].contiguous(), key[:, key_start:key_end].unsqueeze(2).contiguous(), 

158 weights[:, start:end].float().contiguous(), topk, 

159 cmp_residual_k=residue, metadata=metadata, 

160 layout_q="BSND", layout_k="BSND", mask_mode=3, cmp_ratio=ratio, 

161 ) 

162 indices = indices[:, :, 0] 

163 indices = torch.where(indices >= 0, indices + key_start, key.shape[1]) 

164 indices = sort_key_indices(indices, key.shape[1]) 

165 result[:, start:end] = indices.masked_fill(indices == key.shape[1], -1) 

166 return result 

167 

168 

169@lru_cache(maxsize=128) 

170@torch.inference_mode(False) 

171def _tnd_indexer_metadata(plan, device, heads, dim, topk, ratio): 

172 """Cache native selection metadata for the complete packed document partition.""" 

173 options = plan.indexer_options(device, ratio) 

174 metadata = None 

175 if _uses_indexer_metadata(device): 

176 metadata = _load_indexer_metadata_op()( 

177 heads, 1, dim, topk, batch_size=len(plan.segments), 

178 max_seqlen_q=max(plan.q_lengths), max_seqlen_k=max(plan.k_lengths), **options, 

179 ) 

180 return options, metadata 

181 

182 

183@torch.no_grad() 

184def lightning_indexer( 

185 query: torch.Tensor, key: torch.Tensor, weights: torch.Tensor, ratio: int, topk: int, segments: tuple, 

186) -> torch.Tensor: 

187 """Select packed document prefixes in one TND call; single documents use BSND. 

188 

189 Args: 

190 query: Local Indexer queries [batch, queries, heads, 128]. 

191 key: Global Indexer key bank [batch, compressed keys, 128]. 

192 weights: Per-query head weights including model scaling. 

193 ratio: Model compression ratio. 

194 topk: Selection width, including invalid tail slots. 

195 segments: Local Q start/end, global document K start/prefix end and 

196 raw-prefix residual. Packed documents require batch size one. 

197 

198 Returns: 

199 Sorted global key IDs [batch, queries, topk], with invalid slots -1. 

200 """ 

201 if len(segments) == 1: 

202 return _bsnd_lightning_indexer(query, key, weights, ratio, topk, segments[0]) 

203 if query.shape[0] != 1: 

204 raise ValueError("Packed Lightning Indexer requires batch size one") 

205 result = torch.full((*query.shape[:2], topk), -1, device=query.device, dtype=torch.int32) 

206 plan = csa_tnd_plan(segments, ratio) 

207 if plan is None: 

208 return result 

209 options, metadata = _tnd_indexer_metadata(plan, query.device, query.shape[2], query.shape[3], topk, ratio) 

210 with torch.autocast(device_type=query.device.type, enabled=False): 

211 indices, _ = _load_indexer_op()( 

212 plan.query(query), plan.key(key).unsqueeze(1), plan.query(weights).float(), topk, 

213 metadata=metadata, **options, 

214 ) 

215 offsets = plan.lengths(query.device)[3] 

216 indices = torch.where(indices[:, 0] >= 0, indices[:, 0] + offsets, key.shape[1]) 

217 indices = sort_key_indices(indices, key.shape[1]) 

218 indices = indices.masked_fill(indices == key.shape[1], -1) 

219 plan.restore(indices, result) 

220 return result