diff --git a/humming/config/config.py b/humming/config/config.py index b699fadd..8c69c228 100644 --- a/humming/config/config.py +++ b/humming/config/config.py @@ -80,6 +80,7 @@ class LayerConfig(BaseHummingConfig): "is_token_input_scale_2", "is_tensor_input_scale_2", "use_raw_weight", + "use_flat_weight", "use_native_dequant", "use_block_scaled_mma", "mma_scale_group_size", @@ -132,7 +133,7 @@ def is_umma_supported(self): @property def use_raw_weight(self): if self.a_dtype.num_bits == self.b_dtype.num_bits: - return True + return self.is_umma_supported or current_device.is_ppu has_native_mixed_activation = self.a_dtype in ( dtypes.float8e4m3, dtypes.float8e5m2, @@ -145,6 +146,11 @@ def use_raw_weight(self): ) return has_native_mixed_activation and has_native_mixed_weight and self.is_umma_supported + @property + def use_flat_weight(self): + has_same_width = self.a_dtype.num_bits == self.b_dtype.num_bits + return has_same_width and not self.is_umma_supported and not current_device.is_ppu + @property def use_block_scaled_mma(self): if self.mxmma_supported: @@ -168,7 +174,7 @@ def mma_scale_group_size(self): @property def use_native_dequant(self): assert self.sm_version is not None - if self.use_raw_weight or self.sm_version // 10 == 9: + if self.use_raw_weight or self.use_flat_weight or self.sm_version // 10 == 9: return False from humming.jit.runtime import KernelRuntime @@ -379,6 +385,7 @@ def __post_init__(self): can_use_packed_k = ( self.sm_version // 10 == 9 and not self.use_raw_weight + and not self.use_flat_weight and has_packed_k_operands and has_packed_k_input_scale ) @@ -389,7 +396,9 @@ def __post_init__(self): self.use_packed_k_layout = can_use_packed_k and prefer_packed_k elif self.use_packed_k_layout: assert self.sm_version // 10 == 9, "use_packed_k_layout requires SM90" - assert not self.use_raw_weight, "raw weights do not use packed-K layout" + assert not (self.use_raw_weight or self.use_flat_weight), ( + "raw/flat weights do not use packed-K layout" + ) assert self.a_dtype.num_bits == 8, "use_packed_k_layout requires 8-bit activation" assert self.b_dtype.num_bits % 2 == 0, "use_packed_k_layout requires even-bit weight" diff --git a/humming/config/mma.py b/humming/config/mma.py index 77dac209..202cc6c1 100644 --- a/humming/config/mma.py +++ b/humming/config/mma.py @@ -174,10 +174,11 @@ def get_register_budget_error(layer_config, tuning_config, use_f16_accum=False, return None # UMMA accumulators use TMEM and have separate resource checks. if mma_type == MmaType.WGMMA: + use_wgmma_ss = layer_config.a_dtype.num_bits == layer_config.b_dtype.num_bits # ptxas needs 26 (SS) or 32 (RS) registers beyond one N=64 accumulator tuple. # setmaxnreg cannot relax this per-instruction compile limit. instruction_accumulators = warp_m // (4 if use_f16_accum else 2) - instruction_overhead = 26 if layer_config.use_raw_weight else 32 + instruction_overhead = 26 if use_wgmma_ss else 32 instruction_registers = instruction_accumulators + instruction_overhead if layer_config.use_packed_k_layout and tuning_config.wgmma_split_issue_wait: # Split packed-K keeps all RS K slabs live across prefetch. The diff --git a/humming/csrc/launcher/launcher.cpp b/humming/csrc/launcher/launcher.cpp index 40e948d3..bed4e20d 100644 --- a/humming/csrc/launcher/launcher.cpp +++ b/humming/csrc/launcher/launcher.cpp @@ -336,6 +336,7 @@ std::tuple register_kernel(const std::string &cubin_path) reader.getBool("USE_PACKED_K_LAYOUT"), reader.getBool("USE_UMMA_SS"), reader.getBool("USE_RAW_WEIGHT"), + reader.getBool("USE_FLAT_WEIGHT"), reader.getBool("USE_BLOCK_SCALED_MMA")}; std::unique_lock lock(g_kernel_mutex); diff --git a/humming/csrc/launcher/tensor.h b/humming/csrc/launcher/tensor.h index d7c83fa6..c9732674 100644 --- a/humming/csrc/launcher/tensor.h +++ b/humming/csrc/launcher/tensor.h @@ -87,7 +87,7 @@ inline void check_tensor_b(Tensor &tensor, KernelData &kernel_data, int64_t dev) std::vector expected_shape = {}; if (kernel_data.gemm_type_id != 0) expected_shape.push_back(kernel_data.num_experts); - if (kernel_data.use_raw_weight) { + if (kernel_data.use_raw_weight || kernel_data.use_flat_weight) { if (get_dtype_num_bits(kernel_data.a_dtype_id) == 8 && num_bits < 8) problem_shape_k = CEIL_DIV(problem_shape_k, 128) * 128; expected_shape.push_back(problem_shape_n); @@ -390,6 +390,24 @@ inline CUtensorMap make_tma_desc_b(Tensor &tensor, KernelData &kernel_data) { uint32_t block_shape_n = kernel_data.block_shape_n; uint32_t block_shape_k = kernel_data.block_shape_k; + if (kernel_data.use_flat_weight) { + CUtensorMap descriptor{}; + uint64_t warp_n = kernel_data.mma_type == MmaType::WGMMA ? kernel_data.warp_shape_n : 16; + uint64_t shape_n = kernel_data.problem_shape_n; + uint64_t num_experts = std::max(1u, kernel_data.num_experts); + uint64_t row_bytes = uint64_t(tensor.size(-1)) * 4; + uint64_t dimensions[] = {64, 64 / warp_n, shape_n / 64, warp_n / 16, row_bytes / 16 * num_experts}; + uint64_t strides[] = {warp_n * 16, 1024, 256, shape_n * 16}; + uint32_t box[] = {64, uint32_t(std::min(64u, block_shape_n) / warp_n), + (block_shape_n + 63) / 64, uint32_t(warp_n / 16), block_shape_k * num_bits / 128}; + uint32_t element_strides[] = {1, 1, 1, 1, 1}; + CUresult status = cuTensorMapEncodeTiled(&descriptor, CU_TENSOR_MAP_DATA_TYPE_INT32, 5, + tensor.data_ptr(), dimensions, strides, box, element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + ASSERT_CHECK(status == CUDA_SUCCESS, "TMA Encode Failed for flat weights: ", int(status)); + return descriptor; + } + if (kernel_data.use_raw_weight) { // FP4/FP6 in an f8f6f4 instruction uses padded 16-element blocks in SMEM. bool expand_low_bit = num_bits < 8 && get_dtype_num_bits(kernel_data.a_dtype_id) == 8; @@ -408,25 +426,6 @@ inline CUtensorMap make_tma_desc_b(Tensor &tensor, KernelData &kernel_data) { ASSERT_CHECK(status == CUDA_SUCCESS, "TMA Encode Failed for expanded low-bit weights"); return descriptor; } - if (kernel_data.mma_type == MmaType::WGMMA) { - CUtensorMap descriptor{}; - uint64_t w = kernel_data.warp_shape_n; - uint64_t row_bytes = uint64_t(tensor.size(-1)) * 4; - // Fragment planes span all N warps in the CTA, leaving the fifth - // dimension available for K slabs even with multiple N warpgroups. - uint64_t dimensions[] = {swizzle_bytes / 4, 16, - uint64_t(kernel_data.problem_shape_n) * std::max(1u, kernel_data.num_experts) / w, - w / 16, row_bytes / swizzle_bytes}; - uint64_t strides[] = {row_bytes, w * row_bytes, 16 * row_bytes, swizzle_bytes}; - uint32_t box[] = {swizzle_bytes / 4, 16, uint32_t(block_shape_n / w), uint32_t(w / 16), - block_shape_k * num_bits / 8 / swizzle_bytes}; - uint32_t element_strides[] = {1, 1, 1, 1, 1}; - CUresult status = cuTensorMapEncodeTiled(&descriptor, CU_TENSOR_MAP_DATA_TYPE_INT32, 5, - tensor.data_ptr(), dimensions, strides, box, element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, - get_swizzle_enum(swizzle_bytes), CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - ASSERT_CHECK(status == CUDA_SUCCESS, "TMA Encode Failed for WGMMA N permutation: ", int(status)); - return descriptor; - } auto rows = torch_view_shape(tensor, {-1, tensor.size(-1)}); return make_tma_desc(rows, {swizzle_bytes / 4, std::min(256u, kernel_data.use_umma_ss ? std::max(128u, block_shape_n) : block_shape_n)}, swizzle_bytes, "b"); } diff --git a/humming/csrc/launcher/utils.h b/humming/csrc/launcher/utils.h index 92dca757..03ee33c8 100644 --- a/humming/csrc/launcher/utils.h +++ b/humming/csrc/launcher/utils.h @@ -179,6 +179,7 @@ struct KernelData { bool use_packed_k_layout; bool use_umma_ss; bool use_raw_weight; + bool use_flat_weight; bool use_block_scaled_mma; }; diff --git a/humming/include/humming/kernel/humming_ws.cuh b/humming/include/humming/kernel/humming_ws.cuh index 3d4eb908..1599142c 100644 --- a/humming/include/humming/kernel/humming_ws.cuh +++ b/humming/include/humming/kernel/humming_ws.cuh @@ -189,6 +189,9 @@ __global__ __launch_bounds__(TuningConfig::kNumThreads, TuningConfig::kNumCtasPe if (warp_iter_id == Ctx::kWarpIters - 1 && slice_iter + 1 < num_slice_iters) { consumer.wait_stage((stage_id + 1) % kNumStages); } + if constexpr (!Ctx::kUseWgmma) { + if (warp_iter_id == Ctx::kWarpIters - 1) consumer.arrive(stage_id); + } const bool use_late_scale_prefetch = Ctx::kUseWgmmaLateScalePrefetch && warp_iter_id == Ctx::kWarpIters - 1; const bool has_next_fragment = warp_iter_id + 1 < Ctx::kWarpIters || slice_iter + 1 < num_slice_iters; if (has_next_fragment) { @@ -208,7 +211,9 @@ __global__ __launch_bounds__(TuningConfig::kNumThreads, TuningConfig::kNumCtasPe ((warp_iter_id + 1) % Ctx::kWarpIters) % 2, (warp_iter_id + 1) % Ctx::kWarpIters); } - if (warp_iter_id == Ctx::kWarpIters - 1) consumer.arrive(stage_id); + if constexpr (Ctx::kUseWgmma) { + if (warp_iter_id == Ctx::kWarpIters - 1) consumer.arrive(stage_id); + } } }; diff --git a/humming/include/humming/memory/g2s_loader/loader_b.cuh b/humming/include/humming/memory/g2s_loader/loader_b.cuh index 454fa24c..6232f20f 100644 --- a/humming/include/humming/memory/g2s_loader/loader_b.cuh +++ b/humming/include/humming/memory/g2s_loader/loader_b.cuh @@ -54,7 +54,8 @@ public: template CUDA_INLINE void load(int4 *smem_ptr, void *mbar_ptr) { - if constexpr (Ctx::kUseRawWeight && !kUseTma) load_raw(smem_ptr); + if constexpr (Ctx::kUseFlatWeight && !kUseTma) load_flat(smem_ptr); + else if constexpr (Ctx::kUseRawWeight && !kUseTma) load_raw(smem_ptr); else if constexpr (kUseAiu) load_aiu(smem_ptr); else if constexpr (kUseTma) load_tma(smem_ptr, mbar_ptr); else load_legacy(smem_ptr); @@ -76,22 +77,21 @@ public: CUDA_INLINE void load_tma(int4 *smem_ptr, void *mbar_ptr) { if (ctx.load_thread_id() == 0) { - if constexpr (Ctx::kUseRawWeight) { + if constexpr (Ctx::kUseFlatWeight) { + uint32_t expert_k = col_offset / ProblemShape::N * ProblemShape::K + row_offset; + if (cluster_rank == 0) + tma_load_5d(tensor_map_ptr, smem_ptr, mbar_ptr, + 0, col_offset % 64 / 16, col_offset % ProblemShape::N / 64, 0, expert_k * ElementQB::kBits / 128); + } else if constexpr (Ctx::kUseRawWeight) { constexpr uint32_t kCoordinateBits = SharedStorage::kExpandUmmaWeight ? 1 : 32 / ElementQB::kBits; uint32_t start_k = kCoordinateBits == 1 ? row_offset / 128 * 128 : row_offset; - if constexpr (Ctx::kUseWgmmaTmaKPack) { - if (cluster_rank == 0) - tma_load_5d(tensor_map_ptr, smem_ptr, mbar_ptr, - 0, 0, col_offset / Ctx::WarpShape::N, 0, start_k / kSwizzleK); - } else { + PRAGMA_UNROLL + for (uint32_t k = 0; k < SharedStorage::kWeightStageK; k += kSwizzleK) { PRAGMA_UNROLL - for (uint32_t k = 0; k < SharedStorage::kWeightStageK; k += kSwizzleK) { - PRAGMA_UNROLL - for (uint32_t n = 0; n < SharedStorage::kWeightStageN; n += 256) { - uint32_t offset = (k / kSwizzleK * SharedStorage::kWeightStageN + n) * kSwizzleBytes / sizeof(int4); - if (cluster_rank == 0) tma_load_2d( - tensor_map_ptr, smem_ptr + offset, mbar_ptr, (start_k + k) / kCoordinateBits, col_offset + n); - } + for (uint32_t n = 0; n < SharedStorage::kWeightStageN; n += 256) { + uint32_t offset = (k / kSwizzleK * SharedStorage::kWeightStageN + n) * kSwizzleBytes / sizeof(int4); + if (cluster_rank == 0) tma_load_2d( + tensor_map_ptr, smem_ptr + offset, mbar_ptr, (start_k + k) / kCoordinateBits, col_offset + n); } } } else if constexpr (kMultiCastSizeB == 1) { @@ -106,12 +106,12 @@ public: void prefetch_tma() { if constexpr (kUseTma && kMultiCastSizeB == 1) { if (ctx.load_thread_id() == 0) { - if constexpr (!Ctx::kUseRawWeight) tma_prefetch_3d(tensor_map_ptr, 0, col_offset, row_offset); + if constexpr (!Ctx::kUseRawWeight && !Ctx::kUseFlatWeight) tma_prefetch_3d(tensor_map_ptr, 0, col_offset, row_offset); } } } - CUDA_INLINE static uint32_t permute_raw_ss_row(uint32_t row) { + CUDA_INLINE static uint32_t permute_wgmma_ss_row(uint32_t row) { constexpr uint32_t kWarpN = Ctx::WarpShape::N; constexpr uint32_t kGroupN = 4 * kWarpN; uint32_t row_in_group = row % kGroupN; @@ -127,36 +127,34 @@ public: constexpr uint32_t kGlobalRowInt4s = ProblemShape::K * ElementQB::kBits / 128; constexpr uint32_t kSwizzleInt4s = kSwizzleBytes / sizeof(int4); uint32_t smem_base = cast_smem_ptr_to_uint(smem_ptr) / 128; - constexpr uint32_t kRowsPerIteration = kNumLoadThreads / kRowInt4s; - constexpr bool kPowerOfTwoRows = kRowsPerIteration != 0 && (kRowsPerIteration & (kRowsPerIteration - 1)) == 0; - constexpr bool kPowerOfTwoWarpN = (Ctx::WarpShape::N & (Ctx::WarpShape::N - 1)) == 0; - constexpr bool kSplitRowPermutation = - Ctx::kUseWgmmaCpAsyncNPermute && kNumLoadThreads % kRowInt4s == 0 && kPowerOfTwoRows && kPowerOfTwoWarpN; - uint32_t permuted_thread_row = 0; - if constexpr (kSplitRowPermutation) { - // The thread and unrolled iteration own disjoint row bits. Permute - // thread bits once; the iteration contribution folds to constants. - uint32_t thread_row = (ctx.load_thread_id() / kRowInt4s) % kRowsPerIteration; - permuted_thread_row = permute_raw_ss_row(thread_row); - } PRAGMA_UNROLL for (uint32_t i = 0; i < CEIL_DIV(BlockShape::N * kRowInt4s, kNumLoadThreads); ++i) { uint32_t index = i * kNumLoadThreads + ctx.load_thread_id(); uint32_t row = index / kRowInt4s; uint32_t col = index % kRowInt4s; - uint32_t smem_row = row; - if constexpr (Ctx::kUseWgmmaCpAsyncNPermute) { - if constexpr (kSplitRowPermutation) { - smem_row = permuted_thread_row + permute_raw_ss_row(i * kRowsPerIteration); - } else smem_row = permute_raw_ss_row(row); - } - uint32_t offset = (col / kSwizzleInt4s * BlockShape::N + smem_row) * kSwizzleInt4s + col % kSwizzleInt4s; + uint32_t offset = (col / kSwizzleInt4s * BlockShape::N + row) * kSwizzleInt4s + col % kSwizzleInt4s; uint32_t swizzled_offset = offset ^ ((smem_base + offset / 8) % kSwizzleInt4s); legacy_load_pred(gmem_ptr + row * kGlobalRowInt4s + col, smem_ptr + swizzled_offset, row < BlockShape::N); } } + CUDA_INLINE + void load_flat(int4 *smem_ptr) { + constexpr uint32_t kChunkRows = MIN(64u, BlockShape::N); + constexpr uint32_t kTileInt4s = kChunkRows * BlockShape::K * ElementQB::kBits / 128; + PRAGMA_UNROLL + for (uint32_t i = 0; i < CEIL_DIV(kNumInt4s, kNumLoadThreads); ++i) { + uint32_t index = i * kNumLoadThreads + ctx.load_thread_id(); + uint32_t row = index / kTileInt4s * 64 + index % kChunkRows; + uint32_t col = index % kTileInt4s / kChunkRows; + if constexpr (Ctx::kUseWgmmaCpAsyncNPermute) row = permute_wgmma_ss_row(row); + uint32_t offset = col * BlockShape::N + row; + uint32_t gmem_offset = col * ProblemShape::N + index / kTileInt4s * 64 + index % kChunkRows; + legacy_load_pred(gmem_ptr + gmem_offset, smem_ptr + offset, index < kNumInt4s); + } + } + CUDA_INLINE void load_raw_aiu(int4 *smem_ptr) { const uint32_t warp_id = ctx.load_thread_id() / 32; @@ -179,9 +177,9 @@ public: CUDA_INLINE void advance() { - if constexpr (Ctx::kUseRawWeight) { + if constexpr (Ctx::kUseRawWeight || Ctx::kUseFlatWeight) { row_offset += BlockShape::K; - if constexpr (!kUseTma) gmem_ptr += BlockShape::K * ElementQB::kBits / 128; + if constexpr (!kUseTma) gmem_ptr += BlockShape::K * ElementQB::kBits / 128 * (Ctx::kUseFlatWeight ? ProblemShape::N : 1); return; } row_offset += BlockShape::K / kPackSizeK; @@ -190,12 +188,18 @@ public: CUDA_INLINE void seek(uint32_t expert_id, uint32_t n_block_id, uint32_t k_block_id) { - if constexpr (Ctx::kUseRawWeight) { + if constexpr (Ctx::kUseRawWeight || Ctx::kUseFlatWeight) { row_offset = k_block_id * BlockShape::K; col_offset = expert_id * ProblemShape::N + n_block_id * BlockShape::N; if constexpr (!kUseTma) { - uint64_t offset = uint64_t(col_offset) * ProblemShape::K + row_offset; - gmem_ptr = gmem_ptr_raw + offset * ElementQB::kBits / 128; + if constexpr (Ctx::kUseFlatWeight) { + uint64_t expert_offset = uint64_t(expert_id) * ProblemShape::N * ProblemShape::K; + uint64_t offset = expert_offset + uint64_t(row_offset) * ProblemShape::N; + gmem_ptr = gmem_ptr_raw + offset * ElementQB::kBits / 128 + n_block_id * BlockShape::N; + } else { + uint64_t offset = uint64_t(col_offset) * ProblemShape::K + row_offset; + gmem_ptr = gmem_ptr_raw + offset * ElementQB::kBits / 128; + } } return; } diff --git a/humming/include/humming/memory/s2r_loader/loader_b.cuh b/humming/include/humming/memory/s2r_loader/loader_b.cuh index e521fe8c..6d2e6c71 100644 --- a/humming/include/humming/memory/s2r_loader/loader_b.cuh +++ b/humming/include/humming/memory/s2r_loader/loader_b.cuh @@ -62,7 +62,7 @@ public: } CUDA_INLINE - void load_raw(const int4 *smem_ptr, uint32_t *regs_ptr, uint32_t iter_id) { + void load_native(const int4 *smem_ptr, uint32_t *regs_ptr, uint32_t iter_id) { if constexpr (USE_PPU && Ctx::kUseCpAsync && !Ctx::kUseTmaB) return load_raw_aiu(smem_ptr, regs_ptr, iter_id); constexpr uint32_t kSwizzleInt4s = MIN(8u, BlockShape::K * ElementQB::kBits / 128); @@ -76,6 +76,7 @@ public: uint32_t col = (ctx.k_warp_offset() * ElementQB::kBits / 256 + iter_id) * 2 + (lane_id / 8) % 2; uint32_t offset = (col / kSwizzleInt4s * BlockShape::N + row) * kSwizzleInt4s + col % kSwizzleInt4s; uint32_t swizzled_offset = offset ^ ((smem_base + offset / 8) % kSwizzleInt4s); + if constexpr (Ctx::kUseFlatWeight) swizzled_offset = col * BlockShape::N + row; ld_shared(smem_ptr + swizzled_offset, reinterpret_cast(regs_ptr + n * kRowsPerLoad / 4)); } } @@ -94,7 +95,7 @@ public: CUDA_INLINE void load(const int4 *smem_ptr, uint32_t *regs_ptr, uint32_t iter_id) { - if constexpr (Ctx::kUseRawWeight) return load_raw(smem_ptr, regs_ptr, iter_id); + if constexpr (Ctx::kUseRawWeight || Ctx::kUseFlatWeight) return load_native(smem_ptr, regs_ptr, iter_id); if constexpr (Ctx::kUsePackedKLayout) return load_packed_k(smem_ptr, regs_ptr, iter_id); uint32_t warp_id = ctx.warp_id(); uint32_t n_warp_id = ctx.n_warp_id(); diff --git a/humming/include/humming/memory/s2r_pipeline.cuh b/humming/include/humming/memory/s2r_pipeline.cuh index 9b8809c4..d0e6c662 100644 --- a/humming/include/humming/memory/s2r_pipeline.cuh +++ b/humming/include/humming/memory/s2r_pipeline.cuh @@ -85,7 +85,7 @@ public: uint32_t buffer_id = iter_id % 2; auto &smem = ctx.smem; - if constexpr (!(kUseWgmma && Ctx::kUseRawWeight)) + if constexpr (!Ctx::kUseWgmmaSs) loader_b.load(smem.stages[stage_id].b, mma.regs_qb_as_ptr(buffer_id), iter_id); if constexpr (!kUseWgmma && !Ctx::kUseUmma) loader_a.load(smem.stages[stage_id].a, mma.regs_a_as_ptr(buffer_id), iter_id, stage_id); diff --git a/humming/include/humming/mma/mxmma.cuh b/humming/include/humming/mma/mxmma.cuh index bd587201..39b970f1 100644 --- a/humming/include/humming/mma/mxmma.cuh +++ b/humming/include/humming/mma/mxmma.cuh @@ -70,8 +70,7 @@ public: CUDA_INLINE void transform_b(uint32_t buffer_id, uint32_t iter_id) { - if constexpr (Ctx::kUseRawWeight) { - // Raw operands are loaded directly into regs_b, including mixed FP4 formats. + if constexpr (Ctx::kUseRawWeight || Ctx::kUseFlatWeight) { return; } else if constexpr (kNativeMixed) { PRAGMA_UNROLL @@ -180,7 +179,7 @@ public: template CUDA_INLINE T *regs_qb_as_ptr(uint32_t buffer_id) { - if constexpr (Ctx::kUseRawWeight) { + if constexpr (Ctx::kUseRawWeight || Ctx::kUseFlatWeight) { return reinterpret_cast(regs_b[buffer_id]); } else { return reinterpret_cast(regs_qb[buffer_id]); diff --git a/humming/include/humming/mma/wgmma.cuh b/humming/include/humming/mma/wgmma.cuh index d89aa3fd..4a90a252 100644 --- a/humming/include/humming/mma/wgmma.cuh +++ b/humming/include/humming/mma/wgmma.cuh @@ -13,7 +13,7 @@ CUDA_INLINE uint64_t make_wgmma_smem_desc(uint32_t addr) { constexpr uint64_t desc_base = (swizzle_type << 62) | (stride << 32); uint64_t desc = desc_base; - reinterpret_cast(&desc)[0] = (addr >> 4); + reinterpret_cast(&desc)[0] = (addr >> 4) & 0x3fff; return desc; }; @@ -81,7 +81,7 @@ public: CUDA_INLINE void transform_b(uint32_t buffer_id, uint32_t iter_id) { - if constexpr (Ctx::kUseRawWeight || std::is_same::value) return; + if constexpr (Ctx::kUseWgmmaSs) return; if constexpr (kUseFusedE8m0Scale) { uint32_t *regs_b_ptr = reinterpret_cast(regs_b[buffer_id]); @@ -157,19 +157,18 @@ public: PRAGMA_UNROLL for (uint32_t j = 0; j < kNumIters; j++) { - if constexpr (Ctx::kUseRawWeight) { - constexpr uint32_t kSwizzleK = kSwizzleBytes * 8 / ElementQB::kBits; + if constexpr (Ctx::kUseWgmmaSs) { + static_assert(Ctx::kUseFlatWeight); uint32_t row = ctx.n_warp_id() / 4 * (WarpShape::N * 4) + j * MmaShape::N; if constexpr (Ctx::kUseWgmmaTmaNPermute) { constexpr uint32_t kFragmentPlaneRows = BlockShape::N / WarpShape::N * 16; row = ctx.n_warp_id() / 4 * MmaShape::N + j * kFragmentPlaneRows; } uint32_t col = ctx.k_warp_offset() + k_slab * kPartMmaShapeK; - uint32_t offset = (col / kSwizzleK * BlockShape::N + row) * kSwizzleBytes; - offset += col % kSwizzleK * ElementQB::kBits / 8; + uint32_t offset = (col * ElementQB::kBits / 128 * BlockShape::N + row) * 16; const uint32_t b_addr = stage_addr + offsetof(typename SharedStorage::StageStorage, b); - uint64_t b_desc = make_wgmma_smem_desc(b_addr); - reinterpret_cast(&b_desc)[0] += offset / 16; + uint32_t b_start = ((b_addr + offset) >> 4) & 0x3fff; + uint64_t b_desc = b_start | (uint64_t(BlockShape::N) << 16) | (uint64_t(8) << 32); MmaOpClass::fma(desc, b_desc, regs_c[0][delta_j + j][0], scale_d); } else { MmaOpClass::fma(desc, regs_b[buffer_id][j][k], regs_c[0][delta_j + j][0], scale_d); diff --git a/humming/include/humming/scheduler.cuh b/humming/include/humming/scheduler.cuh index 70357f85..5c03afa1 100644 --- a/humming/include/humming/scheduler.cuh +++ b/humming/include/humming/scheduler.cuh @@ -120,7 +120,7 @@ public: streamk_mnk_total_iters = MAX(streamk_mnk_total_iters, kMinSliceIters); constexpr int32_t bpg = kStreamKBlockAlignment; - constexpr int32_t align_iters = bpg / ct_gcd(bpg, (int32_t)kNumStages) * (int32_t)kNumStages; + constexpr int32_t align_iters = bpg / ct_gcd(bpg, kNumStages) * kNumStages; if constexpr (align_iters > 1) { streamk_mnk_total_iters = align_iters * CEIL_DIV(streamk_mnk_total_iters, align_iters); }; diff --git a/humming/include/humming/utils/context.cuh b/humming/include/humming/utils/context.cuh index b5d4f787..9ea411ac 100644 --- a/humming/include/humming/utils/context.cuh +++ b/humming/include/humming/utils/context.cuh @@ -59,9 +59,8 @@ struct KernelContext : LayerConfig_, ComputeConfig_, TuningConfig_ { static constexpr bool kUseWgmma = TuningConfig::kMmaType == MmaType::WGMMA; static constexpr bool kUseCpAsync = LayerConfig::kSmVersion >= 80; static constexpr bool kUseMBarrier = kUseUmma || TuningConfig::kUseWarpSpec || TuningConfig::kUseTma; - static constexpr bool kUseWgmmaSs = kUseWgmma && LayerConfig::kUseRawWeight; + static constexpr bool kUseWgmmaSs = kUseWgmma && std::is_same::value; static constexpr bool kUseWgmmaTmaNPermute = kUseWgmmaSs && TuningConfig::kUseTmaB; - static constexpr bool kUseWgmmaTmaKPack = kUseWgmmaTmaNPermute; static constexpr bool kUseWgmmaTmaAPack = kUseWgmma && TuningConfig::kUseTmaA && ElementA::kBits >= 8 && !kIsIndexedGemm && (ProblemShape::K - PadShape::K) * ElementA::kBits % 1024 == 0 && BlockShape::K * ElementA::kBits > 1024; static constexpr bool kUseWgmmaCpAsyncNPermute = kUseWgmmaSs && !TuningConfig::kUseTmaB && WarpShape::N > 16; diff --git a/humming/include/humming/utils/storage.cuh b/humming/include/humming/utils/storage.cuh index 05f9eb15..1e8ecce5 100644 --- a/humming/include/humming/utils/storage.cuh +++ b/humming/include/humming/utils/storage.cuh @@ -239,7 +239,7 @@ public: static constexpr bool kUseWarpSpec = TuningConfig::kUseWarpSpec; struct StageStorage { alignas(1024) int4 a[kStageSizeA]; - alignas(LayerConfig::kUseRawWeight ? 1024 : 128) int4 b[kStageSizeB]; + alignas(LayerConfig::kUseRawWeight || LayerConfig::kUseFlatWeight ? 1024 : 128) int4 b[kStageSizeB]; IF_HAS_STAGE_INPUT_SCALE(alignas(128) int4 as[kStageSizeAS];) IF_HAS_STAGE_WEIGHT_SCALE(alignas(128) int4 bs[kStageSizeBS];) IF_HAS_STAGE_ZERO_POINT(alignas(128) int4 bzp[kStageSizeBZP];) diff --git a/humming/kernel/humming.py b/humming/kernel/humming.py index 110154a9..f26573e2 100644 --- a/humming/kernel/humming.py +++ b/humming/kernel/humming.py @@ -260,7 +260,8 @@ def select_mma_dtypes(self): and self.a_dtype in (dtypes.float8e4m3, dtypes.float8e5m2, dtypes.float8e3m4) and self.b_dtype in (dtypes.float4e2m1, dtypes.float6e3m2, dtypes.float6e2m3) ) - use_native_weight = self.use_raw_weight or mma_native_mixed + use_native_weight = self.a_dtype.num_bits == self.b_dtype.num_bits or self.use_umma_ss + use_native_weight |= mma_native_mixed use_native_weight |= self.mma_type == MmaType.MXMMA and self.mxmma_native_mixed if use_native_weight: self.mma_b_dtype = self.b_dtype diff --git a/humming/transform.py b/humming/transform.py index d77c7651..483877ce 100644 --- a/humming/transform.py +++ b/humming/transform.py @@ -256,6 +256,7 @@ def transform_humming_weight( use_packed_k_layout: bool = False, use_native_dequant: bool = False, use_raw_weight: bool = False, + use_flat_weight: bool = False, ) -> torch.Tensor: is_moe = weight.ndim == 3 weight = weight.unsqueeze(0) if not is_moe else weight @@ -292,15 +293,19 @@ def transform_humming_weight( weight2 = (weight & 0xF0) - 8 * 16 weight = (weight1 | weight2).view(torch.int32) - if use_raw_weight or a_dtype.num_bits == b_dtype.num_bits: + if use_raw_weight or use_flat_weight or a_dtype.num_bits == b_dtype.num_bits: if a_dtype.num_bits == 8 and b_dtype.num_bits < 8: padded_shape_k = round_up(padded_shape_k, 128) - # K-contiguous rows are consumed directly by the shared-memory MMA operand. if not packed: weight = ops.pack_weight(weight, b_dtype.num_bits) weight = torch.nn.functional.pad( weight, (0, (padded_shape_k - shape_k) * b_dtype.num_bits // 32, 0, padded_shape_n - shape_n) ) + if use_flat_weight: + original_shape = weight.shape + num_k_chunks = padded_shape_k // (packed_block_size_k // 2) + weight = weight.reshape(-1, padded_shape_n // 64, 64, num_k_chunks, 4) + weight = weight.permute(0, 3, 1, 2, 4).contiguous().reshape(original_shape) if current_device.is_ppu: # Match the even-then-odd channel order of PPU MMA's B operand. weight = weight.view(-1, padded_shape_n // 8, 8, weight.size(-1)) @@ -560,6 +565,7 @@ def transform_humming_tensors( use_packed_k_layout=config.use_packed_k_layout, use_native_dequant=config.use_native_dequant, use_raw_weight=config.use_raw_weight, + use_flat_weight=config.use_flat_weight, ) if weight_scale is not None and config.bs_dtype == dtypes.float8e8m0: diff --git a/humming/tune/sm120.py b/humming/tune/sm120.py index c67398f9..b1c992d5 100644 --- a/humming/tune/sm120.py +++ b/humming/tune/sm120.py @@ -244,7 +244,7 @@ def _tune_small_tile_transfers(cls, layer_config, config, gemm_type, shape_m) -> num_a_bits = layer_config.a_dtype.num_bits num_b_bits = layer_config.b_dtype.num_bits if gemm_type == GemmType.DENSE and shape_m <= block_m <= 16: - has_dequant_mainloop = num_a_bits < 16 and not layer_config.use_raw_weight + has_dequant_mainloop = num_a_bits < 16 and num_a_bits != num_b_bits if not config["use_stream_k"] or not has_dequant_mainloop: cls._use_cp_async(config) elif gemm_type == GemmType.INDEXED: diff --git a/humming/tune/sm90_policies.py b/humming/tune/sm90_policies.py index ff0efb21..fe8f4fd2 100644 --- a/humming/tune/sm90_policies.py +++ b/humming/tune/sm90_policies.py @@ -29,6 +29,7 @@ def fit_sm90_tile( gemm_type: GemmType = GemmType.DENSE, use_m_major_input_scale: bool = False, ) -> dict: + use_wgmma_ss = layer_config.a_dtype.num_bits == layer_config.b_dtype.num_bits config = config.copy() requested_ctas = config.get("num_ctas_per_sm", 1) block_m, block_n, block_k = config["block_shape"] @@ -46,7 +47,7 @@ def fit_sm90_tile( has_full_group_accumulator = layer_config.a_dtype.is_integer_type has_full_group_accumulator |= config.get("wgmma_split_issue_wait", False) max_group_warp_m = 80 if has_packed_input_scales and has_full_group_accumulator else 96 - has_unpacked_rs = not layer_config.use_raw_weight and not layer_config.use_packed_k_layout + has_unpacked_rs = not use_wgmma_ss and not layer_config.use_packed_k_layout should_widen_rs = has_unpacked_rs and has_group_accumulator and not use_warp_spec and warp_n < 32 if mma_type == MmaType.WGMMA and should_widen_rs and layer_config.shape_n % 128 == 0: config["block_shape"] = (block_m, max(128, block_n), block_k) @@ -106,9 +107,7 @@ def fit_sm90_tile( and layer_config.is_group_input_scale and layer_config.shape_k <= 512 ) - should_limit_stages = has_packed_group_input or ( - tuning.use_stream_k and layer_config.use_raw_weight - ) + should_limit_stages = has_packed_group_input or (tuning.use_stream_k and use_wgmma_ss) mma_type = tuning.mma_type or get_default_mma_type(layer_config) if should_limit_stages and mma_type == MmaType.WGMMA: k_iters = layer_config.shape_k // config["block_shape"][2] @@ -248,11 +247,12 @@ def get_compact_scaled_moe_config( use_m_major_input_scale: bool = False, expert_probability_cv: float = DeviceHeuristics.expert_probability_cv, ) -> dict | None: + use_wgmma_ss = layer_config.a_dtype.num_bits == layer_config.b_dtype.num_bits if gemm_type == GemmType.DENSE or use_f16_accum or use_batch_invariant: return None if layer_config.a_dtype.num_bits != 8 or layer_config.use_packed_k_layout: return None - if layer_config.use_raw_weight or layer_config.shape_n % 64 or layer_config.shape_k % 64: + if use_wgmma_ss or layer_config.shape_n % 64 or layer_config.shape_k % 64: return None has_short_k = layer_config.shape_k <= 512 has_narrow_n = layer_config.shape_n % 128 != 0 @@ -365,17 +365,16 @@ def get_overlap_config( expert_probability_cv: float = DeviceHeuristics.expert_probability_cv, ) -> dict | None: layer = layer_config + use_wgmma_ss = layer.a_dtype.num_bits == layer.b_dtype.num_bits is_dense = gemm_type == GemmType.DENSE is_masked = gemm_type == GemmType.GROUPED_MASKED is_a16 = layer.a_dtype.num_bits == 16 - is_unpacked_a8 = ( - layer.a_dtype.num_bits == 8 and not layer.use_raw_weight and not layer.use_packed_k_layout - ) + is_unpacked_a8 = layer.a_dtype.num_bits == 8 and not use_wgmma_ss and not layer.use_packed_k_layout can_overlap_a16 = is_a16 and not is_dense and (not is_h20 or is_masked) and layer.b_dtype.num_bits < 16 can_overlap_packed = layer.use_packed_k_layout and not is_h20 and (not is_dense or shape_m <= 64) - can_overlap_block = layer.use_raw_weight and layer.is_block_weight_scale and not is_dense and not is_h20 - can_overlap_small_raw = layer.use_raw_weight and is_dense and not is_h20 and shape_m <= 64 - has_a8_schedule = is_unpacked_a8 or can_overlap_packed or can_overlap_block or can_overlap_small_raw + can_overlap_block = use_wgmma_ss and layer.is_block_weight_scale and not is_dense and not is_h20 + can_overlap_small_ss = use_wgmma_ss and is_dense and not is_h20 and shape_m <= 64 + has_a8_schedule = is_unpacked_a8 or can_overlap_packed or can_overlap_block or can_overlap_small_ss if not (can_overlap_a16 or (layer.a_dtype.num_bits == 8 and has_a8_schedule)): return None if use_f16_accum or use_batch_invariant or layer.shape_k <= 512: @@ -504,6 +503,7 @@ def get_short_k_config( is_h20: bool = False, use_m_major_input_scale: bool = False, ) -> dict | None: + use_wgmma_ss = layer_config.a_dtype.num_bits == layer_config.b_dtype.num_bits if gemm_type != GemmType.DENSE or use_f16_accum or use_batch_invariant: return None if layer_config.a_dtype.num_bits != 8 or layer_config.shape_k > 512: @@ -512,7 +512,7 @@ def get_short_k_config( return None if get_problem_rejection_reasons(layer_config, MmaType.WGMMA): return None - if is_h20 and layer_config.use_raw_weight and layer_config.shape_k > 256 and shape_m > 64: + if is_h20 and use_wgmma_ss and layer_config.shape_k > 256 and shape_m > 64: return None if layer_config.use_fused_e8m0_scale and layer_config.use_packed_k_layout and shape_m > 64: return None @@ -522,7 +522,7 @@ def get_short_k_config( block_m = int(shape_m * n_tiles / current_device.sm_count) // alignment * alignment block_m = min(64, max(16, block_m)) has_unpacked_group_weights = ( - not layer_config.use_raw_weight + not use_wgmma_ss and not layer_config.use_packed_k_layout and layer_config.is_group_weight_scale and not layer_config.use_fused_e8m0_scale @@ -856,6 +856,7 @@ def select_grouped_scale( problem: TuningProblem, ) -> TuningDecision: layer_config = problem.layer_config + use_wgmma_ss = layer_config.a_dtype.num_bits == layer_config.b_dtype.num_bits if problem.use_f16_accum: max_block_m = 256 elif layer_config.input_scale_group_size > 0: @@ -879,7 +880,7 @@ def select_grouped_scale( candidates = [] # Candidate order records measured preference; legality supplies fallbacks. - tile_shapes = ((128, 32), (128, 16), (64, 16)) if layer_config.use_raw_weight else ((128, 32), (64, 16)) + tile_shapes = ((128, 32), (128, 16), (64, 16)) if use_wgmma_ss else ((128, 32), (64, 16)) for block_shape_n, warp_shape_n in tile_shapes: for block_shape_k in block_ks: multicast_values = (True, False) if use_multicast else (False,) @@ -970,15 +971,16 @@ def get_scaled_a8_config( use_m_major_input_scale: bool = False, expert_probability_cv: float = DeviceHeuristics.expert_probability_cv, ) -> dict | None: - has_raw_w8a8 = layer_config.a_dtype.num_bits == 8 and layer_config.a_dtype == layer_config.b_dtype + use_wgmma_ss = layer_config.a_dtype.num_bits == layer_config.b_dtype.num_bits + has_native_w8a8 = layer_config.a_dtype.num_bits == 8 and layer_config.a_dtype == layer_config.b_dtype has_unpacked_weights = ( layer_config.a_dtype.num_bits == 8 - and not layer_config.use_raw_weight + and not use_wgmma_ss and not layer_config.use_packed_k_layout and not layer_config.is_group_input_scale and layer_config.shape_k > 512 ) - if not (has_raw_w8a8 or has_unpacked_weights): + if not (has_native_w8a8 or has_unpacked_weights): return None if use_f16_accum or use_batch_invariant: return None diff --git a/humming/utils/smem.py b/humming/utils/smem.py index 46c678df..258c83f5 100644 --- a/humming/utils/smem.py +++ b/humming/utils/smem.py @@ -56,7 +56,7 @@ def _stage_storage_bytes( expand_low_bit_weight = use_umma_ss and a_bits == 8 and b_bits < 8 weight_k_alignment = math.gcd(block_k, 128) weight_stage_k = round_up(block_k + 128 - weight_k_alignment, 128) if expand_low_bit_weight else block_k - weight_alignment = 1024 if layer_config.use_raw_weight else 128 + weight_alignment = 1024 if layer_config.use_raw_weight or layer_config.use_flat_weight else 128 weight_stage_n = max(block_n, 128) if use_umma_ss else block_n fields.append((weight_stage_n * weight_stage_k * weight_smem_bits // 8, weight_alignment))