[kernel/fix] Performance Optimization for Decoding Kernel and Benchmarking (#5274)

* prevent re-creating intermediate tensors * add singleton class holding intermediate values * fix triton kernel api * add benchmark in pytest * fix kernel api and add benchmark * revise flash decoding triton kernel in/out shapes * fix calling of triton kernel in modeling * fix pytest: extract to util functions
2025-09-01 17:17:05 +00:00 · 2024-01-19 15:47:16 +08:00
parent 9e2342bde2
commit 6e487e7d3c
7 changed files with 382 additions and 152 deletions
--- a/tests/test_infer_ops/triton/kernel_utils.py
+++ b/tests/test_infer_ops/triton/kernel_utils.py
@@ -1,3 +1,5 @@
+from typing import Tuple
+
 import torch
 from torch.nn import functional as F

@@ -17,13 +19,22 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    return hidden_states.reshape(bsz, num_key_value_heads * n_rep, seq_len, head_dim)


+def prepare_padding_mask(kv_lengths: torch.Tensor, bsz: int, kv_seq_len: int, device="cuda"):
+    padding_mask = torch.zeros((bsz, 1, 1, kv_seq_len), dtype=torch.float32, device=device)
+    for i in range(bsz):
+        cur_seq_len = kv_lengths[i].item()
+        assert cur_seq_len <= kv_seq_len
+        padding_mask[i, :, :, : kv_seq_len - cur_seq_len] = float("-inf")
+    return padding_mask
+
+
 # Attention calculation adapted from HuggingFace transformers repository
 # src/transformers/models/llama/modeling_llama.py
 # https://github.com/huggingface/transformers/blob/633215ba58fe5114d8c8d32e415a04600e010701/src/transformers/models/llama/modeling_llama.py#L350
 def torch_attn_ref(
-    q: torch.Tensor,  # [bsz, seq_len, num_heads, head_dim]
-    k: torch.Tensor,  # [bsz, kv_seq_len, num_heads, head_dim]
-    v: torch.Tensor,  # [bsz, kv_seq_len, num_heads, head_dim]
+    q: torch.Tensor,  # [bsz, num_heads, q_len, head_dim]
+    k: torch.Tensor,  # [bsz, num_heads, kv_seq_len, head_dim]
+    v: torch.Tensor,  # [bsz, num_heads, kv_seq_len, head_dim]
    attention_mask: torch.Tensor,  # [bsz, 1, seq_len, kv_seq_len]
    bsz: int,
    seq_len: int,
@@ -31,14 +42,8 @@ def torch_attn_ref(
    num_heads: int,
    num_kv_heads: int,
    head_dim: int,
-):
+) -> torch.Tensor:
    assert q.shape[-1] == k.shape[-1] == v.shape[-1] == head_dim
-    q = q.view(bsz, seq_len, num_heads, head_dim)
-    k = k.view(bsz, kv_seq_len, num_kv_heads, head_dim)
-    v = v.view(bsz, kv_seq_len, num_kv_heads, head_dim)
-    q = q.transpose(1, 2)
-    k = k.transpose(1, 2)
-    v = v.transpose(1, 2)

    # repeat kv for GQA and MQA
    # k/v won't change if kv_group_num is 1
@@ -49,7 +54,6 @@ def torch_attn_ref(

    qk = torch.matmul(q, k.transpose(2, 3))
    attn_scores = qk / (head_dim**0.5)
-
    assert attn_scores.shape == (bsz, num_heads, seq_len, kv_seq_len), "Invalid shape of attention scores"
    # for left-side padding
    if attention_mask.size() != (bsz, 1, seq_len, kv_seq_len):
@@ -77,7 +81,7 @@ def mock_alloc_block_table_and_kvcache(
    num_seqs: int,
    max_num_blocks_per_seq: int,
    block_size: int,
-):
+) -> torch.Tensor:
    """Allocate block tables based on provided context lengths; and copy KV to blocked KV Cache."""
    block_id = 0
    block_tables = torch.full(size=(num_seqs, max_num_blocks_per_seq), fill_value=-1, dtype=torch.int32)
@@ -102,12 +106,10 @@ def mock_alloc_block_table_and_kvcache(
    return block_tables


-def mock_alloc_single_token(block_tables: torch.Tensor, context_lengths: torch.Tensor, block_size: int):
-    """Allocate 1 token on the block table for each seqs in block tables.
-    It won't change provided context_lengths
-    """
-
-    # consider max_block_id as the last physical block allocated
+def mock_alloc_single_token(block_tables: torch.Tensor, context_lengths: torch.Tensor, block_size: int) -> None:
+    # Allocate 1 token on the block table for each seqs in block tables.
+    # It won't change provided context_lengths.
+    # Consider max_block_id as the last physical block allocated
    # NOTE It assumes all the blocks preceding this block have been allocated
    max_block_id = torch.max(block_tables).item()
    # the indices on each block table representing the cache block to be allocated one more token
@@ -126,3 +128,36 @@ def mock_alloc_single_token(block_tables: torch.Tensor, context_lengths: torch.T
    if new_block_ids.numel():
        new_block_alloc_local_indices = alloc_local_block_indices[require_new_block]
        block_tables[require_new_block, new_block_alloc_local_indices] = new_block_ids
+
+
+def generate_caches_and_block_tables(
+    k_unpad, v_unpad, kv_lengths, bsz, max_num_blocks_per_seq, block_size, dtype=torch.float16, device="cuda"
+) -> Tuple[torch.Tensor, ...]:
+    # Mock generation of k/v blocked caches and block tables from providied kv unpad and seq lengths
+    # k_unpad/v_unpad [num_total_tokens, num_kv_heads, head_dim]
+    _, num_kv_heads, head_dim = k_unpad.shape
+    cache_shape = (bsz * max_num_blocks_per_seq, num_kv_heads, head_dim, block_size)
+    k_cache = torch.zeros(size=cache_shape, dtype=dtype, device=device)
+    v_cache = torch.zeros(size=cache_shape, dtype=dtype, device=device)
+    # Mock allocation on block tables as well as blocked kv caches
+    block_tables = mock_alloc_block_table_and_kvcache(
+        k_unpad, v_unpad, k_cache, v_cache, kv_lengths, bsz, max_num_blocks_per_seq, block_size
+    )
+    return k_cache, v_cache, block_tables
+
+
+def convert_kv_unpad_to_padded(
+    k_unpad: torch.Tensor, kv_seq_lengths: torch.Tensor, bsz: int, max_seq_len: int
+) -> torch.Tensor:
+    # Rebuild (batched) k/v with padding to be used by torch attention
+    # input k_unpad/v_unpad [num_total_tokens, num_kv_heads, head_dim]
+    # returns k/v padded    [bsz, num_kv_heads, max_seq_len, head_dim]
+    _, num_kv_heads, head_dim = k_unpad.shape
+    k_torch = torch.zeros((bsz, max_seq_len, num_kv_heads, head_dim), dtype=k_unpad.dtype, device=k_unpad.device)
+    prev_len_sum = 0
+    for i, seq_len in enumerate(kv_seq_lengths.tolist()):
+        # left-side padding
+        k_torch[i, -seq_len:, :, :] = k_unpad[prev_len_sum : prev_len_sum + seq_len]
+        prev_len_sum += seq_len
+    k_torch = k_torch.transpose(1, 2)
+    return k_torch
--- a/tests/test_infer_ops/triton/test_context_attn_unpad.py
+++ b/tests/test_infer_ops/triton/test_context_attn_unpad.py
@@ -4,7 +4,7 @@ from packaging import version

 from colossalai.kernel.triton import context_attention_unpadded
 from colossalai.utils import get_current_device
-from tests.test_infer_ops.triton.kernel_utils import mock_alloc_block_table_and_kvcache, torch_attn_ref
+from tests.test_infer_ops.triton.kernel_utils import generate_caches_and_block_tables, torch_attn_ref

 try:
    import triton  # noqa
@@ -16,6 +16,8 @@ except ImportError:

 TRITON_CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse("11.4")

+HEAD_DIM = 32
+

 def torch_attn_unpad(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, context_lengths: torch.Tensor, num_heads: int, num_kv_heads: int
@@ -34,9 +36,9 @@ def torch_attn_unpad(
        mask[mask == 0.0] = float("-inf")

        torch_attn_ref_out = torch_attn_ref(
-            q[start_idx:end_idx].unsqueeze(0),
-            k[start_idx:end_idx].unsqueeze(0),
-            v[start_idx:end_idx].unsqueeze(0),
+            q[start_idx:end_idx].unsqueeze(0).transpose(1, 2),
+            k[start_idx:end_idx].unsqueeze(0).transpose(1, 2),
+            v[start_idx:end_idx].unsqueeze(0).transpose(1, 2),
            mask,
            1,  # set bsz as 1 as we're processing sequence one by one
            seq_len,
@@ -74,7 +76,6 @@ def test_context_attention(

    num_kv_heads = num_attn_heads // kv_group_num
    assert isinstance(num_kv_heads, int) and num_kv_heads > 0, "Invalid number of kv heads."
-    head_dim = 32
    max_seq_len = max_num_blocks_per_seq * block_size
    dtype = torch.float16
    device = get_current_device()
@@ -85,28 +86,28 @@ def test_context_attention(
        context_lengths = torch.randint(low=1, high=max_seq_len, size=(bsz,), dtype=torch.int32, device=device)
    num_tokens = torch.sum(context_lengths).item()

-    qkv_size = (num_tokens, num_attn_heads + 2 * num_kv_heads, head_dim)
-    qkv = torch.empty(size=qkv_size, dtype=dtype, device=device).normal_(mean=0.0, std=0.5)
-    q, k, v = torch.split(qkv, [num_attn_heads, num_kv_heads, num_kv_heads], dim=-2)
+    qkv_size = (num_tokens, num_attn_heads + 2 * num_kv_heads, HEAD_DIM)
+    qkv_unpad = torch.empty(size=qkv_size, dtype=dtype, device=device).normal_(mean=0.0, std=0.5)
+    q_unpad, k_unpad, v_unpad = torch.split(qkv_unpad, [num_attn_heads, num_kv_heads, num_kv_heads], dim=-2)

-    cache_shape = (bsz * max_num_blocks_per_seq, num_kv_heads, head_dim, block_size)
-    k_cache_torch = torch.zeros(size=cache_shape, dtype=dtype, device=device)
-    k_cache_triton = torch.zeros_like(k_cache_torch)
-    v_cache_torch = torch.zeros(size=cache_shape, dtype=dtype, device=device)
-    v_cache_triton = torch.zeros_like(v_cache_torch)
-
-    # Mock allocation on block tables
-    block_tables = mock_alloc_block_table_and_kvcache(
-        k, v, k_cache_torch, v_cache_torch, context_lengths, bsz, max_num_blocks_per_seq, block_size
+    k_cache_ref, v_cache_ref, block_tables = generate_caches_and_block_tables(
+        k_unpad, v_unpad, context_lengths, bsz, max_num_blocks_per_seq, block_size, dtype, device
    )
    block_tables = block_tables.to(device=device)
+    k_cache_triton = torch.zeros_like(k_cache_ref)
+    v_cache_triton = torch.zeros_like(v_cache_ref)
+
    out_triton = context_attention_unpadded(
-        q, k, v, k_cache_triton, v_cache_triton, context_lengths, block_tables, block_size
+        q_unpad, k_unpad, v_unpad, k_cache_triton, v_cache_triton, context_lengths, block_tables, block_size
    )

-    out_torch = torch_attn_unpad(q, k, v, context_lengths, num_attn_heads, num_kv_heads)
+    out_torch = torch_attn_unpad(q_unpad, k_unpad, v_unpad, context_lengths, num_attn_heads, num_kv_heads)

    assert out_torch.shape == out_triton.shape
-    assert torch.allclose(out_torch, out_triton, atol=1e-3, rtol=1e-4)
-    assert torch.allclose(k_cache_torch, k_cache_triton)
-    assert torch.allclose(v_cache_torch, v_cache_triton)
+    assert torch.allclose(out_torch, out_triton, atol=1e-3)
+    assert torch.equal(k_cache_ref, k_cache_triton)
+    assert torch.equal(v_cache_ref, v_cache_triton)
+
+
+if __name__ == "__main__":
+    test_context_attention(4, 32, 8, 16, 1, True)
--- a/tests/test_infer_ops/triton/test_decoding_attn.py
+++ b/tests/test_infer_ops/triton/test_decoding_attn.py
@@ -2,9 +2,14 @@ import pytest
 import torch
 from packaging import version

-from colossalai.kernel.triton import flash_decoding_fwd
+from colossalai.kernel.triton import flash_decoding_attention
 from colossalai.utils import get_current_device
-from tests.test_infer_ops.triton.kernel_utils import mock_alloc_block_table_and_kvcache, torch_attn_ref
+from tests.test_infer_ops.triton.kernel_utils import (
+    convert_kv_unpad_to_padded,
+    generate_caches_and_block_tables,
+    prepare_padding_mask,
+    torch_attn_ref,
+)

 try:
    import triton  # noqa
@@ -16,23 +21,37 @@ except ImportError:

 TRITON_CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse("11.4")

+Q_LEN = 1
+HEAD_DIM = 128

-def torch_decoding(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, context_lengths: torch.Tensor):
-    assert context_lengths.dim() == 1, "context_lengths should be a 1D tensor"
-    assert q.size(1) == 1, "Only used for decoding"
-    assert k.shape == v.shape

-    bsz, _, num_heads, head_dim = q.shape
-    _, kv_seq_len, num_kv_heads, _ = k.shape
-    assert num_heads % num_kv_heads == 0, "Invalid kv heads and attention heads."
-    padding_mask = torch.zeros((bsz, 1, 1, kv_seq_len), dtype=torch.float32, device=q.device)
-    for i in range(bsz):
-        cur_seq_len = context_lengths[i].item()
-        assert cur_seq_len <= kv_seq_len
-        padding_mask[i, :, :, : kv_seq_len - cur_seq_len] = float("-inf")
+def prepare_data(
+    bsz: int,
+    num_attn_heads: int,
+    num_kv_heads: int,
+    head_dim: int,
+    same_context_len: bool,
+    q_len: int,
+    max_kv_seq_len: int,
+    dtype=torch.float16,
+    device="cuda",
+):
+    # Use the provided maximum sequence length for each sequence when testing with teh same context length,
+    # otherwise generate random context lengths.
+    kv_lengths = (
+        torch.tensor([max_kv_seq_len for _ in range(bsz)], dtype=torch.int32, device=device)
+        if same_context_len
+        else torch.randint(low=1, high=max_kv_seq_len, size=(bsz,), dtype=torch.int32, device=device)
+    )
+    num_tokens = torch.sum(kv_lengths).item()

-    out = torch_attn_ref(q, k, v, padding_mask, bsz, 1, kv_seq_len, num_heads, num_kv_heads, head_dim)
-    return out
+    q_size = (bsz, q_len, num_attn_heads, head_dim)
+    q = torch.empty(size=q_size, dtype=dtype, device=device).normal_(mean=0.0, std=0.5).transpose(1, 2)
+    kv_size = (num_tokens, 2 * num_kv_heads, head_dim)
+    kv_unpad = torch.empty(size=kv_size, dtype=dtype, device=device).normal_(mean=0.0, std=0.5)
+    k_unpad, v_unpad = torch.split(kv_unpad, [num_kv_heads, num_kv_heads], dim=-2)
+
+    return q, k_unpad, v_unpad, kv_lengths


@pytest.mark.skipif(not (HAS_TRITON and TRITON_CUDA_SUPPORT), reason="requires triton")
@@ -57,59 +76,135 @@ def test_flash_decoding(

    num_kv_heads = num_attn_heads // kv_group_num
    assert isinstance(num_kv_heads, int) and num_kv_heads > 0, "Invalid number of kv heads."
-    q_len = 1
-    head_dim = 128
    max_seq_len = block_size * max_num_blocks_per_seq
    dtype = torch.float16
    device = get_current_device()

-    if same_context_len:
-        context_lengths = torch.tensor([max_seq_len for _ in range(bsz)], dtype=torch.int32, device=device)
-    else:
-        context_lengths = torch.randint(low=1, high=max_seq_len, size=(bsz,), dtype=torch.int32, device=device)
-    num_tokens = torch.sum(context_lengths).item()
-
-    q_size = (bsz, q_len, num_attn_heads, head_dim)
-    q = torch.empty(size=q_size, dtype=dtype, device=device).normal_(mean=0.0, std=0.5)
-    kv_size = (num_tokens, 2 * num_kv_heads, head_dim)
-    kv = torch.empty(size=kv_size, dtype=dtype, device=device).normal_(mean=0.0, std=0.5)
-    k, v = torch.split(kv, [num_kv_heads, num_kv_heads], dim=-2)
-
-    cache_shape = (bsz * max_num_blocks_per_seq, num_kv_heads, head_dim, block_size)
-    k_cache = torch.zeros(size=cache_shape, dtype=dtype, device=device)
-    v_cache = torch.zeros(size=cache_shape, dtype=dtype, device=device)
-    # Mock allocation on block tables as well as blocked kv caches
-    block_tables = mock_alloc_block_table_and_kvcache(
-        k, v, k_cache, v_cache, context_lengths, bsz, max_num_blocks_per_seq, block_size
+    q, k_unpad, v_unpad, kv_seq_lengths = prepare_data(
+        bsz, num_attn_heads, num_kv_heads, HEAD_DIM, same_context_len, Q_LEN, max_seq_len, dtype, device
+    )
+    k_cache, v_cache, block_tables = generate_caches_and_block_tables(
+        k_unpad, v_unpad, kv_seq_lengths, bsz, max_num_blocks_per_seq, block_size, dtype, device
    )
    block_tables = block_tables.to(device=device)
-
-    q = q.view(bsz, q_len, num_attn_heads, head_dim)
-    out_triton = flash_decoding_fwd(
+    # The maximum sequence length in the batch (if context lengths randomly generated)
+    max_seq_len_in_b = kv_seq_lengths.max().item()
+    # The maximum block length splitted on kv should be the kv cache block size
+    kv_max_split_num = (max_seq_len_in_b + block_size - 1) // block_size
+    mid_output = torch.empty(
+        size=(bsz, num_attn_heads, kv_max_split_num, HEAD_DIM), dtype=torch.float32, device=q.device
+    )
+    mid_output_lse = torch.empty(size=(bsz, num_attn_heads, kv_max_split_num), dtype=torch.float32, device=q.device)
+    sm_scale = 1.0 / (HEAD_DIM**0.5)
+    out_triton = flash_decoding_attention(
        q,
        k_cache,
        v_cache,
-        context_lengths,
+        kv_seq_lengths,
        block_tables,
        block_size,
-        kv_group_num,
-    )
-    out_triton = out_triton.unsqueeze(1)  # [bsz, 1, num_heads, head_dim]
+        max_seq_len_in_b,
+        mid_output,
+        mid_output_lse,
+        sm_scale=sm_scale,
+        kv_group_num=kv_group_num,
+    )  # [bsz, 1, num_heads, head_dim]

-    # rebuild (batched) kv with padding for torch attention
-    # q   [bsz, 1, num_heads, head_dim]
-    # k/v [num_tokens, num_kv_heads, head_dim]
-    max_seq_len = context_lengths.max().item()
-    k_torch = torch.zeros((bsz, max_seq_len, num_kv_heads, head_dim), dtype=k.dtype, device=k.device)
-    v_torch = torch.zeros_like(k_torch)
-    prev_len_sum = 0
-    for i, seq_len in enumerate(context_lengths.tolist()):
-        # mock left-side padding
-        k_torch[i, -seq_len:, :, :] = k[prev_len_sum : prev_len_sum + seq_len]
-        v_torch[i, -seq_len:, :, :] = v[prev_len_sum : prev_len_sum + seq_len]
-        prev_len_sum += seq_len
-    # k/v [bsz, max_seq_len, num_kv_heads, head_dim]
-    out_torch = torch_decoding(q, k_torch, v_torch, context_lengths)
+    k_torch = convert_kv_unpad_to_padded(k_unpad, kv_seq_lengths, bsz, max_seq_len_in_b)
+    v_torch = convert_kv_unpad_to_padded(v_unpad, kv_seq_lengths, bsz, max_seq_len_in_b)
+    torch_padding_mask = prepare_padding_mask(kv_seq_lengths, bsz, max_seq_len_in_b, q.device)
+    out_torch = torch_attn_ref(
+        q, k_torch, v_torch, torch_padding_mask, bsz, 1, max_seq_len_in_b, num_attn_heads, num_kv_heads, HEAD_DIM
+    )

    assert out_torch.shape == out_triton.shape
    assert torch.allclose(out_torch, out_triton, atol=1e-3, rtol=1e-4)
+
+
+BATCH = 16
+BLOCK_SIZE = 32
+SAME_LEN = True
+WARM_UPS = 10
+REPS = 100
+configs = [
+    triton.testing.Benchmark(
+        x_names=["KV_LEN"],
+        x_vals=[2**i for i in range(8, 14)],
+        # x_vals=[x for x in range(256, 8192, 256)],
+        line_arg="provider",
+        line_vals=["torch", "triton"],
+        line_names=["Torch", "Triton"],
+        styles=[("red", "-"), ("blue", "-")],
+        ylabel="ms",
+        plot_name=f"decoding-block_size-{BLOCK_SIZE}-batch{BATCH}",
+        args={"bsz": BATCH, "block_size": BLOCK_SIZE, "same_context_len": SAME_LEN, "kv_group_num": 1},
+    )
+]
+
+
+@triton.testing.perf_report(configs)
+def bench_kernel(
+    bsz,
+    KV_LEN,
+    provider,
+    block_size: int,
+    kv_group_num: int,
+    same_context_len: bool,
+):
+    num_attn_heads = 16
+    max_num_blocks_per_seq = triton.cdiv(KV_LEN, block_size)
+    max_seq_len = block_size * max_num_blocks_per_seq
+
+    num_kv_heads = num_attn_heads // kv_group_num
+    assert isinstance(num_kv_heads, int) and num_kv_heads > 0, "Invalid number of kv heads."
+    block_size * max_num_blocks_per_seq
+    dtype = torch.float16
+    device = get_current_device()
+
+    q, k_unpad, v_unpad, kv_lengths = prepare_data(
+        bsz, num_attn_heads, num_kv_heads, HEAD_DIM, same_context_len, Q_LEN, max_seq_len, dtype, device
+    )
+    max_seq_len_in_b = kv_lengths.max().item()  # for random lengths
+
+    quantiles = [0.5, 0.2, 0.8]
+    if provider == "torch":
+        k_torch = convert_kv_unpad_to_padded(k_unpad, kv_lengths, bsz, max_seq_len_in_b)
+        v_torch = convert_kv_unpad_to_padded(v_unpad, kv_lengths, bsz, max_seq_len_in_b)
+        torch_padding_mask = prepare_padding_mask(kv_lengths, bsz, max_seq_len_in_b, q.device)
+        fn = lambda: torch_attn_ref(
+            q, k_torch, v_torch, torch_padding_mask, bsz, 1, max_seq_len_in_b, num_attn_heads, num_kv_heads, HEAD_DIM
+        )
+        ms, min_ms, max_ms = triton.testing.do_bench(fn, warmup=WARM_UPS, rep=REPS, quantiles=quantiles)
+    if provider == "triton":
+        k_cache, v_cache, block_tables = generate_caches_and_block_tables(
+            k_unpad, v_unpad, kv_lengths, bsz, max_num_blocks_per_seq, block_size, dtype, device
+        )
+        block_tables = block_tables.to(device=device)
+        # the maximum block length splitted on kv should be the kv cache block size
+        kv_max_split_num = (max_seq_len_in_b + block_size - 1) // block_size
+        mid_output = torch.empty(
+            size=(bsz, num_attn_heads, kv_max_split_num, HEAD_DIM), dtype=torch.float32, device=q.device
+        )
+        mid_output_lse = torch.empty(size=(bsz, num_attn_heads, kv_max_split_num), dtype=torch.float32, device=q.device)
+        sm_scale = 1.0 / (HEAD_DIM**0.5)
+        fn = lambda: flash_decoding_attention(
+            q,
+            k_cache,
+            v_cache,
+            kv_lengths,
+            block_tables,
+            block_size,
+            max_seq_len_in_b,
+            mid_output,
+            mid_output_lse,
+            sm_scale=sm_scale,
+            kv_group_num=kv_group_num,
+        )  # [bsz, 1, num_heads, head_dim]
+        ms, min_ms, max_ms = triton.testing.do_bench(fn, warmup=WARM_UPS, rep=REPS, quantiles=quantiles)
+
+    return ms, min_ms, max_ms
+
+
+if __name__ == "__main__":
+    test_flash_decoding(16, 32, 32, 16, 1, True)
+    # bench_kernel.run(save_path=".", print_data=True)