feat(mhc): migrate operator suite-ops-tilelang-AtomGit

Mdocs/architecture.md+26-0

	@@ -94,6 +94,32 @@ engram_fused_weight(weight_hidden, weight_embed)
94	-> update grad_weight_hidden and grad_weight_embed in place	94	-> update grad_weight_hidden and grad_weight_embed in place
95	```	95	```
96		96
		97	+MHC 训练链路由显式前向和反向原语组成：
		98	+
		99	+```text
		100	+expand_to_mhc_forward(...)
		101	+ -> mhc_fn_normw_merge_forward(...)
		102	+ -> mhc_gemm_with_sqrsum_forward(...)
		103	+ -> mhc_reduce_partials_and_rmsnorm_forward(...)
		104	+ -> mhc_pre_split_mixes_forward(...)
		105	+ -> mhc_sinkhorn_forward(...)
		106	+ -> mhc_pre_apply_mix_forward(...)
		107	+ -> upper-layer computation
		108	+ -> mhc_post_forward(...)
		109	+
		110	+mhc_post_backward(...)
		111	+ -> upper-layer backward
		112	+ -> mhc_pre_apply_mix_backward(...)
		113	+ -> mhc_sinkhorn_backward(...)
		114	+ -> mhc_pre_split_mixes_backward(...)
		115	+ -> mhc_reduce_partials_and_rmsnorm_backward(...)
		116	+ -> mhc_gemm_with_sqrsum_backward(...)
		117	+ -> mhc_fn_normw_merge_backward(...)
		118	+ -> expand_to_mhc_backward(...)
		119	+```
		120	+
		121	+`mhc_head_compute_mix_forward` 和 `mhc_head_compute_mix_backward` 处理 head mixing 系数，`mhc_multilayer_recompute` 融合多层 pre/post mix 重计算。`mhc_pre_big_fuse` 融合 reduce/RMSNorm、split mixes、Sinkhorn 和 pre apply mix 的前向路径。
		122	+
97	## 扩展方向	123	## 扩展方向
98		124
99	- `src/cann_ops_tilelang/quant/`：公共 scaling-factor layout 和 FP8/FP4 类型系统；	125	- `src/cann_ops_tilelang/quant/`：公共 scaling-factor layout 和 FP8/FP4 类型系统；

Msrc/cann_ops_tilelang/__init__.py+42-0

	@@ -9,6 +9,26 @@ __all__ = [
9	'engram_sinkhorn_finalize',	9	'engram_sinkhorn_finalize',
10	'engram_sinkhorn_momentum_update',	10	'engram_sinkhorn_momentum_update',
11	'engram_sinkhorn_step',	11	'engram_sinkhorn_step',
		12	+ 'expand_to_mhc_backward',
		13	+ 'expand_to_mhc_forward',
		14	+ 'mhc_fn_normw_merge_backward',
		15	+ 'mhc_fn_normw_merge_forward',
		16	+ 'mhc_gemm_with_sqrsum_backward',
		17	+ 'mhc_gemm_with_sqrsum_forward',
		18	+ 'mhc_head_compute_mix_backward',
		19	+ 'mhc_head_compute_mix_forward',
		20	+ 'mhc_multilayer_recompute',
		21	+ 'mhc_post_backward',
		22	+ 'mhc_post_forward',
		23	+ 'mhc_pre_apply_mix_backward',
		24	+ 'mhc_pre_apply_mix_forward',
		25	+ 'mhc_pre_big_fuse',
		26	+ 'mhc_pre_split_mixes_backward',
		27	+ 'mhc_pre_split_mixes_forward',
		28	+ 'mhc_reduce_partials_and_rmsnorm_backward',
		29	+ 'mhc_reduce_partials_and_rmsnorm_forward',
		30	+ 'mhc_sinkhorn_backward',
		31	+ 'mhc_sinkhorn_forward',
12	'moe_topk_gate',	32	'moe_topk_gate',
13	'normalize_weight',	33	'normalize_weight',
14	'topk_gate',	34	'topk_gate',
	@@ -24,4 +44,26 @@ from .engram import (
24	engram_sinkhorn_momentum_update,	44	engram_sinkhorn_momentum_update,
25	engram_sinkhorn_step,	45	engram_sinkhorn_step,
26	)	46	)
		47	+from .mhc import (
		48	+ expand_to_mhc_backward,
		49	+ expand_to_mhc_forward,
		50	+ mhc_fn_normw_merge_backward,
		51	+ mhc_fn_normw_merge_forward,
		52	+ mhc_gemm_with_sqrsum_backward,
		53	+ mhc_gemm_with_sqrsum_forward,
		54	+ mhc_head_compute_mix_backward,
		55	+ mhc_head_compute_mix_forward,
		56	+ mhc_multilayer_recompute,
		57	+ mhc_post_backward,
		58	+ mhc_post_forward,
		59	+ mhc_pre_apply_mix_backward,
		60	+ mhc_pre_apply_mix_forward,
		61	+ mhc_pre_big_fuse,
		62	+ mhc_pre_split_mixes_backward,
		63	+ mhc_pre_split_mixes_forward,
		64	+ mhc_reduce_partials_and_rmsnorm_backward,
		65	+ mhc_reduce_partials_and_rmsnorm_forward,
		66	+ mhc_sinkhorn_backward,
		67	+ mhc_sinkhorn_forward,
		68	+)
27	from .moe import moe_topk_gate, normalize_weight, topk_gate	69	from .moe import moe_topk_gate, normalize_weight, topk_gate

Asrc/cann_ops_tilelang/mhc/__init__.py+41-0

@@ -0,0 +1,41 @@
	1	+"""MHC operators."""
	2	+
	3	+__all__ = [
	4	+ 'expand_to_mhc_backward',
	5	+ 'expand_to_mhc_forward',
	6	+ 'mhc_fn_normw_merge_backward',
	7	+ 'mhc_fn_normw_merge_forward',
	8	+ 'mhc_gemm_with_sqrsum_backward',
	9	+ 'mhc_gemm_with_sqrsum_forward',
	10	+ 'mhc_head_compute_mix_backward',
	11	+ 'mhc_head_compute_mix_forward',
	12	+ 'mhc_multilayer_recompute',
	13	+ 'mhc_post_backward',
	14	+ 'mhc_post_forward',
	15	+ 'mhc_pre_apply_mix_backward',
	16	+ 'mhc_pre_apply_mix_forward',
	17	+ 'mhc_pre_big_fuse',
	18	+ 'mhc_pre_split_mixes_backward',
	19	+ 'mhc_pre_split_mixes_forward',
	20	+ 'mhc_reduce_partials_and_rmsnorm_backward',
	21	+ 'mhc_reduce_partials_and_rmsnorm_forward',
	22	+ 'mhc_sinkhorn_backward',
	23	+ 'mhc_sinkhorn_forward',
	24	+]
	25	+
	26	+from .expand import expand_to_mhc_backward, expand_to_mhc_forward
	27	+from .head_compute_mix import mhc_head_compute_mix_backward, mhc_head_compute_mix_forward
	28	+from .multilayer_recompute import mhc_multilayer_recompute
	29	+from .norm_fn import (
	30	+ mhc_fn_normw_merge_backward,
	31	+ mhc_fn_normw_merge_forward,
	32	+ mhc_gemm_with_sqrsum_backward,
	33	+ mhc_gemm_with_sqrsum_forward,
	34	+ mhc_reduce_partials_and_rmsnorm_backward,
	35	+ mhc_reduce_partials_and_rmsnorm_forward,
	36	+)
	37	+from .post import mhc_post_backward, mhc_post_forward
	38	+from .pre_apply_mix import mhc_pre_apply_mix_backward, mhc_pre_apply_mix_forward
	39	+from .pre_big_fuse import mhc_pre_big_fuse
	40	+from .pre_split_mixes import mhc_pre_split_mixes_backward, mhc_pre_split_mixes_forward
	41	+from .sinkhorn import mhc_sinkhorn_backward, mhc_sinkhorn_forward

Asrc/cann_ops_tilelang/mhc/expand.py+128-0

@@ -0,0 +1,128 @@
	1	+"""Public interfaces for expanding activations across MHC streams."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import torch
	6	+
	7	+from cann_ops_tilelang.config import get_num_vector_cores, require_ascend
	8	+
	9	+from .expand_kernel import expand_to_mhc_backward_kernel, expand_to_mhc_forward_kernel
	10	+
	11	+_FORWARD_HIDDEN_ALIGNMENT = 128
	12	+_BACKWARD_HIDDEN_ALIGNMENT = 64
	13	+_FORWARD_NUM_STAGES = 2
	14	+_BACKWARD_NUM_STAGES = 3
	15	+_BACKWARD_HIDDEN_BLOCK = 4096
	16	+
	17	+
	18	+def _check_npu_bfloat16_tensor(
	19	+ name: str,
	20	+ tensor: torch.Tensor,
	21	+ *,
	22	+ shape: tuple[int, ...] \| None = None,
	23	+ device: torch.device \| None = None,
	24	+) -> None:
	25	+ if not isinstance(tensor, torch.Tensor):
	26	+ raise TypeError(f'{name} must be a torch.Tensor.')
	27	+ if tensor.device.type != 'npu':
	28	+ raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')
	29	+ if device is not None and tensor.device != device:
	30	+ raise ValueError(f'{name} must be on {device}, got {tensor.device}.')
	31	+ if tensor.dtype != torch.bfloat16:
	32	+ raise TypeError(f'{name} must use torch.bfloat16, got {tensor.dtype}.')
	33	+ if shape is not None and tuple(tensor.shape) != shape:
	34	+ raise ValueError(f'{name} must have shape {shape}, got {tuple(tensor.shape)}.')
	35	+ if not tensor.is_contiguous():
	36	+ raise ValueError(f'{name} must be contiguous.')
	37	+
	38	+
	39	+def _shares_storage(left: torch.Tensor, right: torch.Tensor) -> bool:
	40	+ return (
	41	+ left.numel() > 0
	42	+ and right.numel() > 0
	43	+ and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr()
	44	+ )
	45	+
	46	+
	47	+def _forward_hidden_block(num_tokens: int, hidden_size: int) -> int:
	48	+ is_short = num_tokens <= 512
	49	+ if is_short and hidden_size >= 7168:
	50	+ return 1024
	51	+ if is_short or hidden_size <= 4096:
	52	+ return 512
	53	+ return 1024
	54	+
	55	+
	56	+def expand_to_mhc_forward(x: torch.Tensor, output: torch.Tensor) -> None:
	57	+ """Copy every BF16 activation row into each MHC stream on Ascend."""
	58	+
	59	+ require_ascend()
	60	+ if not isinstance(x, torch.Tensor):
	61	+ raise TypeError('x must be a torch.Tensor.')
	62	+ if x.ndim != 2:
	63	+ raise ValueError(f'x must have shape (num_tokens, hidden_size), got {tuple(x.shape)}.')
	64	+ num_tokens, hidden_size = x.shape
	65	+ if hidden_size <= 0 or hidden_size % _FORWARD_HIDDEN_ALIGNMENT:
	66	+ raise ValueError(f'hidden_size must be a positive multiple of {_FORWARD_HIDDEN_ALIGNMENT}, got {hidden_size}.')
	67	+ if not isinstance(output, torch.Tensor):
	68	+ raise TypeError('output must be a torch.Tensor.')
	69	+ if output.ndim != 3:
	70	+ raise ValueError('output must have shape (num_tokens, mhc_mult, hidden_size).')
	71	+ mhc_mult = output.shape[1]
	72	+ if mhc_mult <= 0:
	73	+ raise ValueError(f'mhc_mult must be positive, got {mhc_mult}.')
	74	+
	75	+ _check_npu_bfloat16_tensor('x', x)
	76	+ _check_npu_bfloat16_tensor(
	77	+ 'output',
	78	+ output,
	79	+ shape=(num_tokens, mhc_mult, hidden_size),
	80	+ device=x.device,
	81	+ )
	82	+ if _shares_storage(output, x):
	83	+ raise ValueError('output and x must not share storage.')
	84	+ if num_tokens == 0:
	85	+ return
	86	+
	87	+ expand_to_mhc_forward_kernel(
	88	+ x,
	89	+ output,
	90	+ hidden_block=_forward_hidden_block(num_tokens, hidden_size),
	91	+ num_stages=_FORWARD_NUM_STAGES,
	92	+ num_vec_cores=get_num_vector_cores(),
	93	+ )
	94	+
	95	+
	96	+def expand_to_mhc_backward(output_grad: torch.Tensor, x_grad: torch.Tensor) -> None:
	97	+ """Sum BF16 MHC stream gradients into a caller-owned activation gradient."""
	98	+
	99	+ require_ascend()
	100	+ if not isinstance(output_grad, torch.Tensor):
	101	+ raise TypeError('output_grad must be a torch.Tensor.')
	102	+ if output_grad.ndim != 3:
	103	+ raise ValueError('output_grad must have shape (num_tokens, mhc_mult, hidden_size).')
	104	+ num_tokens, mhc_mult, hidden_size = output_grad.shape
	105	+ if mhc_mult <= 0:
	106	+ raise ValueError(f'mhc_mult must be positive, got {mhc_mult}.')
	107	+ if hidden_size <= 0 or hidden_size % _BACKWARD_HIDDEN_ALIGNMENT:
	108	+ raise ValueError(f'hidden_size must be a positive multiple of {_BACKWARD_HIDDEN_ALIGNMENT}, got {hidden_size}.')
	109	+
	110	+ _check_npu_bfloat16_tensor('output_grad', output_grad)
	111	+ _check_npu_bfloat16_tensor(
	112	+ 'x_grad',
	113	+ x_grad,
	114	+ shape=(num_tokens, hidden_size),
	115	+ device=output_grad.device,
	116	+ )
	117	+ if _shares_storage(x_grad, output_grad):
	118	+ raise ValueError('x_grad and output_grad must not share storage.')
	119	+ if num_tokens == 0:
	120	+ return
	121	+
	122	+ expand_to_mhc_backward_kernel(
	123	+ output_grad,
	124	+ x_grad,
	125	+ hidden_block=_BACKWARD_HIDDEN_BLOCK,
	126	+ num_stages=_BACKWARD_NUM_STAGES,
	127	+ num_vec_cores=get_num_vector_cores(),
	128	+ )

Asrc/cann_ops_tilelang/mhc/expand_kernel.py+97-0

@@ -0,0 +1,97 @@
	1	+"""TileLang kernels for expanding activations across MHC streams."""
	2	+
	3	+import tilelang
	4	+from tilelang.ascend import language as T
	5	+
	6	+_VECTOR_WIDTH = 64
	7	+_FORWARD_TOKEN_BLOCK = 32
	8	+
	9	+_PASS_CONFIGS = {
	10	+ tilelang.PassConfigKey.TL_DISABLE_LOOP_UNSWITCHING: True,
	11	+}
	12	+
	13	+
	14	+@tilelang.jit(pass_configs=_PASS_CONFIGS)
	15	+def expand_to_mhc_forward_kernel(x, output, hidden_block: int, num_stages: int, num_vec_cores: int):
	16	+ num_tokens = T.dynamic('num_tokens')
	17	+ hidden_size, mhc_mult = T.const('hidden_size, mhc_mult')
	18	+ x: T.Tensor[(num_tokens, hidden_size), T.bfloat16]
	19	+ output: T.Tensor[(num_tokens, mhc_mult, hidden_size), T.bfloat16]
	20	+
	21	+ with T.Kernel(num_vec_cores) as core_id:
	22	+ x_ub = T.alloc_shared((_FORWARD_TOKEN_BLOCK, hidden_block), T.bfloat16)
	23	+ T.annotate_buffer_versions({x_ub: num_stages})
	24	+
	25	+ for token_index, hidden_index in T.Persistent(
	26	+ [T.ceildiv(num_tokens, _FORWARD_TOKEN_BLOCK), T.ceildiv(hidden_size, hidden_block)],
	27	+ num_vec_cores,
	28	+ core_id,
	29	+ group_size=1,
	30	+ num_stages=num_stages,
	31	+ ):
	32	+ row_base = token_index * _FORWARD_TOKEN_BLOCK
	33	+ column_base = hidden_index * hidden_block
	34	+ T.copy(x[row_base, column_base], x_ub)
	35	+ for mhc_index in T.serial(mhc_mult):
	36	+ T.copy(
	37	+ x_ub,
	38	+ output[
	39	+ row_base : row_base + _FORWARD_TOKEN_BLOCK,
	40	+ mhc_index,
	41	+ column_base : column_base + hidden_block,
	42	+ ],
	43	+ )
	44	+
	45	+
	46	+@tilelang.jit(pass_configs=_PASS_CONFIGS)
	47	+def expand_to_mhc_backward_kernel(
	48	+ output_grad,
	49	+ x_grad,
	50	+ hidden_block: int,
	51	+ num_stages: int,
	52	+ num_vec_cores: int,
	53	+):
	54	+ num_tokens = T.dynamic('num_tokens')
	55	+ hidden_size, mhc_mult = T.const('hidden_size, mhc_mult')
	56	+ output_grad: T.Tensor[(num_tokens, mhc_mult, hidden_size), T.bfloat16]
	57	+ x_grad: T.Tensor[(num_tokens, hidden_size), T.bfloat16]
	58	+
	59	+ with T.Kernel(num_vec_cores) as core_id:
	60	+ output_grad_ub = T.alloc_shared((mhc_mult, hidden_block), T.bfloat16)
	61	+ x_grad_ub = T.alloc_shared((hidden_block,), T.bfloat16)
	62	+ T.annotate_buffer_versions({output_grad_ub: num_stages, x_grad_ub: num_stages})
	63	+
	64	+ for token_index, hidden_index in T.Persistent(
	65	+ [num_tokens, T.ceildiv(hidden_size, hidden_block)],
	66	+ num_vec_cores,
	67	+ core_id,
	68	+ group_size=1,
	69	+ num_stages=num_stages,
	70	+ ):
	71	+ T.copy(output_grad[token_index, 0, hidden_index * hidden_block], output_grad_ub)
	72	+
	73	+ with T.SimdVF():
	74	+ accumulator = T.simd.alloc_var('float32')
	75	+ ones = T.simd.vdup(1.0, T.float32)
	76	+ num_chunks = hidden_block // _VECTOR_WIDTH
	77	+ for chunk in range(num_chunks):
	78	+ column = chunk * _VECTOR_WIDTH
	79	+ accumulator = T.simd.vcvt(
	80	+ T.simd.vld(output_grad_ub[0, column], dist='UNPK_B16'),
	81	+ T.float32,
	82	+ part=0,
	83	+ )
	84	+ for mhc_index in T.unroll(mhc_mult - 1, explicit=True):
	85	+ output_grad_value = T.simd.vcvt(
	86	+ T.simd.vld(output_grad_ub[mhc_index + 1, column], dist='UNPK_B16'),
	87	+ T.float32,
	88	+ part=0,
	89	+ )
	90	+ T.simd.vmula(accumulator, output_grad_value, ones)
	91	+ T.simd.vsts(
	92	+ x_grad_ub[column],
	93	+ T.simd.vcvt(accumulator, T.bfloat16),
	94	+ dist='PK_B32',
	95	+ )
	96	+
	97	+ T.copy(x_grad_ub, x_grad[token_index, hidden_index * hidden_block])

Asrc/cann_ops_tilelang/mhc/head_compute_mix.py+199-0

@@ -0,0 +1,199 @@
	1	+"""Public interfaces for the MHC head mix transform."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import math
	6	+from numbers import Real
	7	+
	8	+import torch
	9	+
	10	+from cann_ops_tilelang.config import get_num_vector_cores, require_ascend
	11	+
	12	+from .head_compute_mix_kernel import (
	13	+ mhc_head_compute_mix_backward_kernel,
	14	+ mhc_head_compute_mix_forward_kernel,
	15	+)
	16	+
	17	+_MHC_MULT = 4
	18	+_MAX_TOKEN_BLOCK_SIZE = 64
	19	+_ASCEND_TOKEN_BLOCK_SIZE = 512
	20	+
	21	+
	22	+def _check_npu_fp32_tensor(
	23	+ name: str,
	24	+ tensor: torch.Tensor,
	25	+ *,
	26	+ shape: tuple[int, ...],
	27	+ device: torch.device \| None = None,
	28	+) -> None:
	29	+ if not isinstance(tensor, torch.Tensor):
	30	+ raise TypeError(f'{name} must be a torch.Tensor.')
	31	+ if tensor.device.type != 'npu':
	32	+ raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')
	33	+ if device is not None and tensor.device != device:
	34	+ raise ValueError(f'{name} must be on {device}, got {tensor.device}.')
	35	+ if tensor.dtype != torch.float32:
	36	+ raise TypeError(f'{name} must use torch.float32, got {tensor.dtype}.')
	37	+ if tuple(tensor.shape) != shape:
	38	+ raise ValueError(f'{name} must have shape {shape}, got {tuple(tensor.shape)}.')
	39	+ if not tensor.is_contiguous():
	40	+ raise ValueError(f'{name} must be contiguous.')
	41	+
	42	+
	43	+def _check_eps(value: float) -> float:
	44	+ if isinstance(value, bool) or not isinstance(value, Real):
	45	+ raise TypeError('mhc_pre_eps must be a real number.')
	46	+ result = float(value)
	47	+ if not math.isfinite(result) or result < 0:
	48	+ raise ValueError('mhc_pre_eps must be finite and non-negative.')
	49	+ return result
	50	+
	51	+
	52	+def _check_num_vec_cores(value: int) -> int:
	53	+ if isinstance(value, bool) or not isinstance(value, int):
	54	+ raise TypeError('num_vec_cores must be an integer.')
	55	+ detected = get_num_vector_cores()
	56	+ if value <= 0 or value > detected:
	57	+ raise ValueError(f'num_vec_cores must be in [1, {detected}], got {value}.')
	58	+ return value
	59	+
	60	+
	61	+def _choose_token_block_size(token_block_size: int) -> int:
	62	+ if isinstance(token_block_size, bool) or not isinstance(token_block_size, int):
	63	+ raise TypeError('token_block_size must be an integer.')
	64	+ if token_block_size < 0:
	65	+ raise ValueError(f'token_block_size must be non-negative, got {token_block_size}.')
	66	+ if token_block_size and (token_block_size > _MAX_TOKEN_BLOCK_SIZE or token_block_size & (token_block_size - 1)):
	67	+ raise ValueError(f'token_block_size must be a power of two no greater than {_MAX_TOKEN_BLOCK_SIZE}.')
	68	+ # Match TileKernels-Nightly's Ascend specialization. The public tuning
	69	+ # argument remains accepted and validated for API compatibility, but the
	70	+ # Ascend implementation always uses its optimized 512-token tile.
	71	+ return _ASCEND_TOKEN_BLOCK_SIZE
	72	+
	73	+
	74	+def _shares_storage(left: torch.Tensor, right: torch.Tensor) -> bool:
	75	+ return (
	76	+ left.numel() > 0
	77	+ and right.numel() > 0
	78	+ and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr()
	79	+ )
	80	+
	81	+
	82	+def _check_output_aliases(
	83	+ inputs: tuple[tuple[str, torch.Tensor], ...],
	84	+ outputs: tuple[tuple[str, torch.Tensor], ...],
	85	+) -> None:
	86	+ for output_index, (output_name, output) in enumerate(outputs):
	87	+ for other_name, other in inputs + outputs[:output_index]:
	88	+ if _shares_storage(output, other):
	89	+ raise ValueError(f'{output_name} and {other_name} must not share storage.')
	90	+
	91	+
	92	+def _validate_inputs(
	93	+ input_mix: torch.Tensor,
	94	+ mhc_scale: torch.Tensor,
	95	+ mhc_base: torch.Tensor,
	96	+) -> tuple[int, torch.device]:
	97	+ if not isinstance(input_mix, torch.Tensor):
	98	+ raise TypeError('input_mix must be a torch.Tensor.')
	99	+ if input_mix.ndim != 2 or input_mix.shape[1] != _MHC_MULT:
	100	+ raise ValueError(f'input_mix must have shape (num_tokens, {_MHC_MULT}).')
	101	+ num_tokens = input_mix.shape[0]
	102	+ device = input_mix.device
	103	+ _check_npu_fp32_tensor('input_mix', input_mix, shape=(num_tokens, _MHC_MULT))
	104	+ _check_npu_fp32_tensor('mhc_scale', mhc_scale, shape=(1,), device=device)
	105	+ _check_npu_fp32_tensor('mhc_base', mhc_base, shape=(_MHC_MULT,), device=device)
	106	+ return num_tokens, device
	107	+
	108	+
	109	+def mhc_head_compute_mix_forward(
	110	+ input_mix: torch.Tensor,
	111	+ mhc_scale: torch.Tensor,
	112	+ mhc_base: torch.Tensor,
	113	+ output_mix: torch.Tensor,
	114	+ mhc_pre_eps: float,
	115	+ token_block_size: int = 0,
	116	+) -> None:
	117	+ """Apply the FP32 affine sigmoid transform for four MHC streams."""
	118	+
	119	+ eps = _check_eps(mhc_pre_eps)
	120	+ require_ascend()
	121	+ num_tokens, device = _validate_inputs(input_mix, mhc_scale, mhc_base)
	122	+ _check_npu_fp32_tensor(
	123	+ 'output_mix',
	124	+ output_mix,
	125	+ shape=(num_tokens, _MHC_MULT),
	126	+ device=device,
	127	+ )
	128	+ _check_output_aliases(
	129	+ (('input_mix', input_mix), ('mhc_scale', mhc_scale), ('mhc_base', mhc_base)),
	130	+ (('output_mix', output_mix),),
	131	+ )
	132	+ selected_token_block_size = _choose_token_block_size(token_block_size)
	133	+ if num_tokens == 0:
	134	+ return
	135	+
	136	+ mhc_head_compute_mix_forward_kernel(
	137	+ input_mix,
	138	+ mhc_scale,
	139	+ mhc_base,
	140	+ output_mix,
	141	+ eps,
	142	+ token_block_size=selected_token_block_size,
	143	+ )
	144	+
	145	+
	146	+def mhc_head_compute_mix_backward(
	147	+ output_mix_grad: torch.Tensor,
	148	+ input_mix: torch.Tensor,
	149	+ mhc_scale: torch.Tensor,
	150	+ mhc_base: torch.Tensor,
	151	+ input_mix_grad: torch.Tensor,
	152	+ mhc_scale_grad_partial: torch.Tensor,
	153	+ mhc_base_grad_partial: torch.Tensor,
	154	+ num_vec_cores: int,
	155	+ token_block_size: int = 0,
	156	+) -> None:
	157	+ """Compute input and partial parameter gradients for the MHC head mix."""
	158	+
	159	+ require_ascend()
	160	+ num_vec_cores = _check_num_vec_cores(num_vec_cores)
	161	+ num_tokens, device = _validate_inputs(input_mix, mhc_scale, mhc_base)
	162	+ tensor_specs = (
	163	+ ('output_mix_grad', output_mix_grad, (num_tokens, _MHC_MULT)),
	164	+ ('input_mix_grad', input_mix_grad, (num_tokens, _MHC_MULT)),
	165	+ ('mhc_scale_grad_partial', mhc_scale_grad_partial, (num_vec_cores, 1)),
	166	+ ('mhc_base_grad_partial', mhc_base_grad_partial, (num_vec_cores, _MHC_MULT)),
	167	+ )
	168	+ for name, tensor, shape in tensor_specs:
	169	+ _check_npu_fp32_tensor(name, tensor, shape=shape, device=device)
	170	+
	171	+ inputs = (
	172	+ ('output_mix_grad', output_mix_grad),
	173	+ ('input_mix', input_mix),
	174	+ ('mhc_scale', mhc_scale),
	175	+ ('mhc_base', mhc_base),
	176	+ )
	177	+ outputs = (
	178	+ ('input_mix_grad', input_mix_grad),
	179	+ ('mhc_scale_grad_partial', mhc_scale_grad_partial),
	180	+ ('mhc_base_grad_partial', mhc_base_grad_partial),
	181	+ )
	182	+ _check_output_aliases(inputs, outputs)
	183	+ selected_token_block_size = _choose_token_block_size(token_block_size)
	184	+ if num_tokens == 0:
	185	+ mhc_scale_grad_partial.zero_()
	186	+ mhc_base_grad_partial.zero_()
	187	+ return
	188	+
	189	+ mhc_head_compute_mix_backward_kernel(
	190	+ output_mix_grad,
	191	+ input_mix,
	192	+ mhc_scale,
	193	+ mhc_base,
	194	+ input_mix_grad,
	195	+ mhc_scale_grad_partial,
	196	+ mhc_base_grad_partial,
	197	+ token_block_size=selected_token_block_size,
	198	+ num_vec_cores=num_vec_cores,
	199	+ )

Asrc/cann_ops_tilelang/mhc/head_compute_mix_kernel.py+230-0

@@ -0,0 +1,230 @@
	1	+"""TileLang kernels for the MHC head mix transform."""
	2	+
	3	+import tilelang
	4	+from tilelang.ascend import language as T
	5	+from tilelang.ascend.language import simd as S
	6	+
	7	+_VECTOR_WIDTH = 64
	8	+_PASS_CONFIGS = {
	9	+ tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
	10	+}
	11	+
	12	+
	13	+@T.macro()
	14	+def _split_accumulators(base_acc, width: int) -> None:
	15	+ """Split interleaved component accumulators once in the reduction tree."""
	16	+ for component in T.unroll(width, explicit=True):
	17	+ low, high = S.vdintlv(base_acc[component], base_acc[component])
	18	+ base_acc[component] = low
	19	+ base_acc[component + width] = high
	20	+
	21	+
	22	+@T.macro
	23	+def _split_accumulator_tree(base_acc, mhc_mult: int) -> None:
	24	+ """Separate interleaved accumulators into one accumulator per component."""
	25	+ if mhc_mult > 2:
	26	+ _split_accumulators(base_acc, 2)
	27	+ if mhc_mult > 4:
	28	+ _split_accumulators(base_acc, 4)
	29	+ if mhc_mult > 8:
	30	+ _split_accumulators(base_acc, 8)
	31	+ if mhc_mult > 16:
	32	+ _split_accumulators(base_acc, 16)
	33	+
	34	+
	35	+@tilelang.jit(pass_configs=_PASS_CONFIGS)
	36	+def mhc_head_compute_mix_forward_kernel(
	37	+ input_mix,
	38	+ mhc_scale,
	39	+ mhc_base,
	40	+ output_mix,
	41	+ mhc_pre_eps: float,
	42	+ token_block_size: int,
	43	+):
	44	+ num_tokens = T.dynamic('num_tokens')
	45	+ mhc_mult = T.const('mhc_mult')
	46	+
	47	+ input_mix: T.Tensor[(num_tokens, mhc_mult), T.float32]
	48	+ mhc_scale: T.Tensor[(1,), T.float32]
	49	+ mhc_base: T.Tensor[(mhc_mult,), T.float32]
	50	+ output_mix: T.Tensor[(num_tokens, mhc_mult), T.float32]
	51	+
	52	+ with T.Kernel(T.ceildiv(num_tokens, token_block_size)) as block_id:
	53	+ input_ub = T.alloc_shared((token_block_size, mhc_mult), T.float32)
	54	+ output_ub = T.alloc_shared((token_block_size, mhc_mult), T.float32)
	55	+ scale_ub = T.alloc_shared((8,), T.float32)
	56	+ base_ub = T.alloc_shared((T.ceildiv(mhc_mult, 8) * 8,), T.float32)
	57	+
	58	+ token_offset = block_id * token_block_size
	59	+ valid_tokens = T.min(token_block_size, num_tokens - token_offset)
	60	+
	61	+ T.copy(mhc_scale[0], scale_ub[0:1])
	62	+ T.copy(mhc_base[0:mhc_mult], base_ub[0:mhc_mult])
	63	+ T.copy(
	64	+ input_mix[token_offset : token_offset + valid_tokens, 0:mhc_mult],
	65	+ input_ub[0:valid_tokens, 0:mhc_mult],
	66	+ )
	67	+
	68	+ with T.SimdVF():
	69	+ zeros = S.vdup(T.float32(0.0), T.float32)
	70	+ ones = S.vdup(T.float32(1.0), T.float32)
	71	+ scale = S.vld(scale_ub[0], dist='BRC_B32')
	72	+ base_indices = S.vand(
	73	+ T.reinterpret(S.vci(T.int32(0), T.int32), 'uint32x64'),
	74	+ S.vdup(T.uint32(mhc_mult - 1), T.uint32),
	75	+ )
	76	+ base_pattern = S.vgather2(base_ub[0], base_indices)
	77	+ for vector_id in T.serial(token_block_size * mhc_mult // _VECTOR_WIDTH):
	78	+ row = vector_id * _VECTOR_WIDTH // mhc_mult
	79	+ x = S.vld(input_ub[row, 0])
	80	+ affine = S.vadd(S.vmul(x, scale), base_pattern)
	81	+ sigmoid = S.vdiv(ones, S.vadds(S.vexpdif(zeros, affine), 1.0))
	82	+ S.vsts(output_ub[row, 0], S.vadds(sigmoid, mhc_pre_eps))
	83	+
	84	+ T.copy(
	85	+ output_ub[0:valid_tokens, 0:mhc_mult],
	86	+ output_mix[token_offset : token_offset + valid_tokens, 0:mhc_mult],
	87	+ )
	88	+
	89	+
	90	+@tilelang.jit(pass_configs=_PASS_CONFIGS)
	91	+def mhc_head_compute_mix_backward_kernel(
	92	+ output_mix_grad,
	93	+ input_mix,
	94	+ mhc_scale,
	95	+ mhc_base,
	96	+ input_mix_grad,
	97	+ mhc_scale_grad_partial,
	98	+ mhc_base_grad_partial,
	99	+ token_block_size: int,
	100	+ num_vec_cores: int,
	101	+):
	102	+ num_tokens = T.dynamic('num_tokens')
	103	+ mhc_mult = T.const('mhc_mult')
	104	+
	105	+ output_mix_grad: T.Tensor[(num_tokens, mhc_mult), T.float32]
	106	+ input_mix: T.Tensor[(num_tokens, mhc_mult), T.float32]
	107	+ mhc_scale: T.Tensor[(1,), T.float32]
	108	+ mhc_base: T.Tensor[(mhc_mult,), T.float32]
	109	+ input_mix_grad: T.Tensor[(num_tokens, mhc_mult), T.float32]
	110	+ mhc_scale_grad_partial: T.Tensor[(num_vec_cores, 1), T.float32]
	111	+ mhc_base_grad_partial: T.Tensor[(num_vec_cores, mhc_mult), T.float32]
	112	+
	113	+ num_tiles = T.ceildiv(num_tokens, token_block_size)
	114	+
	115	+ with T.Kernel(num_vec_cores) as block_id:
	116	+ output_grad_ub = T.alloc_shared((token_block_size, mhc_mult), T.float32)
	117	+ input_ub = T.alloc_shared((token_block_size, mhc_mult), T.float32)
	118	+ scale_ub = T.alloc_shared((8,), T.float32)
	119	+ base_ub = T.alloc_shared((T.ceildiv(mhc_mult, 8) * 8,), T.float32)
	120	+ input_grad_ub = T.alloc_shared((token_block_size, mhc_mult), T.float32)
	121	+ scale_grad_partial_ub = T.alloc_shared((8,), T.float32)
	122	+ base_grad_partial_ub = T.alloc_shared((T.ceildiv(mhc_mult, 8) * 8,), T.float32)
	123	+ scale_acc_ub = T.alloc_shared((_VECTOR_WIDTH,), T.float32)
	124	+ seed_acc_ub = T.alloc_shared((2, _VECTOR_WIDTH), T.float32)
	125	+
	126	+ T.copy(mhc_base[0:mhc_mult], base_ub[0:mhc_mult])
	127	+ T.copy(mhc_scale[0], scale_ub[0:1])
	128	+
	129	+ with T.SimtVF(threads=_VECTOR_WIDTH):
	130	+ T.clear(scale_acc_ub)
	131	+ T.clear(seed_acc_ub)
	132	+
	133	+ for tile_id in T.Persistent(
	134	+ [num_tiles],
	135	+ num_vec_cores,
	136	+ block_id,
	137	+ group_size=1,
	138	+ num_stages=2,
	139	+ annotations={'multi_buffer_eligible': [output_grad_ub, input_ub, input_grad_ub]},
	140	+ ):
	141	+ token_offset = tile_id * token_block_size
	142	+ valid_tokens = T.min(token_block_size, num_tokens - token_offset)
	143	+
	144	+ T.copy(
	145	+ output_mix_grad[token_offset : token_offset + valid_tokens, 0:mhc_mult],
	146	+ output_grad_ub[0:valid_tokens, 0:mhc_mult],
	147	+ )
	148	+ T.copy(
	149	+ input_mix[token_offset : token_offset + valid_tokens, 0:mhc_mult],
	150	+ input_ub[0:valid_tokens, 0:mhc_mult],
	151	+ )
	152	+
	153	+ with T.SimdVF():
	154	+ zeros = S.vdup(T.float32(0.0), T.float32)
	155	+ ones = S.vdup(T.float32(1.0), T.float32)
	156	+ scale = S.vld(scale_ub[0], dist='BRC_B32')
	157	+ base_indices = S.vand(
	158	+ T.reinterpret(S.vci(T.int32(0), T.int32), 'uint32x64'),
	159	+ S.vdup(T.uint32(mhc_mult - 1), T.uint32),
	160	+ )
	161	+ base_pattern = S.vgather2(base_ub[0], base_indices)
	162	+
	163	+ S.mem_bar('VST_VLD')
	164	+ scale_acc = S.alloc_var(T.float32)
	165	+ scale_acc = S.vld(scale_acc_ub[0])
	166	+ seed_acc = S.alloc_local((2,), T.float32)
	167	+ seed_acc[0] = S.vld(seed_acc_ub[0, 0])
	168	+ seed_acc[1] = S.vld(seed_acc_ub[1, 0])
	169	+
	170	+ for vector_id in T.serial(token_block_size * mhc_mult // _VECTOR_WIDTH):
	171	+ row = vector_id * _VECTOR_WIDTH // mhc_mult
	172	+ x = S.alloc_var(T.float32)
	173	+ x = S.vld(input_ub[row, 0])
	174	+ output_grad = S.alloc_var(T.float32)
	175	+ output_grad = S.vld(output_grad_ub[row, 0])
	176	+
	177	+ if valid_tokens < token_block_size:
	178	+ valid_mask = S.vcmp(
	179	+ S.vci(T.int32(row * mhc_mult), T.int32),
	180	+ S.vdup(T.int32((num_tokens % token_block_size) * mhc_mult), T.int32),
	181	+ op='lt',
	182	+ )
	183	+ x = S.vsel(x, zeros, valid_mask)
	184	+ output_grad = S.vsel(output_grad, zeros, valid_mask)
	185	+
	186	+ affine = S.alloc_var(T.float32)
	187	+ affine = x
	188	+ S.vmadd(affine, scale, base_pattern)
	189	+ sigmoid = S.vdiv(ones, S.vadds(S.vexpdif(zeros, affine), 1.0))
	190	+ grad = S.vmul(S.vmul(sigmoid, S.vsub(ones, sigmoid)), output_grad)
	191	+ S.vsts(input_grad_ub[row, 0], S.vmul(grad, scale))
	192	+
	193	+ S.vmula(scale_acc, grad, x)
	194	+ even, odd = S.vdintlv(grad, grad)
	195	+ seed_acc[0] = S.vadd(seed_acc[0], even)
	196	+ seed_acc[1] = S.vadd(seed_acc[1], odd)
	197	+
	198	+ S.vsts(scale_acc_ub[0], scale_acc)
	199	+ S.vsts(seed_acc_ub[0, 0], seed_acc[0])
	200	+ S.vsts(seed_acc_ub[1, 0], seed_acc[1])
	201	+
	202	+ T.copy(
	203	+ input_grad_ub[0:valid_tokens, 0:mhc_mult],
	204	+ input_mix_grad[token_offset : token_offset + valid_tokens, 0:mhc_mult],
	205	+ )
	206	+
	207	+ with T.SimdVF():
	208	+ S.mem_bar('VST_VLD')
	209	+ scale_acc = S.alloc_var(T.float32)
	210	+ scale_acc = S.vld(scale_acc_ub[0])
	211	+ S.vsts(scale_grad_partial_ub[0], S.vcadd(scale_acc), dist='ONEPT_B32')
	212	+
	213	+ seed_acc = S.alloc_local((2,), T.float32)
	214	+ seed_acc[0] = S.vld(seed_acc_ub[0, 0])
	215	+ seed_acc[1] = S.vld(seed_acc_ub[1, 0])
	216	+
	217	+ base_acc = S.alloc_local((mhc_mult,), T.float32)
	218	+ base_acc[0] = seed_acc[0]
	219	+ base_acc[1] = seed_acc[1]
	220	+ _split_accumulator_tree(base_acc, mhc_mult)
	221	+ component_lanes_mask = S.pset(32, f'PAT_VL{_VECTOR_WIDTH // mhc_mult}')
	222	+ for component in T.unroll(mhc_mult, explicit=True):
	223	+ S.vsts(
	224	+ base_grad_partial_ub[component],
	225	+ S.vcadd(base_acc[component], component_lanes_mask),
	226	+ dist='ONEPT_B32',
	227	+ )
	228	+
	229	+ T.copy(scale_grad_partial_ub[0:1], mhc_scale_grad_partial[block_id, 0:1])
	230	+ T.copy(base_grad_partial_ub[0:mhc_mult], mhc_base_grad_partial[block_id, 0:mhc_mult])

Asrc/cann_ops_tilelang/mhc/multilayer_recompute.py+256-0

@@ -0,0 +1,256 @@
	1	+"""Ascend MHC multilayer recompute operator."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+from collections.abc import Sequence
	6	+
	7	+import torch
	8	+
	9	+from cann_ops_tilelang.config import get_num_vector_cores, require_ascend
	10	+from cann_ops_tilelang.utils import align_up
	11	+
	12	+from .multilayer_recompute_kernel import mhc_multilayer_recompute_asc
	13	+
	14	+_HIDDEN_ALIGNMENT = 64
	15	+_MAX_HIDDEN_BLOCK = 8192
	16	+
	17	+
	18	+def _check_tensor(
	19	+ name: str,
	20	+ tensor: torch.Tensor,
	21	+ *,
	22	+ dtype: torch.dtype,
	23	+ shape: tuple[int, ...],
	24	+ device: torch.device,
	25	+) -> None:
	26	+ if not isinstance(tensor, torch.Tensor):
	27	+ raise TypeError(f'{name} must be a torch.Tensor.')
	28	+ if tensor.device.type != 'npu':
	29	+ raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')
	30	+ if tensor.device != device:
	31	+ raise ValueError(f'{name} must be on {device}, got {tensor.device}.')
	32	+ if tensor.dtype != dtype:
	33	+ raise TypeError(f'{name} must use {dtype}, got {tensor.dtype}.')
	34	+ if tensor.shape != shape:
	35	+ raise ValueError(f'{name} must have shape {shape}, got {tuple(tensor.shape)}.')
	36	+ if not tensor.is_contiguous():
	37	+ raise ValueError(f'{name} must be contiguous.')
	38	+
	39	+
	40	+def _check_tensor_list(
	41	+ name: str,
	42	+ tensors: Sequence[torch.Tensor],
	43	+ *,
	44	+ length: int,
	45	+ dtype: torch.dtype,
	46	+ shape: tuple[int, ...],
	47	+ device: torch.device,
	48	+) -> None:
	49	+ if not isinstance(tensors, list):
	50	+ raise TypeError(f'{name} must be a list of torch.Tensor instances.')
	51	+ if len(tensors) != length:
	52	+ raise ValueError(f'{name} must contain {length} tensors, got {len(tensors)}.')
	53	+ for index, tensor in enumerate(tensors):
	54	+ _check_tensor(
	55	+ f'{name}[{index}]',
	56	+ tensor,
	57	+ dtype=dtype,
	58	+ shape=shape,
	59	+ device=device,
	60	+ )
	61	+
	62	+
	63	+def _shares_storage(left: torch.Tensor, right: torch.Tensor) -> bool:
	64	+ return (
	65	+ left.numel() > 0
	66	+ and right.numel() > 0
	67	+ and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr()
	68	+ )
	69	+
	70	+
	71	+def _check_output_aliases(
	72	+ outputs: Sequence[tuple[str, torch.Tensor]],
	73	+ inputs: Sequence[tuple[str, torch.Tensor]],
	74	+) -> None:
	75	+ for output_index, (output_name, output) in enumerate(outputs):
	76	+ for input_name, input_tensor in inputs:
	77	+ if _shares_storage(output, input_tensor):
	78	+ raise ValueError(f'{output_name} must not share storage with {input_name}.')
	79	+ for other_name, other_output in outputs[:output_index]:
	80	+ if _shares_storage(output, other_output):
	81	+ raise ValueError(f'{output_name} must not share storage with {other_name}.')
	82	+
	83	+
	84	+def _make_pointer_tables(
	85	+ tensor_lists: Sequence[Sequence[torch.Tensor]],
	86	+ device: torch.device,
	87	+) -> list[torch.Tensor]:
	88	+ sizes = [len(tensors) for tensors in tensor_lists]
	89	+ pointer_buffer = torch.empty(sum(sizes), dtype=torch.int64, device='cpu', pin_memory=True)
	90	+ offset = 0
	91	+ for tensors in tensor_lists:
	92	+ for index, tensor in enumerate(tensors):
	93	+ pointer_buffer[offset + index] = tensor.data_ptr()
	94	+ offset += len(tensors)
	95	+
	96	+ device_buffer = pointer_buffer.to(device=device, non_blocking=True)
	97	+ pointer_tables: list[torch.Tensor] = []
	98	+ offset = 0
	99	+ for size in sizes:
	100	+ pointer_tables.append(device_buffer[offset : offset + size])
	101	+ offset += size
	102	+ return pointer_tables
	103	+
	104	+
	105	+def mhc_multilayer_recompute(
	106	+ initial_residual: torch.Tensor,
	107	+ pre_mix_list: list[torch.Tensor],
	108	+ layer_output_list: list[torch.Tensor],
	109	+ post_mix_list: list[torch.Tensor],
	110	+ comb_mix_list: list[torch.Tensor],
	111	+ layer_input_list: list[torch.Tensor],
	112	+ residual_list: list[torch.Tensor],
	113	+) -> None:
	114	+ """Recompute MHC layer inputs and residuals into caller-provided buffers."""
	115	+
	116	+ if not isinstance(initial_residual, torch.Tensor):
	117	+ raise TypeError('initial_residual must be a torch.Tensor.')
	118	+ if initial_residual.ndim < 3:
	119	+ raise ValueError(
	120	+ f'initial_residual must have shape (..., mhc_mult, hidden_size), got {tuple(initial_residual.shape)}.'
	121	+ )
	122	+
	123	+ require_ascend()
	124	+ if initial_residual.device.type != 'npu':
	125	+ raise ValueError(f'initial_residual must be an NPU tensor, got {initial_residual.device}.')
	126	+ if initial_residual.dtype != torch.bfloat16:
	127	+ raise TypeError(f'initial_residual must use {torch.bfloat16}, got {initial_residual.dtype}.')
	128	+ if not initial_residual.is_contiguous():
	129	+ raise ValueError('initial_residual must be contiguous.')
	130	+
	131	+ for name, tensors in (
	132	+ ('pre_mix_list', pre_mix_list),
	133	+ ('layer_output_list', layer_output_list),
	134	+ ('post_mix_list', post_mix_list),
	135	+ ('comb_mix_list', comb_mix_list),
	136	+ ('layer_input_list', layer_input_list),
	137	+ ('residual_list', residual_list),
	138	+ ):
	139	+ if not isinstance(tensors, list):
	140	+ raise TypeError(f'{name} must be a list of torch.Tensor instances.')
	141	+
	142	+ num_layers = len(pre_mix_list)
	143	+ num_post_layers = len(layer_output_list)
	144	+ if num_layers <= 0:
	145	+ raise ValueError(f'pre_mix_list must contain at least one tensor, got {num_layers}.')
	146	+ if num_post_layers not in (num_layers - 1, num_layers):
	147	+ raise ValueError(
	148	+ f'layer_output_list must contain {num_layers - 1} or {num_layers} tensors, got {num_post_layers}.'
	149	+ )
	150	+
	151	+ leading_shape = tuple(initial_residual.shape[:-2])
	152	+ mhc_mult, hidden_size = initial_residual.shape[-2:]
	153	+ if mhc_mult <= 0:
	154	+ raise ValueError(f'mhc_mult must be positive, got {mhc_mult}.')
	155	+ if hidden_size <= 0 or hidden_size % _HIDDEN_ALIGNMENT:
	156	+ raise ValueError(f'hidden_size must be a positive multiple of {_HIDDEN_ALIGNMENT}, got {hidden_size}.')
	157	+
	158	+ device = initial_residual.device
	159	+ _check_tensor_list(
	160	+ 'pre_mix_list',
	161	+ pre_mix_list,
	162	+ length=num_layers,
	163	+ dtype=torch.float32,
	164	+ shape=(*leading_shape, mhc_mult, 1),
	165	+ device=device,
	166	+ )
	167	+ _check_tensor_list(
	168	+ 'layer_output_list',
	169	+ layer_output_list,
	170	+ length=num_post_layers,
	171	+ dtype=torch.bfloat16,
	172	+ shape=(*leading_shape, hidden_size),
	173	+ device=device,
	174	+ )
	175	+ _check_tensor_list(
	176	+ 'post_mix_list',
	177	+ post_mix_list,
	178	+ length=num_post_layers,
	179	+ dtype=torch.float32,
	180	+ shape=(*leading_shape, mhc_mult, 1),
	181	+ device=device,
	182	+ )
	183	+ _check_tensor_list(
	184	+ 'comb_mix_list',
	185	+ comb_mix_list,
	186	+ length=num_post_layers,
	187	+ dtype=torch.float32,
	188	+ shape=(*leading_shape, mhc_mult, mhc_mult),
	189	+ device=device,
	190	+ )
	191	+ _check_tensor_list(
	192	+ 'layer_input_list',
	193	+ layer_input_list,
	194	+ length=num_layers,
	195	+ dtype=torch.bfloat16,
	196	+ shape=(*leading_shape, hidden_size),
	197	+ device=device,
	198	+ )
	199	+ _check_tensor_list(
	200	+ 'residual_list',
	201	+ residual_list,
	202	+ length=num_post_layers,
	203	+ dtype=torch.bfloat16,
	204	+ shape=(*leading_shape, mhc_mult, hidden_size),
	205	+ device=device,
	206	+ )
	207	+
	208	+ input_tensors = [('initial_residual', initial_residual)]
	209	+ for list_name, tensors in (
	210	+ ('pre_mix_list', pre_mix_list),
	211	+ ('layer_output_list', layer_output_list),
	212	+ ('post_mix_list', post_mix_list),
	213	+ ('comb_mix_list', comb_mix_list),
	214	+ ):
	215	+ input_tensors.extend((f'{list_name}[{index}]', tensor) for index, tensor in enumerate(tensors))
	216	+ output_tensors = [
	217	+ *((f'layer_input_list[{index}]', tensor) for index, tensor in enumerate(layer_input_list)),
	218	+ *((f'residual_list[{index}]', tensor) for index, tensor in enumerate(residual_list)),
	219	+ ]
	220	+ _check_output_aliases(output_tensors, input_tensors)
	221	+
	222	+ num_tokens = initial_residual.numel() // (mhc_mult * hidden_size)
	223	+ if num_tokens == 0:
	224	+ return
	225	+
	226	+ (
	227	+ pre_mix_ptrs,
	228	+ layer_input_ptrs,
	229	+ residual_ptrs,
	230	+ layer_output_ptrs,
	231	+ post_mix_ptrs,
	232	+ comb_mix_ptrs,
	233	+ ) = _make_pointer_tables(
	234	+ [
	235	+ pre_mix_list,
	236	+ layer_input_list,
	237	+ residual_list,
	238	+ layer_output_list,
	239	+ post_mix_list,
	240	+ comb_mix_list,
	241	+ ],
	242	+ device,
	243	+ )
	244	+
	245	+ hidden_block = min(align_up(hidden_size, _HIDDEN_ALIGNMENT), _MAX_HIDDEN_BLOCK)
	246	+ mhc_multilayer_recompute_asc(
	247	+ initial_residual.view(num_tokens, mhc_mult, hidden_size),
	248	+ pre_mix_ptrs,
	249	+ layer_output_ptrs,
	250	+ post_mix_ptrs,
	251	+ comb_mix_ptrs,
	252	+ layer_input_ptrs,
	253	+ residual_ptrs,
	254	+ h_blk=hidden_block,
	255	+ num_vec_cores=get_num_vector_cores(),
	256	+ )

Asrc/cann_ops_tilelang/mhc/multilayer_recompute_kernel.py+219-0

@@ -0,0 +1,219 @@
	1	+import tilelang
	2	+from tilelang.ascend import language as T
	3	+from tilelang.ascend.language import simd as S
	4	+
	5	+_VEC = 64
	6	+
	7	+_PASS_CONFIGS = {
	8	+ tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: False,
	9	+}
	10	+
	11	+
	12	+@tilelang.jit(pass_configs=_PASS_CONFIGS)
	13	+def mhc_multilayer_recompute_asc(
	14	+ initial_residual,
	15	+ pre_mix_ptrs,
	16	+ layer_output_ptrs,
	17	+ post_mix_ptrs,
	18	+ comb_mix_ptrs,
	19	+ layer_input_ptrs,
	20	+ residual_ptrs,
	21	+ h_blk: int,
	22	+ num_vec_cores: int,
	23	+):
	24	+ num_tokens = T.dynamic('n')
	25	+ mhc = T.const('mhc')
	26	+ hidden = T.const('h')
	27	+ num_layers = T.const('L')
	28	+ num_post_layers = T.const('L_post')
	29	+
	30	+ initial_residual: T.Tensor[(num_tokens, mhc, hidden), T.bfloat16]
	31	+ pre_mix_ptrs: T.Tensor[(num_layers,), T.ptr]
	32	+ layer_output_ptrs: T.Tensor[(num_post_layers,), T.ptr]
	33	+ post_mix_ptrs: T.Tensor[(num_post_layers,), T.ptr]
	34	+ comb_mix_ptrs: T.Tensor[(num_post_layers,), T.ptr]
	35	+ layer_input_ptrs: T.Tensor[(num_layers,), T.ptr]
	36	+ residual_ptrs: T.Tensor[(num_post_layers,), T.ptr]
	37	+
	38	+ num_hidden_tiles = T.ceildiv(hidden, h_blk)
	39	+ num_chunks = h_blk // _VEC
	40	+ mhc_pad = T.ceildiv(mhc, 8) * 8
	41	+
	42	+ with T.Kernel(num_vec_cores) as core_id:
	43	+ residual_ub = T.alloc_shared((2, mhc, h_blk), T.bfloat16)
	44	+ layer_output_ub = T.alloc_shared((2, h_blk), T.bfloat16)
	45	+ layer_input_ub = T.alloc_shared((2, h_blk), T.bfloat16)
	46	+ pre_mix_ub = T.alloc_shared((2, mhc_pad), T.float32)
	47	+ post_mix_ub = T.alloc_shared((2, mhc_pad), T.float32)
	48	+ comb_mix_ub = T.alloc_shared((2, mhc, mhc), T.float32)
	49	+
	50	+ T.ascend_set_flag('MTE3_MTE2', 0)
	51	+
	52	+ for pid_n, pid_h in T.Persistent(
	53	+ [num_tokens, num_hidden_tiles],
	54	+ num_vec_cores,
	55	+ core_id,
	56	+ group_size=1,
	57	+ num_stages=0,
	58	+ ):
	59	+ T.ascend_set_flag('V_MTE2', 0)
	60	+ T.ascend_set_flag('V_MTE2', 1)
	61	+ T.ascend_set_flag('MTE3_V', 0)
	62	+ T.ascend_set_flag('MTE3_V', 1)
	63	+
	64	+ T.ascend_wait_flag('MTE3_MTE2', 0)
	65	+ T.copy(initial_residual[pid_n, 0, pid_h * h_blk], residual_ub[0, :, :])
	66	+ T.ascend_set_flag('MTE2_V', 2)
	67	+ T.ascend_wait_flag('MTE2_V', 2)
	68	+
	69	+ for layer_index in T.serial(num_post_layers):
	70	+ input_slot = layer_index % 2
	71	+ output_slot = 1 - input_slot
	72	+ pre_mix_tensor = T.make_tensor(pre_mix_ptrs[layer_index], (num_tokens, mhc), T.float32)
	73	+ post_mix_tensor = T.make_tensor(post_mix_ptrs[layer_index], (num_tokens, mhc), T.float32)
	74	+ comb_mix_tensor = T.make_tensor(
	75	+ comb_mix_ptrs[layer_index],
	76	+ (num_tokens, mhc, mhc),
	77	+ T.float32,
	78	+ )
	79	+ layer_output_tensor = T.make_tensor(
	80	+ layer_output_ptrs[layer_index],
	81	+ (num_tokens, hidden),
	82	+ T.bfloat16,
	83	+ )
	84	+ layer_input_tensor = T.make_tensor(
	85	+ layer_input_ptrs[layer_index],
	86	+ (num_tokens, hidden),
	87	+ T.bfloat16,
	88	+ )
	89	+ residual_tensor = T.make_tensor(
	90	+ residual_ptrs[layer_index],
	91	+ (num_tokens, mhc, hidden),
	92	+ T.bfloat16,
	93	+ )
	94	+
	95	+ T.ascend_wait_flag('V_MTE2', input_slot)
	96	+ T.copy(pre_mix_tensor[pid_n, 0], pre_mix_ub[input_slot, 0:mhc])
	97	+ T.copy(post_mix_tensor[pid_n, 0], post_mix_ub[input_slot, 0:mhc])
	98	+ T.copy(comb_mix_tensor[pid_n, 0, 0], comb_mix_ub[input_slot, :, :])
	99	+ T.copy(layer_output_tensor[pid_n, pid_h * h_blk], layer_output_ub[input_slot, :])
	100	+ T.ascend_set_flag('MTE2_V', input_slot)
	101	+
	102	+ T.ascend_wait_flag('MTE2_V', input_slot)
	103	+ T.ascend_wait_flag('MTE3_V', output_slot)
	104	+ with T.SimdVF():
	105	+ pre_mix = S.alloc_local((mhc,), T.float32)
	106	+ post_mix = S.alloc_local((mhc,), T.float32)
	107	+ comb_mix = S.alloc_local((mhc * mhc,), T.float32)
	108	+ for mix_index in T.unroll(mhc, explicit=True):
	109	+ pre_mix[mix_index] = S.vld(pre_mix_ub[input_slot, mix_index], dist='BRC_B32')
	110	+ post_mix[mix_index] = S.vld(post_mix_ub[input_slot, mix_index], dist='BRC_B32')
	111	+ for input_index in T.unroll(mhc, explicit=True):
	112	+ for output_index in T.unroll(mhc, explicit=True):
	113	+ comb_mix[input_index * mhc + output_index] = S.vld(
	114	+ comb_mix_ub[input_slot, input_index, output_index],
	115	+ dist='BRC_B32',
	116	+ )
	117	+ residual = S.alloc_local((mhc,), T.float32)
	118	+ output = S.alloc_local((mhc,), T.float32)
	119	+ layer_input = S.alloc_var(T.float32)
	120	+ for chunk_index in T.serial(num_chunks):
	121	+ column_offset = chunk_index * _VEC
	122	+ for input_index in T.unroll(mhc, explicit=True):
	123	+ residual[input_index] = S.vcvt(
	124	+ S.vld(
	125	+ residual_ub[input_slot, input_index, column_offset],
	126	+ dist='UNPK_B16',
	127	+ ),
	128	+ T.float32,
	129	+ part=0,
	130	+ )
	131	+ layer_input = S.vmul(residual[0], pre_mix[0])
	132	+ for input_index in T.unroll(mhc - 1, explicit=True):
	133	+ S.vmula(layer_input, residual[input_index + 1], pre_mix[input_index + 1])
	134	+ S.vsts(
	135	+ layer_input_ub[input_slot, column_offset],
	136	+ S.vcvt(layer_input, T.bfloat16),
	137	+ dist='PK_B32',
	138	+ )
	139	+ layer_output = S.vcvt(
	140	+ S.vld(layer_output_ub[input_slot, column_offset], dist='UNPK_B16'),
	141	+ T.float32,
	142	+ part=0,
	143	+ )
	144	+ for output_index in T.unroll(mhc, explicit=True):
	145	+ output[output_index] = S.vmul(layer_output, post_mix[output_index])
	146	+ for input_index in T.unroll(mhc, explicit=True):
	147	+ for output_index in T.unroll(mhc, explicit=True):
	148	+ S.vmula(
	149	+ output[output_index],
	150	+ residual[input_index],
	151	+ comb_mix[input_index * mhc + output_index],
	152	+ )
	153	+ for output_index in T.unroll(mhc, explicit=True):
	154	+ S.vsts(
	155	+ residual_ub[output_slot, output_index, column_offset],
	156	+ S.vcvt(output[output_index], T.bfloat16),
	157	+ dist='PK_B32',
	158	+ )
	159	+ T.ascend_set_flag('V_MTE2', input_slot)
	160	+ T.ascend_set_flag('V_MTE3', input_slot)
	161	+
	162	+ T.ascend_wait_flag('V_MTE3', input_slot)
	163	+ T.copy(residual_ub[output_slot, :, :], residual_tensor[pid_n, 0, pid_h * h_blk])
	164	+ T.copy(layer_input_ub[input_slot, :], layer_input_tensor[pid_n, pid_h * h_blk])
	165	+ T.ascend_set_flag('MTE3_V', output_slot)
	166	+
	167	+ T.ascend_set_flag('MTE3_MTE2', 0)
	168	+
	169	+ if num_layers > num_post_layers:
	170	+ final_slot = num_post_layers % 2
	171	+ final_output_slot = 1 - final_slot
	172	+ pre_mix_tensor = T.make_tensor(pre_mix_ptrs[num_post_layers], (num_tokens, mhc), T.float32)
	173	+ layer_input_tensor = T.make_tensor(
	174	+ layer_input_ptrs[num_post_layers],
	175	+ (num_tokens, hidden),
	176	+ T.bfloat16,
	177	+ )
	178	+ T.ascend_wait_flag('V_MTE2', final_slot)
	179	+ T.copy(pre_mix_tensor[pid_n, 0], pre_mix_ub[final_slot, 0:mhc])
	180	+ T.ascend_set_flag('MTE2_V', final_slot)
	181	+ T.ascend_wait_flag('MTE2_V', final_slot)
	182	+ T.ascend_wait_flag('MTE3_V', final_output_slot)
	183	+ with T.SimdVF():
	184	+ pre_mix = S.alloc_local((mhc,), T.float32)
	185	+ for mix_index in T.unroll(mhc, explicit=True):
	186	+ pre_mix[mix_index] = S.vld(pre_mix_ub[final_slot, mix_index], dist='BRC_B32')
	187	+ residual = S.alloc_local((mhc,), T.float32)
	188	+ layer_input = S.alloc_var(T.float32)
	189	+ for chunk_index in T.serial(num_chunks):
	190	+ column_offset = chunk_index * _VEC
	191	+ for mix_index in T.unroll(mhc, explicit=True):
	192	+ residual[mix_index] = S.vcvt(
	193	+ S.vld(
	194	+ residual_ub[final_slot, mix_index, column_offset],
	195	+ dist='UNPK_B16',
	196	+ ),
	197	+ T.float32,
	198	+ part=0,
	199	+ )
	200	+ layer_input = S.vmul(residual[0], pre_mix[0])
	201	+ for mix_index in T.unroll(mhc - 1, explicit=True):
	202	+ S.vmula(layer_input, residual[mix_index + 1], pre_mix[mix_index + 1])
	203	+ S.vsts(
	204	+ layer_input_ub[final_slot, column_offset],
	205	+ S.vcvt(layer_input, T.bfloat16),
	206	+ dist='PK_B32',
	207	+ )
	208	+ T.ascend_set_flag('V_MTE2', final_slot)
	209	+ T.ascend_set_flag('V_MTE3', final_slot)
	210	+ T.ascend_wait_flag('V_MTE3', final_slot)
	211	+ T.copy(layer_input_ub[final_slot, :], layer_input_tensor[pid_n, pid_h * h_blk])
	212	+ T.ascend_set_flag('MTE3_V', final_output_slot)
	213	+
	214	+ T.ascend_wait_flag('V_MTE2', 0)
	215	+ T.ascend_wait_flag('V_MTE2', 1)
	216	+ T.ascend_wait_flag('MTE3_V', 0)
	217	+ T.ascend_wait_flag('MTE3_V', 1)
	218	+
	219	+ T.ascend_wait_flag('MTE3_MTE2', 0)

Asrc/cann_ops_tilelang/mhc/norm_fn.py+416-0

@@ -0,0 +1,416 @@
	1	+"""Ascend MHC norm, GEMM, and partial-reduction operators."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import math
	6	+from numbers import Integral, Real
	7	+
	8	+import torch
	9	+
	10	+from cann_ops_tilelang.config import get_num_ai_cores, require_ascend
	11	+
	12	+from .norm_fn_kernel import (
	13	+ mhc_fn_normw_merge_bwd_asc,
	14	+ mhc_fn_normw_merge_fwd_asc,
	15	+ mhc_gemm_with_sqrsum_bwd_asc,
	16	+ mhc_gemm_with_sqrsum_fwd_asc,
	17	+ mhc_reduce_partials_and_rmsnorm_bwd_asc,
	18	+ mhc_reduce_partials_and_rmsnorm_fwd_asc,
	19	+)
	20	+
	21	+_MHC_MULT3 = 24
	22	+
	23	+
	24	+def _check_tensor(
	25	+ name: str,
	26	+ tensor: torch.Tensor,
	27	+ dtype: torch.dtype,
	28	+ ndim: int,
	29	+ *,
	30	+ device: torch.device \| None = None,
	31	+) -> None:
	32	+ if not isinstance(tensor, torch.Tensor):
	33	+ raise TypeError(f'{name} must be a torch.Tensor.')
	34	+ if tensor.device.type != 'npu':
	35	+ raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')
	36	+ if device is not None and tensor.device != device:
	37	+ raise ValueError(f'{name} must be on {device}, got {tensor.device}.')
	38	+ if tensor.dtype != dtype:
	39	+ raise TypeError(f'{name} must use {dtype}, got {tensor.dtype}.')
	40	+ if tensor.ndim != ndim:
	41	+ raise ValueError(f'{name} must be {ndim}-dimensional, got shape {tuple(tensor.shape)}.')
	42	+ if not tensor.is_contiguous():
	43	+ raise ValueError(f'{name} must be contiguous.')
	44	+
	45	+
	46	+def _shares_storage(left: torch.Tensor, right: torch.Tensor) -> bool:
	47	+ return (
	48	+ left.numel() > 0
	49	+ and right.numel() > 0
	50	+ and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr()
	51	+ )
	52	+
	53	+
	54	+def _check_output_aliases(inputs: dict[str, torch.Tensor], outputs: dict[str, torch.Tensor]) -> None:
	55	+ output_items = list(outputs.items())
	56	+ for output_name, output in output_items:
	57	+ for input_name, input_tensor in inputs.items():
	58	+ if _shares_storage(output, input_tensor):
	59	+ raise ValueError(f'{output_name} must not share storage with {input_name}.')
	60	+ for index, (left_name, left) in enumerate(output_items):
	61	+ for right_name, right in output_items[index + 1 :]:
	62	+ if _shares_storage(left, right):
	63	+ raise ValueError(f'{left_name} and {right_name} must not share storage.')
	64	+
	65	+
	66	+def _check_positive_int(name: str, value: int) -> int:
	67	+ if isinstance(value, bool) or not isinstance(value, Integral):
	68	+ raise TypeError(f'{name} must be an integer.')
	69	+ converted = int(value)
	70	+ if converted <= 0:
	71	+ raise ValueError(f'{name} must be positive, got {converted}.')
	72	+ return converted
	73	+
	74	+
	75	+def _check_non_negative_real(name: str, value: float) -> float:
	76	+ if isinstance(value, bool) or not isinstance(value, Real):
	77	+ raise TypeError(f'{name} must be a real number.')
	78	+ converted = float(value)
	79	+ if not math.isfinite(converted) or converted < 0:
	80	+ raise ValueError(f'{name} must be finite and non-negative, got {value}.')
	81	+ return converted
	82	+
	83	+
	84	+def _check_shape(name: str, tensor: torch.Tensor, expected: tuple[int, ...]) -> None:
	85	+ if tensor.shape != expected:
	86	+ raise ValueError(f'{name} must have shape {expected}, got {tuple(tensor.shape)}.')
	87	+
	88	+
	89	+def _ascend_forward_config(mhc_hidden_size: int, mhc_mult3: int, num_tokens: int) -> dict[str, int]:
	90	+ block_num_tokens = 128
	91	+ block_hidden = 256
	92	+ l0b_size_bytes = 64 * 1024
	93	+ l0c_size_bytes = 256 * 1024
	94	+
	95	+ num_ai_cores = get_num_ai_cores()
	96	+ block_mhc_mult3 = (mhc_mult3 + 15) // 16 * 16
	97	+ num_hidden_blocks = mhc_hidden_size // block_hidden
	98	+ num_token_blocks = math.ceil(num_tokens / block_num_tokens)
	99	+ num_l0b_stages = l0b_size_bytes // (block_mhc_mult3 * block_hidden * 4)
	100	+ num_l0c_stages = l0c_size_bytes // (block_num_tokens * block_mhc_mult3 * 4)
	101	+
	102	+ num_k_splits_by_tokens = max(1, num_ai_cores * num_l0c_stages // num_token_blocks)
	103	+ num_k_splits_by_hidden = num_hidden_blocks & -num_hidden_blocks
	104	+ num_k_splits_by_cores = num_ai_cores & -num_ai_cores
	105	+ split_limit = min(
	106	+ num_k_splits_by_tokens,
	107	+ num_k_splits_by_hidden,
	108	+ num_k_splits_by_cores,
	109	+ )
	110	+ num_k_splits = 1 << (split_limit.bit_length() - 1)
	111	+ active_ai_cores = num_k_splits * min(num_ai_cores // num_k_splits, num_token_blocks)
	112	+ cores_per_split = active_ai_cores // num_k_splits
	113	+ num_hidden_blocks_per_split = num_hidden_blocks // num_k_splits
	114	+ num_token_blocks_per_chunk = cores_per_split * num_l0c_stages
	115	+
	116	+ return {
	117	+ 'block_mhc_mult3': block_mhc_mult3,
	118	+ 'num_aic_cores': active_ai_cores,
	119	+ 'num_k_splits': num_k_splits,
	120	+ 'cores_per_split': cores_per_split,
	121	+ 'num_hidden_blocks_per_split': num_hidden_blocks_per_split,
	122	+ 'num_l0b_stages': num_l0b_stages,
	123	+ 'num_l0c_stages': num_l0c_stages,
	124	+ 'num_token_blocks_per_chunk': num_token_blocks_per_chunk,
	125	+ }
	126	+
	127	+
	128	+def mhc_fn_normw_merge_forward(
	129	+ fn: torch.Tensor,
	130	+ normw: torch.Tensor,
	131	+ out_fn: torch.Tensor,
	132	+) -> None:
	133	+ """Multiply an FP32 MHC function matrix by its FP32 normalization weights."""
	134	+
	135	+ if not all(isinstance(tensor, torch.Tensor) for tensor in (fn, normw, out_fn)):
	136	+ raise TypeError('fn, normw and out_fn must be torch.Tensor instances.')
	137	+ require_ascend()
	138	+ _check_tensor('fn', fn, torch.float32, 2)
	139	+ _check_tensor('normw', normw, torch.float32, 1, device=fn.device)
	140	+ _check_tensor('out_fn', out_fn, torch.float32, 2, device=fn.device)
	141	+ if fn.shape[0] != _MHC_MULT3 or fn.shape[1] <= 0 or fn.shape[1] % 256:
	142	+ raise ValueError(f'fn must have shape (24, hidden) with hidden divisible by 256, got {tuple(fn.shape)}.')
	143	+ _check_shape('normw', normw, (fn.shape[1],))
	144	+ _check_shape('out_fn', out_fn, tuple(fn.shape))
	145	+ _check_output_aliases({'fn': fn, 'normw': normw}, {'out_fn': out_fn})
	146	+
	147	+ mhc_fn_normw_merge_fwd_asc(
	148	+ fn,
	149	+ normw,
	150	+ out_fn,
	151	+ num_stages=2,
	152	+ )
	153	+
	154	+
	155	+def mhc_fn_normw_merge_backward(
	156	+ fn: torch.Tensor,
	157	+ normw: torch.Tensor,
	158	+ out_fn_grad: torch.Tensor,
	159	+ fn_grad: torch.Tensor,
	160	+ normw_grad: torch.Tensor,
	161	+) -> None:
	162	+ """Accumulate gradients into the FP32 function and normalization buffers."""
	163	+
	164	+ tensors = (fn, normw, out_fn_grad, fn_grad, normw_grad)
	165	+ if not all(isinstance(tensor, torch.Tensor) for tensor in tensors):
	166	+ raise TypeError('all norm-weight merge inputs and outputs must be torch.Tensor instances.')
	167	+ require_ascend()
	168	+ _check_tensor('fn', fn, torch.float32, 2)
	169	+ device = fn.device
	170	+ _check_tensor('normw', normw, torch.float32, 1, device=device)
	171	+ _check_tensor('out_fn_grad', out_fn_grad, torch.float32, 2, device=device)
	172	+ _check_tensor('fn_grad', fn_grad, torch.float32, 2, device=device)
	173	+ _check_tensor('normw_grad', normw_grad, torch.float32, 1, device=device)
	174	+ if fn.shape[0] != _MHC_MULT3 or fn.shape[1] <= 0 or fn.shape[1] % 128:
	175	+ raise ValueError(f'fn must have shape (24, hidden) with hidden divisible by 128, got {tuple(fn.shape)}.')
	176	+ _check_shape('normw', normw, (fn.shape[1],))
	177	+ _check_shape('out_fn_grad', out_fn_grad, tuple(fn.shape))
	178	+ _check_shape('fn_grad', fn_grad, tuple(fn.shape))
	179	+ _check_shape('normw_grad', normw_grad, (fn.shape[1],))
	180	+ _check_output_aliases(
	181	+ {'fn': fn, 'normw': normw, 'out_fn_grad': out_fn_grad},
	182	+ {'fn_grad': fn_grad, 'normw_grad': normw_grad},
	183	+ )
	184	+
	185	+ mhc_fn_normw_merge_bwd_asc(
	186	+ fn,
	187	+ normw,
	188	+ out_fn_grad,
	189	+ fn_grad,
	190	+ normw_grad,
	191	+ num_stages=2,
	192	+ )
	193	+
	194	+
	195	+def mhc_gemm_with_sqrsum_forward(
	196	+ x: torch.Tensor,
	197	+ fn: torch.Tensor,
	198	+ out_mul_splitted: torch.Tensor,
	199	+ sqrsum_splitted: torch.Tensor,
	200	+ n_splits: int = 16,
	201	+) -> None:
	202	+ """Compute the MHC GEMM and squared sums into one internally reduced partial."""
	203	+
	204	+ tensors = (x, fn, out_mul_splitted, sqrsum_splitted)
	205	+ if not all(isinstance(tensor, torch.Tensor) for tensor in tensors):
	206	+ raise TypeError('all GEMM forward inputs and outputs must be torch.Tensor instances.')
	207	+ require_ascend()
	208	+ _check_tensor('x', x, torch.bfloat16, 2)
	209	+ device = x.device
	210	+ _check_tensor('fn', fn, torch.float32, 2, device=device)
	211	+ _check_tensor('out_mul_splitted', out_mul_splitted, torch.float32, 4, device=device)
	212	+ _check_tensor('sqrsum_splitted', sqrsum_splitted, torch.float32, 3, device=device)
	213	+ _check_positive_int('n_splits', n_splits)
	214	+
	215	+ num_tokens, mhc_hidden_size = x.shape
	216	+ if fn.shape != (_MHC_MULT3, mhc_hidden_size):
	217	+ raise ValueError(f'fn must have shape {(24, mhc_hidden_size)}, got {tuple(fn.shape)}.')
	218	+ if mhc_hidden_size <= 0 or mhc_hidden_size % 256:
	219	+ raise ValueError(f'x hidden size must be a positive multiple of 256, got {mhc_hidden_size}.')
	220	+ _check_shape('out_mul_splitted', out_mul_splitted, (1, num_tokens, 1, _MHC_MULT3))
	221	+ _check_shape('sqrsum_splitted', sqrsum_splitted, (1, num_tokens, 1))
	222	+ _check_output_aliases(
	223	+ {'x': x, 'fn': fn},
	224	+ {'out_mul_splitted': out_mul_splitted, 'sqrsum_splitted': sqrsum_splitted},
	225	+ )
	226	+ if num_tokens == 0:
	227	+ return
	228	+
	229	+ config = _ascend_forward_config(mhc_hidden_size, fn.shape[0], num_tokens)
	230	+ mhc_gemm_with_sqrsum_fwd_asc(
	231	+ x,
	232	+ fn,
	233	+ out_mul_splitted[0, :, 0, :],
	234	+ sqrsum_splitted[0, :, 0],
	235	+ block_mhc_mult3=config['block_mhc_mult3'],
	236	+ num_aic_cores=config['num_aic_cores'],
	237	+ num_k_splits=config['num_k_splits'],
	238	+ cores_per_split=config['cores_per_split'],
	239	+ num_hidden_blocks_per_split=config['num_hidden_blocks_per_split'],
	240	+ num_l0b_stages=config['num_l0b_stages'],
	241	+ num_l0c_stages=config['num_l0c_stages'],
	242	+ num_token_blocks_per_chunk=config['num_token_blocks_per_chunk'],
	243	+ deterministic=torch.are_deterministic_algorithms_enabled(),
	244	+ )
	245	+
	246	+
	247	+def mhc_gemm_with_sqrsum_backward(
	248	+ out_mul_grad: torch.Tensor,
	249	+ sqrsum_grad: torch.Tensor,
	250	+ x: torch.Tensor,
	251	+ fn: torch.Tensor,
	252	+ x_grad: torch.Tensor,
	253	+ fn_grad: torch.Tensor,
	254	+) -> None:
	255	+ """Accumulate the BF16 input gradient and write the FP32 function gradient."""
	256	+
	257	+ tensors = (out_mul_grad, sqrsum_grad, x, fn, x_grad, fn_grad)
	258	+ if not all(isinstance(tensor, torch.Tensor) for tensor in tensors):
	259	+ raise TypeError('all GEMM backward inputs and outputs must be torch.Tensor instances.')
	260	+ require_ascend()
	261	+ _check_tensor('x', x, torch.bfloat16, 2)
	262	+ device = x.device
	263	+ _check_tensor('out_mul_grad', out_mul_grad, torch.float32, 3, device=device)
	264	+ _check_tensor('sqrsum_grad', sqrsum_grad, torch.float32, 2, device=device)
	265	+ _check_tensor('fn', fn, torch.float32, 2, device=device)
	266	+ _check_tensor('x_grad', x_grad, torch.bfloat16, 2, device=device)
	267	+ _check_tensor('fn_grad', fn_grad, torch.float32, 2, device=device)
	268	+
	269	+ num_tokens, mhc_hidden_size = x.shape
	270	+ if fn.shape != (_MHC_MULT3, mhc_hidden_size):
	271	+ raise ValueError(f'fn must have shape {(24, mhc_hidden_size)}, got {tuple(fn.shape)}.')
	272	+ if mhc_hidden_size <= 0 or mhc_hidden_size % 128:
	273	+ raise ValueError(f'x hidden size must be a positive multiple of 128, got {mhc_hidden_size}.')
	274	+ _check_shape('out_mul_grad', out_mul_grad, (num_tokens, 1, _MHC_MULT3))
	275	+ _check_shape('sqrsum_grad', sqrsum_grad, (num_tokens, 1))
	276	+ _check_shape('x_grad', x_grad, tuple(x.shape))
	277	+ _check_shape('fn_grad', fn_grad, tuple(fn.shape))
	278	+ _check_output_aliases(
	279	+ {'out_mul_grad': out_mul_grad, 'sqrsum_grad': sqrsum_grad, 'x': x, 'fn': fn},
	280	+ {'x_grad': x_grad, 'fn_grad': fn_grad},
	281	+ )
	282	+ if num_tokens == 0:
	283	+ fn_grad.zero_()
	284	+ return
	285	+
	286	+ kernel_out_mul_grad = out_mul_grad
	287	+ kernel_sqrsum_grad = sqrsum_grad
	288	+ kernel_x = x
	289	+ kernel_x_grad = x_grad
	290	+ if num_tokens % 128:
	291	+ padded_tokens = (num_tokens + 127) // 128 * 128
	292	+ kernel_out_mul_grad = out_mul_grad.new_zeros((padded_tokens, 1, _MHC_MULT3))
	293	+ kernel_sqrsum_grad = sqrsum_grad.new_zeros((padded_tokens, 1))
	294	+ kernel_x = x.new_zeros((padded_tokens, mhc_hidden_size))
	295	+ kernel_x_grad = x_grad.new_zeros((padded_tokens, mhc_hidden_size))
	296	+ kernel_out_mul_grad[:num_tokens].copy_(out_mul_grad)
	297	+ kernel_sqrsum_grad[:num_tokens].copy_(sqrsum_grad)
	298	+ kernel_x[:num_tokens].copy_(x)
	299	+ kernel_x_grad[:num_tokens].copy_(x_grad)
	300	+
	301	+ mhc_gemm_with_sqrsum_bwd_asc(
	302	+ kernel_out_mul_grad,
	303	+ kernel_sqrsum_grad,
	304	+ kernel_x,
	305	+ fn,
	306	+ kernel_x_grad,
	307	+ fn_grad,
	308	+ token_block=128,
	309	+ hidden_block=128,
	310	+ )
	311	+ if kernel_x_grad is not x_grad:
	312	+ x_grad.copy_(kernel_x_grad[:num_tokens])
	313	+
	314	+
	315	+def mhc_reduce_partials_and_rmsnorm_forward(
	316	+ out_mul_splitted: torch.Tensor,
	317	+ sqrsum_splitted: torch.Tensor,
	318	+ out_mul: torch.Tensor,
	319	+ sqrsum: torch.Tensor,
	320	+ out: torch.Tensor,
	321	+ rms_group_size: int,
	322	+ rms_eps: float,
	323	+ n_splits: int,
	324	+) -> None:
	325	+ """Reduce split partials and apply RMS normalization into caller buffers."""
	326	+
	327	+ tensors = (out_mul_splitted, sqrsum_splitted, out_mul, sqrsum, out)
	328	+ if not all(isinstance(tensor, torch.Tensor) for tensor in tensors):
	329	+ raise TypeError('all reduce/RMSNorm forward inputs and outputs must be torch.Tensor instances.')
	330	+ require_ascend()
	331	+ _check_tensor('out_mul_splitted', out_mul_splitted, torch.float32, 4)
	332	+ device = out_mul_splitted.device
	333	+ _check_tensor('sqrsum_splitted', sqrsum_splitted, torch.float32, 3, device=device)
	334	+ _check_tensor('out_mul', out_mul, torch.float32, 3, device=device)
	335	+ _check_tensor('sqrsum', sqrsum, torch.float32, 2, device=device)
	336	+ _check_tensor('out', out, torch.float32, 2, device=device)
	337	+ n_splits = _check_positive_int('n_splits', n_splits)
	338	+ rms_group_size = _check_positive_int('rms_group_size', rms_group_size)
	339	+ rms_eps = _check_non_negative_real('rms_eps', rms_eps)
	340	+
	341	+ if out_mul_splitted.shape[0] != n_splits:
	342	+ raise ValueError('out_mul_splitted first dimension must equal n_splits.')
	343	+ _, num_tokens, num_rms_groups, mhc_mult3 = out_mul_splitted.shape
	344	+ if num_rms_groups != 1 or mhc_mult3 != _MHC_MULT3:
	345	+ raise ValueError(
	346	+ f'out_mul_splitted must have shape (n_splits, num_tokens, 1, 24), got {tuple(out_mul_splitted.shape)}.'
	347	+ )
	348	+ _check_shape('sqrsum_splitted', sqrsum_splitted, (n_splits, num_tokens, 1))
	349	+ _check_shape('out_mul', out_mul, (num_tokens, 1, _MHC_MULT3))
	350	+ _check_shape('sqrsum', sqrsum, (num_tokens, 1))
	351	+ _check_shape('out', out, (num_tokens, _MHC_MULT3))
	352	+ _check_output_aliases(
	353	+ {'out_mul_splitted': out_mul_splitted, 'sqrsum_splitted': sqrsum_splitted},
	354	+ {'out_mul': out_mul, 'sqrsum': sqrsum, 'out': out},
	355	+ )
	356	+ if num_tokens == 0:
	357	+ return
	358	+
	359	+ mhc_reduce_partials_and_rmsnorm_fwd_asc(
	360	+ out_mul_splitted,
	361	+ sqrsum_splitted,
	362	+ out_mul,
	363	+ sqrsum,
	364	+ out,
	365	+ rms_group_size,
	366	+ rms_eps,
	367	+ n_splits,
	368	+ )
	369	+
	370	+
	371	+def mhc_reduce_partials_and_rmsnorm_backward(
	372	+ out_grad: torch.Tensor,
	373	+ out_mul: torch.Tensor,
	374	+ sqrsum: torch.Tensor,
	375	+ out_mul_grad: torch.Tensor,
	376	+ sqrsum_grad: torch.Tensor,
	377	+ rms_group_size: int,
	378	+ rms_eps: float,
	379	+) -> None:
	380	+ """Backpropagate through the MHC partial reduction and RMS normalization."""
	381	+
	382	+ tensors = (out_grad, out_mul, sqrsum, out_mul_grad, sqrsum_grad)
	383	+ if not all(isinstance(tensor, torch.Tensor) for tensor in tensors):
	384	+ raise TypeError('all reduce/RMSNorm backward inputs and outputs must be torch.Tensor instances.')
	385	+ require_ascend()
	386	+ _check_tensor('out_grad', out_grad, torch.float32, 2)
	387	+ device = out_grad.device
	388	+ _check_tensor('out_mul', out_mul, torch.float32, 3, device=device)
	389	+ _check_tensor('sqrsum', sqrsum, torch.float32, 2, device=device)
	390	+ _check_tensor('out_mul_grad', out_mul_grad, torch.float32, 3, device=device)
	391	+ _check_tensor('sqrsum_grad', sqrsum_grad, torch.float32, 2, device=device)
	392	+ rms_group_size = _check_positive_int('rms_group_size', rms_group_size)
	393	+ rms_eps = _check_non_negative_real('rms_eps', rms_eps)
	394	+
	395	+ num_tokens = out_grad.shape[0]
	396	+ _check_shape('out_grad', out_grad, (num_tokens, _MHC_MULT3))
	397	+ _check_shape('out_mul', out_mul, (num_tokens, 1, _MHC_MULT3))
	398	+ _check_shape('sqrsum', sqrsum, (num_tokens, 1))
	399	+ _check_shape('out_mul_grad', out_mul_grad, (num_tokens, 1, _MHC_MULT3))
	400	+ _check_shape('sqrsum_grad', sqrsum_grad, (num_tokens, 1))
	401	+ _check_output_aliases(
	402	+ {'out_grad': out_grad, 'out_mul': out_mul, 'sqrsum': sqrsum},
	403	+ {'out_mul_grad': out_mul_grad, 'sqrsum_grad': sqrsum_grad},
	404	+ )
	405	+ if num_tokens == 0:
	406	+ return
	407	+
	408	+ mhc_reduce_partials_and_rmsnorm_bwd_asc(
	409	+ out_grad,
	410	+ out_mul,
	411	+ sqrsum,
	412	+ out_mul_grad,
	413	+ sqrsum_grad,
	414	+ rms_group_size,
	415	+ rms_eps,
	416	+ )

Asrc/cann_ops_tilelang/mhc/norm_fn_kernel.py+1310-0

Asrc/cann_ops_tilelang/mhc/post.py+187-0

@@ -0,0 +1,187 @@
	1	+"""Public interfaces for MHC post forward and backward."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import torch
	6	+
	7	+from cann_ops_tilelang.config import get_num_vector_cores, require_ascend, should_print_kernel_source
	8	+from cann_ops_tilelang.utils import align_up, ceil_div
	9	+
	10	+from .post_kernel import get_mhc_post_backward_kernel, get_mhc_post_forward_kernel
	11	+
	12	+_HIDDEN_ALIGNMENT = 64
	13	+
	14	+
	15	+def _check_npu_tensor(
	16	+ name: str,
	17	+ tensor: torch.Tensor,
	18	+ *,
	19	+ dtype: torch.dtype,
	20	+ device: torch.device \| None = None,
	21	+) -> None:
	22	+ if not isinstance(tensor, torch.Tensor):
	23	+ raise TypeError(f'{name} must be a torch.Tensor.')
	24	+ if tensor.device.type != 'npu':
	25	+ raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')
	26	+ if device is not None and tensor.device != device:
	27	+ raise ValueError(f'{name} must be on {device}, got {tensor.device}.')
	28	+ if tensor.dtype != dtype:
	29	+ raise TypeError(f'{name} must use {dtype}, got {tensor.dtype}.')
	30	+ if not tensor.is_contiguous():
	31	+ raise ValueError(f'{name} must be contiguous.')
	32	+
	33	+
	34	+def _shares_storage(left: torch.Tensor, right: torch.Tensor) -> bool:
	35	+ return (
	36	+ left.numel() > 0
	37	+ and right.numel() > 0
	38	+ and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr()
	39	+ )
	40	+
	41	+
	42	+def _validate_common_inputs(
	43	+ x: torch.Tensor,
	44	+ residual: torch.Tensor,
	45	+ post_layer_mix: torch.Tensor,
	46	+ comb_res_mix: torch.Tensor,
	47	+) -> tuple[int, int, int, int]:
	48	+ tensors = (x, residual, post_layer_mix, comb_res_mix)
	49	+ if not all(isinstance(tensor, torch.Tensor) for tensor in tensors):
	50	+ raise TypeError('x, residual, post_layer_mix and comb_res_mix must be torch.Tensor instances.')
	51	+
	52	+ require_ascend()
	53	+ _check_npu_tensor('x', x, dtype=torch.bfloat16)
	54	+ _check_npu_tensor('residual', residual, dtype=torch.bfloat16, device=x.device)
	55	+ _check_npu_tensor('post_layer_mix', post_layer_mix, dtype=torch.float32, device=x.device)
	56	+ _check_npu_tensor('comb_res_mix', comb_res_mix, dtype=torch.float32, device=x.device)
	57	+
	58	+ if residual.ndim != 4:
	59	+ raise ValueError('residual must have shape (num_seqs, num_tokens, mhc_mult, hidden_size).')
	60	+ num_seqs, num_tokens, mhc_mult, hidden_size = residual.shape
	61	+ if mhc_mult <= 0:
	62	+ raise ValueError(f'mhc_mult must be positive, got {mhc_mult}.')
	63	+ if hidden_size <= 0 or hidden_size % _HIDDEN_ALIGNMENT:
	64	+ raise ValueError(f'hidden_size must be a positive multiple of {_HIDDEN_ALIGNMENT}, got {hidden_size}.')
	65	+
	66	+ expected_x_shape = (num_seqs, num_tokens, hidden_size)
	67	+ if x.shape != expected_x_shape:
	68	+ raise ValueError(f'x must have shape {expected_x_shape}, got {tuple(x.shape)}.')
	69	+ expected_post_shape = (num_seqs, num_tokens, mhc_mult, 1)
	70	+ if post_layer_mix.shape != expected_post_shape:
	71	+ raise ValueError(f'post_layer_mix must have shape {expected_post_shape}, got {tuple(post_layer_mix.shape)}.')
	72	+ expected_comb_shape = (num_seqs, num_tokens, mhc_mult, mhc_mult)
	73	+ if comb_res_mix.shape != expected_comb_shape:
	74	+ raise ValueError(f'comb_res_mix must have shape {expected_comb_shape}, got {tuple(comb_res_mix.shape)}.')
	75	+ return num_seqs, num_tokens, mhc_mult, hidden_size
	76	+
	77	+
	78	+def _choose_hidden_block(hidden_size: int) -> int:
	79	+ hidden_block = hidden_size if hidden_size <= 4096 else min(ceil_div(hidden_size, 2), 4096)
	80	+ return align_up(hidden_block, _HIDDEN_ALIGNMENT)
	81	+
	82	+
	83	+def mhc_post_forward(
	84	+ x: torch.Tensor,
	85	+ residual: torch.Tensor,
	86	+ post_layer_mix: torch.Tensor,
	87	+ comb_res_mix: torch.Tensor,
	88	+ out: torch.Tensor \| None = None,
	89	+) -> torch.Tensor:
	90	+ """Apply the MHC post residual mix on Ascend.
	91	+
	92	+ ``x`` has shape ``(num_seqs, num_tokens, hidden_size)`` and ``residual``
	93	+ has shape ``(num_seqs, num_tokens, mhc_mult, hidden_size)``. Activations
	94	+ use BF16, while both mix tensors use FP32.
	95	+ """
	96	+ num_seqs, num_tokens, mhc_mult, hidden_size = _validate_common_inputs(
	97	+ x,
	98	+ residual,
	99	+ post_layer_mix,
	100	+ comb_res_mix,
	101	+ )
	102	+ if out is not None:
	103	+ _check_npu_tensor('out', out, dtype=torch.bfloat16, device=x.device)
	104	+ if out.shape != residual.shape:
	105	+ raise ValueError(f'out must have shape {tuple(residual.shape)}, got {tuple(out.shape)}.')
	106	+ for name, tensor in (
	107	+ ('x', x),
	108	+ ('residual', residual),
	109	+ ('post_layer_mix', post_layer_mix),
	110	+ ('comb_res_mix', comb_res_mix),
	111	+ ):
	112	+ if _shares_storage(out, tensor):
	113	+ raise ValueError(f'out and {name} must not share storage.')
	114	+
	115	+ result = torch.empty_like(residual) if out is None else out
	116	+ if num_seqs == 0 or num_tokens == 0:
	117	+ return result
	118	+
	119	+ hidden_block = _choose_hidden_block(hidden_size)
	120	+ kernel = get_mhc_post_forward_kernel(
	121	+ mhc_mult,
	122	+ hidden_size,
	123	+ hidden_block,
	124	+ get_num_vector_cores(),
	125	+ )
	126	+ if should_print_kernel_source():
	127	+ print(kernel.get_kernel_source())
	128	+ flattened_tokens = num_seqs * num_tokens
	129	+ kernel(
	130	+ comb_res_mix.view(flattened_tokens, mhc_mult, mhc_mult),
	131	+ residual.view(flattened_tokens, mhc_mult, hidden_size),
	132	+ post_layer_mix.view(flattened_tokens, mhc_mult),
	133	+ x.view(flattened_tokens, hidden_size),
	134	+ result.view(flattened_tokens, mhc_mult, hidden_size),
	135	+ )
	136	+ return result
	137	+
	138	+
	139	+def mhc_post_backward(
	140	+ x: torch.Tensor,
	141	+ residual: torch.Tensor,
	142	+ post_layer_mix: torch.Tensor,
	143	+ comb_res_mix: torch.Tensor,
	144	+ d_o: torch.Tensor,
	145	+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
	146	+ """Return gradients for the four inputs of :func:`mhc_post_forward`."""
	147	+ num_seqs, num_tokens, mhc_mult, hidden_size = _validate_common_inputs(
	148	+ x,
	149	+ residual,
	150	+ post_layer_mix,
	151	+ comb_res_mix,
	152	+ )
	153	+ _check_npu_tensor('d_o', d_o, dtype=torch.bfloat16, device=x.device)
	154	+ if d_o.shape != residual.shape:
	155	+ raise ValueError(f'd_o must have shape {tuple(residual.shape)}, got {tuple(d_o.shape)}.')
	156	+
	157	+ grad_x = torch.empty_like(x)
	158	+ grad_residual = torch.empty_like(residual)
	159	+ grad_post_layer_mix = torch.empty_like(post_layer_mix)
	160	+ grad_comb_res_mix = torch.empty_like(comb_res_mix)
	161	+ if num_seqs == 0 or num_tokens == 0:
	162	+ return grad_x, grad_residual, grad_post_layer_mix, grad_comb_res_mix
	163	+
	164	+ hidden_block = _choose_hidden_block(hidden_size)
	165	+ num_stages = 3 if hidden_block <= 2560 else 2
	166	+ kernel = get_mhc_post_backward_kernel(
	167	+ mhc_mult,
	168	+ hidden_size,
	169	+ hidden_block,
	170	+ num_stages,
	171	+ get_num_vector_cores(),
	172	+ )
	173	+ if should_print_kernel_source():
	174	+ print(kernel.get_kernel_source())
	175	+ flattened_tokens = num_seqs * num_tokens
	176	+ kernel(
	177	+ d_o.view(flattened_tokens, mhc_mult, hidden_size),
	178	+ comb_res_mix.view(flattened_tokens, mhc_mult, mhc_mult),
	179	+ residual.view(flattened_tokens, mhc_mult, hidden_size),
	180	+ post_layer_mix.view(flattened_tokens, mhc_mult),
	181	+ x.view(flattened_tokens, hidden_size),
	182	+ grad_comb_res_mix.view(flattened_tokens, mhc_mult, mhc_mult),
	183	+ grad_residual.view(flattened_tokens, mhc_mult, hidden_size),
	184	+ grad_post_layer_mix.view(flattened_tokens, mhc_mult),
	185	+ grad_x.view(flattened_tokens, hidden_size),
	186	+ )
	187	+ return grad_x, grad_residual, grad_post_layer_mix, grad_comb_res_mix

Asrc/cann_ops_tilelang/mhc/post_kernel.py+290-0

@@ -0,0 +1,290 @@
	1	+"""TileLang Ascend kernels for MHC post forward and backward."""
	2	+
	3	+import tilelang
	4	+from tilelang.ascend import language as T
	5	+from tilelang.ascend.language import simd as S
	6	+
	7	+_VECTOR_LANES = 64
	8	+
	9	+
	10	+@tilelang.jit
	11	+def get_mhc_post_forward_kernel(
	12	+ mhc_mult: int,
	13	+ hidden_size: int,
	14	+ hidden_block: int,
	15	+ num_vector_cores: int,
	16	+):
	17	+ """Build the Ascend MHC post forward kernel."""
	18	+ num_tokens = T.dynamic('num_tokens')
	19	+ num_hidden_tiles = T.ceildiv(hidden_size, hidden_block)
	20	+ num_chunks = hidden_block // _VECTOR_LANES
	21	+ num_stages = 3
	22	+
	23	+ @T.prim_func
	24	+ def mhc_post_forward_kernel(
	25	+ comb_res_mix: T.Tensor[(num_tokens, mhc_mult, mhc_mult), T.float32],
	26	+ residual: T.Tensor[(num_tokens, mhc_mult, hidden_size), T.bfloat16],
	27	+ post_layer_mix: T.Tensor[(num_tokens, mhc_mult), T.float32],
	28	+ x: T.Tensor[(num_tokens, hidden_size), T.bfloat16],
	29	+ out: T.Tensor[(num_tokens, mhc_mult, hidden_size), T.bfloat16],
	30	+ ) -> None:
	31	+ with T.Kernel(num_vector_cores) as core_id:
	32	+ residual_ub = T.alloc_shared((mhc_mult, hidden_block), T.bfloat16)
	33	+ x_ub = T.alloc_shared((hidden_block,), T.bfloat16)
	34	+ comb_res_mix_ub = T.alloc_shared((mhc_mult, mhc_mult), T.float32)
	35	+ post_layer_mix_ub = T.alloc_shared((T.ceildiv(mhc_mult, 8) * 8,), T.float32)
	36	+ out_ub = T.alloc_shared((mhc_mult, hidden_block), T.bfloat16)
	37	+ T.annotate_buffer_versions(
	38	+ {
	39	+ residual_ub: num_stages,
	40	+ x_ub: num_stages,
	41	+ comb_res_mix_ub: num_stages,
	42	+ post_layer_mix_ub: num_stages,
	43	+ out_ub: num_stages,
	44	+ }
	45	+ )
	46	+
	47	+ for token_id, hidden_tile in T.Persistent(
	48	+ [num_tokens, num_hidden_tiles],
	49	+ num_vector_cores,
	50	+ core_id,
	51	+ group_size=1,
	52	+ num_stages=num_stages,
	53	+ ):
	54	+ T.copy(comb_res_mix[token_id, 0, 0], comb_res_mix_ub)
	55	+ T.copy(post_layer_mix[token_id, 0], post_layer_mix_ub[0:mhc_mult])
	56	+ T.copy(residual[token_id, 0, hidden_tile * hidden_block], residual_ub)
	57	+ T.copy(x[token_id, hidden_tile * hidden_block], x_ub)
	58	+
	59	+ with T.SimdVF():
	60	+ comb_mix = S.alloc_local((mhc_mult * mhc_mult,), T.float32)
	61	+ post_mix = S.alloc_local((mhc_mult,), T.float32)
	62	+ for input_head in T.unroll(mhc_mult, explicit=True):
	63	+ for output_head in T.unroll(mhc_mult, explicit=True):
	64	+ comb_mix[input_head * mhc_mult + output_head] = S.vld(
	65	+ comb_res_mix_ub[input_head, output_head],
	66	+ dist='BRC_B32',
	67	+ )
	68	+ for output_head in T.unroll(mhc_mult, explicit=True):
	69	+ post_mix[output_head] = S.vld(
	70	+ post_layer_mix_ub[output_head],
	71	+ dist='BRC_B32',
	72	+ )
	73	+ out_reg = S.alloc_local((mhc_mult,), T.float32)
	74	+ residual_reg = S.alloc_local((mhc_mult,), T.float32)
	75	+ for chunk in T.serial(num_chunks):
	76	+ column = chunk * _VECTOR_LANES
	77	+ x_fp32 = S.vcvt(
	78	+ S.vld(x_ub[column], dist='UNPK_B16'),
	79	+ T.float32,
	80	+ part=0,
	81	+ )
	82	+ for input_head in T.unroll(mhc_mult, explicit=True):
	83	+ residual_reg[input_head] = S.vcvt(
	84	+ S.vld(residual_ub[input_head, column], dist='UNPK_B16'),
	85	+ T.float32,
	86	+ part=0,
	87	+ )
	88	+ for output_head in T.unroll(mhc_mult, explicit=True):
	89	+ out_reg[output_head] = S.vmul(x_fp32, post_mix[output_head])
	90	+ for input_head in T.unroll(mhc_mult, explicit=True):
	91	+ for output_head in T.unroll(mhc_mult, explicit=True):
	92	+ S.vmula(
	93	+ out_reg[output_head],
	94	+ residual_reg[input_head],
	95	+ comb_mix[input_head * mhc_mult + output_head],
	96	+ )
	97	+ for output_head in T.unroll(mhc_mult, explicit=True):
	98	+ S.vsts(
	99	+ out_ub[output_head, column],
	100	+ S.vcvt(out_reg[output_head], T.bfloat16),
	101	+ dist='PK_B32',
	102	+ )
	103	+
	104	+ T.copy(out_ub, out[token_id, 0, hidden_tile * hidden_block])
	105	+
	106	+ return mhc_post_forward_kernel
	107	+
	108	+
	109	+@tilelang.jit
	110	+def get_mhc_post_backward_kernel(
	111	+ mhc_mult: int,
	112	+ hidden_size: int,
	113	+ hidden_block: int,
	114	+ num_stages: int,
	115	+ num_vector_cores: int,
	116	+):
	117	+ """Build the Ascend MHC post backward kernel."""
	118	+ num_tokens = T.dynamic('num_tokens')
	119	+ num_hidden_tiles = T.ceildiv(hidden_size, hidden_block)
	120	+ num_chunks = hidden_block // _VECTOR_LANES
	121	+
	122	+ @T.prim_func
	123	+ def mhc_post_backward_kernel(
	124	+ grad_out: T.Tensor[(num_tokens, mhc_mult, hidden_size), T.bfloat16],
	125	+ comb_res_mix: T.Tensor[(num_tokens, mhc_mult, mhc_mult), T.float32],
	126	+ residual: T.Tensor[(num_tokens, mhc_mult, hidden_size), T.bfloat16],
	127	+ post_layer_mix: T.Tensor[(num_tokens, mhc_mult), T.float32],
	128	+ x: T.Tensor[(num_tokens, hidden_size), T.bfloat16],
	129	+ grad_comb_res_mix: T.Tensor[(num_tokens, mhc_mult, mhc_mult), T.float32],
	130	+ grad_residual: T.Tensor[(num_tokens, mhc_mult, hidden_size), T.bfloat16],
	131	+ grad_post_layer_mix: T.Tensor[(num_tokens, mhc_mult), T.float32],
	132	+ grad_x: T.Tensor[(num_tokens, hidden_size), T.bfloat16],
	133	+ ) -> None:
	134	+ with T.Kernel(num_vector_cores) as core_id:
	135	+ grad_out_ub = T.alloc_shared((mhc_mult, hidden_block), T.bfloat16)
	136	+ x_ub = T.alloc_shared((hidden_block,), T.bfloat16)
	137	+ residual_ub = T.alloc_shared((mhc_mult, hidden_block), T.bfloat16)
	138	+ comb_res_mix_ub = T.alloc_shared((mhc_mult, mhc_mult), T.float32)
	139	+ post_layer_mix_ub = T.alloc_shared((T.ceildiv(mhc_mult, 8) * 8,), T.float32)
	140	+ grad_x_ub = T.alloc_shared((hidden_block,), T.bfloat16)
	141	+ grad_residual_ub = T.alloc_shared((mhc_mult, hidden_block), T.bfloat16)
	142	+ grad_post_mix_ub = T.alloc_shared((T.ceildiv(mhc_mult, 8) * 8,), T.float32)
	143	+ grad_comb_mix_ub = T.alloc_shared((mhc_mult, mhc_mult), T.float32)
	144	+ post_mix_acc = T.alloc_shared((mhc_mult, _VECTOR_LANES), T.float32)
	145	+ comb_mix_acc = T.alloc_shared((mhc_mult * mhc_mult, _VECTOR_LANES), T.float32)
	146	+ T.annotate_buffer_versions(
	147	+ {
	148	+ grad_out_ub: num_stages,
	149	+ x_ub: num_stages,
	150	+ residual_ub: num_stages,
	151	+ grad_x_ub: num_stages,
	152	+ grad_residual_ub: num_stages,
	153	+ comb_res_mix_ub: num_stages,
	154	+ post_layer_mix_ub: num_stages,
	155	+ grad_post_mix_ub: 1,
	156	+ grad_comb_mix_ub: 1,
	157	+ post_mix_acc: 1,
	158	+ comb_mix_acc: 1,
	159	+ }
	160	+ )
	161	+
	162	+ for token_id in T.Persistent(
	163	+ [num_tokens],
	164	+ num_vector_cores,
	165	+ core_id,
	166	+ group_size=1,
	167	+ num_stages=num_stages,
	168	+ ):
	169	+ T.copy(comb_res_mix[token_id, 0, 0], comb_res_mix_ub)
	170	+ T.copy(post_layer_mix[token_id, 0], post_layer_mix_ub[0:mhc_mult])
	171	+
	172	+ with T.SimdVF():
	173	+ zero = S.vdup(0.0, T.float32)
	174	+ for output_head in T.unroll(mhc_mult, explicit=True):
	175	+ S.vsts(post_mix_acc[output_head, 0], zero)
	176	+ for index in T.unroll(mhc_mult * mhc_mult, explicit=True):
	177	+ S.vsts(comb_mix_acc[index, 0], zero)
	178	+
	179	+ for hidden_tile in T.serial(num_hidden_tiles):
	180	+ T.copy(grad_out[token_id, 0, hidden_tile * hidden_block], grad_out_ub)
	181	+ T.copy(x[token_id, hidden_tile * hidden_block], x_ub)
	182	+ T.copy(residual[token_id, 0, hidden_tile * hidden_block], residual_ub)
	183	+
	184	+ with T.SimdVF():
	185	+ post_mix = S.alloc_local((mhc_mult,), T.float32)
	186	+ for output_head in T.unroll(mhc_mult, explicit=True):
	187	+ post_mix[output_head] = S.vld(
	188	+ post_layer_mix_ub[output_head],
	189	+ dist='BRC_B32',
	190	+ )
	191	+ comb_mix = S.alloc_local((mhc_mult * mhc_mult,), T.float32)
	192	+ for input_head in T.unroll(mhc_mult, explicit=True):
	193	+ for output_head in T.unroll(mhc_mult, explicit=True):
	194	+ comb_mix[input_head * mhc_mult + output_head] = S.vld(
	195	+ comb_res_mix_ub[input_head, output_head],
	196	+ dist='BRC_B32',
	197	+ )
	198	+
	199	+ grad_out_reg = S.alloc_local((mhc_mult,), T.float32)
	200	+ grad_x_reg = S.alloc_var(T.float32)
	201	+ grad_residual_reg = S.alloc_var(T.float32)
	202	+ post_acc_reg = S.alloc_var(T.float32)
	203	+ grad_comb_mix_input = S.alloc_local((mhc_mult,), T.float32)
	204	+
	205	+ for chunk in T.serial(num_chunks):
	206	+ column = chunk * _VECTOR_LANES
	207	+ x_fp32 = S.vcvt(
	208	+ S.vld(x_ub[column], dist='UNPK_B16'),
	209	+ T.float32,
	210	+ part=0,
	211	+ )
	212	+ for output_head in T.unroll(mhc_mult, explicit=True):
	213	+ grad_out_reg[output_head] = S.vcvt(
	214	+ S.vld(grad_out_ub[output_head, column], dist='UNPK_B16'),
	215	+ T.float32,
	216	+ part=0,
	217	+ )
	218	+
	219	+ grad_x_reg = S.vdup(0.0, T.float32)
	220	+ for output_head in T.unroll(mhc_mult, explicit=True):
	221	+ S.vmula(grad_x_reg, grad_out_reg[output_head], post_mix[output_head])
	222	+ S.vsts(
	223	+ grad_x_ub[column],
	224	+ S.vcvt(grad_x_reg, T.bfloat16),
	225	+ dist='PK_B32',
	226	+ )
	227	+
	228	+ for output_head in T.unroll(mhc_mult, explicit=True):
	229	+ post_acc_reg = S.vld(post_mix_acc[output_head, 0])
	230	+ S.vmula(post_acc_reg, grad_out_reg[output_head], x_fp32)
	231	+ S.vsts(post_mix_acc[output_head, 0], post_acc_reg)
	232	+
	233	+ for input_head in T.unroll(mhc_mult, explicit=True):
	234	+ residual_input = S.vcvt(
	235	+ S.vld(residual_ub[input_head, column], dist='UNPK_B16'),
	236	+ T.float32,
	237	+ part=0,
	238	+ )
	239	+ for output_head in T.unroll(mhc_mult, explicit=True):
	240	+ grad_comb_mix_input[output_head] = S.vld(
	241	+ comb_mix_acc[input_head * mhc_mult + output_head, 0]
	242	+ )
	243	+ grad_residual_reg = S.vdup(0.0, T.float32)
	244	+ for output_head in T.unroll(mhc_mult, explicit=True):
	245	+ S.vmula(
	246	+ grad_residual_reg,
	247	+ grad_out_reg[output_head],
	248	+ comb_mix[input_head * mhc_mult + output_head],
	249	+ )
	250	+ S.vmula(
	251	+ grad_comb_mix_input[output_head],
	252	+ grad_out_reg[output_head],
	253	+ residual_input,
	254	+ )
	255	+ S.vsts(
	256	+ grad_residual_ub[input_head, column],
	257	+ S.vcvt(grad_residual_reg, T.bfloat16),
	258	+ dist='PK_B32',
	259	+ )
	260	+ for output_head in T.unroll(mhc_mult, explicit=True):
	261	+ S.vsts(
	262	+ comb_mix_acc[input_head * mhc_mult + output_head, 0],
	263	+ grad_comb_mix_input[output_head],
	264	+ )
	265	+
	266	+ T.copy(grad_x_ub, grad_x[token_id, hidden_tile * hidden_block])
	267	+ T.copy(
	268	+ grad_residual_ub,
	269	+ grad_residual[token_id, 0, hidden_tile * hidden_block],
	270	+ )
	271	+
	272	+ with T.SimdVF():
	273	+ for output_head in T.unroll(mhc_mult, explicit=True):
	274	+ S.vsts(
	275	+ grad_post_mix_ub[output_head],
	276	+ S.vcadd(S.vld(post_mix_acc[output_head, 0])),
	277	+ dist='ONEPT_B32',
	278	+ )
	279	+ for input_head in T.unroll(mhc_mult, explicit=True):
	280	+ for output_head in T.unroll(mhc_mult, explicit=True):
	281	+ S.vsts(
	282	+ grad_comb_mix_ub[input_head, output_head],
	283	+ S.vcadd(S.vld(comb_mix_acc[input_head * mhc_mult + output_head, 0])),
	284	+ dist='ONEPT_B32',
	285	+ )
	286	+
	287	+ T.copy(grad_post_mix_ub[0:mhc_mult], grad_post_layer_mix[token_id, 0])
	288	+ T.copy(grad_comb_mix_ub, grad_comb_res_mix[token_id, 0, 0])
	289	+
	290	+ return mhc_post_backward_kernel

Asrc/cann_ops_tilelang/mhc/pre_apply_mix.py+138-0

@@ -0,0 +1,138 @@
	1	+"""Ascend MHC pre-apply mix forward and backward operators."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import torch
	6	+
	7	+from cann_ops_tilelang.config import get_num_vector_cores, require_ascend
	8	+from cann_ops_tilelang.utils import align_up, ceil_div
	9	+
	10	+from .pre_apply_mix_kernel import mhc_pre_apply_mix_bwd_asc, mhc_pre_apply_mix_fwd_asc
	11	+
	12	+
	13	+def _check_tensor(
	14	+ name: str,
	15	+ tensor: torch.Tensor,
	16	+ dtype: torch.dtype,
	17	+ ndim: int,
	18	+ *,
	19	+ device: torch.device \| None = None,
	20	+) -> None:
	21	+ if not isinstance(tensor, torch.Tensor):
	22	+ raise TypeError(f'{name} must be a torch.Tensor.')
	23	+ if tensor.device.type != 'npu':
	24	+ raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')
	25	+ if device is not None and tensor.device != device:
	26	+ raise ValueError(f'{name} must be on {device}, got {tensor.device}.')
	27	+ if tensor.dtype != dtype:
	28	+ raise TypeError(f'{name} must use {dtype}, got {tensor.dtype}.')
	29	+ if tensor.ndim != ndim:
	30	+ raise ValueError(f'{name} must be {ndim}-dimensional, got shape {tuple(tensor.shape)}.')
	31	+ if not tensor.is_contiguous():
	32	+ raise ValueError(f'{name} must be contiguous.')
	33	+
	34	+
	35	+def _shares_storage(left: torch.Tensor, right: torch.Tensor) -> bool:
	36	+ return (
	37	+ left.numel() > 0
	38	+ and right.numel() > 0
	39	+ and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr()
	40	+ )
	41	+
	42	+
	43	+def _check_shape(num_tokens: int, mhc: int, hidden: int) -> None:
	44	+ if num_tokens < 0:
	45	+ raise ValueError(f'num_tokens must be non-negative, got {num_tokens}.')
	46	+ if mhc <= 0:
	47	+ raise ValueError(f'mhc must be positive, got {mhc}.')
	48	+ if hidden <= 0 or hidden % 64:
	49	+ raise ValueError(f'hidden must be a positive multiple of 64, got {hidden}.')
	50	+
	51	+
	52	+def mhc_pre_apply_mix_forward(
	53	+ x: torch.Tensor,
	54	+ mix: torch.Tensor,
	55	+ o: torch.Tensor,
	56	+) -> None:
	57	+ """Apply token-local MHC mixing into a preallocated BF16 output.
	58	+
	59	+ ``x`` has shape ``(num_tokens, mhc, hidden)``, ``mix`` has shape
	60	+ ``(num_tokens, mhc)``, and ``o`` has shape
	61	+ ``(num_tokens, hidden)``.
	62	+ """
	63	+
	64	+ if not all(isinstance(tensor, torch.Tensor) for tensor in (x, mix, o)):
	65	+ raise TypeError('x, mix and o must be torch.Tensor instances.')
	66	+ require_ascend()
	67	+ _check_tensor('x', x, torch.bfloat16, 3)
	68	+ _check_tensor('mix', mix, torch.float32, 2, device=x.device)
	69	+ _check_tensor('o', o, torch.bfloat16, 2, device=x.device)
	70	+
	71	+ num_tokens, mhc, hidden = x.shape
	72	+ _check_shape(num_tokens, mhc, hidden)
	73	+ if mix.shape != (num_tokens, mhc):
	74	+ raise ValueError(f'mix must have shape {(num_tokens, mhc)}, got {tuple(mix.shape)}.')
	75	+ if o.shape != (num_tokens, hidden):
	76	+ raise ValueError(f'o must have shape {(num_tokens, hidden)}, got {tuple(o.shape)}.')
	77	+ if _shares_storage(o, x) or _shares_storage(o, mix):
	78	+ raise ValueError('o must not share storage with x or mix.')
	79	+ if num_tokens == 0:
	80	+ return
	81	+
	82	+ hidden_block = hidden if hidden <= 6144 else min(ceil_div(hidden, 2), 6144)
	83	+ hidden_block = align_up(hidden_block, 64)
	84	+ mhc_pre_apply_mix_fwd_asc(
	85	+ x,
	86	+ mix,
	87	+ o,
	88	+ h_blk=hidden_block,
	89	+ num_vec_cores=get_num_vector_cores(),
	90	+ )
	91	+
	92	+
	93	+def mhc_pre_apply_mix_backward(
	94	+ o_grad: torch.Tensor,
	95	+ x: torch.Tensor,
	96	+ mix: torch.Tensor,
	97	+ x_grad: torch.Tensor,
	98	+) -> torch.Tensor:
	99	+ """Accumulate ``x`` gradients and return the FP32 mix gradient."""
	100	+
	101	+ tensors = (o_grad, x, mix, x_grad)
	102	+ if not all(isinstance(tensor, torch.Tensor) for tensor in tensors):
	103	+ raise TypeError('o_grad, x, mix and x_grad must be torch.Tensor instances.')
	104	+ require_ascend()
	105	+ _check_tensor('x', x, torch.bfloat16, 3)
	106	+ _check_tensor('o_grad', o_grad, torch.bfloat16, 2, device=x.device)
	107	+ _check_tensor('mix', mix, torch.float32, 2, device=x.device)
	108	+ _check_tensor('x_grad', x_grad, torch.bfloat16, 3, device=x.device)
	109	+
	110	+ num_tokens, mhc, hidden = x.shape
	111	+ _check_shape(num_tokens, mhc, hidden)
	112	+ if o_grad.shape != (num_tokens, hidden):
	113	+ raise ValueError(f'o_grad must have shape {(num_tokens, hidden)}, got {tuple(o_grad.shape)}.')
	114	+ if mix.shape != (num_tokens, mhc):
	115	+ raise ValueError(f'mix must have shape {(num_tokens, mhc)}, got {tuple(mix.shape)}.')
	116	+ if x_grad.shape != x.shape:
	117	+ raise ValueError(f'x_grad must have shape {tuple(x.shape)}, got {tuple(x_grad.shape)}.')
	118	+ for name, tensor in (('o_grad', o_grad), ('x', x), ('mix', mix)):
	119	+ if _shares_storage(x_grad, tensor):
	120	+ raise ValueError(f'x_grad must not share storage with {name}.')
	121	+ if num_tokens == 0:
	122	+ return torch.empty((0, mhc), dtype=torch.float32, device=x.device)
	123	+
	124	+ if hidden <= 4096:
	125	+ hidden_block = hidden
	126	+ num_stages = 3
	127	+ else:
	128	+ hidden_block = align_up(min(ceil_div(hidden, 4), 3072), 64)
	129	+ num_stages = 4
	130	+ return mhc_pre_apply_mix_bwd_asc(
	131	+ o_grad,
	132	+ x,
	133	+ mix,
	134	+ x_grad,
	135	+ h_blk=hidden_block,
	136	+ num_stages=num_stages,
	137	+ num_vec_cores=get_num_vector_cores(),
	138	+ )

Asrc/cann_ops_tilelang/mhc/pre_apply_mix_kernel.py+128-0

@@ -0,0 +1,128 @@
	1	+import tilelang
	2	+from tilelang.ascend import language as T
	3	+from tilelang.ascend.language import simd as S
	4	+
	5	+# Each bf16 chunk is unpacked into a 64-lane fp32 vector register.
	6	+_VEC = 64
	7	+
	8	+
	9	+@tilelang.jit()
	10	+def mhc_pre_apply_mix_fwd_asc(x, mix, o, h_blk, num_vec_cores: int):
	11	+ """Ascend forward: ``o[n, h] = sum_mhc mix[n, mhc] * x[n, mhc, h]``."""
	12	+ num_tokens = T.dynamic('num_tokens')
	13	+ mhc = T.const('mhc')
	14	+ hidden = T.const('hidden')
	15	+ x: T.Tensor[(num_tokens, mhc, hidden), T.bfloat16]
	16	+ mix: T.Tensor[(num_tokens, mhc), T.float32]
	17	+ o: T.Tensor[(num_tokens, hidden), T.bfloat16]
	18	+
	19	+ num_hidden_tiles = T.ceildiv(hidden, h_blk)
	20	+ num_chunks = h_blk // _VEC
	21	+ num_stages = 3
	22	+
	23	+ with T.Kernel(num_vec_cores) as core_id:
	24	+ x_ub = T.alloc_shared((mhc, h_blk), T.bfloat16)
	25	+ mix_ub = T.alloc_shared((T.ceildiv(mhc, 8) * 8,), T.float32)
	26	+ o_ub = T.alloc_shared((h_blk,), T.bfloat16)
	27	+ T.annotate_buffer_versions({x_ub: num_stages, mix_ub: num_stages, o_ub: num_stages})
	28	+
	29	+ for pid_n, pid_h in T.Persistent(
	30	+ [num_tokens, num_hidden_tiles], num_vec_cores, core_id, group_size=1, num_stages=num_stages
	31	+ ):
	32	+ T.copy(mix[pid_n, 0], mix_ub[0:mhc])
	33	+ T.copy(x[pid_n, 0, pid_h * h_blk], x_ub)
	34	+
	35	+ with T.SimdVF():
	36	+ # Preload the mix scalars as broadcast registers once per token-tile.
	37	+ mix_reg = S.alloc_local((mhc,), T.float32)
	38	+ for m in T.unroll(mhc, explicit=True):
	39	+ mix_reg[m] = S.vld(mix_ub[m], dist='BRC_B32')
	40	+ xs = S.alloc_local((mhc,), T.float32)
	41	+ acc = S.alloc_var(T.float32)
	42	+ for c in T.serial(num_chunks):
	43	+ col = c * _VEC
	44	+ for m in T.unroll(mhc, explicit=True):
	45	+ xs[m] = S.vcvt(S.vld(x_ub[m, col], dist='UNPK_B16'), T.float32, part=0)
	46	+ acc = S.vmul(xs[0], mix_reg[0])
	47	+ for m in T.unroll(mhc - 1, explicit=True):
	48	+ S.vmula(acc, xs[m + 1], mix_reg[m + 1])
	49	+ S.vsts(o_ub[col], S.vcvt(acc, T.bfloat16), dist='PK_B32')
	50	+
	51	+ T.copy(o_ub, o[pid_n, pid_h * h_blk])
	52	+
	53	+
	54	+@tilelang.jit()
	55	+def mhc_pre_apply_mix_bwd_asc(o_grad, x, mix, x_grad, h_blk, num_stages, num_vec_cores: int):
	56	+ """Ascend backward.
	57	+
	58	+ ``x_grad[n, m, h] += mix[n, m] * o_grad[n, h]`` (accumulated onto existing values)
	59	+ ``mix_grad[n, m] = sum_h o_grad[n, h] * x[n, m, h]`` (reduction over h)
	60	+ """
	61	+ num_tokens = T.dynamic('num_tokens')
	62	+ mhc = T.const('mhc')
	63	+ hidden = T.const('hidden')
	64	+ o_grad: T.Tensor[(num_tokens, hidden), T.bfloat16]
	65	+ x: T.Tensor[(num_tokens, mhc, hidden), T.bfloat16]
	66	+ mix: T.Tensor[(num_tokens, mhc), T.float32]
	67	+ x_grad: T.Tensor[(num_tokens, mhc, hidden), T.bfloat16]
	68	+ mix_grad = T.empty((num_tokens, mhc), T.float32)
	69	+
	70	+ num_hidden_tiles = T.ceildiv(hidden, h_blk)
	71	+ num_chunks = h_blk // _VEC
	72	+
	73	+ with T.Kernel(num_vec_cores) as core_id:
	74	+ o_grad_ub = T.alloc_shared((h_blk,), T.bfloat16)
	75	+ x_ub = T.alloc_shared((mhc, h_blk), T.bfloat16)
	76	+ x_grad_ub = T.alloc_shared((mhc, h_blk), T.bfloat16)
	77	+ mix_ub = T.alloc_shared((T.ceildiv(mhc, 8) * 8,), T.float32)
	78	+ mix_grad_acc = T.alloc_shared((mhc, _VEC), T.float32)
	79	+ mix_grad_ub = T.alloc_shared((T.ceildiv(mhc, 8) * 8,), T.float32)
	80	+ T.annotate_buffer_versions(
	81	+ {
	82	+ o_grad_ub: num_stages,
	83	+ x_ub: num_stages,
	84	+ x_grad_ub: num_stages,
	85	+ mix_ub: num_stages,
	86	+ mix_grad_ub: num_stages,
	87	+ }
	88	+ )
	89	+
	90	+ for pid_n in T.Persistent([num_tokens], num_vec_cores, core_id, group_size=1, num_stages=num_stages):
	91	+ T.copy(mix[pid_n, 0], mix_ub[0:mhc])
	92	+
	93	+ with T.SimdVF():
	94	+ zero = S.vdup(0.0, T.float32)
	95	+ for m in T.unroll(mhc, explicit=True):
	96	+ S.vsts(mix_grad_acc[m, 0], zero)
	97	+
	98	+ for pid_h in T.serial(num_hidden_tiles):
	99	+ T.copy(o_grad[pid_n, pid_h * h_blk], o_grad_ub)
	100	+ T.copy(x[pid_n, 0, pid_h * h_blk], x_ub)
	101	+ T.copy(x_grad[pid_n, 0, pid_h * h_blk], x_grad_ub)
	102	+
	103	+ with T.SimdVF():
	104	+ mix_reg = S.alloc_local((mhc,), T.float32)
	105	+ for m in T.unroll(mhc, explicit=True):
	106	+ mix_reg[m] = S.vld(mix_ub[m], dist='BRC_B32')
	107	+ mix_grad_reg = S.alloc_var(T.float32)
	108	+ x_grad_f32 = S.alloc_var(T.float32)
	109	+ for m in T.unroll(mhc, explicit=True):
	110	+ mix_grad_reg = S.vld(mix_grad_acc[m, 0])
	111	+ for c in T.serial(num_chunks):
	112	+ col = c * _VEC
	113	+ o_grad_f32 = S.vcvt(S.vld(o_grad_ub[col], dist='UNPK_B16'), T.float32, part=0)
	114	+ x_f32 = S.vcvt(S.vld(x_ub[m, col], dist='UNPK_B16'), T.float32, part=0)
	115	+ x_grad_f32 = S.vcvt(S.vld(x_grad_ub[m, col], dist='UNPK_B16'), T.float32, part=0)
	116	+ S.vmula(x_grad_f32, o_grad_f32, mix_reg[m])
	117	+ S.vsts(x_grad_ub[m, col], S.vcvt(x_grad_f32, T.bfloat16), dist='PK_B32')
	118	+ S.vmula(mix_grad_reg, o_grad_f32, x_f32)
	119	+ S.vsts(mix_grad_acc[m, 0], mix_grad_reg)
	120	+
	121	+ T.copy(x_grad_ub, x_grad[pid_n, 0, pid_h * h_blk])
	122	+
	123	+ with T.SimdVF():
	124	+ for m in T.unroll(mhc, explicit=True):
	125	+ S.vsts(mix_grad_ub[m], S.vcadd(S.vld(mix_grad_acc[m, 0])), dist='ONEPT_B32')
	126	+ T.copy(mix_grad_ub[0:mhc], mix_grad[pid_n, 0])
	127	+
	128	+ return mix_grad

Asrc/cann_ops_tilelang/mhc/pre_big_fuse.py+291-0

@@ -0,0 +1,291 @@
	1	+"""Ascend fused MHC preprocessing operator."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import math
	6	+from numbers import Integral, Real
	7	+
	8	+import torch
	9	+
	10	+from cann_ops_tilelang.config import get_num_vector_cores, require_ascend
	11	+
	12	+from .pre_big_fuse_kernel import (
	13	+ _mhc_reduce_partials_and_rmsnorm_forward_kernel,
	14	+ mhc_pre_big_fuse_fwd_asc,
	15	+)
	16	+
	17	+_ASCEND_MAX_H_BLK = 6144
	18	+_ASCEND_STATIC_TOKEN_COUNTS = (512, 1024, 2048, 8192)
	19	+
	20	+
	21	+def _ascend_pre_big_fuse_config(num_tokens: int, hidden_size: int) -> dict[str, int \| bool]:
	22	+ """Mirror the current TileKernels main launch policy for Ascend."""
	23	+
	24	+ use_batched_residual = num_tokens == 8192 and hidden_size == 7168
	25	+ if use_batched_residual:
	26	+ hidden_block = 256
	27	+ residual_token_block = 16
	28	+ residual_num_stages = 4
	29	+ else:
	30	+ hidden_block = min(hidden_size, _ASCEND_MAX_H_BLK)
	31	+ hidden_block -= hidden_block % 64
	32	+ while hidden_block > 0 and hidden_size % hidden_block:
	33	+ hidden_block -= 64
	34	+ if hidden_block == 0:
	35	+ raise ValueError(f'No 64-aligned hidden block divides hidden={hidden_size}.')
	36	+ residual_token_block = 1
	37	+ residual_num_stages = 3
	38	+
	39	+ return {
	40	+ 'h_blk': hidden_block,
	41	+ 'residual_token_block': residual_token_block,
	42	+ 'residual_num_stages': residual_num_stages,
	43	+ 'static_num_tokens': num_tokens if num_tokens in _ASCEND_STATIC_TOKEN_COUNTS else 0,
	44	+ 'use_fma_affine': num_tokens == 8192,
	45	+ }
	46	+
	47	+
	48	+def _check_tensor(
	49	+ name: str,
	50	+ tensor: torch.Tensor,
	51	+ dtype: torch.dtype,
	52	+ ndim: int,
	53	+ *,
	54	+ device: torch.device \| None = None,
	55	+) -> None:
	56	+ if not isinstance(tensor, torch.Tensor):
	57	+ raise TypeError(f'{name} must be a torch.Tensor.')
	58	+ if tensor.device.type != 'npu':
	59	+ raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')
	60	+ if device is not None and tensor.device != device:
	61	+ raise ValueError(f'{name} must be on {device}, got {tensor.device}.')
	62	+ if tensor.dtype != dtype:
	63	+ raise TypeError(f'{name} must use {dtype}, got {tensor.dtype}.')
	64	+ if tensor.ndim != ndim:
	65	+ raise ValueError(f'{name} must be {ndim}-dimensional, got shape {tuple(tensor.shape)}.')
	66	+ if not tensor.is_contiguous():
	67	+ raise ValueError(f'{name} must be contiguous.')
	68	+
	69	+
	70	+def _shares_storage(left: torch.Tensor, right: torch.Tensor) -> bool:
	71	+ return (
	72	+ left.numel() > 0
	73	+ and right.numel() > 0
	74	+ and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr()
	75	+ )
	76	+
	77	+
	78	+def _check_real(name: str, value: float, *, non_negative: bool = False) -> float:
	79	+ if isinstance(value, bool) or not isinstance(value, Real):
	80	+ raise TypeError(f'{name} must be a real number.')
	81	+ converted = float(value)
	82	+ if not math.isfinite(converted):
	83	+ raise ValueError(f'{name} must be finite.')
	84	+ if non_negative and converted < 0:
	85	+ raise ValueError(f'{name} must be non-negative.')
	86	+ return converted
	87	+
	88	+
	89	+def _check_positive_int(name: str, value: int) -> int:
	90	+ if isinstance(value, bool) or not isinstance(value, Integral):
	91	+ raise TypeError(f'{name} must be an integer.')
	92	+ converted = int(value)
	93	+ if converted <= 0:
	94	+ raise ValueError(f'{name} must be positive, got {converted}.')
	95	+ return converted
	96	+
	97	+
	98	+def _check_output_aliases(inputs: dict[str, torch.Tensor], outputs: dict[str, torch.Tensor]) -> None:
	99	+ output_items = list(outputs.items())
	100	+ for output_name, output in output_items:
	101	+ for input_name, input_tensor in inputs.items():
	102	+ if _shares_storage(output, input_tensor):
	103	+ raise ValueError(f'{output_name} must not share storage with {input_name}.')
	104	+ for index, (left_name, left) in enumerate(output_items):
	105	+ for right_name, right in output_items[index + 1 :]:
	106	+ if _shares_storage(left, right):
	107	+ raise ValueError(f'{left_name} and {right_name} must not share storage.')
	108	+
	109	+
	110	+def _mhc_reduce_partials_and_rmsnorm_forward(
	111	+ out_mul_splitted: torch.Tensor,
	112	+ sqrsum_splitted: torch.Tensor,
	113	+ out_mul: torch.Tensor,
	114	+ sqrsum: torch.Tensor,
	115	+ out: torch.Tensor,
	116	+ rms_group_size: int,
	117	+ rms_eps: float,
	118	+ n_splits: int,
	119	+) -> None:
	120	+ """Run the private reduce/RMSNorm reference stage used by fused tests."""
	121	+
	122	+ tensors = (out_mul_splitted, sqrsum_splitted, out_mul, sqrsum, out)
	123	+ if not all(isinstance(tensor, torch.Tensor) for tensor in tensors):
	124	+ raise TypeError('all reduce/RMSNorm inputs and outputs must be torch.Tensor instances.')
	125	+ require_ascend()
	126	+ _check_tensor('out_mul_splitted', out_mul_splitted, torch.float32, 4)
	127	+ device = out_mul_splitted.device
	128	+ _check_tensor('sqrsum_splitted', sqrsum_splitted, torch.float32, 3, device=device)
	129	+ _check_tensor('out_mul', out_mul, torch.float32, 3, device=device)
	130	+ _check_tensor('sqrsum', sqrsum, torch.float32, 2, device=device)
	131	+ _check_tensor('out', out, torch.float32, 2, device=device)
	132	+
	133	+ n_splits = _check_positive_int('n_splits', n_splits)
	134	+ rms_group_size = _check_positive_int('rms_group_size', rms_group_size)
	135	+ rms_eps = _check_real('rms_eps', rms_eps, non_negative=True)
	136	+ if out_mul_splitted.shape[0] != n_splits:
	137	+ raise ValueError(
	138	+ f'out_mul_splitted first dimension must equal n_splits={n_splits}, got {out_mul_splitted.shape[0]}.'
	139	+ )
	140	+ _, num_tokens, num_rms_groups, mhc_mult3 = out_mul_splitted.shape
	141	+ if num_rms_groups != 1 or mhc_mult3 != 24:
	142	+ raise ValueError(
	143	+ f'out_mul_splitted must have shape (n_splits, num_tokens, 1, 24), got {tuple(out_mul_splitted.shape)}.'
	144	+ )
	145	+ expected_shapes = {
	146	+ 'sqrsum_splitted': (n_splits, num_tokens, 1),
	147	+ 'out_mul': (num_tokens, 1, 24),
	148	+ 'sqrsum': (num_tokens, 1),
	149	+ 'out': (num_tokens, 24),
	150	+ }
	151	+ actual_tensors = {
	152	+ 'sqrsum_splitted': sqrsum_splitted,
	153	+ 'out_mul': out_mul,
	154	+ 'sqrsum': sqrsum,
	155	+ 'out': out,
	156	+ }
	157	+ for name, expected_shape in expected_shapes.items():
	158	+ if actual_tensors[name].shape != expected_shape:
	159	+ raise ValueError(f'{name} must have shape {expected_shape}, got {tuple(actual_tensors[name].shape)}.')
	160	+ _check_output_aliases(
	161	+ {
	162	+ 'out_mul_splitted': out_mul_splitted,
	163	+ 'sqrsum_splitted': sqrsum_splitted,
	164	+ },
	165	+ {'out_mul': out_mul, 'sqrsum': sqrsum, 'out': out},
	166	+ )
	167	+ if num_tokens == 0:
	168	+ return
	169	+
	170	+ _mhc_reduce_partials_and_rmsnorm_forward_kernel(
	171	+ out_mul_splitted,
	172	+ sqrsum_splitted,
	173	+ out_mul,
	174	+ sqrsum,
	175	+ out,
	176	+ rms_group_size,
	177	+ rms_eps,
	178	+ n_splits,
	179	+ get_num_vector_cores(),
	180	+ )
	181	+
	182	+
	183	+def mhc_pre_big_fuse(
	184	+ gemm_out_mul: torch.Tensor,
	185	+ gemm_out_sqrsum: torch.Tensor,
	186	+ mhc_scale: torch.Tensor,
	187	+ mhc_base: torch.Tensor,
	188	+ residual: torch.Tensor,
	189	+ post_mix: torch.Tensor,
	190	+ comb_mix: torch.Tensor,
	191	+ layer_input: torch.Tensor,
	192	+ rms_eps: float,
	193	+ mhc_pre_eps: float,
	194	+ mhc_sinkhorn_eps: float,
	195	+ mhc_post_mult_value: float,
	196	+ sinkhorn_repeat: int,
	197	+ n_splits: int = 16,
	198	+) -> None:
	199	+ """Run fused MHC preprocessing into caller-provided output buffers."""
	200	+
	201	+ tensors = (
	202	+ gemm_out_mul,
	203	+ gemm_out_sqrsum,
	204	+ mhc_scale,
	205	+ mhc_base,
	206	+ residual,
	207	+ post_mix,
	208	+ comb_mix,
	209	+ layer_input,
	210	+ )
	211	+ if not all(isinstance(tensor, torch.Tensor) for tensor in tensors):
	212	+ raise TypeError('all MHC pre-big-fuse inputs and outputs must be torch.Tensor instances.')
	213	+ require_ascend()
	214	+ _check_tensor('gemm_out_mul', gemm_out_mul, torch.float32, 3)
	215	+ device = gemm_out_mul.device
	216	+ _check_tensor('gemm_out_sqrsum', gemm_out_sqrsum, torch.float32, 2, device=device)
	217	+ _check_tensor('mhc_scale', mhc_scale, torch.float32, 1, device=device)
	218	+ _check_tensor('mhc_base', mhc_base, torch.float32, 1, device=device)
	219	+ _check_tensor('residual', residual, torch.bfloat16, 3, device=device)
	220	+ _check_tensor('post_mix', post_mix, torch.float32, 2, device=device)
	221	+ _check_tensor('comb_mix', comb_mix, torch.float32, 2, device=device)
	222	+ _check_tensor('layer_input', layer_input, torch.bfloat16, 2, device=device)
	223	+
	224	+ rms_eps = _check_real('rms_eps', rms_eps, non_negative=True)
	225	+ mhc_pre_eps = _check_real('mhc_pre_eps', mhc_pre_eps, non_negative=True)
	226	+ mhc_sinkhorn_eps = _check_real('mhc_sinkhorn_eps', mhc_sinkhorn_eps, non_negative=True)
	227	+ mhc_post_mult_value = _check_real('mhc_post_mult_value', mhc_post_mult_value)
	228	+ sinkhorn_repeat = _check_positive_int('sinkhorn_repeat', sinkhorn_repeat)
	229	+ _check_positive_int('n_splits', n_splits)
	230	+
	231	+ num_tokens, mhc, hidden = residual.shape
	232	+ if mhc != 4:
	233	+ raise ValueError(f'residual must use mHC4, got mhc={mhc}.')
	234	+ if hidden <= 0 or hidden % 64:
	235	+ raise ValueError(f'hidden must be a positive multiple of 64, got {hidden}.')
	236	+ expected_shapes = {
	237	+ 'gemm_out_mul': (1, num_tokens, 24),
	238	+ 'gemm_out_sqrsum': (1, num_tokens),
	239	+ 'mhc_scale': (3,),
	240	+ 'mhc_base': (24,),
	241	+ 'post_mix': (num_tokens, 4),
	242	+ 'comb_mix': (num_tokens, 16),
	243	+ 'layer_input': (num_tokens, hidden),
	244	+ }
	245	+ actual_tensors = {
	246	+ 'gemm_out_mul': gemm_out_mul,
	247	+ 'gemm_out_sqrsum': gemm_out_sqrsum,
	248	+ 'mhc_scale': mhc_scale,
	249	+ 'mhc_base': mhc_base,
	250	+ 'post_mix': post_mix,
	251	+ 'comb_mix': comb_mix,
	252	+ 'layer_input': layer_input,
	253	+ }
	254	+ for name, expected_shape in expected_shapes.items():
	255	+ if actual_tensors[name].shape != expected_shape:
	256	+ raise ValueError(f'{name} must have shape {expected_shape}, got {tuple(actual_tensors[name].shape)}.')
	257	+ _check_output_aliases(
	258	+ {
	259	+ 'gemm_out_mul': gemm_out_mul,
	260	+ 'gemm_out_sqrsum': gemm_out_sqrsum,
	261	+ 'mhc_scale': mhc_scale,
	262	+ 'mhc_base': mhc_base,
	263	+ 'residual': residual,
	264	+ },
	265	+ {'post_mix': post_mix, 'comb_mix': comb_mix, 'layer_input': layer_input},
	266	+ )
	267	+ if num_tokens == 0:
	268	+ return
	269	+
	270	+ config = _ascend_pre_big_fuse_config(num_tokens, hidden)
	271	+ mhc_pre_big_fuse_fwd_asc(
	272	+ gemm_out_mul,
	273	+ gemm_out_sqrsum,
	274	+ mhc_scale,
	275	+ mhc_base,
	276	+ residual,
	277	+ post_mix,
	278	+ comb_mix,
	279	+ layer_input,
	280	+ rms_eps,
	281	+ mhc_pre_eps,
	282	+ mhc_sinkhorn_eps,
	283	+ mhc_post_mult_value,
	284	+ sinkhorn_repeat,
	285	+ h_blk=config['h_blk'],
	286	+ residual_token_block=config['residual_token_block'],
	287	+ residual_num_stages=config['residual_num_stages'],
	288	+ use_fma_affine=config['use_fma_affine'],
	289	+ static_num_tokens=config['static_num_tokens'],
	290	+ n_splits=1,
	291	+ )

Asrc/cann_ops_tilelang/mhc/pre_big_fuse_kernel.py+587-0

Asrc/cann_ops_tilelang/mhc/pre_split_mixes.py+234-0

Asrc/cann_ops_tilelang/mhc/pre_split_mixes_kernel.py+604-0

Asrc/cann_ops_tilelang/mhc/sinkhorn.py+138-0

Asrc/cann_ops_tilelang/mhc/sinkhorn_kernel.py+316-0

Atests/mhc/mhc_ref.py+105-0

Atests/mhc/test_expand.py+126-0

Atests/mhc/test_head_compute_mix.py+213-0

Atests/mhc/test_multilayer_recompute.py+250-0

Atests/mhc/test_norm_fn.py+624-0

Atests/mhc/test_post.py+180-0

Atests/mhc/test_pre_apply_mix.py+140-0

Atests/mhc/test_pre_big_fuse.py+272-0

Atests/mhc/test_pre_split_mixes.py+297-0

Atests/mhc/test_sinkhorn.py+194-0