Skip to content

benchmark: add HND KV layout support with NHD vs HND comparison table#2

Closed
hlin99 with Copilot wants to merge 8 commits into
ww30_PR_HNDfrom
copilot/modify-benchmark-files-for-kv-layout-support
Closed

benchmark: add HND KV layout support with NHD vs HND comparison table#2
hlin99 with Copilot wants to merge 8 commits into
ww30_PR_HNDfrom
copilot/modify-benchmark-files-for-kv-layout-support

Conversation

Copilot AI commented Jul 6, 2026

Copy link
Copy Markdown

Extends the existing batch decode and varlen benchmarks to support the HND paged KV cache layout introduced in PR vllm-project#449, enabling direct NHD vs HND performance comparison across all existing shapes.

benchmark_cutlass_flash_attn_decode.py

  • make_decode_with_paged_kv_input(config, kv_layout="NHD"): creates KV cache as NHD then .permute(0, 2, 1, 3) for HND — producing non-contiguous strides the kernel auto-detects via is_paged_kv_hnd_layout()
  • benchmark_batch_decode(..., kv_layout="NHD"): forwards layout to input builder
  • BATCH_DECODE_CONFIGS loop: runs each config under both layouts and emits a side-by-side table:
config                                   | batch  kv_sum |   NHD(us)   HND(us)   ratio winner
8xmixed_128-16384_(32,8)                 |     8   36736 |     123.4      98.7   1.250    HND

benchmark_cutlass_flash_attn_varlen.py

  • make_varlen_with_paged_kv_input(config, kv_layout="NHD"): same permute logic, paged path only
  • benchmark_varlen_with_paged_kv(..., kv_layout="NHD"): forwards layout to input builder

Default behavior (kv_layout="NHD") is unchanged; existing callers are unaffected.

Copilot AI added 2 commits July 6, 2026 04:40
- make_decode_with_paged_kv_input: accept kv_layout="NHD"|"HND" parameter;
  HND path creates NHD tensor then permute(0,2,1,3) to produce
  [num_blocks, num_kv_heads, block_size, head_size] with non-contiguous
  strides that the kernel auto-detects via is_paged_kv_hnd_layout().
- benchmark_batch_decode: forward kv_layout to the input builder.
- BATCH_DECODE_CONFIGS loop: run each config under both NHD and HND,
  then print a side-by-side comparison table with latency, ratio, and winner.
- make_varlen_with_paged_kv_input: same kv_layout parameter for the
  paged path in the varlen benchmark.
- benchmark_varlen_with_paged_kv: forward kv_layout (default "NHD") to
  the input builder so callers can request HND timing.

Signed-off-by: GitHub Copilot <copilot@github.com>
Signed-off-by: GitHub Copilot <copilot@github.com>
Copilot AI changed the title [WIP] Update benchmark files to add HND KV layout support benchmark: add HND KV layout support with NHD vs HND comparison table Jul 6, 2026
Copilot finished work on behalf of hlin99 July 6, 2026 04:41
Copilot AI requested a review from hlin99 July 6, 2026 04:41
hlin99 added 5 commits July 6, 2026 12:54
…th bandwidth

Keep the original "Batch Decode Benchmark" section unchanged (NHD-only,
showing time + BW). Add a new "NHD vs HND KV Layout Comparison" section
after it that runs both layouts and shows latency, bandwidth, ratio, and
winner for each config.
The varlen benchmark previously only added kv_layout parameter support
to make_varlen_with_paged_kv_input and benchmark_varlen_with_paged_kv
but never actually ran HND benchmarks. This adds a "Prefill (varlen
paged): NHD vs HND KV Layout Comparison" section at the end that runs
representative paged prefill configs with both layouts and prints a
comparison table with latency, ratio, and winner.
Remove the hardcoded PREFILL_LAYOUT_CONFIGS. Instead, filter paged
configs from the existing gen_perf_configs() and run each with both
NHD and HND layouts for a direct comparison. No new configs introduced.
Use a readable config name (e.g. "B4_(64,8)_h128_blk64_bf16_causal")
with fixed-width columns: config | seqs kv_sum | NHD(us) HND(us) | ratio winner.
This matches the decode benchmark's output format.
…d_kv

The layout comparison table already shows all necessary info per config,
so the redundant multi-line config dump is removed to keep output clean.
@hlin99 hlin99 closed this Jul 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants