# Adapted from: https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/mamba/ops/mamba_ssm.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

# Copyright (c) 2024, Tri Dao, Albert Gu.
# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/selective_state_update.py

import torch
import triton
import triton.language as tl

from .softplus import softplus

PAD_SLOT_ID = -1


@triton.jit
def convert_rs_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor:
    """Stochastic rounding: fp32 pair → fp16x2 using random bits (PTX cvt.rs.f16x2.f32)."""
    y = tl.inline_asm_elementwise(
        asm="""{
        cvt.rs.f16x2.f32 $0, $2, $1, $3;
        }""",
        constraints=("=r,r,r,r,r"),
        args=(x, rand),
        dtype=tl.float16,
        is_pure=True,
        pack=2,
    )
    return y


@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None})
@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None})
@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None})
@triton.heuristics(
    {
        "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"]
        is not None
    }
)
@triton.heuristics(
    {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}
)
@triton.heuristics(
    {
        "CACHE_INTERMEDIATE_STATES": lambda args: args["intermediate_states_buffer"]
        is not None
    }
)
@triton.heuristics(
    {
        "HAS_EAGLE_TREE_CUSTOM_ATTN_MASK": lambda args: args[
            "retrieve_parent_token_ptr"
        ]
        is not None
    }
)
@triton.heuristics(
    {
        "HAS_INTERMEDIATE_STATE_INDICES": lambda args: args[
            "intermediate_state_indices_ptr"
        ]
        is not None
    }
)
@triton.heuristics(
    {"HAS_STATE_SCALE": lambda args: args["state_scale_ptr"] is not None}
)
@triton.heuristics(
    {
        "HAS_INTERMEDIATE_STATE_SCALES": lambda args: args[
            "intermediate_state_scales_ptr"
        ]
        is not None
    }
)
@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None})
@triton.jit(do_not_specialize=["T"])
def _selective_scan_update_kernel(
    # Pointers to matrices
    state_ptr,
    state_scale_ptr,
    rand_seed_ptr,
    x_ptr,
    dt_ptr,
    dt_bias_ptr,
    A_ptr,
    B_ptr,
    C_ptr,
    D_ptr,
    z_ptr,
    out_ptr,
    state_batch_indices_ptr,
    pad_slot_id,
    intermediate_states_buffer,
    cache_steps,
    retrieve_parent_token_ptr,
    intermediate_state_indices_ptr,
    intermediate_state_scales_ptr,
    # Matrix dimensions
    batch,
    T,
    nheads,
    dim,
    dstate,
    nheads_ngroups_ratio,
    # Strides
    stride_state_batch,
    stride_state_head,
    stride_state_dim,
    stride_state_dstate,
    stride_state_scale_batch,
    stride_state_scale_head,
    stride_state_scale_dim,
    stride_x_batch,
    stride_x_T,
    stride_x_head,
    stride_x_dim,
    stride_dt_batch,
    stride_dt_T,
    stride_dt_head,
    stride_dt_dim,
    stride_dt_bias_head,
    stride_dt_bias_dim,
    stride_A_head,
    stride_A_dim,
    stride_A_dstate,
    stride_B_batch,
    stride_B_T,
    stride_B_group,
    stride_B_dstate,
    stride_C_batch,
    stride_C_T,
    stride_C_group,
    stride_C_dstate,
    stride_D_head,
    stride_D_dim,
    stride_z_batch,
    stride_z_T,
    stride_z_head,
    stride_z_dim,
    stride_out_batch,
    stride_out_T,
    stride_out_head,
    stride_out_dim,
    stride_retrieve_parent_token_batch,
    stride_retrieve_parent_token_T,
    # Meta-parameters
    DT_SOFTPLUS: tl.constexpr,
    TIE_HDIM: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    HAS_DT_BIAS: tl.constexpr,
    HAS_D: tl.constexpr,
    HAS_Z: tl.constexpr,
    HAS_STATE_BATCH_INDICES: tl.constexpr,
    DISABLE_STATE_UPDATE: tl.constexpr,
    CACHE_INTERMEDIATE_STATES: tl.constexpr,
    HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: tl.constexpr,
    HAS_INTERMEDIATE_STATE_INDICES: tl.constexpr,
    HAS_STATE_SCALE: tl.constexpr,
    HAS_INTERMEDIATE_STATE_SCALES: tl.constexpr,
    USE_RS_ROUNDING: tl.constexpr,
    PHILOX_ROUNDS: tl.constexpr,
    BLOCK_SIZE_DSTATE: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)
    pid_b = tl.program_id(axis=1)
    pid_h = tl.program_id(axis=2)

    # If HAS_STATE_BATCH_INDICES is true, then the ssm state's batch coordinate
    # is taken from the state_batch_indices_ptr Otherwise, the state coordinate
    # is the same as the batch id.
    if HAS_STATE_BATCH_INDICES:
        state_batch_indices_ptr += pid_b
        state_batch_idx = tl.load(state_batch_indices_ptr).to(tl.int64)
        state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head
        if HAS_STATE_SCALE:
            state_scale_ptr += (
                state_batch_idx * stride_state_scale_batch
                + pid_h * stride_state_scale_head
            )
    else:
        state_batch_idx = pid_b
        state_ptr += pid_b * stride_state_batch + pid_h * stride_state_head
        if HAS_STATE_SCALE:
            state_scale_ptr += (
                pid_b * stride_state_scale_batch + pid_h * stride_state_scale_head
            )

    x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head
    dt_ptr += pid_b * stride_dt_batch + pid_h * stride_dt_head
    if HAS_DT_BIAS:
        dt_bias_ptr += pid_h * stride_dt_bias_head
    A_ptr += pid_h * stride_A_head
    B_ptr += pid_b * stride_B_batch + (pid_h // nheads_ngroups_ratio) * stride_B_group
    C_ptr += pid_b * stride_C_batch + (pid_h // nheads_ngroups_ratio) * stride_C_group
    if HAS_Z:
        z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head
    out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = tl.arange(0, BLOCK_SIZE_DSTATE)
    state_ptrs = state_ptr + (
        offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate
    )

    mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate)
    if HAS_STATE_BATCH_INDICES:
        mask &= state_batch_idx != pad_slot_id
    state = tl.load(state_ptrs, mask=mask, other=0.0).to(tl.float32)

    if HAS_STATE_SCALE:
        state_scale_ptrs = state_scale_ptr + offs_m[:, None] * stride_state_scale_dim
        scales_mask = offs_m[:, None] < dim
        if HAS_STATE_BATCH_INDICES:
            scales_mask = scales_mask & (state_batch_idx != pad_slot_id)
        decode_scale = tl.load(state_scale_ptrs, mask=scales_mask, other=1.0).to(
            tl.float32
        )
        state = state * decode_scale

    if HAS_DT_BIAS:
        dt_bias_ptrs = dt_bias_ptr + offs_m * stride_dt_bias_dim
    if HAS_D:
        D_ptr += pid_h * stride_D_head
        D_ptrs = D_ptr + offs_m * stride_D_dim
    A_ptrs = A_ptr + offs_m[:, None] * stride_A_dim + offs_n[None, :] * stride_A_dstate

    cache_idx = -1
    if CACHE_INTERMEDIATE_STATES:
        if HAS_INTERMEDIATE_STATE_INDICES:
            intermediate_state_idx = tl.load(intermediate_state_indices_ptr + pid_b).to(
                tl.int64
            )
            cache_idx = intermediate_state_idx
        elif HAS_STATE_BATCH_INDICES:
            cache_idx = state_batch_idx
        else:
            cache_idx = pid_b

    current_step_idx = 0
    for _ in range(T):
        if HAS_EAGLE_TREE_CUSTOM_ATTN_MASK:
            if current_step_idx != 0 and cache_idx >= 0:
                parent_ptr = (
                    retrieve_parent_token_ptr
                    + pid_b * stride_retrieve_parent_token_batch
                    + current_step_idx * stride_retrieve_parent_token_T
                )
                parent_step_idx = tl.load(parent_ptr).to(tl.int32)

                if parent_step_idx >= 0 and parent_step_idx < T:
                    step_offset = parent_step_idx * nheads * dim * dstate
                    cache_ptr = (
                        intermediate_states_buffer
                        + cache_idx * cache_steps * nheads * dim * dstate
                        + step_offset
                        + pid_h * dim * dstate
                        + offs_m[:, None] * dstate
                        + offs_n[None, :]
                    )
                    state = tl.load(cache_ptr, mask=mask, other=0.0).to(tl.float32)
                    if HAS_INTERMEDIATE_STATE_SCALES:
                        # Dequantize using the stored per-step scale
                        iscale_load_ptrs = (
                            intermediate_state_scales_ptr
                            + cache_idx * cache_steps * nheads * dim
                            + parent_step_idx * nheads * dim
                            + pid_h * dim
                            + offs_m
                        )
                        iscale = tl.load(
                            iscale_load_ptrs, mask=offs_m < dim, other=1.0
                        ).to(tl.float32)
                        state = state * iscale[:, None]

        x_ptrs = x_ptr + offs_m * stride_x_dim
        dt_ptrs = dt_ptr + offs_m * stride_dt_dim
        B_ptrs = B_ptr + offs_n * stride_B_dstate
        C_ptrs = C_ptr + offs_n * stride_C_dstate
        if HAS_Z:
            z_ptrs = z_ptr + offs_m * stride_z_dim
        out_ptrs = out_ptr + offs_m * stride_out_dim

        x = tl.load(x_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
        if not TIE_HDIM:
            dt = tl.load(dt_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
            if HAS_DT_BIAS:
                dt += tl.load(dt_bias_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
            if DT_SOFTPLUS:
                dt = softplus(dt)
            A = tl.load(
                A_ptrs,
                mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate),
                other=0.0,
            ).to(tl.float32)
            dA = tl.exp(A * dt[:, None])
        else:
            dt = tl.load(dt_ptr).to(tl.float32)
            if HAS_DT_BIAS:
                dt += tl.load(dt_bias_ptr).to(tl.float32)
            if DT_SOFTPLUS:
                dt = softplus(dt)
            A = tl.load(A_ptr).to(tl.float32)
            dA = tl.exp(A * dt)  # scalar, not a matrix

        B = tl.load(B_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)
        C = tl.load(C_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)
        if HAS_D:
            D = tl.load(D_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
        if HAS_Z:
            z = tl.load(z_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

        dB = B[None, :] * dt[:, None] if not TIE_HDIM else B * dt
        state = state * dA + dB * x[:, None]

        if CACHE_INTERMEDIATE_STATES:
            if state_batch_idx != pad_slot_id:
                cache_ptr_base = (
                    intermediate_states_buffer
                    + cache_idx * cache_steps * nheads * dim * dstate
                    + current_step_idx * nheads * dim * dstate
                    + pid_h * dim * dstate
                )
                cache_ptrs = cache_ptr_base + (
                    offs_m[:, None] * dstate + offs_n[None, :]
                )
                if HAS_INTERMEDIATE_STATE_SCALES:
                    # Quantize intermediate state with per-step block scaling
                    int16_max_i: tl.constexpr = 32767.0
                    amax_i = tl.max(tl.abs(state), axis=-1)
                    encode_scale_i = tl.where(amax_i == 0.0, 1.0, int16_max_i / amax_i)
                    decode_scale_i = 1.0 / encode_scale_i
                    # Store intermediate decode scale
                    iscale_ptrs = (
                        intermediate_state_scales_ptr
                        + cache_idx * cache_steps * nheads * dim
                        + current_step_idx * nheads * dim
                        + pid_h * dim
                        + offs_m
                    )
                    tl.store(iscale_ptrs, decode_scale_i, mask=offs_m < dim)
                    # Quantize and store
                    q_state = state * encode_scale_i[:, None]
                    q_state = tl.extra.cuda.libdevice.round(q_state)
                    q_state = tl.minimum(tl.maximum(q_state, -int16_max_i), int16_max_i)
                    tl.store(
                        cache_ptrs, q_state.to(cache_ptrs.dtype.element_ty), mask=mask
                    )
                elif USE_RS_ROUNDING:
                    rand_seed = tl.load(rand_seed_ptr)
                    if HAS_STATE_BATCH_INDICES:
                        rand_offsets = (
                            state_batch_idx * stride_state_batch
                            + pid_h * stride_state_head
                        )
                    else:
                        rand_offsets = (
                            pid_b * stride_state_batch + pid_h * stride_state_head
                        )
                    rand_offsets += (
                        offs_m[:, None] * stride_state_dim
                        + offs_n[None, :] * stride_state_dstate
                    )
                    if PHILOX_ROUNDS > 0:
                        rand = tl.randint(rand_seed, rand_offsets, PHILOX_ROUNDS)
                    else:
                        rand = tl.randint(rand_seed, rand_offsets)
                    tl.store(cache_ptrs, convert_rs_fp16x2(state, rand), mask=mask)
                else:
                    tl.store(
                        cache_ptrs, state.to(cache_ptrs.dtype.element_ty), mask=mask
                    )

        out = tl.sum(state * C[None, :], axis=1)
        if HAS_D:
            out += x * D
        if HAS_Z:
            out *= z * tl.sigmoid(z)
        tl.store(out_ptrs, out, mask=offs_m < dim)

        current_step_idx += 1  # noqa: SIM113

        x_ptr += stride_x_T
        dt_ptr += stride_dt_T
        B_ptr += stride_B_T
        C_ptr += stride_C_T
        out_ptr += stride_out_T
        if HAS_Z:
            z_ptr += stride_z_T

    if not DISABLE_STATE_UPDATE:
        if HAS_STATE_SCALE:
            # Quantize state back with block scaling (int16 path)
            int16_max: tl.constexpr = 32767.0
            amax = tl.max(tl.abs(state), axis=-1, keep_dims=True)
            encode_scale = tl.where(amax == 0.0, 1.0, int16_max / amax)
            new_decode_scale = 1.0 / encode_scale
            # Store updated decode scales
            dst_scales_mask = offs_m[:, None] < dim
            tl.store(state_scale_ptrs, new_decode_scale, mask=dst_scales_mask)
            # Quantize
            state = state * encode_scale
            state = tl.extra.cuda.libdevice.round(state)
            state = tl.minimum(tl.maximum(state, -int16_max), int16_max)
            tl.store(state_ptrs, state.to(state_ptrs.dtype.element_ty), mask=mask)
        elif USE_RS_ROUNDING:
            # Stochastic rounding for fp16 state
            rand_seed = tl.load(rand_seed_ptr)
            if HAS_STATE_BATCH_INDICES:
                rand_offsets = (
                    state_batch_idx * stride_state_batch + pid_h * stride_state_head
                )
            else:
                rand_offsets = pid_b * stride_state_batch + pid_h * stride_state_head
            rand_offsets += (
                offs_m[:, None] * stride_state_dim
                + offs_n[None, :] * stride_state_dstate
            )
            if PHILOX_ROUNDS > 0:
                rand = tl.randint(rand_seed, rand_offsets, PHILOX_ROUNDS)
            else:
                rand = tl.randint(rand_seed, rand_offsets)
            tl.store(state_ptrs, convert_rs_fp16x2(state, rand), mask=mask)
        else:
            tl.store(state_ptrs, state.to(state_ptrs.dtype.element_ty), mask=mask)


def selective_state_update_triton(
    state,
    x,
    dt,
    A,
    B,
    C,
    D=None,
    z=None,
    dt_bias=None,
    dt_softplus=False,
    state_batch_indices=None,
    pad_slot_id=PAD_SLOT_ID,
    out=None,
    disable_state_update=False,
    intermediate_states_buffer=None,
    cache_steps=None,
    retrieve_parent_token=None,
    intermediate_state_indices=None,
    state_scale=None,
    intermediate_state_scales=None,
    rand_seed=None,
    philox_rounds=0,
):
    """
    Argument:
        state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
        x: (batch, dim) or (batch, nheads, dim) for single-token or (batch, T, nheads, dim) for multi-token
        dt: (batch, dim) or (batch, nheads, dim) for single-token or (batch, T, nheads, dim) for multi-token
        A: (dim, dstate) or (nheads, dim, dstate)
        B: (batch, dstate) or (batch, ngroups, dstate) for single-token or (batch, T, ngroups, dstate) for multi-token
        C: (batch, dstate) or (batch, ngroups, dstate) for single-token or (batch, T, ngroups, dstate) for multi-token
        D: (dim,) or (nheads, dim)
        z: (batch, dim) or (batch, nheads, dim) for single-token or (batch, T, nheads, dim) for multi-token
        dt_bias: (dim,) or (nheads, dim)
        pad_slot_id: int
            if cache_indices is passed, lets the kernel identify padded
            entries that will not be processed,
            for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id]
            in this case, the kernel will not process entries at
            indices 0 and 3
        out: Preallocated ssm output tensor. Assume same shape as x.
             In-place updated.
        disable_state_update: If True, don't write back to state (for speculative verify)
        intermediate_states_buffer: Buffer to cache intermediate states
        cache_steps: Total number of steps in the buffer
        retrieve_parent_token: (batch, T) tensor of parent token indices for EAGLE tree attention
        intermediate_state_indices: (batch,) tensor of indices for intermediate_states_buffer operations.
            If provided, uses these indices instead of state_batch_indices for the buffer.
        state_scale: Optional (batch, nheads, dim, 1) float32 tensor of decode scales for
            block-scaled state quantization. When provided, the kernel dequantizes state
            on load and requantizes on store.
        rand_seed: Optional scalar tensor (int64 or int32) for Philox PRNG seed.
            When provided, enables stochastic rounding for fp16 state stores.
            Mutually exclusive with state_scale (int16 block scaling).
        philox_rounds: Number of Philox PRNG rounds (0 = default, typically 10).
    """
    # Track original x dimensionality to squeeze output appropriately
    x_orig_dim = x.dim()

    if state.dim() == 3:
        state = state.unsqueeze(1)
    if x.dim() == 2:
        x = x.unsqueeze(1)
    if x.dim() == 3:
        x = x.unsqueeze(1)
    if dt.dim() == 2:
        dt = dt.unsqueeze(1)
    if dt.dim() == 3:
        dt = dt.unsqueeze(1)
    if A.dim() == 2:
        A = A.unsqueeze(0)
    if B.dim() == 2:
        B = B.unsqueeze(1)
    if B.dim() == 3:
        B = B.unsqueeze(1)
    if C.dim() == 2:
        C = C.unsqueeze(1)
    if C.dim() == 3:
        C = C.unsqueeze(1)
    if D is not None and D.dim() == 1:
        D = D.unsqueeze(0)
    if z is not None:
        if z.dim() == 2:
            z = z.unsqueeze(1)
        if z.dim() == 3:
            z = z.unsqueeze(1)
    if dt_bias is not None and dt_bias.dim() == 1:
        dt_bias = dt_bias.unsqueeze(0)
    if out is None:
        out = torch.empty_like(x)
    if out.dim() == 2:
        out = out.unsqueeze(1)
    if out.dim() == 3:
        out = out.unsqueeze(1)

    _, nheads, dim, dstate = state.shape
    batch, T, _, _ = x.shape

    assert x.shape == (batch, T, nheads, dim)
    assert dt.shape == x.shape
    assert A.shape == (nheads, dim, dstate)
    ngroups = B.shape[2]
    assert nheads % ngroups == 0, "nheads must be divisible by ngroups"
    assert B.shape == (batch, T, ngroups, dstate)
    assert C.shape == B.shape
    if D is not None:
        assert D.shape == (nheads, dim)
    if z is not None:
        assert z.shape == x.shape
    if dt_bias is not None:
        assert dt_bias.shape == (nheads, dim)
    if state_batch_indices is not None:
        assert state_batch_indices.shape == (batch,)
    assert out.shape == x.shape

    grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads)
    z_strides = (
        (z.stride(0), z.stride(1), z.stride(2), z.stride(3))
        if z is not None
        else (0, 0, 0, 0)
    )
    # We don't want autotune since it will overwrite the state
    # We instead tune by hand.
    BLOCK_SIZE_M, num_warps = (
        (32, 4)
        if dstate <= 16
        else (
            (16, 4)
            if dstate <= 32
            else ((8, 4) if dstate <= 64 else ((4, 4) if dstate <= 128 else ((4, 8))))
        )
    )
    tie_hdim = (
        A.stride(-1) == 0
        and A.stride(-2) == 0
        and dt.stride(-1) == 0
        and (dt_bias is None or dt_bias.stride(-1) == 0)
    )

    if state_scale is not None:
        assert state_scale.dtype == torch.float32
        assert state_scale.shape == (
            state.shape[0],
            nheads,
            dim,
        ) or state_scale.shape == (state.shape[0], nheads, dim, 1)

    state_scale_strides = (
        (state_scale.stride(0), state_scale.stride(1), state_scale.stride(2))
        if state_scale is not None
        else (0, 0, 0)
    )

    retrieve_parent_token_strides = (
        (retrieve_parent_token.stride(0), retrieve_parent_token.stride(1))
        if retrieve_parent_token is not None
        else (0, 0)
    )

    with torch.cuda.device(x.device.index):
        _selective_scan_update_kernel[grid](
            state,
            state_scale,
            rand_seed,
            x,
            dt,
            dt_bias,
            A,
            B,
            C,
            D,
            z,
            out,
            state_batch_indices,
            pad_slot_id,
            intermediate_states_buffer,
            cache_steps if cache_steps is not None else 0,
            retrieve_parent_token,
            intermediate_state_indices,
            intermediate_state_scales,
            batch,
            T,
            nheads,
            dim,
            dstate,
            nheads // ngroups,
            state.stride(0),
            state.stride(1),
            state.stride(2),
            state.stride(3),
            state_scale_strides[0],
            state_scale_strides[1],
            state_scale_strides[2],
            x.stride(0),
            x.stride(1),
            x.stride(2),
            x.stride(3),
            dt.stride(0),
            dt.stride(1),
            dt.stride(2),
            dt.stride(3),
            *(dt_bias.stride(0), dt_bias.stride(1)) if dt_bias is not None else (0, 0),
            A.stride(0),
            A.stride(1),
            A.stride(2),
            B.stride(0),
            B.stride(1),
            B.stride(2),
            B.stride(3),
            C.stride(0),
            C.stride(1),
            C.stride(2),
            C.stride(3),
            *(D.stride(0), D.stride(1)) if D is not None else (0, 0),
            z_strides[0],
            z_strides[1],
            z_strides[2],
            z_strides[3],
            out.stride(0),
            out.stride(1),
            out.stride(2),
            out.stride(3),
            retrieve_parent_token_strides[0],
            retrieve_parent_token_strides[1],
            dt_softplus,
            tie_hdim,
            BLOCK_SIZE_M,
            DISABLE_STATE_UPDATE=disable_state_update,
            PHILOX_ROUNDS=philox_rounds,
            num_warps=num_warps,
        )
    # Squeeze T dimension if original x didn't have it (was 2D or 3D)
    if x_orig_dim < 4:
        out = out.squeeze(1)
    return out
