Skip to content

Commit 572ef63

Browse files
xiaolil1xiaolil1
committed
refine code
1 parent 08fe237 commit 572ef63

1 file changed

Lines changed: 31 additions & 24 deletions

File tree

‎csrc/xpu_cutlass_fusion.cpp‎

Lines changed: 31 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -145,8 +145,8 @@ using atom_load_B = Copy_Atom<traits_load_B, ElementB>;
145145
using val_layout_load_B = decltype(make_layout(shape_div(typename traits_load_B::BlockShape{}, CopyThreadShape{})));
146146
using Copy_B = decltype(make_tiled_copy(atom_load_B{}, Layout<CopyThreadShape>{}, val_layout_load_B{}));
147147

148-
//using GmemTiledCopyScale = XE_2D_U16x1x32_LD_N; //XE_2D_U16x1x16_LD_N;
149-
using GmemTiledCopyScale = XE_2D_U16x1x16_LD_N;
148+
using GmemTiledCopyScale = XE_2D_U16x1x32_LD_N; //XE_2D_U16x1x16_LD_N;
149+
//using GmemTiledCopyScale = XE_2D_U16x1x16_LD_N;
150150
static constexpr auto SG_QNT_WIDTH = Int<SG_N>{};
151151
using StrideScale = cute::Stride<_1, int64_t, int64_t>; //dynamic stride
152152
using traits_load_scale = Copy_Traits<GmemTiledCopyScale, StrideScale>;
@@ -241,11 +241,11 @@ class kgemm_4bit_inference_cutlass_dequant {
241241
//if(syclcompat::global_id::x() == 2 && syclcompat::global_id::y() ==0 && syclcompat::global_id::z() ==0 )
242242
//printf("syclcompat::global_id::x() = %d, syclcompat::global_id::y() = %d, syclcompat::global_id::z() = %d, thread_idx = %d, i = %d, in[i].ptr_ = %x, in[i].idx_=%x, value_bit = %x, value = %d, quant_map[value] = %f, out[i] = %f\n",syclcompat::global_id::x(), syclcompat::global_id::y(), syclcompat::global_id::z(), thread_idx, i, in[i].ptr_, in[i].idx_, value, static_cast<int>(value), quant_map[value], static_cast<float>(out[i]));
243243
}
244-
int scale_number = decltype(size(tCrS_input))::value;
245-
for(int i=0; i<scale_number; i++){
246-
auto s_value = tCrS_input(i);
247-
if(cute::thread0()) printf("scale_number = %d, tCrS_input[%d] = %f\n",scale_number, i, s_value);
248-
}
244+
}
245+
int scale_number = decltype(size(tCrS_input))::value;
246+
for(int i=0; i<scale_number; i++){
247+
auto s_value = tCrS_input(i);
248+
if(cute::thread0()) printf("scale_number = %d, tCrS_input[%d] = %f\n",scale_number, i, s_value);
249249
}
250250
#else
251251
static constexpr auto N = decltype(size<1>(in))::value;
@@ -382,12 +382,14 @@ class kgemm_4bit_inference_cutlass_dequant {
382382
Tensor mma_B = make_tensor<ElementMMA>(make_fragment_layout(tiled_copy_b, tCgB(_,_,_,0).shape()));
383383

384384
Tensor dequant_frag = make_tensor<ElementB>(mma_B.layout());
385-
385+
386+
const int SubgroupSize = 16;
387+
const int SG_QNT_WIDTH = 32;
386388
static constexpr auto scale_traits_size = decltype(size(typename GmemTiledCopyScale::BlockShape{}))::value / SubgroupSize;
387-
static constexpr auto scale_traits_num = SG_QNT_WIDTH / size<1>(typename GmemTiledCopyScale::BlockShape{});
389+
static constexpr auto scale_traits_num = SG_QNT_WIDTH / decltype(size<1>(typename GmemTiledCopyScale::BlockShape{}))::value;
388390
using FragScaleLayout = Layout<Shape<Int<scale_traits_size>, Int<scale_traits_num>, _1>>;
389391
Tensor fragment_scale = make_tensor<ElementScale>(FragScaleLayout{});
390-
if(cute::thread0()) printf("scale_traits_size = %d, scale_traits_num = %d, SG_QNT_WIDTH = %d\n", scale_traits_size, scale_traits_num, SG_QNT_WIDTH);
392+
if(cute::thread0()) printf("scale_traits_size = %d, scale_traits_num = %d, SG_QNT_WIDTH = %d, BlockShape = %d, BlockShape_1= %d\n", scale_traits_size, scale_traits_num, SG_QNT_WIDTH, decltype(size(typename GmemTiledCopyScale::BlockShape{}))::value, decltype(size<1>(typename GmemTiledCopyScale::BlockShape{}))::value);
391393

392394
static_assert(std::is_same_v<typename decltype(dequant_frag)::value_type, ElementQuant>);
393395
static_assert(std::is_same_v<typename decltype(mma_A)::value_type, ElementMMA>);
@@ -403,6 +405,7 @@ class kgemm_4bit_inference_cutlass_dequant {
403405
Tensor tBgB = thr_copy_B.retile_S(tCgB);
404406

405407
//// Prepare for prefetch
408+
const int BLK_K = 64;
406409
auto tiled_prefetch_a = cute::prefetch_selector<Shape<Int<BLK_M>,Int<BLK_K>>, Num_SGs>(tiled_copy_a);;
407410
auto tiled_prefetch_b = cute::prefetch_selector<Shape<Int<BLK_N>,Int<BLK_K>>, Num_SGs>(tiled_copy_b);;
408411
auto thr_prefetch_A = tiled_prefetch_a.get_slice(thread_idx);
@@ -413,22 +416,25 @@ class kgemm_4bit_inference_cutlass_dequant {
413416
auto pBgB = thr_prefetch_B.partition_S(gB);
414417

415418
// Run mainloop
419+
const int BLK_N = 64;
420+
const int ATOM_N = 2;
421+
const int SG_N = 32;
416422
auto [m_idx, n_idx, k_idx, l_idx] = blk_coord_mnkl;
417423
const int n_coord_s = n_idx * BLK_N + (get_sub_group_id() % ATOM_N) * SG_N;
418424
const int l_coord_s = l_idx;
419425

420-
if(cute::thread0()) printf("m_idx = %d, n_idx = %d, k_idx = %d, l_idx = %d, n_coord_s = %d, l_coord_s = %d\n",m_idx, n_idx, k_idx, l_idx, n_coord_s, l_coord_s);
426+
if(cute::thread0()) printf("get_sub_group_id() = %d, m_idx = %d, n_idx = %d, k_idx = %d, l_idx = %d, n_coord_s = %d, l_coord_s = %d\n",get_sub_group_id(), m_idx, n_idx, k_idx, l_idx, n_coord_s, l_coord_s);
421427

422428
auto copy_iter_s = [&](){
423429
return make_tensor(make_inttuple_iter(make_coord(n_coord_s, 0, l_coord_s)),
424430
make_layout(make_shape(Int<scale_traits_size>{}, Int<scale_traits_num>{}, _1{}, k_tile_count),
425-
make_stride(E<0>{} * _16{}, E<0>{} * size<1>(typename GmemTiledCopyScale::BlockShape{}), _0{}, E<1>{} * _1{})));
431+
make_stride(E<0>{} * _16{}, E<0>{} * decltype(size<1>(typename GmemTiledCopyScale::BlockShape{}))::value, _0{}, E<1>{} * _1{})));
426432

427433
}();
428434
#if 1
429435
#define PRINT(x) print(#x ": "); print(x); print("\n");
430436
if (cutlass::thread(LOG_THREAD, LOG_GROUP)) {
431-
print("======================= A: \n");
437+
print("\n\n======================= A: \n");
432438
print(" gA : "); print(gA); print("\n");
433439
print(" tCgA : "); print(tCgA); print("\n");
434440
print(" tAgA : "); print(tAgA); print("\n");
@@ -457,7 +463,7 @@ class kgemm_4bit_inference_cutlass_dequant {
457463
print(" tiled_prefetch_a : "); print(tiled_prefetch_a); print("\n");
458464
print(" tiled_prefetch_b : "); print(tiled_prefetch_b); print("\n");
459465
print(" pAgA : "); print(pAgA); print("\n");
460-
print(" pBgB : "); print(pBgB); print("\n");
466+
print(" pBgB : "); print(pBgB); print("\n\n\n");
461467
}
462468
#undef PRINT
463469
#endif
@@ -477,7 +483,7 @@ class kgemm_4bit_inference_cutlass_dequant {
477483
copy(tiled_copy_a, tAgA(_,_,_,k_tile), frag_copy_A);
478484
copy(tiled_copy_b, tBgB(_,_,_,k_tile), frag_copy_B);
479485

480-
const int k_reload_factor = params.group_size / BLK_K;
486+
const int k_reload_factor = ceil_div(params.group_size, BLK_K);
481487

482488
if(cute::thread0()) printf("params.group_size = %d, BLK_K = %d, k_reload_factor = %d\n",params.group_size, BLK_K, k_reload_factor);
483489

@@ -591,14 +597,6 @@ void gemm_4bit_inference_cutlass_dequant(int m, int n, int k, T *A, unsigned cha
591597
auto mB_nkl = make_tensor(cute::subbyte_iterator<ElementB>(B), make_layout(make_shape(n, k, l), stride_B));
592598
Copy_B tiled_copy_b{Copy_B{}.with(mB_nkl)};
593599

594-
#define PRINT(x) print(#x ": "); print(x); print("\n");
595-
if (cutlass::thread(LOG_THREAD, LOG_GROUP)) {
596-
print("===================== B :\n");
597-
print(" stride_B : "); print(stride_B); print("\n");
598-
print("===================== B :\n");
599-
}
600-
#undef PRINT
601-
602600
params.tiled_copy_a = tiled_copy_a;
603601
params.tiled_copy_b = tiled_copy_b;
604602

@@ -607,12 +605,21 @@ void gemm_4bit_inference_cutlass_dequant(int m, int n, int k, T *A, unsigned cha
607605
std::cout<<"n = "<<n<<" k = "<<k<<" blocksize = "<<blocksize<<" scale_k = "<<scale_k<<std::endl;
608606

609607
auto mScale = make_tensor(
610-
make_gmem_ptr(reinterpret_cast<ElementScale *>(absmax_)),
608+
make_gmem_ptr(absmax_),
611609
make_layout(make_shape(n, scale_k, l), stride_S));
612610
Copy_Scale tiled_copy_scale{Copy_Scale{}.with(mScale)};
613611

614612
params.tiled_copy_scale = tiled_copy_scale;
615613

614+
#define PRINT(x) print(#x ": "); print(x); print("\n");
615+
if (cutlass::thread(LOG_THREAD, LOG_GROUP)) {
616+
print("===================== B :\n");
617+
print(" stride_B : "); print(stride_B); print("\n");
618+
print(" stride_S : "); print(stride_S); print("\n");
619+
print("===================== B :\n");
620+
}
621+
#undef PRINT
622+
616623
cutlass::KernelHardwareInfo hw_info;
617624
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
618625
auto problem_shape_MNKL = append<4>(problem_size, 1);

0 commit comments

Comments
 (0)