ttnn.transformer.ring_mla

ttnn.transformer.ring_mla(input_tensor_q: ttnn.Tensor, input_tensor_kv: ttnn.Tensor, *, persistent_output_buffer_kv: ttnn.Tensor, head_dim_v: int, logical_n: int, program_config: ttnn.SDPAProgramConfig, scale: float = None, compute_kernel_config: ttnn.DeviceComputeKernelConfig = None, dim: int, multi_device_global_semaphore: List[ttnn.GlobalSemaphore], num_links: int, cluster_axis: int, mesh_device: ttnn.MeshDevice, topology: ttnn.ccl.Topology, subdevice_id: tt.tt_metal.SubDeviceId | None = None, ccl_core_grid_offset: ttnn.CoreCoord, use_column_major_ccl: bool = False, is_balanced: bool = False, kv_cache_batch_idx: int = None, kv_actual_isl: int | None)

Causal Ring MLA attention over a single KV tensor.

K and V are represented by one tensor. QK uses the full KV head dimension and QKT@V uses the first head_dim_v columns as V. The V dimension must be smaller than K and tile aligned. The KV tensor must have one shared KV head.

Parameters:

  • input_tensor_q (ttnn.Tensor) – Queries [b x nqh x N/num_devices x dh].

  • input_tensor_kv (ttnn.Tensor) – Shared KV tensor [b x nkv x N/num_devices x dh].

Keyword Arguments:

  • persistent_output_buffer_kv (ttnn.Tensor) – Persistent buffer for gathered KV tensor.

  • head_dim_v (int) – Tile-aligned V hidden dimension, read from KV’s prefix.

  • logical_n (int) – Logical global sequence length before sharding.

  • program_config (ttnn.SDPAProgramConfig) – Program configuration for the operation.

  • scale (float, optional) – Scale factor for QK^T. Defaults to None.

  • compute_kernel_config (ttnn.DeviceComputeKernelConfig, optional) – Defaults to None.

  • dim (int) – Dimension for ring all-gather.

  • multi_device_global_semaphore (List[ttnn.GlobalSemaphore]) – Global semaphores for CCL synchronization.

  • num_links (int) – Number of CCL links.

  • cluster_axis (int) – Mesh axis for all-gather.

  • mesh_device (ttnn.MeshDevice) – Multi-device mesh.

  • topology (ttnn.ccl.Topology) – Communication topology.

  • subdevice_id (Optional[tt.tt_metal.SubDeviceId]) – Sub-device identifier. Defaults to None.

  • ccl_core_grid_offset (ttnn.CoreCoord) – Core grid offset for CCL workers.

  • use_column_major_ccl (bool) – If true, allocate CCL workers column-major. Defaults to False.

  • is_balanced (bool) – Whether to use balanced causal work distribution. Defaults to False.

  • kv_cache_batch_idx (int, optional) – Selects one batch slot from an indexed K/V cache. Defaults to None.

  • kv_actual_isl (int, optional) – Prior valid global KV length before this fixed-size chunk. When passed, enables KV-pad-aware rotation and derives current valid tokens as logical_n - kv_actual_isl.

Returns:

(ttnn.Tensor, ttnn.Tensor) – - Attention output [b x nqh x N/num_devices x head_dim_v]. - Streaming statistics scratch [b x nqh x 2*N/num_devices x 1].