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

@@ -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,

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) "