ttnn.transformer.sparse_sdpa_msa

ttnn.transformer.sparse_sdpa_msa(q: ttnn.Tensor, k: ttnn.Tensor, v: ttnn.Tensor, indices: ttnn.Tensor, *, scale: float | None, block_size: int, cache_batch_idx: int | None, chunk_start_idx: int | None, cluster_axis: int | None) → ttnn.Tensor

MSA block-sparse prefill (MiniMax Sparse Attention), Blackhole single-chip. Attends the block_size-token K/V blocks named in indices; -1 (0xFFFFFFFF) sentinels mask a contiguous tail. Block selection bounds causality to block granularity; pass chunk_start_idx to additionally enforce a token-level causal mask on the diagonal block (required for correct causal prefill). RoPE and QK-norm are applied upstream.

Parameters:

  • q (ttnn.Tensor) – [1, H, S, d] bf16 | fp8_e4m3 ROW_MAJOR. H must be divisible by n_kv; H/n_kv may be 16 or a multiple of 32.

  • k (ttnn.Tensor) – [B, n_kv, T, d] TILE bf16|bfloat8_b (B>1 only when indexed)

  • v (ttnn.Tensor) – [B, n_kv, T, v_dim] TILE bf16|bfloat8_b

  • indices (ttnn.Tensor) – [1, n_kv, S, TOPK] uint32 block-ids (-1 = sentinel, contiguous tail). Each row must contain a valid block, and valid block ids must be < T / block_size.

Keyword Arguments:

  • scale (float, optional) – defaults to d to the power of -0.5.

  • block_size (int) – KV block size in tokens; defaults to 128. Must be a multiple of 32 and divide T.

  • compute_kernel_config (ttnn.DeviceComputeKernelConfig, optional) –

  • cache_batch_idx (int, optional) – select the batch slot of a shared [B, n_kv, T, feature_dim] K/V cache. It is a dynamic runtime arg, so changing it (or T) does not recompile the kernels.

  • chunk_start_idx (int, optional) – global position of query row 0. When set, enforces a token-level causal mask on the diagonal block (the query’s own block); toggling set/unset (None vs int) selects a different cached program. Requires bf16 q; fp8 q with this set is rejected.

  • cluster_axis (int, optional) – SP mesh axis used to derive the per-device chunk_start (chunk_start_idx + rank*S) under sequence parallelism. Host-side only.

Returns:

ttnn.Tensor – [1, H, S, v_dim] ROW-MAJOR, dtype = q.

Additional preconditions: d and v_dim must be multiples of 32; TOPK times 4 and output row bytes must meet device DRAM alignment.