def test_simple_grad(self):
N = 2
H = 4
L = 100
S = 100
E = 32
k = 10
Q = torch.randn(N, H, L, E).to(self.device).requires_grad_(True)
K = torch.randn(N, H, S, E).to(self.device).requires_grad_(True)
topk = torch.round(torch.cumsum(torch.rand(N, H, L, k) * 10,
dim=-1)).long().to(self.device)
self._zero_grad(Q, K)
QK_full = torch.einsum("nhle,nhse->nhls", Q, K)
QK_selected = QK_full[torch.arange(N).view(N, 1, 1, 1).to(self.device),
torch.arange(H).view(1, H, 1, 1).to(self.device),
torch.arange(L).view(1, 1, L, 1).to(self.device),
topk]
QK_selected.sum().backward()
grad = [torch.clone(Q.grad), torch.clone(K.grad)]
self._zero_grad(Q, K)
QK_selected_hat = sparse_dot_product(Q, K, topk)
QK_selected_hat.sum().backward()
grad_hat = [torch.clone(Q.grad), torch.clone(K.grad)]
self.assertLess(torch.abs(QK_selected - QK_selected_hat).max(), 1e-4)
for g1, g2 in zip(grad, grad_hat):
self.assertLess(torch.abs(g1 - g2).max(), 1e-4)
def test_small_benchmark(self):
N = 12
H = 8
L = 1000
S = 1000
E = 32
k = 32
X = torch.randn(N, H, L, E)
Y = torch.randn(N, H, S, E)
topk = (torch.cumsum(torch.rand(N, H, L, k) * 40, dim=-1)).long()
n_runs = 10
s = time.time()
for run in range(n_runs):
products = sparse_dot_product(
X,
Y,
topk,
)
e = time.time()
t_s = (e - s) / n_runs
s = time.time()
for run in range(n_runs):
torch.einsum("nhle,nhse->nhls", X, Y)
e = time.time()
t_f = (e - s) / n_runs
print("Sparse: {}, Full: {}, F/S: {}".format(t_s, t_f, t_f / t_s))
def test_benchmark_forward_backward(self):
N = 12
H = 8
L = 1024
S = 1024
E = 32
k = 32
Q = torch.randn(N, H, L, E).to(self.device).requires_grad_(True)
K = torch.randn(N, H, S, E).to(self.device).requires_grad_(True)
topk = torch.round(
torch.cumsum(torch.rand(N, H, L, k) * (S // k),
dim=-1)).long().to(self.device)
n_runs = 10
self._zero_grad(Q, K)
s = time.time()
for i in range(n_runs):
QK = torch.einsum("nhle,nhse->nhls", Q, K)
QK.sum().backward()
e = time.time()
t_full = (e - s) / n_runs
self._zero_grad(Q, K)
s = time.time()
for i in range(n_runs):
QK = sparse_dot_product(Q, K, topk)
QK.sum().backward()
e = time.time()
t_sparse = (e - s) / n_runs
print("Benchmark Forward-Backward: T_Full: {}, T_Sparse: {}".format(
t_full, t_sparse))
def test_simple_product(self):
X = torch.randn(10, 4, 100, 32).cuda()
Y = torch.randn(10, 4, 100, 32).cuda()
lengths = torch.full((10, ), 100).int().cuda()
topk = (torch.cumsum(torch.rand(10, 4, 100, 10) * 10,
dim=-1)).long().cuda()
A = torch.randn(10, 4, 100, 100).to(X.device).requires_grad_(False)
topk_v, topk = torch.topk(A, 10, dim=-1)
topk = topk.contiguous()
products = sparse_dot_product(
X,
Y,
topk,
)
all_products = torch.einsum("nhle,nhse->nhls", X, Y)
self.assertLess(
torch.max(
torch.abs(products -
all_products[torch.arange(10).view(10, 1, 1, 1),
torch.arange(4).view(1, 4, 1, 1),
torch.arange(100).view(1, 1, 100, 1),
topk])), 1e-4)
def test_benchmark_backward(self):
N = 12
H = 8
L = 1024
S = 1024
E = 32
k = 32
Q = torch.randn(N, H, L, E).to(self.device).requires_grad_(True)
K = torch.randn(N, H, S, E).to(self.device).requires_grad_(True)
topk = torch.round(
torch.cumsum(torch.rand(N, H, L, k) * (S // k),
dim=-1)).long().to(self.device)
self._zero_grad(Q, K)
for i in range(2000):
QK = torch.einsum("nhle,nhse->nhls", Q, K)
QK.sum().backward()
self._zero_grad(Q, K)
s = torch.cuda.Event(enable_timing=True)
e = torch.cuda.Event(enable_timing=True)
QK = torch.einsum("nhle,nhse->nhls", Q, K)
s.record()
QK.sum().backward()
e.record()
torch.cuda.synchronize()
t_full = s.elapsed_time(e)
self._zero_grad(Q, K)
for i in range(2000):
QK = sparse_dot_product(Q, K, topk)
QK.sum().backward()
self._zero_grad(Q, K)
s = torch.cuda.Event(enable_timing=True)
e = torch.cuda.Event(enable_timing=True)
QK = sparse_dot_product(Q, K, topk)
s.record()
QK.sum().backward()
e.record()
torch.cuda.synchronize()
t_sparse = s.elapsed_time(e)
print("Benchmark Backward: T_Full: {}, T_Sparse: {}".format(
t_full, t_sparse))
def test_small_benchmark(self):
N = 12
H = 8
L = 1000
S = 1000
E = 32
k = 32
X = torch.randn(N, H, L, E).cuda()
Y = torch.randn(N, H, S, E).cuda()
A = torch.randn(N, H, L, S).to(X.device).requires_grad_(False)
topk_v, topk = torch.topk(A, k, dim=-1)
topk = topk.contiguous()
for i in range(1000):
products = sparse_dot_product(
X,
Y,
topk,
)
torch.cuda.synchronize()
s = torch.cuda.Event(enable_timing=True)
e = torch.cuda.Event(enable_timing=True)
s.record()
products = sparse_dot_product(
X,
Y,
topk,
)
e.record()
torch.cuda.synchronize()
t_s = s.elapsed_time(e)
for i in range(1000):
torch.einsum("nhle,nhse->nhls", X, Y)
s = torch.cuda.Event(enable_timing=True)
e = torch.cuda.Event(enable_timing=True)
s.record()
torch.einsum("nhle,nhse->nhls", X, Y)
e.record()
torch.cuda.synchronize()
t_f = s.elapsed_time(e)
print("Sparse: {}, Full: {}, F/S: {}".format(t_s, t_f, t_f / t_s))
def test_small_benchmark(self):
N = 12
H = 8
L = 1000
E = 32
S = 1000
k = 32
C = 100
I = 10
B = 32
Q = torch.randn(N, H, L, E).to(self.device)
K = torch.randn(N, H, S, E).to(self.device)
lengths = torch.full((N, ), L).int().to(self.device)
groups, counts = cluster_queries(Q, lengths, C, I, B)
Q_grouped = aggregate(Q, groups, 1 / counts.float())
QK = torch.einsum("nhle,nhse->nhls", Q_grouped, K)
_, topk = torch.topk(QK, k, dim=-1)
topk = topk.contiguous()
products = torch.zeros((N, H, L, k),
dtype=torch.float32).to(self.device)
products = clustered_sparse_dot_product(Q, K, topk, groups, counts,
lengths)
n_runs = 10
s = time.time()
for i in range(n_runs):
products = clustered_sparse_dot_product(Q, K, topk, groups, counts,
lengths)
e = time.time()
t_sc = (e - s) / n_runs
topk_broadcast = broadcast(
topk.float(), groups, torch.ones_like(counts, dtype=torch.float32),
torch.zeros((N, H, L, k), device=Q.device))
s = time.time()
for i in range(n_runs):
products = sparse_dot_product(Q, K, topk_broadcast.long())
e = time.time()
t_s = (e - s) / n_runs
s = time.time()
for i in range(n_runs):
torch.einsum("nhle,nhse->nhls", Q, K)
e = time.time()
t_f = (e - s) / n_runs
print("Sparse_Clustered: {}, Sparse: {}, Full: {}".format(
t_sc, t_s, t_f))
def test_single_query(self):
X = torch.randn(1, 1, 1, 32).cuda()
Y = torch.randn(1, 1, 100, 32).cuda()
lengths = torch.full((1, ), 1).int().cuda()
topk = (torch.cumsum(torch.rand(1, 1, 1, 10) * 10,
dim=-1)).long().cuda()
products = sparse_dot_product(
X,
Y,
topk,
)
all_products = torch.einsum("nhle,nhse->nhls", X, Y)
self.assertLess(
torch.max(
torch.abs(products.squeeze() -
all_products[0, 0, 0, topk[0, 0, 0]])), 1e-4)
def test_simple_product(self):
X = torch.randn(10, 4, 100, 32)
Y = torch.randn(10, 4, 100, 32)
topk = (torch.cumsum(torch.rand(10, 4, 100, 10) * 10, dim=-1)).long()
products = sparse_dot_product(
X,
Y,
topk,
)
all_products = torch.einsum("nhle,nhse->nhls", X, Y)
self.assertLess(
torch.max(
torch.abs(products -
all_products[torch.arange(10).view(10, 1, 1, 1),
torch.arange(4).view(1, 4, 1, 1),
torch.arange(100).view(1, 1, 100, 1),
topk])), 1e-4)
def test_small_benchmark(self):
N = 12
H = 8
L = 1000
E = 32
S = 1000
k = 32
C = 100
I = 10
B = 32
Q = torch.randn(N, H, L, E).to(self.device)
K = torch.randn(N, H, S, E).to(self.device)
lengths = torch.full((N,), L, dtype=torch.int32).to(self.device)
groups, counts = cluster_queries(Q, lengths, C, I, B)
Q_grouped = aggregate(Q, groups, 1/counts.float())
QK = torch.einsum("nhle,nhse->nhls", Q_grouped, K)
_, topk = torch.topk(QK, k, dim=-1)
topk = topk.contiguous()
products = torch.zeros((N, H, L, k), dtype=torch.float32).to(self.device)
products = clustered_sparse_dot_product(Q, K, topk, groups, counts, lengths)
for i in range(1000):
products = clustered_sparse_dot_product(
Q,
K,
topk,
groups,
counts,
lengths
)
torch.cuda.synchronize()
s = torch.cuda.Event(enable_timing=True)
e = torch.cuda.Event(enable_timing=True)
s.record()
products = clustered_sparse_dot_product(
Q,
K,
topk,
groups,
counts,
lengths
)
e.record()
torch.cuda.synchronize()
t_sc = s.elapsed_time(e)
topk_broadcast = broadcast(
topk.float(),
groups,
torch.ones_like(counts, dtype=torch.float32),
torch.zeros((N, H, L, k), device=Q.device)
)
for i in range(1000):
products = sparse_dot_product(
Q,
K,
topk_broadcast.long()
)
torch.cuda.synchronize()
s = torch.cuda.Event(enable_timing=True)
e = torch.cuda.Event(enable_timing=True)
s.record()
products_s = sparse_dot_product(
Q,
K,
topk_broadcast.long(),
)
e.record()
torch.cuda.synchronize()
t_s = s.elapsed_time(e)
for i in range(1000):
torch.einsum("nhle,nhse->nhls", Q, K)
s = torch.cuda.Event(enable_timing=True)
e = torch.cuda.Event(enable_timing=True)
s.record()
torch.einsum("nhle,nhse->nhls", Q, K)
e.record()
torch.cuda.synchronize()
t_f = s.elapsed_time(e)
print("Sparse_Clustered: {}, Sparse: {}, Full: {}".format(t_sc, t_s, t_f))