ttnn.transformer.sparse_sdpa

ttnn.transformer.sparse_sdpa(q: ttnn.Tensor, kv: ttnn.Tensor, indices: ttnn.Tensor, v_dim: int, *, kv_format: SparseKVFormat, attention_sink: ttnn.Tensor | None, scale: float | None, k_chunk_size: int, cache_batch_idx: int | None, block_cyclic_sp_axis: int | None, block_cyclic_chunk_local: int | None, block_cyclic_cache_tp_sharded: bool) → 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/raw fp8, or one packed scaled-FP8 row per token; interpretation is selected only by kv_format, never inferred from width. The packed DSA row is [512 FP8 | 4 FP32 scales | 64 BF16 RoPE] = 656 bytes. When cache_batch_idx is set, B may exceed 1 and kv may be ND-sharded.

  • 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:

  • kv_format (SparseKVFormat)explicit physical/logical format of kv.

  • attention_sink (ttnn.Tensor, optional)[1, 1, 1, H] unpadded interleaved ROW_MAJOR BF16 tensor in DRAM. One scalar per head, shared across tokens. As in classic SDPA, the sink is multiplied by scale and contributes only to the softmax denominator. DeepSeek-V4 sinks are already scaled logits: pass (model_sink / scale).reshape(1, 1, 1, H) (scale != 0) to contribute exp(model_sink). Defaults to None.

  • 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 the slot does not recompile. Changing T also reuses the program for a plain interleaved cache, but recompiles for sharded or block-cyclic caches.

  • 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.

  • block_cyclic_cache_tp_sharded (bool)True = the cache is striped across ALL sp*tp devices (linear chip = sp_coord*tp + tp_coord), so stripes = sp*tp and per-stripe chunk = block_cyclic_chunk_local/tp.

Returns:

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