This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2 

--- a/lib/Analysis/Utility.cpp
+++ b/lib/Analysis/Utility.cpp
@@ -1197,7 +1197,7 @@
     if (k < 256 / aElemTy.getIntOrFloatBitWidth())
       return false;
     if (!(retShapePerCTA[rank - 2] % 64 == 0 &&
-          retShapePerCTA[rank - 1] % 8 == 0))
+          retShapePerCTA[rank - 1] % 16 == 0))
       return false;
     if (aElemTy.isF64() || bElemTy.isF64() ||
         retType.getElementType().isF64()) {
@@ -1222,7 +1222,7 @@
     if (rank == 3)
       return false;
     if (!(numWarps % 4 == 0 && retShapePerCTA[rank - 2] % 64 == 0 &&
-          retShapePerCTA[rank - 1] % 8 == 0 &&
+          retShapePerCTA[rank - 1] % 16 == 0 &&
           (llvm::isa<Float8E5M2Type, Float8E4M3FNType>(aElemTy) ||
            aElemTy.isInteger(8) || aElemTy.isF16() || aElemTy.isBF16() ||
            aElemTy.isF32()))) {

--- a/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp
+++ b/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp
@@ -184,13 +184,21 @@
     auto kReg = StringAttr::get(ctx, "register");
     LinearLayout quot;
     int bestContig = 1;
+    bool isDim0 = false;
     for (int contig = 1; bitwidth * contig <= 32; contig *= 2) {
       auto maybeQuot = divideLeft(
           ll, LinearLayout::identity1D(contig, rowColDims[1], dims[1]));
+      bool currDim0 = false;
+      if (!maybeQuot) {
+        maybeQuot = divideLeft(
+            ll, LinearLayout::identity1D(contig, rowColDims[1], dims[0]));
+        currDim0 = true;
+      }
       if (!maybeQuot)
         break;
       quot = *maybeQuot;
       bestContig = contig;
+      isDim0 = currDim0;
     }
 
     // Pack contiguous elements
@@ -200,7 +208,7 @@
                                                  bitwidth * bestContig);
       if (!ret)
         return ret;
-      auto castbbitwidth = LinearLayout::identity1D(bestContig, kReg, dims[1]);
+      auto castbbitwidth = LinearLayout::identity1D(bestContig, kReg, isDim0 ? dims[0] : dims[1]);
       return castbbitwidth * ret.value();
     }
     if (auto maybeQuot = divideLeft(
@@ -241,10 +249,7 @@
 
   auto nColsTile = tile.getOutDimSize(rowColDims[1]);
   auto nColsLL = ll.getInDimSize(rowColDims[1]);
-  auto nColsMissing = nColsLL / nColsTile;
-  if (nColsMissing == 0) {
-    return std::nullopt;
-  }
+  auto nColsMissing = std::max<int>(nColsLL / nColsTile, 1);
   auto kReg = StringAttr::get(ctx, "register");
   auto kLane = StringAttr::get(ctx, "lane");
   auto kWarp = StringAttr::get(ctx, "warp");
@@ -261,7 +266,9 @@
   // basis as that would disallow the resulting layout to be left-divisible by
   // the tile
   auto trivialBlock = LinearLayout::identity1D(1, kBlock, kBlock);
-  auto comp = (tile * trivialBlock)
+  auto tileMult = tile * trivialBlock;
+  auto tileMultShrunk = tileMult.resizeOutDim(rowColDims[0], std::min<int>(ll.getInDimSize(rowColDims[0]), tileMult.getOutDimSize(rowColDims[0]))).resizeOutDim(rowColDims[1], std::min<int>(ll.getInDimSize(rowColDims[1]), tileMult.getOutDimSize(rowColDims[1])));
+  auto comp = tileMultShrunk
                   .compose(ll)
                   .sublayout({kReg, kLane}, to_vector(ll.getOutDimNames()));
   if (instr32Rows) {
@@ -269,7 +276,7 @@
     // basis
     comp = comp.resizeInDim(kLane, comp.getInDimSize(kLane) / 2);
   }
-  if (!comp.isInjective())
+  if (!comp.isInjective() && nColsLL > 32)
     return std::nullopt;
 
   // Fit the warp bases either tiling on the RHS or in row=16
@@ -324,9 +331,8 @@
   // Add CTAs as a trivial map
   auto nCTAs = ll.getInDimSize(kBlock);
   tile *= LinearLayout::identity1D(nCTAs, kBlock, kBlock);
-  assert(tile.getOutDimSize(rowColDims[1]) == ll.getInDimSize(rowColDims[1]));
-
-  auto ret = tile.compose(ll);
+  auto tileShrunk = tile.resizeOutDim(rowColDims[0], std::min<int>(ll.getInDimSize(rowColDims[0]), tile.getOutDimSize(rowColDims[0]))).resizeOutDim(rowColDims[1], std::min<int>(ll.getInDimSize(rowColDims[1]), tile.getOutDimSize(rowColDims[1])));
+  auto ret = tileShrunk.compose(ll);
   return ret;
 }
 
@@ -349,6 +355,18 @@
   auto *ctx = memType.getContext();
   auto layout = getDistributedLayoutForTmemLdSt(
       memType, TMemAccessAtom::I32x32b, numWarps);
+  if (!layout) {
+    layout = getDistributedLayoutForTmemLdSt(
+        memType, TMemAccessAtom::I16x256b, numWarps);
+  }
+  if (!layout) {
+    layout = getDistributedLayoutForTmemLdSt(
+        memType, TMemAccessAtom::I16x128b, numWarps);
+  }
+  if (!layout) {
+    layout = getDistributedLayoutForTmemLdSt(
+        memType, TMemAccessAtom::I16x64b, numWarps);
+  }
   assert(layout);
   return LinearEncodingAttr::get(ctx, std::move(*layout));
 }
@@ -422,6 +440,10 @@
 bool isDistributedLayoutTMemCompatible(Operation *op,
                                        RankedTensorType tensorType,
                                        gpu::MemDescType memType) {
+  for (auto l : getTmemCompatibleLayouts(op, tensorType, memType)) {
+    if (l == tensorType.getEncoding())
+      return true;
+  }
   auto maxnreg = getContextualMaxNReg(op);
   return succeeded(computeTMemLdStEncodingInfo(tensorType, memType, maxnreg));
 }

--- a/lib/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.cpp
+++ b/lib/Dialect/TritonNvidiaGPU/IR/TensorMemoryUtils.cpp
@@ -117,7 +117,7 @@
 }
 
 FailureOr<TMemLdStEncodingInfo>
-lowerTMemLdSt(const LinearLayout &cvt, int maxnreg, int bitwidth,
+lowerTMemLdSt(const LinearLayout &cvt, int maxnreg, int bitwidth, bool isScales,
               std::function<InFlightDiagnostic()> emitError,
               bool unpacked = false) {
   // We will fill in the returned value recursively (if it exists)
@@ -126,7 +126,7 @@
   auto removeBroadcastSrc = actionRemoveBroadcastedRegs(cvt);
   if (!removeBroadcastSrc.isIdentity())
     return lowerTMemLdSt(removeBroadcastSrc.apply(cvt), maxnreg, bitwidth,
-                         emitError, unpacked);
+                         isScales, emitError, unpacked);
   auto *ctx = cvt.getInDimNames().begin()->getContext();
   auto S = [ctx](StringRef str) { return StringAttr::get(ctx, str); };
   auto kReg = S("register");
@@ -139,6 +139,9 @@
     for (int contig = 1; bitwidth * contig <= 32; contig *= 2) {
       auto maybeQuot =
           divideLeft(cvt, LinearLayout::identity1D(contig, kReg, kCol));
+      if (!maybeQuot) {
+        maybeQuot = divideLeft(cvt, LinearLayout::identity1D(contig, kReg, kRow));
+      }
       if (!maybeQuot)
         break;
       quot = *maybeQuot;
@@ -182,7 +185,8 @@
     if (unpacked) {
       quot = LinearLayout::zeros1D(1, kReg, kCol, 32 / bitwidth) * quot;
     }
-    auto info = lowerTMemLdSt(quot, maxnreg, newBitwidth, emitError, unpacked);
+    auto info = lowerTMemLdSt(quot, maxnreg, newBitwidth, isScales, emitError,
+                              unpacked);
     if (failed(info))
       return failure();
     if (bestContig > 1) {
@@ -231,6 +235,19 @@
       auto row = reps.getBasis(kLane, 4, kRow);
       auto col = reps.getBasis(kLane, 4, kCol);
       secondHalfOffset = (row << 16) | col;
+      if (*secondHalfOffset == 0) {
+        // Workaround for ptxas bug, we cannot use secondHalfOffset = 0 to write
+        // only 16 elements. We use secondHalfOffset = 1 instead and we pad the
+        // allocation.
+        if (!isScales) {
+          if (emitError) {
+            emitError()
+                << "Only supported for scales as we pad the allocation.";
+          }
+          return failure();
+        }
+        secondHalfOffset = 1;
+      }
       // We "quotient it out", meaning we remove the last basis from reps
       auto basis = reps.getBases();
       basis[kLane][4] = {0, 0};
@@ -295,8 +312,9 @@
   cvt = LinearLayout(std::move(bases), cvt.getOutDims(),
                      /*isSurjective=*/cvt.isSurjective());
 
+  bool isScales = isa<TensorMemoryScalesEncodingAttr>(memTy.getEncoding());
   int bitwidth = memTy.getElementTypeBitWidth();
-  return lowerTMemLdSt(cvt, maxnreg, bitwidth, emitError);
+  return lowerTMemLdSt(cvt, maxnreg, bitwidth, isScales, emitError);
 }
 
 bool supportsTMemLoadReduce(RankedTensorType regTy, MemDescType memTy,

--- a/python/test/gluon/test_fpsan.py
+++ b/python/test/gluon/test_fpsan.py
@@ -1870,6 +1870,8 @@
 
 
 def _native_mma_k(type_a: str) -> int:
+    if is_cuda() and type_a == "f32":
+        return 16
     return 256 // _float_dtype_info(type_a)[0]
 
 

--- a/python/test/unit/language/test_compile_only.py
+++ b/python/test/unit/language/test_compile_only.py
@@ -213,7 +213,7 @@
     assert re.search(pattern, str(ttgir)), "The TTGIR does not match the expected pattern."
 
     ptx = k.asm["ptx"]
-    pattern = (r"tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32")
+    pattern = (r"tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X")
     assert re.search(pattern, str(ptx)), "The PTX does not match the expected pattern."
     assert k.asm["cubin"] != b""
 

--- a/python/test/unit/language/test_core.py
+++ b/python/test/unit/language/test_core.py
@@ -3673,10 +3673,7 @@
     if not is_cuda():
         return []
     return [(2, 4, 32, 1, False, False, 'None', 'ieee', 'float16', 'float32', 1, None),
-            (1, 2, 32, 1, False, False, 'None', 'ieee', 'float8e5', 'float32', 1, None),
-            # N=8: TF32 K=8 (wgmma.m64n8k8, sm90+) and FP16 K=16 (wgmma.m64n8k16)
-            (64, 8, 8, 4, False, False, 'None', 'tf32', 'float32', 'float32', 1, None),
-            (64, 8, 16, 4, False, False, 'None', 'ieee', 'float16', 'float32', 1, None)]
+            (1, 2, 32, 1, False, False, 'None', 'ieee', 'float8e5', 'float32', 1, None)]
 
 
 @pytest.mark.interpreter
@@ -3911,21 +3908,19 @@
     # make sure ld/st are vectorized
     ptx = pgm.asm['ptx']
 
-    # XXX: skip small sizes because they are not vectorized; with runtime
-    # strides, v4 needs the contiguous dim >= 16 (K for loads, N for stores).
-    enough_work = (M * N // (num_warps * 32) >= 4) and (K > 16 or N > 16 or M > 16)
-    if enough_work and K >= 16:
-        if 'float64' in in_dtype:
-            assert 'ld.global.v2.b64' in ptx
-        else:
-            assert 'ld.global.v4' in ptx
-    if enough_work and N >= 16:
-        if 'float8' in in_dtype:
-            assert 'st.global.v2' in ptx
-        elif 'float64' in in_dtype:
-            assert 'st.global.v2.b64' in ptx
-        else:
-            assert 'st.global.v4' in ptx
+    if (K > 16 or N > 16 or M > 16) and (M * N // (num_warps * 32) >= 4):
+        if K >= 16:
+            if 'float64' in in_dtype:
+                assert 'ld.global.v2.b64' in ptx
+            else:
+                assert 'ld.global.v4' in ptx
+        if N >= 16:
+            if 'float8' in in_dtype:
+                assert 'st.global.v2' in ptx
+            elif 'float64' in in_dtype:
+                assert 'st.global.v2.b64' in ptx
+            else:
+                assert 'st.global.v4' in ptx
 
     is_tcgen5 = (capability[0] == 10) and (num_warps % 4) == 0 and (M % 64) == 0 and (N % 8) == 0
 

--- a/python/test/unit/language/test_matmul.py
+++ b/python/test/unit/language/test_matmul.py
@@ -402,6 +402,8 @@
             pytest.skip(f"CDNA4 does not support {BLOCK_K=} for scaled mfma {nonKDim=} variants")
         if (BLOCK_M == 256 or BLOCK_N == 256) and BLOCK_K == 256:
             pytest.skip("Config requires too much shared memory")
+    if is_cuda() and torch.cuda.get_device_capability()[0] >= 10 and BLOCK_N < 32:
+        pytest.skip("Blackwell TMEM scaling instructions require aligned BLOCK_N >= 32")
 
     if not is_rubin() and BLOCK_N == 256 and BLOCK_K == 256:
         NUM_STAGES = min(NUM_STAGES, 2)

--- a/python/tutorials/09-persistent-matmul.py
+++ b/python/tutorials/09-persistent-matmul.py
@@ -62,7 +62,7 @@
 
 
 def is_hopper():
-    return torch.cuda.get_device_capability()[0] == 9
+    return torch.cuda.get_device_capability()[0] >= 9
 
 
 def supports_ws():

--- a/test/Conversion/tritongpu_to_llvm_blackwell.mlir
+++ b/test/Conversion/tritongpu_to_llvm_blackwell.mlir
@@ -302,7 +302,7 @@
   // CHECK-LABEL: @tc_gen5_mma_scaled_a_tmem
   // CHECK: %[[TMEM_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
   // CHECK: %[[A_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32
-  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 0 ], $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,r,l,r,r,r,b,b" %[[TMEM_BASE]], %[[A_BASE]]
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 0 ], $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,r,l,r,r,r,b,b" %[[TMEM_BASE]], %[[A_BASE]]
   tt.func @tc_gen5_mma_scaled_a_tmem(
       %a: !ttg.memdesc<128x256xf8E5M2, #tmem, #ttng.tensor_memory>,
       %b: !ttg.memdesc<256x64xf8E5M2, #shared, #ttg.shared_memory>,
@@ -337,10 +337,10 @@
   // CHECK: %[[P1:.+]] = llvm.and %{{.*}}, %[[P0]]  : i1
   // CHECK: llvm.cond_br %[[P1]]
   // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(144708608 : i32) : i32
-  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]], %{{.+}}, %{{.+}}, %arg5
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]], %{{.+}}, %{{.+}}, %arg5
   // CHECK: %[[TRUE:.+]] = llvm.mlir.constant(true) : i1
   // CHECK: %[[DESC1:.+]] = llvm.mlir.constant(681579536 : i32) : i32
-  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC1]], %{{.+}}, %{{.+}}, %[[TRUE]]
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC1]], %{{.+}}, %{{.+}}, %[[TRUE]]
   tt.func @tc_gen5_mma_block_scale(%a: !ttg.memdesc<128x64xf8E4M3FN, #shared, #ttg.shared_memory>,
                        %b: !ttg.memdesc<32x128xi8, #shared1, #ttg.shared_memory>,
                        %c: !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>,
@@ -371,13 +371,13 @@
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
   // CHECK-LABEL: @tc_gen5_mma_block_scale_fp4_a
   // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(144769664 : i32) : i32
-  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC0]]
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC0]]
   // CHECK: %[[DESC1:.+]] = llvm.mlir.constant(681640592 : i32) : i32
-  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC1]]
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC1]]
   // CHECK: %[[DESC2:.+]] = llvm.mlir.constant(1218511520 : i32) : i32
-  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC2]]
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC2]]
   // CHECK: %[[DESC3:.+]] = llvm.mlir.constant(1755382448 : i32) : i32
-  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC3]]
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC3]]
   tt.func @tc_gen5_mma_block_scale_fp4_a(%a: !ttg.memdesc<128x64xi8, #shared1, #ttg.shared_memory>,
                        %b: !ttg.memdesc<128x128xi8, #shared, #ttg.shared_memory>,
                        %c: !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>,
@@ -540,8 +540,8 @@
   // CHECK-LABEL: @tc_gen5_mma_block_scale_nvfp4
   // CHECK: %[[TMEM_BASE:.+]] = llvm.ptrtoint %{{.*}} : !llvm.ptr<3> to i32
   // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(138413184 : i32) : i32
-  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]]
-  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]]
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]]
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]]
   tt.func @tc_gen5_mma_block_scale_nvfp4(%a: !ttg.memdesc<128x64xi8, #shared, #ttg.shared_memory>,
                        %b: !ttg.memdesc<64x256xi8, #shared1, #ttg.shared_memory>,
                        %c: !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>,
@@ -573,9 +573,9 @@
   // CHECK-LABEL: @tc_gen5_mma_block_scale_mxfp4
   // CHECK-DAG: %[[TMEM_BASE:.+]] = llvm.ptrtoint %{{.*}} : !llvm.ptr<3> to i32
   // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(146801792 : i32) : i32
-  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]]
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]]
   // CHECK: %[[DESC1:.+]] = llvm.mlir.constant(1220543648 : i32) : i32
-  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC1]]
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC1]]
   tt.func @tc_gen5_mma_block_scale_mxfp4(%a: !ttg.memdesc<128x64xi8, #shared, #ttg.shared_memory>,
                        %b: !ttg.memdesc<64x256xi8, #shared1, #ttg.shared_memory>,
                        %c: !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>,
@@ -1099,7 +1099,7 @@
 // CHECK-LABEL: @load_store_16x32bx1_broadcast
 tt.func private @load_store_16x32bx1_broadcast(%arg0: !ttg.memdesc<16x8xi8, #tmem_scales, #ttng.tensor_memory, mutable>, %arg1: tensor<16x8xi8, #linear>) {
   %true = arith.constant true
-  // CHECK: @$0 tcgen05.st.sync.aligned.16x32bx2.x1.b32 [$1 + 0], 0, {$2}
+  // CHECK: @$0 tcgen05.st.sync.aligned.16x32bx2.x1.b32 [$1 + 0], 1, {$2}
   ttng.tmem_store %arg1, %arg0, %true : tensor<16x8xi8, #linear> -> !ttg.memdesc<16x8xi8, #tmem_scales, #ttng.tensor_memory, mutable>
   tt.return
 }
@@ -1114,7 +1114,7 @@
 // CHECK-LABEL: @load_store_16x2_scales_uses_zero_second_half
 tt.func private @load_store_16x2_scales_uses_zero_second_half(%arg0: !ttg.memdesc<16x2xi8, #tmem_scales, #ttng.tensor_memory, mutable>, %arg1: tensor<16x2xi8, #linear>) {
   %true = arith.constant true
-  // CHECK: @$0 tcgen05.st.sync.aligned.16x32bx2.x1.b32 [$1 + 0], 0, {$2}
+  // CHECK: @$0 tcgen05.st.sync.aligned.16x32bx2.x1.b32 [$1 + 0], 1, {$2}
   // CHECK: llvm.return
   ttng.tmem_store %arg1, %arg0, %true : tensor<16x2xi8, #linear> -> !ttg.memdesc<16x2xi8, #tmem_scales, #ttng.tensor_memory, mutable>
   tt.return
@@ -1178,22 +1178,22 @@
       %pred: i1) {
     // CHECK: %[[ACC_BASE:.+]] = llvm.ptrtoint %arg4 : !llvm.ptr<3> to i32
     // CHECK: %[[A_E4M3_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 0 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 0 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_E4M3_BASE]]
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 8 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 8 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_E4M3_BASE]]
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 16 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 16 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_E4M3_BASE]]
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 24 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 24 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_E4M3_BASE]]
     // CHECK: %[[A_E5M2_BASE:.+]] = llvm.ptrtoint %arg1 : !llvm.ptr<3> to i32
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 0 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 0 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_E5M2_BASE]]
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 8 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 8 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_E5M2_BASE]]
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 16 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 16 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_E5M2_BASE]]
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 24 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 24 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_E5M2_BASE]]
     ttng.tc_gen5_mma_scaled %a_e4m3, %b_e5m2, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e4m3 rhs = e5m2 :
        !ttg.memdesc<128x128xf8E4M3FN, #tmem_fp8_lhs, #ttng.tensor_memory>,
@@ -1221,13 +1221,13 @@
       %pred: i1) {
     // CHECK-DAG: %[[ACC_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
     // CHECK-DAG: %[[A_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 0 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 0 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]]
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 8 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 8 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]]
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 16 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 16 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]]
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], [ $1 + 24 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], [ $1 + 24 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]]
     ttng.tc_gen5_mma_scaled %a, %b, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e2m1 rhs = e5m2 :
        !ttg.memdesc<128x64xi8, #tmem_fp4_padded_lhs, #ttng.tensor_memory>,
@@ -1257,13 +1257,13 @@
       %pred: i1) {
     // CHECK-DAG: %[[ACC_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
     // CHECK-DAG: %[[A_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], [ $1 + 0 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], [ $1 + 0 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]]
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], [ $1 + 8 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], [ $1 + 8 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]]
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], [ $1 + 16 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], [ $1 + 16 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]]
-    // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], [ $1 + 24 ], $2
+    // CHECK: tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], [ $1 + 24 ], $2
     // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]]
     ttng.tc_gen5_mma_scaled %a, %b, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e2m1 rhs = e2m1 :
        !ttg.memdesc<128x128xi8, #tmem_fp4_dense_lhs, #ttng.tensor_memory>,

--- a/test/Conversion/tritongpu_to_llvm_rubin.mlir
+++ b/test/Conversion/tritongpu_to_llvm_rubin.mlir
@@ -8,7 +8,7 @@
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:107"} {
   // CHECK-LABEL: @tc_gen5_mma_block_scale_fp8_sm107
   // CHECK: %[[TMEM_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
-  // CHECK-COUNT-2: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32
+  // CHECK-COUNT-2: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X
   // CHECK-NOT: tcgen05.mma
   tt.func @tc_gen5_mma_block_scale_fp8_sm107(%a: !ttg.memdesc<128x128xf8E4M3FN, #shared, #ttg.shared_memory>,
                        %b: !ttg.memdesc<64x128xi8, #shared1, #ttg.shared_memory>,
@@ -41,7 +41,7 @@
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:107"} {
   // CHECK-LABEL: @tc_gen5_mma_block_scale_nvfp4_sm107_m256
   // CHECK: %[[DESC:.+]] = llvm.mlir.constant(136316040 : i32) : i32
-  // CHECK-COUNT-4: tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block32
+  // CHECK-COUNT-4: tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::2X
   // CHECK-NOT: tcgen05.mma
   tt.func @tc_gen5_mma_block_scale_nvfp4_sm107_m256(%a: !ttg.memdesc<256x128xi8, #shared, #ttg.shared_memory>,
                        %b: !ttg.memdesc<128x128xi8, #shared1, #ttg.shared_memory>,
@@ -103,7 +103,7 @@
 #tmem_scales1 = #ttng.tensor_memory_scales_encoding<CGALayout = [[1, 0]]>
 module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:107", "ttg.threads-per-warp" = 32 : i32, "ttng.two-ctas" = true} {
   // CHECK-LABEL: @tc_gen5_mma_scaled_fp4_padded_shared_linear_a
-  // CHECK-COUNT-4: tcgen05.mma.cta_group::2.kind::mxf8f6f4.block_scale.block32
+  // CHECK-COUNT-4: tcgen05.mma.cta_group::2.kind::mxf8f6f4.block_scale.scale_vec::1X
   // CHECK-NOT: tcgen05.mma
   tt.func @tc_gen5_mma_scaled_fp4_padded_shared_linear_a(
       %a: !ttg.memdesc<256x64xi8, #shared_fp4_padded_sl, #smem, mutable>,
@@ -158,8 +158,8 @@
 #tmem_scales = #ttng.tensor_memory_scales_encoding<>
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:107"} {
   // CHECK-LABEL: @tc_gen5_mma_block_scale_breuse
-  // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32.collector::b::fill
-  // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32.collector::b::lastuse
+  // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X.collector::b::fill
+  // CHECK: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X.collector::b::lastuse
   tt.func @tc_gen5_mma_block_scale_breuse(%a: !ttg.memdesc<256x64xf8E4M3FN, #shared, #ttg.shared_memory>,
                        %b: !ttg.memdesc<32x128xi8, #shared1, #ttg.shared_memory>,
                        %c: !ttg.memdesc<256x128xf32, #tmem, #ttng.tensor_memory, mutable>,
@@ -195,7 +195,7 @@
   // scale kind in the descriptor for both K=128 sub-instructions of a
   // logical BLOCK_K=256 tile.
   // CHECK: %[[DESC:.+]] = llvm.mlir.constant(153093256 : i32) : i32
-  // CHECK-COUNT-2: tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16
+  // CHECK-COUNT-2: tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X
   // CHECK-NOT: tcgen05.mma
   tt.func @tc_gen5_mma_block_scale_nvfp4_ue5m3_sm107(%a: !ttg.memdesc<128x128xi8, #shared, #ttg.shared_memory>,
                        %b: !ttg.memdesc<128x128xi8, #shared1, #ttg.shared_memory>,

--- a/test/TritonGPU/accelerate-matmul.mlir
+++ b/test/TritonGPU/accelerate-matmul.mlir
@@ -1052,40 +1052,6 @@
 
 // -----
 
-// Verify TF32 dot with N=8, K=8 (native WGMMA tile) selects MMAv3 on sm90.
-// CHECK: #[[$MMA:.+]] = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [4, 1], instrShape = [16, 8, 8]}>
-#blocked_tf32_n8 = #ttg.blocked<{sizePerThread = [2, 2], threadsPerWarp = [8, 4], warpsPerCTA = [4, 1], order = [1, 0]}>
-module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
-  // CHECK-LABEL: tf32_n8k8_sm90
-  tt.func public @tf32_n8k8_sm90(
-      %a: tensor<64x8xf32, #ttg.dot_op<{opIdx = 0, parent = #blocked_tf32_n8}>>,
-      %b: tensor<8x8xf32, #ttg.dot_op<{opIdx = 1, parent = #blocked_tf32_n8}>>) -> tensor<64x8xf32, #blocked_tf32_n8> {
-    %cst = arith.constant dense<0.000000e+00> : tensor<64x8xf32, #blocked_tf32_n8>
-    // CHECK: ttng.warp_group_dot {{.*}} -> tensor<64x8xf32, #[[$MMA]]>
-    %d = tt.dot %a, %b, %cst, inputPrecision = tf32 : tensor<64x8xf32, #ttg.dot_op<{opIdx = 0, parent = #blocked_tf32_n8}>> * tensor<8x8xf32, #ttg.dot_op<{opIdx = 1, parent = #blocked_tf32_n8}>> -> tensor<64x8xf32, #blocked_tf32_n8>
-    tt.return %d : tensor<64x8xf32, #blocked_tf32_n8>
-  }
-}
-
-// -----
-
-// Verify FP16 dot with N=8, K=16 (native WGMMA tile) selects MMAv3 on sm90.
-// CHECK: #[[$MMA:.+]] = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [4, 1], instrShape = [16, 8, 16]}>
-#blocked_fp16_n8 = #ttg.blocked<{sizePerThread = [2, 2], threadsPerWarp = [8, 4], warpsPerCTA = [4, 1], order = [1, 0]}>
-module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
-  // CHECK-LABEL: fp16_n8k16_sm90
-  tt.func public @fp16_n8k16_sm90(
-      %a: tensor<64x16xf16, #ttg.dot_op<{opIdx = 0, parent = #blocked_fp16_n8}>>,
-      %b: tensor<16x8xf16, #ttg.dot_op<{opIdx = 1, parent = #blocked_fp16_n8}>>) -> tensor<64x8xf32, #blocked_fp16_n8> {
-    %cst = arith.constant dense<0.000000e+00> : tensor<64x8xf32, #blocked_fp16_n8>
-    // CHECK: ttng.warp_group_dot {{.*}} -> tensor<64x8xf32, #[[$MMA]]>
-    %d = tt.dot %a, %b, %cst : tensor<64x16xf16, #ttg.dot_op<{opIdx = 0, parent = #blocked_fp16_n8}>> * tensor<16x8xf16, #ttg.dot_op<{opIdx = 1, parent = #blocked_fp16_n8}>> -> tensor<64x8xf32, #blocked_fp16_n8>
-    tt.return %d : tensor<64x8xf32, #blocked_fp16_n8>
-  }
-}
-
-// -----
-
 #blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [4, 1], order = [0, 1]}>
 #blocked1 = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [16, 2], warpsPerCTA = [4, 1], order = [1, 0]}>
 #blocked2 = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [2, 16], warpsPerCTA = [4, 1], order = [1, 0]}>

--- a/third_party/nvidia/backend/compiler.py
+++ b/third_party/nvidia/backend/compiler.py
@@ -29,8 +29,6 @@
             return (1, 1, 32)
         elif lhs_bitwidth == 64:
             return (1, 1, 4)
-        elif lhs_bitwidth == 32:
-            return (1, 1, 8)
         else:
             return (1, 1, 16)
 

--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp
+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp
@@ -398,12 +398,12 @@
   std::string opcode =
       "tcgen05.mma.cta_group::" + std::to_string(twoCTAs ? 2 : 1) + ".kind::";
   if (mxfpInstKind == mxfpKind::mxf8f6f4) {
-    opcode += "mxf8f6f4.block_scale.block32";
+    opcode += "mxf8f6f4.block_scale.scale_vec::1X";
   } else if (mxfpInstKind == mxfpKind::mxf4) {
-    opcode += "mxf4.block_scale.block32";
+    opcode += "mxf4.block_scale.scale_vec::2X";
   } else if (mxfpInstKind == mxfpKind::mxf4nvf4) {
-    opcode += getScaleVecSize(op) == 32 ? "mxf4nvf4.block_scale.block32"
-                                        : "mxf4nvf4.block_scale.block16";
+    opcode += getScaleVecSize(op) == 32 ? "mxf4nvf4.block_scale.scale_vec::2X"
+                                        : "mxf4nvf4.block_scale.scale_vec::4X";
   } else {
     assert(0 && "Unsupported mxfp kind.");
   }
@@ -881,15 +881,18 @@
     int scaleFactorColsPerSet =
         getScaleFactorColsPerSet(mxfpInstKind, op, dot.mmaSizeK);
     int numRepKWords = ceil<int>(numRepK, 4 / scaleFactorColsPerSet);
+    int aMultibuffers = cast<MemDescType>(op.getAScale().getType()).getRank() > 2 ? cast<MemDescType>(op.getAScale().getType()).getDimSize(0) : 1;
+    int bMultibuffers = cast<MemDescType>(op.getBScale().getType()).getRank() > 2 ? cast<MemDescType>(op.getBScale().getType()).getDimSize(0) : 1;
     int numColPerScaleBlockA = ceil<int>(
         ttng::getTmemAllocSizes(cast<MemDescType>(op.getAScale().getType()))
-            .numCols,
+            .numCols / aMultibuffers,
         numRepM * numRepKWords);
     int numColPerScaleBlockB = ceil<int>(
         ttng::getTmemAllocSizes(cast<MemDescType>(op.getBScale().getType()))
-            .numCols,
+            .numCols / bMultibuffers,
         numRepN * numRepKWords);
-    numColPerScaleBlockB = std::max(numColPerScaleBlockB, 2);
+    if (desc.mmaSizeN >= 32)
+      numColPerScaleBlockB = std::max(numColPerScaleBlockB, 2);
     int subWordIdx = k % (4 / scaleFactorColsPerSet);
     int wordIdx = k / (4 / scaleFactorColsPerSet);
     int scaleIdxA = linearizeScaleBlockIdx(op.getAScale(), m, wordIdx, numRepM,

