@triton.jit
def _causal_conv1d_varlen_states(X, CU_SEQLENS, STATES, state_len, dim,
    stride_x_seqlen, stride_x_dim, stride_states_batch,
    stride_states_seqlen, stride_states_dim, BLOCK_M: tl.constexpr, BLOCK_N:
    tl.constexpr):
    batch_idx = tl.program_id(2)
    STATES += batch_idx * stride_states_batch
    end_idx = tl.load(CU_SEQLENS + batch_idx + 1)
    start_idx = tl.maximum(tl.load(CU_SEQLENS + batch_idx), end_idx - state_len
        )
    rows = end_idx - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(0, BLOCK_M)
    cols = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N)
    x = tl.load(X + rows[:, None] * stride_x_seqlen + cols[None, :] *
        stride_x_dim, mask=(rows[:, None] >= start_idx) & (cols[None, :] <
        dim), other=0)
    rows_states = state_len - (tl.program_id(1) + 1) * BLOCK_M + tl.arange(
        0, BLOCK_M)
    tl.store(STATES + rows_states[:, None] * stride_states_seqlen + cols[
        None, :] * stride_states_dim, x, mask=(rows_states[:, None] >= 0) &
        (cols[None, :] < dim))
