Sf only per tok quant

csrc/nv_internal/cpp/kernels/quantization.cu

@@ -413,6 +413,80 @@ void invokeNvfp4QuantAndPerTokenScale(uint32_t m, uint32_t n, T const* input, fl
   });
 }
 
+void invokeScaleOnlyQuantAndPerTokenScale(uint32_t m, uint32_t n, float const* scaleInput,
+                                          float globalScaleInv, int32_t* expandedIdxToPermutedIdx,
+                                          uint8_t* scaleOutput, float* perTokenScaleOutput,
+                                          QuantizationSFLayout sfLayout, cudaStream_t stream) {
+  TLLM_CHECK_WITH_INFO(n % 16 == 0, "n must be a multiple of 16 for NVFP4 quantization");
+  constexpr uint32_t SF_VEC_SIZE = 16;
+
+  // Cache the per-block scales (num_sf_vecs_per_row floats) in shared memory so the quantization
+  // pass reads them back from SMEM instead of issuing a second global load of `scaleInput`.
+  uint32_t const num_sf_vecs_per_row = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE;
+  TLLM_CHECK_WITH_INFO(
+      num_sf_vecs_per_row % 4 == 0,
+      "num_sf_vecs_per_row must be a multiple of 4 for fp32 scale per-token quantization");
+
+  dispatchSFLayout(sfLayout, [&](auto sfLayoutTag) {
+    constexpr QuantizationSFLayout SF_LAYOUT = decltype(sfLayoutTag)::value;
+    constexpr uint32_t MAX_REGISTER_CACHED_SF_VECS = 512;
+    if (num_sf_vecs_per_row > 0 && num_sf_vecs_per_row % 4 == 0 &&
+        num_sf_vecs_per_row <= MAX_REGISTER_CACHED_SF_VECS) {
+#define LAUNCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_WARP_KERNEL(WARPS_PER_BLOCK, SCALES_PER_LOAD) \
+  {                                                                                               \
+    constexpr uint32_t NUM_WARPS = WARPS_PER_BLOCK;                                               \
+    uint32_t const num_blocks = (m + NUM_WARPS - 1) / NUM_WARPS;                                  \
+    scaleOnlyQuantAndPerTokenScaleWarpKernel<NUM_WARPS, SF_LAYOUT, SCALES_PER_LOAD>               \
+        <<<num_blocks, NUM_WARPS * 32, 0, stream>>>(m, n, scaleInput, globalScaleInv,             \
+                                                    expandedIdxToPermutedIdx, scaleOutput,        \
+                                                    perTokenScaleOutput);                         \
+  }
+#define DISPATCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_WARPS(SCALES_PER_LOAD)      \
+  if (m < 8) {                                                                    \
+    LAUNCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_WARP_KERNEL(1, SCALES_PER_LOAD);  \
+  } else if (m < 2048) {                                                          \
+    LAUNCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_WARP_KERNEL(4, SCALES_PER_LOAD);  \
+  } else if (m < 4096) {                                                          \
+    LAUNCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_WARP_KERNEL(8, SCALES_PER_LOAD);  \
+  } else if (m < 8192) {                                                          \
+    LAUNCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_WARP_KERNEL(16, SCALES_PER_LOAD); \
+  } else {                                                                        \
+    LAUNCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_WARP_KERNEL(32, SCALES_PER_LOAD); \
+  }
+      if (num_sf_vecs_per_row % 8 == 0) {
+        DISPATCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_WARPS(8);
+      } else {
+        DISPATCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_WARPS(4);
+      }
+#undef DISPATCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_WARPS
+#undef LAUNCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_WARP_KERNEL
+      return;
+    }
+    constexpr uint32_t BLOCK_SIZE = 128;
+    dim3 block(BLOCK_SIZE);
+    dim3 grid(m);
+    uint32_t smem_size = num_sf_vecs_per_row * sizeof(float);
+    // use an arbitrary large number to avoid exceeding shared memory limit.
+    TLLM_CHECK_WITH_INFO(n <= 32768,
+                         "n must be less than or equal to 32768 for NVFP4 quantization");
+#define LAUNCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_KERNEL(SCALES_PER_LOAD)            \
+  {                                                                                    \
+    scaleOnlyQuantAndPerTokenScaleKernel<BLOCK_SIZE, SF_LAYOUT, SCALES_PER_LOAD, true> \
+        <<<grid, block, smem_size, stream>>>(m, n, scaleInput, globalScaleInv,         \
+                                             expandedIdxToPermutedIdx, scaleOutput,    \
+                                             perTokenScaleOutput);                     \
+  }
+    if (num_sf_vecs_per_row % 8 == 0) {
+      LAUNCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_KERNEL(8);
+    } else if (num_sf_vecs_per_row % 4 == 0) {
+      LAUNCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_KERNEL(4);
+    } else {
+      TLLM_THROW("N must be divisible by 64 for NVFP4 quantization");
+    }
+#undef LAUNCH_SCALE_ONLY_QUANT_AND_PER_TOKEN_SCALE_KERNEL
+  });
+}
+
 // Instantiate the function.
 template void invokeNvfp4QuantAndPerTokenScale<float>(
     uint32_t m, uint32_t n, float const* input, float globalScaleInv,

csrc/nv_internal/tensorrt_llm/kernels/quantization.cuh

@@ -964,5 +964,261 @@ __global__ void nvfp4QuantAndPerTokenScaleFP32Kernel(
   }
 }
 
+// This kernel quantizes fp32 NVFP4 dequant scales to e4m3 and finds the row-wise amax.
+// `scaleInput` contains one fp32 dequant scale (block amax / E2M1 max) per SF_VEC_SIZE block in
+// linear layout. Reconstruct the block amax before deriving the per-token scale. The output layout
+// is linear, 128x4, or 8x4.
+template <uint32_t BLOCK_SIZE, QuantizationSFLayout SF_LAYOUT_OUT, int SCALES_PER_LOAD = 8,
+          bool CACHE_SCALE_IN_SMEM = true>
+__global__ void scaleOnlyQuantAndPerTokenScaleKernel(
+    // input
+    uint32_t m,
+    uint32_t n,  // linear shape of the tokens
+    float const* scaleInput, float globalScaleInv, int32_t* expandedIdxToPermutedIdx,
+    // output
+    uint8_t* scaleOutput, float* perTokenScaleOutput) {
+  static constexpr int SF_VEC_SIZE = 16;
+  static constexpr float E2M1_MAX_VALUE = 6.f;
+  static constexpr float E2M1_MAX_VALUE_INV = 1.f / 6.f;
+  static_assert(SCALES_PER_LOAD == 4 || SCALES_PER_LOAD == 8, "SCALES_PER_LOAD must be 4, or 8");
+  using ScaleVec = std::array<float, SCALES_PER_LOAD>;
+  int rowIdx = blockIdx.x;
+  if (rowIdx >= m) return;
+  if (expandedIdxToPermutedIdx != nullptr) {
+    rowIdx = expandedIdxToPermutedIdx[rowIdx];
+  }
+  if (rowIdx < 0) return;
+
+  // one input scale per SF_VEC_SIZE block, stored linearly per row
+  uint32_t num_sf_vecs_per_row = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE;
+  float const* rowScaleInput = scaleInput + static_cast<int64_t>(rowIdx) * num_sf_vecs_per_row;
+  uint32_t const num_scale_vecs_per_row = num_sf_vecs_per_row / SCALES_PER_LOAD;
+
+  extern __shared__ float scaleSmem[];
+
+  // scaleInput contains nonnegative block scales, so unsigned ordering of their IEEE-754 bits is
+  // identical to floating-point ordering. Reduce the bits first and reconstruct the amax once.
+  uint32_t globalAmaxBits = 0;
+  for (uint32_t scaleVecIdx = threadIdx.x; scaleVecIdx < num_scale_vecs_per_row;
+       scaleVecIdx += blockDim.x) {
+    ScaleVec scales;
+    loadPackedVec(scales, reinterpret_cast<ScaleVec const*>(rowScaleInput) + scaleVecIdx);
+#pragma unroll
+    for (int i = 0; i < SCALES_PER_LOAD; ++i) {
+      uint32_t vecIdx = scaleVecIdx * SCALES_PER_LOAD + i;
+      float localAmax = scales[i] * E2M1_MAX_VALUE;
+      if constexpr (CACHE_SCALE_IN_SMEM) {
+        scaleSmem[vecIdx] = localAmax;
+      }
+      globalAmaxBits = max(globalAmaxBits, __float_as_uint(scales[i]));
+    }
+  }
+  for (uint32_t vecIdx = num_scale_vecs_per_row * SCALES_PER_LOAD + threadIdx.x;
+       vecIdx < num_sf_vecs_per_row; vecIdx += blockDim.x) {
+    float localAmax = rowScaleInput[vecIdx] * E2M1_MAX_VALUE;
+    if constexpr (CACHE_SCALE_IN_SMEM) {
+      scaleSmem[vecIdx] = localAmax;
+    }
+    globalAmaxBits = max(globalAmaxBits, __float_as_uint(rowScaleInput[vecIdx]));
+  }
+  using BlockReduce = cub::BlockReduce<uint32_t, BLOCK_SIZE>;
+  __shared__ typename BlockReduce::TempStorage tempStorage;
+  __shared__ float perTokenScaleSmem;
+  globalAmaxBits = BlockReduce(tempStorage).Reduce(globalAmaxBits, cuda::maximum<>{});
+
+  // save the per-token scale
+  float globalAmax = __fmul_rn(__uint_as_float(globalAmaxBits), E2M1_MAX_VALUE);
+  float perTokenScale = globalAmax * globalScaleInv;
+  if (threadIdx.x == 0) {
+    perTokenScaleOutput[rowIdx] = perTokenScale;
+    perTokenScaleSmem = perTokenScale;
+  }
+  __syncthreads();
+  perTokenScale = perTokenScaleSmem;
+  float const perTokenScaleInv = reciprocal_approximate_ftz(perTokenScale);
+
+  int64_t sfOutputRowBase;
+  if constexpr (SF_LAYOUT_OUT == QuantizationSFLayout::LINEAR) {
+    sfOutputRowBase = static_cast<int64_t>(rowIdx) * num_sf_vecs_per_row;
+  } else if constexpr (SF_LAYOUT_OUT == QuantizationSFLayout::SWIZZLED_128x4) {
+    sfOutputRowBase = get_sf_out_offset_128x4(rowIdx, 0, num_sf_vecs_per_row);
+  } else {
+    sfOutputRowBase = get_sf_out_offset_8x4(rowIdx, 0, num_sf_vecs_per_row);
+  }
+  auto getSfOutputOffset = [&](uint32_t vecIdx) -> int64_t {
+    if constexpr (SF_LAYOUT_OUT == QuantizationSFLayout::LINEAR) {
+      return sfOutputRowBase + vecIdx;
+    } else if constexpr (SF_LAYOUT_OUT == QuantizationSFLayout::SWIZZLED_128x4) {
+      return get_sf_out_offset_128x4(rowIdx, vecIdx, num_sf_vecs_per_row);
+    } else {
+      return get_sf_out_offset_8x4(rowIdx, vecIdx, num_sf_vecs_per_row);
+    }
+  };
+
+  // quantize the per-block scales to e4m3 with the per-token scale
+  for (uint32_t scaleVecIdx = threadIdx.x; scaleVecIdx < num_scale_vecs_per_row;
+       scaleVecIdx += blockDim.x) {
+    ScaleVec scales;
+    if constexpr (!CACHE_SCALE_IN_SMEM) {
+      if constexpr (SCALES_PER_LOAD == 8 || SCALES_PER_LOAD == 4) {
+        loadPackedVec(scales, reinterpret_cast<ScaleVec const*>(rowScaleInput) + scaleVecIdx);
+      } else {
+        scales = reinterpret_cast<ScaleVec const*>(rowScaleInput)[scaleVecIdx];
+      }
+    }
+    if constexpr (SCALES_PER_LOAD >= 4) {
+#pragma unroll
+      for (int i = 0; i < SCALES_PER_LOAD; i += 4) {
+        uint32_t packedOutput = 0;
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+          uint32_t vecIdx = scaleVecIdx * SCALES_PER_LOAD + i + j;
+          float dequantScale;
+          if constexpr (CACHE_SCALE_IN_SMEM) {
+            dequantScale = scaleSmem[vecIdx] * E2M1_MAX_VALUE_INV;
+          } else {
+            dequantScale = scales[i + j];
+          }
+          float fp32Scale = dequantScale == 0.f ? 0 : dequantScale * perTokenScaleInv;
+          packedOutput |= static_cast<uint32_t>(__nv_fp8_e4m3(fp32Scale).__x) << (j * 8);
+        }
+        uint32_t vecIdx = scaleVecIdx * SCALES_PER_LOAD + i;
+        constexpr int64_t SF_GROUP_STRIDE =
+            SF_LAYOUT_OUT == QuantizationSFLayout::LINEAR
+                ? 4
+                : (SF_LAYOUT_OUT == QuantizationSFLayout::SWIZZLED_128x4 ? 512 : 32);
+        int64_t sfOffset = sfOutputRowBase + static_cast<int64_t>(vecIdx / 4) * SF_GROUP_STRIDE;
+        *reinterpret_cast<uint32_t*>(scaleOutput + sfOffset) = packedOutput;
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < SCALES_PER_LOAD; ++i) {
+        uint32_t vecIdx = scaleVecIdx * SCALES_PER_LOAD + i;
+        float dequantScale;
+        if constexpr (CACHE_SCALE_IN_SMEM) {
+          dequantScale = scaleSmem[vecIdx] * E2M1_MAX_VALUE_INV;
+        } else {
+          dequantScale = scales[i];
+        }
+        float fp32Scale = dequantScale == 0.f ? 0 : dequantScale * perTokenScaleInv;
+        scaleOutput[getSfOutputOffset(vecIdx)] = __nv_fp8_e4m3(fp32Scale).__x;
+      }
+    }
+  }
+  for (uint32_t vecIdx = num_scale_vecs_per_row * SCALES_PER_LOAD + threadIdx.x;
+       vecIdx < num_sf_vecs_per_row; vecIdx += blockDim.x) {
+    float dequantScale;
+    if constexpr (CACHE_SCALE_IN_SMEM) {
+      dequantScale = scaleSmem[vecIdx] * E2M1_MAX_VALUE_INV;
+    } else {
+      dequantScale = rowScaleInput[vecIdx];
+    }
+    float fp32Scale = dequantScale == 0.f ? 0 : dequantScale * perTokenScaleInv;
+    scaleOutput[getSfOutputOffset(vecIdx)] = __nv_fp8_e4m3(fp32Scale).__x;
+  }
+}
+
+// Runtime-width warp-per-row fast path. One warp owns one row and retains up to
+// MAX_SF_VECS_PER_ROW input scales in registers across the row reduction. Unlike an exact-width
+// specialization, this kernel accepts n at runtime and supports every output layout. The launcher
+// falls back to scaleOnlyQuantAndPerTokenScaleKernel for wider rows.
+template <uint32_t WARPS_PER_BLOCK, QuantizationSFLayout SF_LAYOUT_OUT, uint32_t SCALES_PER_LOAD,
+          uint32_t MAX_SF_VECS_PER_ROW = 512>
+__global__ __launch_bounds__(WARPS_PER_BLOCK * 32) void scaleOnlyQuantAndPerTokenScaleWarpKernel(
+    // input
+    uint32_t m, uint32_t n, float const* scaleInput, float globalScaleInv,
+    int32_t* expandedIdxToPermutedIdx,
+    // output
+    uint8_t* scaleOutput, float* perTokenScaleOutput) {
+  static constexpr int SF_VEC_SIZE = 16;
+  static_assert(SCALES_PER_LOAD == 4 || SCALES_PER_LOAD == 8, "SCALES_PER_LOAD must be 4 or 8");
+  static_assert(MAX_SF_VECS_PER_ROW % SCALES_PER_LOAD == 0,
+                "MAX_SF_VECS_PER_ROW must be divisible by SCALES_PER_LOAD");
+  static_assert(MAX_SF_VECS_PER_ROW % 128 == 0, "MAX_SF_VECS_PER_ROW must be a multiple of 128");
+  static constexpr int MAX_SCALE_VECS_PER_ROW = MAX_SF_VECS_PER_ROW / SCALES_PER_LOAD;
+  static constexpr int NUM_ROUNDS = (MAX_SCALE_VECS_PER_ROW + 31) / 32;
+  static constexpr float E2M1_MAX_VALUE = 6.f;
+  using ScaleVec = std::array<float, SCALES_PER_LOAD>;
+
+  int const warpIdx = threadIdx.x / warpSize;
+  int const laneIdx = threadIdx.x % warpSize;
+  int rowIdx = static_cast<int>(blockIdx.x) * WARPS_PER_BLOCK + warpIdx;
+  if (rowIdx >= m) return;
+  if (expandedIdxToPermutedIdx != nullptr) {
+    rowIdx = expandedIdxToPermutedIdx[rowIdx];
+  }
+  if (rowIdx < 0) return;
+
+  uint32_t const numSfVecsPerRow = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE;
+  uint32_t const numScaleVecsPerRow = numSfVecsPerRow / SCALES_PER_LOAD;
+  ScaleVec const* rowInput = reinterpret_cast<ScaleVec const*>(
+      scaleInput + static_cast<int64_t>(rowIdx) * numSfVecsPerRow);
+  ScaleVec values[NUM_ROUNDS];
+  // scaleInput contains nonnegative block scales, so unsigned ordering of their IEEE-754 bits is
+  // identical to floating-point ordering. On Ampere and newer this becomes one warp REDUX instead
+  // of five shuffle/FMAX stages.
+  uint32_t globalAmaxBits = 0;
+#pragma unroll
+  for (int i = 0; i < NUM_ROUNDS; ++i) {
+    int const scaleVecIdx = laneIdx + i * warpSize;
+    if (scaleVecIdx < numScaleVecsPerRow) {
+      loadPackedVec(values[i], rowInput + scaleVecIdx);
+#pragma unroll
+      for (int j = 0; j < SCALES_PER_LOAD; ++j) {
+        globalAmaxBits = max(globalAmaxBits, __float_as_uint(values[i][j]));
+      }
+    }
+  }
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+  globalAmaxBits = __reduce_max_sync(0xffffffff, globalAmaxBits);
+#else
+#pragma unroll
+  for (int offset = warpSize / 2; offset > 0; offset /= 2) {
+    globalAmaxBits =
+        max(globalAmaxBits, __shfl_xor_sync(0xffffffff, globalAmaxBits, offset, warpSize));
+  }
+#endif
+  float const globalAmax = __fmul_rn(__uint_as_float(globalAmaxBits), E2M1_MAX_VALUE);
+
+  float const perTokenScale = globalAmax * globalScaleInv;
+  if (laneIdx == 0) {
+    perTokenScaleOutput[rowIdx] = perTokenScale;
+  }
+  float const perTokenScaleInv = reciprocal_approximate_ftz(perTokenScale);
+
+  int64_t sfOutputRowBase;
+  if constexpr (SF_LAYOUT_OUT == QuantizationSFLayout::LINEAR) {
+    sfOutputRowBase = static_cast<int64_t>(rowIdx) * numSfVecsPerRow;
+  } else if constexpr (SF_LAYOUT_OUT == QuantizationSFLayout::SWIZZLED_128x4) {
+    sfOutputRowBase = get_sf_out_offset_128x4(rowIdx, 0, numSfVecsPerRow);
+  } else {
+    sfOutputRowBase = get_sf_out_offset_8x4(rowIdx, 0, numSfVecsPerRow);
+  }
+  constexpr int64_t SF_GROUP_STRIDE =
+      SF_LAYOUT_OUT == QuantizationSFLayout::LINEAR
+          ? 4
+          : (SF_LAYOUT_OUT == QuantizationSFLayout::SWIZZLED_128x4 ? 512 : 32);
+#pragma unroll
+  for (int i = 0; i < NUM_ROUNDS; ++i) {
+    int const scaleVecIdx = laneIdx + i * warpSize;
+    if (scaleVecIdx >= numScaleVecsPerRow) {
+      continue;
+    }
+#pragma unroll
+    for (int groupIdx = 0; groupIdx < SCALES_PER_LOAD; groupIdx += 4) {
+      uint32_t packedOutput = 0;
+#pragma unroll
+      for (int j = 0; j < 4; ++j) {
+        float const dequantScale = values[i][groupIdx + j];
+        float const fp32Scale = dequantScale == 0.f ? 0 : dequantScale * perTokenScaleInv;
+        packedOutput |= static_cast<uint32_t>(__nv_fp8_e4m3(fp32Scale).__x) << (j * 8);
+      }
+      uint32_t const vecIdx = scaleVecIdx * SCALES_PER_LOAD + groupIdx;
+      int64_t sfOffset = sfOutputRowBase + static_cast<int64_t>(vecIdx / 4) * SF_GROUP_STRIDE;
+      *reinterpret_cast<uint32_t*>(scaleOutput + sfOffset) = packedOutput;
+    }
+  }
+}
+
 }  // namespace kernels
 }  // namespace tensorrt_llm

csrc/nv_internal/tensorrt_llm/kernels/quantization.h

@@ -78,6 +78,14 @@ void invokeNvfp4QuantAndPerTokenScale(uint32_t m, uint32_t n, T const* input, fl
                                       uint8_t* scaleOutput, float* perTokenScaleOutput,
                                       QuantizationSFLayout sfLayout, cudaStream_t stream = 0);
 
+// The raw nvfp4 input is not available; the per-block amaxes are provided directly as fp32 scales
+// (one per SF_VEC_SIZE=16 group, linear layout). The row-wise amax is the max over those scales.
+void invokeScaleOnlyQuantAndPerTokenScale(uint32_t m, uint32_t n, float const* scaleInput,
+                                          float globalScaleInv,
+                                          int32_t* expanded_idx_to_permuted_idx,
+                                          uint8_t* scaleOutput, float* perTokenScaleOutput,
+                                          QuantizationSFLayout sfLayout, cudaStream_t stream = 0);
+
 template <typename T>
 void invokeBlockScaleInterleave(int b, int m, int m_padded, int n, int n_padded, T const* SFIn,
                                 T* SFOutput, int multiProcessorCount, cudaStream_t stream = 0);

csrc/trtllm_batched_gemm_runner.cu

@@ -117,14 +117,22 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(
         if (options.mFusedBiasShuffleMode != mOptions.fusedBiasShuffleMode) continue;
         if (options.mBiasDtype != mOptions.biasDtype) continue;
       }
-      if (mOptions.usePerTokenScaling) {
-        if (options.mTransposeMmaOutput && !options.mUsePerTokenSfB) continue;
-        if (!options.mTransposeMmaOutput && !options.mUsePerTokenSfA) continue;
-      }
+      bool const isNonStandardSfC = (options.mDtypeSfC != Dtype::Void &&
+                                     options.mDtypeSfC != dtypeGetBlockSfType(mOptions.dtypeC));
+      // if use non-standard sfC, force sf layout to linear
+      if (isNonStandardSfC && options.mSfLayoutC != SfLayout::Linear) continue;
       if (mOptions.usePerChannelScaling) {
         if (options.mTransposeMmaOutput && !options.mUsePerTokenSfA) continue;
         if (!options.mTransposeMmaOutput && !options.mUsePerTokenSfB) continue;
       }
+      if (mOptions.usePerTokenScaling) {
+        if (options.mTransposeMmaOutput && !options.mUsePerTokenSfB) continue;
+        if (!options.mTransposeMmaOutput && !options.mUsePerTokenSfA) continue;
+      } else {
+        // Avoid using the cubins with non-standard dtype sfc.
+        // They should only be selected for the fast per-token scaling case.
+        if (isNonStandardSfC) continue;
+      }
       if (options.mFusedAct) {
         if (options.mActType != static_cast<batchedGemm::gemmGatedAct::ActType>(mOptions.actType)) {
           continue;