[CK] Add FP8 KV_BLOCKSCALE support for batch prefill

[Jeff-Huang]

Implement per-page K/V quantization for paged attention:

• Add KV_BLOCKSCALE enum to BlockAttentionQuantScaleEnum
• Use exp2 shift trick to eliminate explicit P scaling overhead
• Prefetch physical pages offset for KV cache, overlaps with computations

Proposed changes

Please describe the motivation behind the pull request, whether it enables a new feature or fixes a bug. If there are associated pull requests or issues, please link them to the pull request.

Checklist

Please put an x into the boxes that apply. You can also fill these out after creating the PR. If you're not sure, please don't hesitate to ask.

• I have added tests relevant to the introduced functionality, and the unit tests are passing locally
• I have added the test to REGRESSION_TESTS list defined at the top of CMakeLists.txt in tests/CMakeLists.txt, IF the test takes more than 30 seconds to run.
• I have added inline documentation which enables the maintainers with understanding the motivation
• I have removed the stale documentation which is no longer relevant after this pull request
• (If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request
• I have run clang-format on all changed files
• Any dependent changes have been merged

Discussion

If this is a relatively large or complex change, feel free to start a discussion by explaining why you chose the solution you did and what alternatives you considered

[Copilot]

Pull request overview

This pull request implements FP8 KV_BLOCKSCALE support for batch prefill in the FMHA (Fused Multi-Head Attention) pipeline. The implementation adds per-page K/V quantization for paged attention caches, uses an exp2 shift trick to eliminate explicit P scaling overhead, and prefetches physical page offsets to hide memory latency.

Changes:

• Added KV_BLOCKSCALE enum to BlockAttentionQuantScaleEnum for per-page K/V quantization
• Refactored KV cache offset computation by splitting into load_physical_pages() and kv_offset_array_transform()
• Implemented exp2 shift trick to scale attention probabilities by 2^shift without explicit multiplication
• Added prefetching of physical pages to overlap memory loads with computation

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp	Core pipeline implementation: added load_physical_pages(), refactored kv_offset_array_transform(), implemented KV_BLOCKSCALE logic with exp2 shift trick and prefetching
fmha_batch_prefill_kernel.hpp	Added KV_BLOCKSCALE kernel argument structures (FmhaFwdKVBlockScaleKargs), QScaleKargsSelector template, and operator() overloads
block_attention_quant_scale_enum.hpp	Added KV_BLOCKSCALE enum value (3) and string conversion template specialization
quant.hpp	Added kv_blockscale enum value and decode/encode logic for command-line interface
fmha_fwd.hpp	Added kv_block_descale parameters to fmha_batch_prefill_args structure
fmha_batch_prefill.py	Restricted kv_blockscale to page_size=1024 in code generation
cpp_symbol_map.py	Added kv_blockscale mappings to QSCALE_MAP and QSCALE_CHECK_MAP

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The comment says "For KV_BLOCKSCALE" but these shift constants appear to be general FP8 constants that could be used for both BLOCKSCALE and KV_BLOCKSCALE. In other pipeline files (block_fmha_pipeline_qr_ks_vs.hpp and block_fmha_pipeline_qr_ks_vs_async.hpp), the same constants have comments saying "For BLOCKSCALE". Consider making the comment more general (e.g., "For BLOCKSCALE/KV_BLOCKSCALE") or extracting these constants to a shared location to avoid duplication and inconsistency.

Suggested change

	// For KV_BLOCKSCALE: shift value for exp2(x + shift) to scale P to [0, 2^shift]
	// For BLOCKSCALE/KV_BLOCKSCALE: shift value for exp2(x + shift) to scale P to [0, 2^shift]

[Copilot]

The comment "!kVTileCrossesPages" is misleading. This is the else branch for "if constexpr(kIsKcache)", so it handles V cache regardless of whether the V tile crosses pages or not. The comment should be "!kIsKcache" or "V cache" to accurately describe this branch.

Suggested change

	else // !kVTileCrossesPages
	{
	// V cache: use physical_pages[k0] for each token
	// physical_pages was already populated correctly by load_physical_pages(), handling:
	// - page_size=1: page_idx maps token_idx -> physical_page directly
	// - V tile crosses pages: per-token page lookup
	// - V tile in single page: lane0 lookup with broadcast to all lanes
	else // !kIsKcache (V cache)
	{
	// V cache: use physical_pages[k0] for each token
	// physical_pages was already populated correctly by load_physical_pages(), handling:
	// - page_size=1: page_idx maps token_idx -> physical_page directly
	// - V tile crosses pages: per-token page lookup
	- // - V tile in single page: lane0 lookup with broadcast to all lanes

[asleepzzz]

some minor issues

[poyenc]

Please clarify that the pipeline only supports kPageBlockSize==1024 if QScaleEnum ==BlockAttentionQuantScaleEnum::KV_BLOCKSCALE

[ltqin]

Group mode is not used. Should it be deleted, or should we use static_assert to mask it

[Jeff-Huang]

It seems Aiter still uses batch prefill kernel with group mode in the mha_batch_prefill_kernels.cu:

        float t = aiter::mha_batch_prefill(args,
                                           stream_config,
                                           dtype_str,
                                           true, // is_group_mode
                                           mask.type,
                                           bias_type,
                                           has_lse,
                                           qscale_type,
                                           false);

Do you have a specific concern about group mode in batch prefill kernel, or was there another reason you suggested removing it? I'm happy to discuss further if there's a better way to structure this.

[ltqin]

If using group mode, the calculation method is incorrect. It is recommended to disable group mode

[Jeff-Huang]

Could you clarify which part of the calculation you think is incorrect in group mode? I'd like to make sure I understand your concern correctly.

From my understanding, the current implementation should work for both batch mode and group mode:

1. k_physical_pages[0] is the physical page index obtained from page table lookup. The page_idx pointer is already per-batch adjusted:
   • SGLang: page_idx = kv_page_indices + kv_indptr[i_batch]
   • vLLM: page_idx = block_table_ptr + i_batch * batch_stride
2. kv_head_idx is computed as i_nhead / nhead_ratio_qk, which should be the same for both modes.
3. The scale tensor is indexed by physical page (not logical sequence position), so the offset calculation physical_page * nblock_stride + kv_head_idx * nhead_stride should be correct regardless of batch mode or group mode.

Is there a specific scenario or edge case you have in mind where this would produce incorrect results? Happy to discuss further!

[ammallya]

Imported to ROCm/rocm-libraries