fix(ArcKernels): FP8 threshold was 512 because 512 was the only point measured; cuBLASLt branch never flattened to 2-D

mistralrs-quant/kernels/blockwise_fp8/blockwise_fp8_gemm.cu

@@ -51,8 +51,19 @@ __global__ void fp8_matmul_tiled(const T *__restrict__ input,
                                  T *__restrict__ output, int M, int N, int K,
                                  int scale_row_stride, int block_size_y,
                                  int block_size_x) {
-  __shared__ float s_input[BLOCK_M][BLOCK_K + 4];
-  __shared__ float s_weight[BLOCK_N][BLOCK_K + 4];
+  // Padding is +1, not +4. The inner product reads `s_weight[tx][k]` with `tx`
+  // varying fastest within a warp, so consecutive lanes are `BLOCK_K + pad`
+  // floats apart. At BLOCK_K=32 a +4 pad gives stride 36; 36 mod 32 = 4, so
+  // gcd(4,32)=4 and the warp lands on only 8 of the 32 shared-memory banks --
+  // a 4-way conflict on every FMA of the k loop. A +1 pad gives stride 33,
+  // which is coprime with 32, so all 32 banks are hit and the access is
+  // conflict-free.
+  //
+  // Bit-identical output: the k loop is bounded by `k < BLOCK_K`, so columns
+  // BLOCK_K..BLOCK_K+pad-1 are written by nobody and read by nobody. The pad
+  // exists only to change the row stride.
+  __shared__ float s_input[BLOCK_M][BLOCK_K + 1];
+  __shared__ float s_weight[BLOCK_N][BLOCK_K + 1];
 
   const int bx = blockIdx.x;
   const int by = blockIdx.y;

mistralrs-quant/src/blockwise_fp8/mod.rs

@@ -142,6 +142,34 @@ impl QuantMethod for BlockwiseFP8Linear {
             weight,
             self.bias.clone(),
         )))?;
+
+        // 🔴 FLATTEN TO 2-D FIRST — the native branch above already does this at
+        // the top of `forward`, and this branch did not.
+        //
+        // V4's activation is rank-3 `[B, T, hidden]`. Passed through untouched,
+        // `UnquantLinear::forward` takes `w.broadcast_left(B)` and, because bias
+        // is `None` on every V4 linear, dispatches `cublaslt.batch_matmul` with
+        // `stride_b = 0`. At decode T=1, so that is a **B-way batched GEMM with
+        // m=1** — a batched GEMV wearing a GEMM's name, which is the one shape
+        // cuBLASLt has no advantage in. The documented 27x was measured at
+        // prefill, where `[1, 512, hidden]` collapses to batch=1 and the call is
+        // a genuine `[512, hidden] x [hidden, N]` GEMM.
+        //
+        // Without this, lowering `ARC_FP8_CUBLAS_MIN_M` measures the batched-GEMV
+        // path and concludes "cuBLASLt loses at decode" — a wrong conclusion
+        // drawn from a shape the threshold change was never meant to select.
+        // Flattening makes the two branches comparable, which is the whole
+        // premise of the A/B.
+        let orig_dims = x.dims().to_vec();
+        if orig_dims.len() > 2 {
+            let features = orig_dims[orig_dims.len() - 1];
+            let rows: usize = orig_dims[..orig_dims.len() - 1].iter().product();
+            let out = unquant.forward(&x.reshape((rows, features))?)?;
+            let out_features = out.dim(out.rank() - 1)?;
+            let mut new_dims = orig_dims[..orig_dims.len() - 1].to_vec();
+            new_dims.push(out_features);
+            return out.reshape(new_dims);
+        }
         unquant.forward(x)
     }
 
@@ -555,18 +583,68 @@ pub fn blockwise_fp8_linear_b(
 /// Measured on an H200 against DeepSeek-V4-Flash at 512 prompt tokens: the
 /// cuBLASLt path is 27x faster on the MLA `q_proj`/`o_proj` projections
 /// (706.1 -> 26.2 ms/step) and 21x on the shared expert (253.2 -> 11.9), for
-/// 1.513x on prefill TTFT end to end. The default is set to that measured
-/// point rather than to an interpolated crossover.
+/// 1.513x on prefill TTFT end to end.
+///
+/// # 🔴 THE DEFAULT WAS 512 BECAUSE 512 WAS THE ONLY POINT MEASURED
+///
+/// The previous doc said so outright — *"the default is set to that measured
+/// point rather than to an interpolated crossover"* — and 512 then became a
+/// wall. `fp8_gemv_warp` owns `M <= 4`, so **`M = 5..511` fell through to
+/// `fp8_matmul_tiled`**, the kernel whose own comment says there is *"no
+/// tensor-core instruction anywhere in it"*. A profile of the shipped default
+/// caught it running **7,525 times for 17.6% of GPU time** in one decode
+/// window.
+///
+/// Swept on an H200, V4-Flash, one binary and one env toggle, aggregate tok/s,
+/// **clean rows only** (>=95% achieved concurrency, <15% derived-vs-measured
+/// spread):
+///
+/// ```text
+///   M=1     34.92  vs  25.27   cuBLASLt 0.72x  -- the GEMV floor is REAL
+///   M=8     49.25  vs  64.95   cuBLASLt 1.32x
+///   M=16    45.61  vs  52.55   cuBLASLt 1.15x
+///   M=128  114.21  vs 130.76   cuBLASLt 1.14x
+/// ```
+///
+/// So the crossover is not 512; it is immediately above the GEMV's domain. The
+/// default is now **5** — the first row `fp8_gemv_warp` does not own.
+/// `M = 5..7` is **unmeasured**, bounded below by the GEMV domain and above by
+/// the +32% at M=8.
+///
+/// **Keep the floor.** At M=1 cuBLASLt is *worse* (0.72x): the dedicated
+/// `fp8_gemv_warp` kernel exists for a reason and forcing the library path over
+/// it costs 28% of b=1 decode. That row is why this constant is 5 and not 1.
+///
+/// **Mechanism, not just a number.** The tiled kernel is scalar, so its cost is
+/// instruction-bound and grows with M while the memory controller idles. The
+/// cuBLASLt path dequantizes once and hands a tensor-core GEMM the whole tile.
+/// The signature is visible in the counter: memory-controller utilisation went
+/// **2-3% -> 6-19%** across the sweep, the first intervention this session to
+/// move it at all. A scalar kernel that keeps the card 100% busy and the memory
+/// controller at 2% is not doing the work; it is doing arithmetic *about* the
+/// work.
+///
+/// # A recurring failure in this codebase: one measured point frozen as a gate
+///
+/// This is the **second** gate found this way. `qtip::gather_policy`'s tile-fill
+/// predicate required `6n/256 >= 16`, i.e. n >= 683, and so kept the amortizing
+/// grouped GEMM unreachable in every decode step; removing it was worth 1.45x
+/// at B=128. Both gates were derived from a single working point and neither
+/// was ever swept. **If you are about to set a dispatch threshold from one
+/// measurement, sweep it first and write down which rows were clean** — a
+/// threshold is a claim about every value it excludes, not just the one you
+/// tried.
 ///
-/// The decode regime is untouched: it has its own dedicated `fp8_gemv_warp`
-/// kernel for M <= 4 and never reaches this threshold.
+/// `fp8_gemv_warp` still owns `M <= 4` and is not affected by this constant.
 ///
 /// `ARC_FP8_CUBLAS_MIN_M` overrides; a value above any real batch restores
 /// the previous always-native dispatch.
 #[cfg(feature = "cuda")]
 fn arc_fp8_cublas_min_m() -> usize {
     use std::sync::OnceLock;
-    const DEFAULT_MIN_M: usize = 512;
+    // 5 = the first row `fp8_gemv_warp` does not own. Swept, clean rows only;
+    // M=1 measured 0.72x for cuBLASLt, which is why this is not 1.
+    const DEFAULT_MIN_M: usize = 5;
     static CACHE: OnceLock<usize> = OnceLock::new();
     *CACHE.get_or_init(|| {
         std::env::var("ARC_FP8_CUBLAS_MIN_M")

mistralrs-quant/src/blockwise_fp8/ops.rs

@@ -1034,8 +1034,7 @@ fn fp8_blockwise_matmul_impl(
 
     match input.dtype() {
         DType::F16 => {
-            let output =
-                crate::arc_outbuf::alloc_out_fully_written::<f16>(&dev, (m * n) as usize)?;
+            let output = crate::arc_outbuf::alloc_out_fully_written::<f16>(&dev, (m * n) as usize)?;
 
             let input_s = match &*input_storage {
                 Storage::Cuda(cuda_storage) => cuda_storage.as_cuda_slice::<f16>()?,