Upload benchmark.py

benchmarks/benchmark.py

@@ -0,0 +1,322 @@
+import math
+import torch
+
+from kernels.benchmark import Benchmark
+
+
+def _cdiv(a, b):
+    return (a + b - 1) // b
+
+
+def _extract_output(result):
+    if isinstance(result, tuple):
+        return result[0]
+    return result
+
+
+def _reference_mla_decode(q, blocked_k, block_table, cache_seqlens, head_dim_v, causal=False):
+    b, s_q, h_q, d = q.size()
+    block_size = blocked_k.size(1)
+    h_kv = blocked_k.size(2)
+
+    out = torch.empty(b, s_q, h_q, head_dim_v, dtype=torch.float32, device=q.device)
+
+    for i in range(b):
+        cur_len = int(cache_seqlens[i].item())
+        num_blocks = _cdiv(cur_len, block_size)
+        cur_blocks = block_table[i][:num_blocks]
+        kv = blocked_k[cur_blocks].reshape(-1, h_kv, d)[:cur_len]
+
+        query = q[i].transpose(0, 1).float()  # [h_q, s_q, d]
+        key_val = kv.transpose(0, 1).float()  # [h_kv, s_k, d]
+
+        if h_kv != h_q:
+            key_val = key_val.repeat_interleave(h_q // h_kv, dim=0)
+
+        attn = query @ key_val.transpose(-2, -1) / math.sqrt(d)
+
+        s_k = key_val.size(1)
+        if causal and s_q > 1:
+            mask = torch.ones(s_q, s_k, dtype=torch.bool, device=q.device).tril(
+                diagonal=s_k - s_q
+            )
+            attn.masked_fill_(~mask, float("-inf"))
+
+        attn = torch.softmax(attn, dim=-1)
+        output = attn @ key_val[..., :head_dim_v]
+        out[i] = output.transpose(0, 1)
+
+    return out.to(q.dtype)
+
+
+def _varlen_reference_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, causal=False):
+    batch_size = cu_seqlens_q.shape[0] - 1
+    total_tokens_q = q.shape[0]
+    num_heads = q.shape[1]
+    head_dim_v = v.shape[2]
+    scale = q.shape[-1] ** (-0.5)
+
+    out = torch.zeros(
+        (total_tokens_q, num_heads, head_dim_v), device=q.device, dtype=q.dtype
+    )
+
+    for b in range(batch_size):
+        start_q, end_q = cu_seqlens_q[b], cu_seqlens_q[b + 1]
+        start_k, end_k = cu_seqlens_k[b], cu_seqlens_k[b + 1]
+
+        q_b = q[start_q:end_q].transpose(0, 1).float()  # [H, seq_q, D_qk]
+        k_b = k[start_k:end_k].transpose(0, 1).float()  # [H, seq_k, D_qk]
+        v_b = v[start_k:end_k].transpose(0, 1).float()  # [H, seq_k, D_v]
+
+        attn = q_b @ k_b.transpose(-2, -1) * scale
+
+        if causal:
+            seq_q, seq_k = q_b.size(1), k_b.size(1)
+            mask = torch.ones(seq_q, seq_k, dtype=torch.bool, device=q.device).tril(
+                diagonal=seq_k - seq_q
+            )
+            attn.masked_fill_(~mask, float("-inf"))
+
+        attn = torch.softmax(attn, dim=-1)
+        result = attn @ v_b  # [H, seq_q, D_v]
+        out[start_q:end_q] = result.transpose(0, 1).to(q.dtype)
+
+    return out
+
+
+# MLA decode constants (DeepSeek V3 architecture)
+_HEAD_DIM = 576  # Q/K head dimension
+_HEAD_DIM_V = 512  # V head dimension
+_NUM_HEADS_K = 1  # MLA uses single KV head
+_PAGE_BLOCK_SIZE = 64  # Page block size
+
+
+def _setup_mla_decode(bench, batch_size, seq_k, num_heads_q):
+    max_num_blocks = _cdiv(seq_k, _PAGE_BLOCK_SIZE)
+    total_blocks = batch_size * max_num_blocks
+
+    bench.q = (
+        torch.randn(
+            batch_size, 1, num_heads_q, _HEAD_DIM, device="cuda", dtype=torch.bfloat16
+        )
+        / 10
+    )
+    bench.blocked_k = (
+        torch.randn(
+            total_blocks,
+            _PAGE_BLOCK_SIZE,
+            _NUM_HEADS_K,
+            _HEAD_DIM,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        / 10
+    )
+    bench.block_table = torch.arange(
+        total_blocks, device="cuda", dtype=torch.int32
+    ).view(batch_size, max_num_blocks)
+    bench.cache_seqlens = torch.full(
+        (batch_size,), seq_k, device="cuda", dtype=torch.int32
+    )
+    bench.tile_scheduler_metadata, _ = bench.kernel.get_mla_metadata()
+    bench.out = torch.empty(
+        batch_size, 1, num_heads_q, _HEAD_DIM_V, device="cuda", dtype=torch.bfloat16
+    )
+
+
+def _run_mla_decode(bench, causal=False):
+    out, lse = bench.kernel.flash_mla_with_kvcache(
+        q=bench.q,
+        k_cache=bench.blocked_k,
+        block_table=bench.block_table,
+        cache_seqlens=bench.cache_seqlens,
+        head_dim_v=_HEAD_DIM_V,
+        tile_scheduler_metadata=bench.tile_scheduler_metadata,
+        causal=causal,
+    )
+    bench.out = out
+
+
+def _verify_mla_decode(bench, causal=False):
+    return _reference_mla_decode(
+        bench.q,
+        bench.blocked_k,
+        bench.block_table,
+        bench.cache_seqlens,
+        _HEAD_DIM_V,
+        causal=causal,
+    )
+
+
+class FlashMLABenchmark(Benchmark):
+    seed: int = 42
+
+    # Workload: small (B=2, S_k=256, H_q=64)
+    def setup_small(self):
+        _setup_mla_decode(self, batch_size=2, seq_k=256, num_heads_q=64)
+
+    def benchmark_small(self):
+        _run_mla_decode(self, causal=False)
+
+    def verify_small(self) -> torch.Tensor:
+        return _verify_mla_decode(self, causal=False)
+
+    # Workload: medium (B=4, S_k=1024, H_q=64)
+    def setup_medium(self):
+        _setup_mla_decode(self, batch_size=4, seq_k=1024, num_heads_q=64)
+
+    def benchmark_medium(self):
+        _run_mla_decode(self, causal=False)
+
+    def verify_medium(self) -> torch.Tensor:
+        return _verify_mla_decode(self, causal=False)
+
+    # Workload: large (B=8, S_k=4096, H_q=128)
+    def setup_large(self):
+        _setup_mla_decode(self, batch_size=8, seq_k=4096, num_heads_q=128)
+
+    def benchmark_large(self):
+        _run_mla_decode(self, causal=False)
+
+    def verify_large(self) -> torch.Tensor:
+        return _verify_mla_decode(self, causal=False)
+
+
+class FlashMLACausalBenchmark(Benchmark):
+    seed: int = 42
+
+    # Workload: small (B=2, S_k=256, H_q=64)
+    def setup_small(self):
+        _setup_mla_decode(self, batch_size=2, seq_k=256, num_heads_q=64)
+
+    def benchmark_small(self):
+        _run_mla_decode(self, causal=True)
+
+    def verify_small(self) -> torch.Tensor:
+        return _verify_mla_decode(self, causal=True)
+
+    # Workload: medium (B=4, S_k=1024, H_q=64)
+    def setup_medium(self):
+        _setup_mla_decode(self, batch_size=4, seq_k=1024, num_heads_q=64)
+
+    def benchmark_medium(self):
+        _run_mla_decode(self, causal=True)
+
+    def verify_medium(self) -> torch.Tensor:
+        return _verify_mla_decode(self, causal=True)
+
+    # Workload: large (B=8, S_k=4096, H_q=128)
+    def setup_large(self):
+        _setup_mla_decode(self, batch_size=8, seq_k=4096, num_heads_q=128)
+
+    def benchmark_large(self):
+        _run_mla_decode(self, causal=True)
+
+    def verify_large(self) -> torch.Tensor:
+        return _verify_mla_decode(self, causal=True)
+
+
+class FlashMLAVarlenBenchmark(Benchmark):
+    seed: int = 42
+
+    # Workload: small (3 sequences, max_seqlen=64)
+    def setup_small(self):
+        H, D = 8, 64
+        seqlens = [32, 48, 64]
+        total = sum(seqlens)
+        self.q = torch.randn(total, H, D, device="cuda", dtype=torch.bfloat16)
+        self.k = torch.randn(total, H, D, device="cuda", dtype=torch.bfloat16)
+        self.v = torch.randn(total, H, D, device="cuda", dtype=torch.bfloat16)
+        self.cu_seqlens = torch.tensor(
+            [0] + list(torch.cumsum(torch.tensor(seqlens), 0)),
+            device="cuda",
+            dtype=torch.int32,
+        )
+        self.max_seqlen = max(seqlens)
+        self.out = torch.empty(total, H, D, device="cuda", dtype=torch.bfloat16)
+
+    def benchmark_small(self):
+        self.out = _extract_output(
+            self.kernel.flash_attn_varlen_func(
+                self.q,
+                self.k,
+                self.v,
+                self.cu_seqlens,
+                self.cu_seqlens,
+                self.max_seqlen,
+                self.max_seqlen,
+            )
+        )
+
+    def verify_small(self) -> torch.Tensor:
+        return _varlen_reference_attention(
+            self.q, self.k, self.v, self.cu_seqlens, self.cu_seqlens, causal=False
+        )
+
+    # Workload: medium (5 sequences, max_seqlen=256)
+    def setup_medium(self):
+        H, D = 16, 64
+        seqlens = [128, 192, 256, 200, 150]
+        total = sum(seqlens)
+        self.q = torch.randn(total, H, D, device="cuda", dtype=torch.bfloat16)
+        self.k = torch.randn(total, H, D, device="cuda", dtype=torch.bfloat16)
+        self.v = torch.randn(total, H, D, device="cuda", dtype=torch.bfloat16)
+        self.cu_seqlens = torch.tensor(
+            [0] + list(torch.cumsum(torch.tensor(seqlens), 0)),
+            device="cuda",
+            dtype=torch.int32,
+        )
+        self.max_seqlen = max(seqlens)
+        self.out = torch.empty(total, H, D, device="cuda", dtype=torch.bfloat16)
+
+    def benchmark_medium(self):
+        self.out = _extract_output(
+            self.kernel.flash_attn_varlen_func(
+                self.q,
+                self.k,
+                self.v,
+                self.cu_seqlens,
+                self.cu_seqlens,
+                self.max_seqlen,
+                self.max_seqlen,
+            )
+        )
+
+    def verify_medium(self) -> torch.Tensor:
+        return _varlen_reference_attention(
+            self.q, self.k, self.v, self.cu_seqlens, self.cu_seqlens, causal=False
+        )
+
+    # Workload: large (8 sequences, max_seqlen=512)
+    def setup_large(self):
+        H, D = 32, 128
+        seqlens = [256, 384, 512, 448, 320, 480, 400, 512]
+        total = sum(seqlens)
+        self.q = torch.randn(total, H, D, device="cuda", dtype=torch.bfloat16)
+        self.k = torch.randn(total, H, D, device="cuda", dtype=torch.bfloat16)
+        self.v = torch.randn(total, H, D, device="cuda", dtype=torch.bfloat16)
+        self.cu_seqlens = torch.tensor(
+            [0] + list(torch.cumsum(torch.tensor(seqlens), 0)),
+            device="cuda",
+            dtype=torch.int32,
+        )
+        self.max_seqlen = max(seqlens)
+        self.out = torch.empty(total, H, D, device="cuda", dtype=torch.bfloat16)
+
+    def benchmark_large(self):
+        self.out = _extract_output(
+            self.kernel.flash_attn_varlen_func(
+                self.q,
+                self.k,
+                self.v,
+                self.cu_seqlens,
+                self.cu_seqlens,
+                self.max_seqlen,
+                self.max_seqlen,
+            )
+        )
+
+    def verify_large(self) -> torch.Tensor:
+        return _varlen_reference_attention(
+            self.q, self.k, self.v, self.cu_seqlens, self.cu_seqlens, causal=False
+        )