ttnn.transformer.sparse_sdpa

ttnn.transformer.sparse_sdpa(q: ttnn.Tensor, kv: ttnn.Tensor, indices: ttnn.Tensor, v_dim: int, *, scale: float | None, k_chunk_size: int, cache_batch_idx: int | None, block_cyclic_sp_axis: int | None, block_cyclic_chunk_local: int | None) → ttnn.Tensor

Sparse MLA prefill (DeepSeek DSA), Blackhole single-chip. softmax(Q@Kᵀ·scale masked)@V over the top-k selected latents per query token; V = kv[…, :v_dim]. Masking is baked into indices (0xFFFFFFFF = masked; sentinels are a contiguous tail). Inputs are ROW-MAJOR DRAM-interleaved. K_DIM (the q/kv head dim) is taken from the tensors.

Parameters:

  • q (ttnn.Tensor) – [1, H, S, K_DIM] bf16 or fp8_e4m3 (H a multiple of 32)

  • kv (ttnn.Tensor) – [1, 1, T, K_DIM] bf16 or fp8_e4m3 (fp8 halves the K-gather bytes; tilized to bfp8_b in-op). When cache_batch_idx is set, [B, 1, T, K_DIM] and may be ND-sharded across DRAM banks.

  • indices (ttnn.Tensor) – [1, 1, S, TOPK] uint32

  • v_dim (int) – width of V (leading v_dim cols of the K_DIM-wide cache); the output width.

Keyword Arguments:

  • scale (float, optional) – defaults to K_DIM**-0.5.

  • k_chunk_size (int) – defaults to 128 (must divide TOPK, multiple of 32).

  • compute_kernel_config (ttnn.DeviceComputeKernelConfig, optional) –

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

  • block_cyclic_sp_axis (int, optional) – when set (with block_cyclic_chunk_local), indices are NATURAL token positions and kv is stored block-cyclic across an SP-sharded cache; the kernel remaps each index natural->physical page on the fly (no host reorder needed). This is the MESH axis the cache was striped over: sp is read from the mesh shape on that axis (the op derives it, so it cannot disagree with the device). T % sp == 0 required. Both must be set together.

  • block_cyclic_chunk_local (int, optional) – the per-shard chunk length (chunk_size_global / sp). Required iff block_cyclic_sp_axis is set. Cross-checked against q’s per-chip seq length: must be q_isl or tp*q_isl (tp = mesh_size/sp) — the only two values it can legally take.

Returns:

ttnn.Tensor – [1, H, S, v_dim] ROW-MAJOR, DRAM interleaved; dtype matches q (bf16->bf16, fp8->fp8).