nemo_automodel.components.models.deepseek_v4.kernels.tilelang_sparse_mla_fwd

Module Contents

Functions

Name	Description
`sparse_mqa_fwd`	-
`sparse_mqa_fwd_interface`	Forward interface for V4 sparse MQA attention.

API

nemo_automodel.components.models.deepseek_v4.kernels.tilelang_sparse_mla_fwd.sparse_mqa_fwd(
    heads,
    dim,
    topk,
    sm_scale = None,
    block_I = 64,
    num_stages = 2,
    threads = 256,
    reference_rounding = False
)

nemo_automodel.components.models.deepseek_v4.kernels.tilelang_sparse_mla_fwd.sparse_mqa_fwd_interface(
    q: torch.Tensor,
    kv: torch.Tensor,
    attn_sink: torch.Tensor,
    topk_idxs: torch.Tensor,
    sm_scale: float | None = None,
    block_I: int = 64,
    num_stages: int = 2,
    threads: int = 256,
    reference_rounding: bool = False
) -> tuple[torch.Tensor, torch.Tensor]

Forward interface for V4 sparse MQA attention.

Parameters:

torch.Tensor

Contiguous CUDA BF16 queries [batch, sequence, heads, head_dim].

torch.Tensor

Contiguous CUDA BF16 shared keys/values [batch, kv_sequence, head_dim].

attn_sink

torch.Tensor

CUDA FP32 denominator biases [heads].

topk_idxs

torch.Tensor

Contiguous CUDA integer indices [batch, sequence, slots]. Entries outside [0, kv_sequence) are masked; slots are internally padded to a multiple of block_I.

sm_scale

float | NoneDefaults to None

Score multiplier, defaulting to head_dim**-0.5.

block_I

intDefaults to 64

Sparse-key slots processed by each kernel iteration.

num_stages

intDefaults to 2

Pipeline stages for the generated kernel.

threads

intDefaults to 256

CUDA threads per block.

reference_rounding

boolDefaults to False

Match the original inference kernel’s scaled-logit FP32 arithmetic while retaining log2 LSE for the existing backward.

Returns: torch.Tensor

Independent BF16 output [batch, sequence, heads, head_dim] and FP32