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

FlexAttention execution of Qwen3.8-Flash-Next token-indexed sparse GQA.

One code path serves every training layout: dense right-padded batches,
packed (THD) rows, and context parallelism (local queries against gathered
global K/V) all reduce to "each query row attends exactly to its route IDs".
The routes are scattered into a boolean membership table, FlexAttention's
BlockMask skips fully-masked 128x128 tiles, and the kernel avoids materializing
dense attention scores. Rows whose routes are all `-1` (padding queries)
produce exactly zero output and zero gradients.

## Module Contents

### Functions

| Name                                                                                                                   | Description                                                |
| ---------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------- |
| [`_compiled_flex`](#nemo_automodel-components-models-qwen3_8_flash_next-flex_qsa-_compiled_flex)                       | Compile lazily so CPU-only imports never trigger inductor. |
| [`_routes_to_membership`](#nemo_automodel-components-models-qwen3_8_flash_next-flex_qsa-_routes_to_membership)         | Scatter route IDs into a boolean membership table.         |
| [`flex_sparse_gqa_attention`](#nemo_automodel-components-models-qwen3_8_flash_next-flex_qsa-flex_sparse_gqa_attention) | Run route-sparse GQA through FlexAttention.                |

### Data

[`__all__`](#nemo_automodel-components-models-qwen3_8_flash_next-flex_qsa-__all__)

### API

```python
nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa._compiled_flex()
```

Compile lazily so CPU-only imports never trigger inductor.

```python
nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa._routes_to_membership(
    selected_token_ids: torch.Tensor,
    kv_length: int
) -> tuple[torch.Tensor, torch.Tensor]
```

Scatter route IDs into a boolean membership table.

**Parameters:**

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

Global route IDs `[B, S_q, K]`; negative or
out-of-range entries are padding.

---

**`kv_length`** `int`

Number of physical K/V rows.

---

**Returns:** `torch.Tensor`

Kernel-safe boolean membership `[B, S_q, kv_length]` and a boolean

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

Run route-sparse GQA through FlexAttention.

**Parameters:**

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

BF16 CUDA queries `[B, S_q, Hq, D]`.

---

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

BF16 CUDA keys `[B, S_kv, Hkv, D]`. `S_kv` may differ from
`S_q`; under context parallelism it is the gathered global
length.

---

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

BF16 CUDA values `[B, S_kv, Hkv, D]`.

---

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

int32/int64 route IDs `[B, S_q, K]` in global
K/V coordinates; `-1` and out-of-range entries are padding.

---

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

Positive QK scale, defaulting to `1 / sqrt(D)`.

---

**Returns:** `torch.Tensor`

BF16 attention output `[B, S_q, Hq, D]`. Padding-query rows are

```python
nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa.__all__ = ['flex_sparse_gqa_attention']
```