## Coverage Report

#### Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /home/noah/src/realizar/src/brick/mod.rs

    Line
    Count
    Source
    1
    //! ComputeBrick: Token-centric, self-verifying compute units.
    2
    //!
    3
    //! Per CBTOP-SPEC-001 and SHOWCASE-BRICK-001, every compute operation is a brick with:
    4
    //! - Token budget (tok/sec performance target)
    5
    //! - Assertions (falsifiable correctness claims)
    6
    //! - Verification (self-checking via baseline comparison)
    7
    //!
    8
    //! # Toyota Way Integration
    9
    //! - **Jidoka**: Stop-the-line on budget violation
    10
    //! - **Poka-Yoke**: Type-safe brick composition
    11
    //! - **Genchi Genbutsu**: Real metrics from hardware
    12
    //! - **Mieruka**: Visual control via cbtop TUI
    13
    //!
    14
    //! # References
    15
    //! - Popper, K. (1959). "The Logic of Scientific Discovery."
    16
    //! - Ohno, T. (1988). "Toyota Production System."
    17
    //! - Little, J.D.C. (1961). "A Proof for the Queuing Formula: L = λW."
    18
    19
    use std::fmt;
    20
    use std::time::Instant;
    21
    22
    use crate::quantize::Q8_0Block;
    23
    24
    // PMAT-802: Extracted modules
    25
    #[cfg(feature = "cuda")]
    26
    mod fused;
    27
    #[cfg(feature = "cuda")]
    28
    pub use fused::{CoalescedDp4aBrick, FusedFfnBrick};
    29
    30
    // Phase 14: BrickTracer for GPU/CPU parity debugging
    31
    pub mod tracer;
    32
    pub use tracer::{BrickTracer, TraceEvent, TraceComparison, TraceDiff};
    33
    34
    // ============================================================================
    35
    // Core Types
    36
    // ============================================================================
    37
    38
    /// Performance budget expressed in token terms.
    39
    /// Aligns compute costs with LLM inference metrics.
    40
    ///
    41
    /// Per Little's Law (1961): throughput = 1 / latency
    42
    #[derive(Debug, Clone, Copy, PartialEq)]
    43
    pub struct TokenBudget {
    44
        /// Latency budget per token (microseconds)
    45
        pub us_per_token: f64,
    46
        /// Throughput target (tokens/second)
    47
        pub tokens_per_sec: f64,
    48
        /// Batch size for amortization
    49
        pub batch_size: usize,
    50
    }
    51
    52
    impl TokenBudget {
    53
        /// Create budget from latency target.
    54
        /// 50µs/token = 20,000 tokens/sec
    55
        #[must_use]
    56
    385
        pub fn from_latency(us_per_token: f64) -> Self {
    57
    385
            Self {
    58
    385
                us_per_token,
    59
    385
                tokens_per_sec: 1_000_000.0 / us_per_token,
    60
    385
                batch_size: 1,
    61
    385
            }
    62
    385
        }
    63
    64
        /// Create budget from throughput target.
    65
        /// 20,000 tokens/sec = 50µs/token
    66
        #[must_use]
    67
    4
        pub fn from_throughput(tokens_per_sec: f64) -> Self {
    68
    4
            Self {
    69
    4
                us_per_token: 1_000_000.0 / tokens_per_sec,
    70
    4
                tokens_per_sec,
    71
    4
                batch_size: 1,
    72
    4
            }
    73
    4
        }
    74
    75
        /// Create budget with batch size.
    76
        #[must_use]
    77
    1
        pub fn with_batch_size(mut self, batch_size: usize) -> Self {
    78
    1
            self.batch_size = batch_size;
    79
    1
            self
    80
    1
        }
    81
    82
        /// Check if actual performance meets budget.
    83
        #[must_use]
    84
    65
        pub fn is_met(&self, actual_us_per_token: f64) -> bool {
    85
    65
            actual_us_per_token <= self.us_per_token
    86
    65
        }
    87
    88
        /// Calculate gap factor (actual / budget).
    89
        /// < 1.0 means under budget (good), > 1.0 means over budget (bad).
    90
        #[must_use]
    91
    9
        pub fn gap_factor(&self, actual_us_per_token: f64) -> f64 {
    92
    9
            actual_us_per_token / self.us_per_token
    93
    9
        }
    94
    }
    95
    96
    impl Default for TokenBudget {
    97
    1
        fn default() -> Self {
    98
    1
            Self::from_latency(100.0) // 100µs = 10k tok/s default
    99
    1
        }
    100
    }
    101
    102
    /// Result of ComputeBrick execution with token metrics.
    103
    #[derive(Debug, Clone)]
    104
    pub struct TokenResult<T> {
    105
        /// Computed output
    106
        pub output: T,
    107
        /// Number of tokens processed
    108
        pub tokens_processed: usize,
    109
        /// Actual latency (microseconds/token)
    110
        pub us_per_token: f64,
    111
        /// Actual throughput (tokens/second)
    112
        pub tokens_per_sec: f64,
    113
        /// Did we meet the budget?
    114
        pub budget_met: bool,
    115
    }
    116
    117
    impl<T: Default> Default for TokenResult<T> {
    118
    0
        fn default() -> Self {
    119
    0
            Self {
    120
    0
                output: T::default(),
    121
    0
                tokens_processed: 0,
    122
    0
                us_per_token: 0.0,
    123
    0
                tokens_per_sec: 0.0,
    124
    0
                budget_met: true,
    125
    0
            }
    126
    0
        }
    127
    }
    128
    129
    impl<T> TokenResult<T> {
    130
        /// Create a new token result.
    131
    62
        pub fn new(output: T, tokens: usize, elapsed_us: f64, budget: &TokenBudget) -> Self {
    132
    62
            let us_per_token = elapsed_us / tokens.max(1) as f64;
    133
    62
            let tokens_per_sec = if us_per_token > 0.0 {
    134
    60
                1_000_000.0 / us_per_token
    135
            } else {
    136
    2
                0.0
    137
            };
    138
    139
    62
            Self {
    140
    62
                output,
    141
    62
                tokens_processed: tokens,
    142
    62
                us_per_token,
    143
    62
                tokens_per_sec,
    144
    62
                budget_met: budget.is_met(us_per_token),
    145
    62
            }
    146
    62
        }
    147
    }
    148
    149
    /// Errors from ComputeBrick execution.
    150
    /// Tells you exactly what failed (Jidoka: stop and signal).
    151
    #[derive(Debug)]
    152
    pub enum BrickError {
    153
        /// Assertion failed during verification.
    154
        AssertionFailed {
    155
            /// Assertion name
    156
            name: String,
    157
            /// Expected value
    158
            expected: String,
    159
            /// Actual value
    160
            actual: String,
    161
        },
    162
        /// Performance budget exceeded.
    163
        BudgetExceeded {
    164
            /// Budget limit in µs
    165
            limit_us: f64,
    166
            /// Actual time in µs
    167
            actual_us: f64,
    168
        },
    169
        /// Compute operation failed.
    170
        ComputeError(String),
    171
        /// Invalid input.
    172
        InvalidInput(String),
    173
    }
    174
    175
    impl fmt::Display for BrickError {
    176
    6
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    177
    6
            match self {
    178
                Self::AssertionFailed {
    179
    1
                    name,
    180
    1
                    expected,
    181
    1
                    actual,
    182
                } => {
    183
    1
                    write!(
    184
    1
                        f,
    185
    1
                        "Assertion failed: {name} - expected {expected}, got {actual}"
    186
                    )
    187
                },
    188
                Self::BudgetExceeded {
    189
    2
                    limit_us,
    190
    2
                    actual_us,
    191
                } => {
    192
    2
                    write!(
    193
    2
                        f,
    194
    2
                        "Budget exceeded: {limit_us:.1}µs/tok limit, {actual_us:.1}µs/tok actual"
    195
                    )
    196
                },
    197
    1
                Self::ComputeError(msg) => write!(f, "Compute error: {msg}"),
    198
    2
                Self::InvalidInput(msg) => write!(f, "Invalid input: {msg}"),
    199
            }
    200
    6
        }
    201
    }
    202
    203
    impl std::error::Error for BrickError {}
    204
    205
    // ============================================================================
    206
    // Brick Assertions
    207
    // ============================================================================
    208
    209
    /// Falsifiable assertion for a brick (Popper criterion).
    210
    #[derive(Debug, Clone)]
    211
    pub struct BrickAssertion {
    212
        /// Assertion name (for error messages)
    213
        pub name: String,
    214
        /// Description of what is being asserted
    215
        pub description: String,
    216
        /// Assertion type
    217
        pub kind: AssertionKind,
    218
    }
    219
    220
    /// Types of assertions a brick can make.
    221
    #[derive(Debug, Clone)]
    222
    pub enum AssertionKind {
    223
        /// Output matches scalar baseline within tolerance.
    224
        EquivScalar {
    225
            /// Tolerance for comparison
    226
            tolerance: f64,
    227
        },
    228
        /// Output contains no NaN values.
    229
        NoNaN,
    230
        /// Output contains no Inf values.
    231
        NoInf,
    232
        /// Output values within bounds.
    233
        Bounds {
    234
            /// Minimum allowed value
    235
            min: f64,
    236
            /// Maximum allowed value
    237
            max: f64,
    238
        },
    239
        /// Budget must be met.
    240
        BudgetMet,
    241
        /// Custom assertion (returns true if passed).
    242
        Custom {
    243
            /// Name of the custom check
    244
            check_name: String,
    245
        },
    246
    }
    247
    248
    impl BrickAssertion {
    249
        /// Create assertion that output matches scalar baseline.
    250
    2
        pub fn equiv_scalar(tolerance: f64) -> Self {
    251
    2
            Self {
    252
    2
                name: "equiv_scalar".to_string(),
    253
    2
                description: format!("Output matches scalar baseline within {tolerance}"),
    254
    2
                kind: AssertionKind::EquivScalar { tolerance },
    255
    2
            }
    256
    2
        }
    257
    258
        /// Create assertion that output contains no NaN.
    259
    73
        pub fn no_nan() -> Self {
    260
    73
            Self {
    261
    73
                name: "no_nan".to_string(),
    262
    73
                description: "Output contains no NaN values".to_string(),
    263
    73
                kind: AssertionKind::NoNaN,
    264
    73
            }
    265
    73
        }
    266
    267
        /// Create assertion that output contains no Inf.
    268
    72
        pub fn no_inf() -> Self {
    269
    72
            Self {
    270
    72
                name: "no_inf".to_string(),
    271
    72
                description: "Output contains no Inf values".to_string(),
    272
    72
                kind: AssertionKind::NoInf,
    273
    72
            }
    274
    72
        }
    275
    276
        /// Create assertion that values are within bounds.
    277
    5
        pub fn bounds(min: f64, max: f64) -> Self {
    278
    5
            Self {
    279
    5
                name: "bounds".to_string(),
    280
    5
                description: format!("Output values in [{min}, {max}]"),
    281
    5
                kind: AssertionKind::Bounds { min, max },
    282
    5
            }
    283
    5
        }
    284
    285
        /// Create assertion that budget is met.
    286
    73
        pub fn budget_met() -> Self {
    287
    73
            Self {
    288
    73
                name: "budget_met".to_string(),
    289
    73
                description: "Performance budget is met".to_string(),
    290
    73
                kind: AssertionKind::BudgetMet,
    291
    73
            }
    292
    73
        }
    293
    294
        /// Check assertion against f32 slice output.
    295
    183
        pub fn check_f32(&self, output: &[f32], budget_met: bool) -> Result<(), BrickError> {
    296
    183
            match &self.kind {
    297
                AssertionKind::NoNaN => {
    298
    72.7k
                    if let Some(idx1) = output.iter()61.position61(|x| x.is_nan()) {
    299
    1
                        return Err(BrickError::AssertionFailed {
    300
    1
                            name: self.name.clone(),
    301
    1
                            expected: "no NaN".to_string(),
    302
    1
                            actual: format!("NaN at index {idx}"),
    303
    1
                        });
    304
    60
                    }
    305
                },
    306
                AssertionKind::NoInf => {
    307
    72.7k
                    if let Some(idx0) = output.iter()59.position59(|x| x.is_infinite()) {
    308
    0
                        return Err(BrickError::AssertionFailed {
    309
    0
                            name: self.name.clone(),
    310
    0
                            expected: "no Inf".to_string(),
    311
    0
                            actual: format!("Inf at index {idx}"),
    312
    0
                        });
    313
    59
                    }
    314
                },
    315
    3
                AssertionKind::Bounds { min, max } => {
    316
    9
                    for (idx, &val) in output3.iter3().enumerate3() {
    317
    9
                        if (val as f64) < *min || (val as f64) > *max {
    318
    1
                            return Err(BrickError::AssertionFailed {
    319
    1
                                name: self.name.clone(),
    320
    1
                                expected: format!("value in [{min}, {max}]"),
    321
    1
                                actual: format!("value {val} at index {idx}"),
    322
    1
                            });
    323
    8
                        }
    324
                    }
    325
                },
    326
                AssertionKind::BudgetMet => {
    327
    58
                    if !budget_met {
    328
    0
                        return Err(BrickError::AssertionFailed {
    329
    0
                            name: self.name.clone(),
    330
    0
                            expected: "budget met".to_string(),
    331
    0
                            actual: "budget exceeded".to_string(),
    332
    0
                        });
    333
    58
                    }
    334
                },
    335
    2
                AssertionKind::EquivScalar { .. } | AssertionKind::Custom { .. } => {
    336
    2
                    // These require external comparison, skip for basic check
    337
    2
                },
    338
            }
    339
    181
            Ok(())
    340
    183
        }
    341
    }
    342
    343
    // ============================================================================
    344
    // Brick Verification
    345
    // ============================================================================
    346
    347
    /// Result of brick verification.
    348
    #[derive(Debug, Clone)]
    349
    pub struct BrickVerification {
    350
        /// All assertions passed?
    351
        pub is_valid: bool,
    352
        /// Individual assertion results
    353
        pub results: Vec<(String, bool, String)>,
    354
    }
    355
    356
    impl BrickVerification {
    357
        /// Create a passing verification.
    358
    10
        pub fn pass() -> Self {
    359
    10
            Self {
    360
    10
                is_valid: true,
    361
    10
                results: vec![],
    362
    10
            }
    363
    10
        }
    364
    365
        /// Create a failing verification.
    366
    1
        pub fn fail(name: &str, reason: &str) -> Self {
    367
    1
            Self {
    368
    1
                is_valid: false,
    369
    1
                results: vec![(name.to_string(), false, reason.to_string())],
    370
    1
            }
    371
    1
        }
    372
    373
        /// Add an assertion result.
    374
    2
        pub fn add(&mut self, name: &str, passed: bool, message: &str) {
    375
    2
            self.results
    376
    2
                .push((name.to_string(), passed, message.to_string()));
    377
    2
            if !passed {
    378
    1
                self.is_valid = false;
    379
    1
            }
    380
    2
        }
    381
    }
    382
    383
    // ============================================================================
    384
    // ComputeBrick Trait
    385
    // ============================================================================
    386
    387
    /// Core trait for self-verifying, token-centric compute units.
    388
    ///
    389
    /// Every brick must:
    390
    /// 1. Have at least one assertion (Popper criterion)
    391
    /// 2. Have a non-zero budget (accountability)
    392
    /// 3. Be verifiable against baseline
    393
    ///
    394
    /// # Invariants (PROBAR-SPEC-009 §3)
    395
    /// - `assertions().len() > 0` (at least one falsifiable claim)
    396
    /// - `budget().us_per_token > 0` (performance accountability)
    397
    /// - `verify()` checks ALL assertions
    398
    pub trait ComputeBrick: Send + Sync {
    399
        /// Output type of this brick.
    400
        type Output;
    401
    402
        /// Brick name for identification.
    403
        fn name(&self) -> &'static str;
    404
    405
        /// Token throughput budget.
    406
        fn budget(&self) -> TokenBudget;
    407
    408
        /// Falsifiable assertions for this brick.
    409
        fn assertions(&self) -> Vec<BrickAssertion>;
    410
    411
        /// Verify brick state without running.
    412
    7
        fn verify(&self) -> BrickVerification {
    413
    7
            let assertions = self.assertions();
    414
    7
            if assertions.is_empty() {
    415
    0
                return BrickVerification::fail(
    416
    0
                    self.name(),
    417
    0
                    "No assertions defined (Popper violation)",
    418
                );
    419
    7
            }
    420
    421
    7
            let budget = self.budget();
    422
    7
            if budget.us_per_token <= 0.0 {
    423
    0
                return BrickVerification::fail(self.name(), "Zero or negative budget");
    424
    7
            }
    425
    426
    7
            BrickVerification::pass()
    427
    7
        }
    428
    429
        /// Can this brick run? (Jidoka gate)
    430
    5
        fn can_run(&self) -> bool {
    431
    5
            self.verify().is_valid
    432
    5
        }
    433
    }
    434
    435
    // ============================================================================
    436
    // Transformer Brick Implementations
    437
    // ============================================================================
    438
    439
    /// RMSNorm brick - layer normalization.
    440
    #[derive(Debug)]
    441
    pub struct RmsNormBrick {
    442
        /// Weight vector
    443
        pub weight: Vec<f32>,
    444
        /// Epsilon for numerical stability
    445
        pub eps: f32,
    446
        /// Budget
    447
        budget: TokenBudget,
    448
    }
    449
    450
    impl RmsNormBrick {
    451
        /// Create a new RMSNorm brick.
    452
    113
        pub fn new(weight: Vec<f32>, eps: f32) -> Self {
    453
    113
            Self {
    454
    113
                weight,
    455
    113
                eps,
    456
    113
                budget: TokenBudget::from_latency(1.5), // 1.5µs target
    457
    113
            }
    458
    113
        }
    459
    460
        /// Set custom budget.
    461
        #[must_use]
    462
    4
        pub fn with_budget(mut self, budget: TokenBudget) -> Self {
    463
    4
            self.budget = budget;
    464
    4
            self
    465
    4
        }
    466
    467
        /// Run RMSNorm on input.
    468
    58
        pub fn run(&self, input: &[f32]) -> Result<TokenResult<Vec<f32>>, BrickError> {
    469
    58
            if input.len() != self.weight.len() {
    470
    0
                return Err(BrickError::InvalidInput(format!(
    471
    0
                    "Input len {} != weight len {}",
    472
    0
                    input.len(),
    473
    0
                    self.weight.len()
    474
    0
                )));
    475
    58
            }
    476
    477
    58
            let start = Instant::now();
    478
    479
            // Compute RMS
    480
    72.7k
            let rms58 = (input58.iter58().map58(|x| x * x).sum58::<f32>() / input.len() as f3258 + self.eps58).sqrt58();
    481
    482
            // Normalize and scale
    483
    58
            let output: Vec<f32> = input
    484
    58
                .iter()
    485
    58
                .zip(self.weight.iter())
    486
    72.7k
                .map58(|(x, w)| (x / rms) * w)
    487
    58
                .collect();
    488
    489
    58
            let elapsed_us = start.elapsed().as_micros() as f64;
    490
    58
            let result = TokenResult::new(output, 1, elapsed_us, &self.budget);
    491
    492
            // Check assertions
    493
    174
            for assertion in self58.assertions58() {
    494
    174
                assertion.check_f32(&result.output, result.budget_met)?0;
    495
            }
    496
    497
    58
            Ok(result)
    498
    58
        }
    499
    }
    500
    501
    impl ComputeBrick for RmsNormBrick {
    502
        type Output = Vec<f32>;
    503
    504
    2
        fn name(&self) -> &'static str {
    505
    2
            "rms_norm"
    506
    2
        }
    507
    508
    16
        fn budget(&self) -> TokenBudget {
    509
    16
            self.budget
    510
    16
        }
    511
    512
    65
        fn assertions(&self) -> Vec<BrickAssertion> {
    513
    65
            vec![
    514
    65
                BrickAssertion::no_nan(),
    515
    65
                BrickAssertion::no_inf(),
    516
    65
                BrickAssertion::budget_met(),
    517
            ]
    518
    65
        }
    519
    }
    520
    521
    /// QKV projection brick.
    522
    #[derive(Debug)]
    523
    pub struct QkvBrick {
    524
        /// Hidden dimension
    525
        pub hidden_dim: usize,
    526
        /// Q output dimension
    527
        pub q_dim: usize,
    528
        /// K output dimension
    529
        pub k_dim: usize,
    530
        /// V output dimension
    531
        pub v_dim: usize,
    532
        /// Budget
    533
        budget: TokenBudget,
    534
        /// Has bias (Qwen2 has large biases)
    535
        pub has_bias: bool,
    536
    }
    537
    538
    impl QkvBrick {
    539
        /// Create a new QKV brick.
    540
    5
        pub fn new(hidden_dim: usize, q_dim: usize, k_dim: usize, v_dim: usize) -> Self {
    541
    5
            Self {
    542
    5
                hidden_dim,
    543
    5
                q_dim,
    544
    5
                k_dim,
    545
    5
                v_dim,
    546
    5
                budget: TokenBudget::from_latency(6.0), // 6µs target
    547
    5
                has_bias: false,
    548
    5
            }
    549
    5
        }
    550
    551
        /// Set custom budget.
    552
        #[must_use]
    553
    0
        pub fn with_budget(mut self, budget: TokenBudget) -> Self {
    554
    0
            self.budget = budget;
    555
    0
            self
    556
    0
        }
    557
    558
        /// Mark as having bias.
    559
        #[must_use]
    560
    1
        pub fn with_bias(mut self) -> Self {
    561
    1
            self.has_bias = true;
    562
    1
            self
    563
    1
        }
    564
    565
        /// Total output dimension.
    566
    1
        pub fn total_out_dim(&self) -> usize {
    567
    1
            self.q_dim + self.k_dim + self.v_dim
    568
    1
        }
    569
    }
    570
    571
    impl ComputeBrick for QkvBrick {
    572
        type Output = (Vec<f32>, Vec<f32>, Vec<f32>);
    573
    574
    0
        fn name(&self) -> &'static str {
    575
    0
            "qkv_proj"
    576
    0
        }
    577
    578
    4
        fn budget(&self) -> TokenBudget {
    579
    4
            self.budget
    580
    4
        }
    581
    582
    0
        fn assertions(&self) -> Vec<BrickAssertion> {
    583
    0
            vec![
    584
    0
                BrickAssertion::no_nan(),
    585
    0
                BrickAssertion::no_inf(),
    586
    0
                BrickAssertion::budget_met(),
    587
            ]
    588
    0
        }
    589
    }
    590
    591
    /// RoPE brick - rotary position embedding.
    592
    #[derive(Debug)]
    593
    pub struct RopeBrick {
    594
        /// Head dimension
    595
        pub head_dim: usize,
    596
        /// Number of heads
    597
        pub num_heads: usize,
    598
        /// Base theta
    599
        pub theta: f32,
    600
        /// RoPE type (0=NORM, 2=NEOX)
    601
        pub rope_type: u32,
    602
        /// Budget
    603
        budget: TokenBudget,
    604
    }
    605
    606
    impl RopeBrick {
    607
        /// Create a new RoPE brick.
    608
    3
        pub fn new(head_dim: usize, num_heads: usize, theta: f32, rope_type: u32) -> Self {
    609
    3
            Self {
    610
    3
                head_dim,
    611
    3
                num_heads,
    612
    3
                theta,
    613
    3
                rope_type,
    614
    3
                budget: TokenBudget::from_latency(1.0), // 1µs target
    615
    3
            }
    616
    3
        }
    617
    618
        /// Set custom budget.
    619
        #[must_use]
    620
    0
        pub fn with_budget(mut self, budget: TokenBudget) -> Self {
    621
    0
            self.budget = budget;
    622
    0
            self
    623
    0
        }
    624
    }
    625
    626
    impl ComputeBrick for RopeBrick {
    627
        type Output = Vec<f32>;
    628
    629
    0
        fn name(&self) -> &'static str {
    630
    0
            "rope"
    631
    0
        }
    632
    633
    4
        fn budget(&self) -> TokenBudget {
    634
    4
            self.budget
    635
    4
        }
    636
    637
    0
        fn assertions(&self) -> Vec<BrickAssertion> {
    638
    0
            vec![
    639
    0
                BrickAssertion::no_nan(),
    640
    0
                BrickAssertion::no_inf(),
    641
    0
                BrickAssertion::budget_met(),
    642
            ]
    643
    0
        }
    644
    }
    645
    646
    /// Attention brick.
    647
    #[derive(Debug)]
    648
    pub struct AttentionBrick {
    649
        /// Number of query heads
    650
        pub num_heads: usize,
    651
        /// Number of KV heads (for GQA)
    652
        pub num_kv_heads: usize,
    653
        /// Head dimension
    654
        pub head_dim: usize,
    655
        /// Budget
    656
        budget: TokenBudget,
    657
    }
    658
    659
    impl AttentionBrick {
    660
        /// Create a new attention brick.
    661
    106
        pub fn new(num_heads: usize, num_kv_heads: usize, head_dim: usize) -> Self {
    662
    106
            Self {
    663
    106
                num_heads,
    664
    106
                num_kv_heads,
    665
    106
                head_dim,
    666
    106
                budget: TokenBudget::from_latency(10.0), // 10µs target
    667
    106
            }
    668
    106
        }
    669
    670
        /// Set custom budget.
    671
        #[must_use]
    672
    0
        pub fn with_budget(mut self, budget: TokenBudget) -> Self {
    673
    0
            self.budget = budget;
    674
    0
            self
    675
    0
        }
    676
    677
        /// GQA group size.
    678
    0
        pub fn group_size(&self) -> usize {
    679
    0
            self.num_heads / self.num_kv_heads.max(1)
    680
    0
        }
    681
    }
    682
    683
    impl ComputeBrick for AttentionBrick {
    684
        type Output = Vec<f32>;
    685
    686
    0
        fn name(&self) -> &'static str {
    687
    0
            "attention"
    688
    0
        }
    689
    690
    7
        fn budget(&self) -> TokenBudget {
    691
    7
            self.budget
    692
    7
        }
    693
    694
    2
        fn assertions(&self) -> Vec<BrickAssertion> {
    695
    2
            vec![
    696
    2
                BrickAssertion::no_nan(),
    697
    2
                BrickAssertion::no_inf(),
    698
    2
                BrickAssertion::budget_met(),
    699
                // Attention outputs should be bounded
    700
    2
                BrickAssertion::bounds(-100.0, 100.0),
    701
            ]
    702
    2
        }
    703
    }
    704
    705
    /// FlashAttentionBrick - incremental flash attention for decode (P1 optimization).
    706
    ///
    707
    /// **Algorithm** (FlashAttention-2, Dao et al. 2023):
    708
    /// ```text
    709
    /// For decode (single query token):
    710
    ///   Q: [1, H, D]     (single query)
    711
    ///   K: [S, H_kv, D]  (full KV cache)
    712
    ///   V: [S, H_kv, D]  (full KV cache)
    713
    ///
    714
    ///   Online softmax (no full attention matrix materialization):
    715
    ///   for tile in KV_tiles(TILE_SIZE=128):
    716
    ///       S_tile = Q @ K_tile^T / sqrt(D)    # [1, H, TILE_SIZE]
    717
    ///       m_new = max(m_old, max(S_tile))    # Running max
    718
    ///       P_tile = exp(S_tile - m_new)       # Stable softmax numerator
    719
    ///       O = O * exp(m_old - m_new) + P_tile @ V_tile  # Accumulate
    720
    ///       l = l * exp(m_old - m_new) + sum(P_tile)      # Running denominator
    721
    ///   O = O / l  # Final output
    722
    /// ```
    723
    ///
    724
    /// **Performance vs naive**:
    725
    /// - Naive: O(S) memory for attention matrix
    726
    /// - Flash: O(TILE_SIZE) memory, 2x speedup from better cache locality
    727
    ///
    728
    /// **Reference**: Dao, T., et al. (2023). "FlashAttention-2: Faster Attention
    729
    /// with Better Parallelism and Work Partitioning." arXiv:2307.08691.
    730
    #[derive(Debug, Clone)]
    731
    pub struct FlashAttentionBrick {
    732
        /// Number of query heads
    733
        pub num_heads: usize,
    734
        /// Number of KV heads (for GQA)
    735
        pub num_kv_heads: usize,
    736
        /// Head dimension
    737
        pub head_dim: usize,
    738
        /// Tile size for KV cache (default: 128 for L2 cache fit)
    739
        pub tile_size: usize,
    740
        /// Budget (target: 5.0µs for 2x improvement over naive)
    741
        budget: TokenBudget,
    742
        /// Use online softmax (FlashAttention algorithm)
    743
        pub use_online_softmax: bool,
    744
    }
    745
    746
    impl FlashAttentionBrick {
    747
        /// Create new flash attention brick with default tile size.
    748
        #[must_use]
    749
    15
        pub fn new(num_heads: usize, num_kv_heads: usize, head_dim: usize) -> Self {
    750
    15
            Self {
    751
    15
                num_heads,
    752
    15
                num_kv_heads,
    753
    15
                head_dim,
    754
    15
                tile_size: 128, // Optimal for L2 cache on most GPUs
    755
    15
                budget: TokenBudget::from_latency(5.0), // 5µs target (2x vs 10µs naive)
    756
    15
                use_online_softmax: true,
    757
    15
            }
    758
    15
        }
    759
    760
        /// Create with custom tile size (for tuning).
    761
        #[must_use]
    762
    1
        pub fn with_tile_size(
    763
    1
            num_heads: usize,
    764
    1
            num_kv_heads: usize,
    765
    1
            head_dim: usize,
    766
    1
            tile_size: usize,
    767
    1
        ) -> Self {
    768
    1
            Self {
    769
    1
                num_heads,
    770
    1
                num_kv_heads,
    771
    1
                head_dim,
    772
    1
                tile_size,
    773
    1
                budget: TokenBudget::from_latency(5.0),
    774
    1
                use_online_softmax: true,
    775
    1
            }
    776
    1
        }
    777
    778
        /// Set custom budget.
    779
        #[must_use]
    780
    0
        pub fn with_budget(mut self, budget: TokenBudget) -> Self {
    781
    0
            self.budget = budget;
    782
    0
            self
    783
    0
        }
    784
    785
        /// GQA group size (query heads per KV head).
    786
        #[must_use]
    787
    5
        pub fn group_size(&self) -> usize {
    788
    5
            self.num_heads / self.num_kv_heads.max(1)
    789
    5
        }
    790
    791
        /// Compute FLOPs per token for attention at given sequence length.
    792
        ///
    793
        /// FLOPs = 2 * H * D * S (Q @ K^T) + 2 * H * S * D (attn @ V)
    794
        ///       = 4 * H * D * S
    795
        #[must_use]
    796
    2
        pub fn flops(&self, seq_len: usize) -> u64 {
    797
    2
            4 * self.num_heads as u64 * self.head_dim as u64 * seq_len as u64
    798
    2
        }
    799
    800
        /// Compute memory bandwidth for naive vs flash attention.
    801
        ///
    802
        /// Naive: Reads full KV cache for each head
    803
        /// Flash: Reads KV cache in tiles, better cache reuse
    804
        #[must_use]
    805
    4
        pub fn memory_bytes(&self, seq_len: usize) -> (u64, u64) {
    806
    4
            let kv_bytes = 2 * self.num_kv_heads as u64 * self.head_dim as u64 * seq_len as u64 * 4; // K + V, f32
    807
    4
            let naive = kv_bytes + self.num_heads as u64 * seq_len as u64 * 4; // + attention matrix
    808
    4
            let flash = kv_bytes; // No attention matrix materialized
    809
    4
            (naive, flash)
    810
    4
        }
    811
    812
        /// Compute arithmetic intensity (FLOPs / bytes).
    813
        #[must_use]
    814
    0
        pub fn arithmetic_intensity(&self, seq_len: usize) -> f64 {
    815
    0
            let (_, flash_bytes) = self.memory_bytes(seq_len);
    816
    0
            self.flops(seq_len) as f64 / flash_bytes as f64
    817
    0
        }
    818
    819
        /// Number of tiles needed for given sequence length.
    820
        #[must_use]
    821
    3
        pub fn num_tiles(&self, seq_len: usize) -> usize {
    822
    3
            seq_len.div_ceil(self.tile_size)
    823
    3
        }
    824
    825
        /// Compute flash attention with online softmax algorithm.
    826
        ///
    827
        /// **REAL IMPLEMENTATION** - FlashAttention-2 (Dao et al. 2023)
    828
        ///
    829
        /// # Arguments
    830
        /// * `query` - Query tensor [num_heads, head_dim]
    831
        /// * `keys` - Key cache [seq_len, num_kv_heads, head_dim]
    832
        /// * `values` - Value cache [seq_len, num_kv_heads, head_dim]
    833
        ///
    834
        /// # Returns
    835
        /// * Output tensor [num_heads, head_dim]
    836
    4
        pub fn forward(
    837
    4
            &self,
    838
    4
            query: &[f32],  // [num_heads * head_dim]
    839
    4
            keys: &[f32],   // [seq_len * num_kv_heads * head_dim]
    840
    4
            values: &[f32], // [seq_len * num_kv_heads * head_dim]
    841
    4
            seq_len: usize,
    842
    4
        ) -> Result<Vec<f32>, BrickError> {
    843
    4
            if self.num_heads == 0 || self.head_dim == 0 {
    844
    0
                return Err(BrickError::InvalidInput("Zero dimension".to_string()));
    845
    4
            }
    846
    4
            if query.len() != self.num_heads * self.head_dim {
    847
    0
                return Err(BrickError::InvalidInput(format!(
    848
    0
                    "Query length {} != num_heads * head_dim = {}",
    849
    0
                    query.len(),
    850
    0
                    self.num_heads * self.head_dim
    851
    0
                )));
    852
    4
            }
    853
    4
            let expected_kv_len = seq_len * self.num_kv_heads * self.head_dim;
    854
    4
            if keys.len() != expected_kv_len || values.len() != expected_kv_len {
    855
    0
                return Err(BrickError::InvalidInput(format!(
    856
    0
                    "KV length {} != seq_len * num_kv_heads * head_dim = {}",
    857
    0
                    keys.len(),
    858
    0
                    expected_kv_len
    859
    0
                )));
    860
    4
            }
    861
    862
    4
            let scale = 1.0 / (self.head_dim as f32).sqrt();
    863
    4
            let group_size = self.group_size();
    864
    4
            let mut output = vec![0.0f32; self.num_heads * self.head_dim];
    865
    866
            // Process each query head
    867
    14
            for h in 0..self.num_heads4 {
    868
    14
                let kv_head = h / group_size; // GQA: map query head to KV head
    869
    14
                let q_start = h * self.head_dim;
    870
    871
                // Online softmax variables (FlashAttention-2)
    872
    14
                let mut m = f32::NEG_INFINITY; // Running max
    873
    14
                let mut l = 0.0f32; // Running sum of exp
    874
    14
                let mut o = vec![0.0f32; self.head_dim]; // Running output
    875
    876
                // Process in tiles for cache efficiency
    877
    14
                for tile_start in (0..seq_len).step_by(self.tile_size) {
    878
    14
                    let tile_end = (tile_start + self.tile_size).min(seq_len);
    879
    880
    1.04k
                    for s in tile_start14..tile_end14 {
    881
                        // Compute Q @ K^T for this position
    882
    1.04k
                        let k_start = (s * self.num_kv_heads + kv_head) * self.head_dim;
    883
    1.04k
                        let mut score = 0.0f32;
    884
    65.6k
                        for d in 0..self.head_dim1.04k {
    885
    65.6k
                            score += query[q_start + d] * keys[k_start + d];
    886
    65.6k
                        }
    887
    1.04k
                        score *= scale;
    888
    889
                        // Online softmax update (Milakov & Gimelshein, 2018)
    890
    1.04k
                        let m_new = m.max(score);
    891
    1.04k
                        let exp_old = (m - m_new).exp();
    892
    1.04k
                        let exp_score = (score - m_new).exp();
    893
    894
                        // Update running sum
    895
    1.04k
                        l = l * exp_old + exp_score;
    896
    897
                        // Update running output: O = O * exp(m_old - m_new) + exp(score - m_new) * V
    898
    1.04k
                        let v_start = (s * self.num_kv_heads + kv_head) * self.head_dim;
    899
    65.6k
                        for d in 0..self.head_dim1.04k {
    900
    65.6k
                            o[d] = o[d] * exp_old + exp_score * values[v_start + d];
    901
    65.6k
                        }
    902
    903
    1.04k
                        m = m_new;
    904
                    }
    905
                }
    906
    907
                // Normalize output: O = O / l
    908
    14
                if l > 0.0 {
    909
    552
                    for d in 0..self.head_dim14 {
    910
    552
                        output[h * self.head_dim + d] = o[d] / l;
    911
    552
                    }
    912
    0
                }
    913
            }
    914
    915
    4
            Ok(output)
    916
    4
        }
    917
    918
        /// Execute flash attention with timing (for benchmarking).
    919
    1
        pub fn forward_timed(
    920
    1
            &self,
    921
    1
            query: &[f32],
    922
    1
            keys: &[f32],
    923
    1
            values: &[f32],
    924
    1
            seq_len: usize,
    925
    1
        ) -> Result<TokenResult<Vec<f32>>, BrickError> {
    926
    1
            let start = Instant::now();
    927
    1
            let output = self.forward(query, keys, values, seq_len)?0;
    928
    1
            let elapsed_us = start.elapsed().as_secs_f64() * 1_000_000.0;
    929
    930
    1
            Ok(TokenResult {
    931
    1
                output,
    932
    1
                tokens_processed: 1,
    933
    1
                us_per_token: elapsed_us,
    934
    1
                tokens_per_sec: 1_000_000.0 / elapsed_us,
    935
    1
                budget_met: elapsed_us <= self.budget.us_per_token,
    936
    1
            })
    937
    1
        }
    938
    939
        /// Legacy stub for backward compatibility (prefer `forward()`)
    940
        #[deprecated(note = "Use forward() for real implementation")]
    941
    1
        pub fn execute(&self, _seq_len: usize) -> Result<Vec<f32>, BrickError> {
    942
    1
            if self.num_heads == 0 || self.head_dim == 00 {
    943
    1
                return Err(BrickError::InvalidInput("Zero dimension".to_string()));
    944
    0
            }
    945
    0
            Ok(vec![0.0; self.num_heads * self.head_dim])
    946
    1
        }
    947
    }
    948
    949
    impl ComputeBrick for FlashAttentionBrick {
    950
        type Output = Vec<f32>;
    951
    952
    0
        fn name(&self) -> &'static str {
    953
    0
            "flash_attention"
    954
    0
        }
    955
    956
    1
        fn budget(&self) -> TokenBudget {
    957
    1
            self.budget
    958
    1
        }
    959
    960
    1
        fn assertions(&self) -> Vec<BrickAssertion> {
    961
    1
            vec![
    962
    1
                BrickAssertion::no_nan(),
    963
    1
                BrickAssertion::no_inf(),
    964
    1
                BrickAssertion::budget_met(),
    965
                // Attention outputs should be bounded
    966
    1
                BrickAssertion::bounds(-100.0, 100.0),
    967
                // Custom assertions for flash attention
    968
    1
                BrickAssertion {
    969
    1
                    name: "online_softmax".to_string(),
    970
    1
                    description: "Uses online softmax (no full attention matrix)".to_string(),
    971
    1
                    kind: AssertionKind::Custom {
    972
    1
                        check_name: "online_softmax".to_string(),
    973
    1
                    },
    974
    1
                },
    975
    1
                BrickAssertion {
    976
    1
                    name: "tiled_kv_access".to_string(),
    977
    1
                    description: "KV cache accessed in tiles for cache locality".to_string(),
    978
    1
                    kind: AssertionKind::Custom {
    979
    1
                        check_name: "tiled_kv_access".to_string(),
    980
    1
                    },
    981
    1
                },
    982
            ]
    983
    1
        }
    984
    985
    0
        fn can_run(&self) -> bool {
    986
    0
            self.num_heads > 0 && self.head_dim > 0 && self.tile_size > 0
    987
    0
        }
    988
    }
    989
    990
    /// FFN brick (SwiGLU).
    991
    #[derive(Debug)]
    992
    pub struct FfnBrick {
    993
        /// Hidden dimension
    994
        pub hidden_dim: usize,
    995
        /// Intermediate dimension
    996
        pub intermediate_dim: usize,
    997
        /// Budget
    998
        budget: TokenBudget,
    999
    }
    1000
    1001
    impl FfnBrick {
    1002
        /// Create a new FFN brick.
    1003
    105
        pub fn new(hidden_dim: usize, intermediate_dim: usize) -> Self {
    1004
    105
            Self {
    1005
    105
                hidden_dim,
    1006
    105
                intermediate_dim,
    1007
    105
                budget: TokenBudget::from_latency(12.2), // 12.2µs target
    1008
    105
            }
    1009
    105
        }
    1010
    1011
        /// Set custom budget.
    1012
        #[must_use]
    1013
    0
        pub fn with_budget(mut self, budget: TokenBudget) -> Self {
    1014
    0
            self.budget = budget;
    1015
    0
            self
    1016
    0
        }
    1017
    }
    1018
    1019
    impl ComputeBrick for FfnBrick {
    1020
        type Output = Vec<f32>;
    1021
    1022
    0
        fn name(&self) -> &'static str {
    1023
    0
            "ffn"
    1024
    0
        }
    1025
    1026
    6
        fn budget(&self) -> TokenBudget {
    1027
    6
            self.budget
    1028
    6
        }
    1029
    1030
    2
        fn assertions(&self) -> Vec<BrickAssertion> {
    1031
    2
            vec![
    1032
    2
                BrickAssertion::no_nan(),
    1033
    2
                BrickAssertion::no_inf(),
    1034
    2
                BrickAssertion::budget_met(),
    1035
            ]
    1036
    2
        }
    1037
    }
    1038
    1039
    /// Output projection brick.
    1040
    #[derive(Debug)]
    1041
    pub struct OProjBrick {
    1042
        /// Input dimension (num_heads * head_dim)
    1043
        pub in_dim: usize,
    1044
        /// Output dimension (hidden_dim)
    1045
        pub out_dim: usize,
    1046
        /// Budget
    1047
        budget: TokenBudget,
    1048
    }
    1049
    1050
    impl OProjBrick {
    1051
        /// Create a new O projection brick.
    1052
    2
        pub fn new(in_dim: usize, out_dim: usize) -> Self {
    1053
    2
            Self {
    1054
    2
                in_dim,
    1055
    2
                out_dim,
    1056
    2
                budget: TokenBudget::from_latency(3.5), // 3.5µs target
    1057
    2
            }
    1058
    2
        }
    1059
    1060
        /// Set custom budget.
    1061
        #[must_use]
    1062
    0
        pub fn with_budget(mut self, budget: TokenBudget) -> Self {
    1063
    0
            self.budget = budget;
    1064
    0
            self
    1065
    0
        }
    1066
    }
    1067
    1068
    impl ComputeBrick for OProjBrick {
    1069
        type Output = Vec<f32>;
    1070
    1071
    0
        fn name(&self) -> &'static str {
    1072
    0
            "o_proj"
    1073
    0
        }
    1074
    1075
    3
        fn budget(&self) -> TokenBudget {
    1076
    3
            self.budget
    1077
    3
        }
    1078
    1079
    0
        fn assertions(&self) -> Vec<BrickAssertion> {
    1080
    0
            vec![
    1081
    0
                BrickAssertion::no_nan(),
    1082
    0
                BrickAssertion::no_inf(),
    1083
    0
                BrickAssertion::budget_met(),
    1084
            ]
    1085
    0
        }
    1086
    }
    1087
    1088
    /// ActivationQuantBrick - Q8 activation quantization for memory bandwidth reduction (P2).
    1089
    ///
    1090
    /// **Purpose**: Quantize intermediate activations from f32 to int8 to reduce
    1091
    /// memory bandwidth by 4x during inter-layer communication.
    1092
    ///
    1093
    /// **Pipeline**:
    1094
    /// ```text
    1095
    /// Layer N output (f32) → Q8 quantize → transfer → Q8 dequantize → Layer N+1 input (f32)
    1096
    ///
    1097
    /// Memory bandwidth reduction:
    1098
    ///   f32: 4 bytes/element
    1099
    ///   int8: 1 byte/element + 2 floats (scale, zero_point)
    1100
    ///   Effective: ~4x reduction for large activations
    1101
    /// ```
    1102
    ///
    1103
    /// **Algorithm** (per-tensor affine quantization):
    1104
    /// ```text
    1105
    /// Quantize:
    1106
    ///   scale = (max - min) / 255
    1107
    ///   zero_point = round(-min / scale)
    1108
    ///   q[i] = clamp(round(x[i] / scale + zero_point), 0, 255)
    1109
    ///
    1110
    /// Dequantize:
    1111
    ///   x[i] = (q[i] - zero_point) * scale
    1112
    /// ```
    1113
    ///
    1114
    /// **Performance**: 2x memory BW improvement with ~0.1% accuracy loss
    1115
    ///
    1116
    /// **Reference**: Jacob, B., et al. (2018). "Quantization and Training of
    1117
    /// Neural Networks for Efficient Integer-Arithmetic-Only Inference." CVPR '18.
    1118
    #[derive(Debug, Clone)]
    1119
    pub struct ActivationQuantBrick {
    1120
        /// Activation dimension (e.g., hidden_dim or intermediate_dim)
    1121
        pub dim: usize,
    1122
        /// Budget (target: 0.5µs for quant+dequant overhead)
    1123
        budget: TokenBudget,
    1124
        /// Use per-channel quantization (more accurate but slower)
    1125
        pub per_channel: bool,
    1126
    }
    1127
    1128
    impl ActivationQuantBrick {
    1129
        /// Create new activation quantization brick.
    1130
        #[must_use]
    1131
    14
        pub fn new(dim: usize) -> Self {
    1132
    14
            Self {
    1133
    14
                dim,
    1134
    14
                budget: TokenBudget::from_latency(0.5), // 0.5µs overhead target
    1135
    14
                per_channel: false,
    1136
    14
            }
    1137
    14
        }
    1138
    1139
        /// Create with per-channel quantization.
    1140
        #[must_use]
    1141
    1
        pub fn with_per_channel(dim: usize) -> Self {
    1142
    1
            Self {
    1143
    1
                dim,
    1144
    1
                budget: TokenBudget::from_latency(1.0), // 1.0µs for per-channel
    1145
    1
                per_channel: true,
    1146
    1
            }
    1147
    1
        }
    1148
    1149
        /// Set custom budget.
    1150
        #[must_use]
    1151
    0
        pub fn with_budget(mut self, budget: TokenBudget) -> Self {
    1152
    0
            self.budget = budget;
    1153
    0
            self
    1154
    0
        }
    1155
    1156
        /// Compute memory bandwidth reduction factor.
    1157
        ///
    1158
        /// f32 (4 bytes) → int8 (1 byte) + scale/zero_point = ~4x reduction
    1159
        #[must_use]
    1160
    1
        pub fn bandwidth_reduction(&self) -> f64 {
    1161
            // Original: dim * 4 bytes (f32)
    1162
            // Quantized: dim * 1 byte (int8) + 8 bytes (scale + zero_point)
    1163
    1
            let original_bytes = self.dim * 4;
    1164
    1
            let quantized_bytes = self.dim + 8; // +8 for scale and zero_point (f32 each)
    1165
    1
            original_bytes as f64 / quantized_bytes as f64
    1166
    1
        }
    1167
    1168
        /// Compute quantization error estimate (typical for 8-bit).
    1169
        ///
    1170
        /// Per Jacob et al. 2018, typical Q8 error is ~0.1% for activations.
    1171
        #[must_use]
    1172
    2
        pub fn estimated_error(&self) -> f64 {
    1173
    2
            if self.per_channel {
    1174
    1
                0.0005 // 0.05% for per-channel
    1175
            } else {
    1176
    1
                0.001 // 0.1% for per-tensor
    1177
            }
    1178
    2
        }
    1179
    1180
        /// Compute bytes saved per token.
    1181
        #[must_use]
    1182
    3
        pub fn bytes_saved(&self) -> usize {
    1183
            // f32 (4 bytes) → int8 (1 byte) = 3 bytes saved per element
    1184
    3
            self.dim * 3
    1185
    3
        }
    1186
    1187
        /// Quantize f32 activations to int8 using Q8_0 block format.
    1188
        ///
    1189
        /// **REAL IMPLEMENTATION** - Not a stub.
    1190
        /// Uses symmetric quantization: scale = max(abs(values)) / 127.0
    1191
        ///
    1192
        /// # Arguments
    1193
        /// * `input` - f32 activations to quantize (must be length == self.dim)
    1194
        ///
    1195
        /// # Returns
    1196
        /// * Quantized int8 values and scale factors
    1197
        ///
    1198
        /// # Example
    1199
        /// ```ignore
    1200
        /// let brick = ActivationQuantBrick::new(64);
    1201
        /// let input = vec![1.0f32; 64];
    1202
        /// let (quants, scales) = brick.quantize(&input)?;
    1203
        /// assert_eq!(quants.len(), 64);
    1204
        /// ```
    1205
    4
        pub fn quantize(&self, input: &[f32]) -> Result<(Vec<i8>, Vec<f32>), BrickError> {
    1206
    4
            if input.len() != self.dim {
    1207
    1
                return Err(BrickError::InvalidInput(format!(
    1208
    1
                    "Input length {} != dim {}",
    1209
    1
                    input.len(),
    1210
    1
                    self.dim
    1211
    1
                )));
    1212
    3
            }
    1213
    3
            if self.dim == 0 {
    1214
    0
                return Err(BrickError::InvalidInput("Zero dimension".to_string()));
    1215
    3
            }
    1216
    1217
            // Quantize in blocks of 32 (Q8_0 block size)
    1218
    3
            let num_blocks = self.dim.div_ceil(32);
    1219
    3
            let mut quants = Vec::with_capacity(self.dim);
    1220
    3
            let mut scales = Vec::with_capacity(num_blocks);
    1221
    1222
    35
            for block_idx in 0..num_blocks3 {
    1223
    35
                let start = block_idx * 32;
    1224
    35
                let end = (start + 32).min(self.dim);
    1225
    1226
                // Pad to 32 if needed
    1227
    35
                let mut block_data = [0.0f32; 32];
    1228
    1.12k
                for (i, &v) in input35[start..end].iter35().enumerate35() {
    1229
    1.12k
                    block_data[i] = v;
    1230
    1.12k
                }
    1231
    1232
    35
                let block = Q8_0Block::quantize(&block_data);
    1233
    35
                scales.push(block.scale);
    1234
    1235
                // Only take the actual values (not padding)
    1236
    1.12k
                for &q in &block.quants35[0..(end - start)]35 {
    1237
    1.12k
                    quants.push(q);
    1238
    1.12k
                }
    1239
            }
    1240
    1241
    3
            Ok((quants, scales))
    1242
    4
        }
    1243
    1244
        /// Dequantize int8 back to f32 using stored scales.
    1245
        ///
    1246
        /// **REAL IMPLEMENTATION** - Not a stub.
    1247
    2
        pub fn dequantize(&self, quants: &[i8], scales: &[f32]) -> Result<Vec<f32>, BrickError> {
    1248
    2
            if quants.len() != self.dim {
    1249
    0
                return Err(BrickError::InvalidInput(format!(
    1250
    0
                    "Quants length {} != dim {}",
    1251
    0
                    quants.len(),
    1252
    0
                    self.dim
    1253
    0
                )));
    1254
    2
            }
    1255
    1256
    2
            let mut output = Vec::with_capacity(self.dim);
    1257
    2
            for (block_idx, &scale) in scales.iter().enumerate() {
    1258
    2
                let start = block_idx * 32;
    1259
    2
                let end = (start + 32).min(self.dim);
    1260
    64
                for &q in &quants2[start..end]2 {
    1261
    64
                    output.push(q as f32 * scale);
    1262
    64
                }
    1263
            }
    1264
    1265
    2
            Ok(output)
    1266
    2
        }
    1267
    1268
        /// Compute quantization error vs original input.
    1269
        ///
    1270
        /// **REAL IMPLEMENTATION** - Measures actual error, not estimates.
    1271
    1
        pub fn measure_error(
    1272
    1
            &self,
    1273
    1
            original: &[f32],
    1274
    1
            quants: &[i8],
    1275
    1
            scales: &[f32],
    1276
    1
        ) -> Result<f64, BrickError> {
    1277
    1
            let dequantized = self.dequantize(quants, scales)?0;
    1278
    1279
    1
            let max_error = original
    1280
    1
                .iter()
    1281
    1
                .zip(dequantized.iter())
    1282
    32
                .map1(|(a, b)| (a - b).abs())
    1283
    1
                .fold(0.0f32, f32::max);
    1284
    1285
    32
            let max_val1 = original1.iter1().map1(|v| v.abs()).fold1(0.0f32, f32::max);
    1286
    1
            if max_val < 1e-10 {
    1287
    0
                return Ok(0.0);
    1288
    1
            }
    1289
    1290
    1
            Ok((max_error / max_val) as f64)
    1291
    1
        }
    1292
    1293
        /// Execute quantization with timing (for benchmarking).
    1294
        #[allow(clippy::type_complexity)]
    1295
    1
        pub fn execute_timed(
    1296
    1
            &self,
    1297
    1
            input: &[f32],
    1298
    1
        ) -> Result<TokenResult<(Vec<i8>, Vec<f32>)>, BrickError> {
    1299
    1
            let start = Instant::now();
    1300
    1
            let (quants, scales) = self.quantize(input)?0;
    1301
    1
            let elapsed_us = start.elapsed().as_secs_f64() * 1_000_000.0;
    1302
    1303
    1
            Ok(TokenResult {
    1304
    1
                output: (quants, scales),
    1305
    1
                tokens_processed: 1,
    1306
    1
                us_per_token: elapsed_us,
    1307
    1
                tokens_per_sec: 1_000_000.0 / elapsed_us,
    1308
    1
                budget_met: elapsed_us <= self.budget.us_per_token,
    1309
    1
            })
    1310
    1
        }
    1311
    1312
        /// Legacy stub for backward compatibility (prefer `quantize()`)
    1313
        #[deprecated(note = "Use quantize() for real implementation")]
    1314
    2
        pub fn execute(&self) -> Result<Vec<u8>, BrickError> {
    1315
    2
            if self.dim == 0 {
    1316
    2
                return Err(BrickError::InvalidInput("Zero dimension".to_string()));
    1317
    0
            }
    1318
            // Return zeros for backward compat - use quantize() for real output
    1319
    0
            Ok(vec![128u8; self.dim])
    1320
    2
        }
    1321
    }
    1322
    1323
    impl ComputeBrick for ActivationQuantBrick {
    1324
        type Output = Vec<u8>;
    1325
    1326
    1
        fn name(&self) -> &'static str {
    1327
    1
            "activation_quant"
    1328
    1
        }
    1329
    1330
    1
        fn budget(&self) -> TokenBudget {
    1331
    1
            self.budget
    1332
    1
        }
    1333
    1334
    1
        fn assertions(&self) -> Vec<BrickAssertion> {
    1335
    1
            vec![
    1336
    1
                BrickAssertion::budget_met(),
    1337
    1
                BrickAssertion {
    1338
    1
                    name: "symmetric_range".to_string(),
    1339
    1
                    description: "Q8 values centered around 128 (zero_point)".to_string(),
    1340
    1
                    kind: AssertionKind::Custom {
    1341
    1
                        check_name: "symmetric_range".to_string(),
    1342
    1
                    },
    1343
    1
                },
    1344
    1
                BrickAssertion {
    1345
    1
                    name: "error_bound".to_string(),
    1346
    1
                    description: "Quantization error < 0.1% (per-tensor) or 0.05% (per-channel)"
    1347
    1
                        .to_string(),
    1348
    1
                    kind: AssertionKind::Custom {
    1349
    1
                        check_name: "error_bound".to_string(),
    1350
    1
                    },
    1351
    1
                },
    1352
            ]
    1353
    1
        }
    1354
    1355
    2
        fn can_run(&self) -> bool {
    1356
    2
            self.dim > 0
    1357
    2
        }
    1358
    }
    1359
    1360
    // ============================================================================
    1361
    // Transformer Layer Brick
    1362
    // ============================================================================
    1363
    1364
    /// Full transformer layer as a composed brick.
    1365
    #[derive(Debug)]
    1366
    pub struct TransformerLayerBrick {
    1367
        /// Layer index
    1368
        pub layer_idx: usize,
    1369
        /// Attention layer normalization brick
    1370
        pub attn_norm: RmsNormBrick,
    1371
        /// QKV projection brick
    1372
        pub qkv: QkvBrick,
    1373
        /// Rotary position embedding brick
    1374
        pub rope: RopeBrick,
    1375
        /// Attention computation brick
    1376
        pub attention: AttentionBrick,
    1377
        /// Output projection brick
    1378
        pub o_proj: OProjBrick,
    1379
        /// FFN layer normalization brick
    1380
        pub ffn_norm: RmsNormBrick,
    1381
        /// Feed-forward network brick
    1382
        pub ffn: FfnBrick,
    1383
        /// Timing metrics (updated after each run)
    1384
        pub last_timing: Option<LayerTiming>,
    1385
    }
    1386
    1387
    /// Timing breakdown for a layer.
    1388
    #[derive(Debug, Clone, Default)]
    1389
    pub struct LayerTiming {
    1390
        /// Attention normalization time (µs)
    1391
        pub attn_norm_us: f64,
    1392
        /// QKV projection time (µs)
    1393
        pub qkv_us: f64,
    1394
        /// RoPE application time (µs)
    1395
        pub rope_us: f64,
    1396
        /// Attention computation time (µs)
    1397
        pub attention_us: f64,
    1398
        /// Output projection time (µs)
    1399
        pub o_proj_us: f64,
    1400
        /// FFN normalization time (µs)
    1401
        pub ffn_norm_us: f64,
    1402
        /// FFN computation time (µs)
    1403
        pub ffn_us: f64,
    1404
        /// Total layer time (µs)
    1405
        pub total_us: f64,
    1406
    }
    1407
    1408
    impl LayerTiming {
    1409
        /// Find the bottleneck brick.
    1410
    2
        pub fn bottleneck(&self) -> (&'static str, f64) {
    1411
    2
            let bricks = [
    1412
    2
                ("attn_norm", self.attn_norm_us),
    1413
    2
                ("qkv", self.qkv_us),
    1414
    2
                ("rope", self.rope_us),
    1415
    2
                ("attention", self.attention_us),
    1416
    2
                ("o_proj", self.o_proj_us),
    1417
    2
                ("ffn_norm", self.ffn_norm_us),
    1418
    2
                ("ffn", self.ffn_us),
    1419
    2
            ];
    1420
    1421
    2
            bricks
    1422
    2
                .into_iter()
    1423
    12
                .max_by2(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
    1424
    2
                .unwrap_or(("unknown", 0.0))
    1425
    2
        }
    1426
    }
    1427
    1428
    impl TransformerLayerBrick {
    1429
        /// Create from configuration.
    1430
    2
        pub fn from_config(
    1431
    2
            layer_idx: usize,
    1432
    2
            hidden_dim: usize,
    1433
    2
            num_heads: usize,
    1434
    2
            num_kv_heads: usize,
    1435
    2
            intermediate_dim: usize,
    1436
    2
            eps: f32,
    1437
    2
            rope_theta: f32,
    1438
    2
            rope_type: u32,
    1439
    2
        ) -> Self {
    1440
    2
            let head_dim = hidden_dim / num_heads;
    1441
    2
            let q_dim = num_heads * head_dim;
    1442
    2
            let kv_dim = num_kv_heads * head_dim;
    1443
    1444
    2
            Self {
    1445
    2
                layer_idx,
    1446
    2
                attn_norm: RmsNormBrick::new(vec![1.0; hidden_dim], eps),
    1447
    2
                qkv: QkvBrick::new(hidden_dim, q_dim, kv_dim, kv_dim),
    1448
    2
                rope: RopeBrick::new(head_dim, num_heads, rope_theta, rope_type),
    1449
    2
                attention: AttentionBrick::new(num_heads, num_kv_heads, head_dim),
    1450
    2
                o_proj: OProjBrick::new(q_dim, hidden_dim),
    1451
    2
                ffn_norm: RmsNormBrick::new(vec![1.0; hidden_dim], eps),
    1452
    2
                ffn: FfnBrick::new(hidden_dim, intermediate_dim),
    1453
    2
                last_timing: None,
    1454
    2
            }
    1455
    2
        }
    1456
    1457
        /// Get total budget for this layer.
    1458
    1
        pub fn total_budget_us(&self) -> f64 {
    1459
    1
            self.attn_norm.budget().us_per_token
    1460
    1
                + self.qkv.budget().us_per_token
    1461
    1
                + self.rope.budget().us_per_token
    1462
    1
                + self.attention.budget().us_per_token
    1463
    1
                + self.o_proj.budget().us_per_token
    1464
    1
                + self.ffn_norm.budget().us_per_token
    1465
    1
                + self.ffn.budget().us_per_token
    1466
    1
        }
    1467
    }
    1468
    1469
    impl ComputeBrick for TransformerLayerBrick {
    1470
        type Output = Vec<f32>;
    1471
    1472
    0
        fn name(&self) -> &'static str {
    1473
    0
            "transformer_layer"
    1474
    0
        }
    1475
    1476
    0
        fn budget(&self) -> TokenBudget {
    1477
    0
            TokenBudget::from_latency(self.total_budget_us())
    1478
    0
        }
    1479
    1480
    1
        fn assertions(&self) -> Vec<BrickAssertion> {
    1481
    1
            vec![
    1482
    1
                BrickAssertion::no_nan(),
    1483
    1
                BrickAssertion::no_inf(),
    1484
    1
                BrickAssertion::budget_met(),
    1485
            ]
    1486
    1
        }
    1487
    1488
    1
        fn verify(&self) -> BrickVerification {
    1489
            // Verify all component bricks
    1490
    1
            let mut result = BrickVerification::pass();
    1491
    1492
    2
            for brick in [
    1493
    1
                &self.attn_norm as &dyn ComputeBrick<Output = Vec<f32>>,
    1494
    1
                &self.ffn_norm as &dyn ComputeBrick<Output = Vec<f32>>,
    1495
            ] {
    1496
    2
                let v = brick.verify();
    1497
    2
                if !v.is_valid {
    1498
    0
                    result.is_valid = false;
    1499
    0
                    result.results.extend(v.results);
    1500
    2
                }
    1501
            }
    1502
    1503
    1
            result
    1504
    1
        }
    1505
    }
    1506
    1507
    // ============================================================================
    1508
    // Bottleneck Report
    1509
    // ============================================================================
    1510
    1511
    /// Report identifying pipeline bottleneck.
    1512
    #[derive(Debug, Clone)]
    1513
    pub struct BottleneckReport {
    1514
        /// Layer index containing bottleneck
    1515
        pub layer_idx: usize,
    1516
        /// Brick name
    1517
        pub brick_name: &'static str,
    1518
        /// Actual latency (µs)
    1519
        pub actual_us: f64,
    1520
        /// Budget latency (µs)
    1521
        pub budget_us: f64,
    1522
        /// Gap factor (actual / budget)
    1523
        pub gap_factor: f64,
    1524
    }
    1525
    1526
    impl fmt::Display for BottleneckReport {
    1527
    0
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    1528
    0
            write!(
    1529
    0
                f,
    1530
    0
                "Bottleneck: {} (layer {}) - {:.1}µs actual vs {:.1}µs budget ({:.2}x)",
    1531
                self.brick_name, self.layer_idx, self.actual_us, self.budget_us, self.gap_factor
    1532
            )
    1533
    0
        }
    1534
    }
    1535
    1536
    // ============================================================================
    1537
    // Benchmark Brick
    1538
    // ============================================================================
    1539
    1540
    /// Configuration for benchmark runs.
    1541
    #[derive(Debug, Clone)]
    1542
    pub struct BenchmarkConfig {
    1543
        /// Number of warmup iterations
    1544
        pub warmup: usize,
    1545
        /// Number of sample iterations
    1546
        pub samples: usize,
    1547
        /// Maximum allowed CV (coefficient of variation)
    1548
        pub max_cv: f64,
    1549
    }
    1550
    1551
    impl Default for BenchmarkConfig {
    1552
    7
        fn default() -> Self {
    1553
    7
            Self {
    1554
    7
                warmup: 10,
    1555
    7
                samples: 100,
    1556
    7
                max_cv: 0.05, // 5% per Stabilizer (Curtsinger & Berger 2013)
    1557
    7
            }
    1558
    7
        }
    1559
    }
    1560
    1561
    /// Benchmark report with statistical analysis.
    1562
    #[derive(Debug, Clone)]
    1563
    pub struct BenchmarkReport {
    1564
        /// Brick name
    1565
        pub brick_name: String,
    1566
        /// Mean latency (µs)
    1567
        pub mean_us: f64,
    1568
        /// Standard deviation (µs)
    1569
        pub std_us: f64,
    1570
        /// Coefficient of variation
    1571
        pub cv: f64,
    1572
        /// 50th percentile (µs)
    1573
        pub p50_us: f64,
    1574
        /// 99th percentile (µs)
    1575
        pub p99_us: f64,
    1576
        /// Throughput (tokens/sec)
    1577
        pub tokens_per_sec: f64,
    1578
        /// Budget target (µs)
    1579
        pub budget_us: f64,
    1580
        /// Budget met?
    1581
        pub budget_met: bool,
    1582
        /// Statistical validity (CV < max_cv)
    1583
        pub statistically_valid: bool,
    1584
    }
    1585
    1586
    impl fmt::Display for BenchmarkReport {
    1587
    0
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    1588
    0
            let status = if self.budget_met { "PASS" } else { "FAIL" };
    1589
    0
            write!(
    1590
    0
                f,
    1591
    0
                "{}: {:.1}µs ± {:.1}µs (CV={:.1}%) | {:.0} tok/s | budget: {} ({})",
    1592
                self.brick_name,
    1593
                self.mean_us,
    1594
                self.std_us,
    1595
    0
                self.cv * 100.0,
    1596
                self.tokens_per_sec,
    1597
                self.budget_us,
    1598
                status
    1599
            )
    1600
    0
        }
    1601
    }
    1602
    1603
    /// Calculate percentile from sorted samples.
    1604
    2
    fn percentile(samples: &[f64], p: f64) -> f64 {
    1605
    2
        if samples.is_empty() {
    1606
    0
            return 0.0;
    1607
    2
        }
    1608
    2
        let idx = ((samples.len() as f64) * p).floor() as usize;
    1609
    2
        samples[idx.min(samples.len() - 1)]
    1610
    2
    }
    1611
    1612
    /// Run benchmark on a brick with statistical rigor.
    1613
    1
    pub fn benchmark_brick<B: ComputeBrick>(
    1614
    1
        brick: &B,
    1615
    1
        run_fn: impl Fn() -> f64,
    1616
    1
        config: &BenchmarkConfig,
    1617
    1
    ) -> BenchmarkReport {
    1618
        // Warmup (Jidoka: ensure stable state)
    1619
    5
        for _ in 0..config.warmup1 {
    1620
    5
            let _ = run_fn();
    1621
    5
        }
    1622
    1623
        // Collect samples
    1624
    1
        let mut samples: Vec<f64> = Vec::with_capacity(config.samples);
    1625
    50
        for _ in 0..config.samples1 {
    1626
    50
            samples.push(run_fn());
    1627
    50
        }
    1628
    1629
        // Sort for percentiles
    1630
    293
        samples1.sort_by1(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    1631
    1632
        // Statistical analysis
    1633
    1
        let mean = samples.iter().sum::<f64>() / samples.len() as f64;
    1634
    1
        let std =
    1635
    50
            (samples.iter()1.map1(|x| (x - mean).powi(2)).sum1::<f64>() / samples.len() as f641).sqrt1();
    1636
    1
        let cv = std / mean;
    1637
    1638
    1
        let budget = brick.budget();
    1639
    1640
    1
        BenchmarkReport {
    1641
    1
            brick_name: brick.name().to_string(),
    1642
    1
            mean_us: mean,
    1643
    1
            std_us: std,
    1644
    1
            cv,
    1645
    1
            p50_us: percentile(&samples, 0.50),
    1646
    1
            p99_us: percentile(&samples, 0.99),
    1647
    1
            tokens_per_sec: 1_000_000.0 / mean,
    1648
    1
            budget_us: budget.us_per_token,
    1649
    1
            budget_met: mean <= budget.us_per_token,
    1650
    1
            statistically_valid: cv <= config.max_cv,
    1651
    1
        }
    1652
    1
    }
    1653
    1654
    // ============================================================================
    1655
    // CUDA Graph Brick (Section 5.2 - P0)
    1656
    // ============================================================================
    1657
    1658
    /// CUDA Graph Brick for eliminating kernel launch overhead.
    1659
    ///
    1660
    /// Per spec: docs/specifications/qwen2.5-coder-showcase-demo.md §5.2
    1661
    ///
    1662
    /// Uses CUDA graph capture to reduce ~280 kernel launches to single graph replay.
    1663
    /// Expected impact: 5.6ms overhead → 0.02ms = 280x overhead reduction.
    1664
    ///
    1665
    /// # Implementation
    1666
    ///
    1667
    /// Wraps `CudaExecutor::decode_graph` and `try_graph_capture()` from cuda.rs.
    1668
    /// Uses indirect kernels (KvCacheScatterIndirect, RopeIndirect) for graph compatibility.
    1669
    #[derive(Debug, Clone)]
    1670
    pub struct CudaGraphBrick {
    1671
        /// Number of layers captured in graph
    1672
        pub num_layers: usize,
    1673
        /// Hidden dimension
    1674
        pub hidden_dim: usize,
    1675
        /// Whether graph is currently captured
    1676
        pub captured: bool,
    1677
        /// Token budget (target: 10µs launch overhead vs 5600µs eager)
    1678
        budget: TokenBudget,
    1679
    }
    1680
    1681
    impl CudaGraphBrick {
    1682
        /// Create new CUDA Graph brick for model configuration.
    1683
        #[must_use]
    1684
    4
        pub fn new(num_layers: usize, hidden_dim: usize) -> Self {
    1685
            // Graph overhead should be < 100µs (vs ~5.6ms for 280 launches)
    1686
    4
            let budget_us = 20.0; // Conservative: 20µs for graph replay
    1687
    4
            Self {
    1688
    4
                num_layers,
    1689
    4
                hidden_dim,
    1690
    4
                captured: false,
    1691
    4
                budget: TokenBudget::from_latency(budget_us),
    1692
    4
            }
    1693
    4
        }
    1694
    1695
        /// Set custom budget.
    1696
        #[must_use]
    1697
    0
        pub fn with_budget(mut self, budget: TokenBudget) -> Self {
    1698
    0
            self.budget = budget;
    1699
    0
            self
    1700
    0
        }
    1701
    1702
        /// Mark graph as captured.
    1703
    1
        pub fn set_captured(&mut self, captured: bool) {
    1704
    1
            self.captured = captured;
    1705
    1
        }
    1706
    1707
        /// Check if graph can be used (captured and valid).
    1708
        #[must_use]
    1709
    2
        pub fn can_replay(&self) -> bool {
    1710
    2
            self.captured
    1711
    2
        }
    1712
    1713
        /// Replay the captured graph (stub - actual execution via CudaExecutor).
    1714
    1
        pub fn replay(&self) -> Result<(), BrickError> {
    1715
    1
            if !self.captured {
    1716
    0
                return Err(BrickError::ComputeError(
    1717
    0
                    "CUDA graph not captured yet".to_string(),
    1718
    0
                ));
    1719
    1
            }
    1720
            // Actual replay would be done via CudaExecutor::forward_graphed()
    1721
    1
            Ok(())
    1722
    1
        }
    1723
    }
    1724
    1725
    impl ComputeBrick for CudaGraphBrick {
    1726
        type Output = ();
    1727
    1728
    1
        fn name(&self) -> &'static str {
    1729
    1
            "cuda_graph"
    1730
    1
        }
    1731
    1732
    2
        fn budget(&self) -> TokenBudget {
    1733
    2
            self.budget
    1734
    2
        }
    1735
    1736
    1
        fn assertions(&self) -> Vec<BrickAssertion> {
    1737
    1
            vec![
    1738
    1
                BrickAssertion::budget_met(),
    1739
    1
                BrickAssertion {
    1740
    1
                    name: "graph_speedup".to_string(),
    1741
    1
                    description: "Graph replay faster than eager execution".to_string(),
    1742
    1
                    kind: AssertionKind::Custom {
    1743
    1
                        check_name: "graph_speedup".to_string(),
    1744
    1
                    },
    1745
    1
                },
    1746
            ]
    1747
    1
        }
    1748
    1749
    0
        fn can_run(&self) -> bool {
    1750
    0
            self.num_layers > 0 && self.hidden_dim > 0
    1751
    0
        }
    1752
    }
    1753
    1754
    // ============================================================================
    1755
    // Tests (F001-F020)
    1756
    // ============================================================================
    1757
    1758
    // Tests extracted to tests.rs (PMAT-802)
    1759
    #[cfg(test)]
    1760
    #[path = "tests.rs"]
    1761
    mod brick_tests;