nemo_automodel.components.models.deepseek_v4.optimized_kernels

Optional DeepSeek V4 optimized kernel dispatch.

The torch implementations below are kept as the numerical reference. Optional TileLang-backed paths are sourced from:

Sinkhorn: imported from DeepSeek TileKernels tile_kernels.modeling.mhc.ops.sinkhorn_normalize. No TileKernels source is vendored in AutoModel. Upstream source: https://github.com/deepseek-ai/TileKernels/blob/36d9e45d38e204ebb87e6f6e833821eee0482fe5/tile_kernels/modeling/mhc/ops/sinkhorn.py Upstream license: MIT, copyright 2026 DeepSeek.
Sparse attention and indexer: vendored/adapted Miles DeepSeek V4 ops in nemo_automodel.components.models.deepseek_v4.kernels. Upstream source: https://github.com/yueming-yuan/miles/tree/e561465d0b9bbf06188b7a5e2020dc7fd691f732/miles_plugins/models/deepseek_v4/ops Upstream license: Apache-2.0, copyright 2025 Zhipu AI. See nemo_automodel/components/models/deepseek_v4/kernels/__init__.py for the per-file attribution.

Those packages are imported with safe_import so environments without TileLang still import the model and use the existing torch path.

Module Contents

Classes

Name	Description
`SparseIndexPlan`	Sparse-attention indices prepared once for every layer that consumes the same selection.
`_Dsv4TileKernelsSinkhorn`	TileKernels Sinkhorn wrapper that accepts non-contiguous backward gradients.

Functions

Name	Description
`_all_cuda`	-
`_dsv4_sparse_attention_cudnn`	Run the shared cuDNN sparse attention on batched DSV4 tensors.
`_flatten_sparse_indices`	Turn per-batch `[batch, sequence, slots]` indices into int64 `[batch * sequence, slots]` global coordinates.
`_should_use_tilelang`	-
`_tile_kernels_sinkhorn_contiguous_grad`	-
`build_dsv4_sparse_topk_indices`	Build Miles-style top-k key indices for DSV4 local-window + compressed KV attention.
`dense_attention_topk_torch`	Compute a dense FP32 oracle for top-k sparse attention.
`dsv4_indexer_scores`	Run DSV4 C4 indexer scores through Miles TileLang kernels or torch fallback.
`dsv4_indexer_topk_scores`	Run DSV4 C4 top-k indexer scores through Miles autograd kernels or torch fallback.
`dsv4_sinkhorn_normalize`	Normalize HyperConnection combination logits with torch or TileKernels.
`dsv4_sparse_attention`	Run sparse attention with an optional original-inference rounding mode.
`extract_indexer_topk_scores_torch`	Extract top-k score values, masking `-1` entries with `-inf`.
`indexer_scores_torch`	Compute FP32 indexer scores for compressed key selection.
`is_dsv4_kernel_available`	Return whether the optional TileLang kernel package for `name` is importable.
`plan_dsv4_sparse_indices`	Prepare the per-call index work of `dsv4_sparse_attention` once for a shared selection.
`sinkhorn_normalize_torch`	Torch reference for TileKernels MHC Sinkhorn normalization.
`sparse_attention_torch`	Miles sparse MQA torch reference.

Data

Dsv4IndexerBackend

Dsv4SinkhornBackend

Dsv4SparseAttentionBackend

API

class nemo_automodel.components.models.deepseek_v4.optimized_kernels.SparseIndexPlan(
    backend: nemo_automodel.components.models.deepseek_v4.optimized_kernels.Dsv4SparseAttentionBackend,
    indices: torch.Tensor,
    topk_length: torch.Tensor | None,
    valid_rows: torch.Tensor | None,
    kv_sequence: int
)

Dataclass

Sparse-attention indices prepared once for every layer that consumes the same selection.

backend

Dsv4SparseAttentionBackend

The sparse backend the layout is prepared for ("cudnn" or "tilelang").

indices

Tensor

"cudnn": contiguous int32 [batch * sequence, 1, padded_width] with global flat K/V coordinates (per-batch offsets applied), an ascending valid prefix and a -1 suffix. "tilelang": contiguous int32 [batch, sequence, slots] with -1 masked, the kernel layout.

kv_sequence

int

Number of K/V positions per batch row the coordinates refer to.

topk_length

Tensor | None

"cudnn" only: contiguous int32 [batch * sequence] valid-prefix lengths.

valid_rows

Tensor | None

"cudnn" with padded inputs only: int64 indices of the nonempty query rows, so the backward skips its empty-row scan; None when every row is nonempty.

class nemo_automodel.components.models.deepseek_v4.optimized_kernels._Dsv4TileKernelsSinkhorn()

Bases: Function

TileKernels Sinkhorn wrapper that accepts non-contiguous backward gradients.

The upstream high-level wrapper launches the backward kernel with grad_output as-is. DSV4 consumes HC combinations through transposed matmul sites, so autograd can provide a transposed gradient layout. The low-level TileKernels backward kernel requires contiguous row-major inputs.

nemo_automodel.components.models.deepseek_v4.optimized_kernels._Dsv4TileKernelsSinkhorn.backward(
    ctx: torch.autograd.function.FunctionCtx,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor, None, None]

staticmethod

nemo_automodel.components.models.deepseek_v4.optimized_kernels._Dsv4TileKernelsSinkhorn.forward(
    ctx: torch.autograd.function.FunctionCtx,
    x: torch.Tensor,
    repeat: int,
    eps: float
) -> torch.Tensor

staticmethod

nemo_automodel.components.models.deepseek_v4.optimized_kernels._all_cuda(
    tensors: torch.Tensor = ()
) -> bool

nemo_automodel.components.models.deepseek_v4.optimized_kernels._dsv4_sparse_attention_cudnn(
    q: torch.Tensor,
    kv: torch.Tensor,
    sinks: torch.Tensor,
    topk_idxs: torch.Tensor | None,
    sm_scale: float,
    all_rows_nonempty: bool = False,
    plan: nemo_automodel.components.models.deepseek_v4.optimized_kernels.SparseIndexPlan | None = None
) -> torch.Tensor

Run the shared cuDNN sparse attention on batched DSV4 tensors.

Parameters:

torch.Tensor

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

torch.Tensor

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

sinks

torch.Tensor

FP32 per-head sink logits [heads].

topk_idxs

torch.Tensor | None

Per-batch sparse indices [batch, sequence, slots] with -1 masked; None with plan.

sm_scale

float

Softmax scale.

all_rows_nonempty

boolDefaults to False

Whether every query has at least one valid slot.

plan

SparseIndexPlan | NoneDefaults to None

Optional prepared "cudnn" plan for this kv, replacing topk_idxs.

Returns: torch.Tensor

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

nemo_automodel.components.models.deepseek_v4.optimized_kernels._flatten_sparse_indices(
    topk_idxs: torch.Tensor,
    kv_sequence: int
) -> torch.Tensor

Turn per-batch [batch, sequence, slots] indices into int64 [batch * sequence, slots] global coordinates.

nemo_automodel.components.models.deepseek_v4.optimized_kernels._should_use_tilelang(
    backend: str,
    available: bool,
    kernel_name: str,
    tensors: tuple[torch.Tensor, ...],
    require_bf16: bool = False
) -> bool

nemo_automodel.components.models.deepseek_v4.optimized_kernels._tile_kernels_sinkhorn_contiguous_grad(
    x: torch.Tensor,
    repeat: int,
    eps: float
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v4.optimized_kernels.build_dsv4_sparse_topk_indices(
    batch_size: int,
    seq_len: int,
    key_len: int,
    window_size: int,
    device: torch.device,
    attention_mask: torch.Tensor | None = None,
    compress_ratio: int = 0,
    compressed_topk: torch.Tensor | None = None,
    n_pooled: int = 0,
    vanilla_key_len: int | None = None,
    q_positions: torch.Tensor | None = None,
    vision_token_types: torch.Tensor | None = None,
    max_image_tokens: int = 0
) -> torch.Tensor

Build Miles-style top-k key indices for DSV4 local-window + compressed KV attention.

Parameters:

vision_token_types

torch.Tensor | NoneDefaults to None

Optional pseudo-token types with shape [batch, sequence]. Text positions contain -1; image spans contain IMAGE_START, IMAGE, and IMAGE_END token types. When provided with max_image_tokens > 0, image queries additionally see the full image span.

max_image_tokens

intDefaults to 0

Maximum number of tokens in an image span. 0 disables visual windowing.

nemo_automodel.components.models.deepseek_v4.optimized_kernels.dense_attention_topk_torch(
    q: torch.Tensor,
    kv: torch.Tensor,
    sinks: torch.Tensor,
    topk_idxs: torch.Tensor,
    sm_scale: float
) -> torch.Tensor

Compute a dense FP32 oracle for top-k sparse attention.

Parameters:

torch.Tensor

Query tensor of shape [batch, sequence, heads, head_dim].

torch.Tensor

Shared key/value tensor of shape [batch, keys, head_dim].

sinks

torch.Tensor

Attention-sink logits of shape [heads].

topk_idxs

torch.Tensor

Integer tensor [batch, sequence, topk]; -1 masks an entry.

sm_scale

float

Attention scaling factor.

Returns: torch.Tensor

Tensor [batch, sequence, heads, head_dim] in q’s dtype.

nemo_automodel.components.models.deepseek_v4.optimized_kernels.dsv4_indexer_scores(
    q: torch.Tensor,
    pooled_kv: torch.Tensor,
    weights: torch.Tensor,
    compress_ratio: int,
    softmax_scale: float,
    backend: nemo_automodel.components.models.deepseek_v4.optimized_kernels.Dsv4IndexerBackend,
    query_start: int = 0,
    query_total_len: int | None = None
) -> torch.Tensor

Run DSV4 C4 indexer scores through Miles TileLang kernels or torch fallback.

nemo_automodel.components.models.deepseek_v4.optimized_kernels.dsv4_indexer_topk_scores(
    q: torch.Tensor,
    pooled_kv: torch.Tensor,
    weights: torch.Tensor,
    topk_indices: torch.Tensor,
    compress_ratio: int,
    softmax_scale: float,
    backend: nemo_automodel.components.models.deepseek_v4.optimized_kernels.Dsv4IndexerBackend
) -> torch.Tensor

Run DSV4 C4 top-k indexer scores through Miles autograd kernels or torch fallback.

nemo_automodel.components.models.deepseek_v4.optimized_kernels.dsv4_sinkhorn_normalize(
    x: torch.Tensor,
    backend: nemo_automodel.components.models.deepseek_v4.optimized_kernels.Dsv4SinkhornBackend,
    repeat: int,
    eps: float
) -> torch.Tensor

Normalize HyperConnection combination logits with torch or TileKernels.

nemo_automodel.components.models.deepseek_v4.optimized_kernels.dsv4_sparse_attention(
    q: torch.Tensor,
    kv: torch.Tensor,
    sinks: torch.Tensor,
    topk_idxs: torch.Tensor | None,
    sm_scale: float,
    backend: nemo_automodel.components.models.deepseek_v4.optimized_kernels.Dsv4SparseAttentionBackend,
    reference_rounding: bool = False,
    backward_impl: str = 'atomic',
    backward_workspace_bytes: int = 2 * 1024 ** 3,
    all_rows_nonempty: bool = False,
    plan: nemo_automodel.components.models.deepseek_v4.optimized_kernels.SparseIndexPlan | None = None
) -> torch.Tensor

Run sparse attention with an optional original-inference rounding mode.

Parameters:

torch.Tensor

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

torch.Tensor

Shared keys/values [batch, kv_sequence, head_dim].

sinks

torch.Tensor

FP32 denominator biases [heads].

topk_idxs

torch.Tensor | None

Sparse indices [batch, sequence, slots], with -1 masked; None only together with plan, which then carries the prepared indices.

sm_scale

float

Query/key score multiplier.

backend

Dsv4SparseAttentionBackend

Sparse execution backend.

backward_impl

strDefaults to 'atomic'

Atomic or fixed-order deterministic backward (TileLang only).

backward_workspace_bytes

intDefaults to 2 * 1024 ** 3

Partial-KV workspace cap; auxiliary buffers are extra.

reference_rounding

boolDefaults to False

Scale logits before max reduction and use the original inference exponent arithmetic. Requires TileLang.

all_rows_nonempty

boolDefaults to False

Whether every query has at least one valid slot ("cudnn").

plan

SparseIndexPlan | NoneDefaults to None

Optional SparseIndexPlan from plan_dsv4_sparse_indices for this backend and kv; it replaces the per-call index preparation of topk_idxs.

Returns: torch.Tensor

Attention output [batch, sequence, heads, head_dim], with q’s dtype.

nemo_automodel.components.models.deepseek_v4.optimized_kernels.extract_indexer_topk_scores_torch(
    logits: torch.Tensor,
    topk_indices: torch.Tensor
) -> torch.Tensor

Extract top-k score values, masking -1 entries with -inf.

nemo_automodel.components.models.deepseek_v4.optimized_kernels.indexer_scores_torch(
    q: torch.Tensor,
    pooled_kv: torch.Tensor,
    weights: torch.Tensor,
    softmax_scale: float
) -> torch.Tensor

Compute FP32 indexer scores for compressed key selection.

Parameters:

torch.Tensor

Query tensor of shape [batch, sequence, heads, head_dim].

pooled_kv

torch.Tensor

Compressed keys of shape [batch, compressed, head_dim].

weights

torch.Tensor

Head weights of shape [batch, sequence, heads].

softmax_scale

float

Query/key scaling factor.

Returns: torch.Tensor

FP32 scores of shape [batch, sequence, compressed].

nemo_automodel.components.models.deepseek_v4.optimized_kernels.is_dsv4_kernel_available(
    name: typing.Literal['sinkhorn', 'sparse_attn', 'indexer']
) -> bool

Return whether the optional TileLang kernel package for name is importable.

nemo_automodel.components.models.deepseek_v4.optimized_kernels.plan_dsv4_sparse_indices(
    topk_idxs: torch.Tensor,
    kv_sequence: int,
    backend: nemo_automodel.components.models.deepseek_v4.optimized_kernels.Dsv4SparseAttentionBackend,
    all_rows_nonempty: bool = False
) -> nemo_automodel.components.models.deepseek_v4.optimized_kernels.SparseIndexPlan

Prepare the per-call index work of dsv4_sparse_attention once for a shared selection.

Parameters:

topk_idxs

torch.Tensor

Sparse indices [batch, sequence, slots] with -1 masked.

kv_sequence

int

Number of K/V positions per batch row.

backend

Dsv4SparseAttentionBackend

Sparse execution backend the plan is for.

all_rows_nonempty

boolDefaults to False

Whether every query has at least one valid slot ("cudnn": when false the plan also carries the nonempty rows, one host synchronisation per plan instead of one per layer backward).

Returns: SparseIndexPlan

The plan to pass as plan= to dsv4_sparse_attention together with the same kv.

nemo_automodel.components.models.deepseek_v4.optimized_kernels.sinkhorn_normalize_torch(
    x: torch.Tensor,
    repeat: int,
    eps: float
) -> torch.Tensor

Torch reference for TileKernels MHC Sinkhorn normalization.

nemo_automodel.components.models.deepseek_v4.optimized_kernels.sparse_attention_torch(
    q: torch.Tensor,
    kv: torch.Tensor,
    sinks: torch.Tensor,
    topk_idxs: torch.Tensor,
    sm_scale: float
) -> torch.Tensor

Miles sparse MQA torch reference.

Parameters:

torch.Tensor

Query tensor with shape [B, S, H, D].

torch.Tensor

Single-head KV tensor with shape [B, K, D].

sinks

torch.Tensor

Per-head attention sink logits with shape [H].

topk_idxs

torch.Tensor

Key indices with shape [B, S, K_top]; -1 masks an entry.

sm_scale

float

Attention scaling factor.

nemo_automodel.components.models.deepseek_v4.optimized_kernels.Dsv4IndexerBackend = Literal['torch', 'tilelang', 'auto']

nemo_automodel.components.models.deepseek_v4.optimized_kernels.Dsv4SinkhornBackend = Literal['torch', 'tilelang', 'auto']

nemo_automodel.components.models.deepseek_v4.optimized_kernels.Dsv4SparseAttentionBackend = Literal['torch', 'sparse_torch', 'tilelang', 'cudnn', 'auto']