> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.deepseek_v4.kernels.sparse_attention

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

Attribution:

* Upstream project: Miles, [https://github.com/yueming-yuan/miles](https://github.com/yueming-yuan/miles)
* Upstream revision: e561465d0b9bbf06188b7a5e2020dc7fd691f732, deepseek-v4 branch
* Upstream license: Apache-2.0, copyright 2025 Zhipu AI
* Original source:
  [https://github.com/yueming-yuan/miles/blob/e561465d0b9bbf06188b7a5e2020dc7fd691f732/miles\_plugins/models/deepseek\_v4/ops/attention\_core.py](https://github.com/yueming-yuan/miles/blob/e561465d0b9bbf06188b7a5e2020dc7fd691f732/miles_plugins/models/deepseek_v4/ops/attention_core.py)

## Module Contents

### Classes

| Name                                                                                                                                                  | Description                                                                       |
| ----------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------- |
| [`DeepSeekV4SparseAttention`](#nemo_automodel-components-models-deepseek_v4-kernels-sparse_attention-DeepSeekV4SparseAttention)                       | TileLang sparse MQA attention with custom backward.                               |
| [`DeepSeekV4SparseAttentionHeadChunked`](#nemo_automodel-components-models-deepseek_v4-kernels-sparse_attention-DeepSeekV4SparseAttentionHeadChunked) | TileLang sparse attention with smaller head groups and fp32 KV-grad accumulation. |

### Functions

| Name                                                                                                                                            | Description                                                      |
| ----------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------- |
| [`sparse_attn_tilelang`](#nemo_automodel-components-models-deepseek_v4-kernels-sparse_attention-sparse_attn_tilelang)                           | Run vendored Miles DeepSeek V4 TileLang sparse attention.        |
| [`sparse_attn_tilelang_head_chunked`](#nemo_automodel-components-models-deepseek_v4-kernels-sparse_attention-sparse_attn_tilelang_head_chunked) | Run vendored sparse attention with bounded backward head chunks. |

### API

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

**Bases:** `Function`

TileLang sparse MQA attention with custom backward.

```python
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

```python
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 | None` — default: None

Score multiplier, defaulting to head\_dim\*\*-0.5.

---

**`reference_rounding`** `bool` — default: False

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

---

**Returns:** `torch.Tensor`

Independent CUDA BF16 output \[batch, sequence, heads, head\_dim].

```python
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.

```python
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

```python
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 | None` — default: None

Score multiplier, defaulting to head\_dim\*\*-0.5.

---

**`reference_rounding`** `bool` — default: 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].

```python
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 | None` — default: None

Score multiplier, defaulting to head\_dim\*\*-0.5.

---

**`reference_rounding`** `bool` — default: False

Select released-inference forward exponent arithmetic.

---

**Returns:** `torch.Tensor`

Independent BF16 tensor \[batch, sequence, heads, head\_dim], with the

```python
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 | None` — default: None

Score multiplier, defaulting to head\_dim\*\*-0.5.

---

**`reference_rounding`** `bool` — default: 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].