@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 _forward(X, OUT, LUT, sizemax, stride_zx, stride_zout, stride_hout, **meta
    ):
    TN = meta['TN']
    BLOCK = meta['BLOCK']
    pidhm = tl.program_id(0)
    pidz = tl.program_id(1)
    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 + 0)
    rowid = tl.load(LUT + offset + rbmn * 4 + 2)
    headid = tl.load(LUT + offset + rbmn * 4 + 3)
    px = X + pidz * stride_zx + blockid * BLOCK * BLOCK + rxm * BLOCK + rxn
    x = tl.load(px, mask=check, other=-float('inf'))
    x = x.to(tl.float32)
    c = tl.max(x, axis=0)
    out = tl.log(tl.sum(tl.exp(x - c), axis=0)) + c
    pout = (OUT + pidz * stride_zout + headid * stride_hout + rowid * BLOCK +
        rxm)
    tl.store(pout, out)
