            plan->prepare_mtp = true;
            plan->mtp_bridge  = plan->reuse_base < plan->summary.prompt_tokens
                                    ? MtpBridgeMode::BeforeSuffix
                                    : MtpBridgeMode::AfterExactHit;
        } else if (is_rewrite_checkpoint_restore(plan->reuse) ||
                   plan->reuse == ReusePath::PrivateLongAnchor ||
                   plan->reuse == ReusePath::SharedStablePrefix) {
            plan->prepare_mtp = true;
            plan->mtp_bridge  = plan->reuse_base < plan->summary.prompt_tokens
                                    ? MtpBridgeMode::BeforeSuffix
                                    : MtpBridgeMode::AfterExactHit;
        }
    }

    if (base.vision_control_plan) {
        VisionPrefillPlan vision;
        vision.control_plan = base.vision_control_plan;
        vision.uses.reserve(base.vision_control_plan->items.size());
        std::size_t max_merged = 0;
        for (std::size_t index = 0; index < base.vision_control_plan->items.size(); ++index) {
            const qwen3_6::VisionItemControlPlan& item = base.vision_control_plan->items[index];
            if (item.token_end <= plan->reuse_base) { continue; }
            const std::uint32_t begin = plan->prepare_mtp && item.token_begin != 0
                                            ? item.token_begin - 1
                                            : item.token_begin;
            vision.uses.push_back(VisionUseSpan{
                .begin               = begin,
                .end                 = item.token_end,
                .prepared_item_index = static_cast<std::uint32_t>(index),
            });
            max_merged = std::max(max_merged, item.merged_count);
        }
        if (!vision.uses.empty()) {
            vision.max_merged_count = max_merged;
            plan->vision            = std::move(vision);
        }
    }

    const std::size_t prefill_splits = plan->vision ? plan->vision->uses.size() : 0ULL;
    plan->summary.service_work_quanta =
        projected_service_work(plan->summary, plan->reuse_base, prefill_chunk, prefill_splits,
                               plan->capture_groups, prompt.identity.rewrite_execution_frontiers);
    std::uint64_t remaining_vision_items   = 0;
    std::uint64_t remaining_vision_patches = 0;
    if (plan->vision) {
        std::vector<bool> counted(prompt.vision_items.size(), false);
        for (const VisionUseSpan& use : plan->vision->uses) {
            if (use.prepared_item_index >= counted.size()) {
                throw std::logic_error("Vision cost input references a missing item");
            }
            if (counted[use.prepared_item_index]) { continue; }
            counted[use.prepared_item_index] = true;
            if (remaining_vision_items == std::numeric_limits<std::uint64_t>::max()) {
                throw std::overflow_error("Vision prefill item count exceeds uint64");
            }
            ++remaining_vision_items;
            const std::size_t patches = prompt.vision_items[use.prepared_item_index].patch_count;
            if (patches > std::numeric_limits<std::uint64_t>::max() - remaining_vision_patches) {
                throw std::overflow_error("Vision prefill cost input exceeds uint64");
            }
            remaining_vision_patches += static_cast<std::uint64_t>(patches);
        }
    }
    plan->remaining_prefill_work =
        scheduled_prefill_work(plan->reuse_base, plan->summary.prompt_tokens,
                               remaining_vision_items, remaining_vision_patches, prefill_chunk,
                               plan->capture_groups, prompt.identity.rewrite_execution_frontiers);
    plan->transfer_requirements.reserve(4);
    const auto add_state_transfer = [&](runtime::ContextTransferDirection direction,
                                        bool dflash_local_only = false) {
        const TransferWork work = dflash_local_only
                                      ? dflash_local_transfer_work(state_images->host_layout())
                                      : state_image_transfer_work(state_images->host_layout());
        plan->transfer_requirements.push_back(runtime::ContextTransferRequirement{
            .resource   = runtime::ContextResourceClass::State,
            .direction  = direction,
            .units      = 1,
            .page_count = 0,
            .work       = work,
        });
    };
    const auto add_kv_transfer = [&](runtime::ContextResourceClass resource,
                                     runtime::ContextTransferDirection direction,
                                     const LogicalKVPageStore& pages, std::uint32_t page_count,
                                     std::uint32_t contiguous_runs = 1) {
        if (page_count == 0) { return; }
        const HostKVPageLayout layout = plan_host_kv_page_layout(pages.physical_pool().geometry());
        const TransferWork work =
            direction == runtime::ContextTransferDirection::DeviceToDevice
                ? plan_device_kv_copy_work(layout, page_count)
                : plan_host_kv_transfer_work(layout, page_count, contiguous_runs);
        plan->transfer_requirements.push_back(runtime::ContextTransferRequirement{
            .resource   = resource,
            .direction  = direction,
            .units      = work.payload_bytes,
            .page_count = page_count,
            .work       = work,
        });
    };
    const auto missing_kv_restore = [&](const KVAddressSpaceStore& addresses,
                                        const LogicalKVPageStore& pages,
                                        KVAddressSpaceHandle address, std::uint32_t required) {
        std::pair<std::uint32_t, std::uint32_t> out;
        if (required > addresses.mapped_pages(address)) {
            throw std::logic_error("checkpoint KV requirement exceeds address membership");
        }
        std::optional<HostKVPageReplica> previous;
        for (std::uint32_t page = 0; page < required; ++page) {
            const LogicalKVPageHandle logical = addresses.logical_page(address, page);
            if (pages.device_resident(logical)) { continue; }
            if (!pages.host_resident(logical)) {
                throw std::logic_error("checkpoint KV page has no restorable replica");
            }
            const HostKVPageReplica replica = pages.host_replica(logical);
            if (!previous || previous->extent != replica.extent ||
                previous->page_offset + 1U != replica.page_offset) {
                ++out.second;
            }
            previous = replica;
            ++out.first;
        }
        return out;
    };
    const detail::PhysicalResources source_resources =
        source != nullptr ? resident_resources(*source) : detail::PhysicalResources{};
    plan->source_resources = source_resources;
    if (source != nullptr || shared_source != nullptr) {
        const StateImageHandle selected =
            source != nullptr ? selected_state(*source, plan->reuse, plan->selected_checkpoint)
                              : shared_source->state;
        const StateReplicaResidency selected_state_residency = state_store->residency(selected);
        plan->needs_transfer = selected_state_residency == StateReplicaResidency::HostOnly;
        if (plan->needs_transfer) {
            add_state_transfer(runtime::ContextTransferDirection::HostToDevice);
        }
        const bool splits_private_both =
            source != nullptr && selected_state_residency == StateReplicaResidency::Both;
        const bool forks_device_state =
            (plan->source_disposition == runtime::ClaimDisposition::Retained &&
             !splits_private_both) ||
            (plan->source_disposition == runtime::ClaimDisposition::ConsumedToActive &&
             plan->state_fork_required);
        if (speculative_backend == SpeculativeBackend::DFlash && forks_device_state &&
            selected_state_residency != StateReplicaResidency::HostOnly) {
            // DFlash has lane-local recurrent state which is copied eagerly when a retained
            // checkpoint forks. The copy completes on the compute stream in the publication
            // boundary, so it contributes candidate cost without turning the plan into an
            // asynchronous materialization.
            add_state_transfer(runtime::ContextTransferDirection::DeviceToDevice, true);
        }
        const SequenceKVBundle* source_kv =
            source != nullptr ? (source->kv ? &*source->kv : nullptr)
                              : (shared_source->kv ? &*shared_source->kv : nullptr);
        if (source_kv == nullptr) { throw std::logic_error("checkpoint has no KV address space"); }
        const std::uint32_t backend_frontier =
            backend_frontier_at(speculative_backend, plan->reuse_base);
        if (plan->source_disposition == runtime::ClaimDisposition::Retained) {
            plan->text_prefix_fork_required    = true;
            plan->backend_prefix_fork_required = source_kv->backend && backend_frontier != 0;
        } else if (source != nullptr) {
            plan->text_prefix_fork_required =
                partial_tail_cow_required(*text_kv_addresses, source_kv->text, plan->reuse_base);
            plan->backend_prefix_fork_required =
                source_kv->backend && backend_frontier != 0 &&
                partial_tail_cow_required(*backend_kv_addresses, *source_kv->backend,
                                          backend_frontier);
        }
        const std::uint32_t main_required = pages_for_tokens(plan->reuse_base);
        const std::uint32_t main_device =
            device_kv_prefix_pages(*text_kv_addresses, source_kv->text, plan->reuse_base);
        if (main_device > main_required) {
            throw std::logic_error("Text KV resident prefix exceeds checkpoint requirement");
        }
        const auto [main_missing, main_contiguous_runs] =
            missing_kv_restore(*text_kv_addresses, *text_kv_pages, source_kv->text, main_required);
        if (main_missing != main_required - main_device) {
            throw std::logic_error("Text KV restore inventory is inconsistent");
        }
        plan->needs_transfer = plan->needs_transfer || main_missing != 0;
        if (main_missing != 0) {
            add_kv_transfer(runtime::ContextResourceClass::MainKV,
                            runtime::ContextTransferDirection::HostToDevice, *text_kv_pages,
                            main_missing, main_contiguous_runs);
        }
        if (source_kv->backend && backend_frontier != 0) {
            const std::uint32_t backend_required = pages_for_tokens(backend_frontier);
            const std::uint32_t backend_device   = device_kv_prefix_pages(
                *backend_kv_addresses, *source_kv->backend, backend_frontier);
            if (backend_device > backend_required) {
                throw std::logic_error("Backend KV resident prefix exceeds checkpoint requirement");
            }
            const auto [backend_missing, backend_contiguous_runs] = missing_kv_restore(
                *backend_kv_addresses, *backend_kv_pages, *source_kv->backend, backend_required);
            if (backend_missing != backend_required - backend_device) {
                throw std::logic_error("Backend KV restore inventory is inconsistent");
            }
            plan->needs_transfer = plan->needs_transfer || backend_missing != 0;
            if (backend_missing != 0) {
                add_kv_transfer(runtime::ContextResourceClass::BackendKV,
                                runtime::ContextTransferDirection::HostToDevice, *backend_kv_pages,
                                backend_missing, backend_contiguous_runs);
            }
        }
        if (plan->text_prefix_fork_required) {
            const std::uint32_t page_size = static_cast<std::uint32_t>(kPagedKVPageSize);
            if (plan->reuse_base % page_size != 0) {
                add_kv_transfer(runtime::ContextResourceClass::MainKV,
                                runtime::ContextTransferDirection::DeviceToDevice, *text_kv_pages,
                                1);
                plan->needs_transfer = true;
            }
            if (plan->backend_prefix_fork_required && backend_frontier % page_size != 0) {
                add_kv_transfer(runtime::ContextResourceClass::BackendKV,
                                runtime::ContextTransferDirection::DeviceToDevice,
                                *backend_kv_pages, 1);
                plan->needs_transfer = true;
            }
        }
    }
    detail::PhysicalDeviceResources active = base.root_demand.active_entitlement.device;
    if (plan->active_optional_resources.device.state_slots >
        std::numeric_limits<std::uint32_t>::max() - active.state_slots) {
        throw std::overflow_error("active optional StateImage entitlement overflow");
    }
    active.state_slots += plan->active_optional_resources.device.state_slots;
    if ((source != nullptr || shared_source != nullptr) &&
        plan->source_disposition == runtime::ClaimDisposition::Retained) {
        const SequenceKVBundle* source_kv =
            source != nullptr ? (source->kv ? &*source->kv : nullptr)
                              : (shared_source->kv ? &*shared_source->kv : nullptr);
        if (source_kv == nullptr) {
            throw std::logic_error("retained source has no KV address space");
        }
        const std::uint32_t main_full_pages =
            plan->reuse_base / static_cast<std::uint32_t>(kPagedKVPageSize);
        const std::uint32_t backend_frontier =
            backend_frontier_at(speculative_backend, plan->reuse_base);
        const std::uint32_t backend_full_pages =
            backend_frontier / static_cast<std::uint32_t>(kPagedKVPageSize);
        if (main_full_pages > active.main_kv_pages ||
            backend_full_pages > active.backend_kv_pages) {
            throw std::logic_error("retained prefix exceeds its active KV entitlement");
        }
        detail::PhysicalResources active_resources{
            .device =
                {
                    .active_lanes     = active.active_lanes,
                    .state_slots      = active.state_slots,
                    .main_kv_pages    = active.main_kv_pages - main_full_pages,
                    .backend_kv_pages = active.backend_kv_pages - backend_full_pages,
                },
        };
        detail::PhysicalResources source_replica_additions;
        const std::uint32_t main_required = pages_for_tokens(plan->reuse_base);
        const std::uint32_t main_device =
            device_kv_prefix_pages(*text_kv_addresses, source_kv->text, plan->reuse_base);
        if (main_device > main_required) {
            throw std::logic_error("retained Text KV replica count exceeds requirement");
        }
        source_replica_additions.device.main_kv_pages = main_required - main_device;
        if (source_kv->backend) {
            const std::uint32_t backend_required =
                backend_frontier == 0 ? 0U : pages_for_tokens(backend_frontier);
            const std::uint32_t backend_device = device_kv_prefix_pages(
                *backend_kv_addresses, *source_kv->backend, backend_frontier);
            if (backend_device > backend_required) {
                throw std::logic_error("retained Backend KV replica count exceeds requirement");
            }
            source_replica_additions.device.backend_kv_pages = backend_required - backend_device;
        }
        detail::PhysicalResources retained_tail_added;
        detail::PhysicalResources retained_tail_removed;
        detail::PhysicalResources retained_tail_preparation_peak;
        const auto plan_retained_tail_release =
            [&](KVAddressSpaceStore& addresses, LogicalKVPageStore& pages,
                KVAddressSpaceHandle address, std::uint32_t frontier, std::uint32_t active_pages,
                std::uint32_t missing_source_pages, bool prefix_fork, bool& staged,
                runtime::ContextResourceClass resource) {
                const std::uint32_t required = pages_for_tokens(frontier);
                const std::uint64_t final_without_release =
                    static_cast<std::uint64_t>(required) + active_pages;
                const std::uint32_t capacity = pages.physical_pool().capacity_pages();
                if (final_without_release <= capacity) { return true; }
                if (!prefix_fork || frontier == 0 ||
                    frontier % static_cast<std::uint32_t>(kPagedKVPageSize) == 0 ||
                    final_without_release != static_cast<std::uint64_t>(capacity) + 1U ||
                    required > addresses.mapped_pages(address)) {
                    return false;
                }
                const LogicalKVPageHandle tail = addresses.logical_page(address, required - 1U);
                if (pages.address_references(tail) != 1 || pages.writer_references(tail) != 0 ||
                    addresses.has_active_reference(tail)) {
                    return false;
                }
                staged                 = true;
                const bool backend     = resource == runtime::ContextResourceClass::BackendKV;
                std::uint32_t& removed = backend ? retained_tail_removed.device.backend_kv_pages
                                                 : retained_tail_removed.device.main_kv_pages;
                std::uint32_t& preparation =
                    backend ? retained_tail_preparation_peak.device.backend_kv_pages
                            : retained_tail_preparation_peak.device.main_kv_pages;
                removed     = 1;
                preparation = missing_source_pages + 1U;
                if (!pages.host_resident(tail)) {
                    const std::size_t stride =
                        plan_host_kv_page_layout(pages.physical_pool().geometry()).page_stride;
                    if (stride > std::numeric_limits<std::size_t>::max() -
                                     retained_tail_added.host.kv_bytes) {
                        throw std::overflow_error("retained KV tail Host backup size overflow");
                    }
                    retained_tail_added.host.kv_bytes += stride;
                    add_kv_transfer(resource, runtime::ContextTransferDirection::DeviceToHost,
                                    pages, 1);
                    plan->needs_transfer = true;
                }
                return true;
            };
        if (!plan_retained_tail_release(
                *text_kv_addresses, *text_kv_pages, source_kv->text, plan->reuse_base,
                active_resources.device.main_kv_pages,
                source_replica_additions.device.main_kv_pages, plan->text_prefix_fork_required,
                plan->text_retained_tail_release, runtime::ContextResourceClass::MainKV)) {
            return std::nullopt;
        }
        if (source_kv->backend) {
            if (!plan_retained_tail_release(
                    *backend_kv_addresses, *backend_kv_pages, *source_kv->backend, backend_frontier,
                    active_resources.device.backend_kv_pages,
                    source_replica_additions.device.backend_kv_pages,
                    plan->backend_prefix_fork_required, plan->backend_retained_tail_release,
                    runtime::ContextResourceClass::BackendKV)) {
                return std::nullopt;
            }
        }
        const bool splits_private_state = [&] {
            if (source == nullptr) { return false; }
            const StateImageHandle selected =
                selected_state(*source, plan->reuse, plan->selected_checkpoint);
            return state_store->residency(selected) == StateReplicaResidency::Both;
        }();
        detail::PhysicalResources added = active_resources;
        if (source_replica_additions.device.main_kv_pages >
                std::numeric_limits<std::uint32_t>::max() - added.device.main_kv_pages ||
            source_replica_additions.device.backend_kv_pages >
                std::numeric_limits<std::uint32_t>::max() - added.device.backend_kv_pages) {
            throw std::overflow_error("retained materialization demand overflow");
        }
        added.device.main_kv_pages += source_replica_additions.device.main_kv_pages;
        added.device.backend_kv_pages += source_replica_additions.device.backend_kv_pages;
        added.host.kv_bytes = retained_tail_added.host.kv_bytes;
        detail::PhysicalResources credit;
        detail::PhysicalResources removed;
        detail::PhysicalResources physical_peak = added;
        if (splits_private_state) {
            credit.device.state_slots  = 1;
            removed.device.state_slots = 1;
            if (physical_peak.device.state_slots == 0) {
                throw std::logic_error("StateImage identity split has no active destination");
            }
            --physical_peak.device.state_slots;
        }
        credit.device.main_kv_pages     = retained_tail_removed.device.main_kv_pages;
        credit.device.backend_kv_pages  = retained_tail_removed.device.backend_kv_pages;
        removed.device.main_kv_pages    = retained_tail_removed.device.main_kv_pages;
        removed.device.backend_kv_pages = retained_tail_removed.device.backend_kv_pages;
        const auto stage_peak           = [](std::uint32_t added_pages, std::uint32_t removed_pages,
                                   std::uint32_t preparation_pages) {
            const std::uint32_t settled =
                added_pages > removed_pages ? added_pages - removed_pages : 0U;
            return std::max(settled, preparation_pages);
        };
        if (plan->text_retained_tail_release) {
            physical_peak.device.main_kv_pages =
                stage_peak(added.device.main_kv_pages, retained_tail_removed.device.main_kv_pages,
                           retained_tail_preparation_peak.device.main_kv_pages);
        }
        if (plan->backend_retained_tail_release) {
            physical_peak.device.backend_kv_pages = stage_peak(
                added.device.backend_kv_pages, retained_tail_removed.device.backend_kv_pages,
                retained_tail_preparation_peak.device.backend_kv_pages);
        }
        plan->demand = detail::PhysicalDemand{
            .active_entitlement       = active_resources,
            .reservation_added        = added,
            .reservation_credit       = credit,
            .physical_peak_additional = physical_peak,
            .final_removed            = removed,
            .final_added              = added,
        };
        return AdmissionCandidate(std::move(plan));
    }
    detail::PhysicalDeviceResources exclusive_active = active;
    detail::PhysicalResources shared_replica_additions;
    detail::PhysicalResources transient_source_restores;
    detail::PhysicalDeviceResources conversions;
    detail::PhysicalHostResources retained_host = plan->active_optional_resources.host;
    std::size_t transient_host_bytes            = 0;
    const auto private_tail_restore =
        [&](const KVAddressSpaceStore& addresses, const LogicalKVPageStore& pages,
            KVAddressSpaceHandle address, std::uint32_t frontier, bool prefix_fork) {
            std::pair<std::uint32_t, std::size_t> out;
            if (!prefix_fork || frontier == 0 ||
                frontier % static_cast<std::uint32_t>(kPagedKVPageSize) == 0) {
                return out;
            }
            const std::uint32_t required = pages_for_tokens(frontier);
            if (required > addresses.mapped_pages(address)) {
                throw std::logic_error("checkpoint KV requirement exceeds address membership");
            }
            const LogicalKVPageHandle tail = addresses.logical_page(address, required - 1U);
            if (pages.address_references(tail) != 1 || pages.device_resident(tail)) { return out; }
            if (!pages.host_resident(tail)) {
                throw std::logic_error("private COW tail has no restorable replica");
            }
            out.first  = 1;
            out.second = plan_host_kv_page_layout(pages.physical_pool().geometry()).page_stride;
            return out;
        };
    if (source != nullptr) {
        const StateImageHandle selected_state_handle =
            selected_state(*source, plan->reuse, plan->selected_checkpoint);
        const StateReplicaResidency selected_residency =
            state_store->residency(selected_state_handle);
        conversions.state_slots = plan->active_optional_resources.device.state_slots;
        if (!plan->state_fork_required &&
            (selected_residency == StateReplicaResidency::DeviceOnly ||
             selected_residency == StateReplicaResidency::Both)) {
            if (conversions.state_slots == std::numeric_limits<std::uint32_t>::max()) {
                throw std::overflow_error("active StateImage conversion overflow");
            }
            ++conversions.state_slots;
        }

        if (!source->kv) { throw std::logic_error("private checkpoint has no KV address space"); }
        const std::uint32_t main_shared =
            shared_kv_prefix_pages(*text_kv_addresses, source->kv->text, plan->reuse_base);
        const std::uint32_t main_missing_shared = missing_shared_device_kv_prefix_pages(
            *text_kv_addresses, source->kv->text, plan->reuse_base);
        if (main_shared > exclusive_active.main_kv_pages) {
            throw std::logic_error("shared Main KV exceeds the active entitlement");
        }
        exclusive_active.main_kv_pages -= main_shared;
        shared_replica_additions.device.main_kv_pages = main_missing_shared;
        const std::uint32_t main_shared_device =
            shared_device_kv_prefix_pages(*text_kv_addresses, source->kv->text, plan->reuse_base);
        conversions.main_kv_pages =
            device_kv_prefix_pages(*text_kv_addresses, source->kv->text, plan->reuse_base) -
            main_shared_device;
        if (source->kv->backend) {
            const std::uint32_t backend_frontier =
                backend_frontier_at(speculative_backend, plan->reuse_base);
            const std::uint32_t backend_shared = shared_kv_prefix_pages(
                *backend_kv_addresses, *source->kv->backend, backend_frontier);
            const std::uint32_t backend_missing_shared = missing_shared_device_kv_prefix_pages(
                *backend_kv_addresses, *source->kv->backend, backend_frontier);
            if (backend_shared > exclusive_active.backend_kv_pages) {
                throw std::logic_error("shared Backend KV exceeds the active entitlement");
            }
            exclusive_active.backend_kv_pages -= backend_shared;
            shared_replica_additions.device.backend_kv_pages = backend_missing_shared;
            const std::uint32_t backend_shared_device        = shared_device_kv_prefix_pages(
                *backend_kv_addresses, *source->kv->backend, backend_frontier);
            conversions.backend_kv_pages =
                device_kv_prefix_pages(*backend_kv_addresses, *source->kv->backend,
                                       backend_frontier) -
                backend_shared_device;
        }
        retained_host.kv_bytes =
            host_kv_prefix_bytes(*text_kv_addresses, source->kv->text, plan->reuse_base);
        const auto [main_tail_restore, main_discarded_host_bytes] =
            private_tail_restore(*text_kv_addresses, *text_kv_pages, source->kv->text,
                                 plan->reuse_base, plan->text_prefix_fork_required);
        transient_source_restores.device.main_kv_pages = main_tail_restore;
        transient_host_bytes                           = main_discarded_host_bytes;
        if (source->kv->backend) {
            const std::uint32_t backend_frontier =
                backend_frontier_at(speculative_backend, plan->reuse_base);
            const std::size_t backend_bytes =
                host_kv_prefix_bytes(*backend_kv_addresses, *source->kv->backend, backend_frontier);
            if (backend_bytes > std::numeric_limits<std::size_t>::max() - retained_host.kv_bytes) {
                throw std::overflow_error("active Host KV entitlement overflow");
            }
            retained_host.kv_bytes += backend_bytes;
            const auto [backend_tail_restore, backend_discarded_host_bytes] =
                private_tail_restore(*backend_kv_addresses, *backend_kv_pages, *source->kv->backend,
                                     backend_frontier, plan->backend_prefix_fork_required);
            transient_source_restores.device.backend_kv_pages = backend_tail_restore;
            if (backend_discarded_host_bytes >
                std::numeric_limits<std::size_t>::max() - transient_host_bytes) {
                throw std::overflow_error("private COW Host KV bytes overflow");
            }
            transient_host_bytes += backend_discarded_host_bytes;
        }
    }
    detail::PhysicalHostResources preparation_retained_host = retained_host;
    if (transient_host_bytes >
        std::numeric_limits<std::size_t>::max() - preparation_retained_host.kv_bytes) {
        throw std::overflow_error("private COW preparation Host KV entitlement overflow");
    }
    preparation_retained_host.kv_bytes += transient_host_bytes;
    conversions = convertible_source_resources(exclusive_active, conversions);
    const detail::PhysicalDeviceResources exclusive_additional =
        additional_resources(exclusive_active, conversions);
    detail::PhysicalResources reservation_added{.device = exclusive_additional};
    if (shared_replica_additions.device.main_kv_pages >
            std::numeric_limits<std::uint32_t>::max() - reservation_added.device.main_kv_pages ||
        shared_replica_additions.device.backend_kv_pages >
            std::numeric_limits<std::uint32_t>::max() - reservation_added.device.backend_kv_pages) {
        throw std::overflow_error("shared KV restore demand overflow");
    }
    reservation_added.device.main_kv_pages += shared_replica_additions.device.main_kv_pages;
    reservation_added.device.backend_kv_pages += shared_replica_additions.device.backend_kv_pages;
    if (transient_source_restores.device.main_kv_pages >
            std::numeric_limits<std::uint32_t>::max() - reservation_added.device.main_kv_pages ||
        transient_source_restores.device.backend_kv_pages >
            std::numeric_limits<std::uint32_t>::max() - reservation_added.device.backend_kv_pages) {
        throw std::overflow_error("private COW source restore demand overflow");
    }
    reservation_added.device.main_kv_pages += transient_source_restores.device.main_kv_pages;
    reservation_added.device.backend_kv_pages += transient_source_restores.device.backend_kv_pages;
    const detail::PhysicalResources active_resources{.device = exclusive_active,
                                                     .host   = retained_host};
    detail::PhysicalResources final_added = active_resources;
    final_added.device.main_kv_pages += shared_replica_additions.device.main_kv_pages;
    final_added.device.backend_kv_pages += shared_replica_additions.device.backend_kv_pages;
    const detail::PhysicalResources reused_source{.device = conversions,
                                                  .host   = preparation_retained_host};
    const detail::PhysicalResources early_source_release =
        positive_difference(source_resources, reused_source);
    plan->demand = detail::PhysicalDemand{
        .active_entitlement       = active_resources,
        .reservation_added        = reservation_added,
        .reservation_credit       = {.device = conversions},
        .physical_peak_additional = positive_difference(reservation_added, early_source_release),
        .final_removed            = source_resources,
        .final_added              = final_added,
    };
    return AdmissionCandidate(std::move(plan));
}

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

#include "runtime/engine/context_cost.h"

#include <algorithm>
#include <chrono>
#include <cstdint>
#include <limits>
#include <optional>
#include <span>
#include <tuple>
#include <utility>
#include <vector>

namespace ninfer::runtime {

struct MaterializationCheckpointPolicy {
    std::uint32_t owner_ordinal = 0;
    CheckpointRef checkpoint;
    RetentionClass retention_class   = RetentionClass::RecentPrivate;
    std::uint64_t selected_hit_count = 0;
    std::uint64_t last_hit_epoch     = 0;
};

struct MaterializationOwnerPolicy {
    std::uint32_t ordinal            = 0;
    RetentionClass retention_class   = RetentionClass::RecentPrivate;
    std::uint64_t selected_hit_count = 0;
    std::uint64_t last_hit_epoch     = 0;
};

template <class Package>
class MaterializationPlanner {
public:
    using Program              = typename Package::Program;
    using PreparedPrompt       = typename Package::PreparedPrompt;
    using AdmissionCandidate   = typename Package::AdmissionCandidate;
    using ResourcePlan         = typename Package::ResourcePlan;
    using ContinuationHandle   = typename Package::ContinuationHandle;
    using SharedPrefixHandle   = typename Package::SharedPrefixHandle;
    using PressureTargetHandle = typename Package::PressureTargetHandle;
    using Clock                = std::chrono::steady_clock;

    struct CandidateInput {
        AdmissionCandidate* candidate = nullptr;
        std::uint32_t stable_ordinal  = 0;
        bool current_session_binding  = false;
    };

    struct LogicalGoal {
        std::uint32_t publication_slot = std::numeric_limits<std::uint32_t>::max();
    };

    struct PressureInputs {
        std::span<const ContinuationHandle* const> private_owners;
        std::span<const std::uint32_t> private_owner_ordinals;
        std::span<const SharedPrefixHandle* const> shared_owners;
        std::span<const std::uint32_t> shared_owner_ordinals;
        std::span<const MaterializationOwnerPolicy> owner_policy;
        std::span<const MaterializationCheckpointPolicy> checkpoint_policy;
    };

    struct Result {
        std::optional<ResourcePlan> plan;
        std::uint32_t candidate_index       = 0;
        std::uint32_t publication_slot      = std::numeric_limits<std::uint32_t>::max();
        ClaimDisposition source_disposition = ClaimDisposition::ConsumedToActive;
        std::vector<PressureOwnerOutcome> owner_outcomes;
        MaterializationDiagnostics diagnostics;
    };

    MaterializationPlanner() {
        queue_.reserve(kTargetBudget);
        identity_costs_.reserve(16);
        assessed_.reserve(kTargetBudget);
        expanded_.reserve(kTargetBudget);
        impact_scratch_.reserve(32);
    }

    template <class PressureInputsFn, class LogicalGoalFn>
    [[nodiscard]] std::optional<Result>
    plan(Program& program, const PreparedPrompt& prompt,
         const ContextMachineCostModel& machine_cost, std::span<const CandidateInput> candidates,
         std::uint32_t root_candidate_index, PressureInputsFn&& pressure_inputs,
         LogicalGoalFn&& logical_goal, Clock::time_point planning_started) {
        if (candidates.empty() || root_candidate_index >= candidates.size()) {
            throw std::invalid_argument("materialization planning problem has no root candidate");
        }
        queue_.clear();
        identity_costs_.clear();
        assessed_.clear();
        expanded_.clear();

        std::optional<Incumbent> identity_best;
        std::vector<IdentityRoot> roots;
        roots.reserve(candidates.size());
        std::uint64_t projection_work = 0;
        for (std::size_t index = 0; index < candidates.size(); ++index) {
            const CandidateInput& input = candidates[index];
            if (input.candidate == nullptr) {
                throw std::invalid_argument("materialization candidate is null");
            }
            const IdentityMaterializationAssessment& identity =
                input.candidate->identity_assessment();
            planning_saturating_add(projection_work, identity.projection_work);
            const FoldedCost cost = fold_identity(input, identity);
            identity_costs_.push_back(cost);
            std::optional<LogicalGoal> goal;
            if (identity.physical_status == MaterializationPhysicalStatus::Feasible) {
                goal = logical_goal(static_cast<std::uint32_t>(index), identity.source_disposition,
                                    std::span<const PressureOwnerOutcome>{});
            }
            if (goal && (!identity_best || cost.less(identity_best->cost))) {
                identity_best = Incumbent{
                    .candidate_index    = static_cast<std::uint32_t>(index),
                    .publication_slot   = goal->publication_slot,
                    .source_disposition = identity.source_disposition,
                    .cost               = cost,
                    .assessment_digest  = identity.assessment_digest,
                };
            }
            const bool needs_pressure =
                !goal.has_value() &&
                (identity.physical_status == MaterializationPhysicalStatus::Feasible ||
                 identity.expandable);
            const bool pressure_can_improve =
                goal.has_value() &&
                identity.physical_status == MaterializationPhysicalStatus::Feasible &&
                cost.lower_bound_ns < cost.total_ns;
            roots.push_back(IdentityRoot{
                .candidate_index = static_cast<std::uint32_t>(index),
                .lower_bound_ns  = cost.lower_bound_ns,
                .expandable      = needs_pressure || pressure_can_improve,
            });
        }

        if (identity_best) {
            const bool dominates = std::all_of(roots.begin(), roots.end(), [&](const auto& root) {
                return !root.expandable || root.lower_bound_ns > identity_best->cost.total_ns;
            });
            if (dominates) {
                std::optional<ResourcePlan> sealed = program.seal_identity(
                    *candidates[identity_best->candidate_index].candidate, prompt);
                if (!sealed) { return std::nullopt; }
                MaterializationDiagnostics diagnostics = complete_diagnostics(
                    identity_best->cost, static_cast<std::uint32_t>(candidates.size()),
                    projection_work, planning_started, MaterializationStopReason::NoPressure,
                    false);
                Result result;
                result.plan               = std::move(*sealed);
                result.candidate_index    = identity_best->candidate_index;
                result.publication_slot   = identity_best->publication_slot;
                result.source_disposition = identity_best->source_disposition;
                result.diagnostics        = diagnostics;
                return result;
            }
        }

        std::vector<const AdmissionCandidate*> candidate_handles;
        candidate_handles.reserve(candidates.size());
        for (const CandidateInput& input : candidates) {
            candidate_handles.push_back(input.candidate);
        }
        const PressureInputs pressure = pressure_inputs();
        if (pressure.private_owners.size() != pressure.private_owner_ordinals.size() ||
            pressure.shared_owners.size() != pressure.shared_owner_ordinals.size()) {
            throw std::logic_error("materialization pressure owner arrays are not aligned");
        }
        auto session = program.begin_pressure_planning(
            machine_cost, candidate_handles, pressure.private_owners,
            pressure.private_owner_ordinals, pressure.shared_owners,
            pressure.shared_owner_ordinals);

        Incumbent incumbent;
        std::uint32_t targets_evaluated = static_cast<std::uint32_t>(candidates.size());
        if (identity_best) {
            incumbent = *identity_best;
            incumbent.target =
                session.identity_target(*candidates[incumbent.candidate_index].candidate);
        } else {
            PressureTargetHandle root_maximal =
                session.root_maximal_target(*candidates[root_candidate_index].candidate);
            PressureTargetAssessment assessment = session.assess(root_maximal);
            ++targets_evaluated;
            planning_saturating_add(projection_work, assessment.projection_work);
            std::optional<LogicalGoal> goal;
            if (assessment.physical_status == MaterializationPhysicalStatus::Feasible) {
                goal = logical_goal(root_candidate_index, assessment.source_disposition,
                                    assessment.owner_outcomes);
            }
            if (!goal) { return std::nullopt; }
            incumbent = make_incumbent(root_maximal, assessment, candidates[root_candidate_index],
                                       pressure.owner_policy, pressure.checkpoint_policy, *goal);
        }

        for (const IdentityRoot& root : roots) {
            if (!root.expandable || root.lower_bound_ns > incumbent.cost.total_ns) { continue; }
            QueueEntry entry;
            entry.target = session.identity_target(*candidates[root.candidate_index].candidate);
            entry.candidate_index   = root.candidate_index;
            entry.lower_bound_ns    = root.lower_bound_ns;
            entry.remaining_prefill = identity_costs_[root.candidate_index].remaining_text_prefill;
            entry.remaining_vision_prefill =
                identity_costs_[root.candidate_index].remaining_vision_prefill;
            entry.reused_prompt_tokens = identity_costs_[root.candidate_index].reused_prompt_tokens;
            entry.current_session_binding =
                identity_costs_[root.candidate_index].current_session_binding;
            entry.candidate_ordinal     = identity_costs_[root.candidate_index].candidate_ordinal;
            entry.stable_target_ordinal = root.candidate_index;
            assessed_.push_back(entry.target);
            queue_push(std::move(entry));
        }

        const Clock::time_point search_started = Clock::now();
        const std::uint64_t search_budget_ns =
            std::min<std::uint64_t>(5'000'000ULL, incumbent.cost.total_ns / 20U);
        std::uint64_t maximum_expansion_ns    = 0;
        std::uint32_t optional_targets        = 0;
        MaterializationStopReason stop_reason = MaterializationStopReason::QueueExhausted;
        bool model_optimal                    = true;
        bool budget_exhausted                 = false;
        std::uint64_t interrupted_lower_bound = std::numeric_limits<std::uint64_t>::max();

        for (;;) {
            if (queue_.empty()) {
                stop_reason   = MaterializationStopReason::QueueExhausted;
                model_optimal = true;
                break;
            }
            const QueueEntry& next = queue_.front();
            if (next.lower_bound_ns > incumbent.cost.total_ns) {
                stop_reason   = MaterializationStopReason::ModelOptimal;
                model_optimal = true;
                break;
            }
            if (optional_targets >= kTargetBudget) {
                stop_reason      = MaterializationStopReason::TargetBudget;
                model_optimal    = false;
                budget_exhausted = true;
                break;
            }
            const std::uint64_t elapsed = elapsed_ns(search_started, Clock::now());
            if (elapsed >= search_budget_ns) {
                stop_reason      = MaterializationStopReason::TimeBudget;
                model_optimal    = false;
                budget_exhausted = true;
                break;
            }
            const std::uint64_t possible_improvement =
                incumbent.cost.total_ns > next.lower_bound_ns
                    ? incumbent.cost.total_ns - next.lower_bound_ns
                    : 0;
            if (maximum_expansion_ns != 0 && maximum_expansion_ns >= possible_improvement) {
                stop_reason   = MaterializationStopReason::ValueOfNextExpansion;
                model_optimal = false;
                break;
            }

            const QueueEntry parent = queue_pop();
            if (contains(expanded_, parent.target)) { continue; }
            const Clock::time_point expansion_started = Clock::now();
            auto prepared                             = session.prepare_expansion(parent.target);
            if (prepared.new_canonical_count() > kTargetBudget - optional_targets) {
                session.discard_expansion(std::move(prepared));
                interrupted_lower_bound = parent.lower_bound_ns;
                stop_reason             = MaterializationStopReason::ExpansionCapacity;
                model_optimal           = false;
                budget_exhausted        = true;
                break;
            }
            const auto children = session.commit_expansion(std::move(prepared));
            optional_targets += children.new_canonical_count;
            expanded_.push_back(parent.target);

            for (const PressureTargetHandle child : children.children) {
                if (contains(assessed_, child)) { continue; }
                PressureTargetAssessment assessment = session.assess(child);
                assessed_.push_back(child);
                ++targets_evaluated;
                planning_saturating_add(projection_work, assessment.projection_work);
                const FoldedCost cost =
                    fold_assessment(candidates[assessment.candidate_ordinal], assessment,
                                    pressure.owner_policy, pressure.checkpoint_policy);
                std::optional<LogicalGoal> goal;
                if (assessment.physical_status == MaterializationPhysicalStatus::Feasible) {
                    goal = logical_goal(assessment.candidate_ordinal, assessment.source_disposition,
                                        assessment.owner_outcomes);
                }
                if (goal && cost.less(incumbent.cost)) {
                    incumbent =
                        make_incumbent(child, assessment, candidates[assessment.candidate_ordinal],
                                       pressure.owner_policy, pressure.checkpoint_policy, *goal);
                }
                if (assessment.expandable && cost.lower_bound_ns <= incumbent.cost.total_ns) {
                    queue_push(QueueEntry{
                        .target                    = child,
                        .candidate_index           = assessment.candidate_ordinal,
                        .lower_bound_ns            = cost.lower_bound_ns,
                        .shared_stable_units       = cost.shared_stable_units,
                        .live_session_units        = cost.live_session_units,
                        .recent_private_units      = cost.recent_private_units,
                        .disposable_units          = cost.disposable_units,
                        .affected_selected_hits    = cost.affected_selected_hits,
                        .newest_affected_hit_epoch = cost.newest_affected_hit_epoch,
                        .owner_evictions           = cost.owner_evictions,
                        .checkpoint_drops          = cost.checkpoint_drops,
                        .copy_operations           = cost.copy_operations,
                        .transferred_bytes         = cost.transferred_bytes,
                        .remaining_prefill         = cost.remaining_text_prefill,
                        .remaining_vision_prefill  = cost.remaining_vision_prefill,
                        .reused_prompt_tokens      = cost.reused_prompt_tokens,
                        .current_session_binding   = cost.current_session_binding,
                        .candidate_ordinal         = cost.candidate_ordinal,
                        .stable_target_ordinal     = assessment.stable_target_ordinal,
                    });
                }
            }
            maximum_expansion_ns =
                std::max(maximum_expansion_ns, elapsed_ns(expansion_started, Clock::now()));
        }

        const std::uint64_t search_elapsed_ns        = elapsed_ns(search_started, Clock::now());
        PressureTargetAssessment selected_assessment = session.assess(incumbent.target);
        planning_saturating_add(projection_work, selected_assessment.projection_work);
        const FoldedCost selected_cost =
            fold_assessment(candidates[incumbent.candidate_index], selected_assessment,
                            pressure.owner_policy, pressure.checkpoint_policy);
        if (!equivalent(selected_cost, incumbent.cost) ||
            selected_assessment.assessment_digest != incumbent.assessment_digest) {
            throw std::logic_error("selected pressure target changed before seal");
        }
        std::vector<PressureOwnerOutcome> selected_outcomes(
            selected_assessment.owner_outcomes.begin(), selected_assessment.owner_outcomes.end());
        std::optional<ResourcePlan> sealed = session.seal(incumbent.target, prompt);
        if (!sealed) { throw std::logic_error("selected pressure target could not be sealed"); }

        std::uint64_t best_remaining =
            queue_.empty() ? interrupted_lower_bound : queue_.front().lower_bound_ns;
        best_remaining = std::min(best_remaining, interrupted_lower_bound);
        if (best_remaining == std::numeric_limits<std::uint64_t>::max()) {
            best_remaining = selected_cost.total_ns;
        }
        MaterializationDiagnostics diagnostics = make_diagnostics(
            selected_cost, targets_evaluated, projection_work, planning_started, search_elapsed_ns,
            stop_reason, model_optimal, budget_exhausted, best_remaining,
            selected_assessment.degradation_units, selected_assessment.root_maximal);

        Result result;
        result.plan               = std::move(*sealed);
        result.candidate_index    = incumbent.candidate_index;
        result.publication_slot   = incumbent.publication_slot;
        result.source_disposition = selected_assessment.source_disposition;
        result.owner_outcomes     = std::move(selected_outcomes);
        result.diagnostics        = diagnostics;
        return result;
    }

private:
    static constexpr std::uint32_t kTargetBudget = 4096;

    struct FoldedCost {
        std::uint64_t now_ns                    = 0;
        std::uint64_t future_loss_ns            = 0;
        std::uint64_t total_ns                  = 0;
        std::uint64_t lower_bound_ns            = 0;
        std::uint32_t shared_stable_units       = 0;
        std::uint32_t live_session_units        = 0;
        std::uint32_t recent_private_units      = 0;
        std::uint32_t disposable_units          = 0;
        std::uint64_t affected_selected_hits    = 0;
        std::uint64_t newest_affected_hit_epoch = 0;
        std::uint32_t owner_evictions           = 0;
        std::uint32_t checkpoint_drops          = 0;
        std::uint32_t copy_operations           = 0;
        std::uint64_t transferred_bytes         = 0;
        std::uint64_t remaining_text_prefill    = 0;
        std::uint64_t remaining_vision_prefill  = 0;
        std::uint32_t reused_prompt_tokens      = 0;
        bool current_session_binding            = false;
        std::uint32_t candidate_ordinal         = 0;
        std::uint32_t target_ordinal            = 0;

        [[nodiscard]] auto key() const noexcept {
            return std::tuple{
                total_ns,
                shared_stable_units,
                live_session_units,
                recent_private_units,
                disposable_units,
                affected_selected_hits,
                newest_affected_hit_epoch,
                owner_evictions,
                checkpoint_drops,
                copy_operations,
                transferred_bytes,
                remaining_text_prefill,
                remaining_vision_prefill,
                std::numeric_limits<std::uint32_t>::max() - reused_prompt_tokens,
                current_session_binding ? 0U : 1U,
                candidate_ordinal,
                target_ordinal,
            };
        }

        [[nodiscard]] bool less(const FoldedCost& other) const noexcept {
            return key() < other.key();
        }
    };

    struct Incumbent {
        PressureTargetHandle target{};
        std::uint32_t candidate_index       = 0;
        std::uint32_t publication_slot      = std::numeric_limits<std::uint32_t>::max();
        ClaimDisposition source_disposition = ClaimDisposition::ConsumedToActive;
        FoldedCost cost;
        std::uint64_t assessment_digest = 0;
    };

    struct IdentityRoot {
        std::uint32_t candidate_index = 0;
        std::uint64_t lower_bound_ns  = 0;
        bool expandable               = false;
    };

    struct QueueEntry {
        PressureTargetHandle target{};
        std::uint32_t candidate_index           = 0;
        std::uint64_t lower_bound_ns            = 0;
        std::uint32_t shared_stable_units       = 0;
        std::uint32_t live_session_units        = 0;
        std::uint32_t recent_private_units      = 0;
        std::uint32_t disposable_units          = 0;
        std::uint64_t affected_selected_hits    = 0;
        std::uint64_t newest_affected_hit_epoch = 0;
        std::uint32_t owner_evictions           = 0;
        std::uint32_t checkpoint_drops          = 0;
        std::uint32_t copy_operations           = 0;
        std::uint64_t transferred_bytes         = 0;
        std::uint64_t remaining_prefill         = 0;
        std::uint64_t remaining_vision_prefill  = 0;
        std::uint32_t reused_prompt_tokens      = 0;
        bool current_session_binding            = false;
        std::uint32_t candidate_ordinal         = 0;
        std::uint32_t stable_target_ordinal     = 0;
    };

    struct CombinedImpact {
        std::uint32_t owner_ordinal = 0;
        CheckpointRef checkpoint;
        std::uint64_t baseline_ns = 0;
        std::uint64_t target_ns   = 0;
    };

    [[nodiscard]] static std::uint64_t elapsed_ns(Clock::time_point begin,
                                                  Clock::time_point end) noexcept {
        const auto count =
            std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin).count();
        return count > 0 ? static_cast<std::uint64_t>(count) : 0;
    }

    static void planning_saturating_add(std::uint64_t& value, std::uint64_t add) noexcept {
        value = add > std::numeric_limits<std::uint64_t>::max() - value
                    ? std::numeric_limits<std::uint64_t>::max()
                    : value + add;
    }

    [[nodiscard]] static std::uint64_t retention_weight(RetentionClass retention) noexcept {
        switch (retention) {
        case RetentionClass::Disposable:
            return 1;
        case RetentionClass::RecentPrivate:
            return 4;
        case RetentionClass::LiveSession:
            return 16;
        case RetentionClass::SharedStable:
            return 64;
        }
        return 64;
    }

    [[nodiscard]] static const MaterializationOwnerPolicy*
    owner_policy_for(std::span<const MaterializationOwnerPolicy> policies,
                     std::uint32_t ordinal) noexcept {
        const auto found = std::find_if(policies.begin(), policies.end(), [&](const auto& policy) {
            return policy.ordinal == ordinal;
        });
        return found == policies.end() ? nullptr : &*found;
    }

    [[nodiscard]] static const MaterializationCheckpointPolicy*
    checkpoint_policy_for(std::span<const MaterializationCheckpointPolicy> policies,
                          std::uint32_t owner_ordinal, CheckpointRef checkpoint) noexcept {
        const auto found = std::find_if(policies.begin(), policies.end(), [&](const auto& policy) {
            return policy.owner_ordinal == owner_ordinal && policy.checkpoint == checkpoint;
        });
        return found == policies.end() ? nullptr : &*found;
    }

    [[nodiscard]] FoldedCost
    fold_identity(const CandidateInput& candidate,
                  const IdentityMaterializationAssessment& assessment) const noexcept {
        FoldedCost cost;
        cost.now_ns                   = assessment.machine.immediate_ns;
        cost.total_ns                 = cost.now_ns;
        cost.lower_bound_ns           = assessment.machine.minimum_request_ns;
        cost.copy_operations          = assessment.machine.copy_operations;
        cost.transferred_bytes        = assessment.machine.transferred_bytes;
        cost.remaining_text_prefill   = assessment.machine.remaining_prefill_work.tokens;
        cost.remaining_vision_prefill = assessment.machine.remaining_prefill_work.vision_patches;
        cost.reused_prompt_tokens     = assessment.machine.reused_prompt_tokens;
        cost.current_session_binding  = candidate.current_session_binding;
        cost.candidate_ordinal        = candidate.stable_ordinal;
        cost.target_ordinal           = candidate.stable_ordinal;
        return cost;
    }

    [[nodiscard]] FoldedCost
    fold_assessment(const CandidateInput& candidate, const PressureTargetAssessment& assessment,
                    std::span<const MaterializationOwnerPolicy> owner_policies,
                    std::span<const MaterializationCheckpointPolicy> checkpoint_policies) {
        FoldedCost cost;
        cost.now_ns                   = assessment.machine.immediate_ns;
        cost.copy_operations          = assessment.machine.copy_operations;
        cost.transferred_bytes        = assessment.machine.transferred_bytes;
        cost.remaining_text_prefill   = assessment.machine.remaining_prefill_work.tokens;
        cost.remaining_vision_prefill = assessment.machine.remaining_prefill_work.vision_patches;
        cost.reused_prompt_tokens     = assessment.machine.reused_prompt_tokens;
        cost.current_session_binding  = candidate.current_session_binding;
        cost.candidate_ordinal        = candidate.stable_ordinal;
        cost.target_ordinal           = assessment.stable_target_ordinal;
        cost.checkpoint_drops         = assessment.dropped_checkpoints;

        for (const PressureOwnerOutcome& outcome : assessment.owner_outcomes) {
            const MaterializationOwnerPolicy* policy =
                owner_policy_for(owner_policies, outcome.owner_ordinal);
            if (policy == nullptr) {
                throw std::logic_error("pressure target references an unknown logical owner");
            }
            switch (policy->retention_class) {
            case RetentionClass::SharedStable:
                cost.shared_stable_units += outcome.degradation_units;
                break;
            case RetentionClass::LiveSession:
                cost.live_session_units += outcome.degradation_units;
                break;
            case RetentionClass::RecentPrivate:
                cost.recent_private_units += outcome.degradation_units;
                break;
            case RetentionClass::Disposable:
                cost.disposable_units += outcome.degradation_units;
                break;
            }
            if (outcome.disposition == ClaimDisposition::Evicted) { ++cost.owner_evictions; }
        }

        impact_scratch_.clear();
        for (const PressureCheckpointRecoveryImpact& impact : assessment.checkpoint_impacts) {
            const auto found = std::find_if(impact_scratch_.begin(), impact_scratch_.end(),
                                            [&](const CombinedImpact& item) {
                                                return item.owner_ordinal == impact.owner_ordinal &&
                                                       item.checkpoint == impact.checkpoint;
                                            });
            if (found == impact_scratch_.end()) {
                impact_scratch_.push_back(CombinedImpact{
                    .owner_ordinal = impact.owner_ordinal,
                    .checkpoint    = impact.checkpoint,
                    .baseline_ns   = impact.baseline_recovery_ns,
                    .target_ns     = impact.target_recovery_ns,
                });
            } else {
                found->baseline_ns = std::min(found->baseline_ns, impact.baseline_recovery_ns);
                found->target_ns   = std::max(found->target_ns, impact.target_recovery_ns);
            }
        }
        for (const CombinedImpact& impact : impact_scratch_) {
            if (impact.target_ns <= impact.baseline_ns) { continue; }
            const MaterializationCheckpointPolicy* policy =
                checkpoint_policy_for(checkpoint_policies, impact.owner_ordinal, impact.checkpoint);
            const MaterializationOwnerPolicy* owner =
                owner_policy_for(owner_policies, impact.owner_ordinal);
            if (owner == nullptr) {
                throw std::logic_error("checkpoint impact has no owner policy");
            }
            const RetentionClass retention =
                policy != nullptr ? policy->retention_class : owner->retention_class;
            const std::uint64_t delta  = impact.target_ns - impact.baseline_ns;
            const std::uint64_t weight = retention_weight(retention);
            const std::uint64_t weighted =
                delta > std::numeric_limits<std::uint64_t>::max() / weight
                    ? std::numeric_limits<std::uint64_t>::max()
                    : delta * weight;
            planning_saturating_add(cost.future_loss_ns, weighted);
            planning_saturating_add(cost.affected_selected_hits, policy != nullptr
                                                                     ? policy->selected_hit_count
                                                                     : owner->selected_hit_count);
            cost.newest_affected_hit_epoch =
                std::max(cost.newest_affected_hit_epoch,
                         policy != nullptr ? policy->last_hit_epoch : owner->last_hit_epoch);
        }
        cost.total_ns = cost.now_ns;
        planning_saturating_add(cost.total_ns, cost.future_loss_ns);
        cost.lower_bound_ns = assessment.machine.minimum_request_ns;
        planning_saturating_add(cost.lower_bound_ns, cost.future_loss_ns);
        return cost;
    }

    [[nodiscard]] Incumbent make_incumbent(
        PressureTargetHandle target, const PressureTargetAssessment& assessment,
        const CandidateInput& candidate, std::span<const MaterializationOwnerPolicy> owner_policy,
        std::span<const MaterializationCheckpointPolicy> checkpoint_policy, LogicalGoal goal) {
        return Incumbent{
            .target             = target,
            .candidate_index    = assessment.candidate_ordinal,
            .publication_slot   = goal.publication_slot,
            .source_disposition = assessment.source_disposition,
            .cost = fold_assessment(candidate, assessment, owner_policy, checkpoint_policy),
            .assessment_digest = assessment.assessment_digest,
        };
    }

    [[nodiscard]] static bool equivalent(const FoldedCost& left, const FoldedCost& right) noexcept {
        return left.key() == right.key();
    }

    [[nodiscard]] static auto queue_key(const QueueEntry& entry) noexcept {
        return std::tuple{
            entry.lower_bound_ns,
            entry.shared_stable_units,
            entry.live_session_units,
            entry.recent_private_units,
            entry.disposable_units,
            entry.affected_selected_hits,
            entry.newest_affected_hit_epoch,
            entry.owner_evictions,
            entry.checkpoint_drops,
            entry.copy_operations,
            entry.transferred_bytes,
            entry.remaining_prefill,
            entry.remaining_vision_prefill,
            std::numeric_limits<std::uint32_t>::max() - entry.reused_prompt_tokens,
            entry.current_session_binding ? 0U : 1U,
            entry.candidate_ordinal,
            entry.stable_target_ordinal,
        };
    }

    void queue_push(QueueEntry entry) {
        queue_.push_back(std::move(entry));
        std::push_heap(queue_.begin(), queue_.end(), [](const auto& left, const auto& right) {
            return queue_key(right) < queue_key(left);
        });
    }

    [[nodiscard]] QueueEntry queue_pop() {
        std::pop_heap(queue_.begin(), queue_.end(), [](const auto& left, const auto& right) {
            return queue_key(right) < queue_key(left);
        });
        QueueEntry result = std::move(queue_.back());
        queue_.pop_back();
        return result;
    }

    [[nodiscard]] static bool contains(std::span<const PressureTargetHandle> handles,
                                       PressureTargetHandle target) noexcept {
        return std::find(handles.begin(), handles.end(), target) != handles.end();
    }

    [[nodiscard]] static MaterializationDiagnostics
    complete_diagnostics(const FoldedCost& cost, std::uint32_t targets_evaluated,
                         std::uint64_t projection_work, Clock::time_point planning_started,
                         MaterializationStopReason reason, bool maximal_fallback) noexcept {
        return make_diagnostics(cost, targets_evaluated, projection_work, planning_started, 0,
                                reason, true, false, cost.total_ns, 0, maximal_fallback);
    }

    [[nodiscard]] static MaterializationDiagnostics
    make_diagnostics(const FoldedCost& cost, std::uint32_t targets_evaluated,
                     std::uint64_t projection_work, Clock::time_point planning_started,
                     std::uint64_t search_elapsed_ns, MaterializationStopReason reason,
                     bool model_optimal, bool budget_exhausted,
                     std::uint64_t best_remaining_lower_bound_ns, std::uint32_t degradation_units,
                     bool maximal_fallback) noexcept {
        const std::uint64_t gap = model_optimal || best_remaining_lower_bound_ns >= cost.total_ns
                                      ? 0
                                      : cost.total_ns - best_remaining_lower_bound_ns;
        return MaterializationDiagnostics{
            .predicted_now_ns              = cost.now_ns,
            .predicted_future_loss_ns      = cost.future_loss_ns,
            .predicted_total_ns            = cost.total_ns,
            .targets_evaluated             = targets_evaluated,
            .projection_work               = projection_work,
            .planning_elapsed_ns           = elapsed_ns(planning_started, Clock::now()),
            .search_elapsed_ns             = search_elapsed_ns,
            .stop_reason                   = reason,
            .model_optimal                 = model_optimal,
            .budget_exhausted              = budget_exhausted,
            .best_remaining_lower_bound_ns = best_remaining_lower_bound_ns,
            .absolute_bound_gap_ns         = gap,
            .relative_bound_gap =
                cost.total_ns == 0 ? 0.0
                                   : static_cast<double>(gap) / static_cast<double>(cost.total_ns),
            .selected_degradation_units = degradation_units,
            .selected_maximal_fallback  = maximal_fallback,
        };
    }

    std::vector<QueueEntry> queue_;
    std::vector<FoldedCost> identity_costs_;
    std::vector<PressureTargetHandle> assessed_;
    std::vector<PressureTargetHandle> expanded_;
    std::vector<CombinedImpact> impact_scratch_;
};

} // namespace ninfer::runtime
#pragma once

#include "core/tensor.h"

#include <cuda_runtime.h>

namespace ninfer::ops::detail {

void packed_attention_launch(const Tensor& q, const Tensor& k, const Tensor& v,
                             const Tensor& cu_seqlens, Tensor* tiles, Tensor& out,
                             cudaStream_t stream);

std::int32_t packed_attention_uniform_tile(std::int32_t segment_length);

void packed_attention_uniform_launch(const Tensor& q, const Tensor& k, const Tensor& v,
                                     std::int32_t segment_length, Tensor& out, cudaStream_t stream);

void packed_attention_uniform_launch_with_tile(const Tensor& q, const Tensor& k, const Tensor& v,
                                               std::int32_t segment_length, std::int32_t tile_size,
                                               Tensor& out, cudaStream_t stream);

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

#include "core/tensor.h"

#include <cuda_runtime.h>

namespace ninfer::ops::detail {

void prepare_ragged_prefix_launch(const Tensor& source, const Tensor& lanes, const Tensor& starts,
                                  const Tensor& ends, Tensor& destination, Tensor& positions,
                                  Tensor& counts, cudaStream_t stream);

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

#include <cuda_pipeline.h>
#include <cuda_runtime.h>

namespace ninfer::ops {

enum class Cache { ca, cg };

template <class V, class T>
__device__ __forceinline__ V load_vec(const T* ptr) {
    static_assert(sizeof(V) == 1 || sizeof(V) == 2 || sizeof(V) == 4 || sizeof(V) == 8 ||
                  sizeof(V) == 16);
    return *reinterpret_cast<const V*>(ptr);
}

template <class V, class T>
__device__ __forceinline__ V load_ldg(const T* ptr) {
    static_assert(sizeof(V) == 1 || sizeof(V) == 2 || sizeof(V) == 4 || sizeof(V) == 8 ||
                  sizeof(V) == 16);
    return __ldg(reinterpret_cast<const V*>(ptr));
}

template <class T, class V>
__device__ __forceinline__ void store_vec(T* ptr, V value) {
    static_assert(sizeof(V) == 1 || sizeof(V) == 2 || sizeof(V) == 4 || sizeof(V) == 8 ||
                  sizeof(V) == 16);
    *reinterpret_cast<V*>(ptr) = value;
}

__device__ __forceinline__ unsigned smem_addr(const void* ptr) {
    return static_cast<unsigned>(__cvta_generic_to_shared(ptr));
}

template <int Bytes, Cache Policy = Cache::ca>
__device__ __forceinline__ void cp_async(void* smem_dst, const void* gmem_src) {
    static_assert(Bytes == 4 || Bytes == 8 || Bytes == 16, "cp_async supports 4, 8, or 16 bytes");
    if constexpr (Policy == Cache::cg) {
        static_assert(Bytes == 16, "cp.async.cg requires a 16-byte copy");
        asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n"
                     :
                     : "r"(smem_addr(smem_dst)), "l"(gmem_src));
    } else {
        asm volatile("cp.async.ca.shared.global [%0], [%1], %2;\n"
                     :
                     : "r"(smem_addr(smem_dst)), "l"(gmem_src), "n"(Bytes));
    }
}

template <int Bytes, Cache Policy = Cache::ca>
__device__ __forceinline__ void cp_async_zfill(void* smem_dst, const void* gmem_src,
                                               int src_bytes) {
    static_assert(Bytes == 4 || Bytes == 8 || Bytes == 16,
                  "cp_async_zfill supports 4, 8, or 16 bytes");
    if constexpr (Policy == Cache::cg) {
        static_assert(Bytes == 16, "cp.async.cg requires a 16-byte copy");
        asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n"
                     :
                     : "r"(smem_addr(smem_dst)), "l"(gmem_src), "r"(src_bytes));
    } else {
        asm volatile("cp.async.ca.shared.global [%0], [%1], %2, %3;\n"
                     :
                     : "r"(smem_addr(smem_dst)), "l"(gmem_src), "n"(Bytes), "r"(src_bytes));
    }
}

__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;\n"); }

template <int Groups>
__device__ __forceinline__ void cp_wait() {
    static_assert(Groups >= 0 && Groups <= 7, "cp_wait group count must fit the PTX immediate");
    asm volatile("cp.async.wait_group %0;\n" : : "n"(Groups));
}

template <int Bytes>
__device__ __forceinline__ void pipe_copy(void* smem_dst, const void* gmem_src) {
    static_assert(Bytes == 4 || Bytes == 8 || Bytes == 16, "pipe_copy supports 4, 8, or 16 bytes");
    __pipeline_memcpy_async(smem_dst, gmem_src, Bytes);
}

__device__ __forceinline__ void pipe_commit() { __pipeline_commit(); }

template <int Groups>
__device__ __forceinline__ void pipe_wait() {
    static_assert(Groups >= 0 && Groups <= 7, "pipe_wait group count must fit the PTX immediate");
    __pipeline_wait_prior(Groups);
}

} // namespace ninfer::ops
#pragma once

#include "core/tensor.h"

#include <cuda_runtime.h>

namespace ninfer::ops::detail {

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

void launch_q4_gemv_r4_w1_direct(const Tensor& x, const Weight& w, Tensor& out,
                                 cudaStream_t stream);
void launch_q4_gemv_r1_w8_direct(const Tensor& x, const Weight& w, Tensor& out,
                                 cudaStream_t stream);
void launch_q4_simt_r8_c4(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream);
void launch_q4_simt_r8_c8(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream);
void launch_q4_draft_head_small_t(const Tensor& x, const Weight& w, Tensor& out,
                                  cudaStream_t stream);
void launch_q4_mma_r64_c32(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream);
void launch_q4_mma_r64_c48(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream);
void launch_q4_mma_r64_c56(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream);
void launch_q4_mma_r64_c64(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream);
void launch_q4_mma_r64_c64_endpoint(const Tensor& x, const Weight& w, Tensor& out,
                                    cudaStream_t stream);
void launch_q4_mma_r64_c72(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream);
void launch_q4_mma_r64_c80(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream);
void launch_q4_mma_r64_c96(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream);
void launch_q4_mma_r64_c104_bounded(const Tensor& x, const Weight& w, Tensor& out,
                                    cudaStream_t stream);
void launch_q4_mma_r64_c112_partial(const Tensor& x, const Weight& w, Tensor& out,
                                    cudaStream_t stream);
void launch_q4_mma_r64_c112(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream);
void launch_q4_mma_r64_c120_partial(const Tensor& x, const Weight& w, Tensor& out,
                                    cudaStream_t stream);
void launch_q4_mma_r64_c120(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream);
void launch_q4_mma_r64_c128(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream);

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

// Reusable row-scaled FP8 CUDA-core mainloop for a compile-time number of BF16 activation
// columns. A CTA owns one row tile and one compile-time token tile, reusing each decoded weight
// pair across that token tile before advancing K. Output, epilogue, and row policies let fused
// consumers retain their observable semantics without duplicating the contraction.

#include "ops/linear/fp8/fp8_gemv.cuh"

#include <cuda_bf16.h>

#include <cstdint>

namespace ninfer::ops::detail {

enum class Fp8SmallTFinalization : std::uint8_t {
    Elementwise,
    RowVector,
};

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

template <int Values>
__device__ __forceinline__ Fp8ActivationPack<Values>
load_fp8_activation_pack(const __nv_bfloat16* pointer) {
    Fp8ActivationPack<Values> result;
#pragma unroll
    for (int chunk = 0; chunk < Values / 8; ++chunk) {
        const uint4 packed          = load_vec<uint4>(pointer + chunk * 8);
        result.words[chunk * 4]     = packed.x;
        result.words[chunk * 4 + 1] = packed.y;
        result.words[chunk * 4 + 2] = packed.z;
        result.words[chunk * 4 + 3] = packed.w;
    }
    return result;
}

template <class Schedule>
struct Fp8SmallTSharedStorage {
    static constexpr int kValuesPerPhase = 32 * Schedule::kValuesPerLane;
    static constexpr int kActivationElements =
        Schedule::kActivationAccess == Fp8SmallTActivationAccess::SharedPhase
            ? Schedule::kTokenTile * kValuesPerPhase
            : 8;
    alignas(16) __nv_bfloat16 activation[kActivationElements];
};

template <class Geometry, int ActiveTokens, class Schedule, class Output,
          class Epilogue = Fp8IdentityEpilogue, class RowPolicy = Fp8GemvIdentityRows,
          bool PairRows                      = false,
          Fp8SmallTFinalization Finalization = Fp8SmallTFinalization::Elementwise>
__global__ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void fp8_small_t_kernel(
    const __nv_bfloat16* __restrict__ x, const std::uint8_t* __restrict__ weight_codes,
    const __nv_bfloat16* __restrict__ row_scales, Output output, Epilogue epilogue = {},
    RowPolicy row_policy = {}) {
    static_assert(ActiveTokens >= 2);
    static_assert(Schedule::kTokenTile <= ActiveTokens);
    static_assert(!PairRows || (Schedule::kRowsPerWarp % 2) == 0);
    constexpr int kValuesPerPhase = 32 * Schedule::kValuesPerLane;
    static_assert((Geometry::kInputRows % kValuesPerPhase) == 0);
    static_assert((Geometry::kOutputRows % Schedule::kRowsPerCta) == 0);
    constexpr int kPhases = Geometry::kInputRows / kValuesPerPhase;
    constexpr int kStoredRowsPerWarp =
        PairRows ? Schedule::kRowsPerWarp / 2 : Schedule::kRowsPerWarp;
    constexpr int kStoredRowsPerCta = Schedule::kWarpsPerCta * kStoredRowsPerWarp;
    constexpr int kRowBlocks        = Geometry::kOutputRows / Schedule::kRowsPerCta;
    constexpr int kTokenTiles = (ActiveTokens + Schedule::kTokenTile - 1) / Schedule::kTokenTile;

    const int linear_block = static_cast<int>(blockIdx.x);
    int row_block;
    int token_tile;
    if constexpr (Schedule::kBlockOrder == Fp8SmallTBlockOrder::RowsContiguous) {
        token_tile = linear_block / kRowBlocks;
        row_block  = linear_block - token_tile * kRowBlocks;
    } else {
        row_block  = linear_block / kTokenTiles;
        token_tile = linear_block - row_block * kTokenTiles;
    }
    const int token0 = kTokenTiles == 1 ? 0 : token_tile * Schedule::kTokenTile;

    __shared__ Fp8SmallTSharedStorage<Schedule> shared;
    const int lane      = static_cast<int>(threadIdx.x) & 31;
    const int warp      = static_cast<int>(threadIdx.x) >> 5;
    const int row_begin = row_block * kStoredRowsPerCta + warp * kStoredRowsPerWarp;
    float accumulators[Schedule::kRowsPerWarp][Schedule::kTokenTile][Schedule::kAccumulatorChains] =
        {};

#pragma unroll Schedule::kPhaseUnroll
    for (int phase = 0; phase < kPhases; ++phase) {
        if constexpr (Schedule::kActivationAccess == Fp8SmallTActivationAccess::SharedPhase) {
            constexpr int kPacksPerToken = kValuesPerPhase / 8;
            constexpr int kPacks         = Schedule::kTokenTile * kPacksPerToken;
            auto* destination            = reinterpret_cast<uint4*>(shared.activation);
            for (int task = static_cast<int>(threadIdx.x); task < kPacks;
                 task += Schedule::kThreads) {
                const int local_token = task / kPacksPerToken;
                const int local_pack  = task - local_token * kPacksPerToken;
                const int token       = token0 + local_token;
                if (token < ActiveTokens) {
                    destination[task] = load_vec<uint4>(
                        x + static_cast<std::int64_t>(token) * Geometry::kInputRows +
                        phase * kValuesPerPhase + local_pack * 8);
                }
            }
        }

        const int value_begin = phase * kValuesPerPhase + lane * Schedule::kValuesPerLane;
        Fp8CodePack<Schedule::kValuesPerLane> row_codes[Schedule::kRowsPerWarp];
#pragma unroll
        for (int local_row = 0; local_row < Schedule::kRowsPerWarp; ++local_row) {
            const int weight_row = row_policy.weight_row(row_begin, local_row);
            row_codes[local_row] = load_fp8_codes<Schedule::kCodeCache, Schedule::kValuesPerLane>(
                weight_codes + static_cast<std::int64_t>(weight_row) * Geometry::kInputRows +
                value_begin);
        }

        Fp8ActivationPack<Schedule::kValuesPerLane> activation[Schedule::kTokenTile];
        if constexpr (Schedule::kActivationAccess == Fp8SmallTActivationAccess::SharedPhase) {
            __syncthreads();
#pragma unroll
            for (int local_token = 0; local_token < Schedule::kTokenTile; ++local_token) {
                if (token0 + local_token < ActiveTokens) {
                    activation[local_token] = load_fp8_activation_pack<Schedule::kValuesPerLane>(
                        shared.activation + local_token * kValuesPerPhase +
                        lane * Schedule::kValuesPerLane);
                }
            }
        } else {
#pragma unroll
            for (int local_token = 0; local_token < Schedule::kTokenTile; ++local_token) {
                const int token = token0 + local_token;
                if (token < ActiveTokens) {
                    activation[local_token] = load_fp8_activation_pack<Schedule::kValuesPerLane>(
                        x + static_cast<std::int64_t>(token) * Geometry::kInputRows + value_begin);
                }
            }
        }

        constexpr int kChainMask = Schedule::kAccumulatorChains - 1;
#pragma unroll
        for (int pair = 0; pair < Schedule::kValuesPerLane / 2; ++pair) {
            float2 weights[Schedule::kRowsPerWarp];
#pragma unroll
            for (int local_row = 0; local_row < Schedule::kRowsPerWarp; ++local_row) {
                const std::uint32_t word   = row_codes[local_row].words[pair >> 1];
                const std::uint16_t packed = static_cast<std::uint16_t>(word >> ((pair & 1) * 16));
                weights[local_row]         = decode_fp8_e4m3x2(packed);
            }
#pragma unroll
            for (int local_token = 0; local_token < Schedule::kTokenTile; ++local_token) {
                if (token0 + local_token >= ActiveTokens) { continue; }
                const float2 value = bf16x2_bits_to_float2(activation[local_token].words[pair]);
#pragma unroll
                for (int local_row = 0; local_row < Schedule::kRowsPerWarp; ++local_row) {
                    accumulators[local_row][local_token][(2 * pair) & kChainMask] =
                        fmaf(weights[local_row].x, value.x,
                             accumulators[local_row][local_token][(2 * pair) & kChainMask]);
                    accumulators[local_row][local_token][(2 * pair + 1) & kChainMask] =
                        fmaf(weights[local_row].y, value.y,
                             accumulators[local_row][local_token][(2 * pair + 1) & kChainMask]);
                }
            }
        }

        if constexpr (Schedule::kActivationAccess == Fp8SmallTActivationAccess::SharedPhase) {
            __syncthreads();
        }
    }

    if constexpr (Finalization == Fp8SmallTFinalization::RowVector) {
        static_assert(!PairRows, "row-vector finalization does not pair output rows");
        static_assert(Schedule::kTokenTile == ActiveTokens,
                      "row-vector finalization requires one CTA to own the full token row");
#pragma unroll
        for (int local_row = 0; local_row < Schedule::kRowsPerWarp; ++local_row) {
            const int parent_row = row_policy.weight_row(row_begin, local_row);
            const float scale    = __bfloat162float(row_scales[parent_row]);
            float projected[ActiveTokens];
#pragma unroll
            for (int local_token = 0; local_token < ActiveTokens; ++local_token) {
                float total = 0.0F;
#pragma unroll
                for (int chain = 0; chain < Schedule::kAccumulatorChains; ++chain) {
                    total += accumulators[local_row][local_token][chain];
                }
                total = warp_reduce_sum(total);
                if (lane == 0) {
                    projected[local_token] = epilogue.apply(parent_row, local_token, total * scale);
                }
            }
            if (lane == 0) { output.store_row(parent_row, projected); }
        }
    } else if constexpr (PairRows) {
#pragma unroll
        for (int local_token = 0; local_token < Schedule::kTokenTile; ++local_token) {
            const int token = token0 + local_token;
            if (token >= ActiveTokens) { continue; }
            float totals[Schedule::kRowsPerWarp];
#pragma unroll
            for (int local_row = 0; local_row < Schedule::kRowsPerWarp; ++local_row) {
                float total = 0.0F;
#pragma unroll
                for (int chain = 0; chain < Schedule::kAccumulatorChains; ++chain) {
                    total += accumulators[local_row][local_token][chain];
                }
                totals[local_row] = warp_reduce_sum(total);
            }
            if (lane == 0) {
#pragma unroll
                for (int local_row = 0; local_row < kStoredRowsPerWarp; ++local_row) {
                    const int first_row = row_policy.weight_row(row_begin, local_row);
                    const int second_row =
                        row_policy.weight_row(row_begin, kStoredRowsPerWarp + local_row);
                    const float first =
                        epilogue.apply(first_row, token,
                                       totals[local_row] * __bfloat162float(row_scales[first_row]));
                    const float second =
                        epilogue.apply(second_row, token,
                                       totals[kStoredRowsPerWarp + local_row] *
                                           __bfloat162float(row_scales[second_row]));
                    output.store_pair(row_begin + local_row, token, first, second);
                }
            }
        }
    } else {
#pragma unroll
        for (int local_row = 0; local_row < Schedule::kRowsPerWarp; ++local_row) {
            const int parent_row = row_policy.weight_row(row_begin, local_row);
            const float scale    = __bfloat162float(row_scales[parent_row]);
#pragma unroll
            for (int local_token = 0; local_token < Schedule::kTokenTile; ++local_token) {
                const int token = token0 + local_token;
                if (token >= ActiveTokens) { continue; }
                float total = 0.0F;
#pragma unroll
                for (int chain = 0; chain < Schedule::kAccumulatorChains; ++chain) {
                    total += accumulators[local_row][local_token][chain];
                }
                total = warp_reduce_sum(total);
                if (lane == 0) {
                    output.store(parent_row, token,
                                 epilogue.apply(parent_row, token, total * scale));
                }
            }
        }
    }
}

} // 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_config.h"
#include "ops/linear/fp8/fp8_gemv.cuh"

#include <cuda_bf16.h>

namespace ninfer::ops::detail {

void fp8_gdn_input_decode_launch(const Tensor& x, const Weight& weight, Tensor& qkv, Tensor& z,
                                 cudaStream_t stream) {
    using Geometry        = Fp8GdnInputGeometry;
    using Schedule        = typename Fp8LinearDecodeProductionSchedule<Geometry>::Type;
    constexpr int kBlocks = Geometry::kOutputRows / Schedule::kRowsPerCta;
    const Fp8GdnInputOutput output{static_cast<__nv_bfloat16*>(qkv.data),
                                   static_cast<__nv_bfloat16*>(z.data)};
    fp8_gemv_kernel<Geometry, 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), output);
    CUDA_CHECK(cudaGetLastError());
}

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

#include "ninfer/ops/softmax_attention.h"

namespace ninfer::ops::detail {

enum class ContextAttentionRoute {
    Direct,
    SplitKv,
};

struct ContextAttentionPlan {
    ContextAttentionRoute route = ContextAttentionRoute::SplitKv;
    std::int32_t tokens         = 0;
    std::int32_t warps          = 0;
    std::int32_t key_block      = 0;
    std::int32_t split_capacity = 0;
};

[[nodiscard]] ContextAttentionPlan
context_attention_resolve_plan(std::int32_t tokens, ContextAttentionExecutionEnvelope envelope);

[[nodiscard]] const char* context_attention_route_name(ContextAttentionRoute route);

void context_attention_launch(const Tensor& q, const Tensor& query_k, const Tensor& query_v,
                              const Tensor& context_lengths, const Tensor& valid_columns,
                              const Tensor& table_rows, float scale,
                              const PagedKVBatchLayerView& context,
                              const ContextAttentionPlan& plan, Tensor& partial_acc,
                              Tensor& partial_m, Tensor& partial_l, Tensor& out,
                              cudaStream_t stream);

} // namespace ninfer::ops::detail
#include <ninfer/targets/qwen3_6_35b_a3b/package.h>
#include <ninfer/targets/qwen3_6/frontend_resources.h>
#include <ninfer/targets/qwen3_6/prepared_prompt.h>

#include "artifact/reader.h"
#include "targets/qwen3_6_35b_a3b/impl/load/bindings.h"
#include "targets/qwen3_6_35b_a3b/impl/variant.h"

#include <stdexcept>
#include <utility>

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

class LoadPlan::Impl {
public:
    Impl(WeightsProfile weights_profile_in, ArtifactLoadPlan target_plan)
        : weights_profile(weights_profile_in), plan(std::move(target_plan)) {}

    WeightsProfile weights_profile;
    ArtifactLoadPlan plan;
};

LoadPlan::LoadPlan(std::unique_ptr<Impl> impl) noexcept : impl_(std::move(impl)) {}

LoadPlan::LoadPlan(LoadPlan&&) noexcept            = default;
LoadPlan& LoadPlan::operator=(LoadPlan&&) noexcept = default;
LoadPlan::~LoadPlan()                              = default;

const artifact::MaterializationPlan& LoadPlan::materialization() const {
    if (impl_ == nullptr) { throw std::logic_error("target load plan is empty"); }
    return impl_->plan.materialization;
}

LoadedModel::LoadedModel(std::unique_ptr<Impl> impl) noexcept : impl_(std::move(impl)) {}

LoadedModel::~LoadedModel() = default;

} // namespace ninfer::targets::qwen3_6_35b_a3b::detail

namespace ninfer::targets::qwen3_6_35b_a3b {
namespace {

constexpr ModelSamplingDefaults kQwen3_6_35BA3BDefaults{
    .thinking     = {.temperature       = 1.0F,
                     .top_k             = 20,
                     .top_p             = 0.95F,
                     .min_p             = 0.0F,
                     .presence_penalty  = 1.5F,
                     .frequency_penalty = 0.0F},
    .non_thinking = {.temperature       = 0.7F,
                     .top_k             = 20,
                     .top_p             = 0.80F,
                     .min_p             = 0.0F,
                     .presence_penalty  = 1.5F,
                     .frequency_penalty = 0.0F},
};

} // namespace

ModelSamplingDefaults Package::sampling_defaults(std::string_view model) {
    if (model == model_id) { return kQwen3_6_35BA3BDefaults; }
    throw std::runtime_error("model '" + std::string(model) +
                             "' has no sampling defaults in target package '" +
                             std::string(target_key) + "'");
}

Package::WeightsProfile Package::resolve_weights(const artifact::ArtifactIdentity& identity) {
    if (identity.model_id == model_id && identity.weights_id == "groupwise-int") {
        return WeightsProfile::GroupwiseInt;
    }
    throw std::runtime_error("artifact identity '" + identity.model_id + "/" + identity.weights_id +
                             "' is not supported by target '" + std::string(target_key) + "'");
}

Package::LoadPlan Package::plan_load(artifact::Binder& binder, const EngineOptions& options,
                                     WeightsProfile weights_profile) {
    return LoadPlan(std::make_unique<LoadPlan::Impl>(
        weights_profile, detail::bind_artifact(binder, qwen3_6::startup_features(options))));
}

std::unique_ptr<Package::LoadedModel>
Package::construct_loaded_model(LoadPlan&& plan, artifact::MaterializedArtifact&& materialized) {
    if (plan.impl_ == nullptr) { throw std::invalid_argument("target load plan is empty"); }
    auto impl = std::make_unique<LoadedModel::Impl>(
        plan.impl_->weights_profile, std::move(plan.impl_->plan.bindings), std::move(materialized));
    plan.impl_.reset();
    return std::unique_ptr<LoadedModel>(new LoadedModel(std::move(impl)));
}

Package::Frontend Package::make_frontend(const LoadedModel& model, const EngineOptions& options) {
    if (model.impl_ == nullptr) { throw std::invalid_argument("loaded model is empty"); }
    return qwen3_6::make_frontend(
        model.impl_->data.frontend,
        qwen3_6::FrontendOptions{
            .vision_enabled                = model.impl_->data.runtime.features.vision,
            .max_context                   = options.max_context,
            .media_cache_bytes             = options.media_cache_bytes,
            .media_live_bytes              = options.media_live_bytes,
            .media_preprocess_threads      = options.media_preprocess_threads,
            .max_cache_markers_per_request = *options.context_cache.max_cache_markers_per_request,
        });
}

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

std::unique_ptr<Package::Program>
Package::create_program(const LoadedModel& model, SequencePlan&& plan, DeviceContext& device) {
    if (model.impl_ == nullptr) { throw std::invalid_argument("loaded model is empty"); }
    return qwen3_6::create_program<detail::Variant>(
        model.impl_->data.runtime, model.impl_->weights_profile, std::move(plan), device);
}

} // namespace ninfer::targets::qwen3_6_35b_a3b
// ninfer::ops - causal_conv1d launcher: grid/block/stream configuration + kernel launch.
#include "ops/launcher/causal_conv1d.h"

#include "ops/common/math.h"
#include "ops/kernel/causal_conv1d.cuh"
#include "core/device.h" // CUDA_CHECK

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

namespace ninfer::ops::detail {

// The small-T routes launch one thread per channel and token in one block, so the bound and the
// channel tile together may not exceed the hardware block ceiling.
static_assert(kCausalConvChannelTile * kCausalConvParallelMaxTokens <= 1024,
              "small-T block would exceed 1024 threads");
namespace {

int grid_for(std::int64_t n, int block, const char* label) {
    const std::int64_t grid = div_up(n, static_cast<std::int64_t>(block));
    if (grid > std::numeric_limits<int>::max()) {
        throw std::overflow_error(std::string("causal_conv1d: ") + label +
                                  " grid exceeds CUDA limit");
    }
    return static_cast<int>(std::max<std::int64_t>(1, grid));
}

int prefill_output_grid_for(std::int32_t C, std::int32_t T, int block) {
    const std::int64_t c_blocks =
        div_up(static_cast<std::int64_t>(C), static_cast<std::int64_t>(block));
    const std::int64_t grid = c_blocks * static_cast<std::int64_t>(T);
    if (grid > std::numeric_limits<int>::max()) {
        throw std::overflow_error("causal_conv1d: prefill output grid exceeds CUDA limit");
    }
    return static_cast<int>(std::max<std::int64_t>(1, grid));
}

} // namespace

void causal_conv1d_prefill_launch(const Tensor& x, const Tensor& weight,
                                  const Tensor& conv_state_in, Tensor& conv_state_out, Tensor& out,
                                  cudaStream_t stream) {
    constexpr int kOutputBlock  = 256;
    constexpr int kChannelBlock = 256;
    constexpr int kPairBlock    = 256;
    const std::int32_t C        = x.ne[0];
    const std::int32_t T        = x.ne[1];
    const auto x_addr           = reinterpret_cast<std::uintptr_t>(x.data);
    const auto w_addr           = reinterpret_cast<std::uintptr_t>(weight.data);
    const auto in_addr          = reinterpret_cast<std::uintptr_t>(conv_state_in.data);
    const auto out_state_addr   = reinterpret_cast<std::uintptr_t>(conv_state_out.data);
    const auto out_addr         = reinterpret_cast<std::uintptr_t>(out.data);

    if (((x_addr | w_addr | in_addr | out_state_addr | out_addr) & (alignof(__nv_bfloat162) - 1)) ==
            0 &&
        (C & 1) == 0) {
        const CausalConvContiguousOutput<__nv_bfloat162> out_map{
            reinterpret_cast<__nv_bfloat162*>(out.data), C / 2};
        causal_conv1d_prefill_pairs_kernel<<<prefill_output_grid_for(C / 2, T, kPairBlock),
                                             kPairBlock, 0, stream>>>(
            static_cast<const __nv_bfloat16*>(x.data),
            static_cast<const __nv_bfloat16*>(weight.data),
            static_cast<const __nv_bfloat16*>(conv_state_in.data), out_map, C, T);
        CUDA_CHECK(cudaGetLastError());
    } else {
        const CausalConvContiguousOutput<__nv_bfloat16> out_map{
            static_cast<__nv_bfloat16*>(out.data), C};
        causal_conv1d_prefill_kernel<<<prefill_output_grid_for(C, T, kOutputBlock), kOutputBlock, 0,
                                       stream>>>(
            static_cast<const __nv_bfloat16*>(x.data),
            static_cast<const __nv_bfloat16*>(weight.data),
            static_cast<const __nv_bfloat16*>(conv_state_in.data), out_map, C, T);
        CUDA_CHECK(cudaGetLastError());
    }

    causal_conv1d_prefill_state_kernel<<<grid_for(C, kChannelBlock, "prefill state"), kChannelBlock,
                                         0, stream>>>(
        static_cast<const __nv_bfloat16*>(x.data),
        static_cast<const __nv_bfloat16*>(conv_state_in.data),
        static_cast<__nv_bfloat16*>(conv_state_out.data), C, T);
    CUDA_CHECK(cudaGetLastError());
}

void causal_conv1d_smallt_launch(const Tensor& x, const Tensor& weight, const Tensor& conv_state_in,
                                 Tensor& conv_state_out, Tensor& out, cudaStream_t stream) {
    const std::int32_t C = x.ne[0];
    const std::int32_t T = x.ne[1];
    const dim3 block(kCausalConvChannelTile, static_cast<unsigned int>(T));
    const int grid = grid_for(C, kCausalConvChannelTile, "small-T");

    const CausalConvContiguousOutput<__nv_bfloat16> out_map{static_cast<__nv_bfloat16*>(out.data),
                                                            C};
    causal_conv1d_smallt_kernel<<<grid, block, 0, stream>>>(
        static_cast<const __nv_bfloat16*>(x.data), static_cast<const __nv_bfloat16*>(weight.data),
        static_cast<const __nv_bfloat16*>(conv_state_in.data),
        static_cast<__nv_bfloat16*>(conv_state_out.data), out_map, C, T);
    CUDA_CHECK(cudaGetLastError());
}

template <std::int32_t Rows0, std::int32_t Rows1, std::int32_t Rows2>
void smallt_split_launch(const Tensor& x, const Tensor& weight, const Tensor& conv_state_in,
                         Tensor& conv_state_out, Tensor& out0, Tensor& out1, Tensor& out2,
                         cudaStream_t stream) {
    const std::int32_t C = x.ne[0];
    const std::int32_t T = x.ne[1];
    if (Rows0 + Rows1 + Rows2 != C) {
        throw std::invalid_argument("causal_conv1d: split geometry does not cover the channels");
    }
    const dim3 block(kCausalConvChannelTile, static_cast<unsigned int>(T));
    const int grid = grid_for(C, kCausalConvChannelTile, "small-T");

    const CausalConvSplitOutput3<__nv_bfloat16, Rows0, Rows1, Rows2> out_map{
        static_cast<__nv_bfloat16*>(out0.data), static_cast<__nv_bfloat16*>(out1.data),
        static_cast<__nv_bfloat16*>(out2.data)};
    causal_conv1d_smallt_kernel<<<grid, block, 0, stream>>>(
        static_cast<const __nv_bfloat16*>(x.data), static_cast<const __nv_bfloat16*>(weight.data),
        static_cast<const __nv_bfloat16*>(conv_state_in.data),
        static_cast<__nv_bfloat16*>(conv_state_out.data), out_map, C, T);
    CUDA_CHECK(cudaGetLastError());
}

void causal_conv1d_smallt_split_launch(const Tensor& x, const Tensor& weight,
                                       const Tensor& conv_state_in, Tensor& conv_state_out,
                                       Tensor& out0, Tensor& out1, Tensor& out2,
                                       CausalConvSplitGeometry geometry, cudaStream_t stream) {
    switch (geometry) {
    case CausalConvSplitGeometry::Rows2048x2048x4096:
        smallt_split_launch<2048, 2048, 4096>(x, weight, conv_state_in, conv_state_out, out0, out1,
                                              out2, stream);
        return;
    case CausalConvSplitGeometry::Rows2048x2048x6144:
        smallt_split_launch<2048, 2048, 6144>(x, weight, conv_state_in, conv_state_out, out0, out1,
                                              out2, stream);
        return;
    }
    throw std::logic_error("causal_conv1d: unhandled split geometry");
}

template <std::int32_t Rows0, std::int32_t Rows1, std::int32_t Rows2>
void prefill_split_launch(const Tensor& x, const Tensor& weight, const Tensor& conv_state_in,
                          Tensor& conv_state_out, Tensor& out0, Tensor& out1, Tensor& out2,
                          cudaStream_t stream) {
    static_assert((Rows0 % 2) == 0 && (Rows1 % 2) == 0 && (Rows2 % 2) == 0,
                  "a pair must not straddle two destinations");
    constexpr int kChannelBlock = 256;
    constexpr int kPairBlock    = 256;
    const std::int32_t C        = x.ne[0];
    const std::int32_t T        = x.ne[1];
    if (Rows0 + Rows1 + Rows2 != C) {
        throw std::invalid_argument("causal_conv1d: split geometry does not cover the channels");
    }

    const CausalConvSplitOutput3<__nv_bfloat162, Rows0 / 2, Rows1 / 2, Rows2 / 2> out_map{
        reinterpret_cast<__nv_bfloat162*>(out0.data), reinterpret_cast<__nv_bfloat162*>(out1.data),
        reinterpret_cast<__nv_bfloat162*>(out2.data)};
    causal_conv1d_prefill_pairs_kernel<<<prefill_output_grid_for(C / 2, T, kPairBlock), kPairBlock,
                                         0, stream>>>(
        static_cast<const __nv_bfloat16*>(x.data), static_cast<const __nv_bfloat16*>(weight.data),
        static_cast<const __nv_bfloat16*>(conv_state_in.data), out_map, C, T);
    CUDA_CHECK(cudaGetLastError());

    causal_conv1d_prefill_state_kernel<<<grid_for(C, kChannelBlock, "prefill state"), kChannelBlock,
                                         0, stream>>>(
        static_cast<const __nv_bfloat16*>(x.data),
        static_cast<const __nv_bfloat16*>(conv_state_in.data),
        static_cast<__nv_bfloat16*>(conv_state_out.data), C, T);
    CUDA_CHECK(cudaGetLastError());
}

void causal_conv1d_prefill_split_launch(const Tensor& x, const Tensor& weight,
                                        const Tensor& conv_state_in, Tensor& conv_state_out,
                                        Tensor& out0, Tensor& out1, Tensor& out2,
                                        CausalConvSplitGeometry geometry, cudaStream_t stream) {
    switch (geometry) {
    case CausalConvSplitGeometry::Rows2048x2048x4096:
        prefill_split_launch<2048, 2048, 4096>(x, weight, conv_state_in, conv_state_out, out0, out1,
                                               out2, stream);
        return;
    case CausalConvSplitGeometry::Rows2048x2048x6144:
        prefill_split_launch<2048, 2048, 6144>(x, weight, conv_state_in, conv_state_out, out0, out1,
                                               out2, stream);
        return;
    }
    throw std::logic_error("causal_conv1d: unhandled split geometry");
}

void causal_conv1d_decode_launch(const Tensor& x, const Tensor& weight, const Tensor& conv_state_in,
                                 Tensor& conv_state_out, Tensor& out, cudaStream_t stream) {
    constexpr int kBlock = 256;
    const std::int32_t C = x.ne[0];

    if (conv_state_in.data == conv_state_out.data) {
        causal_conv1d_decode_kernel<<<grid_for(C, kBlock, "decode"), kBlock, 0, stream>>>(
            static_cast<const __nv_bfloat16*>(x.data),
            static_cast<const __nv_bfloat16*>(weight.data),
            static_cast<__nv_bfloat16*>(conv_state_out.data), static_cast<__nv_bfloat16*>(out.data),
            C);
    } else {
        causal_conv1d_decode_distinct_kernel<<<grid_for(C, kBlock, "distinct decode"), kBlock, 0,
                                               stream>>>(
            static_cast<const __nv_bfloat16*>(x.data),
            static_cast<const __nv_bfloat16*>(weight.data),
            static_cast<const __nv_bfloat16*>(conv_state_in.data),
            static_cast<__nv_bfloat16*>(conv_state_out.data), static_cast<__nv_bfloat16*>(out.data),
            C);
    }
    CUDA_CHECK(cudaGetLastError());
}

template <bool Masked>
void launch_batched_sequence_snapshot(const Tensor& x, const Tensor& weight, Tensor& conv_states,
                                      const Tensor& valid_columns,
                                      const Tensor& initial_state_slots,
                                      const Tensor& snapshot_base_slots, Tensor& out,
                                      std::int64_t slot_stride, cudaStream_t stream) {
    const std::int32_t C     = x.ne[0];
    const std::int32_t width = x.ne[1];
    const std::int32_t batch = x.ne[2];
    const int block          = width == 1 ? 256 : 32;
    const dim3 grid(grid_for(C, block, "batched sequence snapshot"),
                    static_cast<unsigned int>(batch));
    causal_conv1d_batched_sequence_snapshot_kernel<Masked><<<grid, block, 0, stream>>>(
        static_cast<const __nv_bfloat16*>(x.data), static_cast<const __nv_bfloat16*>(weight.data),
        static_cast<__nv_bfloat16*>(conv_states.data),
        Masked ? static_cast<const std::int32_t*>(valid_columns.data) : nullptr,
        static_cast<const std::int32_t*>(initial_state_slots.data),
        static_cast<const std::int32_t*>(snapshot_base_slots.data),
        static_cast<__nv_bfloat16*>(out.data), C, width, slot_stride);
}

template <bool Masked>
void launch_batched_smallt_snapshot(const Tensor& x, const Tensor& weight, Tensor& conv_states,
                                    const Tensor& valid_columns, const Tensor& initial_state_slots,
                                    const Tensor& snapshot_base_slots, Tensor& out,
                                    std::int64_t slot_stride, cudaStream_t stream) {
    const std::int32_t C     = x.ne[0];
    const std::int32_t width = x.ne[1];
    const std::int32_t batch = x.ne[2];
    const dim3 block(kCausalConvChannelTile, static_cast<unsigned int>(width));
    const dim3 grid(grid_for(C, kCausalConvChannelTile, "batched small-T snapshot"),
                    static_cast<unsigned int>(batch));
    causal_conv1d_batched_snapshot_smallt_kernel<Masked><<<grid, block, 0, stream>>>(
        static_cast<const __nv_bfloat16*>(x.data), static_cast<const __nv_bfloat16*>(weight.data),
        static_cast<__nv_bfloat16*>(conv_states.data),
        Masked ? static_cast<const std::int32_t*>(valid_columns.data) : nullptr,
        static_cast<const std::int32_t*>(initial_state_slots.data),
        static_cast<const std::int32_t*>(snapshot_base_slots.data),
        static_cast<__nv_bfloat16*>(out.data), C, width, slot_stride);
}

void causal_conv1d_snapshot_launch(const Tensor& x, const Tensor& weight, Tensor& conv_states,
                                   const Tensor& valid_columns, const Tensor& initial_state_slots,
                                   const Tensor& snapshot_base_slots, Tensor& out,
                                   cudaStream_t stream) {
    const std::int32_t C = x.ne[0];
    const std::int32_t T = x.ne[1];
    const std::int32_t B = x.ne[2];
    const std::int64_t slot_stride =
        static_cast<std::int64_t>(conv_states.ne[0]) * static_cast<std::int64_t>(conv_states.ne[1]);

    if (B == 1 && valid_columns.data == nullptr && T == 1) {
        constexpr int kBlock = 256;
        causal_conv1d_snapshot_decode_kernel<<<grid_for(C, kBlock, "snapshot decode"), kBlock, 0,
                                               stream>>>(
            static_cast<const __nv_bfloat16*>(x.data),
            static_cast<const __nv_bfloat16*>(weight.data),
            static_cast<__nv_bfloat16*>(conv_states.data),
            static_cast<const std::int32_t*>(initial_state_slots.data),
            static_cast<const std::int32_t*>(snapshot_base_slots.data),
            static_cast<__nv_bfloat16*>(out.data), C, slot_stride);
    } else if (B == 1 && valid_columns.data == nullptr && T <= kCausalConvParallelMaxTokens) {
        const dim3 block(kCausalConvChannelTile, static_cast<unsigned int>(T));
        const int grid = grid_for(C, kCausalConvChannelTile, "small-T snapshot");
        causal_conv1d_snapshot_smallt_kernel<<<grid, block, 0, stream>>>(
            static_cast<const __nv_bfloat16*>(x.data),
            static_cast<const __nv_bfloat16*>(weight.data),
            static_cast<__nv_bfloat16*>(conv_states.data),
            static_cast<const std::int32_t*>(initial_state_slots.data),
            static_cast<const std::int32_t*>(snapshot_base_slots.data),
            static_cast<__nv_bfloat16*>(out.data), C, T, slot_stride);
    } else if (B == 1 && valid_columns.data == nullptr) {
        constexpr int kBlock = 32;
        causal_conv1d_sequence_snapshot_kernel<<<grid_for(C, kBlock, "sequence snapshot"), kBlock,
                                                 0, stream>>>(
            static_cast<const __nv_bfloat16*>(x.data),
            static_cast<const __nv_bfloat16*>(weight.data),
            static_cast<__nv_bfloat16*>(conv_states.data),
            static_cast<const std::int32_t*>(initial_state_slots.data),
            static_cast<const std::int32_t*>(snapshot_base_slots.data),
            static_cast<__nv_bfloat16*>(out.data), C, T, slot_stride);
    } else if (T == 1) {
        if (valid_columns.data == nullptr) {
            launch_batched_sequence_snapshot<false>(x, weight, conv_states, valid_columns,
                                                    initial_state_slots, snapshot_base_slots, out,
                                                    slot_stride, stream);
        } else {
            launch_batched_sequence_snapshot<true>(x, weight, conv_states, valid_columns,
                                                   initial_state_slots, snapshot_base_slots, out,
                                                   slot_stride, stream);
        }
    } else if (T <= kCausalConvParallelMaxTokens) {
        if (valid_columns.data == nullptr) {
            launch_batched_smallt_snapshot<false>(x, weight, conv_states, valid_columns,
                                                  initial_state_slots, snapshot_base_slots, out,
                                                  slot_stride, stream);
        } else {
            launch_batched_smallt_snapshot<true>(x, weight, conv_states, valid_columns,
                                                 initial_state_slots, snapshot_base_slots, out,
                                                 slot_stride, stream);
        }
    } else if (valid_columns.data == nullptr) {
        launch_batched_sequence_snapshot<false>(x, weight, conv_states, valid_columns,
                                                initial_state_slots, snapshot_base_slots, out,
                                                slot_stride, stream);
    } else {
        launch_batched_sequence_snapshot<true>(x, weight, conv_states, valid_columns,
                                               initial_state_slots, snapshot_base_slots, out,
                                               slot_stride, stream);
    }
    CUDA_CHECK(cudaGetLastError());
}

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

// Typed Qwen3.6 phase-root allocation clusters shared by the real schedule and its startup
// WorkspaceLayoutBuilder simulation. Child Op scratch remains owned by each Op capacity query.

#include "core/arena.h"
#include "core/layout.h"

#include <cstdint>

namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::workspace_recipe {

template <class Allocator>
Tensor matrix(Allocator& allocator, DType dtype, std::int32_t rows, std::int32_t tokens) {
    return allocator.alloc(dtype, {rows, tokens});
}

template <class Allocator>
Tensor vector(Allocator& allocator, DType dtype, std::int32_t elements) {
    return allocator.alloc(dtype, {elements});
}

struct TextPrefillRoots {
    Tensor ids;
    Tensor positions;
    Tensor rope_positions;
    Tensor residual;
    Tensor scatter_indices;
};

template <class Config, class Allocator>
TextPrefillRoots text_prefill_roots(Allocator& allocator, std::int32_t tokens,
                                    std::int32_t rope_axes, std::int32_t scatter_tokens) {
    TextPrefillRoots out;
    out.ids       = vector(allocator, DType::I32, tokens);
    out.positions = vector(allocator, DType::I32, tokens);
    if (rope_axes != 0) { out.rope_positions = matrix(allocator, DType::I32, tokens, rope_axes); }
    out.residual = matrix(allocator, DType::BF16, Config::hidden, tokens);
    if (scatter_tokens != 0) {
        out.scatter_indices = vector(allocator, DType::I32, scatter_tokens);
    }
    return out;
}

template <class Allocator>
Tensor visual_scatter_indices(Allocator& allocator, std::int32_t tokens) {
    return vector(allocator, DType::I32, tokens);
}

struct TextAttentionProjectionRoots {
    Tensor hidden;
    Tensor query;
    Tensor gate;
    Tensor key;
    Tensor value;
};

template <class Config, class Allocator>
TextAttentionProjectionRoots text_attention_projection(Allocator& allocator, std::int32_t tokens) {
    return {
        matrix(allocator, DType::BF16, Config::hidden, tokens),
        matrix(allocator, DType::BF16, Config::query_size, tokens),
        matrix(allocator, DType::BF16, Config::query_size, tokens),
        matrix(allocator, DType::BF16, Config::kv_size, tokens),
        matrix(allocator, DType::BF16, Config::kv_size, tokens),
    };
}

struct TextAttentionResultRoots {
    Tensor normalized_query;
    Tensor normalized_key;
    Tensor attention;
};

template <class Config, class Allocator>
TextAttentionResultRoots text_attention_results(Allocator& allocator, std::int32_t tokens) {
    return {
        matrix(allocator, DType::BF16, Config::query_size, tokens),
        matrix(allocator, DType::BF16, Config::kv_size, tokens),
        matrix(allocator, DType::BF16, Config::query_size, tokens),
    };
}

struct GdnControlRoots {
    Tensor hidden;
    Tensor g;
    Tensor beta;
};

template <class Config, class Allocator>
GdnControlRoots gdn_control(Allocator& allocator, std::int32_t tokens) {
    return {
        matrix(allocator, DType::BF16, Config::hidden, tokens),
        matrix(allocator, DType::FP32, Config::gdn_value_heads, tokens),
        matrix(allocator, DType::FP32, Config::gdn_value_heads, tokens),
    };
}

struct GdnProjectionRoots {
    Tensor output_gate;
    Tensor query;
    Tensor key;
    Tensor value;
};

template <class Config, class Allocator>
GdnProjectionRoots gdn_projection(Allocator& allocator, std::int32_t tokens) {
    return {
        matrix(allocator, DType::BF16, Config::value_dim, tokens),
        matrix(allocator, DType::BF16, Config::key_dim, tokens),
        matrix(allocator, DType::BF16, Config::key_dim, tokens),
        matrix(allocator, DType::BF16, Config::value_dim, tokens),
    };
}

template <class Config, class Allocator>
Tensor gdn_prefill_conv(Allocator& allocator, std::int32_t tokens) {
    return matrix(allocator, DType::BF16, Config::convolution_dim, tokens);
}

template <class Config, class Allocator>
Tensor gdn_recurrent_output(Allocator& allocator, std::int32_t tokens) {
    return matrix(allocator, DType::BF16, Config::value_dim, tokens);
}

template <class Config, class Allocator>
Tensor gdn_normalized_output(Allocator& allocator, std::int32_t tokens) {
    return matrix(allocator, DType::BF16, Config::value_dim, tokens);
}

template <class Config, class Allocator>
Tensor post_mixer_hidden(Allocator& allocator, std::int32_t tokens) {
    return matrix(allocator, DType::BF16, Config::hidden, tokens);
}

struct MtpStemRoots {
    Tensor embedding;
    Tensor normalized_embedding;
    Tensor normalized_hidden;
    Tensor packed_input;
    Tensor residual;
    Tensor attention_hidden;
};

template <class Config, class Allocator>
MtpStemRoots mtp_stem(Allocator& allocator, std::int32_t tokens, bool allocate_embedding) {
    MtpStemRoots out;
    if (allocate_embedding) {
        out.embedding = matrix(allocator, DType::BF16, Config::hidden, tokens);
    }
    out.normalized_embedding = matrix(allocator, DType::BF16, Config::hidden, tokens);
    out.normalized_hidden    = matrix(allocator, DType::BF16, Config::hidden, tokens);
    out.packed_input         = matrix(allocator, DType::BF16, Config::mtp_input_rows, tokens);
    out.residual             = matrix(allocator, DType::BF16, Config::hidden, tokens);
    out.attention_hidden     = matrix(allocator, DType::BF16, Config::hidden, tokens);
    return out;
}

struct MtpAttentionProjectionRoots {
    Tensor query;
    Tensor key;
    Tensor gate;
    Tensor value;
};

template <class Config, class Allocator>
MtpAttentionProjectionRoots mtp_attention_projection(Allocator& allocator, std::int32_t tokens) {
    return {
        matrix(allocator, DType::BF16, Config::query_size, tokens),
        matrix(allocator, DType::BF16, Config::kv_size, tokens),
        matrix(allocator, DType::BF16, Config::query_size, tokens),
        matrix(allocator, DType::BF16, Config::kv_size, tokens),
    };
}

struct MtpAttentionResultRoots {
    Tensor normalized_query;
    Tensor normalized_key;
    Tensor attention;
};

template <class Config, class Allocator>
MtpAttentionResultRoots mtp_attention_results(Allocator& allocator, std::int32_t tokens) {
    return {
        matrix(allocator, DType::BF16, Config::query_size, tokens),
        matrix(allocator, DType::BF16, Config::kv_size, tokens),
        matrix(allocator, DType::BF16, Config::query_size, tokens),
    };
}

struct MtpPostAttentionRoots {
    Tensor output;
    Tensor post_mixer_hidden;
};

template <class Config, class Allocator>
MtpPostAttentionRoots mtp_post_attention(Allocator& allocator, std::int32_t tokens) {
    return {
        matrix(allocator, DType::BF16, Config::hidden, tokens),
        matrix(allocator, DType::BF16, Config::hidden, tokens),
    };
}

struct DFlashContextRoots {
    Tensor projected;
    Tensor normalized;
};

template <class Config, class Allocator>
DFlashContextRoots dflash_context(Allocator& allocator, std::int32_t tokens) {
    return {
        matrix(allocator, DType::BF16, Config::hidden, tokens),
        matrix(allocator, DType::BF16, Config::hidden, tokens),
    };
}

struct DFlashContextLayerRoots {
    Tensor key_raw;
    Tensor value;
    Tensor key;
};

template <class Config, class Allocator>
DFlashContextLayerRoots dflash_context_layer(Allocator& allocator, std::int32_t tokens) {
    return {
        matrix(allocator, DType::BF16, Config::kv_size, tokens),
        matrix(allocator, DType::BF16, Config::kv_size, tokens),
        matrix(allocator, DType::BF16, Config::kv_size, tokens),
    };
}

struct DFlashProposalRoots {
    Tensor ids;
    Tensor positions;
    Tensor residual;
};

template <class Config, class Allocator>
DFlashProposalRoots dflash_proposal(Allocator& allocator, std::int32_t tokens) {
    return {
        vector(allocator, DType::I32, tokens),
        vector(allocator, DType::I32, tokens),
        matrix(allocator, DType::BF16, Config::hidden, tokens),
    };
}

struct DFlashAttentionRoots {
    Tensor hidden;
    Tensor query_raw;
    Tensor key_raw;
    Tensor value;
    Tensor query;
    Tensor key;
    Tensor attention;
};

template <class Config, class Allocator>
DFlashAttentionRoots dflash_attention(Allocator& allocator, std::int32_t tokens) {
    return {
        matrix(allocator, DType::BF16, Config::hidden, tokens),
        matrix(allocator, DType::BF16, Config::query_size, tokens),
        matrix(allocator, DType::BF16, Config::kv_size, tokens),
        matrix(allocator, DType::BF16, Config::kv_size, tokens),
        matrix(allocator, DType::BF16, Config::query_size, tokens),
        matrix(allocator, DType::BF16, Config::kv_size, tokens),
        matrix(allocator, DType::BF16, Config::query_size, tokens),
    };
}

struct DFlashMlpRoots {
    Tensor hidden;
    Tensor intermediate;
};

template <class Config, class Allocator>
DFlashMlpRoots dflash_mlp(Allocator& allocator, std::int32_t tokens) {
    return {
        matrix(allocator, DType::BF16, Config::hidden, tokens),
        matrix(allocator, DType::BF16, Config::intermediate, tokens),
    };
}

} // namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::workspace_recipe
#include "ops/attn_input_proj/q4_q5/q4_q5_attn_input_kernels.h"

#include "core/device.h"
#include "ops/common/math.h"
#include "ops/linear/q4/q4_rowsplit_gemm_simt.cuh"
#include "ops/linear/q4/q4_rowsplit_gemv.cuh"
#include "ops/linear/q5/q5_rowsplit_gemm_simt.cuh"
#include "ops/linear/q5/q5_rowsplit_gemv.cuh"

#include <cuda_bf16.h>

#include <cstdint>
#include <stdexcept>

namespace ninfer::ops::detail {
namespace {

constexpr std::int32_t kParentRows = 7168;
constexpr std::int32_t kSplitRow   = 6144;
constexpr std::int32_t kHidden     = 5120;

using Q4AttnSimtR8C4Schedule = Q4RowSplitSimtGemmSchedule<8, 4, 16, 2, Cache::ca, 1>;
using Q4AttnSimtR8C8Schedule = Q4RowSplitSimtGemmSchedule<8, 8, 16, 2, Cache::ca, 1>;

void launch_q4_gemv(const Tensor& x, const Weight& weight, Tensor& q, Tensor& key,
                    cudaStream_t stream) {
    using Schedule = Q4GemvR1W8DirectSchedule;
    const dim3 grid(static_cast<unsigned>(div_up(kParentRows, Schedule::kRowsPerCta)), 1u, 1u);
    constexpr dim3 block(static_cast<unsigned>(Schedule::kThreads), 1u, 1u);
    q4_rowsplit_gemv_kernel<Schedule, true, kSplitRow><<<grid, block, 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), static_cast<__nv_bfloat16*>(q.data),
        static_cast<__nv_bfloat16*>(key.data), kParentRows, kHidden);
    CUDA_CHECK(cudaGetLastError());
}

template <class Schedule, bool Full>
void launch_q4_simt(const Tensor& x, const Weight& weight, Tensor& q, Tensor& key,
                    cudaStream_t stream) {
    const std::int32_t cols = x.ne[1];
    const dim3 grid(static_cast<unsigned>(div_up(kParentRows, Schedule::kRowsPerCta)),
                    static_cast<unsigned>(div_up(cols, Schedule::kColsPerTile)), 1u);
    q4_rowsplit_gemm_simt_kernel<Schedule, Full, true, kSplitRow>
        <<<grid, 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), static_cast<__nv_bfloat16*>(q.data),
            static_cast<__nv_bfloat16*>(key.data), q.ne[0], key.ne[0], kParentRows, kHidden, cols,
            weight.padded_shape[1]);
    CUDA_CHECK(cudaGetLastError());
}

template <class Schedule>
void launch_q4_simt_route(const Tensor& x, const Weight& weight, Tensor& q, Tensor& key,
                          cudaStream_t stream) {
    const bool full = (kParentRows % Schedule::kRowsPerCta) == 0 &&
                      ((kHidden / Q4RowSplitStorage::kGroupK) % Schedule::kGroupsPerStage) == 0 &&
                      (x.ne[1] % Schedule::kColsPerTile) == 0;
    if (full) {
        launch_q4_simt<Schedule, true>(x, weight, q, key, stream);
    } else {
        launch_q4_simt<Schedule, false>(x, weight, q, key, stream);
    }
}

void launch_q4(const Tensor& x, const Weight& weight, Tensor& q, Tensor& key, cudaStream_t stream) {
    switch (x.ne[1]) {
    case 1:
        launch_q4_gemv(x, weight, q, key, stream);
        return;
    case 2:
    case 3:
    case 4:
    case 5:
    case 6:
    case 7:
    case 9:
    case 10:
    case 11:
    case 12:
    case 13:
    case 14:
    case 15:
        launch_q4_simt_route<Q4AttnSimtR8C4Schedule>(x, weight, q, key, stream);
        return;
    case 8:
    case 16:
        launch_q4_simt_route<Q4AttnSimtR8C8Schedule>(x, weight, q, key, stream);
        return;
    default:
        throw std::invalid_argument("attention Q4 split-output requires T in [1,16]");
    }
}

void launch_q5_gemv(const Tensor& x, const Weight& weight, Tensor& gate, Tensor& value,
                    cudaStream_t stream) {
    constexpr int kRowsPerBlock = 16;
    constexpr int kBlockThreads = kRowsPerBlock * 32;
    constexpr int kGrid         = kParentRows / kRowsPerBlock;
    q5_rowsplit_gemv_kernel<kParentRows, kHidden, kRowsPerBlock, 2, true, false, true, kSplitRow>
        <<<kGrid, kBlockThreads, 0, stream>>>(static_cast<const __nv_bfloat16*>(x.data),
                                              static_cast<const std::uint8_t*>(weight.qdata),
                                              static_cast<const std::uint8_t*>(weight.qhigh),
                                              static_cast<const std::uint8_t*>(weight.scales),
                                              static_cast<__nv_bfloat16*>(gate.data),
                                              static_cast<__nv_bfloat16*>(value.data));
    CUDA_CHECK(cudaGetLastError());
}

template <int Cols>
void launch_q5_split4(const Tensor& x, const Weight& weight, Tensor& gate, Tensor& value,
                      cudaStream_t stream) {
    constexpr int kThreads = 4 * 32;
    const dim3 grid(static_cast<unsigned>(kParentRows), 1u, 1u);
    q5_rowsplit_gemm_simt_split4_kernel<Q5RowSplitSimtSchedule, Cols, 5, kHidden, true, kSplitRow>
        <<<grid, 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.qhigh),
                                        static_cast<const std::uint8_t*>(weight.scales),
                                        static_cast<__nv_bfloat16*>(gate.data),
                                        static_cast<__nv_bfloat16*>(value.data), kParentRows,
                                        gate.ne[0], kHidden, Cols, weight.padded_shape[1], 5);
    CUDA_CHECK(cudaGetLastError());
}

void launch_q5_split4_exact(const Tensor& x, const Weight& weight, Tensor& gate, Tensor& value,
                            cudaStream_t stream) {
    switch (x.ne[1]) {
    case 2:
        launch_q5_split4<2>(x, weight, gate, value, stream);
        return;
    case 3:
        launch_q5_split4<3>(x, weight, gate, value, stream);
        return;
    case 4:
        launch_q5_split4<4>(x, weight, gate, value, stream);
        return;
    case 5:
        launch_q5_split4<5>(x, weight, gate, value, stream);
        return;
    case 6:
        launch_q5_split4<6>(x, weight, gate, value, stream);
        return;
    default:
        throw std::invalid_argument("attention Q5 split4 requires T in [2,6]");
    }
}

template <int ColsPerTile>
void launch_q5_simt(const Tensor& x, const Weight& weight, Tensor& gate, Tensor& value,
                    cudaStream_t stream) {
    constexpr int kRowsPerBlock = 8;
    constexpr int kStages       = 2;
    constexpr int kThreads      = kRowsPerBlock * 32;
    const std::int32_t cols     = x.ne[1];
    const dim3 grid(static_cast<unsigned>(div_up(kParentRows, kRowsPerBlock)),
                    static_cast<unsigned>(div_up(cols, ColsPerTile)), 1u);
    q5_rowsplit_gemm_simt_kernel<Q5RowSplitSimtSchedule, ColsPerTile, kRowsPerBlock, kStages, true,
                                 kSplitRow><<<grid, 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.qhigh),
        static_cast<const std::uint8_t*>(weight.scales), static_cast<__nv_bfloat16*>(gate.data),
        static_cast<__nv_bfloat16*>(value.data), kParentRows, gate.ne[0], kHidden, cols,
        weight.padded_shape[1], 5);
    CUDA_CHECK(cudaGetLastError());
}

void launch_q5(const Tensor& x, const Weight& weight, Tensor& gate, Tensor& value,
               cudaStream_t stream) {
    if (x.ne[1] == 1) {
        launch_q5_gemv(x, weight, gate, value, stream);
        return;
    }
    if (x.ne[1] <= 6) {
        launch_q5_split4_exact(x, weight, gate, value, stream);
        return;
    }
    if (x.ne[1] <= 16) {
        launch_q5_simt<4>(x, weight, gate, value, stream);
        return;
    }
    throw std::invalid_argument("attention Q5 split-output requires T in [1,16]");
}

} // namespace

void q4_q5_attn_input_small_t_launch(const Tensor& x, const Weight& query_key_weight,
                                     const Weight& gate_value_weight, Tensor& q, Tensor& gate,
                                     Tensor& k, Tensor& v, cudaStream_t stream) {
    launch_q4(x, query_key_weight, q, k, stream);
    launch_q5(x, gate_value_weight, gate, v, stream);
}

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

#include "ops/common/mma.cuh"
#include "ops/linear_attention/gated_delta_net/chunked/common.cuh"

// Stage 3: chunk-sequential state passing.
//
// Math + I/O layouts: see the Gated DeltaNet chunked state_passing_config.
// Two fixed geometries share one kernel body:
//   value strip 16: 8 warps, ~44 KB dynamic smem, two CTAs/SM;
//   value strip 32: 16 warps, ~56 KB dynamic smem, one CTA/SM.

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

using ninfer::ops::mma_tf32_bits;
using ninfer::ops::ldmatrix_x2;
using ninfer::ops::ldmatrix_x2_t;
using ninfer::ops::smem_addr;
using ninfer::ops::exp2_approx;

static_assert(kChunkSize == 64, "stage_state_passing: kChunkSize must be 64");
static_assert(kStateDim == 128);

template <int NStrip>
struct kernel_dims;

template <>
struct kernel_dims<16> {
    static constexpr int N_STRIP_PER_BLOCK  = 16;
    static constexpr int D_STRIPS           = kStateDim / N_STRIP_PER_BLOCK;
    static constexpr int DT_TILES_PER_BLOCK = N_STRIP_PER_BLOCK / MMA_N;
    static constexpr int BT_SPLITS          = 4;
    static constexpr int N_WARPS            = DT_TILES_PER_BLOCK * BT_SPLITS; // 8
    static constexpr int THREADS            = N_WARPS * ninfer::ops::kWarpSize;
    static constexpr int MIN_BLOCKS         = 2;
};

template <>
struct kernel_dims<32> {
    static constexpr int N_STRIP_PER_BLOCK  = 32;
    static constexpr int D_STRIPS           = kStateDim / N_STRIP_PER_BLOCK;
    static constexpr int DT_TILES_PER_BLOCK = N_STRIP_PER_BLOCK / MMA_N;
    static constexpr int BT_SPLITS          = 4;
    static constexpr int N_WARPS            = DT_TILES_PER_BLOCK * BT_SPLITS; // 16
    static constexpr int THREADS            = N_WARPS * ninfer::ops::kWarpSize;
    static constexpr int MIN_BLOCKS         = 1;
};

template <int NStrip>
struct smem_layout {
    using D                             = kernel_dims<NStrip>;
    static constexpr int W_STRIDE       = kStateDim;
    static constexpr int W_LOAD_BF16    = BT * W_STRIDE;
    static constexpr int W_STORAGE_FLT  = W_LOAD_BF16 / 2;
    static constexpr int UVD_FLT        = BT * D::N_STRIP_PER_BLOCK; // U-vd alias
    static constexpr int K_LOAD_ROWS    = BT;
    static constexpr int K_LOAD_BF16    = K_LOAD_ROWS * kStateDim;
    static constexpr int K_STORAGE_FLT  = K_LOAD_BF16 / 2;
    static constexpr int M_TILES_H_PW   = (kStateDim / D::BT_SPLITS) / MMA_M;
    static constexpr int SNAP_K_ROWS    = MMA_M * M_TILES_H_PW;
    static constexpr int SNAP_FLT       = SNAP_K_ROWS * D::N_STRIP_PER_BLOCK;
    static constexpr int M_TILES_H_GLOB = kStateDim / MMA_M;
    static constexpr int N_SNAP_ITERS   = M_TILES_H_GLOB / M_TILES_H_PW;
    // One buffer per sit so all warps can scatter their owned h_frag once at
    // chunk start (Phase A) and then read from any sit's buffer in the unified
    // matmul1 / coop_write paths without per-sit re-scatter+sync.
    static constexpr int N_SNAP_BUF = N_SNAP_ITERS;
    static constexpr int SMEM_FLOATS =
        W_STORAGE_FLT + UVD_FLT + K_STORAGE_FLT + SNAP_FLT * N_SNAP_BUF + BT;
};

// Snapshot stores the FP32 state operand as [value, state] with a stride-32
// swizzle shared by the scatter, cooperative h_chunk write, and MM1 LDSM read.
struct SnapView {
    float* __restrict__ base;
    static constexpr int kStride = kStateDim / 4;
    static_assert(kStride == 32);

    __device__ __forceinline__ int swz_xor(int row) const { return (row & 7) << 2; }

    __device__ __forceinline__ float& at(int row, int col) const {
        return base[row * kStride + (col ^ swz_xor(row))];
    }

    __device__ __forceinline__ float4& vec4_at(int row, int col) const {
        return *reinterpret_cast<float4*>(&base[row * kStride + (col ^ swz_xor(row))]);
    }
};

// prepare_wy_wu stores W's private workspace in the native-BF16 fragment
// order {0,4,1,5,2,6,3,7} within each group of eight logical columns.
struct Bf16WView {
    __nv_bfloat16* __restrict__ base;

    __device__ __forceinline__ int swizzled_col(int row, int col) const {
        return (col & ~63) + ((((col & 63) >> 3) ^ (row & 7)) << 3) + (col & 7);
    }

    __device__ __forceinline__ __nv_bfloat16* ptr(int row, int col) const {
        return &base[row * kStateDim + swizzled_col(row, col)];
    }
};

template <int THREADS>
__device__ __forceinline__ void issue_load_w_bf16(Bf16WView view,
                                                  const __nv_bfloat16* gmem_base_row0,
                                                  int64_t gmem_row_stride, int tid) {
    constexpr int ELEMS_PER_COPY = 8;
    constexpr int COPIES_PER_ROW = kStateDim / ELEMS_PER_COPY;
    constexpr int COPIES         = BT * COPIES_PER_ROW;
#pragma unroll
    for (int copy = tid; copy < COPIES; copy += THREADS) {
        const int row = copy / COPIES_PER_ROW;
        const int col = (copy - row * COPIES_PER_ROW) * ELEMS_PER_COPY;
        cp_async<16>(view.ptr(row, col),
                     gmem_base_row0 + static_cast<int64_t>(row) * gmem_row_stride + col);
    }
}

// K is staged in native BF16 because it is already representable exactly as
// a TF32 operand. Within each group of eight token rows, physical rows are
// ordered {0,4,1,5,2,6,3,7}. ldmatrix.x2.trans then packs token columns
// {t,t+4} into the two halves of each register, matching the TF32 A fragment.
struct Bf16KView {
    __nv_bfloat16* __restrict__ base;

    __device__ __forceinline__ int physical_row(int logical_row) const {
        const int row8 = logical_row & 7;
        return (logical_row & ~7) + ((row8 & 3) << 1) + (row8 >> 2);
    }

    __device__ __forceinline__ int swizzled_col(int physical_row_, int col) const {
        return (col & ~63) + ((((col & 63) >> 3) ^ (physical_row_ & 7)) << 3) + (col & 7);
    }

    __device__ __forceinline__ __nv_bfloat16* logical_ptr(int row, int col) const {
        const int row_phys = physical_row(row);
        return &base[row_phys * kStateDim + swizzled_col(row_phys, col)];
    }

    __device__ __forceinline__ __nv_bfloat16* physical_ptr(int row, int col) const {
        return &base[row * kStateDim + swizzled_col(row, col)];
    }
};

template <int THREADS>
__device__ __forceinline__ void issue_load_k_bf16(Bf16KView view,
                                                  const __nv_bfloat16* gmem_base_row0,
                                                  int64_t gmem_row_stride, int tid) {
    constexpr int ELEMS_PER_COPY = 8;
    constexpr int COPIES_PER_ROW = kStateDim / ELEMS_PER_COPY;
    constexpr int COPIES         = BT * COPIES_PER_ROW;
#pragma unroll
    for (int copy = tid; copy < COPIES; copy += THREADS) {
        const int row = copy / COPIES_PER_ROW;
        const int col = (copy - row * COPIES_PER_ROW) * ELEMS_PER_COPY;
        cp_async<16>(view.logical_ptr(row, col),
                     gmem_base_row0 + static_cast<int64_t>(row) * gmem_row_stride + col);
    }
}

__device__ __forceinline__ void unpack_bf16x2_to_fp32_bits(unsigned packed, unsigned& low,
                                                           unsigned& high) {
    low  = packed << 16;
    high = packed & 0xffff0000U;
}

// The narrow geometry targets two 128-register CTAs per SM. The wide geometry
// uses one 512-thread CTA; both expose 16 resident warps without local spills.
template <int NStrip>
__launch_bounds__(kernel_dims<NStrip>::THREADS, kernel_dims<NStrip>::MIN_BLOCKS) __global__
    void state_passing_kernel(const __nv_bfloat16* __restrict__ W_in,
                              const __nv_bfloat16* __restrict__ U_in,
                              const __nv_bfloat16* __restrict__ k_in,
                              const float* __restrict__ g_cumsum, const float* state_in,
                              __nv_bfloat16* __restrict__ v_new,
                              __nv_bfloat16* __restrict__ h_chunk, float* state_out,
                              head_map qk_map, int chunks) {
    using D                         = kernel_dims<NStrip>;
    using L                         = smem_layout<NStrip>;
    constexpr int N_STRIP_PER_BLOCK = D::N_STRIP_PER_BLOCK;
    constexpr int D_STRIPS          = D::D_STRIPS;
    constexpr int BT_SPLITS         = D::BT_SPLITS;
    constexpr int THREADS_K         = D::THREADS;
    constexpr int W_STRIDE          = L::W_STRIDE;

    constexpr int BT_PER_WARP    = BT / BT_SPLITS;
    constexpr int S_PER_WARP     = kStateDim / BT_SPLITS;
    constexpr int M_TILES_MM1_PW = BT_PER_WARP / MMA_M;
    constexpr int M_TILES_H_PW   = S_PER_WARP / MMA_M;
    constexpr int M_TILES_H_GLOB = kStateDim / MMA_M;
    constexpr int K_TILES_MM2    = BT / MMA_K;

    static_assert(M_TILES_H_PW >= 1, "S_PER_WARP must yield >= 1 M-tile per warp");
    static_assert(M_TILES_MM1_PW >= 1, "BT_PER_WARP must yield >= 1 M-tile per warp");
    static_assert(M_TILES_H_GLOB == BT_SPLITS * M_TILES_H_PW,
                  "BT_SPLITS partition of state dim must be exact");
    static_assert(W_STRIDE >= 16, "SmemTile<W_STRIDE> requires stride >= 16");

    constexpr int SNAP_K_ROWS           = L::SNAP_K_ROWS;
    constexpr int N_SNAP_ITERS          = L::N_SNAP_ITERS;
    constexpr int K_TILES_PER_SNAP_ITER = SNAP_K_ROWS / MMA_K;

    static_assert(N_SNAP_ITERS == BT_SPLITS, "design assumes one snap iter per s_idx");

    constexpr int W_STORAGE_FLT = L::W_STORAGE_FLT;
    constexpr int UVD_FLT       = L::UVD_FLT;
    constexpr int K_STORAGE_FLT = L::K_STORAGE_FLT;
    constexpr int SNAP_FLT      = L::SNAP_FLT;
    constexpr int N_SNAP_BUF    = L::N_SNAP_BUF;

    // Smem partition. U and vd alias (`uvd_smem`) in disjoint phases. snap
    // is per-sit so the chunk-start unified scatter feeds every sit's MM1.
    extern __shared__ float smem[];
    auto* const W_smem     = reinterpret_cast<__nv_bfloat16*>(smem); // W_LOAD_BF16
    float* const uvd_smem  = smem + W_STORAGE_FLT;                   // UVD_FLT
    auto* const k_smem     = reinterpret_cast<__nv_bfloat16*>(uvd_smem + UVD_FLT);
    float* const snap_smem = uvd_smem + UVD_FLT + K_STORAGE_FLT; // SNAP_FLT * N_SNAP_BUF
    float* const g_smem    = snap_smem + SNAP_FLT * N_SNAP_BUF;  // BT

    Bf16WView W_view{W_smem};
    SmemTile<N_STRIP_PER_BLOCK> vd_view{uvd_smem};
    Bf16KView k_view{k_smem};
    SmemTile<N_STRIP_PER_BLOCK> U_view{uvd_smem};
    // One SnapView per sit so each owning warp scatters into a unique buffer
    // (Phase 2.a unified-scatter design). The array is sized on N_SNAP_BUF
    // (= N_SNAP_ITERS) rather than hard-coded.
    SnapView snap_views[N_SNAP_BUF];
#pragma unroll
    for (int b_ = 0; b_ < N_SNAP_BUF; ++b_) {
        snap_views[b_] = SnapView{snap_smem + b_ * SNAP_FLT};
    }

    // Block / lane indexing.
    //   grid.x = hd in [0, H_v*D_STRIPS).
    //   warp = dt_idx * BT_SPLITS + s_idx.
    const int tid    = static_cast<int>(threadIdx.x);
    const int lane   = tid & (kWarpSize - 1);
    const int warp   = tid / kWarpSize;
    const int lane_g = lane >> 2;
    const int lane_t = lane & 3;

    const int s_idx        = warp % BT_SPLITS;
    const int dt_idx       = warp / BT_SPLITS;
    const int warp_d_local = dt_idx * MMA_N;

    const int hd            = static_cast<int>(blockIdx.x);
    const int h_v           = hd / D_STRIPS;
    const int strip_idx     = hd - h_v * D_STRIPS;
    const int d_off         = strip_idx * N_STRIP_PER_BLOCK;
    const int warp_d_global = d_off + warp_d_local;
    const std::int64_t H_v  = qk_map.H_v;

    // === Phase 0: load state_in (AR-transposed) -> per-warp h_frag ===
    float h_frag[M_TILES_H_PW][4];
    {
        const int64_t st_base = static_cast<int64_t>(h_v) * kStateDim * kStateDim;
        const int row_off     = s_idx * S_PER_WARP;
#pragma unroll
        for (int m = 0; m < M_TILES_H_PW; ++m) {
            const int row_g0 = row_off + m * MMA_M + lane_g;
            const int row_g1 = row_g0 + 8;
            const int col_d0 = warp_d_global + 2 * lane_t;
            const int col_d1 = col_d0 + 1;
            h_frag[m][0] =
                load_ldg<float>(state_in + st_base + (int64_t)col_d0 * kStateDim + row_g0);
            h_frag[m][1] =
                load_ldg<float>(state_in + st_base + (int64_t)col_d1 * kStateDim + row_g0);
            h_frag[m][2] =
                load_ldg<float>(state_in + st_base + (int64_t)col_d0 * kStateDim + row_g1);
            h_frag[m][3] =
                load_ldg<float>(state_in + st_base + (int64_t)col_d1 * kStateDim + row_g1);
        }
    }

    // Chunk 0 is staged before the loop. Later W/K async groups and converted
    // U are issued together at the preceding chunk's Phase-E boundary. All
    // global bases advance additively across the sequential chunk loop.
    const int64_t W_stride        = H_v * kStateDim;
    const int64_t k_stride        = static_cast<int64_t>(qk_map.H_qk) * kStateDim;
    const int64_t W_chunk_stride  = (int64_t)BT * W_stride;
    const int64_t k_chunk_stride  = (int64_t)BT * k_stride;
    const int64_t hc_chunk_stride = H_v * kStateDim * kStateDim;
    const int64_t vn_stride       = W_stride;
    const int64_t vn_chunk_stride = (int64_t)BT * vn_stride;
    const int64_t g_chunk_step    = (int64_t)BT * H_v;

    const int64_t W_block_base  = static_cast<int64_t>(h_v) * kStateDim;
    const int64_t k_block_base  = static_cast<int64_t>(qk_map.qk_head(h_v)) * kStateDim;
    const int64_t hc_block_base = static_cast<int64_t>(h_v) * kStateDim * kStateDim;
    const int64_t vn_block_base = static_cast<int64_t>(h_v) * kStateDim;
    const int64_t g_block_base  = h_v;
    const int64_t g_thread_base = g_block_base + (int64_t)tid * H_v;

    // W/K stay native BF16 in shared memory. W is committed first for MM1;
    // K is the later group and remains in flight until MM2. U is expanded
    // synchronously while both async groups make progress.
    {
        issue_load_w_bf16<THREADS_K>(W_view, W_in + W_block_base, W_stride, tid);
        cp_commit();
        issue_load_k_bf16<THREADS_K>(k_view, k_in + k_block_base, k_stride, tid);
        cp_commit();
        issue_load_bf16_to_float_vec4<BT, N_STRIP_PER_BLOCK, THREADS_K>(
            U_view, U_in + W_block_base + d_off, W_stride, tid);
    }

    // === Main chunk loop ===
    int64_t W_base      = W_block_base;
    int64_t k_base      = k_block_base;
    int64_t hc_base     = hc_block_base;
    int64_t vn_base     = vn_block_base;
    int64_t g_cs_offset = 0;
    for (int chunk = 0; chunk < chunks; ++chunk) {
        const int64_t W_base_next = W_base + W_chunk_stride;
        const int64_t k_base_next = k_base + k_chunk_stride;

        // === Phase A: drain W + scatter h_frag to snap ===
        //
        // Phase 2.a unified-scatter: every warp scatters its owned h_frag into
        // snap_views[s_idx] BEFORE the Phase A drain sync. The single sync
        // below covers W/U/g_smem AND snap visibility, so Phase B no longer
        // needs per-sit scatter+sync.
        if (tid < BT) { g_smem[tid] = g_cumsum[g_thread_base + g_cs_offset]; }

        {
            SnapView snap = snap_views[s_idx];
#pragma unroll
            for (int m = 0; m < M_TILES_H_PW; ++m) {
                const int k_g0    = m * MMA_M + lane_g;
                const int k_g1    = k_g0 + 8;
                const int d0      = warp_d_local + 2 * lane_t;
                const int d1      = d0 + 1;
                snap.at(d0, k_g0) = h_frag[m][0];
                snap.at(d1, k_g0) = h_frag[m][1];
                snap.at(d0, k_g1) = h_frag[m][2];
                snap.at(d1, k_g1) = h_frag[m][3];
            }
        }

        cp_wait<1>();    // drain W; leave K in flight
        __syncthreads(); // gates W/U/g_smem STS and scatter visibility

        // === Phase B: per-sit coop_write + matmul1 (no per-sit scatter sync) ===
        //
        // The unified scatter in Phase A populated snap_views[s_idx] for all
        // s_idx in one shot, so each sit just consumes snap[sit] without
        // re-scattering or syncing. We retain the per-sit interleave between
        // coop_write (LDS+STG to gmem) and mma (LDS) because that was the
        // crucial latency-hiding pattern in the per-sit baseline -- moving
        // all coop_writes ahead of all mmas costs ~15 us on L=4096.
        float vnew_frag[M_TILES_MM1_PW][4] = {};

#pragma unroll
        for (int sit = 0; sit < N_SNAP_ITERS; ++sit) {
            const int k_row_off = sit * SNAP_K_ROWS;
            SnapView snap       = snap_views[sit];

            // Coop float4 gmem write of h_chunk for this snap block.
            // No sync above: snap was populated by the Phase A
            // unified-scatter, and snap[sit] is read-only from now on.
            {
                constexpr int K_VEC_PER_D = SNAP_K_ROWS / 4;
                constexpr int N_VEC_SNAP  = SNAP_FLT / 4;
#pragma unroll
                for (int v = tid; v < N_VEC_SNAP; v += THREADS_K) {
                    const int d_local  = v / K_VEC_PER_D;
                    const int kvec     = v - d_local * K_VEC_PER_D;
                    const int k_off    = kvec * 4;
                    float4 val         = snap.vec4_at(d_local, k_off);
                    const int d_global = d_off + d_local;
                    __nv_bfloat16* out =
                        &h_chunk[hc_base + (int64_t)d_global * kStateDim + k_row_off + k_off];
                    store_vec(out, __floats2bfloat162_rn(val.x, val.y));
                    store_vec(out + 2, __floats2bfloat162_rn(val.z, val.w));
                }
            }

            // W's producer-interleaved BF16 layout lets ldmatrix.x2 form the
            // four exact FP32/TF32 A operands. The FP32 state snapshot uses
            // ldmatrix.x2 as a b16 transport for the two B operand bit patterns.
            const int lane_in_8   = lane & 7;
            const int matrix_half = (lane >> 3) & 1;
#pragma unroll
            for (int kt = 0; kt < K_TILES_PER_SNAP_ITER; ++kt) {
                const int W_k_local = k_row_off + kt * MMA_K;
                const int snap_k    = kt * MMA_K;

                const int b_row       = warp_d_local + lane_in_8;
                const int b_col       = snap_k + matrix_half * 4;
                const unsigned b_addr = smem_addr(&snap.at(b_row, b_col));
                unsigned ub0, ub1;
                ldmatrix_x2(ub0, ub1, b_addr);

#pragma unroll
                for (int m_mm1 = 0; m_mm1 < M_TILES_MM1_PW; ++m_mm1) {
                    const int row_base    = s_idx * BT_PER_WARP + m_mm1 * MMA_M;
                    const int a_row       = row_base + matrix_half * 8 + lane_in_8;
                    const unsigned a_addr = smem_addr(W_view.ptr(a_row, W_k_local));
                    unsigned packed0, packed1;
                    ldmatrix_x2(packed0, packed1, a_addr);
                    unsigned ua0, ua1, ua2, ua3;
                    unpack_bf16x2_to_fp32_bits(packed0, ua0, ua2);
                    unpack_bf16x2_to_fp32_bits(packed1, ua1, ua3);

                    mma_tf32_bits(vnew_frag[m_mm1][0], vnew_frag[m_mm1][1], vnew_frag[m_mm1][2],
                                  vnew_frag[m_mm1][3], ua0, ua1, ua2, ua3, ub0, ub1);
                }
            }
        }

        // === Phase C: subtract converted FP32 U ===
#pragma unroll
        for (int m_mm1 = 0; m_mm1 < M_TILES_MM1_PW; ++m_mm1) {
            const int row_g0    = s_idx * BT_PER_WARP + m_mm1 * MMA_M + lane_g;
            const int row_g1    = row_g0 + 8;
            const int col_d0    = warp_d_local + 2 * lane_t;
            const float2 u_top  = load_vec<float2>(&U_view.at(row_g0, col_d0));
            const float2 u_bot  = load_vec<float2>(&U_view.at(row_g1, col_d0));
            vnew_frag[m_mm1][0] = u_top.x - vnew_frag[m_mm1][0];
            vnew_frag[m_mm1][1] = u_top.y - vnew_frag[m_mm1][1];
            vnew_frag[m_mm1][2] = u_bot.x - vnew_frag[m_mm1][2];
            vnew_frag[m_mm1][3] = u_bot.y - vnew_frag[m_mm1][3];
        }

        // === Phase D: STG vnew (UNDECAYED), STS v_decay -> vd_view, scale h_frag ===
        const float g_C = g_smem[BT - 1];
        float gamma_C   = 0.0f;
        if (lane == 0) { gamma_C = exp2_approx(g_C * kLog2E); }
        gamma_C = __shfl_sync(0xffffffffU, gamma_C, 0);

#pragma unroll
        for (int m_mm1 = 0; m_mm1 < M_TILES_MM1_PW; ++m_mm1) {
            const int row_g0 = s_idx * BT_PER_WARP + m_mm1 * MMA_M + lane_g;
            const int row_g1 = row_g0 + 8;
            const int col_d0 = warp_d_global + 2 * lane_t;

            float dec_top = 0.0f;
            float dec_bot = 0.0f;
            if (lane_t == 0) {
                dec_top = exp2_approx((g_C - g_smem[row_g0]) * kLog2E);
                dec_bot = exp2_approx((g_C - g_smem[row_g1]) * kLog2E);
            }
            const int decay_src_lane = lane & ~3;
            dec_top                  = __shfl_sync(0xffffffffU, dec_top, decay_src_lane);
            dec_bot                  = __shfl_sync(0xffffffffU, dec_bot, decay_src_lane);

            const float v0 = vnew_frag[m_mm1][0];
            const float v1 = vnew_frag[m_mm1][1];
            const float v2 = vnew_frag[m_mm1][2];
            const float v3 = vnew_frag[m_mm1][3];

            const __nv_bfloat162 out0 = __floats2bfloat162_rn(v0, v1);
            const __nv_bfloat162 out1 = __floats2bfloat162_rn(v2, v3);
            store_vec(&v_new[vn_base + (int64_t)row_g0 * vn_stride + col_d0], out0);
            store_vec(&v_new[vn_base + (int64_t)row_g1 * vn_stride + col_d0], out1);

            const int row_g0_loc = s_idx * BT_PER_WARP + m_mm1 * MMA_M + lane_g;
            const int row_g1_loc = row_g0_loc + 8;
            const int col_d0_loc = warp_d_local + 2 * lane_t;
            store_vec(&vd_view.at(row_g0_loc, col_d0_loc), make_float2(v0 * dec_top, v1 * dec_top));
            store_vec(&vd_view.at(row_g1_loc, col_d0_loc), make_float2(v2 * dec_bot, v3 * dec_bot));
        }

#pragma unroll
        for (int m = 0; m < M_TILES_H_PW; ++m) {
#pragma unroll
            for (int e = 0; e < 4; ++e) { h_frag[m][e] *= gamma_C; }
        }

        // Drain native-BF16 K, then make K and Phase-D v_decay visible to MM2.
        cp_wait<0>();
        __syncthreads();

        // === Phase E: matmul2 over the full chunk of k ===
#pragma unroll
        for (int kt = 0; kt < K_TILES_MM2; ++kt) {
            const int k_off_local = kt * MMA_K;

            const int row_t0 = k_off_local + lane_t;
            const int row_t1 = row_t0 + 4;
            const int col_g  = warp_d_local + lane_g;
            const float b0   = vd_view.at(row_t0, col_g);
            const float b1   = vd_view.at(row_t1, col_g);

#pragma unroll
            for (int m = 0; m < M_TILES_H_PW; ++m) {
                const int col_a_base = s_idx * S_PER_WARP + m * MMA_M;
                const int a_row      = k_off_local + (lane & 7);
                const int a_col      = col_a_base + ((lane >> 3) & 1) * 8;
                unsigned packed0, packed1;
                ldmatrix_x2_t(packed0, packed1, smem_addr(k_view.physical_ptr(a_row, a_col)));
                unsigned ua0, ua1, ua2, ua3;
                unpack_bf16x2_to_fp32_bits(packed0, ua0, ua2);
                unpack_bf16x2_to_fp32_bits(packed1, ua1, ua3);

                mma_tf32_bits(h_frag[m][0], h_frag[m][1], h_frag[m][2], h_frag[m][3], ua0, ua1, ua2,
                              ua3, __float_as_uint(b0), __float_as_uint(b1));
            }
        }

        __syncthreads(); // gates MM2 before the next chunk overwrites W/K/U

        // The next chunk repeats the W-then-K async group order used by the
        // prologue. Its Phase A drains W only; Phase E drains K.
        if (chunk + 1 < chunks) {
            issue_load_w_bf16<THREADS_K>(W_view, W_in + W_base_next, W_stride, tid);
            cp_commit();
            issue_load_k_bf16<THREADS_K>(k_view, k_in + k_base_next, k_stride, tid);
            cp_commit();
            issue_load_bf16_to_float_vec4<BT, N_STRIP_PER_BLOCK, THREADS_K>(
                U_view, U_in + W_base_next + d_off, W_stride, tid);
        }

        // Advance loop-carried bases for the next chunk.
        W_base = W_base_next;
        k_base = k_base_next;
        hc_base += hc_chunk_stride;
        vn_base += vn_chunk_stride;
        g_cs_offset += g_chunk_step;
    }

    // === Phase Z: store h_frag -> state_out (AR-transposed) ===
    const int64_t st_base = static_cast<int64_t>(h_v) * kStateDim * kStateDim;

#pragma unroll
    for (int m = 0; m < M_TILES_H_PW; ++m) {
        const int k_g0 = s_idx * S_PER_WARP + m * MMA_M + lane_g;
        const int k_g1 = k_g0 + 8;
        const int d0   = warp_d_global + 2 * lane_t;
        const int d1   = d0 + 1;
        state_out[st_base + (int64_t)d0 * kStateDim + k_g0] = h_frag[m][0];
        state_out[st_base + (int64_t)d1 * kStateDim + k_g0] = h_frag[m][1];
        state_out[st_base + (int64_t)d0 * kStateDim + k_g1] = h_frag[m][2];
        state_out[st_base + (int64_t)d1 * kStateDim + k_g1] = h_frag[m][3];
    }
}

} // namespace ninfer::ops::detail::gated_delta_net::chunked::state_passing
#pragma once

#include "ninfer/types.h"

#include <chrono>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <iosfwd>
#include <string>
#include <string_view>

namespace ninfer::product {

enum class LoadProgressOutputMode : std::uint8_t {
    Interactive,
    Log,
};

struct LoadProgressRendererOptions {
    LoadProgressOutputMode mode                    = LoadProgressOutputMode::Log;
    std::chrono::milliseconds min_refresh_interval = std::chrono::seconds(10);
    std::function<std::string()> line_prefix;
};

[[nodiscard]] LoadProgressRendererOptions stderr_load_progress_options() noexcept;

// Owns product-side terminal rendering for Engine load progress. The renderer must
// outlive the Engine whose LoadProgress callback references it.
class LoadProgressRenderer {
public:
    LoadProgressRenderer(std::ostream& output, LoadProgressRendererOptions options);
    ~LoadProgressRenderer();

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

    [[nodiscard]] LoadProgress callback();

    // Ends an open interactive line. This is idempotent and is also called by
    // the destructor so an Engine construction failure cannot corrupt later output.
    void finish() noexcept;

private:
    using Clock = std::chrono::steady_clock;

    void update(std::string_view phase, std::uint64_t done, std::uint64_t total);
    void render(std::uint64_t done, std::uint64_t total, Clock::time_point now, bool final);

    std::ostream* output_ = nullptr;
    LoadProgressRendererOptions options_;
    bool has_phase_ = false;
    bool line_open_ = false;
    std::string phase_;
    std::uint64_t last_done_    = 0;
    std::uint64_t last_total_   = 0;
    std::size_t terminal_width_ = 0;
    Clock::time_point phase_started_;
    Clock::time_point last_rendered_;
};

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

#include "ops/linear/fp8/fp8_config.h"

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

namespace ninfer::ops::detail {
namespace {

enum class Fp8GdnInputRoute : std::uint8_t {
    A16,
    A8,
};

Fp8GdnInputRoute resolve_route(LinearPolicy policy, std::int32_t tokens) {
    if (tokens <= 0) { throw std::invalid_argument("fp8 gdn_input_proj: T must be positive"); }
    if (policy == LinearPolicy::A16Only) { return Fp8GdnInputRoute::A16; }
    if (policy != LinearPolicy::AllowA8) {
        throw std::invalid_argument("fp8 gdn_input_proj: unsupported policy");
    }
    return tokens >= 8 ? Fp8GdnInputRoute::A8 : Fp8GdnInputRoute::A16;
}

} // namespace

std::size_t fp8_gdn_input_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("fp8 gdn_input_proj workspace: invalid token interval");
    }
    (void)resolve_route(policy, min_tokens);
    return resolve_route(policy, max_tokens) == Fp8GdnInputRoute::A8
               ? fp8_a8_workspace_capacity_bytes(max_tokens, Fp8GdnInputGeometry::kInputRows)
               : 0;
}

void fp8_gdn_input_a16_dispatch(const Tensor& x, const Weight& weight, Tensor& qkv, Tensor& z,
                                cudaStream_t stream) {
    constexpr std::int32_t kQkvRows = 10240;
    constexpr std::int32_t kZRows   = 6144;
    constexpr std::int32_t kChunk   = kFp8LinearSmallTMax<Fp8GdnInputGeometry>;
    for (std::int32_t token_begin = 0; token_begin < x.ne[1]; token_begin += kChunk) {
        const std::int32_t active = std::min(kChunk, x.ne[1] - token_begin);
        auto* input               = static_cast<std::uint8_t*>(x.data) +
                      static_cast<std::int64_t>(token_begin) * weight.k * sizeof(std::uint16_t);
        auto* qkv_output =
            static_cast<std::uint8_t*>(qkv.data) +
            static_cast<std::int64_t>(token_begin) * kQkvRows * sizeof(std::uint16_t);
        auto* z_output = static_cast<std::uint8_t*>(z.data) +
                         static_cast<std::int64_t>(token_begin) * kZRows * sizeof(std::uint16_t);
        Tensor input_chunk(input, DType::BF16, {weight.k, active});
        Tensor qkv_chunk(qkv_output, DType::BF16, {kQkvRows, active});
        Tensor z_chunk(z_output, DType::BF16, {kZRows, active});
        if (active == 1) {
            fp8_gdn_input_decode_launch(input_chunk, weight, qkv_chunk, z_chunk, stream);
        } else {
            fp8_gdn_input_small_t_launch(input_chunk, weight, qkv_chunk, z_chunk, stream);
        }
    }
}

void fp8_gdn_input_a8_dispatch(const Tensor& x, const Weight& weight, Tensor& qkv, Tensor& z,
                               WorkspaceArena& workspace, cudaStream_t stream) {
    auto scope                   = workspace.scope();
    const Fp8A8Workspace scratch = allocate_fp8_a8_workspace(workspace, x.ne[1], weight.k);
    fp8_gdn_input_a8_launch(x, weight, qkv, z, scratch, stream);
}

void fp8_gdn_input_dispatch(const Tensor& x, const Weight& weight, Tensor& qkv, Tensor& z,
                            LinearPolicy policy, WorkspaceArena* workspace, cudaStream_t stream) {
    if (resolve_route(policy, x.ne[1]) == Fp8GdnInputRoute::A16) {
        fp8_gdn_input_a16_dispatch(x, weight, qkv, z, stream);
        return;
    }
    if (workspace == nullptr) {
        throw std::invalid_argument("fp8 A8 gdn_input_proj requires caller workspace");
    }
    fp8_gdn_input_a8_dispatch(x, weight, qkv, z, *workspace, stream);
}

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

#include <stdexcept>
#include <string>

namespace ninfer::ops {
namespace {

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

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

} // namespace

void mtp_pack_fc_input(const Tensor& embedding_norm, const Tensor& hidden_norm, Tensor& out,
                       cudaStream_t stream) {
    constexpr const char* op = "mtp_pack_fc_input";
    require_bf16_contiguous_nonnull(embedding_norm, op, "embedding_norm");
    require_bf16_contiguous_nonnull(hidden_norm, op, "hidden_norm");
    require_bf16_contiguous_nonnull(out, op, "out");
    const std::int32_t rows   = embedding_norm.ne[0];
    const std::int32_t tokens = embedding_norm.ne[1];
    if (rows <= 0) { throw std::invalid_argument("mtp_pack_fc_input: D must be positive"); }
    if (tokens <= 0) { throw std::invalid_argument("mtp_pack_fc_input: T must be positive"); }
    require_shape(embedding_norm, rows, tokens, op, "embedding_norm");
    require_shape(hidden_norm, rows, tokens, op, "hidden_norm");
    require_shape(out, 2 * rows, tokens, op, "out");

    detail::mtp_pack_fc_input_launch(embedding_norm, hidden_norm, out, stream);
}

void mtp_split_attn_in(const Tensor& attn_in, Tensor& q, Tensor& k, Tensor& gate, Tensor& v,
                       cudaStream_t stream) {
    constexpr const char* op = "mtp_split_attn_in";
    require_bf16_contiguous_nonnull(attn_in, op, "attn_in");
    require_bf16_contiguous_nonnull(q, op, "q");
    require_bf16_contiguous_nonnull(k, op, "k");
    require_bf16_contiguous_nonnull(gate, op, "gate");
    require_bf16_contiguous_nonnull(v, op, "v");
    const std::int32_t tokens = attn_in.ne[1];
    if (tokens <= 0) { throw std::invalid_argument("mtp_split_attn_in: T must be positive"); }
    require_shape(attn_in, 14336, tokens, op, "attn_in");
    if (q.ne[0] != 256 || q.ne[1] != 24 || q.ne[2] != tokens || q.ne[3] != 1) {
        throw std::invalid_argument("mtp_split_attn_in: invalid shape for q");
    }
    if (k.ne[0] != 256 || k.ne[1] != 4 || k.ne[2] != tokens || k.ne[3] != 1) {
        throw std::invalid_argument("mtp_split_attn_in: invalid shape for k");
    }
    if (gate.ne[0] != 256 || gate.ne[1] != 24 || gate.ne[2] != tokens || gate.ne[3] != 1) {
        throw std::invalid_argument("mtp_split_attn_in: invalid shape for gate");
    }
    if (v.ne[0] != 256 || v.ne[1] != 4 || v.ne[2] != tokens || v.ne[3] != 1) {
        throw std::invalid_argument("mtp_split_attn_in: invalid shape for v");
    }

    detail::mtp_split_attn_in_launch(attn_in, q, k, gate, v, stream);
}

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

#include "core/layout.h"

#include <stdexcept>

namespace ninfer::ops::detail {

bool sparse_moe_uses_small_t(std::int32_t tokens) noexcept {
    return tokens >= kSparseMoeSmallTMin && tokens <= kSparseMoeSmallTMax;
}

std::size_t sparse_moe_small_t_workspace_bytes(std::int32_t tokens) {
    if (!sparse_moe_uses_small_t(tokens)) {
        throw std::invalid_argument("sparse_moe small-T: tokens must be in [2,46]");
    }
    WorkspaceLayoutBuilder layout;
    (void)allocate_sparse_moe_small_t_workspace(layout, tokens);
    return layout.peak_bytes(1);
}

SparseMoeSmallTPlan resolve_sparse_moe_small_t_plan(std::int32_t tokens, QType routed_gate_up,
                                                    QType routed_down) {
    if (!sparse_moe_uses_small_t(tokens)) {
        throw std::invalid_argument("sparse_moe small-T: unsupported token count");
    }
    const bool main_profile =
        routed_gate_up == QType::Q4G64_F16S &&
        (routed_down == QType::Q5G64_F16S || routed_down == QType::Q6G64_F16S);
    const bool mtp_profile =
        routed_gate_up == QType::W8G32_F16S && routed_down == QType::W8G32_F16S;
    if (!main_profile && !mtp_profile) {
        throw std::invalid_argument("sparse_moe small-T: unsupported routed codec profile");
    }

    SparseMoeSmallTPlan plan{tokens, sparse_moe_small_t_workspace_bytes(tokens)};
    if (mtp_profile) {
        plan.d3_schedule =
            tokens <= 5 ? SparseMoeSmallTD3Schedule::Paths1 : SparseMoeSmallTD3Schedule::Paths9;
        plan.d4_schedule =
            tokens <= 8 ? SparseMoeSmallTD4Schedule::Rows1 : SparseMoeSmallTD4Schedule::Rows4;
        return plan;
    }

    plan.d3_schedule = SparseMoeSmallTD3Schedule::Paths3;
    if (routed_down == QType::Q5G64_F16S) {
        plan.d4_schedule = tokens <= 2   ? SparseMoeSmallTD4Schedule::Rows1
                           : tokens <= 5 ? SparseMoeSmallTD4Schedule::Rows2
                                         : SparseMoeSmallTD4Schedule::Rows4;
    } else {
        plan.d4_schedule = tokens <= 2    ? SparseMoeSmallTD4Schedule::Rows1
                           : tokens <= 11 ? SparseMoeSmallTD4Schedule::Rows2
                                          : SparseMoeSmallTD4Schedule::Rows4;
    }
    return plan;
}

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

#include "core/tensor.h"

#include <cuda_runtime.h>

namespace ninfer::ops {

/**
 * Depthwise causal width-4 convolution followed by SiLU. Let u[c,-3..-1] be the three values in
 * the input state and u[c,t]=x[c,t] for t>=0. Then
 *
 *   ideal[c,t] = SiLU(sum_{j=0..3} weight[c,j] * u[c,t-3+j]).
 *
 * `x` and `out` are contiguous BF16 [C,T], `weight` is contiguous BF16 [C,4], and a state is
 * contiguous BF16 [C,3] ordered oldest to newest. The oracle evaluates `ideal` naively in FP64 from
 * the represented inputs. The BF16 output is promoted and compared directly with that result;
 * output storage rounding belongs to the Op's numerical criterion, not the oracle. Kernel
 * accumulator and staging precision are implementation choices. Input, weight, output, and state
 * storage do not overlap except for the explicitly allowed exact alias between state input and
 * state output. No caller workspace is used. T may be any positive value.
 */

// Reads conv_state as the initial window and replaces it with the final three values after x.
void causal_conv1d_silu(const Tensor& x, const Tensor& weight, Tensor& conv_state, Tensor& out,
                        cudaStream_t stream);

// Distinct-state form. conv_state_in and conv_state_out may be disjoint or exactly the same
// storage; conv_state_out receives the trailing width-3 window of concat(conv_state_in,x).
void causal_conv1d_silu(const Tensor& x, const Tensor& weight, const Tensor& conv_state_in,
                        Tensor& conv_state_out, Tensor& out, cudaStream_t stream);

// Split-output form. `ideal` is unchanged; its rows are written to three destinations instead of
// one, partitioned in row order as out0, out1, out2. Every destination is contiguous BF16 with the
// column count of x, and every operand is four-byte aligned. The supported row profiles are
// (2048, 2048, 4096) over C = 8192 and (2048, 2048, 6144) over C = 10240; any other profile is
// rejected. conv_state_in and conv_state_out follow the family rule above: disjoint or exactly the
// same storage. No destination may overlap another, x, weight, or either state.
void causal_conv1d_silu_split(const Tensor& x, const Tensor& weight, const Tensor& conv_state_in,
                              Tensor& conv_state_out, Tensor& out0, Tensor& out1, Tensor& out2,
                              cudaStream_t stream);

/**
 * Snapshot form for B independent sequences. `x` and `out` are contiguous BF16 [C,W,B],
 * `conv_states` is contiguous BF16 [C,3,Slots], and `initial_state_slots` and
 * `snapshot_base_slots` are contiguous I32 [B]. `valid_columns` is either contiguous I32 [B],
 * with every value in [1,W], or an empty Tensor meaning every row has W valid columns. B=1
 * accepts every positive W; B=2..8 accepts W=1..16.
 *
 * Row b starts from the window selected by initial_state_slots[b]. After valid column j, the new
 * window is written to snapshot_base_slots[b]+j. Invalid-tail output columns are exact BF16 zero
 * and do not mutate state. The caller reserves the complete [base,base+W) interval for every row;
 * all row reservations are disjoint and no row overwrites another row's initial slot. A row's own
 * initial slot may lie in its destination reservation. `conv_states` is the only persistent state
 * mutated and must not overlap x, weight, metadata, or out.
 */
void causal_conv1d_silu_snapshot(const Tensor& x, const Tensor& weight, Tensor& conv_states,
                                 const Tensor& valid_columns, const Tensor& initial_state_slots,
                                 const Tensor& snapshot_base_slots, Tensor& out,
                                 cudaStream_t stream);

} // namespace ninfer::ops
#pragma once

#include "ops/common/memory.cuh"

#include <cuda_bf16.h>

#include <cstdint>

namespace ninfer::ops::detail {

struct Fp8GdnInputOutput {
    static constexpr std::int32_t kQueryRows = 2048;
    static constexpr std::int32_t kKeyRows   = 2048;
    static constexpr std::int32_t kValueRows = 6144;
    static constexpr std::int32_t kQkvRows   = kQueryRows + kKeyRows + kValueRows;
    static constexpr std::int32_t kZRows     = 6144;
    static constexpr std::int32_t kRows      = kQkvRows + kZRows;

    __nv_bfloat16* qkv;
    __nv_bfloat16* z;

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

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

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

static_assert(Fp8GdnInputOutput::kRows == 16384);
static_assert((Fp8GdnInputOutput::kQkvRows % 128) == 0);
static_assert((Fp8GdnInputOutput::kZRows % 128) == 0);

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

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

#include <cuda_bf16.h>

#include <cstdint>
#include <stdexcept>

namespace ninfer::ops::detail {
namespace {

template <class Geometry, bool FullTokens>
void launch_mma(const Weight& weight, Tensor& residual, Fp8A8Workspace workspace,
                std::int32_t tokens, cudaStream_t stream) {
    using Schedule          = typename Fp8LinearA8ProductionSchedule<Geometry>::Type;
    constexpr int kRowTiles = Geometry::kOutputRows / Schedule::kBlockRows;
    const int token_tiles   = (tokens + Schedule::kBlockTokens - 1) / Schedule::kBlockTokens;
    const int blocks        = kRowTiles * token_tiles;
    auto* output            = static_cast<__nv_bfloat16*>(residual.data);
    const Fp8AddResidualEpilogue epilogue{output, Geometry::kOutputRows};
    const Fp8ContiguousOutput destination{output, Geometry::kOutputRows};

    if constexpr (Schedule::kSharedBytes > 48 * 1024) {
        static const cudaError_t attribute = cudaFuncSetAttribute(
            fp8_mma_kernel<Geometry, Schedule, FullTokens, Fp8AddResidualEpilogue,
                           Fp8ContiguousOutput>,
            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, epilogue, destination);
    CUDA_CHECK(cudaGetLastError());
}

template <class Geometry>
void launch_problem(const Weight& weight, Tensor& residual, Fp8A8Workspace workspace,
                    std::int32_t tokens, cudaStream_t stream) {
    using Schedule = typename Fp8LinearA8ProductionSchedule<Geometry>::Type;
    if ((tokens % Schedule::kBlockTokens) == 0) {
        launch_mma<Geometry, true>(weight, residual, workspace, tokens, stream);
    } else {
        launch_mma<Geometry, false>(weight, residual, workspace, tokens, stream);
    }
}

} // namespace

void fp8_linear_add_a8_launch(const Tensor& x, const Weight& weight, Tensor& residual,
                              WorkspaceArena& workspace, cudaStream_t stream) {
    auto scope                   = workspace.scope();
    const Fp8A8Workspace scratch = allocate_fp8_a8_workspace(workspace, x.ne[1], weight.k);
    launch_fp8_a8_quantize(x, weight, scratch, stream);
    switch (resolve_fp8_problem(weight.n, weight.k)) {
    case Fp8Problem::Residual6144:
        launch_problem<Fp8Residual6144Geometry>(weight, residual, scratch, x.ne[1], stream);
        return;
    case Fp8Problem::Residual17408:
        launch_problem<Fp8Residual17408Geometry>(weight, residual, scratch, x.ne[1], stream);
        return;
    case Fp8Problem::AttnInput:
    case Fp8Problem::GdnInput:
    case Fp8Problem::MlpGateUp:
    case Fp8Problem::Vocabulary:
        break;
    }
    throw std::invalid_argument("fp8 linear_add: unsupported problem");
}

} // namespace ninfer::ops::detail
#include "ops/linear_add/bf16/bf16_linear_add_plan.h"

#include <stdexcept>

namespace ninfer::ops::detail {

bool bf16_linear_add_admits(std::int32_t output_rows, std::int32_t input_rows,
                            std::int32_t tokens) noexcept {
    return output_rows == 5120 && input_rows == 6144 && tokens > 0;
}

Bf16LinearAddScheduleId bf16_linear_add_select(std::int32_t output_rows, std::int32_t input_rows,
                                               std::int32_t tokens) {
    if (!bf16_linear_add_admits(output_rows, input_rows, tokens)) {
        throw std::invalid_argument("bf16 linear_add: unsupported exact problem");
    }
    if (tokens == 1) { return Bf16LinearAddScheduleId::Decode; }
    if (tokens <= kBf16LinearAddSmallTDispatchEnd) { return Bf16LinearAddScheduleId::SmallT; }
    if (tokens <= kBf16LinearAddAggregateMmaEnd) { return Bf16LinearAddScheduleId::AggregateMma; }
    return Bf16LinearAddScheduleId::Mma;
}

const char* bf16_linear_add_schedule_name(Bf16LinearAddScheduleId schedule) noexcept {
    switch (schedule) {
    case Bf16LinearAddScheduleId::Decode:
        return "linear_add.bf16.decode.residual";
    case Bf16LinearAddScheduleId::SmallT:
        return "linear_add.bf16.small_t.residual";
    case Bf16LinearAddScheduleId::AggregateMma:
        return "linear_add.bf16.aggregate_mma.residual";
    case Bf16LinearAddScheduleId::Mma:
        return "linear_add.bf16.mma.residual";
    }
    return "linear_add.bf16.unknown";
}

void bf16_linear_add_dispatch(const Tensor& x, const Weight& weight, Tensor& residual,
                              cudaStream_t stream) {
    switch (bf16_linear_add_select(weight.n, weight.k, x.ne[1])) {
    case Bf16LinearAddScheduleId::Decode:
        bf16_linear_add_decode_launch(x, weight, residual, stream);
        return;
    case Bf16LinearAddScheduleId::SmallT:
        bf16_linear_add_small_t_launch(x, weight, residual, stream);
        return;
    case Bf16LinearAddScheduleId::AggregateMma:
        bf16_linear_add_aggregate_mma_launch(x, weight, residual, stream);
        return;
    case Bf16LinearAddScheduleId::Mma:
        bf16_linear_add_mma_launch(x, weight, residual, stream);
        return;
    }
    throw std::logic_error("bf16 linear_add: unknown schedule");
}

} // namespace ninfer::ops::detail
#include "product/media_acquire/acquire.h"

#include <curl/curl.h>

#include <arpa/inet.h>
#include <netdb.h>
#include <sys/socket.h>

#include <algorithm>
#include <array>
#include <cstddef>
#include <chrono>
#include <cstdint>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <limits>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <string>
#include <string_view>
#include <vector>

namespace ninfer::product::media_acquire {
namespace {

using Clock = std::chrono::steady_clock;

void check_control(const Policy& policy) {
    if (policy.is_cancelled && policy.is_cancelled()) {
        throw Error(ErrorKind::Cancelled, "media acquisition was cancelled");
    }
    if (policy.deadline != Clock::time_point{} && Clock::now() >= policy.deadline) {
        throw Error(ErrorKind::DeadlineExceeded, "media acquisition exceeded request deadline");
    }
}

long bounded_timeout_ms(const Policy& policy, int configured) {
    if (policy.deadline == Clock::time_point{}) { return configured; }
    check_control(policy);
    const auto remaining =
        std::chrono::duration_cast<std::chrono::milliseconds>(policy.deadline - Clock::now())
            .count();
    return std::min<long>(configured, std::max<std::int64_t>(1, remaining));
}

std::vector<std::uint8_t> decode_base64(std::string_view text) {
    static constexpr std::array<std::int8_t, 256> table = [] {
        std::array<std::int8_t, 256> out{};
        out.fill(-1);
        for (int i = 0; i < 26; ++i) {
            out[static_cast<std::size_t>('A' + i)] = static_cast<std::int8_t>(i);
            out[static_cast<std::size_t>('a' + i)] = static_cast<std::int8_t>(26 + i);
        }
        for (int i = 0; i < 10; ++i) {
            out[static_cast<std::size_t>('0' + i)] = static_cast<std::int8_t>(52 + i);
        }
        out[static_cast<std::size_t>('+')] = 62;
        out[static_cast<std::size_t>('/')] = 63;
        return out;
    }();

    std::vector<std::uint8_t> out;
    out.reserve(text.size() * 3 / 4);
    std::uint32_t bits = 0;
    int count          = 0;
    bool padded        = false;
    for (const unsigned char c : text) {
        if (c == '=') {
            padded = true;
            continue;
        }
        if (c == ' ' || c == '\n' || c == '\r' || c == '\t') { continue; }
        if (padded || table[c] < 0) { throw std::invalid_argument("malformed base64 media data"); }
        bits = (bits << 6U) | static_cast<std::uint32_t>(table[c]);
        count += 6;
        if (count >= 8) {
            count -= 8;
            out.push_back(static_cast<std::uint8_t>((bits >> count) & 0xffU));
        }
    }
    if (count >= 6) { throw std::invalid_argument("malformed base64 media padding"); }
    return out;
}

bool private_ipv4(std::uint32_t address) {
    const std::uint32_t a = ntohl(address);
    return (a >> 24U) == 0 || (a >> 24U) == 10 || (a >> 24U) == 127 || (a >> 16U) == 0xa9fe ||
           (a >> 20U) == 0xac1 || (a >> 16U) == 0xc0a8 || (a >> 22U) == 0x0191 ||
           (a >> 17U) == 0x633f || (a >> 24U) >= 224;
}

bool private_address(const sockaddr* address) {
    if (address->sa_family == AF_INET) {
        return private_ipv4(reinterpret_cast<const sockaddr_in*>(address)->sin_addr.s_addr);
    }
    if (address->sa_family != AF_INET6) { return true; }
    const in6_addr& a = reinterpret_cast<const sockaddr_in6*>(address)->sin6_addr;
    if (IN6_IS_ADDR_UNSPECIFIED(&a) || IN6_IS_ADDR_LOOPBACK(&a) || IN6_IS_ADDR_LINKLOCAL(&a) ||
        IN6_IS_ADDR_MULTICAST(&a) || (a.s6_addr[0] & 0xfeU) == 0xfcU) {
        return true;
    }
    if (IN6_IS_ADDR_V4MAPPED(&a)) {
        std::uint32_t v4 = 0;
        std::memcpy(&v4, &a.s6_addr[12], sizeof(v4));
        return private_ipv4(v4);
    }
    return false;
}

struct UrlParts {
    std::string scheme;
    std::string host;
    std::string port;
};

std::string curlu_part(CURLU* url, CURLUPart part, unsigned flags = 0) {
    char* raw            = nullptr;
    const CURLUcode code = curl_url_get(url, part, &raw, flags);
    if (code != CURLUE_OK || raw == nullptr) { return {}; }
    std::string out(raw);
    curl_free(raw);
    return out;
}

UrlParts parse_url(std::string_view value) {
    std::unique_ptr<CURLU, decltype(&curl_url_cleanup)> url(curl_url(), curl_url_cleanup);
    if (!url ||
        curl_url_set(url.get(), CURLUPART_URL, std::string(value).c_str(), 0) != CURLUE_OK) {
        throw std::invalid_argument("invalid media URL");
    }
    UrlParts out;
    out.scheme             = curlu_part(url.get(), CURLUPART_SCHEME);
    out.host               = curlu_part(url.get(), CURLUPART_HOST);
    out.port               = curlu_part(url.get(), CURLUPART_PORT, CURLU_DEFAULT_PORT);
    const std::string user = curlu_part(url.get(), CURLUPART_USER);
    if ((out.scheme != "http" && out.scheme != "https") || out.host.empty() || out.port.empty() ||
        !user.empty()) {
        throw std::invalid_argument("media URL must be credential-free HTTP(S)");
    }
    return out;
}

std::string resolve_public(const UrlParts& url, bool allow_private) {
    addrinfo hints{};
    hints.ai_family   = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    addrinfo* raw     = nullptr;
    const int rc      = getaddrinfo(url.host.c_str(), url.port.c_str(), &hints, &raw);
    if (rc != 0) {
        throw Error(ErrorKind::RemoteUnavailable,
                    "failed to resolve media URL host: " + std::string(gai_strerror(rc)));
    }
    std::unique_ptr<addrinfo, decltype(&freeaddrinfo)> addresses(raw, freeaddrinfo);
    std::string selected;
    for (const addrinfo* it = raw; it != nullptr; it = it->ai_next) {
        if (!allow_private && private_address(it->ai_addr)) { continue; }
        std::array<char, INET6_ADDRSTRLEN> text{};
        const void* bytes =
            it->ai_family == AF_INET
                ? static_cast<const void*>(
                      &reinterpret_cast<const sockaddr_in*>(it->ai_addr)->sin_addr)
                : static_cast<const void*>(
                      &reinterpret_cast<const sockaddr_in6*>(it->ai_addr)->sin6_addr);
        if (inet_ntop(it->ai_family, bytes, text.data(), text.size()) != nullptr) {
            selected = text.data();
            if (it->ai_family == AF_INET) { break; }
        }
    }
    if (selected.empty()) {
        throw std::invalid_argument("media URL resolves only to disallowed network addresses");
    }
    return selected;
}

struct CurlBuffer {
    std::vector<std::uint8_t> bytes;
    std::size_t limit = 0;
};

int curl_progress(void* opaque, curl_off_t, curl_off_t, curl_off_t, curl_off_t) noexcept {
    const auto& policy = *static_cast<const Policy*>(opaque);
    try {
        return (policy.is_cancelled && policy.is_cancelled()) ||
                       (policy.deadline != Clock::time_point{} && Clock::now() >= policy.deadline)
                   ? 1
                   : 0;
    } catch (...) { return 1; }
}

std::size_t curl_write(char* data, std::size_t size, std::size_t count, void* opaque) {
    auto& out = *static_cast<CurlBuffer*>(opaque);
    if (size != 0 && count > std::numeric_limits<std::size_t>::max() / size) { return 0; }
    const std::size_t amount = size * count;
    if (amount > out.limit - std::min(out.limit, out.bytes.size())) { return 0; }
    out.bytes.insert(out.bytes.end(), reinterpret_cast<std::uint8_t*>(data),
                     reinterpret_cast<std::uint8_t*>(data) + amount);
    return amount;
}

std::vector<std::uint8_t> fetch_url(std::string url, const Policy& policy) {
    if (!policy.allow_remote) { throw std::invalid_argument("remote media URLs are disabled"); }
    static std::once_flag init;
    std::call_once(init, [] {
        if (curl_global_init(CURL_GLOBAL_DEFAULT) != CURLE_OK) {
            throw std::runtime_error("failed to initialize libcurl");
        }
    });

    for (int redirect = 0; redirect <= policy.max_redirects; ++redirect) {
        check_control(policy);
        const UrlParts parts = parse_url(url);
        const std::string ip = resolve_public(parts, policy.allow_private_network);
        check_control(policy);
        std::string resolve = parts.host + ":" + parts.port + ":";
        resolve += ip.find(':') == std::string::npos ? ip : "[" + ip + "]";
        curl_slist* resolve_list = curl_slist_append(nullptr, resolve.c_str());
        if (resolve_list == nullptr) { throw std::bad_alloc(); }
        std::unique_ptr<curl_slist, decltype(&curl_slist_free_all)> resolve_guard(
            resolve_list, curl_slist_free_all);
        std::unique_ptr<CURL, decltype(&curl_easy_cleanup)> curl(curl_easy_init(),
                                                                 curl_easy_cleanup);
        if (!curl) { throw std::runtime_error("failed to create libcurl handle"); }
        CurlBuffer buffer{.bytes = {}, .limit = policy.max_bytes};
        curl_easy_setopt(curl.get(), CURLOPT_URL, url.c_str());
        curl_easy_setopt(curl.get(), CURLOPT_PROTOCOLS_STR, "http,https");
        curl_easy_setopt(curl.get(), CURLOPT_REDIR_PROTOCOLS_STR, "http,https");
        curl_easy_setopt(curl.get(), CURLOPT_FOLLOWLOCATION, 0L);
        curl_easy_setopt(curl.get(), CURLOPT_CONNECTTIMEOUT_MS,
                         bounded_timeout_ms(policy, policy.connect_timeout_ms));
        curl_easy_setopt(curl.get(), CURLOPT_TIMEOUT_MS,
                         bounded_timeout_ms(policy, policy.timeout_ms));
        curl_easy_setopt(curl.get(), CURLOPT_NOSIGNAL, 1L);
        curl_easy_setopt(curl.get(), CURLOPT_SSL_VERIFYPEER, 1L);
        curl_easy_setopt(curl.get(), CURLOPT_SSL_VERIFYHOST, 2L);
        curl_easy_setopt(curl.get(), CURLOPT_PROXY, "");
        curl_easy_setopt(curl.get(), CURLOPT_RESOLVE, resolve_list);
        curl_easy_setopt(curl.get(), CURLOPT_MAXFILESIZE_LARGE,
                         static_cast<curl_off_t>(policy.max_bytes));
        curl_easy_setopt(curl.get(), CURLOPT_WRITEFUNCTION, curl_write);
        curl_easy_setopt(curl.get(), CURLOPT_WRITEDATA, &buffer);
        curl_easy_setopt(curl.get(), CURLOPT_NOPROGRESS, 0L);
        curl_easy_setopt(curl.get(), CURLOPT_XFERINFOFUNCTION, curl_progress);
        curl_easy_setopt(curl.get(), CURLOPT_XFERINFODATA, &policy);
        curl_easy_setopt(curl.get(), CURLOPT_USERAGENT, "ninfer/vision");
        const CURLcode code = curl_easy_perform(curl.get());
        if (code != CURLE_OK) {
            check_control(policy);
            if (code == CURLE_OPERATION_TIMEDOUT) {
                throw Error(ErrorKind::RemoteTimeout,
                            "media URL fetch timed out: " + std::string(curl_easy_strerror(code)));
            }
            if (code == CURLE_FILESIZE_EXCEEDED || code == CURLE_WRITE_ERROR) {
                throw Error(ErrorKind::BudgetExceeded, "media URL exceeds byte limit");
            }
            throw Error(ErrorKind::RemoteUnavailable,
                        "failed to fetch media URL: " + std::string(curl_easy_strerror(code)));
        }
        long status = 0;
        curl_easy_getinfo(curl.get(), CURLINFO_RESPONSE_CODE, &status);
        if (status >= 200 && status < 300) {
            check_control(policy);
            return buffer.bytes;
        }
        if (status < 300 || status >= 400 || redirect == policy.max_redirects) {
            throw Error(ErrorKind::RemoteUnavailable,
                        "media URL returned HTTP " + std::to_string(status));
        }
        char* next = nullptr;
        curl_easy_getinfo(curl.get(), CURLINFO_REDIRECT_URL, &next);
        if (next == nullptr || *next == '\0') {
            throw Error(ErrorKind::RemoteUnavailable, "media URL redirect has no location");
        }
        url = next;
    }
    throw Error(ErrorKind::RemoteUnavailable, "too many media URL redirects");
}

std::vector<std::uint8_t> read_path(const Source& source, const Policy& policy) {
    check_control(policy);
    std::error_code ec;
    std::filesystem::path path = std::filesystem::weakly_canonical(source.value, ec);
    if (ec || !std::filesystem::is_regular_file(path, ec)) {
        throw std::invalid_argument("media path is not a regular file: " + source.value);
    }
    if (!policy.media_root.empty()) {
        const std::filesystem::path root = std::filesystem::weakly_canonical(policy.media_root, ec);
        const auto relative              = std::filesystem::relative(path, root, ec);
        if (ec || relative.empty() || relative.native().starts_with("..")) {
            throw std::invalid_argument("media path is outside configured media root");
        }
    }
    const std::uintmax_t size = std::filesystem::file_size(path, ec);
    check_control(policy);
    if (ec) { throw std::invalid_argument("failed to inspect media file: " + source.value); }
    if (size > policy.max_bytes) {
        throw Error(ErrorKind::BudgetExceeded, "media file exceeds byte limit");
    }

    std::ifstream stream(path, std::ios::binary);
    if (!stream) { throw std::invalid_argument("failed to open media path: " + path.string()); }
    stream.seekg(0, std::ios::end);
    const std::streamoff end = stream.tellg();
    if (end < 0 || static_cast<std::uintmax_t>(end) > policy.max_bytes) {
        throw Error(ErrorKind::BudgetExceeded, "media file exceeds byte limit");
    }
    stream.seekg(0, std::ios::beg);
    std::vector<std::uint8_t> bytes(static_cast<std::size_t>(end));
    if (!bytes.empty()) {
        stream.read(reinterpret_cast<char*>(bytes.data()),
                    static_cast<std::streamsize>(bytes.size()));
        if (!stream) { throw std::invalid_argument("failed to read media path: " + path.string()); }
    }
    check_control(policy);
    return bytes;
}

} // namespace

std::vector<std::uint8_t> acquire_bytes(const Source& source, const Policy& policy) {
    if (policy.max_bytes == 0) { throw std::invalid_argument("media byte limit must be positive"); }
    check_control(policy);
    if (source.kind == SourceKind::Bytes) {
        if (source.bytes.empty()) { throw std::invalid_argument("media source is empty"); }
        if (source.bytes.size() > policy.max_bytes) {
            throw Error(ErrorKind::BudgetExceeded, "media bytes exceed byte limit");
        }
        std::vector<std::uint8_t> bytes = source.bytes;
        check_control(policy);
        return bytes;
    }
    if (source.value.empty()) { throw std::invalid_argument("media source is empty"); }

    if (source.kind == SourceKind::Url) {
        std::vector<std::uint8_t> bytes = fetch_url(source.value, policy);
        if (bytes.empty()) { throw std::invalid_argument("media source contains no data"); }
        return bytes;
    }
    if (source.kind == SourceKind::Data) {
        const std::size_t comma = source.value.find(',');
        if (!source.value.starts_with("data:") || comma == std::string::npos ||
            source.value.substr(0, comma).find(";base64") == std::string::npos) {
            throw std::invalid_argument("media data source must be a base64 data URI");
        }
        const std::size_t encoded = source.value.size() - comma - 1;
        if (encoded > (policy.max_bytes / 3 + 1) * 4) {
            throw Error(ErrorKind::BudgetExceeded, "media data exceeds byte limit");
        }
        std::vector<std::uint8_t> bytes =
            decode_base64(std::string_view(source.value).substr(comma + 1));
        check_control(policy);
        if (bytes.size() > policy.max_bytes) {
            throw Error(ErrorKind::BudgetExceeded, "media data exceeds byte limit");
        }
        if (bytes.empty()) { throw std::invalid_argument("media source contains no data"); }
        return bytes;
    }
    return read_path(source, policy);
}

} // namespace ninfer::product::media_acquire
#pragma once

#include "core/tensor.h"

#include <cuda_runtime.h>

namespace ninfer::ops {

/**
 * Adds a four-corner interpolated position embedding. For d,p:
 *
 *   position[d,p] = sum_{c=0..3} table[d,indices[c,p]] * weights[c,p]
 *   ideal[d,p]    = x[d,p] + position[d,p].
 *
 * table is contiguous BF16 [D,R], indices is contiguous I32 [4,P] with values in [0,R), weights
 * is contiguous FP32 [4,P], and x is contiguous BF16 [D,P]. The oracle evaluates `ideal` naively 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.
 * Production routes may choose their natural accumulator and staging precision. Inputs must not
 * overlap x. The Op updates all of x in place and uses no workspace or other persistent state.
 * D=1152 with aligned table/x storage selects the optimized small-P warp or large-P CTA route;
 * other contiguous dimensions use the scalar route.
 */
void vision_pos_embed_add(const Tensor& table, const Tensor& indices, const Tensor& weights,
                          Tensor& x, cudaStream_t stream);

} // namespace ninfer::ops
// Implements: include/ninfer/ops/argmax.h
// Match: validated contiguous BF16 logits and I32 output.
// Algorithm assumptions: one tile uses a direct reduction; larger domains use
// zero-initialized atomic winners across route-selected row tiles.
#include "ops/launcher/argmax.h"

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

#include <cstdint>

namespace ninfer::ops::detail {
namespace {

constexpr std::int32_t kFullPhysicalRows      = 248320;
constexpr std::int32_t kFullValidRows         = 248077;
constexpr std::int32_t kShortlistRows         = 131072;
constexpr std::int32_t kSmallAggregateColumns = 8;
constexpr int kFullAggregateBlock             = 128;
constexpr int kShortlistAggregateBlock        = 256;

int tiled_block_for(std::int32_t physical_rows, std::int32_t valid_rows, std::int32_t t_count) {
    if (t_count <= kSmallAggregateColumns) { return kArgmaxBlock; }
    if (physical_rows == kFullPhysicalRows && valid_rows == kFullValidRows) {
        return kFullAggregateBlock;
    }
    if (physical_rows == kShortlistRows && valid_rows == kShortlistRows) {
        return kShortlistAggregateBlock;
    }
    return kArgmaxBlock;
}

void argmax_tiled_atomic_launch(const Tensor& logits, Tensor& out, std::int32_t valid_rows,
                                int block, cudaStream_t stream);

} // namespace

void argmax_launch(const Tensor& logits, Tensor& out, std::int32_t valid_rows,
                   cudaStream_t stream) {
    const std::int32_t physical_rows = logits.ne[0];
    const std::int32_t t_count       = logits.ne[1];
    if (t_count == 0) { return; }

    constexpr int kTileElems = kArgmaxBlock * kArgmaxItemsPerThread;
    const int tiled_blocks   = div_up(valid_rows, kTileElems);
    if (tiled_blocks < 2) {
        argmax_kernel<<<static_cast<unsigned int>(t_count), kArgmaxBlock, 0, stream>>>(
            static_cast<const __nv_bfloat16*>(logits.data), static_cast<std::int32_t*>(out.data),
            valid_rows, physical_rows);
        CUDA_CHECK(cudaGetLastError());
        return;
    }

    argmax_tiled_atomic_launch(logits, out, valid_rows,
                               tiled_block_for(physical_rows, valid_rows, t_count), stream);
}

namespace {

void argmax_tiled_atomic_launch(const Tensor& logits, Tensor& out, std::int32_t valid_rows,
                                int block, cudaStream_t stream) {
    const std::int32_t physical_rows = logits.ne[0];
    const std::int32_t t_count       = logits.ne[1];
    const int tiled_blocks           = div_up(valid_rows, block * kArgmaxItemsPerThread);
    for_each_token_slice(t_count, 1, [&](int token_offset, int token_count) {
        const Tensor logits_slice = logits.slice(1, token_offset, token_count);
        Tensor out_slice          = out.slice(0, token_offset, token_count);
        CUDA_CHECK(cudaMemsetAsync(out_slice.data, 0,
                                   static_cast<std::size_t>(token_count) * sizeof(std::int32_t),
                                   stream));
        const dim3 grid(static_cast<unsigned int>(tiled_blocks),
                        static_cast<unsigned int>(token_count));
        argmax_tiled_atomic_kernel<<<grid, block, 0, stream>>>(
            static_cast<const __nv_bfloat16*>(logits_slice.data),
            static_cast<std::int32_t*>(out_slice.data), valid_rows, physical_rows);
        CUDA_CHECK(cudaGetLastError());
    });
}

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

#include "ninfer/ops/attention_geometry.h"

#include "core/arena.h"
#include "core/paged_kv_cache.h"
#include "core/tensor.h"

#include <cuda_runtime.h>

#include <cstddef>
#include <cstdint>

namespace ninfer::ops {

inline constexpr std::uint32_t kCausalAttentionMaximumVisibleKeys = 262144;

struct CausalAttentionExecutionEnvelope {
    std::uint32_t min_visible_keys = 0;
    std::uint32_t max_visible_keys = 0;
};

struct ContextAttentionExecutionEnvelope {
    std::uint32_t min_context = 0;
    std::uint32_t max_context = 0;
};

/**
 * Shared numerical contract.
 *
 * Every entry computes stable scaled dot-product Softmax Attention. Query head h reads KV head
 * floor(h / (Hq/Hkv)). Public BF16 inputs and BF16 cache rows are interpreted after their storage
 * boundary. For a declared visible key set J, the independent mathematical oracle is
 *
 *   score[j]       = scale * dot(FP64(q[:,h,i]), FP64(k[:,kvh,j]))
 *   probability[j] = exp(score[j] - max(score)) / sum_x exp(score[x] - max(score))
 *   ideal[:,h,i]   = sum_j probability[j] * FP64(v[:,kvh,j]).
 *
 * Dot products, stable Softmax, and the value reduction are evaluated naively in FP64. The BF16
 * Op output is promoted to FP64 for comparison; storage rounding belongs to the Op criterion and
 * is not reproduced by the oracle.
 *
 * A causal INT8-G64 V row is interpreted by decoding each signed code c with its stored FP16 scale
 * bits: FP32(c) * FP32(scale_bits). A causal FP8-E4M3FN V row has one stored FP16 scale for D256;
 * each finite code e is decoded as FP32(e) * FP32(scale_bits). Quantized K uses the paired physical
 * representation written by kv_cache_append; its original-coordinate logical row is consumed
 * through the matching private Q/K profile. The fixed orthogonal preparation and transient Q
 * quantization are implementation details, not intermediate values in the ideal oracle above.
 *
 * The qualified FP8 compute profile uses native E4M3FN QK MMA with FP32 accumulation, FP32
 * Softmax, exact E4M3FN-to-FP16 V-code conversion followed by one represented FP16 scale multiply,
 * FP16 P/V MMA with FP32 accumulation, FP32 split merge/normalization, and a final BF16 output
 * store. This arithmetic path is an implementation profile rather than an extra public semantic
 * boundary. BF16, INT8-G64, and FP8-E4M3FN routes have separate numerical criteria and are each
 * checked directly against the ideal oracle; route-to-route parity is only supplementary evidence.
 * Those criteria apply to the registered geometries, tested extents, conformance matrix, and
 * target-representative activation range; they are not universal error bounds for arbitrary
 * adversarial BF16 tensors.
 */

/**
 * Dense, non-causal single-segment attention.
 *
 * The registered profile is D=72, Hq=Hkv=16, scale=1/sqrt(72). q/k/v are BF16 [72,16,T]
 * with contiguous feature and head dimensions; their token stride may be padded. out is contiguous
 * BF16 [72,16,T]. Every query attends all T keys. q/k/v/out are mutually non-overlapping, inputs
 * are unchanged, out is completely overwritten, and the Op has no persistent state side effect.
 * The single segment needs no transient workspace.
 */
void softmax_attention(const Tensor& q, const Tensor& k, const Tensor& v,
                       AttentionHeadGeometry geometry, float scale, WorkspaceArena& workspace,
                       Tensor& out, cudaStream_t stream);

/**
 * Packed block-diagonal dense attention for the same D72/H16 profile.
 *
 * cu_seqlens is contiguous device I32 [S+1], starts at 0, ends at T, and is strictly increasing.
 * Each range [cu_seqlens[s],cu_seqlens[s+1]) is an independent non-causal segment; no score crosses
 * a segment boundary. q/k/v/out have the dtype, layout, alias, numerical, mutation, and state
 * contract of softmax_attention. Opaque tile descriptors are allocated from workspace for the
 * duration of the call; a single segment consumes no capacity.
 */
void packed_softmax_attention(const Tensor& q, const Tensor& k, const Tensor& v,
                              AttentionHeadGeometry geometry, float scale, const Tensor& cu_seqlens,
                              WorkspaceArena& workspace, Tensor& out, cudaStream_t stream);

/**
 * Equal-length form of packed_softmax_attention. T is divisible by segment_length and consecutive
 * ranges [s*segment_length,(s+1)*segment_length) are the independent segments. Segment descriptors
 * are derived directly, so this form needs no workspace or descriptor-setup launch.
 */
void packed_softmax_attention(const Tensor& q, const Tensor& k, const Tensor& v,
                              AttentionHeadGeometry geometry, float scale,
                              std::int32_t segment_length, Tensor& out, cudaStream_t stream);

/**
 * Return caller-owned transient capacity for every legal (T,S) pair in the inclusive envelope.
 * A pair is legal when 1 <= S <= T. An envelope with no legal pair throws; a legal single-segment
 * envelope may return zero.
 */
[[nodiscard]] std::size_t packed_softmax_attention_workspace_capacity_bytes(
    AttentionHeadGeometry geometry, std::int32_t min_tokens, std::int32_t max_tokens,
    std::int32_t min_segments, std::int32_t max_segments);

/**
 * Append K/V for B independent rows and compute causal grouped-query attention.
 *
 * The registered profiles are [D,Hq,Hkv]=[256,24,4] (group 6) and [256,16,2] (group 8), with
 * scale=1/sqrt(256). q/out are contiguous BF16 [D,Hq,W,B], k/v are contiguous BF16
 * [D,Hkv,W,B], positions are contiguous device I32 [W,B], kv_table_rows is contiguous device I32
 * [B], and the cache is BF16, INT8-G64, or row-scaled FP8-E4M3FN. valid_columns is either
 * contiguous device I32 [B] or an empty Tensor meaning every row has W live columns. This
 * dense/masked topology is chosen by the caller and never inferred by copying device metadata to
 * the host. B=1 accepts every positive W in the current prompt/decode domain; B=2..8 accepts
 * W=1..16.
 *
 * Let Vb be W for dense input or valid_columns[b] otherwise. For live column j<Vb with absolute
 * position p=positions[j,b], query head h attends cache rows [0,p] through table row
 * kv_table_rows[b]. The current k/v row is appended before it is observed, so the formula is the
 * shared oracle above over J=[0,p]. Each masked row has a live prefix [0,Vb); its live positions
 * are sequential and address populated histories. A nonempty row repeats its last live position
 * through the inert tail; an empty row uses zero positions. Other tail values are safe dummies.
 * Tail columns do not mutate cache and produce exact BF16 zero.
 *
 * The caller guarantees that the maximum p+1 over live rows lies within envelope. The envelope is
 * a host launch/workspace resource promise over that batch maximum, not a mask and not persistent
 * state. Inputs, output, every cache plane/table, and live workspace suballocations are pairwise
 * non-overlapping. The Op overwrites every addressed cache row but owns no cache allocation,
 * frontier, request identity, or commit authority.
 */
void causal_softmax_attention(const Tensor& q, const Tensor& k, const Tensor& v,
                              const Tensor& positions, const Tensor& valid_columns,
                              const Tensor& kv_table_rows, AttentionHeadGeometry geometry,
                              float scale, PagedKVBatchLayerView cache,
                              CausalAttentionExecutionEnvelope envelope, WorkspaceArena& workspace,
                              Tensor& out, cudaStream_t stream);

/**
 * Read-only single-sequence causal attention over an already populated cache.
 *
 * q/out are contiguous BF16 [256,24|16,T], positions is contiguous sequential device I32 [T],
 * and cache geometry, visible rows, scale, numerical oracle, envelope, alias, and workspace rules
 * are the same as causal_softmax_attention. The Op accepts no new K/V and leaves every cache byte
 * unchanged; out is completely overwritten.
 */
void causal_softmax_attention_cached(const Tensor& q, const Tensor& positions,
                                     AttentionHeadGeometry geometry, float scale,
                                     const PagedKVLayerView& cache,
                                     CausalAttentionExecutionEnvelope envelope,
                                     WorkspaceArena& workspace, Tensor& out, cudaStream_t stream);

/**
 * Return transient capacity for every W in the inclusive interval at one exact batch size. The
 * head geometry, cache dtype, and execution envelope are fixed implementation-profile inputs.
 * Invalid profiles or intervals throw; a legal prompt route may return zero.
 */
[[nodiscard]] std::size_t causal_softmax_attention_workspace_capacity_bytes(
    AttentionHeadGeometry geometry, DType cache_dtype, CausalAttentionExecutionEnvelope envelope,
    std::int32_t batch_size, std::int32_t min_tokens, std::int32_t max_tokens);

/**
 * Non-causal grouped-query attention over persistent context plus one live query block.
 *
 * The registered profile is D=128, Hq=32, Hkv=8 (group 4), scale=1/sqrt(128), T=1..16, and
 * B=1..8. q/out are contiguous BF16 [128,32,T,B], query_k/query_v are contiguous BF16
 * [128,8,T,B], and context_lengths, valid_columns, and table_rows are contiguous device I32 [B].
 * The read-only paged BF16 context uses head-major page planes [128,64,Nphysical,8].
 *
 * For row b, let L=context_lengths[b] and V=valid_columns[b]. Every live query i<V attends the
 * complete logical set consisting of context rows [0,L) followed by every temporary query K/V row
 * [0,V). There is no causal triangle. Columns i>=V are inert and produce exact BF16 zero. The
 * context and temporary query K/V remain separate physical segments; every input and cache byte is
 * unchanged, and out is the only observable mutation and is completely overwritten.
 *
 * The caller guarantees envelope.min_context <= L <= envelope.max_context and materialized block
 * table entries for every logical page intersecting [0,L). The envelope may affect finite launch
 * selection and workspace capacity, never the admitted key set or numerical result. Inputs,
 * context, output, and live workspace allocations are pairwise non-overlapping.
 */
void context_softmax_attention(const Tensor& q, const Tensor& query_k, const Tensor& query_v,
                               const Tensor& context_lengths, const Tensor& valid_columns,
                               const Tensor& table_rows, AttentionHeadGeometry geometry,
                               float scale, const PagedKVBatchLayerView& context,
                               ContextAttentionExecutionEnvelope envelope,
                               WorkspaceArena& workspace, Tensor& out, cudaStream_t stream);

/**
 * Return transient capacity for every T in the inclusive optimized interval at the exact batch
 * size. The execution envelope is fixed; invalid profiles or intervals throw.
 */
[[nodiscard]] std::size_t context_softmax_attention_workspace_capacity_bytes(
    AttentionHeadGeometry geometry, ContextAttentionExecutionEnvelope envelope,
    std::int32_t min_tokens, std::int32_t max_tokens, std::int32_t batch_size);

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

#include "artifact/materializer.h"
#include "artifact/typed_binding.h"

#include <cstddef>
#include <string>

namespace ninfer::targets::qwen3_6 {
namespace {

std::string take_string(artifact::MaterializedArtifact& materialized,
                        artifact::ObjectHandle handle) {
    const auto bytes = materialized.take_resource_bytes(handle);
    return std::string(reinterpret_cast<const char*>(bytes.data()), bytes.size());
}

} // namespace

FrontendResourcePlan bind_frontend_resources(artifact::Binder& binder) {
    return FrontendResourcePlan{
        .tokenizer_json = artifact::bind_raw_resource(binder, "frontend/tokenizer.json"),
        .tokenizer_config_json =
            artifact::bind_raw_resource(binder, "frontend/tokenizer_config.json"),
        .chat_template_jinja = artifact::bind_raw_resource(binder, "frontend/chat_template.jinja"),
        .generation_config_json =
            artifact::bind_raw_resource(binder, "frontend/generation_config.json"),
        .preprocessor_config_json =
            artifact::bind_raw_resource(binder, "frontend/preprocessor_config.json"),
        .video_preprocessor_config_json =
            artifact::bind_raw_resource(binder, "frontend/video_preprocessor_config.json"),
    };
}

FrontendResources take_frontend_resources(artifact::MaterializedArtifact& materialized,
                                          const FrontendResourcePlan& plan) {
    return FrontendResources{
        .tokenizer_json           = take_string(materialized, plan.tokenizer_json),
        .tokenizer_config_json    = take_string(materialized, plan.tokenizer_config_json),
        .chat_template_jinja      = take_string(materialized, plan.chat_template_jinja),
        .generation_config_json   = take_string(materialized, plan.generation_config_json),
        .preprocessor_config_json = take_string(materialized, plan.preprocessor_config_json),
        .video_preprocessor_config_json =
            take_string(materialized, plan.video_preprocessor_config_json),
    };
}

} // namespace ninfer::targets::qwen3_6
#pragma once

#include "targets/qwen3_6/impl/runtime/instance.h"
// Qwen3.6 family runtime implementation; instantiated only by exact variants.

#include "core/arena.h"
#include "core/device.h"
#include "core/tensor.h"
#include "core/weight.h"
#include <ninfer/targets/qwen3_6/vision_control.h>
#include "targets/qwen3_6/impl/runtime/layouts.h"
#include "targets/qwen3_6/impl/runtime/vision_prefill.h"

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

namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::schedule {

struct VisionItemView {
    std::span<const std::uint16_t> patches;
    const qwen3_6::VisionItemControl* control = nullptr;
};

struct VisionScheduleConfig {
    static constexpr int layers              = VisionConfig::layers;
    static constexpr int hidden              = VisionConfig::hidden;
    static constexpr int intermediate        = VisionConfig::intermediate;
    static constexpr int out_hidden          = VisionConfig::output_hidden;
    static constexpr int heads               = VisionConfig::heads;
    static constexpr int head_dim            = VisionConfig::head_dim;
    static constexpr int patch_dim           = VisionConfig::patch_dim;
    static constexpr int merge_unit          = VisionConfig::merge_unit;
    static constexpr int merger_hidden       = VisionConfig::merger_hidden;
    static constexpr int position_embeddings = VisionConfig::position_embeddings;
    static constexpr int rotary_dim          = VisionConfig::rotary_dim;
    static constexpr float rope_theta        = VisionConfig::rope_theta;
    static constexpr float norm_eps          = VisionConfig::norm_epsilon;
    static constexpr float attention_scale   = 0.11785113019775792F;
};

class VisionContext {
public:
    VisionContext(DeviceContext& device, const LoadedModelData& model);

    [[nodiscard]] static std::size_t workspace_bytes(const qwen3_6::VisionItemControl& item);
    [[nodiscard]] static std::size_t workspace_bytes(std::size_t patches,
                                                     std::size_t merged_tokens);
    [[nodiscard]] static VisionWorkspacePlan plan_workspace(std::uint32_t max_merged_tokens,
                                                            std::size_t general_capacity_bytes);
    [[nodiscard]] static Tensor bind_output(DeviceSpan backing, const VisionWorkspacePlan& plan,
                                            std::size_t merged_tokens);
    void encode(const VisionItemView& item, Tensor& output, DeviceSpan backing,
                const VisionWorkspacePlan& plan) const;

private:
    struct BlockW {
        const Tensor* norm1_weight    = nullptr;
        const Tensor* norm1_bias      = nullptr;
        const Weight* qkv             = nullptr;
        const Tensor* qkv_bias        = nullptr;
        const Weight* projection      = nullptr;
        const Tensor* projection_bias = nullptr;
        const Tensor* norm2_weight    = nullptr;
        const Tensor* norm2_bias      = nullptr;
        const Weight* fc1             = nullptr;
        const Tensor* fc1_bias        = nullptr;
        const Weight* fc2             = nullptr;
        const Tensor* fc2_bias        = nullptr;
    };

    struct MergerW {
        const Tensor* norm_weight = nullptr;
        const Tensor* norm_bias   = nullptr;
        const Weight* fc1         = nullptr;
        const Tensor* fc1_bias    = nullptr;
        const Weight* fc2         = nullptr;
        const Tensor* fc2_bias    = nullptr;
    };

    DeviceContext& ctx_;
    const Weight* patch_embed_      = nullptr;
    const Tensor* patch_embed_bias_ = nullptr;
    const Tensor* position_embed_   = nullptr;
    std::array<BlockW, VisionScheduleConfig::layers> blocks_{};
    MergerW merger_{};
};

struct VisionChunk {
    std::int32_t length                       = 0;
    const qwen3_6::VisionItemControl* control = nullptr;
    Tensor embeddings;
};

class VisionPrefillSession {
public:
    VisionPrefillSession(DeviceContext& device, const LoadedModelData& model, DeviceSpan workspace,
                         const VisionWorkspacePlan& workspace_plan,
                         qwen3_6::PreparedPromptData& prompt, const VisionPrefillPlan& plan,
                         std::size_t& handoff_peak_bytes);

    [[nodiscard]] VisionChunk prepare_chunk(std::uint32_t begin, std::uint32_t nominal_length);
    void release_encoded_media_payloads() noexcept;
    void retire_handoff() noexcept;
    [[nodiscard]] double elapsed_seconds() const;

    [[nodiscard]] std::size_t active_handoff_bytes() const noexcept {
        return active_handoff_bytes_;
    }

private:
    DeviceContext& device_;
    DeviceSpan workspace_;
    const VisionWorkspacePlan& workspace_plan_;
    qwen3_6::PreparedPromptData& prompt_;
    const VisionPrefillPlan& plan_;
    std::size_t& handoff_peak_bytes_;
    VisionContext context_;
    std::size_t next_use_ = 0;
    std::optional<std::uint32_t> active_item_;
    std::size_t active_handoff_bytes_ = 0;
    std::vector<std::uint32_t> encoded_payloads_pending_release_;
    std::vector<CudaEventTimer> timers_;
};

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

#include "ops/common/memory.cuh"
#include "ops/common/warp.cuh"
#include "ops/kernel/paged_kv_address.cuh"
#include "ops/kv_cache/fp8_e4m3_row_codec.cuh"
#include "ops/kv_cache/int8_g64_codec.cuh"

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

#include <cstdint>

namespace ninfer::ops {

inline constexpr int kKVCacheAppendFullHeadDim = 256;

template <int KVHeadsValue>
struct KVCacheAppendFullGeometry {
    static_assert(KVHeadsValue == 4 || KVHeadsValue == 2);
    static constexpr int KVHeads = KVHeadsValue;
};

using KVCacheAppendD256Kv4 = KVCacheAppendFullGeometry<4>;
using KVCacheAppendD256Kv2 = KVCacheAppendFullGeometry<2>;

struct KVCacheAppendDirectMetadata {
    const std::int32_t* table;

    __device__ __forceinline__ std::int32_t valid_tokens(std::int32_t width) const { return width; }

    __device__ __forceinline__ const std::int32_t* block_table() const { return table; }
};

template <bool Masked>
struct KVCacheAppendBatchMetadata {
    const std::int32_t* tables;
    const std::int32_t* valid_columns;
    const std::int32_t* table_rows;
    std::int32_t table_stride;

    __device__ __forceinline__ std::int32_t valid_tokens(std::int32_t width) const {
        if constexpr (Masked) {
            const std::int32_t valid = valid_columns[0];
            return valid <= 0 ? 0 : (valid < width ? valid : width);
        }
        return width;
    }

    __device__ __forceinline__ const std::int32_t* block_table() const {
        return tables + static_cast<std::int64_t>(table_rows[0]) * table_stride;
    }
};

template <typename Geometry>
__device__ __forceinline__ void
kv_cache_append_full_fp8_row(const __nv_bfloat16* __restrict__ k,
                             const __nv_bfloat16* __restrict__ v,
                             std::uint8_t* __restrict__ cache_k, std::uint8_t* __restrict__ cache_v,
                             __half* __restrict__ scale_k, __half* __restrict__ scale_v, int token,
                             int kv_head, int physical_page, int page_off, int lane) {
    constexpr unsigned FullMask = 0xffffffffU;
    float values[8];
    float local_absmax = 0.0F;
#pragma unroll
    for (int r = 0; r < 8; ++r) {
        const int d = lane + 32 * r;
        values[r]   = __bfloat162float(k[kv_cache_fp8_src_index<Geometry>(kv_head, d, token)]);
    }
    normalized_hadamard_d256_inplace(values, lane);
#pragma unroll
    for (float value : values) { local_absmax = fmaxf(local_absmax, fabsf(value)); }
    const auto k_quant = kv_cache_fp8_quant_params(warp_max(local_absmax, FullMask));
#pragma unroll
    for (int r = 0; r < 8; ++r) {
        const int d = lane + 32 * r;
        cache_k[kv_cache_fp8_code_index<Geometry>(physical_page, kv_head, d, page_off)] =
            kv_cache_fp8_quant_code(values[r], k_quant.inverse_scale);
    }
    if (lane == 0) {
        scale_k[kv_cache_fp8_scale_index<Geometry>(physical_page, kv_head, page_off)] =
            k_quant.scale;
    }

    local_absmax = 0.0F;
#pragma unroll
    for (int r = 0; r < 8; ++r) {
        const int d  = lane + 32 * r;
        values[r]    = __bfloat162float(v[kv_cache_fp8_src_index<Geometry>(kv_head, d, token)]);
        local_absmax = fmaxf(local_absmax, fabsf(values[r]));
    }
    const auto v_quant = kv_cache_fp8_quant_params(warp_max(local_absmax, FullMask));
#pragma unroll
    for (int r = 0; r < 8; ++r) {
        const int d = lane + 32 * r;
        cache_v[kv_cache_fp8_code_index<Geometry>(physical_page, kv_head, d, page_off)] =
            kv_cache_fp8_quant_code(values[r], v_quant.inverse_scale);
    }
    if (lane == 0) {
        scale_v[kv_cache_fp8_scale_index<Geometry>(physical_page, kv_head, page_off)] =
            v_quant.scale;
    }
}

template <typename Geometry, typename Metadata>
__global__ void kv_cache_append_full_bf16_kernel(
    const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v,
    const std::int32_t* __restrict__ positions, Metadata metadata,
    __nv_bfloat16* __restrict__ cache_k, __nv_bfloat16* __restrict__ cache_v, std::int32_t width) {
    constexpr int VecElems = 8;
    const int tokens       = metadata.valid_tokens(width);
    const std::int64_t idx = static_cast<std::int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
    const std::int64_t n   = static_cast<std::int64_t>(tokens) * Geometry::KVHeads *
                           (kKVCacheAppendFullHeadDim / VecElems);
    if (idx >= n) return;

    const int vec      = static_cast<int>(idx % (kKVCacheAppendFullHeadDim / VecElems));
    const int tmp      = static_cast<int>(idx / (kKVCacheAppendFullHeadDim / VecElems));
    const int kv_head  = tmp % Geometry::KVHeads;
    const int token    = tmp / Geometry::KVHeads;
    const int d        = vec * VecElems;
    const int position = positions[0] + token;
    const int lane     = static_cast<int>(threadIdx.x) & 31;
    const std::int32_t* block_table = metadata.block_table();
    int physical_page               = lane == 0 ? paged_kv_physical_page(block_table, position) : 0;
    const std::int64_t src_off =
        static_cast<std::int64_t>(d) + static_cast<std::int64_t>(kKVCacheAppendFullHeadDim) *
                                           (kv_head + Geometry::KVHeads * token);
    const int4 k_value = load_vec<int4>(&k[src_off]);
    const int4 v_value = load_vec<int4>(&v[src_off]);
    physical_page      = __shfl_sync(0xffffffffu, physical_page, 0);
    const std::int64_t cache_off =
        paged_kv_element_offset<kKVCacheAppendFullHeadDim, Geometry::KVHeads>(
            physical_page, kv_head, position & kPagedKVPageMask, d);
    store_vec(&cache_k[cache_off], k_value);
    store_vec(&cache_v[cache_off], v_value);
}

template <typename Geometry, typename Metadata>
__launch_bounds__(256) __global__
    void kv_cache_append_full_fp8_kernel(const __nv_bfloat16* __restrict__ k,
                                         const __nv_bfloat16* __restrict__ v,
                                         const std::int32_t* __restrict__ positions,
                                         Metadata metadata, std::uint8_t* __restrict__ cache_k,
                                         std::uint8_t* __restrict__ cache_v,
                                         __half* __restrict__ scale_k, __half* __restrict__ scale_v,
                                         std::int32_t width) {
    constexpr int Warps         = 8;
    constexpr unsigned FullMask = 0xffffffffU;
    const int tokens            = metadata.valid_tokens(width);
    const int warp              = static_cast<int>(threadIdx.x) >> 5;
    const int lane              = static_cast<int>(threadIdx.x) & 31;
    const int unit              = static_cast<int>(blockIdx.x) * Warps + warp;
    const int units             = tokens * Geometry::KVHeads;
    if (unit >= units) return;

    const int kv_head               = unit % Geometry::KVHeads;
    const int token                 = unit / Geometry::KVHeads;
    const int position              = positions[0] + token;
    const std::int32_t* block_table = metadata.block_table();
    int physical_page               = lane == 0 ? paged_kv_physical_page(block_table, position) : 0;
    physical_page                   = __shfl_sync(FullMask, physical_page, 0);
    kv_cache_append_full_fp8_row<Geometry>(k, v, cache_k, cache_v, scale_k, scale_v, token, kv_head,
                                           physical_page, position & kPagedKVPageMask, lane);
}

template <typename Geometry, typename Metadata>
__launch_bounds__(256) __global__
    void kv_cache_append_full_fp8_page_kernel(const __nv_bfloat16* __restrict__ k,
                                              const __nv_bfloat16* __restrict__ v,
                                              const std::int32_t* __restrict__ positions,
                                              Metadata metadata, std::uint8_t* __restrict__ cache_k,
                                              std::uint8_t* __restrict__ cache_v,
                                              __half* __restrict__ scale_k,
                                              __half* __restrict__ scale_v, std::int32_t width) {
    constexpr int TokensPerTile = 8;
    constexpr unsigned FullMask = 0xffffffffU;
    const int tokens            = metadata.valid_tokens(width);
    const int warp              = static_cast<int>(threadIdx.x) >> 5;
    const int lane              = static_cast<int>(threadIdx.x) & 31;
    const int kv_head           = static_cast<int>(blockIdx.y);
    const int tile_delta        = static_cast<int>(blockIdx.x);
    const int base_position     = positions[0];
    const int tile_position     = (base_position / TokensPerTile + tile_delta) * TokensPerTile;
    const int logical_page      = tile_position >> kPagedKVPageShift;
    const int token_begin       = max(0, tile_position - base_position);
    const int token_end         = min(tokens, tile_position + TokensPerTile - base_position);
    if (token_begin >= token_end) return;

    const int token = token_begin + warp;
    if (token >= token_end) return;
    const std::int32_t* block_table = metadata.block_table();
    int physical_page               = lane == 0 ? block_table[logical_page] : 0;
    physical_page                   = __shfl_sync(FullMask, physical_page, 0);
    const int position              = base_position + token;
    kv_cache_append_full_fp8_row<Geometry>(k, v, cache_k, cache_v, scale_k, scale_v, token, kv_head,
                                           physical_page, position & kPagedKVPageMask, lane);
}

template <typename Geometry, typename Metadata>
__launch_bounds__(256) __global__
    void kv_cache_append_full_i8_kernel(const __nv_bfloat16* __restrict__ k,
                                        const __nv_bfloat16* __restrict__ v,
                                        const std::int32_t* __restrict__ positions,
                                        Metadata metadata, std::int8_t* __restrict__ cache_k,
                                        std::int8_t* __restrict__ cache_v,
                                        __half* __restrict__ scale_k, __half* __restrict__ scale_v,
                                        std::int32_t width) {
    constexpr int Warps         = 8;
    constexpr unsigned FullMask = 0xffffffffu;
    const int tokens            = metadata.valid_tokens(width);
    const int warp              = static_cast<int>(threadIdx.x) >> 5;
    const int lane              = static_cast<int>(threadIdx.x) & 31;
    const int unit              = static_cast<int>(blockIdx.x) * Warps + warp;
    const int units             = tokens * Geometry::KVHeads;
    if (unit >= units) return;

    const int kv_head               = unit % Geometry::KVHeads;
    const int token                 = unit / Geometry::KVHeads;
    const int position              = positions[0] + token;
    const std::int32_t* block_table = metadata.block_table();
    int page                        = lane == 0 ? paged_kv_physical_page(block_table, position) : 0;
    const int page_off              = position & kPagedKVPageMask;
    page                            = __shfl_sync(FullMask, page, 0);

    float k_values[8];
#pragma unroll
    for (int r = 0; r < 8; ++r) {
        const int d = lane + 32 * r;
        k_values[r] =
            __bfloat162float(k[kv_cache_int8_quant_src_index<Geometry>(kv_head, d, token)]);
    }
    normalized_hadamard_d256_inplace(k_values, lane);

#pragma unroll
    for (int group = 0; group < kKVCacheInt8Groups; ++group) {
        const int d0                 = group * kKVCacheInt8Group + lane;
        const int d1                 = d0 + 32;
        const float k0               = k_values[2 * group];
        const float k1               = k_values[2 * group + 1];
        const std::int64_t src0      = kv_cache_int8_quant_src_index<Geometry>(kv_head, d0, token);
        const std::int64_t src1      = kv_cache_int8_quant_src_index<Geometry>(kv_head, d1, token);
        const float v0               = __bfloat162float(v[src0]);
        const float v1               = __bfloat162float(v[src1]);
        const float k_abs            = warp_max(fmaxf(fabsf(k0), fabsf(k1)), FullMask);
        const float v_abs            = warp_max(fmaxf(fabsf(v0), fabsf(v1)), FullMask);
        const auto k_quant           = kv_cache_int8_quant_params(k_abs);
        const auto v_quant           = kv_cache_int8_quant_params(v_abs);
        const std::int64_t code_base = kv_cache_int8_quant_code_index<Geometry>(
            page, kv_head, group * kKVCacheInt8Group, page_off);
        cache_k[code_base + lane]      = kv_cache_int8_quant_code(k0, k_quant.inverse_scale);
        cache_k[code_base + lane + 32] = kv_cache_int8_quant_code(k1, k_quant.inverse_scale);
        cache_v[code_base + lane]      = kv_cache_int8_quant_code(v0, v_quant.inverse_scale);
        cache_v[code_base + lane + 32] = kv_cache_int8_quant_code(v1, v_quant.inverse_scale);
        if (lane == 0) {
            const std::int64_t scale_off =
                kv_cache_int8_quant_scale_index<Geometry>(page, kv_head, group, page_off);
            scale_k[scale_off] = k_quant.scale;
            scale_v[scale_off] = v_quant.scale;
        }
    }
}

template <typename Geometry, typename Metadata>
__launch_bounds__(256) __global__
    void kv_cache_append_full_i8_page_kernel(const __nv_bfloat16* __restrict__ k,
                                             const __nv_bfloat16* __restrict__ v,
                                             const std::int32_t* __restrict__ positions,
                                             Metadata metadata, std::int8_t* __restrict__ cache_k,
                                             std::int8_t* __restrict__ cache_v,
                                             __half* __restrict__ scale_k,
                                             __half* __restrict__ scale_v, std::int32_t width) {
    constexpr int TokensPerTile = 8;
    constexpr unsigned FullMask = 0xffffffffu;
    const int tokens            = metadata.valid_tokens(width);
    const int warp              = static_cast<int>(threadIdx.x) >> 5;
    const int lane              = static_cast<int>(threadIdx.x) & 31;
    const int kv_head           = static_cast<int>(blockIdx.y);
    const int tile_delta        = static_cast<int>(blockIdx.x);
    const int base_position     = positions[0];
    const int tile_position     = (base_position / TokensPerTile + tile_delta) * TokensPerTile;
    const int logical_page      = tile_position >> kPagedKVPageShift;
    const int token_begin       = max(0, tile_position - base_position);
    const int token_end         = min(tokens, tile_position + TokensPerTile - base_position);
    if (token_begin >= token_end) return;

    const int token = token_begin + warp;
    if (token >= token_end) return;

    const std::int32_t* block_table = metadata.block_table();
    int physical_page               = lane == 0 ? block_table[logical_page] : 0;
    physical_page                   = __shfl_sync(FullMask, physical_page, 0);

    const int position = base_position + token;
    const int page_off = position & kPagedKVPageMask;

    float k_values[8];
#pragma unroll
    for (int r = 0; r < 8; ++r) {
        const int d = lane + 32 * r;
        k_values[r] =
            __bfloat162float(k[kv_cache_int8_quant_src_index<Geometry>(kv_head, d, token)]);
    }
    normalized_hadamard_d256_inplace(k_values, lane);

#pragma unroll
    for (int group = 0; group < kKVCacheInt8Groups; ++group) {
        const int d0                 = group * kKVCacheInt8Group + lane;
        const int d1                 = d0 + 32;
        const float k0               = k_values[2 * group];
        const float k1               = k_values[2 * group + 1];
        const std::int64_t src0      = kv_cache_int8_quant_src_index<Geometry>(kv_head, d0, token);
        const std::int64_t src1      = kv_cache_int8_quant_src_index<Geometry>(kv_head, d1, token);
        const float v0               = __bfloat162float(v[src0]);
        const float v1               = __bfloat162float(v[src1]);
        const float k_abs            = warp_max(fmaxf(fabsf(k0), fabsf(k1)), FullMask);
        const float v_abs            = warp_max(fmaxf(fabsf(v0), fabsf(v1)), FullMask);
        const auto k_quant           = kv_cache_int8_quant_params(k_abs);
        const auto v_quant           = kv_cache_int8_quant_params(v_abs);
        const std::int64_t code_base = kv_cache_int8_quant_code_index<Geometry>(
            physical_page, kv_head, group * kKVCacheInt8Group, page_off);
        cache_k[code_base + lane]      = kv_cache_int8_quant_code(k0, k_quant.inverse_scale);
        cache_k[code_base + lane + 32] = kv_cache_int8_quant_code(k1, k_quant.inverse_scale);
        cache_v[code_base + lane]      = kv_cache_int8_quant_code(v0, v_quant.inverse_scale);
        cache_v[code_base + lane + 32] = kv_cache_int8_quant_code(v1, v_quant.inverse_scale);
        if (lane == 0) {
            const std::int64_t scale_offset =
                kv_cache_int8_quant_scale_index<Geometry>(physical_page, kv_head, group, page_off);
            scale_k[scale_offset] = k_quant.scale;
            scale_v[scale_offset] = v_quant.scale;
        }
    }
}

inline constexpr int kKVCacheAppendPrefixHeadDim = 128;
inline constexpr int kKVCacheAppendPrefixHeads   = 8;
inline constexpr int kKVCacheAppendPrefixWindow  = 4096;
inline constexpr int kKVCacheAppendPrefixPage    = 64;

__device__ __forceinline__ void kv_cache_append_prefix_copy_cyclic_unit(
    const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v,
    __nv_bfloat16* __restrict__ cache_k, __nv_bfloat16* __restrict__ cache_v, int token,
    int unit_in_token, int slot, int padded_capacity) {
    constexpr int Bf16PerUnit  = 16;
    constexpr int UnitsPerHead = kKVCacheAppendPrefixHeadDim / Bf16PerUnit;
    const int kv_head          = unit_in_token / UnitsPerHead;
    const int d                = (unit_in_token - kv_head * UnitsPerHead) * Bf16PerUnit;
    const std::int64_t src =
        static_cast<std::int64_t>(d) + static_cast<std::int64_t>(kKVCacheAppendPrefixHeadDim) *
                                           (kv_head + kKVCacheAppendPrefixHeads * token);
    const std::int64_t dst = static_cast<std::int64_t>(d) +
                             static_cast<std::int64_t>(kKVCacheAppendPrefixHeadDim) *
                                 (slot + static_cast<std::int64_t>(padded_capacity) * kv_head);

    const int4 k0                               = *reinterpret_cast<const int4*>(&k[src]);
    const int4 v0                               = *reinterpret_cast<const int4*>(&v[src]);
    *reinterpret_cast<int4*>(&cache_k[dst])     = k0;
    *reinterpret_cast<int4*>(&cache_v[dst])     = v0;
    const int4 k1                               = *reinterpret_cast<const int4*>(&k[src + 8]);
    const int4 v1                               = *reinterpret_cast<const int4*>(&v[src + 8]);
    *reinterpret_cast<int4*>(&cache_k[dst + 8]) = k1;
    *reinterpret_cast<int4*>(&cache_v[dst + 8]) = v1;
}

__device__ __forceinline__ void kv_cache_append_prefix_copy_paged_unit(
    const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v,
    __nv_bfloat16* __restrict__ cache_k, __nv_bfloat16* __restrict__ cache_v, int token,
    int unit_in_token, int page_offset, int physical_page, int physical_pages) {
    constexpr int Bf16PerUnit  = 16;
    constexpr int UnitsPerHead = kKVCacheAppendPrefixHeadDim / Bf16PerUnit;
    const int kv_head          = unit_in_token / UnitsPerHead;
    const int d                = (unit_in_token - kv_head * UnitsPerHead) * Bf16PerUnit;
    const std::int64_t src =
        static_cast<std::int64_t>(d) + static_cast<std::int64_t>(kKVCacheAppendPrefixHeadDim) *
                                           (kv_head + kKVCacheAppendPrefixHeads * token);
    const std::int64_t dst =
        static_cast<std::int64_t>(d) +
        static_cast<std::int64_t>(kKVCacheAppendPrefixHeadDim) *
            (page_offset + kKVCacheAppendPrefixPage * (physical_page + physical_pages * kv_head));

    const int4 k0                               = *reinterpret_cast<const int4*>(&k[src]);
    const int4 v0                               = *reinterpret_cast<const int4*>(&v[src]);
    *reinterpret_cast<int4*>(&cache_k[dst])     = k0;
    *reinterpret_cast<int4*>(&cache_v[dst])     = v0;
    const int4 k1                               = *reinterpret_cast<const int4*>(&k[src + 8]);
    const int4 v1                               = *reinterpret_cast<const int4*>(&v[src + 8]);
    *reinterpret_cast<int4*>(&cache_k[dst + 8]) = k1;
    *reinterpret_cast<int4*>(&cache_v[dst + 8]) = v1;
}

__global__ void kv_cache_append_prefix_cyclic_kernel(
    const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v,
    const std::int32_t* __restrict__ positions, const std::int32_t* __restrict__ counts,
    const std::int32_t* __restrict__ lanes, __nv_bfloat16* __restrict__ cache_k,
    __nv_bfloat16* __restrict__ cache_v, int min_count, int max_count, int width,
    int padded_capacity) {
    constexpr int UnitsPerToken  = kKVCacheAppendPrefixHeads * 8;
    constexpr int TokensPerBlock = 256 / UnitsPerToken;
    static_assert(TokensPerBlock * UnitsPerToken == 256);
    const int batch = static_cast<int>(blockIdx.y);
    const int count = counts[batch];
    if (count < min_count || count > max_count) return;

    constexpr std::int64_t ElementsPerToken =
        kKVCacheAppendPrefixHeadDim * kKVCacheAppendPrefixHeads;
    const std::int64_t input_offset = ElementsPerToken * width * batch;
    const std::int64_t cache_offset =
        ElementsPerToken * static_cast<std::int64_t>(padded_capacity) * lanes[batch];
    k += input_offset;
    v += input_offset;
    positions += static_cast<std::int64_t>(width) * batch;
    cache_k += cache_offset;
    cache_v += cache_offset;

    const int local         = static_cast<int>(threadIdx.x);
    const int local_token   = local / UnitsPerToken;
    const int unit_in_token = local - local_token * UnitsPerToken;
    const int token         = static_cast<int>(blockIdx.x) * TokensPerBlock + local_token;
    if (token >= count) return;
    const int position = positions[token];
    const int slot     = position & (kKVCacheAppendPrefixWindow - 1);
    kv_cache_append_prefix_copy_cyclic_unit(k, v, cache_k, cache_v, token, unit_in_token, slot,
                                            padded_capacity);
}

__global__ void kv_cache_append_prefix_paged_kernel(
    const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v,
    const std::int32_t* __restrict__ positions, const std::int32_t* __restrict__ counts,
    const std::int32_t* __restrict__ table_rows, __nv_bfloat16* __restrict__ cache_k,
    __nv_bfloat16* __restrict__ cache_v, const std::int32_t* __restrict__ block_tables,
    int physical_pages, int logical_pages, int min_count, int max_count, int width) {
    constexpr int UnitsPerToken  = kKVCacheAppendPrefixHeads * 8;
    constexpr int TokensPerBlock = 256 / UnitsPerToken;
    static_assert(TokensPerBlock * UnitsPerToken == 256);
    const int batch = static_cast<int>(blockIdx.y);
    const int count = counts[batch];
    if (count < min_count || count > max_count) return;

    constexpr std::int64_t ElementsPerToken =
        kKVCacheAppendPrefixHeadDim * kKVCacheAppendPrefixHeads;
    const std::int64_t input_offset = ElementsPerToken * width * batch;
    k += input_offset;
    v += input_offset;
    positions += static_cast<std::int64_t>(width) * batch;
    const std::int32_t* block_table =
        block_tables + static_cast<std::int64_t>(logical_pages) * table_rows[batch];

    const int local         = static_cast<int>(threadIdx.x);
    const int local_token   = local / UnitsPerToken;
    const int unit_in_token = local - local_token * UnitsPerToken;
    const int lane          = local & 31;
    const int token         = static_cast<int>(blockIdx.x) * TokensPerBlock + local_token;
    int position            = 0;
    int physical_page       = 0;
    if (lane == 0 && token < count) {
        position      = positions[token];
        physical_page = block_table[position >> 6];
    }
    position      = __shfl_sync(0xffffffffu, position, 0);
    physical_page = __shfl_sync(0xffffffffu, physical_page, 0);
    if (token < count) {
        kv_cache_append_prefix_copy_paged_unit(k, v, cache_k, cache_v, token, unit_in_token,
                                               position & (kKVCacheAppendPrefixPage - 1),
                                               physical_page, physical_pages);
    }
}

} // namespace ninfer::ops
#pragma once

#include <ninfer/targets/qwen3_6/startup_features.h>
#include <ninfer/targets/qwen3_6/vision.h>

#include "core/tensor.h"

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

namespace ninfer {

class DeviceArena;

namespace targets::qwen3_6 {

template <class ProjectionPayload, class PostMixerPayload>
struct FullAttentionWeights {
    Tensor input_norm;
    ProjectionPayload projection;
    Tensor query_norm;
    Tensor key_norm;
    Weight output;
    Tensor post_attention_norm;
    PostMixerPayload post_mixer;
};

template <class ProjectionPayload, class PostMixerPayload>
struct GdnWeights {
    Tensor input_norm;
    ProjectionPayload projection;
    Tensor convolution;
    Tensor norm;
    Weight output;
    Tensor post_attention_norm;
    PostMixerPayload post_mixer;
};

template <class AttentionPayload, class PostMixerPayload>
struct MtpWeights {
    Weight input_projection;
    Tensor embedding_norm;
    Tensor hidden_norm;
    Tensor input_norm;
    AttentionPayload attention;
    Tensor query_norm;
    Tensor key_norm;
    Weight output;
    Tensor post_attention_norm;
    PostMixerPayload post_mixer;
    Tensor final_norm;
};

struct OptimizedProposalWeights {
    Weight head;
    Tensor token_ids;
};

struct DFlashLayerWeights {
    Tensor input_norm;
    Weight query_key_value;
    Weight context_key;
    Weight context_value;
    Tensor query_norm;
    Tensor key_norm;
    Weight attention_output;
    Tensor post_attention_norm;
    Weight gate_up;
    Weight down;
};

template <std::size_t Layers>
struct DFlashWeights {
    Weight feature_projection;
    Tensor context_norm;
    std::array<DFlashLayerWeights, Layers> layers;
    Tensor final_norm;
};

template <class FullProjectionPayload, class GdnProjectionPayload, class MainPostMixerPayload,
          class MtpAttentionPayload, class MtpPostMixerPayload, class DFlashPayload,
          std::size_t FullAttentionLayers, std::size_t GdnLayers>
struct ModelView {
    using FullLayer = FullAttentionWeights<FullProjectionPayload, MainPostMixerPayload>;
    using GdnLayer  = GdnWeights<GdnProjectionPayload, MainPostMixerPayload>;
    using MtpLayer  = MtpWeights<MtpAttentionPayload, MtpPostMixerPayload>;
    using DFlash    = DFlashPayload;

    DeviceArena* weights_arena = nullptr;
    Weight token_embedding;
    std::array<FullLayer, FullAttentionLayers> full_layers;
    std::array<GdnLayer, GdnLayers> gdn_layers;
    Tensor final_norm;
    Weight output_head;
    StartupFeatures features;
    std::optional<OptimizedProposalWeights> optimized_proposal;
    std::optional<MtpLayer> mtp;
    std::optional<DFlashPayload> dflash;
    std::optional<VisionWeights> vision;
};

} // namespace targets::qwen3_6
} // namespace ninfer
#pragma once

#include <cuda_bf16.h>

#include <cstdint>

namespace ninfer::ops {

struct LayerNormMoments {
    float mean;
    float m2;
    int count;
};

__device__ __forceinline__ LayerNormMoments layer_norm_merge(LayerNormMoments a,
                                                             LayerNormMoments b) {
    if (b.count == 0) { return a; }
    if (a.count == 0) { return b; }
    const float delta = b.mean - a.mean;
    const int count   = a.count + b.count;
    const float ratio = static_cast<float>(b.count) / static_cast<float>(count);
    return LayerNormMoments{
        a.mean + delta * ratio,
        a.m2 + b.m2 +
            delta * delta *
                (static_cast<float>(a.count) * static_cast<float>(b.count) /
                 static_cast<float>(count)),
        count,
    };
}

__device__ __forceinline__ LayerNormMoments layer_norm_warp_reduce(LayerNormMoments value) {
    constexpr unsigned mask = 0xffffffffu;
    for (int offset = 16; offset > 0; offset >>= 1) {
        const LayerNormMoments other{
            __shfl_down_sync(mask, value.mean, offset),
            __shfl_down_sync(mask, value.m2, offset),
            __shfl_down_sync(mask, value.count, offset),
        };
        value = layer_norm_merge(value, other);
    }
    return value;
}

__device__ __forceinline__ void layer_norm_update(LayerNormMoments& moments, float value) {
    ++moments.count;
    const float delta = value - moments.mean;
    moments.mean += delta / static_cast<float>(moments.count);
    moments.m2 += delta * (value - moments.mean);
}

template <int Block>
__launch_bounds__(Block) __global__
    void layer_norm_d1152_bf16x2_kernel(const __nv_bfloat162* x, const __nv_bfloat162* weight,
                                        const __nv_bfloat162* bias, __nv_bfloat162* out,
                                        std::int64_t rows, float eps) {
    constexpr int d        = 1152;
    constexpr int pairs    = d / 2;
    constexpr int warps    = Block / 32;
    const std::int64_t row = static_cast<std::int64_t>(blockIdx.x);
    if (row >= rows) { return; }
    const std::int64_t base = row * pairs;

    LayerNormMoments local{0.0f, 0.0f, 0};
    for (int pair = static_cast<int>(threadIdx.x); pair < pairs; pair += Block) {
        const float2 values = __bfloat1622float2(x[base + pair]);
        layer_norm_update(local, values.x);
        layer_norm_update(local, values.y);
    }
    local = layer_norm_warp_reduce(local);

    __shared__ LayerNormMoments warp_moments[warps];
    __shared__ float final_mean;
    __shared__ float final_inv;
    const int lane = static_cast<int>(threadIdx.x) & 31;
    const int warp = static_cast<int>(threadIdx.x) >> 5;
    if (lane == 0) { warp_moments[warp] = local; }
    __syncthreads();

    if (warp == 0) {
        LayerNormMoments block_value =
            lane < warps ? warp_moments[lane] : LayerNormMoments{0.0f, 0.0f, 0};
        block_value = layer_norm_warp_reduce(block_value);
        if (lane == 0) {
            final_mean = block_value.mean;
            final_inv  = rsqrtf(block_value.m2 / static_cast<float>(d) + eps);
        }
    }
    __syncthreads();

    for (int pair = static_cast<int>(threadIdx.x); pair < pairs; pair += Block) {
        const float2 values = __bfloat1622float2(x[base + pair]);
        const float2 scales = __bfloat1622float2(weight[pair]);
        const float2 shifts = __bfloat1622float2(bias[pair]);
        out[base + pair] =
            __floats2bfloat162_rn((values.x - final_mean) * final_inv * scales.x + shifts.x,
                                  (values.y - final_mean) * final_inv * scales.y + shifts.y);
    }
}

template <int Block>
__launch_bounds__(Block) __global__
    void layer_norm_d1152_warp_kernel(const __nv_bfloat162* x, const __nv_bfloat162* weight,
                                      const __nv_bfloat162* bias, __nv_bfloat162* out,
                                      std::int64_t rows, float eps) {
    constexpr int d        = 1152;
    constexpr int pairs    = d / 2;
    constexpr int warps    = Block / 32;
    const int lane         = static_cast<int>(threadIdx.x) & 31;
    const int warp         = static_cast<int>(threadIdx.x) >> 5;
    const std::int64_t row = static_cast<std::int64_t>(blockIdx.x) * warps + warp;
    if (row >= rows) { return; }
    const std::int64_t base = row * pairs;

    LayerNormMoments local{0.0f, 0.0f, 0};
    for (int pair = lane; pair < pairs; pair += 32) {
        const float2 values = __bfloat1622float2(x[base + pair]);
        layer_norm_update(local, values.x);
        layer_norm_update(local, values.y);
    }
    local                   = layer_norm_warp_reduce(local);
    constexpr unsigned mask = 0xffffffffu;
    const float mean        = __shfl_sync(mask, local.mean, 0);
    const float m2          = __shfl_sync(mask, local.m2, 0);
    const float inv         = rsqrtf(m2 / static_cast<float>(d) + eps);

    for (int pair = lane; pair < pairs; pair += 32) {
        const float2 values = __bfloat1622float2(x[base + pair]);
        const float2 scales = __bfloat1622float2(weight[pair]);
        const float2 shifts = __bfloat1622float2(bias[pair]);
        out[base + pair]    = __floats2bfloat162_rn((values.x - mean) * inv * scales.x + shifts.x,
                                                    (values.y - mean) * inv * scales.y + shifts.y);
    }
}

template <int Block>
__global__ void layer_norm_kernel(const __nv_bfloat16* x, const __nv_bfloat16* weight,
                                  const __nv_bfloat16* bias, __nv_bfloat16* out, std::int32_t d,
                                  std::int64_t rows, float eps) {
    const std::int64_t row = static_cast<std::int64_t>(blockIdx.x);
    if (row >= rows) { return; }
