biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix graph replay: pass q_a from Graph A output to forward_attention

Browse Source 
- 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)
This commit is contained in:
biondizzle
2026-06-04 08:09:30 +00:00
parent 26da6d33af
commit 64f547058e
1 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
single_shot_inference.py
View File

@@ -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,
)