biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Switch compressor + indexer weights_proj to BF16 F.linear

Browse Source 
Only the CSA indexer QK path (q_b_proj) is explicitly FP4-QATed.
The rest of the compressor/indexer projections are NOT, so use BF16:

- Compressor kv_proj, gate_proj: dequantize NVFP4 → BF16, F.linear
- Indexer weights_proj: dequantize NVFP4 → BF16, F.linear
- Indexer q_b_proj: KEEP as NVFP4 (this IS the FP4-QATed path)
- Indexer compressor: inherits Compressor's BF16 path
This commit is contained in:
biondizzle
2026-06-03 14:19:41 +00:00
parent 95e45a87e3
commit 2dd16d5789
1 changed files with 30 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
single_shot_inference.py
View File

@@ -302,6 +302,8 @@ class Compressor:
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.
@@ -312,26 +314,29 @@ class Compressor:
self._buf_len = 0
def load(self, w, pfx, dev=None):
"""Load weights and build production Nvfp4Linear instances."""
"""Load weights and build BF16 projections (dequantized from NVFP4)."""
if dev is None: dev = self.device
# Build production NVFP4 GEMM instances for the two projections
# kv_proj: in=7168, out=kv_dim (1024 for CSA, 512 for HCA)
# gate_proj: same shapes
# Compressor projections are NOT explicitly FP4-QATed — use BF16 F.linear
from dsv4.ops.quantize import dequantize_nvfp4
kv_w, kv_ws, kv_ws2, kv_isc = get_nvfp4_weight(w, pfx, 'kv_proj')
gate_w, gate_ws, gate_ws2, gate_isc = get_nvfp4_weight(w, pfx, 'gate_proj')
if kv_w is not None:
kv_out = kv_w.shape[0] # N_packed
kv_in = kv_w.shape[1] * 2 # K_packed * 2
self.kv_lin = make_nvfp4_linear(kv_in, kv_out, dev, w, pfx, 'kv_proj')
ws2_v = kv_ws2.float().item() if kv_ws2 is not None else 1.0
gsb = 1.0 * ws2_v
gsa = torch.tensor([gsb] * kv_w.shape[0], device=dev, dtype=torch.float32)
kv_bf16 = dequantize_nvfp4(kv_w.to(dev), kv_ws.to(dev), gsa) # (out, in)
self._kv_bf16 = kv_bf16.to(dev).contiguous()
if gate_w is not None:
gate_out = gate_w.shape[0]
gate_in = gate_w.shape[1] * 2
self.gate_lin = make_nvfp4_linear(gate_in, gate_out, dev, w, pfx, 'gate_proj')
ws2_v = gate_ws2.float().item() if gate_ws2 is not None else 1.0
gsb = 1.0 * ws2_v
gsa = torch.tensor([gsb] * gate_w.shape[0], device=dev, dtype=torch.float32)
gate_bf16 = dequantize_nvfp4(gate_w.to(dev), gate_ws.to(dev), gsa) # (out, in)
self._gate_bf16 = gate_bf16.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_lin is None: return None, None, None
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.
@@ -358,9 +363,9 @@ class Compressor:
n_complete = T // r
if n_complete == 0: return None, None, None
# Step 1-2: NVFP4 GEMM projections → FP32 for compress
kv = self.kv_lin(hidden_states).float() # (T, kv_dim) FP32
gate = self.gate_lin(hidden_states).float() # (T, kv_dim) FP32
# 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.
@@ -398,22 +403,27 @@ class Indexer:
"""
def __init__(self, n_ih, ihd, top_k, device):
self.n_ih, self.ihd, self.top_k, self.device = n_ih, ihd, top_k, device
self.q_b_lin = None # production Nvfp4Linear for q_b_proj
self.wp_lin = None # production Nvfp4Linear for weights_proj
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 — use BF16 F.linear
if wp_w is not None:
wp_out = wp_w.shape[0]
wp_in = wp_w.shape[1] * 2
self.wp_lin = make_nvfp4_linear(wp_in, wp_out, dev, w, pfx, 'weights_proj')
from dsv4.ops.quantize import dequantize_nvfp4
ws2_v = wp_ws2.float().item() if wp_ws2 is not None else 1.0
gsb = 1.0 * ws2_v
gsa = torch.tensor([gsb] * wp_w.shape[0], device=dev, dtype=torch.float32)
wp_bf16 = dequantize_nvfp4(wp_w.to(dev), wp_ws.to(dev), gsa)
self._wp_bf16 = wp_bf16.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:
@@ -436,7 +446,7 @@ class Indexer:
li = layer_idx
q_idx = self.q_b_lin(q_lora).reshape(T, self.n_ih, self.ihd) # (T, n_ih, ihd)
w_h = self.wp_lin(hidden_states) # (T, n_ih)
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.