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

CUDA graph: Implement eager-break-at-attention decoder with sub-graph A/B split

Browse Source 
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.
This commit is contained in:
biondizzle
2026-06-03 19:24:26 +00:00
parent 5ea3aa3406
commit 486f74d900
1 changed files with 356 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files

441
single_shot_inference.py
View File

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