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

Classes

Name	Description
`FlexQSAMask`	Route-derived FlexAttention mask for one QSA layer call.

Functions

Name	Description
`_block_mask_from_membership`	Build the FlexAttention `BlockMask` directly from the dense membership table.
`_compiled_flex`	Compile lazily so CPU-only imports never trigger inductor.
`_dense_blocks_to_ordered`	Per-row block counts and column order, matching `torch.nn.attention.flex_attention`.
`_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.
`build_flex_qsa_mask`	Build the FlexAttention mask for a set of routed token IDs.
`flex_sparse_gqa_attention`	Run route-sparse GQA through FlexAttention.

Data

_FLEX_BLOCK_SIZE

API

class nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa.FlexQSAMask(
    block_mask: torch.nn.attention.flex_attention.BlockMask,
    has_routes: torch.Tensor,
    query_length: int,
    kv_length: int
)

Dataclass

Route-derived FlexAttention mask for one QSA layer call.

block_mask

BlockMask

BlockMask over [batch, query_length, kv_length] built from the routed token IDs and shared by all heads.

has_routes

Tensor

Boolean tensor of shape [batch, query_length] marking query rows with at least one real route; rows without routes are zeroed after attention.

kv_length

int

Number of physical K/V rows the mask was built for.

query_length

int

Number of query rows the mask was built for.

nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa._block_mask_from_membership(
    membership: torch.Tensor,
    mask_mod,
    block_size: int = _FLEX_BLOCK_SIZE
) -> torch.nn.attention.flex_attention.BlockMask

Build the FlexAttention BlockMask directly from the dense membership table.

create_block_mask would evaluate mask_mod through vmap at every [B, S_q, KV] coordinate only to recover the membership table this module already holds, then reduce it to block occupancy. The vmap evaluation is a long chain of small host-side ops (about 8 ms per QSA layer on GB200); reducing the existing table costs a handful of launches and yields the same block structure. mask_mod is still attached so the kernel can resolve partial blocks element-wise.

Parameters:

membership

torch.Tensor

Boolean [B, S_q, KV] table; True where a query may attend a K/V row. S_q and KV need not be block multiples.

mask_mod

Element-wise FlexAttention mask function consistent with membership inside [0, S_q) x [0, KV) and False outside.

block_size

intDefaults to _FLEX_BLOCK_SIZE

Square sparse block size of the BlockMask.

Returns: BlockMask

A BlockMask over [B, S_q, KV] with heads broadcast (H = 1),

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._dense_blocks_to_ordered(
    dense_blocks: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Per-row block counts and column order, matching torch.nn.attention.flex_attention.

Parameters:

dense_blocks

torch.Tensor

Boolean block occupancy [B, H, num_q_blocks, num_kv_blocks].

Returns: torch.Tensor

int32 counts [B, H, num_q_blocks] and int32 kv-block indices

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.build_flex_qsa_mask(
    selected_token_ids: torch.Tensor,
    kv_length: int,
    device: torch.device
) -> nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa.FlexQSAMask

Build the FlexAttention mask for a set of routed token IDs.

Building the mask is the CPU-heavy part of the flex QSA path (a scatter into a [B, S_q, kv_length] table followed by create_block_mask). Callers that evaluate the same routes twice, checkpoint forward and recompute, build it once and hand it to flex_sparse_gqa_attention.

Parameters:

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.

kv_length

int

Number of physical K/V rows.

device

torch.device

CUDA device the attention will run on.

Returns: FlexQSAMask

The block mask, the [B, S_q] has-routes marker, and the extents it was built for.

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,
    mask: nemo_automodel.components.models.qwen3_8_flash_next.flex_qsa.FlexQSAMask | 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).

mask

FlexQSAMask | NoneDefaults to None

Optional mask previously built from selected_token_ids with build_flex_qsa_mask; built here when omitted.

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._FLEX_BLOCK_SIZE = 128