Update logbook: Repro - KnapSpec: Self-Speculative Decoding via Adaptive Layer Selection as a Knapsack Problem

logbook.json

@@ -11,7 +11,7 @@
     "icml2026-repro",
     "paper-k5nKHWp9VC"
   ],
-  "updated_at": "2026-07-16T14:48:59+00:00",
+  "updated_at": "2026-07-16T14:52:30+00:00",
   "root": {
     "slug": "index",
     "title": "Repro - KnapSpec: Self-Speculative Decoding via Adaptive Layer Selection as a Knapsack Problem",
@@ -55,6 +55,6 @@
       }
     ]
   },
-  "agent_view_tokens": 2070,
-  "revision": "1784213339262066106"
+  "agent_view_tokens": 2186,
+  "revision": "1784213550515826534"
 }

pages/claim-3-tpt-and-two-stage-optimization/page.md

@@ -217,3 +217,172 @@ https://huggingface.co/buckets/nickvdw/repro-knapspec-self-speculative-decoding-
 {"type": "markdown", "id": "cell_e9c31dd5cf16", "created_at": "2026-07-16T14:46:37+00:00", "title": "The implementation computes the paper TPT equation over gamma=1..8 after DP ove…"}
 -->
 The implementation computes the paper TPT equation over gamma=1..8 after DP over latency budgets. At n=128 the proxy TPT decreased monotonically from 7.42 (gamma=1) to 2.49 (gamma=8). The paper-level PCC values 0.837 (TPT) and 0.538 (acceptance rate) require the Llama-3.1-8B GovReport sweep and were not reproduced here.
+
+
+---
+<!-- trackio-cell
+{"type": "code", "id": "cell_659c963163ca", "created_at": "2026-07-16T14:52:27+00:00", "title": "Run: uv ssd_benchmark.py (exit 1)", "command": ["uv", "run", "python", "repro_knapspec/ssd_benchmark.py", "--steps", "4", "--gamma", "2"], "exit_code": 1, "duration_s": 15.977}
+-->
+````bash
+$ uv run python repro_knapspec/ssd_benchmark.py --steps 4 --gamma 2
+````
+
+exit 1 · 16.0s
+
+
+````python title=ssd_benchmark.py
+#!/usr/bin/env -S uv run --script
+# /// script
+# dependencies = ["torch", "transformers"]
+# ///
+"""End-to-end self-speculative draft/verify benchmark for Llama 3.2 1B.
+
+This prototype uses direct sublayer wrappers (no forward hooks). It measures
+the algorithmic draft/verify loop and reports its current Python/model-wrapper
+overhead separately from the claimed optimized-kernel speedup.
+"""
+import argparse, json, time
+from pathlib import Path
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+MODEL = "meta-llama/Llama-3.2-1B"
+PROMPT = "The committee reviewed the proposal and requested a detailed analysis of the expected costs, benefits, risks, and implementation timeline."
+
+class Switchable(torch.nn.Module):
+    def __init__(self, inner):
+        super().__init__(); self.inner = inner; self.skip = False
+    def forward(self, *args, **kwargs):
+        if not self.skip:
+            return self.inner(*args, **kwargs)
+        # Attention returns (hidden, optional weights, optional cache).
+        hidden = args[0]
+        zeros = torch.zeros_like(hidden)
+        if isinstance(self.inner, torch.nn.Module) and self.inner.__class__.__name__.endswith("Attention"):
+            return (zeros, None, None)
+        return zeros
+
+def install_switches(model):
+    switches = []
+    for layer in model.model.layers:
+        attn = Switchable(layer.self_attn); mlp = Switchable(layer.mlp)
+        layer.self_attn = attn; layer.mlp = mlp
+        switches.extend([("attn", attn), ("mlp", mlp)])
+    return switches
+
+def set_mask(switches, mask):
+    for name, module in switches: module.skip = (name, module.layer_idx) in mask if hasattr(module, "layer_idx") else False
+
+def set_mask_by_index(switches, mask):
+    for idx, (_, module) in enumerate(switches): module.skip = idx in mask
+
+def forward(model, ids, **kwargs):
+    with torch.inference_mode(): return model(input_ids=ids, use_cache=True, **kwargs)
+
+def greedy(model, ids, steps):
+    past = None; seq = ids; generated = 0; t0=time.perf_counter()
+    for _ in range(steps):
+        out = forward(model, seq if past is None else seq[:,-1:], past_key_values=past)
+        past = out.past_key_values; nxt=out.logits[:,-1].argmax(-1,keepdim=True); seq=torch.cat([seq,nxt],1); generated+=1
+    if torch.cuda.is_available(): torch.cuda.synchronize()
+    return seq, generated/(time.perf_counter()-t0)
+
+def spec_decode(model, switches, prompt, steps, gamma):
+    # Prototype draft/verify: draft incrementally, then verify the whole block.
+    # Full verification is intentionally explicit so acceptance is auditable.
+    seq=prompt.clone(); accepted=0; proposed=0; rounds=0; t0=time.perf_counter()
+    for _ in range((steps+gamma-1)//gamma):
+        rounds += 1; draft=seq.clone(); draft_tokens=[]
+        set_mask_by_index(switches, set(range(len(switches))))
+        # Use a conservative fixed mask: skip every other MLP, preserving all attention.
+        set_mask_by_index(switches, {i for i,(n,_) in enumerate(switches) if n == "mlp" and (i//2)%2==0})
+        past=None
+        for _ in range(min(gamma, steps-proposed)):
+            out=forward(model, draft if past is None else draft[:,-1:], past_key_values=past)
+            past=out.past_key_values; tok=out.logits[:,-1].argmax(-1,keepdim=True); draft=torch.cat([draft,tok],1); draft_tokens.append(tok)
+        block=torch.cat(draft_tokens,1); proposed += block.shape[1]
+        # Verify all proposed tokens with the unmasked target in one forward.
+        set_mask_by_index(switches, set())
+        verify=forward(model, torch.cat([seq,block],1), use_cache=False).logits
+        start=seq.shape[1]-1
+        target_next=verify[:,start:start+block.shape[1]].argmax(-1)
+        n=0
+        for j in range(block.shape[1]):
+            if target_next[0,j].item() != block[0,j].item(): break
+            n += 1
+        accepted += n
+        # Preserve lossless behavior: append accepted draft tokens, then target correction.
+        seq=torch.cat([seq,block[:,:n]],1)
+        if n < block.shape[1]: seq=torch.cat([seq,target_next[:,n:n+1]],1)
+        proposed = min(proposed, steps)
+        if seq.shape[1] >= prompt.shape[1]+steps: break
+    if torch.cuda.is_available(): torch.cuda.synchronize()
+    elapsed=time.perf_counter()-t0
+    return seq, {'accepted':accepted,'proposed':proposed,'rounds':rounds,'acceptance':accepted/max(1,proposed),'tokens_per_s':(seq.shape[1]-prompt.shape[1])/elapsed,'elapsed_s':elapsed}
+
+def main():
+    ap=argparse.ArgumentParser(); ap.add_argument('--steps',type=int,default=16); ap.add_argument('--gamma',type=int,default=4); args=ap.parse_args()
+    device='cuda' if torch.cuda.is_available() else 'cpu'; dtype=torch.float16 if device=='cuda' else torch.float32
+    tok=AutoTokenizer.from_pretrained(MODEL); model=AutoModelForCausalLM.from_pretrained(MODEL,dtype=dtype).to(device).eval(); switches=install_switches(model)
+    prompt=tok(PROMPT,return_tensors='pt').input_ids.to(device)
+    _, baseline_tps=greedy(model,prompt,args.steps)
+    _, spec=spec_decode(model,switches,prompt,args.steps,args.gamma)
+    result={'model':MODEL,'device':device,'prompt_tokens':prompt.shape[1],'requested_steps':args.steps,'gamma':args.gamma,'baseline_tokens_per_s':baseline_tps,**spec,'speedup_vs_greedy':spec['tokens_per_s']/baseline_tps}
+    Path('repro_knapspec/outputs').mkdir(parents=True,exist_ok=True); Path('repro_knapspec/outputs/ssd_benchmark.json').write_text(json.dumps(result,indent=2)); print(json.dumps(result,indent=2))
+if __name__=='__main__': main()
+
+````
+
+
+````output
+
+Loading weights:   0%|          | 0/146 [00:00<?, ?it/s]
+Loading weights:   1%|          | 1/146 [00:00<00:37,  3.85it/s]
+Loading weights:  64%|██████▎   | 93/146 [00:00<00:00, 319.09it/s]
+Loading weights:  98%|█████████▊| 143/146 [00:00<00:00, 247.26it/s]
+Loading weights: 100%|██████████| 146/146 [00:00<00:00, 231.69it/s]
+Traceback (most recent call last):
+  File "/home/nick/workspace/icml26-hf-challenge-knapspec/repro_knapspec/ssd_benchmark.py", line 99, in <module>
+    if __name__=='__main__': main()
+                             ~~~~^^
+  File "/home/nick/workspace/icml26-hf-challenge-knapspec/repro_knapspec/ssd_benchmark.py", line 96, in main
+    _, spec=spec_decode(model,switches,prompt,args.steps,args.gamma)
+            ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/nick/workspace/icml26-hf-challenge-knapspec/repro_knapspec/ssd_benchmark.py", line 73, in spec_decode
+    verify=forward(model, torch.cat([seq,block],1), use_cache=False).logits
+           ~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/nick/workspace/icml26-hf-challenge-knapspec/repro_knapspec/ssd_benchmark.py", line 47, in forward
+    with torch.inference_mode(): return model(input_ids=ids, use_cache=True, **kwargs)
+                                        ~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+TypeError: LlamaForCausalLM(
+  (model): LlamaModel(
+    (embed_tokens): Embedding(128256, 2048)
+    (layers): ModuleList(
+      (0-15): 16 x LlamaDecoderLayer(
+        (self_attn): Switchable(
+          (inner): LlamaAttention(
+            (q_proj): Linear(in_features=2048, out_features=2048, bias=False)
+            (k_proj): Linear(in_features=2048, out_features=512, bias=False)
+            (v_proj): Linear(in_features=2048, out_features=512, bias=False)
+            (o_proj): Linear(in_features=2048, out_features=2048, bias=False)
+          )
+        )
+        (mlp): Switchable(
+          (inner): LlamaMLP(
+            (gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
+            (up_proj): Linear(in_features=2048, out_features=8192, bias=False)
+            (down_proj): Linear(in_features=8192, out_features=2048, bias=False)
+            (act_fn): SiLUActivation()
+          )
+        )
+        (input_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
+        (post_attention_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
+      )
+    )
+    (norm): LlamaRMSNorm((2048,), eps=1e-05)
+    (rotary_emb): LlamaRotaryEmbedding()
+  )
+  (lm_head): Linear(in_features=2048, out_features=128256, bias=False)
+) got multiple values for keyword argument 'use_cache'
+
+````