From 2dd16d578940b47ad76d2a06c6df30ac09688313 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 14:19:41 +0000 Subject: [PATCH] Switch compressor + indexer weights_proj to BF16 F.linear MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- single_shot_inference.py | 50 ++++++++++++++++++++++++---------------- 1 file changed, 30 insertions(+), 20 deletions(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index bd6a376d..5e00618f 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -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.