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

### Classes

| Name                                                                                       | Description                                              |
| ------------------------------------------------------------------------------------------ | -------------------------------------------------------- |
| [`FlexQSAMask`](#nemo_automodel-components-models-qwen3_8_flash_next-flex_qsa-FlexQSAMask) | Route-derived FlexAttention mask for one QSA layer call. |

### Functions

| Name                                                                                                                       | Description                                                                          |
| -------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------ |
| [`_block_mask_from_membership`](#nemo_automodel-components-models-qwen3_8_flash_next-flex_qsa-_block_mask_from_membership) | Build the FlexAttention `BlockMask` directly from the dense membership table.        |
| [`_compiled_flex`](#nemo_automodel-components-models-qwen3_8_flash_next-flex_qsa-_compiled_flex)                           | Compile lazily so CPU-only imports never trigger inductor.                           |
| [`_dense_blocks_to_ordered`](#nemo_automodel-components-models-qwen3_8_flash_next-flex_qsa-_dense_blocks_to_ordered)       | Per-row block counts and column order, matching `torch.nn.attention.flex_attention`. |
| [`_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.                                   |
| [`build_flex_qsa_mask`](#nemo_automodel-components-models-qwen3_8_flash_next-flex_qsa-build_flex_qsa_mask)                 | Build the FlexAttention mask for a set of routed token IDs.                          |
| [`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

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

### API

```python
class nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa.FlexQSAMask(
    block_mask: torch.nn.attention.flex_attention.BlockMask,
    has_routes: torch.Tensor,
    query_length: int,
    kv_length: int
)
```

Dataclass

Route-derived FlexAttention mask for one QSA layer call.

**`block_mask`** `BlockMask`

`BlockMask` over `[batch, query_length, kv_length]` built from the
routed token IDs and shared by all heads.

---

**`has_routes`** `Tensor`

Boolean tensor of shape `[batch, query_length]` marking query rows with
at least one real route; rows without routes are zeroed after attention.

---

**`kv_length`** `int`

Number of physical K/V rows the mask was built for.

---

**`query_length`** `int`

Number of query rows the mask was built for.

---

```python
nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa._block_mask_from_membership(
    membership: torch.Tensor,
    mask_mod,
    block_size: int = _FLEX_BLOCK_SIZE
) -> torch.nn.attention.flex_attention.BlockMask
```

Build the FlexAttention `BlockMask` directly from the dense membership table.

`create_block_mask` would evaluate `mask_mod` through vmap at every
`[B, S_q, KV]` coordinate only to recover the membership table this module
already holds, then reduce it to block occupancy. The vmap evaluation is a
long chain of small host-side ops (about 8 ms per QSA layer on GB200);
reducing the existing table costs a handful of launches and yields the same
block structure. `mask_mod` is still attached so the kernel can resolve
partial blocks element-wise.

**Parameters:**

**`membership`** `torch.Tensor`

Boolean `[B, S_q, KV]` table; `True` where a query may
attend a K/V row. `S_q` and `KV` need not be block multiples.

---

**`mask_mod`**

Element-wise FlexAttention mask function consistent with
`membership` inside `[0, S_q) x [0, KV)` and `False` outside.

---

**`block_size`** `int` — default: \_FLEX\_BLOCK\_SIZE

Square sparse block size of the BlockMask.

---

**Returns:** `BlockMask`

A `BlockMask` over `[B, S_q, KV]` with heads broadcast (`H = 1`),

```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._dense_blocks_to_ordered(
    dense_blocks: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Per-row block counts and column order, matching `torch.nn.attention.flex_attention`.

**Parameters:**

**`dense_blocks`** `torch.Tensor`

Boolean block occupancy `[B, H, num_q_blocks, num_kv_blocks]`.

---

**Returns:** `torch.Tensor`

int32 counts `[B, H, num_q_blocks]` and int32 kv-block indices

```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.build_flex_qsa_mask(
    selected_token_ids: torch.Tensor,
    kv_length: int,
    device: torch.device
) -> nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa.FlexQSAMask
```

Build the FlexAttention mask for a set of routed token IDs.

Building the mask is the CPU-heavy part of the flex QSA path (a scatter into a
`[B, S_q, kv_length]` table followed by `create_block_mask`). Callers that
evaluate the same routes twice, checkpoint forward and recompute, build it once
and hand it to `flex_sparse_gqa_attention`.

**Parameters:**

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

---

**`kv_length`** `int`

Number of physical K/V rows.

---

**`device`** `torch.device`

CUDA device the attention will run on.

---

**Returns:** `FlexQSAMask`

The block mask, the `[B, S_q]` has-routes marker, and the extents it was built for.

```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,
    mask: nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa.FlexQSAMask | 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)`.

---

**`mask`** `FlexQSAMask | None` — default: None

Optional mask previously built from `selected_token_ids` with
`build_flex_qsa_mask`; built here when omitted.

---

**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._FLEX_BLOCK_SIZE = 128
```