File size: 3,193 Bytes
eafbe80 | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """
Dependencies:
$ pip install mamba-ssm==2.2.2 triton==2.3.1
For correctness check, see:
https://github.com/sustcsonglin/flash-linear-attention/pull/49
"""
import torch
import triton
from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined
from fla.ops.simple_gla import chunk_simple_gla
@triton.testing.perf_report(
triton.testing.Benchmark(
# argument names to use as an x-axis for the plot
x_names=['T'],
# different possible values for `x_name`
x_vals=[64] + [128 * 2 ** i for i in range(0, 8)],
# argument name whose value corresponds to a different line in the plot
line_arg='provider',
# possible values for `line_arg``
line_vals=["chunk_simple_gla", "mamba2_ssd"],
# label name for the lines
line_names=["chunk_simple_gla", "mamba2_ssd"],
# line styles
styles=[('blue', '-'), ('red', '-')],
ylabel="Execution Time (ms)", # label name for the y-axis
# name for the plot. Used also as a file name for saving the plot.
plot_name="Performance",
args={},
),
)
def benchmark(T, provider):
# TODO: also add bwd pass benchmark
from fla.utils import device
dtype = torch.bfloat16
B, H, D = 16, 8, 128
# TODO: test more shapes
# TODO: different values for D_V and D_QK
# TODO: different values for H_Q and H_KV
final_state = False # does not impact performance
# initialize Mamba2-format inputs
X_mamba = 0.1 * torch.randn(B, T, H, D, dtype=dtype, device=device)
dt_mamba = torch.ones(B, T, H, dtype=dtype, device=device)
A_mamba = -0.1 * torch.rand(H, dtype=dtype, device=device)
B_mamba = 0.1 * torch.randn(B, T, H, D, dtype=dtype, device=device)
C_mamba = 0.1 * torch.randn(B, T, H, D, dtype=dtype, device=device)
quantiles = [0.5, 0.2, 0.8]
if provider == 'chunk_simple_gla':
# mapping inputs Mamba2 -> FLA
# C, B, X: [B, T, H, D] -> [B, H, T, D]
# g: [B, T, H] -> [B, H, T]
q = C_mamba.transpose(1, 2).contiguous()
k = B_mamba.transpose(1, 2).contiguous()
v = X_mamba.transpose(1, 2).contiguous()
g = (A_mamba * dt_mamba).transpose(1, 2).contiguous()
# NOTE: whether to include the memory-copy cost of `contiguous()`?
# this depends on the memory layout used by surrounding non-SSM layers
results = triton.testing.do_bench(
lambda: chunk_simple_gla(
q, k, v, g, scale=1.0, output_final_state=final_state,
), quantiles=quantiles,
)
elif provider == 'mamba2_ssd':
# NOTE: `chunk_size` is configurable in mamba2 kernel
# here sets to the same hard-coded `BT = 64` as in simple_gla kernel
# TODO: benchmark different chunk sizes
results = triton.testing.do_bench(
lambda: mamba_chunk_scan_combined(
X_mamba, dt_mamba, A_mamba, B_mamba, C_mamba,
chunk_size=64, D=None, return_final_states=final_state,
),
quantiles=quantiles,
)
return results
if __name__ == '__main__':
benchmark.run(print_data=True, save_path='.')
|