ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsDeepseek V4Kernels

nemo_automodel.components.models.deepseek_v4.kernels.sparse_attention

View as Markdown

Autograd wrapper for vendored Miles DeepSeek V4 sparse-attention kernels.

Attribution:

Module Contents

Classes

NameDescription
DeepSeekV4SparseAttentionTileLang sparse MQA attention with custom backward.
DeepSeekV4SparseAttentionHeadChunkedTileLang sparse attention with smaller head groups and fp32 KV-grad accumulation.

Functions

NameDescription
sparse_attn_tilelangRun vendored Miles DeepSeek V4 TileLang sparse attention.
sparse_attn_tilelang_head_chunkedRun vendored sparse attention with bounded backward head chunks.

API

class nemo_automodel.components.models.deepseek_v4.kernels.sparse_attention.DeepSeekV4SparseAttention()

Bases: Function

TileLang sparse MQA attention with custom backward.

nemo_automodel.components.models.deepseek_v4.kernels.sparse_attention.DeepSeekV4SparseAttention.backward(
    ctx: torch.autograd.function.FunctionCtx,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor | None, ...]
staticmethod

Run the existing backward using the saved base-two LSE.

Parameters:

ctx
torch.autograd.function.FunctionCtx

Saved forward tensors and score multiplier.

grad_output
torch.Tensor

CUDA output gradient [batch, sequence, heads, head_dim]. Non-contiguous gradients are copied before the kernel call.

Returns: torch.Tensor | None

Query gradient [batch, sequence, heads, head_dim], shared-KV gradient

nemo_automodel.components.models.deepseek_v4.kernels.sparse_attention.DeepSeekV4SparseAttention.forward(
    ctx: torch.autograd.function.FunctionCtx,
    q: torch.Tensor,
    kv: torch.Tensor,
    attn_sink: torch.Tensor,
    topk_idxs: torch.Tensor,
    sm_scale: float | None = None,
    reference_rounding: bool = False
) -> torch.Tensor
staticmethod

Run the vendored sparse attention forward kernel.

Parameters:

ctx
torch.autograd.function.FunctionCtx

Autograd context retaining inputs, output and base-two LSE.

q
torch.Tensor

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

kv
torch.Tensor

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

attn_sink
torch.Tensor

CUDA FP32 softmax denominator biases [heads].

topk_idxs
torch.Tensor

Contiguous CUDA integer indices [batch, sequence, slots], with -1 for masked slots.

sm_scale
float | NoneDefaults to None

Score multiplier, defaulting to head_dim**-0.5.

reference_rounding
boolDefaults to False

Use scaled-logit exponent arithmetic matching the released inference kernel.

Returns: torch.Tensor

Independent CUDA BF16 output [batch, sequence, heads, head_dim].

class nemo_automodel.components.models.deepseek_v4.kernels.sparse_attention.DeepSeekV4SparseAttentionHeadChunked()

Bases: Function

TileLang sparse attention with smaller head groups and fp32 KV-grad accumulation.

nemo_automodel.components.models.deepseek_v4.kernels.sparse_attention.DeepSeekV4SparseAttentionHeadChunked.backward(
    ctx: torch.autograd.function.FunctionCtx,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor | None, ...]
staticmethod

Run chunked backward and accumulate shared KV gradients in FP32.

Parameters:

ctx
torch.autograd.function.FunctionCtx

Saved forward tensors, score multiplier and backward chunk size.

grad_output
torch.Tensor

CUDA output gradient [batch, sequence, heads, head_dim]. Non-contiguous gradients are copied before the kernel call.

Returns: torch.Tensor | None

Query gradient [batch, sequence, heads, head_dim], shared-KV gradient

nemo_automodel.components.models.deepseek_v4.kernels.sparse_attention.DeepSeekV4SparseAttentionHeadChunked.forward(
    ctx: torch.autograd.function.FunctionCtx,
    q: torch.Tensor,
    kv: torch.Tensor,
    attn_sink: torch.Tensor,
    topk_idxs: torch.Tensor,
    max_heads_per_kernel: int,
    sm_scale: float | None = None,
    reference_rounding: bool = False
) -> torch.Tensor
staticmethod

Run forward over head chunks while preserving backward chunk size.

Parameters:

ctx
torch.autograd.function.FunctionCtx

Autograd context retaining tensors and backward chunk size.

q
torch.Tensor

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

kv
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], with -1 for masked slots.

max_heads_per_kernel
int

Maximum heads per backward kernel. Forward uses the same count unless reference mode requires at least 64.

sm_scale
float | NoneDefaults to None

Score multiplier, defaulting to head_dim**-0.5.

reference_rounding
boolDefaults to False

Preserve the original 64-head forward layout and scaled-logit exponent arithmetic.

Returns: torch.Tensor

Independent CUDA BF16 output [batch, sequence, heads, head_dim].

nemo_automodel.components.models.deepseek_v4.kernels.sparse_attention.sparse_attn_tilelang(
    q: torch.Tensor,
    kv: torch.Tensor,
    attn_sink: torch.Tensor,
    topk_idxs: torch.Tensor,
    sm_scale: float | None = None,
    reference_rounding: bool = False
) -> torch.Tensor

Run vendored Miles DeepSeek V4 TileLang sparse attention.

Parameters:

q
torch.Tensor

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

kv
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], with -1 for masked slots.

sm_scale
float | NoneDefaults to None

Score multiplier, defaulting to head_dim**-0.5.

reference_rounding
boolDefaults to False

Select released-inference forward exponent arithmetic.

Returns: torch.Tensor

Independent BF16 tensor [batch, sequence, heads, head_dim], with the

nemo_automodel.components.models.deepseek_v4.kernels.sparse_attention.sparse_attn_tilelang_head_chunked(
    q: torch.Tensor,
    kv: torch.Tensor,
    attn_sink: torch.Tensor,
    topk_idxs: torch.Tensor,
    max_heads_per_kernel: int,
    sm_scale: float | None = None,
    reference_rounding: bool = False
) -> torch.Tensor

Run vendored sparse attention with bounded backward head chunks.

Parameters:

q
torch.Tensor

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

kv
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], with -1 for masked slots.

max_heads_per_kernel
int

Backward head chunk size; forward uses at least 64 heads per chunk when reference_rounding is enabled.

sm_scale
float | NoneDefaults to None

Score multiplier, defaulting to head_dim**-0.5.

reference_rounding
boolDefaults to False

Select original-inference forward rounding while retaining smaller backward chunks and FP32 shared-KV accumulation.

Returns: torch.Tensor

Independent CUDA BF16 tensor [batch, sequence, heads, head_dim].