feat(VT-MATMUL-FP8-BLOCK-REF): add vt::MatmulFp8BlockScaled, the block-scaled fp8 GEMM whose scales cannot live in an epilogue

include/vt/ops.h

@@ -433,6 +433,14 @@ enum class OpId : uint8_t {
   // See vt::QuantFp8Group below for the contract.
   // Appended before kCount so no existing op's id shifts.
   kQuantFp8Group,
+  // --- Block-wise FP8 (VT-MATMUL-FP8-BLOCK-REF, #1189 milestone M2). The
+  // 128x128 block-scaled fp8 GEMM that consumes kQuantFp8Group's output. It is
+  // NOT a parameter of kMatmulFp8Cutlass and cannot be: that op folds ONE
+  // scalar alpha into the epilogue, which has exactly one degree of freedom per
+  // output element, while this scheme has cdiv(K, block_k) of them and applies
+  // them in the MAINLOOP. See vt::MatmulFp8BlockScaled below for the contract.
+  // Appended before kCount so no existing op's id shifts.
+  kMatmulFp8BlockScaled,
   kCount
 };
 
@@ -936,6 +944,13 @@ using QuantFp8StaticFn = void (*)(Queue&, Tensor&, const Tensor&, float);
 // and the f32 [M, K/group_size] per-group scale.
 using QuantFp8GroupFn = void (*)(Queue&, Tensor& /*out_fp8*/, Tensor& /*out_scale*/,
                                  const Tensor& /*x*/, int /*group_size*/);
+// Two scale streams and the block geometry, where the per-tensor fp8 GEMMs above
+// carry one scalar alpha. The geometry is not a convenience: it is what selects
+// which scale pair each K-block multiplies by.
+using MatmulFp8BlockScaledFn = void (*)(Queue&, Tensor& /*out*/, const Tensor& /*a_fp8*/,
+                                        const Tensor& /*a_scale*/, const Tensor& /*b_fp8*/,
+                                        const Tensor& /*b_scale*/, int /*block_n*/,
+                                        int /*block_k*/);
 using RmsNormQuantFp8Fn = void (*)(Queue&, Tensor& /*out_fp8*/, Tensor* /*out_bf16*/,
                                    const Tensor& /*x*/, const Tensor& /*weight*/,
                                    const RmsNormArgs&, Tensor* /*residual*/, float /*input_scale*/);
@@ -1582,6 +1597,69 @@ void QuantFp8Static(Queue& q, Tensor& out_fp8, const Tensor& x, float input_scal
 void QuantFp8Group(Queue& q, Tensor& out_fp8, Tensor& out_scale, const Tensor& x,
                    int group_size);
 
+// MatmulFp8BlockScaled (VT-MATMUL-FP8-BLOCK-REF, #1189 M2,
+// .agents/specs/vt-matmul-fp8-block-ref.md) — the 128x128 block-scaled fp8 GEMM,
+// mirroring native_w8a8_block_matmul (tests/kernels/quant_utils.py:91-154):
+//
+//   for each (m, n):
+//     acc = 0                                              f32
+//     for kt in [0, cdiv(K, block_k)):
+//       part = 0                                           f32, a SEPARATE register
+//       for k in the k-tile:  part += f8(a[m,k]) * f8(b[n,k])
+//       acc += part * ( a_scale[m, kt] * b_scale[n / block_n, kt] )
+//     out[m, n] = acc                                      stored to out's dtype
+//
+// THE SCALES APPLY IN THE MAINLOOP, ONCE PER K-BLOCK, INTO AN F32 ACCUMULATOR —
+// NOT IN THE EPILOGUE, and that is a correctness constraint rather than an
+// optimisation choice. MatmulFp8Cutlass above folds one scalar alpha after the
+// whole K reduction. An epilogue has exactly ONE degree of freedom per output
+// element; this scheme has cdiv(K, block_k) of them. An epilogue-only
+// application therefore cannot express a per-K-block scale AT ALL, which is why
+// this is a separate op. DO NOT "simplify" `part` away into `acc`: that IS the
+// epilogue form, and tests/vt/test_ops_matmul_fp8_block_cpu.cpp G4 is built so
+// that no single-alpha implementation can pass it.
+//
+// WHICH UPSTREAM ARM THIS MIRRORS. The Triton kernel at fp8_utils.py:826-836 is
+// not what executes on the target architecture; CUTLASS is
+// (vllm/model_executor/kernels/linear/__init__.py:355-377 ranks it third,
+// DeepGEMM is auto-disabled for qwen3_5_text on family 120 at
+// vllm/utils/deep_gemm.py:27-46, and Marlin is excluded at cc >= 89). Unlike the
+// QuantFp8Group case above, the two AGREE: csrc/.../c3x/
+// scaled_mm_blockwise_sm120_fp8_dispatch.cuh:56-58,218-235 hands both scale
+// pointers to the MAINLOOP arguments over an `ElementAccumulator = float`, and
+// cutlass 4.5.0's sm120_mma_tma_blockwise_scaling.hpp:714-717 is literally
+// `accum(i) += tmp_accum(i) * tCrScaleAViewAsC(i) * tCrScaleBViewAsC(i)`.
+// CUTLASS associates the two scale multiplies left to right where the reference
+// forms their product first (quant_utils.py:150-151); the difference is at most
+// one f32 ULP per K-block and upstream's own gate admits it, comparing the two
+// at rel_diff < 0.001 (test_block_fp8.py:194-200). We mirror the reference's
+// association, because this op IS the reference port and the CUDA kernel that
+// #1189 milestone M5 lands will be measured against it.
+//
+// SHAPES, with CEIL on every tiling, so a ragged final block is legal and must
+// work (upstream asserts exactly this at fp8_utils.py:935-936):
+//   a_fp8   [M,K]                              i8, raw fp8-e4m3fn bytes
+//   a_scale [M, cdiv(K, block_k)]              F32
+//   b_fp8   [N,K]                              i8, raw fp8-e4m3fn bytes
+//   b_scale [cdiv(N, block_n), cdiv(K, block_k)] F32
+//   out     [M,N]                              f32 or bf16
+// The scales are f32 because upstream refuses any other dtype on this path
+// (csrc/.../c3x/scaled_mm_helper.hpp:15-18) and the accumulator is f32.
+// a_scale's K axis is a CEIL too, so this op accepts a K that vt::QuantFp8Group
+// would refuse; that asymmetry is upstream's own (fp8_utils.py:930 uses cdiv
+// where fp8_utils.py:596-599 demands divisibility).
+//
+// No bias: upstream refuses one outright on the blockwise path
+// (scaled_mm_helper.hpp:54), so this mirrors a refusal rather than deferring a
+// feature.
+//
+// CPU only. A CORRECTNESS REFERENCE, NOT A PERFORMANCE PATH — it is the
+// numerical oracle #1189 milestone M5's CUTLASS kernel is measured against, and
+// it makes no speed claim. M5 owns the CUDA arm.
+void MatmulFp8BlockScaled(Queue& q, Tensor& out, const Tensor& a_fp8, const Tensor& a_scale,
+                          const Tensor& b_fp8, const Tensor& b_scale, int block_n,
+                          int block_k);
+
 // RmsNormQuantFp8 (fused fp8 RMSNorm -> static per-tensor activation quant). One
 // HBM pass mirrors vLLM's Inductor `fused_add_rms_norm_static_fp8_quant`
 // (vllm/compilation/passes/fusion/rms_quant_fusion.py:124) — the RMSNorm producer

src/vt/cpu/cpu_ops.cpp

@@ -628,6 +628,79 @@ void QuantFp8GroupKernel(Queue&, Tensor& out_fp8, Tensor& out_scale, const Tenso
   });
 }
 
+// --- Block-wise FP8 (VT-MATMUL-FP8-BLOCK-REF, #1189 M2). MatmulFp8BlockScaled
+// CPU kernel: the 128x128 block-scaled fp8 GEMM, mirroring
+// native_w8a8_block_matmul (tests/kernels/quant_utils.py:91-154).
+//
+// THE SCALES APPLY IN THE MAINLOOP, ONCE PER K-BLOCK. `part` is a SEPARATE f32
+// accumulator that is scaled and only then folded into `acc`. That separation is
+// the whole point of the op: MatmulFp8CutlassKernel fifty lines below folds one
+// scalar alpha AFTER the whole K reduction, which has exactly one degree of
+// freedom per output element while this scheme has cdiv(K, block_k) of them. An
+// epilogue-only application cannot express a per-K-block scale AT ALL. Collapsing
+// `part` into `acc` here IS that epilogue form, and
+// tests/vt/test_ops_matmul_fp8_block_cpu.cpp G4 is constructed so it cannot pass.
+//
+// The scale PRODUCT is formed first, `part * (a_s * b_s)`, mirroring
+// `s = As_tiles[i] * Bs[j][i]` then `c += matmul(a, b.t()) * s`
+// (quant_utils.py:150-151). The kernel that executes upstream is CUTLASS, not
+// Triton (vllm/model_executor/kernels/linear/__init__.py:355-377,
+// vllm/utils/deep_gemm.py:27-46), and it associates left to right --
+// `accum(i) += tmp_accum(i) * tCrScaleAViewAsC(i) * tCrScaleBViewAsC(i)`, cutlass
+// 4.5.0 sm120_mma_tma_blockwise_scaling.hpp:714-717. The two differ by at most one
+// f32 ULP per K-block, and upstream's own gate is what admits it: it compares the
+// CUTLASS arm against THIS reference at rel_diff < 0.001
+// (test_block_fp8.py:194-200). We mirror the reference, because this op is the
+// reference port.
+//
+// `col / block_n` indexes the b_scale ROW by OUTPUT COLUMN, mirroring
+// `offs_bsn = offs_bn // group_n` (fp8_utils.py:823) and the reference's tiling
+// (quant_utils.py:131-143). For a round N that agrees with a tile counter; for a
+// ragged N it does not, which is why N=576 is in the ported grid.
+//
+// CEIL tiling and a short final K-tile: `k1 = min(k0 + block_k, k)`
+// (quant_utils.py:131-141). Upstream's wrapper asserts the ceil shapes
+// (fp8_utils.py:935-936), so a ragged block is legal rather than tolerated.
+//
+// A CORRECTNESS REFERENCE, NOT A PERFORMANCE PATH, in the same sense as
+// MatmulFp8CutlassKernel below: a naive nest that makes the block-fp8 seam
+// resolvable on a CPU queue so #1189 milestones M3 and M4 can be gated without a
+// GPU. It makes no speed claim. M5 owns the CUDA arm and will be measured
+// against this one. Parallel over ROWS; each output row is independent and the
+// reduction order inside a row is fixed, so the result does not depend on the
+// thread count.
+void MatmulFp8BlockScaledKernel(Queue&, Tensor& out, const Tensor& a_fp8, const Tensor& a_scale,
+                                const Tensor& b_fp8, const Tensor& b_scale, int block_n,
+                                int block_k) {
+  const int64_t m = a_fp8.shape[0], k = a_fp8.shape[1], n = b_fp8.shape[0];
+  const int64_t k_tiles = (k + block_k - 1) / block_k;
+  const auto* ap = a_fp8.Ptr<uint8_t>();
+  const auto* bp = b_fp8.Ptr<uint8_t>();
+  const auto* asp = a_scale.Ptr<float>();
+  const auto* bsp = b_scale.Ptr<float>();
+  ForRows(m, [&](int64_t r0, int64_t r1) {
+    std::vector<float> arow(static_cast<size_t>(k));
+    for (int64_t i = r0; i < r1; ++i) {
+      // Decode the A row once and reuse it across N, as MatmulNvfp4Fp4Kernel does.
+      for (int64_t kk = 0; kk < k; ++kk)
+        arow[static_cast<size_t>(kk)] = Fp8ToF32(ap[i * k + kk]);
+      for (int64_t col = 0; col < n; ++col) {
+        const int64_t nb = col / block_n;   // fp8_utils.py:823
+        float acc = 0.0F;
+        for (int64_t kt = 0; kt < k_tiles; ++kt) {
+          const int64_t k0 = kt * block_k;
+          const int64_t k1 = std::min(k0 + static_cast<int64_t>(block_k), k);
+          float part = 0.0F;                // the MAINLOOP register, kept separate
+          for (int64_t kk = k0; kk < k1; ++kk)
+            part += arow[static_cast<size_t>(kk)] * Fp8ToF32(bp[col * k + kk]);
+          acc += part * (asp[i * k_tiles + kt] * bsp[nb * k_tiles + kt]);
+        }
+        StoreF32(out, i * n + col, acc);
+      }
+    }
+  });
+}
+
 // MatmulFp8Cutlass CPU kernel: out[m,n] = alpha * Sum_k f8val(a[m,k])*f8val(b[n,k]),
 // f32 accumulate, ONE folded alpha (= input_scale*weight_scale — our recorded
 // deviation from upstream's two epilogue scalars, see include/vt/ops.h).
@@ -3256,6 +3329,9 @@ struct Registrar {
                reinterpret_cast<void*>(static_cast<QuantFp8GroupFn>(&QuantFp8GroupKernel)));
     RegisterOp(OpId::kMatmulFp8Cutlass, DeviceType::kCPU,
                reinterpret_cast<void*>(static_cast<MatmulFp8CutlassFn>(&MatmulFp8CutlassKernel)));
+    RegisterOp(OpId::kMatmulFp8BlockScaled, DeviceType::kCPU,
+               reinterpret_cast<void*>(
+                   static_cast<MatmulFp8BlockScaledFn>(&MatmulFp8BlockScaledKernel)));
     RegisterOp(OpId::kSiluAndMul, DeviceType::kCPU,
                reinterpret_cast<void*>(static_cast<SiluAndMulFn>(&SiluAndMulKernel)));
     RegisterOp(OpId::kGeluAndMul, DeviceType::kCPU,

src/vt/op_provider.cpp

@@ -491,6 +491,8 @@ const char* OpNameImpl(OpId op) {
       return "AttentionRelPos";
     case OpId::kQuantFp8Group:
       return "QuantFp8Group";
+    case OpId::kMatmulFp8BlockScaled:
+      return "MatmulFp8BlockScaled";
     case OpId::kCount:
       break;
   }

src/vt/ops.cpp

@@ -673,6 +673,50 @@ void QuantFp8Group(Queue& q, Tensor& out_fp8, Tensor& out_scale, const Tensor& x
   reinterpret_cast<QuantFp8GroupFn>(GetOp(OpId::kQuantFp8Group, q.device.type))(
       q, out_fp8, out_scale, x, group_size);
 }
+void MatmulFp8BlockScaled(Queue& q, Tensor& out, const Tensor& a_fp8, const Tensor& a_scale,
+                          const Tensor& b_fp8, const Tensor& b_scale, int block_n,
+                          int block_k) {
+  VT_CHECK(out.rank == 2 && a_fp8.rank == 2 && a_scale.rank == 2 && b_fp8.rank == 2 &&
+               b_scale.rank == 2,
+           "matmul_fp8_block_scaled: out/a_fp8/a_scale/b_fp8/b_scale must be rank-2");
+  // Validated BEFORE either one divides anything: `x / 0` and `x % 0` are
+  // undefined behaviour, so a zero must refuse rather than trap.
+  VT_CHECK(block_n > 0 && block_k > 0,
+           "matmul_fp8_block_scaled: block_n and block_k must be positive");
+  const int64_t m = a_fp8.shape[0], k = a_fp8.shape[1];
+  const int64_t n = b_fp8.shape[0];
+  // Upstream: `assert A.shape[-1] == B.shape[-1]` (quant_utils.py:111).
+  VT_CHECK(b_fp8.shape[1] == k,
+           "matmul_fp8_block_scaled: a_fp8 [M,K] and b_fp8 [N,K] must share K");
+  VT_CHECK(out.shape[0] == m && out.shape[1] == n, "matmul_fp8_block_scaled: out must be [M,N]");
+  VT_CHECK(a_fp8.dtype == DType::kI8 && b_fp8.dtype == DType::kI8,
+           "matmul_fp8_block_scaled: a_fp8/b_fp8 must be i8 (raw fp8-e4m3fn bytes)");
+  // f32, not the model dtype: upstream refuses any other scale dtype on this
+  // path (csrc/.../w8a8/cutlass/c3x/scaled_mm_helper.hpp:15-18) and the
+  // accumulator these multiply into is f32.
+  VT_CHECK(a_scale.dtype == DType::kF32 && b_scale.dtype == DType::kF32,
+           "matmul_fp8_block_scaled: a_scale/b_scale must be f32");
+  VT_CHECK(out.dtype == DType::kF32 || out.dtype == DType::kBF16,
+           "matmul_fp8_block_scaled: out must be f32 or bf16");
+  // CEIL on every tiling, so a ragged final block is legal: upstream asserts
+  // `triton.cdiv(N, block_n) == Bs.shape[0]` and
+  // `triton.cdiv(K, block_k) == Bs.shape[1]` (fp8_utils.py:935-936), and
+  // `triton.cdiv(A.shape[-1], block_k) == As.shape[-1]` (fp8_utils.py:930).
+  const int64_t k_tiles = (k + block_k - 1) / block_k;
+  const int64_t n_tiles = (n + block_n - 1) / block_n;
+  VT_CHECK(a_scale.shape[0] == m && a_scale.shape[1] == k_tiles,
+           "matmul_fp8_block_scaled: a_scale must be [M, cdiv(K, block_k)]");
+  VT_CHECK(b_scale.shape[0] == n_tiles && b_scale.shape[1] == k_tiles,
+           "matmul_fp8_block_scaled: b_scale must be [cdiv(N, block_n), cdiv(K, block_k)]");
+  VT_CHECK(out.IsContiguous() && a_fp8.IsContiguous() && a_scale.IsContiguous() &&
+               b_fp8.IsContiguous() && b_scale.IsContiguous(),
+           "matmul_fp8_block_scaled: contiguous tensors required");
+  VT_CHECK(out.device == q.device && a_fp8.device == q.device && a_scale.device == q.device &&
+               b_fp8.device == q.device && b_scale.device == q.device,
+           "matmul_fp8_block_scaled: device mismatch (out/a_fp8/a_scale/b_fp8/b_scale/queue)");
+  reinterpret_cast<MatmulFp8BlockScaledFn>(GetOp(OpId::kMatmulFp8BlockScaled, q.device.type))(
+      q, out, a_fp8, a_scale, b_fp8, b_scale, block_n, block_k);
+}
 void RmsNormQuantFp8(Queue& q, Tensor& out_fp8, Tensor* out_bf16, const Tensor& x,
                      const Tensor& weight, const RmsNormArgs& args, Tensor* residual,
                      float input_scale) {

tests/CMakeLists.txt

@@ -1743,6 +1743,11 @@ vllm_cpp_add_test(test_ops_fp8_cpu vt/test_ops_fp8_cpu.cpp)
 # quant. CPU-gateable by construction; its CPU-vs-CUDA byte-identity arm (G6) is
 # CUDA-gated and reports PENDING rather than skipping where no device exists.
 vllm_cpp_add_test(test_ops_quant_fp8_group_cpu vt/test_ops_quant_fp8_group_cpu.cpp)
+# VT-MATMUL-FP8-BLOCK-REF (#1189 M2): the 128x128 block-scaled fp8 GEMM, CPU
+# reference arm — the numerical oracle every later block-fp8 kernel is measured
+# against. Its G4 holds the constraint the row exists for: the scales apply in
+# the MAINLOOP, once per K-block, which an epilogue-folded alpha cannot express.
+vllm_cpp_add_test(test_ops_matmul_fp8_block_cpu vt/test_ops_matmul_fp8_block_cpu.cpp)
 # Opt-in arm: run the fp8 plan-cache byte-exact case with the cache ENABLED
 # (VT_FP8_PLAN_CACHE=1 -> first MatmulFp8CublasLt call builds the plan fresh,
 # later calls hit the cache). Proves the cached-plan GEMM is BYTE-identical to the