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.
`_membership_flat_offset`	Flat offset into a `[B, S_q, kv_length]` membership table, evaluated in int64.
`_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._membership_flat_offset(
    batch_idx: torch.Tensor,
    query_idx: torch.Tensor,
    kv_idx: torch.Tensor,
    query_length: int,
    kv_length: int
) -> torch.Tensor

Flat offset into a [B, S_q, kv_length] membership table, evaluated in int64.

FlexAttention inlines mask_mod into its Triton template and emits that inlined index arithmetic in int32. The membership table crosses INT32_MAX once B * S_q * kv_length > 2**31 — a square 46341-token sequence is already past it — and from that point query_idx * kv_length wraps negative inside the kernel. The wrapped address still lands in mapped memory for a while, so the tail queries silently read a wrong mask before the failure escalates to CUDA error: an illegal memory access at larger sequence lengths. Widening the operands here keeps the generated address arithmetic in int64.

Parameters:

batch_idx

torch.Tensor

Scalar batch coordinate supplied by FlexAttention.

query_idx

torch.Tensor

Scalar query coordinate, already clamped in range.

kv_idx

torch.Tensor

Scalar key/value coordinate, already clamped in range.

query_length

int

Number of query rows in the membership table.

kv_length

int

Number of physical K/V rows in the membership table.

Returns: torch.Tensor

int64 offset of [batch_idx, query_idx, kv_idx] in the flattened table.

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