Specialize batch prefill for paged KV layout

[rlrs]

What this changes

Specialize the JIT batch-prefill build for the page size, KV memory layout, and
page-table layout used by the call. It also adds 2048-token pages to the batch
prefill tests and command-line test runner.

This lets vLLM send a shuffled paged KV cache directly to CK batch prefill
instead of gathering a long prefix into a temporary contiguous buffer.

The signature order is aligned with the registered operator schema so the
generated module and its fake implementation expose the same optional
arguments.

This draft depends on ROCm/composable_kernel#3760. I will update AITER's CK
submodule after that change is available on a compatible CK revision.

Results

MI250X, BF16 MQA, four query heads, one KV head, head dimension 128:

Context	Existing Triton paged prefill	CK batch prefill	Speedup
128K	141.54 ms	83.24 ms	1.70x
256K	284.62 ms	168.19 ms	1.69x

The isolated 2048-page correctness run covered both contiguous and strided
caches:

python op_tests/test_batch_prefill.py \
  -p 2048 -c true -l 0 -d bf16 -s 8192 \
  -q 4 -k 1 -t vllm --kv_layout vectorized \
  --input_dtype bf16 --quant_method none \
  --head_dim 128 --return_lse false

Total: 2, Passed: 2, Skipped: 0

I also ran an end-to-end 128K TP8 serving test with NVIDIA Nemotron 3 Super.
All requests completed successfully, including a 40K prompt that exercises a
paged continuation.

The open-PR search did not find another batch-prefill change for 2048-token
pages.

The change was prepared with AI assistance and reviewed and tested by the
submitter.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4388 --add-label <label>