add flash mla with kvcache-cannbot-dsl-AtomGit

已合并

add flash mla with kvcache #112

yangyibin创建于 17 天前从yangyibin/cannbot-dsl： master929 合入到 cann/cannbot-dsl： master

add flash mla with kvcache #112

已合并

yangyibin创建于 17 天前

从yangyibin/cannbot-dsl master929合入到cann/cannbot-dsl master

共 6 个文件变更+5734-0

Afigures/flash_mla_with_kvcache.png+0-0

Asamples/flash_mla_with_kvcache/README.md+201-0

@@ -0,0 +1,201 @@
	1	+# Flash MLA with KV Cache
	2	+
	3	+基于 CANNBot-DSL 实现的 Flash MLA 推理算子，使用 paged KV cache 完成注意力计算：
	4	+
	5	+$$
	6	+O = \operatorname{softmax}(QK^T \cdot \text{softmax\_scale} + M)V
	7	+$$
	8	+
	9	+## 支持范围
	10	+
	11	+\| 特性与约束 \| 说明 \|
	12	+\| :--- \| :--- \|
	13	+\| 数据类型 \| `float16`、`bfloat16`，Q 与 KV cache 类型一致 \|
	14	+\| Query \| `TND`，shape 为 `[T, N, 576]`，`N` 支持 64 或 96 \|
	15	+\| KV cache \| `PA_NZ` 或 `PA_BBND`，page size 固定为 128，KV head 数固定为 1 \|
	16	+\| 输出布局 \| `NTD`，Attention 输出 shape 为 `[N, T, 512]` \|
	17	+\| Mask \| `mask_mode=0`（无 mask）或 `mask_mode=3`（right-down causal） \|
	18	+\| Value head dim \| 512 \|
	19	+
	20	+实现由三个 Python 文件组成：
	21	+
	22	+- `flash_mla_with_kvcache.py`：主算子和 host launch。
	23	+- `flash_mla_with_kvcache_metadata.py`：生成设备侧调度 metadata。
	24	+- `checker.py`：公开接口参数校验。
	25	+
	26	+## 快速开始
	27	+
	28	+在已加载 CANN 环境并将当前 sample 目录加入 `PYTHONPATH` 后，可运行下面的
	29	+`TND + PA_BBND` 示例。它包含一个 batch、2 个 Query token 和 129 个 KV token：
	30	+
	31	+```python
	32	+import math
	33	+
	34	+import torch
	35	+import torch_npu
	36	+
	37	+from flash_mla_with_kvcache import flash_mla_with_kvcache
	38	+from flash_mla_with_kvcache_metadata import flash_mla_with_kvcache_metadata
	39	+
	40	+dtype = torch.bfloat16
	41	+num_heads_q = 64
	42	+num_heads_kv = 1
	43	+query_tokens = 2
	44	+kv_tokens = 129
	45	+page_size = 128
	46	+page_count = 2
	47	+
	48	+q = torch.randn(
	49	+ query_tokens, num_heads_q, 576, dtype=dtype, device="npu"
	50	+)
	51	+k_cache = torch.randn(
	52	+ page_count, page_size, num_heads_kv, 576,
	53	+ dtype=dtype,
	54	+ device="npu",
	55	+)
	56	+block_table = torch.tensor([[0, 1]], dtype=torch.int32, device="npu")
	57	+cache_seqlens = torch.tensor([kv_tokens], dtype=torch.int32, device="npu")
	58	+cu_seqlens_q = torch.tensor([0, query_tokens], dtype=torch.int32, device="npu")
	59	+
	60	+metadata = flash_mla_with_kvcache_metadata(
	61	+ cache_seqlens,
	62	+ num_heads_q,
	63	+ num_heads_kv,
	64	+ cu_seqlens_q=cu_seqlens_q,
	65	+ max_seqlen_q=query_tokens,
	66	+ max_seqlen_kv=kv_tokens,
	67	+ layout_q="TND",
	68	+)
	69	+
	70	+attention_out, softmax_lse = flash_mla_with_kvcache(
	71	+ q,
	72	+ k_cache,
	73	+ block_table=block_table,
	74	+ cache_seqlens=cache_seqlens,
	75	+ cu_seqlens_q=cu_seqlens_q,
	76	+ metadata=metadata,
	77	+ softmax_scale=1.0 / math.sqrt(576),
	78	+ layout_q="TND",
	79	+ layout_kv="PA_BBND",
	80	+ layout_out="NTD",
	81	+ return_softmax_lse=True,
	82	+)
	83	+torch.npu.synchronize()
	84	+```
	85	+
	86	+调用顺序固定为：先生成 metadata，再将其与同一组序列信息传给主算子。调度配置不变时，metadata 可以复用。
	87	+
	88	+metadata 调度和主算子启动均使用当前 stream 的有效 AIC/AIV 核数，支持 stream 控核；主算子的 block 数和 workspace 随有效 AIC 核数缩放。生成与使用 metadata 时须保持相同的设备与核数配置，修改 stream 核数配额后应重新生成 metadata。此路径不将设备侧 metadata 回读到 host，支持 ACLGraph capture。
	89	+
	90	+## Metadata 接口
	91	+
	92	+```python
	93	+flash_mla_with_kvcache_metadata(
	94	+ cache_seqlens,
	95	+ num_heads_q: int,
	96	+ num_heads_kv: int,
	97	+ cu_seqlens_q=None,
	98	+ seqused_q=None,
	99	+ max_seqlen_q=-1,
	100	+ max_seqlen_kv=-1,
	101	+ head_dim_qk=576,
	102	+ head_dim_v=512,
	103	+ mask_mode=0,
	104	+ layout_q="BSND",
	105	+)
	106	+```
	107	+
	108	+\| 参数 \| 说明 \|
	109	+\| :--- \| :--- \|
	110	+\| `cache_seqlens` \| 必选，NPU 上连续的一维 `int32` Tensor，shape 为 `[B]`，表示每个 batch 的有效 KV 长度 \|
	111	+\| `num_heads_q` \| Query head 数；配合当前主算子使用时为 64 或 96 \|
	112	+\| `num_heads_kv` \| KV head 数，当前固定为 1 \|
	113	+\| `cu_seqlens_q` \| `TND` 必选，NPU 上连续的一维 `int32` Tensor，shape 为 `[B+1]` \|
	114	+\| `seqused_q` \| 可选，NPU 上连续的一维 `int32` Tensor，shape 为 `[B]`，表示每个 batch 实际使用的 Query 数 \|
	115	+\| `max_seqlen_q` \| 最大 Query 长度，`-1` 表示未指定；也可传非负整数供调度使用 \|
	116	+\| `max_seqlen_kv` \| 最大 KV 长度，`-1` 表示未指定；也可传非负整数供调度使用 \|
	117	+\| `head_dim_qk` \| Q/K head dim，当前固定为 576 \|
	118	+\| `head_dim_v` \| Value head dim，当前固定为 512 \|
	119	+\| `mask_mode` \| 0 表示无 mask，3 表示 right-down causal \|
	120	+\| `layout_q` \| metadata 支持 `TND`、`BSND`、`BNSD`；当前主算子使用 `TND` \|
	121	+\| 返回值 \| NPU 上的一维 `int32` metadata Tensor，供主算子直接使用 \|
	122	+
	123	+## 主算子接口
	124	+
	125	+```python
	126	+flash_mla_with_kvcache(
	127	+ q,
	128	+ k_cache,
	129	+ block_table=None,
	130	+ cache_seqlens=None,
	131	+ cu_seqlens_q=None,
	132	+ seqused_q=None,
	133	+ attn_mask=None,
	134	+ metadata=None,
	135	+ head_dim_v=512,
	136	+ softmax_scale=1.0,
	137	+ mask_mode=0,
	138	+ max_seqlen_q=-1,
	139	+ max_seqlen_kv=-1,
	140	+ layout_q="BSND",
	141	+ layout_kv="PA_BBND",
	142	+ layout_out="BSND",
	143	+ return_softmax_lse=False,
	144	+)
	145	+```
	146	+
	147	+\| 参数 \| 说明 \|
	148	+\| :--- \| :--- \|
	149	+\| `q` \| 必选，`float16` 或 `bfloat16`，当前为 `TND [T, N, 576]` \|
	150	+\| `k_cache` \| 必选，与 Q 同 dtype；`PA_BBND [blocks, 128, 1, 576]` 或 `PA_NZ [blocks, 1, 36, 128, 16]` \|
	151	+\| `block_table` \| 必选，`int32 [B, max_pages]`，记录各 batch 使用的物理 page 编号 \|
	152	+\| `cache_seqlens` \| 必选，`int32 [B]`，每个 batch 的有效 KV 长度 \|
	153	+\| `cu_seqlens_q` \| `TND` 必选，`int32 [B+1]` Query 前缀和 \|
	154	+\| `seqused_q` \| 可选，`int32 [B]`，每个 batch 实际参与计算的 Query 数 \|
	155	+\| `attn_mask` \| `mask_mode=3` 时必选，shape 为 `[2048, 2048]`、dtype 为 `int8`；`mask_mode=0` 时必须为 `None` \|
	156	+\| `metadata` \| 必选，由 `flash_mla_with_kvcache_metadata()` 生成的一维 `int32` Tensor \|
	157	+\| `head_dim_v` \| 当前固定为 512 \|
	158	+\| `softmax_scale` \| Softmax 缩放系数，默认值和传入 `None` 时均为 `1.0` \|
	159	+\| `mask_mode` \| 0 表示无 mask，3 表示 right-down causal \|
	160	+\| `max_seqlen_q` \| 当前主算子要求为 `-1` 或 `None` \|
	161	+\| `max_seqlen_kv` \| 当前主算子要求为 `-1` 或 `None` \|
	162	+\| `layout_q` \| 当前必须为 `TND` \|
	163	+\| `layout_kv` \| `PA_NZ` 或 `PA_BBND` \|
	164	+\| `layout_out` \| 当前必须为 `NTD` \|
	165	+\| `return_softmax_lse` \| 是否返回 Softmax LSE \|
	166	+
	167	+返回 `(attention_out, softmax_lse)`：
	168	+
	169	+- `attention_out`：shape 为 `[N, T, 512]`，dtype 与 Q 一致。
	170	+- `softmax_lse`：启用时 shape 为 `[N, T]`、dtype 为 `float32`；关闭时返回空的 `float32` Tensor。
	171	+
	172	+注意：虽然函数签名保留了 `BSND` 默认值，当前 checker 开放的主算子组合为 `layout_q="TND"` 和 `layout_out="NTD"`，调用时应显式传入。
	173	+
	174	+## 性能对比
	175	+
	176	+下图为 CANNBot-DSL 与 Flash Attn With Kvcache（CANN built-in ASC）在 8 个代表性 case 上的 kernel device time。
	177	+使用 msprof 采集 `Task Duration`，每个实现执行两组、每组 1 次 warmup 和
	178	+20 次正式采样；丢弃 warmup 后，从 40 个样本中取最小值。首次编译和 host
	179	+开销不计入结果。加速比定义为 `Flash Attn With Kvcache / CANNBot-DSL`，大于 1 表示
	180	+CANNBot-DSL 更快。
	181	+
	182	+![Flash MLA with KV Cache 与 Flash Attn With Kvcache 性能对比](../../figures/flash_mla_with_kvcache.png)
	183	+
	184	+测试环境为 CANN 9.2.0 两种实现使用同一
	185	+设备、相同输入配置和相同重复次数。
	186	+
	187	+## 功能测试
	188	+
	189	+在仓库根目录执行：
	190	+
	191	+```bash
	192	+python -m pytest test/flash_mla_with_kvcache/test_flash_mla_with_kvcache.py -v
	193	+```
	194	+
	195	+ACLGraph 模式：
	196	+
	197	+```bash
	198	+python -m pytest test/flash_mla_with_kvcache/test_flash_mla_with_kvcache.py -v -o mode=aclgraph
	199	+```
	200	+
	201	+测试覆盖两种 KV layout、两种数据类型、普通与 causal attention、ragged/empty batch、尾块和非连续外层 stride。

Asamples/flash_mla_with_kvcache/checker.py+397-0

@@ -0,0 +1,397 @@
	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	+"""Host-side validation for Flash MLA with KV cache."""
	10	+
	11	+from dataclasses import dataclass
	12	+
	13	+import torch
	14	+
	15	+
	16	+HEAD_DIM_QK = 576
	17	+HEAD_DIM_V = 512
	18	+MASK_TEMPLATE_SIZE = 2048
	19	+KV_BLOCK_SIZE = 128
	20	+PA_NZ_D0 = 16
	21	+PA_NZ_D1 = HEAD_DIM_QK // PA_NZ_D0
	22	+MAX_BATCH_SIZE = 65536
	23	+UNSPECIFIED_SEQLEN = -1
	24	+NO_MASK_MODE = 0
	25	+CAUSAL_MASK_MODE = 3
	26	+SUPPORTED_Q_LAYOUTS = ("TND",)
	27	+SUPPORTED_KV_LAYOUTS = ("PA_NZ", "PA_BBND")
	28	+SUPPORTED_HEAD_COUNTS = (64, 96)
	29	+SUPPORTED_MASK_MODES = (NO_MASK_MODE, CAUSAL_MASK_MODE)
	30	+SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
	31	+
	32	+
	33	+@dataclass(frozen=True, slots=True)
	34	+class CheckedInputs:
	35	+ """Normalized values and dimensions needed by the host launcher."""
	36	+
	37	+ layout_q: str
	38	+ layout_kv: str
	39	+ layout_out: str
	40	+ batch: int
	41	+ total_q: int
	42	+ num_heads_q: int
	43	+ logical_rows: int
	44	+ block_size: int
	45	+
	46	+
	47	+@dataclass(frozen=True, slots=True)
	48	+class InputRelations:
	49	+ """Dimensions established by relationships between public inputs."""
	50	+
	51	+ batch: int
	52	+
	53	+
	54	+@dataclass(frozen=True, slots=True)
	55	+class FeatureSelection:
	56	+ """Normalized values derived from the supported feature set."""
	57	+
	58	+ layout_q: str
	59	+ layout_kv: str
	60	+ layout_out: str
	61	+ total_q: int
	62	+ num_heads_q: int
	63	+ block_size: int
	64	+
	65	+
	66	+def _check_dtype(tensor, name, allowed):
	67	+ display_name = name[0].upper() + name[1:]
	68	+ if tensor.dtype not in allowed:
	69	+ if allowed == SUPPORTED_DTYPES:
	70	+ raise TypeError(
	71	+ f"{display_name} must be float16 or bfloat16; got {tensor.dtype}"
	72	+ )
	73	+ raise TypeError(f"{display_name} must be {allowed}; got {tensor.dtype}")
	74	+
	75	+
	76	+def _require_inner_contiguous(tensor, name):
	77	+ """Accept outer strides while keeping each innermost row DMA-compatible."""
	78	+ display_name = name[0].upper() + name[1:]
	79	+ if tensor.dim() > 0 and tensor.stride(-1) != 1:
	80	+ raise ValueError(f"{display_name} must be contiguous in its last dimension")
	81	+
	82	+
	83	+def _first_noncontiguous_dim(tensor):
	84	+ expected_stride = 1
	85	+ for dim in range(tensor.dim() - 1, -1, -1):
	86	+ if tensor.stride(dim) != expected_stride:
	87	+ return dim
	88	+ expected_stride *= tensor.shape[dim]
	89	+ return -1
	90	+
	91	+
	92	+def _normalize_layout(value, allowed, name):
	93	+ display_name = name[0].upper() + name[1:]
	94	+ if value not in allowed:
	95	+ raise ValueError(
	96	+ f"{display_name} must be one of {sorted(allowed)}; got {value!r}"
	97	+ )
	98	+ return value
	99	+
	100	+
	101	+def _check_required_inputs(*, metadata, block_table, cache_seqlens):
	102	+ """Check that unconditional public inputs are present."""
	103	+ if metadata is None:
	104	+ raise ValueError("Metadata is required")
	105	+ if cache_seqlens is None or block_table is None:
	106	+ raise ValueError("Cache_seqlens and block_table are required")
	107	+
	108	+
	109	+def _check_single_arguments(
	110	+ *,
	111	+ q,
	112	+ k_cache,
	113	+ block_table,
	114	+ cache_seqlens,
	115	+ cu_seqlens_q,
	116	+ seqused_q,
	117	+ attn_mask,
	118	+ metadata,
	119	+ head_dim_v,
	120	+ mask_mode,
	121	+ max_seqlen_q,
	122	+ max_seqlen_kv,
	123	+ layout_q,
	124	+ layout_kv,
	125	+ layout_out,
	126	+ return_softmax_lse,
	127	+):
	128	+ """Check constraints that can be decided from one argument alone."""
	129	+ max_seqlen_q = (
	130	+ UNSPECIFIED_SEQLEN if max_seqlen_q is None else max_seqlen_q
	131	+ )
	132	+ max_seqlen_kv = (
	133	+ UNSPECIFIED_SEQLEN if max_seqlen_kv is None else max_seqlen_kv
	134	+ )
	135	+ if type(max_seqlen_q) is not int:
	136	+ raise TypeError("Max_seqlen_q must be an int")
	137	+ if type(max_seqlen_kv) is not int:
	138	+ raise TypeError("Max_seqlen_kv must be an int")
	139	+ if max_seqlen_q != UNSPECIFIED_SEQLEN:
	140	+ raise ValueError(
	141	+ f"Max_seqlen_q must be {UNSPECIFIED_SEQLEN} or omitted"
	142	+ )
	143	+ if max_seqlen_kv != UNSPECIFIED_SEQLEN:
	144	+ raise ValueError(
	145	+ f"Max_seqlen_kv must be {UNSPECIFIED_SEQLEN} or omitted"
	146	+ )
	147	+ if type(head_dim_v) is not int or head_dim_v != HEAD_DIM_V:
	148	+ raise ValueError(f"Head_dim_v must be {HEAD_DIM_V}; got {head_dim_v}")
	149	+ if not isinstance(return_softmax_lse, bool):
	150	+ raise TypeError("Return_softmax_lse must be bool")
	151	+ if type(mask_mode) is not int:
	152	+ raise TypeError("Mask_mode must be an int")
	153	+
	154	+ for value, name in (
	155	+ (layout_q, "layout_q"),
	156	+ (layout_kv, "layout_kv"),
	157	+ (layout_out, "layout_out"),
	158	+ ):
	159	+ if value is not None and not isinstance(value, str):
	160	+ display_name = name[0].upper() + name[1:]
	161	+ raise ValueError(f"{display_name} must be a string; got {value!r}")
	162	+
	163	+ _check_dtype(q, "q", SUPPORTED_DTYPES)
	164	+ _check_dtype(k_cache, "k_cache", SUPPORTED_DTYPES)
	165	+ if block_table.dtype != torch.int32:
	166	+ raise TypeError("Block_table must be int32")
	167	+ if cache_seqlens.dtype != torch.int32:
	168	+ raise TypeError("Cache_seqlens must be int32")
	169	+ if cache_seqlens.dim() != 1:
	170	+ raise ValueError("Cache_seqlens must be one-dimensional")
	171	+ if cu_seqlens_q is not None and cu_seqlens_q.dtype != torch.int32:
	172	+ raise TypeError("Cu_seqlens_q must be int32")
	173	+ if seqused_q is not None and seqused_q.dtype != torch.int32:
	174	+ raise TypeError("Seqused_q must be int32")
	175	+ if attn_mask is not None:
	176	+ mask_shape = (MASK_TEMPLATE_SIZE, MASK_TEMPLATE_SIZE)
	177	+ if tuple(attn_mask.shape) != mask_shape:
	178	+ raise ValueError(f"Attn_mask must be {mask_shape}")
	179	+ if attn_mask.dtype != torch.int8:
	180	+ raise TypeError("Attn_mask must be int8")
	181	+
	182	+ if not isinstance(metadata, torch.Tensor):
	183	+ raise ValueError(
	184	+ "Metadata tile_table must be replaced by one flat int32 tensor"
	185	+ )
	186	+ if metadata.dtype != torch.int32 or metadata.dim() != 1:
	187	+ raise ValueError("Metadata must be one flat int32 tensor")
	188	+ if metadata.numel() <= 0:
	189	+ raise ValueError("Metadata first dimension must be greater than 0")
	190	+
	191	+ for tensor, name in (
	192	+ (q, "q"),
	193	+ (block_table, "block_table"),
	194	+ (cache_seqlens, "cache_seqlens"),
	195	+ (metadata, "metadata"),
	196	+ (cu_seqlens_q, "cu_seqlens_q"),
	197	+ (seqused_q, "seqused_q"),
	198	+ (attn_mask, "attn_mask"),
	199	+ ):
	200	+ if tensor is not None:
	201	+ _require_inner_contiguous(tensor, name)
	202	+
	203	+
	204	+def _check_cross_arguments(
	205	+ *,
	206	+ q,
	207	+ k_cache,
	208	+ block_table,
	209	+ cache_seqlens,
	210	+ cu_seqlens_q,
	211	+ seqused_q,
	212	+ attn_mask,
	213	+ mask_mode,
	214	+ layout_q,
	215	+):
	216	+ """Check relationships between two or more public inputs."""
	217	+ if k_cache.dtype != q.dtype:
	218	+ raise TypeError("Q and k_cache must have the same dtype")
	219	+ if (
	220	+ mask_mode in SUPPORTED_MASK_MODES
	221	+ and (mask_mode == NO_MASK_MODE) != (attn_mask is None)
	222	+ ):
	223	+ raise ValueError(
	224	+ f"Mask_mode = {NO_MASK_MODE} forbids attn_mask; "
	225	+ f"mask_mode = {CAUSAL_MASK_MODE} requires it"
	226	+ )
	227	+
	228	+ batch = cache_seqlens.shape[0]
	229	+ if not 0 < batch < MAX_BATCH_SIZE:
	230	+ raise ValueError(f"Batch size must be in (0, {MAX_BATCH_SIZE})")
	231	+ if block_table.dim() != 2 or block_table.shape[0] != batch:
	232	+ raise ValueError("Block_table must have shape (B, max_pages)")
	233	+ if block_table.shape[1] <= 0:
	234	+ raise ValueError("Block_table max_pages must be positive")
	235	+ if layout_q == "TND":
	236	+ if cu_seqlens_q is None or cu_seqlens_q.shape != (batch + 1,):
	237	+ raise ValueError("TND requires cu_seqlens_q with shape (B+1,)")
	238	+ if seqused_q is not None and seqused_q.shape != (batch,):
	239	+ raise ValueError("Seqused_q must have shape (B,)")
	240	+ return InputRelations(batch=batch)
	241	+
	242	+
	243	+def _check_supported_features(
	244	+ *,
	245	+ q,
	246	+ k_cache,
	247	+ layout_q,
	248	+ layout_kv,
	249	+ layout_out,
	250	+ mask_mode,
	251	+):
	252	+ """Check limits imposed by the currently implemented feature set."""
	253	+ layout_q = _normalize_layout(layout_q, SUPPORTED_Q_LAYOUTS, "layout_q")
	254	+ layout_kv = _normalize_layout(layout_kv, SUPPORTED_KV_LAYOUTS, "layout_kv")
	255	+ layout_out = "BSND" if layout_out is None else _normalize_layout(
	256	+ layout_out, ("NTD",), "layout_out"
	257	+ )
	258	+ if layout_out != "NTD":
	259	+ raise ValueError("Layout_out must be NTD")
	260	+ if mask_mode not in SUPPORTED_MASK_MODES:
	261	+ raise ValueError(
	262	+ f"Mask_mode must be {NO_MASK_MODE} or {CAUSAL_MASK_MODE}; "
	263	+ f"got {mask_mode}"
	264	+ )
	265	+
	266	+ if q.dim() != 3 or q.shape[2] != HEAD_DIM_QK:
	267	+ raise ValueError(f"TND q must have shape (T, N, {HEAD_DIM_QK})")
	268	+ if q.stride(1) != HEAD_DIM_QK:
	269	+ raise ValueError(
	270	+ "TND q supports non-contiguous storage only between tokens; "
	271	+ "heads inside each token must be contiguous"
	272	+ )
	273	+ total_q, num_heads_q = q.shape[0], q.shape[1]
	274	+ if total_q <= 0:
	275	+ raise ValueError("TND q token count must be positive")
	276	+ if num_heads_q not in SUPPORTED_HEAD_COUNTS:
	277	+ supported_heads = " or ".join(map(str, SUPPORTED_HEAD_COUNTS))
	278	+ raise ValueError(f"Q head count must be {supported_heads}")
	279	+
	280	+ noncontiguous_dim = _first_noncontiguous_dim(k_cache)
	281	+ max_noncontiguous_dim = 0 if layout_kv == "PA_BBND" else 1
	282	+ if noncontiguous_dim > max_noncontiguous_dim:
	283	+ allowed = (
	284	+ "dimension 0" if max_noncontiguous_dim == 0 else "dimensions 0 or 1"
	285	+ )
	286	+ raise ValueError(
	287	+ f"{layout_kv} k_cache supports non-contiguous storage only in "
	288	+ f"{allowed}; got dimension {noncontiguous_dim}"
	289	+ )
	290	+
	291	+ if layout_kv == "PA_NZ":
	292	+ if (
	293	+ k_cache.dim() != 5
	294	+ or k_cache.shape[1] != 1
	295	+ or k_cache.shape[2] != PA_NZ_D1
	296	+ or k_cache.shape[4] != PA_NZ_D0
	297	+ ):
	298	+ raise ValueError(
	299	+ "PA_NZ k_cache must be "
	300	+ f"(blocks, 1, {PA_NZ_D1}, block, {PA_NZ_D0})"
	301	+ )
	302	+ block_size = k_cache.shape[3]
	303	+ else:
	304	+ if k_cache.dim() != 4 or k_cache.shape[2:] != (1, HEAD_DIM_QK):
	305	+ raise ValueError(
	306	+ "PA_BBND k_cache must be "
	307	+ f"(blocks, block, 1, {HEAD_DIM_QK})"
	308	+ )
	309	+ block_size = k_cache.shape[1]
	310	+ if k_cache.shape[0] <= 0:
	311	+ raise ValueError("K_cache block count must be positive")
	312	+ if block_size != KV_BLOCK_SIZE:
	313	+ raise ValueError(f"Paged KV block size must be {KV_BLOCK_SIZE}")
	314	+
	315	+ return FeatureSelection(
	316	+ layout_q=layout_q,
	317	+ layout_kv=layout_kv,
	318	+ layout_out=layout_out,
	319	+ total_q=total_q,
	320	+ num_heads_q=num_heads_q,
	321	+ block_size=block_size,
	322	+ )
	323	+
	324	+
	325	+def check_flash_mla_inputs(
	326	+ *,
	327	+ q,
	328	+ k_cache,
	329	+ block_table,
	330	+ cache_seqlens,
	331	+ cu_seqlens_q,
	332	+ seqused_q,
	333	+ attn_mask,
	334	+ metadata,
	335	+ head_dim_v,
	336	+ mask_mode,
	337	+ max_seqlen_q,
	338	+ max_seqlen_kv,
	339	+ layout_q,
	340	+ layout_kv,
	341	+ layout_out,
	342	+ return_softmax_lse,
	343	+):
	344	+ """Validate inputs by category and derive host launch metadata."""
	345	+ _check_required_inputs(
	346	+ metadata=metadata,
	347	+ block_table=block_table,
	348	+ cache_seqlens=cache_seqlens,
	349	+ )
	350	+ _check_single_arguments(
	351	+ q=q,
	352	+ k_cache=k_cache,
	353	+ block_table=block_table,
	354	+ cache_seqlens=cache_seqlens,
	355	+ cu_seqlens_q=cu_seqlens_q,
	356	+ seqused_q=seqused_q,
	357	+ attn_mask=attn_mask,
	358	+ metadata=metadata,
	359	+ head_dim_v=head_dim_v,
	360	+ mask_mode=mask_mode,
	361	+ max_seqlen_q=max_seqlen_q,
	362	+ max_seqlen_kv=max_seqlen_kv,
	363	+ layout_q=layout_q,
	364	+ layout_kv=layout_kv,
	365	+ layout_out=layout_out,
	366	+ return_softmax_lse=return_softmax_lse,
	367	+ )
	368	+ relations = _check_cross_arguments(
	369	+ q=q,
	370	+ k_cache=k_cache,
	371	+ block_table=block_table,
	372	+ cache_seqlens=cache_seqlens,
	373	+ cu_seqlens_q=cu_seqlens_q,
	374	+ seqused_q=seqused_q,
	375	+ attn_mask=attn_mask,
	376	+ mask_mode=mask_mode,
	377	+ layout_q=layout_q,
	378	+ )
	379	+ features = _check_supported_features(
	380	+ q=q,
	381	+ k_cache=k_cache,
	382	+ layout_q=layout_q,
	383	+ layout_kv=layout_kv,
	384	+ layout_out=layout_out,
	385	+ mask_mode=mask_mode,
	386	+ )
	387	+
	388	+ return CheckedInputs(
	389	+ layout_q=features.layout_q,
	390	+ layout_kv=features.layout_kv,
	391	+ layout_out=features.layout_out,
	392	+ batch=relations.batch,
	393	+ total_q=features.total_q,
	394	+ num_heads_q=features.num_heads_q,
	395	+ logical_rows=features.total_q * features.num_heads_q,
	396	+ block_size=features.block_size,
	397	+ )

Asamples/flash_mla_with_kvcache/flash_mla_with_kvcache.py+3282-0

Asamples/flash_mla_with_kvcache/flash_mla_with_kvcache_metadata.py+1327-0

Atest/flash_mla_with_kvcache/test_flash_mla_with_kvcache.py+527-0

@@ -0,0 +1,527 @@
	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	+"""Representative NPU accuracy coverage for Flash MLA with paged KV cache.
	10	+
	11	+This is a compact functional subset of the historical 204-case regression
	12	+suite, restricted to the public mainline interface. Cases deliberately cross
	13	+multiple behaviors instead of repeating geometry: dtype, both supported KV
	14	+layouts, tails, ragged/empty batches, seqused padding, causal masking,
	15	+FlashDecode and metadata.
	16	+Every end-to-end case runs metadata followed by the main operator and compares
	17	+against an independent CPU reference.
	18	+"""
	19	+
	20	+import importlib
	21	+import math
	22	+from pathlib import Path
	23	+import sys
	24	+
	25	+import pytest
	26	+import torch
	27	+
	28	+REPO_ROOT = Path(__file__).resolve().parents[2]
	29	+SAMPLES = REPO_ROOT / "samples"
	30	+OP_DIR = SAMPLES / "flash_mla_with_kvcache"
	31	+sys.path.insert(0, str(OP_DIR))
	32	+
	33	+flash_mla_module = importlib.import_module("flash_mla_with_kvcache")
	34	+metadata_module = importlib.import_module("flash_mla_with_kvcache_metadata")
	35	+
	36	+
	37	+@pytest.fixture(autouse=True)
	38	+def _automatic_buffer_sync(monkeypatch):
	39	+ """Exercise the sample with compiler-generated Buffer synchronization."""
	40	+ monkeypatch.setenv("CANNBOTDSL_AUTO_BUFFER_SYNC", "1")
	41	+
	42	+
	43	+DEFAULT_RTOL = 0.005
	44	+BFLOAT16_RTOL = 0.0078125
	45	+DEFAULT_ATOL = 0.000025
	46	+BFLOAT16_ATOL = BFLOAT16_RTOL
	47	+FAIL_RATIO = 0.005
	48	+MAX_RELATIVE_ERROR = 10.0
	49	+RELATIVE_FLOOR = (1.0 / (1 << 14)) / 0.005
	50	+RELATIVE_EPSILON = 2e-9
	51	+FULL_D_FAILURE_RATIO = 0.8
	52	+LSE_ATOL = 0.1
	53	+
	54	+
	55	+def _torch_tolerances(dtype):
	56	+ if dtype == torch.float16:
	57	+ return DEFAULT_ATOL, DEFAULT_RTOL
	58	+ if dtype == torch.bfloat16:
	59	+ return BFLOAT16_ATOL, BFLOAT16_RTOL
	60	+ raise ValueError(f"unsupported precision dtype: {dtype}")
	61	+
	62	+
	63	+def _compare_stats(actual, golden, atol, rtol):
	64	+ if actual.shape != golden.shape:
	65	+ return {"shape_ok": False, "fail_ratio": 1.0, "max_rel": 0.0,
	66	+ "full_d_failure": False}
	67	+ actual = actual.float()
	68	+ golden = golden.float()
	69	+ difference = (actual - golden).abs()
	70	+ exact = actual == golden
	71	+ absolute_pass = exact \| (difference <= atol)
	72	+ relative_error = difference / (
	73	+ golden.abs().clamp(min=RELATIVE_FLOOR) + RELATIVE_EPSILON
	74	+ )
	75	+ bad = ~absolute_pass & ~(exact \| (relative_error <= rtol))
	76	+ fail_ratio = float(bad.sum()) / difference.numel()
	77	+ needs_relative = ~absolute_pass
	78	+ max_relative = (
	79	+ relative_error[needs_relative].max().item()
	80	+ if needs_relative.any()
	81	+ else 0.0
	82	+ )
	83	+ bad_rows = bad.reshape(-1, actual.shape[-1])
	84	+ row_failure_ratio = bad_rows.float().mean(dim=1)
	85	+ full_d_failure = bool(
	86	+ bad_rows.all(dim=1).any()
	87	+ or (row_failure_ratio > FULL_D_FAILURE_RATIO).any()
	88	+ )
	89	+ return {
	90	+ "shape_ok": True,
	91	+ "fail_ratio": fail_ratio,
	92	+ "max_rel": max_relative,
	93	+ "full_d_failure": full_d_failure,
	94	+ }
	95	+
	96	+
	97	+def _precision_pass(stats):
	98	+ return (
	99	+ stats["shape_ok"]
	100	+ and not stats["full_d_failure"]