Open
feat(quant): support torch.compile for per-block quantization ops #66
daisysjAdd on 1 day ago
feat(quant): support torch.compile for per-block quantization ops #66
Open
Merged fromdaisysj/ops-tilelanggraph_adaptintocann/ops-tilelangmaster
Total: 17 files changed+857-60
| @@ -4,7 +4,7 @@ from tilelang.ascend import language as T | |||
| 4 | from cann_ops_tilelang.config import require_ascend, should_print_kernel_source | 4 | from cann_ops_tilelang.config import require_ascend, should_print_kernel_source |
| 5 | 5 | ||
| 6 | from .cast_back_kernel import get_cast_back_kernel_ascend | 6 | from .cast_back_kernel import get_cast_back_kernel_ascend |
| 7 | -from .common import get_cast_input_and_config, get_logical_hidden | 7 | +from .common import get_cast_input_and_config, get_logical_hidden, get_quant_vector_cores |
| 8 | from .types import QuantTensor | 8 | from .types import QuantTensor |
| 9 | 9 | ||
| 10 | 10 | ||
| @@ -24,30 +24,55 @@ def cast_back( | |||
| 24 | 24 | ||
| 25 | Returns: | 25 | Returns: |
| 26 | Dequantized tensor in the requested format. | 26 | Dequantized tensor in the requested format. |
| 27 | + | ||
| 28 | + Fixed-shape npugraph_ex inference is supported. No autograd formula is registered. | ||
| 27 | """ | 29 | """ |
| 28 | - require_ascend() | 30 | + num_cores = get_quant_vector_cores() if x[0].shape[0] else 0 |
| 31 | + return _cast_back_impl(x, fmt, x_block_size, x_special_fmt, num_cores) | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +def _validate_cast_back(x, x_sf, fmt, x_block_size, x_special_fmt): | ||
| 29 | assert fmt in ('bf16', 'fp32') | 35 | assert fmt in ('bf16', 'fp32') |
| 30 | out_dtype = torch.bfloat16 if fmt == 'bf16' else torch.float32 | 36 | out_dtype = torch.bfloat16 if fmt == 'bf16' else torch.float32 |
| 31 | 37 | ||
| 32 | - x, x_sf = x | ||
| 33 | if x.device.type != 'npu' or x_sf.device.type != 'npu': | 38 | if x.device.type != 'npu' or x_sf.device.type != 'npu': |
| 34 | raise ValueError('x and its scaling factors must be NPU tensors.') | 39 | raise ValueError('x and its scaling factors must be NPU tensors.') |
| 35 | 40 | ||
| 36 | assert x_special_fmt is None, f'Unsupported special format: {x_special_fmt}' | 41 | assert x_special_fmt is None, f'Unsupported special format: {x_special_fmt}' |
| 37 | 42 | ||
| 38 | - num_tokens, hidden = x.shape | 43 | + _, hidden = x.shape |
| 39 | assert x_sf.dim() == 2 | 44 | assert x_sf.dim() == 2 |
| 40 | assert x_sf.dtype in (torch.int16, torch.int32, torch.float32) | 45 | assert x_sf.dtype in (torch.int16, torch.int32, torch.float32) |
| 41 | 46 | ||
| 42 | hidden = get_logical_hidden(hidden, x.dtype) | 47 | hidden = get_logical_hidden(hidden, x.dtype) |
| 43 | assert hidden % 64 == 0 | 48 | assert hidden % 64 == 0 |
| 44 | 49 | ||
| 45 | - x, x_sf, in_config = get_cast_input_and_config((x, x_sf), x_block_size) | 50 | + _, x_sf, in_config = get_cast_input_and_config((x, x_sf), x_block_size) |
| 51 | + # Mirror the kernel's static block/layout restrictions in fake execution. | ||
| 52 | + assert in_config.sf_block in ((1, 32), (1, 128), (32, 1), (32, 32), (128, 1), (128, 128)) | ||
| 53 | + assert not (in_config.sf_block == (32, 1) and in_config.use_tma_aligned_col_major_sf) | ||
| 54 | + return x_sf, in_config, hidden, out_dtype | ||
| 55 | + | ||
| 56 | + | ||
| 57 | +@torch.library.custom_op( | ||
| 58 | + 'cann_ops_tilelang::cast_back', | ||
| 59 | + mutates_args=(), | ||
| 60 | + schema='(Tensor[] x, str fmt, SymInt[] x_block_size, str? x_special_fmt, int num_cores) -> Tensor', | ||
| 61 | +) | ||
| 62 | +def _cast_back_impl(x, fmt, x_block_size, x_special_fmt, num_cores): | ||
| 63 | + x_block_size = tuple(x_block_size) | ||
| 64 | + require_ascend() | ||
| 65 | + data, scale = x | ||
| 66 | + x_sf, in_config, hidden, out_dtype = _validate_cast_back(data, scale, fmt, x_block_size, x_special_fmt) | ||
| 67 | + x = data | ||
| 68 | + num_tokens = x.shape[0] | ||
| 46 | out = torch.empty((num_tokens, hidden), dtype=out_dtype, device=x.device) | 69 | out = torch.empty((num_tokens, hidden), dtype=out_dtype, device=x.device) |
| 47 | if num_tokens == 0: | 70 | if num_tokens == 0: |
| 48 | return out | 71 | return out |
| 49 | 72 | ||
| 50 | - kernel = get_cast_back_kernel_ascend(hidden=hidden, in_config=in_config, out_dtype=T.dtype(out_dtype)) | 73 | + kernel = get_cast_back_kernel_ascend( |
| 74 | + hidden=hidden, in_config=in_config, out_dtype=T.dtype(out_dtype), n_cores=num_cores | ||
| 75 | + ) | ||
| 51 | 76 | ||
| 52 | if should_print_kernel_source(): | 77 | if should_print_kernel_source(): |
| 53 | print(kernel.get_kernel_source()) | 78 | print(kernel.get_kernel_source()) |
| @@ -73,5 +98,14 @@ def per_token_cast_back( | |||
| 73 | 98 | ||
| 74 | Returns: | 99 | Returns: |
| 75 | Dequantized tensor in the requested format. | 100 | Dequantized tensor in the requested format. |
| 101 | + | ||
| 102 | + Fixed-shape npugraph_ex inference is supported. No autograd formula is registered. | ||
| 76 | """ | 103 | """ |
| 77 | return cast_back(x, fmt, (1, num_per_channels), x_special_fmt=x_special_fmt) | 104 | return cast_back(x, fmt, (1, num_per_channels), x_special_fmt=x_special_fmt) |
| 105 | + | ||
| 106 | + | ||
| 107 | +@_cast_back_impl.register_fake | ||
| 108 | +def _cast_back_fake(x, fmt, x_block_size, x_special_fmt, num_cores): | ||
| 109 | + data, scale = x | ||
| 110 | + _, _, hidden, out_dtype = _validate_cast_back(data, scale, fmt, tuple(x_block_size), x_special_fmt) | ||
| 111 | + return torch.empty((data.shape[0], hidden), dtype=out_dtype, device=data.device) | ||
| @@ -4,7 +4,6 @@ import tilelang | |||
| 4 | from tilelang.ascend import language as T | 4 | from tilelang.ascend import language as T |
| 5 | from tilelang.ascend.language import simd as S | 5 | from tilelang.ascend.language import simd as S |
| 6 | 6 | ||
| 7 | -from cann_ops_tilelang.config import get_num_vector_cores | ||
| 8 | from cann_ops_tilelang.utils import align_up as align | 7 | from cann_ops_tilelang.utils import align_up as align |
| 9 | from cann_ops_tilelang.utils import ceil_div | 8 | from cann_ops_tilelang.utils import ceil_div |
| 10 | 9 | ||
| @@ -16,9 +15,9 @@ def get_cast_back_kernel_ascend( | |||
| 16 | hidden: int, | 15 | hidden: int, |
| 17 | in_config: CastInputConfig, | 16 | in_config: CastInputConfig, |
| 18 | out_dtype: T.dtype = T.bfloat16, | 17 | out_dtype: T.dtype = T.bfloat16, |
| 18 | + *, | ||
| 19 | + n_cores: int, | ||
| 19 | ): | 20 | ): |
| 20 | - n_cores = get_num_vector_cores() | ||
| 21 | - | ||
| 22 | num_per_tokens, num_per_channels = in_config.sf_block | 21 | num_per_tokens, num_per_channels = in_config.sf_block |
| 23 | pack_factor = get_packed_ue8m0_pack_factor() if in_config.use_packed_ue8m0 else 1 | 22 | pack_factor = get_packed_ue8m0_pack_factor() if in_config.use_packed_ue8m0 else 1 |
| 24 | assert (num_per_tokens, num_per_channels) in ((1, 32), (1, 128), (32, 1), (32, 32), (128, 1), (128, 128)) | 23 | assert (num_per_tokens, num_per_channels) in ((1, 32), (1, 128), (32, 1), (32, 32), (128, 1), (128, 128)) |
| @@ -1,16 +1,32 @@ | |||
| 1 | import functools | 1 | import functools |
| 2 | +import os | ||
| 2 | from dataclasses import dataclass | 3 | from dataclasses import dataclass |
| 3 | from typing import Union # noqa: UP035 - TileLang macro annotations require typing.Union | 4 | from typing import Union # noqa: UP035 - TileLang macro annotations require typing.Union |
| 4 | 5 | ||
| 5 | import torch | 6 | import torch |
| 6 | from tilelang.ascend import language as T | 7 | from tilelang.ascend import language as T |
| 7 | 8 | ||
| 8 | -from cann_ops_tilelang.config import get_amax_clamp_for_quant | 9 | +from cann_ops_tilelang.config import get_amax_clamp_for_quant, get_num_ai_cores |
| 9 | from cann_ops_tilelang.utils import ceil_div, str_to_dtype | 10 | from cann_ops_tilelang.utils import ceil_div, str_to_dtype |
| 10 | 11 | ||
| 11 | from .types import QuantTensor | 12 | from .types import QuantTensor |
| 12 | 13 | ||
| 13 | 14 | ||
| 15 | +def get_quant_vector_cores() -> int: | ||
| 16 | + # Compiled wrappers need this metadata without tracing require_ascend. | ||
| 17 | + # Keep mutable tuning settings visible to Dynamo so they remain guarded. | ||
| 18 | + detected = get_num_ai_cores() * 2 | ||
| 19 | + # Guard the key set as well, including an unset-to-set environment change. | ||
| 20 | + key = 'OPS_TILELANG_NUM_VECTOR_CORES' | ||
| 21 | + override = os.environ[key] if key in tuple(os.environ) else None | ||
| 22 | + if override is None: | ||
| 23 | + return detected | ||
| 24 | + requested = int(override) | ||
| 25 | + if requested <= 0 or requested > detected: | ||
| 26 | + raise ValueError(f'OPS_TILELANG_NUM_VECTOR_CORES must be in [1, {detected}], got {requested}.') | ||
| 27 | + return requested | ||
| 28 | + | ||
| 29 | + | ||
| 14 | @dataclass(frozen=True) | 30 | @dataclass(frozen=True) |
| 15 | class BaseCastConfig: | 31 | class BaseCastConfig: |
| 16 | torch_dtype: torch.dtype = torch.float8_e4m3fn | 32 | torch_dtype: torch.dtype = torch.float8_e4m3fn |
| @@ -303,20 +319,20 @@ def unpack_from_e2m1fn_x2(x: torch.Tensor, out_dtype: torch.dtype = torch.float3 | |||
| 303 | if x.ndim == 0: | 319 | if x.ndim == 0: |
| 304 | raise ValueError('x must have at least 1 dimension so the last dim can be doubled') | 320 | raise ValueError('x must have at least 1 dimension so the last dim can be doubled') |
| 305 | 321 | ||
| 322 | + # Device-side fills avoid temporary host scalar copies during graph capture. | ||
| 323 | + high_shift = torch.full((), 4, dtype=x.dtype, device=x.device) | ||
| 324 | + sign_shift = torch.full((), 3, dtype=torch.int16, device=x.device) | ||
| 325 | + exponent_shift = torch.full((), 1, dtype=torch.int16, device=x.device) | ||
| 306 | lo = (x & 0x0F).to(torch.int16) | 326 | lo = (x & 0x0F).to(torch.int16) |
| 307 | - hi = ((x >> 4) & 0x0F).to(torch.int16) | 327 | + hi = ((x >> high_shift) & 0x0F).to(torch.int16) |
| 308 | 328 | ||
| 309 | def decode_fp4_e2m1(n: torch.Tensor) -> torch.Tensor: | 329 | def decode_fp4_e2m1(n: torch.Tensor) -> torch.Tensor: |
| 310 | # n in [0..15], layout: s(1) | e(2) | m(1) | 330 | # n in [0..15], layout: s(1) | e(2) | m(1) |
| 311 | - s = (n >> 3) & 0x1 | 331 | + s = (n >> sign_shift) & 0x1 |
| 312 | - e = (n >> 1) & 0x3 | 332 | + e = (n >> exponent_shift) & 0x3 |
| 313 | m = n & 0x1 | 333 | m = n & 0x1 |
| 314 | 334 | ||
| 315 | - sign = torch.where( | 335 | + sign = torch.where(s == 1, -1.0, 1.0) |
| 316 | - s == 1, | ||
| 317 | - torch.tensor(-1.0, device=n.device), | ||
| 318 | - torch.tensor(1.0, device=n.device), | ||
| 319 | - ) | ||
| 320 | 336 | ||
| 321 | bias = 1 | 337 | bias = 1 |
| 322 | 338 | ||
| @@ -326,7 +342,7 @@ def unpack_from_e2m1fn_x2(x: torch.Tensor, out_dtype: torch.dtype = torch.float3 | |||
| 326 | 342 | ||
| 327 | # normal: e in {1,2,3} -> value = sign * 2^(e-bias) * (1 + m/2) | 343 | # normal: e in {1,2,3} -> value = sign * 2^(e-bias) * (1 + m/2) |
| 328 | norm = (1.0 + m.to(torch.float32) * 0.5) * torch.pow( | 344 | norm = (1.0 + m.to(torch.float32) * 0.5) * torch.pow( |
| 329 | - torch.tensor(2.0, device=n.device), | 345 | + 2.0, |
| 330 | (e - bias).to(torch.float32), | 346 | (e - bias).to(torch.float32), |
| 331 | ) | 347 | ) |
| 332 | 348 | ||
| @@ -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) |
| @@ -1,12 +1,13 @@ | |||
| 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 | get_sf_shape, | 11 | get_sf_shape, |
| 11 | ) | 12 | ) |
| 12 | 13 | ||
| @@ -35,11 +36,33 @@ def per_channel_cast( | |||
| 35 | 36 | ||
| 36 | Returns: | 37 | Returns: |
| 37 | A tuple ``(out, out_sf)`` with FP8 output and row-major sf-factor tensor. | 38 | A tuple ``(out, out_sf)`` with FP8 output and row-major sf-factor tensor. |
| 39 | + | ||
| 40 | + Fixed-shape npugraph_ex inference is supported. Clamp configuration is an | ||
| 41 | + explicit internal graph input; no autograd formula is registered. | ||
| 38 | """ | 42 | """ |
| 39 | - require_ascend() | 43 | + clamp_min_value = get_amax_clamp_for_quant(fmt) |
| 44 | + data, scale = x if isinstance(x, tuple) else (x, None) | ||
| 45 | + return _per_channel_cast_impl( | ||
| 46 | + data, | ||
| 47 | + scale, | ||
| 48 | + fmt, | ||
| 49 | + num_per_tokens, | ||
| 50 | + round_sf, | ||
| 51 | + num_per_channels, | ||
| 52 | + use_packed_ue8m0, | ||
| 53 | + clamp_min_value, | ||
| 54 | + get_quant_vector_cores(), | ||
| 55 | + ) | ||
| 56 | + | ||
| 57 | + | ||
| 58 | +def _validate_per_channel_cast(x, fmt, num_per_tokens, round_sf, num_per_channels, use_packed_ue8m0, clamp_min_value): | ||
| 40 | x, x_sf_invs, in_config = get_cast_input_and_config(x, None if num_per_channels is None else (1, num_per_channels)) | 59 | x, x_sf_invs, in_config = get_cast_input_and_config(x, None if num_per_channels is None else (1, num_per_channels)) |
| 41 | 60 | ||
| 42 | assert fmt == 'e4m3' | 61 | assert fmt == 'e4m3' |
| 62 | + if in_config.with_sf: | ||
| 63 | + assert x.dtype == torch.float8_e4m3fn, 'per_channel_cast with scaling factors requires e4m3 input' | ||
| 64 | + else: | ||
| 65 | + assert x.dtype == torch.bfloat16, 'per_channel_cast without scaling factors requires bf16 input' | ||
| 43 | assert x.device.type == 'npu' | 66 | assert x.device.type == 'npu' |
| 44 | assert x.dim() == 2 and x.is_contiguous() | 67 | assert x.dim() == 2 and x.is_contiguous() |
| 45 | assert num_per_tokens in (32, 128) | 68 | assert num_per_tokens in (32, 128) |
| @@ -53,11 +76,44 @@ def per_channel_cast( | |||
| 53 | if in_config.use_tma_aligned_col_major_sf and in_config.use_packed_ue8m0: | 76 | if in_config.use_tma_aligned_col_major_sf and in_config.use_packed_ue8m0: |
| 54 | x_sf_invs = x_sf_invs.view(torch.int16) | 77 | x_sf_invs = x_sf_invs.view(torch.int16) |
| 55 | 78 | ||
| 56 | - out_config = get_cast_output_config(fmt, (num_per_tokens, 1), False, round_sf, use_packed_ue8m0) | 79 | + out_config = get_cast_output_config_impl( |
| 80 | + fmt, (num_per_tokens, 1), False, round_sf, use_packed_ue8m0, clamp_min_value | ||
| 81 | + ) | ||
| 82 | + quant_max = 448.0 | ||
| 83 | + assert out_config.clamp_min_value >= quant_max * (2 ** (-126)), ( | ||
| 84 | + 'Ascend scale exponent fast path requires clamp to normal number' | ||
| 85 | + ) | ||
| 57 | assert num_tokens % 128 == 0 | 86 | assert num_tokens % 128 == 0 |
| 87 | + return x, x_sf_invs, in_config, out_config | ||
| 88 | + | ||
| 89 | + | ||
| 90 | +@torch.library.custom_op('cann_ops_tilelang::per_channel_cast', mutates_args=()) | ||
| 91 | +def _per_channel_cast_impl( | ||
| 92 | + data: torch.Tensor, | ||
| 93 | + scale: torch.Tensor | None, | ||
| 94 | + fmt: str, | ||
| 95 | + num_per_tokens: int, | ||
| 96 | + round_sf: bool, | ||
| 97 | + num_per_channels: int | None, | ||
| 98 | + use_packed_ue8m0: bool, | ||
| 99 | + clamp_min_value: float, | ||
| 100 | + num_cores: int, | ||
| 101 | +) -> tuple[torch.Tensor, torch.Tensor]: | ||
| 102 | + source = data if scale is None else (data, scale) | ||
| 103 | + require_ascend() | ||
| 104 | + x, x_sf_invs, in_config, out_config = _validate_per_channel_cast( | ||
| 105 | + source, | ||
| 106 | + fmt, | ||
| 107 | + num_per_tokens, | ||
| 108 | + round_sf, | ||
| 109 | + num_per_channels, | ||
| 110 | + use_packed_ue8m0, | ||
| 111 | + clamp_min_value, | ||
| 112 | + ) | ||
| 113 | + num_tokens, hidden = x.shape | ||
| 58 | out = torch.empty((num_tokens, hidden), dtype=out_config.torch_dtype, device=x.device) | 114 | out = torch.empty((num_tokens, hidden), dtype=out_config.torch_dtype, device=x.device) |
| 59 | out_sf = alloc_scaling_factors((num_tokens, hidden), out_config, device=x.device) | 115 | out_sf = alloc_scaling_factors((num_tokens, hidden), out_config, device=x.device) |
| 60 | - kernel = get_per_channel_cast_kernel_asc(hidden, in_config=in_config, out_config=out_config) | 116 | + kernel = get_per_channel_cast_kernel_asc(hidden, in_config=in_config, out_config=out_config, num_cores=num_cores) |
| 61 | 117 | ||
| 62 | if should_print_kernel_source(): | 118 | if should_print_kernel_source(): |
| 63 | print(kernel.get_kernel_source()) | 119 | print(kernel.get_kernel_source()) |
| @@ -66,3 +122,22 @@ def per_channel_cast( | |||
| 66 | kernel(x, out, out_sf, x_sf_invs) | 122 | kernel(x, out, out_sf, x_sf_invs) |
| 67 | out_sf = cast_epilogue(out_sf, out_config) | 123 | out_sf = cast_epilogue(out_sf, out_config) |
| 68 | return out, out_sf | 124 | return out, out_sf |
| 125 | + | ||
| 126 | + | ||
| 127 | +@_per_channel_cast_impl.register_fake | ||
| 128 | +def _per_channel_cast_fake( | ||
| 129 | + data, scale, fmt, num_per_tokens, round_sf, num_per_channels, use_packed_ue8m0, clamp_min_value, num_cores | ||
| 130 | +): | ||
| 131 | + source = data if scale is None else (data, scale) | ||
| 132 | + _, _, _, config = _validate_per_channel_cast( | ||
| 133 | + source, | ||
| 134 | + fmt, | ||
| 135 | + num_per_tokens, | ||
| 136 | + round_sf, | ||
| 137 | + num_per_channels, | ||
| 138 | + use_packed_ue8m0, | ||
| 139 | + clamp_min_value, | ||
| 140 | + ) | ||
| 141 | + output = torch.empty(data.shape, dtype=config.torch_dtype, device=data.device) | ||
| 142 | + scales = alloc_scaling_factors(tuple(data.shape), config, data.device) | ||
| 143 | + return output, cast_epilogue(scales, config) | ||
| @@ -4,7 +4,6 @@ import tilelang | |||
| 4 | from tilelang.ascend import language as T | 4 | from tilelang.ascend import language as T |
| 5 | from tilelang.ascend.language import simd as S | 5 | from tilelang.ascend.language import simd as S |
| 6 | 6 | ||
| 7 | -from cann_ops_tilelang.config import get_num_vector_cores | ||
| 8 | from cann_ops_tilelang.quant.common import ( | 7 | from cann_ops_tilelang.quant.common import ( |
| 9 | CastInputConfig, | 8 | CastInputConfig, |
| 10 | CastOutputConfig, | 9 | CastOutputConfig, |
| @@ -19,6 +18,7 @@ def get_per_channel_cast_kernel_asc( | |||
| 19 | hidden: int, | 18 | hidden: int, |
| 20 | in_config: CastInputConfig, | 19 | in_config: CastInputConfig, |
| 21 | out_config: CastOutputConfig, | 20 | out_config: CastOutputConfig, |
| 21 | + num_cores: int, | ||
| 22 | ): | 22 | ): |
| 23 | # bf16 hidden >= 2048 ę¶ēØ nbg=1ļ¼čæ代ę°ēæ»åļ¼MTE2 čÆ·ę±ę“é¢ē¹ļ¼bf16 2048 ę¹å ~4%ļ¼ļ¼ | 23 | # bf16 hidden >= 2048 ę¶ēØ nbg=1ļ¼čæ代ę°ēæ»åļ¼MTE2 čÆ·ę±ę“é¢ē¹ļ¼bf16 2048 ę¹å ~4%ļ¼ļ¼ |
| 24 | # å° hidden(128) äæę nbg=2 仄čå pack_rowļ¼éæå éåć | 24 | # å° hidden(128) äæę nbg=2 仄čå pack_rowļ¼éæå éåć |
| @@ -91,8 +91,6 @@ def get_per_channel_cast_kernel_asc( | |||
| 91 | else: | 91 | else: |
| 92 | num_blocks_per_group = 1 | 92 | num_blocks_per_group = 1 |
| 93 | 93 | ||
| 94 | - num_cores = get_num_vector_cores() | ||
| 95 | - | ||
| 96 | num_tokens = T.dynamic('num_tokens') | 94 | num_tokens = T.dynamic('num_tokens') |
| 97 | sf_stride = T.dynamic('sf_stride') | 95 | sf_stride = T.dynamic('sf_stride') |
| 98 | sf_shape = get_sf_shape((num_tokens, hidden), out_config) | 96 | sf_shape = get_sf_shape((num_tokens, hidden), out_config) |
| @@ -1,9 +1,9 @@ | |||
| 1 | import torch | 1 | import torch |
| 2 | from tilelang.ascend import language as T | 2 | from tilelang.ascend import language as T |
| 3 | 3 | ||
| 4 | -from cann_ops_tilelang.config import require_ascend, should_print_kernel_source | 4 | +from cann_ops_tilelang.config import get_amax_clamp_for_quant, require_ascend, should_print_kernel_source |
| 5 | 5 | ||
| 6 | -from .common import get_cast_output_config | 6 | +from .common import get_cast_output_config_impl, get_quant_vector_cores |
| 7 | from .per_token_cast_and_cast_back_kernel import get_per_token_cast_and_cast_back_kernel_asc | 7 | from .per_token_cast_and_cast_back_kernel import get_per_token_cast_and_cast_back_kernel_asc |
| 8 | 8 | ||
| 9 | 9 | ||
| @@ -21,26 +21,56 @@ def per_token_cast_and_cast_back( | |||
| 21 | fmt: Quantized format used for cast-and-cast-back (``'e4m3'`` or ``'e2m1'``). | 21 | fmt: Quantized format used for cast-and-cast-back (``'e4m3'`` or ``'e2m1'``). |
| 22 | num_per_channels: Number of channels in each scaling block (32, 64 or 128). | 22 | num_per_channels: Number of channels in each scaling block (32, 64 or 128). |
| 23 | round_sf: Whether to round scaling factors to powers of two. | 23 | round_sf: Whether to round scaling factors to powers of two. |
| 24 | + | ||
| 25 | + Returns None and mutates only x. Fixed-shape npugraph_ex inference is | ||
| 26 | + supported; the registered operator declares the same mutation contract. | ||
| 24 | """ | 27 | """ |
| 25 | - require_ascend() | 28 | + clamp_min_value = get_amax_clamp_for_quant(fmt) |
| 29 | + _cast_and_back_impl(x, fmt, num_per_channels, round_sf, clamp_min_value, get_quant_vector_cores()) | ||
| 30 | + return None | ||
| 31 | + | ||
| 32 | + | ||
| 33 | +def _validate_cast_and_back(x, fmt, num_per_channels, round_sf, clamp_min_value): | ||
| 26 | if x.device.type != 'npu': | 34 | if x.device.type != 'npu': |
| 27 | raise ValueError(f'x must be an NPU tensor, got {x.device}.') | 35 | raise ValueError(f'x must be an NPU tensor, got {x.device}.') |
| 28 | assert x.dim() == 2 | 36 | assert x.dim() == 2 |
| 37 | + assert x.dtype in (torch.bfloat16, torch.float16, torch.float32), ( | ||
| 38 | + 'Ascend per-token cast and cast-back only supports bf16, fp16 or fp32 input' | ||
| 39 | + ) | ||
| 29 | assert x.stride(1) == 1 | 40 | assert x.stride(1) == 1 |
| 30 | assert fmt in ('e4m3', 'e2m1') | 41 | assert fmt in ('e4m3', 'e2m1') |
| 31 | - assert num_per_channels in (32, 64, 128) | 42 | + assert num_per_channels in (32, 64, 128), ( |
| 43 | + 'Ascend per-token cast and cast-back only supports num_per_channels=32/64/128' | ||
| 44 | + ) | ||
| 32 | 45 | ||
| 33 | - num_tokens, hidden = x.shape | 46 | + _, hidden = x.shape |
| 47 | + assert hidden % 64 == 0, 'Ascend per-token cast and cast-back hidden must be a multiple of 64' | ||
| 34 | assert hidden % num_per_channels == 0 | 48 | assert hidden % num_per_channels == 0 |
| 35 | 49 | ||
| 36 | - out_config = get_cast_output_config(fmt, (1, num_per_channels), round_sf=round_sf) | 50 | + out_config = get_cast_output_config_impl(fmt, (1, num_per_channels), False, round_sf, False, clamp_min_value) |
| 37 | assert out_config.sf_block[0] == 1 | 51 | assert out_config.sf_block[0] == 1 |
| 38 | 52 | ||
| 53 | + return out_config | ||
| 54 | + | ||
| 55 | + | ||
| 56 | +@torch.library.custom_op('cann_ops_tilelang::per_token_cast_and_cast_back', mutates_args=('x',)) | ||
| 57 | +def _cast_and_back_impl( | ||
| 58 | + x: torch.Tensor, | ||
| 59 | + fmt: str, | ||
| 60 | + num_per_channels: int, | ||
| 61 | + round_sf: bool, | ||
| 62 | + clamp_min_value: float, | ||
| 63 | + num_cores: int, | ||
| 64 | +) -> None: | ||
| 65 | + require_ascend() | ||
| 66 | + out_config = _validate_cast_and_back(x, fmt, num_per_channels, round_sf, clamp_min_value) | ||
| 67 | + num_tokens, hidden = x.shape | ||
| 39 | kernel = get_per_token_cast_and_cast_back_kernel_asc( | 68 | kernel = get_per_token_cast_and_cast_back_kernel_asc( |
| 40 | hidden=hidden, | 69 | hidden=hidden, |
| 41 | token_stride=x.stride(0), | 70 | token_stride=x.stride(0), |
| 42 | dtype=T.dtype(x.dtype), | 71 | dtype=T.dtype(x.dtype), |
| 43 | out_config=out_config, | 72 | out_config=out_config, |
| 73 | + n_cores=num_cores, | ||
| 44 | ) | 74 | ) |
| 45 | 75 | ||
| 46 | if should_print_kernel_source(): | 76 | if should_print_kernel_source(): |
| @@ -48,3 +78,8 @@ def per_token_cast_and_cast_back( | |||
| 48 | 78 | ||
| 49 | if num_tokens > 0: | 79 | if num_tokens > 0: |
| 50 | kernel(x) | 80 | kernel(x) |
| 81 | + | ||
| 82 | + | ||
| 83 | +@_cast_and_back_impl.register_fake | ||
| 84 | +def _cast_and_back_fake(x, fmt, num_per_channels, round_sf, clamp_min_value, num_cores): | ||
| 85 | + _validate_cast_and_back(x, fmt, num_per_channels, round_sf, clamp_min_value) | ||
| @@ -2,7 +2,6 @@ import tilelang | |||
| 2 | from tilelang.ascend import language as T | 2 | from tilelang.ascend import language as T |
| 3 | from tilelang.ascend.language import simd as S | 3 | from tilelang.ascend.language import simd as S |
| 4 | 4 | ||
| 5 | -from cann_ops_tilelang.config import get_num_vector_cores | ||
| 6 | from cann_ops_tilelang.utils import align_up as align | 5 | from cann_ops_tilelang.utils import align_up as align |
| 7 | 6 | ||
| 8 | from .common import CastOutputConfig | 7 | from .common import CastOutputConfig |
| @@ -102,6 +101,7 @@ def get_per_token_cast_and_cast_back_kernel_asc( | |||
| 102 | token_stride: int, | 101 | token_stride: int, |
| 103 | dtype: T.dtype, | 102 | dtype: T.dtype, |
| 104 | out_config: CastOutputConfig, | 103 | out_config: CastOutputConfig, |
| 104 | + n_cores: int, | ||
| 105 | ): | 105 | ): |
| 106 | num_per_channels = out_config.sf_block[1] | 106 | num_per_channels = out_config.sf_block[1] |
| 107 | is_bf16_in = dtype == T.bfloat16 | 107 | is_bf16_in = dtype == T.bfloat16 |
| @@ -120,7 +120,6 @@ def get_per_token_cast_and_cast_back_kernel_asc( | |||
| 120 | round_carry = 0x003FFFFF if is_fp4_out else 0x001FFFFF | 120 | round_carry = 0x003FFFFF if is_fp4_out else 0x001FFFFF |
| 121 | quant_max_exp = 0x01000000 if is_fp4_out else 0x04000000 | 121 | quant_max_exp = 0x01000000 if is_fp4_out else 0x04000000 |
| 122 | 122 | ||
| 123 | - n_cores = get_num_vector_cores() | ||
| 124 | num_stages = 4 if hidden >= 2560 else 3 | 123 | num_stages = 4 if hidden >= 2560 else 3 |
| 125 | 124 | ||
| 126 | block_m, block_k, tile_width = _get_block_shape(hidden, dtype, num_stages, num_per_channels, n_cores) | 125 | block_m, block_k, tile_width = _get_block_shape(hidden, dtype, num_stages, num_per_channels, n_cores) |
| @@ -189,3 +189,120 @@ def test_per_token_cast_compile_then_cast_back(fmt): | |||
| 189 | restored = cann_ops_tilelang.per_token_cast_back((data, sf), 'fp32', num_per_channels=32) | 189 | restored = cann_ops_tilelang.per_token_cast_back((data, sf), 'fp32', num_per_channels=32) |
| 190 | assert_equal(restored, cast_back(reference, 'fp32', (1, 32))) | 190 | assert_equal(restored, cast_back(reference, 'fp32', (1, 32))) |
| 191 | del data, sf | 191 | del data, sf |
| 192 | + | ||
| 193 | + | ||
| 194 | +@pytest.mark.parametrize( | ||
| 195 | + ('entry', 'fmt', 'out_dtype', 'block_size', 'packed'), | ||
| 196 | + [('block', 'e4m3', 'fp32', (32, 32), False), ('token', 'e2m1', 'bf16', (1, 32), True)], | ||
| 197 | +) | ||
| 198 | +@torch.inference_mode() | ||
| 199 | +def test_cast_back_compile(entry, fmt, out_dtype, block_size, packed): | ||
| 200 | + from tests.testing.quant import call_quant_graph, check_quant_graph_replay, compile_quant_graph | ||
| 201 | + | ||
| 202 | + compiled = compile_quant_graph( | ||
| 203 | + cann_ops_tilelang.cast_back if entry == 'block' else cann_ops_tilelang.per_token_cast_back | ||
| 204 | + ) | ||
| 205 | + source = None | ||
| 206 | + records = [] | ||
| 207 | + for step in range(3): | ||
| 208 | + raw = torch.randn((33, 256), dtype=torch.bfloat16, device=get_device()) | ||
| 209 | + refreshed = cast(raw, fmt, block_size, round_sf=True, use_packed_ue8m0=packed) | ||
| 210 | + if step == 1: | ||
| 211 | + source[0].view(torch.uint8).copy_(refreshed[0].view(torch.uint8)) | ||
| 212 | + source[1].copy_(refreshed[1]) | ||
| 213 | + else: | ||
| 214 | + source = refreshed | ||
| 215 | + snapshots = tuple(value.clone() for value in source) | ||
| 216 | + expected = cast_back(source, out_dtype, block_size) | ||
| 217 | + scale_arg = block_size if entry == 'block' else block_size[1] | ||
| 218 | + result, events = call_quant_graph(compiled, source, out_dtype, scale_arg) | ||
| 219 | + records.append(events) | ||
| 220 | + assert_equal(result, expected) | ||
| 221 | + for value, before in zip(source, snapshots, strict=True): | ||
| 222 | + assert_equal(value, before) | ||
| 223 | + del result | ||
| 224 | + check_quant_graph_replay(records) | ||
| 225 | + from cann_ops_tilelang.quant.common import get_quant_vector_cores | ||
| 226 | + | ||
| 227 | + torch.library.opcheck( | ||
| 228 | + torch.ops.cann_ops_tilelang.cast_back.default, | ||
| 229 | + (source, out_dtype, block_size, None, get_quant_vector_cores()), | ||
| 230 | + # TorchNPU schema_check uses allclose, which does not support FP8 inputs. | ||
| 231 | + # FP4 runs the same schema check; both cases check read-only inputs above. | ||
| 232 | + test_utils=('test_faketensor',) if fmt == 'e4m3' else ('test_schema', 'test_faketensor'), | ||
| 233 | + ) | ||
| 234 | + | ||
| 235 | + | ||
| 236 | +@torch.inference_mode() | ||
| 237 | +def test_unpack_fp4_compile(): | ||
| 238 | + from tests.testing.quant import call_quant_graph, check_quant_graph_replay, compile_quant_graph | ||
| 239 | + | ||
| 240 | + # Independent table for all sixteen E2M1 codes, including signed zero. | ||
| 241 | + levels = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6], device=get_device()) | ||
| 242 | + compiled = compile_quant_graph(cann_ops_tilelang.unpack_from_e2m1fn_x2) | ||
| 243 | + values = torch.arange(256, device=get_device(), dtype=torch.int32) | ||
| 244 | + x = values.to(torch.uint8).view(torch.int8).reshape(16, 16) | ||
| 245 | + records = [] | ||
| 246 | + for step in range(3): | ||
| 247 | + if step == 1: | ||
| 248 | + x.copy_(((values + 17) % 256).to(torch.uint8).view(torch.int8).reshape_as(x)) | ||
| 249 | + elif step == 2: | ||
| 250 | + x = torch.flip(x, dims=[0]).clone() | ||
| 251 | + before = x.clone() | ||
| 252 | + codes = x.to(torch.int32) | ||
| 253 | + expected = torch.stack((levels[codes & 15], levels[(codes >> 4) & 15]), dim=-1).reshape(16, 32) | ||
| 254 | + if step == 1: | ||
| 255 | + # Reuse temporary host scalar buffers before replaying the captured graph. | ||
| 256 | + host_buffers = [torch.full((), 7, dtype=torch.int32, pin_memory=True) for _ in range(128)] | ||
| 257 | + result, events = call_quant_graph(compiled, x) | ||
| 258 | + if step == 1: | ||
| 259 | + del host_buffers | ||
| 260 | + records.append(events) | ||
| 261 | + assert_equal(result, expected) | ||
| 262 | + assert_equal(x, before) | ||
| 263 | + del result | ||
| 264 | + check_quant_graph_replay(records) | ||
| 265 | + | ||
| 266 | + | ||
| 267 | +@pytest.mark.parametrize('entry', ['block', 'token']) | ||
| 268 | +@torch.inference_mode() | ||
| 269 | +def test_cast_back_compile_special_values(entry): | ||
| 270 | + from tests.testing.quant import call_quant_graph, check_quant_graph_replay, compile_quant_graph | ||
| 271 | + | ||
| 272 | + block = (32, 32) if entry == 'block' else (1, 32) | ||
| 273 | + compiled = compile_quant_graph( | ||
| 274 | + cann_ops_tilelang.cast_back if entry == 'block' else cann_ops_tilelang.per_token_cast_back | ||
| 275 | + ) | ||
| 276 | + source = None | ||
| 277 | + records = [] | ||
| 278 | + for step in range(3): | ||
| 279 | + if entry == 'block': | ||
| 280 | + data = torch.full((33, 256), float(step), dtype=torch.bfloat16, device=get_device()) | ||
| 281 | + data = data.to(torch.float8_e4m3fn) | ||
| 282 | + else: | ||
| 283 | + data = torch.full((33, 128), 0x22 * step, dtype=torch.int8, device=get_device()) | ||
| 284 | + scale = torch.full(((33 + block[0] - 1) // block[0], 8), 1.0, dtype=torch.float32, device=get_device()) | ||
| 285 | + if step: | ||
| 286 | + scale.fill_(2.0 if step == 1 else 0.5) | ||
| 287 | + row = 0 if step == 1 else scale.shape[0] - 1 | ||
| 288 | + scale[row, 0] = float('nan') | ||
| 289 | + scale[row, 1] = float('inf') | ||
| 290 | + scale[row, 2] = -float('inf') | ||
| 291 | + if step == 1: | ||
| 292 | + source[0].view(torch.uint8).copy_(data.view(torch.uint8)) | ||
| 293 | + source[1].copy_(scale) | ||
| 294 | + else: | ||
| 295 | + source = (data, scale) | ||
| 296 | + snapshots = tuple(value.clone() for value in source) | ||
| 297 | + expected = cast_back(source, 'fp32', block) | ||
| 298 | + scale_arg = block if entry == 'block' else block[1] | ||
| 299 | + result, events = call_quant_graph(compiled, source, 'fp32', scale_arg) | ||
| 300 | + records.append(events) | ||
| 301 | + # NaN payloads are unspecified; finite values, signed zeros and infinities remain bit-exact. | ||
| 302 | + nan_mask = torch.isnan(expected) | ||
| 303 | + assert_equal(torch.isnan(result), nan_mask) | ||
| 304 | + assert_equal(result[~nan_mask], expected[~nan_mask]) | ||
| 305 | + for value, before in zip(source, snapshots, strict=True): | ||
| 306 | + assert_equal(value, before) | ||
| 307 | + del result | ||
| 308 | + check_quant_graph_replay(records) | ||
| @@ -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 | + ) | ||
| @@ -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 |
| @@ -155,3 +156,122 @@ def test_per_channel_cast_benchmark(benchmark_timer, benchmark_record, params): | |||
| 155 | time_us=t_us, | 156 | time_us=t_us, |
| 156 | bandwidth_gbs=num_bytes / t_us / 1e3, | 157 | bandwidth_gbs=num_bytes / t_us / 1e3, |
| 157 | ) | 158 | ) |
| 159 | + | ||
| 160 | + | ||
| 161 | +def _refresh_per_channel_cast_input(source, raw, step, prequantized): | ||
| 162 | + refreshed = cast(raw, 'e4m3', (1, 128), round_sf=True) if prequantized else raw | ||
| 163 | + if step != 1: | ||
| 164 | + return refreshed | ||
| 165 | + if prequantized: | ||
| 166 | + source[0].view(torch.uint8).copy_(refreshed[0].view(torch.uint8)) | ||
| 167 | + source[1].copy_(refreshed[1]) | ||
| 168 | + else: | ||
| 169 | + source.copy_(refreshed) | ||
| 170 | + return source | ||
| 171 | + | ||
| 172 | + | ||
| 173 | +def _check_per_channel_cast_registration(source, token_group, kw): | ||
| 174 | + from cann_ops_tilelang.config import get_amax_clamp_for_quant | ||
| 175 | + | ||
| 176 | + data, scale = source if isinstance(source, tuple) else (source, None) | ||
| 177 | + from cann_ops_tilelang.quant.common import get_quant_vector_cores | ||
| 178 | + | ||
| 179 | + torch.library.opcheck( | ||
| 180 | + torch.ops.cann_ops_tilelang.per_channel_cast.default, | ||
| 181 | + ( | ||
| 182 | + data, | ||
| 183 | + scale, | ||
| 184 | + 'e4m3', | ||
| 185 | + token_group, | ||
| 186 | + kw['round_sf'], | ||
| 187 | + kw['num_per_channels'], | ||
| 188 | + kw['use_packed_ue8m0'], | ||
| 189 | + get_amax_clamp_for_quant('e4m3'), | ||
| 190 | + get_quant_vector_cores(), | ||
| 191 | + ), | ||
| 192 | + # FP8 allclose is unavailable in TorchNPU's standard schema checker. | ||
| 193 | + # The BF16 case covers schema; both cases verify input bytes in the caller. | ||
| 194 | + test_utils=('test_faketensor',) if scale is not None else ('test_schema', 'test_faketensor'), | ||
| 195 | + ) | ||
| 196 | + | ||
| 197 | + | ||
| 198 | +@pytest.mark.parametrize('prequantized', [False, True]) | ||
| 199 | +@torch.inference_mode() | ||
| 200 | +def test_per_channel_cast_compile(prequantized): | ||
| 201 | + from cann_ops_tilelang.config import get_device | ||
| 202 | + from tests.testing.quant import call_quant_graph, check_quant_graph_replay, compile_quant_graph | ||
| 203 | + | ||
| 204 | + packed = prequantized | ||
| 205 | + token_group = 128 if prequantized else 32 | ||
| 206 | + # Two token groups keep every packed scale byte meaningful in this core case. | ||
| 207 | + tokens = 256 if prequantized else 128 | ||
| 208 | + channel_group = 128 if prequantized else None | ||
| 209 | + kw = dict(round_sf=packed, num_per_channels=channel_group, use_packed_ue8m0=packed) | ||
| 210 | + compiled = compile_quant_graph(cann_ops_tilelang.per_channel_cast) | ||
| 211 | + source = None | ||
| 212 | + records = [] | ||
| 213 | + for step in range(3): | ||
| 214 | + raw = torch.randn((tokens, 256), dtype=torch.bfloat16, device=get_device()) | ||
| 215 | + source = _refresh_per_channel_cast_input(source, raw, step, prequantized) | ||
| 216 | + inputs = source if isinstance(source, tuple) else (source,) | ||
| 217 | + snapshots = tuple(value.clone() for value in inputs) | ||
| 218 | + expected = cast( | ||
| 219 | + source, | ||
| 220 | + 'e4m3', | ||
| 221 | + (token_group, 1), | ||
| 222 | + x_block_size=(1, 128) if prequantized else None, | ||
| 223 | + round_sf=packed, | ||
| 224 | + use_packed_ue8m0=packed, | ||
| 225 | + ) | ||
| 226 | + result, events = call_quant_graph(compiled, source, 'e4m3', token_group, **kw) | ||
| 227 | + records.append(events) | ||
| 228 | + for actual, reference in zip(result, expected, strict=True): | ||
| 229 | + assert_equal(actual, reference) | ||
| 230 | + for value, before in zip(inputs, snapshots, strict=True): | ||
| 231 | + assert_equal(value, before) | ||
| 232 | + del result, actual | ||
| 233 | + check_quant_graph_replay(records) | ||
| 234 | + _check_per_channel_cast_registration(source, token_group, kw) | ||
| 235 | + | ||
| 236 | + | ||
| 237 | +@torch.inference_mode() | ||
| 238 | +def test_per_channel_cast_compile_special_values(): | ||
| 239 | + from cann_ops_tilelang.config import get_device | ||
| 240 | + from tests.testing.quant import call_quant_graph, check_quant_graph_replay, compile_quant_graph | ||
| 241 | + | ||
| 242 | + compiled = compile_quant_graph(cann_ops_tilelang.per_channel_cast) | ||
| 243 | + x = torch.zeros((128, 256), dtype=torch.bfloat16, device=get_device()) | ||
| 244 | + finite_columns = torch.ones(256, dtype=torch.bool, device=x.device) | ||
| 245 | + finite_columns[0] = finite_columns[32] = finite_columns[64] = False | ||
| 246 | + records = [] | ||
| 247 | + for step in range(3): | ||
| 248 | + if step == 1: | ||
| 249 | + x.fill_(0.5) | ||
| 250 | + elif step == 2: | ||
| 251 | + x = torch.full_like(x, -1.25) | ||
| 252 | + if step: | ||
| 253 | + x[:, 96:128].zero_() | ||
| 254 | + row = 0 if step == 1 else 32 | ||
| 255 | + x[row, 0] = float('nan') | ||
| 256 | + x[row, 32] = float('inf') | ||
| 257 | + x[row, 64] = -float('inf') | ||
| 258 | + before = x.clone() | ||
| 259 | + reference = cast(x, 'e4m3', (32, 1), round_sf=True) | ||
| 260 | + eager = cann_ops_tilelang.per_channel_cast(x, 'e4m3', 32, round_sf=True) | ||
| 261 | + result, events = call_quant_graph(compiled, x, 'e4m3', 32, round_sf=True) | ||
| 262 | + records.append(events) | ||
| 263 | + for actual, expected in zip(result, eager, strict=True): | ||
| 264 | + assert_equal(actual, expected) | ||
| 265 | + # Existing non-finite quantization differs from cast_ref. Check graph/eager above; | ||
| 266 | + # independently check every column for zero input, and finite-only columns otherwise. | ||
| 267 | + if step == 0: | ||
| 268 | + for actual, expected in zip(result, reference, strict=True): | ||
| 269 | + assert_equal(actual, expected) | ||
| 270 | + else: | ||
| 271 | + assert_equal( | ||
| 272 | + result[0].view(torch.uint8)[:, finite_columns], reference[0].view(torch.uint8)[:, finite_columns] | ||
| 273 | + ) | ||
| 274 | + assert_equal(result[1][:, finite_columns], reference[1][:, finite_columns]) | ||
| 275 | + assert_equal(x, before) | ||
| 276 | + del result, actual, eager, reference | ||
| 277 | + check_quant_graph_replay(records) | ||
| @@ -103,3 +103,74 @@ def test_per_token_cast_and_cast_back_benchmark(benchmark_timer, benchmark_recor | |||
| 103 | time_us=t_us, | 103 | time_us=t_us, |
| 104 | bandwidth_gbs=num_bytes / t_us / 1e3, | 104 | bandwidth_gbs=num_bytes / t_us / 1e3, |
| 105 | ) | 105 | ) |
| 106 | + | ||
| 107 | + | ||
| 108 | +@torch.inference_mode() | ||
| 109 | +def test_per_token_cast_and_cast_back_compile(): | ||
| 110 | + from cann_ops_tilelang.config import get_amax_clamp_for_quant | ||
| 111 | + from tests.testing.quant import call_quant_graph, check_quant_graph_replay, compile_quant_graph | ||
| 112 | + | ||
| 113 | + compiled = compile_quant_graph(cann_ops_tilelang.per_token_cast_and_cast_back) | ||
| 114 | + backing = torch.randn((31, 512), dtype=torch.bfloat16, device=get_device()) | ||
| 115 | + records = [] | ||
| 116 | + for step in range(3): | ||
| 117 | + if step == 1: | ||
| 118 | + backing.neg_().add_(0.125) | ||
| 119 | + elif step == 2: | ||
| 120 | + backing = torch.randn_like(backing) | ||
| 121 | + x = backing[:, :256] | ||
| 122 | + before = backing.clone() | ||
| 123 | + expected = cast_back(cast(x, 'e2m1', (1, 32), round_sf=True), 'bf16', (1, 32)) | ||
| 124 | + pointer, stride, offset = x.data_ptr(), x.stride(), x.storage_offset() | ||
| 125 | + result, events = call_quant_graph(compiled, x, 'e2m1', 32, round_sf=True) | ||
| 126 | + records.append(events) | ||
| 127 | + assert result is None | ||
| 128 | + assert_equal(x, expected, check_stride=False) | ||
| 129 | + assert_equal(backing[:, 256:], before[:, 256:]) | ||
| 130 | + assert (x.data_ptr(), x.stride(), x.storage_offset()) == (pointer, stride, offset) | ||
| 131 | + check_quant_graph_replay(records) | ||
| 132 | + from cann_ops_tilelang.quant.common import get_quant_vector_cores | ||
| 133 | + | ||
| 134 | + torch.library.opcheck( | ||
| 135 | + torch.ops.cann_ops_tilelang.per_token_cast_and_cast_back.default, | ||
| 136 | + (x, 'e2m1', 32, True, get_amax_clamp_for_quant('e2m1'), get_quant_vector_cores()), | ||
| 137 | + test_utils=('test_schema', 'test_faketensor'), | ||
| 138 | + ) | ||
| 139 | + | ||
| 140 | + | ||
| 141 | +@torch.inference_mode() | ||
| 142 | +def test_per_token_cast_and_cast_back_compile_special_values(): | ||
| 143 | + from tests.testing.quant import call_quant_graph, check_quant_graph_replay, compile_quant_graph | ||
| 144 | + | ||
| 145 | + compiled = compile_quant_graph(cann_ops_tilelang.per_token_cast_and_cast_back) | ||
| 146 | + backing = torch.zeros((31, 512), dtype=torch.bfloat16, device=get_device()) | ||
| 147 | + records = [] | ||
| 148 | + for step in range(3): | ||
| 149 | + if step == 1: | ||
| 150 | + backing.fill_(0.5) | ||
| 151 | + elif step == 2: | ||
| 152 | + backing = torch.full_like(backing, -1.25) | ||
| 153 | + x = backing[:, :256] | ||
| 154 | + finite_mask = torch.ones_like(x, dtype=torch.bool) | ||
| 155 | + if step: | ||
| 156 | + x[:, 96:128].zero_() | ||
| 157 | + row = 0 if step == 1 else x.shape[0] - 1 | ||
| 158 | + x[row, 0] = float('nan') | ||
| 159 | + x[row, 32] = float('inf') | ||
| 160 | + x[row, 64] = -float('inf') | ||
| 161 | + finite_mask[row, :96] = False | ||
| 162 | + before = backing.clone() | ||
| 163 | + reference = cast_back(cast(x, 'e2m1', (1, 32), round_sf=True), 'bf16', (1, 32)) | ||
| 164 | + eager = x.clone() | ||
| 165 | + cann_ops_tilelang.per_token_cast_and_cast_back(eager, 'e2m1', 32, round_sf=True) | ||
| 166 | + pointer, stride, offset = x.data_ptr(), x.stride(), x.storage_offset() | ||
| 167 | + result, events = call_quant_graph(compiled, x, 'e2m1', 32, round_sf=True) | ||
| 168 | + records.append(events) | ||
| 169 | + assert result is None | ||
| 170 | + assert_equal(x, eager, check_stride=False) | ||
| 171 | + # Non-finite quantization has an existing cast_ref discrepancy. Keep the | ||
| 172 | + # independent check bit-exact on zero input and all finite-only groups. | ||
| 173 | + assert_equal(x[finite_mask], reference[finite_mask]) | ||
| 174 | + assert_equal(backing[:, 256:], before[:, 256:]) | ||
| 175 | + assert (x.data_ptr(), x.stride(), x.storage_offset()) == (pointer, stride, offset) | ||
| 176 | + check_quant_graph_replay(records) | ||
| @@ -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}') | ||