[LLVMGPU] Track paged KV attention lowering for NVIDIA GPUs

[weimin023]

Summary

Following the discussion on Discord, I'd like to use this issue to track my investigation into paged KV attention lowering and performance on NVIDIA GPUs.

I've looked through the existing paged-attention work in #21222.
The current LLVMGPU infrastructure and ROCm coverage seem like a helpful starting point for exploring equivalent NVVM coverage.

Rather than trying to cover compiler lowering, CUDA runtime correctness, and performance work in one large change, I plan to work through this in three stages. Depending on the results and review feedback, each stage may be split into one or more focused PRs.

Motivation

Paged KV attention is important for decode-time LLM inference. Unlike regular attention, the KV cache is accessed indirectly through page tables instead of as one contiguous sequence tensor.

Stage 1: NVVM compile-time paged-attention coverage

• NVVM compile-time paged-attention coverage [LLVMGPU] Add NVVM paged attention pipeline test #24770

This stage is compiler-only coverage, so it does not require NVIDIA hardware in CI.

• Add an NVVM VectorDistribute pipeline test for paged decode attention.
• Cover dynamic page tables, paged K/V gathers, causal masking, and online_attention lowering.
• Start with an explicit CUDA lowering configuration and verify that the high-level operations lower to loops and subgroup reductions.

[Ongoing] Stage 2: CUDA runtime correctness and reproducible baseline

• Add CUDA runtime correctness coverage for representative paged-KV shapes.

  • Runtime correctness test [LLVMGPU] Add dense and paged attention runtime correctness tests #24829

• Establish a reproducible benchmark baseline, including the hardware, workload shape, page size, command line, and measurement method.

The goal here is to validate correctness and establish a baseline before making
performance claims or changing the lowering strategy.

Stage 3: Profile-driven NVIDIA / SM120 optimization

• Profile the generated CUDA kernel and focus on bottlenecks shown by the measurements.
• Explore layout, data movement, prefetching, software pipelining, and other NVIDIA-specific improvements where they are justified by profiling.
• Evaluate the existing FP8 mma.sync lowering as an optional path for FP8-compatible workloads.
• Compare candidate changes against the reproducible baseline from Stage 2.

If the scope and available test infrastructure allow, I would also like to explore an e2e lowering path from a Torch-level input and include an
Nsight Compute (ncu) profile of the generated CUDA kernel. I see this as a stretch goal for Stage 3, and it could become a separate follow-up PR if that keeps the optimization work more manageable.

Although the initial focus is the NVIDIA backend, comparable AMD measurements would also be useful where suitable hardware is available.
They could help separate generic paged-attention improvements from NVIDIA-specific ones and provide useful comparison points.

Any feedback or suggestions would be very welcome. Thanks!