@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a3a51d20e7
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 14:30 UTC

[Benchmark] Improvements to attention benchmark script (#37115)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>

Wei Zhao committed Mar 16, 2026 at 18:22 UTC a3a51d20e7d040542118f04f5089c57a27bc7aca
6 files changed +311 -68
benchmarks/attention_benchmarks/benchmark.py
+54 -16
@@ -47,6 +47,8 @@ from common import (
47 is_mla_backend,
48 )
49
50 +from vllm.v1.worker.workspace import init_workspace_manager
51 +
52
53 def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
54 """Run standard attention benchmark (Flash/Triton/FlashInfer)."""
@@ -462,7 +464,7 @@ def main():
464 parser.add_argument(
465 "--batch-specs",
466 nargs="+",
465 - default=["q2k", "8q1s1k"],
467 + default=None,
468 help="Batch specifications using extended grammar",
469 )
470
@@ -478,6 +480,21 @@ def main():
480 parser.add_argument("--repeats", type=int, default=1, help="Repetitions")
481 parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations")
482 parser.add_argument("--profile-memory", action="store_true", help="Profile memory")
483 + parser.add_argument(
484 + "--kv-cache-dtype",
485 + default="auto",
486 + choices=["auto", "fp8"],
487 + help="KV cache dtype: auto or fp8",
488 + )
489 + parser.add_argument(
490 + "--cuda-graphs",
491 + action=argparse.BooleanOptionalAction,
492 + default=True,
493 + help=(
494 + "Launch kernels with CUDA graphs to eliminate CPU overhead"
495 + "in measurements (default: True)"
496 + ),
497 + )
498
499 # Parameter sweep (use YAML config for advanced sweeps)
500 parser.add_argument(
@@ -536,21 +553,24 @@ def main():
553
554 # Batch specs and sizes
555 # Support both explicit batch_specs and generated batch_spec_ranges
539 - if "batch_spec_ranges" in yaml_config:
540 - # Generate batch specs from ranges
541 - generated_specs = generate_batch_specs_from_ranges(
542 - yaml_config["batch_spec_ranges"]
543 - )
544 - # Combine with any explicit batch_specs
545 - if "batch_specs" in yaml_config:
546 - args.batch_specs = yaml_config["batch_specs"] + generated_specs
547 - else:
548 - args.batch_specs = generated_specs
549 - console.print(
550 - f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
551 - )
552 - elif "batch_specs" in yaml_config:
553 - args.batch_specs = yaml_config["batch_specs"]
556 + # CLI --batch-specs takes precedence over YAML when provided.
557 + cli_batch_specs_provided = args.batch_specs is not None
558 + if not cli_batch_specs_provided:
559 + if "batch_spec_ranges" in yaml_config:
560 + # Generate batch specs from ranges
561 + generated_specs = generate_batch_specs_from_ranges(
562 + yaml_config["batch_spec_ranges"]
563 + )
564 + # Combine with any explicit batch_specs
565 + if "batch_specs" in yaml_config:
566 + args.batch_specs = yaml_config["batch_specs"] + generated_specs
567 + else:
568 + args.batch_specs = generated_specs
569 + console.print(
570 + f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
571 + )
572 + elif "batch_specs" in yaml_config:
573 + args.batch_specs = yaml_config["batch_specs"]
574
575 if "batch_sizes" in yaml_config:
576 args.batch_sizes = yaml_config["batch_sizes"]
@@ -575,6 +595,10 @@ def main():
595 args.warmup_iters = yaml_config["warmup_iters"]
596 if "profile_memory" in yaml_config:
597 args.profile_memory = yaml_config["profile_memory"]
598 + if "kv_cache_dtype" in yaml_config:
599 + args.kv_cache_dtype = yaml_config["kv_cache_dtype"]
600 + if "cuda_graphs" in yaml_config:
601 + args.cuda_graphs = yaml_config["cuda_graphs"]
602
603 # Parameter sweep configuration
604 if "parameter_sweep" in yaml_config:
@@ -629,12 +653,18 @@ def main():
653 # Determine backends
654 backends = args.backends or ([args.backend] if args.backend else ["flash"])
655 prefill_backends = getattr(args, "prefill_backends", None)
656 + if not args.batch_specs:
657 + args.batch_specs = ["q2k", "8q1s1k"]
658 console.print(f"Backends: {', '.join(backends)}")
659 if prefill_backends:
660 console.print(f"Prefill backends: {', '.join(prefill_backends)}")
661 console.print(f"Batch specs: {', '.join(args.batch_specs)}")
662 + console.print(f"KV cache dtype: {args.kv_cache_dtype}")
663 + console.print(f"CUDA graphs: {args.cuda_graphs}")
664 console.print()
665
666 + init_workspace_manager(args.device)
667 +
668 # Run benchmarks
669 all_results = []
670
@@ -687,6 +717,8 @@ def main():
717 repeats=args.repeats,
718 warmup_iters=args.warmup_iters,
719 profile_memory=args.profile_memory,
720 + kv_cache_dtype=args.kv_cache_dtype,
721 + use_cuda_graphs=args.cuda_graphs,
722 )
723
724 # Add decode pipeline config
@@ -839,6 +871,8 @@ def main():
871 "repeats": args.repeats,
872 "warmup_iters": args.warmup_iters,
873 "profile_memory": args.profile_memory,
874 + "kv_cache_dtype": args.kv_cache_dtype,
875 + "use_cuda_graphs": args.cuda_graphs,
876 }
877 all_results = run_model_parameter_sweep(
878 backends,
@@ -861,6 +895,8 @@ def main():
895 "repeats": args.repeats,
896 "warmup_iters": args.warmup_iters,
897 "profile_memory": args.profile_memory,
898 + "kv_cache_dtype": args.kv_cache_dtype,
899 + "use_cuda_graphs": args.cuda_graphs,
900 }
901 all_results = run_parameter_sweep(
902 backends, args.batch_specs, base_config_args, args.parameter_sweep, console
@@ -891,6 +927,8 @@ def main():
927 repeats=args.repeats,
928 warmup_iters=args.warmup_iters,
929 profile_memory=args.profile_memory,
930 + kv_cache_dtype=args.kv_cache_dtype,
931 + use_cuda_graphs=args.cuda_graphs,
932 )
933
934 result = run_benchmark(config)
benchmarks/attention_benchmarks/common.py
+5
@@ -213,6 +213,9 @@ class BenchmarkConfig:
213 profile_memory: bool = False
214 use_cuda_graphs: bool = False
215
216 + # "auto" or "fp8"
217 + kv_cache_dtype: str = "auto"
218 +
219 # MLA-specific
220 prefill_backend: str | None = None
221 kv_lora_rank: int | None = None
@@ -369,6 +372,7 @@ class ResultsFormatter:
372 "backend",
373 "batch_spec",
374 "num_layers",
375 + "kv_cache_dtype",
376 "mean_time",
377 "std_time",
378 "throughput",
@@ -382,6 +386,7 @@ class ResultsFormatter:
386 "backend": r.config.backend,
387 "batch_spec": r.config.batch_spec,
388 "num_layers": r.config.num_layers,
389 + "kv_cache_dtype": r.config.kv_cache_dtype,
390 "mean_time": r.mean_time,
391 "std_time": r.std_time,
392 "throughput": r.throughput_tokens_per_sec or 0,
benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml
+3 -3
@@ -30,9 +30,9 @@ batch_specs:
30 - "2q16k_32q1s4k" # 2 very large prefill + 32 decode
31
32 # Context extension + decode
33 - - "2q1kkv2k_16q1s1k" # 2 extend + 16 decode
34 - - "4q2kkv4k_32q1s2k" # 4 extend + 32 decode
35 - - "2q1kkv8k_32q1s2k" # 2 large extend + 32 decode
33 + - "2q1ks2k_16q1s1k" # 2 extend + 16 decode
34 + - "4q2ks4k_32q1s2k" # 4 extend + 32 decode
35 + - "2q1ks8k_32q1s2k" # 2 large extend + 32 decode
36
37 # Explicitly chunked prefill
38 - "q8k" # 8k prefill with chunking hint
benchmarks/attention_benchmarks/configs/mla_sparse_decode.yaml new
+58
@@ -0,0 +1,58 @@
1 +# MLA decode-only benchmark configuration
2 +
3 +model:
4 + name: "deepseek-v3"
5 + num_layers: 60
6 + num_q_heads: 128 # Base value, can be swept for TP simulation
7 + num_kv_heads: 1 # MLA uses single latent KV
8 + head_dim: 576
9 + kv_lora_rank: 512
10 + qk_nope_head_dim: 128
11 + qk_rope_head_dim: 64
12 + v_head_dim: 128
13 + block_size: 128 # CUTLASS MLA and FlashAttn MLA use 128
14 +
15 +# Model parameter sweep: simulate tensor parallelism by varying num_q_heads
16 +# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads
17 +model_parameter_sweep:
18 + param_name: "num_q_heads"
19 + values: [128, 64, 32, 16]
20 + label_format: "{backend}_{value}h"
21 +
22 +batch_specs:
23 + # Small batches, varying sequence lengths
24 + - "16q1s512" # 16 requests, 512 KV cache
25 + - "16q1s1k" # 16 requests, 1k KV cache
26 + - "16q1s2k" # 16 requests, 2k KV cache
27 + - "16q1s4k" # 16 requests, 4k KV cache
28 +
29 + # Medium batches
30 + - "32q1s1k" # 32 requests, 1k KV cache
31 + - "32q1s2k" # 32 requests, 2k KV cache
32 + - "32q1s4k" # 32 requests, 4k KV cache
33 + - "32q1s8k" # 32 requests, 8k KV cache
34 +
35 + # Large batches
36 + - "64q1s1k" # 64 requests, 1k KV cache
37 + - "64q1s2k" # 64 requests, 2k KV cache
38 + - "64q1s4k" # 64 requests, 4k KV cache
39 + - "64q1s8k" # 64 requests, 8k KV cache
40 +
41 + # Very large batches
42 + - "128q1s1k" # 128 requests, 1k KV cache
43 + - "128q1s2k" # 128 requests, 2k KV cache
44 + - "128q1s4k" # 128 requests, 4k KV cache
45 + - "128q1s8k" # 128 requests, 8k KV cache
46 +
47 + # Long context
48 + - "32q1s16k" # 32 requests, 16k KV cache
49 + - "32q1s32k" # 32 requests, 32k KV cache
50 +
51 +backends:
52 + - FLASHMLA_SPARSE
53 + - FLASHINFER_MLA_SPARSE
54 +
55 +device: "cuda:0"
56 +repeats: 100
57 +warmup_iters: 10
58 +profile_memory: true
benchmarks/attention_benchmarks/mla_runner.py
+132 -33
@@ -60,9 +60,11 @@ def create_minimal_vllm_config(
60 model_name: str = "deepseek-v3",
61 block_size: int = 128,
62 max_num_seqs: int = 256,
63 + max_num_batched_tokens: int = 8192,
64 mla_dims: dict | None = None,
65 index_topk: int | None = None,
66 prefill_backend: str | None = None,
67 + kv_cache_dtype: str = "auto",
68 ) -> VllmConfig:
69 """
70 Create minimal VllmConfig for MLA benchmarks.
@@ -149,13 +151,13 @@ def create_minimal_vllm_config(
151 cache_config = CacheConfig(
152 block_size=block_size,
153 gpu_memory_utilization=0.9,
152 - cache_dtype="auto",
154 + cache_dtype=kv_cache_dtype,
155 enable_prefix_caching=False,
156 )
157
158 scheduler_config = SchedulerConfig(
159 max_num_seqs=max_num_seqs,
158 - max_num_batched_tokens=8192,
160 + max_num_batched_tokens=max(max_num_batched_tokens, max_num_seqs),
161 max_model_len=32768,
162 is_encoder_decoder=False,
163 enable_chunked_prefill=True,
@@ -535,6 +537,7 @@ def _create_backend_impl(
537 device: torch.device,
538 max_num_tokens: int = 8192,
539 index_topk: int | None = None,
540 + kv_cache_dtype: str = "auto",
541 ):
542 """
543 Create backend implementation instance.
@@ -583,7 +586,7 @@ def _create_backend_impl(
586 "num_kv_heads": mla_dims["num_kv_heads"],
587 "alibi_slopes": None,
588 "sliding_window": None,
586 - "kv_cache_dtype": "auto",
589 + "kv_cache_dtype": kv_cache_dtype,
590 "logits_soft_cap": None,
591 "attn_type": "decoder",
592 "kv_sharing_target_layer_name": None,
@@ -701,6 +704,7 @@ def _run_single_benchmark(
704 mla_dims: dict,
705 device: torch.device,
706 indexer=None,
707 + kv_cache_dtype: str | None = None,
708 ) -> BenchmarkResult:
709 """
710 Run a single benchmark iteration.
@@ -734,49 +738,124 @@ def _run_single_benchmark(
738 )
739
740 # Create KV cache
737 - kv_cache = torch.zeros(
738 - num_blocks,
739 - block_size,
740 - mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"],
741 - device=device,
742 - dtype=torch.bfloat16,
743 - )
741 + if kv_cache_dtype is None:
742 + kv_cache_dtype = getattr(config, "kv_cache_dtype", "auto")
743 + head_size = mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"]
744 + if kv_cache_dtype == "fp8_ds_mla":
745 + # FlashMLA sparse custom format: 656 bytes per token, stored as uint8.
746 + # Layout: kv_lora_rank fp8 bytes + 4 float32 tile scales
747 + # + 2*rope_dim bf16 bytes
748 + # = 512 + 16 + 128 = 656 bytes for DeepSeek dims.
749 + kv_cache = torch.zeros(
750 + num_blocks,
751 + block_size,
752 + 656,
753 + device=device,
754 + dtype=torch.uint8,
755 + )
756 + elif kv_cache_dtype == "fp8":
757 + from vllm.platforms import current_platform
758
745 - # Create input tensors for both decode and prefill modes
746 - decode_inputs, prefill_inputs = _create_input_tensors(
747 - total_q,
748 - mla_dims,
749 - backend_cfg["query_format"],
750 - device,
751 - torch.bfloat16,
752 - )
759 + kv_cache = torch.zeros(
760 + num_blocks,
761 + block_size,
762 + head_size,
763 + device=device,
764 + dtype=torch.uint8,
765 + ).view(current_platform.fp8_dtype())
766 + else:
767 + kv_cache = torch.zeros(
768 + num_blocks,
769 + block_size,
770 + head_size,
771 + device=device,
772 + dtype=torch.bfloat16,
773 + )
774
775 # Fill indexer with random indices for sparse backends
776 is_sparse = backend_cfg.get("is_sparse", False)
777 if is_sparse and indexer is not None:
778 indexer.fill_random_indices(total_q, max_kv_len)
779
759 - # Determine which forward method to use based on metadata
760 - if metadata.decode is not None:
761 - forward_fn = lambda: impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)
762 - elif metadata.prefill is not None:
763 - forward_fn = lambda: impl.forward_mha(
764 - prefill_inputs["q"],
765 - prefill_inputs["k_c_normed"],
766 - prefill_inputs["k_pe"],
767 - kv_cache,
768 - metadata,
769 - prefill_inputs["k_scale"],
770 - prefill_inputs["output"],
771 - )
772 - else:
780 + # Determine which forward methods to use based on metadata.
781 + # Sparse MLA backends always use forward_mqa
782 + has_decode = is_sparse or getattr(metadata, "decode", None) is not None
783 + has_prefill = not is_sparse and getattr(metadata, "prefill", None) is not None
784 + if not has_decode and not has_prefill:
785 raise RuntimeError("Metadata has neither decode nor prefill metadata")
786
787 + num_decode = (
788 + metadata.num_decode_tokens
789 + if (has_decode and has_prefill)
790 + else total_q
791 + if has_decode
792 + else 0
793 + )
794 + num_prefill = total_q - num_decode
795 +
796 + # Some backends requires fp8 queries when using fp8 KV cache.
797 + is_fp8_kvcache = kv_cache_dtype.startswith("fp8")
798 + quantize_query = is_fp8_kvcache and getattr(
799 + impl, "supports_quant_query_input", False
800 + )
801 +
802 + # quantize_query forces concat format
803 + query_fmt = "concat" if quantize_query else backend_cfg["query_format"]
804 +
805 + # Create decode query tensors
806 + if has_decode:
807 + decode_inputs, _ = _create_input_tensors(
808 + num_decode, mla_dims, query_fmt, device, torch.bfloat16
809 + )
810 + # Cast decode query to fp8 if the backend supports it
811 + if quantize_query:
812 + from vllm.platforms import current_platform
813 +
814 + if isinstance(decode_inputs, tuple):
815 + decode_inputs = torch.cat(list(decode_inputs), dim=-1)
816 + decode_inputs = decode_inputs.to(current_platform.fp8_dtype())
817 +
818 + # Create prefill input tensors
819 + if has_prefill:
820 + _, prefill_inputs = _create_input_tensors(
821 + num_prefill, mla_dims, query_fmt, device, torch.bfloat16
822 + )
823 +
824 + # Build forward function
825 + def forward_fn():
826 + results = []
827 + if has_decode:
828 + results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer))
829 + if has_prefill:
830 + results.append(
831 + impl.forward_mha(
832 + prefill_inputs["q"],
833 + prefill_inputs["k_c_normed"],
834 + prefill_inputs["k_pe"],
835 + kv_cache,
836 + metadata,
837 + prefill_inputs["k_scale"],
838 + prefill_inputs["output"],
839 + )
840 + )
841 + return results[0] if len(results) == 1 else tuple(results)
842 +
843 # Warmup
844 for _ in range(config.warmup_iters):
845 forward_fn()
846 torch.accelerator.synchronize()
847
848 + # Optionally capture a CUDA graph after warmup.
849 + # Graph replay eliminates CPU launch overhead so timings reflect pure
850 + # kernel time.
851 + if config.use_cuda_graphs:
852 + graph = torch.cuda.CUDAGraph()
853 + with torch.cuda.graph(graph):
854 + forward_fn()
855 + benchmark_fn = graph.replay
856 + else:
857 + benchmark_fn = forward_fn
858 +
859 # Benchmark
860 times = []
861 for _ in range(config.repeats):
@@ -785,7 +864,7 @@ def _run_single_benchmark(
864
865 start.record()
866 for _ in range(config.num_layers):
788 - forward_fn()
867 + benchmark_fn()
868 end.record()
869
870 torch.accelerator.synchronize()
@@ -852,13 +931,30 @@ def _run_mla_benchmark_batched(
931 # Determine if this is a sparse backend
932 is_sparse = backend_cfg.get("is_sparse", False)
933
934 + # Extract kv_cache_dtype from the first config
935 + kv_cache_dtype = getattr(first_config, "kv_cache_dtype", "auto")
936 +
937 + # FlashMLA sparse only supports "fp8_ds_mla" internally (not generic "fp8").
938 + # Remap here so the user can pass --kv-cache-dtype fp8 regardless of backend.
939 + if backend.upper() == "FLASHMLA_SPARSE" and kv_cache_dtype == "fp8":
940 + kv_cache_dtype = "fp8_ds_mla"
941 +
942 + # Compute max total_q across all configs so the metadata builder buffer
943 + # and scheduler config are large enough for all batch specs.
944 + max_total_q = max(
945 + sum(r.q_len for r in parse_batch_spec(cfg.batch_spec))
946 + for cfg, *_ in configs_with_params
947 + )
948 +
949 # Create and set vLLM config for MLA (reused across all benchmarks)
950 vllm_config = create_minimal_vllm_config(
951 model_name="deepseek-v3", # Used only for model path
952 block_size=block_size,
953 + max_num_batched_tokens=max_total_q,
954 mla_dims=mla_dims, # Use custom dims from config or default
955 index_topk=index_topk if is_sparse else None,
956 prefill_backend=prefill_backend,
957 + kv_cache_dtype=kv_cache_dtype,
958 )
959
960 results = []
@@ -883,7 +979,9 @@ def _run_mla_benchmark_batched(
979 mla_dims,
980 vllm_config,
981 device,
982 + max_num_tokens=max_total_q,
983 index_topk=index_topk if is_sparse else None,
984 + kv_cache_dtype=kv_cache_dtype,
985 )
986
987 # Verify the actual prefill backend matches what was requested
@@ -942,6 +1040,7 @@ def _run_mla_benchmark_batched(
1040 mla_dims,
1041 device,
1042 indexer=indexer,
1043 + kv_cache_dtype=kv_cache_dtype,
1044 )
1045 results.append(result)
1046
benchmarks/attention_benchmarks/runner.py
+59 -16
@@ -140,7 +140,7 @@ def _create_vllm_config(
140
141 cache_config = CacheConfig(
142 block_size=config.block_size,
143 - cache_dtype="auto",
143 + cache_dtype=config.kv_cache_dtype,
144 )
145 cache_config.num_gpu_blocks = max_num_blocks
146 cache_config.num_cpu_blocks = 0
@@ -215,7 +215,7 @@ def _create_backend_impl(
215 num_kv_heads=config.num_kv_heads,
216 alibi_slopes=None,
217 sliding_window=None,
218 - kv_cache_dtype="auto",
218 + kv_cache_dtype=config.kv_cache_dtype,
219 )
220
221 kv_cache_spec = FullAttentionSpec(
@@ -288,12 +288,22 @@ def _create_input_tensors(
288 total_q: int,
289 device: torch.device,
290 dtype: torch.dtype,
291 + quantize_query: bool = False,
292 ) -> tuple:
292 - """Create Q, K, V input tensors for all layers."""
293 + """Create Q, K, V input tensors for all layers.
294 +
295 + When quantize_query is True, queries are cast to fp8 to match backends
296 + that require query/key/value dtype consistency.
297 + """
298 + q_dtype = dtype
299 + if quantize_query:
300 + from vllm.platforms import current_platform
301 +
302 + q_dtype = current_platform.fp8_dtype()
303 q_list = [
304 torch.randn(
305 total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
296 - )
306 + ).to(q_dtype)
307 for _ in range(config.num_layers)
308 ]
309 k_list = [
@@ -344,10 +354,17 @@ def _create_kv_cache(
354 # Compute inverse permutation to get back to logical view
355 inv_order = [stride_order.index(i) for i in range(len(stride_order))]
356
357 + # Use fp8 dtype for cache when requested.
358 + cache_dtype = dtype
359 + if config.kv_cache_dtype == "fp8":
360 + from vllm.platforms import current_platform
361 +
362 + cache_dtype = current_platform.fp8_dtype()
363 +
364 cache_list = []
365 for _ in range(config.num_layers):
366 # Allocate in physical layout order (contiguous in memory)
350 - cache = torch.zeros(*physical_shape, device=device, dtype=dtype)
367 + cache = torch.zeros(*physical_shape, device=device, dtype=cache_dtype)
368 # Permute to logical view
369 cache = cache.permute(*inv_order)
370 cache_list.append(cache)
@@ -392,6 +409,37 @@ def _run_single_benchmark(
409 )
410 torch.accelerator.synchronize()
411
412 + # Optionally capture a CUDA graph after warmup.
413 + # Graph replay eliminates CPU launch overhead so timings reflect pure
414 + # kernel time.
415 + if config.use_cuda_graphs:
416 + graph = torch.cuda.CUDAGraph()
417 + with torch.cuda.graph(graph):
418 + for i in range(config.num_layers):
419 + impl.forward(
420 + layer,
421 + q_list[i],
422 + k_list[i],
423 + v_list[i],
424 + cache_list[i],
425 + attn_metadata,
426 + output=out,
427 + )
428 + benchmark_fn = graph.replay
429 + else:
430 +
431 + def benchmark_fn():
432 + for i in range(config.num_layers):
433 + impl.forward(
434 + layer,
435 + q_list[i],
436 + k_list[i],
437 + v_list[i],
438 + cache_list[i],
439 + attn_metadata,
440 + output=out,
441 + )
442 +
443 # Benchmark
444 times = []
445 for _ in range(config.repeats):
@@ -399,16 +447,7 @@ def _run_single_benchmark(
447 end = torch.cuda.Event(enable_timing=True)
448
449 start.record()
402 - for i in range(config.num_layers):
403 - impl.forward(
404 - layer,
405 - q_list[i],
406 - k_list[i],
407 - v_list[i],
408 - cache_list[i],
409 - attn_metadata,
410 - output=out,
411 - )
450 + benchmark_fn()
451 end.record()
452
453 torch.accelerator.synchronize()
@@ -502,8 +541,12 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
541 common_attn_metadata=common_metadata,
542 )
543
544 + # Only quantize queries when the impl supports it
545 + quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr(
546 + impl, "supports_quant_query_input", False
547 + )
548 q_list, k_list, v_list = _create_input_tensors(
506 - config, total_q, device, dtype
549 + config, total_q, device, dtype, quantize_query=quantize_query
550 )
551
552 cache_list = _create_kv_cache(