> 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.                     |
| [`_membership_flat_offset`](#nemo_automodel-components-models-qwen3_8_flash_next-flex_qsa-_membership_flat_offset)     | Flat offset into a `[B, S_q, kv_length]` membership table, evaluated in int64. |
| [`_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._membership_flat_offset(
    batch_idx: torch.Tensor,
    query_idx: torch.Tensor,
    kv_idx: torch.Tensor,
    query_length: int,
    kv_length: int
) -> torch.Tensor
```

Flat offset into a `[B, S_q, kv_length]` membership table, evaluated in int64.

FlexAttention inlines `mask_mod` into its Triton template and emits that
inlined index arithmetic in int32. The membership table crosses
`INT32_MAX` once `B * S_q * kv_length > 2**31` -- a square 46341-token
sequence is already past it -- and from that point `query_idx * kv_length`
wraps negative inside the kernel. The wrapped address still lands in mapped
memory for a while, so the tail queries silently read a wrong mask before
the failure escalates to `CUDA error: an illegal memory access` at larger
sequence lengths. Widening the operands here keeps the generated address
arithmetic in int64.

**Parameters:**

**`batch_idx`** `torch.Tensor`

Scalar batch coordinate supplied by FlexAttention.

---

**`query_idx`** `torch.Tensor`

Scalar query coordinate, already clamped in range.

---

**`kv_idx`** `torch.Tensor`

Scalar key/value coordinate, already clamped in range.

---

**`query_length`** `int`

Number of query rows in the membership table.

---

**`kv_length`** `int`

Number of physical K/V rows in the membership table.

---

**Returns:** `torch.Tensor`

int64 offset of `[batch_idx, query_idx, kv_idx]` in the flattened table.

```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']
```