feat(engram): migrate remaining operators-ops-tilelang-AtomGit

M.gitcode/config/operator_tests.conf+1-1

	@@ -12,7 +12,7 @@ src/cann_ops_tilelang/engram/engram_fused_weight*.py tests/engram/test_engram_fu
12	src/cann_ops_tilelang/engram/engram_gate*.py tests/engram/test_engram_gate_fwd.py tests/engram/test_engram_gate_bwd.py	12	src/cann_ops_tilelang/engram/engram_gate*.py tests/engram/test_engram_gate_fwd.py tests/engram/test_engram_gate_bwd.py
13	src/cann_ops_tilelang/engram/engram_grad_w_reduce*.py tests/engram/test_engram_grad_w_reduce.py	13	src/cann_ops_tilelang/engram/engram_grad_w_reduce*.py tests/engram/test_engram_grad_w_reduce.py
14	src/cann_ops_tilelang/engram/engram_hash*.py tests/engram/test_engram_hash.py	14	src/cann_ops_tilelang/engram/engram_hash*.py tests/engram/test_engram_hash.py
15	-src/cann_ops_tilelang/engram/engram_sinkhorn_step*.py tests/engram/test_engram_sinkhorn_step.py	15	+src/cann_ops_tilelang/engram/engram_sinkhorn*.py tests/engram/test_engram_sinkhorn.py
16		16
17	# MHC	17	# MHC
18	src/cann_ops_tilelang/mhc/expand*.py tests/mhc/test_expand.py	18	src/cann_ops_tilelang/mhc/expand*.py tests/mhc/test_expand.py

Msrc/cann_ops_tilelang/__init__.py+13-1

	@@ -5,11 +5,23 @@ __all__ = [
5	'engram_gate_bwd',	5	'engram_gate_bwd',
6	'engram_gate_fwd',	6	'engram_gate_fwd',
7	'engram_grad_w_reduce',	7	'engram_grad_w_reduce',
		8	+ 'engram_hash',
		9	+ 'engram_sinkhorn_finalize',
		10	+ 'engram_sinkhorn_momentum_update',
8	'engram_sinkhorn_step',	11	'engram_sinkhorn_step',
9	'moe_topk_gate',	12	'moe_topk_gate',
10	'normalize_weight',	13	'normalize_weight',
11	'topk_gate',	14	'topk_gate',
12	]	15	]
13		16
14	-from .engram import engram_fused_weight, engram_gate_bwd, engram_gate_fwd, engram_grad_w_reduce, engram_sinkhorn_step	17	+from .engram import (
		18	+ engram_fused_weight,
		19	+ engram_gate_bwd,
		20	+ engram_gate_fwd,
		21	+ engram_grad_w_reduce,
		22	+ engram_hash,
		23	+ engram_sinkhorn_finalize,
		24	+ engram_sinkhorn_momentum_update,
		25	+ engram_sinkhorn_step,
		26	+)
15	from .moe import moe_topk_gate, normalize_weight, topk_gate	27	from .moe import moe_topk_gate, normalize_weight, topk_gate

Msrc/cann_ops_tilelang/engram/__init__.py+5-1

	@@ -5,10 +5,14 @@ __all__ = [
5	'engram_gate_bwd',	5	'engram_gate_bwd',
6	'engram_gate_fwd',	6	'engram_gate_fwd',
7	'engram_grad_w_reduce',	7	'engram_grad_w_reduce',
		8	+ 'engram_hash',
		9	+ 'engram_sinkhorn_finalize',
		10	+ 'engram_sinkhorn_momentum_update',
8	'engram_sinkhorn_step',	11	'engram_sinkhorn_step',
9	]	12	]
10		13
11	from .engram_fused_weight import engram_fused_weight	14	from .engram_fused_weight import engram_fused_weight
12	from .engram_gate import engram_gate_bwd, engram_gate_fwd	15	from .engram_gate import engram_gate_bwd, engram_gate_fwd
13	from .engram_grad_w_reduce import engram_grad_w_reduce	16	from .engram_grad_w_reduce import engram_grad_w_reduce
14	-from .engram_sinkhorn_step import engram_sinkhorn_step	17	+from .engram_hash import engram_hash
		18	+from .engram_sinkhorn import engram_sinkhorn_finalize, engram_sinkhorn_momentum_update, engram_sinkhorn_step

Asrc/cann_ops_tilelang/engram/engram_hash.py+110-0

@@ -0,0 +1,110 @@
	1	+"""Public interface for the Engram n-gram hash operator."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import math
	6	+
	7	+import torch
	8	+
	9	+from cann_ops_tilelang.config import get_num_vector_cores, require_ascend, should_print_kernel_source
	10	+
	11	+from .engram_hash_kernel import get_engram_hash_kernel
	12	+
	13	+
	14	+def _check_npu_tensor(
	15	+ name: str,
	16	+ tensor: torch.Tensor,
	17	+ *,
	18	+ dtype: torch.dtype,
	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 not tensor.is_contiguous():
	30	+ raise ValueError(f'{name} must be contiguous.')
	31	+
	32	+
	33	+def engram_hash(
	34	+ ngram_token_ids: torch.Tensor,
	35	+ multipliers: torch.Tensor,
	36	+ vocab_sizes: torch.Tensor,
	37	+ offsets: torch.Tensor,
	38	+) -> torch.Tensor:
	39	+ """Compute Engram n-gram hash embedding indices on Ascend.
	40	+
	41	+ ``ngram_token_ids`` has shape ``(num_tokens, max_ngram_size)`` and uses
	42	+ INT32. ``multipliers`` uses INT64, while ``vocab_sizes`` and ``offsets``
	43	+ use INT32. The caller must provide strictly positive vocabulary sizes.
	44	+ """
	45	+
	46	+ if not all(isinstance(tensor, torch.Tensor) for tensor in (ngram_token_ids, multipliers, vocab_sizes, offsets)):
	47	+ raise TypeError('ngram_token_ids, multipliers, vocab_sizes and offsets must be torch.Tensor instances.')
	48	+
	49	+ require_ascend()
	50	+ _check_npu_tensor('ngram_token_ids', ngram_token_ids, dtype=torch.int32)
	51	+ _check_npu_tensor('multipliers', multipliers, dtype=torch.int64, device=ngram_token_ids.device)
	52	+ _check_npu_tensor('vocab_sizes', vocab_sizes, dtype=torch.int32, device=ngram_token_ids.device)
	53	+ _check_npu_tensor('offsets', offsets, dtype=torch.int32, device=ngram_token_ids.device)
	54	+
	55	+ if ngram_token_ids.ndim != 2:
	56	+ raise ValueError('ngram_token_ids must have shape (num_tokens, max_ngram_size).')
	57	+ num_tokens, max_ngram_size = ngram_token_ids.shape
	58	+ if max_ngram_size < 2:
	59	+ raise ValueError(f'max_ngram_size must be at least 2, got {max_ngram_size}.')
	60	+ if vocab_sizes.ndim != 3:
	61	+ raise ValueError(
	62	+ 'vocab_sizes must have shape (num_ngram_layers, max_ngram_size - 1, num_embed_table_per_ngram).'
	63	+ )
	64	+ num_ngram_layers, vocab_ngram_size, num_embed_table_per_ngram = vocab_sizes.shape
	65	+ if num_ngram_layers == 0:
	66	+ raise ValueError('num_ngram_layers must be positive.')
	67	+ if vocab_ngram_size != max_ngram_size - 1:
	68	+ raise ValueError(f'vocab_sizes dimension 1 must be {max_ngram_size - 1}, got {vocab_ngram_size}.')
	69	+ if num_embed_table_per_ngram == 0:
	70	+ raise ValueError('num_embed_table_per_ngram must be positive.')
	71	+ if multipliers.shape != (num_ngram_layers, max_ngram_size):
	72	+ raise ValueError(
	73	+ f'multipliers must have shape {(num_ngram_layers, max_ngram_size)}, got {tuple(multipliers.shape)}.'
	74	+ )
	75	+
	76	+ num_out_cols = (max_ngram_size - 1) * num_embed_table_per_ngram
	77	+ if offsets.shape != (num_ngram_layers, num_out_cols):
	78	+ raise ValueError(f'offsets must have shape {(num_ngram_layers, num_out_cols)}, got {tuple(offsets.shape)}.')
	79	+ token_alignment = math.lcm(
	80	+ 8 // math.gcd(max_ngram_size, 8),
	81	+ 8 // math.gcd(num_out_cols, 8),
	82	+ )
	83	+ if 1024 // num_out_cols < token_alignment:
	84	+ raise ValueError(
	85	+ 'the n-gram configuration exceeds the 1024-thread Ascend SimtVF limit: '
	86	+ f'max_ngram_size={max_ngram_size}, num_out_cols={num_out_cols}.'
	87	+ )
	88	+
	89	+ output = torch.empty(
	90	+ (num_ngram_layers, num_tokens, num_out_cols),
	91	+ dtype=torch.int32,
	92	+ device=ngram_token_ids.device,
	93	+ )
	94	+ if num_tokens == 0:
	95	+ return output
	96	+
	97	+ available_cores = get_num_vector_cores()
	98	+ if num_ngram_layers > available_cores:
	99	+ raise ValueError(
	100	+ f'num_ngram_layers must not exceed the Ascend vector core count {available_cores}, got {num_ngram_layers}.'
	101	+ )
	102	+ kernel = get_engram_hash_kernel(
	103	+ max_ngram_size,
	104	+ num_ngram_layers,
	105	+ num_embed_table_per_ngram,
	106	+ )
	107	+ if should_print_kernel_source():
	108	+ print(kernel.get_kernel_source())
	109	+ kernel(ngram_token_ids, multipliers, vocab_sizes, offsets, output)
	110	+ return output

Asrc/cann_ops_tilelang/engram/engram_hash_kernel.py+107-0

@@ -0,0 +1,107 @@
	1	+"""TileLang Ascend kernel for Engram n-gram hashing."""
	2	+
	3	+import math
	4	+
	5	+import tilelang
	6	+from tilelang import language as T
	7	+
	8	+from cann_ops_tilelang.config import get_num_vector_cores
	9	+
	10	+
	11	+@tilelang.jit
	12	+def get_engram_hash_kernel(
	13	+ max_ngram_size: int = 3,
	14	+ num_ngram_layers: int = 2,
	15	+ num_embed_table_per_ngram: int = 8,
	16	+):
	17	+ """Build the Engram hash kernel for Ascend vector cores."""
	18	+
	19	+ num_tokens = T.dynamic('num_tokens')
	20	+ available_cores = get_num_vector_cores()
	21	+ cores_per_layer = available_cores // num_ngram_layers
	22	+ num_out_cols = (max_ngram_size - 1) * num_embed_table_per_ngram
	23	+ max_simt_threads = 1024
	24	+ input_token_alignment = 8 // math.gcd(max_ngram_size, 8)
	25	+ output_token_alignment = 8 // math.gcd(num_out_cols, 8)
	26	+ token_alignment = math.lcm(input_token_alignment, output_token_alignment)
	27	+ max_tile_tokens = max_simt_threads // num_out_cols
	28	+ assert max_tile_tokens >= token_alignment
	29	+ tile_tokens = max_tile_tokens // token_alignment * token_alignment
	30	+ simt_threads = tile_tokens * num_out_cols
	31	+ num_cores = cores_per_layer * num_ngram_layers
	32	+
	33	+ @T.prim_func
	34	+ def engram_hash_kernel(
	35	+ ngram_token_ids: T.Tensor[(num_tokens, max_ngram_size), T.int32],
	36	+ multipliers: T.Tensor[(num_ngram_layers, max_ngram_size), T.int64],
	37	+ vocab_sizes: T.Tensor[
	38	+ (num_ngram_layers, max_ngram_size - 1, num_embed_table_per_ngram),
	39	+ T.int32,
	40	+ ],
	41	+ offsets: T.Tensor[(num_ngram_layers, num_out_cols), T.int32],
	42	+ output: T.Tensor[(num_ngram_layers, num_tokens, num_out_cols), T.int32],
	43	+ ):
	44	+ with T.Kernel(num_cores) as core_id:
	45	+ layer_idx = core_id // cores_per_layer
	46	+ layer_core_id = core_id % cores_per_layer
	47	+ token_ids_ub = T.alloc_shared((tile_tokens, max_ngram_size), T.int32)
	48	+ multipliers_ub = T.alloc_shared((max_ngram_size,), T.int64)
	49	+ vocab_sizes_ub = T.alloc_shared(
	50	+ (max_ngram_size - 1, num_embed_table_per_ngram),
	51	+ T.int32,
	52	+ )
	53	+ offsets_ub = T.alloc_shared((num_out_cols,), T.int32)
	54	+ output_ub = T.alloc_shared((tile_tokens, num_out_cols), T.int32)
	55	+
	56	+ T.copy(multipliers[layer_idx, :], multipliers_ub)
	57	+ T.copy(vocab_sizes[layer_idx, :, :], vocab_sizes_ub)
	58	+ T.copy(offsets[layer_idx, :], offsets_ub)
	59	+
	60	+ for tile_idx in T.Persistent(
	61	+ [T.ceildiv(num_tokens, tile_tokens)],
	62	+ cores_per_layer,
	63	+ layer_core_id,
	64	+ group_size=1,
	65	+ num_stages=0,
	66	+ ):
	67	+ token_start = tile_idx * tile_tokens
	68	+ valid_tokens = T.min(tile_tokens, num_tokens - token_start)
	69	+ T.copy(
	70	+ ngram_token_ids[token_start : token_start + valid_tokens, :],
	71	+ token_ids_ub[:valid_tokens, :],
	72	+ )
	73	+
	74	+ with T.SimtVF(threads=simt_threads):
	75	+ thread_idx = T.get_thread_binding()
	76	+ token_in_tile = thread_idx // num_out_cols
	77	+ out_col = thread_idx % num_out_cols
	78	+ if token_in_tile < valid_tokens:
	79	+ last_ngram_idx = out_col // num_embed_table_per_ngram + 1
	80	+ out_table_idx = out_col % num_embed_table_per_ngram
	81	+ first_token_id = T.cast(token_ids_ub[token_in_tile, 0], T.int64)
	82	+ first_product = first_token_id * multipliers_ub[0]
	83	+ hash_value = T.alloc_var(T.int64, init=first_product)
	84	+ for hash_ngram_idx in T.serial(1, max_ngram_size):
	85	+ if hash_ngram_idx <= last_ngram_idx:
	86	+ token_id = T.cast(token_ids_ub[token_in_tile, hash_ngram_idx], T.int64)
	87	+ product = token_id * multipliers_ub[hash_ngram_idx]
	88	+ hash_value = T.bitwise_xor(hash_value, product)
	89	+
	90	+ vocab_size = T.cast(
	91	+ vocab_sizes_ub[last_ngram_idx - 1, out_table_idx],
	92	+ T.int64,
	93	+ )
	94	+ remainder = T.truncmod(hash_value, vocab_size)
	95	+ hash_index = T.Select(
	96	+ remainder < 0,
	97	+ remainder + vocab_size,
	98	+ remainder,
	99	+ )
	100	+ output_ub[token_in_tile, out_col] = T.cast(hash_index, T.int32) + offsets_ub[out_col]
	101	+
	102	+ T.copy(
	103	+ output_ub[:valid_tokens, :],
	104	+ output[layer_idx, token_start : token_start + valid_tokens, :],
	105	+ )
	106	+
	107	+ return engram_hash_kernel

Rsrc/cann_ops_tilelang/engram/engram_sinkhorn_step.py→src/cann_ops_tilelang/engram/engram_sinkhorn.py+179-49

	@@ -1,4 +1,4 @@
1	-"""Public interface for one Engram Sinkhorn normalization step."""	1	+"""Public interfaces for the Engram Sinkhorn operator family."""
2		2
3	from __future__ import annotations	3	from __future__ import annotations
4		4
	@@ -6,19 +6,27 @@ import math
6	from numbers import Real	6	from numbers import Real
7		7
8	import torch	8	import torch
		9	+from tilelang import language as T
9		10
10	from cann_ops_tilelang.config import get_num_vector_cores, require_ascend, should_print_kernel_source	11	from cann_ops_tilelang.config import get_num_vector_cores, require_ascend, should_print_kernel_source
11		12
12	-from .engram_sinkhorn_step_kernel import get_engram_sinkhorn_step_kernel	13	+from .engram_sinkhorn_kernel import (
		14	+ get_engram_sinkhorn_finalize_kernel,
		15	+ get_engram_sinkhorn_momentum_update_kernel,
		16	+ get_engram_sinkhorn_step_kernel,
		17	+)
13		18
14	-# SimtVF reserves 32 KiB of the physical 248 KiB UB for thread context.	19	+_PHYSICAL_UB_CAPACITY_BYTES = 248 * 1024
15	-_SHARED_UB_CAPACITY_BYTES = (248 - 32) * 1024	20	+# SimtVF reserves 32 KiB of the physical UB for thread context.
		21	+_SIMT_SHARED_UB_CAPACITY_BYTES = (248 - 32) * 1024
		22	+_MOMENTUM_TILE_SIZE = 4096
16		23
17		24
18	-def _check_npu_fp32_tensor(	25	+def _check_npu_tensor(
19	name: str,	26	name: str,
20	tensor: torch.Tensor,	27	tensor: torch.Tensor,
21	*,	28	*,
		29	+ dtypes: tuple[torch.dtype, ...],
22	device: torch.device \| None = None,	30	device: torch.device \| None = None,
23	) -> None:	31	) -> None:
24	if not isinstance(tensor, torch.Tensor):	32	if not isinstance(tensor, torch.Tensor):
	@@ -27,12 +35,23 @@ def _check_npu_fp32_tensor(
27	raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')	35	raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')
28	if device is not None and tensor.device != device:	36	if device is not None and tensor.device != device:
29	raise ValueError(f'{name} must be on {device}, got {tensor.device}.')	37	raise ValueError(f'{name} must be on {device}, got {tensor.device}.')
30	- if tensor.dtype != torch.float32:	38	+ if tensor.dtype not in dtypes:
31	- raise TypeError(f'{name} must use torch.float32, got {tensor.dtype}.')	39	+ expected = ' or '.join(str(dtype) for dtype in dtypes)
		40	+ raise TypeError(f'{name} must use {expected}, got {tensor.dtype}.')
32	if not tensor.is_contiguous():	41	if not tensor.is_contiguous():
33	raise ValueError(f'{name} must be contiguous.')	42	raise ValueError(f'{name} must be contiguous.')
34		43
35		44
		45	+def _check_real(name: str, value: float, *, non_negative: bool = False) -> float:
		46	+ if not isinstance(value, Real):
		47	+ raise TypeError(f'{name} must be a real number.')
		48	+ value = float(value)
		49	+ if not math.isfinite(value) or (non_negative and value < 0):
		50	+ requirement = 'finite and non-negative' if non_negative else 'finite'
		51	+ raise ValueError(f'{name} must be {requirement}.')
		52	+ return value
		53	+
		54	+
36	def _shares_storage(left: torch.Tensor, right: torch.Tensor) -> bool:	55	def _shares_storage(left: torch.Tensor, right: torch.Tensor) -> bool:
37	return (	56	return (
38	left.numel() > 0	57	left.numel() > 0
	@@ -41,41 +60,98 @@ def _shares_storage(left: torch.Tensor, right: torch.Tensor) -> bool:
41	)	60	)
42		61
43		62
44	-def _validate_inputs(	63	+def _reject_aliases(inputs: dict[str, torch.Tensor]) -> None:
45	- matrix: torch.Tensor,
46	- row_scale: torch.Tensor,
47	- col_scale: torch.Tensor,
48	- eps: float,
49	-) -> tuple[int, int]:
50	- if not all(isinstance(tensor, torch.Tensor) for tensor in (matrix, row_scale, col_scale)):
51	- raise TypeError('matrix, row_scale and col_scale must be torch.Tensor instances.')
52	- if not isinstance(eps, Real):
53	- raise TypeError('eps must be a real number.')
54	- if not math.isfinite(float(eps)) or eps < 0:
55	- raise ValueError('eps must be finite and non-negative.')
56	-
57	- require_ascend()
58	- _check_npu_fp32_tensor('matrix', matrix)
59	- _check_npu_fp32_tensor('row_scale', row_scale, device=matrix.device)
60	- _check_npu_fp32_tensor('col_scale', col_scale, device=matrix.device)
61	-
62	- if matrix.ndim != 2:
63	- raise ValueError('matrix must have shape (num_rows, head_dim).')
64	- num_rows, head_dim = matrix.shape
65	- if head_dim == 0 or head_dim % 256:
66	- raise ValueError(f'head_dim must be a positive multiple of 256, got {head_dim}.')
67	- if row_scale.ndim != 1 or row_scale.shape[0] != num_rows:
68	- raise ValueError(f'row_scale must have shape ({num_rows},), got {tuple(row_scale.shape)}.')
69	- if col_scale.ndim != 1 or col_scale.shape[0] != head_dim:
70	- raise ValueError(f'col_scale must have shape ({head_dim},), got {tuple(col_scale.shape)}.')
71	-
72	- inputs = {'matrix': matrix, 'row_scale': row_scale, 'col_scale': col_scale}
73	input_items = list(inputs.items())	64	input_items = list(inputs.items())
74	for index, (left_name, left) in enumerate(input_items):	65	for index, (left_name, left) in enumerate(input_items):
75	for right_name, right in input_items[index + 1 :]:	66	for right_name, right in input_items[index + 1 :]:
76	if _shares_storage(left, right):	67	if _shares_storage(left, right):
77	raise ValueError(f'{left_name} and {right_name} must not share storage.')	68	raise ValueError(f'{left_name} and {right_name} must not share storage.')
78	- return num_rows, head_dim	69	+
		70	+
		71	+def _validate_matrix_inputs(
		72	+ matrix: torch.Tensor,
		73	+ row_scale: torch.Tensor,
		74	+ col_scale: torch.Tensor,
		75	+ eps: float,
		76	+ *,
		77	+ head_dim_alignment: int,
		78	+) -> tuple[int, int, float]:
		79	+ if not all(isinstance(tensor, torch.Tensor) for tensor in (matrix, row_scale, col_scale)):
		80	+ raise TypeError('matrix, row_scale and col_scale must be torch.Tensor instances.')
		81	+ eps = _check_real('eps', eps, non_negative=True)
		82	+
		83	+ require_ascend()
		84	+ fp32 = (torch.float32,)
		85	+ _check_npu_tensor('matrix', matrix, dtypes=fp32)
		86	+ _check_npu_tensor('row_scale', row_scale, dtypes=fp32, device=matrix.device)
		87	+ _check_npu_tensor('col_scale', col_scale, dtypes=fp32, device=matrix.device)
		88	+
		89	+ if matrix.ndim != 2:
		90	+ raise ValueError('matrix must have shape (num_rows, head_dim).')
		91	+ num_rows, head_dim = matrix.shape
		92	+ if head_dim == 0 or head_dim % head_dim_alignment:
		93	+ raise ValueError(f'head_dim must be a positive multiple of {head_dim_alignment}, got {head_dim}.')
		94	+ if row_scale.ndim != 1 or row_scale.shape[0] != num_rows:
		95	+ raise ValueError(f'row_scale must have shape ({num_rows},), got {tuple(row_scale.shape)}.')
		96	+ if col_scale.ndim != 1 or col_scale.shape[0] != head_dim:
		97	+ raise ValueError(f'col_scale must have shape ({head_dim},), got {tuple(col_scale.shape)}.')
		98	+
		99	+ _reject_aliases({'matrix': matrix, 'row_scale': row_scale, 'col_scale': col_scale})
		100	+ return num_rows, head_dim, eps
		101	+
		102	+
		103	+def engram_sinkhorn_momentum_update(
		104	+ grad: torch.Tensor,
		105	+ momentum_buffer: torch.Tensor,
		106	+ beta1: float,
		107	+ nesterov: bool = True,
		108	+) -> torch.Tensor:
		109	+ """Update Engram momentum in place and return the FP32 Sinkhorn input.
		110	+
		111	+ ``grad`` and ``momentum_buffer`` must be contiguous one-dimensional NPU
		112	+ tensors with equal shape. Each tensor may use BF16 or FP32, independently,
		113	+ and its length must be a multiple of 256. Computation is performed in FP32;
		114	+ the updated momentum is converted back to ``momentum_buffer.dtype``.
		115	+ """
		116	+
		117	+ if not isinstance(grad, torch.Tensor) or not isinstance(momentum_buffer, torch.Tensor):
		118	+ raise TypeError('grad and momentum_buffer must be torch.Tensor instances.')
		119	+ beta1 = _check_real('beta1', beta1)
		120	+ if not isinstance(nesterov, bool):
		121	+ raise TypeError('nesterov must be a bool.')
		122	+
		123	+ require_ascend()
		124	+ supported_dtypes = (torch.bfloat16, torch.float32)
		125	+ _check_npu_tensor('grad', grad, dtypes=supported_dtypes)
		126	+ _check_npu_tensor('momentum_buffer', momentum_buffer, dtypes=supported_dtypes, device=grad.device)
		127	+ if grad.ndim != 1 or momentum_buffer.ndim != 1:
		128	+ raise ValueError('grad and momentum_buffer must be one-dimensional.')
		129	+ if grad.shape != momentum_buffer.shape:
		130	+ raise ValueError(
		131	+ f'grad and momentum_buffer must have the same shape, got {tuple(grad.shape)} and '
		132	+ f'{tuple(momentum_buffer.shape)}.'
		133	+ )
		134	+ if grad.numel() % 256:
		135	+ raise ValueError(f'grad length must be a multiple of 256, got {grad.numel()}.')
		136	+ _reject_aliases({'grad': grad, 'momentum_buffer': momentum_buffer})
		137	+
		138	+ sinkhorn_input = torch.empty_like(momentum_buffer, dtype=torch.float32)
		139	+ if grad.numel() == 0:
		140	+ return sinkhorn_input
		141	+
		142	+ num_tiles = (grad.numel() + _MOMENTUM_TILE_SIZE - 1) // _MOMENTUM_TILE_SIZE
		143	+ num_cores = min(get_num_vector_cores(), num_tiles)
		144	+ kernel = get_engram_sinkhorn_momentum_update_kernel(
		145	+ T.dtype(grad.dtype),
		146	+ T.dtype(momentum_buffer.dtype),
		147	+ nesterov,
		148	+ beta1,
		149	+ num_cores,
		150	+ )
		151	+ if should_print_kernel_source():
		152	+ print(kernel.get_kernel_source())
		153	+ kernel(grad, momentum_buffer, sinkhorn_input)
		154	+ return sinkhorn_input
79		155
80		156
81	def engram_sinkhorn_step(	157	def engram_sinkhorn_step(
	@@ -86,18 +162,20 @@ def engram_sinkhorn_step(
86	) -> tuple[torch.Tensor, torch.Tensor]:	162	) -> tuple[torch.Tensor, torch.Tensor]:
87	"""Normalize one Engram Sinkhorn iterate over rows.	163	"""Normalize one Engram Sinkhorn iterate over rows.
88		164
89	- ``matrix`` has shape ``(num_rows, head_dim)`` and both scale tensors are	165	+ All inputs must be contiguous FP32 NPU tensors and ``head_dim`` must be a
90	- one-dimensional. All tensors must be contiguous FP32 NPU tensors;	166	+ positive multiple of 256. The function updates ``row_scale`` in place and
91	- ``head_dim`` must be a positive multiple of 256 and ``row_scale`` must	167	+ returns the pre-normalization row norms and normalized column squared sums.
92	- contain non-negative Sinkhorn scaling factors. The function updates	168	+ The caller must provide non-negative Sinkhorn scaling factors in
93	- ``row_scale`` in place and returns the pre-normalization ``row_norm`` and	169	+ ``row_scale``.
94	- normalized per-column squared sums ``col_sumsq``.
95	-
96	- The three inputs must not share storage. ``col_sumsq`` is initialized to
97	- zero before the cross-core atomic reduction.
98	"""	170	"""
99		171
100	- num_rows, head_dim = _validate_inputs(matrix, row_scale, col_scale, eps)	172	+ num_rows, head_dim, eps = _validate_matrix_inputs(
		173	+ matrix,
		174	+ row_scale,
		175	+ col_scale,
		176	+ eps,
		177	+ head_dim_alignment=256,
		178	+ )
101	if num_rows == 0:	179	if num_rows == 0:
102	return torch.empty_like(row_scale), torch.zeros_like(col_scale)	180	return torch.empty_like(row_scale), torch.zeros_like(col_scale)
103		181
	@@ -115,7 +193,7 @@ def engram_sinkhorn_step(
115	+ 3 * rows_per_core * row_scale.element_size()	193	+ 3 * rows_per_core * row_scale.element_size()
116	+ (row_bytes if head_dim > 512 else 0)	194	+ (row_bytes if head_dim > 512 else 0)
117	)	195	)
118	- if shared_ub_bytes > _SHARED_UB_CAPACITY_BYTES:	196	+ if shared_ub_bytes > _SIMT_SHARED_UB_CAPACITY_BYTES:
119	raise ValueError(	197	raise ValueError(
120	'the requested num_rows/head_dim combination exceeds the 216 KiB usable shared UB capacity: '	198	'the requested num_rows/head_dim combination exceeds the 216 KiB usable shared UB capacity: '
121	f'num_rows={num_rows}, head_dim={head_dim}, required_bytes={shared_ub_bytes}.'	199	f'num_rows={num_rows}, head_dim={head_dim}, required_bytes={shared_ub_bytes}.'
	@@ -124,7 +202,7 @@ def engram_sinkhorn_step(
124	col_sumsq = torch.zeros_like(col_scale)	202	col_sumsq = torch.zeros_like(col_scale)
125	kernel = get_engram_sinkhorn_step_kernel(	203	kernel = get_engram_sinkhorn_step_kernel(
126	head_dim,	204	head_dim,
127	- float(eps),	205	+ eps,
128	num_cores,	206	num_cores,
129	rows_per_core,	207	rows_per_core,
130	chunk_rows,	208	chunk_rows,
	@@ -134,3 +212,55 @@ def engram_sinkhorn_step(
134	print(kernel.get_kernel_source())	212	print(kernel.get_kernel_source())
135	kernel(matrix, row_scale, col_scale, row_norm, col_sumsq)	213	kernel(matrix, row_scale, col_scale, row_norm, col_sumsq)
136	return row_norm, col_sumsq	214	return row_norm, col_sumsq
		215	+
		216	+
		217	+def engram_sinkhorn_finalize(
		218	+ matrix: torch.Tensor,
		219	+ row_scale: torch.Tensor,
		220	+ col_scale: torch.Tensor,
		221	+ eps: float,
		222	+ alpha: float,
		223	+ out: torch.Tensor \| None = None,
		224	+) -> torch.Tensor:
		225	+ """Materialize the final row-normalized Engram Sinkhorn matrix.
		226	+
		227	+ The three inputs must be contiguous FP32 NPU tensors and ``head_dim`` must
		228	+ be a positive multiple of 64. ``out`` may be supplied as a separate,
		229	+ contiguous FP32 output tensor with the same shape and device as ``matrix``.
		230	+ """
		231	+
		232	+ num_rows, head_dim, eps = _validate_matrix_inputs(
		233	+ matrix,
		234	+ row_scale,
		235	+ col_scale,
		236	+ eps,
		237	+ head_dim_alignment=64,
		238	+ )
		239	+ alpha = _check_real('alpha', alpha)
		240	+ if out is not None:
		241	+ if not isinstance(out, torch.Tensor):
		242	+ raise TypeError('out must be a torch.Tensor when provided.')
		243	+ _check_npu_tensor('out', out, dtypes=(torch.float32,), device=matrix.device)
		244	+ if out.shape != matrix.shape:
		245	+ raise ValueError(f'out must have shape {tuple(matrix.shape)}, got {tuple(out.shape)}.')
		246	+ _reject_aliases({'matrix': matrix, 'row_scale': row_scale, 'col_scale': col_scale, 'out': out})
		247	+ else:
		248	+ out = torch.empty_like(matrix)
		249	+
		250	+ if num_rows == 0:
		251	+ return out
		252	+
		253	+ num_stages = 6 if head_dim <= 512 else 3
		254	+ row_bytes = head_dim * matrix.element_size()
		255	+ shared_ub_bytes = (2 * num_stages + 1) * row_bytes + num_stages * row_scale.element_size()
		256	+ if shared_ub_bytes > _PHYSICAL_UB_CAPACITY_BYTES:
		257	+ raise ValueError(
		258	+ 'the requested head_dim exceeds the 248 KiB shared UB capacity: '
		259	+ f'head_dim={head_dim}, required_bytes={shared_ub_bytes}.'
		260	+ )
		261	+ num_cores = min(get_num_vector_cores(), num_rows)
		262	+ kernel = get_engram_sinkhorn_finalize_kernel(head_dim, eps, alpha, num_cores)
		263	+ if should_print_kernel_source():
		264	+ print(kernel.get_kernel_source())
		265	+ kernel(matrix, row_scale, col_scale, out)
		266	+ return out

Rsrc/cann_ops_tilelang/engram/engram_sinkhorn_step_kernel.py→src/cann_ops_tilelang/engram/engram_sinkhorn_kernel.py+191-1

	@@ -1,10 +1,117 @@
1	-"""TileLang Ascend implementation of one Engram Sinkhorn step."""	1	+"""TileLang Ascend kernels for the Engram Sinkhorn operator family."""
2		2
3	import tilelang	3	import tilelang
4	from tilelang import language as T	4	from tilelang import language as T
5	from tilelang.language import simd as S	5	from tilelang.language import simd as S
6		6
7		7
		8	+@tilelang.jit
		9	+def get_engram_sinkhorn_momentum_update_kernel(
		10	+ grad_dtype: T.dtype,
		11	+ momentum_dtype: T.dtype,
		12	+ nesterov: bool,
		13	+ beta1: float,
		14	+ num_cores: int,
		15	+):
		16	+ """Build the vectorized Ascend Sinkhorn momentum-update kernel."""
		17	+
		18	+ tile_size = 4096
		19	+ num_stages = 2
		20	+ vector_size = 64
		21	+ grad_is_bf16 = grad_dtype == T.bfloat16
		22	+ momentum_is_bf16 = momentum_dtype == T.bfloat16
		23	+ numel = T.dynamic('numel')
		24	+
		25	+ @T.prim_func
		26	+ def engram_sinkhorn_momentum_update_kernel(
		27	+ grad: T.Tensor[(numel,), grad_dtype],
		28	+ momentum_buffer: T.Tensor[(numel,), momentum_dtype],
		29	+ sinkhorn_input: T.Tensor[(numel,), T.float32],
		30	+ ):
		31	+ total_tiles = T.ceildiv(numel, tile_size)
		32	+
		33	+ with T.Kernel(num_cores) as core_id:
		34	+ grad_ub = T.alloc_shared((tile_size,), grad_dtype)
		35	+ momentum_in_ub = T.alloc_shared((tile_size,), momentum_dtype)
		36	+ momentum_out_ub = T.alloc_shared((tile_size,), momentum_dtype)
		37	+ sinkhorn_ub = T.alloc_shared((tile_size,), T.float32)
		38	+ T.annotate_buffer_versions(
		39	+ {
		40	+ grad_ub: num_stages,
		41	+ momentum_in_ub: num_stages,
		42	+ momentum_out_ub: num_stages,
		43	+ sinkhorn_ub: num_stages,
		44	+ }
		45	+ )
		46	+
		47	+ for tile_id in T.Persistent([total_tiles], num_cores, core_id, num_stages=num_stages):
		48	+ base = tile_id * tile_size
		49	+ valid_elements = T.min(tile_size, numel - base)
		50	+ T.copy(grad[base : base + valid_elements], grad_ub[:valid_elements])
		51	+ T.copy(
		52	+ momentum_buffer[base : base + valid_elements],
		53	+ momentum_in_ub[:valid_elements],
		54	+ )
		55	+
		56	+ with T.SimdVF():
		57	+ full = S.pset(32, 'PAT_ALL')
		58	+ for vector in range(tile_size // vector_size):
		59	+ offset = vector * vector_size
		60	+ if grad_is_bf16:
		61	+ grad_vec = S.vcvt(
		62	+ S.vld(grad_ub[offset], dist='UNPK_B16'),
		63	+ T.float32,
		64	+ part=0,
		65	+ )
		66	+ else:
		67	+ grad_vec = S.vld(grad_ub[offset])
		68	+
		69	+ if momentum_is_bf16:
		70	+ momentum_vec = S.vcvt(
		71	+ S.vld(momentum_in_ub[offset], dist='UNPK_B16'),
		72	+ T.float32,
		73	+ part=0,
		74	+ )
		75	+ else:
		76	+ momentum_vec = S.vld(momentum_in_ub[offset])
		77	+
		78	+ updated_momentum = S.vadd(
		79	+ S.vmuls(momentum_vec, beta1, full),
		80	+ S.vmuls(grad_vec, 1.0 - beta1, full),
		81	+ full,
		82	+ )
		83	+ if nesterov:
		84	+ sinkhorn_vec = S.vadd(
		85	+ S.vmuls(updated_momentum, beta1, full),
		86	+ S.vmuls(grad_vec, 1.0 - beta1, full),
		87	+ full,
		88	+ )
		89	+ else:
		90	+ sinkhorn_vec = updated_momentum
		91	+
		92	+ if momentum_is_bf16:
		93	+ S.vsts(
		94	+ momentum_out_ub[offset],
		95	+ S.vcvt(updated_momentum, T.bfloat16),
		96	+ full,
		97	+ dist='PK_B32',
		98	+ )
		99	+ else:
		100	+ S.vsts(momentum_out_ub[offset], updated_momentum, full)
		101	+ S.vsts(sinkhorn_ub[offset], sinkhorn_vec, full)
		102	+
		103	+ T.copy(
		104	+ momentum_out_ub[:valid_elements],
		105	+ momentum_buffer[base : base + valid_elements],
		106	+ )
		107	+ T.copy(
		108	+ sinkhorn_ub[:valid_elements],
		109	+ sinkhorn_input[base : base + valid_elements],
		110	+ )
		111	+
		112	+ return engram_sinkhorn_momentum_update_kernel
		113	+
		114	+
8	@tilelang.jit(	115	@tilelang.jit(
9	pass_configs={	116	pass_configs={
10	tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: False,	117	tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: False,
	@@ -156,3 +263,86 @@ def get_engram_sinkhorn_step_kernel(
156	T.atomic_add(col_sumsq[index], col_acc_ub[index])	263	T.atomic_add(col_sumsq[index], col_acc_ub[index])
157		264
158	return engram_sinkhorn_step_kernel	265	return engram_sinkhorn_step_kernel
		266	+
		267	+
		268	+@tilelang.jit
		269	+def get_engram_sinkhorn_finalize_kernel(
		270	+ head_dim: int,
		271	+ eps: float,
		272	+ alpha: float,
		273	+ num_cores: int,
		274	+):
		275	+ """Build the vectorized Ascend Sinkhorn finalize kernel."""
		276	+
		277	+ vector_size = 64
		278	+ assert head_dim % vector_size == 0
		279	+ num_vectors = head_dim // vector_size
		280	+ num_stages = 6 if head_dim <= 512 else 3
		281	+ num_rows = T.dynamic('num_rows')
		282	+
		283	+ @T.prim_func
		284	+ def engram_sinkhorn_finalize_kernel(
		285	+ matrix: T.Tensor[(num_rows, head_dim), T.float32],
		286	+ row_scale: T.Tensor[(num_rows,), T.float32],
		287	+ col_scale: T.Tensor[(head_dim,), T.float32],
		288	+ out: T.Tensor[(num_rows, head_dim), T.float32],
		289	+ ):
		290	+ with T.Kernel(num_cores) as core_id:
		291	+ x_ub = T.alloc_shared((head_dim,), T.float32)
		292	+ out_ub = T.alloc_shared((head_dim,), T.float32)
		293	+ col_scale_ub = T.alloc_shared((head_dim,), T.float32)
		294	+ row_scale_ub = T.alloc_shared((1,), T.float32)
		295	+ T.annotate_buffer_versions(
		296	+ {
		297	+ x_ub: num_stages,
		298	+ out_ub: num_stages,
		299	+ row_scale_ub: num_stages,
		300	+ }
		301	+ )
		302	+
		303	+ T.copy(col_scale[:head_dim], col_scale_ub[:head_dim])
		304	+ rows_per_core = T.ceildiv(num_rows, num_cores)
		305	+ row_begin = core_id * rows_per_core
		306	+ row_end = T.min(row_begin + rows_per_core, num_rows)
		307	+ iterations = row_end - row_begin
		308	+ for row_iter in T.Pipelined(
		309	+ iterations,
		310	+ num_stages=num_stages,
		311	+ annotations={'multi_buffer_eligible': [x_ub, out_ub, row_scale_ub]},
		312	+ ):
		313	+ row_id = row_begin + row_iter
		314	+ T.copy(matrix[row_id, :head_dim], x_ub[:head_dim])
		315	+ row_scale_ub[0] = row_scale[row_id]
		316	+
		317	+ with T.SimdVF():
		318	+ full = S.pset(32, 'PAT_ALL')
		319	+ zero = S.vdup(0.0, T.float32, full)
		320	+ alpha_vec = S.vdup(alpha, T.float32, full)
		321	+ scale = S.vld(row_scale_ub[0], dist='BRC_B32')
		322	+ sum_x2 = S.alloc_var(T.float32)
		323	+ sum_x2 = zero
		324	+
		325	+ for vector in range(num_vectors):
		326	+ offset = vector * vector_size
		327	+ x_vec = S.vmul(
		328	+ S.vmul(S.vld(x_ub[offset]), scale, full),
		329	+ S.vld(col_scale_ub[offset]),
		330	+ full,
		331	+ )
		332	+ S.vmula(sum_x2, x_vec, x_vec, full)
		333	+
		334	+ norm = S.vsqrt(S.vdupv(S.vcadd(sum_x2, full), full))
		335	+ output_scale = S.vdiv(alpha_vec, S.vadds(norm, eps))
		336	+ final_scale = S.vmul(scale, output_scale, full)
		337	+ for vector in range(num_vectors):
		338	+ offset = vector * vector_size
		339	+ out_vec = S.vmul(
		340	+ S.vmul(S.vld(x_ub[offset]), final_scale, full),
		341	+ S.vld(col_scale_ub[offset]),
		342	+ full,
		343	+ )
		344	+ S.vsts(out_ub[offset], out_vec, full)
		345	+
		346	+ T.copy(out_ub[:head_dim], out[row_id, :head_dim])
		347	+
		348	+ return engram_sinkhorn_finalize_kernel

Dtests/engram/benchmark/test_engram_grad_w_reduce.jsonl+0-6

	@@ -1,6 +0,0 @@
1	-{"kernel": "engram_grad_w_reduce", "operation": "fwd", "params": {"hidden": 2048, "num_partials": 30}, "param_id": "hidden=2048,num_partials=30,", "cluster": "a5", "time_us": 2.15, "bandwidth_gbs": 503.2415}
2	-{"kernel": "engram_grad_w_reduce", "operation": "fwd", "params": {"hidden": 2560, "num_partials": 30}, "param_id": "hidden=2560,num_partials=30,", "cluster": "a5", "time_us": 2.33, "bandwidth_gbs": 581.0937}
3	-{"kernel": "engram_grad_w_reduce", "operation": "fwd", "params": {"hidden": 3072, "num_partials": 30}, "param_id": "hidden=3072,num_partials=30,", "cluster": "a5", "time_us": 2.33, "bandwidth_gbs": 695.5042}
4	-{"kernel": "engram_grad_w_reduce", "operation": "fwd", "params": {"hidden": 4096, "num_partials": 30}, "param_id": "hidden=4096,num_partials=30,", "cluster": "a5", "time_us": 2.58, "bandwidth_gbs": 836.9596}
5	-{"kernel": "engram_grad_w_reduce", "operation": "fwd", "params": {"hidden": 6144, "num_partials": 30}, "param_id": "hidden=6144,num_partials=30,", "cluster": "a5", "time_us": 2.83, "bandwidth_gbs": 1148.1818}
6	-{"kernel": "engram_grad_w_reduce", "operation": "fwd", "params": {"hidden": 7168, "num_partials": 30}, "param_id": "hidden=7168,num_partials=30,", "cluster": "a5", "time_us": 2.87, "bandwidth_gbs": 1319.1494}

Mtests/engram/engram_gate_ref.py+31-22

	@@ -1,7 +1,3 @@
1	-"""Pure PyTorch reference for the Engram gate operator."""
2	-
3	-from __future__ import annotations
4	-
5	import torch	1	import torch
6		2
7		3
	@@ -14,34 +10,47 @@ def engram_gate_ref(
14	eps: float,	10	eps: float,
15	save_for_backward: bool = False,	11	save_for_backward: bool = False,
16	) -> torch.Tensor \| tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:	12	) -> torch.Tensor \| tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
17	- """Compute the Engram gate with eager PyTorch operations.	13	+ """Pure PyTorch reference implementation of engram gate (vectorized, supports autograd).
18		14
19	- This implementation is intentionally device-agnostic and supports	15	+ Computes: output = x + sigmoid(signed_sqrt(dot(RMSNorm(x, wh), RMSNorm(k, we)) * scalar)) * v
20	- autograd, making it the numerical oracle for both forward and backward	16	+
21	- TileLang kernels.	17	+ Args:
		18	+ hidden_states: Input of shape (num_tokens, hc_mult, hidden_size), bfloat16.
		19	+ kv: Packed key/value embeddings of shape (num_tokens, (hc_mult + 1) * hidden_size), bfloat16.
		20	+ weight_hidden: RMSNorm weight for hidden states, shape (hc_mult, hidden_size), bfloat16.
		21	+ weight_embed: RMSNorm weight for key embeddings, shape (hc_mult, hidden_size), bfloat16.
		22	+ clamp_value: Clamp threshold for signed-sqrt gate activation.
		23	+ eps: Epsilon for RMSNorm numerical stability.
		24	+ save_for_backward: If True, also return (dot, gate_score, rstd_x, rstd_k).
		25	+
		26	+ Returns:
		27	+ If save_for_backward is False: output tensor of shape (num_tokens, hc_mult, hidden_size), bfloat16.
		28	+ If save_for_backward is True: tuple of (output, dot, gate_score, rstd_x, rstd_k).
22	"""	29	"""
23	-
24	num_tokens, hc_mult, hidden_size = hidden_states.shape	30	num_tokens, hc_mult, hidden_size = hidden_states.shape
25	scalar = hidden_size**-0.5	31	scalar = hidden_size**-0.5
		32	+
26	k, v = kv.split([hc_mult * hidden_size, hidden_size], dim=-1)	33	k, v = kv.split([hc_mult * hidden_size, hidden_size], dim=-1)
27	k = k.view(num_tokens, hc_mult, hidden_size)	34	k = k.view(num_tokens, hc_mult, hidden_size)
28		35
29	x = hidden_states.float()	36	x = hidden_states.float()
30	- k_float = k.float()	37	+ k_f = k.float()
31	- weight_hidden_float = weight_hidden.float().unsqueeze(0)	38	+ wh = weight_hidden.float().unsqueeze(0)
32	- weight_embed_float = weight_embed.float().unsqueeze(0)	39	+ we = weight_embed.float().unsqueeze(0)
33		40
34	- rstd_x = torch.rsqrt(x.square().mean(-1) + eps)	41	+ # RMSNorm
35	- rstd_k = torch.rsqrt(k_float.square().mean(-1) + eps)	42	+ rstd_x = torch.rsqrt(x.pow(2).mean(-1) + eps)
36	- raw_dot = torch.einsum(	43	+ rstd_k = torch.rsqrt(k_f.pow(2).mean(-1) + eps)
37	- '...d,...d->...',	44	+
38	- x * weight_hidden_float,	45	+ # Dot -> sqrt-gate -> sigmoid
39	- k_float * weight_embed_float,	46	+ # raw_dot is the unnormalized sum(x * wh * k * we), matching the kernel's dot_out
40	- )	47	+ raw_dot = torch.einsum('...d,...d->...', x * wh, k_f * we)
41	- normalized_dot = raw_dot * rstd_x * rstd_k * scalar	48	+ dot = raw_dot * rstd_x * rstd_k * scalar
42	- signed_sqrt = normalized_dot.abs().clamp_min(clamp_value).sqrt() * normalized_dot.sign()	49	+ signed_sqrt = dot.abs().clamp_min(clamp_value).sqrt() * dot.sign()
43	gate_score = signed_sqrt.sigmoid()	50	gate_score = signed_sqrt.sigmoid()
44	- output = (x + gate_score.unsqueeze(-1) * v.unsqueeze(-2)).bfloat16()	51	+
		52	+ output = x + gate_score.unsqueeze(-1) * v.unsqueeze(-2)
		53	+ output = output.bfloat16()
45		54
46	if save_for_backward:	55	if save_for_backward:
47	return output, raw_dot, gate_score, rstd_x, rstd_k	56	return output, raw_dot, gate_score, rstd_x, rstd_k

Atests/engram/engram_hash_ref.py+69-0

@@ -0,0 +1,69 @@
	1	+import torch
	2	+
	3	+
	4	+def make_offsets(vocab_sizes: torch.Tensor) -> torch.Tensor:
	5	+ """Compute exclusive prefix-sum offsets from vocab_sizes.
	6	+
	7	+ Args:
	8	+ vocab_sizes: Per-layer per-ngram embedding table sizes of shape
	9	+ (num_ngram_layers, max_ngram_size - 1, num_embed_table_per_ngram), int32.
	10	+
	11	+ Returns:
	12	+ Offsets of shape (num_ngram_layers, (max_ngram_size - 1) * num_embed_table_per_ngram), int32.
	13	+ """
	14	+ num_ngram_layers = vocab_sizes.shape[0]
	15	+ offsets_list = []
	16	+ for layer_idx in range(num_ngram_layers):
	17	+ flat = vocab_sizes[layer_idx].view(-1)
	18	+ prefix = torch.cat(
	19	+ [
	20	+ torch.zeros(1, dtype=torch.int32, device=flat.device),
	21	+ flat[:-1].cumsum(0, dtype=torch.int32),
	22	+ ]
	23	+ )
	24	+ offsets_list.append(prefix)
	25	+ return torch.stack(offsets_list, dim=0)
	26	+
	27	+
	28	+def engram_hash_ref(
	29	+ ngram_token_ids: torch.Tensor,
	30	+ multipliers: torch.Tensor,
	31	+ vocab_sizes: torch.Tensor,
	32	+ offsets: torch.Tensor,
	33	+) -> torch.Tensor:
	34	+ """Pure PyTorch reference implementation of engram hash.
	35	+
	36	+ Args:
	37	+ ngram_token_ids: N-gram token IDs of shape (num_tokens, max_ngram_size), int32.
	38	+ multipliers: Per-layer hash multipliers of shape (num_ngram_layers, max_ngram_size), int64.
	39	+ vocab_sizes: Per-layer per-ngram embedding table sizes of shape
	40	+ (num_ngram_layers, max_ngram_size - 1, num_embed_table_per_ngram), int32.
	41	+ offsets: Per-layer embedding table offsets of shape
	42	+ (num_ngram_layers, (max_ngram_size - 1) * num_embed_table_per_ngram), int32.
	43	+
	44	+ Returns:
	45	+ Embedding indices of shape
	46	+ (num_ngram_layers, num_tokens, (max_ngram_size - 1) * num_embed_table_per_ngram), int32.
	47	+ """
	48	+
	49	+ num_ngram_layers = multipliers.shape[0]
	50	+ max_ngram_size = multipliers.shape[1]
	51	+
	52	+ prod = ngram_token_ids.to(torch.int64).unsqueeze(0) * multipliers.unsqueeze(1)
	53	+
	54	+ ans = [[] for _ in range(num_ngram_layers)]
	55	+ hashes = prod[:, :, 0].clone()
	56	+ for i in range(1, max_ngram_size):
	57	+ hashes.bitwise_xor_(prod[:, :, i])
	58	+ for layer_idx in range(num_ngram_layers):
	59	+ ans[layer_idx].append(
	60	+ (hashes[layer_idx].unsqueeze(-1) % vocab_sizes[layer_idx, i - 1].to(torch.int64).unsqueeze(0)).to(
	61	+ torch.int32
	62	+ )
	63	+ )
	64	+
	65	+ for layer_idx in range(num_ngram_layers):
	66	+ ans[layer_idx] = torch.cat(ans[layer_idx], dim=-1)
	67	+
	68	+ output = torch.stack(ans, dim=0)
	69	+ return output + offsets.unsqueeze(1)

Rtests/engram/engram_sinkhorn_step_ref.py→tests/engram/engram_sinkhorn_ref.py+38-1

@@ @@ -1,10 +1,27 @@ @@
++"""Pure PyTorch references for the Engram Sinkhorn operator family."""
 from __future__ import annotations
 import torch
++def engram_sinkhorn_momentum_update_ref(
++    grad: torch.Tensor,
++    momentum_buffer: torch.Tensor,
++    beta1: float,
++    nesterov: bool = True,
++) -> torch.Tensor:
++    """Update momentum in place and return the FP32 Sinkhorn input."""
++
++    grad_fp32 = grad.float()
++    momentum_fp32 = momentum_buffer.float()
++    momentum_fp32.mul_(beta1).add_(grad_fp32, alpha=1 - beta1)
++    momentum_buffer.copy_(momentum_fp32)
++    if nesterov:
++        return torch.lerp(momentum_fp32, grad_fp32, 1 - beta1)
++    return momentum_fp32
++
++
 def engram_sinkhorn_step_ref(
     matrix: torch.Tensor,
     row_scale: torch.Tensor,
     row_scale.div_(row_norm + eps)
     normalized = row_scale.unsqueeze(1) * matrix * col_scale
     return row_norm, normalized.square().sum(dim=0)
++
++
++def engram_sinkhorn_finalize_ref(
++    matrix: torch.Tensor,
++    row_scale: torch.Tensor,
++    col_scale: torch.Tensor,
++    eps: float,
++    alpha: float,
++    out: torch.Tensor | None = None,
++) -> torch.Tensor:
++    """Materialize the final row-normalized Sinkhorn matrix."""
++
++    iterate = row_scale.unsqueeze(1) * matrix * col_scale
++    row_norm = torch.linalg.vector_norm(iterate, dim=1)
++    normalized_row_scale = row_scale / (row_norm + eps)
++    result = alpha * normalized_row_scale.unsqueeze(1) * matrix * col_scale
++    if out is None:
++        return result
++    out.copy_(result)
++    return out

Dtests/engram/npu_mem_profiles/test_engram_grad_w_reduce.json+0-20

	@@ -1,20 +0,0 @@
1	-{
2	- "engram/test_engram_grad_w_reduce.py::test_engram_grad_w_reduce[hidden=2048,num_partials=30,]": {
3	- "peak_npu_mb": 33.0
4	- },
5	- "engram/test_engram_grad_w_reduce.py::test_engram_grad_w_reduce[hidden=2560,num_partials=30,]": {
6	- "peak_npu_mb": 33.3
7	- },
8	- "engram/test_engram_grad_w_reduce.py::test_engram_grad_w_reduce[hidden=3072,num_partials=30,]": {
9	- "peak_npu_mb": 33.5
10	- },
11	- "engram/test_engram_grad_w_reduce.py::test_engram_grad_w_reduce[hidden=4096,num_partials=30,]": {
12	- "peak_npu_mb": 34.1
13	- },
14	- "engram/test_engram_grad_w_reduce.py::test_engram_grad_w_reduce[hidden=6144,num_partials=30,]": {
15	- "peak_npu_mb": 35.1
16	- },
17	- "engram/test_engram_grad_w_reduce.py::test_engram_grad_w_reduce[hidden=7168,num_partials=30,]": {
18	- "peak_npu_mb": 35.6
19	- }
20	-}

Mtests/engram/test_engram_fused_weight.py+0-2

	@@ -22,7 +22,6 @@ def generate_test_data(params):
22		22
23		23
24	def generate_test_params(level: int) -> list[dict]:	24	def generate_test_params(level: int) -> list[dict]:
25	- del level
26	return [{'hc': hc, 'hidden': hidden_size} for hc in (4,) for hidden_size in generate_hidden_sizes(256)]	25	return [{'hc': hc, 'hidden': hidden_size} for hc in (4,) for hidden_size in generate_hidden_sizes(256)]
27		26
28		27
	@@ -33,7 +32,6 @@ def test_engram_fused_weight(params):
33	ref = wh_data.float() * we_data.float()	32	ref = wh_data.float() * we_data.float()
34	out = cann_ops_tilelang.engram_fused_weight(wh_data, we_data)	33	out = cann_ops_tilelang.engram_fused_weight(wh_data, we_data)
35		34
36	- assert out.dtype == torch.float32
37	assert_equal(out, ref)	35	assert_equal(out, ref)
38		36
39		37

Atests/engram/test_engram_gate_bwd.py+150-0

@@ -0,0 +1,150 @@
	1	+import os
	2	+
	3	+import pytest
	4	+import torch
	5	+
	6	+import cann_ops_tilelang
	7	+from cann_ops_tilelang.config import get_device
	8	+from tests.engram.engram_gate_ref import engram_gate_ref
	9	+from tests.testing.bench import make_param_id
	10	+from tests.testing.generator import generate_hidden_sizes, generate_num_tokens, get_test_level
	11	+from tests.testing.numeric import calc_diff, count_bytes
	12	+
	13	+os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
	14	+pytestmark = pytest.mark.npu
	15	+
	16	+
	17	+def generate_test_data(params):
	18	+ if get_device().type == 'npu':
	19	+ torch.manual_seed(0)
	20	+ num_tokens = params['num_tokens']
	21	+ hc_mult = params['hc']
	22	+ hidden_size = params['hidden']
	23	+ eps = 1e-20
	24	+ clamp_value = 1e-6
	25	+ x_data = torch.randn(num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=get_device())
	26	+ kv_data = torch.randn(num_tokens, (hc_mult + 1) * hidden_size, dtype=torch.bfloat16, device=get_device())
	27	+ wh_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=get_device())
	28	+ we_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=get_device())
	29	+ weight_fused = wh_data.float() * we_data.float()
	30	+ grad_out = torch.randn(num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=get_device())
	31	+ return x_data, kv_data, wh_data, we_data, weight_fused, grad_out, eps, clamp_value
	32	+
	33	+
	34	+def generate_test_params(level: int) -> list[dict]:
	35	+ return [
	36	+ {'num_tokens': t, 'hc': hc, 'hidden': hidden_size}
	37	+ for t in generate_num_tokens(level)
	38	+ for hc in (4,)
	39	+ for hidden_size in generate_hidden_sizes(256)
	40	+ ]
	41	+
	42	+
	43	+@pytest.mark.parametrize('params', generate_test_params(get_test_level()), ids=make_param_id)
	44	+def test_engram_gate_bwd(params):
	45	+ x_data, kv_data, wh_data, we_data, weight_fused, grad_out, eps, clamp_value = generate_test_data(params)
	46	+
	47	+ x_ref = x_data.clone().requires_grad_(True)
	48	+ kv_ref = kv_data.clone().requires_grad_(True)
	49	+ wh_ref = wh_data.float().requires_grad_(True)
	50	+ we_ref = we_data.float().requires_grad_(True)
	51	+ output_ref, dot_ref, gate_score_ref, rstd_x_ref, rstd_k_ref = engram_gate_ref(
	52	+ x_ref,
	53	+ kv_ref,
	54	+ wh_ref,
	55	+ we_ref,
	56	+ clamp_value,
	57	+ eps,
	58	+ save_for_backward=True,
	59	+ )
	60	+ output_ref.backward(grad_out)
	61	+
	62	+ grad_x, grad_kv, grad_w_partial = cann_ops_tilelang.engram_gate_bwd(
	63	+ grad_out,
	64	+ x_data,
	65	+ kv_data,
	66	+ weight_fused,
	67	+ dot_ref,
	68	+ gate_score_ref,
	69	+ rstd_x_ref,
	70	+ rstd_k_ref,
	71	+ clamp_value,
	72	+ )
	73	+ grad_w_fused = grad_w_partial.sum(0)
	74	+ grad_wh = grad_w_fused * we_data.float()
	75	+ grad_we = grad_w_fused * wh_data.float()
	76	+
	77	+ diff_x = calc_diff(grad_x, x_ref.grad)
	78	+ assert diff_x < 1e-8, f'grad_x mismatch: {diff_x:.6e}'
	79	+ diff_kv = calc_diff(grad_kv, kv_ref.grad)
	80	+ assert diff_kv < 1e-8, f'grad_kv mismatch: {diff_kv:.6e}'
	81	+ diff_wh = calc_diff(grad_wh, wh_ref.grad)
	82	+ assert diff_wh < 1e-8, f'grad_wh mismatch: {diff_wh:.6e}'
	83	+ diff_we = calc_diff(grad_we, we_ref.grad)
	84	+ assert diff_we < 1e-8, f'grad_we mismatch: {diff_we:.6e}'
	85	+
	86	+
	87	+@pytest.mark.benchmark
	88	+@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)
	89	+def test_engram_gate_bwd_benchmark(benchmark_timer, benchmark_record, params):
	90	+ x_data, kv_data, wh_data, we_data, weight_fused, grad_out, eps, clamp_value = generate_test_data(params)
	91	+
	92	+ _, dot_ref, gate_score_ref, rstd_x_ref, rstd_k_ref = engram_gate_ref(
	93	+ x_data,
	94	+ kv_data,
	95	+ wh_data,
	96	+ we_data,
	97	+ clamp_value,
	98	+ eps,
	99	+ save_for_backward=True,
	100	+ )
	101	+ grad_x, grad_kv, grad_w_partial = cann_ops_tilelang.engram_gate_bwd(
	102	+ grad_out,
	103	+ x_data,
	104	+ kv_data,
	105	+ weight_fused,
	106	+ dot_ref,
	107	+ gate_score_ref,
	108	+ rstd_x_ref,
	109	+ rstd_k_ref,
	110	+ clamp_value,
	111	+ )
	112	+ grad_w_fused = grad_w_partial.sum(0)
	113	+ grad_wh = grad_w_fused * we_data.float()
	114	+ grad_we = grad_w_fused * wh_data.float()
	115	+
	116	+ def run_backward():
	117	+ return cann_ops_tilelang.engram_gate_bwd(
	118	+ grad_out,
	119	+ x_data,
	120	+ kv_data,
	121	+ weight_fused,
	122	+ dot_ref,
	123	+ gate_score_ref,
	124	+ rstd_x_ref,
	125	+ rstd_k_ref,
	126	+ clamp_value,
	127	+ )
	128	+
	129	+ t_us = benchmark_timer(run_backward)
	130	+ num_bytes = count_bytes(
	131	+ grad_out,
	132	+ x_data,
	133	+ kv_data,
	134	+ weight_fused,
	135	+ dot_ref,
	136	+ gate_score_ref,
	137	+ rstd_x_ref,
	138	+ rstd_k_ref,
	139	+ grad_x,
	140	+ grad_kv,
	141	+ grad_wh,
	142	+ grad_we,
	143	+ )
	144	+ benchmark_record(
	145	+ kernel='engram_gate_bwd',
	146	+ operation='bwd',
	147	+ params=params,
	148	+ time_us=t_us,
	149	+ bandwidth_gbs=num_bytes / t_us / 1e3,
	150	+ )

Atests/engram/test_engram_gate_fwd.py+147-0

@@ -0,0 +1,147 @@
	1	+import os
	2	+
	3	+import pytest
	4	+import torch
	5	+
	6	+import cann_ops_tilelang
	7	+from cann_ops_tilelang.config import get_device
	8	+from tests.engram.engram_gate_ref import engram_gate_ref
	9	+from tests.testing.bench import make_param_id
	10	+from tests.testing.generator import generate_hidden_sizes, generate_num_tokens, get_test_level
	11	+from tests.testing.numeric import assert_equal, calc_diff, count_bytes
	12	+
	13	+os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
	14	+pytestmark = pytest.mark.npu
	15	+
	16	+
	17	+def generate_test_data(params):
	18	+ if get_device().type == 'npu':
	19	+ torch.manual_seed(0)
	20	+ num_tokens = params['num_tokens']
	21	+ hc_mult = params['hc']
	22	+ hidden_size = params['hidden']
	23	+ eps = 1e-20
	24	+ clamp_value = 1e-6
	25	+ x_data = torch.randn(num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=get_device())
	26	+ kv_data = torch.randn(num_tokens, (hc_mult + 1) * hidden_size, dtype=torch.bfloat16, device=get_device())
	27	+ wh_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=get_device())
	28	+ we_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=get_device())
	29	+ weight_fused = wh_data.float() * we_data.float()
	30	+ return x_data, kv_data, wh_data, we_data, weight_fused, eps, clamp_value
	31	+
	32	+
	33	+def generate_test_params(level: int) -> list[dict]:
	34	+ return [
	35	+ {'num_tokens': t, 'hc': hc, 'hidden': hidden_size}
	36	+ for t in generate_num_tokens(level)
	37	+ for hc in (4,)
	38	+ for hidden_size in generate_hidden_sizes(256)
	39	+ ]
	40	+
	41	+
	42	+@pytest.mark.parametrize('params', generate_test_params(get_test_level()), ids=make_param_id)
	43	+def test_engram_gate_fwd(params):
	44	+ x_data, kv_data, wh_data, we_data, weight_fused, eps, clamp_value = generate_test_data(params)
	45	+
	46	+ out_ref, dot_ref, gate_score_ref, rstd_x_ref, rstd_k_ref = engram_gate_ref(
	47	+ x_data,
	48	+ kv_data,
	49	+ wh_data,
	50	+ we_data,
	51	+ clamp_value,
	52	+ eps,
	53	+ save_for_backward=True,
	54	+ )
	55	+
	56	+ out_save, dot, gate_score, rstd_x, rstd_k = cann_ops_tilelang.engram_gate_fwd(
	57	+ x_data,
	58	+ kv_data,
	59	+ weight_fused,
	60	+ eps,
	61	+ clamp_value,
	62	+ save_for_backward=True,
	63	+ )
	64	+ assert dot is not None and gate_score is not None and rstd_x is not None and rstd_k is not None
	65	+ diff_out = calc_diff(out_save, out_ref)
	66	+ assert diff_out < 2e-10, f'out_save mismatch: {diff_out:.6e}'
	67	+ diff_dot = calc_diff(dot, dot_ref)
	68	+ assert diff_dot < 2e-10, f'dot mismatch: {diff_dot:.6e}'
	69	+ diff_gate = calc_diff(gate_score, gate_score_ref)
	70	+ assert diff_gate < 2e-10, f'gate_score mismatch: {diff_gate:.6e}'
	71	+ diff_rstd_x = calc_diff(rstd_x, rstd_x_ref)
	72	+ assert diff_rstd_x < 2e-10, f'rstd_x mismatch: {diff_rstd_x:.6e}'
	73	+ diff_rstd_k = calc_diff(rstd_k, rstd_k_ref)
	74	+ assert diff_rstd_k < 2e-10, f'rstd_k mismatch: {diff_rstd_k:.6e}'
	75	+
	76	+ out_no_save, dot_n, gate_score_n, rstd_x_n, rstd_k_n = cann_ops_tilelang.engram_gate_fwd(
	77	+ x_data,
	78	+ kv_data,
	79	+ weight_fused,
	80	+ eps,
	81	+ clamp_value,
	82	+ save_for_backward=False,
	83	+ )
	84	+ assert dot_n is None and gate_score_n is None and rstd_x_n is None and rstd_k_n is None
	85	+ diff_out = calc_diff(out_no_save, out_ref)
	86	+ assert diff_out < 2e-10, f'out_no_save mismatch: {diff_out:.6e}'
	87	+ assert_equal(out_no_save, out_save)
	88	+
	89	+
	90	+@pytest.mark.benchmark
	91	+@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)
	92	+def test_engram_gate_fwd_benchmark(benchmark_timer, benchmark_record, params):
	93	+ x_data, kv_data, _, _, weight_fused, eps, clamp_value = generate_test_data(params)
	94	+
	95	+ out_save, dot, gate_score, rstd_x, rstd_k = cann_ops_tilelang.engram_gate_fwd(
	96	+ x_data,
	97	+ kv_data,
	98	+ weight_fused,
	99	+ eps,
	100	+ clamp_value,
	101	+ save_for_backward=True,
	102	+ )
	103	+ t_save_us = benchmark_timer(
	104	+ lambda: cann_ops_tilelang.engram_gate_fwd(
	105	+ x_data,
	106	+ kv_data,
	107	+ weight_fused,
	108	+ eps,
	109	+ clamp_value,
	110	+ save_for_backward=True,
	111	+ )
	112	+ )
	113	+ num_bytes_save = count_bytes(x_data, kv_data, weight_fused, out_save, dot, gate_score, rstd_x, rstd_k)
	114	+ benchmark_record(
	115	+ kernel='engram_gate_fwd',
	116	+ operation='fwd',
	117	+ params={**params, 'save': True},
	118	+ time_us=t_save_us,
	119	+ bandwidth_gbs=num_bytes_save / t_save_us / 1e3,
	120	+ )
	121	+
	122	+ out_no_save = cann_ops_tilelang.engram_gate_fwd(
	123	+ x_data,
	124	+ kv_data,
	125	+ weight_fused,
	126	+ eps,
	127	+ clamp_value,
	128	+ save_for_backward=False,
	129	+ )[0]
	130	+ t_no_save_us = benchmark_timer(
	131	+ lambda: cann_ops_tilelang.engram_gate_fwd(
	132	+ x_data,
	133	+ kv_data,
	134	+ weight_fused,
	135	+ eps,
	136	+ clamp_value,
	137	+ save_for_backward=False,
	138	+ )
	139	+ )
	140	+ num_bytes_no_save = count_bytes(x_data, kv_data, weight_fused, out_no_save)
	141	+ benchmark_record(
	142	+ kernel='engram_gate_fwd',
	143	+ operation='fwd',
	144	+ params={**params, 'save': False},
	145	+ time_us=t_no_save_us,
	146	+ bandwidth_gbs=num_bytes_no_save / t_no_save_us / 1e3,
	147	+ )

Mtests/engram/test_engram_grad_w_reduce.py+64-72

	@@ -3,17 +3,17 @@ import os
3	import pytest	3	import pytest
4	import torch	4	import torch
5		5
6	-import cann_ops_tilelang	6	+from cann_ops_tilelang import engram_grad_w_reduce
7	-from cann_ops_tilelang.config import get_device, get_num_vector_cores	7	+from cann_ops_tilelang.config import clear_runtime_caches, get_device
8	from tests.testing.bench import make_param_id	8	from tests.testing.bench import make_param_id
9	-from tests.testing.generator import generate_hidden_sizes, get_test_level	9	+from tests.testing.generator import generate_hidden_sizes, generate_num_sms, get_test_level
10	from tests.testing.numeric import calc_diff, count_bytes	10	from tests.testing.numeric import calc_diff, count_bytes
11		11
12	os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'	12	os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
13	pytestmark = pytest.mark.npu	13	pytestmark = pytest.mark.npu
14		14
15		15
16	-def engram_grad_w_reduce_ref(	16	+def grad_w_reduce_ref(
17	grad_w_partial,	17	grad_w_partial,
18	weight_hidden,	18	weight_hidden,
19	weight_embed,	19	weight_embed,
	@@ -21,82 +21,78 @@ def engram_grad_w_reduce_ref(
21	grad_weight_embed,	21	grad_weight_embed,
22	):	22	):
23	grad_w_sum = grad_w_partial.sum(0)	23	grad_w_sum = grad_w_partial.sum(0)
24	- grad_weight_hidden.add_(grad_w_sum * weight_embed.float())	24	+ grad_weight_hidden += grad_w_sum * weight_embed.float()
25	- grad_weight_embed.add_(grad_w_sum * weight_hidden.float())	25	+ grad_weight_embed += grad_w_sum * weight_hidden.float()
26	-
27	-
28	-def generate_num_partials(level: int) -> list[int]:
29	- gate_num_partials = (get_num_vector_cores() - 12) // 2
30	- values = [gate_num_partials]
31	- if level >= 1:
32	- values.insert(0, max(gate_num_partials - 9, 1))
33	- if level >= 2:
34	- values.insert(0, 1)
35	- return list(dict.fromkeys(values))
36	-
37	-
38	-def generate_test_params(level: int) -> list[dict]:
39	- return [
40	- {'hidden': hidden_size, 'num_partials': num_partials}
41	- for hidden_size in generate_hidden_sizes(256)
42	- for num_partials in generate_num_partials(level)
43	- ]
44		26
45		27
46	def generate_test_data(params):	28	def generate_test_data(params):
47	- torch.manual_seed(0)
48	hidden_size = params['hidden']	29	hidden_size = params['hidden']
49	- num_partials = params['num_partials']
50	hc_mult = 4	30	hc_mult = 4
		31	+ num_persistent_blocks = params['num_sms']
51	device = get_device()	32	device = get_device()
52	- grad_w_partial = torch.randn(num_partials, hc_mult, hidden_size, dtype=torch.float32, device=device)	33	+ grad_w_partial = torch.randn(
		34	+ num_persistent_blocks,
		35	+ hc_mult,
		36	+ hidden_size,
		37	+ dtype=torch.float32,
		38	+ device=device,
		39	+ )
53	weight_hidden = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=device)	40	weight_hidden = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=device)
54	weight_embed = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=device)	41	weight_embed = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=device)
55	return grad_w_partial, weight_hidden, weight_embed	42	return grad_w_partial, weight_hidden, weight_embed
56		43
57		44
58	-@pytest.mark.parametrize('params', generate_test_params(get_test_level()), ids=make_param_id)	45	+def generate_test_params(level: int) -> list[dict]:
59	-def test_engram_grad_w_reduce(params):	46	+ return [
60	- grad_w_partial, weight_hidden, weight_embed = generate_test_data(params)	47	+ {
61	- output_shape = weight_hidden.shape	48	+ 'hidden': hidden_size,
62	- device = weight_hidden.device	49	+ 'num_sms': num_sms,
63	- grad_weight_hidden_ref = torch.randn(output_shape, dtype=torch.float32, device=device)	50	+ }
64	- grad_weight_embed_ref = torch.randn(output_shape, dtype=torch.float32, device=device)	51	+ for hidden_size in generate_hidden_sizes(256)
65	- grad_weight_hidden = grad_weight_hidden_ref.clone()	52	+ for num_sms in generate_num_sms(level)
66	- grad_weight_embed = grad_weight_embed_ref.clone()	53	+ ]
67		54
68	- engram_grad_w_reduce_ref(	55	+
69	- grad_w_partial,	56	+def _set_num_sms(monkeypatch, num_sms: int) -> None:
70	- weight_hidden,	57	+ monkeypatch.setenv('OPS_TILELANG_NUM_VECTOR_CORES', str(num_sms * 2))
71	- weight_embed,	58	+ clear_runtime_caches()
72	- grad_weight_hidden_ref,	59	+
73	- grad_weight_embed_ref,	60	+
74	- )	61	+@pytest.mark.parametrize('params', generate_test_params(get_test_level()), ids=make_param_id)
75	- cann_ops_tilelang.engram_grad_w_reduce(	62	+def test_engram_grad_w_reduce(params, monkeypatch):
76	- grad_w_partial,	63	+ hidden_size = params['hidden']
77	- weight_hidden,	64	+ device = get_device()
78	- weight_embed,	65	+
79	- grad_weight_hidden,	66	+ _set_num_sms(monkeypatch, params['num_sms'])
80	- grad_weight_embed,	67	+ grad_w_partial, weight_hidden, weight_embed = generate_test_data(params)
81	- )	68	+ hc_mult = grad_w_partial.shape[1]
82	- for name, actual, expected in (	69	+
83	- ('grad_weight_hidden', grad_weight_hidden, grad_weight_hidden_ref),	70	+ grad_wh_ref = torch.randn(hc_mult, hidden_size, dtype=torch.float32, device=device)
84	- ('grad_weight_embed', grad_weight_embed, grad_weight_embed_ref),	71	+ grad_we_ref = torch.randn(hc_mult, hidden_size, dtype=torch.float32, device=device)
85	- ):	72	+ grad_weight_hidden = grad_wh_ref.clone()
86	- diff = calc_diff(actual, expected)	73	+ grad_weight_embed = grad_we_ref.clone()
87	- assert diff < 1e-10, f'{name} mismatch: {diff:.6e}'	74	+ grad_w_reduce_ref(grad_w_partial, weight_hidden, weight_embed, grad_wh_ref, grad_we_ref)
		75	+ engram_grad_w_reduce(grad_w_partial, weight_hidden, weight_embed, grad_weight_hidden, grad_weight_embed)
		76	+ diff_wh = calc_diff(grad_weight_hidden.cpu(), grad_wh_ref.cpu())
		77	+ assert diff_wh < 1e-10, f'grad_wh mismatch: {diff_wh:.6e}'
		78	+ diff_we = calc_diff(grad_weight_embed.cpu(), grad_we_ref.cpu())
		79	+ assert diff_we < 1e-10, f'grad_we mismatch: {diff_we:.6e}'
88		80
89		81
90	@pytest.mark.benchmark	82	@pytest.mark.benchmark
91	@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)	83	@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)
92	-def test_engram_grad_w_reduce_benchmark(benchmark_timer, benchmark_record, params):	84	+def test_engram_grad_w_reduce_benchmark(benchmark_timer, benchmark_record, params, monkeypatch):
		85	+ hidden_size = params['hidden']
		86	+ device = get_device()
		87	+
		88	+ _set_num_sms(monkeypatch, params['num_sms'])
93	grad_w_partial, weight_hidden, weight_embed = generate_test_data(params)	89	grad_w_partial, weight_hidden, weight_embed = generate_test_data(params)
94	- output_shape = weight_hidden.shape	90	+ hc_mult = grad_w_partial.shape[1]
95	- device = weight_hidden.device	91	+ grad_weight_hidden = torch.randn(hc_mult, hidden_size, dtype=torch.float32, device=device)
96	- grad_weight_hidden = torch.randn(output_shape, dtype=torch.float32, device=device)	92	+ grad_weight_embed = torch.randn(hc_mult, hidden_size, dtype=torch.float32, device=device)
97	- grad_weight_embed = torch.randn(output_shape, dtype=torch.float32, device=device)	93	+
98	- time_us = benchmark_timer(	94	+ t_us = benchmark_timer(
99	- lambda: cann_ops_tilelang.engram_grad_w_reduce(	95	+ lambda: engram_grad_w_reduce(
100	grad_w_partial,	96	grad_w_partial,
101	weight_hidden,	97	weight_hidden,
102	weight_embed,	98	weight_embed,
	@@ -104,17 +100,13 @@ def test_engram_grad_w_reduce_benchmark(benchmark_timer, benchmark_record, param
104	grad_weight_embed,	100	grad_weight_embed,
105	)	101	)
106	)	102	)
107	- num_bytes = count_bytes(	103	+
108	- grad_w_partial,	104	+ num_bytes = count_bytes(grad_w_partial, weight_hidden, weight_embed, grad_weight_hidden, grad_weight_embed)
109	- weight_hidden,	105	+ bandwidth_gbs = num_bytes / t_us / 1e3
110	- weight_embed,
111	- grad_weight_hidden,
112	- grad_weight_embed,
113	- )
114	benchmark_record(	106	benchmark_record(
115	kernel='engram_grad_w_reduce',	107	kernel='engram_grad_w_reduce',
116	operation='fwd',	108	operation='fwd',
117	params=params,	109	params=params,
118	- time_us=time_us,	110	+ time_us=t_us,
119	- bandwidth_gbs=num_bytes / time_us / 1e3,	111	+ bandwidth_gbs=bandwidth_gbs,
120	)	112	)

Atests/engram/test_engram_hash.py+105-0

@@ -0,0 +1,105 @@
	1	+import os
	2	+
	3	+import pytest
	4	+import torch
	5	+
	6	+from cann_ops_tilelang import engram_hash
	7	+from cann_ops_tilelang.config import get_device
	8	+from tests.engram.engram_hash_ref import engram_hash_ref, make_offsets
	9	+from tests.testing.bench import make_param_id
	10	+from tests.testing.generator import generate_num_tokens, get_test_level
	11	+from tests.testing.numeric import assert_equal, count_bytes
	12	+
	13	+os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
	14	+pytestmark = pytest.mark.npu
	15	+
	16	+
	17	+def generate_test_data(params):
	18	+ num_tokens = params['num_tokens']
	19	+ max_ngram_size = params['ngram']
	20	+ num_ngram_layers = params['layers']
	21	+ num_embed_table_per_ngram = params['tables']
	22	+ device = get_device()
	23	+ ngram_token_ids = torch.randint(
	24	+ 0,
	25	+ 100000,
	26	+ (num_tokens, max_ngram_size),
	27	+ dtype=torch.int32,
	28	+ device=device,
	29	+ )
	30	+ multipliers = torch.randint(
	31	+ 0,
	32	+ 100000,
	33	+ (num_ngram_layers, max_ngram_size),
	34	+ dtype=torch.int64,
	35	+ device=device,
	36	+ )
	37	+ vocab_sizes = torch.randint(
	38	+ 100000,
	39	+ 1000000,
	40	+ (num_ngram_layers, max_ngram_size - 1, num_embed_table_per_ngram),
	41	+ dtype=torch.int32,
	42	+ device=device,
	43	+ )
	44	+ offsets = make_offsets(vocab_sizes)
	45	+ return ngram_token_ids, multipliers, vocab_sizes, offsets
	46	+
	47	+
	48	+def generate_test_params(level: int) -> list[dict]:
	49	+ return [{'num_tokens': t} for t in generate_num_tokens(level)]
	50	+
	51	+
	52	+@pytest.mark.parametrize('params', generate_test_params(get_test_level()), ids=make_param_id)
	53	+def test_engram_hash(params):
	54	+ num_tokens = params['num_tokens']
	55	+ max_ngram_size = 3
	56	+ num_ngram_layers = 2
	57	+ num_embed_table_per_ngram = 8
	58	+
	59	+ ngram_token_ids, multipliers, vocab_sizes, offsets = generate_test_data(
	60	+ {
	61	+ 'num_tokens': num_tokens,
	62	+ 'ngram': max_ngram_size,
	63	+ 'layers': num_ngram_layers,
	64	+ 'tables': num_embed_table_per_ngram,
	65	+ }
	66	+ )
	67	+
	68	+ output = engram_hash(ngram_token_ids, multipliers, vocab_sizes, offsets)
	69	+ output_ref = engram_hash_ref(ngram_token_ids, multipliers, vocab_sizes, offsets)
	70	+ assert_equal(output, output_ref)
	71	+
	72	+
	73	+@pytest.mark.benchmark
	74	+@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)
	75	+def test_engram_hash_benchmark(benchmark_timer, benchmark_record, params):
	76	+ max_ngram_size = 3
	77	+ num_ngram_layers = 2
	78	+ num_embed_table_per_ngram = 8
	79	+
	80	+ ngram_token_ids, multipliers, vocab_sizes, offsets = generate_test_data(
	81	+ {
	82	+ **params,
	83	+ 'ngram': max_ngram_size,
	84	+ 'layers': num_ngram_layers,
	85	+ 'tables': num_embed_table_per_ngram,
	86	+ }
	87	+ )
	88	+ output = engram_hash(ngram_token_ids, multipliers, vocab_sizes, offsets)
	89	+
	90	+ t_us = benchmark_timer(lambda: engram_hash(ngram_token_ids, multipliers, vocab_sizes, offsets))
	91	+
	92	+ num_bytes = count_bytes(ngram_token_ids, multipliers, vocab_sizes, offsets, output)
	93	+ bandwidth_gbs = num_bytes / t_us / 1e3
	94	+ benchmark_record(
	95	+ kernel='engram_hash',
	96	+ operation='fwd',
	97	+ params={
	98	+ **params,
	99	+ 'ngram': max_ngram_size,
	100	+ 'layers': num_ngram_layers,
	101	+ 'tables': num_embed_table_per_ngram,
	102	+ },
	103	+ time_us=t_us,
	104	+ bandwidth_gbs=bandwidth_gbs,
	105	+ )

Atests/engram/test_engram_sinkhorn.py+174-0

@@ -0,0 +1,174 @@
	1	+import os
	2	+
	3	+import pytest
	4	+import torch
	5	+
	6	+from cann_ops_tilelang import (
	7	+ engram_sinkhorn_finalize,
	8	+ engram_sinkhorn_momentum_update,
	9	+ engram_sinkhorn_step,
	10	+)
	11	+from cann_ops_tilelang.config import get_device
	12	+from tests.engram.engram_sinkhorn_ref import (
	13	+ engram_sinkhorn_finalize_ref,
	14	+ engram_sinkhorn_momentum_update_ref,
	15	+ engram_sinkhorn_step_ref,
	16	+)
	17	+from tests.testing.bench import make_param_id
	18	+from tests.testing.generator import get_test_level
	19	+from tests.testing.numeric import calc_diff, count_bytes
	20	+
	21	+os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
	22	+pytestmark = pytest.mark.npu
	23	+
	24	+_DTYPES = {
	25	+ 'bf16': torch.bfloat16,
	26	+ 'fp32': torch.float32,
	27	+}
	28	+
	29	+
	30	+def generate_momentum_test_data(params):
	31	+ numel = params['numel']
	32	+ grad_dtype = _DTYPES[params['grad_dtype']]
	33	+ momentum_dtype = _DTYPES[params['momentum_dtype']]
	34	+ nesterov = params['nesterov']
	35	+ beta1 = params['beta1']
	36	+
	37	+ grad = torch.randn(numel, dtype=grad_dtype, device=get_device())
	38	+ momentum_buffer = torch.randn(numel, dtype=momentum_dtype, device=get_device())
	39	+ return grad, momentum_buffer, beta1, nesterov
	40	+
	41	+
	42	+def generate_momentum_test_params(level: int) -> list[dict]:
	43	+ sizes = [1001 * 256, 10001 * 256, 100001 * 512, 1000001 * 512]
	44	+ grad_dtypes = ('fp32',) if level == 0 else ('bf16', 'fp32')
	45	+ return [
	46	+ {
	47	+ 'numel': numel,
	48	+ 'grad_dtype': grad_dtype,
	49	+ 'momentum_dtype': momentum_dtype,
	50	+ 'nesterov': nesterov,
	51	+ 'beta1': beta1,
	52	+ }
	53	+ for numel in sizes
	54	+ for grad_dtype in grad_dtypes
	55	+ for momentum_dtype in ('bf16', 'fp32')
	56	+ for nesterov in (False, True)
	57	+ for beta1 in (0.95,)
	58	+ ]
	59	+
	60	+
	61	+@pytest.mark.parametrize('params', generate_momentum_test_params(get_test_level()), ids=make_param_id)
	62	+def test_engram_sinkhorn_momentum_update(params):
	63	+ grad, momentum_buffer, beta1, nesterov = generate_momentum_test_data(params)
	64	+ momentum_buffer_ref = momentum_buffer.clone()
	65	+
	66	+ sinkhorn_input_ref = engram_sinkhorn_momentum_update_ref(grad, momentum_buffer_ref, beta1, nesterov)
	67	+ sinkhorn_input = engram_sinkhorn_momentum_update(grad, momentum_buffer, beta1, nesterov)
	68	+
	69	+ diff_momentum = calc_diff(momentum_buffer, momentum_buffer_ref)
	70	+ diff_sinkhorn = calc_diff(sinkhorn_input, sinkhorn_input_ref)
	71	+ assert diff_momentum < 5e-8, f'momentum_buffer mismatch: {diff_momentum:.6e}'
	72	+ assert diff_sinkhorn < 5e-8, f'sinkhorn_input mismatch: {diff_sinkhorn:.6e}'
	73	+
	74	+
	75	+@pytest.mark.benchmark
	76	+@pytest.mark.parametrize('params', generate_momentum_test_params(0), ids=make_param_id)
	77	+def test_engram_sinkhorn_momentum_update_benchmark(benchmark_timer, benchmark_record, params):
	78	+ grad, momentum_buffer, beta1, nesterov = generate_momentum_test_data(params)
	79	+ sinkhorn_input = engram_sinkhorn_momentum_update(grad, momentum_buffer, beta1, nesterov)
	80	+
	81	+ time_us = benchmark_timer(lambda: engram_sinkhorn_momentum_update(grad, momentum_buffer, beta1, nesterov))
	82	+
	83	+ num_bytes = count_bytes(grad, momentum_buffer, momentum_buffer, sinkhorn_input)
	84	+ bandwidth_gbs = num_bytes / time_us / 1e3
	85	+ benchmark_record(
	86	+ kernel='engram_sinkhorn',
	87	+ operation='momentum_update',
	88	+ params=params,
	89	+ time_us=time_us,
	90	+ bandwidth_gbs=bandwidth_gbs,
	91	+ )
	92	+
	93	+
	94	+def generate_row_step_test_data(params):
	95	+ num_rows = params['num_rows']
	96	+ head_dim = params['head_dim']
	97	+ matrix = torch.randn(num_rows, head_dim, dtype=torch.float32, device=get_device())
	98	+ row_scale = torch.rand(num_rows, dtype=torch.float32, device=get_device())
	99	+ col_scale = torch.rand(head_dim, dtype=torch.float32, device=get_device())
	100	+ return matrix, row_scale, col_scale, params['eps']
	101	+
	102	+
	103	+def generate_row_step_test_params(level: int) -> list[dict]:
	104	+ return [
	105	+ {'num_rows': num_rows, 'head_dim': head_dim, 'eps': eps}
	106	+ for num_rows in [1001, 10001, 100001]
	107	+ for head_dim in [256, 5120]
	108	+ for eps in (1e-20,)
	109	+ ]
	110	+
	111	+
	112	+@pytest.mark.parametrize('params', generate_row_step_test_params(get_test_level()), ids=make_param_id)
	113	+def test_engram_sinkhorn_step(params):
	114	+ matrix, row_scale, col_scale, eps = generate_row_step_test_data(params)
	115	+ row_scale_ref = row_scale.clone()
	116	+
	117	+ row_norm_ref, col_sumsq_ref = engram_sinkhorn_step_ref(matrix, row_scale_ref, col_scale, eps)
	118	+ row_norm, col_sumsq = engram_sinkhorn_step(matrix, row_scale, col_scale, eps)
	119	+
	120	+ assert calc_diff(row_norm, row_norm_ref) < 1e-10
	121	+ assert calc_diff(row_scale, row_scale_ref) < 1e-10
	122	+ assert calc_diff(col_sumsq, col_sumsq_ref) < 1e-10
	123	+
	124	+
	125	+@pytest.mark.parametrize('params', generate_row_step_test_params(get_test_level()), ids=make_param_id)
	126	+def test_engram_sinkhorn_finalize(params):
	127	+ matrix, row_scale, col_scale, eps = generate_row_step_test_data(params)
	128	+ alpha = params['head_dim'] ** 0.5
	129	+
	130	+ out_ref = engram_sinkhorn_finalize_ref(matrix, row_scale, col_scale, eps, alpha)
	131	+ out = engram_sinkhorn_finalize(matrix, row_scale, col_scale, eps, alpha)
	132	+
	133	+ assert calc_diff(out, out_ref) < 1e-10
	134	+
	135	+
	136	+@pytest.mark.benchmark
	137	+@pytest.mark.parametrize('params', generate_row_step_test_params(0), ids=make_param_id)
	138	+def test_engram_sinkhorn_step_benchmark(benchmark_timer, benchmark_record, params):
	139	+ matrix, row_scale, col_scale, eps = generate_row_step_test_data(params)
	140	+
	141	+ row_norm, col_sumsq = engram_sinkhorn_step(matrix, row_scale, col_scale, eps)
	142	+
	143	+ time_us = benchmark_timer(lambda: engram_sinkhorn_step(matrix, row_scale, col_scale, eps))
	144	+
	145	+ num_bytes = count_bytes(matrix, row_scale, col_scale, row_scale, row_norm, col_sumsq)
	146	+ bandwidth_gbs = num_bytes / time_us / 1e3
	147	+ benchmark_record(
	148	+ kernel='engram_sinkhorn',
	149	+ operation='step',
	150	+ params=params,
	151	+ time_us=time_us,
	152	+ bandwidth_gbs=bandwidth_gbs,
	153	+ )
	154	+
	155	+
	156	+@pytest.mark.benchmark
	157	+@pytest.mark.parametrize('params', generate_row_step_test_params(0), ids=make_param_id)
	158	+def test_engram_sinkhorn_finalize_benchmark(benchmark_timer, benchmark_record, params):
	159	+ matrix, row_scale, col_scale, eps = generate_row_step_test_data(params)
	160	+ alpha = params['head_dim'] ** 0.5
	161	+ out = torch.empty_like(matrix)
	162	+ engram_sinkhorn_finalize(matrix, row_scale, col_scale, eps, alpha, out)
	163	+
	164	+ time_us = benchmark_timer(lambda: engram_sinkhorn_finalize(matrix, row_scale, col_scale, eps, alpha, out))
	165	+
	166	+ num_bytes = count_bytes(matrix, row_scale, col_scale, out)
	167	+ bandwidth_gbs = num_bytes / time_us / 1e3
	168	+ benchmark_record(
	169	+ kernel='engram_sinkhorn',
	170	+ operation='finalize',
	171	+ params=params,
	172	+ time_us=time_us,
	173	+ bandwidth_gbs=bandwidth_gbs,
	174	+ )

Dtests/engram/test_engram_sinkhorn_step.py+0-165

	@@ -1,165 +0,0 @@
1	-import os
2	-
3	-import pytest
4	-import torch
5	-
6	-import cann_ops_tilelang
7	-from cann_ops_tilelang.config import get_device
8	-from tests.engram.engram_sinkhorn_step_ref import engram_sinkhorn_step_ref
9	-from tests.testing.bench import make_param_id
10	-from tests.testing.generator import get_test_level
11	-from tests.testing.numeric import calc_diff, count_bytes
12	-
13	-os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
14	-pytestmark = pytest.mark.npu
15	-
16	-
17	-def generate_test_params(level: int) -> list[dict]:
18	- source_cases = [
19	- {'num_rows': num_rows, 'head_dim': head_dim, 'eps': 1e-20}
20	- for num_rows in (1001, 10001, 100001)
21	- for head_dim in (256, 5120)
22	- ]
23	- if level == 0:
24	- return [source_cases[0], source_cases[3]]
25	- if level == 1:
26	- return source_cases
27	- return [
28	- {'num_rows': 0, 'head_dim': 256, 'eps': 1e-20},
29	- {'num_rows': 1, 'head_dim': 256, 'eps': 0.0},
30	- {'num_rows': 17, 'head_dim': 512, 'eps': 1e-6},
31	- *source_cases,
32	- ]
33	-
34	-
35	-def generate_test_data(params):
36	- device = get_device()
37	- matrix = torch.randn(
38	- params['num_rows'],
39	- params['head_dim'],
40	- dtype=torch.float32,
41	- device=device,
42	- )
43	- row_scale = torch.rand(params['num_rows'], dtype=torch.float32, device=device)
44	- col_scale = torch.rand(params['head_dim'], dtype=torch.float32, device=device)
45	- return matrix, row_scale, col_scale
46	-
47	-
48	-@pytest.mark.parametrize('params', generate_test_params(get_test_level()), ids=make_param_id)
49	-def test_engram_sinkhorn_step(params):
50	- matrix, row_scale, col_scale = generate_test_data(params)
51	- row_scale_ref = row_scale.clone()
52	-
53	- row_norm_ref, col_sumsq_ref = engram_sinkhorn_step_ref(
54	- matrix,
55	- row_scale_ref,
56	- col_scale,
57	- params['eps'],
58	- )
59	- row_norm, col_sumsq = cann_ops_tilelang.engram_sinkhorn_step(
60	- matrix,
61	- row_scale,
62	- col_scale,
63	- params['eps'],
64	- )
65	-
66	- assert row_norm.shape == row_scale.shape
67	- assert col_sumsq.shape == col_scale.shape
68	- assert row_norm.dtype == col_sumsq.dtype == torch.float32
69	- assert row_norm.device == col_sumsq.device == matrix.device
70	- assert calc_diff(row_norm, row_norm_ref) < 1e-10
71	- assert calc_diff(row_scale, row_scale_ref) < 1e-10
72	- assert calc_diff(col_sumsq, col_sumsq_ref) < 1e-10
73	-
74	-
75	-def test_engram_sinkhorn_step_empty_rows():
76	- params = {'num_rows': 0, 'head_dim': 256, 'eps': 1e-20}
77	- matrix, row_scale, col_scale = generate_test_data(params)
78	-
79	- row_norm, col_sumsq = cann_ops_tilelang.engram_sinkhorn_step(matrix, row_scale, col_scale, params['eps'])
80	-
81	- assert row_norm.shape == (0,)
82	- assert torch.count_nonzero(col_sumsq).item() == 0
83	-
84	-
85	-def test_engram_sinkhorn_step_rejects_invalid_inputs():
86	- device = get_device()
87	- matrix = torch.randn((17, 256), dtype=torch.float32, device=device)
88	- row_scale = torch.rand((17,), dtype=torch.float32, device=device)
89	- col_scale = torch.rand((256,), dtype=torch.float32, device=device)
90	-
91	- with pytest.raises(TypeError, match='torch.Tensor'):
92	- cann_ops_tilelang.engram_sinkhorn_step(None, row_scale, col_scale, 1e-20)
93	- with pytest.raises(ValueError, match='NPU tensor'):
94	- cann_ops_tilelang.engram_sinkhorn_step(matrix.cpu(), row_scale, col_scale, 1e-20)
95	- with pytest.raises(ValueError, match='NPU tensor'):
96	- cann_ops_tilelang.engram_sinkhorn_step(matrix, row_scale.cpu(), col_scale, 1e-20)
97	- with pytest.raises(TypeError, match='torch.float32'):
98	- cann_ops_tilelang.engram_sinkhorn_step(matrix.bfloat16(), row_scale, col_scale, 1e-20)
99	- with pytest.raises(TypeError, match='torch.float32'):
100	- cann_ops_tilelang.engram_sinkhorn_step(matrix, row_scale, col_scale.bfloat16(), 1e-20)
101	- with pytest.raises(ValueError, match='shape'):
102	- cann_ops_tilelang.engram_sinkhorn_step(matrix.view(-1), row_scale, col_scale, 1e-20)
103	- with pytest.raises(ValueError, match='row_scale'):
104	- cann_ops_tilelang.engram_sinkhorn_step(matrix, row_scale[:-1], col_scale, 1e-20)
105	- with pytest.raises(ValueError, match='col_scale'):
106	- cann_ops_tilelang.engram_sinkhorn_step(matrix, row_scale, col_scale[:-1], 1e-20)
107	- with pytest.raises(ValueError, match='multiple of 256'):
108	- cann_ops_tilelang.engram_sinkhorn_step(matrix[:, :-1].contiguous(), row_scale, col_scale[:-1], 1e-20)
109	- wide_matrix = torch.randn((1, 6144), dtype=torch.float32, device=device)
110	- with pytest.raises(ValueError, match='shared UB capacity'):
111	- cann_ops_tilelang.engram_sinkhorn_step(
112	- wide_matrix,
113	- torch.rand((1,), device=device),
114	- torch.rand((6144,), device=device),
115	- 1e-20,
116	- )
117	- with pytest.raises(ValueError, match='contiguous'):
118	- cann_ops_tilelang.engram_sinkhorn_step(matrix.T, torch.rand(256, device=device), row_scale, 1e-20)
119	- with pytest.raises(ValueError, match='contiguous'):
120	- non_contiguous_row_scale = torch.rand((17, 2), device=device)[:, 0]
121	- cann_ops_tilelang.engram_sinkhorn_step(matrix, non_contiguous_row_scale, col_scale, 1e-20)
122	- with pytest.raises(ValueError, match='contiguous'):
123	- non_contiguous_col_scale = torch.rand((256, 2), device=device)[:, 0]
124	- cann_ops_tilelang.engram_sinkhorn_step(matrix, row_scale, non_contiguous_col_scale, 1e-20)
125	- with pytest.raises(TypeError, match='real number'):
126	- cann_ops_tilelang.engram_sinkhorn_step(matrix, row_scale, col_scale, None)
127	- with pytest.raises(ValueError, match='finite and non-negative'):
128	- cann_ops_tilelang.engram_sinkhorn_step(matrix, row_scale, col_scale, -1.0)
129	- for invalid_eps in (float('inf'), float('nan')):
130	- with pytest.raises(ValueError, match='finite and non-negative'):
131	- cann_ops_tilelang.engram_sinkhorn_step(matrix, row_scale, col_scale, invalid_eps)
132	-
133	-
134	-def test_engram_sinkhorn_step_rejects_input_aliases():
135	- device = get_device()
136	- col_scale = torch.rand((256,), dtype=torch.float32, device=device)
137	- storage = torch.empty((17 * 256,), dtype=torch.float32, device=device)
138	- aliased_matrix = storage.view(17, 256)
139	- aliased_row_scale = storage[:17]
140	- with pytest.raises(ValueError, match='must not share storage'):
141	- cann_ops_tilelang.engram_sinkhorn_step(aliased_matrix, aliased_row_scale, col_scale, 1e-20)
142	-
143	-
144	-@pytest.mark.benchmark
145	-@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)
146	-def test_engram_sinkhorn_step_benchmark(benchmark_timer, benchmark_record, params):
147	- matrix, row_scale, col_scale = generate_test_data(params)
148	- row_norm, col_sumsq = cann_ops_tilelang.engram_sinkhorn_step(
149	- matrix,
150	- row_scale,
151	- col_scale,
152	- params['eps'],
153	- )
154	-
155	- time_us = benchmark_timer(
156	- lambda: cann_ops_tilelang.engram_sinkhorn_step(matrix, row_scale, col_scale, params['eps'])
157	- )
158	- num_bytes = count_bytes(matrix, row_scale, col_scale, row_scale, row_norm, col_sumsq)
159	- benchmark_record(
160	- kernel='engram_sinkhorn_step',
161	- operation='step',
162	- params=params,
163	- time_us=time_us,
164	- bandwidth_gbs=num_bytes / time_us / 1e3,
165	- )