A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Perf] Use upstream CUTLASS for SM90 Block FP8 kernel (#23280)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Michael Goin committed
Sep 11, 2025 at 18:43 UTC
c3aea10dc8b83645523843a27999a95340020ddf
13 files changed
+222
-1261
benchmarks/kernels/bench_block_fp8_gemm.py
+44
-13
@@ -4,7 +4,10 @@
4
import torch
5
6
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
7
- w8a8_block_fp8_matmul,
7
+ apply_w8a8_block_fp8_linear,
8
+)
9
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
10
+ CUTLASS_BLOCK_FP8_SUPPORTED,
11
)
12
from vllm.platforms import current_platform
13
from vllm.triton_utils import triton as vllm_triton
@@ -29,7 +32,7 @@ DEEPSEEK_V3_SHAPES = [
32
]
33
34
32
-def build_w8a8_block_fp8_runner(M, N, K, block_size, device):
35
+def build_w8a8_block_fp8_runner(M, N, K, block_size, device, use_cutlass):
36
"""Build runner function for w8a8 block fp8 matmul."""
37
factor_for_scale = 1e-2
38
@@ -37,37 +40,54 @@ def build_w8a8_block_fp8_runner(M, N, K, block_size, device):
40
fp8_max, fp8_min = fp8_info.max, fp8_info.min
41
42
# Create random FP8 tensors
40
- A_fp32 = (torch.rand(M, K, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max
41
- A = A_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
43
+ A_ref = (torch.rand(M, K, dtype=torch.bfloat16, device=device) - 0.5) * 2 * fp8_max
44
43
- B_fp32 = (torch.rand(N, K, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max
44
- B = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
45
+ B_ref = (torch.rand(N, K, dtype=torch.bfloat16, device=device) - 0.5) * 2 * fp8_max
46
+ B = B_ref.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
47
48
# Create scales
49
block_n, block_k = block_size[0], block_size[1]
50
n_tiles = (N + block_n - 1) // block_n
51
k_tiles = (K + block_k - 1) // block_k
52
51
- As = torch.rand(M, k_tiles, dtype=torch.float32, device=device) * factor_for_scale
53
Bs = (
54
torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device)
55
* factor_for_scale
56
)
57
58
+ # SM90 CUTLASS requires row-major format for scales
59
+ if use_cutlass and current_platform.is_device_capability(90):
60
+ Bs = Bs.T.contiguous()
61
+
62
def run():
58
- return w8a8_block_fp8_matmul(A, B, As, Bs, block_size, torch.bfloat16)
63
+ if use_cutlass:
64
+ return apply_w8a8_block_fp8_linear(
65
+ A_ref, B, block_size, Bs, cutlass_block_fp8_supported=True
66
+ )
67
+ else:
68
+ return apply_w8a8_block_fp8_linear(
69
+ A_ref, B, block_size, Bs, cutlass_block_fp8_supported=False
70
+ )
71
72
return run
73
74
75
+# Determine available providers
76
+available_providers = ["torch-bf16", "w8a8-block-fp8-triton"]
77
+plot_title = "BF16 vs W8A8 Block FP8 GEMMs"
78
+
79
+if CUTLASS_BLOCK_FP8_SUPPORTED:
80
+ available_providers.append("w8a8-block-fp8-cutlass")
81
+
82
+
83
@vllm_triton.testing.perf_report(
84
vllm_triton.testing.Benchmark(
85
x_names=["batch_size"],
86
x_vals=[1, 16, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384],
87
x_log=False,
88
line_arg="provider",
69
- line_vals=["torch-bf16", "w8a8-block-fp8"],
70
- line_names=["torch-bf16", "w8a8-block-fp8"],
89
+ line_vals=available_providers,
90
+ line_names=available_providers,
91
ylabel="TFLOP/s (larger is better)",
92
plot_name="BF16 vs W8A8 Block FP8 GEMMs",
93
args={},
@@ -85,11 +105,22 @@ def benchmark_tflops(batch_size, provider, N, K, block_size=(128, 128)):
105
ms, min_ms, max_ms = vllm_triton.testing.do_bench_cudagraph(
106
lambda: torch.nn.functional.linear(a, b), quantiles=quantiles
107
)
88
- else: # w8a8-block-fp8
89
- run_w8a8 = build_w8a8_block_fp8_runner(M, N, K, block_size, device)
108
+ elif provider == "w8a8-block-fp8-triton":
109
+ run_w8a8_triton = build_w8a8_block_fp8_runner(
110
+ M, N, K, block_size, device, use_cutlass=False
111
+ )
112
+ ms, min_ms, max_ms = vllm_triton.testing.do_bench_cudagraph(
113
+ lambda: run_w8a8_triton(), quantiles=quantiles
114
+ )
115
+ elif provider == "w8a8-block-fp8-cutlass":
116
+ run_w8a8_cutlass = build_w8a8_block_fp8_runner(
117
+ M, N, K, block_size, device, use_cutlass=True
118
+ )
119
ms, min_ms, max_ms = vllm_triton.testing.do_bench_cudagraph(
91
- lambda: run_w8a8(), quantiles=quantiles
120
+ lambda: run_w8a8_cutlass(), quantiles=quantiles
121
)
122
+ else:
123
+ raise ValueError(f"Unknown provider: {provider}")
124
125
to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3)
126
return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms)
csrc/cutlass_extensions/gemm/collective/collective_builder.hpp
deleted
-123
@@ -1,123 +0,0 @@
1
-// Modified from: cutlass/gemm/collective/builders/sm90_gmma_builder.inl
2
-// clang-format off
3
-#pragma once
4
-
5
-#include "cutlass/gemm/collective/builders/sm90_gmma_builder.inl"
6
-
7
-#include "cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp"
8
-
9
-
10
-/////////////////////////////////////////////////////////////////////////////////////////////////
11
-
12
-namespace cutlass::gemm::collective {
13
-
14
-/////////////////////////////////////////////////////////////////////////////////////////////////
15
-
16
-// GMMA_TMA_WS_SS (BlockScaled Builders)
17
-template <
18
- class ElementA,
19
- class GmemLayoutATag,
20
- int AlignmentA,
21
- class ElementB,
22
- class GmemLayoutBTag,
23
- int AlignmentB,
24
- class ElementAccumulator,
25
- class TileShape_MNK,
26
- class ClusterShape_MNK,
27
- class StageCountType,
28
- int ScaleGranularityM
29
->
30
-struct CollectiveBuilder<
31
- arch::Sm90,
32
- arch::OpClassTensorOp,
33
- ElementA,
34
- GmemLayoutATag,
35
- AlignmentA,
36
- ElementB,
37
- GmemLayoutBTag,
38
- AlignmentB,
39
- ElementAccumulator,
40
- TileShape_MNK,
41
- ClusterShape_MNK,
42
- StageCountType,
43
- KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<ScaleGranularityM>,
44
- cute::enable_if_t<
45
- not detail::is_use_rmem_A<ElementA, GmemLayoutATag, ElementB, GmemLayoutBTag>()>
46
-> {
47
- using KernelScheduleType = KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<ScaleGranularityM>;
48
-
49
- static_assert(is_static<TileShape_MNK>::value);
50
- static_assert(is_static<ClusterShape_MNK>::value);
51
-#ifndef CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED
52
- static_assert(cutlass::detail::dependent_false<ElementA>, "Unsupported Toolkit for SM90 Collective Builder\n");
53
-#endif
54
- static_assert(detail::is_aligned<ElementA, AlignmentA, ElementB, AlignmentB, detail::tma_alignment_bytes>(),
55
- "Should meet TMA alignment requirement\n");
56
-
57
- static constexpr bool IsArrayOfPointersGemm = (cute::is_any_of_v<KernelScheduleType,
58
- KernelPtrArrayTmaWarpSpecializedCooperative,
59
- KernelPtrArrayTmaWarpSpecializedPingpong>);
60
- static constexpr bool IsFP8Input = detail::is_input_fp8<ElementA, ElementB>();
61
- static_assert((!IsFP8Input || !IsArrayOfPointersGemm),
62
- "KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum is only compatible with FP8 Blocked Scaled version right now.");
63
-
64
- // For fp32 types, map to tf32 MMA value type
65
- using ElementAMma = cute::conditional_t<cute::is_same_v<ElementA, float>, tfloat32_t, ElementA>;
66
- using ElementBMma = cute::conditional_t<cute::is_same_v<ElementB, float>, tfloat32_t, ElementB>;
67
-
68
- static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_ss_tag_to_major_A<ElementAMma, GmemLayoutATag>();
69
- static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_ss_tag_to_major_B<ElementBMma, GmemLayoutBTag>();
70
-
71
- static constexpr bool IsCooperative = cute::is_any_of_v<KernelScheduleType,
72
- KernelTmaWarpSpecializedCooperative,
73
- KernelPtrArrayTmaWarpSpecializedCooperative,
74
- KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<ScaleGranularityM>>;
75
- using AtomLayoutMNK = cute::conditional_t<IsCooperative,
76
- Layout<Shape<_2,_1,_1>>, Layout<Shape<_1,_1,_1>>>;
77
-
78
- using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::ss_op_selector<
79
- ElementAMma, ElementBMma, ElementAccumulator, TileShape_MNK, GmmaMajorA, GmmaMajorB>(), AtomLayoutMNK{}));
80
-
81
- using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<1>(ClusterShape_MNK{})));
82
- using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape_MNK{})));
83
-
84
- using SmemLayoutAtomA = decltype(detail::ss_smem_selector<
85
- GmmaMajorA, ElementAMma, decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
86
- using SmemLayoutAtomB = decltype(detail::ss_smem_selector<
87
- GmmaMajorB, ElementBMma, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
88
-
89
- static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0;
90
- static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage);
91
-
92
- static constexpr int PipelineStages = detail::compute_stage_count_or_override<detail::sm90_smem_capacity_bytes - KernelSmemCarveout,
93
- ElementAMma, ElementBMma, TileShape_MNK>(StageCountType{});
94
- using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8<PipelineStages, ClusterShape_MNK, KernelScheduleType, ScaleGranularityM>;
95
-
96
- using SmemCopyAtomA = void;
97
- using SmemCopyAtomB = void;
98
-
99
- using CollectiveOp = CollectiveMma<
100
- DispatchPolicy,
101
- TileShape_MNK,
102
- ElementA,
103
- TagToStrideA_t<GmemLayoutATag>,
104
- ElementB,
105
- TagToStrideB_t<GmemLayoutBTag>,
106
- TiledMma,
107
- GmemTiledCopyA,
108
- SmemLayoutAtomA,
109
- SmemCopyAtomA,
110
- cute::identity,
111
- GmemTiledCopyB,
112
- SmemLayoutAtomB,
113
- SmemCopyAtomB,
114
- cute::identity
115
- >;
116
-};
117
-
118
-
119
-/////////////////////////////////////////////////////////////////////////////////////////////////
120
-
121
-} // namespace cutlass::gemm::collective
122
-
123
-/////////////////////////////////////////////////////////////////////////////////////////////////
csrc/cutlass_extensions/gemm/collective/fp8_accumulation.hpp
deleted
-183
@@ -1,183 +0,0 @@
1
-// clang-format off
2
-// adapted from: https://github.com/soundOfDestiny/cutlass/blob/a4208aa6958864923505cade9c63eb2a6daf16e5/include/cutlass/gemm/collective/fp8_accumulation.hpp
3
-
4
-/***************************************************************************************************
5
- * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
6
- * SPDX-License-Identifier: BSD-3-Clause
7
- *
8
- * Redistribution and use in source and binary forms, with or without
9
- * modification, are permitted provided that the following conditions are met:
10
- *
11
- * 1. Redistributions of source code must retain the above copyright notice, this
12
- * list of conditions and the following disclaimer.
13
- *
14
- * 2. Redistributions in binary form must reproduce the above copyright notice,
15
- * this list of conditions and the following disclaimer in the documentation
16
- * and/or other materials provided with the distribution.
17
- *
18
- * 3. Neither the name of the copyright holder nor the names of its
19
- * contributors may be used to endorse or promote products derived from
20
- * this software without specific prior written permission.
21
- *
22
- * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
23
- * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
24
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
25
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
26
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
27
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
28
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
29
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
30
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
31
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
32
- *
33
- **************************************************************************************************/
34
-
35
-#pragma once
36
-
37
-#include "cute/algorithm/clear.hpp"
38
-#include "cute/tensor.hpp"
39
-
40
-//////////////////////////////////////////////////////////////////////////////
41
-///////////////////////////////////FP8 Accumulation///////////////////////////
42
-//////////////////////////////////////////////////////////////////////////////
43
-/// This class provides API to promote (add) or scale (multiply_add) the results
44
-/// from the tensor core accumulators to the main accumulators when the number
45
-/// of MMAs reaches the max number of MMA interval specified by user, after that
46
-/// the tensor core accumulators are zeroed.
47
-//////////////////////////////////////////////////////////////////////////////
48
-
49
-namespace cutlass::gemm::collective {
50
-
51
-template <
52
- class EngineAccum,
53
- class LayoutAccum>
54
-struct GmmaFP8AccumulationWithScale {
55
- using TensorAccum = cute::Tensor<EngineAccum, LayoutAccum>;
56
- using ElementAccumulator = typename EngineAccum::value_type;
57
-
58
- static_assert(is_static<LayoutAccum>::value, "Accumulator Layout should be static");
59
- static_assert(is_rmem<TensorAccum>::value , "Accumulator tensor must be rmem resident.");
60
-
61
-private:
62
- TensorAccum& accum_;
63
- TensorAccum accum_temp_;
64
-
65
- uint32_t accum_promotion_interval_; // defines the max num of executed MMAs after which accum should be promoted.
66
- uint32_t mma_count_per_mainloop_iteration_; // num of MMAs per k_tile of mainloop
67
- uint32_t mma_count_; // current executed MMAs
68
- uint32_t reset_accum_flag_; // accum needs to be zeroed or not.
69
-
70
- // promote or `add` the partial accumulators to main accumulator (FADD).
71
- CUTLASS_DEVICE
72
- void promote_core() {
73
- warpgroup_wait<0>();
74
- CUTLASS_PRAGMA_UNROLL
75
- for (int i = 0; i < size(accum_); ++i) {
76
- accum_(i) += accum_temp_(i);
77
- }
78
- }
79
-
80
- // `multiply` scale the partial accumulators and `add` to main accumulator (FFMA).
81
- template <
82
- class EngineScale,
83
- class LayoutScale>
84
- CUTLASS_DEVICE
85
- void scale_core(const cute::Tensor<EngineScale, LayoutScale> &scale) {
86
- using TensorScale = cute::Tensor<EngineScale, LayoutScale>;
87
-
88
- static_assert(is_static<LayoutScale>::value, "Scale Layout should be static");
89
- static_assert(is_rmem<TensorScale>::value , "Scale tensor must be rmem resident.");
90
-
91
- static_assert(LayoutAccum{}.shape() == LayoutScale{}.shape(), "Accumulator and scale must have same shape.");
92
-
93
- warpgroup_wait<0>();
94
- CUTLASS_PRAGMA_UNROLL
95
- for (int i = 0; i < size(accum_); ++i) {
96
- accum_(i) += accum_temp_(i) * scale(i);
97
- }
98
- }
99
-
100
-public:
101
- CUTLASS_DEVICE
102
- GmmaFP8AccumulationWithScale(
103
- TensorAccum &accum,
104
- uint32_t accum_promotion_interval,
105
- uint32_t mma_count_per_mainloop_iteration)
106
- : accum_(accum),
107
- accum_promotion_interval_(accum_promotion_interval),
108
- mma_count_per_mainloop_iteration_(mma_count_per_mainloop_iteration),
109
- mma_count_(0),
110
- reset_accum_flag_(0)
111
- {
112
- accum_temp_ = cute::make_fragment_like(accum);
113
- }
114
-
115
- //
116
- // Methods (Common)
117
- //
118
-
119
- CUTLASS_DEVICE
120
- TensorAccum& operator()() {
121
- return accum_temp_;
122
- }
123
-
124
- /// prepare the MMA accumulators when initialization or zeroing is required.
125
- CUTLASS_DEVICE
126
- bool prepare_if_needed() {
127
- return reset_accum_flag_;
128
- }
129
-
130
- //
131
- // Methods (for FADD version)
132
- //
133
-
134
- /// promote (add) the results from the MMA accumulators to main accumulator if needed.
135
- CUTLASS_DEVICE
136
- void promote_if_needed() {
137
- mma_count_ += mma_count_per_mainloop_iteration_;
138
- reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
139
- if (reset_accum_flag_) {
140
- promote_core();
141
- mma_count_ = 0;
142
- }
143
- }
144
-
145
- /// promote (add) the residue results from the MMA accumulators to main accumulator if needed.
146
- CUTLASS_DEVICE
147
- void promote_residue_if_needed() {
148
- if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
149
- promote_core();
150
- }
151
- }
152
-
153
- //
154
- // Methods (for FFMA version)
155
- //
156
-
157
- /// scale (multiply_add) the results from the MMA accumulators to main accumulator if needed.
158
- template <
159
- class EngineScale,
160
- class LayoutScale>
161
- CUTLASS_DEVICE
162
- void scale_if_needed(const cute::Tensor<EngineScale, LayoutScale> &scale) {
163
- mma_count_ += mma_count_per_mainloop_iteration_;
164
- reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
165
- if (reset_accum_flag_) {
166
- scale_core(scale);
167
- mma_count_ = 0;
168
- }
169
- }
170
-
171
- /// scale (multiply_add) the residue results from the MMA accumulators to main accumulator if needed.
172
- template <
173
- class EngineScale,
174
- class LayoutScale>
175
- CUTLASS_DEVICE
176
- void scale_residue_if_needed(const cute::Tensor<EngineScale, LayoutScale> &scale) {
177
- if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
178
- scale_core(scale);
179
- }
180
- }
181
-};
182
-
183
-} // namespace cutlass::gemm::collective
csrc/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
deleted
-729
@@ -1,729 +0,0 @@
1
-// clang-format off
2
-// Adapted (Heavily) from: https://github.com/soundOfDestiny/cutlass/blob/9d997ce0dea4c5fa1a617db6b7ff29aa9235822c/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
3
-
4
-/***************************************************************************************************
5
- * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
6
- * SPDX-License-Identifier: BSD-3-Clause
7
- *
8
- * Redistribution and use in source and binary forms, with or without
9
- * modification, are permitted provided that the following conditions are met:
10
- *
11
- * 1. Redistributions of source code must retain the above copyright notice, this
12
- * list of conditions and the following disclaimer.
13
- *
14
- * 2. Redistributions in binary form must reproduce the above copyright notice,
15
- * this list of conditions and the following disclaimer in the documentation
16
- * and/or other materials provided with the distribution.
17
- *
18
- * 3. Neither the name of the copyright holder nor the names of its
19
- * contributors may be used to endorse or promote products derived from
20
- * this software without specific prior written permission.
21
- *
22
- * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
23
- * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
24
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
25
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
26
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
27
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
28
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
29
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
30
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
31
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
32
- *
33
- **************************************************************************************************/
34
-
35
-#pragma once
36
-
37
-#include "cutlass/cutlass.h"
38
-#include "cutlass/gemm/dispatch_policy.hpp"
39
-#include "cutlass/trace.h"
40
-#include "cutlass/numeric_types.h"
41
-
42
-#include "cute/arch/cluster_sm90.hpp"
43
-#include "cute/arch/copy_sm80.hpp"
44
-#include "cute/arch/copy_sm90.hpp"
45
-#include "cute/algorithm/functional.hpp"
46
-#include "cute/atom/mma_atom.hpp"
47
-#include "cute/algorithm/gemm.hpp"
48
-#include "cute/numeric/arithmetic_tuple.hpp"
49
-
50
-#include "cutlass_extensions/gemm/dispatch_policy.hpp"
51
-#include "cutlass_extensions/gemm/collective/fp8_accumulation.hpp"
52
-
53
-/////////////////////////////////////////////////////////////////////////////////////////////////
54
-
55
-namespace cutlass::gemm::collective {
56
-using namespace cute;
57
-
58
-/////////////////////////////////////////////////////////////////////////////////////////////////
59
-
60
-// WarpSpecialized Mainloop
61
-template <
62
- int Stages,
63
- class ClusterShape,
64
- class KernelSchedule,
65
- int ScaleGranularityM_,
66
- class TileShape_,
67
- class ElementA_,
68
- class StrideA_,
69
- class ElementB_,
70
- class StrideB_,
71
- class TiledMma_,
72
- class GmemTiledCopyA_,
73
- class SmemLayoutAtomA_,
74
- class SmemCopyAtomA_,
75
- class TransformA_,
76
- class GmemTiledCopyB_,
77
- class SmemLayoutAtomB_,
78
- class SmemCopyAtomB_,
79
- class TransformB_>
80
-struct CollectiveMma<
81
- MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8<Stages, ClusterShape, KernelSchedule, ScaleGranularityM_>,
82
- TileShape_,
83
- ElementA_,
84
- StrideA_,
85
- ElementB_,
86
- StrideB_,
87
- TiledMma_,
88
- GmemTiledCopyA_,
89
- SmemLayoutAtomA_,
90
- SmemCopyAtomA_,
91
- TransformA_,
92
- GmemTiledCopyB_,
93
- SmemLayoutAtomB_,
94
- SmemCopyAtomB_,
95
- TransformB_>
96
-{
97
- //
98
- // Type Aliases
99
- //
100
- using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8<Stages, ClusterShape, KernelSchedule, ScaleGranularityM_>;
101
- using TileShape = TileShape_;
102
- using ElementA = ElementA_;
103
- using StrideA = StrideA_;
104
- using ElementB = ElementB_;
105
- using StrideB = StrideB_;
106
- using TiledMma = TiledMma_;
107
- using ElementAccumulator = typename TiledMma::ValTypeC;
108
- using ElementBlockScale = ElementAccumulator;
109
- using GmemTiledCopyA = GmemTiledCopyA_;
110
- using GmemTiledCopyB = GmemTiledCopyB_;
111
- using SmemLayoutAtomA = SmemLayoutAtomA_;
112
- using SmemLayoutAtomB = SmemLayoutAtomB_;
113
- using SmemCopyAtomA = SmemCopyAtomA_;
114
- using SmemCopyAtomB = SmemCopyAtomB_;
115
- using TransformA = TransformA_;
116
- using TransformB = TransformB_;
117
- using ArchTag = typename DispatchPolicy::ArchTag;
118
-
119
- using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{}));
120
- using MainloopPipeline = cutlass::PipelineTmaAsync<DispatchPolicy::Stages>;
121
- using PipelineState = cutlass::PipelineState<DispatchPolicy::Stages>;
122
- using PipelineParams = typename MainloopPipeline::Params;
123
-
124
- // Two threads per CTA are producers (1 for operand tile and 32 for scales)
125
- static constexpr int NumProducerThreadEvents = 33;
126
-
127
- static constexpr int ScaleGranularityM = ScaleGranularityM_ == 0 ? size<0>(TileShape{}) : ScaleGranularityM_;
128
- static constexpr int ScaleMsPerTile = size<0>(TileShape{}) / ScaleGranularityM;
129
-
130
- static_assert(cute::rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
131
- static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
132
- static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
133
-
134
- static_assert(cute::rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
135
- static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
136
- static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
137
-
138
- static_assert((size<0>(TileShape{}) % ScaleGranularityM) == 0, "FP8 scaling granularity must evenly divide tile shape along M.");
139
-
140
- // Tile along modes in a way that maximizes the TMA box size.
141
- using SmemLayoutA = decltype(tile_to_shape(
142
- SmemLayoutAtomA{},
143
- make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
144
- cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
145
- using SmemLayoutB = decltype(tile_to_shape(
146
- SmemLayoutAtomB{},
147
- make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
148
- cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
149
-
150
- // Block scaling gmem-to-smem copy atom
151
- using SmemBlockScalingCopyAtomA = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
152
- using SmemBlockScalingCopyAtomB = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
153
-
154
- // Block scaling smem layout
155
- using SmemLayoutScaleA = Layout<Shape<Int<ScaleMsPerTile>, Int<DispatchPolicy::Stages>>>;
156
- using SmemLayoutScaleB = Layout<Shape<Int<DispatchPolicy::Stages>>, Stride<_1>>; // `ScaleNsPerTile` is always 1.
157
-
158
- static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more.");
159
- static_assert(cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
160
- cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
161
- "MMA atom must source both A and B operand from smem_desc for this mainloop.");
162
- static_assert(cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>,
163
- "GmemTiledCopy - invalid SM90 TMA copy atom specified.");
164
- static_assert(cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>,
165
- "GmemTiledCopy - invalid SM90 TMA copy atom specified.");
166
- static_assert(cute::is_same_v<ElementAccumulator, ElementBlockScale>,
167
- "ElementAccumulator and ElementBlockScale should be same datatype");
168
-
169
- struct SharedStorage
170
- {
171
- struct TensorStorage : cute::aligned_struct<128> {
172
- cute::array_aligned<typename TiledMma::ValTypeA, cute::cosize_v<SmemLayoutA>> smem_A; // mxk
173
- cute::array_aligned<typename TiledMma::ValTypeB, cute::cosize_v<SmemLayoutB>> smem_B; // nxk
174
- cute::array_aligned<ElementBlockScale, cute::cosize_v<SmemLayoutScaleA>> smem_scale_A; // ScaleMsPerTile x k
175
- cute::array_aligned<ElementBlockScale, cute::cosize_v<SmemLayoutScaleB>> smem_scale_B; // 1xk
176
- } tensors;
177
-
178
- using PipelineStorage = typename MainloopPipeline::SharedStorage;
179
- PipelineStorage pipeline;
180
- };
181
- using TensorStorage = typename SharedStorage::TensorStorage;
182
- using PipelineStorage = typename SharedStorage::PipelineStorage;
183
-
184
- // Host side kernel arguments
185
- struct Arguments {
186
- ElementA const* ptr_A;
187
- StrideA dA;
188
- ElementB const* ptr_B;
189
- StrideB dB;
190
- ElementBlockScale const* ptr_scale_A;
191
- ElementBlockScale const* ptr_scale_B;
192
- };
193
-
194
- // Device side kernel params
195
- struct Params {
196
- // Assumption: StrideA is congruent with Problem_MK
197
- using TMA_A = decltype(make_tma_copy_A_sm90(
198
- GmemTiledCopyA{},
199
- make_tensor(static_cast<ElementA const*>(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}),
200
- SmemLayoutA{}(_,_,0),
201
- TileShape{},
202
- ClusterShape{}));
203
- // Assumption: StrideB is congruent with Problem_NK
204
- using TMA_B = decltype(make_tma_copy_B_sm90(
205
- GmemTiledCopyB{},
206
- make_tensor(static_cast<ElementB const*>(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}),
207
- SmemLayoutB{}(_,_,0),
208
- TileShape{},
209
- ClusterShape{}));
210
- TMA_A tma_load_a;
211
- TMA_B tma_load_b;
212
- uint32_t tma_transaction_bytes = TmaTransactionBytes;
213
- uint32_t tma_transaction_bytes_mk = TmaTransactionBytesMK;
214
- uint32_t tma_transaction_bytes_nk = TmaTransactionBytesNK;
215
- // Block scaling factors for A and B
216
- ElementBlockScale const* ptr_scale_A;
217
- ElementBlockScale const* ptr_scale_B;
218
- };
219
-
220
- //
221
- // Methods
222
- //
223
-
224
- template <class ProblemShape>
225
- static constexpr Params
226
- to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
227
- (void) workspace;
228
-
229
- // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
230
- auto problem_shape_MNKL = append<4>(problem_shape, 1);
231
- auto [M,N,K,L] = problem_shape_MNKL;
232
-
233
- auto ptr_A = reinterpret_cast<ElementA const*>(args.ptr_A);
234
- auto ptr_B = reinterpret_cast<ElementB const*>(args.ptr_B);
235
-
236
- Tensor tensor_a = make_tensor(ptr_A, make_layout(make_shape(M,K,L), args.dA));
237
- Tensor tensor_b = make_tensor(ptr_B, make_layout(make_shape(N,K,L), args.dB));
238
- typename Params::TMA_A tma_load_a = make_tma_copy_A_sm90(
239
- GmemTiledCopyA{},
240
- tensor_a,
241
- SmemLayoutA{}(_,_,cute::Int<0>{}),
242
- TileShape{},
243
- ClusterShape{});
244
- typename Params::TMA_B tma_load_b = make_tma_copy_B_sm90(
245
- GmemTiledCopyB{},
246
- tensor_b,
247
- SmemLayoutB{}(_,_,cute::Int<0>{}),
248
- TileShape{},
249
- ClusterShape{});
250
- uint32_t transaction_bytes_mk = TmaTransactionBytesMK;
251
- uint32_t transaction_bytes_nk = TmaTransactionBytesNK;
252
- uint32_t transaction_bytes = transaction_bytes_mk + transaction_bytes_nk;
253
-
254
- return {
255
- tma_load_a,
256
- tma_load_b,
257
- transaction_bytes,
258
- transaction_bytes_mk,
259
- transaction_bytes_nk,
260
- args.ptr_scale_A,
261
- args.ptr_scale_B
262
- };
263
- }
264
-
265
- template<class ProblemShape>
266
- static bool
267
- can_implement(
268
- ProblemShape const& problem_shape,
269
- [[maybe_unused]] Arguments const& args) {
270
- constexpr int tma_alignment_bits = 128;
271
- auto problem_shape_MNKL = append<4>(problem_shape, 1);
272
- auto [M,N,K,L] = problem_shape_MNKL;
273
-
274
- bool implementable = true;
275
- constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
276
- implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
277
- constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value;
278
- implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
279
-
280
- if (!implementable) {
281
- CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
282
- }
283
- return implementable;
284
- }
285
-
286
- static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
287
- static constexpr int K_PIPE_MMAS = 1;
288
- static constexpr uint32_t TmaTransactionBytesMK =
289
- cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof_bits<ElementA>::value));
290
- static constexpr uint32_t TmaTransactionBytesNK =
291
- cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof_bits<ElementB>::value));
292
- static constexpr uint32_t TmaTransactionBytes = TmaTransactionBytesMK + TmaTransactionBytesNK;
293
-
294
- /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
295
- CUTLASS_DEVICE
296
- static void prefetch_tma_descriptors(Params const& mainloop_params)
297
- {
298
- cute::prefetch_tma_descriptor(mainloop_params.tma_load_a.get_tma_descriptor());
299
- cute::prefetch_tma_descriptor(mainloop_params.tma_load_b.get_tma_descriptor());
300
- }
301
-
302
- /// Set up the data needed by this collective for load and mma.
303
- /// Returns a tuple of tensors. The collective and the kernel layer have the contract
304
- /// Returned tuple must contain at least two elements, with the first two elements being:
305
- /// gA_mkl - The tma tensor, A after a local tile so it has shape (BLK_M,BLK_K,m,k,l)
306
- /// gB_nkl - The tma tensor, B after a local tile so it has shape (BLK_N,BLK_K,n,k,l)
307
- template <class ProblemShape_MNKL>
308
- CUTLASS_DEVICE auto
309
- load_init(ProblemShape_MNKL const& problem_shape_MNKL, Params const& mainloop_params) const {
310
- using X = Underscore;
311
- // Separate out problem shape for convenience
312
- auto [M,N,K,L] = problem_shape_MNKL;
313
-
314
- // TMA requires special handling of strides to deal with coord codomain mapping
315
- // Represent the full tensors -- get these from TMA
316
- Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l)
317
- Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l)
318
-
319
- // Make tiled views, defer the slice
320
- Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
321
- Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
322
-
323
- constexpr auto scales_m = Int<ScaleMsPerTile>{};
324
- auto tM = get<2>(gA_mkl.shape());
325
- auto tN = get<2>(gB_nkl.shape());
326
- auto tK = get<3>(gA_mkl.shape());
327
-
328
- // Make the tiled views of scale tensors
329
- auto scaleA_shape = make_shape(M / ScaleGranularityM, tK, L); // (scale_m,k,l)
330
- auto scaleA_layout = make_ordered_layout(scaleA_shape, Step<_0, _1, _2>{});
331
- auto scaleB_shape = make_shape(tN, tK, L); // (n,k,l)
332
- auto scaleB_layout = make_ordered_layout(scaleB_shape, Step<_1, _0, _2>{});
333
-
334
- // Note that mScaleA_mkl and mScaleB_nkl are already blocked tiled in the `m` host and
335
- // gScaleA_mkl and gScaleB_nkl in `g` global memory are same as mScaleA_mkl and mScaleB_nkl.
336
- Tensor mScaleA_mkl = make_tensor(make_gmem_ptr(mainloop_params.ptr_scale_A), scaleA_layout); // (scale_m,k,l)
337
- Tensor mScaleB_nkl = make_tensor(make_gmem_ptr(mainloop_params.ptr_scale_B), scaleB_layout); // (n,k,l)
338
-
339
- return cute::make_tuple(gA_mkl, gB_nkl, mScaleA_mkl, mScaleB_nkl);
340
- }
341
-
342
- /// Perform a collective-scoped matrix multiply-accumulate
343
- /// Producer Perspective
344
- template <
345
- class TensorA, class TensorB,
346
- class TensorScaleA, class TensorScaleB,
347
- class KTileIterator, class BlockCoord
348
- >
349
- CUTLASS_DEVICE void
350
- load(
351
- Params const& mainloop_params,
352
- MainloopPipeline pipeline,
353
- PipelineState smem_pipe_write,
354
- cute::tuple<TensorA, TensorB, TensorScaleA, TensorScaleB> const& load_inputs,
355
- BlockCoord const& blk_coord,
356
- KTileIterator k_tile_iter, int k_tile_count,
357
- int thread_idx,
358
- uint32_t block_rank_in_cluster,
359
- TensorStorage& shared_tensors) {
360
- int lane_predicate = cute::elect_one_sync();
361
-
362
- // Blockscaling: Tma loads for load_input and CpAsync for load_scale
363
- Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
364
- Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
365
- Tensor sScaleA = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_A.data()), SmemLayoutScaleA{}); // (ScaleMsPerTile,k)
366
- Tensor sScaleB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_B.data()), SmemLayoutScaleB{}); // (k)
367
-
368
- //
369
- // Prepare the TMA loads for A and B
370
- //
371
-
372
- constexpr uint32_t cluster_shape_x = get<0>(ClusterShape());
373
- uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
374
-
375
- Tensor gA_mkl = get<0>(load_inputs);
376
- Tensor gB_nkl = get<1>(load_inputs);
377
-
378
- auto block_tma_a = mainloop_params.tma_load_a.get_slice(cluster_local_block_id.y);
379
- auto block_tma_b = mainloop_params.tma_load_b.get_slice(cluster_local_block_id.x);
380
-
381
- // Partition the inputs based on the current block coordinates.
382
- auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
383
- Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
384
- Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
385
-
386
-
387
- // Block scaling: load_scale has scaling tensors in global memory which are not tiled
388
- Tensor mScaleA_mkl = get<2>(load_inputs);
389
- Tensor mScaleB_nkl = get<3>(load_inputs);
390
- auto scales_m = get<0>(mScaleA_mkl.shape());
391
-
392
- Tensor cScaleA_mkl = make_identity_tensor(mScaleA_mkl.shape());
393
-
394
- Tensor gScaleA = local_tile(
395
- mScaleA_mkl, make_tile(Int<ScaleMsPerTile>{}),
396
- make_coord(m_coord,_,l_coord)); // (ScaleMsPerTile,k,1)
397
- Tensor cScaleA = local_tile(
398
- cScaleA_mkl, make_tile(Int<ScaleMsPerTile>{}),
399
- make_coord(m_coord,_,l_coord));
400
- Tensor gScaleB = mScaleB_nkl(n_coord,_,l_coord); // (1,k,1)
401
-
402
- // TODO: test `scale_copy_a` with `ScaleMsPerTile` < 128
403
- TiledCopy scale_copy_a = make_tiled_copy(SmemBlockScalingCopyAtomA{},
404
- Layout<Shape<_32>>{}, Layout<Shape<_1>>{}); // (1,1,1)
405
- TiledCopy scale_copy_b = make_tiled_copy(SmemBlockScalingCopyAtomB{},
406
- Layout<Shape<_1>>{}, Layout<Shape<_1>>{}); // (1,1,1)
407
- ThrCopy thr_scale_copy_a = scale_copy_a.get_slice(threadIdx.x);
408
- ThrCopy thr_scale_copy_b = scale_copy_b.get_slice(threadIdx.x);
409
-
410
- Tensor tAgA_ScaleA = thr_scale_copy_a.partition_S(gScaleA);
411
- Tensor tAcA_ScaleA = thr_scale_copy_a.partition_S(cScaleA);
412
- Tensor tAsA_ScaleA = thr_scale_copy_a.partition_D(sScaleA);
413
-
414
- Tensor tBgB_ScaleB = thr_scale_copy_b.partition_S(gScaleB);
415
- Tensor tBsB_ScaleB = thr_scale_copy_b.partition_D(sScaleB);
416
-
417
- // Applies the mapping from block_tma_a
418
- Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k)
419
- Tensor tAsA = block_tma_a.partition_D(sA); // (TMA,TMA_M,TMA_K,PIPE)
420
-
421
- Tensor tBgB = block_tma_b.partition_S(gB); // (TMA,TMA_N,TMA_K,k)
422
- Tensor tBsB = block_tma_b.partition_D(sB); // (TMA,TMA_N,TMA_K,PIPE)
423
-
424
- uint16_t mcast_mask_a = 0;
425
- uint16_t mcast_mask_b = 0;
426
-
427
- // Issue TmaLoads for GEMM operands A/B and CpAsync for scale tensors
428
- // Maps the tile -> block, value
429
- if constexpr (cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>) {
430
- auto block_layout = Layout<typename DispatchPolicy::ClusterShape>{}; // (m,n) -> block_id
431
- for (int n = 0; n < size<1>(block_layout); ++n) {
432
- mcast_mask_a |= (uint16_t(1) << block_layout(cluster_local_block_id.x,n,Int<0>{}));
433
- }
434
- }
435
-
436
- if constexpr (cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>) {
437
- auto block_layout = Layout<typename DispatchPolicy::ClusterShape>{}; // (m,n) -> block_id
438
- for (int m = 0; m < size<0>(block_layout); ++m) {
439
- mcast_mask_b |= (uint16_t(1) << block_layout(m,cluster_local_block_id.y,Int<0>{}));
440
- }
441
- }
442
-
443
- // Allocate predicate tensors for a_scales (since we can't guarantee that
444
- // all scales are valid, since we could have a partial tiles along M)
445
- Tensor tApA_ScaleA = make_tensor<bool>(shape(tAsA_ScaleA(_,_,0)));
446
- #pragma unroll
447
- for (int i = 0; i < size(tApA_ScaleA); ++i) {
448
- tApA_ScaleA(i) = get<0>(tAcA_ScaleA(i)) < scales_m;
449
- }
450
-
451
- // Mainloop
452
- CUTLASS_PRAGMA_NO_UNROLL
453
- for ( ; k_tile_count > 0; --k_tile_count) {
454
- // LOCK smem_pipe_write for _writing_
455
- pipeline.producer_acquire(smem_pipe_write);
456
-
457
- //
458
- // Copy gmem to smem for *k_tile_iter
459
- //
460
- int write_stage = smem_pipe_write.index();
461
- using BarrierType = typename MainloopPipeline::ProducerBarrierType;
462
- BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write);
463
-
464
- // Copy operands A and B from global memory to shared memory
465
- if (lane_predicate) copy(mainloop_params.tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
466
- if (lane_predicate) copy(mainloop_params.tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
467
-
468
- // Copy scale tensors from global memory to shared memory
469
- copy_if(scale_copy_a, tApA_ScaleA, tAgA_ScaleA(_,_,*k_tile_iter), tAsA_ScaleA(_,_,write_stage));
470
- copy(scale_copy_b, tBgB_ScaleB(_,*k_tile_iter), tBsB_ScaleB(_,write_stage));
471
- pipeline.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive_noinc);
472
-
473
- ++k_tile_iter;
474
-
475
- // Advance smem_pipe_write
476
- ++smem_pipe_write;
477
- }
478
- }
479
-
480
- /// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster
481
- CUTLASS_DEVICE void
482
- load_tail(
483
- MainloopPipeline pipeline,
484
- PipelineState smem_pipe_write) {
485
- int lane_predicate = cute::elect_one_sync();
486
-
487
- // Issue the epilogue waits
488
- if (lane_predicate) {
489
- /* This helps avoid early exit of blocks in Cluster
490
- * Waits for all stages to either be released (all
491
- * Consumer UNLOCKs), or if the stage was never used
492
- * then would just be acquired since the phase was
493
- * still inverted from make_producer_start_state
494
- */
495
- pipeline.producer_tail(smem_pipe_write);
496
- }
497
- }
498
-
499
- /// Perform a collective-scoped matrix multiply-accumulate
500
- /// Consumer Perspective
501
- template <
502
- class FrgTensorC
503
- >
504
- CUTLASS_DEVICE void
505
- mma(MainloopPipeline pipeline,
506
- PipelineState smem_pipe_read,
507
- FrgTensorC& accum,
508
- int k_tile_count,
509
- int thread_idx,
510
- TensorStorage& shared_tensors,
511
- Params const& mainloop_params) {
512
-
513
-
514
- static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
515
- static_assert(cute::rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
516
- static_assert(cute::rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
517
- static_assert(cute::is_void_v<SmemCopyAtomA>,
518
- "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
519
- static_assert(cute::is_void_v<SmemCopyAtomB>,
520
- "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
521
-
522
- Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
523
- Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
524
-
525
- // Block scaling
526
- Tensor sScaleAViewAsC = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_A.data()),
527
- Layout<
528
- Shape<Shape<Int<ScaleGranularityM>, Int<ScaleMsPerTile>>, cute::tuple_element_t<1, TileShape>, Int<DispatchPolicy::Stages>>,
529
- Stride<Stride<_0, _1>, _0, Int<ScaleMsPerTile>>
530
- >{}); // ((ScaleGranularityM,ScaleMsPerTile),n,k)
531
- Tensor sScaleB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_B.data()), SmemLayoutScaleB{}); // (k)
532
-
533
- //
534
- // Define C accumulators and A/B partitioning
535
- //
536
-
537
- // Layout of warp group to thread mapping
538
-
539
- static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
540
- stride<0>(typename TiledMma::BLayout{}) == 0 and
541
- size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
542
- size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
543
- "Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
544
-
545
- constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
546
- Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
547
- Int<NumThreadsPerWarpGroup>{});
548
-
549
- int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
550
-
551
- TiledMma tiled_mma;
552
- auto thread_mma = tiled_mma.get_slice(warp_group_thread_layout(warp_group_idx));
553
-
554
- Tensor tCsScaleAViewAsC = tiled_mma.get_slice(thread_idx).partition_C(sScaleAViewAsC); // (MMA,MMA_M,MMA_N,PIPE), `thread_mma` above is correct when partitioning A and B, but it is not correct when partitioning C.
555
-
556
- Tensor tCsA = thread_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
557
- Tensor tCsB = thread_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
558
-
559
- // Allocate "fragments/descriptors"
560
- Tensor tCrA = thread_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)
561
- Tensor tCrB = thread_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE)
562
-
563
- CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum)); // M
564
- CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N
565
- CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K
566
- CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE
567
- CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA)); // PIPE
568
- CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB)); // PIPE
569
-
570
- //
571
- // PIPELINED MAIN LOOP
572
- //
573
- static_assert((0 <= K_PIPE_MMAS) && (K_PIPE_MMAS < K_PIPE_MAX),
574
- "ERROR : Incorrect number of MMAs in flight");
575
-
576
- // We release buffers to producer warps(dma load) with some mmas in flight
577
- PipelineState smem_pipe_release = smem_pipe_read;
578
-
579
- // Per block scale values for operand A and B
580
-
581
- using RegLayoutScaleAViewAsC = decltype(make_layout_like(tCsScaleAViewAsC(_, _, _, 0).layout())); // `make_layout_like` makes a compact layout.
582
- using RegLayoutScaleAEssential = decltype(filter_zeros(RegLayoutScaleAViewAsC{}.stride(), RegLayoutScaleAViewAsC{}.shape())); // an interface to traverse the underlying storage for the compact layout mentioned above
583
-
584
- Tensor tCrScaleAViewAsC = make_tensor<ElementBlockScale>(RegLayoutScaleAViewAsC{}); // (MMA,MMA_M,MMA_N)
585
- ElementBlockScale scale_b;
586
-
587
- // Prologue GMMAs
588
- int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
589
-
590
- tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
591
-
592
- GmmaFP8AccumulationWithScale accumulation(accum, size<2>(TileShape{}) / size<2>(typename TiledMma::AtomShape_MNK{}), size<2>(tCrA));
593
- warpgroup_fence_operand(accumulation());
594
- CUTLASS_PRAGMA_UNROLL
595
- for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue)
596
- {
597
- // WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
598
- auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
599
- pipeline.consumer_wait(smem_pipe_read, barrier_token);
600
-
601
- if (accumulation.prepare_if_needed()) {
602
- tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
603
- }
604
-
605
- int read_stage = smem_pipe_read.index();
606
-
607
- // Load per block scale values from shared memory to registers.
608
- scale_b = sScaleB[read_stage];
609
- CUTLASS_PRAGMA_UNROLL
610
- for (int i = 0; i < size(RegLayoutScaleAEssential{}); i++) {
611
- tCrScaleAViewAsC.data()[i] = tCsScaleAViewAsC(_, _, _, read_stage)(idx2crd(i, RegLayoutScaleAEssential{}));
612
- }
613
- if constexpr (ScaleMsPerTile == 1) {
614
- static_assert(size(RegLayoutScaleAEssential{}) == 1);
615
- tCrScaleAViewAsC.data()[0] = __shfl_sync(0xffffffff, tCrScaleAViewAsC.data()[0] * scale_b, 0); // `tCrScaleAViewAsC.data()[0]` are all same in a warp group when `ScaleMsPerTile == 1`.
616
- } else {
617
- CUTLASS_PRAGMA_UNROLL
618
- for (int i = 0; i < size(RegLayoutScaleAEssential{}); i++) {
619
- tCrScaleAViewAsC.data()[i] = tCrScaleAViewAsC.data()[i] * scale_b;
620
- }
621
- }
622
-
623
- warpgroup_arrive();
624
- // Unroll the K mode manually to set scale D to 1
625
- CUTLASS_PRAGMA_UNROLL
626
- for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
627
- // (V,M,K) x (V,N,K) => (V,M,N)
628
- cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
629
- tiled_mma.accumulate_ = GMMA::ScaleOut::One;
630
- }
631
- warpgroup_commit_batch();
632
-
633
- // Block scale the accumulators with reg tensor `tCrScaleAViewAsC`
634
- accumulation.scale_if_needed(tCrScaleAViewAsC);
635
-
636
- ++smem_pipe_read;
637
- }
638
-
639
- warpgroup_fence_operand(accumulation());
640
- // Mainloop GMMAs
641
- k_tile_count -= prologue_mma_count;
642
-
643
- CUTLASS_PRAGMA_NO_UNROLL
644
- for ( ; k_tile_count > 0; --k_tile_count)
645
- {
646
- // WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
647
- auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
648
- pipeline.consumer_wait(smem_pipe_read, barrier_token);
649
-
650
- //
651
- // Compute on k_tile
652
- //
653
-
654
- int read_stage = smem_pipe_read.index();
655
-
656
- // Load per block scale values from shared memory to registers (at most twice per block along M and exactly once per block along N)
657
- scale_b = sScaleB[read_stage];
658
- CUTLASS_PRAGMA_UNROLL
659
- for (int i = 0; i < size(RegLayoutScaleAEssential{}); i++) {
660
- tCrScaleAViewAsC.data()[i] = tCsScaleAViewAsC(_, _, _, read_stage)(idx2crd(i, RegLayoutScaleAEssential{}));
661
- }
662
- if constexpr (ScaleMsPerTile == 1) {
663
- static_assert(size(RegLayoutScaleAEssential{}) == 1);
664
- tCrScaleAViewAsC.data()[0] = __shfl_sync(0xffffffff, tCrScaleAViewAsC.data()[0] * scale_b, 0); // `tCrScaleAViewAsC.data()[0]` are all same in a warp group when `ScaleMsPerTile == 1`.
665
- } else {
666
- CUTLASS_PRAGMA_UNROLL
667
- for (int i = 0; i < size(RegLayoutScaleAEssential{}); i++) {
668
- tCrScaleAViewAsC.data()[i] = tCrScaleAViewAsC.data()[i] * scale_b;
669
- }
670
- }
671
-
672
- if (accumulation.prepare_if_needed()) {
673
- tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
674
- }
675
-
676
- warpgroup_fence_operand(accumulation());
677
- warpgroup_arrive();
678
- // Unroll the K mode manually to set scale D to 1
679
- CUTLASS_PRAGMA_UNROLL
680
- for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
681
- // (V,M,K) x (V,N,K) => (V,M,N)
682
- cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
683
- tiled_mma.accumulate_ = GMMA::ScaleOut::One;
684
- }
685
- warpgroup_commit_batch();
686
-
687
- /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
688
- warpgroup_wait<K_PIPE_MMAS>();
689
- warpgroup_fence_operand(accumulation());
690
-
691
- // Block scale the accumulators with reg tensor `tCrScaleAViewAsC`
692
- accumulation.scale_if_needed(tCrScaleAViewAsC);
693
-
694
- pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
695
-
696
- // Advance smem_pipe_read and smem_pipe_release
697
- ++smem_pipe_read;
698
- ++smem_pipe_release;
699
- }
700
-
701
- accumulation.scale_residue_if_needed(tCrScaleAViewAsC);
702
-
703
- warpgroup_fence_operand(accumulation());
704
- }
705
-
706
- /// Perform a Consumer Epilogue to release all buffers
707
- CUTLASS_DEVICE void
708
- mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) {
709
- // Prologue GMMAs
710
- int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
711
- k_tile_count -= prologue_mma_count;
712
-
713
- smem_pipe_release.advance(k_tile_count);
714
-
715
- // Wait on all GMMAs to complete
716
- warpgroup_wait<0>();
717
-
718
- for (int count = 0; count < prologue_mma_count; ++count) {
719
- pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
720
- ++smem_pipe_release;
721
- }
722
- }
723
-};
724
-
725
-/////////////////////////////////////////////////////////////////////////////////////////////////
726
-
727
-} // namespace cutlass::gemm::collective
728
-
729
-/////////////////////////////////////////////////////////////////////////////////////////////////
csrc/cutlass_extensions/gemm/dispatch_policy.hpp
deleted
-39
@@ -1,39 +0,0 @@
1
-#pragma once
2
-
3
-#include "cutlass/gemm/dispatch_policy.hpp"
4
-
5
-namespace cutlass::gemm {
6
-
7
-//////////////////////////////////////////////////////////////////////////////
8
-
9
-// FP8 related policies (including Blocked Scaled Accumulation)
10
-// `ScaleGranularityM` specifies scaling granularity along M, while zero-value
11
-// `ScaleGranularityM` indicates that scaling granularity is
12
-// `size<0>(TileShape_MNK{})` along M.
13
-template <int ScaleGranularityM = 0>
14
-struct KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum
15
- : KernelTmaWarpSpecializedCooperative {};
16
-
17
-// n-buffer in smem (Hopper TMA), pipelined with Hopper GMMA and TMA, Warp
18
-// specialized dynamic schedule For FP8 kernels with Block Scaling
19
-template <int Stages_, class ClusterShape_ = Shape<_1, _1, _1>,
20
- class KernelSchedule = KernelTmaWarpSpecialized,
21
- int ScaleGranularityM =
22
- 0 // `ScaleGranularityM` specifies scaling granularity along M,
23
- // while zero-value `ScaleGranularityM` indicates that scaling
24
- // granularity is `size<0>(TileShape_MNK{})` along M.
25
- >
26
-struct MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8
27
- : MainloopSm90TmaGmmaWarpSpecialized<Stages_, ClusterShape_,
28
- KernelSchedule> {
29
- static_assert(
30
- cute::is_same_v<
31
- KernelSchedule,
32
- KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<
33
- ScaleGranularityM>>,
34
- "KernelSchedule must be one of the warp specialized policies");
35
-};
36
-
37
-//////////////////////////////////////////////////////////////////////////////
38
-
39
-} // namespace cutlass::gemm
\ No newline at end of file
csrc/cutlass_extensions/vllm_collective_builder.cuh
+1
-1
@@ -1,6 +1,6 @@
1
#pragma once
2
3
-#include "cutlass_extensions/gemm/collective/collective_builder.hpp"
3
+#include "cutlass/gemm/collective/collective_builder.hpp"
4
5
namespace cutlass::gemm::collective {
6
using namespace cute;
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh
-3
@@ -14,9 +14,6 @@
14
#include "cutlass/epilogue/dispatch_policy.hpp"
15
#include "cutlass/epilogue/collective/collective_builder.hpp"
16
17
-#include "cutlass_extensions/gemm/dispatch_policy.hpp"
18
-#include "cutlass_extensions/gemm/collective/collective_builder.hpp"
19
-
17
#include "cutlass_gemm_caller.cuh"
18
19
namespace vllm {
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh
-3
@@ -14,9 +14,6 @@
14
#include "cutlass/epilogue/dispatch_policy.hpp"
15
#include "cutlass/epilogue/collective/collective_builder.hpp"
16
17
-#include "cutlass_extensions/gemm/dispatch_policy.hpp"
18
-#include "cutlass_extensions/gemm/collective/collective_builder.hpp"
19
-
17
#include "cutlass_gemm_caller.cuh"
18
19
namespace vllm {
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh
+90
-112
@@ -13,27 +13,18 @@
13
#include "cutlass/epilogue/dispatch_policy.hpp"
14
#include "cutlass/epilogue/collective/collective_builder.hpp"
15
16
-#include "cutlass_extensions/gemm/dispatch_policy.hpp"
17
-#include "cutlass_extensions/gemm/collective/collective_builder.hpp"
18
-
16
#include "cutlass_gemm_caller.cuh"
17
18
namespace vllm {
19
20
using namespace cute;
21
25
-template <typename SchedulerType, typename OutType, int GroupSizeM_,
26
- int GroupSizeN_, int GroupSizeK_, int TileSizeM_ = 128,
27
- class ClusterShape = Shape<_1, _2, _1>>
22
+// clang-format off
23
+template <class OutType, int ScaleGranularityM,
24
+ int ScaleGranularityN, int ScaleGranularityK,
25
+ class MmaTileShape, class ClusterShape,
26
+ class EpilogueScheduler, class MainloopScheduler>
27
struct cutlass_3x_gemm_fp8_blockwise {
29
- using GroupSizeM = Int<GroupSizeM_>;
30
- using GroupSizeN = Int<GroupSizeN_>;
31
- using GroupSizeK = Int<GroupSizeK_>;
32
- using TileSizeM = Int<TileSizeM_>;
33
-
34
- static_assert(TileSizeM_ % GroupSizeM_ == 0,
35
- "TileSizeM must be a multiple of GroupSizeM");
36
-
28
using ElementAB = cutlass::float_e4m3_t;
29
30
using ElementA = ElementAB;
@@ -45,52 +36,67 @@ struct cutlass_3x_gemm_fp8_blockwise {
36
static constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
37
38
using ElementD = OutType;
48
- using StrideD = Stride<int64_t, Int<1>, Int<0>>;
39
+ using LayoutD = cutlass::layout::RowMajor;
40
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
41
51
- using ElementC = void;
52
- using StrideC = StrideD;
42
+ using ElementC = void; // TODO: support bias
43
+ using LayoutC = LayoutD;
44
static constexpr int AlignmentC = AlignmentD;
45
46
using ElementAccumulator = float;
56
- using ElementBlockScale = float;
47
using ElementCompute = float;
48
+ using ElementBlockScale = float;
49
+
50
+ using ScaleConfig = cutlass::detail::Sm90BlockwiseScaleConfig<
51
+ ScaleGranularityM, ScaleGranularityN, ScaleGranularityK>;
52
+
53
+ using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());
54
+ using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB());
55
+
56
using ArchTag = cutlass::arch::Sm90;
57
using OperatorClass = cutlass::arch::OpClassTensorOp;
60
- using TileShape = Shape<TileSizeM, GroupSizeN, GroupSizeK>;
61
-
62
- using KernelSchedule = cutlass::gemm::
63
- KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<
64
- GroupSizeM_>;
65
- using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
66
- using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
67
-
68
- using StoreEpilogueCompute = typename cutlass::epilogue::fusion::Sm90EVT<
69
- cutlass::epilogue::fusion::Sm90AccFetch>;
70
-
71
- using CollectiveEpilogue =
72
- typename cutlass::epilogue::collective::CollectiveBuilder<
73
- ArchTag, OperatorClass, TileShape, ClusterShape, EpilogueTileType,
74
- ElementAccumulator, ElementCompute, ElementC, StrideC, AlignmentC,
75
- ElementD, StrideD, AlignmentD, EpilogueSchedule,
76
- StoreEpilogueCompute>::CollectiveOp;
77
-
78
- using CollectiveMainloop =
79
- typename cutlass::gemm::collective::CollectiveBuilder<
80
- ArchTag, OperatorClass, ElementA, LayoutA, AlignmentA, ElementB,
81
- LayoutB, AlignmentB, ElementAccumulator, TileShape, ClusterShape,
82
- cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
83
- sizeof(typename CollectiveEpilogue::SharedStorage))>,
84
- KernelSchedule>::CollectiveOp;
58
+
59
+ static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
60
+ using ElementScalar = float;
61
+ using DefaultOperation = cutlass::epilogue::fusion::LinearCombination<ElementD, ElementCompute, ElementC, ElementScalar, RoundStyle>;
62
+ using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
63
+ ArchTag,
64
+ OperatorClass,
65
+ MmaTileShape,
66
+ ClusterShape,
67
+ cutlass::epilogue::collective::EpilogueTileAuto,
68
+ ElementAccumulator,
69
+ ElementCompute,
70
+ ElementC,
71
+ LayoutC,
72
+ AlignmentC,
73
+ ElementD,
74
+ LayoutD,
75
+ AlignmentD,
76
+ EpilogueScheduler,
77
+ DefaultOperation
78
+ >::CollectiveOp;
79
+
80
+ using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
81
+ ArchTag,
82
+ OperatorClass,
83
+ ElementA,
84
+ cute::tuple<LayoutA, LayoutSFA>,
85
+ AlignmentA,
86
+ ElementB,
87
+ cute::tuple<LayoutB, LayoutSFB>,
88
+ AlignmentB,
89
+ ElementAccumulator,
90
+ MmaTileShape,
91
+ ClusterShape,
92
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
93
+ MainloopScheduler
94
+ >::CollectiveOp;
95
96
using KernelType = enable_sm90_or_later<cutlass::gemm::kernel::GemmUniversal<
87
- Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
88
- SchedulerType>>;
97
+ Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue>>;
98
99
struct GemmKernel : public KernelType {};
91
-
92
- using StrideA = typename GemmKernel::StrideA;
93
- using StrideB = typename GemmKernel::StrideB;
100
};
101
102
template <typename Gemm>
@@ -99,76 +105,54 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
105
torch::Tensor const& a_scales,
106
torch::Tensor const& b_scales) {
107
using GemmKernel = typename Gemm::GemmKernel;
108
+ using StrideA = typename Gemm::GemmKernel::StrideA;
109
+ using StrideB = typename Gemm::GemmKernel::StrideB;
110
+ using StrideD = typename Gemm::GemmKernel::StrideD;
111
+ using StrideC = typename Gemm::GemmKernel::StrideC;
112
+ using LayoutSFA = typename Gemm::LayoutSFA;
113
+ using LayoutSFB = typename Gemm::LayoutSFB;
114
+ using ScaleConfig = typename Gemm::ScaleConfig;
115
116
using ElementAB = typename Gemm::ElementAB;
117
using ElementD = typename Gemm::ElementD;
118
106
- auto prob_shape = c3x::get_problem_shape(a, b);
107
- int32_t m = get<0>(prob_shape), n = get<1>(prob_shape),
108
- k = get<2>(prob_shape);
119
+ int32_t m = a.size(0), n = b.size(1), k = a.size(1);
120
110
- int64_t lda = a.stride(0);
111
- int64_t ldb = b.stride(1);
112
- int64_t ldc = out.stride(0);
121
+ TORCH_CHECK(m % 4 == 0, "m must be divisible by 4");
122
114
- using StrideA = Stride<int64_t, Int<1>, int64_t>;
115
- using StrideB = Stride<int64_t, Int<1>, int64_t>;
116
- using StrideC = typename Gemm::StrideC;
123
+ StrideA a_stride;
124
+ StrideB b_stride;
125
+ StrideC c_stride;
126
+ a_stride =
127
+ cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1));
128
+ b_stride =
129
+ cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1));
130
+ c_stride =
131
+ cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1));
132
118
- StrideA a_stride{lda, Int<1>{}, 0};
119
- StrideB b_stride{ldb, Int<1>{}, 0};
120
- StrideC c_stride{ldc, Int<1>{}, Int<0>{}};
133
+ LayoutSFA layout_SFA =
134
+ ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1));
135
+ LayoutSFB layout_SFB =
136
+ ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
137
138
auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
139
auto b_ptr = static_cast<ElementAB*>(b.data_ptr());
140
auto a_scales_ptr = static_cast<float*>(a_scales.data_ptr());
141
auto b_scales_ptr = static_cast<float*>(b_scales.data_ptr());
142
127
- // Check is the t is contiguous and is 1D or 2D with one of the dimensions
128
- // being 1 (i.e. a row or column vector)
129
- auto is_contiguous_vector = [](const torch::Tensor& t) {
130
- auto t_sizes = t.sizes();
131
- return t.is_contiguous() &&
132
- (t.dim() == 1 ||
133
- (t.dim() == 2 &&
134
- *std::min_element(t_sizes.begin(), t_sizes.end()) == 1));
135
- };
136
-
137
- // TODO(lucas): lets clean-up the kernel so that we pass in Strides so
138
- // we don't have to deal with enforcing implicit layouts
139
- TORCH_CHECK(a_scales.size(0) == m / Gemm::GroupSizeM::value);
140
- TORCH_CHECK(a_scales.size(1) == k / Gemm::GroupSizeK::value);
141
- TORCH_CHECK(a_scales.stride(0) == 1 || is_contiguous_vector(a_scales),
142
- "a_scales must be M major");
143
- TORCH_CHECK(b_scales.size(0) == k / Gemm::GroupSizeK::value);
144
- TORCH_CHECK(b_scales.size(1) == n / Gemm::GroupSizeN::value);
145
- TORCH_CHECK(b_scales.stride(0) == 1 || is_contiguous_vector(b_scales),
146
- "b_scales must be K major");
147
- typename GemmKernel::MainloopArguments mainloop_args{
148
- a_ptr, a_stride, b_ptr, b_stride, a_scales_ptr, b_scales_ptr};
143
+ auto mainloop_args = [&](){
144
+ return typename GemmKernel::MainloopArguments{
145
+ a_ptr, a_stride, b_ptr, b_stride,
146
+ a_scales_ptr, layout_SFA, b_scales_ptr, layout_SFB
147
+ };
148
+ }();
149
+ auto prob_shape = cute::make_shape(m, n, k, 1);
150
151
auto c_ptr = static_cast<ElementD*>(out.data_ptr());
152
typename GemmKernel::EpilogueArguments epilogue_args{
153
{}, c_ptr, c_stride, c_ptr, c_stride};
153
-
154
- typename GemmKernel::TileSchedulerArguments scheduler;
155
-
156
- static constexpr bool UsesStreamKScheduler =
157
- cute::is_same_v<typename GemmKernel::TileSchedulerTag,
158
- cutlass::gemm::StreamKScheduler>;
159
-
160
- if constexpr (UsesStreamKScheduler) {
161
- using DecompositionMode = typename cutlass::gemm::kernel::detail::
162
- PersistentTileSchedulerSm90StreamKParams::DecompositionMode;
163
- using ReductionMode = typename cutlass::gemm::kernel::detail::
164
- PersistentTileSchedulerSm90StreamKParams::ReductionMode;
165
-
166
- scheduler.decomposition_mode = DecompositionMode::StreamK;
167
- scheduler.reduction_mode = ReductionMode::Nondeterministic;
168
- }
169
-
154
c3x::cutlass_gemm_caller<GemmKernel>(a.device(), prob_shape, mainloop_args,
171
- epilogue_args, scheduler);
155
+ epilogue_args);
156
}
157
158
template <typename OutType>
@@ -177,18 +161,12 @@ void cutlass_gemm_blockwise_sm90_fp8_dispatch(torch::Tensor& out,
161
torch::Tensor const& b,
162
torch::Tensor const& a_scales,
163
torch::Tensor const& b_scales) {
180
- auto k = a.size(1);
181
- auto n = b.size(1);
182
-
183
- if (k > 3 * n) {
184
- cutlass_gemm_caller_blockwise<cutlass_3x_gemm_fp8_blockwise<
185
- cutlass::gemm::StreamKScheduler, OutType, 1, 128, 128>>(
186
- out, a, b, a_scales, b_scales);
187
- } else {
188
- cutlass_gemm_caller_blockwise<cutlass_3x_gemm_fp8_blockwise<
189
- cutlass::gemm::PersistentScheduler, OutType, 1, 128, 128>>(
190
- out, a, b, a_scales, b_scales);
191
- }
164
+ // TODO: better heuristics
165
+ cutlass_gemm_caller_blockwise<cutlass_3x_gemm_fp8_blockwise<
166
+ OutType, 1, 128, 128, Shape<_128, _128, _128>,
167
+ Shape<_1, _2, _1>, cutlass::epilogue::TmaWarpSpecializedCooperative,
168
+ cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum>>(
169
+ out, a, b, a_scales, b_scales);
170
}
171
172
} // namespace vllm
\ No newline at end of file
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_helper.hpp
+1
-27
@@ -32,7 +32,7 @@ void dispatch_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
32
TORCH_CHECK(a_scales.dim() == 2, "a scale must be 2d tensor.");
33
TORCH_CHECK(b_scales.dim() == 2, "b scale must be 2d tensor.");
34
int32_t version_num = get_sm_version_num();
35
- if (version_num >= 100) {
35
+ if (version_num >= 90) {
36
TORCH_CHECK(
37
a.size(0) == a_scales.size(0) &&
38
cuda_utils::ceil_div(a.size(1), int64_t(128)) == a_scales.size(1),
@@ -41,32 +41,6 @@ void dispatch_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
41
cuda_utils::ceil_div(b.size(0), int64_t(128)) == b_scales.size(0) &&
42
cuda_utils::ceil_div(b.size(1), int64_t(128)) == b_scales.size(1),
43
"b_scale_group_shape must be [128, 128].");
44
- } else {
45
- // TODO: Remove this after using cutlass sm90 blockwise scaling gemm
46
- // kernel, or introducing ceil_div to the load_init() of mainloop.
47
- using GroupShape = std::array<int64_t, 2>;
48
- auto make_group_shape = [](torch::Tensor const& x,
49
- torch::Tensor const& s) -> GroupShape {
50
- TORCH_CHECK(s.dim() == 2, "cutlass_scaled_mm group scales must be 2D");
51
- return {cuda_utils::ceil_div(x.size(0), s.size(0)),
52
- cuda_utils::ceil_div(x.size(1), s.size(1))};
53
- };
54
-
55
- GroupShape a_scale_group_shape = make_group_shape(a, a_scales);
56
- GroupShape b_scale_group_shape = make_group_shape(b, b_scales);
57
-
58
- // 1x128 per-token group scales for activations
59
- // 128x128 blockwise scales for weights
60
- TORCH_CHECK((a_scale_group_shape == GroupShape{1, 128} &&
61
- b_scale_group_shape == GroupShape{128, 128} &&
62
- a.dtype() == torch::kFloat8_e4m3fn &&
63
- b.dtype() == torch::kFloat8_e4m3fn),
64
- "cutlass_scaled_mm only supports datatype float8_e4m3fn.\n"
65
- "a_scale_group_shape must be [1, 128]. Got: [",
66
- a_scale_group_shape[0], ", ", a_scale_group_shape[1],
67
- "]\n"
68
- "b_scale_group_shape must be [128, 128]. Got: [",
69
- b_scale_group_shape[0], ", ", b_scale_group_shape[1], "]");
44
}
45
46
TORCH_CHECK(!bias, "Bias not yet supported blockwise scaled_mm");
tests/kernels/quantization/test_block_fp8.py
+50
-2
@@ -11,8 +11,8 @@ from tests.kernels.quant_utils import (native_per_token_group_quant_fp8,
11
native_w8a8_block_matmul)
12
from vllm.config import VllmConfig
13
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
14
- get_col_major_tma_aligned_tensor, per_token_group_quant_fp8,
15
- w8a8_block_fp8_matmul)
14
+ cutlass_scaled_mm, get_col_major_tma_aligned_tensor,
15
+ per_token_group_quant_fp8, w8a8_block_fp8_matmul)
16
from vllm.platforms import current_platform
17
from vllm.utils import has_deep_gemm
18
from vllm.utils.deep_gemm import fp8_gemm_nt, per_block_cast_to_fp8
@@ -98,6 +98,54 @@ def test_w8a8_block_fp8_matmul(M, N, K, block_size, out_dtype, seed):
98
assert rel_diff < 0.001
99
100
101
+@torch.inference_mode()
102
+def test_w8a8_block_fp8_cutlass_matmul():
103
+ # Test simple case where weight.shape % 128 != 0,
104
+ # like in DSV3 kv_a_proj_with_mqa
105
+ M = 32
106
+ N = 576
107
+ K = 7168
108
+ block_size = [128, 128]
109
+ out_dtype = torch.bfloat16
110
+ seed = 0
111
+
112
+ torch.manual_seed(seed)
113
+ factor_for_scale = 1e-2
114
+ fp8_info = torch.finfo(torch.float8_e4m3fn)
115
+ fp8_max, fp8_min = fp8_info.max, fp8_info.min
116
+
117
+ A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
118
+
119
+ B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
120
+ B_fp8 = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
121
+
122
+ block_n, block_k = block_size[0], block_size[1]
123
+ n_tiles = (N + block_n - 1) // block_n
124
+ k_tiles = (K + block_k - 1) // block_k
125
+
126
+ Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale
127
+ # Hopper requires row-major format for scales
128
+ Bs_cutlass = Bs.T.contiguous() if current_platform.is_device_capability(
129
+ 90) else Bs
130
+
131
+ A_fp8, As = per_token_group_quant_fp8(A_fp32,
132
+ block_size[1],
133
+ column_major_scales=False)
134
+ # CUTLASS uses column-major format for scales
135
+ A_fp8_cutlass, As_cutlass = per_token_group_quant_fp8(
136
+ A_fp32, block_size[1], column_major_scales=True)
137
+
138
+ ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size,
139
+ out_dtype)
140
+ out = cutlass_scaled_mm(A_fp8_cutlass, B_fp8, As_cutlass, Bs_cutlass,
141
+ block_size, out_dtype)
142
+
143
+ rel_diff = (torch.mean(
144
+ torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
145
+ torch.mean(torch.abs(ref_out.to(torch.float32))))
146
+ assert rel_diff < 0.001
147
+
148
+
149
@pytest.mark.parametrize(
150
"M,N,K,block_size,out_dtype,seed",
151
itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS))
vllm/model_executor/layers/quantization/fp8.py
+11
-1
@@ -30,7 +30,8 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
30
register_moe_scaling_factors, rotate_flashinfer_fp8_moe_weights,
31
select_cutlass_fp8_gemm_impl, swap_w13_to_w31)
32
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
33
- get_col_major_tma_aligned_tensor, requant_weight_ue8m0_inplace)
33
+ get_col_major_tma_aligned_tensor, requant_weight_ue8m0_inplace,
34
+ should_use_deepgemm_for_fp8_linear)
35
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
36
apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin,
37
prepare_moe_fp8_layer_for_marlin)
@@ -462,6 +463,15 @@ class Fp8LinearMethod(LinearMethodBase):
463
block_sz,
464
)
465
466
+ # SM90 Block FP8 CUTLASS requires row-major weight scales
467
+ if (self.block_quant and current_platform.is_device_capability(90)
468
+ and self.cutlass_block_fp8_supported
469
+ and not should_use_deepgemm_for_fp8_linear(
470
+ torch.bfloat16, layer.weight)):
471
+ layer.weight_scale_inv = Parameter(
472
+ layer.weight_scale_inv.data.T.contiguous(),
473
+ requires_grad=False)
474
+
475
def apply(self,
476
layer: torch.nn.Module,
477
x: torch.Tensor,
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+25
-25
@@ -40,11 +40,14 @@ def cutlass_scaled_mm(
40
block_size: list[int],
41
output_dtype: torch.dtype = torch.float16,
42
) -> torch.Tensor:
43
- return ops.cutlass_scaled_mm(A,
44
- B.T,
45
- out_dtype=output_dtype,
46
- scale_a=As,
47
- scale_b=Bs.T)
43
+ return ops.cutlass_scaled_mm(
44
+ A,
45
+ B.T,
46
+ out_dtype=output_dtype,
47
+ scale_a=As,
48
+ # SM90 block FP8 requires row-major scale_b, which we do ahead of time
49
+ scale_b=Bs if block_size is not None
50
+ and current_platform.is_device_capability(90) else Bs.T)
51
52
53
def rocm_aiter_gemm_w8a8_blockscale_impl(
@@ -152,35 +155,32 @@ def apply_w8a8_block_fp8_linear(
155
output += bias
156
return output.to(dtype=output_dtype).view(*output_shape)
157
155
- if current_platform.is_cuda():
156
- if current_platform.has_device_capability(100):
157
-
158
- use_cutlass = cutlass_block_fp8_supported and (
159
- cdiv(weight.shape[0], 128) == weight_scale.shape[0]
160
- and cdiv(weight.shape[1], 128) == weight_scale.shape[1])
161
- else:
162
- # TODO: update this after switching to public sm90 block scale gemm
163
- # as it also supports weight.shape % 128 != 0
164
- use_cutlass = cutlass_block_fp8_supported and (
165
- weight.shape[0] % 128 == 0 and weight.shape[1] % 128 == 0)
166
- else:
167
- use_cutlass = False
168
-
158
w8a8_blockscale_func = dispatch_w8a8_blockscale_func(
170
- use_cutlass, use_aiter_and_is_supported)
171
- if use_cutlass:
172
- q_input, x_scale = per_token_group_quant_fp8(
173
- input_2d, block_size[1], column_major_scales=use_cutlass)
159
+ cutlass_block_fp8_supported, use_aiter_and_is_supported)
160
+ if cutlass_block_fp8_supported:
161
+ num_pad = 0
162
+ if current_platform.is_device_capability(90):
163
+ # pad first dimension to be divisible by 4 due to
164
+ # cutlass blockwise gemm limitation for hopper
165
+ num_pad = 4 - (input_2d.shape[0] % 4)
166
+ if num_pad > 0:
167
+ input_2d = torch.nn.functional.pad(input_2d,
168
+ (0, 0, 0, num_pad),
169
+ "constant", 0)
170
+ q_input, x_scale = per_token_group_quant_fp8(input_2d,
171
+ block_size[1],
172
+ column_major_scales=True)
173
output = w8a8_blockscale_func(q_input, weight, x_scale, weight_scale,
174
block_size, input.dtype)
176
-
175
+ if num_pad > 0:
176
+ output = output[:-num_pad]
177
else:
178
if use_aiter_and_is_supported:
179
q_input, x_scale = aiter_per1x128_quant(
180
input_2d.contiguous(), quant_dtype=rocm_aiter.dtypes.fp8)
181
else:
182
q_input, x_scale = per_token_group_quant_fp8(
183
- input_2d, block_size[1], column_major_scales=use_cutlass)
183
+ input_2d, block_size[1], column_major_scales=False)
184
185
output = w8a8_blockscale_func(q_input, weight, x_scale, weight_scale,
186
block_size, input.dtype)