Merged
feat(quant): support packed row-major sf stores #18
wangliweiAdd on Sep 24
feat(quant): support packed row-major sf stores #18
Merged
Total: 2 files changed+166-73
| @@ -69,11 +69,12 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 69 | assert not in_config.use_tma_aligned_col_major_sf or in_config.use_packed_ue8m0 | 69 | assert not in_config.use_tma_aligned_col_major_sf or in_config.use_packed_ue8m0 |
| 70 | assert not out_config.use_tma_aligned_col_major_sf or out_config.use_packed_ue8m0 | 70 | assert not out_config.use_tma_aligned_col_major_sf or out_config.use_packed_ue8m0 |
| 71 | assert hidden % out_config.sf_block[1] == 0 | 71 | assert hidden % out_config.sf_block[1] == 0 |
| 72 | - # TMA (col-major) output sf is only supported for the block-level (32, 32) layout: | 72 | + is_tma_packed = out_config.use_tma_aligned_col_major_sf and out_config.use_packed_ue8m0 |
| 73 | - # per-data-row sf (1, 128) and the single-sf-per-block (128, 128) case cannot be | 73 | + # TMA+packed (col-major UE8M0) packs two out-sf values into one uint16, so each |
| 74 | - # packed into UE8M0 within one block (tests filter those combinations on Ascend). | 74 | + # persistent task must span 2*out_sf_k columns. For out_sf_k==128: per-row uses |
| 75 | - assert not out_config.use_tma_aligned_col_major_sf or out_config.sf_block == (32, 32) | 75 | + # block_k=256 (row-local or-shift pack); (128, 128) keeps block_k=128 for the UB |
| 76 | - | 76 | + # budget and iterates num_k_subs sub-blocks, staging the first sf in a padding |
| 77 | + # row and or-shifting it with the second. | ||
| 77 | out_sf_m, out_sf_k = out_config.sf_block | 78 | out_sf_m, out_sf_k = out_config.sf_block |
| 78 | if out_sf_m == 1: | 79 | if out_sf_m == 1: |
| 79 | # (1, 128): every data row produces its own out-sf. block_m=64 halves the | 80 | # (1, 128): every data row produces its own out-sf. block_m=64 halves the |
| @@ -85,14 +86,29 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 85 | else: | 86 | else: |
| 86 | block_m = out_sf_m | 87 | block_m = out_sf_m |
| 87 | if out_sf_k == 128: | 88 | if out_sf_k == 128: |
| 88 | - # Output sf covers 128 columns; fix block_k at 128 so the 4-slot column | 89 | + # Output sf covers 128 columns; block_k stays 128 so the 4-slot merge is a |
| 89 | - # merge reduces to a single PAT_ALL group (verified) and every block | 90 | + # single PAT_ALL group (verified) and (128, 128) stays in UB budget. |
| 90 | - # produces exactly one out-sf column. block_k = 128 also keeps block_m = | 91 | + if is_tma_packed: |
| 91 | - # 128 (i.e. (128, 128)) inside the shared-memory budget. | 92 | + if out_sf_m == 1: |
| 92 | - block_k = 128 | 93 | + # per-row: cover 256 columns in one task -> two sf packed in-row. |
| 94 | + block_k = 256 if hidden % 256 == 0 else 128 | ||
| 95 | + num_k_subs = 1 | ||
| 96 | + else: | ||
| 97 | + # (128, 128): 2 sub-blocks of 128 columns span 256 columns. | ||
| 98 | + block_k = 128 | ||
| 99 | + num_k_subs = 2 if hidden % 256 == 0 else 1 | ||
| 100 | + else: | ||
| 101 | + block_k = 128 | ||
| 102 | + num_k_subs = 1 | ||
| 93 | else: | 103 | else: |
| 94 | block_k = math.gcd(2048, hidden) | 104 | block_k = math.gcd(2048, hidden) |
| 95 | - assert block_k % out_sf_k == 0 and hidden % block_k == 0 | 105 | + num_k_subs = 1 |
| 106 | + k_span = block_k * num_k_subs | ||
| 107 | + sf_per_pid = k_span // out_sf_k | ||
| 108 | + # Packed UE8M0 needs an even number of sf values per persistent task to form a | ||
| 109 | + # uint16; a task producing exactly one sf stores (sf | 0). | ||
| 110 | + assert not is_tma_packed or sf_per_pid % get_packed_ue8m0_pack_factor() == 0 or sf_per_pid == 1 | ||
| 111 | + assert block_k % out_sf_k == 0 and hidden % k_span == 0 | ||
| 96 | if block_k <= 1024 and block_m <= 64: | 112 | if block_k <= 1024 and block_m <= 64: |
| 97 | num_stages = 3 | 113 | num_stages = 3 |
| 98 | else: | 114 | else: |
| @@ -120,7 +136,18 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 120 | sf_load_ub_shape = (block_m, vec_size * in_pack_factor) | 136 | sf_load_ub_shape = (block_m, vec_size * in_pack_factor) |
| 121 | if out_config.use_tma_aligned_col_major_sf: | 137 | if out_config.use_tma_aligned_col_major_sf: |
| 122 | num_batched_tiles = 2 | 138 | num_batched_tiles = 2 |
| 123 | - out_sf_shape_ub = (max(num_sf_store_cols, 4 * out_pack_factor), num_batched_tiles * out_pack_factor) | 139 | + if is_tma_packed: |
| 140 | + if out_sf_m == 1: | ||
| 141 | + # per-row: flat staging buffer, slot = row*num_batched_tiles + batch; | ||
| 142 | + # padded to 256 so a NORM_B16 vsts footprint (slot + 128) stays in bounds. | ||
| 143 | + out_sf_shape_ub = (max(block_m * num_batched_tiles, 256),) | ||
| 144 | + else: | ||
| 145 | + # block: (num_sf_store_cols, num_batched_tiles) padded to >= 128 rows | ||
| 146 | + # (NORM_B16 vsts footprint); one extra row stages the first sf of a | ||
| 147 | + # (128,128) 2-sub-block pair. | ||
| 148 | + out_sf_shape_ub = (max(num_sf_store_cols + (1 if num_k_subs == 2 else 0), 128), num_batched_tiles) | ||
| 149 | + else: | ||
| 150 | + out_sf_shape_ub = (max(num_sf_store_cols, 4 * out_pack_factor), num_batched_tiles * out_pack_factor) | ||
| 124 | else: | 151 | else: |
| 125 | num_batched_tiles = 1 | 152 | num_batched_tiles = 1 |
| 126 | out_sf_shape_ub = (num_out_sf_rows, max(num_sf_store_cols, 8 * out_pack_factor) * out_pack_factor) | 153 | out_sf_shape_ub = (num_out_sf_rows, max(num_sf_store_cols, 8 * out_pack_factor) * out_pack_factor) |
| @@ -154,23 +181,37 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 154 | return S.vpack(S.vshrs(T.reinterpret(values, 'uint32x64'), 23)) | 181 | return S.vpack(S.vshrs(T.reinterpret(values, 'uint32x64'), 23)) |
| 155 | 182 | ||
| 156 | @T.macro | 183 | @T.macro |
| 157 | - def copy_sf(x_sf, sf_load_ub, row_start, pid_k): | 184 | + def copy_sf(x_sf, sf_load_ub, row_start, pid_k, sub=0): |
| 158 | if in_config.use_tma_aligned_col_major_sf: | 185 | if in_config.use_tma_aligned_col_major_sf: |
| 159 | - col_start = T.alloc_var(T.int32, init=pid_k * num_sf_load_cols) | 186 | + col_start = T.alloc_var(T.int32, init=pid_k * num_sf_load_cols * num_k_subs + sub * num_sf_load_cols) |
| 160 | T.copy( | 187 | T.copy( |
| 161 | x_sf[col_start, row_start * in_pack_factor], sf_load_ub[:num_sf_load_cols, : block_m * in_pack_factor] | 188 | x_sf[col_start, row_start * in_pack_factor], sf_load_ub[:num_sf_load_cols, : block_m * in_pack_factor] |
| 162 | ) | 189 | ) |
| 163 | else: | 190 | else: |
| 164 | - col_start = T.alloc_var(T.int32, init=pid_k * num_sf_cols_per_block) | 191 | + col_start = T.alloc_var( |
| 192 | + T.int32, init=pid_k * num_sf_cols_per_block * num_k_subs + sub * num_sf_cols_per_block | ||
| 193 | + ) | ||
| 165 | T.copy(x_sf[row_start, col_start], sf_load_ub[:block_m, :num_sf_cols_per_block]) | 194 | T.copy(x_sf[row_start, col_start], sf_load_ub[:block_m, :num_sf_cols_per_block]) |
| 166 | 195 | ||
| 167 | @T.macro | 196 | @T.macro |
| 168 | - def col_merge(v): | 197 | + def col_merge(v, num_groups): |
| 169 | - # Merge every 4 consecutive 32-column slots into a single 128-column | 198 | + # Merge groups of 4 consecutive 32-column slots into a single 128-column |
| 170 | - # out-sf value. block_k is fixed at 128 for out_sf_k == 128, so the merge | 199 | + # out-sf value. num_groups=1 (block_k=128): single PAT_ALL reduce + broadcast. |
| 171 | - # is a single PAT_ALL group reduce + broadcast (verified semantics). | 200 | + # num_groups=2 (block_k=256): slots 0-3 -> SF0 (lanes 0-3), slots 4-7 -> SF1 |
| 201 | + # (lanes 4+), so the or-shift pack yields (SF1<<8)|SF0 at lane 0. | ||
| 172 | pat_all = S.pset(16, 'PAT_ALL') | 202 | pat_all = S.pset(16, 'PAT_ALL') |
| 173 | - return S.vdupv(S.vcmax(v, pat_all), pat_all) | 203 | + if num_groups == 1: |
| 204 | + return S.vdupv(S.vcmax(v, pat_all), pat_all) | ||
| 205 | + lanes = lane_indices(T.uint16) | ||
| 206 | + zero = S.vdup(0, T.uint16) | ||
| 207 | + mask_lt4 = S.vcmps(lanes, T.uint16(4), op='lt') | ||
| 208 | + idx_shift = S.vsel(S.vadds(lanes, 4), zero, mask_lt4) | ||
| 209 | + v_hi_raw = S.vselr(v, idx_shift) | ||
| 210 | + v_hi = S.vsel(v_hi_raw, zero, mask_lt4) | ||
| 211 | + v_lo = S.vsel(v, zero, mask_lt4) | ||
| 212 | + sf0 = S.vdupv(S.vcmax(v_lo, pat_all), pat_all) | ||
| 213 | + sf1 = S.vdupv(S.vcmax(v_hi, pat_all), pat_all) | ||
| 214 | + return S.vsel(sf0, sf1, mask_lt4) | ||
| 174 | 215 | ||
| 175 | @T.macro | 216 | @T.macro |
| 176 | def emit_out_sf( | 217 | def emit_out_sf( |
| @@ -185,6 +226,7 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 185 | mask, | 226 | mask, |
| 186 | zero, | 227 | zero, |
| 187 | six, | 228 | six, |
| 229 | + sf_sub=0, | ||
| 188 | ): | 230 | ): |
| 189 | full_mask = vector_mask(T.uint16, max(num_sf_cols_per_block, 16)) | 231 | full_mask = vector_mask(T.uint16, max(num_sf_cols_per_block, 16)) |
| 190 | S.vsts(sf_exp_ub[sf_exp_row, 0], S.vadds(out_exp_row, -5), full_mask, dist='NORM_B16') | 232 | S.vsts(sf_exp_ub[sf_exp_row, 0], S.vadds(out_exp_row, -5), full_mask, dist='NORM_B16') |
| @@ -193,6 +235,41 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 193 | acc_invalid = S.vmax(max_invalid, S.vld(invalid_exp_ub[0], dist='BRC_B16')) | 235 | acc_invalid = S.vmax(max_invalid, S.vld(invalid_exp_ub[0], dist='BRC_B16')) |
| 194 | assert_mask = vector_mask(T.uint16, 16) | 236 | assert_mask = vector_mask(T.uint16, 16) |
| 195 | S.vsts(invalid_exp_ub[0], acc_invalid, assert_mask, dist='NORM_B16') | 237 | S.vsts(invalid_exp_ub[0], acc_invalid, assert_mask, dist='NORM_B16') |
| 238 | + one_mask = S.pset(16, 'PAT_VL1') | ||
| 239 | + if is_tma_packed: | ||
| 240 | + if out_sf_k == 128: | ||
| 241 | + if sf_per_pid % out_pack_factor == 0: | ||
| 242 | + if num_k_subs == 2: | ||
| 243 | + # (128, 128): two sub-blocks share one uint16; stage SF0 in a | ||
| 244 | + # padding row, or-shift SF1 on the second sub-block. | ||
| 245 | + exp_lo = S.vand(out_exp_row, S.vdup(T.uint16(0x00FF), T.uint16)) | ||
| 246 | + if sf_sub == 0: | ||
| 247 | + S.vsts(out_sf_ub[num_sf_store_cols, store_row], exp_lo, one_mask, dist='ONEPT_B16') | ||
| 248 | + else: | ||
| 249 | + prev = S.vld(out_sf_ub[num_sf_store_cols, store_row], dist='BRC_B16') | ||
| 250 | + packed = S.vor(prev, S.vshls(S.vdupv(out_exp_row, pos='POS_HIGHEST'), 8)) | ||
| 251 | + S.vsts(out_sf_ub[0, store_row], packed, one_mask, dist='ONEPT_B16') | ||
| 252 | + return | ||
| 253 | + # per-row 256-col: col_merge leaves SF0 in lanes 0-3 and SF1 in | ||
| 254 | + # lanes 4+, so the highest lane carries the second sf's exponent. | ||
| 255 | + exponent_high = S.vdupv(out_exp_row, pos='POS_HIGHEST') | ||
| 256 | + out_sf_value = S.vor(out_exp_row, S.vshls(exponent_high, 8)) | ||
| 257 | + S.vsts(out_sf_ub[store_row * block_m + sf_ub_row], out_sf_value, one_mask, dist='ONEPT_B16') | ||
| 258 | + return | ||
| 259 | + # single sf per task: (sf | 0), trailing byte zero-filled. | ||
| 260 | + out_sf_value = S.vand(out_exp_row, S.vdup(T.uint16(0x00FF), T.uint16)) | ||
| 261 | + if out_sf_m == 1: | ||
| 262 | + S.vsts(out_sf_ub[store_row * block_m + sf_ub_row], out_sf_value, one_mask, dist='ONEPT_B16') | ||
| 263 | + else: | ||
| 264 | + S.vsts(out_sf_ub[0, store_row], out_sf_value, one_mask, dist='ONEPT_B16') | ||
| 265 | + return | ||
| 266 | + # out_sf_k == 32: vpack produces adjacent sf pairs (SF1<<8)|SF0 at | ||
| 267 | + # lanes 0..num_sf_store_cols-1; store each pair via ONEPT_B16. | ||
| 268 | + out_sf_value = T.reinterpret(S.vpack(out_exp_row), 'uint16x128') | ||
| 269 | + for s in range(num_sf_store_cols): | ||
| 270 | + pair = S.vselr(out_sf_value, S.vdup(s, T.uint16)) | ||
| 271 | + S.vsts(out_sf_ub[s, store_row], pair, one_mask, dist='ONEPT_B16') | ||
| 272 | + return | ||
| 196 | out_sf_value = S.alloc_var(out_sf_dtype) | 273 | out_sf_value = S.alloc_var(out_sf_dtype) |
| 197 | if out_sf_dtype == T.uint16: | 274 | if out_sf_dtype == T.uint16: |
| 198 | out_sf_value = T.reinterpret(S.vpack(out_exp_row), 'uint16x128') | 275 | out_sf_value = T.reinterpret(S.vpack(out_exp_row), 'uint16x128') |
| @@ -200,17 +277,13 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 200 | wide_exp, _ = S.vintlv(out_exp_row, zero) | 277 | wide_exp, _ = S.vintlv(out_exp_row, zero) |
| 201 | out_sf_value = T.reinterpret(S.vshls(T.reinterpret(wide_exp, 'uint32x64'), 23), 'float32x64') | 278 | out_sf_value = T.reinterpret(S.vshls(T.reinterpret(wide_exp, 'uint32x64'), 23), 'float32x64') |
| 202 | out_store_mask = vector_mask(out_sf_dtype, num_sf_store_cols) | 279 | out_store_mask = vector_mask(out_sf_dtype, num_sf_store_cols) |
| 203 | - if out_config.use_tma_aligned_col_major_sf: | 280 | + if out_sf_dtype == T.uint16: |
| 204 | - store_indices = lane_indices(out_sf_dtype) | ||
| 205 | - offsets = S.vadds(S.vmuls(store_indices, num_batched_tiles), store_row) | ||
| 206 | - S.vscatter(out_sf_value, out_sf_ub[0, 0], offsets, out_store_mask) | ||
| 207 | - elif out_sf_dtype == T.uint16: | ||
| 208 | S.vsts(out_sf_ub[sf_ub_row, 0], out_sf_value, out_store_mask, dist='NORM_B16') | 281 | S.vsts(out_sf_ub[sf_ub_row, 0], out_sf_value, out_store_mask, dist='NORM_B16') |
| 209 | else: | 282 | else: |
| 210 | S.vsts(out_sf_ub[sf_ub_row, 0], out_sf_value, out_store_mask, dist='NORM_B32') | 283 | S.vsts(out_sf_ub[sf_ub_row, 0], out_sf_value, out_store_mask, dist='NORM_B32') |
| 211 | 284 | ||
| 212 | @T.macro | 285 | @T.macro |
| 213 | - def reduce_sf(sf_load_ub, sf_exp_ub, out_sf_ub, invalid_exp_ub, store_row): | 286 | + def reduce_sf(sf_load_ub, sf_exp_ub, out_sf_ub, invalid_exp_ub, store_row, sf_sub=0): |
| 214 | with T.SimdVF(): | 287 | with T.SimdVF(): |
| 215 | mask = vector_mask(T.uint16, num_sf_cols_per_block) | 288 | mask = vector_mask(T.uint16, num_sf_cols_per_block) |
| 216 | full_mask = vector_mask(T.uint16, max(num_sf_cols_per_block, 16)) | 289 | full_mask = vector_mask(T.uint16, max(num_sf_cols_per_block, 16)) |
| @@ -247,8 +320,8 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 247 | # max_exp/min_exp cross-row reduction is only needed for the | 320 | # max_exp/min_exp cross-row reduction is only needed for the |
| 248 | # block-level (out_sf_m != 1) path, so skip it here to save 4 | 321 | # block-level (out_sf_m != 1) path, so skip it here to save 4 |
| 249 | # vector ops per 2 rows (compile-time branch). | 322 | # vector ops per 2 rows (compile-time branch). |
| 250 | - out_exp0 = col_merge(S.vsub(S.vmax(in_exp0_clean, six), six)) | 323 | + out_exp0 = col_merge(S.vsub(S.vmax(in_exp0_clean, six), six), block_k // 128) |
| 251 | - out_exp1 = col_merge(S.vsub(S.vmax(in_exp1_clean, six), six)) | 324 | + out_exp1 = col_merge(S.vsub(S.vmax(in_exp1_clean, six), six), block_k // 128) |
| 252 | emit_out_sf( | 325 | emit_out_sf( |
| 253 | sf_exp_ub, | 326 | sf_exp_ub, |
| 254 | out_sf_ub, | 327 | out_sf_ub, |
| @@ -261,6 +334,7 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 261 | mask, | 334 | mask, |
| 262 | zero, | 335 | zero, |
| 263 | six, | 336 | six, |
| 337 | + sf_sub, | ||
| 264 | ) | 338 | ) |
| 265 | emit_out_sf( | 339 | emit_out_sf( |
| 266 | sf_exp_ub, | 340 | sf_exp_ub, |
| @@ -274,6 +348,7 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 274 | mask, | 348 | mask, |
| 275 | zero, | 349 | zero, |
| 276 | six, | 350 | six, |
| 351 | + sf_sub, | ||
| 277 | ) | 352 | ) |
| 278 | else: | 353 | else: |
| 279 | max_exp0 = S.vmax(max_exp0, in_exp0_clean, mask) | 354 | max_exp0 = S.vmax(max_exp0, in_exp0_clean, mask) |
| @@ -286,9 +361,20 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 286 | out_exp = S.alloc_var(T.uint16) | 361 | out_exp = S.alloc_var(T.uint16) |
| 287 | out_exp = S.vsub(S.vmax(max_exp, six), six) | 362 | out_exp = S.vsub(S.vmax(max_exp, six), six) |
| 288 | if out_sf_k == 128: | 363 | if out_sf_k == 128: |
| 289 | - out_exp = col_merge(out_exp) | 364 | + out_exp = col_merge(out_exp, 1) |
| 290 | emit_out_sf( | 365 | emit_out_sf( |
| 291 | - sf_exp_ub, out_sf_ub, invalid_exp_ub, out_exp, min_exp, block_m, 0, store_row, mask, zero, six | 366 | + sf_exp_ub, |
| 367 | + out_sf_ub, | ||
| 368 | + invalid_exp_ub, | ||
| 369 | + out_exp, | ||
| 370 | + min_exp, | ||
| 371 | + block_m, | ||
| 372 | + 0, | ||
| 373 | + store_row, | ||
| 374 | + mask, | ||
| 375 | + zero, | ||
| 376 | + six, | ||
| 377 | + sf_sub, | ||
| 292 | ) | 378 | ) |
| 293 | 379 | ||
| 294 | @T.macro | 380 | @T.macro |
| @@ -347,11 +433,17 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 347 | ): | 433 | ): |
| 348 | with T.Kernel(num_cores) as core_id: | 434 | with T.Kernel(num_cores) as core_id: |
| 349 | x_ub = T.alloc_shared((block_m, quant_cols), in_config.dtype) | 435 | x_ub = T.alloc_shared((block_m, quant_cols), in_config.dtype) |
| 436 | + # tilelang-new's PlanMTECopy requires sub-byte (4-bit) rows to be | ||
| 437 | + # byte-aligned; x's physical storage is packed uint8 (2 fp4 per byte). | ||
| 438 | + # Copying through a uint8 byte view keeps the same layout while making | ||
| 439 | + # row_bits = n*8 provably a multiple of 8, bypassing that check. | ||
| 440 | + x_bytes = T.Tensor((num_tokens, hidden // 2), T.uint8, x.data) | ||
| 441 | + x_ub_bytes = T.Tensor((block_m, quant_cols // 2), T.uint8, x_ub.data) | ||
| 350 | sf_exp_ub = T.alloc_shared((block_m + num_out_sf_rows, max(num_sf_cols_per_block, 16)), T.uint16) | 442 | sf_exp_ub = T.alloc_shared((block_m + num_out_sf_rows, max(num_sf_cols_per_block, 16)), T.uint16) |
| 351 | out_ub = T.alloc_shared((block_m, quant_cols), out_config.dtype) | 443 | out_ub = T.alloc_shared((block_m, quant_cols), out_config.dtype) |
| 352 | fp4_to_fp8_map_ub = T.alloc_shared((256,), T.uint8) | 444 | fp4_to_fp8_map_ub = T.alloc_shared((256,), T.uint8) |
| 353 | sf_load_ub = T.alloc_shared(sf_load_ub_shape, in_config.sf_dtype) | 445 | sf_load_ub = T.alloc_shared(sf_load_ub_shape, in_config.sf_dtype) |
| 354 | - out_sf_ub = T.alloc_shared(out_sf_shape_ub, out_config.sf_dtype) | 446 | + out_sf_ub = T.alloc_shared(out_sf_shape_ub, T.uint16 if is_tma_packed else out_config.sf_dtype) |
| 355 | invalid_exp_ub = T.alloc_shared((16,), T.uint16) | 447 | invalid_exp_ub = T.alloc_shared((16,), T.uint16) |
| 356 | T.annotate_buffer_versions( | 448 | T.annotate_buffer_versions( |
| 357 | { | 449 | { |
| @@ -365,7 +457,7 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 365 | ) | 457 | ) |
| 366 | init_fp4_to_fp8_map(fp4_to_fp8_map_ub, invalid_exp_ub) | 458 | init_fp4_to_fp8_map(fp4_to_fp8_map_ub, invalid_exp_ub) |
| 367 | for pid_m, pid_k in T.Persistent( | 459 | for pid_m, pid_k in T.Persistent( |
| 368 | - [num_tokens // (block_m * num_batched_tiles), hidden // block_k], | 460 | + [T.ceildiv(num_tokens, block_m * num_batched_tiles), hidden // k_span], |
| 369 | num_cores, | 461 | num_cores, |
| 370 | core_id, | 462 | core_id, |
| 371 | group_size=1, | 463 | group_size=1, |
| @@ -374,26 +466,51 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 374 | for batch in T.serial(num_batched_tiles): | 466 | for batch in T.serial(num_batched_tiles): |
| 375 | tile_m = T.alloc_var(T.int32, init=pid_m * num_batched_tiles + batch) | 467 | tile_m = T.alloc_var(T.int32, init=pid_m * num_batched_tiles + batch) |
| 376 | row_start = T.alloc_var(T.int32, init=tile_m * block_m) | 468 | row_start = T.alloc_var(T.int32, init=tile_m * block_m) |
| 377 | - if block_k < 256: | 469 | + if row_start < num_tokens: |
| 378 | - # Zero-fill the padded 256-column tile so the padded FP4 bytes | 470 | + if block_k < 256: |
| 379 | - # decode to code 0 (legal lookup) during quant. S.vsts writes | 471 | + # Zero-fill the padded 256-column tile so the padded FP4 bytes |
| 380 | - # 256 bytes (two 128-byte rows) per call. | 472 | + # decode to code 0 (legal lookup) during quant. S.vsts writes |
| 381 | - with T.SimdVF(): | 473 | + # 256 bytes (two 128-byte rows) per call. |
| 382 | - zero8 = S.vdup(0, T.uint8) | 474 | + with T.SimdVF(): |
| 383 | - for pad_row in T.serial(0, block_m, 2): | 475 | + zero8 = S.vdup(0, T.uint8) |
| 384 | - S.vsts(x_ub[pad_row, 0], zero8, dist='NORM_B8') | 476 | + for pad_row in T.serial(0, block_m, 2): |
| 385 | - T.copy(x[row_start, pid_k * block_k], x_ub[:block_m, :block_k]) | 477 | + S.vsts(x_ub[pad_row, 0], zero8, dist='NORM_B8') |
| 386 | - copy_sf(x_sf, sf_load_ub, row_start, pid_k) | 478 | + for sub in range(num_k_subs): |
| 387 | - reduce_sf(sf_load_ub, sf_exp_ub, out_sf_ub, invalid_exp_ub, batch) | 479 | + k_col = T.alloc_var(T.int32, init=pid_k * k_span + sub * block_k) |
| 388 | - quant(sf_exp_ub, x_ub, out_ub, fp4_to_fp8_map_ub) | 480 | + T.copy(x_bytes[row_start, k_col // 2], x_ub_bytes[:block_m, : block_k // 2]) |
| 389 | - T.copy(out_ub[:block_m, :block_k], out[row_start, pid_k * block_k]) | 481 | + copy_sf(x_sf, sf_load_ub, row_start, pid_k, sub) |
| 482 | + reduce_sf(sf_load_ub, sf_exp_ub, out_sf_ub, invalid_exp_ub, batch, sf_sub=sub) | ||
| 483 | + quant(sf_exp_ub, x_ub, out_ub, fp4_to_fp8_map_ub) | ||
| 484 | + T.copy(out_ub[:block_m, :block_k], out[row_start, k_col]) | ||
| 390 | first_tile = T.alloc_var(T.int32, init=pid_m * num_batched_tiles) | 485 | first_tile = T.alloc_var(T.int32, init=pid_m * num_batched_tiles) |
| 391 | store_base = T.alloc_var(T.int32, init=pid_k * num_sf_store_cols) | 486 | store_base = T.alloc_var(T.int32, init=pid_k * num_sf_store_cols) |
| 392 | if out_config.use_tma_aligned_col_major_sf: | 487 | if out_config.use_tma_aligned_col_major_sf: |
| 393 | - T.copy( | 488 | + if is_tma_packed: |
| 394 | - out_sf_ub[:num_sf_store_cols, : num_batched_tiles * out_pack_factor], | 489 | + # out_sf is physically uint8 (packed UE8M0); view it as uint16 |
| 395 | - out_sf[store_base, first_tile * out_pack_factor], | 490 | + # so the staging copy is uint16 -> uint16. |
| 396 | - ) | 491 | + output_sf = T.Tensor( |
| 492 | + (out_sf_shape[0], out_sf_shape[1] // out_pack_factor), T.uint16, out_sf.data | ||
| 493 | + ) | ||
| 494 | + if out_sf_m == 1: | ||
| 495 | + # per-row: flat (N,) staging copied into the k-slot row. | ||
| 496 | + T.copy( | ||
| 497 | + out_sf_ub[0 : block_m * num_batched_tiles], | ||
| 498 | + output_sf[ | ||
| 499 | + pid_k * num_sf_store_cols, | ||
| 500 | + first_tile * block_m : first_tile * block_m + block_m * num_batched_tiles, | ||
| 501 | + ], | ||
| 502 | + ) | ||
| 503 | + else: | ||
| 504 | + # block: phys[k_slot, token_block]. | ||
| 505 | + T.copy( | ||
| 506 | + out_sf_ub[:num_sf_store_cols, :num_batched_tiles], | ||
| 507 | + output_sf[pid_k * num_sf_store_cols, first_tile], | ||
| 508 | + ) | ||
| 509 | + else: | ||
| 510 | + T.copy( | ||
| 511 | + out_sf_ub[:num_sf_store_cols, : num_batched_tiles * out_pack_factor], | ||
| 512 | + out_sf[store_base, first_tile * out_pack_factor], | ||
| 513 | + ) | ||
| 397 | else: | 514 | else: |
| 398 | # Both buffers expose packed UE8M0 as uint8, so copy and index by | 515 | # Both buffers expose packed UE8M0 as uint8, so copy and index by |
| 399 | # logical output-sf columns. This also handles a one-sf tile without | 516 | # logical output-sf columns. This also handles a one-sf tile without |
| @@ -4,7 +4,6 @@ import pytest | |||
| 4 | import torch | 4 | import torch |
| 5 | 5 | ||
| 6 | import cann_ops_tilelang | 6 | import cann_ops_tilelang |
| 7 | -from cann_ops_tilelang.config import get_device | ||
| 8 | from tests.quant.cast_ref import cast, cast_back | 7 | from tests.quant.cast_ref import cast, cast_back |
| 9 | from tests.testing.bench import make_param_id | 8 | from tests.testing.bench import make_param_id |
| 10 | from tests.testing.generator import ( | 9 | from tests.testing.generator import ( |
| @@ -97,7 +96,6 @@ def generate_test_params(level: int) -> list[dict]: | |||
| 97 | for out_sf_block_m, out_sf_block_k in out_sf_blocks | 96 | for out_sf_block_m, out_sf_block_k in out_sf_blocks |
| 98 | for in_sf_block_m, in_sf_block_k in ((1, 32),) | 97 | for in_sf_block_m, in_sf_block_k in ((1, 32),) |
| 99 | if hidden_size % out_sf_block_k == 0 | 98 | if hidden_size % out_sf_block_k == 0 |
| 100 | - if not (out_sf_block_m != 32 and out_use_tma_aligned_col_major_sf) | ||
| 101 | ] | 99 | ] |
| 102 | return params | 100 | return params |
| 103 | 101 | ||
| @@ -119,28 +117,6 @@ def test_per_block_cast_lossless(params): | |||
| 119 | assert_equal(x_fp8_fp32, x_fp8_fp32_ref) | 117 | assert_equal(x_fp8_fp32, x_fp8_fp32_ref) |
| 120 | 118 | ||
| 121 | 119 | ||
| 122 | -def test_per_block_cast_lossless_row_major_packed_ue8m0(): | ||
| 123 | - x = torch.empty((128, 256), dtype=torch.bfloat16, device=get_device()) | ||
| 124 | - x[:, :128] = 1.0 | ||
| 125 | - x[:, 128:] = 32.0 | ||
| 126 | - x_fp4 = cast(x, 'e2m1', (1, 32), round_sf=True) | ||
| 127 | - | ||
| 128 | - x_fp8 = cann_ops_tilelang.per_block_cast_lossless( | ||
| 129 | - x_fp4, | ||
| 130 | - 'e4m3', | ||
| 131 | - x_block_size=(1, 32), | ||
| 132 | - out_block_size=(128, 128), | ||
| 133 | - round_sf=True, | ||
| 134 | - use_packed_ue8m0=True, | ||
| 135 | - ) | ||
| 136 | - x_fp8_ref = cast(x, 'e4m3', (128, 128), round_sf=True, use_packed_ue8m0=True) | ||
| 137 | - | ||
| 138 | - x_fp8_fp32_ref = cast_back(x_fp4, 'fp32', (1, 32)) | ||
| 139 | - x_fp8_fp32 = cast_back(x_fp8, 'fp32', (128, 128)) | ||
| 140 | - assert torch.equal(x_fp8[1], x_fp8_ref[1]) | ||
| 141 | - assert_equal(x_fp8_fp32, x_fp8_fp32_ref) | ||
| 142 | - | ||
| 143 | - | ||
| 144 | @pytest.mark.benchmark | 120 | @pytest.mark.benchmark |
| 145 | @pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id) | 121 | @pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id) |
| 146 | def test_per_block_cast_lossless_benchmark(benchmark_timer, benchmark_record, params): | 122 | def test_per_block_cast_lossless_benchmark(benchmark_timer, benchmark_record, params): |