lightning_indexer

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

• Description:

  The lightning_indexer_metadata API generates a task list that contains the Batch, Head, and the block indices of Q and K at the start and end points of the Attention computation task for each AI core, for use by the subsequent lightning_indexer operator.

  The lightning_indexer API obtains the top-k positions corresponding to each token through a series of operations. The main computation process is as follows:

  1. Multiply the input parameter q ($Q_{index}\in\R^{g\times d}$) corresponding to a token by the given context k ($K_{index}\in\R^{S_{k}\times d}$) to obtain the correlation.
  2. After filtering out invalid negative correlation signals through the activation function $ReLU$, obtain the correlation score vector between the current token and all preceding tokens.
  3. Multiply it by the weight coefficient w ($W$), then along the direction of g, select the top $Top-k$ indices to obtain the output $sparseIndices$, and output the corresponding $sparseValues$ as the input to Attention.

• Formula:

  $$Top-k \left\{ \left[ 1 \left] \mathop{{}}\nolimits_{{1 \times \text{ }g}}\text{@} \left[ \left( W\text{@} \left[ 1 \left] \mathop{{}}\nolimits_{{1\text{ } \times \text{ }S\mathop{{}}\nolimits_{{k}}}} \left) \text{ } \odot \text{ }ReLU \left( Q\mathop{{}}\nolimits_{{index}}\text{@}K\mathop{{}}\nolimits_{{T}}^{{index}} \left) \left] \right\} \right. \right. \right. \right. \right. \right. \right. \right. \right. \right.$$

Prototype

Before calling the lightning_indexer API, call the preceding API lightning_indexer_metadata to complete the load balancing computation for lightning_indexer.

lightning_indexer_metadata(num_heads_q, num_heads_k, head_dim, topk, *, cu_seqlens_q=None, cu_seqlens_k=None, seqused_q=None, seqused_k=None, cmp_residual_k=None, batch_size=0, max_seqlen_q=-1, max_seqlen_k=-1, layout_q="BSND", layout_k="BSND", mask_mode=0, cmp_ratio=1) -> Tensor

lightning_indexer(q, k, w, topk, *, cu_seqlens_q=None, cu_seqlens_k=None, seqused_q=None, seqused_k=None, cmp_residual_k=None, block_table=None, output_idx_offset=None, metadata=None, max_seqlen_q=-1, layout_q="BSND", layout_k="BSND", mask_mode=0, cmp_ratio=1, return_value=0) -> (Tensor, Tensor)

Parameters

Note:

The dimension meanings of the q, k, and w parameters are as follows: B (Batch Size) indicates the batch size of input samples, S1 indicates the input sample sequence length of q, S2 indicates the input sample sequence length of k, N1 indicates the number of heads of q, N2 indicates the number of heads of k, D (Head Dim) indicates the dimension of the attention head, T1 indicates the cumulative sum of the input sample sequence length of q, and T2 indicates the cumulative sum of the input sample sequence length of k. The D value in parameter q is equal to the D value in parameter k, and currently only 128 is supported.

lightning_indexer_metadata

lightning_indexer

• Atlas A3 training products/Atlas A3 inference products, Atlas A2 training products/Atlas A2 inference products:
  • The supported topk value range is [1, 2048], plus 3072, 4096, 5120, 6144, 7168, and 8192.
  • The sequsedQOptional, outputIdxOffsetOptional, and maxSeqlenQ features are not supported. Do not pass these parameters.
  • When layout_k is PA_BBND, sequsedKOptional must be passed. When layout_k is not PA_BBND, the sequsedKOptional feature is not supported. Do not pass this parameter.

Return Value

lightning_indexer_metadata

lightning_indexer

Constraints

• This API supports use in training and inference use cases.
• This API supports calls in single-operator mode and TorchAir graph mode (aclgraph).
• The lightning_indexer_metadata API must be used together with the lightning_indexer operator.
• B (Batch) indicates the batch size of the input samples.
• The values of cu_seqlens_q and cu_seqlens_k must be the cumulative number of valid tokens of the current batch and preceding batches. The value of each subsequent element must be greater than or equal to that of the preceding element.
• The values of seqused_q and seqused_k must indicate the number of valid tokens in each batch.
• cmp_residual_k must satisfy cmp_residual_k[i] < cmp_ratio.
• For details about the mask mode indicated by mask_mode, see sparse_mode parameter description.
• The N of parameter q supports 1 to 64, the N of k supports 1, and headdim supports 128.
• pa_kv_cache supports non-contiguity on axis 0. pa_block_size supports 1 to 1024 and meets the 32B block size alignment.
• The data types of parameters q and k must be consistent.
• Fill the invalid part of sparse_indices with -1 and the invalid part of sparse_values with -inf.
• Atlas A3 training products/Atlas A3 inference products, Atlas A2 training products/Atlas A2 inference products:
  • The topk value range currently supports only [1, 2048], as well as 3072, 4096, 5120, 6144, 7168, and 8192.
  • The sequsedQOptional, outputIdxOffsetOptional, and maxSeqlenQ functions are not supported. Passing these parameters is not recommended.
  • When layout_k is PA_BBND, sequsedKOptional must be passed. When layout_k is not PA_BBND, the sequsedKOptional function is not supported, and passing this parameter is not recommended.

Deterministic Computation

• Defaults to a deterministic computation.

Example

• Single-operator mode call:

  import torch
  import torch_npu
  import numpy as np
  from cann_ops_transformer.ops import lightning_indexer, lightning_indexer_metadata
  
  B = 2
  S1 = 64
  S2 = 130
  N1 = 16
  N2 = 1
  D = 128
  topk = 32
  cmp_ratio = 4
  
  # Calculate the compressed S2 length and remainder.
  S2_compressed = S2 // cmp_ratio
  res_k = S2 % cmp_ratio
  
  # Construct the input.
  q = torch.randn(B, S1, N1, D, dtype=torch.float16).npu()
  k = torch.randn(B, S2_compressed, N2, D, dtype=torch.float16).npu()
  w = torch.randn(B, S1, N1, dtype=torch.float32).npu()
  cmp_residual_k = torch.tensor([res_k] * B, dtype=torch.int32).npu()
  
  # Generate metadata.
  metadata = lightning_indexer_metadata(
      num_heads_q=N1,
      num_heads_k=N2,
      head_dim=D,
      topk=topk,
      cmp_residual_k=cmp_residual_k,
      batch_size=B,
      max_seqlen_q=S1,
      max_seqlen_k=S2_compressed,
      layout_q="BSND",
      layout_k="BSND",
      mask_mode=3,
      cmp_ratio=cmp_ratio
  )
  
  # Execute lightning_indexer.
  sparse_indices, sparse_values = lightning_indexer(
      q, k, w, topk,
      cmp_residual_k=cmp_residual_k,
      metadata=metadata,
      max_seqlen_q=-1,
      layout_q="BSND",
      layout_k="BSND",
      mask_mode=3,
      cmp_ratio=cmp_ratio,
      return_value=0
  )
  print(f"sparse_indices shape: {sparse_indices.shape}")

• aclgraph mode call:

  import torch
  import torch_npu
  import numpy as np
  import torchair
  from cann_ops_transformer.ops import lightning_indexer, lightning_indexer_metadata
  
  B = 2
  S1 = 64
  S2 = 130
  N1 = 16
  N2 = 1
  D = 128
  topk = 32
  cmp_ratio = 4
  
  S2_compressed = S2 // cmp_ratio
  res_k = S2 % cmp_ratio
  
  q = torch.randn(B, S1, N1, D, dtype=torch.float16).npu()
  k = torch.randn(B, S2_compressed, N2, D, dtype=torch.float16).npu()
  w = torch.randn(B, S1, N1, dtype=torch.float32).npu()
  cmp_residual_k = torch.tensor([res_k] * B, dtype=torch.int32).npu()
  
  class LightningIndexerNetwork(torch.nn.Module):
      def __init__(self):
          super(LightningIndexerNetwork, self).__init__()
  
      def forward(self, q, k, w, cmp_residual_k):
          metadata = torch.ops.cann_ops_transformer.lightning_indexer_metadata(
              num_heads_q=N1,
              num_heads_k=N2,
              head_dim=D,
              topk=topk,
              cmp_residual_k=cmp_residual_k,
              batch_size=B,
              max_seqlen_q=S1,
              max_seqlen_k=S2_compressed,
              layout_q="BSND",
              layout_k="BSND",
              mask_mode=3,
              cmp_ratio=cmp_ratio
          )
  
          return torch.ops.cann_ops_transformer.lightning_indexer(
              q, k, w, topk,
              cmp_residual_k=cmp_residual_k,
              metadata=metadata,
              max_seqlen_q=-1,
              layout_q="BSND",
              layout_k="BSND",
              mask_mode=3,
              cmp_ratio=cmp_ratio,
              return_value=1
          )
  
  from torchair.configs.compiler_config import CompilerConfig
  config = CompilerConfig()
  config.mode = "reduce-overhead"
  npu_backend = torchair.get_npu_backend(compiler_config=config)
  torch._dynamo.reset()
  npu_mode = torch.compile(LightningIndexerNetwork(), fullgraph=True, backend=npu_backend, dynamic=False)
  sparse_indices, sparse_values = npu_mode(q, k, w, cmp_residual_k)
  print(f"sparse_indices shape: {sparse_indices.shape}")