Merge branch 'main' into process-input-refactor

Author: jinzhen-lin

humming/include/humming/datatype/dequant_fused.cuh

@@ -17,9 +17,10 @@ CUDA_INLINE uint2 fused_dequant_single_for_mxfp4<Float8E4M3>(const uint32_t qb,
   uint32_t exp_offset_buffer1 = (exp_offset * 0x08080800) + ((0x03020100 << 2) - 0x00000400);
   uint32_t exp_offset_buffer2 = (exp_offset * 0x08080808) + (0x07060504 << 2);
 
+  uint32_t qb_selectors = qb & 0x77777777;
   uint32_t exp_offsets[2] = {
-      __byte_perm(exp_offset_buffer1, exp_offset_buffer2, qb),
-      __byte_perm(exp_offset_buffer1, exp_offset_buffer2, qb >> 16)};
+      prmt(exp_offset_buffer1, exp_offset_buffer2, qb_selectors),
+      prmt(exp_offset_buffer1, exp_offset_buffer2, qb_selectors >> 16)};
 
   uint32_t res[2] = {
       lop3_and_or(qb << 4, 0x80808080, exp_offsets[0]),
@@ -36,13 +37,14 @@ CUDA_INLINE uint2 fused_dequant_single_for_mxfp4<Int8>(const uint32_t qb, const
 
   uint32_t res[2];
   uint32_t signs[2] = {qb >> 3, qb >> 7};
+  uint32_t qb_selectors = qb & 0x77777777;
   uint32_t int8s[2] = {
-      __byte_perm(buffer1, buffer2, qb),
-      __byte_perm(buffer1, buffer2, qb >> 16)};
+      prmt(buffer1, buffer2, qb_selectors),
+      prmt(buffer1, buffer2, qb_selectors >> 16)};
 
   PRAGMA_UNROLL
   for (uint32_t i = 0; i < 2; i++) {
-    uint32_t val = __byte_perm(int8s[0], int8s[1], 0x6420 + 0x1111 * i);
+    uint32_t val = prmt(int8s[0], int8s[1], 0x6420 + 0x1111 * i);
     uint32_t flag = signs[i] & 0x01010101;
     uint32_t mask = flag * 0xFF;
     res[i] = (val - flag) ^ mask;
@@ -55,7 +57,9 @@ template <class TargetType, uint32_t kCount, bool kUseWgmma>
 CUDA_INLINE void fused_dequant_for_mxfp4(const uint32_t *qb_ptrs, uint32_t *res_ptrs, uint32_t *scales_ptr) {
   PRAGMA_UNROLL
   for (uint32_t i = 0; i < kCount * 2; i++) {
-    uint32_t exp_offset = reinterpret_cast<uint8_t *>(scales_ptr)[i];
+    uint32_t scale_word = scales_ptr[i / 4];
+    uint32_t byte_selector = 1u << (8 * (i % 4));
+    uint32_t exp_offset = __dp4a(scale_word, byte_selector, 0u);
     uint2 res = fused_dequant_single_for_mxfp4<TargetType>(qb_ptrs[i], exp_offset);
     res_ptrs[i * 2] = res.x;
     res_ptrs[i * 2 + 1] = res.y;

humming/include/humming/utils/ptx/math.cuh

@@ -2,6 +2,14 @@
 
 #include <humming/utils/base.cuh>
 
+CUDA_INLINE uint32_t prmt(uint32_t a, uint32_t b, uint32_t s) {
+  uint32_t res;
+  asm volatile("prmt.b32 %0, %1, %2, %3;\n"
+               : "=r"(res)
+               : "r"(a), "r"(b), "r"(s));
+  return res;
+};
+
 template <uint32_t lut>
 CUDA_INLINE uint32_t lop3(uint32_t a, uint32_t b, uint32_t c) {
   uint32_t res;

humming/tune/sm90_policies.py

@@ -115,7 +115,12 @@ def build_sm90_seed_config(problem: TuningProblem) -> dict:
     )
     if use_wide_indexed_tile:
         warp_shape_n = 64
-        if layer_config.shape_k <= 512 and layer_config.shape_n >= 2048:
+        # N=512 spills its accumulator at two-CTA residency from M=48 onward.
+        if (
+            layer_config.shape_k <= 512
+            and layer_config.shape_n >= 2048
+            and block_shape_m < 48
+        ):
             block_shape_n = 512
             block_shape_k = 64
         else:

tests/test_sm90_heuristics.py

@@ -288,6 +288,34 @@ def test_mxfp4_a16_indexed_preserves_warp_k_and_fits_grid(
     assert config["use_stream_k"] is expected_stream_k
 
 
+@pytest.mark.parametrize("shape_k", [512, 256])
+def test_mxfp4_a16_short_k_limits_wide_n_tile_before_block_m48(shape_k):
+    layer = _layer(
+        6144,
+        shape_k,
+        num_experts=256,
+        a_dtype=dtypes.bfloat16,
+        as_dtype=None,
+        input_scale_group_size=0,
+    )
+
+    small = Sm90Heuristics.get_config(
+        layer,
+        shape_m=6144,
+        gemm_type=GemmType.INDEXED,
+    )
+    large = Sm90Heuristics.get_config(
+        layer,
+        shape_m=8192,
+        gemm_type=GemmType.INDEXED,
+    )
+
+    assert small["block_shape"] == (32, 512, 64)
+    assert small["num_ctas_per_sm"] == 2
+    assert large["block_shape"] == (48, 256, 64)
+    assert large["num_ctas_per_sm"] == 2
+
+
 def test_nvfp4_a16_uses_narrow_first_wave_only():
     layer = _layer(
         5376,