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

Qwen3.8-Flash-Next QSA routing, FlexAttention sparse GQA, and its PyTorch oracle.

The supplied Qwen3.8-Flash-Next reference is an inference implementation: its indexer
returns integer top-k IDs and defines neither an auxiliary indexer loss nor a
straight-through gradient.  This module therefore freezes the indexer weights
explicitly.  Routing is still recomputed from the current hidden states on
every forward, while gradients flow through the main attention Q/K/V path.

## Module Contents

### Classes

| Name                                                                                                                  | Description                                                   |
| --------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------- |
| [`QSARouteSelection`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-QSARouteSelection)                     | Routes chosen by the frozen indexer for one QSA layer call.   |
| [`Qwen3_8_FlashNextQSAIndexer`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-Qwen3_8_FlashNextQSAIndexer) | Frozen, hookable Qwen3.8-Flash-Next compressed-block indexer. |

### Functions

| Name                                                                                                                              | Description                                                                           |
| --------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------- |
| [`_gathered_qsa_gqa_attention_chunk`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-_gathered_qsa_gqa_attention_chunk) | Evaluate one bounded-workspace chunk of the PyTorch QSA oracle.                       |
| [`_qsa_query_chunk_rows`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-_qsa_query_chunk_rows)                         | Rows scored per chunk: the explicit setting, or as many as fit the fp32 score budget. |
| [`apply_qsa_rope`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-apply_qsa_rope)                                       | Apply the Qwen3.8-Flash-Next attention RoPE to indexer states.                        |
| [`gathered_qsa_gqa_attention`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-gathered_qsa_gqa_attention)               | Run the differentiable PyTorch QSA oracle without expanding K/V heads.                |
| [`qsa_gqa_attention`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-qsa_gqa_attention)                                 | Dispatch QSA to the selected CUDA kernel or the PyTorch CPU oracle.                   |
| [`qsa_route_selection_region`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-qsa_route_selection_region)               | Run route selection outside any active TorchDispatch mode.                            |
| [`right_padded_sequence_lengths`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-right_padded_sequence_lengths)         | Validate a non-packed right-tail mask and return logical lengths.                     |
| [`select_qsa_token_ids`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-select_qsa_token_ids)                           | Score compressed blocks and expand gold QSA top-k IDs.                                |

### Data

[`_PYTORCH_ORACLE_QUERY_CHUNK_SIZE`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-_PYTORCH_ORACLE_QUERY_CHUNK_SIZE)

[`_QSA_SCORE_BUDGET_BYTES`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-_QSA_SCORE_BUDGET_BYTES)

[`qsa_route_replay`](#nemo_automodel-components-models-qwen3_8_flash_next-qsa-qsa_route_replay)

### API

```python
class nemo_automodel.components.models.qwen3_8_flash_next.qsa.QSARouteSelection(
    selected_token_ids: torch.Tensor,
    flex_mask: nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa.FlexQSAMask | None
)
```

Dataclass

Routes chosen by the frozen indexer for one QSA layer call.

**`flex_mask`** `FlexQSAMask | None`

FlexAttention mask built from the same routes when the layer runs
on CUDA with the `flex` backend, otherwise `None`.

---

**`selected_token_ids`** `Tensor`

int32 route IDs `[B, S_q, attention_width]` in global
K/V coordinates; `-1` marks padding.

---

```python
class nemo_automodel.components.models.qwen3_8_flash_next.qsa.Qwen3_8_FlashNextQSAIndexer(
    config: object,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

Frozen, hookable Qwen3.8-Flash-Next compressed-block indexer.

`forward` has no cache or mutable routing state and returns the complete
logical-ID tensor `[B, S, indexer_budget + compress_ratio - 1]`.  A normal
PyTorch forward hook can therefore capture the exact routing artifact used
by the subsequent sparse attention.

The fused projection produces raw query/key layout
`[B, S, (H_index + 1) * D_index]`.  Queries become
`[B, S, H_index, D_index]`.  Raw keys become `[B, S, 1, D_index]` and
complete consecutive groups are averaged in FP32 into
`[B, floor(S / c), 1, D_index]` before K RMSNorm and group-start RoPE.

**`compress_ratio`** `= int(getattr(config, 'indexer_compress_ratio'))`

---

**`head_dim`** `= int(getattr(config, 'indexer_head_dim'))`

---

**`hidden_size`** `= int(getattr(config, 'hidden_size'))`

---

**`index_qk_proj`**

---

**`k_layernorm`** `= Qwen3NextRMSNorm(self.head_dim, eps=eps)`

---

**`num_key_heads`** `= int(getattr(config, 'indexer_kv_heads'))`

---

**`num_query_heads`** `= int(getattr(config, 'indexer_n_heads'))`

---

**`q_layernorm`** `= Qwen3NextRMSNorm(self.head_dim, eps=eps)`

---

**`query_chunk_size`**

---

**`token_budget`** `= int(getattr(config, 'indexer_budget'))`

---

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa.Qwen3_8_FlashNextQSAIndexer._forward_packed(
    hidden_states: torch.Tensor,
    freqs_cis: torch.Tensor,
    cu_seqlens: torch.Tensor,
    cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None = None
) -> torch.Tensor
```

Select document-confined routes for one packed THD row.

Each document is unpacked into its own right-padded pseudo-batch row so
the per-row selection runs unchanged: compression groups restart at
every document start and visibility is bounded by the document length.
Local per-document IDs are then offset by the document start and packed
back into the flattened layout.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Packed decoder block input `[1, T, hidden_size]`.

---

**`freqs_cis`** `torch.Tensor`

Document-relative rotary values `[1, T, D_rope]`.

---

**`cu_seqlens`** `torch.Tensor`

Strictly increasing boundaries `[num_docs + 1]`
starting at zero and ending at `T`.

---

**Returns:** `torch.Tensor`

int32 global flattened IDs \`\`\[1, T, token\_budget +

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa.Qwen3_8_FlashNextQSAIndexer._forward_packed_cp(
    hidden_states: torch.Tensor,
    freqs_cis: torch.Tensor,
    boundaries: torch.Tensor,
    cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext
) -> torch.Tensor
```

Select document-confined routes for one rank's packed query shard.

The frozen indexer needs no gradients, so raw index keys and rotary
values are gathered globally with plain collectives and every document
is compressed identically on all ranks. Documents intersecting the
local shard are scored one at a time with the existing per-row
selection; queries in the CP padding tail (beyond the final document
boundary) keep all `-1` routes.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Local packed shard `[1, S_local, hidden_size]`.

---

**`freqs_cis`** `torch.Tensor`

Local document-relative rotary values `[1, S_local, D_rope]`.

---

**`boundaries`** `torch.Tensor`

Validated global boundaries `[num_docs + 1]`.

---

**`cp_context`** `Qwen3_8_FlashNextCPContext`

Contiguous CP metadata for the packed global row.

---

**Returns:** `torch.Tensor`

int32 global flattened IDs \`\`\[1, S\_local, token\_budget +

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa.Qwen3_8_FlashNextQSAIndexer.forward(
    hidden_states: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None = None,
    cu_seqlens: torch.Tensor | None = None
) -> torch.Tensor
```

Return selected logical token IDs for every physical query row.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Decoder block input `[B, S, hidden_size]`.

---

**`freqs_cis`** `torch.Tensor`

Composed Qwen3.8-Flash-Next rotary values `[B, S, D_rope]` stored
as concatenated cosine/sine halves.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional binary right-tail mask `[B, S]`.

---

**`cp_context`** `Qwen3_8_FlashNextCPContext | None` — default: None

Optional contiguous CP metadata. When present,
`hidden_states` and `freqs_cis` contain the local query
shard while raw logical lengths and selected IDs use global
sequence coordinates.

---

**`cu_seqlens`** `torch.Tensor | None` — default: None

Optional packed-document boundaries `[num_docs + 1]`
for a THD batch of one row. Routes then stay confined to each
query's document and use global flattened token IDs.

---

**Returns:** `torch.Tensor`

int32 logical IDs `[B, S, token_budget + compress_ratio - 1]`.

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa.Qwen3_8_FlashNextQSAIndexer.init_weights(
    init_std: float = 0.02
) -> None
```

Initialize frozen indexer parameters for scratch-model construction.

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa._gathered_qsa_gqa_attention_chunk(
    grouped_query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    selected_token_ids: torch.Tensor,
    softmax_scale: float
) -> torch.Tensor
```

Evaluate one bounded-workspace chunk of the PyTorch QSA oracle.

`grouped_query` uses `[B, Q, Hkv, G, D]`; K/V remain
`[B, S_global, Hkv, D]` and IDs use `[B, Q, K]`.

**Parameters:**

**`grouped_query`** `torch.Tensor`

Local grouped queries `[B, Q, Hkv, G, D]`.

---

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

Global keys `[B, S_global, Hkv, D]`.

---

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

Global values `[B, S_global, Hkv, D]`.

---

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

Global logical IDs `[B, Q, K]`; `-1` marks
invalid fixed-width slots.

---

**`softmax_scale`** `float`

Score multiplier.

---

**Returns:** `torch.Tensor`

Chunk output `[B, Q, Hq, D]` in the query dtype.

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa._qsa_query_chunk_rows(
    query_chunk_size: int | None,
    num_heads: int,
    num_blocks: int,
    rows: int
) -> int
```

Rows scored per chunk: the explicit setting, or as many as fit the fp32 score budget.

**Parameters:**

**`query_chunk_size`** `int | None`

Explicit maximum rows per chunk, or `None` for the memory-based default.

---

**`num_heads`** `int`

Index query heads (`H_index`); the score tensor is `[rows, H_index, num_blocks]` fp32.

---

**`num_blocks`** `int`

Compressed key blocks visible to this batch row.

---

**`rows`** `int`

Number of valid query rows to cover.

---

**Returns:** `int`

A positive chunk size no larger than `rows`.

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa.apply_qsa_rope(
    states: torch.Tensor,
    freqs_cis: torch.Tensor
) -> torch.Tensor
```

Apply the Qwen3.8-Flash-Next attention RoPE to indexer states.

**Parameters:**

**`states`** `torch.Tensor`

Index query or compressed-key states `[B, N, H, D_index]`.

---

**`freqs_cis`** `torch.Tensor`

Model-composed rotary values `[B, N, D_rope]` laid out as
`cat(cos[..., D_rope/2], sin[..., D_rope/2])`.  `D_rope` may be
smaller than `D_index`; the remaining index dimensions pass
through unchanged.

---

**Returns:** `torch.Tensor`

Rotated states in the same `[B, N, H, D_index]` layout and dtype.

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

Run the differentiable PyTorch QSA oracle without expanding K/V heads.

This implementation is retained for CPU execution and numerical parity.
CUDA training with `backend.attn='flex'` dispatches to FlexAttention
instead. The oracle uses a private fixed query chunk solely to bound
temporary gathered K/V storage; it has no public model configuration.

**Parameters:**

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

Main normalized/rotated local queries `[B, S_query, Hq, D]`.

---

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

Main normalized/rotated global keys `[B, S_global, Hkv, D]`.

---

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

Main global values `[B, S_global, Hkv, D]`.

---

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

Indexer output `[B, S_query, K]`. IDs are global
logical positions in the same batch row; `-1` marks fixed-width
padding.

---

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

Score multiplier, defaulting to `1 / sqrt(D)`.

---

**Returns:** `torch.Tensor`

Sparse attention output `[B, S_query, Hq, D]`. Rows with no selected

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa.qsa_gqa_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    selected_token_ids: torch.Tensor,
    backend: str,
    softmax_scale: float | None = None,
    flex_mask: nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa.FlexQSAMask | None = None
) -> torch.Tensor
```

Dispatch QSA to the selected CUDA kernel or the PyTorch CPU oracle.

CPU execution always uses the oracle so model construction, checkpoint
inspection, and distributed CPU parity tests need no compiled kernels.
CUDA execution is strict: unsupported backends or dtypes are reported
rather than silently falling back to the gathered implementation.

**Parameters:**

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

Tensor \[batch, local\_queries, query\_heads, head\_dim].

---

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

Tensor \[batch, global\_keys, kv\_heads, head\_dim].

---

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

Tensor with key's layout. Q/K/V share one dtype and device.

---

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

Signed IDs \[batch, local\_queries, routes] in global
K/V coordinates; invalid CUDA IDs are padding and duplicates collapse.

---

**`backend`** `str`

CUDA backend, "flex" or "cute". FA4 requires SM90 BF16 D256.

---

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

Optional positive QK score multiplier.

---

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

Optional FlexAttention mask already built from `selected_token_ids`
(see `build_flex_qsa_mask`); used by the flex backend only.

---

**Returns:** `torch.Tensor`

Independent tensor \[batch, local\_queries, query\_heads, head\_dim] with

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa.qsa_route_selection_region() -> collections.abc.Iterator[None]
```

Run route selection outside any active TorchDispatch mode.

Selective activation checkpointing records every op of the checkpoint forward
through a TorchDispatch mode and expects the recompute to issue the same op
sequence. Replaying recorded routes skips the indexer's ops during recompute, so
those ops must be invisible to the mode during forward as well. The indexer runs
under `torch.no_grad` and the mask is integer-only, so nothing here is needed by
autograd; hiding the region only affects checkpoint bookkeeping.

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa.right_padded_sequence_lengths(
    attention_mask: torch.Tensor | None,
    batch_size: int,
    sequence_length: int,
    device: torch.device
) -> torch.Tensor
```

Validate a non-packed right-tail mask and return logical lengths.

**Parameters:**

**`attention_mask`** `torch.Tensor | None`

`None` or a binary tensor `[B, S]`.  Each row must
be exactly `1 ** L + 0 ** (S - L)`; left/interior padding and
packed document-ID masks are rejected.

---

**`batch_size`** `int`

Expected batch dimension `B`.

---

**`sequence_length`** `int`

Expected physical sequence dimension `S`.

---

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

Device on which to return the lengths.

---

**Returns:** `torch.Tensor`

Logical sequence lengths as int64 `[B]`.

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa.select_qsa_token_ids(
    index_queries: torch.Tensor,
    compressed_keys: torch.Tensor,
    sequence_lengths: torch.Tensor,
    token_budget: int,
    compress_ratio: int,
    query_chunk_size: int | None = None,
    query_position_offset: int = 0,
    global_sequence_length: int | None = None
) -> torch.Tensor
```

Score compressed blocks and expand gold QSA top-k IDs.

For query position `t`, only `floor((t + 1) / compress_ratio)` complete
blocks are visible.  Each block score is
`sum_h(relu(dot(q[t,h], k[block,0]))) / sqrt(D)`.  The best
`token_budget / compress_ratio` blocks are expanded to token IDs, then
the 0--`compress_ratio - 1` tokens in the current incomplete causal tail
are appended.  Invalid slots are `-1`.

The selection is vectorized over query rows: the only host synchronization
is one read of `sequence_lengths`; chunk boundaries, the first sparse row,
and the tail slots follow from positions known on the host, so no
data-dependent branch waits on the device.

**Parameters:**

**`index_queries`** `torch.Tensor`

Normalized and rotated local index queries
`[B, S_query, H_index, D_index]`.

---

**`compressed_keys`** `torch.Tensor`

FP32-mean-pooled, normalized, rotated index keys
`[B, floor(S_global / compress_ratio), 1, D_index]`.

---

**`sequence_lengths`** `torch.Tensor`

Right-padded logical lengths `[B]`.

---

**`token_budget`** `int`

Maximum number of tokens contributed by complete blocks.

---

**`compress_ratio`** `int`

Number of consecutive tokens represented by one block.

---

**`query_chunk_size`** `int | None` — default: None

Maximum query rows scored together. `None` sizes
chunks so the temporary FP32 score tensor `[rows, H_index,
blocks]` stays within 256 MiB (one chunk for 4k-16k sequences).
Chunking bounds memory without changing top-k semantics.

---

**`query_position_offset`** `int` — default: 0

Global position represented by local query row
zero. It is zero without CP and `cp_rank * S_query` for the
contiguous CP layout.

---

**`global_sequence_length`** `int | None` — default: None

Padded global physical sequence length. It
defaults to the local query length for the non-CP path.

---

**Returns:** `torch.Tensor`

Global logical token IDs \`\`\[B, S\_query, token\_budget +

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa._PYTORCH_ORACLE_QUERY_CHUNK_SIZE = 16
```

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa._QSA_SCORE_BUDGET_BYTES = 256 * 1024 * 1024
```

```python
nemo_automodel.components.models.qwen3_8_flash_next.qsa.qsa_route_replay: RecomputeReplay[QSARouteSelection] = RecomputeReplay('QSA route')
```