//! Phase 50 - GpuModel Coverage Tests Part 02 (gpu/scheduler/model.rs)
//!
//! Additional tests for improving coverage on model.rs (~24% uncovered):
//! - from_apr_weights constructor
//! - generate_gpu wrapper
//! - Edge cases for GQA attention paths
//! - SwiGLU FFN activation paths
//! - Error handling edge cases
//! - GpuModelConfig Debug/Clone traits
//! - Large vocab CPU fallback paths

use crate::gpu::executor::MockExecutor;
use crate::gpu::scheduler::{
    AttentionBuffers, BlockWeights, GpuGenerateConfig, GpuModel, GpuModelConfig, WeightType,
};
use crate::gpu::StreamingKVCache;

// ============================================================================
// Test Helpers
// ============================================================================

/// Create a minimal test model configuration
fn create_minimal_config() -> GpuModelConfig {
    GpuModelConfig {
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 1,
        num_heads: 2,
        num_kv_heads: 2,
        vocab_size: 50,
        eps: 1e-5,
        rope_theta: 10000.0,
    }
}

/// Create a GQA test configuration (num_heads > num_kv_heads)
fn create_gqa_config() -> GpuModelConfig {
    GpuModelConfig {
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 2,
        num_heads: 8,
        num_kv_heads: 2, // GQA: 4 Q heads per KV head
        vocab_size: 100,
        eps: 1e-5,
        rope_theta: 10000.0,
    }
}

/// Create config with large vocab for CPU fallback path testing
fn create_large_vocab_config() -> GpuModelConfig {
    GpuModelConfig {
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 32000, // Large vocab triggers CPU fallback
        eps: 1e-6,
        rope_theta: 10000.0,
    }
}

/// Create BlockWeights with SwiGLU gate (for SwiGLU FFN path)
fn create_block_weights_with_swiglu(config: &GpuModelConfig) -> BlockWeights {
    let hidden_dim = config.hidden_dim;
    let intermediate_dim = config.intermediate_dim;
    let qkv_dim = config.qkv_dim();

    BlockWeights {
        attn_norm_weight: vec![1.0; hidden_dim],
        attn_norm_bias: vec![0.0; hidden_dim],
        qkv_weight: vec![0.01; hidden_dim * qkv_dim],
        qkv_bias: vec![0.0; qkv_dim],
        out_weight: vec![0.01; hidden_dim * hidden_dim],
        out_bias: vec![0.0; hidden_dim],
        ffn_norm_weight: vec![1.0; hidden_dim],
        ffn_norm_bias: vec![0.0; hidden_dim],
        ffn_fc1_weight: vec![0.01; hidden_dim * intermediate_dim],
        ffn_fc1_bias: vec![0.0; intermediate_dim],
        ffn_fc2_weight: vec![0.01; intermediate_dim * hidden_dim],
        ffn_fc2_bias: vec![0.0; hidden_dim],
        ffn_gate_weight: Some(vec![0.01; hidden_dim * intermediate_dim]), // Enable SwiGLU
    }
}

// ============================================================================
// from_apr_weights Tests
// ============================================================================

#[test]
fn test_gpu_model_from_apr_weights_basic() {
    let config = create_minimal_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;
    let intermediate_dim = config.intermediate_dim;
    let qkv_dim = config.qkv_dim();

    // Create minimal weight vectors
    let embedding_weights = vec![0.01f32; vocab_size * hidden_dim];
    let block_weights = vec![BlockWeights {
        attn_norm_weight: vec![1.0; hidden_dim],
        attn_norm_bias: vec![0.0; hidden_dim],
        qkv_weight: vec![0.01; hidden_dim * qkv_dim],
        qkv_bias: vec![0.0; qkv_dim],
        out_weight: vec![0.01; hidden_dim * hidden_dim],
        out_bias: vec![0.0; hidden_dim],
        ffn_norm_weight: vec![1.0; hidden_dim],
        ffn_norm_bias: vec![0.0; hidden_dim],
        ffn_fc1_weight: vec![0.01; hidden_dim * intermediate_dim],
        ffn_fc1_bias: vec![0.0; intermediate_dim],
        ffn_fc2_weight: vec![0.01; intermediate_dim * hidden_dim],
        ffn_fc2_bias: vec![0.0; hidden_dim],
        ffn_gate_weight: None,
    }];
    let final_norm_weight = vec![1.0; hidden_dim];
    let final_norm_bias = vec![0.0; hidden_dim];
    let lm_head_weight = vec![0.01; hidden_dim * vocab_size];
    let lm_head_weight_t = vec![0.01; vocab_size * hidden_dim];
    let lm_head_bias = vec![0.0; vocab_size];

    let result = GpuModel::from_apr_weights(
        config.clone(),
        embedding_weights,
        block_weights,
        final_norm_weight,
        final_norm_bias,
        lm_head_weight,
        lm_head_weight_t,
        lm_head_bias,
    );

    assert!(result.is_ok(), "from_apr_weights should succeed");
    let model = result.unwrap();
    assert_eq!(model.config().hidden_dim, hidden_dim);
    assert_eq!(model.config().vocab_size, vocab_size);
}

#[test]
fn test_gpu_model_from_apr_weights_with_swiglu() {
    let config = create_minimal_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;
    let intermediate_dim = config.intermediate_dim;

    let embedding_weights = vec![0.01f32; vocab_size * hidden_dim];
    let block_weights = vec![create_block_weights_with_swiglu(&config)];
    let final_norm_weight = vec![1.0; hidden_dim];
    let final_norm_bias = vec![0.0; hidden_dim];
    let lm_head_weight = vec![0.01; hidden_dim * vocab_size];
    let lm_head_weight_t = vec![0.01; vocab_size * hidden_dim];
    let lm_head_bias = vec![0.0; vocab_size];

    let result = GpuModel::from_apr_weights(
        config,
        embedding_weights,
        block_weights,
        final_norm_weight,
        final_norm_bias,
        lm_head_weight,
        lm_head_weight_t,
        lm_head_bias,
    );

    assert!(result.is_ok());
    let model = result.unwrap();
    // Verify the model has gate weights
    assert!(model.block_weights[0].ffn_gate_weight.is_some());
    assert_eq!(
        model.block_weights[0].ffn_gate_weight.as_ref().unwrap().len(),
        hidden_dim * intermediate_dim
    );
}

#[test]
fn test_gpu_model_from_apr_weights_gqa() {
    let config = create_gqa_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;
    let intermediate_dim = config.intermediate_dim;
    let qkv_dim = config.qkv_dim();

    let embedding_weights = vec![0.01f32; vocab_size * hidden_dim];
    let mut block_weights = Vec::new();
    for _ in 0..config.num_layers {
        block_weights.push(BlockWeights {
            attn_norm_weight: vec![1.0; hidden_dim],
            attn_norm_bias: vec![0.0; hidden_dim],
            qkv_weight: vec![0.01; hidden_dim * qkv_dim],
            qkv_bias: vec![0.0; qkv_dim],
            out_weight: vec![0.01; hidden_dim * hidden_dim],
            out_bias: vec![0.0; hidden_dim],
            ffn_norm_weight: vec![1.0; hidden_dim],
            ffn_norm_bias: vec![0.0; hidden_dim],
            ffn_fc1_weight: vec![0.01; hidden_dim * intermediate_dim],
            ffn_fc1_bias: vec![0.0; intermediate_dim],
            ffn_fc2_weight: vec![0.01; intermediate_dim * hidden_dim],
            ffn_fc2_bias: vec![0.0; hidden_dim],
            ffn_gate_weight: None,
        });
    }
    let final_norm_weight = vec![1.0; hidden_dim];
    let final_norm_bias = vec![0.0; hidden_dim];
    let lm_head_weight = vec![0.01; hidden_dim * vocab_size];
    let lm_head_weight_t = vec![0.01; vocab_size * hidden_dim];
    let lm_head_bias = vec![0.0; vocab_size];

    let result = GpuModel::from_apr_weights(
        config.clone(),
        embedding_weights,
        block_weights,
        final_norm_weight,
        final_norm_bias,
        lm_head_weight,
        lm_head_weight_t,
        lm_head_bias,
    );

    assert!(result.is_ok());
    let model = result.unwrap();
    assert!(model.config().is_gqa());
    assert_eq!(model.config().num_kv_heads, 2);
}

// ============================================================================
// generate_gpu Tests
// ============================================================================

#[test]
fn test_gpu_model_generate_gpu_basic() {
    let config = create_minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    let mock = MockExecutor::new("generate_gpu_test");
    model.with_test_executor(Box::new(mock));

    let prompt = vec![1, 2];
    let max_tokens = 3;
    let result = model.generate_gpu(&prompt, max_tokens);

    assert!(result.is_ok());
    let tokens = result.unwrap();
    // Should include prompt + generated tokens
    assert!(tokens.len() >= prompt.len());
}

#[test]
fn test_gpu_model_generate_gpu_single_token() {
    let config = create_minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    let mock = MockExecutor::new("generate_gpu_single");
    model.with_test_executor(Box::new(mock));

    let prompt = vec![5];
    let max_tokens = 2;
    let result = model.generate_gpu(&prompt, max_tokens);

    assert!(result.is_ok());
}

#[test]
fn test_gpu_model_generate_gpu_max_tokens_zero() {
    let config = create_minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    let mock = MockExecutor::new("generate_gpu_zero");
    model.with_test_executor(Box::new(mock));

    let prompt = vec![1];
    let max_tokens = 0;
    let result = model.generate_gpu(&prompt, max_tokens);

    assert!(result.is_ok());
    // Should return at least the prompt
    let tokens = result.unwrap();
    assert!(!tokens.is_empty());
}

// ============================================================================
// SwiGLU FFN Path Tests
// ============================================================================

#[test]
fn test_gpu_model_forward_with_swiglu() {
    let config = create_minimal_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;
    let intermediate_dim = config.intermediate_dim;

    let embedding_weights = vec![0.01f32; vocab_size * hidden_dim];
    let block_weights = vec![create_block_weights_with_swiglu(&config)];
    let final_norm_weight = vec![1.0; hidden_dim];
    let final_norm_bias = vec![0.0; hidden_dim];
    let lm_head_weight = vec![0.01; hidden_dim * vocab_size];
    let lm_head_weight_t = vec![0.01; vocab_size * hidden_dim];
    let lm_head_bias = vec![0.0; vocab_size];

    let mut model = GpuModel::from_apr_weights(
        config.clone(),
        embedding_weights,
        block_weights,
        final_norm_weight,
        final_norm_bias,
        lm_head_weight,
        lm_head_weight_t,
        lm_head_bias,
    )
    .unwrap();

    let mock = MockExecutor::new("swiglu_forward");
    model.with_test_executor(Box::new(mock));

    let result = model.forward_gpu(&[1, 2]);
    assert!(result.is_ok());
    let logits = result.unwrap();
    assert_eq!(logits.len(), 2 * vocab_size);
}

#[test]
fn test_gpu_model_incremental_with_swiglu() {
    let config = create_minimal_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;

    let embedding_weights = vec![0.01f32; vocab_size * hidden_dim];
    let block_weights = vec![create_block_weights_with_swiglu(&config)];
    let final_norm_weight = vec![1.0; hidden_dim];
    let final_norm_bias = vec![0.0; hidden_dim];
    let lm_head_weight = vec![0.01; hidden_dim * vocab_size];
    let lm_head_weight_t = vec![0.01; vocab_size * hidden_dim];
    let lm_head_bias = vec![0.0; vocab_size];

    let mut model = GpuModel::from_apr_weights(
        config.clone(),
        embedding_weights,
        block_weights,
        final_norm_weight,
        final_norm_bias,
        lm_head_weight,
        lm_head_weight_t,
        lm_head_bias,
    )
    .unwrap();

    let mock = MockExecutor::new("swiglu_incremental");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    let result = model.forward_gpu_incremental_optimized(1, &mut kv_cache);
    assert!(result.is_ok());
}

// ============================================================================
// GpuModelConfig Debug/Clone Tests
// ============================================================================

#[test]
fn test_gpu_model_config_debug() {
    let config = create_minimal_config();
    let debug_str = format!("{:?}", config);

    assert!(debug_str.contains("GpuModelConfig"));
    assert!(debug_str.contains("hidden_dim"));
    assert!(debug_str.contains("vocab_size"));
    assert!(debug_str.contains("num_heads"));
}

#[test]
fn test_gpu_model_config_clone() {
    let config = create_minimal_config();
    let cloned = config.clone();

    assert_eq!(config.hidden_dim, cloned.hidden_dim);
    assert_eq!(config.vocab_size, cloned.vocab_size);
    assert_eq!(config.num_layers, cloned.num_layers);
    assert_eq!(config.eps, cloned.eps);
    assert_eq!(config.rope_theta, cloned.rope_theta);
}

#[test]
fn test_gpu_generate_config_debug() {
    let config = GpuGenerateConfig::with_sampling(64, 0.8, 40);
    let debug_str = format!("{:?}", config);

    assert!(debug_str.contains("GpuGenerateConfig"));
    assert!(debug_str.contains("max_tokens"));
    assert!(debug_str.contains("temperature"));
}

#[test]
fn test_gpu_generate_config_clone() {
    let config = GpuGenerateConfig::with_sampling(64, 0.8, 40).with_stop_tokens(vec![1, 2]);
    let cloned = config.clone();

    assert_eq!(config.max_tokens, cloned.max_tokens);
    assert_eq!(config.temperature, cloned.temperature);
    assert_eq!(config.top_k, cloned.top_k);
    assert_eq!(config.stop_tokens, cloned.stop_tokens);
}

// ============================================================================
// AttentionBuffers Debug Tests
// ============================================================================

#[test]
fn test_attention_buffers_debug() {
    let config = create_minimal_config();
    let buffers = AttentionBuffers::new(&config, 64);
    let debug_str = format!("{:?}", buffers);

    assert!(debug_str.contains("AttentionBuffers"));
    assert!(debug_str.contains("max_seq_len"));
}

// ============================================================================
// Large Vocab CPU Fallback Tests
// ============================================================================

#[test]
fn test_gpu_model_large_vocab_forward() {
    let config = create_large_vocab_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("large_vocab_forward");
    model.with_test_executor(Box::new(mock));

    // This should hit the CPU fallback path for LM head
    let result = model.forward_gpu(&[1]);
    assert!(result.is_ok());
    let logits = result.unwrap();
    assert_eq!(logits.len(), config.vocab_size);
}

#[test]
fn test_gpu_model_large_vocab_incremental() {
    let config = create_large_vocab_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("large_vocab_incremental");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    // This should hit the CPU fallback path for LM head in incremental
    let result = model.forward_gpu_incremental_optimized(1, &mut kv_cache);
    assert!(result.is_ok());
    let logits = result.unwrap();
    assert_eq!(logits.len(), config.vocab_size);
}

// ============================================================================
// Multi-layer Forward Tests
// ============================================================================

#[test]
fn test_gpu_model_multi_layer_forward() {
    let config = GpuModelConfig {
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 4, // Multiple layers
        num_heads: 2,
        num_kv_heads: 2,
        vocab_size: 50,
        eps: 1e-5,
        rope_theta: 10000.0,
    };
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("multi_layer_forward");
    model.with_test_executor(Box::new(mock));

    let result = model.forward_gpu(&[1, 2, 3]);
    assert!(result.is_ok());
    let logits = result.unwrap();
    assert_eq!(logits.len(), 3 * config.vocab_size);
}

#[test]
fn test_gpu_model_all_blocks_forward_idx() {
    let config = GpuModelConfig {
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 3,
        num_heads: 2,
        num_kv_heads: 2,
        vocab_size: 50,
        eps: 1e-5,
        rope_theta: 10000.0,
    };
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("all_blocks");
    model.with_test_executor(Box::new(mock));

    let mut hidden = vec![0.1f32; config.hidden_dim];
    for block_idx in 0..config.num_layers {
        let result = model.forward_block_idx(&hidden, 1, block_idx);
        assert!(result.is_ok(), "Block {} should succeed", block_idx);
        hidden = result.unwrap();
        assert_eq!(hidden.len(), config.hidden_dim);
    }
}

// ============================================================================
// GQA Edge Cases Tests
// ============================================================================

#[test]
fn test_gpu_model_gqa_multiple_q_per_kv() {
    // Test GQA where multiple Q heads share one KV head
    let config = GpuModelConfig {
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 16,  // 16 Q heads
        num_kv_heads: 4, // 4 KV heads -> 4 Q heads per KV head
        vocab_size: 50,
        eps: 1e-5,
        rope_theta: 10000.0,
    };
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("gqa_multiple");
    model.with_test_executor(Box::new(mock));

    let result = model.forward_gpu(&[1]);
    assert!(result.is_ok());
}

#[test]
fn test_gpu_model_gqa_incremental_with_cache() {
    let config = create_gqa_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("gqa_cache");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        128,
        config.num_kv_heads,
        config.head_dim(),
    );

    // Process multiple tokens to populate cache
    let _ = model.forward_gpu_with_cache(&[1, 2, 3], &mut kv_cache);

    // Do incremental forward with populated cache
    let result = model.forward_gpu_incremental_optimized(4, &mut kv_cache);
    assert!(result.is_ok());
}

// ============================================================================
// Error Edge Cases Tests
// ============================================================================

#[test]
fn test_gpu_model_empty_generate() {
    let config = create_minimal_config();
    let mut model = GpuModel::new(config).unwrap();

    let gen_config = GpuGenerateConfig::deterministic(5);
    let result = model.generate(&[], &gen_config);

    assert!(result.is_err(), "Generate with empty prompt should fail");
}

#[test]
fn test_gpu_model_token_at_vocab_boundary() {
    let config = create_minimal_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("boundary");
    model.with_test_executor(Box::new(mock));

    // Test token at vocab boundary (vocab_size - 1)
    let result = model.forward_gpu(&[config.vocab_size - 1]);
    assert!(result.is_ok());

    // Test token exactly at vocab_size (should fail)
    let result = model.forward_gpu(&[config.vocab_size]);
    assert!(result.is_err());
}

// ============================================================================
// WeightType All Variants Tests
// ============================================================================

#[test]
fn test_weight_type_copy_semantics() {
    let types = [
        WeightType::Qkv,
        WeightType::Output,
        WeightType::FfnFc1,
        WeightType::FfnFc2,
        WeightType::LmHead,
    ];

    for wt in types {
        let copied = wt; // Copy
        assert!(matches!(
            (wt, copied),
            (WeightType::Qkv, WeightType::Qkv)
                | (WeightType::Output, WeightType::Output)
                | (WeightType::FfnFc1, WeightType::FfnFc1)
                | (WeightType::FfnFc2, WeightType::FfnFc2)
                | (WeightType::LmHead, WeightType::LmHead)
        ));
    }
}

// ============================================================================
// BlockWeights Field Access Tests
// ============================================================================

#[test]
fn test_block_weights_all_fields() {
    let config = create_minimal_config();
    let block = create_block_weights_with_swiglu(&config);

    // Verify all fields can be accessed
    assert!(!block.attn_norm_weight.is_empty());
    assert!(!block.attn_norm_bias.is_empty());
    assert!(!block.qkv_weight.is_empty());
    assert!(!block.qkv_bias.is_empty());
    assert!(!block.out_weight.is_empty());
    assert!(!block.out_bias.is_empty());
    assert!(!block.ffn_norm_weight.is_empty());
    assert!(!block.ffn_norm_bias.is_empty());
    assert!(!block.ffn_fc1_weight.is_empty());
    assert!(!block.ffn_fc1_bias.is_empty());
    assert!(!block.ffn_fc2_weight.is_empty());
    assert!(!block.ffn_fc2_bias.is_empty());
    assert!(block.ffn_gate_weight.is_some());
}

// ============================================================================
// RoPE Theta Variation Tests
// ============================================================================

#[test]
fn test_gpu_model_custom_rope_theta() {
    // Test with non-default rope_theta
    let config = GpuModelConfig {
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 1,
        num_heads: 2,
        num_kv_heads: 2,
        vocab_size: 50,
        eps: 1e-5,
        rope_theta: 500000.0, // Different rope_theta (like Llama 3)
    };
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("custom_rope");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    let result = model.forward_gpu_incremental_optimized(1, &mut kv_cache);
    assert!(result.is_ok());
}

// ============================================================================
// Epsilon Variation Tests
// ============================================================================

#[test]
fn test_gpu_model_different_eps() {
    let configs = [1e-5, 1e-6, 1e-8];

    for eps in configs {
        let config = GpuModelConfig {
            hidden_dim: 32,
            intermediate_dim: 64,
            num_layers: 1,
            num_heads: 2,
            num_kv_heads: 2,
            vocab_size: 50,
            eps,
            rope_theta: 10000.0,
        };
        let mut model = GpuModel::new(config.clone()).unwrap();

        let mock = MockExecutor::new(&format!("eps_{}", eps));
        model.with_test_executor(Box::new(mock));

        let result = model.forward_gpu(&[1]);
        assert!(result.is_ok(), "Forward with eps={} should succeed", eps);
    }
}

// ============================================================================
// Sequential Forward Tests (Cache State)
// ============================================================================

#[test]
fn test_gpu_model_sequential_incremental_forwards() {
    let config = create_minimal_config();
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("sequential");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        128,
        config.num_kv_heads,
        config.head_dim(),
    );

    // Process multiple incremental forwards
    for token_id in 0..10 {
        let token = token_id % config.vocab_size;
        let result = model.forward_gpu_incremental_optimized(token, &mut kv_cache);
        assert!(result.is_ok(), "Token {} should succeed", token_id);
    }
}

// ============================================================================
// Forward Block Incremental Optimized with All Layers
// ============================================================================

#[test]
fn test_gpu_model_block_incremental_all_layers() {
    let config = GpuModelConfig {
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 3,
        num_heads: 2,
        num_kv_heads: 2,
        vocab_size: 50,
        eps: 1e-5,
        rope_theta: 10000.0,
    };
    let mut model = GpuModel::new(config.clone()).unwrap();

    let mock = MockExecutor::new("block_all_layers");
    model.with_test_executor(Box::new(mock));

    let mut kv_cache = StreamingKVCache::new(
        config.num_layers,
        64,
        config.num_kv_heads,
        config.head_dim(),
    );

    let mut hidden = vec![0.1f32; config.hidden_dim];
    for block_idx in 0..config.num_layers {
        let result = model.forward_block_incremental_optimized(&hidden, block_idx, &mut kv_cache);
        assert!(result.is_ok(), "Block {} incremental should succeed", block_idx);
        hidden = result.unwrap();
    }
}
