kgrabko commited on
Commit
531dee3
·
verified ·
1 Parent(s): b5299b4

Upload 3 files

Browse files
BRE_memory_routing.md ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Buffered Routing Embedding (BRE) Algorithm
2
+ **Inventor:** Konstantin Vladimirovich Grabko
3
+
4
+ ### Problem Statement
5
+ Ultra-scale models (140B+) suffer from "Memory Wall" bottlenecks where the GPU waits for embedding weights to be fetched from HBM.
6
+
7
+ ### The BRE Solution
8
+ BRE implements a predictive pre-fetching ring buffer.
9
+ 1. **Token Prediction Window:** A lightweight heuristic monitors the last $N$ tokens to predict high-probability future embeddings.
10
+ 2. **HBM Routing:** Predicted weights are moved from standard HBM to a specialized "High-Speed Buffer" partition (L3 Cache/Shared Memory) *before* the attention computation begins.
11
+ 3. **Synchronous Paging:** BRE uses Peer-to-Peer (P2P) DMA transfers across the ROCm/Infinity Fabric to ensure that weights for the next 4 layers are already local to the current GPU.
JiRackPyTorch_GPT5_class_13b.py ADDED
@@ -0,0 +1,144 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # ==============================================================================
2
+ # COPYRIGHT (C) 2025 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
3
+ # PATENT PENDING | CMS MANHATTAN JIRACK TECHNOLOGY
4
+ #
5
+ # This software is licensed under the Commercial License Agreement V.1.2.
6
+ # Any use, modification, or distribution of this code requires compliance with
7
+ # the terms found in the LICENSE.md file in the root directory.
8
+ #
9
+ # NO PATENTING RIGHTS: Users are strictly prohibited from filing patent claims
10
+ # based on the BRE or SWA architectures disclosed herein.
11
+ # Contact: grabko@cmsmanhattan.com | +1 (516) 777-0945
12
+ # ==============================================================================
13
+ # Version: 13B Scale-Up (GQA + SwiGLU + RoPE)
14
+
15
+ import torch
16
+ import torch.nn as nn
17
+ import torch.nn.functional as F
18
+ import math
19
+
20
+ # --- КОНФИГУРАЦИЯ 13B ---
21
+ VOCAB_SIZE = 50257
22
+ MODEL_DIM = 5120 # Масштаб 13B
23
+ NUM_HEADS = 40 # Головы для Queries
24
+ NUM_KV_HEADS = 8 # GQA (коэффициент 5:1)
25
+ NUM_LAYERS = 40 # Глубина 13B
26
+ MAX_SEQ_LEN = 2048
27
+ FFN_HIDDEN_DIM = 13824 # SwiGLU DIM для 13B
28
+ HEAD_DIM = MODEL_DIM // NUM_HEADS
29
+ EPSILON = 1e-5
30
+ WINDOW_SIZE = 1024
31
+
32
+ class RMSNorm(nn.Module):
33
+ def __init__(self, dim, eps=EPSILON):
34
+ super().__init__()
35
+ self.eps = eps
36
+ self.weight = nn.Parameter(torch.ones(dim))
37
+ def forward(self, x):
38
+ return (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)) * self.weight
39
+
40
+ def precompute_freqs_cis(dim, seq_len, theta=10000.0):
41
+ freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
42
+ t = torch.arange(seq_len)
43
+ freqs = torch.outer(t, freqs).float()
44
+ return torch.polar(torch.ones_like(freqs), freqs)
45
+
46
+ def apply_rotary_emb(xq, xk, freqs_cis):
47
+ xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
48
+ xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
49
+ freqs_cis = freqs_cis.view(1, xq_.size(1), 1, xq_.size(3))
50
+ xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
51
+ xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
52
+ return xq_out.type_as(xq), xk_out.type_as(xk)
53
+
54
+ def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
55
+ if n_rep == 1:
56
+ return x
57
+ bs, slen, n_kv_heads, head_dim = x.shape
58
+ return (
59
+ x[:, :, :, None, :]
60
+ .expand(bs, slen, n_kv_heads, n_rep, head_dim)
61
+ .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
62
+ )
63
+
64
+ class MultiHeadAttention(nn.Module):
65
+ def __init__(self):
66
+ super().__init__()
67
+ self.n_kv_heads = NUM_KV_HEADS
68
+ self.n_rep = NUM_HEADS // NUM_KV_HEADS
69
+
70
+ self.wq = nn.Linear(MODEL_DIM, NUM_HEADS * HEAD_DIM, bias=False)
71
+ self.wk = nn.Linear(MODEL_DIM, NUM_KV_HEADS * HEAD_DIM, bias=False)
72
+ self.wv = nn.Linear(MODEL_DIM, NUM_KV_HEADS * HEAD_DIM, bias=False)
73
+ self.wo = nn.Linear(NUM_HEADS * HEAD_DIM, MODEL_DIM, bias=False)
74
+
75
+ def forward(self, x, freqs_cis, past_kv=None):
76
+ b, t, _ = x.shape
77
+ q, k, v = self.wq(x), self.wk(x), self.wv(x)
78
+
79
+ q = q.view(b, t, NUM_HEADS, HEAD_DIM)
80
+ k = k.view(b, t, self.n_kv_heads, HEAD_DIM)
81
+ v = v.view(b, t, self.n_kv_heads, HEAD_DIM)
82
+
83
+ q, k = apply_rotary_emb(q, k, freqs_cis[:t])
84
+
85
+ if past_kv is not None:
86
+ pk, pv = past_kv
87
+ k = torch.cat([pk, k], dim=1)
88
+ v = torch.cat([pv, v], dim=1)
89
+ if k.size(1) > WINDOW_SIZE:
90
+ k, v = k[:, -WINDOW_SIZE:], v[:, -WINDOW_SIZE:]
91
+
92
+ current_kv = (k.detach(), v.detach())
93
+ k = repeat_kv(k, self.n_rep)
94
+ v = repeat_kv(v, self.n_rep)
95
+
96
+ q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
97
+ out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
98
+ return self.wo(out.transpose(1, 2).contiguous().view(b, t, MODEL_DIM)), current_kv
99
+
100
+ class SwiGLU(nn.Module):
101
+ def __init__(self):
102
+ super().__init__()
103
+ self.w1 = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
104
+ self.w2 = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
105
+ self.w3 = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
106
+ def forward(self, x):
107
+ return self.w2(F.silu(self.w1(x)) * self.w3(x))
108
+
109
+ class JiRackPyTorch(nn.Module):
110
+ def __init__(self):
111
+ super().__init__()
112
+ self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
113
+ self.blocks = nn.ModuleList([nn.ModuleDict({
114
+ 'norm1': RMSNorm(MODEL_DIM),
115
+ 'attn': MultiHeadAttention(),
116
+ 'norm2': RMSNorm(MODEL_DIM),
117
+ 'ffn': SwiGLU()
118
+ }) for _ in range(NUM_LAYERS)])
119
+ self.norm_f = RMSNorm(MODEL_DIM)
120
+ self.head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
121
+ self.head.weight = self.token_emb.weight
122
+
123
+ self.register_buffer("freqs_cis", precompute_freqs_cis(HEAD_DIM, MAX_SEQ_LEN * 2))
124
+
125
+ signature = "Author: Konstantin Vladimirovich Grabko (CMS Manhattan) 2025"
126
+ self.register_buffer("proof_of_authorship", torch.tensor([ord(c) for c in signature], dtype=torch.uint8))
127
+
128
+ def get_author_info(self):
129
+ return "".join([chr(c) for c in self.proof_of_authorship.tolist()])
130
+
131
+ def forward(self, idx, targets=None, past_kv=None):
132
+ x = self.token_emb(idx)
133
+ new_kvs = []
134
+ for i, block in enumerate(self.blocks):
135
+ h, kv = block['attn'](block['norm1'](x), self.freqs_cis, past_kv[i] if past_kv else None)
136
+ x = x + h
137
+ x = x + block['ffn'](block['norm2'](x))
138
+ new_kvs.append(kv)
139
+ x = self.norm_f(x)
140
+ logits = self.head(x)
141
+ if targets is not None:
142
+ loss = F.cross_entropy(logits.view(-1, VOCAB_SIZE), targets.view(-1))
143
+ return logits, loss, new_kvs
144
+ return logits, new_kvs
SWA_fusion_spec.md ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Technical Specification: SwiGLU-Attention (SWA) Fusion Kernel
2
+ **Document ID:** CMS-SWA-2025-001
3
+ **Status:** Proprietary / High Performance
4
+
5
+ ### Abstract
6
+ In standard transformers, the Multi-Head Attention (MHA) and Feed-Forward Network (FFN) are executed as discrete kernels, forcing multiple round-trips to Global Memory (VRAM). SWA Fusion merges these into a single computational pass.
7
+
8
+ ### Computational Logic
9
+ The kernel pipelines the $Q, K, V$ projections simultaneously with the $W_1$ and $W_3$ projections of the SwiGLU layer.
10
+ 1. **Shared Input Latches:** Input $x$ is loaded into Shared Memory (SRAM) once.
11
+ 2. **Parallel Projections:** $$Y_{attn} = \text{Softmax}(\frac{QK^T}{\sqrt{d_k}})V$$
12
+ $$Y_{ffn} = (SiLU(xW_1) \otimes xW_3)W_2$$
13
+ 3. **Fused Accumulation:** $x_{out} = x + Y_{attn} + Y_{ffn}$ is computed in a single thread block before writing back to HBM.
14
+
15
+
16
+
17
+ ### Performance Target
18
+ - **Memory Bandwidth Reduction:** 30% lower VRAM traffic.
19
+ - **Hardware Target:** Optimized for AMD CDNA3 (MI300X) Matrix Cores.