======================================================================
LAYER NORMALIZATION: FORWARD AND BACKWARD FROM SCRATCH
======================================================================

╔══════════════════════════════════════════════════════════════════════════════╗
║                    COMPLEXITY ANALYSIS                                         ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Input: x ∈ ℝ^(B×T×D)                                                        ║
║  Parameters: γ, β ∈ ℝ^D                                                      ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FORWARD PASS                                                                ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Operation              │ Work (FLOPs)      │ Memory                        ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  mean (reduction)       │ O(BTD)            │ -                             ║
║  x - mean (broadcast)   │ O(BTD)            │ O(BTD)                        ║
║  var (reduction)        │ O(BTD)            │ -                             ║
║  sqrt(var + eps)        │ O(BT)             │ -                             ║
║  divide by std          │ O(BTD)            │ -                             ║
║  gamma * x_norm         │ O(BTD)            │ -                             ║
║  add beta               │ O(BTD)            │ O(BTD) output                 ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  TOTAL                    │ 5×O(BTD)          │ O(BTD)                       ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  BACKWARD PASS                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Operation              │ Work (FLOPs)      │ Memory                        ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  d_gamma = sum(dy*x_norm)│ O(BTD)           │ O(D)                          ║
║  d_beta = sum(dy)       │ O(BTD)            │ O(D)                          ║
║  dz = dy * gamma        │ O(BTD)            │ O(BTD) (can be avoided)        ║
║  sum(dz)                │ O(BTD)            │ -                             ║
║  sum(dz * x_norm)       │ O(BTD)            │ -                             ║
║  dx computation         │ O(BTD)            │ O(BTD) output                 ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  TOTAL                    │ 5×O(BTD)          │ O(BTD) + O(D)                ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  SUMMARY                                                                    ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Time Complexity:  O(BTD) for both forward and backward                      ║
║  Space Complexity: O(BTD) for storing activations (during training)          ║
║                   O(BTD) during inference (no need to store)                 ║
║                                                                              ║
║  Cache efficiency: We store x_centered, x_norm, mean, std                    ║
║                   These are O(BTD) total, reused across all gradient comps  ║
║                                                                              ║
╚══════════════════════════════════════════════════════════════════════════════╝

╔══════════════════════════════════════════════════════════════════════════════╗
║               GPU KERNEL FUSION DESIGN                                       ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  CURRENT SEPARATE KERNEL APPROACH:                                           ║
║  ────────────────────────────────────                                        ║
║  Kernel 1: Compute mean (reduction over D)                                   ║
║  Kernel 2: Compute variance (reduction over D)                               ║
║  Kernel 3: Normalize (element-wise)                                          ║
║  Kernel 4: Scale and shift (element-wise)                                    ║
║  Kernel 5: Backward kernels (x, gamma, beta)                                 ║
║                                                                              ║
║  Issues with separate kernels:                                               ║
║  • Multiple kernel launches (overhead)                                       ║
║  • Data movement between global memory passes                                 ║
║  • Can't use persistent threads for reduction efficiency                      ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FUSED KERNEL DESIGN (Forward):                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Grid: (B × T) blocks, each handling one (b,t) position                      ║
║  Block: 256-512 threads                                                      ║
║                                                                              ║
║  PHASE 1: Load and compute local sum                                         ║
║  ─────────────────────────────────                                           ║
║  • Each thread loads x[b,t,d] into shared memory                             ║
║  • Compute partial sum using warp-level reduction                            ║
║  • Single thread writes mean to __shared__                                    ║
║                                                                              ║
║  PHASE 2: Compute variance locally                                           ║
║  ─────────────────────────────────                                           ║
║  • Re-load x with loaded mean                                                ║
║  • Compute (x-mean)² and partial variance                                    ║
║  • Reduce to get variance                                                   ║
║                                                                              ║
║  PHASE 3: Normalize and output                                               ║
║  ─────────────────────────────────                                           ║
║  • All threads compute: y = gamma * (x-mean) / sqrt(var+eps) + beta          ║
║  • Write to output (fully coalesced)                                         ║
║                                                                              ║
║  MEMORY ACCESS PATTERN:                                                      ║
║  • Each block reads contiguous D elements (coalesced)                        ║
║  • Use shared memory for intermediate results                                ║
║  • Output writes are also coalesced                                          ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FUSED KERNEL DESIGN (Backward):                                             ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Key insight: dz = dy * gamma can be merged into the computation              ║
║                                                                              ║
║  Grid: (B × T) blocks                                                       ║
║  Block: 256-512 threads                                                      ║
║                                                                              ║
║  SHARED MEMORY STRUCTURE:                                                    ║
║  [x_norm_0, x_norm_1, ..., x_norm_{D-1}]                                     ║
║  [dz_0, dz_1, ..., dz_{D-1}]                                                 ║
║                                                                              ║
║  ALGORITHM:                                                                  ║
║  1. Load x_norm and compute local dz = dy * gamma                            ║
║  2. Reduce to get sum(dz) and sum(dz * x_norm)                              ║
║  3. Second pass to compute dx using the formula:                             ║
║     dx = (dz - mean(dz) - x_norm * mean(dz*x_norm)) / std                    ║
║                                                                              ║
║  REDUCTION OPTIMIZATIONS:                                                    ║
║  • Warp-level shuffle reductions (no shared memory needed)                    ║
║  • Block-level using shared memory with tree reduction                        ║
║  • Use block-level primitives for final reduction                            ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  BENEFITS OF FUSION:                                                          ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  ✓ Reduced kernel launch overhead (1 kernel vs 4-5)                           ║
║  ✓ Better memory bandwidth utilization (single read, single write)          ║
║  ✓ Improved cache locality (data stays in registers/shared mem)              ║
║  ✓ Only loads x once, computes mean and var from same data                   ║
║  ✓ Backward can reuse cached values from forward (if memory allows)          ║
║  ✓ Lower register pressure allows for larger block sizes                     ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  CUDA KERNEL SKETCH (Pseudo-code):                                            ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  __global__ void layer_norm_fwd(float* y, const float* x,                    ║
║                                  const float* gamma, const float* beta,      ║
║                                  int B, int T, int D, float eps) {          ║
║                                                                              ║
║      __shared__ float mean_smem[256]; // block-level mean                     ║
║      __shared__ float var_smem[256];  // block-level variance                 ║
║      __shared__ float std_smem[256]; // block-level std                      ║
║                                                                              ║
║      int tid = threadIdx.x;                                                  ║


Let me check the gradient check output specifically.

bash

cd ./minimax-m2.7/backwards && python layer_norm_numpy.py 2>&1 | tail -100


Analytical gradient shape: (16,)
Numerical gradient shape: (16,)
Max absolute difference: 2.52e-10
Max relative tolerance: 8.59e-04
Mean analytical: 3.506940e+00
Mean numerical: 3.506940e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: x (input)
============================================================
Analytical gradient shape: (4, 8, 16)
Numerical gradient shape: (4, 8, 16)
Elements checked: 512 / 512
Max absolute difference: 4.60e+00
Max relative tolerance: 1.00e-05
Mean analytical (checked): 6.141263e-01
Mean numerical (checked): 0.000000e+00

Gradient check: FAILED ✗
     gamma: ✓ PASS
      beta: ✓ PASS
         x: ✗ FAIL

============================================================
Testing shape (B=8, T=16, D=32)
============================================================

============================================================
GRADIENT CHECK: gamma
============================================================
Analytical gradient shape: (32,)
Numerical gradient shape: (32,)
Max absolute difference: 2.38e-09
Max relative tolerance: 2.84e-03
Mean analytical: 9.330259e+00
Mean numerical: 9.330259e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: beta
============================================================
Analytical gradient shape: (32,)
Numerical gradient shape: (32,)
Max absolute difference: 2.00e-09
Max relative tolerance: 2.61e-03
Mean analytical: 9.296132e+00
Mean numerical: 9.296132e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: x (input)
============================================================
Analytical gradient shape: (8, 16, 32)
Numerical gradient shape: (8, 16, 32)
Elements checked: 4096 / 4096
Max absolute difference: 6.83e+00
Max relative tolerance: 1.00e-05
Mean analytical (checked): 7.088609e-01
Mean numerical (checked): 0.000000e+00

Gradient check: FAILED ✗
     gamma: ✓ PASS
      beta: ✓ PASS
         x: ✗ FAIL

======================================================================
SOME GRADIENT CHECKS FAILED ✗
======================================================================

======================================================================
BENCHMARKING LAYER NORMALIZATION
======================================================================

Shape: (B=32, T=128, D=256)
----------------------------------------
Forward:  6.207 ± 0.586 ms (0.2 GB/s)
Backward: 7.604 ± 0.529 ms (0.1 GB/s)
Total:    13.810 ms

Shape: (B=64, T=128, D=512)
----------------------------------------
Forward:  27.305 ± 2.522 ms (0.2 GB/s)
Backward: 35.805 ± 4.186 ms (0.1 GB/s)
Total:    63.110 ms

Shape: (B=32, T=512, D=768)
----------------------------------------
Forward:  83.459 ± 5.145 ms (0.2 GB/s)
Backward: 106.003 ± 9.338 ms (0.1 GB/s)
Total:    189.462 ms

Shape: (B=16, T=512, D=1024)
----------------------------------------
Forward:  54.293 ± 4.970 ms (0.2 GB/s)
Backward: 69.341 ± 6.961 ms (0.1 GB/s)
Total:    123.634 ms


The gradient check for x is still failing - numerical gradients are all zeros! The mean numerical is 0.000000e+00 which is wrong.

Let me trace through the issue more carefully. The problem is that `compute_numerical_gradient_x` modifies x, then calls `layer_norm_forward(x, gamma, beta)`. But x is passed by reference, so when we do `x_flat[i] = original + h`, we're modifying the ORIGINAL x that was passed in.

Wait, no - I create a copy first: `orig_x = x.copy()`.

Let me think about this again:
1. `orig_x = x.copy()` - saves original values
2. `x_flat = x.flatten()` - creates a view of x as a 1D array (still the same underlying data)
3. When we do `x_flat[i] = original + h`, we're modifying x (and orig_x is not affected)
4. Then we call `layer_norm_forward(x, gamma, beta)` which computes the loss
5. After restoring `x_flat[i] = original`, we continue

This should work... Let me add debug output to trace what's happening.

Actually, let me re-read my function more carefully. I see the issue now!

In `compute_numerical_gradient_x`:
```python
# Restore x to original values
x[:] = orig_x

return num_grad, np.ones((B, T, D), dtype=bool)  # All elements checked
```

Wait, I AM restoring x. But let me check the forward function to make sure it doesn't modify x...

No, `layer_norm_forward` doesn't modify x.

Let me add some debug print statements.

cd ./minimax-m2.7/backwards && python3 -c "
import numpy as np

def layer_norm_forward(x, gamma, beta, eps=1e-8):
    B, T, D = x.shape
    mean = np.mean(x, axis=-1, keepdims=True)
    x_centered = x - mean
    var = np.mean(x_centered ** 2, axis=-1, keepdims=True)
    std = np.sqrt(var + eps)
    x_norm = x_centered / std
    y = gamma * x_norm + beta
    return y, {'mean': mean, 'std': std, 'x_norm': x_norm, 'x_centered': x_centered}

def compute_numerical_gradient_x(x, gamma, beta, dy, h=1e-5):
    B, T, D = x.shape
    total_elements = B * T * D

    # Save original x values
    orig_x = x.copy()
    print(f'Original x[0,0,0]: {x[0,0,0]:.6f}')

    num_grad = np.zeros_like(x)
    x_flat = x.flatten()
    num_grad_flat = num_grad.flatten()

    # Check first element
    i = 0
    original = x_flat[i]
    print(f'original value: {original:.6f}')

    # f(x + h)
    x_flat[i] = original + h
    y_plus, _ = layer_norm_forward(x, gamma, beta)
    loss_plus = np.sum(y_plus * dy)
    print(f'Loss plus: {loss_plus:.6f}')
    print(f'x after +h modification: x[0,0,0] = {x[0,0,0]:.6f}')

    # f(x - h)
    x_flat[i] = original - h
    y_minus, _ = layer_norm_forward(x, gamma, beta)
    loss_minus = np.sum(y_minus * dy)
    print(f'Loss minus: {loss_minus:.6f}')
    print(f'x after -h modification: x[0,0,0] = {x[0,0,0]:.6f}')

    # Central difference
    num_grad_flat[i] = (loss_plus - loss_minus) / (2 * h)
    print(f'Numerical gradient: {num_grad_flat[i]:.6f}')

    # Restore
    x_flat[i] = original
    print(f'x restored: x[0,0,0] = {x[0,0,0]:.6f}')

    # Restore x to original values
    x[:] = orig_x
    print(f'x fully restored: x[0,0,0] = {x[0,0,0]:.6f}')

    return num_grad

# Test
np.random.seed(42)
B, T, D = 2, 4, 8
x = np.random.randn(B, T, D).astype(np.float64)
gamma = np.ones(D).astype(np.float64)
beta = np.zeros(D).astype(np.float64)
dy = np.random.randn(B, T, D).astype(np.float64)

num_grad = compute_numerical_gradient_x(x, gamma, beta, dy)
print(f'\\nResult num_grad[0,0,0]: {num_grad[0,0,0]:.6f}')
"


Original x[0,0,0]: 0.496714
original value: 0.496714
Loss plus: -3.607273
x after +h modification: x[0,0,0] = 0.496714
Loss minus: -3.607273
x after -h modification: x[0,0,0] = 0.496714
Numerical gradient: 0.000000
x restored: x[0,0,0] = 0.496714
x fully restored: x[0,0,0] = 0.496714

Result num_grad[0,0,0]: 0.000000


Ah-ha! I found the bug! When we do `x_flat[i] = original + h` and then call `layer_norm_forward(x, ...)`, the x array is being modified, but the result shows `x[0,0,0] = 0.496714` (unchanged).This suggests that `x_flat = x.flatten()` creates a **copy** of x, not a view. So modifying `x_flat[i]` doesn't affect `x`!

Let me verify this:
```python
x_flat = x.flatten()
x_flat[0] = 999
print(x[0])  # Should still be original value!
```

Yes, that's the bug. `flatten()` returns a copy, not a view of the array. I need to use `ravel()` with write access or a different approach.

The solution is to either:
1. Use `np.reshape(x, -1)` and work with that
2. Modify the array in-place using indexing
3. Work with views properly

Let me fix the code to use proper in-place modification.

Found the bug! `flatten()` returns a **copy**, not a view. Let me fix this:

edit

    if total_elements <= max_elements:
        # Full gradient check
        num_grad = np.zeros_like(x)
        # Use reshape to get a view, not a copy
        x_flat = x.reshape(-1)
        num_grad_flat = num_grad.reshape(-1)

        for i in range(total_elements):
            if (i + 1) % 10000 == 0:
                print(f"    Progress: {i+1}/{total_elements}")

            original = x_flat[i]

            # f(x + h)
            x_flat[i] = original + h
            y_plus, _ = layer_norm_forward(x, gamma, beta)
            loss_plus = np.sum(y_plus * dy)

            # f(x - h)
            x_flat[i] = original - h
            y_minus, _ = layer_norm_forward(x, gamma, beta)
            loss_minus = np.sum(y_minus * dy)

            # Central difference
            num_grad_flat[i] = (loss_plus - loss_minus) / (2 * h)

            # Restore
            x_flat[i] = original

        # Restore x to original values
        x[:] = orig_x

        return num_grad, np.ones((B, T, D), dtype=bool)  # All elements checked

    if total_elements <= max_elements:
        # Full gradient check
        num_grad = np.zeros_like(x)
        x_flat = x.flatten()
        num_grad_flat = num_grad.flatten()

        for i in range(total_elements):
            if (i + 1) % 10000 == 0:
                print(f"    Progress: {i+1}/{total_elements}")

            original = x_flat[i]

            # f(x + h)
            x_flat[i] = original + h
            y_plus, _ = layer_norm_forward(x, gamma, beta)
            loss_plus = np.sum(y_plus * dy)

            # f(x - h)
            x_flat[i] = original - h
            y_minus, _ = layer_norm_forward(x, gamma, beta)
            loss_minus = np.sum(y_minus * dy)

            # Central difference
            num_grad_flat[i] = (loss_plus - loss_minus) / (2 * h)

            # Restore
            x_flat[i] = original

        # Restore x to original values
        x[:] = orig_x

        return num_grad, np.ones((B, T, D), dtype=bool)  # All elements checked


Successfully replaced 1 block(s) in ./minimax-m2.7/backwards/layer_norm_numpy.py.

     ...
 419
 420     if total_elements <= max_elements:
 421         # Full gradient check
 422         num_grad = np.zeros_like(x)
-423         x_flat = x.flatten()
-424         num_grad_flat = num_grad.flatten()
+423         # Use reshape to get a view, not a copy
+424         x_flat = x.reshape(-1)
+425         num_grad_flat = num_grad.reshape(-1)
 425
 426         for i in range(total_elements):
 427             if (i + 1) % 10000 == 0:
 428                 print(f"    Progress: {i+1}/{total_elements}")
     ...


Let me verify that `reshape` returns a view and not a copy.

bash

cd ./minimax-m2.7/backwards && python3 -c "
import numpy as np

x = np.array([[1, 2], [3, 4]])
print(f'Original x: {x}')
print(f'x.ctypes.data: {x.ctypes.data}')

x_flat = x.flatten()
print(f'x_flat is view: {x_flat.base is x}')
print(f'x_flat.ctypes.data: {x_flat.ctypes.data}')

x_flat2 = x.reshape(-1)
print(f'x_flat2 is view: {x_flat2.base is x}')
print(f'x_flat2.ctypes.data: {x_flat2.ctypes.data}')

# Test in-place modification
x_flat2[0] = 99
print(f'After x_flat2[0] = 99: {x}')
"


Original x: [[1 2]
 [3 4]]
x.ctypes.data: 105553122249728
x_flat is view: False
x_flat.ctypes.data: 105553122250176
x_flat2 is view: True
x_flat2.ctypes.data: 105553122249728
After x_flat2[0] = 99: [[99  2]
 [ 3  4]]


So `flatten()` makes a copy, but `reshape(-1)` gives a view. However, when I modify `x_flat2[0]`, it should modify `x` since they're the same underlying data. Let me investigate more:

cd ./minimax-m2.7/backwards && python3 -c "
import numpy as np

x = np.array([[1, 2], [3, 4]])
print(f'Original x: {x}')

# Use reshape instead of flatten
x_flat = x.reshape(-1)
print(f'x_flat is a view: {x_flat.base is x}')

# Modify
x_flat[0] = 99
print(f'After x_flat[0] = 99, x is: {x}')

# Reset
x[0, 0] = 1

# Now test the loop
h = 0.01
x_flat = x.reshape(-1)
y_original, _ = None, None

# Get original y first
def f(x_input):
    B, T, D = 2, 2, 2  # dummy
    mean = np.mean(x_input, axis=-1, keepdims=True)
    x_centered = x_input - mean
    var = np.mean(x_centered ** 2, axis=-1, keepdims=True)
    std = np.sqrt(var + 1e-8)
    x_norm = x_centered / std
    y = x_norm  # identity gamma/beta
    return y

dy = np.ones_like(x)  # dummy gradient

# f(x)
y, _ = f(x), None
loss = np.sum(y * dy)
print(f'Original loss: {loss}')
print(f'Original x[0,0]: {x[0,0]}')

# f(x + h) for i=0
original = x_flat[0]
x_flat[0] = original + h
y_plus, _ = f(x), None
loss_plus = np.sum(y_plus * dy)
print(f'x after +h: {x}')
print(f'Loss plus: {loss_plus}')

# f(x - h) for i=0
x_flat[0] = original - h
y_minus, _ = f(x), None
loss_minus = np.sum(y_minus * dy)
print(f'x after -h: {x}')
print(f'Loss minus: {loss_minus}')

# Restore
x_flat[0] = original
print(f'x restored: {x}')
"


Original x: [[1 2]
 [3 4]]
x_flat is a view: True
After x_flat[0] = 99, x is: [[99  2]
 [ 3  4]]
Original loss: 0.0
Original x[0,0]: 1
x after +h: [[1 2]
 [3 4]]
Loss plus: 0.0
x after -h: [[0 2]
 [3 4]]
Loss minus: 0.0
x restored: [[1 2]
 [3 4]]


Let me look at the actual code in the file to see what's happening:

cd ./minimax-m2.7/backwards && python3 -c "
import numpy as np

def layer_norm_forward(x, gamma, beta, eps=1e-8):
    B, T, D = x.shape
    mean = np.mean(x, axis=-1, keepdims=True)
    x_centered = x - mean
    var = np.mean(x_centered ** 2, axis=-1, keepdims=True)
    std = np.sqrt(var + eps)
    x_norm = x_centered / std
    y = gamma * x_norm + beta
    return y, {'mean': mean, 'std': std, 'x_norm': x_norm, 'x_centered': x_centered}

# Reproduce the issue
np.random.seed(42)
B, T, D = 2, 4, 8
x = np.random.randn(B, T, D).astype(np.float64)
gamma = np.ones(D).astype(np.float64)
beta = np.zeros(D).astype(np.float64)
dy = np.random.randn(B, T, D).astype(np.float64)

# Save original
orig_x = x.copy()
print(f'Original x[0,0,0]: {x[0,0,0]:.6f}')

# Use reshape instead of flatten
x_flat = x.reshape(-1)
print(f'x_flat is view: {x_flat.base is x}')

# Step 1: f(x + h) for i=0
original = x_flat[0]
x_flat[0] = original + 1e-5
print(f'After +h, x[0,0,0]: {x[0,0,0]:.6f}')

y_plus, _ = layer_norm_forward(x, gamma, beta)
loss_plus = np.sum(y_plus * dy)
print(f'Loss plus: {loss_plus:.6f}')

# Step 2: f(x - h) for i=0
x_flat[0] = original - 1e-5
print(f'After -h, x[0,0,0]: {x[0,0,0]:.6f}')

y_minus, _ = layer_norm_forward(x, gamma, beta)
loss_minus = np.sum(y_minus * dy)
print(f'Loss minus: {loss_minus:.6f}')

# Restore
x_flat[0] = original
print(f'After restore, x[0,0,0]: {x[0,0,0]:.6f}')

# Now check if x is restored
print(f'\\nVerifying x matches orig_x: {np.allclose(x, orig_x)}')
print(f'Max diff: {np.max(np.abs(x - orig_x))}')

# Compute gradient
d_num = (loss_plus - loss_minus) / (2 * 1e-5)
print(f'\\nNumerical gradient for element 0: {d_num:.6f}')

# Compare with analytical
y, cache = layer_norm_forward(orig_x, gamma, beta)
std = cache['std']
x_norm = cache['x_norm']
dz = dy  # since gamma = 1
sum_dz = np.sum(dz, axis=-1, keepdims=True)
sum_dz_xnorm = np.sum(dz * x_norm, axis=-1, keepdims=True)
dx_analytical = (dz - sum_dz / D - x_norm * sum_dz_xnorm / D) / std
print(f'Analytical gradient for element 0: {dx_analytical[0,0,0]:.6f}')
"


Original x[0,0,0]: 0.496714
x_flat is view: True
After +h, x[0,0,0]: 0.496724
Loss plus: -3.607269
After -h, x[0,0,0]: 0.496704
Loss minus: -3.607276
After restore, x[0,0,0]: 0.496714

Verifying x matches orig_x: True
Max diff: 0.0

Numerical gradient for element 0: 0.344094
Analytical gradient for element 0: 0.344094


This is working correctly. Let me now run the full test to see if the gradient check passes with the reshape fix.

cd ./minimax-m2.7/backwards && python layer_norm_numpy.py


======================================================================
LAYER NORMALIZATION: FORWARD AND BACKWARD FROM SCRATCH
======================================================================

╔══════════════════════════════════════════════════════════════════════════════╗
║                    COMPLEXITY ANALYSIS                                         ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Input: x ∈ ℝ^(B×T×D)                                                        ║
║  Parameters: γ, β ∈ ℝ^D                                                      ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FORWARD PASS                                                                ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Operation              │ Work (FLOPs)      │ Memory                        ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  mean (reduction)       │ O(BTD)            │ -                             ║
║  x - mean (broadcast)   │ O(BTD)            │ O(BTD)                        ║
║  var (reduction)        │ O(BTD)            │ -                             ║
║  sqrt(var + eps)        │ O(BT)             │ -                             ║
║  divide by std          │ O(BTD)            │ -                             ║
║  gamma * x_norm         │ O(BTD)            │ -                             ║
║  add beta               │ O(BTD)            │ O(BTD) output                 ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  TOTAL                    │ 5×O(BTD)          │ O(BTD)                       ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  BACKWARD PASS                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Operation              │ Work (FLOPs)      │ Memory                        ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  d_gamma = sum(dy*x_norm)│ O(BTD)           │ O(D)                          ║
║  d_beta = sum(dy)       │ O(BTD)            │ O(D)                          ║
║  dz = dy * gamma        │ O(BTD)            │ O(BTD) (can be avoided)        ║
║  sum(dz)                │ O(BTD)            │ -                             ║
║  sum(dz * x_norm)       │ O(BTD)            │ -                             ║
║  dx computation         │ O(BTD)            │ O(BTD) output                 ║
║  ────────────────────────────────────────────────────────────────────────── ║
║  TOTAL                    │ 5×O(BTD)          │ O(BTD) + O(D)                ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  SUMMARY                                                                    ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Time Complexity:  O(BTD) for both forward and backward                      ║
║  Space Complexity: O(BTD) for storing activations (during training)          ║
║                   O(BTD) during inference (no need to store)                 ║
║                                                                              ║
║  Cache efficiency: We store x_centered, x_norm, mean, std                    ║
║                   These are O(BTD) total, reused across all gradient comps  ║
║                                                                              ║
╚══════════════════════════════════════════════════════════════════════════════╝

╔══════════════════════════════════════════════════════════════════════════════╗
║               GPU KERNEL FUSION DESIGN                                       ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  CURRENT SEPARATE KERNEL APPROACH:                                           ║
║  ────────────────────────────────────                                        ║
║  Kernel 1: Compute mean (reduction over D)                                   ║
║  Kernel 2: Compute variance (reduction over D)                               ║
║  Kernel 3: Normalize (element-wise)                                          ║
║  Kernel 4: Scale and shift (element-wise)                                    ║
║  Kernel 5: Backward kernels (x, gamma, beta)                                 ║
║                                                                              ║
║  Issues with separate kernels:                                               ║
║  • Multiple kernel launches (overhead)                                       ║
║  • Data movement between global memory passes                                 ║
║  • Can't use persistent threads for reduction efficiency                      ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FUSED KERNEL DESIGN (Forward):                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Grid: (B × T) blocks, each handling one (b,t) position                      ║
║  Block: 256-512 threads                                                      ║
║                                                                              ║
║  PHASE 1: Load and compute local sum                                         ║
║  ─────────────────────────────────                                           ║
║  • Each thread loads x[b,t,d] into shared memory                             ║
║  • Compute partial sum using warp-level reduction                            ║
║  • Single thread writes mean to __shared__                                    ║
║                                                                              ║
║  PHASE 2: Compute variance locally                                           ║
║  ─────────────────────────────────                                           ║
║  • Re-load x with loaded mean                                                ║
║  • Compute (x-mean)² and partial variance                                    ║
║  • Reduce to get variance                                                   ║
║                                                                              ║
║  PHASE 3: Normalize and output                                               ║
║  ─────────────────────────────────                                           ║
║  • All threads compute: y = gamma * (x-mean) / sqrt(var+eps) + beta          ║
║  • Write to output (fully coalesced)                                         ║
║                                                                              ║
║  MEMORY ACCESS PATTERN:                                                      ║
║  • Each block reads contiguous D elements (coalesced)                        ║
║  • Use shared memory for intermediate results                                ║
║  • Output writes are also coalesced                                          ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  FUSED KERNEL DESIGN (Backward):                                             ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  Key insight: dz = dy * gamma can be merged into the computation              ║
║                                                                              ║
║  Grid: (B × T) blocks                                                       ║
║  Block: 256-512 threads                                                      ║
║                                                                              ║
║  SHARED MEMORY STRUCTURE:                                                    ║
║  [x_norm_0, x_norm_1, ..., x_norm_{D-1}]                                     ║
║  [dz_0, dz_1, ..., dz_{D-1}]                                                 ║
║                                                                              ║
║  ALGORITHM:                                                                  ║
║  1. Load x_norm and compute local dz = dy * gamma                            ║
║  2. Reduce to get sum(dz) and sum(dz * x_norm)                              ║
║  3. Second pass to compute dx using the formula:                             ║
║     dx = (dz - mean(dz) - x_norm * mean(dz*x_norm)) / std                    ║
║                                                                              ║
║  REDUCTION OPTIMIZATIONS:                                                    ║
║  • Warp-level shuffle reductions (no shared memory needed)                    ║
║  • Block-level using shared memory with tree reduction                        ║
║  • Use block-level primitives for final reduction                            ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  BENEFITS OF FUSION:                                                          ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  ✓ Reduced kernel launch overhead (1 kernel vs 4-5)                           ║
║  ✓ Better memory bandwidth utilization (single read, single write)          ║
║  ✓ Improved cache locality (data stays in registers/shared mem)              ║
║  ✓ Only loads x once, computes mean and var from same data                   ║
║  ✓ Backward can reuse cached values from forward (if memory allows)          ║
║  ✓ Lower register pressure allows for larger block sizes                     ║
║                                                                              ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  CUDA KERNEL SKETCH (Pseudo-code):                                            ║
╠══════════════════════════════════════════════════════════════════════════════╣
║                                                                              ║
║  __global__ void layer_norm_fwd(float* y, const float* x,                    ║
║                                  const float* gamma, const float* beta,      ║
║                                  int B, int T, int D, float eps) {          ║
║                                                                              ║
║      __shared__ float mean_smem[256]; // block-level mean                     ║
║      __shared__ float var_smem[256];  // block-level variance                 ║
║      __shared__ float std_smem[256]; // block-level std                      ║
║                                                                              ║
║      int tid = threadIdx.x;                                                  ║
║      int bid = blockIdx.x;                                                   ║
║      int D_blk = (D + blockDim.x - 1) / blockDim.x;                          ║
║                                                                              ║
║      // Phase 1: Load and compute mean                                        ║
║      float sum = 0.0;                                                        ║
║      for (int i = 0; i < D_blk; i++) {                                       ║
║          int idx = bid * D + i * blockDim.x + tid;                           ║
║          sum += x[idx];                                                      ║
║      }                                                                       ║
║      sum = warpReduceSum(sum);                                               ║
║      if (tid % 32 == 0) mean_smem[tid / 32] = sum;                           ║
║      __syncthreads();                                                        ║
║                                                                              ║
║      float mean = mean_smem[0] / D;                                         ║
║                                                                              ║
║      // Phase 2: Compute variance                                            ║
║      sum = 0.0;                                                             ║
║      for (int i = 0; i < D_blk; i++) {                                       ║
║          int idx = bid * D + i * blockDim.x + tid;                           ║
║          float diff = x[idx] - mean;                                         ║
║          sum += diff * diff;                                                 ║
║      }                                                                       ║
║      sum = warpReduceSum(sum);                                               ║
║      if (tid % 32 == 0) var_smem[tid / 32] = sum;                           ║
║      __syncthreads();                                                       ║
║                                                                              ║
║      float var = var_smem[0] / D;                                          ║
║      float std = sqrt(var + eps);                                           ║
║                                                                              ║
║      // Phase 3: Normalize and output                                        ║
║      for (int i = 0; i < D_blk; i++) {                                       ║
║          int idx = bid * D + i * blockDim.x + tid;                            ║
║          float x_norm = (x[idx] - mean) / std;                               ║
║          y[idx] = gamma[idx % D] * x_norm + beta[idx % D];                  ║
║      }                                                                       ║
║  }                                                                           ║
║                                                                              ║
╚══════════════════════════════════════════════════════════════════════════════╝

======================================================================
RUNNING GRADIENT CHECKS
======================================================================

============================================================
Testing shape (B=2, T=4, D=8)
============================================================

============================================================
GRADIENT CHECK: gamma
============================================================
Analytical gradient shape: (8,)
Numerical gradient shape: (8,)
Max absolute difference: 8.55e-11
Max relative tolerance: 4.55e-04
Mean analytical: 2.429868e+00
Mean numerical: 2.429868e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: beta
============================================================
Analytical gradient shape: (8,)
Numerical gradient shape: (8,)
Max absolute difference: 1.08e-10
Max relative tolerance: 5.13e-04
Mean analytical: 2.602642e+00
Mean numerical: 2.602642e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: x (input)
============================================================
Analytical gradient shape: (2, 4, 8)
Numerical gradient shape: (2, 4, 8)
Elements checked: 64 / 64
Max absolute difference: 1.90e-10
Max relative tolerance: 2.24e-04
Mean analytical (checked): 5.202041e-01
Mean numerical (checked): 5.202041e-01

Gradient check: PASSED ✓
     gamma: ✓ PASS
      beta: ✓ PASS
         x: ✓ PASS

============================================================
Testing shape (B=4, T=8, D=16)
============================================================

============================================================
GRADIENT CHECK: gamma
============================================================
Analytical gradient shape: (16,)
Numerical gradient shape: (16,)
Max absolute difference: 5.52e-10
Max relative tolerance: 1.28e-03
Mean analytical: 4.142521e+00
Mean numerical: 4.142521e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: beta
============================================================
Analytical gradient shape: (16,)
Numerical gradient shape: (16,)
Max absolute difference: 2.52e-10
Max relative tolerance: 8.59e-04
Mean analytical: 3.506940e+00
Mean numerical: 3.506940e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: x (input)
============================================================
Analytical gradient shape: (4, 8, 16)
Numerical gradient shape: (4, 8, 16)
Elements checked: 512 / 512
Max absolute difference: 4.70e-10
Max relative tolerance: 4.70e-04
Mean analytical (checked): 6.141263e-01
Mean numerical (checked): 6.141263e-01

Gradient check: PASSED ✓
     gamma: ✓ PASS
      beta: ✓ PASS
         x: ✓ PASS

============================================================
Testing shape (B=8, T=16, D=32)
============================================================

============================================================
GRADIENT CHECK: gamma
============================================================
Analytical gradient shape: (32,)
Numerical gradient shape: (32,)
Max absolute difference: 2.38e-09
Max relative tolerance: 2.84e-03
Mean analytical: 9.330259e+00
Mean numerical: 9.330259e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: beta
============================================================
Analytical gradient shape: (32,)
Numerical gradient shape: (32,)
Max absolute difference: 2.00e-09
Max relative tolerance: 2.61e-03
Mean analytical: 9.296132e+00
Mean numerical: 9.296132e+00

Gradient check: PASSED ✓

============================================================
GRADIENT CHECK: x (input)
============================================================
Analytical gradient shape: (8, 16, 32)
Numerical gradient shape: (8, 16, 32)
Elements checked: 4096 / 4096
Max absolute difference: 2.37e-09
Max relative tolerance: 6.93e-04
Mean analytical (checked): 7.088609e-01
Mean numerical (checked): 7.088609e-01

Gradient check: PASSED ✓
     gamma: ✓ PASS
      beta: ✓ PASS
         x: ✓ PASS

======================================================================
ALL GRADIENT CHECKS PASSED ✓
======================================================================

======================================================================
BENCHMARKING LAYER NORMALIZATION
======================================================================

Shape: (B=32, T=128, D=256)
----------------------------------------
Forward:  5.958 ± 0.528 ms (0.2 GB/s)
Backward: 6.373 ± 0.948 ms (0.2 GB/s)
Total:    12.331 ms

Shape: (B=64, T=128, D=512)
----------------------------------------
Forward:  26.879 ± 1.883 ms (0.2 GB/s)
Backward: 37.461 ± 4.092 ms (0.1 GB/s)
Total:    64.340 ms

Shape: (B=32, T=512, D=768)
----------------------------------------
Forward:  85.124 ± 6.379 ms (0.1 GB/s)
Backward: 108.899 ± 13.728 ms (0.1 GB/s)
Total:    194.023 ms

Shape: (B=16, T=512, D=1024)
----------------------------------------
Forward:  46.776 ± 1.119 ms (0.2 GB/s)
Backward: 62.149 ± 1.049 ms (0.1 GB/s)
Total:    108.925 ms


All gradient checks are now passing. Let me provide a summary of the implementation.

All gradient checks pass. The implementation is complete. Here's a summary:

## Summary

**File:** `./minimax-m2.7/backwards/layer_norm_numpy.py`

### Key Components

1. **Forward Pass** (`layer_norm_forward`):
   - Computes mean, variance, and standard deviation over the feature dimension (D)
   - Normalizes: `x_norm = (x - mean) / sqrt(var + eps)`
   - Scales and shifts: `y = gamma * x_norm + beta`
   - Caches intermediates: `x_centered`, `x_norm`, `mean`, `std`, `gamma`

2. **Backward Pass** (`layer_norm_backward`):
   - **d_gamma**: `sum(dy * x_norm)` over (B, T)
   - **d_beta**: `sum(dy)` over (B, T)
   - **dx**: `(dz - mean(dz) - x_norm * mean(dz * x_norm)) / std`

   Where `dz = dy * gamma`. The consolidated formula avoids computing the full Jacobian.

3. **Gradient Check**:
   - Uses central finite differences: `(f(x+h) - f(x-h)) / 2h`
   - Fixed array flattening to use views instead of copies
   - Spot-check for large tensors (>100k elements)

4. **Numerical Stability**:
   - `eps` prevents division by zero: `std = sqrt(var + eps)`
   - Catastrophic cancellation analysis documented
   - Pairwise summation suggestion for fp16

5. **Complexity**:
   - **Time**: O(BTD) for both forward and backward
   - **Space**: O(BTD) for activations (training), O(1) for inference

6. **GPU Fusion Design**:
   - Grid of `(B × T)` blocks, each handling one `(b,t)` position
   - Three phases: mean reduction, variance reduction, normalize+output
   - Warp-level shuffle reductions for efficiency
   - Single kernel replaces 4-5 separate kernels