sparse_flash_mla

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:

  sparse_flash_mla is a cann_ops_transformer extension API based on TorchNPU. It calls the SparseFlashMla operator to perform sparse attention computation on shared KV (Key and Value use the same input). The API supports the following three computation modes:

  • SWA (Sliding Window Attention): uses only ori_kv to perform sliding window attention on the original KV.
  • CSA (Compressed Sparse Attention): uses ori_kv, cmp_kv, and cmp_sparse_indices together to perform attention on the original KV window and the compressed KV selected by TopK.
  • HCA (Heavily Compressed Attention): Uses both ori_kv and cmp_kv to perform attention on the original KV window and the contiguous compressed KV segments together.

  sparse_flash_mla_metadata is the metadata pre-processing API of sparse_flash_mla. It generates metadata before the main API runs. The metadata records the task splitting result of AICore/AIVCore. The main API must be passed in this metadata. The typical call flow is as follows:

  1. Prepare the inputs such as q, ori_kv, cmp_kv, sequence length, block table, and sinks.
  2. Call sparse_flash_mla_metadata to generate metadata.
  3. Call sparse_flash_mla and pass the metadata obtained in the previous step into the main operator.

• Formula:

  $$O = \text{softmax}(Q \cdot \tilde{K}^{T} \cdot \text{softmax\_scale}) \cdot \tilde{V}$$

  $\tilde{K}=\tilde{V}$, which consists of the sliding window part of ori_kv and the compressed part of cmp_kv. The KV range actually participating in the calculation is determined by cmp_ratio, ori_mask_mode, cmp_mask_mode, ori_win_left, ori_win_right, and cmp_sparse_indices.

Prototype

Before calling the sparse_flash_mla API, call the sparse_flash_mla_metadata API first to complete the load balancing calculation.

cann_ops_transformer.sparse_flash_mla_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=0,
    max_seqlen_q=0,
    max_seqlen_ori_kv=0,
    max_seqlen_cmp_kv=0,
    ori_topk=0,
    cmp_topk=0,
    cmp_ratio=1,
    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

cann_ops_transformer.sparse_flash_mla(
    q,
    *,
    ori_kv=None,
    cmp_kv=None,
    ori_sparse_indices=None,
    cmp_sparse_indices=None,
    ori_block_table=None,
    cmp_block_table=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=1.0,
    cmp_ratio=1,
    ori_mask_mode=0,
    cmp_mask_mode=0,
    ori_win_left=-1,
    ori_win_right=-1,
    layout_q="BSND",
    layout_kv="BSND",
    topk_value_mode=1,
    return_softmax_lse=False
) -> (Tensor, Tensor)

Parameters

sparse_flash_mla_metadata

sparse_flash_mla

Return Value

sparse_flash_mla_metadata

• metadata: Output of sparse_flash_mla_metadata. The shape is fixed to [1024] and the dtype is int32.

sparse_flash_mla

• attention_out: The first output of sparse_flash_mla. Its shape is the same as q, and its dtype is the same as q.
• softmax_lse: The second output of sparse_flash_mla. When return_softmax_lse=False, a FLOAT32 scalar placeholder tensor is returned. When return_softmax_lse=True, the FLOAT32 log-sum-exp result is returned.

Constraints

• Common constraints:

  • Applicable use case: This API supports both training and inference use cases.
  • Call mode: This API supports single-operator mode and TorchAir graph mode calls.
  • The data types of q, ori_kv, and cmp_kv must be the same. FLOAT16 and BFLOAT16 are supported.
  • The layout_q and layout_kv combinations only support "BSND"/"BSND", "TND"/"TND", "BSND"/"PA_BBND", and "TND"/"PA_BBND". In non-PA_BBND use cases, layout_q and layout_kv must be the same.
  • layout_q supports "BSND" and "TND". When layout_q="BSND", q must be 4-dimensional. When layout_q="TND", q must be 3-dimensional and cu_seqlens_q must be passed in.
  • layout_kv supports "BSND", "TND", and "PA_BBND". When layout_kv="BSND" or layout_kv="PA_BBND", ori_kv and cmp_kv must be 4-dimensional. When layout_kv="TND", ori_kv and cmp_kv must be 3-dimensional.
  • When layout_kv="TND", cu_seqlens_ori_kv must be passed in. When cmp_kv is passed in, cu_seqlens_cmp_kv must also be passed in.
  • When layout_kv="PA_BBND", seqused_ori_kv and ori_block_table must be passed in. When cmp_kv is passed in, cmp_block_table must also be passed in.
  • seqused_cmp_kv is an optional input under all layout_kv values. When explicitly passed in, it overrides the effective length of the cmp-side logic.
  • metadata is fixed to 1024 INT32 elements. topk_value_mode only supports 1. ori_sparse_indices, ori_topk_length, and cmp_topk_length do not support passing in a non-empty Tensor in the current version.
  • ori_kv and cmp_kv allow non-contiguous memory with inter-row padding. The API obtains the stride information through aclnn and passes it to the underlying operator.

• Specification constraints:

  • Common parameter constraints:
    • head_dim only supports 512, and num_heads_kv only supports 1.
    • num_heads_q / num_heads_kv supports 1, 2, 4, 8, 16, 32, 64, and 128.
    • ori_mask_mode only supports 4, ori_win_left only supports 127, and ori_win_right only supports 0.
    • The block_size of PageAttention supports multiples of 16 and does not exceed 1024.
  • SWA:
    • When only ori_kv is passed in, cmp_ratio does not participate in the compressed KV calculation and must keep the default value 1.
    • Do not pass in cmp_kv, cmp_sparse_indices, or cmp_block_table.
    • Pass 0 for cmp_topk and 0 for cmp_mask_mode.
  • CSA:
    • cmp_ratio only supports 4, and cmp_mask_mode only supports 3.
    • cmp_sparse_indices must be passed in, and its last dimension supports 512 or 1024; pass 512 or 1024 for cmp_topk accordingly.
    • cmp_residual_kv must be passed in, and its length must be equal to the batch size.
  • HCA:
    • cmp_ratio only supports 128, and cmp_mask_mode only supports 3.
    • Do not pass in cmp_sparse_indices; pass 0 to cmp_topk.
    • cmp_residual_kv must be passed in, and its length must be equal to the batch size.

Deterministic Computation

• Defaults to a deterministic computation.
• Batch invariance is supported by default.

Example

SWA, BSND Input

import math
import torch
import torch_npu
import cann_ops_transformer

torch_npu.npu.set_device(0)

dtype = torch.bfloat16
B = 1
S1 = 16
S2 = 64
N1 = 64
N2 = 1
D = 512
cmp_ratio = 1  # The SWA example only passes in ori_kv, so cmp_ratio does not participate in the compressed KV calculation and keeps the default value 1.

q = torch.randn(B, S1, N1, D, dtype=dtype, device="npu")
ori_kv = torch.randn(B, S2, N2, D, dtype=dtype, device="npu")
sinks = torch.zeros(N1, dtype=torch.float32, device="npu")

metadata = cann_ops_transformer.sparse_flash_mla_metadata(
    N1,
    N2,
    D,
    batch_size=B,
    max_seqlen_q=S1,
    max_seqlen_ori_kv=S2,
    ori_topk=0,
    cmp_topk=0,
    cmp_ratio=cmp_ratio,
    ori_mask_mode=4,
    cmp_mask_mode=3,
    ori_win_left=127,
    ori_win_right=0,
    layout_q="BSND",
    layout_kv="BSND",
    has_ori_kv=True,
    has_cmp_kv=False,
)

attn_out, softmax_lse = cann_ops_transformer.sparse_flash_mla(
    q,
    ori_kv=ori_kv,
    sinks=sinks,
    metadata=metadata,
    softmax_scale=1.0 / math.sqrt(D),
    cmp_ratio=cmp_ratio,
    ori_mask_mode=4,
    cmp_mask_mode=3,
    ori_win_left=127,
    ori_win_right=0,
    layout_q="BSND",
    layout_kv="BSND",
    return_softmax_lse=False,
)
torch_npu.npu.synchronize()
assert attn_out.shape == q.shape
assert attn_out.dtype == q.dtype
assert softmax_lse.shape == torch.Size([])
assert torch.isfinite(attn_out.float()).all().item()

HCA, BSND input

import math
import torch
import torch_npu
import cann_ops_transformer

torch_npu.npu.set_device(0)

dtype = torch.bfloat16
B = 1
S1 = 16
S2 = 128
S3 = 1
N1 = 64
N2 = 1
D = 512
cmp_ratio = 128

q = torch.randn(B, S1, N1, D, dtype=dtype, device="npu")
ori_kv = torch.randn(B, S2, N2, D, dtype=dtype, device="npu")
cmp_kv = torch.randn(B, S3, N2, D, dtype=dtype, device="npu")
cmp_residual_kv = torch.zeros(B, dtype=torch.int32, device="npu")
sinks = torch.zeros(N1, dtype=torch.float32, device="npu")

metadata = cann_ops_transformer.sparse_flash_mla_metadata(
    N1,
    N2,
    D,
    batch_size=B,
    max_seqlen_q=S1,
    max_seqlen_ori_kv=S2,
    max_seqlen_cmp_kv=S3,
    cmp_residual_kv=cmp_residual_kv,
    ori_topk=0,
    cmp_topk=0,
    cmp_ratio=cmp_ratio,
    ori_mask_mode=4,
    cmp_mask_mode=3,
    ori_win_left=127,
    ori_win_right=0,
    layout_q="BSND",
    layout_kv="BSND",
    has_ori_kv=True,
    has_cmp_kv=True,
)

attn_out, softmax_lse = cann_ops_transformer.sparse_flash_mla(
    q,
    ori_kv=ori_kv,
    cmp_kv=cmp_kv,
    cmp_residual_kv=cmp_residual_kv,
    sinks=sinks,
    metadata=metadata,
    softmax_scale=1.0 / math.sqrt(D),
    cmp_ratio=cmp_ratio,
    ori_mask_mode=4,
    cmp_mask_mode=3,
    ori_win_left=127,
    ori_win_right=0,
    layout_q="BSND",
    layout_kv="BSND",
    return_softmax_lse=False,
)
torch_npu.npu.synchronize()
assert attn_out.shape == q.shape
assert attn_out.dtype == q.dtype
assert softmax_lse.shape == torch.Size([])
assert torch.isfinite(attn_out.float()).all().item()

CSA, TND input with cmp_residual_kv enabled

import math
import torch
import torch_npu
import cann_ops_transformer

torch_npu.npu.set_device(0)

dtype = torch.float16
B = 1
q_lens = [1]
ori_lens = [6]
cmp_lens = [1]
N1 = 64
N2 = 1
D = 512
K = 512
cmp_ratio = 4

cu_q = torch.tensor([0, 1], dtype=torch.int32, device="npu")
cu_ori = torch.tensor([0, 6], dtype=torch.int32, device="npu")
cu_cmp = torch.tensor([0, 1], dtype=torch.int32, device="npu")
cmp_residual_kv = torch.tensor([2], dtype=torch.int32, device="npu")

q = torch.randn(sum(q_lens), N1, D, dtype=dtype, device="npu")
ori_kv = torch.randn(sum(ori_lens), N2, D, dtype=dtype, device="npu")
cmp_kv = torch.randn(sum(cmp_lens), N2, D, dtype=dtype, device="npu")
sinks = torch.zeros(N1, dtype=torch.float32, device="npu")

cmp_sparse_indices = torch.full((sum(q_lens), N2, K), -1, dtype=torch.int32, device="npu")
cmp_sparse_indices[:, :, :1] = torch.arange(1, dtype=torch.int32, device="npu").view(1, 1, 1)

metadata = cann_ops_transformer.sparse_flash_mla_metadata(
    N1,
    N2,
    D,
    cu_seqlens_q=cu_q,
    cu_seqlens_ori_kv=cu_ori,
    cu_seqlens_cmp_kv=cu_cmp,
    max_seqlen_q=max(q_lens),
    max_seqlen_ori_kv=max(ori_lens),
    max_seqlen_cmp_kv=max(cmp_lens),
    ori_topk=0,
    cmp_topk=K,
    cmp_ratio=cmp_ratio,
    ori_mask_mode=4,
    cmp_mask_mode=3,
    cmp_residual_kv=cmp_residual_kv,
    ori_win_left=127,
    ori_win_right=0,
    layout_q="TND",
    layout_kv="TND",
    has_ori_kv=True,
    has_cmp_kv=True,
)

attn_out, softmax_lse = cann_ops_transformer.sparse_flash_mla(
    q,
    ori_kv=ori_kv,
    cmp_kv=cmp_kv,
    cmp_sparse_indices=cmp_sparse_indices,
    cu_seqlens_q=cu_q,
    cu_seqlens_ori_kv=cu_ori,
    cu_seqlens_cmp_kv=cu_cmp,
    cmp_residual_kv=cmp_residual_kv,
    sinks=sinks,
    metadata=metadata,
    softmax_scale=1.0 / math.sqrt(D),
    cmp_ratio=cmp_ratio,
    ori_mask_mode=4,
    cmp_mask_mode=3,
    ori_win_left=127,
    ori_win_right=0,
    layout_q="TND",
    layout_kv="TND",
    return_softmax_lse=False,
)
torch_npu.npu.synchronize()
assert attn_out.shape == q.shape
assert attn_out.dtype == q.dtype
assert softmax_lse.shape == torch.Size([])
assert torch.isfinite(attn_out.float()).all().item()

CP Split Example, TND + CSA, rank0 Splits the Second seq

The following example uses a single process to sequentially simulate two CP ranks, illustrating the relationship between the global TND data and the input parameters of each rank. Assume there are two sequences globally and cmp_ratio=4:

Although rank1 only calculates [8,18) of seq1, ori_kv and cmp_kv must be passed in as the prefix up to the current position. In this case, ori_prefix_len - q_len = 18 - 10 = 8, and the q start derived by the kernel is exactly the CP split point. Each local batch must satisfy cmp_len * cmp_ratio + cmp_residual_kv[b] == ori_prefix_len.

import math
import torch
import torch_npu
import cann_ops_transformer


torch_npu.npu.set_device(0)

dtype = torch.float16
cmp_ratio = 4
K = 512
N1 = 64
N2 = 1
D = 512

# Global packed TND perspective: seq0 length is 16, seq1 length is 18.
global_q_lens = [16, 18]
global_ori_lens = [16, 18]
global_cmp_lens = [4, 4]
global_cmp_residual = [0, 2]

q_global = torch.randn(sum(global_q_lens), N1, D, dtype=dtype, device="npu")
ori_global = torch.randn(sum(global_ori_lens), N2, D, dtype=dtype, device="npu")
cmp_global = torch.randn(sum(global_cmp_lens), N2, D, dtype=dtype, device="npu")
sinks = torch.zeros(N1, dtype=torch.float32, device="npu")


def make_cu(lengths):
    cu = [0]
    for length in lengths:
        cu.append(cu[-1] + length)
    return torch.tensor(cu, dtype=torch.int32, device="npu")


def make_cmp_sparse_indices(q_lens, ori_prefix_lens, cmp_lens):
    indices = torch.full((sum(q_lens), N2, K), -1, dtype=torch.int32, device="npu")
    q_base = 0
    for q_len, ori_prefix_len, cmp_len in zip(q_lens, ori_prefix_lens, cmp_lens):
        q_start = ori_prefix_len - q_len
        for row in range(q_len):
            q_pos = q_start + row
            cmp_end = min(cmp_len, (q_pos + 1) // cmp_ratio)
            if cmp_end > 0:
                indices[q_base + row, :, :cmp_end] = torch.arange(
                    cmp_end, dtype=torch.int32, device="npu"
                ).view(1, cmp_end)
        q_base += q_len
    return indices


def run_one_rank(name, q, ori_kv, cmp_kv, q_lens, ori_prefix_lens, cmp_lens, residuals):
    cu_q = make_cu(q_lens)
    cu_ori = make_cu(ori_prefix_lens)
    cu_cmp = make_cu(cmp_lens)
    cmp_residual_kv = torch.tensor(residuals, dtype=torch.int32, device="npu")
    cmp_sparse_indices = make_cmp_sparse_indices(q_lens, ori_prefix_lens, cmp_lens)

    metadata = cann_ops_transformer.sparse_flash_mla_metadata(
        N1,
        N2,
        D,
        cu_seqlens_q=cu_q,
        cu_seqlens_ori_kv=cu_ori,
        cu_seqlens_cmp_kv=cu_cmp,
        max_seqlen_q=max(q_lens),
        max_seqlen_ori_kv=max(ori_prefix_lens),
        max_seqlen_cmp_kv=max(cmp_lens),
        ori_topk=0,
        cmp_topk=K,
        cmp_ratio=cmp_ratio,
        ori_mask_mode=4,
        cmp_mask_mode=3,
        cmp_residual_kv=cmp_residual_kv,
        ori_win_left=127,
        ori_win_right=0,
        layout_q="TND",
        layout_kv="TND",
        has_ori_kv=True,
        has_cmp_kv=True,
    )

    attn_out, softmax_lse = cann_ops_transformer.sparse_flash_mla(
        q,
        ori_kv=ori_kv,
        cmp_kv=cmp_kv,
        cmp_sparse_indices=cmp_sparse_indices,
        cu_seqlens_q=cu_q,
        cu_seqlens_ori_kv=cu_ori,
        cu_seqlens_cmp_kv=cu_cmp,
        cmp_residual_kv=cmp_residual_kv,
        sinks=sinks,
        metadata=metadata,
        softmax_scale=1.0 / math.sqrt(D),
        cmp_ratio=cmp_ratio,
        ori_mask_mode=4,
        cmp_mask_mode=3,
        ori_win_left=127,
        ori_win_right=0,
        layout_q="TND",
        layout_kv="TND",
        return_softmax_lse=False,
    )
    torch_npu.npu.synchronize()
    assert attn_out.shape == q.shape, name
    assert attn_out.dtype == q.dtype, name
    assert softmax_lse.shape == torch.Size([]), name
    assert torch.isfinite(attn_out.float()).all().item(), name
    return attn_out


# rank0: contains the complete seq0 and splits to the first 8 tokens of seq1.
rank0_q = torch.cat([q_global[0:16], q_global[16:24]], dim=0)
rank0_ori = torch.cat([ori_global[0:16], ori_global[16:24]], dim=0)
rank0_cmp = torch.cat([cmp_global[0:4], cmp_global[4:6]], dim=0)
rank0_out = run_one_rank(
    "rank0", rank0_q, rank0_ori, rank0_cmp,
    q_lens=[16, 8], ori_prefix_lens=[16, 8], cmp_lens=[4, 2], residuals=[0, 0]
)

# rank1: only calculates the last 10 tokens of seq1, but ori_kv/cmp_kv pass in the prefix of seq1 up to 18.
rank1_q = q_global[24:34]
rank1_ori = ori_global[16:34]
rank1_cmp = cmp_global[4:8]
rank1_out = run_one_rank(
    "rank1", rank1_q, rank1_ori, rank1_cmp,
    q_lens=[10], ori_prefix_lens=[18], cmp_lens=[4], residuals=[2]
)

seq0_out = rank0_out[:16]
seq1_out = torch.cat([rank0_out[16:24], rank1_out], dim=0)
assert seq0_out.shape == (16, N1, D)
assert seq1_out.shape == (18, N1, D)