Add B2 minimal debug test to find hang point

tests/unit/test_b2_minimal.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+"""B2 minimal debug: test if the FP8 indexer kernel even completes.
+
+Start with n_comp=1 (trivial case) and work up.
+Also test: does the kernel return at all? Or does it hang?
+"""
+import sys
+import math
+import torch
+import torch.nn.functional as F
+import time
+
+
+def quantize_fp8_e4m3(x_fp32):
+    amax = x_fp32.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
+    scale = amax / 448.0
+    fp8 = (x_fp32 / scale).clamp(-448, 448).to(torch.float8_e4m3fn)
+    return fp8.view(torch.uint8), scale.squeeze(-1)
+
+
+def main():
+    N_IH = 64; IHD = 128; TOP_K = 1
+    
+    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",
+                          ])
+
+    # Test with increasing n_comp
+    for n_comp in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]:
+        print(f"\n--- n_comp={n_comp} ---", flush=True)
+        torch.manual_seed(42)
+        q_idx = torch.randn(N_IH, IHD, dtype=torch.bfloat16).cuda() * 0.5
+        k_fp32 = torch.randn(n_comp, IHD, dtype=torch.float32) * 0.5
+        w_h = torch.randn(N_IH, dtype=torch.bfloat16).cuda() * 0.3
+        k_fp8, k_scale = quantize_fp8_e4m3(k_fp32)
+        k_fp8 = k_fp8.cuda(); k_scale = k_scale.cuda()
+
+        tk = min(TOP_K, n_comp)
+        topk_indices = torch.empty(tk, dtype=torch.int32, device='cuda')
+
+        # Add a CUDA event to time the kernel
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        try:
+            mod.indexer_fp8_score_topk(
+                q_idx, k_fp8, k_scale, w_h, topk_indices,
+                N_IH, IHD, tk)
+            end.record()
+            torch.cuda.synchronize()
+            ms = start.elapsed_time(end)
+            print(f"  OK: {ms:.2f}ms indices={topk_indices.cpu().tolist()}", flush=True)
+        except Exception as e:
+            torch.cuda.synchronize()
+            print(f"  FAILED: {e}", flush=True)
+            break
+        
+        # Check for CUDA errors
+        err = torch.cuda.current_device()
+        if n_comp >= 128:
+            # After a certain size, try timing it
+            # Also check if the GPU is responsive
+            print(f"  GPU responsive: yes", flush=True)
+
+    print("\nDone.")
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()