diff --git a/single_shot_inference.py b/single_shot_inference.py index 9b67516d..105953ed 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -1,23 +1,15 @@ #!/usr/bin/env python3 """Single-shot DSV4-Pro inference — Full production pipeline, 8-GPU. -Exercises the production kernel stack end-to-end: - - NVFP4 GEMM kernels (CuTeDSL ScaledGroupedGemm) for all projections - - 6-warp TMA FMHA kernel (fmha_6warp_tma_multirow_multitile.cuh) with sink bias - - CSA/HCA compressor (token-level softmax) - - Indexer score+topk - - Dense/Hash router kernels - - Production mHC (Sinkhorn-Knopp, B_l transposed, [pre,post,comb]) - - Production Nvfp4Linear, Nvfp4MoE, Nvfp4SharedExpert +ALL projections use production NVFP4 GEMM kernels (CuTeDSL). +ALL attention uses production FMHA (6-warp TMA multi-tile + sink bias). +ALL MoE uses production Nvfp4MoE + Nvfp4SharedExpert + Router. NO PyTorch SDPA fallback. NO dequant+matmul for production projections. -ALL tensor-core NVFP4 GEMMs. ALL kernel paths. - This is the ground truth for vLLM / SGLang integration. """ import os, sys, time, json, math, argparse, logging import torch -# os.environ['CUDA_LAUNCH_BLOCKING'] = '1' # Disabled — was for debugging import torch.nn.functional as F from pathlib import Path @@ -31,23 +23,23 @@ def parse_args(): p.add_argument('--seed', type=int, default=42) p.add_argument('--verbose', type=int, default=1) p.add_argument('--prefill-only', action='store_true') - p.add_argument('--debug-layer', type=int, default=-1) + p.add_argument('--num-gpus', type=int, default=8) + p.add_argument('--checkpoint', type=str, default="/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4") return p.parse_args() _args = parse_args() -CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" +CHECKPOINT_DIR = _args.checkpoint MAX_NEW_TOKENS = _args.max_tokens PROMPT = _args.prompt or "The capital of France is" -NUM_GPUS = 8 +NUM_GPUS = _args.num_gpus SEED = _args.seed VERBOSE = _args.verbose THINK_START, THINK_END = 128821, 128822 USER_TOKEN, ASSISTANT_TOKEN = 128803, 128804 - FP4_LUT = torch.tensor([0., 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) # ===================================================================== -# RoPE cache (FP32 — BF16 destroys cos²+sin²=1) +# RoPE (FP32 — BF16 destroys cos²+sin²=1) # ===================================================================== def build_rope_cache(max_pos, rope_dim, device, theta=10000., rope_type="default", rope_factor=1., orig_max=4096, beta_fast=32, beta_slow=1): @@ -66,23 +58,30 @@ def build_rope_cache(max_pos, rope_dim, device, theta=10000., rope_type="default angles = torch.outer(torch.arange(max_pos, dtype=torch.float32), freqs) return torch.cos(angles).to(device), torch.sin(angles).to(device) +def _apply_rope(x, pos, cos, sin, rope_dim, inverse=False): + T, nh, hd = x.shape; nope = hd - rope_dim + if pos.device != cos.device: pos = pos.to(cos.device) + c, s = cos[pos].unsqueeze(1), sin[pos].unsqueeze(1) + xr = x[:, :, nope:].float(); ev, od = xr[..., 0::2], xr[..., 1::2] + if inverse: rev, rod = ev*c + od*s, -ev*s + od*c + else: rev, rod = ev*c - od*s, ev*s + od*c + out = x.clone(); ro = torch.empty_like(xr) + ro[..., 0::2], ro[..., 1::2] = rev, rod + out[:, :, nope:] = ro.bfloat16(); return out + # ===================================================================== # Weight loading # ===================================================================== def load_all_weights(checkpoint_dir): from safetensors.torch import load_file - cdir = Path(checkpoint_dir) - wmap = {} + cdir = Path(checkpoint_dir); wmap = {} idx = cdir / "model.safetensors.index.json" if idx.exists(): with open(idx) as f: wmap = json.load(f).get("weight_map", {}) - shards = set(wmap.values()) if wmap else set() - all_w = {} + shards = set(wmap.values()) if wmap else set(); all_w = {} for sn in sorted(shards): - if (cdir / sn).exists(): - all_w.update(load_file(str(cdir / sn))) - log.info(f"Loaded {len(all_w)} tensors from {len(shards)} shards") - return all_w + if (cdir / sn).exists(): all_w.update(load_file(str(cdir / sn))) + log.info(f"Loaded {len(all_w)} tensors from {len(shards)} shards"); return all_w # ===================================================================== # RMSNorm @@ -92,18 +91,14 @@ def rmsnorm(x, weight, eps=1e-6): return (xf * xf.pow(2).mean(-1, keepdim=True).add(eps).rsqrt() * weight.float()).bfloat16() def unweighted_rmsnorm(x, eps=1e-6): - xf = x.float() - return xf * xf.pow(2).mean(-1, keepdim=True).add(eps).rsqrt() + xf = x.float(); return xf * xf.pow(2).mean(-1, keepdim=True).add(eps).rsqrt() # ===================================================================== -# NVFP4 dequant — used ONLY for compressor/indexer projections -# (these don't go through the CuTeDSL GEMM kernel yet) +# NVFP4 ref dequant — compressor/indexer ONLY # ===================================================================== def dequant_nvfp4(weight, weight_scale, weight_scale_2=None, input_scale=None): - O, I2 = weight.shape - I = I2 * 2 - lo = (weight & 0x0F).to(torch.int8) - hi = (weight >> 4).to(torch.int8) + O, I2 = weight.shape; I = I2 * 2 + lo = (weight & 0x0F).to(torch.int8); hi = (weight >> 4).to(torch.int8) lut = FP4_LUT.to(device=weight.device, dtype=torch.float32) lo_f = lut[(lo & 0x07).long()] * torch.where((lo >> 3).bool(), -1., 1.) hi_f = lut[(hi & 0x07).long()] * torch.where((hi >> 3).bool(), -1., 1.) @@ -129,49 +124,39 @@ def do_nvfp4_linear_ref(x, w, pfx, proj_name): isc.to(d) if isc is not None else None) # ===================================================================== -# Production Nvfp4Linear wrapper +# Production Nvfp4Linear factory # ===================================================================== -def make_nvfp4_linear(in_features, out_features, device, weight, weight_scale, - weight_scale_2=None, input_scale=None): - """Create a production Nvfp4Linear with weights loaded from checkpoint.""" +def make_nvfp4_linear(in_features, out_features, device, all_w, pfx, proj_name): from dsv4.layers.linear import Nvfp4Linear d = device lin = Nvfp4Linear(in_features, out_features, max_num_tokens=8192, device=d) - lin.fp4 = [weight.to(d)] - lin.sf = [weight_scale.to(d)] - gs = input_scale.float().item() if input_scale is not None else 1.0 / (6.0 * 448.0) - lin.gs = [gs] - return lin + weight, ws, ws2, isc = get_nvfp4_weight(all_w, pfx, proj_name) + assert weight is not None, f"{pfx}.{proj_name}.weight not found" + lin.fp4 = [weight.to(d)]; lin.sf = [ws.to(d)] + gs = isc.float().item() if isc is not None else 1.0 / (6.0 * 448.0) + lin.gs = [gs]; lin.finalize_weights(); return lin # ===================================================================== -# Compressor — CSA (ratio=4) and HCA (ratio=128) -# (Reference PyTorch — compressor not yet on tensor cores) +# Compressor — CSA (ratio=4) and HCA (ratio=128) [PyTorch ref] # ===================================================================== class Compressor: def __init__(self, ratio, head_dim, hidden_size, device): self.ratio, self.hd, self.H, self.device = ratio, head_dim, hidden_size, device - self.is_csa = (ratio == 4) - self.kv_dim = 2 * head_dim if self.is_csa else head_dim + self.is_csa = (ratio == 4); self.kv_dim = 2 * head_dim if self.is_csa else head_dim self.wkv_w = self.wkv_ws = self.wkv_ws2 = self.wkv_isc = None self.wgate_w = self.wgate_ws = self.wgate_ws2 = self.wgate_isc = None - self.ape = None - self.kv_norm_w = None + self.ape = None; self.kv_norm_w = None def load(self, w, pfx): self.wkv_w, self.wkv_ws, self.wkv_ws2, self.wkv_isc = get_nvfp4_weight(w, pfx, 'kv_proj') self.wgate_w, self.wgate_ws, self.wgate_ws2, self.wgate_isc = get_nvfp4_weight(w, pfx, 'gate_proj') - self.ape = w.get(f"{pfx}.position_bias") - self.kv_norm_w = w.get(f"{pfx}.kv_norm.weight") + self.ape = w.get(f"{pfx}.position_bias"); self.kv_norm_w = w.get(f"{pfx}.kv_norm.weight") def forward(self, hidden_states, positions): - if self.ratio == 0 or self.wkv_w is None: - return None, None, None - T = hidden_states.shape[0] - r = self.ratio - dev = hidden_states.device + if self.ratio == 0 or self.wkv_w is None: return None, None, None + T = hidden_states.shape[0]; r = self.ratio; dev = hidden_states.device n_complete = T // r - if n_complete == 0: - return None, None, None + if n_complete == 0: return None, None, None kv = nvfp4_linear_ref(hidden_states, self.wkv_w.to(dev), self.wkv_ws.to(dev), self.wkv_ws2.to(dev) if self.wkv_ws2 is not None else None, self.wkv_isc.to(dev) if self.wkv_isc is not None else None) @@ -182,50 +167,40 @@ class Compressor: ape = self.ape.to(dev) for bi in range(T // r): s, e = bi * r, (bi + 1) * r - kv[s:e] += ape.to(kv.dtype) - gate[s:e] += ape.to(gate.dtype) - T_comp = n_complete * r - comp_list, comp_pos_list = [], [] + kv[s:e] += ape.to(kv.dtype); gate[s:e] += ape.to(gate.dtype) + T_comp = n_complete * r; comp_list, comp_pos_list = [], [] if self.is_csa: Ca = kv[:T_comp, :self.hd].reshape(n_complete, r, self.hd) Cb = kv[:T_comp, self.hd:].reshape(n_complete, r, self.hd) Ga = gate[:T_comp, :self.hd].reshape(n_complete, r, self.hd) Gb = gate[:T_comp, self.hd:].reshape(n_complete, r, self.hd) for bi in range(n_complete): - if bi > 0: - block_kv = torch.cat([Ca[bi-1], Cb[bi]], dim=0) - block_gate = torch.cat([Ga[bi-1], Gb[bi]], dim=0) - else: - block_kv = Cb[bi]; block_gate = Gb[bi] - probs = torch.softmax(block_gate.float(), dim=0) - compressed = (probs * block_kv.float()).sum(0) + if bi > 0: block_kv = torch.cat([Ca[bi-1], Cb[bi]], dim=0); block_gate = torch.cat([Ga[bi-1], Gb[bi]], dim=0) + else: block_kv = Cb[bi]; block_gate = Gb[bi] + probs = torch.softmax(block_gate.float(), dim=0); compressed = (probs * block_kv.float()).sum(0) if self.kv_norm_w is not None: nw = self.kv_norm_w.to(dev).float() compressed = compressed * compressed.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() * nw - comp_list.append(compressed.bfloat16()) - comp_pos_list.append(positions[(bi+1)*r - 1]) + comp_list.append(compressed.bfloat16()); comp_pos_list.append(positions[(bi+1)*r - 1]) else: kv_blocks = kv[:T_comp].reshape(n_complete, r, self.hd) gate_blocks = gate[:T_comp].reshape(n_complete, r, self.hd) for bi in range(n_complete): - probs = torch.softmax(gate_blocks[bi].float(), dim=0) - compressed = (probs * kv_blocks[bi].float()).sum(0) + probs = torch.softmax(gate_blocks[bi].float(), dim=0); compressed = (probs * kv_blocks[bi].float()).sum(0) if self.kv_norm_w is not None: nw = self.kv_norm_w.to(dev).float() compressed = compressed * compressed.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() * nw - comp_list.append(compressed.bfloat16()) - comp_pos_list.append(positions[(bi+1)*r - 1]) + comp_list.append(compressed.bfloat16()); comp_pos_list.append(positions[(bi+1)*r - 1]) return torch.stack(comp_list), torch.stack(comp_pos_list), torch.zeros(1, T, n_complete, dtype=torch.float32, device=dev) # ===================================================================== -# Indexer — CSA top-k (Reference PyTorch) +# Indexer — CSA top-k [PyTorch ref] # ===================================================================== class Indexer: def __init__(self, n_ih, ihd, top_k, device): self.n_ih, self.ihd, self.top_k, self.device = n_ih, ihd, top_k, device self.q_b_w = self.q_b_ws = self.q_b_ws2 = self.q_b_isc = None - self.wp_w = self.wp_ws = self.wp_ws2 = self.wp_isc = None - self.compressor = None + self.wp_w = self.wp_ws = self.wp_ws2 = self.wp_isc = None; self.compressor = None def load(self, w, pfx): self.q_b_w, self.q_b_ws, self.q_b_ws2, self.q_b_isc = get_nvfp4_weight(w, pfx, 'q_b_proj') @@ -235,11 +210,8 @@ class Indexer: self.compressor.load(w, f"{pfx}.compressor") def forward(self, q_lora, hidden_states, comp_indexer_kv, positions): - if self.q_b_w is None or comp_indexer_kv is None or comp_indexer_kv.shape[0] == 0: - return None - dev = q_lora.device - T = q_lora.shape[0] - n_comp = comp_indexer_kv.shape[0] + if self.q_b_w is None or comp_indexer_kv is None or comp_indexer_kv.shape[0] == 0: return None + dev = q_lora.device; T = q_lora.shape[0]; n_comp = comp_indexer_kv.shape[0] q_idx = nvfp4_linear_ref(q_lora, self.q_b_w.to(dev), self.q_b_ws.to(dev), self.q_b_ws2.to(dev) if self.q_b_ws2 is not None else None, self.q_b_isc.to(dev) if self.q_b_isc is not None else None) @@ -249,11 +221,8 @@ class Indexer: self.wp_isc.to(dev) if self.wp_isc is not None else None) k_idx = comp_indexer_kv.reshape(n_comp, self.n_ih, self.ihd) scores = torch.einsum('tnd,cnd->tnc', q_idx.float(), k_idx.float()) - scores = F.relu(scores) - total = (scores * w_h.unsqueeze(-1).float()).sum(1) - tk = min(self.top_k, n_comp) - _, idx = total.topk(tk, -1) - return idx + scores = F.relu(scores); total = (scores * w_h.unsqueeze(-1).float()).sum(1) + tk = min(self.top_k, n_comp); _, idx = total.topk(tk, -1); return idx # ===================================================================== # KV Cache @@ -264,16 +233,13 @@ class KVCache: self.swa = torch.zeros(window_size, head_dim, dtype=torch.bfloat16, device=device) self.swa_pos = torch.zeros(window_size, dtype=torch.long, device=device) self.swa_len, self.swa_head = 0, 0 - self.comp_kv, self.comp_pos, self.n_comp = None, None, 0 - self.comp_idx_kv = None + self.comp_kv, self.comp_pos, self.n_comp = None, None, 0; self.comp_idx_kv = None def append_swa(self, kv, pos): T = kv.shape[0] for i in range(T): - idx = (self.swa_head + i) % self.ws - self.swa[idx], self.swa_pos[idx] = kv[i], pos[i] - self.swa_head = (self.swa_head + T) % self.ws - self.swa_len = min(self.swa_len + T, self.ws) + idx = (self.swa_head + i) % self.ws; self.swa[idx], self.swa_pos[idx] = kv[i], pos[i] + self.swa_head = (self.swa_head + T) % self.ws; self.swa_len = min(self.swa_len + T, self.ws) def add_compressed(self, ckv, cpos, idx_kv=None): if ckv is None: return @@ -285,136 +251,67 @@ class KVCache: def get_swa(self): if self.swa_len == 0: - return torch.zeros(0, self.hd, device=self.dev, dtype=torch.bfloat16), \ - torch.zeros(0, device=self.dev, dtype=torch.long) - if self.swa_len < self.ws: - return self.swa[:self.swa_len].clone(), self.swa_pos[:self.swa_len].clone() + return torch.zeros(0, self.hd, device=self.dev, dtype=torch.bfloat16), torch.zeros(0, device=self.dev, dtype=torch.long) + if self.swa_len < self.ws: return self.swa[:self.swa_len].clone(), self.swa_pos[:self.swa_len].clone() idx = torch.arange(self.swa_head, self.swa_head + self.ws) % self.ws return self.swa[idx].clone(), self.swa_pos[idx].clone() # ===================================================================== -# RoPE apply (FP32 cache, partial, GPT-J interleaved) -# ===================================================================== -def _apply_rope(x, pos, cos, sin, rope_dim, inverse=False): - T, nh, hd = x.shape - nope = hd - rope_dim - if pos.device != cos.device: - pos = pos.to(cos.device) - c, s = cos[pos].unsqueeze(1), sin[pos].unsqueeze(1) - xr = x[:, :, nope:].float() - ev, od = xr[..., 0::2], xr[..., 1::2] - if inverse: rev, rod = ev*c + od*s, -ev*s + od*c - else: rev, rod = ev*c - od*s, ev*s + od*c - out = x.clone() - ro = torch.empty_like(xr) - ro[..., 0::2], ro[..., 1::2] = rev, rod - out[:, :, nope:] = ro.bfloat16() - return out - -# ===================================================================== -# HcHead — FP32 projection, read out from mHC state +# HcHead # ===================================================================== HC_EPS = 1e-6 - class HcHead: def __init__(self, hidden_dim=7168, n_hc=4, device='cuda:0'): self.K, self.device, self.n_hc = n_hc * hidden_dim, device, n_hc - def load(self, fn, base, scale=None): self.fn = fn.to(self.device, torch.float32).contiguous() self.base = base.to(self.device, torch.float32).contiguous() self.scale = scale.to(self.device, torch.float32).item() if scale is not None else 1.0 - def forward(self, X): - T = X.shape[0] - Xn = unweighted_rmsnorm(X.reshape(T, self.K).bfloat16()) + T = X.shape[0]; Xn = unweighted_rmsnorm(X.reshape(T, self.K).bfloat16()) mix = F.linear(Xn, self.fn[:self.n_hc]).float() pre = torch.sigmoid(mix * self.scale + self.base[:self.n_hc].unsqueeze(0)) + HC_EPS return (pre.unsqueeze(-1) * X.float()).sum(1).bfloat16() # ===================================================================== -# Production FMHA — 6-warp TMA multi-tile kernel with sink bias +# Production FMHA # ===================================================================== def _run_production_fmha(q_heads, all_kv, n_h, hd, T, seq_len, scale, dev, li, w, pfx): - """Run production FMHA kernel with sink bias support. - - The kernel handles: - - N < 128: K/V padded to 128, kernel uses N_orig for softmax masking - - Multi-tile KV for N > 128 - - Attention sinks via per-head logit bias (D5c: single softmax) - """ - from dsv4.kernels.attention.production import dsv4_attention - - q = q_heads.permute(1, 0, 2).contiguous() # (n_h, T, hd) - k = all_kv.unsqueeze(0).contiguous() # (1, seq_len, hd) — MQA - v = all_kv.unsqueeze(0).contiguous() - - sinks = w.get(f"{pfx}.sinks") - sink_bias = None - if sinks is not None: - sink_bias = sinks.to(device=dev).float().reshape(n_h) - - attn_out = dsv4_attention( - q=q, k=k, v=v, scale=scale, - n_comp=0, sink_bias=sink_bias, - ) # (n_h, T, hd) - return attn_out.permute(1, 0, 2) # (T, n_h, hd) + from dsv4.kernels.attention.production import dsv4_attention_per_head + q = q_heads.permute(1, 0, 2).contiguous(); k = all_kv.unsqueeze(0).contiguous(); v = k.clone() + sinks = w.get(f"{pfx}.sinks"); sink_bias = None + if sinks is not None: sink_bias = sinks.to(device=dev).float().reshape(n_h) + attn_out = dsv4_attention_per_head(q=q, k=k, v=v, scale=scale, n_comp=0, sink_bias=sink_bias) + return attn_out.permute(1, 0, 2) # ===================================================================== -# Attention forward — production FMHA + production Nvfp4Linear +# Attention — ALL production kernels # ===================================================================== def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, - kv_cache, positions, compressor, indexer, - prod_lin=None): - """Attention sub-block using production kernels. - - All projections go through Nvfp4Linear (CuTeDSL GEMM). - FMHA goes through 6-warp TMA multi-tile kernel with sink bias. - Inverse RoPE applied after FMHA. - """ - 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) + kv_cache, positions, compressor, indexer, prod_lin): + 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) ratio = compressor.ratio if compressor is not None else 0 - scale = 1.0 / math.sqrt(hd) - pfx = f"model.layers.{li}.self_attn" - if positions.device != rope_cos.device: - positions = positions.to(rope_cos.device) + scale = 1.0 / math.sqrt(hd); pfx = f"model.layers.{li}.self_attn" + if positions.device != rope_cos.device: positions = positions.to(rope_cos.device) - # 1. Q projection: q_a → q_a_norm → q_b → q_b_norm - q_a = prod_lin['q_a'](x_normed) if prod_lin and 'q_a' in prod_lin else \ - do_nvfp4_linear_ref(x_normed, w, pfx, 'q_a_proj') - if q_a is None: - log.warning(f" L{li}: q_a_proj not found") - return torch.zeros(T, cfg["hidden_size"], dtype=torch.bfloat16, device=dev), None + # 1. Q: q_a (NVFP4 GEMM) → q_a_norm → q_b (NVFP4 GEMM) → q_b_norm + q_a = prod_lin['q_a'](x_normed) q_norm_w = w.get(f"{pfx}.q_a_norm.weight") if q_norm_w is not None: q_a = rmsnorm(q_a, q_norm_w.to(dev, torch.float32)) - q = prod_lin['q_b'](q_a) if prod_lin and 'q_b' in prod_lin else \ - do_nvfp4_linear_ref(q_a, w, pfx, 'q_b_proj') - q = unweighted_rmsnorm(q).bfloat16() - q_heads = q.reshape(T, n_h, hd) - q_heads = _apply_rope(q_heads, positions, rope_cos, rope_sin, rd) + q = prod_lin['q_b'](q_a); q = unweighted_rmsnorm(q).bfloat16() + q_heads = q.reshape(T, n_h, hd); q_heads = _apply_rope(q_heads, positions, rope_cos, rope_sin, rd) - # 2. KV projection (MQA, single KV head, hd dim) - kv = prod_lin['kv'](x_normed) if prod_lin and 'kv' in prod_lin else \ - do_nvfp4_linear_ref(x_normed, w, pfx, 'kv_proj') - if kv is None: - log.warning(f" L{li}: kv_proj not found") - return torch.zeros(T, cfg["hidden_size"], dtype=torch.bfloat16, device=dev), q_a + # 2. KV (NVFP4 GEMM, MQA, single KV head) + kv = prod_lin['kv'](x_normed) 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(T, 1, hd) - kv_3d = _apply_rope(kv_3d, positions, rope_cos, rope_sin, rd) - kv_roped = kv_3d.reshape(T, hd) - kv_cache.append_swa(kv_roped, positions) + kv_3d = kv.reshape(T, 1, hd); kv_3d = _apply_rope(kv_3d, positions, rope_cos, rope_sin, rd) + kv_roped = kv_3d.reshape(T, hd); kv_cache.append_swa(kv_roped, positions) # 3. Compressor → compressed KV - comp_kv, comp_pos, block_bias = None, None, None - comp_idx_kv = None + comp_kv, comp_pos, block_bias = None, None, None; comp_idx_kv = None if compressor is not None and compressor.ratio > 0: comp_kv, comp_pos, block_bias = compressor.forward(x_normed, positions) if comp_kv is not None: @@ -425,64 +322,47 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, comp_idx_kv, _, _ = indexer.compressor.forward(x_normed, positions) kv_cache.add_compressed(comp_kv, comp_pos, comp_idx_kv) - # 4. Indexer top-k (CSA only) + # 4. 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.comp_idx_kv, positions) - # 5. Gather full KV: [compressed, swa] + # 5. Gather KV swa_kv, swa_pos = kv_cache.get_swa() if kv_cache.comp_kv is not None and kv_cache.n_comp > 0: if ratio == 4 and topk_idx is not None: tk = topk_idx[0].clamp(0, kv_cache.n_comp - 1) - sel_comp = kv_cache.comp_kv[tk] - all_kv = torch.cat([sel_comp, swa_kv], dim=0) - elif ratio > 4: - all_kv = torch.cat([kv_cache.comp_kv, swa_kv], dim=0) - else: - all_kv = swa_kv - else: - all_kv = swa_kv - + all_kv = torch.cat([kv_cache.comp_kv[tk], swa_kv], dim=0) + elif ratio > 4: all_kv = torch.cat([kv_cache.comp_kv, swa_kv], dim=0) + else: all_kv = swa_kv + else: all_kv = swa_kv seq_len = all_kv.shape[0] - if seq_len == 0: - return torch.zeros(T, cfg["hidden_size"], dtype=torch.bfloat16, device=dev), q_a + if seq_len == 0: return torch.zeros(T, cfg["hidden_size"], dtype=torch.bfloat16, device=dev), q_a - # 6. Production FMHA kernel (6-warp TMA multi-tile) with sink bias + # 6. Production FMHA attn_out = _run_production_fmha(q_heads, all_kv, n_h, hd, T, seq_len, scale, dev, li, w, pfx) - - # 7. Inverse RoPE (FP32 cache) + # 7. Inverse RoPE attn_out = _apply_rope(attn_out, positions, rope_cos, rope_sin, rd, inverse=True) - # 8. Output projection: wo_a (BF16 grouped BMM) + wo_b (NVFP4 GEMM) - hpg = n_h // o_groups - gid = hpg * hd + # 8. Output: wo_a (BF16 grouped BMM) + wo_b (NVFP4 GEMM) + hpg = n_h // o_groups; gid = hpg * hd oa_w = w.get(f"{pfx}.o_a_proj.weight") if oa_w is not None: - oa_bf = oa_w.bfloat16().to(dev) - a_flat = attn_out.reshape(T, n_h * hd) - a_grp = a_flat.reshape(T, o_groups, gid) - oa_3d = oa_bf.reshape(o_groups, o_rank, gid) + oa_bf = oa_w.bfloat16().to(dev); a_flat = attn_out.reshape(T, n_h * hd) + a_grp = a_flat.reshape(T, o_groups, gid); oa_3d = oa_bf.reshape(o_groups, o_rank, gid) g_out = torch.bmm(a_grp.permute(1, 0, 2), oa_3d.transpose(1, 2)) g_flat = g_out.permute(1, 0, 2).reshape(T, o_groups * o_rank) - F_attn = prod_lin['o_b'](g_flat) if prod_lin and 'o_b' in prod_lin else \ - do_nvfp4_linear_ref(g_flat, w, pfx, 'o_b_proj') + F_attn = prod_lin['o_b'](g_flat) else: - F_attn = prod_lin['o_a'](attn_out.reshape(T, n_h * hd)) if prod_lin and 'o_a' in prod_lin else \ - do_nvfp4_linear_ref(attn_out.reshape(T, n_h * hd), w, pfx, 'o_a_proj') + F_attn = prod_lin['o_a'](attn_out.reshape(T, n_h * hd)) return F_attn, q_a # ===================================================================== -# MoE forward — production Nvfp4MoE + Nvfp4SharedExpert + Router +# MoE — production kernels # ===================================================================== def moe_forward(x, li, moe_runner, se_runner, router, token_id): - """MoE forward using production NVFP4 GEMM kernels. - - NO fallback to reference. Production kernels ONLY. - """ topk_w, topk_ids = router(x, token_ids=token_id) - routed_out = moe_runner(x, topk_w, topk_ids) - shared_out = se_runner(x) + routed_out = moe_runner(x, topk_w, topk_ids); shared_out = se_runner(x) return routed_out + shared_out # ===================================================================== @@ -494,147 +374,97 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, compressor=None, indexer=None, moe_runner=None, se_runner=None, router=None, prod_lin=None): - dev = X_l.device - # Attention sub-block - x_in, ctx_a = attn_mhc.pre_block(X_l) - x_normed = rmsnorm(x_in, attn_norm_w) + x_in, ctx_a = attn_mhc.pre_block(X_l); x_normed = rmsnorm(x_in, attn_norm_w) F_attn, _ = forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, - kv_cache, positions, compressor, indexer, - prod_lin=prod_lin) + kv_cache, positions, compressor, indexer, prod_lin) X_mid = attn_mhc.post_block(X_l, F_attn, ctx_a) - # FFN sub-block - x_in_f, ctx_f = ffn_mhc.pre_block(X_mid) - x_ffn = rmsnorm(x_in_f, ffn_norm_w) + x_in_f, ctx_f = ffn_mhc.pre_block(X_mid); x_ffn = rmsnorm(x_in_f, ffn_norm_w) F_ffn = moe_forward(x_ffn, li, moe_runner, se_runner, router, token_id) X_next = ffn_mhc.post_block(X_mid, F_ffn, ctx_f) if VERBOSE >= 1: - print(f" L{li}: |X|={X_l.abs().max().item():.1f}→{X_next.abs().max().item():.1f} " + print(f" L{li}: |X|={X_l.abs().max().item():.1f}->{X_next.abs().max().item():.1f} " f"|Fa|={F_attn.abs().max().item():.1f} |Ff|={F_ffn.abs().max().item():.1f}", flush=True) return X_next # ===================================================================== -# MoE weight loading (stacked path for production GEMM) +# MoE weight loading # ===================================================================== def _load_moe_weights_stacked(all_w, li, pfx, dev, moe, cfg): n_e = cfg["n_routed_experts"] - # Load expert weights directly to GPU, one at a time - # to avoid creating huge CPU tensors l1_fp4_list, l1_sf_list, l1_gs_list = [], [], [] l2_fp4_list, l2_sf_list, l2_gs_list = [], [], [] for eid in range(n_e): ep = f"{pfx}.experts.{eid}" - # L1: gate + up gw, gws, _, gisc = get_nvfp4_weight(all_w, ep, 'gate_proj') uw, uws, _, uisc = get_nvfp4_weight(all_w, ep, 'up_proj') if gw is not None and uw is not None: - # Stack gate and up along dim 0 → (2*N, K) uint8 l1_fp4_list.append(torch.cat([gw, uw], dim=0).to(dev)) - if gws is not None and uws is not None: - l1_sf_list.append(torch.cat([gws, uws], dim=0).to(dev)) + if gws is not None and uws is not None: l1_sf_list.append(torch.cat([gws, uws], dim=0).to(dev)) gs = gisc.float().item() if gisc is not None else 1.0 / (6.0 * 448.0) l1_gs_list.append(gs) - # L2: down dw, dws, _, disc = get_nvfp4_weight(all_w, ep, 'down_proj') if dw is not None: l2_fp4_list.append(dw.to(dev)) - if dws is not None: - l2_sf_list.append(dws.to(dev)) + if dws is not None: l2_sf_list.append(dws.to(dev)) gs2 = disc.float().item() if disc is not None else 1.0 / (6.0 * 448.0) l2_gs_list.append(gs2) - - if not l1_fp4_list: - log.warning(f"L{li}: No expert weights found") - return - - # Stack into (E, N, K) tensors on GPU + if not l1_fp4_list: log.warning(f"L{li}: No expert weights found"); return l1_stacked = torch.stack(l1_fp4_list).to(dev) l1_sf_stacked = torch.stack(l1_sf_list).to(dev) if l1_sf_list else None l2_stacked = torch.stack(l2_fp4_list).to(dev) if l2_fp4_list else None l2_sf_stacked = torch.stack(l2_sf_list).to(dev) if l2_sf_list else None del l1_fp4_list, l1_sf_list, l2_fp4_list, l2_sf_list - - moe.prepare_weights_from_stacked( - l1_stacked, l1_sf_stacked, l1_gs_list, - l2_stacked, l2_sf_stacked, l2_gs_list, - ) - + moe.prepare_weights_from_stacked(l1_stacked, l1_sf_stacked, l1_gs_list, l2_stacked, l2_sf_stacked, l2_gs_list) def _load_shared_expert_weights(all_w, li, pfx, dev, se, cfg): - moe_inter = cfg.get('moe_intermediate_size', 3072) - # Shared expert: gate_proj + up_proj → L1, down_proj → L2 gw, gws, _, gisc = get_nvfp4_weight(all_w, f"{pfx}.shared_experts", 'gate_proj') uw, uws, _, uisc = get_nvfp4_weight(all_w, f"{pfx}.shared_experts", 'up_proj') dw, dws, _, disc = get_nvfp4_weight(all_w, f"{pfx}.shared_experts", 'down_proj') - if gw is not None and uw is not None: se.l1_fp4 = [torch.cat([gw, uw], dim=0).to(dev)] - if gws is not None and uws is not None: - se.l1_sf = [torch.cat([gws, uws], dim=0).to(dev)] - else: - se.l1_sf = [torch.zeros(1, device=dev, dtype=torch.float8_e4m3fn)] - gs = gisc.float().item() if gisc is not None else 1.0 / (6.0 * 448.0) - se.l1_gs = [gs] + se.l1_sf = [torch.cat([gws, uws], dim=0).to(dev)] if gws is not None and uws is not None else [torch.zeros(1, device=dev, dtype=torch.float8_e4m3fn)] + se.l1_gs = [gisc.float().item() if gisc is not None else 1.0 / (6.0 * 448.0)] if dw is not None: se.l2_fp4 = [dw.to(dev)] se.l2_sf = [dws.to(dev)] if dws is not None else [torch.zeros(1, device=dev, dtype=torch.float8_e4m3fn)] - gs2 = disc.float().item() if disc is not None else 1.0 / (6.0 * 448.0) - se.l2_gs = [gs2] - # finalize_weights called lazily by Nvfp4SharedExpert._ensure_initialized() - + se.l2_gs = [disc.float().item() if disc is not None else 1.0 / (6.0 * 448.0)] def _cache_layer_weights_no_experts(all_w, n_layers, devices): - """Cache per-layer weights to GPUs, EXCLUDING MoE expert weights.""" cached = {} for li in range(n_layers): - dev = devices[li % len(devices)] - pfx = f"model.layers.{li}." - w = {k: v.to(device=dev, non_blocking=True) - for k, v in all_w.items() + dev = devices[li % len(devices)]; pfx = f"model.layers.{li}." + w = {k: v.to(device=dev, non_blocking=True) for k, v in all_w.items() if k.startswith(pfx) and '.experts.' not in k and '.shared_experts.' not in k} cached[li] = w - if (li+1) % 10 == 0: print(f" Cached {li+1}/{n_layers} layers") + if (li+1) % 10 == 0: log.info(f" Cached {li+1}/{n_layers} layers") return cached - -def load_weights(checkpoint_dir): - from safetensors.torch import load_file - cdir = Path(checkpoint_dir) - wmap = {} - idx = cdir / "model.safetensors.index.json" - if idx.exists(): - with open(idx) as f: wmap = json.load(f).get("weight_map", {}) - shards = set(wmap.values()) if wmap else set() - all_w = {} - for sn in sorted(shards): - if (cdir / sn).exists(): - all_w.update(load_file(str(cdir / sn))) - return all_w - - +# ===================================================================== +# Main +# ===================================================================== def main(): - t0 = time.time() - torch.manual_seed(SEED) + t0 = time.time(); torch.manual_seed(SEED) print("=" * 70) - print("DSV4 Single-Shot Inference — PRODUCTION KERNEL STACK") + print("DSV4 Single-Shot Inference - PRODUCTION KERNEL STACK") print(" FMHA: 6-warp TMA multi-tile + sink bias") - print(" NVFP4 GEMM (CuTeDSL) | Router kernels | Production mHC") + print(" NVFP4 GEMM (CuTeDSL) for ALL projections") + print(" Production MoE + Router | Production mHC") print(" NO PyTorch SDPA | NO dequant+matmul | NO reference fallback") print("=" * 70) with open(os.path.join(CHECKPOINT_DIR, "config.json")) as f: cfg = json.load(f) - n_layers = cfg["num_hidden_layers"] - H = cfg["hidden_size"] - hd = cfg["head_dim"] + n_layers = cfg["num_hidden_layers"]; H = cfg["hidden_size"] + hd = cfg["head_dim"]; n_h = cfg["num_attention_heads"] rd = cfg.get("qk_rope_head_dim", 64) cr = cfg.get("compress_ratios", [128] * n_layers) - print(f"Model: {n_layers} layers, {cfg['num_attention_heads']} heads, hd={hd}, rope_dim={rd}") + o_groups = cfg.get("o_groups", 16); o_rank = cfg.get("o_lora_rank", 1024) + print(f"Model: {n_layers} layers, {n_h} heads, hd={hd}, rope_dim={rd}") print(f"Compress ratios: first5={cr[:5]} len={len(cr)}") print(f"Experts: {cfg['n_routed_experts']}, top-{cfg.get('num_experts_per_tok', 6)}") # ---- Phase 1: Load weights ---- - print(f"\nPhase 1: Loading weights...") - all_w = load_all_weights(CHECKPOINT_DIR) + print(f"\nPhase 1: Loading weights..."); all_w = load_all_weights(CHECKPOINT_DIR) print(f" {time.time()-t0:.1f}s") # ---- Phase 2: Build production components ---- @@ -643,12 +473,9 @@ def main(): from dsv4.layers.router import Router from dsv4.layers.moe import Nvfp4MoE from dsv4.layers.shared_expert import Nvfp4SharedExpert - from dsv4.layers.linear import Nvfp4Linear - # Kill stale GPU processes (safety) - for g in range(NUM_GPUS): - torch.cuda.set_device(g) - torch.cuda.empty_cache() + # Kill stale GPU processes + for g in range(NUM_GPUS): torch.cuda.set_device(g); torch.cuda.empty_cache() torch.cuda.set_device(0) # mHC + norms @@ -656,89 +483,77 @@ def main(): for li in range(n_layers): dev = f"cuda:{li % NUM_GPUS}" for tag, blocks, fn_s, base_s, scale_s in [ - ("attn", attn_mhcs, f"model.layers.{li}.attn_hc.fn", - f"model.layers.{li}.attn_hc.base", f"model.layers.{li}.attn_hc.scale"), - ("ffn", ffn_mhcs, f"model.layers.{li}.ffn_hc.fn", - f"model.layers.{li}.ffn_hc.base", f"model.layers.{li}.ffn_hc.scale"), + ("attn", attn_mhcs, f"model.layers.{li}.attn_hc.fn", f"model.layers.{li}.attn_hc.base", f"model.layers.{li}.attn_hc.scale"), + ("ffn", ffn_mhcs, f"model.layers.{li}.ffn_hc.fn", f"model.layers.{li}.ffn_hc.base", f"model.layers.{li}.ffn_hc.scale"), ]: fn, base, scale = all_w.get(fn_s), all_w.get(base_s), all_w.get(scale_s) if fn is not None and base is not None and scale is not None: m = mHCLayer(hidden_dim=H, n_hc=4, t_max_sinkhorn=20, device=dev) - # Split fn/base/scale into pre/post/comb n = 4 m.load_weights( - W_pre=fn[0:n].to(dev, torch.float32), - W_post=fn[n:2*n].to(dev, torch.float32), + W_pre=fn[0:n].to(dev, torch.float32), W_post=fn[n:2*n].to(dev, torch.float32), W_comb=fn[2*n:].to(dev, torch.float32), S_pre=base[0:n].reshape(1, n).to(dev, torch.float32), S_post=base[n:2*n].reshape(n, 1).to(dev, torch.float32), S_comb=base[2*n:].reshape(n, n).to(dev, torch.float32), - alpha_pre=scale[0].item(), - alpha_post=scale[1].item(), - alpha_comb=scale[2].item(), + alpha_pre=scale[0].item(), alpha_post=scale[1].item(), alpha_comb=scale[2].item(), ) blocks[li] = m - an_k = f"model.layers.{li}.input_layernorm.weight" if an_k in all_w: attn_norms[li] = all_w[an_k].to(dev, torch.float32) fn_k = f"model.layers.{li}.post_attention_layernorm.weight" if fn_k in all_w: ffn_norms[li] = all_w[fn_k].to(dev, torch.float32) # Production Nvfp4Linear for attention projections - # Nvfp4Linear for attention projections (deferred — use reference for now) - # Production MoE + Router + FMHA are the critical paths. - # Nvfp4Linear for small projections can be enabled once MoE is validated. - prod_lins = {} # Empty = use reference dequant path - print(" Using reference dequant for attention projections") + print(" Building production Nvfp4Linear for attention projections...") + prod_lins = {} + # Weight dimensions (from checkpoint): + # q_a_proj: (1536, 3584) uint8 -> in=7168, out=1536 + # q_b_proj: (65536, 768) uint8 -> in=1536, out=65536 + # kv_proj: (512, 3584) uint8 -> in=7168, out=512 + # o_b_proj: (7168, 8192) uint8 -> in=16384, out=7168 + for li in range(n_layers): + dev = f"cuda:{li % NUM_GPUS}"; pfx = f"model.layers.{li}.self_attn" + torch.cuda.set_device(li % NUM_GPUS) + pl = {} + pl['q_a'] = make_nvfp4_linear(7168, 1536, dev, all_w, pfx, 'q_a_proj') + pl['q_b'] = make_nvfp4_linear(1536, 65536, dev, all_w, pfx, 'q_b_proj') + pl['kv'] = make_nvfp4_linear(7168, 512, dev, all_w, pfx, 'kv_proj') + pl['o_b'] = make_nvfp4_linear(16384, 7168, dev, all_w, pfx, 'o_b_proj') + prod_lins[li] = pl + if (li+1) % 10 == 0: print(f" {li+1}/{n_layers} layers") + print(" All attention projections: production NVFP4 GEMM") # Routers, MoE, shared experts routers, moe_runners, se_runners = {}, {}, {} for li in range(n_layers): - dev = f"cuda:{li % NUM_GPUS}" - pfx = f"model.layers.{li}.mlp" - torch.cuda.set_device(li % NUM_GPUS) - # Verify GPU is in good state before MoE loading - torch.cuda.synchronize() + dev = f"cuda:{li % NUM_GPUS}"; pfx = f"model.layers.{li}.mlp" + torch.cuda.set_device(li % NUM_GPUS); torch.cuda.synchronize() is_hash = (li < cfg.get("num_hash_layers", 3)) and (f"{pfx}.gate.tid2eid" in all_w) - - router = Router( - hidden_size=H, num_experts=cfg["n_routed_experts"], - top_k=cfg.get("num_experts_per_tok", 6), - routed_scaling_factor=cfg.get("routed_scaling_factor", 2.5), - mode="hash" if is_hash else "dense", - vocab_size=cfg.get("vocab_size", 128000) if is_hash else None, - device=dev, - ) + router = Router(hidden_size=H, num_experts=cfg["n_routed_experts"], + top_k=cfg.get("num_experts_per_tok", 6), + routed_scaling_factor=cfg.get("routed_scaling_factor", 2.5), + mode="hash" if is_hash else "dense", + vocab_size=cfg.get("vocab_size", 128000) if is_hash else None, device=dev) if is_hash: router.load_weights(hash_lut=all_w[f"{pfx}.gate.tid2eid"].to(dev, torch.int32)) else: - gw = all_w.get(f"{pfx}.gate.weight") - eb = all_w.get(f"{pfx}.gate.e_score_correction_bias") + gw = all_w.get(f"{pfx}.gate.weight"); eb = all_w.get(f"{pfx}.gate.e_score_correction_bias") if gw is not None and eb is not None: - # Checkpoint may store gate weight transposed: (n_experts, hidden) vs (hidden, n_experts) - if gw.shape == (cfg["n_routed_experts"], H): - gw = gw.T.contiguous() + if gw.shape == (cfg["n_routed_experts"], H): gw = gw.T.contiguous() router.load_weights(W_gate=gw.bfloat16().to(dev), e_bias=eb.to(dev, torch.float32)) - router.finalize_weights() - routers[li] = router + router.finalize_weights(); routers[li] = router - moe = Nvfp4MoE( - num_experts=cfg["n_routed_experts"], hidden_size=H, - intermediate_size=cfg.get("moe_intermediate_size", 3072), - top_k=cfg.get("num_experts_per_tok", 6), device=dev, - ) + moe = Nvfp4MoE(num_experts=cfg["n_routed_experts"], hidden_size=H, + intermediate_size=cfg.get("moe_intermediate_size", 3072), + top_k=cfg.get("num_experts_per_tok", 6), device=dev) moe.set_swiglu_limit(cfg.get("swiglu_limit", 10.0)) - _load_moe_weights_stacked(all_w, li, pfx, dev, moe, cfg) - moe_runners[li] = moe + _load_moe_weights_stacked(all_w, li, pfx, dev, moe, cfg); moe_runners[li] = moe - se = Nvfp4SharedExpert( - hidden_size=H, intermediate_size=cfg.get("moe_intermediate_size", 3072), - device=dev, swiglu_limit=cfg.get("swiglu_limit", 10.0), - ) - _load_shared_expert_weights(all_w, li, pfx, dev, se, cfg) - se_runners[li] = se - - if (li+1) % 10 == 0: print(f" Built {li+1}/{n_layers} layers") + se = Nvfp4SharedExpert(hidden_size=H, intermediate_size=cfg.get("moe_intermediate_size", 3072), + device=dev, swiglu_limit=cfg.get("swiglu_limit", 10.0)) + _load_shared_expert_weights(all_w, li, pfx, dev, se, cfg); se_runners[li] = se + if (li+1) % 10 == 0: print(f" Built {li+1}/{n_layers} MoE layers") # Global weights torch.cuda.set_device(0) @@ -749,42 +564,31 @@ def main(): if final_norm_w is not None: final_norm_w = final_norm_w.to('cuda:0', torch.float32) hc_head = HcHead(H, 4, 'cuda:0') - hc_fn = all_w.get("model.hc_head.hc_fn") - hc_base = all_w.get("model.hc_head.hc_base") - hc_scale = all_w.get("model.hc_head.hc_scale") - if hc_fn is not None and hc_base is not None: - hc_head.load(hc_fn, hc_base, hc_scale) - print(" hc_head loaded") + hc_fn = all_w.get("model.hc_head.hc_fn"); hc_base = all_w.get("model.hc_head.hc_base"); hc_scale = all_w.get("model.hc_head.hc_scale") + if hc_fn is not None and hc_base is not None: hc_head.load(hc_fn, hc_base, hc_scale); print(" hc_head loaded") # RoPE (FP32) rp = cfg.get("rope_scaling", cfg.get("rope_parameters", {})) - rt = rp.get("type", rp.get("rope_type", "yarn")) - rf = rp.get("factor", 16.0) - rtheta = cfg.get("rope_theta", 10000.) - romax = rp.get("original_max_position_embeddings", 65536) + rt = rp.get("type", rp.get("rope_type", "yarn")); rf = rp.get("factor", 16.0) + rtheta = cfg.get("rope_theta", 10000.); romax = rp.get("original_max_position_embeddings", 65536) rbfast, rbslow = rp.get("beta_fast", 32), rp.get("beta_slow", 1) - rope_caches = {g: build_rope_cache(8192, rd, f"cuda:{g}", rtheta, rt, rf, romax, rbfast, rbslow) - for g in range(NUM_GPUS)} + rope_caches = {g: build_rope_cache(8192, rd, f"cuda:{g}", rtheta, rt, rf, romax, rbfast, rbslow) for g in range(NUM_GPUS)} # KV caches, compressors, indexers kv_caches, compressors, indexers = {}, {}, {} - n_ih = cfg.get("index_n_heads", 64) - ihd = cfg.get("index_head_dim", 128) - itk = cfg.get("index_topk", 1024) + n_ih = cfg.get("index_n_heads", 64); ihd = cfg.get("index_head_dim", 128); itk = cfg.get("index_topk", 1024) for li in range(n_layers): - dev = f"cuda:{li % NUM_GPUS}" - ratio = cr[li] if li < len(cr) else 128 + dev = f"cuda:{li % NUM_GPUS}"; ratio = cr[li] if li < len(cr) else 128 kv_caches[li] = KVCache(hd, cfg.get("sliding_window", 128), dev) if ratio > 0: compressors[li] = Compressor(ratio, hd, H, dev) if ratio == 4: indexers[li] = Indexer(n_ih, ihd, itk, dev) - # Cache layer weights (EXCLUDE MoE/SE expert weights — handled by production runners) + # Cache layer weights (no MoE/SE) print("Caching layer weights to GPUs (excluding MoE expert weights)...") devs = [f"cuda:{g}" for g in range(NUM_GPUS)] layer_w = _cache_layer_weights_no_experts(all_w, n_layers, devs) del all_w; import gc; gc.collect() - for g in range(NUM_GPUS): - torch.cuda.set_device(g); torch.cuda.empty_cache() + for g in range(NUM_GPUS): torch.cuda.set_device(g); torch.cuda.empty_cache() torch.cuda.set_device(0) print(f" {time.time()-t0:.1f}s") @@ -829,9 +633,7 @@ def main(): if pi % 10 == 0: print(f" Token {pi}/{len(generated)}: {time.time()-t1:.2f}s", flush=True) print(f" Prefill done ({time.time()-t0:.1f}s)") - if _args.prefill_only: - print("Prefill-only mode, stopping.") - return + if _args.prefill_only: print("Prefill-only mode, stopping."); return # Decode print(f"\nDecoding (max {MAX_NEW_TOKENS} tokens)...") @@ -856,14 +658,12 @@ def main(): 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 = F.linear(x_out, lm_w) - next_id = torch.argmax(logits, -1).item() - all_tokens.append(next_id) + next_id = torch.argmax(logits, -1).item(); all_tokens.append(next_id) dt = time.time() - t1 has_nan = torch.isnan(logits.float()).any().item() if step % 5 == 0 or has_nan: tv, ti = torch.topk(logits[0], 5) - top5 = ' '.join(f'{tokenizer.decode([t.item()])}({v.item():.1f})' - for t, v in zip(ti[:5], tv[:5])) + top5 = ' '.join(f'{tokenizer.decode([t.item()])}({v.item():.1f})' for t, v in zip(ti[:5], tv[:5])) print(f" Step {step}: {next_id} '{tokenizer.decode([next_id])}' ({dt:.2f}s) " f"logits=[{logits.float().min().item():.1f},{logits.float().max().item():.1f}] " f"nan={has_nan} |X|={X.abs().max().item():.1f} top5: {top5}", flush=True) @@ -877,6 +677,5 @@ def main(): print(f"Total: {time.time()-t0:.1f}s") print(f"{'='*70}") - if __name__ == "__main__": - main() \ No newline at end of file + main()