@triton.heuristics({'num_warps': lambda *args, **meta: num_warps(args[3] *
    meta['BLOCK'])})
@triton.heuristics({'TN': lambda *args, **meta: next_power_of_2(args[3]) *
    meta['BLOCK']})
@triton.jit
def _backward(DX, DOUT, LUT, sizemax, stride_zdx, stride_zdout,
    stride_hdout, **meta):
    pidhm = tl.program_id(0)
    pidz = tl.program_id(1)
    TN = meta['TN']
    BLOCK = meta['BLOCK']
    rxm = pidhm % BLOCK
    rbm = pidhm // BLOCK
    rxn = tl.arange(0, TN) % BLOCK
    rbn = tl.arange(0, TN) // BLOCK
    header = LUT + rbm * 2
    size = tl.load(header + 0)
    offset = tl.load(header + 1)
    check = rbn < size
    rbmn = tl.where(check, rbn, size - 1)
    blockid = tl.load(LUT + offset + rbmn * 4)
    rowid = tl.load(LUT + offset + rbmn * 4 + 2)
    headid = tl.load(LUT + offset + rbmn * 4 + 3)
    pdx = DX + pidz * stride_zdx + blockid * BLOCK * BLOCK + rxm * BLOCK + rxn
    pdout = (DOUT + pidz * stride_zdout + headid * stride_hdout + rowid *
        BLOCK + rxm)
    dx_zeros = tl.load(pdx, mask=check, other=0)
    dout = tl.load(pdout)
    dx = dout - dx_zeros
    tl.store(pdx, dx, mask=check)
