ttnn.situ_glu

ttnn.situ_glu(gate: ttnn.Tensor, up: ttnn.Tensor, beta1: float, beta2: float, *, memory_config: ttnn.MemoryConfig = None, sub_core_grids: ttnn.CoreRangeSet = None, sub_device_id: ttnn.SubDeviceId | None) → ttnn.Tensor

Computes Moonshot’s SiTU-GLU activation over the pre-split gate and up tensors.

\[\mathrm{output\_tensor}_i = \left(\verb|beta1| \cdot \tanh(\mathrm{gate}_i / \verb|beta1|) \cdot \sigma(\mathrm{gate}_i)\right) \cdot \left(\verb|beta2| \cdot \tanh(\mathrm{up}_i / \verb|beta2|)\right)\]
Parameters:

  • gate (ttnn.Tensor)the gate input tensor.

  • up (ttnn.Tensor)the up input tensor.

  • beta1 (float)the softcap beta applied to the gate half. Must be non-zero.

  • beta2 (float)the softcap beta applied to the up half. Must be non-zero.

Keyword Arguments:

  • memory_config (ttnn.MemoryConfig, optional)memory configuration for the operation. Defaults to None.

  • sub_core_grids (ttnn.CoreRangeSet, optional)the cores every composed step runs on. Defaults to None.

  • sub_device_id (ttnn.SubDeviceId, optional)sub-device whose worker cores to run on, as an alternative to spelling them out in sub_core_grids. Mutually exclusive with it. Defaults to None.

Returns:

ttnn.Tensorthe output tensor.

Note

Supported dtypes and layouts:

Dtypes

Layouts

BFLOAT16, BFLOAT8_B

TILE

Implemented for Blackhole only.

Restricting the cores forces the intermediates to the output’s memory space, because the L1 placement this picks on a full grid is unsafe next to a concurrently running op. For the same reason a core restriction rejects an interleaved-L1 output, whether asked for through memory_config or inherited from an interleaved-L1 input_tensor_a: such a buffer takes L1 on the cores restricted away. Sharded L1 is accepted – its shard spec confines it.