#pragma once

// Implements: include/ninfer/ops/residual_add.h
// Match: contiguous BF16 inputs. Registered aligned/eight-element domains use
// one 16-byte pack per thread; BF16x2 and scalar routes preserve correctness for
// smaller alignments and odd tails.

#include "ops/common/bf16_vector.cuh"
#include "ops/common/memory.cuh"

#include <cuda_bf16.h>

#include <cstdint>

namespace ninfer::ops {

inline constexpr int kResidualAddPairsPerThread = 4;

__device__ __forceinline__ __nv_bfloat162 residual_add_pair(__nv_bfloat162 y, __nv_bfloat162 x) {
    const float r0 = __low2float(x) + __low2float(y);
    const float r1 = __high2float(x) + __high2float(y);
    return __floats2bfloat162_rn(r0, r1);
}

__global__ void residual_add_scalar_kernel(const __nv_bfloat16* y, __nv_bfloat16* x,
                                           std::int64_t n) {
    const std::int64_t start  = blockIdx.x * static_cast<std::int64_t>(blockDim.x) + threadIdx.x;
    const std::int64_t stride = static_cast<std::int64_t>(gridDim.x) * blockDim.x;
    for (std::int64_t i = start; i < n; i += stride) {
        x[i] = __float2bfloat16_rn(__bfloat162float(x[i]) + __bfloat162float(y[i]));
    }
}

__launch_bounds__(256) __global__
    void residual_add_bf16x8_kernel(const Bf16x8Pack* y, Bf16x8Pack* x, std::int64_t packs) {
    const std::int64_t start  = blockIdx.x * static_cast<std::int64_t>(blockDim.x) + threadIdx.x;
    const std::int64_t stride = static_cast<std::int64_t>(gridDim.x) * blockDim.x;
    for (std::int64_t i = start; i < packs; i += stride) {
        const Bf16x8Pack yv = load_vec<Bf16x8Pack>(y + i);
        Bf16x8Pack xv       = load_vec<Bf16x8Pack>(x + i);
#pragma unroll
        for (int pair = 0; pair < 4; ++pair) {
            xv.pair[pair] = residual_add_pair(yv.pair[pair], xv.pair[pair]);
        }
        store_vec(x + i, xv);
    }
}

__launch_bounds__(256) __global__
    void residual_add_bf16x2_kernel(const __nv_bfloat16* y, __nv_bfloat16* x, std::int64_t n) {
    const std::int64_t tid = blockIdx.x * static_cast<std::int64_t>(blockDim.x) + threadIdx.x;
    const std::int64_t stride =
        static_cast<std::int64_t>(gridDim.x) * blockDim.x * kResidualAddPairsPerThread;
    const std::int64_t n2 = n / 2;

    const auto* y2 = reinterpret_cast<const __nv_bfloat162*>(y);
    auto* x2       = reinterpret_cast<__nv_bfloat162*>(x);
    for (std::int64_t j = tid * kResidualAddPairsPerThread; j < n2; j += stride) {
        const __nv_bfloat162 y0 = y2[j];
        const __nv_bfloat162 x0 = x2[j];
        if (j + 3 < n2) {
            const __nv_bfloat162 y1  = y2[j + 1];
            const __nv_bfloat162 x1  = x2[j + 1];
            const __nv_bfloat162 y2v = y2[j + 2];
            const __nv_bfloat162 x2v = x2[j + 2];
            const __nv_bfloat162 y3  = y2[j + 3];
            const __nv_bfloat162 x3  = x2[j + 3];
            x2[j]                    = residual_add_pair(y0, x0);
            x2[j + 1]                = residual_add_pair(y1, x1);
            x2[j + 2]                = residual_add_pair(y2v, x2v);
            x2[j + 3]                = residual_add_pair(y3, x3);
        } else {
            x2[j] = residual_add_pair(y0, x0);
            if (j + 1 < n2) { x2[j + 1] = residual_add_pair(y2[j + 1], x2[j + 1]); }
            if (j + 2 < n2) { x2[j + 2] = residual_add_pair(y2[j + 2], x2[j + 2]); }
        }
    }

    if (tid == 0 && (n & 1) != 0) {
        const std::int64_t i = n - 1;
        x[i]                 = __float2bfloat16_rn(__bfloat162float(x[i]) + __bfloat162float(y[i]));
    }
}

} // namespace ninfer::ops
#pragma once

#include <ninfer/targets/qwen3_6_27b/package.h>
#include <ninfer/targets/qwen3_6/frontend_resources.h>
#include <ninfer/targets/qwen3_6/model_view.h>
#include <ninfer/targets/qwen3_6/startup_features.h>
#include <ninfer/targets/qwen3_6/vision.h>

#include "artifact/binder.h"
#include "artifact/materializer.h"
#include "core/tensor.h"

#include <array>
#include <cstddef>
#include <cstdint>
#include <utility>
#include <variant>

namespace ninfer::targets::qwen3_6_27b::detail {

inline constexpr std::size_t kTextLayers          = 64;
inline constexpr std::size_t kFullAttentionLayers = 16;
inline constexpr std::size_t kGdnLayers           = 48;

struct WeightPlan {
    artifact::ObjectHandle object;
    artifact::NumericFormat format          = artifact::NumericFormat::BF16;
    std::uint32_t weight_scale_divisor_bits = 0;
    std::uint32_t input_scale_divisor_bits  = 0;
};

struct MlpPlan {
    WeightPlan gate_up;
    WeightPlan down;
};

struct SplitAttentionProjectionPlan {
    WeightPlan query_key;
    WeightPlan gate_value;
};

struct FusedAttentionProjectionPlan {
    WeightPlan query_key_gate_value;
};

struct FullAttentionPlan {
    std::variant<SplitAttentionProjectionPlan, FusedAttentionProjectionPlan> projection;
    artifact::ObjectHandle query_norm;
    artifact::ObjectHandle key_norm;
    WeightPlan output;
};

struct SplitGdnInputProjectionPlan {
    WeightPlan query_key;
    WeightPlan value_z;
};

struct FusedGdnInputProjectionPlan {
    WeightPlan query_key_value_z;
};

struct SplitGdnControlProjectionPlan {
    WeightPlan a_projection;
    WeightPlan b_projection;
};

struct FusedGdnControlProjectionPlan {
    WeightPlan a_b_projection;
};

using GdnControlProjectionPlan =
    std::variant<SplitGdnControlProjectionPlan, FusedGdnControlProjectionPlan>;

struct GdnPlan {
    artifact::ObjectHandle a_log;
    artifact::ObjectHandle dt_bias;
    artifact::ObjectHandle convolution;
    GdnControlProjectionPlan control_projection;
    std::variant<SplitGdnInputProjectionPlan, FusedGdnInputProjectionPlan> input_projection;
    artifact::ObjectHandle norm;
    WeightPlan output;
};

struct TextLayerPlan {
    artifact::ObjectHandle input_norm;
    FullAttentionPlan attention{};
    GdnPlan gdn{};
    bool is_full_attention = false;
    artifact::ObjectHandle post_attention_norm;
    MlpPlan mlp;
};

struct MtpPlan {
    artifact::ObjectHandle input_projection;
    artifact::ObjectHandle embedding_norm;
    artifact::ObjectHandle hidden_norm;
    artifact::ObjectHandle input_norm;
    artifact::ObjectHandle query_key_gate_value;
    artifact::ObjectHandle query_norm;
    artifact::ObjectHandle key_norm;
    artifact::ObjectHandle output;
    artifact::ObjectHandle post_attention_norm;
    MlpPlan mlp;
    artifact::ObjectHandle final_norm;
};

struct BindingPlan {
    qwen3_6::FrontendResourcePlan frontend;
    qwen3_6::StartupFeatures features;

    WeightPlan token_embedding;
    std::array<TextLayerPlan, kTextLayers> text_layers;
    artifact::ObjectHandle final_norm;
    WeightPlan output_head;
    artifact::ObjectHandle draft_head;
    artifact::ObjectHandle draft_head_token_ids;
    MtpPlan mtp;

    qwen3_6::VisionBackbonePlan vision_backbone;
    qwen3_6::VisionMergerInputPlan vision_merger_input;
    artifact::ObjectHandle vision_merger_fc2;
    artifact::ObjectHandle vision_merger_fc2_bias;
    qwen3_6::VisionMergerNormPlan vision_merger_norm;
};

struct ArtifactLoadPlan {
    BindingPlan bindings;
    artifact::MaterializationPlan materialization;
};

ArtifactLoadPlan bind_artifact(artifact::Binder& binder, WeightsProfile weights_profile,
                               qwen3_6::StartupFeatures features);

struct DensePostMixerPayload {
    Weight gate_up;
    Weight down;
};

struct SplitAttentionProjectionPayload {
    Weight query_key;
    Weight gate_value;
};

struct FusedAttentionProjectionPayload {
    Weight query_key_gate_value;
};

using FullAttentionProjectionPayload =
    std::variant<SplitAttentionProjectionPayload, FusedAttentionProjectionPayload>;

struct SplitGdnInputProjectionPayload {
    Weight query_key;
    Weight value_z;
};

struct FusedGdnInputProjectionPayload {
    Weight query_key_value_z;
};

using GdnInputProjectionPayload =
    std::variant<SplitGdnInputProjectionPayload, FusedGdnInputProjectionPayload>;

struct SplitGdnControlProjectionPayload {
    Weight a_projection;
    Weight b_projection;
};

struct FusedGdnControlProjectionPayload {
    Weight a_b_projection;
};

using GdnControlProjectionPayload =
    std::variant<SplitGdnControlProjectionPayload, FusedGdnControlProjectionPayload>;

struct GdnProjectionPayload {
    Tensor a_log;
    Tensor dt_bias;
    GdnControlProjectionPayload control_projection;
    GdnInputProjectionPayload input_projection;
};

struct MtpAttentionPayload {
    Weight packed;
    Weight query;
    Weight key;
    Weight output_gate;
    Weight value;
};

using RuntimeModelView =
    qwen3_6::ModelView<FullAttentionProjectionPayload, GdnProjectionPayload, DensePostMixerPayload,
                       MtpAttentionPayload, DensePostMixerPayload, qwen3_6::DFlashWeights<6>,
                       kFullAttentionLayers, kGdnLayers>;
using FullAttentionWeights = RuntimeModelView::FullLayer;
using GdnWeights           = RuntimeModelView::GdnLayer;
using MtpWeights           = RuntimeModelView::MtpLayer;

class LoadedModelData {
public:
    LoadedModelData(BindingPlan plan, artifact::MaterializedArtifact materialized);

    LoadedModelData(const LoadedModelData&)            = delete;
    LoadedModelData& operator=(const LoadedModelData&) = delete;
    LoadedModelData(LoadedModelData&&)                 = delete;
    LoadedModelData& operator=(LoadedModelData&&)      = delete;

    artifact::MaterializedArtifact backing;
    qwen3_6::FrontendResources frontend;
    RuntimeModelView runtime;
};

class LoadedModel::Impl {
public:
    Impl(WeightsProfile weights_profile_in, BindingPlan plan,
         artifact::MaterializedArtifact materialized)
        : weights_profile(weights_profile_in), data(std::move(plan), std::move(materialized)) {}

    WeightsProfile weights_profile;
    LoadedModelData data;
};

} // namespace ninfer::targets::qwen3_6_27b::detail
#include "ops/attn_input_proj/bf16/bf16_attn_input_plan.h"

#include "core/device.h"
#include "ops/linear/bf16/bf16_config.h"
#include "ops/linear/bf16/bf16_small_t.cuh"

#include <array>
#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include <utility>

namespace ninfer::ops::detail {
namespace {

using Launch = void (*)(const Tensor&, const Weight&, Tensor&, Tensor&, Tensor&, Tensor&,
                        cudaStream_t);

struct Bf16AttentionInputSmallTOutput {
    __nv_bfloat16* query;
    __nv_bfloat16* key;
    __nv_bfloat16* gate;
    __nv_bfloat16* value;

    __device__ __forceinline__ void store(std::int32_t parent_row, std::int32_t token,
                                          float result) const {
        constexpr std::int32_t kQueryRows  = 6144;
        constexpr std::int32_t kKeyRows    = 1024;
        constexpr std::int32_t kGateRows   = 6144;
        constexpr std::int32_t kKeyBegin   = kQueryRows;
        constexpr std::int32_t kGateBegin  = kKeyBegin + kKeyRows;
        constexpr std::int32_t kValueBegin = kGateBegin + kGateRows;
        const __nv_bfloat16 value_bf16     = __float2bfloat16_rn(result);

        if (parent_row < kKeyBegin) {
            query[static_cast<std::int64_t>(token) * kQueryRows + parent_row] = value_bf16;
        } else if (parent_row < kGateBegin) {
            key[static_cast<std::int64_t>(token) * kKeyRows + parent_row - kKeyBegin] = value_bf16;
        } else if (parent_row < kValueBegin) {
            gate[static_cast<std::int64_t>(token) * kGateRows + parent_row - kGateBegin] =
                value_bf16;
        } else {
            value[static_cast<std::int64_t>(token) * kKeyRows + parent_row - kValueBegin] =
                value_bf16;
        }
    }
};

template <int ActiveTokens>
struct Bf16AttentionSmallTProductionSchedule {
    static_assert(ActiveTokens >= 2 && ActiveTokens <= 32);
    // The Attention epilogue is measured separately from Linear, so it owns its exact-T winner
    // mapping while sharing the Linear computation body.
    static constexpr int kRowsPerWarp = ActiveTokens <= 4 ? 8 : (ActiveTokens <= 8 ? 4 : 2);
    static constexpr Bf16SmallTActivationAccess kActivationAccess =
        ActiveTokens <= 8 ? Bf16SmallTActivationAccess::WarpPacked
                          : Bf16SmallTActivationAccess::DirectStream;
    static constexpr bool kSequential = ActiveTokens <= 9 || ActiveTokens >= 17;
    static constexpr bool kUnroll2 = ActiveTokens == 4 || ActiveTokens == 5 || ActiveTokens == 8 ||
                                     (ActiveTokens >= 10 && ActiveTokens <= 18) ||
                                     ActiveTokens >= 23;
    static constexpr Bf16WeightCache kWeightCache =
        ActiveTokens == 7 ? Bf16WeightCache::Streaming : Bf16WeightCache::Default;
    static constexpr Bf16PhaseOrder kPhaseOrder =
        kSequential ? Bf16PhaseOrder::Sequential : Bf16PhaseOrder::RowSwizzled;
    using Type = Bf16SmallTInnerSchedule<4, 1, kRowsPerWarp, 8, 1, 4, kActivationAccess,
                                         kWeightCache, kPhaseOrder, 1, kUnroll2 ? 2 : 1, 1, 2>;
};

template <int ActiveTokens>
void launch_exact(const Tensor& x, const Weight& weight, Tensor& q, Tensor& gate, Tensor& k,
                  Tensor& v, cudaStream_t stream) {
    using Geometry = Bf16GemvGeometry<14336, 5120>;
    using Schedule = typename Bf16AttentionSmallTProductionSchedule<ActiveTokens>::Type;
    static_assert((Geometry::kOutputRows % Schedule::kRowsPerCta) == 0);
    static_assert((6144 % Schedule::kRowsPerCta) == 0);
    static_assert((1024 % Schedule::kRowsPerCta) == 0);

    const Bf16AttentionInputSmallTOutput output{
        static_cast<__nv_bfloat16*>(q.data),
        static_cast<__nv_bfloat16*>(k.data),
        static_cast<__nv_bfloat16*>(gate.data),
        static_cast<__nv_bfloat16*>(v.data),
    };
    constexpr int kBlocks = Geometry::kOutputRows / Schedule::kRowsPerCta;
    bf16_small_t_inner_kernel<Geometry, ActiveTokens, Schedule>
        <<<kBlocks, Schedule::kThreads, 0, stream>>>(
            static_cast<const __nv_bfloat16*>(x.data),
            static_cast<const __nv_bfloat16*>(weight.qdata), output);
    CUDA_CHECK(cudaGetLastError());
}

template <std::size_t... Offsets>
constexpr auto make_launchers(std::index_sequence<Offsets...>) {
    return std::array<Launch, sizeof...(Offsets)>{
        &launch_exact<kBf16AttnInputSmallTMinTokens + static_cast<int>(Offsets)>...};
}

constexpr auto kLaunchers = make_launchers(
    std::make_index_sequence<kBf16AttnInputSmallTMaxTokens - kBf16AttnInputSmallTMinTokens + 1>{});

} // namespace

void bf16_attn_input_small_t_launch(const Tensor& x, const Weight& weight, Tensor& q, Tensor& gate,
                                    Tensor& k, Tensor& v, cudaStream_t stream) {
    if (x.ne[1] < kBf16AttnInputSmallTMinTokens || x.ne[1] > kBf16AttnInputSmallTMaxTokens) {
        throw std::invalid_argument("bf16 attn_input_proj small-T requires T in [2,32]");
    }
    kLaunchers[x.ne[1] - kBf16AttnInputSmallTMinTokens](x, weight, q, gate, k, v, stream);
}

} // namespace ninfer::ops::detail
#include "ops/attn_input_proj/w8/w8_attn_input_kernels.h"

#include "core/device.h"
#include "ops/linear/w8/w8_k2048_decode.cuh"

namespace ninfer::ops::detail {

namespace {

template <int RowsPerCta>
void launch_companion_decode(const Tensor& x, const Weight& weight, Tensor& q, Tensor& k, Tensor& v,
                             cudaStream_t stream) {
    constexpr int kRows = 6144;
    static_assert((4096 % RowsPerCta) == 0 && (1024 % RowsPerCta) == 0);
    using Output = W8SplitOutput3<4096, 1024, 1024>;
    const Output output{static_cast<__nv_bfloat16*>(q.data), static_cast<__nv_bfloat16*>(k.data),
                        static_cast<__nv_bfloat16*>(v.data)};
    w8_k2048_decode_kernel<kRows, RowsPerCta><<<kRows / RowsPerCta, RowsPerCta * 32, 0, stream>>>(
        static_cast<const __nv_bfloat16*>(x.data), static_cast<const std::uint8_t*>(weight.qdata),
        static_cast<const std::uint8_t*>(weight.scales), output);
    CUDA_CHECK(cudaGetLastError());
}

} // namespace

void w8_attn_input_decode_launch(const Tensor& x, const Weight& weight, Tensor& q, Tensor& gate,
                                 Tensor& k, Tensor& v, cudaStream_t stream) {
    constexpr int kRows       = 9216;
    constexpr int kRowsPerCta = 8;
    static_assert((4096 % kRowsPerCta) == 0 && (512 % kRowsPerCta) == 0);
    using Output = W8SplitOutput4<4096, 512, 4096, 512>;
    const Output output{static_cast<__nv_bfloat16*>(q.data), static_cast<__nv_bfloat16*>(k.data),
                        static_cast<__nv_bfloat16*>(gate.data),
                        static_cast<__nv_bfloat16*>(v.data)};
    w8_k2048_decode_kernel<kRows, kRowsPerCta>
        <<<kRows / kRowsPerCta, kRowsPerCta * 32, 0, stream>>>(
            static_cast<const __nv_bfloat16*>(x.data),
            static_cast<const std::uint8_t*>(weight.qdata),
            static_cast<const std::uint8_t*>(weight.scales), output);
    CUDA_CHECK(cudaGetLastError());
}

void w8_attn_input_decode_launch(const Tensor& x, const Weight& weight, Tensor& q, Tensor& k,
                                 Tensor& v, cudaStream_t stream) {
    launch_companion_decode<8>(x, weight, q, k, v, stream);
}

void w8_companion_attn_input_decode_r4_launch(const Tensor& x, const Weight& weight, Tensor& q,
                                              Tensor& k, Tensor& v, cudaStream_t stream) {
    launch_companion_decode<4>(x, weight, q, k, v, stream);
}

void w8_companion_attn_input_decode_r16_launch(const Tensor& x, const Weight& weight, Tensor& q,
                                               Tensor& k, Tensor& v, cudaStream_t stream) {
    launch_companion_decode<16>(x, weight, q, k, v, stream);
}

} // namespace ninfer::ops::detail
#include "ninfer/ops/vision_pos_embed.h"

#include "ops/launcher/vision_pos_embed.h"

#include <stdexcept>

namespace ninfer::ops {

void vision_pos_embed_add(const Tensor& table, const Tensor& indices, const Tensor& weights,
                          Tensor& x, cudaStream_t stream) {
    if (table.dtype != DType::BF16 || x.dtype != DType::BF16 || indices.dtype != DType::I32 ||
        weights.dtype != DType::FP32) {
        throw std::invalid_argument(
            "vision_pos_embed_add: table/x must be BF16, indices I32, weights FP32");
    }
    if (table.ne[2] != 1 || table.ne[3] != 1 || x.ne[2] != 1 || x.ne[3] != 1 ||
        indices.ne[0] != 4 || indices.ne[1] != x.ne[1] || indices.ne[2] != 1 ||
        indices.ne[3] != 1 || weights.ne[0] != 4 || weights.ne[1] != x.ne[1] ||
        weights.ne[2] != 1 || weights.ne[3] != 1 || table.ne[0] != x.ne[0]) {
        throw std::invalid_argument(
            "vision_pos_embed_add: expected table [D,R], controls [4,P], x [D,P]");
    }
    if (x.ne[0] < 0 || x.ne[1] < 0 || table.ne[1] <= 0) {
        throw std::invalid_argument("vision_pos_embed_add: invalid dimensions");
    }
    if (x.ne[0] == 0 || x.ne[1] == 0) { return; }
    if (!table.is_contiguous() || !indices.is_contiguous() || !weights.is_contiguous() ||
        !x.is_contiguous()) {
        throw std::invalid_argument("vision_pos_embed_add: tensors must be contiguous");
    }
    if (table.data == nullptr || indices.data == nullptr || weights.data == nullptr ||
        x.data == nullptr) {
        throw std::invalid_argument("vision_pos_embed_add: tensor data must be non-null");
    }
    detail::vision_pos_embed_add_launch(table, indices, weights, x, stream);
}

} // namespace ninfer::ops
#pragma once

// ninfer::ops::detail - private launch prototype for residual_add. Included by the wrapper
// (host) and defined by the launcher (.cu). Not part of the public api.
// See docs/op-development.md §2.

#include "core/tensor.h"

#include <cuda_runtime.h>

namespace ninfer::ops::detail {

// Host entry; assumes inputs already validated by the wrapper.
void residual_add_launch(const Tensor& y, Tensor& x, cudaStream_t stream);

} // namespace ninfer::ops::detail
#pragma once

#include "core/dtype.h"

#include <cstddef>
#include <cstdint>
#include <initializer_list>

namespace ninfer {

struct Tensor {
    void* data         = nullptr;
    DType dtype        = DType::BF16;
    std::int32_t ne[4] = {1, 1, 1, 1};
    std::int64_t nb[4] = {0, 0, 0, 0};

    Tensor() noexcept = default;
    Tensor(void* data, DType dtype, std::initializer_list<std::int32_t> shape);

    std::int64_t numel() const;
    std::size_t bytes() const;
    bool is_contiguous() const;

    Tensor view(std::initializer_list<std::int32_t> shape) const;
    Tensor reshape(std::initializer_list<std::int32_t> shape) const;
    Tensor slice(int dim, std::int32_t start, std::int32_t len) const;
    Tensor permute(std::initializer_list<int> order) const;
};

enum class QType : std::uint16_t {
    Q4G64_F16S           = 0,
    Q5G64_F16S           = 1,
    Q6G64_F16S           = 2,
    W8G32_F16S           = 3,
    BF16_CTRL            = 4,
    FP32_CTRL            = 5,
    I32_CTRL             = 6,
    NVFP4                = 7,
    FP8_E4M3FN_ROW_BF16S = 8,
};

enum class QuantLayout : std::uint16_t {
    RowSplit            = 0,
    Contiguous          = 1,
    BlockScaleK16M128x4 = 2,
    RowScale            = 3,
};

struct Weight {
    const void* payload            = nullptr;
    std::uint64_t payload_bytes    = 0;
    std::uint64_t high_plane_bytes = 0;
    QType qtype                    = QType::Q4G64_F16S;
    std::uint32_t group_size       = 0;
    std::int32_t shape[4]          = {1, 1, 1, 1};
    std::int32_t padded_shape[4]   = {1, 1, 1, 1};
    std::uint32_t ndim             = 0;

    const void* qdata          = nullptr;
    const void* qhigh          = nullptr;
    const void* scales         = nullptr;
    std::int32_t n             = 0;
    std::int32_t k             = 0;
    std::int32_t group         = 0;
    QuantLayout layout         = QuantLayout::RowSplit;
    DType scale_dtype          = DType::FP32;
    std::int32_t scale_ne[4]   = {1, 1, 1, 1};
    std::int64_t scale_nb[4]   = {0, 0, 0, 0};
    float weight_scale_divisor = 0.0F;
    float input_scale_divisor  = 0.0F;
};

} // namespace ninfer
#include "ninfer/ops/prepare_masked_block.h"

#include "ops/launcher/prepare_masked_block.h"

#include <cstdint>
#include <stdexcept>
#include <string>

namespace ninfer::ops {
namespace {

void require_i32_vector(const Tensor& tensor, std::int32_t size, const char* name) {
    if (tensor.dtype != DType::I32 || tensor.ne[0] != size || tensor.ne[1] != 1 ||
        tensor.ne[2] != 1 || tensor.ne[3] != 1 || !tensor.is_contiguous() ||
        tensor.data == nullptr) {
        throw std::invalid_argument("prepare_masked_block: " + std::string(name) +
                                    " must be a contiguous I32 vector");
    }
}

void require_i32_matrix(const Tensor& tensor, std::int32_t width, std::int32_t batch,
                        const char* name) {
    if (tensor.dtype != DType::I32 || tensor.ne[0] != width || tensor.ne[1] != batch ||
        tensor.ne[2] != 1 || tensor.ne[3] != 1 || !tensor.is_contiguous() ||
        tensor.data == nullptr) {
        throw std::invalid_argument("prepare_masked_block: " + std::string(name) +
                                    " must be a contiguous I32 matrix");
    }
}

bool overlaps(const Tensor& lhs, const Tensor& rhs) {
    const auto lhs_begin = reinterpret_cast<std::uintptr_t>(lhs.data);
    const auto rhs_begin = reinterpret_cast<std::uintptr_t>(rhs.data);
    return lhs_begin < rhs_begin + rhs.bytes() && rhs_begin < lhs_begin + lhs.bytes();
}

} // namespace

void prepare_masked_block(const Tensor& anchors, const Tensor& lengths, const Tensor& valid_columns,
                          std::int32_t mask_id, Tensor& ids, Tensor& positions,
                          cudaStream_t stream) {
    const std::int32_t block_size = ids.ne[0];
    const std::int32_t batch_size = ids.ne[1];
    if (block_size < 1 || block_size > 16) {
        throw std::invalid_argument("prepare_masked_block: W must be 1..16");
    }
    if (batch_size < 1 || batch_size > 8) {
        throw std::invalid_argument("prepare_masked_block: B must be 1..8");
    }
    if (mask_id < 0) {
        throw std::invalid_argument("prepare_masked_block: mask_id must be nonnegative");
    }
    require_i32_vector(anchors, batch_size, "anchors");
    require_i32_vector(lengths, batch_size, "lengths");
    require_i32_vector(valid_columns, batch_size, "valid_columns");
    require_i32_matrix(ids, block_size, batch_size, "ids");
    require_i32_matrix(positions, block_size, batch_size, "positions");
    if (overlaps(ids, positions) || overlaps(anchors, ids) || overlaps(anchors, positions) ||
        overlaps(lengths, ids) || overlaps(lengths, positions) || overlaps(valid_columns, ids) ||
        overlaps(valid_columns, positions)) {
        throw std::invalid_argument("prepare_masked_block: inputs and outputs must not overlap");
    }
    detail::prepare_masked_block_launch(anchors, lengths, valid_columns, mask_id, ids, positions,
                                        block_size, stream);
}

} // namespace ninfer::ops
#pragma once

#include "core/layout.h"
#include "ops/linear_attention/gated_delta_net/common.h"

#include <cuda_bf16.h>
#include <cuda_runtime.h>

#include <cstddef>
#include <cstdint>
#include <cstdio>

#define NINFER_GATED_DELTA_NET_PROPAGATE(expr)                                                     \
    do {                                                                                           \
        const cudaError_t ninfer_gated_delta_net_error = (expr);                                   \
        if (ninfer_gated_delta_net_error != cudaSuccess) { return ninfer_gated_delta_net_error; }  \
    } while (0)

namespace ninfer::ops::detail::gated_delta_net::chunked {

inline constexpr std::size_t kWorkspaceAlign = 256;

struct workspace_layout {
    TensorRegion g_cumsum;
    TensorRegion W;
    TensorRegion U;
    TensorRegion v_new;
    TensorRegion h_chunk;
    std::size_t total_bytes = 0;
};

inline workspace_layout compute_workspace_layout(std::int32_t value_heads, std::int32_t tokens) {
    const std::int32_t chunks = tokens / kChunkSize;
    LayoutBuilder builder;
    workspace_layout w{};
    w.g_cumsum =
        builder.add_tensor(DType::FP32, {value_heads, tokens}, kWorkspaceAlign, "g_cumsum");
    w.W = builder.add_tensor(DType::BF16, {kStateDim, value_heads, tokens}, kWorkspaceAlign, "W");
    w.U = builder.add_tensor(DType::BF16, {kStateDim, value_heads, tokens}, kWorkspaceAlign, "U");
    w.v_new =
        builder.add_tensor(DType::BF16, {kStateDim, value_heads, tokens}, kWorkspaceAlign, "v_new");
    w.h_chunk     = builder.add_tensor(DType::BF16, {kStateDim, kStateDim, value_heads, chunks},
                                       kWorkspaceAlign, "h_chunk");
    w.total_bytes = builder.finish(kWorkspaceAlign, "Gated DeltaNet chunk workspace");
    return w;
}

inline std::size_t workspace_bytes(std::int32_t value_heads, std::int32_t tokens) {
    return compute_workspace_layout(value_heads, tokens).total_bytes;
}

struct prepare_wy_wu_config {
    std::int32_t H_qk = 0;
    std::int32_t H_v  = 0;
    std::int32_t L    = 0;

    const __nv_bfloat16* k = nullptr;
    const __nv_bfloat16* v = nullptr;
    const float* g_in      = nullptr;
    const float* beta      = nullptr;

    __nv_bfloat16* W    = nullptr;
    __nv_bfloat16* U    = nullptr;
    float* g_cumsum_out = nullptr;

    cudaStream_t stream = nullptr;
};

struct state_passing_config {
    std::int32_t H_qk = 0;
    std::int32_t H_v  = 0;
    std::int32_t L    = 0;

    const __nv_bfloat16* W = nullptr;
    const __nv_bfloat16* U = nullptr;
    const __nv_bfloat16* k = nullptr;
    const float* g_cumsum  = nullptr;
    const float* state_in  = nullptr;

    __nv_bfloat16* v_new   = nullptr;
    __nv_bfloat16* h_chunk = nullptr;
    float* state_out       = nullptr;

    cudaStream_t stream = nullptr;
};

struct chunk_output_config {
    std::int32_t H_qk = 0;
    std::int32_t H_v  = 0;
    std::int32_t L    = 0;

    const __nv_bfloat16* q       = nullptr;
    const __nv_bfloat16* k       = nullptr;
    const __nv_bfloat16* v_new   = nullptr;
    const float* g_cumsum        = nullptr;
    const __nv_bfloat16* h_chunk = nullptr;

    __nv_bfloat16* attn_out = nullptr;

    float scale = 0.0f;

    cudaStream_t stream = nullptr;
};

struct stage_validator {
    const char* name;
    std::int32_t H_qk;
    std::int32_t H_v;
    std::int32_t T;

    cudaError_t check_shape() const {
        if (T <= 0 || !are_head_counts_valid(H_qk, H_v)) {
            std::fprintf(stderr, "%s: invalid shape (H_qk=%d H_v=%d T=%d)\n", name, H_qk, H_v, T);
            return cudaErrorInvalidValue;
        }
        return cudaSuccess;
    }

    cudaError_t check_full_chunks() const {
        if ((T % kChunkSize) != 0) {
            std::fprintf(stderr,
                         "%s: Gated DeltaNet chunked path requires T to be a multiple of %d; "
                         "route tail tokens through AR instead (T=%lld)\n",
                         name, kChunkSize, static_cast<long long>(T));
            return cudaErrorInvalidValue;
        }
        return cudaSuccess;
    }

    cudaError_t check_grid(std::int64_t grid_x, std::int64_t grid_y,
                           std::int64_t grid_z = 1) const {
        if (grid_x > static_cast<std::int64_t>(0xffffffff)) {
            std::fprintf(stderr, "%s: grid.x too large (%lld)\n", name,
                         static_cast<long long>(grid_x));
            return cudaErrorInvalidConfiguration;
        }
        if (grid_y > static_cast<std::int64_t>(0xffff)) {
            std::fprintf(stderr, "%s: grid.y too large (%lld)\n", name,
                         static_cast<long long>(grid_y));
            return cudaErrorInvalidConfiguration;
        }
        if (grid_z > static_cast<std::int64_t>(0xffff)) {
            std::fprintf(stderr, "%s: grid.z too large (%lld)\n", name,
                         static_cast<long long>(grid_z));
            return cudaErrorInvalidConfiguration;
        }
        return cudaSuccess;
    }
};

cudaError_t launch_prepare_wy_wu(const prepare_wy_wu_config& cfg);
cudaError_t launch_state_passing(const state_passing_config& cfg);
cudaError_t launch_output(const chunk_output_config& cfg);

} // namespace ninfer::ops::detail::gated_delta_net::chunked
// ninfer::ops - embedding wrapper: public api validation and qtype dispatch.
#include "ninfer/ops/embedding.h"

#include "ops/common/math.h"
#include "ops/linear/fp8/fp8_format.h"
#include "ops/launcher/embed_gather.h" // detail::embed_gather_*_launch
#include "core/weight.h"

#include <cstdint>
#include <limits>
#include <stdexcept>
#include <string>

namespace ninfer::ops {
namespace {

std::int64_t numel_allow_zero(const Tensor& t, const char* label) {
    bool has_zero = false;
    for (int d = 0; d < 4; ++d) {
        if (t.ne[d] < 0) {
            throw std::invalid_argument(std::string("embedding: ") + label +
                                        " dimensions must be nonnegative");
        }
        if (t.ne[d] == 0) { has_zero = true; }
    }
    if (has_zero) { return 0; }

    std::int64_t total = 1;
    for (int d = 0; d < 4; ++d) {
        if (total > std::numeric_limits<std::int64_t>::max() / t.ne[d]) {
            throw std::overflow_error("embedding: tensor size overflows int64");
        }
        total *= t.ne[d];
    }
    return total;
}

std::uint64_t checked_mul_u64(std::uint64_t a, std::uint64_t b) {
    if (b != 0 && a > std::numeric_limits<std::uint64_t>::max() / b) {
        throw std::overflow_error("embedding: weight payload size overflows uint64");
    }
    return a * b;
}

std::int32_t align_up_i32(std::int32_t x, std::int32_t m) {
    const std::int64_t y = round_up(static_cast<std::int64_t>(x), static_cast<std::int64_t>(m));
    if (y > std::numeric_limits<std::int32_t>::max()) {
        throw std::overflow_error("embedding: padded shape overflows int32");
    }
    return static_cast<std::int32_t>(y);
}

void require_ids_shape(const Tensor& ids) {
    if (ids.ne[1] != 1 || ids.ne[2] != 1 || ids.ne[3] != 1) {
        throw std::invalid_argument("embedding: ids must have shape [T]");
    }
}

void require_out_shape(const Tensor& ids, const Tensor& out) {
    if (out.ne[2] != 1 || out.ne[3] != 1) {
        throw std::invalid_argument("embedding: out must have shape [d,T]");
    }
    if (out.ne[1] != ids.ne[0]) {
        throw std::invalid_argument("embedding: out T dimension must match ids");
    }
}

void require_weight_2d(const Weight& table) {
    if (table.ndim != 2) { throw std::invalid_argument("embedding: table must be 2-D [vocab,d]"); }
    if (table.shape[0] <= 0 || table.shape[1] <= 0) {
        throw std::invalid_argument("embedding: table shape must be positive");
    }
}

void require_dense_metadata(const Weight& table, const Tensor& out) {
    if (table.layout != QuantLayout::Contiguous) {
        throw std::invalid_argument("embedding: BF16_CTRL table must be Contiguous");
    }
    require_weight_2d(table);
    if (table.shape[1] != out.ne[0]) {
        throw std::invalid_argument("embedding: dense table d must match out.ne[0]");
    }
    if (table.qhigh != nullptr || table.high_plane_bytes != 0) {
        throw std::invalid_argument("embedding: dense table high plane must be null");
    }
    const std::uint64_t expected =
        checked_mul_u64(checked_mul_u64(static_cast<std::uint64_t>(table.shape[0]),
                                        static_cast<std::uint64_t>(table.shape[1])),
                        2);
    if (table.payload_bytes != 0 && table.payload_bytes < expected) {
        throw std::invalid_argument("embedding: dense payload is too small");
    }
}

void require_q6_metadata(const Weight& table, const Tensor& out) {
    if (table.layout != QuantLayout::RowSplit) {
        throw std::invalid_argument("embedding: Q6G64_F16S table must be RowSplit");
    }
    require_weight_2d(table);
    if (table.group_size != 64 || table.group != 64) {
        throw std::invalid_argument("embedding: Q6G64_F16S table group must be 64");
    }
    if (table.scale_dtype != DType::FP16) {
        throw std::invalid_argument("embedding: Q6G64_F16S table scale dtype must be FP16");
    }
    if (table.padded_shape[0] != table.shape[0] ||
        table.padded_shape[1] != align_up_i32(table.shape[1], 128)) {
        throw std::invalid_argument("embedding: Q6G64_F16S padded shape is invalid");
    }
    if (table.shape[1] != out.ne[0]) {
        throw std::invalid_argument("embedding: Q6G64_F16S table d must match out.ne[0]");
    }
    const std::uint64_t kg = static_cast<std::uint64_t>(table.padded_shape[1] / 64);
    const std::uint64_t nibble_plane_bytes =
        checked_mul_u64(checked_mul_u64(static_cast<std::uint64_t>(table.shape[0]), kg), 32);
    const std::uint64_t high_plane_bytes =
        checked_mul_u64(checked_mul_u64(static_cast<std::uint64_t>(table.shape[0]), kg), 16);
    const std::uint64_t scale_plane_bytes =
        checked_mul_u64(checked_mul_u64(static_cast<std::uint64_t>(table.shape[0]), kg), 2);
    const std::uint64_t high_plane_off = ((nibble_plane_bytes + 255u) / 256u) * 256u;
    const std::uint64_t scale_plane_off =
        high_plane_off + ((high_plane_bytes + 255u) / 256u) * 256u;
    const std::uint64_t expected = scale_plane_off + scale_plane_bytes;
    if (table.payload_bytes != 0 && table.payload_bytes < expected) {
        throw std::invalid_argument("embedding: Q6G64_F16S payload is too small");
    }
    if (table.qdata == nullptr || table.qhigh == nullptr || table.scales == nullptr) {
        throw std::invalid_argument("embedding: Q6G64_F16S planes must be non-null");
    }
    if (table.high_plane_bytes < high_plane_bytes) {
        throw std::invalid_argument("embedding: Q6G64_F16S high plane is too small");
    }
}

void require_w8_metadata(const Weight& table, const Tensor& out) {
    if (table.layout != QuantLayout::RowSplit) {
        throw std::invalid_argument("embedding: W8G32_F16S table must be RowSplit");
    }
    require_weight_2d(table);
    if (table.group_size != 32 || table.group != 32) {
        throw std::invalid_argument("embedding: W8G32_F16S table group must be 32");
    }
    if (table.scale_dtype != DType::FP16) {
        throw std::invalid_argument("embedding: W8G32_F16S table scale dtype must be FP16");
    }
    if (table.padded_shape[0] != table.shape[0] ||
        table.padded_shape[1] != align_up_i32(table.shape[1], 128)) {
        throw std::invalid_argument("embedding: W8G32_F16S padded shape is invalid");
    }
    if (table.shape[1] != out.ne[0]) {
        throw std::invalid_argument("embedding: W8G32_F16S table d must match out.ne[0]");
    }
    const std::uint64_t kg = static_cast<std::uint64_t>(table.padded_shape[1] / 32);
    const std::uint64_t code_plane_bytes =
        checked_mul_u64(checked_mul_u64(static_cast<std::uint64_t>(table.shape[0]), kg), 32);
    const std::uint64_t scale_plane_bytes =
        checked_mul_u64(checked_mul_u64(static_cast<std::uint64_t>(table.shape[0]), kg), 2);
    const std::uint64_t scale_plane_off = ((code_plane_bytes + 255u) / 256u) * 256u;
    const std::uint64_t expected        = scale_plane_off + scale_plane_bytes;
    if (table.payload_bytes != 0 && table.payload_bytes < expected) {
        throw std::invalid_argument("embedding: W8G32_F16S payload is too small");
    }
    if (table.qdata == nullptr || table.scales == nullptr) {
        throw std::invalid_argument("embedding: W8G32_F16S planes must be non-null");
    }
    if (table.qhigh != nullptr || table.high_plane_bytes != 0) {
        throw std::invalid_argument("embedding: W8G32_F16S high plane must be empty");
    }
}

void require_fp8_metadata(const Weight& table, const Tensor& out) {
    constexpr std::int32_t kVocabulary = 248320;
    constexpr std::int32_t kHidden     = 5120;
    if (table.n != kVocabulary || table.k != kHidden || out.ne[0] != kHidden) {
        throw std::invalid_argument("embedding: unsupported FP8 table shape");
    }
    if ((reinterpret_cast<std::uintptr_t>(out.data) &
         (static_cast<std::uintptr_t>(alignof(std::uint32_t)) - 1)) != 0) {
        throw std::invalid_argument("embedding: FP8 output must be 4-byte aligned");
    }
    (void)detail::validate_fp8_weight(table, "embedding");
}

bool is_empty_T(const Tensor& ids, const Tensor& out) { return ids.ne[0] == 0 || out.ne[1] == 0; }

void require_non_empty_tensors(const Tensor& ids, const Tensor& out) {
    if (!ids.is_contiguous() || !out.is_contiguous()) {
        throw std::invalid_argument("embedding: ids/out must be contiguous");
    }
    if (ids.data == nullptr || out.data == nullptr) {
        throw std::invalid_argument("embedding: ids/out data must be non-null");
    }
}

} // namespace

void embedding(const Tensor& ids, const Weight& table, Tensor& out, cudaStream_t stream) {
    if (ids.dtype != DType::I32) { throw std::invalid_argument("embedding: ids must be I32"); }
    if (out.dtype != DType::BF16) { throw std::invalid_argument("embedding: out must be BF16"); }

    (void)numel_allow_zero(ids, "ids");
    (void)numel_allow_zero(out, "out");
    require_ids_shape(ids);
    require_out_shape(ids, out);

    switch (table.qtype) {
    case QType::BF16_CTRL: {
        require_dense_metadata(table, out);
        if (is_empty_T(ids, out)) { return; }
        require_non_empty_tensors(ids, out);
        if (table.qdata == nullptr) {
            throw std::invalid_argument("embedding: dense table data must be non-null");
        }
        const Tensor dense = as_dense(table);
        detail::embed_gather_dense_launch(ids, dense, out, stream);
    } break;
    case QType::Q6G64_F16S:
        require_q6_metadata(table, out);
        if (is_empty_T(ids, out)) { return; }
        require_non_empty_tensors(ids, out);
        detail::embed_gather_q6_launch(ids, table, out, stream);
        break;
    case QType::W8G32_F16S:
        require_w8_metadata(table, out);
        if (is_empty_T(ids, out)) { return; }
        require_non_empty_tensors(ids, out);
        detail::embed_gather_w8_launch(ids, table, out, stream);
        break;
    case QType::FP8_E4M3FN_ROW_BF16S:
        require_fp8_metadata(table, out);
        if (is_empty_T(ids, out)) { return; }
        require_non_empty_tensors(ids, out);
        detail::embed_gather_fp8_launch(ids, table, out, stream);
        break;
    default:
        throw std::invalid_argument("embedding: unsupported table qtype");
    }
}

} // namespace ninfer::ops
#include "ops/linear_add/q5/q5_linear_add_kernels.h"

#include "core/device.h"
#include "ops/linear/q5/q5_rowsplit_gemm_simt.cuh"

#include <cuda_bf16.h>

#include <cstdint>
#include <stdexcept>

namespace ninfer::ops::detail {
namespace {

template <int Cols, int FullSlabs, int Stride>
void launch_split2(const Tensor& x, const Weight& w, Tensor& residual_out, cudaStream_t stream) {
    constexpr int kThreads = 2 * 32;
    const dim3 grid(static_cast<unsigned>(residual_out.ne[0]), 1u, 1u);
    q5_rowsplit_gemm_simt_split2_kernel<Q5RowSplitSimtSchedule, Cols, FullSlabs, Stride, false, 0,
                                        true><<<grid, kThreads, 0, stream>>>(
        static_cast<const __nv_bfloat16*>(x.data), static_cast<const std::uint8_t*>(w.qdata),
        static_cast<const std::uint8_t*>(w.qhigh), static_cast<const std::uint8_t*>(w.scales),
        static_cast<__nv_bfloat16*>(residual_out.data), residual_out.ne[0], x.ne[0], x.ne[1],
        w.padded_shape[1], FullSlabs);
}

template <int Cols>
void dispatch_shape(const Tensor& x, const Weight& w, Tensor& residual_out, cudaStream_t stream) {
    if (w.k == 6144) {
        launch_split2<Cols, 6, 6144>(x, w, residual_out, stream);
    } else if (w.k == 17408) {
        launch_split2<Cols, 17, 17408>(x, w, residual_out, stream);
    } else {
        throw std::invalid_argument("q5 linear_add split2: unsupported exact K");
    }
}

template <class Launch>
void dispatch_cols(std::int32_t cols, Launch&& launch) {
    switch (cols) {
#define NINFER_Q5_LINEAR_ADD_EXACT(COLS)                                                           \
    case COLS:                                                                                     \
        launch.template operator()<COLS>();                                                        \
        return
        NINFER_Q5_LINEAR_ADD_EXACT(2);
        NINFER_Q5_LINEAR_ADD_EXACT(3);
        NINFER_Q5_LINEAR_ADD_EXACT(4);
        NINFER_Q5_LINEAR_ADD_EXACT(5);
        NINFER_Q5_LINEAR_ADD_EXACT(6);
        NINFER_Q5_LINEAR_ADD_EXACT(7);
        NINFER_Q5_LINEAR_ADD_EXACT(8);
        NINFER_Q5_LINEAR_ADD_EXACT(9);
        NINFER_Q5_LINEAR_ADD_EXACT(10);
        NINFER_Q5_LINEAR_ADD_EXACT(11);
        NINFER_Q5_LINEAR_ADD_EXACT(12);
        NINFER_Q5_LINEAR_ADD_EXACT(13);
        NINFER_Q5_LINEAR_ADD_EXACT(14);
        NINFER_Q5_LINEAR_ADD_EXACT(15);
        NINFER_Q5_LINEAR_ADD_EXACT(16);
#undef NINFER_Q5_LINEAR_ADD_EXACT
    default:
        throw std::invalid_argument("q5 linear_add split2: T must be in [2,16]");
    }
}

} // namespace

void q5_linear_add_split2_exact_launch(const Tensor& x, const Weight& w, Tensor& residual_out,
                                       cudaStream_t stream) {
    dispatch_cols(x.ne[1], [&]<int Cols>() { dispatch_shape<Cols>(x, w, residual_out, stream); });
    CUDA_CHECK(cudaGetLastError());
}

} // namespace ninfer::ops::detail
#pragma once

#include "core/arena.h"
#include "core/tensor.h"
#include "ninfer/ops/linear.h"
#include "ops/linear/nvfp4/nvfp4_w4a4_plan.h"

#include <cuda_runtime.h>

#include <cstddef>
#include <cstdint>

namespace ninfer::ops::detail {

[[nodiscard]] std::size_t nvfp4_linear_add_workspace_capacity_bytes(std::int32_t output_rows,
                                                                    std::int32_t input_rows,
                                                                    LinearPolicy policy,
                                                                    std::int32_t min_tokens,
                                                                    std::int32_t max_tokens);

void nvfp4_linear_add_decode_launch(const Tensor& x, const Weight& weight, Tensor& residual,
                                    cudaStream_t stream);

void nvfp4_linear_add_small_t_launch(const Tensor& x, const Weight& weight, Tensor& residual,
                                     cudaStream_t stream);

void nvfp4_linear_add_w4a4_launch(const Tensor& x, const Weight& weight, Tensor& residual,
                                  Nvfp4W4a4Workspace workspace, cudaStream_t stream);

void nvfp4_linear_add_dispatch(const Tensor& x, const Weight& weight, Tensor& residual,
                               LinearPolicy policy, WorkspaceArena& workspace, cudaStream_t stream);

} // namespace ninfer::ops::detail
#include "ops/softmax_attention/sliding_window/launch.h"

#include "core/device.h"
#include "ops/softmax_attention/sliding_window/kernel.cuh"

#include <algorithm>
#include <cstdint>
#include <stdexcept>

namespace ninfer::ops::detail {
namespace {

template <int Tokens, class Launch>
void dispatch_token_case(Launch&& launch) {
    constexpr int Warps = (Tokens + 3) / 4;
    launch.template operator()<Tokens, Warps>();
}

template <class Launch>
void dispatch_tokens(std::int32_t tokens, Launch&& launch) {
    switch (tokens) {
#define NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(TOKENS)                                         \
    case TOKENS:                                                                                   \
        dispatch_token_case<TOKENS>(launch);                                                       \
        return
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(1);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(2);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(3);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(4);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(5);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(6);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(7);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(8);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(9);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(10);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(11);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(12);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(13);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(14);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(15);
        NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE(16);
#undef NINFER_SLIDING_WINDOW_ATTENTION_TOKEN_CASE
    default:
        throw std::invalid_argument("sliding_window_attention: unsupported T");
    }
}

} // namespace

SlidingWindowAttentionPlan
sliding_window_attention_resolve_plan(std::int32_t tokens,
                                      SlidingWindowAttentionExecutionEnvelope envelope) {
    if (tokens < 1 || tokens > 16) {
        throw std::invalid_argument("sliding_window_attention plan: T must be 1..16");
    }
    if (envelope.min_context > envelope.max_context) {
        throw std::invalid_argument("sliding_window_attention plan: invalid envelope");
    }
    // Graph envelopes whose longest context fits three key tiles avoid the second kernel and
    // workspace round trip. At four tiles, split-KV is already faster for every qualified T.
    constexpr std::uint32_t direct_context_limit = 96;
    const bool direct                            = envelope.max_context <= direct_context_limit;
    constexpr std::int32_t key_block             = 32;
    const std::uint32_t context_rows             = std::min(envelope.max_context, 4095u);
    const std::int32_t context_tiles =
        static_cast<std::int32_t>((context_rows + key_block - 1u) / key_block);
    constexpr std::int32_t split_limit = 32;
    return {
        .route =
            direct ? SlidingWindowAttentionRoute::Direct : SlidingWindowAttentionRoute::SplitKv,
        .tokens         = tokens,
        .warps          = (tokens + 3) / 4,
        .split_capacity = direct ? 1 : std::min(split_limit, std::max(1, context_tiles)),
        .max_context    = static_cast<std::int32_t>(envelope.max_context),
    };
}

const char* sliding_window_attention_route_name(SlidingWindowAttentionRoute route) {
    switch (route) {
    case SlidingWindowAttentionRoute::Direct:
        return "direct";
    case SlidingWindowAttentionRoute::SplitKv:
        return "split_kv";
    }
    return "unknown";
}

void sliding_window_attention_launch(const Tensor& q, const Tensor& query_k, const Tensor& query_v,
                                     const Tensor& positions, const Tensor& valid_columns,
                                     const Tensor& lanes, float scale,
                                     const CyclicKVCacheLayerView& context,
                                     const SlidingWindowAttentionPlan& plan, Tensor& partial_acc,
                                     Tensor& partial_m, Tensor& partial_l, Tensor& out,
                                     cudaStream_t stream) {
    dispatch_tokens(q.ne[2], [&]<int Tokens, int Warps>() {
        const bool direct = plan.route == SlidingWindowAttentionRoute::Direct;
        if (plan.warps != Warps || plan.split_capacity < 1 ||
            plan.split_capacity > kSlidingWindowMaxCandidateSplit ||
            (direct && plan.split_capacity != 1)) {
            throw std::invalid_argument("sliding_window_attention: inconsistent plan");
        }
        constexpr int KeyBlock = 32;
        constexpr std::size_t SmemBytes =
            2u * KeyBlock * kContextQueryHeadDim * sizeof(__nv_bfloat16);
        if (direct) {
            const dim3 direct_grid(kContextQueryKVHeads, 1, q.ne[3]);
            sliding_window_attention_split_partial_kernel<Tokens, Warps, KeyBlock, true>
                <<<direct_grid, Warps * 32, SmemBytes, stream>>>(
                    static_cast<const __nv_bfloat16*>(q.data),
                    static_cast<const __nv_bfloat16*>(query_k.data),
                    static_cast<const __nv_bfloat16*>(query_v.data),
                    static_cast<const std::int32_t*>(positions.data),
                    static_cast<const std::int32_t*>(valid_columns.data),
                    static_cast<const std::int32_t*>(lanes.data),
                    static_cast<const __nv_bfloat16*>(context.k.data),
                    static_cast<const __nv_bfloat16*>(context.v.data),
                    static_cast<int>(context.padded_capacity), plan.max_context, 1, scale,
                    static_cast<__nv_bfloat16*>(partial_acc.data),
                    static_cast<float*>(partial_m.data), static_cast<float*>(partial_l.data),
                    static_cast<__nv_bfloat16*>(out.data));
            CUDA_CHECK(cudaGetLastError());
            return;
        }

        const dim3 partial_grid(kContextQueryKVHeads, plan.split_capacity, q.ne[3]);
        sliding_window_attention_split_partial_kernel<Tokens, Warps, KeyBlock, false>
            <<<partial_grid, Warps * 32, SmemBytes, stream>>>(
                static_cast<const __nv_bfloat16*>(q.data),
                static_cast<const __nv_bfloat16*>(query_k.data),
                static_cast<const __nv_bfloat16*>(query_v.data),
                static_cast<const std::int32_t*>(positions.data),
                static_cast<const std::int32_t*>(valid_columns.data),
                static_cast<const std::int32_t*>(lanes.data),
                static_cast<const __nv_bfloat16*>(context.k.data),
                static_cast<const __nv_bfloat16*>(context.v.data),
                static_cast<int>(context.padded_capacity), plan.max_context, plan.split_capacity,
                scale, static_cast<__nv_bfloat16*>(partial_acc.data),
                static_cast<float*>(partial_m.data), static_cast<float*>(partial_l.data),
                static_cast<__nv_bfloat16*>(out.data));
        CUDA_CHECK(cudaGetLastError());

        constexpr int ReduceWarps = 1;
        constexpr int ReduceRows  = kContextQueryQHeads * Tokens;
        const dim3 reduce_grid((ReduceRows + ReduceWarps - 1) / ReduceWarps, 1, q.ne[3]);
        sliding_window_attention_reduce_kernel<Tokens, KeyBlock, ReduceWarps>
            <<<reduce_grid, ReduceWarps * 32, 0, stream>>>(
                static_cast<const __nv_bfloat16*>(partial_acc.data),
                static_cast<const float*>(partial_m.data),
                static_cast<const float*>(partial_l.data),
                static_cast<const std::int32_t*>(positions.data),
                static_cast<const std::int32_t*>(valid_columns.data), plan.max_context,
                plan.split_capacity, static_cast<__nv_bfloat16*>(out.data));
        CUDA_CHECK(cudaGetLastError());
    });
}

} // namespace ninfer::ops::detail
#include "ninfer/ops/sliding_window_attention.h"

#include "core/layout.h"
#include "ops/softmax_attention/sliding_window/launch.h"

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <limits>
#include <stdexcept>
#include <string>

namespace ninfer::ops {
namespace {

constexpr std::int32_t kHeadDim = 128;
constexpr std::int32_t kQHeads  = 32;
constexpr std::int32_t kKVHeads = 8;
constexpr std::int32_t kWindow  = 4096;
constexpr float kExpectedScale  = 0.08838834764831844055f;

void require_profile(AttentionHeadGeometry geometry, std::uint32_t window, const char* op) {
    if (!valid_attention_head_geometry(geometry) || geometry.head_dim != kHeadDim ||
        geometry.query_heads != kQHeads || geometry.kv_heads != kKVHeads) {
        throw std::invalid_argument(std::string(op) + ": unsupported head geometry");
    }
    if (window != static_cast<std::uint32_t>(kWindow)) {
        throw std::invalid_argument(std::string(op) + ": supported window is 4096");
    }
}

void require_shape(const Tensor& tensor, std::int32_t n0, std::int32_t n1, std::int32_t n2,
                   std::int32_t n3, const char* op, const char* name) {
    if (tensor.ne[0] != n0 || tensor.ne[1] != n1 || tensor.ne[2] != n2 || tensor.ne[3] != n3) {
        throw std::invalid_argument(std::string(op) + ": invalid shape for " + name);
    }
}

void require_contiguous_nonnull(const Tensor& tensor, const char* op, const char* name) {
    if (!tensor.is_contiguous()) {
        throw std::invalid_argument(std::string(op) + ": " + name + " must be contiguous");
    }
    if (tensor.data == nullptr) {
        throw std::invalid_argument(std::string(op) + ": " + name + " data must be non-null");
    }
}

void validate_context(const CyclicKVCacheLayerView& context, const char* op) {
    if (context.num_kv_heads != kKVHeads || context.head_dim != kHeadDim ||
        context.capacity != kWindow || context.padded_capacity < context.capacity ||
        context.lane_capacity <= 0) {
        throw std::invalid_argument(std::string(op) + ": invalid cyclic context");
    }
    if (context.padded_capacity >
        static_cast<std::uint32_t>(std::numeric_limits<std::int32_t>::max())) {
        throw std::overflow_error(std::string(op) + ": padded capacity exceeds int32");
    }
    const auto padded = static_cast<std::int32_t>(context.padded_capacity);
    if (context.k.dtype != DType::BF16 || context.v.dtype != DType::BF16) {
        throw std::invalid_argument(std::string(op) + ": context K/V must be BF16");
    }
    require_shape(context.k, kHeadDim, padded, kKVHeads, context.lane_capacity, op, "context k");
    require_shape(context.v, kHeadDim, padded, kKVHeads, context.lane_capacity, op, "context v");
    require_contiguous_nonnull(context.k, op, "context k");
    require_contiguous_nonnull(context.v, op, "context v");
}

struct PartialWorkspace {
    Tensor acc;
    Tensor m;
    Tensor l;
};

template <class Allocator>
PartialWorkspace allocate_workspace(Allocator& workspace, std::int32_t tokens, std::int32_t splits,
                                    std::int32_t batch_size) {
    return {
        workspace.alloc(DType::BF16, {kHeadDim, kQHeads, tokens, splits * batch_size}),
        workspace.alloc(DType::FP32, {kQHeads, tokens, splits * batch_size}),
        workspace.alloc(DType::FP32, {kQHeads, tokens, splits * batch_size}),
    };
}

} // namespace

std::size_t sliding_window_attention_workspace_capacity_bytes(
    AttentionHeadGeometry geometry, std::uint32_t window,
    SlidingWindowAttentionExecutionEnvelope envelope, std::int32_t min_tokens,
    std::int32_t max_tokens, std::int32_t batch_size) {
    require_profile(geometry, window, "sliding_window_attention workspace");
    if (min_tokens < 1 || max_tokens < min_tokens || max_tokens > 16 || batch_size < 1 ||
        batch_size > 8 || envelope.min_context > envelope.max_context ||
        envelope.max_context >
            static_cast<std::uint32_t>(std::numeric_limits<std::int32_t>::max())) {
        throw std::invalid_argument(
            "sliding_window_attention workspace: invalid envelope or token interval");
    }
    const auto plan = detail::sliding_window_attention_resolve_plan(max_tokens, envelope);
    WorkspaceLayoutBuilder layout;
    (void)allocate_workspace(layout, max_tokens, plan.split_capacity, batch_size);
    return layout.peak_bytes(1);
}

void sliding_window_attention(const Tensor& q, const Tensor& query_k, const Tensor& query_v,
                              const Tensor& positions, const Tensor& valid_columns,
                              const Tensor& lanes, AttentionHeadGeometry geometry,
                              std::uint32_t window, float scale,
                              const CyclicKVCacheLayerView& context,
                              SlidingWindowAttentionExecutionEnvelope envelope,
                              WorkspaceArena& workspace, Tensor& out, cudaStream_t stream) {
    constexpr const char* op = "sliding_window_attention";
    require_profile(geometry, window, op);
    if (q.dtype != DType::BF16 || query_k.dtype != DType::BF16 || query_v.dtype != DType::BF16 ||
        out.dtype != DType::BF16) {
        throw std::invalid_argument("sliding_window_attention: q/k/v/out must be BF16");
    }
    if (positions.dtype != DType::I32 || valid_columns.dtype != DType::I32 ||
        lanes.dtype != DType::I32) {
        throw std::invalid_argument(
            "sliding_window_attention: positions/valid_columns/lanes must be I32");
    }
    const std::int32_t tokens = q.ne[2];
    const std::int32_t batch  = q.ne[3];
    if (tokens < 1 || tokens > 16) {
        throw std::invalid_argument("sliding_window_attention: optimized domain is T=1..16");
    }
    if (batch < 1 || batch > 8) {
        throw std::invalid_argument("sliding_window_attention: B must be 1..8");
    }
    require_shape(q, kHeadDim, kQHeads, tokens, batch, op, "q");
    require_shape(query_k, kHeadDim, kKVHeads, tokens, batch, op, "query k");
    require_shape(query_v, kHeadDim, kKVHeads, tokens, batch, op, "query v");
    require_shape(positions, tokens, batch, 1, 1, op, "positions");
    require_shape(valid_columns, batch, 1, 1, 1, op, "valid columns");
    require_shape(lanes, batch, 1, 1, 1, op, "lanes");
    require_shape(out, kHeadDim, kQHeads, tokens, batch, op, "out");
    require_contiguous_nonnull(q, op, "q");
    require_contiguous_nonnull(query_k, op, "query k");
    require_contiguous_nonnull(query_v, op, "query v");
    require_contiguous_nonnull(positions, op, "positions");
    require_contiguous_nonnull(valid_columns, op, "valid columns");
    require_contiguous_nonnull(lanes, op, "lanes");
    require_contiguous_nonnull(out, op, "out");
    validate_context(context, op);
    if (envelope.min_context > envelope.max_context ||
        envelope.max_context >
            static_cast<std::uint32_t>(std::numeric_limits<std::int32_t>::max())) {
        throw std::invalid_argument("sliding_window_attention: invalid execution envelope");
    }
    if (!std::isfinite(scale) || std::abs(scale - kExpectedScale) > 1e-7f) {
        throw std::invalid_argument("sliding_window_attention: scale must be 1/sqrt(128)");
    }

    auto scope               = workspace.scope();
    const auto plan          = detail::sliding_window_attention_resolve_plan(tokens, envelope);
    PartialWorkspace partial = allocate_workspace(workspace, tokens, plan.split_capacity, batch);
    detail::sliding_window_attention_launch(q, query_k, query_v, positions, valid_columns, lanes,
                                            scale, context, plan, partial.acc, partial.m, partial.l,
                                            out, stream);
}

} // namespace ninfer::ops
#include "ops/linear_pair/w8/w8_pair_kernels.h"
#include "ops/linear_pair/w8/w8_pair_plan.h"

#include "core/device.h"
#include "ops/common/math.h"
#include "ops/linear/w8/w8_rowsplit_gemm_mma.cuh"

#include <cstdint>
#include <stdexcept>

namespace ninfer::ops::detail {
namespace {

constexpr int kRows   = 1024;
constexpr int kHidden = 2048;

struct PairOutputTile {
    __nv_bfloat16* data;
    std::int32_t parent_row_begin;
    std::int32_t parent_row_end;

    __device__ __forceinline__ __nv_bfloat16* at(std::int32_t parent_row, std::int32_t col) const {
        return data + static_cast<std::int64_t>(col) * kRows + parent_row - parent_row_begin;
    }

    __device__ __forceinline__ bool valid(std::int32_t parent_row,
                                          std::int32_t /*total_rows*/) const {
        return parent_row < parent_row_end;
    }
};

struct PairOutput {
    __nv_bfloat16* first;
    __nv_bfloat16* second;

    __device__ __forceinline__ std::int32_t row_begin(std::int32_t block,
                                                      std::int32_t tile_rows) const {
        const std::int32_t blocks_per_output = (kRows + tile_rows - 1) / tile_rows;
        const std::int32_t output_index      = block / blocks_per_output;
        return output_index * kRows + (block - output_index * blocks_per_output) * tile_rows;
    }

    __device__ __forceinline__ PairOutputTile tile(std::int32_t parent_row_begin) const {
        return parent_row_begin < kRows ? PairOutputTile{first, 0, kRows}
                                        : PairOutputTile{second, kRows, 2 * kRows};
    }
};

template <class Schedule, bool Full>
void launch_variant(const Tensor& x, const Weight& first_weight, Tensor& first_out,
                    Tensor& second_out, cudaStream_t stream) {
    const PairOutput output{static_cast<__nv_bfloat16*>(first_out.data),
                            static_cast<__nv_bfloat16*>(second_out.data)};
    const dim3 grid(static_cast<unsigned>(2 * div_up(kRows, Schedule::BM)),
                    static_cast<unsigned>(div_up(x.ne[1], Schedule::BN)), 1u);
    w8_rowsplit_gemm_mma_kernel<Schedule, Full, W8Epilogue::Store, PairOutput>
        <<<grid, Schedule::THREADS, 0, stream>>>(
            static_cast<const __nv_bfloat16*>(x.data),
            static_cast<const std::uint8_t*>(first_weight.qdata),
            static_cast<const std::uint8_t*>(first_weight.scales), output, 2 * kRows, kHidden,
            x.ne[1], kHidden);
}

template <class Schedule>
void dispatch_variant(bool full, const Tensor& x, const Weight& first_weight, Tensor& first_out,
                      Tensor& second_out, cudaStream_t stream) {
    if (full) {
        launch_variant<Schedule, true>(x, first_weight, first_out, second_out, stream);
    } else {
        launch_variant<Schedule, false>(x, first_weight, first_out, second_out, stream);
    }
    CUDA_CHECK(cudaGetLastError());
}

void require_adjacent(const Weight& first_weight, const Weight& second_weight) {
    if (static_cast<const std::uint8_t*>(second_weight.qdata) !=
            static_cast<const std::uint8_t*>(first_weight.qdata) + kRows * kHidden ||
        static_cast<const std::uint8_t*>(second_weight.scales) !=
            static_cast<const std::uint8_t*>(first_weight.scales) + kRows * (kHidden / 32) * 2) {
        throw std::invalid_argument("W8 concatenated pair MMA requires adjacent K/V row views");
    }
}

} // namespace

void w8_pair_concat_mma_launch(W8PairScheduleId schedule, bool full, const Tensor& x,
                               const Weight& first_weight, const Weight& second_weight,
                               Tensor& first_out, Tensor& second_out, cudaStream_t stream) {
    require_adjacent(first_weight, second_weight);
    switch (schedule) {
    case W8PairScheduleId::ConcatMmaR32C64:
        dispatch_variant<W8RowSplitMmaGemmSchedule<32, 64, 32, 16, 3>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR32C80:
        dispatch_variant<W8RowSplitMmaGemmSchedule<32, 80, 32, 16, 3>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR32C96:
        dispatch_variant<W8RowSplitMmaGemmSchedule<32, 96, 32, 16, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR32C112:
        dispatch_variant<W8RowSplitMmaGemmSchedule<32, 112, 32, 16, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR32C128:
        dispatch_variant<W8RowSplitMmaGemmSchedule<32, 128, 32, 16, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR48C64:
        dispatch_variant<W8RowSplitMmaGemmSchedule<48, 64, 48, 16, 3, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR48C96:
        dispatch_variant<W8RowSplitMmaGemmSchedule<48, 96, 48, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR48C112:
        dispatch_variant<W8RowSplitMmaGemmSchedule<48, 112, 48, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR48C128:
        dispatch_variant<W8RowSplitMmaGemmSchedule<48, 128, 48, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR64C64:
        dispatch_variant<W8RowSplitMmaGemmSchedule<64, 64, 64, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR64C80:
        dispatch_variant<W8RowSplitMmaGemmSchedule<64, 80, 64, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR64C96:
        dispatch_variant<W8RowSplitMmaGemmSchedule<64, 96, 64, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR64C128:
        dispatch_variant<W8RowSplitMmaGemmSchedule<64, 128, 64, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR96C64:
        dispatch_variant<W8RowSplitMmaGemmSchedule<96, 64, 48, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR96C80:
        dispatch_variant<W8RowSplitMmaGemmSchedule<96, 80, 48, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR96C96:
        dispatch_variant<W8RowSplitMmaGemmSchedule<96, 96, 48, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR96C112:
        dispatch_variant<W8RowSplitMmaGemmSchedule<96, 112, 48, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR128C64:
        dispatch_variant<W8RowSplitMmaGemmSchedule<128, 64, 64, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    case W8PairScheduleId::ConcatMmaR128C80:
        dispatch_variant<W8RowSplitMmaGemmSchedule<128, 80, 64, 16, 2, 2>>(
            full, x, first_weight, first_out, second_out, stream);
        return;
    default:
        throw std::invalid_argument("W8 concatenated pair MMA schedule is not supported");
    }
}

} // namespace ninfer::ops::detail
#pragma once

#include "core/tensor.h"

#include <cuda_runtime.h>

namespace ninfer::ops {

enum class GeluMode {
    Exact,
    Tanh,
};

/**
 * Applies GELU elementwise in place. For z=x[i]:
 *
 *   Exact: ideal[i] = 0.5*z*(1 + erf(z/sqrt(2)))
 *   Tanh:  ideal[i] = 0.5*z*(1 + tanh(sqrt(2/pi)*(z + 0.044715*z^3))).
 *
 * `x` is an arbitrary-rank contiguous BF16 tensor. The oracle evaluates `ideal` in FP64 from the
 * represented input. The updated BF16 x is promoted and compared directly with that result; output
 * storage rounding belongs to the Op's numerical criterion, not the oracle. Private kernel
 * arithmetic is implementation-defined. The Op mutates only x and uses no workspace or persistent
 * state.
 */
void gelu(Tensor& x, GeluMode mode, cudaStream_t stream);

} // namespace ninfer::ops
#include "ninfer/ops/sparse_moe.h"

#include "core/nvtx.h"
#include "ops/sparse_moe/decode/sparse_moe_decode.h"
#include "ops/sparse_moe/prefill/sparse_moe_prefill.h"
#include "ops/sparse_moe/small_t/sparse_moe_small_t.h"

#include <algorithm>
#include <cstdint>
#include <limits>
#include <stdexcept>
#include <string>
#include <vector>

namespace ninfer::ops {
namespace {

constexpr std::int32_t kHidden         = 2048;
constexpr std::int32_t kExperts        = 256;
constexpr std::int32_t kRouterRows     = kExperts + 1;
constexpr std::int32_t kExpertRows     = 1024;
constexpr std::int32_t kIntermediate   = 512;
constexpr std::int32_t kRoutedGateRows = kExperts * kExpertRows;
constexpr std::int32_t kRoutedDownRows = kExperts * kHidden;
constexpr std::int32_t kSharedGateRows = 2 * kIntermediate;

struct AddressRange {
    std::uintptr_t begin = 0;
    std::uintptr_t end   = 0;
    std::string name;
};

bool aligned_to(const void* pointer, std::uintptr_t alignment) {
    return pointer != nullptr && (reinterpret_cast<std::uintptr_t>(pointer) & (alignment - 1)) == 0;
}

AddressRange address_range(const void* pointer, std::size_t bytes, std::string name) {
    if (pointer == nullptr || bytes == 0) {
        throw std::invalid_argument("sparse_moe: " + name + " storage must be non-empty");
    }
    const std::uintptr_t begin = reinterpret_cast<std::uintptr_t>(pointer);
    if (bytes > std::numeric_limits<std::uintptr_t>::max() - begin) {
        throw std::overflow_error("sparse_moe: " + name + " address range overflows");
    }
    return {begin, begin + bytes, std::move(name)};
}

void require_disjoint(const std::vector<AddressRange>& ranges) {
    for (std::size_t i = 0; i < ranges.size(); ++i) {
        for (std::size_t j = i + 1; j < ranges.size(); ++j) {
            if (ranges[i].begin < ranges[j].end && ranges[j].begin < ranges[i].end) {
                throw std::invalid_argument("sparse_moe: " + ranges[i].name + " overlaps " +
                                            ranges[j].name);
            }
        }
    }
}

std::int32_t require_tensor(const Tensor& tensor, const char* name) {
    if (tensor.dtype != DType::BF16 || tensor.ne[0] != kHidden || tensor.ne[1] < 1 ||
        tensor.ne[2] != 1 || tensor.ne[3] != 1 || !tensor.is_contiguous() ||
        tensor.data == nullptr || !aligned_to(tensor.data, 16)) {
        throw std::invalid_argument(std::string("sparse_moe: invalid ") + name);
    }
    return tensor.ne[1];
}

void require_matrix_metadata(const Weight& weight, std::int32_t n, std::int32_t k,
                             const char* name) {
    if (weight.ndim != 2 || weight.n != n || weight.k != k || weight.shape[0] != n ||
        weight.shape[1] != k || weight.shape[2] != 1 || weight.shape[3] != 1 ||
        weight.padded_shape[0] != n || weight.padded_shape[1] != k || weight.padded_shape[2] != 1 ||
        weight.padded_shape[3] != 1) {
        throw std::invalid_argument(std::string("sparse_moe: invalid shape for ") + name);
    }
}

void require_router(const Weight& weight, std::vector<AddressRange>& ranges) {
    require_matrix_metadata(weight, kRouterRows, kHidden, "router_shared_gate");
    const std::size_t bytes = static_cast<std::size_t>(kRouterRows) * kHidden * 2;
    if (weight.qtype != QType::BF16_CTRL || weight.layout != QuantLayout::Contiguous ||
        weight.qdata == nullptr || weight.qhigh != nullptr || weight.scales != nullptr ||
        weight.payload_bytes < bytes || !aligned_to(weight.qdata, 16)) {
        throw std::invalid_argument(
            "sparse_moe: router_shared_gate must be aligned contiguous BF16");
    }
    ranges.push_back(address_range(weight.qdata, bytes, "router_shared_gate"));
}

struct QuantGeometry {
    std::int32_t group_size;
    std::size_t code_bytes_per_group;
    std::size_t high_bytes_per_group;
};

QuantGeometry quant_geometry(QType qtype) {
    switch (qtype) {
    case QType::Q4G64_F16S:
        return {64, 32, 0};
    case QType::Q5G64_F16S:
        return {64, 32, 8};
    case QType::Q6G64_F16S:
        return {64, 32, 16};
    case QType::W8G32_F16S:
        return {32, 32, 0};
    default:
        throw std::invalid_argument("sparse_moe: unsupported quantized weight format");
    }
}

void require_quantized(const Weight& weight, std::int32_t n, std::int32_t k, const char* name,
                       std::vector<AddressRange>& ranges) {
    require_matrix_metadata(weight, n, k, name);
    const QuantGeometry geometry       = quant_geometry(weight.qtype);
    const std::size_t groups           = static_cast<std::size_t>(n) * k / geometry.group_size;
    const std::size_t code_bytes       = groups * geometry.code_bytes_per_group;
    const std::size_t high_bytes       = groups * geometry.high_bytes_per_group;
    const std::size_t scale_bytes      = groups * 2;
    const std::size_t required_payload = code_bytes + high_bytes + scale_bytes;
    if (weight.layout != QuantLayout::RowSplit || weight.scale_dtype != DType::FP16 ||
        weight.group_size != static_cast<std::uint32_t>(geometry.group_size) ||
        weight.group != geometry.group_size || weight.qdata == nullptr ||
        weight.scales == nullptr || weight.payload_bytes < required_payload ||
        weight.high_plane_bytes < high_bytes || !aligned_to(weight.qdata, 16) ||
        !aligned_to(weight.scales, 16)) {
        throw std::invalid_argument(std::string("sparse_moe: invalid row-split ") + name);
    }
    if ((high_bytes == 0 && weight.qhigh != nullptr) ||
        (high_bytes != 0 && (weight.qhigh == nullptr || !aligned_to(weight.qhigh, 16)))) {
        throw std::invalid_argument(std::string("sparse_moe: invalid high plane for ") + name);
    }
    ranges.push_back(address_range(weight.qdata, code_bytes, std::string(name) + " code"));
    if (high_bytes != 0) {
        ranges.push_back(address_range(weight.qhigh, high_bytes, std::string(name) + " high"));
    }
    ranges.push_back(address_range(weight.scales, scale_bytes, std::string(name) + " scales"));
}

void validate_weights(const SparseMoeWeights& weights, std::vector<AddressRange>& ranges) {
    require_router(weights.router_shared_gate, ranges);
    if (weights.routed_gate_up.qtype != QType::Q4G64_F16S &&
        weights.routed_gate_up.qtype != QType::W8G32_F16S) {
        throw std::invalid_argument("sparse_moe: routed_gate_up must be Q4 or W8");
    }
    if (weights.routed_down.qtype != QType::Q5G64_F16S &&
        weights.routed_down.qtype != QType::Q6G64_F16S &&
        weights.routed_down.qtype != QType::W8G32_F16S) {
        throw std::invalid_argument("sparse_moe: routed_down must be Q5, Q6, or W8");
    }
    if (weights.shared_gate_up.qtype != QType::W8G32_F16S ||
        weights.shared_down.qtype != QType::W8G32_F16S) {
        throw std::invalid_argument("sparse_moe: shared weights must be W8");
    }
    require_quantized(weights.routed_gate_up, kRoutedGateRows, kHidden, "routed_gate_up", ranges);
    require_quantized(weights.routed_down, kRoutedDownRows, kIntermediate, "routed_down", ranges);
    require_quantized(weights.shared_gate_up, kSharedGateRows, kHidden, "shared_gate_up", ranges);
    require_quantized(weights.shared_down, kHidden, kIntermediate, "shared_down", ranges);
}

} // namespace

std::size_t sparse_moe_workspace_capacity_bytes(QType routed_gate_up, QType routed_down,
                                                std::int32_t min_tokens, std::int32_t max_tokens) {
    if (min_tokens <= 0 || max_tokens < min_tokens) {
        throw std::invalid_argument("sparse_moe workspace: invalid token interval");
    }
    (void)detail::resolve_sparse_moe_decode_plan(routed_gate_up, routed_down);

    const bool w8_profile = routed_gate_up == QType::W8G32_F16S && routed_down == QType::W8G32_F16S;
    const std::int32_t prefill_first =
        w8_profile ? detail::kSparseMoePrefillW8W8Min
                   : (routed_down == QType::Q5G64_F16S ? detail::kSparseMoePrefillQ4Q5Min
                                                       : detail::kSparseMoePrefillQ4Q6Min);

    std::size_t required = 0;
    if (min_tokens == 1) { required = detail::sparse_moe_decode_workspace_bytes(); }

    const std::int32_t small_first = std::max(min_tokens, detail::kSparseMoeSmallTMin);
    const std::int32_t small_last =
        std::min({max_tokens, detail::kSparseMoeSmallTMax, prefill_first - 1});
    if (small_first <= small_last) {
        required = std::max(required, detail::sparse_moe_small_t_workspace_bytes(small_last));
    }

    const std::int32_t prefill_interval_first = std::max(min_tokens, prefill_first);
    if (prefill_interval_first <= max_tokens) {
        required = std::max(required, detail::sparse_moe_prefill_workspace_bytes(max_tokens));
    }
    return required;
}

void sparse_moe(const Tensor& x, const SparseMoeWeights& weights, SparseMoeEpilogue epilogue,
                Tensor& destination, WorkspaceArena& workspace, cudaStream_t stream) {
    if (epilogue != SparseMoeEpilogue::AddResidual) {
        throw std::invalid_argument("sparse_moe: unsupported epilogue");
    }
    const std::int32_t tokens = require_tensor(x, "x");
    if (require_tensor(destination, "destination") != tokens) {
        throw std::invalid_argument("sparse_moe: x and destination token counts must match");
    }

    std::vector<AddressRange> ranges;
    ranges.reserve(16);
    ranges.push_back(address_range(x.data, x.bytes(), "x"));
    ranges.push_back(address_range(destination.data, destination.bytes(), "destination"));
    validate_weights(weights, ranges);

    const bool use_small_t = detail::sparse_moe_uses_small_t(tokens);
    const bool use_prefill = detail::sparse_moe_uses_prefill(tokens, weights.routed_gate_up.qtype,
                                                             weights.routed_down.qtype);
    nvtx::ScopedRange moe_range(use_prefill   ? nvtx::Name::SparseMoePrefill
                                : use_small_t ? nvtx::Name::SparseMoeSmallT
                                              : nvtx::Name::SparseMoeDecode,
                                nvtx::Category::Moe, static_cast<std::uint64_t>(tokens));
    std::size_t required = 0;
    if (use_prefill) {
        required = detail::resolve_sparse_moe_prefill_plan(tokens, weights.routed_gate_up.qtype,
                                                           weights.routed_down.qtype)
                       .workspace_bytes;
    } else if (use_small_t) {
        required = detail::resolve_sparse_moe_small_t_plan(tokens, weights.routed_gate_up.qtype,
                                                           weights.routed_down.qtype)
                       .workspace_bytes;
    } else {
        required = detail::resolve_sparse_moe_decode_plan(weights.routed_gate_up.qtype,
                                                          weights.routed_down.qtype)
                       .workspace_bytes;
    }
    if (workspace.base() == nullptr || workspace.capacity() < required ||
        workspace.used() > workspace.capacity() - required) {
        throw std::invalid_argument("sparse_moe: insufficient workspace capacity");
    }
    ranges.push_back(address_range(workspace.base(), workspace.capacity(), "workspace"));
    require_disjoint(ranges);

    auto scope = workspace.scope();
    if (use_prefill) {
        const detail::SparseMoePrefillPlan plan = detail::resolve_sparse_moe_prefill_plan(
            tokens, weights.routed_gate_up.qtype, weights.routed_down.qtype);
        const detail::SparseMoePrefillWorkspace views =
            detail::allocate_sparse_moe_prefill_workspace(workspace, plan.slice_tokens);
        detail::sparse_moe_prefill_launch(x, weights, destination, plan, views, stream);
        return;
    }
    if (use_small_t) {
        const detail::SparseMoeSmallTPlan plan = detail::resolve_sparse_moe_small_t_plan(
            tokens, weights.routed_gate_up.qtype, weights.routed_down.qtype);
        const detail::SparseMoeSmallTWorkspace views =
            detail::allocate_sparse_moe_small_t_workspace(workspace, tokens);
        detail::sparse_moe_small_t_launch(x, weights, destination, plan, views, stream);
        return;
    }

    const detail::SparseMoeDecodePlan plan = detail::resolve_sparse_moe_decode_plan(
        weights.routed_gate_up.qtype, weights.routed_down.qtype);
    const detail::SparseMoeDecodeWorkspace views =
        detail::allocate_sparse_moe_decode_workspace(workspace);
    for (std::int32_t token = 0; token < tokens; ++token) {
        const Tensor x_column     = x.slice(1, token, 1);
        Tensor destination_column = destination.slice(1, token, 1);
        detail::sparse_moe_decode_launch(x_column, weights, destination_column, views, stream);
    }
}

} // namespace ninfer::ops
#pragma once

// Adapter between the wire-independent protocol request and the public engine API.

#include "ninfer/types.h"
#include "serve/request.h"
#include "serve/serve_options.h"

#include <functional>

namespace ninfer::serve {

// Media acquisition is a product-layer concern. Translation preserves part order
// and asks the caller to turn each wire source into owning bytes before the
// target frontend sees it.
using MediaAcquirer = std::function<ninfer::OwnedMedia(const ContentPart&)>;

struct ResolvedPromptSemantics {
    bool enable_thinking = true;
    std::optional<ninfer::ReasoningEffort> reasoning_effort;
    std::optional<ninfer::ReasoningEffort> effective_reasoning_effort;
    bool preserve_thinking = false;
};

ResolvedPromptSemantics resolve_prompt_semantics(const GenerationRequest& req,
                                                 const ServeOptions& server,
                                                 const ninfer::PromptCapabilities& capabilities);

ninfer::PromptInput to_prompt_input(const GenerationRequest& req,
                                    const ResolvedPromptSemantics& semantics,
                                    const MediaAcquirer& acquire_media);

// Build public request options (output budget, thinking, stop policy, sampler). The
// sampler is resolved from the request's SamplingParams over the server defaults;
// --greedy on the server forces exact argmax regardless of the request. Prefix-reuse
// participation is resolved by GenerationService and supplied explicitly because
// operational requests do not inherit the external-traffic policy.
ninfer::RequestOptions to_request_options(const GenerationRequest& req, const ServeOptions& server,
                                          const ResolvedPromptSemantics& semantics,
                                          bool allow_prefix_reuse);

} // namespace ninfer::serve
#include "core/gdn_replay_records.h"

#include <array>
#include <limits>
#include <stdexcept>
#include <string>

namespace ninfer {
namespace {

constexpr std::size_t kRecordAlignment        = 256;
constexpr std::int32_t kMaximumRecordCapacity = 8;

void require_positive(std::int32_t value, const char* label) {
    if (value <= 0) {
        throw std::invalid_argument(std::string("GDN replay ") + label + " must be positive");
    }
}

std::int32_t checked_outer_extent(const GdnReplayRecordSpec& spec) {
    const auto layers   = static_cast<std::int64_t>(spec.layers);
    const auto capacity = static_cast<std::int64_t>(spec.record_capacity);
    if (layers > std::numeric_limits<std::int32_t>::max() / capacity) {
        throw std::overflow_error("GDN replay layer/row extent exceeds int32");
    }
    return static_cast<std::int32_t>(layers * capacity);
}

void validate_spec(const GdnReplayRecordSpec& spec) {
    require_positive(spec.layers, "layer count");
    require_positive(spec.record_capacity, "record capacity");
    require_positive(spec.width, "width");
    require_positive(spec.conv_channels, "convolution channel count");
    require_positive(spec.qk_heads, "Q/K head count");
    require_positive(spec.value_heads, "value head count");
    require_positive(spec.key_dim, "key dimension");
    require_positive(spec.value_dim, "value dimension");
    if (spec.record_capacity > kMaximumRecordCapacity) {
        throw std::invalid_argument("GDN replay record capacity exceeds eight rows");
    }
    if (spec.value_heads % spec.qk_heads != 0) {
        throw std::invalid_argument("GDN replay value heads must be grouped by Q/K heads");
    }
    (void)checked_outer_extent(spec);
}

void require_region(const TensorRegion& region, DType dtype,
                    const std::array<std::int32_t, 4>& shape, const char* label) {
    if (region.dtype != dtype || region.shape != shape ||
        region.region.alignment != kRecordAlignment) {
        throw std::logic_error(std::string("GDN replay ") + label +
                               " layout does not match its spec");
    }
    const Tensor expected(nullptr, dtype, {shape[0], shape[1], shape[2], shape[3]});
    if (region.region.bytes != expected.bytes()) {
        throw std::logic_error(std::string("GDN replay ") + label +
                               " layout has an inconsistent byte size");
    }
}

void require_disjoint(const TensorRegion& a, const TensorRegion& b) {
    const std::size_t a_begin = a.region.offset;
    const std::size_t b_begin = b.region.offset;
    if (a.region.bytes > std::numeric_limits<std::size_t>::max() - a_begin ||
        b.region.bytes > std::numeric_limits<std::size_t>::max() - b_begin) {
        throw std::overflow_error("GDN replay region end overflows size_t");
    }
    const std::size_t a_end = a_begin + a.region.bytes;
    const std::size_t b_end = b_begin + b.region.bytes;
    if (a_begin < b_end && b_begin < a_end) {
        throw std::logic_error("GDN replay record planes overlap");
    }
}

void validate_layout(const GdnReplayRecordLayout& layout) {
    validate_spec(layout.spec);
    const std::int32_t outer = checked_outer_extent(layout.spec);
    require_region(layout.conv, DType::BF16,
                   {layout.spec.conv_channels, layout.spec.width, outer, 1}, "conv");
    require_region(layout.key, DType::BF16,
                   {layout.spec.key_dim, layout.spec.qk_heads, layout.spec.width, outer}, "key");
    require_region(layout.value, DType::BF16,
                   {layout.spec.value_dim, layout.spec.value_heads, layout.spec.width, outer},
                   "value");
    require_region(layout.gate, DType::FP32, {2, layout.spec.value_heads, layout.spec.width, outer},
                   "gate");

    const TensorRegion* regions[] = {&layout.conv, &layout.key, &layout.value, &layout.gate};
    for (std::size_t i = 0; i < std::size(regions); ++i) {
        for (std::size_t j = 0; j < i; ++j) { require_disjoint(*regions[i], *regions[j]); }
    }
}

} // namespace

GdnReplayRecordLayout plan_gdn_replay_records(LayoutBuilder& builder,
                                              const GdnReplayRecordSpec& spec) {
    validate_spec(spec);
    const std::int32_t outer = checked_outer_extent(spec);

    GdnReplayRecordLayout layout;
    layout.spec = spec;
    layout.conv = builder.add_tensor(DType::BF16, {spec.conv_channels, spec.width, outer},
                                     kRecordAlignment, "GDN replay conv records");
    layout.key  = builder.add_tensor(DType::BF16, {spec.key_dim, spec.qk_heads, spec.width, outer},
                                     kRecordAlignment, "GDN replay key records");
    layout.value =
        builder.add_tensor(DType::BF16, {spec.value_dim, spec.value_heads, spec.width, outer},
                           kRecordAlignment, "GDN replay value records");
    layout.gate = builder.add_tensor(DType::FP32, {2, spec.value_heads, spec.width, outer},
                                     kRecordAlignment, "GDN replay gate records");
    return layout;
}

std::size_t GdnReplayRecordLayout::payload_bytes() const noexcept {
    return conv.region.bytes + key.region.bytes + value.region.bytes + gate.region.bytes;
}

GdnReplayRecords::GdnReplayRecords(DeviceSpan backing, const GdnReplayRecordLayout& layout)
    : conv(layout.conv.bind(backing)), key(layout.key.bind(backing)),
      value(layout.value.bind(backing)), gate(layout.gate.bind(backing)), spec(layout.spec) {
    validate_layout(layout);
}

GdnReplayRecordLayer GdnReplayRecords::layer(std::int32_t layer_index, std::int32_t rows) const {
    validate_spec(spec);
    if (layer_index < 0 || layer_index >= spec.layers) {
        throw std::out_of_range("GDN replay layer index out of range");
    }
    if (rows <= 0 || rows > spec.record_capacity) {
        throw std::out_of_range("GDN replay active row count out of range");
    }
    const std::int32_t outer_begin = layer_index * spec.record_capacity;
    return GdnReplayRecordLayer{
        .conv  = conv.slice(2, outer_begin, rows).view({spec.conv_channels, spec.width, rows}),
        .key   = key.slice(3, outer_begin, rows),
        .value = value.slice(3, outer_begin, rows),
        .gate  = gate.slice(3, outer_begin, rows),
    };
}

} // namespace ninfer
// ninfer::ops - gdn_gating wrapper: public api validation and launcher dispatch.
#include "ninfer/ops/gdn_gating.h"

#include "ops/launcher/gdn_gating.h" // detail::gdn_gating_launch

#include <cstdint>
#include <limits>
#include <stdexcept>
#include <string>

namespace ninfer::ops {
namespace {

std::int64_t numel_allow_zero(const Tensor& t, const char* label) {
    bool has_zero = false;
    for (int d = 0; d < 4; ++d) {
        if (t.ne[d] < 0) {
            throw std::invalid_argument(std::string("gdn_gating: ") + label +
                                        " dimensions must be nonnegative");
        }
        if (t.ne[d] == 0) { has_zero = true; }
    }
    if (has_zero) { return 0; }

    std::int64_t total = 1;
    for (int d = 0; d < 4; ++d) {
        if (total > std::numeric_limits<std::int64_t>::max() / t.ne[d]) {
            throw std::overflow_error("gdn_gating: tensor size overflows int64");
        }
        total *= t.ne[d];
    }
    return total;
}

void require_gate_shape(const Tensor& t, const char* label) {
    if (t.ne[0] != 48 || t.ne[2] != 1 || t.ne[3] != 1) {
        throw std::invalid_argument(std::string("gdn_gating: ") + label +
                                    " must have shape [48,T]");
    }
}

void require_vector48_shape(const Tensor& t, const char* label) {
    if (t.ne[0] != 48 || t.ne[1] != 1 || t.ne[2] != 1 || t.ne[3] != 1) {
        throw std::invalid_argument(std::string("gdn_gating: ") + label + " must have shape [48]");
    }
}

void require_same_gate_shape(const Tensor& ref, const Tensor& t, const char* label) {
    for (int d = 0; d < 4; ++d) {
        if (ref.ne[d] != t.ne[d]) {
            throw std::invalid_argument(std::string("gdn_gating: ") + label +
                                        " shape must match a");
        }
    }
}

} // namespace

void gdn_gating(const Tensor& a, const Tensor& b, const Tensor& A_log, const Tensor& dt_bias,
                Tensor& g, Tensor& beta, cudaStream_t stream) {
    if (a.dtype != DType::BF16 || b.dtype != DType::BF16) {
        throw std::invalid_argument("gdn_gating: a/b must be BF16");
    }
    if (A_log.dtype != DType::FP32 || dt_bias.dtype != DType::FP32) {
        throw std::invalid_argument("gdn_gating: A_log/dt_bias must be FP32");
    }
    if (g.dtype != DType::FP32 || beta.dtype != DType::FP32) {
        throw std::invalid_argument("gdn_gating: g/beta must be FP32");
    }

    const std::int64_t n = numel_allow_zero(a, "a");
    (void)numel_allow_zero(b, "b");
    (void)numel_allow_zero(A_log, "A_log");
    (void)numel_allow_zero(dt_bias, "dt_bias");
    (void)numel_allow_zero(g, "g");
    (void)numel_allow_zero(beta, "beta");

    require_gate_shape(a, "a");
    require_gate_shape(b, "b");
    require_gate_shape(g, "g");
    require_gate_shape(beta, "beta");
    require_vector48_shape(A_log, "A_log");
    require_vector48_shape(dt_bias, "dt_bias");
    require_same_gate_shape(a, b, "b");
    require_same_gate_shape(a, g, "g");
    require_same_gate_shape(a, beta, "beta");
    if (n == 0) { return; }

    if (!a.is_contiguous() || !b.is_contiguous() || !A_log.is_contiguous() ||
        !dt_bias.is_contiguous() || !g.is_contiguous() || !beta.is_contiguous()) {
        throw std::invalid_argument("gdn_gating: all tensors must be contiguous");
    }
    if (a.data == nullptr || b.data == nullptr || A_log.data == nullptr ||
        dt_bias.data == nullptr || g.data == nullptr || beta.data == nullptr) {
        throw std::invalid_argument("gdn_gating: all tensor data pointers must be non-null");
    }

    detail::gdn_gating_launch(a, b, A_log, dt_bias, g, beta, stream);
}

} // namespace ninfer::ops
#pragma once

#include "ops/common/memory.cuh"

#include <cuda_bf16.h>

#include <cstdint>

namespace ninfer::ops::detail {

struct Nvfp4GdnInputOutput {
    static constexpr std::int32_t kQkvRows = 10240;
    static constexpr std::int32_t kZRows   = 6144;

    __nv_bfloat16* qkv;
    __nv_bfloat16* z;

    __device__ __forceinline__ __nv_bfloat16* destination(std::int32_t parent_row,
                                                          std::int32_t token) const {
        if (parent_row < kQkvRows) {
            return qkv + static_cast<std::int64_t>(token) * kQkvRows + parent_row;
        }
        return z + static_cast<std::int64_t>(token) * kZRows + parent_row - kQkvRows;
    }

    __device__ __forceinline__ void store(std::int32_t parent_row, std::int32_t token,
                                          float value) const {
        *destination(parent_row, token) = __float2bfloat16_rn(value);
    }

    __device__ __forceinline__ void store_vector(std::int32_t parent_row, std::int32_t token,
                                                 uint4 values) const {
        store_vec(destination(parent_row, token), values);
    }
};

static_assert((Nvfp4GdnInputOutput::kQkvRows % 128) == 0);
static_assert((Nvfp4GdnInputOutput::kZRows % 128) == 0);

} // namespace ninfer::ops::detail
#pragma once

#include "ops/linear/nvfp4/nvfp4_config.h"
#include "ops/linear/nvfp4/nvfp4_codec.cuh"
#include "ops/linear/nvfp4/nvfp4_output.cuh"
#include "ops/common/math.cuh"
#include "ops/common/memory.cuh"
#include "ops/common/warp.cuh"

#include <cuda_bf16.h>
#include <cuda_runtime.h>

#include <cstdint>

namespace ninfer::ops::detail {

template <int Values>
struct alignas(Values / 2) Nvfp4CodePack {
    static_assert(Values == 8 || Values == 16 || Values == 32);
    std::uint32_t words[Values / 8];
};

static_assert(sizeof(Nvfp4CodePack<8>) == 4);
static_assert(sizeof(Nvfp4CodePack<16>) == 8);
static_assert(sizeof(Nvfp4CodePack<32>) == 16);

template <Nvfp4CodeCache Cache, int Values>
__device__ __forceinline__ Nvfp4CodePack<Values> load_nvfp4_codes(const std::uint8_t* pointer) {
    if constexpr (Cache == Nvfp4CodeCache::Default) {
        return load_vec<Nvfp4CodePack<Values>>(pointer);
    } else if constexpr (Values == 8) {
        Nvfp4CodePack<Values> result;
        asm volatile("ld.global.cg.u32 %0, [%1];\n" : "=r"(result.words[0]) : "l"(pointer));
        return result;
    } else if constexpr (Values == 16) {
        Nvfp4CodePack<Values> result;
        asm volatile("ld.global.cg.v2.u32 {%0, %1}, [%2];\n"
                     : "=r"(result.words[0]), "=r"(result.words[1])
                     : "l"(pointer));
        return result;
    } else {
        Nvfp4CodePack<Values> result;
        asm volatile("ld.global.cg.v4.u32 {%0, %1, %2, %3}, [%4];\n"
                     : "=r"(result.words[0]), "=r"(result.words[1]), "=r"(result.words[2]),
                       "=r"(result.words[3])
                     : "l"(pointer));
        return result;
    }
}

template <class Geometry, class Schedule>
struct Nvfp4GemvSharedStorage {
    static constexpr int kRawScaleBytes = Schedule::kScaleAccess == Nvfp4ScaleAccess::StagedRaw
                                              ? Schedule::kRowsPerCta * Geometry::kGroupsPerRow
                                              : 16;
    alignas(16) std::uint8_t raw_scales[kRawScaleBytes];
};

template <class Geometry, class Schedule>
__device__ __forceinline__ void
stage_nvfp4_scales(const std::uint8_t* __restrict__ scales,
                   Nvfp4GemvSharedStorage<Geometry, Schedule>& shared, int m_tile, int rmod_base) {
    if constexpr (Schedule::kScaleAccess == Nvfp4ScaleAccess::StagedRaw) {
        constexpr int kQuartetsPerCta = Schedule::kRowsPerCta / 4;
        constexpr int kTasks          = Geometry::kScaleTilesPerRow * kQuartetsPerCta;
        for (int task = static_cast<int>(threadIdx.x); task < kTasks; task += Schedule::kThreads) {
            const int scale_tile = task / kQuartetsPerCta;
            const int quartet    = task - scale_tile * kQuartetsPerCta;
            const std::int64_t source_offset =
                static_cast<std::int64_t>(m_tile * Geometry::kScaleTilesPerRow + scale_tile) * 512 +
                static_cast<std::int64_t>(rmod_base + quartet) * 16;
            const uint4 packed           = load_vec<uint4>(scales + source_offset);
            const std::uint32_t words[4] = {packed.x, packed.y, packed.z, packed.w};
#pragma unroll
            for (int quartile = 0; quartile < 4; ++quartile) {
                const int local_row = quartet * 4 + quartile;
                auto* destination   = reinterpret_cast<std::uint32_t*>(
                    shared.raw_scales + local_row * Geometry::kGroupsPerRow + scale_tile * 4);
                *destination = words[quartile];
            }
        }
        __syncthreads();
    }
}

template <class Geometry>
__device__ __forceinline__ std::int64_t nvfp4_scale_offset(int parent_row, int group) {
    const int m_tile       = parent_row / 128;
    const int row_inner    = parent_row - m_tile * 128;
    const int scale_tile   = group / 4;
    const int scale_lane   = group & 3;
    const int row_mod32    = row_inner & 31;
    const int row_quartile = row_inner >> 5;
    return static_cast<std::int64_t>(m_tile * Geometry::kScaleTilesPerRow + scale_tile) * 512 +
           row_mod32 * 16 + row_quartile * 4 + scale_lane;
}

template <class Geometry, class Schedule>
__device__ __forceinline__ std::uint32_t
load_staged_scale_word(const Nvfp4GemvSharedStorage<Geometry, Schedule>& shared, int local_row,
                       int phase, int lane) {
    constexpr int kSubgroupWidth  = 64 / Schedule::kValuesPerLane;
    constexpr int kGroupsPerPhase = (32 * Schedule::kValuesPerLane) / 16;
    const int subgroup_lane       = lane & (kSubgroupWidth - 1);
    const int group_base          = phase * kGroupsPerPhase + (lane / kSubgroupWidth) * 4;
    std::uint32_t word            = 0;
    if (subgroup_lane == 0) {
        word = *reinterpret_cast<const std::uint32_t*>(
            shared.raw_scales + local_row * Geometry::kGroupsPerRow + group_base);
    }
    return __shfl_sync(0xffffffffU, word, 0, kSubgroupWidth);
}

template <class Geometry, class Schedule>
__device__ __forceinline__ void load_nvfp4_coefficients(
    const std::uint8_t* __restrict__ scales,
    const Nvfp4GemvSharedStorage<Geometry, Schedule>& shared, int parent_row, int local_row,
    int phase, int lane, float inverse_weight_divisor,
    float (&coefficients)[Schedule::kValuesPerLane < 16 ? 1 : Schedule::kValuesPerLane / 16]) {
    constexpr int kGroupsPerLane =
        Schedule::kValuesPerLane < 16 ? 1 : Schedule::kValuesPerLane / 16;
    constexpr int kLanesPerGroup =
        Schedule::kValuesPerLane < 16 ? 16 / Schedule::kValuesPerLane : 1;
    const int value_begin = phase * 32 * Schedule::kValuesPerLane + lane * Schedule::kValuesPerLane;
    const int group_begin = value_begin / 16;

    if constexpr (Schedule::kScaleAccess == Nvfp4ScaleAccess::StagedRaw) {
        constexpr int kSubgroupWidth = 64 / Schedule::kValuesPerLane;
        const int subgroup_lane      = lane & (kSubgroupWidth - 1);
        const int first_byte         = (subgroup_lane / kLanesPerGroup) * kGroupsPerLane;
        const std::uint32_t word =
            load_staged_scale_word<Geometry, Schedule>(shared, local_row, phase, lane);
#pragma unroll
        for (int group = 0; group < kGroupsPerLane; ++group) {
            const auto scale    = static_cast<std::uint8_t>(word >> (8 * (first_byte + group)));
            coefficients[group] = decode_nvfp4_e4m3(scale) * inverse_weight_divisor;
        }
    } else {
#pragma unroll
        for (int group = 0; group < kGroupsPerLane; ++group) {
            const std::uint8_t scale =
                scales[nvfp4_scale_offset<Geometry>(parent_row, group_begin + group)];
            coefficients[group] = decode_nvfp4_e4m3(scale) * inverse_weight_divisor;
        }
    }
}

template <class Geometry, class Schedule>
__device__ __forceinline__ void
compute_nvfp4_rows(const __nv_bfloat16* __restrict__ x, const std::uint8_t* __restrict__ codes,
                   const std::uint8_t* __restrict__ scales,
                   const Nvfp4GemvSharedStorage<Geometry, Schedule>& shared,
                   float inverse_weight_divisor, const int (&parent_rows)[Schedule::kRowsPerWarp],
                   int flat_row0, int lane,
                   float (&accumulators)[Schedule::kRowsPerWarp][Schedule::kAccumulatorChains]) {
    constexpr int kValuesPerPhase = 32 * Schedule::kValuesPerLane;
    constexpr int kPhases         = Geometry::kInputRows / kValuesPerPhase;
    constexpr int kGroupsPerLane =
        Schedule::kValuesPerLane < 16 ? 1 : Schedule::kValuesPerLane / 16;
    static_assert((Geometry::kInputRows % kValuesPerPhase) == 0);
    const auto* activation_pairs = reinterpret_cast<const std::uint32_t*>(x);

#pragma unroll
    for (int phase = 0; phase < kPhases; ++phase) {
        float coefficients[Schedule::kRowsPerWarp][kGroupsPerLane];
        Nvfp4CodePack<Schedule::kValuesPerLane> row_codes[Schedule::kRowsPerWarp];
#pragma unroll
        for (int local_row = 0; local_row < Schedule::kRowsPerWarp; ++local_row) {
            load_nvfp4_coefficients<Geometry, Schedule>(
                scales, shared, parent_rows[local_row], flat_row0 + local_row, phase, lane,
                inverse_weight_divisor, coefficients[local_row]);
            const std::int64_t code_offset =
                static_cast<std::int64_t>(parent_rows[local_row]) * Geometry::kCodeBytesPerRow +
                phase * (kValuesPerPhase / 2) + lane * (Schedule::kValuesPerLane / 2);
            row_codes[local_row] = load_nvfp4_codes<Schedule::kCodeCache, Schedule::kValuesPerLane>(
                codes + code_offset);
        }

#pragma unroll
        for (int pair = 0; pair < Schedule::kPairsPerLane; ++pair) {
            const int activation_index =
                phase * (kValuesPerPhase / 2) + lane * Schedule::kPairsPerLane + pair;
            const float2 activation = bf16x2_bits_to_float2(activation_pairs[activation_index]);
            const int group         = ((lane * Schedule::kValuesPerLane & 15) + pair * 2) / 16;
#pragma unroll
            for (int local_row = 0; local_row < Schedule::kRowsPerWarp; ++local_row) {
                const std::uint32_t word  = row_codes[local_row].words[pair / 4];
                const std::uint8_t packed = static_cast<std::uint8_t>(word >> (8 * (pair & 3)));
                const float2 code         = decode_nvfp4_e2m1x2(packed);
                const float coefficient   = coefficients[local_row][group];
                constexpr int kChainMask  = Schedule::kAccumulatorChains - 1;
                accumulators[local_row][(2 * pair) & kChainMask] =
                    fmaf(code.x * coefficient, activation.x,
                         accumulators[local_row][(2 * pair) & kChainMask]);
                accumulators[local_row][(2 * pair + 1) & kChainMask] =
                    fmaf(code.y * coefficient, activation.y,
                         accumulators[local_row][(2 * pair + 1) & kChainMask]);
            }
        }
    }
}

template <class Geometry, class Schedule, class Epilogue, class Output>
__global__ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_gemv_kernel(
    const __nv_bfloat16* __restrict__ x, const std::uint8_t* __restrict__ codes,
    const std::uint8_t* __restrict__ scales, float inverse_weight_divisor, Epilogue epilogue,
    Output output) {
    static_assert((Geometry::kOutputRows % 128) == 0);
    static_assert((Schedule::kRowsPerCta % 4) == 0);
    static_assert((128 % Schedule::kRowsPerCta) == 0);

    __shared__ Nvfp4GemvSharedStorage<Geometry, Schedule> shared;
    constexpr int kCtasPerM128 = 128 / Schedule::kRowsPerCta;
    const int m_tile           = static_cast<int>(blockIdx.x) / kCtasPerM128;
    const int cta_in_tile      = static_cast<int>(blockIdx.x) - m_tile * kCtasPerM128;
    const int rmod_base        = cta_in_tile * (Schedule::kRowsPerCta / 4);
    stage_nvfp4_scales<Geometry, Schedule>(scales, shared, m_tile, rmod_base);

    const int lane      = static_cast<int>(threadIdx.x) & 31;
    const int warp      = static_cast<int>(threadIdx.x) >> 5;
    const int flat_row0 = warp * Schedule::kRowsPerWarp;
    int parent_rows[Schedule::kRowsPerWarp];
#pragma unroll
    for (int local_row = 0; local_row < Schedule::kRowsPerWarp; ++local_row) {
        const int flat_row     = flat_row0 + local_row;
        const int rmod         = rmod_base + flat_row / 4;
        const int quartile     = flat_row & 3;
        parent_rows[local_row] = m_tile * 128 + rmod + quartile * 32;
    }

    float accumulators[Schedule::kRowsPerWarp][Schedule::kAccumulatorChains] = {};
    compute_nvfp4_rows<Geometry, Schedule>(x, codes, scales, shared, inverse_weight_divisor,
                                           parent_rows, flat_row0, lane, accumulators);

#pragma unroll
    for (int local_row = 0; local_row < Schedule::kRowsPerWarp; ++local_row) {
        float total = 0.0F;
#pragma unroll
        for (int chain = 0; chain < Schedule::kAccumulatorChains; ++chain) {
            total += accumulators[local_row][chain];
        }
        total = warp_reduce_sum(total);
        if (lane == 0) {
            const int parent_row = parent_rows[local_row];
            output.store(parent_row, 0, epilogue.apply(parent_row, 0, total));
        }
    }
}

} // namespace ninfer::ops::detail
#include "ops/linear_swiglu/q4/q4_linear_swiglu_kernels.h"

#include "ops/common/math.cuh"
#include "ops/common/memory.cuh"
#include "ops/common/warp.cuh"
#include "core/device.h" // CUDA_CHECK
#include "ops/linear/q4/q4_small_t_mma.cuh"

#include <cuda_bf16.h>
#include <cuda_fp16.h>

#include <cstdint>
#include <stdexcept>
#include <array>
#include <utility>

namespace ninfer::ops::detail {
namespace {

constexpr int kN                 = 34816;
constexpr int kK                 = 5120;
constexpr int kIntermediate      = kN / 2;
constexpr int kGroupK            = 64;
constexpr int kGroups            = kK / kGroupK;
constexpr int kBytesPerGroup     = 32;
constexpr int kVecBytes          = 16;
constexpr int kGroupsPerWarpTile = 16;
constexpr int kVecsPerWarpTile   = kGroupsPerWarpTile * kBytesPerGroup / kVecBytes;
constexpr int kWarpsPerBlock     = 4;
constexpr int kBlockThreads      = kWarpsPerBlock * 32;
constexpr int kPairsPerBlock     = kWarpsPerBlock;
constexpr int kXVecs             = kK / 8; // x as uint4 (8 bf16 each)
constexpr int kTiles             = kGroups / kGroupsPerWarpTile;
static_assert(kIntermediate % kPairsPerBlock == 0);
static_assert(kBytesPerGroup == 2 * kVecBytes);
static_assert(kGroups % kGroupsPerWarpTile == 0);
static_assert(kVecsPerWarpTile == 32);

struct Q4SwiGluSmallTGeometry {
    static constexpr int kInputRows    = kK;
    static constexpr int kGroupsPerRow = kK / kGroupK;
};

struct Q4SwiGluSmallTRows {
    static constexpr int kOutputRowsPerCta = 8;

    __device__ __forceinline__ int weight_row(int output_row0, int local_row) const {
        return output_row0 + (local_row & 7) + (local_row >= 8 ? kIntermediate : 0);
    }
};

struct Q4SwiGluSmallTEpilogue {
    __nv_bfloat16* out;

    template <int ActiveCols>
    __device__ __forceinline__ void store(int row, int col0, float4 projected) const {
        if (col0 < ActiveCols) {
            out[static_cast<std::int64_t>(col0) * kIntermediate + row] =
                __float2bfloat16_rn(silu(projected.x) * projected.z);
        }
        if (col0 + 1 < ActiveCols) {
            out[static_cast<std::int64_t>(col0 + 1) * kIntermediate + row] =
                __float2bfloat16_rn(silu(projected.y) * projected.w);
        }
    }
};

using SmallTLauncher = void (*)(const Tensor&, const Weight&, Tensor&, cudaStream_t);

template <int ActiveCols>
void launch_small_t_active(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream) {
    constexpr int TileCols =
        ActiveCols <= 8 ? 8 : (ActiveCols <= 16 ? 16 : (ActiveCols <= 24 ? 24 : 32));
    constexpr int kBlocks = kIntermediate / Q4SwiGluSmallTRows::kOutputRowsPerCta;
    const Q4SwiGluSmallTEpilogue epilogue{static_cast<__nv_bfloat16*>(out.data)};
    q4_small_t_mma_kernel<Q4SwiGluSmallTGeometry, TileCols, ActiveCols, Q4SwiGluSmallTEpilogue,
                          Q4SwiGluSmallTRows>
        <<<kBlocks, Q4DraftSmallTSchedule::kThreads, 0, stream>>>(
            static_cast<const __nv_bfloat16*>(x.data), static_cast<const std::uint8_t*>(w.qdata),
            static_cast<const std::uint8_t*>(w.scales), static_cast<__nv_bfloat16*>(out.data),
            epilogue, Q4SwiGluSmallTRows{});
    CUDA_CHECK(cudaGetLastError());
}

template <std::size_t... Offsets>
constexpr auto make_small_t_launchers(std::index_sequence<Offsets...>) {
    return std::array<SmallTLauncher, sizeof...(Offsets)>{
        &launch_small_t_active<2 + static_cast<int>(Offsets)>...};
}

constexpr auto kSmallTLaunchers = make_small_t_launchers(std::make_index_sequence<31>{});

__device__ __forceinline__ void q4_issue_pair_tile(uint4 (*__restrict__ s_code)[kVecsPerWarpTile],
                                                   uint4 (*__restrict__ s_scale)[2],
                                                   const std::uint8_t* __restrict__ gate_code_row,
                                                   const std::uint8_t* __restrict__ gate_scale_row,
                                                   const std::uint8_t* __restrict__ up_code_row,
                                                   const std::uint8_t* __restrict__ up_scale_row,
                                                   int tile, int lane) {
    const int g0 = tile * kGroupsPerWarpTile;
    pipe_copy<16>(&s_code[0][lane],
                  reinterpret_cast<const uint4*>(gate_code_row + g0 * kBytesPerGroup) + lane);
    pipe_copy<16>(&s_code[1][lane],
                  reinterpret_cast<const uint4*>(up_code_row + g0 * kBytesPerGroup) + lane);
    if (lane < 2) {
        pipe_copy<16>(&s_scale[0][lane],
                      reinterpret_cast<const uint4*>(gate_scale_row + g0 * 2) + lane);
        pipe_copy<16>(&s_scale[1][lane],
                      reinterpret_cast<const uint4*>(up_scale_row + g0 * 2) + lane);
    }
    pipe_commit();
}

__global__ void q4_linear_swiglu_gemv_pair_kernel(const __nv_bfloat16* __restrict__ x,
                                                  const std::uint8_t* __restrict__ codes,
                                                  const std::uint8_t* __restrict__ scales,
                                                  __nv_bfloat16* __restrict__ out) {
    constexpr int kStages   = 3;
    constexpr int kPrefetch = kStages - 1;
    __shared__ __align__(16) __nv_bfloat16 x_sh[kK];
    __shared__ uint4 code_tile[kWarpsPerBlock][kStages][2][kVecsPerWarpTile];
    __shared__ uint4 scale_tile[kWarpsPerBlock][kStages][2][2];

    auto* x_sh_v    = reinterpret_cast<uint4*>(x_sh);
    const auto* x_g = reinterpret_cast<const uint4*>(x);
    for (int i = static_cast<int>(threadIdx.x); i < kXVecs; i += static_cast<int>(blockDim.x)) {
        x_sh_v[i] = x_g[i];
    }
    __syncthreads();

    const int lane    = static_cast<int>(threadIdx.x) & 31;
    const int warp    = static_cast<int>(threadIdx.x) >> 5;
    const int out_row = static_cast<int>(blockIdx.x) * kPairsPerBlock + warp;

    const std::uint8_t* gate_code_row =
        codes + static_cast<std::int64_t>(out_row) * kGroups * kBytesPerGroup;
    const std::uint8_t* gate_scale_row = scales + static_cast<std::int64_t>(out_row) * kGroups * 2;
    const std::uint8_t* up_code_row =
        codes + static_cast<std::int64_t>(out_row + kIntermediate) * kGroups * kBytesPerGroup;
    const std::uint8_t* up_scale_row =
        scales + static_cast<std::int64_t>(out_row + kIntermediate) * kGroups * 2;
    const auto* x2 = reinterpret_cast<const __nv_bfloat162*>(x_sh);

    float gate_acc = 0.0f;
    float up_acc   = 0.0f;
#pragma unroll
    for (int p = 0; p < kPrefetch; ++p) {
        if (p < kTiles) {
            q4_issue_pair_tile(code_tile[warp][p], scale_tile[warp][p], gate_code_row,
                               gate_scale_row, up_code_row, up_scale_row, p, lane);
        } else {
            pipe_commit();
        }
    }

#pragma unroll 1
    for (int tile = 0; tile < kTiles; ++tile) {
        const int fetch = tile + kPrefetch;
        if (fetch < kTiles) {
            const int buf = fetch % kStages;
            q4_issue_pair_tile(code_tile[warp][buf], scale_tile[warp][buf], gate_code_row,
                               gate_scale_row, up_code_row, up_scale_row, fetch, lane);
        } else {
            pipe_commit();
        }
        pipe_wait<kPrefetch>();
        __syncwarp();

        const int buf           = tile % kStages;
        const auto* gate_codes  = reinterpret_cast<const std::uint8_t*>(code_tile[warp][buf][0]);
        const auto* up_codes    = reinterpret_cast<const std::uint8_t*>(code_tile[warp][buf][1]);
        const auto* gate_scales = reinterpret_cast<const std::uint16_t*>(scale_tile[warp][buf][0]);
        const auto* up_scales   = reinterpret_cast<const std::uint16_t*>(scale_tile[warp][buf][1]);
#pragma unroll
        for (int tile_group = 0; tile_group < kGroupsPerWarpTile; ++tile_group) {
            const float gate_scale =
                __half2float(__ushort_as_half(static_cast<std::uint16_t>(gate_scales[tile_group])));
            const float up_scale =
                __half2float(__ushort_as_half(static_cast<std::uint16_t>(up_scales[tile_group])));

            const int gate_packed =
                static_cast<int>(gate_codes[tile_group * kBytesPerGroup + lane]);
            const int gate_q0   = sign_extend<4>(gate_packed & 0x0f);
            const int gate_q1   = sign_extend<4>(gate_packed >> 4);
            const int up_packed = static_cast<int>(up_codes[tile_group * kBytesPerGroup + lane]);
            const int up_q0     = sign_extend<4>(up_packed & 0x0f);
            const int up_q1     = sign_extend<4>(up_packed >> 4);
            const int k0        = (tile * kGroupsPerWarpTile + tile_group) * kGroupK + lane * 2;
            const float2 xv     = __bfloat1622float2(x2[k0 >> 1]);
            gate_acc            = fmaf(static_cast<float>(gate_q0) * gate_scale, xv.x, gate_acc);
            gate_acc            = fmaf(static_cast<float>(gate_q1) * gate_scale, xv.y, gate_acc);
            up_acc              = fmaf(static_cast<float>(up_q0) * up_scale, xv.x, up_acc);
            up_acc              = fmaf(static_cast<float>(up_q1) * up_scale, xv.y, up_acc);
        }
        __syncwarp();
    }

    gate_acc = warp_reduce_sum(gate_acc);
    up_acc   = warp_reduce_sum(up_acc);
    if (lane == 0) { out[out_row] = __float2bfloat16(silu(gate_acc) * up_acc); }
}

} // namespace

void q4_linear_swiglu_gemv_pair_launch(const Tensor& x, const Weight& w, Tensor& out,
                                       cudaStream_t stream) {
    if (w.n != kN || w.k != kK || w.padded_shape[1] != kK) {
        throw std::invalid_argument("q4 linear_swiglu GEMV requires weight [34816,5120]");
    }
    const int grid = kIntermediate / kPairsPerBlock;
    q4_linear_swiglu_gemv_pair_kernel<<<grid, kBlockThreads, 0, stream>>>(
        static_cast<const __nv_bfloat16*>(x.data), static_cast<const std::uint8_t*>(w.qdata),
        static_cast<const std::uint8_t*>(w.scales), static_cast<__nv_bfloat16*>(out.data));
    CUDA_CHECK(cudaGetLastError());
}

void q4_linear_swiglu_small_t_exact_launch(const Tensor& x, const Weight& w, Tensor& out,
                                           cudaStream_t stream) {
    if (x.ne[1] < 2 || x.ne[1] > 32) {
        throw std::invalid_argument("Q4 LinearSwiGLU exact small-T requires T=2..32");
    }
    kSmallTLaunchers[static_cast<std::size_t>(x.ne[1] - 2)](x, w, out, stream);
}

} // namespace ninfer::ops::detail
#pragma once

#include "ninfer/ops/linear.h"
#include "ops/linear/bf16/bf16_launch.h"

#include <cuda_runtime.h>

namespace ninfer::ops::detail {

Bf16Launch select_bf16_a16_launch(std::int32_t n, std::int32_t k, std::int32_t t);
Bf16Launch select_bf16_launch(std::int32_t n, std::int32_t k, std::int32_t t, LinearPolicy policy);

void bf16_dispatch(const Tensor& x, const Weight& weight, Tensor& out, LinearPolicy policy,
                   cudaStream_t stream);

} // namespace ninfer::ops::detail
#pragma once

// ninfer::ops::detail - private launch prototype for rope. Included by the wrapper
// and defined by the CUDA launcher.

#include "core/tensor.h"

#include <cuda_runtime.h>

namespace ninfer::ops::detail {

void rope_launch(const Tensor& positions, int rotary_dim, float theta, Tensor& q, Tensor& k,
                 cudaStream_t stream);

void rope_single_launch(const Tensor& positions, int rotary_dim, float theta, Tensor& x,
                        cudaStream_t stream);

} // namespace ninfer::ops::detail
#pragma once

#include "core/tensor.h"

#include <cuda_runtime.h>

namespace ninfer::ops {

/**
 * Op: scatter
 *
 * Math / indexing:
 *   destination[:, indices[i]] = source[:, i] for 0 <= i < V.
 *
 * Logical shapes:
 *   BF16 source [D,V], I32 indices [V], BF16 destination [D,T], all contiguous.
 *
 * Numeric:
 *   Exact BF16 element copies. Callers provide indices in [0,T) with no duplicates when
 *   deterministic overwrite order is required.
 *
 * Effects:
 *   Overwrites only the selected destination columns; all other columns retain their old values.
 *   Inputs and destination must not alias.
 *
 * Workspace:
 *   None. The Op has no persistent state side effect beyond the selected destination writes.
 *
 * Supported routes:
 *   D divisible by eight with 16-byte-aligned source/destination uses BF16x8 copies; other aligned
 *   even D uses BF16x2; remaining contiguous dimensions use scalar copies.
 */
void scatter(const Tensor& src, const Tensor& indices, Tensor& dst, cudaStream_t stream);

/**
 * Scatter compact request-major BF16 blocks into lane-owned fixed storage.
 *
 * source is contiguous [D,W,B], lanes and valid_columns are contiguous I32 [B], and destination
 * is a possibly row-sliced [D,W,C] view. For each b and j<valid_columns[b]:
 *
 *   destination[:,j,lanes[b]] = source[:,j,b]
 *
 * Columns j>=valid_columns[b] and every unselected lane remain unchanged. The caller guarantees
 * 0<=valid_columns[b]<=W and 0<=lanes[b]<C. D is divisible by eight and all column starts are
 * 16-byte aligned. Inputs and destination do not alias.
 */
void scatter_bf16_batch(const Tensor& source, const Tensor& lanes, const Tensor& valid_columns,
                        Tensor& destination, cudaStream_t stream);

/**
 * Op: extract_bf16_columns
 *
 * Math / indexing:
 *   For source [D,T] and destination [D',T]:
 *   destination[d,t] = source[source_column + d,t], 0<=d<D', 0<=t<T.
 *
 * Logical shapes:
 *   Contiguous BF16 source [D,T] and destination [D',T]. Despite the historical parameter name,
 *   source_column is an offset in the fastest (ne[0]) dimension; 0<=source_column and
 *   source_column+D'<=D.
 *
 * Numeric:
 *   Exact BF16 element copies.
 *
 * Effects:
 *   Writes the full destination; source and destination must not alias.
 *
 * Workspace:
 *   None. The Op has no persistent state side effect.
 */
void extract_bf16_columns(const Tensor& source, std::int32_t source_column, Tensor& destination,
                          cudaStream_t stream);

} // namespace ninfer::ops
#pragma once

#include <chrono>
#include <string>
#include <string_view>

namespace ninfer::serve {

enum class ConsoleLogLevel {
    Info,
    Warning,
    Error,
};

[[nodiscard]] std::string format_console_log_prefix(std::chrono::system_clock::time_point timestamp,
                                                    ConsoleLogLevel level);

[[nodiscard]] std::string current_console_log_prefix(ConsoleLogLevel level);

void write_console_log(ConsoleLogLevel level, std::string_view message);

} // namespace ninfer::serve
#pragma once

#include "ninfer/ops/linear.h"
#include "ops/linear/q6/q6_launch.h"

#include <cstdint>

namespace ninfer::ops::detail {

Q6Launch select_q6_a16_launch(std::int32_t n, std::int32_t k, std::int32_t t);
Q6Launch select_q6_launch(std::int32_t n, std::int32_t k, std::int32_t t, LinearPolicy policy);

void q6_dispatch(const Tensor& x, const Weight& w, Tensor& out, LinearPolicy policy,
                 cudaStream_t stream);

} // namespace ninfer::ops::detail
#include "targets/registry.h"

#include "artifact/binder.h"
#include "artifact/materializer.h"
#include "artifact/reader.h"
#include "core/device.h"
#include "runtime/engine/kv_capacity.h"
#include "runtime/engine/context_cost.h"

#include <chrono>
#include <cstdint>
#include <stdexcept>
#include <string>
#include <string_view>
#include <utility>

namespace ninfer::targets {
namespace {

using Clock = std::chrono::steady_clock;

void validate_options(const EngineOptions& options) {
    if (options.artifact_path.empty()) {
        throw std::invalid_argument("Engine artifact_path must not be empty");
    }
    if (options.artifact_path.extension() != ".ninfer") {
        throw std::invalid_argument("NInfer accepts only .ninfer artifacts");
    }
    if (options.max_context == 0) {
        throw std::invalid_argument("Engine max_context must be nonzero");
    }
    switch (options.kv_capacity.mode) {
    case KvCapacityMode::Explicit:
        if (options.kv_capacity.explicit_tokens == 0) {
            throw std::invalid_argument("Engine explicit kv_capacity must be nonzero");
        }
        if (options.kv_capacity.automatic_headroom_bytes != 0) {
            throw std::invalid_argument(
                "Engine explicit kv_capacity must not carry automatic headroom");
        }
        break;
    case KvCapacityMode::Automatic:
        if (options.kv_capacity.explicit_tokens != 0) {
            throw std::invalid_argument(
                "Engine automatic kv_capacity must not carry explicit tokens");
        }
        break;
    default:
        throw std::invalid_argument("Engine kv_capacity mode is invalid");
    }
    if (options.max_concurrency == 0 || options.max_concurrency > kMaximumConcurrency) {
        throw std::invalid_argument("Engine max_concurrency must be in [1,8]");
    }
    if (options.max_pending_requests == 0 || options.pending_timeout_ms == 0) {
        throw std::invalid_argument("Engine pending request capacity and timeout must be nonzero");
    }
    if (options.enable_vision && options.media_live_bytes == 0) {
        throw std::invalid_argument(
            "Engine media_live_bytes must be nonzero when Vision is enabled");
    }
    if (options.media_preprocess_threads > 64) {
        throw std::invalid_argument("Engine media_preprocess_threads must be in [0,64]");
    }
}

artifact::LoadProgress artifact_progress(const LoadProgress& progress) {
    return artifact::LoadProgress{.callback = progress.callback};
}

std::size_t runtime_bytes_after_planned_weights(std::uint64_t weight_bytes) {
    std::size_t free_bytes  = 0;
    std::size_t total_bytes = 0;
    CUDA_CHECK(cudaMemGetInfo(&free_bytes, &total_bytes));
    if (weight_bytes > free_bytes) {
        throw std::invalid_argument("model weights require " + std::to_string(weight_bytes) +
                                    " bytes of device memory, but only " +
                                    std::to_string(free_bytes) +
                                    " bytes are free before loading weights");
    }
    return free_bytes - static_cast<std::size_t>(weight_bytes);
}

std::size_t current_free_device_bytes() {
    std::size_t free_bytes  = 0;
    std::size_t total_bytes = 0;
    CUDA_CHECK(cudaMemGetInfo(&free_bytes, &total_bytes));
    return free_bytes;
}

template <class Target, class Loaded, class Instance>
ConstructedTarget construct_registered(const EngineOptions& options, DeviceContext& device,
                                       artifact::Reader& reader, Clock::time_point load_start,
                                       std::string_view target_key) {
    const auto& identity                          = reader.identity();
    const auto weights_profile                    = Target::resolve_weights(identity);
    const ModelSamplingDefaults sampling_defaults = Target::sampling_defaults(identity.model_id);
    const runtime::ContextCostIdentity context_cost_identity{
        .hardware_class = runtime::context_cost_hardware_class(
            device.props.name, device.props.major, device.props.minor),
        .model_id   = identity.model_id,
        .weights_id = identity.weights_id,
    };
    runtime::ResolvedContextMachineCost context_cost = runtime::resolve_context_machine_cost(
        context_cost_identity, options.context_cost.preset_path);

    artifact::Binder binder(reader);
    auto load_plan        = Target::plan_load(binder, options, weights_profile);
    auto sequence_planner = Target::make_sequence_planner(device, options, weights_profile);
    const runtime::SequenceCapacityCurve curve = sequence_planner.capacity_curve();
    const std::size_t preflight_runtime_bytes =
        runtime_bytes_after_planned_weights(load_plan.materialization().device_capacity_bytes);
    (void)runtime::resolve_kv_capacity(options.kv_capacity, curve, preflight_runtime_bytes);

    auto progress     = artifact_progress(options.load_progress);
    auto materialized = artifact::materialize(reader, load_plan.materialization(), device,
                                              progress.callback ? &progress : nullptr);
    const artifact::MaterializationStats stats = materialized.stats();

    auto model = Target::construct_loaded_model(std::move(load_plan), std::move(materialized));
    device.synchronize();
    runtime::KvCapacityResolution capacity_resolution =
        runtime::resolve_kv_capacity(options.kv_capacity, curve, current_free_device_bytes());
    auto sequence_plan = std::move(sequence_planner).finalize(capacity_resolution.main_page_groups);
    if (sequence_plan.device_reservation_bytes() != capacity_resolution.runtime_reservation_bytes ||
        sequence_plan.kv_capacity() != capacity_resolution.resolved_tokens) {
        throw std::logic_error("resolved KV capacity does not match the finalized target plan");
    }
    auto loaded   = std::make_unique<Loaded>(std::move(model), options);
    auto instance = std::make_unique<Instance>(std::move(loaded), capacity_resolution,
                                               std::move(sequence_plan), device);
    device.synchronize();
    instance->kv_capacity_resolution.available_after_startup_bytes = current_free_device_bytes();

    LoadSummary summary;
    summary.target               = std::string(target_key);
    summary.model_id             = identity.model_id;
    summary.weights_id           = identity.weights_id;
    summary.load_seconds         = std::chrono::duration<double>(Clock::now() - load_start).count();
    summary.upload_seconds       = stats.upload_seconds;
    summary.artifact_bytes_read  = stats.file_bytes;
    summary.host_to_device_bytes = stats.h2d_bytes;
    summary.peak_staging_bytes   = stats.peak_staging_bytes;
    summary.tensor_count         = stats.tensor_count;
    summary.resource_count       = stats.resource_count;
    summary.context_cost         = context_cost.summary;
    return ConstructedTarget{.active            = ActiveTarget(std::move(instance)),
                             .load              = std::move(summary),
                             .sampling_defaults = sampling_defaults,
                             .context_cost      = std::move(context_cost.model)};
}

} // namespace

LoadedQwen3_6_27B::LoadedQwen3_6_27B(std::unique_ptr<Qwen3_6_27B::LoadedModel> stable_model,
                                     const EngineOptions& options)
    : model(std::move(stable_model)), frontend(Qwen3_6_27B::make_frontend(*model, options)) {}

LoadedQwen3_6_27B::~LoadedQwen3_6_27B() = default;

Qwen3_6_27BInstance::Qwen3_6_27BInstance(std::unique_ptr<LoadedQwen3_6_27B> stable_loaded,
                                         runtime::KvCapacityResolution resolution,
                                         Qwen3_6_27B::SequencePlan sequence_plan,
                                         DeviceContext& device)
    : loaded(std::move(stable_loaded)), kv_capacity_resolution(resolution),
      capacity(sequence_plan.capacity()),
      program(Qwen3_6_27B::create_program(*loaded->model, std::move(sequence_plan), device)) {}

Qwen3_6_27BInstance::~Qwen3_6_27BInstance() = default;

LoadedQwen3_6_35BA3B::LoadedQwen3_6_35BA3B(
    std::unique_ptr<Qwen3_6_35BA3B::LoadedModel> stable_model, const EngineOptions& options)
    : model(std::move(stable_model)), frontend(Qwen3_6_35BA3B::make_frontend(*model, options)) {}

LoadedQwen3_6_35BA3B::~LoadedQwen3_6_35BA3B() = default;

Qwen3_6_35BA3BInstance::Qwen3_6_35BA3BInstance(std::unique_ptr<LoadedQwen3_6_35BA3B> stable_loaded,
                                               runtime::KvCapacityResolution resolution,
                                               Qwen3_6_35BA3B::SequencePlan sequence_plan,
                                               DeviceContext& device)
    : loaded(std::move(stable_loaded)), kv_capacity_resolution(resolution),
      capacity(sequence_plan.capacity()),
      program(Qwen3_6_35BA3B::create_program(*loaded->model, std::move(sequence_plan), device)) {}

Qwen3_6_35BA3BInstance::~Qwen3_6_35BA3BInstance() = default;

ConstructedTarget construct_target(const EngineOptions& options, DeviceContext& device) {
    validate_options(options);
    const auto load_start = Clock::now();

    artifact::Reader reader(options.artifact_path);
    const auto& identity = reader.identity();
    if (identity.model_id == Qwen3_6_27B::model_id) {
        return construct_registered<Qwen3_6_27B, LoadedQwen3_6_27B, Qwen3_6_27BInstance>(
            options, device, reader, load_start, Qwen3_6_27B::target_key);
    }
    if (identity.model_id == Qwen3_6_27B::qwen3_8_model_id) {
        return construct_registered<Qwen3_6_27B, LoadedQwen3_6_27B, Qwen3_6_27BInstance>(
            options, device, reader, load_start, Qwen3_6_27B::qwen3_8_target_key);
    }
    if (identity.model_id == Qwen3_6_35BA3B::model_id) {
        return construct_registered<Qwen3_6_35BA3B, LoadedQwen3_6_35BA3B, Qwen3_6_35BA3BInstance>(
            options, device, reader, load_start, Qwen3_6_35BA3B::target_key);
    }
    throw std::runtime_error("artifact identity '" + identity.model_id + "/" + identity.weights_id +
                             "' has no registered target for this device");
}

} // namespace ninfer::targets
#pragma once

#include "core/arena.h"
#include "core/tensor.h"
#include "ninfer/ops/linear.h"

#include <cuda_runtime.h>

#include <cstddef>
#include <cstdint>

namespace ninfer::ops::detail {

[[nodiscard]] std::size_t fp8_gdn_snapshot_workspace_capacity_bytes(LinearPolicy policy,
                                                                    std::int32_t batch_size,
                                                                    std::int32_t min_width,
                                                                    std::int32_t max_width);

[[nodiscard]] std::size_t fp8_gdn_record_workspace_capacity_bytes(LinearPolicy policy,
                                                                  std::int32_t batch_size,
                                                                  std::int32_t min_width,
                                                                  std::int32_t max_width);

void fp8_gdn_snapshot_fused_launch(const Tensor& x, const Weight& weight, const Tensor& conv_weight,
                                   Tensor& conv_states, const Tensor& valid_columns,
                                   const Tensor& initial_slot, const Tensor& snapshot_base_slot,
                                   Tensor& query, Tensor& key, Tensor& value, Tensor& z,
                                   cudaStream_t stream);

void fp8_gdn_record_fused_launch(const Tensor& x, const Weight& weight, const Tensor& conv_weight,
                                 const Tensor& conv_states, const Tensor& valid_columns,
                                 const Tensor& initial_slot, Tensor& conv_record, Tensor& query,
                                 Tensor& key, Tensor& value, Tensor& z, cudaStream_t stream);

void fp8_gdn_snapshot_dispatch(const Tensor& x, const Weight& weight, const Tensor& conv_weight,
                               Tensor& conv_states, const Tensor& valid_columns,
                               const Tensor& initial_slot, const Tensor& snapshot_base_slot,
                               Tensor& query, Tensor& key, Tensor& value, Tensor& z,
                               LinearPolicy policy, WorkspaceArena& workspace, cudaStream_t stream);

void fp8_gdn_record_dispatch(const Tensor& x, const Weight& weight, const Tensor& conv_weight,
                             const Tensor& conv_states, const Tensor& valid_columns,
                             const Tensor& initial_slot, Tensor& conv_record, Tensor& query,
                             Tensor& key, Tensor& value, Tensor& z, LinearPolicy policy,
                             WorkspaceArena& workspace, cudaStream_t stream);

} // namespace ninfer::ops::detail
#include "ops/linear/w8/w8_rowsplit_gemm_mma.cuh"

#include "core/device.h"
#include "ops/common/math.h"
#include "ops/common/token_slices.h"
#include "ops/linear/w8/w8_launch.h"

#include <cstdint>
#include <stdexcept>

namespace ninfer::ops::detail {
namespace {

template <class Schedule, bool Full>
void launch_slice(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream) {
    const std::int32_t rows     = w.n;
    const std::int32_t k        = w.k;
    const std::int32_t cols     = x.ne[1];
    const std::int32_t padded_k = w.padded_shape[1];
    const dim3 grid(static_cast<unsigned>(div_up(rows, Schedule::BM)),
                    static_cast<unsigned>(div_up(cols, Schedule::BN)), 1u);
    const W8ContiguousOutput output{static_cast<__nv_bfloat16*>(out.data), rows};
    w8_rowsplit_gemm_mma_kernel<Schedule, Full><<<grid, Schedule::THREADS, 0, stream>>>(
        static_cast<const __nv_bfloat16*>(x.data), static_cast<const std::uint8_t*>(w.qdata),
        static_cast<const std::uint8_t*>(w.scales), output, rows, k, cols, padded_k);
    CUDA_CHECK(cudaGetLastError());
}

template <class Schedule>
void launch_route(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream) {
    const bool full = (w.n % Schedule::BM) == 0 && (x.ne[1] % Schedule::BN) == 0 &&
                      w.k == w.padded_shape[1] && (w.k % Schedule::BK) == 0;
    for_each_token_slice(x.ne[1], Schedule::BN, [&](std::int32_t offset, std::int32_t count) {
        const Tensor x_slice = x.slice(1, offset, count);
        Tensor out_slice     = out.slice(1, offset, count);
        if (full) {
            launch_slice<Schedule, true>(x_slice, w, out_slice, stream);
        } else {
            launch_slice<Schedule, false>(x_slice, w, out_slice, stream);
        }
    });
}

template <std::int32_t TileCols>
void launch_exact_tail(W8Launch prefix_launch, const Tensor& x, const Weight& w, Tensor& out,
                       cudaStream_t stream) {
    const std::int32_t full_cols = (x.ne[1] / TileCols) * TileCols;
    if (full_cols <= 0) {
        throw std::invalid_argument("w8 exact-tail route requires a non-empty MMA prefix");
    }

    const Tensor x_prefix = x.slice(1, 0, full_cols);
    Tensor out_prefix     = out.slice(1, 0, full_cols);
    prefix_launch(x_prefix, w, out_prefix, stream);

    const std::int32_t tail = x.ne[1] - full_cols;
    if (tail < 1 || tail > 65) {
        throw std::invalid_argument("w8 exact-tail route requires tail=1..65");
    }
    const Tensor x_tail = x.slice(1, full_cols, tail);
    Tensor out_tail     = out.slice(1, full_cols, tail);
    if (tail == 1) {
        launch_w8_decode_r4(x_tail, w, out_tail, stream);
    } else if (tail <= 32) {
        launch_w8_exact_t_splitk(x_tail, w, out_tail, stream);
    } else {
        launch_w8_exact_t_composite(x_tail, w, out_tail, stream);
    }
}

using MmaR32C64          = W8RowSplitMmaGemmSchedule<32, 64, 32, 16, 3>;
using MmaR32C96          = W8RowSplitMmaGemmSchedule<32, 96, 32, 16, 2>;
using MmaR32C128         = W8RowSplitMmaGemmSchedule<32, 128, 32, 16, 2>;
using MmaR48C64          = W8RowSplitMmaGemmSchedule<48, 64, 48, 16, 3>;
using MmaR48C96          = W8RowSplitMmaGemmSchedule<48, 96, 48, 16, 2>;
using MmaR48C112         = W8RowSplitMmaGemmSchedule<48, 112, 48, 16, 2>;
using MmaR48C128         = W8RowSplitMmaGemmSchedule<48, 128, 48, 16, 2>;
using MmaR64C96          = W8RowSplitMmaGemmSchedule<64, 96, 64, 16, 2>;
using MmaR64C112         = W8RowSplitMmaGemmSchedule<64, 112, 64, 16, 2>;
using MmaR64C128         = W8RowSplitMmaGemmSchedule<64, 128, 64, 16, 2, 2>;
using MmaR96C96          = W8RowSplitMmaGemmSchedule<96, 96, 48, 16, 2>;
using MmaR128C64         = W8RowSplitMmaGemmSchedule<128, 64, 64, 16, 2>;
using MmaR128C80         = W8RowSplitMmaGemmSchedule<128, 80, 64, 16, 2>;
using MmaR64x16C48K128A1 = W8RowSplitMmaGemmSchedule<64, 48, 16, 24, 2, 2, 128, 1>;

} // namespace

#define NINFER_W8_MMA_LAUNCHER(Name, Schedule)                                                     \
    void Name(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream) {                \
        launch_route<Schedule>(x, w, out, stream);                                                 \
    }

NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r32_c64, MmaR32C64)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r32_c96, MmaR32C96)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r32_c128, MmaR32C128)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r48_c64, MmaR48C64)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r48_c96, MmaR48C96)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r48_c112, MmaR48C112)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r48_c128, MmaR48C128)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r64_c96, MmaR64C96)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r64_c112, MmaR64C112)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r64_c128, MmaR64C128)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r96_c96, MmaR96C96)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r128_c64, MmaR128C64)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r128_c80, MmaR128C80)
NINFER_W8_MMA_LAUNCHER(launch_w8_mma_r64x16_c48_k128_a1, MmaR64x16C48K128A1)

#undef NINFER_W8_MMA_LAUNCHER

#define NINFER_W8_EXACT_LAUNCHER(Name, Prefix, TileCols)                                           \
    void Name(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream) {                \
        launch_exact_tail<TileCols>(Prefix, x, w, out, stream);                                    \
    }

NINFER_W8_EXACT_LAUNCHER(launch_w8_exact_mma_r32_c96, launch_w8_mma_r32_c96, 96)
NINFER_W8_EXACT_LAUNCHER(launch_w8_exact_mma_r32_c128, launch_w8_mma_r32_c128, 128)
NINFER_W8_EXACT_LAUNCHER(launch_w8_exact_mma_r48_c96, launch_w8_mma_r48_c96, 96)
NINFER_W8_EXACT_LAUNCHER(launch_w8_exact_mma_r48_c128, launch_w8_mma_r48_c128, 128)
NINFER_W8_EXACT_LAUNCHER(launch_w8_exact_mma_r64_c96, launch_w8_mma_r64_c96, 96)
NINFER_W8_EXACT_LAUNCHER(launch_w8_exact_mma_r64_c128, launch_w8_mma_r64_c128, 128)
NINFER_W8_EXACT_LAUNCHER(launch_w8_exact_mma_r96_c96, launch_w8_mma_r96_c96, 96)
NINFER_W8_EXACT_LAUNCHER(launch_w8_exact_mma_r128_c80, launch_w8_mma_r128_c80, 80)

#undef NINFER_W8_EXACT_LAUNCHER

} // namespace ninfer::ops::detail
#pragma once

// ninfer::ops - RMSNorm kernels over contiguous BF16 rows.

#include "ops/common/math.cuh"
#include "ops/common/warp.cuh"

#include <cuda_bf16.h>

#include <cstdint>

namespace ninfer::ops {

enum class RmsEpilogue {
    Offset,
    Plain,
    Gated,
};

template <RmsEpilogue Epilogue>
__device__ __forceinline__ float rmsnorm_epilogue(float x, float inv, float weight, float z) {
    if constexpr (Epilogue == RmsEpilogue::Offset) { weight += 1.0f; }
    float value = x * inv * weight;
    if constexpr (Epilogue == RmsEpilogue::Gated) { value *= silu(z); }
    return value;
}

// Fast geometry for D in {64, 128, 192, 256}. One warp owns one row, keeps the input in
// registers, and uses only warp shuffles for the reduction. Block is a scheduling choice rather
// than part of the row geometry.
template <RmsEpilogue Epilogue, int Block>
__launch_bounds__(Block) __global__
    void rmsnorm_warp_bf16x2_kernel(const __nv_bfloat162* x, const __nv_bfloat162* weight,
                                    const __nv_bfloat162* z, __nv_bfloat162* out, std::int32_t d,
                                    std::int64_t rows, float eps) {
    static_assert(Block % kWarpSize == 0);
    constexpr int kWarpsPerBlock   = Block / kWarpSize;
    constexpr int kMaxPairsPerLane = 4;
    const int lane                 = static_cast<int>(threadIdx.x) & (kWarpSize - 1);
    const int warp                 = static_cast<int>(threadIdx.x) / kWarpSize;
    const std::int64_t row         = static_cast<std::int64_t>(blockIdx.x) * kWarpsPerBlock + warp;
    if (row >= rows) { return; }

    const int pairs             = d / 2;
    const std::int64_t row_base = row * static_cast<std::int64_t>(pairs);
    __nv_bfloat162 values[kMaxPairsPerLane];
    float sum = 0.0f;

#pragma unroll
    for (int k = 0; k < kMaxPairsPerLane; ++k) {
        const int pair = lane + k * kWarpSize;
        if (pair < pairs) {
            values[k]       = x[row_base + pair];
            const float2 xf = __bfloat1622float2(values[k]);
            sum += xf.x * xf.x + xf.y * xf.y;
        }
    }

    sum       = warp_reduce_sum(sum);
    float inv = lane == 0 ? rsqrtf(sum / static_cast<float>(d) + eps) : 0.0f;
    inv       = __shfl_sync(kFullWarpMask, inv, 0);

#pragma unroll
    for (int k = 0; k < kMaxPairsPerLane; ++k) {
        const int pair = lane + k * kWarpSize;
        if (pair < pairs) {
            const float2 xf = __bfloat1622float2(values[k]);
            const float2 wf = __bfloat1622float2(weight[pair]);
            float2 zf{0.0f, 0.0f};
            if constexpr (Epilogue == RmsEpilogue::Gated) {
                zf = __bfloat1622float2(z[row_base + pair]);
            }
            out[row_base + pair] =
                __floats2bfloat162_rn(rmsnorm_epilogue<Epilogue>(xf.x, inv, wf.x, zf.x),
                                      rmsnorm_epilogue<Epilogue>(xf.y, inv, wf.y, zf.y));
        }
    }
}

// Implements: include/ninfer/ops/rmsnorm.h
// Match: aligned contiguous BF16, plain epilogue, D=128, sm_120a.
// Algorithm assumptions: exactly two BF16x2 values per lane; one warp owns one logical row.
template <RmsEpilogue Epilogue, int Block>
__launch_bounds__(Block) __global__
    void rmsnorm_d128_bf16x2_kernel(const __nv_bfloat162* x, const __nv_bfloat162* weight,
                                    const __nv_bfloat162* z, __nv_bfloat162* out, std::int64_t rows,
                                    float eps) {
    static_assert(Block % kWarpSize == 0);
    constexpr int kWarpsPerBlock = Block / kWarpSize;
    constexpr int kPairsPerRow   = 64;
    const int lane               = static_cast<int>(threadIdx.x) & (kWarpSize - 1);
    const int warp               = static_cast<int>(threadIdx.x) / kWarpSize;
    const std::int64_t row       = static_cast<std::int64_t>(blockIdx.x) * kWarpsPerBlock + warp;
    if (row >= rows) { return; }

    const std::int64_t row_base = row * kPairsPerRow;
    const int pair0             = lane;
    const int pair1             = lane + kWarpSize;
    const __nv_bfloat162 value0 = x[row_base + pair0];
    const __nv_bfloat162 value1 = x[row_base + pair1];
    const float2 x0             = __bfloat1622float2(value0);
    const float2 x1             = __bfloat1622float2(value1);
    float sum                   = x0.x * x0.x + x0.y * x0.y + x1.x * x1.x + x1.y * x1.y;
    sum                         = warp_reduce_sum(sum);
    float inv                   = lane == 0 ? rsqrtf(sum * (1.0f / 128.0f) + eps) : 0.0f;
    inv                         = __shfl_sync(kFullWarpMask, inv, 0);

    const float2 w0 = __bfloat1622float2(weight[pair0]);
    const float2 w1 = __bfloat1622float2(weight[pair1]);
    float2 z0{0.0f, 0.0f};
    float2 z1{0.0f, 0.0f};
    if constexpr (Epilogue == RmsEpilogue::Gated) {
        z0 = __bfloat1622float2(z[row_base + pair0]);
        z1 = __bfloat1622float2(z[row_base + pair1]);
    }
    out[row_base + pair0] =
        __floats2bfloat162_rn(rmsnorm_epilogue<Epilogue>(x0.x, inv, w0.x, z0.x),
                              rmsnorm_epilogue<Epilogue>(x0.y, inv, w0.y, z0.y));
    out[row_base + pair1] =
        __floats2bfloat162_rn(rmsnorm_epilogue<Epilogue>(x1.x, inv, w1.x, z1.x),
                              rmsnorm_epilogue<Epilogue>(x1.y, inv, w1.y, z1.y));
}

// Fast geometry for wide rows. One CTA owns one row and keeps up to MaxPairsPerThread BF16x2
// values per lane. The launcher admits only widths evenly divisible by the CTA vector span.
template <RmsEpilogue Epilogue, int Block, int MaxPairsPerThread>
__launch_bounds__(Block) __global__
    void rmsnorm_cta_bf16x2_kernel(const __nv_bfloat162* x, const __nv_bfloat162* weight,
                                   const __nv_bfloat162* z, __nv_bfloat162* out, std::int32_t d,
                                   std::int64_t rows, float eps) {
    static_assert(Block % kWarpSize == 0);
    const std::int64_t row = static_cast<std::int64_t>(blockIdx.x);
    if (row >= rows) { return; }

    const int pairs             = d / 2;
    const int pairs_per_thread  = pairs / Block;
    const std::int64_t row_base = row * static_cast<std::int64_t>(pairs);
    __nv_bfloat162 values[MaxPairsPerThread];
    float sum = 0.0f;

#pragma unroll
    for (int k = 0; k < MaxPairsPerThread; ++k) {
        if (k < pairs_per_thread) {
            const int pair  = static_cast<int>(threadIdx.x) + k * Block;
            values[k]       = x[row_base + pair];
            const float2 xf = __bfloat1622float2(values[k]);
            sum += xf.x * xf.x + xf.y * xf.y;
        }
    }

    __shared__ float warp_sums[Block / kWarpSize];
    __shared__ float inv_shared;
    const float block_sum = block_reduce_sum<Block>(sum, warp_sums);
    if (threadIdx.x == 0) { inv_shared = rsqrtf(block_sum / static_cast<float>(d) + eps); }
    __syncthreads();
    const float inv = inv_shared;

#pragma unroll
    for (int k = 0; k < MaxPairsPerThread; ++k) {
        if (k < pairs_per_thread) {
            const int pair  = static_cast<int>(threadIdx.x) + k * Block;
            const float2 xf = __bfloat1622float2(values[k]);
            const float2 wf = __bfloat1622float2(weight[pair]);
            float2 zf{0.0f, 0.0f};
            if constexpr (Epilogue == RmsEpilogue::Gated) {
                zf = __bfloat1622float2(z[row_base + pair]);
            }
            out[row_base + pair] =
                __floats2bfloat162_rn(rmsnorm_epilogue<Epilogue>(xf.x, inv, wf.x, zf.x),
                                      rmsnorm_epilogue<Epilogue>(xf.y, inv, wf.y, zf.y));
        }
    }
}

// Implements: include/ninfer/ops/rmsnorm.h
// Match: aligned contiguous BF16, plain epilogue, D=2048, sm_120a.
// Algorithm assumptions: exactly two BF16x2 values per thread; one 512-thread CTA owns one row.
template <RmsEpilogue Epilogue>
__launch_bounds__(512) __global__
    void rmsnorm_d2048_bf16x2_kernel(const __nv_bfloat162* x, const __nv_bfloat162* weight,
                                     const __nv_bfloat162* z, __nv_bfloat162* out,
                                     std::int64_t rows, float eps) {
    constexpr int kBlock       = 512;
    constexpr int kPairsPerRow = 1024;
    const std::int64_t row     = static_cast<std::int64_t>(blockIdx.x);
    if (row >= rows) { return; }

    const std::int64_t row_base = row * kPairsPerRow;
    const int pair0             = static_cast<int>(threadIdx.x);
    const int pair1             = pair0 + kBlock;
    const __nv_bfloat162 value0 = x[row_base + pair0];
    const __nv_bfloat162 value1 = x[row_base + pair1];
    const float2 x0             = __bfloat1622float2(value0);
    const float2 x1             = __bfloat1622float2(value1);
    const float local_sum       = x0.x * x0.x + x0.y * x0.y + x1.x * x1.x + x1.y * x1.y;

    __shared__ float warp_sums[kBlock / kWarpSize];
    __shared__ float inv_shared;
    const float sum = block_reduce_sum<kBlock>(local_sum, warp_sums);
    if (threadIdx.x == 0) { inv_shared = rsqrtf(sum * (1.0f / 2048.0f) + eps); }
    __syncthreads();
    const float inv = inv_shared;

    const float2 w0 = __bfloat1622float2(weight[pair0]);
    const float2 w1 = __bfloat1622float2(weight[pair1]);
    float2 z0{0.0f, 0.0f};
    float2 z1{0.0f, 0.0f};
    if constexpr (Epilogue == RmsEpilogue::Gated) {
        z0 = __bfloat1622float2(z[row_base + pair0]);
        z1 = __bfloat1622float2(z[row_base + pair1]);
    }
    out[row_base + pair0] =
        __floats2bfloat162_rn(rmsnorm_epilogue<Epilogue>(x0.x, inv, w0.x, z0.x),
                              rmsnorm_epilogue<Epilogue>(x0.y, inv, w0.y, z0.y));
    out[row_base + pair1] =
        __floats2bfloat162_rn(rmsnorm_epilogue<Epilogue>(x1.x, inv, w1.x, z1.x),
                              rmsnorm_epilogue<Epilogue>(x1.y, inv, w1.y, z1.y));
}

// Functional fallback outside the aligned fast domains. It intentionally favors a simple complete
// implementation over another family of shape-specific paths.
template <RmsEpilogue Epilogue>
__launch_bounds__(256) __global__
    void rmsnorm_generic_kernel(const __nv_bfloat16* x, const __nv_bfloat16* weight,
                                const __nv_bfloat16* z, __nv_bfloat16* out, std::int32_t d,
                                std::int64_t rows, float eps) {
    const std::int64_t row = static_cast<std::int64_t>(blockIdx.x);
    if (row >= rows) { return; }

    const std::int64_t base = row * static_cast<std::int64_t>(d);
    float sum               = 0.0f;
    for (std::int64_t i = threadIdx.x; i < static_cast<std::int64_t>(d); i += blockDim.x) {
        const float xv = __bfloat162float(x[base + i]);
        sum += xv * xv;
    }

    __shared__ float scratch[256];
    scratch[threadIdx.x] = sum;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) { scratch[threadIdx.x] += scratch[threadIdx.x + stride]; }
        __syncthreads();
    }

    const float inv = rsqrtf(scratch[0] / static_cast<float>(d) + eps);
    for (std::int64_t i = threadIdx.x; i < static_cast<std::int64_t>(d); i += blockDim.x) {
        const std::int64_t index = base + i;
        const float xv           = __bfloat162float(x[index]);
        const float wv           = __bfloat162float(weight[i]);
        float zv                 = 0.0f;
        if constexpr (Epilogue == RmsEpilogue::Gated) { zv = __bfloat162float(z[index]); }
        out[index] = __float2bfloat16_rn(rmsnorm_epilogue<Epilogue>(xv, inv, wv, zv));
    }
}

} // namespace ninfer::ops
#pragma once

// BF16-cache causal prompt kernel. INT8 has an independent kernel body and resource
// policy in causal_attention_prompt_i8.cuh.
//
//   * Br = 64 query rows and Bc = 64 key columns per CTA tile.
//   * 4 warps / 128 threads; each warp owns 16 query rows of the tile.
//   * Q, K, V staged in 96 KiB of dynamic shared memory (single-buffered), with
//     the cp.async of the next K/V tile overlapped against the current
//     QK / PV tensor-core work (exactly FA's single-buffer overlap pattern).
//   * m16n8k16 bf16 MMA for both S = Q Kᵀ and O += P V, online softmax in exp2.
//
// The op first writes the new chunk K/V into absolute positions in the paged cache,
// then computes causal attention for
// every chunk token over all cached history using bottom-right causal alignment
// (query row i attends to keys [0, base_pos + i]).

#include <math_constants.h>

#include "ops/softmax_attention/dense/causal_cache/prompt_common.cuh"

namespace ninfer::ops {

// Stage one [Bc, D] K or V tile from the per-kv-head contiguous cache into the
// swizzled smem buffer. Keys beyond max_query_abs (which the causal mask always
// drops) are zeroed so the padded/uninitialized cache tail never feeds NaNs into
// the tensor cores. Mirrors FA's predicated K/V cp.async + Clear_OOB path.
template <typename Geometry>
__device__ __forceinline__ void
causal_prompt_stage_kv(__nv_bfloat16* dst, const __nv_bfloat16* cache, int kv_head, int k0,
                       int max_query_abs, int physical_page, int tid) {
    constexpr int D         = kCausalPromptHeadDim;
    constexpr int Bc        = kCausalPromptBc;
    constexpr int Threads   = kCausalPromptThreads;
    constexpr int VecPerRow = D / 8; // 8 bf16 per 16B cp.async
    const bool full_tile    = (k0 + Bc - 1) <= max_query_abs;
    // Block base pointer computed once (int64); per-element offsets stay 32-bit.
    const __nv_bfloat16* cache_block =
        cache + paged_kv_element_offset<kCausalPromptHeadDim, Geometry::KVHeads>(
                    physical_page, kv_head, k0 & kPagedKVPageMask, 0);
    if (full_tile) {
#pragma unroll
        for (int chunk = tid; chunk < Bc * VecPerRow; chunk += Threads) {
            const int key_l  = chunk >> 5;        // / VecPerRow (32)
            const int d      = (chunk & 31) << 3; // (chunk % 32) * 8
            __nv_bfloat16* p = &dst[key_l * D + causal_prompt_swz(key_l, d)];
            cp_async<16, Cache::cg>(p, &cache_block[key_l * D + d]);
        }
    } else {
#pragma unroll
        for (int chunk = tid; chunk < Bc * VecPerRow; chunk += Threads) {
            const int key_l  = chunk >> 5;        // / VecPerRow (32)
            const int d      = (chunk & 31) << 3; // (chunk % 32) * 8
            __nv_bfloat16* p = &dst[key_l * D + causal_prompt_swz(key_l, d)];
            if ((k0 + key_l) <= max_query_abs) {
                cp_async<16, Cache::cg>(p, &cache_block[key_l * D + d]);
            } else {
                store_vec(p, make_int4(0, 0, 0, 0));
            }
        }
    }
}

// FlashAttention-2 forward, one CTA per (query 64-row block, query head). Grid is
// (ceil(tokens/64), q_heads). seqlen_q = tokens, seqlen_k = base_pos + tokens, with
// bottom-right causal alignment (query row i sees keys [0, base_pos + i]).
template <typename Geometry, typename Metadata>
__launch_bounds__(kCausalPromptThreads, 1) __global__
    void causal_attention_prompt_bf16_kernel(const __nv_bfloat16* __restrict__ q,
                                             const __nv_bfloat16* __restrict__ cache_k,
                                             const __nv_bfloat16* __restrict__ cache_v,
                                             Metadata metadata,
                                             const std::int32_t* __restrict__ positions,
                                             float scale, __nv_bfloat16* __restrict__ out,
                                             std::int32_t width) {
    constexpr int D             = kCausalPromptHeadDim; // 256
    constexpr int Br            = kCausalPromptBr;      // 64 query rows
    constexpr int Bc            = kCausalPromptBc;      // 64 key cols
    constexpr int Threads       = kCausalPromptThreads; // 128
    constexpr int QKNt          = Bc / 8;               // 8  QK score n-tiles
    constexpr int QKKs          = D / 16;               // 16 QK contraction steps over head_dim
    constexpr int PVNt          = D / 8;                // 32 PV output n-tiles
    constexpr int PVKs          = Bc / 16;              // 4  PV contraction steps over keys
    constexpr float Log2E       = 1.4426950408889634074f;
    constexpr unsigned FullMask = 0xffffffffu;

    static_assert(Threads == 128);

    extern __shared__ __align__(16) __nv_bfloat16 causal_smem[];
    __nv_bfloat16* q_s = causal_smem;  // [Br, D] swizzled
    __nv_bfloat16* k_s = q_s + Br * D; // [Bc, D] swizzled
    __nv_bfloat16* v_s = k_s + Bc * D; // [Bc, D] swizzled

    const int q_block = static_cast<int>(blockIdx.x);
    const int q_head  = static_cast<int>(blockIdx.y);
    const int tid     = static_cast<int>(threadIdx.x);
    const int warp    = tid >> 5;
    const int lane    = tid & 31;
    const int q0      = q_block * Br;
    const int kv_head = q_head / Geometry::GroupSize;
    const int tokens  = metadata.valid_tokens(width);

    if (q_head >= Geometry::QHeads || q0 >= width) { return; }
    if (q0 >= tokens) {
        causal_prompt_zero_output_rows<Geometry>(out, q_head, q0, min(q0 + Br, width), tid,
                                                 Threads);
        return;
    }
    const int base_pos              = positions[0];
    const std::int32_t* block_table = metadata.block_table();

    const int gid = lane >> 2;
    const int lid = lane & 3;

    const int a_mat     = lane >> 3;
    const int a_rin     = lane & 7;
    const int a_rowoff  = a_rin + ((a_mat & 1) << 3);
    const int b_rin     = lane & 7;
    const int b_koff    = ((lane >> 3) & 1) << 3;
    const int warp_row0 = warp * 16; // this warp owns rows [warp_row0, warp_row0+16)

    // Per-lane precomputed swizzled ldmatrix base addresses (see causal_prompt_swz_addr).
    const unsigned q_sbase = smem_addr(q_s);
    const unsigned k_sbase = smem_addr(k_s);
    const unsigned v_sbase = smem_addr(v_s);
    // Q A-fragment: row = warp_row0 + a_rowoff, col = k*16 + a_coloff.
    const unsigned q_lane_base = q_sbase + static_cast<unsigned>((warp_row0 + a_rowoff) * 512);
    const unsigned q_as        = static_cast<unsigned>((a_mat >> 1) << 4);
    const unsigned q_r         = static_cast<unsigned>(a_rin << 4);
    // K B-fragment via ldmatrix.x4 (2 n-tiles/instr): lanes 16-31 fetch the +8-key
    // half (extra 4096 bytes), lanes with bit3 set fetch the +8 d-contract half.
    const unsigned k_lane_base =
        k_sbase + static_cast<unsigned>(b_rin * 512) + (static_cast<unsigned>(lane >> 4) << 12);
    const unsigned k_as = static_cast<unsigned>((b_koff >> 3) << 4);
    const unsigned k_r  = static_cast<unsigned>(b_rin << 4);
    // V B-fragment via ldmatrix.x4.trans (2 n-tiles/instr): row = k*16 + (bit3)*8 + b_rin,
    // col = n*8 + (lane>>4)*8.
    const unsigned v_lane_base = v_sbase + static_cast<unsigned>(((lane >> 3) & 1) * 4096) +
                                 static_cast<unsigned>(b_rin * 512);
    const unsigned v_as = static_cast<unsigned>((lane >> 4) << 4);
    const unsigned v_r  = static_cast<unsigned>(b_rin << 4);

    // Stage Q into smem once via cp.async (overlaps with the K(0) prologue load
    // below); it stays resident for the whole key loop. Global Q rows are 256 bf16
    // contiguous, with a token stride of 256*QHeads.
    {
        constexpr int VecPerRow      = D / 8;
        constexpr int QRowStride     = D * Geometry::QHeads; // global stride between tokens
        const __nv_bfloat16* q_block = q + causal_prompt_q_index<Geometry>(q_head, 0, q0);
        if (q0 + Br <= tokens) {
#pragma unroll
            for (int chunk = tid; chunk < Br * VecPerRow; chunk += Threads) {
                const int row    = chunk >> 5;
                const int d      = (chunk & 31) << 3;
                __nv_bfloat16* p = &q_s[row * D + causal_prompt_swz(row, d)];
                cp_async<16, Cache::cg>(p, &q_block[row * QRowStride + d]);
            }
        } else {
#pragma unroll
            for (int chunk = tid; chunk < Br * VecPerRow; chunk += Threads) {
                const int row    = chunk >> 5;
                const int d      = (chunk & 31) << 3;
                __nv_bfloat16* p = &q_s[row * D + causal_prompt_swz(row, d)];
                if (q0 + row < tokens) {
                    cp_async<16, Cache::cg>(p, &q_block[row * QRowStride + d]);
                } else {
                    store_vec(p, make_int4(0, 0, 0, 0));
                }
            }
        }
    }

    float acc[PVNt][4];
#pragma unroll
    for (int n = 0; n < PVNt; ++n) {
#pragma unroll
        for (int i = 0; i < 4; ++i) { acc[n][i] = 0.0f; }
    }
    float m0 = -CUDART_INF_F, m1 = -CUDART_INF_F, l0 = 0.0f, l1 = 0.0f;

    const int tile_rows     = min(Br, tokens - q0);
    const int max_query_abs = base_pos + q0 + tile_rows - 1;
    const int n_block_max   = (max_query_abs / Bc) + 1; // n_block_min == 0

    // Fold softmax_scale into the exp2 (FA-style): scores stay raw, so the
    // per-element "* scale" multiply drops out of the QK epilogue entirely.
    const float scale_l2 = scale * Log2E;
    int physical_page    = block_table[0];

    // Prologue: commit Q, then kick off K(0). The loop's wait<0> below drains both.
    ninfer::ops::cp_commit();
    causal_prompt_stage_kv<Geometry>(k_s, cache_k, kv_head, 0, max_query_abs, physical_page, tid);
    ninfer::ops::cp_commit();

    for (int kb = 0; kb < n_block_max; ++kb) {
        const int k0                 = kb * Bc;
        const int next_physical_page = (kb + 1 < n_block_max) ? block_table[kb + 1] : physical_page;

        ninfer::ops::cp_wait<0>(); // K(kb) landed (also publishes q_s / prev PV done)
        __syncthreads();

        // Overlap V(kb) load against the QK MMA below.
        causal_prompt_stage_kv<Geometry>(v_s, cache_v, kv_head, k0, max_query_abs, physical_page,
                                         tid);
        ninfer::ops::cp_commit();

        // S = Q Kᵀ for this warp's 16 rows over all Bc keys, in registers.
        // Software-pipelined like cute's gemm: issue the ldmatrix for contraction
        // step k+1 while the m16n8k16 MMAs for step k run, so the LSU (ldmatrix)
        // and tensor pipes overlap instead of stalling on each other.
        float score[QKNt][4];
#pragma unroll
        for (int nt = 0; nt < QKNt; ++nt) {
            score[nt][0] = score[nt][1] = score[nt][2] = score[nt][3] = 0.0f;
        }
        // Swizzled ldmatrix addresses via precomputed per-lane bases + immediates.
        unsigned af[2][4];
        unsigned bf[2][QKNt][2];
        {
            ldmatrix_x4(af[0][0], af[0][1], af[0][2], af[0][3],
                        causal_prompt_swz_addr(q_lane_base, 0u, q_as, q_r));
#pragma unroll
            for (int nt2 = 0; nt2 < QKNt; nt2 += 2) {
                ldmatrix_x4(bf[0][nt2][0], bf[0][nt2][1], bf[0][nt2 + 1][0], bf[0][nt2 + 1][1],
                            causal_prompt_swz_addr(k_lane_base + static_cast<unsigned>(nt2 * 4096),
                                                   0u, k_as, k_r));
            }
        }
#pragma unroll
        for (int k = 0; k < QKKs; ++k) {
            const int cur = k & 1;
            const int nxt = cur ^ 1;
            if (k + 1 < QKKs) {
                const unsigned ck = static_cast<unsigned>((k + 1) << 5);
                ldmatrix_x4(af[nxt][0], af[nxt][1], af[nxt][2], af[nxt][3],
                            causal_prompt_swz_addr(q_lane_base, ck, q_as, q_r));
#pragma unroll
                for (int nt2 = 0; nt2 < QKNt; nt2 += 2) {
                    ldmatrix_x4(
                        bf[nxt][nt2][0], bf[nxt][nt2][1], bf[nxt][nt2 + 1][0], bf[nxt][nt2 + 1][1],
                        causal_prompt_swz_addr(k_lane_base + static_cast<unsigned>(nt2 * 4096), ck,
                                               k_as, k_r));
                }
            }
#pragma unroll
            for (int nt = 0; nt < QKNt; ++nt) {
                mma_bf16(score[nt][0], score[nt][1], score[nt][2], score[nt][3], af[cur][0],
                         af[cur][1], af[cur][2], af[cur][3], bf[cur][nt][0], bf[cur][nt][1]);
            }
        }

        const int row0             = warp_row0 + gid;
        const int row1             = warp_row0 + gid + 8;
        const int qrow0            = q0 + row0;
        const int qrow1            = q0 + row1;
        const int qabs0            = (qrow0 < tokens) ? base_pos + qrow0 : -1;
        const int qabs1            = (qrow1 < tokens) ? base_pos + qrow1 : -1;
        const bool full_score_tile = (q0 + Br <= tokens) && ((k0 + Bc - 1) <= (base_pos + q0));

        // block row-max on raw (unscaled) scores; scale is folded into exp2 below
        float bm0 = -CUDART_INF_F, bm1 = -CUDART_INF_F;
        if (full_score_tile) {
#pragma unroll
            for (int nt = 0; nt < QKNt; ++nt) {
                bm0 = fmaxf(bm0, fmaxf(score[nt][0], score[nt][1]));
                bm1 = fmaxf(bm1, fmaxf(score[nt][2], score[nt][3]));
            }
        } else {
#pragma unroll
            for (int nt = 0; nt < QKNt; ++nt) {
                const int key0 = k0 + nt * 8 + 2 * lid;
                const int key1 = key0 + 1;
                score[nt][0]   = (qrow0 < tokens && key0 <= qabs0) ? score[nt][0] : -CUDART_INF_F;
                score[nt][1]   = (qrow0 < tokens && key1 <= qabs0) ? score[nt][1] : -CUDART_INF_F;
                score[nt][2]   = (qrow1 < tokens && key0 <= qabs1) ? score[nt][2] : -CUDART_INF_F;
                score[nt][3]   = (qrow1 < tokens && key1 <= qabs1) ? score[nt][3] : -CUDART_INF_F;
                bm0            = fmaxf(bm0, fmaxf(score[nt][0], score[nt][1]));
                bm1            = fmaxf(bm1, fmaxf(score[nt][2], score[nt][3]));
            }
        }
        bm0 = warp_max<4>(bm0, FullMask);
        bm1 = warp_max<4>(bm1, FullMask);

        const float nm0        = fmaxf(m0, bm0);
        const float nm1        = fmaxf(m1, bm1);
        const float nm0_scaled = nm0 * scale_l2;
        const float nm1_scaled = nm1 * scale_l2;
        const float alpha0     = exp2_approx(__fmaf_rn(m0, scale_l2, -nm0_scaled));
        const float alpha1     = exp2_approx(__fmaf_rn(m1, scale_l2, -nm1_scaled));

        // P = exp2(S - m), repacked into the PV A-fragment layout, plus local block row-sum.
        // The row-sum allreduce is deferred to the epilogue; only row max must be reduced per tile.
        float bl0 = 0.0f, bl1 = 0.0f;
        unsigned p_frag[PVKs][4];
        if (full_score_tile) {
#pragma unroll
            for (int nt = 0; nt < QKNt; ++nt) {
                const float p00 = exp2_approx(__fmaf_rn(score[nt][0], scale_l2, -nm0_scaled));
                const float p01 = exp2_approx(__fmaf_rn(score[nt][1], scale_l2, -nm0_scaled));
                const float p10 = exp2_approx(__fmaf_rn(score[nt][2], scale_l2, -nm1_scaled));
                const float p11 = exp2_approx(__fmaf_rn(score[nt][3], scale_l2, -nm1_scaled));
                bl0 += p00 + p01;
                bl1 += p10 + p11;
                const int pk = nt >> 1;
                if ((nt & 1) == 0) {
                    p_frag[pk][0] = pack_bf16x2(p00, p01);
                    p_frag[pk][1] = pack_bf16x2(p10, p11);
                } else {
                    p_frag[pk][2] = pack_bf16x2(p00, p01);
                    p_frag[pk][3] = pack_bf16x2(p10, p11);
                }
            }
        } else {
#pragma unroll
            for (int nt = 0; nt < QKNt; ++nt) {
                const float p00 = (score[nt][0] > -CUDART_INF_F)
                                      ? exp2_approx(__fmaf_rn(score[nt][0], scale_l2, -nm0_scaled))
                                      : 0.0f;
                const float p01 = (score[nt][1] > -CUDART_INF_F)
                                      ? exp2_approx(__fmaf_rn(score[nt][1], scale_l2, -nm0_scaled))
                                      : 0.0f;
                const float p10 = (score[nt][2] > -CUDART_INF_F)
                                      ? exp2_approx(__fmaf_rn(score[nt][2], scale_l2, -nm1_scaled))
                                      : 0.0f;
                const float p11 = (score[nt][3] > -CUDART_INF_F)
                                      ? exp2_approx(__fmaf_rn(score[nt][3], scale_l2, -nm1_scaled))
                                      : 0.0f;
                bl0 += p00 + p01;
                bl1 += p10 + p11;
                const int pk = nt >> 1;
                if ((nt & 1) == 0) {
                    p_frag[pk][0] = pack_bf16x2(p00, p01);
                    p_frag[pk][1] = pack_bf16x2(p10, p11);
                } else {
                    p_frag[pk][2] = pack_bf16x2(p00, p01);
                    p_frag[pk][3] = pack_bf16x2(p10, p11);
                }
            }
        }

        l0 = __fmaf_rn(l0, alpha0, bl0);
        l1 = __fmaf_rn(l1, alpha1, bl1);
        m0 = nm0;
        m1 = nm1;
#pragma unroll
        for (int n = 0; n < PVNt; ++n) {
            acc[n][0] *= alpha0;
            acc[n][1] *= alpha0;
            acc[n][2] *= alpha1;
            acc[n][3] *= alpha1;
        }

        ninfer::ops::cp_wait<0>(); // V(kb) landed; QK done reading k_s
        __syncthreads();

        // Prefetch K(kb+1) into the (now-free) K buffer, overlapping the PV MMA.
        if (kb + 1 < n_block_max) {
            physical_page = next_physical_page;
            causal_prompt_stage_kv<Geometry>(k_s, cache_k, kv_head, (kb + 1) * Bc, max_query_abs,
                                             physical_page, tid);
            ninfer::ops::cp_commit();
        }

        // O += P V, contracting over the Bc keys. The (k, n) iteration space is
        // flattened and software-pipelined: the transposed ldmatrix for the next
        // V fragment is issued while the current MMA runs.
        // Each x4.trans load covers 2 output n-tiles (16 dims); pipeline the next
        // load against the current pair of MMAs.
        constexpr int PVHalf  = PVNt / 2;      // 16 n-tile pairs
        constexpr int PVLoads = PVKs * PVHalf; // 64 x4.trans loads
        // Swizzled V x4.trans addresses via precomputed per-lane base + immediates.
        unsigned vf[2][4];
        {
            ldmatrix_x4_t(vf[0][0], vf[0][1], vf[0][2], vf[0][3],
                          causal_prompt_swz_addr(v_lane_base, 0u, v_as, v_r));
        }
#pragma unroll
        for (int li = 0; li < PVLoads; ++li) {
            const int k   = li / PVHalf;
            const int n2  = (li % PVHalf) * 2;
            const int cur = li & 1;
            const int nxt = cur ^ 1;
            if (li + 1 < PVLoads) {
                const int k2       = (li + 1) / PVHalf;
                const int n2b      = ((li + 1) % PVHalf) * 2;
                const unsigned ckv = static_cast<unsigned>(n2b << 4);
                ldmatrix_x4_t(vf[nxt][0], vf[nxt][1], vf[nxt][2], vf[nxt][3],
                              causal_prompt_swz_addr(v_lane_base + static_cast<unsigned>(k2 * 8192),
                                                     ckv, v_as, v_r));
            }
            mma_bf16(acc[n2][0], acc[n2][1], acc[n2][2], acc[n2][3], p_frag[k][0], p_frag[k][1],
                     p_frag[k][2], p_frag[k][3], vf[cur][0], vf[cur][1]);
            mma_bf16(acc[n2 + 1][0], acc[n2 + 1][1], acc[n2 + 1][2], acc[n2 + 1][3], p_frag[k][0],
                     p_frag[k][1], p_frag[k][2], p_frag[k][3], vf[cur][2], vf[cur][3]);
        }
    }

    l0 = warp_sum<4>(l0, FullMask);
    l1 = warp_sum<4>(l1, FullMask);

    // Normalize once per row via reciprocal-multiply instead of 128 IEEE divides.
    const float inv_l0 = (l0 > 0.0f) ? __frcp_rn(l0) : 0.0f;
    const float inv_l1 = (l1 > 0.0f) ? __frcp_rn(l1) : 0.0f;
#pragma unroll
    for (int n = 0; n < PVNt; ++n) {
        const int d0    = n * 8 + 2 * lid;
        const int qrow0 = q0 + warp_row0 + gid;
        const int qrow1 = q0 + warp_row0 + gid + 8;
        if (qrow0 < tokens) {
            *reinterpret_cast<unsigned*>(&out[causal_prompt_q_index<Geometry>(q_head, d0, qrow0)]) =
                pack_bf16x2(acc[n][0] * inv_l0, acc[n][1] * inv_l0);
        }
        if (qrow1 < tokens) {
            *reinterpret_cast<unsigned*>(&out[causal_prompt_q_index<Geometry>(q_head, d0, qrow1)]) =
                pack_bf16x2(acc[n][2] * inv_l1, acc[n][3] * inv_l1);
        }
    }
    causal_prompt_zero_output_rows<Geometry>(out, q_head, tokens, min(q0 + Br, width), tid,
                                             Threads);
}

} // namespace ninfer::ops
#pragma once

#include "core/tensor.h"

#include <cuda_runtime.h>

namespace ninfer::ops::detail {

void launch_fp8_decode(const Tensor& x, const Weight& weight, Tensor& out, cudaStream_t stream);
void launch_fp8_small_t(const Tensor& x, const Weight& weight, Tensor& out, cudaStream_t stream);
void launch_fp8_vocabulary_a16_mma(const Tensor& x, const Weight& weight, Tensor& out,
                                   cudaStream_t stream);

} // namespace ninfer::ops::detail
#include "targets/qwen3_6/impl/runtime/instance.h"

#include <ninfer/targets/qwen3_6/prepared_prompt.h>

#include "targets/qwen3_6/impl/runtime/layouts.h"
#include "targets/qwen3_6/impl/runtime/program.h"

#include <stdexcept>
#include <utility>

namespace ninfer::targets::qwen3_6 {

using detail::NINFER_QWEN36_RUNTIME_NS::Variant;

template <>
SequencePlan<Variant>::SequencePlan(
    std::unique_ptr<detail::SequencePlanImpl<Variant>> impl) noexcept
    : impl_(std::move(impl)) {}

template <>
SequencePlan<Variant>::SequencePlan(SequencePlan&&) noexcept = default;
template <>
SequencePlan<Variant>& SequencePlan<Variant>::operator=(SequencePlan&&) noexcept = default;
template <>
SequencePlan<Variant>::~SequencePlan() = default;

template <>
std::uint32_t SequencePlan<Variant>::capacity() const noexcept {
    return impl_ != nullptr ? impl_->capacity : 0;
}

template <>
std::uint32_t SequencePlan<Variant>::kv_capacity() const noexcept {
    return impl_ != nullptr ? impl_->kv_capacity : 0;
}

template <>
std::uint32_t SequencePlan<Variant>::max_concurrency() const noexcept {
    return impl_ != nullptr ? impl_->max_concurrency : 0;
}

template <>
std::size_t SequencePlan<Variant>::device_reservation_bytes() const noexcept {
    return impl_ != nullptr ? impl_->device_reservation_bytes : 0;
}

template <>
std::size_t SequencePlan<Variant>::workspace_capacity_bytes() const noexcept {
    return impl_ != nullptr ? impl_->workspace.capacity : 0;
}

template <>
SequencePlanner<Variant>::SequencePlanner(
    std::unique_ptr<detail::SequencePlannerImpl<Variant>> impl) noexcept
    : impl_(std::move(impl)) {}

template <>
SequencePlanner<Variant>::SequencePlanner(SequencePlanner&&) noexcept = default;
template <>
SequencePlanner<Variant>& SequencePlanner<Variant>::operator=(SequencePlanner&&) noexcept = default;
template <>
SequencePlanner<Variant>::~SequencePlanner() = default;

template <>
const runtime::SequenceCapacityCurve& SequencePlanner<Variant>::capacity_curve() const noexcept {
    static const runtime::SequenceCapacityCurve empty;
    return impl_ != nullptr ? impl_->curve : empty;
}

template <>
SequencePlan<Variant> SequencePlanner<Variant>::finalize(std::uint32_t main_page_groups) && {
    if (impl_ == nullptr) { throw std::logic_error("sequence planner is empty"); }
    return SequencePlan<Variant>(detail::NINFER_QWEN36_RUNTIME_NS::finalize_sequence_plan_impl(
        std::move(impl_), main_page_groups));
}

template <>
RequestBasePlan<Variant>::RequestBasePlan(
    std::unique_ptr<detail::RequestBasePlanImpl<Variant>> impl) noexcept
    : impl_(std::move(impl)) {}

template <>
RequestBasePlan<Variant>::RequestBasePlan(RequestBasePlan&&) noexcept = default;
template <>
RequestBasePlan<Variant>& RequestBasePlan<Variant>::operator=(RequestBasePlan&&) noexcept = default;
template <>
RequestBasePlan<Variant>::~RequestBasePlan() = default;

template <>
const runtime::RequestPlanSummary& RequestBasePlan<Variant>::summary() const noexcept {
    static const runtime::RequestPlanSummary empty;
    return impl_ != nullptr ? impl_->summary : empty;
}

template <>
const PreparedContextCache& RequestBasePlan<Variant>::context_cache() const noexcept {
    static const PreparedContextCache empty;
    return impl_ != nullptr ? impl_->context_cache : empty;
}

template <>
std::optional<PrefixShortlistKey>
RequestBasePlan<Variant>::prefix_shortlist_key(std::uint32_t frontier) const noexcept {
    if (impl_ == nullptr || frontier == 0 || frontier > impl_->prefix_digests.size()) {
        return std::nullopt;
    }
    return PrefixShortlistKey{
        .digest       = impl_->prefix_digests.at(frontier),
        .frontier     = frontier,
        .identity_tag = impl_->prefix_identity_tag,
    };
}

template <>
PressurePlanningSession<Variant>::PressurePlanningSession(
    std::unique_ptr<detail::PressurePlanningSessionImpl<Variant>> impl) noexcept
    : impl_(std::move(impl)) {}

template <>
PressurePlanningSession<Variant>::PressurePlanningSession(PressurePlanningSession&&) noexcept =
    default;

template <>
PressurePlanningSession<Variant>&
PressurePlanningSession<Variant>::operator=(PressurePlanningSession&&) noexcept = default;

template <>
PressurePlanningSession<Variant>::~PressurePlanningSession() = default;

template <>
PressureTargetHandle PressurePlanningSession<Variant>::identity_target(
    const AdmissionCandidate<Variant>& candidate) const {
    if (impl_ == nullptr) { throw std::logic_error("pressure planning session is empty"); }
    return impl_->identity_target(candidate);
}

template <>
PressureTargetHandle PressurePlanningSession<Variant>::root_maximal_target(
    const AdmissionCandidate<Variant>& root_candidate) {
    if (impl_ == nullptr) { throw std::logic_error("pressure planning session is empty"); }
    return impl_->root_maximal_target(root_candidate);
}

template <>
runtime::PressureTargetAssessment
PressurePlanningSession<Variant>::assess(PressureTargetHandle target) {
    if (impl_ == nullptr) { throw std::logic_error("pressure planning session is empty"); }
    return impl_->assess(target);
}

template <>
PreparedPressureExpansion<Variant>
PressurePlanningSession<Variant>::prepare_expansion(PressureTargetHandle parent) {
    if (impl_ == nullptr) { throw std::logic_error("pressure planning session is empty"); }
    return impl_->prepare_expansion(parent);
}

template <>
PressureExpansionView
PressurePlanningSession<Variant>::commit_expansion(PreparedPressureExpansion<Variant>&& prepared) {
    if (impl_ == nullptr) { throw std::logic_error("pressure planning session is empty"); }
    return impl_->commit_expansion(std::move(prepared));
}

template <>
void PressurePlanningSession<Variant>::discard_expansion(
    PreparedPressureExpansion<Variant>&& prepared) noexcept {
    if (impl_ != nullptr) { impl_->discard_expansion(std::move(prepared)); }
}

template <>
std::optional<ResourcePlan<Variant>>
PressurePlanningSession<Variant>::seal(PressureTargetHandle target, const PreparedPrompt& prompt) {
    if (impl_ == nullptr) { throw std::logic_error("pressure planning session is empty"); }
    std::optional<AdmissionCandidate<Variant>> sealed =
        impl_->seal(target, PreparedPromptAccess::view(prompt));
    if (!sealed) { return std::nullopt; }
    const bool needs_transfer = sealed->impl_->needs_transfer;
    return ResourcePlan<Variant>(std::move(*sealed), impl_->resource_revision, needs_transfer);
}

template <>
Program<Variant>::Program(std::unique_ptr<detail::ProgramImpl<Variant>> impl) noexcept
    : impl_(std::move(impl)) {}

template <>
Program<Variant>::~Program() noexcept = default;

template <>
RequestBasePlan<Variant>
Program<Variant>::plan_request(const PreparedPrompt& prompt,
                               const runtime::ResolvedExecutionOptions& options) {
    return impl_->plan_request(PreparedPromptAccess::view(prompt), options);
}

template <>
std::optional<AdmissionCandidate<Variant>> Program<Variant>::inspect_admission(
    const PreparedPrompt& prompt, const RequestBasePlan<Variant>& base, runtime::LaneId destination,
    const ContinuationHandle<Variant>* source, const SharedPrefixHandle<Variant>* shared_source,
    std::optional<runtime::CheckpointRef> checkpoint, bool must_retain_private_source,
    const runtime::ContextMachineCostModel& machine_cost) {
    return impl_->inspect_admission(PreparedPromptAccess::view(prompt), base, destination, source,
                                    shared_source, checkpoint, must_retain_private_source,
                                    machine_cost);
}

template <>
std::optional<ResourcePlan<Variant>>
Program<Variant>::seal_identity(const AdmissionCandidate<Variant>& admission,
                                const PreparedPrompt& prompt) {
    std::optional<AdmissionCandidate<Variant>> sealed =
        impl_->seal_materialization(admission, PreparedPromptAccess::view(prompt), {}, {}, {}, {});
    if (!sealed) { return std::nullopt; }
    const bool needs_transfer = sealed->impl_->needs_transfer;
    return ResourcePlan<Variant>(std::move(*sealed), impl_->resource_revision(), needs_transfer);
}

template <>
PressurePlanningSession<Variant> Program<Variant>::begin_pressure_planning(
    const runtime::ContextMachineCostModel& machine_cost,
    std::span<const AdmissionCandidate<Variant>* const> candidates,
    std::span<const ContinuationHandle<Variant>* const> private_owners,
    std::span<const std::uint32_t> private_owner_ordinals,
    std::span<const SharedPrefixHandle<Variant>* const> shared_owners,
    std::span<const std::uint32_t> shared_owner_ordinals) {
    return PressurePlanningSession<Variant>(
        std::make_unique<detail::PressurePlanningSessionImpl<Variant>>(
            *impl_, machine_cost, candidates, private_owners, private_owner_ordinals, shared_owners,
            shared_owner_ordinals));
}

template <>
runtime::ContextTransactionReserveStatus
Program<Variant>::start_resource_transaction(ResourcePlan<Variant>&& plan, PreparedPrompt&& prompt,
                                             runtime::CancellationFlagView cancellation) {
    if (plan.revision_ == 0 || plan.revision_ != impl_->resource_revision()) {
        return runtime::ContextTransactionReserveStatus::Aborted;
    }
    return impl_->reserve_materialization(
        std::move(plan.admission_), PreparedPromptAccess::take(std::move(prompt)), cancellation);
}

template <>
std::optional<PersistentBackfillProof<Variant>> Program<Variant>::prove_persistent_backfill(
    const RequestBasePlan<Variant>& blocked_head, const ResourcePlan<Variant>& candidate,
    std::span<const SequenceHandle<Variant>> persistent_borrowers) const {
    if (candidate.revision_ == 0 || candidate.revision_ != impl_->resource_revision() ||
        !impl_->persistent_backfill_safe(blocked_head, candidate.admission_,
                                         persistent_borrowers)) {
        return std::nullopt;
    }
    return PersistentBackfillProof<Variant>(candidate.revision_);
}

template <>
ContextTransactionProgress<Variant>
Program<Variant>::progress_context_transaction(runtime::CancellationFlagView cancellation) {
    return impl_->progress_context_transaction(cancellation);
}

template <>
void Program<Variant>::finalize_context_transaction() noexcept {
    impl_->finalize_context_transaction();
}

template <>
bool Program<Variant>::has_context_transaction() const noexcept {
    return impl_->has_context_transaction();
}

template <>
PrefillProgress<Variant>
Program<Variant>::advance_prefill(SequenceHandle<Variant> sequence,
                                  runtime::ExecutionTiming* failed_timing) {
    return impl_->advance_prefill(sequence, failed_timing);
}

template <>
CaptureAssessment
Program<Variant>::inspect_capture(const CaptureOffer<Variant>& offer,
                                  const SharedPrefixHandle<Variant>* exact_shared,
                                  const SharedPrefixHandle<Variant>* replacement,
                                  std::optional<runtime::CheckpointRef> private_replacement) const {
    return impl_->inspect_capture(offer, exact_shared, replacement, private_replacement);
}

template <>
bool Program<Variant>::shared_capture_matches(const CaptureOffer<Variant>& offer,
                                              const SharedPrefixHandle<Variant>& shared) const {
    return impl_->shared_capture_matches(offer, shared);
}

template <>
void Program<Variant>::skip_capture(CaptureOffer<Variant>&& offer) {
    impl_->skip_capture(std::move(offer));
}

template <>
runtime::ContextTransactionReserveStatus
Program<Variant>::reserve_active_capture(CaptureOffer<Variant>&& offer,
                                         const SharedPrefixHandle<Variant>* exact_shared,
                                         const SharedPrefixHandle<Variant>* replacement,
                                         std::optional<runtime::CheckpointRef> private_replacement,
                                         runtime::CancellationFlagView cancellation) {
    return impl_->reserve_active_capture(std::move(offer), exact_shared, replacement,
                                         private_replacement, cancellation);
}

template <>
PendingBatch<Variant> Program<Variant>::decode(std::span<const SequenceHandle<Variant>> sequences,
                                               std::span<const runtime::RoundBudget> budgets,
                                               runtime::ExecutionTiming* failed_timing) {
    return impl_->decode(sequences, budgets, failed_timing);
}

template <>
runtime::ExecutionTiming Program<Variant>::append_forced_tokens(
    std::span<const SequenceHandle<Variant>> sequences, std::span<const TokenId> row_major_tokens,
    std::uint32_t row_stride, runtime::ExecutionTiming* failed_timing) {
    return impl_->append_forced_tokens(sequences, row_major_tokens, row_stride, failed_timing);
}

template <>
CommitResult<Variant> Program<Variant>::commit(PendingBatch<Variant>&& pending,
                                               std::span<const runtime::CommitDecision> decisions,
                                               runtime::CommitObservation observation,
                                               runtime::ExecutionTiming* failed_timing) {
    return impl_->commit(std::move(pending), decisions, observation, failed_timing);
}

template <>
DiscardResult<Variant> Program<Variant>::abort_pending(PendingBatch<Variant>&& pending) noexcept {
    return impl_->abort_pending(std::move(pending));
}

template <>
FinishResult<Variant> Program<Variant>::finish(SequenceHandle<Variant> sequence) noexcept {
    return impl_->finish(sequence);
}

template <>
AbortResult<Variant> Program<Variant>::abort(SequenceHandle<Variant> sequence) noexcept {
    return impl_->abort(sequence);
}

template <>
ReleaseResult<Variant>
Program<Variant>::release_continuation(ContinuationHandle<Variant>&& continuation) noexcept {
    return impl_->release_continuation(std::move(continuation));
}

template <>
ReleaseResult<Variant>
Program<Variant>::release_shared_prefix(SharedPrefixHandle<Variant>&& shared) noexcept {
    return impl_->release_shared_prefix(std::move(shared));
}

template <>
void Program<Variant>::fail_all_cleanup() noexcept {
    impl_->fail_all_cleanup();
}

template <>
bool Program<Variant>::isolated_request_feasible(
    const RequestBasePlan<Variant>& base) const noexcept {
    return impl_->isolated_request_feasible(base);
}

template <>
std::uint64_t Program<Variant>::resource_revision() const noexcept {
    return impl_->resource_revision();
}

template <>
PhysicalUsageSnapshot Program<Variant>::physical_usage() const noexcept {
    return impl_->physical_usage();
}

template <>
MemorySummary Program<Variant>::memory_summary() const noexcept {
    return impl_->memory_summary();
}

template <>
void Program<Variant>::reset_memory_peaks() noexcept {
    impl_->reset_memory_peaks();
}

template <>
SequencePlanner<Variant> make_sequence_planner<Variant>(DeviceContext& device,
                                                        const EngineOptions& options,
                                                        Variant::WeightsProfile weights_profile) {
    return SequencePlanner<Variant>(detail::NINFER_QWEN36_RUNTIME_NS::make_sequence_planner_impl(
        device, options, weights_profile));
}

template <>
std::unique_ptr<Program<Variant>>
create_program<Variant>(const Variant::ModelView& model, Variant::WeightsProfile weights_profile,
                        SequencePlan<Variant>&& plan, DeviceContext& device) {
    if (plan.impl_ == nullptr) { throw std::invalid_argument("sequence plan is empty"); }
    if (plan.impl_->weights_profile != weights_profile) {
        throw std::invalid_argument(
            "loaded model weights profile does not match the sequence plan");
    }
    auto impl = std::make_unique<detail::ProgramImpl<Variant>>(model, *plan.impl_, device);
    plan.impl_.reset();
    return std::unique_ptr<Program<Variant>>(new Program<Variant>(std::move(impl)));
}

} // namespace ninfer::targets::qwen3_6
#include "targets/qwen3_6/impl/runtime/instance.h"
#include "targets/qwen3_6/impl/runtime/schedule.h"

#include "ninfer/ops/mtp_round.h"
#include "ninfer/ops/scatter.h"
#include "ninfer/ops/scalar.h"

#include <cuda_runtime.h>

#include <stdexcept>

namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::schedule {
void mtp_bridge_and_propose(PrefillContext& state, const Tensor& next_token,
                            const Tensor& previous_hidden, std::int32_t position,
                            std::span<const std::int32_t> rope_position, bool build_proposal,
                            const Tensor* next_embedding) {
    if (!state.mtp_kv.valid() || !state.execution.io.mtp) {
        throw std::logic_error("MTP bridge requires MTP storage");
    }
    if (rope_position.size() != 3) {
        throw std::invalid_argument("MTP bridge requires one three-axis rope position");
    }
    state.execution.work.reset();
    TextContext card(state.execution.device, state.execution.model, state.execution.work,
                     state.text_kv, state.execution.linear_attention, state.execution.io,
                     state.execution.prefill_hidden, state.execution.prefill_chunk,
                     state.text_kv_base, state.mtp_kv, &state.text_cache, state.mtp_cache);
    configure_text_card(card, state.execution, state.sampling, state.state_source_slot,
                        state.state_destination_slot, state.mtp_proposal_extent);

    Tensor position_view = state.execution.io.mtp->target_positions.slice(0, 0, 1);
    ops::set_i32_scalar(position_view, position, state.execution.device.stream);
    Tensor mtp_hidden         = state.execution.io.mtp->ar_hidden;
    Tensor logits             = state.execution.io.logits.slice(1, 0, 1);
    Tensor draft0             = state.execution.io.mtp->draft_tokens.slice(0, 0, 1);
    Tensor rope_position_view = state.execution.work.alloc(DType::I32, {1, 3});
    CUDA_CHECK(cudaMemcpyAsync(rope_position_view.data, rope_position.data(),
                               rope_position.size_bytes(), cudaMemcpyHostToDevice,
                               state.execution.device.stream));
    const auto bridge_visible = static_cast<std::uint32_t>(position + 1);
    const ops::CausalAttentionExecutionEnvelope bridge_envelope{bridge_visible, bridge_visible};
    card.mtp_forward_batch(next_token, previous_hidden, position_view, bridge_envelope, mtp_hidden,
                           build_proposal ? 0 : -1, build_proposal ? &logits : nullptr,
                           build_proposal ? &draft0 : nullptr, &rope_position_view, next_embedding);
    if (!build_proposal) { return; }

    if (state.mtp_proposal_extent == 0 ||
        state.mtp_proposal_extent >
            static_cast<std::uint32_t>(state.execution.io.mtp->draft_tokens.ne[0])) {
        throw std::logic_error("MTP bridge proposal extent is outside the configured window");
    }

    Tensor ar_position = state.execution.io.mtp->position.slice(0, 0, 1);
    ops::set_i32_scalar(ar_position, position + 1, state.execution.device.stream);
    for (int i = 1; i < static_cast<int>(state.mtp_proposal_extent); ++i) {
        Tensor previous_token = state.execution.io.mtp->draft_tokens.slice(0, i - 1, 1);
        Tensor next_draft     = state.execution.io.mtp->draft_tokens.slice(0, i, 1);
        Tensor next_hidden    = state.execution.prefill_hidden.slice(1, i, 1);
        const auto visible    = static_cast<std::uint32_t>(position + i + 1);
        const ops::CausalAttentionExecutionEnvelope envelope{visible, visible};
        card.mtp_forward_ar_step(previous_token, state.execution.io.mtp->ar_hidden, ar_position,
                                 envelope, next_hidden, logits, next_draft);
        CUDA_CHECK(cudaMemcpyAsync(state.execution.io.mtp->ar_hidden.data, next_hidden.data,
                                   state.execution.io.mtp->ar_hidden.bytes(),
                                   cudaMemcpyDeviceToDevice, state.execution.device.stream));
        ops::increment_i32_scalar(ar_position, state.execution.device.stream);
    }
}

auto mtp_decode_batch_body(MtpBatchContext& state, std::int32_t batch_size, std::uint32_t k,
                           MtpCausalAttentionEnvelopes envelopes) {
    return [&state, batch_size, k, envelopes] {
        if (batch_size <= 0 || batch_size > static_cast<std::int32_t>(kMaximumConcurrency) ||
            k == 0 || k > kMtpDecodeMaximumDrafts) {
            throw std::logic_error("MTP decode batch state is incomplete");
        }

        qwen3_6::MtpDecodeState& frame = state.frame;
        const std::int32_t width       = static_cast<std::int32_t>(k) + 1;
        CUDA_CHECK(cudaMemcpyAsync(frame.ingress.data, &state.host_ingress,
                                   sizeof(qwen3_6::MtpDecodeIngress), cudaMemcpyHostToDevice,
                                   state.execution.device.stream));

        TextContext card(state.execution.device, state.execution.model, state.execution.work, {},
                         state.execution.linear_attention, state.execution.io,
                         state.execution.prefill_hidden, state.execution.prefill_chunk, 0, {},
                         &state.text_cache, &state.mtp_cache);
        Tensor anchors            = frame.anchors.slice(0, 0, batch_size);
        Tensor frontiers          = frame.base_frontiers.slice(0, 0, batch_size);
        Tensor budgets            = frame.remaining_budgets.slice(0, 0, batch_size);
        Tensor current_extents    = frame.current_extents.slice(0, 0, batch_size);
        Tensor target_valid       = frame.target_valid_columns.slice(0, 0, batch_size);
        Tensor current_drafts     = frame.current_drafts.slice(1, 0, batch_size);
        Tensor target_rope        = frame.target_rope_positions.slice(1, 0, batch_size);
        Tensor text_rows          = frame.text_kv_table_rows.slice(0, 0, batch_size);
        Tensor mtp_rows           = frame.mtp_kv_table_rows.slice(0, 0, batch_size);
        Tensor state_sources      = frame.state_source_slots.slice(0, 0, batch_size);
        Tensor state_destinations = frame.state_destination_slots.slice(0, 0, batch_size);
        Tensor rope_deltas        = frame.rope_deltas.slice(0, 0, batch_size);
        Tensor verify_ids         = frame.verify_ids.slice(1, 0, batch_size);
        Tensor target_positions   = frame.target_positions.slice(1, 0, batch_size);
        Tensor target_tokens      = frame.target_argmax.slice(1, 0, batch_size);
        Tensor target_logits      = frame.target_logits.slice(2, 0, batch_size);
        Tensor target_hidden      = frame.target_hidden.slice(2, 0, batch_size);
        Tensor selected_hidden    = frame.target_continuation_hidden.slice(1, 0, batch_size);
        Tensor licensed_tokens    = frame.licensed_tokens.slice(1, 0, batch_size);
        Tensor licensed_counts    = frame.licensed_counts.slice(0, 0, batch_size);
        Tensor accepted           = frame.accepted_drafts.slice(0, 0, batch_size);
        Tensor next_extents       = frame.next_extents.slice(0, 0, batch_size);
        Tensor alignment_ids      = frame.alignment_ids.slice(1, 0, batch_size);
        Tensor alignment_hidden   = frame.alignment_hidden.slice(2, 0, batch_size);
        Tensor ar_hidden          = frame.ar_hidden.slice(1, 0, batch_size);
        Tensor next_hidden        = frame.next_hidden.slice(1, 0, batch_size);
        Tensor ar_positions       = frame.ar_positions.slice(0, 0, batch_size);
        Tensor ar_rope_positions  = frame.ar_rope_positions.slice(0, 0, batch_size);
        Tensor ar_valid_columns   = frame.ar_valid_columns.slice(0, 0, batch_size);
        Tensor next_drafts        = frame.next_drafts.slice(0, 0, batch_size);

        ops::speculative_prepare_verify_inputs(anchors, current_drafts, frontiers, current_extents,
                                               verify_ids, target_positions,
                                               state.execution.device.stream);
        target_verify_accept(state.execution, state.continuation_hidden_store, card,
                             TargetVerifyFrameView{
                                 .ids                     = verify_ids,
                                 .cache_positions         = target_positions,
                                 .rope_positions          = target_rope,
                                 .valid_columns           = target_valid,
                                 .kv_table_rows           = text_rows,
                                 .state_source_slots      = state_sources,
                                 .state_destination_slots = state_destinations,
                                 .target_hidden           = target_hidden,
                                 .target_logits           = target_logits,
                                 .target_tokens           = target_tokens,
                                 .drafts                  = current_drafts,
                                 .current_extents         = current_extents,
                                 .frontiers               = frontiers,
                                 .anchors                 = anchors,
                                 .licensed_tokens         = licensed_tokens,
                                 .licensed_counts         = licensed_counts,
                                 .accepted_drafts         = accepted,
                                 .selected_hidden         = selected_hidden,
                                 .replay_records          = state.execution.replay_records,
                                 .sampling                = frame.sampling,
                             },
                             envelopes.target_verify);

        ops::mtp_prepare_next_round(verify_ids, anchors, accepted, frontiers, budgets,
                                    licensed_counts, rope_deltas, alignment_ids, next_extents,
                                    ar_positions, ar_rope_positions, ar_valid_columns,
                                    static_cast<std::int32_t>(state.text_cache.max_context()),
                                    state.execution.device.stream);
        card.mtp_forward_decode_batch(alignment_ids, target_hidden, target_positions, target_rope,
                                      licensed_counts, mtp_rows, envelopes.batch, alignment_hidden);
        ops::speculative_select_accepted_hidden(alignment_hidden, accepted, ar_hidden,
                                                state.execution.device.stream);

        Tensor proposal_logits = frame.proposal_logits.slice(1, 0, batch_size);
        Tensor draft0          = next_drafts.slice(1, 0, 1).view({batch_size});
        card.mtp_propose_batch(ar_hidden, proposal_logits, draft0);
        for (std::uint32_t step = 0; step + 1 < k; ++step) {
            Tensor previous =
                next_drafts.slice(1, static_cast<std::int32_t>(step), 1).view({batch_size});
            Tensor next =
                next_drafts.slice(1, static_cast<std::int32_t>(step + 1), 1).view({batch_size});
            Tensor position =
                ar_positions.slice(1, static_cast<std::int32_t>(step), 1).view({1, batch_size});
            Tensor rope = ar_rope_positions.slice(1, static_cast<std::int32_t>(step), 1)
                              .view({1, batch_size});
            Tensor valid =
                ar_valid_columns.slice(1, static_cast<std::int32_t>(step), 1).view({batch_size});
            Tensor previous_batch    = previous.view({1, batch_size});
            Tensor hidden_batch      = ar_hidden.view({TextConfig::hidden, 1, batch_size});
            Tensor next_hidden_batch = next_hidden.view({TextConfig::hidden, 1, batch_size});
            card.mtp_forward_decode_batch(previous_batch, hidden_batch, position, rope, valid,
                                          mtp_rows, envelopes.ar[step], next_hidden_batch);
            card.mtp_propose_batch(next_hidden, proposal_logits, next);
            CUDA_CHECK(cudaMemcpyAsync(ar_hidden.data, next_hidden.data, ar_hidden.bytes(),
                                       cudaMemcpyDeviceToDevice, state.execution.device.stream));
        }

        CUDA_CHECK(cudaMemcpyAsync(&state.host_egress, frame.egress.data,
                                   sizeof(qwen3_6::MtpDecodeEgress), cudaMemcpyDeviceToHost,
                                   state.execution.device.stream));
    };
}

void capture_mtp_decode_batch(MtpBatchContext& state, std::int32_t batch_size, std::uint32_t k,
                              MtpCausalAttentionEnvelopes envelopes,
                              DecodeGraphDefinition& definition) {
    auto body = mtp_decode_batch_body(state, batch_size, k, envelopes);
    capture_graph(state, definition, body);
}

void mtp_decode_batch(MtpBatchContext& state, std::int32_t batch_size, std::uint32_t k,
                      MtpCausalAttentionEnvelopes envelopes, DecodeGraphExecutable* executable) {
    auto body = mtp_decode_batch_body(state, batch_size, k, envelopes);
    run_prepared(state, executable, body);
}

} // namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::schedule
#pragma once

#include "core/tensor.h"

#include <cuda_runtime.h>

namespace ninfer::ops {

/**
 * Adds one bias value to every element whose fastest-dimension index is the same:
 *
 *   ideal[d,r] = x[d,r] + bias[d],  0 <= d < D.
 *
 * Here r flattens all dimensions above ne[0]. `bias` is contiguous BF16 [D] and `x` is a
 * contiguous BF16 tensor with ne[0]=D. The oracle evaluates `ideal` in FP64 from the represented
 * inputs. The updated BF16 x is promoted and compared directly with that result; output storage
 * rounding belongs to the Op's numerical criterion, not the oracle. Private kernel arithmetic is
 * implementation-defined. `bias` must not overlap `x`. There is no workspace or other state side
 * effect.
 */
void add_bias(const Tensor& bias, Tensor& x, cudaStream_t stream);

} // namespace ninfer::ops
#include "ops/launcher/mtp_pack.h"

#include "ops/common/math.h"
#include "ops/common/token_slices.h"
#include "ops/kernel/mtp_pack.cuh"
#include "core/device.h"

#include <algorithm>
#include <cstdint>

namespace ninfer::ops::detail {

void mtp_pack_fc_input_launch(const Tensor& embedding_norm, const Tensor& hidden_norm, Tensor& out,
                              cudaStream_t stream) {
    constexpr int kBlock = 256;
    for_each_token_slice(embedding_norm.ne[1], 1, [&](int token_offset, int token_count) {
        const Tensor embedding_slice = embedding_norm.slice(1, token_offset, token_count);
        const Tensor hidden_slice    = hidden_norm.slice(1, token_offset, token_count);
        Tensor out_slice             = out.slice(1, token_offset, token_count);
        const dim3 grid(static_cast<unsigned int>(div_up(embedding_norm.ne[0], kBlock)),
                        static_cast<unsigned int>(token_count));
        mtp_pack_fc_input_kernel<<<grid, kBlock, 0, stream>>>(
            static_cast<const __nv_bfloat16*>(embedding_slice.data),
            static_cast<const __nv_bfloat16*>(hidden_slice.data),
            static_cast<__nv_bfloat16*>(out_slice.data), embedding_norm.ne[0]);
        CUDA_CHECK(cudaGetLastError());
    });
}

void mtp_split_attn_in_launch(const Tensor& attn_in, Tensor& q, Tensor& k, Tensor& gate, Tensor& v,
                              cudaStream_t stream) {
    constexpr int kBlock = 256;
    const std::int64_t n = static_cast<std::int64_t>(attn_in.ne[0]) * attn_in.ne[1];
    const int grid =
        static_cast<int>(std::max<std::int64_t>(1, div_up(n, static_cast<std::int64_t>(kBlock))));
    mtp_split_attn_in_kernel<<<grid, kBlock, 0, stream>>>(
        static_cast<const __nv_bfloat16*>(attn_in.data), static_cast<__nv_bfloat16*>(q.data),
        static_cast<__nv_bfloat16*>(k.data), static_cast<__nv_bfloat16*>(gate.data),
        static_cast<__nv_bfloat16*>(v.data), attn_in.ne[1]);
    CUDA_CHECK(cudaGetLastError());
}

} // namespace ninfer::ops::detail
#pragma once

// Q4G64 RowSplit x BF16 Tensor Core GEMM.
//
// out[Rows, Cols] = W[Rows, K] * x[K, Cols]
//
// Each CTA owns a BlockRows x BlockCols output tile and advances K in one
// 64-value quant group at a time. Raw Q4 codes and FP16 scales are staged
// independently from the BF16 activation tile. Q4 values are decoded to BF16
// in shared memory, then consumed by m16n8k16 BF16 MMA with FP32 accumulation.
//
// The template describes only physical kernel structure. Rows, Cols, and K
// remain runtime dimensions; registered shapes select a closed schedule and a
// statically compiled boundary variant outside the kernel.

#include "ops/common/mma.cuh"
#include "ops/linear/q4/q4_rowsplit_storage.cuh"

#include <cuda_bf16.h>
#include <cuda_fp16.h>

#include <cstdint>

namespace ninfer::ops::detail {

enum class Q4FragmentPipeline {
    Serial,
    PingPong,
};

enum class Q4ScaleLoad {
    Scalar16,
    Pair32,
};

template <int BlockRows_, int BlockCols_, int BlockK_, int WarpRows_, int WarpCols_,
          int PipelineStages_, int LaunchBoundsMinBlocks_, Q4FragmentPipeline FragmentPipeline_,
          Cache QuantCache_, Cache ActivationCache_, Q4ScaleLoad ScaleLoadMode_>
struct Q4RowSplitMmaGemmSchedule {
    static constexpr int kBlockRows = BlockRows_;
    static constexpr int kBlockCols = BlockCols_;
    static constexpr int kBlockK    = BlockK_;
    static constexpr int kWarpRows  = WarpRows_;
    static constexpr int kWarpCols  = WarpCols_;

    static constexpr int kPipelineStages                  = PipelineStages_;
    static constexpr int kLaunchBoundsMinBlocks           = LaunchBoundsMinBlocks_;
    static constexpr Q4FragmentPipeline kFragmentPipeline = FragmentPipeline_;
    static constexpr Cache kQuantCache                    = QuantCache_;
    static constexpr Cache kActivationCache               = ActivationCache_;
    static constexpr Q4ScaleLoad kScaleLoadMode           = ScaleLoadMode_;

    static constexpr int kWarpGridRows = kBlockRows / kWarpRows;
    static constexpr int kWarpGridCols = kBlockCols / kWarpCols;
    static constexpr int kWarps        = kWarpGridRows * kWarpGridCols;
    static constexpr int kThreads      = kWarps * 32;
    static constexpr int kMmaRows      = kWarpRows / 16;
    static constexpr int kMmaCols      = kWarpCols / 8;
    static constexpr int kMmaKSteps    = kBlockK / 16;
    static constexpr int kGroupsPerK   = kBlockK / Q4RowSplitStorage::kGroupK;
    static constexpr int kScaleBytes =
        kScaleLoadMode == Q4ScaleLoad::Pair32 ? 4 : Q4RowSplitStorage::kScaleBytesPerGroup;

    static constexpr int kSharedBytes =
        kBlockRows * kBlockK * static_cast<int>(sizeof(__nv_bfloat16)) +
        kPipelineStages * kBlockCols * kBlockK * static_cast<int>(sizeof(__nv_bfloat16)) +
        kPipelineStages * kBlockRows * kGroupsPerK * Q4RowSplitStorage::kCodeBytesPerGroup +
        kPipelineStages * kBlockRows * kGroupsPerK * kScaleBytes;

    static_assert(kBlockRows > 0 && kBlockCols > 0);
    static_assert(kBlockK == Q4RowSplitStorage::kGroupK,
                  "Q4 MMA currently stages exactly one quant group per K tile");
    static_assert(kBlockRows % kWarpRows == 0 && kBlockCols % kWarpCols == 0,
                  "Q4 MMA block tile must divide into warp tiles");
    static_assert(kWarpRows % 16 == 0 && kWarpCols % 8 == 0,
                  "Q4 MMA warp tile must be composed of m16n8 MMA tiles");
    static_assert(kPipelineStages >= 2 && kPipelineStages <= 8,
                  "Q4 MMA cp.async pipeline depth must fit cp_wait");
    static_assert(kLaunchBoundsMinBlocks >= 1);
    static_assert(kWarps >= 1 && kThreads <= 1024);
    static_assert(kSharedBytes <= 48 * 1024,
                  "Q4 MMA staged shared memory exceeds the static 48 KiB budget");
};

// XOR swizzle for a [rows][64] BF16 tile. The eight 16-byte column groups are
// permuted by the low row bits so ldmatrix reads do not repeatedly hit the same
// shared-memory bank group.
__device__ __forceinline__ int q4_mma_swizzle_k64(int row, int col) {
    return (((col >> 3) ^ (row & 7)) << 3) | (col & 7);
}

// clang-format off
template <class Schedule_, bool Full>
__global__ __launch_bounds__(Schedule_::kThreads, Schedule_::kLaunchBoundsMinBlocks)
void q4_rowsplit_gemm_mma_kernel(
    const __nv_bfloat16* __restrict__ x,
    const std::uint8_t* __restrict__ codes,
    const std::uint8_t* __restrict__ scales,
    __nv_bfloat16* __restrict__ out,
    std::int32_t rows,
    std::int32_t k,
    std::int32_t cols,
    std::int32_t padded_k) {
    // clang-format on
    using Schedule       = Schedule_;
    constexpr bool kFull = Full;
    constexpr int BM     = Schedule::kBlockRows;
    constexpr int BN     = Schedule::kBlockCols;
    constexpr int BK     = Schedule::kBlockK;
    constexpr int WM     = Schedule::kWarpRows;
    constexpr int WN     = Schedule::kWarpCols;
    constexpr int MT     = Schedule::kMmaRows;
    constexpr int NT     = Schedule::kMmaCols;
    constexpr int KSUB   = Schedule::kMmaKSteps;
    constexpr int S      = Schedule::kPipelineStages;
    constexpr int GPB    = Schedule::kGroupsPerK;
    constexpr int SB     = Schedule::kScaleBytes;

    __shared__ __align__(16) __nv_bfloat16 As[BM * BK];
    __shared__ __align__(16) __nv_bfloat16 Bs[S][BN * BK];
    __shared__ __align__(16) std::uint8_t Cr[S][BM * GPB * Q4RowSplitStorage::kCodeBytesPerGroup];
    __shared__ __align__(16) std::uint8_t Sr[S][BM * GPB * SB];

    const int groups_per_row = padded_k / Q4RowSplitStorage::kGroupK;
    const int tid            = static_cast<int>(threadIdx.x);
    const int warp           = tid >> 5;
    const int lane           = tid & 31;
    const int warp_row       = warp / Schedule::kWarpGridCols;
    const int warp_col       = warp % Schedule::kWarpGridCols;
    const int mma_row        = lane >> 2;
    const int mma_col        = lane & 3;

    const int row0 = static_cast<int>(blockIdx.x) * BM;
    const int col0 = static_cast<int>(blockIdx.y) * BN;

    float accum[MT][NT][4];
#pragma unroll
    for (int mi = 0; mi < MT; ++mi) {
#pragma unroll
        for (int ni = 0; ni < NT; ++ni) {
            accum[mi][ni][0] = 0.0f;
            accum[mi][ni][1] = 0.0f;
            accum[mi][ni][2] = 0.0f;
            accum[mi][ni][3] = 0.0f;
        }
    }

    const int k_tiles = padded_k / BK;

    const int a_matrix     = lane >> 3;
    const int a_inner_row  = lane & 7;
    const int a_row_offset = a_inner_row + ((a_matrix & 1) << 3);
    const int a_col_offset = (a_matrix >> 1) << 3;
    const int b_inner_row  = lane & 7;
    const int b_k_offset   = ((lane >> 3) & 1) << 3;

    auto stage_activation = [&](int stage, int k_tile) {
        const int k0 = k_tile * BK;
#pragma unroll 1
        for (int item = tid; item < BN * (BK / 8); item += Schedule::kThreads) {
            const int local_col = item / (BK / 8);
            const int k8        = item - local_col * (BK / 8);
            const int kk        = k0 + k8 * 8;
            const int col       = col0 + local_col;
            auto* dst = &Bs[stage][local_col * BK + q4_mma_swizzle_k64(local_col, k8 * 8)];
            if constexpr (kFull) {
                cp_async<16, Schedule::kActivationCache>(
                    dst, &x[static_cast<std::int64_t>(col) * k + kk]);
            } else {
                if (col < cols && kk + 8 <= k) {
                    cp_async<16, Schedule::kActivationCache>(
                        dst, &x[static_cast<std::int64_t>(col) * k + kk]);
                } else {
                    store_vec(dst, make_int4(0, 0, 0, 0));
                }
            }
        }
    };

    auto stage_quant = [&](int stage, int k_tile) {
        const int group0 = (k_tile * BK) / Q4RowSplitStorage::kGroupK;
#pragma unroll 1
        for (int item = tid; item < BM * GPB * 2; item += Schedule::kThreads) {
            const int row_group = item >> 1;
            const int half      = item & 1;
            const int local_row = row_group / GPB;
            const int group     = row_group - local_row * GPB;
            const int row       = row0 + local_row;
            auto* dst = &Cr[stage][row_group * Q4RowSplitStorage::kCodeBytesPerGroup + half * 16];
            if constexpr (kFull) {
                const std::int64_t group_index =
                    static_cast<std::int64_t>(row) * groups_per_row + group0 + group;
                cp_async<16, Schedule::kQuantCache>(
                    dst, &codes[group_index * Q4RowSplitStorage::kCodeBytesPerGroup + half * 16]);
            } else {
                if (row < rows) {
                    const std::int64_t group_index =
                        static_cast<std::int64_t>(row) * groups_per_row + group0 + group;
                    cp_async<16, Schedule::kQuantCache>(
                        dst,
                        &codes[group_index * Q4RowSplitStorage::kCodeBytesPerGroup + half * 16]);
                } else {
                    store_vec(dst, make_int4(0, 0, 0, 0));
                }
            }
        }

#pragma unroll 1
        for (int row_group = tid; row_group < BM * GPB; row_group += Schedule::kThreads) {
            const int local_row   = row_group / GPB;
            const int group       = row_group - local_row * GPB;
            const int row         = row0 + local_row;
            const int scale_group = group0 + group;
            auto* dst             = &Sr[stage][row_group * SB];
            if constexpr (kFull) {
                const std::int64_t group_index =
                    static_cast<std::int64_t>(row) * groups_per_row + scale_group;
                if constexpr (Schedule::kScaleLoadMode == Q4ScaleLoad::Pair32) {
                    const int aligned_group = scale_group & ~1;
                    const std::int64_t aligned_index =
                        static_cast<std::int64_t>(row) * groups_per_row + aligned_group;
                    if (aligned_group + 1 < groups_per_row) {
                        cp_async<4>(
                            dst, &scales[aligned_index * Q4RowSplitStorage::kScaleBytesPerGroup]);
                    } else {
                        *reinterpret_cast<std::uint16_t*>(dst) =
                            *reinterpret_cast<const std::uint16_t*>(
                                &scales[group_index * Q4RowSplitStorage::kScaleBytesPerGroup]);
                        *reinterpret_cast<std::uint16_t*>(dst + 2) = 0;
                    }
                } else {
                    *reinterpret_cast<std::uint16_t*>(dst) =
                        *reinterpret_cast<const std::uint16_t*>(
                            &scales[group_index * Q4RowSplitStorage::kScaleBytesPerGroup]);
                }
            } else {
                if (row < rows) {
                    const std::int64_t group_index =
                        static_cast<std::int64_t>(row) * groups_per_row + scale_group;
                    if constexpr (Schedule::kScaleLoadMode == Q4ScaleLoad::Pair32) {
                        const int aligned_group = scale_group & ~1;
                        const std::int64_t aligned_index =
                            static_cast<std::int64_t>(row) * groups_per_row + aligned_group;
                        if (aligned_group + 1 < groups_per_row) {
                            cp_async<4>(
                                dst,
                                &scales[aligned_index * Q4RowSplitStorage::kScaleBytesPerGroup]);
                        } else {
                            *reinterpret_cast<std::uint16_t*>(dst) =
                                *reinterpret_cast<const std::uint16_t*>(
                                    &scales[group_index * Q4RowSplitStorage::kScaleBytesPerGroup]);
                            *reinterpret_cast<std::uint16_t*>(dst + 2) = 0;
                        }
                    } else {
                        *reinterpret_cast<std::uint16_t*>(dst) =
                            *reinterpret_cast<const std::uint16_t*>(
                                &scales[group_index * Q4RowSplitStorage::kScaleBytesPerGroup]);
                    }
                } else {
                    dst[0] = 0;
                    dst[1] = 0;
                    if constexpr (Schedule::kScaleLoadMode == Q4ScaleLoad::Pair32) {
                        dst[2] = 0;
                        dst[3] = 0;
                    }
                }
            }
        }
    };

    auto stage_inputs = [&](int stage, int k_tile) {
        stage_activation(stage, k_tile);
        stage_quant(stage, k_tile);
    };

    auto decode_weight = [&](int stage, int k_tile) {
        const int scale_group = (k_tile * BK) / Q4RowSplitStorage::kGroupK;
        for (int local_row = warp; local_row < BM; local_row += Schedule::kWarps) {
            auto* dst = &As[local_row * BK];
#pragma unroll
            for (int group = 0; group < GPB; ++group) {
                const int staged_group = local_row * GPB + group;
                const std::uint8_t* scale_ptr =
                    &Sr[stage][staged_group * SB + (Schedule::kScaleLoadMode == Q4ScaleLoad::Pair32
                                                        ? ((scale_group + group) & 1) *
                                                              Q4RowSplitStorage::kScaleBytesPerGroup
                                                        : 0)];
                const __nv_bfloat162 weights =
                    Q4MmaDecodeAtom::decode_pair(Cr[stage], scale_ptr, staged_group, lane);
                const int shared_col =
                    q4_mma_swizzle_k64(local_row, group * Q4RowSplitStorage::kGroupK + 2 * lane);
                store_vec(&dst[shared_col], weights);
            }
        }
    };

#pragma unroll
    for (int stage = 0; stage < S; ++stage) {
        if (stage < k_tiles) { stage_inputs(stage, stage); }
        cp_commit();
    }

    for (int k_tile = 0; k_tile < k_tiles; ++k_tile) {
        const int stage = k_tile % S;
        cp_wait<S - 1>();
        __syncthreads();

        decode_weight(stage, k_tile);
        __syncthreads();

        auto load_fragments = [&](int k_step, unsigned(&a_frag)[MT][4], unsigned(&b_frag)[NT][2]) {
#pragma unroll
            for (int mi = 0; mi < MT; ++mi) {
                const int row = warp_row * WM + mi * 16 + a_row_offset;
                const int col = k_step + a_col_offset;
                ldmatrix_x4(a_frag[mi][0], a_frag[mi][1], a_frag[mi][2], a_frag[mi][3],
                            smem_addr(&As[row * BK + q4_mma_swizzle_k64(row, col)]));
            }
#pragma unroll
            for (int ni = 0; ni < NT; ++ni) {
                const int row = warp_col * WN + ni * 8 + b_inner_row;
                const int col = k_step + b_k_offset;
                ldmatrix_x2(b_frag[ni][0], b_frag[ni][1],
                            smem_addr(&Bs[stage][row * BK + q4_mma_swizzle_k64(row, col)]));
            }
        };

        if constexpr (Schedule::kFragmentPipeline == Q4FragmentPipeline::PingPong) {
            unsigned a_frag[2][MT][4];
            unsigned b_frag[2][NT][2];
            load_fragments(0, a_frag[0], b_frag[0]);
#pragma unroll
            for (int ki = 0; ki < KSUB; ++ki) {
                const int current = ki & 1;
                const int next    = (ki + 1) & 1;
                if (ki + 1 < KSUB) { load_fragments((ki + 1) * 16, a_frag[next], b_frag[next]); }
#pragma unroll
                for (int mi = 0; mi < MT; ++mi) {
#pragma unroll
                    for (int ni = 0; ni < NT; ++ni) {
                        mma_bf16(accum[mi][ni][0], accum[mi][ni][1], accum[mi][ni][2],
                                 accum[mi][ni][3], a_frag[current][mi][0], a_frag[current][mi][1],
                                 a_frag[current][mi][2], a_frag[current][mi][3],
                                 b_frag[current][ni][0], b_frag[current][ni][1]);
                    }
                }
            }
        } else {
            unsigned a_frag[MT][4];
            unsigned b_frag[NT][2];
#pragma unroll
            for (int ki = 0; ki < KSUB; ++ki) {
                load_fragments(ki * 16, a_frag, b_frag);
#pragma unroll
                for (int mi = 0; mi < MT; ++mi) {
#pragma unroll
                    for (int ni = 0; ni < NT; ++ni) {
                        mma_bf16(accum[mi][ni][0], accum[mi][ni][1], accum[mi][ni][2],
                                 accum[mi][ni][3], a_frag[mi][0], a_frag[mi][1], a_frag[mi][2],
                                 a_frag[mi][3], b_frag[ni][0], b_frag[ni][1]);
                    }
                }
            }
        }

        __syncthreads();
        const int prefetch_tile = k_tile + S;
        if (prefetch_tile < k_tiles) { stage_inputs(stage, prefetch_tile); }
        cp_commit();
    }

#pragma unroll
    for (int mi = 0; mi < MT; ++mi) {
        const int output_row0 = row0 + warp_row * WM + mi * 16 + mma_row;
        const int output_row1 = output_row0 + 8;
#pragma unroll
        for (int ni = 0; ni < NT; ++ni) {
            const int output_col0 = col0 + warp_col * WN + ni * 8 + 2 * mma_col;
            const int output_col1 = output_col0 + 1;
            const float* values   = accum[mi][ni];
            if constexpr (kFull) {
                out[static_cast<std::int64_t>(output_col0) * rows + output_row0] =
                    __float2bfloat16_rn(values[0]);
                out[static_cast<std::int64_t>(output_col1) * rows + output_row0] =
                    __float2bfloat16_rn(values[1]);
                out[static_cast<std::int64_t>(output_col0) * rows + output_row1] =
                    __float2bfloat16_rn(values[2]);
                out[static_cast<std::int64_t>(output_col1) * rows + output_row1] =
                    __float2bfloat16_rn(values[3]);
            } else {
                if (output_row0 < rows) {
                    if (output_col0 < cols) {
                        out[static_cast<std::int64_t>(output_col0) * rows + output_row0] =
                            __float2bfloat16_rn(values[0]);
                    }
                    if (output_col1 < cols) {
                        out[static_cast<std::int64_t>(output_col1) * rows + output_row0] =
                            __float2bfloat16_rn(values[1]);
                    }
                }
                if (output_row1 < rows) {
                    if (output_col0 < cols) {
                        out[static_cast<std::int64_t>(output_col0) * rows + output_row1] =
                            __float2bfloat16_rn(values[2]);
                    }
                    if (output_col1 < cols) {
                        out[static_cast<std::int64_t>(output_col1) * rows + output_row1] =
                            __float2bfloat16_rn(values[3]);
                    }
                }
            }
        }
    }
}

} // namespace ninfer::ops::detail
#include "ops/linear/w8/w8_launch.h"

#include "core/device.h"
#include "ops/common/math.cuh"
#include "ops/common/memory.cuh"
#include "ops/common/warp.cuh"

#include <cuda_bf16.h>
#include <cuda_fp16.h>

#include <cstdint>
#include <stdexcept>

namespace ninfer::ops::detail {
namespace {

constexpr int kRows          = 2048;
constexpr int kHidden        = 16384;
constexpr int kGroups        = kHidden / 32;
constexpr int kValuesPerLane = 8;

template <int RowsPerCta>
__global__ __launch_bounds__(RowsPerCta * 32, 2) void w8_rowsplit_k16384_decode_kernel(
    const __nv_bfloat16* __restrict__ x, const std::uint8_t* __restrict__ codes,
    const std::uint8_t* __restrict__ scales, __nv_bfloat16* __restrict__ out) {
    constexpr int kValuesPerPhase = 32 * kValuesPerLane;
    constexpr int kGroupsPerPhase = kValuesPerPhase / 32;
    constexpr int kPhases         = kHidden / kValuesPerPhase;
    constexpr unsigned kMask      = 0xffffffffu;

    const int lane       = static_cast<int>(threadIdx.x) & 31;
    const int warp       = static_cast<int>(threadIdx.x) >> 5;
    const int row        = static_cast<int>(blockIdx.x) * RowsPerCta + warp;
    const auto* code_row = codes + static_cast<std::int64_t>(row) * kHidden;
    const auto* scale_row =
        scales + static_cast<std::int64_t>(row) * kGroups * sizeof(std::uint16_t);

    float acc = 0.0f;
#pragma unroll 8
    for (int phase = 0; phase < kPhases; ++phase) {
        unsigned scale_bits = 0;
        if (lane < kGroupsPerPhase) {
            scale_bits = *reinterpret_cast<const std::uint16_t*>(
                scale_row + static_cast<std::int64_t>(phase * kGroupsPerPhase + lane) * 2);
        }
        scale_bits        = __shfl_sync(kMask, scale_bits, lane >> 2);
        const float scale = __half2float(__ushort_as_half(scale_bits));

        const int phase_k  = phase * kValuesPerPhase + lane * kValuesPerLane;
        const uint2 packed = load_vec<uint2>(code_row + phase_k);
        const uint4 values = load_vec<uint4>(x + phase_k);
        const float2 xv[4] = {
            bf16x2_bits_to_float2(values.x),
            bf16x2_bits_to_float2(values.y),
            bf16x2_bits_to_float2(values.z),
            bf16x2_bits_to_float2(values.w),
        };
#pragma unroll
        for (int word_index = 0; word_index < 2; ++word_index) {
            const std::uint32_t word = (&packed.x)[word_index];
#pragma unroll
            for (int byte = 0; byte < 4; ++byte) {
                const float2 activation_pair = xv[word_index * 2 + (byte >> 1)];
