> 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                                                   |
| --------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------- |
| [`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.        |
| [`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 FlexAttention on CUDA or the PyTorch oracle elsewhere. |
| [`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)

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

### API

```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.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
) -> torch.Tensor
```

Dispatch QSA to FlexAttention on CUDA or the PyTorch oracle elsewhere.

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.

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

**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` — default: 128

Query rows scored together.  This bounds the temporary
FP32 score tensor 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.__all__ = ['Qwen3_8_FlashNextQSAIndexer', 'apply_qsa_rope', 'gathered_qsa_gqa_attention', ...
```