#!/usr/bin/env python3 """Single-shot DSV4-Pro inference — Full 61-layer pipeline, 8-GPU. This is a reference implementation that exercises the production kernel stack end-to-end. It should be usable as ground truth when integrating into vLLM or SGLang. Architecture (paper §2): X_l → mHC.pre_block → RMSNorm → Attention → F_attn → mHC.post_block → X_mid X_mid → mHC.pre_block → RMSNorm → FFN(MoE) → F_ffn → mHC.post_block → X_{l+1} Components exercised: - mHC (Manifold-Constrained Hyper-Connections) — proper Sinkhorn-Knopp - Low-rank Q projection (q_a → q_b) + KV projection (MQA, 1 KV head) - Partial RoPE (last 64 dims, GPT-J interleaved) - Production FMHA kernel (6-warp multi-tile, C API + ctypes) - Inverse RoPE on attention output (paper §2.3.3) - Grouped output projection (wo_a BMM + wo_b NVFP4) - Routed MoE (384 experts, top-6, hash + dense routing, SwiGLU clamp) - Shared expert (NVFP4 gate/up/down) - RMSNorm (pre-norm before each sub-block) - KV cache across decode steps Attention type simplification for this single-shot test: For short sequences (seq_len ≤ sliding_window=128), ALL attention types (CSA/HCA/SWA) reduce to dense attention over the full KV cache. CSA's compressed branch and indexer are only needed for long sequences where seq_len > sliding_window. HCA is dense over compressed entries, but at short sequence lengths, the compressed sequence is trivially small. So we use dense MQA attention over the full KV for all layers. This is mathematically correct for short sequences and exercises the FMHA kernel properly. Usage (on B200): source /root/dsv4-nvfp4-workspace/venv/bin/activate cd /root/dsv4-nvfp4-workspace/kernel python3 single_shot_inference.py """ import os, sys, time, json, math, argparse import torch from pathlib import Path # ===================================================================== # Configuration # ===================================================================== def parse_args(): p = argparse.ArgumentParser(description='DSV4 Single-Shot Inference') p.add_argument('--no-inverse-rope', action='store_true', help='Skip inverse RoPE on attention output') p.add_argument('--skip-moe', action='store_true', help='Only use shared expert (skip routed)') p.add_argument('--no-thinking', action='store_true', help='Force model to skip thinking (use <|EOT|> instead of thinking tokens)') p.add_argument('--max-tokens', type=int, default=512, help='Max new tokens to generate') p.add_argument('--prompt', type=str, default=None, help='Override prompt') return p.parse_args() _args = parse_args() CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" MAX_NEW_TOKENS = _args.max_tokens SYSTEM_PROMPT = "" # Empty system prompt for testing PROMPT = _args.prompt or "The capital of France is" NUM_GPUS = 8 SKIP_ROUTED_MOE = _args.skip_moe # If True, only use shared expert (debug) INVERSE_ROPE = not _args.no_inverse_rope # If False, skip inverse RoPE on attention output (diagnostic) SKIP_MHC = _args.skip_mhc # If True, bypass mHC and use simple residual connections (diagnostic) MHC_DIAG = True # If True, print per-layer mHC diagnostics (B_l row/col sums, C_l values) # When True: applies inverse RoPE at query position → converts absolute→relative # When False: leaves relative position encoding intact for output projection # DSV4 partial RoPE only affects last 64/512 dims; first 448 are always un-RoPE'd print(f"Config: INVERSE_ROPE={INVERSE_ROPE}, SKIP_ROUTED_MOE={SKIP_ROUTED_MOE}, MAX_NEW_TOKENS={MAX_NEW_TOKENS}") # ===================================================================== # NVFP4 dequantization — matches checkpoint format exactly # ===================================================================== FP4_LUT = torch.tensor([0., 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) # E2M1 magnitudes def dequant_nvfp4_weight(weight, weight_scale, weight_scale_2): """Dequantize NVFP4 weight to BF16. weight: (out_dim, in_dim//2) uint8 — 2 FP4 values per byte weight_scale: (out_dim, in_dim//16) E4M3 — per-16-element block scale weight_scale_2: (out_dim, 1) float32 — per-row global scale """ out_dim = weight.shape[0] in_packed = weight.shape[1] in_features = in_packed * 2 low = (weight & 0x0F).to(torch.int8) high = (weight >> 4).to(torch.int8) low_sign, low_idx = (low >> 3).bool(), (low & 0x07).long() high_sign, high_idx = (high >> 3).bool(), (high & 0x07).long() lut = FP4_LUT.to(device=weight.device, dtype=torch.float32) low_f = lut[low_idx] * torch.where(low_sign, -1.0, 1.0) high_f = lut[high_idx] * torch.where(high_sign, -1.0, 1.0) w_f = torch.stack([low_f, high_f], dim=-1).reshape(out_dim, in_features) scale_f = weight_scale.float() * weight_scale_2.float() scale_expanded = scale_f.repeat_interleave(16, dim=1) return (w_f * scale_expanded).bfloat16() def nvfp4_linear(x, weight, weight_scale, weight_scale_2): """BF16 linear with NVFP4 dequant.""" w = dequant_nvfp4_weight(weight, weight_scale, weight_scale_2) return torch.nn.functional.linear(x, w) # ===================================================================== # RMSNorm — matches dsv4/layers/norm.py # ===================================================================== class RMSNorm: def __init__(self, hidden_size, eps=1e-6, device='cuda:0'): self.eps = eps self.weight = torch.ones(hidden_size, dtype=torch.float32, device=device) def forward(self, x): """x: (T, H) BF16 → (T, H) BF16""" x_f = x.float() rms = x_f.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt() return (x_f * rms * self.weight).to(torch.bfloat16) # ===================================================================== # mHC — proper Sinkhorn-Knopp implementation # ===================================================================== class mHCBlock: """Wrapper around dsv4.layers.mhc.mHCLayer for single-shot inference. Uses the production mHCLayer implementation with proper Sinkhorn-Knopp. """ def __init__(self, hidden_dim=7168, n_hc=4, sinkhorn_iters=20, device='cuda:0'): from dsv4.layers.mhc import mHCLayer self._impl = mHCLayer( hidden_dim=hidden_dim, n_hc=n_hc, t_max_sinkhorn=sinkhorn_iters, device=device, dtype=torch.bfloat16) self.device = device self.n_hc = n_hc self.hidden_dim = hidden_dim def load_from_checkpoint(self, fn, base, scale): """Load from checkpoint tensors. Checkpoint layout (verified against HuggingFace DeepseekV4HyperConnection): fn: (24, 28672) — rows ordered [pre(4), post(4), comb(16)] base: (24,) — ordered [pre(4), post(4), comb(16)] scale: (3,) — [alpha_pre, alpha_post, alpha_comb] The HuggingFace model does: pre_w, post_w, comb_w = F.linear(flat, fn).split([4, 4, 16]) pre_b, post_b, comb_b = base.split([4, 4, 16]) pre_scale, post_scale, comb_scale = scale.unbind(0) """ n = self.n_hc dev = self.device # fn rows: [pre(4), post(4), comb(16)] — matches HuggingFace W_pre = fn[0:n].to(device=dev, dtype=torch.float32).contiguous() # fn[0:4] W_post = fn[n:2*n].to(device=dev, dtype=torch.float32).contiguous() # fn[4:8] W_comb = fn[2*n:].to(device=dev, dtype=torch.float32).contiguous() # fn[8:24] # base: [S_pre(4), S_post(4), S_comb(16)] — same ordering as fn S_pre = base[0:n].reshape(1, n).to(device=dev, dtype=torch.bfloat16).contiguous() # base[0:4] S_post = base[n:2*n].reshape(n, 1).to(device=dev, dtype=torch.bfloat16).contiguous() # base[4:8] S_comb = base[2*n:].reshape(n, n).to(device=dev, dtype=torch.bfloat16).contiguous() # base[8:24] # scale: [alpha_pre, alpha_post, alpha_comb] alpha_pre = scale[0].item() alpha_post = scale[1].item() alpha_comb = scale[2].item() self._impl.load_weights( W_pre=W_pre, W_post=W_post, W_comb=W_comb, S_pre=S_pre, S_post=S_post, S_comb=S_comb, alpha_pre=alpha_pre, alpha_post=alpha_post, alpha_comb=alpha_comb) @staticmethod def init_state(embeddings, n_hc=4): from dsv4.layers.mhc import mHCLayer return mHCLayer.init_state(embeddings, n_hc) def pre_block(self, X_l): return self._impl.pre_block(X_l) def post_block(self, X_l, F_out, ctx): return self._impl.post_block(X_l, F_out, ctx) # ===================================================================== # RoPE — partial, GPT-J interleaved, last rope_dim dims # ===================================================================== def build_rope_cache(max_pos, rope_dim, device, theta=10000.0, rope_type="default", rope_factor=1.0, original_max_pos=4096, beta_fast=32, beta_slow=1): """Build cos/sin caches for partial RoPE. CRITICAL: FP32, not BF16! BF16 quantization destroys cos²+sin²=1 identity needed for inverse RoPE. BF16 cos²+sin² can be 0.996, causing ~3% round-trip error that accumulates across 61 layers. Supports YaRN (Yet another RoPE extensioN) scaling for long context. The DSV4 Pro model uses rope_type='yarn' with factor=16. Returns: (cos_cache, sin_cache) each (max_pos, rope_dim//2) FP32 """ half = rope_dim // 2 # Base frequencies: 1 / theta^(2i/d) freqs = 1.0 / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float32) / rope_dim)) if rope_type == "yarn" and rope_factor > 1.0: # YaRN frequency scaling # Compute wavelength thresholds low_freq_wavelen = original_max_pos / (beta_fast * 2.0) # High-freq cutoff high_freq_wavelen = original_max_pos / (beta_slow * 2.0) # Low-freq cutoff new_freqs = [] for freq in freqs: wavelen = 2 * math.pi / freq if wavelen < low_freq_wavelen: # High frequency: no scaling new_freqs.append(freq) elif wavelen > high_freq_wavelen: # Low frequency: scale by 1/factor new_freqs.append(freq / rope_factor) else: # Medium frequency: smooth interpolation smooth = (original_max_pos / (wavelen * beta_slow) - rope_factor) / ( rope_factor * (beta_fast / beta_slow - 1) ) new_freqs.append((1 - smooth) * freq / rope_factor + smooth * freq) freqs = torch.tensor(new_freqs, dtype=torch.float32) 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_partial(x, positions, cos_cache, sin_cache, head_dim, rope_dim): """Apply partial GPT-J interleaved RoPE to the last rope_dim dims of each head. Computes in FP32 for numerical stability (inverse RoPE requires cos²+sin²=1).""" T, n_h, hd = x.shape nope = hd - rope_dim cos = cos_cache[positions].unsqueeze(1) # (T, 1, half) FP32 sin = sin_cache[positions].unsqueeze(1) x_rope = x[:, :, nope:].float() # FP32 for accurate rotation x_even = x_rope[..., 0::2] x_odd = x_rope[..., 1::2] rot_even = x_even * cos - x_odd * sin rot_odd = x_even * sin + x_odd * cos result = x.clone() rope_out = torch.empty_like(x_rope) rope_out[..., 0::2] = rot_even rope_out[..., 1::2] = rot_odd result[:, :, nope:] = rope_out.to(torch.bfloat16) return result def apply_inverse_rope(o, positions, cos_cache, sin_cache, head_dim, rope_dim): """Apply inverse RoPE (conjugate rotation) to attention output. Computes in FP32 for numerical stability.""" T, n_h, hd = o.shape nope = hd - rope_dim cos = cos_cache[positions].unsqueeze(1) sin = sin_cache[positions].unsqueeze(1) o_rope = o[:, :, nope:].float() o_even = o_rope[..., 0::2] o_odd = o_rope[..., 1::2] inv_even = o_even * cos + o_odd * sin inv_odd = -o_even * sin + o_odd * cos result = o.clone() rope_out = torch.empty_like(o_rope) rope_out[..., 0::2] = inv_even rope_out[..., 1::2] = inv_odd result[:, :, nope:] = rope_out.to(torch.bfloat16) return result class SimpleKVCache: """Per-layer KV cache for decode. Stores BF16 K,V accumulated across steps. MQA: 1 KV head, so cache is (1, seq_len, hd) per layer.""" def __init__(self, head_dim, max_seq=8192, device='cuda:0'): self.hd = head_dim self.max_seq = max_seq self.device = device self.k = torch.zeros(1, max_seq, head_dim, dtype=torch.bfloat16, device=device) self.v = torch.zeros(1, max_seq, head_dim, dtype=torch.bfloat16, device=device) self.len = 0 def append(self, k_new, v_new): """Append K,V. k_new: (1, T, hd), v_new: (1, T, hd).""" T = k_new.shape[1] self.k[0, self.len:self.len + T] = k_new[0] self.v[0, self.len:self.len + T] = v_new[0] self.len += T def get(self): """Get K,V up to current length. Returns (1, seq_len, hd) each.""" return self.k[:, :self.len], self.v[:, :self.len] # ===================================================================== # Weight loading — streams safetensors shards, distributes to 8 GPUs # ===================================================================== def load_weights_to_cpu(checkpoint_dir): """Load all weights from checkpoint to CPU memory. Weights stay on CPU; we move per-layer to GPU on demand during inference. This avoids OOM from 285K GPU allocations and allows streaming. Returns: all_weights: dict[key] → tensor on CPU """ from safetensors.torch import load_file cdir = Path(checkpoint_dir) index_path = cdir / "model.safetensors.index.json" weight_map = {} if index_path.exists(): with open(index_path) as f: weight_map = json.load(f).get("weight_map", {}) shard_names = set(weight_map.values()) if weight_map else { f"model-{i:05d}-of-00095.safetensors" for i in range(1, 96) } print(f"Loading {len(shard_names)} shards to CPU...") all_weights = {} loaded = 0 for shard_name in sorted(shard_names): if not (cdir / shard_name).exists(): continue data = load_file(str(cdir / shard_name)) all_weights.update(data) loaded += 1 if loaded % 20 == 0: print(f" {loaded}/{len(shard_names)} shards, {len(all_weights)} tensors") print(f" Done: {len(all_weights)} tensors on CPU") return all_weights def get_layer_weights(all_weights, li, device): """Get weights for layer li, moved to target device. Returns dict of key→tensor on device. Filters by model.layers.{li} prefix. """ prefix = f"model.layers.{li}." w = {} for key in all_weights: if key.startswith(prefix): w[key] = all_weights[key].to(device=device, non_blocking=True) return w def cache_all_layer_weights(all_weights, n_layers, devices): """Pre-load ALL layer weights to their target GPUs. This avoids the per-token CPU→GPU transfer bottleneck. Each layer's weights stay on its target GPU for the entire inference run. """ print(f" Caching layer weights to GPUs...") cached = {} for li in range(n_layers): gpu = li % len(devices) dev = devices[gpu] cached[li] = get_layer_weights(all_weights, li, dev) if (li + 1) % 10 == 0: print(f" {li+1}/{n_layers} layers cached") print(f" All {n_layers} layers cached to GPUs") return cached # ===================================================================== # Single layer forward # ===================================================================== def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, attn_mhc, ffn_mhc, attn_norm, ffn_norm, kv_cache, token_id, positions): """Forward one layer with mHC + Attention + FFN. Architecture (paper §2): X_l → mHC.pre_block(attn) → RMSNorm → Attention → F_attn → mHC.post_block → X_mid X_mid → mHC.pre_block(ffn) → RMSNorm → MoE → F_ffn → mHC.post_block → X_{l+1} X_l: (T, n_hc, H) BF16 — mHC residual state Returns: X_next (T, n_hc, H) BF16 """ device = X_l.device H = cfg["hidden_size"] n_h = cfg["num_attention_heads"] hd = cfg["head_dim"] rd = cfg.get("qk_rope_head_dim", cfg.get("rope_dim", 64)) o_rank = cfg.get("output_group_dim", 1024) o_groups = cfg.get("num_output_groups", 16) n_hc = 4 pre = f"model.layers.{li}.self_attn" T = X_l.shape[0] heads_per_group = n_h // o_groups group_input_dim = heads_per_group * hd # ================================================================== # ATTENTION SUB-BLOCK # ================================================================== if SKIP_MHC: # Simple residual: skip mHC, use direct input x_in = X_l[:, 0, :] # Just take stream 0 attn_ctx = None else: # -- mHC pre_block (attention) -- x_in, attn_ctx = attn_mhc.pre_block(X_l) # x_in: (T, H) if MHC_DIAG and attn_ctx is not None: # mHC diagnostics B_l, C_l = attn_ctx.B_l, attn_ctx.C_l print(f" L{li} pre_attn: |X_l|={X_l.abs().max().item():.2f} |x_in|={x_in.abs().max().item():.2f}", flush=True) # -- RMSNorm (pre-norm before attention) -- x_normed = attn_norm.forward(x_in) # (T, H) BF16 # -- Q projection: q_a (low-rank down) → q_a_norm → q_b (low-rank up) -- c_Q = nvfp4_linear(x_normed, w[f"{pre}.q_a_proj.weight"], w[f"{pre}.q_a_proj.weight_scale"], w[f"{pre}.q_a_proj.weight_scale_2"]) # (T, dc) # Q norm (RMSNorm after q_a, before q_b) q_norm_w = w.get(f"{pre}.q_a_norm.weight") if q_norm_w is not None: c_Q_f = c_Q.float() c_Q_rms = c_Q_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() c_Q = (c_Q_f * c_Q_rms * q_norm_w.float()).bfloat16() q = nvfp4_linear(c_Q, w[f"{pre}.q_b_proj.weight"], w[f"{pre}.q_b_proj.weight_scale"], w[f"{pre}.q_b_proj.weight_scale_2"]) # (T, n_h * hd) # q_b_norm — unweighted RMSNorm after q_b_proj (paper §2.3.1) # This is critical: normalizes Q before attention, preventing score collapse. # No learnable parameters — just q / RMS(q). q_f = q.float() q_rms = q_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() q = (q_f * q_rms).bfloat16() # -- KV projection (MQA: 1 KV head) + KV norm -- kv = nvfp4_linear(x_normed, w[f"{pre}.kv_proj.weight"], w[f"{pre}.kv_proj.weight_scale"], w[f"{pre}.kv_proj.weight_scale_2"]) # (T, hd) # KV norm (RMSNorm after kv_proj) kv_norm_w = w.get(f"{pre}.kv_norm.weight") if kv_norm_w is not None: kv_f = kv.float() kv_rms = kv_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() kv = (kv_f * kv_rms * kv_norm_w.float()).bfloat16() # -- Reshape for attention -- q_heads = q.reshape(T, n_h, hd) # (T, n_h, hd) kv_new = kv.reshape(T, 1, hd) # (T, 1, hd) — 1 KV head # Diagnostic: Q/KV norms if MHC_DIAG and li < 3: print(f" L{li} Q: |q|={q_heads.abs().max().item():.2f} mean={q_heads.float().abs().mean().item():.4f}") print(f" L{li} KV: |kv|={kv_new.abs().max().item():.2f} mean={kv_new.float().abs().mean().item():.4f}") # -- Apply RoPE to Q (at current positions) -- positions_dev = positions.to(device) q_heads = apply_rope_partial(q_heads, positions_dev, rope_cos, rope_sin, hd, rd) # -- Apply RoPE to KV (at current positions) BEFORE caching -- # DSV4 convention: RoPE applied to KV before writing to cache. # K = V in DSV4 MQA (same projection, same RoPE'd tensor). kv_new = apply_rope_partial(kv_new, positions_dev, rope_cos, rope_sin, hd, rd) # -- KV cache: append RoPE'd KV (K=V) -- k_new = kv_new # (T, 1, hd) — RoPE'd v_new = kv_new # K = V in DSV4 MQA kv_cache.append(k_new.permute(1, 0, 2), v_new.permute(1, 0, 2)) # (1, T, hd) # -- Get full KV from cache (already RoPE'd) -- k_full, v_full = kv_cache.get() # (1, seq_len, hd) each — RoPE'd, K=V seq_len = k_full.shape[1] # -- Attention: SDPA for short seqs (avoids FMHA padding bug), FMHA for long -- q_input = q_heads.permute(1, 0, 2) # (n_h, T, hd) scale = 1.0 / math.sqrt(hd) # FMHA pads N to next multiple of 128. For N<<128, padded zero-K entries # contribute exp(0)=1 to softmax, diluting real attention weights by ~128/N. # Use SDPA for short sequences where padding dominates. if seq_len < 120: k_expanded = k_full.expand(n_h, -1, -1).contiguous() v_expanded = v_full.expand(n_h, -1, -1).contiguous() # Attention: compute raw scores, add sinks as logit bias, softmax, multiply by V # (paper D5c, matching HuggingFace reference implementation) # Sinks are added as a logit column, softmaxed together, then DROPPED # before V multiplication — NOT as a dummy KV entry. sink_key = f"{pre}.sinks" scores_raw = torch.matmul(q_input, k_expanded.transpose(-1, -2)) * scale # (n_h, T, seq_len) if sink_key in w and seq_len > 0: sinks = w[sink_key].to(device=device) # (n_h,) BF16 # sinks: (n_h,) → reshape to (n_h, 1, 1) for broadcasting with (n_h, T, seq_len) sink_logits = sinks.float().reshape(n_h, 1, 1).expand(-1, T, 1) combined_logits = torch.cat([scores_raw, sink_logits], dim=-1) # (n_h, T, seq_len+1) # Stable softmax combined_logits = combined_logits - combined_logits.max(dim=-1, keepdim=True).values probs = torch.softmax(combined_logits.float(), dim=-1).to(torch.bfloat16) attn_weights = probs[..., :-1] # Drop sink column (n_h, T, seq_len) else: attn_weights = torch.softmax(scores_raw.float(), dim=-1).to(torch.bfloat16) attn_out = torch.matmul(attn_weights, v_expanded) # (n_h, T, hd) attn_out = attn_out.permute(1, 0, 2) # (T, n_h, hd) # Diagnostic: check attention entropy (how spread out the attention is) if MHC_DIAG and li < 3: with torch.no_grad(): scores = torch.matmul(q_input, k_expanded.transpose(-1, -2)) * scale # (n_h, T, seq_len) weights = torch.softmax(scores.float(), dim=-1) # (n_h, 1, seq_len) # For head 0: what positions get the most weight? w0 = weights[0, 0] # (seq_len,) top3_pos = torch.topk(w0, min(3, seq_len)) entropy = -(w0 * (w0 + 1e-10).log()).sum().item() print(f" L{li} attn: seq_len={seq_len} entropy={entropy:.2f} top3_pos={top3_pos.indices.tolist()} top3_w={top3_pos.values.tolist()}") else: # Use FMHA kernel for longer sequences (padding effect is negligible) from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw q_4d = q_input.unsqueeze(0).contiguous() k_4d = k_full.unsqueeze(0).contiguous() v_4d = v_full.unsqueeze(0).transpose(-1, -2).contiguous() o_4d, lse = fmha_multitile_decode_raw(q_4d, k_4d, v_4d, scale) attn_out = o_4d.squeeze(0).permute(1, 0, 2) # Sink correction sink_key = f"{pre}.sinks" if sink_key in w and seq_len > 0: sinks = w[sink_key].to(device=device) lse_2d = lse.squeeze(0).t() sink_exp = torch.exp(sinks.float()) attn_exp = torch.exp(lse_2d.float()) correction = attn_exp / (attn_exp + sink_exp.unsqueeze(0) + 1e-10) attn_out = (attn_out.float() * correction.unsqueeze(-1)).bfloat16() attn_out = attn_out.bfloat16() # -- Inverse RoPE on attention output (paper §2.3.3) -- # DSV4 uses K=V in MQA; both get RoPE'd. Inverse RoPE on the output # at query position q converts: R(q)⁻¹ Σ softmax(R(q)Q·R(p)K) R(p)V # For single KV entry at p: R(p-q)V (relative position encoding) # This only affects the last 64 dims (partial RoPE); first 448 unchanged. # The relative encoding in those 64 dims may be INTENTIONAL — the # output projection can use it for position-dependent computation. # Test both modes via INVERSE_ROPE flag. if INVERSE_ROPE: attn_out = apply_inverse_rope(attn_out, positions_dev, rope_cos, rope_sin, hd, rd) # -- Output projection: wo_a (grouped BMM) + wo_b (NVFP4) -- # wo_a: grouped linear, (n_h, hd) → (n_groups, o_rank) via BMM attn_flat = attn_out.reshape(T, n_h * hd) # (T, n_h * hd) attn_grouped = attn_flat.reshape(T, o_groups, heads_per_group * hd) # (T, groups, group_dim) oa_w = w[f"{pre}.o_a_proj.weight"].bfloat16() # (n_groups * o_rank, group_input_dim) BF16 oa_3d = oa_w.reshape(o_groups, o_rank, group_input_dim) # (groups, o_rank, group_dim) attn_for_bmm = attn_grouped.permute(1, 0, 2) # (groups, T, group_dim) grouped_out = torch.bmm(attn_for_bmm, oa_3d.transpose(1, 2)) # (groups, T, o_rank) grouped_flat = grouped_out.permute(1, 0, 2).reshape(T, o_groups * o_rank) # (T, groups*o_rank) F_attn = nvfp4_linear(grouped_flat, w[f"{pre}.o_b_proj.weight"], w[f"{pre}.o_b_proj.weight_scale"], w[f"{pre}.o_b_proj.weight_scale_2"]) # (T, H) if SKIP_MHC: X_mid = X_l[:, 0, :].unsqueeze(1).expand(-1, 4, -1) + F_attn.unsqueeze(1) * 0.1 else: # -- mHC post_block (attention) -- X_mid = attn_mhc.post_block(X_l, F_attn, attn_ctx) # (T, n_hc, H) # Diagnostic: check mHC is stabilizing the residual if MHC_DIAG and attn_ctx is not None: # mHC diagnostics B_l, C_l = attn_ctx.B_l, attn_ctx.C_l print(f" L{li} attn: |X_l|={X_l.abs().max().item():.2f} |F_attn|={F_attn.abs().max().item():.2f} |B|={B_l.abs().max().item():.4f} |C|={C_l.abs().max().item():.4f} |X_mid|={X_mid.abs().max().item():.2f}") # Check B_l is doubly stochastic (rows sum to 1.0) B_row_sums = B_l.sum(dim=-1) # (T, n_hc) B_col_sums = B_l.sum(dim=-2) # (T, n_hc) print(f" B row_sums={B_row_sums[0].tolist()} col_sums={B_col_sums[0].tolist()}") print(f" C_l={C_l[0].tolist()}") # ================================================================== # FFN SUB-BLOCK # ================================================================== if SKIP_MHC: x_ffn = X_mid[:, 0, :] # Just take stream 0 ffn_ctx = None else: # -- mHC pre_block (FFN) -- x_ffn, ffn_ctx = ffn_mhc.pre_block(X_mid) # (T, H) # -- RMSNorm (pre-norm before FFN) -- x_ffn_normed = ffn_norm.forward(x_ffn) # (T, H) BF16 # -- MoE + shared expert -- F_ffn = moe_forward(x_ffn_normed, w, li, cfg, token_id, device) if SKIP_MHC: X_next = X_mid + F_ffn.unsqueeze(1) * 0.1 else: # -- mHC post_block (FFN) -- X_next = ffn_mhc.post_block(X_mid, F_ffn, ffn_ctx) # (T, n_hc, H) if MHC_DIAG and ffn_ctx is not None: # ffn mHC diagnostics B_l_ffn, C_l_ffn = ffn_ctx.B_l, ffn_ctx.C_l print(f" L{li} ffn: |X_mid|={X_mid.abs().max().item():.2f} |F_ffn|={F_ffn.abs().max().item():.2f} |B|={B_l_ffn.abs().max().item():.4f} |C|={C_l_ffn.abs().max().item():.4f} |X_next|={X_next.abs().max().item():.2f}", flush=True) return X_next # ===================================================================== # MoE forward — hash + dense routing, SwiGLU with clamping # ===================================================================== def moe_forward(x, w, li, cfg, token_id, device): """Run routed MoE + shared expert. x: (T, H) BF16 — post-RMSNorm FFN input Returns: (T, H) BF16 """ H = cfg["hidden_size"] n_experts = cfg["n_routed_experts"] top_k = cfg.get("num_experts_per_tok", 6) routed_scaling = cfg.get("routed_scaling_factor", 2.5) swiglu_limit = cfg.get("swiglu_limit", 10.0) mlp_inter = cfg["moe_intermediate_size"] # ---- Routing ---- # Layers 0-2: hash routing (tid2eid lookup) # Layers 3+: noaux_tc (sqrt(softplus) scoring + e_score_correction_bias for selection only) # Config: topk_method='noaux_tc', scoring_func='sqrtsoftplus' expert_ids = None expert_weights = None tid2eid_key = f"model.layers.{li}.mlp.gate.tid2eid" e_bias_key = f"model.layers.{li}.mlp.gate.e_score_correction_bias" is_hash = tid2eid_key in w and e_bias_key not in w if is_hash: # Hash routing: deterministic per-token lookup, uniform weights tid2eid = w[tid2eid_key] tid = token_id.item() if token_id.numel() == 1 else token_id[0].item() expert_ids = tid2eid[tid] # (top_k,) int64 expert_weights = torch.ones(top_k, dtype=torch.float32, device=x.device) / top_k else: # Dense routing: sqrt(softplus(logits)) scoring gate_w = w[f"model.layers.{li}.mlp.gate.weight"] # (H, n_experts) BF16 logits = torch.nn.functional.linear(x, gate_w.bfloat16()) # (T, n_experts) # Scoring: sqrt(softplus(logits)) scores = torch.sqrt(torch.nn.functional.softplus(logits.float()) + 1e-6) # e_score_correction_bias: per-expert bias for SELECTION ONLY selection_logits = scores.clone() if e_bias_key in w: selection_logits = selection_logits + w[e_bias_key].float().unsqueeze(0) _, indices = selection_logits.topk(top_k, dim=-1) # (T, top_k) # Weights from UNBIASED scores (no e_bias) expert_weights = torch.gather(scores, -1, indices) expert_weights = expert_weights / expert_weights.sum(dim=-1, keepdim=True) # For T=1 decode, squeeze if x.shape[0] == 1: expert_ids = indices[0] expert_weights = expert_weights[0] else: raise NotImplementedError("Multi-token MoE routing") # ---- Run selected experts ---- T = x.shape[0] expert_outputs = [] if not SKIP_ROUTED_MOE: for i, eid in enumerate(expert_ids): eid_int = eid.item() epre = f"model.layers.{li}.mlp.experts.{eid_int}" gate = nvfp4_linear(x, w[f"{epre}.gate_proj.weight"], w[f"{epre}.gate_proj.weight_scale"], w[f"{epre}.gate_proj.weight_scale_2"]) up = nvfp4_linear(x, w[f"{epre}.up_proj.weight"], w[f"{epre}.up_proj.weight_scale"], w[f"{epre}.up_proj.weight_scale_2"]) # SwiGLU with clamping (paper §4.2.3) silu_out = torch.nn.functional.silu(gate.float()) if swiglu_limit is not None: silu_out = silu_out.clamp(-swiglu_limit, swiglu_limit) up_clamped = up.float().clamp(-swiglu_limit, swiglu_limit) else: up_clamped = up.float() hidden = (silu_out * up_clamped).bfloat16() down = nvfp4_linear(hidden, w[f"{epre}.down_proj.weight"], w[f"{epre}.down_proj.weight_scale"], w[f"{epre}.down_proj.weight_scale_2"]) expert_outputs.append(down) # Weighted combine + scaling routed_out = torch.zeros_like(x) for i, (out, wt) in enumerate(zip(expert_outputs, expert_weights)): routed_out = routed_out + (out.float() * wt.item()).bfloat16() routed_out = (routed_out.float() * routed_scaling).bfloat16() # ---- Shared expert ---- se_pre = f"model.layers.{li}.mlp.shared_experts" se_gate_key = f"{se_pre}.gate_proj.weight" if se_gate_key in w: gate = nvfp4_linear(x, w[se_gate_key], w[f"{se_pre}.gate_proj.weight_scale"], w[f"{se_pre}.gate_proj.weight_scale_2"]) up = nvfp4_linear(x, w[f"{se_pre}.up_proj.weight"], w[f"{se_pre}.up_proj.weight_scale"], w[f"{se_pre}.up_proj.weight_scale_2"]) silu_out = torch.nn.functional.silu(gate.float()) if swiglu_limit is not None: silu_out = silu_out.clamp(-swiglu_limit, swiglu_limit) up_clamped = up.float().clamp(-swiglu_limit, swiglu_limit) else: up_clamped = up.float() hidden = (silu_out * up_clamped).bfloat16() shared_out = nvfp4_linear(hidden, w[f"{se_pre}.down_proj.weight"], w[f"{se_pre}.down_proj.weight_scale"], w[f"{se_pre}.down_proj.weight_scale_2"]) else: shared_out = torch.zeros_like(x) return routed_out + shared_out # ===================================================================== # Main # ===================================================================== def main(): t_start = time.time() print("=" * 70) print("DSV4 Single-Shot Inference — Full Pipeline (mHC+Attn+MoE)") print(" Proper Sinkhorn mHC, RMSNorm, inverse RoPE, production FMHA") 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"] n_h = cfg["num_attention_heads"] hd = cfg["head_dim"] rd = cfg.get("qk_rope_head_dim", cfg.get("rope_dim", 64)) n_hc = 4 print(f"Model: {n_layers} layers, {n_h} heads, hd={hd}, rope_dim={rd}") print(f"Experts: {cfg['n_routed_experts']}, top-{cfg.get('num_experts_per_tok', 6)}") # ==== Phase 1: Load weights to CPU ==== print(f"\n{'='*70}\nPhase 1: Loading weights to CPU\n{'='*70}") all_weights = load_weights_to_cpu(CHECKPOINT_DIR) t_loaded = time.time() print(f"Weight loading: {t_loaded - t_start:.1f}s") # ==== Build mHC blocks + RMSNorms (small weights, keep on GPU) ==== print("Building mHC blocks and RMSNorms...") attn_mhc_blocks = {} ffn_mhc_blocks = {} attn_norms = {} ffn_norms = {} for li in range(n_layers): gpu = li % NUM_GPUS dev = f"cuda:{gpu}" # mHC blocks (small weights: fn (24, 28672) FP32 ≈ 2.6MB each) for prefix, blocks in [(f"model.layers.{li}.attn_hc", attn_mhc_blocks), (f"model.layers.{li}.ffn_hc", ffn_mhc_blocks)]: fn_key = f"{prefix}.fn" base_key = f"{prefix}.base" scale_key = f"{prefix}.scale" if fn_key in all_weights and base_key in all_weights and scale_key in all_weights: mhc = mHCBlock(hidden_dim=H, n_hc=n_hc, device=dev) mhc.load_from_checkpoint( all_weights[fn_key], all_weights[base_key], all_weights[scale_key]) blocks[li] = mhc else: print(f" WARNING: no mHC weights for {prefix}, using identity fallback") # Fallback: near-identity mHC (small alphas, identity comb) mhc = mHCBlock(hidden_dim=H, n_hc=n_hc, device=dev) n = n_hc K = n * H mhc._impl.W_pre = torch.zeros(n, K, dtype=torch.float32, device=dev) mhc._impl.W_post = torch.zeros(n, K, dtype=torch.float32, device=dev) mhc._impl.W_comb = torch.zeros(n*n, K, dtype=torch.float32, device=dev) mhc._impl.S_pre = torch.zeros(1, n, dtype=torch.bfloat16, device=dev) mhc._impl.S_post = torch.ones(n, 1, dtype=torch.bfloat16, device=dev) * 0.5 mhc._impl.S_comb = torch.eye(n, dtype=torch.bfloat16, device=dev) mhc._impl.alpha_pre = torch.tensor(0.01, dtype=torch.float32, device=dev) mhc._impl.alpha_post = torch.tensor(0.01, dtype=torch.float32, device=dev) mhc._impl.alpha_comb = torch.tensor(0.01, dtype=torch.float32, device=dev) blocks[li] = mhc # RMSNorms attn_norm = RMSNorm(H, eps=cfg.get('rms_norm_eps', 1e-6), device=dev) an_key = f"model.layers.{li}.input_layernorm.weight" if an_key in all_weights: attn_norm.weight = all_weights[an_key].to(device=dev, dtype=torch.float32) attn_norms[li] = attn_norm ffn_norm = RMSNorm(H, eps=cfg.get('rms_norm_eps', 1e-6), device=dev) fn_key = f"model.layers.{li}.post_attention_layernorm.weight" if fn_key in all_weights: ffn_norm.weight = all_weights[fn_key].to(device=dev, dtype=torch.float32) ffn_norms[li] = ffn_norm print(f" attn mHC: {len(attn_mhc_blocks)}, ffn mHC: {len(ffn_mhc_blocks)}") # ==== Global weights (small, keep on gpu0) ==== torch.cuda.set_device(0) embed_w = all_weights.get("model.embed_tokens.weight") embed = torch.nn.Embedding.from_pretrained(embed_w.bfloat16().to('cuda:0')) lm_w = all_weights.get("lm_head.weight", embed_w).bfloat16().to('cuda:0') final_norm_w = all_weights.get("model.norm.weight") if final_norm_w is not None: final_norm_w = final_norm_w.to('cuda:0') # Build RoPE caches with YaRN scaling from model config rope_params = cfg.get("rope_parameters", {}) rope_type = rope_params.get("rope_type", "default") rope_factor = rope_params.get("factor", 1.0) rope_theta = rope_params.get("rope_theta", cfg.get("rope_theta", 10000.0)) original_max_pos = rope_params.get("original_max_position_embeddings", 4096) beta_fast = rope_params.get("beta_fast", 32) beta_slow = rope_params.get("beta_slow", 1) print(f"RoPE: type={rope_type} factor={rope_factor} theta={rope_theta} " f"orig_max_pos={original_max_pos} beta_fast={beta_fast} beta_slow={beta_slow}", flush=True) rope_caches = {g: build_rope_cache( 8192, rd, f"cuda:{g}", theta=rope_theta, rope_type=rope_type, rope_factor=rope_factor, original_max_pos=original_max_pos, beta_fast=beta_fast, beta_slow=beta_slow ) for g in range(NUM_GPUS)} # ==== KV caches (one per layer on its GPU) ==== kv_caches = {} for li in range(n_layers): kv_caches[li] = SimpleKVCache(head_dim=hd, max_seq=8192, device=f"cuda:{li % NUM_GPUS}") # ==== Cache ALL layer weights to GPUs (avoids per-token CPU→GPU transfer) ==== print(f"\n Caching layer weights to GPUs (one-time transfer)...", flush=True) devices = [f"cuda:{g}" for g in range(NUM_GPUS)] layer_weights = cache_all_layer_weights(all_weights, n_layers, devices) print(f" Done. Freeing CPU weights...", flush=True) del all_weights import gc; gc.collect() # ==== Phase 2: Compile FMHA ==== print(f"\n{'='*70}\nPhase 2: JIT compiling\n{'='*70}") from dsv4.kernels.attention.production import dsv4_attention torch.cuda.set_device(0) dummy_q = torch.randn(n_h, 1, hd, dtype=torch.bfloat16, device='cuda:0') dummy_k = torch.randn(1, 1, hd, dtype=torch.bfloat16, device='cuda:0') try: _ = dsv4_attention(dummy_q, dummy_k, dummy_k.clone()) print(" FMHA: compiled OK") except Exception as e: print(f" FMHA error: {e}") t_compiled = time.time() print(f"Compile: {t_compiled - t_loaded:.1f}s") # ==== Phase 2.5: Minimal E2E test ==== print(f"\n{'='*70}\nPhase 2.5: Minimal E2E Test (single token 'The')\n{'='*70}") from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR) minimal_e2e_test(layer_weights, cfg, rope_caches, attn_mhc_blocks, ffn_mhc_blocks, attn_norms, ffn_norms, embed, lm_w, final_norm_w, tokenizer) # ==== Phase 2.6: Single-layer trace ==== if True: # Always run the trace print(f"\n{'='*70}\nPhase 2.6: Single-Layer Trace (layer 0, first prefill token)\n{'='*70}", flush=True) li = 0 dev = f"cuda:0" w = layer_weights[li] pre = f"model.layers.{li}.self_attn" T_dim = 1 positions = torch.tensor([0], dtype=torch.long, device=dev) rope_cos, rope_sin = rope_caches[0] # Start from the embedding tid = torch.tensor([tokenizer.encode("The")[-1]], dtype=torch.long, device=dev) emb = embed(tid) # (1, H) X = mHCBlock.init_state(emb, 4) # (1, 4, H) print(f" X after init_state: |X|={X.abs().max().item():.4f} stream0_mean={X[:,0,:].float().abs().mean().item():.6f}", flush=True) # mHC pre_block attn_mhc = attn_mhc_blocks[0] x_in, ctx = attn_mhc.pre_block(X) print(f" x_in (mHC pre_block): |x_in|={x_in.abs().max().item():.4f} mean={x_in.float().abs().mean().item():.6f}", flush=True) B_l = ctx.B_l C_l = ctx.C_l print(f" B_l row_sums={B_l[0].sum(dim=-1).tolist()}", flush=True) print(f" C_l={C_l[0].tolist()}", flush=True) # RMSNorm a_norm = attn_norms[0] x_normed = a_norm.forward(x_in) print(f" x_normed: |x|={x_normed.abs().max().item():.4f} mean={x_normed.float().abs().mean().item():.6f}", flush=True) # Q projection c_Q = nvfp4_linear(x_normed, w[f"{pre}.q_a_proj.weight"], w[f"{pre}.q_a_proj.weight_scale"], w[f"{pre}.q_a_proj.weight_scale_2"]) print(f" c_Q (q_a_proj): |c_Q|={c_Q.abs().max().item():.4f} mean={c_Q.float().abs().mean().item():.6f}", flush=True) # q_a_norm q_norm_w = w.get(f"{pre}.q_a_norm.weight") if q_norm_w is not None: c_Q_f = c_Q.float() c_Q_rms = c_Q_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() c_Q = (c_Q_f * c_Q_rms * q_norm_w.float()).bfloat16() print(f" c_Q after q_a_norm: |c_Q|={c_Q.abs().max().item():.4f}", flush=True) q = nvfp4_linear(c_Q, w[f"{pre}.q_b_proj.weight"], w[f"{pre}.q_b_proj.weight_scale"], w[f"{pre}.q_b_proj.weight_scale_2"]) q_heads = q.reshape(T_dim, n_h, hd) print(f" q_heads: |q|={q_heads.abs().max().item():.4f} mean={q_heads.float().abs().mean().item():.6f}", flush=True) # KV projection kv = nvfp4_linear(x_normed, w[f"{pre}.kv_proj.weight"], w[f"{pre}.kv_proj.weight_scale"], w[f"{pre}.kv_proj.weight_scale_2"]) print(f" kv (kv_proj): |kv|={kv.abs().max().item():.4f} mean={kv.float().abs().mean().item():.6f}", flush=True) # kv_norm kv_norm_w = w.get(f"{pre}.kv_norm.weight") if kv_norm_w is not None: kv_f = kv.float() kv_rms = kv_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() kv = (kv_f * kv_rms * kv_norm_w.float()).bfloat16() print(f" kv after kv_norm: |kv|={kv.abs().max().item():.4f}", flush=True) kv_new = kv.reshape(T_dim, 1, hd) # (1, 1, hd) print(f" kv_new shape: {kv_new.shape}", flush=True) # Apply RoPE q_heads = apply_rope_partial(q_heads, positions, rope_cos, rope_sin, hd, rd) kv_new = apply_rope_partial(kv_new, positions, rope_cos, rope_sin, hd, rd) print(f" After RoPE: |q|={q_heads.abs().max().item():.4f} |kv|={kv_new.abs().max().item():.4f}", flush=True) # Self-attention (single token, trivially weight=1.0) q_input = q_heads.permute(1, 0, 2) # (n_h, 1, hd) k_input = kv_new.permute(1, 0, 2) # (1, 1, hd) -> expand k_expanded = k_input.expand(n_h, -1, -1).contiguous() v_expanded = k_expanded.clone() # K=V in DSV4 MQA attn_out = torch.nn.functional.scaled_dot_product_attention( q_input, k_expanded, v_expanded, scale=1.0/math.sqrt(hd)) attn_out = attn_out.permute(1, 0, 2) # (1, n_h, hd) print(f" attn_out: |o|={attn_out.abs().max().item():.4f} mean={attn_out.float().abs().mean().item():.6f}", flush=True) # Inverse RoPE if INVERSE_ROPE: attn_out = apply_inverse_rope(attn_out, positions, rope_cos, rope_sin, hd, rd) print(f" After inverse RoPE: |o|={attn_out.abs().max().item():.4f}", flush=True) # Output projection o_groups = cfg.get("num_output_groups", 16) o_rank = cfg.get("output_group_dim", 1024) heads_per_group = n_h // o_groups group_input_dim = heads_per_group * hd attn_flat = attn_out.reshape(T_dim, n_h * hd) attn_grouped = attn_flat.reshape(T_dim, o_groups, heads_per_group * hd) oa_w = w[f"{pre}.o_a_proj.weight"].bfloat16() oa_3d = oa_w.reshape(o_groups, o_rank, group_input_dim) attn_for_bmm = attn_grouped.permute(1, 0, 2) grouped_out = torch.bmm(attn_for_bmm, oa_3d.transpose(1, 2)) grouped_flat = grouped_out.permute(1, 0, 2).reshape(T_dim, o_groups * o_rank) print(f" grouped_out (wo_a): |o|={grouped_flat.abs().max().item():.4f} mean={grouped_flat.float().abs().mean().item():.6f}", flush=True) F_attn = nvfp4_linear(grouped_flat, w[f"{pre}.o_b_proj.weight"], w[f"{pre}.o_b_proj.weight_scale"], w[f"{pre}.o_b_proj.weight_scale_2"]) print(f" F_attn (wo_b): |F|={F_attn.abs().max().item():.4f} mean={F_attn.float().abs().mean().item():.6f}", flush=True) # mHC post_block X_mid = attn_mhc.post_block(X, F_attn, ctx) print(f" X_mid: |X|={X_mid.abs().max().item():.4f} stream0_mean={X_mid[:,0,:].float().abs().mean().item():.6f}", flush=True) print(f" Layer 0 trace complete.", flush=True) # ==== Phase 3: Inference ==== print(f"\n{'='*70}\nPhase 3: Inference\n{'='*70}") # DeepSeek V4 chat format: <｜begin▁of▁sentence｜><｜User｜>prompt<｜Assistant｜> # For reasoning models: <｜User｜>prompt<｜Assistant｜>fithinking...flanswer # Special token IDs: <｜User｜>=128803, <｜Assistant｜>=128804, <|EOT|>=128805 # Thinking tokens: fi=128821, fl=128822 USER_TOKEN = 128803 ASSISTANT_TOKEN = 128804 EOT_TOKEN = 128805 THINK_START = 128821 # fi THINK_END = 128822 # fl # Build input with proper DeepSeek chat format bos_id = tokenizer.bos_token_id or 0 # <｜User｜> System prompt \n\n User prompt <｜Assistant｜> input_ids_list = [bos_id, USER_TOKEN] input_ids_list += tokenizer.encode(SYSTEM_PROMPT, add_special_tokens=False) input_ids_list += tokenizer.encode('\n\n' + PROMPT, add_special_tokens=False) input_ids_list.append(ASSISTANT_TOKEN) input_ids = torch.tensor([input_ids_list], dtype=torch.long).cuda() print(f"DeepSeek chat format. Input: {input_ids.shape[1]} tokens", flush=True) print(f"Decoded start: '{tokenizer.decode(input_ids[0][:20])}...'", flush=True) print(f"Decoded end: '...{tokenizer.decode(input_ids[0][-5:])}'", flush=True) generated = input_ids[0].tolist() # ==== Prefill: process prompt tokens to fill KV cache ==== print(f"Prefilling {len(generated)} prompt tokens...", flush=True) for prefill_idx, tid_val in enumerate(generated): t0 = time.time() tid = torch.tensor([tid_val], dtype=torch.long, device='cuda:0') positions = torch.tensor([prefill_idx], dtype=torch.long, device='cuda:0') emb = embed(tid) # (1, H) on gpu0 X = mHCBlock.init_state(emb, n_hc) # (1, n_hc, H) for li in range(n_layers): gpu = li % NUM_GPUS dev = f"cuda:{gpu}" if X.device != torch.device(dev): X = X.to(dev) torch.cuda.set_device(gpu) w = layer_weights[li] attn_mhc = attn_mhc_blocks.get(li) ffn_mhc = ffn_mhc_blocks.get(li) a_norm = attn_norms[li] f_norm = ffn_norms[li] rc, rs = rope_caches[gpu] X = forward_layer(X, w, li, cfg, rc, rs, attn_mhc, ffn_mhc, a_norm, f_norm, kv_caches[li], tid, positions) X = X.to('cuda:0') torch.cuda.set_device(0) if prefill_idx % 10 == 0: print(f" Token {prefill_idx}/{len(generated)}: {time.time()-t0:.2f}s", flush=True) print(f" Prefill done ({len(generated)} tokens, {time.time()-t_compiled:.1f}s)") # ==== Decode: generate new tokens ==== print(f"\nDecoding (max {MAX_NEW_TOKENS} new tokens)...") all_tokens = generated.copy() for step in range(MAX_NEW_TOKENS): t0 = time.time() tid = torch.tensor([all_tokens[-1]], dtype=torch.long, device='cuda:0') decode_pos = len(all_tokens) - 1 positions = torch.tensor([decode_pos], dtype=torch.long, device='cuda:0') emb = embed(tid) # (1, H) on gpu0 X = mHCBlock.init_state(emb, n_hc) # (1, n_hc, H) for li in range(n_layers): gpu = li % NUM_GPUS dev = f"cuda:{gpu}" if X.device != torch.device(dev): X = X.to(dev) torch.cuda.set_device(gpu) w = layer_weights[li] attn_mhc = attn_mhc_blocks.get(li) ffn_mhc = ffn_mhc_blocks.get(li) a_norm = attn_norms[li] f_norm = ffn_norms[li] rc, rs = rope_caches[gpu] X = forward_layer(X, w, li, cfg, rc, rs, attn_mhc, ffn_mhc, a_norm, f_norm, kv_caches[li], tid, positions) X = X.to('cuda:0') torch.cuda.set_device(0) # Read out stream 0 → RMSNorm → lm_head x_out = X[:, 0, :] # (1, H) if final_norm_w is not None: xf = x_out.float() rms = xf.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() x_out = (xf * rms * final_norm_w.float()).bfloat16() logits = torch.nn.functional.linear(x_out, lm_w) # Top-5 predictions for debugging # Top-20 predictions for debugging (includes thinking tokens) top20_vals, top20_ids = torch.topk(logits[0], 20) top5_str = ' '.join([f'{tokenizer.decode([tid.item()])}({val.item():.1f})' for tid, val in zip(top20_ids[:5], top20_vals[:5])]) # Check if thinking tokens are in top-20 thinking_in_top20 = any(tid.item() in [128821, 128822] for tid in top20_ids) top20_ids_set = set(top20_ids.tolist()) next_id = torch.argmax(logits, dim=-1).item() generated.append(next_id) all_tokens.append(next_id) tok_str = tokenizer.decode([next_id]) dt = time.time() - t0 has_nan = torch.isnan(logits.float()).any().item() has_inf = torch.isinf(logits.float()).any().item() lmin, lmax = logits.float().min().item(), logits.float().max().item() x_max = X.abs().max().item() print(f" Step {step}: {next_id} '{tok_str}' ({dt:.2f}s) " f"logits=[{lmin:.1f},{lmax:.1f}] nan={has_nan} inf={has_inf} " f"|X|={x_max:.3f} top5: {top5_str}", flush=True) if thinking_in_top20: for tid_t, val_t in zip(top20_ids, top20_vals): if tid_t.item() in [128821, 128822]: print(f" THINK TOKEN: {tid_t.item()} logit={val_t.item():.3f}", flush=True) if step % 5 == 0: print(f" Top-20: {[(tokenizer.decode([t.item()]), f'{v.item():.2f}') for t, v in zip(top20_ids, top20_vals)]}", flush=True) if has_nan or has_inf: print(" Numerical issue — stopping") break if next_id == tokenizer.eos_token_id: break out = tokenizer.decode(generated, skip_special_tokens=True) total = time.time() - t_start print(f"\n{'='*70}") print(f"Input: '{PROMPT}'") print(f"Output: '{out}'") print(f"Total: {total:.1f}s") print(f"{'='*70}") # ===================================================================== # Minimal end-to-end test — single token "The" through the model # ===================================================================== def minimal_e2e_test(layer_weights, cfg, rope_caches, attn_mhc_blocks, ffn_mhc_blocks, attn_norms, ffn_norms, embed, lm_w, final_norm_w, tokenizer): """Process a single token 'The' through the model and check output logits. This is a focused diagnostic: if the model can't even produce reasonable logits for a single token, something is fundamentally wrong in the pipeline. We check: 1. No NaN/Inf in any layer output 2. Residual stream magnitude stays bounded 3. Top-5 logits are sensible (not all Chinese tokens for English) 4. Logit spread (max - min) is > 1.0 (not uniform) """ n_layers = cfg["num_hidden_layers"] H = cfg["hidden_size"] n_h = cfg["num_attention_heads"] hd = cfg["head_dim"] rd = cfg.get("qk_rope_head_dim", cfg.get("rope_dim", 64)) n_hc = 4 # Tokenize just "The" tid = torch.tensor(tokenizer.encode("The"), dtype=torch.long, device='cuda:0') if tid.numel() > 1: # If tokenizer adds BOS, take last token print(f" Note: 'The' tokenized to {tid.numel()} tokens, using last one") tid = tid[-1:] print(f" Token ID: {tid.item()} = '{tokenizer.decode(tid.tolist())}'") # Setup positions = torch.tensor([0], dtype=torch.long, device='cuda:0') emb = embed(tid) # (1, H) X = mHCBlock.init_state(emb, n_hc) # (1, n_hc, H) # Track per-layer diagnostics layer_diags = [] for li in range(n_layers): gpu = li % NUM_GPUS dev = f"cuda:{gpu}" if X.device != torch.device(dev): X = X.to(dev) torch.cuda.set_device(gpu) w = layer_weights[li] attn_mhc = attn_mhc_blocks.get(li) ffn_mhc = ffn_mhc_blocks.get(li) a_norm = attn_norms[li] f_norm = ffn_norms[li] rc, rs = rope_caches[gpu] kv_cache = SimpleKVCache(head_dim=hd, max_seq=8192, device=dev) X = forward_layer(X, w, li, cfg, rc, rs, attn_mhc, ffn_mhc, a_norm, f_norm, kv_cache, tid, positions) # Per-layer diagnostic x_max = X.abs().max().item() has_nan = torch.isnan(X.float()).any().item() has_inf = torch.isinf(X.float()).any().item() # Stream 0 (primary) x0 = X[:, 0, :] x0_mean = x0.float().abs().mean().item() x0_std = x0.float().std().item() layer_diags.append({ 'layer': li, 'gpu': gpu, 'x_max': x_max, 'x0_mean': x0_mean, 'x0_std': x0_std, 'nan': has_nan, 'inf': has_inf }) if has_nan or has_inf: print(f" ❌ Layer {li}: NaN={has_nan} Inf={has_inf} — STOPPING") break X = X.to('cuda:0') torch.cuda.set_device(0) # Final norm + lm_head x_out = X[:, 0, :] if final_norm_w is not None: xf = x_out.float() rms = xf.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() x_out = (xf * rms * final_norm_w.float()).bfloat16() logits = torch.nn.functional.linear(x_out, lm_w) # Results print(f"\n === Minimal E2E Test Results ===") print(f" Logits: min={logits.float().min().item():.2f} max={logits.float().max().item():.2f} " f"spread={logits.float().max().item() - logits.float().min().item():.2f}") print(f" NaN={torch.isnan(logits.float()).any().item()} " f"Inf={torch.isinf(logits.float()).any().item()}") top10_vals, top10_ids = torch.topk(logits[0], 10) print(f" Top-10 predictions:") for i, (tid_v, val) in enumerate(zip(top10_ids, top10_vals)): tok_str = tokenizer.decode([tid_v.item()]) print(f" {i+1}. '{tok_str}' (id={tid_v.item()}, logit={val.item():.3f})") # Print residual stream evolution print(f"\n Residual stream evolution (stream 0):") for d in layer_diags[::5]: # Every 5th layer print(f" L{d['layer']:2d}: |X|={d['x_max']:.1f} " f"mean={d['x0_mean']:.1f} std={d['x0_std']:.1f} " f"nan={d['nan']} inf={d['inf']}") # Always print last if layer_diags: d = layer_diags[-1] print(f" L{d['layer']:2d}: |X|={d['x_max']:.1f} " f"mean={d['x0_mean']:.1f} std={d['x0_std']:.1f} " f"nan={d['nan']} inf={d['inf']}") # Check for reasonable output spread = logits.float().max().item() - logits.float().min().item() if spread < 1.0: print(f" ⚠️ Logit spread {spread:.2f} is very low — model is essentially uniform") else: print(f" ✓ Logit spread {spread:.2f} looks reasonable") return logits, layer_diags if __name__ == "__main__": main()