biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
91c370360ae33fadaec2db933d3cc733f0bb8c2e
nvfp4-megamoe-kernel/single_shot_inference.py

2058 lines
108 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/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
os.environ['CUDA_LAUNCH_BLOCKING'] = '1' # Catch async CUDA errors immediately
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=512)
p.add_argument('--temperature', type=float, default=0.6, help='Sampling temperature (0=greedy)')
p.add_argument('--repetition-penalty', type=float, default=1.1, help='Repetition penalty factor (>1 penalizes repeats)')
p.add_argument('--top-k', type=int, default=50, help='Top-k filtering (0=disabled)')
p.add_argument('--top-p', type=float, default=0.95, help='Top-p (nucleus) filtering (1.0=disabled)')
p.add_argument('--prompt', type=str, default=None)
p.add_argument('--thinking-mode', choices=['thinking', 'chat'], default='thinking',
help='Thinking mode: "thinking" = model reasons first, "chat" = model generates directly')
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('--no-fused-rmsnorm', action='store_true', help='Disable P4 fused RMSNorm+quantize (use unfused path)')
p.add_argument('--warmup-gsa', action='store_true', help='Fix gsa values after first decode step (eliminates amax kernel launches)')
p.add_argument('--profile', action='store_true', help='Profile per-component GPU time using CUDA events')
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")')
p.add_argument('--cuda-graph', action='store_true', help='Capture CUDA graph per layer for decode (eliminates Python dispatch overhead)')
p.add_argument('--max-context', type=int, default=8192, help='Target max context length (determines KV cache pre-allocation)')
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
# Special token IDs — derived from official encoding module strings + tokenizer.
# Do NOT hardcode these; the encoding module defines the canonical token strings.
from encoding.deepseek_v4_encoding import (
thinking_start_token as _THINK_START_STR,
thinking_end_token as _THINK_END_STR,
USER_SP_TOKEN as _USER_STR,
ASSISTANT_SP_TOKEN as _ASSISTANT_STR,
eos_token as _EOS_STR,
bos_token as _BOS_STR,
)
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):
"""In-place RoPE — uses CUDA kernel (1 launch) instead of PyTorch ops (5-6 launches).
P3: Eliminates ~732 kernel launches per token across 61 layers.
"""
try:
from dsv4.ops.rope_cuda import apply_rope
return apply_rope(x, pos, cos, sin, rope_dim, inverse=inverse)
except Exception:
# Fallback to PyTorch (should never happen in production)
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:]
ev = xr[..., 0::2].clone()
od = xr[..., 1::2]
if inverse:
xr[..., 0::2] = (ev * c + od * s).bfloat16()
xr[..., 1::2] = (-ev * s + od * c).bfloat16()
else:
xr[..., 0::2] = (ev * c - od * s).bfloat16()
xr[..., 1::2] = (ev * s + od * c).bfloat16()
return x
# =====================================================================
# 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()
# =====================================================================
# CUDA Graph Decoder — capture per-layer graphs for zero-dispatch decode
# =====================================================================
class CUDAGraphDecoder:
"""Captures and replays CUDA graphs for the decode loop.
Architecture: Eager-break-at-attention (Phase 1)
Each layer's forward is split into two captured sub-graphs with an eager
section between them:
Sub-graph A (pre-attention compute):
mHC pre_block → RMSNorm + quantize → Q/KV projections → RoPE
Outputs: q_heads (roped), kv_roped, x_normed (for compressor), q_a (for indexer)
EAGER section (not captured — dynamic shapes from compressor/FMHA):
KV append → Compressor → Indexer → KV gather → FMHA → Inverse RoPE → o_proj
Outputs: F_attn (for mHC post_block), attn_out (intermediate, for o_proj)
Sub-graph B (post-attention compute + FFN):
mHC post_block (attn) → mHC pre_block (FFN) → RMSNorm → Router → MoE → SE → mHC post_block (FFN)
This eliminates ~80% of Python dispatch overhead while keeping the dynamic-shape
attention path eager. Phase 2 will add paged KV for full capture.
Cross-GPU transfers (X.to(cuda:N)) happen OUTSIDE graphs between layers.
The hc_head + norm + lm_head are captured as a separate graph on cuda:0.
Constraints:
- All tensors must have fixed addresses (pre-allocated)
- No CPU-GPU syncs inside the graph
- The only per-step sync is argmax for sampling (outside graph)
"""
def __init__(self, n_layers, num_gpus, hidden_size, devices):
self.n_layers = n_layers
self.num_gpus = num_gpus
self.hidden_size = hidden_size
self.devices = devices
self.captured = False
# Sub-graphs: 2 per layer (A = pre-attention, B = post-attention + FFN)
self.graphs_a = {} # li -> torch.cuda.CUDAGraph (pre-attention compute)
self.graphs_b = {} # li -> torch.cuda.CUDAGraph (post-attention + FFN compute)
self.lm_graph = None # single graph for hc_head + norm + lm_head on cuda:0
# Pre-allocated I/O buffers — fixed addresses for graph capture
# Sub-graph A I/O (on layer's device)
self.x_in_a = {} # li -> (1, 4, H) BF16 — input X_l to both sub-graphs
self.q_heads_buf = {} # li -> (1, n_h, hd) BF16 — Q heads after RoPE
self.kv_roped_buf = {} # li -> (1, hd) BF16 — KV after RoPE (for KV append)
self.q_a_buf = {} # li -> (1, 1536) BF16 — q_a output (for indexer)
self.x_normed_buf = {} # li -> (1, H) BF16 — normalized input (for compressor)
# Eager → sub-graph B I/O (on layer's device)
self.attn_out_buf = {} # li -> (1, n_h, hd) BF16 — FMHA output after inverse RoPE
# Sub-graph B output (on layer's device)
self.x_next_buf = {} # li -> (1, 4, H) BF16 — output X_next
# lm_head graph buffers (on cuda:0)
self.x_lm_in = None # (1, 4, H) BF16 on cuda:0
self.logits_buf = None # (1, vocab_size) BF16 on cuda:0
def pre_allocate(self, cfg):
"""Pre-allocate all I/O buffers with fixed addresses."""
H = self.hidden_size
n_h = cfg["num_attention_heads"]
hd = cfg["head_dim"]
V = cfg.get("vocab_size", 129280)
q_a_dim = 1536 # q_a_proj output dimension
for li in range(self.n_layers):
dev = self.devices[li % self.num_gpus]
# Sub-graph A I/O
self.x_in_a[li] = torch.zeros(1, 4, H, dtype=torch.bfloat16, device=dev)
self.q_heads_buf[li] = torch.zeros(1, n_h, hd, dtype=torch.bfloat16, device=dev)
self.kv_roped_buf[li] = torch.zeros(1, hd, dtype=torch.bfloat16, device=dev)
self.q_a_buf[li] = torch.zeros(1, q_a_dim, dtype=torch.bfloat16, device=dev)
self.x_normed_buf[li] = torch.zeros(1, H, dtype=torch.bfloat16, device=dev)
# Eager → sub-graph B
self.attn_out_buf[li] = torch.zeros(1, n_h, hd, dtype=torch.bfloat16, device=dev)
# Sub-graph B output
self.x_next_buf[li] = torch.zeros(1, 4, H, dtype=torch.bfloat16, device=dev)
# lm_head graph I/O (cuda:0 only)
self.x_lm_in = torch.zeros(1, 4, H, dtype=torch.bfloat16, device='cuda:0')
self.logits_buf = torch.zeros(1, V, dtype=torch.bfloat16, device='cuda:0')
def capture(self, cfg, attn_mhcs, ffn_mhcs, attn_norms, ffn_norms,
kv_caches, compressors, indexers, moe_runners, se_runners,
routers, prod_lins, layer_w, rope_caches, hc_head,
final_norm_w, lm_w, dec_pos_per_gpu, dec_tid32_per_gpu, comp_rope_caches=None):
"""Capture CUDA graphs for all layers + lm_head.
Must be called after one warmup step so that:
1. All CuTeDSL kernels are compiled and cached
2. gsa values are fixed (from warmup_gsa)
3. CUDA kernels are warmed up (first launch is often slower)
"""
from dsv4.ops.quantize import (
rmsnorm_quantize_nvfp4, mhc_rmsnorm_quantize_nvfp4,
dequantize_nvfp4,
)
from dsv4.layers.mhc import mHCContext
H = self.hidden_size
n_h = cfg["num_attention_heads"]
hd = cfg["head_dim"]
rd = cfg.get("qk_rope_head_dim", 64)
print(" Capturing CUDA graphs for decode (eager-break-at-attention)...", flush=True)
# Capture each layer's two sub-graphs
for li in range(self.n_layers):
gpu = li % self.num_gpus
dev = self.devices[gpu]
torch.cuda.set_device(gpu)
attn_mhc = attn_mhcs.get(li)
ffn_mhc = ffn_mhcs.get(li)
attn_norm_w = attn_norms.get(li)
ffn_norm_w = ffn_norms.get(li)
pl = prod_lins.get(li)
w = layer_w[li]
pfx = f"model.layers.{li}.self_attn"
rope_cos, rope_sin = rope_caches[gpu]
router = routers.get(li)
moe = moe_runners.get(li)
se = se_runners.get(li)
# ---- Sub-graph A: mHC pre + rmsnorm + Q/KV projections ----
graph_a = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph_a):
X_l = self.x_in_a[li]
# mHC pre_block + rmsnorm + NVFP4 quantize (fused P5 path)
A_l_a, B_l_a, C_l_a = attn_mhc._dynamic_params(X_l)
x_quant_attn = mhc_rmsnorm_quantize_nvfp4(
X_l, A_l_a, attn_norm_w.to(dev, torch.float32))
x_normed = dequantize_nvfp4(
x_quant_attn.x_fp4, x_quant_attn.x_sf, x_quant_attn.gsa)
# Q projections: q_a → q_a_norm → q_b → q_b_norm
q_a = pl['q_a'].run_from_quantized(x_quant_attn)
q_norm_w = w.get(f"{pfx}.q_a_norm.weight")
if q_norm_w is not None:
q_a_quant = rmsnorm_quantize_nvfp4(q_a, q_norm_w.to(dev, torch.float32))
q_a_deq = dequantize_nvfp4(q_a_quant.x_fp4, q_a_quant.x_sf, q_a_quant.gsa)
q = pl['q_b'].run_from_quantized(q_a_quant)
else:
q_a_deq = q_a
q = pl['q_b'](q_a)
q = unweighted_rmsnorm(q).bfloat16()
q_heads = q.reshape(1, n_h, hd)
q_heads = _apply_rope(q_heads, dec_pos_per_gpu[gpu], rope_cos, rope_sin, rd)
# KV projection
kv = pl['kv'].run_from_quantized(x_quant_attn)
kv_norm_w = w.get(f"{pfx}.kv_norm.weight")
if kv_norm_w is not None:
kv = rmsnorm(kv, kv_norm_w.to(dev, torch.float32))
kv_3d = kv.reshape(1, 1, hd)
kv_3d = _apply_rope(kv_3d, dec_pos_per_gpu[gpu], rope_cos, rope_sin, rd)
kv_roped = kv_3d.reshape(1, hd)
# Write to output buffers (fixed-address copies for graph output)
self.q_heads_buf[li].copy_(q_heads)
self.kv_roped_buf[li].copy_(kv_roped)
self.q_a_buf[li].copy_(q_a_deq if q_norm_w is not None else q_a)
self.x_normed_buf[li].copy_(x_normed)
self.graphs_a[li] = graph_a
# ---- Sub-graph B: o_proj → mHC post(attn) → FFN → mHC post(FFN) ----
graph_b = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph_b):
X_l = self.x_in_a[li] # same X_l (both sub-graphs share this buffer)
attn_out = self.attn_out_buf[li] # set by eager section (FMHA + inv RoPE)
# Output projection: attn_out → o_a (grouped) → o_b
wo_a_lin = pl.get('o_a')
o_groups = cfg.get('o_groups', 16)
o_rank = cfg.get('o_lora_rank', 1024)
if wo_a_lin is not None:
g_3d = wo_a_lin.run(attn_out) # (1, n_groups, o_rank)
g_flat = g_3d.reshape(1, -1) # (1, n_groups * o_rank)
F_attn = pl['o_b'](g_flat)
else:
F_attn = torch.zeros(1, H, dtype=torch.bfloat16, device=dev)
# mHC post_block (attention)
A_l_a2, B_l_a2, C_l_a2 = attn_mhc._dynamic_params(X_l)
ctx_a = mHCContext(B_l=B_l_a2, C_l=C_l_a2)
X_mid = attn_mhc.post_block(X_l, F_attn, ctx_a)
# mHC pre_block (FFN) + rmsnorm + NVFP4 quantize
A_l_f, B_l_f, C_l_f = ffn_mhc._dynamic_params(X_mid)
x_quant_ffn = mhc_rmsnorm_quantize_nvfp4(
X_mid, A_l_f, ffn_norm_w.to(dev, torch.float32))
x_ffn = dequantize_nvfp4(
x_quant_ffn.x_fp4, x_quant_ffn.x_sf, x_quant_ffn.gsa)
# Router + MoE + SE
F_ffn = moe_forward(x_ffn, li, moe, se, router, dec_tid32_per_gpu[gpu])
# mHC post_block (FFN)
ctx_f = mHCContext(B_l=B_l_f, C_l=C_l_f)
X_next = ffn_mhc.post_block(X_mid, F_ffn, ctx_f)
self.x_next_buf[li].copy_(X_next)
self.graphs_b[li] = graph_b
if (li + 1) % 10 == 0:
print(f" Captured {li+1}/{self.n_layers} layer graphs (A+B)", flush=True)
# ---- Capture hc_head + norm + lm_head on cuda:0 ----
torch.cuda.set_device(0)
self.lm_graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(self.lm_graph):
x_out = hc_head.forward(self.x_lm_in) if hc_head is not None else self.x_lm_in[:, 0, :]
if final_norm_w is not None:
x_out = rmsnorm(x_out, final_norm_w)
logits = torch.nn.functional.linear(x_out, lm_w)
self.logits_buf.copy_(logits)
self.captured = True
print(f" Captured {len(self.graphs_a)} sub-graph A + {len(self.graphs_b)} sub-graph B + lm_head", flush=True)
# =====================================================================
def dequant_nvfp4(weight, weight_scale, weight_scale_2=None, input_scale=None):
O, I2 = weight.shape; I = I2 * 2
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"
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)]
lin.gs = [1.0] # base gs — finalize_weights will multiply by ws2
lin.ws2 = [ws2.to(d) if ws2 is not None else None]
# CRITICAL FIX: Compute gsa at RUNTIME from actual input magnitude.
# The checkpoint's input_scale is for training-time FP8 quantization.
# Using it as gsa causes E4M3 block scale overflow when x/gsa > 2688.
# We set a placeholder and override in the forward pass.
lin._activation_global_scale = 1.0 / (6.0 * 448.0) # placeholder
lin._use_runtime_gsa = True # flag to compute gsa at runtime
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._kv_bf16 = None # BF16 weight for kv_proj (dequantized from NVFP4)
self._gate_bf16 = None # BF16 weight for gate_proj (dequantized from NVFP4)
self.ape = None; self.kv_norm_w = None
self._reduce_loaded = False
# P7: Decode buffering — accumulate hidden_states until we have a complete block.
# HCA (r=128): skip GEMMs entirely at T=1 decode (n_complete=0 every time).
# CSA (r=4): buffer 4 decode steps, run GEMMs once per 4 tokens.
self._hs_buffer = None # (buf_len, H) BF16
self._pos_buffer = None # (buf_len,) long
self._buf_len = 0
def load(self, w, pfx, dev=None):
"""Load weights and build BF16 projections (dequantized from NVFP4)."""
if dev is None: dev = self.device
# Compressor projections are NOT explicitly FP4-QATed — dequant to BF16, use F.linear
# CRITICAL: Use the PyTorch dequant_nvfp4 (defined in this file), NOT the CUDA
# dequantize_nvfp4 from dsv4/ops/quantize.py. The CUDA kernel assumes
# activation/KV scale layout (row-major (M, N/16)) and crashes on weight scales
# that don't match — async illegal memory access surfaces at next sync.
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:
self._kv_bf16 = dequant_nvfp4(kv_w.to(dev), kv_ws.to(dev), kv_ws2, kv_isc).to(dev).contiguous()
if gate_w is not None:
self._gate_bf16 = dequant_nvfp4(gate_w.to(dev), gate_ws.to(dev), gate_ws2, gate_isc).to(dev).contiguous()
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_bf16 is None: return None, None, None
T = hidden_states.shape[0]; r = self.ratio; dev = hidden_states.device
# P7: Buffer decode steps until we have a complete block.
# For HCA (r=128) at T=1 decode: n_complete is always 0, so we skip
# the 2 NVFP4 GEMM launches entirely. No wasted compute.
# For CSA (r=4): accumulate 4 tokens, run GEMMs once.
if T < r:
# Buffer this token's hidden_states + position
if self._hs_buffer is None:
self._hs_buffer = torch.zeros(r, self.H, dtype=torch.bfloat16, device=dev)
self._pos_buffer = torch.zeros(r, dtype=torch.long, device=dev)
if self._buf_len < r:
self._hs_buffer[self._buf_len] = hidden_states[0] if T == 1 else hidden_states[self._buf_len]
self._pos_buffer[self._buf_len] = positions[0] if positions.numel() == 1 else positions[self._buf_len]
self._buf_len += 1
if self._buf_len < r:
return None, None, None # Not enough tokens yet
# We have a full buffer — use it
hidden_states = self._hs_buffer[:self._buf_len]
positions = self._pos_buffer[:self._buf_len]
T = self._buf_len
self._buf_len = 0 # Reset for next block
n_complete = T // r
if n_complete == 0: return None, None, None
# Step 1-2: BF16 F.linear projections → FP32 for compress
kv = torch.nn.functional.linear(hidden_states, self._kv_bf16).float() # (T, kv_dim) FP32
gate = torch.nn.functional.linear(hidden_states, self._gate_bf16).float() # (T, kv_dim) FP32
# Step 3: CUDA softmax/reduce kernel → FP32
# KV-1/KV-2: Return FP32. Caller applies RoPE, then quantizes to NVFP4.
from dsv4.kernels.compressor.production_compress import csa_compress_production_fp32, hca_compress_production_fp32
if self.is_csa:
compressed = csa_compress_production_fp32(
kv, gate, self.ape, self.kv_norm_w, m=r)
else:
compressed = hca_compress_production_fp32(
kv, gate, self.ape, self.kv_norm_w, m=r)
if compressed.shape[0] == 0: return None, None, None
n_comp = compressed.shape[0]
# Vectorized position computation — no Python loop, no .item()
# Block-aligned: use FIRST position of each block (vLLM cross-check confirmed)
# Wrong: ((bi+1)*r - 1) uses LAST position → off by r-1 (3 for CSA, 127 for HCA)
bi = torch.arange(n_comp, device=dev)
pos_idx = (bi * r).clamp(max=positions.numel() - 1)
comp_pos = positions[pos_idx]
# Return FP32 compressed output — caller handles RoPE + NVFP4 quantize
return compressed, comp_pos, torch.zeros(1, T, n_comp, 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 (FP4-QATed)
self._wp_bf16 = None # BF16 weight for weights_proj (dequantized from NVFP4)
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')
# q_b_proj IS the FP4-QATed QK path — keep as NVFP4
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')
# weights_proj is NOT FP4-QATed — dequant to BF16 via PyTorch reference
# CRITICAL: Use PyTorch dequant_nvfp4, NOT CUDA dequantize_nvfp4 (see Compressor.load)
if wp_w is not None:
self._wp_bf16 = dequant_nvfp4(wp_w.to(dev), wp_ws.to(dev), wp_ws2, wp_isc).to(dev).contiguous()
# Indexer compressor weights are directly under the indexer prefix
# (e.g. *.indexer.kv_proj.weight), NOT nested under *.indexer.compressor.
if f"{pfx}.kv_proj.weight" in w:
self.compressor = Compressor(4, self.ihd, 7168, dev)
self.compressor.load(w, pfx, dev)
def forward(self, q_lora, hidden_states, kv_cache, positions, layer_idx=None):
"""B2 FP8 tensor-core indexer scoring + weighted ReLU + top-k.
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. FP8 GEMM + ReLU + weighted sum + top-k (CUDA kernel)
Indexer keys are consumed directly in FP8_E4M3 format — no BF16 dequant.
"""
if self.q_b_lin is None or kv_cache is None or not kv_cache._has_idx or kv_cache.n_comp == 0:
return None
dev = q_lora.device; T = q_lora.shape[0]
li = layer_idx
q_idx = self.q_b_lin(q_lora).reshape(T, self.n_ih, self.ihd) # (T, n_ih, ihd)
w_h = torch.nn.functional.linear(hidden_states, self._wp_bf16) # (T, n_ih) BF16
# B2: FP8 tensor-core scoring path.
# Indexer keys are stored as FP8_E4M3 in the KV cache.
# No BF16 dequantization — the CUDA kernel consumes FP8 directly.
k_fp8 = kv_cache.comp_idx_fp8[:kv_cache.n_comp] # (n_comp, ihd) uint8
k_scale = kv_cache.comp_idx_scale[:kv_cache.n_comp] # (n_comp,) FP32
n_comp = kv_cache.n_comp
if li == 0:
print(f"\n=== INDEXER PROBE L0 (B2 FP8) ===", flush=True)
print(f" q_idx: shape={tuple(q_idx.shape)} dtype={q_idx.dtype}", flush=True)
print(f" k_fp8: shape={tuple(k_fp8.shape)} dtype={k_fp8.dtype}", flush=True)
print(f" k_scale: shape={tuple(k_scale.shape)} dtype={k_scale.dtype}", flush=True)
print(f" w_h: shape={tuple(w_h.shape)} dtype={w_h.dtype}", flush=True)
# For T=1 decode: use the B2 FP8 CUDA kernel
if T == 1 and self.ihd == 128 and self.n_ih == 64:
from dsv4.kernels.cuda.loader import get_cuda_module
mod = get_cuda_module("indexer_fp8_score_topk", ["indexer_fp8_score_topk.cu"],
extra_cuda_cflags=[
"-gencode=arch=compute_100a,code=sm_100a",
"-O3", "--use_fast_math", "--expt-relaxed-constexpr",
])
q_2d = q_idx.squeeze(0).contiguous() # (n_ih, ihd) BF16
w_1d = w_h.squeeze(0).contiguous() # (n_ih,) BF16
tk = min(self.top_k, n_comp)
topk_indices = torch.empty(tk, dtype=torch.int32, device=dev)
mod.indexer_fp8_score_topk(
q_2d, k_fp8, k_scale, w_1d, topk_indices,
self.n_ih, self.ihd, tk)
return topk_indices.unsqueeze(0) # (1, top_k)
# Fallback for T>1 or non-standard dimensions — FP32 einsum
k_idx = k_fp8 # still FP8, need dequant for einsum
if k_idx.dtype == torch.uint8 or str(k_idx.dtype) == 'torch.float8_e4m3fn':
from dsv4.kernels.cuda.loader import get_cuda_module
kv_mod = get_cuda_module("kv_quantize", ["kv_quantize.cu"])
k_idx = kv_mod.dequant_fp8_e4m3(k_fp8, k_scale) # (n_comp, ihd) BF16
scores = torch.einsum('tnd,cd->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:
"""KV Cache with mixed-precision compressed KV (DeepSeek V4 paper format).
KV-1/KV-2: Compressed KV uses mixed storage:
- Non-RoPE dims (448 of 512): FP8_E4M3 → ~50% size reduction
- RoPE dims (64 of 512): BF16 (RoPE applied directly, stored as BF16)
KV-3: Indexer keys stored as FP8_E4M3 (ihd=128, no RoPE).
SWA: BF16 (128 tokens × 512 × 61 layers = 8MB, fits in L2).
This matches the DeepSeek V4 paper: "BF16 for RoPE dims, FP8 for remaining dims.
This hybrid representation reduces the KV cache size by nearly half."
WHY NOT NVFP4 (native Blackwell FP4)?
─────────────────────────────────────
We *really* wanted to use NVFP4 (E2M1 + E4M3 block scales + FP32 global scale)
for compressed KV storage. Blackwell's native FP4→MMA path would have given us
3.5× memory savings and direct tensor-core consumption — the dream pipeline.
We tried. Hard. Three separate approaches:
1. Fused compressor_reduce_quant.cu — single-kernel compress→NVFP4. Bugs in
cross-warp block amax reduction and shared memory corruption (s_scratch
stomping adjacent variables). Best cos=0.703. Dead.
2. Proven two-kernel path (amax_gsa → quantize_from_buffer) using kv_quantize.cu's
compute_amax_gsa_fp32 + quantize_nvfp4_from_fp32. cos=0.995 on random data,
but that's the *quantize/dequant* round-trip in isolation. In the full pipeline,
the 4-bit precision on 448 non-RoPE dimensions accumulated error across 61 layers
of mHC — residual |X| already grows to 300-500, and NVFP4's 16-element block
quantization (4.5 bits effective) added ~0.5% per layer on top of that.
3. FP32 RoPE kernel (rope_fp32 in kv_quantize.cu) to avoid BF16 RoPE intermediate.
Had an indexing bug (cos=0.977 for M>1). Fixed but the real issue was NVFP4,
not RoPE.
The verdict: NVFP4's 4.5 effective bits per element is simply too coarse for
compressed KV values that get summed in attention softmax. FP8_E4M3's 5.3 effective
bits gives cos=0.9997 round-trip (vs NVFP4's 0.995) — that 0.4% difference compounds
fatally across 61 layers.
So we settled on FP8_E4M3 for non-RoPE + BF16 for RoPE — exactly what DeepSeek V4
ships in production. Not because we couldn't build the NVFP4 path (we did, it compiled
and ran), but because the math didn't hold up. Sometimes 4 bits isn't enough.
If Blackwell adds a finer-grained FP4 variant (8-element blocks, 6 effective bits),
revisit this. The kernels exist. The quantize/dequant path is proven. The precision
just isn't there yet for attention-sensitive KV values.
Storage per compressed entry at hd=512:
nope (448) × FP8 = 448 bytes + 4 bytes (scale) = 452
rope (64) × BF16 = 128 bytes
Total = 580 bytes vs 1024 bytes BF16 → 1.76× savings
"""
def __init__(self, head_dim, window_size=128, max_comp=65536, device='cuda:0',
indexer_key_dim=128, compress_ratio=4, indexer_top_k=1024, rope_dim=64):
self.hd, self.ws, self.dev = head_dim, window_size, device
self.idx_key_dim = indexer_key_dim
self.ratio = compress_ratio
self.max_comp = max_comp
self.rope_dim = rope_dim
self.nope_dim = head_dim - rope_dim # 448
# SWA: BF16 (small, fits in L2)
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
# Compressed KV: mixed FP8 (nope) + BF16 (rope)
self.comp_nope_fp8 = torch.zeros(max_comp, self.nope_dim, dtype=torch.uint8, device=device)
self.comp_nope_scale = torch.zeros(max_comp, dtype=torch.float32, device=device)
self.comp_rope_bf16 = torch.zeros(max_comp, rope_dim, dtype=torch.bfloat16, device=device)
self.comp_pos_buf = torch.zeros(max_comp, dtype=torch.long, device=device)
# Indexer compressed keys: FP8_E4M3
self.comp_idx_fp8 = torch.zeros(max_comp, indexer_key_dim, dtype=torch.uint8, device=device)
self.comp_idx_scale = torch.zeros(max_comp, dtype=torch.float32, device=device)
# Pre-allocated mixed gather buffers.
# CSA needs top_k + SWA; HCA is dense over compressed blocks, so it needs
# max_comp + SWA. These buffers preserve the paper/native storage layout:
# noPE stays FP8_E4M3 + scale, RoPE stays BF16.
if compress_ratio > 4:
self.mixed_gather_cap = max_comp + window_size
elif compress_ratio == 4:
self.mixed_gather_cap = indexer_top_k + window_size
else:
self.mixed_gather_cap = window_size
self.gather_nope_fp8 = torch.zeros(self.mixed_gather_cap, self.nope_dim, dtype=torch.uint8, device=device)
self.gather_nope_scale = torch.zeros(self.mixed_gather_cap, dtype=torch.float32, device=device)
self.gather_rope_bf16 = torch.zeros(self.mixed_gather_cap, rope_dim, dtype=torch.bfloat16, device=device)
# Legacy BF16 gather buffer kept only for non-B1 experiments; the live
# B1 path below does not materialize noPE KV as global BF16.
self.gather_buf = torch.zeros(indexer_top_k + window_size, head_dim, dtype=torch.bfloat16, device=device)
self.n_comp = 0
self._has_idx = False
# Cache extension modules (loaded once)
self._kv_quant_mod = None
self._fp8_attn_io_mod = None
def _get_kv_quant_mod(self):
if self._kv_quant_mod is None:
from dsv4.kernels.cuda.loader import get_cuda_module
self._kv_quant_mod = get_cuda_module("kv_quantize", ["kv_quantize.cu"])
return self._kv_quant_mod
def _get_fp8_attn_io_mod(self):
if self._fp8_attn_io_mod is None:
from dsv4.kernels.cuda.loader import get_cuda_module
self._fp8_attn_io_mod = get_cuda_module(
"fp8_attention_io", ["fp8_attention_io.cu"],
extra_cuda_cflags=[
"-gencode=arch=compute_100a,code=sm_100a",
"-O3", "--use_fast_math", "--expt-relaxed-constexpr",
],
)
return self._fp8_attn_io_mod
def append_swa(self, kv, pos):
"""Vectorized SWA append — 2 kernel launches instead of 2T."""
T = kv.shape[0]
idx = (self.swa_head + torch.arange(T, device=self.dev)) % self.ws
self.swa.index_copy_(0, idx, kv)
self.swa_pos.index_copy_(0, idx, pos)
self.swa_head = (self.swa_head + T) % self.ws
self.swa_len = min(self.swa_len + T, self.ws)
def set_compressed_mixed(self, nope_fp8, nope_scale, rope_bf16, comp_pos=None):
"""Add compressed KV entries (mixed FP8 nope + BF16 rope)."""
T = nope_fp8.shape[0]
end = self.n_comp
self.comp_nope_fp8[end:end+T] = nope_fp8.view(torch.uint8) if nope_fp8.dtype != torch.uint8 else nope_fp8
self.comp_nope_scale[end:end+T] = nope_scale
self.comp_rope_bf16[end:end+T] = rope_bf16
if comp_pos is not None:
self.comp_pos_buf[end:end+T] = comp_pos
self.n_comp = end + T
def set_indexer_keys_fp8(self, idx_kv):
"""Add indexer compressed keys. idx_kv is BF16 (n_comp, ihd) or FP8 (fp8, scale)."""
if idx_kv is None: return
T = idx_kv[0].shape[0] if isinstance(idx_kv, tuple) else idx_kv.shape[0]
end = self.n_comp - T
if isinstance(idx_kv, tuple) and len(idx_kv) == 2:
fp8, scale = idx_kv
self.comp_idx_fp8[end:end+T] = fp8.view(torch.uint8) if fp8.dtype != torch.uint8 else fp8
self.comp_idx_scale[end:end+T] = scale
elif isinstance(idx_kv, torch.Tensor):
mod = self._get_kv_quant_mod()
fp8, scale = mod.quantize_fp8_e4m3_from_fp32(idx_kv.float().contiguous())
self.comp_idx_fp8[end:end+T] = fp8.view(torch.uint8)
self.comp_idx_scale[end:end+T] = scale
self._has_idx = True
def comp_nope_selective(self, indices):
"""Dequant FP8 nope for selected entries → BF16."""
mod = self._get_kv_quant_mod()
return mod.dequant_fp8_e4m3_selective(
self.comp_nope_fp8, self.comp_nope_scale, indices.int())
def comp_rope_selective(self, indices):
"""Gather BF16 rope for selected entries."""
return self.comp_rope_bf16[indices.long()]
@property
def comp_nope_all(self):
"""Dequant all FP8 nope → BF16."""
mod = self._get_kv_quant_mod()
return mod.dequant_fp8_e4m3(
self.comp_nope_fp8[:self.n_comp],
self.comp_nope_scale[:self.n_comp])
@property
def comp_rope_all(self):
"""Return all BF16 rope entries."""
return self.comp_rope_bf16[:self.n_comp]
@property
def comp_pos(self):
return self.comp_pos_buf[:self.n_comp] if self.n_comp > 0 else None
@property
def comp_idx_kv(self):
"""Dequant FP8 indexer keys → BF16 for scoring."""
if not self._has_idx or self.n_comp == 0: return None
mod = self._get_kv_quant_mod()
return mod.dequant_fp8_e4m3(
self.comp_idx_fp8[:self.n_comp],
self.comp_idx_scale[:self.n_comp])
def gather_mixed_selective(self, indices):
"""Gather selected compressed KV + SWA into mixed FP8/BF16 buffers.
Returns (nope_fp8, nope_scale, rope_bf16), each sliced to total length.
noPE is not dequantized to global BF16.
"""
mod = self._get_fp8_attn_io_mod()
swa_kv, _ = self.get_swa()
idx = indices.int().contiguous()
total = idx.numel() + swa_kv.shape[0]
if total > self.mixed_gather_cap:
raise RuntimeError(f"mixed gather capacity {self.mixed_gather_cap} < requested {total}")
mod.gather_mixed_selective_(
self.comp_nope_fp8, self.comp_nope_scale, self.comp_rope_bf16,
swa_kv, idx, self.gather_nope_fp8, self.gather_nope_scale, self.gather_rope_bf16)
return (self.gather_nope_fp8[:total],
self.gather_nope_scale[:total],
self.gather_rope_bf16[:total])
def gather_mixed_all(self):
"""Gather all compressed KV + SWA in mixed FP8/BF16 storage for HCA."""
mod = self._get_fp8_attn_io_mod()
swa_kv, _ = self.get_swa()
n_comp = int(self.n_comp)
total = n_comp + swa_kv.shape[0]
if total > self.mixed_gather_cap:
raise RuntimeError(f"mixed gather capacity {self.mixed_gather_cap} < requested {total}")
mod.gather_mixed_all_(
self.comp_nope_fp8[:n_comp], self.comp_nope_scale[:n_comp], self.comp_rope_bf16[:n_comp],
swa_kv, self.gather_nope_fp8, self.gather_nope_scale, self.gather_rope_bf16)
return (self.gather_nope_fp8[:total],
self.gather_nope_scale[:total],
self.gather_rope_bf16[:total])
def gather_mixed_swa_only(self):
"""Quantize SWA noPE tail to FP8 and keep SWA RoPE as BF16."""
mod = self._get_fp8_attn_io_mod()
swa_kv, _ = self.get_swa()
total = swa_kv.shape[0]
if total > self.mixed_gather_cap:
raise RuntimeError(f"mixed gather capacity {self.mixed_gather_cap} < requested {total}")
mod.gather_mixed_swa_only_(
swa_kv, self.gather_nope_fp8, self.gather_nope_scale, self.gather_rope_bf16, self.rope_dim)
return (self.gather_nope_fp8[:total],
self.gather_nope_scale[:total],
self.gather_rope_bf16[:total])
def get_swa(self):
"""Return SWA KV and positions as views (no clone)."""
if self.swa_len == 0:
return self.swa[:0], self.swa_pos[:0]
if self.swa_len < self.ws:
return self.swa[:self.swa_len], self.swa_pos[:self.swa_len]
idx = torch.arange(self.swa_head, self.swa_head + self.ws, device=self.dev) % self.ws
return self.swa[idx], self.swa_pos[idx]
# =====================================================================
# 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
# K and V are the same in MQA — V = K transposed to (hd, N) format.
# .transpose(-1,-2).contiguous() creates a new tensor (no clone needed).
# This saves one full KV copy (~256KB per layer per decode step).
v = k
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)
def _run_production_fmha_mixed(q_heads, kv_nope_fp8, kv_nope_scale, kv_rope_bf16,
n_h, hd, T, seq_len, scale, dev, li, w, pfx, rope_dim):
"""B1 storage-native mixed FP8/BF16 FMHA. Supports decode (T=1) and prefill (T>1)."""
from dsv4.kernels.attention.production import dsv4_attention_mixed_fp8_decode, dsv4_attention_mixed_fp8_prefill
q = q_heads.permute(1, 0, 2).contiguous() # (n_h, T, hd)
sinks = w.get(f"{pfx}.sinks"); sink_bias = None
if sinks is not None:
sink_bias = sinks.to(device=dev).float().reshape(n_h)
if T == 1:
attn_out = dsv4_attention_mixed_fp8_decode(
q=q,
k_nope_fp8=kv_nope_fp8,
k_nope_scale=kv_nope_scale,
k_rope_bf16=kv_rope_bf16,
scale=scale,
sink_bias=sink_bias,
rope_dim=rope_dim,
)
else:
attn_out = dsv4_attention_mixed_fp8_prefill(
q=q,
k_nope_fp8=kv_nope_fp8,
k_nope_scale=kv_nope_scale,
k_rope_bf16=kv_rope_bf16,
scale=scale,
sink_bias=sink_bias,
rope_dim=rope_dim,
)
return attn_out.permute(1, 0, 2) # (T, n_h, hd)
# =====================================================================
# Attention — ALL production kernels
# =====================================================================
def eager_attention(q_heads, kv_roped, x_normed, q_a, w, li, cfg,
rope_cos, rope_sin, kv_cache, positions,
compressor, indexer, comp_rope_cos=None, comp_rope_sin=None):
"""Eager attention section — runs OUTSIDE CUDA graph capture.
This function handles the dynamic-shape parts of attention:
KV append → Compressor → Indexer → KV gather → FMHA → Inverse RoPE
Returns: attn_out (1, n_h, hd) — output of FMHA after inverse RoPE.
The caller (sub-graph B) will apply o_proj and mHC post_block.
"""
dev = x_normed.device; T = q_heads.shape[0]
n_h = cfg["num_attention_heads"]; hd = cfg["head_dim"]; rd = cfg.get("qk_rope_head_dim", 64)
ratio = compressor.ratio if compressor is not None else 0
scale = 1.0 / math.sqrt(hd); pfx = f"model.layers.{li}.self_attn"
nope_dim = hd - rd
if positions.device != rope_cos.device: positions = positions.to(rope_cos.device)
# KV append (already roped from sub-graph A)
kv_cache.append_swa(kv_roped, positions)
# Compressor → compressed KV (mixed storage: FP8 + BF16 RoPE)
comp_pos, block_bias = None, None; comp_idx_kv = None
if compressor is not None and compressor.ratio > 0:
comp_kv_fp32, comp_pos, block_bias = compressor.forward(x_normed, positions)
if comp_kv_fp32 is not None:
from dsv4.kernels.cuda.loader import get_cuda_module
kv_mod = get_cuda_module("kv_quantize", ["kv_quantize.cu"])
nope_fp32 = comp_kv_fp32[:, :nope_dim].contiguous()
rope_bf16 = comp_kv_fp32[:, nope_dim:].bfloat16().contiguous()
rope_3d = rope_bf16.unsqueeze(1)
crc = comp_rope_cos if comp_rope_cos is not None else rope_cos
crs = comp_rope_sin if comp_rope_sin is not None else rope_sin
rope_3d = _apply_rope(rope_3d, comp_pos, crc, crs, rd)
rope_bf16 = rope_3d.squeeze(1)
nope_fp8, nope_scale = kv_mod.quantize_fp8_e4m3_from_fp32(nope_fp32)
kv_cache.set_compressed_mixed(nope_fp8, nope_scale, rope_bf16, comp_pos)
if compressor.is_csa and indexer is not None and indexer.compressor is not None:
comp_idx_kv, _, _ = indexer.compressor.forward(x_normed, positions)
kv_cache.set_indexer_keys_fp8(comp_idx_kv)
# Indexer top-k (CSA)
topk_idx = None
if indexer is not None and ratio == 4:
topk_idx = indexer.forward(q_a, x_normed, kv_cache, positions, layer_idx=li)
# Gather KV — B1 storage-native mixed path
swa_kv, _swa_pos = kv_cache.get_swa()
swa_len = swa_kv.shape[0]
if kv_cache.n_comp > 0:
if ratio == 4:
assert topk_idx is not None, f"CSA layer {li}: indexer returned no top-k"
tk = topk_idx[0].clamp(0, kv_cache.n_comp - 1).int()
kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_selective(tk)
elif ratio > 4:
kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_all()
else:
kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_swa_only()
else:
kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_swa_only()
seq_len = kv_nope_scale.shape[0]
if seq_len == 0:
return torch.zeros(T, n_h, hd, dtype=torch.bfloat16, device=dev)
# Production FMHA — B1 mixed FP8/BF16 decode path
attn_out = _run_production_fmha_mixed(
q_heads, kv_nope_fp8, kv_nope_scale, kv_rope_bf16,
n_h, hd, T, seq_len, scale, dev, li, w, pfx, rd)
# Inverse RoPE
attn_out = _apply_rope(attn_out, positions, rope_cos, rope_sin, rd, inverse=True)
return attn_out
def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
kv_cache, positions, compressor, indexer, prod_lin,
x_quant=None,
_profile_detail=False, _profile_times=None,
comp_rope_cos=None, comp_rope_sin=None):
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"
nope_dim = hd - rd # 448 — used by both compress and gather
if positions.device != rope_cos.device: positions = positions.to(rope_cos.device)
def _pt(tag):
"""Profile timing helper — records CUDA-sync'd timestamp."""
if _profile_detail and _profile_times is not None:
torch.cuda.synchronize()
_profile_times.append((tag, li, time.perf_counter()))
_pt('q_a_start')
# 1. Q: q_a (NVFP4 GEMM) → q_a_norm → q_b (NVFP4 GEMM) → q_b_norm
q_a = prod_lin['q_a'].run_from_quantized(x_quant) if x_quant is not None else prod_lin['q_a'](x_normed)
_pt('q_a_end')
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")
# B3: Fused rmsnorm+quant for q_a_norm → q_b path
# Replaces: rmsnorm(q_a, w) → BF16 → q_b quantizes internally
# With: fused rmsnorm+NVFP4 quantize → QuantizedActivation → q_b.run_from_quantized
# Saves: ~6 kernel launches per layer (rmsnorm 4+ + quantize 2 vs fused 2)
if q_norm_w is not None:
from dsv4.ops.quantize import rmsnorm_quantize_nvfp4 as _rmsnorm_quantize, dequantize_nvfp4 as _dequantize_nvfp4
q_a_quant = _rmsnorm_quantize(q_a, q_norm_w.to(dev, torch.float32))
q_a = _dequantize_nvfp4(q_a_quant.x_fp4, q_a_quant.x_sf, q_a_quant.gsa)
_pt('q_b_start')
if q_norm_w is not None:
q = prod_lin['q_b'].run_from_quantized(q_a_quant)
else:
q = prod_lin['q_b'](q_a)
q = unweighted_rmsnorm(q).bfloat16()
_pt('q_b_end')
q_heads = q.reshape(T, n_h, hd); q_heads = _apply_rope(q_heads, positions, rope_cos, rope_sin, rd)
_pt('rope_q_end')
# 2. KV (NVFP4 GEMM, MQA, single KV head)
_pt('kv_start')
kv = prod_lin['kv'].run_from_quantized(x_quant) if x_quant is not None else prod_lin['kv'](x_normed)
_pt('kv_end')
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)
_pt('rope_kv_end')
kv_roped = kv_3d.reshape(T, hd); kv_cache.append_swa(kv_roped, positions)
# 3. Compressor → compressed KV (mixed storage: FP8 + BF16 RoPE)
# DeepSeek V4 paper: "BF16 for RoPE dims, FP8 for remaining dims"
_pt('compress_start')
comp_pos, block_bias = None, None; comp_idx_kv = None
if compressor is not None and compressor.ratio > 0:
comp_kv_fp32, comp_pos, block_bias = compressor.forward(x_normed, positions)
if comp_kv_fp32 is not None:
from dsv4.kernels.cuda.loader import get_cuda_module
kv_mod = get_cuda_module("kv_quantize", ["kv_quantize.cu"])
# Split into non-RoPE (FP8) and RoPE (BF16) parts
nope_fp32 = comp_kv_fp32[:, :nope_dim].contiguous() # (n_comp, 448) FP32
rope_bf16 = comp_kv_fp32[:, nope_dim:].bfloat16().contiguous() # (n_comp, 64) BF16
# Apply RoPE on BF16 rope dims (existing BF16 RoPE kernel)
rope_3d = rope_bf16.unsqueeze(1) # (n_comp, 1, 64)
# Use compress_rope_theta cache for compressed entries if available
crc = comp_rope_cos if comp_rope_cos is not None else rope_cos
crs = comp_rope_sin if comp_rope_sin is not None else rope_sin
rope_3d = _apply_rope(rope_3d, comp_pos, crc, crs, rd)
rope_bf16 = rope_3d.squeeze(1) # (n_comp, 64) BF16
# Quantize non-RoPE part FP32 → FP8_E4M3
nope_fp8, nope_scale = kv_mod.quantize_fp8_e4m3_from_fp32(nope_fp32)
# Store mixed-format compressed KV + positions
kv_cache.set_compressed_mixed(nope_fp8, nope_scale, rope_bf16, comp_pos)
if compressor.is_csa and indexer is not None and indexer.compressor is not None:
comp_idx_kv, _, _ = indexer.compressor.forward(x_normed, positions)
# Indexer keys: FP8_E4M3 (ihd=128, no RoPE)
kv_cache.set_indexer_keys_fp8(comp_idx_kv)
_pt('compress_end')
# 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, positions, layer_idx=li)
# 5. Gather KV — B1 storage-native mixed path.
# noPE remains FP8_E4M3 + per-row scale; RoPE remains BF16.
# There is no global FP8->BF16 noPE materialization before FMHA.
_pt('gather_start')
swa_kv, _swa_pos = kv_cache.get_swa()
swa_len = swa_kv.shape[0]
if kv_cache.n_comp > 0:
if ratio == 4:
# CSA: gather top-k compressed rows + SWA tail without dequantizing noPE.
assert topk_idx is not None, f"CSA layer {li}: indexer returned no top-k — indexer is broken"
tk = topk_idx[0].clamp(0, kv_cache.n_comp - 1).int()
kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_selective(tk)
elif ratio > 4:
# HCA: dense over compressed rows, still mixed storage.
kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_all()
else:
kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_swa_only()
else:
kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_swa_only()
seq_len = kv_nope_scale.shape[0]
if seq_len == 0: return torch.zeros(T, cfg["hidden_size"], dtype=torch.bfloat16, device=dev), q_a
# 6. Production FMHA — B1 mixed FP8/BF16 decode path.
_pt('fmha_start')
if li == 0:
if VERBOSE >= 2:
print(f" L0 B1 verify: kv_nope_fp8 dtype={kv_nope_fp8.dtype} shape={tuple(kv_nope_fp8.shape)} "
f"kv_nope_scale dtype={kv_nope_scale.dtype} shape={tuple(kv_nope_scale.shape)} "
f"kv_rope_bf16 dtype={kv_rope_bf16.dtype} shape={tuple(kv_rope_bf16.shape)}", flush=True)
assert kv_nope_fp8.dtype in (torch.uint8, torch.float8_e4m3fn), f"kv_nope_fp8 wrong dtype: {kv_nope_fp8.dtype}"
assert kv_nope_scale.dtype == torch.float32, f"kv_nope_scale wrong dtype: {kv_nope_scale.dtype}"
assert kv_rope_bf16.dtype == torch.bfloat16, f"kv_rope_bf16 wrong dtype: {kv_rope_bf16.dtype}"
assert kv_nope_fp8.shape[-1] == nope_dim, f"kv_nope_fp8 dim={kv_nope_fp8.shape[-1]} != nope_dim={nope_dim}"
assert kv_rope_bf16.shape[-1] == rd, f"kv_rope_bf16 dim={kv_rope_bf16.shape[-1]} != rope_dim={rd}"
if VERBOSE >= 2 and li < 3:
print(f" L{li} FMHA mixed input: T={T} seq_len={seq_len} hd={hd} n_h={n_h} n_comp={kv_cache.n_comp} swa_len={swa_len}", flush=True)
attn_out = _run_production_fmha_mixed(
q_heads, kv_nope_fp8, kv_nope_scale, kv_rope_bf16,
n_h, hd, T, seq_len, scale, dev, li, w, pfx, rd)
_pt('fmha_end')
if VERBOSE >= 2 and li < 3:
print(f" L{li} FMHA mixed: |prod|={attn_out.abs().max().item():.6f} (reference disabled: B1 forbids global BF16 KV staging)", flush=True)
# 7. Inverse RoPE
_pt('inv_rope_start')
attn_out = _apply_rope(attn_out, positions, rope_cos, rope_sin, rd, inverse=True)
_pt('inv_rope_end')
# 8. Output: wo_a (NVFP4 grouped GEMM) + wo_b (NVFP4 GEMM)
_pt('o_proj_start')
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)
_pt('o_proj_end')
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)
# DEBUG: check topk_ids validity (only for first 3 and last 3 layers)
if VERBOSE >= 2 and (li < 3 or li >= 58):
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 >= 58:
print(f" L{li} MoE DIAG: topk_ids={topk_ids[0].tolist()} topk_w=[{','.join(f'{w:.3f}' for w in topk_w[0].tolist())}]", flush=True)
# Also print gate logits for debugging
if hasattr(router, '_gate_lin') and router._gate_lin is not None:
gate_logits = router._gate_lin(x).float()
print(f" L{li} gate logits: [{gate_logits.min().item():.3f}, {gate_logits.max().item():.3f}] mean={gate_logits.mean().item():.3f}", 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)
shared_out = se_runner.run(x)
if VERBOSE >= 2 and li >= 58:
print(f" L{li} MoE DIAG: |routed|={routed_out.abs().max().item():.1f} |shared|={shared_out.abs().max().item():.1f} |x|={x.abs().max().item():.1f}", flush=True)
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, _profile_detail=False, _profile_times=None,
_use_fused_rmsnorm_quantize=True,
comp_rope_cos=None, comp_rope_sin=None,
):
"""Forward one transformer layer.
"""
# P5: Fused mHC pre_block + RMSNorm + NVFP4 quantize
# Replaces: pre_block (bmm) + rmsnorm (~4 launches) + quantize (2 launches)
# With: 2 kernel launches total (mhc_rmsnorm_amax_gsa + mhc_rmsnorm_quantize_nvfp4)
# Savings: ~5 launches per site × 2 sites × 61 layers = 610 launches/token
from dsv4.ops.quantize import (
rmsnorm_quantize_nvfp4, mhc_rmsnorm_quantize_nvfp4,
QuantizedActivation, dequantize_nvfp4,
)
from dsv4.layers.mhc import mHCContext
# Attention mHC: fused pre_block + rmsnorm + NVFP4 quantize
A_l_a, B_l_a, C_l_a = attn_mhc._dynamic_params(X_l)
ctx_a = mHCContext(B_l=B_l_a, C_l=C_l_a)
if _use_fused_rmsnorm_quantize:
# P5 fused: X_l + A_l → bmm + rmsnorm + NVFP4 quantize in 2 kernel launches
x_quant_attn = mhc_rmsnorm_quantize_nvfp4(
X_l, A_l_a, attn_norm_w.to(X_l.device, torch.float32))
# Dequantize for compressor/indexer (1 kernel launch)
x_normed = dequantize_nvfp4(x_quant_attn.x_fp4, x_quant_attn.x_sf, x_quant_attn.gsa)
else:
x_in = torch.bmm(A_l_a.unsqueeze(1).float(), X_l.float()).squeeze(1).bfloat16()
x_normed = rmsnorm(x_in, attn_norm_w)
x_quant_attn = None
if _profile_detail: torch.cuda.synchronize(); t_attn0 = time.perf_counter()
F_attn, _ = forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
kv_cache, positions, compressor, indexer, prod_lin,
x_quant=x_quant_attn,
_profile_detail=_profile_detail, _profile_times=_profile_times,
comp_rope_cos=comp_rope_cos, comp_rope_sin=comp_rope_sin)
if _profile_detail: torch.cuda.synchronize(); t_attn1 = time.perf_counter()
X_mid = attn_mhc.post_block(X_l, F_attn, ctx_a)
# FFN mHC: fused pre_block + rmsnorm + NVFP4 quantize
A_l_f, B_l_f, C_l_f = ffn_mhc._dynamic_params(X_mid)
ctx_f = mHCContext(B_l=B_l_f, C_l=C_l_f)
if _use_fused_rmsnorm_quantize:
# P5 fused: X_mid + A_l → bmm + rmsnorm + NVFP4 quantize in 2 kernel launches
x_quant_ffn = mhc_rmsnorm_quantize_nvfp4(
X_mid, A_l_f, ffn_norm_w.to(X_mid.device, torch.float32))
# Dequantize for MoE (BF16 input required by MoE quantize path)
x_ffn = dequantize_nvfp4(x_quant_ffn.x_fp4, x_quant_ffn.x_sf, x_quant_ffn.gsa)
else:
x_in_f = torch.bmm(A_l_f.unsqueeze(1).float(), X_mid.float()).squeeze(1).bfloat16()
x_ffn = rmsnorm(x_in_f, ffn_norm_w)
if _profile_detail: torch.cuda.synchronize(); t_ffn0 = time.perf_counter()
F_ffn = moe_forward(x_ffn, li, moe_runner, se_runner, router, token_id)
if _profile_detail: torch.cuda.synchronize(); t_ffn1 = time.perf_counter()
X_next = ffn_mhc.post_block(X_mid, F_ffn, ctx_f)
if VERBOSE >= 2 and (li < 3 or li >= 58):
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)
# Detailed diagnostics — only with VERBOSE >= 2 to avoid .item() syncs on hot path
if VERBOSE >= 2 and (li >= 58 or (li > 0 and X_next.abs().max().item() > 200)):
A_a, B_a, C_a = attn_mhc._dynamic_params(X_l)
A_f, B_f, C_f = ffn_mhc._dynamic_params(X_mid)
print(f" L{li} DIAG: A_attn=[{A_a.min().item():.4f},{A_a.max().item():.4f}] "
f"C_attn=[{C_a.min().item():.4f},{C_a.max().item():.4f}] "
f"A_ffn=[{A_f.min().item():.4f},{A_f.max().item():.4f}] "
f"C_ffn=[{C_f.min().item():.4f},{C_f.max().item():.4f}]", flush=True)
print(f" L{li} DIAG: B_attn row_sum=[{B_a.sum(-1).min().item():.4f},{B_a.sum(-1).max().item():.4f}] "
f"col_sum=[{B_a.sum(-2).min().item():.4f},{B_a.sum(-2).max().item():.4f}] "
f"B_ffn row_sum=[{B_f.sum(-1).min().item():.4f},{B_f.sum(-1).max().item():.4f}] "
f"col_sum=[{B_f.sum(-2).min().item():.4f},{B_f.sum(-2).max().item():.4f}]", flush=True)
print(f" L{li} DIAG: |x_in_attn|={x_in.abs().max().item():.1f} "
f"|x_in_ffn|={x_in_f.abs().max().item():.1f} "
f"|X_l|={X_l.abs().max().item():.1f} "
f"|X_mid|={X_mid.abs().max().item():.1f} "
f"|X_next|={X_next.abs().max().item():.1f}", flush=True)
if _profile_detail and (li < 3 or li == 30 or li >= 58):
torch.cuda.synchronize()
attn_ms = (t_attn1 - t_attn0) * 1000
ffn_ms = (t_ffn1 - t_ffn0) * 1000
print(f" L{li}: attn={attn_ms:.2f}ms ffn={ffn_ms:.2f}ms", 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 (no dequant/re-quant)
wo_a.load_nvfp4_weight(oa_w_nvfp4.to(dev), oa_ws.to(dev),
oa_ws2.to(dev) if oa_ws2 is not None else None,
oa_isc.to(dev) if oa_isc is not None else None)
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
wo_a._use_runtime_gsa = True # compute gsa from actual input to avoid E4M3 overflow
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")
# BF16 router gate — dequantize NVFP4 to BF16, use F.linear
E = cfg["n_routed_experts"]
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:
# Checkpoint has NVFP4 gate weight — dequantize to BF16
# CRITICAL: Use PyTorch dequant_nvfp4, NOT CUDA dequantize_nvfp4
# (same fix as Compressor.load — CUDA kernel crashes on weight scale layouts)
gate_bf16 = dequant_nvfp4(gate_w.to(dev), gate_ws.to(dev), gate_ws2, gate_isc)
router.W_gate = gate_bf16.T.contiguous().to(dev) # (H, E) for F.linear(x, W_gate.T)
else:
# BF16 gate weight from checkpoint
gw = all_w.get(f"{pfx}.gate.weight")
gate_bf16 = gw.bfloat16().to(dev)
if gate_bf16.shape[0] != H:
gate_bf16 = gate_bf16.T.contiguous() # ensure (H, E)
router.W_gate = gate_bf16.contiguous()
# No gate_lin — force BF16 dispatch path
router.gate_lin = None
router.load_weights(e_bias=eb.to(dev, torch.float32))
if li < 5: print(f" L{li}: BF16 router gate (dequantized from NVFP4)", flush=True)
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))
# P0: ENABLE fused SwiGLU — NVFP4 GEMM + SiLU in kernel registers.
# Saves 240+ unfused BF16 kernel launches per token (gate_silu, clamp, mul, quantize).
moe.set_fused_swiglu(True)
_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)
# FIX: Do NOT use checkpoint input_scale as gsa — causes E4M3 overflow.
# Instead, compute gsa at runtime from actual activation magnitude.
# The MoE runner's compute_activation_global_scales() does this correctly.
# We enable runtime gsa for both MoE and SharedExpert.
moe._use_runtime_gsa = True
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)
# P1: ENABLE fused SwiGLU for shared expert (1-group variant of MoE fused kernel)
se.set_fused_swiglu(True)
# EAGERLY process shared expert weights
se._ensure_initialized()
# P1: Eagerly warmup fused SwiGLU compilation for SE (1-group)
if se._fused_swiglu:
from dsv4.ops.gemm_runner import warmup_fused_swiglu_compilation
K_packed = H // 2
N_packed_l1 = (2 * cfg.get("moe_intermediate_size", 3072)) // 2 # gate+up
warmup_fused_swiglu_compilation(
1, K_packed, N_packed_l1, dev,
swiglu_limit=cfg.get("swiglu_limit", 10.0),
)
# Fix activation global scales — _ensure_initialized sets gsa from l1_gs (which is 1.0)
# FIX: Same runtime gsa for SharedExpert
se._use_runtime_gsa = True
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_head: BF16 GEMM (checkpoint weight is BF16, no quantization)
lm_w_raw = all_w.get("lm_head.weight", embed_w).bfloat16().to('cuda:0')
lm_head_lin = None # Use raw BF16 F.linear for lm_head
lm_w = lm_w_raw # Keep as (V, H) BF16 for F.linear
print(" lm_head: BF16 GEMM (checkpoint weight, no quantization)")
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(romax, rd, f"cuda:{g}", rtheta, rt, rf, romax, rbfast, rbslow) for g in range(NUM_GPUS)}
# Compressed-entry RoPE uses separate theta (vLLM cross-check: compress_rope_theta)
# If compress_rope_theta differs from rope_theta, compressed KV entries need their own cache
comp_rtheta = cfg.get("compress_rope_theta", rtheta)
if comp_rtheta != rtheta:
comp_rope_caches = {g: build_rope_cache(romax, rd, f"cuda:{g}", comp_rtheta, rt, rf, romax, rbfast, rbslow) for g in range(NUM_GPUS)}
print(f" Compressed RoPE theta: {comp_rtheta} (different from normal: {rtheta})")
else:
comp_rope_caches = rope_caches # Same theta, reuse normal cache
# 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)
max_ctx = _args.max_context
print(f" Max context: {max_ctx} tokens (governs KV cache pre-allocation)")
for li in range(n_layers):
dev = f"cuda:{li % NUM_GPUS}"; ratio = cr[li] if li < len(cr) else 128
# C1: max_comp derived from target context and compress ratio
max_comp = (max_ctx + ratio - 1) // ratio if ratio > 0 else 0
kv_caches[li] = KVCache(hd, cfg.get("sliding_window", 128), max_comp=max_comp, device=dev,
indexer_key_dim=ihd, compress_ratio=ratio, indexer_top_k=itk, rope_dim=rd)
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)
# Derive special token IDs from official encoding strings + tokenizer.
# This is the ONLY source of truth — never hardcode these IDs.
THINK_START = tokenizer.convert_tokens_to_ids(_THINK_START_STR)
THINK_END = tokenizer.convert_tokens_to_ids(_THINK_END_STR)
USER_TOKEN = tokenizer.convert_tokens_to_ids(_USER_STR)
ASSISTANT_TOKEN = tokenizer.convert_tokens_to_ids(_ASSISTANT_STR)
bos = tokenizer.bos_token_id or 0
# A1: Build explicit stop set — DSV4 uses special turn-end tokens beyond eos
STOP_IDS = set()
eos_id = tokenizer.eos_token_id
if eos_id is not None:
STOP_IDS.add(eos_id)
for tok_name in ("<|end_of_sentence|>",):
tid = tokenizer.convert_tokens_to_ids(tok_name)
if tid is not None and tid >= 0 and tid != tokenizer.unk_token_id:
STOP_IDS.add(tid)
# If model emits USER_TOKEN it's trying to open a new user turn = it's done
STOP_IDS.add(USER_TOKEN)
print(f" Stop set: {STOP_IDS} (eos={eos_id}, eos_token={tokenizer.eos_token})")
print(f" Special tokens: {tokenizer.special_tokens_map}")
print(f" THINK_START={THINK_START} THINK_END={THINK_END} USER={USER_TOKEN} ASST={ASSISTANT_TOKEN}")
if _args.prefill_tokens:
generated = [int(x) for x in _args.prefill_tokens.split(',')]
else:
# Official DeepSeek V4 encoding — canonical path, no hand-rolled alternatives.
# Uses encoding/deepseek_v4_encoding.py (copied from vLLM tree) to build
# the prompt. This is the ONLY way to construct prompts — the official
# encoder handles BOS, User/Assistant tokens, thinking mode, and all
# special token placement. It can't drift because it's the same code
# the inference engines will use.
from encoding.deepseek_v4_encoding import encode_messages
messages = [{"role": "user", "content": PROMPT}]
thinking_mode = _args.thinking_mode # 'thinking' or 'chat'
encoded_str = encode_messages(messages, thinking_mode=thinking_mode)
generated = tokenizer.encode(encoded_str, add_special_tokens=False)
# Ensure BOS token is present at the start
if generated[0] != bos:
generated = [bos] + generated
all_tokens = generated.copy()
print(f"Input: {len(generated)} tokens (thinking_mode={_args.thinking_mode})")
# Batched prefill — process tokens in chunks of up to 128 (FMHA T≤128 constraint)
PREFILL_CHUNK = 128 # max T per FMHA launch; split larger prefills into chunks
n_prefill = len(generated)
print(f"Batched prefill: {n_prefill} tokens, chunk_size={PREFILL_CHUNK}")
prefill_ids = torch.tensor(generated, dtype=torch.long, device='cuda:0')
prefill_ids32 = prefill_ids.to(torch.int32)
all_positions = torch.arange(n_prefill, dtype=torch.long, device='cuda:0')
# Process chunks: each chunk goes through ALL 61 layers before the next chunk.
# This ensures KV cache is populated correctly for each layer.
chunk_starts = list(range(0, n_prefill, PREFILL_CHUNK))
X = None # will be set by first chunk's embedding
for ci, cs in enumerate(chunk_starts):
ce = min(cs + PREFILL_CHUNK, n_prefill)
chunk_len = ce - cs
t1 = time.time()
# Embed chunk tokens: (chunk_len, d)
chunk_ids = prefill_ids[cs:ce]
chunk_ids32 = prefill_ids32[cs:ce]
chunk_positions = all_positions[cs:ce]
chunk_embed = embed(chunk_ids) # (chunk_len, d) BF16
X = mHCLayer.init_state(chunk_embed) # (chunk_len, n_hc, d) BF16
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], chunk_positions, chunk_ids32,
compressors.get(li), indexers.get(li),
moe_runners.get(li), se_runners.get(li), routers.get(li),
prod_lin=prod_lins.get(li),
_use_fused_rmsnorm_quantize=not _args.no_fused_rmsnorm,
comp_rope_cos=comp_rope_caches[gpu][0], comp_rope_sin=comp_rope_caches[gpu][1],
)
except Exception as e:
torch.cuda.synchronize()
print(f" CRASH at chunk {ci} (tokens {cs}-{ce-1}) layer {li} gpu {gpu}: {e}", flush=True)
raise
if VERBOSE >= 2 and ci == 0 and li < 3:
torch.cuda.synchronize(gpu)
print(f" Chunk {ci} L{li}: OK |X|={X.abs().max().item():.1f}", flush=True)
X = X.to('cuda:0'); torch.cuda.set_device(0)
print(f" Chunk {ci+1}/{len(chunk_starts)} tokens {cs}-{ce-1} ({chunk_len} tok): {time.time()-t1:.2f}s", flush=True)
print(f" Batched prefill done ({time.time()-t0:.1f}s)")
if _args.prefill_only: print("Prefill-only mode, stopping."); return
# ---- Build sampler ----
from dsv4.model.sampler import CUDASampler
sampler = CUDASampler(device='cuda:0', max_penalty_tokens=256)
sample_temp = _args.temperature
sample_topk = _args.top_k
sample_topp = _args.top_p
sample_rep_pen = _args.repetition_penalty
is_greedy = (sample_temp == 0.0)
print(f" Sampler: temp={sample_temp} top_k={sample_topk} top_p={sample_topp} "
f"rep_pen={sample_rep_pen} greedy={is_greedy}")
print(f" DSV4 reasoning model: thinking_start={THINK_START} thinking_end={THINK_END}")
print(f" Thinking tokens are NOT garbage — model uses ）、... format")
# Pre-allocate decode buffers — zero per-step allocation
dec_tid_buf = torch.zeros(1, dtype=torch.long, device='cuda:0')
dec_pos_buf = torch.zeros(1, dtype=torch.long, device='cuda:0')
dec_tid32_buf = torch.zeros(1, dtype=torch.int32, device='cuda:0')
# Per-GPU token ID buffers — each GPU needs its own copy for graph capture
# (cross-device .to() inside a CUDA graph is not reliable)
dec_tid32_per_gpu = {g: torch.zeros(1, dtype=torch.int32, device=f'cuda:{g}') for g in range(NUM_GPUS)}
dec_pos_per_gpu = {g: torch.zeros(1, dtype=torch.long, device=f'cuda:{g}') for g in range(NUM_GPUS)}
# Decode
print(f"\nDecoding (max {MAX_NEW_TOKENS} tokens)...")
in_thinking = False
profile = _args.profile
warmup_gsa = _args.warmup_gsa
prof_embed_layers = 0.0
prof_lm_head = 0.0
prof_sample = 0.0
prof_sample_start = 0.0
# CUDA event profiling — measures ACTUAL GPU time, not wall clock
# Only profile steps 1-3 (after warmup) to get stable results
cuda_events = {}
if profile:
for tag in ['embed', 'layers', 'hc_norm_lm', 'sample', 'diagnostics']:
cuda_events[tag] = (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True))
# Per-layer category events (sampled on step 1 only)
layer_event_tags = ['mhc_pre', 'attn_proj', 'rope_kv', 'compress_idx', 'fmha', 'inv_rope', 'o_proj',
'mhc_post', 'mhc_pre_ffn', 'router', 'moe', 'shared_expert', 'mhc_post_ffn']
cuda_layer_events = {}
for tag in layer_event_tags:
cuda_layer_events[tag] = (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True))
layer_event_accum = {tag: 0.0 for tag in layer_event_tags}
layer_event_count = 0
cuda_layer_events = [] # list of (tag, li, timestamp) for fine-grained profiling
# Pre-allocate decode X buffer — zero per-step allocation
# init_state writes to this buffer in-place (no .clone() allocation)
dec_X_buf = torch.zeros(1, 4, H, dtype=torch.bfloat16, device='cuda:0')
dec_embed_buf = torch.zeros(1, H, dtype=torch.bfloat16, device='cuda:0')
# Pre-allocate pinned CPU buffer for token ID transfer (graph-capturable)
dec_tid_pinned = torch.zeros(1, dtype=torch.long, device='cpu').pin_memory()
dec_tid32_pinned = torch.zeros(1, dtype=torch.int32, device='cpu').pin_memory()
dec_pos_pinned = torch.zeros(1, dtype=torch.long, device='cpu').pin_memory()
# ---- CUDA Graph Setup ----
graph_decoder = None
if _args.cuda_graph:
print(" CUDA graph capture requested — will capture after warmup step")
graph_decoder = CUDAGraphDecoder(n_layers, NUM_GPUS, H, [f'cuda:{g}' for g in range(NUM_GPUS)])
graph_decoder.pre_allocate(cfg)
for step in range(MAX_NEW_TOKENS):
t1 = time.time()
# Write token/position to pinned CPU buffers, then async copy to GPU
dec_tid_pinned[0] = all_tokens[-1]
dec_tid_buf.copy_(dec_tid_pinned)
dec_tid32_pinned[0] = all_tokens[-1]
dec_tid32_buf.copy_(dec_tid32_pinned)
dec_pos_pinned[0] = len(all_tokens) - 1
dec_pos_buf.copy_(dec_pos_pinned)
# Copy token/position to per-GPU buffers for graph capture
for g in range(NUM_GPUS):
dec_tid32_per_gpu[g].copy_(dec_tid32_pinned)
dec_pos_per_gpu[g].copy_(dec_pos_pinned)
t_e = time.perf_counter()
X = mHCLayer.init_state(embed(dec_tid_buf), out_buf=dec_X_buf)
# ---- Forward: graph replay or eager ----
if graph_decoder is not None and graph_decoder.captured:
# CUDA graph replay path
for li in range(n_layers):
gpu = li % NUM_GPUS
torch.cuda.set_device(gpu)
# Copy X into graph input buffer (copy_ handles cross-GPU transfer)
graph_decoder.x_in_a[li].copy_(X)
# Replay sub-graph A (mHC pre + rmsnorm + Q/KV projections)
graph_decoder.graphs_a[li].replay()
# Eager attention section (dynamic shapes)
attn_out = eager_attention(
graph_decoder.q_heads_buf[li],
graph_decoder.kv_roped_buf[li],
graph_decoder.x_normed_buf[li],
graph_decoder.q_a_buf[li],
layer_w[li], li, cfg, *rope_caches[gpu],
kv_caches[li], dec_pos_per_gpu[gpu],
compressors.get(li), indexers.get(li),
comp_rope_cos=comp_rope_caches[gpu][0] if comp_rope_caches else None,
comp_rope_sin=comp_rope_caches[gpu][1] if comp_rope_caches else None,
)
# Copy attn_out into graph buffer for sub-graph B
graph_decoder.attn_out_buf[li].copy_(attn_out)
# Replay sub-graph B (o_proj + mHC post + FFN)
graph_decoder.graphs_b[li].replay()
# Read output from sub-graph B (no clone — buffer is stable until next step)
X = graph_decoder.x_next_buf[li]
# Transfer last layer output to cuda:0 for lm_head graph
graph_decoder.x_lm_in.copy_(X) # cross-GPU copy if needed
# lm_head graph replay
graph_decoder.lm_graph.replay()
logits = graph_decoder.logits_buf
else:
# Eager forward path (warmup or no --cuda-graph)
for li in range(n_layers):
gpu = li % NUM_GPUS
if X.device != torch.device(f"cuda:{gpu}"): X = X.to(f"cuda:{gpu}")
torch.cuda.set_device(gpu)
X = forward_layer(X, layer_w[li], li, cfg, *rope_caches[gpu],
attn_mhcs.get(li), ffn_mhcs.get(li),
attn_norms.get(li), ffn_norms.get(li),
kv_caches[li], dec_pos_buf, dec_tid32_buf,
compressors.get(li), indexers.get(li),
moe_runners.get(li), se_runners.get(li), routers.get(li),
prod_lin=prod_lins.get(li),
_profile_detail=(profile and step == 1),
_profile_times=cuda_layer_events if (profile and step == 1) else None,
_use_fused_rmsnorm_quantize=not _args.no_fused_rmsnorm,
comp_rope_cos=comp_rope_caches[gpu][0], comp_rope_sin=comp_rope_caches[gpu][1],
)
X = X.to('cuda:0'); torch.cuda.set_device(0)
t_layers = time.perf_counter()
# After first decode step: fix gsa values from runtime amax
# This eliminates amax_gsa kernel launches on subsequent steps
# Only applies to attention linears and router gate (fixed per-projection gsa)
# MoE/SE keep runtime gsa (gsa varies per token)
if warmup_gsa and step == 0:
torch.cuda.synchronize()
n_fixed = 0
for li in range(n_layers):
pl = prod_lins.get(li)
if pl is None: continue
for key, lin in pl.items():
if hasattr(lin, '_gsa_buf') and hasattr(lin, '_use_runtime_gsa') and lin._use_runtime_gsa:
# Nvfp4GroupedLinear has per-group gsa; reduce to scalar (max) for fixed gsa
fixed_gsa = lin._gsa_buf.max().item() if lin._gsa_buf.numel() > 1 else lin._gsa_buf.item()
lin._activation_global_scale = fixed_gsa
lin._use_runtime_gsa = False
n_fixed += 1
# Router gate
router = routers.get(li)
if router and hasattr(router, '_gate_lin') and router._gate_lin is not None:
gl = router._gate_lin
if hasattr(gl, '_gsa_buf') and hasattr(gl, '_use_runtime_gsa') and gl._use_runtime_gsa:
fixed_gsa = gl._gsa_buf.item()
gl._activation_global_scale = fixed_gsa
gl._use_runtime_gsa = False
n_fixed += 1
# lm_head (BF16 — no gsa needed)
if lm_head_lin is not None and hasattr(lm_head_lin, '_gsa_buf') and hasattr(lm_head_lin, '_use_runtime_gsa') and lm_head_lin._use_runtime_gsa:
fixed_gsa = lm_head_lin._gsa_buf.item()
lm_head_lin._activation_global_scale = fixed_gsa
lm_head_lin._use_runtime_gsa = False
n_fixed += 1
print(f" Warmup gsa: fixed {n_fixed} projection gsa values from step 0 (MoE/SE keep runtime gsa)", flush=True)
# ---- lm_head: graph replay or eager ----
if graph_decoder is not None and graph_decoder.captured:
# logits already computed by lm_head graph replay above
pass
else:
x_out = hc_head.forward(X) if hc_head is not None else X[:, 0, :]
if final_norm_w is not None: x_out = rmsnorm(x_out, final_norm_w)
logits = torch.nn.functional.linear(x_out, lm_w) if lm_head_lin is None else lm_head_lin(x_out)
# ---- CUDA graph capture after warmup ----
if graph_decoder is not None and not graph_decoder.captured and step == 0:
print(" Step 0 warmup done. Capturing CUDA graphs...", flush=True)
torch.cuda.synchronize()
graph_decoder.capture(
cfg, attn_mhcs, ffn_mhcs, attn_norms, ffn_norms,
kv_caches, compressors, indexers, moe_runners, se_runners,
routers, prod_lins, layer_w, rope_caches, hc_head,
final_norm_w, lm_w, dec_pos_per_gpu, dec_tid32_per_gpu,
comp_rope_caches=comp_rope_caches,
)
print(f" CUDA graphs captured. Graph replay starts on step 1.", flush=True)
if profile: torch.cuda.synchronize()
t_lm = time.perf_counter()
# Check thinking start token logit on first step
if step == 0:
ls = logits.float()
for tid, name in [(THINK_START, 'think_start'), (THINK_END, 'think_end'), (USER_TOKEN, 'user'), (ASSISTANT_TOKEN, 'assistant')]:
print(f" {name}({tid}) logit={ls[0, tid].item():.2f}", flush=True)
# Paris token check — only check known token IDs, no 129K iteration
for t in [11111, 51119, 60107]:
if t < ls.shape[-1]:
print(f" Paris-candidate({t}) logit={ls[0, t].item():.2f}", flush=True)
# Sync for profiling and error check
if profile: torch.cuda.synchronize()
t_sample_start = time.perf_counter()
# Only sync + validate on first 3 steps and every 20th step (reduces pipeline stalls)
if step < 3 or (step + 1) % 20 == 0:
torch.cuda.synchronize() # catch CUDA errors at source
ls = logits.float()
if step < 3 or (step + 1) % 20 == 0:
has_nan = torch.isnan(ls).any().item()
has_inf = torch.isinf(ls).any().item()
print(f" logits: shape={list(logits.shape)} dtype={logits.dtype} "
f"min={ls.min().item():.1f} max={ls.max().item():.1f} "
f"nan={has_nan} inf={has_inf}", flush=True)
if has_nan or has_inf:
print(f" NaN/Inf in logits at step {step}, aborting", flush=True)
break
# Sampling — fused CUDA kernel (or greedy argmax for temp=0)
if is_greedy:
next_id = torch.argmax(logits, -1).item()
else:
sampled = sampler(
logits,
temperature=sample_temp,
top_k=sample_topk,
top_p=sample_topp,
repetition_penalty=sample_rep_pen,
recent_tokens=all_tokens[-256:],
seed=SEED,
)
# Check for async CUDA errors from sampler
if step < 3:
torch.cuda.synchronize()
next_id = sampled[0].item()
all_tokens.append(next_id)
dt = time.time() - t1
if profile: torch.cuda.synchronize()
t_s = time.perf_counter()
# Track thinking state
if next_id == THINK_START: in_thinking = True
elif next_id == THINK_END: in_thinking = False
if profile:
prof_embed_layers += (t_layers - t_e)
prof_lm_head += (t_lm - t_layers)
prof_sample_start = t_sample_start
prof_sample += (t_s - t_sample_start)
# Diagnostics — every step for first 20, then every 5th
if step < 20 or step % 5 == 0:
tv, ti = torch.topk(logits[0].float(), 5)
top5 = ' '.join(f'{tokenizer.decode([t.item()])}({v.item():.1f})' for t, v in zip(ti[:5], tv[:5]))
think_tag = " [THINKING]" if in_thinking else ""
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"|X|={X.abs().max().item():.1f} top5: {top5}{think_tag}", flush=True)
# NaN safety — periodic check only
if step == 0 or (step+1) % 20 == 0:
if torch.isnan(logits.float()).any().item():
print(f" NaN at step {step}", flush=True); break
if next_id in STOP_IDS:
print(f" STOP ({next_id}) at step {step} — token='{tokenizer.decode([next_id])}'", flush=True); break
if profile and MAX_NEW_TOKENS > 0:
n = MAX_NEW_TOKENS
print(f"\n PROFILE (sync'd wall clock, {n} steps):")
print(f" Embed + 61 layers: {prof_embed_layers:.3f}s total, {prof_embed_layers/n*1000:.1f}ms/token")
print(f" hc_head + norm + lm_head: {prof_lm_head:.3f}s total, {prof_lm_head/n*1000:.1f}ms/token")
print(f" Sampling: {prof_sample:.3f}s total, {prof_sample/n*1000:.1f}ms/token")
# Fine-grained attention profile (from step 1)
if hasattr(cuda_layer_events, '__len__') and len(cuda_layer_events) >= 2:
print(f"\n FINE-GRAINED ATTENTION PROFILE (step 1, CUDA-sync'd):")
prev_t = None
for tag, li, t in cuda_layer_events:
if prev_t is not None:
dt_ms = (t - prev_t) * 1000
if li <= 2 or li >= 58: # Only print for first/last layers
print(f" L{li} {tag}: {dt_ms:.2f}ms")
prev_t = t
out_raw = tokenizer.decode(all_tokens, skip_special_tokens=False)
# Use official DSV4 parser for structured output
try:
from encoding.deepseek_v4_encoding import parse_message_from_completion_text
# Find the assistant portion — after the last ASSISTANT token
assistant_start = out_raw.find(_ASSISTANT_STR)
if assistant_start >= 0:
assistant_text = out_raw[assistant_start + len(_ASSISTANT_STR):]
else:
assistant_text = out_raw
parsed = parse_message_from_completion_text(assistant_text, thinking_mode=_args.thinking_mode)
reasoning = parsed.get('reasoning', '')
content = parsed.get('content', '')
print(f"\n{'='*70}")
print(f"Input: '{PROMPT}'")
if reasoning:
print(f"Reasoning: {reasoning[:500]}{'...' if len(reasoning) > 500 else ''}")
print(f"Content: {content}")
print(f"Total: {time.time()-t0:.1f}s")
print(f"{'='*70}")
except Exception as e:
# Fallback: raw decode (shouldn't happen with correct output)
out = tokenizer.decode(all_tokens, skip_special_tokens=True)
print(f"\n{'='*70}")
print(f"Input: '{PROMPT}'")
print(f"Output (raw): '{out}'")
print(f"Parse error: {e}")
print(f"Total: {time.time()-t0:.1f}s")
print(f"{'='*70}")
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink