nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa

cuDNN Frontend indexed QSA training with the model’s existing routes.

Module Contents

Classes

Name	Description
`_CudnnQSA`	Own saved FP32 output and LSE for one independent GQA head group.

Functions

Name	Description
`_load_cudnn`	Load optional public entry points lazily; retain no tensor state.
`_pack_indices`	Translate batch-local QSA IDs to packed coordinates without changing selection.
`cudnn_sparse_gqa_attention`	Evaluate model-produced four-token QSA routes with cuDNN on Blackwell.

API

class nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa._CudnnQSA()

Bases: Function

Own saved FP32 output and LSE for one independent GQA head group.

nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa._CudnnQSA.backward(
    ctx: typing.Any,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor | None, ...]

staticmethod

Return BF16 dQ [T,H,256], dK/dV [T,1,256], and two None entries.

Parameters:

ctx

Any

Context holding immutable forward inputs and saved state.

grad_output

torch.Tensor

BF16 CUDA upstream gradient [T,H,256]; may be strided.

Returns: torch.Tensor | None

Gradients for the five forward arguments. Discrete indices and

nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa._CudnnQSA.forward(
    ctx: typing.Any,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    indices: torch.Tensor,
    scale: float
) -> torch.Tensor

staticmethod

Run forward on contiguous packed Q [T,H,256], K/V [T,1,256].

Parameters:

ctx

Any

PyTorch autograd context.

torch.Tensor

BF16 CUDA query [T,12,256].

torch.Tensor

BF16 CUDA keys [T,1,256].

torch.Tensor

BF16 CUDA values [T,1,256].

indices

torch.Tensor

CUDA int32 selection [T,W], in packed token coordinates.

scale

float

Finite positive score multiplier.

Returns: torch.Tensor

Independent BF16 CUDA output [T,H,256]. Saves native FP32 output

nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa._load_cudnn() -> tuple[collections.abc.Callable, collections.abc.Callable]

Load optional public entry points lazily; retain no tensor state.

nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa._pack_indices(
    routes: torch.Tensor,
    sequence_length: int
) -> torch.Tensor

Translate batch-local QSA IDs to packed coordinates without changing selection.

Parameters:

routes

torch.Tensor

Signed integer tensor [batch, sequence, width], produced by the model’s compression-ratio-four indexer. Four-slot groups contain distinct complete blocks or a consecutive causal tail and -1 padding.

sequence_length

int

Physical sequence extent shared by query and K/V.

Returns: torch.Tensor

Independent contiguous int32 tensor [batch * sequence, ceil(width/4)*4].

nemo_automodel.components.models.qwen3_8_flash_next.cudnn_qsa.cudnn_sparse_gqa_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    selected_token_ids: torch.Tensor,
    softmax_scale: float | None = None
) -> torch.Tensor

Evaluate model-produced four-token QSA routes with cuDNN on Blackwell.

Selection is unchanged: this adapter does not run the cuDNN indexer. Each four-slot group must be a distinct complete block or a consecutive causal tail padded with -1, as emitted by the model’s ratio-four indexer. Blocks restart at each packed document origin. Arbitrary/duplicate token sets are not supported by this backend. These device-value invariants are the indexer’s contract, not host-synchronized checks in the attention hot path.

CP with local Q and gathered global K/V is unsupported. First-order backward uses cuDNN FP32 atomic accumulation and is not bitwise deterministic.

Parameters:

query

torch.Tensor

BF16 CUDA tensor [batch, sequence, query_heads, 256].

key

torch.Tensor

BF16 CUDA tensor [batch, sequence, kv_heads, 256]. Supported (query_heads, kv_heads) are (24,2) and (12,1), matching the published cuDNN Frontend 1.32.1 indexed API.

value

torch.Tensor

BF16 CUDA tensor with key’s shape and device. Strided Q/K/V are accepted and copied to contiguous per-KV-head packed buffers.

selected_token_ids

torch.Tensor

Int32/int64 CUDA tensor [batch, sequence, width], with 1 <= width <= 2052. IDs address the physical sequence within each batch row; invalid IDs are padding. Packed documents are supported through their existing absolute-within-row route IDs.

softmax_scale

float | NoneDefaults to None

Finite positive host scalar; defaults to 1/sqrt(256).

Returns: torch.Tensor

Independent BF16 CUDA output [batch, sequence, query_heads, 256].

Raises:

ValueError: For unsupported shapes, devices, heads, scale or CP layout.
TypeError: For non-BF16 compute tensors or unsigned/noninteger routes.
RuntimeError: For unsupported GPU or deterministic backward requests.
ImportError: When the optional cuDNN Frontend QSA runtime is absent.