## Diagram: Two-Stage Transformer Processing with Caching Mechanism
### Overview
The diagram illustrates a two-stage processing pipeline for a transformer-based model, featuring a **prefill** stage (a) and a **decode** stage (b). The system employs a **group of 2 K/V caches** to manage cached and non-cached data, with explicit separation of computational paths for efficiency.
---
### Components/Axes
#### Legend
- **Blue**: Non-Cached (e.g., L1, L2, Concatenated K/V)
- **Red**: Cached (e.g., B1, B2, Shared Token Basis A)
#### Prefill Stage (a)
1. **Input Layers**:
- `g₁` (Layer 1) → `L1` (Non-Cached)
- `g₁` (Layer 2) → `L2` (Non-Cached)
- `gₖ` (Layer N-1) → `L_{N-1}` (Non-Cached)
- `gₖ` (Layer N) → `L_N` (Non-Cached)
2. **Caching Mechanism**:
- Group of 2 K/V Caches (blue boxes) store intermediate K/V pairs.
3. **Processing Flow**:
- **Concatenated K/V**: Combines K/V from L1 and L2.
- **Shared Token Basis (A)**: Red box, acts as a reference for reconstruction.
- **SVD (Singular Value Decomposition)**: Reduces dimensionality of concatenated K/V.
- **Reconstruction Matrices (B1, B2)**: Red boxes, derived from SVD for cached data.
#### Decode Stage (b)
1. **Input**:
- Shared Token Basis (A) → Matmul (Matrix Multiplication) → B1/B2.
2. **Processing Flow**:
- **B1** → Matmul → Fetch K/V of Layer 1 → `L1` (Non-Cached).
- **B2** → Matmul → Fetch K/V of Layer 2 → `L2` (Non-Cached).
---
### Detailed Analysis
#### Prefill Stage
- **Layer Hierarchy**: Input layers (`g₁`, `gₖ`) feed into sequential non-cached layers (L1–LN).
- **Caching Logic**:
- K/V pairs from L1 and L2 are cached in a group of 2 K/V caches.
- Concatenated K/V is processed via SVD to generate low-rank approximations (B1, B2).
- **Shared Token Basis (A)**: Central to both stages, likely representing a fixed embedding or reference for reconstruction.
#### Decode Stage
- **Reconstruction Path**:
- B1 and B2 (from Prefill) are multiplied with the Shared Token Basis (A) to reconstruct K/V pairs for specific layers.
- Reconstructed K/V is fed back into L1 and L2 for decoding.
- **Layer-Specific Fetching**:
- B1 targets Layer 1 K/V, B2 targets Layer 2 K/V, ensuring selective retrieval.
---
### Key Observations
1. **Caching Efficiency**:
- Only L1 and L2 K/V pairs are cached, reducing memory overhead.
- B1/B2 matrices enable efficient reconstruction without storing full K/V.
2. **Stage Separation**:
- Prefill focuses on encoding and caching.
- Decode focuses on reconstructing and utilizing cached data.
3. **Color Consistency**:
- Blue (non-cached) dominates Prefill; red (cached) dominates Decode.
- Shared Token Basis (A) is red, indicating its role in caching.
---
### Interpretation
This diagram represents a **memory-efficient transformer architecture** where:
- **Prefill** optimizes computation by caching critical K/V pairs and compressing them via SVD.
- **Decode** leverages cached data (via B1/B2) to reconstruct layer-specific K/V, reducing redundant computations.
- The **Shared Token Basis (A)** acts as a bridge between cached and non-cached data, ensuring consistency across stages.
The design prioritizes **selective caching** (only L1/L2) and **dimensionality reduction** (SVD), which are critical for handling long sequences in transformers. The absence of numerical values suggests a conceptual flow rather than empirical data, emphasizing architectural logic over performance metrics.
