{
  "revision": "e652d55e2652137a4abcf1312357abbf3eb9d692",
  "file_sha256": "a3c74af27a7045c12f2893a8b7a91724c00d8a4148315c3165f3480c83016cf3",
  "header": "\n#include <metal_stdlib>\nusing namespace metal;\n\n// Total order for the block/dense output network: selected blocks by\n// ascending id, followed by invalid/padded lanes.\ninline bool qsa_index_before(\n    uint a_index, bool a_valid, uint b_index, bool b_valid) {\n    if (a_valid != b_valid) {\n        return a_valid;\n    }\n    return a_index < b_index;\n}\n\n// MLX's Metal ArgPartition currently delegates to its stable ascending merge\n// sort.  The v2.10 rows-gather lane consumes the final K indices directly, so\n// its selected valid blocks appear in ascending adjusted-score order.  Equal\n// values preserve input order, which is ascending block id.  Keep invalid\n// lanes after valid lanes inside the network; the epilogue places the exact\n// number of selected masked fillers before the winners.\ninline bool qsa_row_score_before(\n    float a_score,\n    uint a_index,\n    bool a_valid,\n    float b_score,\n    uint b_index,\n    bool b_valid) {\n    if (a_valid != b_valid) {\n        return a_valid;\n    }\n    if (!a_valid) {\n        return a_index < b_index;\n    }\n    bool a_nan = metal::isnan(a_score);\n    bool b_nan = metal::isnan(b_score);\n    if (a_nan || b_nan) {\n        if (a_nan != b_nan) {\n            return !a_nan;\n        }\n        return a_index < b_index;\n    }\n    if (a_score < b_score) {\n        return true;\n    }\n    if (b_score < a_score) {\n        return false;\n    }\n    return a_index < b_index;\n}\n\n// Monotonic IEEE-754 mapping: numerically larger finite floats produce larger\n// unsigned keys. Appending the block id makes the 64-bit key unique and mirrors\n// v2.10's stable ascending GPU sort followed by a final-K slice: if the\n// 1e-12 adjustment itself rounds away, the later/higher id wins the cutoff.\ninline uint qsa_float_order_key(float value) {\n    const uint bits = as_type<uint>(value);\n    return (bits & 0x80000000u) != 0 ? ~bits : (bits ^ 0x80000000u);\n}\n\ninline ulong qsa_composite_key(float adjusted_score, uint block_id) {\n    return (ulong(qsa_float_order_key(adjusted_score)) << 32) |\n           ulong(block_id);\n}\n\n// MLX's NAX float32 GEMM path truncates each operand to a 10-bit mantissa\n// before fp32 accumulation, but the legacy GEMV route remains full fp32.\n// A uniform runtime mask mirrors that eager-matmul contract without adding a\n// query-row or logical-history specialization to the Python kernel cache.\ninline float qsa_mlx_gemm_operand(float value, uint operand_mask) {\n    return as_type<float>(as_type<uint>(value) & operand_mask);\n}\n",
  "body": "\n        const uint row = threadgroup_position_in_grid.x;\n        const uint lane = thread_position_in_threadgroup.x;\n        const int qpos = pos_start[0] + int(row);\n        const int total = total_tokens[0];\n        const int logical_value = logical_blocks[0];\n        const uint logical = logical_value > 0\n            ? metal::min(uint(logical_value), BACKING_BLOCKS)\n            : 0u;\n        // Mirror Matmul::eval_gpu's check_transpose + batch-collapse route.\n        // A non-f32 astype materializes contiguous storage. For an existing\n        // f32 view, check_transpose preserves recognized row/transposed\n        // layouts and copies anything else. A copied operand is contiguous;\n        // a copied broadcast B, however, no longer has a zero S-batch stride\n        // and prevents folding S into M. When folding does not happen, M is\n        // HEADS. MLX sends min(M,N)==1 to full-fp32 GEMV before NAX/TF32.\n        const bool q_is_vector = HEADS == 1u;\n        const bool q_kept_untransposed = Q_INPUT_IS_FLOAT32 &&\n            q_strides[3] == 1u &&\n            (!q_is_vector || q_strides[2] == HEAD_DIM);\n        const bool q_kept_transposed = Q_INPUT_IS_FLOAT32 &&\n            !q_kept_untransposed && q_strides[2] == 1u &&\n            (!q_is_vector || q_strides[3] == HEADS);\n        const bool q_copied =\n            !Q_INPUT_IS_FLOAT32 ||\n            (!q_kept_untransposed && !q_kept_transposed);\n        const bool q_batch_contiguous = q_copied ||\n            (q_kept_untransposed && q_strides[2] == HEAD_DIM &&\n             q_strides[1] == size_t(HEADS) * HEAD_DIM);\n\n        // pooled:[1,N,D] is swapped to B:[1,1,D,N]. For N>1 its two\n        // recognized matrix layouts correspond to either original stride\n        // being one. Otherwise check_transpose copies the broadcast view.\n        const bool pooled_kept = !POOLED_INPUT_IS_FLOAT32 ||\n            pooled_strides[1] == 1u || pooled_strides[2] == 1u;\n        const bool collapse_s_into_m = q_shape[1] > 1u &&\n            !q_kept_transposed && q_batch_contiguous && pooled_kept;\n        const bool effective_m_gt_one =\n            HEADS > 1u || collapse_s_into_m;\n        const bool use_tf32_operands = ENABLE_TF32 &&\n            effective_m_gt_one && logical > 1u;\n        const uint gemm_operand_mask =\n            use_tf32_operands ? 0xffffe000u : 0xffffffffu;\n        const int complete_value = (qpos + 1) / int(RATIO);\n        const uint complete = complete_value > 0 ? uint(complete_value) : 0u;\n        const uint valid_count = metal::min(logical, complete);\n\n        threadgroup float exchange_scores[WIDTH];\n        threadgroup uint exchange_indices[WIDTH];\n        threadgroup uchar exchange_valid[WIDTH];\n        threadgroup atomic_uint radix_histogram[RADIX_BINS];\n        threadgroup atomic_uint selected_count;\n        threadgroup ulong radix_prefix;\n        threadgroup uint radix_rank;\n        threadgroup ulong threshold_key;\n\n        // Score every visible complete block once. The scratch output keeps\n        // those fp32 adjusted scores available to all eight radix passes\n        // without recomputing HEADS*HEAD_DIM dot products. Only the visible\n        // prefix is written/read; backing capacity does not tax an early row.\n        const size_t scratch_base = (size_t)row * BACKING_BLOCKS;\n        for (uint block = lane; block < valid_count; block += WIDTH) {\n            float score_sum = 0.0f;\n            for (uint head = 0; head < HEADS; ++head) {\n                float dot = 0.0f;\n                const size_t q_base =\n                    (size_t)row * q_strides[1] +\n                    (size_t)head * q_strides[2];\n                const size_t pooled_base =\n                    (size_t)block * pooled_strides[1];\n                for (uint dim = 0; dim < HEAD_DIM; ++dim) {\n                    const float q_value = qsa_mlx_gemm_operand(\n                        float(q[q_base + (size_t)dim * q_strides[3]]),\n                        gemm_operand_mask);\n                    const float k_value = qsa_mlx_gemm_operand(\n                        float(pooled[\n                            pooled_base + (size_t)dim * pooled_strides[2]]),\n                        gemm_operand_mask);\n                    dot += q_value * k_value;\n                }\n                score_sum += metal::max(dot, 0.0f);\n            }\n            const float score = score_sum / SQRT_HEAD_DIM;\n            const float adjusted = score - float(block) * 1.0e-12f;\n            score_scratch[scratch_base + block] = adjusted;\n        }\n        threadgroup_barrier(\n            mem_flags::mem_threadgroup | mem_flags::mem_device);\n\n        // Select the Kth largest strict composite key a byte at a time. Each\n        // pass scans only candidates matching the already-selected high-byte\n        // prefix. The rank is relative to that prefix bucket.\n        const uint k_eff = metal::min(TOP_K, valid_count);\n        if (lane == 0) {\n            radix_prefix = 0ul;\n            radix_rank = k_eff > 0 ? k_eff - 1 : 0;\n            threshold_key = 0xfffffffffffffffful;\n        }\n        threadgroup_barrier(mem_flags::mem_threadgroup);\n\n        if (k_eff > 0) {\n            for (uint pass = 0; pass < 8; ++pass) {\n                if (lane < RADIX_BINS) {\n                    atomic_store_explicit(\n                        &radix_histogram[lane], 0u, memory_order_relaxed);\n                }\n                threadgroup_barrier(mem_flags::mem_threadgroup);\n\n                const uint shift = 56u - pass * 8u;\n                const ulong prefix = radix_prefix;\n                for (uint block = lane; block < valid_count; block += WIDTH) {\n                    const float adjusted = score_scratch[scratch_base + block];\n                    const ulong key = qsa_composite_key(adjusted, block);\n                    bool prefix_matches = true;\n                    if (pass > 0) {\n                        prefix_matches = (key >> (shift + 8u)) == prefix;\n                    }\n                    if (prefix_matches) {\n                        const uint digit = uint((key >> shift) & 0xfful);\n                        atomic_fetch_add_explicit(\n                            &radix_histogram[digit], 1u, memory_order_relaxed);\n                    }\n                }\n                threadgroup_barrier(mem_flags::mem_threadgroup);\n\n                if (lane == 0) {\n                    uint rank = radix_rank;\n                    uint chosen = 0;\n                    for (int digit = 255; digit >= 0; --digit) {\n                        const uint count = atomic_load_explicit(\n                            &radix_histogram[uint(digit)], memory_order_relaxed);\n                        if (rank < count) {\n                            chosen = uint(digit);\n                            break;\n                        }\n                        rank -= count;\n                    }\n                    radix_prefix = (radix_prefix << 8) | ulong(chosen);\n                    radix_rank = rank;\n                }\n                threadgroup_barrier(mem_flags::mem_threadgroup);\n            }\n            if (lane == 0) {\n                threshold_key = radix_prefix;\n            }\n        }\n        threadgroup_barrier(mem_flags::mem_threadgroup);\n\n        // Compact the exact winners into WIDTH-sized threadgroup storage. The\n        // atomic arrival order is irrelevant: the following bitonic network\n        // establishes the mode's deterministic output order.\n        exchange_scores[lane] = -INFINITY;\n        exchange_indices[lane] = 0xffffffffu;\n        exchange_valid[lane] = 0;\n        if (lane == 0) {\n            atomic_store_explicit(\n                &selected_count, 0u, memory_order_relaxed);\n        }\n        threadgroup_barrier(mem_flags::mem_threadgroup);\n        if (k_eff > 0) {\n            const ulong threshold = threshold_key;\n            for (uint block = lane; block < valid_count; block += WIDTH) {\n                const float adjusted = score_scratch[scratch_base + block];\n                if (qsa_composite_key(adjusted, block) >= threshold) {\n                    const uint slot = atomic_fetch_add_explicit(\n                        &selected_count, 1u, memory_order_relaxed);\n                    if (slot < TOP_K) {\n                        exchange_scores[slot] = adjusted;\n                        exchange_indices[slot] = block;\n                        exchange_valid[slot] = 1;\n                    }\n                }\n            }\n        }\n        threadgroup_barrier(mem_flags::mem_threadgroup);\n\n        uint my_index = exchange_indices[lane];\n        bool my_valid = exchange_valid[lane] != 0;\n        float my_score = exchange_scores[lane];\n        for (uint sequence = 2; sequence <= WIDTH; sequence <<= 1) {\n            for (uint stride = sequence >> 1; stride > 0; stride >>= 1) {\n                exchange_scores[lane] = my_score;\n                exchange_indices[lane] = my_index;\n                exchange_valid[lane] = my_valid ? 1 : 0;\n                threadgroup_barrier(mem_flags::mem_threadgroup);\n\n                const uint partner = lane ^ stride;\n                const float other_score = exchange_scores[partner];\n                const uint other_index = exchange_indices[partner];\n                const bool other_valid = exchange_valid[partner] != 0;\n                threadgroup_barrier(mem_flags::mem_threadgroup);\n\n                const bool is_lower = (lane & stride) == 0;\n                const float a_score = is_lower ? my_score : other_score;\n                const uint a_index = is_lower ? my_index : other_index;\n                const bool a_valid = is_lower ? my_valid : other_valid;\n                const float b_score = is_lower ? other_score : my_score;\n                const uint b_index = is_lower ? other_index : my_index;\n                const bool b_valid = is_lower ? other_valid : my_valid;\n\n                const bool lower_wants_before = (lane & sequence) == 0;\n                const bool b_before_a = ROW_TOKEN_MODE\n                    ? qsa_row_score_before(\n                        b_score, b_index, b_valid,\n                        a_score, a_index, a_valid)\n                    : qsa_index_before(\n                        b_index, b_valid, a_index, a_valid);\n                const bool a_before_b = ROW_TOKEN_MODE\n                    ? qsa_row_score_before(\n                        a_score, a_index, a_valid,\n                        b_score, b_index, b_valid)\n                    : qsa_index_before(\n                        a_index, a_valid, b_index, b_valid);\n                const bool swap = lower_wants_before ? b_before_a : a_before_b;\n                if (swap) {\n                    my_score = is_lower ? b_score : a_score;\n                    my_index = is_lower ? b_index : a_index;\n                    my_valid = is_lower ? b_valid : a_valid;\n                }\n            }\n        }\n    ",
  "header_sha256": "d4fc7c3a7aabf30fb070dfbcd08c88dbd58203671ed08a5d7c2ba7d255641221",
  "body_sha256": "4af0dfa74e39b513470bc59980ededebfb081860676ed4a7197b8bd22ab3e8a8",
  "epilogues": {
    "blocks": {
      "outputs": [
        "block_ids",
        "block_valid",
        "adjusted_scores"
      ],
      "source": "\n        if (lane < TOP_K) {\n            const bool ok = my_valid;\n            const size_t out_at = (size_t)row * TOP_K + lane;\n            block_ids[out_at] = ok ? int(my_index) : 0;\n            block_valid[out_at] = ok;\n            adjusted_scores[out_at] = ok ? my_score : -INFINITY;\n        }\n        ",
      "sha256": "dea9615acddef6faf699262289e4c52c075cfc4bc9b9ecfa12676cb6d8dccf44"
    },
    "dense_mask": {
      "outputs": [
        "dense_mask"
      ],
      "source": "\n        // First establish the visible incomplete tail (and clear both skipped\n        // blocks and capacity beyond runtime T) in O(T/WIDTH) work per lane.\n        for (int token = int(lane); token < int(OUTPUT_TOKENS); token += WIDTH) {\n            const bool in_tail = token >= int(complete * RATIO);\n            const bool causal = token < total && token <= qpos;\n            dense_mask[(size_t)row * OUTPUT_TOKENS + token] =\n                in_tail && causal;\n        }\n        threadgroup_barrier(mem_flags::mem_device);\n\n        // Then mark each selected complete block. Sorted selected lanes own\n        // disjoint blocks, so these writes cannot race with one another.\n        if (lane < TOP_K && my_valid) {\n            const uint token0 = my_index * RATIO;\n            for (uint within = 0; within < RATIO; ++within) {\n                const uint token = token0 + within;\n                if (token < OUTPUT_TOKENS && int(token) < total &&\n                    int(token) <= qpos) {\n                    dense_mask[(size_t)row * OUTPUT_TOKENS + token] = true;\n                }\n            }\n        }\n        ",
      "sha256": "b36dd77cf4726d26236ebfd61924c198a1588e1ce3e04078c192d20931264073"
    },
    "row_tokens": {
      "outputs": [
        "token_ids",
        "token_valid"
      ],
      "source": "\n        exchange_indices[lane] = my_index;\n        exchange_valid[lane] = my_valid ? 1 : 0;\n        threadgroup_barrier(mem_flags::mem_threadgroup);\n\n        constexpr uint BLOCK_TOKEN_SLOTS = TOP_K * RATIO;\n        constexpr uint ROW_TOKEN_SLOTS = BLOCK_TOKEN_SLOTS + RATIO;\n        const size_t out_base = (size_t)row * ROW_TOKEN_SLOTS;\n        for (uint slot = lane; slot < ROW_TOKEN_SLOTS; slot += WIDTH) {\n            int token = 0;\n            bool ok = false;\n            if (slot < BLOCK_TOKEN_SLOTS) {\n                const uint block_slot = slot / RATIO;\n                const uint in_block = slot % RATIO;\n                // The eager GPU argpartition is a full stable ascending sort,\n                // sliced to its final min(K, logical) entries.  When a row\n                // sees fewer than K complete blocks, selected -inf fillers\n                // therefore precede the finite winners.  If logical<K, the\n                // Python path appends the remaining padding after them.\n                const uint selected_slots = metal::min(TOP_K, logical);\n                const uint invalid_prefix = selected_slots - k_eff;\n                const bool winner_slot =\n                    block_slot >= invalid_prefix &&\n                    block_slot < invalid_prefix + k_eff;\n                const uint winner = block_slot >= invalid_prefix\n                    ? block_slot - invalid_prefix\n                    : 0u;\n                ok = winner_slot && exchange_valid[winner] != 0;\n                if (ok) {\n                    token = int(exchange_indices[winner] * RATIO + in_block);\n                }\n            } else {\n                const uint in_tail = slot - BLOCK_TOKEN_SLOTS;\n                token = int(complete * RATIO + in_tail);\n                ok = token <= qpos;\n            }\n            token_ids[out_base + slot] = ok ? token : 0;\n            token_valid[out_base + slot] = ok;\n        }\n        ",
      "sha256": "e64f3e99e480faab4ad2cafdf39e594a5a9766beba367f349ec8c47d57b71ca3"
    }
  },
  "prefill": {
    "file_sha256": "4d6fd428243c001746f69f8aed45991356772c2bd4a45586eb3c6813c91998d3",
    "mpp_header": "\n#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>\nusing namespace metal;\n\nconstant constexpr uint QSA_SCORE_HEADS = 4;\nconstant constexpr uint QSA_SCORE_HEAD_DIM = 128;\nconstant constexpr uint QSA_SCORE_QUERY_TILE = 16;\nconstant constexpr uint QSA_SCORE_KEY_TILE = 32;\nconstant constexpr uint QSA_SCORE_QUERIES_PER_SIMDGROUP = 4;\nconstant constexpr uint QSA_SCORE_SIMDGROUPS = 4;\nconstant constexpr uint QSA_SCORE_THREADS = 128;\nconstant constexpr uint QSA_SCORE_K_FRAGMENTS = 8;\nconstant constexpr float QSA_SCORE_SQRT_HEAD_DIM = 11.313708498984761f;\n\n// Metal 4's 16x32x16 cooperative fragment maps each lane to two rows eight\n// apart and four adjacent columns. This is the same proven map used by MLX's\n// NAX tiles and mtplx/kernels/sdpa_nax_tile.py.\ninline short2 qsa_score_nax_coord(ushort lane) {\n    const short qid = short(lane >> 2);\n    const short fragment_row =\n        ((qid & 4) | ((short(lane) >> 1) & 3));\n    const short fragment_col =\n        ((qid & 2) | (short(lane) & 1)) * 4;\n    return short2{fragment_col, fragment_row};\n}\n",
    "mpp_body": "\n    const uint tid = thread_position_in_threadgroup.x;\n    const uint simdgroup = tid >> 5;\n    const ushort lane = ushort(tid & 31u);\n    const uint rows = q_shape[1];\n    const uint blocks = pooled_shape[1];\n    const uint key_tiles = (blocks + QSA_SCORE_KEY_TILE - 1u) /\n        QSA_SCORE_KEY_TILE;\n    const uint tile = threadgroup_position_in_grid.x;\n    const uint query_tile = tile / key_tiles;\n    const uint key_tile = tile - query_tile * key_tiles;\n    const uint query0 =\n        query_tile * QSA_SCORE_QUERY_TILE +\n        simdgroup * QSA_SCORE_QUERIES_PER_SIMDGROUP;\n    const uint block0 = key_tile * QSA_SCORE_KEY_TILE;\n\n    // Four simdgroups share one [32,128] pooled-key tile (8 KiB for fp16 or\n    // bf16). Each key is therefore fetched once per sixteen output query rows,\n    // rather than once per head or per query row.\n    threadgroup InT pooled_tile[QSA_SCORE_KEY_TILE * QSA_SCORE_HEAD_DIM];\n    constexpr uint VECTORS_PER_KEY = QSA_SCORE_HEAD_DIM / 4u;\n    constexpr uint TILE_VECTORS = QSA_SCORE_KEY_TILE * VECTORS_PER_KEY;\n    for (uint item = tid; item < TILE_VECTORS; item += QSA_SCORE_THREADS) {\n        const uint key_local = item / VECTORS_PER_KEY;\n        const uint dim4 = item - key_local * VECTORS_PER_KEY;\n        const uint block = block0 + key_local;\n        vec<InT, 4> values = vec<InT, 4>(InT(0));\n        if (block < blocks) {\n            const int64_t source =\n                int64_t(block) * pooled_strides[1] +\n                int64_t(dim4 * 4u) * pooled_strides[2];\n            // Source layout is intentionally not guessed on the host.  MLX\n            // injects the real strides for this invocation; scalar gathers\n            // also avoid imposing an unobservable vec4 base-alignment\n            // contract on sliced/as_strided views.  Contiguous lanes remain\n            // adjacent and can still be coalesced by the compiler/hardware.\n            for (uint elem = 0u; elem < 4u; ++elem) {\n                values[elem] = pooled[\n                    source + int64_t(elem) * pooled_strides[2]];\n            }\n        }\n        const uint destination =\n            key_local * QSA_SCORE_HEAD_DIM + dim4 * 4u;\n        for (uint elem = 0u; elem < 4u; ++elem) {\n            pooled_tile[destination + elem] = values[elem];\n        }\n    }\n    threadgroup_barrier(mem_flags::mem_threadgroup);\n\n    constexpr auto descriptor = mpp::tensor_ops::matmul2d_descriptor(\n        16, 32, 16, false, true, true,\n        mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate);\n    mpp::tensor_ops::matmul2d<descriptor, metal::execution_simdgroup> matmul;\n    auto left = matmul.get_left_input_cooperative_tensor<InT, InT, float>();\n    auto right = matmul.get_right_input_cooperative_tensor<InT, InT, float>();\n    auto accumulator = matmul.get_destination_cooperative_tensor<\n        decltype(left), decltype(right), float>();\n\n    constexpr short ELEMENTS_PER_FRAGMENT = 8;\n    constexpr short ELEMENT_COLUMNS = 4;\n    constexpr short ELEMENT_ROW_JUMP = 8;\n    const short2 coordinate = qsa_score_nax_coord(lane);\n    for (short item = 0; item < 2 * ELEMENTS_PER_FRAGMENT; ++item) {\n        accumulator[item] = 0.0f;\n    }\n\n    // A's sixteen rows are head-major [h0:q0..q3, h1:q0..q3, ...].\n    // This layout makes lane^16 pair the h0/h2 carrier with h1/h3 for the\n    // same query and output columns after the TensorOp completes.\n    for (uint k_frag = 0u;\n         k_frag < QSA_SCORE_K_FRAGMENTS;\n         ++k_frag) {\n        for (short row_part = 0; row_part < 2; ++row_part) {\n            const uint matrix_row =\n                uint(coordinate.y + row_part * ELEMENT_ROW_JUMP);\n            const uint head = matrix_row / QSA_SCORE_QUERIES_PER_SIMDGROUP;\n            const uint query_local =\n                matrix_row - head * QSA_SCORE_QUERIES_PER_SIMDGROUP;\n            const uint query = query0 + query_local;\n            vec<InT, 4> values = vec<InT, 4>(InT(0));\n            if (query < rows) {\n                const int64_t source =\n                    int64_t(query) * q_strides[1] +\n                    int64_t(head) * q_strides[2] +\n                    int64_t(k_frag * 16u + uint(coordinate.x)) *\n                        q_strides[3];\n                for (uint elem = 0u; elem < 4u; ++elem) {\n                    values[elem] = q[\n                        source + int64_t(elem) * q_strides[3]];\n                }\n            }\n            for (short elem = 0; elem < ELEMENT_COLUMNS; ++elem) {\n                left[row_part * ELEMENT_COLUMNS + elem] = values[elem];\n            }\n        }\n\n        for (short key_half = 0; key_half < 2; ++key_half) {\n            for (short row_part = 0; row_part < 2; ++row_part) {\n                const uint key_local = uint(\n                    key_half * 16 + coordinate.y +\n                    row_part * ELEMENT_ROW_JUMP);\n                const uint source =\n                    key_local * QSA_SCORE_HEAD_DIM +\n                    k_frag * 16u + uint(coordinate.x);\n                const threadgroup vec<InT, 4>* source4 =\n                    reinterpret_cast<const threadgroup vec<InT, 4>*>(\n                        pooled_tile + source);\n                const vec<InT, 4> values = source4[0];\n                for (short elem = 0; elem < ELEMENT_COLUMNS; ++elem) {\n                    right[\n                        key_half * ELEMENTS_PER_FRAGMENT +\n                        row_part * ELEMENT_COLUMNS + elem] = values[elem];\n                }\n            }\n        }\n        matmul.run(left, right, accumulator);\n    }\n\n    // For a lane with bit 4 clear, row_part 0 carries h0 and row_part 1\n    // carries h2. lane^16 carries h1 and h3 for exactly the same query and\n    // four output columns. Apply ReLU before the strictly h0,h1,h2,h3-ordered\n    // float32 sum, matching the eager expression's reduction contract.\n    for (short key_half = 0; key_half < 2; ++key_half) {\n        for (short elem = 0; elem < ELEMENT_COLUMNS; ++elem) {\n            const float head0_or_1 = metal::max(\n                accumulator[key_half * ELEMENTS_PER_FRAGMENT + elem], 0.0f);\n            const float head2_or_3 = metal::max(\n                accumulator[\n                    key_half * ELEMENTS_PER_FRAGMENT +\n                    ELEMENT_COLUMNS + elem],\n                0.0f);\n            const float paired_head1_or_0 =\n                simd_shuffle_xor(head0_or_1, ushort(16));\n            const float paired_head3_or_2 =\n                simd_shuffle_xor(head2_or_3, ushort(16));\n            if ((lane & 16u) == 0u) {\n                const uint query = query0 + uint(coordinate.y);\n                const uint block =\n                    block0 + uint(key_half * 16 + coordinate.x + elem);\n                if (query < rows && block < blocks) {\n                    const float head_sum =\n                        ((head0_or_1 + paired_head1_or_0) + head2_or_3) +\n                        paired_head3_or_2;\n                    scores[size_t(query) * blocks + block] =\n                        head_sum / QSA_SCORE_SQRT_HEAD_DIM;\n                }\n            }\n        }\n    }\n",
    "topk_header": "\nconstant constexpr uint BACKING_BLOCKS = @blocks@;\nconstant constexpr uint TOP_K = @topk@;\nconstant constexpr uint RATIO = @ratio@;\nconstant constexpr uint WIDTH = @width@;\nconstant constexpr uint RADIX_BINS = @_RADIX_BINS@;\nconstant constexpr uint FINAL_CANDIDATES = @_FINAL_CANDIDATES@;\nconstant constexpr uint INSERTION_CANDIDATES = @_INSERTION_CANDIDATES@;\nconstant constexpr uint RADIX_PASSES = 6;\nconstant constexpr uint OUTPUT_TOKENS = @output_tokens@;\nconstant constexpr bool ROW_TOKEN_MODE = @str(mode == 'row_tokens').lower()@;\n\ninline uint qsa_prefill_radix_shift(uint pass) {\n    return pass == 0u ? 53u\n        : pass == 1u ? 42u\n        : pass == 2u ? 31u\n        : pass == 3u ? 20u\n        : pass == 4u ? 9u\n        : 0u;\n}\n\ninline uint qsa_prefill_radix_bits(uint pass) {\n    return pass < 5u ? 11u : 9u;\n}\n",
    "topk_body": "\n        const uint row = threadgroup_position_in_grid.x;\n        const uint lane = thread_position_in_threadgroup.x;\n        const int qpos = pos_start[0] + int(row);\n        const int total = total_tokens[0];\n        const int logical_value = logical_blocks[0];\n        const uint logical = logical_value > 0\n            ? metal::min(uint(logical_value), BACKING_BLOCKS)\n            : 0u;\n        const int complete_value = (qpos + 1) / int(RATIO);\n        const uint complete = complete_value > 0 ? uint(complete_value) : 0u;\n        const uint valid_count = metal::min(logical, complete);\n        const uint k_eff = metal::min(TOP_K, valid_count);\n        const size_t score_base = (size_t)row * BACKING_BLOCKS;\n\n        threadgroup float exchange_scores[WIDTH];\n        threadgroup uint exchange_indices[WIDTH];\n        threadgroup uchar exchange_valid[WIDTH];\n        threadgroup atomic_uint radix_histogram[RADIX_BINS];\n        threadgroup ulong final_candidate_keys[FINAL_CANDIDATES];\n        threadgroup atomic_uint selected_count;\n        threadgroup atomic_uint final_candidate_count;\n        threadgroup uint final_candidate_total;\n        threadgroup ulong radix_prefix;\n        threadgroup uint radix_rank;\n        threadgroup uint radix_done;\n        threadgroup uint radix_final_shift;\n        threadgroup uint radix_next_pass;\n        threadgroup ulong threshold_key;\n\n        // Find a <=2048-item bucket containing the Kth-largest strict\n        // (adjusted-float, block-id) key. Six high-to-low radix digits cover\n        // all 64 bits as 11/11/11/11/11/9. Usually the first 11-bit pass is\n        // already selective enough; unlike the decode selector, prefill never\n        // performs eight unconditional full-row byte scans.\n        if (lane == 0) {\n            radix_prefix = 0ul;\n            radix_rank = k_eff > 0 ? k_eff - 1 : 0;\n            radix_done = k_eff > 0 ? 0u : 1u;\n            radix_final_shift = 0u;\n            radix_next_pass = 0u;\n            threshold_key = 0xfffffffffffffffful;\n        }\n        threadgroup_barrier(mem_flags::mem_threadgroup);\n\n        if (k_eff > 0) {\n            for (uint pass = 0; pass < RADIX_PASSES; ++pass) {\n                if (radix_done == 0u) {\n                    for (uint bin = lane; bin < RADIX_BINS; bin += WIDTH) {\n                        atomic_store_explicit(\n                            &radix_histogram[bin], 0u, memory_order_relaxed);\n                    }\n                    threadgroup_barrier(mem_flags::mem_threadgroup);\n\n                    const uint shift = qsa_prefill_radix_shift(pass);\n                    const uint bits = qsa_prefill_radix_bits(pass);\n                    const uint digit_mask = (1u << bits) - 1u;\n                    const ulong prefix = radix_prefix;\n                    for (uint block = lane; block < valid_count; block += WIDTH) {\n                        const float adjusted =\n                            scores[score_base + block] - float(block) * 1.0e-12f;\n                        const ulong key = qsa_composite_key(adjusted, block);\n                        const bool prefix_matches = pass == 0\n                            ? true\n                            : (key >> (shift + bits)) == prefix;\n                        if (prefix_matches) {\n                            const uint digit = uint((key >> shift) & digit_mask);\n                            atomic_fetch_add_explicit(\n                                &radix_histogram[digit], 1u,\n                                memory_order_relaxed);\n                        }\n                    }\n                    threadgroup_barrier(mem_flags::mem_threadgroup);\n\n                    if (lane == 0) {\n                        uint rank = radix_rank;\n                        uint chosen = 0;\n                        uint chosen_count = 0;\n                        const int max_digit = int(digit_mask);\n                        for (int digit = max_digit; digit >= 0; --digit) {\n                            const uint count = atomic_load_explicit(\n                                &radix_histogram[uint(digit)],\n                                memory_order_relaxed);\n                            if (rank < count) {\n                                chosen = uint(digit);\n                                chosen_count = count;\n                                break;\n                            }\n                            rank -= count;\n                        }\n                        radix_prefix = (radix_prefix << bits) | ulong(chosen);\n                        radix_rank = rank;\n                        radix_final_shift = shift;\n                        radix_next_pass = pass + 1u;\n                        if (chosen_count <= FINAL_CANDIDATES) {\n                            radix_done = 1u;\n                        }\n                    }\n                    threadgroup_barrier(mem_flags::mem_threadgroup);\n                }\n                // All lanes observe radix_done uniformly before the next pass.\n                threadgroup_barrier(mem_flags::mem_threadgroup);\n            }\n\n            // Collect only the final threshold bucket. The adaptive descent\n            // guarantees that its cardinality is at most FINAL_CANDIDATES.\n            if (lane == 0) {\n                atomic_store_explicit(\n                    &final_candidate_count, 0u, memory_order_relaxed);\n            }\n            threadgroup_barrier(mem_flags::mem_threadgroup);\n\n            const ulong final_prefix = radix_prefix;\n            const uint final_shift = radix_final_shift;\n            for (uint block = lane; block < valid_count; block += WIDTH) {\n                const float adjusted =\n                    scores[score_base + block] - float(block) * 1.0e-12f;\n                const ulong key = qsa_composite_key(adjusted, block);\n                if ((key >> final_shift) == final_prefix) {\n                    const uint slot = atomic_fetch_add_explicit(\n                        &final_candidate_count, 1u, memory_order_relaxed);\n                    if (slot < FINAL_CANDIDATES) {\n                        final_candidate_keys[slot] = key;\n                    }\n                }\n            }\n            threadgroup_barrier(mem_flags::mem_threadgroup);\n\n            if (lane == 0) {\n                final_candidate_total = atomic_load_explicit(\n                    &final_candidate_count, memory_order_relaxed);\n            }\n            threadgroup_barrier(mem_flags::mem_threadgroup);\n\n            // vLLM's prefill selector uses insertion for the common small-row\n            // regime and radix work for the long-row regime. Make the same\n            // choice from the actual threshold-bucket cardinality: insertion\n            // is capped at 64 candidates (at most 4096 key comparisons), while\n            // a larger bucket finishes the remaining 64-bit radix digits over\n            // at most 2048 resident keys. This avoids both an O(2048^2) tail\n            // and a padded 2048-key bitonic network's 66 barriers.\n            const uint candidate_count = final_candidate_total;\n            if (candidate_count <= INSERTION_CANDIDATES) {\n                // Composite keys are strict because block_id occupies their\n                // low 32 bits, so exactly one candidate has this rank.\n                for (uint item = lane; item < candidate_count; item += WIDTH) {\n                    const ulong key = final_candidate_keys[item];\n                    uint greater_rank = 0u;\n                    for (uint other = 0u; other < candidate_count; ++other) {\n                        greater_rank +=\n                            final_candidate_keys[other] > key ? 1u : 0u;\n                    }\n                    if (greater_rank == radix_rank) {\n                        threshold_key = key;\n                    }\n                }\n            } else {\n                // The first phase has already consumed [0, radix_next_pass)\n                // and all resident candidates share radix_prefix. Complete\n                // every remaining digit using the same 2048-bin histogram,\n                // but scan only the bounded candidate array, never the full\n                // score row. After pass 5, radix_prefix is the exact unique\n                // 64-bit threshold key.\n                const uint first_refine_pass = radix_next_pass;\n                for (uint pass = 0u; pass < RADIX_PASSES; ++pass) {\n                    if (pass >= first_refine_pass) {\n                        for (uint bin = lane; bin < RADIX_BINS; bin += WIDTH) {\n                            atomic_store_explicit(\n                                &radix_histogram[bin], 0u,\n                                memory_order_relaxed);\n                        }\n                        threadgroup_barrier(mem_flags::mem_threadgroup);\n\n                        const uint shift = qsa_prefill_radix_shift(pass);\n                        const uint bits = qsa_prefill_radix_bits(pass);\n                        const uint digit_mask = (1u << bits) - 1u;\n                        const ulong prefix = radix_prefix;\n                        for (uint item = lane;\n                             item < candidate_count;\n                             item += WIDTH) {\n                            const ulong key = final_candidate_keys[item];\n                            if ((key >> (shift + bits)) == prefix) {\n                                const uint digit =\n                                    uint((key >> shift) & digit_mask);\n                                atomic_fetch_add_explicit(\n                                    &radix_histogram[digit], 1u,\n                                    memory_order_relaxed);\n                            }\n                        }\n                        threadgroup_barrier(mem_flags::mem_threadgroup);\n\n                        if (lane == 0) {\n                            uint rank = radix_rank;\n                            uint chosen = 0u;\n                            const int max_digit = int(digit_mask);\n                            for (int digit = max_digit; digit >= 0; --digit) {\n                                const uint count = atomic_load_explicit(\n                                    &radix_histogram[uint(digit)],\n                                    memory_order_relaxed);\n                                if (rank < count) {\n                                    chosen = uint(digit);\n                                    break;\n                                }\n                                rank -= count;\n                            }\n                            radix_prefix =\n                                (radix_prefix << bits) | ulong(chosen);\n                            radix_rank = rank;\n                        }\n                        threadgroup_barrier(mem_flags::mem_threadgroup);\n                    }\n                    threadgroup_barrier(mem_flags::mem_threadgroup);\n                }\n                if (lane == 0) {\n                    threshold_key = radix_prefix;\n                }\n            }\n        }\n        threadgroup_barrier(mem_flags::mem_threadgroup);\n\n        // Compact the exact winners. Atomic arrival order is canonicalized by\n        // the same bitonic output network as the decode-oriented selector.\n        exchange_scores[lane] = -INFINITY;\n        exchange_indices[lane] = 0xffffffffu;\n        exchange_valid[lane] = 0;\n        if (lane == 0) {\n            atomic_store_explicit(&selected_count, 0u, memory_order_relaxed);\n        }\n        threadgroup_barrier(mem_flags::mem_threadgroup);\n        if (k_eff > 0) {\n            const ulong threshold = threshold_key;\n            for (uint block = lane; block < valid_count; block += WIDTH) {\n                const float adjusted =\n                    scores[score_base + block] - float(block) * 1.0e-12f;\n                if (qsa_composite_key(adjusted, block) >= threshold) {\n                    const uint slot = atomic_fetch_add_explicit(\n                        &selected_count, 1u, memory_order_relaxed);\n                    if (slot < TOP_K) {\n                        exchange_scores[slot] = adjusted;\n                        exchange_indices[slot] = block;\n                        exchange_valid[slot] = 1;\n                    }\n                }\n            }\n        }\n        threadgroup_barrier(mem_flags::mem_threadgroup);\n\n        uint my_index = exchange_indices[lane];\n        bool my_valid = exchange_valid[lane] != 0;\n        float my_score = exchange_scores[lane];\n        for (uint sequence = 2; sequence <= WIDTH; sequence <<= 1) {\n            for (uint stride = sequence >> 1; stride > 0; stride >>= 1) {\n                exchange_scores[lane] = my_score;\n                exchange_indices[lane] = my_index;\n                exchange_valid[lane] = my_valid ? 1 : 0;\n                threadgroup_barrier(mem_flags::mem_threadgroup);\n\n                const uint partner = lane ^ stride;\n                const float other_score = exchange_scores[partner];\n                const uint other_index = exchange_indices[partner];\n                const bool other_valid = exchange_valid[partner] != 0;\n                threadgroup_barrier(mem_flags::mem_threadgroup);\n\n                const bool is_lower = (lane & stride) == 0;\n                const float a_score = is_lower ? my_score : other_score;\n                const uint a_index = is_lower ? my_index : other_index;\n                const bool a_valid = is_lower ? my_valid : other_valid;\n                const float b_score = is_lower ? other_score : my_score;\n                const uint b_index = is_lower ? other_index : my_index;\n                const bool b_valid = is_lower ? other_valid : my_valid;\n\n                const bool lower_wants_before = (lane & sequence) == 0;\n                const bool b_before_a = ROW_TOKEN_MODE\n                    ? qsa_row_score_before(\n                        b_score, b_index, b_valid,\n                        a_score, a_index, a_valid)\n                    : qsa_index_before(\n                        b_index, b_valid, a_index, a_valid);\n                const bool a_before_b = ROW_TOKEN_MODE\n                    ? qsa_row_score_before(\n                        a_score, a_index, a_valid,\n                        b_score, b_index, b_valid)\n                    : qsa_index_before(\n                        a_index, a_valid, b_index, b_valid);\n                const bool swap = lower_wants_before ? b_before_a : a_before_b;\n                if (swap) {\n                    my_score = is_lower ? b_score : a_score;\n                    my_index = is_lower ? b_index : a_index;\n                    my_valid = is_lower ? b_valid : a_valid;\n                }\n            }\n        }\n        ",
    "mpp_header_sha256": "eddbff46fac2d4a09e1a1bf1e07b4d10afa2c2bb761304b4ea8683181131176a",
    "mpp_body_sha256": "138bfe4485b8c9b979278d6b112fbe45658b37827bd90135594f99dc5328484d",
    "topk_header_sha256": "69173e1a65f9e19aa6d684e218168cb9eb62dfe4b91c5ecdafcfc3d9199c1a5f",
    "topk_body_sha256": "32da29054cab798ad08bdb982b746bd30c6f130abb6dd10e438b97058ece870d"
  }
}