Skip to content

Commit 0aa000a

Browse files
jinzhen-lincodex
jinzhen-lin
and
codex
authored
perf(cuda): release MMA stages earlier and flatten native weight tiles (#137)
Signed-off-by: jinzhen.ljz <jinzhen.ljz@antgroup.com> Co-authored-by: Codex (GPT-6) <codex@openai.com>
1 parent c6756f7 commit 0aa000a

19 files changed

Lines changed: 132 additions & 105 deletions

File tree

‎humming/config/config.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ class LayerConfig(BaseHummingConfig):
8080
"is_token_input_scale_2",
8181
"is_tensor_input_scale_2",
8282
"use_raw_weight",
83+
"use_flat_weight",
8384
"use_native_dequant",
8485
"use_block_scaled_mma",
8586
"mma_scale_group_size",
@@ -132,7 +133,7 @@ def is_umma_supported(self):
132133
@property
133134
def use_raw_weight(self):
134135
if self.a_dtype.num_bits == self.b_dtype.num_bits:
135-
return True
136+
return self.is_umma_supported or current_device.is_ppu
136137
has_native_mixed_activation = self.a_dtype in (
137138
dtypes.float8e4m3,
138139
dtypes.float8e5m2,
@@ -145,6 +146,11 @@ def use_raw_weight(self):
145146
)
146147
return has_native_mixed_activation and has_native_mixed_weight and self.is_umma_supported
147148

149+
@property
150+
def use_flat_weight(self):
151+
has_same_width = self.a_dtype.num_bits == self.b_dtype.num_bits
152+
return has_same_width and not self.is_umma_supported and not current_device.is_ppu
153+
148154
@property
149155
def use_block_scaled_mma(self):
150156
if self.mxmma_supported:
@@ -168,7 +174,7 @@ def mma_scale_group_size(self):
168174
@property
169175
def use_native_dequant(self):
170176
assert self.sm_version is not None
171-
if self.use_raw_weight or self.sm_version // 10 == 9:
177+
if self.use_raw_weight or self.use_flat_weight or self.sm_version // 10 == 9:
172178
return False
173179

174180
from humming.jit.runtime import KernelRuntime
@@ -379,6 +385,7 @@ def __post_init__(self):
379385
can_use_packed_k = (
380386
self.sm_version // 10 == 9
381387
and not self.use_raw_weight
388+
and not self.use_flat_weight
382389
and has_packed_k_operands
383390
and has_packed_k_input_scale
384391
)
@@ -389,7 +396,9 @@ def __post_init__(self):
389396
self.use_packed_k_layout = can_use_packed_k and prefer_packed_k
390397
elif self.use_packed_k_layout:
391398
assert self.sm_version // 10 == 9, "use_packed_k_layout requires SM90"
392-
assert not self.use_raw_weight, "raw weights do not use packed-K layout"
399+
assert not (self.use_raw_weight or self.use_flat_weight), (
400+
"raw/flat weights do not use packed-K layout"
401+
)
393402
assert self.a_dtype.num_bits == 8, "use_packed_k_layout requires 8-bit activation"
394403
assert self.b_dtype.num_bits % 2 == 0, "use_packed_k_layout requires even-bit weight"
395404

‎humming/config/mma.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,10 +174,11 @@ def get_register_budget_error(layer_config, tuning_config, use_f16_accum=False,
174174
return None # UMMA accumulators use TMEM and have separate resource checks.
175175

176176
if mma_type == MmaType.WGMMA:
177+
use_wgmma_ss = layer_config.a_dtype.num_bits == layer_config.b_dtype.num_bits
177178
# ptxas needs 26 (SS) or 32 (RS) registers beyond one N=64 accumulator tuple.
178179
# setmaxnreg cannot relax this per-instruction compile limit.
179180
instruction_accumulators = warp_m // (4 if use_f16_accum else 2)
180-
instruction_overhead = 26 if layer_config.use_raw_weight else 32
181+
instruction_overhead = 26 if use_wgmma_ss else 32
181182
instruction_registers = instruction_accumulators + instruction_overhead
182183
if layer_config.use_packed_k_layout and tuning_config.wgmma_split_issue_wait:
183184
# Split packed-K keeps all RS K slabs live across prefetch. The

‎humming/csrc/launcher/launcher.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -336,6 +336,7 @@ std::tuple<int64_t, std::string> register_kernel(const std::string &cubin_path)
336336
reader.getBool("USE_PACKED_K_LAYOUT"),
337337
reader.getBool("USE_UMMA_SS"),
338338
reader.getBool("USE_RAW_WEIGHT"),
339+
reader.getBool("USE_FLAT_WEIGHT"),
339340
reader.getBool("USE_BLOCK_SCALED_MMA")};
340341

341342
std::unique_lock lock(g_kernel_mutex);

‎humming/csrc/launcher/tensor.h‎

Lines changed: 19 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ inline void check_tensor_b(Tensor &tensor, KernelData &kernel_data, int64_t dev)
8787

8888
std::vector<int64_t> expected_shape = {};
8989
if (kernel_data.gemm_type_id != 0) expected_shape.push_back(kernel_data.num_experts);
90-
if (kernel_data.use_raw_weight) {
90+
if (kernel_data.use_raw_weight || kernel_data.use_flat_weight) {
9191
if (get_dtype_num_bits(kernel_data.a_dtype_id) == 8 && num_bits < 8)
9292
problem_shape_k = CEIL_DIV(problem_shape_k, 128) * 128;
9393
expected_shape.push_back(problem_shape_n);
@@ -390,6 +390,24 @@ inline CUtensorMap make_tma_desc_b(Tensor &tensor, KernelData &kernel_data) {
390390
uint32_t block_shape_n = kernel_data.block_shape_n;
391391
uint32_t block_shape_k = kernel_data.block_shape_k;
392392

393+
if (kernel_data.use_flat_weight) {
394+
CUtensorMap descriptor{};
395+
uint64_t warp_n = kernel_data.mma_type == MmaType::WGMMA ? kernel_data.warp_shape_n : 16;
396+
uint64_t shape_n = kernel_data.problem_shape_n;
397+
uint64_t num_experts = std::max(1u, kernel_data.num_experts);
398+
uint64_t row_bytes = uint64_t(tensor.size(-1)) * 4;
399+
uint64_t dimensions[] = {64, 64 / warp_n, shape_n / 64, warp_n / 16, row_bytes / 16 * num_experts};
400+
uint64_t strides[] = {warp_n * 16, 1024, 256, shape_n * 16};
401+
uint32_t box[] = {64, uint32_t(std::min(64u, block_shape_n) / warp_n),
402+
(block_shape_n + 63) / 64, uint32_t(warp_n / 16), block_shape_k * num_bits / 128};
403+
uint32_t element_strides[] = {1, 1, 1, 1, 1};
404+
CUresult status = cuTensorMapEncodeTiled(&descriptor, CU_TENSOR_MAP_DATA_TYPE_INT32, 5,
405+
tensor.data_ptr(), dimensions, strides, box, element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE,
406+
CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
407+
ASSERT_CHECK(status == CUDA_SUCCESS, "TMA Encode Failed for flat weights: ", int(status));
408+
return descriptor;
409+
}
410+
393411
if (kernel_data.use_raw_weight) {
394412
// FP4/FP6 in an f8f6f4 instruction uses padded 16-element blocks in SMEM.
395413
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) {
408426
ASSERT_CHECK(status == CUDA_SUCCESS, "TMA Encode Failed for expanded low-bit weights");
409427
return descriptor;
410428
}
411-
if (kernel_data.mma_type == MmaType::WGMMA) {
412-
CUtensorMap descriptor{};
413-
uint64_t w = kernel_data.warp_shape_n;
414-
uint64_t row_bytes = uint64_t(tensor.size(-1)) * 4;
415-
// Fragment planes span all N warps in the CTA, leaving the fifth
416-
// dimension available for K slabs even with multiple N warpgroups.
417-
uint64_t dimensions[] = {swizzle_bytes / 4, 16,
418-
uint64_t(kernel_data.problem_shape_n) * std::max(1u, kernel_data.num_experts) / w,
419-
w / 16, row_bytes / swizzle_bytes};
420-
uint64_t strides[] = {row_bytes, w * row_bytes, 16 * row_bytes, swizzle_bytes};
421-
uint32_t box[] = {swizzle_bytes / 4, 16, uint32_t(block_shape_n / w), uint32_t(w / 16),
422-
block_shape_k * num_bits / 8 / swizzle_bytes};
423-
uint32_t element_strides[] = {1, 1, 1, 1, 1};
424-
CUresult status = cuTensorMapEncodeTiled(&descriptor, CU_TENSOR_MAP_DATA_TYPE_INT32, 5,
425-
tensor.data_ptr(), dimensions, strides, box, element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE,
426-
get_swizzle_enum(swizzle_bytes), CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
427-
ASSERT_CHECK(status == CUDA_SUCCESS, "TMA Encode Failed for WGMMA N permutation: ", int(status));
428-
return descriptor;
429-
}
430429
auto rows = torch_view_shape(tensor, {-1, tensor.size(-1)});
431430
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");
432431
}

‎humming/csrc/launcher/utils.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,7 @@ struct KernelData {
179179
bool use_packed_k_layout;
180180
bool use_umma_ss;
181181
bool use_raw_weight;
182+
bool use_flat_weight;
182183
bool use_block_scaled_mma;
183184
};
184185

‎humming/include/humming/kernel/humming_ws.cuh‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,9 @@ __global__ __launch_bounds__(TuningConfig::kNumThreads, TuningConfig::kNumCtasPe
189189
if (warp_iter_id == Ctx::kWarpIters - 1 && slice_iter + 1 < num_slice_iters) {
190190
consumer.wait_stage((stage_id + 1) % kNumStages);
191191
}
192+
if constexpr (!Ctx::kUseWgmma) {
193+
if (warp_iter_id == Ctx::kWarpIters - 1) consumer.arrive(stage_id);
194+
}
192195
const bool use_late_scale_prefetch = Ctx::kUseWgmmaLateScalePrefetch && warp_iter_id == Ctx::kWarpIters - 1;
193196
const bool has_next_fragment = warp_iter_id + 1 < Ctx::kWarpIters || slice_iter + 1 < num_slice_iters;
194197
if (has_next_fragment) {
@@ -208,7 +211,9 @@ __global__ __launch_bounds__(TuningConfig::kNumThreads, TuningConfig::kNumCtasPe
208211
((warp_iter_id + 1) % Ctx::kWarpIters) % 2,
209212
(warp_iter_id + 1) % Ctx::kWarpIters);
210213
}
211-
if (warp_iter_id == Ctx::kWarpIters - 1) consumer.arrive(stage_id);
214+
if constexpr (Ctx::kUseWgmma) {
215+
if (warp_iter_id == Ctx::kWarpIters - 1) consumer.arrive(stage_id);
216+
}
212217
}
213218
};
214219

‎humming/include/humming/memory/g2s_loader/loader_b.cuh‎

Lines changed: 44 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,8 @@ public:
5454

5555
template <bool kShouldAdvance = true>
5656
CUDA_INLINE void load(int4 *smem_ptr, void *mbar_ptr) {
57-
if constexpr (Ctx::kUseRawWeight && !kUseTma) load_raw(smem_ptr);
57+
if constexpr (Ctx::kUseFlatWeight && !kUseTma) load_flat(smem_ptr);
58+
else if constexpr (Ctx::kUseRawWeight && !kUseTma) load_raw(smem_ptr);
5859
else if constexpr (kUseAiu) load_aiu(smem_ptr);
5960
else if constexpr (kUseTma) load_tma(smem_ptr, mbar_ptr);
6061
else load_legacy(smem_ptr);
@@ -76,22 +77,21 @@ public:
7677
CUDA_INLINE
7778
void load_tma(int4 *smem_ptr, void *mbar_ptr) {
7879
if (ctx.load_thread_id() == 0) {
79-
if constexpr (Ctx::kUseRawWeight) {
80+
if constexpr (Ctx::kUseFlatWeight) {
81+
uint32_t expert_k = col_offset / ProblemShape::N * ProblemShape::K + row_offset;
82+
if (cluster_rank == 0)
83+
tma_load_5d<kMultiCastSizeB>(tensor_map_ptr, smem_ptr, mbar_ptr,
84+
0, col_offset % 64 / 16, col_offset % ProblemShape::N / 64, 0, expert_k * ElementQB::kBits / 128);
85+
} else if constexpr (Ctx::kUseRawWeight) {
8086
constexpr uint32_t kCoordinateBits = SharedStorage::kExpandUmmaWeight ? 1 : 32 / ElementQB::kBits;
8187
uint32_t start_k = kCoordinateBits == 1 ? row_offset / 128 * 128 : row_offset;
82-
if constexpr (Ctx::kUseWgmmaTmaKPack) {
83-
if (cluster_rank == 0)
84-
tma_load_5d<kMultiCastSizeB>(tensor_map_ptr, smem_ptr, mbar_ptr,
85-
0, 0, col_offset / Ctx::WarpShape::N, 0, start_k / kSwizzleK);
86-
} else {
88+
PRAGMA_UNROLL
89+
for (uint32_t k = 0; k < SharedStorage::kWeightStageK; k += kSwizzleK) {
8790
PRAGMA_UNROLL
88-
for (uint32_t k = 0; k < SharedStorage::kWeightStageK; k += kSwizzleK) {
89-
PRAGMA_UNROLL
90-
for (uint32_t n = 0; n < SharedStorage::kWeightStageN; n += 256) {
91-
uint32_t offset = (k / kSwizzleK * SharedStorage::kWeightStageN + n) * kSwizzleBytes / sizeof(int4);
92-
if (cluster_rank == 0) tma_load_2d<kMultiCastSizeB, kEvictWeightsFirst, Ctx::kUseUmmaCooperativeTma ? 2 : 1>(
93-
tensor_map_ptr, smem_ptr + offset, mbar_ptr, (start_k + k) / kCoordinateBits, col_offset + n);
94-
}
91+
for (uint32_t n = 0; n < SharedStorage::kWeightStageN; n += 256) {
92+
uint32_t offset = (k / kSwizzleK * SharedStorage::kWeightStageN + n) * kSwizzleBytes / sizeof(int4);
93+
if (cluster_rank == 0) tma_load_2d<kMultiCastSizeB, kEvictWeightsFirst, Ctx::kUseUmmaCooperativeTma ? 2 : 1>(
94+
tensor_map_ptr, smem_ptr + offset, mbar_ptr, (start_k + k) / kCoordinateBits, col_offset + n);
9595
}
9696
}
9797
} else if constexpr (kMultiCastSizeB == 1) {
@@ -106,12 +106,12 @@ public:
106106
void prefetch_tma() {
107107
if constexpr (kUseTma && kMultiCastSizeB == 1) {
108108
if (ctx.load_thread_id() == 0) {
109-
if constexpr (!Ctx::kUseRawWeight) tma_prefetch_3d(tensor_map_ptr, 0, col_offset, row_offset);
109+
if constexpr (!Ctx::kUseRawWeight && !Ctx::kUseFlatWeight) tma_prefetch_3d(tensor_map_ptr, 0, col_offset, row_offset);
110110
}
111111
}
112112
}
113113

114-
CUDA_INLINE static uint32_t permute_raw_ss_row(uint32_t row) {
114+
CUDA_INLINE static uint32_t permute_wgmma_ss_row(uint32_t row) {
115115
constexpr uint32_t kWarpN = Ctx::WarpShape::N;
116116
constexpr uint32_t kGroupN = 4 * kWarpN;
117117
uint32_t row_in_group = row % kGroupN;
@@ -127,36 +127,34 @@ public:
127127
constexpr uint32_t kGlobalRowInt4s = ProblemShape::K * ElementQB::kBits / 128;
128128
constexpr uint32_t kSwizzleInt4s = kSwizzleBytes / sizeof(int4);
129129
uint32_t smem_base = cast_smem_ptr_to_uint(smem_ptr) / 128;
130-
constexpr uint32_t kRowsPerIteration = kNumLoadThreads / kRowInt4s;
131-
constexpr bool kPowerOfTwoRows = kRowsPerIteration != 0 && (kRowsPerIteration & (kRowsPerIteration - 1)) == 0;
132-
constexpr bool kPowerOfTwoWarpN = (Ctx::WarpShape::N & (Ctx::WarpShape::N - 1)) == 0;
133-
constexpr bool kSplitRowPermutation =
134-
Ctx::kUseWgmmaCpAsyncNPermute && kNumLoadThreads % kRowInt4s == 0 && kPowerOfTwoRows && kPowerOfTwoWarpN;
135-
uint32_t permuted_thread_row = 0;
136-
if constexpr (kSplitRowPermutation) {
137-
// The thread and unrolled iteration own disjoint row bits. Permute
138-
// thread bits once; the iteration contribution folds to constants.
139-
uint32_t thread_row = (ctx.load_thread_id() / kRowInt4s) % kRowsPerIteration;
140-
permuted_thread_row = permute_raw_ss_row(thread_row);
141-
}
142130
PRAGMA_UNROLL
143131
for (uint32_t i = 0; i < CEIL_DIV(BlockShape::N * kRowInt4s, kNumLoadThreads); ++i) {
144132
uint32_t index = i * kNumLoadThreads + ctx.load_thread_id();
145133
uint32_t row = index / kRowInt4s;
146134
uint32_t col = index % kRowInt4s;
147-
uint32_t smem_row = row;
148-
if constexpr (Ctx::kUseWgmmaCpAsyncNPermute) {
149-
if constexpr (kSplitRowPermutation) {
150-
smem_row = permuted_thread_row + permute_raw_ss_row(i * kRowsPerIteration);
151-
} else smem_row = permute_raw_ss_row(row);
152-
}
153-
uint32_t offset = (col / kSwizzleInt4s * BlockShape::N + smem_row) * kSwizzleInt4s + col % kSwizzleInt4s;
135+
uint32_t offset = (col / kSwizzleInt4s * BlockShape::N + row) * kSwizzleInt4s + col % kSwizzleInt4s;
154136
uint32_t swizzled_offset = offset ^ ((smem_base + offset / 8) % kSwizzleInt4s);
155137
legacy_load_pred<kUseCpAsync>(gmem_ptr + row * kGlobalRowInt4s + col,
156138
smem_ptr + swizzled_offset, row < BlockShape::N);
157139
}
158140
}
159141

142+
CUDA_INLINE
143+
void load_flat(int4 *smem_ptr) {
144+
constexpr uint32_t kChunkRows = MIN(64u, BlockShape::N);
145+
constexpr uint32_t kTileInt4s = kChunkRows * BlockShape::K * ElementQB::kBits / 128;
146+
PRAGMA_UNROLL
147+
for (uint32_t i = 0; i < CEIL_DIV(kNumInt4s, kNumLoadThreads); ++i) {
148+
uint32_t index = i * kNumLoadThreads + ctx.load_thread_id();
149+
uint32_t row = index / kTileInt4s * 64 + index % kChunkRows;
150+
uint32_t col = index % kTileInt4s / kChunkRows;
151+
if constexpr (Ctx::kUseWgmmaCpAsyncNPermute) row = permute_wgmma_ss_row(row);
152+
uint32_t offset = col * BlockShape::N + row;
153+
uint32_t gmem_offset = col * ProblemShape::N + index / kTileInt4s * 64 + index % kChunkRows;
154+
legacy_load_pred<kUseCpAsync>(gmem_ptr + gmem_offset, smem_ptr + offset, index < kNumInt4s);
155+
}
156+
}
157+
160158
CUDA_INLINE
161159
void load_raw_aiu(int4 *smem_ptr) {
162160
const uint32_t warp_id = ctx.load_thread_id() / 32;
@@ -179,9 +177,9 @@ public:
179177

180178
CUDA_INLINE
181179
void advance() {
182-
if constexpr (Ctx::kUseRawWeight) {
180+
if constexpr (Ctx::kUseRawWeight || Ctx::kUseFlatWeight) {
183181
row_offset += BlockShape::K;
184-
if constexpr (!kUseTma) gmem_ptr += BlockShape::K * ElementQB::kBits / 128;
182+
if constexpr (!kUseTma) gmem_ptr += BlockShape::K * ElementQB::kBits / 128 * (Ctx::kUseFlatWeight ? ProblemShape::N : 1);
185183
return;
186184
}
187185
row_offset += BlockShape::K / kPackSizeK;
@@ -190,12 +188,18 @@ public:
190188

191189
CUDA_INLINE
192190
void seek(uint32_t expert_id, uint32_t n_block_id, uint32_t k_block_id) {
193-
if constexpr (Ctx::kUseRawWeight) {
191+
if constexpr (Ctx::kUseRawWeight || Ctx::kUseFlatWeight) {
194192
row_offset = k_block_id * BlockShape::K;
195193
col_offset = expert_id * ProblemShape::N + n_block_id * BlockShape::N;
196194
if constexpr (!kUseTma) {
197-
uint64_t offset = uint64_t(col_offset) * ProblemShape::K + row_offset;
198-
gmem_ptr = gmem_ptr_raw + offset * ElementQB::kBits / 128;
195+
if constexpr (Ctx::kUseFlatWeight) {
196+
uint64_t expert_offset = uint64_t(expert_id) * ProblemShape::N * ProblemShape::K;
197+
uint64_t offset = expert_offset + uint64_t(row_offset) * ProblemShape::N;
198+
gmem_ptr = gmem_ptr_raw + offset * ElementQB::kBits / 128 + n_block_id * BlockShape::N;
199+
} else {
200+
uint64_t offset = uint64_t(col_offset) * ProblemShape::K + row_offset;
201+
gmem_ptr = gmem_ptr_raw + offset * ElementQB::kBits / 128;
202+
}
199203
}
200204
return;
201205
}

‎humming/include/humming/memory/s2r_loader/loader_b.cuh‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ public:
6262
}
6363

6464
CUDA_INLINE
65-
void load_raw(const int4 *smem_ptr, uint32_t *regs_ptr, uint32_t iter_id) {
65+
void load_native(const int4 *smem_ptr, uint32_t *regs_ptr, uint32_t iter_id) {
6666
if constexpr (USE_PPU && Ctx::kUseCpAsync && !Ctx::kUseTmaB)
6767
return load_raw_aiu(smem_ptr, regs_ptr, iter_id);
6868
constexpr uint32_t kSwizzleInt4s = MIN(8u, BlockShape::K * ElementQB::kBits / 128);
@@ -76,6 +76,7 @@ public:
7676
uint32_t col = (ctx.k_warp_offset() * ElementQB::kBits / 256 + iter_id) * 2 + (lane_id / 8) % 2;
7777
uint32_t offset = (col / kSwizzleInt4s * BlockShape::N + row) * kSwizzleInt4s + col % kSwizzleInt4s;
7878
uint32_t swizzled_offset = offset ^ ((smem_base + offset / 8) % kSwizzleInt4s);
79+
if constexpr (Ctx::kUseFlatWeight) swizzled_offset = col * BlockShape::N + row;
7980
ld_shared<kRowsPerLoad / 4>(smem_ptr + swizzled_offset, reinterpret_cast<int4 *>(regs_ptr + n * kRowsPerLoad / 4));
8081
}
8182
}
@@ -94,7 +95,7 @@ public:
9495

9596
CUDA_INLINE
9697
void load(const int4 *smem_ptr, uint32_t *regs_ptr, uint32_t iter_id) {
97-
if constexpr (Ctx::kUseRawWeight) return load_raw(smem_ptr, regs_ptr, iter_id);
98+
if constexpr (Ctx::kUseRawWeight || Ctx::kUseFlatWeight) return load_native(smem_ptr, regs_ptr, iter_id);
9899
if constexpr (Ctx::kUsePackedKLayout) return load_packed_k(smem_ptr, regs_ptr, iter_id);
99100
uint32_t warp_id = ctx.warp_id();
100101
uint32_t n_warp_id = ctx.n_warp_id();

‎humming/include/humming/memory/s2r_pipeline.cuh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ public:
8585
uint32_t buffer_id = iter_id % 2;
8686
auto &smem = ctx.smem;
8787

88-
if constexpr (!(kUseWgmma && Ctx::kUseRawWeight))
88+
if constexpr (!Ctx::kUseWgmmaSs)
8989
loader_b.load(smem.stages[stage_id].b, mma.regs_qb_as_ptr(buffer_id), iter_id);
9090
if constexpr (!kUseWgmma && !Ctx::kUseUmma)
9191
loader_a.load(smem.stages[stage_id].a, mma.regs_a_as_ptr(buffer_id), iter_id, stage_id);

‎humming/include/humming/mma/mxmma.cuh‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -70,8 +70,7 @@ public:
7070

7171
CUDA_INLINE
7272
void transform_b(uint32_t buffer_id, uint32_t iter_id) {
73-
if constexpr (Ctx::kUseRawWeight) {
74-
// Raw operands are loaded directly into regs_b, including mixed FP4 formats.
73+
if constexpr (Ctx::kUseRawWeight || Ctx::kUseFlatWeight) {
7574
return;
7675
} else if constexpr (kNativeMixed) {
7776
PRAGMA_UNROLL
@@ -180,7 +179,7 @@ public:
180179

181180
template <class T = uint32_t>
182181
CUDA_INLINE T *regs_qb_as_ptr(uint32_t buffer_id) {
183-
if constexpr (Ctx::kUseRawWeight) {
182+
if constexpr (Ctx::kUseRawWeight || Ctx::kUseFlatWeight) {
184183
return reinterpret_cast<T *>(regs_b[buffer_id]);
185184
} else {
186185
return reinterpret_cast<T *>(regs_qb[buffer_id]);

0 commit comments

Comments
 (0)