nemo_automodel.components.models.qwen3_5.packing

Shared packed-sequence metadata for dense and MoE Qwen3.5 models.

Module Contents

Classes

Name	Description
`GatedDeltaPackedMetadata`	Packed-sequence metadata shared by every GatedDeltaNet layer.

Functions

Name	Description
`prepare_gated_delta_packed_metadata`	Build shared GatedDeltaNet metadata once for a model forward.

API

class nemo_automodel.components.models.qwen3_5.packing.GatedDeltaPackedMetadata(
    document_ids: torch.Tensor,
    indices: torch.Tensor,
    cu_seqlens: torch.Tensor,
    cu_seqlens_cpu: torch.Tensor
)

Dataclass

Packed-sequence metadata shared by every GatedDeltaNet layer.

Parameters:

document_ids

torch.Tensor

Indexed document mask of shape [batch, sequence] on the compute device, with zero denoting padding.

indices

torch.Tensor

Flattened valid-token indices of shape [tokens] on the compute device.

cu_seqlens

torch.Tensor

Cumulative document lengths of shape [documents + 1] on the compute device.

cu_seqlens_cpu

torch.Tensor

CPU mirror of cu_seqlens with shape [documents + 1] for FLA host-side chunk planning.

cu_seqlens

Tensor

cu_seqlens_cpu

Tensor

document_ids

Tensor

indices

Tensor

nemo_automodel.components.models.qwen3_5.packing.prepare_gated_delta_packed_metadata(
    attention_mask: torch.Tensor | None,
    packed_seq_ids: torch.Tensor | None
) -> nemo_automodel.components.models.qwen3_5.packing.GatedDeltaPackedMetadata | None

Build shared GatedDeltaNet metadata once for a model forward.

Parameters:

attention_mask

torch.Tensor | None

Optional indexed document mask of shape [batch, sequence] or a backend-specific attention mask.

packed_seq_ids

torch.Tensor | None

Optional indexed document IDs of shape [batch, sequence] supplied beside a backend-specific attention mask.

Returns: GatedDeltaPackedMetadata | None

Device and CPU packed-sequence metadata whose tensor layouts are