sparse_flash_mla_grad

Applicable Products

• Ascend 950PR/Ascend 950DT: Not supported
• Atlas A3 training products/Atlas A3 inference products: Supported
• Atlas A2 training products/Atlas A2 inference products: Supported
• Atlas 200I/500 A2 inference products: Not supported
• Atlas inference products: Not supported
• Atlas training products: Not supported

Function

• API Function:

  • sparse_flash_mla_grad: Computes the backward output of attention in the SparseFlashMla training use case. Supports Sliding Window Attention, Compressed Attention, and Sparse Compressed Attention.
  • sparse_flash_mla_grad_metadata: Generates metadata before the sparse_flash_mla_grad main operator. This API is reserved for future extension and currently has no actual effect. It outputs None.

• Formula:

  Phase 1: Select between the input ori_kv and cmp_kv based on different cmp_ratio use cases.

  • When cmp_ratio = 1 (SWA):

    $$selectedKv\text{ }=\text{ }orikv$$

  • When cmp_ratio = 4 (SCFA):

    $$selectedKv\text{ }=concat(oriKv, \text{ }Gather \left( cmpkv,topkIndices \left[ i \left] \left)) ,\text{ }0\text{ } < =i < \text{ }selectBlockCount\right. \right. \right. \right.$$

  • else (CFA):

    $$selectedKv\text{ }=concat(oriKv, \text{ }cmpkv)$$

    Phase 2: Calculate P, dP, and dS.

    $$P = SimpleSoftmax(Mask(Q \text{ }@\text{ } selectedKv^{{T}} \cdot \text{ } scale), lse)$$ $$dP = dO \text{ }@\text{ } selectedKv^{{T}}$$ $$dS = P \times (dP\text{ } -\text{ } SoftmaxGrad(dO, O))$$

    Phase 3: Calculate dQ, dKV, and dSinks.

    $$dQ = dS \text{ } @ \text{ } selectedKv \text{ } \cdot \text{ } scale$$ $$dKV = dS^{{T}} \text{ } @ \text{ } Q \text{ } \cdot \text{ } scale + P^{{T}} @ \text{ } dO$$ $$dSinks = ReduceSum(-P \text{ }\times\text{ } dP \text{ }\times\text{ } SimpleSoftmax(sinks, lse), dim=-1)$$

Prototype

Before calling the sparse_flash_mla_grad API, call the sparse_flash_mla_grad_metadata API to complete load balancing calculation.

sparse_flash_mla_grad_metadata(num_heads_q, num_heads_kv, head_dim,
                            cu_seqlens_q=None, cu_seqlens_ori_kv=None, cu_seqlens_cmp_kv=None,
                            seqused_q=None, seqused_ori_kv=None, seqused_cmp_kv=None,
                            cmp_residual_kv=None, ori_topk_length=None, cmp_topk_length=None,
                            batch_size=None, max_seqlen_q=None, max_seqlen_ori_kv=None,
                            max_seqlen_cmp_kv=None, ori_topk=None, cmp_topk=None, 
                            cmp_ratio=None, ori_mask_mode=0, cmp_mask_mode=0,
                            ori_win_left=-1, ori_win_right=-1,
                            layout_q="BSND", layout_kv="BSND", has_ori_kv=True, has_cmp_kv=True) -> Tensor

sparse_flash_mla_grad(q, dout, attn_out, softmax_lse, ori_kv=None, cmp_kv=None,
                        ori_sparse_indices=None, cmp_sparse_indices=None, cu_seqlens_q=None,
                        cu_seqlens_ori_kv=None, cu_seqlens_cmp_kv=None, seqused_q=None,
                        seqused_ori_kv=None, seqused_cmp_kv=None, cmp_residual_kv=None, 
                        ori_topk_length=None, cmp_topk_length=None, 
                        sinks=None, metadata=None,
                        softmax_scale=None, cmp_ratio=None, ori_mask_mode=0, cmp_mask_mode=0,
                        ori_win_left=-1, ori_win_right=-1,
                        layout_q="BSND", layout_kv="BSND") -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)

Parameters

sparse_flash_mla_grad_metadata

sparse_flash_mla_grad

Return Value

sparse_flash_mla_grad_metadata

sparse_flash_mla_grad

Constraints

• This API is supported only in training use cases.
• This API supports only single-operator mode calls.
• The data types of q, dout, attn_out, ori_kv, and cmp_kv must be consistent.
• The usage of cmp_kv and cmp_sparse_indices in each use case is as follows:
  • SWA use case: cmp_kv == None && cmp_sparse_indices == None is required.
  • SCFA use case: cmp_kv != None && cmp_sparse_indices != None is required.
  • CFA use case: cmp_kv != None && cmp_sparse_indices == None is required.

Deterministic Computation

• Deterministic computation is not supported by default.

Example

• Single-operator call

  import math
  import random
  import numpy as np
  
  import torch
  import torch_npu
  import cann_ops_transformer
  
  S1 = 256
  S2 = 1024
  cmp_ratio = 4
  actual_seq_q = [S1]
  actual_seq_ori_kv = [S2]
  actual_seq_cmp_kv = [S2 // cmp_ratio]
  cmp_residual = [i % cmp_ratio for i in actual_seq_ori_kv]
  
  T1 = sum(actual_seq_q)
  T2 = sum(actual_seq_ori_kv)
  T3 = sum(actual_seq_cmp_kv)
  
  B = 1
  N1 = 64
  N2 = 1
  D = 512
  K = 512
  scaleValue = 1 / math.sqrt(D)
  dtype = torch.float16
  input_layout = "TND"
  
  q_shape = tuple((T1, N1, D))
  ori_kv_shape = tuple((T2, N2, D))
  cmp_kv_shape = tuple((T3, N2, D))
  out_shape = tuple((T1, N1, D))
  lse_shape = tuple((N2, T1, N1 // N2))
  
  cu_seq_qlen = [0] + [sum(actual_seq_q[:x+1]) for x in range(len(actual_seq_q))]
  cu_seq_ori_kvlen = [0] + [sum(actual_seq_ori_kv[:x+1]) for x in range(len(actual_seq_ori_kv))]
  cu_seq_cmp_kvlen = [0] + [sum(actual_seq_cmp_kv[:x+1]) for x in range(len(actual_seq_cmp_kv))]
  
  q = (torch.rand(q_shape).to(dtype)) * 2
  ori_kv = (torch.rand(ori_kv_shape).to(dtype)) * 2
  cmp_kv = (torch.rand(cmp_kv_shape).to(dtype)) * 2 if cmp_ratio != 1 else None
  dy = (torch.rand(out_shape).to(dtype)) * 2 
  out = (torch.rand(out_shape).to(dtype)) * 2 # Use the forward output attn_out in actual use cases.
  sinks = (torch.rand(N1).to(torch.float32)) * (128)
  lse = (torch.rand(lse_shape).to(torch.float32)) # Use the forward output lse in actual use cases.
  
  if cmp_ratio == 4:
      cmp_sparse_indices = torch.ones([T1, N2, K], dtype=torch.int32) * -1
      accum_t = 0
      for i in range(B):
          cur_s1 = actual_seq_q[i]
          cur_ori_s2 = actual_seq_cmp_kv[i] * cmp_ratio + cmp_residual[i]
          delta_s = cur_ori_s2 - cur_s1
          for s1_idx in range(cur_s1):
              threshold = delta_s + s1_idx + 1 if delta_s + s1_idx + 1 >= 0 else 0
              max_cmp_s2 = threshold // cmp_ratio
              cur_s2_idxs = np.arange(max_cmp_s2)
              if max_cmp_s2 < K:
                  cmp_sparse_indices[accum_t, :, :max_cmp_s2] = torch.tensor(np.random.choice(cur_s2_idxs, size=max_cmp_s2, replace=False))
              else:
                  cmp_sparse_indices[accum_t, :, :] = torch.tensor(np.random.choice(cur_s2_idxs, size=K, replace=False))
              accum_t += 1
  else:
      cmp_sparse_indices = None
  
  cmp_kv_npu = cmp_kv.npu() if cmp_kv != None else None
  cmp_sparse_indices_npu = cmp_sparse_indices.npu() if cmp_sparse_indices != None else None
  cu_seq_qlen = None if input_layout != "TND" else torch.tensor(cu_seq_qlen).to(torch.int32).npu()
  cu_seq_ori_kvlen = None if input_layout != "TND" else torch.tensor(cu_seq_ori_kvlen).to(torch.int32).npu()
  cu_seq_cmp_kvlen = None if input_layout != "TND" else torch.tensor(cu_seq_cmp_kvlen).to(torch.int32).npu()
  cmp_residual_kv = torch.tensor(cmp_residual).to(torch.int32).npu() if cmp_residual != None else None
  
  metadata = cann_ops_transformer.ops.sparse_flash_mla_grad_metadata(N1, N2, D)
  npu_out = cann_ops_transformer.ops.sparse_flash_mla_grad(
          q.npu(), dy.npu(), out.npu(), lse.npu(),
          ori_kv=ori_kv.npu(), cmp_kv=cmp_kv_npu, 
          cmp_sparse_indices=cmp_sparse_indices_npu,
          cu_seqlens_q=cu_seq_qlen, 
          cu_seqlens_ori_kv=cu_seq_ori_kvlen, 
          cu_seqlens_cmp_kv=cu_seq_cmp_kvlen,
          cmp_residual_kv=cmp_residual_kv,
          sinks=sinks.npu(),
          softmax_scale=scaleValue,
          cmp_ratio=cmp_ratio,
          layout_q=input_layout,
          layout_kv=input_layout,
          ori_win_left=127,
          ori_win_right=0,
          ori_mask_mode=4,
          cmp_mask_mode=3,
          metadata=metadata
      )