ferrum-kernels 0.8.2

//! CUDA providers for backend-neutral vNext operation contracts.

use std::collections::BTreeSet;
use std::sync::Arc;

use cudarc::driver::{CudaFunction, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use ferrum_interfaces::vnext::{
    causal_paged_attention_contract, dense_linear_contract, dense_swiglu_contract,
    gated_delta_recurrent_attention_contract, last_token_dense_linear_contract,
    last_token_masked_argmax_contract, residual_add_contract, rms_norm_contract,
    token_embedding_contract, AttributeId, BatchedOperationInvocation, CapabilityCatalog,
    CapabilityId, ContractVersion, DeviceBatchingForm, DeviceId,
    DeviceReusableExecutionTopologyFingerprint, DeviceRuntime, DynamicStorageAllocator,
    DynamicStorageProfile, DynamicStorageRequirement, DynamicStorageView, ElementType,
    EncodedDeviceOperation, EngineProviderDescriptor, OperationContract, OperationFailure,
    OperationInvocation, OperationProvider, OperationProviderDescriptor, OperationResourceEstimate,
    OperationResourceEstimateRequest, OperationResourceEstimator, OperationRuntimeRegistry,
    ProfilePhase, ProviderId, ProviderStorageBindingRequirement, ProviderWorkspaceRequirement,
    ProviderWorkspaceReusePolicy, ProviderWorkspaceScope, ProviderWorkspaceSizeFormula,
    ResolvedTensorLayout, ResolvedValueBinding, ResolvedValueRole, ReusableExecutionTopology,
    ReusableExecutionTopologyRequest, SemanticValue, VNextError, WeightFormatId,
    WeightMaterializerId, WeightMaterializerRegistry, CAUSAL_PAGED_ATTENTION_F16_CAPABILITY_ID,
    DENSE_LINEAR_F16_CAPABILITY_ID, DENSE_SWIGLU_F16_CAPABILITY_ID,
    DEVICE_NATIVE_ADAPTIVE_ATTENTION_CAPABILITY_ID, DEVICE_REUSABLE_EXECUTION_CAPABILITY_ID,
    GATED_DELTA_RECURRENT_ATTENTION_F16_CAPABILITY_ID, IDENTITY_WEIGHT_MATERIALIZER_ID,
    LAST_TOKEN_DENSE_LINEAR_F16_CAPABILITY_ID, LAST_TOKEN_DENSE_LINEAR_OPERATION_ID,
    LAST_TOKEN_MASKED_ARGMAX_F16_CAPABILITY_ID, LAST_TOKEN_MASKED_ARGMAX_OPERATION_ID,
    RESIDUAL_ADD_F16_CAPABILITY_ID, RMS_NORM_F16_CAPABILITY_ID, TOKEN_EMBEDDING_F16_CAPABILITY_ID,
    TOKEN_EMBEDDING_OPERATION_ID,
};
#[cfg(feature = "vllm-moe-marlin")]
use ferrum_interfaces::vnext::{
    routed_shared_swiglu_moe_contract, routed_swiglu_moe_contract,
    ROUTED_SHARED_SWIGLU_MOE_F16_CAPABILITY_ID, ROUTED_SWIGLU_MOE_F16_CAPABILITY_ID,
};
use ferrum_types::{
    AttentionExecutionPolicy, NativeOperatorBackend, NativeOperatorProviderCatalog,
};
use sha2::{Digest, Sha256};

use super::vnext_replay::CudaCommandReplayKeyBuilder;
use super::vnext_runtime::{
    CudaBufferRegion, CudaDeviceBuffer, CudaDeviceCommand, CudaDeviceRuntime,
    CudaDeviceRuntimeConfig, CudaDeviceRuntimeError,
};

mod transformer;

const TOKEN_EMBEDDING_PROVIDER_ID: &str = "provider.cuda.token_embedding.f16";
const TOKEN_EMBEDDING_ESTIMATOR_ID: &str = "resource-estimator.cuda.token_embedding.f16";
const LAST_TOKEN_DENSE_LINEAR_PROVIDER_ID: &str =
    "provider.cuda.last_token_dense_linear.f16.cublas";
const LAST_TOKEN_DENSE_LINEAR_ESTIMATOR_ID: &str =
    "resource-estimator.cuda.last_token_dense_linear.f16.cublas";
const LAST_TOKEN_MASKED_ARGMAX_PROVIDER_ID: &str = "provider.cuda.last_token_masked_argmax.f16";
const LAST_TOKEN_MASKED_ARGMAX_ESTIMATOR_ID: &str =
    "resource-estimator.cuda.last_token_masked_argmax.f16";
const CUDA_ENGINE_PROVIDER_ID: &str = "provider.engine.cuda.vnext";
const DENSE_SAFETENSORS_FORMAT_ID: &str = "weight-format.safetensors.dense";
const EMBEDDING_FUNCTION_NAME: &str = "vnext_embedding_lookup_f16";
const MASKED_ARGMAX_PRESERVING_LOGITS_FUNCTION_NAME: &str =
    "last_token_masked_argmax_preserving_logits_f16";
const VALUE_ALIGNMENT_BYTES: u64 = 16;
const THREADS_PER_BLOCK: u32 = 256;
const MAXIMUM_TOKENS_PER_LAUNCH: u64 = u16::MAX as u64;
pub(super) const VNEXT_KV_PAGE_BYTES: u64 = 64 * 1024;

/// Typed CUDA runtime input for the currently installed vNext provider bundle.
pub fn cuda_vnext_runtime_config(
    ordinal: usize,
    device_id: DeviceId,
    requested_attention_policy: AttentionExecutionPolicy,
) -> Result<CudaDeviceRuntimeConfig, VNextError> {
    let fingerprint_parts: Vec<&[u8]> = vec![
        include_str!("vnext_runtime.rs").as_bytes(),
        include_str!("vnext_replay.rs").as_bytes(),
        include_str!("vnext_ops.rs").as_bytes(),
        include_str!("vnext_ops/transformer.rs").as_bytes(),
        include_str!("vnext_ops/transformer/attention.rs").as_bytes(),
        include_str!("vnext_ops/transformer/causal_attention.rs").as_bytes(),
        crate::ptx::EMBEDDING_LOOKUP.as_bytes(),
        crate::ptx::ARGMAX_ROWS.as_bytes(),
        crate::ptx::RMS_NORM.as_bytes(),
        crate::ptx::FUSED_SILU_MUL.as_bytes(),
        crate::ptx::RESIDUAL_ADD.as_bytes(),
        crate::ptx::SANDWICH_NORM.as_bytes(),
        crate::ptx::LINEAR_ATTENTION.as_bytes(),
        crate::ptx::GATED_DELTA_RULE.as_bytes(),
        crate::ptx::VNEXT_CAUSAL_ATTENTION.as_bytes(),
    ];
    #[cfg(feature = "vllm-moe-marlin")]
    let fingerprint_parts = {
        let mut fingerprint_parts = fingerprint_parts;
        fingerprint_parts.extend([
            include_str!("vnext_ops/transformer/moe.rs").as_bytes(),
            include_str!("vnext_ops/transformer/moe_launch.rs").as_bytes(),
            include_str!("vnext_ops/transformer/moe_routed.rs").as_bytes(),
            include_str!("vnext_ops/transformer/moe_weights.rs").as_bytes(),
            include_str!("vnext_ops/transformer/moe_workspace.rs").as_bytes(),
            crate::ptx::MOE_ROUTER.as_bytes(),
            crate::ptx::MOE_ALIGN_BLOCK_SIZE_PAIR_IDS.as_bytes(),
            crate::ptx::MOE_COMBINE.as_bytes(),
        ]);
        fingerprint_parts
    };
    #[cfg(feature = "vllm-marlin")]
    let fingerprint_parts = {
        let mut fingerprint_parts = fingerprint_parts;
        fingerprint_parts.push(include_str!("../../marlin_fp8_materializer.rs").as_bytes());
        fingerprint_parts
    };
    let capabilities = cuda_vnext_capabilities()?;
    let attention_execution_policy = requested_attention_policy
        .resolve(capabilities.iter().any(|capability| {
            capability.as_str() == DEVICE_NATIVE_ADAPTIVE_ATTENTION_CAPABILITY_ID
        }))
        .map_err(|reason| VNextError::InvalidExecutionPlan { reason })?;
    Ok(CudaDeviceRuntimeConfig {
        ordinal,
        device_id,
        attention_execution_policy,
        runtime_implementation_fingerprint: implementation_fingerprint(&fingerprint_parts),
        capabilities,
        dynamic_storage_profiles: BTreeSet::from([
            DynamicStorageProfile::new(
                DynamicStorageAllocator::LinearArena,
                DynamicStorageView::Contiguous,
            )?,
            DynamicStorageProfile::new(
                DynamicStorageAllocator::FixedBlockArena {
                    block_bytes: VNEXT_KV_PAGE_BYTES,
                },
                DynamicStorageView::PagedRegions {
                    block_bytes: VNEXT_KV_PAGE_BYTES,
                },
            )?,
        ]),
    })
}

pub fn cuda_vnext_capabilities() -> Result<BTreeSet<CapabilityId>, VNextError> {
    let capabilities = [
        TOKEN_EMBEDDING_F16_CAPABILITY_ID,
        LAST_TOKEN_DENSE_LINEAR_F16_CAPABILITY_ID,
        LAST_TOKEN_MASKED_ARGMAX_F16_CAPABILITY_ID,
        RMS_NORM_F16_CAPABILITY_ID,
        DENSE_LINEAR_F16_CAPABILITY_ID,
        DENSE_SWIGLU_F16_CAPABILITY_ID,
        RESIDUAL_ADD_F16_CAPABILITY_ID,
        GATED_DELTA_RECURRENT_ATTENTION_F16_CAPABILITY_ID,
        CAUSAL_PAGED_ATTENTION_F16_CAPABILITY_ID,
        DEVICE_REUSABLE_EXECUTION_CAPABILITY_ID,
    ]
    .into_iter()
    .map(CapabilityId::new)
    .collect::<Result<BTreeSet<_>, _>>()?;
    #[cfg(feature = "vllm-moe-marlin")]
    let capabilities = {
        let mut capabilities = capabilities;
        capabilities.insert(CapabilityId::new(
            ROUTED_SHARED_SWIGLU_MOE_F16_CAPABILITY_ID,
        )?);
        capabilities.insert(CapabilityId::new(ROUTED_SWIGLU_MOE_F16_CAPABILITY_ID)?);
        capabilities
    };
    #[cfg(feature = "vllm-marlin")]
    let capabilities = {
        let mut capabilities = capabilities;
        capabilities.insert(CapabilityId::new(
            crate::marlin_fp8_materializer::MARLIN_FP8_CAPABILITY_ID,
        )?);
        capabilities.insert(CapabilityId::new(transformer::GPTQ_MARLIN_CAPABILITY_ID)?);
        capabilities.insert(CapabilityId::new(
            transformer::COMPRESSED_TENSORS_MARLIN_CAPABILITY_ID,
        )?);
        capabilities
    };
    #[cfg(feature = "vllm-paged-attn-v2")]
    let capabilities = {
        let mut capabilities = capabilities;
        capabilities.insert(CapabilityId::new(
            DEVICE_NATIVE_ADAPTIVE_ATTENTION_CAPABILITY_ID,
        )?);
        capabilities
    };
    Ok(capabilities)
}

/// Build the exact composition root used for both planning and dispatch.
pub fn cuda_vnext_operation_registry(
    runtime: &CudaDeviceRuntime,
) -> Result<OperationRuntimeRegistry<CudaDeviceRuntime>, CudaDeviceRuntimeError> {
    let contracts: Vec<Box<dyn OperationContract>> = vec![
        Box::new(token_embedding_contract().map_err(contract_error)?),
        Box::new(last_token_dense_linear_contract().map_err(contract_error)?),
        Box::new(last_token_masked_argmax_contract().map_err(contract_error)?),
        Box::new(rms_norm_contract().map_err(contract_error)?),
        Box::new(dense_linear_contract().map_err(contract_error)?),
        Box::new(dense_swiglu_contract().map_err(contract_error)?),
        Box::new(residual_add_contract().map_err(contract_error)?),
        Box::new(gated_delta_recurrent_attention_contract().map_err(contract_error)?),
        Box::new(causal_paged_attention_contract().map_err(contract_error)?),
    ];
    #[cfg(feature = "vllm-moe-marlin")]
    let contracts = {
        let mut contracts = contracts;
        contracts.push(Box::new(
            routed_shared_swiglu_moe_contract().map_err(contract_error)?,
        ));
        contracts.push(Box::new(
            routed_swiglu_moe_contract().map_err(contract_error)?,
        ));
        contracts
    };
    let providers: Vec<Box<dyn OperationProvider<CudaDeviceRuntime>>> = vec![
        Box::new(CudaTokenEmbeddingProvider::new(runtime)?),
        Box::new(CudaLastTokenDenseLinearProvider::new(runtime)?),
        Box::new(CudaLastTokenMaskedArgmaxProvider::new(runtime)?),
        Box::new(transformer::CudaRmsNormProvider::new(runtime)?),
        Box::new(transformer::CudaDenseLinearProvider::new(runtime)?),
        Box::new(transformer::CudaDenseSwiGluProvider::new(runtime)?),
        Box::new(transformer::CudaResidualAddProvider::new(runtime)?),
        Box::new(transformer::CudaGatedDeltaRecurrentAttentionProvider::new(
            runtime,
        )?),
        Box::new(transformer::CudaCausalPagedAttentionProvider::new(
            runtime,
            runtime.attention_execution_policy(),
        )?),
    ];
    #[cfg(feature = "vllm-marlin")]
    let providers = {
        let mut providers = providers;
        providers.push(Box::new(
            transformer::CudaMarlinFp8DenseLinearProvider::new(runtime)?,
        ));
        providers
    };
    #[cfg(feature = "vllm-moe-marlin")]
    let providers = {
        let mut providers = providers;
        providers.push(Box::new(
            transformer::CudaRoutedSharedSwiGluMoeProvider::new(runtime)?,
        ));
        providers.push(Box::new(transformer::CudaRoutedSwiGluMoeProvider::new(
            runtime,
        )?));
        providers
    };
    OperationRuntimeRegistry::new(contracts, providers).map_err(contract_error)
}

/// One CUDA composition root shared by planning, provisioning, and dispatch.
/// The capability catalog is derived from the retained registry objects rather
/// than reconstructed from a second descriptor list.
pub struct CudaVNextComposition {
    runtime: Arc<CudaDeviceRuntime>,
    registry: OperationRuntimeRegistry<CudaDeviceRuntime>,
    weight_materializers: WeightMaterializerRegistry,
    weight_materializer_id: WeightMaterializerId,
    catalog: CapabilityCatalog,
}

impl CudaVNextComposition {
    fn prepare(
        ordinal: usize,
        device_id: DeviceId,
        requested_attention_policy: AttentionExecutionPolicy,
    ) -> Result<Self, CudaDeviceRuntimeError> {
        let config = cuda_vnext_runtime_config(ordinal, device_id, requested_attention_policy)
            .map_err(contract_error)?;
        let runtime = Arc::new(CudaDeviceRuntime::new(config)?);
        let registry = cuda_vnext_operation_registry(&runtime)?;
        #[cfg(feature = "vllm-marlin")]
        let weight_materializers = WeightMaterializerRegistry::new(vec![
            crate::marlin_fp8_materializer::marlin_fp8_weight_materializer()
                .map_err(contract_error)?,
        ])
        .map_err(contract_error)?;
        #[cfg(not(feature = "vllm-marlin"))]
        let weight_materializers =
            WeightMaterializerRegistry::identity_only().map_err(contract_error)?;
        let weight_materializer_id =
            WeightMaterializerId::new(IDENTITY_WEIGHT_MATERIALIZER_ID).map_err(contract_error)?;
        let engine = EngineProviderDescriptor::new(
            ProviderId::new(CUDA_ENGINE_PROVIDER_ID).map_err(contract_error)?,
            ContractVersion::new(1, 0),
            implementation_fingerprint(&[
                include_str!("vnext_ops.rs").as_bytes(),
                include_str!("vnext_runtime.rs").as_bytes(),
                CUDA_ENGINE_PROVIDER_ID.as_bytes(),
            ]),
            runtime.descriptor().id.clone(),
            runtime.descriptor().capabilities.clone(),
        )
        .map_err(contract_error)?;
        let catalog = registry
            .capability_catalog(runtime.descriptor().clone(), vec![engine])
            .map_err(contract_error)?;
        let catalog = weight_materializers
            .augment_catalog(catalog)
            .map_err(contract_error)?;
        Ok(Self {
            runtime,
            registry,
            weight_materializers,
            weight_materializer_id,
            catalog,
        })
    }

    pub fn create(
        ordinal: usize,
        device_id: DeviceId,
        requested_attention_policy: AttentionExecutionPolicy,
    ) -> Result<Self, CudaDeviceRuntimeError> {
        let composition = Self::prepare(ordinal, device_id, requested_attention_policy)?;
        composition.validate_compiled_native_operators()?;
        Ok(composition)
    }

    fn validate_compiled_native_operators(&self) -> Result<(), CudaDeviceRuntimeError> {
        let catalog = &self.catalog;
        let native_provider_catalog = catalog
            .native_operator_provider_catalog(NativeOperatorBackend::Cuda)
            .map_err(contract_error)?;
        crate::native_ops::validate_compiled_native_operator_provider_catalog(
            &native_provider_catalog,
            crate::native_ops::compiled_native_operator_artifacts(),
        )
        .map_err(CudaDeviceRuntimeError::contract)
    }

    pub fn runtime(&self) -> &Arc<CudaDeviceRuntime> {
        &self.runtime
    }

    pub fn registry(&self) -> &OperationRuntimeRegistry<CudaDeviceRuntime> {
        &self.registry
    }

    pub fn catalog(&self) -> &CapabilityCatalog {
        &self.catalog
    }

    pub fn into_parts(
        self,
    ) -> (
        Arc<CudaDeviceRuntime>,
        OperationRuntimeRegistry<CudaDeviceRuntime>,
        WeightMaterializerRegistry,
        WeightMaterializerId,
        CapabilityCatalog,
    ) {
        (
            self.runtime,
            self.registry,
            self.weight_materializers,
            self.weight_materializer_id,
            self.catalog,
        )
    }
}

/// Immutable catalog input used to rebuild native artifacts after provider
/// identity changes. It deliberately exposes no executable runtime.
pub struct CudaNativeOperatorCatalogInput {
    provider_catalog: NativeOperatorProviderCatalog,
    capability_catalog: CapabilityCatalog,
}

impl CudaNativeOperatorCatalogInput {
    pub fn provider_catalog(&self) -> &NativeOperatorProviderCatalog {
        &self.provider_catalog
    }

    pub fn capability_catalog(&self) -> &CapabilityCatalog {
        &self.capability_catalog
    }

    pub fn into_parts(self) -> (NativeOperatorProviderCatalog, CapabilityCatalog) {
        (self.provider_catalog, self.capability_catalog)
    }
}

/// Capture the exact provider identities needed to package a new native
/// artifact set. Product composition still uses [`CudaVNextComposition::create`]
/// and cannot bypass compiled-artifact validation.
pub fn cuda_native_operator_catalog_input(
    ordinal: usize,
    device_id: DeviceId,
    requested_attention_policy: AttentionExecutionPolicy,
) -> Result<CudaNativeOperatorCatalogInput, CudaDeviceRuntimeError> {
    let composition =
        CudaVNextComposition::prepare(ordinal, device_id, requested_attention_policy)?;
    let provider_catalog = composition
        .catalog
        .native_operator_provider_catalog(NativeOperatorBackend::Cuda)
        .map_err(contract_error)?;
    Ok(CudaNativeOperatorCatalogInput {
        provider_catalog,
        capability_catalog: composition.catalog,
    })
}

pub struct CudaTokenEmbeddingProvider {
    descriptor: OperationProviderDescriptor,
    function: CudaFunction,
}

impl CudaTokenEmbeddingProvider {
    pub fn new(runtime: &CudaDeviceRuntime) -> Result<Self, CudaDeviceRuntimeError> {
        let contract = token_embedding_contract().map_err(contract_error)?;
        let capability =
            CapabilityId::new(TOKEN_EMBEDDING_F16_CAPABILITY_ID).map_err(contract_error)?;
        if !runtime.descriptor().capabilities.contains(&capability) {
            return Err(CudaDeviceRuntimeError::contract(
                "CUDA runtime does not advertise the token embedding capability",
            ));
        }

        let provider_fingerprint = implementation_fingerprint(&[
            include_str!("vnext_ops.rs").as_bytes(),
            crate::ptx::EMBEDDING_LOOKUP.as_bytes(),
            EMBEDDING_FUNCTION_NAME.as_bytes(),
        ]);
        let estimator_fingerprint = implementation_fingerprint(&[
            include_str!("vnext_ops.rs").as_bytes(),
            TOKEN_EMBEDDING_ESTIMATOR_ID.as_bytes(),
        ]);
        let descriptor = OperationProviderDescriptor::new(
            ProviderId::new(TOKEN_EMBEDDING_PROVIDER_ID).map_err(contract_error)?,
            contract.descriptor().id.clone(),
            contract
                .descriptor()
                .fingerprint()
                .map_err(contract_error)?,
            provider_fingerprint,
            ferrum_interfaces::vnext::ProviderExecutionSemantics::bitwise_eager_and_replay(),
            contract.descriptor().version,
            runtime.descriptor().id.clone(),
            BTreeSet::from([capability]),
            BTreeSet::from([
                WeightFormatId::new(DENSE_SAFETENSORS_FORMAT_ID).map_err(contract_error)?
            ]),
            BTreeSet::new(),
            contiguous_bindings(),
            TOKEN_EMBEDDING_ESTIMATOR_ID,
            ContractVersion::new(1, 0),
            estimator_fingerprint,
        )
        .map_err(contract_error)?;
        let module = runtime
            .context()
            .load_module(Ptx::from_src(crate::ptx::EMBEDDING_LOOKUP.to_owned()))
            .map_err(|error| CudaDeviceRuntimeError::driver("embedding module load", error))?;
        let function = module
            .load_function(EMBEDDING_FUNCTION_NAME)
            .map_err(|error| CudaDeviceRuntimeError::driver("embedding function load", error))?;
        Ok(Self {
            descriptor,
            function,
        })
    }
}

impl OperationResourceEstimator for CudaTokenEmbeddingProvider {
    fn descriptor(&self) -> &OperationProviderDescriptor {
        &self.descriptor
    }

    fn estimate_resources(
        &self,
        request: OperationResourceEstimateRequest<'_>,
    ) -> Result<OperationResourceEstimate, VNextError> {
        if request.operation().id.as_str() != TOKEN_EMBEDDING_OPERATION_ID
            || request.operation().fingerprint()? != self.descriptor.operation_fingerprint()
        {
            return Err(VNextError::InvalidExecutionPlan {
                reason: "CUDA token embedding estimator received another operation".to_owned(),
            });
        }
        Ok(OperationResourceEstimate::new(
            self.descriptor.resource_estimator_id(),
            self.descriptor.resource_estimator_version(),
            self.descriptor
                .resource_estimator_implementation_fingerprint(),
            request.input_fingerprint(),
            VALUE_ALIGNMENT_BYTES,
            None,
            None,
        ))
    }
}

impl OperationProvider<CudaDeviceRuntime> for CudaTokenEmbeddingProvider {
    fn reusable_execution_topology(
        &self,
        request: ReusableExecutionTopologyRequest<'_>,
    ) -> Result<ReusableExecutionTopology, VNextError> {
        reusable_token_topology(
            &request,
            b"ferrum.cuda.token-embedding.reusable-topology.v2\0",
        )
    }

    fn encode_selected(
        &self,
        invocation: BatchedOperationInvocation<'_, CudaDeviceBuffer>,
    ) -> Result<EncodedDeviceOperation<CudaDeviceCommand>, OperationFailure> {
        let identity = invocation.participants()[0].identity().clone();
        encode_token_embedding(
            &self.function,
            self.descriptor.provider_implementation_fingerprint(),
            invocation,
        )
        .map(EncodedDeviceOperation::compute)
        .map_err(|message| provider_failure(identity, "cuda.token_embedding.encode", message))
    }
}

pub struct CudaLastTokenDenseLinearProvider {
    descriptor: OperationProviderDescriptor,
}

impl CudaLastTokenDenseLinearProvider {
    pub fn new(runtime: &CudaDeviceRuntime) -> Result<Self, CudaDeviceRuntimeError> {
        let contract = last_token_dense_linear_contract().map_err(contract_error)?;
        let descriptor = transformer::provider_descriptor(
            runtime,
            &contract,
            LAST_TOKEN_DENSE_LINEAR_PROVIDER_ID,
            LAST_TOKEN_DENSE_LINEAR_F16_CAPABILITY_ID,
            LAST_TOKEN_DENSE_LINEAR_ESTIMATOR_ID,
            transformer::contiguous_bindings(2),
            implementation_fingerprint(&[
                include_str!("vnext_ops.rs").as_bytes(),
                LAST_TOKEN_DENSE_LINEAR_PROVIDER_ID.as_bytes(),
            ]),
        )?;
        Ok(Self { descriptor })
    }
}

impl OperationResourceEstimator for CudaLastTokenDenseLinearProvider {
    fn descriptor(&self) -> &OperationProviderDescriptor {
        &self.descriptor
    }

    fn estimate_resources(
        &self,
        request: OperationResourceEstimateRequest<'_>,
    ) -> Result<OperationResourceEstimate, VNextError> {
        transformer::ensure_estimator_request(
            &self.descriptor,
            &request,
            LAST_TOKEN_DENSE_LINEAR_OPERATION_ID,
        )?;
        Ok(transformer::estimate(
            &self.descriptor,
            request.input_fingerprint(),
            None,
        ))
    }
}

impl OperationProvider<CudaDeviceRuntime> for CudaLastTokenDenseLinearProvider {
    fn reusable_execution_topology(
        &self,
        request: ReusableExecutionTopologyRequest<'_>,
    ) -> Result<ReusableExecutionTopology, VNextError> {
        reusable_token_topology(
            &request,
            b"ferrum.cuda.last-token-linear.reusable-topology.v2\0",
        )
    }

    fn encode_selected(
        &self,
        invocation: BatchedOperationInvocation<'_, CudaDeviceBuffer>,
    ) -> Result<EncodedDeviceOperation<CudaDeviceCommand>, OperationFailure> {
        let identity = invocation.participants()[0].identity().clone();
        encode_last_token_dense_linear(
            self.descriptor.provider_implementation_fingerprint(),
            invocation,
        )
        .map(EncodedDeviceOperation::compute)
        .map_err(|message| {
            provider_failure(identity, "cuda.last_token_dense_linear.encode", message)
        })
    }
}

pub struct CudaLastTokenMaskedArgmaxProvider {
    descriptor: OperationProviderDescriptor,
    function: CudaFunction,
}

impl CudaLastTokenMaskedArgmaxProvider {
    pub fn new(runtime: &CudaDeviceRuntime) -> Result<Self, CudaDeviceRuntimeError> {
        let contract = last_token_masked_argmax_contract().map_err(contract_error)?;
        let descriptor = transformer::provider_descriptor(
            runtime,
            &contract,
            LAST_TOKEN_MASKED_ARGMAX_PROVIDER_ID,
            LAST_TOKEN_MASKED_ARGMAX_F16_CAPABILITY_ID,
            LAST_TOKEN_MASKED_ARGMAX_ESTIMATOR_ID,
            transformer::contiguous_bindings(5),
            implementation_fingerprint(&[
                include_str!("vnext_ops.rs").as_bytes(),
                crate::ptx::ARGMAX_ROWS.as_bytes(),
                MASKED_ARGMAX_PRESERVING_LOGITS_FUNCTION_NAME.as_bytes(),
            ]),
        )?;
        let module = runtime
            .context()
            .load_module(Ptx::from_src(crate::ptx::ARGMAX_ROWS.to_owned()))
            .map_err(|error| CudaDeviceRuntimeError::driver("masked argmax module load", error))?;
        let function = module
            .load_function(MASKED_ARGMAX_PRESERVING_LOGITS_FUNCTION_NAME)
            .map_err(|error| {
                CudaDeviceRuntimeError::driver(
                    "masked argmax preserving-logits function load",
                    error,
                )
            })?;
        Ok(Self {
            descriptor,
            function,
        })
    }
}

impl OperationResourceEstimator for CudaLastTokenMaskedArgmaxProvider {
    fn descriptor(&self) -> &OperationProviderDescriptor {
        &self.descriptor
    }

    fn estimate_resources(
        &self,
        request: OperationResourceEstimateRequest<'_>,
    ) -> Result<OperationResourceEstimate, VNextError> {
        transformer::ensure_estimator_request(
            &self.descriptor,
            &request,
            LAST_TOKEN_MASKED_ARGMAX_OPERATION_ID,
        )?;
        let vocabulary_size = unsigned_attribute(request.attributes(), "vocab_size")
            .map_err(|reason| VNextError::InvalidExecutionPlan { reason })?;
        let scratch_bytes = masked_argmax_scratch_stride(vocabulary_size)
            .map_err(|reason| VNextError::InvalidExecutionPlan { reason })?;
        let scratch = ProviderWorkspaceRequirement::from_formula(
            ProviderWorkspaceSizeFormula::actual_sequences(scratch_bytes)?,
            VALUE_ALIGNMENT_BYTES,
            ProviderWorkspaceScope::Invocation,
            ProviderWorkspaceReusePolicy::OverwriteBeforeRead,
            DynamicStorageRequirement::contiguous(),
        )?;
        Ok(transformer::estimate(
            &self.descriptor,
            request.input_fingerprint(),
            Some(scratch),
        ))
    }
}

impl OperationProvider<CudaDeviceRuntime> for CudaLastTokenMaskedArgmaxProvider {
    fn reusable_execution_topology(
        &self,
        request: ReusableExecutionTopologyRequest<'_>,
    ) -> Result<ReusableExecutionTopology, VNextError> {
        transformer::static_contiguous_reusable_topology(
            &request,
            5,
            &[transformer::CapturedProviderWorkspace::Scratch],
        )
    }

    fn encode_selected(
        &self,
        invocation: BatchedOperationInvocation<'_, CudaDeviceBuffer>,
    ) -> Result<EncodedDeviceOperation<CudaDeviceCommand>, OperationFailure> {
        let identity = invocation.participants()[0].identity().clone();
        encode_last_token_masked_argmax(
            &self.function,
            self.descriptor.provider_implementation_fingerprint(),
            invocation,
        )
        .map(EncodedDeviceOperation::compute)
        .map_err(|message| {
            provider_failure(identity, "cuda.last_token_masked_argmax.encode", message)
        })
    }
}

#[derive(Debug, Clone, Copy)]
struct MaskedArgmaxLaunch {
    first_region: usize,
    scratch_offset_bytes: u64,
    vocabulary_size: i32,
    repetition_capacity: i32,
}

fn encode_last_token_masked_argmax(
    function: &CudaFunction,
    provider_fingerprint: &str,
    invocation: BatchedOperationInvocation<'_, CudaDeviceBuffer>,
) -> Result<CudaDeviceCommand, String> {
    if invocation.operation().id.as_str() != LAST_TOKEN_MASKED_ARGMAX_OPERATION_ID
        || invocation.participants().is_empty()
    {
        return Err("CUDA masked argmax received another or empty operation".to_owned());
    }

    let first_vocabulary_size =
        unsigned_attribute(invocation.participants()[0].attributes(), "vocab_size")?;
    let scratch_stride = masked_argmax_scratch_stride(first_vocabulary_size)?;
    let required_scratch_bytes = scratch_stride
        .checked_mul(invocation.participants().len() as u64)
        .ok_or_else(|| "CUDA masked argmax scratch size overflows".to_owned())?;
    let mut regions = Vec::with_capacity(invocation.participants().len() * 6 + 1);
    let mut launches = Vec::with_capacity(invocation.participants().len());
    for (participant_index, participant) in invocation.participants().iter().enumerate() {
        let logits = binding(participant.bindings(), ResolvedValueRole::Input, 0)?;
        let valid_mask = binding(participant.bindings(), ResolvedValueRole::Input, 1)?;
        let repetition_token_ids = binding(participant.bindings(), ResolvedValueRole::Input, 2)?;
        let repetition_offsets = binding(participant.bindings(), ResolvedValueRole::Input, 3)?;
        let repetition_penalty = binding(participant.bindings(), ResolvedValueRole::Input, 4)?;
        let output = binding(participant.bindings(), ResolvedValueRole::Output, 0)?;
        let vocabulary_size = unsigned_attribute(participant.attributes(), "vocab_size")?;
        if vocabulary_size != first_vocabulary_size {
            return Err("CUDA masked argmax participants disagree on vocabulary size".to_owned());
        }
        let repetition_capacity = validate_masked_argmax_signature(
            logits,
            valid_mask,
            repetition_token_ids,
            repetition_offsets,
            repetition_penalty,
            output,
            vocabulary_size,
        )?;

        let first_region = regions.len();
        regions.push(contiguous_region(participant, logits, ElementType::F16)?);
        regions.push(contiguous_region(participant, valid_mask, ElementType::U8)?);
        regions.push(contiguous_region(
            participant,
            repetition_token_ids,
            ElementType::U32,
        )?);
        regions.push(contiguous_region(
            participant,
            repetition_offsets,
            ElementType::U32,
        )?);
        regions.push(contiguous_region(
            participant,
            repetition_penalty,
            ElementType::F32,
        )?);
        regions.push(contiguous_region(participant, output, ElementType::U32)?);
        launches.push(MaskedArgmaxLaunch {
            first_region,
            scratch_offset_bytes: scratch_stride
                .checked_mul(participant_index as u64)
                .ok_or_else(|| "CUDA masked argmax scratch offset overflows".to_owned())?,
            vocabulary_size: i32::try_from(vocabulary_size)
                .map_err(|_| "masked argmax vocabulary exceeds i32".to_owned())?,
            repetition_capacity,
        });
    }
    let scratch_region = regions.len();
    regions.push(transformer::shared_scratch_region(
        &invocation,
        required_scratch_bytes,
    )?);

    let participant_count = u32::try_from(invocation.participants().len())
        .map_err(|_| "masked argmax participant count exceeds u32".to_owned())?;
    let mut replay_key =
        CudaCommandReplayKeyBuilder::new(provider_fingerprint, "vnext_last_token_masked_argmax")
            .u64(launches.len() as u64);
    for launch in &launches {
        replay_key = replay_key
            .u64(launch.first_region as u64)
            .u64(launch.scratch_offset_bytes)
            .i32(launch.vocabulary_size)
            .i32(launch.repetition_capacity);
    }
    let function = function.clone();
    CudaDeviceCommand::replayable_operation(
        "vnext_last_token_masked_argmax",
        regions,
        replay_key.finish(),
        move |stream, regions| {
            for launch in &launches {
                let logits = regions[launch.first_region].device_ptr();
                let valid_mask = regions[launch.first_region + 1].device_ptr();
                let repetition_token_ids = regions[launch.first_region + 2].device_ptr();
                let repetition_offsets = regions[launch.first_region + 3].device_ptr();
                let repetition_penalty = regions[launch.first_region + 4].device_ptr();
                let output = regions[launch.first_region + 5].device_ptr();
                let scratch = regions[scratch_region]
                    .device_ptr()
                    .checked_add(launch.scratch_offset_bytes)
                    .ok_or_else(|| {
                        CudaDeviceRuntimeError::contract(
                            "vNext masked argmax scratch pointer overflows",
                        )
                    })?;
                let mut builder = stream.launch_builder(&function);
                builder.arg(&logits);
                builder.arg(&scratch);
                builder.arg(&launch.vocabulary_size);
                builder.arg(&valid_mask);
                builder.arg(&launch.vocabulary_size);
                builder.arg(&repetition_offsets);
                builder.arg(&repetition_token_ids);
                builder.arg(&repetition_penalty);
                builder.arg(&launch.repetition_capacity);
                builder.arg(&output);
                unsafe {
                    builder.launch(LaunchConfig {
                        grid_dim: (1, 1, 1),
                        block_dim: (THREADS_PER_BLOCK, 1, 1),
                        shared_mem_bytes: 0,
                    })
                }
                .map_err(|error| {
                    CudaDeviceRuntimeError::driver("vNext masked argmax launch", error)
                })?;
            }
            Ok(())
        },
    )
    .and_then(|command| {
        command.with_work_attribution(
            if participant_count == 1 {
                DeviceBatchingForm::Scalar
            } else {
                DeviceBatchingForm::ParticipantLoop
            },
            participant_count,
            u64::from(participant_count),
            u64::from(participant_count),
            0,
        )
    })
    .map_err(|error| error.to_string())
}

fn masked_argmax_scratch_stride(vocabulary_size: u64) -> Result<u64, String> {
    let bytes = vocabulary_size
        .checked_mul(ElementType::F16.size_bytes())
        .ok_or_else(|| "CUDA masked argmax scratch size overflows".to_owned())?;
    bytes
        .checked_add(VALUE_ALIGNMENT_BYTES - 1)
        .map(|value| value & !(VALUE_ALIGNMENT_BYTES - 1))
        .filter(|value| *value != 0)
        .ok_or_else(|| "CUDA masked argmax scratch alignment overflows".to_owned())
}

fn validate_masked_argmax_signature(
    logits: &ResolvedValueBinding,
    valid_mask: &ResolvedValueBinding,
    repetition_token_ids: &ResolvedValueBinding,
    repetition_offsets: &ResolvedValueBinding,
    repetition_penalty: &ResolvedValueBinding,
    output: &ResolvedValueBinding,
    vocabulary_size: u64,
) -> Result<i32, String> {
    let contiguous = |binding: &ResolvedValueBinding| {
        matches!(binding.tensor().layout(), ResolvedTensorLayout::Contiguous)
    };
    if logits.tensor().element_type() != ElementType::F16
        || valid_mask.tensor().element_type() != ElementType::U8
        || repetition_token_ids.tensor().element_type() != ElementType::U32
        || repetition_offsets.tensor().element_type() != ElementType::U32
        || repetition_penalty.tensor().element_type() != ElementType::F32
        || output.tensor().element_type() != ElementType::U32
        || logits.tensor().dimensions() != [1, vocabulary_size]
        || valid_mask.tensor().dimensions() != [vocabulary_size]
        || repetition_token_ids.tensor().dimensions().len() != 1
        || repetition_token_ids.tensor().dimensions()[0] == 0
        || repetition_offsets.tensor().dimensions() != [2]
        || repetition_penalty.tensor().dimensions() != [1]
        || output.tensor().dimensions() != [1]
        || !contiguous(logits)
        || !contiguous(valid_mask)
        || !contiguous(repetition_token_ids)
        || !contiguous(repetition_offsets)
        || !contiguous(repetition_penalty)
        || !contiguous(output)
    {
        return Err("CUDA masked argmax invocation differs from its resolved signature".to_owned());
    }
    i32::try_from(repetition_token_ids.tensor().dimensions()[0])
        .map_err(|_| "CUDA masked argmax repetition capacity exceeds i32".to_owned())
}

#[derive(Debug, Clone, Copy)]
struct LastTokenDenseLinearLaunch {
    input_region: usize,
    output_region: usize,
}

fn encode_last_token_dense_linear(
    provider_fingerprint: &str,
    invocation: BatchedOperationInvocation<'_, CudaDeviceBuffer>,
) -> Result<CudaDeviceCommand, String> {
    if invocation.operation().id.as_str() != LAST_TOKEN_DENSE_LINEAR_OPERATION_ID
        || invocation.participants().is_empty()
    {
        return Err("CUDA last-token dense-linear received another or empty operation".to_owned());
    }
    let token_ranges = invocation.participant_token_ranges();
    if token_ranges.len() != invocation.participants().len() {
        return Err("CUDA last-token dense-linear participant ranges are incomplete".to_owned());
    }
    let first = &invocation.participants()[0];
    let hidden_size = unsigned_attribute(first.attributes(), "hidden_size")?;
    let out_features = unsigned_attribute(first.attributes(), "out_features")?;
    let input_packed =
        transformer::token_binding_is_packed(&invocation, ResolvedValueRole::Input, 0)?;
    let mut regions = vec![transformer::shared_full_region(
        &invocation,
        ResolvedValueRole::Input,
        1,
        ElementType::F16,
    )?];
    let mut launches = Vec::with_capacity(invocation.participants().len());
    for (participant, token_range) in invocation.participants().iter().zip(token_ranges) {
        let input = binding(participant.bindings(), ResolvedValueRole::Input, 0)?;
        let participant_weight = binding(participant.bindings(), ResolvedValueRole::Input, 1)?;
        let output = binding(participant.bindings(), ResolvedValueRole::Output, 0)?;
        if unsigned_attribute(participant.attributes(), "hidden_size")? != hidden_size
            || unsigned_attribute(participant.attributes(), "out_features")? != out_features
        {
            return Err("CUDA last-token dense-linear participant attributes disagree".to_owned());
        }
        validate_last_token_dense_linear_signature(
            input,
            participant_weight,
            output,
            hidden_size,
            out_features,
        )?;
        let source_range = token_range.source_token_range();
        let packed_range = token_range.immediate_token_range();
        let selected_range = if input_packed {
            packed_range
        } else {
            source_range
        };
        if selected_range.is_empty() {
            return Err("CUDA last-token dense-linear cannot select from an empty span".to_owned());
        }
        let last_token = selected_range.end - 1;
        let input_region = regions.len();
        let source = contiguous_token_region(participant, input, ElementType::F16, last_token, 1)?;
        regions.push(source);
        let output_region = regions.len();
        let destination = contiguous_region(participant, output, ElementType::F16)?;
        regions.push(destination);
        launches.push(LastTokenDenseLinearLaunch {
            input_region,
            output_region,
        });
    }

    let participant_count = u32::try_from(invocation.participants().len())
        .map_err(|_| "last-token dense-linear participant count exceeds u32".to_owned())?;
    let token_count = u64::from(participant_count);
    let compute_dispatch_count = launches.len() as u64;
    let rows = 1_i32;
    let hidden_size = i32::try_from(hidden_size)
        .map_err(|_| "last-token dense-linear hidden size exceeds i32".to_owned())?;
    let out_features = i32::try_from(out_features)
        .map_err(|_| "last-token dense-linear output width exceeds i32".to_owned())?;
    let mut replay_key =
        CudaCommandReplayKeyBuilder::new(provider_fingerprint, "vnext_last_token_dense_linear")
            .i32(rows)
            .i32(hidden_size)
            .i32(out_features)
            .u64(launches.len() as u64);
    for launch in &launches {
        replay_key = replay_key
            .u64(launch.input_region as u64)
            .u64(launch.output_region as u64);
    }
    CudaDeviceCommand::replayable_operation_with_blas(
        "vnext_last_token_dense_linear",
        regions,
        replay_key.finish(),
        move |_stream, blas, regions| {
            let weight = regions[0].device_ptr();
            for launch in &launches {
                transformer::launch_gemm_f16(
                    blas,
                    regions[launch.input_region].device_ptr(),
                    weight,
                    regions[launch.output_region].device_ptr(),
                    rows,
                    out_features,
                    hidden_size,
                    "vNext last-token dense-linear GEMM",
                )?;
            }
            Ok(())
        },
    )
    .and_then(|command| {
        command.with_work_attribution(
            DeviceBatchingForm::ParticipantLoop,
            participant_count,
            token_count,
            compute_dispatch_count,
            0,
        )
    })
    .map_err(|error| error.to_string())
}

fn validate_last_token_dense_linear_signature(
    input: &ResolvedValueBinding,
    weight: &ResolvedValueBinding,
    output: &ResolvedValueBinding,
    hidden_size: u64,
    out_features: u64,
) -> Result<(), String> {
    let contiguous = |binding: &ResolvedValueBinding| {
        matches!(binding.tensor().layout(), ResolvedTensorLayout::Contiguous)
    };
    let input_dimensions = input.tensor().dimensions();
    if input.tensor().element_type() != ElementType::F16
        || weight.tensor().element_type() != ElementType::F16
        || output.tensor().element_type() != ElementType::F16
        || input_dimensions.len() != 2
        || input_dimensions[0] == 0
        || input_dimensions[1] != hidden_size
        || weight.tensor().dimensions() != [out_features, hidden_size]
        || output.tensor().dimensions() != [1, out_features]
        || !contiguous(input)
        || !contiguous(weight)
        || !contiguous(output)
    {
        return Err(
            "CUDA last-token dense-linear invocation differs from its resolved signature"
                .to_owned(),
        );
    }
    Ok(())
}

#[derive(Debug, Clone, Copy)]
struct EmbeddingLaunch {
    first_region: usize,
    token_count: u64,
    vocabulary_size: u32,
    hidden_size: i32,
    grid_x: u32,
}

fn encode_token_embedding(
    function: &CudaFunction,
    provider_fingerprint: &str,
    invocation: BatchedOperationInvocation<'_, CudaDeviceBuffer>,
) -> Result<CudaDeviceCommand, String> {
    if invocation.operation().id.as_str() != TOKEN_EMBEDDING_OPERATION_ID
        || invocation.participants().is_empty()
    {
        return Err("CUDA token embedding received another or empty operation".to_owned());
    }

    let token_ranges = invocation.participant_token_ranges();
    if token_ranges.len() != invocation.participants().len() {
        return Err("CUDA token embedding participant ranges are incomplete".to_owned());
    }
    let input_packed =
        transformer::token_binding_is_packed(&invocation, ResolvedValueRole::Input, 0)?;
    let mut regions = Vec::with_capacity(invocation.participants().len() * 3);
    let mut launches = Vec::with_capacity(invocation.participants().len());
    for (participant, token_range) in invocation.participants().iter().zip(token_ranges) {
        let token_ids = binding(participant.bindings(), ResolvedValueRole::Input, 0)?;
        let table = binding(participant.bindings(), ResolvedValueRole::Input, 1)?;
        let output = binding(participant.bindings(), ResolvedValueRole::Output, 0)?;
        let hidden_size = unsigned_attribute(participant.attributes(), "hidden_size")?;
        let vocabulary_size = unsigned_attribute(participant.attributes(), "vocab_size")?;
        validate_signature(token_ids, table, output, vocabulary_size, hidden_size)?;
        let source_range = token_range.source_token_range();
        let packed_range = token_range.immediate_token_range();
        let token_count = token_range.immediate_tokens();
        let grid_x = hidden_size
            .div_ceil(THREADS_PER_BLOCK as u64)
            .try_into()
            .map_err(|_| "embedding launch grid exceeds u32".to_owned())?;

        let first_region = regions.len();
        regions.push(contiguous_region(participant, table, ElementType::F16)?);
        regions.push(contiguous_token_region(
            participant,
            token_ids,
            ElementType::U32,
            if input_packed {
                packed_range.start
            } else {
                source_range.start
            },
            token_count,
        )?);
        regions.push(contiguous_token_region(
            participant,
            output,
            ElementType::F16,
            packed_range.start,
            token_count,
        )?);
        launches.push(EmbeddingLaunch {
            first_region,
            token_count,
            vocabulary_size: vocabulary_size
                .try_into()
                .map_err(|_| "embedding vocabulary size exceeds u32".to_owned())?,
            hidden_size: hidden_size
                .try_into()
                .map_err(|_| "embedding hidden size exceeds i32".to_owned())?,
            grid_x,
        });
    }

    let participant_count = u32::try_from(invocation.participants().len())
        .map_err(|_| "embedding participant count exceeds u32".to_owned())?;
    let token_count = invocation.work_shape().immediate_tokens();
    let compute_dispatch_count = launches
        .iter()
        .map(|launch| launch.token_count.div_ceil(MAXIMUM_TOKENS_PER_LAUNCH))
        .sum();
    let mut replay_key =
        CudaCommandReplayKeyBuilder::new(provider_fingerprint, "vnext_token_embedding")
            .u64(launches.len() as u64);
    for launch in &launches {
        replay_key = replay_key
            .u64(launch.first_region as u64)
            .u64(launch.token_count)
            .u32(launch.vocabulary_size)
            .i32(launch.hidden_size)
            .u32(launch.grid_x);
    }
    let function = function.clone();
    CudaDeviceCommand::replayable_operation(
        "vnext_token_embedding",
        regions,
        replay_key.finish(),
        move |stream, regions| {
            for launch in &launches {
                let table = regions[launch.first_region].device_ptr();
                let token_ids_base = regions[launch.first_region + 1].device_ptr();
                let output_base = regions[launch.first_region + 2].device_ptr();
                let mut token_offset = 0_u64;
                while token_offset < launch.token_count {
                    let chunk_tokens =
                        (launch.token_count - token_offset).min(MAXIMUM_TOKENS_PER_LAUNCH);
                    let token_ids = checked_pointer_offset(
                        token_ids_base,
                        token_offset,
                        ElementType::U32.size_bytes(),
                        "token id",
                    )?;
                    let output_element_offset = token_offset
                        .checked_mul(launch.hidden_size as u64)
                        .ok_or_else(|| {
                            CudaDeviceRuntimeError::contract(
                                "vNext embedding output element offset overflows",
                            )
                        })?;
                    let output = checked_pointer_offset(
                        output_base,
                        output_element_offset,
                        ElementType::F16.size_bytes(),
                        "embedding output",
                    )?;
                    let batch = chunk_tokens as i32;
                    let mut builder = stream.launch_builder(&function);
                    builder.arg(&table);
                    builder.arg(&token_ids);
                    builder.arg(&output);
                    builder.arg(&batch);
                    builder.arg(&launch.hidden_size);
                    builder.arg(&launch.vocabulary_size);
                    unsafe {
                        builder.launch(LaunchConfig {
                            grid_dim: (launch.grid_x, chunk_tokens as u32, 1),
                            block_dim: (THREADS_PER_BLOCK, 1, 1),
                            shared_mem_bytes: 0,
                        })
                    }
                    .map_err(|error| {
                        CudaDeviceRuntimeError::driver("vNext token embedding launch", error)
                    })?;
                    token_offset += chunk_tokens;
                }
            }
            Ok(())
        },
    )
    .and_then(|command| {
        command.with_work_attribution(
            DeviceBatchingForm::ParticipantLoop,
            participant_count,
            token_count,
            compute_dispatch_count,
            0,
        )
    })
    .map_err(|error| error.to_string())
}

fn checked_pointer_offset(
    base: cudarc::driver::sys::CUdeviceptr,
    elements: u64,
    element_bytes: u64,
    context: &'static str,
) -> Result<cudarc::driver::sys::CUdeviceptr, CudaDeviceRuntimeError> {
    elements
        .checked_mul(element_bytes)
        .and_then(|bytes| base.checked_add(bytes))
        .ok_or_else(|| CudaDeviceRuntimeError::contract(format!("{context} pointer overflows")))
}

fn validate_signature(
    token_ids: &ResolvedValueBinding,
    table: &ResolvedValueBinding,
    output: &ResolvedValueBinding,
    vocabulary_size: u64,
    hidden_size: u64,
) -> Result<u64, String> {
    let token_dimensions = token_ids.tensor().dimensions();
    let table_dimensions = table.tensor().dimensions();
    let output_dimensions = output.tensor().dimensions();
    let contiguous = |binding: &ResolvedValueBinding| {
        matches!(binding.tensor().layout(), ResolvedTensorLayout::Contiguous)
    };
    if token_ids.tensor().element_type() != ElementType::U32
        || table.tensor().element_type() != ElementType::F16
        || output.tensor().element_type() != ElementType::F16
        || token_dimensions.len() != 1
        || table_dimensions != [vocabulary_size, hidden_size]
        || output_dimensions != [token_dimensions[0], hidden_size]
        || !contiguous(token_ids)
        || !contiguous(table)
        || !contiguous(output)
    {
        return Err(
            "CUDA token embedding invocation differs from its resolved signature".to_owned(),
        );
    }
    Ok(token_dimensions[0])
}

fn binding(
    bindings: &[ResolvedValueBinding],
    role: ResolvedValueRole,
    ordinal: u32,
) -> Result<&ResolvedValueBinding, String> {
    bindings
        .iter()
        .find(|binding| binding.role() == role && binding.ordinal() == ordinal)
        .ok_or_else(|| format!("CUDA operation lacks {role:?} binding {ordinal}"))
}

fn unsigned_attribute(
    attributes: &std::collections::BTreeMap<AttributeId, SemanticValue>,
    name: &str,
) -> Result<u64, String> {
    match attributes
        .iter()
        .find(|(attribute, _)| attribute.as_str() == name)
        .map(|(_, value)| value)
    {
        Some(SemanticValue::Unsigned(value)) => Ok(*value),
        _ => Err(format!("CUDA operation lacks unsigned attribute {name:?}")),
    }
}

fn reusable_token_topology(
    request: &ReusableExecutionTopologyRequest<'_>,
    domain: &'static [u8],
) -> Result<ReusableExecutionTopology, VNextError> {
    for (role, ordinal) in [
        (ResolvedValueRole::Input, 0),
        (ResolvedValueRole::Input, 1),
        (ResolvedValueRole::Output, 0),
    ] {
        if request
            .binding_reusable_address_scope(role, ordinal)?
            .is_none()
        {
            return Ok(ReusableExecutionTopology::EagerBoundary);
        }
    }

    let bind_source_ranges =
        !request.binding_uses_packed_batch_coordinates(ResolvedValueRole::Input, 0)?;
    let ranges = request.work_shape().participant_token_ranges();
    let mut digest = Sha256::new();
    digest.update(domain);
    digest.update((ranges.len() as u64).to_le_bytes());
    digest.update(request.work_shape().immediate_tokens().to_le_bytes());
    for range in ranges {
        let source = range.source_token_range();
        let packed = range.immediate_token_range();
        digest.update(range.immediate_tokens().to_le_bytes());
        if bind_source_ranges {
            digest.update(source.start.to_le_bytes());
            digest.update(source.end.to_le_bytes());
        }
        digest.update(packed.start.to_le_bytes());
        digest.update(packed.end.to_le_bytes());
    }
    Ok(ReusableExecutionTopology::Dynamic(
        DeviceReusableExecutionTopologyFingerprint::from_sha256(digest.finalize().into()),
    ))
}

fn contiguous_region(
    participant: &OperationInvocation<'_, CudaDeviceBuffer>,
    binding: &ResolvedValueBinding,
    element_type: ElementType,
) -> Result<CudaBufferRegion, String> {
    let [component] = binding.storage().components() else {
        return Err("CUDA operation requires one storage component per value".to_owned());
    };
    if component.element_type() != element_type {
        return Err("CUDA operation storage element type differs from its contract".to_owned());
    }
    contiguous_region_range(
        participant,
        binding,
        element_type,
        component.offset_bytes(),
        component.length_bytes(),
    )
}

fn contiguous_token_region(
    participant: &OperationInvocation<'_, CudaDeviceBuffer>,
    binding: &ResolvedValueBinding,
    element_type: ElementType,
    token_start: u64,
    token_count: u64,
) -> Result<CudaBufferRegion, String> {
    let [component] = binding.storage().components() else {
        return Err("CUDA operation requires one storage component per value".to_owned());
    };
    let projection = participant
        .work()
        .token_projection(binding.role(), binding.ordinal())
        .ok_or_else(|| "CUDA operation binding has no token work projection".to_owned())?;
    let dimensions = binding.tensor().dimensions();
    if projection.axis() != 0
        || projection.rank() as usize != dimensions.len()
        || dimensions.first() != Some(&projection.canonical_extent())
        || component.offset_bytes() != 0
        || component.length_bytes() % projection.canonical_extent() != 0
    {
        return Err(
            "CUDA contiguous token projection is not a canonical leading-axis tensor".to_owned(),
        );
    }
    let bytes_per_token = component.length_bytes() / projection.canonical_extent();
    let logical_offset = token_start
        .checked_mul(bytes_per_token)
        .ok_or_else(|| "CUDA token region offset overflows".to_owned())?;
    let logical_length = token_count
        .checked_mul(bytes_per_token)
        .ok_or_else(|| "CUDA token region length overflows".to_owned())?;
    contiguous_region_range(
        participant,
        binding,
        element_type,
        logical_offset,
        logical_length,
    )
}

fn contiguous_region_range(
    participant: &OperationInvocation<'_, CudaDeviceBuffer>,
    binding: &ResolvedValueBinding,
    element_type: ElementType,
    logical_offset_bytes: u64,
    logical_length_bytes: u64,
) -> Result<CudaBufferRegion, String> {
    let [component] = binding.storage().components() else {
        return Err("CUDA operation requires one storage component per value".to_owned());
    };
    if component.element_type() != element_type {
        return Err("CUDA operation storage element type differs from its contract".to_owned());
    }
    let view = participant
        .views()
        .iter()
        .find(|view| view.resource_id() == component.resource_id())
        .ok_or_else(|| "CUDA operation value has no resource view".to_owned())?;
    // The invocation construction has already proved the resource view and
    // component describe the same allocation. The provider still translates
    // the logical slice instead of assuming an arena base pointer.
    let translated = view
        .translate(logical_offset_bytes, logical_length_bytes)
        .map_err(|error| error.to_string())?;
    let mut physical = translated.iter();
    let region = physical
        .next()
        .ok_or_else(|| "CUDA operation translated to no physical region".to_owned())?;
    if physical.next().is_some() {
        return Err("CUDA operation requires contiguous physical storage".to_owned());
    }
    let (buffer, range, retention) = region.buffer_and_physical_range();
    let region = buffer
        .retained_region(range, retention)
        .map_err(|error| error.to_string())?;
    if region.element_type() != element_type || region.length_bytes() != logical_length_bytes {
        return Err(
            "CUDA operation physical region differs from its resolved component".to_owned(),
        );
    }
    Ok(region)
}

fn same_physical_region(left: &CudaBufferRegion, right: &CudaBufferRegion) -> bool {
    left.device_ptr() == right.device_ptr()
        && left.length_bytes() == right.length_bytes()
        && left.element_type() == right.element_type()
}

fn contiguous_bindings() -> Vec<ProviderStorageBindingRequirement> {
    [
        (ResolvedValueRole::Input, 0),
        (ResolvedValueRole::Input, 1),
        (ResolvedValueRole::Output, 0),
    ]
    .into_iter()
    .map(|(role, ordinal)| {
        ProviderStorageBindingRequirement::new(
            role,
            ordinal,
            DynamicStorageRequirement::contiguous(),
        )
    })
    .collect()
}

fn provider_failure(
    identity: ferrum_interfaces::vnext::ExecutionIdentityEnvelope,
    stage: &'static str,
    message: String,
) -> OperationFailure {
    let message = message.chars().take(2048).collect::<String>();
    OperationFailure::new(identity, ProfilePhase::Forward, stage, message, false)
        .expect("core-issued CUDA operation identity must form a valid provider failure")
}

fn implementation_fingerprint(parts: &[&[u8]]) -> String {
    let mut digest = Sha256::new();
    for part in parts {
        digest.update((part.len() as u64).to_le_bytes());
        digest.update(part);
    }
    format!("{:x}", digest.finalize())
}

fn contract_error(error: VNextError) -> CudaDeviceRuntimeError {
    CudaDeviceRuntimeError::contract(error.to_string())
}