CUDA graph: Fix per-step allocations in decode loop

1. mHCLayer.init_state: Add out_buf parameter for in-place write
   - Pre-allocated dec_X_buf (1, 4, 7168) on cuda:0
   - Eliminates .unsqueeze().expand().clone() allocation each step

2. single_shot_inference.py: Pre-allocate dec_embed_buf
   - Placeholder for embedding output (graph capture will use this)

3. Note: Cross-GPU X.to() transfers still allocate per step
   - This requires per-GPU X buffers (part of graph capture architecture)

dsv4/layers/mhc.py

@@ -431,12 +431,23 @@ class mHCLayer:
     def init_state(
         embeddings: torch.Tensor,   # (T, d) BF16 — token embeddings
         n_hc: int = 4,
+        out_buf: torch.Tensor = None,  # (T, n_hc, d) BF16 — pre-allocated output buffer
     ) -> torch.Tensor:
         """
         Initialise X_0 for the first layer.
 
         Returns: (T, n_hc, d) BF16
+        
+        When out_buf is provided, writes to it in-place (no allocation).
+        This is required for CUDA graph capture where per-step
+        allocations are forbidden.
         """
+        if out_buf is not None:
+            # In-place: copy embeddings to all n_hc streams
+            out_buf[:, 0, :].copy_(embeddings)  # Stream 0 gets the embedding
+            for h in range(1, n_hc):
+                out_buf[:, h, :].copy_(embeddings)  # All other streams too
+            return out_buf
         return embeddings.unsqueeze(1).expand(-1, n_hc, -1).clone()
 
     @staticmethod

single_shot_inference.py

@@ -1576,6 +1576,13 @@ def main():
         layer_event_count = 0
         cuda_layer_events = []  # list of (tag, li, timestamp) for fine-grained profiling
 
+    # Pre-allocate decode X buffer — zero per-step allocation
+    # init_state writes to this buffer in-place (no .clone() allocation)
+    dec_X_buf = torch.zeros(1, 4, H, dtype=torch.bfloat16, device='cuda:0')
+    # Pre-allocate embedding output buffer — embed() returns a new tensor each call.
+    # For graph capture, we'd copy into this buffer. For now, used as reference.
+    dec_embed_buf = torch.zeros(1, H, dtype=torch.bfloat16, device='cuda:0')
+
     for step in range(MAX_NEW_TOKENS):
         t1 = time.time()
         dec_tid_buf[0] = all_tokens[-1]
@@ -1583,7 +1590,7 @@ def main():
         dec_pos_buf[0] = len(all_tokens) - 1
 
         t_e = time.perf_counter()
-        X = mHCLayer.init_state(embed(dec_tid_buf))
+        X = mHCLayer.init_state(embed(dec_tid_buf), out_buf=dec_X_buf)
         for li in range(n_layers):
             gpu = li % NUM_GPUS
             if X.device != torch.device(f"cuda:{gpu}"): X = X.to(f"cuda:{gpu}")