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;