@triton.jit
def _ld_softmax_bwd_kernel(ds_ptr, p_ptr, dp_ptr, ds_row_stride,
    p_row_stride, dp_row_stride, n_rows, n_cols, BLOCK_SIZE: tl.constexpr):
    row_idx = tl.program_id(0)
    p_start_ptr = p_ptr + row_idx * p_row_stride
    dp_start_ptr = dp_ptr + row_idx * dp_row_stride
    col_offsets = tl.arange(0, BLOCK_SIZE)
    p_ptrs = p_start_ptr + col_offsets
    dp_ptrs = dp_start_ptr + col_offsets
    mask = col_offsets < n_cols
    p_row = tl.load(p_ptrs, mask=mask, other=0)
    dp_row = tl.load(dp_ptrs, mask=mask, other=0)
    ds_row = p_row * (dp_row - tl.sum(p_row * dp_row, axis=0))
    ds_start_ptr = ds_ptr + row_idx * ds_row_stride
    ds_ptrs = ds_start_ptr + col_offsets
    tl.store(ds_ptrs, ds_row, mask=mask)
