// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Example 32: FMHA Forward with Sink Tokens
//
// Demonstrates forward kernel with sink tokens enabled. Sink tokens keep the
// first K positions always visible to all queries (StreamingLLM-style). Used
// with causal mask: positions [0, sink_size) are never masked. Planning only.

#include <iostream>

#include "ck_tile/dispatcher.hpp"
#include "ck_tile/dispatcher/example_args.hpp"

using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::utils;

DECL_FMHA_KERNEL_SET(sink_tokens_fmha_kernels,
                     // Forward with sink: first K tokens always visible
                     .add(FmhaSignature()
                              .family("fwd")
                              .dtype("fp16")
                              .mode("batch")
                              .vlayout("r")
                              .hdim(128)
                              .mask("top_left") // causal required with sink
                              .bias("no")
                              .lse(false)
                              .dropout(false)
                              .qscale("no")
                              .sink(true), // enables sink token path
                          FmhaAlgorithm()
                              .tile_m0(128)
                              .tile_n0(128)
                              .tile_k0(32)
                              .tile_n1(128)
                              .tile_k1(32)
                              .tile_k0max(128)
                              .wave_m0(4)
                              .wave_n0(1)
                              .wave_k0(1)
                              .wave_m1(4)
                              .wave_n1(1)
                              .wave_k1(1)
                              .warp_m0(32)
                              .warp_n0(32)
                              .warp_k0(16)
                              .warp_m1(32)
                              .warp_n1(32)
                              .warp_k1(16)
                              .pipeline("qr_async")
                              .padding(true, true, true, true)
                              .alignments(128, 128)
                              .selection_rank(0),
                          "gfx950"));

int main(int argc, char* argv[])
{
    ExampleArgs args("Example 32: FMHA Sink Tokens", "Forward with first K tokens always visible");
    args.add_option("--arch", "gfx950", "GPU architecture");
    args.add_option("--batch", "2", "Batch size");
    args.add_option("--nhead", "4", "Number of heads");
    args.add_option("--seqlen", "128", "Sequence length");
    args.add_option("--hdim", "128", "Head dimension");
    args.add_option("--sink", "4", "Number of sink tokens");
    if(!args.parse(argc, argv))
        return 0;

    const std::string gfx_arch = args.get("--arch", "gfx950");
    const int batch            = args.get_int("--batch", 2);
    const int nhead            = args.get_int("--nhead", 4);
    const int seqlen           = args.get_int("--seqlen", 128);
    const int hdim             = args.get_int("--hdim", 128);
    const int sink_size        = args.get_int("--sink", 4);

    print_header("Example 32: FMHA Sink Tokens");

    std::cout << "\nStep 1: Register Kernels\n";
    FmhaRegistry registry;
    registry.set_name("sink_tokens_fmha");
    REGISTER_GENERATED_KERNELS(registry, gfx_arch);
    std::cout << "  Registered " << registry.size() << " kernel(s)\n";

    std::cout << "\nStep 2: Plan\n";
    FmhaDispatcher dispatcher(&registry);
    fmha_fwd_traits traits{};
    traits.hdim_q        = hdim;
    traits.hdim_v        = hdim;
    traits.data_type     = "fp16";
    traits.is_group_mode = false;
    traits.is_v_rowmajor = true;
    traits.has_sink      = true;
    traits.mask_type     = mask_enum::mask_top_left;
    traits.bias_type     = bias_enum::no_bias;
    traits.has_lse       = false;
    traits.has_dropout   = false;
    traits.qscale_type   = quant_scale_enum::no_scale;

    fmha_fwd_args fwd_args{};
    fwd_args.batch     = batch;
    fwd_args.seqlen_q  = seqlen;
    fwd_args.seqlen_k  = seqlen;
    fwd_args.nhead_q   = nhead;
    fwd_args.nhead_k   = nhead;
    fwd_args.hdim_q    = hdim;
    fwd_args.hdim_v    = hdim;
    fwd_args.sink_size = sink_size;

    auto plan = dispatcher.plan(
        FmhaProblem::from_invocation(FmhaInvocation::make(traits, fwd_args), gfx_arch));
    std::cout << "  Plan valid: " << (plan.is_valid() ? "yes" : "no") << "\n";

    std::cout << "\nStep 3: Sink Tokens\n";
    std::cout << "  First " << sink_size << " tokens always visible to all queries.\n";
    std::cout << "  Used with causal mask for StreamingLLM-style long-context.\n";

    print_separator();
    return 0;
}