[CUDA] Enable native SM90, block_size=32, and fused bias for fpA_intB MatMulNBits

docs/ContribOperators.md

@@ -3159,7 +3159,7 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>block_size</tt> : int (required)</dt>
 <dd>Size of each quantization block along the K (input feature) dimension. Must be a power of two and ≥ 16 (e.g., 16, 32, 64, 128).</dd>
 <dt><tt>weight_prepacked</tt> : int</dt>
-<dd>If set, input B is already prepacked into an EP-specific layout and the EP skips runtime weight prepacking. 0 (default): not prepacked. 1: prepacked in the CUDA SM80 fpA_intB layout. 2: reserved for a future SM90 layout (currently rejected at kernel construction).</dd>
+<dd>If set, input B is already prepacked into an EP-specific layout and the EP skips runtime weight prepacking. 0 (default): not prepacked. 1: prepacked in the CUDA SM80 fpA_intB layout. 2: prepacked in the CUDA SM90 (Hopper) fpA_intB layout, consumed by the native SM90 kernel (requires a compute capability 9.0 device and block_size in {64, 128}).</dd>
 </dl>
 
 #### Inputs (3 - 6)

docs/contrib_ops/cuda/matmul_nbits.md

@@ -44,7 +44,7 @@ Source files:
 | `bits` | Quantization bit width: `4` or `8`. |
 | `block_size` | Quantization group size along `K` (16 / 32 / 64 / 128). One scale (and optional zero point) per group. |
 | `accuracy_level` | Minimum accuracy level for internal handling of `A`; default `0` means unset. |
-| `weight_prepacked` | CUDA fpA_intB weight-layout selector. `0` (default): `B` is in standard MatMulNBits layout and may be runtime-prepacked. `1`: `B` is already prepacked in the CUDA SM80 fpA_intB layout. `2`: reserved SM90 layout; currently rejected. |
+| `weight_prepacked` | CUDA fpA_intB weight-layout selector. `0` (default): `B` is in standard MatMulNBits layout and may be runtime-prepacked. `1`: `B` is already prepacked in the CUDA SM80 fpA_intB layout. `2`: `B` is prepacked in the CUDA SM90 (Hopper) fpA_intB layout, consumed by the native SM90 kernel (requires an SM90 device and `block_size` in {64, 128}). |
 
 | Input | Index | Notes |
 |-------|-------|-------|
@@ -95,13 +95,18 @@ prepacked_flat = _pybind.pack_weights_for_cuda_mixed_gemm(
 prepacked_b = np.asarray(prepacked_flat, dtype=np.int8).view(np.uint8).reshape(q_weight.shape)
 ```
 
-The final argument is the target packing architecture. For MatMulNBits v1, use
-`80`: both runtime preprocessing and offline preprocessing force the SM80
-layout. On SM90 devices, the supported mixed FP16/BF16 activation + int4/int8
-weight path routes to the SM80 CUTLASS kernel/layout for this operator.
+The final argument is the target packing architecture. Use `80` for the SM80
+layout (consumed by the SM80 CUTLASS kernel, including on newer GPUs via the
+compatibility path) and set `weight_prepacked=1` on the node. Use `90` for the
+native SM90 (Hopper) layout and set `weight_prepacked=2` on the node.
 
-`weight_prepacked=2` is reserved for a future SM90/Hopper-specific layout and is
-currently rejected during kernel construction.
+`weight_prepacked=2` selects the native SM90 (Hopper TMA/WGMMA) mixed-GEMM
+kernel and its Hopper weight layout. It requires a compute capability 9.0 device
+and `block_size` in `{64, 128}` (the SM90 kernel needs `group_size` to be a
+multiple of the 64-element Hopper K tile, so `block_size=32` is SM80-only). On
+SM90 devices, runtime-prepacked (`weight_prepacked=0`) and SM80-prepacked
+(`weight_prepacked=1`) weights continue to route to the SM80 CUTLASS
+kernel/layout.
 
 ---
 
@@ -112,7 +117,7 @@ falls through to progressively more general ones:
 
 ```mermaid
 flowchart TD
-  A[ComputeInternal] --> F{has_fpA_intB_gemm_?<br/>FP16/BF16, ORT-prepackable weights,<br/>block 64/128, sm>=75}
+  A[ComputeInternal] --> F{has_fpA_intB_gemm_?<br/>FP16/BF16, ORT-prepackable weights,<br/>block 32/64/128, sm>=75}
   F -- yes --> FP[fpA_intB CUDA GEMV<br/>or CUTLASS grouped GEMM] --> R[return]
   F -- no --> G{reorder_idx == null<br/>and zero_points not typed-T?}
   G -- no --> DQ
@@ -236,11 +241,15 @@ and enabled via `ORT_FPA_INTB_GEMM`, FP16/BF16 MatMulNBits can use the
 TensorRT-LLM-derived CUTLASS weight-only kernels. The constructor sets
 `has_fpA_intB_gemm_` only when:
 
-- dtype is FP16 or BF16, `bits ∈ {4, 8}`, `block_size ∈ {64, 128}`,
-- no `g_idx`, no `bias`, `N % (bits==8 ? 32 : 64) == 0`, `K % block_size == 0`,
+- dtype is FP16 or BF16, `bits ∈ {4, 8}`, `block_size ∈ {32, 64, 128}`,
+- no `g_idx`, `N % (bits==8 ? 32 : 64) == 0`, `K % block_size == 0`,
 - `sm_ >= 75`, and weight/scale/zero-point inputs are constant initializers that
   ORT can prepack.
 
+`block_size=32` is served by the SM80/Ampere-class fine-grained kernel (and its
+SM90 compatibility path); the native SM90 kernel (`weight_prepacked=2`) supports
+only `block_size ∈ {64, 128}` — see §2.1.
+
 At run time a profiler picks the best tactic; small `M` may use a dedicated CUDA
 GEMV kernel (`bestTactic->enableCudaKernel`), otherwise a CUTLASS grouped GEMM.
 This path takes precedence over everything in §3 when active.
@@ -258,8 +267,10 @@ Prepacked weights are intentionally strict:
   throws instead of silently falling back to a raw-layout path.
 - Nonzero `weight_prepacked` requires FP16 or BF16 input `A`, because only the
   CUDA fpA_intB path consumes this layout.
-- `weight_prepacked=1` must match the currently required SM80 layout; `2` is
-  reserved and rejected.
+- `weight_prepacked` must match the layout the selected kernel expects: `1` is
+  the SM80 layout, `2` is the native SM90 (Hopper) layout. `2` additionally
+  requires a compute-capability 9.0 device and `block_size ∈ {64, 128}` and is
+  rejected otherwise.
 
 ---
 

onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h

@@ -247,7 +247,7 @@ struct GemmFpAIntB {
     }
 
     if constexpr (isFinegrained(Mma::QuantOp)) {
-      if (args.group_size != 64 && args.group_size != 128) {
+      if (args.group_size != 32 && args.group_size != 64 && args.group_size != 128) {
         return Status::kErrorNotSupported;
       }
     }
@@ -442,11 +442,16 @@ struct GemmFpAIntB {
     run_kernel<arch::Sm80>(params, shared_storage);
 #elif (__CUDA_ARCH__ == 890)
     run_kernel<arch::Sm89>(params, shared_storage);
-#elif (__CUDA_ARCH__ >= 1000)
-    // Use SM80 implementation for GB10x, GB20x.
+#elif (__CUDA_ARCH__ >= 900)
+    // Reuse the SM80 (Ampere) mixed-GEMM kernel on Hopper (SM90) and Blackwell (SM100+, e.g.
+    // GB10x/GB20x). The Ampere tensor-core mixed-input path is valid on these archs, so callers
+    // that pack the SM80 weight layout and dispatch KernelArch=Sm80 (see MatMulNBits) get a real
+    // kernel here instead of a stub. The dedicated CUTLASS 3.x Hopper TMA warp-specialized
+    // mixed-input kernels are reached through the separate sm90_dispatch path, not this operator.
+    // Mirrors MoeFCGemm::operator() which already reuses the SM80 kernel for SM90+.
     run_kernel<arch::Sm80>(params, shared_storage);
 #else
-    CUTLASS_NOT_IMPLEMENTED();  // Don't compile these for Hopper or later. Use CUTLASS 3.x kernels.
+    CUTLASS_NOT_IMPLEMENTED();
 #endif
 #else
     CUTLASS_NOT_IMPLEMENTED();

onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm/fpA_intB_gemm.h

@@ -74,6 +74,19 @@ class CutlassFpAIntBGemmRunnerInterface {
 
   virtual std::vector<tkc::CutlassGemmConfig> getConfigs() const = 0;
 
+  // Overrides the SM architecture used for tactic/config enumeration, workspace sizing and kernel
+  // dispatch. By default the runner targets the detected device SM. On SM90 the half/bf16
+  // weight-only path dispatches the SM80 (Ampere) mixed-GEMM kernel (which now runs on Hopper, see
+  // GemmFpAIntB::operator()), so MatMulNBits forces the runner to SM80 to keep the enumerated
+  // tactics (tile_config_sm80) and workspace sizing consistent with the dispatched kernel.
+  // Default: no-op (keep detected SM).
+  virtual void setArch(int /*sm*/) {}
+
+  // Opts in to the native SM90 (Hopper TMA/WGMMA) mixed-GEMM kernel instead of the SM80
+  // compatibility path. Only meaningful when the runner targets SM90 (setArch is left at 90) and
+  // the weights were prepacked for the Hopper layout. Default: no-op (keep the SM80 kernel).
+  virtual void setUseSm90Native(bool /*use*/) {}
+
  protected:
   static constexpr int SPLIT_K_LIMIT = 7;
   static constexpr int MIN_M_TILE = 16;
@@ -118,6 +131,10 @@ class CutlassFpAIntBGemmRunner : public virtual CutlassFpAIntBGemmRunnerInterfac
 
   std::vector<tkc::CutlassGemmConfig> getConfigs() const override;
 
+  void setArch(int sm) override { sm_ = sm; }
+
+  void setUseSm90Native(bool use) override { use_sm90_native_ = use; }
+
  private:
   template <typename EpilogueTag>
   void dispatch_to_arch(ActivationType const* A, WeightType const* B, ScaleZeroType const* weight_scales,
@@ -128,6 +145,7 @@ class CutlassFpAIntBGemmRunner : public virtual CutlassFpAIntBGemmRunnerInterfac
  private:
   int sm_;
   int multi_processor_count_;
+  bool use_sm90_native_{false};
 };
 
 }  // namespace cutlass_kernels

onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm/fpA_intB_gemm_template.h

@@ -134,8 +134,8 @@ void generic_mixed_gemm_kernelLauncher(ActivationType const* A, WeightType const
       }
     }
 
-    if (group_size != 64 && group_size != 128) {
-      ORT_THROW("Only group size 64 and 128 supported for fine grained kernels.");
+    if (group_size != 32 && group_size != 64 && group_size != 128) {
+      ORT_THROW("Only group size 32, 64 and 128 supported for fine grained kernels.");
     }
 
     if constexpr (QuantOp == cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY) {
@@ -393,9 +393,20 @@ void CutlassFpAIntBGemmRunner<ActivationType, WeightType, QuantOp, ScaleZeroType
                   cutlass::platform::is_same<ActivationType, ScaleZeroType>::value &&
                   cutlass::platform::is_same<ActivationType, BiasType>::value &&
                   cutlass::platform::is_same<ActivationType, OutputType>::value) {
-      dispatch_gemm_to_cutlass<ActivationType, WeightType, ScaleZeroType, BiasType, OutputType, cutlass::arch::Sm80,
-                               QuantOp, EpilogueTag>(A, B, weight_scales, weight_zero_points, biases, alpha, C, m, n, k, group_size,
-                                                     workspace_ptr, workspace_bytes, gemm_config, stream, occupancy);
+      // For half/bf16 weight-only GEMM on Hopper we support two paths:
+      //   - use_sm90_native_ == false: reuse the SM80 (Ampere) mixed-GEMM kernel (compat path,
+      //     consumes the SM80 column-interleaved weight layout).
+      //   - use_sm90_native_ == true: the native SM90 TMA/WGMMA mixed-GEMM kernel, which consumes
+      //     the Hopper weight layout (prepacked with arch=90, no column interleave).
+      if (use_sm90_native_) {
+        sm90_dispatch_gemm_to_cutlass<ActivationType, WeightType, ScaleZeroType, BiasType, OutputType, QuantOp,
+                                      EpilogueTag>(A, B, weight_scales, weight_zero_points, biases, alpha, C, m, n, k, group_size, workspace_ptr,
+                                                   workspace_bytes, gemm_config, stream, occupancy);
+      } else {
+        dispatch_gemm_to_cutlass<ActivationType, WeightType, ScaleZeroType, BiasType, OutputType, cutlass::arch::Sm80,
+                                 QuantOp, EpilogueTag>(A, B, weight_scales, weight_zero_points, biases, alpha, C, m, n, k, group_size,
+                                                       workspace_ptr, workspace_bytes, gemm_config, stream, occupancy);
+      }
     } else {
       static_assert(!cutlass::platform::is_same<ActivationType, __nv_fp8_e4m3>::value || cutlass::platform::is_same<ScaleZeroType, half>::value,
                     "ScaleZeroType must be half for activation=fp8");

onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm/launchers/fpA_intB_launcher_sm90.inl

@@ -59,7 +59,13 @@ using namespace cute;
 template <typename ActivationType, typename WeightType, typename ScaleZeroType, typename BiasType, typename OutputType,
           cutlass::WeightOnlyQuantOp QuantOp, typename EpilogueTag, typename CTAShape, typename ClusterShape,
           typename MainloopScheduleType, typename EpilogueScheduleType>
-#if defined(COMPILE_HOPPER_TMA_GEMMS) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 900) && defined(__NV_SASS_VERSION__)
+// Gate the real launcher on the COMPILE_HOPPER_TMA_GEMMS preprocessor macro only (matching the MoE
+// TMA launcher in moe_gemm_tma_ws_launcher.inl). The host-callable launcher symbol is emitted from
+// the host compilation pass, so it must NOT be gated on __CUDA_ARCH__/__NV_SASS_VERSION__ (which are
+// only set during device passes) — otherwise the host links the stub below and every SM90 tactic
+// fails at runtime with "recompile ... 90a". The device kernel body is guarded internally by the
+// collective (CUTE_ARCH_MMA_SM90A_ENABLED); these files are compiled at sm_90a-real.
+#if defined(COMPILE_HOPPER_TMA_GEMMS)
 void sm90_generic_mixed_gemm_kernelLauncher(
     ActivationType const* A, WeightType const* B,
     ScaleZeroType const* weight_scales, ScaleZeroType const* weight_zero_points, BiasType const* biases,

onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher.h

@@ -444,7 +444,10 @@ void check_pointer(Params& params, cudaStream_t s) {
 template <bool isGroupwise, typename Details>
 void select_gs(Params& params, cudaStream_t s) {
   if constexpr (isGroupwise) {
-    if (params.groupsize == 64) {
+    if (params.groupsize == 32) {
+      check_pointer<Details, 32>(params, s);
+      return;
+    } else if (params.groupsize == 64) {
       check_pointer<Details, 64>(params, s);
       return;
     } else if (params.groupsize == 128) {

onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h

@@ -57,20 +57,22 @@ struct GemmDims {
 };
 
 // Unique ID of GEMM
-// In our case GEMM is uniqly identified by N and K
+// In our case GEMM is uniquely identified by N and K, plus the target SM architecture (so the
+// SM80-compatibility and native SM90 kernels for the same shape do not share profiled configs).
 class GemmIdCore {
  public:
   int n;
   int k;
   nvinfer::DataType dtype;
+  int sm;
 
-  GemmIdCore(int n_, int k_, nvinfer::DataType const& dtype_)
-      : n(n_), k(k_), dtype(dtype_) {
+  GemmIdCore(int n_, int k_, nvinfer::DataType const& dtype_, int sm_ = 0)
+      : n(n_), k(k_), dtype(dtype_), sm(sm_) {
   }
 
   GemmIdCore()
-      : n(-1), k(-1), dtype(nvinfer::DataType::kFLOAT)  // dtype does not matter here
-  {
+      : n(-1), k(-1), dtype(nvinfer::DataType::kFLOAT),  // dtype does not matter here
+        sm(0) {
   }
 
   bool operator==(GemmIdCore const& id) const {
@@ -80,12 +82,13 @@ class GemmIdCore {
   friend std::ostream& operator<<(std::ostream& out, GemmIdCore const& id) {
     out << "(N;K)=(" << id.n << ";" << id.k << "),";
     out << " type=" << static_cast<int>(id.dtype);
+    out << " sm=" << id.sm;
     return out;
   }
 
  protected:
   bool isEqual(GemmIdCore const& id) const {
-    return n == id.n && k == id.k && dtype == id.dtype;
+    return n == id.n && k == id.k && dtype == id.dtype && sm == id.sm;
   }
 };
 
@@ -95,7 +98,8 @@ struct GemmIdCoreHash {
     auto h1 = std::hash<int>{}(id.n);
     auto h2 = std::hash<int>{}(id.k);
     auto h3 = std::hash<int>{}(static_cast<int>(id.dtype));
-    return h1 ^ h2 ^ h3;
+    auto h4 = std::hash<int>{}(id.sm);
+    return h1 ^ h2 ^ h3 ^ h4;
   }
 };