softmax with EVT #177
softmax with EVT #177
Conversation
jiyang1011
requested review from
aacostadiaz,
muhammad-tanvir-1211,
taozha2 and
t4c1
jiyang1011
force-pushed
the
jiyang/softmax
branch
3 times, most recently
from
a435631 to
ad00c03
Compare
jiyang1011
force-pushed
the
jiyang/softmax
branch
from
ad00c03 to
bf7c5e0
Compare
jiyang1011
force-pushed
the
jiyang/softmax
branch
3 times, most recently
from
0651b94 to
7fdf576
Compare
jiyang1011
force-pushed
the
jiyang/softmax
branch
from
7fdf576 to
a84227b
Compare
jiyang1011
force-pushed
the
jiyang/softmax
branch
from
f9d5154 to
783d9ab
Compare
| @@ -312,7 +313,7 @@ class CollectiveEpilogue< | |||
| bool is_C_load_needed = is_source_supported && fusion_callbacks.is_C_load_needed(); | |||
|
|
|||
| Tensor trC = make_tensor<typename TiledMma::ValTypeC>(Shape<Int<FragmentSize>>{}); | |||
| Tensor trD = make_tensor<typename TiledMma::ValTypeD>(Shape<Int<FragmentSize>>{}); | |||
| Tensor trD = make_tensor<typename TiledMma::ValTypeD>(Shape<Int<FragmentSize>, Int<FragsM>, Int<FragsN>>{}); | |||
There was a problem hiding this comment.
If a callback needs to store all the values for Int and Int, that should be done inside the callback instead of in the generic path.
| } | ||
| copy(params.xe_store_d, trD, rw_coord(_, epi_m, epi_n)); | ||
| cst_callbacks.reduce(nullptr, synchronize, epi_m, epi_n, (epi_m == FragsM - 1 && epi_n == FragsN - 1), trD(_, epi_m, epi_n)); |
| CUTLASS_PRAGMA_UNROLL | ||
| for (int epi_n = 0; epi_n < FragsN; epi_n++) { | ||
| CUTLASS_PRAGMA_UNROLL | ||
| for (int epi_m = 0; epi_m < FragsM; epi_m++) { | ||
| copy(params.xe_store_d, trD(_, epi_m, epi_n), rw_coord(_, epi_m, epi_n)); | ||
| } | ||
| } |
There was a problem hiding this comment.
This change will group all the store operations to run at the end.
In the SM90 implementation, there’s a condition around the copy operation to check if the D matrix is available. We can apply the same logic here.
If D isn't available, the softmax output should be stored in a different matrix, which will be added directly to the callback. This way, the generic path won’t know about that matrix. Since the reduce method is the only one that knows when the data is ready to be copied, I suggest we handle the copy operations within that method.
There was a problem hiding this comment.
I've added the if condition to skip the copy in PR#199. Now, you should be able to revert the changes in the generic path for the epilogue and implement all the softmax functionalities inside the callback.
| float cute_time = timer.seconds() / options.iterations; | ||
| double tflops = (2.0 * options.m * options.n * options.k * options.l) * 1e-12; | ||
| std::cout << "Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << 'x' << options.l << std::endl; | ||
| printf("Cutlass GEMM Performance: [%4.3f]GB/s (%6.4f)ms\n", io / cute_time, cute_time*1000); |
There was a problem hiding this comment.
| printf("Cutlass GEMM Performance: [%4.3f]GB/s (%6.4f)ms\n", io / cute_time, cute_time*1000); | |
| printf("Cutlass GEMM Performance: %4.3f GB/s , %4.3f TF/s , %6.4f ms\n", io / cute_time, tflops/cute_time, cute_time*1000); |
| SYCL_DEVICE_OCL(float sub_group_reduce_add(float i)); | ||
| SYCL_DEVICE_OCL(float sub_group_reduce_max(float i)); |
There was a problem hiding this comment.
the reduce operation is not part of mma
No description provided.