143 lines
6.9 KiB
Python
143 lines
6.9 KiB
Python
# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/fmha.py
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
import stream_attn_cuda
|
|
|
|
|
|
def convert_blockmask(blockmask, causal):
|
|
"""Convert from the 0-1 format to the format used by the CUDA code.
|
|
0 means the block is skipped.
|
|
nonzero means the block is not skipped.
|
|
Argument:
|
|
blockmask: (row, col): a 0-1 tensor
|
|
Return:
|
|
blockmask_converted: (col, row), dtype torch.int32: for each column, it contains the row
|
|
indices of the nonzero blocks, padded with -1 to reach length @row.
|
|
The indices are multiplied by 4, with the smallest bit used to encode whether
|
|
it is the first nonzero in its row, and the 2nd smallest bit to encode whether it is
|
|
the last nonzero in its row..
|
|
"""
|
|
assert not causal
|
|
# TD [2022-05-13]: The indexing and sorting is very tricky
|
|
nrow, ncol = blockmask.shape
|
|
# Sort does not support bool on CUDA
|
|
blockmask = blockmask.to(dtype=torch.uint8)
|
|
nonzero_val, nonzero_sorted_rowidx = blockmask.sort(dim=0, stable=True, descending=True)
|
|
nonzero_unsorted_rowidx = nonzero_sorted_rowidx.argsort(dim=0)
|
|
last_nonzero_col_per_row = blockmask.sort(dim=-1, stable=True).indices[:, -1]
|
|
last_nonzero_col_per_row_after_sort = nonzero_unsorted_rowidx[
|
|
torch.arange(nrow, device=blockmask.device), last_nonzero_col_per_row
|
|
]
|
|
first_nonzero_col_per_row = blockmask.sort(dim=-1, stable=True, descending=True).indices[:, 0]
|
|
first_nonzero_col_per_row_after_sort = nonzero_unsorted_rowidx[
|
|
torch.arange(nrow, device=blockmask.device), first_nonzero_col_per_row
|
|
]
|
|
nonzero_idx = nonzero_sorted_rowidx * 4
|
|
nonzero_idx[last_nonzero_col_per_row_after_sort, last_nonzero_col_per_row] += 2
|
|
nonzero_idx[first_nonzero_col_per_row_after_sort, first_nonzero_col_per_row] += 1
|
|
nonzero_idx[nonzero_val == 0] = -1
|
|
return nonzero_idx.T.contiguous().to(dtype=torch.int32)
|
|
|
|
|
|
def _stream_blocksparse_attn_forward(qkv, cu_seqlens, blockmask, dropout_p, max_s, softmax_scale,
|
|
causal, return_softmax):
|
|
context, softmax_lse, *rest = stream_attn_cuda.fwd_block(qkv, cu_seqlens, blockmask, dropout_p,
|
|
max_s, softmax_scale, causal,
|
|
return_softmax, None)
|
|
# if context.isnan().any() or softmax_lse.isnan().any():
|
|
# breakpoint()
|
|
S_dmask = rest[0] if return_softmax else None
|
|
return context, softmax_lse, S_dmask
|
|
|
|
|
|
def _stream_blocksparse_attn_backward(dout, qkv, out, S_dmask, softmax_lse, cu_seqlens, blockmask,
|
|
dropout_p, max_s, softmax_scale, causal):
|
|
dqkv, dp, softmax_d = stream_attn_cuda.bwd_block(dout, qkv, out, S_dmask, softmax_lse, cu_seqlens,
|
|
blockmask, dropout_p, softmax_scale, max_s,
|
|
causal, None)
|
|
# if dqkv.isnan().any() or softmax_d.isnan().any():
|
|
# breakpoint()
|
|
return dqkv
|
|
|
|
|
|
class StreamBlocksparseAttnFun(torch.autograd.Function):
|
|
|
|
@staticmethod
|
|
def forward(ctx, qkv, cu_seqlens, blockmask, dropout_p, max_s, softmax_scale, causal):
|
|
# Save rng_state because the backward pass will regenerate the dropout mask
|
|
rng_state = torch.cuda.get_rng_state() if dropout_p > 0 else None
|
|
if softmax_scale is None:
|
|
softmax_scale = qkv.shape[-1] ** (-0.5)
|
|
context, softmax_lse, S_dmask = _stream_blocksparse_attn_forward(
|
|
qkv, cu_seqlens, blockmask, dropout_p, max_s, softmax_scale, causal=causal,
|
|
return_softmax=False
|
|
)
|
|
ctx.save_for_backward(qkv, context, S_dmask, softmax_lse, cu_seqlens, blockmask, rng_state)
|
|
ctx.dropout_p = dropout_p
|
|
ctx.max_s = max_s
|
|
ctx.softmax_scale = softmax_scale
|
|
ctx.causal = causal
|
|
return context
|
|
|
|
@staticmethod
|
|
def backward(ctx, dout):
|
|
qkv, context, S_dmask, softmax_lse, cu_seqlens, blockmask, rng_state = ctx.saved_tensors
|
|
if rng_state is not None:
|
|
cur_rng_state = torch.cuda.get_rng_state()
|
|
torch.cuda.set_rng_state(rng_state)
|
|
# S_dmask is None, temporarily use another tensor just to get it running
|
|
dqkv = _stream_blocksparse_attn_backward(
|
|
dout, qkv, context, context, softmax_lse, cu_seqlens, blockmask, ctx.dropout_p,
|
|
ctx.max_s, ctx.softmax_scale, ctx.causal
|
|
)
|
|
if rng_state is not None:
|
|
torch.cuda.set_rng_state(cur_rng_state)
|
|
return dqkv, None, None, None, None, None, None, None
|
|
|
|
|
|
# We duplicate code to return both the output and the softmax for testing
|
|
# Returning both makes backward a bit slower, so we want to keep using the other version for speed.
|
|
class StreamBlocksparseAttnFunWithS(torch.autograd.Function):
|
|
|
|
@staticmethod
|
|
def forward(ctx, qkv, cu_seqlens, blockmask, dropout_p, max_s, softmax_scale, causal):
|
|
# Save rng_state because the backward pass is gonna regenerate the dropout mask
|
|
rng_state = torch.cuda.get_rng_state() if dropout_p > 0 else None
|
|
if softmax_scale is None:
|
|
softmax_scale = qkv.shape[-1] ** (-0.5)
|
|
context, softmax_lse, S_dmask = _stream_blocksparse_attn_forward(
|
|
qkv, cu_seqlens, blockmask, dropout_p, max_s, softmax_scale, causal=causal,
|
|
return_softmax=True
|
|
)
|
|
ctx.save_for_backward(qkv, context, S_dmask, softmax_lse, cu_seqlens, blockmask, rng_state)
|
|
ctx.dropout_p = dropout_p
|
|
ctx.max_s = max_s
|
|
ctx.softmax_scale = softmax_scale
|
|
ctx.causal = causal
|
|
return context, S_dmask, softmax_lse
|
|
|
|
@staticmethod
|
|
def backward(ctx, dout, _dS_dmask_ignored, _dsoftmax_sum_ignored):
|
|
qkv, context, S_dmask, softmax_lse, cu_seqlens, blockmask, rng_state = ctx.saved_tensors
|
|
if rng_state is not None:
|
|
cur_rng_state = torch.cuda.get_rng_state()
|
|
torch.cuda.set_rng_state(rng_state)
|
|
dqkv = _stream_blocksparse_attn_backward(
|
|
dout, qkv, context, S_dmask, softmax_lse, cu_seqlens, blockmask, ctx.dropout_p,
|
|
ctx.max_s, ctx.softmax_scale, ctx.causal
|
|
)
|
|
if rng_state is not None:
|
|
torch.cuda.set_rng_state(cur_rng_state)
|
|
return dqkv, None, None, None, None, None, None
|
|
|
|
|
|
def stream_blocksparse_attn_func(qkv, cu_seqlens, blockmask, dropout_p, max_s, softmax_scale=None,
|
|
causal=False, return_attn_probs=False, convert_mask=True):
|
|
"""dropout_p should be set to 0.0 during evaluation
|
|
"""
|
|
func = StreamBlocksparseAttnFun if not return_attn_probs else StreamBlocksparseAttnFunWithS
|
|
if convert_mask:
|
|
blockmask = convert_blockmask(blockmask, causal=causal)
|
|
return func.apply(qkv, cu_seqlens, blockmask, dropout_p, max_s, softmax_scale, causal)
|