Skip to content

Commit d2c1023

Browse files
xiaolil1xiaolil1
committed
refine code
1 parent 78948e4 commit d2c1023

3 files changed

Lines changed: 101 additions & 39 deletions

File tree

‎csrc/xpu_cutlass_fusion.cpp‎

Lines changed: 97 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -64,28 +64,28 @@ using TiledMma =
6464
Layout<Shape<_1, _2, _1>, Stride<_2, _1, _0>>>::TiledMMA;
6565

6666
using WorkgroupTileShape = TileShape;
67-
static constexpr auto BLK_M = get<0>(WorkgroupTileShape{});
68-
static constexpr auto BLK_N = get<1>(WorkgroupTileShape{});
69-
static constexpr auto BLK_K = get<2>(WorkgroupTileShape{});
67+
static constexpr auto BLK_M = get<0>(WorkgroupTileShape{}); //16
68+
static constexpr auto BLK_N = get<1>(WorkgroupTileShape{}); //64
69+
static constexpr auto BLK_K = get<2>(WorkgroupTileShape{}); //64
7070

7171
//Threads number
72-
static constexpr auto ATOM_M = get<1>(typename TiledMma::ThrLayoutVMNK{}.shape());
73-
static constexpr auto ATOM_N = get<2>(typename TiledMma::ThrLayoutVMNK{}.shape());
74-
static constexpr auto ATOM_K = get<3>(typename TiledMma::ThrLayoutVMNK{}.shape());
72+
static constexpr auto ATOM_M = get<1>(typename TiledMma::ThrLayoutVMNK{}.shape()); //1
73+
static constexpr auto ATOM_N = get<2>(typename TiledMma::ThrLayoutVMNK{}.shape()); //2
74+
static constexpr auto ATOM_K = get<3>(typename TiledMma::ThrLayoutVMNK{}.shape()); //1
7575

7676
static_assert(BLK_M % TiledMma{}.template tile_size_mnk<0>() == 0, "TiledMma permutation size must match block size.");
7777
static_assert(BLK_N % TiledMma{}.template tile_size_mnk<1>() == 0, "TiledMma permutation size must match block size.");
7878
static_assert(BLK_K % TiledMma{}.template tile_size_mnk<2>() == 0, "TiledMma permutation size must match block size.");
7979

8080
//sub-tile shape
81-
static constexpr auto SG_M = ceil_div(BLK_M, ATOM_M);
82-
static constexpr auto SG_N = ceil_div(BLK_N, ATOM_N);
83-
static constexpr auto SG_K = ceil_div(BLK_K, ATOM_K);
84-
using SubgroupTileShape = Shape<decltype(SG_M), decltype(SG_N), decltype(SG_K)>;
81+
static constexpr auto SG_M = ceil_div(BLK_M, ATOM_M); //16
82+
static constexpr auto SG_N = ceil_div(BLK_N, ATOM_N); //32
83+
static constexpr auto SG_K = ceil_div(BLK_K, ATOM_K); //64
84+
using SubgroupTileShape = Shape<decltype(SG_M), decltype(SG_N), decltype(SG_K)>; //<16, 32, 64>
8585

8686
//Total Threads number
87-
static constexpr auto Num_SGs = ATOM_N * ATOM_M * ATOM_K;
88-
static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{});
87+
static constexpr auto Num_SGs = ATOM_N * ATOM_M * ATOM_K; //2
88+
static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{}); //1*2*1*16=32
8989

9090
// Define Mainloop dispatch policy
9191
constexpr int PipelineStages = 3;
@@ -227,7 +227,16 @@ class kgemm_4bit_inference_cutlass_dequant {
227227
using SrcType = typename EngineIn::value_type;
228228
using DstType = typename EngineOut::value_type;
229229
//using ScaleType = typename EngineScales::value_type;
230-
230+
#if 1
231+
int numbers = decltype(size(in))::value;
232+
for(int i=0; i<numbers; i++){
233+
//auto in_ptr_8 = (uint8_t*)(raw_pointer_cast(in.data()));
234+
//out[i] = static_cast<DstType>(quant_map[in_ptr_8[i].data()]);
235+
uint8_t value = in[i].get();
236+
out[i] = static_cast<DstType>(quant_map[value]);
237+
if(cute::thread0()) printf("thread_idx = %d, i = %d, value_bit = %x, value = %d, quant_map[value] = %f, out[i] = %f\n",int(ThreadIdxX()), i, value, static_cast<int>(value), quant_map[value], static_cast<float>(out[i]));
238+
}
239+
#else
231240
static constexpr auto N = decltype(size<1>(in))::value;
232241
233242
using format_type = ushort; //16
@@ -245,37 +254,29 @@ class kgemm_4bit_inference_cutlass_dequant {
245254
auto s_tensor = make_tensor((format_type*)(raw_pointer_cast(in.data())), Shape<Int<loop_cnt / scalar>, Int<N>>{});
246255
auto d_tensor = make_tensor(out.data(), Shape<Int<vec_size>, Int<splits>, Int<N>>{});
247256
248-
//if(cute::thread0())
257+
if(cute::thread0())
249258
printf("thread_idx = %d, decltype(size(in))::value = %d, K = %d, N = %d, L = %d, src_bits = %d, sizeof_bits_v<format_type> = %d, scalar = %d, decltype(size(out))::value = %d, loop_cnt = %d, splits = %d\n",int(ThreadIdxX()), decltype(size(in))::value, decltype(size<0>(in))::value, N, decltype(size<2>(in))::value, src_bits, sizeof_bits_v<format_type>, scalar, decltype(size(out))::value, loop_cnt, splits);
250259
251-
CUTLASS_PRAGMA_UNROLL
252260
for (int n = 0; n < N; n++) {
253261
//const auto ts = tCrS_input(n);
254262
255263
auto& src = *(cute::array<format_type, loop_cnt / scalar>*)(s_tensor(_, n).data());
256264
257-
CUTLASS_PRAGMA_UNROLL
258265
for (int s = 0; s < splits; s++) {
259266
auto idx = vec_size * s / scalar;
260267
auto format_data = src[idx];
261268
262269
auto& dst = *(cute::array<DstType, vec_size>*)(d_tensor(_, s, n).data());
263270
264-
CUTLASS_PRAGMA_UNROLL
265271
for (int i = 0; i < vec_size; i++) {
266-
auto data = [&](){
267-
if constexpr (cutlass::platform::numeric_limits<SrcType>::is_signed) {
268-
return static_cast<SrcType>((format_data >> (src_bits * i)) & 0xf);
269-
} else {
270-
return (format_data >> (src_bits * i)) & 0xf;
271-
}
272-
}();
273-
274-
int8_t minus(data);
275-
dst[i] = (static_cast<DstType>(quant_map[minus]));// * ts;
272+
uint8_t value = (format_data >> (src_bits * i)) & 0xf;
273+
dst[i] = (static_cast<DstType>(quant_map[value]));// * ts;
274+
//if(cute::thread0())
275+
printf("n = %d, s = %d, i = %d, src = %d, dst = %f\n", n, s, i, static_cast<int>(value), static_cast<float>(dst[i]));
276276
}
277277
}
278278
}
279+
#endif
279280
}
280281

281282
CUTLASS_DEVICE
@@ -329,8 +330,12 @@ class kgemm_4bit_inference_cutlass_dequant {
329330
l_coord = BlockIdxZ();
330331
}
331332
auto blk_coord_mnkl = make_coord(m_coord, n_coord, _, l_coord);
332-
if(cute::thread0()) printf("M = %d, N=%d, K=%d, L=%d, m_coord = %d, n_coord = %d, l_coord = %d, BlockIdxX() = %d, BlockIdxY() = %d, BlockIdxZ() = %d\n",M, N, K, L, m_coord, n_coord, l_coord, BlockIdxX(), BlockIdxY(), BlockIdxZ());
333+
if(cute::thread0()) {
334+
printf("M = %d, N=%d, K=%d, L=%d\n", M, N, K, L);
335+
//}
336+
printf("thread_idx = %d, m_coord = %d, n_coord = %d, l_coord = %d, BlockIdxX() = %d, BlockIdxY() = %d, BlockIdxZ() = %d\n",thread_idx, m_coord, n_coord, l_coord, BlockIdxX(), BlockIdxY(), BlockIdxZ());
333337

338+
}
334339
constexpr auto workgroup_shape = WorkgroupTileShape{}; //256, 256, 32
335340
constexpr auto subgroup_tile_shape = SubgroupTileShape{}; //32, 64, 32 (number of atom level workgroup: 256/8=32, 256/4=64, 32/2=32)
336341

@@ -368,7 +373,7 @@ class kgemm_4bit_inference_cutlass_dequant {
368373
Tensor mma_A = make_tensor<ElementMMA>(make_fragment_layout(tiled_copy_a, tCgA(_,_,_,0).shape()));
369374
Tensor mma_B = make_tensor<ElementMMA>(make_fragment_layout(tiled_copy_b, tCgB(_,_,_,0).shape()));
370375

371-
Tensor dequant_frag = make_tensor<ElementB>(mma_B.layout());
376+
Tensor dequant_frag = make_tensor<ElementB>(mma_B.layout());
372377

373378
//static constexpr auto scale_traits_size = decltype(size(typename GmemTiledCopyScale::BlockShape{}))::value / SubgroupSize;
374379
//static constexpr auto scale_traits_num = SG_QNT_WIDTH / size<1>(typename GmemTiledCopyScale::BlockShape{});
@@ -410,6 +415,34 @@ class kgemm_4bit_inference_cutlass_dequant {
410415
//
411416
// }();
412417

418+
#define PRINT(x) print(#x ": "); print(x); print("\n");
419+
if (cutlass::thread(LOG_THREAD, LOG_GROUP)) {
420+
print("======================= A: \n");
421+
print(" gA : "); print(gA); print("\n");
422+
print(" tCgA : "); print(tCgA); print("\n");
423+
print(" tAgA : "); print(tAgA); print("\n");
424+
print(" mma_A : "); print(mma_A); print("\n");
425+
print(" frag_copy_A : "); print(frag_copy_A); print("\n");
426+
427+
print("===================== B :\n");
428+
print(" gB : "); print(gB); print("\n");
429+
print(" tCgB : "); print(tCgB); print("\n");
430+
print(" tBgB : "); print(tBgB); print("\n");
431+
print(" mma_B : "); print(mma_B); print("\n");
432+
print(" frag_copy_B : "); print(frag_copy_B); print("\n");
433+
print(" dequant_frag : "); print(dequant_frag); print("\n");
434+
435+
print("===================== Config: \n");
436+
print(" threads per workgroup : "); print(MaxThreadsPerBlock); print("\n");
437+
print(" SubgroupTileShape : "); print(SubgroupTileShape{}); print("\n");
438+
439+
print(" tiled_prefetch_a : "); print(tiled_prefetch_a); print("\n");
440+
print(" tiled_prefetch_b : "); print(tiled_prefetch_b); print("\n");
441+
print(" pAgA : "); print(pAgA); print("\n");
442+
print(" pBgB : "); print(pBgB); print("\n");
443+
}
444+
#undef PRINT
445+
413446
const int k_start_idx = crd2idx((*k_tile_iter), make_shape(K));
414447
int prefetch_k = k_start_idx;
415448

@@ -420,9 +453,7 @@ class kgemm_4bit_inference_cutlass_dequant {
420453
}
421454

422455
for (int k_tile = k_start_idx; k_tile < k_tile_count + k_start_idx; k_tile++, prefetch_k++) {
423-
constexpr int barrier_scope = 2;
424-
425-
barrier_arrive(barrier_scope);
456+
barrier_arrive(2);
426457

427458
// Copy gmem to rmem for the first k_tile
428459
copy(tiled_copy_a, tAgA(_,_,_,k_tile), frag_copy_A);
@@ -434,13 +465,41 @@ class kgemm_4bit_inference_cutlass_dequant {
434465

435466
if(prefetch_k < k_tile_count) {
436467
prefetch(tiled_prefetch_a, pAgA(_,_,_,prefetch_k));
468+
}
469+
if(prefetch_k < k_tile_count / 2) {
437470
prefetch(tiled_prefetch_b, pBgB(_,_,_,prefetch_k));
438471
}
439472

440473
dequant(dequant_frag, mma_B, /*fragment_scale,*/ quant_map);
441474

442475
cute::gemm(tiled_mma, mma_A, mma_B, accumulators);
443-
barrier_wait(barrier_scope);
476+
477+
//// 在调用gemm前后添加打印逻辑
478+
//auto debug_print = [&](const char* name, auto& tensor) {
479+
// if (thread_idx == 0) {
480+
// printf("----- %s -----\n", name);
481+
// for (int i = 0; i < size<0>(tensor); ++i) {
482+
// for (int j = 0; j < size<1>(tensor); ++j) {
483+
// printf("%6.2f ", static_cast<float>(tensor(i, j)));
484+
// }
485+
// printf("\n");
486+
// }
487+
// }
488+
// barrier_wait(2);
489+
//};
490+
//
491+
//// 打印输入
492+
//debug_print("Input A (mma_A)", mma_A);
493+
//debug_print("Input B (mma_B)", mma_B);
494+
//debug_print("Accumulators (Before GEMM)", accumulators);
495+
//
496+
//// 执行GEMM
497+
//cute::gemm(tiled_mma, mma_A, mma_B, accumulators);
498+
//
499+
//// 打印输出
500+
//debug_print("Accumulators (After GEMM)", accumulators);
501+
502+
barrier_wait(2);
444503
}
445504

446505
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>((char*)nullptr);
@@ -458,11 +517,12 @@ class kgemm_4bit_inference_cutlass_dequant {
458517
};
459518

460519
template <typename T, int BITS>
461-
void gemm_4bit_inference_cutlass_dequant(int m, int n, int k, T *A, unsigned char *B,
520+
void gemm_4bit_inference_cutlass_dequant(int m, int n, int k_, T *A, unsigned char *B,
462521
T *absmax_, float *datatype, float *out, int lda,
463522
int ldb, int ldc, int blocksize, sycl::queue *stream) {
464523
std::cout<<"this is gemm_4bit_inference_cutlass_dequant ......................!!!!!!\n";
465524

525+
int k = k_;
466526

467527

468528
sycl::queue q = *stream;
@@ -508,7 +568,7 @@ void gemm_4bit_inference_cutlass_dequant(int m, int n, int k, T *A, unsigned cha
508568
// (n * 4 ) / 8,
509569
// (n * k * 4 ) / 8
510570
// );
511-
auto mB_nkl = make_tensor(cute::subbyte_iterator<const ElementB>(B), make_layout(make_shape(n, k, l), stride_B));
571+
auto mB_nkl = make_tensor(cute::subbyte_iterator<uint4_t>(B), make_layout(make_shape(n, k, l), stride_B));
512572
Copy_B tiled_copy_b{Copy_B{}.with(mB_nkl)};
513573

514574
params.tiled_copy_a = tiled_copy_a;
@@ -546,8 +606,8 @@ void gemm_4bit_inference_cutlass_dequant(int m, int n, int k, T *A, unsigned cha
546606
dim3 const block = GemmKernel::get_block_shape();
547607
dim3 const grid = GemmKernel::get_grid_shape(params);
548608

549-
const syclcompat::dim3 sycl_block(block.x, block.y, block.z); //workgroup_size: 8*4*1*16, 1, 1
550-
const syclcompat::dim3 sycl_grid(grid.x, grid.y, grid.z); //workgroup_number (problem_size / tile_size): N/256, M/256, 1
609+
const syclcompat::dim3 sycl_block(block.x, block.y, block.z); //workgroup_size: 1*2*1*16, 1, 1
610+
const syclcompat::dim3 sycl_grid(grid.x, grid.y, grid.z); //workgroup_number (problem_size / tile_size): N/64, M/16, 1
551611
printf("Host Grid: (%d, %d, %d)\n", grid.x, grid.y, grid.z);
552612
printf("Host Block: (%d, %d, %d)\n", block.x, block.y, block.z);
553613

‎run_case.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828

2929

3030
#gdb -args python -m pytest -vs tests/test_xpu.py::TestXPU::test_gemm_4bit
31-
pytest -vs tests/test_xpu.py::TestXPU::test_gemm_4bit
32-
#python tests/test_xpu_db.py
31+
#pytest -vs tests/test_xpu.py::TestXPU::test_gemm_4bit
32+
python tests/test_xpu_db.py
3333
#gdb -args python tests/test_xpu_db.py
3434
#pytest tests/test_functional.py::TestQuantize4BitFunctional::test_gemv_4bit[dim=256-uint8-bf16-fc1-nf4-DQ_True-xpu]

‎tests/test_xpu.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,8 @@ def test_gemm_4bit(self, device, dim, dtype, storage_type, quant_storage, double
9191
#pdb.set_trace()
9292
C2 = F.gemv_4bit(A, qB.t(), state=state)
9393
#pdb.set_trace()
94+
print("C3.sum() = ", C3.sum())
95+
print("C2.sum() = ", C2.sum())
9496
diff = abs(C2-C3)
9597
print("diff = ", diff.sum())
9698
print(C3[0])

0 commit comments

Comments
 (0)