biondizzle
6c28c57b6a
The attention output projection first half (wo_a) was using BF16 grouped BMM (torch.bmm). Now uses production Nvfp4GroupedLinear which performs the same grouped GEMM with NVFP4 tensor-core acceleration on Blackwell. The weight is loaded from NVFP4 checkpoint if available, otherwise quantized from BF16 via set_bf16_weight(). Also includes: - NVFP4 gate projection for router (from previous commit) - Compressor position_bias in CUDA kernel (from earlier fix)
879 lines
47 KiB
Python
879 lines
47 KiB
Python
#!/usr/bin/env python3
|
||
"""Single-shot DSV4-Pro inference — Full production pipeline, 8-GPU.
|
||
|
||
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.
|
||
This is the ground truth for vLLM / SGLang integration.
|
||
"""
|
||
import os, sys, time, json, math, argparse, logging
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from pathlib import Path
|
||
|
||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||
log = logging.getLogger("single_shot")
|
||
|
||
def parse_args():
|
||
p = argparse.ArgumentParser()
|
||
p.add_argument('--max-tokens', type=int, default=8192)
|
||
p.add_argument('--prompt', type=str, default=None)
|
||
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('--num-gpus', type=int, default=8)
|
||
p.add_argument('--checkpoint', type=str, default="/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4")
|
||
p.add_argument('--prefill-tokens', type=str, default=None,
|
||
help='Override prompt tokens as comma-separated IDs (e.g. "1,128803,313,128804")')
|
||
return p.parse_args()
|
||
|
||
_args = parse_args()
|
||
CHECKPOINT_DIR = _args.checkpoint
|
||
MAX_NEW_TOKENS = _args.max_tokens
|
||
PROMPT = _args.prompt or "The capital of France is"
|
||
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 (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):
|
||
freqs = 1. / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float32) / rope_dim))
|
||
if rope_type == "yarn" and rope_factor > 1.:
|
||
nf = []
|
||
for f in freqs:
|
||
wl = 2 * math.pi / f
|
||
lo, hi = orig_max / (beta_fast * 2.), orig_max / (beta_slow * 2.)
|
||
if wl < lo: nf.append(f)
|
||
elif wl > hi: nf.append(f / rope_factor)
|
||
else:
|
||
sm = (orig_max / (wl * beta_slow) - rope_factor) / (rope_factor * (beta_fast / beta_slow - 1))
|
||
nf.append((1 - sm) * f / rope_factor + sm * f)
|
||
freqs = torch.tensor(nf, 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(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 = {}
|
||
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)))
|
||
log.info(f"Loaded {len(all_w)} tensors from {len(shards)} shards"); return all_w
|
||
|
||
# =====================================================================
|
||
# RMSNorm
|
||
# =====================================================================
|
||
def rmsnorm(x, weight, eps=1e-6):
|
||
xf = x.float()
|
||
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()
|
||
|
||
# =====================================================================
|
||
# 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)
|
||
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.)
|
||
w = torch.stack([lo_f, hi_f], -1).reshape(O, I)
|
||
s = weight_scale.float().repeat_interleave(16, 1)
|
||
if weight_scale_2 is not None: s = s * weight_scale_2.float()
|
||
return (w * s).bfloat16()
|
||
|
||
def nvfp4_linear_ref(x, weight, weight_scale, weight_scale_2=None, input_scale=None):
|
||
return F.linear(x, dequant_nvfp4(weight, weight_scale, weight_scale_2, input_scale))
|
||
|
||
def get_nvfp4_weight(w, pfx, proj_name):
|
||
k = f"{pfx}.{proj_name}"
|
||
return (w.get(f"{k}.weight"), w.get(f"{k}.weight_scale"),
|
||
w.get(f"{k}.weight_scale_2"), w.get(f"{k}.input_scale"))
|
||
|
||
def do_nvfp4_linear_ref(x, w, pfx, proj_name):
|
||
weight, ws, ws2, isc = get_nvfp4_weight(w, pfx, proj_name)
|
||
if weight is None: return None
|
||
d = x.device
|
||
return nvfp4_linear_ref(x, weight.to(d), ws.to(d),
|
||
ws2.to(d) if ws2 is not None else None,
|
||
isc.to(d) if isc is not None else None)
|
||
|
||
# =====================================================================
|
||
# Production Nvfp4Linear factory
|
||
# =====================================================================
|
||
def make_nvfp4_linear(in_features, out_features, device, all_w, pfx, proj_name):
|
||
from dsv4.layers.linear import Nvfp4Linear
|
||
d = device
|
||
weight, ws, ws2, isc = get_nvfp4_weight(all_w, pfx, proj_name)
|
||
assert weight is not None, f"{pfx}.{proj_name}.weight not found"
|
||
# Checkpoint weight is (N_packed, K_packed) uint8
|
||
# NVFP4 GEMM output dim = N_packed BF16 elements
|
||
# Activation buffer needs K_packed FP4 columns = in_features BF16
|
||
# So: in_features = K_packed * 2, out_features = N_packed
|
||
actual_out = weight.shape[0] # N_packed = GEMM output dimension
|
||
actual_in = weight.shape[1] * 2 # K_packed * 2 = BF16 input dim (for buffer allocation)
|
||
lin = Nvfp4Linear(actual_in, actual_out, max_num_tokens=8192, device=d)
|
||
lin.fp4 = [weight.to(d)]; lin.sf = [ws.to(d)]
|
||
# Global scales for NVFP4 GEMM:
|
||
# gsb (weight global scale) = weight_scale_2 (NOT input_scale * weight_scale_2)
|
||
# gsa (activation global scale) = input_scale from checkpoint
|
||
# Dequant: w = lut[w_packed] * weight_scale * weight_scale_2
|
||
# GEMM: y = (x * scale_a * gsa) @ (w * scale_b * gsb)
|
||
# Nvfp4Linear.finalize_weights does: gsb = gs * ws2_val
|
||
# So to get gsb = ws2_val, set gs = 1.0 and let ws2 do its job
|
||
lin.gs = [1.0] # base gs — finalize_weights will multiply by ws2
|
||
lin.ws2 = [ws2.to(d) if ws2 is not None else None]
|
||
# Set activation global scale from checkpoint input_scale
|
||
isc_val = isc.float().item() if isc is not None else 1.0 / (6.0 * 448.0)
|
||
lin._activation_global_scale = isc_val # gsa = input_scale
|
||
lin.finalize_weights(); return lin
|
||
|
||
# =====================================================================
|
||
# Compressor — CSA (ratio=4) and HCA (ratio=128) [PRODUCTION KERNELS]
|
||
# =====================================================================
|
||
class Compressor:
|
||
"""Production compressor: NVFP4 GEMM projections + CUDA softmax/reduce.
|
||
|
||
Pipeline:
|
||
1. NVFP4 GEMM: hidden_states @ kv_proj → (T, kv_dim) BF16
|
||
2. NVFP4 GEMM: hidden_states @ gate_proj → (T, kv_dim) BF16
|
||
3. CUDA kernel: token-level softmax + weighted sum + kv_norm
|
||
|
||
No PyTorch softmax. No reference fallback.
|
||
"""
|
||
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.kv_lin = None # production Nvfp4Linear for kv_proj
|
||
self.gate_lin = None # production Nvfp4Linear for gate_proj
|
||
self.ape = None; self.kv_norm_w = None
|
||
self._reduce_loaded = False
|
||
|
||
def load(self, w, pfx, dev=None):
|
||
"""Load weights and build production Nvfp4Linear instances."""
|
||
if dev is None: dev = self.device
|
||
# Build production NVFP4 GEMM instances for the two projections
|
||
# kv_proj: in=7168, out=kv_dim (1024 for CSA, 512 for HCA)
|
||
# gate_proj: same shapes
|
||
kv_w, kv_ws, kv_ws2, kv_isc = get_nvfp4_weight(w, pfx, 'kv_proj')
|
||
gate_w, gate_ws, gate_ws2, gate_isc = get_nvfp4_weight(w, pfx, 'gate_proj')
|
||
if kv_w is not None:
|
||
kv_out = kv_w.shape[0] # N_packed
|
||
kv_in = kv_w.shape[1] * 2 # K_packed * 2
|
||
self.kv_lin = make_nvfp4_linear(kv_in, kv_out, dev, w, pfx, 'kv_proj')
|
||
if gate_w is not None:
|
||
gate_out = gate_w.shape[0]
|
||
gate_in = gate_w.shape[1] * 2
|
||
self.gate_lin = make_nvfp4_linear(gate_in, gate_out, dev, w, pfx, 'gate_proj')
|
||
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.kv_lin 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
|
||
|
||
# Step 1-2: NVFP4 GEMM projections → BF16, then cast to FP32 for reduce
|
||
kv = self.kv_lin(hidden_states).float() # (T, kv_dim) FP32
|
||
gate = self.gate_lin(hidden_states).float() # (T, kv_dim) FP32
|
||
|
||
# Position bias is handled inside the CUDA kernel (added to both kv and gate)
|
||
# Step 3: CUDA softmax/reduce kernel
|
||
from dsv4.kernels.compressor.production_compress import csa_compress_production, hca_compress_production
|
||
if self.is_csa:
|
||
compressed = csa_compress_production(
|
||
kv, gate, self.ape, self.kv_norm_w, m=r)
|
||
else:
|
||
compressed = hca_compress_production(
|
||
kv, gate, self.ape, self.kv_norm_w, m=r)
|
||
|
||
if compressed.shape[0] == 0: return None, None, None
|
||
comp_pos = torch.tensor([positions[(bi+1)*r - 1].item() if positions.numel() > (bi+1)*r - 1 else 0
|
||
for bi in range(n_complete)],
|
||
dtype=torch.long, device=dev)
|
||
return compressed, comp_pos, torch.zeros(1, T, n_complete, dtype=torch.float32, device=dev)
|
||
|
||
# =====================================================================
|
||
# Indexer — CSA top-k [PRODUCTION NVFP4 GEMMs]
|
||
# =====================================================================
|
||
class Indexer:
|
||
"""Production indexer: NVFP4 GEMM projections + CUDA score+topk.
|
||
|
||
Pipeline:
|
||
1. NVFP4 GEMM: q_a (lora) @ q_b_proj → (T, n_ih * ihd) BF16
|
||
2. NVFP4 GEMM: hidden_states @ weights_proj → (T, n_ih) BF16
|
||
3. CUDA kernel: ReLU(Q·K) * w_head → score, top-k selection
|
||
"""
|
||
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_lin = None # production Nvfp4Linear for q_b_proj
|
||
self.wp_lin = None # production Nvfp4Linear for weights_proj
|
||
self.compressor = None
|
||
|
||
def load(self, w, pfx, dev=None):
|
||
if dev is None: dev = self.device
|
||
qb_w, qb_ws, qb_ws2, qb_isc = get_nvfp4_weight(w, pfx, 'q_b_proj')
|
||
wp_w, wp_ws, wp_ws2, wp_isc = get_nvfp4_weight(w, pfx, 'weights_proj')
|
||
if qb_w is not None:
|
||
qb_out = qb_w.shape[0]
|
||
qb_in = qb_w.shape[1] * 2
|
||
self.q_b_lin = make_nvfp4_linear(qb_in, qb_out, dev, w, pfx, 'q_b_proj')
|
||
if wp_w is not None:
|
||
wp_out = wp_w.shape[0]
|
||
wp_in = wp_w.shape[1] * 2
|
||
self.wp_lin = make_nvfp4_linear(wp_in, wp_out, dev, w, pfx, 'weights_proj')
|
||
if f"{pfx}.compressor.kv_proj.weight" in w:
|
||
self.compressor = Compressor(4, self.ihd, 7168, dev)
|
||
self.compressor.load(w, f"{pfx}.compressor", dev)
|
||
|
||
def forward(self, q_lora, hidden_states, comp_indexer_kv, positions):
|
||
if self.q_b_lin 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 = self.q_b_lin(q_lora).reshape(T, self.n_ih, self.ihd)
|
||
w_h = self.wp_lin(hidden_states) # (T, n_ih)
|
||
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
|
||
|
||
# =====================================================================
|
||
# KV Cache
|
||
# =====================================================================
|
||
class KVCache:
|
||
def __init__(self, head_dim, window_size=128, device='cuda:0'):
|
||
self.hd, self.ws, self.dev = head_dim, window_size, device
|
||
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
|
||
|
||
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)
|
||
|
||
def add_compressed(self, ckv, cpos, idx_kv=None):
|
||
if ckv is None: return
|
||
self.comp_kv = ckv if self.comp_kv is None else torch.cat([self.comp_kv, ckv])
|
||
self.comp_pos = cpos if self.comp_pos is None else torch.cat([self.comp_pos, cpos])
|
||
self.n_comp = self.comp_kv.shape[0]
|
||
if idx_kv is not None:
|
||
self.comp_idx_kv = idx_kv if self.comp_idx_kv is None else torch.cat([self.comp_idx_kv, idx_kv])
|
||
|
||
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()
|
||
idx = torch.arange(self.swa_head, self.swa_head + self.ws) % self.ws
|
||
return self.swa[idx].clone(), self.swa_pos[idx].clone()
|
||
|
||
# =====================================================================
|
||
# 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())
|
||
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
|
||
# =====================================================================
|
||
def _run_production_fmha(q_heads, all_kv, n_h, hd, T, seq_len, scale, dev, li, w, pfx):
|
||
from dsv4.kernels.attention.production import dsv4_attention
|
||
# Head-packed dispatch: single kernel launch for all 128 heads (MQA: 1 KV head shared)
|
||
q = q_heads.permute(1, 0, 2).contiguous() # (n_h, T, hd)
|
||
k = all_kv.unsqueeze(0).contiguous() # (1, N, hd) — MQA single KV head
|
||
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(q=q, k=k, v=v, scale=scale, n_comp=0, sink_bias=sink_bias)
|
||
return attn_out.permute(1, 0, 2) # (T, n_h, hd)
|
||
|
||
# =====================================================================
|
||
# Attention — ALL production kernels
|
||
# =====================================================================
|
||
def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
|
||
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)
|
||
|
||
# 1. Q: q_a (NVFP4 GEMM) → q_a_norm → q_b (NVFP4 GEMM) → q_b_norm
|
||
q_a = prod_lin['q_a'](x_normed)
|
||
if VERBOSE >= 2 and li < 3:
|
||
# Compare q_a with PyTorch reference
|
||
q_a_ref = do_nvfp4_linear_ref(x_normed, w, pfx, 'q_a_proj')
|
||
if q_a_ref is not None:
|
||
cos_qa = torch.nn.functional.cosine_similarity(q_a.flatten().float(), q_a_ref.flatten().float(), dim=0).item()
|
||
print(f" L{li} q_a: |prod|={q_a.abs().max().item():.6f} |ref|={q_a_ref.abs().max().item():.6f} cos={cos_qa:.6f}", flush=True)
|
||
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); 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 (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)
|
||
|
||
# 3. Compressor → compressed KV
|
||
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:
|
||
comp_kv_3d = comp_kv.unsqueeze(1)
|
||
comp_kv_3d = _apply_rope(comp_kv_3d, comp_pos, rope_cos, rope_sin, rd)
|
||
comp_kv = comp_kv_3d.squeeze(1)
|
||
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.add_compressed(comp_kv, comp_pos, comp_idx_kv)
|
||
|
||
# 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 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)
|
||
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
|
||
|
||
# 6. Production FMHA
|
||
attn_out = _run_production_fmha(q_heads, all_kv, n_h, hd, T, seq_len, scale, dev, li, w, pfx)
|
||
if VERBOSE >= 2 and li < 3:
|
||
# Compare with PyTorch reference
|
||
k_exp = all_kv.unsqueeze(0).expand(n_h, -1, -1).contiguous()
|
||
v_exp = k_exp.clone()
|
||
q_in = q_heads.permute(1, 0, 2)
|
||
ref_scores = torch.matmul(q_in, k_exp.transpose(-1, -2)) * scale
|
||
ref_attn = torch.matmul(torch.softmax(ref_scores.float(), -1).bfloat16(), v_exp).permute(1, 0, 2)
|
||
cos_sim = torch.nn.functional.cosine_similarity(attn_out.flatten().float(), ref_attn.flatten().float(), dim=0).item()
|
||
print(f" L{li} FMHA: |prod|={attn_out.abs().max().item():.6f} |ref|={ref_attn.abs().max().item():.6f} cos={cos_sim:.6f}", flush=True)
|
||
# 7. Inverse RoPE
|
||
attn_out = _apply_rope(attn_out, positions, rope_cos, rope_sin, rd, inverse=True)
|
||
|
||
# 8. Output: wo_a (NVFP4 grouped GEMM) + wo_b (NVFP4 GEMM)
|
||
wo_a_lin = prod_lin.get('o_a')
|
||
if wo_a_lin is not None:
|
||
# Nvfp4GroupedLinear: (T, n_h, hd) → (T, n_groups, o_rank) → flatten for o_b
|
||
g_3d = wo_a_lin.run(attn_out) # (T, n_groups, o_rank) BF16
|
||
g_flat = g_3d.reshape(T, -1) # (T, n_groups * o_rank) BF16
|
||
F_attn = prod_lin['o_b'](g_flat)
|
||
else:
|
||
# BF16 grouped BMM fallback (should not happen in production)
|
||
hpg_fb = n_h // o_groups; gid_fb = hpg_fb * hd
|
||
oa_full = w.get(f"{pfx}.o_a_proj.weight")
|
||
if oa_full is not None:
|
||
oa_bf = oa_full.bfloat16().to(dev); a_flat = attn_out.reshape(T, n_h * hd)
|
||
a_grp = a_flat.reshape(T, o_groups, gid_fb); oa_3d = oa_bf.reshape(o_groups, o_rank, gid_fb)
|
||
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)
|
||
else:
|
||
log.warning(f"L{li}: No o_a_proj weight, zero attention output")
|
||
F_attn = torch.zeros(T, cfg["hidden_size"], dtype=torch.bfloat16, device=dev)
|
||
if VERBOSE >= 2 and li < 3:
|
||
print(f" L{li} F_attn: |F_attn|={F_attn.abs().max().item():.6f}", flush=True)
|
||
return F_attn, q_a
|
||
|
||
# =====================================================================
|
||
# MoE — production kernels
|
||
# =====================================================================
|
||
def moe_forward(x, li, moe_runner, se_runner, router, token_id):
|
||
# Ensure token_id is on same GPU as router
|
||
token_id_dev = token_id.to(x.device) if token_id.device != x.device else token_id
|
||
topk_w, topk_ids = router(x, token_ids=token_id_dev)
|
||
torch.cuda.synchronize(x.device)
|
||
if topk_ids.max().item() >= 384 or topk_ids.min().item() < 0:
|
||
print(f" L{li} BAD topk_ids: min={topk_ids.min().item()} max={topk_ids.max().item()}", flush=True)
|
||
if VERBOSE >= 2 and li < 3:
|
||
print(f" L{li} MoE input: |x|={x.abs().max().item():.4f} has_nan={torch.isnan(x).any().item()}", flush=True)
|
||
routed_out = moe_runner.run(x, topk_w, topk_ids)
|
||
if VERBOSE >= 2 and li < 3:
|
||
print(f" L{li} MoE routed: |out|={routed_out.abs().max().item():.4f} has_nan={torch.isnan(routed_out).any().item()}", flush=True)
|
||
shared_out = se_runner.run(x)
|
||
if VERBOSE >= 2 and li < 3:
|
||
has_nan = torch.isnan(shared_out).any().item()
|
||
out_max = shared_out.abs().max().item() if not has_nan else float('nan')
|
||
print(f" L{li} MoE shared: |out|={out_max:.4f} has_nan={has_nan}", flush=True)
|
||
# Check weight integrity
|
||
if hasattr(se_runner, '_l1_mat_b') and se_runner._l1_mat_b is not None:
|
||
wb = se_runner._l1_mat_b.view(torch.uint8)
|
||
print(f" L{li} SE l1 weight: shape={list(se_runner._l1_mat_b.shape)} dtype={se_runner._l1_mat_b.dtype} uint8_range=[{wb.min().item()},{wb.max().item()}]", flush=True)
|
||
if hasattr(se_runner, '_l1_scale_b') and se_runner._l1_scale_b is not None:
|
||
sb = se_runner._l1_scale_b.float()
|
||
print(f" L{li} SE l1 scale: shape={list(se_runner._l1_scale_b.shape)} dtype={se_runner._l1_scale_b.dtype} float_range=[{sb.min().item():.6f},{sb.max().item():.6f}] has_nan={torch.isnan(sb).any().item()}", flush=True)
|
||
print(f" L{li} SE gsa: l1={se_runner._l1_activation_global_scale:.6f} l2={se_runner._l2_activation_global_scale:.6f} gsb: l1={se_runner._l1_gsb[0].item():.6f} l2={se_runner._l2_gsb[0].item():.6f}", flush=True)
|
||
return routed_out + shared_out
|
||
|
||
# =====================================================================
|
||
# Layer forward
|
||
# =====================================================================
|
||
def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
|
||
attn_mhc, ffn_mhc, attn_norm_w, ffn_norm_w,
|
||
kv_cache, positions, token_id,
|
||
compressor=None, indexer=None,
|
||
moe_runner=None, se_runner=None, router=None,
|
||
prod_lin=None):
|
||
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)
|
||
X_mid = attn_mhc.post_block(X_l, F_attn, ctx_a)
|
||
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} "
|
||
f"|Fa|={F_attn.abs().max().item():.1f} |Ff|={F_ffn.abs().max().item():.1f}", flush=True)
|
||
return X_next
|
||
|
||
# =====================================================================
|
||
# MoE weight loading
|
||
# =====================================================================
|
||
def _load_moe_weights_stacked(all_w, li, pfx, dev, moe, cfg):
|
||
n_e = cfg["n_routed_experts"]
|
||
l1_fp4_list, l1_sf_list, l1_gs_list, l1_ws2_list, l1_gsa_list = [], [], [], [], []
|
||
l2_fp4_list, l2_sf_list, l2_gs_list, l2_ws2_list, l2_gsa_list = [], [], [], [], []
|
||
for eid in range(n_e):
|
||
ep = f"{pfx}.experts.{eid}"
|
||
gw, gws, gws2, gisc = get_nvfp4_weight(all_w, ep, 'gate_proj')
|
||
uw, uws, uws2, uisc = get_nvfp4_weight(all_w, ep, 'up_proj')
|
||
if gw is not None and uw is not None:
|
||
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))
|
||
gs = gisc.float().item() if gisc is not None else 1.0 / (6.0 * 448.0)
|
||
l1_gs_list.append(1.0) # gsb base — ws2 will be folded in by _ensure_stacked
|
||
l1_gsa_list.append(gs) # gsa = input_scale
|
||
# weight_scale_2: scalar, folded into global_scale_b
|
||
l1_ws2_list.append(gws2.to(dev) if gws2 is not None else None)
|
||
dw, dws, dws2, 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))
|
||
gs2 = disc.float().item() if disc is not None else 1.0 / (6.0 * 448.0)
|
||
l2_gs_list.append(1.0) # gsb base
|
||
l2_gsa_list.append(gs2) # gsa = input_scale
|
||
l2_ws2_list.append(dws2.to(dev) if dws2 is not None else None)
|
||
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)
|
||
# Save activation global scales — _ensure_stacked will override them from l1_gs (which is 1.0)
|
||
# We must re-set them AFTER _ensure_stacked
|
||
moe._saved_l1_gsa = l1_gsa_list[0] if l1_gsa_list else 1.0 / (6.0 * 448.0)
|
||
moe._saved_l2_gsa = l2_gsa_list[0] if l2_gsa_list else 1.0 / (6.0 * 448.0)
|
||
moe.l1_ws2 = l1_ws2_list
|
||
moe.l2_ws2 = l2_ws2_list
|
||
|
||
def _load_shared_expert_weights(all_w, li, pfx, dev, se, cfg):
|
||
gw, gws, gws2, gisc = get_nvfp4_weight(all_w, f"{pfx}.shared_experts", 'gate_proj')
|
||
uw, uws, uws2, uisc = get_nvfp4_weight(all_w, f"{pfx}.shared_experts", 'up_proj')
|
||
dw, dws, dws2, 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)]
|
||
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)]
|
||
l1_isc = gisc.float().item() if gisc is not None else 1.0 / (6.0 * 448.0)
|
||
se.l1_gs = [1.0] # gsb base — ws2 will be folded in by finalize_weights
|
||
se.l1_ws2 = [gws2.to(dev) if gws2 is not None else None]
|
||
se._l1_activation_global_scale = l1_isc # Will be overridden by _ensure_initialized
|
||
se._saved_l1_gsa = l1_isc # Save for after _ensure_initialized
|
||
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)]
|
||
l2_isc = disc.float().item() if disc is not None else 1.0 / (6.0 * 448.0)
|
||
se.l2_gs = [1.0] # gsb base
|
||
se.l2_ws2 = [dws2.to(dev) if dws2 is not None else None]
|
||
se._l2_activation_global_scale = l2_isc # Will be overridden by _ensure_initialized
|
||
se._saved_l2_gsa = l2_isc # Save for after _ensure_initialized
|
||
|
||
def _cache_layer_weights_no_experts(all_w, n_layers, devices):
|
||
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()
|
||
if k.startswith(pfx) and '.experts.' not in k and '.shared_experts.' not in k}
|
||
cached[li] = w
|
||
if (li+1) % 10 == 0: log.info(f" Cached {li+1}/{n_layers} layers")
|
||
return cached
|
||
|
||
# =====================================================================
|
||
# Main
|
||
# =====================================================================
|
||
def kill_stale_gpu_processes():
|
||
"""Kill any leftover python processes on all GPUs before starting."""
|
||
import subprocess
|
||
try:
|
||
result = subprocess.run(['nvidia-smi', '--query-compute-apps=pid', '--format=csv,noheader'],
|
||
capture_output=True, text=True, timeout=5)
|
||
if result.returncode == 0 and result.stdout.strip():
|
||
pids = [p.strip() for p in result.stdout.strip().split('\n') if p.strip()]
|
||
for pid in pids:
|
||
try:
|
||
import os; os.kill(int(pid), 9)
|
||
log.info(f" Killed stale GPU process {pid}")
|
||
except (ValueError, ProcessLookupError):
|
||
pass
|
||
except Exception as e:
|
||
log.warning(f"Could not check GPU processes: {e}")
|
||
|
||
def main():
|
||
t0 = time.time(); torch.manual_seed(SEED)
|
||
print("=" * 70)
|
||
print("DSV4 Single-Shot Inference - PRODUCTION KERNEL STACK")
|
||
print(" FMHA: 6-warp TMA multi-tile + sink bias")
|
||
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_h = cfg["num_attention_heads"]
|
||
rd = cfg.get("qk_rope_head_dim", 64)
|
||
cr = cfg.get("compress_ratios", [128] * n_layers)
|
||
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" {time.time()-t0:.1f}s")
|
||
|
||
# ---- Phase 2: Build production components ----
|
||
print("Building production components...")
|
||
from dsv4.layers.mhc import mHCLayer
|
||
from dsv4.layers.router import Router
|
||
from dsv4.layers.moe import Nvfp4MoE
|
||
from dsv4.layers.shared_expert import Nvfp4SharedExpert
|
||
|
||
# Kill stale GPU processes from prior runs (OOM, crash, etc.)
|
||
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
|
||
attn_mhcs, ffn_mhcs, attn_norms, ffn_norms = {}, {}, {}, {}
|
||
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"),
|
||
]:
|
||
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)
|
||
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_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(),
|
||
)
|
||
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
|
||
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_a_proj: (16384, 4096) BF16 -> Nvfp4GroupedLinear (16 groups, 1024×4096 each)
|
||
# o_b_proj: (7168, 8192) uint8 -> in=16384, out=7168
|
||
from dsv4.layers.grouped_linear import Nvfp4GroupedLinear
|
||
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')
|
||
# o_a_proj: Nvfp4GroupedLinear (NVFP4 grouped GEMM)
|
||
n_local_groups = cfg.get('o_groups', 16)
|
||
heads_per_group = n_h // n_local_groups
|
||
o_rank_val = cfg.get('o_lora_rank', 1024)
|
||
wo_a = Nvfp4GroupedLinear(
|
||
n_local_groups=n_local_groups,
|
||
heads_per_group=heads_per_group,
|
||
head_dim=hd,
|
||
o_lora_rank=o_rank_val,
|
||
max_num_tokens=8192,
|
||
device=dev,
|
||
)
|
||
oa_w_nvfp4, oa_ws, oa_ws2, oa_isc = get_nvfp4_weight(all_w, pfx, 'o_a_proj')
|
||
if oa_w_nvfp4 is not None and oa_ws is not None:
|
||
# Checkpoint has NVFP4 weights — load directly
|
||
# TODO: Nvfp4GroupedLinear needs a load_nvfp4_weight method
|
||
# For now, dequant and re-quantize via set_bf16_weight
|
||
oa_bf16 = dequant_nvfp4(oa_w_nvfp4, oa_ws, oa_ws2, oa_isc).to(dev)
|
||
wo_a.set_bf16_weight(oa_bf16)
|
||
else:
|
||
# BF16 checkpoint weight
|
||
oa_bf = all_w.get(f"{pfx}.o_a_proj.weight")
|
||
if oa_bf is not None:
|
||
wo_a.set_bf16_weight(oa_bf.bfloat16().to(dev))
|
||
pl['o_a'] = wo_a
|
||
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 (o_a now NVFP4 grouped)")
|
||
|
||
# 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); 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)
|
||
if is_hash:
|
||
router.load_weights(hash_lut=all_w[f"{pfx}.gate.tid2eid"].to(dev, torch.int32))
|
||
else:
|
||
eb = all_w.get(f"{pfx}.gate.e_score_correction_bias")
|
||
# Try NVFP4 gate weights first (production path)
|
||
gate_w, gate_ws, gate_ws2, gate_isc = get_nvfp4_weight(all_w, pfx, 'gate')
|
||
if gate_w is not None and gate_ws is not None:
|
||
# NVFP4 gate: build production Nvfp4Linear
|
||
router.load_weights(e_bias=eb.to(dev, torch.float32))
|
||
gate_lin = make_nvfp4_linear(H, cfg["n_routed_experts"], dev, all_w, pfx, 'gate')
|
||
router.load_nvfp4_gate(gate_lin)
|
||
else:
|
||
# BF16 fallback
|
||
gw = all_w.get(f"{pfx}.gate.weight")
|
||
if gw is not None:
|
||
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
|
||
|
||
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)
|
||
# EAGERLY process stacked weights → K-major + swizzle, free raw tensors
|
||
moe._ensure_stacked()
|
||
# Fix activation global scales — _ensure_stacked sets gsa from l1_gs (which is 1.0)
|
||
if hasattr(moe, '_saved_l1_gsa'):
|
||
moe._l1_activation_global_scale = moe._saved_l1_gsa
|
||
if hasattr(moe, '_saved_l2_gsa'):
|
||
moe._l2_activation_global_scale = moe._saved_l2_gsa
|
||
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)
|
||
# EAGERLY process shared expert weights
|
||
se._ensure_initialized()
|
||
# Fix activation global scales — _ensure_initialized sets gsa from l1_gs (which is 1.0)
|
||
# The correct gsa is the input_scale from the checkpoint, saved in _saved_l1_gsa
|
||
if hasattr(se, '_saved_l1_gsa'):
|
||
se._l1_activation_global_scale = se._saved_l1_gsa
|
||
if hasattr(se, '_saved_l2_gsa'):
|
||
se._l2_activation_global_scale = se._saved_l2_gsa
|
||
se_runners[li] = se
|
||
if (li+1) % 10 == 0: print(f" Built {li+1}/{n_layers} MoE layers")
|
||
torch.cuda.empty_cache()
|
||
|
||
# Global weights
|
||
torch.cuda.set_device(0)
|
||
embed_w = all_w.get("model.embed_tokens.weight")
|
||
embed = torch.nn.Embedding.from_pretrained(embed_w.bfloat16().to('cuda:0'))
|
||
lm_w = all_w.get("lm_head.weight", embed_w).bfloat16().to('cuda:0')
|
||
final_norm_w = all_w.get("model.norm.weight")
|
||
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")
|
||
|
||
# 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)
|
||
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)}
|
||
|
||
# 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)
|
||
for li in range(n_layers):
|
||
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 (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()
|
||
torch.cuda.set_device(0)
|
||
print(f" {time.time()-t0:.1f}s")
|
||
|
||
# Load compressor/indexer weights
|
||
for li in range(n_layers):
|
||
pfx = f"model.layers.{li}.self_attn.compressor"
|
||
if li in compressors: compressors[li].load(layer_w[li], pfx, dev=f"cuda:{li % NUM_GPUS}")
|
||
if li in indexers: indexers[li].load(layer_w[li], f"{pfx}.indexer", dev=f"cuda:{li % NUM_GPUS}")
|
||
print(" Compressors/indexers loaded")
|
||
|
||
# ---- Phase 3: Inference ----
|
||
print(f"\nPhase 3: Inference")
|
||
from transformers import AutoTokenizer
|
||
tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR)
|
||
|
||
bos = tokenizer.bos_token_id or 0
|
||
if _args.prefill_tokens:
|
||
generated = [int(x) for x in _args.prefill_tokens.split(',')]
|
||
else:
|
||
input_ids = [bos, USER_TOKEN]
|
||
input_ids += tokenizer.encode('\n\n' + PROMPT, add_special_tokens=False)
|
||
input_ids.append(ASSISTANT_TOKEN)
|
||
generated = input_ids
|
||
all_tokens = generated.copy()
|
||
print(f"Input: {len(generated)} tokens")
|
||
|
||
# Prefill — one token at a time (decode-style; TODO: batched prefill)
|
||
print(f"Prefilling {len(generated)} tokens...")
|
||
for pi, tid_val in enumerate(generated):
|
||
t1 = time.time()
|
||
tid_int64 = torch.tensor([tid_val], dtype=torch.long, device='cuda:0')
|
||
tid = tid_int64.to(torch.int32) # hash router needs int32
|
||
pos = torch.tensor([pi], dtype=torch.long, device='cuda:0')
|
||
X = mHCLayer.init_state(embed(tid_int64))
|
||
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)
|
||
try:
|
||
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], pos, tid,
|
||
compressors.get(li), indexers.get(li),
|
||
moe_runners.get(li), se_runners.get(li), routers.get(li),
|
||
prod_lin=prod_lins.get(li))
|
||
except Exception as e:
|
||
torch.cuda.synchronize()
|
||
err = torch.cuda.current_stream(gpu).query()
|
||
print(f" CRASH at token {pi} layer {li} gpu {gpu}: {e}", flush=True)
|
||
raise
|
||
if VERBOSE >= 2 and pi == 0 and li < 3:
|
||
torch.cuda.synchronize(gpu)
|
||
print(f" Token {pi} L{li}: OK |X|={X.abs().max().item():.1f}", flush=True)
|
||
X = X.to('cuda:0'); torch.cuda.set_device(0)
|
||
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
|
||
|
||
# Decode
|
||
print(f"\nDecoding (max {MAX_NEW_TOKENS} tokens)...")
|
||
for step in range(MAX_NEW_TOKENS):
|
||
t1 = time.time()
|
||
tid_int64 = torch.tensor([all_tokens[-1]], dtype=torch.long, device='cuda:0')
|
||
tid = tid_int64.to(torch.int32) # hash router needs int32
|
||
dec_pos = torch.tensor([len(all_tokens)-1], dtype=torch.long, device='cuda:0')
|
||
X = mHCLayer.init_state(embed(tid_int64))
|
||
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, tid,
|
||
compressors.get(li), indexers.get(li),
|
||
moe_runners.get(li), se_runners.get(li), routers.get(li),
|
||
prod_lin=prod_lins.get(li))
|
||
X = X.to('cuda:0'); torch.cuda.set_device(0)
|
||
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)
|
||
dt = time.time() - t1
|
||
has_nan = torch.isnan(logits.float()).any().item()
|
||
if step % 1 == 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]))
|
||
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)
|
||
if has_nan: break
|
||
if next_id == tokenizer.eos_token_id: break
|
||
|
||
out = tokenizer.decode(all_tokens, skip_special_tokens=True)
|
||
print(f"\n{'='*70}")
|
||
print(f"Input: '{PROMPT}'")
|
||
print(f"Output: '{out}'")
|
||
print(f"Total: {time.time()-t0:.1f}s")
|
||
print(f"{'='*70}")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|