update 3.8 v2 (#2112)

* update 3.8 v2

* update 3.8

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
This commit is contained in:
Yujia Zhai
2025-02-19 19:03:14 -08:00
committed by GitHub
parent e9627ce55b
commit b84e9802d8
166 changed files with 3986 additions and 4037 deletions

View File

@@ -285,7 +285,9 @@ Layout, and Dispatch Policy combinations for each row of [Table 1](#legacy_gemm_
| 1/2 SM | Epilogue Dispatch Policy |
|--------|------------------------------------------|
| 1SM | cutlass::epilogue::TmaWarpSpecialized1Sm |
| 1SM | cutlass::epilogue::NoSmemWarpSpecialized1Sm |
| 2SM | cutlass::epilogue::TmaWarpSpecialized2Sm |
| 2SM | cutlass::epilogue::NoSmemWarpSpecialized2Sm |
**Table 15: Epilogue PerSmTileShape_MNK** <a id="epi_persmtileshape" name="epi_persmtileshape"></a>
| 1/2 SM | MMA tile Shape | PerSmTileShape_MNK |
@@ -442,7 +444,7 @@ PerSmTileShape_MNK should be deduced from the mainloop setup. For example, in ab
It means each CTA is doing (256 / 2sm) x 256 x 128 output, so the PerSmTileShape_MNK is 128x256x128. The possible PerSmTileShape_MNK
is listed in [Table 15](#epi_persmtileshape)
The epilogue scheduling policy is configurable, and it is common to set `cutlass::epilogue::TmaWarpSpecialized2Sm`
The epilogue scheduling policy is configurable, and it is common to set `cutlass::epilogue::collective::EpilogueScheduleAuto`
to allow the epilogue builder to automatically select the appropriate policy. However, it can also be explicitly defined to
use other policies based on the 1sm or 2sm MMA instruction. The available policies are listed in [Table 14](#epi_dispatch).
@@ -458,10 +460,6 @@ use other policies based on the 1sm or 2sm MMA instruction. The available polici
using ElementAccumulator = float;
// Epilogue computation's precision type
using ElementCompute = float;
// Cluster size for multicast
using ClusterShape_MNK = Shape<_4,_4,_1>;
// Collective Epilogue takes the output tile shape for 1 CTA
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
//
// Construct CollectiveEpilogue
@@ -469,7 +467,7 @@ use other policies based on the 1sm or 2sm MMA instruction. The available polici
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp, // Arch and Tensorop spec
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
MmaTileShape_MNK, ClusterShape_MNK, // MMA tile shape, and cluster shape
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
ElementC, GmemLayoutC, AlignC, // C tensor description
@@ -499,12 +497,12 @@ Typically, GmemLayoutSFD would be same as the GmemLayoutD.
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp, // Arch and Tensorop spec
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
MmaTileShape_MNK, ClusterShape_MNK, // MMA tile shape, and cluster shape
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
ElementC, GmemLayoutC, AlignC, // C tensor description
ElementD, GmemLayoutD, AlignD, // D tensor description
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
FusionOperation // <================================== Pass the fusion config into epilogue builder.
>::CollectiveOp;
```