> 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.tilelang_sparse_mla_fwd

## Module Contents

### Functions

| Name                                                                                                                                 | Description                                    |
| ------------------------------------------------------------------------------------------------------------------------------------ | ---------------------------------------------- |
| [`sparse_mqa_fwd`](#nemo_automodel-components-models-deepseek_v4-kernels-tilelang_sparse_mla_fwd-sparse_mqa_fwd)                     | -                                              |
| [`sparse_mqa_fwd_interface`](#nemo_automodel-components-models-deepseek_v4-kernels-tilelang_sparse_mla_fwd-sparse_mqa_fwd_interface) | Forward interface for V4 sparse MQA attention. |

### API

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

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

**`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].
Entries outside \[0, kv\_sequence) are masked; slots are internally
padded to a multiple of block\_I.

---

**`sm_scale`** `float | None` — default: None

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

---

**`block_I`** `int` — default: 64

Sparse-key slots processed by each kernel iteration.

---

**`num_stages`** `int` — default: 2

Pipeline stages for the generated kernel.

---

**`threads`** `int` — default: 256

CUDA threads per block.

---

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