feat: add standalone ReplaySSM verification-cannbot-dsl-AtomGit

已合并

feat: add standalone ReplaySSM verification #190

chuqihang创建于 19 天前从chuqihang/cannbot-dsl-cqh： feat/replayssm 合入到 cann/cannbot-dsl： master

feat: add standalone ReplaySSM verification #190

已合并

chuqihang创建于 19 天前

从chuqihang/cannbot-dsl-cqh feat/replayssm合入到cann/cannbot-dsl master

共 7 个文件变更+1240-0

MREADME.md+1-0

	@@ -191,6 +191,7 @@ python3 -m pytest test/sparse_flash_attention/test_sparse_flash_attention.py -v
191	\| [mixed_quant_sparse_flash_mla](samples/mixed_quant_sparse_flash_mla) \| `mixed_quant_sparse_flash_mla()` \| 混合量化稀疏 MLA：BF16 Q + FP8 原始 KV + 可选 FP4 压缩 KV \|	191	\| [mixed_quant_sparse_flash_mla](samples/mixed_quant_sparse_flash_mla) \| `mixed_quant_sparse_flash_mla()` \| 混合量化稀疏 MLA：BF16 Q + FP8 原始 KV + 可选 FP4 压缩 KV \|
192	\| [mixed_quant_sparse_flash_mla_metadata](samples/mixed_quant_sparse_flash_mla_metadata) \| `mixed_quant_sparse_flash_mla_metadata()` \| MQSMLA 的 AICPU 分核调度，输出整行 / FlashDecode 分片计划 \|	192	\| [mixed_quant_sparse_flash_mla_metadata](samples/mixed_quant_sparse_flash_mla_metadata) \| `mixed_quant_sparse_flash_mla_metadata()` \| MQSMLA 的 AICPU 分核调度，输出整行 / FlashDecode 分片计划 \|
193	\| [flash_kda](samples/flash_kda) \| `flash_kda()`、`flash_kda_metadata()` \| Kimi Delta Attention prefill 融合算子，含 AICPU 调度 metadata \|	193	\| [flash_kda](samples/flash_kda) \| `flash_kda()`、`flash_kda_metadata()` \| Kimi Delta Attention prefill 融合算子，含 AICPU 调度 metadata \|
		194	+\| [fused_recurrent_kda_replayssm](samples/fused_recurrent_kda_replayssm) \| `fused_recurrent_kda_replayssm()` \| 独立 ReplaySSM verify，读取 FP32 checkpoint，写出 FP16 U/K 和 FP32 decay；支持 BSND/BNSD/TND 和 8/24/28/32 控核 \|
194	\| [attn_prologue](samples/attn_prologue) \| `attn_prologue()` \| MXFP8 的 QA/KV 投影、RMSNorm、QR 量化、QB 投影、RoPE 与 KV cache 写回 \|	195	\| [attn_prologue](samples/attn_prologue) \| `attn_prologue()` \| MXFP8 的 QA/KV 投影、RMSNorm、QR 量化、QB 投影、RoPE 与 KV cache 写回 \|
195	\| [attn_epilogue](samples/attn_epilogue) \| `attn_epilogue()` \| Attention 输出投影：inverse RoPE、两次 MXFP8 激活量化与 WoA/WoB 两级矩阵乘；BF16 输入/输出，E4M3FN 权重与 E8M0 scale \|	196	\| [attn_epilogue](samples/attn_epilogue) \| `attn_epilogue()` \| Attention 输出投影：inverse RoPE、两次 MXFP8 激活量化与 WoA/WoB 两级矩阵乘；BF16 输入/输出，E4M3FN 权重与 E8M0 scale \|
196	\| [indexer_prologue_qw](samples/indexer_prologue_qw) \| `indexer_prologue_qw()` \| MXFP8 Q GEMM、尾部 RoPE、MXFP4 量化，以及 BF16 W GEMM \|	197	\| [indexer_prologue_qw](samples/indexer_prologue_qw) \| `indexer_prologue_qw()` \| MXFP8 Q GEMM、尾部 RoPE、MXFP4 量化，以及 BF16 W GEMM \|

Mnet/native_package/operator_groups.toml+2-0

	@@ -20,6 +20,7 @@ engram_gate = ["engram_gate.py"]
20	block_attn_res_update_rms_norm = ["block_attn_res_update_rms_norm.py"]	20	block_attn_res_update_rms_norm = ["block_attn_res_update_rms_norm.py"]
21	flash_kda = ["flash_kda.py", "flash_kda_metadata.py"]	21	flash_kda = ["flash_kda.py", "flash_kda_metadata.py"]
22	fused_recurrent_kda = ["fused_recurrent_kda_snapshot.py"]	22	fused_recurrent_kda = ["fused_recurrent_kda_snapshot.py"]
		23	+fused_recurrent_kda_replayssm = ["fused_recurrent_kda_replayssm.py"]
23	mega_recurrent_kda = ["mega_recurrent_kda.py"]	24	mega_recurrent_kda = ["mega_recurrent_kda.py"]
24	mega_recurrent_kda_replayssm = ["mega_recurrent_kda_replayssm.py"]	25	mega_recurrent_kda_replayssm = ["mega_recurrent_kda_replayssm.py"]
25	commit_recurrent_kda_replayssm = ["commit_recurrent_kda_replayssm.py"]	26	commit_recurrent_kda_replayssm = ["commit_recurrent_kda_replayssm.py"]
	@@ -33,6 +34,7 @@ k3 = [
33	"block_attn_res_update_rms_norm",	34	"block_attn_res_update_rms_norm",
34	"flash_kda",	35	"flash_kda",
35	"fused_recurrent_kda",	36	"fused_recurrent_kda",
		37	+ "fused_recurrent_kda_replayssm",
36	"mega_recurrent_kda",	38	"mega_recurrent_kda",
37	"mega_recurrent_kda_replayssm",	39	"mega_recurrent_kda_replayssm",
38	"commit_recurrent_kda_replayssm",	40	"commit_recurrent_kda_replayssm",

Asamples/fused_recurrent_kda_replayssm/README.md+92-0

@@ -0,0 +1,92 @@
	1	+# fused_recurrent_kda_replayssm
	2	+
	3	+## 产品支持情况
	4	+
	5	+- Ascend 950：支持。
	6	+- 支持 8、24、28、32 AIC 配额；默认选择当前 stream 配额内最大的支持值。
	7	+
	8	+## 功能说明
	9	+
	10	+独立 ReplaySSM verify 算子从只读 FP32 checkpoint 计算候选 token 的 recurrent
	11	+KDA 输出，并写出后续 Commit 使用的 FP16 U/K 和 FP32 decay 记录。输入 Q/K/V
	12	+已经完成投影与卷积；算子内部完成 Q/K 归一化、门控激活和 FP32 状态递推。
	13	+
	14	+每个 token 的计算为：
	15	+
	16	+```text
	17	+q = Q / sqrt(sum(Q²) + 1e-6)
	18	+k = K / sqrt(sum(K²) + 1e-6)
	19	+decay = exp(lower_bound * sigmoid(exp(A_log) * (g + dt_bias)))
	20	+state = state * decay
	21	+u = sigmoid(beta) * (V - state @ k)
	22	+state = state + u ⊗ k
	23	+output = state @ q * scale
	24	+```
	25	+
	26	+递推中的 `state`、`u` 和 `k` 使用 FP32；输出转换为 BF16，保存的 U/K 记录
	27	+转换为 FP16。算子不会把最终状态写回 `recurrent_state`。接受候选 token 后，
	28	+调用 `commit_recurrent_kda_replayssm` 消费这些记录，按 accepted 前缀更新 checkpoint。
	29	+
	30	+## 接口说明
	31	+
	32	+```python
	33	+fused_recurrent_kda_replayssm(
	34	+ mixedqkv, recurrent_state, beta, g, scale, A_log, dt_bias,
	35	+ lower_bound=-5.0, layout_qkv="BSND", *,
	36	+ replay_u, replay_k, replay_decay, cu_seqlens=None, block_num=None,
	37	+)
	38	+```
	39	+
	40	+Q/K/V 头数相同，按 Q、K、V 顺序打包，`D=128`。返回值与 `g` 具有相同 shape，
	41	+数据类型为 BF16。
	42	+
	43	+\| 参数 \| 类型 \| shape / 说明 \|
	44	+\|:---\|:---\|:---\|
	45	+\| `mixedqkv` \| BF16 \| BSND：`[B,S,3N128]`；BNSD：`[B,3N,S,128]`；TND：`[T,3N*128]` \|
	46	+\| `recurrent_state` \| FP32 \| 只读 `[B,N,128,128]`，最后两维依次为 value 和 key \|
	47	+\| `beta` \| BF16 \| BSND：`[B,S,N,1]`；BNSD：`[B,N,S,1]`；TND：`[T,N,1]`，为 sigmoid 前的 logits \|
	48	+\| `g` \| BF16 \| BSND：`[B,S,N,128]`；BNSD：`[B,N,S,128]`；TND：`[T,N,128]` \|
	49	+\| `scale` \| float 或 None \| 有限值；None 使用 `128**-0.5` \|
	50	+\| `A_log` \| FP32 \| `[N]` \|
	51	+\| `dt_bias` \| FP32 \| `[N,128]` \|
	52	+\| `lower_bound` \| float \| 默认 -5.0，范围 `[-5,0]` \|
	53	+\| `layout_qkv` \| str \| BSND、BNSD 或 TND，默认 BSND \|
	54	+\| `replay_u`、`replay_k` \| FP16 \| 原地写入的独立 `[B,C,N,128]` 记录 \|
	55	+\| `replay_decay` \| FP32 \| 原地写入 `[B,C,N,128]`，保存 decay 的指数值 \|
	56	+\| `cu_seqlens` \| INT32 \| TND 必填 `[B+1]` 前缀和；其余布局无需提供 \|
	57	+\| `block_num` \| int 或 None \| 可选 8、24、28、32；显式值不得超过当前 stream 的有效配额 \|
	58	+
	59	+## 约束说明
	60	+
	61	+- `B>0`，`1≤N≤96`，`D=128`，序列长度为正数。
	62	+- 所有 Tensor 连续，并位于同一 NPU。TND prefix sum 从 0 严格递增到 T，
	63	+ 每个序列的长度不得超过记录容量 C；调用方保证 NPU 上的 prefix sum 数值合法。
	64	+- BSND/BNSD 要求 `C≥S`；TND 记录按 batch 从位置 0 保存，未使用的后缀保持原值。
	65	+- verify 计算全部候选 token，不接受 accepted 参数；Commit 单独选择接受前缀。
	66	+- 一个动态 compiled callable 复用布局、shape 与启动配额。
	67	+
	68	+## 调用示例
	69	+
	70	+```python
	71	+output = fused_recurrent_kda_replayssm(
	72	+ mixedqkv, checkpoint, beta, g, None, A_log, dt_bias,
	73	+ layout_qkv="BSND", replay_u=replay_u, replay_k=replay_k,
	74	+ replay_decay=replay_decay,
	75	+)
	76	+commit_recurrent_kda_replayssm(
	77	+ checkpoint, replay_u, replay_k, replay_decay, accepted_tokens,
	78	+)
	79	+```
	80	+
	81	+## 精度测试
	82	+
	83	+```bash
	84	+TORCH_DEVICE_BACKEND_AUTOLOAD=0 OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 \
	85	+python -m pytest -q test/fused_recurrent_kda_replayssm/test_fused_recurrent_kda_replayssm.py
	86	+```
	87	+
	88	+3 个精度 case 覆盖 BSND/BNSD/TND、32/64/128 行分块、跨 wave 调度、
	89	+变长序列、不同 accepted 前缀及记录后缀保留；另有 `B=1,N=1,S=1,D=128`
	90	+的最小全量精度 case。
	91	+独立 CPU FP32 reference 检查输出、U/K/decay 记录和 Commit 后的状态，精度容差
	92	+为 `atol=rtol=0.005`；只读 checkpoint、accepted=0 和未使用的记录后缀使用零容差。

Asamples/fused_recurrent_kda_replayssm/fused_recurrent_kda_replayssm.py+927-0

Mscripts/ci/operator_list.yaml+1-0文件内容审核中，请稍后刷新重试

Atest/fused_recurrent_kda_replayssm/test_fused_recurrent_kda_replayssm.py+209-0

@@ -0,0 +1,209 @@
	1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	2	+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	3	+# CANN Open Software License Agreement Version 2.0 (the "License").
	4	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	5	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	6	+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	7	+# See LICENSE in the root of the software repository for the full text of the License.
	8	+
	9	+"""Independent FP32 oracle for standalone ReplaySSM verify and commit."""
	10	+
	11	+from dataclasses import dataclass
	12	+
	13	+import pytest
	14	+import torch
	15	+from cannbotdsl import get_platform_info
	16	+
	17	+from _samples_path import load_sample
	18	+
	19	+SAMPLE = "fused_recurrent_kda_replayssm/fused_recurrent_kda_replayssm.py"
	20	+COMMIT_SAMPLE = "commit_recurrent_kda_replayssm/commit_recurrent_kda_replayssm.py"
	21	+DIM = 128
	22	+SENTINEL = -7.0
	23	+
	24	+
	25	+@dataclass(frozen=True)
	26	+class _PrecisionInputs:
	27	+ """Dense CPU inputs shared by the oracle and layout packing."""
	28	+
	29	+ query: torch.Tensor
	30	+ key: torch.Tensor
	31	+ value: torch.Tensor
	32	+ gate: torch.Tensor
	33	+ beta: torch.Tensor
	34	+ checkpoint: torch.Tensor
	35	+ a_log: torch.Tensor
	36	+ dt_bias: torch.Tensor
	37	+ lengths: tuple[int, ...]
	38	+
	39	+
	40	+def _make_inputs(heads, lengths):
	41	+ batch, capacity = len(lengths), max(lengths) + 2
	42	+ generator = torch.Generator().manual_seed(20261014 + heads + sum(lengths))
	43	+
	44	+ def random(shape, dtype=torch.bfloat16, scale=0.2):
	45	+ return (torch.randn(shape, generator=generator) * scale).to(dtype)
	46	+
	47	+ shape = (batch, capacity, heads, DIM)
	48	+ return _PrecisionInputs(
	49	+ *(random(shape) for _ in range(4)),
	50	+ random((batch, capacity, heads, 1)),
	51	+ random((batch, heads, DIM, DIM), torch.float32, 0.02),
	52	+ random((heads,), torch.float32),
	53	+ random((heads, DIM), torch.float32),
	54	+ tuple(lengths),
	55	+ )
	56	+
	57	+
	58	+def _reference(inputs):
	59	+ query = inputs.query
	60	+ key = inputs.key
	61	+ value = inputs.value
	62	+ gate = inputs.gate
	63	+ beta = inputs.beta
	64	+ checkpoint = inputs.checkpoint
	65	+ a_log = inputs.a_log
	66	+ dt_bias = inputs.dt_bias
	67	+ lengths = inputs.lengths
	68	+ query = query.float()
	69	+ key = key.float()
	70	+ query = query * torch.rsqrt(query.square().sum(-1, keepdim=True) + 1e-6)
	71	+ key = key * torch.rsqrt(key.square().sum(-1, keepdim=True) + 1e-6)
	72	+ decay = torch.exp(-5.0 * torch.sigmoid(a_log.exp()[None, None, :, None] * (gate.float() + dt_bias)))
	73	+ output = torch.zeros_like(value)
	74	+ records_u = torch.full_like(value, SENTINEL, dtype=torch.float16)
	75	+ records_k = torch.full_like(value, SENTINEL, dtype=torch.float16)
	76	+ records_decay = torch.full_like(value, SENTINEL, dtype=torch.float32)
	77	+ snapshots = []
	78	+ for batch, length in enumerate(lengths):
	79	+ state = checkpoint[batch].clone()
	80	+ states = [state.clone()]
	81	+ for token in range(length):
	82	+ state *= decay[batch, token].unsqueeze(-2)
	83	+ delta = (value[batch, token].float() - (state * key[batch, token].unsqueeze(-2)).sum(-1)) * beta[
	84	+ batch, token
	85	+ ].float().sigmoid()
	86	+ state += delta.unsqueeze(-1) * key[batch, token].unsqueeze(-2)
	87	+ output[batch, token] = ((state * query[batch, token].unsqueeze(-2)).sum(-1) * DIM**-0.5).bfloat16()
	88	+ records_u[batch, token] = delta.half()
	89	+ records_k[batch, token] = key[batch, token].half()
	90	+ records_decay[batch, token] = decay[batch, token]
	91	+ states.append(state.clone())
	92	+ snapshots.append(states)
	93	+ return output, records_u, records_k, records_decay, snapshots
	94	+
	95	+
	96	+def _pack_layout(inputs, layout):
	97	+ query, key, value = inputs.query, inputs.key, inputs.value
	98	+ gate, beta, lengths = inputs.gate, inputs.beta, inputs.lengths
	99	+ # Dense layouts contain the common sequence; records keep a sentinel suffix.
	100	+ if layout != "TND":
	101	+ query, key, value, gate, beta = (tensor[:, : lengths[0]] for tensor in (query, key, value, gate, beta))
	102	+ mixed = torch.cat((query, key, value), dim=2)
	103	+ if layout == "BSND":
	104	+ return mixed.flatten(2).contiguous(), gate.contiguous(), beta.contiguous(), None
	105	+ if layout == "BNSD":
	106	+ return (
	107	+ mixed.transpose(1, 2).contiguous(),
	108	+ gate.transpose(1, 2).contiguous(),
	109	+ beta.transpose(1, 2).contiguous(),
	110	+ None,
	111	+ )
	112	+ packed = [torch.cat([x[b, :length] for b, length in enumerate(lengths)], dim=0) for x in (mixed, gate, beta)]
	113	+ cu = torch.tensor([0, *torch.tensor(lengths).cumsum(0).tolist()], dtype=torch.int32)
	114	+ return packed[0].flatten(1).contiguous(), packed[1].contiguous(), packed[2].contiguous(), cu
	115	+
	116	+
	117	+def _check_output(output, expected, layout, lengths):
	118	+ if layout == "TND":
	119	+ expected = torch.cat([expected[batch, :length] for batch, length in enumerate(lengths)], dim=0)
	120	+ else:
	121	+ expected = expected[:, : lengths[0]]
	122	+ if layout == "BNSD":
	123	+ expected = expected.transpose(1, 2)
	124	+ assert output.dtype == torch.bfloat16
	125	+ torch.testing.assert_close(output.cpu().float(), expected.float(), atol=0.005, rtol=0.005)
	126	+
	127	+
	128	+def _check_records(records, references, lengths):
	129	+ for actual, reference in zip(records, references):
	130	+ torch.testing.assert_close(actual.cpu().float(), reference.float(), atol=0.005, rtol=0.005)
	131	+ for batch, length in enumerate(lengths):
	132	+ torch.testing.assert_close(actual.cpu()[batch, length:], reference[batch, length:], atol=0, rtol=0)
	133	+
	134	+
	135	+def _check_commit(commit, state, records, snapshots, control):
	136	+ # Commit consumes exactly the FP16 U/K and FP32 decay records from verify.
	137	+ accepted, cores = control
	138	+ checkpoint = state.cpu()
	139	+ committed = state.clone()
	140	+ commit.commit_recurrent_kda_replayssm(
	141	+ committed, *records, torch.tensor(accepted, dtype=torch.int32, device="npu"), block_num=cores
	142	+ )
	143	+ torch.npu.synchronize()
	144	+ expected_state = torch.stack([snapshots[batch][count] for batch, count in enumerate(accepted)])
	145	+ torch.testing.assert_close(committed.cpu(), expected_state, atol=0.005, rtol=0.005)
	146	+ for batch, count in enumerate(accepted):
	147	+ if count == 0:
	148	+ torch.testing.assert_close(committed.cpu()[batch], checkpoint[batch], atol=0, rtol=0)
	149	+
	150	+
	151	+def _precision_case(layout, heads, lengths, accepted, cores):
	152	+ pytest.importorskip("torch_npu")
	153	+ if not hasattr(torch, "npu") or not torch.npu.is_available():
	154	+ pytest.skip("requires Ascend NPU")
	155	+ module = load_sample(SAMPLE)
	156	+ commit = load_sample(COMMIT_SAMPLE)
	157	+ inputs = _make_inputs(heads, lengths)
	158	+ mixed, gate, beta, cu = _pack_layout(inputs, layout)
	159	+ mixed = mixed.npu()
	160	+ if cores is not None:
	161	+ stream = torch.npu.current_stream(mixed.device)
	162	+ available_cores = get_platform_info(stream=stream).cube_core_num
	163	+ if cores > available_cores:
	164	+ pytest.skip(f"requires {cores} AIC cores; current stream exposes {available_cores}")
	165	+ expected = _reference(inputs)
	166	+ records = [
	167	+ torch.full(inputs.query.shape, SENTINEL, dtype=dtype, device="npu")
	168	+ for dtype in (torch.float16, torch.float16, torch.float32)
	169	+ ]
	170	+ state = inputs.checkpoint.npu()
	171	+ output = module.fused_recurrent_kda_replayssm(
	172	+ mixed,
	173	+ state,
	174	+ beta.npu(),
	175	+ gate.npu(),
	176	+ None,
	177	+ A_log=inputs.a_log.npu(),
	178	+ dt_bias=inputs.dt_bias.npu(),
	179	+ layout_qkv=layout,
	180	+ replay_u=records[0],
	181	+ replay_k=records[1],
	182	+ replay_decay=records[2],
	183	+ cu_seqlens=None if cu is None else cu.npu(),
	184	+ block_num=cores,
	185	+ )
	186	+ torch.npu.synchronize()
	187	+ torch.testing.assert_close(state.cpu(), inputs.checkpoint, atol=0, rtol=0)
	188	+ _check_output(output, expected[0], layout, lengths)
	189	+ _check_records(records, expected[1:4], lengths)
	190	+ _check_commit(commit, state, records, expected[4], (accepted, cores))
	191	+ torch.testing.assert_close(state.cpu(), inputs.checkpoint, atol=0, rtol=0)
	192	+
	193	+
	194	+@pytest.mark.npu
	195	+@pytest.mark.parametrize(
	196	+ "layout,heads,lengths,accepted,cores",
	197	+ [
	198	+ pytest.param("BSND", 1, (8,), (8,), 24, id="bsnd-small-row-24-cores"),
	199	+ pytest.param("BNSD", 17, (9, 9), (0, 9), 8, id="bnsd-full-row-multiwave"),
	200	+ pytest.param("TND", 8, (5, 1, 3), (5, 0, 2), None, id="tnd-medium-row-variable-length"),
	201	+ ],
	202	+)
	203	+def test_precision_layouts_and_commit(layout, heads, lengths, accepted, cores):
	204	+ _precision_case(layout, heads, lengths, accepted, cores)
	205	+
	206	+
	207	+@pytest.mark.npu
	208	+def test_minimal_full_operator():
	209	+ _precision_case("BSND", 1, (1,), (1,), None)

Mtest/test_config.yaml+8-0文件内容审核中，请稍后刷新重试