From 64f547058ebc5de997df80684002de8ab3c02d88 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 4 Jun 2026 08:09:30 +0000 Subject: [PATCH] Fix graph replay: pass q_a from Graph A output to forward_attention - q_a is needed by the indexer in CSA layers - When q_heads/kv_3d are provided (graph replay), the projection code is skipped so q_a is never computed - Fix: add q_a_bufs to CUDAGraphDecoder, write q_a during Graph A capture, pass q_a as kwarg to forward_attention during graph replay - Also: forward_attention now accepts q_a kwarg (default None) --- single_shot_inference.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index a355ba1d..02de7db1 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -190,6 +190,7 @@ class CUDAGraphDecoder: self.x_normed_bufs = {} # li -> (1, H) BF16 — for compressor/indexer self.q_heads_bufs = {} # li -> (1, n_h, hd) BF16 — for FMHA self.kv_3d_bufs = {} # li -> (1, 1, hd) BF16 — for FMHA (pre-RoPE) + self.q_a_bufs = {} # li -> (1, q_a_dim) BF16 — q_a for indexer self.ctx_a_B_bufs = {} # li -> (1, 4, 4) FP32 — B_l for post_block self.ctx_a_C_bufs = {} # li -> (1, 4) BF16 — C_l for post_block self.X_mid_bufs = {} # li -> (1, 4, H) BF16 — X_l for post_block @@ -216,6 +217,7 @@ class CUDAGraphDecoder: self.x_normed_bufs[li] = torch.zeros(1, H, dtype=torch.bfloat16, device=dev) self.q_heads_bufs[li] = torch.zeros(1, n_h, hd, dtype=torch.bfloat16, device=dev) self.kv_3d_bufs[li] = torch.zeros(1, 1, hd, dtype=torch.bfloat16, device=dev) + self.q_a_bufs[li] = torch.zeros(1, 1536, dtype=torch.bfloat16, device=dev) # q_a for indexer self.ctx_a_B_bufs[li] = torch.zeros(1, 4, 4, dtype=torch.float32, device=dev) self.ctx_a_C_bufs[li] = torch.zeros(1, 4, dtype=torch.bfloat16, device=dev) self.X_mid_bufs[li] = torch.zeros(1, 4, H, dtype=torch.bfloat16, device=dev) @@ -302,6 +304,7 @@ class CUDAGraphDecoder: self.x_normed_bufs[li].copy_(x_normed) self.q_heads_bufs[li].copy_(q_heads) self.kv_3d_bufs[li].copy_(kv_3d) + self.q_a_bufs[li].copy_(q_a) self.ctx_a_B_bufs[li].copy_(B_l_a) self.ctx_a_C_bufs[li].copy_(C_l_a) self.X_mid_bufs[li].copy_(X_l) @@ -1003,7 +1006,7 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, x_quant=None, _profile_detail=False, _profile_times=None, comp_rope_cos=None, comp_rope_sin=None, - q_heads=None, kv_3d=None): + q_heads=None, kv_3d=None, q_a=None): dev = x_normed.device; T = x_normed.shape[0] n_h = cfg["num_attention_heads"]; hd = cfg["head_dim"]; rd = cfg.get("qk_rope_head_dim", 64) o_groups = cfg.get("o_groups", 16); o_rank = cfg.get("o_lora_rank", 1024) @@ -1044,6 +1047,9 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, q = unweighted_rmsnorm(q).bfloat16() _pt('q_b_end') q_heads = q.reshape(T, n_h, hd) + else: + # Graph replay: q_a provided from pre-allocated buffer + q_a = q_a # use the passed q_a from graph A output q_heads = _apply_rope(q_heads, positions, rope_cos, rope_sin, rd) _pt('rope_q_end') @@ -1843,7 +1849,7 @@ def main(): x_normed, layer_w[li], li, cfg, *rope_caches[gpu], kv_caches[li], dec_pos_per_gpu[gpu], compressors.get(li), indexers.get(li), prod_lins.get(li), - q_heads=q_heads, kv_3d=kv_3d, # pass pre-computed q/kv from graph A + q_heads=q_heads, kv_3d=kv_3d, q_a=graph_decoder.q_a_bufs[li], comp_rope_cos=comp_rope_caches[gpu][0] if comp_rope_caches else None, comp_rope_sin=comp_rope_caches[gpu][1] if comp_rope_caches else None, )