@triton.autotune(configs=_k_configs, key=['N'])
@triton.jit
def k_dropout_fw(Y, X, BIAS, SEEDS, stride, N, p, **META):
    """
    Apply dropout on an input tensor
    Y : Output (M, N)
    X : Input (M, N)
    S : Seeds (M,)
    p : dropout probability
    """
    BLOCK_SIZE = META['BLOCK_SIZE']
    row = tl.program_id(axis=0)
    col = tl.program_id(axis=1)
    offsets = row * stride + col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < N
    x_ptrs = X + offsets
    x = tl.load(x_ptrs, mask=mask)
    if META['USE_BIAS']:
        b_ptrs = BIAS + col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
        b = tl.load(b_ptrs, mask=mask)
        x += b
    if META['ACTIVATION']:
        x = META['ACTIVATION'](x)
    if p > 0.0:
        output = _drop_and_scale(SEEDS, row, p, offsets, x)
    else:
        output = x
    y_ptrs = Y + offsets
    tl.store(y_ptrs, output, mask=mask)
