[Perf] Optimize Hopper MXFP4 × FP8_BLOCK indexed MoE prefill

humming/include/humming/memory/g2s_loader/loader_as.cuh

@@ -48,6 +48,13 @@ private:
   static_assert(BlockShape::M % kScaleMAlignment == 0);
   static_assert(!kUseTma || kScaleBlockM <= 256);
 
+  // One uint4 gather per routed row covers four single-group stages. Stream-K is
+  // excluded so every tile starts on a vector boundary; the plain smem stores
+  // rely on the CTA barrier of the non-mbarrier pipeline.
+  static constexpr bool kUseIndexedScaleCache =
+      kIsIndexedGemm && kIsGroupScale && !kUseMxScale && Ctx::kWgmmaUseLateAS && !Ctx::kUseMBarrier &&
+      !Ctx::kUseStreamK && kNumGroups == 1 && kProblemNumGroups % 4 == 0;
+
   using LoadType = typename LoadTypeChooser<kNumGroups * 4>::Type;
 
 public:
@@ -64,6 +71,7 @@ public:
   uint32_t load_row_index[kRowLoadIters];
   uint32_t col_offset = 0;
   uint32_t counter = 0;
+  uint4 indexed_scale_cache[kRowLoadIters];
 
   CUDA_INLINE
   G2SMemoryLoaderAS(Ctx &ctx)
@@ -87,6 +95,7 @@ public:
       else if constexpr (SharedStorage::kUseUmmaRowMajorSmemInputScale) load_mx_legacy_row_major(smem_ptr);
       else load_mx_legacy(smem_ptr);
     } else if constexpr (kUseTma) load_tma(smem_ptr, mbar_ptr);
+    else if constexpr (kUseIndexedScaleCache) load_legacy_indexed_cached(smem_ptr);
     else load_legacy(smem_ptr);
     if constexpr (kShouldAdvance) advance();
   }
@@ -198,6 +207,23 @@ public:
     }
   }
 
+  CUDA_INLINE void load_legacy_indexed_cached(void *smem_ptr) {
+    uint32_t thread_id = ctx.load_thread_id();
+    uint32_t *smem_ptr_load = reinterpret_cast<uint32_t *>(smem_ptr);
+    const uint32_t component = col_offset % 4;
+    PRAGMA_UNROLL
+    for (uint32_t i = 0; i < kRowLoadIters; i++) {
+      uint32_t row = i * kNumLoadThreads + thread_id;
+      uint32_t source_row = load_row_index[i];
+      if (row < BlockShape::M && source_row < shape_m) {
+        if (component == 0)
+          indexed_scale_cache[i] = *reinterpret_cast<const uint4 *>(gmem_ptr + source_row * kProblemNumGroups);
+        const uint4 v = indexed_scale_cache[i];
+        smem_ptr_load[row] = component == 0 ? v.x : component == 1 ? v.y : component == 2 ? v.z : v.w;
+      }
+    }
+  }
+
   CUDA_INLINE void load_legacy(void *smem_ptr) {
     uint32_t thread_id = ctx.load_thread_id();
     if constexpr (!kIsIndexedGemm && (kMMajorInputScale || kIsChannelScale)) {

humming/tune/__init__.py

@@ -18,7 +18,11 @@
 from humming.tune.sm75 import Sm75Heuristics
 from humming.tune.sm90 import Sm90Heuristics
 from humming.tune.sm90_h20 import Sm90H20Heuristics
-from humming.tune.sm90_policies import apply_w4a8_config, specialize_w4a8_ranges
+from humming.tune.sm90_policies import (
+    apply_indexed_w4a8_config,
+    apply_w4a8_config,
+    specialize_w4a8_ranges,
+)
 from humming.tune.sm100 import Sm100Heuristics
 from humming.tune.sm120 import Sm120Heuristics
 from humming.tune.sm121 import Sm121Heuristics
@@ -129,6 +133,7 @@ def _get_heuristics_config(
         _disable_indexed_input_scale_tma(config, gemm_type)
         _apply_raster_group_m(config, layer_config, gemm_type)
         apply_w4a8_config(config, layer_config, use_m_major_input_scale, gemm_type, shape_m)
+        apply_indexed_w4a8_config(config, layer_config, gemm_type)
         return config
     else:
         configs = heuristics_cls.get_configs(
@@ -142,6 +147,7 @@ def _get_heuristics_config(
             _apply_m_major_input_scale(entry[2], use_m_major_input_scale, layer_config, gemm_type)
             _disable_indexed_input_scale_tma(entry[2], gemm_type)
             _apply_raster_group_m(entry[2], layer_config, gemm_type)
+            apply_indexed_w4a8_config(entry[2], layer_config, gemm_type)
         return specialize_w4a8_ranges(configs, layer_config, use_m_major_input_scale, gemm_type)
 
 

humming/tune/sm90_policies.py

@@ -322,16 +322,10 @@ def expected_work(tile: int) -> float:
 )
 
 
-def _w4a8_enabled(
-    layer_config: LayerConfig,
-    use_m_major_input_scale: bool,
-    gemm_type: GemmType,
-) -> bool:
+def _is_w4a8_moe_layer(layer_config: LayerConfig) -> bool:
     return (
-        gemm_type == GemmType.GROUPED_CONTIGUOUS
-        and layer_config.sm_version == 90
+        layer_config.sm_version == 90
         and layer_config.shape_n % 128 == 0
-        and use_m_major_input_scale
         and get_default_mma_type(layer_config) == MmaType.WGMMA
         and layer_config.a_dtype == dtypes.float8e4m3
         and layer_config.b_dtype == dtypes.float4e2m1
@@ -345,6 +339,47 @@ def _w4a8_enabled(
     )
 
 
+def _is_grouped_w4a8(
+    layer_config: LayerConfig,
+    use_m_major_input_scale: bool,
+    gemm_type: GemmType,
+) -> bool:
+    return (
+        gemm_type == GemmType.GROUPED_CONTIGUOUS
+        and use_m_major_input_scale
+        and _is_w4a8_moe_layer(layer_config)
+    )
+
+
+def apply_indexed_w4a8_config(
+    config: dict,
+    layer_config: LayerConfig,
+    gemm_type: GemmType,
+) -> None:
+    """Select the packed GS128 mainloop while preserving indexed row gathering."""
+    if (
+        gemm_type != GemmType.INDEXED
+        or not _is_w4a8_moe_layer(layer_config)
+        or config.get("use_f16_accum", False)
+    ):
+        return
+
+    block_m = min(config["block_shape"][0], 128)
+    # Preserve the small-M schedule, where its wider K stage remains faster.
+    if block_m < 64:
+        return
+    config.update(
+        block_shape=(block_m, 128, 128),
+        warp_shape=(block_m, 16, 128),
+        num_stages=4,
+        use_warp_spec=False,
+        wgmma_use_late_as=True,
+        wgmma_split_issue_wait=True,
+        use_tma=False,
+        use_stream_k=False,
+    )
+
+
 def _w4a8_block_m(layer_config: LayerConfig, shape_m: int) -> int:
     rows_per_expert = (shape_m + layer_config.num_experts - 1) // layer_config.num_experts
     return _w4a8_tile_m_for_expert_rows(rows_per_expert)
@@ -375,7 +410,7 @@ def apply_w4a8_config(
     gemm_type: GemmType,
     shape_m: int,
 ) -> None:
-    if not _w4a8_enabled(layer_config, use_m_major_input_scale, gemm_type):
+    if not _is_grouped_w4a8(layer_config, use_m_major_input_scale, gemm_type):
         return
     use_h200_policy = _w4a8_uses_variable_m_tiles()
     block_m = _w4a8_block_m(layer_config, shape_m) if use_h200_policy else _W4A8_DEFAULT_TILE_M
@@ -388,7 +423,7 @@ def specialize_w4a8_ranges(
     use_m_major_input_scale: bool,
     gemm_type: GemmType,
 ) -> list:
-    if not _w4a8_enabled(layer_config, use_m_major_input_scale, gemm_type):
+    if not _is_grouped_w4a8(layer_config, use_m_major_input_scale, gemm_type):
         return configs
 
     if not _w4a8_uses_variable_m_tiles():