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

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 # Rebuild singleton strides; contiguous() alone may retain the full-bank batch stride. 

156 native_key = key[:, key_start:key_end].contiguous().view(batch, key_end-key_start, 1, dim) 

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

158 # Singleton batch strides may survive contiguous(); view canonicalizes the native BSND layout. 

159 native_key = key[:, key_start:key_end].contiguous().view(batch, key_end-key_start, 1, dim) 

160 indices, _ = indexer( 

161 query[:, start:end].contiguous(), native_key, 

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

163 cmp_residual_k=residue, metadata=metadata, 

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

165 ) 

166 indices = indices[:, :, 0] 

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

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

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

170 return result 

171 

172 

173@lru_cache(maxsize=128) 

174@torch.inference_mode(False) 

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

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

177 options = plan.indexer_options(device, ratio) 

178 metadata = None 

179 if _uses_indexer_metadata(device): 

180 metadata = _load_indexer_metadata_op()( 

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

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

183 ) 

184 return options, metadata 

185 

186 

187@torch.no_grad() 

188def lightning_indexer( 

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

190) -> torch.Tensor: 

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

192 

193 Args: 

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

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

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

197 ratio: Model compression ratio. 

198 topk: Selection width, including invalid tail slots. 

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

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

201 

202 Returns: 

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

204 """ 

205 if len(segments) == 1: 

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

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

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

209 plan = csa_tnd_plan(segments, ratio) 

210 if plan is None: 

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

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

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

214 indices, _ = _load_indexer_op()( 

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

216 metadata=metadata, **options, 

217 ) 

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

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

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

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

222 if plan.covers_queries(query.shape[1]): 

223 return indices.unsqueeze(0) 

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

225 plan.restore(indices, result) 

226 return result