nemo_automodel.components.models.deepseek_v41.indexer

Fused frozen V4.1 indexer scoring with the released BF16 rounding boundaries.

The tiled GEMM/head-reduction structure follows the vendored V4 TileLang indexer. V4.1 rounds the GEMM and weighted head scores to BF16 before reducing in FP32. Visibility is supplied by the existing attention metadata path, including packed document boundaries and global CP positions; top-k/candidate selection stays there.

Module Contents

Functions

Name	Description
`_indexer_fwd`	Compile a query/key tile; per-head intermediates remain on chip.
`indexer_scores`	Fuse frozen indexer scoring without materializing per-head global scores.

API

nemo_automodel.components.models.deepseek_v41.indexer._indexer_fwd(
    heads: int,
    index_dim: int
)

Compile a query/key tile; per-head intermediates remain on chip.

nemo_automodel.components.models.deepseek_v41.indexer.indexer_scores(
    queries: torch.Tensor,
    keys: torch.Tensor,
    weights: torch.Tensor,
    allowed: torch.Tensor
) -> torch.Tensor

Fuse frozen indexer scoring without materializing per-head global scores.

Parameters:

queries

torch.Tensor

BF16 CUDA queries [batch, local_sequence, heads, index_dim]. Heads must be a power of two no larger than 128; index_dim must be divisible by 16. Quantize/dequantize and RoPE are already applied.

keys

torch.Tensor

BF16 CUDA keys [batch, global_compressed, index_dim], gathered across CP ranks before this call.

weights

torch.Tensor

BF16 CUDA weights [batch, local_sequence, heads], including the released head/dimension scaling.

allowed

torch.Tensor

Boolean CUDA visibility [batch, local_sequence, global_compressed], including global causality, valid compression groups and packed document isolation. All inputs must reside on the same CUDA device.

Returns: torch.Tensor

Independent BF16 scores [batch, local_sequence, global_compressed],