Feat: HyperParallel支持LLM推理Generate流程 #2101-hyper-parallel-AtomGit

Aexamples/generate/README.md+192-0

@@ -0,0 +1,192 @@
	1	+# Generate example
	2	+
	3	+This directory contains a minimal validation path for
	4	+`hyper_parallel.infer`. The example covers greedy/top-k/top-p generation,
	5	+KV-cache decode, no-cache fallback, HuggingFace alignment, and a small
	6	+performance baseline script.
	7	+
	8	+## Functional tests
	9	+
	10	+```bash
	11	+python -m pytest tests/torch/generate -q
	12	+```
	13	+
	14	+Expected result:
	15	+
	16	+```text
	17	+44 passed
	18	+```
	19	+
	20	+## Python API
	21	+
	22	+```python
	23	+import torch
	24	+
	25	+from hyper_parallel.infer import GenerationConfig, generate
	26	+
	27	+model.eval()
	28	+input_ids = torch.tensor([[1, 2, 3]], dtype=torch.long)
	29	+output_ids = generate(
	30	+ model,
	31	+ input_ids=input_ids,
	32	+ generation_config=GenerationConfig(
	33	+ max_new_tokens=16,
	34	+ do_sample=False,
	35	+ eos_token_id=None,
	36	+ ),
	37	+)
	38	+```
	39	+
	40	+If the model returns `past_key_values`, decode uses KV cache. If the model
	41	+does not return `past_key_values`, generation falls back to full-sequence
	42	+recompute so existing project models can still run.
	43	+
	44	+Sampling supports greedy, top-k, and top-p modes through `GenerationConfig`.
	45	+Custom `logits_processor` and `stopping_criteria` callables can be supplied as
	46	+lists. Each logits processor receives `(input_ids, scores)` and returns updated
	47	+scores; each stopping criterion receives `(input_ids, scores)` and returns a
	48	+scalar boolean or scalar tensor.
	49	+
	50	+## Benchmark
	51	+
	52	+Run a small deterministic baseline:
	53	+
	54	+```bash
	55	+python examples/generate/benchmark_generate.py \
	56	+ --device npu \
	57	+ --batch-size 4 \
	58	+ --prompt-len 32 \
	59	+ --max-new-tokens 64 \
	60	+ --warmup 2 \
	61	+ --repeat 5 \
	62	+ --output /tmp/generate_baseline.json
	63	+```
	64	+
	65	+The JSON output records latency and token throughput for:
	66	+
	67	+- `prefill`: prompt-only forward latency.
	68	+- `with_cache`: model returns `past_key_values`; decode uses single-token steps.
	69	+- `no_cache`: model does not return `past_key_values`; decode recomputes the
	70	+ whole sequence.
	71	+
	72	+Key output fields:
	73	+
	74	+- top-level: `device`, `batch_size`, `prompt_len`, `max_new_tokens`,
	75	+ `vocab_size`, `warmup`, `repeat`.
	76	+- `prefill`: `mean_seconds`, `min_seconds`, `max_seconds`,
	77	+ `tokens_per_second`.
	78	+- `with_cache` and `no_cache`: `calls`, `first_call`, `last_call`,
	79	+ `output_shape`, `mean_seconds`, `min_seconds`, `max_seconds`,
	80	+ `tokens_per_second`.
	81	+
	82	+Use the same command and hardware when comparing future generate changes.
	83	+
	84	+## HuggingFace alignment
	85	+
	86	+Run a real causal-LM alignment check with a local HuggingFace model:
	87	+
	88	+```bash
	89	+export MODEL_PATH=/data/models/Qwen3-4B-Instruct-2507
	90	+
	91	+python examples/generate/hf_alignment_check.py \
	92	+ --model "$MODEL_PATH" \
	93	+ --prompt "用一段话简单介绍Mindspore。" \
	94	+ --max-new-tokens 128 \
	95	+ --logits-compare-steps 128 \
	96	+ --device npu \
	97	+ --trust-remote-code \
	98	+ --output /tmp/generate_hf_alignment_128.json
	99	+```
	100	+
	101	+The helper runs HuggingFace greedy `model.generate(...)`,
	102	+`hyper_parallel.infer.generate(..., use_cache=True)`, and
	103	+`hyper_parallel.infer.generate(..., use_cache=False)` on the same prompt. It
	104	+fails if:
	105	+
	106	+- HuggingFace ids and HyperParallel cache ids differ;
	107	+- HyperParallel cache and no-cache ids differ;
	108	+- cache/no-cache decode logits cosine similarity falls below the configured
	109	+ threshold.
	110	+
	111	+The output JSON should contain the following validation fields:
	112	+
	113	+```json
	114	+{
	115	+ "prompt": "用一段话简单介绍Mindspore。",
	116	+ "generated_new_tokens": 128,
	117	+ "hf_vs_hyper_cache_ids_match": true,
	118	+ "hyper_cache_vs_no_cache_ids_match": true,
	119	+ "logits_compare_steps": 128,
	120	+ "logits_cosine_min": 0.9996622800827026,
	121	+ "logits_cosine_pass": true,
	122	+ "logits_cosine_threshold": 0.999
	123	+}
	124	+```
	125	+
	126	+The same JSON also includes decoded `hf_text`, `hyper_cache_text`, and
	127	+`hyper_no_cache_text` fields, which can be inspected to confirm the generated
	128	+text for the prompt.
	129	+
	130	+## Project model smoke
	131	+
	132	+The CPU test suite covers no-cache generation against real repository model
	133	+classes with small random configurations:
	134	+
	135	+- `Qwen3_5ForCausalLM` through the functional `generate(model, ...)` API.
	136	+- `Qwen3_5ForCausalLM` with `GenerateMixin`, using the `model.generate(...)`
	137	+ method form.
	138	+- `Qwen3_5MoeForCausalLM` through the functional API and the
	139	+ `model.generate(...)` method form.
	140	+
	141	+These tests validate that at least one repository model can use the generate
	142	+flow without requiring checkpoint downloads. The current Qwen project-model
	143	+checks use no-cache fallback because those model forwards do not return
	144	+`past_key_values`; cache behavior is covered by cache-capable test models and
	145	+the KV cache unit tests.
	146	+
	147	+## Distributed boundary
	148	+
	149	+The current implementation provides the model-agnostic generate loop and KV
	150	+cache container. Tensor-parallel inference can gather vocab-sharded logits
	151	+before sampling by enabling `GenerationConfig(gather_logits=True)` and passing
	152	+the tensor-parallel process group through `logits_process_group` when needed.
	153	+
	154	+Context-parallel prefill can select final-token logits from the rank that owns
	155	+the global last prompt token by setting `context_logits_rank` and passing the
	156	+context-parallel process group through `context_process_group`. This keeps
	157	+sampling decisions identical across ranks after sequence-sharded prefill.
	158	+
	159	+Context-parallel KV cache decode is available as an opt-in path by setting
	160	+`GenerationConfig(context_parallel_cache=True)`. The model forward may either
	161	+return a full prefill cache, which `generate` shards into a
	162	+`ContextParallelKVCache`, or return a local cache together with
	163	+`sequence_shard_info`. During cached decode, CP-aware models must accept the
	164	+local `past_key_values`, the current `sequence_shard_info`, and
	165	+`global_seq_len`, then return updated local K/V tensors plus the next
	166	+`sequence_shard_info`.
	167	+
	168	+Contiguous prefix cache reuse is supported by passing
	169	+`prefix_past_key_values` and `prefix_attention_mask` in `GenerationConfig`.
	170	+The prefix cache is treated as a continuous history immediately before
	171	+`input_ids`; returned token ids contain only `input_ids` and newly generated
	172	+tokens, not the reused prefix tokens. For CP prefix cache, callers may pass a
	173	+full prefix cache or a local prefix cache with `prefix_sequence_shard_info`.
	174	+
	175	+Supported CP cache boundary:
	176	+
	177	+- contiguous sequence shards only;
	178	+- append-only decode positions;
	179	+- contiguous prefix cache reuse before `input_ids`;
	180	+- `context_logits_rank` is local to `context_process_group`;
	181	+- packed/non-contiguous cache layouts are not part of this generic generate
	182	+ path because they require explicit per-token position or segment metadata
	183	+ and model-side attention support.
	184	+
	185	+Run the end-to-end CP cache check with:
	186	+
	187	+```bash
	188	+python -m torch.distributed.run --nproc_per_node=2 examples/generate/cp_generate_check.py
	189	+```
	190	+
	191	+For models that already return full replicated logits on each rank, the
	192	+single-process API can be used directly without enabling logits gathering.

Aexamples/generate/benchmark_generate.py+247-0

@@ -0,0 +1,247 @@

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 +"""Generate micro-benchmark."""

16 +import argparse

17 +import importlib

18 +import json

19 +import sys

20 +import time

21 +import warnings

22 +from pathlib import Path

23 +from statistics import mean

24 +

25 +import torch

26 +from torch import nn

27 +

28 +_REPO_ROOT = Path(__file__).resolve().parents[2]

29 +if str(_REPO_ROOT) not in sys.path:

30 + sys.path.insert(0, str(_REPO_ROOT))

31 +

32 +_infer = importlib.import_module("hyper_parallel.infer")

33 +GenerationConfig = _infer.GenerationConfig

34 +build_causal_mask = _infer.build_causal_mask

35 +build_position_ids = _infer.build_position_ids

36 +generate = _infer.generate

37 +

38 +

39 +class CacheLengthLM(nn.Module):

40 + """Small deterministic LM for measuring generate loop overhead."""

41 +

42 + def __init__(self, vocab_size: int = 32000, use_cache: bool = True):

43 + super().__init__()

44 + self.vocab_size = vocab_size

45 + self.use_cache = use_cache

46 + self.calls = []

47 +

48 + def forward(

49 + self,

50 + input_ids,

51 + position_ids=None,

52 + attention_mask=None,

53 + past_key_values=None,

54 + use_cache=True,

55 + **kwargs,

56 + ):

57 + """Run deterministic forward with optional cache output."""

58 + del position_ids, attention_mask, kwargs

59 + batch_size, seq_len = input_ids.shape

60 + past_len = 0

61 + if past_key_values is not None:

62 + past_len = past_key_values[0][0].shape[-2]

63 + total_len = past_len + seq_len

64 + token_id = total_len % self.vocab_size

65 + logits = torch.full(

66 + (batch_size, seq_len, self.vocab_size),

67 + -1000.0,

68 + device=input_ids.device,

69 + dtype=torch.float32,

70 + )

71 + logits[:, -1, token_id] = 1000.0

72 + self.calls.append((seq_len, past_len))

73 + past = None

74 + if self.use_cache and use_cache:

75 + key = torch.zeros(batch_size, 1, total_len, 4, device=input_ids.device)

76 + value = torch.zeros(batch_size, 1, total_len, 4, device=input_ids.device)

77 + past = [(key, value)]

78 + return {"logits": logits, "past_key_values": past}

79 +

80 +

81 +def _synchronize(device: torch.device) -> None:

82 + if device.type == "cuda" and torch.cuda.is_available():

83 + torch.cuda.synchronize()

84 + elif device.type == "npu" and hasattr(torch, "npu"):

85 + torch.npu.synchronize()

86 +

87 +

88 +def _resolve_device(name: str) -> torch.device:

89 + if name == "auto":

90 + if hasattr(torch, "npu") and torch.npu.is_available():

91 + return torch.device("npu")

92 + if torch.cuda.is_available():

93 + return torch.device("cuda")

94 + return torch.device("cpu")

95 + return torch.device(name)

96 +

97 +

98 +def _run_case(args, device: torch.device, use_cache: bool) -> dict:

99 + """Measure one full generate case."""

100 + model = CacheLengthLM(vocab_size=args.vocab_size, use_cache=use_cache).to(device)

101 + input_ids = (

102 + torch.arange(args.prompt_len, device=device, dtype=torch.long)

103 + .view(1, -1)

104 + .expand(args.batch_size, -1)

105 + .contiguous()

106 + % args.vocab_size

107 + )

108 + config = GenerationConfig(

109 + max_new_tokens=args.max_new_tokens,

110 + do_sample=False,

111 + eos_token_id=None,

112 + use_cache=use_cache,

113 + )

atomgit-bot6月22日

🔵 Low Priority

在 benchmark_generate.py 第 105-110 行，_run_case 函数中 GenerationConfig 的 use_cache 被硬编码为 True：

但在 main() 中，_run_case 被调用两次：

当 use_cache=False 时，虽然模型的 self.use_cache=False 阻止了实际返回 past_key_values，但 config.use_cache=True 会让 generate() 仍然请求模型使用 cache。当前行为碰巧正确，但语义不明确：既然意图是测量 no-cache 场景，config.use_cache 应设为 False 使语义一致。如果未来 generate() 函数的行为发生改变（例如在 config.use_cache=True 时做额外优化假设），此处可能产生不一致。

修复：将 config.use_cache 改为使用传入的 use_cache 参数。

建议：将 config.use_cache 从硬编码 True 改为使用 use_cache 参数，使 benchmark 语义明确：no_cache 场景下配置与模型行为一致。

改动建议

+    config = GenerationConfig(
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+        eos_token_id=None,
+        use_cache=use_cache,
     )

应用建议

不准确？

114 +

115 + for _ in range(args.warmup):

116 + generate(model, input_ids, config)

117 + _synchronize(device)

118 +

119 + elapsed = []

120 + last_output = None

121 + for _ in range(args.repeat):

122 + model.calls.clear()

123 + start = time.perf_counter()

124 + last_output = generate(model, input_ids, config)

125 + _synchronize(device)

126 + elapsed.append(time.perf_counter() - start)

127 +

128 + total_generated = args.batch_size * args.max_new_tokens

129 + mean_seconds = mean(elapsed)

130 + return {

131 + "mean_seconds": mean_seconds,

132 + "min_seconds": min(elapsed),

133 + "max_seconds": max(elapsed),

134 + "tokens_per_second": total_generated / mean_seconds,

135 + "output_shape": list(last_output.shape),

136 + "calls": len(model.calls),

137 + "first_call": list(model.calls[0]),

138 + "last_call": list(model.calls[-1]),

139 + }

140 +

141 +

142 +def _measure_prefill(args, device: torch.device) -> dict:

143 + """Measure the standalone prefill forward path."""

144 + model = CacheLengthLM(vocab_size=args.vocab_size, use_cache=True).to(device)

145 + input_ids = (

146 + torch.arange(args.prompt_len, device=device, dtype=torch.long)

147 + .view(1, -1)

148 + .expand(args.batch_size, -1)

149 + .contiguous()

150 + % args.vocab_size

151 + )

152 + position_ids = build_position_ids(input_ids)

153 + attention_mask = build_causal_mask(input_ids)

154 +

155 + for _ in range(args.warmup):

156 + model(

157 + input_ids=input_ids,

158 + position_ids=position_ids,

159 + attention_mask=attention_mask,

160 + past_key_values=None,

161 + use_cache=True,

162 + )

163 + _synchronize(device)

164 +

165 + elapsed = []

166 + for _ in range(args.repeat):

167 + start = time.perf_counter()

168 + model(

169 + input_ids=input_ids,

170 + position_ids=position_ids,

171 + attention_mask=attention_mask,

172 + past_key_values=None,

173 + use_cache=True,

174 + )

175 + _synchronize(device)

176 + elapsed.append(time.perf_counter() - start)

177 +

178 + mean_seconds = mean(elapsed)

179 + return {

180 + "mean_seconds": mean_seconds,

181 + "min_seconds": min(elapsed),

182 + "max_seconds": max(elapsed),

183 + "tokens_per_second": (args.batch_size * args.prompt_len) / mean_seconds,

184 + }

185 +

186 +

187 +def parse_args():

188 + """Parse benchmark command-line arguments."""

189 + parser = argparse.ArgumentParser(description=__doc__)

190 + parser.add_argument("--batch-size", type=int, default=4)

191 + parser.add_argument("--prompt-len", type=int, default=32)

192 + parser.add_argument("--max-new-tokens", type=int, default=64)

193 + parser.add_argument("--vocab-size", type=int, default=32000)

atomgit-bot6月29日

🔵 Low Priority

Changed lines: parse_args() 中 line 193 声明了 --vocab-size 参数，但 _validate_args() (line 201-218) 未对其做正数校验。如果用户传入 --vocab-size 0 或负数，将导致：

ZeroDivisionError — input_ids 构造中的 % args.vocab_size (lines 106, 150) 对 vocab_size=0 取模直接崩溃。
无效模型 — CacheLengthLM(vocab_size=...) (lines 100, 144) 中 nn.Embedding(0, ...) / nn.Linear(..., 0) 会抛出 RuntimeError。
token_id 越界 — total_len % 0 同样崩溃 (line 64)。

Affected behavior: 用户通过 CLI 传入非法 vocab_size 时脚本直接崩溃，且错误信息不是友好的校验提示。

Failure mode: 用户输入 --vocab-size 0 时，脚本在模型构造或取模运算处抛出 ZeroDivisionError 或 RuntimeError，而非明确的参数校验错误。

建议：在 _validate_args() 中添加 if args.vocab_size <= 0: raise ValueError("vocab-size must be positive")

不准确？

194 + parser.add_argument("--warmup", type=int, default=2)

195 + parser.add_argument("--repeat", type=int, default=5)

196 + parser.add_argument("--device", default="auto")

197 + parser.add_argument("--output", default=None)

atomgit-bot6月22日

🔵 Low Priority

建议：对 --warmup 参数增加 warmup < 0 的校验，避免用户传入负数后静默跳过预热循环。同时可考虑增加 warmup == 0 时的警告，提醒用户没有预热可能影响测量准确性。

不准确？

198 + return parser.parse_args()

199 +

200 +

201 +def _validate_args(args) -> None:

202 + """Validate benchmark arguments."""

203 + if args.batch_size <= 0:

204 + raise ValueError("batch-size must be positive")

205 + if args.prompt_len <= 0:

206 + raise ValueError("prompt-len must be positive")

207 + if args.max_new_tokens <= 0:

208 + raise ValueError("max-new-tokens must be positive")

209 + if args.vocab_size <= 0:

210 + raise ValueError("vocab-size must be positive")

211 + if args.warmup < 0:

212 + raise ValueError("warmup must be >= 0")

213 + if args.warmup == 0:

214 + warnings.warn(

215 + "warmup is 0; benchmark results may include first-run overhead",

216 + RuntimeWarning,

217 + stacklevel=2,

218 + )

219 + if args.repeat <= 0:

220 + raise ValueError("repeat must be positive")

221 +

222 +

223 +def main():

224 + """Run the generate benchmark and optionally write JSON output."""

225 + args = parse_args()

226 + _validate_args(args)

227 + device = _resolve_device(args.device)

228 + result = {

229 + "device": str(device),

230 + "batch_size": args.batch_size,

231 + "prompt_len": args.prompt_len,

232 + "max_new_tokens": args.max_new_tokens,

233 + "vocab_size": args.vocab_size,

234 + "warmup": args.warmup,

235 + "repeat": args.repeat,

236 + "prefill": _measure_prefill(args, device),

237 + "with_cache": _run_case(args, device, use_cache=True),

238 + "no_cache": _run_case(args, device, use_cache=False),

239 + }

240 + text = json.dumps(result, indent=2, sort_keys=True)

241 + print(text)

242 + if args.output:

243 + Path(args.output).write_text(text + "\n", encoding="utf-8")

244 +

245 +

246 +if __name__ == "__main__":

247 + main()

Aexamples/generate/cp_generate_check.py+288-0

@@ -0,0 +1,288 @@
	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	+"""End-to-end context-parallel generate check."""
	16	+import argparse
	17	+import json
	18	+import math
	19	+from pathlib import Path
	20	+
	21	+import torch
	22	+import torch.distributed as dist
	23	+from torch import nn
	24	+
	25	+from hyper_parallel.infer import GenerationConfig, generate, get_sequence_shard_info
	26	+
	27	+
	28	+def _init_weights(module: nn.Module, vocab_size: int) -> None:
	29	+ with torch.no_grad():
	30	+ module.embedding.weight.zero_()
	31	+ module.lm_head.weight.zero_()
	32	+ for token_id in range(vocab_size):
	33	+ scale = 20.0 if token_id == 7 else 1.0
	34	+ module.embedding.weight[token_id, token_id] = scale
	35	+ module.lm_head.weight[token_id, token_id] = 1.0
	36	+
	37	+
	38	+class FullTinyAttentionLM(nn.Module):
	39	+ """Single-process tiny attention model used as the baseline."""
	40	+
	41	+ def __init__(self, vocab_size: int):
	42	+ super().__init__()
	43	+ self.vocab_size = vocab_size
	44	+ self.hidden_size = vocab_size
	45	+ self.embedding = nn.Embedding(vocab_size, self.hidden_size)
	46	+ self.lm_head = nn.Linear(self.hidden_size, vocab_size, bias=False)
	47	+ _init_weights(self, vocab_size)
	48	+
	49	+ def forward(
	50	+ self,
	51	+ input_ids,
	52	+ position_ids=None,
	53	+ attention_mask=None,
	54	+ past_key_values=None,
	55	+ use_cache=True,
	56	+ **kwargs,
	57	+ ):
	58	+ """Run full-cache attention forward."""
	59	+ del position_ids, kwargs
	60	+ hidden = self.embedding(input_ids)
	61	+ key = hidden.unsqueeze(1)
	62	+ value = key
	63	+ if past_key_values is not None:
	64	+ past_key, past_value = past_key_values[0]
	65	+ key = torch.cat([past_key, key], dim=-2)
	66	+ value = torch.cat([past_value, value], dim=-2)
	67	+ logits = self._attention_logits(hidden, key, value, attention_mask)
	68	+ past = [(key.detach(), value.detach())] if use_cache else None
	69	+ return {"logits": logits, "past_key_values": past}
	70	+
	71	+ def _attention_logits(self, hidden, key, value, attention_mask):
	72	+ query = torch.ones_like(hidden).unsqueeze(1)
	73	+ scores = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(self.hidden_size)
	74	+ if attention_mask is not None:
	75	+ scores = scores + attention_mask.to(dtype=scores.dtype)
	76	+ probs = torch.softmax(scores, dim=-1)
	77	+ context = torch.matmul(probs, value).squeeze(1)
	78	+ return self.lm_head(context)
	79	+
	80	+
	81	+class ContextParallelTinyAttentionLM(FullTinyAttentionLM):
	82	+ """Tiny attention model whose KV cache is sharded by sequence across ranks."""
	83	+
	84	+ def __init__(self, vocab_size: int, rank: int, world_size: int):
	85	+ super().__init__(vocab_size)
	86	+ self.rank = rank
	87	+ self.world_size = world_size
	88	+
	89	+ def forward(
	90	+ self,
	91	+ input_ids,
	92	+ position_ids=None,
	93	+ attention_mask=None,
	94	+ past_key_values=None,
	95	+ use_cache=True,
	96	+ sequence_shard_info=None,
	97	+ global_seq_len=None,
	98	+ **kwargs,
	99	+ ):
	100	+ del position_ids, kwargs
	101	+ hidden = self.embedding(input_ids)
	102	+ new_key = hidden.unsqueeze(1)
	103	+ new_value = new_key
	104	+ if past_key_values is None:
	105	+ full_key = new_key
	106	+ full_value = new_value
	107	+ else:
	108	+ full_key, full_value = self._gather_full_cache(
	109	+ past_key_values,
	110	+ sequence_shard_info,
	111	+ )
	112	+ full_key = torch.cat([full_key, new_key], dim=-2)
	113	+ full_value = torch.cat([full_value, new_value], dim=-2)
	114	+ if global_seq_len is not None and full_key.shape[-2] != global_seq_len:
	115	+ raise ValueError("global_seq_len does not match rebuilt CP cache")
	116	+ logits = self._attention_logits(hidden, full_key, full_value, attention_mask)
	117	+ shard_info = get_sequence_shard_info(full_key.shape[-2], self.rank, self.world_size)
	118	+ local_key = full_key.narrow(-2, shard_info.start, shard_info.local_seq_len).contiguous()
	119	+ local_value = full_value.narrow(-2, shard_info.start, shard_info.local_seq_len).contiguous()
	120	+ past = [(local_key.detach(), local_value.detach())] if use_cache else None
	121	+ return {
	122	+ "logits": logits,
	123	+ "past_key_values": past,
	124	+ "sequence_shard_info": shard_info,
	125	+ }
	126	+
	127	+ def _gather_full_cache(self, past_key_values, sequence_shard_info):
	128	+ """Gather local CP cache shards into full-sequence cache."""
	129	+ if sequence_shard_info is None:
	130	+ raise ValueError("sequence_shard_info is required for CP cached decode")
	131	+ local_key, local_value = past_key_values[0]
	132	+ start = torch.tensor([sequence_shard_info.start], device=local_key.device)
	133	+ local_len = torch.tensor([local_key.shape[-2]], device=local_key.device)
	134	+ starts = [torch.empty_like(start) for _ in range(self.world_size)]
	135	+ lengths = [torch.empty_like(local_len) for _ in range(self.world_size)]
	136	+ dist.all_gather(starts, start)
	137	+ dist.all_gather(lengths, local_len)
	138	+
	139	+ max_len = int(torch.stack(lengths).max().item())
	140	+ padded_key = _pad_sequence_shard(local_key, max_len)
	141	+ padded_value = _pad_sequence_shard(local_value, max_len)
	142	+ gathered_keys = [torch.empty_like(padded_key) for _ in range(self.world_size)]
	143	+ gathered_values = [torch.empty_like(padded_value) for _ in range(self.world_size)]
	144	+ dist.all_gather(gathered_keys, padded_key)
	145	+ dist.all_gather(gathered_values, padded_value)
	146	+
	147	+ gathered = sorted(
	148	+ (
	149	+ int(rank_start.item()),
	150	+ key.narrow(-2, 0, int(rank_len.item())),
	151	+ value.narrow(-2, 0, int(rank_len.item())),
	152	+ )
	153	+ for rank_start, rank_len, key, value in zip(
	154	+ starts,
	155	+ lengths,
	156	+ gathered_keys,
	157	+ gathered_values,
	158	+ )
	159	+ )
	160	+ full_key = torch.cat([item[1] for item in gathered], dim=-2)
	161	+ full_value = torch.cat([item[2] for item in gathered], dim=-2)
	162	+ return full_key, full_value
	163	+
	164	+
	165	+def _pad_sequence_shard(tensor: torch.Tensor, max_len: int) -> torch.Tensor:
	166	+ pad_len = max_len - tensor.shape[-2]
	167	+ if pad_len == 0:
	168	+ return tensor.contiguous()
	169	+ pad_shape = list(tensor.shape)
	170	+ pad_shape[-2] = pad_len
	171	+ padding = tensor.new_zeros(pad_shape)
	172	+ return torch.cat([tensor, padding], dim=-2).contiguous()
	173	+
	174	+
	175	+def _gather_token_outputs(output: torch.Tensor, world_size: int):
	176	+ gathered = [torch.empty_like(output) for _ in range(world_size)]
	177	+ dist.all_gather(gathered, output)
	178	+ return [item.cpu().tolist() for item in gathered]
	179	+
	180	+
	181	+def _parse_args():
	182	+ parser = argparse.ArgumentParser(description=__doc__)
	183	+ parser.add_argument("--backend", default="gloo")
	184	+ parser.add_argument("--vocab-size", type=int, default=64)
	185	+ parser.add_argument("--max-new-tokens", type=int, default=4)
	186	+ parser.add_argument("--output", default=None)
	187	+ return parser.parse_args()
	188	+
	189	+
	190	+def main():
	191	+ """Run the distributed CP generate check."""
	192	+ args = _parse_args()
	193	+ if not dist.is_initialized():
	194	+ dist.init_process_group(backend=args.backend)
	195	+ rank = dist.get_rank()
	196	+ world_size = dist.get_world_size()
	197	+ if world_size < 2:
	198	+ raise ValueError("CP generate check requires world_size >= 2")
	199	+
	200	+ input_ids = torch.tensor([[7, 7, 1], [0, 7, 1]], dtype=torch.long)
	201	+ attention_mask = torch.tensor([[0, 0, 1], [0, 1, 1]], dtype=torch.long)
	202	+ config_kwargs = {
	203	+ "max_new_tokens": args.max_new_tokens,
	204	+ "do_sample": False,
	205	+ "eos_token_id": None,
	206	+ "pad_token_id": 0,
	207	+ }
	208	+ cp_output = generate(
	209	+ ContextParallelTinyAttentionLM(args.vocab_size, rank, world_size),
	210	+ input_ids,
	211	+ GenerationConfig(
	212	+ **config_kwargs,
	213	+ use_cache=True,
	214	+ context_parallel_cache=True,
	215	+ context_process_group=dist.group.WORLD,
	216	+ context_logits_rank=world_size - 1,
	217	+ ),
	218	+ attention_mask=attention_mask,
	219	+ )
	220	+ gathered_outputs = _gather_token_outputs(cp_output, world_size)
	221	+
	222	+ prefix_ids = torch.tensor([[7, 2], [3, 7]], dtype=torch.long)
	223	+ prefix_mask = torch.ones_like(prefix_ids)
	224	+ suffix_ids = torch.tensor([[1], [1]], dtype=torch.long)
	225	+ suffix_mask = torch.ones_like(suffix_ids)
	226	+ prefix_outputs = FullTinyAttentionLM(args.vocab_size)(
	227	+ input_ids=prefix_ids,
	228	+ attention_mask=torch.zeros(prefix_ids.size(0), 1, prefix_ids.size(1), prefix_ids.size(1)),
	229	+ use_cache=True,
	230	+ )
	231	+ cp_prefix_output = generate(
	232	+ ContextParallelTinyAttentionLM(args.vocab_size, rank, world_size),
	233	+ suffix_ids,
	234	+ GenerationConfig(
	235	+ **config_kwargs,
	236	+ use_cache=True,
	237	+ prefix_past_key_values=prefix_outputs["past_key_values"],
	238	+ prefix_attention_mask=prefix_mask,
	239	+ context_parallel_cache=True,
	240	+ context_process_group=dist.group.WORLD,
	241	+ context_logits_rank=world_size - 1,
	242	+ ),
	243	+ attention_mask=suffix_mask,
	244	+ )
	245	+ gathered_prefix_outputs = _gather_token_outputs(cp_prefix_output, world_size)
	246	+
	247	+ if rank == 0:
	248	+ baseline_output = generate(
	249	+ FullTinyAttentionLM(args.vocab_size),
	250	+ input_ids,
	251	+ GenerationConfig(**config_kwargs, use_cache=True),
	252	+ attention_mask=attention_mask,
	253	+ )
	254	+ prefix_baseline_output = generate(
	255	+ FullTinyAttentionLM(args.vocab_size),
	256	+ torch.cat([prefix_ids, suffix_ids], dim=-1),
	257	+ GenerationConfig(**config_kwargs, use_cache=True),
	258	+ attention_mask=torch.cat([prefix_mask, suffix_mask], dim=-1),
	259	+ )
	260	+ expected_prefix_output = prefix_baseline_output[:, prefix_ids.shape[-1]:]
	261	+ result = {
	262	+ "world_size": world_size,
	263	+ "max_new_tokens": args.max_new_tokens,
	264	+ "cp_outputs": gathered_outputs,
	265	+ "baseline_output": baseline_output.tolist(),
	266	+ "all_ranks_match_baseline": all(
	267	+ output == baseline_output.tolist()
	268	+ for output in gathered_outputs
	269	+ ),
	270	+ "cp_prefix_outputs": gathered_prefix_outputs,
	271	+ "prefix_baseline_output": expected_prefix_output.tolist(),
	272	+ "prefix_all_ranks_match_baseline": all(
	273	+ output == expected_prefix_output.tolist()
	274	+ for output in gathered_prefix_outputs
	275	+ ),
	276	+ }
	277	+ print(json.dumps(result, indent=2))
	278	+ if args.output:
	279	+ output_path = Path(args.output)
	280	+ output_path.parent.mkdir(parents=True, exist_ok=True)
	281	+ output_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
	282	+ if not result["all_ranks_match_baseline"] or not result["prefix_all_ranks_match_baseline"]:
	283	+ raise SystemExit(1)
	284	+ dist.barrier()
	285	+
	286	+
	287	+if __name__ == "__main__":
	288	+ main()

Aexamples/generate/hf_alignment_check.py+267-0

@@ -0,0 +1,267 @@
	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	+"""Compare hyper generate output with HuggingFace greedy generate."""
	16	+import argparse
	17	+import importlib
	18	+import json
	19	+import sys
	20	+from pathlib import Path
	21	+
	22	+import torch
	23	+
	24	+_REPO_ROOT = Path(__file__).resolve().parents[2]
	25	+if str(_REPO_ROOT) not in sys.path:
	26	+ sys.path.insert(0, str(_REPO_ROOT))
	27	+
	28	+_infer = importlib.import_module("hyper_parallel.infer")
	29	+GenerationConfig = _infer.GenerationConfig
	30	+generate = _infer.generate
	31	+
	32	+
	33	+class HFGenerateAdapter(torch.nn.Module):
	34	+ """Adapter exposing HuggingFace causal LM output to hyper generate."""
	35	+
	36	+ def __init__(self, model):
	37	+ super().__init__()
	38	+ self.model = model
	39	+
	40	+ def forward(
	41	+ self,
	42	+ input_ids,
	43	+ position_ids=None,
	44	+ attention_mask=None,
	45	+ past_key_values=None,
	46	+ use_cache=True,
	47	+ ):
	48	+ """Forward inputs to a HuggingFace causal LM."""
	49	+ del attention_mask
	50	+ outputs = self.model(
	51	+ input_ids=input_ids,
	52	+ position_ids=position_ids,
	53	+ past_key_values=past_key_values,
	54	+ use_cache=use_cache,
	55	+ )
	56	+ return {
	57	+ "logits": outputs.logits,
	58	+ "past_key_values": outputs.past_key_values,
	59	+ }
	60	+
	61	+
	62	+def _resolve_device(name: str) -> torch.device:
	63	+ if name != "auto":
	64	+ return torch.device(name)
	65	+ if hasattr(torch, "npu") and torch.npu.is_available():
	66	+ return torch.device("npu")
	67	+ if torch.cuda.is_available():
	68	+ return torch.device("cuda")
	69	+ return torch.device("cpu")
	70	+
	71	+
	72	+def _compare_cache_no_cache_logits(model, input_ids, generated_ids, max_steps: int):
	73	+ """Compare cache and no-cache decode logits with cosine similarity."""
	74	+ steps = min(max_steps, generated_ids.size(1))
	75	+ if steps <= 0:
	76	+ return []
	77	+
	78	+ values = []
	79	+ context_ids = input_ids
	80	+ cache_outputs = model(input_ids=context_ids, use_cache=True)
	81	+ past_key_values = cache_outputs.past_key_values
	82	+ cache_logits = cache_outputs.logits[:, -1, :]
	83	+
	84	+ for step in range(steps):
	85	+ no_cache_outputs = model(input_ids=context_ids, use_cache=False)
	86	+ no_cache_logits = no_cache_outputs.logits[:, -1, :]
	87	+ similarity = torch.nn.functional.cosine_similarity(
	88	+ cache_logits.float(),
	89	+ no_cache_logits.float(),
	90	+ dim=-1,
	91	+ )
	92	+ values.extend(similarity.detach().cpu().tolist())
	93	+
	94	+ next_token = generated_ids[:, step:step + 1]
	95	+ context_ids = torch.cat([context_ids, next_token], dim=-1)
	96	+ if step == steps - 1:
	97	+ break
	98	+ cache_outputs = model(
	99	+ input_ids=next_token,
	100	+ past_key_values=past_key_values,
	101	+ use_cache=True,
	102	+ )
	103	+ past_key_values = cache_outputs.past_key_values
	104	+ cache_logits = cache_outputs.logits[:, -1, :]
	105	+ return values
	106	+
	107	+
	108	+def parse_args():
	109	+ """Parse alignment check command-line arguments."""
	110	+ parser = argparse.ArgumentParser(description=__doc__)
	111	+ parser.add_argument("--model", required=True, help="Local HF model path or model id")
	112	+ parser.add_argument("--prompt", default="Hello, my name is")
	113	+ parser.add_argument("--max-new-tokens", type=int, default=8)
	114	+ parser.add_argument("--logits-compare-steps", type=int, default=16)
	115	+ parser.add_argument("--logits-cosine-threshold", type=float, default=0.999)
	116	+ parser.add_argument("--device", default="auto")
	117	+ parser.add_argument("--trust-remote-code", action="store_true")
	118	+ parser.add_argument("--output", default=None)
	119	+ return parser.parse_args()
	120	+
	121	+
	122	+def _validate_args(args) -> None:
	123	+ if args.max_new_tokens <= 0:
	124	+ raise ValueError("max-new-tokens must be positive")
	125	+ if args.logits_compare_steps < 0:
	126	+ raise ValueError("logits-compare-steps must be >= 0")
	127	+
	128	+
	129	+def _load_model_and_tokenizer(args, device):
	130	+ """Load the HuggingFace baseline model and tokenizer."""
	131	+ transformers = importlib.import_module("transformers")
	132	+
	133	+ tokenizer = transformers.AutoTokenizer.from_pretrained(
	134	+ args.model,
	135	+ trust_remote_code=args.trust_remote_code,
	136	+ )
	137	+ model = transformers.AutoModelForCausalLM.from_pretrained(
	138	+ args.model,
	139	+ trust_remote_code=args.trust_remote_code,
	140	+ ).to(device)
	141	+ model.eval()
	142	+ return model, tokenizer, HFGenerateAdapter(model)
	143	+
	144	+
	145	+def _run_alignment(args, model, adapter, input_ids, pad_token_id, eos_token_id):
	146	+ """Run HF generate, hyper cache generate, and hyper no-cache generate."""
	147	+ hf_output = model.generate(
	148	+ input_ids=input_ids,
	149	+ max_new_tokens=args.max_new_tokens,
	150	+ do_sample=False,
	151	+ use_cache=True,
	152	+ pad_token_id=pad_token_id,
	153	+ eos_token_id=eos_token_id,
	154	+ )
	155	+ hyper_cache_output = generate(
	156	+ adapter,
	157	+ input_ids,
	158	+ GenerationConfig(
	159	+ max_new_tokens=args.max_new_tokens,
	160	+ do_sample=False,
	161	+ eos_token_id=eos_token_id,
	162	+ pad_token_id=pad_token_id,
	163	+ use_cache=True,
	164	+ ),
	165	+ )
	166	+ hyper_no_cache_output = generate(
	167	+ adapter,
	168	+ input_ids,
	169	+ GenerationConfig(
	170	+ max_new_tokens=args.max_new_tokens,
	171	+ do_sample=False,
	172	+ eos_token_id=eos_token_id,
	173	+ pad_token_id=pad_token_id,
	174	+ use_cache=False,
	175	+ ),
	176	+ )
	177	+ generated_ids = hyper_cache_output[:, input_ids.shape[1]:]
	178	+ logits_cosine = _compare_cache_no_cache_logits(
	179	+ model,
	180	+ input_ids,
	181	+ generated_ids,
	182	+ args.logits_compare_steps,
	183	+ )
	184	+ return hf_output, hyper_cache_output, hyper_no_cache_output, logits_cosine
	185	+
	186	+
	187	+def _build_result(
	188	+ args,
	189	+ tokenizer,
	190	+ input_ids,
	191	+ outputs,
	192	+ logits_cosine,
	193	+ device,
	194	+):
	195	+ """Build the JSON result for the alignment check."""
	196	+ hf_output, hyper_cache_output, hyper_no_cache_output = outputs
	197	+ generated_ids = hyper_cache_output[:, input_ids.shape[1]:]
	198	+ logits_cosine_min = min(logits_cosine) if logits_cosine else None
	199	+ logits_cosine_pass = (
	200	+ logits_cosine_min is None
	201	+ or logits_cosine_min >= args.logits_cosine_threshold
	202	+ )
	203	+ return {
	204	+ "model": args.model,
	205	+ "prompt": args.prompt,
	206	+ "device": str(device),
	207	+ "max_new_tokens": args.max_new_tokens,
	208	+ "logits_compare_steps": min(args.logits_compare_steps, generated_ids.size(1)),
	209	+ "logits_cosine_similarity": logits_cosine,
	210	+ "logits_cosine_min": logits_cosine_min,
	211	+ "logits_cosine_threshold": args.logits_cosine_threshold,
	212	+ "logits_cosine_pass": logits_cosine_pass,
	213	+ "hf_text": tokenizer.decode(hf_output[0], skip_special_tokens=True),
	214	+ "hyper_cache_text": tokenizer.decode(hyper_cache_output[0], skip_special_tokens=True),
	215	+ "hyper_no_cache_text": tokenizer.decode(hyper_no_cache_output[0], skip_special_tokens=True),
	216	+ "hf_new_text": tokenizer.decode(
	217	+ hf_output[0, input_ids.shape[1]:],
	218	+ skip_special_tokens=True,
	219	+ ),
	220	+ "generated_new_tokens": int(hf_output.shape[1] - input_ids.shape[1]),
	221	+ "hf_vs_hyper_cache_ids_match": torch.equal(hf_output.cpu(), hyper_cache_output.cpu()),
	222	+ "hyper_cache_vs_no_cache_ids_match": torch.equal(
	223	+ hyper_cache_output.cpu(),
	224	+ hyper_no_cache_output.cpu(),
	225	+ ),
	226	+ "hf_ids": hf_output[0].detach().cpu().tolist(),
	227	+ "hyper_cache_ids": hyper_cache_output[0].detach().cpu().tolist(),
	228	+ "hyper_no_cache_ids": hyper_no_cache_output[0].detach().cpu().tolist(),
	229	+ }
	230	+
	231	+
	232	+def _check_result(result) -> None:
	233	+ if result["generated_new_tokens"] <= 0:
	234	+ raise AssertionError("model did not generate new tokens")
	235	+ if not result["hf_vs_hyper_cache_ids_match"]:
	236	+ raise AssertionError("hyper generate cache output differs from HuggingFace generate")
	237	+ if not result["hyper_cache_vs_no_cache_ids_match"]:
	238	+ raise AssertionError("hyper generate cache and no-cache outputs differ")
	239	+ if not result["logits_cosine_pass"]:
	240	+ raise AssertionError("cache/no-cache logits cosine similarity is below threshold")
	241	+
	242	+
	243	+def main():
	244	+ """Run HuggingFace and HyperParallel generate alignment."""
	245	+ args = parse_args()
	246	+ _validate_args(args)
	247	+ device = _resolve_device(args.device)
	248	+ model, tokenizer, adapter = _load_model_and_tokenizer(args, device)
	249	+ encoded = tokenizer(args.prompt, return_tensors="pt")
	250	+ input_ids = encoded["input_ids"].to(device)
	251	+ pad_token_id = tokenizer.pad_token_id
	252	+ if pad_token_id is None:
	253	+ pad_token_id = tokenizer.eos_token_id or 0
	254	+ eos_token_id = tokenizer.eos_token_id
	255	+
	256	+ with torch.no_grad():
	257	+ outputs = _run_alignment(args, model, adapter, input_ids, pad_token_id, eos_token_id)
	258	+ result = _build_result(args, tokenizer, input_ids, outputs[:3], outputs[3], device)
	259	+ text = json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True)
	260	+ print(text)
	261	+ if args.output:
	262	+ Path(args.output).write_text(text + "\n", encoding="utf-8")
	263	+ _check_result(result)
	264	+
	265	+
	266	+if __name__ == "__main__":
	267	+ main()

Ahyper_parallel/infer/__init__.py+65-0

@@ -0,0 +1,65 @@
	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.
	阿苏6月22日目录名改成叫infer Moy 6月23日评论： 6月23日评论：
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Autoregressive generation utilities."""
	16	+
	17	+from hyper_parallel.infer.generation import generate
	18	+from hyper_parallel.infer.kv_cache import (
	19	+ ContextParallelKVCache,
	20	+ KVCache,
	21	+ SequenceShardInfo,
	22	+ get_sequence_shard_info,
	23	+ shard_past_key_values,
	24	+)
	25	+from hyper_parallel.infer.mixin import GenerateMixin
	26	+from hyper_parallel.infer.sampler import (
	27	+ apply_repetition_penalty,
	28	+ greedy_sample,
	29	+ sample_next_token,
	30	+ top_k_sample,
	31	+ top_p_sample,
	32	+)
	33	+from hyper_parallel.infer.utils import (
	34	+ GenerationConfig,
	35	+ apply_logits_processors,
	36	+ build_causal_mask,
	37	+ build_position_ids,
	38	+ gather_context_parallel_logits,
	39	+ gather_tensor_parallel_logits,
	40	+ prepare_logits_for_sampling,
	41	+ should_stop_generation,
	42	+)
	43	+
	44	+__all__ = [
	45	+ "GenerationConfig",
	46	+ "GenerateMixin",
	47	+ "KVCache",
	48	+ "ContextParallelKVCache",
	49	+ "SequenceShardInfo",
	50	+ "apply_logits_processors",
	51	+ "apply_repetition_penalty",
	52	+ "build_causal_mask",
	53	+ "build_position_ids",
	54	+ "gather_context_parallel_logits",
	55	+ "gather_tensor_parallel_logits",
	56	+ "get_sequence_shard_info",
	57	+ "generate",
	58	+ "greedy_sample",
	59	+ "prepare_logits_for_sampling",
	60	+ "sample_next_token",
	61	+ "shard_past_key_values",
	62	+ "should_stop_generation",
	63	+ "top_k_sample",
	64	+ "top_p_sample",
	65	+]

Ahyper_parallel/infer/generation.py+649-0

@@ -0,0 +1,649 @@
	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	+"""Prefill + decode generation loop."""
	16	+import inspect
	17	+from typing import Optional
	18	+
	19	+import torch
	20	+import torch.distributed as dist
	21	+
	22	+from hyper_parallel.infer.kv_cache import (
	23	+ ContextParallelKVCache,
	24	+ KVCache,
	25	+ detach_and_validate_past_key_values,
	26	+)
	27	+from hyper_parallel.infer.sampler import sample_next_token
	28	+from hyper_parallel.infer.utils import (
	29	+ GenerationConfig,
	30	+ append_attention_mask,
	31	+ apply_logits_processors,
	32	+ build_causal_mask,
	33	+ build_position_ids,
	34	+ prepare_logits_for_sampling,
	35	+ should_stop_generation,
	36	+)
	37	+
	38	+
	39	+def _get_output(outputs, name: str):
	40	+ """Read an output field from dict-like or object-like model outputs."""
	41	+ if isinstance(outputs, dict):
	42	+ return outputs.get(name)
	43	+ return getattr(outputs, name, None)
	44	+
	45	+
	46	+def _model_forward(
	47	+ model,
	48	+ input_ids: torch.Tensor,
	49	+ position_ids: torch.Tensor,
	50	+ attention_mask: Optional[torch.Tensor],
	51	+ past_key_values,
	52	+ use_cache: bool,
	53	+ sequence_shard_info=None,
	54	+ global_seq_len: Optional[int] = None,
	55	+):
	56	+ """Call model.forward with only the keyword arguments it accepts."""
	57	+ kwargs = {
	58	+ "input_ids": input_ids,
	59	+ "position_ids": position_ids,
	60	+ "attention_mask": attention_mask,
	61	+ "past_key_values": past_key_values,
	62	+ "use_cache": use_cache,
	63	+ "sequence_shard_info": sequence_shard_info,
	64	+ "global_seq_len": global_seq_len,
	65	+ }
	66	+ forward = getattr(model, "forward", model)
	67	+ try:
	68	+ signature = inspect.signature(forward)
	69	+ except (TypeError, ValueError):
	70	+ return forward(**kwargs)
	71	+ parameters = signature.parameters
	72	+ accepts_kwargs = any(
	73	+ param.kind == inspect.Parameter.VAR_KEYWORD
	74	+ for param in parameters.values()
	75	+ )
	76	+ if not accepts_kwargs:
	77	+ for name in list(kwargs):
	78	+ if name not in parameters:
	79	+ kwargs.pop(name)
	80	+ return forward(**kwargs)
	81	+
	82	+
	83	+def _resolve_context_parallel_rank_world(config: GenerationConfig) -> tuple[int, int]:
	84	+ if config.context_parallel_rank is not None:
	85	+ return config.context_parallel_rank, config.context_parallel_world_size
	86	+ if not dist.is_available() or not dist.is_initialized():
	87	+ raise ValueError(
	88	+ "context_parallel_cache requires initialized torch.distributed "
	89	+ "or explicit context_parallel_rank/context_parallel_world_size",
	90	+ )
	91	+ return (
	92	+ dist.get_rank(group=config.context_process_group),
	93	+ dist.get_world_size(group=config.context_process_group),
	94	+ )
	95	+
	96	+
	97	+def _init_cache(config: GenerationConfig) -> KVCache:
	98	+ if not config.context_parallel_cache:
	99	+ return KVCache()
	100	+ rank, world_size = _resolve_context_parallel_rank_world(config)

Ahyper_parallel/infer/kv_cache.py+283-0

@@ -0,0 +1,283 @@
	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	+"""KV cache container for generation."""
	16	+from dataclasses import dataclass
	17	+from typing import Any, Iterable, List, Optional, Tuple
	18	+
	19	+import torch
	20	+
	21	+PastKeyValues = List[Tuple[torch.Tensor, torch.Tensor]]
	22	+
	23	+
	24	+def _validate_pair_shapes(key: torch.Tensor, value: torch.Tensor) -> None:
	25	+ """Validate one key/value cache tensor pair."""
	26	+ if not isinstance(key, torch.Tensor) or not isinstance(value, torch.Tensor):
	27	+ raise ValueError("key and value must be tensors")
	28	+ if key.ndim != 4 or value.ndim != 4:
	29	+ raise ValueError("key and value must have shape (batch, heads, seq, dim)")
	30	+ if key.shape != value.shape:
	31	+ raise ValueError("key and value batch/heads/seq/dim dimensions must match")
	32	+
	33	+
	34	+def detach_and_validate_past_key_values(past_key_values: Iterable) -> PastKeyValues:
	35	+ """Return detached tuple KV tensors after validating their shapes."""
	36	+ values = []
	37	+ for item in past_key_values:
	38	+ if not isinstance(item, (tuple, list)) or len(item) != 2:
	39	+ raise ValueError("each cache entry must be a (key, value) pair")
	40	+ key, value = item
	41	+ _validate_pair_shapes(key, value)
	42	+ values.append((key.detach(), value.detach()))
	43	+ return values
	44	+
	45	+
	46	+@dataclass(frozen=True)
	47	+class SequenceShardInfo:
	48	+ """Sequence range held by one context-parallel rank."""
	49	+
	50	+ rank: int
	51	+ world_size: int
	52	+ start: int
	53	+ end: int
	54	+ global_seq_len: int
	55	+
	56	+ @property
	57	+ def local_seq_len(self) -> int:
	58	+ """Return the sequence length stored by this rank."""
	59	+ return self.end - self.start
	60	+
	61	+
	62	+def get_sequence_shard_info(
	63	+ global_seq_len: int,
	64	+ rank: int,
	65	+ world_size: int,
	66	+) -> SequenceShardInfo:
	67	+ """Return the contiguous sequence shard range for a CP rank."""
	68	+ if global_seq_len < 0:
	69	+ raise ValueError("global_seq_len must be >= 0")
	70	+ if world_size <= 0:
	71	+ raise ValueError("world_size must be > 0")
	72	+ if rank < 0 or rank >= world_size:
	73	+ raise ValueError("rank must be in [0, world_size)")
	74	+ base = global_seq_len // world_size
	75	+ remainder = global_seq_len % world_size
	76	+ start = rank * base + min(rank, remainder)
	77	+ end = start + base + (1 if rank < remainder else 0)
	78	+ return SequenceShardInfo(
	79	+ rank=rank,
	80	+ world_size=world_size,
	81	+ start=start,
	82	+ end=end,
	83	+ global_seq_len=global_seq_len,
	84	+ )
	85	+
	86	+
	87	+def shard_past_key_values(
	88	+ past_key_values: Iterable,
	89	+ rank: int,
	90	+ world_size: int,
	91	+ global_seq_len: Optional[int] = None,
	92	+) -> Tuple[PastKeyValues, SequenceShardInfo]:
	93	+ """Shard full past key values on the sequence dimension for CP cache."""
	94	+ values = detach_and_validate_past_key_values(past_key_values)
	95	+ if not values:
	96	+ seq_len = 0 if global_seq_len is None else global_seq_len
	97	+ else:
	98	+ seq_len = values[0][0].shape[-2]
	99	+ if global_seq_len is None:
	100	+ global_seq_len = seq_len
	101	+ if seq_len != global_seq_len:
	102	+ raise ValueError("global_seq_len must match full cache sequence length")
	103	+ shard_info = get_sequence_shard_info(global_seq_len, rank, world_size)
	104	+ sharded = [
	105	+ (
	106	+ key.narrow(-2, shard_info.start, shard_info.local_seq_len).contiguous(),
	107	+ value.narrow(-2, shard_info.start, shard_info.local_seq_len).contiguous(),
	108	+ )
	109	+ for key, value in values
	110	+ ]
	111	+ return sharded, shard_info
	112	+
	113	+
	114	+class KVCache:
	115	+ """Stores per-layer key/value tensors."""
	116	+
	117	+ def __init__(self):
	118	+ self.past_key_values: Optional[Any] = None
	119	+
	120	+ @property
	121	+ def is_empty(self) -> bool:
	122	+ """Check whether no usable KV cache is stored."""
	123	+ return self.past_key_values is None or (
	124	+ isinstance(self.past_key_values, list) and len(self.past_key_values) == 0
	125	+ )
	126	+
	127	+ def clear(self) -> None:
	128	+ """Drop all cached tensors."""
	129	+ self.past_key_values = None
	130	+
	131	+ def update(self, past_key_values: Optional[Iterable]) -> None:
	132	+ """Replace the cache with detached past key values."""
	133	+ if past_key_values is None:
	134	+ return
	135	+ if self._is_opaque_cache(past_key_values):
	136	+ self.past_key_values = past_key_values
	137	+ return
	138	+ values = self._detach_and_validate(past_key_values)
	139	+ self.past_key_values = None if not values else values
	140	+
	141	+ def merge(self, past_key_values: Optional[Iterable]) -> None:
	142	+ """Append incremental key/value tensors on the sequence dimension."""
	143	+ if past_key_values is None:
	144	+ return
	145	+ if self._is_opaque_cache(past_key_values):
	146	+ self.past_key_values = past_key_values
	147	+ return
	148	+ new_values = self._detach_and_validate(past_key_values)
	149	+ if not new_values:
	150	+ return
	151	+ if self.past_key_values is None:
	152	+ self.past_key_values = new_values
	153	+ return
	154	+ if len(self.past_key_values) != len(new_values):
	155	+ raise ValueError("past_key_values layer count mismatch")
	156	+ merged = []
	157	+ for (old_k, old_v), (new_k, new_v) in zip(self.past_key_values, new_values):
	158	+ self._validate_pair_shapes(old_k, old_v)
	159	+ self._validate_pair_shapes(new_k, new_v)
	160	+ if old_k.shape[:-2] != new_k.shape[:-2] or old_k.shape[-1] != new_k.shape[-1]:
	161	+ raise ValueError("key cache shape mismatch")
	162	+ if old_v.shape[:-2] != new_v.shape[:-2] or old_v.shape[-1] != new_v.shape[-1]:
	163	+ raise ValueError("value cache shape mismatch")
	164	+ merged.append((
	165	+ torch.cat([old_k, new_k], dim=-2),
	166	+ torch.cat([old_v, new_v], dim=-2),
	167	+ ))
	168	+ self.past_key_values = merged
	169	+
	170	+ @classmethod
	171	+ def _detach_and_validate(cls, past_key_values: Iterable) -> PastKeyValues:
	172	+ return detach_and_validate_past_key_values(past_key_values)
	173	+
	174	+ @staticmethod
	175	+ def _validate_pair_shapes(key: torch.Tensor, value: torch.Tensor) -> None:
	176	+ _validate_pair_shapes(key, value)
	177	+
	178	+ @staticmethod
	179	+ def _is_opaque_cache(past_key_values) -> bool:
	180	+ return hasattr(past_key_values, "get_seq_length") and not isinstance(
	181	+ past_key_values, (list, tuple),
	182	+ )
	183	+
	184	+
	185	+class ContextParallelKVCache(KVCache):
	X X xuxinglei6月23日【功能/接口】ContextParallelKVCache 及 shard_past_key_values/update_full/update_local/merge_local 实现完整、单测齐全,但从未接入 generate 主循环(generation.py:167 只用 KVCache())。它们已从 init 导出为公共 API,会让用户以为支持 CP 缓存推理,实际不可达。建议:要么接线进 generate() + 补 e2e 测试(断言 CP 下 logits 一致),要么先从公共导出移除,待真正使用它的 PR 一起落地,避免发布 dead public API。 Moy 6月24日评论： 6月24日评论：
	186	+ """Stores a local sequence shard of generation KV cache."""
	187	+
	188	+ def __init__(self, rank: int, world_size: int):
	189	+ super().__init__()
	190	+ if world_size <= 0:
	191	+ raise ValueError("world_size must be > 0")
	192	+ if rank < 0 or rank >= world_size:
	193	+ raise ValueError("rank must be in [0, world_size)")
	194	+ self.rank = rank
	195	+ self.world_size = world_size
	196	+ self.shard_info = get_sequence_shard_info(0, rank, world_size)
	197	+
	198	+ def update_full(self, past_key_values: Optional[Iterable]) -> None:
	199	+ """Shard full prefill K/V cache and store only this rank's sequence slice."""
	200	+ if past_key_values is None:
	201	+ return
	202	+ sharded, shard_info = shard_past_key_values(
	203	+ past_key_values,
	204	+ rank=self.rank,
	205	+ world_size=self.world_size,
	206	+ )
	207	+ self.past_key_values = sharded
	208	+ self.shard_info = shard_info
	209	+
	210	+ def update_local(
	211	+ self,
	212	+ past_key_values: Optional[Iterable],
	213	+ shard_info: SequenceShardInfo,
	214	+ ) -> None:
	215	+ """Store K/V tensors that are already local to this CP rank."""
	216	+ if past_key_values is None:
	217	+ return
	218	+ self._validate_shard_info(shard_info)
	219	+ values = self._detach_and_validate(past_key_values)
	220	+ self._validate_local_seq_len(values, shard_info.local_seq_len)
	221	+ self.past_key_values = values
	222	+ self.shard_info = shard_info
	223	+
	224	+ def merge_local(
	225	+ self,
	226	+ past_key_values: Optional[Iterable],
	227	+ global_seq_len: Optional[int] = None,
	228	+ ) -> None:
	229	+ """Append local incremental K/V tensors and advance global sequence metadata."""
	230	+ if past_key_values is None:
	231	+ return
	232	+ new_values = self._detach_and_validate(past_key_values)
	233	+ if self.past_key_values is None:
	234	+ if global_seq_len is None and self.world_size > 1:
	235	+ raise ValueError("global_seq_len is required for initial CP local cache")
	236	+ inferred_global = (
	237	+ self.shard_info.global_seq_len + new_values[0][0].shape[-2]
	238	+ if global_seq_len is None and new_values
	239	+ else global_seq_len
	240	+ )
	241	+ shard_info = get_sequence_shard_info(
	242	+ 0 if inferred_global is None else inferred_global,
	243	+ self.rank,
	244	+ self.world_size,
	245	+ )
	246	+ self._validate_local_seq_len(new_values, shard_info.local_seq_len)
	247	+ self.past_key_values = new_values
	248	+ self.shard_info = shard_info
	249	+ return
	atomgit-bot6月22日 🟡 Medium Priority 在 `ContextParallelKVCache.merge_local` 首次调用路径（`self.past_key_values is None` 且 `global_seq_len=None`）中，第 217-219 行将 global_seq_len 推断为 `self.shard_info.global_seq_len + new_values[0][0].shape[-2]`，即 `0 + 本地序列长度`。但在 world_size > 1 时，本地序列长度并不等于全局序列长度。例如 rank=0、world_size=2、local_seq_len=3 时，正确的 global 应为 5 或 6（取决于 other rank），而非 3。`get_sequence_shard_info(3, 0, 2).local_seq_len` 返回 2，与 local_seq_len=3 不匹配，会在第 227 行的 `_validate_local_seq_len` 中触发 ValueError。建议：当 `self.past_key_values is None` 且 `global_seq_len` 未提供时，不要尝试从 local_seq_len 反推 global_seq_len（这在 world_size > 1 时不正确），而是要求调用者必须提供 `global_seq_len`，或仅在 world_size == 1 时允许推断（此时 local == global）。不准确？
	250	+ old_local_seq_len = self.shard_info.local_seq_len
	251	+ next_global_seq_len = (
	252	+ self.shard_info.global_seq_len + new_values[0][0].shape[-2]
	253	+ if global_seq_len is None and new_values
	254	+ else global_seq_len
	255	+ )
	256	+ if next_global_seq_len is None:
	257	+ raise ValueError("global_seq_len is required for empty incremental cache")
	258	+ shard_info = get_sequence_shard_info(
	259	+ next_global_seq_len,
	260	+ self.rank,
	261	+ self.world_size,
	262	+ )
	263	+ expected_growth = shard_info.local_seq_len - old_local_seq_len
	264	+ actual_growth = new_values[0][0].shape[-2] if new_values else 0
	265	+ if expected_growth != actual_growth:
	266	+ raise ValueError("local cache growth does not match CP shard metadata")
	267	+ super().merge(new_values)
	268	+ self.shard_info = shard_info
	269	+
	270	+ def clear(self) -> None:
	271	+ """Drop all cached tensors and reset CP sequence metadata."""
	272	+ super().clear()
	273	+ self.shard_info = get_sequence_shard_info(0, self.rank, self.world_size)
	274	+
	275	+ def _validate_shard_info(self, shard_info: SequenceShardInfo) -> None:
	276	+ if shard_info.rank != self.rank or shard_info.world_size != self.world_size:
	277	+ raise ValueError("shard_info does not match this CP cache")
	278	+
	279	+ @staticmethod
	280	+ def _validate_local_seq_len(values: PastKeyValues, local_seq_len: int) -> None:
	281	+ for key, value in values:
	282	+ if key.shape[-2] != local_seq_len or value.shape[-2] != local_seq_len:
	283	+ raise ValueError("local cache sequence length does not match shard_info")

Ahyper_parallel/infer/mixin.py+31-0

@@ -0,0 +1,31 @@
	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	+"""Mixin exposing a model.generate style API."""
	16	+
	17	+from hyper_parallel.infer.generation import generate
	18	+
	19	+
	20	+class GenerateMixin:
	21	+ """Mixin that forwards to :func:`hyper_parallel.infer.generate`."""
	22	+
	23	+ def generate(self, input_ids, generation_config=None, attention_mask=None, **kwargs):
	24	+ """Generate token ids with the common HyperParallel generation loop."""
	25	+ return generate(
	26	+ self,
	27	+ input_ids=input_ids,
	28	+ generation_config=generation_config,
	29	+ attention_mask=attention_mask,
	30	+ **kwargs,
	31	+ )

Ahyper_parallel/infer/sampler.py+109-0

@@ -0,0 +1,109 @@
	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	+"""Sampling helpers for autoregressive generation."""
	16	+import torch
	17	+from torch.nn import functional as F
	18	+
	19	+from hyper_parallel.infer.utils import GenerationConfig
	20	+
	21	+
	22	+def _filter_top_k(logits: torch.Tensor, top_k: int) -> torch.Tensor:
	23	+ if top_k <= 0 or top_k >= logits.size(-1):
	24	+ return logits
	25	+ values, _ = torch.topk(logits, k=top_k, dim=-1)
	26	+ threshold = values[:, -1:].contiguous()
	27	+ return logits.masked_fill(logits < threshold, float("-inf"))
	28	+
	29	+
	30	+def greedy_sample(logits: torch.Tensor) -> torch.Tensor:
	31	+ """Select the highest-logit token for each batch item."""
	32	+ if logits.ndim != 2:
	33	+ raise ValueError("logits must have shape (batch, vocab)")
	34	+ return logits.argmax(dim=-1, keepdim=True)
	35	+
	36	+
	37	+def top_k_sample(
	38	+ logits: torch.Tensor,
	39	+ top_k: int,
	40	+ temperature: float = 1.0,
	41	+) -> torch.Tensor:
	42	+ """Sample from the top-k logits."""
	43	+ if logits.ndim != 2:
	44	+ raise ValueError("logits must have shape (batch, vocab)")
	45	+ filtered = _filter_top_k(logits, top_k)
	46	+ probs = F.softmax(filtered / temperature, dim=-1)
	47	+ sampled = torch.multinomial(probs, num_samples=1)
	48	+ return sampled
	49	+
	50	+
	51	+def top_p_sample(
	52	+ logits: torch.Tensor,
	53	+ top_p: float,
	54	+ temperature: float = 1.0,
	55	+) -> torch.Tensor:
	56	+ """Sample from the nucleus token set."""
	57	+ if logits.ndim != 2:
	58	+ raise ValueError("logits must have shape (batch, vocab)")
	59	+ if top_p >= 1.0:
	60	+ probs = F.softmax(logits / temperature, dim=-1)
	61	+ return torch.multinomial(probs, num_samples=1)
	62	+ sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
	63	+ sorted_probs = F.softmax(sorted_logits / temperature, dim=-1)
	64	+ cumulative = sorted_probs.cumsum(dim=-1)
	65	+ remove = cumulative - sorted_probs > top_p
	66	+ filtered = sorted_logits.masked_fill(remove, float("-inf"))
	67	+ probs = F.softmax(filtered / temperature, dim=-1)
	68	+ sampled = torch.multinomial(probs, num_samples=1)
	69	+ return sorted_indices.gather(dim=-1, index=sampled)
	70	+
	71	+
	72	+def apply_repetition_penalty(
	73	+ logits: torch.Tensor,
	74	+ input_ids: torch.Tensor,
	75	+ penalty: float,
	76	+) -> torch.Tensor:
	77	+ """Apply per-item repetition penalty to seen token ids."""
	78	+ if penalty == 1.0:
	79	+ return logits
	80	+ if logits.ndim != 2 or input_ids.ndim != 2:
	81	+ raise ValueError("logits and input_ids must be 2-D tensors")
	82	+ adjusted = logits.clone()
	83	+ vocab_size = logits.size(-1)
	84	+ valid = (input_ids >= 0) & (input_ids < vocab_size)
	85	+ seen_mask = torch.zeros_like(adjusted, dtype=torch.bool)
	86	+ token_ids = input_ids.to(dtype=torch.long).clamp(min=0, max=max(vocab_size - 1, 0))
	87	+ seen_mask.scatter_(dim=1, index=token_ids, src=valid)
	88	+ penalized = torch.where(adjusted < 0, adjusted * penalty, adjusted / penalty)
	89	+ return torch.where(seen_mask, penalized, adjusted)
	90	+
	91	+
	92	+def sample_next_token(
	93	+ logits: torch.Tensor,
	94	+ input_ids: torch.Tensor,
	95	+ config: GenerationConfig,
	96	+) -> torch.Tensor:
	97	+ """Apply repetition penalty and select the next token."""
	98	+ logits = apply_repetition_penalty(
	99	+ logits,
	100	+ input_ids=input_ids,
	101	+ penalty=config.repetition_penalty,
	102	+ )
	103	+ if not config.do_sample:
	104	+ return greedy_sample(logits)
	105	+ logits = _filter_top_k(logits, config.top_k)
	106	+ if config.top_p < 1.0:
	107	+ return top_p_sample(logits, top_p=config.top_p, temperature=config.temperature)
	X X xuxinglei6月23日【正确性】当 top_p<1.0 时直接 return,config.top_k 被静默忽略,不存在 top_k+top_p 组合路径。HF 默认 top_k=50, top_p=0.9 的配置在这里只生效 top_p,采样分布与配置预期不符;issue 又点名了 greedy/top-k。建议:do_sample 且 top_k>0 时先 top_k 过滤再 top_p nucleus(对齐 HF TopK+TopP warper 顺序),或显式文档声明二者互斥。 Moy 6月24日评论： 6月24日评论：
	108	+ probs = F.softmax(logits / config.temperature, dim=-1)
	109	+ return torch.multinomial(probs, num_samples=1)

Ahyper_parallel/infer/utils.py+280-0

@@ -0,0 +1,280 @@
	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	+"""Generation configuration and mask helpers."""
	16	+from dataclasses import dataclass
	17	+from typing import Any, Callable, List, Optional
	18	+
	19	+import torch
	20	+import torch.distributed as dist
	21	+
	22	+
	23	+@dataclass
	24	+class GenerationConfig:
	25	+ """Runtime options for autoregressive generation."""
	26	+
	27	+ max_new_tokens: int = 256
	28	+ temperature: float = 1.0
	29	+ top_k: int = 50
	30	+ top_p: float = 1.0
	31	+ do_sample: bool = False
	32	+ eos_token_id: Optional[int] = 2
	33	+ pad_token_id: int = 0
	34	+ repetition_penalty: float = 1.0
	35	+ use_cache: bool = True
	36	+ prefix_past_key_values: Optional[Any] = None
	37	+ prefix_attention_mask: Optional[torch.Tensor] = None
	38	+ prefix_sequence_shard_info: Optional[Any] = None
	39	+ prefix_cache_length: Optional[int] = None
	40	+ context_parallel_cache: bool = False
	41	+ context_parallel_rank: Optional[int] = None
	42	+ context_parallel_world_size: Optional[int] = None
	43	+ # context_logits_rank is local to context_process_group.
	44	+ context_logits_rank: Optional[Any] = None
	45	+ context_process_group: Optional[Any] = None
	46	+ gather_logits: bool = False
	47	+ logits_process_group: Optional[Any] = None
	48	+ logits_gather_dim: int = -1
	49	+ mask_dtype: Optional[torch.dtype] = None
	50	+ logits_processor: Optional[List[Callable]] = None
	51	+ stopping_criteria: Optional[List[Callable]] = None
	52	+
	53	+ def __post_init__(self):
	54	+ self._validate_sampling()
	55	+ self._validate_prefix()
	56	+ self._validate_context_parallel()
	57	+ if self.logits_gather_dim >= 0:
	58	+ raise ValueError("logits_gather_dim must be negative")
	59	+ if self.mask_dtype is not None and not isinstance(self.mask_dtype, torch.dtype):
	60	+ raise ValueError("mask_dtype must be a torch.dtype")
	61	+ self._validate_callables(self.logits_processor, "logits_processor")
	62	+ self._validate_callables(self.stopping_criteria, "stopping_criteria")
	63	+
	64	+ def _validate_sampling(self) -> None:
	65	+ """Validate scalar sampling and stopping options."""
	66	+ if self.max_new_tokens < 0:
	67	+ raise ValueError("max_new_tokens must be >= 0")
	68	+ if self.temperature <= 0:
	69	+ raise ValueError("temperature must be > 0")
	70	+ if self.top_k < 0:
	71	+ raise ValueError("top_k must be >= 0")
	72	+ if not 0 < self.top_p <= 1.0:
	73	+ raise ValueError("top_p must be in (0, 1]")
	74	+ if self.repetition_penalty <= 0:
	75	+ raise ValueError("repetition_penalty must be > 0")
	76	+
	77	+ def _validate_prefix(self) -> None:
	78	+ """Validate optional prefix cache metadata."""
	79	+ if self.prefix_past_key_values is None:
	80	+ if self.prefix_attention_mask is not None:
	81	+ raise ValueError("prefix_attention_mask requires prefix_past_key_values")
	82	+ if self.prefix_sequence_shard_info is not None:
	83	+ raise ValueError("prefix_sequence_shard_info requires prefix_past_key_values")
	84	+ if self.prefix_cache_length is not None:
	85	+ raise ValueError("prefix_cache_length requires prefix_past_key_values")
	86	+ else:
	87	+ if not self.use_cache:
	88	+ raise ValueError("prefix_past_key_values requires use_cache=True")
	89	+ if self.prefix_cache_length is not None and self.prefix_cache_length < 0:
	90	+ raise ValueError("prefix_cache_length must be >= 0")
	91	+ if self.prefix_sequence_shard_info is not None and not self.context_parallel_cache:
	92	+ raise ValueError("prefix_sequence_shard_info requires context_parallel_cache=True")
	93	+
	94	+ def _validate_context_parallel(self) -> None:
	95	+ """Validate context-parallel cache and logits metadata."""
	96	+ if self.context_parallel_cache and not self.use_cache:
	97	+ raise ValueError("context_parallel_cache requires use_cache=True")
	98	+ if (self.context_parallel_rank is None) != (self.context_parallel_world_size is None):
	99	+ raise ValueError(
	100	+ "context_parallel_rank and context_parallel_world_size must be set together",
	101	+ )
	102	+ if self.context_parallel_world_size is not None:
	103	+ if self.context_parallel_world_size <= 0:
	104	+ raise ValueError("context_parallel_world_size must be > 0")
	105	+ if (
	106	+ self.context_parallel_rank < 0
	107	+ or self.context_parallel_rank >= self.context_parallel_world_size
	108	+ ):
	109	+ raise ValueError("context_parallel_rank must be in [0, context_parallel_world_size)")
	110	+
	111	+ @staticmethod
	112	+ def _validate_callables(values: Optional[List[Callable]], field_name: str) -> None:
	113	+ """Validate optional generation extension hooks."""
	114	+ if values is None:
	115	+ return
	116	+ if not isinstance(values, list):
	117	+ raise ValueError(f"{field_name} must be a list of callables")
	118	+ if not all(callable(item) for item in values):
	119	+ raise ValueError(f"{field_name} must contain only callables")
	120	+
	121	+
	122	+def build_position_ids(
	123	+ input_ids: torch.Tensor,
	124	+ attention_mask: Optional[torch.Tensor] = None,
	125	+) -> torch.Tensor:
	126	+ """Build left-padding aware position ids."""
	127	+ if input_ids.ndim != 2:
	128	+ raise ValueError("input_ids must have shape (batch, seq)")
	129	+ if attention_mask is None:
	130	+ seq_len = input_ids.size(1)
	131	+ return torch.arange(
	132	+ seq_len, device=input_ids.device, dtype=torch.long,
	133	+ ).view(1, -1).expand(input_ids.size(0), -1)
	134	+ if attention_mask.shape != input_ids.shape:
	135	+ raise ValueError("attention_mask must match input_ids shape")
	136	+ position_ids = attention_mask.long().cumsum(dim=-1) - 1
	137	+ return position_ids.clamp_min_(0)
	138	+
	139	+
	140	+def gather_context_parallel_logits(logits: torch.Tensor, config: GenerationConfig) -> torch.Tensor:
	141	+ """Select final-token logits from the owning CP rank before sampling."""
	142	+ if config.context_logits_rank is None:
	143	+ return logits
	144	+ if not dist.is_available() or not dist.is_initialized():
	145	+ return logits
	146	+ world_size = dist.get_world_size(group=config.context_process_group)
	147	+ if world_size == 1:
	148	+ return logits
	149	+ gathered = [torch.empty_like(logits) for _ in range(world_size)]
	150	+ dist.all_gather(gathered, logits, group=config.context_process_group)
	151	+ stacked = torch.stack(gathered, dim=0)
	152	+ owner = torch.as_tensor(
	153	+ config.context_logits_rank,
	154	+ device=logits.device,
	155	+ dtype=torch.long,
	156	+ )
	157	+ if owner.ndim == 0:
	158	+ owner_rank = int(owner.item())
	159	+ if owner_rank < 0 or owner_rank >= world_size:
	160	+ raise ValueError("context_logits_rank contains an invalid rank")
	161	+ return stacked[owner_rank]
	162	+ if owner.shape != (logits.shape[0],):
	163	+ raise ValueError("context_logits_rank must be a scalar or a batch-sized tensor")
	164	+ if torch.any((owner < 0) \| (owner >= world_size)):
	165	+ raise ValueError("context_logits_rank contains an invalid rank")
	166	+ batch_indices = torch.arange(logits.shape[0], device=logits.device)
	167	+ return stacked[owner, batch_indices]
	168	+
	169	+
	170	+def gather_tensor_parallel_logits(logits: torch.Tensor, config: GenerationConfig) -> torch.Tensor:
	171	+ """Gather vocab-sharded logits before sampling when TP inference is active."""
	172	+ if not config.gather_logits:
	173	+ return logits
	174	+ if not dist.is_available() or not dist.is_initialized():
	175	+ return logits
	176	+ world_size = dist.get_world_size(group=config.logits_process_group)
	177	+ if world_size == 1:
	178	+ return logits
	179	+ gather_dim = logits.ndim + config.logits_gather_dim
	180	+ if gather_dim < 0 or gather_dim >= logits.ndim:
	181	+ raise ValueError("logits_gather_dim is out of range for logits")
	182	+ local_shard = torch.tensor(
	183	+ [logits.shape[gather_dim]],
	184	+ device=logits.device,
	185	+ dtype=torch.long,
	186	+ )
	187	+ shard_sizes = [torch.empty_like(local_shard) for _ in range(world_size)]
	188	+ dist.all_gather(shard_sizes, local_shard, group=config.logits_process_group)
	189	+ shard_sizes = torch.cat(shard_sizes)
	190	+ if torch.any(shard_sizes != shard_sizes[0]):
	191	+ raise ValueError(
	192	+ "tensor-parallel logits gather requires equal local vocab shard sizes; "
	193	+ "pad vocab shards before generation",
	194	+ )
	195	+ gathered = [torch.empty_like(logits) for _ in range(world_size)]
	X X xuxinglei6月23日【正确性】所有 gather slot 用 `torch.empty_like(logits)`,即调用 rank 自己的分片形状。当 vocab_size % tp_world != 0 时各 rank 分片大小不同,all_gather 形状不一致直接崩(底层 collective 报错很晦涩)。建议:gather 前显式校验各 rank vocab 分片一致,或要求 vocab pad 到 tp_world 的整数倍并在不满足时抛带说明的 ValueError。 Moy 6月24日评论： 6月24日评论：
	196	+ dist.all_gather(gathered, logits, group=config.logits_process_group)
	197	+ return torch.cat(gathered, dim=config.logits_gather_dim)
	198	+
	199	+
	200	+def prepare_logits_for_sampling(logits: torch.Tensor, config: GenerationConfig) -> torch.Tensor:
	201	+ """Apply distributed logits handoffs before sampling."""
	202	+ logits = gather_context_parallel_logits(logits, config)
	203	+ return gather_tensor_parallel_logits(logits, config)
	204	+
	205	+
	206	+def apply_logits_processors(
	207	+ input_ids: torch.Tensor,
	208	+ logits: torch.Tensor,
	209	+ config: GenerationConfig,
	210	+) -> torch.Tensor:
	211	+ """Apply user-supplied logits processors in order."""
	212	+ if config.logits_processor is None:
	213	+ return logits
	214	+ processed = logits
	215	+ for processor in config.logits_processor:
	216	+ processed = processor(input_ids, processed)
	217	+ if not isinstance(processed, torch.Tensor):
	218	+ raise ValueError("logits_processor must return a tensor")
	219	+ return processed
	220	+
	221	+
	222	+def should_stop_generation(
	223	+ input_ids: torch.Tensor,
	224	+ logits: torch.Tensor,
	225	+ config: GenerationConfig,
	226	+) -> bool:
	227	+ """Return whether any configured stopping criterion requests termination."""
	228	+ if config.stopping_criteria is None:
	229	+ return False
	230	+ for criterion in config.stopping_criteria:
	231	+ result = criterion(input_ids, logits)
	232	+ if isinstance(result, torch.Tensor):
	233	+ if result.numel() != 1:
	234	+ raise ValueError("stopping_criteria tensor output must be scalar")
	235	+ result = bool(result.item())
	236	+ if bool(result):
	237	+ return True
	238	+ return False
	239	+
	240	+
	241	+def build_causal_mask(
	242	+ input_ids: torch.Tensor,
	243	+ attention_mask: Optional[torch.Tensor] = None,
	244	+ dtype: torch.dtype = torch.float32,
	245	+) -> torch.Tensor:
	246	+ """Build additive causal + padding mask for prefill."""
	247	+ if input_ids.ndim != 2:
	248	+ raise ValueError("input_ids must have shape (batch, seq)")
	249	+ batch_size, seq_len = input_ids.shape
	250	+ device = input_ids.device
	251	+ mask = torch.zeros(
	252	+ batch_size, 1, seq_len, seq_len, device=device, dtype=dtype,
	253	+ )
	254	+ causal = torch.triu(
	255	+ torch.full((seq_len, seq_len), float("-inf"), device=device, dtype=dtype),
	256	+ diagonal=1,
	257	+ )
	258	+ mask = mask + causal.view(1, 1, seq_len, seq_len)
	259	+ if attention_mask is not None:
	260	+ if attention_mask.shape != input_ids.shape:
	261	+ raise ValueError("attention_mask must match input_ids shape")
	262	+ padding = attention_mask.to(device=device) == 0
	263	+ mask = mask.masked_fill(padding.view(batch_size, 1, 1, seq_len), float("-inf"))
	264	+ mask = mask.masked_fill(padding.view(batch_size, 1, seq_len, 1), 0.0)
	265	+ return mask
	266	+
	267	+
	268	+def append_attention_mask(
	269	+ attention_mask: Optional[torch.Tensor],
	270	+ next_tokens: torch.Tensor,
	271	+) -> Optional[torch.Tensor]:
	272	+ """Append valid-token mask entries for generated tokens."""
	273	+ if attention_mask is None:
	274	+ return None
	275	+ ones = torch.ones(
	276	+ next_tokens.shape,
	277	+ device=attention_mask.device,
	278	+ dtype=attention_mask.dtype,
	279	+ )
	280	+ return torch.cat([attention_mask, ones], dim=-1)

Atests/torch/generate/__init__.py+15-0

@@ -0,0 +1,15 @@
	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	+"""Generate tests."""

Atests/torch/generate/stub_model.py+73-0

@@ -0,0 +1,73 @@
	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	+"""Stub causal language models for generate tests."""
	16	+import torch
	17	+from torch import nn
	18	+
	19	+
	20	+class CacheLengthLM(nn.Module):
	21	+ """Predicts the current total sequence length modulo vocab size."""
	22	+
	23	+ def __init__(self, vocab_size: int = 32, eos_at_total_len: int \| None = None):
	24	+ super().__init__()
	25	+ self.vocab_size = vocab_size
	26	+ self.eos_at_total_len = eos_at_total_len
	27	+ self.calls = []
	28	+
	29	+ def forward(
	30	+ self,
	31	+ input_ids,
	32	+ position_ids=None,
	33	+ attention_mask=None,
	34	+ past_key_values=None,
	35	+ use_cache=True,
	36	+ **kwargs,
	37	+ ):
	38	+ """Run deterministic forward with an optional tuple KV cache."""
	39	+ del position_ids, attention_mask, kwargs
	40	+ past_len = 0
	41	+ if past_key_values is not None:
	42	+ past_len = past_key_values[0][0].shape[-2]
	43	+ batch_size, seq_len = input_ids.shape
	44	+ total_len = past_len + seq_len
	45	+ token_id = total_len % self.vocab_size
	46	+ if self.eos_at_total_len is not None and total_len >= self.eos_at_total_len:
	47	+ token_id = 2
	48	+ logits = input_ids.new_full(
	49	+ (batch_size, seq_len, self.vocab_size),
	50	+ -1000,
	51	+ dtype=torch.float32,
	52	+ )
	53	+ logits[:, -1, token_id] = 1000
	54	+ self.calls.append({
	55	+ "seq_len": seq_len,
	56	+ "past_len": past_len,
	57	+ "use_cache": use_cache,
	58	+ })
	59	+ past = None
	60	+ if use_cache:
	61	+ key = torch.zeros(batch_size, 1, total_len, 4)
	62	+ value = torch.zeros(batch_size, 1, total_len, 4)
	63	+ past = [(key, value)]
	64	+ return {"logits": logits, "past_key_values": past}
	65	+
	66	+
	67	+class NoCacheLengthLM(CacheLengthLM):
	68	+ """Same logits as CacheLengthLM but never returns past_key_values."""
	69	+
	70	+ def forward(self, args, *kwargs):
	71	+ output = super().forward(args, *kwargs)
	72	+ output["past_key_values"] = None
	73	+ return output

Atests/torch/generate/test_generate.py+725-0

Atests/torch/generate/test_kv_cache.py+446-0

@@ -0,0 +1,446 @@
	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	+"""KV cache tests."""
	阿苏6月22日 kv cache需要测试和没开kv cache推理生成下的输出结果。结果评估两种方式，一种看实际的token输出，检查语义是否相同。一种要对比两个logit的余弦相似度之类的指标。 Moy 6月23日评论： 6月23日评论：
	16	+import pytest
	17	+import torch
	18	+from torch.nn import functional as F
	19	+
	20	+from hyper_parallel.infer import (
	21	+ GenerationConfig,
	22	+ KVCache,
	23	+ generate,
	24	+)
	25	+from hyper_parallel.infer.kv_cache import (
	26	+ ContextParallelKVCache,
	27	+ SequenceShardInfo,
	28	+ get_sequence_shard_info,
	29	+ shard_past_key_values,
	30	+)
	31	+from tests.torch.generate.stub_model import CacheLengthLM, NoCacheLengthLM
	32	+
	33	+
	34	+def _past(seq_len: int, requires_grad: bool = False, batch_size: int = 2):
	35	+ key = torch.ones(batch_size, 3, seq_len, 4, requires_grad=requires_grad)
	36	+ value = torch.ones(batch_size, 3, seq_len, 4, requires_grad=requires_grad)
	37	+ return [(key, value)]
	38	+
	39	+
	40	+class OpaqueCache:
	41	+ """HF-style cache object with its own sequence-length API."""
	42	+
	43	+ def get_seq_length(self):
	44	+ """Return the cached sequence length."""
	45	+ return 1
	46	+
	47	+
	48	+class TinyAttentionCacheLM(torch.nn.Module):
	49	+ """Small causal-attention model with KV cache support."""
	50	+
	51	+ def __init__(self, vocab_size: int = 64):
	52	+ super().__init__()
	53	+ self.vocab_size = vocab_size
	54	+ self.hidden_size = vocab_size
	55	+ self.embedding = torch.nn.Embedding(vocab_size, self.hidden_size)
	56	+ self.lm_head = torch.nn.Linear(self.hidden_size, vocab_size, bias=False)
	57	+ with torch.no_grad():
	58	+ self.embedding.weight.zero_()
	59	+ self.lm_head.weight.zero_()
	60	+ for token_id in range(vocab_size):
	61	+ scale = 20.0 if token_id == 7 else 1.0
	62	+ self.embedding.weight[token_id, token_id] = scale
	63	+ self.lm_head.weight[token_id, token_id] = 1.0
	64	+
	65	+ def forward(
	66	+ self,
	67	+ input_ids,
	68	+ position_ids=None,
	69	+ attention_mask=None,
	70	+ past_key_values=None,
	71	+ use_cache=True,
	72	+ **kwargs,
	73	+ ):
	74	+ """Run causal attention with tuple KV cache."""
	75	+ del position_ids, kwargs
	76	+ hidden = self.embedding(input_ids)
	77	+ key = hidden.unsqueeze(1)
	78	+ value = key
	79	+ if past_key_values is not None:
	80	+ past_key, past_value = past_key_values[0]
	81	+ key = torch.cat([past_key, key], dim=-2)
	82	+ value = torch.cat([past_value, value], dim=-2)
	83	+ query = torch.ones_like(hidden).unsqueeze(1)
	84	+ scores = torch.matmul(query, key.transpose(-1, -2)) / (self.hidden_size ** 0.5)
	85	+ if attention_mask is not None:
	86	+ scores = scores + attention_mask.to(dtype=scores.dtype)
	87	+ probs = torch.softmax(scores, dim=-1)
	88	+ context = torch.matmul(probs, value).squeeze(1)
	89	+ projected = self.lm_head(context)
	90	+ logits = input_ids.new_full(
	91	+ input_ids.shape + (self.vocab_size,),
	92	+ -1000,
	93	+ dtype=torch.float32,
	94	+ )
	95	+ logits.copy_(projected)
	96	+ past = [(key.detach(), value.detach())] if use_cache else None
	97	+ return {"logits": logits, "past_key_values": past}
	98	+
	99	+
	100	+def test_cache_update_detaches_tensors():
	101	+ """
	102	+ Feature: KV cache update
	103	+ Description: Cache update stores detached key/value tensors.
	104	+ Expectation: Cached tensors do not require grad.
	105	+ """
	106	+ cache = KVCache()
	107	+ cache.update(_past(5, requires_grad=True))
	108	+
	109	+ assert not cache.is_empty
	110	+ key, value = cache.past_key_values[0]
	111	+ assert key.shape[-2] == 5
	112	+ assert value.shape[-2] == 5
	113	+ assert key.requires_grad is False
	114	+ assert value.requires_grad is False
	115	+
	116	+
	117	+def test_cache_update_preserves_opaque_cache_object():
	118	+ """
	119	+ Feature: KV cache update
	120	+ Description: Some model backends return cache objects instead of K/V lists.
	121	+ Expectation: Opaque cache objects are preserved for the next model call.
	122	+ """
	123	+ opaque = OpaqueCache()
	124	+ cache = KVCache()
	125	+
	126	+ cache.update(opaque)
	127	+
	128	+ assert cache.past_key_values is opaque
	129	+
	130	+
	131	+def test_cache_update_normalizes_empty_list_to_empty_cache():
	132	+ """
	133	+ Feature: KV cache update
	134	+ Description: Some model backends may return an empty cache list.
	135	+ Expectation: Empty lists are treated as no available cache.
	136	+ """
	137	+ cache = KVCache()
	138	+
	139	+ cache.update([])
	140	+
	141	+ assert cache.is_empty
	142	+ assert cache.past_key_values is None
	143	+
	144	+
	145	+def test_cache_merge_appends_sequence_dimension():
	146	+ """
	147	+ Feature: KV cache merge
	148	+ Description: Incremental K/V tensors append on the sequence dimension.
	149	+ Expectation: Sequence length grows after merge.
	150	+ """
	151	+ cache = KVCache()
	152	+ cache.update(_past(5))
	153	+ cache.merge(_past(1))
	154	+
	155	+ key, value = cache.past_key_values[0]
	156	+ assert key.shape[-2] == 6
	157	+ assert value.shape[-2] == 6
	158	+
	159	+
	160	+def test_cache_clear_and_shape_validation():
	161	+ """
	162	+ Feature: KV cache validation
	163	+ Description: Invalid cache entries should fail fast.
	164	+ Expectation: Clear empties cache; invalid K/V rank raises ValueError.
	165	+ """
	166	+ cache = KVCache()
	167	+ cache.update(_past(2))
	168	+ cache.clear()
	169	+ assert cache.is_empty
	170	+
	171	+ with pytest.raises(ValueError, match="shape"):
	172	+ cache.update([(torch.ones(2, 3), torch.ones(2, 3))])
	173	+
	174	+
	175	+def test_generate_cache_and_no_cache_outputs_match():
	176	+ """
	177	+ Feature: cache and no-cache generation consistency
	178	+ Description: Run the same deterministic prompt with and without KV cache.
	179	+ Expectation: Generated token ids match exactly.
	180	+ """
	181	+ input_ids = torch.tensor([[7, 8, 9]])
	182	+ config = GenerationConfig(max_new_tokens=4, do_sample=False, eos_token_id=None)
	183	+
	184	+ cache_output = generate(CacheLengthLM(vocab_size=32), input_ids, config)
	185	+ no_cache_output = generate(NoCacheLengthLM(vocab_size=32), input_ids, config)
	186	+
	187	+ assert torch.equal(cache_output, no_cache_output)
	188	+
	189	+
	190	+def test_left_padded_attention_cache_decode_matches_no_cache_decode():
	191	+ """
	192	+ Feature: left-padded KV cache decode
	193	+ Description: Cached attention decode should keep padding keys masked after prefill.
	194	+ Expectation: Cache and no-cache outputs match for left-padded batches.
	195	+ """
	196	+ input_ids = torch.tensor([[7, 7, 1], [0, 7, 1]])
	197	+ attention_mask = torch.tensor([[0, 0, 1], [0, 1, 1]])
	198	+ base_config = {
	199	+ "max_new_tokens": 3,
	200	+ "do_sample": False,
	201	+ "eos_token_id": None,
	202	+ "pad_token_id": 0,
	203	+ }
	204	+
	205	+ cache_output = generate(
	206	+ TinyAttentionCacheLM(vocab_size=64),
	207	+ input_ids,
	208	+ GenerationConfig(**base_config, use_cache=True),
	209	+ attention_mask=attention_mask,
	210	+ )
	211	+ no_cache_output = generate(
	212	+ TinyAttentionCacheLM(vocab_size=64),
	213	+ input_ids,
	214	+ GenerationConfig(**base_config, use_cache=False),
	215	+ attention_mask=attention_mask,
	216	+ )
	217	+
	218	+ assert torch.equal(cache_output, no_cache_output)
	219	+
	220	+
	221	+def test_prefix_cache_reuse_matches_full_prompt_generation():
	222	+ """
	223	+ Feature: prefix cache reuse
	224	+ Description: Continue generation from a contiguous prefix KV cache.
	225	+ Expectation: Reused prefix cache matches full prompt generation after the prefix.
	226	+ """
	227	+ prefix_ids = torch.tensor([[7, 2], [3, 7]])
	228	+ prefix_mask = torch.ones_like(prefix_ids)
	229	+ suffix_ids = torch.tensor([[1], [1]])
	230	+ suffix_mask = torch.ones_like(suffix_ids)
	231	+ config_kwargs = {
	232	+ "max_new_tokens": 3,
	233	+ "do_sample": False,
	234	+ "eos_token_id": None,
	235	+ "pad_token_id": 0,
	236	+ }
	237	+ prefix_model = TinyAttentionCacheLM(vocab_size=64)
	238	+ prefix_outputs = prefix_model(
	239	+ input_ids=prefix_ids,
	240	+ attention_mask=torch.zeros(2, 1, 2, 2),
	241	+ use_cache=True,
	242	+ )
	243	+
	244	+ reused = generate(
	245	+ TinyAttentionCacheLM(vocab_size=64),
	246	+ suffix_ids,
	247	+ GenerationConfig(
	248	+ **config_kwargs,
	249	+ prefix_past_key_values=prefix_outputs["past_key_values"],
	250	+ prefix_attention_mask=prefix_mask,
	251	+ ),
	252	+ attention_mask=suffix_mask,
	253	+ )
	254	+ full = generate(
	255	+ TinyAttentionCacheLM(vocab_size=64),
	256	+ torch.cat([prefix_ids, suffix_ids], dim=-1),
	257	+ GenerationConfig(**config_kwargs),
	258	+ attention_mask=torch.cat([prefix_mask, suffix_mask], dim=-1),
	259	+ )
	260	+
	261	+ assert torch.equal(reused, full[:, prefix_ids.shape[-1]:])
	262	+
	263	+
	264	+def test_prefix_cache_reuse_requires_model_cache_output():
	265	+ """
	266	+ Feature: prefix cache reuse validation
	267	+ Description: Prefix reuse cannot fall back to no-cache recompute without prefix ids.
	268	+ Expectation: Missing updated model cache raises a clear ValueError.
	269	+ """
	270	+ with pytest.raises(ValueError, match="prefix_past_key_values"):
	271	+ generate(
	272	+ NoCacheLengthLM(vocab_size=32),
	273	+ torch.tensor([[1]]),
	274	+ GenerationConfig(
	275	+ max_new_tokens=1,
	276	+ eos_token_id=None,
	277	+ prefix_past_key_values=_past(2, batch_size=1),
	278	+ prefix_attention_mask=torch.ones(1, 2),
	279	+ ),
	280	+ )
	281	+
	282	+
	283	+def test_prefix_cache_batch_size_must_match_input_ids():
	284	+ """
	285	+ Feature: prefix cache validation
	286	+ Description: Prefix cache batch size should match the new prompt batch size.
	287	+ Expectation: Mismatch raises a clear ValueError before model forward.
	288	+ """
	289	+ with pytest.raises(ValueError, match="batch size"):
	290	+ generate(
	291	+ TinyAttentionCacheLM(vocab_size=64),
	292	+ torch.tensor([[1]]),
	293	+ GenerationConfig(
	294	+ max_new_tokens=1,
	295	+ eos_token_id=None,
	296	+ prefix_past_key_values=_past(2),
	297	+ prefix_attention_mask=torch.ones(1, 2),
	298	+ ),
	299	+ )
	300	+
	301	+
	302	+def test_cache_and_no_cache_decode_logits_cosine_similarity():
	303	+ """
	304	+ Feature: cache and no-cache decode consistency
	305	+ Description: Compare next-step logits from cached and full-sequence decode.
	306	+ Expectation: Cosine similarity is close to 1.0 for deterministic logits.
	307	+ """
	308	+ input_ids = torch.tensor([[7, 8, 9]])
	309	+ next_token = torch.tensor([[3]])
	310	+
	311	+ cache_model = CacheLengthLM(vocab_size=32)
	312	+ no_cache_model = NoCacheLengthLM(vocab_size=32)
	313	+
	314	+ prefill = cache_model(input_ids=input_ids, use_cache=True)
	315	+ cached_decode = cache_model(
	316	+ input_ids=next_token,
	317	+ past_key_values=prefill["past_key_values"],
	318	+ use_cache=True,
	319	+ )
	320	+ no_cache_decode = no_cache_model(
	321	+ input_ids=torch.cat([input_ids, next_token], dim=-1),
	322	+ use_cache=False,
	323	+ )
	324	+
	325	+ cached_logits = cached_decode["logits"][:, -1, :]
	326	+ no_cache_logits = no_cache_decode["logits"][:, -1, :]
	327	+ similarity = F.cosine_similarity(cached_logits, no_cache_logits, dim=-1)
	328	+
	329	+ assert torch.all(similarity > 0.9999)
	330	+
	331	+
	332	+def test_sequence_shard_info_handles_uneven_lengths():
	333	+ """
	334	+ Feature: CP sequence shard metadata
	335	+ Description: Split a global sequence across CP ranks with uneven remainder.
	336	+ Expectation: Shards are contiguous, non-overlapping, and cover the full sequence.
	337	+ """
	338	+ shards = [get_sequence_shard_info(5, rank, 2) for rank in range(2)]
	339	+
	340	+ assert [(shard.start, shard.end) for shard in shards] == [(0, 3), (3, 5)]
	341	+ assert sum(shard.local_seq_len for shard in shards) == 5
	342	+ with pytest.raises(ValueError, match="rank"):
	343	+ get_sequence_shard_info(5, rank=2, world_size=2)
	344	+
	345	+
	346	+def test_shard_past_key_values_slices_sequence_dimension():
	347	+ """
	348	+ Feature: CP full-cache sharding
	349	+ Description: Convert full prefill K/V cache into this rank's local sequence shard.
	350	+ Expectation: The returned tensors match the contiguous global sequence slice.
	351	+ """
	352	+ key = torch.arange(2 * 3 * 5 * 4).view(2, 3, 5, 4)
	353	+ value = key + 1000
	354	+
	355	+ sharded, shard_info = shard_past_key_values([(key, value)], rank=1, world_size=2)
	356	+
	357	+ assert shard_info == SequenceShardInfo(
	358	+ rank=1,
	359	+ world_size=2,
	360	+ start=3,
	361	+ end=5,
	362	+ global_seq_len=5,
	363	+ )
	364	+ assert torch.equal(sharded[0][0], key[:, :, 3:5, :])
	365	+ assert torch.equal(sharded[0][1], value[:, :, 3:5, :])
	366	+
	367	+
	368	+def test_context_parallel_cache_update_full_and_clear():
	369	+ """
	370	+ Feature: CP KV cache prefill
	371	+ Description: Store only the local shard of a full prefill cache.
	372	+ Expectation: Cache metadata tracks this rank's global sequence range.
	373	+ """
	374	+ cache = ContextParallelKVCache(rank=0, world_size=2)
	375	+ cache.update_full(_past(5))
	376	+
	377	+ key, value = cache.past_key_values[0]
	378	+ assert cache.shard_info.start == 0
	379	+ assert cache.shard_info.end == 3
	380	+ assert key.shape[-2] == 3
	381	+ assert value.shape[-2] == 3
	382	+
	383	+ cache.clear()
	384	+ assert cache.is_empty
	385	+ assert cache.shard_info.global_seq_len == 0
	386	+
	387	+
	388	+def test_context_parallel_cache_merge_local_validates_shard_growth():
	389	+ """
	390	+ Feature: CP KV cache local decode
	391	+ Description: Append local decode K/V only when this rank's sequence shard grows.
	392	+ Expectation: Valid local growth updates metadata; invalid growth raises ValueError.
	393	+ """
	394	+ cache = ContextParallelKVCache(rank=1, world_size=2)
	395	+ cache.update_full(_past(5))
	396	+ cache.merge_local(_past(1), global_seq_len=6)
	397	+
	398	+ key, value = cache.past_key_values[0]
	399	+ assert cache.shard_info.start == 3
	400	+ assert cache.shard_info.end == 6
	401	+ assert key.shape[-2] == 3
	402	+ assert value.shape[-2] == 3
	403	+
	404	+ cache = ContextParallelKVCache(rank=0, world_size=2)
	405	+ cache.update_full(_past(5))
	406	+ before = cache.past_key_values[0][0].clone()
	407	+ with pytest.raises(ValueError, match="growth"):
	408	+ cache.merge_local(_past(1), global_seq_len=6)
	409	+ assert torch.equal(cache.past_key_values[0][0], before)
	410	+
	411	+
	412	+def test_context_parallel_initial_local_merge_requires_global_seq_len():
	413	+ """
	414	+ Feature: CP KV cache local decode
	415	+ Description: Initial local cache merge cannot infer global length in CP mode.
	416	+ Expectation: Missing global_seq_len raises ValueError when world_size > 1.
	417	+ """
	418	+ cache = ContextParallelKVCache(rank=0, world_size=2)
	419	+
	420	+ with pytest.raises(ValueError, match="global_seq_len"):
	421	+ cache.merge_local(_past(3))
	422	+
	423	+
	424	+def test_context_parallel_cache_update_local_validates_metadata():
	425	+ """
	426	+ Feature: CP KV cache local update
	427	+ Description: Accept K/V tensors that are already local to the current CP rank.
	428	+ Expectation: Mismatched shard metadata or local length raises ValueError.
	429	+ """
	430	+ cache = ContextParallelKVCache(rank=1, world_size=2)
	431	+ cache.update_local(
	432	+ _past(2),
	433	+ SequenceShardInfo(rank=1, world_size=2, start=3, end=5, global_seq_len=5),
	434	+ )
	435	+
	436	+ assert cache.past_key_values[0][0].shape[-2] == 2
	437	+ with pytest.raises(ValueError, match="does not match"):
	438	+ cache.update_local(
	439	+ _past(2),
	440	+ SequenceShardInfo(rank=0, world_size=2, start=0, end=3, global_seq_len=5),
	441	+ )
	442	+ with pytest.raises(ValueError, match="sequence length"):
	443	+ cache.update_local(
	444	+ _past(1),
	445	+ SequenceShardInfo(rank=1, world_size=2, start=3, end=5, global_seq_len=5),
	446	+ )