ferrum-kernels 0.8.5

//! CUDA provider for the standard causal paged-attention operation.

use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;

use cudarc::cublas::CudaBlas;
use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use ferrum_interfaces::vnext::{
    causal_paged_attention_contract, AttributeId, BatchedOperationInvocation, CapabilityId,
    ContractVersion, DeviceBatchingForm, DeviceReusableExecutionTopologyFingerprint, DeviceRuntime,
    DynamicStorageAllocator, DynamicStorageProfile, DynamicStorageRequirement, DynamicStorageView,
    ElementType, EncodedDeviceOperation, EncodedReusableExecutionBindings,
    OperationBufferStorageKind, OperationContract, OperationFailure, OperationInvocation,
    OperationProvider, OperationProviderDescriptor, OperationResourceEstimate,
    OperationResourceEstimateRequest, OperationResourceEstimator, ProfilePhase, ProviderId,
    ProviderStorageBindingRequirement, ProviderWorkspaceRequirement, ProviderWorkspaceReusePolicy,
    ProviderWorkspaceScope, ProviderWorkspaceSizeFormula, QuantizationFormatId,
    ResolvedTensorLayout, ResolvedValueBinding, ResolvedValueRole, ReusableExecutionTopology,
    ReusableExecutionTopologyRequest, ReusableExecutionValueAddress,
    ReusableExecutionWorkspaceAddress, SemanticValue, VNextError, WeightFormatId,
    CAUSAL_PAGED_ATTENTION_F16_CAPABILITY_ID,
};
use ferrum_types::{AttentionExecutionPolicy, CUDA_NATIVE_ADAPTIVE_V1_MAX_SEQUENCE_TOKENS};
use sha2::{Digest, Sha256};

use super::{attach_invocation_binding, ensure_estimator_request, estimate, launch_gemm_f16};
#[cfg(feature = "vllm-marlin")]
use super::{
    marlin_fp8_weights::resolve_marlin_fp8_weight,
    moe_weights::{
        resolve_compressed_tensors_marlin_matrix_weight,
        resolve_compressed_tensors_symmetric_marlin_matrix_weight,
        resolve_gptq_marlin_matrix_weight, COMPRESSED_TENSORS_MARLIN_CAPABILITY_ID,
        COMPRESSED_TENSORS_MARLIN_QUANTIZATION_FORMAT_ID,
        COMPRESSED_TENSORS_MARLIN_SYMMETRIC_CAPABILITY_ID,
        COMPRESSED_TENSORS_MARLIN_SYMMETRIC_QUANTIZATION_FORMAT_ID,
        COMPRESSED_TENSORS_MARLIN_SYMMETRIC_WEIGHT_FORMAT_ID,
        COMPRESSED_TENSORS_MARLIN_WEIGHT_FORMAT_ID, GPTQ_MARLIN_CAPABILITY_ID,
        GPTQ_MARLIN_QUANTIZATION_FORMAT_ID, GPTQ_MARLIN_WEIGHT_FORMAT_ID,
    },
    MarlinProjectionRuntime,
};
#[cfg(feature = "vllm-marlin")]
use crate::backend::cuda::vllm_marlin::MarlinF16WeightType;
#[cfg(feature = "vllm-paged-attn-v2")]
use crate::backend::cuda::vllm_paged_attn::{
    dispatch_vnext_addressed_paged_attention_raw, VnextAddressedPagedAttentionKernel,
};
use crate::backend::cuda::vnext_ops::{
    binding, contiguous_token_region, contract_error, implementation_fingerprint,
    DENSE_SAFETENSORS_FORMAT_ID, THREADS_PER_BLOCK, VNEXT_KV_PAGE_BYTES,
};
use crate::backend::cuda::vnext_replay::CudaCommandReplayKeyBuilder;
use crate::backend::cuda::vnext_runtime::{
    CudaBufferRegion, CudaDeviceBuffer, CudaDeviceCommand, CudaDeviceRuntime,
    CudaDeviceRuntimeError,
};
#[cfg(feature = "vllm-marlin")]
use crate::marlin_fp8_materializer::{
    MARLIN_FP8_CAPABILITY_ID, MARLIN_FP8_GROUP128_QUANTIZATION_FORMAT_ID,
    MARLIN_FP8_GROUP128_WEIGHT_FORMAT_ID, MARLIN_FP8_QUANTIZATION_FORMAT_ID,
    MARLIN_FP8_WEIGHT_FORMAT_ID,
};

const PROVIDER_ID: &str = "provider.cuda.causal_paged_attention.f16";
const ESTIMATOR_ID: &str = "resource-estimator.cuda.causal_paged_attention.f16";
const GEMMA4_PROVIDER_ID: &str = "provider.cuda.gemma4_causal_paged_attention.f16";
const GEMMA4_ESTIMATOR_ID: &str = "resource-estimator.cuda.gemma4_causal_paged_attention.f16";
const RMS_NORM_FUNCTION: &str = "rms_norm_f16";
const PREPARE_FUNCTION: &str = "vnext_causal_prepare_f16";
const ATTENTION_FUNCTION: &str = "vnext_causal_attention_f16";
const GROUPED_ATTENTION_FUNCTION: &str = "vnext_causal_attention_grouped_f16";
const VARLEN_ADDRESSED_FUNCTION: &str = "vnext_paged_varlen_attn_vllm_addressed_f16";
const VARLEN_TILED_ADDRESSED_FUNCTION: &str = "vnext_paged_varlen_attn_vllm_tiled_q4_addressed_f16";
const ATTENTION_GATE_FUNCTION: &str = "qwen35_apply_attention_gate_f16";
const RESIDUAL_ADD_FUNCTION: &str = "residual_add_f16";
const RESIDUAL_ADD_INPLACE_FUNCTION: &str = "residual_add_inplace_f16";
const COMPUTE_TOKEN_MAJOR_OPERATION: &str = "vnext.causal_attention.token_major_fallback";
const COMPUTE_VLLM_FALLBACK_OPERATION: &str = "vnext.causal_attention.vllm_addressed_fallback";
const COMPUTE_VLLM_VARLEN_OPERATION: &str = "vnext.causal_attention.vllm_varlen_addressed";
const COMPUTE_VLLM_VARLEN_TILED_OPERATION: &str = "vnext.causal_attention.vllm_varlen_q4_addressed";
const COMPUTE_VLLM_DECODE_V1_OPERATION: &str =
    "vnext.causal_attention.vllm_paged_attention_v1_addressed";
const COMPUTE_VLLM_DECODE_V2_OPERATION: &str =
    "vnext.causal_attention.vllm_paged_attention_v2_addressed";
const COMPUTE_MIXED_OPERATION: &str = "vnext.causal_attention.mixed_native_paths";
const SCRATCH_ALIGNMENT: u64 = 16;
const POINTER_BYTES: u64 = std::mem::size_of::<u64>() as u64;
// table entries, source start, active tokens, sequence tokens, packed token
// start, and one reserved word keep the following u64 page table aligned.
const BINDING_CONTROL_WORDS: usize = 6;
const BINDING_CONTROL_BYTES: u64 = (BINDING_CONTROL_WORDS * std::mem::size_of::<i32>()) as u64;
const BINDING_SEQUENCE_LENGTH_OFFSET: u64 = 3 * std::mem::size_of::<i32>() as u64;
const WARP_THREADS: u32 = 32;
const GROUPED_FALLBACK_DEFAULT_KV_TILE_TOKENS: u64 = 16;
const GROUPED_FALLBACK_GEMMA_LOCAL_KV_TILE_TOKENS: u64 = 32;
const MAXIMUM_GROUPED_QUERY_HEADS_PER_KV: u32 = 16;
const MAXIMUM_HEAD_DIM: u64 = 512;
const MAXIMUM_STANDARD_HEAD_DIM: u64 = 256;
const MAXIMUM_VARLEN_HEAD_DIM: u64 = 256;
const VLLM_BLOCK_TOKENS: u64 = 16;
const VLLM_PARTITION_TOKENS: u64 = CUDA_NATIVE_ADAPTIVE_V1_MAX_SEQUENCE_TOKENS;
const VARLEN_DEFAULT_SHARED_LIMIT_BYTES: u64 = 48 * 1024;
const VARLEN_STATIC_SHARED_RESERVE_BYTES: u64 = 1024;
const VARLEN_DYNAMIC_SHARED_BUDGET_BYTES: u64 =
    VARLEN_DEFAULT_SHARED_LIMIT_BYTES - VARLEN_STATIC_SHARED_RESERVE_BYTES;
const VARLEN_TILED_QUERY_TOKENS: u64 = 4;

pub(in crate::backend::cuda::vnext_ops) struct CudaCausalPagedAttentionProvider {
    descriptor: OperationProviderDescriptor,
    functions: CausalAttentionFunctions,
    attention_policy: AttentionExecutionPolicy,
    semantics: CausalAttentionSemantics,
    #[cfg(feature = "vllm-marlin")]
    projection_runtime: MarlinProjectionRuntime,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CausalAttentionSemantics {
    Standard,
    Gemma4,
}

impl CausalAttentionSemantics {
    const fn provider_id(self) -> &'static str {
        match self {
            Self::Standard => PROVIDER_ID,
            Self::Gemma4 => GEMMA4_PROVIDER_ID,
        }
    }

    const fn estimator_id(self) -> &'static str {
        match self {
            Self::Standard => ESTIMATOR_ID,
            Self::Gemma4 => GEMMA4_ESTIMATOR_ID,
        }
    }

    const fn input_count(self) -> u32 {
        match self {
            Self::Standard => 9,
            Self::Gemma4 => 10,
        }
    }

    const fn has_post_attention_norm(self) -> bool {
        matches!(self, Self::Gemma4)
    }

    const fn fingerprint_tag(self) -> &'static [u8] {
        match self {
            Self::Standard => b"standard-causal-attention-v2",
            Self::Gemma4 => b"gemma4-causal-attention-v1",
        }
    }
}

#[derive(Clone)]
struct CausalAttentionFunctions {
    rms_norm: CudaFunction,
    prepare: CudaFunction,
    attention: CudaFunction,
    grouped_attention: CudaFunction,
    varlen_addressed: CudaFunction,
    varlen_tiled_addressed: CudaFunction,
    attention_gate: CudaFunction,
    residual_add: CudaFunction,
    residual_add_inplace: CudaFunction,
}

impl CudaCausalPagedAttentionProvider {
    pub(in crate::backend::cuda::vnext_ops) fn new(
        runtime: &CudaDeviceRuntime,
        attention_policy: AttentionExecutionPolicy,
    ) -> Result<Self, CudaDeviceRuntimeError> {
        if !attention_policy.is_resolved() {
            return Err(CudaDeviceRuntimeError::contract(
                "CUDA causal attention policy must be resolved before provider construction",
            ));
        }
        let contract = causal_paged_attention_contract().map_err(contract_error)?;
        Self::new_for_contract(
            runtime,
            attention_policy,
            &contract,
            CAUSAL_PAGED_ATTENTION_F16_CAPABILITY_ID,
            CausalAttentionSemantics::Standard,
        )
    }

    /// Builds the CUDA provider for Gemma 4's sandwich-normalized attention.
    ///
    /// The contract and capability are supplied by the registration layer so
    /// this implementation module does not depend on a model-family module.
    /// The expected Gemma contract keeps the standard inputs 0..=8 and adds
    /// post-attention RMSNorm weight at input 9.
    pub(in crate::backend::cuda::vnext_ops) fn new_gemma4(
        runtime: &CudaDeviceRuntime,
        attention_policy: AttentionExecutionPolicy,
        contract: &dyn OperationContract,
        capability_id: &str,
    ) -> Result<Self, CudaDeviceRuntimeError> {
        Self::new_for_contract(
            runtime,
            attention_policy,
            contract,
            capability_id,
            CausalAttentionSemantics::Gemma4,
        )
    }

    fn new_for_contract(
        runtime: &CudaDeviceRuntime,
        attention_policy: AttentionExecutionPolicy,
        contract: &dyn OperationContract,
        capability_id: &str,
        semantics: CausalAttentionSemantics,
    ) -> Result<Self, CudaDeviceRuntimeError> {
        if !attention_policy.is_resolved() {
            return Err(CudaDeviceRuntimeError::contract(
                "CUDA causal attention policy must be resolved before provider construction",
            ));
        }
        let capability = CapabilityId::new(capability_id).map_err(contract_error)?;
        if !runtime.descriptor().capabilities.contains(&capability) {
            return Err(CudaDeviceRuntimeError::contract(
                "CUDA runtime does not advertise causal paged attention",
            ));
        }

        let source = include_str!("causal_attention.rs");
        let mut provider_sources = vec![
            source.as_bytes(),
            crate::ptx::RMS_NORM.as_bytes(),
            crate::ptx::VNEXT_CAUSAL_ATTENTION.as_bytes(),
            crate::ptx::PAGED_VARLEN_ATTENTION_VLLM.as_bytes(),
            crate::ptx::QK_NORM_ROPE.as_bytes(),
            crate::ptx::RESIDUAL_ADD.as_bytes(),
        ];
        #[cfg(feature = "vllm-marlin")]
        provider_sources.extend([
            include_str!("marlin_fp8_weights.rs").as_bytes(),
            include_str!("moe_weights.rs").as_bytes(),
            include_str!("../../vllm_marlin.rs").as_bytes(),
            MARLIN_FP8_CAPABILITY_ID.as_bytes(),
            MARLIN_FP8_WEIGHT_FORMAT_ID.as_bytes(),
            MARLIN_FP8_GROUP128_WEIGHT_FORMAT_ID.as_bytes(),
            MARLIN_FP8_QUANTIZATION_FORMAT_ID.as_bytes(),
            MARLIN_FP8_GROUP128_QUANTIZATION_FORMAT_ID.as_bytes(),
            GPTQ_MARLIN_QUANTIZATION_FORMAT_ID.as_bytes(),
            COMPRESSED_TENSORS_MARLIN_QUANTIZATION_FORMAT_ID.as_bytes(),
            COMPRESSED_TENSORS_MARLIN_SYMMETRIC_QUANTIZATION_FORMAT_ID.as_bytes(),
        ]);
        #[cfg(feature = "vllm-paged-attn-v2")]
        provider_sources.extend([
            include_str!("../../vllm_paged_attn.rs").as_bytes(),
            crate::native_ops::CUDA_NATIVE_SOURCE_BUNDLE_ID.as_bytes(),
        ]);
        provider_sources.push(attention_policy.as_runtime_value().as_bytes());
        provider_sources.push(semantics.fingerprint_tag());
        let provider_fingerprint = implementation_fingerprint(&provider_sources);
        let estimator_fingerprint = implementation_fingerprint(&[
            source.as_bytes(),
            semantics.estimator_id().as_bytes(),
            semantics.fingerprint_tag(),
        ]);
        let mut provider_capabilities = BTreeSet::from([capability]);
        let mut accepted_weight_formats =
            BTreeSet::from([
                WeightFormatId::new(DENSE_SAFETENSORS_FORMAT_ID).map_err(contract_error)?
            ]);
        let mut accepted_quantization_formats = BTreeSet::new();
        #[cfg(feature = "vllm-marlin")]
        {
            let fp8_marlin_capability =
                CapabilityId::new(MARLIN_FP8_CAPABILITY_ID).map_err(contract_error)?;
            if !runtime
                .descriptor()
                .capabilities
                .contains(&fp8_marlin_capability)
            {
                return Err(CudaDeviceRuntimeError::contract(
                    "CUDA runtime does not advertise causal-attention Marlin FP8",
                ));
            }
            provider_capabilities.insert(fp8_marlin_capability);
            let marlin_capability =
                CapabilityId::new(GPTQ_MARLIN_CAPABILITY_ID).map_err(contract_error)?;
            if !runtime
                .descriptor()
                .capabilities
                .contains(&marlin_capability)
            {
                return Err(CudaDeviceRuntimeError::contract(
                    "CUDA runtime does not advertise causal-attention GPTQ-Marlin",
                ));
            }
            provider_capabilities.insert(marlin_capability);
            let compressed_tensors_capability =
                CapabilityId::new(COMPRESSED_TENSORS_MARLIN_CAPABILITY_ID)
                    .map_err(contract_error)?;
            if !runtime
                .descriptor()
                .capabilities
                .contains(&compressed_tensors_capability)
            {
                return Err(CudaDeviceRuntimeError::contract(
                    "CUDA runtime does not advertise causal-attention compressed-tensors Marlin",
                ));
            }
            provider_capabilities.insert(compressed_tensors_capability);
            let compressed_tensors_symmetric_capability =
                CapabilityId::new(COMPRESSED_TENSORS_MARLIN_SYMMETRIC_CAPABILITY_ID)
                    .map_err(contract_error)?;
            if !runtime
                .descriptor()
                .capabilities
                .contains(&compressed_tensors_symmetric_capability)
            {
                return Err(CudaDeviceRuntimeError::contract(
                    "CUDA runtime does not advertise causal-attention symmetric compressed-tensors Marlin",
                ));
            }
            provider_capabilities.insert(compressed_tensors_symmetric_capability);
            accepted_weight_formats
                .insert(WeightFormatId::new(MARLIN_FP8_WEIGHT_FORMAT_ID).map_err(contract_error)?);
            accepted_weight_formats.insert(
                WeightFormatId::new(MARLIN_FP8_GROUP128_WEIGHT_FORMAT_ID)
                    .map_err(contract_error)?,
            );
            accepted_weight_formats
                .insert(WeightFormatId::new(GPTQ_MARLIN_WEIGHT_FORMAT_ID).map_err(contract_error)?);
            accepted_weight_formats.insert(
                WeightFormatId::new(COMPRESSED_TENSORS_MARLIN_WEIGHT_FORMAT_ID)
                    .map_err(contract_error)?,
            );
            accepted_weight_formats.insert(
                WeightFormatId::new(COMPRESSED_TENSORS_MARLIN_SYMMETRIC_WEIGHT_FORMAT_ID)
                    .map_err(contract_error)?,
            );
            accepted_quantization_formats.insert(
                QuantizationFormatId::new(MARLIN_FP8_QUANTIZATION_FORMAT_ID)
                    .map_err(contract_error)?,
            );
            accepted_quantization_formats.insert(
                QuantizationFormatId::new(MARLIN_FP8_GROUP128_QUANTIZATION_FORMAT_ID)
                    .map_err(contract_error)?,
            );
            accepted_quantization_formats.insert(
                QuantizationFormatId::new(GPTQ_MARLIN_QUANTIZATION_FORMAT_ID)
                    .map_err(contract_error)?,
            );
            accepted_quantization_formats.insert(
                QuantizationFormatId::new(COMPRESSED_TENSORS_MARLIN_QUANTIZATION_FORMAT_ID)
                    .map_err(contract_error)?,
            );
            accepted_quantization_formats.insert(
                QuantizationFormatId::new(
                    COMPRESSED_TENSORS_MARLIN_SYMMETRIC_QUANTIZATION_FORMAT_ID,
                )
                .map_err(contract_error)?,
            );
        }
        let descriptor = OperationProviderDescriptor::new(
            ProviderId::new(semantics.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(),
            provider_capabilities,
            accepted_weight_formats,
            accepted_quantization_formats,
            storage_bindings(semantics).map_err(contract_error)?,
            semantics.estimator_id(),
            ContractVersion::new(1, 0),
            estimator_fingerprint,
        )
        .map_err(contract_error)?;

        let rms_module = runtime
            .context()
            .load_module(Ptx::from_src(crate::ptx::RMS_NORM.to_owned()))
            .map_err(|error| {
                CudaDeviceRuntimeError::driver("causal attention RMSNorm module", error)
            })?;
        let attention_module = runtime
            .context()
            .load_module(Ptx::from_src(crate::ptx::VNEXT_CAUSAL_ATTENTION.to_owned()))
            .map_err(|error| CudaDeviceRuntimeError::driver("causal attention module", error))?;
        let varlen_module = runtime
            .context()
            .load_module(Ptx::from_src(
                crate::ptx::PAGED_VARLEN_ATTENTION_VLLM.to_owned(),
            ))
            .map_err(|error| {
                CudaDeviceRuntimeError::driver("causal attention varlen module", error)
            })?;
        let gate_module = runtime
            .context()
            .load_module(Ptx::from_src(crate::ptx::QK_NORM_ROPE.to_owned()))
            .map_err(|error| {
                CudaDeviceRuntimeError::driver("causal attention gate module", error)
            })?;
        let residual_module = runtime
            .context()
            .load_module(Ptx::from_src(crate::ptx::RESIDUAL_ADD.to_owned()))
            .map_err(|error| {
                CudaDeviceRuntimeError::driver("causal attention residual module", error)
            })?;
        let functions = CausalAttentionFunctions {
            rms_norm: load_function(&rms_module, RMS_NORM_FUNCTION, "causal attention RMSNorm")?,
            prepare: load_function(
                &attention_module,
                PREPARE_FUNCTION,
                "causal attention prepare",
            )?,
            attention: load_function(&attention_module, ATTENTION_FUNCTION, "causal attention")?,
            grouped_attention: load_function(
                &attention_module,
                GROUPED_ATTENTION_FUNCTION,
                "causal grouped-query attention",
            )?,
            varlen_addressed: load_function(
                &varlen_module,
                VARLEN_ADDRESSED_FUNCTION,
                "causal attention addressed varlen",
            )?,
            varlen_tiled_addressed: load_function(
                &varlen_module,
                VARLEN_TILED_ADDRESSED_FUNCTION,
                "causal attention addressed tiled varlen",
            )?,
            attention_gate: load_function(
                &gate_module,
                ATTENTION_GATE_FUNCTION,
                "causal attention output gate",
            )?,
            residual_add: load_function(
                &residual_module,
                RESIDUAL_ADD_FUNCTION,
                "causal attention residual",
            )?,
            residual_add_inplace: load_function(
                &residual_module,
                RESIDUAL_ADD_INPLACE_FUNCTION,
                "causal attention in-place residual",
            )?,
        };
        Ok(Self {
            descriptor,
            functions,
            attention_policy,
            semantics,
            #[cfg(feature = "vllm-marlin")]
            projection_runtime: MarlinProjectionRuntime::query(runtime)?,
        })
    }
}

fn load_function(
    module: &Arc<cudarc::driver::CudaModule>,
    name: &str,
    operation: &'static str,
) -> Result<CudaFunction, CudaDeviceRuntimeError> {
    module
        .load_function(name)
        .map_err(|error| CudaDeviceRuntimeError::driver(operation, error))
}

fn storage_bindings(
    semantics: CausalAttentionSemantics,
) -> Result<Vec<ProviderStorageBindingRequirement>, VNextError> {
    let paged = DynamicStorageRequirement::new(vec![DynamicStorageProfile::new(
        DynamicStorageAllocator::FixedBlockArena {
            block_bytes: VNEXT_KV_PAGE_BYTES,
        },
        DynamicStorageView::PagedRegions {
            block_bytes: VNEXT_KV_PAGE_BYTES,
        },
    )?])?;
    Ok((0..semantics.input_count())
        .map(|ordinal| {
            ProviderStorageBindingRequirement::new(
                ResolvedValueRole::Input,
                ordinal,
                if ordinal == 8 {
                    paged.clone()
                } else {
                    DynamicStorageRequirement::contiguous()
                },
            )
        })
        .chain(std::iter::once(ProviderStorageBindingRequirement::new(
            ResolvedValueRole::Output,
            0,
            DynamicStorageRequirement::contiguous(),
        )))
        .collect())
}

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

    fn estimate_resources(
        &self,
        request: OperationResourceEstimateRequest<'_>,
    ) -> Result<OperationResourceEstimate, VNextError> {
        ensure_estimator_request(
            &self.descriptor,
            &request,
            self.descriptor.operation_id().as_str(),
        )?;
        let shape = CausalAttentionShape::from_attributes_for(request.attributes(), self.semantics)
            .map_err(invalid_plan)?;
        #[cfg(feature = "vllm-marlin")]
        let projection = CausalProjection::from_values(request.values(), self.projection_runtime)
            .map_err(invalid_plan)?;
        #[cfg(not(feature = "vllm-marlin"))]
        let projection = CausalProjection::F16;
        let scratch = ProviderWorkspaceRequirement::from_formula(
            ProviderWorkspaceSizeFormula::affine(
                shape
                    .attention_policy_scratch_bytes(self.attention_policy)
                    .and_then(|bytes| {
                        bytes
                            .checked_add(projection.workspace_reservation_bytes()?)
                            .ok_or_else(|| {
                                "causal attention fixed scratch size overflows".to_owned()
                            })
                    })
                    .map_err(invalid_plan)?,
                0,
                shape.scratch_bytes_per_token().map_err(invalid_plan)?,
            )?,
            SCRATCH_ALIGNMENT,
            ProviderWorkspaceScope::Invocation,
            ProviderWorkspaceReusePolicy::OverwriteBeforeRead,
            DynamicStorageRequirement::contiguous(),
        )?;
        let binding = ProviderWorkspaceRequirement::from_formula(
            ProviderWorkspaceSizeFormula::actual_sequences(
                shape.binding_slot_bytes().map_err(invalid_plan)?,
            )?,
            SCRATCH_ALIGNMENT,
            ProviderWorkspaceScope::Invocation,
            ProviderWorkspaceReusePolicy::OverwriteBeforeRead,
            DynamicStorageRequirement::contiguous(),
        )?;
        Ok(
            estimate(&self.descriptor, request.input_fingerprint(), Some(scratch))
                .with_binding(binding),
        )
    }
}

impl OperationProvider<CudaDeviceRuntime> for CudaCausalPagedAttentionProvider {
    fn reusable_execution_topology(
        &self,
        request: ReusableExecutionTopologyRequest<'_>,
    ) -> Result<ReusableExecutionTopology, VNextError> {
        let mut values = (0..self.semantics.input_count())
            .filter(|ordinal| *ordinal != 8)
            .map(|ordinal| {
                ReusableExecutionValueAddress::captured(ResolvedValueRole::Input, ordinal)
            })
            .collect::<Vec<_>>();
        values.extend([
            ReusableExecutionValueAddress::program_binding(ResolvedValueRole::Input, 8),
            ReusableExecutionValueAddress::captured(ResolvedValueRole::Output, 0),
        ]);
        if request
            .reusable_address_scope(
                &values,
                &[
                    ReusableExecutionWorkspaceAddress::Scratch,
                    ReusableExecutionWorkspaceAddress::Binding,
                ],
            )?
            .is_none()
        {
            return Ok(ReusableExecutionTopology::EagerBoundary);
        }
        reusable_attention_topology(&request, self.attention_policy, self.semantics)
            .map_err(invalid_plan)
    }

    fn encode_selected(
        &self,
        invocation: BatchedOperationInvocation<'_, CudaDeviceBuffer>,
    ) -> Result<EncodedDeviceOperation<CudaDeviceCommand>, OperationFailure> {
        let identity = invocation.participants()[0].identity().clone();
        encode_attention(
            &self.functions,
            self.descriptor.provider_implementation_fingerprint(),
            self.attention_policy,
            self.semantics,
            self.descriptor.operation_id().as_str(),
            #[cfg(feature = "vllm-marlin")]
            self.projection_runtime,
            invocation,
        )
        .map_err(|message| {
            OperationFailure::new(
                identity,
                ProfilePhase::Forward,
                "cuda.causal_paged_attention.encode",
                message.chars().take(2048).collect::<String>(),
                false,
            )
            .expect("core-issued CUDA causal attention identity must be valid")
        })
    }

    fn encode_reusable_execution_bindings(
        &self,
        invocation: BatchedOperationInvocation<'_, CudaDeviceBuffer>,
    ) -> Result<EncodedReusableExecutionBindings<CudaDeviceCommand>, OperationFailure> {
        let identity = invocation.participants()[0].identity().clone();
        encode_reusable_attention_bindings(
            invocation,
            self.semantics,
            self.descriptor.operation_id().as_str(),
        )
        .map_err(|message| {
            OperationFailure::new(
                identity,
                ProfilePhase::Forward,
                "cuda.causal_paged_attention.encode_reusable_bindings",
                message.chars().take(2048).collect::<String>(),
                false,
            )
            .expect("core-issued CUDA causal attention identity must be valid")
        })
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
struct CausalAttentionShape {
    hidden_size: u64,
    query_heads: u64,
    key_value_heads: u64,
    head_dim: u64,
    query_features: u64,
    query_projection_features: u64,
    kv_features: u64,
    rope_dim: u64,
    rope_frequency_denominator: u64,
    maximum_context_tokens: u64,
    epsilon: f32,
    rope_theta: f32,
    attention_scale: f32,
    sliding_window_tokens: u64,
    rope_interleaved: bool,
    output_gate: bool,
    value_rms_norm: bool,
    attention_k_eq_v: bool,
    post_attention_norm: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CausalKvLayout {
    TokenMajorPages,
    VllmBlocks16 {
        combined_block_bytes: u64,
        blocks_per_page: u64,
    },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CausalAttentionKernelPath {
    TokenMajorFallback,
    VllmAddressedFallback,
    VllmAddressedVarlen,
    VllmAddressedVarlenTiled,
    VllmAddressedDecodeV1,
    VllmAddressedDecodeV2,
}

impl CausalAttentionKernelPath {
    fn select(
        attention_policy: AttentionExecutionPolicy,
        shape: CausalAttentionShape,
        active_tokens: u64,
        sequence_tokens: u64,
    ) -> Result<Self, String> {
        if matches!(shape.kv_layout()?, CausalKvLayout::TokenMajorPages) {
            return Ok(Self::TokenMajorFallback);
        }
        if active_tokens == 1 && shape.tiled_vllm_supported()? {
            return match attention_policy {
                AttentionExecutionPolicy::Portable => Ok(Self::VllmAddressedFallback),
                AttentionExecutionPolicy::NativeAdaptive => {
                    #[cfg(feature = "vllm-paged-attn-v2")]
                    {
                        Ok(
                            match VnextAddressedPagedAttentionKernel::for_sequence_length(
                                sequence_tokens,
                            ) {
                                VnextAddressedPagedAttentionKernel::V1 => {
                                    Self::VllmAddressedDecodeV1
                                }
                                VnextAddressedPagedAttentionKernel::V2 => {
                                    Self::VllmAddressedDecodeV2
                                }
                            },
                        )
                    }
                    #[cfg(not(feature = "vllm-paged-attn-v2"))]
                    {
                        Err(
                            "native-adaptive causal attention requires the compiled vLLM paged-attention provider"
                                .to_owned(),
                        )
                    }
                }
                AttentionExecutionPolicy::Auto => {
                    Err("causal attention received an unresolved auto policy".to_owned())
                }
            };
        }
        // The addressed varlen kernels currently support arbitrary scale but
        // only head dimensions through 256 and full-prefix attention. Gemma 4
        // local layers and 512-wide full layers therefore stay on the generic
        // addressed fallback; this also keeps them away from the vLLM decode
        // kernels whose scale and head-size ABI are fixed.
        if !shape.addressed_varlen_supported() {
            return Ok(Self::VllmAddressedFallback);
        }
        let score_bytes = sequence_tokens
            .checked_mul(std::mem::size_of::<f32>() as u64)
            .ok_or_else(|| "causal attention varlen score bytes overflow".to_owned())?;
        if active_tokens >= VARLEN_TILED_QUERY_TOKENS
            && score_bytes
                .checked_mul(VARLEN_TILED_QUERY_TOKENS)
                .is_some_and(|bytes| bytes <= VARLEN_DYNAMIC_SHARED_BUDGET_BYTES)
        {
            Ok(Self::VllmAddressedVarlenTiled)
        } else if score_bytes <= VARLEN_DYNAMIC_SHARED_BUDGET_BYTES {
            Ok(Self::VllmAddressedVarlen)
        } else {
            Ok(Self::VllmAddressedFallback)
        }
    }

    fn operation(self) -> &'static str {
        match self {
            Self::TokenMajorFallback => COMPUTE_TOKEN_MAJOR_OPERATION,
            Self::VllmAddressedFallback => COMPUTE_VLLM_FALLBACK_OPERATION,
            Self::VllmAddressedVarlen => COMPUTE_VLLM_VARLEN_OPERATION,
            Self::VllmAddressedVarlenTiled => COMPUTE_VLLM_VARLEN_TILED_OPERATION,
            Self::VllmAddressedDecodeV1 => COMPUTE_VLLM_DECODE_V1_OPERATION,
            Self::VllmAddressedDecodeV2 => COMPUTE_VLLM_DECODE_V2_OPERATION,
        }
    }

    fn native_kernel_id(self) -> &'static str {
        match self {
            Self::TokenMajorFallback => "ferrum.vnext_causal_attention.token_major",
            Self::VllmAddressedFallback => "ferrum.vnext_causal_attention.vllm_addressed",
            Self::VllmAddressedVarlen => "ferrum.paged_varlen_attention.vllm_addressed",
            Self::VllmAddressedVarlenTiled => "ferrum.paged_varlen_attention.vllm_q4_addressed",
            Self::VllmAddressedDecodeV1 => "vllm.paged_attention_v1.addressed",
            Self::VllmAddressedDecodeV2 => "vllm.paged_attention_v2.addressed",
        }
    }

    fn replay_id(self) -> u64 {
        match self {
            Self::TokenMajorFallback => 0,
            Self::VllmAddressedFallback => 1,
            Self::VllmAddressedVarlen => 2,
            Self::VllmAddressedVarlenTiled => 3,
            Self::VllmAddressedDecodeV1 => 4,
            Self::VllmAddressedDecodeV2 => 5,
        }
    }

    fn attention_dispatch_count(self) -> u64 {
        match self {
            Self::VllmAddressedDecodeV2 => 2,
            _ => 1,
        }
    }

    fn uses_vllm_layout(self) -> bool {
        !matches!(self, Self::TokenMajorFallback)
    }

    fn is_fallback(self) -> bool {
        matches!(self, Self::TokenMajorFallback | Self::VllmAddressedFallback)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct CausalAttentionReplayEnvelope {
    sequence_capacity_tokens: u64,
    table_capacity_entries: i32,
}

impl CausalAttentionReplayEnvelope {
    fn new(
        shape: CausalAttentionShape,
        path: CausalAttentionKernelPath,
        sequence_tokens: u64,
    ) -> Result<Self, String> {
        let sequence_capacity_tokens = match path {
            CausalAttentionKernelPath::VllmAddressedDecodeV1 => {
                VLLM_PARTITION_TOKENS.min(shape.maximum_context_tokens)
            }
            CausalAttentionKernelPath::VllmAddressedDecodeV2 => sequence_tokens
                .div_ceil(VLLM_PARTITION_TOKENS)
                .checked_mul(VLLM_PARTITION_TOKENS)
                .map(|capacity| capacity.min(shape.maximum_context_tokens))
                .ok_or_else(|| "causal attention replay sequence capacity overflows".to_owned())?,
            CausalAttentionKernelPath::TokenMajorFallback
            | CausalAttentionKernelPath::VllmAddressedFallback => shape.maximum_context_tokens,
            CausalAttentionKernelPath::VllmAddressedVarlen
            | CausalAttentionKernelPath::VllmAddressedVarlenTiled => sequence_tokens,
        };
        if sequence_tokens == 0
            || sequence_tokens > sequence_capacity_tokens
            || sequence_capacity_tokens > shape.maximum_context_tokens
        {
            return Err("causal attention replay sequence capacity is invalid".to_owned());
        }
        Ok(Self {
            sequence_capacity_tokens,
            table_capacity_entries: checked_i32(
                shape.table_entries(sequence_capacity_tokens)?,
                "causal attention replay table capacity",
            )?,
        })
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CausalAttentionReplayTopology {
    PartitionStable(CausalAttentionReplayEnvelope),
    ExactShapeEager(CausalAttentionReplayEnvelope),
}

impl CausalAttentionReplayTopology {
    fn new(
        shape: CausalAttentionShape,
        path: CausalAttentionKernelPath,
        sequence_tokens: u64,
    ) -> Result<Self, String> {
        let envelope = CausalAttentionReplayEnvelope::new(shape, path, sequence_tokens)?;
        Ok(match path {
            CausalAttentionKernelPath::TokenMajorFallback
            | CausalAttentionKernelPath::VllmAddressedFallback
            | CausalAttentionKernelPath::VllmAddressedDecodeV1
            | CausalAttentionKernelPath::VllmAddressedDecodeV2 => Self::PartitionStable(envelope),
            CausalAttentionKernelPath::VllmAddressedVarlen
            | CausalAttentionKernelPath::VllmAddressedVarlenTiled => {
                Self::ExactShapeEager(envelope)
            }
        })
    }

    const fn envelope(self) -> CausalAttentionReplayEnvelope {
        match self {
            Self::PartitionStable(envelope) | Self::ExactShapeEager(envelope) => envelope,
        }
    }

    const fn is_partition_stable(self) -> bool {
        matches!(self, Self::PartitionStable(_))
    }
}

fn reusable_attention_topology(
    request: &ReusableExecutionTopologyRequest<'_>,
    attention_policy: AttentionExecutionPolicy,
    semantics: CausalAttentionSemantics,
) -> Result<ReusableExecutionTopology, String> {
    let shape = CausalAttentionShape::from_attributes_for(request.attributes(), semantics)?;
    let ranges = request.work_shape().participant_token_ranges();
    if ranges.is_empty() {
        return Err("CUDA causal topology has no participant token ranges".to_owned());
    }

    reusable_attention_topology_from_rows(
        attention_policy,
        shape,
        ranges.len(),
        ranges.iter().map(|range| {
            let source = range.source_token_range();
            if source.end > range.full_input_tokens()
                || range.full_input_tokens() > shape.maximum_context_tokens
            {
                return Err("causal topology token range exceeds its admitted context".to_owned());
            }
            Ok(CausalAttentionTopologyRow {
                active_tokens: range.immediate_tokens(),
                sequence_tokens: source.end,
            })
        }),
    )
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct CausalAttentionTopologyRow {
    active_tokens: u64,
    sequence_tokens: u64,
}

fn reusable_attention_topology_from_rows<I>(
    attention_policy: AttentionExecutionPolicy,
    shape: CausalAttentionShape,
    row_count: usize,
    rows: I,
) -> Result<ReusableExecutionTopology, String>
where
    I: IntoIterator<Item = Result<CausalAttentionTopologyRow, String>>,
{
    if row_count == 0 {
        return Err("CUDA causal topology has no participants".to_owned());
    }

    const DOMAIN: &[u8] = b"ferrum.cuda.causal-attention.reusable-topology.v1\0";
    let mut digest = Sha256::new();
    digest.update(DOMAIN);
    digest.update((row_count as u64).to_le_bytes());
    let mut observed_rows = 0_usize;
    let mut total_tokens = 0_u64;
    let mut partition_stable = true;
    for row in rows {
        let row = row?;
        observed_rows = observed_rows
            .checked_add(1)
            .ok_or_else(|| "CUDA causal topology participant count overflows".to_owned())?;
        total_tokens = total_tokens
            .checked_add(row.active_tokens)
            .ok_or_else(|| "CUDA causal topology token count overflows".to_owned())?;
        let path = CausalAttentionKernelPath::select(
            attention_policy,
            shape,
            row.active_tokens,
            row.sequence_tokens,
        )?;
        let topology = CausalAttentionReplayTopology::new(shape, path, row.sequence_tokens)?;
        partition_stable &= topology.is_partition_stable();
        let envelope = topology.envelope();
        digest.update(row.active_tokens.to_le_bytes());
        digest.update(path.replay_id().to_le_bytes());
        digest.update(envelope.sequence_capacity_tokens.to_le_bytes());
        digest.update(envelope.table_capacity_entries.to_le_bytes());
        digest.update([u8::from(topology.is_partition_stable())]);
    }
    if observed_rows != row_count {
        return Err("CUDA causal topology participant count changed while hashing".to_owned());
    }
    if !partition_stable {
        return Ok(ReusableExecutionTopology::EagerBoundary);
    }
    digest.update(total_tokens.to_le_bytes());
    Ok(ReusableExecutionTopology::Dynamic(
        DeviceReusableExecutionTopologyFingerprint::from_sha256(digest.finalize().into()),
    ))
}

impl CausalAttentionShape {
    fn from_attributes(attributes: &BTreeMap<AttributeId, SemanticValue>) -> Result<Self, String> {
        Self::from_attributes_for(attributes, CausalAttentionSemantics::Standard)
    }

    fn from_attributes_for(
        attributes: &BTreeMap<AttributeId, SemanticValue>,
        semantics: CausalAttentionSemantics,
    ) -> Result<Self, String> {
        let head_dim = unsigned_attribute(attributes, "head_dim")?;
        let rope_dim = unsigned_attribute(attributes, "rope_dim")?;
        let (
            rope_frequency_denominator,
            attention_scale,
            sliding_window_tokens,
            output_gate,
            value_rms_norm,
            attention_k_eq_v,
            post_attention_norm,
        ) = match semantics {
            CausalAttentionSemantics::Standard => (
                rope_dim,
                1.0_f32 / (head_dim as f32).sqrt(),
                0,
                bool_attribute(attributes, "output_gate")?,
                false,
                false,
                false,
            ),
            CausalAttentionSemantics::Gemma4 => (
                unsigned_attribute(attributes, "rope_frequency_denominator")?,
                rational_attribute(attributes, "attention_scale")?,
                unsigned_attribute(attributes, "sliding_window_tokens")?,
                false,
                bool_attribute(attributes, "value_rms_norm")?,
                bool_attribute(attributes, "attention_k_eq_v")?,
                true,
            ),
        };
        let shape = Self {
            hidden_size: unsigned_attribute(attributes, "hidden_size")?,
            query_heads: unsigned_attribute(attributes, "query_heads")?,
            key_value_heads: unsigned_attribute(attributes, "key_value_heads")?,
            head_dim,
            query_features: unsigned_attribute(attributes, "query_features")?,
            query_projection_features: unsigned_attribute(attributes, "query_projection_features")?,
            kv_features: unsigned_attribute(attributes, "kv_features")?,
            rope_dim,
            rope_frequency_denominator,
            maximum_context_tokens: unsigned_attribute(attributes, "maximum_context_tokens")?,
            epsilon: rational_attribute(attributes, "epsilon")?,
            rope_theta: rational_attribute(attributes, "rope_theta")?,
            attention_scale,
            sliding_window_tokens,
            rope_interleaved: bool_attribute(attributes, "rope_interleaved")?,
            output_gate,
            value_rms_norm,
            attention_k_eq_v,
            post_attention_norm,
        };
        if !bool_attribute(attributes, "causal")? {
            return Err("causal attention requires causal=true".to_owned());
        }
        let query_features = shape
            .query_heads
            .checked_mul(shape.head_dim)
            .ok_or_else(|| "causal attention query width overflows".to_owned())?;
        let kv_features = shape
            .key_value_heads
            .checked_mul(shape.head_dim)
            .ok_or_else(|| "causal attention KV width overflows".to_owned())?;
        let projection_multiplier = if shape.output_gate { 2 } else { 1 };
        let query_projection_features = query_features
            .checked_mul(projection_multiplier)
            .ok_or_else(|| "causal attention query projection width overflows".to_owned())?;
        if shape.hidden_size == 0
            || shape.query_heads == 0
            || shape.key_value_heads == 0
            || shape.head_dim == 0
            || shape.rope_dim == 0
            || shape.rope_frequency_denominator == 0
            || shape.maximum_context_tokens == 0
            || !shape.attention_scale.is_finite()
            || shape.attention_scale <= 0.0
        {
            return Err("causal attention dimensions and context must be non-zero".to_owned());
        }
        checked_i32(
            shape.maximum_context_tokens,
            "causal attention maximum context",
        )?;
        let maximum_head_dim = match semantics {
            CausalAttentionSemantics::Standard => MAXIMUM_STANDARD_HEAD_DIM,
            CausalAttentionSemantics::Gemma4 => MAXIMUM_HEAD_DIM,
        };
        if shape.query_heads % shape.key_value_heads != 0
            || shape.head_dim > maximum_head_dim
            || shape.rope_dim > shape.head_dim
            || shape.rope_dim % 2 != 0
            || shape.rope_frequency_denominator < shape.rope_dim
            || shape.rope_frequency_denominator > shape.head_dim
            || shape.rope_frequency_denominator % 2 != 0
            || shape.sliding_window_tokens > shape.maximum_context_tokens
            || shape.query_features != query_features
            || shape.kv_features != kv_features
            || shape.query_projection_features != query_projection_features
        {
            return Err("causal attention attributes are inconsistent".to_owned());
        }
        shape.cuda_shape()?;
        shape.maximum_pages()?;
        Ok(shape)
    }

    fn state_bytes_per_token(self) -> Result<u64, String> {
        self.kv_features
            .checked_mul(2)
            .and_then(|elements| elements.checked_mul(ElementType::F16.size_bytes()))
            .ok_or_else(|| "causal attention KV bytes per token overflow".to_owned())
    }

    fn kv_layout(self) -> Result<CausalKvLayout, String> {
        let combined_block_bytes = self
            .state_bytes_per_token()?
            .checked_mul(VLLM_BLOCK_TOKENS)
            .ok_or_else(|| "causal attention vLLM block size overflows".to_owned())?;
        if self.head_dim % 8 != 0
            || combined_block_bytes > VNEXT_KV_PAGE_BYTES
            || VNEXT_KV_PAGE_BYTES % combined_block_bytes != 0
        {
            return Ok(CausalKvLayout::TokenMajorPages);
        }
        Ok(CausalKvLayout::VllmBlocks16 {
            combined_block_bytes,
            blocks_per_page: VNEXT_KV_PAGE_BYTES / combined_block_bytes,
        })
    }

    fn table_entries(self, tokens: u64) -> Result<u64, String> {
        if tokens == 0 {
            return Err("causal attention table cannot describe zero tokens".to_owned());
        }
        match self.kv_layout()? {
            CausalKvLayout::TokenMajorPages => {
                Ok(self.physical_state_bytes(tokens)? / VNEXT_KV_PAGE_BYTES)
            }
            CausalKvLayout::VllmBlocks16 { .. } => Ok(tokens.div_ceil(VLLM_BLOCK_TOKENS)),
        }
    }

    fn physical_state_bytes(self, tokens: u64) -> Result<u64, String> {
        if tokens == 0 {
            return Err("causal attention state cannot describe zero tokens".to_owned());
        }
        match self.kv_layout()? {
            CausalKvLayout::TokenMajorPages => {
                let logical = self
                    .state_bytes_per_token()?
                    .checked_mul(tokens)
                    .ok_or_else(|| "causal attention KV state size overflows".to_owned())?;
                align_up(logical, VNEXT_KV_PAGE_BYTES)
            }
            CausalKvLayout::VllmBlocks16 {
                blocks_per_page, ..
            } => tokens
                .div_ceil(VLLM_BLOCK_TOKENS)
                .div_ceil(blocks_per_page)
                .checked_mul(VNEXT_KV_PAGE_BYTES)
                .ok_or_else(|| "causal attention vLLM-layout state size overflows".to_owned()),
        }
    }

    fn physical_state_bytes_for_source_frontier(
        self,
        source_end_tokens: u64,
        full_input_tokens: u64,
    ) -> Result<u64, String> {
        if source_end_tokens == 0 || source_end_tokens > full_input_tokens {
            return Err("causal attention source frontier exceeds its full input".to_owned());
        }
        self.physical_state_bytes(source_end_tokens)
    }

    fn maximum_pages(self) -> Result<u64, String> {
        Ok(self.physical_state_bytes(self.maximum_context_tokens)? / VNEXT_KV_PAGE_BYTES)
    }

    fn binding_slot_bytes(self) -> Result<u64, String> {
        BINDING_CONTROL_BYTES
            .checked_add(aligned_bytes(
                self.table_entries(self.maximum_context_tokens)?,
                POINTER_BYTES,
            )?)
            .ok_or_else(|| "causal attention binding slot size overflows".to_owned())
    }

    fn scratch_bytes_per_token(self) -> Result<u64, String> {
        [
            self.hidden_size,
            self.query_projection_features,
            self.kv_features,
            self.kv_features,
            self.query_features,
            self.query_features,
            self.hidden_size,
        ]
        .into_iter()
        .try_fold(0_u64, |total, elements| {
            total
                .checked_add(aligned_bytes(elements, ElementType::F16.size_bytes())?)
                .ok_or_else(|| "causal attention token scratch size overflows".to_owned())
        })
    }

    fn vllm_scratch_bytes(self) -> Result<u64, String> {
        if !self.tiled_vllm_supported()? {
            return Ok(0);
        }
        let partitions = self.maximum_context_tokens.div_ceil(VLLM_PARTITION_TOKENS);
        let rows = self
            .query_heads
            .checked_mul(partitions)
            .ok_or_else(|| "causal attention vLLM partition rows overflow".to_owned())?;
        let statistics = aligned_bytes(rows, std::mem::size_of::<f32>() as u64)?;
        let temporary = aligned_bytes(
            rows.checked_mul(self.head_dim)
                .ok_or_else(|| "causal attention vLLM temporary rows overflow".to_owned())?,
            ElementType::F16.size_bytes(),
        )?;
        statistics
            .checked_mul(2)
            .and_then(|bytes| bytes.checked_add(temporary))
            .ok_or_else(|| "causal attention vLLM scratch size overflows".to_owned())
    }

    fn attention_policy_scratch_bytes(
        self,
        attention_policy: AttentionExecutionPolicy,
    ) -> Result<u64, String> {
        match attention_policy {
            AttentionExecutionPolicy::Portable => Ok(0),
            AttentionExecutionPolicy::NativeAdaptive => self.vllm_scratch_bytes(),
            AttentionExecutionPolicy::Auto => {
                Err("causal attention scratch received an unresolved auto policy".to_owned())
            }
        }
    }

    fn tiled_vllm_supported(self) -> Result<bool, String> {
        Ok(cfg!(feature = "vllm-paged-attn-v2")
            && matches!(self.kv_layout()?, CausalKvLayout::VllmBlocks16 { .. })
            && matches!(self.head_dim, 128 | 256)
            && self.sliding_window_tokens == 0
            && self.attention_scale.to_bits()
                == (1.0_f32 / (self.head_dim as f32).sqrt()).to_bits())
    }

    fn addressed_varlen_supported(self) -> bool {
        self.head_dim <= MAXIMUM_VARLEN_HEAD_DIM && self.sliding_window_tokens == 0
    }

    fn cuda_shape(self) -> Result<CudaCausalAttentionShape, String> {
        Ok(CudaCausalAttentionShape {
            hidden_size: checked_i32(self.hidden_size, "causal attention hidden size")?,
            query_heads: checked_i32(self.query_heads, "causal attention query heads")?,
            key_value_heads: checked_i32(self.key_value_heads, "causal attention key/value heads")?,
            head_dim: checked_i32(self.head_dim, "causal attention head dimension")?,
            query_features: checked_i32(self.query_features, "causal attention query width")?,
            query_projection_features: checked_i32(
                self.query_projection_features,
                "causal attention query projection width",
            )?,
            kv_features: checked_i32(self.kv_features, "causal attention KV width")?,
            rope_dim: checked_i32(self.rope_dim, "causal attention RoPE width")?,
            rope_frequency_denominator: checked_i32(
                self.rope_frequency_denominator,
                "causal attention RoPE frequency denominator",
            )?,
            epsilon: self.epsilon,
            rope_theta: self.rope_theta,
            attention_scale: self.attention_scale,
            sliding_window_tokens: checked_i32(
                self.sliding_window_tokens,
                "causal attention sliding window",
            )?,
            rope_interleaved: i32::from(self.rope_interleaved),
            output_gate: i32::from(self.output_gate),
            value_rms_norm: i32::from(self.value_rms_norm),
        })
    }
}

#[derive(Debug, Clone, Copy)]
struct CudaCausalAttentionShape {
    hidden_size: i32,
    query_heads: i32,
    key_value_heads: i32,
    head_dim: i32,
    query_features: i32,
    query_projection_features: i32,
    kv_features: i32,
    rope_dim: i32,
    rope_frequency_denominator: i32,
    epsilon: f32,
    rope_theta: f32,
    attention_scale: f32,
    sliding_window_tokens: i32,
    rope_interleaved: i32,
    output_gate: i32,
    value_rms_norm: i32,
}

#[cfg(feature = "vllm-marlin")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CausalProjectionWeightAbi {
    F16,
    MarlinFp8,
    GptqMarlinInt4,
    CompressedTensorsMarlinInt4,
    CompressedTensorsMarlinSymmetricInt4,
}

#[cfg(feature = "vllm-marlin")]
fn causal_projection_weight_abi(
    quantization_formats: &BTreeSet<QuantizationFormatId>,
) -> Result<CausalProjectionWeightAbi, String> {
    let mut formats = quantization_formats.iter();
    let Some(format) = formats.next() else {
        return Ok(CausalProjectionWeightAbi::F16);
    };
    if formats.next().is_some() {
        return Err("causal attention projection has more than one quantization format".to_owned());
    }
    match format.as_str() {
        MARLIN_FP8_QUANTIZATION_FORMAT_ID | MARLIN_FP8_GROUP128_QUANTIZATION_FORMAT_ID => {
            Ok(CausalProjectionWeightAbi::MarlinFp8)
        }
        GPTQ_MARLIN_QUANTIZATION_FORMAT_ID => Ok(CausalProjectionWeightAbi::GptqMarlinInt4),
        COMPRESSED_TENSORS_MARLIN_QUANTIZATION_FORMAT_ID => {
            Ok(CausalProjectionWeightAbi::CompressedTensorsMarlinInt4)
        }
        COMPRESSED_TENSORS_MARLIN_SYMMETRIC_QUANTIZATION_FORMAT_ID => {
            Ok(CausalProjectionWeightAbi::CompressedTensorsMarlinSymmetricInt4)
        }
        unknown => Err(format!(
            "causal attention projection quantization format `{unknown}` is not admitted"
        )),
    }
}

#[cfg(feature = "vllm-marlin")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CausalProjectionAbi {
    F16,
    MarlinFp8,
    MarlinInt4,
}

#[cfg(feature = "vllm-marlin")]
fn causal_projection_abi(
    weights: impl IntoIterator<Item = CausalProjectionWeightAbi>,
) -> Result<CausalProjectionAbi, String> {
    let mut uses_fp8 = false;
    let mut uses_int4 = false;
    for weight in weights {
        match weight {
            CausalProjectionWeightAbi::F16 => {}
            CausalProjectionWeightAbi::MarlinFp8 => uses_fp8 = true,
            CausalProjectionWeightAbi::GptqMarlinInt4
            | CausalProjectionWeightAbi::CompressedTensorsMarlinInt4
            | CausalProjectionWeightAbi::CompressedTensorsMarlinSymmetricInt4 => uses_int4 = true,
        }
    }
    if uses_fp8 && uses_int4 {
        return Err(
            "causal attention cannot mix FP8 and INT4 projection ABIs in one operation".to_owned(),
        );
    }
    Ok(if uses_fp8 {
        CausalProjectionAbi::MarlinFp8
    } else if uses_int4 {
        CausalProjectionAbi::MarlinInt4
    } else {
        CausalProjectionAbi::F16
    })
}

#[derive(Debug, Clone, Copy)]
enum CausalProjection {
    F16,
    #[cfg(feature = "vllm-marlin")]
    MarlinFp8 {
        runtime: MarlinProjectionRuntime,
    },
    #[cfg(feature = "vllm-marlin")]
    MarlinInt4 {
        runtime: MarlinProjectionRuntime,
    },
}

impl CausalProjection {
    #[cfg(feature = "vllm-marlin")]
    fn from_values(
        values: &[ResolvedValueBinding],
        runtime: MarlinProjectionRuntime,
    ) -> Result<Self, String> {
        let mut weights = Vec::with_capacity(4);
        for ordinal in 2..=5 {
            let value = binding(values, ResolvedValueRole::Input, ordinal)?;
            let weight = value.weight().ok_or_else(|| {
                format!(
                    "causal attention projection input {ordinal} lacks its physical weight layout"
                )
            })?;
            weights.push(causal_projection_weight_abi(
                &weight.quantization_formats(),
            )?);
        }
        Ok(match causal_projection_abi(weights)? {
            CausalProjectionAbi::F16 => Self::F16,
            CausalProjectionAbi::MarlinFp8 => Self::MarlinFp8 { runtime },
            CausalProjectionAbi::MarlinInt4 => Self::MarlinInt4 { runtime },
        })
    }

    fn workspace_bytes(self) -> Result<u64, String> {
        match self {
            Self::F16 => Ok(0),
            #[cfg(feature = "vllm-marlin")]
            Self::MarlinFp8 { runtime } | Self::MarlinInt4 { runtime } => runtime.workspace_bytes(),
        }
    }

    fn workspace_reservation_bytes(self) -> Result<u64, String> {
        super::aligned_projection_workspace_bytes(
            self.workspace_bytes()?,
            SCRATCH_ALIGNMENT,
            "causal attention projection workspace",
        )
    }

    fn replay_tag(self) -> &'static str {
        match self {
            Self::F16 => "f16-cublas",
            #[cfg(feature = "vllm-marlin")]
            Self::MarlinFp8 { .. } => "mixed-fp8-marlin-f16-reduce",
            #[cfg(feature = "vllm-marlin")]
            Self::MarlinInt4 { .. } => "mixed-int4-marlin-f16-reduce",
        }
    }
}

#[derive(Debug, Clone, Copy)]
struct ScratchLayout {
    required_bytes: u64,
    projection_workspace: Option<u64>,
    normalized: u64,
    query_raw: u64,
    key_raw: u64,
    value_raw: u64,
    query: u64,
    context: u64,
    projected: u64,
    vllm: Option<VllmScratchLayout>,
}

#[derive(Debug, Clone, Copy)]
struct VllmScratchLayout {
    exp_sums: u64,
    max_logits: u64,
    temporary_output: u64,
}

impl ScratchLayout {
    fn new(
        shape: CausalAttentionShape,
        total_tokens: u64,
        projection: CausalProjection,
        attention_policy: AttentionExecutionPolicy,
    ) -> Result<Self, String> {
        if total_tokens == 0 {
            return Err("causal attention scratch cannot be sized for empty work".to_owned());
        }
        let mut offset = 0;
        let projection_workspace_bytes = projection.workspace_bytes()?;
        let projection_workspace_reservation_bytes = projection.workspace_reservation_bytes()?;
        let projection_workspace = (projection_workspace_bytes > 0)
            .then(|| reserve_elements(&mut offset, projection_workspace_bytes, 1))
            .transpose()?;
        let normalized = reserve_tokens(&mut offset, shape.hidden_size, total_tokens)?;
        let query_raw = reserve_tokens(&mut offset, shape.query_projection_features, total_tokens)?;
        let key_raw = reserve_tokens(&mut offset, shape.kv_features, total_tokens)?;
        let value_raw = reserve_tokens(&mut offset, shape.kv_features, total_tokens)?;
        let query = reserve_tokens(&mut offset, shape.query_features, total_tokens)?;
        let context = reserve_tokens(&mut offset, shape.query_features, total_tokens)?;
        let projected = reserve_tokens(&mut offset, shape.hidden_size, total_tokens)?;
        let attention_policy_scratch_bytes =
            shape.attention_policy_scratch_bytes(attention_policy)?;
        let vllm = if attention_policy_scratch_bytes == 0 {
            None
        } else {
            let partitions = shape.maximum_context_tokens.div_ceil(VLLM_PARTITION_TOKENS);
            let rows = shape
                .query_heads
                .checked_mul(partitions)
                .ok_or_else(|| "causal attention vLLM partition rows overflow".to_owned())?;
            let exp_sums = reserve_elements(&mut offset, rows, std::mem::size_of::<f32>() as u64)?;
            let max_logits =
                reserve_elements(&mut offset, rows, std::mem::size_of::<f32>() as u64)?;
            let temporary_output = reserve_elements(
                &mut offset,
                rows.checked_mul(shape.head_dim)
                    .ok_or_else(|| "causal attention vLLM temporary rows overflow".to_owned())?,
                ElementType::F16.size_bytes(),
            )?;
            Some(VllmScratchLayout {
                exp_sums,
                max_logits,
                temporary_output,
            })
        };
        let token_bytes = shape
            .scratch_bytes_per_token()?
            .checked_mul(total_tokens)
            .ok_or_else(|| "causal attention scratch size overflows".to_owned())?;
        let expected = token_bytes
            .checked_add(attention_policy_scratch_bytes)
            .and_then(|bytes| bytes.checked_add(projection_workspace_reservation_bytes))
            .ok_or_else(|| "causal attention scratch size overflows".to_owned())?;
        if offset != expected {
            return Err("causal attention scratch layout differs from its estimate".to_owned());
        }
        Ok(Self {
            required_bytes: offset,
            projection_workspace,
            normalized,
            query_raw,
            key_raw,
            value_raw,
            query,
            context,
            projected,
            vllm,
        })
    }

    fn token_offset(self, base: u64, token_start: u64, width: u64) -> Result<u64, String> {
        base.checked_add(
            width
                .checked_mul(ElementType::F16.size_bytes())
                .ok_or_else(|| "causal attention token scratch stride overflows".to_owned())?
                .checked_mul(token_start)
                .ok_or_else(|| "causal attention token scratch offset overflows".to_owned())?,
        )
        .filter(|offset| *offset < self.required_bytes)
        .ok_or_else(|| "causal attention token scratch range is invalid".to_owned())
    }
}

#[derive(Debug, Clone, Copy)]
struct BindingLayout {
    required_bytes: u64,
    slot_bytes: u64,
}

impl BindingLayout {
    fn new(shape: CausalAttentionShape, participant_count: usize) -> Result<Self, String> {
        if participant_count == 0 {
            return Err("causal attention binding cannot be sized for empty work".to_owned());
        }
        let participant_count = u64::try_from(participant_count)
            .map_err(|_| "causal attention participant count exceeds u64".to_owned())?;
        let slot_bytes = shape.binding_slot_bytes()?;
        let required_bytes = slot_bytes
            .checked_mul(participant_count)
            .ok_or_else(|| "causal attention binding workspace size overflows".to_owned())?;
        Ok(Self {
            required_bytes,
            slot_bytes,
        })
    }

    fn binding_offset(self, participant: usize) -> Result<u64, String> {
        self.slot_bytes
            .checked_mul(
                u64::try_from(participant)
                    .map_err(|_| "causal attention participant index exceeds u64".to_owned())?,
            )
            .filter(|offset| *offset < self.required_bytes)
            .ok_or_else(|| "causal attention binding offset exceeds its workspace".to_owned())
    }
}

#[derive(Debug, Clone, Copy)]
struct SharedRegions {
    input_norm: usize,
    query_weight: SharedProjectionWeight,
    key_weight: SharedProjectionWeight,
    value_weight: SharedProjectionWeight,
    output_weight: SharedProjectionWeight,
    query_norm: usize,
    key_norm: usize,
    post_attention_norm: Option<usize>,
    scratch: usize,
    binding: usize,
}

#[derive(Debug, Clone, Copy)]
enum SharedProjectionWeight {
    F16 {
        region: usize,
    },
    #[cfg(feature = "vllm-marlin")]
    Marlin {
        packed_region: usize,
        scales_region: usize,
        zero_points_region: Option<usize>,
        group_size: i32,
        weight_type: MarlinF16WeightType,
    },
}

impl SharedProjectionWeight {
    fn replay_tag(self) -> &'static str {
        match self {
            Self::F16 { .. } => "f16",
            #[cfg(feature = "vllm-marlin")]
            Self::Marlin { weight_type, .. } => marlin_projection_replay_tag(weight_type),
        }
    }

    fn bind_replay_abi(
        self,
        replay_key: CudaCommandReplayKeyBuilder,
    ) -> CudaCommandReplayKeyBuilder {
        let replay_key = replay_key.bytes(self.replay_tag().as_bytes());
        match self {
            Self::F16 { .. } => replay_key.i32(0),
            #[cfg(feature = "vllm-marlin")]
            Self::Marlin { group_size, .. } => replay_key.i32(group_size),
        }
    }
}

#[cfg(feature = "vllm-marlin")]
const fn marlin_projection_replay_tag(weight_type: MarlinF16WeightType) -> &'static str {
    match weight_type {
        MarlinF16WeightType::E4M3Fn => "marlin-fp8-e4m3fn",
        MarlinF16WeightType::U4 => "compressed-tensors-marlin-u4-zp",
        MarlinF16WeightType::U4B8 => "gptq-marlin-u4b8",
    }
}

#[derive(Debug, Clone, Copy)]
struct CausalAttentionLaunch {
    input_region: usize,
    output_region: usize,
    binding_offset: u64,
    packed_token_start: u64,
    packed_query_raw: u64,
    packed_key_raw: u64,
    packed_value_raw: u64,
    packed_query: u64,
    packed_context: u64,
    tokens: u64,
    tokens_i32: i32,
    sequence_tokens: u64,
    sequence_tokens_i32: i32,
    table_entries_i32: i32,
    replay_topology: CausalAttentionReplayTopology,
    path: CausalAttentionKernelPath,
}

#[derive(Debug, Clone, Copy)]
struct PackedCausalAttentionLaunch {
    input_region: usize,
    output_region: usize,
    tokens: u64,
    tokens_i32: i32,
}

#[derive(Debug, Clone, Copy)]
struct PackedFallbackLaunch {
    token_grid: u64,
    packed_token_grid: u64,
    participant_grid: u32,
    participant_count_i32: i32,
    binding_slot_bytes: u64,
    path: CausalAttentionKernelPath,
}

fn packed_fallback_launch(
    launches: &[CausalAttentionLaunch],
    binding_layout: BindingLayout,
) -> Result<Option<PackedFallbackLaunch>, String> {
    let Some(first) = launches.first() else {
        return Err("packed causal attention has no participants".to_owned());
    };
    if launches.len() <= 1
        || !first.path.is_fallback()
        || launches.iter().any(|launch| launch.path != first.path)
    {
        return Ok(None);
    }
    let token_grid = launches
        .iter()
        .map(|launch| launch.tokens)
        .max()
        .ok_or_else(|| "packed causal fallback has no token grid".to_owned())?;
    let packed_token_grid = launches.iter().try_fold(0_u64, |total, launch| {
        total
            .checked_add(launch.tokens)
            .ok_or_else(|| "packed causal fallback token grid overflows".to_owned())
    })?;
    let participant_grid = u32::try_from(launches.len())
        .map_err(|_| "packed causal fallback participant grid exceeds u32".to_owned())?;
    Ok(Some(PackedFallbackLaunch {
        token_grid,
        packed_token_grid,
        participant_grid,
        participant_count_i32: i32::try_from(participant_grid)
            .map_err(|_| "packed causal fallback participant count exceeds i32".to_owned())?,
        binding_slot_bytes: binding_layout.slot_bytes,
        path: first.path,
    }))
}

#[derive(Debug, Clone, Copy)]
struct CausalAttentionBinding {
    first_page_region: usize,
    page_count: usize,
    host_binding: usize,
    binding_offset: u64,
}

fn encode_attention(
    functions: &CausalAttentionFunctions,
    provider_fingerprint: &str,
    attention_policy: AttentionExecutionPolicy,
    semantics: CausalAttentionSemantics,
    operation_id: &str,
    #[cfg(feature = "vllm-marlin")] projection_runtime: MarlinProjectionRuntime,
    invocation: BatchedOperationInvocation<'_, CudaDeviceBuffer>,
) -> Result<EncodedDeviceOperation<CudaDeviceCommand>, String> {
    if invocation.participants().is_empty() || invocation.operation().id.as_str() != operation_id {
        return Err("CUDA causal attention received another or empty operation".to_owned());
    }
    let first = &invocation.participants()[0];
    let shape = CausalAttentionShape::from_attributes_for(first.attributes(), semantics)?;
    #[cfg(feature = "vllm-marlin")]
    let projection = CausalProjection::from_values(first.bindings(), projection_runtime)?;
    #[cfg(not(feature = "vllm-marlin"))]
    let projection = CausalProjection::F16;
    validate_signature(first, shape, semantics)?;
    for participant in &invocation.participants()[1..] {
        if CausalAttentionShape::from_attributes_for(participant.attributes(), semantics)? != shape
        {
            return Err("CUDA causal attention participant attributes disagree".to_owned());
        }
        #[cfg(feature = "vllm-marlin")]
        if CausalProjection::from_values(participant.bindings(), projection_runtime)?.replay_tag()
            != projection.replay_tag()
        {
            return Err(
                "CUDA causal attention participants use different projection ABIs".to_owned(),
            );
        }
        validate_signature(participant, shape, semantics)?;
    }
    let program_binding = invocation.program_binding().cloned();

    let total_tokens = invocation.work_shape().immediate_tokens();
    let layout = ScratchLayout::new(shape, total_tokens, projection, attention_policy)?;
    let binding_layout = BindingLayout::new(shape, invocation.participants().len())?;
    let cuda = shape.cuda_shape()?;
    let token_ranges = invocation.participant_token_ranges();
    if token_ranges.len() != invocation.participants().len() {
        return Err("CUDA causal attention participant ranges are incomplete".to_owned());
    }
    let input_packed = super::token_binding_is_packed(&invocation, ResolvedValueRole::Input, 0)?;
    let output_packed = super::token_binding_is_packed(&invocation, ResolvedValueRole::Output, 0)?;

    let mut compute_regions = Vec::new();
    let shared = SharedRegions {
        input_norm: push_shared_weight(&mut compute_regions, &invocation, 1)?,
        query_weight: push_shared_projection_weight(
            &mut compute_regions,
            &invocation,
            2,
            &[shape.query_projection_features, shape.hidden_size],
        )?,
        key_weight: push_shared_projection_weight(
            &mut compute_regions,
            &invocation,
            3,
            &[shape.kv_features, shape.hidden_size],
        )?,
        value_weight: push_shared_projection_weight(
            &mut compute_regions,
            &invocation,
            4,
            &[shape.kv_features, shape.hidden_size],
        )?,
        output_weight: push_shared_projection_weight(
            &mut compute_regions,
            &invocation,
            5,
            &[shape.hidden_size, shape.query_features],
        )?,
        query_norm: push_shared_weight(&mut compute_regions, &invocation, 6)?,
        key_norm: push_shared_weight(&mut compute_regions, &invocation, 7)?,
        post_attention_norm: semantics
            .has_post_attention_norm()
            .then(|| push_shared_weight(&mut compute_regions, &invocation, 9))
            .transpose()?,
        scratch: {
            let index = compute_regions.len();
            compute_regions.push(super::shared_scratch_region(
                &invocation,
                layout.required_bytes,
            )?);
            index
        },
        binding: {
            let index = compute_regions.len();
            compute_regions.push(super::shared_binding_region(
                &invocation,
                binding_layout.required_bytes,
            )?);
            index
        },
    };

    let packed = if input_packed && output_packed && invocation.participants().len() > 1 {
        let input_region = compute_regions.len();
        compute_regions.push(super::shared_token_region(
            &invocation,
            ResolvedValueRole::Input,
            0,
            ElementType::F16,
            total_tokens,
        )?);
        let output_region = compute_regions.len();
        compute_regions.push(super::shared_token_region(
            &invocation,
            ResolvedValueRole::Output,
            0,
            ElementType::F16,
            total_tokens,
        )?);
        Some(PackedCausalAttentionLaunch {
            input_region,
            output_region,
            tokens: total_tokens,
            tokens_i32: checked_i32(total_tokens, "packed causal attention token count")?,
        })
    } else {
        None
    };

    let mut binding_regions = vec![compute_regions[shared.binding].clone()];
    let mut compute_fence_dependencies = Vec::new();
    let mut host_storage = Vec::with_capacity(invocation.participants().len());
    let mut launches = Vec::with_capacity(invocation.participants().len());
    let mut bindings = Vec::with_capacity(invocation.participants().len());
    for (participant_index, (participant, token_range)) in invocation
        .participants()
        .iter()
        .zip(token_ranges)
        .enumerate()
    {
        let tokens = token_range.immediate_tokens();
        let source = token_range.source_token_range();
        let packed_range = token_range.immediate_token_range();
        if source.end > token_range.full_input_tokens()
            || token_range.full_input_tokens() > shape.maximum_context_tokens
        {
            return Err("causal attention token range exceeds its admitted context".to_owned());
        }
        let input_region = if let Some(packed) = packed {
            packed.input_region
        } else {
            let input_region = compute_regions.len();
            compute_regions.push(contiguous_token_region(
                participant,
                binding(participant.bindings(), ResolvedValueRole::Input, 0)?,
                ElementType::F16,
                if input_packed {
                    packed_range.start
                } else {
                    source.start
                },
                tokens,
            )?);
            input_region
        };
        let output_region = if let Some(packed) = packed {
            packed.output_region
        } else {
            let output_region = compute_regions.len();
            compute_regions.push(contiguous_token_region(
                participant,
                binding(participant.bindings(), ResolvedValueRole::Output, 0)?,
                ElementType::F16,
                if output_packed {
                    packed_range.start
                } else {
                    source.start
                },
                tokens,
            )?);
            output_region
        };

        let first_page_region = binding_regions.len();
        let state = binding(participant.bindings(), ResolvedValueRole::Input, 8)?;
        let pages = paged_state_regions(
            participant,
            state,
            shape.physical_state_bytes_for_source_frontier(
                source.end,
                token_range.full_input_tokens(),
            )?,
        )?;
        let page_count = u64::try_from(pages.len())
            .map_err(|_| "causal attention page count exceeds u64".to_owned())?;
        if page_count > shape.maximum_pages()? {
            return Err("causal attention page table exceeds its admitted maximum".to_owned());
        }
        let table_entries = shape.table_entries(source.end)?;
        if table_entries > shape.table_entries(shape.maximum_context_tokens)? {
            return Err("causal attention address table exceeds its admitted maximum".to_owned());
        }
        let tokens_i32 = checked_i32(tokens, "causal attention participant token count")?;
        let position_start = checked_i32(source.start, "causal attention source position")?;
        let sequence_tokens_i32 = checked_i32(source.end, "causal attention sequence token count")?;
        let table_entries_i32 =
            checked_i32(table_entries, "causal attention address-table entry count")?;
        let path = CausalAttentionKernelPath::select(attention_policy, shape, tokens, source.end)?;
        let replay_topology = CausalAttentionReplayTopology::new(shape, path, source.end)?;
        let host_binding = host_storage.len();
        host_storage.push(binding_payload(
            shape.kv_layout()?,
            table_entries_i32,
            position_start,
            tokens_i32,
            sequence_tokens_i32,
            checked_i32(packed_range.start, "causal attention packed token start")?,
            &pages,
        )?);
        let page_count = pages.len();
        compute_fence_dependencies.extend(pages.iter().cloned());
        binding_regions.extend(pages);
        let binding_offset = binding_layout.binding_offset(participant_index)?;
        bindings.push(CausalAttentionBinding {
            first_page_region,
            page_count,
            host_binding,
            binding_offset,
        });
        launches.push(CausalAttentionLaunch {
            input_region,
            output_region,
            binding_offset,
            packed_token_start: packed_range.start,
            packed_query_raw: layout.token_offset(
                layout.query_raw,
                packed_range.start,
                shape.query_projection_features,
            )?,
            packed_key_raw: layout.token_offset(
                layout.key_raw,
                packed_range.start,
                shape.kv_features,
            )?,
            packed_value_raw: layout.token_offset(
                layout.value_raw,
                packed_range.start,
                shape.kv_features,
            )?,
            packed_query: layout.token_offset(
                layout.query,
                packed_range.start,
                shape.query_features,
            )?,
            packed_context: layout.token_offset(
                layout.context,
                packed_range.start,
                shape.query_features,
            )?,
            tokens,
            tokens_i32,
            sequence_tokens: source.end,
            sequence_tokens_i32,
            table_entries_i32,
            replay_topology,
            path,
        });
    }
    if packed.is_some() {
        validate_packed_token_ranges(
            launches
                .iter()
                .map(|launch| (launch.packed_token_start, launch.tokens)),
            total_tokens,
        )?;
    }

    let participant_count = u32::try_from(invocation.participants().len())
        .map_err(|_| "CUDA causal attention participant count exceeds u32".to_owned())?;
    let (binding_command, has_compiled_program_slot) =
        if let Some(program_binding) = program_binding {
            let mut regions = binding_regions.into_iter();
            let destination = regions
                .next()
                .ok_or_else(|| "CUDA causal binding destination is missing".to_owned())?;
            let fence_dependencies = regions.collect::<Vec<_>>();
            let mut writes = Vec::with_capacity(bindings.len());
            for (index, (binding, payload)) in bindings.into_iter().zip(host_storage).enumerate() {
                if binding.host_binding != index {
                    return Err("CUDA causal binding payload order is not canonical".to_owned());
                }
                writes.push(
                    super::CudaProgramBindingWrite::new(binding.binding_offset, payload)
                        .map_err(|error| error.to_string())?,
                );
            }
            (
                CudaDeviceCommand::program_binding_patch(
                    "vnext_causal_paged_attention_bindings",
                    program_binding,
                    destination,
                    writes,
                    fence_dependencies,
                ),
                true,
            )
        } else {
            (
                CudaDeviceCommand::operation_with_host_storage_and_blas(
                    "vnext_causal_paged_attention_bindings",
                    binding_regions,
                    host_storage,
                    move |stream, _blas, regions, host_storage| {
                        enqueue_bindings(stream, binding_layout, &bindings, regions, host_storage)
                    },
                ),
                false,
            )
        };
    let binding_command = binding_command
        .and_then(|command| {
            command.with_work_attribution(
                DeviceBatchingForm::ParticipantLoop,
                participant_count,
                total_tokens,
                0,
                u64::from(participant_count),
            )
        })
        .map_err(|error| error.to_string())?;

    let compute_operation = launches
        .first()
        .map(|launch| launch.path)
        .filter(|path| launches.iter().all(|launch| launch.path == *path))
        .map(CausalAttentionKernelPath::operation)
        .unwrap_or(COMPUTE_MIXED_OPERATION);
    let packed_enabled = packed.is_some();
    let compute_dispatch_count = physical_dispatch_count(
        launches.iter().map(|launch| launch.path),
        shape.output_gate,
        shape.post_attention_norm,
        packed_enabled,
    );
    let replay_key = launches
        .iter()
        .all(|launch| launch.replay_topology.is_partition_stable())
        .then(|| {
            let replay_key =
                CudaCommandReplayKeyBuilder::new(provider_fingerprint, compute_operation)
                    .bytes(projection.replay_tag().as_bytes());
            let replay_key = shared.query_weight.bind_replay_abi(replay_key);
            let replay_key = shared.key_weight.bind_replay_abi(replay_key);
            let replay_key = shared.value_weight.bind_replay_abi(replay_key);
            let mut replay_key = shared
                .output_weight
                .bind_replay_abi(replay_key)
                .u64(shape.hidden_size)
                .u64(shape.query_heads)
                .u64(shape.key_value_heads)
                .u64(shape.head_dim)
                .u64(shape.query_features)
                .u64(shape.query_projection_features)
                .u64(shape.kv_features)
                .u64(shape.rope_dim)
                .u64(shape.rope_frequency_denominator)
                .u64(shape.maximum_context_tokens)
                .f32(shape.epsilon)
                .f32(shape.rope_theta)
                .f32(shape.attention_scale)
                .u64(shape.sliding_window_tokens)
                .boolean(shape.rope_interleaved)
                .boolean(shape.output_gate)
                .boolean(shape.value_rms_norm)
                .boolean(shape.attention_k_eq_v)
                .boolean(shape.post_attention_norm)
                .u64(total_tokens)
                .boolean(packed_enabled)
                .u64(
                    packed
                        .map(|launch| launch.input_region as u64)
                        .unwrap_or(u64::MAX),
                )
                .u64(
                    packed
                        .map(|launch| launch.output_region as u64)
                        .unwrap_or(u64::MAX),
                )
                .u64(layout.required_bytes)
                .u64(layout.projection_workspace.unwrap_or(u64::MAX))
                .u64(layout.normalized)
                .u64(layout.query_raw)
                .u64(layout.key_raw)
                .u64(layout.value_raw)
                .u64(layout.query)
                .u64(layout.context)
                .u64(layout.projected)
                .u64(layout.vllm.map_or(0, |vllm| vllm.exp_sums))
                .u64(layout.vllm.map_or(0, |vllm| vllm.max_logits))
                .u64(layout.vllm.map_or(0, |vllm| vllm.temporary_output))
                .u64(binding_layout.required_bytes)
                .u64(binding_layout.slot_bytes)
                .u64(launches.len() as u64);
            for launch in &launches {
                let replay_envelope = launch.replay_topology.envelope();
                replay_key = replay_key
                    .u64(launch.input_region as u64)
                    .u64(launch.output_region as u64)
                    .u64(launch.binding_offset)
                    .u64(launch.packed_token_start)
                    .u64(launch.packed_query_raw)
                    .u64(launch.packed_key_raw)
                    .u64(launch.packed_value_raw)
                    .u64(launch.packed_query)
                    .u64(launch.packed_context)
                    .u64(launch.tokens)
                    .i32(launch.tokens_i32)
                    .u64(replay_envelope.sequence_capacity_tokens)
                    .i32(replay_envelope.table_capacity_entries)
                    .u64(launch.path.replay_id());
            }
            replay_key.finish()
        });

    let functions = functions.clone();
    let enqueue_compute =
        move |stream: &CudaStream, blas: &CudaBlas, regions: &[CudaBufferRegion]| {
            if let Some(packed) = packed {
                enqueue_packed_attention(
                    stream,
                    blas,
                    &functions,
                    projection,
                    shape,
                    cuda,
                    layout,
                    binding_layout,
                    shared,
                    packed,
                    &launches,
                    regions,
                )?;
            } else {
                for launch in &launches {
                    enqueue_attention(
                        stream, blas, &functions, projection, shape, cuda, layout, shared, *launch,
                        regions,
                    )?;
                }
            }
            Ok(())
        };
    let compute_command = match replay_key {
        Some(replay_key) => {
            CudaDeviceCommand::replayable_operation_with_blas_and_fence_dependencies(
                compute_operation,
                compute_regions,
                compute_fence_dependencies,
                replay_key,
                enqueue_compute,
            )
        }
        None => CudaDeviceCommand::operation_with_blas_and_fence_dependencies(
            compute_operation,
            compute_regions,
            compute_fence_dependencies,
            enqueue_compute,
        ),
    }
    .and_then(|command| {
        command.with_work_attribution(
            if packed_enabled {
                DeviceBatchingForm::Packed
            } else if participant_count == 1 {
                DeviceBatchingForm::Scalar
            } else {
                DeviceBatchingForm::ParticipantLoop
            },
            participant_count,
            total_tokens,
            compute_dispatch_count,
            0,
        )
    })
    .map_err(|error| error.to_string())?;

    Ok(attach_invocation_binding(
        EncodedDeviceOperation::compute(compute_command),
        binding_command,
        has_compiled_program_slot,
    ))
}

fn physical_dispatch_count(
    paths: impl IntoIterator<Item = CausalAttentionKernelPath>,
    output_gate: bool,
    post_attention_norm: bool,
    packed: bool,
) -> u64 {
    let paths = paths.into_iter().collect::<Vec<_>>();
    let packed_fallback = packed
        && paths.len() > 1
        && paths.first().is_some_and(|path| path.is_fallback())
        && paths.iter().all(|path| Some(path) == paths.first());
    if packed_fallback {
        return (6 + u64::from(post_attention_norm))
            .saturating_add(2)
            .saturating_add(if output_gate { paths.len() as u64 } else { 0 });
    }
    paths.into_iter().fold(
        if packed {
            6 + u64::from(post_attention_norm)
        } else {
            0
        },
        |total, path| {
            total.saturating_add(
                (if packed { 1 } else { 7 })
                    + path.attention_dispatch_count()
                    + u64::from(output_gate)
                    + u64::from(post_attention_norm && !packed),
            )
        },
    )
}

fn validate_packed_token_ranges(
    ranges: impl IntoIterator<Item = (u64, u64)>,
    total_tokens: u64,
) -> Result<(), String> {
    let mut next_token = 0_u64;
    for (token_start, tokens) in ranges {
        if tokens == 0 || token_start != next_token {
            return Err(
                "packed causal attention token ranges are not one canonical dense span".to_owned(),
            );
        }
        next_token = next_token
            .checked_add(tokens)
            .ok_or_else(|| "packed causal attention token range overflows".to_owned())?;
    }
    if next_token != total_tokens {
        return Err("packed causal attention token ranges do not cover the work shape".to_owned());
    }
    Ok(())
}

fn encode_reusable_attention_bindings(
    invocation: BatchedOperationInvocation<'_, CudaDeviceBuffer>,
    semantics: CausalAttentionSemantics,
    operation_id: &str,
) -> Result<EncodedReusableExecutionBindings<CudaDeviceCommand>, String> {
    if invocation.participants().is_empty() || invocation.operation().id.as_str() != operation_id {
        return Err("CUDA causal attention received another or empty operation".to_owned());
    }
    let program_binding = invocation.program_binding().cloned().ok_or_else(|| {
        "CUDA causal direct execution requires a compiled program binding".to_owned()
    })?;
    let shape = CausalAttentionShape::from_attributes_for(
        invocation.participants()[0].attributes(),
        semantics,
    )?;
    let total_tokens = invocation.work_shape().immediate_tokens();
    let participant_count = u32::try_from(invocation.participants().len())
        .map_err(|_| "CUDA causal attention participant count exceeds u32".to_owned())?;
    let binding_layout = BindingLayout::new(shape, invocation.participants().len())?;
    let token_ranges = invocation.participant_token_ranges();
    if token_ranges.len() != invocation.participants().len() {
        return Err("CUDA causal attention participant ranges are incomplete".to_owned());
    }
    let destination = super::shared_binding_region(&invocation, binding_layout.required_bytes)?;
    let kv_layout = shape.kv_layout()?;
    let maximum_pages = shape.maximum_pages()?;
    let maximum_table_entries = shape.table_entries(shape.maximum_context_tokens)?;
    let mut writes = Vec::with_capacity(invocation.participants().len());
    let mut fence_dependencies = Vec::new();

    // Static weights, scratch, kernel selection, and replay topology are
    // sealed into the exact program identity. Only the current KV page table
    // and sequence frontier belong on this hot path.
    for (participant_index, (participant, token_range)) in invocation
        .participants()
        .iter()
        .zip(token_ranges)
        .enumerate()
    {
        let tokens = token_range.immediate_tokens();
        let source = token_range.source_token_range();
        let packed_range = token_range.immediate_token_range();
        if source.end > token_range.full_input_tokens()
            || token_range.full_input_tokens() > shape.maximum_context_tokens
        {
            return Err("causal attention token range exceeds its admitted context".to_owned());
        }
        let state = binding(participant.bindings(), ResolvedValueRole::Input, 8)?;
        let pages = paged_state_regions(
            participant,
            state,
            shape.physical_state_bytes_for_source_frontier(
                source.end,
                token_range.full_input_tokens(),
            )?,
        )?;
        let page_count = u64::try_from(pages.len())
            .map_err(|_| "causal attention page count exceeds u64".to_owned())?;
        if page_count > maximum_pages {
            return Err("causal attention page table exceeds its admitted maximum".to_owned());
        }
        let table_entries = shape.table_entries(source.end)?;
        if table_entries > maximum_table_entries {
            return Err("causal attention address table exceeds its admitted maximum".to_owned());
        }
        let payload = binding_payload(
            kv_layout,
            checked_i32(table_entries, "causal attention address-table entry count")?,
            checked_i32(source.start, "causal attention source position")?,
            checked_i32(tokens, "causal attention participant token count")?,
            checked_i32(source.end, "causal attention sequence token count")?,
            checked_i32(packed_range.start, "causal attention packed token start")?,
            &pages,
        )?;
        writes.push(
            super::CudaProgramBindingWrite::new(
                binding_layout.binding_offset(participant_index)?,
                payload,
            )
            .map_err(|error| error.to_string())?,
        );
        fence_dependencies.extend(pages);
    }

    let binding_command = CudaDeviceCommand::program_binding_patch(
        "vnext_causal_paged_attention_bindings",
        program_binding,
        destination,
        writes,
        fence_dependencies,
    )
    .and_then(|command| {
        command.with_work_attribution(
            DeviceBatchingForm::ParticipantLoop,
            participant_count,
            total_tokens,
            0,
            u64::from(participant_count),
        )
    })
    .map_err(|error| error.to_string())?;
    Ok(EncodedReusableExecutionBindings::empty().with_program_binding(binding_command))
}

fn enqueue_bindings(
    stream: &CudaStream,
    layout: BindingLayout,
    bindings: &[CausalAttentionBinding],
    regions: &[CudaBufferRegion],
    host_storage: &[Box<[u8]>],
) -> Result<(), CudaDeviceRuntimeError> {
    let binding_workspace = &regions[0];
    if binding_workspace.length_bytes() < layout.required_bytes {
        return Err(CudaDeviceRuntimeError::contract(
            "causal attention binding workspace is smaller than its admitted estimate",
        ));
    }
    for binding in bindings {
        let page_region_end = binding
            .first_page_region
            .checked_add(binding.page_count)
            .ok_or_else(|| {
                CudaDeviceRuntimeError::contract("causal attention page region range overflows")
            })?;
        if regions
            .get(binding.first_page_region..page_region_end)
            .is_none_or(|pages| {
                pages.iter().any(|page| {
                    page.length_bytes() != VNEXT_KV_PAGE_BYTES
                        || page.element_type() != ElementType::F16
                })
            })
        {
            return Err(CudaDeviceRuntimeError::contract(
                "causal attention page regions changed after encoding",
            ));
        }
        let payload = host_storage.get(binding.host_binding).ok_or_else(|| {
            CudaDeviceRuntimeError::contract("causal attention binding payload is missing")
        })?;
        if payload.len() as u64 > layout.slot_bytes {
            return Err(CudaDeviceRuntimeError::contract(
                "causal attention binding payload exceeds its admitted slot",
            ));
        }
        let destination = scratch_pointer(binding_workspace.device_ptr(), binding.binding_offset)?;
        unsafe {
            cudarc::driver::result::memcpy_htod_async(
                destination,
                payload.as_ref(),
                stream.cu_stream(),
            )
        }
        .map_err(|error| {
            CudaDeviceRuntimeError::driver("causal attention binding upload", error)
        })?;
    }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn enqueue_packed_attention(
    stream: &CudaStream,
    blas: &CudaBlas,
    functions: &CausalAttentionFunctions,
    projection: CausalProjection,
    logical: CausalAttentionShape,
    cuda: CudaCausalAttentionShape,
    layout: ScratchLayout,
    binding_layout: BindingLayout,
    shared: SharedRegions,
    packed: PackedCausalAttentionLaunch,
    launches: &[CausalAttentionLaunch],
    regions: &[CudaBufferRegion],
) -> Result<(), CudaDeviceRuntimeError> {
    let scratch = &regions[shared.scratch];
    if scratch.length_bytes() < layout.required_bytes {
        return Err(CudaDeviceRuntimeError::contract(
            "packed causal attention scratch is smaller than its admitted estimate",
        ));
    }
    let scratch_base = scratch.device_ptr();
    let binding = &regions[shared.binding];
    let input = regions[packed.input_region].device_ptr();
    let output = regions[packed.output_region].device_ptr();
    let normalized = scratch_pointer(scratch_base, layout.normalized)?;
    let query_raw = scratch_pointer(scratch_base, layout.query_raw)?;
    let key_raw = scratch_pointer(scratch_base, layout.key_raw)?;
    let value_raw = scratch_pointer(scratch_base, layout.value_raw)?;
    let context = scratch_pointer(scratch_base, layout.context)?;
    let projected = scratch_pointer(scratch_base, layout.projected)?;

    launch_rms_norm(
        stream,
        &functions.rms_norm,
        input,
        regions[shared.input_norm].device_ptr(),
        normalized,
        packed.tokens,
        cuda.hidden_size,
        cuda.epsilon,
    )?;
    for (weight, destination, out_features, operation) in [
        (
            shared.query_weight,
            query_raw,
            cuda.query_projection_features,
            "packed causal attention Q GEMM",
        ),
        (
            shared.key_weight,
            key_raw,
            cuda.kv_features,
            "packed causal attention K GEMM",
        ),
        (
            shared.value_weight,
            value_raw,
            cuda.kv_features,
            "packed causal attention V GEMM",
        ),
    ] {
        launch_causal_projection(
            stream,
            blas,
            projection,
            weight,
            normalized,
            destination,
            scratch,
            layout,
            regions,
            packed.tokens_i32,
            out_features,
            cuda.hidden_size,
            operation,
        )?;
    }

    if let Some(packed_fallback) = packed_fallback_launch(launches, binding_layout)
        .map_err(CudaDeviceRuntimeError::contract)?
    {
        let launch = launches[0];
        let control = binding.device_ptr();
        let page_table = control.checked_add(BINDING_CONTROL_BYTES).ok_or_else(|| {
            CudaDeviceRuntimeError::contract("packed causal attention page-table pointer overflows")
        })?;
        launch_prepare(
            stream,
            &functions.prepare,
            query_raw,
            key_raw,
            value_raw,
            regions[shared.query_norm].device_ptr(),
            regions[shared.key_norm].device_ptr(),
            scratch_pointer(scratch_base, layout.query)?,
            control,
            page_table,
            launch,
            cuda,
            i32::from(packed_fallback.path.uses_vllm_layout()),
            Some(packed_fallback),
        )?;
        launch_fallback_attention(
            stream,
            functions,
            scratch_pointer(scratch_base, layout.query)?,
            query_raw,
            control,
            page_table,
            context,
            launch,
            cuda,
            i32::from(packed_fallback.path.uses_vllm_layout()),
            Some(packed_fallback),
        )?;
        if logical.output_gate {
            for launch in launches {
                launch_attention_gate(
                    stream,
                    &functions.attention_gate,
                    scratch_pointer(scratch_base, launch.packed_context)?,
                    scratch_pointer(scratch_base, launch.packed_query_raw)?,
                    *launch,
                    cuda,
                )?;
            }
        }
    } else {
        for launch in launches {
            let control = scratch_pointer(binding.device_ptr(), launch.binding_offset)?;
            let page_table = control.checked_add(BINDING_CONTROL_BYTES).ok_or_else(|| {
                CudaDeviceRuntimeError::contract(
                    "packed causal attention page-table pointer overflows",
                )
            })?;
            let participant_query_raw = scratch_pointer(scratch_base, launch.packed_query_raw)?;
            let participant_key_raw = scratch_pointer(scratch_base, launch.packed_key_raw)?;
            let participant_value_raw = scratch_pointer(scratch_base, launch.packed_value_raw)?;
            let participant_query = scratch_pointer(scratch_base, launch.packed_query)?;
            let participant_context = scratch_pointer(scratch_base, launch.packed_context)?;

            launch_prepare(
                stream,
                &functions.prepare,
                participant_query_raw,
                participant_key_raw,
                participant_value_raw,
                regions[shared.query_norm].device_ptr(),
                regions[shared.key_norm].device_ptr(),
                participant_query,
                control,
                page_table,
                *launch,
                cuda,
                i32::from(launch.path.uses_vllm_layout()),
                None,
            )?;
            launch_selected_attention(
                stream,
                functions,
                participant_query,
                participant_query_raw,
                control,
                page_table,
                participant_context,
                *launch,
                cuda,
                layout,
                scratch_base,
            )?;
            if logical.output_gate {
                launch_attention_gate(
                    stream,
                    &functions.attention_gate,
                    participant_context,
                    participant_query_raw,
                    *launch,
                    cuda,
                )?;
            }
        }
    }

    launch_causal_projection(
        stream,
        blas,
        projection,
        shared.output_weight,
        context,
        projected,
        scratch,
        layout,
        regions,
        packed.tokens_i32,
        cuda.hidden_size,
        cuda.query_features,
        "packed causal attention output GEMM",
    )?;
    let residual_branch = if logical.post_attention_norm {
        let norm_region = shared.post_attention_norm.ok_or_else(|| {
            CudaDeviceRuntimeError::contract(
                "packed Gemma4 causal attention lacks post-attention RMSNorm weight",
            )
        })?;
        launch_rms_norm(
            stream,
            &functions.rms_norm,
            projected,
            regions[norm_region].device_ptr(),
            normalized,
            packed.tokens,
            cuda.hidden_size,
            cuda.epsilon,
        )?;
        normalized
    } else {
        projected
    };
    let elements = packed
        .tokens
        .checked_mul(logical.hidden_size)
        .ok_or_else(|| CudaDeviceRuntimeError::contract("packed causal residual size overflows"))?;
    launch_residual(
        stream,
        &functions.residual_add,
        &functions.residual_add_inplace,
        input,
        residual_branch,
        output,
        elements,
    )
}

#[allow(clippy::too_many_arguments)]
fn enqueue_attention(
    stream: &CudaStream,
    blas: &CudaBlas,
    functions: &CausalAttentionFunctions,
    projection: CausalProjection,
    logical: CausalAttentionShape,
    cuda: CudaCausalAttentionShape,
    layout: ScratchLayout,
    shared: SharedRegions,
    launch: CausalAttentionLaunch,
    regions: &[CudaBufferRegion],
) -> Result<(), CudaDeviceRuntimeError> {
    let scratch = &regions[shared.scratch];
    if scratch.length_bytes() < layout.required_bytes {
        return Err(CudaDeviceRuntimeError::contract(
            "causal attention scratch is smaller than its admitted estimate",
        ));
    }
    let scratch_base = scratch.device_ptr();
    let binding = &regions[shared.binding];
    let control = scratch_pointer(binding.device_ptr(), launch.binding_offset)?;
    let page_table = control.checked_add(BINDING_CONTROL_BYTES).ok_or_else(|| {
        CudaDeviceRuntimeError::contract("causal attention page-table pointer overflows")
    })?;

    let input = regions[launch.input_region].device_ptr();
    let output = regions[launch.output_region].device_ptr();
    let normalized = scratch_pointer(scratch_base, layout.normalized)?;
    let query_raw = scratch_pointer(scratch_base, layout.query_raw)?;
    let key_raw = scratch_pointer(scratch_base, layout.key_raw)?;
    let value_raw = scratch_pointer(scratch_base, layout.value_raw)?;
    let query = scratch_pointer(scratch_base, layout.query)?;
    let context = scratch_pointer(scratch_base, layout.context)?;
    let projected = scratch_pointer(scratch_base, layout.projected)?;

    launch_rms_norm(
        stream,
        &functions.rms_norm,
        input,
        regions[shared.input_norm].device_ptr(),
        normalized,
        launch.tokens,
        cuda.hidden_size,
        cuda.epsilon,
    )?;
    for (weight, destination, out_features, operation) in [
        (
            shared.query_weight,
            query_raw,
            cuda.query_projection_features,
            "causal attention Q GEMM",
        ),
        (
            shared.key_weight,
            key_raw,
            cuda.kv_features,
            "causal attention K GEMM",
        ),
        (
            shared.value_weight,
            value_raw,
            cuda.kv_features,
            "causal attention V GEMM",
        ),
    ] {
        launch_causal_projection(
            stream,
            blas,
            projection,
            weight,
            normalized,
            destination,
            scratch,
            layout,
            regions,
            launch.tokens_i32,
            out_features,
            cuda.hidden_size,
            operation,
        )?;
    }
    launch_prepare(
        stream,
        &functions.prepare,
        query_raw,
        key_raw,
        value_raw,
        regions[shared.query_norm].device_ptr(),
        regions[shared.key_norm].device_ptr(),
        query,
        control,
        page_table,
        launch,
        cuda,
        i32::from(launch.path.uses_vllm_layout()),
        None,
    )?;
    launch_selected_attention(
        stream,
        functions,
        query,
        query_raw,
        control,
        page_table,
        context,
        launch,
        cuda,
        layout,
        scratch_base,
    )?;
    if logical.output_gate {
        launch_attention_gate(
            stream,
            &functions.attention_gate,
            context,
            query_raw,
            launch,
            cuda,
        )?;
    }
    launch_causal_projection(
        stream,
        blas,
        projection,
        shared.output_weight,
        context,
        projected,
        scratch,
        layout,
        regions,
        launch.tokens_i32,
        cuda.hidden_size,
        cuda.query_features,
        "causal attention output GEMM",
    )?;
    let residual_branch = if logical.post_attention_norm {
        let norm_region = shared.post_attention_norm.ok_or_else(|| {
            CudaDeviceRuntimeError::contract(
                "Gemma4 causal attention lacks post-attention RMSNorm weight",
            )
        })?;
        launch_rms_norm(
            stream,
            &functions.rms_norm,
            projected,
            regions[norm_region].device_ptr(),
            normalized,
            launch.tokens,
            cuda.hidden_size,
            cuda.epsilon,
        )?;
        normalized
    } else {
        projected
    };
    let elements = launch
        .tokens
        .checked_mul(logical.hidden_size)
        .ok_or_else(|| CudaDeviceRuntimeError::contract("causal residual size overflows"))?;
    launch_residual(
        stream,
        &functions.residual_add,
        &functions.residual_add_inplace,
        input,
        residual_branch,
        output,
        elements,
    )
}

#[allow(clippy::too_many_arguments)]
fn launch_causal_projection(
    stream: &CudaStream,
    blas: &CudaBlas,
    projection: CausalProjection,
    weight: SharedProjectionWeight,
    input: u64,
    output: u64,
    scratch: &CudaBufferRegion,
    layout: ScratchLayout,
    regions: &[CudaBufferRegion],
    rows: i32,
    output_features: i32,
    input_features: i32,
    operation: &'static str,
) -> Result<(), CudaDeviceRuntimeError> {
    match weight {
        SharedProjectionWeight::F16 { region } => launch_gemm_f16(
            blas,
            input,
            regions[region].device_ptr(),
            output,
            rows,
            output_features,
            input_features,
            operation,
        ),
        #[cfg(feature = "vllm-marlin")]
        SharedProjectionWeight::Marlin {
            packed_region,
            scales_region,
            zero_points_region,
            group_size,
            weight_type,
        } => {
            let runtime = match (projection, weight_type) {
                (CausalProjection::MarlinFp8 { runtime }, MarlinF16WeightType::E4M3Fn)
                | (CausalProjection::MarlinInt4 { runtime }, MarlinF16WeightType::U4)
                | (CausalProjection::MarlinInt4 { runtime }, MarlinF16WeightType::U4B8) => runtime,
                _ => {
                    return Err(CudaDeviceRuntimeError::contract(format!(
                        "{operation} Marlin weight type differs from its admitted projection ABI"
                    )))
                }
            };
            let workspace_offset = layout.projection_workspace.ok_or_else(|| {
                CudaDeviceRuntimeError::contract(format!(
                    "{operation} lacks its admitted Marlin lock workspace"
                ))
            })?;
            let workspace_bytes = runtime
                .workspace_bytes()
                .map_err(CudaDeviceRuntimeError::contract)?;
            let workspace_end = workspace_offset
                .checked_add(workspace_bytes)
                .ok_or_else(|| {
                    CudaDeviceRuntimeError::contract(format!(
                        "{operation} Marlin workspace range overflows"
                    ))
                })?;
            if workspace_end > scratch.length_bytes() {
                return Err(CudaDeviceRuntimeError::contract(format!(
                    "{operation} Marlin workspace exceeds causal-attention scratch"
                )));
            }
            runtime.launch(
                weight_type,
                stream,
                input,
                regions[packed_region].device_ptr(),
                regions[scales_region].device_ptr(),
                zero_points_region.map(|region| regions[region].device_ptr()),
                output,
                scratch_pointer(scratch.device_ptr(), workspace_offset)?,
                workspace_bytes,
                rows,
                output_features,
                input_features,
                group_size,
                operation,
            )
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn launch_prepare(
    stream: &CudaStream,
    function: &CudaFunction,
    query_raw: u64,
    key_raw: u64,
    value_raw: u64,
    query_norm: u64,
    key_norm: u64,
    query: u64,
    control: u64,
    page_table: u64,
    launch: CausalAttentionLaunch,
    shape: CudaCausalAttentionShape,
    kv_layout: i32,
    packed: Option<PackedFallbackLaunch>,
) -> Result<(), CudaDeviceRuntimeError> {
    let page_elements = checked_i32_runtime(
        VNEXT_KV_PAGE_BYTES / ElementType::F16.size_bytes(),
        "causal page elements",
    )?;
    let query_head_stride = shape
        .head_dim
        .checked_mul(if shape.output_gate != 0 { 2 } else { 1 })
        .ok_or_else(|| CudaDeviceRuntimeError::contract("causal query head stride overflows"))?;
    let combined_heads = shape
        .key_value_heads
        .checked_mul(2)
        .and_then(|heads| heads.checked_add(shape.query_heads))
        .ok_or_else(|| CudaDeviceRuntimeError::contract("causal prepare head count overflows"))?;
    let mut builder = stream.launch_builder(function);
    let pointers = [
        query_raw, key_raw, value_raw, query_norm, key_norm, query, control, page_table,
    ];
    for pointer in &pointers {
        builder.arg(pointer);
    }
    let dimensions = [
        page_elements,
        kv_layout,
        shape.query_heads,
        shape.key_value_heads,
        shape.head_dim,
        shape.rope_dim,
        shape.rope_frequency_denominator,
        shape.query_projection_features,
        query_head_stride,
        shape.kv_features,
    ];
    for dimension in &dimensions {
        builder.arg(dimension);
    }
    builder.arg(&shape.epsilon);
    builder.arg(&shape.rope_theta);
    builder.arg(&shape.rope_interleaved);
    builder.arg(&shape.value_rms_norm);
    let binding_slot_bytes = packed.map_or(0, |packed| packed.binding_slot_bytes);
    builder.arg(&binding_slot_bytes);
    let token_grid = packed.map_or(launch.tokens, |packed| packed.token_grid);
    let participant_grid = packed.map_or(1, |packed| packed.participant_grid);
    unsafe {
        builder.launch(LaunchConfig {
            grid_dim: (
                checked_u32_runtime(token_grid, "causal prepare token grid")?,
                u32::try_from(combined_heads).map_err(|_| {
                    CudaDeviceRuntimeError::contract("causal prepare head grid exceeds u32")
                })?,
                participant_grid,
            ),
            block_dim: (WARP_THREADS, 1, 1),
            shared_mem_bytes: 0,
        })
    }
    .map(|_| ())
    .map_err(|error| CudaDeviceRuntimeError::driver("causal attention prepare launch", error))
}

#[allow(clippy::too_many_arguments)]
fn launch_selected_attention(
    stream: &CudaStream,
    functions: &CausalAttentionFunctions,
    query: u64,
    query_raw: u64,
    control: u64,
    page_table: u64,
    output: u64,
    launch: CausalAttentionLaunch,
    shape: CudaCausalAttentionShape,
    layout: ScratchLayout,
    scratch_base: u64,
) -> Result<(), CudaDeviceRuntimeError> {
    match launch.path {
        CausalAttentionKernelPath::TokenMajorFallback => launch_fallback_attention(
            stream, functions, query, query_raw, control, page_table, output, launch, shape, 0,
            None,
        ),
        CausalAttentionKernelPath::VllmAddressedFallback => launch_fallback_attention(
            stream, functions, query, query_raw, control, page_table, output, launch, shape, 1,
            None,
        ),
        CausalAttentionKernelPath::VllmAddressedVarlen => launch_addressed_varlen_attention(
            stream,
            &functions.varlen_addressed,
            query,
            control,
            page_table,
            output,
            launch,
            shape,
            false,
        ),
        CausalAttentionKernelPath::VllmAddressedVarlenTiled => launch_addressed_varlen_attention(
            stream,
            &functions.varlen_tiled_addressed,
            query,
            control,
            page_table,
            output,
            launch,
            shape,
            true,
        ),
        CausalAttentionKernelPath::VllmAddressedDecodeV1
        | CausalAttentionKernelPath::VllmAddressedDecodeV2 => {
            #[cfg(feature = "vllm-paged-attn-v2")]
            {
                let scratch = layout.vllm.ok_or_else(|| {
                    CudaDeviceRuntimeError::contract(
                        "vLLM addressed decode has no caller-owned scratch layout",
                    )
                })?;
                let expected = match launch.path {
                    CausalAttentionKernelPath::VllmAddressedDecodeV1 => {
                        VnextAddressedPagedAttentionKernel::V1
                    }
                    CausalAttentionKernelPath::VllmAddressedDecodeV2 => {
                        VnextAddressedPagedAttentionKernel::V2
                    }
                    _ => unreachable!(),
                };
                let sequence_length_device = control
                    .checked_add(BINDING_SEQUENCE_LENGTH_OFFSET)
                    .ok_or_else(|| {
                        CudaDeviceRuntimeError::contract(
                            "causal attention sequence-length pointer overflows",
                        )
                    })?;
                let actual = unsafe {
                    dispatch_vnext_addressed_paged_attention_raw(
                        stream,
                        output,
                        query,
                        page_table,
                        sequence_length_device,
                        launch.replay_topology.envelope().sequence_capacity_tokens,
                        Some(scratch_pointer(scratch_base, scratch.exp_sums)?),
                        Some(scratch_pointer(scratch_base, scratch.max_logits)?),
                        Some(scratch_pointer(scratch_base, scratch.temporary_output)?),
                        shape.query_heads,
                        shape.key_value_heads,
                        shape.head_dim,
                        launch.replay_topology.envelope().table_capacity_entries,
                    )
                }
                .map_err(|error| CudaDeviceRuntimeError::contract(error.to_string()))?;
                if actual != expected || actual.native_kernel_id() != launch.path.native_kernel_id()
                {
                    return Err(CudaDeviceRuntimeError::contract(
                        "vLLM addressed decode selected a different native kernel",
                    ));
                }
                Ok(())
            }
            #[cfg(not(feature = "vllm-paged-attn-v2"))]
            {
                let _ = layout;
                Err(CudaDeviceRuntimeError::contract(
                    "vLLM addressed decode path was selected without its compiled feature",
                ))
            }
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn launch_addressed_varlen_attention(
    stream: &CudaStream,
    function: &CudaFunction,
    query: u64,
    control: u64,
    page_table: u64,
    output: u64,
    launch: CausalAttentionLaunch,
    shape: CudaCausalAttentionShape,
    tiled: bool,
) -> Result<(), CudaDeviceRuntimeError> {
    if shape.sliding_window_tokens != 0 {
        return Err(CudaDeviceRuntimeError::contract(
            "addressed varlen attention does not admit a sliding window",
        ));
    }
    let scale = shape.attention_scale;
    let score_rows = if tiled { VARLEN_TILED_QUERY_TOKENS } else { 1 };
    let shared_mem_bytes = launch
        .sequence_tokens
        .checked_mul(score_rows)
        .and_then(|values| values.checked_mul(std::mem::size_of::<f32>() as u64))
        .and_then(|bytes| u32::try_from(bytes).ok())
        .ok_or_else(|| CudaDeviceRuntimeError::contract("varlen shared memory size overflows"))?;
    if u64::from(shared_mem_bytes) > VARLEN_DYNAMIC_SHARED_BUDGET_BYTES {
        return Err(CudaDeviceRuntimeError::contract(
            "varlen shared memory exceeds the selected kernel path",
        ));
    }
    let grid_y = if tiled {
        launch.tokens.div_ceil(VARLEN_TILED_QUERY_TOKENS)
    } else {
        launch.tokens
    };
    let mut builder = stream.launch_builder(function);
    let pointers = [query, control, page_table, output];
    for pointer in &pointers {
        builder.arg(pointer);
    }
    let dimensions = [shape.query_heads, shape.key_value_heads, shape.head_dim];
    for dimension in &dimensions {
        builder.arg(dimension);
    }
    if tiled {
        builder.arg(&launch.sequence_tokens_i32);
    }
    builder.arg(&scale);
    unsafe {
        builder.launch(LaunchConfig {
            grid_dim: (
                u32::try_from(shape.query_heads).map_err(|_| {
                    CudaDeviceRuntimeError::contract("varlen query-head grid exceeds u32")
                })?,
                checked_u32_runtime(grid_y, "varlen query-token grid")?,
                1,
            ),
            block_dim: (128, 1, 1),
            shared_mem_bytes,
        })
    }
    .map(|_| ())
    .map_err(|error| {
        CudaDeviceRuntimeError::driver("causal attention addressed varlen launch", error)
    })
}

fn launch_attention_gate(
    stream: &CudaStream,
    function: &CudaFunction,
    context: u64,
    query_raw: u64,
    launch: CausalAttentionLaunch,
    shape: CudaCausalAttentionShape,
) -> Result<(), CudaDeviceRuntimeError> {
    let elements = launch
        .tokens
        .checked_mul(shape.query_features as u64)
        .ok_or_else(|| CudaDeviceRuntimeError::contract("attention gate size overflows"))?;
    let grid = checked_u32_runtime(
        elements.div_ceil(u64::from(THREADS_PER_BLOCK)),
        "attention gate grid",
    )?;
    let mut builder = stream.launch_builder(function);
    let pointers = [context, query_raw];
    for pointer in &pointers {
        builder.arg(pointer);
    }
    let dimensions = [
        launch.tokens_i32,
        shape.query_features,
        shape.query_projection_features,
        shape.head_dim,
    ];
    for dimension in &dimensions {
        builder.arg(dimension);
    }
    unsafe {
        builder.launch(LaunchConfig {
            grid_dim: (grid, 1, 1),
            block_dim: (THREADS_PER_BLOCK, 1, 1),
            shared_mem_bytes: 0,
        })
    }
    .map(|_| ())
    .map_err(|error| CudaDeviceRuntimeError::driver("causal attention gate launch", error))
}

#[allow(clippy::too_many_arguments)]
fn launch_fallback_attention(
    stream: &CudaStream,
    functions: &CausalAttentionFunctions,
    query: u64,
    query_raw: u64,
    control: u64,
    page_table: u64,
    output: u64,
    launch: CausalAttentionLaunch,
    shape: CudaCausalAttentionShape,
    kv_layout: i32,
    packed: Option<PackedFallbackLaunch>,
) -> Result<(), CudaDeviceRuntimeError> {
    // A scalar launch already exposes its Q heads as independent blocks and
    // has no cross-request dispatch overhead to amortize. Keep that latency
    // path unchanged; group Q heads only when multiple packed participants can
    // fill the wider KV-head grid and reuse each loaded K/V tile.
    if packed.is_some() {
        if let Some(block_threads) = grouped_fallback_block_threads(shape) {
            return launch_grouped_fallback_attention(
                stream,
                &functions.grouped_attention,
                query,
                query_raw,
                control,
                page_table,
                output,
                launch,
                shape,
                kv_layout,
                packed,
                block_threads,
            );
        }
    }

    let page_elements = checked_i32_runtime(
        VNEXT_KV_PAGE_BYTES / ElementType::F16.size_bytes(),
        "causal page elements",
    )?;
    let mut builder = stream.launch_builder(&functions.attention);
    let pointers = [query, query_raw, control, page_table, output];
    for pointer in &pointers {
        builder.arg(pointer);
    }
    let dimensions = [
        page_elements,
        kv_layout,
        shape.query_heads,
        shape.key_value_heads,
        shape.head_dim,
        shape.query_projection_features,
        0,
    ];
    for dimension in &dimensions {
        builder.arg(dimension);
    }
    builder.arg(&shape.attention_scale);
    builder.arg(&shape.sliding_window_tokens);
    let binding_slot_bytes = packed.map_or(0, |packed| packed.binding_slot_bytes);
    builder.arg(&binding_slot_bytes);
    let token_grid = packed.map_or(launch.tokens, |packed| packed.token_grid);
    let participant_grid = packed.map_or(1, |packed| packed.participant_grid);
    unsafe {
        builder.launch(LaunchConfig {
            grid_dim: (
                checked_u32_runtime(token_grid, "causal attention token grid")?,
                u32::try_from(shape.query_heads).map_err(|_| {
                    CudaDeviceRuntimeError::contract("causal attention head grid exceeds u32")
                })?,
                participant_grid,
            ),
            block_dim: (WARP_THREADS, 1, 1),
            shared_mem_bytes: 0,
        })
    }
    .map(|_| ())
    .map_err(|error| CudaDeviceRuntimeError::driver("causal attention launch", error))
}

fn grouped_fallback_block_threads(shape: CudaCausalAttentionShape) -> Option<u32> {
    if shape.query_heads <= 0
        || shape.key_value_heads <= 0
        || shape.query_heads % shape.key_value_heads != 0
    {
        return None;
    }
    let queries_per_kv = u32::try_from(shape.query_heads / shape.key_value_heads).ok()?;
    if queries_per_kv == 0 || queries_per_kv > MAXIMUM_GROUPED_QUERY_HEADS_PER_KV {
        return None;
    }
    queries_per_kv.checked_mul(WARP_THREADS)
}

#[allow(clippy::too_many_arguments)]
fn launch_grouped_fallback_attention(
    stream: &CudaStream,
    function: &CudaFunction,
    query: u64,
    query_raw: u64,
    control: u64,
    page_table: u64,
    output: u64,
    launch: CausalAttentionLaunch,
    shape: CudaCausalAttentionShape,
    kv_layout: i32,
    packed: Option<PackedFallbackLaunch>,
    block_threads: u32,
) -> Result<(), CudaDeviceRuntimeError> {
    let page_elements = checked_i32_runtime(
        VNEXT_KV_PAGE_BYTES / ElementType::F16.size_bytes(),
        "causal grouped-query page elements",
    )?;
    let mut builder = stream.launch_builder(function);
    let pointers = [query, query_raw, control, page_table, output];
    for pointer in &pointers {
        builder.arg(pointer);
    }
    let dimensions = [
        page_elements,
        kv_layout,
        shape.query_heads,
        shape.key_value_heads,
        shape.head_dim,
        shape.query_projection_features,
        0,
    ];
    for dimension in &dimensions {
        builder.arg(dimension);
    }
    builder.arg(&shape.attention_scale);
    builder.arg(&shape.sliding_window_tokens);
    let binding_slot_bytes = packed.map_or(0, |packed| packed.binding_slot_bytes);
    let participant_count = packed.map_or(1, |packed| packed.participant_count_i32);
    builder.arg(&binding_slot_bytes);
    builder.arg(&participant_count);

    let token_grid = packed.map_or(launch.tokens, |packed| packed.packed_token_grid);
    let head_dim = u64::try_from(shape.head_dim).map_err(|_| {
        CudaDeviceRuntimeError::contract("causal grouped-query head dimension is negative")
    })?;
    let kv_tile_tokens = if shape.head_dim == 256 && shape.sliding_window_tokens == 1_024 {
        GROUPED_FALLBACK_GEMMA_LOCAL_KV_TILE_TOKENS
    } else {
        GROUPED_FALLBACK_DEFAULT_KV_TILE_TOKENS
    };
    let shared_mem_bytes = kv_tile_tokens
        .checked_mul(head_dim)
        .and_then(|elements| elements.checked_mul(ElementType::F16.size_bytes()))
        .and_then(|bytes| bytes.checked_mul(2))
        .ok_or_else(|| {
            CudaDeviceRuntimeError::contract(
                "causal grouped-query attention shared memory overflows",
            )
        })?;
    unsafe {
        builder.launch(LaunchConfig {
            grid_dim: (
                checked_u32_runtime(token_grid, "causal grouped-query token grid")?,
                u32::try_from(shape.key_value_heads).map_err(|_| {
                    CudaDeviceRuntimeError::contract(
                        "causal grouped-query KV-head grid exceeds u32",
                    )
                })?,
                1,
            ),
            block_dim: (block_threads, 1, 1),
            shared_mem_bytes: checked_u32_runtime(
                shared_mem_bytes,
                "causal grouped-query shared memory",
            )?,
        })
    }
    .map(|_| ())
    .map_err(|error| CudaDeviceRuntimeError::driver("causal grouped-query attention launch", error))
}

fn launch_rms_norm(
    stream: &CudaStream,
    function: &CudaFunction,
    input: u64,
    weight: u64,
    output: u64,
    tokens: u64,
    hidden_size: i32,
    epsilon: f32,
) -> Result<(), CudaDeviceRuntimeError> {
    let rows = checked_u32_runtime(tokens, "causal RMSNorm rows")?;
    let mut builder = stream.launch_builder(function);
    builder.arg(&input);
    builder.arg(&weight);
    builder.arg(&output);
    builder.arg(&hidden_size);
    builder.arg(&epsilon);
    unsafe {
        builder.launch(LaunchConfig {
            grid_dim: (rows, 1, 1),
            block_dim: ((hidden_size as u32).min(1024), 1, 1),
            shared_mem_bytes: 0,
        })
    }
    .map(|_| ())
    .map_err(|error| CudaDeviceRuntimeError::driver("causal attention RMSNorm launch", error))
}

fn launch_residual(
    stream: &CudaStream,
    function: &CudaFunction,
    inplace_function: &CudaFunction,
    input: u64,
    branch: u64,
    output: u64,
    elements: u64,
) -> Result<(), CudaDeviceRuntimeError> {
    let elements_i32 = checked_i32_runtime(elements, "causal residual elements")?;
    let grid = checked_u32_runtime(
        elements.div_ceil(u64::from(THREADS_PER_BLOCK)),
        "causal residual grid",
    )?;
    let config = LaunchConfig {
        grid_dim: (grid, 1, 1),
        block_dim: (THREADS_PER_BLOCK, 1, 1),
        shared_mem_bytes: 0,
    };
    let result = if input == output {
        let mut builder = stream.launch_builder(inplace_function);
        builder.arg(&output);
        builder.arg(&branch);
        builder.arg(&elements_i32);
        unsafe { builder.launch(config) }
    } else {
        let mut builder = stream.launch_builder(function);
        builder.arg(&input);
        builder.arg(&branch);
        builder.arg(&output);
        builder.arg(&elements_i32);
        unsafe { builder.launch(config) }
    };
    result
        .map(|_| ())
        .map_err(|error| CudaDeviceRuntimeError::driver("causal attention residual launch", error))
}

fn paged_state_regions(
    participant: &OperationInvocation<'_, CudaDeviceBuffer>,
    state: &ResolvedValueBinding,
    expected_physical_bytes: u64,
) -> Result<Vec<CudaBufferRegion>, String> {
    let [component] = state.storage().components() else {
        return Err("causal attention state requires one logical storage component".to_owned());
    };
    let view = participant
        .views()
        .iter()
        .find(|view| view.resource_id() == component.resource_id())
        .ok_or_else(|| "causal attention state has no resource view".to_owned())?;
    if component.offset_bytes() != 0
        || component.element_type() != ElementType::F16
        || view.descriptor().element_type != ElementType::F16
        || view.storage_kind() != OperationBufferStorageKind::DynamicPaged
        || view.descriptor().size_bytes != expected_physical_bytes
        || expected_physical_bytes == 0
        || expected_physical_bytes % VNEXT_KV_PAGE_BYTES != 0
    {
        return Err("causal attention state is not its admitted fixed-block paged view".to_owned());
    }
    let translated = view
        .translate(0, expected_physical_bytes)
        .map_err(|error| error.to_string())?;
    let page_capacity = usize::try_from(expected_physical_bytes / VNEXT_KV_PAGE_BYTES)
        .map_err(|_| "causal attention page capacity exceeds usize".to_owned())?;
    let mut pages = Vec::with_capacity(page_capacity);
    let mut next_logical = 0_u64;
    for physical in translated.iter() {
        if physical.logical_offset_bytes() != next_logical
            || physical.length_bytes() == 0
            || physical.length_bytes() % VNEXT_KV_PAGE_BYTES != 0
        {
            return Err("causal attention paged translation lost block geometry".to_owned());
        }
        let (buffer, range, retention) = physical.buffer_and_physical_range();
        let mut offset = 0_u64;
        while offset < physical.length_bytes() {
            let start = range
                .start
                .checked_add(offset)
                .ok_or_else(|| "causal attention page offset overflows".to_owned())?;
            let end = start
                .checked_add(VNEXT_KV_PAGE_BYTES)
                .ok_or_else(|| "causal attention page range overflows".to_owned())?;
            let page = buffer
                .retained_region(start..end, retention.clone())
                .map_err(|error| error.to_string())?;
            if page.length_bytes() != VNEXT_KV_PAGE_BYTES || page.element_type() != ElementType::F16
            {
                return Err("causal attention physical page differs from its contract".to_owned());
            }
            pages.push(page);
            offset += VNEXT_KV_PAGE_BYTES;
        }
        next_logical = next_logical
            .checked_add(physical.length_bytes())
            .ok_or_else(|| "causal attention logical page coverage overflows".to_owned())?;
    }
    if next_logical != expected_physical_bytes || pages.is_empty() {
        return Err("causal attention pages do not cover the admitted state".to_owned());
    }
    Ok(pages)
}

fn binding_payload(
    layout: CausalKvLayout,
    table_entries: i32,
    position_start: i32,
    active_tokens: i32,
    sequence_tokens: i32,
    packed_token_start: i32,
    pages: &[CudaBufferRegion],
) -> Result<Box<[u8]>, String> {
    let page_addresses = pages
        .iter()
        .map(CudaBufferRegion::device_ptr)
        .collect::<Vec<_>>();
    let addresses = binding_addresses(layout, table_entries, &page_addresses)?;
    let mut payload = Vec::with_capacity(
        BINDING_CONTROL_BYTES as usize + addresses.len() * std::mem::size_of::<u64>(),
    );
    for value in [
        table_entries,
        position_start,
        active_tokens,
        sequence_tokens,
        packed_token_start,
        0,
    ] {
        payload.extend_from_slice(&value.to_ne_bytes());
    }
    for address in addresses {
        payload.extend_from_slice(&address.to_ne_bytes());
    }
    Ok(payload.into_boxed_slice())
}

fn binding_addresses(
    layout: CausalKvLayout,
    table_entries: i32,
    page_addresses: &[u64],
) -> Result<Vec<u64>, String> {
    let table_entries_usize = usize::try_from(table_entries)
        .map_err(|_| "causal attention address-table count is negative".to_owned())?;
    let mut addresses = Vec::with_capacity(table_entries_usize);
    match layout {
        CausalKvLayout::TokenMajorPages => {
            if table_entries_usize != page_addresses.len() {
                return Err(
                    "token-major causal attention table does not match its retained pages"
                        .to_owned(),
                );
            }
            addresses.extend_from_slice(page_addresses);
        }
        CausalKvLayout::VllmBlocks16 {
            combined_block_bytes,
            blocks_per_page,
        } => {
            let blocks_per_page = usize::try_from(blocks_per_page)
                .map_err(|_| "causal attention blocks per page exceed usize".to_owned())?;
            if blocks_per_page == 0
                || combined_block_bytes == 0
                || combined_block_bytes > VNEXT_KV_PAGE_BYTES
            {
                return Err("causal attention vLLM block geometry is invalid".to_owned());
            }
            for logical_block in 0..table_entries_usize {
                let page = *page_addresses
                    .get(logical_block / blocks_per_page)
                    .ok_or_else(|| {
                        "causal attention vLLM address table exceeds retained pages".to_owned()
                    })?;
                let offset = u64::try_from(logical_block % blocks_per_page)
                    .ok()
                    .and_then(|block| block.checked_mul(combined_block_bytes))
                    .filter(|offset| {
                        offset
                            .checked_add(combined_block_bytes)
                            .is_some_and(|end| end <= VNEXT_KV_PAGE_BYTES)
                    })
                    .ok_or_else(|| "causal attention vLLM block offset overflows".to_owned())?;
                addresses.push(
                    page.checked_add(offset).ok_or_else(|| {
                        "causal attention vLLM block address overflows".to_owned()
                    })?,
                );
            }
        }
    }
    Ok(addresses)
}

fn validate_signature(
    participant: &OperationInvocation<'_, CudaDeviceBuffer>,
    shape: CausalAttentionShape,
    semantics: CausalAttentionSemantics,
) -> Result<(), String> {
    let value = |ordinal| binding(participant.bindings(), ResolvedValueRole::Input, ordinal);
    let hidden = value(0)?;
    let output = binding(participant.bindings(), ResolvedValueRole::Output, 0)?;
    let [tokens, hidden_width] = hidden.tensor().dimensions() else {
        return Err("causal attention hidden input is not two-dimensional".to_owned());
    };
    let mut expected = vec![
        (value(1)?, vec![shape.hidden_size]),
        (
            value(2)?,
            vec![shape.query_projection_features, shape.hidden_size],
        ),
        (value(3)?, vec![shape.kv_features, shape.hidden_size]),
        (value(4)?, vec![shape.kv_features, shape.hidden_size]),
        (value(5)?, vec![shape.hidden_size, shape.query_features]),
        (value(6)?, vec![shape.head_dim]),
        (value(7)?, vec![shape.head_dim]),
        (value(8)?, vec![2, shape.key_value_heads, shape.head_dim]),
    ];
    if semantics.has_post_attention_norm() {
        expected.push((value(9)?, vec![shape.hidden_size]));
    }
    if shape.attention_k_eq_v && value(3)?.value_id() != value(4)?.value_id() {
        return Err(
            "causal attention attention_k_eq_v=true requires K and V to bind one typed value"
                .to_owned(),
        );
    }
    if *tokens == 0
        || *hidden_width != shape.hidden_size
        || output.tensor().dimensions() != [*tokens, shape.hidden_size]
        || !f16_contiguous(hidden)
        || !f16_contiguous(output)
        || expected.iter().any(|(binding, dimensions)| {
            binding.tensor().dimensions() != dimensions.as_slice() || !f16_contiguous(binding)
        })
    {
        return Err("causal attention signature differs from its resolved shape".to_owned());
    }
    Ok(())
}

fn push_shared_projection_weight(
    regions: &mut Vec<CudaBufferRegion>,
    invocation: &BatchedOperationInvocation<'_, CudaDeviceBuffer>,
    ordinal: u32,
    logical_dimensions: &[u64],
) -> Result<SharedProjectionWeight, String> {
    let [expected_output_features, expected_input_features] = logical_dimensions else {
        return Err(format!(
            "causal attention projection input {ordinal} must have two logical dimensions"
        ));
    };
    #[cfg(feature = "vllm-marlin")]
    {
        let first_participant = &invocation.participants()[0];
        let first_binding = binding(
            first_participant.bindings(),
            ResolvedValueRole::Input,
            ordinal,
        )?;
        let first_layout = first_binding.weight().ok_or_else(|| {
            format!("causal attention projection input {ordinal} lacks its physical weight layout")
        })?;
        let quantization_formats = first_layout.quantization_formats();
        let weight_abi = causal_projection_weight_abi(&quantization_formats)?;
        if weight_abi == CausalProjectionWeightAbi::MarlinFp8 {
            let first =
                resolve_marlin_fp8_weight(first_participant, first_binding, logical_dimensions)?;
            if first.output_features() != *expected_output_features
                || first.input_features() != *expected_input_features
            {
                return Err(format!(
                    "causal attention projection input {ordinal} resolved inconsistent Marlin FP8 dimensions"
                ));
            }
            for participant in &invocation.participants()[1..] {
                let candidate = resolve_marlin_fp8_weight(
                    participant,
                    binding(participant.bindings(), ResolvedValueRole::Input, ordinal)?,
                    logical_dimensions,
                )?;
                if candidate.output_features() != first.output_features()
                    || candidate.input_features() != first.input_features()
                    || candidate.group_size() != first.group_size()
                    || !super::same_physical_region(
                        first.packed_region(),
                        candidate.packed_region(),
                    )
                    || !super::same_physical_region(
                        first.scales_region(),
                        candidate.scales_region(),
                    )
                {
                    return Err(format!(
                        "causal attention projection input {ordinal} is not one shared physical Marlin FP8 matrix"
                    ));
                }
            }
            let group_size = first.group_size();
            let [packed, scales] = first.into_regions();
            let packed_region = regions.len();
            regions.push(packed);
            let scales_region = regions.len();
            regions.push(scales);
            return Ok(SharedProjectionWeight::Marlin {
                packed_region,
                scales_region,
                zero_points_region: None,
                group_size,
                weight_type: MarlinF16WeightType::E4M3Fn,
            });
        }
        if weight_abi == CausalProjectionWeightAbi::CompressedTensorsMarlinInt4 {
            let first = resolve_compressed_tensors_marlin_matrix_weight(
                first_participant,
                first_binding,
                logical_dimensions,
            )?;
            for participant in &invocation.participants()[1..] {
                let candidate = resolve_compressed_tensors_marlin_matrix_weight(
                    participant,
                    binding(participant.bindings(), ResolvedValueRole::Input, ordinal)?,
                    logical_dimensions,
                )?;
                if candidate.logical_dimensions() != first.logical_dimensions()
                    || candidate.packed_physical_dimensions() != first.packed_physical_dimensions()
                    || candidate.scales_physical_dimensions() != first.scales_physical_dimensions()
                    || candidate.zero_points_physical_dimensions()
                        != first.zero_points_physical_dimensions()
                    || candidate.group_size() != first.group_size()
                    || !super::same_physical_region(
                        first.packed_region(),
                        candidate.packed_region(),
                    )
                    || !super::same_physical_region(
                        first.scales_region(),
                        candidate.scales_region(),
                    )
                    || !super::same_physical_region(
                        first.zero_points_region(),
                        candidate.zero_points_region(),
                    )
                {
                    return Err(format!(
                        "causal attention projection input {ordinal} is not one shared compressed-tensors Marlin matrix"
                    ));
                }
            }
            let group_size = i32::try_from(first.group_size()).map_err(|_| {
                format!("causal attention projection input {ordinal} group size exceeds i32")
            })?;
            let [packed, scales, zero_points] = first.into_regions();
            let packed_region = regions.len();
            regions.push(packed);
            let scales_region = regions.len();
            regions.push(scales);
            let zero_points_region = regions.len();
            regions.push(zero_points);
            return Ok(SharedProjectionWeight::Marlin {
                packed_region,
                scales_region,
                zero_points_region: Some(zero_points_region),
                group_size,
                weight_type: MarlinF16WeightType::U4,
            });
        }
        if weight_abi == CausalProjectionWeightAbi::CompressedTensorsMarlinSymmetricInt4 {
            let first = resolve_compressed_tensors_symmetric_marlin_matrix_weight(
                first_participant,
                first_binding,
                logical_dimensions,
            )?;
            if first.logical_dimensions() != logical_dimensions
                || first.logical_dimensions()
                    != [*expected_output_features, *expected_input_features]
            {
                return Err(format!(
                    "causal attention projection input {ordinal} resolved inconsistent symmetric compressed-tensors dimensions"
                ));
            }
            for participant in &invocation.participants()[1..] {
                let candidate = resolve_compressed_tensors_symmetric_marlin_matrix_weight(
                    participant,
                    binding(participant.bindings(), ResolvedValueRole::Input, ordinal)?,
                    logical_dimensions,
                )?;
                if candidate.logical_dimensions() != first.logical_dimensions()
                    || candidate.packed_physical_dimensions() != first.packed_physical_dimensions()
                    || candidate.scales_physical_dimensions() != first.scales_physical_dimensions()
                    || candidate.group_size() != first.group_size()
                    || !super::same_physical_region(
                        first.packed_region(),
                        candidate.packed_region(),
                    )
                    || !super::same_physical_region(
                        first.scales_region(),
                        candidate.scales_region(),
                    )
                {
                    return Err(format!(
                        "causal attention projection input {ordinal} is not one shared symmetric compressed-tensors Marlin matrix"
                    ));
                }
            }
            let group_size = i32::try_from(first.group_size()).map_err(|_| {
                format!("causal attention projection input {ordinal} group size exceeds i32")
            })?;
            let [packed, scales] = first.into_regions();
            let packed_region = regions.len();
            regions.push(packed);
            let scales_region = regions.len();
            regions.push(scales);
            return Ok(SharedProjectionWeight::Marlin {
                packed_region,
                scales_region,
                zero_points_region: None,
                group_size,
                weight_type: MarlinF16WeightType::U4B8,
            });
        }
        if weight_abi == CausalProjectionWeightAbi::GptqMarlinInt4 {
            let first = resolve_gptq_marlin_matrix_weight(
                first_participant,
                first_binding,
                logical_dimensions,
            )?;
            if first.logical_dimensions() != logical_dimensions
                || first.logical_dimensions()
                    != [*expected_output_features, *expected_input_features]
            {
                return Err(format!(
                    "causal attention projection input {ordinal} resolved inconsistent GPTQ-Marlin dimensions"
                ));
            }
            for participant in &invocation.participants()[1..] {
                let candidate = resolve_gptq_marlin_matrix_weight(
                    participant,
                    binding(participant.bindings(), ResolvedValueRole::Input, ordinal)?,
                    logical_dimensions,
                )?;
                if candidate.logical_dimensions() != first.logical_dimensions()
                    || candidate.packed_physical_dimensions() != first.packed_physical_dimensions()
                    || candidate.scales_physical_dimensions() != first.scales_physical_dimensions()
                    || candidate.group_size() != first.group_size()
                    || !super::same_physical_region(
                        first.packed_region(),
                        candidate.packed_region(),
                    )
                    || !super::same_physical_region(
                        first.scales_region(),
                        candidate.scales_region(),
                    )
                {
                    return Err(format!(
                        "causal attention projection input {ordinal} is not one shared physical GPTQ-Marlin matrix"
                    ));
                }
            }
            let group_size = i32::try_from(first.group_size()).map_err(|_| {
                format!("causal attention projection input {ordinal} group size exceeds i32")
            })?;
            let [packed, scales] = first.into_regions();
            let packed_region = regions.len();
            regions.push(packed);
            let scales_region = regions.len();
            regions.push(scales);
            return Ok(SharedProjectionWeight::Marlin {
                packed_region,
                scales_region,
                zero_points_region: None,
                group_size,
                weight_type: MarlinF16WeightType::U4B8,
            });
        }
        if weight_abi != CausalProjectionWeightAbi::F16 {
            return Err(format!(
                "causal attention projection input {ordinal} has an unhandled resolved weight ABI"
            ));
        }
    }

    Ok(SharedProjectionWeight::F16 {
        region: push_shared_weight(regions, invocation, ordinal)?,
    })
}

fn push_shared_weight(
    regions: &mut Vec<CudaBufferRegion>,
    invocation: &BatchedOperationInvocation<'_, CudaDeviceBuffer>,
    ordinal: u32,
) -> Result<usize, String> {
    let index = regions.len();
    regions.push(super::shared_full_region(
        invocation,
        ResolvedValueRole::Input,
        ordinal,
        ElementType::F16,
    )?);
    Ok(index)
}

fn f16_contiguous(binding: &ResolvedValueBinding) -> bool {
    binding.tensor().element_type() == ElementType::F16
        && matches!(binding.tensor().layout(), ResolvedTensorLayout::Contiguous)
}

fn reserve_tokens(offset: &mut u64, elements: u64, tokens: u64) -> Result<u64, String> {
    let start = *offset;
    let stride = aligned_bytes(elements, ElementType::F16.size_bytes())?;
    *offset = offset
        .checked_add(
            stride
                .checked_mul(tokens)
                .ok_or_else(|| "causal attention scratch span overflows".to_owned())?,
        )
        .ok_or_else(|| "causal attention scratch offset overflows".to_owned())?;
    Ok(start)
}

fn reserve_elements(offset: &mut u64, elements: u64, element_bytes: u64) -> Result<u64, String> {
    let start = *offset;
    *offset = offset
        .checked_add(aligned_bytes(elements, element_bytes)?)
        .ok_or_else(|| "causal attention scratch offset overflows".to_owned())?;
    Ok(start)
}

fn aligned_bytes(elements: u64, element_bytes: u64) -> Result<u64, String> {
    let bytes = elements
        .checked_mul(element_bytes)
        .ok_or_else(|| "causal attention byte count overflows".to_owned())?;
    align_up(bytes, SCRATCH_ALIGNMENT)
}

fn align_up(bytes: u64, alignment: u64) -> Result<u64, String> {
    bytes
        .checked_add(alignment - 1)
        .map(|value| value & !(alignment - 1))
        .filter(|value| *value > 0)
        .ok_or_else(|| "causal attention alignment overflows".to_owned())
}

fn scratch_pointer(base: u64, offset: u64) -> Result<u64, CudaDeviceRuntimeError> {
    base.checked_add(offset).ok_or_else(|| {
        CudaDeviceRuntimeError::contract("causal attention scratch pointer overflows")
    })
}

fn checked_i32(value: u64, context: &str) -> Result<i32, String> {
    i32::try_from(value).map_err(|_| format!("{context} exceeds i32"))
}

fn checked_i32_runtime(value: u64, context: &'static str) -> Result<i32, CudaDeviceRuntimeError> {
    i32::try_from(value)
        .map_err(|_| CudaDeviceRuntimeError::contract(format!("{context} exceeds i32")))
}

fn checked_u32_runtime(value: u64, context: &'static str) -> Result<u32, CudaDeviceRuntimeError> {
    u32::try_from(value)
        .map_err(|_| CudaDeviceRuntimeError::contract(format!("{context} exceeds u32")))
}

fn unsigned_attribute(
    attributes: &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 causal attention lacks unsigned attribute {name:?}"
        )),
    }
}

fn bool_attribute(
    attributes: &BTreeMap<AttributeId, SemanticValue>,
    name: &str,
) -> Result<bool, String> {
    match attributes
        .iter()
        .find(|(attribute, _)| attribute.as_str() == name)
        .map(|(_, value)| value)
    {
        Some(SemanticValue::Bool(value)) => Ok(*value),
        _ => Err(format!(
            "CUDA causal attention lacks boolean attribute {name:?}"
        )),
    }
}

fn rational_attribute(
    attributes: &BTreeMap<AttributeId, SemanticValue>,
    name: &str,
) -> Result<f32, String> {
    let rational = match attributes
        .iter()
        .find(|(attribute, _)| attribute.as_str() == name)
        .map(|(_, value)| value)
    {
        Some(SemanticValue::Rational(value)) => *value,
        _ => {
            return Err(format!(
                "CUDA causal attention lacks rational attribute {name:?}"
            ))
        }
    };
    let value = (rational.numerator() as f64 / rational.denominator() as f64) as f32;
    if !value.is_finite() || value <= 0.0 {
        return Err(format!(
            "CUDA causal attention rational attribute {name:?} is not a positive f32"
        ));
    }
    Ok(value)
}

fn invalid_plan(reason: impl Into<String>) -> VNextError {
    VNextError::InvalidExecutionPlan {
        reason: reason.into(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ferrum_interfaces::vnext::CanonicalRational;

    fn attributes(output_gate: bool) -> BTreeMap<AttributeId, SemanticValue> {
        BTreeMap::from([
            (
                AttributeId::new("hidden_size").unwrap(),
                SemanticValue::Unsigned(2048),
            ),
            (
                AttributeId::new("query_heads").unwrap(),
                SemanticValue::Unsigned(16),
            ),
            (
                AttributeId::new("key_value_heads").unwrap(),
                SemanticValue::Unsigned(2),
            ),
            (
                AttributeId::new("head_dim").unwrap(),
                SemanticValue::Unsigned(128),
            ),
            (
                AttributeId::new("query_features").unwrap(),
                SemanticValue::Unsigned(2048),
            ),
            (
                AttributeId::new("query_projection_features").unwrap(),
                SemanticValue::Unsigned(if output_gate { 4096 } else { 2048 }),
            ),
            (
                AttributeId::new("kv_features").unwrap(),
                SemanticValue::Unsigned(256),
            ),
            (
                AttributeId::new("rope_dim").unwrap(),
                SemanticValue::Unsigned(64),
            ),
            (
                AttributeId::new("maximum_context_tokens").unwrap(),
                SemanticValue::Unsigned(4096),
            ),
            (
                AttributeId::new("epsilon").unwrap(),
                SemanticValue::Rational(CanonicalRational::new(1, 1_000_000).unwrap()),
            ),
            (
                AttributeId::new("rope_theta").unwrap(),
                SemanticValue::Rational(CanonicalRational::new(10_000, 1).unwrap()),
            ),
            (
                AttributeId::new("rope_interleaved").unwrap(),
                SemanticValue::Bool(false),
            ),
            (
                AttributeId::new("output_gate").unwrap(),
                SemanticValue::Bool(output_gate),
            ),
            (
                AttributeId::new("causal").unwrap(),
                SemanticValue::Bool(true),
            ),
            (
                AttributeId::new("layer_index").unwrap(),
                SemanticValue::Unsigned(0),
            ),
        ])
    }

    fn gemma4_attributes(full_attention: bool) -> BTreeMap<AttributeId, SemanticValue> {
        let (key_value_heads, head_dim, rope_dim, rope_denominator, rope_theta, window, k_eq_v) =
            if full_attention {
                (1, 512, 128, 512, 1_000_000, 0, true)
            } else {
                (8, 256, 256, 256, 10_000, 1_024, false)
            };
        let query_heads = 16;
        BTreeMap::from([
            (
                AttributeId::new("hidden_size").unwrap(),
                SemanticValue::Unsigned(3_840),
            ),
            (
                AttributeId::new("query_heads").unwrap(),
                SemanticValue::Unsigned(query_heads),
            ),
            (
                AttributeId::new("key_value_heads").unwrap(),
                SemanticValue::Unsigned(key_value_heads),
            ),
            (
                AttributeId::new("head_dim").unwrap(),
                SemanticValue::Unsigned(head_dim),
            ),
            (
                AttributeId::new("query_features").unwrap(),
                SemanticValue::Unsigned(query_heads * head_dim),
            ),
            (
                AttributeId::new("query_projection_features").unwrap(),
                SemanticValue::Unsigned(query_heads * head_dim),
            ),
            (
                AttributeId::new("kv_features").unwrap(),
                SemanticValue::Unsigned(key_value_heads * head_dim),
            ),
            (
                AttributeId::new("rope_dim").unwrap(),
                SemanticValue::Unsigned(rope_dim),
            ),
            (
                AttributeId::new("rope_frequency_denominator").unwrap(),
                SemanticValue::Unsigned(rope_denominator),
            ),
            (
                AttributeId::new("maximum_context_tokens").unwrap(),
                SemanticValue::Unsigned(262_144),
            ),
            (
                AttributeId::new("epsilon").unwrap(),
                SemanticValue::Rational(CanonicalRational::new(1, 1_000_000).unwrap()),
            ),
            (
                AttributeId::new("rope_theta").unwrap(),
                SemanticValue::Rational(CanonicalRational::new(rope_theta, 1).unwrap()),
            ),
            (
                AttributeId::new("attention_scale").unwrap(),
                SemanticValue::Rational(CanonicalRational::new(1, 1).unwrap()),
            ),
            (
                AttributeId::new("sliding_window_tokens").unwrap(),
                SemanticValue::Unsigned(window),
            ),
            (
                AttributeId::new("rope_interleaved").unwrap(),
                SemanticValue::Bool(false),
            ),
            (
                AttributeId::new("value_rms_norm").unwrap(),
                SemanticValue::Bool(true),
            ),
            (
                AttributeId::new("attention_k_eq_v").unwrap(),
                SemanticValue::Bool(k_eq_v),
            ),
            (
                AttributeId::new("causal").unwrap(),
                SemanticValue::Bool(true),
            ),
            (
                AttributeId::new("layer_index").unwrap(),
                SemanticValue::Unsigned(0),
            ),
        ])
    }

    fn goal_shape(
        query_heads: u64,
        key_value_heads: u64,
        head_dim: u64,
        maximum_context_tokens: u64,
    ) -> CausalAttentionShape {
        let query_features = query_heads * head_dim;
        let kv_features = key_value_heads * head_dim;
        CausalAttentionShape {
            hidden_size: query_features,
            query_heads,
            key_value_heads,
            head_dim,
            query_features,
            query_projection_features: query_features,
            kv_features,
            rope_dim: head_dim,
            rope_frequency_denominator: head_dim,
            maximum_context_tokens,
            epsilon: 1e-6,
            rope_theta: 10_000.0,
            attention_scale: 1.0_f32 / (head_dim as f32).sqrt(),
            sliding_window_tokens: 0,
            rope_interleaved: false,
            output_gate: false,
            value_rms_norm: false,
            attention_k_eq_v: false,
            post_attention_norm: false,
        }
    }

    fn select_native_path(
        shape: CausalAttentionShape,
        active_tokens: u64,
        sequence_tokens: u64,
    ) -> Result<CausalAttentionKernelPath, String> {
        CausalAttentionKernelPath::select(
            AttentionExecutionPolicy::NativeAdaptive,
            shape,
            active_tokens,
            sequence_tokens,
        )
    }

    #[cfg(feature = "vllm-marlin")]
    fn quantization_formats(ids: &[&str]) -> BTreeSet<QuantizationFormatId> {
        ids.iter()
            .map(|id| QuantizationFormatId::new(*id).unwrap())
            .collect()
    }

    #[cfg(feature = "vllm-marlin")]
    #[test]
    fn causal_projection_weight_abi_is_exact_and_rejects_source_fp8() {
        assert_eq!(
            causal_projection_weight_abi(&BTreeSet::new()).unwrap(),
            CausalProjectionWeightAbi::F16
        );
        for (format, expected) in [
            (
                MARLIN_FP8_QUANTIZATION_FORMAT_ID,
                CausalProjectionWeightAbi::MarlinFp8,
            ),
            (
                MARLIN_FP8_GROUP128_QUANTIZATION_FORMAT_ID,
                CausalProjectionWeightAbi::MarlinFp8,
            ),
            (
                GPTQ_MARLIN_QUANTIZATION_FORMAT_ID,
                CausalProjectionWeightAbi::GptqMarlinInt4,
            ),
            (
                COMPRESSED_TENSORS_MARLIN_QUANTIZATION_FORMAT_ID,
                CausalProjectionWeightAbi::CompressedTensorsMarlinInt4,
            ),
            (
                COMPRESSED_TENSORS_MARLIN_SYMMETRIC_QUANTIZATION_FORMAT_ID,
                CausalProjectionWeightAbi::CompressedTensorsMarlinSymmetricInt4,
            ),
        ] {
            assert_eq!(
                causal_projection_weight_abi(&quantization_formats(&[format])).unwrap(),
                expected
            );
        }

        let source_fp8 =
            quantization_formats(&["quantization.safetensors.fp8-e4m3-block-grid-inverse-scale"]);
        let error = causal_projection_weight_abi(&source_fp8).unwrap_err();
        assert!(error.contains("is not admitted"), "{error}");

        let multiple = quantization_formats(&[
            MARLIN_FP8_QUANTIZATION_FORMAT_ID,
            GPTQ_MARLIN_QUANTIZATION_FORMAT_ID,
        ]);
        let error = causal_projection_weight_abi(&multiple).unwrap_err();
        assert!(error.contains("more than one"), "{error}");
    }

    #[cfg(feature = "vllm-marlin")]
    #[test]
    fn causal_projection_abi_rejects_fp8_int4_mixing() {
        assert_eq!(
            causal_projection_abi([CausalProjectionWeightAbi::F16; 4]).unwrap(),
            CausalProjectionAbi::F16
        );
        assert_eq!(
            causal_projection_abi([
                CausalProjectionWeightAbi::MarlinFp8,
                CausalProjectionWeightAbi::F16,
                CausalProjectionWeightAbi::MarlinFp8,
                CausalProjectionWeightAbi::F16,
            ])
            .unwrap(),
            CausalProjectionAbi::MarlinFp8
        );
        assert_eq!(
            causal_projection_abi([
                CausalProjectionWeightAbi::GptqMarlinInt4,
                CausalProjectionWeightAbi::CompressedTensorsMarlinSymmetricInt4,
                CausalProjectionWeightAbi::F16,
                CausalProjectionWeightAbi::F16,
            ])
            .unwrap(),
            CausalProjectionAbi::MarlinInt4
        );
        let error = causal_projection_abi([
            CausalProjectionWeightAbi::MarlinFp8,
            CausalProjectionWeightAbi::GptqMarlinInt4,
        ])
        .unwrap_err();
        assert!(error.contains("cannot mix FP8 and INT4"), "{error}");
    }

    #[cfg(feature = "vllm-marlin")]
    #[test]
    fn causal_marlin_replay_identity_distinguishes_all_weight_abis() {
        let fp8 = marlin_projection_replay_tag(MarlinF16WeightType::E4M3Fn);
        let compressed_tensors = marlin_projection_replay_tag(MarlinF16WeightType::U4);
        let gptq = marlin_projection_replay_tag(MarlinF16WeightType::U4B8);
        assert_eq!(fp8, "marlin-fp8-e4m3fn");
        assert_eq!(compressed_tensors, "compressed-tensors-marlin-u4-zp");
        assert_eq!(gptq, "gptq-marlin-u4b8");
        assert_ne!(fp8, compressed_tensors);
        assert_ne!(fp8, gptq);
        assert_ne!(compressed_tensors, gptq);
    }

    #[cfg(feature = "vllm-marlin")]
    #[test]
    fn causal_marlin_replay_identity_includes_fp8_group_size() {
        let weight = |group_size| SharedProjectionWeight::Marlin {
            packed_region: 0,
            scales_region: 1,
            zero_points_region: None,
            group_size,
            weight_type: MarlinF16WeightType::E4M3Fn,
        };
        let channelwise = weight(-1)
            .bind_replay_abi(CudaCommandReplayKeyBuilder::new("test", "causal"))
            .finish();
        let group128 = weight(128)
            .bind_replay_abi(CudaCommandReplayKeyBuilder::new("test", "causal"))
            .finish();
        assert_ne!(channelwise, group128);
    }

    #[test]
    fn scratch_estimator_and_layout_are_identical() {
        let shape = CausalAttentionShape::from_attributes(&attributes(true)).unwrap();
        let layout = ScratchLayout::new(
            shape,
            17,
            CausalProjection::F16,
            AttentionExecutionPolicy::NativeAdaptive,
        )
        .unwrap();
        let bindings = BindingLayout::new(shape, 3).unwrap();
        assert_eq!(
            layout.required_bytes,
            17 * shape.scratch_bytes_per_token().unwrap() + shape.vllm_scratch_bytes().unwrap()
        );
        assert_eq!(
            bindings.required_bytes,
            3 * shape.binding_slot_bytes().unwrap()
        );
        assert_eq!(bindings.binding_offset(0).unwrap(), 0);
        assert_eq!(
            bindings.binding_offset(2).unwrap(),
            2 * shape.binding_slot_bytes().unwrap()
        );
        assert_eq!(shape.maximum_pages().unwrap(), 64);
    }

    #[cfg(feature = "vllm-marlin")]
    #[test]
    fn l40s_marlin_workspace_keeps_causal_scratch_estimator_and_layout_identical() {
        let shape = CausalAttentionShape::from_attributes(&attributes(true)).unwrap();
        let runtime = MarlinProjectionRuntime {
            multiprocessor_count: 142,
            device_ordinal: 0,
        };
        let raw_workspace_bytes = runtime.workspace_bytes().unwrap();
        assert_eq!(raw_workspace_bytes, 568);
        let reserved_workspace_bytes = aligned_bytes(raw_workspace_bytes, 1).unwrap();
        assert_eq!(reserved_workspace_bytes, 576);

        let projection = CausalProjection::MarlinFp8 { runtime };
        assert_eq!(
            projection.workspace_reservation_bytes().unwrap(),
            reserved_workspace_bytes
        );
        let total_tokens = 1;
        let layout = ScratchLayout::new(
            shape,
            total_tokens,
            projection,
            AttentionExecutionPolicy::Portable,
        )
        .unwrap();
        let expected =
            reserved_workspace_bytes + shape.scratch_bytes_per_token().unwrap() * total_tokens;
        assert_eq!(layout.required_bytes, expected);
    }

    #[test]
    fn portable_policy_omits_native_partition_scratch() {
        let shape = CausalAttentionShape::from_attributes(&attributes(true)).unwrap();
        let portable = ScratchLayout::new(
            shape,
            17,
            CausalProjection::F16,
            AttentionExecutionPolicy::Portable,
        )
        .unwrap();
        let native = ScratchLayout::new(
            shape,
            17,
            CausalProjection::F16,
            AttentionExecutionPolicy::NativeAdaptive,
        )
        .unwrap();
        assert!(portable.vllm.is_none());
        assert_eq!(
            portable.required_bytes,
            17 * shape.scratch_bytes_per_token().unwrap()
        );
        assert_eq!(
            native.required_bytes - portable.required_bytes,
            shape.vllm_scratch_bytes().unwrap()
        );
    }

    #[test]
    fn chunked_prefill_state_tracks_the_source_frontier() {
        let shape = CausalAttentionShape::from_attributes(&attributes(true)).unwrap();
        let chunk_state = shape.physical_state_bytes(3).unwrap();
        let full_prompt_state = shape.physical_state_bytes(128).unwrap();

        assert_ne!(chunk_state, full_prompt_state);
        assert_eq!(
            shape
                .physical_state_bytes_for_source_frontier(3, 128)
                .unwrap(),
            chunk_state
        );
        assert!(shape
            .physical_state_bytes_for_source_frontier(129, 128)
            .is_err());
    }

    #[test]
    fn output_gate_changes_the_typed_query_projection_width() {
        let gated = CausalAttentionShape::from_attributes(&attributes(true)).unwrap();
        let plain = CausalAttentionShape::from_attributes(&attributes(false)).unwrap();
        assert_eq!(
            gated.query_projection_features,
            2 * plain.query_projection_features
        );
        let mut invalid = attributes(true);
        invalid.insert(
            AttributeId::new("query_projection_features").unwrap(),
            SemanticValue::Unsigned(2048),
        );
        assert!(CausalAttentionShape::from_attributes(&invalid).is_err());
    }

    #[test]
    fn packed_dispatch_count_keeps_only_sequence_local_work_per_participant() {
        let v1 = CausalAttentionKernelPath::VllmAddressedDecodeV1;
        let v2 = CausalAttentionKernelPath::VllmAddressedDecodeV2;
        let token_fallback = CausalAttentionKernelPath::TokenMajorFallback;
        let addressed_fallback = CausalAttentionKernelPath::VllmAddressedFallback;

        assert_eq!(physical_dispatch_count([v1], false, false, false), 8);
        assert_eq!(physical_dispatch_count([v1; 4], false, false, false), 32);
        assert_eq!(physical_dispatch_count([v1; 4], false, false, true), 14);
        assert_eq!(physical_dispatch_count([v2; 4], true, false, false), 40);
        assert_eq!(physical_dispatch_count([v2; 4], true, false, true), 22);
        assert_eq!(physical_dispatch_count([v1; 32], true, false, true), 102);
        assert_eq!(physical_dispatch_count([v2; 32], true, false, true), 134);
        assert_eq!(physical_dispatch_count([v1], false, true, false), 9);
        assert_eq!(physical_dispatch_count([v1; 4], false, true, true), 15);
        assert_eq!(
            physical_dispatch_count([token_fallback; 4], false, false, false),
            32
        );
        assert_eq!(
            physical_dispatch_count([token_fallback; 4], false, false, true),
            8
        );
        assert_eq!(
            physical_dispatch_count([token_fallback; 32], false, true, true),
            9
        );
        assert_eq!(
            physical_dispatch_count([addressed_fallback; 32], true, false, true),
            40
        );
        assert_eq!(
            physical_dispatch_count([token_fallback, addressed_fallback], false, false, true,),
            10
        );
    }

    #[test]
    fn packed_fallback_launch_requires_one_shared_fallback_path() {
        let layout = BindingLayout {
            required_bytes: 1_024,
            slot_bytes: 256,
        };
        let launch = |path, tokens| CausalAttentionLaunch {
            input_region: 0,
            output_region: 1,
            binding_offset: 0,
            packed_token_start: 0,
            packed_query_raw: 0,
            packed_key_raw: 0,
            packed_value_raw: 0,
            packed_query: 0,
            packed_context: 0,
            tokens,
            tokens_i32: tokens as i32,
            sequence_tokens: tokens,
            sequence_tokens_i32: tokens as i32,
            table_entries_i32: 1,
            replay_topology: CausalAttentionReplayTopology::PartitionStable(
                CausalAttentionReplayEnvelope {
                    sequence_capacity_tokens: tokens,
                    table_capacity_entries: 1,
                },
            ),
            path,
        };

        let packed = packed_fallback_launch(
            &[
                launch(CausalAttentionKernelPath::TokenMajorFallback, 1),
                launch(CausalAttentionKernelPath::TokenMajorFallback, 4),
            ],
            layout,
        )
        .unwrap()
        .unwrap();
        assert_eq!(packed.token_grid, 4);
        assert_eq!(packed.packed_token_grid, 5);
        assert_eq!(packed.participant_grid, 2);
        assert_eq!(packed.participant_count_i32, 2);
        assert_eq!(packed.binding_slot_bytes, 256);
        assert_eq!(packed.path, CausalAttentionKernelPath::TokenMajorFallback);

        let local = CausalAttentionShape::from_attributes_for(
            &gemma4_attributes(false),
            CausalAttentionSemantics::Gemma4,
        )
        .unwrap()
        .cuda_shape()
        .unwrap();
        let global = CausalAttentionShape::from_attributes_for(
            &gemma4_attributes(true),
            CausalAttentionSemantics::Gemma4,
        )
        .unwrap()
        .cuda_shape()
        .unwrap();
        assert_eq!(grouped_fallback_block_threads(local), Some(64));
        assert_eq!(grouped_fallback_block_threads(global), Some(512));

        assert!(packed_fallback_launch(
            &[
                launch(CausalAttentionKernelPath::TokenMajorFallback, 1),
                launch(CausalAttentionKernelPath::VllmAddressedFallback, 1),
            ],
            layout,
        )
        .unwrap()
        .is_none());
        assert!(packed_fallback_launch(
            &[
                launch(CausalAttentionKernelPath::VllmAddressedDecodeV1, 1),
                launch(CausalAttentionKernelPath::VllmAddressedDecodeV1, 1),
            ],
            layout,
        )
        .unwrap()
        .is_none());
    }

    #[test]
    fn packed_token_ranges_require_one_dense_canonical_span() {
        assert!(validate_packed_token_ranges([(0, 3), (3, 1), (4, 8)], 12).is_ok());
        assert!(validate_packed_token_ranges([(0, 3), (4, 8)], 12).is_err());
        assert!(validate_packed_token_ranges([(0, 3), (3, 0), (3, 9)], 12).is_err());
        assert!(validate_packed_token_ranges([(0, 3), (3, 8)], 12).is_err());
    }

    #[test]
    fn packed_token_offsets_follow_dense_matrix_rows_not_allocation_padding() {
        let shape = goal_shape(3, 1, 10, 128);
        let layout = ScratchLayout::new(
            shape,
            3,
            CausalProjection::F16,
            AttentionExecutionPolicy::NativeAdaptive,
        )
        .unwrap();
        assert_eq!(
            layout
                .token_offset(layout.query_raw, 1, shape.query_projection_features)
                .unwrap(),
            layout.query_raw + 60
        );
        assert_eq!(
            aligned_bytes(
                shape.query_projection_features,
                ElementType::F16.size_bytes()
            )
            .unwrap(),
            64
        );
    }

    #[test]
    fn vllm_page_geometry_covers_goal_model_families() {
        let cases = [
            ("qwen35-dense-or-moe", goal_shape(16, 2, 256, 32_768), 2),
            ("qwen3-moe", goal_shape(32, 4, 128, 32_768), 2),
            ("llama-8b-dense", goal_shape(32, 8, 128, 32_768), 1),
        ];
        for (name, shape, expected_blocks_per_page) in cases {
            let CausalKvLayout::VllmBlocks16 {
                combined_block_bytes,
                blocks_per_page,
            } = shape.kv_layout().unwrap()
            else {
                panic!("{name} did not select the vLLM block layout");
            };
            assert_eq!(
                blocks_per_page, expected_blocks_per_page,
                "{name} blocks/page"
            );
            assert!(combined_block_bytes <= VNEXT_KV_PAGE_BYTES, "{name}");
            assert_eq!(shape.table_entries(17).unwrap(), 2, "{name}");
            for tokens in [1, 15, 16, 17, 31, 32, 33, 127, 128, 129] {
                assert_eq!(
                    shape.physical_state_bytes(tokens).unwrap(),
                    align_up(
                        shape.state_bytes_per_token().unwrap() * tokens,
                        VNEXT_KV_PAGE_BYTES
                    )
                    .unwrap(),
                    "{name} tokens={tokens}"
                );
            }
        }
    }

    #[test]
    fn non_divisible_vllm_block_geometry_uses_token_major_pages() {
        let shape = goal_shape(12, 3, 128, 32_768);
        assert_eq!(shape.kv_layout().unwrap(), CausalKvLayout::TokenMajorPages);
        assert_eq!(
            shape.physical_state_bytes(33).unwrap(),
            align_up(
                shape.state_bytes_per_token().unwrap() * 33,
                VNEXT_KV_PAGE_BYTES
            )
            .unwrap()
        );
        assert_eq!(
            select_native_path(shape, 1, 33).unwrap(),
            CausalAttentionKernelPath::TokenMajorFallback
        );
    }

    #[test]
    fn addressed_blocks_expand_inside_retained_pages() {
        let qwen3 = goal_shape(32, 4, 128, 32_768);
        let layout = qwen3.kv_layout().unwrap();
        let addresses = binding_addresses(layout, 3, &[0x10_0000, 0x20_0000]).unwrap();
        assert_eq!(addresses, vec![0x10_0000, 0x10_0000 + 32 * 1024, 0x20_0000]);
        assert!(binding_addresses(layout, 5, &[0x10_0000, 0x20_0000]).is_err());
    }

    #[test]
    fn qwen35_block_49_to_50_reuses_the_second_half_of_page_25() {
        let qwen35 = goal_shape(16, 2, 256, 4_096);
        let layout = qwen35.kv_layout().unwrap();
        assert_eq!(
            layout,
            CausalKvLayout::VllmBlocks16 {
                combined_block_bytes: 32 * 1024,
                blocks_per_page: 2,
            }
        );
        assert_eq!(qwen35.table_entries(784).unwrap(), 49);
        assert_eq!(qwen35.table_entries(785).unwrap(), 50);
        assert_eq!(
            qwen35.physical_state_bytes(784).unwrap(),
            25 * VNEXT_KV_PAGE_BYTES
        );
        assert_eq!(
            qwen35.physical_state_bytes(785).unwrap(),
            25 * VNEXT_KV_PAGE_BYTES
        );

        let pages = (0..25)
            .map(|index| 0x10_0000 + index * 2 * VNEXT_KV_PAGE_BYTES)
            .collect::<Vec<_>>();
        let before = binding_addresses(layout, 49, &pages).unwrap();
        let after = binding_addresses(layout, 50, &pages).unwrap();
        assert_eq!(before[48], pages[24]);
        assert_eq!(after[48], pages[24]);
        assert_eq!(after[49], pages[24] + 32 * 1024);
        assert_eq!(after[..49], before);

        let payload_bytes =
            BINDING_CONTROL_BYTES + u64::try_from(after.len()).unwrap() * POINTER_BYTES;
        assert_eq!(payload_bytes, 424);
        assert!(payload_bytes <= qwen35.binding_slot_bytes().unwrap());
    }

    #[test]
    fn decode_replay_envelope_is_stable_within_native_partition_topology() {
        let shape = goal_shape(32, 4, 128, 32_768);
        let v1_first = CausalAttentionReplayEnvelope::new(
            shape,
            CausalAttentionKernelPath::VllmAddressedDecodeV1,
            1,
        )
        .unwrap();
        let v1_last = CausalAttentionReplayEnvelope::new(
            shape,
            CausalAttentionKernelPath::VllmAddressedDecodeV1,
            512,
        )
        .unwrap();
        assert_eq!(v1_first, v1_last);
        assert_eq!(v1_first.sequence_capacity_tokens, 512);
        assert_eq!(v1_first.table_capacity_entries, 32);
        assert!(CausalAttentionReplayEnvelope::new(
            shape,
            CausalAttentionKernelPath::VllmAddressedDecodeV1,
            513,
        )
        .is_err());

        let v2_first = CausalAttentionReplayEnvelope::new(
            shape,
            CausalAttentionKernelPath::VllmAddressedDecodeV2,
            513,
        )
        .unwrap();
        let v2_last = CausalAttentionReplayEnvelope::new(
            shape,
            CausalAttentionKernelPath::VllmAddressedDecodeV2,
            1_024,
        )
        .unwrap();
        assert_eq!(v2_first, v2_last);
        assert_eq!(v2_first.sequence_capacity_tokens, 1_024);
        assert_eq!(v2_first.table_capacity_entries, 64);
        assert_ne!(
            v2_last,
            CausalAttentionReplayEnvelope::new(
                shape,
                CausalAttentionKernelPath::VllmAddressedDecodeV2,
                1_025,
            )
            .unwrap()
        );
    }

    #[cfg(feature = "vllm-paged-attn-v2")]
    #[test]
    fn reusable_topology_changes_only_at_native_decode_partitions() {
        let shape = goal_shape(16, 2, 256, 4_096);
        let topology = |sequence_tokens| match reusable_attention_topology_from_rows(
            AttentionExecutionPolicy::NativeAdaptive,
            shape,
            1,
            std::iter::once(Ok(CausalAttentionTopologyRow {
                active_tokens: 1,
                sequence_tokens,
            })),
        )
        .unwrap()
        {
            ReusableExecutionTopology::Dynamic(fingerprint) => fingerprint,
            topology => panic!("decode topology must be replayable, got {topology:?}"),
        };

        assert_ne!(topology(512), topology(513));
        assert_eq!(topology(784), topology(785));
        assert_eq!(topology(513), topology(1_024));
        assert_ne!(topology(1_024), topology(1_025));
    }

    #[cfg(feature = "vllm-paged-attn-v2")]
    #[test]
    fn varlen_paths_are_eager_while_portable_fallback_is_replayable() {
        let shape = goal_shape(16, 2, 256, 4_096);
        let topology = |rows: &[(u64, u64)]| {
            reusable_attention_topology_from_rows(
                AttentionExecutionPolicy::NativeAdaptive,
                shape,
                rows.len(),
                rows.iter().map(|&(active_tokens, sequence_tokens)| {
                    Ok(CausalAttentionTopologyRow {
                        active_tokens,
                        sequence_tokens,
                    })
                }),
            )
            .unwrap()
        };

        assert_eq!(
            topology(&[(8, 64)]),
            ReusableExecutionTopology::EagerBoundary
        );
        assert_eq!(
            topology(&[(1, 64), (8, 64)]),
            ReusableExecutionTopology::EagerBoundary
        );
        assert!(matches!(
            reusable_attention_topology_from_rows(
                AttentionExecutionPolicy::Portable,
                shape,
                1,
                std::iter::once(Ok(CausalAttentionTopologyRow {
                    active_tokens: 1,
                    sequence_tokens: 64,
                })),
            )
            .unwrap(),
            ReusableExecutionTopology::Dynamic(_)
        ));
    }

    #[test]
    fn decode_and_fallback_topologies_are_partition_stable() {
        let shape = goal_shape(32, 4, 128, 32_768);
        for path in [
            CausalAttentionKernelPath::VllmAddressedDecodeV1,
            CausalAttentionKernelPath::VllmAddressedDecodeV2,
        ] {
            let sequence_tokens = if path == CausalAttentionKernelPath::VllmAddressedDecodeV1 {
                512
            } else {
                513
            };
            assert!(
                CausalAttentionReplayTopology::new(shape, path, sequence_tokens)
                    .unwrap()
                    .is_partition_stable()
            );
        }

        for path in [
            CausalAttentionKernelPath::TokenMajorFallback,
            CausalAttentionKernelPath::VllmAddressedFallback,
        ] {
            let topology = CausalAttentionReplayTopology::new(shape, path, 64).unwrap();
            assert!(topology.is_partition_stable());
            assert_eq!(
                topology.envelope().sequence_capacity_tokens,
                shape.maximum_context_tokens
            );
        }

        for path in [
            CausalAttentionKernelPath::VllmAddressedVarlen,
            CausalAttentionKernelPath::VllmAddressedVarlenTiled,
        ] {
            assert!(!CausalAttentionReplayTopology::new(shape, path, 64)
                .unwrap()
                .is_partition_stable());
        }
    }

    #[test]
    fn gemma4_fallback_replay_fingerprint_is_stable_across_sequence_frontiers() {
        let fingerprint = |shape, sequence_tokens| match reusable_attention_topology_from_rows(
            AttentionExecutionPolicy::NativeAdaptive,
            shape,
            1,
            std::iter::once(Ok(CausalAttentionTopologyRow {
                active_tokens: 1,
                sequence_tokens,
            })),
        )
        .unwrap()
        {
            ReusableExecutionTopology::Dynamic(fingerprint) => fingerprint,
            topology => panic!("Gemma4 fallback must be replayable, got {topology:?}"),
        };

        for full_attention in [false, true] {
            let shape = CausalAttentionShape::from_attributes_for(
                &gemma4_attributes(full_attention),
                CausalAttentionSemantics::Gemma4,
            )
            .unwrap();
            assert_eq!(fingerprint(shape, 64), fingerprint(shape, 2_048));
        }
    }

    #[test]
    fn kernel_path_records_exact_native_implementation() {
        let shape = goal_shape(16, 4, 256, 32_768);
        assert_eq!(
            select_native_path(shape, 8, 2_048).unwrap(),
            CausalAttentionKernelPath::VllmAddressedVarlenTiled
        );
        assert_eq!(
            select_native_path(shape, 2, 4_000).unwrap(),
            CausalAttentionKernelPath::VllmAddressedVarlen
        );
        assert_eq!(
            select_native_path(shape, 8, 13_000).unwrap(),
            CausalAttentionKernelPath::VllmAddressedFallback
        );
        assert_eq!(
            CausalAttentionKernelPath::VllmAddressedVarlenTiled.native_kernel_id(),
            "ferrum.paged_varlen_attention.vllm_q4_addressed"
        );
        assert_eq!(
            select_native_path(shape, 4, 3_008).unwrap(),
            CausalAttentionKernelPath::VllmAddressedVarlenTiled
        );
        assert_eq!(
            select_native_path(shape, 4, 3_009).unwrap(),
            CausalAttentionKernelPath::VllmAddressedVarlen
        );
        assert_eq!(
            select_native_path(shape, 2, 12_032).unwrap(),
            CausalAttentionKernelPath::VllmAddressedVarlen
        );
        assert_eq!(
            select_native_path(shape, 2, 12_033).unwrap(),
            CausalAttentionKernelPath::VllmAddressedFallback
        );

        let oversized = goal_shape(16, 8, 256, 32_768);
        assert_eq!(
            select_native_path(oversized, 1, 1).unwrap(),
            CausalAttentionKernelPath::TokenMajorFallback
        );
    }

    #[test]
    fn portable_policy_never_selects_optional_vllm_decode() {
        let shape = goal_shape(16, 4, 256, 32_768);
        for sequence_tokens in [1, 512, 513, 32_768] {
            let path = CausalAttentionKernelPath::select(
                AttentionExecutionPolicy::Portable,
                shape,
                1,
                sequence_tokens,
            )
            .unwrap();
            assert!(!matches!(
                path,
                CausalAttentionKernelPath::VllmAddressedDecodeV1
                    | CausalAttentionKernelPath::VllmAddressedDecodeV2
            ));
            assert!(path.native_kernel_id().starts_with("ferrum."));
        }
    }

    #[cfg(feature = "vllm-paged-attn-v2")]
    #[test]
    fn decode_path_selects_vllm_v1_and_v2_without_runtime_env() {
        let shape = goal_shape(16, 4, 256, 32_768);
        let v1 = select_native_path(shape, 1, 512).unwrap();
        let v2 = select_native_path(shape, 1, 513).unwrap();
        assert_eq!(v1, CausalAttentionKernelPath::VllmAddressedDecodeV1);
        assert_eq!(v2, CausalAttentionKernelPath::VllmAddressedDecodeV2);
        assert_eq!(v1.operation(), COMPUTE_VLLM_DECODE_V1_OPERATION);
        assert_eq!(v2.operation(), COMPUTE_VLLM_DECODE_V2_OPERATION);
        assert_eq!(v1.native_kernel_id(), "vllm.paged_attention_v1.addressed");
        assert_eq!(v2.native_kernel_id(), "vllm.paged_attention_v2.addressed");
    }

    #[test]
    fn gemma4_local_and_full_attention_are_typed_and_avoid_incompatible_native_paths() {
        let local = CausalAttentionShape::from_attributes_for(
            &gemma4_attributes(false),
            CausalAttentionSemantics::Gemma4,
        )
        .unwrap();
        assert_eq!(local.head_dim, 256);
        assert_eq!(local.rope_dim, 256);
        assert_eq!(local.rope_frequency_denominator, 256);
        assert_eq!(local.attention_scale, 1.0);
        assert_eq!(local.sliding_window_tokens, 1_024);
        assert!(local.value_rms_norm);
        assert!(local.post_attention_norm);
        assert_eq!(local.kv_layout().unwrap(), CausalKvLayout::TokenMajorPages);
        assert_eq!(
            CausalAttentionKernelPath::select(
                AttentionExecutionPolicy::NativeAdaptive,
                local,
                8,
                2_048,
            )
            .unwrap(),
            CausalAttentionKernelPath::TokenMajorFallback
        );

        let full = CausalAttentionShape::from_attributes_for(
            &gemma4_attributes(true),
            CausalAttentionSemantics::Gemma4,
        )
        .unwrap();
        assert_eq!(full.head_dim, 512);
        assert_eq!(full.rope_dim, 128);
        assert_eq!(full.rope_frequency_denominator, 512);
        assert!(full.attention_k_eq_v);
        assert!(matches!(
            full.kv_layout().unwrap(),
            CausalKvLayout::VllmBlocks16 { .. }
        ));
        for active_tokens in [1, 8] {
            assert_eq!(
                CausalAttentionKernelPath::select(
                    AttentionExecutionPolicy::NativeAdaptive,
                    full,
                    active_tokens,
                    2_048,
                )
                .unwrap(),
                CausalAttentionKernelPath::VllmAddressedFallback
            );
        }
    }

    #[test]
    fn gemma4_attention_rejects_missing_or_inconsistent_semantic_attributes() {
        let mut missing_denominator = gemma4_attributes(true);
        missing_denominator.remove(&AttributeId::new("rope_frequency_denominator").unwrap());
        assert!(CausalAttentionShape::from_attributes_for(
            &missing_denominator,
            CausalAttentionSemantics::Gemma4,
        )
        .is_err());

        let mut denominator_exceeds_head = gemma4_attributes(true);
        denominator_exceeds_head.insert(
            AttributeId::new("rope_frequency_denominator").unwrap(),
            SemanticValue::Unsigned(1_024),
        );
        assert!(CausalAttentionShape::from_attributes_for(
            &denominator_exceeds_head,
            CausalAttentionSemantics::Gemma4,
        )
        .is_err());

        let mut denominator_shorter_than_rope = gemma4_attributes(true);
        denominator_shorter_than_rope.insert(
            AttributeId::new("rope_frequency_denominator").unwrap(),
            SemanticValue::Unsigned(64),
        );
        assert!(CausalAttentionShape::from_attributes_for(
            &denominator_shorter_than_rope,
            CausalAttentionSemantics::Gemma4,
        )
        .is_err());

        let mut oversized_window = gemma4_attributes(false);
        oversized_window.insert(
            AttributeId::new("sliding_window_tokens").unwrap(),
            SemanticValue::Unsigned(262_145),
        );
        assert!(CausalAttentionShape::from_attributes_for(
            &oversized_window,
            CausalAttentionSemantics::Gemma4,
        )
        .is_err());
    }

    #[test]
    fn gemma4_cuda_source_carries_frequency_vnorm_window_and_512_fallback_contracts() {
        let source = include_str!("../../../../../kernels/vnext_causal_attention.cu");
        assert!(source.contains("#define VNEXT_MAX_HEAD_CHUNKS 16"));
        assert!(source.contains("rope_frequency_denominator"));
        assert!(source.contains("value_rms_norm"));
        assert!(source.contains("absolute_position - sliding_window + 1"));
        assert!(source.contains("const float attention_scale"));
        assert!(source.contains("const int neox_half = head_dim / 2"));
        assert!(source.contains("const int high = pair + neox_half"));
        assert!(source.contains("dim < neox_half + half_rope"));
    }
}