Skip to content

Commit 492705a

Browse files
xiaolil1xiaolil1
committed
effective change?
1 parent cb27ec8 commit 492705a

3 files changed

Lines changed: 33 additions & 28 deletions

File tree

‎bitsandbytes/backends/xpu/ops.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,8 @@ def _(
185185
blocksize: int,
186186
) -> torch.Tensor:
187187
shape = (*A.shape[:-1], shapeB[0])
188+
import pdb
189+
pdb.set_trace()
188190
out = torch.zeros(shape, device=A.device, dtype=torch.float32)
189191
_gemv_4bit_impl(A, B, shapeB, absmax.bfloat16(), code, blocksize, out=out)
190192
return out

‎csrc/xpu_cutlass_fusion.cpp‎

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

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

7171
//Threads number
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
72+
static constexpr auto ATOM_M = get<1>(typename TiledMma::ThrLayoutVMNK{}.shape()); //8 //1
73+
static constexpr auto ATOM_N = get<2>(typename TiledMma::ThrLayoutVMNK{}.shape()); //4 //2
74+
static constexpr auto ATOM_K = get<3>(typename TiledMma::ThrLayoutVMNK{}.shape()); //1 //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); //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
81+
static constexpr auto SG_M = ceil_div(BLK_M, ATOM_M); //32 //16
82+
static constexpr auto SG_N = ceil_div(BLK_N, ATOM_N); //64 //32
83+
static constexpr auto SG_K = ceil_div(BLK_K, ATOM_K); //32 //64
8484
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; //2
88-
static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{}); //1*2*1*16=32
87+
static constexpr auto Num_SGs = ATOM_N * ATOM_M * ATOM_K; //32 //2
88+
static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{}); //8*4*1*16=512 //1*2*1*16=32
8989

9090
// Define Mainloop dispatch policy
91-
constexpr int PipelineStages = 3;
91+
constexpr int PipelineStages = 2;
9292
using DispatchPolicy = cutlass::gemm::MainloopIntelPVCMixedPrecision<PipelineStages>;
93-
static constexpr int SubgroupSize = DispatchPolicy::SubgroupSize; // sub_group size
93+
static constexpr int SubgroupSize = DispatchPolicy::SubgroupSize; // 16
9494

9595
// Design Epilogue
9696
using EpilogueDispatchPolicy = cutlass::epilogue::IntelPVCEpilogue;
97-
using EpilogueOp = cutlass::epilogue::fusion::LinearCombination<float /*data_type of GEMM output*/, ElementComputeEpilogue, ElementAccumulator, ElementAccumulator, cutlass::FloatRoundStyle::round_to_nearest>;
97+
using EpilogueOp = cutlass::epilogue::fusion::LinearCombination<ElementAccumulator, ElementComputeEpilogue, ElementAccumulator, ElementAccumulator, cutlass::FloatRoundStyle::round_to_nearest>;
9898
using FusionCallBacks = cutlass::epilogue::fusion::FusionCallbacks<EpilogueDispatchPolicy, EpilogueOp, TileShape, decltype(tile_shape(TiledMma()))>;
9999
using SharedStorage = FusionCallBacks::SharedStorage;
100100

@@ -112,15 +112,14 @@ using CollectiveEpilogue = cutlass::epilogue::collective::CollectiveEpilogue<
112112
TileShape,
113113
ElementAccumulator,
114114
cutlass::gemm::TagToStrideC_t<cutlass::layout::RowMajor>, // Convert CUTLASS 2.x to CUTLASS 3.x representation
115-
float, // data_type of output: out
115+
ElementOutput,
116116
cutlass::gemm::TagToStrideC_t<cutlass::layout::RowMajor>, // Convert CUTLASS 2.x to CUTLASS 3.x representation
117117
FusionCallBacks,
118118
XE_2D_U32x8x16_LD_N, // The copy atom used to load matrix C
119119
void, void,
120120
XE_2D_U32x8x16_ST_N, // The copy atom used to store matrix D
121121
void, void>;
122122
using EpilogueParams = typename CollectiveEpilogue::Params;
123-
using EpilogueArguments = typename CollectiveEpilogue::Arguments;
124123

125124
using ClusterShape = typename DispatchPolicy::ClusterShape;
126125

@@ -287,7 +286,7 @@ class kgemm_4bit_inference_cutlass_dequant {
287286
for (int i = 0; i < vec_size; i++) {
288287
uint8_t value = (format_data >> (src_bits * i)) & 0xf;
289288
dst[i] = static_cast<DstType>(quant_map[value] * static_cast<float>(ts));
290-
if(cute::thread0()) printf("n = %d, s = %d, i = %d, src = %d, quant_map[value] = %f, ts = %f, dst = %f\n", n, s, i, static_cast<int>(value), quant_map[value], static_cast<float>(ts), static_cast<float>(dst[i]));
289+
//if(cute::thread0()) printf("n = %d, s = %d, i = %d, src = %d, quant_map[value] = %f, ts = %f, dst = %f\n", n, s, i, static_cast<int>(value), quant_map[value], static_cast<float>(ts), static_cast<float>(dst[i]));
291290
}
292291
}
293292
}
@@ -665,23 +664,26 @@ std::cout << std::endl;
665664

666665
params.tiled_copy_scale = tiled_copy_scale;
667666

667+
cutlass::KernelHardwareInfo hw_info;
668+
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
669+
auto problem_shape_MNKL = problem_size; //append<4>(problem_size, 1);
670+
float alpha=1.0f;
671+
float beta=0.f;
672+
StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, l));
673+
StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(m, n, l));
674+
668675
#define PRINT(x) print(#x ": "); print(x); print("\n");
669676
if (cutlass::thread(LOG_THREAD, LOG_GROUP)) {
670-
print("===================== B :\n");
677+
print("===================== stride :\n");
678+
print(" stride_A : "); print(stride_A); print("\n");
671679
print(" stride_B : "); print(stride_B); print("\n");
680+
print(" stride_C : "); print(stride_C); print("\n");
681+
print(" stride_D : "); print(stride_D); print("\n");
672682
print(" stride_S : "); print(stride_S); print("\n");
673-
print("===================== B :\n");
683+
print("===================== stride :\n");
674684
}
675685
#undef PRINT
676686

677-
cutlass::KernelHardwareInfo hw_info;
678-
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
679-
auto problem_shape_MNKL = append<4>(problem_size, 1);
680-
float alpha=1.0;
681-
float beta=0.f;
682-
StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, l));
683-
StrideC stride_D = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, l));
684-
685687
params.hw_info = hw_info;
686688
params.epilogue = CollectiveEpilogue::to_underlying_arguments(problem_size, {{alpha, beta}, nullptr, stride_C, out, stride_D}, nullptr);
687689

‎tests/test_xpu.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,8 @@ def test_gemm_4bit(self, device, dim, dtype, storage_type, quant_storage, double
6767
#pdb.set_trace()
6868
if kind == "fc1":
6969
A = torch.randn(32, dim, dtype=dtype, device=device)
70-
B = torch.randn(dim, dim, dtype=dtype, device=device) / math.sqrt(dim)
70+
#A = torch.arange(1, 32 * 256 + 1).reshape(32, 256).bfloat16().xpu()
71+
B = torch.ones(dim, dim, dtype=dtype, device=device) / math.sqrt(dim)
7172
elif kind == "fc2":
7273
A = torch.randn(1, 4 * dim, dtype=dtype, device=device)
7374
B = torch.randn(dim, 4 * dim, dtype=dtype, device=device) / math.sqrt(dim)

0 commit comments

Comments
 (0)