> 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.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/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.qwen3_8_flash_next.cudnn_qsa

cuDNN Frontend indexed QSA training with the model's existing routes.

## Module Contents

### Classes

| Name                                                                                    | Description                                                       |
| --------------------------------------------------------------------------------------- | ----------------------------------------------------------------- |
| [`_CudnnQSA`](#nemo_automodel-components-models-qwen3_8_flash_next-cudnn_qsa-_CudnnQSA) | Own saved FP32 output and LSE for one independent GQA head group. |

### Functions

| Name                                                                                                                      | Description                                                                     |
| ------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- |
| [`_load_cudnn`](#nemo_automodel-components-models-qwen3_8_flash_next-cudnn_qsa-_load_cudnn)                               | Load optional public entry points lazily; retain no tensor state.               |
| [`_pack_indices`](#nemo_automodel-components-models-qwen3_8_flash_next-cudnn_qsa-_pack_indices)                           | Translate batch-local QSA IDs to packed coordinates without changing selection. |
| [`cudnn_sparse_gqa_attention`](#nemo_automodel-components-models-qwen3_8_flash_next-cudnn_qsa-cudnn_sparse_gqa_attention) | Evaluate model-produced four-token QSA routes with cuDNN on Blackwell.          |

### API

```python
class nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa._CudnnQSA()
```

**Bases:** `Function`

Own saved FP32 output and LSE for one independent GQA head group.

```python
nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa._CudnnQSA.backward(
    ctx: typing.Any,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor | None, ...]
```

staticmethod

Return BF16 dQ \[T,H,256], dK/dV \[T,1,256], and two None entries.

**Parameters:**

**`ctx`** `Any`

Context holding immutable forward inputs and saved state.

---

**`grad_output`** `torch.Tensor`

BF16 CUDA upstream gradient \[T,H,256]; may be strided.

---

**Returns:** `torch.Tensor | None`

Gradients for the five forward arguments. Discrete indices and

```python
nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa._CudnnQSA.forward(
    ctx: typing.Any,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    indices: torch.Tensor,
    scale: float
) -> torch.Tensor
```

staticmethod

Run forward on contiguous packed Q \[T,H,256], K/V \[T,1,256].

**Parameters:**

**`ctx`** `Any`

PyTorch autograd context.

---

**`q`** `torch.Tensor`

BF16 CUDA query \[T,12,256].

---

**`k`** `torch.Tensor`

BF16 CUDA keys \[T,1,256].

---

**`v`** `torch.Tensor`

BF16 CUDA values \[T,1,256].

---

**`indices`** `torch.Tensor`

CUDA int32 selection \[T,W], in packed token coordinates.

---

**`scale`** `float`

Finite positive score multiplier.

---

**Returns:** `torch.Tensor`

Independent BF16 CUDA output \[T,H,256]. Saves native FP32 output

```python
nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa._load_cudnn() -> tuple[collections.abc.Callable, collections.abc.Callable]
```

Load optional public entry points lazily; retain no tensor state.

```python
nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa._pack_indices(
    routes: torch.Tensor,
    sequence_length: int
) -> torch.Tensor
```

Translate batch-local QSA IDs to packed coordinates without changing selection.

**Parameters:**

**`routes`** `torch.Tensor`

Signed integer tensor \[batch, sequence, width], produced by the
model's compression-ratio-four indexer. Four-slot groups contain
distinct complete blocks or a consecutive causal tail and -1 padding.

---

**`sequence_length`** `int`

Physical sequence extent shared by query and K/V.

---

**Returns:** `torch.Tensor`

Independent contiguous int32 tensor \[batch \* sequence, ceil(width/4)\*4].

```python
nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa.cudnn_sparse_gqa_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    selected_token_ids: torch.Tensor,
    softmax_scale: float | None = None
) -> torch.Tensor
```

Evaluate model-produced four-token QSA routes with cuDNN on Blackwell.

Selection is unchanged: this adapter does not run the cuDNN indexer. Each
four-slot group must be a distinct complete block or a consecutive causal
tail padded with -1, as emitted by the model's ratio-four indexer. Blocks
restart at each packed document origin. Arbitrary/duplicate token sets are
not supported by this backend. These device-value invariants are the
indexer's contract, not host-synchronized checks in the attention hot path.

CP with local Q and gathered global K/V is unsupported. First-order
backward uses cuDNN FP32 atomic accumulation and is not bitwise deterministic.

**Parameters:**

**`query`** `torch.Tensor`

BF16 CUDA tensor \[batch, sequence, query_heads, 256].

---

**`key`** `torch.Tensor`

BF16 CUDA tensor \[batch, sequence, kv_heads, 256]. Supported
(query_heads, kv_heads) are (24,2) and (12,1), matching the published
cuDNN Frontend 1.32.1 indexed API.

---

**`value`** `torch.Tensor`

BF16 CUDA tensor with key's shape and device. Strided Q/K/V
are accepted and copied to contiguous per-KV-head packed buffers.

---

**`selected_token_ids`** `torch.Tensor`

Int32/int64 CUDA tensor \[batch, sequence, width],
with 1 \<= width \<= 2052. IDs address the physical sequence within
each batch row; invalid IDs are padding. Packed documents are
supported through their existing absolute-within-row route IDs.

---

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

Finite positive host scalar; defaults to 1/sqrt(256).

---

**Returns:** `torch.Tensor`

Independent BF16 CUDA output \[batch, sequence, query_heads, 256].

**Raises:**

* `ValueError`: For unsupported shapes, devices, heads, scale or CP layout.
* `TypeError`: For non-BF16 compute tensors or unsigned/noninteger routes.
* `RuntimeError`: For unsupported GPU or deterministic backward requests.
* `ImportError`: When the optional cuDNN Frontend QSA runtime is absent.