From 7d60093554329d82475490a1136c2d63ad18e22f Mon Sep 17 00:00:00 2001 From: xuqianmamba <60861262+xuqianmamba@users.noreply.github.com> Date: Mon, 10 Aug 2026 10:45:07 +0800 Subject: [PATCH 1/2] Add DeepSeek-V4 H20 kernel benchmark data --- .../dsa/prefill/h20/attn-16-64-c128.csv | 2 + bench_data/dsa/prefill/h20/attn-16-64-c4.csv | 2 + .../dsa/prefill/h20/attn-64-64-c128.csv | 3 + bench_data/dsa/prefill/h20/attn-64-64-c4.csv | 3 + bench_data/grouped_gemm/prefill/h20/data.csv | 1 + .../flashmla_dsv4_sparse_prefill.py | 181 ++++++++++++++++++ 6 files changed, 192 insertions(+) create mode 100644 bench_data/dsa/prefill/h20/attn-16-64-c128.csv create mode 100644 bench_data/dsa/prefill/h20/attn-16-64-c4.csv create mode 100644 bench_data/dsa/prefill/h20/attn-64-64-c128.csv create mode 100644 bench_data/dsa/prefill/h20/attn-64-64-c4.csv create mode 100644 kernel_benchmark/flashmla_dsv4_sparse_prefill.py diff --git a/bench_data/dsa/prefill/h20/attn-16-64-c128.csv b/bench_data/dsa/prefill/h20/attn-16-64-c128.csv new file mode 100644 index 0000000..0d83586 --- /dev/null +++ b/bench_data/dsa/prefill/h20/attn-16-64-c128.csv @@ -0,0 +1,2 @@ +dtype,s_q,s_kv,latency_us,mfu +bf16,4096,4128,1163.779,0.12018 diff --git a/bench_data/dsa/prefill/h20/attn-16-64-c4.csv b/bench_data/dsa/prefill/h20/attn-16-64-c4.csv new file mode 100644 index 0000000..37bd60c --- /dev/null +++ b/bench_data/dsa/prefill/h20/attn-16-64-c4.csv @@ -0,0 +1,2 @@ +dtype,s_q,s_kv,latency_us,mfu +bf16,4096,5120,2912.725,0.21342 diff --git a/bench_data/dsa/prefill/h20/attn-64-64-c128.csv b/bench_data/dsa/prefill/h20/attn-64-64-c128.csv new file mode 100644 index 0000000..bf944d5 --- /dev/null +++ b/bench_data/dsa/prefill/h20/attn-64-64-c128.csv @@ -0,0 +1,3 @@ +dtype,s_q,s_kv,latency_us,mfu +bf16,4096,4128,1162.481,0.48127 +bf16,8192,8256,2316.29,0.53502 diff --git a/bench_data/dsa/prefill/h20/attn-64-64-c4.csv b/bench_data/dsa/prefill/h20/attn-64-64-c4.csv new file mode 100644 index 0000000..e3df46f --- /dev/null +++ b/bench_data/dsa/prefill/h20/attn-64-64-c4.csv @@ -0,0 +1,3 @@ +dtype,s_q,s_kv,latency_us,mfu +bf16,4096,5120,2914.016,0.8533 +bf16,8192,10240,7455.167,0.93296 diff --git a/bench_data/grouped_gemm/prefill/h20/data.csv b/bench_data/grouped_gemm/prefill/h20/data.csv index 3cd42c7..c00d520 100644 --- a/bench_data/grouped_gemm/prefill/h20/data.csv +++ b/bench_data/grouped_gemm/prefill/h20/data.csv @@ -135,6 +135,7 @@ num_experts,num_gpus,num_local_experts,topk,hidden_size,intermediate_size,seq_le 256,1,256,8,2048,512,16384,512,2554,0.830889,1235,0.838585 256,1,256,8,2048,512,32768,1024,4727,0.830712,2316,0.841093 256,1,256,8,2048,512,65536,2048,9283,0.832533,4635,0.833505 +256,8,256,6,7168,384,4096,96,4491.317,0.43598,3085.757,0.31729 256,2,256,8,3072,512,4096,128,1218.0,0.819763,570.515,0.910827 512,1,512,10,2048,512,4096,80,1233.0,0.753156,597.824,0.776685 512,1,512,10,2048,512,8192,160,2201.0,0.772143,1023.0,0.820001 diff --git a/kernel_benchmark/flashmla_dsv4_sparse_prefill.py b/kernel_benchmark/flashmla_dsv4_sparse_prefill.py new file mode 100644 index 0000000..d99fa1d --- /dev/null +++ b/kernel_benchmark/flashmla_dsv4_sparse_prefill.py @@ -0,0 +1,181 @@ +"""Benchmark the sparse-prefill FlashMLA API used by DeepSeek-V4. + +The benchmark does not reimplement attention. It constructs a V4-shaped +workspace and calls the same ``flash_mla_sparse_fwd`` and index-combiner APIs +as SGLang's DeepSeek-V4 attention backend. +""" + +import argparse +import csv +import json +import statistics +from pathlib import Path + +CSV_FIELDS = ("dtype", "s_q", "s_kv", "latency_us", "mfu") +ATTENTION_COMPUTE_HEAD_DIM = 64 + + +def load_shape(config_path: Path, tp_size: int) -> dict[str, int]: + if config_path.is_dir(): + config_path = config_path / "config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + num_heads = int(config["num_attention_heads"]) + if num_heads % tp_size: + raise ValueError("num_attention_heads must be divisible by tp_size") + return { + "num_heads": num_heads // tp_size, + # The DSV4 sparse attention core computes 64 elements per head. + # This is deliberately not the wider pre-attention projection size. + "head_dim": ATTENTION_COMPUTE_HEAD_DIM, + "value_dim": ATTENTION_COMPUTE_HEAD_DIM, + "sliding_window": int(config.get("sliding_window", 128)), + "index_topk": int(config.get("index_topk", 1024)), + } + + +def median_latency_us(run, warmup: int, repeats: int) -> float: + import torch + + for _ in range(warmup): + run() + torch.cuda.synchronize() + + samples = [] + for _ in range(repeats): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run() + end.record() + end.synchronize() + samples.append(float(start.elapsed_time(end)) * 1000) + return statistics.median(samples) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config-path", type=Path, required=True) + parser.add_argument("--ratio", type=int, choices=(4, 128), required=True) + parser.add_argument("--q-len", type=int, default=4096) + parser.add_argument("--tp-size", type=int, default=8) + parser.add_argument("--bf16-tflops", type=float, required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--repeats", type=int, default=50) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + import torch + + if min(args.q_len, args.tp_size, args.repeats) <= 0 or args.warmup < 0: + parser.error("q-len, tp-size and repeats must be positive") + if args.bf16_tflops <= 0: + parser.error("bf16-tflops must be positive") + if not torch.cuda.is_available(): + parser.error("CUDA is required") + + # These are the APIs called by DeepseekV4AttnBackend._forward_prefill_sparse. + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import ( + combine_topk_swa_indices, + ) + + shape = load_shape(args.config_path, args.tp_size) + device = torch.device(args.device) + torch.cuda.set_device(device) + torch.manual_seed(0) + + q_len = args.q_len + compressed_len = (q_len + args.ratio - 1) // args.ratio + compressed_topk = ( + min(shape["index_topk"], compressed_len) if args.ratio == 4 else compressed_len + ) + + q = torch.randn( + (q_len, shape["num_heads"], shape["head_dim"]), + dtype=torch.bfloat16, + device=device, + ) + # SGLang dequantizes compressed KV followed by SWA KV into this workspace. + kv = torch.randn( + (compressed_len + q_len, 1, shape["head_dim"]), + dtype=torch.bfloat16, + device=device, + ) + + # Entries beyond each query's causal valid length are ignored by the + # combiner. The deterministic pattern keeps every consumed index valid. + topk_indices = ( + torch.arange(compressed_topk, dtype=torch.int32, device=device) + .unsqueeze(0) + .expand(q_len, -1) + .contiguous() + ) + combined_indices, topk_length = combine_topk_swa_indices( + topk_indices=topk_indices, + query_start_loc=torch.tensor([0, q_len], dtype=torch.int32, device=device), + seq_lens=torch.tensor([q_len], dtype=torch.int32, device=device), + gather_lens=torch.tensor([q_len], dtype=torch.int32, device=device), + compressed_base=torch.zeros(1, dtype=torch.int32, device=device), + swa_base=torch.tensor([compressed_len], dtype=torch.int32, device=device), + window_size=shape["sliding_window"], + compress_ratio=args.ratio, + topk=compressed_topk, + ) + indices = combined_indices.unsqueeze(1) + attn_sink = torch.zeros(shape["num_heads"], dtype=torch.float32, device=device) + + def run(): + return flash_mla_sparse_fwd( + q=q, + kv=kv, + indices=indices, + sm_scale=shape["head_dim"] ** -0.5, + d_v=shape["value_dim"], + attn_sink=attn_sink, + topk_length=topk_length, + ) + + # Trigger JIT compilation/autotuning before collecting samples. + run() + torch.cuda.synchronize() + latency_us = median_latency_us(run, args.warmup, args.repeats) + + mean_keys = float(topk_length.float().mean().item()) + # QK and PV each contribute 2 FLOPs per multiply-accumulate: + # 4 * queries * local_heads * compute_head_dim * effective_keys. + # For attention TP2, local_heads = 128 / 2 = 64. For TP8 it is 16. + flops = ( + 4 + * q_len + * shape["num_heads"] + * ATTENTION_COMPUTE_HEAD_DIM + * mean_keys + ) + achieved_tflops = flops / latency_us / 1e6 + # Follow the existing bench_data attention schema. s_kv is the physical + # KV workspace length. The file name records local heads, the 64-element + # compute dimension and the compression group, for example + # attn-64-64-c4.csv for attention TP2. + row = { + "dtype": "bf16", + "s_q": q_len, + "s_kv": kv.shape[0], + "latency_us": round(latency_us, 3), + "mfu": round(achieved_tflops / args.bf16_tflops, 3), + } + + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", newline="", encoding="utf-8") as output: + writer = csv.DictWriter(output, fieldnames=CSV_FIELDS) + writer.writeheader() + writer.writerow(row) + print( + f"ratio={args.ratio}, mean_keys_per_query={mean_keys:.3f}, " + f"achieved_tflops={achieved_tflops:.6f}" + ) + print(json.dumps(row, indent=2)) + + +if __name__ == "__main__": + main() From eb67f9cd37e46d8f42c5fb778fdebbc6497723b4 Mon Sep 17 00:00:00 2001 From: xuqianmamba <60861262+xuqianmamba@users.noreply.github.com> Date: Mon, 17 Aug 2026 17:27:35 +0800 Subject: [PATCH 2/2] Fix DeepSeek-V4 sparse attention head dimension --- .../dsa/prefill/h20/attn-16-64-c128.csv | 2 -- bench_data/dsa/prefill/h20/attn-16-64-c4.csv | 2 -- .../dsa/prefill/h20/attn-64-64-c128.csv | 3 --- bench_data/dsa/prefill/h20/attn-64-64-c4.csv | 3 --- .../flashmla_dsv4_sparse_prefill.py | 26 +++++++++++++------ 5 files changed, 18 insertions(+), 18 deletions(-) delete mode 100644 bench_data/dsa/prefill/h20/attn-16-64-c128.csv delete mode 100644 bench_data/dsa/prefill/h20/attn-16-64-c4.csv delete mode 100644 bench_data/dsa/prefill/h20/attn-64-64-c128.csv delete mode 100644 bench_data/dsa/prefill/h20/attn-64-64-c4.csv diff --git a/bench_data/dsa/prefill/h20/attn-16-64-c128.csv b/bench_data/dsa/prefill/h20/attn-16-64-c128.csv deleted file mode 100644 index 0d83586..0000000 --- a/bench_data/dsa/prefill/h20/attn-16-64-c128.csv +++ /dev/null @@ -1,2 +0,0 @@ -dtype,s_q,s_kv,latency_us,mfu -bf16,4096,4128,1163.779,0.12018 diff --git a/bench_data/dsa/prefill/h20/attn-16-64-c4.csv b/bench_data/dsa/prefill/h20/attn-16-64-c4.csv deleted file mode 100644 index 37bd60c..0000000 --- a/bench_data/dsa/prefill/h20/attn-16-64-c4.csv +++ /dev/null @@ -1,2 +0,0 @@ -dtype,s_q,s_kv,latency_us,mfu -bf16,4096,5120,2912.725,0.21342 diff --git a/bench_data/dsa/prefill/h20/attn-64-64-c128.csv b/bench_data/dsa/prefill/h20/attn-64-64-c128.csv deleted file mode 100644 index bf944d5..0000000 --- a/bench_data/dsa/prefill/h20/attn-64-64-c128.csv +++ /dev/null @@ -1,3 +0,0 @@ -dtype,s_q,s_kv,latency_us,mfu -bf16,4096,4128,1162.481,0.48127 -bf16,8192,8256,2316.29,0.53502 diff --git a/bench_data/dsa/prefill/h20/attn-64-64-c4.csv b/bench_data/dsa/prefill/h20/attn-64-64-c4.csv deleted file mode 100644 index e3df46f..0000000 --- a/bench_data/dsa/prefill/h20/attn-64-64-c4.csv +++ /dev/null @@ -1,3 +0,0 @@ -dtype,s_q,s_kv,latency_us,mfu -bf16,4096,5120,2914.016,0.8533 -bf16,8192,10240,7455.167,0.93296 diff --git a/kernel_benchmark/flashmla_dsv4_sparse_prefill.py b/kernel_benchmark/flashmla_dsv4_sparse_prefill.py index d99fa1d..0d286ec 100644 --- a/kernel_benchmark/flashmla_dsv4_sparse_prefill.py +++ b/kernel_benchmark/flashmla_dsv4_sparse_prefill.py @@ -12,7 +12,7 @@ from pathlib import Path CSV_FIELDS = ("dtype", "s_q", "s_kv", "latency_us", "mfu") -ATTENTION_COMPUTE_HEAD_DIM = 64 +DSV4_ATTENTION_HEAD_DIM = 512 def load_shape(config_path: Path, tp_size: int) -> dict[str, int]: @@ -20,14 +20,20 @@ def load_shape(config_path: Path, tp_size: int) -> dict[str, int]: config_path = config_path / "config.json" config = json.loads(config_path.read_text(encoding="utf-8")) num_heads = int(config["num_attention_heads"]) + head_dim = int(config["head_dim"]) if num_heads % tp_size: raise ValueError("num_attention_heads must be divisible by tp_size") + if head_dim != DSV4_ATTENTION_HEAD_DIM: + raise ValueError( + "This benchmark targets DeepSeek-V4 with a 512-wide attention " + f"head, but config head_dim is {head_dim}" + ) return { "num_heads": num_heads // tp_size, - # The DSV4 sparse attention core computes 64 elements per head. - # This is deliberately not the wider pre-attention projection size. - "head_dim": ATTENTION_COMPUTE_HEAD_DIM, - "value_dim": ATTENTION_COMPUTE_HEAD_DIM, + # DeepSeek-V4 head_dim is the complete QK width: 448 NoPE + 64 RoPE. + # FlashMLA sparse prefill receives this full width, not only RoPE. + "head_dim": head_dim, + "value_dim": head_dim, "sliding_window": int(config.get("sliding_window", 128)), "index_topk": int(config.get("index_topk", 1024)), } @@ -102,6 +108,10 @@ def main() -> None: dtype=torch.bfloat16, device=device, ) + if q.shape[-1] != DSV4_ATTENTION_HEAD_DIM: + raise AssertionError(f"unexpected Q head dimension: {q.shape[-1]}") + if kv.shape[-1] != DSV4_ATTENTION_HEAD_DIM: + raise AssertionError(f"unexpected KV head dimension: {kv.shape[-1]}") # Entries beyond each query's causal valid length are ignored by the # combiner. The deterministic pattern keeps every consumed index valid. @@ -149,14 +159,14 @@ def run(): 4 * q_len * shape["num_heads"] - * ATTENTION_COMPUTE_HEAD_DIM + * shape["head_dim"] * mean_keys ) achieved_tflops = flops / latency_us / 1e6 # Follow the existing bench_data attention schema. s_kv is the physical - # KV workspace length. The file name records local heads, the 64-element + # KV workspace length. The file name records local heads, the 512-element # compute dimension and the compression group, for example - # attn-64-64-c4.csv for attention TP2. + # attn-64-512-c4.csv for attention TP2. row = { "dtype": "bf16", "s_q": q_len,