Version: 0.10.16-alpha.2

Refactor DNN/transformer stack: BuildContext, chunked prefill

- Major refactor: explicit BuildContext, TrainingMode, buffer mgmt
- Enable efficient chunked prefill/inference in transformer/attention
- Replace std::vector shapes with fixed-capacity TensorShape types
- Add modular C++/CUDA attention, positional, and tensor op interfaces

BREAKING CHANGE:
Updates all component, tensor, and operator APIs to use BuildContext, TrainingMode, and shape_t;
removes legacy interfaces and changes model/component build and training lifecycle. Existing code must be updated for new APIs.

Author: ToddThomson

CMakePresets.json

@@ -67,7 +67,8 @@
         "strategy": "external"
       },
       "cacheVariables": {
-        "CMAKE_BUILD_TYPE": "Debug"
+        "CMAKE_BUILD_TYPE": "Debug",
+        "CMAKE_CUDA_FLAGS_DEBUG": "-G"
       }
     },
     {

Data/Scripts/Gpt2/gpt2_output.py

@@ -0,0 +1,68 @@
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
+model = GPT2LMHeadModel.from_pretrained('gpt2')
+model.eval()
+tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+
+# Same input as your Mila test
+input_text = "Once upon a time"
+input_ids = tokenizer.encode(input_text, return_tensors='pt')
+print(f"Token ids: {input_ids}")
+
+# Storage for hook outputs
+captured = {}
+
+def make_hook(name):
+    def fn(module, input, output):
+        t = output if isinstance(output, torch.Tensor) else output[0]
+        captured[name] = (t.min().item(), t.max().item(), t.shape)
+    return fn
+
+hooks = []
+
+# Encoder output (wte + wpe)
+hooks.append(model.transformer.drop.register_forward_hook(make_hook('lenc_out')))
+
+for i, block in enumerate(model.transformer.h):
+    # ln_1
+    hooks.append(block.ln_1.register_forward_hook(make_hook(f'layer_{i}.ln_1')))
+    # fc_qkv_proj (c_attn)
+    hooks.append(block.attn.c_attn.register_forward_hook(make_hook(f'layer_{i}.fc_qkv_proj')))
+    # fc_out_proj (c_proj)
+    hooks.append(block.attn.c_proj.register_forward_hook(make_hook(f'layer_{i}.fc_out_proj')))
+    # ln_2
+    hooks.append(block.ln_2.register_forward_hook(make_hook(f'layer_{i}.ln_2')))
+    # mlp.fc_1 (c_fc)
+    hooks.append(block.mlp.c_fc.register_forward_hook(make_hook(f'layer_{i}.mlp.fc_1')))
+    # gelu
+    hooks.append(block.mlp.act.register_forward_hook(make_hook(f'layer_{i}.mlp.gelu')))
+    # mlp.fc_2 (c_proj)
+    hooks.append(block.mlp.c_proj.register_forward_hook(make_hook(f'layer_{i}.mlp.fc_2')))
+    # full block output (residual)
+    hooks.append(block.register_forward_hook(make_hook(f'layer_{i}.residual_out')))
+
+with torch.no_grad():
+    model(input_ids)
+
+for h in hooks:
+    h.remove()
+
+# Print in same format as your Mila debug output
+print(f"\nlenc out: [{captured['lenc_out'][0]:.3f}, {captured['lenc_out'][1]:.3f}]")
+print()
+
+for i in range(12):
+    for key in [f'layer_{i}.ln_1', f'layer_{i}.fc_qkv_proj', f'layer_{i}.fc_out_proj',
+                f'layer_{i}.ln_2', f'layer_{i}.mlp.fc_1', f'layer_{i}.mlp.gelu',
+                f'layer_{i}.mlp.fc_2', f'layer_{i}.residual_out']:
+        if key in captured:
+            mn, mx, shape = captured[key]
+            print(f"{key}: [{mn:.3f}, {mx:.3f}] shape={list(shape)}")
+    print()
+
+# Print peak residual across all layers to set kResidualAbsLimit
+peak = max(abs(captured[f'layer_{i}.residual_out'][1]) for i in range(12))
+peak_min = min(captured[f'layer_{i}.residual_out'][0] for i in range(12))
+print(f"Peak residual: min={peak_min:.3f}, max={peak:.3f}")
+print(f"Suggested kResidualAbsLimit: {peak * 1.5:.1f}")

Data/Scripts/Gpt2/gpt2_tokenizer_test.py

@@ -0,0 +1,9 @@
+from transformers import GPT2Tokenizer
+
+tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+
+text = "You are a helpful AI Assistant. Your name is Mila"
+tokens = tokenizer.encode(text)
+print(f"Token count: {len(tokens)}")
+print(f"Token IDs: {tokens}")
+print(f"\nDecoded back: '{tokenizer.decode(tokens)}'")

Data/Scripts/Gpt2/gpt2_weights_output.py

@@ -0,0 +1,9 @@
+import torch
+from transformers import GPT2LMHeadModel
+
+model = GPT2LMHeadModel.from_pretrained('gpt2')
+
+for i in [0, 1]:
+    w = model.state_dict()[f'transformer.h.{i}.mlp.c_proj.weight'].T
+    print(f"Layer {i} fc_2 after .T: min={w.min():.6f} max={w.max():.6f} mean={w.mean():.6f}")
+    print(f"First 5x5:\n{w[:5, :5]}\n")

Data/Scripts/Gpt2/hf_decode.py

@@ -0,0 +1,27 @@
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
+model     = GPT2LMHeadModel.from_pretrained( 'gpt2' )
+tokenizer = GPT2Tokenizer.from_pretrained( 'gpt2' )
+model.eval()
+
+input_ids = tokenizer.encode( "Once upon a time", return_tensors='pt' )
+print( f"Token ids: {input_ids.tolist()}" )
+
+with torch.no_grad():
+    out = model( input_ids )
+
+logits = out.logits[ 0 ]  # [T, V]
+
+print( "\n=== Token 11 logit at every position ===" )
+for pos in range( 4 ):
+    print( f"HF token 11 (',') at pos {pos}: {logits[ pos, 11 ].item():.4f}" )
+
+print( "\n=== Top token at every position ===" )
+for pos in range( 4 ):
+    top_token = logits[ pos ].argmax().item()
+    top_logit = logits[ pos, top_token ].item()
+    print( f"HF top token at pos {pos}: token={top_token} '{tokenizer.decode([top_token])}' logit={top_logit:.4f}" )
+
+print( f"\n=== Token 11 at pos 3 (expected ~-50.47) ===" )
+print( f"HF token 11 at pos 3: {logits[ 3, 11 ].item():.4f}" )

Data/Scripts/Gpt2/hf_greedy_validation.py

@@ -0,0 +1,9 @@
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+import torch
+
+model = GPT2LMHeadModel.from_pretrained("gpt2")
+tok   = GPT2Tokenizer.from_pretrained("gpt2")
+
+ids = tok.encode("Once upon a time", return_tensors="pt")
+out = model.generate(ids, max_new_tokens=64, do_sample=False)  # greedy
+print(tok.decode(out[0]))

Data/Scripts/hf_mila_equivalency.py Data/Scripts/Gpt2/hf_mila_equivalency.pyData/Scripts/hf_mila_equivalency.py renamed to Data/Scripts/Gpt2/hf_mila_equivalency.py

@@ -0,0 +1,63 @@
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
+model = GPT2LMHeadModel.from_pretrained( 'gpt2' )
+model.eval()
+tokenizer = GPT2Tokenizer.from_pretrained( 'gpt2' )
+
+# This matches Mila's decode step 0:
+# - Prefill: "Once upon a time" (4 tokens)
+# - First decoded token: ',' (token 11)
+input_text = "Once upon a time,"
+input_ids = tokenizer.encode( input_text, return_tensors='pt' )
+print( f"Token ids: {input_ids}" )
+print( f"Seq len: {input_ids.shape[1]}" )
+
+captured = {}
+
+def make_hook_last( name ):
+    """Captures min/max of the LAST token position only."""
+    def fn( module, input, output ):
+        t = output if isinstance( output, torch.Tensor ) else output[0]
+        last = t[0, -1, :]
+        captured[name] = ( last.min().item(), last.max().item() )
+    return fn
+
+hooks = []
+for i, block in enumerate( model.transformer.h ):
+    hooks.append( block.ln_1.register_forward_hook( make_hook_last( f'layer_{i}.ln_1' ) ) )
+    hooks.append( block.attn.c_attn.register_forward_hook( make_hook_last( f'layer_{i}.fc_qkv_proj' ) ) )
+    hooks.append( block.attn.c_proj.register_forward_hook( make_hook_last( f'layer_{i}.fc_out_proj' ) ) )
+    hooks.append( block.ln_2.register_forward_hook( make_hook_last( f'layer_{i}.ln_2' ) ) )
+    hooks.append( block.mlp.c_fc.register_forward_hook( make_hook_last( f'layer_{i}.mlp.fc_1' ) ) )
+    hooks.append( block.mlp.act.register_forward_hook( make_hook_last( f'layer_{i}.mlp.gelu' ) ) )
+    hooks.append( block.mlp.c_proj.register_forward_hook( make_hook_last( f'layer_{i}.mlp.fc_2' ) ) )
+    hooks.append( block.register_forward_hook( make_hook_last( f'layer_{i}.residual_out' ) ) )
+
+with torch.no_grad():
+    out = model( input_ids )
+    logits = out.logits[0, -1]
+    top5 = torch.topk( logits, 5 )
+    print( f"\nTop 5 predictions after '{input_text}':" )
+    for v, idx in zip( top5.values, top5.indices ):
+        print( f"  {tokenizer.decode([idx.item()])!r:15} {v.item():.4f}" )
+
+for h in hooks:
+    h.remove()
+
+print( "\n=== Per-layer values at LAST token position ===" )
+for i in range( 12 ):
+    print( f"\nlayer_{i}:" )
+    for key in [
+        f'layer_{i}.ln_1',
+        f'layer_{i}.fc_qkv_proj',
+        f'layer_{i}.fc_out_proj',
+        f'layer_{i}.ln_2',
+        f'layer_{i}.mlp.fc_1',
+        f'layer_{i}.mlp.gelu',
+        f'layer_{i}.mlp.fc_2',
+        f'layer_{i}.residual_out',
+    ]:
+        if key in captured:
+            mn, mx = captured[key]
+            print( f"  {key}: [{mn:.3f}, {mx:.3f}]" )

Data/Scripts/Gpt2/hf_mila_test.py

@@ -0,0 +1,51 @@
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
+model = GPT2LMHeadModel.from_pretrained('gpt2')
+model.eval()
+tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+
+input_text = "Once upon a time"
+input_ids = tokenizer.encode(input_text, return_tensors='pt')
+print(f"Token ids: {input_ids}")
+print(f"Seq len: {input_ids.shape[1]}")
+
+captured = {}
+
+def make_hook(name):
+    def fn(module, input, output):
+        captured[name] = output.detach()
+    return fn
+
+hooks = []
+hooks.append(model.transformer.h[0].attn.c_attn.register_forward_hook(make_hook('layer_0.qkv')))
+hooks.append(model.transformer.h[1].attn.c_attn.register_forward_hook(make_hook('layer_1.qkv')))
+
+with torch.no_grad():
+    model(input_ids)
+
+for h in hooks:
+    h.remove()
+
+def extract_v( qkv, layer_name, num_heads=12, head_size=64 ):
+    # qkv shape: [1, T, 2304]
+    T = qkv.shape[1]
+    q, k, v = qkv.split( 768, dim=-1 )  # each [1, T, 768]
+    # Reshape to [B, NH, T, HS]
+    v = v.view( 1, T, num_heads, head_size ).permute( 0, 2, 1, 3 )  # [1, 12, T, 64]
+    k = k.view( 1, T, num_heads, head_size ).permute( 0, 2, 1, 3 )  # [1, 12, T, 64]
+    print( f"\n=== {layer_name} ===" )
+    for head in range( 2 ):  # show head 0 and head 1
+        print( f"\n  V head {head}, positions 0..{T-1} (first 8 elements each):" )
+        for pos in range( T ):
+            vals = v[0, head, pos, :8].tolist()
+            formatted = "  ".join( f"{x:10.6f}" for x in vals )
+            print( f"    pos {pos}: [ {formatted} ... ]" )
+        print( f"\n  K head {head}, positions 0..{T-1} (first 8 elements each):" )
+        for pos in range( T ):
+            vals = k[0, head, pos, :8].tolist()
+            formatted = "  ".join( f"{x:10.6f}" for x in vals )
+            print( f"    pos {pos}: [ {formatted} ... ]" )
+
+extract_v( captured['layer_0.qkv'], 'Layer 0' )
+extract_v( captured['layer_1.qkv'], 'Layer 1' )