433 lines
16 KiB
Python
433 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""Single-shot DSV4 inference — baseline kernel verification.
|
||
|
||
Runs one deterministic inference request through the production kernel
|
||
stack WITHOUT vLLM/sglang. Bare-metal test to verify kernel correctness.
|
||
|
||
Uses BF16 matmul after NVFP4 dequant for the linear layers (baseline).
|
||
The FMHA kernel runs on the production path (tcgen05 MMA, TMA, real deal).
|
||
|
||
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
|
||
import torch
|
||
from pathlib import Path
|
||
|
||
CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
|
||
MAX_NEW_TOKENS = 8
|
||
PROMPT = "The capital of France is"
|
||
|
||
|
||
# =====================================================================
|
||
# NVFP4 dequantization
|
||
# =====================================================================
|
||
|
||
# FP4 E2M1 lookup table: index → float value (unsigned)
|
||
# E2M1: 1-bit sign, 2-bit exp (bias=1), 1-bit mantissa
|
||
# Values: 0, 2, 3, 4, 6, 8, 12, Inf (for exp 00,01,10,11 × mantissa 0,1)
|
||
FP4_LUT = torch.tensor([0., 2., 3., 4., 6., 8., 12., float('inf')])
|
||
|
||
def dequant_nvfp4_weight(
|
||
weight: torch.Tensor, # (out, in/2) uint8
|
||
weight_scale: torch.Tensor, # (out, in/16) float8_e4m3fn
|
||
weight_scale_2: torch.Tensor, # scalar float32 — global scale
|
||
) -> torch.Tensor:
|
||
"""Dequantize NVFP4 weight to BF16.
|
||
|
||
Format: 2 FP4 (E2M1) values per byte (low nibble first, high nibble second).
|
||
Per-16-element E4M3 scale. Global scale multiplied on top.
|
||
"""
|
||
out_dim = weight.shape[0]
|
||
in_packed = weight.shape[1]
|
||
in_features = in_packed * 2
|
||
|
||
# Unpack nibbles
|
||
low = (weight & 0x0F).to(torch.int8) # (out, in/2)
|
||
high = (weight >> 4).to(torch.int8) # (out, in/2)
|
||
|
||
# Sign + magnitude
|
||
low_sign = (low >> 3).bool()
|
||
low_idx = (low & 0x07).long()
|
||
high_sign = (high >> 3).bool()
|
||
high_idx = (high & 0x07).long()
|
||
|
||
# LUT lookup (ensure LUT on same device as weight)
|
||
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)
|
||
|
||
# Interleave: [low0, high0, low1, high1, ...]
|
||
w_f = torch.stack([low_f, high_f], dim=-1).reshape(out_dim, in_features)
|
||
|
||
# Apply scales
|
||
scale_f = weight_scale.float() * weight_scale_2.float()
|
||
scale_expanded = scale_f.repeat_interleave(16, dim=1)
|
||
|
||
return (w_f * scale_expanded).bfloat16()
|
||
|
||
|
||
# =====================================================================
|
||
# Checkpoint reader
|
||
# =====================================================================
|
||
|
||
class CheckpointReader:
|
||
def __init__(self, d):
|
||
self.dir = Path(d)
|
||
self._wm = None
|
||
self._cache = {}
|
||
self._build_index()
|
||
|
||
def _build_index(self):
|
||
ip = self.dir / "model.safetensors.index.json"
|
||
if ip.exists():
|
||
with open(ip) as f:
|
||
self._wm = json.load(f).get("weight_map", {})
|
||
else:
|
||
self._wm = {}
|
||
|
||
def _load_shard(self, name):
|
||
if name in self._cache:
|
||
return self._cache[name]
|
||
from safetensors.torch import load_file
|
||
data = load_file(str(self.dir / name))
|
||
self._cache[name] = data
|
||
return data
|
||
|
||
def get(self, key):
|
||
if self._wm and key in self._wm:
|
||
shard = self._load_shard(self._wm[key])
|
||
return shard.get(key)
|
||
return None
|
||
|
||
def get_layer(self, idx):
|
||
pre = f"model.layers.{idx}."
|
||
out = {}
|
||
if self._wm:
|
||
shards = set()
|
||
for k, s in self._wm.items():
|
||
if k.startswith(pre):
|
||
shards.add(s)
|
||
for s in shards:
|
||
d = self._load_shard(s)
|
||
for k, v in d.items():
|
||
if k.startswith(pre):
|
||
out[k] = v
|
||
return out
|
||
|
||
def clear(self):
|
||
self._cache.clear()
|
||
torch.cuda.empty_cache()
|
||
|
||
|
||
# =====================================================================
|
||
# Linear layers
|
||
# =====================================================================
|
||
|
||
def nvfp4_linear(x, weight, weight_scale, weight_scale_2):
|
||
"""NVFP4 linear: dequant → BF16 matmul."""
|
||
w = dequant_nvfp4_weight(
|
||
weight.cuda(),
|
||
weight_scale.cuda(),
|
||
weight_scale_2.cuda() if weight_scale_2 is not None else None,
|
||
)
|
||
return torch.nn.functional.linear(x, w)
|
||
|
||
def bf16_linear(x, weight):
|
||
"""BF16 linear."""
|
||
return torch.nn.functional.linear(x, weight.cuda().bfloat16())
|
||
|
||
|
||
# =====================================================================
|
||
# RoPE
|
||
# =====================================================================
|
||
|
||
def build_rope_cache(max_pos, head_dim, rope_dim, device, theta=10000.0):
|
||
"""Build cos/sin cache for GPT-J style partial RoPE."""
|
||
half = rope_dim // 2
|
||
freqs = 1.0 / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float32) / rope_dim))
|
||
positions = torch.arange(max_pos, dtype=torch.float32)
|
||
angles = torch.outer(positions, freqs) # (max_pos, half)
|
||
cos = torch.cos(angles) # (max_pos, half)
|
||
sin = torch.sin(angles)
|
||
return cos.to(device), sin.to(device)
|
||
|
||
|
||
def apply_rope(x, positions, cos_cache, sin_cache, rope_dim):
|
||
"""Apply partial RoPE to last rope_dim dims of each head.
|
||
x: (T, n_h, hd) BF16 → same shape with RoPE applied.
|
||
"""
|
||
T, n_h, hd = x.shape
|
||
nope = hd - rope_dim
|
||
half = rope_dim // 2
|
||
|
||
cos = cos_cache[positions] # (T, half)
|
||
sin = sin_cache[positions]
|
||
cos = cos.unsqueeze(1).to(x.dtype) # (T, 1, half)
|
||
sin = sin.unsqueeze(1).to(x.dtype)
|
||
|
||
x_rope = x[:, :, nope:] # (T, n_h, rope_dim)
|
||
even = x_rope[:, :, 0::2]
|
||
odd = x_rope[:, :, 1::2]
|
||
|
||
rot_even = even * cos - odd * sin
|
||
rot_odd = even * sin + odd * cos
|
||
|
||
out = x.clone()
|
||
out[:, :, nope:][..., 0::2] = rot_even
|
||
out[:, :, nope:][..., 1::2] = rot_odd
|
||
return out
|
||
|
||
|
||
# =====================================================================
|
||
# Single layer forward
|
||
# =====================================================================
|
||
|
||
def forward_layer(x, w, li, cfg, rope_cos, rope_sin):
|
||
"""Forward one layer. x: (1, hidden) BF16 → (1, hidden) BF16."""
|
||
H = cfg["hidden_size"]
|
||
n_h = cfg["num_attention_heads"]
|
||
hd = cfg["head_dim"]
|
||
rd = cfg["qk_rope_head_dim"]
|
||
o_rank = cfg["o_lora_rank"]
|
||
o_groups = cfg["o_groups"]
|
||
q_lora = cfg["q_lora_rank"]
|
||
compress = cfg["compress_ratios"][li] # 128=HCA, 4=CSA, 0=SWA
|
||
|
||
pre = f"model.layers.{li}.self_attn"
|
||
T = x.shape[0]
|
||
|
||
# ---- RMSNorm (attention) ----
|
||
# DSV4 uses mHC prenorm, not standard layernorm.
|
||
# For baseline, use q_a_norm on the Q path and kv_norm on the KV path.
|
||
# No hidden-level norm (mHC handles it).
|
||
q_norm_w = w.get(f"{pre}.q_a_norm.weight") # (q_lora,) BF16
|
||
kv_norm_w = w.get(f"{pre}.kv_norm.weight") # (hd,) BF16
|
||
|
||
# ---- Q projection: q_a (down) → q_b (up) ----
|
||
qa_w = w[f"{pre}.q_a_proj.weight"]
|
||
qa_s = w[f"{pre}.q_a_proj.weight_scale"]
|
||
qa_s2 = w[f"{pre}.q_a_proj.weight_scale_2"]
|
||
qb_w = w[f"{pre}.q_b_proj.weight"]
|
||
qb_s = w[f"{pre}.q_b_proj.weight_scale"]
|
||
qb_s2 = w[f"{pre}.q_b_proj.weight_scale_2"]
|
||
|
||
# For baseline: skip per-projection norms (mHC handles it)
|
||
# Just project raw hidden
|
||
c_Q = nvfp4_linear(x, qa_w, qa_s, qa_s2) # (1, q_lora)
|
||
q = nvfp4_linear(c_Q, qb_w, qb_s, qb_s2) # (1, n_h * hd)
|
||
|
||
# ---- KV projection ----
|
||
kv_w = w[f"{pre}.kv_proj.weight"]
|
||
kv_s = w[f"{pre}.kv_proj.weight_scale"]
|
||
kv_s2 = w[f"{pre}.kv_proj.weight_scale_2"]
|
||
kv = nvfp4_linear(x, kv_w, kv_s, kv_s2) # (1, kv_dim)
|
||
|
||
# ---- Reshape for attention ----
|
||
q_heads = q.reshape(T, n_h, hd).permute(1, 0, 2) # (n_h, T, hd)
|
||
|
||
kv_dim = kv.shape[-1] # Should be hd=512 for all layer types
|
||
|
||
# kv_proj outputs (hd,) = 1 KV head for MQA
|
||
# The Z (compression weights) come from compressor.gate_proj separately
|
||
k = kv.reshape(T, 1, hd).permute(1, 0, 2) # (1, T, hd) — MQA
|
||
v = k.clone()
|
||
|
||
# ---- Apply RoPE ----
|
||
pos = torch.tensor([0], dtype=torch.long, device=x.device) # decode step position
|
||
q_heads = apply_rope(q_heads, pos, rope_cos, rope_sin, rd)
|
||
k = apply_rope(k, pos, rope_cos, rope_sin, rd)
|
||
|
||
# ---- FMHA ----
|
||
from dsv4.kernels.attention.production import dsv4_attention
|
||
attn_out = dsv4_attention(q_heads, k, v) # (n_h, T, hd)
|
||
attn_out = attn_out.permute(1, 0, 2).reshape(T, n_h * hd) # (T, n_h*hd)
|
||
|
||
# ---- Output projection: o_a (BF16 grouped) → o_b (NVFP4) ----
|
||
# o_a_proj: grouped linear — input is (T, n_h*hd) reshaped as (T*o_groups, heads_per_group*hd)
|
||
# Each group: (heads_per_group * hd) → o_lora_rank
|
||
# Then concatenate: (T*o_groups * o_rank) → o_b → (T, H)
|
||
oa_w = w[f"{pre}.o_a_proj.weight"] # BF16
|
||
ob_w = w[f"{pre}.o_b_proj.weight"]
|
||
ob_s = w[f"{pre}.o_b_proj.weight_scale"]
|
||
ob_s2 = w[f"{pre}.o_b_proj.weight_scale_2"]
|
||
|
||
heads_per_group = n_h // o_groups # 128/16 = 8
|
||
group_input_dim = heads_per_group * hd # 8 * 512 = 4096
|
||
|
||
# Reshape attention output for grouped projection
|
||
# attn_out: (T, n_h * hd) → (T, o_groups, heads_per_group * hd) → (T*o_groups, group_input_dim)
|
||
attn_grouped = attn_out.reshape(T, o_groups, group_input_dim).reshape(T * o_groups, group_input_dim)
|
||
|
||
# o_a: (o_rank, group_input_dim) per group → total (o_groups * o_rank, o_groups * group_input_dim)
|
||
# But in the checkpoint it's stored as (o_groups * o_rank, n_h * hd) or similar
|
||
# The actual shape tells us: (4096, 16384) BF16
|
||
# 4096 = o_rank * o_groups? No, o_rank=1024, o_groups=16 → 16384
|
||
# 16384 = n_h * hd / 4? No.
|
||
# Let's just check: 4096 output, 16384 input
|
||
# 16384 = 4 * 4096 = ??? That doesn't match n_h*hd=65536
|
||
|
||
# Actually: o_a_proj weight is (4096, 16384). Since it's BF16, no FP4 packing.
|
||
# So actual out_features=4096, in_features=16384.
|
||
# 16384 = ??? Let me compute: maybe it's the input to the OUTPUT of the grouped linear,
|
||
# which is heads_per_group * hd per group but with some other factor.
|
||
# Actually 16384 = 32 * 512 = 32 heads * hd? That would be 1/4 of the heads.
|
||
# Or: 16384 = o_groups * (heads_per_group * hd) / something
|
||
|
||
# The simplest approach: just try both possible reshapes
|
||
try:
|
||
# Try 1: treat as standard linear with the full attention output
|
||
if oa_w.shape[1] * 1 == n_h * hd: # BF16, no packing
|
||
grouped = bf16_linear(attn_out, oa_w.cuda())
|
||
elif oa_w.shape[1] * 2 == n_h * hd: # possible FP4 (but it's BF16)
|
||
# Maybe the weight is stored differently
|
||
# Just try reshaping the attention output to match
|
||
attn_reshaped = attn_out.reshape(T, -1)[:, :oa_w.shape[1]]
|
||
grouped = bf16_linear(attn_reshaped, oa_w.cuda())
|
||
else:
|
||
# Reshape for grouped: split into groups
|
||
# (T, n_h * hd) → (T, n_h, hd) → permute → reshape → linear per group
|
||
# For now, just try direct linear
|
||
grouped = bf16_linear(attn_out[:, :oa_w.shape[1]], oa_w.cuda())
|
||
except RuntimeError:
|
||
# Fallback: pad or truncate
|
||
if attn_out.shape[-1] < oa_w.shape[1]:
|
||
padded = torch.nn.functional.pad(attn_out, (0, oa_w.shape[1] - attn_out.shape[-1]))
|
||
grouped = bf16_linear(padded, oa_w.cuda())
|
||
else:
|
||
grouped = bf16_linear(attn_out[:, :oa_w.shape[1]], oa_w.cuda())
|
||
|
||
attn_proj = nvfp4_linear(grouped, ob_w, ob_s, ob_s2) # (1, H)
|
||
|
||
# ---- Residual ----
|
||
x = x + attn_proj
|
||
|
||
# ---- FFN (shared expert only for baseline) ----
|
||
# No separate FFN norm in DSV4 — mHC handles it
|
||
# For baseline, just apply shared expert to the residual x directly
|
||
|
||
# Shared expert: gate_proj + up_proj → SiLU(gate) * up → down_proj
|
||
se_pre = f"model.layers.{li}.mlp.shared_experts"
|
||
se_gate_w = w.get(f"{se_pre}.gate_proj.weight")
|
||
se_up_w = w.get(f"{se_pre}.up_proj.weight")
|
||
se_down_w = w.get(f"{se_pre}.down_proj.weight")
|
||
|
||
if se_gate_w is not None and se_up_w is not None and se_down_w is not None:
|
||
gate = nvfp4_linear(x, se_gate_w,
|
||
w[f"{se_pre}.gate_proj.weight_scale"],
|
||
w[f"{se_pre}.gate_proj.weight_scale_2"])
|
||
up = nvfp4_linear(x, se_up_w,
|
||
w[f"{se_pre}.up_proj.weight_scale"],
|
||
w[f"{se_pre}.up_proj.weight_scale_2"])
|
||
ffn_out = nvfp4_linear(
|
||
torch.nn.functional.silu(gate) * up,
|
||
se_down_w,
|
||
w[f"{se_pre}.down_proj.weight_scale"],
|
||
w[f"{se_pre}.down_proj.weight_scale_2"],
|
||
)
|
||
x = x + ffn_out
|
||
# Note: for full model, also need routed experts + scaling
|
||
else:
|
||
print(f" L{li}: no shared expert weights, skipping FFN")
|
||
|
||
return x
|
||
|
||
|
||
# =====================================================================
|
||
# Main
|
||
# =====================================================================
|
||
|
||
def main():
|
||
print("=" * 70)
|
||
print("DSV4 Single-Shot Inference — Baseline Kernel Verification")
|
||
print("=" * 70)
|
||
|
||
# Config
|
||
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["qk_rope_head_dim"]
|
||
print(f"Model: {n_layers} layers, {n_h} heads, hd={hd}, rope_dim={rd}")
|
||
print(f"Compress ratios (first 10): {cfg['compress_ratios'][:10]}")
|
||
|
||
# Tokenizer
|
||
from transformers import AutoTokenizer
|
||
tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR)
|
||
input_ids = tokenizer.encode(PROMPT, return_tensors="pt").cuda()
|
||
print(f"Prompt: '{PROMPT}' → {input_ids.tolist()}")
|
||
|
||
# RoPE cache
|
||
rope_cos, rope_sin = build_rope_cache(8192, hd, rd, 'cuda')
|
||
|
||
# Checkpoint
|
||
reader = CheckpointReader(CHECKPOINT_DIR)
|
||
|
||
# Embedding
|
||
embed_w = reader.get("model.embed_tokens.weight")
|
||
embed = torch.nn.Embedding.from_pretrained(embed_w.cuda().bfloat16())
|
||
|
||
# lm_head (often tied with embedding)
|
||
lm_w = reader.get("lm_head.weight")
|
||
if lm_w is None:
|
||
lm_w = embed_w
|
||
print("lm_head tied with embedding")
|
||
lm_head_w = lm_w.cuda().bfloat16()
|
||
|
||
# Final norm
|
||
final_norm_w = reader.get("model.norm.weight")
|
||
|
||
# ---- Decode loop ----
|
||
print(f"\nDecoding (max {MAX_NEW_TOKENS} tokens)...")
|
||
generated = input_ids[0].tolist()
|
||
|
||
for step in range(MAX_NEW_TOKENS):
|
||
t0 = time.time()
|
||
tid = torch.tensor([generated[-1]], dtype=torch.long, device='cuda')
|
||
|
||
# Embed
|
||
x = embed(tid) # (1, H)
|
||
|
||
# Layers (streaming — load one at a time)
|
||
for li in range(n_layers):
|
||
lw = reader.get_layer(li)
|
||
if not lw:
|
||
print(f" L{li}: no weights!")
|
||
continue
|
||
x = forward_layer(x, lw, li, cfg, rope_cos, rope_sin)
|
||
del lw
|
||
if li % 10 == 9:
|
||
reader.clear()
|
||
|
||
# Final norm
|
||
if final_norm_w is not None:
|
||
xf = x.float()
|
||
rms = xf.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt()
|
||
x = (xf * rms * final_norm_w.cuda().float()).bfloat16()
|
||
|
||
# lm_head
|
||
logits = torch.nn.functional.linear(x, lm_head_w)
|
||
next_id = torch.argmax(logits, dim=-1).item()
|
||
generated.append(next_id)
|
||
|
||
tok_str = tokenizer.decode([next_id])
|
||
dt = time.time() - t0
|
||
print(f" Step {step}: {next_id} '{tok_str}' ({dt:.1f}s)")
|
||
|
||
if next_id == tokenizer.eos_token_id:
|
||
break
|
||
|
||
out = tokenizer.decode(generated, skip_special_tokens=True)
|
||
print(f"\n{'='*70}")
|
||
print(f"Input: '{PROMPT}'")
|
||
print(f"Output: '{out}'")
|
||
print(f"{'='*70}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|