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

__all__

API

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

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

intDefaults to 128

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