diff --git a/dsv4/layers/mhc.py b/dsv4/layers/mhc.py index 6b99a816..8e361a2f 100644 --- a/dsv4/layers/mhc.py +++ b/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 diff --git a/single_shot_inference.py b/single_shot_inference.py index 725d9c76..308f9f38 100644 --- a/single_shot_inference.py +++ b/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}")