SparseFlashMlaGrad

English | 简体中文

Product Support Status

Product	Supported
Ascend 950PR&950DT products	√
Atlas A3 products	√
Atlas A2 products	√
Atlas 200I/500 A2 inference products	×
Atlas inference products	×
Atlas training products	×

Function Description

• Operator function: Computes the backward output of attention in the SparseFlashMla training scenario, and supports Sliding Window Attention, Compressed Attention, and Sparse Compressed Attention.

• Computation formula:

  Stage 1: Selects the input ori_kv and cmp_kv according to different cmp_ratio scenarios

  • When cmp_ratio = 1 (SWA):

  $$selectedKv=orikv$$

  • When cmp_ratio = 4 (SCFA):

  $$selectedKv=concat(oriKv,Gather(cmpkv,topkIndices[i])),0<=i<selectBlockCount$$

  • else (CFA):

  $$selectedKv=concat(oriKv,cmpkv)$$

  Stage 2: Computes P, dP, and dS

  $$P=SimpleSoftmax(Mask(Q@selectedK{v}^T\cdot scale),lse)$$

  $$dP=dO@selectedK{v}^T$$

  $$dS=P\times (dP-SoftmaxGrad(dO,O))$$

  Stage 3: Computes dQ, dKV, and dSinks

  $$dQ=dS@selectedKv\cdot scale$$

  $$dKV=d{S}^T@Q\cdot scale+P^T@dO$$

  $$dSinks=ReduceSum(-P\times dP\times SimpleSoftmax(sinks,lse),dim=-1)$$

Parameter Description

Parameter Name	Input/Output/Attribute	Description	Data Type	Data Format
query	Input	Input Q of the attention structure.	BFLOAT16, FLOAT16	ND
dOut	Input	Gradient of the attention output matrix.	BFLOAT16, FLOAT16	ND
out	Input	Attention output matrix.	BFLOAT16, FLOAT16	ND
lse	Input	Output lse of the forward attention computation. For the computation formula, see the forward document.	FLOAT32	ND
oriKvOptional	Input	Original input K(V) of the attention structure.	BFLOAT16, FLOAT16	ND
cmpKvOptional	Input	K(V) compressed by the Compressor.	BFLOAT16, FLOAT16	ND
oriSparseIndicesOptional	Input	Attention indices with higher weights selected from oriKvOptional in the sparse scenario.	INT32	ND
cmpSparseIndicesOptional	Input	Attention indices with higher weights selected from cmpKvOptional in the sparse scenario.	INT32	ND
cuSeqlensQOptional	Input	Number of valid Query tokens in each batch.	INT32	ND
cuSeqlensOriKvOptional	Input	Number of valid oriKvOptional tokens in each batch.	INT32	ND
cuSeqlensCmpKvOptional	Input	Number of valid cmpKvOptional tokens in each batch.	INT32	ND
sequsedQOptional	Input	Number of query tokens that actually participate in the computation in different batches.	INT32	ND
sequsedOriKvOptional	Input	Number of oriKvOptional tokens that actually participate in the computation in different batches.	INT32	ND
sequsedCmpKvOptional	Input	Number of cmpKvOptional tokens that actually participate in the computation in different batches.	INT32	ND
cmpResidualKvOptional	Input	Remainder of S2 // cmpRatio for each batch.	INT32	ND
oriTopkLengthOptional	Input	Actual selectable topk length of oriKvOptional corresponding to each row of query.	INT32	ND
cmpTopkLengthOptional	Input	Actual selectable topk length of cmpKvOptional corresponding to each row of query.	INT32	ND
sinksOptional	Input	Attention sink tensor.	FLOAT32	ND
metadataOptional	Input	Output result of the AICPU operator with tiling offloaded.	INT32	ND
scaleValue	Attribute	Scaling factor.	FLOAT32	-
cmpRatio	Attribute	Compression ratio of oriKvOptional.	INT64	-
oriMaskMode	Attribute	Mask mode for the computation between query and oriKvOptional.	INT64	-
cmpMaskMode	Attribute	Mask mode for the computation between query and cmpKvOptional.	INT64	-
oriWinLeft	Attribute	Number of past tokens computed by query in the computation between query and oriKvOptional.	INT64	-
oriWinRight	Attribute	Number of future tokens computed by query in the computation between query and oriKvOptional.	INT64	-
layoutQOptional	Attribute	Data layout format of the input query.	STRING	-
layoutKvOptional	Attribute	Data layout format of the inputs ori_kv and cmp_kv.	STRING	-
dQueryOut	Output	Gradient of query.	BFLOAT16, FLOAT16	ND
dOriKvOutOptional	Output	Gradient of oriKvOptional.	BFLOAT16, FLOAT16	ND
dCmpKvOptional	Output	Gradient of cmpKvOptional.	BFLOAT16, FLOAT16	ND
dSinksOutOptional	Output	Gradient of sinksOptional.	FLOAT32	ND
oriSoftmaxL1NormOptional	Output	Softmax result computed from query and oriKvOptional.	FLOAT32	ND
cmpSoftmaxL1NormOptional	Output	Softmax result computed from query and cmpKvOptional.	FLOAT32	ND

• Atlas A2 products, Atlas A3 products: The oriTopkLengthOptional, cmpTopkLengthOptional, and metadataOptional fields are not supported yet.

Constraints

• Only the BSND or TND layout is supported. The constraints on the data shape are as follows:
  • B: Generalized support.
  • S1, S2: Generalized support; S1 and S2 can have different lengths.
  • N1: Supports 1 to 128.
  • N2: Supports only 1.
  • D: Supports only 512.

Calling Example

Calling Method	Sample Code	Description
aclnn API	test_aclnn_sparse_flash_mla_grad	Calls the operator through the aclnnSparseFlashMlaGrad API
	PyTorch API		-	Calls the sparse_flash_mla_grad operator through the sparse_flash_mla_grad API.