    const std::int64_t base = row * d;

    LayerNormMoments local{0.0f, 0.0f, 0};
    for (std::int32_t i = static_cast<std::int32_t>(threadIdx.x); i < d; i += Block) {
        const float value = __bfloat162float(x[base + i]);
        ++local.count;
        const float delta = value - local.mean;
        local.mean += delta / static_cast<float>(local.count);
        local.m2 += delta * (value - local.mean);
    }

    __shared__ float means[Block];
    __shared__ float m2s[Block];
    __shared__ int counts[Block];
    means[threadIdx.x]  = local.mean;
    m2s[threadIdx.x]    = local.m2;
    counts[threadIdx.x] = local.count;
    __syncthreads();

    for (int offset = Block / 2; offset > 0; offset >>= 1) {
        if (static_cast<int>(threadIdx.x) < offset) {
            const LayerNormMoments a{means[threadIdx.x], m2s[threadIdx.x], counts[threadIdx.x]};
            const LayerNormMoments b{means[threadIdx.x + offset], m2s[threadIdx.x + offset],
                                     counts[threadIdx.x + offset]};
            const LayerNormMoments merged = layer_norm_merge(a, b);
            means[threadIdx.x]            = merged.mean;
            m2s[threadIdx.x]              = merged.m2;
            counts[threadIdx.x]           = merged.count;
        }
        __syncthreads();
    }

    const float mean = means[0];
    const float inv  = rsqrtf(m2s[0] / static_cast<float>(d) + eps);
    for (std::int32_t i = static_cast<std::int32_t>(threadIdx.x); i < d; i += Block) {
        const float value = (__bfloat162float(x[base + i]) - mean) * inv;
        out[base + i] =
            __float2bfloat16_rn(value * __bfloat162float(weight[i]) + __bfloat162float(bias[i]));
    }
}

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

#include "ops/launcher/argmax.h" // detail::argmax_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("argmax: ") + 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("argmax: tensor size overflows int64");
        }
        total *= t.ne[d];
    }
    return total;
}

} // namespace

void argmax(const Tensor& logits, Tensor& out, std::int32_t valid_rows, cudaStream_t stream) {
    if (logits.dtype != DType::BF16) { throw std::invalid_argument("argmax: logits must be BF16"); }
    if (out.dtype != DType::I32) { throw std::invalid_argument("argmax: out must be I32"); }

    const std::int64_t logits_n = numel_allow_zero(logits, "logits");
    (void)numel_allow_zero(out, "out");

    if (logits.ne[2] != 1 || logits.ne[3] != 1) {
        throw std::invalid_argument("argmax: logits must be rank-2 [vocab,T]");
    }
    if (out.ne[1] != 1 || out.ne[2] != 1 || out.ne[3] != 1) {
        throw std::invalid_argument("argmax: out must be rank-1 [T]");
    }
    if (logits.ne[0] <= 0) {
        throw std::invalid_argument("argmax: physical rows must be positive");
    }
    if (valid_rows <= 0 || valid_rows > logits.ne[0]) {
        throw std::invalid_argument("argmax: valid_rows must be in [1, logits.ne[0]]");
    }
    if (out.ne[0] != logits.ne[1]) {
        throw std::invalid_argument("argmax: out shape must be [logits.ne[1]]");
    }
    if (logits_n == 0) { return; }

    if (!logits.is_contiguous() || !out.is_contiguous()) {
        throw std::invalid_argument("argmax: logits/out must be contiguous");
    }
    if (logits.data == nullptr || out.data == nullptr) {
        throw std::invalid_argument("argmax: logits/out data must be non-null");
    }

    detail::argmax_launch(logits, out, valid_rows, stream);
}

} // namespace ninfer::ops
#pragma once

#include "core/tensor.h"
#include "ops/gdn_input_proj/gdn_conv.cuh"

#include <cuda_bf16.h>

#include <cstdint>

namespace ninfer::ops::detail {

inline constexpr std::int32_t kGdnQueryRows  = 2048;
inline constexpr std::int32_t kGdnKeyRows    = 2048;
inline constexpr std::int32_t kGdnValueRows  = 6144;
inline constexpr std::int32_t kGdnChannels   = kGdnQueryRows + kGdnKeyRows + kGdnValueRows;
inline constexpr std::int32_t kGdnZRows      = 6144;
inline constexpr std::int32_t kGdnParentRows = kGdnChannels + kGdnZRows;

// Format-neutral output policy for a single-parent [query,key,value,z] contraction. The
// contraction owns the represented projected values; this adapter owns only the final GDN
// convolution/state semantics and the physical split between Q/K/V and Z.
template <int Tokens, class Publish>
struct GdnConvOutput {
    GdnConvEpilogue<Publish> conv;
    __nv_bfloat16* z;

    __device__ __forceinline__ void store_row(std::int32_t parent_row,
                                              const float (&projected)[Tokens]) const {
        if (parent_row < kGdnChannels) {
            conv.store(parent_row, projected);
            return;
        }
#pragma unroll
        for (int token = 0; token < Tokens; ++token) {
            z[static_cast<std::int64_t>(token) * kGdnZRows + parent_row - kGdnChannels] =
                __float2bfloat16_rn(projected[token]);
        }
    }

    __device__ __forceinline__ void store(std::int32_t parent_row, std::int32_t token,
                                          float projected) const {
        static_assert(Tokens == 1);
        (void)token;
        const float row[1]{projected};
        store_row(parent_row, row);
    }
};

template <int Tokens, class Publish>
GdnConvOutput<Tokens, Publish>
make_gdn_conv_output(const Tensor& conv_weight, const Tensor& conv_states,
                     const Tensor& valid_columns, const Tensor& initial_slot, Tensor& query,
                     Tensor& key, Tensor& value, Tensor& z, Publish publish) {
    return {
        {
            static_cast<const __nv_bfloat16*>(conv_weight.data),
            static_cast<const __nv_bfloat16*>(conv_states.data),
            static_cast<const std::int32_t*>(initial_slot.data),
            valid_columns.data == nullptr ? nullptr
                                          : static_cast<const std::int32_t*>(valid_columns.data),
            static_cast<__nv_bfloat16*>(query.data),
            static_cast<__nv_bfloat16*>(key.data),
            static_cast<__nv_bfloat16*>(value.data),
            kGdnChannels,
            kGdnQueryRows,
            kGdnKeyRows,
            kGdnValueRows,
            0,
            Tokens,
            0,
            publish,
        },
        static_cast<__nv_bfloat16*>(z.data),
    };
}

static_assert(kGdnParentRows == 16384);

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

#include "core/gdn_replay_records.h"
#include "core/linear_attention_state.h"

#include <cuda_runtime.h>

#include <cstdint>
#include <span>

namespace ninfer::ops {

struct GdnReplayFoldRow {
    std::int32_t source_state_slot;
    std::int32_t destination_state_slot;
    std::int32_t commit_columns;
};

/**
 * Prepared Op: gdn_replay_fold
 *
 * Replays each row's [0,commit_columns) record prefix across every registered GDN layer and
 * updates the caller-selected absolute linear-attention destination state slot. rows[b] always maps
 * to physical record row b; rows are not filtered, compressed, or reordered. The active row count
 * is rows.size() and must be in [1,records.spec.record_capacity].
 *
 * source_state_slot and destination_state_slot are in [0,states.spec.slot_count). A row may be
 * in-place. Destinations are distinct and cannot overwrite another active row's source.
 * source_state_slot is the absolute slot used to produce that row's records. commit_columns is in
 * [0,T]. Zero is a
 * strict no-op for the row: no record or state is read and neither recurrent state nor convolution
 * history is written. For a positive extent, the Op consumes raw key/value/{g,beta} records in
 * order, writes the final FP32 recurrent state, and sets convolution history to
 * tail_3(old_history || conv_record[0:commit_columns]).
 *
 * Construction validates the immutable record/state geometry, layout, address ranges and
 * disjointness once. execute() validates only active rows, state selectors and commit extents.
 * The bound record/state addresses must remain stable for the plan lifetime.
 *
 * The Op admits the two registered all-layer geometries only, owns no workspace or device
 * allocation, and does not read query or generate token output. The four record planes are
 * read-only, disjoint, and do not overlap either state region.
 */
class GdnReplayFoldPlan {
public:
    GdnReplayFoldPlan(const GdnReplayRecords& records, LinearAttentionStateAllLayersView states);

    void execute(std::span<const GdnReplayFoldRow> rows, cudaStream_t stream) const;

private:
    GdnReplayRecords records_;
    LinearAttentionStateAllLayersView states_;
};

} // namespace ninfer::ops
#pragma once

// Signed int8, per-token G64 KV-cache codec shared by append and causal-attention kernels.
// This header owns index math, vectorized decode, and scalar encode; there is deliberately no
// standalone transcode kernel in the production path.

#include "ops/common/math.cuh"
#include "ops/common/memory.cuh"
#include "ops/kernel/paged_kv_address.cuh"
#include "ops/kv_cache/hadamard_d256.cuh"

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

#include <cstdint>

namespace ninfer::ops {

inline constexpr int kKVCacheInt8HeadDim = 256;
inline constexpr int kKVCacheInt8Group   = 64;
inline constexpr int kKVCacheInt8Groups  = kKVCacheInt8HeadDim / kKVCacheInt8Group;

template <typename Geometry>
__device__ __forceinline__ std::int64_t
kv_cache_int8_quant_code_index(int physical_page, int kv_head, int d, int page_offset) {
    return paged_kv_element_offset<kKVCacheInt8HeadDim, Geometry::KVHeads>(physical_page, kv_head,
                                                                           page_offset, d);
}

template <typename Geometry>
__device__ __forceinline__ std::int64_t
kv_cache_int8_quant_scale_index(int physical_page, int kv_head, int group, int page_offset) {
    return paged_kv_element_offset<kKVCacheInt8Groups, Geometry::KVHeads>(physical_page, kv_head,
                                                                          page_offset, group);
}

template <typename Geometry>
__device__ __forceinline__ std::int64_t kv_cache_int8_quant_src_index(int kv_head, int d,
                                                                      int token) {
    return static_cast<std::int64_t>(d) +
           static_cast<std::int64_t>(kKVCacheInt8HeadDim) *
               (static_cast<std::int64_t>(kv_head) +
                static_cast<std::int64_t>(Geometry::KVHeads) * token);
}

struct KVCacheInt8QuantParams {
    __half scale;
    float inverse_scale;
};

// Exact persistent group-scale boundary shared by standalone and fused append. The stored scale is
// FP16-RNE(absmax/127); codes always use the reciprocal of that represented FP16 value.
__device__ __forceinline__ KVCacheInt8QuantParams kv_cache_int8_quant_params(float absmax) {
    const __half scale            = __float2half_rn(absmax > 0.0f ? absmax / 127.0f : 0.0f);
    const float represented_scale = __half2float(scale);
    return {
        .scale         = scale,
        .inverse_scale = represented_scale > 0.0f ? 1.0f / represented_scale : 0.0f,
    };
}

__device__ __forceinline__ std::int8_t kv_cache_int8_quant_code(float x, float inv_scale) {
    if (inv_scale == 0.0f) { return static_cast<std::int8_t>(0); }
    int q = __float2int_rn(x * inv_scale);
    q     = max(-127, min(127, q));
    return static_cast<std::int8_t>(q);
}

__device__ __forceinline__ int4 kv_cache_int8_dequant_i8x8_from(const std::int8_t* codes8,
                                                                float s) {
    const int2 raw       = load_vec<int2>(codes8);
    const std::int8_t* c = reinterpret_cast<const std::int8_t*>(&raw);
    unsigned packed[4];
#pragma unroll
    for (int i = 0; i < 4; ++i) {
        const float x0 = static_cast<float>(c[2 * i]) * s;
        const float x1 = static_cast<float>(c[2 * i + 1]) * s;
        packed[i]      = pack_bf16x2(x0, x1);
    }
    return make_int4(static_cast<int>(packed[0]), static_cast<int>(packed[1]),
                     static_cast<int>(packed[2]), static_cast<int>(packed[3]));
}

} // namespace ninfer::ops
#pragma once

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

#include <cuda_bf16.h>
#include <math_constants.h>

#include <cstdint>

namespace ninfer::ops {

inline constexpr int kContextQueryHeadDim  = 128;
inline constexpr int kContextQueryQHeads   = 32;
inline constexpr int kContextQueryKVHeads  = 8;
inline constexpr int kContextQueryGroup    = 4;
inline constexpr int kContextQueryMaxSplit = 85;

__device__ __forceinline__ int context_query_swz(int row, int col) {
    return (((col >> 3) ^ (row & 7)) << 3) | (col & 7);
}

__device__ __forceinline__ unsigned context_query_swz_addr(unsigned lane_base, unsigned ck,
                                                           unsigned as, unsigned r) {
    return lane_base + ((ck | as) ^ r);
}

__device__ __forceinline__ std::int64_t context_query_q_index(int q_head, int d, int token) {
    return static_cast<std::int64_t>(d) +
           static_cast<std::int64_t>(kContextQueryHeadDim) *
               (static_cast<std::int64_t>(q_head) +
                static_cast<std::int64_t>(kContextQueryQHeads) * token);
}

__device__ __forceinline__ std::int64_t context_query_query_kv_index(int kv_head, int d,
                                                                     int token) {
    return static_cast<std::int64_t>(d) +
           static_cast<std::int64_t>(kContextQueryHeadDim) *
               (static_cast<std::int64_t>(kv_head) +
                static_cast<std::int64_t>(kContextQueryKVHeads) * token);
}

template <int Tokens>
__device__ __forceinline__ std::int64_t context_query_partial_index(int q_head, int d, int token,
                                                                    int split) {
    return static_cast<std::int64_t>(d) +
           static_cast<std::int64_t>(kContextQueryHeadDim) *
               (static_cast<std::int64_t>(q_head) +
                static_cast<std::int64_t>(kContextQueryQHeads) *
                    (static_cast<std::int64_t>(token) + static_cast<std::int64_t>(Tokens) * split));
}

template <int Tokens>
__device__ __forceinline__ std::int64_t context_query_stat_index(int q_head, int token, int split) {
    return static_cast<std::int64_t>(q_head) +
           static_cast<std::int64_t>(kContextQueryQHeads) *
               (static_cast<std::int64_t>(token) + static_cast<std::int64_t>(Tokens) * split);
}

__device__ __forceinline__ void context_query_row_to_qt(int row, int kv_head, int& q_head,
                                                        int& token) {
    token             = row / kContextQueryGroup;
    const int q_local = row - token * kContextQueryGroup;
    q_head            = kv_head * kContextQueryGroup + q_local;
}

template <typename ContextPolicy, int KeyBlock, int Threads>
__device__ __forceinline__ void
context_query_stage_tile(__nv_bfloat16* dst, const __nv_bfloat16* context,
                         const __nv_bfloat16* query, int key0, int valid_keys, bool query_tile,
                         int kv_head, int physical_page, ContextPolicy& policy, int tid) {
    constexpr int VecsPerRow        = kContextQueryHeadDim / 8;
    const std::int64_t context_tile = policy.context_tile(kv_head, key0, physical_page);
    for (int chunk = tid; chunk < KeyBlock * VecsPerRow; chunk += Threads) {
        const int row      = chunk / VecsPerRow;
        const int d        = (chunk - row * VecsPerRow) * 8;
        const bool live    = row < valid_keys;
        const int safe_row = live ? row : 0;
        const std::int64_t src_index =
            query_tile ? context_query_query_kv_index(kv_head, d, safe_row)
                       : policy.context_index(context_tile, d, live ? key0 + row : 0, safe_row);
        const __nv_bfloat16* src = query_tile ? query + src_index : context + src_index;
        __nv_bfloat16* smem      = &dst[row * kContextQueryHeadDim + context_query_swz(row, d)];
        cp_async_zfill<16, Cache::cg>(smem, src, live ? 16 : 0);
    }
}

template <typename ContextPolicy, int Tokens, int WarpsPerCta, int KeyBlock, bool DirectOutput>
__device__ __forceinline__ void context_query_split_partial_body(
    const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ query_k,
    const __nv_bfloat16* __restrict__ query_v, const std::int32_t* __restrict__ valid_columns,
    const __nv_bfloat16* __restrict__ context_k, const __nv_bfloat16* __restrict__ context_v,
    ContextPolicy policy, int length, int max_context, int split_capacity, float scale,
    __nv_bfloat16* __restrict__ partial_acc, float* __restrict__ partial_m,
    float* __restrict__ partial_l, __nv_bfloat16* __restrict__ out) {
    static_assert(Tokens >= 1 && Tokens <= 16);
    static_assert(WarpsPerCta == (Tokens + 3) / 4);
    static_assert(KeyBlock == 32 || KeyBlock == 64);

    constexpr int D             = kContextQueryHeadDim;
    constexpr int Wc            = WarpsPerCta;
    constexpr int Threads       = Wc * 32;
    constexpr int Br            = Wc * 16;
    constexpr int RowCount      = Tokens * kContextQueryGroup;
    constexpr int QKNt          = KeyBlock / 8;
    constexpr int QKKs          = D / 16;
    constexpr int PVNt          = D / 8;
    constexpr int PVKs          = KeyBlock / 16;
    constexpr int RowBytes      = D * static_cast<int>(sizeof(__nv_bfloat16));
    constexpr float Log2E       = 1.4426950408889634074f;
    constexpr unsigned FullMask = 0xffffffffu;

    static_assert(RowCount <= Br);
    static_assert(Br <= 2 * KeyBlock);
    const int kv_head = static_cast<int>(blockIdx.x);
    const int split   = static_cast<int>(blockIdx.y);
    const int batch   = static_cast<int>(blockIdx.z);
    const int tid     = static_cast<int>(threadIdx.x);
    const int warp    = tid >> 5;
    const int lane    = tid & 31;

    constexpr std::int64_t QueryElements =
        static_cast<std::int64_t>(D) * kContextQueryQHeads * Tokens;
    constexpr std::int64_t QueryKvElements =
        static_cast<std::int64_t>(D) * kContextQueryKVHeads * Tokens;
    constexpr std::int64_t PartialElements = QueryElements;
    constexpr std::int64_t StatElements = static_cast<std::int64_t>(kContextQueryQHeads) * Tokens;
    q += QueryElements * batch;
    query_k += QueryKvElements * batch;
    query_v += QueryKvElements * batch;
    out += QueryElements * batch;
    partial_acc += PartialElements * split_capacity * batch;
    partial_m += StatElements * split_capacity * batch;
    partial_l += StatElements * split_capacity * batch;
    const int valid = valid_columns[batch];
    if (kv_head >= kContextQueryKVHeads || split >= split_capacity || length < 0 ||
        length > max_context || valid < 0 || valid > Tokens) {
        return;
    }
    if (valid == 0) {
        if constexpr (DirectOutput) {
            constexpr int VecsPerKVHead = D * kContextQueryGroup / 8;
            constexpr int OutputVecs    = Tokens * VecsPerKVHead;
            constexpr int TokenStride   = D * kContextQueryQHeads;
            for (int unit = tid; unit < OutputVecs; unit += Threads) {
                const int token  = unit / VecsPerKVHead;
                const int vector = unit - token * VecsPerKVHead;
                const int offset =
                    token * TokenStride + kv_head * (D * kContextQueryGroup) + vector * 8;
                store_vec(&out[offset], make_int4(0, 0, 0, 0));
            }
        }
        // Split routes delegate physical-tail zeroing to the reduce kernel.
        return;
    }

    const int context_count = policy.context_count(length);
    const int context_start = length - context_count;
    const int context_tiles = (context_count + KeyBlock - 1) / KeyBlock;
    const int active_splits = context_tiles > 0 ? min(context_tiles, split_capacity) : 1;
    if (split >= active_splits) { return; }

    const int tile_begin =
        static_cast<int>((static_cast<std::int64_t>(context_tiles) * split) / active_splits);
    const int tile_end =
        static_cast<int>((static_cast<std::int64_t>(context_tiles) * (split + 1)) / active_splits);
    const bool owns_query        = split == active_splits - 1;
    const int context_tile_count = tile_end - tile_begin;
    const int iterations         = context_tile_count + (owns_query ? 1 : 0);

    policy.template prime<Tokens, KeyBlock>(context_tile_count, context_start, tile_begin, lane);

    extern __shared__ __align__(16) __nv_bfloat16 shared[];
    __nv_bfloat16* k_s = shared;
    __nv_bfloat16* v_s = shared + KeyBlock * D;

    // The two K/V buffers together hold at least Br rows. Use them once as Q staging, then retain
    // all Q MMA fragments in registers for the complete split.
    for (int chunk = tid; chunk < Br * (D / 8); chunk += Threads) {
        const int row = chunk / (D / 8);
        const int d   = (chunk - row * (D / 8)) * 8;
        int q_head = 0, token = 0;
        context_query_row_to_qt(row, kv_head, q_head, token);
        const bool live = row < RowCount && token < valid;
        const __nv_bfloat16* src =
            q + context_query_q_index(live ? q_head : 0, d, live ? token : 0);
        __nv_bfloat16* dst = &shared[row * D + context_query_swz(row, d)];
        cp_async_zfill<16, Cache::cg>(dst, src, live ? 16 : 0);
    }
    cp_commit();
    cp_wait<0>();
    __syncthreads();

    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 a_coloff = (a_mat >> 1) << 3;
    const int b_rin    = lane & 7;
    const int b_koff   = ((lane >> 3) & 1) << 3;

    const int warp_row0   = warp * 16;
    const int row0        = warp_row0 + gid;
    const int row1        = row0 + 8;
    const int q_position0 = policy.query_position(row0 < RowCount ? row0 / kContextQueryGroup : 0);
    const int q_position1 = policy.query_position(row1 < RowCount ? row1 / kContextQueryGroup : 0);
    unsigned af_q[QKKs][4];
#pragma unroll
    for (int ks = 0; ks < QKKs; ++ks) {
        const int row = warp_row0 + a_rowoff;
        const int col = ks * 16 + a_coloff;
        ldmatrix_x4(af_q[ks][0], af_q[ks][1], af_q[ks][2], af_q[ks][3],
                    smem_addr(&shared[row * D + context_query_swz(row, col)]));
    }
    __syncthreads();

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

    const unsigned v_sbase     = smem_addr(v_s);
    const unsigned v_lane_base = v_sbase + static_cast<unsigned>(((lane >> 3) & 1) * 8 * RowBytes) +
                                 static_cast<unsigned>(b_rin * RowBytes);
    const unsigned v_as = static_cast<unsigned>((lane >> 4) << 4);
    const unsigned v_r  = static_cast<unsigned>(b_rin << 4);

    auto tile_metadata = [&](int iteration, bool& is_query, int& key0, int& valid_keys) {
        is_query = iteration >= context_tile_count;
        if (is_query) {
            key0       = 0;
            valid_keys = valid;
        } else {
            key0       = context_start + (tile_begin + iteration) * KeyBlock;
            valid_keys = min(KeyBlock, length - key0);
        }
    };
    bool current_is_query = false;
    int current_key0      = 0;
    int current_valid     = 0;
    tile_metadata(0, current_is_query, current_key0, current_valid);
    int current_page = 0;
    if constexpr (ContextPolicy::PageMapped) {
        current_page =
            current_is_query ? 0 : policy.template page<Tokens>(current_key0, lane, FullMask);
    }
    context_query_stage_tile<ContextPolicy, KeyBlock, Threads>(
        k_s, context_k, query_k, current_key0, current_valid, current_is_query, kv_head,
        current_page, policy, tid);
    cp_commit();

    for (int iteration = 0; iteration < iterations; ++iteration) {
        cp_wait<0>();
        __syncthreads();

        context_query_stage_tile<ContextPolicy, KeyBlock, Threads>(
            v_s, context_v, query_v, current_key0, current_valid, current_is_query, kv_head,
            current_page, policy, tid);
        cp_commit();

        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;
#pragma unroll
            for (int ks = 0; ks < QKKs; ++ks) {
                unsigned bf[2];
                const int brow = nt * 8 + b_rin;
                const int bcol = ks * 16 + b_koff;
                ldmatrix_x2(bf[0], bf[1],
                            smem_addr(&k_s[brow * D + context_query_swz(brow, bcol)]));
                mma_bf16(score[nt][0], score[nt][1], score[nt][2], score[nt][3], af_q[ks][0],
                         af_q[ks][1], af_q[ks][2], af_q[ks][3], bf[0], bf[1]);
            }
        }

        cp_wait<0>();
        __syncthreads();

        bool next_is_query = false;
        int next_key0      = 0;
        int next_valid     = 0;
        int next_page      = 0;
        if (iteration + 1 < iterations) {
            tile_metadata(iteration + 1, next_is_query, next_key0, next_valid);
            if constexpr (ContextPolicy::PageMapped) {
                next_page = next_is_query
                                ? 0
                                : (!current_is_query && (next_key0 >> 6) == (current_key0 >> 6)
                                       ? current_page
                                       : policy.template page<Tokens>(next_key0, lane, FullMask));
            }
            context_query_stage_tile<ContextPolicy, KeyBlock, Threads>(
                k_s, context_k, query_k, next_key0, next_valid, next_is_query, kv_head, next_page,
                policy, tid);
            cp_commit();
        }

        float block_m0 = -CUDART_INF_F;
        float block_m1 = -CUDART_INF_F;
#pragma unroll
        for (int nt = 0; nt < QKNt; ++nt) {
            const int col0       = nt * 8 + 2 * lid;
            const int col1       = col0 + 1;
            const bool row0_live = row0 < RowCount && row0 / kContextQueryGroup < valid;
            const bool row1_live = row1 < RowCount && row1 / kContextQueryGroup < valid;
            const bool allow00 =
                row0_live && col0 < current_valid &&
                (current_is_query || policy.allow_context(q_position0, current_key0 + col0));
            const bool allow01 =
                row0_live && col1 < current_valid &&
                (current_is_query || policy.allow_context(q_position0, current_key0 + col1));
            const bool allow10 =
                row1_live && col0 < current_valid &&
                (current_is_query || policy.allow_context(q_position1, current_key0 + col0));
            const bool allow11 =
                row1_live && col1 < current_valid &&
                (current_is_query || policy.allow_context(q_position1, current_key0 + col1));
            score[nt][0] = allow00 ? score[nt][0] * scale : -CUDART_INF_F;
            score[nt][1] = allow01 ? score[nt][1] * scale : -CUDART_INF_F;
            score[nt][2] = allow10 ? score[nt][2] * scale : -CUDART_INF_F;
            score[nt][3] = allow11 ? score[nt][3] * scale : -CUDART_INF_F;
            block_m0     = fmaxf(block_m0, fmaxf(score[nt][0], score[nt][1]));
            block_m1     = fmaxf(block_m1, fmaxf(score[nt][2], score[nt][3]));
        }
        block_m0 = warp_max<4>(block_m0, FullMask);
        block_m1 = warp_max<4>(block_m1, FullMask);

        const float next_m0 = fmaxf(m0, block_m0);
        const float next_m1 = fmaxf(m1, block_m1);
        const float alpha0  = m0 == -CUDART_INF_F ? 0.0f : exp2_approx((m0 - next_m0) * Log2E);
        const float alpha1  = m1 == -CUDART_INF_F ? 0.0f : exp2_approx((m1 - next_m1) * Log2E);

        unsigned p_frag[PVKs][4];
        float block_l0 = 0.0f;
        float block_l1 = 0.0f;
#pragma unroll
        for (int nt = 0; nt < QKNt; ++nt) {
            const float p00 =
                score[nt][0] > -CUDART_INF_F ? exp2_approx((score[nt][0] - next_m0) * Log2E) : 0.0f;
            const float p01 =
                score[nt][1] > -CUDART_INF_F ? exp2_approx((score[nt][1] - next_m0) * Log2E) : 0.0f;
            const float p10 =
                score[nt][2] > -CUDART_INF_F ? exp2_approx((score[nt][2] - next_m1) * Log2E) : 0.0f;
            const float p11 =
                score[nt][3] > -CUDART_INF_F ? exp2_approx((score[nt][3] - next_m1) * Log2E) : 0.0f;
            block_l0 += p00 + p01;
            block_l1 += 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);
            }
        }
        block_l0 = warp_sum<4>(block_l0, FullMask);
        block_l1 = warp_sum<4>(block_l1, FullMask);

        l0 = l0 * alpha0 + block_l0;
        l1 = l1 * alpha1 + block_l1;
        m0 = next_m0;
        m1 = next_m1;
#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;
        }

        constexpr int PVTilePairs = (PVNt + 1) / 2;
        constexpr int PVLoads     = PVKs * PVTilePairs;
        unsigned vf[2][4];
        ldmatrix_x4_t(vf[0][0], vf[0][1], vf[0][2], vf[0][3],
                      context_query_swz_addr(v_lane_base, 0u, v_as, v_r));
#pragma unroll
        for (int load = 0; load < PVLoads; ++load) {
            const int pk   = load / PVTilePairs;
            const int n2   = (load % PVTilePairs) * 2;
            const int cur  = load & 1;
            const int next = cur ^ 1;
            if (load + 1 < PVLoads) {
                const int next_pk = (load + 1) / PVTilePairs;
                const int next_n2 = ((load + 1) % PVTilePairs) * 2;
                ldmatrix_x4_t(vf[next][0], vf[next][1], vf[next][2], vf[next][3],
                              context_query_swz_addr(
                                  v_lane_base + static_cast<unsigned>(next_pk * 16 * RowBytes),
                                  static_cast<unsigned>(next_n2 << 4), v_as, v_r));
            }
            mma_bf16(acc[n2][0], acc[n2][1], acc[n2][2], acc[n2][3], p_frag[pk][0], p_frag[pk][1],
                     p_frag[pk][2], p_frag[pk][3], vf[cur][0], vf[cur][1]);
            if (n2 + 1 < PVNt) {
                mma_bf16(acc[n2 + 1][0], acc[n2 + 1][1], acc[n2 + 1][2], acc[n2 + 1][3],
                         p_frag[pk][0], p_frag[pk][1], p_frag[pk][2], p_frag[pk][3], vf[cur][2],
                         vf[cur][3]);
            }
        }

        current_is_query = next_is_query;
        current_key0     = next_key0;
        current_valid    = next_valid;
        current_page     = next_page;
    }

    if constexpr (!DirectOutput) {
        if (lid == 0) {
            const int row0 = warp_row0 + gid;
            const int row1 = row0 + 8;
            if (row0 < RowCount) {
                int q_head = 0, token = 0;
                context_query_row_to_qt(row0, kv_head, q_head, token);
                partial_m[context_query_stat_index<Tokens>(q_head, token, split)] = m0;
                partial_l[context_query_stat_index<Tokens>(q_head, token, split)] = l0;
            }
            if (row1 < RowCount) {
                int q_head = 0, token = 0;
                context_query_row_to_qt(row1, kv_head, q_head, token);
                partial_m[context_query_stat_index<Tokens>(q_head, token, split)] = m1;
                partial_l[context_query_stat_index<Tokens>(q_head, token, split)] = l1;
            }
        }
    }

#pragma unroll
    for (int n = 0; n < PVNt; ++n) {
        const int d0   = n * 8 + 2 * lid;
        const int row0 = warp_row0 + gid;
        const int row1 = row0 + 8;
        if (row0 < RowCount) {
            int q_head = 0, token = 0;
            context_query_row_to_qt(row0, kv_head, q_head, token);
            if constexpr (DirectOutput) {
                const float inv_l = l0 > 0.0f ? 1.0f / l0 : 0.0f;
                const auto dst    = context_query_q_index(q_head, d0, token);
                store_vec(&out[dst], pack_bf16x2(acc[n][0] * inv_l, acc[n][1] * inv_l));
            } else {
                const auto dst = context_query_partial_index<Tokens>(q_head, d0, token, split);
                store_vec(&partial_acc[dst], pack_bf16x2(acc[n][0], acc[n][1]));
            }
        }
        if (row1 < RowCount) {
            int q_head = 0, token = 0;
            context_query_row_to_qt(row1, kv_head, q_head, token);
            if constexpr (DirectOutput) {
                const float inv_l = l1 > 0.0f ? 1.0f / l1 : 0.0f;
                const auto dst    = context_query_q_index(q_head, d0, token);
                store_vec(&out[dst], pack_bf16x2(acc[n][2] * inv_l, acc[n][3] * inv_l));
            } else {
                const auto dst = context_query_partial_index<Tokens>(q_head, d0, token, split);
                store_vec(&partial_acc[dst], pack_bf16x2(acc[n][2], acc[n][3]));
            }
        }
    }
}

template <int Tokens, int KeyBlock>
__device__ __forceinline__ void
context_query_reduce_body(const __nv_bfloat16* __restrict__ partial_acc,
                          const float* __restrict__ partial_m, const float* __restrict__ partial_l,
                          int length, int context_count,
                          const std::int32_t* __restrict__ valid_columns, int max_context,
                          int split_capacity, __nv_bfloat16* __restrict__ out) {
    const int q_head = static_cast<int>(blockIdx.x);
    const int token  = static_cast<int>(blockIdx.y);
    const int batch  = static_cast<int>(blockIdx.z);
    const int tid    = static_cast<int>(threadIdx.x);
    constexpr std::int64_t QueryElements =
        static_cast<std::int64_t>(kContextQueryHeadDim) * kContextQueryQHeads * Tokens;
    constexpr std::int64_t StatElements = static_cast<std::int64_t>(kContextQueryQHeads) * Tokens;
    partial_acc += QueryElements * split_capacity * batch;
    partial_m += StatElements * split_capacity * batch;
    partial_l += StatElements * split_capacity * batch;
    out += QueryElements * batch;
    if (q_head >= kContextQueryQHeads || token >= Tokens) { return; }
    if (length < 0 || length > max_context || token >= valid_columns[batch]) {
        if (tid < kContextQueryHeadDim) {
            out[context_query_q_index(q_head, tid, token)] = __float2bfloat16(0.0f);
        }
        return;
    }

    const int context_tiles = (context_count + KeyBlock - 1) / KeyBlock;
    const int active_splits = context_tiles > 0 ? min(context_tiles, split_capacity) : 1;
    __shared__ float reduce[128];

    float local_m = -CUDART_INF_F;
    for (int split = tid; split < active_splits; split += blockDim.x) {
        local_m = fmaxf(local_m, partial_m[context_query_stat_index<Tokens>(q_head, token, split)]);
    }
    reduce[tid] = local_m;
    __syncthreads();
    for (int stride = 64; stride > 0; stride >>= 1) {
        if (tid < stride) { reduce[tid] = fmaxf(reduce[tid], reduce[tid + stride]); }
        __syncthreads();
    }
    const float global_m = reduce[0];
    __syncthreads();

    float local_l = 0.0f;
    for (int split = tid; split < active_splits; split += blockDim.x) {
        const auto idx = context_query_stat_index<Tokens>(q_head, token, split);
        local_l += partial_l[idx] * expf(partial_m[idx] - global_m);
    }
    reduce[tid] = local_l;
    __syncthreads();
    for (int stride = 64; stride > 0; stride >>= 1) {
        if (tid < stride) { reduce[tid] += reduce[tid + stride]; }
        __syncthreads();
    }
    const float global_l = reduce[0];

    if (tid < kContextQueryHeadDim) {
        float numerator = 0.0f;
        for (int split = 0; split < active_splits; ++split) {
            const auto stat    = context_query_stat_index<Tokens>(q_head, token, split);
            const float weight = expf(partial_m[stat] - global_m);
            numerator +=
                __bfloat162float(
                    partial_acc[context_query_partial_index<Tokens>(q_head, tid, token, split)]) *
                weight;
        }
        const float value = global_l > 0.0f ? numerator / global_l : 0.0f;
        out[context_query_q_index(q_head, tid, token)] = __float2bfloat16(value);
    }
}

} // namespace ninfer::ops
#pragma once

#include "runtime/contract/types.h"

#include <cstdint>
#include <cstdlib>

namespace ninfer::runtime {

class GenerationBudget {
public:
    GenerationBudget(std::uint32_t effective_tokens, FinishReason limit_reason)
        : remaining_(effective_tokens), limit_reason_(limit_reason) {
        if (limit_reason_ != FinishReason::OutputLimit &&
            limit_reason_ != FinishReason::ContextCapacity) {
            std::abort();
        }
    }

    [[nodiscard]] std::uint32_t remaining() const noexcept { return remaining_; }

    [[nodiscard]] FinishReason limit_reason() const noexcept { return limit_reason_; }

    void commit(std::uint32_t tokens) noexcept {
        if (tokens > remaining_) { std::abort(); }
        remaining_ -= tokens;
    }

private:
    std::uint32_t remaining_   = 0;
    FinishReason limit_reason_ = FinishReason::None;
};

} // namespace ninfer::runtime
#pragma once

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

#include <cuda_bf16.h>

#include <cstdint>

namespace ninfer::ops::detail {

template <int RowsPerBranch, int IntermediateRows>
struct Fp8SwiGluRows {
    static_assert(RowsPerBranch > 0 && (RowsPerBranch & (RowsPerBranch - 1)) == 0);

    __device__ __forceinline__ int weight_row(int row_begin, int local_row) const {
        return row_begin + (local_row & (RowsPerBranch - 1)) +
               (local_row >= RowsPerBranch ? IntermediateRows : 0);
    }
};

union Fp8SwiGluBf16Pair {
    unsigned bits;
    __nv_bfloat162 values;
};

struct Fp8SwiGluOutput {
    __nv_bfloat16* data;
    std::int32_t rows;

    __device__ __forceinline__ unsigned combine(unsigned gate_bits, unsigned up_bits) const {
        Fp8SwiGluBf16Pair gate{gate_bits};
        Fp8SwiGluBf16Pair up{up_bits};
        const float2 gate_values = __bfloat1622float2(gate.values);
        const float2 up_values   = __bfloat1622float2(up.values);
        Fp8SwiGluBf16Pair result;
        result.values = __floats2bfloat162_rn(silu(gate_values.x) * up_values.x,
                                              silu(gate_values.y) * up_values.y);
        return result.bits;
    }

    __device__ __forceinline__ void store_pair_vector(std::int32_t row, std::int32_t token,
                                                      uint4 gate, uint4 up) const {
        const uint4 values = make_uint4(combine(gate.x, up.x), combine(gate.y, up.y),
                                        combine(gate.z, up.z), combine(gate.w, up.w));
        store_vec(data + static_cast<std::int64_t>(token) * rows + row, values);
    }

    __device__ __forceinline__ void store_pair(std::int32_t row, std::int32_t token, float gate,
                                               float up) const {
        data[static_cast<std::int64_t>(token) * rows + row] = __float2bfloat16_rn(silu(gate) * up);
    }
};

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

#include <nlohmann/json.hpp>

#include <array>
#include <chrono>
#include <cstdio>
#include <random>
#include <string_view>
#include <utility>

namespace ninfer::serve {

namespace {

std::string chat_identifier(std::string_view prefix) {
    static thread_local std::mt19937_64 random{std::random_device{}()};
    std::uniform_int_distribution<std::uint64_t> distribution;
    std::array<char, 32> buffer{};
    std::snprintf(buffer.data(), buffer.size(), "%016llx",
                  static_cast<unsigned long long>(distribution(random)));
    return std::string(prefix) + buffer.data();
}

std::string responses_identifier(std::string_view prefix) {
    static thread_local std::mt19937_64 random{std::random_device{}()};
    std::uniform_int_distribution<std::uint64_t> distribution;
    std::array<char, 48> buffer{};
    std::snprintf(buffer.data(), buffer.size(), "%016llx%016llx",
                  static_cast<unsigned long long>(distribution(random)),
                  static_cast<unsigned long long>(distribution(random)));
    return std::string(prefix) + "_" + buffer.data();
}

} // namespace

using Json = nlohmann::json;

std::string make_models_list(const std::string& model_id, std::int64_t created,
                             std::uint32_t max_model_len) {
    // vLLM/llama.cpp-compatible discovery metadata for the configured per-request context limit.
    const Json payload = {{"object", "list"},
                          {"data", Json::array({Json{{"id", model_id},
                                                     {"object", "model"},
                                                     {"created", created},
                                                     {"owned_by", "ninfer"},
                                                     {"max_model_len", max_model_len}}})}};
    return payload.dump();
}

std::string make_model_object(const std::string& model_id, std::int64_t created,
                              std::uint32_t max_model_len) {
    // vLLM/llama.cpp-compatible discovery metadata for the configured per-request context limit.
    const Json payload = {{"id", model_id},
                          {"object", "model"},
                          {"created", created},
                          {"owned_by", "ninfer"},
                          {"max_model_len", max_model_len}};
    return payload.dump();
}

std::string make_error_body(const ApiError& error) {
    Json rendered     = {{"message", error.message}, {"type", error.type}};
    rendered["param"] = error.param.empty() ? Json(nullptr) : Json(error.param);
    rendered["code"]  = error.code.empty() ? Json(nullptr) : Json(error.code);
    return Json{{"error", std::move(rendered)}}.dump();
}

std::int64_t unix_time_now() {
    return std::chrono::duration_cast<std::chrono::seconds>(
               std::chrono::system_clock::now().time_since_epoch())
        .count();
}

void validate_openai_model(std::string_view requested, std::string_view available) {
    if (requested == available) { return; }
    ApiError error;
    error.status  = 404;
    error.type    = "invalid_request_error";
    error.param   = "model";
    error.code    = "model_not_found";
    error.message = "model '" + std::string(requested) + "' not found";
    throw ApiException(std::move(error));
}

std::string new_openai_chat_completion_id() { return chat_identifier("chatcmpl-"); }

std::string new_openai_chat_tool_call_id() { return chat_identifier("call_"); }

std::string new_openai_request_id() { return responses_identifier("req"); }

std::string new_openai_response_id() { return responses_identifier("resp"); }

std::string new_openai_response_item_id(std::string_view prefix) {
    return responses_identifier(prefix);
}

} // namespace ninfer::serve
#pragma once

#include "targets/qwen3_6_35b_a3b/impl/config.h"
#include "targets/qwen3_6_35b_a3b/impl/load/bindings.h"
#include <ninfer/targets/qwen3_6/runtime.h>

#include <cstddef>
#include <cstdint>
#include <vector>

namespace ninfer::targets::qwen3_6_35b_a3b::detail {

using GraphExecutionProfile = qwen3_6::GraphExecutionProfile;

struct Variant {
    using WeightsProfile                 = detail::WeightsProfile;
    using TextConfig                     = detail::TextConfig;
    using VisionConfig                   = detail::VisionConfig;
    using DFlashConfig                   = detail::DFlashConfig;
    using ModelView                      = detail::RuntimeModelView;
    using FullAttentionProjectionWeights = detail::AttentionProjectionPayload;
    using GdnProjectionWeights           = detail::GdnProjectionPayload;
    using PostMixerWeights               = detail::SparseMoePayload;
    using MtpAttentionProjectionWeights  = detail::AttentionProjectionPayload;
    using MtpPostMixerWeights            = detail::SparseMoePayload;
    using VisionWeights                  = qwen3_6::VisionWeights;
    using GraphExecutionProfile          = detail::GraphExecutionProfile;

    static constexpr float attention_scale                     = kAttentionScale;
    static constexpr float gdn_scale                           = kGdnScale;
    static constexpr std::uint32_t prefill_chunk_alignment     = kPrefillChunkAlignment;
    static constexpr std::uint32_t maximum_mtp_draft_tokens    = kMaximumMtpDraftTokens;
    static constexpr std::uint32_t maximum_dflash_draft_tokens = kMaximumDFlashDraftTokens;
    static constexpr std::uint32_t maximum_context             = kNativeContext;
    static constexpr bool supports_dflash                      = DFlashConfig::supported;
    static constexpr std::int32_t draft_head_rows              = 131072;

    [[nodiscard]] static std::vector<GraphExecutionProfile>
    ordinary_graph_profiles(std::uint32_t capacity);
    [[nodiscard]] static std::vector<GraphExecutionProfile>
    mtp_graph_profiles(std::uint32_t capacity, std::uint32_t draft_window);
    [[nodiscard]] static std::vector<GraphExecutionProfile>
    dflash_graph_profiles(std::uint32_t capacity, std::uint32_t draft_window,
                          std::uint32_t batch_size);

    static void attention_projection(const Tensor& hidden,
                                     const FullAttentionProjectionWeights& weights, Tensor& query,
                                     Tensor& gate, Tensor& key, Tensor& value,
                                     qwen3_6::TextPhase phase, WorkspaceArena& workspace,
                                     cudaStream_t stream);
    static void attention_output_projection(const Tensor& attention, const Weight& weight,
                                            Tensor& residual, qwen3_6::TextPhase phase,
                                            WorkspaceArena& workspace, cudaStream_t stream);
    static void mtp_attention_projection(const Tensor& hidden,
                                         const MtpAttentionProjectionWeights& weights,
                                         Tensor& query, Tensor& gate, Tensor& key, Tensor& value,
                                         WorkspaceArena& workspace, cudaStream_t stream);
    static void mtp_kv_projection(const Tensor& hidden,
                                  const MtpAttentionProjectionWeights& weights, Tensor& key,
                                  Tensor& value, WorkspaceArena& workspace, cudaStream_t stream);
    static void mtp_q_gate_projection(const Tensor& hidden,
                                      const MtpAttentionProjectionWeights& weights, Tensor& query,
                                      Tensor& gate, WorkspaceArena& workspace, cudaStream_t stream);
    static void gdn_input_projection(const Tensor& hidden, const GdnProjectionWeights& weights,
                                     Tensor& qkv, Tensor& output_gate, qwen3_6::TextPhase phase,
                                     WorkspaceArena& workspace, cudaStream_t stream);
    static void
    gdn_input_projection_snapshot(const Tensor& hidden, const GdnProjectionWeights& weights,
                                  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& output_gate, qwen3_6::TextPhase phase,
                                  WorkspaceArena& workspace, cudaStream_t stream);
    static void gdn_input_projection_record(
        const Tensor& hidden, const GdnProjectionWeights& weights, const Tensor& conv_weight,
        const Tensor& conv_states, const Tensor& valid_columns, const Tensor& initial_slots,
        Tensor& conv_record, Tensor& query, Tensor& key, Tensor& value, Tensor& output_gate,
        qwen3_6::TextPhase phase, WorkspaceArena& workspace, cudaStream_t stream);
    static void gdn_output_projection(const Tensor& hidden, const Weight& weight, Tensor& residual,
                                      qwen3_6::TextPhase phase, WorkspaceArena& workspace,
                                      cudaStream_t stream);
    static void gdn_norm_control_projection(const Tensor& residual, const Tensor& norm_weight,
                                            float eps, const GdnProjectionWeights& weights,
                                            Tensor& hidden, Tensor& g, Tensor& beta,
                                            WorkspaceArena& workspace, cudaStream_t stream);
    static void post_mixer(const Tensor& hidden, const PostMixerWeights& weights, Tensor& residual,
                           qwen3_6::TextPhase phase, WorkspaceArena& workspace,
                           cudaStream_t stream);
    static void mtp_post_mixer(const Tensor& hidden, const MtpPostMixerWeights& weights,
                               Tensor& residual, WorkspaceArena& workspace, cudaStream_t stream);

    [[nodiscard]] static std::size_t
    mtp_attention_projection_workspace_capacity_bytes(std::int32_t first, std::int32_t last);
    [[nodiscard]] static std::size_t mtp_kv_projection_workspace_capacity_bytes(std::int32_t first,
                                                                                std::int32_t last);
    [[nodiscard]] static std::size_t
    mtp_q_gate_projection_workspace_capacity_bytes(std::int32_t first, std::int32_t last);
    [[nodiscard]] static std::size_t
    attention_projection_workspace_capacity_bytes(WeightsProfile weights_profile,
                                                  qwen3_6::TextPhase phase, std::int32_t first,
                                                  std::int32_t last);
    [[nodiscard]] static std::size_t
    attention_output_projection_workspace_capacity_bytes(WeightsProfile weights_profile,
                                                         qwen3_6::TextPhase phase,
                                                         std::int32_t first, std::int32_t last);
    [[nodiscard]] static std::size_t
    gdn_input_projection_workspace_capacity_bytes(WeightsProfile weights_profile,
                                                  qwen3_6::TextPhase phase, std::int32_t first,
                                                  std::int32_t last);
    [[nodiscard]] static std::size_t gdn_input_projection_snapshot_workspace_capacity_bytes(
        WeightsProfile weights_profile, qwen3_6::TextPhase phase, std::int32_t batch_size,
        std::int32_t first, std::int32_t last);
    [[nodiscard]] static std::size_t gdn_input_projection_record_workspace_capacity_bytes(
        WeightsProfile weights_profile, qwen3_6::TextPhase phase, std::int32_t batch_size,
        std::int32_t first, std::int32_t last);
    [[nodiscard]] static std::size_t
    gdn_output_projection_workspace_capacity_bytes(WeightsProfile weights_profile,
                                                   qwen3_6::TextPhase phase, std::int32_t first,
                                                   std::int32_t last);
    [[nodiscard]] static std::size_t
    gdn_norm_control_projection_workspace_capacity_bytes(std::int32_t first, std::int32_t last);
    [[nodiscard]] static std::size_t
    post_mixer_workspace_capacity_bytes(WeightsProfile weights_profile, qwen3_6::TextPhase phase,
                                        std::int32_t first, std::int32_t last);
    [[nodiscard]] static std::size_t mtp_post_mixer_workspace_capacity_bytes(std::int32_t first,
                                                                             std::int32_t last);
};

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

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

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

namespace ninfer::ops::detail {
namespace {

constexpr int kTargetRows             = 9216;
constexpr int kCompanionRows          = 6144;
constexpr int kHidden                 = 2048;
constexpr int kRowsPerCta             = 16;
constexpr int kFirstExactCols         = 2;
constexpr int kLastTargetExactCols    = 48;
constexpr int kLastCompanionExactCols = 32;
using TargetOutput                    = W8SplitOutput4<4096, 512, 4096, 512>;
using CompanionOutput                 = W8SplitOutput3<4096, 1024, 1024>;
using TargetLauncher    = void (*)(const Tensor&, const Weight&, Tensor&, Tensor&, Tensor&, Tensor&,
                                cudaStream_t);
using CompanionLauncher = void (*)(const Tensor&, const Weight&, Tensor&, Tensor&, Tensor&,
                                   cudaStream_t);

template <int ActiveCols, int Rows, class Output>
void launch_output(const Tensor& x, const Weight& weight, Output output, cudaStream_t stream) {
    constexpr int TileCols = ActiveCols <= 8    ? 8
                             : ActiveCols <= 16 ? 16
                             : ActiveCols <= 24 ? 24
                             : ActiveCols <= 32 ? 32
                             : ActiveCols <= 40 ? 40
                                                : 48;
    using Geometry         = W8LinearGeometry<Rows, kHidden>;
    using Schedule         = W8SmallTMmaDefaultSchedule<TileCols, ActiveCols>;
    w8_small_t_mma_kernel<Geometry, ActiveCols, Schedule>
        <<<Rows / kRowsPerCta, Schedule::kThreads, 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);
}

template <int ActiveCols>
void launch_target_active_cols(const Tensor& x, const Weight& weight, Tensor& q, Tensor& gate,
                               Tensor& k, Tensor& v, cudaStream_t stream) {
    static_assert((4096 % kRowsPerCta) == 0 && (512 % kRowsPerCta) == 0);
    const TargetOutput 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)};
    launch_output<ActiveCols, kTargetRows>(x, weight, output, stream);
}

template <int ActiveCols>
void launch_companion_active_cols(const Tensor& x, const Weight& weight, Tensor& q, Tensor& k,
                                  Tensor& v, cudaStream_t stream) {
    static_assert((4096 % kRowsPerCta) == 0 && (1024 % kRowsPerCta) == 0);
    const CompanionOutput output{static_cast<__nv_bfloat16*>(q.data),
                                 static_cast<__nv_bfloat16*>(k.data),
                                 static_cast<__nv_bfloat16*>(v.data)};
    launch_output<ActiveCols, kCompanionRows>(x, weight, output, stream);
}

template <std::size_t... Offsets>
constexpr auto make_target_launchers(std::index_sequence<Offsets...>) {
    return std::array<TargetLauncher, sizeof...(Offsets)>{
        &launch_target_active_cols<kFirstExactCols + static_cast<int>(Offsets)>...};
}

template <std::size_t... Offsets>
constexpr auto make_companion_launchers(std::index_sequence<Offsets...>) {
    return std::array<CompanionLauncher, sizeof...(Offsets)>{
        &launch_companion_active_cols<kFirstExactCols + static_cast<int>(Offsets)>...};
}

constexpr auto kTargetLaunchers =
    make_target_launchers(std::make_index_sequence<kLastTargetExactCols - kFirstExactCols + 1>{});
constexpr auto kCompanionLaunchers = make_companion_launchers(
    std::make_index_sequence<kLastCompanionExactCols - kFirstExactCols + 1>{});

template <int TileCols, int KSplits, int NGroups, int MinBlocks>
void launch_target_medium_cols(const Tensor& x, const Weight& weight, Tensor& q, Tensor& gate,
                               Tensor& k, Tensor& v, cudaStream_t stream) {
    static_assert((4096 % kRowsPerCta) == 0 && (512 % kRowsPerCta) == 0);
    const TargetOutput 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_rowsplit_medium_t_splitk_kernel<kHidden, TileCols, KSplits, NGroups, MinBlocks>
        <<<kTargetRows / kRowsPerCta, KSplits * NGroups * 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, x.ne[1]);
}

template <int TileCols, int KSplits, int NGroups, int MinBlocks>
void launch_companion_medium_cols(const Tensor& x, const Weight& weight, Tensor& q, Tensor& k,
                                  Tensor& v, cudaStream_t stream) {
    static_assert((4096 % kRowsPerCta) == 0 && (1024 % kRowsPerCta) == 0);
    const CompanionOutput output{static_cast<__nv_bfloat16*>(q.data),
                                 static_cast<__nv_bfloat16*>(k.data),
                                 static_cast<__nv_bfloat16*>(v.data)};
    w8_rowsplit_medium_t_splitk_kernel<kHidden, TileCols, KSplits, NGroups, MinBlocks>
        <<<kCompanionRows / kRowsPerCta, KSplits * NGroups * 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, x.ne[1]);
}

} // namespace

void w8_attn_input_splitk_mma_launch(const Tensor& x, const Weight& weight, Tensor& q, Tensor& gate,
                                     Tensor& k, Tensor& v, cudaStream_t stream) {
    if (x.ne[1] < kFirstExactCols || x.ne[1] > 64) {
        throw std::invalid_argument("W8 attention input split-K MMA requires T=2..64");
    }
    if (x.ne[1] <= kLastTargetExactCols) {
        kTargetLaunchers[x.ne[1] - kFirstExactCols](x, weight, q, gate, k, v, stream);
    } else {
        launch_target_medium_cols<64, 4, 2, 2>(x, weight, q, gate, k, v, stream);
    }
    CUDA_CHECK(cudaGetLastError());
}

void w8_attn_input_splitk_mma_launch(const Tensor& x, const Weight& weight, Tensor& q, Tensor& k,
                                     Tensor& v, cudaStream_t stream) {
    if (x.ne[1] < kFirstExactCols || x.ne[1] > 96) {
        throw std::invalid_argument("W8 companion attention input split-K MMA requires T=2..96");
    }
    if (x.ne[1] <= kLastCompanionExactCols) {
        kCompanionLaunchers[x.ne[1] - kFirstExactCols](x, weight, q, k, v, stream);
    } else if (x.ne[1] <= 48) {
        launch_companion_medium_cols<48, 4, 2, 3>(x, weight, q, k, v, stream);
    } else if (x.ne[1] <= 64) {
        launch_companion_medium_cols<64, 4, 2, 2>(x, weight, q, k, v, stream);
    } else {
        launch_companion_medium_cols<96, 2, 4, 3>(x, weight, q, k, v, stream);
    }
    CUDA_CHECK(cudaGetLastError());
}

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

// ninfer::ops::detail - private launch prototypes for embedding variants.

#include "core/tensor.h"

#include <cuda_runtime.h>

namespace ninfer::ops::detail {

enum class W8EmbedRoute {
    Auto,
    Grouped,
    Row,
};

void embed_gather_dense_launch(const Tensor& ids, const Tensor& table, Tensor& out,
                               cudaStream_t stream);
void embed_gather_q6_launch(const Tensor& ids, const Weight& table, Tensor& out,
                            cudaStream_t stream);
void embed_gather_w8_launch(const Tensor& ids, const Weight& table, Tensor& out,
                            cudaStream_t stream);
void embed_gather_fp8_launch(const Tensor& ids, const Weight& table, Tensor& out,
                             cudaStream_t stream);
void embed_gather_w8_2048_launch(const Tensor& ids, const Weight& table, Tensor& out,
                                 W8EmbedRoute route, cudaStream_t stream);
const char* w8_embed_route_name(W8EmbedRoute route);

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

#include "core/layout.h"
#include "core/tensor.h"
#include "ninfer/ops/sampling.h"
#include "ninfer/types.h"

#include <array>
#include <cstddef>
#include <cstdint>
#include <optional>

namespace ninfer::targets::qwen3_6 {

inline constexpr std::uint32_t kMtpDecodeMaximumDrafts    = 5;
inline constexpr std::uint32_t kMtpDecodeMaximumWidth     = kMtpDecodeMaximumDrafts + 1;
inline constexpr std::uint32_t kDFlashDecodeMaximumDrafts = 15;
inline constexpr std::uint32_t kDFlashDecodeMaximumWidth  = kDFlashDecodeMaximumDrafts + 1;

struct RoundStateSpec {
    std::int32_t hidden          = 0;
    std::int32_t output_rows     = 0;
    std::uint32_t batch_capacity = 1;
    std::uint32_t draft_window   = 0;
    bool enable_mtp              = false;
    bool enable_dflash           = false;
};

// Stable pinned/device transfer format for ordinary decode. The full fixed-size object is copied
// once per round; only its exact-B prefixes are consumed by the model schedule.
struct OrdinaryDecodeIngress {
    std::array<TokenId, kMaximumConcurrency> tokens{};
    std::array<std::int32_t, kMaximumConcurrency> cache_positions{};
    std::array<std::int32_t, kMaximumConcurrency> rope_positions{};
    std::array<std::int32_t, kMaximumConcurrency> text_kv_table_rows{};
    std::array<std::int32_t, kMaximumConcurrency> state_source_slots{};
    std::array<std::int32_t, kMaximumConcurrency> state_destination_slots{};
    std::array<ops::SamplingConfig, kMaximumConcurrency> sampling{};
};

struct OrdinaryDecodeEgress {
    std::array<TokenId, kMaximumConcurrency> sampled_tokens{};
};

// Stable pinned/device transfer formats for concurrent MTP decode. The arrays use the maximum
// product domain; RoundState binds only the configured [K,C] and [K+1,C] prefixes.
struct MtpDecodeIngress {
    std::array<TokenId, kMaximumConcurrency> anchors{};
    std::array<std::int32_t, kMaximumConcurrency> base_frontiers{};
    std::array<std::int32_t, kMaximumConcurrency> remaining_budgets{};
    std::array<std::int32_t, kMaximumConcurrency> current_extents{};
    std::array<std::int32_t, kMaximumConcurrency> target_valid_columns{};
    std::array<TokenId, kMaximumConcurrency * kMtpDecodeMaximumDrafts> current_drafts{};
    std::array<std::int32_t, kMaximumConcurrency * kMtpDecodeMaximumWidth> target_rope_positions{};
    std::array<std::int32_t, kMaximumConcurrency> text_kv_table_rows{};
    std::array<std::int32_t, kMaximumConcurrency> mtp_kv_table_rows{};
    std::array<std::int32_t, kMaximumConcurrency> state_source_slots{};
    std::array<std::int32_t, kMaximumConcurrency> state_destination_slots{};
    std::array<std::int32_t, kMaximumConcurrency> rope_deltas{};
    std::array<ops::SamplingConfig, kMaximumConcurrency> sampling{};
};

struct MtpDecodeEgress {
    std::array<TokenId, kMaximumConcurrency * kMtpDecodeMaximumWidth> licensed_tokens{};
    std::array<std::int32_t, kMaximumConcurrency> licensed_counts{};
    std::array<std::int32_t, kMaximumConcurrency> accepted_drafts{};
    // Step-major: all B rows for proposal step 0, followed by all B rows for step 1, etc.
    std::array<TokenId, kMaximumConcurrency * kMtpDecodeMaximumDrafts> next_drafts{};
    std::array<std::int32_t, kMaximumConcurrency> next_extents{};
};

// Stable pinned/device transfer formats for one exact-B DFlash transaction. The proposal is
// produced and verified in the same round, so no draft state crosses the round boundary.
struct DFlashDecodeIngress {
    std::array<TokenId, kMaximumConcurrency> anchors{};
    std::array<std::int32_t, kMaximumConcurrency> execution_frontiers{};
    std::array<std::int32_t, kMaximumConcurrency> context_frontiers{};
    std::array<std::int32_t, kMaximumConcurrency> proposal_extents{};
    std::array<std::int32_t, kMaximumConcurrency> target_valid_columns{};
    std::array<std::int32_t, kMaximumConcurrency> text_kv_table_rows{};
    std::array<std::int32_t, kMaximumConcurrency> dflash_kv_table_rows{};
    std::array<std::int32_t, kMaximumConcurrency> active_lanes{};
    std::array<std::int32_t, kMaximumConcurrency> state_source_slots{};
    std::array<std::int32_t, kMaximumConcurrency> state_destination_slots{};
    std::array<ops::SamplingConfig, kMaximumConcurrency> sampling{};
};

struct DFlashDecodeEgress {
    std::array<TokenId, kMaximumConcurrency * kDFlashDecodeMaximumWidth> licensed_tokens{};
    std::array<std::int32_t, kMaximumConcurrency> licensed_counts{};
    std::array<std::int32_t, kMaximumConcurrency> accepted_drafts{};
};

struct OrdinaryDecodeStateLayout {
    LayoutRegion ingress;
    LayoutRegion egress;
    TensorRegion logits;
    TensorRegion hidden;
};

struct MtpPrefillStateLayout {
    TensorRegion position;
    TensorRegion ar_hidden;
    TensorRegion draft_tokens;
    TensorRegion target_input_ids;
    TensorRegion target_positions;
};

struct DFlashPrefillStateLayout {
    TensorRegion produced_count;
};

struct MtpDecodeStateLayout {
    LayoutRegion ingress;
    LayoutRegion egress;
    TensorRegion verify_ids;
    TensorRegion target_positions;
    TensorRegion target_argmax;
    TensorRegion target_logits;
    TensorRegion target_hidden;
    TensorRegion target_continuation_hidden;
    TensorRegion proposal_logits;
    TensorRegion alignment_ids;
    TensorRegion alignment_hidden;
    TensorRegion ar_hidden;
    TensorRegion next_hidden;
    TensorRegion ar_positions;
    TensorRegion ar_rope_positions;
    TensorRegion ar_valid_columns;
};

struct DFlashDecodeStateLayout {
    LayoutRegion ingress;
    LayoutRegion egress;
    TensorRegion proposal_ids;
    TensorRegion proposal_positions;
    TensorRegion append_positions;
    TensorRegion append_counts;
    TensorRegion draft_tokens;
    TensorRegion verify_ids;
    TensorRegion target_argmax;
    TensorRegion target_logits;
    TensorRegion target_hidden;
    TensorRegion target_continuation_hidden;
};

struct RoundStateLayout {
    RoundStateSpec spec;
    std::optional<OrdinaryDecodeStateLayout> ordinary;
    TensorRegion token;
    TensorRegion pos;
    TensorRegion rope_pos;
    TensorRegion rope_delta;
    TensorRegion logits;
    TensorRegion text_kv_table_row;
    TensorRegion backend_kv_table_row;
    std::optional<MtpPrefillStateLayout> mtp;
    std::optional<DFlashPrefillStateLayout> dflash_prefill;
    std::optional<MtpDecodeStateLayout> mtp_decode;
    std::optional<DFlashDecodeStateLayout> dflash_decode;
    bool complete = false;
};

struct OrdinaryDecodeState {
    DeviceSpan ingress;
    DeviceSpan egress;
    Tensor tokens;
    Tensor cache_positions;
    Tensor rope_positions;
    Tensor text_kv_table_rows;
    Tensor state_source_slots;
    Tensor state_destination_slots;
    const ops::SamplingConfig* sampling = nullptr;
    Tensor sampled_tokens;
    Tensor logits;
    Tensor hidden;

    OrdinaryDecodeState() = default;
    OrdinaryDecodeState(DeviceSpan backing, const OrdinaryDecodeStateLayout& layout,
                        std::uint32_t batch_capacity);
};

// The two planning calls expose one deliberate exact-target extension seam after scalar logits.
// This lets a target retain its schedule-sized prefill activation at the established physical
// address without making that activation part of the family round contract.
[[nodiscard]] RoundStateLayout begin_round_state_layout(LayoutBuilder& builder,
                                                        const RoundStateSpec& spec);
void complete_round_state_layout(LayoutBuilder& builder, RoundStateLayout& layout);

struct MtpPrefillState {
    Tensor position;
    Tensor ar_hidden;
    Tensor draft_tokens;
    Tensor target_input_ids;
    Tensor target_positions;

    MtpPrefillState() = default;
    MtpPrefillState(DeviceSpan backing, const MtpPrefillStateLayout& layout);
};

struct DFlashPrefillState {
    Tensor produced_count;

    DFlashPrefillState() = default;
    DFlashPrefillState(DeviceSpan backing, const DFlashPrefillStateLayout& layout);
};

struct MtpDecodeState {
    DeviceSpan ingress;
    DeviceSpan egress;
    Tensor anchors;
    Tensor base_frontiers;
    Tensor remaining_budgets;
    Tensor current_extents;
    Tensor target_valid_columns;
    Tensor current_drafts;
    Tensor target_rope_positions;
    Tensor text_kv_table_rows;
    Tensor mtp_kv_table_rows;
    Tensor state_source_slots;
    Tensor state_destination_slots;
    Tensor rope_deltas;
    const ops::SamplingConfig* sampling = nullptr;
    Tensor licensed_tokens;
    Tensor licensed_counts;
    Tensor accepted_drafts;
    Tensor next_drafts;
    Tensor next_extents;
    Tensor verify_ids;
    Tensor target_positions;
    Tensor target_argmax;
    Tensor target_logits;
    Tensor target_hidden;
    Tensor target_continuation_hidden;
    Tensor proposal_logits;
    Tensor alignment_ids;
    Tensor alignment_hidden;
    Tensor ar_hidden;
    Tensor next_hidden;
    Tensor ar_positions;
    Tensor ar_rope_positions;
    Tensor ar_valid_columns;

    MtpDecodeState() = default;
    MtpDecodeState(DeviceSpan backing, const MtpDecodeStateLayout& layout,
                   std::uint32_t batch_capacity, std::uint32_t draft_window);
};

struct DFlashDecodeState {
    DeviceSpan ingress;
    DeviceSpan egress;
    Tensor anchors;
    Tensor execution_frontiers;
    Tensor context_frontiers;
    Tensor proposal_extents;
    Tensor target_valid_columns;
    Tensor text_kv_table_rows;
    Tensor dflash_kv_table_rows;
    Tensor active_lanes;
    Tensor state_source_slots;
    Tensor state_destination_slots;
    const ops::SamplingConfig* sampling = nullptr;
    Tensor licensed_tokens;
    Tensor licensed_counts;
    Tensor accepted_drafts;
    Tensor proposal_ids;
    Tensor proposal_positions;
    Tensor append_positions;
    Tensor append_counts;
    Tensor draft_tokens;
    Tensor verify_ids;
    Tensor target_argmax;
    Tensor target_logits;
    Tensor target_hidden;
    Tensor target_continuation_hidden;

    DFlashDecodeState() = default;
    DFlashDecodeState(DeviceSpan backing, const DFlashDecodeStateLayout& layout,
                      std::uint32_t batch_capacity, std::uint32_t draft_window);
};

struct RoundState {
    std::optional<OrdinaryDecodeState> ordinary;
    Tensor token;
    Tensor pos;
    Tensor rope_pos;
    Tensor rope_delta;
    Tensor logits;
    Tensor text_kv_table_row;
    Tensor backend_kv_table_row;
    std::optional<MtpPrefillState> mtp;
    std::optional<DFlashPrefillState> dflash_prefill;
    std::optional<MtpDecodeState> mtp_decode;
    std::optional<DFlashDecodeState> dflash_decode;

    RoundState() = default;
    RoundState(DeviceSpan backing, const RoundStateLayout& layout);
};

} // namespace ninfer::targets::qwen3_6
#include "serve/openai_responses.h"
#include "serve/openai_common.h"
#include "serve/request_validation.h"

#include <algorithm>
#include <iterator>
#include <optional>
#include <stdexcept>
#include <string>
#include <string_view>
#include <unordered_map>
#include <unordered_set>
#include <utility>

namespace ninfer::serve {
namespace {

using Json = nlohmann::json;

void require_object(const Json& value, std::string_view name = "request body") {
    if (!value.is_object()) { bad_request(std::string(name) + " must be a JSON object"); }
}

void require_string_hint(const Json& body, const char* key, std::size_t max_bytes = 0) {
    if (!body.contains(key) || body.at(key).is_null()) { return; }
    if (!body.at(key).is_string()) { bad_request(std::string(key) + " must be a string", key); }
    if (max_bytes != 0 && body.at(key).get_ref<const std::string&>().size() > max_bytes) {
        bad_request(std::string(key) + " exceeds its maximum length", key, "invalid_value");
    }
}

std::string require_function_name(const Json& object, const char* param) {
    if (!object.contains("name") || !object.at("name").is_string()) {
        bad_request("function name must be a string", param);
    }
    const std::string name = object.at("name").get<std::string>();
    if (!valid_tool_name(name, 64)) {
        bad_request("function name must match [A-Za-z0-9_-]{1,64}", param);
    }
    return name;
}

std::optional<std::string> optional_namespace_name(const Json& object, const char* param) {
    if (!object.contains("namespace") || object.at("namespace").is_null()) { return std::nullopt; }
    if (!object.at("namespace").is_string()) {
        bad_request("function namespace must be a string or null", param);
    }
    const std::string name = object.at("namespace").get<std::string>();
    if (!valid_tool_name(name, 64)) {
        bad_request("function namespace must match [A-Za-z0-9_-]{1,64}", param,
                    "invalid_tool_name");
    }
    return name;
}

std::string require_namespace_tool_name(const Json& object) {
    if (!object.contains("name") || !object.at("name").is_string()) {
        bad_request("namespace name must be a string", "tools");
    }
    const std::string name = object.at("name").get<std::string>();
    if (!valid_tool_name(name, 64)) {
        bad_request("namespace name must match [A-Za-z0-9_-]{1,64}", "tools", "invalid_tool_name");
    }
    return name;
}

OpenAIResponsesFunctionIdentity function_identity(const Json& object, const char* param) {
    return OpenAIResponsesFunctionIdentity{.name = require_function_name(object, param),
                                           .wire_namespace =
                                               optional_namespace_name(object, param)};
}

std::string lower_function_identity(
    const OpenAIResponsesFunctionIdentity& identity,
    std::unordered_map<std::string, OpenAIResponsesFunctionIdentity>& identities,
    const char* param) {
    const std::string engine_name =
        identity.wire_namespace ? *identity.wire_namespace + "__" + identity.name : identity.name;
    if (!valid_tool_name(engine_name, 64)) {
        bad_request("flattened function identity '" + engine_name +
                        "' exceeds the Engine tool-name contract [A-Za-z0-9_-]{1,64}",
                    param, "invalid_tool_name");
    }
    const auto [position, inserted] = identities.emplace(engine_name, identity);
    if (!inserted && position->second != identity) {
        bad_request("function identity collision after namespace translation: '" + engine_name +
                        "'",
                    param, "duplicate_tool_name");
    }
    return engine_name;
}

void add_wire_function_identity(Json& object, const OpenAIResponsesFunctionIdentity& identity) {
    object["name"] = identity.name;
    if (identity.wire_namespace) { object["namespace"] = *identity.wire_namespace; }
}

std::string item_id(const Json& item, const char* prefix) {
    if (!item.contains("id") || item.at("id").is_null()) {
        return new_openai_response_item_id(prefix);
    }
    if (!item.at("id").is_string() || item.at("id").get_ref<const std::string&>().empty()) {
        bad_request("input Item id must be a non-empty string", "input");
    }
    return item.at("id").get<std::string>();
}

bool prompt_cache_breakpoint(const Json& value) {
    if (!value.contains("prompt_cache_breakpoint") ||
        value.at("prompt_cache_breakpoint").is_null()) {
        return false;
    }
    const Json& breakpoint = value.at("prompt_cache_breakpoint");
    if (!breakpoint.is_object() || breakpoint.size() != 1 || !breakpoint.contains("mode") ||
        !breakpoint.at("mode").is_string() ||
        breakpoint.at("mode").get<std::string>() != "explicit") {
        bad_request("prompt_cache_breakpoint must be {mode:'explicit'}", "input",
                    "invalid_cache_breakpoint");
    }
    return true;
}

ninfer::product::media_acquire::Source parse_image_source(const Json& part) {
    if (part.contains("file_id") && !part.at("file_id").is_null()) {
        bad_request("input_image.file_id requires a Files API, which NInfer does not provide",
                    "input", "file_inputs_not_supported");
    }
    if (!part.contains("image_url") || !part.at("image_url").is_string() ||
        part.at("image_url").get_ref<const std::string&>().empty()) {
        bad_request("input_image must contain a non-empty image_url", "input");
    }
    if (part.contains("detail") && !part.at("detail").is_null()) {
        if (!part.at("detail").is_string()) {
            bad_request("input_image.detail must be a string", "input");
        }
        if (part.at("detail").get<std::string>() != "auto") {
            bad_request("only input_image detail 'auto' is supported", "input",
                        "image_detail_not_supported");
        }
    }

    ninfer::product::media_acquire::Source source;
    source.value = part.at("image_url").get<std::string>();
    if (source.value.starts_with("data:")) {
        source.kind = ninfer::product::media_acquire::SourceKind::Data;
    } else if (source.value.starts_with("http://") || source.value.starts_with("https://")) {
        source.kind = ninfer::product::media_acquire::SourceKind::Url;
    } else {
        bad_request("input_image.image_url must use HTTP(S) or a data URI", "input");
    }
    return source;
}

ninfer::product::media_acquire::Source parse_video_source(const Json& part) {
    if (!part.contains("video_url") || !part.at("video_url").is_string() ||
        part.at("video_url").get_ref<const std::string&>().empty()) {
        bad_request("input_video must contain a non-empty video_url", "input");
    }
    ninfer::product::media_acquire::Source source;
    source.value = part.at("video_url").get<std::string>();
    if (source.value.starts_with("data:")) {
        source.kind = ninfer::product::media_acquire::SourceKind::Data;
    } else if (source.value.starts_with("http://") || source.value.starts_with("https://")) {
        source.kind = ninfer::product::media_acquire::SourceKind::Url;
    } else {
        bad_request("input_video.video_url must use HTTP(S) or a data URI", "input");
    }
    return source;
}

void apply_shared_breakpoint(ContentPart& part, const Json& wire, std::size_t& breakpoint_count) {
    if (!prompt_cache_breakpoint(wire)) { return; }
    ++breakpoint_count;
    if (breakpoint_count > 4) {
        bad_request("at most four prompt_cache_breakpoint values are supported per request",
                    "input", "too_many_cache_breakpoints");
    }
    part.cache_boundary_after = ninfer::PromptCacheMarkerKind::SharedStablePrefix;
}

struct ParsedMessage {
    ChatTurn turn;
    Json canonical;
};

void append_message_text(ParsedMessage& parsed, Json& content, const std::string& text,
                         const std::string& wire_type, const Json* wire,
                         std::size_t& breakpoint_count) {
    ContentPart part;
    part.kind     = ContentKind::Text;
    part.text     = text;
    part.type_raw = wire_type;
    if (wire != nullptr) { apply_shared_breakpoint(part, *wire, breakpoint_count); }
    parsed.turn.content.push_back(std::move(part));

    Json canonical;
    if (wire_type == "refusal") {
        canonical = Json{{"type", "refusal"}, {"refusal", text}};
    } else {
        canonical = Json{{"type", wire_type}, {"text", text}};
        if (wire_type == "output_text") {
            canonical["annotations"] = wire != nullptr && wire->contains("annotations")
                                           ? wire->at("annotations")
                                           : Json::array();
            if (wire != nullptr && wire->contains("logprobs")) {
                canonical["logprobs"] = wire->at("logprobs");
            }
        }
    }
    if (wire != nullptr && wire->contains("prompt_cache_breakpoint")) {
        canonical["prompt_cache_breakpoint"] = wire->at("prompt_cache_breakpoint");
    }
    content.push_back(std::move(canonical));
}

void parse_message_content_part(const Json& value, ChatRole role, ParsedMessage& parsed,
                                Json& content, std::size_t& breakpoint_count) {
    if (!value.is_object() || !value.contains("type") || !value.at("type").is_string()) {
        bad_request("input message content parts must have a string type", "input");
    }
    const std::string type = value.at("type").get<std::string>();
    if (type == "input_text") {
        if (!value.contains("text") || !value.at("text").is_string()) {
            bad_request("input_text must contain a string text", "input");
        }
        append_message_text(parsed, content, value.at("text").get<std::string>(), type, &value,
                            breakpoint_count);
        return;
    }
    if (type == "output_text") {
        if (role != ChatRole::Assistant) {
            bad_request("output_text is only valid on assistant messages", "input");
        }
        if (!value.contains("text") || !value.at("text").is_string()) {
            bad_request("output_text must contain a string text", "input");
        }
        if (value.contains("annotations") && !value.at("annotations").is_null() &&
            !value.at("annotations").is_array()) {
            bad_request("output_text.annotations must be an array", "input");
        }
        if (value.contains("logprobs") && !value.at("logprobs").is_null() &&
            !value.at("logprobs").is_array()) {
            bad_request("output_text.logprobs must be an array", "input");
        }
        append_message_text(parsed, content, value.at("text").get<std::string>(), type, &value,
                            breakpoint_count);
        return;
    }
    if (type == "refusal") {
        if (role != ChatRole::Assistant) {
            bad_request("refusal is only valid on assistant messages", "input");
        }
        if (!value.contains("refusal") || !value.at("refusal").is_string()) {
            bad_request("refusal must contain a string refusal", "input");
        }
        append_message_text(parsed, content, value.at("refusal").get<std::string>(), type, &value,
                            breakpoint_count);
        return;
    }
    if (type == "input_image") {
        if (role != ChatRole::User && role != ChatRole::Assistant) {
            bad_request("input_image is only supported on user or assistant messages", "input");
        }
        ContentPart part;
        part.kind     = ContentKind::Image;
        part.type_raw = type;
        part.source   = parse_image_source(value);
        apply_shared_breakpoint(part, value, breakpoint_count);
        parsed.turn.content.push_back(std::move(part));
        Json canonical{
            {"type", "input_image"}, {"image_url", value.at("image_url")}, {"detail", "auto"}};
        if (value.contains("prompt_cache_breakpoint")) {
            canonical["prompt_cache_breakpoint"] = value.at("prompt_cache_breakpoint");
        }
        content.push_back(std::move(canonical));
        return;
    }
    if (type == "input_video") {
        if (role != ChatRole::User) {
            bad_request("input_video is only supported on user messages", "input");
        }
        ContentPart part;
        part.kind     = ContentKind::Video;
        part.type_raw = type;
        part.source   = parse_video_source(value);
        parsed.turn.content.push_back(std::move(part));
        content.push_back(Json{{"type", "input_video"}, {"video_url", value.at("video_url")}});
        return;
    }
    if (type == "input_file") {
        bad_request("input_file requires a Files API, which NInfer does not provide", "input",
                    "file_inputs_not_supported");
    }
    if (type == "input_audio") {
        bad_request("input_audio is not supported by the Engine", "input",
                    "audio_inputs_not_supported");
    }
    bad_request("unsupported message content type: " + type, "input", "modality_not_supported");
}

ParsedMessage parse_message_item(const Json& item, std::size_t index,
                                 std::size_t& breakpoint_count) {
    if (!item.contains("role") || !item.at("role").is_string()) {
        bad_request("input message " + std::to_string(index) + " must contain a string role",
                    "input");
    }
    const std::string role = item.at("role").get<std::string>();
    ChatRole parsed_role;
    if (role == "user") {
        parsed_role = ChatRole::User;
    } else if (role == "assistant") {
        parsed_role = ChatRole::Assistant;
    } else if (role == "system") {
        parsed_role = ChatRole::System;
    } else if (role == "developer") {
        parsed_role = ChatRole::Developer;
    } else {
        bad_request("unsupported input message role: " + role, "input", "unsupported_role");
    }

    if (item.contains("status") && !item.at("status").is_null() && !item.at("status").is_string()) {
        bad_request("input message status must be a string", "input");
    }
    if (item.contains("phase") && !item.at("phase").is_null() && !item.at("phase").is_string()) {
        bad_request("input message phase must be a string", "input");
    }
    if (!item.contains("content") || item.at("content").is_null()) {
        bad_request("input message " + std::to_string(index) + " must contain content", "input");
    }

    ParsedMessage parsed;
    parsed.turn.role = parsed_role;
    Json content     = Json::array();

    if (item.at("content").is_string()) {
        append_message_text(parsed, content, item.at("content").get<std::string>(),
                            parsed_role == ChatRole::Assistant ? "output_text" : "input_text",
                            nullptr, breakpoint_count);
    } else if (item.at("content").is_array()) {
        for (const Json& value : item.at("content")) {
            parse_message_content_part(value, parsed_role, parsed, content, breakpoint_count);
        }
    } else {
        bad_request("input message content must be a string or array", "input");
    }
    if (parsed.turn.content.empty()) {
        bad_request("input message content must not be empty", "input");
    }

    parsed.canonical = {{"id", item_id(item, "msg")},
                        {"type", "message"},
                        {"role", role},
                        {"content", std::move(content)}};
    if (item.contains("status") && !item.at("status").is_null()) {
        parsed.canonical["status"] = item.at("status");
    }
    if (item.contains("phase") && !item.at("phase").is_null()) {
        parsed.canonical["phase"] = item.at("phase");
    }
    return parsed;
}

std::string parse_reasoning_item(const Json& item, Json& canonical) {
    const bool has_encrypted =
        item.contains("encrypted_content") && !item.at("encrypted_content").is_null();
    if (has_encrypted && !item.at("encrypted_content").is_string()) {
        bad_request("reasoning encrypted_content must be a string", "input");
    }
    if (item.contains("summary") && !item.at("summary").is_null() &&
        !item.at("summary").is_array()) {
        bad_request("reasoning summary must be an array", "input");
    }
    if (!item.contains("content") || !item.at("content").is_array()) {
        bad_request("reasoning Item must contain a content array", "input");
    }

    std::string text;
    Json content = Json::array();
    for (const Json& part : item.at("content")) {
        if (!part.is_object() || !part.contains("type") || !part.at("type").is_string() ||
            part.at("type").get<std::string>() != "reasoning_text" || !part.contains("text") ||
            !part.at("text").is_string()) {
            bad_request("reasoning content only supports reasoning_text parts", "input");
        }
        text += part.at("text").get<std::string>();
        content.push_back(Json{{"type", "reasoning_text"}, {"text", part.at("text")}});
    }
    if (text.empty() &&
        (has_encrypted || (item.contains("summary") && !item.at("summary").is_null() &&
                           !item.at("summary").empty()))) {
        bad_request("reasoning Items require raw reasoning_text; summary or encrypted content "
                    "cannot reconstruct the model context",
                    "input", "reasoning_content_not_supported");
    }

    canonical = {{"id", item_id(item, "rs")},
                 {"type", "reasoning"},
                 {"summary", item.contains("summary") && !item.at("summary").is_null()
                                 ? item.at("summary")
                                 : Json::array()},
                 {"content", std::move(content)}};
    if (has_encrypted) { canonical["encrypted_content"] = item.at("encrypted_content"); }
    return text;
}

void reject_nonnull_unknown_members(const Json& object,
                                    const std::unordered_set<std::string>& allowed,
                                    const char* param) {
    for (auto iterator = object.begin(); iterator != object.end(); ++iterator) {
        if (!allowed.contains(iterator.key()) && !iterator.value().is_null()) {
            bad_request("unsupported " + std::string(param) + " member: " + iterator.key(), param,
                        "parameter_not_supported");
        }
    }
}

ToolCall parse_function_call_item(
    const Json& item, Json& canonical,
    std::unordered_map<std::string, OpenAIResponsesFunctionIdentity>& identities) {
    static const std::unordered_set<std::string> allowed = {
        "id", "type", "call_id", "name", "arguments", "status", "caller", "namespace"};
    reject_nonnull_unknown_members(item, allowed, "input");
    if (item.contains("caller") && !item.at("caller").is_null()) {
        bad_request("function_call.caller is not supported", "input",
                    "tool_relationship_not_supported");
    }

    ToolCall call;
    if (!item.contains("call_id") || !item.at("call_id").is_string() ||
        item.at("call_id").get_ref<const std::string&>().empty()) {
        bad_request("function_call must contain a non-empty call_id", "input");
    }
    call.id                                        = item.at("call_id").get<std::string>();
    const OpenAIResponsesFunctionIdentity identity = function_identity(item, "input");
    call.name = lower_function_identity(identity, identities, "input");
    if (!item.contains("arguments") || !item.at("arguments").is_string()) {
        bad_request("function_call arguments must be a JSON string", "input");
    }
    call.arguments_json  = item.at("arguments").get<std::string>();
    const Json arguments = Json::parse(call.arguments_json, nullptr, false);
    if (arguments.is_discarded() || !arguments.is_object()) {
        bad_request("function_call arguments must encode a JSON object", "input");
    }
    if (item.contains("status") && !item.at("status").is_null() &&
        (!item.at("status").is_string() || item.at("status").get<std::string>() != "completed")) {
        bad_request("partial function_call Items cannot be represented in model history", "input",
                    "partial_tool_call_not_supported");
    }
    canonical = {{"id", item_id(item, "fc")},
                 {"type", "function_call"},
                 {"status", "completed"},
                 {"call_id", call.id},
                 {"arguments", call.arguments_json}};
    add_wire_function_identity(canonical, identity);
    return call;
}

ContentPart tool_output_text(std::string text, const Json* wire, std::size_t& breakpoint_count) {
    ContentPart part;
    part.kind     = ContentKind::Text;
    part.type_raw = "input_text";
    part.text     = std::move(text);
    if (wire != nullptr) { apply_shared_breakpoint(part, *wire, breakpoint_count); }
    return part;
}

ChatTurn parse_function_call_output_item(
    const Json& item, Json& canonical, std::size_t& breakpoint_count,
    std::unordered_map<std::string, OpenAIResponsesFunctionIdentity>& identities) {
    static const std::unordered_set<std::string> allowed = {
        "id", "type", "call_id", "output", "status", "caller", "name", "namespace"};
    reject_nonnull_unknown_members(item, allowed, "input");
    if (item.contains("caller") && !item.at("caller").is_null()) {
        bad_request("function_call_output.caller is not supported", "input",
                    "tool_relationship_not_supported");
    }
    if (!item.contains("call_id") || !item.at("call_id").is_string() ||
        item.at("call_id").get_ref<const std::string&>().empty()) {
        bad_request("function_call_output must contain a non-empty call_id", "input");
    }
    if (!item.contains("output")) {
        bad_request("function_call_output must contain output", "input");
    }
    if (item.contains("status") && !item.at("status").is_null() &&
        (!item.at("status").is_string() || item.at("status").get<std::string>() != "completed")) {
        bad_request("partial function_call_output Items cannot be represented in model history",
                    "input", "partial_tool_result_not_supported");
    }

    ChatTurn turn;
    turn.role                = ChatRole::Tool;
    turn.tool_call_id        = item.at("call_id").get<std::string>();
    const bool has_name      = item.contains("name") && !item.at("name").is_null();
    const bool has_namespace = item.contains("namespace") && !item.at("namespace").is_null();
    if (has_namespace && !has_name) {
        bad_request("function_call_output.namespace requires function_call_output.name", "input",
                    "invalid_tool_history");
    }
    std::optional<OpenAIResponsesFunctionIdentity> asserted_identity;
    if (has_name) {
        asserted_identity     = function_identity(item, "input");
        turn.tool_result_name = lower_function_identity(*asserted_identity, identities, "input");
    }
    if (item.at("output").is_string()) {
        turn.content.push_back(
            tool_output_text(item.at("output").get<std::string>(), nullptr, breakpoint_count));
    } else if (item.at("output").is_array()) {
        for (const Json& value : item.at("output")) {
            if (!value.is_object() || !value.contains("type") || !value.at("type").is_string()) {
                bad_request("function_call_output content parts must have a string type", "input");
            }
            const std::string type = value.at("type").get<std::string>();
            if (type == "input_text") {
                if (!value.contains("text") || !value.at("text").is_string()) {
                    bad_request("tool result input_text must contain a string text", "input");
                }
                turn.content.push_back(tool_output_text(value.at("text").get<std::string>(), &value,
                                                        breakpoint_count));
            } else if (type == "input_image") {
                ContentPart part;
                part.kind     = ContentKind::Image;
                part.type_raw = type;
                part.source   = parse_image_source(value);
                apply_shared_breakpoint(part, value, breakpoint_count);
                turn.content.push_back(std::move(part));
            } else if (type == "input_file") {
                bad_request("tool result input_file requires a Files API", "input",
                            "file_inputs_not_supported");
            } else {
                bad_request("unsupported function_call_output content type: " + type, "input",
                            "modality_not_supported");
            }
        }
        if (turn.content.empty()) {
            bad_request("function_call_output content must not be empty", "input");
        }
    } else {
        bad_request("function_call_output output must be a string or content array", "input");
    }

    canonical = {{"id", item_id(item, "fco")},
                 {"type", "function_call_output"},
                 {"status", "completed"},
                 {"call_id", turn.tool_call_id},
                 {"output", item.at("output")}};
    if (asserted_identity) { add_wire_function_identity(canonical, *asserted_identity); }
    return turn;
}

enum class AssistantInputPhase {
    Empty,
    Reasoning,
    Content,
    Calls,
};

[[noreturn]] void invalid_assistant_history(std::size_t index, std::string message) {
    bad_request("input Item " + std::to_string(index) + " " + std::move(message) +
                    "; supported assistant Item order is reasoning, message content, then "
                    "function calls",
                "input", "invalid_assistant_history");
}

struct AssistantInputRun {
    AssistantInputRun() { reset(); }

    void append_reasoning(std::string reasoning, std::size_t index) {
        switch (phase) {
        case AssistantInputPhase::Empty:
            turn.reasoning_content = std::move(reasoning);
            phase                  = AssistantInputPhase::Reasoning;
            return;
        case AssistantInputPhase::Reasoning:
            invalid_assistant_history(index, "reasoning cannot follow another reasoning Item");
        case AssistantInputPhase::Content:
            invalid_assistant_history(index, "reasoning cannot follow assistant message content");
        case AssistantInputPhase::Calls:
            invalid_assistant_history(index, "reasoning cannot follow function_call Items");
        }
        throw std::logic_error("unreachable assistant input phase");
    }

    void append_message(ChatTurn message, std::size_t index) {
        if (message.role != ChatRole::Assistant) {
            throw std::logic_error("assistant input run received a non-assistant message");
        }
        if (phase == AssistantInputPhase::Calls) {
            invalid_assistant_history(
                index, "assistant message content cannot follow function_call Items");
        }
        turn.content.insert(turn.content.end(), std::make_move_iterator(message.content.begin()),
                            std::make_move_iterator(message.content.end()));
        phase = AssistantInputPhase::Content;
    }

    void append_call(ToolCall call) {
        turn.tool_calls.push_back(std::move(call));
        phase = AssistantInputPhase::Calls;
    }

    void flush(std::vector<ChatTurn>& turns) {
        if (phase == AssistantInputPhase::Empty) { return; }
        if (!turn.reasoning_content.empty() || !turn.content.empty() || !turn.tool_calls.empty()) {
            turns.push_back(std::move(turn));
        }
        reset();
    }

private:
    void reset() {
        turn      = ChatTurn{};
        turn.role = ChatRole::Assistant;
        phase     = AssistantInputPhase::Empty;
    }

    ChatTurn turn;
    AssistantInputPhase phase = AssistantInputPhase::Empty;
};

void parse_input(const Json& input, OpenAIResponsesPromptRequest& out,
                 std::unordered_map<std::string, OpenAIResponsesFunctionIdentity>& identities) {
    Json values;
    if (input.is_string()) {
        values = Json::array({Json{{"type", "message"}, {"role", "user"}, {"content", input}}});
    } else if (input.is_array()) {
        values = input;
    } else {
        bad_request("input must be a string or an array of Items", "input");
    }

    AssistantInputRun assistant;
    std::size_t breakpoint_count = 0;
    std::unordered_set<std::string> item_ids;
    for (std::size_t index = 0; index < values.size(); ++index) {
        const Json& item = values.at(index);
        if (!item.is_object()) {
            bad_request("input Item " + std::to_string(index) + " must be an object", "input");
        }
        std::string type;
        if (item.contains("type") && !item.at("type").is_null()) {
            if (!item.at("type").is_string()) {
                bad_request("input Item type must be a string", "input");
            }
            type = item.at("type").get<std::string>();
        } else if (item.contains("role")) {
            type = "message";
        } else {
            bad_request("input Item must contain type", "input");
        }

        Json canonical;
        if (type == "message") {
            ParsedMessage message = parse_message_item(item, index, breakpoint_count);
            canonical             = std::move(message.canonical);
            if (message.turn.role == ChatRole::Assistant) {
                assistant.append_message(std::move(message.turn), index);
            } else {
                assistant.flush(out.input_turns);
                out.input_turns.push_back(std::move(message.turn));
            }
        } else if (type == "reasoning") {
            assistant.append_reasoning(parse_reasoning_item(item, canonical), index);
        } else if (type == "function_call") {
            assistant.append_call(parse_function_call_item(item, canonical, identities));
        } else if (type == "function_call_output") {
            ChatTurn result =
                parse_function_call_output_item(item, canonical, breakpoint_count, identities);
            assistant.flush(out.input_turns);
            out.input_turns.push_back(std::move(result));
        } else if (type == "input_file") {
            bad_request("input_file requires a Files API, which NInfer does not provide", "input",
                        "file_inputs_not_supported");
        } else {
            bad_request("unsupported input Item type: " + type, "input", "item_type_not_supported");
        }

        const std::string id = canonical.at("id").get<std::string>();
        if (!item_ids.insert(id).second) { bad_request("duplicate input Item id: " + id, "input"); }
        out.input_items.push_back(std::move(canonical));
    }
    assistant.flush(out.input_turns);
}

struct ParsedPromptFields {
    OpenAIResponsesPromptRequest prompt;
    Json wire_tools          = Json::array();
    Json wire_tool_choice    = "auto";
    bool parallel_tool_calls = true;
    std::unordered_map<std::string, OpenAIResponsesFunctionIdentity> tool_identities;
};

struct ParsedFunctionTool {
    ToolDefinition definition;
    Json canonical;
    std::string engine_name;
};

ParsedFunctionTool
parse_function_tool(const Json& item, std::optional<std::string> wire_namespace,
                    std::string_view namespace_description,
                    std::unordered_map<std::string, OpenAIResponsesFunctionIdentity>& identities) {
    static const std::unordered_set<std::string> allowed_members = {
        "type",          "name",         "description", "parameters", "strict", "allowed_callers",
        "defer_loading", "output_schema"};
    reject_nonnull_unknown_members(item, allowed_members, "tools");

    const OpenAIResponsesFunctionIdentity identity{.name = require_function_name(item, "tools"),
                                                   .wire_namespace = std::move(wire_namespace)};
    ParsedFunctionTool parsed;
    parsed.engine_name     = lower_function_identity(identity, identities, "tools");
    parsed.definition.name = parsed.engine_name;

    std::string function_description;
    if (item.contains("description") && !item.at("description").is_null()) {
        if (!item.at("description").is_string()) {
            bad_request("function description must be a string", "tools");
        }
        function_description = item.at("description").get<std::string>();
    }
    if (!namespace_description.empty()) {
        parsed.definition.description = std::string(namespace_description);
        if (!function_description.empty()) {
            parsed.definition.description += "\n\n" + function_description;
        }
    } else {
        parsed.definition.description = function_description;
    }

    Json parameters = Json{{"type", "object"}, {"properties", Json::object()}};
    if (item.contains("parameters") && !item.at("parameters").is_null()) {
        if (!item.at("parameters").is_object()) {
            bad_request("function parameters must be a JSON object", "tools");
        }
        parameters = item.at("parameters");
    }
    if (item.contains("strict") && !item.at("strict").is_null()) {
        if (!item.at("strict").is_boolean()) {
            bad_request("function strict must be a boolean", "tools");
        }
        if (item.at("strict").get<bool>()) {
            bad_request("strict function schema enforcement requires constrained decoding, "
                        "which the Engine does not provide",
                        "tools", "strict_tools_not_supported");
        }
    }
    if (item.contains("defer_loading") && !item.at("defer_loading").is_null()) {
        if (!item.at("defer_loading").is_boolean()) {
            bad_request("function defer_loading must be a boolean", "tools");
        }
        if (item.at("defer_loading").get<bool>()) {
            bad_request("deferred tool loading is not supported", "tools",
                        "deferred_tools_not_supported");
        }
    }
    if (item.contains("allowed_callers") && !item.at("allowed_callers").is_null()) {
        const Json& callers = item.at("allowed_callers");
        if (!callers.is_array()) {
            bad_request("function allowed_callers must be an array", "tools");
        }
        bool direct = false;
        for (const Json& caller : callers) {
            if (!caller.is_string()) {
                bad_request("function allowed_callers entries must be strings", "tools");
            }
            direct = direct || caller.get<std::string>() == "direct";
        }
        if (!direct) {
            bad_request("function allowed_callers must permit direct invocation", "tools",
                        "tool_caller_not_supported");
        }
    }
    if (item.contains("output_schema") && !item.at("output_schema").is_null()) {
        bad_request("function output_schema cannot be enforced", "tools",
                    "tool_output_schema_not_supported");
    }

    parsed.definition.input_schema_json = parameters.dump();
    parsed.canonical                    = {{"type", "function"},
                                           {"name", identity.name},
                                           {"parameters", parameters},
                                           {"strict", false}};
    if (!function_description.empty()) {
        parsed.canonical["description"] = std::move(function_description);
    }
    if (item.contains("allowed_callers") && !item.at("allowed_callers").is_null()) {
        parsed.canonical["allowed_callers"] = item.at("allowed_callers");
    }
    if (item.contains("defer_loading") && !item.at("defer_loading").is_null()) {
        parsed.canonical["defer_loading"] = false;
    }
    return parsed;
}

void parse_tools(const Json& body, ParsedPromptFields& out) {
    if (!body.contains("tools") || body.at("tools").is_null()) { return; }
    if (!body.at("tools").is_array()) { bad_request("tools must be an array", "tools"); }

    static const std::unordered_set<std::string> namespace_members = {"type", "name", "description",
                                                                      "tools"};
    std::unordered_set<std::string> declared_names;
    std::unordered_set<std::string> namespace_names;
    const auto append_function = [&](ParsedFunctionTool parsed) {
        if (!declared_names.insert(parsed.engine_name).second) {
            bad_request("duplicate function tool identity: " + parsed.engine_name, "tools",
                        "duplicate_tool_name");
        }
        out.prompt.generation.tools.push_back(std::move(parsed.definition));
        return std::move(parsed.canonical);
    };

    for (const Json& item : body.at("tools")) {
        if (!item.is_object() || !item.contains("type") || !item.at("type").is_string()) {
            bad_request("tools entries must be objects with a string type", "tools");
        }
        const std::string type = item.at("type").get<std::string>();
        if (type == "function") {
            out.wire_tools.push_back(
                append_function(parse_function_tool(item, std::nullopt, {}, out.tool_identities)));
            continue;
        }
        if (type != "namespace") {
            bad_request("tool type '" + type +
                            "' requires an executor that NInfer does not provide",
                        "tools", "tool_type_not_supported");
        }

        // OpenAI Responses beta groups functions/custom tools under a namespace. NInfer lowers
        // only nested functions because custom tools require unsupported free-form decoding.
        reject_nonnull_unknown_members(item, namespace_members, "tools");
        const std::string namespace_name = require_namespace_tool_name(item);
        if (!namespace_names.insert(namespace_name).second) {
            bad_request("duplicate namespace tool name: " + namespace_name, "tools",
                        "duplicate_tool_name");
        }
        std::string namespace_description;
        if (item.contains("description") && !item.at("description").is_null()) {
            if (!item.at("description").is_string()) {
                bad_request("namespace description must be a string", "tools");
            }
            namespace_description = item.at("description").get<std::string>();
        }
        if (!item.contains("tools") || !item.at("tools").is_array()) {
            bad_request("namespace tools must contain a tools array", "tools");
        }
        Json canonical = {
            {"type", "namespace"}, {"name", namespace_name}, {"tools", Json::array()}};
        if (!namespace_description.empty()) { canonical["description"] = namespace_description; }
        for (const Json& nested : item.at("tools")) {
            if (!nested.is_object() || !nested.contains("type") || !nested.at("type").is_string()) {
                bad_request("namespace tool entries must be objects with a string type", "tools");
            }
            const std::string nested_type = nested.at("type").get<std::string>();
            if (nested_type != "function") {
                bad_request("nested tool type '" + nested_type +
                                "' cannot be represented by the Engine",
                            "tools", "tool_type_not_supported");
            }
            canonical["tools"].push_back(append_function(parse_function_tool(
                nested, namespace_name, namespace_description, out.tool_identities)));
        }
        out.wire_tools.push_back(std::move(canonical));
    }
}

void filter_allowed_tools(const Json& choice, ParsedPromptFields& out) {
    static const std::unordered_set<std::string> allowed_choice = {"type", "mode", "tools"};
    reject_nonnull_unknown_members(choice, allowed_choice, "tool_choice");
    if (!choice.contains("mode") || !choice.at("mode").is_string()) {
        bad_request("allowed_tools tool_choice must contain a string mode", "tool_choice");
    }
    if (choice.at("mode").get<std::string>() != "auto") {
        bad_request("allowed_tools mode 'required' cannot be enforced", "tool_choice",
                    "tool_choice_not_supported");
    }
    if (!choice.contains("tools") || !choice.at("tools").is_array()) {
        bad_request("allowed_tools tool_choice must contain a tools array", "tool_choice");
    }

    std::unordered_set<std::string> declared;
    for (const ToolDefinition& tool : out.prompt.generation.tools) { declared.insert(tool.name); }
    std::unordered_set<std::string> selected;
    for (const Json& item : choice.at("tools")) {
        if (!item.is_object() || !item.contains("type") || !item.at("type").is_string() ||
            item.at("type").get<std::string>() != "function") {
            bad_request("allowed_tools only supports function entries", "tool_choice",
                        "tool_choice_not_supported");
        }
        static const std::unordered_set<std::string> allowed_entry = {"type", "name", "namespace"};
        reject_nonnull_unknown_members(item, allowed_entry, "tool_choice");
        const OpenAIResponsesFunctionIdentity identity = function_identity(item, "tool_choice");
        const std::string name =
            lower_function_identity(identity, out.tool_identities, "tool_choice");
        if (!declared.contains(name)) {
            const std::string wire_name = identity.wire_namespace
                                              ? *identity.wire_namespace + "." + identity.name
                                              : identity.name;
            bad_request("allowed_tools references undeclared function '" + wire_name + "'",
                        "tool_choice", "invalid_tool_choice");
        }
        selected.insert(name);
    }

    std::vector<ToolDefinition> effective;
    effective.reserve(selected.size());
    for (ToolDefinition& tool : out.prompt.generation.tools) {
        if (selected.contains(tool.name)) { effective.push_back(std::move(tool)); }
    }
    out.prompt.generation.tools = std::move(effective);
}

void parse_tool_choice(const Json& body, ParsedPromptFields& out) {
    if (!body.contains("tool_choice") || body.at("tool_choice").is_null()) {
        out.wire_tool_choice = "auto";
        return;
    }
    const Json& choice = body.at("tool_choice");
    if (choice.is_string()) {
        const std::string value = choice.get<std::string>();
        if (value == "auto") {
            out.prompt.generation.tool_choice.mode = ToolChoiceMode::Auto;
        } else if (value == "none") {
            out.prompt.generation.tool_choice.mode = ToolChoiceMode::None;
        } else if (value == "required") {
            bad_request("tool_choice 'required' cannot be guaranteed by the Engine", "tool_choice",
                        "tool_choice_not_supported");
        } else {
            bad_request("tool_choice must be 'auto', 'none', or a supported object", "tool_choice");
        }
        out.wire_tool_choice = value;
        return;
    }
    if (!choice.is_object() || !choice.contains("type") || !choice.at("type").is_string()) {
        bad_request("tool_choice must be a string or typed object", "tool_choice");
    }
    if (choice.at("type").get<std::string>() != "allowed_tools") {
        bad_request("named or hosted tool_choice cannot be enforced", "tool_choice",
                    "tool_choice_not_supported");
    }
    filter_allowed_tools(choice, out);
    out.wire_tool_choice = choice;
}

void parse_reasoning(const Json& body, OpenAIResponsesPromptRequest& out) {
    if (!body.contains("reasoning") || body.at("reasoning").is_null()) { return; }
    const Json& reasoning = body.at("reasoning");
    if (!reasoning.is_object()) { bad_request("reasoning must be an object", "reasoning"); }
    static const std::unordered_set<std::string> allowed = {"effort", "context", "summary",
                                                            "generate_summary", "mode"};
    reject_nonnull_unknown_members(reasoning, allowed, "reasoning");
    for (const char* key : {"context", "summary", "generate_summary", "mode"}) {
        if (reasoning.contains(key) && !reasoning.at(key).is_null()) {
            bad_request("reasoning." + std::string(key) +
                            " changes reasoning input or output and is not supported",
                        "reasoning", "reasoning_option_not_supported");
        }
    }
    if (!reasoning.contains("effort") || reasoning.at("effort").is_null()) { return; }
    if (!reasoning.at("effort").is_string()) {
        bad_request("reasoning.effort must be a string", "reasoning");
    }
    const std::string value = reasoning.at("effort").get<std::string>();
    const std::optional<RequestedReasoningEffort> effort = parse_requested_reasoning_effort(value);
    if (!effort) {
        bad_request("reasoning.effort must be one of none, minimal, low, medium, high, xhigh, or "
                    "max",
                    "reasoning");
    }
    out.generation.reasoning_effort = *effort;
}

void parse_text(const Json& body) {
    if (!body.contains("text") || body.at("text").is_null()) { return; }
    const Json& text = body.at("text");
    if (!text.is_object()) { bad_request("text must be an object", "text"); }
    static const std::unordered_set<std::string> allowed = {"format", "verbosity"};
    reject_nonnull_unknown_members(text, allowed, "text");
    if (text.contains("format") && !text.at("format").is_null()) {
        const Json& format = text.at("format");
        if (!format.is_object() || !format.contains("type") || !format.at("type").is_string()) {
            bad_request("text.format must be a typed object", "text");
        }
        if (format.at("type").get<std::string>() != "text" || format.size() != 1) {
            bad_request("structured text output requires constrained decoding, which the Engine "
                        "does not provide",
                        "text", "structured_outputs_not_supported");
        }
    }
    if (text.contains("verbosity") && !text.at("verbosity").is_null()) {
        if (!text.at("verbosity").is_string()) {
            bad_request("text.verbosity must be a string", "text");
        }
        const std::string verbosity = text.at("verbosity").get<std::string>();
        if (verbosity != "medium") {
            bad_request("text.verbosity '" + verbosity + "' cannot be enforced by the Engine",
                        "text", "verbosity_not_supported");
        }
    }
}

void parse_preserve_thinking(const Json& body, OpenAIResponsesPromptRequest& out) {
    if (body.contains("preserve_thinking") && !body.at("preserve_thinking").is_null()) {
        if (!body.at("preserve_thinking").is_boolean()) {
            bad_request("preserve_thinking must be a boolean or null", "preserve_thinking");
        }
        out.generation.preserve_thinking = body.at("preserve_thinking").get<bool>();
    }
    if (!body.contains("chat_template_kwargs") || body.at("chat_template_kwargs").is_null()) {
        return;
    }
    const Json& kwargs = body.at("chat_template_kwargs");
    if (!kwargs.is_object()) {
        bad_request("chat_template_kwargs must be an object", "chat_template_kwargs");
    }
    for (auto iterator = kwargs.begin(); iterator != kwargs.end(); ++iterator) {
        if (iterator.key() != "preserve_thinking" && !iterator.value().is_null()) {
            bad_request("chat_template_kwargs." + iterator.key() + " is not supported",
                        "chat_template_kwargs", "chat_template_option_not_supported");
        }
    }
    if (!kwargs.contains("preserve_thinking") || kwargs.at("preserve_thinking").is_null()) {
        return;
    }
    if (!kwargs.at("preserve_thinking").is_boolean()) {
        bad_request("chat_template_kwargs.preserve_thinking must be a boolean or null",
                    "chat_template_kwargs");
    }
    const bool nested = kwargs.at("preserve_thinking").get<bool>();
    if (out.generation.preserve_thinking && *out.generation.preserve_thinking != nested) {
        bad_request("conflicting preserve_thinking values", "preserve_thinking",
                    "conflicting_template_option");
    }
    out.generation.preserve_thinking = nested;
}

void parse_truncation(const Json& body) {
    if (!body.contains("truncation") || body.at("truncation").is_null()) { return; }
    if (!body.at("truncation").is_string()) {
        bad_request("truncation must be a string", "truncation");
    }
    if (body.at("truncation").get<std::string>() != "disabled") {
        bad_request("truncation 'auto' would discard input Items and is not supported",
                    "truncation", "truncation_not_supported");
    }
}

ParsedPromptFields parse_prompt_fields(const Json& body, const RequestLimits& limits) {
    ParsedPromptFields out;
    if (!body.contains("model") || !body.at("model").is_string() ||
        body.at("model").get_ref<const std::string&>().empty()) {
        bad_request("missing required field: model", "model");
    }
    out.prompt.model = body.at("model").get<std::string>();
    if (body.contains("input") && !body.at("input").is_null()) {
        parse_input(body.at("input"), out.prompt, out.tool_identities);
    }
    if (body.contains("instructions") && !body.at("instructions").is_null()) {
        if (!body.at("instructions").is_string()) {
            bad_request("instructions must be a string", "instructions");
        }
        out.prompt.instructions = body.at("instructions").get<std::string>();
    }
    if (body.contains("previous_response_id") && !body.at("previous_response_id").is_null()) {
        if (!body.at("previous_response_id").is_string() ||
            body.at("previous_response_id").get_ref<const std::string&>().empty()) {
            bad_request("previous_response_id must be a non-empty string", "previous_response_id");
        }
        out.prompt.previous_response_id = body.at("previous_response_id").get<std::string>();
    }

    parse_tools(body, out);
    parse_tool_choice(body, out);
    out.parallel_tool_calls = optional_bool(body, "parallel_tool_calls", true);
    if (!out.parallel_tool_calls && out.prompt.generation.uses_tools()) {
        bad_request("parallel_tool_calls=false cannot be guaranteed when callable tools are "
                    "present",
                    "parallel_tool_calls", "parallel_tool_calls_not_supported");
    }
    parse_reasoning(body, out.prompt);
    parse_text(body);
    parse_truncation(body);
    parse_preserve_thinking(body, out.prompt);
    out.prompt.generation.max_tokens = limits.default_max_tokens;
    return out;
}

void validate_metadata(const Json& body, Json& metadata) {
    if (!body.contains("metadata") || body.at("metadata").is_null()) { return; }
    if (!body.at("metadata").is_object()) { bad_request("metadata must be an object", "metadata"); }
    if (body.at("metadata").size() > 16) {
        bad_request("metadata supports at most 16 entries", "metadata");
    }
    for (auto iterator = body.at("metadata").begin(); iterator != body.at("metadata").end();
         ++iterator) {
        if (iterator.key().size() > 64 || !iterator.value().is_string() ||
            iterator.value().get_ref<const std::string&>().size() > 512) {
            bad_request("metadata keys must be at most 64 characters and string values at most "
                        "512 characters",
                        "metadata");
        }
    }
    metadata = body.at("metadata");
}

void parse_prompt_cache_hints(const Json& body) {
    require_string_hint(body, "prompt_cache_key");
    require_string_hint(body, "safety_identifier", 64);
    require_string_hint(body, "user");

    if (body.contains("prompt_cache_retention") && !body.at("prompt_cache_retention").is_null()) {
        if (!body.at("prompt_cache_retention").is_string()) {
            bad_request("prompt_cache_retention must be a string", "prompt_cache_retention");
        }
        const std::string value = body.at("prompt_cache_retention").get<std::string>();
        if (value != "in_memory" && value != "24h") {
            bad_request("prompt_cache_retention must be 'in_memory' or '24h'",
                        "prompt_cache_retention");
        }
    }
    if (!body.contains("prompt_cache_options") || body.at("prompt_cache_options").is_null()) {
        return;
    }
    const Json& options = body.at("prompt_cache_options");
    if (!options.is_object()) {
        bad_request("prompt_cache_options must be an object", "prompt_cache_options");
    }
    static const std::unordered_set<std::string> allowed = {"mode", "ttl"};
    reject_nonnull_unknown_members(options, allowed, "prompt_cache_options");
    if (options.contains("mode") && !options.at("mode").is_null()) {
        if (!options.at("mode").is_string()) {
            bad_request("prompt_cache_options.mode must be a string", "prompt_cache_options");
        }
        const std::string mode = options.at("mode").get<std::string>();
        if (mode != "implicit" && mode != "explicit") {
            bad_request("prompt_cache_options.mode must be 'implicit' or 'explicit'",
                        "prompt_cache_options");
        }
    }
    if (options.contains("ttl") && !options.at("ttl").is_null() &&
        (!options.at("ttl").is_string() || options.at("ttl").get<std::string>() != "30m")) {
        bad_request("prompt_cache_options.ttl must be '30m'", "prompt_cache_options");
    }
}

void reject_unsupported_platform_fields(const Json& body) {
    const struct {
        const char* field;
        const char* code;
        const char* reason;
    } unsupported[] = {
        {"conversation", "conversations_not_supported",
         "conversation requires an OpenAI Conversations resource"},
        {"prompt", "prompt_templates_not_supported",
         "prompt requires an OpenAI prompt-template resource"},
        {"context_management", "context_management_not_supported",
         "context_management requires an API compaction pipeline"},
        {"moderation", "moderation_not_supported",
         "moderation changes request acceptance and output but no moderator is configured"},
    };

    for (const auto& entry : unsupported) {
        if (body.contains(entry.field) && !body.at(entry.field).is_null()) {
            bad_request(entry.reason, entry.field, entry.code);
        }
    }
}

void validate_common_top_level(const Json& body, bool create) {
    static const std::unordered_set<std::string> create_fields = {"background",
                                                                  "chat_template_kwargs",
                                                                  "client_metadata",
                                                                  "context_management",
                                                                  "conversation",
                                                                  "include",
                                                                  "input",
                                                                  "instructions",
                                                                  "max_output_tokens",
                                                                  "max_tool_calls",
                                                                  "metadata",
                                                                  "model",
                                                                  "moderation",
                                                                  "parallel_tool_calls",
                                                                  "previous_response_id",
                                                                  "preserve_thinking",
                                                                  "prompt",
                                                                  "prompt_cache_key",
                                                                  "prompt_cache_options",
                                                                  "prompt_cache_retention",
                                                                  "reasoning",
                                                                  "safety_identifier",
                                                                  "service_tier",
                                                                  "store",
                                                                  "stream",
                                                                  "stream_options",
                                                                  "temperature",
                                                                  "text",
                                                                  "tool_choice",
                                                                  "tools",
                                                                  "top_logprobs",
                                                                  "top_p",
                                                                  "truncation",
                                                                  "user"};
    static const std::unordered_set<std::string> count_fields  = {"chat_template_kwargs",
                                                                  "conversation",
                                                                  "input",
                                                                  "instructions",
                                                                  "model",
                                                                  "parallel_tool_calls",
                                                                  "personality",
                                                                  "previous_response_id",
                                                                  "preserve_thinking",
                                                                  "reasoning",
                                                                  "text",
                                                                  "tool_choice",
                                                                  "tools",
                                                                  "truncation"};
    const auto& allowed = create ? create_fields : count_fields;
    for (auto iterator = body.begin(); iterator != body.end(); ++iterator) {
        if (!allowed.contains(iterator.key())) {
            bad_request("unknown parameter: " + iterator.key(), iterator.key(),
                        "unknown_parameter");
        }
    }
}

} // namespace

OpenAIResponsesCreateRequest parse_openai_responses_create_request(const Json& body,
                                                                   const RequestLimits& limits) {
    require_object(body);
    validate_common_top_level(body, true);
    reject_unsupported_platform_fields(body);
    parse_prompt_cache_hints(body);

    ParsedPromptFields parsed = parse_prompt_fields(body, limits);
    OpenAIResponsesCreateRequest out;
    out.prompt              = std::move(parsed.prompt);
    out.tools               = std::move(parsed.wire_tools);
    out.tool_choice         = std::move(parsed.wire_tool_choice);
    out.tool_identities     = std::move(parsed.tool_identities);
    out.parallel_tool_calls = parsed.parallel_tool_calls;
    out.store               = optional_bool(body, "store", true);
    out.stream              = optional_bool(body, "stream", false);
    validate_metadata(body, out.metadata);

    // Codex attaches per-request tracing information here. It is an opaque client hint and has no
    // Engine, cache-identity, response-store, or response-body semantics.
    if (body.contains("client_metadata") && !body.at("client_metadata").is_null() &&
        !body.at("client_metadata").is_object()) {
        bad_request("client_metadata must be an object or null", "client_metadata", "invalid_type");
    }

    if (body.contains("background") && !body.at("background").is_null()) {
        if (!body.at("background").is_boolean()) {
            bad_request("background must be a boolean", "background");
        }
        if (body.at("background").get<bool>()) {
            bad_request("background execution is not supported", "background",
                        "background_not_supported");
        }
    }
    if (body.contains("include") && !body.at("include").is_null()) {
        if (!body.at("include").is_array()) { bad_request("include must be an array", "include"); }
        if (!body.at("include").empty()) {
            bad_request("the requested additional response fields have no available response "
                        "representation",
                        "include", "include_not_supported");
        }
    }
    if (body.contains("stream_options") && !body.at("stream_options").is_null()) {
        const Json& options = body.at("stream_options");
        if (!options.is_object()) {
            bad_request("stream_options must be an object", "stream_options");
        }
        static const std::unordered_set<std::string> allowed = {"include_obfuscation"};
        reject_nonnull_unknown_members(options, allowed, "stream_options");
        if (options.contains("include_obfuscation") &&
            !options.at("include_obfuscation").is_null() &&
            !options.at("include_obfuscation").is_boolean()) {
            bad_request("stream_options.include_obfuscation must be a boolean", "stream_options");
        }
    }
    if (body.contains("service_tier") && !body.at("service_tier").is_null()) {
        if (!body.at("service_tier").is_string()) {
            bad_request("service_tier must be a string", "service_tier");
        }
        const std::string tier = body.at("service_tier").get<std::string>();
        if (tier != "auto" && tier != "default") {
            bad_request("the requested service tier is not provided by this local server",
                        "service_tier", "service_tier_not_supported");
        }
    }
    if (const std::optional<int> top_logprobs = optional_int(body, "top_logprobs")) {
        if (*top_logprobs < 0 || *top_logprobs > 20) {
            bad_request("top_logprobs must be in [0,20]", "top_logprobs");
        }
        if (*top_logprobs != 0) {
            bad_request("the Engine does not return token log probabilities", "top_logprobs",
                        "logprobs_not_supported");
        }
    }
    if (const std::optional<int> max_tool_calls = optional_int(body, "max_tool_calls")) {
        if (*max_tool_calls < 0) {
            bad_request("max_tool_calls must be non-negative", "max_tool_calls");
        }
        out.max_tool_calls = *max_tool_calls;
    }

    if (const std::optional<double> temperature = optional_number(body, "temperature")) {
        if (*temperature < 0.0 || *temperature > 2.0) {
            bad_request("temperature must be in [0,2]", "temperature");
        }
        out.prompt.generation.sampling.temperature = *temperature;
    }
    if (const std::optional<double> top_p = optional_number(body, "top_p")) {
        if (*top_p < 0.0 || *top_p > 1.0) { bad_request("top_p must be in [0,1]", "top_p"); }
        out.prompt.generation.sampling.top_p = *top_p;
    }
    if (const std::optional<int> max_output = optional_int(body, "max_output_tokens")) {
        if (*max_output < 0) {
            bad_request("max_output_tokens must be non-negative", "max_output_tokens");
        }
        out.requested_max_output_tokens  = *max_output;
        out.prompt.generation.max_tokens = *max_output;
    }
    return out;
}

OpenAIResponsesPromptRequest
parse_openai_responses_input_tokens_request(const Json& body, const RequestLimits& limits) {
    require_object(body);
    validate_common_top_level(body, false);
    reject_unsupported_platform_fields(body);
    if (body.contains("personality") && !body.at("personality").is_null()) {
        bad_request("personality changes prompt construction and is not supported", "personality",
                    "personality_not_supported");
    }
    return std::move(parse_prompt_fields(body, limits).prompt);
}

} // namespace ninfer::serve
#pragma once

// Dense BF16 x BF16 Tensor Core GEMM shared by pure Linear and semantic Ops.
//
//   C[M,N] = A[M,K] * B[K,N]
//
// A is row-major with contiguous K. Public NInfer activations store one contiguous
// K vector per token, which is the column-major B representation consumed by the
// row.col MMA atom. M and K come from a compiled geometry; N remains runtime.

#include "ops/common/mma.cuh"

#include <cuda_bf16.h>

#include <algorithm>
#include <cstdint>

namespace ninfer::ops::detail {

enum class Bf16MmaFragmentPipeline : std::uint8_t {
    Serial,
    PingPong,
};

enum class Bf16MmaRaster : std::uint8_t {
    TokenFast,
    RowFast,
    Grouped,
};

enum class Bf16MmaSwizzle : std::uint8_t {
    Plain,
    Xor64,
};

template <int BlockRows, int BlockCols, int BlockK, int WarpRows, int WarpCols, int PipelineStages,
          int MinBlocks, Cache WeightCache, Cache ActivationCache,
          Bf16MmaFragmentPipeline FragmentPipeline, Bf16MmaRaster Raster,
          Bf16MmaSwizzle Swizzle = Bf16MmaSwizzle::Xor64, int RasterGroupRows = 1>
struct Bf16MmaSchedule {
    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 kMinBlocks                            = MinBlocks;
    static constexpr Cache kWeightCache                        = WeightCache;
    static constexpr Cache kActivationCache                    = ActivationCache;
    static constexpr Bf16MmaFragmentPipeline kFragmentPipeline = FragmentPipeline;
    static constexpr Bf16MmaRaster kRaster                     = Raster;
    static constexpr Bf16MmaSwizzle kSwizzle                   = Swizzle;
    static constexpr int kRasterGroupRows                      = RasterGroupRows;

    static constexpr int kWarpsM         = kBlockRows / kWarpRows;
    static constexpr int kWarpsN         = kBlockCols / kWarpCols;
    static constexpr int kWarps          = kWarpsM * kWarpsN;
    static constexpr int kThreads        = kWarps * 32;
    static constexpr int kMmaRows        = kWarpRows / 16;
    static constexpr int kMmaCols        = kWarpCols / 8;
    static constexpr int kMmaK           = kBlockK / 16;
    static constexpr int kSharedElements = kPipelineStages * (kBlockRows + kBlockCols) * kBlockK;
    static constexpr int kSharedBytes = kSharedElements * static_cast<int>(sizeof(__nv_bfloat16));

    static_assert(kBlockRows > 0 && kBlockCols > 0 && kBlockK > 0);
    static_assert(kBlockRows % kWarpRows == 0 && kBlockCols % kWarpCols == 0);
    static_assert(kWarpRows % 16 == 0 && kWarpCols % 8 == 0);
    static_assert(kBlockK % 64 == 0);
    static_assert(kPipelineStages >= 2 && kPipelineStages <= 8);
    static_assert(kMinBlocks >= 1);
    static_assert(kWarps >= 1 && kThreads <= 1024);
    static_assert(kSharedBytes <= 99 * 1024);
    static_assert(kRaster != Bf16MmaRaster::Grouped || kRasterGroupRows > 0);
};

struct Bf16MmaOutputTile {
    __nv_bfloat16* data;
    std::int32_t leading_dim;
    std::int32_t parent_row_begin;

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

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

struct Bf16MmaContiguousOutput {
    __nv_bfloat16* data;
    std::int32_t leading_dim;

    __device__ __forceinline__ Bf16MmaOutputTile tile(std::int32_t) const {
        return {data, leading_dim, 0};
    }
};

template <class Schedule>
__device__ __forceinline__ int bf16_mma_shared_col(int row, int col) {
    if constexpr (Schedule::kSwizzle == Bf16MmaSwizzle::Plain) {
        return col;
    } else {
        return (col & ~63) + ((((col & 63) >> 3) ^ (row & 7)) << 3) + (col & 7);
    }
}

template <class Schedule>
__device__ __forceinline__ void
bf16_mma_tile_coordinates(std::int32_t linear, std::int32_t tiles_m, std::int32_t tiles_n,
                          std::int32_t& tile_m, std::int32_t& tile_n) {
    if constexpr (Schedule::kRaster == Bf16MmaRaster::TokenFast) {
        tile_m = linear / tiles_n;
        tile_n = linear - tile_m * tiles_n;
    } else if constexpr (Schedule::kRaster == Bf16MmaRaster::RowFast) {
        tile_n = linear / tiles_m;
        tile_m = linear - tile_n * tiles_m;
    } else {
        constexpr int group_rows      = Schedule::kRasterGroupRows;
        const std::int32_t group_span = group_rows * tiles_n;
        const std::int32_t group      = linear / group_span;
        const std::int32_t first_m    = group * group_rows;
        const std::int32_t active_m   = min(group_rows, tiles_m - first_m);
        const std::int32_t within     = linear - group * group_span;
        tile_m                        = first_m + within % active_m;
        tile_n                        = within / active_m;
    }
}

template <class Geometry, class Schedule, bool FullTokens, class Output>
__global__ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocks) void bf16_gemm_mma_kernel(
    const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ weight, Output output,
    std::int32_t tokens) {
    constexpr int M       = Geometry::kOutputRows;
    constexpr int K       = Geometry::kInputRows;
    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::kMmaK;
    constexpr int S       = Schedule::kPipelineStages;
    constexpr int WARPS_N = Schedule::kWarpsN;
    constexpr int THREADS = Schedule::kThreads;
    static_assert(M % BM == 0);
    static_assert(K % BK == 0);
    static_assert(K / BK >= S);

    extern __shared__ __align__(16) unsigned char shared_raw[];
    auto* As = reinterpret_cast<__nv_bfloat16*>(shared_raw);
    auto* Bs = As + S * BM * BK;

    const int tid  = static_cast<int>(threadIdx.x);
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int wm   = warp / WARPS_N;
    const int wn   = warp - wm * WARPS_N;
    const int gid  = lane >> 2;
    const int lid  = lane & 3;

    constexpr int tiles_m = M / BM;
    const int tiles_n     = tokens / BN + static_cast<int>(tokens % BN != 0);
    int tile_m            = 0;
    int tile_n            = 0;
    bf16_mma_tile_coordinates<Schedule>(static_cast<int>(blockIdx.x), tiles_m, tiles_n, tile_m,
                                        tile_n);
    const int m0           = tile_m * BM;
    const int n0           = tile_n * BN;
    const auto output_tile = output.tile(m0);

    float accum[MT][NT][4] = {};

    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_inputs = [&](int stage, int k_tile) {
        const int k0  = k_tile * BK;
        auto* a_stage = As + stage * BM * BK;
        auto* b_stage = Bs + stage * BN * BK;

#pragma unroll 1
        for (int item = tid; item < BM * (BK / 8); item += THREADS) {
            const int row = item / (BK / 8);
            const int k8  = item - row * (BK / 8);
            const int kk  = k8 * 8;
            cp_async<16, Schedule::kWeightCache>(
                &a_stage[row * BK + bf16_mma_shared_col<Schedule>(row, kk)],
                &weight[static_cast<std::int64_t>(m0 + row) * K + k0 + kk]);
        }

#pragma unroll 1
        for (int item = tid; item < BN * (BK / 8); item += THREADS) {
            const int col   = item / (BK / 8);
            const int k8    = item - col * (BK / 8);
            const int kk    = k8 * 8;
            auto* dst       = &b_stage[col * BK + bf16_mma_shared_col<Schedule>(col, kk)];
            const int token = n0 + col;
            if constexpr (FullTokens) {
                cp_async<16, Schedule::kActivationCache>(
                    dst, &x[static_cast<std::int64_t>(token) * K + k0 + kk]);
            } else {
                const bool valid = token < tokens;
                cp_async_zfill<16, Schedule::kActivationCache>(
                    dst, &x[static_cast<std::int64_t>(valid ? token : 0) * K + k0 + kk],
                    valid ? 16 : 0);
            }
        }
    };

    constexpr int kTiles = K / BK;
#pragma unroll
    for (int stage = 0; stage < S; ++stage) {
        stage_inputs(stage, stage);
        cp_commit();
    }

#pragma unroll 1
    for (int k_tile = 0; k_tile < kTiles; ++k_tile) {
        const int stage = k_tile % S;
        if (k_tile + S <= kTiles) {
            cp_wait<S - 1>();
        } else {
            // Once the producer stops refilling the ring, fewer than S groups remain. Drain them
            // instead of applying the steady-state wait count to the final consumer stage.
            cp_wait<0>();
        }
        __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 = wm * WM + mi * 16 + a_row_offset;
                const int col = k_step * 16 + a_col_offset;
                ldmatrix_x4(
                    a_frag[mi][0], a_frag[mi][1], a_frag[mi][2], a_frag[mi][3],
                    smem_addr(
                        &As[stage * BM * BK + row * BK + bf16_mma_shared_col<Schedule>(row, col)]));
            }
#pragma unroll
            for (int ni = 0; ni < NT; ++ni) {
                const int row = wn * WN + ni * 8 + b_inner_row;
                const int col = k_step * 16 + b_k_offset;
                ldmatrix_x2(
                    b_frag[ni][0], b_frag[ni][1],
                    smem_addr(
                        &Bs[stage * BN * BK + row * BK + bf16_mma_shared_col<Schedule>(row, col)]));
            }
        };

        if constexpr (Schedule::kFragmentPipeline == Bf16MmaFragmentPipeline::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 k_step = 0; k_step < KSUB; ++k_step) {
                const int slot = k_step & 1;
                if (k_step + 1 < KSUB) {
                    load_fragments(k_step + 1, a_frag[slot ^ 1], b_frag[slot ^ 1]);
                }
#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[slot][mi][0], a_frag[slot][mi][1],
                                 a_frag[slot][mi][2], a_frag[slot][mi][3], b_frag[slot][ni][0],
                                 b_frag[slot][ni][1]);
                    }
                }
            }
        } else {
            unsigned a_frag[MT][4];
            unsigned b_frag[NT][2];
#pragma unroll
            for (int k_step = 0; k_step < KSUB; ++k_step) {
                load_fragments(k_step, 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 next = k_tile + S;
        if (next < kTiles) {
            stage_inputs(stage, next);
            cp_commit();
        }
    }

#pragma unroll
    for (int mi = 0; mi < MT; ++mi) {
        const int row0 = m0 + wm * WM + mi * 16 + gid;
        const int row1 = row0 + 8;
#pragma unroll
        for (int ni = 0; ni < NT; ++ni) {
            const int token0   = n0 + wn * WN + ni * 8 + 2 * lid;
            const int token1   = token0 + 1;
            const float* value = accum[mi][ni];
            if constexpr (FullTokens) {
                output_tile.store(row0, token0, value[0]);
                output_tile.store(row0, token1, value[1]);
                output_tile.store(row1, token0, value[2]);
                output_tile.store(row1, token1, value[3]);
            } else {
                if (token0 < tokens) {
                    output_tile.store(row0, token0, value[0]);
                    output_tile.store(row1, token0, value[2]);
                }
                if (token1 < tokens) {
                    output_tile.store(row0, token1, value[1]);
                    output_tile.store(row1, token1, value[3]);
                }
            }
        }
    }
}

} // namespace ninfer::ops::detail
#include "ops/linear_add/fp8/fp8_linear_add_plan.h"

#include "core/device.h"
#include "ops/linear/fp8/fp8_config.h"
#include "ops/linear/fp8/fp8_output.cuh"
#include "ops/linear/fp8/fp8_small_t.cuh"
#include "ops/linear_add/fp8/fp8_linear_add_epilogue.cuh"

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

namespace ninfer::ops::detail {
namespace {

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

template <class Geometry, int ActiveTokens>
struct Fp8LinearAddSmallTProductionSchedule;

// LinearAdd's residual read and fused epilogue shift the register/occupancy crossovers relative to
// contiguous Linear. These are the measured RTX 5090 ranges for [5120,6144].
template <int ActiveTokens>
struct Fp8LinearAddSmallTProductionSchedule<Fp8Residual6144Geometry, ActiveTokens> {
    static_assert(ActiveTokens >= kFp8FirstSmallT && ActiveTokens <= kFp8LastSmallT);
    static constexpr int kWarpsPerCta =
        (ActiveTokens >= 15 && ActiveTokens <= 19) || ActiveTokens == 22 ? 4 : 8;
    static constexpr int kRowsPerWarp   = ActiveTokens <= 5 ? 1 : 2;
    static constexpr int kValuesPerLane = ActiveTokens >= 20 && ActiveTokens <= 23 ? 8 : 16;
    static constexpr int kTokenTile     = ActiveTokens == 24 ? 12 : ActiveTokens;
    static constexpr auto kCodeCache =
        (ActiveTokens >= 12 && ActiveTokens <= 14) || ActiveTokens == 23 ? Fp8CodeCache::Streaming
                                                                         : Fp8CodeCache::Default;
    static constexpr auto kBlockOrder = ActiveTokens == 24
                                            ? Fp8SmallTBlockOrder::TokenTilesContiguous
                                            : Fp8SmallTBlockOrder::RowsContiguous;
    using Type =
        Fp8SmallTSchedule<kWarpsPerCta, kRowsPerWarp, kValuesPerLane, kTokenTile, 1,
                          Fp8SmallTActivationAccess::TokenPacked, kCodeCache, 1, kBlockOrder, 1>;
};

// The longer K of [5120,17408] favors streaming code loads around its two register crossovers;
// keeping the default v16 schedule at T=9..16 avoids an isolated fast point followed by a cliff.
template <int ActiveTokens>
struct Fp8LinearAddSmallTProductionSchedule<Fp8Residual17408Geometry, ActiveTokens> {
    static_assert(ActiveTokens >= kFp8FirstSmallT && ActiveTokens <= kFp8LastSmallT);
    static constexpr int kWarpsPerCta   = ActiveTokens >= 21 && ActiveTokens <= 22 ? 4 : 8;
    static constexpr int kRowsPerWarp   = ActiveTokens <= 5 ? 1 : 2;
    static constexpr int kValuesPerLane = ActiveTokens >= 20 && ActiveTokens <= 23 ? 8 : 16;
    static constexpr int kTokenTile     = ActiveTokens == 24 ? 12 : ActiveTokens;
    static constexpr auto kCodeCache    = (ActiveTokens >= 6 && ActiveTokens <= 8) ||
                                               (ActiveTokens >= 17 && ActiveTokens <= 20) ||
                                               ActiveTokens == 23
                                              ? Fp8CodeCache::Streaming
                                              : Fp8CodeCache::Default;
    static constexpr auto kBlockOrder   = ActiveTokens == 24
                                              ? Fp8SmallTBlockOrder::TokenTilesContiguous
                                              : Fp8SmallTBlockOrder::RowsContiguous;
    using Type =
        Fp8SmallTSchedule<kWarpsPerCta, kRowsPerWarp, kValuesPerLane, kTokenTile, 1,
                          Fp8SmallTActivationAccess::TokenPacked, kCodeCache, 1, kBlockOrder, 1>;
};

template <class Geometry, int ActiveTokens>
void launch_exact(const Tensor& x, const Weight& weight, Tensor& residual, cudaStream_t stream) {
    using Schedule = typename Fp8LinearAddSmallTProductionSchedule<Geometry, ActiveTokens>::Type;
    constexpr int kTokenTiles = (ActiveTokens + Schedule::kTokenTile - 1) / Schedule::kTokenTile;
    constexpr int kBlocks     = (Geometry::kOutputRows / Schedule::kRowsPerCta) * kTokenTiles;
    auto* output              = static_cast<__nv_bfloat16*>(residual.data);
    fp8_small_t_kernel<Geometry, ActiveTokens, Schedule>
        <<<kBlocks, Schedule::kThreads, 0, stream>>>(
            static_cast<const __nv_bfloat16*>(x.data),
            static_cast<const std::uint8_t*>(weight.qdata),
            static_cast<const __nv_bfloat16*>(weight.scales),
            Fp8ContiguousOutput{output, Geometry::kOutputRows},
            Fp8AddResidualEpilogue{output, Geometry::kOutputRows});
    CUDA_CHECK(cudaGetLastError());
}

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

template <class Geometry>
const auto& launchers() {
    static constexpr auto kLaunchers =
        make_launchers<Geometry>(std::make_index_sequence<kFp8LastSmallT - kFp8FirstSmallT + 1>{});
    return kLaunchers;
}

} // namespace

void fp8_linear_add_small_t_launch(const Tensor& x, const Weight& weight, Tensor& residual,
                                   cudaStream_t stream) {
    if (x.ne[1] < kFp8FirstSmallT || x.ne[1] > kFp8LastSmallT) {
        throw std::invalid_argument("fp8 linear_add small-T: unsupported T");
    }
    const std::size_t index = static_cast<std::size_t>(x.ne[1] - kFp8FirstSmallT);
    switch (resolve_fp8_problem(weight.n, weight.k)) {
    case Fp8Problem::Residual6144:
        launchers<Fp8Residual6144Geometry>()[index](x, weight, residual, stream);
        return;
    case Fp8Problem::Residual17408:
        launchers<Fp8Residual17408Geometry>()[index](x, weight, residual, stream);
        return;
    case Fp8Problem::AttnInput:
    case Fp8Problem::GdnInput:
    case Fp8Problem::MlpGateUp:
    case Fp8Problem::Vocabulary:
        break;
    }
    throw std::invalid_argument("fp8 linear_add small-T: unsupported problem");
}

} // namespace ninfer::ops::detail
#include "ops/attn_input_proj/fp8/fp8_attn_input_plan.h"

#include "core/device.h"
#include "ops/attn_input_proj/fp8/fp8_attn_input_output.cuh"
#include "ops/linear/fp8/fp8_a8_schedule.cuh"
#include "ops/linear/fp8/fp8_config.h"
#include "ops/linear/fp8/fp8_output.cuh"

#include <cuda_bf16.h>

#include <cstdint>

namespace ninfer::ops::detail {
namespace {

using Geometry = Fp8AttnInputGeometry;
using Schedule = typename Fp8LinearA8ProductionSchedule<Geometry>::Type;

static_assert((kFp8AttnInputQueryRows % Schedule::kBlockRows) == 0);
static_assert((kFp8AttnInputKeyRows % Schedule::kBlockRows) == 0);
static_assert((kFp8AttnInputGateRows % Schedule::kBlockRows) == 0);

template <bool FullTokens>
void launch_mma(const Weight& weight, Tensor& q, Tensor& gate, Tensor& k, Tensor& v,
                Fp8A8Workspace workspace, std::int32_t tokens, cudaStream_t stream) {
    constexpr int kRowTiles = Geometry::kOutputRows / Schedule::kBlockRows;
    const int token_tiles   = (tokens + Schedule::kBlockTokens - 1) / Schedule::kBlockTokens;
    const int blocks        = kRowTiles * token_tiles;
    const Fp8AttentionInputOutput 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),
    };

    if constexpr (Schedule::kSharedBytes > 48 * 1024) {
        static const cudaError_t attribute = cudaFuncSetAttribute(
            fp8_mma_kernel<Geometry, Schedule, FullTokens, Fp8IdentityEpilogue,
                           Fp8AttentionInputOutput>,
            cudaFuncAttributeMaxDynamicSharedMemorySize, Schedule::kSharedBytes);
        CUDA_CHECK(attribute);
    }
    fp8_mma_kernel<Geometry, Schedule, FullTokens>
        <<<blocks, Schedule::kThreads, Schedule::kSharedBytes, stream>>>(
            workspace.codes, workspace.scales, static_cast<const std::uint8_t*>(weight.qdata),
            static_cast<const __nv_bfloat16*>(weight.scales), tokens, Fp8IdentityEpilogue{},
            output);
    CUDA_CHECK(cudaGetLastError());
}

} // namespace

void fp8_attn_input_a8_launch(const Tensor& x, const Weight& weight, Tensor& q, Tensor& gate,
                              Tensor& k, Tensor& v, Fp8A8Workspace workspace, cudaStream_t stream) {
    launch_fp8_a8_quantize(x, weight, workspace, stream);
    if ((x.ne[1] % Schedule::kBlockTokens) == 0) {
        launch_mma<true>(weight, q, gate, k, v, workspace, x.ne[1], stream);
    } else {
        launch_mma<false>(weight, q, gate, k, v, workspace, x.ne[1], stream);
    }
}

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

#include "core/tensor.h"

#include <cuda_runtime.h> // cudaStream_t

namespace ninfer::ops {

/**
 * Elementwise residual update:
 *
 *   ideal[i] = x[i] + y[i].
 *
 * `y` and `x` are non-overlapping, same-shaped contiguous BF16 tensors. The Op updates all of x
 * in place and leaves y unchanged. 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. The Op uses no workspace or other persistent state.
 */
void residual_add(const Tensor& y, Tensor& x, cudaStream_t stream);

} // namespace ninfer::ops
#include "ops/gdn_gating_proj/bf16/bf16_gdn_gating_proj_plan.h"

#include "ninfer/ops/rmsnorm.h"

#include <algorithm>
#include <array>
#include <limits>
#include <stdexcept>

namespace ninfer::ops::detail {
namespace {

inline constexpr std::int32_t kAnyCols = std::numeric_limits<std::int32_t>::max();

struct ColsSet {
    std::int32_t first;
    std::int32_t last;

    constexpr bool contains(std::int32_t cols) const noexcept {
        return cols >= first && cols <= last;
    }
};

struct RouteSpec {
    ColsSet cols;
    Bf16GdnGatingScheduleId schedule;
};

constexpr std::array<RouteSpec, 6> k27Routes{{
    {{1, 1}, Bf16GdnGatingScheduleId::GemvPairedRows},
    {{2, 8}, Bf16GdnGatingScheduleId::SmallTSplit10},
    // As token tiles double, halve SplitK. This keeps the cooperative grid near 192 CTAs instead
    // of making T a launch limit. Once the unsplit grid has enough independent work, it also
    // removes the cooperative-residency constraint.
    {{9, 1024}, Bf16GdnGatingScheduleId::MmaCooperativeSplit8},
    {{1025, 2048}, Bf16GdnGatingScheduleId::MmaCooperativeSplit4},
    {{2049, 4096}, Bf16GdnGatingScheduleId::MmaCooperativeSplit2},
    {{4097, kAnyCols}, Bf16GdnGatingScheduleId::MmaUnsplit},
}};

constexpr std::array<RouteSpec, 5> k35Routes{{
    // The same progression keeps the long-range cooperative routes near 256 CTAs.
    {{1, 127}, Bf16GdnGatingScheduleId::MmaCooperativeSplit16},
    {{128, 1024}, Bf16GdnGatingScheduleId::MmaCooperativeSplit8},
    {{1025, 2048}, Bf16GdnGatingScheduleId::MmaCooperativeSplit4},
    {{2049, 4096}, Bf16GdnGatingScheduleId::MmaCooperativeSplit2},
    {{4097, kAnyCols}, Bf16GdnGatingScheduleId::MmaUnsplit},
}};

template <std::size_t N>
constexpr bool catalog_is_closed(const std::array<RouteSpec, N>& routes,
                                 std::int32_t last) noexcept {
    std::int64_t expected = 1;
    for (const RouteSpec& route : routes) {
        if (route.cols.first != expected || route.cols.last < route.cols.first) { return false; }
        expected = static_cast<std::int64_t>(route.cols.last) + 1;
    }
    return routes.back().cols.last == last && expected == static_cast<std::int64_t>(last) + 1;
}

static_assert(catalog_is_closed(k27Routes, kAnyCols));
static_assert(catalog_is_closed(k35Routes, kAnyCols));

bool is_27(const Bf16GdnGatingProblem& problem) noexcept {
    return problem.heads == 48 && problem.input_rows == 5120;
}

bool is_35(const Bf16GdnGatingProblem& problem) noexcept {
    return problem.heads == 32 && problem.input_rows == 2048;
}

bool schedule_uses_mma(Bf16GdnGatingScheduleId schedule) noexcept {
    switch (schedule) {
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit32:
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit16:
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit8:
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit4:
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit2:
    case Bf16GdnGatingScheduleId::MmaUnsplit:
        return true;
    case Bf16GdnGatingScheduleId::GemvPairedRows:
    case Bf16GdnGatingScheduleId::SmallTSplit10:
    case Bf16GdnGatingScheduleId::SimtWarpRowC4:
    case Bf16GdnGatingScheduleId::SimtWarpRowC8:
        return false;
    }
    return false;
}

std::int32_t mma_tile_cols(const Bf16GdnGatingProblem& problem) noexcept {
    return is_35(problem) ? 64 : 128;
}

std::int32_t schedule_split_k(Bf16GdnGatingScheduleId schedule) {
    switch (schedule) {
    case Bf16GdnGatingScheduleId::SmallTSplit10:
        return 10;
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit32:
        return 32;
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit16:
        return 16;
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit8:
        return 8;
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit4:
        return 4;
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit2:
        return 2;
    case Bf16GdnGatingScheduleId::GemvPairedRows:
    case Bf16GdnGatingScheduleId::SimtWarpRowC4:
    case Bf16GdnGatingScheduleId::SimtWarpRowC8:
    case Bf16GdnGatingScheduleId::MmaUnsplit:
        return 1;
    }
    throw std::logic_error("BF16 GDN gating: unknown schedule");
}

bool cooperative_grid_is_resident(Bf16GdnGatingScheduleId schedule, std::int32_t cols,
                                  std::int32_t tile_cols, std::int32_t row_tiles,
                                  std::int32_t resident_ctas) noexcept {
    const std::int64_t column_tiles = (static_cast<std::int64_t>(cols) + tile_cols - 1) / tile_cols;
    const std::int64_t grid_ctas =
        column_tiles * row_tiles * static_cast<std::int64_t>(schedule_split_k(schedule));
    return grid_ctas <= resident_ctas;
}

bool cooperative_27_grid_is_resident(Bf16GdnGatingScheduleId schedule, std::int32_t cols) noexcept {
    // BN128 uses 40 KiB of dynamic shared memory. Split8 uses 71 registers with 256 threads;
    // split4/2 use 62 registers with 512 threads. Each specialization admits two CTAs/SM, hence
    // 340 resident CTAs device-wide. There are three 16-row tiles per token tile.
    return cooperative_grid_is_resident(schedule, cols, 128, 3, 340);
}

bool cooperative_35_grid_is_resident(Bf16GdnGatingScheduleId schedule, std::int32_t cols) noexcept {
    // BN64 uses 24 KiB of dynamic shared memory and two 16-row tiles. With the registered CUDA
    // 13.1/sm_120a build, split32 uses 91/93 registers per thread and admits two CTAs/SM;
    // split16/8/4/2 use at most 62 registers and admit four CTAs/SM. Across 170 SMs the
    // device-wide limits are 340 and 680 CTAs respectively.
    const std::int32_t resident_ctas =
        schedule == Bf16GdnGatingScheduleId::MmaCooperativeSplit32 ? 340 : 680;
    return cooperative_grid_is_resident(schedule, cols, 64, 2, resident_ctas);
}

bool candidate_is_legal(Bf16GdnGatingScheduleId schedule,
                        const Bf16GdnGatingProblem& problem) noexcept {
    if (!bf16_gdn_gating_admits(problem)) { return false; }
    if (is_27(problem)) {
        switch (schedule) {
        case Bf16GdnGatingScheduleId::GemvPairedRows:
            return problem.cols == 1;
        case Bf16GdnGatingScheduleId::SmallTSplit10:
            return problem.cols >= 2 && problem.cols <= 8;
        case Bf16GdnGatingScheduleId::MmaCooperativeSplit8:
        case Bf16GdnGatingScheduleId::MmaCooperativeSplit4:
        case Bf16GdnGatingScheduleId::MmaCooperativeSplit2:
            return cooperative_27_grid_is_resident(schedule, problem.cols);
        case Bf16GdnGatingScheduleId::MmaUnsplit:
            return true;
        case Bf16GdnGatingScheduleId::SimtWarpRowC4:
        case Bf16GdnGatingScheduleId::SimtWarpRowC8:
        case Bf16GdnGatingScheduleId::MmaCooperativeSplit32:
        case Bf16GdnGatingScheduleId::MmaCooperativeSplit16:
            return false;
        }
    }

    switch (schedule) {
    case Bf16GdnGatingScheduleId::SimtWarpRowC4:
        return problem.cols <= 4 * 65'535;
    case Bf16GdnGatingScheduleId::SimtWarpRowC8:
        return problem.cols <= 8 * 65'535;
    case Bf16GdnGatingScheduleId::MmaUnsplit:
        return true;
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit32:
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit16:
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit8:
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit4:
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit2:
        return cooperative_35_grid_is_resident(schedule, problem.cols);
    case Bf16GdnGatingScheduleId::GemvPairedRows:
    case Bf16GdnGatingScheduleId::SmallTSplit10:
        return false;
    }
    return false;
}

std::size_t checked_partial_bytes(std::int32_t heads, std::int32_t split_k, std::int32_t cols) {
    const std::size_t logical_rows = static_cast<std::size_t>(2 * heads);
    const std::size_t split        = static_cast<std::size_t>(split_k);
    const std::size_t tokens       = static_cast<std::size_t>(cols);
    if (tokens > std::numeric_limits<std::size_t>::max() / logical_rows ||
        split > std::numeric_limits<std::size_t>::max() / (tokens * logical_rows)) {
        throw std::overflow_error("BF16 GDN gating workspace element count overflows size_t");
    }
    const std::size_t elements = split * tokens * logical_rows;
    if (elements > std::numeric_limits<std::size_t>::max() / sizeof(float)) {
        throw std::overflow_error("BF16 GDN gating workspace byte count overflows size_t");
    }
    return elements * sizeof(float);
}

void execute_resolved(const Bf16GdnGatingPlan& plan, const Bf16GdnGatingProblem& problem,
                      const Tensor& x, const Weight& a_weight, const Weight& b_weight,
                      const Tensor& A_log, const Tensor& dt_bias, WorkspaceArena& ws, Tensor& g,
                      Tensor& beta, cudaStream_t stream) {
    auto scratch_scope = ws.scope();
    DeviceSpan scratch{};
    if (plan.workspace_bytes != 0) { scratch = ws.alloc_bytes(plan.workspace_bytes); }

    switch (plan.schedule) {
    case Bf16GdnGatingScheduleId::GemvPairedRows:
        bf16_gdn_gating_proj_gemv_launch(x, a_weight, b_weight, A_log, dt_bias, g, beta, stream);
        return;
    case Bf16GdnGatingScheduleId::SmallTSplit10:
        bf16_gdn_gating_proj_small_t_split10_launch(x, a_weight, b_weight, A_log, dt_bias,
                                                    scratch.data, scratch.bytes, g, beta, stream);
        return;
    case Bf16GdnGatingScheduleId::SimtWarpRowC4:
        bf16_gdn_gating_proj_35_simt_c4_launch(x, a_weight, b_weight, A_log, dt_bias, g, beta,
                                               stream);
        return;
    case Bf16GdnGatingScheduleId::SimtWarpRowC8:
        bf16_gdn_gating_proj_35_simt_c8_launch(x, a_weight, b_weight, A_log, dt_bias, g, beta,
                                               stream);
        return;
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit32:
        bf16_gdn_gating_proj_35_mma_split32_launch(plan.token_variant, x, a_weight, b_weight, A_log,
                                                   dt_bias, scratch.data, g, beta, stream);
        return;
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit16:
        bf16_gdn_gating_proj_35_mma_split16_launch(plan.token_variant, x, a_weight, b_weight, A_log,
                                                   dt_bias, scratch.data, g, beta, stream);
        return;
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit8:
        if (is_35(problem)) {
            bf16_gdn_gating_proj_35_mma_split8_launch(plan.token_variant, x, a_weight, b_weight,
                                                      A_log, dt_bias, scratch.data, g, beta,
                                                      stream);
        } else {
            bf16_gdn_gating_proj_mma_split8_launch(plan.token_variant, x, a_weight, b_weight, A_log,
                                                   dt_bias, scratch.data, g, beta, stream);
        }
        return;
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit4:
        if (is_35(problem)) {
            bf16_gdn_gating_proj_35_mma_split4_launch(plan.token_variant, x, a_weight, b_weight,
                                                      A_log, dt_bias, scratch.data, g, beta,
                                                      stream);
        } else {
            bf16_gdn_gating_proj_mma_split4_launch(plan.token_variant, x, a_weight, b_weight, A_log,
                                                   dt_bias, scratch.data, g, beta, stream);
        }
        return;
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit2:
        if (is_35(problem)) {
            bf16_gdn_gating_proj_35_mma_split2_launch(plan.token_variant, x, a_weight, b_weight,
                                                      A_log, dt_bias, scratch.data, g, beta,
                                                      stream);
        } else {
            bf16_gdn_gating_proj_mma_split2_launch(plan.token_variant, x, a_weight, b_weight, A_log,
                                                   dt_bias, scratch.data, g, beta, stream);
        }
        return;
    case Bf16GdnGatingScheduleId::MmaUnsplit:
        if (is_35(problem)) {
            bf16_gdn_gating_proj_35_mma_unsplit_launch(plan.token_variant, x, a_weight, b_weight,
                                                       A_log, dt_bias, g, beta, stream);
        } else {
            bf16_gdn_gating_proj_mma_unsplit_launch(plan.token_variant, x, a_weight, b_weight,
                                                    A_log, dt_bias, g, beta, stream);
        }
        return;
    }
    throw std::logic_error("BF16 GDN gating: unknown schedule");
}

template <std::size_t N>
std::size_t route_capacity(const std::array<RouteSpec, N>& routes, const Bf16GdnGatingProblem& base,
                           std::int32_t min_cols, std::int32_t max_cols) {
    std::size_t maximum = 0;
    for (const RouteSpec& route : routes) {
        if (route.cols.last < min_cols || route.cols.first > max_cols) { continue; }
        const std::int32_t endpoint = std::min(route.cols.last, max_cols);
        maximum                     = std::max(
            maximum,
            bf16_gdn_gating_resolve_plan({base.heads, base.input_rows, endpoint}).workspace_bytes);
    }
    return maximum;
}

} // namespace

const char* bf16_gdn_gating_schedule_name(Bf16GdnGatingScheduleId schedule) noexcept {
    switch (schedule) {
    case Bf16GdnGatingScheduleId::GemvPairedRows:
        return "gdn_gating_proj.bf16.gemv.paired_rows";
    case Bf16GdnGatingScheduleId::SmallTSplit10:
        return "gdn_gating_proj.bf16.small_t.split10";
    case Bf16GdnGatingScheduleId::SimtWarpRowC4:
        return "gdn_gating_proj.bf16.simt.warp_row.c4";
    case Bf16GdnGatingScheduleId::SimtWarpRowC8:
        return "gdn_gating_proj.bf16.simt.warp_row.c8";
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit32:
        return "gdn_gating_proj.bf16.mma.cooperative_split32";
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit16:
        return "gdn_gating_proj.bf16.mma.cooperative_split16";
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit8:
        return "gdn_gating_proj.bf16.mma.cooperative_split8";
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit4:
        return "gdn_gating_proj.bf16.mma.cooperative_split4";
    case Bf16GdnGatingScheduleId::MmaCooperativeSplit2:
        return "gdn_gating_proj.bf16.mma.cooperative_split2";
    case Bf16GdnGatingScheduleId::MmaUnsplit:
        return "gdn_gating_proj.bf16.mma.unsplit";
    }
    return "gdn_gating_proj.bf16.unknown";
}

const char* bf16_gdn_norm_gating_schedule_name(Bf16GdnNormGatingScheduleId schedule) noexcept {
    switch (schedule) {
    case Bf16GdnNormGatingScheduleId::Composed:
        return "gdn_norm_gating_proj.bf16.composed";
    case Bf16GdnNormGatingScheduleId::MmaCooperativeSplit32:
        return "gdn_norm_gating_proj.bf16.mma.cooperative_split32";
    }
    return "gdn_norm_gating_proj.bf16.unknown";
}

bool bf16_gdn_gating_admits(const Bf16GdnGatingProblem& problem) noexcept {
    if (problem.cols < 1) { return false; }
    return is_27(problem) || is_35(problem);
}

Bf16GdnGatingPlan bf16_gdn_gating_resolve_candidate(Bf16GdnGatingScheduleId schedule,
                                                    const Bf16GdnGatingProblem& problem) {
    if (!candidate_is_legal(schedule, problem)) {
        throw std::invalid_argument("BF16 GDN gating: candidate is not legal for exact problem");
    }
    const bool mma                          = schedule_uses_mma(schedule);
    const Bf16GdnGatingTokenVariant variant = !mma ? Bf16GdnGatingTokenVariant::None
                                                   : ((problem.cols % mma_tile_cols(problem)) == 0
                                                          ? Bf16GdnGatingTokenVariant::Full
                                                          : Bf16GdnGatingTokenVariant::Predicated);
    const std::int32_t split_k              = schedule_split_k(schedule);
    const std::size_t workspace =
        split_k > 1 ? checked_partial_bytes(problem.heads, split_k, problem.cols) : 0;
    return {schedule, variant, workspace};
}

Bf16GdnGatingPlan bf16_gdn_gating_resolve_plan(const Bf16GdnGatingProblem& problem) {
    if (!bf16_gdn_gating_admits(problem)) {
        throw std::invalid_argument(
            "BF16 GDN gating: exact problem or column count is not admitted");
    }
    if (is_27(problem)) {
        for (const RouteSpec& route : k27Routes) {
            if (route.cols.contains(problem.cols)) {
                return bf16_gdn_gating_resolve_candidate(route.schedule, problem);
            }
        }
    } else {
        for (const RouteSpec& route : k35Routes) {
            if (route.cols.contains(problem.cols)) {
                return bf16_gdn_gating_resolve_candidate(route.schedule, problem);
            }
        }
    }
    throw std::logic_error("BF16 GDN gating: admitted problem has no covering route");
}

std::size_t bf16_gdn_gating_capacity_workspace_bytes(std::int32_t heads, std::int32_t input_rows,
                                                     std::int32_t min_cols, std::int32_t max_cols) {
    if (min_cols <= 0 || max_cols < min_cols) {
        throw std::invalid_argument("BF16 GDN gating: invalid column interval");
    }
    const Bf16GdnGatingProblem base{heads, input_rows, 1};
    (void)bf16_gdn_gating_resolve_plan({heads, input_rows, min_cols});
    (void)bf16_gdn_gating_resolve_plan({heads, input_rows, max_cols});
    return is_27(base) ? route_capacity(k27Routes, base, min_cols, max_cols)
                       : route_capacity(k35Routes, base, min_cols, max_cols);
}

Bf16GdnNormGatingPlan bf16_gdn_norm_gating_resolve_plan(const Bf16GdnGatingProblem& problem) {
    Bf16GdnGatingPlan control            = bf16_gdn_gating_resolve_plan(problem);
    Bf16GdnNormGatingScheduleId schedule = Bf16GdnNormGatingScheduleId::Composed;
    std::int32_t norm_splits             = 0;
    if (is_35(problem) && problem.cols <= 16) {
        control  = bf16_gdn_gating_resolve_candidate(Bf16GdnGatingScheduleId::MmaCooperativeSplit32,
                                                     problem);
        schedule = Bf16GdnNormGatingScheduleId::MmaCooperativeSplit32;
        norm_splits = 32;
    }
    const std::size_t norm_partial_bytes =
        static_cast<std::size_t>(norm_splits) * problem.cols * sizeof(float);
    return {schedule, control, control.workspace_bytes + norm_partial_bytes};
}

std::size_t bf16_gdn_norm_gating_capacity_workspace_bytes(std::int32_t heads,
                                                          std::int32_t input_rows,
                                                          std::int32_t min_cols,
                                                          std::int32_t max_cols) {
    std::size_t maximum =
        bf16_gdn_gating_capacity_workspace_bytes(heads, input_rows, min_cols, max_cols);
    if (heads == 32 && input_rows == 2048 && min_cols <= 16) {
        const std::int32_t fused_cols = std::min<std::int32_t>(max_cols, 16);
        maximum                       = std::max(
            maximum,
            bf16_gdn_norm_gating_resolve_plan({heads, input_rows, fused_cols}).workspace_bytes);
    }
    return maximum;
}

void bf16_gdn_gating_execute_plan(const Bf16GdnGatingPlan& plan, const Tensor& x,
                                  const Weight& a_weight, const Weight& b_weight,
                                  const Tensor& A_log, const Tensor& dt_bias, WorkspaceArena& ws,
                                  Tensor& g, Tensor& beta, cudaStream_t stream) {
    const Bf16GdnGatingProblem problem{g.ne[0], x.ne[0], x.ne[1]};
    const Bf16GdnGatingPlan resolved = bf16_gdn_gating_resolve_plan(problem);
    if (resolved.schedule != plan.schedule || resolved.token_variant != plan.token_variant ||
        resolved.workspace_bytes != plan.workspace_bytes) {
        throw std::invalid_argument("BF16 GDN gating: plan does not match the exact problem");
    }
    execute_resolved(plan, problem, x, a_weight, b_weight, A_log, dt_bias, ws, g, beta, stream);
}

void bf16_gdn_gating_execute_candidate(Bf16GdnGatingScheduleId schedule, const Tensor& x,
                                       const Weight& a_weight, const Weight& b_weight,
                                       const Tensor& A_log, const Tensor& dt_bias,
                                       WorkspaceArena& ws, Tensor& g, Tensor& beta,
                                       cudaStream_t stream) {
    const Bf16GdnGatingProblem problem{g.ne[0], x.ne[0], x.ne[1]};
    const Bf16GdnGatingPlan plan = bf16_gdn_gating_resolve_candidate(schedule, problem);
    execute_resolved(plan, problem, x, a_weight, b_weight, A_log, dt_bias, ws, g, beta, stream);
}

void bf16_gdn_gating_dispatch(const Tensor& x, const Weight& a_weight, const Weight& b_weight,
                              const Tensor& A_log, const Tensor& dt_bias, WorkspaceArena& ws,
                              Tensor& g, Tensor& beta, cudaStream_t stream) {
    const Bf16GdnGatingPlan plan = bf16_gdn_gating_resolve_plan({g.ne[0], x.ne[0], x.ne[1]});
    bf16_gdn_gating_execute_plan(plan, x, a_weight, b_weight, A_log, dt_bias, ws, g, beta, stream);
}

void bf16_gdn_norm_gating_dispatch(const Tensor& x, const Tensor& norm_weight, float eps, Tensor& h,
                                   const Weight& a_weight, const Weight& b_weight,
                                   const Tensor& A_log, const Tensor& dt_bias, WorkspaceArena& ws,
                                   Tensor& g, Tensor& beta, cudaStream_t stream) {
    const Bf16GdnGatingProblem problem{g.ne[0], x.ne[0], x.ne[1]};
    const Bf16GdnNormGatingPlan plan = bf16_gdn_norm_gating_resolve_plan(problem);
    if (plan.schedule == Bf16GdnNormGatingScheduleId::Composed) {
        rmsnorm(x, norm_weight, eps, true, h, stream);
        execute_resolved(plan.control, problem, h, a_weight, b_weight, A_log, dt_bias, ws, g, beta,
                         stream);
        return;
    }

    auto scratch_scope = ws.scope();
    DeviceSpan scratch{};
    if (plan.workspace_bytes != 0) { scratch = ws.alloc_bytes(plan.workspace_bytes); }
    bf16_gdn_norm_gating_proj_35_mma_split32_launch(plan.control.token_variant, x, norm_weight, eps,
                                                    h, a_weight, b_weight, A_log, dt_bias,
                                                    scratch.data, g, beta, stream);
}

} // namespace ninfer::ops::detail
#include <ninfer/targets/qwen3_6/round_state.h>

#include <algorithm>
#include <limits>
#include <stdexcept>
#include <type_traits>

namespace ninfer::targets::qwen3_6 {
namespace {

constexpr std::size_t kArenaAlign = 256;

TensorRegion add_tensor(LayoutBuilder& builder, DType dtype,
                        std::initializer_list<std::int32_t> shape, const char* label) {
    return builder.add_tensor(dtype, shape, kArenaAlign, label);
}

std::int32_t checked_i32(std::uint64_t value, const char* label) {
    if (value > static_cast<std::uint64_t>(std::numeric_limits<std::int32_t>::max())) {
        throw std::invalid_argument(label);
    }
    return static_cast<std::int32_t>(value);
}

void validate_spec(const RoundStateSpec& spec) {
    if (spec.enable_mtp && spec.enable_dflash) {
        throw std::invalid_argument("RoundState speculative extensions are mutually exclusive");
    }
    if (spec.hidden <= 0) { throw std::invalid_argument("RoundState hidden must be positive"); }
    if (spec.output_rows <= 0) {
        throw std::invalid_argument("RoundState output_rows must be positive");
    }
    if (spec.enable_mtp && spec.draft_window == 0) {
        throw std::invalid_argument("RoundState cannot enable MTP with an empty draft window");
    }
    if (spec.enable_mtp && spec.draft_window > kMtpDecodeMaximumDrafts) {
        throw std::invalid_argument("RoundState MTP draft window exceeds the decode frame domain");
    }
    if (spec.enable_dflash &&
        (spec.draft_window == 0 || spec.draft_window > kDFlashDecodeMaximumDrafts)) {
        throw std::invalid_argument(
            "RoundState DFlash draft window exceeds the decode frame domain");
    }
    if (spec.batch_capacity == 0 || spec.batch_capacity > kMaximumConcurrency) {
        throw std::invalid_argument("RoundState batch capacity must be in [1,8]");
    }
    (void)checked_i32(static_cast<std::uint64_t>(spec.draft_window) + 1ULL,
                      "RoundState draft window exceeds int32");
}

} // namespace

RoundStateLayout begin_round_state_layout(LayoutBuilder& builder, const RoundStateSpec& spec) {
    validate_spec(spec);
    RoundStateLayout layout;
    layout.spec = spec;
    if (!spec.enable_mtp && !spec.enable_dflash) {
        OrdinaryDecodeStateLayout& ordinary = layout.ordinary.emplace();
        ordinary.ingress =
            builder.add(sizeof(OrdinaryDecodeIngress), 256, "ordinary decode ingress");
        ordinary.egress = builder.add(sizeof(OrdinaryDecodeEgress), 256, "ordinary decode egress");
        ordinary.logits = add_tensor(
            builder, DType::BF16,
            {spec.output_rows, checked_i32(spec.batch_capacity, "RoundState batch capacity")},
            "ordinary decode logits");
        ordinary.hidden =
            add_tensor(builder, DType::BF16,
                       {spec.hidden, checked_i32(spec.batch_capacity, "RoundState batch capacity")},
                       "ordinary decode hidden");
    }
    layout.token      = add_tensor(builder, DType::I32, {1}, "step token");
    layout.pos        = add_tensor(builder, DType::I32, {1}, "step position");
    layout.rope_pos   = add_tensor(builder, DType::I32, {1}, "step rope position");
    layout.rope_delta = add_tensor(builder, DType::I32, {1}, "step rope delta");
    layout.logits     = add_tensor(builder, DType::BF16, {spec.output_rows, 1}, "step logits");
    layout.text_kv_table_row    = add_tensor(builder, DType::I32, {1}, "step Text KV table row");
    layout.backend_kv_table_row = add_tensor(builder, DType::I32, {1}, "step backend KV table row");
    return layout;
}

OrdinaryDecodeState::OrdinaryDecodeState(DeviceSpan backing,
                                         const OrdinaryDecodeStateLayout& layout,
                                         std::uint32_t batch_capacity) {
    if (batch_capacity == 0 || batch_capacity > kMaximumConcurrency) {
        throw std::invalid_argument("ordinary decode batch capacity must be in [1,8]");
    }
    static_assert(std::is_standard_layout_v<OrdinaryDecodeIngress>);
    static_assert(std::is_standard_layout_v<OrdinaryDecodeEgress>);

    ingress                   = layout.ingress.bind(backing);
    egress                    = layout.egress.bind(backing);
    const auto count          = static_cast<std::int32_t>(batch_capacity);
    const auto ingress_tensor = [&](std::size_t offset, DType dtype) {
        return Tensor(static_cast<unsigned char*>(ingress.data) + offset, dtype, {count});
    };
    tokens          = ingress_tensor(offsetof(OrdinaryDecodeIngress, tokens), DType::I32);
    cache_positions = ingress_tensor(offsetof(OrdinaryDecodeIngress, cache_positions), DType::I32);
    rope_positions  = ingress_tensor(offsetof(OrdinaryDecodeIngress, rope_positions), DType::I32);
    text_kv_table_rows =
        ingress_tensor(offsetof(OrdinaryDecodeIngress, text_kv_table_rows), DType::I32);
    state_source_slots =
        ingress_tensor(offsetof(OrdinaryDecodeIngress, state_source_slots), DType::I32);
    state_destination_slots =
        ingress_tensor(offsetof(OrdinaryDecodeIngress, state_destination_slots), DType::I32);
    sampling = reinterpret_cast<const ops::SamplingConfig*>(
        static_cast<const unsigned char*>(ingress.data) +
        offsetof(OrdinaryDecodeIngress, sampling));
    sampled_tokens = Tensor(static_cast<unsigned char*>(egress.data) +
                                offsetof(OrdinaryDecodeEgress, sampled_tokens),
                            DType::I32, {count});
    logits         = layout.logits.bind(backing);
    hidden         = layout.hidden.bind(backing);
}

void complete_round_state_layout(LayoutBuilder& builder, RoundStateLayout& layout) {
    if (layout.complete) { throw std::logic_error("RoundState layout is already complete"); }
    validate_spec(layout.spec);
    const std::int32_t columns =
        checked_i32(static_cast<std::uint64_t>(layout.spec.draft_window) + 1ULL,
                    "RoundState columns exceed int32");
    const std::int32_t drafts = checked_i32(std::max<std::uint64_t>(1ULL, layout.spec.draft_window),
                                            "RoundState drafts exceed int32");
    const auto i32            = [&](std::int32_t count, const char* label) {
        return add_tensor(builder, DType::I32, {count}, label);
    };
    if (layout.spec.enable_mtp) {
        layout.mtp.emplace();
        const auto ar_steps =
            checked_i32(std::max<std::uint64_t>(1ULL, layout.spec.draft_window - 1ULL),
                        "RoundState MTP AR steps exceed int32");
        layout.mtp->position         = i32(1, "MTP prefill autoregressive position");
        layout.mtp->ar_hidden        = add_tensor(builder, DType::BF16, {layout.spec.hidden, 1},
                                                  "MTP prefill autoregressive hidden");
        layout.mtp->draft_tokens     = i32(drafts, "MTP prefill draft tokens");
        layout.mtp->target_input_ids = i32(columns, "MTP prefill target input ids");
        layout.mtp->target_positions = i32(columns, "MTP prefill target positions");

        layout.mtp_decode.emplace();
        MtpDecodeStateLayout& decode = *layout.mtp_decode;
        decode.ingress = builder.add(sizeof(MtpDecodeIngress), kArenaAlign, "MTP decode ingress");
        decode.egress  = builder.add(sizeof(MtpDecodeEgress), kArenaAlign, "MTP decode egress");
        const auto batch =
            checked_i32(layout.spec.batch_capacity, "RoundState MTP batch capacity exceeds int32");
        decode.verify_ids =
            add_tensor(builder, DType::I32, {columns, batch}, "MTP decode verify ids");
        decode.target_positions =
            add_tensor(builder, DType::I32, {columns, batch}, "MTP decode target positions");
        decode.target_argmax =
            add_tensor(builder, DType::I32, {columns, batch}, "MTP decode target argmax");
        decode.target_logits =
            add_tensor(builder, DType::BF16, {layout.spec.output_rows, columns, batch},
                       "MTP decode target logits");
        decode.target_hidden = add_tensor(
            builder, DType::BF16, {layout.spec.hidden, columns, batch}, "MTP decode target hidden");
        decode.target_continuation_hidden =
            add_tensor(builder, DType::BF16, {layout.spec.hidden, batch},
                       "MTP decode target continuation hidden");
        decode.proposal_logits = add_tensor(builder, DType::BF16, {layout.spec.output_rows, batch},
                                            "MTP decode proposal logits");
        decode.alignment_ids =
            add_tensor(builder, DType::I32, {columns, batch}, "MTP decode alignment ids");
        decode.alignment_hidden =
            add_tensor(builder, DType::BF16, {layout.spec.hidden, columns, batch},
                       "MTP decode alignment hidden");
        decode.ar_hidden = add_tensor(builder, DType::BF16, {layout.spec.hidden, batch},
                                      "MTP decode autoregressive hidden");
        decode.next_hidden =
            add_tensor(builder, DType::BF16, {layout.spec.hidden, batch}, "MTP decode next hidden");
        decode.ar_positions      = add_tensor(builder, DType::I32, {batch, ar_steps},
                                              "MTP decode autoregressive positions");
        decode.ar_rope_positions = add_tensor(builder, DType::I32, {batch, ar_steps},
                                              "MTP decode autoregressive rope positions");
        decode.ar_valid_columns  = add_tensor(builder, DType::I32, {batch, ar_steps},
                                              "MTP decode autoregressive valid columns");
    }
    if (layout.spec.enable_dflash) {
        layout.dflash_prefill.emplace().produced_count = i32(1, "DFlash prefill produced count");
        DFlashDecodeStateLayout& decode                = layout.dflash_decode.emplace();
        decode.ingress =
            builder.add(sizeof(DFlashDecodeIngress), kArenaAlign, "DFlash decode ingress");
        decode.egress =
            builder.add(sizeof(DFlashDecodeEgress), kArenaAlign, "DFlash decode egress");
        const auto batch = checked_i32(layout.spec.batch_capacity,
                                       "RoundState DFlash batch capacity exceeds int32");
        decode.proposal_ids =
            add_tensor(builder, DType::I32, {columns, batch}, "DFlash proposal ids");
        decode.proposal_positions =
            add_tensor(builder, DType::I32, {columns, batch}, "DFlash proposal positions");
        decode.append_positions =
            add_tensor(builder, DType::I32, {columns, batch}, "DFlash append positions");
        decode.append_counts = add_tensor(builder, DType::I32, {batch}, "DFlash append counts");
        decode.draft_tokens =
            add_tensor(builder, DType::I32, {columns - 1, batch}, "DFlash proposal draft tokens");
        decode.verify_ids =
            add_tensor(builder, DType::I32, {columns, batch}, "DFlash target verify ids");
        decode.target_argmax =
            add_tensor(builder, DType::I32, {columns, batch}, "DFlash target argmax");
        decode.target_logits =
            add_tensor(builder, DType::BF16, {layout.spec.output_rows, columns, batch},
                       "DFlash target logits");
        decode.target_hidden = add_tensor(
            builder, DType::BF16, {layout.spec.hidden, columns, batch}, "DFlash target hidden");
        decode.target_continuation_hidden = add_tensor(
            builder, DType::BF16, {layout.spec.hidden, batch}, "DFlash target continuation hidden");
    }
    layout.complete = true;
}

MtpPrefillState::MtpPrefillState(DeviceSpan backing, const MtpPrefillStateLayout& layout)
    : position(layout.position.bind(backing)), ar_hidden(layout.ar_hidden.bind(backing)),
      draft_tokens(layout.draft_tokens.bind(backing)),
      target_input_ids(layout.target_input_ids.bind(backing)),
      target_positions(layout.target_positions.bind(backing)) {}

DFlashPrefillState::DFlashPrefillState(DeviceSpan backing, const DFlashPrefillStateLayout& layout)
    : produced_count(layout.produced_count.bind(backing)) {}

MtpDecodeState::MtpDecodeState(DeviceSpan backing, const MtpDecodeStateLayout& layout,
                               std::uint32_t batch_capacity, std::uint32_t draft_window) {
    if (batch_capacity == 0 || batch_capacity > kMaximumConcurrency || draft_window == 0 ||
        draft_window > kMtpDecodeMaximumDrafts) {
        throw std::invalid_argument("MTP decode state dimensions are outside the supported domain");
    }
    static_assert(std::is_standard_layout_v<MtpDecodeIngress>);
    static_assert(std::is_standard_layout_v<MtpDecodeEgress>);
    const auto batch          = static_cast<std::int32_t>(batch_capacity);
    const auto drafts         = static_cast<std::int32_t>(draft_window);
    const auto width          = drafts + 1;
    const auto steps          = std::max(drafts - 1, 1);
    ingress                   = layout.ingress.bind(backing);
    egress                    = layout.egress.bind(backing);
    const auto ingress_tensor = [&](std::size_t offset, DType dtype,
                                    std::initializer_list<std::int32_t> shape) {
        return Tensor(static_cast<unsigned char*>(ingress.data) + offset, dtype, shape);
    };
    const auto egress_tensor = [&](std::size_t offset, DType dtype,
                                   std::initializer_list<std::int32_t> shape) {
        return Tensor(static_cast<unsigned char*>(egress.data) + offset, dtype, shape);
    };
    anchors = ingress_tensor(offsetof(MtpDecodeIngress, anchors), DType::I32, {batch});
    base_frontiers =
        ingress_tensor(offsetof(MtpDecodeIngress, base_frontiers), DType::I32, {batch});
    remaining_budgets =
        ingress_tensor(offsetof(MtpDecodeIngress, remaining_budgets), DType::I32, {batch});
    current_extents =
        ingress_tensor(offsetof(MtpDecodeIngress, current_extents), DType::I32, {batch});
    target_valid_columns =
        ingress_tensor(offsetof(MtpDecodeIngress, target_valid_columns), DType::I32, {batch});
    current_drafts =
        ingress_tensor(offsetof(MtpDecodeIngress, current_drafts), DType::I32, {drafts, batch});
    target_rope_positions = ingress_tensor(offsetof(MtpDecodeIngress, target_rope_positions),
                                           DType::I32, {width, batch});
    text_kv_table_rows =
        ingress_tensor(offsetof(MtpDecodeIngress, text_kv_table_rows), DType::I32, {batch});
    mtp_kv_table_rows =
        ingress_tensor(offsetof(MtpDecodeIngress, mtp_kv_table_rows), DType::I32, {batch});
    state_source_slots =
        ingress_tensor(offsetof(MtpDecodeIngress, state_source_slots), DType::I32, {batch});
    state_destination_slots =
        ingress_tensor(offsetof(MtpDecodeIngress, state_destination_slots), DType::I32, {batch});
    rope_deltas = ingress_tensor(offsetof(MtpDecodeIngress, rope_deltas), DType::I32, {batch});
    sampling    = reinterpret_cast<const ops::SamplingConfig*>(
        static_cast<const unsigned char*>(ingress.data) + offsetof(MtpDecodeIngress, sampling));

    licensed_tokens =
        egress_tensor(offsetof(MtpDecodeEgress, licensed_tokens), DType::I32, {width, batch});
    licensed_counts =
        egress_tensor(offsetof(MtpDecodeEgress, licensed_counts), DType::I32, {batch});
    accepted_drafts =
        egress_tensor(offsetof(MtpDecodeEgress, accepted_drafts), DType::I32, {batch});
    next_drafts =
        egress_tensor(offsetof(MtpDecodeEgress, next_drafts), DType::I32, {batch, drafts});
    next_extents     = egress_tensor(offsetof(MtpDecodeEgress, next_extents), DType::I32, {batch});
    verify_ids       = layout.verify_ids.bind(backing);
    target_positions = layout.target_positions.bind(backing);
    target_argmax    = layout.target_argmax.bind(backing);
    target_logits    = layout.target_logits.bind(backing);
    target_hidden    = layout.target_hidden.bind(backing);
    target_continuation_hidden = layout.target_continuation_hidden.bind(backing);
    proposal_logits            = layout.proposal_logits.bind(backing);
    alignment_ids              = layout.alignment_ids.bind(backing);
    alignment_hidden           = layout.alignment_hidden.bind(backing);
    ar_hidden                  = layout.ar_hidden.bind(backing);
    next_hidden                = layout.next_hidden.bind(backing);
    ar_positions               = layout.ar_positions.bind(backing);
    ar_rope_positions          = layout.ar_rope_positions.bind(backing);
    ar_valid_columns           = layout.ar_valid_columns.bind(backing);
    if (ar_positions.ne[0] != batch || ar_positions.ne[1] != steps) {
        throw std::logic_error("MTP decode AR layout does not match its configured dimensions");
    }
}

DFlashDecodeState::DFlashDecodeState(DeviceSpan backing, const DFlashDecodeStateLayout& layout,
                                     std::uint32_t batch_capacity, std::uint32_t draft_window) {
    if (batch_capacity == 0 || batch_capacity > kMaximumConcurrency || draft_window == 0 ||
        draft_window > kDFlashDecodeMaximumDrafts) {
        throw std::invalid_argument(
            "DFlash decode state dimensions are outside the supported domain");
    }
    static_assert(std::is_standard_layout_v<DFlashDecodeIngress>);
    static_assert(std::is_standard_layout_v<DFlashDecodeEgress>);
    const auto batch          = static_cast<std::int32_t>(batch_capacity);
    const auto drafts         = static_cast<std::int32_t>(draft_window);
    const auto width          = drafts + 1;
    ingress                   = layout.ingress.bind(backing);
    egress                    = layout.egress.bind(backing);
    const auto ingress_tensor = [&](std::size_t offset, DType dtype,
                                    std::initializer_list<std::int32_t> shape) {
        return Tensor(static_cast<unsigned char*>(ingress.data) + offset, dtype, shape);
    };
    const auto egress_tensor = [&](std::size_t offset, DType dtype,
                                   std::initializer_list<std::int32_t> shape) {
        return Tensor(static_cast<unsigned char*>(egress.data) + offset, dtype, shape);
    };
    anchors = ingress_tensor(offsetof(DFlashDecodeIngress, anchors), DType::I32, {batch});
    execution_frontiers =
        ingress_tensor(offsetof(DFlashDecodeIngress, execution_frontiers), DType::I32, {batch});
    context_frontiers =
        ingress_tensor(offsetof(DFlashDecodeIngress, context_frontiers), DType::I32, {batch});
    proposal_extents =
        ingress_tensor(offsetof(DFlashDecodeIngress, proposal_extents), DType::I32, {batch});
    target_valid_columns =
        ingress_tensor(offsetof(DFlashDecodeIngress, target_valid_columns), DType::I32, {batch});
    text_kv_table_rows =
        ingress_tensor(offsetof(DFlashDecodeIngress, text_kv_table_rows), DType::I32, {batch});
    dflash_kv_table_rows =
        ingress_tensor(offsetof(DFlashDecodeIngress, dflash_kv_table_rows), DType::I32, {batch});
    active_lanes = ingress_tensor(offsetof(DFlashDecodeIngress, active_lanes), DType::I32, {batch});
    state_source_slots =
        ingress_tensor(offsetof(DFlashDecodeIngress, state_source_slots), DType::I32, {batch});
    state_destination_slots =
        ingress_tensor(offsetof(DFlashDecodeIngress, state_destination_slots), DType::I32, {batch});
    sampling = reinterpret_cast<const ops::SamplingConfig*>(
        static_cast<const unsigned char*>(ingress.data) + offsetof(DFlashDecodeIngress, sampling));
    licensed_tokens =
        egress_tensor(offsetof(DFlashDecodeEgress, licensed_tokens), DType::I32, {width, batch});
    licensed_counts =
        egress_tensor(offsetof(DFlashDecodeEgress, licensed_counts), DType::I32, {batch});
    accepted_drafts =
        egress_tensor(offsetof(DFlashDecodeEgress, accepted_drafts), DType::I32, {batch});
    proposal_ids               = layout.proposal_ids.bind(backing);
    proposal_positions         = layout.proposal_positions.bind(backing);
    append_positions           = layout.append_positions.bind(backing);
    append_counts              = layout.append_counts.bind(backing);
    draft_tokens               = layout.draft_tokens.bind(backing);
    verify_ids                 = layout.verify_ids.bind(backing);
    target_argmax              = layout.target_argmax.bind(backing);
    target_logits              = layout.target_logits.bind(backing);
    target_hidden              = layout.target_hidden.bind(backing);
    target_continuation_hidden = layout.target_continuation_hidden.bind(backing);
}

RoundState::RoundState(DeviceSpan backing, const RoundStateLayout& layout) {
    if (!layout.complete) { throw std::invalid_argument("RoundState layout is incomplete"); }
    if (layout.ordinary) {
        ordinary.emplace(backing, *layout.ordinary, layout.spec.batch_capacity);
    }
    token                = layout.token.bind(backing);
    pos                  = layout.pos.bind(backing);
    rope_pos             = layout.rope_pos.bind(backing);
    rope_delta           = layout.rope_delta.bind(backing);
    logits               = layout.logits.bind(backing);
    text_kv_table_row    = layout.text_kv_table_row.bind(backing);
    backend_kv_table_row = layout.backend_kv_table_row.bind(backing);
    if (layout.mtp) { mtp.emplace(backing, *layout.mtp); }
    if (layout.dflash_prefill) { dflash_prefill.emplace(backing, *layout.dflash_prefill); }
    if (layout.mtp_decode) {
        mtp_decode.emplace(backing, *layout.mtp_decode, layout.spec.batch_capacity,
                           layout.spec.draft_window);
    }
    if (layout.dflash_decode) {
        dflash_decode.emplace(backing, *layout.dflash_decode, layout.spec.batch_capacity,
                              layout.spec.draft_window);
    }
}

} // namespace ninfer::targets::qwen3_6
#include "serve/anthropic_thinking_signature.h"

#include <sys/random.h>

#include <array>
#include <bit>
#include <cerrno>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <span>
#include <stdexcept>
#include <string>
#include <system_error>
#include <vector>

namespace ninfer::serve {
namespace {

using Digest = std::array<std::uint8_t, 32>;

constexpr std::string_view kSignaturePrefix = "sig_ninfer_v1_";
constexpr std::string_view kDomain          = "ninfer.anthropic.thinking.v1";
constexpr std::array<std::uint32_t, 64> kRound{
    0x428a2f98U, 0x71374491U, 0xb5c0fbcfU, 0xe9b5dba5U, 0x3956c25bU, 0x59f111f1U, 0x923f82a4U,
    0xab1c5ed5U, 0xd807aa98U, 0x12835b01U, 0x243185beU, 0x550c7dc3U, 0x72be5d74U, 0x80deb1feU,
    0x9bdc06a7U, 0xc19bf174U, 0xe49b69c1U, 0xefbe4786U, 0x0fc19dc6U, 0x240ca1ccU, 0x2de92c6fU,
    0x4a7484aaU, 0x5cb0a9dcU, 0x76f988daU, 0x983e5152U, 0xa831c66dU, 0xb00327c8U, 0xbf597fc7U,
    0xc6e00bf3U, 0xd5a79147U, 0x06ca6351U, 0x14292967U, 0x27b70a85U, 0x2e1b2138U, 0x4d2c6dfcU,
    0x53380d13U, 0x650a7354U, 0x766a0abbU, 0x81c2c92eU, 0x92722c85U, 0xa2bfe8a1U, 0xa81a664bU,
    0xc24b8b70U, 0xc76c51a3U, 0xd192e819U, 0xd6990624U, 0xf40e3585U, 0x106aa070U, 0x19a4c116U,
    0x1e376c08U, 0x2748774cU, 0x34b0bcb5U, 0x391c0cb3U, 0x4ed8aa4aU, 0x5b9cca4fU, 0x682e6ff3U,
    0x748f82eeU, 0x78a5636fU, 0x84c87814U, 0x8cc70208U, 0x90befffaU, 0xa4506cebU, 0xbef9a3f7U,
    0xc67178f2U,
};

std::uint32_t load_be32(const std::uint8_t* bytes) {
    return (static_cast<std::uint32_t>(bytes[0]) << 24U) |
           (static_cast<std::uint32_t>(bytes[1]) << 16U) |
           (static_cast<std::uint32_t>(bytes[2]) << 8U) | static_cast<std::uint32_t>(bytes[3]);
}

void store_be32(std::uint32_t value, std::uint8_t* bytes) {
    bytes[0] = static_cast<std::uint8_t>(value >> 24U);
    bytes[1] = static_cast<std::uint8_t>(value >> 16U);
    bytes[2] = static_cast<std::uint8_t>(value >> 8U);
    bytes[3] = static_cast<std::uint8_t>(value);
}

void process_block(std::array<std::uint32_t, 8>& state, const std::uint8_t* block) {
    std::array<std::uint32_t, 64> words{};
    for (std::size_t index = 0; index < 16; ++index) {
        words[index] = load_be32(block + 4U * index);
    }
    for (std::size_t index = 16; index < words.size(); ++index) {
        const std::uint32_t s0 = std::rotr(words[index - 15], 7) ^
                                 std::rotr(words[index - 15], 18) ^ (words[index - 15] >> 3U);
        const std::uint32_t s1 = std::rotr(words[index - 2], 17) ^ std::rotr(words[index - 2], 19) ^
                                 (words[index - 2] >> 10U);
        words[index] = words[index - 16] + s0 + words[index - 7] + s1;
    }

    std::uint32_t a = state[0];
    std::uint32_t b = state[1];
    std::uint32_t c = state[2];
    std::uint32_t d = state[3];
    std::uint32_t e = state[4];
    std::uint32_t f = state[5];
    std::uint32_t g = state[6];
    std::uint32_t h = state[7];
    for (std::size_t index = 0; index < words.size(); ++index) {
        const std::uint32_t sum1     = std::rotr(e, 6) ^ std::rotr(e, 11) ^ std::rotr(e, 25);
        const std::uint32_t choose   = (e & f) ^ (~e & g);
        const std::uint32_t t1       = h + sum1 + choose + kRound[index] + words[index];
        const std::uint32_t sum0     = std::rotr(a, 2) ^ std::rotr(a, 13) ^ std::rotr(a, 22);
        const std::uint32_t majority = (a & b) ^ (a & c) ^ (b & c);
        const std::uint32_t t2       = sum0 + majority;
        h                            = g;
        g                            = f;
        f                            = e;
        e                            = d + t1;
        d                            = c;
        c                            = b;
        b                            = a;
        a                            = t1 + t2;
    }
    state[0] += a;
    state[1] += b;
    state[2] += c;
    state[3] += d;
    state[4] += e;
    state[5] += f;
    state[6] += g;
    state[7] += h;
}

Digest sha256(std::span<const std::uint8_t> input) {
    if (input.size() > std::numeric_limits<std::uint64_t>::max() / 8ULL) {
        throw std::invalid_argument("Thinking signature payload is too large");
    }
    std::array<std::uint32_t, 8> state{0x6a09e667U, 0xbb67ae85U, 0x3c6ef372U, 0xa54ff53aU,
                                       0x510e527fU, 0x9b05688cU, 0x1f83d9abU, 0x5be0cd19U};
    std::size_t offset = 0;
    while (input.size() - offset >= 64) {
        process_block(state, input.data() + offset);
        offset += 64;
    }

    std::array<std::uint8_t, 128> tail{};
    const std::size_t remaining = input.size() - offset;
    for (std::size_t index = 0; index < remaining; ++index) { tail[index] = input[offset + index]; }
    tail[remaining]             = 0x80U;
    const std::size_t tail_size = remaining < 56 ? 64 : 128;
    const std::uint64_t bits    = static_cast<std::uint64_t>(input.size()) * 8ULL;
    for (std::size_t index = 0; index < 8; ++index) {
        tail[tail_size - 1U - index] = static_cast<std::uint8_t>(bits >> (8U * index));
    }
    process_block(state, tail.data());
    if (tail_size == 128) { process_block(state, tail.data() + 64); }

    Digest result{};
    for (std::size_t index = 0; index < state.size(); ++index) {
        store_be32(state[index], result.data() + 4U * index);
    }
    return result;
}

void append_be64(std::vector<std::uint8_t>& output, std::uint64_t value) {
    for (int shift = 56; shift >= 0; shift -= 8) {
        output.push_back(static_cast<std::uint8_t>(value >> static_cast<unsigned int>(shift)));
    }
}

Digest thinking_mac(const AnthropicThinkingSigner::Key& key, std::string_view thinking,
                    std::size_t block_index) {
    constexpr std::size_t kBlockBytes = 64;
    std::vector<std::uint8_t> inner;
    inner.reserve(kBlockBytes + kDomain.size() + 1U + 16U + thinking.size());
    for (std::size_t index = 0; index < kBlockBytes; ++index) {
        const std::uint8_t value = index < key.size() ? key[index] : 0;
        inner.push_back(static_cast<std::uint8_t>(value ^ 0x36U));
    }
    inner.insert(inner.end(), kDomain.begin(), kDomain.end());
    inner.push_back(0);
    append_be64(inner, static_cast<std::uint64_t>(block_index));
    append_be64(inner, static_cast<std::uint64_t>(thinking.size()));
    inner.insert(inner.end(), thinking.begin(), thinking.end());
    const Digest inner_digest = sha256(inner);

    std::array<std::uint8_t, kBlockBytes + 32U> outer{};
    for (std::size_t index = 0; index < kBlockBytes; ++index) {
        const std::uint8_t value = index < key.size() ? key[index] : 0;
        outer[index]             = static_cast<std::uint8_t>(value ^ 0x5cU);
    }
    for (std::size_t index = 0; index < inner_digest.size(); ++index) {
        outer[kBlockBytes + index] = inner_digest[index];
    }
    return sha256(outer);
}

std::string hex_digest(const Digest& digest) {
    constexpr char kHex[] = "0123456789abcdef";
    std::string result(digest.size() * 2U, '\0');
    for (std::size_t index = 0; index < digest.size(); ++index) {
        result[2U * index]     = kHex[digest[index] >> 4U];
        result[2U * index + 1] = kHex[digest[index] & 0x0fU];
    }
    return result;
}

int hex_value(char value) {
    if (value >= '0' && value <= '9') { return value - '0'; }
    if (value >= 'a' && value <= 'f') { return value - 'a' + 10; }
    return -1;
}

bool decode_digest(std::string_view encoded, Digest& digest) {
    if (encoded.size() != digest.size() * 2U) { return false; }
    for (std::size_t index = 0; index < digest.size(); ++index) {
        const int high = hex_value(encoded[2U * index]);
        const int low  = hex_value(encoded[2U * index + 1]);
        if (high < 0 || low < 0) { return false; }
        digest[index] = static_cast<std::uint8_t>((high << 4) | low);
    }
    return true;
}

AnthropicThinkingSigner::Key random_key() {
    AnthropicThinkingSigner::Key key{};
    std::size_t offset = 0;
    while (offset < key.size()) {
        const ssize_t count = ::getrandom(key.data() + offset, key.size() - offset, 0);
        if (count < 0) {
            if (errno == EINTR) { continue; }
            throw std::system_error(errno, std::generic_category(),
                                    "failed to initialize Anthropic Thinking signer");
        }
        if (count == 0) {
            throw std::runtime_error("operating-system random source returned no key bytes");
        }
        offset += static_cast<std::size_t>(count);
    }
    return key;
}

} // namespace

AnthropicThinkingSigner::AnthropicThinkingSigner() : key_(random_key()) {}

AnthropicThinkingSigner::AnthropicThinkingSigner(Key key) : key_(key) {}

std::string AnthropicThinkingSigner::sign(std::string_view thinking,
                                          std::size_t block_index) const {
    return std::string(kSignaturePrefix) + hex_digest(thinking_mac(key_, thinking, block_index));
}

bool AnthropicThinkingSigner::verify(std::string_view thinking, std::size_t block_index,
                                     std::string_view signature) const {
    if (!signature.starts_with(kSignaturePrefix)) { return false; }
    Digest supplied{};
    if (!decode_digest(signature.substr(kSignaturePrefix.size()), supplied)) { return false; }
    const Digest expected   = thinking_mac(key_, thinking, block_index);
    std::uint8_t difference = 0;
    for (std::size_t index = 0; index < expected.size(); ++index) {
        difference |= static_cast<std::uint8_t>(expected[index] ^ supplied[index]);
    }
    return difference == 0;
}

} // namespace ninfer::serve
// ninfer::ops::detail - FP8 E4M3FN split-KV small-T launch ownership.
#include "ops/softmax_attention/dense/causal_cache/launch.h"

#include "core/device.h"
#include "ops/common/math.h"
#include "ops/softmax_attention/dense/causal_cache/small_t_fp8.cuh"

#include <cstdint>
#include <stdexcept>

namespace ninfer::ops::detail {
namespace {

PagedKVBatchLayerView fp8_single_row_batch_view(const PagedKVLayerView& cache) {
    return {
        .k_pages       = cache.k_pages,
        .v_pages       = cache.v_pages,
        .k_scale_pages = cache.k_scale_pages,
        .v_scale_pages = cache.v_scale_pages,
        .block_tables  = cache.block_table.view({cache.block_table.ne[0], 1}),
        .head_dim      = cache.head_dim,
        .num_kv_heads  = cache.num_kv_heads,
        .dtype         = cache.dtype,
        .quant_group   = cache.quant_group,
    };
}

template <typename Geometry, int TokenTile, bool MultiBatch, bool Masked, typename CacheInput>
void launch_fp8_partial(const Tensor& q, CacheInput input, const Tensor& positions, float scale,
                        PagedKVBatchLayerView cache, const CausalSmallTInvocation& invocation,
                        std::int32_t logical_capacity, std::int32_t splits, Tensor& partial_acc,
                        Tensor& partial_m, Tensor& partial_l, cudaStream_t stream) {
    constexpr int RowCount             = TokenTile * Geometry::GroupSize;
    constexpr int RowTiles             = (RowCount + 15) / 16;
    constexpr int Warps                = RowTiles == 3 ? 12 : 8;
    constexpr int KeyBlock             = TokenTile == 1 ? 32 : 64;
    constexpr int MinBlocks            = TokenTile == 1 ? 2 : 1;
    constexpr std::size_t DynamicBytes = 4u * KeyBlock * kCausalHeadDim;
    using KernelInput                  = CacheInput;
    const dim3 grid(Geometry::KVHeads, splits, invocation.batch_size);
    const auto launch = [&]() {
        auto kernel = causal_attention_small_t_fp8_tiled_kernel<
            Geometry, TokenTile, Warps, MinBlocks, KeyBlock, true, MultiBatch, Masked, KernelInput>;
        static const cudaError_t attr = cudaFuncSetAttribute(
            kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(DynamicBytes));
        CUDA_CHECK(attr);

        const auto q_ptr         = static_cast<const __nv_bfloat16*>(q.data);
        const auto positions_ptr = static_cast<const std::int32_t*>(positions.data);
        const auto cache_k_ptr   = static_cast<std::uint8_t*>(cache.k_pages.data);
        const auto cache_v_ptr   = static_cast<std::uint8_t*>(cache.v_pages.data);
        const auto k_scale_ptr   = static_cast<__half*>(cache.k_scale_pages.data);
        const auto v_scale_ptr   = static_cast<__half*>(cache.v_scale_pages.data);
        const auto tables_ptr    = static_cast<const std::int32_t*>(cache.block_tables.data);
        const auto valid_ptr =
            invocation.valid_columns == nullptr
                ? nullptr
                : static_cast<const std::int32_t*>(invocation.valid_columns->data);
        const auto rows_ptr   = invocation.table_rows == nullptr
                                    ? nullptr
                                    : static_cast<const std::int32_t*>(invocation.table_rows->data);
        auto* partial_acc_ptr = static_cast<float*>(partial_acc.data);
        auto* partial_m_ptr   = static_cast<float*>(partial_m.data);
        auto* partial_l_ptr   = static_cast<float*>(partial_l.data);
        kernel<<<grid, Warps * 32, DynamicBytes, stream>>>(
            q_ptr, input, positions_ptr, cache_k_ptr, cache_v_ptr, k_scale_ptr, v_scale_ptr,
            tables_ptr, valid_ptr, rows_ptr, cache.block_tables.ne[0], invocation.full_width,
            invocation.column_begin, logical_capacity, scale, partial_acc_ptr, partial_m_ptr,
            partial_l_ptr);
        CUDA_CHECK(cudaGetLastError());
    };

    launch();
}

template <typename Geometry, bool MultiBatch, bool Masked>
void launch_fp8_reduce(const Tensor& positions, const CausalSmallTInvocation& invocation,
                       std::int32_t splits, const Tensor& partial_acc, const Tensor& partial_m,
                       const Tensor& partial_l, Tensor& out, cudaStream_t stream) {
    constexpr int Block  = 256;
    constexpr int DChunk = 64;
    const dim3 grid(Geometry::QHeads, div_up(kCausalHeadDim, DChunk),
                    invocation.width * invocation.batch_size);
    const auto launch = [&]<bool Offset>() {
        causal_attention_small_t_fp8_reduce_output_kernel<Geometry, DChunk, MultiBatch, Masked,
                                                          Offset><<<grid, Block, 0, stream>>>(
            static_cast<const float*>(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),
            invocation.valid_columns == nullptr
                ? nullptr
                : static_cast<const std::int32_t*>(invocation.valid_columns->data),
            invocation.width, invocation.full_width, invocation.column_begin, invocation.batch_size,
            splits, static_cast<__nv_bfloat16*>(out.data));
    };
    if (invocation.column_begin == 0) {
        launch.template operator()<false>();
    } else {
        launch.template operator()<true>();
    }
    CUDA_CHECK(cudaGetLastError());
}

template <typename Geometry, typename CacheInput>
void causal_attention_small_t_fp8_launch_for(const Tensor& q, CacheInput input,
                                             const Tensor& positions, float scale,
                                             PagedKVBatchLayerView cache,
                                             const CausalSmallTInvocation& invocation,
                                             CausalAttentionExecutionEnvelope envelope,
                                             Tensor& partial_acc, Tensor& partial_m,
                                             Tensor& partial_l, Tensor& out, cudaStream_t stream) {
    const auto logical_capacity = static_cast<std::int32_t>(envelope.max_visible_keys);
    const auto splits =
        causal_attention_split_capacity(Geometry::QHeads, invocation.width, cache.dtype, envelope);

    const auto launch_partial = [&]<int Tokens, bool MultiBatch, bool Masked>() {
        launch_fp8_partial<Geometry, Tokens, MultiBatch, Masked>(
            q, input, positions, scale, cache, invocation, logical_capacity, splits, partial_acc,
            partial_m, partial_l, stream);
    };
    const auto dispatch_metadata = [&]<int Tokens>() {
        const bool masked = invocation.valid_columns != nullptr;
        if (invocation.batch_size == 1) {
            if (masked) {
                launch_partial.template operator()<Tokens, false, true>();
            } else {
                launch_partial.template operator()<Tokens, false, false>();
            }
        } else if (masked) {
            launch_partial.template operator()<Tokens, true, true>();
        } else {
            launch_partial.template operator()<Tokens, true, false>();
        }
    };

    switch (invocation.width) {
    case 1:
        dispatch_metadata.template operator()<1>();
        break;
    case 2:
        dispatch_metadata.template operator()<2>();
        break;
    case 3:
        dispatch_metadata.template operator()<3>();
        break;
    case 4:
        dispatch_metadata.template operator()<4>();
        break;
    case 5:
        dispatch_metadata.template operator()<5>();
        break;
    case 6:
        dispatch_metadata.template operator()<6>();
        break;
    default:
        throw std::invalid_argument("causal_attention_small_t_fp8_launch: unsupported T");
    }

    const bool masked = invocation.valid_columns != nullptr;
    if (invocation.batch_size == 1) {
        if (masked) {
            launch_fp8_reduce<Geometry, false, true>(positions, invocation, splits, partial_acc,
                                                     partial_m, partial_l, out, stream);
        } else {
            launch_fp8_reduce<Geometry, false, false>(positions, invocation, splits, partial_acc,
                                                      partial_m, partial_l, out, stream);
        }
    } else if (masked) {
        launch_fp8_reduce<Geometry, true, true>(positions, invocation, splits, partial_acc,
                                                partial_m, partial_l, out, stream);
    } else {
        launch_fp8_reduce<Geometry, true, false>(positions, invocation, splits, partial_acc,
                                                 partial_m, partial_l, out, stream);
    }
}

} // namespace

void causal_attention_small_t_fp8_launch(
    const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& positions,
    const Tensor& valid_columns, const Tensor& table_rows, float scale, PagedKVBatchLayerView cache,
    CausalAttentionExecutionEnvelope envelope, std::int32_t column_begin, std::int32_t width,
    Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, Tensor& out, cudaStream_t stream) {
    const CausalAppendInput input{static_cast<const __nv_bfloat16*>(k.data),
                                  static_cast<const __nv_bfloat16*>(v.data)};
    const CausalSmallTInvocation invocation{
        .valid_columns = valid_columns.data == nullptr ? nullptr : &valid_columns,
        .table_rows    = &table_rows,
        .full_width    = q.ne[2],
        .column_begin  = column_begin,
        .width         = width,
        .batch_size    = q.ne[3],
    };
    if (q.ne[1] == CausalD256H24Kv4::QHeads) {
        causal_attention_small_t_fp8_launch_for<CausalD256H24Kv4>(
            q, input, positions, scale, cache, invocation, envelope, partial_acc, partial_m,
            partial_l, out, stream);
        return;
    }
    causal_attention_small_t_fp8_launch_for<CausalD256H16Kv2>(q, input, positions, scale, cache,
                                                              invocation, envelope, partial_acc,
                                                              partial_m, partial_l, out, stream);
}

void causal_attention_cached_small_t_fp8_launch(const Tensor& q, const Tensor& positions,
                                                float scale, const PagedKVLayerView& cache,
                                                CausalAttentionExecutionEnvelope envelope,
                                                Tensor& partial_acc, Tensor& partial_m,
                                                Tensor& partial_l, Tensor& out,
                                                cudaStream_t stream) {
    const CausalCachedInput input{};
    const CausalSmallTInvocation invocation{
        .valid_columns = nullptr,
        .table_rows    = nullptr,
        .full_width    = q.ne[2],
        .column_begin  = 0,
        .width         = q.ne[2],
        .batch_size    = 1,
    };
    PagedKVBatchLayerView batch_cache = fp8_single_row_batch_view(cache);
    if (q.ne[1] == CausalD256H24Kv4::QHeads) {
        causal_attention_small_t_fp8_launch_for<CausalD256H24Kv4>(
            q, input, positions, scale, batch_cache, invocation, envelope, partial_acc, partial_m,
            partial_l, out, stream);
        return;
    }
    causal_attention_small_t_fp8_launch_for<CausalD256H16Kv2>(
        q, input, positions, scale, batch_cache, invocation, envelope, partial_acc, partial_m,
        partial_l, out, stream);
}

} // namespace ninfer::ops::detail
#include "ops/gdn_input_proj/fp8/fp8_gdn_input_plan.h"

#include "core/device.h"
#include "ops/gdn_input_proj/fp8/fp8_gdn_input_output.cuh"
#include "ops/linear/fp8/fp8_a8_schedule.cuh"
#include "ops/linear/fp8/fp8_config.h"
#include "ops/linear/fp8/fp8_output.cuh"

#include <cuda_bf16.h>

#include <cstdint>

namespace ninfer::ops::detail {
namespace {

using Geometry = Fp8GdnInputGeometry;
using Schedule = typename Fp8LinearA8ProductionSchedule<Geometry>::Type;

static_assert((Fp8GdnInputOutput::kQkvRows % Schedule::kBlockRows) == 0);
static_assert((Fp8GdnInputOutput::kZRows % Schedule::kBlockRows) == 0);

template <bool FullTokens>
void launch_mma(const Weight& weight, Tensor& qkv, Tensor& z, Fp8A8Workspace workspace,
                std::int32_t tokens, cudaStream_t stream) {
    constexpr int kRowTiles = Geometry::kOutputRows / Schedule::kBlockRows;
    const int token_tiles   = (tokens + Schedule::kBlockTokens - 1) / Schedule::kBlockTokens;
    const int blocks        = kRowTiles * token_tiles;
    const Fp8GdnInputOutput output{static_cast<__nv_bfloat16*>(qkv.data),
                                   static_cast<__nv_bfloat16*>(z.data)};

    if constexpr (Schedule::kSharedBytes > 48 * 1024) {
        static const cudaError_t attribute = cudaFuncSetAttribute(
            fp8_mma_kernel<Geometry, Schedule, FullTokens, Fp8IdentityEpilogue, Fp8GdnInputOutput>,
            cudaFuncAttributeMaxDynamicSharedMemorySize, Schedule::kSharedBytes);
        CUDA_CHECK(attribute);
    }
    fp8_mma_kernel<Geometry, Schedule, FullTokens>
        <<<blocks, Schedule::kThreads, Schedule::kSharedBytes, stream>>>(
            workspace.codes, workspace.scales, static_cast<const std::uint8_t*>(weight.qdata),
            static_cast<const __nv_bfloat16*>(weight.scales), tokens, Fp8IdentityEpilogue{},
            output);
    CUDA_CHECK(cudaGetLastError());
}

} // namespace

void fp8_gdn_input_a8_launch(const Tensor& x, const Weight& weight, Tensor& qkv, Tensor& z,
                             Fp8A8Workspace workspace, cudaStream_t stream) {
    launch_fp8_a8_quantize(x, weight, workspace, stream);
    if ((x.ne[1] % Schedule::kBlockTokens) == 0) {
        launch_mma<true>(weight, qkv, z, workspace, x.ne[1], stream);
    } else {
        launch_mma<false>(weight, qkv, z, workspace, x.ne[1], stream);
    }
}

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

// ninfer::ops - split-KV causal small-T attention shared scaffolding. The BF16 and
// int8 partial kernels live in causal_attention_small_t_bf16.cuh and
// causal_attention_small_t_i8.cuh respectively; they are fully separate kernels (no
// shared body) so each KV format can be optimized independently. This header owns
// only what both share: layout constants, device helpers, and the split reducer.

#include "ops/common/math.cuh"
#include "ops/common/mma.cuh"
#include "ops/common/warp.cuh"
#include "ops/softmax_attention/dense/causal_cache/geometry.cuh"
#include "ops/kernel/paged_kv_address.cuh"

#include <cuda_bf16.h>
#include <math_constants.h>

#include <cstdint>

namespace ninfer::ops {

inline constexpr int kCausalHeadDim = 256;

struct CausalAppendInput {
    static constexpr bool writes_cache = true;
    const __nv_bfloat16* k;
    const __nv_bfloat16* v;
};

struct CausalCachedInput {
    static constexpr bool writes_cache = false;
};

template <typename Geometry>
__device__ __forceinline__ std::int64_t causal_cache_index(int physical_page, int kv_head, int d,
                                                           int page_offset) {
    return paged_kv_element_offset<kCausalHeadDim, Geometry::KVHeads>(physical_page, kv_head,
                                                                      page_offset, d);
}

template <typename Geometry>
__device__ __forceinline__ std::int64_t causal_q_index(int q_head, int d, int token = 0) {
    return static_cast<std::int64_t>(d) + static_cast<std::int64_t>(kCausalHeadDim) *
                                              (static_cast<std::int64_t>(q_head) +
                                               static_cast<std::int64_t>(Geometry::QHeads) * token);
}

template <typename Geometry>
__device__ __forceinline__ std::int64_t kv_cache_int8_new_index(int kv_head, int d, int token = 0) {
    return static_cast<std::int64_t>(d) +
           static_cast<std::int64_t>(kCausalHeadDim) *
               (static_cast<std::int64_t>(kv_head) +
                static_cast<std::int64_t>(Geometry::KVHeads) * token);
}

template <typename Geometry>
__device__ __forceinline__ std::int64_t causal_partial_acc_index(int q_head, int d, int token,
                                                                 int split, int tokens) {
    return static_cast<std::int64_t>(d) +
           static_cast<std::int64_t>(kCausalHeadDim) *
               (static_cast<std::int64_t>(q_head) +
                static_cast<std::int64_t>(Geometry::QHeads) *
                    (static_cast<std::int64_t>(token) + static_cast<std::int64_t>(tokens) * split));
}

template <typename Geometry>
__device__ __forceinline__ std::int64_t causal_partial_stat_index(int q_head, int token, int split,
                                                                  int tokens) {
    return static_cast<std::int64_t>(q_head) +
           static_cast<std::int64_t>(Geometry::QHeads) *
               (static_cast<std::int64_t>(token) + static_cast<std::int64_t>(tokens) * split);
}

template <typename Geometry>
__device__ __forceinline__ bool causal_valid_q_head(int kv_head, int q_head) {
    return kv_head >= 0 && kv_head < Geometry::KVHeads && q_head >= kv_head * Geometry::GroupSize &&
           q_head < (kv_head + 1) * Geometry::GroupSize && q_head < Geometry::QHeads;
}

template <typename Geometry>
__device__ __forceinline__ int causal_small_t_default_splits(int window) {
    int target_keys_per_split = 480 / Geometry::SmallTSplitScale;
    if (window <= 4096) {
        target_keys_per_split = 64 / Geometry::SmallTSplitScale;
    } else if (window <= 8198) {
        target_keys_per_split = 128 / Geometry::SmallTSplitScale;
    } else if (window <= 16390) {
        target_keys_per_split = 256 / Geometry::SmallTSplitScale;
    }
    constexpr int kMinSplits = 4 * Geometry::SmallTSplitScale;
    int splits               = div_up(window, target_keys_per_split);
    splits                   = splits > kMinSplits ? splits : kMinSplits;
    return splits < Geometry::SmallTMaximumSplits ? splits : Geometry::SmallTMaximumSplits;
}

template <typename Geometry, bool Int8>
__device__ __forceinline__ int causal_small_t_active_splits(int window, int launch_capacity,
                                                            int tokens) {
    if (window <= 0) { return launch_capacity; }
    int splits = 0;
    if constexpr (Int8) {
        if (tokens == 5 && window > 128 && window <= 512) {
            splits = div_up(window, 32 / Geometry::SmallTSplitScale);
        } else if (tokens == 6 && window > 128 && window <= 160) {
            constexpr int kKeysPerSplit = Geometry::SmallTSplitScale == 2 ? 17 : 24;
            splits                      = div_up(window, kKeysPerSplit);
        } else if (tokens == 6 && window > 5000 && window <= 8198) {
            splits             = div_up(window, 192 / Geometry::SmallTSplitScale);
            constexpr int kMin = 4 * Geometry::SmallTSplitScale;
            constexpr int kMax = 42 * Geometry::SmallTSplitScale;
            splits             = splits > kMin ? splits : kMin;
            splits             = splits < kMax ? splits : kMax;
        } else {
            splits = causal_small_t_default_splits<Geometry>(window);
        }
    } else {
        splits = causal_small_t_default_splits<Geometry>(window);
    }
    return splits < launch_capacity ? splits : launch_capacity;
}

template <typename Geometry>
__device__ __forceinline__ int causal_small_t_fp8_active_splits(int window, int launch_capacity,
                                                                int tokens) {
    int splits = causal_small_t_default_splits<Geometry>(window);
    if constexpr (Geometry::SmallTSplitScale == 1) {
        if (tokens == 1 && window > 8198) { splits = Geometry::SmallTMaximumSplits; }
    }
    return splits < launch_capacity ? splits : launch_capacity;
}

__device__ __forceinline__ int causal_small_t_tc_swz(int row, int col) {
    return (((col >> 3) ^ (row & 7)) << 3) | (col & 7);
}

__device__ __forceinline__ int causal_small_t_tc_swz32(int row, int col) {
    return (((col >> 3) ^ (row & 3)) << 3) | (col & 7);
}

// Signed int8 QK MMA, k=32 contraction. A = 16x32 s8 (4 regs/thread, 4 s8 each),
// B = 8x32 s8 col-major (2 regs/thread), D = 16x8 s32 (4 regs/thread). The A/B
// register byte layout is identical to the m16n8k16 bf16 fragments loaded by
// ldmatrix_x4/x2 over a d-contiguous int8 tile reinterpreted as
// b16 (two packed int8 per 16-bit lane), so the same ldmatrix helpers and XOR
// swizzle feed this MMA. The s32 accumulator layout matches the bf16 f32
// accumulator (c0/c1 -> row groupID, c2/c3 -> row groupID+8), so score
// consumption is unchanged; only per-64-group scale rescale differs.
template <typename Geometry>
__device__ __forceinline__ void causal_small_t_tc_row_to_qt(int row, int tokens, int kv_head,
                                                            int& q_head, int& token) {
    token             = row / Geometry::GroupSize;
    const int local_q = row - token * Geometry::GroupSize;
    q_head            = kv_head * Geometry::GroupSize + local_q;
}

template <typename Geometry, int DChunk, bool Int8, bool MultiBatch, bool Masked, bool Offset>
__launch_bounds__(256) __global__ void causal_attention_small_t_reduce_output_kernel(
    const __nv_bfloat16* partial_acc, const float* partial_m, const float* partial_l,
    const std::int32_t* positions, const std::int32_t* valid_columns, std::int32_t tokens,
    std::int32_t full_width, std::int32_t column_begin, std::int32_t batch_size,
    std::int32_t split_count, __nv_bfloat16* out) {
    static_assert(DChunk > 0 && DChunk <= kCausalHeadDim);

    const int q_head      = static_cast<int>(blockIdx.x);
    const int d_start     = static_cast<int>(blockIdx.y) * DChunk;
    const int flat_column = static_cast<int>(blockIdx.z);
    int batch             = 0;
    int token             = flat_column;
    if constexpr (MultiBatch) {
        batch = flat_column / tokens;
        token = flat_column - batch * tokens;
    }
    const int tid = threadIdx.x;
    if (q_head >= Geometry::QHeads || token >= tokens) { return; }
    if constexpr (MultiBatch) {
        if (batch >= batch_size) { return; }
    }

    if constexpr (Offset) { positions += column_begin; }
    if constexpr (MultiBatch) { positions += batch * full_width; }
    const int last_pos = positions[tokens - 1];
    int output_column  = token;
    if constexpr (Offset) { output_column += column_begin; }
    if constexpr (MultiBatch) { output_column += batch * full_width; }

    if constexpr (MultiBatch) {
        const std::int64_t partial_acc_row = static_cast<std::int64_t>(batch) * kCausalHeadDim *
                                             Geometry::QHeads * tokens * split_count;
        const std::int64_t partial_stat_row =
            static_cast<std::int64_t>(batch) * Geometry::QHeads * tokens * split_count;
        partial_acc += partial_acc_row;
        partial_m += partial_stat_row;
        partial_l += partial_stat_row;
    }

    const int window = last_pos + 1;
    const int active_split_count =
        causal_small_t_active_splits<Geometry, Int8>(window, split_count, tokens);

    __shared__ float reduce[256];

    float local_m = -CUDART_INF_F;
    for (int split = tid; split < active_split_count; split += blockDim.x) {
        local_m = fmaxf(
            local_m, partial_m[causal_partial_stat_index<Geometry>(q_head, token, split, tokens)]);
    }
    reduce[tid] = local_m;
    __syncthreads();

    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) { reduce[tid] = fmaxf(reduce[tid], reduce[tid + stride]); }
        __syncthreads();
    }
    const float head_m = reduce[0];
    __syncthreads();

    if (head_m == -CUDART_INF_F) {
        const int d = d_start + tid;
        if (tid < DChunk && d < kCausalHeadDim) {
            out[causal_q_index<Geometry>(q_head, d, output_column)] = __float2bfloat16(0.0f);
        }
        return;
    }

    float local_l = 0.0f;
    for (int split = tid; split < active_split_count; split += blockDim.x) {
        const float tile_l =
            partial_l[causal_partial_stat_index<Geometry>(q_head, token, split, tokens)];
        if (tile_l > 0.0f) {
            local_l +=
                tile_l *
                expf(partial_m[causal_partial_stat_index<Geometry>(q_head, token, split, tokens)] -
                     head_m);
        }
    }
    reduce[tid] = local_l;
    __syncthreads();

    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) { reduce[tid] += reduce[tid + stride]; }
        __syncthreads();
    }
    const float head_l = reduce[0];

    const int d = d_start + tid;
    if (tid >= DChunk || d >= kCausalHeadDim) { return; }

    float numerator = 0.0f;
    if (head_l > 0.0f) {
        for (int split = 0; split < active_split_count; ++split) {
            const float tile_l =
                partial_l[causal_partial_stat_index<Geometry>(q_head, token, split, tokens)];
            if (tile_l <= 0.0f) { continue; }
            const float weight =
                expf(partial_m[causal_partial_stat_index<Geometry>(q_head, token, split, tokens)] -
                     head_m);
            numerator += __bfloat162float(partial_acc[causal_partial_acc_index<Geometry>(
                             q_head, d, token, split, tokens)]) *
                         weight;
        }
    }
    bool valid = true;
    if constexpr (Masked) {
        int absolute_column = token;
        if constexpr (Offset) { absolute_column += column_begin; }
        valid = absolute_column < valid_columns[batch];
    }
    const float value = (valid && head_l > 0.0f) ? numerator / head_l : 0.0f;
    out[causal_q_index<Geometry>(q_head, d, output_column)] = __float2bfloat16(value);
}

} // namespace ninfer::ops
#include "ops/gdn_input_proj/nvfp4/nvfp4_gdn_snapshot_plan.h"

#include "core/layout.h"
#include "ops/gdn_input_proj/nvfp4/nvfp4_gdn_input_plan.h"

#include <stdexcept>

namespace ninfer::ops::detail {
namespace {

struct Nvfp4GdnProjectedWorkspace {
    Tensor projected;
    DeviceSpan projection;
};

template <class Allocator>
Nvfp4GdnProjectedWorkspace allocate_workspace(Allocator& allocator, std::int32_t tokens) {
    Nvfp4GdnProjectedWorkspace out;
    out.projected = allocator.alloc(DType::BF16, {10240, tokens}, 256);
    const std::size_t projection_bytes =
        nvfp4_gdn_input_workspace_capacity_bytes(LinearPolicy::AllowA4, tokens, tokens);
    out.projection = allocator.alloc_bytes(projection_bytes, 256);
    return out;
}

} // namespace

Nvfp4GdnConvPlan nvfp4_gdn_conv_resolve_plan(LinearPolicy policy, std::int32_t tokens,
                                             std::int32_t batch_size) {
    if (tokens <= 0 || batch_size <= 0 || batch_size > 8) {
        throw std::invalid_argument("nvfp4 gdn conv: invalid B/T domain");
    }
    if (policy != LinearPolicy::A16Only && policy != LinearPolicy::AllowA4) {
        throw std::invalid_argument("nvfp4 gdn conv admits only A16 or A4");
    }
    if (batch_size > 1) { return {Nvfp4GdnConvScheduleId::Materialized}; }
    if (policy == LinearPolicy::A16Only) {
        if (tokens == 1) { return {Nvfp4GdnConvScheduleId::DecodeFusedA16}; }
        if (tokens <= 16) { return {Nvfp4GdnConvScheduleId::SmallTFusedA16}; }
        throw std::invalid_argument("nvfp4 gdn conv A16 is registered only through T=16");
    }
    if (tokens == 1) { return {Nvfp4GdnConvScheduleId::DecodeFusedA16}; }
    if (tokens <= 3) { return {Nvfp4GdnConvScheduleId::SmallTFusedA16}; }
    return {Nvfp4GdnConvScheduleId::Materialized};
}

std::size_t nvfp4_gdn_snapshot_workspace_capacity_bytes(LinearPolicy policy,
                                                        std::int32_t min_tokens,
                                                        std::int32_t max_tokens) {
    if (min_tokens <= 0 || max_tokens < min_tokens) {
        throw std::invalid_argument("nvfp4 gdn snapshot workspace: invalid token interval");
    }
    (void)nvfp4_gdn_conv_resolve_plan(policy, min_tokens, 1);
    const Nvfp4GdnConvPlan maximum_plan = nvfp4_gdn_conv_resolve_plan(policy, max_tokens, 1);
    if (maximum_plan.schedule != Nvfp4GdnConvScheduleId::Materialized) { return 0; }

    WorkspaceLayoutBuilder layout;
    (void)allocate_workspace(layout, max_tokens);
    return layout.peak_bytes(1);
}

void nvfp4_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) {
    switch (nvfp4_gdn_conv_resolve_plan(policy, x.ne[1], 1).schedule) {
    case Nvfp4GdnConvScheduleId::DecodeFusedA16:
        nvfp4_gdn_snapshot_decode_launch(x, weight, conv_weight, conv_states, valid_columns,
                                         initial_slot, snapshot_base_slot, query, key, value, z,
                                         stream);
        return;
    case Nvfp4GdnConvScheduleId::SmallTFusedA16:
        nvfp4_gdn_snapshot_small_t_launch(x, weight, conv_weight, conv_states, valid_columns,
                                          initial_slot, snapshot_base_slot, query, key, value, z,
                                          stream);
        return;
    case Nvfp4GdnConvScheduleId::Materialized:
        break;
    }

    auto scope                         = workspace.scope();
    Nvfp4GdnProjectedWorkspace scratch = allocate_workspace(workspace, x.ne[1]);
    WorkspaceArena projection_workspace(scratch.projection);
    nvfp4_gdn_input_dispatch(x, weight, scratch.projected, z, LinearPolicy::AllowA4,
                             &projection_workspace, stream);
    nvfp4_gdn_snapshot_post_launch(scratch.projected, conv_weight, conv_states, valid_columns,
                                   initial_slot, snapshot_base_slot, query, key, value, stream);
