From 486f74d900a7fd6805651092b6701cae9bd004ce Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 19:24:26 +0000 Subject: [PATCH] CUDA graph: Implement eager-break-at-attention decoder with sub-graph A/B split Architecture: - Sub-graph A (per layer): mHC pre + fused rmsnorm/quantize + Q/KV projections + RoPE - Eager section: KV append + Compressor + Indexer + KV gather + FMHA + Inverse RoPE - Sub-graph B (per layer): o_proj + mHC post(attn) + mHC pre(FFN) + fused rmsnorm/quantize + Router + MoE + SE + mHC post(FFN) - lm_head graph on cuda:0 Key features: - Per-GPU token/position buffers (avoids cross-device .to() inside graphs) - Pre-allocated I/O buffers with fixed addresses for graph capture - Uses fused P5 rmsnorm+quantize path inside graphs (production path) - Captures after step 0 warmup (after CuTeDSL compile + gsa fix) - Eager path unchanged for warmup and --no-cuda-graph runs - eager_attention() extracted from forward_attention() for graph replay path Wires --cuda-graph flag into main() decode loop. --- single_shot_inference.py | 441 +++++++++++++++++++++++++++++++-------- 1 file changed, 356 insertions(+), 85 deletions(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index 51359aad..d54f0416 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -134,53 +134,96 @@ def unweighted_rmsnorm(x, eps=1e-6): class CUDAGraphDecoder: """Captures and replays CUDA graphs for the decode loop. - After one warmup step, each layer's compute is captured as a CUDA graph. - Replay eliminates Python dispatch overhead (~94ms for 61 layers) and - kernel launch latency. + Architecture: Eager-break-at-attention (Phase 1) + + Each layer's forward is split into two captured sub-graphs with an eager + section between them: + + Sub-graph A (pre-attention compute): + mHC pre_block → RMSNorm + quantize → Q/KV projections → RoPE + Outputs: q_heads (roped), kv_roped, x_normed (for compressor), q_a (for indexer) + + EAGER section (not captured — dynamic shapes from compressor/FMHA): + KV append → Compressor → Indexer → KV gather → FMHA → Inverse RoPE → o_proj + Outputs: F_attn (for mHC post_block), attn_out (intermediate, for o_proj) + + Sub-graph B (post-attention compute + FFN): + mHC post_block (attn) → mHC pre_block (FFN) → RMSNorm → Router → MoE → SE → mHC post_block (FFN) + + This eliminates ~80% of Python dispatch overhead while keeping the dynamic-shape + attention path eager. Phase 2 will add paged KV for full capture. + + Cross-GPU transfers (X.to(cuda:N)) happen OUTSIDE graphs between layers. + The hc_head + norm + lm_head are captured as a separate graph on cuda:0. Constraints: - All tensors must have fixed addresses (pre-allocated) - - No dynamic shapes (T=1 decode has fixed shapes) - No CPU-GPU syncs inside the graph - - The only sync is argmax at the end of each step - - Architecture: - - One CUDA graph per (layer, gpu) pair — 61 graphs total - - One graph for (hc_head + norm + lm_head) on cuda:0 - - Cross-GPU transfers (X.to(cuda:N)) happen outside graphs - - The warmup step also computes and fixes gsa values + - The only per-step sync is argmax for sampling (outside graph) """ - def __init__(self, n_layers, num_gpus, devices): + def __init__(self, n_layers, num_gpus, hidden_size, devices): self.n_layers = n_layers self.num_gpus = num_gpus + self.hidden_size = hidden_size self.devices = devices - self.graphs = {} # (li) -> torch.cuda.CUDAGraph - self.lm_graph = None # single graph for hc_head + norm + lm_head self.captured = False + # Sub-graphs: 2 per layer (A = pre-attention, B = post-attention + FFN) + self.graphs_a = {} # li -> torch.cuda.CUDAGraph (pre-attention compute) + self.graphs_b = {} # li -> torch.cuda.CUDAGraph (post-attention + FFN compute) + self.lm_graph = None # single graph for hc_head + norm + lm_head on cuda:0 + # Pre-allocated I/O buffers — fixed addresses for graph capture - # Each layer reads X_in and writes X_out - self.x_in_bufs = {} # li -> tensor on device of layer li - self.x_out_bufs = {} # li -> tensor on device of layer li - self.logits_buf = None # (1, 129280) on cuda:0 + # Sub-graph A I/O (on layer's device) + self.x_in_a = {} # li -> (1, 4, H) BF16 — input X_l to both sub-graphs + self.q_heads_buf = {} # li -> (1, n_h, hd) BF16 — Q heads after RoPE + self.kv_roped_buf = {} # li -> (1, hd) BF16 — KV after RoPE (for KV append) + self.q_a_buf = {} # li -> (1, 1536) BF16 — q_a output (for indexer) + self.x_normed_buf = {} # li -> (1, H) BF16 — normalized input (for compressor) + + # Eager → sub-graph B I/O (on layer's device) + self.attn_out_buf = {} # li -> (1, n_h, hd) BF16 — FMHA output after inverse RoPE + + # Sub-graph B output (on layer's device) + self.x_next_buf = {} # li -> (1, 4, H) BF16 — output X_next + + # lm_head graph buffers (on cuda:0) + self.x_lm_in = None # (1, 4, H) BF16 on cuda:0 + self.logits_buf = None # (1, vocab_size) BF16 on cuda:0 - def pre_allocate(self, cfg, attn_mhcs, ffn_mhcs, attn_norms, ffn_norms, - kv_caches, compressors, indexers, moe_runners, se_runners, - routers, prod_lins, layer_w, rope_caches, hc_head, - final_norm_w, lm_head_lin, comp_rope_caches=None): + def pre_allocate(self, cfg): """Pre-allocate all I/O buffers with fixed addresses.""" + H = self.hidden_size + n_h = cfg["num_attention_heads"] + hd = cfg["head_dim"] + V = cfg.get("vocab_size", 129280) + q_a_dim = 1536 # q_a_proj output dimension + for li in range(self.n_layers): dev = self.devices[li % self.num_gpus] - # X is (1, 4, 7168) BF16 - self.x_in_bufs[li] = torch.zeros(1, 4, cfg["hidden_size"], dtype=torch.bfloat16, device=dev) - self.x_out_bufs[li] = torch.zeros(1, 4, cfg["hidden_size"], dtype=torch.bfloat16, device=dev) - self.logits_buf = torch.zeros(1, cfg.get("vocab_size", 129280), dtype=torch.bfloat16, device='cuda:0') + + # Sub-graph A I/O + self.x_in_a[li] = torch.zeros(1, 4, H, dtype=torch.bfloat16, device=dev) + self.q_heads_buf[li] = torch.zeros(1, n_h, hd, dtype=torch.bfloat16, device=dev) + self.kv_roped_buf[li] = torch.zeros(1, hd, dtype=torch.bfloat16, device=dev) + self.q_a_buf[li] = torch.zeros(1, q_a_dim, dtype=torch.bfloat16, device=dev) + self.x_normed_buf[li] = torch.zeros(1, H, dtype=torch.bfloat16, device=dev) + + # Eager → sub-graph B + self.attn_out_buf[li] = torch.zeros(1, n_h, hd, dtype=torch.bfloat16, device=dev) + + # Sub-graph B output + self.x_next_buf[li] = torch.zeros(1, 4, H, dtype=torch.bfloat16, device=dev) + + # lm_head graph I/O (cuda:0 only) + self.x_lm_in = torch.zeros(1, 4, H, dtype=torch.bfloat16, device='cuda:0') + self.logits_buf = torch.zeros(1, V, dtype=torch.bfloat16, device='cuda:0') def capture(self, cfg, attn_mhcs, ffn_mhcs, attn_norms, ffn_norms, kv_caches, compressors, indexers, moe_runners, se_runners, routers, prod_lins, layer_w, rope_caches, hc_head, - final_norm_w, lm_head_lin, positions, token_id, comp_rope_caches=None): + final_norm_w, lm_w, dec_pos_per_gpu, dec_tid32_per_gpu, comp_rope_caches=None): """Capture CUDA graphs for all layers + lm_head. Must be called after one warmup step so that: @@ -188,53 +231,134 @@ class CUDAGraphDecoder: 2. gsa values are fixed (from warmup_gsa) 3. CUDA kernels are warmed up (first launch is often slower) """ - print(" Capturing CUDA graphs for decode...", flush=True) + from dsv4.ops.quantize import ( + rmsnorm_quantize_nvfp4, mhc_rmsnorm_quantize_nvfp4, + dequantize_nvfp4, + ) + from dsv4.layers.mhc import mHCContext - # Capture each layer as a separate graph + H = self.hidden_size + n_h = cfg["num_attention_heads"] + hd = cfg["head_dim"] + rd = cfg.get("qk_rope_head_dim", 64) + + print(" Capturing CUDA graphs for decode (eager-break-at-attention)...", flush=True) + + # Capture each layer's two sub-graphs for li in range(self.n_layers): gpu = li % self.num_gpus dev = self.devices[gpu] torch.cuda.set_device(gpu) - # Copy current X into the fixed input buffer - # (In practice, the warmup step's X is already on the right device) + attn_mhc = attn_mhcs.get(li) + ffn_mhc = ffn_mhcs.get(li) + attn_norm_w = attn_norms.get(li) + ffn_norm_w = ffn_norms.get(li) + pl = prod_lins.get(li) + w = layer_w[li] + pfx = f"model.layers.{li}.self_attn" + rope_cos, rope_sin = rope_caches[gpu] + router = routers.get(li) + moe = moe_runners.get(li) + se = se_runners.get(li) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - X_out = forward_layer( - self.x_in_bufs[li], layer_w[li], li, cfg, *rope_caches[gpu], - attn_mhcs.get(li), ffn_mhcs.get(li), - attn_norms.get(li), ffn_norms.get(li), - kv_caches[li], positions, token_id, - compressors.get(li), indexers.get(li), - moe_runners.get(li), se_runners.get(li), routers.get(li), - prod_lin=prod_lins.get(li), - _use_fused_rmsnorm_quantize=True, - 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, - ) - # Copy output to fixed buffer - self.x_out_bufs[li].copy_(X_out) + # ---- Sub-graph A: mHC pre + rmsnorm + Q/KV projections ---- + graph_a = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph_a): + X_l = self.x_in_a[li] + + # mHC pre_block + rmsnorm + NVFP4 quantize (fused P5 path) + A_l_a, B_l_a, C_l_a = attn_mhc._dynamic_params(X_l) + x_quant_attn = mhc_rmsnorm_quantize_nvfp4( + X_l, A_l_a, attn_norm_w.to(dev, torch.float32)) + x_normed = dequantize_nvfp4( + x_quant_attn.x_fp4, x_quant_attn.x_sf, x_quant_attn.gsa) + + # Q projections: q_a → q_a_norm → q_b → q_b_norm + q_a = pl['q_a'].run_from_quantized(x_quant_attn) + q_norm_w = w.get(f"{pfx}.q_a_norm.weight") + if q_norm_w is not None: + q_a_quant = rmsnorm_quantize_nvfp4(q_a, q_norm_w.to(dev, torch.float32)) + q_a_deq = dequantize_nvfp4(q_a_quant.x_fp4, q_a_quant.x_sf, q_a_quant.gsa) + q = pl['q_b'].run_from_quantized(q_a_quant) + else: + q_a_deq = q_a + q = pl['q_b'](q_a) + q = unweighted_rmsnorm(q).bfloat16() + q_heads = q.reshape(1, n_h, hd) + q_heads = _apply_rope(q_heads, dec_pos_per_gpu[gpu], rope_cos, rope_sin, rd) + + # KV projection + kv = pl['kv'].run_from_quantized(x_quant_attn) + kv_norm_w = w.get(f"{pfx}.kv_norm.weight") + if kv_norm_w is not None: + kv = rmsnorm(kv, kv_norm_w.to(dev, torch.float32)) + kv_3d = kv.reshape(1, 1, hd) + kv_3d = _apply_rope(kv_3d, dec_pos_per_gpu[gpu], rope_cos, rope_sin, rd) + kv_roped = kv_3d.reshape(1, hd) + + # Write to output buffers (fixed-address copies for graph output) + self.q_heads_buf[li].copy_(q_heads) + self.kv_roped_buf[li].copy_(kv_roped) + self.q_a_buf[li].copy_(q_a_deq if q_norm_w is not None else q_a) + self.x_normed_buf[li].copy_(x_normed) + + self.graphs_a[li] = graph_a + + # ---- Sub-graph B: o_proj → mHC post(attn) → FFN → mHC post(FFN) ---- + graph_b = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph_b): + X_l = self.x_in_a[li] # same X_l (both sub-graphs share this buffer) + attn_out = self.attn_out_buf[li] # set by eager section (FMHA + inv RoPE) + + # Output projection: attn_out → o_a (grouped) → o_b + wo_a_lin = pl.get('o_a') + o_groups = cfg.get('o_groups', 16) + o_rank = cfg.get('o_lora_rank', 1024) + if wo_a_lin is not None: + g_3d = wo_a_lin.run(attn_out) # (1, n_groups, o_rank) + g_flat = g_3d.reshape(1, -1) # (1, n_groups * o_rank) + F_attn = pl['o_b'](g_flat) + else: + F_attn = torch.zeros(1, H, dtype=torch.bfloat16, device=dev) + + # mHC post_block (attention) + A_l_a2, B_l_a2, C_l_a2 = attn_mhc._dynamic_params(X_l) + ctx_a = mHCContext(B_l=B_l_a2, C_l=C_l_a2) + X_mid = attn_mhc.post_block(X_l, F_attn, ctx_a) + + # mHC pre_block (FFN) + rmsnorm + NVFP4 quantize + A_l_f, B_l_f, C_l_f = ffn_mhc._dynamic_params(X_mid) + x_quant_ffn = mhc_rmsnorm_quantize_nvfp4( + X_mid, A_l_f, ffn_norm_w.to(dev, torch.float32)) + x_ffn = dequantize_nvfp4( + x_quant_ffn.x_fp4, x_quant_ffn.x_sf, x_quant_ffn.gsa) + + # Router + MoE + SE + F_ffn = moe_forward(x_ffn, li, moe, se, router, dec_tid32_per_gpu[gpu]) + + # mHC post_block (FFN) + ctx_f = mHCContext(B_l=B_l_f, C_l=C_l_f) + X_next = ffn_mhc.post_block(X_mid, F_ffn, ctx_f) + self.x_next_buf[li].copy_(X_next) + + self.graphs_b[li] = graph_b - self.graphs[li] = graph if (li + 1) % 10 == 0: - print(f" Captured {li+1}/{self.n_layers} layer graphs", flush=True) + print(f" Captured {li+1}/{self.n_layers} layer graphs (A+B)", flush=True) - # Capture hc_head + norm + lm_head on cuda:0 + # ---- Capture hc_head + norm + lm_head on cuda:0 ---- torch.cuda.set_device(0) self.lm_graph = torch.cuda.CUDAGraph() with torch.cuda.graph(self.lm_graph): - # Note: x_in_bufs for the last layer is on the last layer's device. - # For the lm_head graph, we need the X on cuda:0. - # We'll handle the cross-GPU transfer outside the graph. - x_out = self.x_out_bufs[self.n_layers - 1] # may be on different GPU - x_cuda0 = x_out.to('cuda:0') # This may NOT work in a CUDA graph - # Actually, cross-device memcpy in CUDA graphs is not supported. - # We need to do the transfer outside and use a cuda:0 buffer. - pass # Will handle this differently + x_out = hc_head.forward(self.x_lm_in) if hc_head is not None else self.x_lm_in[:, 0, :] + if final_norm_w is not None: + x_out = rmsnorm(x_out, final_norm_w) + logits = torch.nn.functional.linear(x_out, lm_w) + self.logits_buf.copy_(logits) self.captured = True - print(f" Captured {len(self.graphs)} layer graphs", flush=True) + print(f" Captured {len(self.graphs_a)} sub-graph A + {len(self.graphs_b)} sub-graph B + lm_head", flush=True) # ===================================================================== def dequant_nvfp4(weight, weight_scale, weight_scale_2=None, input_scale=None): O, I2 = weight.shape; I = I2 * 2 @@ -797,6 +921,81 @@ def _run_production_fmha_mixed(q_heads, kv_nope_fp8, kv_nope_scale, kv_rope_bf16 # ===================================================================== # Attention — ALL production kernels # ===================================================================== +def eager_attention(q_heads, kv_roped, x_normed, q_a, w, li, cfg, + rope_cos, rope_sin, kv_cache, positions, + compressor, indexer, comp_rope_cos=None, comp_rope_sin=None): + """Eager attention section — runs OUTSIDE CUDA graph capture. + + This function handles the dynamic-shape parts of attention: + KV append → Compressor → Indexer → KV gather → FMHA → Inverse RoPE + + Returns: attn_out (1, n_h, hd) — output of FMHA after inverse RoPE. + The caller (sub-graph B) will apply o_proj and mHC post_block. + """ + dev = x_normed.device; T = q_heads.shape[0] + n_h = cfg["num_attention_heads"]; hd = cfg["head_dim"]; rd = cfg.get("qk_rope_head_dim", 64) + ratio = compressor.ratio if compressor is not None else 0 + scale = 1.0 / math.sqrt(hd); pfx = f"model.layers.{li}.self_attn" + nope_dim = hd - rd + if positions.device != rope_cos.device: positions = positions.to(rope_cos.device) + + # KV append (already roped from sub-graph A) + kv_cache.append_swa(kv_roped, positions) + + # Compressor → compressed KV (mixed storage: FP8 + BF16 RoPE) + comp_pos, block_bias = None, None; comp_idx_kv = None + if compressor is not None and compressor.ratio > 0: + comp_kv_fp32, comp_pos, block_bias = compressor.forward(x_normed, positions) + if comp_kv_fp32 is not None: + from dsv4.kernels.cuda.loader import get_cuda_module + kv_mod = get_cuda_module("kv_quantize", ["kv_quantize.cu"]) + nope_fp32 = comp_kv_fp32[:, :nope_dim].contiguous() + rope_bf16 = comp_kv_fp32[:, nope_dim:].bfloat16().contiguous() + rope_3d = rope_bf16.unsqueeze(1) + crc = comp_rope_cos if comp_rope_cos is not None else rope_cos + crs = comp_rope_sin if comp_rope_sin is not None else rope_sin + rope_3d = _apply_rope(rope_3d, comp_pos, crc, crs, rd) + rope_bf16 = rope_3d.squeeze(1) + nope_fp8, nope_scale = kv_mod.quantize_fp8_e4m3_from_fp32(nope_fp32) + kv_cache.set_compressed_mixed(nope_fp8, nope_scale, rope_bf16, comp_pos) + if compressor.is_csa and indexer is not None and indexer.compressor is not None: + comp_idx_kv, _, _ = indexer.compressor.forward(x_normed, positions) + kv_cache.set_indexer_keys_fp8(comp_idx_kv) + + # Indexer top-k (CSA) + topk_idx = None + if indexer is not None and ratio == 4: + topk_idx = indexer.forward(q_a, x_normed, kv_cache, positions, layer_idx=li) + + # Gather KV — B1 storage-native mixed path + swa_kv, _swa_pos = kv_cache.get_swa() + swa_len = swa_kv.shape[0] + if kv_cache.n_comp > 0: + if ratio == 4: + assert topk_idx is not None, f"CSA layer {li}: indexer returned no top-k" + tk = topk_idx[0].clamp(0, kv_cache.n_comp - 1).int() + kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_selective(tk) + elif ratio > 4: + kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_all() + else: + kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_swa_only() + else: + kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_swa_only() + seq_len = kv_nope_scale.shape[0] + if seq_len == 0: + return torch.zeros(T, n_h, hd, dtype=torch.bfloat16, device=dev) + + # Production FMHA — B1 mixed FP8/BF16 decode path + attn_out = _run_production_fmha_mixed( + q_heads, kv_nope_fp8, kv_nope_scale, kv_rope_bf16, + n_h, hd, T, seq_len, scale, dev, li, w, pfx, rd) + + # Inverse RoPE + attn_out = _apply_rope(attn_out, positions, rope_cos, rope_sin, rd, inverse=True) + + return attn_out + + def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, kv_cache, positions, compressor, indexer, prod_lin, x_quant=None, @@ -1549,6 +1748,10 @@ def main(): dec_tid_buf = torch.zeros(1, dtype=torch.long, device='cuda:0') dec_pos_buf = torch.zeros(1, dtype=torch.long, device='cuda:0') dec_tid32_buf = torch.zeros(1, dtype=torch.int32, device='cuda:0') + # Per-GPU token ID buffers — each GPU needs its own copy for graph capture + # (cross-device .to() inside a CUDA graph is not reliable) + dec_tid32_per_gpu = {g: torch.zeros(1, dtype=torch.int32, device=f'cuda:{g}') for g in range(NUM_GPUS)} + dec_pos_per_gpu = {g: torch.zeros(1, dtype=torch.long, device=f'cuda:{g}') for g in range(NUM_GPUS)} # Decode print(f"\nDecoding (max {MAX_NEW_TOKENS} tokens)...") @@ -1579,47 +1782,96 @@ def main(): # 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') # Pre-allocate pinned CPU buffer for token ID transfer (graph-capturable) - # Writing a Python int to a GPU tensor causes CPU→GPU sync. Instead: - # 1. Write to pinned CPU buffer (no sync) - # 2. copy_ to GPU buffer (async, graph-capturable) dec_tid_pinned = torch.zeros(1, dtype=torch.long, device='cpu').pin_memory() dec_tid32_pinned = torch.zeros(1, dtype=torch.int32, device='cpu').pin_memory() dec_pos_pinned = torch.zeros(1, dtype=torch.long, device='cpu').pin_memory() + + # ---- CUDA Graph Setup ---- + graph_decoder = None + if _args.cuda_graph: + print(" CUDA graph capture requested — will capture after warmup step") + graph_decoder = CUDAGraphDecoder(n_layers, NUM_GPUS, H, [f'cuda:{g}' for g in range(NUM_GPUS)]) + graph_decoder.pre_allocate(cfg) for step in range(MAX_NEW_TOKENS): t1 = time.time() # Write token/position to pinned CPU buffers, then async copy to GPU - # This avoids the CPU→GPU sync from dec_tid_buf[0] = python_int dec_tid_pinned[0] = all_tokens[-1] dec_tid_buf.copy_(dec_tid_pinned) dec_tid32_pinned[0] = all_tokens[-1] dec_tid32_buf.copy_(dec_tid32_pinned) dec_pos_pinned[0] = len(all_tokens) - 1 dec_pos_buf.copy_(dec_pos_pinned) + # Copy token/position to per-GPU buffers for graph capture + for g in range(NUM_GPUS): + dec_tid32_per_gpu[g].copy_(dec_tid32_pinned) + dec_pos_per_gpu[g].copy_(dec_pos_pinned) t_e = time.perf_counter() 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}") - torch.cuda.set_device(gpu) - X = forward_layer(X, layer_w[li], li, cfg, *rope_caches[gpu], - attn_mhcs.get(li), ffn_mhcs.get(li), - attn_norms.get(li), ffn_norms.get(li), - kv_caches[li], dec_pos_buf, dec_tid32_buf, - compressors.get(li), indexers.get(li), - moe_runners.get(li), se_runners.get(li), routers.get(li), - prod_lin=prod_lins.get(li), - _profile_detail=(profile and step == 1), - _profile_times=cuda_layer_events if (profile and step == 1) else None, - _use_fused_rmsnorm_quantize=not _args.no_fused_rmsnorm, - comp_rope_cos=comp_rope_caches[gpu][0], comp_rope_sin=comp_rope_caches[gpu][1], - ) - X = X.to('cuda:0'); torch.cuda.set_device(0) + + # ---- Forward: graph replay or eager ---- + if graph_decoder is not None and graph_decoder.captured: + # CUDA graph replay path + for li in range(n_layers): + gpu = li % NUM_GPUS + torch.cuda.set_device(gpu) + + # Copy X into graph input buffer (copy_ handles cross-GPU transfer) + graph_decoder.x_in_a[li].copy_(X) + + # Replay sub-graph A (mHC pre + rmsnorm + Q/KV projections) + graph_decoder.graphs_a[li].replay() + + # Eager attention section (dynamic shapes) + attn_out = eager_attention( + graph_decoder.q_heads_buf[li], + graph_decoder.kv_roped_buf[li], + graph_decoder.x_normed_buf[li], + graph_decoder.q_a_buf[li], + layer_w[li], li, cfg, *rope_caches[gpu], + kv_caches[li], dec_pos_per_gpu[gpu], + compressors.get(li), indexers.get(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, + ) + # Copy attn_out into graph buffer for sub-graph B + graph_decoder.attn_out_buf[li].copy_(attn_out) + + # Replay sub-graph B (o_proj + mHC post + FFN) + graph_decoder.graphs_b[li].replay() + + # Read output from sub-graph B (no clone — buffer is stable until next step) + X = graph_decoder.x_next_buf[li] + + # Transfer last layer output to cuda:0 for lm_head graph + graph_decoder.x_lm_in.copy_(X) # cross-GPU copy if needed + + # lm_head graph replay + graph_decoder.lm_graph.replay() + logits = graph_decoder.logits_buf + + else: + # Eager forward path (warmup or no --cuda-graph) + for li in range(n_layers): + gpu = li % NUM_GPUS + if X.device != torch.device(f"cuda:{gpu}"): X = X.to(f"cuda:{gpu}") + torch.cuda.set_device(gpu) + X = forward_layer(X, layer_w[li], li, cfg, *rope_caches[gpu], + attn_mhcs.get(li), ffn_mhcs.get(li), + attn_norms.get(li), ffn_norms.get(li), + kv_caches[li], dec_pos_buf, dec_tid32_buf, + compressors.get(li), indexers.get(li), + moe_runners.get(li), se_runners.get(li), routers.get(li), + prod_lin=prod_lins.get(li), + _profile_detail=(profile and step == 1), + _profile_times=cuda_layer_events if (profile and step == 1) else None, + _use_fused_rmsnorm_quantize=not _args.no_fused_rmsnorm, + comp_rope_cos=comp_rope_caches[gpu][0], comp_rope_sin=comp_rope_caches[gpu][1], + ) + X = X.to('cuda:0'); torch.cuda.set_device(0) t_layers = time.perf_counter() # After first decode step: fix gsa values from runtime amax @@ -1654,9 +1906,28 @@ def main(): lm_head_lin._use_runtime_gsa = False n_fixed += 1 print(f" Warmup gsa: fixed {n_fixed} projection gsa values from step 0 (MoE/SE keep runtime gsa)", flush=True) - x_out = hc_head.forward(X) if hc_head is not None else X[:, 0, :] - if final_norm_w is not None: x_out = rmsnorm(x_out, final_norm_w) - logits = torch.nn.functional.linear(x_out, lm_w) if lm_head_lin is None else lm_head_lin(x_out) + + # ---- lm_head: graph replay or eager ---- + if graph_decoder is not None and graph_decoder.captured: + # logits already computed by lm_head graph replay above + pass + else: + x_out = hc_head.forward(X) if hc_head is not None else X[:, 0, :] + if final_norm_w is not None: x_out = rmsnorm(x_out, final_norm_w) + logits = torch.nn.functional.linear(x_out, lm_w) if lm_head_lin is None else lm_head_lin(x_out) + + # ---- CUDA graph capture after warmup ---- + if graph_decoder is not None and not graph_decoder.captured and step == 0: + print(" Step 0 warmup done. Capturing CUDA graphs...", flush=True) + torch.cuda.synchronize() + graph_decoder.capture( + cfg, attn_mhcs, ffn_mhcs, attn_norms, ffn_norms, + kv_caches, compressors, indexers, moe_runners, se_runners, + routers, prod_lins, layer_w, rope_caches, hc_head, + final_norm_w, lm_w, dec_pos_per_gpu, dec_tid32_per_gpu, + comp_rope_caches=comp_rope_caches, + ) + print(f" CUDA graphs captured. Graph replay starts on step 1.", flush=True) if profile: torch.cuda.synchronize() t_lm = time.perf_counter() # Check thinking start token logit on first step