ttnn.transformer.chunk_gated_delta_rule

ttnn.transformer.chunk_gated_delta_rule(q: ttnn.Tensor, k: ttnn.Tensor, v: ttnn.Tensor, g: ttnn.Tensor, beta: ttnn.Tensor, *, scale: float | None, initial_state: ttnn.Tensor | None, output_final_state: bool, chunk_size: int, use_qk_l2norm: bool, output_head_major: bool, masks: ttnn.Tensor | None) → tuple[ttnn.Tensor, ttnn.Tensor | None]

Standalone chunked Gated Delta Rule forward (flash-linear-attention algorithm).

Parameters:

  • q (ttnn.Tensor) – [B, T, H, K]

  • k (ttnn.Tensor) – [B, T, H, K]

  • v (ttnn.Tensor) – [B, T, HV, V]

  • g (ttnn.Tensor) – [B, T, HV] log-space decay

  • beta (ttnn.Tensor) – [B, T, HV]

Keyword Arguments:

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

  • initial_state (ttnn.Tensor, optional) – [B, HV, K, V].

  • output_final_state (bool) – default False.

  • chunk_size (int) – default 64.

  • use_qk_l2norm (bool) – default False.

  • output_head_major (bool) – default False. When True, o is returned head-major as [B*HV, T, V] in TILE layout (skips the token<->head permute round-trip); otherwise token-major [B, T, HV, V] ROW_MAJOR.

  • memory_config (ttnn.MemoryConfig, optional) –

  • compute_kernel_config (ttnn.DeviceComputeKernelConfig, optional) –

  • ones (eye, tril,) – [1,1,C,C] fp32 TILE constant tiles (identity, lower-triangular ones, all-ones). Caller-supplied so they are device-resident before trace capture and their lifetime is device-scoped. Traced callers MUST pass these (an internal build does a host upload, illegal under trace); if omitted they are built eagerly.

  • masks (ttnn.Tensor, optional) – [1,1,32,96] fp32 TILE quadrant masks; supplied with eye/ tril/ones.

Returns:

tuple[ttnn.Tensor, Optional[ttnn.Tensor]] – o [B, T, HV, V] (or [B*HV, T, V] if output_head_major), final_state [B, HV, K, V] (if output_final_state).