/*
 * SPDX-FileCopyrightText: Copyright (c) 2011-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: NVIDIA TensorRT Source Code License Agreement
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */

#pragma once
#include <fmha/traits.h>
#include <fused_multihead_attention.h>

namespace fmha
{
namespace v2
{

////////////////////////////////////////////////////////////////////////////////////////////////////

template <int USE_LDGSTS>
struct Ldgsts_helper
{
    template <typename This, typename Smem_tile, int LDGS>
    static inline __device__ void load(
        This* this_, Smem_tile& smem_tile, void const* (&ptrs)[LDGS], uint32_t (&preds)[LDGS])
    {
        fmha::pack_predicates(this_->preds_, preds);
        smem_tile.store(ptrs, this_->preds_);
    }
};

////////////////////////////////////////////////////////////////////////////////////////////////////

template <>
struct Ldgsts_helper<0>
{
    template <typename This, typename Smem_tile, int LDGS>
    static inline __device__ void load(
        This* this_, Smem_tile& smem_tile, void const* (&ptrs)[LDGS], uint32_t (&preds)[LDGS])
    {
#if 0
        fmha::pack_predicates(this_->preds_, preds);
        fmha::ldg(this_->fetch_, ptrs, this_->preds_);
#else
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            this_->fetch_[ii] = make_uint4(0u, 0u, 0u, 0u);
        }
        // not packing predicates removes restrictions (e.g. FP16 384, 4 warps)
        Ldg_functor<uint4, LDGS> fct(this_->fetch_, ptrs);
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            fct.ldgsts(ii, preds[ii]);
        }
#endif
    }
};

////////////////////////////////////////////////////////////////////////////////////////////////////

template <
    // The instruction traits.
    typename Traits,
    // The dimensions of the tile computed by the CTA.
    typename Cta_tile,
    // The number of bits per element.
    int BITS_PER_ELEMENT_,
    // The number of rows of Q, K or V loaded by this tile.
    int ROWS_,
    // The number of columns (padded, e.g 64).
    int COLS,
    // The actual number of columns (unpadded, e.g 40)
    int VALID_COLS_,
    // Do we use LDGSTS?
    bool USE_LDGSTS_,
    // Are attention heads interleaved?
    bool HEADS_INTERLEAVED,
    // The number of matrices
    int NUM_MATS = 3,
    // Is sliding window attention used ?
    bool SLIDING_WINDOW_ATTENTION = false>
struct Gmem_tile_qkv
{

    // The size of each LDG.
    enum
    {
        BYTES_PER_LDG = 16
    };

    // The number of bits/bytes of element
    enum
    {
        BITS_PER_ELEMENT = BITS_PER_ELEMENT_
    };

    enum
    {
        BYTES_PER_ELEMENT = BITS_PER_ELEMENT_ / 8
    };

    // The size of a row in bytes.
    enum
    {
        BYTES_PER_ROW = COLS * BITS_PER_ELEMENT / 8
    };

    // The number of threads to load a "row" of the matrix.
    enum
    {
        THREADS_PER_ROW = BYTES_PER_ROW / BYTES_PER_LDG
    };

    // The valid size of a row in bytes (without paddings).
    enum
    {
        VALID_COLS = VALID_COLS_
    };

    // The amount of bytes that are valid per row.
    enum
    {
        VALID_BYTES_PER_ROW = VALID_COLS * BITS_PER_ELEMENT / 8
    };

    // The number of "rows" loaded per LDG.
    enum
    {
        ROWS_PER_LDG = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW
    };

    // The number of rows.
    enum
    {
        ROWS = ROWS_
    };

    // The number of LDGs needed to load a chunk of the Q matrix.
    enum
    {
        LDGS = fmha::Div_up<ROWS, ROWS_PER_LDG>::VALUE
    };

    // The number of predicate registers.
    enum
    {
        PRED_REGS = fmha::Compute_number_of_pred_regs<LDGS>::VALUE
    };

    // Is it Hopper?
    enum
    {
        IS_HOPPER = std::is_same<typename Traits::Gpu_arch, typename fmha::Hopper>::value == true
    };

    // Make sure we use a single register to store predicates. Do not throw for Hopper for now.
    static_assert(!USE_LDGSTS_ || PRED_REGS == 1 || IS_HOPPER, "");

    // We do not use LDGSTS (for the moment).
    enum
    {
        USE_LDGSTS = USE_LDGSTS_
    };

    // Ctor for bert::Fused_multihead_attention_params_v2 class
    template <typename Block_info>
    inline __device__ Gmem_tile_qkv(bert::Fused_multihead_attention_params_v2 const& params, int qkv_offset,
        Block_info const& binfo, int tidx, int cta_row_offset = 0, int cta_col_offset_in_bytes = 0)
        : Gmem_tile_qkv(params.qkv_ptr, params.q_stride_in_bytes, params.d, params.dv, params.h, qkv_offset, binfo,
            tidx, params.h_kv, cta_row_offset, cta_col_offset_in_bytes)
    {
    }

    // Ctor for other param classes (such as Qkv_params in train_ops)
    template <typename Params, typename Block_info>
    inline __device__ Gmem_tile_qkv(Params const& params, int qkv_offset, Block_info const& binfo, int tidx,
        int cta_row_offset = 0, int cta_col_offset_in_bytes = 0)
        : Gmem_tile_qkv(params.qkv_ptr, params.q_stride_in_bytes, params.d, params.dv, params.h, qkv_offset, binfo,
            tidx, cta_row_offset, cta_col_offset_in_bytes)
    {
    }

    // Ctor.
    template <typename Block_info>
    inline __device__ Gmem_tile_qkv(void* qkv_ptr, size_t qkv_stride_in_bytes, int d, int dv, int num_heads,
        int qkv_offset, Block_info const& binfo, int tidx, int num_kv_heads = 0, int cta_row_offset = 0,
        int cta_col_offset_in_bytes = 0)
        : params_qkv_stride_in_bytes_(qkv_stride_in_bytes)
        , actual_seqlen_(binfo.actual_seqlen)
        , qkv_ptr_(reinterpret_cast<char*>(qkv_ptr))
    {

        // Compute the position in the sequence (within the CTA for the moment).
        int row = tidx / THREADS_PER_ROW;
        // Compute the position of the thread in the row.
        int col = tidx % THREADS_PER_ROW;

        // We must store the value to update the predicates in "load".
        row_ = row;
        // Do not load/store if the thread is in the padded area
        col_in_bytes_ = cta_col_offset_in_bytes + col * BYTES_PER_LDG;

        // The row offset in the batched GEMM. For each seq element, we store QKV in that order.
        int64_t row_offset = (int64_t) (row + cta_row_offset) * params_qkv_stride_in_bytes_;
        // Add the byte index.
        int64_t idx;

        // Both MQA and GQA will use non HEADS_INTERLEAVED layout
        if (num_kv_heads < num_heads)
        {
            int const head_id = binfo.bidh;
            int const kv_head_id = binfo.bidh / (num_heads / num_kv_heads);
            // QKV layout [b, s, [q_hd, k_h'd, v_h'd]]
            idx = binfo.sum_s * params_qkv_stride_in_bytes_;
            if (qkv_offset == 0)
            { // Q tensor
                idx += head_id * VALID_BYTES_PER_ROW;
            }
            else if (qkv_offset == 1)
            { // K tensor
                idx += (num_heads + kv_head_id) * VALID_BYTES_PER_ROW;
            }
            else if (qkv_offset == 2)
            { // V tensor
                /*  When qkv_offset == 2, this is an instance of Gmem_tile_v defined in Kernel_traits:
                        using Gmem_tile_v = Gmem_tile_v_<Traits_o,
                                Cta_tile_o,
                                Traits_o::BITS_PER_ELEMENT_B,
                                CTA_O_TILE_K,
                                CTA_O_TILE_N,
                                VALID_DV,   // instead of VALID_D
                                USE_LDGSTS_V,
                                HEADS_INTERLEAVED,
                                3, // NUM_MATS
                                SLIDING_WINDOW_ATTENTION>;
                    the 6th template argument is VALID_DV instead of VALID_D.
                    Thus, here VALID_COLS equals VALID_DV, and
                    VALID_BYTES_PER_ROW equals VALID_DV * BYTES_PER_ELEMENT,
                    and `kv_head_id * dv * BYTES_PER_ELEMENT` can be optimized to
                    `kv_head_id * VALID_BYTES_PER_ROW`. */
                idx += (num_heads + num_kv_heads) * d * BYTES_PER_ELEMENT + kv_head_id * VALID_BYTES_PER_ROW;
            }
        }
        else if (HEADS_INTERLEAVED)
        {
            // [b, s, h, [q_d, k_d, v_d]] aka bsh3d
            // bidx = sum_s * params.h + bidh;
            idx = (binfo.bidx * (2 * d + dv) + qkv_offset * d) * BYTES_PER_ELEMENT;
        }
        else
        {
            // [b, s, [q_hd, k_hd, v_hd]] aka bs3hd
            idx = binfo.sum_s * params_qkv_stride_in_bytes_ + qkv_offset * num_heads * d * BYTES_PER_ELEMENT
                + binfo.bidh * VALID_BYTES_PER_ROW;
        }

        // Assemble the final pointer.
        qkv_ptr_ += row_offset + idx + col_in_bytes_;

        // Take the CTA offset to modify the sequence length.
        actual_seqlen_ -= cta_row_offset;

        // Set the initial seq_len and qkv_offset in case of reinterating
        actual_seqlen_init_ = actual_seqlen_;
        qkv_ptr_init_ = qkv_ptr_;
    }

    // Store data to shared memory.
    template <typename Smem_tile>
    inline __device__ void commit(Smem_tile& smem_tile)
    {
        if (!USE_LDGSTS)
        {
            smem_tile.store(fetch_);
        }
    }

    // Load data from memory.
    template <typename Smem_tile>
    inline __device__ void load(Smem_tile& smem_tile)
    {
        uint32_t preds[LDGS];
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            preds[ii] = row_ + ii * (int) ROWS_PER_LDG < min((int) ROWS, actual_seqlen_);
            preds[ii] &= col_in_bytes_ < VALID_BYTES_PER_ROW;
        }

        // Prepare the load pointers.
        void const* ptrs[LDGS];
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            ptrs[ii] = qkv_ptr_ + (int64_t) ii * ROWS_PER_LDG * params_qkv_stride_in_bytes_;
        }

        // Trigger LDGSTS or the LDGs.
        // The predicates protect against out-of-bound access in rows and cols
        Ldgsts_helper<USE_LDGSTS>::load(this, smem_tile, ptrs, preds);
    }

    // Load data from memory.
    inline __device__ void load()
    {
        uint32_t preds[LDGS];
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            preds[ii] = row_ + ii * (int) ROWS_PER_LDG < min((int) ROWS, actual_seqlen_);
        }

        // Prepare the load pointers.
        void const* ptrs[LDGS];
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            ptrs[ii] = qkv_ptr_ + (int64_t) ii * ROWS_PER_LDG * params_qkv_stride_in_bytes_;
        }

        // Trigger the LDGs.
        if (col_in_bytes_ < VALID_BYTES_PER_ROW)
        {
            fmha::pack_predicates(preds_, preds);
            fmha::ldg(fetch_, ptrs, preds_);
        }
        else
        {
#pragma unroll
            for (int ii = 0; ii < LDGS; ++ii)
            {
                fetch_[ii] = make_uint4(0u, 0u, 0u, 0u);
            }
        }
    }

    // Move the pointer to the next row location.
    inline __device__ void move(int const steps = 1)
    {
        qkv_ptr_ += (int64_t) ROWS * params_qkv_stride_in_bytes_ * steps;
        actual_seqlen_ -= (int) ROWS * steps;
    }

    // Move the pointer to the next row location by the offset (not step).
    inline __device__ void move_by_offset(int const offset)
    {
        qkv_ptr_ = qkv_ptr_init_ + (int64_t) offset * params_qkv_stride_in_bytes_;
        actual_seqlen_ = actual_seqlen_init_ - (int) offset;
    }

    // Move the pointer to the next column location
    inline __device__ void move_col(int const steps = 1)
    {
        qkv_ptr_ += (int64_t) COLS * (BITS_PER_ELEMENT / 8) * steps;
        // Update col_in_bytes_ to ensure load predicates work
        col_in_bytes_ += THREADS_PER_ROW * BYTES_PER_LDG * steps;
    }

    inline __device__ void reset()
    {
        qkv_ptr_ = qkv_ptr_init_;
        actual_seqlen_ = actual_seqlen_init_;
    }

    // Rewind the pointer back to previous column location
    inline __device__ void rewind_col(int const steps)
    {
        qkv_ptr_ -= COLS * (BITS_PER_ELEMENT / 8) * steps;
        // Update col_in_bytes_ to ensure load predicates work
        col_in_bytes_ -= THREADS_PER_ROW * BYTES_PER_LDG * steps;
    }

    inline __device__ void move_to(int const step)
    {
        qkv_ptr_ = qkv_ptr_init_ + (int64_t) ROWS * params_qkv_stride_in_bytes_ * step;
        actual_seqlen_ = actual_seqlen_init_ - (int) ROWS * step;
    }

    // Store data to memory.
    inline __device__ void store(uint4 const (&data)[LDGS])
    {
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            char* ptr = qkv_ptr_ + (int64_t) ii * ROWS_PER_LDG * params_qkv_stride_in_bytes_;
            if (((row_ + ii * ROWS_PER_LDG) < min(ROWS, actual_seqlen_))
                && col_in_bytes_ < VALID_BYTES_PER_ROW /*TODO: double check*/)
            {
                fmha::stg(ptr, data[ii]);
            }
        }
    }

    // The stride between rows for the QKV matrice.
    int64_t params_qkv_stride_in_bytes_;
    // The pointer.
    char* qkv_ptr_;
    char* qkv_ptr_init_;
    // The register to store predicates.
    uint32_t preds_[PRED_REGS];
    // The fetch registers.
    uint4 fetch_[LDGS];
    // Keep track of the row and col the thread is processing as we move the tile.
    int row_;
    int col_in_bytes_;
    // The sequence length.
    int actual_seqlen_;
    int actual_seqlen_init_;
};

////////////////////////////////////////////////////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////////////////////////////////////

// We expect the Q layout to be [B, S, H, D] with variable sequence length support.
template <
    // The instruction traits.
    typename Traits,
    // The dimensions of the tile computed by the CTA.
    typename Cta_tile,
    // The number of bits per element.
    int BITS_PER_ELEMENT_,
    // The number of rows of Q, K or V loaded by this tile.
    int ROWS_,
    // The number of columns (padded, e.g 64).
    int COLS,
    // The actual number of columns (unpadded, e.g 40)
    int VALID_COLS_,
    // Do we use LDGSTS?
    bool USE_LDGSTS_,
    // Are attention heads interleaved? (not used)
    bool HEADS_INTERLEAVED = false,
    // The number of matrices (not used)
    int NUM_MATS = 1,
    // Is sliding window attention used ?
    bool SLIDING_WINDOW_ATTENTION = false>
struct Gmem_tile_q
{

    // The size of each LDG.
    enum
    {
        BYTES_PER_LDG = 16
    };

    // The number of bits/bytes of element
    enum
    {
        BITS_PER_ELEMENT = BITS_PER_ELEMENT_
    };

    enum
    {
        BYTES_PER_ELEMENT = BITS_PER_ELEMENT_ / 8
    };

    // The size of a row in bytes.
    enum
    {
        BYTES_PER_ROW = COLS * BITS_PER_ELEMENT / 8
    };

    // The number of threads to load a "row" of the matrix.
    enum
    {
        THREADS_PER_ROW = BYTES_PER_ROW / BYTES_PER_LDG
    };

    // The valid size of a row in bytes (without paddings).
    enum
    {
        VALID_COLS = VALID_COLS_
    };

    // The amount of bytes that are valid per row.
    enum
    {
        VALID_BYTES_PER_ROW = VALID_COLS * BITS_PER_ELEMENT / 8
    };

    // The number of "rows" loaded per LDG.
    enum
    {
        ROWS_PER_LDG = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW
    };

    // The number of rows.
    enum
    {
        ROWS = ROWS_
    };

    // The number of LDGs needed to load a chunk of the Q matrix.
    enum
    {
        LDGS = fmha::Div_up<ROWS, ROWS_PER_LDG>::VALUE
    };

    // The number of predicate registers.
    enum
    {
        PRED_REGS = fmha::Compute_number_of_pred_regs<LDGS>::VALUE
    };

    // Is it Hopper?
    enum
    {
        IS_HOPPER = std::is_same<typename Traits::Gpu_arch, typename fmha::Hopper>::value == true
    };

    // Make sure we use a single register to store predicates. Do not throw for Hopper for now.
    static_assert(!USE_LDGSTS_ || PRED_REGS == 1 || IS_HOPPER, "");

    // We do not use LDGSTS (for the moment).
    enum
    {
        USE_LDGSTS = USE_LDGSTS_
    };

    // Ctor (keep qkv_offset for compatibility)
    template <typename Block_info>
    inline __device__ Gmem_tile_q(bert::Fused_multihead_attention_params_v2 const& params, int qkv_offset,
        Block_info const& binfo, int tidx, int cta_row_offset = 0, int cta_col_offset_in_bytes = 0)
        : Gmem_tile_q(params, binfo, tidx, cta_row_offset, cta_col_offset_in_bytes)
    {
    }

    // Ctor.
    template <typename Block_info>
    inline __device__ Gmem_tile_q(bert::Fused_multihead_attention_params_v2 const& params, Block_info const& binfo,
        int tidx, int cta_row_offset = 0, int cta_col_offset_in_bytes = 0)
        : params_q_stride_in_bytes_(params.q_stride_in_bytes)
        , actual_seqlen_(binfo.actual_q_seqlen)
        , q_ptr_(reinterpret_cast<char*>(params.q_ptr))
    {

        // Compute the position in the sequence (within the CTA for the moment).
        int row = tidx / THREADS_PER_ROW;
        // Compute the position of the thread in the row.
        int col = tidx % THREADS_PER_ROW;

        // We must store the value to update the predicates in "load".
        row_ = row;
        // Do not load/store if the thread is in the padded area
        col_in_bytes_ = cta_col_offset_in_bytes + col * BYTES_PER_LDG;

        // The row offset in the batched GEMM. For each seq element, we store QKV in that order.
        // We won't consider past_q_length when loading from gmem_q.
        int64_t row_offset = (int64_t) (row + cta_row_offset) * params_q_stride_in_bytes_;
        // Add the block index. (sum_s * h + hidx).
        int64_t idx = binfo.bidx;

        // Assemble the final pointer.
        q_ptr_ += row_offset + idx * VALID_BYTES_PER_ROW + col_in_bytes_;

        // Take the CTA offset to modify the sequence length.
        actual_seqlen_ -= cta_row_offset;
    }

    // Store data to shared memory.
    template <typename Smem_tile>
    inline __device__ void commit(Smem_tile& smem_tile)
    {
        if (!USE_LDGSTS)
        {
            smem_tile.store(fetch_);
        }
    }

    // Load data from memory.
    template <typename Smem_tile>
    inline __device__ void load(Smem_tile& smem_tile)
    {
        uint32_t preds[LDGS];
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            preds[ii] = row_ + ii * (int) ROWS_PER_LDG < min((int) ROWS, actual_seqlen_);
            preds[ii] &= col_in_bytes_ < VALID_BYTES_PER_ROW;
        }

        // Prepare the load pointers.
        void const* ptrs[LDGS];
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            ptrs[ii] = q_ptr_ + (int64_t) ii * ROWS_PER_LDG * params_q_stride_in_bytes_;
        }

        // Trigger LDGSTS or the LDGs.
        // The predicates protect against out-of-bound access in rows and cols
        Ldgsts_helper<USE_LDGSTS>::load(this, smem_tile, ptrs, preds);
    }

    // Move the pointer to the next column location
    inline __device__ void move_col()
    {
        q_ptr_ += (int64_t) COLS * (BITS_PER_ELEMENT / 8);
        // Update col_in_bytes_ to ensure load predicates work
        col_in_bytes_ += THREADS_PER_ROW * BYTES_PER_LDG;
    }

    // Rewind the pointer back to previous column location
    inline __device__ void rewind_col(int const steps)
    {
        q_ptr_ -= COLS * (BITS_PER_ELEMENT / 8) * steps;
        // Update col_in_bytes_ to ensure load predicates work
        col_in_bytes_ -= THREADS_PER_ROW * BYTES_PER_LDG * steps;
    }

    // The stride between rows for the QKV matrice.
    int64_t params_q_stride_in_bytes_;
    // The pointer.
    char* q_ptr_;
    // The register to store predicates.
    uint32_t preds_[PRED_REGS];
    // The fetch registers.
    uint4 fetch_[LDGS];
    // Keep track of the row and col the thread is processing as we move the tile.
    int row_;
    int64_t col_in_bytes_;
    // The sequence length.
    int actual_seqlen_;
};

////////////////////////////////////////////////////////////////////////////////////////////////////

// Shape [B, S, 2, H, D] where S can be variable sequence length.
template <
    // The instruction traits.
    typename Traits,
    // The dimensions of the tile computed by the CTA.
    typename Cta_tile,
    // The number of bits per element.
    int BITS_PER_ELEMENT_,
    // The number of rows of Q, K or V loaded by this tile.
    int ROWS_,
    // The number of columns (padded, e.g 64).
    int COLS,
    // The actual number of columns (unpadded, e.g 40)
    int VALID_COLS_,
    // Do we use LDGSTS?
    bool USE_LDGSTS_,
    // Are attention heads interleaved? (Not used)
    bool HEADS_INTERLEAVED,
    // The number of matrices (Not used)
    int NUM_MATS = 2,
    // Is sliding window attention used ?
    bool SLIDING_WINDOW_ATTENTION = false>
struct Gmem_tile_contiguous_kv
{

    // The size of each LDG.
    enum
    {
        BYTES_PER_LDG = 16
    };

    // The number of bits/bytes of element
    enum
    {
        BITS_PER_ELEMENT = BITS_PER_ELEMENT_
    };

    enum
    {
        BYTES_PER_ELEMENT = BITS_PER_ELEMENT_ / 8
    };

    // The size of a row in bytes.
    enum
    {
        BYTES_PER_ROW = COLS * BITS_PER_ELEMENT / 8
    };

    // The number of threads to load a "row" of the matrix.
    enum
    {
        THREADS_PER_ROW = BYTES_PER_ROW / BYTES_PER_LDG
    };

    // The valid size of a row in bytes (without paddings).
    enum
    {
        VALID_COLS = VALID_COLS_
    };

    // The amount of bytes that are valid per row.
    enum
    {
        VALID_BYTES_PER_ROW = VALID_COLS * BITS_PER_ELEMENT / 8
    };

    // The number of "rows" loaded per LDG.
    enum
    {
        ROWS_PER_LDG = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW
    };

    // The number of rows.
    enum
    {
        ROWS = ROWS_
    };

    // The number of LDGs needed to load a chunk of the Q matrix.
    enum
    {
        LDGS = fmha::Div_up<ROWS, ROWS_PER_LDG>::VALUE
    };

    // The number of predicate registers.
    enum
    {
        PRED_REGS = fmha::Compute_number_of_pred_regs<LDGS>::VALUE
    };

    // Is it Hopper?
    enum
    {
        IS_HOPPER = std::is_same<typename Traits::Gpu_arch, typename fmha::Hopper>::value == true
    };

    // Make sure we use a single register to store predicates. Do not throw for Hopper for now.
    static_assert(!USE_LDGSTS_ || PRED_REGS == 1 || IS_HOPPER, "");

    // We do not use LDGSTS (for the moment).
    enum
    {
        USE_LDGSTS = USE_LDGSTS_
    };

    // Ctor for bert::Fused_multihead_attention_params_v2 class
    template <typename Block_info>
    inline __device__ Gmem_tile_contiguous_kv(bert::Fused_multihead_attention_params_v2 const& params,
        int qkv_offset, // q = 0, k = 1, v = 2.
        Block_info const& binfo, int tidx, int cta_row_offset = 0, int cta_col_offset_in_bytes = 0)
        : Gmem_tile_contiguous_kv(params.kv_ptr, params.k_stride_in_bytes, params.h_kv, params.h_q_per_kv, qkv_offset,
            binfo, tidx, cta_row_offset, cta_col_offset_in_bytes)
    {
    }

    // Ctor.
    template <typename Block_info>
    inline __device__ Gmem_tile_contiguous_kv(void* kv_ptr, size_t kv_stride_in_bytes, int num_kv_heads,
        int head_group_size, int qkv_offset, Block_info const& binfo, int tidx, int cta_row_offset = 0,
        int cta_col_offset_in_bytes = 0)
        : params_kv_stride_in_bytes_(kv_stride_in_bytes)
        , actual_seqlen_(binfo.actual_kv_seqlen)
        , kv_ptr_(reinterpret_cast<char*>(kv_ptr))
    {

        // Compute the position in the sequence (within the CTA for the moment).
        int row = tidx / THREADS_PER_ROW;
        // Compute the position of the thread in the row.
        int col = tidx % THREADS_PER_ROW;

        // We must store the value to update the predicates in "load".
        row_ = row;
        // Do not load/store if the thread is in the padded area
        col_in_bytes_ = cta_col_offset_in_bytes + col * BYTES_PER_LDG;

        // The row offset in the batched GEMM.
        int64_t row_offset = (int64_t) (row + cta_row_offset) * params_kv_stride_in_bytes_;
        // [b, s, 2, h_kv, d].
        int64_t idx = (binfo.sum_s_kv * 2 + qkv_offset - 1) * num_kv_heads + (binfo.bidh / head_group_size);

        // Assemble the final pointer.
        kv_ptr_ += row_offset + idx * VALID_BYTES_PER_ROW + col_in_bytes_;

        // Take the CTA offset to modify the sequence length.
        actual_seqlen_ -= cta_row_offset;

        // Set the initial seq_len and qkv_offset in case of reinterating
        actual_seqlen_init_ = actual_seqlen_;
        kv_ptr_init_ = kv_ptr_;
    }

    // Store data to shared memory.
    template <typename Smem_tile>
    inline __device__ void commit(Smem_tile& smem_tile)
    {
        if (!USE_LDGSTS)
        {
            smem_tile.store(fetch_);
        }
    }

    // Load data from memory.
    template <typename Smem_tile>
    inline __device__ void load(Smem_tile& smem_tile)
    {
        uint32_t preds[LDGS];
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            preds[ii] = row_ + ii * (int) ROWS_PER_LDG < min((int) ROWS, actual_seqlen_);
            preds[ii] &= col_in_bytes_ < VALID_BYTES_PER_ROW;
        }

        // Prepare the load pointers.
        void const* ptrs[LDGS];
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            ptrs[ii] = kv_ptr_ + (int64_t) ii * ROWS_PER_LDG * params_kv_stride_in_bytes_;
        }

        // Trigger LDGSTS or the LDGs.
        // The predicates protect against out-of-bound access in rows and cols
        Ldgsts_helper<USE_LDGSTS>::load(this, smem_tile, ptrs, preds);
    }

    // Load data from memory.
    inline __device__ void load()
    {
        uint32_t preds[LDGS];
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            preds[ii] = row_ + ii * (int) ROWS_PER_LDG < min((int) ROWS, actual_seqlen_);
        }

        // Prepare the load pointers.
        void const* ptrs[LDGS];
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            ptrs[ii] = kv_ptr_ + (int64_t) ii * ROWS_PER_LDG * params_kv_stride_in_bytes_;
        }

        // Trigger the LDGs.
        if (col_in_bytes_ < VALID_BYTES_PER_ROW)
        {
            fmha::pack_predicates(preds_, preds);
            fmha::ldg(fetch_, ptrs, preds_);
        }
        else
        {
#pragma unroll
            for (int ii = 0; ii < LDGS; ++ii)
            {
                fetch_[ii] = make_uint4(0u, 0u, 0u, 0u);
            }
        }
    }

    // Move the pointer to the next row location.
    inline __device__ void move(int const steps = 1)
    {
        kv_ptr_ += (int64_t) ROWS * params_kv_stride_in_bytes_ * steps;
        actual_seqlen_ -= (int) ROWS * steps;
    }

    // Move the pointer to the next row location by the offset (not step).
    inline __device__ void move_by_offset(int const offset)
    {
        kv_ptr_ = kv_ptr_init_ + (int64_t) offset * params_kv_stride_in_bytes_;
        actual_seqlen_ = actual_seqlen_init_ - (int) offset;
    }

    // Move the pointer to the next column location
    inline __device__ void move_col(int const steps = 1)
    {
        kv_ptr_ += (int64_t) COLS * (BITS_PER_ELEMENT / 8) * steps;
        // Update col_in_bytes_ to ensure load predicates work
        col_in_bytes_ += THREADS_PER_ROW * BYTES_PER_LDG * steps;
    }

    inline __device__ void reset()
    {
        kv_ptr_ = kv_ptr_init_;
        actual_seqlen_ = actual_seqlen_init_;
    }

    // Rewind the pointer back to previous column location
    inline __device__ void rewind_col(int const steps)
    {
        kv_ptr_ -= COLS * (BITS_PER_ELEMENT / 8) * steps;
        // Update col_in_bytes_ to ensure load predicates work
        col_in_bytes_ -= THREADS_PER_ROW * BYTES_PER_LDG * steps;
    }

    inline __device__ void move_to(int const step)
    {
        kv_ptr_ = kv_ptr_init_ + (int64_t) ROWS * params_kv_stride_in_bytes_ * step;
        actual_seqlen_ = actual_seqlen_init_ - (int) ROWS * step;
    }

    // Store data to memory.
    inline __device__ void store(uint4 const (&data)[LDGS])
    {
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            char* ptr = kv_ptr_ + (int64_t) ii * ROWS_PER_LDG * params_kv_stride_in_bytes_;
            if (((row_ + ii * ROWS_PER_LDG) < min(ROWS, actual_seqlen_))
                && col_in_bytes_ < VALID_BYTES_PER_ROW /*TODO: double check*/)
            {
                fmha::stg(ptr, data[ii]);
            }
        }
    }

    // The stride between rows for the QKV matrice.
    int64_t params_kv_stride_in_bytes_;
    // The pointer.
    char* kv_ptr_;
    char* kv_ptr_init_;
    // The register to store predicates.
    uint32_t preds_[PRED_REGS];
    // The fetch registers.
    uint4 fetch_[LDGS];
    // Keep track of the row and col the thread is processing as we move the tile.
    int row_;
    int col_in_bytes_;
    // The sequence length.
    int actual_seqlen_;
    int actual_seqlen_init_;
};

////////////////////////////////////////////////////////////////////////////////////////////////////

// We expect the paged KV layout to be blocks of indices with shape of [B, 2, Blocks_per_Seq],
// and the indice tells the memory distance to the pool ptr in global memory.

template <
    // The instruction traits.
    typename Traits,
    // The dimensions of the tile computed by the CTA.
    typename Cta_tile,
    // The number of bits per element.
    int BITS_PER_ELEMENT_,
    // The number of rows of Q, K or V loaded by this tile.
    int ROWS_,
    // The number of columns (padded, e.g 64).
    int COLS,
    // The actual number of columns (unpadded, e.g 40)
    int VALID_COLS_,
    // Do we use LDGSTS?
    bool USE_LDGSTS_,
    // Are attention heads interleaved? (not used)
    bool HEADS_INTERLEAVED = false,
    // The number of matrices (not used)
    int NUM_MATS = 2,
    // Is sliding window attention used ?
    bool SLIDING_WINDOW_ATTENTION_ = false>
struct Gmem_tile_paged_kv
{

    // The size of each LDG.
    enum
    {
        BYTES_PER_LDG = 16
    };

    // The number of bits/bytes of element
    enum
    {
        BITS_PER_ELEMENT = BITS_PER_ELEMENT_
    };

    enum
    {
        BYTES_PER_ELEMENT = BITS_PER_ELEMENT_ / 8
    };

    // The size of a row in bytes.
    enum
    {
        BYTES_PER_ROW = COLS * BITS_PER_ELEMENT / 8
    };

    // The number of threads to load a "row" of the matrix.
    enum
    {
        THREADS_PER_ROW = BYTES_PER_ROW / BYTES_PER_LDG
    };

    // The valid size of a row in bytes (without paddings).
    enum
    {
        VALID_COLS = VALID_COLS_
    };

    // The amount of bytes that are valid per row.
    enum
    {
        VALID_BYTES_PER_ROW = VALID_COLS * BITS_PER_ELEMENT / 8
    };

    // The number of "rows" loaded per LDG.
    enum
    {
        ROWS_PER_LDG = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW
    };

    // The number of rows.
    enum
    {
        ROWS = ROWS_
    };

    // The number of LDGs needed to load a chunk of the Q matrix.
    enum
    {
        LDGS = fmha::Div_up<ROWS, ROWS_PER_LDG>::VALUE
    };

    // The number of predicate registers.
    enum
    {
        PRED_REGS = fmha::Compute_number_of_pred_regs<LDGS>::VALUE
    };

    // Is sliding window attention used ?
    enum
    {
        SLIDING_WINDOW_ATTENTION = SLIDING_WINDOW_ATTENTION_
    };

    // Is it Hopper?
    enum
    {
        IS_HOPPER = std::is_same<typename Traits::Gpu_arch, typename fmha::Hopper>::value == true
    };

    // Make sure we use a single register to store predicates. Do not throw for Hopper for now.
    static_assert(!USE_LDGSTS_ || PRED_REGS == 1 || IS_HOPPER, "");

    // We do not use LDGSTS (for the moment).
    enum
    {
        USE_LDGSTS = USE_LDGSTS_
    };

    // Ctor.
    template <typename Block_info>
    inline __device__ Gmem_tile_paged_kv(bert::Fused_multihead_attention_params_v2 const& params,
        int qkv_offset, // q = 0, k = 1, v = 2.
        Block_info const& binfo, int tidx, int cta_row_offset = 0, int cta_col_offset_in_bytes = 0)
        : actual_seqlen_(binfo.actual_seqlen)
        , past_seqlen_(binfo.actual_seqlen - binfo.actual_q_seqlen)
        , sliding_window_size_(params.sliding_window_size)
        , paged_kv_log2_block_size_(params.paged_kv_cache.mTokensPerBlockLog2)
        , paged_kv_block_pool_ptr_(reinterpret_cast<char*>(params.paged_kv_cache.mPoolPtr))
        , paged_kv_global_block_offsets_(params.paged_kv_cache.mBlockOffsets)
        , params_kv_block_size_in_bytes_(params.paged_kv_cache.mBytesPerBlock)
    {

        // Handle Paged KV with shape [S, Dh], by offsetting it to the target batch.
        int32_t const paged_kv_block_offset
            = (binfo.bidb * 2 + qkv_offset - 1) * params.paged_kv_cache.mMaxBlocksPerSeq;
        paged_kv_global_block_offsets_ += paged_kv_block_offset;

        // Compute the position in the sequence (within the CTA for the moment).
        int row = tidx / THREADS_PER_ROW;
        // Compute the position of the thread in the row.
        int col = tidx % THREADS_PER_ROW;

        // We must store the value to update the predicates in "load".
        row_ = row;
        // Do not load/store if the thread is in the padded area
        col_in_bytes_ = cta_col_offset_in_bytes + col * BYTES_PER_LDG;

        int64_t kv_stride_in_bytes = qkv_offset == 1 ? params.k_stride_in_bytes : params.v_stride_in_bytes;
        // The head offset.
        head_stride_in_bytes_ = (int64_t) (binfo.bidh / params.h_q_per_kv) * kv_stride_in_bytes;
        // When V is padded (like MLA), we cannot use VALID_BYTES_PER_ROW
        token_stride_in_bytes_ = kv_stride_in_bytes >> paged_kv_log2_block_size_;

        // Take the CTA offset to modify the sequence length.
        // Actually we don't need that for flash attention.
        actual_seqlen_ -= cta_row_offset;
    }

    // Store data to shared memory.
    template <typename Smem_tile>
    inline __device__ void commit(Smem_tile& smem_tile)
    {
        if (!USE_LDGSTS)
        {
            smem_tile.store(fetch_);
        }
    }

    // Load data from memory.
    template <typename Smem_tile>
    inline __device__ void load(Smem_tile& smem_tile)
    {
        // Prepare the predicates.
        uint32_t preds[LDGS];
        // Prepare the load pointers.
        void const* ptrs[LDGS];

        // Offset for the new paged kv pointer.
        uint64_t const head_col_in_bytes = head_stride_in_bytes_ + col_in_bytes_;

// Update paged_kv ptr for each LDG (reuse is possible).
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            int row_idx = row_ + ii * (int) ROWS_PER_LDG;
            int paged_kv_block_idx = (row_idx >> paged_kv_log2_block_size_);
            char const* local_kv_ptr = reinterpret_cast<char*>(paged_kv_block_pool_ptr_
                + params_kv_block_size_in_bytes_ * paged_kv_global_block_offsets_[paged_kv_block_idx]);

            // Predicates.
            // TODO: do we need to make sure row_idx < ROWS ?
            preds[ii] = row_idx < actual_seqlen_;
            preds[ii] &= col_in_bytes_ < VALID_BYTES_PER_ROW;

            // Pointers.
            int row_idx_in_block = row_idx & ((1 << paged_kv_log2_block_size_) - 1);
            ptrs[ii] = local_kv_ptr + head_col_in_bytes + (int64_t) row_idx_in_block * token_stride_in_bytes_;
        }

        // Trigger LDGSTS or the LDGs.
        // The predicates protect against out-of-bound access in rows and cols
        Ldgsts_helper<USE_LDGSTS>::load(this, smem_tile, ptrs, preds);
    }

    // Move the pointer to the next row location.
    inline __device__ void move()
    {
        row_ += ROWS;
    }

    // Move the pointer to the next row location by the offset (not step).
    inline __device__ void move_by_offset(int const offset)
    {
        row_ += offset;
    }

    // Move the pointer to the next column location
    inline __device__ void move_col()
    {
        col_in_bytes_ += THREADS_PER_ROW * BYTES_PER_LDG;
    }

    // Rewind the pointer back to previous column location
    inline __device__ void rewind_col(int const steps)
    {
        // Update col_in_bytes_ to ensure load predicates work
        col_in_bytes_ -= THREADS_PER_ROW * BYTES_PER_LDG * steps;
    }

    // The stride between rows for the KV matrice.
    int64_t params_kv_block_size_in_bytes_;
    // The paged cache pool pointer.
    char* paged_kv_block_pool_ptr_;
    // The paged block offsets.
    int32_t* paged_kv_global_block_offsets_;
    // The paged block size.
    int paged_kv_log2_block_size_;
    // The register to store predicates.
    uint32_t preds_[PRED_REGS];
    // The fetch registers.
    uint4 fetch_[LDGS];
    // Keep track of the row and col the thread is processing as we move the tile.
    int row_;
    int64_t col_in_bytes_;
    // Keep track of the head offset.
    int64_t head_stride_in_bytes_;
    // // for DeepSeek MLA, the stride of V tokens != VALID_BYTES_PER_ROW
    int32_t token_stride_in_bytes_;
    // The sequence length.
    int actual_seqlen_;
    // The past sequence length (kv_seqlen - q_seqlen) considering chunked context.
    int past_seqlen_;
    // The sliding attention window size.
    int sliding_window_size_;
};

////////////////////////////////////////////////////////////////////////////////////////////////////

template <
    // The instruction traits.
    typename Traits,
    // The dimensions of the tile computed by the CTA.
    typename Cta_tile,
    // The number of bits per element.
    int BITS_PER_ELEMENT,
    // The number of rows of Q loaded by this tile.
    int ROWS_,
    // The number of columns.
    int COLS,
    // Do we use LDGSTS?
    bool USE_LDGSTS_,
    // Are attention heads interleaved?
    bool HEADS_INTERLEAVED,
    // The number of matrices
    int NUM_MATS = 1>
struct Gmem_tile_q_kv
{

    // The size of each LDG.
    enum
    {
        BYTES_PER_LDG = 16
    };

    // The padded to the next power of 2 number of columns
    enum
    {
        COLS_PADDED = Next_power_of_two<COLS>::VALUE
    };

    // The padded size of a row in bytes.
    enum
    {
        BYTES_PER_ROW_PADDED = COLS_PADDED * BITS_PER_ELEMENT / 8
    };

    // The size of a row in bytes.
    enum
    {
        BYTES_PER_ROW = COLS * BITS_PER_ELEMENT / 8
    };

    // The number of threads to load a padded "row" of the matrix.
    enum
    {
        THREADS_PER_ROW_PADDED = BYTES_PER_ROW_PADDED / BYTES_PER_LDG
    };

    // The number of threads to load a "row" of the matrix.
    enum
    {
        THREADS_PER_ROW = BYTES_PER_ROW / BYTES_PER_LDG
    };

    // The number of "rows" loaded per LDG.
    enum
    {
        ROWS_PER_LDG = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW_PADDED
    };

    // The number of rows.
    enum
    {
        ROWS = ROWS_
    };

    // The number of LDGs needed to load a chunk of the Q matrix.
    enum
    {
        LDGS = fmha::Div_up<ROWS, ROWS_PER_LDG>::VALUE
    };

    // The number of predicate registers.
    enum
    {
        PRED_REGS = fmha::Compute_number_of_pred_regs<LDGS>::VALUE
    };

    // Is it Hopper?
    enum
    {
        IS_HOPPER = std::is_same<typename Traits::Gpu_arch, typename fmha::Hopper>::value == true
    };

    // Make sure we use a single register to store predicates. Do not throw for Hopper for now.
    static_assert(!USE_LDGSTS_ || PRED_REGS == 1 || IS_HOPPER, "");

    // We do not use LDGSTS (for the moment).
    enum
    {
        USE_LDGSTS = USE_LDGSTS_
    };

    // Ctor.
    template <typename Params, typename Block_info>
    inline __device__ Gmem_tile_q_kv(
        Params const& params, int offset, Block_info const& binfo, int tidx, int cta_row_offset = 0)
        : params_stride_in_bytes_(params.stride_in_bytes)
        , actual_seqlen_(binfo.actual_seqlen)
        , ptr_(reinterpret_cast<char*>(params.ptr))
    {

        // Compute the position in the sequence (within the CTA for the moment).
        int row = tidx / THREADS_PER_ROW_PADDED;
        // Compute the position of the thread in the row.
        int col = tidx % THREADS_PER_ROW_PADDED;

        // We must store the value to update the predicates in "load".
        row_ = row;
        // Mask for predicate if the channels are in the padded area
        int const bytes_per_row_non_padded = params.d * BITS_PER_ELEMENT / 8;
        mask_ = col < bytes_per_row_non_padded / BYTES_PER_LDG;

        // The row offset in the batched GEMM. For each seq element, we store QKV in that order.
        int64_t row_offset = (int64_t) (row + cta_row_offset) * params.stride_in_bytes;
        // Add the block index.
        int64_t idx;
        if (HEADS_INTERLEAVED)
        {
            idx = binfo.bidx * NUM_MATS + offset;
        }
        else
        {
            idx = (binfo.sum_s * NUM_MATS + offset) * params.h + binfo.bidh;
        }
        // Assemble the final pointer.
        ptr_ += row_offset + idx * bytes_per_row_non_padded + col * BYTES_PER_LDG;

        // Take the CTA offset to modify the sequence length.
        actual_seqlen_ -= cta_row_offset;
    }

    // Store data to shared memory.
    template <typename Smem_tile>
    inline __device__ void commit(Smem_tile& smem_tile)
    {
        if (!USE_LDGSTS)
        {
            smem_tile.store(fetch_);
        }
    }

    // Load data from memory.
    template <typename Smem_tile>
    inline __device__ void load(Smem_tile& smem_tile)
    {
        uint32_t preds[LDGS];
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            preds[ii] = (row_ + ii * (int) ROWS_PER_LDG < min((int) ROWS, actual_seqlen_)) && mask_;
        }

        // Prepare the load pointers.
        void const* ptrs[LDGS];
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            ptrs[ii] = ptr_ + (int64_t) ii * ROWS_PER_LDG * params_stride_in_bytes_;
        }

        // Trigger LDGSTS or the LDGs.
        Ldgsts_helper<USE_LDGSTS>::load(this, smem_tile, ptrs, preds);
    }

    inline __device__ void move(int const steps = 1)
    {
        ptr_ += (int64_t) ROWS * params_stride_in_bytes_ * steps;
        actual_seqlen_ -= (int) ROWS * steps;
    }

    // Store data to memory.
    inline __device__ void store(uint4 const (&data)[LDGS])
    {
#pragma unroll
        for (int ii = 0; ii < LDGS; ++ii)
        {
            char* ptr = ptr_ + (int64_t) ii * ROWS_PER_LDG * params_stride_in_bytes_;
            if ((row_ + ii * ROWS_PER_LDG) < min(ROWS, actual_seqlen_))
            {
                fmha::stg(ptr, data[ii]);
            }
        }
    }

    // The stride between rows for the matrix.
    int64_t params_stride_in_bytes_;
    // The pointer.
    char* ptr_;
    // The register to store predicates.
    uint32_t preds_[PRED_REGS];
    // The fetch registers.
    uint4 fetch_[LDGS];
    // Keep track of the row and col the thread is processing as we move the tile.
    int row_;
    // Keep track of predicate state that depends only on the initialization state.
    int mask_;
    // The sequence length.
    int actual_seqlen_;
};

////////////////////////////////////////////////////////////////////////////////////////////////////

template <
    // The instruction traits.
    typename Traits,
    // The dimensions of the tile computed by the CTA.
    typename Cta_tile,
    // The number of bits per element.
    int BITS_PER_ELEMENT,
    // The number of rows of Q, K or V loaded by this tile.
    int ROWS_,
    // The number of columns.
    int COLS,
    // Do we use LDGSTS?
    bool USE_LDGSTS_>
struct Gmem_tile_qkv_interleaved
{

    // The vectorization width for NC/32HW32.
    enum
    {
        VEC = 32
    };

    // The size of each LDG.
    enum
    {
        BYTES_PER_LDG = 16
    };

    // The size of a row in bytes.
    enum
    {
        BYTES_PER_ROW = VEC * BITS_PER_ELEMENT / 8
    };

    // DEBUG.
    static_assert(BYTES_PER_ROW == 32, "");

    // END OF DEBUG.

    // The number of threads to load a "row" of the matrix.
    enum
    {
        THREADS_PER_ROW = BYTES_PER_ROW / BYTES_PER_LDG
    };

    // DEBUG.
    static_assert(THREADS_PER_ROW == 2, "");

    // END OF DEBUG.

    // The number of "rows" loaded per LDG.
    enum
    {
        ROWS_PER_LDG = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW
    };

    // The number of slices. It is either 1 for DIM_PER_HEAD == 32 and 2 for DIM_PER_HEAD == 64.
    enum
    {
        NUM_SLICES = COLS / VEC
    };

    // DEBUG.
    static_assert(NUM_SLICES == 1 || NUM_SLICES == 2, "");

    // END OF DEBUG.

    // The number of rows in a slice.
    enum
    {
        ROWS = ROWS_
    };

    // The number of LDGs needed to load a chunk of the Q matrix.
    enum
    {
        LDGS = fmha::Div_up<ROWS * NUM_SLICES, ROWS_PER_LDG>::VALUE
    };

    // The number of predicate registers.
    enum
    {
        PRED_REGS = fmha::Compute_number_of_pred_regs<LDGS>::VALUE
    };

    // Make sure we use a single register to store predicates.
    static_assert(PRED_REGS == 1, "");

    // Do we use LDGSTS on Ampere?
    enum
    {
        USE_LDGSTS = USE_LDGSTS_
    };

    // Ctor.
    template <typename Params, typename Block_info>
    inline __device__ Gmem_tile_qkv_interleaved(
        Params const& params, int qkv_select, Block_info const& block_info, int tidx, int cta_row_offset = 0)
        : actual_seqlen_(block_info.actual_seqlen - cta_row_offset)
        , total_(params.q_stride_in_bytes)
        , kv_ptr_(reinterpret_cast<char const*>(params.qkv_ptr))
    {

        int bidh = block_info.bidh;
        int sum_s = block_info.sum_s;

        // We must keep track of the row to repack predicates in load.
        row_ = tidx / THREADS_PER_ROW;
        // The column.
        int col = tidx % THREADS_PER_ROW;

        // h is N
        // d is H
        // we get the data in as: 3 x h x (d/32) x total x 32 (think 3 x h x (d/32)
        // x b x s x 32)

        // Loading qkv: ignore slice for now.
        int qkv_offset = qkv_select * params.h * NUM_SLICES * total_;
        // bidh * GROUPS * B * S + b * S.
        int block_offset = bidh * NUM_SLICES * total_ + sum_s;
        // The row offset.
        int row_offset = (qkv_offset + block_offset + cta_row_offset) * BYTES_PER_ROW;

        // That's the pointer to load from (see "load").
        kv_ptr_ += row_offset + col * BYTES_PER_LDG;

        init_actual_seqlen_ = actual_seqlen_;
        init_kv_ptr_ = kv_ptr_;
    }

    // Store data to shared memory.
    template <typename Smem_tile>
    inline __device__ void commit(Smem_tile& smem_tile)
    {
        if (!USE_LDGSTS)
        {
            smem_tile.store(fetch_);
        }
    }

    // Load data from memory.
    template <typename Smem_tile>
    inline __device__ void load(Smem_tile& smem_tile)
    {
        void const* ptrs[LDGS];
        uint32_t preds[LDGS];

// We precompute slice offsets and predicates
#pragma unroll
        for (int ii = 0; ii < LDGS; ii++)
        {
            // the next row
            int row_i = row_ + ii * ROWS_PER_LDG;

            // Decompose the current row in slice and original row
            int slice = row_i / ROWS;
            // The position in the slice.
            int row_in_slice = row_i % ROWS;

            // Update the predicate.
            preds[ii] = row_in_slice < min(actual_seqlen_, ROWS);
            // Compute the pointer.
            ptrs[ii] = &kv_ptr_[(slice * total_ + row_in_slice) * BYTES_PER_ROW];
        }

        // Update the predicate register.
        fmha::pack_predicates(preds_, preds);

        // Trigger the loads.
        if (USE_LDGSTS)
        {
            smem_tile.store(ptrs, preds_);
        }
        else
        {
            fmha::ldg(fetch_, ptrs, preds_);
        }
    }

    // Move the pointer to the next location.
    inline __device__ void move(int const steps = 1)
    {
        kv_ptr_ += (int64_t) ROWS * BYTES_PER_ROW * steps;
        actual_seqlen_ -= ROWS * steps;
    }

    // Reset to the initial location.
    inline __device__ void reset()
    {
        kv_ptr_ = init_kv_ptr_;
        actual_seqlen_ = init_actual_seqlen_;
    }

    // The pointer.
    char const* kv_ptr_;
    char const* init_kv_ptr_;
    // The register to store predicates.
    uint32_t preds_[PRED_REGS];
    // The fetch registers.
    uint4 fetch_[LDGS];
    // keep track of the row the thread is processing as we move the tile
    int row_;
    // The sequence length.
    int actual_seqlen_;
    int init_actual_seqlen_;
    // The number of rows per slice??
    int total_;
};

////////////////////////////////////////////////////////////////////////////////////////////////////

} // namespace v2
} // namespace fmha
