feat(quant): migrate quant operators-ops-tilelang-AtomGit

Msrc/cann_ops_tilelang/__init__.py+32-0

	@@ -32,6 +32,21 @@ __all__ = [
32	'moe_topk_gate',	32	'moe_topk_gate',
33	'normalize_weight',	33	'normalize_weight',
34	'topk_gate',	34	'topk_gate',
		35	+ 'cast_back',
		36	+ 'per_block_cast',
		37	+ 'per_block_cast_lossless',
		38	+ 'per_block_cast_with_precomputed_sf',
		39	+ 'per_block_cast_with_sf_only',
		40	+ 'per_channel_cast',
		41	+ 'per_token_cast',
		42	+ 'per_token_cast_and_cast_back',
		43	+ 'per_token_cast_back',
		44	+ 'per_token_cast_with_precomputed_sf',
		45	+ 'per_token_cast_with_sf_only',
		46	+ 'swiglu_backward',
		47	+ 'swiglu_forward',
		48	+ 'swiglu_forward_and_per_token_cast',
		49	+ 'unpack_from_e2m1fn_x2',
35	]	50	]
36		51
37	from .engram import (	52	from .engram import (
	@@ -67,3 +82,20 @@ from .mhc import (
67	mhc_sinkhorn_forward,	82	mhc_sinkhorn_forward,
68	)	83	)
69	from .moe import moe_topk_gate, normalize_weight, topk_gate	84	from .moe import moe_topk_gate, normalize_weight, topk_gate
		85	+from .quant import (
		86	+ cast_back,
		87	+ per_block_cast,
		88	+ per_block_cast_lossless,
		89	+ per_block_cast_with_precomputed_sf,
		90	+ per_block_cast_with_sf_only,
		91	+ per_channel_cast,
		92	+ per_token_cast,
		93	+ per_token_cast_and_cast_back,
		94	+ per_token_cast_back,
		95	+ per_token_cast_with_precomputed_sf,
		96	+ per_token_cast_with_sf_only,
		97	+ swiglu_backward,
		98	+ swiglu_forward,
		99	+ swiglu_forward_and_per_token_cast,
		100	+ unpack_from_e2m1fn_x2,
		101	+)

Msrc/cann_ops_tilelang/config.py+82-1

	@@ -13,6 +13,15 @@ import torch
13	_ASCEND_DEVICE_NODE = '/dev/davinci_manager'	13	_ASCEND_DEVICE_NODE = '/dev/davinci_manager'
14	_NUM_VECTOR_CORES_ENV = 'OPS_TILELANG_NUM_VECTOR_CORES'	14	_NUM_VECTOR_CORES_ENV = 'OPS_TILELANG_NUM_VECTOR_CORES'
15	_PRINT_KERNEL_SOURCE_ENV = 'OPS_TILELANG_PRINT_KERNEL_SOURCE'	15	_PRINT_KERNEL_SOURCE_ENV = 'OPS_TILELANG_PRINT_KERNEL_SOURCE'
		16	+_DEFAULT_TOKEN_ALIGNMENT = 128
		17	+_DEFAULT_QUANT_CLAMP_MIN_VALUES = {
		18	+ 'e4m3': 1e-4,
		19	+ 'e2m1': 6.0 * (2.0**-126),
		20	+}
		21	+
		22	+_num_ai_cores_override: int \| None = None
		23	+_token_alignment = _DEFAULT_TOKEN_ALIGNMENT
		24	+_quant_clamp_min_values = _DEFAULT_QUANT_CLAMP_MIN_VALUES.copy()
16		25
17		26
18	@lru_cache(maxsize=1)	27	@lru_cache(maxsize=1)
	@@ -49,7 +58,30 @@ def get_num_ai_cores(device_index: int \| None = None) -> int:
49	cube_cores = getattr(properties, 'cube_core_num', None)	58	cube_cores = getattr(properties, 'cube_core_num', None)
50	if cube_cores is None:	59	if cube_cores is None:
51	raise RuntimeError('The current torch_npu build does not expose cube_core_num.')	60	raise RuntimeError('The current torch_npu build does not expose cube_core_num.')
52	- return int(cube_cores)	61	+ detected = int(cube_cores)
		62	+ if device_index is None and _num_ai_cores_override is not None:
		63	+ return _num_ai_cores_override
		64	+ return detected
		65	+
		66	+
		67	+def set_num_ai_cores(num_ai_cores: int) -> None:
		68	+ """Override the AI core count used for kernel tuning in tests and experiments."""
		69	+
		70	+ if isinstance(num_ai_cores, bool) or not isinstance(num_ai_cores, int):
		71	+ raise TypeError('num_ai_cores must be an integer.')
		72	+
		73	+ global _num_ai_cores_override
		74	+ previous = _num_ai_cores_override
		75	+ _num_ai_cores_override = None
		76	+ get_num_ai_cores.cache_clear()
		77	+ detected = get_num_ai_cores()
		78	+ if not 0 < num_ai_cores <= detected:
		79	+ _num_ai_cores_override = previous
		80	+ get_num_ai_cores.cache_clear()
		81	+ raise ValueError(f'num_ai_cores must be in [1, {detected}], got {num_ai_cores}.')
		82	+ _num_ai_cores_override = num_ai_cores
		83	+ get_num_ai_cores.cache_clear()
		84	+ get_num_vector_cores.cache_clear()
53		85
54		86
55	@cache	87	@cache
	@@ -77,6 +109,55 @@ def should_print_kernel_source() -> bool:
77	return os.getenv(_PRINT_KERNEL_SOURCE_ENV, '0').lower() in {'1', 'true', 'yes', 'on'}	109	return os.getenv(_PRINT_KERNEL_SOURCE_ENV, '0').lower() in {'1', 'true', 'yes', 'on'}
78		110
79		111
		112	+def get_token_alignment() -> int:
		113	+ """Return the token alignment used by fused quantization kernels."""
		114	+
		115	+ return _token_alignment
		116	+
		117	+
		118	+def set_token_alignment(token_alignment: int) -> None:
		119	+ """Set the token alignment used by fused quantization kernels."""
		120	+
		121	+ if isinstance(token_alignment, bool) or not isinstance(token_alignment, int):
		122	+ raise TypeError('token_alignment must be an integer.')
		123	+ if token_alignment <= 0:
		124	+ raise ValueError(f'token_alignment must be positive, got {token_alignment}.')
		125	+ global _token_alignment
		126	+ _token_alignment = token_alignment
		127	+
		128	+
		129	+def _validate_quant_format(fmt: str) -> None:
		130	+ if fmt not in _quant_clamp_min_values:
		131	+ raise ValueError(f'Unsupported quantization format: {fmt}.')
		132	+
		133	+
		134	+def get_amax_clamp_for_quant(fmt: str) -> float:
		135	+ """Return the minimum magnitude used when computing quantization scales."""
		136	+
		137	+ _validate_quant_format(fmt)
		138	+ return _quant_clamp_min_values[fmt]
		139	+
		140	+
		141	+def set_amax_clamp_for_quant(fmt: str, clamp_min_value: float) -> None:
		142	+ """Set the minimum magnitude used when computing quantization scales."""
		143	+
		144	+ _validate_quant_format(fmt)
		145	+ value = float(clamp_min_value)
		146	+ if value <= 0:
		147	+ raise ValueError(f'clamp_min_value must be positive, got {clamp_min_value}.')
		148	+ _quant_clamp_min_values[fmt] = value
		149	+
		150	+
		151	+def reset_runtime_config() -> None:
		152	+ """Reset mutable quantization tuning settings to their defaults."""
		153	+
		154	+ global _num_ai_cores_override, _token_alignment, _quant_clamp_min_values
		155	+ _num_ai_cores_override = None
		156	+ _token_alignment = _DEFAULT_TOKEN_ALIGNMENT
		157	+ _quant_clamp_min_values = _DEFAULT_QUANT_CLAMP_MIN_VALUES.copy()
		158	+ clear_runtime_caches()
		159	+
		160	+
80	def clear_runtime_caches() -> None:	161	def clear_runtime_caches() -> None:
81	"""Clear lazily cached hardware queries, mainly for tests."""	162	"""Clear lazily cached hardware queries, mainly for tests."""
82		163

Asrc/cann_ops_tilelang/quant/__init__.py+29-0

@@ -0,0 +1,29 @@
	1	+"""Quantization operators for Ascend."""
	2	+
	3	+__all__ = [
	4	+ 'cast_back',
	5	+ 'per_block_cast',
	6	+ 'per_block_cast_lossless',
	7	+ 'per_block_cast_with_precomputed_sf',
	8	+ 'per_block_cast_with_sf_only',
	9	+ 'per_channel_cast',
	10	+ 'per_token_cast',
	11	+ 'per_token_cast_and_cast_back',
	12	+ 'per_token_cast_back',
	13	+ 'per_token_cast_with_precomputed_sf',
	14	+ 'per_token_cast_with_sf_only',
	15	+ 'swiglu_backward',
	16	+ 'swiglu_forward',
	17	+ 'swiglu_forward_and_per_token_cast',
	18	+ 'unpack_from_e2m1fn_x2',
	19	+]
	20	+
	21	+from .cast_back import cast_back, per_token_cast_back
	22	+from .common import unpack_from_e2m1fn_x2
	23	+from .per_block_cast import per_block_cast, per_block_cast_with_precomputed_sf, per_block_cast_with_sf_only
	24	+from .per_block_cast_lossless import per_block_cast_lossless
	25	+from .per_channel_cast import per_channel_cast
	26	+from .per_token_cast import per_token_cast, per_token_cast_with_precomputed_sf, per_token_cast_with_sf_only
	27	+from .per_token_cast_and_cast_back import per_token_cast_and_cast_back
	28	+from .swiglu_backward import swiglu_backward
	29	+from .swiglu_forward import swiglu_forward, swiglu_forward_and_per_token_cast

Asrc/cann_ops_tilelang/quant/cast_back.py+77-0

@@ -0,0 +1,77 @@
	1	+import torch
	2	+from tilelang.ascend import language as T
	3	+
	4	+from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
	5	+
	6	+from .cast_back_kernel import get_cast_back_kernel_ascend
	7	+from .common import get_cast_input_and_config, get_logical_hidden
	8	+from .types import QuantTensor
	9	+
	10	+
	11	+def cast_back(
	12	+ x: QuantTensor,
	13	+ fmt: str,
	14	+ x_block_size: tuple[int, int],
	15	+ x_special_fmt: str \| None = None,
	16	+) -> torch.Tensor:
	17	+ """Dequantize an FP8/FP4 tensor back to BF16 or FP32.
	18	+
	19	+ Args:
	20	+ x: Quantized tensor pair ``(data, sf_factors)``.
	21	+ fmt: Target output format, either ``'bf16'`` or ``'fp32'``.
	22	+ x_block_size: Scaling block size as ``(num_per_tokens, num_per_channels)``.
	23	+ x_special_fmt: Reserved, must be ``None``.
	24	+
	25	+ Returns:
	26	+ Dequantized tensor in the requested format.
	27	+ """
	28	+ require_ascend()
	29	+ assert fmt in ('bf16', 'fp32')
	30	+ out_dtype = torch.bfloat16 if fmt == 'bf16' else torch.float32
	31	+
	32	+ x, x_sf = x
	33	+ if x.device.type != 'npu' or x_sf.device.type != 'npu':
	34	+ raise ValueError('x and its scaling factors must be NPU tensors.')
	35	+
	36	+ assert x_special_fmt is None, f'Unsupported special format: {x_special_fmt}'
	37	+
	38	+ num_tokens, hidden = x.shape
	39	+ assert x_sf.dim() == 2
	40	+ assert x_sf.dtype in (torch.int16, torch.int32, torch.float32)
	41	+
	42	+ hidden = get_logical_hidden(hidden, x.dtype)
	43	+ assert hidden % 64 == 0
	44	+
	45	+ x, x_sf, in_config = get_cast_input_and_config((x, x_sf), x_block_size)
	46	+ out = torch.empty((num_tokens, hidden), dtype=out_dtype, device=x.device)
	47	+ if num_tokens == 0:
	48	+ return out
	49	+
	50	+ kernel = get_cast_back_kernel_ascend(hidden=hidden, in_config=in_config, out_dtype=T.dtype(out_dtype))
	51	+
	52	+ if should_print_kernel_source():
	53	+ print(kernel.get_kernel_source())
	54	+
	55	+ kernel(x, x_sf, out)
	56	+
	57	+ return out
	58	+
	59	+
	60	+def per_token_cast_back(
	61	+ x: QuantTensor,
	62	+ fmt: str,
	63	+ num_per_channels: int,
	64	+ x_special_fmt: str \| None = None,
	65	+) -> torch.Tensor:
	66	+ """Dequantize an FP8/FP4 tensor back to BF16 or FP32 with per-token scaling.
	67	+
	68	+ Args:
	69	+ x: Quantized tensor pair ``(data, sf_factors)``.
	70	+ fmt: Target output format, either ``'bf16'`` or ``'fp32'``.
	71	+ num_per_channels: Number of channels per scaling block.
	72	+ x_special_fmt: Reserved, must be ``None``.
	73	+
	74	+ Returns:
	75	+ Dequantized tensor in the requested format.
	76	+ """
	77	+ return cast_back(x, fmt, (1, num_per_channels), x_special_fmt=x_special_fmt)

Asrc/cann_ops_tilelang/quant/cast_back_kernel.py+356-0

@@ -0,0 +1,356 @@
	1	+import math
	2	+
	3	+import tilelang
	4	+from tilelang.ascend import language as T
	5	+from tilelang.ascend.language import simd as S
	6	+
	7	+from cann_ops_tilelang.config import get_num_vector_cores
	8	+from cann_ops_tilelang.utils import align_up as align
	9	+from cann_ops_tilelang.utils import ceil_div
	10	+
	11	+from .common import CastInputConfig, get_packed_ue8m0_pack_factor, get_sf_shape
	12	+
	13	+
	14	+@tilelang.jit()
	15	+def get_cast_back_kernel_ascend(
	16	+ hidden: int,
	17	+ in_config: CastInputConfig,
	18	+ out_dtype: T.dtype = T.bfloat16,
	19	+):
	20	+ n_cores = get_num_vector_cores()
	21	+
	22	+ 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
	24	+ assert (num_per_tokens, num_per_channels) in ((1, 32), (1, 128), (32, 1), (32, 32), (128, 1), (128, 128))
	25	+ assert not ((num_per_tokens, num_per_channels) == (32, 1) and in_config.use_tma_aligned_col_major_sf)
	26	+
	27	+ use_e2b_scale = (
	28	+ num_per_tokens == 1
	29	+ and num_per_channels == 32
	30	+ and in_config.use_packed_ue8m0
	31	+ and hidden % 256 == 0
	32	+ and in_config.dtype == T.float4_e2m1fn
	33	+ and out_dtype != T.float32
	34	+ )
	35	+ block_k_align = 128 if in_config.dtype == T.float4_e2m1fn else 64
	36	+ if num_per_tokens == 1:
	37	+ block_k = max(math.gcd(hidden, 256 if use_e2b_scale else 512 // out_dtype.bytes), block_k_align)
	38	+ block_m_limit = 64 if use_e2b_scale or out_dtype != T.float32 else 128
	39	+ else:
	40	+ block_k = max(math.gcd(hidden, 256), block_k_align)
	41	+ block_m_limit = 256 // out_dtype.bytes
	42	+ if num_per_tokens == 128 and out_dtype == T.float32:
	43	+ block_k = max(math.gcd(hidden, 128), block_k_align)
	44	+ block_m = 128 if num_per_tokens == 128 else min(block_m_limit, align(8 * ceil_div(hidden, block_k), 64))
	45	+
	46	+ num_sf_rows = ceil_div(block_m, num_per_tokens)
	47	+ num_sf_cols = ceil_div(block_k, num_per_channels)
	48	+
	49	+ sf_raw_rows, sf_raw_cols = get_sf_shape((block_m, block_k), in_config)
	50	+ sf_raw_cols //= pack_factor # uint8 -> uint16 if packed
	51	+ sf_raw_dtype = T.uint16 if in_config.use_packed_ue8m0 else in_config.sf_dtype
	52	+ sf_ub_dtype = T.bfloat16 if use_e2b_scale else T.float32
	53	+
	54	+ if num_per_tokens == 128:
	55	+ if hidden >= 512 and block_k == 64 and out_dtype != T.float32:
	56	+ num_stages = 3
	57	+ else:
	58	+ num_stages = 2 if hidden >= 512 else (2 if num_per_channels == 128 and out_dtype == T.float32 else 1)
	59	+ else:
	60	+ num_stages = 3 if num_per_tokens == 1 and out_dtype != T.float32 else 2
	61	+
	62	+ num_tokens = T.dynamic('num_tokens')
	63	+ sf_shape = get_sf_shape((num_tokens, hidden), in_config)
	64	+ sf_stride = T.dynamic('sf_stride')
	65	+ num_token_tiles = T.ceildiv(num_tokens, block_m)
	66	+ num_hidden_tiles = T.ceildiv(hidden, block_k)
	67	+ num_hidden_sf = ceil_div(hidden, num_per_channels)
	68	+ hidden_tail_cols = hidden % block_k
	69	+
	70	+ @T.macro
	71	+ def load_scale_value(x_sf_raw: T.Tensor, token_index: int, channel_index: int):
	72	+ if in_config.use_packed_ue8m0:
	73	+ if in_config.sf_col_pack:
	74	+ packed = x_sf_raw[token_index // pack_factor, channel_index]
	75	+ packed_index = token_index
	76	+ elif in_config.use_tma_aligned_col_major_sf:
	77	+ packed = x_sf_raw[channel_index // pack_factor, token_index]
	78	+ packed_index = channel_index
	79	+ else:
	80	+ packed = x_sf_raw[token_index, channel_index // pack_factor]
	81	+ packed_index = channel_index
	82	+ exponent = T.if_then_else(packed_index % pack_factor == 0, packed & 0xFF, packed >> 8)
	83	+ return T.reinterpret(T.uint32(exponent) << 23, T.float32)
	84	+ elif in_config.use_tma_aligned_col_major_sf:
	85	+ return x_sf_raw[channel_index, token_index]
	86	+ else:
	87	+ return x_sf_raw[token_index, channel_index]
	88	+
	89	+ @T.macro
	90	+ def load_scale_tile(x_sf_raw: T.Tensor, sf_raw_ub, sf_ub, token_base: int, channel_base: int):
	91	+ if not in_config.use_packed_ue8m0 and not in_config.use_tma_aligned_col_major_sf:
	92	+ # Scale rows can be shorter than one MTE block and their dynamic
	93	+ # GM stride is not necessarily 32-byte aligned (for example,
	94	+ # hidden=576 and sf_block=(32, 32)). Load them on the vector cores
	95	+ # exactly as the 128-channel path already does.
	96	+ with T.SimtVF(threads=num_sf_rows * num_sf_cols):
	97	+ for idx in T.Parallel(num_sf_rows * num_sf_cols):
	98	+ sf_i = idx // num_sf_cols
	99	+ sf_j = idx % num_sf_cols
	100	+ if token_base + sf_i < sf_shape[0]:
	101	+ sf_ub[sf_i, sf_j] = x_sf_raw[token_base + sf_i, channel_base + sf_j]
	102	+ elif num_per_channels == 128:
	103	+ with T.SimtVF(threads=num_sf_rows * num_sf_cols):
	104	+ for idx in T.Parallel(num_sf_rows * num_sf_cols):
	105	+ sf_i = idx // num_sf_cols
	106	+ sf_j = idx % num_sf_cols
	107	+ sf_ub[sf_i, sf_j] = load_scale_value(x_sf_raw, token_base + sf_i, channel_base + sf_j)
	108	+ else:
	109	+ if in_config.use_tma_aligned_col_major_sf:
	110	+ source = x_sf_raw[channel_base // pack_factor, token_base]
	111	+ T.copy(source, sf_raw_ub[:sf_raw_rows, :sf_raw_cols])
	112	+ else:
	113	+ source_token = token_base // pack_factor if in_config.sf_col_pack else token_base
	114	+ source_channel = (
	115	+ channel_base // pack_factor
	116	+ if in_config.use_packed_ue8m0 and not in_config.sf_col_pack
	117	+ else channel_base
	118	+ )
	119	+ T.copy(x_sf_raw[source_token, source_channel], sf_raw_ub)
	120	+
	121	+ with T.SimtVF(threads=num_sf_rows * num_sf_cols):
	122	+ for idx in T.Parallel(num_sf_rows * num_sf_cols):
	123	+ sf_i = idx // num_sf_cols
	124	+ sf_j = idx % num_sf_cols
	125	+ raw_i, raw_j = (sf_j, sf_i) if in_config.use_tma_aligned_col_major_sf else (sf_i, sf_j)
	126	+ if in_config.use_packed_ue8m0:
	127	+ if in_config.sf_col_pack or in_config.use_tma_aligned_col_major_sf:
	128	+ packed = sf_raw_ub[raw_i // pack_factor, raw_j]
	129	+ else:
	130	+ packed = sf_raw_ub[raw_i, raw_j // pack_factor]
	131	+ packed_index = token_base + sf_i if in_config.sf_col_pack else channel_base + sf_j
	132	+ odd = packed_index % pack_factor == 1
	133	+ exponent = T.if_then_else(odd, packed >> 8, packed & 0xFF)
	134	+ if use_e2b_scale:
	135	+ scale_bits = T.uint16(exponent) << 7
	136	+ sf_ub[sf_i, sf_j] = T.reinterpret(scale_bits, T.bfloat16)
	137	+ else:
	138	+ scale_bits = T.uint32(exponent) << 23
	139	+ sf_ub[sf_i, sf_j] = T.reinterpret(scale_bits, T.float32)
	140	+ else:
	141	+ sf_ub[sf_i, sf_j] = sf_raw_ub[raw_i, raw_j]
	142	+
	143	+ @T.macro
	144	+ def load_scale_tail(
	145	+ x_sf_raw: T.Tensor,
	146	+ sf_ub,
	147	+ token_base: int,
	148	+ channel_base: int,
	149	+ valid_sf_rows: int,
	150	+ valid_sf_cols: int,
	151	+ ):
	152	+ with T.SimtVF(threads=num_sf_rows * num_sf_cols):
	153	+ for idx in T.Parallel(num_sf_rows * num_sf_cols):
	154	+ sf_i = idx // num_sf_cols
	155	+ sf_j = idx % num_sf_cols
	156	+ if sf_i < valid_sf_rows and sf_j < valid_sf_cols:
	157	+ sf_ub[sf_i, sf_j] = load_scale_value(
	158	+ x_sf_raw,
	159	+ token_base + sf_i,
	160	+ channel_base + sf_j,
	161	+ )
	162	+ else:
	163	+ sf_ub[sf_i, sf_j] = 0.0
	164	+
	165	+ @T.macro
	166	+ def load_scale(sf_ub: T.SharedBuffer, group: int, col: int, mask_low_f32):
	167	+ if num_per_channels == 1:
	168	+ return S.vld(sf_ub[group, col])
	169	+ elif num_per_channels == 32:
	170	+ scale_low = S.vld(sf_ub[group, col // 32], dist='BRC_B32')
	171	+ scale_high = S.vld(sf_ub[group, col // 32 + 1], dist='BRC_B32')
	172	+ return S.vsel(scale_low, scale_high, mask_low_f32)
	173	+ else: # num_per_channels % 64 == 0
	174	+ return S.vld(sf_ub[group, col // num_per_channels], dist='BRC_B32')
	175	+
	176	+ @T.macro
	177	+ def store_out(out_ub: T.SharedBuffer, row: int, col: int, out_f32):
	178	+ if out_dtype == T.float32:
	179	+ S.vsts(out_ub[row, col], out_f32)
	180	+ else:
	181	+ out_bf16 = S.vcvt(out_f32, T.bfloat16)
	182	+ S.vsts(out_ub[row, col], out_bf16, dist='PK_B32')
	183	+
	184	+ @T.macro
	185	+ def dequant_fp8_chunk(x_ub, out_ub, row, col, scale):
	186	+ x_fp8 = S.vld(x_ub[row, col], dist='UNPK4_B8')
	187	+ x_f32 = S.vcvt(x_fp8, T.float32)
	188	+ store_out(out_ub, row, col, S.vmul(x_f32, scale))
	189	+
	190	+ @T.macro
	191	+ def dequant_fp4_chunk(x_ub, out_ub, row, col, scale_low, scale_high, zero_bf16):
	192	+ x_fp4 = S.vld(x_ub[row, col], dist='UNPK4_B8')
	193	+ x_bf16 = S.vcvt(x_fp4, T.bfloat16)
	194	+ x_uint_low, x_uint_high = S.vintlv(zero_bf16, x_bf16)
	195	+ x_low = T.reinterpret(x_uint_low, 'float32x64')
	196	+ x_high = T.reinterpret(x_uint_high, 'float32x64')
	197	+ store_out(out_ub, row, col, S.vmul(x_low, scale_low))
	198	+ store_out(out_ub, row, col + 64, S.vmul(x_high, scale_high))
	199	+
	200	+ @T.macro
	201	+ def transform_tile(x_ub: T.SharedBuffer, sf_ub: T.SharedBuffer, out_ub: T.SharedBuffer):
	202	+ with T.SimdVF():
	203	+ mask_low_f32 = S.pset(32, 'PAT_VL32') if num_per_channels == 32 else None
	204	+
	205	+ if use_e2b_scale:
	206	+ for row in T.serial(block_m):
	207	+ for chunk in T.serial(block_k // 256):
	208	+ col = chunk * 256
	209	+ scale_e2b = S.vld(sf_ub[row, chunk * 8], dist='E2B_B16')
	210	+ scale_bf16_low, scale_bf16_high = S.vintlv(scale_e2b, scale_e2b)
	211	+ x_low = S.vcvt(S.vld(x_ub[row, col], dist='UNPK4_B8'), T.bfloat16, part=0)
	212	+ x_high = S.vcvt(S.vld(x_ub[row, col + 128], dist='UNPK4_B8'), T.bfloat16, part=0)
	213	+ out_low = S.vmul(x_low, scale_bf16_low)
	214	+ out_high = S.vmul(x_high, scale_bf16_high)
	215	+ S.vsts(out_ub[row, col], out_low, dist='NORM_B16')
	216	+ S.vsts(out_ub[row, col + 128], out_high, dist='NORM_B16')
	217	+ elif in_config.dtype == T.float4_e2m1fn:
	218	+ zero_bf16 = S.vdup(0.0, T.bfloat16)
	219	+ for group in T.serial(num_sf_rows):
	220	+ row_base = group * num_per_tokens
	221	+ for chunk in T.serial(block_k // 128):
	222	+ col = chunk * 128
	223	+ scale_low = load_scale(sf_ub, group, col, mask_low_f32)
	224	+ scale_high = load_scale(sf_ub, group, col + 64, mask_low_f32)
	225	+ for row in T.serial(num_per_tokens):
	226	+ dequant_fp4_chunk(x_ub, out_ub, row_base + row, col, scale_low, scale_high, zero_bf16)
	227	+ else:
	228	+ for group in T.serial(num_sf_rows):
	229	+ row_base = group * num_per_tokens
	230	+ for chunk in T.serial(block_k // 64):
	231	+ col = chunk * 64
	232	+ scale = load_scale(sf_ub, group, col, mask_low_f32)
	233	+ for row in T.serial(num_per_tokens):
	234	+ dequant_fp8_chunk(x_ub, out_ub, row_base + row, col, scale)
	235	+
	236	+ @T.prim_func
	237	+ def cast_back_kernel_ascend(
	238	+ x: T.Tensor[(num_tokens, hidden), in_config.dtype],
	239	+ x_sf: T.StridedTensor[sf_shape, (sf_stride, 1), in_config.sf_dtype],
	240	+ out: T.Tensor[(num_tokens, hidden), out_dtype],
	241	+ ):
	242	+ with T.Kernel(n_cores) as core_id:
	243	+ x_ub = T.alloc_shared((block_m, block_k), in_config.dtype)
	244	+ sf_raw_ub = T.alloc_shared(
	245	+ (
	246	+ sf_raw_rows,
	247	+ align(
	248	+ sf_raw_cols,
	249	+ 32 // sf_raw_dtype.bytes if in_config.use_tma_aligned_col_major_sf else 1,
	250	+ ),
	251	+ ),
	252	+ sf_raw_dtype,
	253	+ )
	254	+ sf_ub = T.alloc_shared((num_sf_rows, align(num_sf_cols, 32 // sf_ub_dtype.bytes)), sf_ub_dtype)
	255	+ out_ub = T.alloc_shared((block_m, block_k), out_dtype)
	256	+ x_sf_raw = (
	257	+ T.Tensor((sf_shape[0], sf_shape[1] // pack_factor), T.uint16, x_sf.data)
	258	+ if in_config.use_packed_ue8m0
	259	+ else x_sf
	260	+ )
	261	+ T.annotate_buffer_versions(
	262	+ {
	263	+ x_ub: num_stages,
	264	+ sf_raw_ub: num_stages if num_per_tokens == 1 else 1,
	265	+ sf_ub: num_stages if num_per_tokens == 1 else 1,
	266	+ out_ub: num_stages,
	267	+ }
	268	+ )
	269	+
	270	+ for pid_token, pid_hidden in T.Persistent(
	271	+ [num_token_tiles, num_hidden_tiles],
	272	+ n_cores,
	273	+ core_id,
	274	+ group_size=1,
	275	+ num_stages=num_stages,
	276	+ ):
	277	+ token_base = T.alloc_var(T.int32, init=pid_token * block_m // num_per_tokens)
	278	+ channel_base = T.alloc_var(T.int32, init=pid_hidden * block_k // num_per_channels)
	279	+ token_offset = pid_token * block_m
	280	+ hidden_offset = pid_hidden * block_k
	281	+ valid_rows = T.min(block_m, num_tokens - token_offset)
	282	+ valid_sf_cols = T.min(num_sf_cols, num_hidden_sf - channel_base)
	283	+ if hidden_tail_cols != 0:
	284	+ if pid_hidden == num_hidden_tiles - 1:
	285	+ T.copy(
	286	+ x[
	287	+ token_offset : token_offset + valid_rows,
	288	+ hidden_offset : hidden_offset + hidden_tail_cols,
	289	+ ],
	290	+ x_ub[:valid_rows, :hidden_tail_cols],
	291	+ )
	292	+ else:
	293	+ if valid_rows == block_m:
	294	+ T.copy(x[token_offset, hidden_offset], x_ub)
	295	+ else:
	296	+ T.copy(
	297	+ x[
	298	+ token_offset : token_offset + valid_rows,
	299	+ hidden_offset : hidden_offset + block_k,
	300	+ ],
	301	+ x_ub[:valid_rows, :block_k],
	302	+ )
	303	+ else:
	304	+ if valid_rows == block_m:
	305	+ T.copy(x[token_offset, hidden_offset], x_ub)
	306	+ else:
	307	+ T.copy(
	308	+ x[
	309	+ token_offset : token_offset + valid_rows,
	310	+ hidden_offset : hidden_offset + block_k,
	311	+ ],
	312	+ x_ub[:valid_rows, :block_k],
	313	+ )
	314	+ if valid_rows == block_m and valid_sf_cols == num_sf_cols:
	315	+ load_scale_tile(x_sf_raw, sf_raw_ub, sf_ub, token_base, channel_base)
	316	+ else:
	317	+ load_scale_tail(
	318	+ x_sf_raw,
	319	+ sf_ub,
	320	+ token_base,
	321	+ channel_base,
	322	+ T.ceildiv(valid_rows, num_per_tokens),
	323	+ valid_sf_cols,
	324	+ )
	325	+ transform_tile(x_ub, sf_ub, out_ub)
	326	+ if hidden_tail_cols != 0:
	327	+ if pid_hidden == num_hidden_tiles - 1:
	328	+ T.copy(
	329	+ out_ub[:valid_rows, :hidden_tail_cols],
	330	+ out[
	331	+ token_offset : token_offset + valid_rows,
	332	+ hidden_offset : hidden_offset + hidden_tail_cols,
	333	+ ],
	334	+ )
	335	+ elif valid_rows == block_m:
	336	+ T.copy(out_ub, out[token_offset, hidden_offset])
	337	+ else:
	338	+ T.copy(
	339	+ out_ub[:valid_rows, :block_k],
	340	+ out[
	341	+ token_offset : token_offset + valid_rows,
	342	+ hidden_offset : hidden_offset + block_k,
	343	+ ],
	344	+ )
	345	+ elif valid_rows == block_m:
	346	+ T.copy(out_ub, out[token_offset, hidden_offset])
	347	+ else:
	348	+ T.copy(
	349	+ out_ub[:valid_rows, :block_k],
	350	+ out[
	351	+ token_offset : token_offset + valid_rows,
	352	+ hidden_offset : hidden_offset + block_k,
	353	+ ],
	354	+ )
	355	+
	356	+ return cast_back_kernel_ascend

Asrc/cann_ops_tilelang/quant/common.py+341-0

@@ -0,0 +1,341 @@
	1	+import functools
	2	+from dataclasses import dataclass
	3	+from typing import Union # noqa: UP035 - TileLang macro annotations require typing.Union
	4	+
	5	+import torch
	6	+from tilelang.ascend import language as T
	7	+
	8	+from cann_ops_tilelang.config import get_amax_clamp_for_quant
	9	+from cann_ops_tilelang.utils import ceil_div, str_to_dtype
	10	+
	11	+from .types import QuantTensor
	12	+
	13	+
	14	+@dataclass(frozen=True)
	15	+class BaseCastConfig:
	16	+ torch_dtype: torch.dtype = torch.float8_e4m3fn
	17	+ sf_block: tuple[int, int] = (1, 1)
	18	+ use_tma_aligned_col_major_sf: bool = False
	19	+ use_packed_ue8m0: bool = False
	20	+
	21	+ @property
	22	+ def with_sf(self) -> bool:
	23	+ return self.torch_dtype in [torch.float8_e4m3fn, torch.int8, torch.uint32]
	24	+
	25	+ @property
	26	+ def dtype(self) -> T.dtype:
	27	+ return T.dtype(self.torch_dtype) if self.torch_dtype != torch.int8 else T.float4_e2m1fn
	28	+
	29	+ @property
	30	+ def sf_torch_dtype(self) -> torch.dtype:
	31	+ return torch.uint8 if self.use_packed_ue8m0 else torch.float32
	32	+
	33	+ @property
	34	+ def sf_dtype(self) -> T.dtype:
	35	+ return T.dtype(self.sf_torch_dtype)
	36	+
	37	+ @property
	38	+ def sf_col_pack(self) -> bool:
	39	+ return self.use_packed_ue8m0 and self.sf_block[1] == 1
	40	+
	41	+
	42	+@dataclass(frozen=True)
	43	+class CastInputConfig(BaseCastConfig):
	44	+ pass
	45	+
	46	+
	47	+@dataclass(frozen=True)
	48	+class CastOutputConfig(BaseCastConfig):
	49	+ round_sf: bool = False
	50	+ clamp_min_value: float \| None = None
	51	+
	52	+
	53	+def get_sf_inv_dtype(value_dtype: T.dtype, out_config: CastOutputConfig) -> T.dtype:
	54	+ if value_dtype == T.bfloat16 and out_config.round_sf:
	55	+ return T.bfloat16
	56	+ return T.float32
	57	+
	58	+
	59	+def get_packed_ue8m0_pack_factor() -> int:
	60	+ return 2
	61	+
	62	+
	63	+def get_packed_ue8m0_torch_dtype() -> torch.dtype:
	64	+ return torch.int16
	65	+
	66	+
	67	+@functools.lru_cache
	68	+def get_cast_input_config_impl(
	69	+ dtype: torch.dtype,
	70	+ sf_block: tuple[int, int] \| None = None,
	71	+ use_tma_aligned_col_major_sf: bool = False,
	72	+ use_packed_ue8m0: bool = False,
	73	+) -> CastInputConfig:
	74	+ if sf_block is None:
	75	+ assert dtype in (torch.bfloat16, torch.float32)
	76	+ return CastInputConfig(dtype)
	77	+ else:
	78	+ assert dtype in (torch.float8_e4m3fn, torch.int8, torch.uint8)
	79	+ return CastInputConfig(dtype, sf_block, use_tma_aligned_col_major_sf, use_packed_ue8m0)
	80	+
	81	+
	82	+def get_cast_input_and_config(
	83	+ x: torch.Tensor \| QuantTensor,
	84	+ sf_block: tuple[int, int] \| None,
	85	+) -> tuple[torch.Tensor, torch.Tensor \| None, CastInputConfig]:
	86	+ if isinstance(x, tuple):
	87	+ x, x_sf = x
	88	+ assert isinstance(sf_block, tuple) and isinstance(x, torch.Tensor) and isinstance(x_sf, torch.Tensor)
	89	+ assert x_sf.dtype in (torch.float32, get_packed_ue8m0_torch_dtype())
	90	+
	91	+ use_tma_aligned_col_major_sf = False
	92	+ use_packed_ue8m0 = False
	93	+ if x_sf.stride(0) == 1 and x_sf.stride(1) != 1:
	94	+ use_tma_aligned_col_major_sf = True
	95	+ x_sf = x_sf.T
	96	+ else:
	97	+ assert x_sf.stride(1) == 1
	98	+ if x_sf.dtype == get_packed_ue8m0_torch_dtype():
	99	+ use_packed_ue8m0 = True
	100	+ x_sf = x_sf.view(torch.uint8)
	101	+
	102	+ config = get_cast_input_config_impl(x.dtype, sf_block, use_tma_aligned_col_major_sf, use_packed_ue8m0)
	103	+ return x, x_sf, config
	104	+ else:
	105	+ assert sf_block is None and isinstance(x, torch.Tensor)
	106	+ config = get_cast_input_config_impl(x.dtype)
	107	+ return x, None, config
	108	+
	109	+
	110	+@functools.lru_cache
	111	+def get_cast_output_config_impl(
	112	+ fmt: str,
	113	+ sf_block: tuple[int, int],
	114	+ use_tma_aligned_col_major_sf: bool,
	115	+ round_sf: bool,
	116	+ use_packed_ue8m0: bool,
	117	+ clamp_min_value: float \| None,
	118	+) -> CastOutputConfig:
	119	+ assert fmt in ('fp32', 'bf16', 'e4m3', 'e2m1')
	120	+ assert not use_packed_ue8m0 or round_sf
	121	+ config = CastOutputConfig(
	122	+ torch_dtype=str_to_dtype(fmt),
	123	+ sf_block=sf_block,
	124	+ use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf,
	125	+ round_sf=round_sf,
	126	+ use_packed_ue8m0=use_packed_ue8m0,
	127	+ clamp_min_value=clamp_min_value,
	128	+ )
	129	+ return config
	130	+
	131	+
	132	+def get_cast_output_config(
	133	+ fmt: str,
	134	+ sf_block: tuple[int, int],
	135	+ use_tma_aligned_col_major_sf: bool = False,
	136	+ round_sf: bool = False,
	137	+ use_packed_ue8m0: bool = False,
	138	+) -> CastOutputConfig:
	139	+ return get_cast_output_config_impl(
	140	+ fmt,
	141	+ sf_block,
	142	+ use_tma_aligned_col_major_sf,
	143	+ round_sf,
	144	+ use_packed_ue8m0,
	145	+ get_amax_clamp_for_quant(fmt) if fmt in ('e4m3', 'e2m1') else None,
	146	+ )
	147	+
	148	+
	149	+def get_logical_hidden(hidden: int, dtype: torch.dtype) -> int:
	150	+ """
	151	+ Compute hidden size when `torch.int8` is used for packing FP4
	152	+ """
	153	+ return hidden if dtype != torch.int8 else hidden * 2
	154	+
	155	+
	156	+def get_physical_hidden(hidden: int, dtype: torch.dtype) -> int:
	157	+ """
	158	+ Compute hidden size when `torch.int8` is used for packing FP4
	159	+ """
	160	+ return hidden if dtype != torch.int8 else hidden // 2
	161	+
	162	+
	163	+def get_sf_shape(shape: tuple[int, int], config: BaseCastConfig) -> tuple[int, int]:
	164	+ num_block_m = ceil_div(shape[0], config.sf_block[0])
	165	+ num_block_k = ceil_div(shape[1], config.sf_block[1])
	166	+
	167	+ if not config.sf_col_pack:
	168	+ if config.use_packed_ue8m0:
	169	+ num_block_k = ceil_div(num_block_k, get_packed_ue8m0_pack_factor())
	170	+ if config.use_tma_aligned_col_major_sf:
	171	+ num_block_m, num_block_k = num_block_k, num_block_m
	172	+ if config.use_packed_ue8m0:
	173	+ num_block_k *= get_packed_ue8m0_pack_factor()
	174	+ else:
	175	+ num_block_m = ceil_div(num_block_m, get_packed_ue8m0_pack_factor())
	176	+ num_block_k *= get_packed_ue8m0_pack_factor()
	177	+
	178	+ return num_block_m, num_block_k
	179	+
	180	+
	181	+def alloc_scaling_factors(
	182	+ shape: tuple[int, int],
	183	+ out_config: BaseCastConfig,
	184	+ device: torch.device \| str = 'npu',
	185	+) -> torch.Tensor:
	186	+ """
	187	+ Allocate scaling factors for quantization.
	188	+ """
	189	+ sf_shape = get_sf_shape(shape, out_config)
	190	+ if out_config.use_tma_aligned_col_major_sf:
	191	+ scaling_factor = torch.empty(
	192	+ size=sf_shape,
	193	+ dtype=out_config.sf_torch_dtype,
	194	+ device=device,
	195	+ )
	196	+ # Torch gives unexpected stride when zero shape is placed at the last dimension
	197	+ if sf_shape[1] == 0:
	198	+ scaling_factor.as_strided_(scaling_factor.shape, (sf_shape[1], 1))
	199	+ return scaling_factor
	200	+ else:
	201	+ return torch.empty(
	202	+ size=sf_shape,
	203	+ dtype=out_config.sf_torch_dtype,
	204	+ device=device,
	205	+ )
	206	+
	207	+
	208	+def cast_epilogue(out_sf: torch.Tensor, config: BaseCastConfig) -> torch.Tensor:
	209	+ """Post-process the sf-factor tensor after a cast kernel launch.
	210	+
	211	+ Args:
	212	+ out_sf: Raw sf-factor tensor produced by the kernel.
	213	+ config: Cast configuration used during the kernel launch.
	214	+
	215	+ Returns:
	216	+ Corrected sf-factor tensor with proper layout and shape.
	217	+ """
	218	+ # Make corrected SF tensor
	219	+ if config.use_packed_ue8m0:
	220	+ out_sf = out_sf.view(dtype=get_packed_ue8m0_torch_dtype())
	221	+ return out_sf.T if config.use_tma_aligned_col_major_sf else out_sf
	222	+
	223	+
	224	+@T.macro
	225	+def get_sf_and_inv(amax: float, out_config: CastOutputConfig, sf_inv_dtype: T.dtype = T.float32):
	226	+ # Clamp with min value
	227	+ clamped_amax = T.max(amax, out_config.clamp_min_value)
	228	+
	229	+ max_value = T.max_value(out_config.dtype)
	230	+ sf = T.alloc_var(T.float32)
	231	+ if out_config.round_sf: # noqa: SIM108 - keep explicit control flow in the TileLang macro
	232	+ # round_sf only uses ceil(log2(sf)); for supported formats this
	233	+ # reciprocal multiply preserves the same power-of-two bin as division.
	234	+ sf = clamped_amax * (T.float32(1) / max_value)
	235	+ else:
	236	+ sf = clamped_amax / max_value
	237	+ if not out_config.round_sf:
	238	+ return sf, max_value / clamped_amax
	239	+
	240	+ # Round into 2's power
	241	+ bits = T.reinterpret(sf, T.uint32)
	242	+ # clamp_min_value ensures sf is positive and normal before exponent extraction.
	243	+ # `(bits - 1) >> 23 + 1` gives ceil(log2).
	244	+ exp_sf = ((bits - 1) >> 23) + 1 - 127
	245	+ if sf_inv_dtype == T.bfloat16:
	246	+ sf_inv = T.reinterpret(T.uint16(127 - exp_sf) << 7, T.bfloat16)
	247	+ else:
	248	+ sf_inv = T.reinterpret((127 - exp_sf) << 23, T.float32)
	249	+ if out_config.use_packed_ue8m0:
	250	+ return T.uint8(exp_sf + 127), sf_inv
	251	+ else:
	252	+ return T.reinterpret((127 + exp_sf) << 23, T.float32), sf_inv
	253	+
	254	+
	255	+@T.macro
	256	+def get_sf_index(m_idx: int, k_idx: int, config: BaseCastConfig):
	257	+ pack_factor = get_packed_ue8m0_pack_factor()
	258	+ if config.sf_col_pack:
	259	+ return m_idx // pack_factor, k_idx * pack_factor + m_idx % pack_factor
	260	+ if config.use_packed_ue8m0 and config.use_tma_aligned_col_major_sf:
	261	+ return k_idx // pack_factor, m_idx * pack_factor + k_idx % pack_factor
	262	+ if config.use_tma_aligned_col_major_sf:
	263	+ return k_idx, m_idx
	264	+ return m_idx, k_idx
	265	+
	266	+
	267	+@T.macro
	268	+def load_sf(tensor: T.Tensor, m_idx: int, k_idx: int, config: BaseCastConfig):
	269	+ row, col = get_sf_index(m_idx, k_idx, config)
	270	+ return tensor[row, col]
	271	+
	272	+
	273	+@T.macro
	274	+def transform_sf(sf: Union[T.float32, T.uint8], config: BaseCastConfig) -> T.float32: # noqa: UP007
	275	+ if config.use_packed_ue8m0:
	276	+ return T.reinterpret(T.uint32(sf) << 23, T.float32)
	277	+ else:
	278	+ return sf
	279	+
	280	+
	281	+@T.macro
	282	+def store_sf(
	283	+ tensor: T.Tensor,
	284	+ sf: Union[T.float32, T.uint8], # noqa: UP007
	285	+ m_idx: int,
	286	+ k_idx: int,
	287	+ config: BaseCastConfig,
	288	+):
	289	+ row, col = get_sf_index(m_idx, k_idx, config)
	290	+ tensor[row, col] = sf
	291	+
	292	+
	293	+def unpack_from_e2m1fn_x2(x: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
	294	+ """
	295	+ Decode a uint8/int8 tensor of packed fp4 values to float32. Used mainly for debugging.
	296	+
	297	+ The input tensor is expected to have shape [..., 2 * K] where the last dimension contains
	298	+ pairs of packed fp4 values.
	299	+ The output tensor will have shape [..., K] with the decoded float32 values.
	300	+ """
	301	+ assert x.dtype == torch.int8 or x.dtype == torch.uint8
	302	+
	303	+ if x.ndim == 0:
	304	+ raise ValueError('x must have at least 1 dimension so the last dim can be doubled')
	305	+
	306	+ lo = (x & 0x0F).to(torch.int16)
	307	+ hi = ((x >> 4) & 0x0F).to(torch.int16)
	308	+
	309	+ def decode_fp4_e2m1(n: torch.Tensor) -> torch.Tensor:
	310	+ # n in [0..15], layout: s(1) \| e(2) \| m(1)
	311	+ s = (n >> 3) & 0x1
	312	+ e = (n >> 1) & 0x3
	313	+ m = n & 0x1
	314	+
	315	+ sign = torch.where(
	316	+ s == 1,
	317	+ torch.tensor(-1.0, device=n.device),
	318	+ torch.tensor(1.0, device=n.device),
	319	+ )
	320	+
	321	+ bias = 1
	322	+
	323	+ # subnormal/zero: e==0 -> value = sign * 2^(1-bias) * (m/2)
	324	+ # bias=1 => 2^(0)=1 => {0, 0.5}
	325	+ sub = (m.to(torch.float32) * 0.5) * (2.0 ** (1 - bias))
	326	+
	327	+ # 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(
	329	+ torch.tensor(2.0, device=n.device),
	330	+ (e - bias).to(torch.float32),
	331	+ )
	332	+
	333	+ val = torch.where(e == 0, sub, norm)
	334	+ return (val * sign).to(out_dtype)
	335	+
	336	+ flo = decode_fp4_e2m1(lo)
	337	+ fhi = decode_fp4_e2m1(hi)
	338	+
	339	+ # (..., L) -> (..., L, 2) -> (..., 2L)
	340	+ y = torch.stack([flo, fhi], dim=-1).reshape(x.shape[:-1], x.shape[-1] 2)
	341	+ return y

Asrc/cann_ops_tilelang/quant/per_block_cast.py+157-0

@@ -0,0 +1,157 @@
	1	+import torch
	2	+
	3	+from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
	4	+from cann_ops_tilelang.quant.common import (
	5	+ QuantTensor,
	6	+ alloc_scaling_factors,
	7	+ cast_epilogue,
	8	+ get_cast_input_and_config,
	9	+ get_cast_output_config,
	10	+)
	11	+
	12	+from .per_block_cast_kernel import get_per_block_cast_kernel_asc
	13	+
	14	+
	15	+def per_block_cast_impl(
	16	+ x: torch.Tensor,
	17	+ fmt: str,
	18	+ block_size: tuple[int, int],
	19	+ use_tma_aligned_col_major_sf: bool = False,
	20	+ round_sf: bool = False,
	21	+ use_packed_ue8m0: bool = False,
	22	+ sf_only: bool = False,
	23	+ sf: torch.Tensor \| None = None,
	24	+) -> torch.Tensor \| QuantTensor:
	25	+ """
	26	+ sf_only: If True, only compute and return sf factors.
	27	+ sf: Pre-computed sf factors; if provided, cast-only mode is used.
	28	+
	29	+ Returns:
	30	+ Scale-factor tensor when ``sf_only=True``, casted tensor when ``sf``
	31	+ is provided, or a tuple ``(out, out_sf)`` with quantized output and sf factors.
	32	+ """
	33	+ require_ascend()
	34	+ assert x.device.type == 'npu'
	35	+ assert x.is_contiguous() and x.dim() == 2
	36	+ num_tokens, hidden = x.shape
	37	+
	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(
	47	+ hidden=hidden,
	48	+ in_config=in_config,
	49	+ out_config=out_config,
	50	+ sf_only=sf_only,
	51	+ cast_only=(sf is not None),
	52	+ )
	53	+
	54	+ if should_print_kernel_source():
	55	+ print(kernel.get_kernel_source())
	56	+
	57	+ out = torch.empty(
	58	+ (num_tokens, hidden if fmt == 'e4m3' else hidden // 2), dtype=out_config.torch_dtype, device=x.device
	59	+ )
	60	+ out_sf = (
	61	+ sf
	62	+ if sf is not None
	63	+ else alloc_scaling_factors(
	64	+ (num_tokens, hidden),
	65	+ out_config=out_config,
	66	+ device=x.device,
	67	+ )
	68	+ )
	69	+
	70	+ if num_tokens > 0:
	71	+ kernel(x, out, out_sf)
	72	+
	73	+ out_sf = cast_epilogue(out_sf, out_config)
	74	+
	75	+ if sf_only:
	76	+ return out_sf
	77	+
	78	+ if sf is not None:
	79	+ return out
	80	+
	81	+ return out, out_sf
	82	+
	83	+
	84	+def per_block_cast(
	85	+ x: torch.Tensor,
	86	+ fmt: str,
	87	+ block_size: tuple[int, int],
	88	+ use_tma_aligned_col_major_sf: bool = False,
	89	+ round_sf: bool = False,
	90	+ use_packed_ue8m0: bool = False,
	91	+) -> QuantTensor:
	92	+ """Cast a matrix to FP8/FP4 with per-block scaling factors.
	93	+
	94	+ Args:
	95	+ x: Input 2D contiguous tensor of shape (num_tokens, hidden).
	96	+ fmt: Target quantized format (``'e4m3'`` or ``'e2m1'``).
	97	+ block_size: Scaling block size as (num_per_tokens, num_per_channels).
	98	+ use_tma_aligned_col_major_sf: Whether to use TMA-aligned column-major sf factors.
	99	+ round_sf: Whether to round scaling factors to powers of two.
	100	+ use_packed_ue8m0: Whether to use packed UE8M0 format for sf factors.
	101	+
	102	+ Returns:
	103	+ Tuple of quantized output and sf factors.
	104	+ """
	105	+ return per_block_cast_impl(x, fmt, block_size, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0)
	106	+
	107	+
	108	+def per_block_cast_with_sf_only(
	109	+ x: torch.Tensor,
	110	+ fmt: str,
	111	+ block_size: tuple[int, int],
	112	+ use_tma_aligned_col_major_sf: bool = False,
	113	+ round_sf: bool = False,
	114	+ use_packed_ue8m0: bool = False,
	115	+) -> torch.Tensor:
	116	+ """Cast a matrix to FP8/FP4, only output the scaling factors.
	117	+
	118	+ Args:
	119	+ x: Input 2D contiguous tensor of shape (num_tokens, hidden).
	120	+ fmt: Target quantized format (``'e4m3'`` or ``'e2m1'``).
	121	+ block_size: Scaling block size as (num_per_tokens, num_per_channels).
	122	+ use_tma_aligned_col_major_sf: Whether to use TMA-aligned column-major sf factors.
	123	+ round_sf: Whether to round scaling factors to powers of two.
	124	+ use_packed_ue8m0: Whether to use packed UE8M0 format for sf factors.
	125	+
	126	+ Returns:
	127	+ Scaling factors only
	128	+ """
	129	+ return per_block_cast_impl(
	130	+ x, fmt, block_size, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0, sf_only=True
	131	+ )
	132	+
	133	+
	134	+def per_block_cast_with_precomputed_sf(
	135	+ x: torch.Tensor,
	136	+ fmt: str,
	137	+ block_size: tuple[int, int],
	138	+ sf: torch.Tensor,
	139	+ use_tma_aligned_col_major_sf: bool = False,
	140	+ round_sf: bool = False,
	141	+ use_packed_ue8m0: bool = False,
	142	+) -> torch.Tensor:
	143	+ """Cast a matrix to FP8/FP4 with per-block scaling using precomputed sf factors.
	144	+
	145	+ Args:
	146	+ x: Input 2D contiguous tensor of shape (num_tokens, hidden).
	147	+ fmt: Target quantized format (``'e4m3'`` or ``'e2m1'``).
	148	+ block_size: Scaling block size as (num_per_tokens, num_per_channels).
	149	+ sf: Pre-computed sf factors.
	150	+ use_tma_aligned_col_major_sf: Whether to use TMA-aligned column-major sf factors.
	151	+ round_sf: Whether to round scaling factors to powers of two.
	152	+ use_packed_ue8m0: Whether to use packed UE8M0 format for sf factors.
	153	+
	154	+ Returns:
	155	+ Casted tensor using the precomputed sf factors.
	156	+ """
	157	+ return per_block_cast_impl(x, fmt, block_size, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0, sf=sf)

Asrc/cann_ops_tilelang/quant/per_block_cast_kernel.py+413-0

@@ -0,0 +1,413 @@
	1	+# Ruff's simplification rules do not model TileLang's compile-time branches.
	2	+# ruff: noqa: SIM108
	3	+
	4	+import math
	5	+
	6	+import tilelang
	7	+from tilelang.ascend import language as T
	8	+from tilelang.ascend.language import simd as S
	9	+
	10	+from cann_ops_tilelang.config import get_num_vector_cores
	11	+from cann_ops_tilelang.quant.common import (
	12	+ CastInputConfig,
	13	+ CastOutputConfig,
	14	+ get_packed_ue8m0_pack_factor,
	15	+ get_sf_shape,
	16	+)
	17	+from cann_ops_tilelang.utils import align_up, ceil_div
	18	+
	19	+
	20	+@tilelang.jit()
	21	+def get_per_block_cast_kernel_asc(
	22	+ hidden: int,
	23	+ in_config: CastInputConfig,
	24	+ out_config: CastOutputConfig,
	25	+ sf_only: bool = False,
	26	+ cast_only: bool = False,
	27	+):
	28	+ n_cores = get_num_vector_cores()
	29	+ 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)
	31	+ assert out_config.dtype in (T.float4_e2m1fn, T.float8_e4m3fn)
	32	+ quant_max = 6.0 if out_config.dtype == T.float4_e2m1fn else 448.0
	33	+
	34	+ # Larger block_m needs a smaller block_k to stay inside the shared-memory budget.
	35	+ block_m = num_per_tokens
	36	+ if block_m == 32:
	37	+ block_k_limit = 512
	38	+ else:
	39	+ if out_config.use_packed_ue8m0 and num_per_channels == 128:
	40	+ block_k_limit = 256
	41	+ else:
	42	+ block_k_limit = 128 if in_config.dtype == T.float32 else 256
	43	+ block_k = align_up(math.gcd(block_k_limit, hidden), 128)
	44	+ if block_m == 128 and block_k == 256:
	45	+ num_stages = 1 if in_config.dtype == T.float32 else 2
	46	+ else:
	47	+ num_stages = 3
	48	+ assert block_k >= num_per_channels
	49	+ assert hidden % block_k == 0
	50	+
	51	+ num_sf_cols_per_block = block_k // num_per_channels
	52	+ pack_factor = get_packed_ue8m0_pack_factor()
	53	+ assert not out_config.use_tma_aligned_col_major_sf or num_sf_cols_per_block % pack_factor == 0 or hidden == block_k
	54	+ if out_config.use_packed_ue8m0:
	55	+ packed_sf_cols_per_block = ceil_div(num_sf_cols_per_block, pack_factor)
	56	+ else:
	57	+ packed_sf_cols_per_block = num_sf_cols_per_block
	58	+ num_stored_sf_cols_per_block = packed_sf_cols_per_block if out_config.use_packed_ue8m0 else num_sf_cols_per_block
	59	+ sf_ub_dtype = T.uint16 if out_config.use_packed_ue8m0 else T.float32
	60	+
	61	+ # Rows processed per persistent task. Larger tiles carry more rows, so the
	62	+ # group shrinks to keep small token counts mapped to more pid_x tiles;
	63	+ # block_m=32 grows to 4 only with TMA+packed sf.
	64	+ token_group = (
	65	+ 2
	66	+ if block_m == 128 and block_k == 128
	67	+ else 1
	68	+ if block_m == 128
	69	+ else 4
	70	+ if out_config.use_tma_aligned_col_major_sf and out_config.use_packed_ue8m0
	71	+ else 2
	72	+ )
	73	+
	74	+ num_tokens = T.dynamic('num_tokens')
	75	+ sf_shape = get_sf_shape((num_tokens, hidden), out_config)
	76	+ sf_stride = T.dynamic('sf_stride')
	77	+
	78	+ @T.macro
	79	+ def f32_pair_to_bf16(x0, x1, round_odd=False): # round_odd=False means round_to_zero here
	80	+ # in this kernel, we ensure that round_odd=False is used only when the low 16bit of fp32 inputs are all zero.
	81	+ low, high = S.vdintlv(T.reinterpret(x0, 'uint16x128'), T.reinterpret(x1, 'uint16x128'))
	82	+ if round_odd:
	83	+ high = S.vor(high, S.vmins(low, 1))
	84	+ return T.reinterpret(high, 'bfloat16x128')
	85	+
	86	+ @T.macro
	87	+ def get_scale_and_inv(amax, fp8_max, inv_fp8_max):
	88	+ if not out_config.round_sf:
	89	+ scale = S.vdiv(amax, fp8_max)
	90	+ return scale, S.vdiv(fp8_max, amax)
	91	+
	92	+ scale_raw = S.vmul(amax, inv_fp8_max)
	93	+ bits = T.reinterpret(scale_raw, 'uint32x64')
	94	+ exp = S.vadds(S.vshrs(S.vsub(bits, S.vdup(1, T.uint32)), 23), 1)
	95	+ scale_bits = S.vshls(exp, 23)
	96	+ scale = T.reinterpret(scale_bits, 'float32x64')
	97	+ inv_exp = S.vsub(S.vdup(254, T.uint32), exp)
	98	+ inv_bits = S.vshls(inv_exp, 23)
	99	+ return scale, T.reinterpret(inv_bits, 'float32x64')
	100	+
	101	+ @T.macro
	102	+ def pack_two_block_amax(amax_vec, m_low, m_high):
	103	+ a0 = S.vcmax(amax_vec, m_low)
	104	+ a1 = S.vcmax(amax_vec, m_high)
	105	+ return S.vsel(S.vdupv(a0), S.vdupv(a1), m_low)
	106	+
	107	+ @T.macro
	108	+ def store_scale_pair(sf_ub, sf_row, sf_j, scale01):
	109	+ if out_config.use_packed_ue8m0:
	110	+ exponents = S.vshrs(T.reinterpret(scale01, 'uint32x64'), 23)
	111	+ exponent_high = S.vdupv(exponents, pos='POS_HIGHEST')
	112	+ packed_exponents = S.vor(exponents, S.vshls(exponent_high, 8))
	113	+ if out_config.use_tma_aligned_col_major_sf:
	114	+ S.vsts(
	115	+ sf_ub[sf_j // pack_factor, sf_row], T.reinterpret(packed_exponents, 'uint16x128'), dist='ONEPT_B16'
	116	+ )
	117	+ else:
	118	+ S.vsts(
	119	+ sf_ub[sf_row, sf_j // pack_factor], T.reinterpret(packed_exponents, 'uint16x128'), dist='ONEPT_B16'
	120	+ )
	121	+ elif out_config.use_tma_aligned_col_major_sf:
	122	+ S.vsts(sf_ub[sf_j, sf_row], scale01, dist='ONEPT_B32')
	123	+ S.vsts(sf_ub[sf_j + 1, sf_row], S.vdupv(scale01, pos='POS_HIGHEST'), dist='ONEPT_B32')
	124	+ else:
	125	+ S.vsts(sf_ub[sf_row, sf_j], scale01, dist='ONEPT_B32')
	126	+ S.vsts(sf_ub[sf_row, sf_j + 1], S.vdupv(scale01, pos='POS_HIGHEST'), dist='ONEPT_B32')
	127	+
	128	+ @T.macro
	129	+ def load_scale_pair(sf_ub, sf_j, m_low):
	130	+ scale0 = S.vld(sf_ub[0, sf_j], dist='BRC_B32')
	131	+ scale1 = S.vld(sf_ub[0, sf_j + 1], dist='BRC_B32')
	132	+ return S.vsel(scale0, scale1, m_low)
	133	+
	134	+ @T.macro
	135	+ def store_scale_single(sf_ub, sf_row, sf_j, scale, saved_exp):
	136	+ if out_config.use_packed_ue8m0:
	137	+ exponents = S.vshrs(T.reinterpret(scale, 'uint32x64'), 23)
	138	+ if block_k == 128:
	139	+ exponent_high = S.vdupv(exponents, pos='POS_HIGHEST')
	140	+ packed_exponents = S.vor(exponents, S.vshls(exponent_high, 8))
	141	+ if out_config.use_tma_aligned_col_major_sf:
	142	+ S.vsts(
	143	+ sf_ub[sf_j // pack_factor, sf_row],
	144	+ T.reinterpret(packed_exponents, 'uint16x128'),
	145	+ dist='ONEPT_B16',
	146	+ )
	147	+ else:
	148	+ S.vsts(
	149	+ sf_ub[sf_row, sf_j // pack_factor],
	150	+ T.reinterpret(packed_exponents, 'uint16x128'),
	151	+ dist='ONEPT_B16',
	152	+ )
	153	+ else:
	154	+ if sf_j % 2 == 0:
	155	+ saved_exp[0] = exponents
	156	+ else:
	157	+ packed_exponents = S.vor(saved_exp[0], S.vshls(exponents, 8))
	158	+ if out_config.use_tma_aligned_col_major_sf:
	159	+ S.vsts(
	160	+ sf_ub[sf_j // pack_factor, sf_row],
	161	+ T.reinterpret(packed_exponents, 'uint16x128'),
	162	+ dist='ONEPT_B16',
	163	+ )
	164	+ else:
	165	+ S.vsts(
	166	+ sf_ub[sf_row, sf_j // pack_factor],
	167	+ T.reinterpret(packed_exponents, 'uint16x128'),
	168	+ dist='ONEPT_B16',
	169	+ )
	170	+ elif out_config.use_tma_aligned_col_major_sf:
	171	+ S.vsts(sf_ub[sf_j, sf_row], scale, dist='ONEPT_B32')
	172	+ else:
	173	+ S.vsts(sf_ub[sf_row, sf_j], scale, dist='ONEPT_B32')
	174	+
	175	+ @T.macro
	176	+ def load_scale_single(sf_ub, sf_j):
	177	+ return S.vld(sf_ub[0, sf_j], dist='BRC_B32')
	178	+
	179	+ @T.macro
	180	+ def quant_block(x_ub, sf_ub, out_ub, sf_row):
	181	+ with T.SimdVF():
	182	+ m_low = S.pset(32, 'PAT_VL32')
	183	+ m_high = S.vcmps(S.vci(0, T.int32), 31, op='gt')
	184	+ mask_all = S.pset(32, 'PAT_ALL')
	185	+ one_reg = S.vdup(1.0, T.float32)
	186	+ eps = S.vdup(out_config.clamp_min_value, T.float32)
	187	+ fp8_max = S.vdup(quant_max, T.float32)
	188	+ inv_fp8_max = S.vdup(1.0 / quant_max, T.float32)
	189	+ zero_bf16 = S.vdup(0.0, T.bfloat16)
	190	+ abs_mask = T.reinterpret(S.vdup(0x7FFF, T.uint16), 'bfloat16x128')
	191	+
	192	+ saved_exp = S.alloc_local((1,), T.uint32)
	193	+
	194	+ # Explicit unrolling keeps packed C128 exponent pairs in vector registers
	195	+ # without creating an unsupported PTO loop-carried local array.
	196	+ for chunk in T.unroll(block_k // 128, explicit=True):
	197	+ col = chunk * 128
	198	+ sf_col = chunk * (128 // num_per_channels)
	199	+ inverse = S.alloc_local((2,), T.float32)
	200	+ if cast_only:
	201	+ if num_per_channels == 128:
	202	+ scale = load_scale_single(sf_ub, sf_col)
	203	+ if out_config.round_sf:
	204	+ exponent = S.vshrs(T.reinterpret(scale, 'uint32x64'), 23)
	205	+ inverse_bits = S.vshls(S.vsub(S.vdup(254, T.uint32), exponent), 23)
	206	+ inverse[0] = T.reinterpret(inverse_bits, 'float32x64')
	207	+ inverse[1] = inverse[0]
	208	+ else:
	209	+ inverse[0] = S.vdiv(one_reg, scale)
	210	+ inverse[1] = inverse[0]
	211	+ else:
	212	+ for half in T.unroll(2, explicit=True):
	213	+ scale = load_scale_pair(sf_ub, sf_col + half * 2, m_low)
	214	+ if out_config.round_sf:
	215	+ exponent = S.vshrs(T.reinterpret(scale, 'uint32x64'), 23)
	216	+ inverse_bits = S.vshls(S.vsub(S.vdup(254, T.uint32), exponent), 23)
	217	+ inverse[half] = T.reinterpret(inverse_bits, 'float32x64')
	218	+ else:
	219	+ inverse[half] = S.vdiv(one_reg, scale)
	220	+ else:
	221	+ amax = S.alloc_local((2,), T.float32)
	222	+ if in_config.dtype == T.bfloat16:
	223	+ amax_bf16 = S.alloc_var(T.bfloat16)
	224	+ amax_bf16 = S.vdup(0.0, T.bfloat16)
	225	+ if num_per_channels == 128:
	226	+ for row in T.serial(num_per_tokens):
	227	+ amax_bf16 = S.vmax(amax_bf16, S.vand(S.vld(x_ub[row, col]), abs_mask))
	228	+ amax_low, amax_high = S.vintlv(zero_bf16, amax_bf16)
	229	+ amax0 = S.vdupv(S.vcmax(T.reinterpret(amax_low, 'float32x64'), mask_all), mask_all)
	230	+ amax1 = S.vdupv(S.vcmax(T.reinterpret(amax_high, 'float32x64'), mask_all), mask_all)
	231	+ amax[0] = S.vmax(amax0, amax1)
	232	+ amax[1] = amax[0]
	233	+ else:
	234	+ for row in T.serial(num_per_tokens):
	235	+ amax_bf16 = S.vmax(amax_bf16, S.vand(S.vld(x_ub[row, col]), abs_mask))
	236	+ amax_low, amax_high = S.vintlv(zero_bf16, amax_bf16)
	237	+ amax[0] = T.reinterpret(amax_low, 'float32x64')
	238	+ amax[1] = T.reinterpret(amax_high, 'float32x64')
	239	+ else:
	240	+ if num_per_channels == 128:
	241	+ amax_f32 = S.alloc_local((1,), T.float32)
	242	+ amax_f32[0] = S.vdup(0.0, T.float32)
	243	+ for row in T.serial(num_per_tokens):
	244	+ amax_f32[0] = S.vmax(amax_f32[0], S.vabs(S.vld(x_ub[row, col])))
	245	+ amax_f32[0] = S.vmax(amax_f32[0], S.vabs(S.vld(x_ub[row, col + 64])))
	246	+ amax[0] = S.vdupv(S.vcmax(amax_f32[0], mask_all), mask_all)
	247	+ amax[1] = amax[0]
	248	+ else:
	249	+ for half in T.unroll(2, explicit=True):
	250	+ amax[half] = S.vdup(0.0, T.float32)
	251	+ for row in T.serial(num_per_tokens):
	252	+ for half in T.unroll(2, explicit=True):
	253	+ amax[half] = S.vmax(amax[half], S.vabs(S.vld(x_ub[row, col + half * 64])))
	254	+ if num_per_channels == 128:
	255	+ amax[0] = S.vmax(amax[0], eps)
	256	+ scale, inverse[0] = get_scale_and_inv(amax[0], fp8_max, inv_fp8_max)
	257	+ inverse[1] = inverse[0]
	258	+ store_scale_single(sf_ub, sf_row, chunk, scale, saved_exp)
	259	+ else:
	260	+ for half in T.unroll(2, explicit=True):
	261	+ amax[half] = S.vmax(amax[half], eps)
	262	+ scale, inverse[half] = get_scale_and_inv(
	263	+ pack_two_block_amax(amax[half], m_low, m_high), fp8_max, inv_fp8_max
	264	+ )
	265	+ store_scale_pair(sf_ub, sf_row, sf_col + half * 2, scale)
	266	+
	267	+ if in_config.dtype == T.bfloat16 and out_config.round_sf:
	268	+ sf_inv = f32_pair_to_bf16(inverse[0], inverse[1])
	269	+
	270	+ if not sf_only:
	271	+ for row in T.serial(num_per_tokens):
	272	+ if in_config.dtype == T.bfloat16 and out_config.round_sf:
	273	+ quantized_bf16 = S.vmul(S.vld(x_ub[row, col]), sf_inv)
	274	+ if out_config.dtype == T.float8_e4m3fn:
	275	+ interleaved_low, interleaved_high = S.vintlv(zero_bf16, quantized_bf16)
	276	+ quantized_low = T.reinterpret(interleaved_low, 'float32x64')
	277	+ quantized_high = T.reinterpret(interleaved_high, 'float32x64')
	278	+ else:
	279	+ if in_config.dtype == T.bfloat16:
	280	+ interleaved_low, interleaved_high = S.vintlv(zero_bf16, S.vld(x_ub[row, col]))
	281	+ input_low = T.reinterpret(interleaved_low, 'float32x64')
	282	+ input_high = T.reinterpret(interleaved_high, 'float32x64')
	283	+ else:
	284	+ input_low = S.vld(x_ub[row, col])
	285	+ input_high = S.vld(x_ub[row, col + 64])
	286	+ quantized_low = S.vmul(input_low, inverse[0])
	287	+ quantized_high = S.vmul(input_high, inverse[1])
	288	+ if out_config.dtype == T.float4_e2m1fn:
	289	+ quantized_bf16 = f32_pair_to_bf16(quantized_low, quantized_high, round_odd=True)
	290	+
	291	+ if out_config.dtype == T.float4_e2m1fn:
	292	+ S.vsts(out_ub[row, col], S.vcvt(quantized_bf16, T.float4_e2m1fn), dist='PK4_B32')
	293	+ else:
	294	+ S.vsts(out_ub[row, col], S.vcvt(quantized_low, T.float8_e4m3fn), dist='PK4_B32')
	295	+ S.vsts(out_ub[row, col + 64], S.vcvt(quantized_high, T.float8_e4m3fn), dist='PK4_B32')
	296	+
	297	+ @T.prim_func
	298	+ def per_block_cast_kernel(
	299	+ x: T.Tensor[(num_tokens, hidden), in_config.dtype],
	300	+ out: T.Tensor[(num_tokens, hidden), out_config.dtype],
	301	+ out_sf: T.StridedTensor[sf_shape, (sf_stride, 1), out_config.sf_dtype],
	302	+ ):
	303	+ with T.Kernel(n_cores) as core_id:
	304	+ x_ub = T.alloc_shared((block_m, block_k), in_config.dtype)
	305	+ sf_cols = max(num_stored_sf_cols_per_block, 16 // token_group if out_config.use_packed_ue8m0 else 8)
	306	+ if out_config.use_tma_aligned_col_major_sf:
	307	+ sf_ub = T.alloc_shared((sf_cols, token_group), sf_ub_dtype)
	308	+ else:
	309	+ sf_ub = T.alloc_shared((token_group, sf_cols), sf_ub_dtype)
	310	+ out_ub = T.alloc_shared((block_m, block_k), out_config.dtype)
	311	+
	312	+ T.annotate_buffer_versions({x_ub: num_stages, sf_ub: num_stages, out_ub: num_stages})
	313	+
	314	+ for pid_x, pid_y in T.Persistent(
	315	+ [T.ceildiv(num_tokens, block_m * token_group), T.ceildiv(hidden, block_k)],
	316	+ n_cores,
	317	+ core_id,
	318	+ group_size=1,
	319	+ num_stages=num_stages,
	320	+ ):
	321	+ sf_m_base = T.alloc_var(T.int32, init=pid_x * token_group)
	322	+ sf_k_base = T.alloc_var(T.int32, init=pid_y * num_sf_cols_per_block)
	323	+ stored_sf_k_base = T.alloc_var(T.int32, init=pid_y * num_stored_sf_cols_per_block)
	324	+ if block_m == 128:
	325	+ # (128, 128): num_tokens may not be a multiple of block_m *
	326	+ # token_group, so guard the tail block and clip the sf store.
	327	+ num_blocks_total = T.alloc_var(T.int32, init=T.ceildiv(num_tokens, block_m))
	328	+ num_valid_blocks = T.alloc_var(
	329	+ T.int32,
	330	+ init=T.if_then_else(
	331	+ sf_m_base + token_group <= num_blocks_total,
	332	+ token_group,
	333	+ num_blocks_total - sf_m_base,
	334	+ ),
	335	+ )
	336	+ for token_offset in T.serial(token_group):
	337	+ token_block = sf_m_base + token_offset
	338	+ if token_block < num_blocks_total:
	339	+ T.copy(
	340	+ x[token_block * block_m, pid_y * block_k],
	341	+ x_ub[0:block_m, 0:block_k],
	342	+ l2_cache_ctrl='NOTALLOC_KEEP',
	343	+ )
	344	+ if cast_only:
	345	+ T.copy(
	346	+ out_sf[token_block, sf_k_base],
	347	+ sf_ub[0:1, 0:num_sf_cols_per_block],
	348	+ l2_cache_ctrl='NOTALLOC_KEEP',
	349	+ )
	350	+ quant_block(x_ub, sf_ub, out_ub, token_offset)
	351	+ if not sf_only:
	352	+ T.copy(
	353	+ out_ub[0:block_m, 0:block_k],
	354	+ out[token_block * block_m, pid_y * block_k],
	355	+ l2_cache_ctrl='NOTALLOC_KEEP',
	356	+ )
	357	+ if not cast_only:
	358	+ if out_config.use_packed_ue8m0:
	359	+ output_sf = T.Tensor((sf_shape[0], sf_shape[1] // pack_factor), T.uint16, out_sf.data)
	360	+ else:
	361	+ output_sf = out_sf
	362	+ if out_config.use_tma_aligned_col_major_sf:
	363	+ T.copy(
	364	+ sf_ub[0:num_stored_sf_cols_per_block, 0:num_valid_blocks],
	365	+ output_sf[stored_sf_k_base, sf_m_base],
	366	+ l2_cache_ctrl='NOTALLOC_KEEP',
	367	+ )
	368	+ else:
	369	+ T.copy(
	370	+ sf_ub[0:num_valid_blocks, 0:num_stored_sf_cols_per_block],
	371	+ output_sf[sf_m_base, stored_sf_k_base],
	372	+ l2_cache_ctrl='NOTALLOC_KEEP',
	373	+ )
	374	+ else:
	375	+ for token_offset in T.serial(token_group):
	376	+ token_block = sf_m_base + token_offset
	377	+ T.copy(
	378	+ x[token_block * block_m, pid_y * block_k],
	379	+ x_ub[0:block_m, 0:block_k],
	380	+ l2_cache_ctrl='NOTALLOC_KEEP',
	381	+ )
	382	+ if cast_only:
	383	+ T.copy(
	384	+ out_sf[token_block, sf_k_base],
	385	+ sf_ub[0:1, 0:num_sf_cols_per_block],
	386	+ l2_cache_ctrl='NOTALLOC_KEEP',
	387	+ )
	388	+ quant_block(x_ub, sf_ub, out_ub, token_offset)
	389	+ if not sf_only:
	390	+ T.copy(
	391	+ out_ub[0:block_m, 0:block_k],
	392	+ out[token_block * block_m, pid_y * block_k],
	393	+ l2_cache_ctrl='NOTALLOC_KEEP',
	394	+ )
	395	+ if not cast_only:
	396	+ if out_config.use_packed_ue8m0:
	397	+ output_sf = T.Tensor((sf_shape[0], sf_shape[1] // pack_factor), T.uint16, out_sf.data)
	398	+ else:
	399	+ output_sf = out_sf
	400	+ if out_config.use_tma_aligned_col_major_sf:
	401	+ T.copy(
	402	+ sf_ub[0:num_stored_sf_cols_per_block, 0:token_group],
	403	+ output_sf[stored_sf_k_base, sf_m_base],
	404	+ l2_cache_ctrl='NOTALLOC_KEEP',
	405	+ )
	406	+ else:
	407	+ T.copy(
	408	+ sf_ub[0:token_group, 0:num_stored_sf_cols_per_block],
	409	+ output_sf[sf_m_base, stored_sf_k_base],
	410	+ l2_cache_ctrl='NOTALLOC_KEEP',
	411	+ )
	412	+
	413	+ return per_block_cast_kernel

Asrc/cann_ops_tilelang/quant/per_block_cast_lossless.py+81-0

@@ -0,0 +1,81 @@
	1	+import torch
	2	+
	3	+from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
	4	+from cann_ops_tilelang.quant.common import (
	5	+ QuantTensor,
	6	+ alloc_scaling_factors,
	7	+ cast_epilogue,
	8	+ get_cast_input_and_config,
	9	+ get_cast_output_config,
	10	+ get_logical_hidden,
	11	+)
	12	+
	13	+from .per_block_cast_lossless_kernel import get_per_block_cast_lossless_kernel_asc
	14	+
	15	+
	16	+def per_block_cast_lossless(
	17	+ x: QuantTensor,
	18	+ fmt: str,
	19	+ x_block_size: tuple[int, int],
	20	+ out_block_size: tuple[int, int],
	21	+ use_tma_aligned_col_major_sf: bool = False,
	22	+ round_sf: bool = False,
	23	+ use_packed_ue8m0: bool = False,
	24	+) -> QuantTensor:
	25	+ """Losslessly re-quantize an FP4 (E2M1) tensor to FP8 (E4M3) with new block sizes.
	26	+
	27	+ Args:
	28	+ x: Quantized tensor pair ``(data, sf_factors)`` in E2M1 format.
	29	+ fmt: Target format (must be ``'e4m3'``).
	30	+ x_block_size: Input scaling block size as (num_per_tokens, num_per_channels).
	31	+ out_block_size: Output scaling block size as (num_per_tokens, num_per_channels).
	32	+ use_tma_aligned_col_major_sf: Whether to use TMA-aligned column-major sf factors.
	33	+ round_sf: Whether to round scaling factors to powers of two.
	34	+ use_packed_ue8m0: Whether to use packed UE8M0 format for sf factors.
	35	+
	36	+ Returns:
	37	+ A tuple ``(out, out_sf)`` with FP8 output and sf-factor tensor.
	38	+ """
	39	+ require_ascend()
	40	+ x, x_sf, in_config = get_cast_input_and_config(x, x_block_size)
	41	+
	42	+ assert fmt == 'e4m3' and x.dtype == torch.int8, 'lossless cast only supports e2m1 -> e4m3 input'
	43	+ assert x.device.type == 'npu'
	44	+ assert x.is_contiguous() and x.dim() == 2
	45	+ num_tokens, hidden = x.shape
	46	+ hidden = get_logical_hidden(hidden, x.dtype)
	47	+ assert num_tokens % 128 == 0, f'lossless cast requires num_tokens % 128 == 0, got {num_tokens}'
	48	+ assert hidden % 32 == 0
	49	+
	50	+ out_config = get_cast_output_config(fmt, out_block_size, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0)
	51	+
	52	+ assert x_block_size == (1, 32), 'ascend lossless only supports in (1, 32)'
	53	+ assert out_block_size in (
	54	+ (32, 32),
	55	+ (1, 128),
	56	+ (128, 128),
	57	+ ), 'ascend lossless only supports out (32, 32) / (1, 128) / (128, 128)'
	58	+ # FP4 is stored as packed int8; reinterpret as the 2x-packed fp4 dtype the kernel expects.
	59	+ x = x.view(torch.float4_e2m1fn_x2)
	60	+ kernel = get_per_block_cast_lossless_kernel_asc(
	61	+ hidden=hidden,
	62	+ in_config=in_config,
	63	+ out_config=out_config,
	64	+ )
	65	+
	66	+ if should_print_kernel_source():
	67	+ print(kernel.get_kernel_source())
	68	+
	69	+ out = torch.empty((num_tokens, hidden), dtype=out_config.torch_dtype, device=x.device)
	70	+ out_sf = alloc_scaling_factors(
	71	+ (num_tokens, hidden),
	72	+ out_config=out_config,
	73	+ device=x.device,
	74	+ )
	75	+
	76	+ if num_tokens > 0:
	77	+ kernel(x, x_sf, out, out_sf)
	78	+
	79	+ out_sf = cast_epilogue(out_sf, out_config)
	80	+
	81	+ return out, out_sf

Asrc/cann_ops_tilelang/quant/per_block_cast_lossless_kernel.py+408-0

@@ -0,0 +1,408 @@
	1	+# Ruff's simplification rules do not model TileLang's compile-time branches.
	2	+# ruff: noqa: SIM108
	3	+
	4	+import math
	5	+
	6	+import tilelang
	7	+from tilelang.ascend import language as T
	8	+from tilelang.ascend.language import simd as S
	9	+
	10	+from cann_ops_tilelang.config import get_num_vector_cores
	11	+from cann_ops_tilelang.quant.common import (
	12	+ CastInputConfig,
	13	+ CastOutputConfig,
	14	+ get_packed_ue8m0_pack_factor,
	15	+ get_sf_shape,
	16	+)
	17	+from cann_ops_tilelang.utils import ceil_div
	18	+
	19	+
	20	+def vector_mask(dtype, num_lanes):
	21	+ return S.pset(dtype.bits, 'PAT_VL' + str(num_lanes))
	22	+
	23	+
	24	+def lane_indices(dtype):
	25	+ index_dtype, vector_type = {8: (T.int8, 'uint8x256'), 16: (T.int16, 'uint16x128'), 32: (T.int32, 'uint32x64')}[
	26	+ dtype.bits
	27	+ ]
	28	+ return T.reinterpret(S.vci(0, index_dtype), vector_type)
	29	+
	30	+
	31	+def get_sf_storage(config, dtype):
	32	+ if config.use_packed_ue8m0:
	33	+ return T.uint16, get_packed_ue8m0_pack_factor()
	34	+ return dtype, 1
	35	+
	36	+
	37	+@T.macro
	38	+def init_fp4_to_fp8_map(fp4_to_fp8_map_ub, invalid_exp_ub):
	39	+ with T.SimdVF():
	40	+ # The mapping index stores exp_delta in the high bits and fp4_code in the low 4 bits;
	41	+ # each entry contains the equivalent FP8 bit pattern.
	42	+ lanes = lane_indices(T.int8)
	43	+ fp8_base_u32 = S.alloc_local((2,), T.uint32)
	44	+ fp8_base_u32[0], fp8_base_u32[1] = S.vintlv(
	45	+ S.vdup(T.uint32(0x3C383000), T.uint32), S.vdup(T.uint32(0x4C484440), T.uint32)
	46	+ )
	47	+ code = S.vand(lanes, S.vdup(T.uint8(0x0F), T.uint8))
	48	+ magnitude = S.vand(code, S.vdup(T.uint8(0x07), T.uint8))
	49	+ base = S.vselr(T.reinterpret(fp8_base_u32[0], 'uint8x256'), magnitude)
	50	+ delta = S.vadds(S.vand(S.vshrs(lanes, 1), S.vdup(T.uint8(0x78), T.uint8)), T.uint8(216))
	51	+ scaled = S.vadd(base, delta)
	52	+ nonzero_mask = S.vcmps(magnitude, T.uint8(0), op='ne')
	53	+ nonzero = S.vsel(scaled, S.vdup(T.uint8(0), T.uint8), nonzero_mask)
	54	+ S.vsts(
	55	+ fp4_to_fp8_map_ub[0],
	56	+ S.vor(nonzero, S.vshls(S.vand(code, S.vdup(T.uint8(0x08), T.uint8)), 4)),
	57	+ dist='NORM_B8',
	58	+ )
	59	+ assert_mask = vector_mask(T.uint16, 16)
	60	+ S.vsts(invalid_exp_ub[0], S.vdup(0, T.uint16), assert_mask, dist='NORM_B16')
	61	+
	62	+
	63	+@tilelang.jit()
	64	+def get_per_block_cast_lossless_kernel_asc(hidden: int, in_config: CastInputConfig, out_config: CastOutputConfig):
	65	+ assert in_config.dtype == T.float4_e2m1fn and out_config.dtype == T.float8_e4m3fn
	66	+ assert in_config.sf_block == (1, 32)
	67	+ assert out_config.sf_block in ((32, 32), (1, 128), (128, 128))
	68	+ assert out_config.round_sf
	69	+ assert not in_config.use_tma_aligned_col_major_sf or in_config.use_packed_ue8m0
	70	+ assert not out_config.use_tma_aligned_col_major_sf or out_config.use_packed_ue8m0
	71	+ assert hidden % out_config.sf_block[1] == 0
	72	+ # TMA (col-major) output sf is only supported for the block-level (32, 32) layout:
	73	+ # per-data-row sf (1, 128) and the single-sf-per-block (128, 128) case cannot be
	74	+ # packed into UE8M0 within one block (tests filter those combinations on Ascend).
	75	+ assert not out_config.use_tma_aligned_col_major_sf or out_config.sf_block == (32, 32)
	76	+
	77	+ out_sf_m, out_sf_k = out_config.sf_block
	78	+ if out_sf_m == 1:
	79	+ # (1, 128): every data row produces its own out-sf. block_m=64 halves the
	80	+ # token-direction tile count (8064/32=252 -> 126 tiles) and cuts per-tile
	81	+ # fixed overhead (kernel loop / T.copy launch / quant prologue) in half.
	82	+ # UB stays well inside budget (~42KB/stage vs 248KB). out_sf_m != 1 uses
	83	+ # out_sf_m rows per tile (128) which is already UB-bound for block_k=128.
	84	+ block_m = 64
	85	+ else:
	86	+ block_m = out_sf_m
	87	+ if out_sf_k == 128:
	88	+ # Output sf covers 128 columns; fix block_k at 128 so the 4-slot column
	89	+ # merge reduces to a single PAT_ALL group (verified) and every block
	90	+ # produces exactly one out-sf column. block_k = 128 also keeps block_m =
	91	+ # 128 (i.e. (128, 128)) inside the shared-memory budget.
	92	+ block_k = 128
	93	+ else:
	94	+ block_k = math.gcd(2048, hidden)
	95	+ assert block_k % out_sf_k == 0 and hidden % block_k == 0
	96	+ if block_k <= 1024 and block_m <= 64:
	97	+ num_stages = 3
	98	+ else:
	99	+ num_stages = 2
	100	+ num_per_channels = in_config.sf_block[1]
	101	+ num_sf_cols_per_block = block_k // num_per_channels
	102	+ vec_size = 64
	103	+ assert vec_size * num_per_channels >= block_k
	104	+
	105	+ num_cores = get_num_vector_cores()
	106	+ sf_dtype, in_pack_factor = get_sf_storage(in_config, T.uint32)
	107	+ num_sf_load_cols = num_sf_cols_per_block // in_pack_factor
	108	+ out_sf_dtype, out_pack_factor = get_sf_storage(out_config, T.float32)
	109	+ # Output sf granularity is out_sf_k columns; the reduce pipeline works on
	110	+ # 32-column slots and stores only one slot per out sf group.
	111	+ num_out_sf_cols_per_block = block_k // out_sf_k
	112	+ if out_config.use_packed_ue8m0:
	113	+ num_sf_store_cols = ceil_div(num_out_sf_cols_per_block, out_pack_factor)
	114	+ else:
	115	+ num_sf_store_cols = num_out_sf_cols_per_block
	116	+ num_out_sf_rows = block_m if out_sf_m == 1 else 1
	117	+ if in_config.use_tma_aligned_col_major_sf:
	118	+ sf_load_ub_shape = (num_sf_load_cols, block_m * in_pack_factor)
	119	+ else:
	120	+ sf_load_ub_shape = (block_m, vec_size * in_pack_factor)
	121	+ if out_config.use_tma_aligned_col_major_sf:
	122	+ num_batched_tiles = 2
	123	+ out_sf_shape_ub = (max(num_sf_store_cols, 4 * out_pack_factor), num_batched_tiles * out_pack_factor)
	124	+ else:
	125	+ num_batched_tiles = 1
	126	+ out_sf_shape_ub = (num_out_sf_rows, max(num_sf_store_cols, 8 * out_pack_factor) * out_pack_factor)
	127	+ num_tokens = T.dynamic('num_tokens')
	128	+ in_sf_shape = get_sf_shape((num_tokens, hidden), in_config)
	129	+ out_sf_shape = get_sf_shape((num_tokens, hidden), out_config)
	130	+ if in_config.use_tma_aligned_col_major_sf:
	131	+ in_sf_stride = T.dynamic('in_sf_stride')
	132	+ else:
	133	+ in_sf_stride = in_sf_shape[1]
	134	+ if out_config.use_tma_aligned_col_major_sf:
	135	+ out_sf_stride = T.dynamic('out_sf_stride')
	136	+ else:
	137	+ out_sf_stride = out_sf_shape[1]
	138	+
	139	+ # For block_k <= 256 the quant tile always runs on 256-column tiles; x_ub/out_ub
	140	+ # are padded to 256 columns (zero-filled) so a single 256-column code path covers
	141	+ # block_k in {64, 128, 256}.
	142	+ quant_cols = 256 if block_k <= 256 else block_k
	143	+
	144	+ @T.macro
	145	+ def load_sf(sf_load_ub, row, gather_base, gather_mask):
	146	+ if in_config.use_tma_aligned_col_major_sf:
	147	+ values = S.vgather2(sf_load_ub[0, 0], S.vadds(gather_base, row * in_pack_factor), gather_mask)
	148	+ elif sf_dtype == T.uint16:
	149	+ values = S.vld(sf_load_ub[row, 0], dist='UNPK_B8')
	150	+ else:
	151	+ values = S.vld(sf_load_ub[row, 0], dist='NORM')
	152	+ if sf_dtype == T.uint16:
	153	+ return T.reinterpret(values, 'uint16x128')
	154	+ return S.vpack(S.vshrs(T.reinterpret(values, 'uint32x64'), 23))
	155	+
	156	+ @T.macro
	157	+ def copy_sf(x_sf, sf_load_ub, row_start, pid_k):
	158	+ if in_config.use_tma_aligned_col_major_sf:
	159	+ col_start = T.alloc_var(T.int32, init=pid_k * num_sf_load_cols)
	160	+ T.copy(
	161	+ x_sf[col_start, row_start * in_pack_factor], sf_load_ub[:num_sf_load_cols, : block_m * in_pack_factor]
	162	+ )
	163	+ else:
	164	+ col_start = T.alloc_var(T.int32, init=pid_k * num_sf_cols_per_block)
	165	+ T.copy(x_sf[row_start, col_start], sf_load_ub[:block_m, :num_sf_cols_per_block])
	166	+
	167	+ @T.macro
	168	+ def col_merge(v):
	169	+ # Merge every 4 consecutive 32-column slots into a single 128-column
	170	+ # out-sf value. block_k is fixed at 128 for out_sf_k == 128, so the merge
	171	+ # is a single PAT_ALL group reduce + broadcast (verified semantics).
	172	+ pat_all = S.pset(16, 'PAT_ALL')
	173	+ return S.vdupv(S.vcmax(v, pat_all), pat_all)
	174	+
	175	+ @T.macro
	176	+ def emit_out_sf(
	177	+ sf_exp_ub,
	178	+ out_sf_ub,
	179	+ invalid_exp_ub,
	180	+ out_exp_row,
	181	+ min_exp_row,
	182	+ sf_exp_row,
	183	+ sf_ub_row,
	184	+ store_row,
	185	+ mask,
	186	+ zero,
	187	+ six,
	188	+ ):
	189	+ full_mask = vector_mask(T.uint16, max(num_sf_cols_per_block, 16))
	190	+ S.vsts(sf_exp_ub[sf_exp_row, 0], S.vadds(out_exp_row, -5), full_mask, dist='NORM_B16')
	191	+ invalid = S.vsub(out_exp_row, S.vmin(out_exp_row, S.vadds(min_exp_row, 5)))
	192	+ max_invalid = S.vcmax(invalid, mask)
	193	+ acc_invalid = S.vmax(max_invalid, S.vld(invalid_exp_ub[0], dist='BRC_B16'))
	194	+ assert_mask = vector_mask(T.uint16, 16)
	195	+ S.vsts(invalid_exp_ub[0], acc_invalid, assert_mask, dist='NORM_B16')
	196	+ out_sf_value = S.alloc_var(out_sf_dtype)
	197	+ if out_sf_dtype == T.uint16:
	198	+ out_sf_value = T.reinterpret(S.vpack(out_exp_row), 'uint16x128')
	199	+ else:
	200	+ wide_exp, _ = S.vintlv(out_exp_row, zero)
	201	+ out_sf_value = T.reinterpret(S.vshls(T.reinterpret(wide_exp, 'uint32x64'), 23), 'float32x64')
	202	+ out_store_mask = vector_mask(out_sf_dtype, num_sf_store_cols)
	203	+ if out_config.use_tma_aligned_col_major_sf:
	204	+ store_indices = lane_indices(out_sf_dtype)
	205	+ offsets = S.vadds(S.vmuls(store_indices, num_batched_tiles), store_row)
	206	+ S.vscatter(out_sf_value, out_sf_ub[0, 0], offsets, out_store_mask)
	207	+ elif out_sf_dtype == T.uint16:
	208	+ S.vsts(out_sf_ub[sf_ub_row, 0], out_sf_value, out_store_mask, dist='NORM_B16')
	209	+ else:
	210	+ S.vsts(out_sf_ub[sf_ub_row, 0], out_sf_value, out_store_mask, dist='NORM_B32')
	211	+
	212	+ @T.macro
	213	+ def reduce_sf(sf_load_ub, sf_exp_ub, out_sf_ub, invalid_exp_ub, store_row):
	214	+ with T.SimdVF():
	215	+ mask = vector_mask(T.uint16, num_sf_cols_per_block)
	216	+ full_mask = vector_mask(T.uint16, max(num_sf_cols_per_block, 16))
	217	+ zero = S.vdup(0, T.uint16)
	218	+ six = S.vdup(6, T.uint16)
	219	+ gather_base = S.alloc_var(sf_dtype)
	220	+ gather_mask = vector_mask(in_config.sf_dtype, num_sf_cols_per_block * in_pack_factor)
	221	+ if in_config.use_tma_aligned_col_major_sf:
	222	+ lanes = lane_indices(sf_dtype)
	223	+ if in_config.use_packed_ue8m0:
	224	+ pair_index = S.vshrs(lanes, 1)
	225	+ byte_index = S.vand(lanes, S.vdup(1, T.uint16))
	226	+ gather_base = S.vadd(S.vmuls(pair_index, block_m * in_pack_factor), byte_index)
	227	+ else:
	228	+ gather_base = S.vmuls(lanes, block_m)
	229	+
	230	+ max_exp0 = S.alloc_var(T.uint16)
	231	+ max_exp1 = S.alloc_var(T.uint16)
	232	+ min_exp0 = S.alloc_var(T.uint16)
	233	+ min_exp1 = S.alloc_var(T.uint16)
	234	+ max_exp0 = zero
	235	+ max_exp1 = zero
	236	+ min_exp0 = S.vdup(T.max_value(T.uint16), T.uint16)
	237	+ min_exp1 = S.vdup(T.max_value(T.uint16), T.uint16)
	238	+ for row in T.serial(0, block_m, 2):
	239	+ in_exp0 = load_sf(sf_load_ub, row, gather_base, gather_mask)
	240	+ in_exp0_clean = S.vsel(in_exp0, zero, mask)
	241	+ S.vsts(sf_exp_ub[row, 0], in_exp0_clean, full_mask, dist='NORM_B16')
	242	+ in_exp1 = load_sf(sf_load_ub, row + 1, gather_base, gather_mask)
	243	+ in_exp1_clean = S.vsel(in_exp1, zero, mask)
	244	+ S.vsts(sf_exp_ub[row + 1, 0], in_exp1_clean, full_mask, dist='NORM_B16')
	245	+ if out_sf_m == 1:
	246	+ # (1, 128): every data row produces its own out-sf row.
	247	+ # max_exp/min_exp cross-row reduction is only needed for the
	248	+ # block-level (out_sf_m != 1) path, so skip it here to save 4
	249	+ # vector ops per 2 rows (compile-time branch).
	250	+ out_exp0 = col_merge(S.vsub(S.vmax(in_exp0_clean, six), six))
	251	+ out_exp1 = col_merge(S.vsub(S.vmax(in_exp1_clean, six), six))
	252	+ emit_out_sf(
	253	+ sf_exp_ub,
	254	+ out_sf_ub,
	255	+ invalid_exp_ub,
	256	+ out_exp0,
	257	+ in_exp0_clean,
	258	+ block_m + row,
	259	+ row,
	260	+ store_row,
	261	+ mask,
	262	+ zero,
	263	+ six,
	264	+ )
	265	+ emit_out_sf(
	266	+ sf_exp_ub,
	267	+ out_sf_ub,
	268	+ invalid_exp_ub,
	269	+ out_exp1,
	270	+ in_exp1_clean,
	271	+ block_m + row + 1,
	272	+ row + 1,
	273	+ store_row,
	274	+ mask,
	275	+ zero,
	276	+ six,
	277	+ )
	278	+ else:
	279	+ max_exp0 = S.vmax(max_exp0, in_exp0_clean, mask)
	280	+ min_exp0 = S.vmin(min_exp0, in_exp0_clean, mask)
	281	+ max_exp1 = S.vmax(max_exp1, in_exp1_clean, mask)
	282	+ min_exp1 = S.vmin(min_exp1, in_exp1_clean, mask)
	283	+ if out_sf_m != 1:
	284	+ max_exp = S.vmax(max_exp0, max_exp1, mask)
	285	+ min_exp = S.vmin(min_exp0, min_exp1, mask)
	286	+ out_exp = S.alloc_var(T.uint16)
	287	+ out_exp = S.vsub(S.vmax(max_exp, six), six)
	288	+ if out_sf_k == 128:
	289	+ out_exp = col_merge(out_exp)
	290	+ emit_out_sf(
	291	+ sf_exp_ub, out_sf_ub, invalid_exp_ub, out_exp, min_exp, block_m, 0, store_row, mask, zero, six
	292	+ )
	293	+
	294	+ @T.macro
	295	+ def quant_chunk(sf_exp_ub, out_ub, fp4_to_fp8_map, fp4_codes_u8, out_exp, row, chunk):
	296	+ in_exp = S.vld(sf_exp_ub[row, chunk * 8], dist='E2B_B16')
	297	+ delta = S.vmuls(S.vsub(in_exp, out_exp), T.uint16(0x1010))
	298	+ index = S.vadd(fp4_codes_u8, T.reinterpret(delta, 'uint8x256'))
	299	+ quantized = T.reinterpret(S.vselr(fp4_to_fp8_map, index), 'float8_e4m3fnx256')
	300	+ S.vsts(out_ub[row, chunk * 256], quantized, dist='NORM_B8')
	301	+
	302	+ @T.macro
	303	+ def quant(sf_exp_ub, x_ub, out_ub, fp4_to_fp8_map_ub):
	304	+ with T.SimdVF():
	305	+ mask = S.vdup(T.uint8(0x0F), T.uint8)
	306	+ fp4_to_fp8_map = S.vld(fp4_to_fp8_map_ub[0])
	307	+ if block_k <= 256:
	308	+ # x_ub/out_ub are padded to 256 columns; the 256-column tile reads two
	309	+ # rows and writes the effective block_k columns (trailing padded code
	310	+ # is zero so the padded output columns hold valid FP8 zeros that get
	311	+ # truncated on store-back).
	312	+ for row in T.serial(0, block_m, 2):
	313	+ if out_sf_m == 1:
	314	+ out_exp0 = S.vld(sf_exp_ub[block_m + row, 0], dist='E2B_B16')
	315	+ out_exp1 = S.vld(sf_exp_ub[block_m + row + 1, 0], dist='E2B_B16')
	316	+ else:
	317	+ out_exp0 = S.vld(sf_exp_ub[block_m, 0], dist='E2B_B16')
	318	+ out_exp1 = out_exp0
	319	+ packed = T.reinterpret(S.vld(x_ub[row, 0]), 'uint8x256')
	320	+ fp4_codes_u8_0, fp4_codes_u8_1 = S.vintlv(S.vand(packed, mask), S.vshrs(packed, 4))
	321	+ quant_chunk(sf_exp_ub, out_ub, fp4_to_fp8_map, fp4_codes_u8_0, out_exp0, row, 0)
	322	+ quant_chunk(sf_exp_ub, out_ub, fp4_to_fp8_map, fp4_codes_u8_1, out_exp1, row + 1, 0)
	323	+ else:
	324	+ for chunk in T.serial(0, block_k // 256, 2):
	325	+ for row in T.serial(0, block_m, 2):
	326	+ if out_sf_m == 1:
	327	+ out_exp0 = S.vld(sf_exp_ub[block_m + row, chunk * 8], dist='E2B_B16')
	328	+ out_exp1 = S.vld(sf_exp_ub[block_m + row, (chunk + 1) * 8], dist='E2B_B16')
	329	+ else:
	330	+ out_exp0 = S.vld(sf_exp_ub[block_m, chunk * 8], dist='E2B_B16')
	331	+ out_exp1 = S.vld(sf_exp_ub[block_m, (chunk + 1) * 8], dist='E2B_B16')
	332	+ packed0 = T.reinterpret(S.vld(x_ub[row, chunk * 256]), 'uint8x256')
	333	+ packed1 = T.reinterpret(S.vld(x_ub[row + 1, chunk * 256]), 'uint8x256')
	334	+ fp4_codes_u8_0_0, fp4_codes_u8_0_1 = S.vintlv(S.vand(packed0, mask), S.vshrs(packed0, 4))
	335	+ fp4_codes_u8_1_0, fp4_codes_u8_1_1 = S.vintlv(S.vand(packed1, mask), S.vshrs(packed1, 4))
	336	+ quant_chunk(sf_exp_ub, out_ub, fp4_to_fp8_map, fp4_codes_u8_0_0, out_exp0, row, chunk)
	337	+ quant_chunk(sf_exp_ub, out_ub, fp4_to_fp8_map, fp4_codes_u8_0_1, out_exp1, row, chunk + 1)
	338	+ quant_chunk(sf_exp_ub, out_ub, fp4_to_fp8_map, fp4_codes_u8_1_0, out_exp0, row + 1, chunk)
	339	+ quant_chunk(sf_exp_ub, out_ub, fp4_to_fp8_map, fp4_codes_u8_1_1, out_exp1, row + 1, chunk + 1)
	340	+
	341	+ @T.prim_func
	342	+ def kernel(
	343	+ x: T.Tensor[(num_tokens, hidden), in_config.dtype],
	344	+ x_sf: T.StridedTensor[in_sf_shape, (in_sf_stride, 1), in_config.sf_dtype],
	345	+ out: T.Tensor[(num_tokens, hidden), out_config.dtype],
	346	+ out_sf: T.StridedTensor[out_sf_shape, (out_sf_stride, 1), out_config.sf_dtype],
	347	+ ):
	348	+ with T.Kernel(num_cores) as core_id:
	349	+ x_ub = T.alloc_shared((block_m, quant_cols), in_config.dtype)
	350	+ sf_exp_ub = T.alloc_shared((block_m + num_out_sf_rows, max(num_sf_cols_per_block, 16)), T.uint16)
	351	+ out_ub = T.alloc_shared((block_m, quant_cols), out_config.dtype)
	352	+ fp4_to_fp8_map_ub = T.alloc_shared((256,), T.uint8)
	353	+ sf_load_ub = T.alloc_shared(sf_load_ub_shape, in_config.sf_dtype)
	354	+ out_sf_ub = T.alloc_shared(out_sf_shape_ub, out_config.sf_dtype)
	355	+ invalid_exp_ub = T.alloc_shared((16,), T.uint16)
	356	+ T.annotate_buffer_versions(
	357	+ {
	358	+ x_ub: num_stages,
	359	+ sf_load_ub: num_stages,
	360	+ sf_exp_ub: num_stages,
	361	+ out_ub: num_stages,
	362	+ out_sf_ub: num_stages,
	363	+ invalid_exp_ub: 1,
	364	+ }
	365	+ )
	366	+ init_fp4_to_fp8_map(fp4_to_fp8_map_ub, invalid_exp_ub)
	367	+ for pid_m, pid_k in T.Persistent(
	368	+ [num_tokens // (block_m * num_batched_tiles), hidden // block_k],
	369	+ num_cores,
	370	+ core_id,
	371	+ group_size=1,
	372	+ num_stages=num_stages,
	373	+ ):
	374	+ for batch in T.serial(num_batched_tiles):
	375	+ tile_m = T.alloc_var(T.int32, init=pid_m * num_batched_tiles + batch)
	376	+ row_start = T.alloc_var(T.int32, init=tile_m * block_m)
	377	+ if block_k < 256:
	378	+ # Zero-fill the padded 256-column tile so the padded FP4 bytes
	379	+ # decode to code 0 (legal lookup) during quant. S.vsts writes
	380	+ # 256 bytes (two 128-byte rows) per call.
	381	+ with T.SimdVF():
	382	+ zero8 = S.vdup(0, T.uint8)
	383	+ for pad_row in T.serial(0, block_m, 2):
	384	+ S.vsts(x_ub[pad_row, 0], zero8, dist='NORM_B8')
	385	+ T.copy(x[row_start, pid_k * block_k], x_ub[:block_m, :block_k])
	386	+ copy_sf(x_sf, sf_load_ub, row_start, pid_k)
	387	+ reduce_sf(sf_load_ub, sf_exp_ub, out_sf_ub, invalid_exp_ub, batch)
	388	+ quant(sf_exp_ub, x_ub, out_ub, fp4_to_fp8_map_ub)
	389	+ T.copy(out_ub[:block_m, :block_k], out[row_start, pid_k * block_k])
	390	+ first_tile = T.alloc_var(T.int32, init=pid_m * num_batched_tiles)
	391	+ store_base = T.alloc_var(T.int32, init=pid_k * num_sf_store_cols)
	392	+ if out_config.use_tma_aligned_col_major_sf:
	393	+ T.copy(
	394	+ out_sf_ub[:num_sf_store_cols, : num_batched_tiles * out_pack_factor],
	395	+ out_sf[store_base, first_tile * out_pack_factor],
	396	+ )
	397	+ else:
	398	+ # Both buffers expose packed UE8M0 as uint8, so copy and index by
	399	+ # logical output-sf columns. This also handles a one-sf tile without
	400	+ # writing its duplicated uint16 padding byte into the next tile.
	401	+ store_col = T.alloc_var(T.int32, init=pid_k * num_out_sf_cols_per_block)
	402	+ T.copy(
	403	+ out_sf_ub[:num_out_sf_rows, :num_out_sf_cols_per_block],
	404	+ out_sf[first_tile * num_out_sf_rows, store_col],
	405	+ )
	406	+ T.device_assert(invalid_exp_ub[0] == 0)
	407	+
	408	+ return kernel

Asrc/cann_ops_tilelang/quant/per_channel_cast.py+68-0

@@ -0,0 +1,68 @@
	1	+import torch
	2	+
	3	+from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
	4	+from cann_ops_tilelang.quant.common import (
	5	+ QuantTensor,
	6	+ alloc_scaling_factors,
	7	+ cast_epilogue,
	8	+ get_cast_input_and_config,
	9	+ get_cast_output_config,
	10	+ get_sf_shape,
	11	+)
	12	+
	13	+from .per_channel_cast_kernel import get_per_channel_cast_kernel_asc
	14	+
	15	+
	16	+def per_channel_cast(
	17	+ x: torch.Tensor \| QuantTensor,
	18	+ fmt: str,
	19	+ num_per_tokens: int,
	20	+ round_sf: bool = False,
	21	+ num_per_channels: int \| None = None,
	22	+ use_packed_ue8m0: bool = False,
	23	+) -> QuantTensor:
	24	+ """Per-channel cast to FP8 (row-major SF).
	25	+
	26	+ Args:
	27	+ x: Input 2D tensor of shape ``(num_tokens, hidden)``, or a ``QuantTensor``
	28	+ ``(data, sf_invs)`` for rescaling pre-quantized FP8 inputs.
	29	+ fmt: Target FP8 format (must be ``'e4m3'``).
	30	+ num_per_tokens: Number of tokens sharing each scaling factor block (``32`` or ``128``).
	31	+ round_sf: Whether to round scaling factors to powers of two.
	32	+ num_per_channels: Number of channels in each input scaling block. Must be
	33	+ ``32`` or ``128`` when ``x`` is a ``QuantTensor``, ``None`` otherwise.
	34	+ use_packed_ue8m0: Whether to store sf factors as packed UE8M0 bytes.
	35	+
	36	+ Returns:
	37	+ A tuple ``(out, out_sf)`` with FP8 output and row-major sf-factor tensor.
	38	+ """
	39	+ require_ascend()
	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))
	41	+
	42	+ assert fmt == 'e4m3'
	43	+ assert x.device.type == 'npu'
	44	+ assert x.dim() == 2 and x.is_contiguous()
	45	+ assert num_per_tokens in (32, 128)
	46	+ num_tokens, hidden = x.shape
	47	+ assert hidden % 128 == 0, f'Per channel cast requires hidden % 128 == 0, got {hidden}'
	48	+ if x_sf_invs is not None:
	49	+ assert num_per_channels in (32, 128)
	50	+ assert x.dtype == torch.float8_e4m3fn
	51	+ assert x_sf_invs.dim() == 2 and x_sf_invs.is_contiguous()
	52	+ assert tuple(x_sf_invs.shape) == get_sf_shape((num_tokens, hidden), in_config)
	53	+ 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)
	55	+
	56	+ out_config = get_cast_output_config(fmt, (num_per_tokens, 1), False, round_sf, use_packed_ue8m0)
	57	+ assert num_tokens % 128 == 0
	58	+ 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)
	60	+ kernel = get_per_channel_cast_kernel_asc(hidden, in_config=in_config, out_config=out_config)
	61	+
	62	+ if should_print_kernel_source():
	63	+ print(kernel.get_kernel_source())
	64	+
	65	+ if num_tokens > 0:
	66	+ kernel(x, out, out_sf, x_sf_invs)
	67	+ out_sf = cast_epilogue(out_sf, out_config)
	68	+ return out, out_sf

Asrc/cann_ops_tilelang/quant/per_channel_cast_kernel.py+421-0

@@ -0,0 +1,421 @@
	1	+import math
	2	+
	3	+import tilelang
	4	+from tilelang.ascend import language as T
	5	+from tilelang.ascend.language import simd as S
	6	+
	7	+from cann_ops_tilelang.config import get_num_vector_cores
	8	+from cann_ops_tilelang.quant.common import (
	9	+ CastInputConfig,
	10	+ CastOutputConfig,
	11	+ get_packed_ue8m0_pack_factor,
	12	+ get_sf_shape,
	13	+)
	14	+from cann_ops_tilelang.utils import ceil_div
	15	+
	16	+
	17	+@tilelang.jit
	18	+def get_per_channel_cast_kernel_asc(
	19	+ hidden: int,
	20	+ in_config: CastInputConfig,
	21	+ out_config: CastOutputConfig,
	22	+):
	23	+ # bf16 hidden >= 2048 时用 nbg=1（迭代数翻倍，MTE2 请求更频繁，bf16 2048 改善 ~4%）；
	24	+ # 小 hidden(128) 保持 nbg=2 以聚合 pack_row，避免退化。
	25	+ bf16_nbg1_min_hidden = 2048
	26	+ # hidden tile 数 > 4 时 hidden 维度置外（L2 复用）；e4m3 4096 不 swap 快 2.5%，
	27	+ # 6144+ 仍需 swap 防退化（去掉后 6144 时延暴涨 ~31%）。
	28	+ swap_min_num_hidden_tiles = 4
	29	+
	30	+ num_per_tokens = out_config.sf_block[0]
	31	+ assert num_per_tokens in (32, 128)
	32	+ if in_config.with_sf:
	33	+ num_per_channels = in_config.sf_block[1]
	34	+ assert num_per_channels in (32, 128), 'Ascend rescale supports num_per_channels in (32, 128) only'
	35	+ else:
	36	+ assert in_config.dtype == T.bfloat16, 'per_channel_cast supports bf16 or e4m3 (rescale) input only'
	37	+
	38	+ quant_max = 448.0
	39	+ assert out_config.clamp_min_value >= quant_max * (2 ** (-126)), (
	40	+ 'Ascend scale exponent fast path requires clamp to normal number'
	41	+ )
	42	+
	43	+ # num_per_tokens=128 packed 路径用 block_k=256（而非 1024），使 block_m=128，
	44	+ # groups_per_tile=1；num_per_tokens=32 保持 block_k=1024 高性能路径不受影响
	45	+ use_large_block_k = in_config.with_sf and in_config.use_packed_ue8m0 and num_per_tokens != 128
	46	+ # num_per_channels=128 non-tma non-packed: dequant_ub=(num_per_tokens, block_k) float32 超出共享内存，
	47	+ # 减小 block_k 到 128
	48	+ if not use_large_block_k and in_config.with_sf and num_per_tokens == 128 and not in_config.use_packed_ue8m0:
	49	+ max_block_k = 128
	50	+ else:
	51	+ max_block_k = 1024 if use_large_block_k else 256
	52	+ block_k = math.gcd(max_block_k, hidden)
	53	+ block_m = min(128, 32768 // block_k)
	54	+ # num_per_tokens=128 时 block_m 需保持 128（=num_per_tokens），不能因 hidden==block_k 而降到 64
	55	+ if hidden // block_k == 1 and num_per_tokens != 128:
	56	+ block_m = min(64, 32768 // block_k)
	57	+ if not in_config.with_sf:
	58	+ # bf16 block_m 不得小于 num_per_tokens，否则 groups_per_tile=0 使量化循环跳过、
	59	+ # 未初始化的 out_ub 被写回（npt=128 时触发）；block_m=128 时 UB 容不下 3 级流水，故 stages=2。
	60	+ block_m = min(64, block_m)
	61	+ if block_m < num_per_tokens:
	62	+ block_m = num_per_tokens
	63	+ num_stages = 3 if block_m <= 64 else 2
	64	+ else:
	65	+ # num_per_channels=128 non-tma non-packed: block_k=128 时双缓冲仍超出 UB，降为单缓冲
	66	+ num_stages = 1 if (num_per_tokens == 128 and not in_config.use_packed_ue8m0) else 2
	67	+ groups_per_tile = block_m // num_per_tokens
	68	+ # block_m 需为 num_per_tokens 的整数倍，否则 groups_per_tile 为 0 或非整数，
	69	+ # 会导致量化循环跳过输出未初始化（num_per_tokens=128 触发）。
	70	+ assert block_m >= num_per_tokens and block_m % num_per_tokens == 0, (
	71	+ f'Unsupported Ascend tile: block_m={block_m}, num_per_tokens={num_per_tokens}'
	72	+ )
	73	+
	74	+ assert block_k % 128 == 0
	75	+ num_hidden_tiles = hidden // block_k
	76	+
	77	+ pack_factor = get_packed_ue8m0_pack_factor()
	78	+ assert pack_factor == 2
	79	+ if in_config.with_sf:
	80	+ num_packed_per_channels = num_per_channels * pack_factor if in_config.use_packed_ue8m0 else num_per_channels
	81	+ if out_config.use_packed_ue8m0:
	82	+ # bf16 2048+ 用 nbg=1 提升 MTE2 请求频率；e4m3 保持 nbg=2（packed e4m3 的 pack_row 依赖 nbg=2 的聚合布局）。
	83	+ # nbg=1 要求 groups_per_tile 为偶数（待 pack SF 行数 = groups_per_tile*1 须能被 pack_factor=2 整除），
	84	+ # 否则 num_packed_rows = groups_per_tile*1//2 不正确（num_per_tokens=128 时 groups_per_tile=1 触发）。
	85	+ num_blocks_per_group = (
	86	+ 2
	87	+ if in_config.with_sf
	88	+ else (1 if (hidden >= bf16_nbg1_min_hidden and groups_per_tile % 2 == 0 and groups_per_tile > 0) else 2)
	89	+ )
	90	+ num_packed_rows = groups_per_tile * num_blocks_per_group // pack_factor
	91	+ else:
	92	+ num_blocks_per_group = 1
	93	+
	94	+ num_cores = get_num_vector_cores()
	95	+
	96	+ num_tokens = T.dynamic('num_tokens')
	97	+ sf_stride = T.dynamic('sf_stride')
	98	+ sf_shape = get_sf_shape((num_tokens, hidden), out_config)
	99	+ x_sf_shape = get_sf_shape((num_tokens, hidden), in_config)
	100	+ packed_col_major_input = in_config.use_tma_aligned_col_major_sf and in_config.use_packed_ue8m0
	101	+ if packed_col_major_input:
	102	+ x_sf_shape = (x_sf_shape[0], x_sf_shape[1] // 2)
	103	+ x_sf_dtype = T.int16 if packed_col_major_input else in_config.sf_dtype
	104	+
	105	+ @T.macro
	106	+ def compute_scale(amax):
	107	+ amax = S.vmax(amax, S.vdup(out_config.clamp_min_value, T.float32))
	108	+ if not out_config.round_sf:
	109	+ scale = S.vdiv(amax, S.vdup(quant_max, T.float32))
	110	+ sf_inv = S.vdiv(S.vdup(quant_max, T.float32), amax)
	111	+ return scale, sf_inv
	112	+ scale_raw = S.vmul(amax, S.vdup(1.0 / quant_max, T.float32))
	113	+ scale_bits = T.reinterpret(scale_raw, 'uint32x64')
	114	+ scale_exponent = S.vadds(S.vshrs(S.vsub(scale_bits, S.vdup(1, T.uint32)), 23), 1)
	115	+ inverse_exponent = S.vsub(S.vdup(254, T.uint32), scale_exponent)
	116	+ inverse = T.reinterpret(S.vshls(inverse_exponent, 23), 'float32x64')
	117	+ if out_config.use_packed_ue8m0:
	118	+ return scale_exponent, inverse
	119	+ scale = T.reinterpret(S.vshls(scale_exponent, 23), 'float32x64')
	120	+ return scale, inverse
	121	+
	122	+ @T.macro
	123	+ def quantize_block_bf16(x_ub, o_ub, sf_dst, sf_input, dequant_ub, dst_base=0):
	124	+ with T.SimdVF():
	125	+ sf_dist = 'PK4_B32' if out_config.use_packed_ue8m0 else 'NORM_B32'
	126	+ zero_bf16 = S.vdup(0.0, T.bfloat16)
	127	+ abs_mask = T.reinterpret(S.vdup(0x7FFF, T.uint16), 'bfloat16x128')
	128	+ amax_values = S.alloc_local((2,), T.bfloat16)
	129	+ inverse = S.alloc_local((2,), T.float32)
	130	+ if in_config.with_sf:
	131	+ lane_channel = T.reinterpret(S.vshrs(S.vci(0, T.int16), 5), 'uint16x128') # 0x32 1x32 2x32 3x32
	132	+ # num_per_channels=32 时，两行 packed SF 用 vsel 按 lane 组对选择；
	133	+ # num_per_channels=128 时，单行 packed SF 用 chunk*8 选字节。
	134	+ if in_config.use_tma_aligned_col_major_sf and num_per_channels != 128:
	135	+ use_pair1 = S.vcmps(lane_channel, 2, op='ge') # 0x64 1x64
	136	+ byte_shift = T.reinterpret(
	137	+ S.vshls(S.vand(lane_channel, S.vdup(1, T.uint16)), 3), 'int16x128'
	138	+ ) # 0x32 8x32 0x32 8x32
	139	+ byte_mask = S.vdup(0x00FF, T.int16)
	140	+ elif in_config.use_tma_aligned_col_major_sf:
	141	+ byte_mask = S.vdup(0x00FF, T.int16)
	142	+ for group in T.serial(groups_per_tile):
	143	+ row_base = group * num_per_tokens
	144	+ for chunk in T.serial(block_k // 128):
	145	+ col = chunk * 128
	146	+ amax = S.alloc_var(T.bfloat16)
	147	+ if in_config.with_sf:
	148	+ amax = S.vdup(0.0, T.bfloat16)
	149	+ # num_per_channels=128 时，每个 chunk 对应一个 SF 组，直接用 chunk 索引；
	150	+ # num_per_channels=32 时，用 lane_channel 派生索引。
	151	+ if num_per_channels == 128:
	152	+ scale_index = S.vdup(chunk, T.uint16)
	153	+ else:
	154	+ scale_index = S.vadds(lane_channel, chunk * 4)
	155	+ for row in T.serial(num_per_tokens):
	156	+ raw0 = S.vcvt(S.vld(x_ub[row_base + row, col], dist='UNPK4_B8'), T.float32)
	157	+ raw1 = S.vcvt(S.vld(x_ub[row_base + row, col + 64], dist='UNPK4_B8'), T.float32)
	158	+ _, raw_bf16 = S.vdintlv(
	159	+ T.reinterpret(raw0, 'bfloat16x128'), T.reinterpret(raw1, 'bfloat16x128')
	160	+ )
	161	+ if in_config.use_tma_aligned_col_major_sf:
	162	+ sf_row = row_base + row
	163	+ # num_per_channels=128 时，sf_in_ub 仅一行，从 sf_input[0]
	164	+ # 读取并用 chunk*8 选字节；num_per_channels=32 时使用双行 vsel。
	165	+ if num_per_channels == 128:
	166	+ packed = S.vld(sf_input[0, sf_row], dist='BRC_B16')
	167	+ byte_shift = S.vdup(chunk * 8, T.int16)
	168	+ exponents = S.vand(S.vshr(packed, byte_shift), byte_mask)
	169	+ scale_bits = S.vshls(exponents, 7)
	170	+ else:
	171	+ packed0 = S.vld(sf_input[chunk * 2, sf_row], dist='BRC_B16')
	172	+ packed1 = S.vld(sf_input[chunk * 2 + 1, sf_row], dist='BRC_B16')
	173	+ packed = S.vsel(packed1, packed0, use_pair1)
	174	+ exponents = S.vand(S.vshr(packed, byte_shift), byte_mask)
	175	+ scale_bits = S.vshls(exponents, 7)
	176	+ else:
	177	+ exponents = T.reinterpret(
	178	+ S.vld(sf_input[row_base + row, 0], dist='UNPK_B8'), 'uint16x128'
	179	+ )
	180	+ scale_bits = S.vshls(S.vselr(exponents, scale_index), 7)
	181	+ values = S.vmul(raw_bf16, T.reinterpret(scale_bits, 'bfloat16x128'))
	182	+ S.vsts(dequant_ub[row, 0], values, dist='NORM_B16')
	183	+ amax = S.vmax(amax, S.vand(values, abs_mask))
	184	+ else:
	185	+ amax_acc = S.alloc_local((4,), T.bfloat16)
	186	+ for acc_i in T.unroll(4, explicit=True):
	187	+ amax_acc[acc_i] = S.vdup(0.0, T.bfloat16)
	188	+ for row4 in T.serial(num_per_tokens // 4):
	189	+ for acc_i in T.unroll(4, explicit=True):
	190	+ row = row4 * 4 + acc_i
	191	+ values = S.vld(x_ub[row_base + row, col])
	192	+ amax_acc[acc_i] = S.vmax(amax_acc[acc_i], S.vand(values, abs_mask))
	193	+ amax = S.vmax(S.vmax(amax_acc[0], amax_acc[1]), S.vmax(amax_acc[2], amax_acc[3]))
	194	+ amax_values[0], amax_values[1] = S.vintlv(zero_bf16, amax)
	195	+ # NOTE: round sf/not round sf path can use different instruction order for faster execution
	196	+ # If manual instruction reorder is required, refer to commit aa73858.
	197	+ for half in T.unroll(2, explicit=True):
	198	+ scale, inverse[half] = compute_scale(T.reinterpret(amax_values[half], 'float32x64'))
	199	+ S.vsts(sf_dst[dst_base + group, col + half * 64], scale, dist=sf_dist)
	200	+ for half in T.unroll(2, explicit=True):
	201	+ for row in T.serial(num_per_tokens):
	202	+ if in_config.with_sf:
	203	+ values = S.vcvt(S.vld(dequant_ub[row, half * 64], dist='UNPK_B16'), T.float32)
	204	+ else:
	205	+ values = S.vcvt(
	206	+ S.vld(x_ub[row_base + row, col + half * 64], dist='UNPK_B16'), T.float32
	207	+ )
	208	+ quantized = S.vmul(values, inverse[half])
	209	+ S.vsts(
	210	+ o_ub[row_base + row, col + half * 64],
	211	+ S.vcvt(quantized, T.float8_e4m3fn),
	212	+ dist='PK4_B32',
	213	+ )
	214	+
	215	+ def dequant_f32_scaled(x_ub, row, col, scale_vec):
	216	+ values = S.vcvt(S.vld(x_ub[row, col], dist='UNPK4_B8'), T.float32)
	217	+ return S.vmul(values, scale_vec)
	218	+
	219	+ @T.macro
	220	+ def quantize_block_f32(x_ub, o_ub, sf_dst, sf_ub, dequant_ub, dst_base=0):
	221	+ with T.SimdVF():
	222	+ sf_dist = 'PK4_B32' if out_config.use_packed_ue8m0 else 'NORM_B32'
	223	+ abs_mask = T.reinterpret(S.vdup(0x7FFFFFFF, T.uint32), 'float32x64')
	224	+ num_subvectors = block_k // 64
	225	+ amax_values = S.alloc_local((num_subvectors,), T.float32)
	226	+ input_scale_values = S.alloc_local((num_subvectors,), T.float32)
	227	+ scale_values = S.alloc_local((num_subvectors,), T.uint32 if out_config.use_packed_ue8m0 else T.float32)
	228	+ if in_config.use_tma_aligned_col_major_sf:
	229	+ sf_lane_channel = T.reinterpret(S.vshrs(S.vci(0, T.int32), 5), 'uint32x64') # 0x32 1x32
	230	+ for g in T.serial(groups_per_tile):
	231	+ row_base = g * num_per_tokens
	232	+ for j in T.unroll(num_subvectors, explicit=True):
	233	+ amax_values[j] = S.vdup(0.0, T.float32)
	234	+ for r in T.serial(num_per_tokens):
	235	+ if in_config.use_tma_aligned_col_major_sf:
	236	+ # num_per_channels=128 时，sf_in_ub 布局为
	237	+ # (block_k//num_per_channels, block_m)，每两个 subvector
	238	+ # 共享一行 SF，并通过 BRC_B32 广播。
	239	+ if num_per_channels == 128:
	240	+ for subvector in T.unroll(num_subvectors, explicit=True):
	241	+ sf_group = subvector // 2
	242	+ input_scale_values[subvector] = S.vld(sf_ub[sf_group, row_base + r], dist='BRC_B32')
	243	+ else:
	244	+ sf_gather_base = S.vadds(S.vmuls(sf_lane_channel, block_m), row_base + r)
	245	+ for subvector in T.unroll(num_subvectors, explicit=True):
	246	+ input_scale_values[subvector] = S.vgather2(
	247	+ sf_ub[0, 0], S.vadds(sf_gather_base, subvector * 2 * block_m)
	248	+ )
	249	+ else:
	250	+ # num_per_channels=128 时，每个 SF 组覆盖 128 列，按
	251	+ # sv//subvectors_per_group 索引行；num_per_channels=32 时用 E2B 广播。
	252	+ if num_per_channels == 128:
	253	+ num_groups = block_k // num_per_channels
	254	+ subvectors_per_group = num_subvectors // num_groups
	255	+ for sv in T.unroll(num_subvectors, explicit=True):
	256	+ input_scale_values[sv] = S.vld(
	257	+ sf_ub[row_base + r, sv // subvectors_per_group], dist='BRC_B32'
	258	+ )
	259	+ else:
	260	+ expanded = S.vld(sf_ub[row_base + r, 0], dist='E2B_B32')
	261	+ low, high = S.vintlv(expanded, expanded)
	262	+ input_scale_values[0], input_scale_values[1] = S.vintlv(low, low)
	263	+ if block_k == 256:
	264	+ input_scale_values[2], input_scale_values[3] = S.vintlv(high, high)
	265	+ for subvector in T.unroll(num_subvectors, explicit=True):
	266	+ col = subvector * 64
	267	+ values = dequant_f32_scaled(x_ub, row_base + r, col, input_scale_values[subvector])
	268	+ S.vsts(dequant_ub[r, col], values)
	269	+ amax_values[subvector] = S.vmax(amax_values[subvector], S.vand(values, abs_mask))
	270	+ if out_config.round_sf:
	271	+ for pair in T.unroll(block_k // 128, explicit=True):
	272	+ for half in T.unroll(2, explicit=True):
	273	+ subvector = pair * 2 + half
	274	+ scale_values[subvector], amax_values[subvector] = compute_scale(amax_values[subvector])
	275	+ S.vsts(sf_dst[dst_base + g, subvector * 64], scale_values[subvector], dist=sf_dist)
	276	+ for r in T.serial(num_per_tokens):
	277	+ for half in T.unroll(2, explicit=True):
	278	+ subvector = pair * 2 + half
	279	+ quantized = S.vmul(S.vld(dequant_ub[r, subvector * 64]), amax_values[subvector])
	280	+ S.vsts(
	281	+ o_ub[row_base + r, subvector * 64],
	282	+ S.vcvt(quantized, T.float8_e4m3fn),
	283	+ dist='PK4_B32',
	284	+ )
	285	+ else:
	286	+ for subvector in T.unroll(num_subvectors, explicit=True):
	287	+ scale, inverse = compute_scale(amax_values[subvector])
	288	+ S.vsts(sf_dst[g, subvector * 64], scale)
	289	+ for r in T.serial(num_per_tokens):
	290	+ q = S.vmul(S.vld(dequant_ub[r, subvector * 64]), inverse)
	291	+ S.vsts(o_ub[row_base + r, subvector * 64], S.vcvt(q, T.float8_e4m3fn), dist='PK4_B32')
	292	+
	293	+ @T.macro
	294	+ def pack_row(exp_src, pk_dst, exp_row, dst_row):
	295	+ with T.SimdVF():
	296	+ for base in T.serial(0, block_k, 256):
	297	+ lo, hi = S.vintlv(S.vld(exp_src[exp_row, base]), S.vld(exp_src[exp_row + 1, base]))
	298	+ S.vsts(pk_dst[dst_row, base * 2], lo, dist='NORM_B8')
	299	+ if block_k - base >= 256:
	300	+ S.vsts(pk_dst[dst_row, base * 2 + 256], hi, dist='NORM_B8')
	301	+
	302	+ @T.prim_func
	303	+ def per_channel_cast_kernel(
	304	+ x: T.Tensor[(num_tokens, hidden), in_config.dtype],
	305	+ out: T.Tensor[(num_tokens, hidden), out_config.dtype],
	306	+ out_sf: T.StridedTensor[sf_shape, (sf_stride, 1), out_config.sf_dtype],
	307	+ x_sf_invs: T.Tensor[x_sf_shape, x_sf_dtype],
	308	+ ):
	309	+ with T.Kernel(num_cores) as core_id:
	310	+ x_ub = T.alloc_shared((block_m, block_k), in_config.dtype)
	311	+ out_ub = T.alloc_shared((block_m, block_k), out_config.dtype)
	312	+ versions = {x_ub: num_stages, out_ub: num_stages}
	313	+
	314	+ if in_config.with_sf:
	315	+ if in_config.use_tma_aligned_col_major_sf:
	316	+ if in_config.use_packed_ue8m0:
	317	+ # block_k 可能小于 num_per_channels*pack_factor（例如
	318	+ # num_per_tokens=128 且 block_k=128），需用 ceil_div 避免分配零行。
	319	+ sf_in_ub = T.alloc_shared((ceil_div(block_k, num_per_channels * pack_factor), block_m), T.int16)
	320	+ else:
	321	+ sf_in_ub = T.alloc_shared((block_k // num_per_channels, block_m), in_config.sf_dtype)
	322	+ elif in_config.use_packed_ue8m0:
	323	+ sf_in_ub = T.alloc_shared((block_m, max(block_k // num_per_channels, 64)), in_config.sf_dtype)
	324	+ else:
	325	+ sf_in_ub = T.alloc_shared((block_m, max(block_k // num_per_channels, 8)), in_config.sf_dtype)
	326	+ dequant_dtype = T.bfloat16 if in_config.use_packed_ue8m0 else T.float32
	327	+ dequant_columns = 128 if in_config.use_packed_ue8m0 else block_k
	328	+ dequant_ub = T.alloc_shared((num_per_tokens, dequant_columns), dequant_dtype)
	329	+ versions[sf_in_ub] = num_stages
	330	+
	331	+ if out_config.use_packed_ue8m0:
	332	+ sf_out_ub = T.alloc_shared((groups_per_tile * num_blocks_per_group, max(block_k, 256)), T.uint8)
	333	+ pk_ub = T.alloc_shared((num_packed_rows, block_k * pack_factor), T.uint8)
	334	+ # num_per_tokens=128 时 UB 紧凑（block_k 减小），版本数从 num_stages+2 降为 num_stages
	335	+ versions[sf_out_ub] = num_stages
	336	+ versions[pk_ub] = num_stages
	337	+ else:
	338	+ sf_out_ub = T.alloc_shared((groups_per_tile, block_k), T.float32)
	339	+ # 同 packed 路径：num_per_tokens=128 UB 紧凑，版本数从 num_stages+2 降为 num_stages
	340	+ versions[sf_out_ub] = num_stages
	341	+ T.annotate_buffer_versions(versions)
	342	+
	343	+ if in_config.with_sf:
	344	+ use_swap = num_hidden_tiles > swap_min_num_hidden_tiles
	345	+ if use_swap:
	346	+ dim_outer, dim_inner = num_hidden_tiles, T.ceildiv(num_tokens, block_m * num_blocks_per_group)
	347	+ else:
	348	+ dim_outer, dim_inner = T.ceildiv(num_tokens, block_m * num_blocks_per_group), num_hidden_tiles
	349	+ for iter_id, tile_id in T.Persistent(
	350	+ [dim_outer, dim_inner],
	351	+ num_cores,
	352	+ core_id,
	353	+ group_size=4,
	354	+ num_stages=num_stages,
	355	+ ):
	356	+ col_offset = (iter_id if use_swap else tile_id) * block_k
	357	+ sf_idx = tile_id if use_swap else iter_id
	358	+ for bg in T.serial(num_blocks_per_group):
	359	+ tile_row = sf_idx * num_blocks_per_group + bg
	360	+ if num_blocks_per_group == 1 or tile_row * block_m < num_tokens:
	361	+ T.copy(x[tile_row * block_m, col_offset], x_ub, l2_cache_ctrl='NOTALLOC_KEEP')
	362	+ if in_config.use_tma_aligned_col_major_sf:
	363	+ T.copy(
	364	+ x_sf_invs[col_offset // num_packed_per_channels, tile_row * block_m],
	365	+ sf_in_ub[: ceil_div(block_k, num_packed_per_channels), :block_m],
	366	+ ) # ceil_div: 防止 block_k < num_per_channels*pack_factor 时拷贝 0 行
	367	+ else:
	368	+ T.copy(
	369	+ x_sf_invs[tile_row * block_m, col_offset // num_per_channels],
	370	+ sf_in_ub[:, : block_k // num_per_channels],
	371	+ )
	372	+ if in_config.use_packed_ue8m0:
	373	+ quantize_block_bf16(
	374	+ x_ub, out_ub, sf_out_ub, sf_in_ub, dequant_ub, dst_base=bg * groups_per_tile
	375	+ )
	376	+ else:
	377	+ quantize_block_f32(
	378	+ x_ub, out_ub, sf_out_ub, sf_in_ub, dequant_ub, dst_base=bg * groups_per_tile
	379	+ )
	380	+ T.copy(out_ub, out[tile_row * block_m, col_offset])
	381	+ if not out_config.use_packed_ue8m0:
	382	+ T.copy(sf_out_ub, out_sf[sf_idx * groups_per_tile, col_offset])
	383	+ else:
	384	+ for p in T.serial(num_packed_rows):
	385	+ if sf_idx * (block_m * num_blocks_per_group) + 2 * p * num_per_tokens < num_tokens:
	386	+ pack_row(sf_out_ub, pk_ub, 2 * p, p)
	387	+ pk_base = T.alloc_var(T.int32, init=sf_idx * num_packed_rows)
	388	+ valid_rows = T.alloc_var(T.int32, init=T.min(num_packed_rows, sf_shape[0] - pk_base))
	389	+ T.copy(pk_ub[0:valid_rows, :], out_sf[pk_base, col_offset * 2])
	390	+ else:
	391	+ use_swap = num_hidden_tiles > swap_min_num_hidden_tiles
	392	+ if use_swap:
	393	+ dim_outer, dim_inner = num_hidden_tiles, T.ceildiv(num_tokens, block_m * num_blocks_per_group)
	394	+ else:
	395	+ dim_outer, dim_inner = T.ceildiv(num_tokens, block_m * num_blocks_per_group), num_hidden_tiles
	396	+ for iter_id, tile_id in T.Persistent(
	397	+ [dim_outer, dim_inner],
	398	+ num_cores,
	399	+ core_id,
	400	+ group_size=4,
	401	+ num_stages=num_stages,
	402	+ ):
	403	+ col_offset = (iter_id if use_swap else tile_id) * block_k
	404	+ sf_idx = tile_id if use_swap else iter_id
	405	+ for bg in T.serial(num_blocks_per_group):
	406	+ tile_row = sf_idx * num_blocks_per_group + bg
	407	+ if num_blocks_per_group == 1 or tile_row * block_m < num_tokens:
	408	+ T.copy(x[tile_row * block_m, col_offset], x_ub, l2_cache_ctrl='NOTALLOC_KEEP')
	409	+ quantize_block_bf16(x_ub, out_ub, sf_out_ub, None, None, dst_base=bg * groups_per_tile)
	410	+ T.copy(out_ub, out[tile_row * block_m, col_offset])
	411	+ if not out_config.use_packed_ue8m0:
	412	+ T.copy(sf_out_ub, out_sf[sf_idx * groups_per_tile, col_offset])
	413	+ else:
	414	+ for p in T.serial(num_packed_rows):
	415	+ if sf_idx * (block_m * num_blocks_per_group) + 2 * p * num_per_tokens < num_tokens:
	416	+ pack_row(sf_out_ub, pk_ub, 2 * p, p)
	417	+ pk_base = T.alloc_var(T.int32, init=sf_idx * num_packed_rows)
	418	+ valid_rows = T.alloc_var(T.int32, init=T.min(num_packed_rows, sf_shape[0] - pk_base))
	419	+ T.copy(pk_ub[0:valid_rows, :], out_sf[pk_base, col_offset * 2])
	420	+
	421	+ return per_channel_cast_kernel

Asrc/cann_ops_tilelang/quant/per_token_cast.py+186-0

@@ -0,0 +1,186 @@
	1	+import torch
	2	+
	3	+from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
	4	+
	5	+from .common import (
	6	+ alloc_scaling_factors,
	7	+ cast_epilogue,
	8	+ get_cast_input_and_config,
	9	+ get_cast_output_config,
	10	+ get_logical_hidden,
	11	+ get_physical_hidden,
	12	+ get_sf_shape,
	13	+)
	14	+from .per_token_cast_kernel import get_per_token_cast_kernel_asc, tall_prequant_tile_pays
	15	+from .types import QuantTensor
	16	+
	17	+
	18	+def per_token_cast_impl(
	19	+ x: torch.Tensor \| QuantTensor,
	20	+ fmt: str,
	21	+ num_per_channels: int,
	22	+ x_block_size: tuple[int, int] \| None = None,
	23	+ use_tma_aligned_col_major_sf: bool = False,
	24	+ round_sf: bool = False,
	25	+ use_packed_ue8m0: bool = False,
	26	+ sf_only: bool = False,
	27	+ stochastic_cast: bool = False,
	28	+ sf: torch.Tensor \| None = None,
	29	+ out_x: torch.Tensor \| None = None,
	30	+ out_sf: torch.Tensor \| None = None,
	31	+) -> torch.Tensor \| QuantTensor:
	32	+ require_ascend()
	33	+ assert fmt in ('e4m3', 'e2m1')
	34	+
	35	+ x, x_sf, in_config = get_cast_input_and_config(x, x_block_size)
	36	+ if x.device.type != 'npu':
	37	+ raise ValueError(f'x must be an NPU tensor, got {x.device}.')
	38	+ out_config = get_cast_output_config(
	39	+ fmt, (1, num_per_channels), use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0
	40	+ )
	41	+
	42	+ assert x.dim() == 2
	43	+ assert x.stride(1) == 1
	44	+ num_tokens, hidden = x.shape
	45	+ hidden = get_logical_hidden(hidden, x.dtype)
	46	+ assert num_per_channels in (16, 32, 64, 128) or (num_per_channels == hidden and hidden % 64 == 0)
	47	+ assert sf is None or out_sf is None
	48	+
	49	+ # The kernel is built for a hidden size but not for a token count, so
	50	+ # the tile-shape question it cannot answer is answered here.
	51	+ kernel = get_per_token_cast_kernel_asc(
	52	+ hidden=hidden,
	53	+ token_stride=get_logical_hidden(x.stride(0), x.dtype),
	54	+ in_config=in_config,
	55	+ out_config=out_config,
	56	+ sf_only=sf_only,
	57	+ cast_only=sf is not None,
	58	+ stochastic_cast=stochastic_cast,
	59	+ small_batch=not tall_prequant_tile_pays(num_tokens, hidden),
	60	+ )
	61	+
	62	+ if should_print_kernel_source():
	63	+ print(kernel.get_kernel_source())
	64	+
	65	+ if out_x is None:
	66	+ out_x = torch.empty(
	67	+ (num_tokens, get_physical_hidden(hidden, out_config.torch_dtype)),
	68	+ dtype=out_config.torch_dtype,
	69	+ device=x.device,
	70	+ )
	71	+ if out_sf is None:
	72	+ out_sf = (
	73	+ sf
	74	+ if sf is not None
	75	+ else alloc_scaling_factors(
	76	+ (num_tokens, hidden),
	77	+ out_config=out_config,
	78	+ device=x.device,
	79	+ )
	80	+ )
	81	+ else:
	82	+ sf_shape = get_sf_shape((num_tokens, hidden), out_config)
	83	+ assert not use_tma_aligned_col_major_sf, 'Not supported yet'
	84	+ if use_packed_ue8m0:
	85	+ out_sf = out_sf.view(dtype=torch.uint8)
	86	+ assert out_sf.shape == sf_shape
	87	+
	88	+ if num_tokens > 0:
	89	+ kernel(x, x_sf, out_x, out_sf)
	90	+
	91	+ out_sf = cast_epilogue(out_sf, out_config)
	92	+
	93	+ if sf_only:
	94	+ return out_sf
	95	+ if sf is not None:
	96	+ return out_x
	97	+ return out_x, out_sf
	98	+
	99	+
	100	+def per_token_cast(
	101	+ x: torch.Tensor \| QuantTensor,
	102	+ fmt: str,
	103	+ num_per_channels: int,
	104	+ x_block_size: tuple[int, int] \| None = None,
	105	+ use_tma_aligned_col_major_sf: bool = False,
	106	+ round_sf: bool = False,
	107	+ use_packed_ue8m0: bool = False,
	108	+ stochastic_cast: bool = False,
	109	+ out: QuantTensor \| None = None,
	110	+) -> QuantTensor:
	111	+ """Cast a matrix to FP8/FP4 with per-token (row-wise) scaling factors.
	112	+
	113	+ Args:
	114	+ x: Input 2D tensor or a tuple of 2D tensor and 2D scaling factors.
	115	+ fmt: Target quantized format (``'e4m3'`` or ``'e2m1'``).
	116	+ num_per_channels: Number of channels in each scaling block.
	117	+ x_block_size: Input scaling block size for pre-quantized inputs.
	118	+ use_tma_aligned_col_major_sf: Whether to use TMA-aligned column-major sf factors.
	119	+ round_sf: Whether to round scaling factors to powers of two.
	120	+ use_packed_ue8m0: Whether to use packed UE8M0 format for sf factors.
	121	+ stochastic_cast: Whether to use stochastic rounding for the narrowing cast.
	122	+ out: Optional pre-allocated QuantTensor (data, sf).
	123	+
	124	+ Returns:
	125	+ A tuple ``(out, out_sf)`` with quantized output and sf factors.
	126	+ """
	127	+ out_x, out_sf = (None, None) if out is None else out
	128	+ return per_token_cast_impl(
	129	+ x,
	130	+ fmt,
	131	+ num_per_channels,
	132	+ x_block_size,
	133	+ use_tma_aligned_col_major_sf,
	134	+ round_sf,
	135	+ use_packed_ue8m0,
	136	+ stochastic_cast=stochastic_cast,
	137	+ out_x=out_x,
	138	+ out_sf=out_sf,
	139	+ )
	140	+
	141	+
	142	+def per_token_cast_with_sf_only(
	143	+ x: torch.Tensor,
	144	+ fmt: str,
	145	+ num_per_channels: int,
	146	+ use_tma_aligned_col_major_sf: bool = False,
	147	+ round_sf: bool = False,
	148	+ use_packed_ue8m0: bool = False,
	149	+ out_sf: torch.Tensor \| None = None,
	150	+) -> torch.Tensor:
	151	+ """Return only the per-token scaling factors for an FP8/FP4 cast."""
	152	+ return per_token_cast_impl(
	153	+ x,
	154	+ fmt,
	155	+ num_per_channels,
	156	+ None,
	157	+ use_tma_aligned_col_major_sf,
	158	+ round_sf,
	159	+ use_packed_ue8m0,
	160	+ sf_only=True,
	161	+ out_sf=out_sf,
	162	+ )
	163	+
	164	+
	165	+def per_token_cast_with_precomputed_sf(
	166	+ x: torch.Tensor,
	167	+ fmt: str,
	168	+ num_per_channels: int,
	169	+ sf: torch.Tensor,
	170	+ use_tma_aligned_col_major_sf: bool = False,
	171	+ round_sf: bool = False,
	172	+ use_packed_ue8m0: bool = False,
	173	+ out: torch.Tensor \| None = None,
	174	+) -> torch.Tensor:
	175	+ """Cast with caller-provided per-token scaling factors."""
	176	+ return per_token_cast_impl(
	177	+ x,
	178	+ fmt,
	179	+ num_per_channels,
	180	+ None,
	181	+ use_tma_aligned_col_major_sf,
	182	+ round_sf,
	183	+ use_packed_ue8m0,
	184	+ sf=sf,
	185	+ out_x=out,
	186	+ )

Asrc/cann_ops_tilelang/quant/per_token_cast_and_cast_back.py+50-0

@@ -0,0 +1,50 @@
	1	+import torch
	2	+from tilelang.ascend import language as T
	3	+
	4	+from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
	5	+
	6	+from .common import get_cast_output_config
	7	+from .per_token_cast_and_cast_back_kernel import get_per_token_cast_and_cast_back_kernel_asc
	8	+
	9	+
	10	+def per_token_cast_and_cast_back(
	11	+ x: torch.Tensor,
	12	+ fmt: str,
	13	+ num_per_channels: int,
	14	+ round_sf: bool = False,
	15	+):
	16	+ """Apply an in-place lossy per-token cast-and-cast-back to ``x``.
	17	+
	18	+ Args:
	19	+ x: Input 2D tensor of shape (num_tokens, hidden). The innermost stride must
	20	+ be 1; ``x.stride(0)`` may exceed ``hidden`` (non-contiguous rows allowed).
	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).
	23	+ round_sf: Whether to round scaling factors to powers of two.
	24	+ """
	25	+ require_ascend()
	26	+ if x.device.type != 'npu':
	27	+ raise ValueError(f'x must be an NPU tensor, got {x.device}.')
	28	+ assert x.dim() == 2
	29	+ assert x.stride(1) == 1
	30	+ assert fmt in ('e4m3', 'e2m1')
	31	+ assert num_per_channels in (32, 64, 128)
	32	+
	33	+ num_tokens, hidden = x.shape
	34	+ assert hidden % num_per_channels == 0
	35	+
	36	+ out_config = get_cast_output_config(fmt, (1, num_per_channels), round_sf=round_sf)
	37	+ assert out_config.sf_block[0] == 1
	38	+
	39	+ kernel = get_per_token_cast_and_cast_back_kernel_asc(
	40	+ hidden=hidden,
	41	+ token_stride=x.stride(0),
	42	+ dtype=T.dtype(x.dtype),
	43	+ out_config=out_config,
	44	+ )
	45	+
	46	+ if should_print_kernel_source():
	47	+ print(kernel.get_kernel_source())
	48	+
	49	+ if num_tokens > 0:
	50	+ kernel(x)