Open
feat(quant): support torch.compile for SwiGLU operators #72
daisysjAdd on 16 hours ago
feat(quant): support torch.compile for SwiGLU operators #72
Open
Total: 22 files changed+1347-85
Msrc/cann_ops_tilelang/quant/cast_back_kernel.py+2-3File content is under review, please try again later
| @@ -1,17 +1,38 @@ | |||
| 1 | import torch | 1 | import torch |
| 2 | 2 | ||
| 3 | -from cann_ops_tilelang.config import require_ascend, should_print_kernel_source | 3 | +from cann_ops_tilelang.config import get_amax_clamp_for_quant, require_ascend, should_print_kernel_source |
| 4 | from cann_ops_tilelang.quant.common import ( | 4 | from cann_ops_tilelang.quant.common import ( |
| 5 | QuantTensor, | 5 | QuantTensor, |
| 6 | alloc_scaling_factors, | 6 | alloc_scaling_factors, |
| 7 | cast_epilogue, | 7 | cast_epilogue, |
| 8 | get_cast_input_and_config, | 8 | get_cast_input_and_config, |
| 9 | - get_cast_output_config, | 9 | + get_cast_output_config_impl, |
| 10 | + get_quant_vector_cores, | ||
| 10 | ) | 11 | ) |
| 11 | 12 | ||
| 12 | from .per_block_cast_kernel import get_per_block_cast_kernel_asc | 13 | from .per_block_cast_kernel import get_per_block_cast_kernel_asc |
| 13 | 14 | ||
| 14 | 15 | ||
| 16 | +def _validate_per_block_cast(x, fmt, block_size, options, clamp_min_value, sf=None, sf_only=False): | ||
| 17 | + assert x.device.type == 'npu' | ||
| 18 | + assert x.is_contiguous() and x.dim() == 2 | ||
| 19 | + num_tokens, hidden = x.shape | ||
| 20 | + | ||
| 21 | + assert num_tokens % 128 == 0 | ||
| 22 | + assert hidden % 128 == 0 | ||
| 23 | + _validate_per_block_scale(sf, sf_only, options) | ||
| 24 | + | ||
| 25 | + x, _, in_config = get_cast_input_and_config(x, None) | ||
| 26 | + out_config = get_cast_output_config_impl(fmt, block_size, *options, clamp_min_value) | ||
| 27 | + return in_config, out_config | ||
| 28 | + | ||
| 29 | + | ||
| 30 | +def _validate_per_block_scale(sf, sf_only, options): | ||
| 31 | + use_tma_aligned_col_major_sf, _, use_packed_ue8m0 = options | ||
| 32 | + if sf is not None: | ||
| 33 | + assert not sf_only and not use_tma_aligned_col_major_sf and not use_packed_ue8m0 | ||
| 34 | + | ||
| 35 | + | ||
| 15 | def per_block_cast_impl( | 36 | def per_block_cast_impl( |
| 16 | x: torch.Tensor, | 37 | x: torch.Tensor, |
| 17 | fmt: str, | 38 | fmt: str, |
| @@ -21,6 +42,9 @@ def per_block_cast_impl( | |||
| 21 | use_packed_ue8m0: bool = False, | 42 | use_packed_ue8m0: bool = False, |
| 22 | sf_only: bool = False, | 43 | sf_only: bool = False, |
| 23 | sf: torch.Tensor | None = None, | 44 | sf: torch.Tensor | None = None, |
| 45 | + *, | ||
| 46 | + clamp_min_value: float, | ||
| 47 | + num_cores: int, | ||
| 24 | ) -> torch.Tensor | QuantTensor: | 48 | ) -> torch.Tensor | QuantTensor: |
| 25 | """ | 49 | """ |
| 26 | sf_only: If True, only compute and return sf factors. | 50 | sf_only: If True, only compute and return sf factors. |
| @@ -31,24 +55,19 @@ def per_block_cast_impl( | |||
| 31 | is provided, or a tuple ``(out, out_sf)`` with quantized output and sf factors. | 55 | is provided, or a tuple ``(out, out_sf)`` with quantized output and sf factors. |
| 32 | """ | 56 | """ |
| 33 | require_ascend() | 57 | require_ascend() |
| 34 | - assert x.device.type == 'npu' | 58 | + options = [use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0] |
| 35 | - assert x.is_contiguous() and x.dim() == 2 | 59 | + in_config, out_config = _validate_per_block_cast( |
| 60 | + x, fmt, block_size, options, clamp_min_value, sf=sf, sf_only=sf_only | ||
| 61 | + ) | ||
| 36 | num_tokens, hidden = x.shape | 62 | num_tokens, hidden = x.shape |
| 37 | 63 | ||
| 38 | - assert num_tokens % 128 == 0 | ||
| 39 | - assert hidden % 128 == 0 | ||
| 40 | - if sf is not None: | ||
| 41 | - assert not sf_only and not use_tma_aligned_col_major_sf and not use_packed_ue8m0 | ||
| 42 | - | ||
| 43 | - x, _, in_config = get_cast_input_and_config(x, None) | ||
| 44 | - out_config = get_cast_output_config(fmt, block_size, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0) | ||
| 45 | - | ||
| 46 | kernel = get_per_block_cast_kernel_asc( | 64 | kernel = get_per_block_cast_kernel_asc( |
| 47 | hidden=hidden, | 65 | hidden=hidden, |
| 48 | in_config=in_config, | 66 | in_config=in_config, |
| 49 | out_config=out_config, | 67 | out_config=out_config, |
| 50 | sf_only=sf_only, | 68 | sf_only=sf_only, |
| 51 | cast_only=(sf is not None), | 69 | cast_only=(sf is not None), |
| 70 | + n_cores=num_cores, | ||
| 52 | ) | 71 | ) |
| 53 | 72 | ||
| 54 | if should_print_kernel_source(): | 73 | if should_print_kernel_source(): |
| @@ -102,7 +121,14 @@ def per_block_cast( | |||
| 102 | Returns: | 121 | Returns: |
| 103 | Tuple of quantized output and sf factors. | 122 | Tuple of quantized output and sf factors. |
| 104 | """ | 123 | """ |
| 105 | - return per_block_cast_impl(x, fmt, block_size, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0) | 124 | + return _per_block_cast_op( |
| 125 | + x, | ||
| 126 | + fmt, | ||
| 127 | + list(block_size), | ||
| 128 | + [use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0], | ||
| 129 | + get_amax_clamp_for_quant(fmt), | ||
| 130 | + get_quant_vector_cores(), | ||
| 131 | + ) | ||
| 106 | 132 | ||
| 107 | 133 | ||
| 108 | def per_block_cast_with_sf_only( | 134 | def per_block_cast_with_sf_only( |
| @@ -126,8 +152,13 @@ def per_block_cast_with_sf_only( | |||
| 126 | Returns: | 152 | Returns: |
| 127 | Scaling factors only | 153 | Scaling factors only |
| 128 | """ | 154 | """ |
| 129 | - return per_block_cast_impl( | 155 | + return _per_block_cast_op_with_sf_only( |
| 130 | - x, fmt, block_size, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0, sf_only=True | 156 | + x, |
| 157 | + fmt, | ||
| 158 | + list(block_size), | ||
| 159 | + [use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0], | ||
| 160 | + get_amax_clamp_for_quant(fmt), | ||
| 161 | + get_quant_vector_cores(), | ||
| 131 | ) | 162 | ) |
| 132 | 163 | ||
| 133 | 164 | ||
| @@ -154,4 +185,94 @@ def per_block_cast_with_precomputed_sf( | |||
| 154 | Returns: | 185 | Returns: |
| 155 | Casted tensor using the precomputed sf factors. | 186 | Casted tensor using the precomputed sf factors. |
| 156 | """ | 187 | """ |
| 157 | - return per_block_cast_impl(x, fmt, block_size, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0, sf=sf) | 188 | + return _per_block_cast_op_with_precomputed_sf( |
| 189 | + x, | ||
| 190 | + fmt, | ||
| 191 | + list(block_size), | ||
| 192 | + [use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0], | ||
| 193 | + get_amax_clamp_for_quant(fmt), | ||
| 194 | + get_quant_vector_cores(), | ||
| 195 | + sf, | ||
| 196 | + ) | ||
| 197 | + | ||
| 198 | + | ||
| 199 | +def _registered_block_kwargs(x, fmt, block_size, options, clamp_min_value, num_cores): | ||
| 200 | + return dict( | ||
| 201 | + x=x, | ||
| 202 | + fmt=fmt, | ||
| 203 | + block_size=tuple(block_size), | ||
| 204 | + use_tma_aligned_col_major_sf=options[0], | ||
| 205 | + round_sf=options[1], | ||
| 206 | + use_packed_ue8m0=options[2], | ||
| 207 | + clamp_min_value=clamp_min_value, | ||
| 208 | + num_cores=num_cores, | ||
| 209 | + ) | ||
| 210 | + | ||
| 211 | + | ||
| 212 | +@torch.library.custom_op('cann_ops_tilelang::per_block_cast', mutates_args=()) | ||
| 213 | +def _per_block_cast_op( | ||
| 214 | + x: torch.Tensor, | ||
| 215 | + fmt: str, | ||
| 216 | + block_size: list[int], | ||
| 217 | + options: list[bool], | ||
| 218 | + clamp_min_value: float, | ||
| 219 | + num_cores: int, | ||
| 220 | +) -> tuple[torch.Tensor, torch.Tensor]: | ||
| 221 | + return per_block_cast_impl(**_registered_block_kwargs(x, fmt, block_size, options, clamp_min_value, num_cores)) | ||
| 222 | + | ||
| 223 | + | ||
| 224 | +@_per_block_cast_op.register_fake | ||
| 225 | +def _per_block_cast_fake(x, fmt, block_size, options, clamp_min_value, num_cores): | ||
| 226 | + _, config = _validate_per_block_cast(x, fmt, tuple(block_size), options, clamp_min_value) | ||
| 227 | + data = torch.empty( | ||
| 228 | + (x.shape[0], x.shape[1] if fmt == 'e4m3' else x.shape[1] // 2), | ||
| 229 | + dtype=config.torch_dtype, | ||
| 230 | + device=x.device, | ||
| 231 | + ) | ||
| 232 | + scale = cast_epilogue(alloc_scaling_factors(x.shape, config, device=x.device), config) | ||
| 233 | + return data, scale | ||
| 234 | + | ||
| 235 | + | ||
| 236 | +@torch.library.custom_op('cann_ops_tilelang::per_block_cast_with_sf_only', mutates_args=()) | ||
| 237 | +def _per_block_cast_op_with_sf_only( | ||
| 238 | + x: torch.Tensor, | ||
| 239 | + fmt: str, | ||
| 240 | + block_size: list[int], | ||
| 241 | + options: list[bool], | ||
| 242 | + clamp_min_value: float, | ||
| 243 | + num_cores: int, | ||
| 244 | +) -> torch.Tensor: | ||
| 245 | + return per_block_cast_impl( | ||
| 246 | + **_registered_block_kwargs(x, fmt, block_size, options, clamp_min_value, num_cores), sf_only=True | ||
| 247 | + ) | ||
| 248 | + | ||
| 249 | + | ||
| 250 | +@_per_block_cast_op_with_sf_only.register_fake | ||
| 251 | +def _per_block_cast_with_sf_only_fake(x, fmt, block_size, options, clamp_min_value, num_cores): | ||
| 252 | + _, config = _validate_per_block_cast(x, fmt, tuple(block_size), options, clamp_min_value, sf_only=True) | ||
| 253 | + return cast_epilogue(alloc_scaling_factors(x.shape, config, device=x.device), config) | ||
| 254 | + | ||
| 255 | + | ||
| 256 | +@torch.library.custom_op('cann_ops_tilelang::per_block_cast_with_precomputed_sf', mutates_args=()) | ||
| 257 | +def _per_block_cast_op_with_precomputed_sf( | ||
| 258 | + x: torch.Tensor, | ||
| 259 | + fmt: str, | ||
| 260 | + block_size: list[int], | ||
| 261 | + options: list[bool], | ||
| 262 | + clamp_min_value: float, | ||
| 263 | + num_cores: int, | ||
| 264 | + sf: torch.Tensor, | ||
| 265 | +) -> torch.Tensor: | ||
| 266 | + return per_block_cast_impl( | ||
| 267 | + **_registered_block_kwargs(x, fmt, block_size, options, clamp_min_value, num_cores), sf=sf | ||
| 268 | + ) | ||
| 269 | + | ||
| 270 | + | ||
| 271 | +@_per_block_cast_op_with_precomputed_sf.register_fake | ||
| 272 | +def _per_block_cast_with_precomputed_sf_fake(x, fmt, block_size, options, clamp_min_value, num_cores, sf): | ||
| 273 | + _, config = _validate_per_block_cast(x, fmt, tuple(block_size), options, clamp_min_value, sf=sf) | ||
| 274 | + return torch.empty( | ||
| 275 | + (x.shape[0], x.shape[1] if fmt == 'e4m3' else x.shape[1] // 2), | ||
| 276 | + dtype=config.torch_dtype, | ||
| 277 | + device=x.device, | ||
| 278 | + ) | ||
| @@ -7,7 +7,6 @@ import tilelang | |||
| 7 | from tilelang.ascend import language as T | 7 | from tilelang.ascend import language as T |
| 8 | from tilelang.ascend.language import simd as S | 8 | from tilelang.ascend.language import simd as S |
| 9 | 9 | ||
| 10 | -from cann_ops_tilelang.config import get_num_vector_cores | ||
| 11 | from cann_ops_tilelang.quant.common import ( | 10 | from cann_ops_tilelang.quant.common import ( |
| 12 | CastInputConfig, | 11 | CastInputConfig, |
| 13 | CastOutputConfig, | 12 | CastOutputConfig, |
| @@ -22,10 +21,10 @@ def get_per_block_cast_kernel_asc( | |||
| 22 | hidden: int, | 21 | hidden: int, |
| 23 | in_config: CastInputConfig, | 22 | in_config: CastInputConfig, |
| 24 | out_config: CastOutputConfig, | 23 | out_config: CastOutputConfig, |
| 24 | + n_cores: int, | ||
| 25 | sf_only: bool = False, | 25 | sf_only: bool = False, |
| 26 | cast_only: bool = False, | 26 | cast_only: bool = False, |
| 27 | ): | 27 | ): |
| 28 | - n_cores = get_num_vector_cores() | ||
| 29 | num_per_tokens, num_per_channels = out_config.sf_block | 28 | num_per_tokens, num_per_channels = out_config.sf_block |
| 30 | assert num_per_tokens in (32, 128) and num_per_channels in (32, 128) | 29 | assert num_per_tokens in (32, 128) and num_per_channels in (32, 128) |
| 31 | assert out_config.dtype in (T.float4_e2m1fn, T.float8_e4m3fn) | 30 | assert out_config.dtype in (T.float4_e2m1fn, T.float8_e4m3fn) |
| @@ -1,13 +1,14 @@ | |||
| 1 | import torch | 1 | import torch |
| 2 | 2 | ||
| 3 | -from cann_ops_tilelang.config import require_ascend, should_print_kernel_source | 3 | +from cann_ops_tilelang.config import get_amax_clamp_for_quant, require_ascend, should_print_kernel_source |
| 4 | from cann_ops_tilelang.quant.common import ( | 4 | from cann_ops_tilelang.quant.common import ( |
| 5 | QuantTensor, | 5 | QuantTensor, |
| 6 | alloc_scaling_factors, | 6 | alloc_scaling_factors, |
| 7 | cast_epilogue, | 7 | cast_epilogue, |
| 8 | get_cast_input_and_config, | 8 | get_cast_input_and_config, |
| 9 | - get_cast_output_config, | 9 | + get_cast_output_config_impl, |
| 10 | get_logical_hidden, | 10 | get_logical_hidden, |
| 11 | + get_quant_vector_cores, | ||
| 11 | ) | 12 | ) |
| 12 | 13 | ||
| 13 | from .per_block_cast_lossless_kernel import get_per_block_cast_lossless_kernel_asc | 14 | from .per_block_cast_lossless_kernel import get_per_block_cast_lossless_kernel_asc |
| @@ -36,7 +37,13 @@ def per_block_cast_lossless( | |||
| 36 | Returns: | 37 | Returns: |
| 37 | A tuple ``(out, out_sf)`` with FP8 output and sf-factor tensor. | 38 | A tuple ``(out, out_sf)`` with FP8 output and sf-factor tensor. |
| 38 | """ | 39 | """ |
| 39 | - require_ascend() | 40 | + options = [use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0] |
| 41 | + return _per_block_cast_lossless_impl( | ||
| 42 | + x, fmt, x_block_size, out_block_size, options, get_amax_clamp_for_quant(fmt), get_quant_vector_cores() | ||
| 43 | + ) | ||
| 44 | + | ||
| 45 | + | ||
| 46 | +def _validate_per_block_cast_lossless(x, fmt, x_block_size, out_block_size, options, clamp_min_value): | ||
| 40 | x, x_sf, in_config = get_cast_input_and_config(x, x_block_size) | 47 | x, x_sf, in_config = get_cast_input_and_config(x, x_block_size) |
| 41 | 48 | ||
| 42 | assert fmt == 'e4m3' and x.dtype == torch.int8, 'lossless cast only supports e2m1 -> e4m3 input' | 49 | assert fmt == 'e4m3' and x.dtype == torch.int8, 'lossless cast only supports e2m1 -> e4m3 input' |
| @@ -47,7 +54,7 @@ def per_block_cast_lossless( | |||
| 47 | assert num_tokens % 128 == 0, f'lossless cast requires num_tokens % 128 == 0, got {num_tokens}' | 54 | assert num_tokens % 128 == 0, f'lossless cast requires num_tokens % 128 == 0, got {num_tokens}' |
| 48 | assert hidden % 32 == 0 | 55 | assert hidden % 32 == 0 |
| 49 | 56 | ||
| 50 | - out_config = get_cast_output_config(fmt, out_block_size, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0) | 57 | + out_config = get_cast_output_config_impl(fmt, out_block_size, *options, clamp_min_value) |
| 51 | 58 | ||
| 52 | assert x_block_size == (1, 32), 'ascend lossless only supports in (1, 32)' | 59 | assert x_block_size == (1, 32), 'ascend lossless only supports in (1, 32)' |
| 53 | assert out_block_size in ( | 60 | assert out_block_size in ( |
| @@ -55,12 +62,32 @@ def per_block_cast_lossless( | |||
| 55 | (1, 128), | 62 | (1, 128), |
| 56 | (128, 128), | 63 | (128, 128), |
| 57 | ), 'ascend lossless only supports out (32, 32) / (1, 128) / (128, 128)' | 64 | ), 'ascend lossless only supports out (32, 32) / (1, 128) / (128, 128)' |
| 65 | + return x, x_sf, in_config, out_config | ||
| 66 | + | ||
| 67 | + | ||
| 68 | +@torch.library.custom_op('cann_ops_tilelang::per_block_cast_lossless', mutates_args=()) | ||
| 69 | +def _per_block_cast_lossless_impl( | ||
| 70 | + x: list[torch.Tensor], | ||
| 71 | + fmt: str, | ||
| 72 | + x_block_size: list[int], | ||
| 73 | + out_block_size: list[int], | ||
| 74 | + options: list[bool], | ||
| 75 | + clamp_min_value: float, | ||
| 76 | + num_cores: int, | ||
| 77 | +) -> tuple[torch.Tensor, torch.Tensor]: | ||
| 78 | + require_ascend() | ||
| 79 | + x, x_sf, in_config, out_config = _validate_per_block_cast_lossless( | ||
| 80 | + tuple(x), fmt, tuple(x_block_size), tuple(out_block_size), options, clamp_min_value | ||
| 81 | + ) | ||
| 82 | + num_tokens, physical_hidden = x.shape | ||
| 83 | + hidden = get_logical_hidden(physical_hidden, x.dtype) | ||
| 58 | # FP4 is stored as packed int8; reinterpret as the 2x-packed fp4 dtype the kernel expects. | 84 | # FP4 is stored as packed int8; reinterpret as the 2x-packed fp4 dtype the kernel expects. |
| 59 | x = x.view(torch.float4_e2m1fn_x2) | 85 | x = x.view(torch.float4_e2m1fn_x2) |
| 60 | kernel = get_per_block_cast_lossless_kernel_asc( | 86 | kernel = get_per_block_cast_lossless_kernel_asc( |
| 61 | hidden=hidden, | 87 | hidden=hidden, |
| 62 | in_config=in_config, | 88 | in_config=in_config, |
| 63 | out_config=out_config, | 89 | out_config=out_config, |
| 90 | + num_cores=num_cores, | ||
| 64 | ) | 91 | ) |
| 65 | 92 | ||
| 66 | if should_print_kernel_source(): | 93 | if should_print_kernel_source(): |
| @@ -79,3 +106,14 @@ def per_block_cast_lossless( | |||
| 79 | out_sf = cast_epilogue(out_sf, out_config) | 106 | out_sf = cast_epilogue(out_sf, out_config) |
| 80 | 107 | ||
| 81 | return out, out_sf | 108 | return out, out_sf |
| 109 | + | ||
| 110 | + | ||
| 111 | +@_per_block_cast_lossless_impl.register_fake | ||
| 112 | +def _per_block_cast_lossless_fake(x, fmt, x_block_size, out_block_size, options, clamp_min_value, num_cores): | ||
| 113 | + x, _, _, config = _validate_per_block_cast_lossless( | ||
| 114 | + tuple(x), fmt, tuple(x_block_size), tuple(out_block_size), options, clamp_min_value | ||
| 115 | + ) | ||
| 116 | + shape = (x.shape[0], get_logical_hidden(x.shape[1], x.dtype)) | ||
| 117 | + out = torch.empty(shape, dtype=config.torch_dtype, device=x.device) | ||
| 118 | + scale = cast_epilogue(alloc_scaling_factors(shape, config, device=x.device), config) | ||
| 119 | + return out, scale | ||
| @@ -7,7 +7,6 @@ import tilelang | |||
| 7 | from tilelang.ascend import language as T | 7 | from tilelang.ascend import language as T |
| 8 | from tilelang.ascend.language import simd as S | 8 | from tilelang.ascend.language import simd as S |
| 9 | 9 | ||
| 10 | -from cann_ops_tilelang.config import get_num_vector_cores | ||
| 11 | from cann_ops_tilelang.quant.common import ( | 10 | from cann_ops_tilelang.quant.common import ( |
| 12 | CastInputConfig, | 11 | CastInputConfig, |
| 13 | CastOutputConfig, | 12 | CastOutputConfig, |
| @@ -61,7 +60,9 @@ def init_fp4_to_fp8_map(fp4_to_fp8_map_ub, invalid_exp_ub): | |||
| 61 | 60 | ||
| 62 | 61 | ||
| 63 | @tilelang.jit() | 62 | @tilelang.jit() |
| 64 | -def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConfig, out_config: CastOutputConfig): | 63 | +def get_per_block_cast_lossless_kernel_asc( |
| 64 | + hidden: int, in_config: CastInputConfig, out_config: CastOutputConfig, num_cores: int | ||
| 65 | +): | ||
| 65 | assert in_config.dtype == T.float4_e2m1fn and out_config.dtype == T.float8_e4m3fn | 66 | assert in_config.dtype == T.float4_e2m1fn and out_config.dtype == T.float8_e4m3fn |
| 66 | assert in_config.sf_block == (1, 32) | 67 | assert in_config.sf_block == (1, 32) |
| 67 | assert out_config.sf_block in ((32, 32), (1, 128), (128, 128)) | 68 | assert out_config.sf_block in ((32, 32), (1, 128), (128, 128)) |
| @@ -118,7 +119,6 @@ def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConf | |||
| 118 | vec_size = 64 | 119 | vec_size = 64 |
| 119 | assert vec_size * num_per_channels >= block_k | 120 | assert vec_size * num_per_channels >= block_k |
| 120 | 121 | ||
| 121 | - num_cores = get_num_vector_cores() | ||
| 122 | sf_dtype, in_pack_factor = get_sf_storage(in_config, T.uint32) | 122 | sf_dtype, in_pack_factor = get_sf_storage(in_config, T.uint32) |
| 123 | num_sf_load_cols = num_sf_cols_per_block // in_pack_factor | 123 | num_sf_load_cols = num_sf_cols_per_block // in_pack_factor |
| 124 | out_sf_dtype, out_pack_factor = get_sf_storage(out_config, T.float32) | 124 | out_sf_dtype, out_pack_factor = get_sf_storage(out_config, T.float32) |
Msrc/cann_ops_tilelang/quant/per_channel_cast.py+80-5File content is under review, please try again later
Msrc/cann_ops_tilelang/quant/per_channel_cast_kernel.py+1-3File content is under review, please try again later
Msrc/cann_ops_tilelang/quant/per_token_cast_and_cast_back.py+41-6File content is under review, please try again later
Msrc/cann_ops_tilelang/quant/per_token_cast_and_cast_back_kernel.py+1-2File content is under review, please try again later
Msrc/cann_ops_tilelang/quant/swiglu_backward.py+97-12File content is under review, please try again later
Msrc/cann_ops_tilelang/quant/swiglu_backward_kernel.py+2-2File content is under review, please try again later
Msrc/cann_ops_tilelang/quant/swiglu_forward.py+214-11File content is under review, please try again later
| @@ -1,6 +1,7 @@ | |||
| 1 | import os | 1 | import os |
| 2 | 2 | ||
| 3 | import pytest | 3 | import pytest |
| 4 | +import torch | ||
| 4 | 5 | ||
| 5 | import cann_ops_tilelang | 6 | import cann_ops_tilelang |
| 6 | from cann_ops_tilelang.utils import str_to_dtype | 7 | from cann_ops_tilelang.utils import str_to_dtype |
| @@ -120,3 +121,86 @@ def test_per_block_cast_benchmark(benchmark_timer, benchmark_record, params): | |||
| 120 | time_us=t_us, | 121 | time_us=t_us, |
| 121 | bandwidth_gbs=num_bytes / t_us / 1e3, | 122 | bandwidth_gbs=num_bytes / t_us / 1e3, |
| 122 | ) | 123 | ) |
| 124 | + | ||
| 125 | + | ||
| 126 | +@pytest.mark.parametrize('entry', ['cast', 'sf_only']) | ||
| 127 | +@pytest.mark.parametrize( | ||
| 128 | + ('fmt', 'dtype', 'block_size', 'tma', 'packed'), | ||
| 129 | + [ | ||
| 130 | + ('e4m3', torch.bfloat16, (32, 32), False, False), | ||
| 131 | + ('e2m1', torch.float32, (128, 128), True, True), | ||
| 132 | + ], | ||
| 133 | +) | ||
| 134 | +@torch.inference_mode() | ||
| 135 | +def test_per_block_cast_compile(entry, fmt, dtype, block_size, tma, packed): | ||
| 136 | + from cann_ops_tilelang.config import get_amax_clamp_for_quant | ||
| 137 | + from tests.testing.quant import call_quant_graph, check_quant_graph_replay, compile_quant_graph | ||
| 138 | + | ||
| 139 | + op = cann_ops_tilelang.per_block_cast if entry == 'cast' else cann_ops_tilelang.per_block_cast_with_sf_only | ||
| 140 | + compiled = compile_quant_graph(op) | ||
| 141 | + kw = dict(use_tma_aligned_col_major_sf=tma, round_sf=True, use_packed_ue8m0=packed) | ||
| 142 | + x = torch.zeros((128, 256), dtype=dtype, device='npu') | ||
| 143 | + records = [] | ||
| 144 | + for step in range(3): | ||
| 145 | + if step == 1: | ||
| 146 | + x.normal_() | ||
| 147 | + elif step == 2: | ||
| 148 | + x = torch.randn_like(x) | ||
| 149 | + before = x.clone() | ||
| 150 | + expected = cast(x, fmt, block_size, **kw) | ||
| 151 | + result, events = call_quant_graph(compiled, x, fmt, block_size, **kw) | ||
| 152 | + records.append(events) | ||
| 153 | + outputs = result if entry == 'cast' else (result,) | ||
| 154 | + references = expected if entry == 'cast' else (expected[1],) | ||
| 155 | + for actual, reference in zip(outputs, references, strict=True): | ||
| 156 | + assert_equal(actual, reference) | ||
| 157 | + assert_equal(x, before) | ||
| 158 | + del result, outputs, actual | ||
| 159 | + check_quant_graph_replay(records) | ||
| 160 | + registered = getattr( | ||
| 161 | + torch.ops.cann_ops_tilelang, 'per_block_cast' if entry == 'cast' else 'per_block_cast_with_sf_only' | ||
| 162 | + ) | ||
| 163 | + from cann_ops_tilelang.quant.common import get_quant_vector_cores | ||
| 164 | + | ||
| 165 | + torch.library.opcheck( | ||
| 166 | + registered.default, | ||
| 167 | + (x, fmt, list(block_size), [tma, True, packed], get_amax_clamp_for_quant(fmt), get_quant_vector_cores()), | ||
| 168 | + test_utils=('test_schema', 'test_faketensor'), | ||
| 169 | + ) | ||
| 170 | + | ||
| 171 | + | ||
| 172 | +@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1']) | ||
| 173 | +@torch.inference_mode() | ||
| 174 | +def test_per_block_cast_precomputed_compile(fmt): | ||
| 175 | + from cann_ops_tilelang.config import get_amax_clamp_for_quant | ||
| 176 | + from tests.testing.quant import call_quant_graph, check_quant_graph_replay, compile_quant_graph | ||
| 177 | + | ||
| 178 | + compiled = compile_quant_graph(cann_ops_tilelang.per_block_cast_with_precomputed_sf) | ||
| 179 | + block_size = (32, 32) | ||
| 180 | + x = torch.randn((128, 256), dtype=torch.bfloat16, device='npu') | ||
| 181 | + sf = cast(x, fmt, block_size, round_sf=True)[1] | ||
| 182 | + records = [] | ||
| 183 | + for step in range(3): | ||
| 184 | + if step == 1: | ||
| 185 | + x.neg_() | ||
| 186 | + sf.mul_(2) | ||
| 187 | + elif step == 2: | ||
| 188 | + x = torch.randn_like(x) | ||
| 189 | + sf = cast(x, fmt, block_size, round_sf=True)[1] | ||
| 190 | + before = (x.clone(), sf.clone()) | ||
| 191 | + expected = cast(x, fmt, block_size, sf=sf) | ||
| 192 | + result, events = call_quant_graph(compiled, x, fmt, block_size, sf) | ||
| 193 | + records.append(events) | ||
| 194 | + assert_equal(result.view(torch.uint8), expected.view(torch.uint8)) | ||
| 195 | + assert result.shape == expected.shape and result.dtype == expected.dtype | ||
| 196 | + assert_equal(x, before[0]) | ||
| 197 | + assert_equal(sf, before[1]) | ||
| 198 | + del result | ||
| 199 | + check_quant_graph_replay(records) | ||
| 200 | + from cann_ops_tilelang.quant.common import get_quant_vector_cores | ||
| 201 | + | ||
| 202 | + torch.library.opcheck( | ||
| 203 | + torch.ops.cann_ops_tilelang.per_block_cast_with_precomputed_sf.default, | ||
| 204 | + (x, fmt, list(block_size), [False, False, False], get_amax_clamp_for_quant(fmt), get_quant_vector_cores(), sf), | ||
| 205 | + test_utils=('test_schema', 'test_faketensor'), | ||
| 206 | + ) | ||
| @@ -133,3 +133,62 @@ def test_per_block_cast_lossless_benchmark(benchmark_timer, benchmark_record, pa | |||
| 133 | time_us=t_us, | 133 | time_us=t_us, |
| 134 | bandwidth_gbs=num_bytes / t_us / 1e3, | 134 | bandwidth_gbs=num_bytes / t_us / 1e3, |
| 135 | ) | 135 | ) |
| 136 | + | ||
| 137 | + | ||
| 138 | +@pytest.mark.parametrize( | ||
| 139 | + ('out_block', 'tma', 'packed'), | ||
| 140 | + [ | ||
| 141 | + ((32, 32), False, False), | ||
| 142 | + ((1, 128), True, True), | ||
| 143 | + ((128, 128), False, True), | ||
| 144 | + ], | ||
| 145 | +) | ||
| 146 | +@torch.inference_mode() | ||
| 147 | +def test_per_block_cast_lossless_compile(out_block, tma, packed): | ||
| 148 | + from cann_ops_tilelang.config import get_amax_clamp_for_quant | ||
| 149 | + from tests.testing.quant import call_quant_graph, check_quant_graph_replay, compile_quant_graph | ||
| 150 | + | ||
| 151 | + compiled = compile_quant_graph(cann_ops_tilelang.per_block_cast_lossless) | ||
| 152 | + kw = dict(use_tma_aligned_col_major_sf=tma, round_sf=True, use_packed_ue8m0=packed) | ||
| 153 | + source = None | ||
| 154 | + records = [] | ||
| 155 | + for step in range(3): | ||
| 156 | + raw = clamp_abs_ratio(torch.randn((128, 256), dtype=torch.float32, device='npu')) | ||
| 157 | + refreshed = cast(raw, 'e2m1', (1, 32), **kw) | ||
| 158 | + if step == 1: | ||
| 159 | + source[0].copy_(refreshed[0]) | ||
| 160 | + source[1].copy_(refreshed[1]) | ||
| 161 | + else: | ||
| 162 | + source = refreshed | ||
| 163 | + snapshots = tuple(value.clone() for value in source) | ||
| 164 | + strides = tuple(value.stride() for value in source) | ||
| 165 | + expected = cast_back(source, 'fp32', (1, 32)) | ||
| 166 | + reference = cann_ops_tilelang.per_block_cast_lossless(source, 'e4m3', (1, 32), out_block, **kw) | ||
| 167 | + result, events = call_quant_graph(compiled, source, 'e4m3', (1, 32), out_block, **kw) | ||
| 168 | + records.append(events) | ||
| 169 | + assert_equal(cast_back(result, 'fp32', out_block), expected) | ||
| 170 | + assert_equal(result[0].view(torch.uint8), reference[0].view(torch.uint8)) | ||
| 171 | + assert_equal(result[1], reference[1]) | ||
| 172 | + for actual, ref in zip(result, reference, strict=True): | ||
| 173 | + assert actual.shape == ref.shape and actual.stride() == ref.stride() | ||
| 174 | + assert actual.dtype == ref.dtype and actual.device == ref.device | ||
| 175 | + for value, before, stride in zip(source, snapshots, strides, strict=True): | ||
| 176 | + assert_equal(value, before, check_stride=False) | ||
| 177 | + assert value.stride() == stride | ||
| 178 | + del result, actual | ||
| 179 | + check_quant_graph_replay(records) | ||
| 180 | + from cann_ops_tilelang.quant.common import get_quant_vector_cores | ||
| 181 | + | ||
| 182 | + torch.library.opcheck( | ||
| 183 | + torch.ops.cann_ops_tilelang.per_block_cast_lossless.default, | ||
| 184 | + ( | ||
| 185 | + source, | ||
| 186 | + 'e4m3', | ||
| 187 | + [1, 32], | ||
| 188 | + list(out_block), | ||
| 189 | + [tma, True, packed], | ||
| 190 | + get_amax_clamp_for_quant('e4m3'), | ||
| 191 | + get_quant_vector_cores(), | ||
| 192 | + ), | ||
| 193 | + test_utils=('test_schema', 'test_faketensor'), | ||
| 194 | + ) | ||
Mtests/quant/test_per_token_cast_and_cast_back.py+71-0File content is under review, please try again later
| @@ -1,3 +1,5 @@ | |||
| 1 | +import logging | ||
| 2 | + | ||
| 1 | import torch | 3 | import torch |
| 2 | 4 | ||
| 3 | from cann_ops_tilelang.quant.common import get_packed_ue8m0_pack_factor, unpack_from_e2m1fn_x2 | 5 | from cann_ops_tilelang.quant.common import get_packed_ue8m0_pack_factor, unpack_from_e2m1fn_x2 |
| @@ -65,3 +67,33 @@ def clear_unused_sf_col_pack(sf: torch.Tensor, num_tokens: int, num_per_tokens: | |||
| 65 | sf_unpacked = sf_uint8.view(aligned_num_token_blocks // pack_factor, -1, pack_factor) | 67 | sf_unpacked = sf_uint8.view(aligned_num_token_blocks // pack_factor, -1, pack_factor) |
| 66 | sf_unpacked[num_token_blocks // pack_factor :, :, num_token_blocks % pack_factor :] = 0 | 68 | sf_unpacked[num_token_blocks // pack_factor :, :, num_token_blocks % pack_factor :] = 0 |
| 67 | return sf_uint8.view(aligned_num_token_blocks // pack_factor, -1) | 69 | return sf_uint8.view(aligned_num_token_blocks // pack_factor, -1) |
| 70 | + | ||
| 71 | + | ||
| 72 | +def compile_quant_graph(op): | ||
| 73 | + """Build an independent fixed-shape graph for a small correctness test.""" | ||
| 74 | + torch.compiler.reset() | ||
| 75 | + return torch.compile(op, backend='npugraph_ex', fullgraph=True, dynamic=False) | ||
| 76 | + | ||
| 77 | + | ||
| 78 | +def call_quant_graph(compiled, *args, **kwargs): | ||
| 79 | + """Return one call's output and evidence of capture/replay, not just Dynamo tracing.""" | ||
| 80 | + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as profile: | ||
| 81 | + result = compiled(*args, **kwargs) | ||
| 82 | + torch.npu.synchronize() | ||
| 83 | + events = {event.key: event.count for event in profile.key_averages()} | ||
| 84 | + names = ('acl_graph_capture', 'acl_graph_replay', 'fx_run_eagerly') | ||
| 85 | + counts = {name: events.get(name, 0) for name in names} | ||
| 86 | + logging.getLogger(__name__).info('NPU graph events: %s', counts) | ||
| 87 | + return result, counts | ||
| 88 | + | ||
| 89 | + | ||
| 90 | +def check_quant_graph_replay(records): | ||
| 91 | + """Same-storage input updates must replay; mutated state may recapture after replacement.""" | ||
| 92 | + if not any(record['acl_graph_capture'] for record in records): | ||
| 93 | + raise AssertionError(f'No NPU graph capture recorded: {records}') | ||
| 94 | + if records[1]['acl_graph_replay'] <= 0: | ||
| 95 | + raise AssertionError(f'The second call did not replay the NPU graph: {records}') | ||
| 96 | + if records[1]['acl_graph_capture'] != 0: | ||
| 97 | + raise AssertionError(f'The second call recaptured the NPU graph: {records}') | ||
| 98 | + if not all(record['fx_run_eagerly'] == 0 for record in records): | ||
| 99 | + raise AssertionError(f'Unexpected eager fallback: {records}') | ||