Rope dim

[JaredforReal]

No description provided.

[Copilot]

Pull request overview

This PR introduces an only_qv mode for Hopper (SM90) FlashAttention forward, plumbing a new flag from the Python interface through the C++/CUDA launch path into the SM90 mainloop so it can skip Q/K work and run the Qv path only (likely to support a “rope dim”/Qv-only score path).

Changes:

• Add only_qv to the forward parameter structs and public torch/Python APIs, and validate it in C++ (only_qv requires q_v).
• Extend the SM90 forward mainloop template with OnlyQv to conditionally skip loading Q/K, skip K pipeline usage, and use Qv GEMM with zero_init when QK is omitted.
• Update forward launch template dispatch to specialize kernels on the new OnlyQv compile-time flag.

Reviewed changes

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

Show a summary per file

File	Description
hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp	Adds OnlyQv template flag and guards Q/K loads, barriers, pipeline usage, and QK GEMMs accordingly.
hopper/flash.h	Adds only_qv to Flash_fwd_params.
hopper/flash_fwd_launch_template.h	Threads OnlyQv into the SM90 collective and adds runtime dispatch on params.only_qv.
hopper/flash_attn_interface.py	Exposes only_qv in Python APIs and forwards it into the extension call; adjusts autograd wrappers.
hopper/flash_api.cpp	Adds only_qv argument, validates it requires q_v, and stores it into params.
hopper/flash_api_torch_lib.cpp	Extends torch library schema to include only_qv.
hopper/flash_api_stable.cpp	Extends stable API to include only_qv for fwd; also adds it to scheduler-metadata signature (currently inconsistent with stable schema/wrapper).

Comments suppressed due to low confidence (1)

hopper/flash_attn_interface.py:450

• FlashAttnVarlenFunc.backward must return one entry per input to forward. After adding only_qv (and with the existing cp_rank / cp_tot_seqused_k args), this return tuple is now too short and will raise at runtime. Add 2 more None entries at the end to match the forward signature.

        return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None

---

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

[simon-veitner-redhat]

Under OnlyQv the kernel never issues the Q, K or K_new TMA copies, but to_underlying_arguments still encodes their descriptors.
This may result in cuTensorMapEncodeTiled to fail if tensor does not met 16 byte multiple requirement of TMA (note that we only pad shape and strides here but pointer stays the same).
This in turn will lead for cutlass under NDEBUG to print a dump and return zeroed descriptor which the kernel than prefetches.
In my micro benchmarks it showed that we roughly incur 50-150 us on the host for H200.

Can we use similar pattern as here:

    using TMA_Qv_ = decltype(make_tma_copy_A_sm90( /* ... */ ));
    using TMA_Qv = std::conditional_t<HasQv, TMA_Qv_, std::nullptr_t>;

        TMA_Qv tma_load_Qv = [&] {
            if constexpr (HasQv) {
                return make_tma_copy_A_sm90(GmemTiledCopyQ{}, mQv, SmemLayoutQv{}, TileShape_MNK_QV{}, ClusterShape{});
            } else {
                return nullptr;
            }
        }();

we can apply that pattern to Q, K and K_new with !OnlyQv as the condition. When OnlyQv is true, make_tma_copy_* is never called for them, so cuTensorMapEncodeTiled is never called, and the host error is gone regardless of what the caller's zero-width tensor looks like. For every other configuration the code is byte-identical to the merge base.

using TMA_Q = std::conditional_t<!OnlyQv, TMA_Q_, std::nullptr_t>;
using TMA_K = std::conditional_t<!OnlyQv, TMA_K_, std::nullptr_t>;

TMA_Q tma_load_Q = [&] {
    if constexpr (!OnlyQv) { return make_tma_copy_A_sm90(GmemTiledCopyQ{}, mQ, SmemLayoutQ{}, TileShape_MNK{}, ClusterShape{}); }
    else { return nullptr; }
}();
// same for tma_load_K (mK) and tma_load_K_new (conditional_return<AppendKV>(mKnew, mK))

What needs to be done

1. Type aliases (mainloop:262-274): rename TMA_Q / TMA_K to TMA_Q_ / TMA_K_ and add the two conditional_t aliases. Params (mainloop:458-461) needs no change.
2. to_underlying_arguments (mainloop:500-521, mainloop:531-539): wrap the three make_tma_copy_* calls in if constexpr (!OnlyQv) lambdas. Delete headdim_qk_tma, shape_Q_tma, shape_K_tma, shape_K_new_tma; mQ, mK, mKnew go back to their pre-PR shapes.
3. prefetch_tma_descriptors (mainloop:613-627): guard the tma_load_Q, tma_load_K and tma_load_K_new prefetches with !OnlyQv. Keep the Qv and V prefetches as they are.
4. Remaining device-side uses, which the compiler will now flag: the Q and K tile setup in load (mainloop:699-723), the body of the load_K lambda (mainloop:800-805), and load_kv_new (mainloop:1606, mainloop:1613, mainloop:1628-1632). Use the Qv tuple lambda for the tiles and if constexpr (!OnlyQv) around the copies.
5. Test: call with fresh torch.empty(s, h, 0) Q and K (not views of a wider tensor), capture stderr, assert it is empty.

[simon-veitner-redhat]

Please add tests for the newly introduced kernel

[simon-veitner-redhat]

Please switch to the new flash_api.cpp, flash_api_stable.cpp can be removed

[simon-veitner-redhat]

Can you update to STD_TORCH_CHECK please

[simon-veitner-redhat]

could you please fix remaining small issue and rebase onto current main? otherwise LGTM

[simon-veitner-redhat]

LGTM @MatthewBonanni could you approve if it looks good to you aswell?

[MatthewBonanni]

LGTM, thanks!