ttnn.experimental.indexer_score_msa

ttnn.experimental.indexer_score_msa(q, k, chunk_start_idx, scale, num_groups, block_size, program_config, compute_kernel_config, cache_batch_idx, kv_len, cluster_axis) → None

MiniMax-M3 MSA lightning-indexer scorer.

score[b, g, s, t] = sum_{h in group g} (q[b,h,s,:] . k[b,t,:]) * scale

Raw dot product (no ReLU), no learned gates – only a 1/sqrt(d) scale (applied as a constant gate). The Hi index heads are partitioned into num_groups GQA groups and summed WITHIN each group only, giving one output plane per group (no cross-group sum) for per-group selection. For DeepSeek-V3.2 / GLM-5’s relu + gated + head-summed indexer, use indexer_score_dsa 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

  • chunk_start_idx – absolute global position of rank 0’s query row 0 (causality: key t visible to query s iff t <= chunk_start + s). OMIT on a mesh -> deduced as T - sp_ring*Sq; single device: set to history + rank*Sq per rank (same semantics as indexer_score_dsa).

  • scale – constant gate applied to every head (e.g. 1/sqrt(d))

  • num_groups – REQUIRED (no default). 1 sums all Hi heads into one plane (the TP=4 group-aligned deployment, Hi=1/device). G>1 partitions the heads into G groups of Hi/G and sums within each group -> output [B, G, Sq, T] (multiple GQA groups on one chip). G>1 needs all heads resident (head_group_size 0 or Hi) and k_chunk_size >= 64.

  • block_size – 0 (default) = no pooling -> score [B, G, Sq, T]. >0 = block-max-pool over block_size keys -> score [B, G, Sq, T/block_size] (block selection; the downstream topk then picks per-group top-k blocks). Requires block_size a multiple of 32, T % block_size == 0, and k_chunk_size % block_size == 0.

  • program_config – work-unit knobs (q_chunk_size, k_chunk_size, head_group_size; elements, tile-aligned). Defaults always fit L1.

  • 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. Same semantics as indexer_score_dsa.

  • 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. When block-max-pooling (block_size > 0) it must also be a multiple of block_size (whole blocks are written). Re-applied each dispatch (a serving loop growing kv_len reuses one program). Only columns/blocks within the valid prefix 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. None = linear device order.

  • block_cyclic_chunk_local (block_cyclic_sp_axis /) – same semantics as indexer_score_dsa. block_cyclic_sp_axis = the MESH AXIS the K cache was striped over (sp derived from the mesh shape on that axis) and block_cyclic_chunk_local = the per-shard chunk length; the reader reads the permuted cache back in natural token order so the per-group scores (and the block-max-pool, which pools token-contiguous blocks) come out correct. Unset = contiguous K.

Returns: score [B, num_groups, Sq, T_out] bf16 row-major (T_out = T, or

T/block_size when block-max-pooling); future/pad columns/blocks -inf.