From 5493a8727e286e2dacce54d455e0b163c82fb61d Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Mon, 1 Jun 2026 22:29:56 +0000
Subject: [PATCH] P7: compressor early return + decode buffering (skip GEMMs
 when n_complete=0); sampler SMEM fix (LK=24 fits 48KB default); topk on float
 not bf16

---
 dsv4/kernels/cuda/sampler.cu | 15 ++++++++++++++-
 single_shot_inference.py     | 30 +++++++++++++++++++++++++++++-
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/dsv4/kernels/cuda/sampler.cu b/dsv4/kernels/cuda/sampler.cu
index 77fef23d..8c773b56 100644
--- a/dsv4/kernels/cuda/sampler.cu
+++ b/dsv4/kernels/cuda/sampler.cu
@@ -34,7 +34,7 @@
 #include <curand_kernel.h>
 
 static constexpr int BDIM = 256;
-static constexpr int LK = 32;  // per-thread local top-k
+static constexpr int LK = 24;  // per-thread local top-k (SMEM budget: 256*24*8=48KB fits default)
 
 // ---------------------------------------------------------------------------
 // Insert into sorted descending array (register-resident, k small)
@@ -173,6 +173,19 @@ torch::Tensor sample_cuda(
     auto out = torch::empty({B}, options);
     int smem = BDIM * LK * (sizeof(float) + sizeof(int));
 
+    // Request enough shared memory for 48KB+ per block
+    cudaFuncSetAttribute(
+        fused_sampler_kernel,
+        cudaFuncAttributeMaxDynamicSharedMemorySize,
+        smem
+    );
+    // Carveout: prefer more shared memory over L1
+    cudaFuncSetAttribute(
+        fused_sampler_kernel,
+        cudaFuncAttributePreferredSharedMemoryCarveout,
+        cudaSharedmemCarveoutMaxShared
+    );
+
     fused_sampler_kernel<<<B, BDIM, smem, c10::cuda::getCurrentCUDAStream()>>>(
         logits.data_ptr<float>(), pi, pv,
         B, V, logits.stride(0), mp,
diff --git a/single_shot_inference.py b/single_shot_inference.py
index c52a9a72..ba747902 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -181,6 +181,12 @@ class Compressor:
         self.gate_lin = None  # production Nvfp4Linear for gate_proj
         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 production Nvfp4Linear instances."""
@@ -204,6 +210,28 @@ class Compressor:
     def forward(self, hidden_states, positions):
         if self.ratio == 0 or self.kv_lin is None: return None, None, None
         T = hidden_states.shape[0]; r = self.ratio; dev = hidden_states.device
+
+        # 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
 
@@ -1017,7 +1045,7 @@ def main():
 
         # Diagnostics — reduce CPU syncs, only top-5 every 5 steps
         if step % 5 == 0 or step < 5:
-            tv, ti = torch.topk(logits[0], 5)
+            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) "