ttnn.experimental.indexer_score_dsa

ttnn.experimental.indexer_score_dsa(q, k, weights, chunk_start_idx, program_config, compute_kernel_config, cache_batch_idx, kv_len, cluster_axis, seq_subshard_axis, block_cyclic_sp_axis, block_cyclic_chunk_local) → None

DeepSeek-V3.2 DSA / GLM-5 lightning-indexer scorer.

score[b, 0, s, t] = sum_h relu(q[b,h,s,:] . k[b,t,:]) * weights[b,h,s]

ReLU(q.kT) gated by the learned per-head weights and summed over ALL Hi index heads into one shared selection row [B, 1, Sq, T]. For MiniMax M3’s raw-dot, per-GQA-group indexer, use indexer_score_msa instead.

Parameters:

  • q – [B, Hi, Sq, D] bf16 or bfp8_b tiled (post non-interleaved RoPE)

  • k – [B, 1, T, D] bf16 or bfp8_b tiled, single shared head

  • weights – [B, Hi, Sq, 1] bf16 tiled learned per-head gates (scale pre-folded)

  • chunk_start_idx – absolute global position of rank 0’s query row 0 (rank 0 = lowest cluster_axis coord; causality: key t visible to query s iff t <= chunk_start + s). OMIT on a mesh -> deduced as T - sp_ring*Sq (sp_ring = mesh extent along cluster_axis, whole mesh if unset). Single device: set to history + rank*Sq per rank.

  • program_config – work-unit knobs (q_chunk_size, k_chunk_size, head_group_size; elements, tile-aligned). Defaults always fit L1; raise head_group_size (0 = all resident) for performance.

  • compute_kernel_config – optional DeviceComputeKernelConfig. Only math_fidelity is honored (default: HiFi2, or LoFi when q and k are both bfloat8_b); fp32_dest_acc_en / dst_full_sync_en must stay false (the custom LLK is validated for bf16 DEST half-sync).

  • cache_batch_idx – optional int. Selects the batch slot of a shared [B, 1, T, D] k cache (which may then be ND-sharded across DRAM banks). Re-applied each dispatch, so switching slots does not recompile.

  • kv_len – optional int. Valid prefix of a k allocated at its full T; the rest is masked out. Tile-aligned, in (0, T], with chunk_start_idx + Sq <= kv_len. Re-applied each dispatch, so a serving loop growing kv_len (<= T) reuses one program – no recompile. Only output columns [0, kv_len) are written.

  • cluster_axis – mesh axis that is the SP ring. On a mesh, device r uses chunk_start = chunk_start_idx + r*Sq, where r is its linearized index along this axis (Sq = q seq len). None = linear device order (1 on a single device, so chunk_start_idx is used as-is).

  • seq_subshard_axis – optional int, the SECOND mesh axis (TP) the query seq is block-cyclically sub-sharded over, on top of the SP block-cyclic layout. Set it alongside a named cluster_axis + block_cyclic when the query rows are further split across TP (block_cyclic_chunk_local == tp*q_isl): the exact geometry adds tp_rank*Sq to each device’s block-cyclic position, so the causal mask stays correct under rotated (mid-slab) starts — unlike the cluster_axis=None flat path, which is linear-approximate there. Requires a block-cyclic layout and an axis distinct from cluster_axis.

  • block_cyclic_sp_axis – optional int. The MESH AXIS the K cache was striped over (chunked prefill + SP all-gather leaves [B,1,T,D] k in per-SP-shard slab / block-cyclic physical order). sp is DERIVED from the mesh shape on that axis, so a caller cannot pass an sp that disagrees with the device. The reader reads the permuted cache back in natural token order (invP per tile) so scores come out token-ordered and the causal mask/pool stay exact. Unset (or sp==1) = contiguous K (no remap). Pair with block_cyclic_chunk_local. Interface matches ttnn.transformer.sparse_sdpa.

  • block_cyclic_chunk_local – optional int, REQUIRED with block_cyclic_sp_axis. The per-shard chunk length (chunk_size_global / sp). Cross-checked against q: must equal q_isl (Sq, seq sharded only on the SP axis) or tp*q_isl (tp = mesh_size/sp). The tp*q_isl case with tp>1 (seq sharded across BOTH axes) has two forms: with cluster_axis=None it uses flat row-major linearization over all devices (linear-approximate under a mid-slab start); with a named cluster_axis it must also pass seq_subshard_axis naming the TP axis, which gives the rotation-exact 2D geometry (see seq_subshard_axis above).

Returns: score [B, 1, Sq, T] bf16 row-major; future/pad columns -inf.