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`	Routes chosen by the frozen indexer for one QSA layer call.
`Qwen3_8_FlashNextQSAIndexer`	Frozen, hookable Qwen3.8-Flash-Next compressed-block indexer.

Functions

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

Data

_PYTORCH_ORACLE_QUERY_CHUNK_SIZE

_QSA_SCORE_BUDGET_BYTES

qsa_route_replay

API

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.

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'))

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 +

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 +

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 | NoneDefaults to None

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

cp_context

Qwen3_8_FlashNextCPContext | NoneDefaults to 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 | NoneDefaults to 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].

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.

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.

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.

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.

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 | NoneDefaults to None

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

Returns: torch.Tensor

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

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 | NoneDefaults to None

Optional positive QK score multiplier.

flex_mask

FlexQSAMask | NoneDefaults to 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

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.

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

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 | NoneDefaults to 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

intDefaults to 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 | NoneDefaults to 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 +

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

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

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