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

Functions

Name	Description
`_compiled_flex`	Compile lazily so CPU-only imports never trigger inductor.
`_routes_to_membership`	Scatter route IDs into a boolean membership table.
`flex_sparse_gqa_attention`	Run route-sparse GQA through FlexAttention.

Data

__all__

API

nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa._compiled_flex()

Compile lazily so CPU-only imports never trigger inductor.

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

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

Positive QK scale, defaulting to 1 / sqrt(D).

Returns: torch.Tensor

BF16 attention output [B, S_q, Hq, D]. Padding-query rows are

nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa.__all__ = ['flex_sparse_gqa_attention']