nemo_automodel.components.models.deepseek_v4.dspark_target

Packed DSpark target forward for DeepSeek V4.

Module Contents

Functions

Name	Description
`dspark_packed_forward`	Run a frozen DeepSeek V4 forward on packed rows for DSpark feature capture.

API

nemo_automodel.components.models.deepseek_v4.dspark_target.dspark_packed_forward(
    model: torch.nn.Module,
    input_ids: torch.Tensor,
    seq_lens: torch.Tensor
) -> nemo_automodel.components.models.deepseek_v41.packing.PackedSequenceLayout

Run a frozen DeepSeek V4 forward on packed rows for DSpark feature capture.

The THD path derives document IDs from its seq_lens, so sliding-window attention, compressed KV, and the indexer stay inside each document. A document pools the same compressed groups as it does alone only if it starts on a compression-group boundary, so every document is moved to one. The alignment padding runs as its own THD segments, which keeps a legal key in every query row. With the HybridEP dispatcher every rank pads to the longest aligned row of all ranks. DSpark reads hooked activations, so only the last position’s logits are projected.

Parameters:

model

nn.Module

DeepseekV4ForCausalLM target.

input_ids

torch.Tensor

Token IDs of shape [batch, sequence] holding packed documents.

seq_lens

torch.Tensor

Document lengths of shape [batch, documents]; each row sums to sequence and zero entries pad the document axis.

Returns: PackedSequenceLayout

Layout whose restore maps a tensor of shape [batch, aligned_sequence, …]