|
6 | 6 | #include <fmha_fwd.hpp> |
7 | 7 | #include <mask.hpp> |
8 | 8 |
|
| 9 | +#include <cstdio> |
| 10 | + |
| 11 | +namespace { |
| 12 | +/* Debug traces. Comment in to aid in debug of 'Invalid argument to fmha_fwd' error |
| 13 | +const char* mask_enum_str(mask_enum m) { |
| 14 | + switch (m) { |
| 15 | + case mask_enum::no_mask: return "no_mask"; |
| 16 | + case mask_enum::mask_top_left: return "mask_top_left"; |
| 17 | + case mask_enum::mask_bottom_right:return "mask_bottom_right"; |
| 18 | + case mask_enum::window_generic: return "window_generic"; |
| 19 | + default: return "unknown"; |
| 20 | + } |
| 21 | +} |
| 22 | +
|
| 23 | +const char* bias_enum_str(bias_enum b) { |
| 24 | + switch (b) { |
| 25 | + case bias_enum::no_bias: return "no_bias"; |
| 26 | + case bias_enum::elementwise_bias: return "elementwise_bias"; |
| 27 | + case bias_enum::alibi: return "alibi"; |
| 28 | + default: return "unknown"; |
| 29 | + } |
| 30 | +} |
| 31 | +
|
| 32 | +const char* qscale_enum_str(quant_scale_enum q) { |
| 33 | + switch (q) { |
| 34 | + case quant_scale_enum::no_scale: return "no_scale"; |
| 35 | + case quant_scale_enum::pertensor: return "pertensor"; |
| 36 | + case quant_scale_enum::blockscale: return "blockscale"; |
| 37 | + case quant_scale_enum::kv_blockscale: return "kv_blockscale"; |
| 38 | + default: return "unknown"; |
| 39 | + } |
| 40 | +} |
| 41 | +
|
| 42 | +void dump_fmha_fwd_kernel_selection_params(const fmha_fwd_traits& t, |
| 43 | + const fmha_fwd_args& a) { |
| 44 | + std::printf("=== fmha_fwd kernel selection params ===\n"); |
| 45 | + std::printf("--- traits ---\n"); |
| 46 | + std::printf(" data_type : %s\n", t.data_type.c_str()); |
| 47 | + std::printf(" hdim_q : %d\n", t.hdim_q); |
| 48 | + std::printf(" hdim_v : %d\n", t.hdim_v); |
| 49 | + std::printf(" is_group_mode : %s\n", t.is_group_mode ? "true" : "false"); |
| 50 | + std::printf(" is_v_rowmajor : %s\n", t.is_v_rowmajor ? "true" : "false"); |
| 51 | + std::printf(" has_logits_soft_cap: %s\n", t.has_logits_soft_cap ? "true" : "false"); |
| 52 | + std::printf(" mask_type : %s (%d)\n", mask_enum_str(t.mask_type), |
| 53 | + static_cast<int>(t.mask_type)); |
| 54 | + std::printf(" bias_type : %s (%d)\n", bias_enum_str(t.bias_type), |
| 55 | + static_cast<int>(t.bias_type)); |
| 56 | + std::printf(" has_lse : %s\n", t.has_lse ? "true" : "false"); |
| 57 | + std::printf(" has_dropout : %s\n", t.has_dropout ? "true" : "false"); |
| 58 | + std::printf(" qscale_type : %s (%d)\n", qscale_enum_str(t.qscale_type), |
| 59 | + static_cast<int>(t.qscale_type)); |
| 60 | + std::printf(" skip_min_seqlen_q : %s\n", t.skip_min_seqlen_q ? "true" : "false"); |
| 61 | + std::printf(" has_sink : %s\n", t.has_sink ? "true" : "false"); |
| 62 | + std::printf("--- args ---\n"); |
| 63 | + std::printf(" hdim_q : %d\n", static_cast<int>(a.hdim_q)); |
| 64 | + std::printf(" hdim_v : %d\n", static_cast<int>(a.hdim_v)); |
| 65 | + std::printf(" seqlen_k : %d\n", static_cast<int>(a.seqlen_k)); |
| 66 | + std::printf(" cu_seqlen_k_ptr : %p\n", a.cu_seqlen_k_ptr); |
| 67 | + std::printf(" batch : %d\n", static_cast<int>(a.batch)); |
| 68 | + std::printf(" nhead_q : %d\n", static_cast<int>(a.nhead_q)); |
| 69 | + std::printf(" max_seqlen_q : %d\n", static_cast<int>(a.max_seqlen_q)); |
| 70 | + std::printf("========================================\n"); |
| 71 | + std::fflush(stdout); |
| 72 | +} |
| 73 | +*/ |
| 74 | +} // anonymous namespace |
9 | 75 |
|
10 | 76 | namespace pytorch_flash { |
11 | 77 |
|
| 78 | +// SFINAE trait to detect block-scale quantization fields in fmha_fwd_args. |
| 79 | +// Newer versions of composable_kernel add these fields; older versions don't. |
| 80 | +// This lets the PyTorch integration layer work with either version. |
| 81 | +template <typename T, typename = void> |
| 82 | +struct has_fmha_block_scale_fields : std::false_type {}; |
| 83 | + |
| 84 | +template <typename T> |
| 85 | +struct has_fmha_block_scale_fields<T, |
| 86 | + std::void_t<decltype(std::declval<T&>().block_scale_size_q)>> : std::true_type {}; |
| 87 | + |
| 88 | +template <typename Args> |
| 89 | +void set_fmha_fwd_block_scale_fields([[maybe_unused]] Args& args) { |
| 90 | + if constexpr (has_fmha_block_scale_fields<Args>::value) { |
| 91 | + args.block_scale_seqstart_q_ptr = nullptr; |
| 92 | + args.block_scale_seqstart_k_ptr = nullptr; |
| 93 | + args.nhead_stride_q_descale = 0; |
| 94 | + args.nhead_stride_k_descale = 0; |
| 95 | + args.nhead_stride_v_descale = 0; |
| 96 | + args.batch_stride_q_descale = 0; |
| 97 | + args.batch_stride_k_descale = 0; |
| 98 | + args.batch_stride_v_descale = 0; |
| 99 | + args.block_scale_size_q = 0; |
| 100 | + args.block_scale_size_kv = 0; |
| 101 | + } |
| 102 | +} |
12 | 103 |
|
13 | 104 | fmha_fwd_traits get_ck_fmha_fwd_traits(const mask_info &mask, |
14 | 105 | std::string dtype, |
@@ -96,61 +187,58 @@ fmha_fwd_args get_ck_fmha_fwd_args(bool has_lse, |
96 | 187 | nhead_stride_bias = a_b.stride(1); |
97 | 188 | batch_stride_bias = a_b.stride(0); |
98 | 189 | } |
99 | | - return fmha_fwd_args{q.data_ptr(), |
100 | | - k.data_ptr(), |
101 | | - v.data_ptr(), |
102 | | - attn_bias_ptr, // bias |
103 | | - nullptr, // q_descale_ptr |
104 | | - nullptr, // k_descale_ptr |
105 | | - nullptr, // v_descale_ptr |
106 | | - has_dropout_randval ? dropout_randval.data_ptr() : nullptr, |
107 | | - has_lse ? softmax_lse.data_ptr() : nullptr, |
108 | | - out.data_ptr(), |
109 | | - nullptr, // seqstart_q |
110 | | - nullptr, // seqstart_k |
111 | | - nullptr, // seqlen_q_ptr |
112 | | - nullptr, // seqlen_k_ptr |
113 | | - nullptr, // cu_seqlen_q_ptr |
114 | | - nullptr, // cu_seqlen_k_ptr |
115 | | - nullptr, // sink_ptr |
116 | | - seqlen_q, |
117 | | - seqlen_k, |
118 | | - b, |
119 | | - seqlen_q, // max_seqlen_q |
120 | | - d, // hdim_q |
121 | | - d, // hdim_v |
122 | | - h, // nhead |
123 | | - h_k, // nhead_k |
124 | | - softmax_scale, // scale_s |
125 | | - 0.0f, // logits_soft_cap |
126 | | - stride_q, |
127 | | - stride_k, |
128 | | - stride_v, |
129 | | - stride_attn_bias, |
130 | | - stride_randval, |
131 | | - stride_o, |
132 | | - nhead_stride_q, |
133 | | - nhead_stride_k, |
134 | | - nhead_stride_v, |
135 | | - nhead_stride_bias, // nhead_stride_bias |
136 | | - nhead_stride_randval, |
137 | | - nhead_stride_lse, |
138 | | - nhead_stride_o, |
139 | | - batch_stride_q, |
140 | | - batch_stride_k, |
141 | | - batch_stride_v, |
142 | | - batch_stride_bias, // batch_stride_bias |
143 | | - batch_stride_randval, |
144 | | - batch_stride_lse, |
145 | | - batch_stride_o, |
146 | | - mask.left, |
147 | | - mask.right, |
148 | | - 0, // sink_size |
149 | | - static_cast<ck_tile::index_t>(mask.type), |
150 | | - -1, // min_seqlen_q |
151 | | - p_dropout, |
152 | | - has_dropout_randval, |
153 | | - drop_seed_offset}; |
| 190 | + fmha_fwd_args args{}; |
| 191 | + args.q_ptr = q.data_ptr(); |
| 192 | + args.k_ptr = k.data_ptr(); |
| 193 | + args.v_ptr = v.data_ptr(); |
| 194 | + args.bias_ptr = attn_bias_ptr; |
| 195 | + args.q_descale_ptr = nullptr; |
| 196 | + args.k_descale_ptr = nullptr; |
| 197 | + args.v_descale_ptr = nullptr; |
| 198 | + args.rand_val_ptr = has_dropout_randval ? dropout_randval.data_ptr() : nullptr; |
| 199 | + args.lse_ptr = has_lse ? softmax_lse.data_ptr() : nullptr; |
| 200 | + args.o_ptr = out.data_ptr(); |
| 201 | + args.sink_ptr = nullptr; |
| 202 | + args.seqlen_q = seqlen_q; |
| 203 | + args.seqlen_k = seqlen_k; |
| 204 | + args.batch = b; |
| 205 | + args.max_seqlen_q = seqlen_q; |
| 206 | + args.hdim_q = d; |
| 207 | + args.hdim_v = d; |
| 208 | + args.nhead_q = h; |
| 209 | + args.nhead_k = h_k; |
| 210 | + args.scale_s = softmax_scale; |
| 211 | + args.logits_soft_cap = 0.0f; |
| 212 | + args.stride_q = stride_q; |
| 213 | + args.stride_k = stride_k; |
| 214 | + args.stride_v = stride_v; |
| 215 | + args.stride_bias = stride_attn_bias; |
| 216 | + args.stride_randval = stride_randval; |
| 217 | + args.stride_o = stride_o; |
| 218 | + args.nhead_stride_q = nhead_stride_q; |
| 219 | + args.nhead_stride_k = nhead_stride_k; |
| 220 | + args.nhead_stride_v = nhead_stride_v; |
| 221 | + args.nhead_stride_bias = nhead_stride_bias; |
| 222 | + args.nhead_stride_randval = nhead_stride_randval; |
| 223 | + args.nhead_stride_lse = nhead_stride_lse; |
| 224 | + args.nhead_stride_o = nhead_stride_o; |
| 225 | + args.batch_stride_q = batch_stride_q; |
| 226 | + args.batch_stride_k = batch_stride_k; |
| 227 | + args.batch_stride_v = batch_stride_v; |
| 228 | + args.batch_stride_bias = batch_stride_bias; |
| 229 | + args.batch_stride_randval = batch_stride_randval; |
| 230 | + args.batch_stride_lse = batch_stride_lse; |
| 231 | + args.batch_stride_o = batch_stride_o; |
| 232 | + args.window_size_left = mask.left; |
| 233 | + args.window_size_right = mask.right; |
| 234 | + args.sink_size = 0; |
| 235 | + args.mask_type = static_cast<ck_tile::index_t>(mask.type); |
| 236 | + args.min_seqlen_q = -1; |
| 237 | + args.p_drop = p_dropout; |
| 238 | + args.s_randval = has_dropout_randval; |
| 239 | + args.drop_seed_offset = drop_seed_offset; |
| 240 | + set_fmha_fwd_block_scale_fields(args); |
| 241 | + return args; |
154 | 242 | } |
155 | 243 |
|
156 | 244 | std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> |
|
0 commit comments