File size: 2,984 Bytes
b3a3b15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | # python run_sparse_attn.py --in_features 8192 --batch_size 16 --seq_len 1920 --attn_topk 0.5
import math
import torch
from HybridTensor.modules.SelectiveMHA import SMHA, _update_kv_cache
from HybridTensor.utils.utils import arg_parser, generate_random_BH_index
from HybridTensor.utils.profiling import cuda_profiler
from HybridTensor.utils.generation import InferenceParams
if __name__ == "__main__":
args = arg_parser()
max_seqlen = args.seq_len + 128
max_batch_size = args.batch_size
device = torch.device(f"cuda:{args.device}")
# simulates SelectiveMHA inference generation stage
inference_params = InferenceParams(max_seqlen=max_seqlen, max_batch_size=max_batch_size)
nheads = args.in_features // 128
softmax_scale = 1 / (128 ** 0.5)
rotary_emb_dim = 0
mha = SMHA(
embed_dim=args.in_features,
num_heads=nheads,
num_heads_kv=None,
causal=True,
layer_idx=0,
use_flash_attn=True,
softmax_scale=softmax_scale,
return_residual=False,
rotary_emb_dim=rotary_emb_dim,
device=device,
dtype=torch.float16,
)
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
with torch.no_grad():
# prefill stage to generate kv cache for all batches
og_x = torch.randn(args.batch_size, args.seq_len, args.in_features, device=device, dtype=torch.float16, requires_grad=False)
# simulate kv cache, bug in flash_attn for larger batches
kv = torch.randn(args.batch_size, args.seq_len, 2, nheads, 128, device=device, dtype=torch.float16, requires_grad=False)
_ = _update_kv_cache(kv, inference_params, 0)
# increment the sequence length to move to the generation stage
inference_params.seqlen_offset += args.seq_len
input_x = torch.randn(args.batch_size, 1, args.in_features, device=device, dtype=torch.float16, requires_grad=False)
selected_heads = math.ceil(nheads * args.head_density)
# generate batch_head_idx for SelectiveMHA
# batch_head_index = generate_BH_index(args.batch_size, nheads, selected_heads, device=device)
batch_head_index = generate_random_BH_index(args.batch_size, nheads, selected_heads, device=device)
# generatation stage Standard MHA, batch_head_idx=None uses dense attention
out, standard_time_ms = cuda_profiler(mha, input_x, inference_params=inference_params, batch_head_idx=None)
print(f"Standard MHA time: {standard_time_ms:.3f} ms")
# generatation stage SelectiveMHA, pass batch_head_idx to use selective head attention
out, select_time_ms = cuda_profiler(mha, input_x, inference_params=inference_params, batch_head_idx=batch_head_index)
print(f"SelectMHA time: {select_time_ms:.3f} ms")
speedup = standard_time_ms / select_time_ms
print(f"Speedup: {speedup:.3f}")
|