[Feature][op][eqbsa-03] expose plugin API and golden test-MindIE-SD-AtomGit

已合并

[Feature][op][eqbsa-03] expose plugin API and golden test #515

lanwangli创建于 24 天前从lanwangli/MindIE-SD_0813： eqbsa-plugin-golden 合入到 Ascend/MindIE-SD： dev

[Feature][op][eqbsa-03] expose plugin API and golden test #515

已合并

lanwangli创建于 24 天前

从lanwangli/MindIE-SD_0813 eqbsa-plugin-golden合入到Ascend/MindIE-SD dev

共 6 个文件变更+591-0

Mcsrc/CMakeLists.txt+1-0

	@@ -55,6 +55,7 @@ add_library(PTAExtensionOPS SHARED
55	./plugin/frequency_regulator.cpp	55	./plugin/frequency_regulator.cpp
56	./plugin/frequency_regulator_register.cpp	56	./plugin/frequency_regulator_register.cpp
57	./plugin/block_sparse_attention.cpp	57	./plugin/block_sparse_attention.cpp
		58	+ ./plugin/eagle_quant_block_sparse_attention.cpp
58	./plugin/quant_flash_attn.cpp	59	./plugin/quant_flash_attn.cpp
59	./plugin/quant_flash_attn_metadata.cpp	60	./plugin/quant_flash_attn_metadata.cpp
60	./plugin/fused_infer_attention_score.cpp	61	./plugin/fused_infer_attention_score.cpp

Acsrc/plugin/eagle_quant_block_sparse_attention.cpp+137-0

@@ -0,0 +1,137 @@
	1	+/**
	2	+ * Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
	3	+ * MindIE is licensed under Mulan PSL v2.
	4	+ * You can use this software according to the terms and conditions of the Mulan PSL v2.
	5	+ * You may obtain a copy of Mulan PSL v2 at:
	6	+ * http://license.coscl.org.cn/MulanPSL2
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
	8	+ * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
	9	+ * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
	10	+ * See the Mulan PSL v2 for more details.
	11	+ */
	12	+
	13	+#include <torch/library.h>
	14	+
	15	+#include "torch_npu/csrc/framework/utils/OpAdapter.h"
	16	+#include "torch_npu/csrc/core/npu/NPUFormat.h"
	17	+#include "pytorch_npu_helper.h"
	18	+#include "eagle_quant_block_sparse_attention.h"
	19	+
	20	+using namespace at;
	21	+
	22	+namespace {
	23	+// V2 kernel supports BF16/FP16/FP8 natively.
	24	+constexpr std::string_view EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_NAME = "aclnnEagleQuantBlockSparseAttention";
	25	+
	26	+constexpr int64_t MASK_TYPE = 0; // no attention mask
	27	+constexpr int64_t PRE_TOKENS = 2147483647; // full context window
	28	+constexpr int64_t NEXT_TOKENS = 2147483647;
	29	+
	30	+inline at::ScalarType ResolveOutputDtype(const at::Tensor &query, const c10::optional<at::Tensor> &query_scale,
	31	+ const c10::optional<at::ScalarType> &output_dtype)
	32	+{
	33	+ if (output_dtype.has_value()) {
	34	+ return output_dtype.value();
	35	+ }
	36	+ // Quant path default: BF16. Non-quant: match query.
	37	+ return query_scale.has_value() ? at::kBFloat16 : query.scalar_type();
	38	+}
	39	+
	40	+// Validate optional *_dtype against tensor storage. INT8 (Char) storage is the
	41	+// quant bitcast path (source/golden: value.view(int8) + value_dtype=fp8); the
	42	+// logical FP8 dtype legitimately differs from int8 storage and its ScalarType
	43	+// enum is not portable (torch_npu.float8_e4m3fn maps to different slots across
	44	+// versions, e.g. prints as "UInt7"), so accept ANY hint on int8 tensors instead
	45	+// of enumerating FP8 enums. *_dtype never reaches the kernel — the kernel derives
	46	+// dtype from the tensor + tilingKey (DT_INT8 / DT_FLOAT8_E4M3FN V share one key).
	47	+// Do NOT rewrite ACL dtype via TensorWrapper here — custom ACL_DTYPE_FLOAT8_*
	48	+// values mismatch opdev and trigger "Key/Value datatype mismatch with query".
	49	+inline void CheckOptionalInputDtype(const char *name, const at::Tensor &tensor,
	50	+ const c10::optional<at::ScalarType> &dtype)
	51	+{
	52	+ if (!dtype.has_value() \|\| dtype.value() == tensor.scalar_type() \|\| tensor.scalar_type() == at::kChar) {
	53	+ return;
	54	+ }
	55	+ TORCH_CHECK(false, "eagle_quant_block_sparse_attention: ", name, "_dtype (", dtype.value(),
	56	+ ") is incompatible with tensor dtype (", tensor.scalar_type(), ")");
	57	+}
	58	+} // namespace
	59	+
	60	+std::tuple<at::Tensor, at::Tensor> eagle_quant_block_sparse_attention_impl_npu(const at::Tensor &query, const at::Tensor &key,
	61	+ const at::Tensor &value, const c10::optional<at::Tensor> &block_sparse_mask, at::IntArrayRef block_shape,
	62	+ std::string q_input_layout, std::string kv_input_layout, int64_t num_key_value_heads, double scale_value,
	63	+ int64_t inner_precise, c10::OptionalIntArrayRef actual_seq_lengths, c10::OptionalIntArrayRef actual_seq_lengths_kv,
	64	+ int64_t softmax_lse_flag, const c10::optional<at::Tensor> &query_scale,
	65	+ const c10::optional<at::Tensor> &key_scale, const c10::optional<at::Tensor> &value_scale,
	66	+ const c10::optional<at::ScalarType> &query_dtype, const c10::optional<at::ScalarType> &key_dtype,
	67	+ const c10::optional<at::ScalarType> &value_dtype, const c10::optional<at::ScalarType> &output_dtype) {
	68	+ TORCH_CHECK(q_input_layout == "TND" \|\| q_input_layout == "BNSD",
	69	+ "eagle_quant_block_sparse_attention: q_input_layout only supports 'TND' and 'BNSD', got ", q_input_layout);
	70	+ TORCH_CHECK(kv_input_layout == "TND" \|\| kv_input_layout == "BNSD",
	71	+ "eagle_quant_block_sparse_attention: kv_input_layout only supports 'TND' and 'BNSD', got ", kv_input_layout);
	72	+ TORCH_CHECK(q_input_layout == kv_input_layout,
	73	+ "eagle_quant_block_sparse_attention: q_input_layout and kv_input_layout must be consistent.");
	74	+ TORCH_CHECK(q_input_layout != "TND" \|\| (actual_seq_lengths.has_value() && actual_seq_lengths_kv.has_value()),
	75	+ "eagle_quant_block_sparse_attention: actual_seq_lengths and actual_seq_lengths_kv are required for TND layout.");
	76	+
	77	+ const char *qLayoutPtr = q_input_layout.c_str();
	78	+ const char *kvLayoutPtr = kv_input_layout.c_str();
	79	+
	80	+ // attenMaskOptional and blockTableOptional must be nullptr.
	81	+ c10::optional<at::Tensor> nulltensor = c10::nullopt;
	82	+
	83	+ /* EXEC_NPU_CMD has ConvertType for c10::optional<at::IntArrayRef> only, not
	84	+ c10::OptionalIntArrayRef. Convert explicitly: nullopt -> nullptr (op tiling
	85	+ skips batch check), has_value() -> AclIntArray*. Do not use .value_or({})
	86	+ — empty array is interpreted as batch=0, conflicting with query batch dim. */
	87	+ c10::optional<at::IntArrayRef> optSeqLen =
	88	+ actual_seq_lengths.has_value() ? c10::optional<at::IntArrayRef>(actual_seq_lengths.value()) : c10::nullopt;
	89	+ c10::optional<at::IntArrayRef> optSeqLenKv = actual_seq_lengths_kv.has_value()
	90	+ ? c10::optional<at::IntArrayRef>(actual_seq_lengths_kv.value())
	91	+ : c10::nullopt;
	92	+
	93	+ // blockSize=0: PagedAttention not supported.
	94	+ constexpr int64_t blockSize = 0;
	95	+
	96	+ auto outOptions = query.options().dtype(ResolveOutputDtype(query, query_scale, output_dtype));
	97	+ at::Tensor attentionOut =
	98	+ at_npu::native::empty_with_format(query.sizes(), outOptions, at_npu::native::get_npu_format(query));
	99	+
	100	+ // TND: [T, N, 1], BNSD: [B, N, S, 1]
	101	+ at::Tensor softmaxLse;
	102	+ if (q_input_layout == "TND") {
	103	+ softmaxLse = at_npu::native::empty_with_format({query.size(0), query.size(1), 1},
	104	+ query.options().dtype(at::kFloat), at_npu::native::get_npu_format(query));
	105	+ } else {
	106	+ softmaxLse = at_npu::native::empty_with_format({query.size(0), query.size(1), query.size(2), 1},
	107	+ query.options().dtype(at::kFloat), at_npu::native::get_npu_format(query));
	108	+ }
	109	+ // Pass nullptr when flag=0 (op skips lse write).
	110	+ c10::optional<at::Tensor> softmaxLseOpt =
	111	+ (softmax_lse_flag != 0) ? c10::optional<at::Tensor>(softmaxLse) : c10::nullopt;
	112	+
	113	+ CheckOptionalInputDtype("query", query, query_dtype);
	114	+ CheckOptionalInputDtype("key", key, key_dtype);
	115	+ CheckOptionalInputDtype("value", value, value_dtype);
	116	+
	117	+ // Pass storage tensors as-is (same as pre-interface-fix path that passed
	118	+ // CheckDataType). query/key/value_dtype are API-compatible validators for
	119	+ // bitcast callers; output_dtype selects attentionOut allocation dtype.
	120	+ // query/key/value_scale inserted after blockTableOptional.
	121	+ // BF16/FP16 path: pass nulltensor (nullptr) for all three scales.
	122	+ // Quant path: pass FLOAT32 scale tensors for Q/K/V (OpDef valueScale is DT_FLOAT).
	123	+ EXEC_NPU_CMD<EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_NAME>(query, key, value, block_sparse_mask,
	124	+ nulltensor, // attenMaskOptional (nullptr)
	125	+ block_shape,
	126	+ optSeqLen, // nullptr when not set
	127	+ optSeqLenKv, // nullptr when not set
	128	+ nulltensor, // blockTableOptional (nullptr)
	129	+ query_scale, // nullptr for BF16/FP16, FLOAT32 for quant
	130	+ key_scale, // nullptr for BF16/FP16, FLOAT32 for quant
	131	+ value_scale, // nullptr for BF16/FP16, FLOAT32 for quant
	132	+ qLayoutPtr, kvLayoutPtr, num_key_value_heads, MASK_TYPE, scale_value, inner_precise, blockSize, PRE_TOKENS,
	133	+ NEXT_TOKENS, softmax_lse_flag, attentionOut,
	134	+ softmaxLseOpt); // nullptr when flag=0
	135	+
	136	+ return std::make_tuple(attentionOut, softmaxLse);
	137	+}

Acsrc/plugin/eagle_quant_block_sparse_attention.h+41-0

@@ -0,0 +1,41 @@
	1	+/**
	2	+ * Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
	3	+ * MindIE is licensed under Mulan PSL v2.
	4	+ * You can use this software according to the terms and conditions of the Mulan PSL v2.
	5	+ * You may obtain a copy of Mulan PSL v2 at:
	6	+ * http://license.coscl.org.cn/MulanPSL2
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
	8	+ * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
	9	+ * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
	10	+ * See the Mulan PSL v2 for more details.
	11	+ */
	12	+
	13	+#ifndef EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_MINDIE_SD_IMPL_H
	14	+#define EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_MINDIE_SD_IMPL_H
	15	+
	16	+#include <ATen/Tensor.h>
	17	+#include <c10/core/ScalarType.h>
	18	+#include <c10/util/Optional.h>
	19	+#include <string>
	20	+#include <tuple>
	21	+
	22	+// Block sparse attention via aclnnEagleQuantBlockSparseAttention (BF16/FP16/FP8).
	23	+// When dequant scales are not provided (BF16/FP16), nullptr is passed to the kernel.
	24	+// When dequant scales are provided (quant path), FLOAT32 scale tensors are used.
	25	+// query/key/value_dtype: optional API validators (INT8 storage + FP8 logical dtype allowed).
	26	+// output_dtype selects attention_out dtype; default bf16 on quant path, else query.dtype.
	27	+// Takes block_sparse_mask (int8/bool as mask, int32 as index). Supports TND and BNSD layouts.
	28	+// Returns (attention_out, softmax_lse).
	29	+std::tuple<at::Tensor, at::Tensor> eagle_quant_block_sparse_attention_impl_npu(const at::Tensor &query, const at::Tensor &key,
	30	+ const at::Tensor &value, const c10::optional<at::Tensor> &block_sparse_mask, at::IntArrayRef block_shape,
	31	+ std::string q_input_layout, std::string kv_input_layout, int64_t num_key_value_heads, double scale_value,
	32	+ int64_t inner_precise, c10::OptionalIntArrayRef actual_seq_lengths, c10::OptionalIntArrayRef actual_seq_lengths_kv,
	33	+ int64_t softmax_lse_flag, const c10::optional<at::Tensor> &query_scale = c10::nullopt,
	34	+ const c10::optional<at::Tensor> &key_scale = c10::nullopt,
	35	+ const c10::optional<at::Tensor> &value_scale = c10::nullopt,
	36	+ const c10::optional<at::ScalarType> &query_dtype = c10::nullopt,
	37	+ const c10::optional<at::ScalarType> &key_dtype = c10::nullopt,
	38	+ const c10::optional<at::ScalarType> &value_dtype = c10::nullopt,
	39	+ const c10::optional<at::ScalarType> &output_dtype = c10::nullopt);
	40	+
	41	+#endif // EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_MINDIE_SD_IMPL_H

Mcsrc/plugin/register_ops.cpp+12-0

	@@ -20,6 +20,7 @@
20	#include "sparse_block_estimate.h"	20	#include "sparse_block_estimate.h"
21	#include "layernorm.h"	21	#include "layernorm.h"
22	#include "block_sparse_attention.h"	22	#include "block_sparse_attention.h"
		23	+#include "eagle_quant_block_sparse_attention.h"
23	#include "quant_flash_attn.h"	24	#include "quant_flash_attn.h"
24	#include "quant_flash_attn_metadata.h"	25	#include "quant_flash_attn_metadata.h"
25	#include "fused_infer_attention_score.h"	26	#include "fused_infer_attention_score.h"
	@@ -69,6 +70,16 @@ TORCH_LIBRARY(mindiesd, m) {
69	int softmax_lse_flag=0, \	70	int softmax_lse_flag=0, \
70	Tensor? q_dequant_scale=None, Tensor? k_dequant_scale=None, \	71	Tensor? q_dequant_scale=None, Tensor? k_dequant_scale=None, \
71	Tensor? v_dequant_scale=None) -> (Tensor, Tensor)");	72	Tensor? v_dequant_scale=None) -> (Tensor, Tensor)");
		73	+ m.def("eagle_quant_block_sparse_attention(Tensor query, Tensor key, Tensor value, \
		74	+ Tensor? block_sparse_mask=None, int[] block_shape=[128,128], \
		75	+ str q_input_layout='BNSD', str kv_input_layout='BNSD', \
		76	+ int num_key_value_heads=1, float scale_value=1.0, int inner_precise=0, \
		77	+ int[]? actual_seq_lengths=None, int[]? actual_seq_lengths_kv=None, \
		78	+ int softmax_lse_flag=0, \
		79	+ Tensor? query_scale=None, Tensor? key_scale=None, \
		80	+ Tensor? value_scale=None, \
		81	+ ScalarType? query_dtype=None, ScalarType? key_dtype=None, \
		82	+ ScalarType? value_dtype=None, ScalarType? output_dtype=None) -> (Tensor, Tensor)");
72	m.def("quant_flash_attn(Tensor query, Tensor key, Tensor value, \	83	m.def("quant_flash_attn(Tensor query, Tensor key, Tensor value, \
73	Tensor q_descale, Tensor k_descale, Tensor v_descale, \	84	Tensor q_descale, Tensor k_descale, Tensor v_descale, \
74	int q_quant_mode, int k_quant_mode, int v_quant_mode, \	85	int q_quant_mode, int k_quant_mode, int v_quant_mode, \
	@@ -129,6 +140,7 @@ TORCH_LIBRARY_IMPL(mindiesd, PrivateUse1, m) {
129	m.impl("sparse_block_estimate", &sparse_block_estimate_mindie_sd_impl_npu);	140	m.impl("sparse_block_estimate", &sparse_block_estimate_mindie_sd_impl_npu);
130	m.impl("layernorm", &layernorm_mindie_sd_impl_npu);	141	m.impl("layernorm", &layernorm_mindie_sd_impl_npu);
131	m.impl("block_sparse_attention", &block_sparse_attention_impl_npu);	142	m.impl("block_sparse_attention", &block_sparse_attention_impl_npu);
		143	+ m.impl("eagle_quant_block_sparse_attention", &eagle_quant_block_sparse_attention_impl_npu);
132	m.impl("quant_flash_attn", &quant_flash_attn_impl_npu);	144	m.impl("quant_flash_attn", &quant_flash_attn_impl_npu);
133	m.impl("quant_flash_attn_metadata", &quant_flash_attn_metadata_impl_npu);	145	m.impl("quant_flash_attn_metadata", &quant_flash_attn_metadata_impl_npu);
134	m.impl("fused_infer_attention_score_v2", &fused_infer_attention_score_v2_impl_npu);	146	m.impl("fused_infer_attention_score_v2", &fused_infer_attention_score_v2_impl_npu);

Mmindiesd/layers/_custom_ops.py+95-0

	@@ -403,6 +403,101 @@ def block_sparse_attention_fake(
403	return attention_out, softmax_lse	403	return attention_out, softmax_lse
404		404
405		405
		406	+def eagle_quant_block_sparse_attention(
		407	+ query: torch.Tensor,
		408	+ key: torch.Tensor,
		409	+ value: torch.Tensor,
		410	+ block_sparse_mask: Optional[torch.Tensor] = None,
		411	+ block_shape: List[int] = None,
		412	+ q_input_layout: str = "BNSD",
		413	+ kv_input_layout: str = "BNSD",
		414	+ num_key_value_heads: int = 1,
		415	+ scale_value: float = 1.0,
		416	+ inner_precise: int = 4,
		417	+ actual_seq_lengths: Optional[List[int]] = None,
		418	+ actual_seq_lengths_kv: Optional[List[int]] = None,
		419	+ softmax_lse_flag: int = 0,
		420	+ query_scale: Optional[torch.Tensor] = None,
		421	+ key_scale: Optional[torch.Tensor] = None,
		422	+ value_scale: Optional[torch.Tensor] = None,
		423	+ query_dtype: Optional[torch.dtype] = None,
		424	+ key_dtype: Optional[torch.dtype] = None,
		425	+ value_dtype: Optional[torch.dtype] = None,
		426	+ output_dtype: Optional[torch.dtype] = None,
		427	+) -> Tuple[torch.Tensor, torch.Tensor]:
		428	+ if block_shape is None:
		429	+ block_shape = [128, 128]
		430	+ kwargs = dict(
		431	+ query=query,
		432	+ key=key,
		433	+ value=value,
		434	+ block_sparse_mask=block_sparse_mask,
		435	+ block_shape=block_shape,
		436	+ q_input_layout=q_input_layout,
		437	+ kv_input_layout=kv_input_layout,
		438	+ num_key_value_heads=num_key_value_heads,
		439	+ scale_value=scale_value,
		440	+ inner_precise=inner_precise,
		441	+ softmax_lse_flag=softmax_lse_flag,
		442	+ )
		443	+ if actual_seq_lengths is not None:
		444	+ kwargs["actual_seq_lengths"] = actual_seq_lengths
		445	+ if actual_seq_lengths_kv is not None:
		446	+ kwargs["actual_seq_lengths_kv"] = actual_seq_lengths_kv
		447	+ if query_scale is not None:
		448	+ kwargs["query_scale"] = query_scale
		449	+ kwargs["key_scale"] = key_scale
		450	+ kwargs["value_scale"] = value_scale
		451	+ if query_dtype is not None:
		452	+ kwargs["query_dtype"] = query_dtype
		453	+ if key_dtype is not None:
		454	+ kwargs["key_dtype"] = key_dtype
		455	+ if value_dtype is not None:
		456	+ kwargs["value_dtype"] = value_dtype
		457	+ if output_dtype is not None:
		458	+ kwargs["output_dtype"] = output_dtype
		459	+ return getattr(torch.ops.mindiesd, "eagle_quant_block_sparse_attention")(**kwargs)
		460	+
		461	+
		462	+@register_ops.register_mindie_fake_op("eagle_quant_block_sparse_attention")
		463	+def eagle_quant_block_sparse_attention_fake(
		464	+ query: torch.Tensor,
		465	+ key: torch.Tensor,
		466	+ value: torch.Tensor,
		467	+ block_sparse_mask: Optional[torch.Tensor] = None,
		468	+ block_shape: List[int] = None,
		469	+ q_input_layout: str = "BNSD",
		470	+ kv_input_layout: str = "BNSD",
		471	+ num_key_value_heads: int = 1,
		472	+ scale_value: float = 1.0,
		473	+ inner_precise: int = 4,
		474	+ actual_seq_lengths: Optional[List[int]] = None,
		475	+ actual_seq_lengths_kv: Optional[List[int]] = None,
		476	+ softmax_lse_flag: int = 0,
		477	+ query_scale: Optional[torch.Tensor] = None,
		478	+ key_scale: Optional[torch.Tensor] = None,
		479	+ value_scale: Optional[torch.Tensor] = None,
		480	+ query_dtype: Optional[torch.dtype] = None,
		481	+ key_dtype: Optional[torch.dtype] = None,
		482	+ value_dtype: Optional[torch.dtype] = None,
		483	+ output_dtype: Optional[torch.dtype] = None,
		484	+) -> Tuple[torch.Tensor, torch.Tensor]:
		485	+ if output_dtype is not None:
		486	+ out_dtype = output_dtype
		487	+ elif query_scale is not None:
		488	+ out_dtype = torch.bfloat16
		489	+ else:
		490	+ out_dtype = query.dtype
		491	+ attention_out = torch.empty(query.shape, device=query.device, dtype=out_dtype)
		492	+ # softmax_lse shape: TND -> [T, N, 1], BNSD -> [B, N, S, 1]
		493	+ if q_input_layout == "TND":
		494	+ lse_shape = [query.shape[0], query.shape[1], 1]
		495	+ else:
		496	+ lse_shape = [query.shape[0], query.shape[1], query.shape[2], 1]
		497	+ softmax_lse = torch.empty(lse_shape, device=query.device, dtype=torch.float32)
		498	+ return attention_out, softmax_lse
		499	+
		500	+
406	def adaln(	501	def adaln(
407	x: torch.Tensor,	502	x: torch.Tensor,
408	scale: torch.Tensor,	503	scale: torch.Tensor,

Atests/ops/eagle_quant_block_sparse_attention/eagle_quant_block_sparse_attention_golden.py+305-0

@@ -0,0 +1,305 @@
	1	+#!/usr/bin/env python
	2	+# coding=utf-8
	3	+# Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
	4	+# MindIE is licensed under Mulan PSL v2.
	5	+# You can use this software according to the terms and conditions of the Mulan PSL v2.
	6	+# You may obtain a copy of Mulan PSL v2 at:
	7	+# http://license.coscl.org.cn/MulanPSL2
	8	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
	9	+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
	10	+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
	11	+# See the Mulan PSL v2 for more details.
	12	+
	13	+# Single-operator precision test for eagle_quant_block_sparse_attention.
	14	+# Ported from the standalone qbsa test.py; keeps wiki/source public kwargs
	15	+# (query_scale / *_dtype / output_dtype) on torch.ops.mindiesd.eagle_quant_block_sparse_attention.
	16	+#
	17	+# NOTE: this operator is an Ascend 950 (A5/arch35) operator. It must be run on a
	18	+# 950 device. On other devices the op kernel is not available.
	19	+#
	20	+# Usage:
	21	+# python eagle_quant_block_sparse_attention_golden.py
	22	+
	23	+import math
	24	+import sys
	25	+
	26	+import numpy as np
	27	+import torch
	28	+import torch.nn.functional as F
	29	+import torch_npu
	30	+
	31	+# Load the MindIE-SD custom op library (registers torch.ops.mindiesd.*).
	32	+from mindiesd.layers.register_ops import _load_mindie_ops_library
	33	+
	34	+_load_mindie_ops_library()
	35	+
	36	+DEVICE_ID = 0
	37	+torch_npu.npu.set_device(int(DEVICE_ID))
	38	+device = "npu:" + str(DEVICE_ID)
	39	+
	40	+
	41	+def check_nan_inf(x, name, max_print=100):
	42	+ """Check tensor / scalar for NaN or Inf; raise if found."""
	43	+ if isinstance(x, float):
	44	+ if math.isnan(x):
	45	+ print(f"\n[NaN FOUND] {name} is NaN", flush=True)
	46	+ raise RuntimeError(f"NaN detected in {name}")
	47	+ if math.isinf(x):
	48	+ print(f"\n[Inf FOUND] {name} is Inf: {x}", flush=True)
	49	+ raise RuntimeError(f"Inf detected in {name}")
	50	+ return x
	51	+
	52	+ if not torch.is_tensor(x):
	53	+ return x
	54	+
	55	+ if not (x.is_floating_point() or x.is_complex()):
	56	+ return x
	57	+
	58	+ x_cpu = x.detach().cpu()
	59	+ nan_mask = torch.isnan(x_cpu)
	60	+ inf_mask = torch.isinf(x_cpu)
	61	+ has_nan = nan_mask.any().item()
	62	+ has_inf = inf_mask.any().item()
	63	+
	64	+ if has_nan or has_inf:
	65	+ print(f"\n[INVALID VALUE FOUND] tensor name: {name}", flush=True)
	66	+ print(f"shape = {tuple(x.shape)}", flush=True)
	67	+ print(f"dtype = {x.dtype}", flush=True)
	68	+ print(f"device = {x.device}", flush=True)
	69	+ if has_nan:
	70	+ nan_idx = nan_mask.nonzero(as_tuple=False)
	71	+ nan_count = nan_idx.shape[0]
	72	+ print(f"\nNaN count = {nan_count}", flush=True)
	73	+ print(f"NaN positions, first {min(max_print, nan_count)}:", flush=True)
	74	+ print(nan_idx[:max_print].tolist(), flush=True)
	75	+ if has_inf:
	76	+ inf_idx = inf_mask.nonzero(as_tuple=False)
	77	+ inf_count = inf_idx.shape[0]
	78	+ print(f"\nInf count = {inf_count}", flush=True)
	79	+ print(f"Inf positions, first {min(max_print, inf_count)}:", flush=True)
	80	+ print(inf_idx[:max_print].tolist(), flush=True)
	81	+ print(f"Inf values, first {min(max_print, inf_count)}:", flush=True)
	82	+ print(x_cpu[inf_mask][:max_print].tolist(), flush=True)
	83	+ raise RuntimeError(f"NaN or Inf detected in {name}")
	84	+
	85	+ return x
	86	+
	87	+
	88	+def block_sparse_attention_cpu(query, key, value, smask, causal=False, blocksize=128):
	89	+ """CPU float reference (non-quant) for block sparse attention."""
	90	+ bs, nq, seq, dim = query.shape
	91	+ nkv = key.shape[1]
	92	+ gqa = nq // nkv
	93	+
	94	+ output = torch.zeros(bs, nq, seq, dim, dtype=torch.float)
	95	+ query = query.float().cpu().numpy()
	96	+ key = key.float().cpu().numpy()
	97	+ value = value.float().cpu().numpy()
	98	+ smask = smask.cpu().numpy()
	99	+
	100	+ for bi in range(bs):
	101	+ for ni in range(nq):
	102	+ num_blocks = (seq + blocksize - 1) // blocksize
	103	+ for s1 in range(num_blocks):
	104	+ mask_block = smask[bi, ni, s1, :num_blocks]
	105	+ mask_seq = np.repeat(mask_block, blocksize)[:seq].astype(bool)
	106	+ start = s1 * blocksize
	107	+ end = min((s1 + 1) * blocksize, seq)
	108	+ q = query[bi, ni, start:end]
	109	+
	110	+ k_head = ni // gqa
	111	+ k = key[bi, k_head][mask_seq]
	112	+ kt = k.T
	113	+
	114	+ p = q @ kt
	115	+ p = p / np.sqrt(dim)
	116	+ if causal:
	117	+ t = end - start
	118	+ cm = np.triu(np.ones((t, t)), k=1) * (-10000.0)
	119	+ p[:, -t:] += cm
	120	+
	121	+ p = p - p.max(axis=-1, keepdims=True)
	122	+ exp_p = np.exp(p)
	123	+ exp_sum = exp_p.sum(axis=-1, keepdims=True)
	124	+ attn = exp_p / (exp_sum + 1e-12)
	125	+ v = value[bi, k_head][mask_seq]
	126	+ out = attn @ v
	127	+ output[bi, ni, start:end] = torch.from_numpy(out)
	128	+ return output
	129	+
	130	+
	131	+def mask_to_indices_4d_for_loop(mask: torch.Tensor) -> torch.Tensor:
	132	+ """Convert a 4D bool mask into a left-packed int32 index tensor (pad -1)."""
	133	+ _, _, _, W = mask.shape
	134	+ result = torch.full_like(mask, -1, dtype=torch.long)
	135	+ mask_flat = mask.view(-1, W)
	136	+ result_flat = result.view(-1, W)
	137	+ for i in range(mask_flat.size(0)):
	138	+ current_row_mask = mask_flat[i]
	139	+ valid_indices = current_row_mask.nonzero(as_tuple=True)[0]
	140	+ num_valid = valid_indices.numel()
	141	+ if num_valid > 0:
	142	+ result_flat[i, :num_valid] = valid_indices
	143	+ return result.to(torch.int32)
	144	+
	145	+
	146	+def ref_compare1(golden: torch.Tensor, actual: torch.Tensor, err=None, print_flag=False):
	147	+ """Single-baseline float compare: \|actual - expected\| <= err * max(1, \|expected\|)."""
	148	+ if err is None:
	149	+ if actual.dtype == torch.float16:
	150	+ err = 2 ** (-10)
	151	+ else:
	152	+ err = 2 ** (-7)
	153	+ golden = golden.to(torch.float32)
	154	+ golden_nmax = torch.clamp(torch.abs(golden), min=1)
	155	+ abs_error = torch.abs(actual.to(torch.float32) - golden)
	156	+ result = (abs_error <= err * golden_nmax).all()
	157	+ EB = torch.mean(abs_error / golden_nmax)
	158	+ if print_flag:
	159	+ print(f"----> EB: {EB.item():.3e} \| max err: {abs_error.max().item():.3e}")
	160	+ return result.item(), EB.item(), abs_error.max().item()
	161	+
	162	+
	163	+@torch.no_grad()
	164	+def perblock_quant(input_tensor, block_size=128, dst_type=torch_npu.float8_e4m3fn, smooth=False, **kwargs):
	165	+ """Per-block quant preprocess for Q/K. Input layout 'BNSD' or 'BSND'."""
	166	+ assert len(input_tensor.shape) == 4, (
	167	+ f"fa block quant preprocess only support qkv quant, dim = 4, but got {len(input_tensor.shape)}."
	168	+ )
	169	+
	170	+ layout = kwargs.get("layout", "BNSD")
	171	+ if layout == "BNSD":
	172	+ b, n, s, d = input_tensor.shape
	173	+ elif layout == "BSND":
	174	+ input_tensor = input_tensor.transpose(1, 2)
	175	+ b, n, s, d = input_tensor.shape
	176	+ else:
	177	+ raise ValueError("unsupport layout")
	178	+
	179	+ if smooth:
	180	+ input_tensor = input_tensor - input_tensor.mean(dim=2, keepdim=True)
	181	+
	182	+ if not s % block_size == 0:
	183	+ padding_length = (block_size - (s % block_size)) % block_size
	184	+ input_tensor = F.pad(input_tensor, (0, 0, 0, padding_length))
	185	+
	186	+ input_tensor = input_tensor.reshape(b, n, math.ceil(s / block_size), -1)
	187	+ input_quant, input_scale = torch_npu.npu_dynamic_quant(input_tensor, dst_type=dst_type)
	188	+
	189	+ if layout == "BNSD":
	190	+ input_quant = input_quant.reshape(b, n, -1, d)[:, :, :s, :]
	191	+ elif layout == "BSND":
	192	+ input_quant = input_quant.transpose(1, 2).reshape(b, -1, n, d)[:, :s, :, :]
	193	+
	194	+ return input_quant, input_scale
	195	+
	196	+
	197	+def test_quant_eagle_block_sparse_attention(
	198	+ b=1, n1=1, s1=1024, d=128, n2=None, s2=None, sparsity=0.5, dtype=torch.bfloat16
	199	+):
	200	+ if not n2:
	201	+ n2 = n1
	202	+ if not s2:
	203	+ s2 = s1
	204	+ causal = False
	205	+ sparse_size = 128
	206	+ sn1 = (s1 + sparse_size - 1) // sparse_size
	207	+ query = torch.randn(b, n1, s1, d, dtype=dtype).npu()
	208	+ key = torch.randn(b, n2, s2, d, dtype=dtype).npu()
	209	+ value = torch.randn(b, n2, s2, d, dtype=dtype).npu()
	210	+
	211	+ smask = torch.rand(b, n1, sn1, sn1) > sparsity
	212	+ smask[:, :, :, 0] = True
	213	+ smask[:, :, :, sn1:] = False
	214	+ smask[:, :, sn1 - 1 : sn1, :] = True
	215	+ smask[:, :, :, sn1 - 1 : sn1] = True
	216	+ smask = smask.npu()
	217	+
	218	+ sn1 = smask.shape[2]
	219	+
	220	+ q_block = 64
	221	+ q_q, q_scales = perblock_quant(query, block_size=q_block, dst_type=torch.int8, smooth=False)
	222	+ k_q, k_scales = perblock_quant(key, block_size=q_block, dst_type=torch.int8, smooth=False)
	223	+ v_q, v_scales = torch_npu.npu_dynamic_quant(value.transpose(-1, -2), dst_type=torch_npu.float8_e4m3fn)
	224	+ v_q = v_q.transpose(-1, -2)
	225	+
	226	+
	227	+ bsa_cpu = block_sparse_attention_cpu(query.cpu(), key.cpu(), value.cpu(), smask.cpu(), causal=causal, blocksize=128)
	228	+ check_nan_inf(bsa_cpu, "CPU-no-quant")
	229	+
	230	+ # ---- run with mask (int8 block sparse mask) ----
	231	+ out, _ = torch.ops.mindiesd.eagle_quant_block_sparse_attention(
	232	+ query=q_q,
	233	+ key=k_q,
	234	+ value=v_q.view(torch.int8),
	235	+ block_sparse_mask=smask.view(torch.int8),
	236	+ block_shape=[128, 128],
	237	+ q_input_layout="BNSD",
	238	+ kv_input_layout="BNSD",
	239	+ num_key_value_heads=n2,
	240	+ scale_value=128 ** -0.5,
	241	+ inner_precise=4,
	242	+ softmax_lse_flag=0,
	243	+ actual_seq_lengths=[s1] * b,
	244	+ actual_seq_lengths_kv=[s1] * b,
	245	+ query_scale=q_scales,
	246	+ key_scale=k_scales,
	247	+ value_scale=v_scales,
	248	+ query_dtype=torch.int8,
	249	+ key_dtype=torch.int8,
	250	+ value_dtype=torch_npu.float8_e4m3fn,
	251	+ output_dtype=torch.bfloat16,
	252	+ )
	253	+ check_nan_inf(out, "npu quant eagle bsa (mask)")
	254	+ print("Compare with CPU-no-quant (mask mode):")
	255	+ _, eb_mask, err_mask = ref_compare1(bsa_cpu.ravel().cpu().float(), out.ravel().cpu().float(), print_flag=True)
	256	+ # ---- run with index (int32 indices derived from mask) ----
	257	+ sindex = mask_to_indices_4d_for_loop(smask)
	258	+ out, _ = torch.ops.mindiesd.eagle_quant_block_sparse_attention(
	259	+ query=q_q,
	260	+ key=k_q,
	261	+ value=v_q.view(torch.int8),
	262	+ block_sparse_mask=sindex,
	263	+ block_shape=[128, 128],
	264	+ q_input_layout="BNSD",
	265	+ kv_input_layout="BNSD",
	266	+ num_key_value_heads=n2,
	267	+ scale_value=128 ** -0.5,
	268	+ inner_precise=4,
	269	+ softmax_lse_flag=0,
	270	+ actual_seq_lengths=[s1] * b,
	271	+ actual_seq_lengths_kv=[s1] * b,
	272	+ query_scale=q_scales,
	273	+ key_scale=k_scales,
	274	+ value_scale=v_scales,
	275	+ query_dtype=torch.int8,
	276	+ key_dtype=torch.int8,
	277	+ value_dtype=torch_npu.float8_e4m3fn,
	278	+ output_dtype=torch.bfloat16,
	279	+ )
	280	+ check_nan_inf(out, "npu quant eagle bsa (index)")
	281	+ print("Compare with CPU-no-quant (index mode):")
	282	+ _, eb_index, err_index = ref_compare1(bsa_cpu.ravel().cpu().float(), out.ravel().cpu().float(), print_flag=True)
	283	+
	284	+ # Pass criterion: mean relative error (EB) within tolerance for both modes
	285	+ # (NaN/Inf already raise above). bf16 + INT8/FP8 quant vs float golden:
	286	+ # EB_TOL = 1e-2 is a comfortable bound for this op.
	287	+ eb_tol = 1e-2
	288	+ passed = (eb_mask <= eb_tol) and (eb_index <= eb_tol)
	289	+ print("\n" + "=" * 60)
	290	+ print(f"[mask mode] EB={eb_mask:.3e} max_err={err_mask:.3e} (EB_TOL={eb_tol:.0e})")
	291	+ print(f"[index mode] EB={eb_index:.3e} max_err={err_index:.3e} (EB_TOL={eb_tol:.0e})")
	292	+ print("=" * 60)
	293	+ if passed:
	294	+ print(">>> eagle_quant_block_sparse_attention TEST PASSED")
	295	+ else:
	296	+ print(">>> eagle_quant_block_sparse_attention TEST FAILED (EB exceeds tolerance)")
	297	+ print("=" * 60)
	298	+ return passed
	299	+
	300	+
	301	+if __name__ == "__main__":
	302	+ np.random.seed(42)
	303	+ torch.manual_seed(42)
	304	+ ok = test_quant_eagle_block_sparse_attention(dtype=torch.bfloat16)
	305	+ sys.exit(0 if ok else 1)