fix(sampling): add ops-ops-tilelang-AtomGit

Msrc/cann_ops_tilelang/__init__.py+106-0

	@@ -47,8 +47,84 @@ __all__ = [
47	'swiglu_forward',	47	'swiglu_forward',
48	'swiglu_forward_and_per_token_cast',	48	'swiglu_forward_and_per_token_cast',
49	'unpack_from_e2m1fn_x2',	49	'unpack_from_e2m1fn_x2',
		50	+ 'apply_token_bitmask_inplace',
		51	+ 'attn_rel_h_rel_w',
		52	+ 'batch_find_value_by_index',
		53	+ 'bmp_to_patches',
		54	+ 'compress_and_update_state_decode',
		55	+ 'compress_and_update_state_prefill',
		56	+ 'compute_confidence_error',
		57	+ 'compute_seqlens_k_zos',
		58	+ 'compute_seqlens_q_by_confidence_rates',
		59	+ 'confidence_head_fused',
		60	+ 'convert_kv_index',
		61	+ 'dflash_input_store',
		62	+ 'find_value_by_index',
		63	+ 'fused_sample',
		64	+ 'gather',
		65	+ 'gather_engram_tokens_zos',
		66	+ 'gather_input_ids_zos',
		67	+ 'get_accept_result',
		68	+ 'get_dflash_mainx_slot_mapping',
		69	+ 'get_fused_dflash_get_input',
		70	+ 'get_kv_ca_index_decode',
		71	+ 'get_kv_ca_index_prefill',
		72	+ 'get_kv_swa_index_decode',
		73	+ 'get_load_kv_index_kernel',
		74	+ 'get_logits_cur_by_step',
		75	+ 'get_prompt_logprobs_tensor',
		76	+ 'get_sum_smaller_probs',
		77	+ 'kv_index_cat',
		78	+ 'layer_norm_fn',
		79	+ 'load_index_kv_cache',
		80	+ 'load_kv_cache',
		81	+ 'load_kv_cache_llama',
		82	+ 'load_state_to_cache',
		83	+ 'log_softmax',
		84	+ 'norm_fn',
		85	+ 'prepare_for_mirror_cp_tilelang',
		86	+ 'prepare_image_attn_pos_id',
		87	+ 'prepare_llama_decode_metadata',
		88	+ 'prepare_v33_decode_metadata',
		89	+ 'random_sampling',
		90	+ 'replace_input_embeds',
		91	+ 'rotary_embedding',
		92	+ 'sampled_tokens_update',
		93	+ 'scatter_verified',
		94	+ 'select_state_slice',
		95	+ 'shuffle_image_tokens',
		96	+ 'softmax',
		97	+ 'softmax_and_logsoftmax',
		98	+ 'store_index_kv_cache',
		99	+ 'store_kv_cache',
		100	+ 'store_kv_cache_llama',
		101	+ 'unbind',
		102	+ 'update_verify_states',
50	]	103	]
51		104
		105	+from .attention import (
		106	+ convert_kv_index,
		107	+ get_dflash_mainx_slot_mapping,
		108	+ get_fused_dflash_get_input,
		109	+ get_kv_ca_index_decode,
		110	+ get_kv_ca_index_prefill,
		111	+ get_kv_swa_index_decode,
		112	+ get_load_kv_index_kernel,
		113	+ kv_index_cat,
		114	+ load_index_kv_cache,
		115	+ load_kv_cache,
		116	+ load_kv_cache_llama,
		117	+ load_state_to_cache,
		118	+ prepare_for_mirror_cp_tilelang,
		119	+ prepare_llama_decode_metadata,
		120	+ prepare_v33_decode_metadata,
		121	+ rotary_embedding,
		122	+ select_state_slice,
		123	+ store_index_kv_cache,
		124	+ store_kv_cache,
		125	+ store_kv_cache_llama,
		126	+)
		127	+from .compression import compress_and_update_state_decode, compress_and_update_state_prefill
52	from .engram import (	128	from .engram import (
53	engram_fused_weight,	129	engram_fused_weight,
54	engram_gate_bwd,	130	engram_gate_bwd,
	@@ -82,6 +158,14 @@ from .mhc import (
82	mhc_sinkhorn_forward,	158	mhc_sinkhorn_forward,
83	)	159	)
84	from .moe import moe_topk_gate, normalize_weight, topk_gate	160	from .moe import moe_topk_gate, normalize_weight, topk_gate
		161	+from .multimodal import (
		162	+ attn_rel_h_rel_w,
		163	+ bmp_to_patches,
		164	+ prepare_image_attn_pos_id,
		165	+ replace_input_embeds,
		166	+ shuffle_image_tokens,
		167	+)
		168	+from .normalization import layer_norm_fn, norm_fn
85	from .quant import (	169	from .quant import (
86	cast_back,	170	cast_back,
87	per_block_cast,	171	per_block_cast,
	@@ -99,3 +183,25 @@ from .quant import (
99	swiglu_forward_and_per_token_cast,	183	swiglu_forward_and_per_token_cast,
100	unpack_from_e2m1fn_x2,	184	unpack_from_e2m1fn_x2,
101	)	185	)
		186	+from .sampling import (
		187	+ apply_token_bitmask_inplace,
		188	+ batch_find_value_by_index,
		189	+ compute_confidence_error,
		190	+ compute_seqlens_q_by_confidence_rates,
		191	+ confidence_head_fused,
		192	+ dflash_input_store,
		193	+ find_value_by_index,
		194	+ fused_sample,
		195	+ get_accept_result,
		196	+ get_logits_cur_by_step,
		197	+ get_prompt_logprobs_tensor,
		198	+ get_sum_smaller_probs,
		199	+ log_softmax,
		200	+ random_sampling,
		201	+ sampled_tokens_update,
		202	+ scatter_verified,
		203	+ softmax,
		204	+ softmax_and_logsoftmax,
		205	+ update_verify_states,
		206	+)
		207	+from .tensor import compute_seqlens_k_zos, gather, gather_engram_tokens_zos, gather_input_ids_zos, unbind

Asrc/cann_ops_tilelang/attention/__init__.py+48-0

@@ -0,0 +1,48 @@
	1	+"""Attention and KV-cache operators for Ascend."""
	2	+
	3	+__all__ = [
	4	+ 'convert_kv_index',
	5	+ 'get_dflash_mainx_slot_mapping',
	6	+ 'get_fused_dflash_get_input',
	7	+ 'get_kv_ca_index_decode',
	8	+ 'get_kv_ca_index_prefill',
	9	+ 'get_kv_swa_index_decode',
	10	+ 'get_load_kv_index_kernel',
	11	+ 'kv_index_cat',
	12	+ 'load_index_kv_cache',
	13	+ 'load_kv_cache',
	14	+ 'load_kv_cache_llama',
	15	+ 'load_state_to_cache',
	16	+ 'prepare_for_mirror_cp_tilelang',
	17	+ 'prepare_llama_decode_metadata',
	18	+ 'prepare_v33_decode_metadata',
	19	+ 'rotary_embedding',
	20	+ 'select_state_slice',
	21	+ 'store_index_kv_cache',
	22	+ 'store_kv_cache',
	23	+ 'store_kv_cache_llama',
	24	+]
	25	+
	26	+from .cache import (
	27	+ load_index_kv_cache,
	28	+ load_kv_cache,
	29	+ load_kv_cache_llama,
	30	+ store_index_kv_cache,
	31	+ store_kv_cache,
	32	+ store_kv_cache_llama,
	33	+)
	34	+from .flash_mla_utils import convert_kv_index
	35	+from .get_kv_index import (
	36	+ get_dflash_mainx_slot_mapping,
	37	+ get_fused_dflash_get_input,
	38	+ get_kv_ca_index_decode,
	39	+ get_kv_ca_index_prefill,
	40	+ get_kv_swa_index_decode,
	41	+ get_load_kv_index_kernel,
	42	+)
	43	+from .kv_index_cat import kv_index_cat
	44	+from .prepare_cp_mirror import prepare_for_mirror_cp_tilelang
	45	+from .prepare_llama_decode_metadata import prepare_llama_decode_metadata
	46	+from .prepare_v33_decode_metadata import prepare_v33_decode_metadata
	47	+from .rope import rotary_embedding
	48	+from .state_cache import load_state_to_cache, select_state_slice

Asrc/cann_ops_tilelang/attention/_simd.py+31-0

@@ -0,0 +1,31 @@
	1	+"""Small Ascend SIMD helpers shared by migrated cache kernels."""
	2	+
	3	+from typing import Any
	4	+
	5	+import tilelang
	6	+from tilelang.ascend import language as T
	7	+from tilelang.backend.target import determine_target
	8	+
	9	+ASCEND_VSTSX2_BF16_PRELUDE = r"""
	10	+__simd_callee__ inline void tl_vstsx2_bf16(
	11	+ __ubuf__ bfloat16_t *dst, vector_bf16 even, vector_bf16 odd,
	12	+ vector_bool mask) {
	13	+ ::vsts(even, odd, dst, 0, INTLV_B16, mask);
	14	+}
	15	+"""
	16	+
	17	+
	18	+def is_pto_target() -> bool:
	19	+ target = determine_target(tilelang.env.get_default_target(), return_object=True)
	20	+ keys = target.keys
	21	+ if isinstance(keys, str):
	22	+ keys = (keys,)
	23	+ return any(str(key) == 'pto' for key in keys)
	24	+
	25	+
	26	+def vstsx2_bf16(dst: Any, even: Any, odd: Any, mask: Any) -> None:
	27	+ dst_ptr = T.access_ptr(dst, 'w', 256)
	28	+ if is_pto_target():
	29	+ T.evaluate(T.call_extern('void', 'pto.vstsx2', even, odd, dst_ptr, T.int32(0), 'INTLV_B16', mask))
	30	+ else:
	31	+ T.evaluate(T.call_extern('void', 'tl_vstsx2_bf16', dst_ptr, even, odd, mask))

Asrc/cann_ops_tilelang/attention/cache.py+401-0

@@ -0,0 +1,401 @@
	1	+import torch
	2	+
	3	+from cann_ops_tilelang.attention.cache_kernel import (
	4	+ _load_index_kv_cache_fp4_tl_ascend,
	5	+ _load_index_kv_cache_fp8_per32_tl_ascend,
	6	+ _load_index_kv_cache_tl_ascend,
	7	+ _load_kv_cache_llama_tl_ascend,
	8	+ _load_kv_cache_per32_tl_ascend,
	9	+ _load_kv_cache_tl_ascend,
	10	+ _store_index_kv_cache_fp4_tl_ascend,
	11	+ _store_index_kv_cache_fp8_per32_tl_ascend,
	12	+ _store_index_kv_cache_tl_ascend,
	13	+ _store_kv_cache_llama_tl_ascend,
	14	+ _store_kv_cache_per32_tl_ascend,
	15	+ _store_kv_cache_tl_ascend,
	16	+)
	17	+from cann_ops_tilelang.config import require_ascend
	18	+from cann_ops_tilelang.quant import per_token_cast
	19	+
	20	+
	21	+def _validate_npu_tensor(tensor: torch.Tensor, name: str) -> None:
	22	+ if tensor.device.type != 'npu':
	23	+ raise ValueError(f'{name} must be on an NPU, got {tensor.device}.')
	24	+
	25	+
	26	+def _validate_slot_mapping(slot_mapping: torch.Tensor, expected_tokens: int \| None = None) -> None:
	27	+ _validate_npu_tensor(slot_mapping, 'slot_mapping')
	28	+ if slot_mapping.ndim != 1 or slot_mapping.dtype != torch.int32:
	29	+ raise ValueError('slot_mapping must be a 1D torch.int32 tensor.')
	30	+ if expected_tokens is not None and slot_mapping.numel() != expected_tokens:
	31	+ raise ValueError(f'slot_mapping must contain {expected_tokens} entries, got {slot_mapping.numel()}.')
	32	+
	33	+
	34	+def _validate_cache(cache: torch.Tensor, slot_mapping: torch.Tensor, *, dtype: torch.dtype) -> None:
	35	+ _validate_npu_tensor(cache, 'kv_cache')
	36	+ if cache.device != slot_mapping.device:
	37	+ raise ValueError('kv_cache and slot_mapping must be on the same device.')
	38	+ if cache.ndim != 4 or cache.dtype != dtype:
	39	+ raise ValueError(f'kv_cache must be a 4D {dtype} tensor.')
	40	+
	41	+
	42	+def store_kv_cache(
	43	+ kv: torch.Tensor,
	44	+ kv_cache: torch.Tensor \| None,
	45	+ slot_mapping: torch.Tensor,
	46	+ use_fp8_rope: bool = False,
	47	+ num_per_channels: int = 64,
	48	+) -> None:
	49	+ """Store KV to cache with fp8 quantization.
	50	+
	51	+ Supported configurations:
	52	+ - (use_fp8_rope=False, num_per_channels=64): nope(448) per-64 fp8 + rope(64) bf16.
	53	+ bytes_per_token = 448 + 64*2 + 8 = 584.
	54	+ - (use_fp8_rope=True, num_per_channels=32): full 512 per-32 fp8.
	55	+ bytes_per_token = 512 + 16 = 528.
	56	+
	57	+ Args:
	58	+ kv: [n, 1, 512], bf16
	59	+ kv_cache: [num_blocks, block_size, 1, bytes_per_token], uint8
	60	+ slot_mapping: [n], int32, -1 means padding
	61	+ use_fp8_rope: if True, rope is also stored as fp8 (per32 mode)
	62	+ num_per_channels: quantization group size (64 or 32)
	63	+ """
	64	+ require_ascend()
	65	+ _validate_npu_tensor(kv, 'kv')
	66	+ _validate_slot_mapping(slot_mapping, kv.shape[0])
	67	+ if kv.ndim != 3 or kv.shape[1:] != (1, 512) or kv.dtype != torch.bfloat16:
	68	+ raise ValueError('kv must have shape [num_tokens, 1, 512] and dtype torch.bfloat16.')
	69	+ if kv_cache is None:
	70	+ return
	71	+ _validate_cache(kv_cache, slot_mapping, dtype=torch.uint8)
	72	+ if kv.device != kv_cache.device:
	73	+ raise ValueError('kv and kv_cache must be on the same device.')
	74	+ if not use_fp8_rope and num_per_channels == 64:
	75	+ num_blocks, block_size, constant_1, bytes_per_token = kv_cache.shape
	76	+ assert constant_1 == 1
	77	+ assert bytes_per_token == 448 + 64 * 2 + 8
	78	+ kv_cache = kv_cache.view(num_blocks, block_size * bytes_per_token)
	79	+ kv_cache_value = kv_cache[:, : block_size * (448 + 64 * 2)].view(num_blocks, block_size, 448 + 64 * 2)
	80	+ kv_cache_sf = kv_cache[:, block_size * (448 + 64 * 2) :]
	81	+ kv_cache_fp8_raw = kv_cache_value[:, :, :448]
	82	+ kv_cache_fp8 = kv_cache_fp8_raw.view(torch.float8_e4m3fn)
	83	+ kv_cache_bf16 = kv_cache_value[:, :, 448:]
	84	+ kv_cache_bf16 = kv_cache_bf16.view(*kv_cache_bf16.shape[:-1], -1).view(torch.bfloat16)
	85	+ kv_cache_sf = kv_cache_sf.view(num_blocks, block_size, 8)
	86	+ assert kv_cache_fp8.stride(0) == kv_cache_bf16.stride(0) * 2
	87	+ assert kv_cache_fp8.stride(0) == kv_cache_sf.stride(0)
	88	+ num_stages = 2
	89	+ n_tokens = kv.shape[0]
	90	+ if n_tokens % 64 == 0:
	91	+ store_block_m = 64
	92	+ elif n_tokens % 32 == 0:
	93	+ store_block_m = 32
	94	+ else:
	95	+ store_block_m = 16
	96	+ full_tiles = n_tokens % store_block_m == 0
	97	+ store_kernel = _store_kv_cache_tl_ascend(
	98	+ block_size,
	99	+ num_stages=num_stages,
	100	+ block_m=store_block_m,
	101	+ cache_block_stride=kv_cache.stride(0),
	102	+ static_n=n_tokens if n_tokens == 65536 else 0,
	103	+ static_num_blocks=num_blocks if block_size in (64, 129) else 0,
	104	+ full_tiles=full_tiles,
	105	+ )
	106	+ store_kernel(kv[:, 0, :], kv_cache, slot_mapping)
	107	+ elif use_fp8_rope and num_per_channels == 32:
	108	+ num_blocks, block_size, constant_1, bytes_per_token = kv_cache.shape
	109	+ assert constant_1 == 1
	110	+ assert bytes_per_token == 512 + 16
	111	+ kv_cache = kv_cache.view(num_blocks, block_size * bytes_per_token)
	112	+ kv_cache_fp8 = kv_cache[:, : block_size * 512].view(num_blocks, block_size, 512).view(torch.float8_e4m3fn)
	113	+ kv_cache_sf = kv_cache[:, block_size * 512 :].view(num_blocks, block_size, 16)
	114	+ assert kv_cache_fp8.stride(0) == kv_cache_sf.stride(0)
	115	+ store_block_m = 64
	116	+ store_kernel = _store_kv_cache_per32_tl_ascend(
	117	+ block_size, full_tiles=kv.shape[0] % store_block_m == 0, block_m=store_block_m
	118	+ )
	119	+ store_kernel(kv[:, 0, :], kv_cache_fp8, kv_cache_sf, slot_mapping)
	120	+ else:
	121	+ raise ValueError(
	122	+ f'Unsupported combination: use_fp8_rope={use_fp8_rope}, '
	123	+ f'num_per_channels={num_per_channels}. Supported: (False, 64) and (True, 32).'
	124	+ )
	125	+
	126	+
	127	+def load_kv_cache(
	128	+ kv_cache: torch.Tensor \| None,
	129	+ slot_mapping: torch.Tensor,
	130	+ out: torch.Tensor \| None = None,
	131	+ use_fp8_rope: bool = False,
	132	+ num_per_channels: int = 64,
	133	+) -> torch.Tensor:
	134	+ """Load KV from cache with fp8 dequantization.
	135	+
	136	+ Supported configurations:
	137	+ - (use_fp8_rope=False, num_per_channels=64): nope(448) per-64 fp8 + rope(64) bf16.
	138	+ - (use_fp8_rope=True, num_per_channels=32): full 512 per-32 fp8.
	139	+
	140	+ Args:
	141	+ kv_cache: [num_blocks, block_size, 1, bytes_per_token], uint8
	142	+ slot_mapping: [n], int32, -1 means padding
	143	+ out: optional output tensor [n, 1, 512], bf16
	144	+ use_fp8_rope: if True, rope is also stored as fp8 (per32 mode)
	145	+ num_per_channels: quantization group size (64 or 32)
	146	+
	147	+ Returns:
	148	+ out: [n, 1, 512], bf16
	149	+ """
	150	+ require_ascend()
	151	+ _validate_slot_mapping(slot_mapping)
	152	+ if out is not None:
	153	+ _validate_npu_tensor(out, 'out')
	154	+ if out.device != slot_mapping.device or out.dtype != torch.bfloat16:
	155	+ raise ValueError('out must be a BF16 tensor on the slot_mapping device.')
	156	+ if kv_cache is None:
	157	+ o = torch.randn(slot_mapping.shape[0], 1, 512, dtype=torch.bfloat16, device=slot_mapping.device)
	158	+ if out is None:
	159	+ out = o
	160	+ else:
	161	+ out.copy_(o)
	162	+ return out
	163	+ _validate_cache(kv_cache, slot_mapping, dtype=torch.uint8)
	164	+ n = slot_mapping.shape[0]
	165	+ if not use_fp8_rope and num_per_channels == 64:
	166	+ num_blocks, block_size, constant_1, bytes_per_token = kv_cache.shape
	167	+ assert constant_1 == 1
	168	+ assert bytes_per_token == 448 + 64 * 2 + 8
	169	+ kv_cache = kv_cache.view(num_blocks, block_size * bytes_per_token)
	170	+ kv_cache_value = kv_cache[:, : block_size * (448 + 64 * 2)].view(num_blocks, block_size, 448 + 64 * 2)
	171	+ kv_cache_sf = kv_cache[:, block_size * (448 + 64 * 2) :]
	172	+ kv_cache_fp8 = kv_cache_value[:, :, :448].view(torch.float8_e4m3fn)
	173	+ kv_cache_bf16 = kv_cache_value[:, :, 448:]
	174	+ kv_cache_bf16 = kv_cache_bf16.view(*kv_cache_bf16.shape[:-1], -1).view(torch.bfloat16)
	175	+ kv_cache_sf = kv_cache_sf.view(num_blocks, block_size, 8)
	176	+ assert kv_cache_fp8.stride(0) == kv_cache_bf16.stride(0) * 2
	177	+ assert kv_cache_fp8.stride(0) == kv_cache_sf.stride(0)
	178	+ if out is None:
	179	+ out = torch.empty((n, 1, 512), dtype=torch.bfloat16, device=slot_mapping.device)
	180	+ else:
	181	+ assert out.shape == (n, 1, 512)
	182	+ if n % 64 == 0:
	183	+ load_kernel = _load_kv_cache_tl_ascend(block_size, num_stages=2, full_tiles=True, block_m=64)
	184	+ else:
	185	+ load_kernel = _load_kv_cache_tl_ascend(block_size, num_stages=4, full_tiles=False, block_m=32)
	186	+ load_kernel(out[:, 0, :], kv_cache_value, kv_cache_sf, slot_mapping)
	187	+ elif use_fp8_rope and num_per_channels == 32:
	188	+ num_blocks, block_size, constant_1, bytes_per_token = kv_cache.shape
	189	+ assert constant_1 == 1
	190	+ assert bytes_per_token == 512 + 16
	191	+ kv_cache = kv_cache.view(num_blocks, block_size * bytes_per_token)
	192	+ kv_cache_fp8 = kv_cache[:, : block_size * 512].view(num_blocks, block_size, 512).view(torch.float8_e4m3fn)
	193	+ kv_cache_sf = kv_cache[:, block_size * 512 :].view(num_blocks, block_size, 16)
	194	+ cache_block_stride = kv_cache_fp8.stride(0)
	195	+ assert cache_block_stride == kv_cache_sf.stride(0)
	196	+ assert cache_block_stride >= block_size * (512 + 16)
	197	+ assert cache_block_stride % 32 == 0
	198	+ assert kv_cache_fp8.stride(1) == 512
	199	+ assert kv_cache_sf.stride(1) == 16
	200	+ if out is None:
	201	+ out = torch.empty((n, 1, 512), dtype=torch.bfloat16, device=slot_mapping.device)
	202	+ else:
	203	+ assert out.shape == (n, 1, 512)
	204	+ load_block_m = 64 if n % 64 == 0 else 32
	205	+ load_kernel = _load_kv_cache_per32_tl_ascend(
	206	+ block_size,
	207	+ full_tiles=n % load_block_m == 0,
	208	+ block_m=load_block_m,
	209	+ cache_block_stride=cache_block_stride,
	210	+ prefetch_slot_waves=n == 65536 and load_block_m == 64,
	211	+ )
	212	+ load_kernel(out[:, 0, :], kv_cache_fp8, kv_cache_sf, slot_mapping)
	213	+ else:
	214	+ raise ValueError(
	215	+ f'Unsupported combination: use_fp8_rope={use_fp8_rope}, '
	216	+ f'num_per_channels={num_per_channels}. Supported: (False, 64) and (True, 32).'
	217	+ )
	218	+ return out
	219	+
	220	+
	221	+def store_index_kv_cache(
	222	+ kv: torch.Tensor,
	223	+ kv_cache: torch.Tensor \| None,
	224	+ slot_mapping: torch.Tensor,
	225	+ is_fp4: bool = False,
	226	+ is_mxfp8: bool = False,
	227	+) -> None:
	228	+ """
	229	+ Args:
	230	+ kv: [n, 1, dim], bf16
	231	+ kv_cache: [num_blocks, block_size, 1, bytes_per_token], uint8
	232	+ slot_mapping: [n], int32, -1 means padding
	233	+ """
	234	+ require_ascend()
	235	+ _validate_npu_tensor(kv, 'kv')
	236	+ _validate_slot_mapping(slot_mapping, kv.shape[0])
	237	+ if kv.ndim != 3 or kv.shape[1:] != (1, 128) or kv.dtype != torch.bfloat16:
	238	+ raise ValueError('kv must have shape [num_tokens, 1, 128] and dtype torch.bfloat16.')
	239	+ if kv_cache is None:
	240	+ return
	241	+ _validate_cache(kv_cache, slot_mapping, dtype=torch.uint8)
	242	+ if kv.device != kv_cache.device:
	243	+ raise ValueError('kv and kv_cache must be on the same device.')
	244	+ num_blocks, block_size, constant_1, bytes_per_token = kv_cache.shape
	245	+ assert constant_1 == 1
	246	+ if is_fp4:
	247	+ assert bytes_per_token == 64 + 4
	248	+ kv_cache = kv_cache.view(num_blocks, block_size * bytes_per_token)
	249	+ kv_cache_fp4 = kv_cache[:, : block_size * 64].view(num_blocks, block_size, 64).view(torch.int8)
	250	+ kv_cache_sf = kv_cache[:, block_size * 64 :].view(num_blocks, block_size, 4).view(torch.uint8)
	251	+ assert kv_cache_fp4.stride(0) == kv_cache_sf.stride(0)
	252	+ store_kernel = _store_index_kv_cache_fp4_tl_ascend(
	253	+ block_size, kv_cache_fp4.stride(1) * 2, kv_cache_sf.stride(1)
	254	+ )
	255	+ store_kernel(kv[:, 0, :], kv_cache_fp4, kv_cache_sf, slot_mapping)
	256	+ elif is_mxfp8:
	257	+ assert bytes_per_token == 128 + 4
	258	+ kv_cache = kv_cache.view(num_blocks, block_size * bytes_per_token)
	259	+ kv_cache_fp8 = kv_cache[:, : block_size * 128].view(num_blocks, block_size, 128).view(torch.float8_e4m3fn)
	260	+ kv_cache_sf = kv_cache[:, block_size * 128 :].view(num_blocks, block_size, 4).view(torch.uint8)
	261	+ assert kv_cache_fp8.stride(0) == kv_cache_sf.stride(0)
	262	+ store_kernel = _store_index_kv_cache_fp8_per32_tl_ascend(
	263	+ block_size, kv_cache_fp8.stride(1), kv_cache_sf.stride(1)
	264	+ )
	265	+ store_kernel(kv[:, 0, :], kv_cache_fp8, kv_cache_sf, slot_mapping)
	266	+ else:
	267	+ assert bytes_per_token == 128 + 4
	268	+ kv_cache = kv_cache.view(num_blocks, block_size * bytes_per_token)
	269	+ kv_cache_fp8 = kv_cache[:, : block_size * 128].view(num_blocks, block_size, 128).view(torch.float8_e4m3fn)
	270	+ kv_cache_sf = kv_cache[:, block_size * 128 :].view(num_blocks, block_size, 4)
	271	+ kv_cache_sf = kv_cache_sf.view(*kv_cache_sf.shape[:-1], -1).view(torch.float)
	272	+ assert kv_cache_fp8.stride(0) == kv_cache_sf.stride(0) * 4
	273	+ store_kernel = _store_index_kv_cache_tl_ascend(block_size, kv_cache_fp8.stride(1), kv_cache_sf.stride(1))
	274	+ store_kernel(kv[:, 0, :], kv_cache_fp8, kv_cache_sf, slot_mapping)
	275	+
	276	+
	277	+def load_index_kv_cache(
	278	+ kv_cache: torch.Tensor \| None, slot_mapping: torch.Tensor, is_fp4: bool = False, is_mxfp8: bool = False
	279	+) -> tuple[torch.Tensor, torch.Tensor]:
	280	+ """
	281	+ Args:
	282	+ kv_cache: [num_blocks, block_size, 1, bytes_per_token], uint8
	283	+ slot_mapping: [n], int32, -1 means padding
	284	+ is_fp4: bool, whether to return fp4 or fp8
	285	+
	286	+ Returns:
	287	+ if is_fp4:
	288	+ kv_fp4: [n, dim // 2], int8
	289	+ kv_sf: [n], int32
	290	+ if is_mxfp8 (fp8 per-32 ue8m0):
	291	+ kv_fp8: [n, dim], float8_e4m3
	292	+ kv_sf: [n], int32 (packed 4 ue8m0)
	293	+ else (fp8 per-token float32):
	294	+ kv_fp8: [n, dim], float8_e4m3
	295	+ kv_sf: [n], float32
	296	+ """
	297	+ require_ascend()
	298	+ _validate_slot_mapping(slot_mapping)
	299	+ if kv_cache is None:
	300	+ if is_fp4:
	301	+ k_bf16 = torch.randn(slot_mapping.shape[0], 128, dtype=torch.bfloat16, device=slot_mapping.device)
	302	+ k_fp4, k_sf = per_token_cast(
	303	+ k_bf16, 'e2m1', 32, round_sf=True, use_tma_aligned_col_major_sf=True, use_packed_ue8m0=True
	304	+ )
	305	+ return (k_fp4, k_sf.squeeze(-1))
	306	+ elif is_mxfp8:
	307	+ k_bf16 = torch.randn(slot_mapping.shape[0], 128, dtype=torch.bfloat16, device=slot_mapping.device)
	308	+ k_fp8, k_sf = per_token_cast(
	309	+ k_bf16, 'e4m3', 32, round_sf=True, use_tma_aligned_col_major_sf=False, use_packed_ue8m0=True
	310	+ )
	311	+ return (k_fp8, k_sf.squeeze(-1))
	312	+ else:
	313	+ k_bf16 = torch.randn(slot_mapping.shape[0], 128, dtype=torch.bfloat16, device=slot_mapping.device)
	314	+ k_fp8, k_sf = per_token_cast(k_bf16, 'e4m3', 128, round_sf=True)
	315	+ return (k_fp8, k_sf.squeeze(-1))
	316	+ _validate_cache(kv_cache, slot_mapping, dtype=torch.uint8)
	317	+ num_blocks, block_size, constant_1, bytes_per_token = kv_cache.shape
	318	+ n = slot_mapping.shape[0]
	319	+ assert constant_1 == 1
	320	+ if is_fp4:
	321	+ assert bytes_per_token == 64 + 4
	322	+ kv_cache = kv_cache.view(num_blocks, block_size * bytes_per_token)
	323	+ kv_cache_fp4 = kv_cache[:, : block_size * 64].view(num_blocks, block_size, 64).view(torch.int8)
	324	+ kv_cache_sf = kv_cache[:, block_size * 64 :].view(num_blocks, block_size, 4)
	325	+ kv_cache_sf = kv_cache_sf.view(*kv_cache_sf.shape[:-1], -1).view(torch.int32)
	326	+ assert kv_cache_fp4.stride(0) == kv_cache_sf.stride(0) * 4
	327	+ kv_fp4 = torch.empty((n, 64), dtype=torch.int8, device=slot_mapping.device)
	328	+ kv_sf = torch.empty((n,), dtype=torch.int32, device=slot_mapping.device)
	329	+ load_kernel = _load_index_kv_cache_fp4_tl_ascend(block_size, kv_cache_fp4.stride(1), kv_cache_sf.stride(1))
	330	+ load_kernel(kv_fp4.view(torch.uint8), kv_sf, kv_cache_fp4.view(torch.uint8), kv_cache_sf, slot_mapping)
	331	+ return (kv_fp4, kv_sf)
	332	+ elif is_mxfp8:
	333	+ assert bytes_per_token == 128 + 4
	334	+ kv_cache = kv_cache.view(num_blocks, block_size * bytes_per_token)
	335	+ kv_cache_fp8 = kv_cache[:, : block_size * 128].view(num_blocks, block_size, 128).view(torch.float8_e4m3fn)
	336	+ kv_cache_sf = kv_cache[:, block_size * 128 :].view(num_blocks, block_size, 4)
	337	+ kv_cache_sf = kv_cache_sf.view(*kv_cache_sf.shape[:-1], -1).view(torch.int32)
	338	+ assert kv_cache_fp8.stride(0) == kv_cache_sf.stride(0) * 4
	339	+ kv_fp8 = torch.empty((n, 128), dtype=torch.float8_e4m3fn, device=slot_mapping.device)
	340	+ kv_scale = torch.empty((n,), dtype=torch.int32, device=slot_mapping.device)
	341	+ load_kernel = _load_index_kv_cache_fp8_per32_tl_ascend(
	342	+ block_size, kv_cache_fp8.stride(1), kv_cache_sf.stride(1)
	343	+ )
	344	+ load_kernel(kv_fp8.view(torch.uint8), kv_scale, kv_cache_fp8.view(torch.uint8), kv_cache_sf, slot_mapping)
	345	+ return (kv_fp8, kv_scale)
	346	+ else:
	347	+ assert bytes_per_token == 128 + 4
	348	+ kv_cache = kv_cache.view(num_blocks, block_size * bytes_per_token)
	349	+ kv_cache_fp8 = kv_cache[:, : block_size * 128].view(num_blocks, block_size, 128).view(torch.float8_e4m3fn)
	350	+ kv_cache_sf = kv_cache[:, block_size * 128 :].view(num_blocks, block_size, 4)
	351	+ kv_cache_sf = kv_cache_sf.view(*kv_cache_sf.shape[:-1], -1).view(torch.float)
	352	+ assert kv_cache_fp8.stride(0) == kv_cache_sf.stride(0) * 4
	353	+ kv_fp8 = torch.empty((n, 128), dtype=torch.float8_e4m3fn, device=slot_mapping.device)
	354	+ kv_sf = torch.empty((n,), dtype=torch.float, device=slot_mapping.device)
	355	+ load_kernel = _load_index_kv_cache_tl_ascend(block_size, kv_cache_fp8.stride(1), kv_cache_sf.stride(1))
	356	+ load_kernel(
	357	+ kv_fp8.view(torch.uint8),
	358	+ kv_sf.view(torch.int32),
	359	+ kv_cache_fp8.view(torch.uint8),
	360	+ kv_cache_sf.view(torch.int32),
	361	+ slot_mapping,
	362	+ )
	363	+ return (kv_fp8, kv_sf)
	364	+
	365	+
	366	+def store_kv_cache_llama(store_kv: torch.Tensor, kv_cache: torch.Tensor \| None, slot_mapping: torch.Tensor) -> None:
	367	+ require_ascend()
	368	+ _validate_npu_tensor(store_kv, 'store_kv')
	369	+ _validate_slot_mapping(slot_mapping, store_kv.shape[0])
	370	+ if store_kv.ndim != 3 or store_kv.shape[2] != 128 or store_kv.dtype != torch.bfloat16:
	371	+ raise ValueError('store_kv must have shape [num_tokens, num_heads, 128] and dtype torch.bfloat16.')
	372	+ if kv_cache is None:
	373	+ return
	374	+ _validate_cache(kv_cache, slot_mapping, dtype=torch.bfloat16)
	375	+ if store_kv.device != kv_cache.device:
	376	+ raise ValueError('store_kv and kv_cache must be on the same device.')
	377	+ _, block_size, n_head_kv, head_size = kv_cache.shape
	378	+ assert head_size == 128
	379	+ store_kv = store_kv.view(-1, n_head_kv * head_size)
	380	+ kv_cache = kv_cache.view(-1, block_size, n_head_kv * head_size)
	381	+ store_kernel = _store_kv_cache_llama_tl_ascend(block_size, n_head_kv * head_size)
	382	+ store_kernel(store_kv, kv_cache, slot_mapping)
	383	+
	384	+
	385	+def load_kv_cache_llama(kv_cache: torch.Tensor \| None, n_head_kv: int, slot_mapping: torch.Tensor) -> torch.Tensor:
	386	+ require_ascend()
	387	+ _validate_slot_mapping(slot_mapping)
	388	+ if n_head_kv <= 0:
	389	+ raise ValueError('n_head_kv must be positive.')
	390	+ if kv_cache is None:
	391	+ return torch.randn((slot_mapping.shape[0], n_head_kv, 128), dtype=torch.bfloat16, device=slot_mapping.device)
	392	+ _validate_cache(kv_cache, slot_mapping, dtype=torch.bfloat16)
	393	+ _, block_size, num_head_kv, head_size = kv_cache.shape
	394	+ assert head_size == 128
	395	+ assert num_head_kv == n_head_kv
	396	+ n = slot_mapping.shape[0]
	397	+ kv_cache = kv_cache.view(-1, block_size, n_head_kv * head_size)
	398	+ kv = torch.empty((n, n_head_kv * head_size), dtype=torch.bfloat16, device=slot_mapping.device)
	399	+ load_kernel = _load_kv_cache_llama_tl_ascend(block_size, n_head_kv * head_size)
	400	+ load_kernel(kv, kv_cache, slot_mapping)
	401	+ return kv.view(n, n_head_kv, head_size)

Asrc/cann_ops_tilelang/attention/cache_kernel.py+1342-0

Asrc/cann_ops_tilelang/attention/flash_mla_utils.py+37-0

@@ -0,0 +1,37 @@
	1	+import torch
	2	+
	3	+from cann_ops_tilelang.attention.flash_mla_utils_kernel import _convert_kv_index_tl_ascend
	4	+from cann_ops_tilelang.config import require_ascend
	5	+
	6	+
	7	+def convert_kv_index(
	8	+ topk_idxs: torch.Tensor, block_size: int, context_lens: torch.Tensor, block_tables: torch.Tensor
	9	+) -> None:
	10	+ """
	11	+ Args:
	12	+ topk_idxs: [bsz, seqlen_q, topk_tokens]
	13	+ context_lens: [bsz, seqlen_q]
	14	+ block_tables: [bsz, max_num_blocks]
	15	+ """
	16	+ require_ascend()
	17	+ if any(tensor.device.type != 'npu' for tensor in (topk_idxs, context_lens, block_tables)):
	18	+ raise ValueError('topk_idxs, context_lens, and block_tables must be NPU tensors.')
	19	+ if not (topk_idxs.device == context_lens.device == block_tables.device):
	20	+ raise ValueError('topk_idxs, context_lens, and block_tables must be on the same device.')
	21	+ if any(tensor.dtype != torch.int32 for tensor in (topk_idxs, context_lens, block_tables)):
	22	+ raise TypeError('topk_idxs, context_lens, and block_tables must have dtype torch.int32.')
	23	+ if block_size <= 0:
	24	+ raise ValueError('block_size must be positive.')
	25	+ assert topk_idxs.is_contiguous()
	26	+ assert context_lens.is_contiguous()
	27	+ assert block_tables.is_contiguous()
	28	+ assert topk_idxs.ndim == 3
	29	+ assert context_lens.ndim == 2
	30	+ assert block_tables.ndim == 2
	31	+ b, s, k = topk_idxs.shape
	32	+ _b, n = block_tables.shape
	33	+ assert b == _b
	34	+ if context_lens.shape != (b, s):
	35	+ raise ValueError(f'context_lens must have shape {(b, s)}, got {tuple(context_lens.shape)}.')
	36	+ kernel = _convert_kv_index_tl_ascend(s, k, n, block_size)
	37	+ kernel(topk_idxs, context_lens, block_tables)

Asrc/cann_ops_tilelang/attention/flash_mla_utils_kernel.py+262-0

@@ -0,0 +1,262 @@
	1	+# Ruff's simplification rules do not model TileLang's compile-time branches.
	2	+# ruff: noqa: SIM108
	3	+
	4	+import tilelang
	5	+from tilelang.ascend import language as T
	6	+
	7	+from cann_ops_tilelang.config import get_num_vector_cores
	8	+
	9	+_UB_BUDGET_BYTES = 200 * 1024
	10	+
	11	+
	12	+@tilelang.jit(
	13	+ out_idx=None,
	14	+ target='ascend',
	15	+ pass_configs={
	16	+ tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True,
	17	+ tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
	18	+ },
	19	+)
	20	+def _convert_kv_index_tl_ascend(
	21	+ s: int,
	22	+ k: int,
	23	+ n: int,
	24	+ block_size: int,
	25	+) -> None:
	26	+ """Single kernel with compile-time specialization: layouts where k is a
	27	+ multiple of 64 and n a multiple of 8 use BR-task batches with flat
	28	+ contiguous DMAs (the batch stride B == k, so DMA and compute offsets
	29	+ match); anything else uses BR=1 per-task 2D copies (unaligned-safe).
	30	+ Both share the same SIMD vgather2 core; blk_ub is padded to >=VL
	31	+ elements so vgather2's 64-lane span stays inside one version slot."""
	32	+ b = T.dynamic('b')
	33	+ num_cores = get_num_vector_cores()
	34	+ vector_length = 64
	35	+ index_tile_size = (k + vector_length - 1) // vector_length * vector_length
	36	+ table_tile_size = max(n, vector_length)
	37	+ batched = k % vector_length == 0 and n % 8 == 0
	38	+ chunked = False
	39	+ rows_per_batch = 5
	40	+ while (
	41	+ rows_per_batch > 1
	42	+ and (rows_per_batch * index_tile_size + rows_per_batch * table_tile_size) * 4 * 2 > _UB_BUDGET_BYTES
	43	+ ):
	44	+ rows_per_batch -= 1
	45	+ if batched and (rows_per_batch * index_tile_size + rows_per_batch * table_tile_size) * 4 * 2 > _UB_BUDGET_BYTES:
	46	+ batched = False
	47	+ if batched:
	48	+ num_stages = 2
	49	+ else:
	50	+ rows_per_batch = 1
	51	+ one_stage_bytes = (index_tile_size + table_tile_size) * 4
	52	+ if one_stage_bytes > _UB_BUDGET_BYTES:
	53	+ chunked = True
	54	+ num_stages = 1
	55	+ else:
	56	+ num_stages = min(4, _UB_BUDGET_BYTES // one_stage_bytes)
	57	+ is_pow2 = (block_size & (block_size - 1)) == 0
	58	+ block_bits = block_size.bit_length() - 1 if is_pow2 else 0
	59	+
	60	+ if chunked:
	61	+ index_tile = 4096
	62	+ block_table_tile = 32768
	63	+ num_index_blocks = (k + index_tile - 1) // index_tile
	64	+
	65	+ @T.prim_func
	66	+ def _convert_kv_index_chunked(
	67	+ topk_idxs: T.Tensor([b, s, k], T.int32), # type: ignore
	68	+ context_lens: T.Tensor([b, s], T.int32), # type: ignore
	69	+ block_tables: T.Tensor([b, n], T.int32), # type: ignore
	70	+ ):
	71	+ with T.Kernel(num_cores) as pid:
	72	+ idx_ub = T.alloc_shared((index_tile,), 'int32')
	73	+ out_ub = T.alloc_shared((index_tile,), 'int32')
	74	+ blk_ub = T.alloc_shared((block_table_tile,), 'int32')
	75	+
	76	+ for task in T.Persistent([b * s * num_index_blocks], num_cores, pid, group_size=1, num_stages=1):
	77	+ row_task = task // num_index_blocks
	78	+ index_block = task % num_index_blocks
	79	+ bid = row_task // s
	80	+ sid = row_task % s
	81	+ cl = context_lens[bid, sid]
	82	+ index_begin = index_block * index_tile
	83	+ index_count = T.min(index_tile, k - index_begin)
	84	+ T.copy(
	85	+ topk_idxs[bid, sid, index_begin : index_begin + index_count],
	86	+ idx_ub[:index_count],
	87	+ )
	88	+
	89	+ with T.SimdVF():
	90	+ mask = T.simd.pset(32)
	91	+ neg1 = T.simd.vdup(-1, 'int32', mask)
	92	+ for vector in T.serial(T.ceildiv(index_count, vector_length)):
	93	+ T.simd.vsts(out_ub[vector * vector_length], neg1, mask)
	94	+
	95	+ for table_block in T.serial(T.ceildiv(n, block_table_tile)):
	96	+ table_begin = table_block * block_table_tile
	97	+ table_count = T.min(block_table_tile, n - table_begin)
	98	+ T.copy(
	99	+ block_tables[bid, table_begin : table_begin + table_count],
	100	+ blk_ub[:table_count],
	101	+ )
	102	+
	103	+ with T.SimdVF():
	104	+ mask = T.simd.pset(32)
	105	+ zero = T.simd.vdup(0, 'int32', mask)
	106	+ table_begin_vec = T.simd.vdup(table_begin, 'int32', mask)
	107	+ table_end_vec = T.simd.vdup(table_begin + table_count, 'int32', mask)
	108	+ if is_pow2:
	109	+ bs_mask = T.simd.vdup(block_size - 1, 'int32', mask)
	110	+ else:
	111	+ bs_vec = T.simd.vdup(block_size, 'int32', mask)
	112	+
	113	+ for vector in T.serial(T.ceildiv(index_count, vector_length)):
	114	+ global_index = T.simd.vci(index_begin + vector * vector_length, 'int32')
	115	+ x_vec = T.simd.vld(idx_ub[vector * vector_length])
	116	+ if is_pow2:
	117	+ blk_idx = T.simd.vshrs(x_vec, block_bits, mask)
	118	+ rem = T.simd.vand(x_vec, bs_mask, mask)
	119	+ else:
	120	+ blk_idx = T.simd.vdiv(x_vec, bs_vec, mask)
	121	+ rem = T.simd.vsub(x_vec, T.simd.vmul(blk_idx, bs_vec, mask), mask)
	122	+
	123	+ valid_index = T.simd.vcmps(global_index, k, mask, 'lt')
	124	+ valid_context = T.simd.vcmps(global_index, cl, mask, 'lt')
	125	+ in_table_lower = T.simd.vcmp(blk_idx, table_begin_vec, mask, 'ge')
	126	+ in_table_upper = T.simd.vcmp(blk_idx, table_end_vec, mask, 'lt')
	127	+ valid = T.simd.pand(valid_index, valid_context, mask)
	128	+ valid = T.simd.pand(valid, in_table_lower, mask)
	129	+ valid = T.simd.pand(valid, in_table_upper, mask)
	130	+
	131	+ local_blk_idx = T.simd.vsub(blk_idx, table_begin_vec, mask)
	132	+ local_blk_idx = T.simd.vsel(local_blk_idx, zero, valid)
	133	+ blk_val = T.simd.vgather2(
	134	+ blk_ub[0],
	135	+ T.reinterpret('uint32x64', local_blk_idx),
	136	+ valid,
	137	+ )
	138	+ if is_pow2:
	139	+ converted = T.simd.vadd(T.simd.vshls(blk_val, block_bits, mask), rem, mask)
	140	+ else:
	141	+ converted = T.simd.vadd(T.simd.vmul(blk_val, bs_vec, mask), rem, mask)
	142	+ previous = T.simd.vld(out_ub[vector * vector_length])
	143	+ T.simd.vsts(
	144	+ out_ub[vector * vector_length],
	145	+ T.simd.vsel(converted, previous, valid),
	146	+ mask,
	147	+ )
	148	+
	149	+ T.copy(
	150	+ out_ub[:index_count],
	151	+ topk_idxs[bid, sid, index_begin : index_begin + index_count],
	152	+ )
	153	+
	154	+ return _convert_kv_index_chunked
	155	+
	156	+ @T.prim_func
	157	+ def _convert_kv_index(
	158	+ topk_idxs: T.Tensor([b, s, k], T.int32), # type: ignore
	159	+ context_lens: T.Tensor([b, s], T.int32), # type: ignore
	160	+ block_tables: T.Tensor([b, n], T.int32), # type: ignore
	161	+ ):
	162	+ topk_flat = T.reshape(topk_idxs, (b * s * k,))
	163	+ cl_flat = T.reshape(context_lens, (b * s,))
	164	+ bt_flat = T.reshape(block_tables, (b * n,))
	165	+ ntask = b * s
	166	+ with T.Kernel(num_cores) as pid:
	167	+ idx_ub = T.alloc_shared((rows_per_batch * index_tile_size,), 'int32')
	168	+ blk_ub = T.alloc_shared((rows_per_batch * table_tile_size,), 'int32')
	169	+ T.annotate_buffer_versions({idx_ub: num_stages, blk_ub: num_stages})
	170	+
	171	+ t0 = (pid * ntask) // num_cores
	172	+ t1 = ((pid + 1) * ntask) // num_cores
	173	+ my = t1 - t0
	174	+
	175	+ for j in T.Pipelined(T.ceildiv(my, rows_per_batch), num_stages=num_stages):
	176	+ lo = t0 + j * rows_per_batch
	177	+ hi = T.min(lo + rows_per_batch, t1)
	178	+ ln = hi - lo
	179	+ if ln > 0:
	180	+ bid_lo = lo // s
	181	+
	182	+ if batched:
	183	+ bid_hi = (hi - 1) // s
	184	+ nrows = bid_hi - bid_lo + 1
	185	+
	186	+ T.copy(topk_flat[lo * k : hi * k], idx_ub[0 : ln * k])
	187	+ T.copy(bt_flat[bid_lo * n : bid_lo * n + nrows * n], blk_ub[0 : nrows * n])
	188	+ else:
	189	+ for i in T.serial(rows_per_batch):
	190	+ if i < ln:
	191	+ task = lo + i
	192	+ bid = task // s
	193	+ sid = task % s
	194	+ T.copy(topk_idxs[bid, sid, 0:k], idx_ub[i * index_tile_size : i * index_tile_size + k])
	195	+ T.copy(block_tables[bid, 0:n], blk_ub[i * table_tile_size : i * table_tile_size + n])
	196	+
	197	+ for i in T.serial(rows_per_batch):
	198	+ if i < ln:
	199	+ task = lo + i
	200	+ bid = task // s
	201	+ cl = cl_flat[task]
	202	+ if batched:
	203	+ row = (bid - bid_lo) * n
	204	+ else:
	205	+ row = i * table_tile_size
	206	+
	207	+ with T.SimdVF():
	208	+ mask = T.simd.pset(32)
	209	+ neg1 = T.simd.vdup(-1, 'int32', mask)
	210	+
	211	+ if is_pow2:
	212	+ bs_mask = T.simd.vdup(block_size - 1, 'int32', mask)
	213	+ for v in T.serial(index_tile_size // vector_length):
	214	+ gidx = T.simd.vci(v * vector_length, 'int32')
	215	+ x_vec = T.simd.vld(idx_ub[i * index_tile_size + v * vector_length])
	216	+ blk_idx = T.simd.vshrs(x_vec, block_bits, mask)
	217	+ rem = T.simd.vand(x_vec, bs_mask, mask)
	218	+
	219	+ mask_k = T.simd.vcmps(gidx, k, mask, 'lt')
	220	+ mask_cl = T.simd.vcmps(gidx, cl, mask, 'lt')
	221	+ mask_gather = T.simd.pand(mask_k, mask_cl, mask)
	222	+
	223	+ blk_val = T.simd.vgather2(
	224	+ blk_ub[row],
	225	+ T.reinterpret('uint32x64', blk_idx),
	226	+ mask_gather,
	227	+ )
	228	+ val = T.simd.vadd(T.simd.vshls(blk_val, block_bits, mask), rem, mask)
	229	+ res = T.simd.vsel(val, neg1, mask_gather)
	230	+ T.simd.vsts(idx_ub[i * index_tile_size + v * vector_length], res, mask)
	231	+ else:
	232	+ bs_vec = T.simd.vdup(block_size, 'int32', mask)
	233	+ for v in T.serial(index_tile_size // vector_length):
	234	+ gidx = T.simd.vci(v * vector_length, 'int32')
	235	+ x_vec = T.simd.vld(idx_ub[i * index_tile_size + v * vector_length])
	236	+ blk_idx = T.simd.vdiv(x_vec, bs_vec, mask)
	237	+ rem = T.simd.vsub(x_vec, T.simd.vmul(blk_idx, bs_vec, mask), mask)
	238	+
	239	+ mask_k = T.simd.vcmps(gidx, k, mask, 'lt')
	240	+ mask_cl = T.simd.vcmps(gidx, cl, mask, 'lt')
	241	+ mask_gather = T.simd.pand(mask_k, mask_cl, mask)
	242	+
	243	+ blk_val = T.simd.vgather2(
	244	+ blk_ub[row],
	245	+ T.reinterpret('uint32x64', blk_idx),
	246	+ mask_gather,
	247	+ )
	248	+ val = T.simd.vadd(T.simd.vmul(blk_val, bs_vec, mask), rem, mask)
	249	+ res = T.simd.vsel(val, neg1, mask_gather)
	250	+ T.simd.vsts(idx_ub[i * index_tile_size + v * vector_length], res, mask)
	251	+
	252	+ if batched:
	253	+ T.copy(idx_ub[0 : ln * k], topk_flat[lo * k : hi * k])
	254	+ else:
	255	+ for i in T.serial(rows_per_batch):
	256	+ if i < ln:
	257	+ task = lo + i
	258	+ bid = task // s
	259	+ sid = task % s
	260	+ T.copy(idx_ub[i * index_tile_size : i * index_tile_size + k], topk_idxs[bid, sid, 0:k])
	261	+
	262	+ return _convert_kv_index

Asrc/cann_ops_tilelang/attention/get_kv_index.py+183-0

@@ -0,0 +1,183 @@
	1	+import torch
	2	+
	3	+from cann_ops_tilelang.attention.get_kv_index_kernel import (
	4	+ _get_dflash_mainx_slot_mapping_tl_ascend,
	5	+ _get_fused_dflash_get_input_tl_ascend,
	6	+ _load_kv_index_tl_ascend,
	7	+)
	8	+from cann_ops_tilelang.config import require_ascend
	9	+
	10	+
	11	+def _validate_int32_npu_tensor(tensor: torch.Tensor, name: str, ndim: int \| None = None) -> None:
	12	+ if tensor.device.type != 'npu':
	13	+ raise ValueError(f'{name} must be on an NPU, got {tensor.device}.')
	14	+ if tensor.dtype != torch.int32:
	15	+ raise TypeError(f'{name} must have dtype torch.int32, got {tensor.dtype}.')
	16	+ if ndim is not None and tensor.ndim != ndim:
	17	+ raise ValueError(f'{name} must be {ndim}D, got shape {tuple(tensor.shape)}.')
	18	+
	19	+
	20	+def _validate_offset_len(tensor: torch.Tensor, name: str) -> None:
	21	+ _validate_int32_npu_tensor(tensor, name, ndim=2)
	22	+ if tensor.shape[1] != 2:
	23	+ raise ValueError(f'{name} must have shape [num_sequences, 2], got {tuple(tensor.shape)}.')
	24	+
	25	+
	26	+def get_load_kv_index_kernel(index_length: int, block_size: int, has_block_tables: bool) -> object:
	27	+ require_ascend()
	28	+ if index_length <= 0:
	29	+ raise ValueError('index_length must be positive.')
	30	+ if block_size < 0 or (has_block_tables and block_size == 0):
	31	+ raise ValueError('block_size must be positive when block tables are used, otherwise non-negative.')
	32	+ return _load_kv_index_tl_ascend(index_length, block_size, has_block_tables)
	33	+
	34	+
	35	+def get_kv_ca_index_prefill(kv_ca_index_offset_len: torch.Tensor, kv_ca_index_length: int) -> torch.Tensor:
	36	+ require_ascend()
	37	+ _validate_offset_len(kv_ca_index_offset_len, 'kv_ca_index_offset_len')
	38	+ if kv_ca_index_length <= 0:
	39	+ raise ValueError('kv_ca_index_length must be positive.')
	40	+ kv_ca_index = torch.empty(
	41	+ (kv_ca_index_offset_len.shape[0], kv_ca_index_length),
	42	+ dtype=torch.int32,
	43	+ device=kv_ca_index_offset_len.device,
	44	+ )
	45	+ kernel = get_load_kv_index_kernel(kv_ca_index_length, 0, has_block_tables=False)
	46	+ kernel(kv_ca_index_offset_len, kv_ca_index, None)
	47	+ return kv_ca_index
	48	+
	49	+
	50	+def get_kv_ca_index_decode(
	51	+ kv_ca_index_offset_len: torch.Tensor, kv_ca_index_length: int, token_per_block: int, block_tables: torch.Tensor
	52	+) -> torch.Tensor:
	53	+ require_ascend()
	54	+ _validate_offset_len(kv_ca_index_offset_len, 'kv_ca_index_offset_len')
	55	+ _validate_int32_npu_tensor(block_tables, 'block_tables')
	56	+ if block_tables.device != kv_ca_index_offset_len.device:
	57	+ raise ValueError('block_tables and kv_ca_index_offset_len must be on the same device.')
	58	+ if kv_ca_index_length <= 0 or token_per_block <= 0:
	59	+ raise ValueError('kv_ca_index_length and token_per_block must be positive.')
	60	+ kv_ca_index = torch.empty(
	61	+ (kv_ca_index_offset_len.shape[0], kv_ca_index_length),
	62	+ dtype=torch.int32,
	63	+ device=kv_ca_index_offset_len.device,
	64	+ )
	65	+ kernel = get_load_kv_index_kernel(kv_ca_index_length, token_per_block, has_block_tables=True)
	66	+ kernel(kv_ca_index_offset_len, kv_ca_index, block_tables.flatten())
	67	+ return kv_ca_index
	68	+
	69	+
	70	+def get_kv_swa_index_decode(
	71	+ kv_swa_index_offset_len: torch.Tensor, window_size: int, swa_block_size: int
	72	+) -> torch.Tensor:
	73	+ require_ascend()
	74	+ _validate_offset_len(kv_swa_index_offset_len, 'kv_swa_index_offset_len')
	75	+ if window_size <= 0 or swa_block_size <= 0:
	76	+ raise ValueError('window_size and swa_block_size must be positive.')
	77	+ kv_swa_index = torch.empty(
	78	+ (kv_swa_index_offset_len.shape[0], window_size),
	79	+ dtype=torch.int32,
	80	+ device=kv_swa_index_offset_len.device,
	81	+ )
	82	+ kernel = get_load_kv_index_kernel(window_size, swa_block_size, has_block_tables=False)
	83	+ kernel(kv_swa_index_offset_len, kv_swa_index, None)
	84	+ return kv_swa_index
	85	+
	86	+
	87	+def get_fused_dflash_get_input(
	88	+ slot_mapping_swa: torch.Tensor,
	89	+ positions: torch.Tensor,
	90	+ ftp: torch.Tensor,
	91	+ accepted_pos_real: torch.Tensor,
	92	+ swa_block_size: int,
	93	+ dflash_block_size: int,
	94	+ window_size: int,
	95	+ has_valid_token_mask: bool = False,
	96	+ has_gate_force_random: bool = False,
	97	+ input_ids_dtype: torch.dtype = torch.int32,
	98	+) -> tuple[
	99	+ torch.Tensor,
	100	+ torch.Tensor,
	101	+ torch.Tensor,
	102	+ torch.Tensor \| None,
	103	+ torch.Tensor \| None,
	104	+ torch.Tensor,
	105	+]:
	106	+ """Return slot mapping, positions, IDs, optional masks, and KV indices."""
	107	+ require_ascend()
	108	+ for name, tensor in (
	109	+ ('slot_mapping_swa', slot_mapping_swa),
	110	+ ('positions', positions),
	111	+ ('ftp', ftp),
	112	+ ('accepted_pos_real', accepted_pos_real),
	113	+ ):
	114	+ _validate_int32_npu_tensor(tensor, name, ndim=1)
	115	+ devices = {tensor.device for tensor in (slot_mapping_swa, positions, ftp, accepted_pos_real)}
	116	+ if len(devices) != 1:
	117	+ raise ValueError('all inputs must be on the same NPU device.')
	118	+ if not 0 < dflash_block_size <= 32:
	119	+ raise ValueError(f'dflash_block_size must be in [1, 32], got {dflash_block_size}.')
	120	+ if swa_block_size <= 0 or window_size <= 0:
	121	+ raise ValueError('swa_block_size and window_size must be positive.')
	122	+ real_cnt = accepted_pos_real.shape[0]
	123	+ if real_cnt > ftp.numel():
	124	+ raise ValueError('accepted_pos_real cannot contain more entries than ftp.')
	125	+ kv_len = window_size + dflash_block_size
	126	+ padded_kv_len = (kv_len + 63) // 64 * 64
	127	+ device = accepted_pos_real.device
	128	+ newx_slot_mapping_out = torch.empty((real_cnt, dflash_block_size), dtype=torch.int32, device=device)
	129	+ positions_out = torch.empty((real_cnt, dflash_block_size), dtype=torch.int32, device=device)
	130	+ input_ids_out = torch.empty((real_cnt, dflash_block_size), dtype=torch.int32, device=device)
	131	+ valid_token_mask_out = (
	132	+ torch.empty((real_cnt, dflash_block_size), dtype=torch.bool, device=device) if has_valid_token_mask else None
	133	+ )
	134	+ gate_force_random_out = (
	135	+ torch.empty((real_cnt, dflash_block_size), dtype=torch.bool, device=device) if has_gate_force_random else None
	136	+ )
	137	+ kv_index_out = torch.empty((real_cnt * dflash_block_size, padded_kv_len), dtype=torch.int32, device=device)
	138	+ kernel = _get_fused_dflash_get_input_tl_ascend(
	139	+ swa_block_size, padded_kv_len, dflash_block_size, has_valid_token_mask, has_gate_force_random
	140	+ )
	141	+ kernel(
	142	+ slot_mapping_swa,
	143	+ positions,
	144	+ ftp[:real_cnt],
	145	+ accepted_pos_real,
	146	+ newx_slot_mapping_out,
	147	+ positions_out,
	148	+ input_ids_out,
	149	+ valid_token_mask_out,
	150	+ gate_force_random_out,
	151	+ kv_index_out,
	152	+ kv_len,
	153	+ window_size,
	154	+ )
	155	+ valid_token_mask = valid_token_mask_out.view(-1) if has_valid_token_mask else None
	156	+ gate_force_random = gate_force_random_out.view(-1) if has_gate_force_random else None
	157	+ return (
	158	+ newx_slot_mapping_out.view(-1),
	159	+ positions_out.view(-1),
	160	+ input_ids_out.to(input_ids_dtype).view(-1),
	161	+ valid_token_mask,
	162	+ gate_force_random,
	163	+ kv_index_out,
	164	+ )
	165	+
	166	+
	167	+def get_dflash_mainx_slot_mapping(
	168	+ ftp: torch.Tensor, valid_offsets: torch.Tensor, slot_mapping_swa: torch.Tensor, real_cnt: int
	169	+) -> torch.Tensor:
	170	+ """Fused tilelang kernel for computing dflash_mainx_slot_mapping_store."""
	171	+ require_ascend()
	172	+ for name, tensor in (('ftp', ftp), ('valid_offsets', valid_offsets), ('slot_mapping_swa', slot_mapping_swa)):
	173	+ _validate_int32_npu_tensor(tensor, name, ndim=1)
	174	+ if not (ftp.device == valid_offsets.device == slot_mapping_swa.device):
	175	+ raise ValueError('ftp, valid_offsets, and slot_mapping_swa must be on the same device.')
	176	+ if not 0 <= real_cnt <= min(ftp.numel(), valid_offsets.numel()):
	177	+ raise ValueError('real_cnt must fit within ftp and valid_offsets.')
	178	+ total_q = slot_mapping_swa.shape[0]
	179	+ out = torch.empty(total_q, dtype=torch.int32, device=slot_mapping_swa.device)
	180	+ get_kernel = _get_dflash_mainx_slot_mapping_tl_ascend
	181	+ kernel = get_kernel()
	182	+ kernel(ftp[:real_cnt], valid_offsets[:real_cnt], slot_mapping_swa, out)
	183	+ return out

Asrc/cann_ops_tilelang/attention/get_kv_index_kernel.py+405-0

@@ -0,0 +1,405 @@
	1	+# Ruff's simplification rules do not model TileLang's compile-time branches.
	2	+# ruff: noqa: SIM102
	3	+
	4	+import tilelang
	5	+from tilelang.ascend import language as T
	6	+
	7	+from cann_ops_tilelang.config import get_num_vector_cores
	8	+
	9	+
	10	+@tilelang.jit(pass_configs={tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True})
	11	+def _load_kv_index_tl_ascend(
	12	+ index_length: int,
	13	+ block_size: int,
	14	+ has_block_tables: bool,
	15	+) -> None:
	16	+ n = T.dynamic('n')
	17	+ num_blocks = T.dynamic('num_blocks') if has_block_tables else 1
	18	+ num_cores = get_num_vector_cores()
	19	+ vector_length = 64
	20	+ index_tile_size = T.ceildiv(index_length, vector_length) * vector_length if index_length <= 8192 else 8192 * 2
	21	+ block_table_tile_size = T.ceildiv((index_tile_size // block_size + 1) if has_block_tables else 0, 8) * 8
	22	+ num_stages = 2
	23	+ is_pow2 = (block_size & (block_size - 1)) == 0
	24	+ assert not has_block_tables or is_pow2, 'swa or 2^k block size'
	25	+ if block_size > 0 and is_pow2:
	26	+ block_bits = block_size.bit_length() - 1
	27	+ inc_r = vector_length % block_size if block_size > 0 else 0
	28	+
	29	+ blocks_per_token = T.ceildiv(index_length, index_tile_size)
	30	+
	31	+ @T.prim_func
	32	+ def _load_kv_index(
	33	+ kv_index_offset_len: T.Tensor((n, 2), 'int32'), # type: ignore
	34	+ kv_index: T.Tensor((n, index_length), 'int32'), # type: ignore
	35	+ flatten_block_tables: T.Tensor((num_blocks), 'int32'), # type: ignore
	36	+ ):
	37	+ with T.Kernel(num_cores) as pid:
	38	+ index = T.alloc_shared((index_tile_size,), 'int32') # type: ignore
	39	+ T.annotate_buffer_versions({index: num_stages})
	40	+ if has_block_tables:
	41	+ blk_ub = T.alloc_shared((block_table_tile_size,), 'int32') # type: ignore
	42	+ T.annotate_buffer_versions({blk_ub: num_stages})
	43	+ for bid in T.Persistent([n * blocks_per_token], num_cores, pid, group_size=1, num_stages=num_stages):
	44	+ bx = bid // blocks_per_token
	45	+ i = bid % blocks_per_token
	46	+ copy_length = T.min(index_tile_size, index_length - i * index_tile_size)
	47	+ offset = kv_index_offset_len[bx, 0]
	48	+ valid_length = kv_index_offset_len[bx, 1]
	49	+
	50	+ if has_block_tables:
	51	+ valid_begin = i * index_tile_size
	52	+ valid_count = T.max(0, T.min(copy_length, valid_length - valid_begin))
	53	+ base_blk = T.alloc_var(T.int32, 0)
	54	+ slice_len = T.alloc_var(T.int32, 0)
	55	+ if valid_count > 0:
	56	+ base_blk = (offset + valid_begin) // block_size
	57	+ last_x = valid_begin + valid_count - 1
	58	+ slice_len = (offset + last_x) // block_size - base_blk + 1
	59	+ T.copy(flatten_block_tables[base_blk : base_blk + slice_len], blk_ub[:slice_len])
	60	+
	61	+ with T.SimdVF():
	62	+ mask = T.simd.pset(32)
	63	+ neg1 = T.simd.vdup(-1, 'int32', mask)
	64	+
	65	+ if has_block_tables:
	66	+ mask_const = T.simd.vdup(block_size - 1, 'int32', mask)
	67	+ inc_q = vector_length // block_size
	68	+ q_var = T.simd.alloc_var('int32')
	69	+ rem_var = T.simd.alloc_var('int32')
	70	+ pos0 = T.simd.vadds(T.simd.vci(i * index_tile_size, 'int32'), offset, mask)
	71	+ q_var = T.simd.vshrs(pos0, block_bits, mask)
	72	+ rem_var = T.simd.vand(pos0, mask_const, mask)
	73	+ q_var = T.simd.vadds(q_var, -base_blk, mask)
	74	+ for v in T.serial(index_tile_size // vector_length):
	75	+ gidx = T.simd.vci(i * index_tile_size + v * vector_length, 'int32')
	76	+ mask_valid = T.simd.vcmps(q_var, slice_len, mask, 'lt')
	77	+ blk = T.simd.vgather2(blk_ub[0], T.reinterpret('uint32x64', q_var), mask_valid)
	78	+ val = T.simd.vadd(T.simd.vshls(blk, block_bits, mask), rem_var, mask)
	79	+ res = T.simd.vsel(val, neg1, T.simd.vcmps(gidx, valid_length, mask, 'lt'))
	80	+ T.simd.vsts(index[v * vector_length], res, mask)
	81	+
	82	+ q_var = T.simd.vadds(q_var, inc_q, mask)
	83	+ if inc_r > 0:
	84	+ rem_var = T.simd.vadds(rem_var, inc_r, mask)
	85	+ ge = T.simd.vcmps(rem_var, block_size, mask, 'ge')
	86	+ rem_var = T.simd.vand(rem_var, mask_const, mask)
	87	+ q_var = T.simd.vsel(T.simd.vadds(q_var, 1, mask), q_var, ge)
	88	+ elif block_size == 0:
	89	+ for v in T.serial(index_tile_size // vector_length):
	90	+ gidx = T.simd.vci(offset + i * index_tile_size + v * vector_length, 'int32')
	91	+ res = T.simd.vsel(gidx, neg1, T.simd.vcmps(gidx, offset + valid_length, mask, 'lt'))
	92	+ T.simd.vsts(index[v * vector_length], res, mask)
	93	+ else:
	94	+ off_base = offset // block_size * block_size
	95	+ rem_var = T.simd.alloc_var('int32')
	96	+ pos0 = T.simd.vadds(T.simd.vci(i * index_tile_size, 'int32'), offset, mask)
	97	+ if is_pow2:
	98	+ mask_const = T.simd.vdup(block_size - 1, 'int32', mask)
	99	+ rem_var = T.simd.vand(pos0, mask_const, mask)
	100	+ else:
	101	+ bs_vec = T.simd.vdup(block_size, 'int32', mask)
	102	+ pos0_q = T.simd.vdiv(pos0, bs_vec, mask)
	103	+ rem_var = T.simd.vsub(pos0, T.simd.vmul(pos0_q, bs_vec, mask), mask)
	104	+ rem_var = T.simd.vadds(rem_var, off_base, mask)
	105	+ for v in T.serial(index_tile_size // vector_length):
	106	+ gidx = T.simd.vci(i * index_tile_size + v * vector_length, 'int32')
	107	+ res = T.simd.vsel(rem_var, neg1, T.simd.vcmps(gidx, valid_length, mask, 'lt'))
	108	+ T.simd.vsts(index[v * vector_length], res, mask)
	109	+ if inc_r > 0:
	110	+ rem_var = T.simd.vadds(rem_var, inc_r, mask)
	111	+ ge = T.simd.vcmps(rem_var, block_size + off_base, mask, 'ge')
	112	+ rem_var = T.simd.vsel(T.simd.vadds(rem_var, -block_size, mask), rem_var, ge)
	113	+ T.copy(index[:copy_length], kv_index[bx, i * index_tile_size : i * index_tile_size + copy_length])
	114	+
	115	+ return _load_kv_index
	116	+
	117	+
	118	+def get_ascend_disable_smem_reuse_config() -> dict:
	119	+ pass_configs = {tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True}
	120	+ key = getattr(tilelang.PassConfigKey, 'TL_DISABLE_SHARED_MEMORY_REUSE', None)
	121	+ if key is not None:
	122	+ pass_configs[key] = True
	123	+ return pass_configs
	124	+
	125	+
	126	+@tilelang.jit(
	127	+ out_idx=None,
	128	+ pass_configs=get_ascend_disable_smem_reuse_config(),
	129	+)
	130	+def _get_fused_dflash_get_input_tl_ascend(
	131	+ swa_block_size: int,
	132	+ padded_kv_len: int,
	133	+ dflash_block_size: int,
	134	+ has_valid_token_mask: bool,
	135	+ has_gate_force_random: bool,
	136	+) -> None:
	137	+ real_cnt = T.dynamic('real_cnt')
	138	+ total_q = T.dynamic('total_q')
	139	+ kv_index_out_shape0 = T.dynamic('kv_index_out_shape0')
	140	+ num_cores = get_num_vector_cores()
	141	+ vector_length = 64
	142	+ bs = dflash_block_size
	143	+ compact_kv = bs * padded_kv_len * 4 > 64 * 1024
	144	+ num_stages = 1 if compact_kv and padded_kv_len * 4 > 96 * 1024 else 2
	145	+ # requests per persistent iteration: keep the kv chunk UB under 64KB (x2 for double buffering)
	146	+ kv_elems_per_request = padded_kv_len if compact_kv else bs * padded_kv_len
	147	+ requests_per_group = max(1, min(8, (64 * 1024) // (kv_elems_per_request * 4)))
	148	+ rows = requests_per_group * bs
	149	+ token_span = 1024 # token-span staging window (elements) for the per-request scalar gather
	150	+ is_pow2 = (swa_block_size & (swa_block_size - 1)) == 0
	151	+
	152	+ @T.prim_func
	153	+ def _get_fused_dflash_get_input(
	154	+ slot_mapping_swa: T.Tensor((total_q,), 'int32'), # type: ignore
	155	+ positions: T.Tensor((total_q,), 'int32'), # type: ignore
	156	+ ftp: T.Tensor((real_cnt,), 'int32'), # type: ignore
	157	+ accepted_pos_real: T.Tensor((real_cnt,), 'int32'), # type: ignore
	158	+ newx_slot_mapping_out: T.Tensor((real_cnt, bs), 'int32'), # type: ignore
	159	+ positions_out: T.Tensor((real_cnt, bs), 'int32'), # type: ignore
	160	+ input_ids_out: T.Tensor((real_cnt, bs), 'int32'), # type: ignore
	161	+ valid_token_mask_out: T.Tensor((real_cnt, bs), 'bool'), # type: ignore
	162	+ gate_force_random_out: T.Tensor((real_cnt, bs), 'bool'), # type: ignore
	163	+ kv_index_out: T.Tensor((kv_index_out_shape0, padded_kv_len), 'int32'), # type: ignore
	164	+ kv_len: T.int32,
	165	+ window_size: T.int32,
	166	+ ):
	167	+ total_rows = real_cnt * bs
	168	+ ids_1d = T.reshape(input_ids_out, (total_rows,))
	169	+ vtm_1d = T.reshape(valid_token_mask_out, (total_rows,))
	170	+ gfr_1d = T.reshape(gate_force_random_out, (total_rows,))
	171	+ newx_1d = T.reshape(newx_slot_mapping_out, (total_rows,))
	172	+ pos_1d = T.reshape(positions_out, (total_rows,))
	173	+ kvout_1d = T.reshape(kv_index_out, (kv_index_out_shape0 * padded_kv_len,))
	174	+
	175	+ with T.Kernel(num_cores) as pid:
	176	+ # constant-fill sources: vector-filled int32 patterns (z: 0, t: 0x01010101)
	177	+ z_ub = T.alloc_shared((256,), 'int32') # type: ignore
	178	+ t32_ub = T.alloc_shared((256,), 'int32')
	179	+ z8 = T.view(z_ub, (1024,), T.bool) # type: ignore # bool view for gfr fill
	180	+ t_ub = T.view(t32_ub, (1024,), T.bool) # type: ignore # bool view for vtm fill
	181	+ with T.SimdVF():
	182	+ mask = T.simd.pset(32)
	183	+ z_vec = T.simd.vdup(0, 'int32', mask)
	184	+ t_vec = T.simd.vdup(0x01010101, 'int32', mask)
	185	+ for c in T.serial(4):
	186	+ T.simd.vsts(z_ub[c * 64], z_vec, mask)
	187	+ T.simd.vsts(t32_ub[c * 64], t_vec, mask)
	188	+
	189	+ ftp_ub = T.alloc_shared((requests_per_group + 1,), 'int32') # type: ignore
	190	+ ap_ub = T.alloc_shared((64,), 'int32') # type: ignore
	191	+ so_ub = T.alloc_shared((64,), 'int32') # type: ignore
	192	+ start_ub = T.alloc_shared((64,), 'int32') # type: ignore
	193	+ pos_span = T.alloc_shared((token_span,), 'int32') # type: ignore
	194	+ slot_span = T.alloc_shared((token_span,), 'int32') # type: ignore
	195	+ kv2 = T.alloc_shared((requests_per_group * kv_elems_per_request,), 'int32') # type: ignore
	196	+ nx_ub = T.alloc_shared(((rows + 63) // 64 * 64,), 'int32') # type: ignore
	197	+ pos_out_ub = T.alloc_shared(((rows + 63) // 64 * 64,), 'int32') # type: ignore
	198	+ T.annotate_buffer_versions(
	199	+ {
	200	+ ftp_ub: num_stages,
	201	+ ap_ub: num_stages,
	202	+ so_ub: num_stages,
	203	+ start_ub: num_stages,
	204	+ pos_span: num_stages,
	205	+ slot_span: num_stages,
	206	+ kv2: num_stages,
	207	+ nx_ub: num_stages,
	208	+ pos_out_ub: num_stages,
	209	+ }
	210	+ )
	211	+
	212	+ for blk in T.Persistent(
	213	+ [T.ceildiv(real_cnt, requests_per_group)], num_cores, pid, group_size=1, num_stages=num_stages
	214	+ ):
	215	+ g0 = blk * requests_per_group
	216	+ glen = T.min(requests_per_group, real_cnt - g0)
	217	+ rows_g = glen * bs
	218	+
	219	+ # 1) request metadata: ftp (+1 boundary), ap
	220	+ load_len = T.min(requests_per_group + 1, real_cnt - g0)
	221	+ T.copy(ftp[g0 : g0 + load_len], ftp_ub[:load_len])
	222	+ T.copy(accepted_pos_real[g0 : g0 + glen], ap_ub[:glen])
	223	+ s0 = ftp_ub[0]
	224	+ s1 = T.Select(g0 + glen < real_cnt, ftp_ub[glen], total_q)
	225	+
	226	+ # 2) per-request so/start via span windows (few large MTE2 instead of 2G tiny ones)
	227	+ nwin = T.ceildiv(s1 - s0, token_span)
	228	+ for w in T.serial(nwin):
	229	+ w0 = s0 + w * token_span
	230	+ wlen = T.min(token_span, s1 - w0)
	231	+ T.copy(positions[w0 : w0 + wlen], pos_span[:wlen])
	232	+ T.copy(slot_mapping_swa[w0 : w0 + wlen], slot_span[:wlen])
	233	+ for r in T.serial(requests_per_group):
	234	+ if r < glen:
	235	+ if ftp_ub[r] >= w0:
	236	+ if ftp_ub[r] < w0 + wlen:
	237	+ off = ftp_ub[r] - w0
	238	+ so_ub[r] = slot_span[off] - pos_span[off] % swa_block_size
	239	+ start_ub[r] = ap_ub[r] - (window_size - 1)
	240	+
	241	+ # 3) kv rows: compute one row per request, replicate to bs rows, one contiguous write
	242	+ with T.SimdVF():
	243	+ mask = T.simd.pset(32)
	244	+ neg1 = T.simd.vdup(-1, 'int32', mask)
	245	+ swa_vec = T.simd.vdup(swa_block_size, 'int32', mask)
	246	+ swa_mask_vec = T.simd.vdup(swa_block_size - 1, 'int32', mask)
	247	+ kv_len_vec = T.simd.vdup(kv_len, 'int32', mask)
	248	+ col = T.simd.alloc_var('int32')
	249	+ for r in T.serial(requests_per_group):
	250	+ if r < glen:
	251	+ so_vec = T.simd.vdup(so_ub[r], 'int32', mask)
	252	+ start_vec = T.simd.vdup(start_ub[r], 'int32', mask)
	253	+ col = T.simd.vci(0, 'int32')
	254	+ for v in T.serial(T.ceildiv(padded_kv_len, vector_length)):
	255	+ pos = T.simd.vadd(col, start_vec, mask)
	256	+ if is_pow2:
	257	+ rem = T.simd.vand(pos, swa_mask_vec, mask)
	258	+ else:
	259	+ pos_q = T.simd.vdiv(pos, swa_vec, mask)
	260	+ rem = T.simd.vsub(pos, T.simd.vmul(pos_q, swa_vec, mask), mask)
	261	+ val = T.simd.vadd(rem, so_vec, mask)
	262	+ # static-operand compares only (avoid scalar->vector syncs)
	263	+ mask_len = T.simd.vcmps(T.simd.vsub(col, kv_len_vec, mask), 0, mask, 'lt')
	264	+ mask_ge = T.simd.vcmps(pos, 0, mask, 'ge')
	265	+ res = T.simd.vsel(val, neg1, T.simd.pand(mask_len, mask_ge, mask))
	266	+ if compact_kv:
	267	+ T.simd.vsts(kv2[r * padded_kv_len + v * vector_length], res, mask)
	268	+ else:
	269	+ for bx in T.serial(bs):
	270	+ T.simd.vsts(kv2[(r * bs + bx) * padded_kv_len + v * vector_length], res, mask)
	271	+ col = T.simd.vadds(col, 64, mask)
	272	+
	273	+ # 4) newx / positions: gather pass over the group's rows, 64-lane chunks
	274	+ with T.SimdVF():
	275	+ mask = T.simd.pset(32)
	276	+ swa_vec = T.simd.vdup(swa_block_size, 'int32', mask)
	277	+ swa_mask_vec = T.simd.vdup(swa_block_size - 1, 'int32', mask)
	278	+ bs_vec = T.simd.vdup(bs, 'int32', mask)
	279	+ for ch in T.serial((rows + 63) // 64):
	280	+ lane = T.simd.vci(ch * 64, 'int32')
	281	+ r_lane = T.simd.vdiv(lane, bs_vec, mask)
	282	+ r_lane = T.simd.vmins(r_lane, requests_per_group - 1, mask)
	283	+ bx_lane = T.simd.vsub(lane, T.simd.vmul(r_lane, bs_vec, mask), mask)
	284	+ r_idx = T.reinterpret('uint32x64', r_lane)
	285	+ ap_g = T.simd.vgather2(ap_ub[0], r_idx, mask)
	286	+ so_g = T.simd.vgather2(so_ub[0], r_idx, mask)
	287	+ posv = T.simd.vadd(T.simd.vadds(bx_lane, 1, mask), ap_g, mask)
	288	+ if is_pow2:
	289	+ rem2 = T.simd.vand(posv, swa_mask_vec, mask)
	290	+ else:
	291	+ q2 = T.simd.vdiv(posv, swa_vec, mask)
	292	+ rem2 = T.simd.vsub(posv, T.simd.vmul(q2, swa_vec, mask), mask)
	293	+ newx = T.simd.vadd(rem2, so_g, mask)
	294	+ mact = T.simd.vcmps(lane, rows, mask, 'lt')
	295	+ T.simd.vsts(nx_ub[ch * 64], newx, mact)
	296	+ T.simd.vsts(pos_out_ub[ch * 64], posv, mact)
	297	+ T.copy(nx_ub[:rows_g], newx_1d[g0 * bs : g0 * bs + rows_g])
	298	+ T.copy(pos_out_ub[:rows_g], pos_1d[g0 * bs : g0 * bs + rows_g])
	299	+
	300	+ # 5) constant fills for the group's rows
	301	+ off1 = g0 * bs
	302	+ T.copy(z_ub[:rows_g], ids_1d[off1 : off1 + rows_g])
	303	+ if has_valid_token_mask:
	304	+ T.copy(t_ub[:rows_g], vtm_1d[off1 : off1 + rows_g])
	305	+ if has_gate_force_random:
	306	+ T.copy(z8[:rows_g], gfr_1d[off1 : off1 + rows_g])
	307	+
	308	+ # 6) kv_index: keep the one-copy path for normal shapes and
	309	+ # replicate from one staged row when the expanded group would
	310	+ # exceed UB capacity.
	311	+ if compact_kv:
	312	+ for r in T.serial(requests_per_group):
	313	+ if r < glen:
	314	+ for bx in T.serial(bs):
	315	+ out_row = (g0 + r) * bs + bx
	316	+ T.copy(
	317	+ kv2[r * padded_kv_len : (r + 1) * padded_kv_len],
	318	+ kvout_1d[out_row * padded_kv_len : (out_row + 1) * padded_kv_len],
	319	+ )
	320	+ else:
	321	+ T.copy(
	322	+ kv2[: rows_g * padded_kv_len],
	323	+ kvout_1d[g0 * bs * padded_kv_len : g0 * bs * padded_kv_len + rows_g * padded_kv_len],
	324	+ )
	325	+
	326	+ return _get_fused_dflash_get_input
	327	+
	328	+
	329	+@tilelang.jit(
	330	+ out_idx=None,
	331	+ pass_configs={
	332	+ tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True,
	333	+ },
	334	+)
	335	+def _get_dflash_mainx_slot_mapping_tl_ascend() -> None:
	336	+ real_cnt = T.dynamic('real_cnt')
	337	+ total_q = T.dynamic('total_q')
	338	+ num_cores = get_num_vector_cores()
	339	+ num_stages = 4
	340	+ threads = 32
	341	+ requests_per_block = 16
	342	+ slot_tile_size = requests_per_block * 64
	343	+
	344	+ @T.prim_func
	345	+ def _get_dflash_mainx_slot_mapping(
	346	+ ftp: T.Tensor((real_cnt,), 'int32'), # type: ignore
	347	+ valid_offsets: T.Tensor((real_cnt,), 'int32'), # type: ignore
	348	+ slot_mapping_swa: T.Tensor((total_q,), 'int32'), # type: ignore
	349	+ out: T.Tensor((total_q,), 'int32'), # type: ignore
	350	+ ):
	351	+ with T.Kernel(num_cores) as pid:
	352	+ valid_ub = T.alloc_shared((requests_per_block,), 'int32') # type: ignore
	353	+ ftp_ub = T.alloc_shared((requests_per_block,), 'int32') # type: ignore
	354	+ out_ub = T.alloc_shared((slot_tile_size,), 'int32') # type: ignore
	355	+ T.annotate_buffer_versions({ftp_ub: num_stages, valid_ub: num_stages, out_ub: num_stages})
	356	+
	357	+ for bid in T.Persistent([T.ceildiv(real_cnt, requests_per_block)], num_cores, pid):
	358	+ bx_start = bid * requests_per_block
	359	+ valid_cnt = T.min(requests_per_block, real_cnt - bx_start)
	360	+ total_start = ftp[bx_start]
	361	+ total_end = T.Select(bx_start + valid_cnt < real_cnt, ftp[bx_start + valid_cnt], total_q)
	362	+ range_len = total_end - total_start
	363	+
	364	+ T.copy(ftp[bx_start : bx_start + valid_cnt], ftp_ub[:valid_cnt])
	365	+ T.copy(valid_offsets[bx_start : bx_start + valid_cnt], valid_ub[:valid_cnt])
	366	+ if range_len <= slot_tile_size:
	367	+ T.copy(slot_mapping_swa[total_start:total_end], out_ub[:range_len])
	368	+
	369	+ with T.SimtVF(threads=requests_per_block * threads):
	370	+ for r, j in T.Parallel(requests_per_block, threads):
	371	+ if r < valid_cnt:
	372	+ local_start = ftp_ub[r]
	373	+ local_end = T.Select(r + 1 < valid_cnt, ftp_ub[r + 1], total_end)
	374	+ valid_end = local_start + valid_ub[r] + 1
	375	+ valid_cut = T.min(valid_end, local_end)
	376	+ for j0 in T.serial(valid_cut, local_end, threads):
	377	+ pos = j0 + j
	378	+ if pos < local_end:
	379	+ out_ub[pos - total_start] = -1
	380	+
	381	+ T.copy(out_ub[:range_len], out[total_start:total_end])
	382	+ else:
	383	+ for chunk_id in T.serial(T.ceildiv(range_len, slot_tile_size)):
	384	+ chunk_start = total_start + chunk_id * slot_tile_size
	385	+ chunk_len = T.min(slot_tile_size, total_end - chunk_start)
	386	+ chunk_end = chunk_start + chunk_len
	387	+ T.copy(slot_mapping_swa[chunk_start:chunk_end], out_ub[:chunk_len])
	388	+
	389	+ with T.SimtVF(threads=requests_per_block * threads):
	390	+ for r, j in T.Parallel(requests_per_block, threads):
	391	+ if r < valid_cnt:
	392	+ local_start = ftp_ub[r]
	393	+ local_end = T.Select(r + 1 < valid_cnt, ftp_ub[r + 1], total_end)
	394	+ valid_end = local_start + valid_ub[r] + 1
	395	+ valid_cut = T.min(valid_end, local_end)
	396	+ invalid_start = T.max(valid_cut, chunk_start)
	397	+ invalid_end = T.min(local_end, chunk_end)
	398	+ for j0 in T.serial(invalid_start, invalid_end, threads):
	399	+ pos = j0 + j
	400	+ if pos < invalid_end:
	401	+ out_ub[pos - chunk_start] = -1
	402	+
	403	+ T.copy(out_ub[:chunk_len], out[chunk_start:chunk_end])
	404	+
	405	+ return _get_dflash_mainx_slot_mapping

Asrc/cann_ops_tilelang/attention/kv_index_cat.py+47-0

@@ -0,0 +1,47 @@
	1	+import torch
	2	+
	3	+from cann_ops_tilelang.attention.kv_index_cat_kernel import _kv_index_cat_tl_ascend
	4	+from cann_ops_tilelang.config import require_ascend
	5	+
	6	+
	7	+def kv_index_cat(kv_index: torch.Tensor, kv_ca_index: torch.Tensor, total_k_swa: int, total_k: int) -> torch.Tensor:
	8	+ """
	9	+ Args:
	10	+ kv_index: [total_q, window_size], int32
	11	+ kv_ca_index: [total_q, topk_tokens], int32
	12	+ total_k_swa: scalar int32
	13	+ total_k: scalar int32
	14	+
	15	+ Returns:
	16	+ out: [total_q, window_size + topk_tokens], int32
	17	+ """
	18	+ require_ascend()
	19	+ assert kv_index.ndim == 2, 'kv_index must be a 2D tensor'
	20	+ assert kv_ca_index.ndim == 2, 'kv_ca_index must be a 2D tensor'
	21	+ total_q, window_size = kv_index.shape
	22	+ _, topk_tokens = kv_ca_index.shape
	23	+ assert total_q > 0, 'total_q must be positive'
	24	+ assert window_size > 0, 'window_size must be positive'
	25	+ assert topk_tokens > 0, 'topk_tokens must be positive'
	26	+ assert kv_ca_index.shape[0] == total_q, 'both inputs must have the same total_q'
	27	+ assert kv_index.is_contiguous(), 'kv_index must be contiguous'
	28	+ assert kv_ca_index.is_contiguous(), 'kv_ca_index must be contiguous'
	29	+ assert kv_index.dtype == torch.int32, 'kv_index must use int32'
	30	+ assert kv_ca_index.dtype == torch.int32, 'kv_ca_index must use int32'
	31	+ assert kv_index.device == kv_ca_index.device, 'both inputs must be on the same device'
	32	+ if kv_index.device.type != 'npu':
	33	+ raise ValueError(f'input tensors must be on an NPU, got {kv_index.device}.')
	34	+ assert kv_index.stride(0) % 4 == 0, 'window_size must be aligned to 4 elements'
	35	+ out_shape = window_size + topk_tokens
	36	+ out_stride = (out_shape + 3) // 4 * 4
	37	+ out = torch.empty((total_q, out_stride), dtype=kv_index.dtype, device=kv_index.device)[:, :out_shape]
	38	+ kernel = _kv_index_cat_tl_ascend('int32', window_size, topk_tokens, out_stride)
	39	+ kernel(
	40	+ kv_index,
	41	+ kv_ca_index,
	42	+ out[:, :window_size],
	43	+ out[:, window_size : window_size + topk_tokens],
	44	+ total_k_swa,
	45	+ total_k,
	46	+ )
	47	+ return out

Asrc/cann_ops_tilelang/attention/kv_index_cat_kernel.py+128-0

@@ -0,0 +1,128 @@
	1	+import tilelang
	2	+from tilelang.ascend import language as T
	3	+
	4	+from cann_ops_tilelang.config import get_num_vector_cores
	5	+
	6	+
	7	+@tilelang.jit(
	8	+ out_idx=None,
	9	+ pass_configs={
	10	+ tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True,
	11	+ tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
	12	+ tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: False,
	13	+ },
	14	+)
	15	+def _kv_index_cat_tl_ascend(
	16	+ dtype: str,
	17	+ window_size: int,
	18	+ topk_tokens: int,
	19	+ out_stride: int,
	20	+) -> None:
	21	+ total_q = T.dynamic('total_q')
	22	+
	23	+ if max(window_size, topk_tokens) <= 128:
	24	+ block_size = 128
	25	+ rows_per_task = 64
	26	+ num_stages = 2
	27	+ else:
	28	+ block_size = 6336
	29	+ rows_per_task = 5
	30	+ num_stages = 2
	31	+ vec_len = 64
	32	+ blocks_l = (window_size + block_size - 1) // block_size
	33	+ blocks_r = (topk_tokens + block_size - 1) // block_size
	34	+ num_cores = get_num_vector_cores()
	35	+
	36	+ @T.prim_func
	37	+ def _kv_index_cat(
	38	+ kv_index: T.Tensor((total_q, window_size), dtype), # type: ignore
	39	+ kv_ca_index: T.Tensor((total_q, topk_tokens), dtype), # type: ignore
	40	+ kv_index_out: T.StridedTensor((total_q, window_size), (out_stride, 1), dtype), # type: ignore
	41	+ kv_ca_index_out: T.StridedTensor((total_q, topk_tokens), (out_stride, 1), dtype), # type: ignore
	42	+ total_k_swa: T.int32,
	43	+ total_k: T.int32,
	44	+ ) -> None:
	45	+ with T.Kernel(num_cores) as pid:
	46	+ data_ub = T.alloc_shared((num_stages, rows_per_task, block_size), dtype)
	47	+
	48	+ for stage in T.serial(num_stages):
	49	+ T.ascend_set_flag('MTE3_MTE2', stage)
	50	+
	51	+ for bid_l in T.Persistent(
	52	+ [T.ceildiv(total_q, rows_per_task) * blocks_l],
	53	+ num_cores,
	54	+ pid,
	55	+ group_size=1,
	56	+ num_stages=0,
	57	+ ):
	58	+ row_l = (bid_l // blocks_l) * rows_per_task
	59	+ block_l = bid_l % blocks_l
	60	+ offset_l = block_l * block_size
	61	+ copy_rows_l = T.min(rows_per_task, total_q - row_l)
	62	+ copy_length_l = T.min(block_size, window_size - offset_l)
	63	+ slot_l = (bid_l // num_cores) % num_stages
	64	+
	65	+ T.ascend_wait_flag('MTE3_MTE2', slot_l)
	66	+ T.copy(
	67	+ kv_index[row_l : row_l + copy_rows_l, offset_l : offset_l + copy_length_l],
	68	+ data_ub[slot_l, :copy_rows_l, :copy_length_l],
	69	+ l2_cache_ctrl='NOTALLOC_KEEP',
	70	+ )
	71	+ T.ascend_set_flag('MTE2_MTE3', slot_l)
	72	+ T.ascend_wait_flag('MTE2_MTE3', slot_l)
	73	+ T.copy(
	74	+ data_ub[slot_l, :copy_rows_l, :copy_length_l],
	75	+ kv_index_out[row_l : row_l + copy_rows_l, offset_l : offset_l + copy_length_l],
	76	+ l2_cache_ctrl='NOTALLOC_KEEP',
	77	+ )
	78	+ T.ascend_set_flag('MTE3_MTE2', slot_l)
	79	+
	80	+ for stage in T.serial(num_stages):
	81	+ T.ascend_wait_flag('MTE3_MTE2', stage)
	82	+ T.ascend_set_flag('MTE3_MTE2', stage)
	83	+
	84	+ for bid_r in T.Persistent(
	85	+ [T.ceildiv(total_q, rows_per_task) * blocks_r],
	86	+ num_cores,
	87	+ pid,
	88	+ group_size=1,
	89	+ num_stages=0,
	90	+ ):
	91	+ row_r = (bid_r // blocks_r) * rows_per_task
	92	+ block_r = bid_r % blocks_r
	93	+ offset_r = block_r * block_size
	94	+ copy_rows_r = T.min(rows_per_task, total_q - row_r)
	95	+ copy_length_r = T.min(block_size, topk_tokens - offset_r)
	96	+ slot_r = (bid_r // num_cores) % num_stages
	97	+
	98	+ T.ascend_wait_flag('MTE3_MTE2', slot_r)
	99	+ T.copy(
	100	+ kv_ca_index[row_r : row_r + copy_rows_r, offset_r : offset_r + copy_length_r],
	101	+ data_ub[slot_r, :copy_rows_r, :copy_length_r],
	102	+ l2_cache_ctrl='NOTALLOC_KEEP',
	103	+ )
	104	+ T.ascend_set_flag('MTE2_V', slot_r)
	105	+ T.ascend_wait_flag('MTE2_V', slot_r)
	106	+
	107	+ with T.SimdVF():
	108	+ mask = T.simd.pset(32)
	109	+ for tile_row in T.serial(copy_rows_r):
	110	+ for v in T.serial(T.ceildiv(copy_length_r, vec_len)):
	111	+ value = T.simd.vld(data_ub[slot_r, tile_row, v * vec_len])
	112	+ value_added = T.simd.vadds(value, total_k + total_k_swa, mask)
	113	+ value_out = T.simd.vsel(value, value_added, T.simd.vcmps(value, -1, mask, 'eq'))
	114	+ T.simd.vsts(data_ub[slot_r, tile_row, v * vec_len], value_out, mask)
	115	+
	116	+ T.ascend_set_flag('V_MTE3', slot_r)
	117	+ T.ascend_wait_flag('V_MTE3', slot_r)
	118	+ T.copy(
	119	+ data_ub[slot_r, :copy_rows_r, :copy_length_r],
	120	+ kv_ca_index_out[row_r : row_r + copy_rows_r, offset_r : offset_r + copy_length_r],
	121	+ l2_cache_ctrl='NOTALLOC_KEEP',
	122	+ )
	123	+ T.ascend_set_flag('MTE3_MTE2', slot_r)
	124	+
	125	+ for stage in T.serial(num_stages):
	126	+ T.ascend_wait_flag('MTE3_MTE2', stage)
	127	+
	128	+ return _kv_index_cat

Asrc/cann_ops_tilelang/attention/prepare_cp_mirror.py+136-0

@@ -0,0 +1,136 @@
	1	+"""Public Ascend interface for mirror context-parallel metadata."""
	2	+
	3	+import torch
	4	+
	5	+from cann_ops_tilelang.config import require_ascend
	6	+
	7	+from .prepare_cp_mirror_kernel import _precompute_tl_ascend, _scatter_tl_ascend
	8	+
	9	+
	10	+def prepare_for_mirror_cp_tilelang(
	11	+ cu_seqlens_q: torch.Tensor,
	12	+ cu_seqlens_k: torch.Tensor,
	13	+ cp_size: int,
	14	+ cp_rank: int,
	15	+ compress_ratio: int,
	16	+ total_q: int,
	17	+ output_len: int,
	18	+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
	19	+ """
	20	+ TileLang Ascend implementation of prepare_for_mirror_cp.
	21	+ This function is sync-free: it does not call .item() or perform a
	22	+ device-to-host synchronization. The caller must provide pre-computed sizes.
	23	+
	24	+ Args:
	25	+ cu_seqlens_q: [batch_size + 1], contiguous int32 device tensor
	26	+ cu_seqlens_k: [batch_size + 1], contiguous int32 device tensor
	27	+ cp_size: context parallelism size
	28	+ cp_rank: current rank
	29	+ compress_ratio: KV compression ratio
	30	+ total_q: total number of query tokens (= cu_seqlens_q[-1], pre-synced to CPU)
	31	+ output_len: total output length for cp_rank (= sum of 2*partial_seqlen, pre-synced to CPU)
	32	+
	33	+ Returns:
	34	+ mirror_cp_index_tensor,
	35	+ mirror_cp_reverse_index_tensor,
	36	+ ks_tensor,
	37	+ ke_tensor,
	38	+ cu_seqlens_q_cp_tensor,
	39	+ cu_seqlens_k_cp_tensor
	40	+ """
	41	+ assert cu_seqlens_q.ndim == 1, 'cu_seqlens_q must be a 1D tensor'
	42	+ assert cu_seqlens_k.ndim == 1, 'cu_seqlens_k must be a 1D tensor'
	43	+ assert cu_seqlens_q.dtype == torch.int32, 'cu_seqlens_q must use int32'
	44	+ assert cu_seqlens_k.dtype == torch.int32, 'cu_seqlens_k must use int32'
	45	+ assert cu_seqlens_q.is_contiguous(), 'cu_seqlens_q must be contiguous'
	46	+ assert cu_seqlens_k.is_contiguous(), 'cu_seqlens_k must be contiguous'
	47	+ assert cu_seqlens_q.device == cu_seqlens_k.device, 'both inputs must be on the same device'
	48	+ assert cp_size > 0, 'cp_size must be positive'
	49	+ assert 0 <= cp_rank < cp_size, 'cp_rank must be in [0, cp_size)'
	50	+ assert compress_ratio > 0, 'compress_ratio must be positive'
	51	+ assert total_q >= 0, 'total_q must be non-negative'
	52	+ assert output_len >= 0, 'output_len must be non-negative'
	53	+
	54	+ require_ascend()
	55	+ assert cu_seqlens_q.device.type == 'npu', 'input tensors must be on an NPU device'
	56	+ device = cu_seqlens_q.device
	57	+ batch_size = cu_seqlens_q.numel() - 1
	58	+ assert batch_size >= 0, 'cu_seqlens_q must contain at least one element'
	59	+ assert cu_seqlens_k.numel() == batch_size + 1
	60	+
	61	+ if batch_size == 0 or output_len == 0 or total_q == 0:
	62	+ empty = torch.empty(0, dtype=torch.int32, device=device)
	63	+ z1 = torch.zeros(1, dtype=torch.int32, device=device)
	64	+ return (
	65	+ empty,
	66	+ torch.zeros(total_q, dtype=torch.int32, device=device),
	67	+ empty,
	68	+ empty,
	69	+ z1,
	70	+ z1,
	71	+ )
	72	+
	73	+ # ---- Phase 1: precompute per-batch quantities (NPU kernel, no sync) ---- #
	74	+ k_c4_lens = torch.empty(batch_size, dtype=torch.int32, device=device)
	75	+ partial_seqlens = torch.empty(batch_size, dtype=torch.int32, device=device)
	76	+ elems_per_batch = torch.empty(batch_size, dtype=torch.int32, device=device)
	77	+ delta_sq = torch.empty(batch_size, 4, dtype=torch.int32, device=device)
	78	+ delta_sk = torch.empty(batch_size, 4, dtype=torch.int32, device=device)
	79	+
	80	+ precompute_kern = _precompute_tl_ascend(cp_size=cp_size, cp_rank=cp_rank, compress_ratio=compress_ratio)
	81	+ precompute_kern(
	82	+ cu_seqlens_q,
	83	+ cu_seqlens_k,
	84	+ k_c4_lens,
	85	+ partial_seqlens,
	86	+ elems_per_batch,
	87	+ delta_sq,
	88	+ delta_sk,
	89	+ )
	90	+
	91	+ # ---- Phase 2: cumulative sums (NPU ops, no sync) ---- #
	92	+ cu_seqlens_k_c4 = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
	93	+ cu_seqlens_k_c4[1:] = torch.cumsum(k_c4_lens, dim=0)
	94	+
	95	+ cu_elems = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
	96	+ cu_elems[1:] = torch.cumsum(elems_per_batch, dim=0)
	97	+
	98	+ total_per_rank = cu_elems[batch_size : batch_size + 1].contiguous() # [1] tensor on NPU
	99	+
	100	+ # ---- Phase 3: main scatter kernel (NPU kernel, no sync) ---- #
	101	+ mirror_cp_reverse_index = torch.zeros(total_q, dtype=torch.int32, device=device)
	102	+ mirror_cp_index = torch.empty(output_len, dtype=torch.int32, device=device)
	103	+ ks_out = torch.empty(output_len, dtype=torch.int32, device=device)
	104	+ ke_out = torch.empty(output_len, dtype=torch.int32, device=device)
	105	+
	106	+ scatter_kern = _scatter_tl_ascend(cp_size=cp_size, cp_rank=cp_rank, compress_ratio=compress_ratio)
	107	+ scatter_kern(
	108	+ cu_seqlens_q,
	109	+ cu_seqlens_k,
	110	+ cu_seqlens_k_c4,
	111	+ partial_seqlens,
	112	+ cu_elems,
	113	+ total_per_rank,
	114	+ mirror_cp_reverse_index,
	115	+ mirror_cp_index,
	116	+ ks_out,
	117	+ ke_out,
	118	+ )
	119	+
	120	+ # ---- Phase 4: cu_seqlens_q_cp / cu_seqlens_k_cp (NPU ops, no sync) ---- #
	121	+ delta_sq_flat = delta_sq.reshape(-1)
	122	+ delta_sk_flat = delta_sk.reshape(-1)
	123	+
	124	+ cu_seqlens_q_cp = torch.zeros(batch_size * 4 + 1, dtype=torch.int32, device=device)
	125	+ cu_seqlens_k_cp = torch.zeros(batch_size * 4 + 1, dtype=torch.int32, device=device)
	126	+ cu_seqlens_q_cp[1:] = torch.cumsum(delta_sq_flat, dim=0)
	127	+ cu_seqlens_k_cp[1:] = torch.cumsum(delta_sk_flat, dim=0)
	128	+
	129	+ return (
	130	+ mirror_cp_index,
	131	+ mirror_cp_reverse_index,
	132	+ ks_out,
	133	+ ke_out,
	134	+ cu_seqlens_q_cp,
	135	+ cu_seqlens_k_cp,
	136	+ )

Asrc/cann_ops_tilelang/attention/prepare_cp_mirror_kernel.py+265-0

@@ -0,0 +1,265 @@
	1	+"""TileLang Ascend kernels for mirror context-parallel metadata."""
	2	+
	3	+import tilelang
	4	+from tilelang.ascend import language as T
	5	+
	6	+from cann_ops_tilelang.config import get_num_vector_cores
	7	+
	8	+
	9	+@tilelang.jit(
	10	+ pass_configs={
	11	+ tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True,
	12	+ tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
	13	+ },
	14	+)
	15	+def _precompute_tl_ascend(cp_size: int, cp_rank: int, compress_ratio: int) -> None:
	16	+ """Ascend AIV implementation of the per-batch precomputation."""
	17	+ batch_size = T.dynamic('batch_size')
	18	+ num_cores = get_num_vector_cores()
	19	+
	20	+ @T.prim_func
	21	+ def _precompute(
	22	+ cu_seqlens_q: T.Tensor([batch_size + 1], 'int32'), # type: ignore
	23	+ cu_seqlens_k: T.Tensor([batch_size + 1], 'int32'), # type: ignore
	24	+ k_c4_lens: T.Tensor([batch_size], 'int32'), # type: ignore
	25	+ partial_seqlens: T.Tensor([batch_size], 'int32'), # type: ignore
	26	+ elems_per_batch: T.Tensor([batch_size], 'int32'), # type: ignore
	27	+ delta_sq: T.Tensor([batch_size, 4], 'int32'), # type: ignore
	28	+ delta_sk: T.Tensor([batch_size, 4], 'int32'), # type: ignore
	29	+ ) -> None:
	30	+ with T.Kernel(num_cores) as pid:
	31	+ for i in T.Persistent([batch_size], num_cores, pid, group_size=1):
	32	+ sq = T.int32(cu_seqlens_q[i + 1]) - T.int32(cu_seqlens_q[i])
	33	+ sk = T.int32(cu_seqlens_k[i + 1]) - T.int32(cu_seqlens_k[i])
	34	+
	35	+ k_c4_val = sk // compress_ratio
	36	+ k_c4_lens[i] = T.max(1, k_c4_val)
	37	+
	38	+ ps = (sq + cp_size * 2 - 1) // (cp_size * 2)
	39	+ partial_seqlens[i] = ps
	40	+ elems_per_batch[i] = 2 * ps
	41	+
	42	+ nv_h0 = T.min(ps, T.max(0, sq - ps * cp_rank))
	43	+ mirror_rank = cp_size * 2 - 1 - cp_rank
	44	+ nv_h1 = T.min(ps, T.max(0, sq - ps * mirror_rank))
	45	+ kl_h0 = sk - sq + ps * cp_rank + nv_h0
	46	+ kl_h1 = sk - sq + ps * mirror_rank + nv_h1
	47	+
	48	+ delta_sk[i, 0] = kl_h0 // compress_ratio
	49	+ delta_sq[i, 0] = kl_h0 // compress_ratio - (kl_h0 - nv_h0) // compress_ratio
	50	+
	51	+ delta_sk[i, 1] = -(kl_h0 // compress_ratio)
	52	+ delta_sq[i, 1] = -(kl_h0 // compress_ratio) + (kl_h0 + ps - nv_h0) // compress_ratio
	53	+
	54	+ delta_sk[i, 2] = kl_h1 // compress_ratio
	55	+ delta_sq[i, 2] = kl_h1 // compress_ratio - (kl_h1 - nv_h1) // compress_ratio
	56	+
	57	+ sk_arg3 = sq - ps * mirror_rank - nv_h1
	58	+ diff3 = sk_arg3 - (ps - nv_h1)
	59	+ td_sk3 = T.if_then_else(sk_arg3 >= 0, sk_arg3 // compress_ratio, -((-sk_arg3) // compress_ratio))
	60	+ td_diff3 = T.if_then_else(diff3 >= 0, diff3 // compress_ratio, -((-diff3) // compress_ratio))
	61	+ delta_sk[i, 3] = td_sk3
	62	+ delta_sq[i, 3] = td_sk3 - td_diff3
	63	+
	64	+ return _precompute
	65	+
	66	+
	67	+@tilelang.jit(
	68	+ pass_configs={
	69	+ tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True,
	70	+ tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
	71	+ },
	72	+)
	73	+def _scatter_tl_ascend(cp_size: int, cp_rank: int, compress_ratio: int) -> None:
	74	+ """Ascend AIV scatter with bounded UB tiles and independent work domains."""
	75	+ batch_size = T.dynamic('batch_size')
	76	+ total_q = T.dynamic('total_q')
	77	+ output_len = T.dynamic('output_len')
	78	+
	79	+ block_size = 8192
	80	+ vec_len = 64
	81	+ num_cores = get_num_vector_cores()
	82	+ ratio_is_pow2 = (compress_ratio & (compress_ratio - 1)) == 0
	83	+ ratio_shift = compress_ratio.bit_length() - 1
	84	+
	85	+ @T.prim_func
	86	+ def _scatter(
	87	+ cu_seqlens_q: T.Tensor([batch_size + 1], 'int32'), # type: ignore
	88	+ cu_seqlens_k: T.Tensor([batch_size + 1], 'int32'), # type: ignore
	89	+ cu_seqlens_k_c4: T.Tensor([batch_size + 1], 'int32'), # type: ignore
	90	+ partial_seqlens: T.Tensor([batch_size], 'int32'), # type: ignore
	91	+ cu_elems: T.Tensor([batch_size + 1], 'int32'), # type: ignore
	92	+ total_per_rank: T.Tensor([1], 'int32'), # type: ignore
	93	+ mirror_cp_reverse_index: T.Tensor([total_q], 'int32'), # type: ignore
	94	+ mirror_cp_index: T.Tensor([output_len], 'int32'), # type: ignore
	95	+ ks_out: T.Tensor([output_len], 'int32'), # type: ignore
	96	+ ke_out: T.Tensor([output_len], 'int32'), # type: ignore
	97	+ ) -> None:
	98	+ with T.Kernel(num_cores) as pid:
	99	+ total_per_rank_value = T.int32(total_per_rank[0])
	100	+ reverse_ub = T.alloc_shared((block_size,), 'int32')
	101	+ ks_ub = T.alloc_shared((block_size,), 'int32')
	102	+ ke_ub = T.alloc_shared((block_size,), 'int32')
	103	+
	104	+ for i_rev0, rank_rev0 in T.Persistent([batch_size, cp_size], num_cores, pid, group_size=1):
	105	+ sq_rev0 = T.int32(cu_seqlens_q[i_rev0 + 1]) - T.int32(cu_seqlens_q[i_rev0])
	106	+ cu_q_rev0 = T.int32(cu_seqlens_q[i_rev0])
	107	+ cu_elem_rev0 = T.int32(cu_elems[i_rev0])
	108	+ ps_rev0 = T.int32(partial_seqlens[i_rev0])
	109	+ local_start_rev0 = ps_rev0 * rank_rev0
	110	+ valid_count_rev0 = T.min(ps_rev0, T.max(0, sq_rev0 - local_start_rev0))
	111	+
	112	+ for block_rev0 in T.serial(T.ceildiv(valid_count_rev0, block_size)):
	113	+ offset_rev0 = block_rev0 * block_size
	114	+ copy_length_rev0 = T.min(block_size, valid_count_rev0 - offset_rev0)
	115	+ with T.SimdVF():
	116	+ mask_rev0 = T.simd.pset(32)
	117	+ for v_rev0 in T.serial(T.ceildiv(copy_length_rev0, vec_len)):
	118	+ value_rev0 = T.simd.vci(
	119	+ rank_rev0 * total_per_rank_value + cu_elem_rev0 + offset_rev0 + v_rev0 * vec_len,
	120	+ 'int32',
	121	+ )
	122	+ T.simd.vsts(reverse_ub[v_rev0 * vec_len], value_rev0, mask_rev0)
	123	+ reverse_start_rev0 = cu_q_rev0 + local_start_rev0 + offset_rev0
	124	+ T.copy(
	125	+ reverse_ub[:copy_length_rev0],
	126	+ mirror_cp_reverse_index[reverse_start_rev0 : reverse_start_rev0 + copy_length_rev0],
	127	+ )
	128	+
	129	+ for i_rev1, rank_rev1 in T.Persistent([batch_size, cp_size], num_cores, pid, group_size=1):
	130	+ sq_rev1 = T.int32(cu_seqlens_q[i_rev1 + 1]) - T.int32(cu_seqlens_q[i_rev1])
	131	+ cu_q_rev1 = T.int32(cu_seqlens_q[i_rev1])
	132	+ cu_elem_rev1 = T.int32(cu_elems[i_rev1])
	133	+ ps_rev1 = T.int32(partial_seqlens[i_rev1])
	134	+ local_start_rev1 = ps_rev1 * (cp_size * 2 - 1 - rank_rev1)
	135	+ valid_count_rev1 = T.min(ps_rev1, T.max(0, sq_rev1 - local_start_rev1))
	136	+
	137	+ for block_rev1 in T.serial(T.ceildiv(valid_count_rev1, block_size)):
	138	+ offset_rev1 = block_rev1 * block_size
	139	+ copy_length_rev1 = T.min(block_size, valid_count_rev1 - offset_rev1)
	140	+ with T.SimdVF():
	141	+ mask_rev1 = T.simd.pset(32)
	142	+ for v_rev1 in T.serial(T.ceildiv(copy_length_rev1, vec_len)):
	143	+ value_rev1 = T.simd.vci(
	144	+ rank_rev1 * total_per_rank_value
	145	+ + cu_elem_rev1
	146	+ + ps_rev1
	147	+ + offset_rev1
	148	+ + v_rev1 * vec_len,
	149	+ 'int32',
	150	+ )
	151	+ T.simd.vsts(reverse_ub[v_rev1 * vec_len], value_rev1, mask_rev1)
	152	+ reverse_start_rev1 = cu_q_rev1 + local_start_rev1 + offset_rev1
	153	+ T.copy(
	154	+ reverse_ub[:copy_length_rev1],
	155	+ mirror_cp_reverse_index[reverse_start_rev1 : reverse_start_rev1 + copy_length_rev1],
	156	+ )
	157	+
	158	+ # Forward outputs include valid and padded positions for cp_rank.
	159	+ for i_fwd0 in T.Persistent([batch_size], num_cores, pid, group_size=1):
	160	+ sq_fwd0 = T.int32(cu_seqlens_q[i_fwd0 + 1]) - T.int32(cu_seqlens_q[i_fwd0])
	161	+ sk_fwd0 = T.int32(cu_seqlens_k[i_fwd0 + 1]) - T.int32(cu_seqlens_k[i_fwd0])
	162	+ cu_q_fwd0 = T.int32(cu_seqlens_q[i_fwd0])
	163	+ cu_k_c4_fwd0 = T.int32(cu_seqlens_k_c4[i_fwd0])
	164	+ ps_fwd0 = T.int32(partial_seqlens[i_fwd0])
	165	+ local_start_fwd0 = ps_fwd0 * cp_rank
	166	+ out_start_fwd0 = T.int32(cu_elems[i_fwd0])
	167	+
	168	+ for block_fwd0 in T.serial(T.ceildiv(ps_fwd0, block_size)):
	169	+ offset_fwd0 = block_fwd0 * block_size
	170	+ copy_length_fwd0 = T.min(block_size, ps_fwd0 - offset_fwd0)
	171	+ with T.SimdVF():
	172	+ mask_fwd0 = T.simd.pset(32)
	173	+ neg1_fwd0 = T.simd.vdup(-1, 'int32', mask_fwd0)
	174	+ ks_vec_fwd0 = T.simd.vdup(cu_k_c4_fwd0, 'int32', mask_fwd0)
	175	+ if compress_ratio > 1 and not ratio_is_pow2:
	176	+ ratio_vec_fwd0 = T.simd.vdup(compress_ratio, 'int32', mask_fwd0)
	177	+ for v_fwd0 in T.serial(T.ceildiv(copy_length_fwd0, vec_len)):
	178	+ local_idx_vec_fwd0 = T.simd.vci(local_start_fwd0 + offset_fwd0 + v_fwd0 * vec_len, 'int32')
	179	+ valid_fwd0 = T.simd.vcmps(local_idx_vec_fwd0, sq_fwd0, mask_fwd0, 'lt')
	180	+ reverse_vec_fwd0 = T.simd.vadds(local_idx_vec_fwd0, cu_q_fwd0, mask_fwd0)
	181	+ ke_numerator_fwd0 = T.simd.vadds(local_idx_vec_fwd0, sk_fwd0 - sq_fwd0 + 1, mask_fwd0)
	182	+ if compress_ratio == 1:
	183	+ ke_delta_fwd0 = ke_numerator_fwd0
	184	+ elif ratio_is_pow2:
	185	+ ke_delta_fwd0 = T.simd.vshrs(ke_numerator_fwd0, ratio_shift, mask_fwd0)
	186	+ else:
	187	+ ke_delta_fwd0 = T.simd.vdiv(ke_numerator_fwd0, ratio_vec_fwd0, mask_fwd0)
	188	+ ke_vec_fwd0 = T.simd.vadds(ke_delta_fwd0, cu_k_c4_fwd0, mask_fwd0)
	189	+ T.simd.vsts(
	190	+ reverse_ub[v_fwd0 * vec_len],
	191	+ T.simd.vsel(reverse_vec_fwd0, neg1_fwd0, valid_fwd0),
	192	+ mask_fwd0,
	193	+ )
	194	+ T.simd.vsts(ks_ub[v_fwd0 * vec_len], ks_vec_fwd0, mask_fwd0)
	195	+ T.simd.vsts(
	196	+ ke_ub[v_fwd0 * vec_len], T.simd.vsel(ke_vec_fwd0, ks_vec_fwd0, valid_fwd0), mask_fwd0
	197	+ )
	198	+
	199	+ T.copy(
	200	+ reverse_ub[:copy_length_fwd0],
	201	+ mirror_cp_index[out_start_fwd0 + offset_fwd0 : out_start_fwd0 + offset_fwd0 + copy_length_fwd0],
	202	+ )
	203	+ T.copy(
	204	+ ks_ub[:copy_length_fwd0],
	205	+ ks_out[out_start_fwd0 + offset_fwd0 : out_start_fwd0 + offset_fwd0 + copy_length_fwd0],
	206	+ )
	207	+ T.copy(
	208	+ ke_ub[:copy_length_fwd0],
	209	+ ke_out[out_start_fwd0 + offset_fwd0 : out_start_fwd0 + offset_fwd0 + copy_length_fwd0],
	210	+ )
	211	+
	212	+ for i_fwd1 in T.Persistent([batch_size], num_cores, pid, group_size=1):
	213	+ sq_fwd1 = T.int32(cu_seqlens_q[i_fwd1 + 1]) - T.int32(cu_seqlens_q[i_fwd1])
	214	+ sk_fwd1 = T.int32(cu_seqlens_k[i_fwd1 + 1]) - T.int32(cu_seqlens_k[i_fwd1])
	215	+ cu_q_fwd1 = T.int32(cu_seqlens_q[i_fwd1])
	216	+ cu_k_c4_fwd1 = T.int32(cu_seqlens_k_c4[i_fwd1])
	217	+ ps_fwd1 = T.int32(partial_seqlens[i_fwd1])
	218	+ local_start_fwd1 = ps_fwd1 * (cp_size * 2 - 1 - cp_rank)
	219	+ out_start_fwd1 = T.int32(cu_elems[i_fwd1]) + ps_fwd1
	220	+
	221	+ for block_fwd1 in T.serial(T.ceildiv(ps_fwd1, block_size)):
	222	+ offset_fwd1 = block_fwd1 * block_size
	223	+ copy_length_fwd1 = T.min(block_size, ps_fwd1 - offset_fwd1)
	224	+ with T.SimdVF():
	225	+ mask_fwd1 = T.simd.pset(32)
	226	+ neg1_fwd1 = T.simd.vdup(-1, 'int32', mask_fwd1)
	227	+ ks_vec_fwd1 = T.simd.vdup(cu_k_c4_fwd1, 'int32', mask_fwd1)
	228	+ if compress_ratio > 1 and not ratio_is_pow2:
	229	+ ratio_vec_fwd1 = T.simd.vdup(compress_ratio, 'int32', mask_fwd1)
	230	+ for v_fwd1 in T.serial(T.ceildiv(copy_length_fwd1, vec_len)):
	231	+ local_idx_vec_fwd1 = T.simd.vci(local_start_fwd1 + offset_fwd1 + v_fwd1 * vec_len, 'int32')
	232	+ valid_fwd1 = T.simd.vcmps(local_idx_vec_fwd1, sq_fwd1, mask_fwd1, 'lt')
	233	+ reverse_vec_fwd1 = T.simd.vadds(local_idx_vec_fwd1, cu_q_fwd1, mask_fwd1)
	234	+ ke_numerator_fwd1 = T.simd.vadds(local_idx_vec_fwd1, sk_fwd1 - sq_fwd1 + 1, mask_fwd1)
	235	+ if compress_ratio == 1:
	236	+ ke_delta_fwd1 = ke_numerator_fwd1
	237	+ elif ratio_is_pow2:
	238	+ ke_delta_fwd1 = T.simd.vshrs(ke_numerator_fwd1, ratio_shift, mask_fwd1)
	239	+ else:
	240	+ ke_delta_fwd1 = T.simd.vdiv(ke_numerator_fwd1, ratio_vec_fwd1, mask_fwd1)
	241	+ ke_vec_fwd1 = T.simd.vadds(ke_delta_fwd1, cu_k_c4_fwd1, mask_fwd1)
	242	+ T.simd.vsts(
	243	+ reverse_ub[v_fwd1 * vec_len],
	244	+ T.simd.vsel(reverse_vec_fwd1, neg1_fwd1, valid_fwd1),
	245	+ mask_fwd1,
	246	+ )
	247	+ T.simd.vsts(ks_ub[v_fwd1 * vec_len], ks_vec_fwd1, mask_fwd1)
	248	+ T.simd.vsts(
	249	+ ke_ub[v_fwd1 * vec_len], T.simd.vsel(ke_vec_fwd1, ks_vec_fwd1, valid_fwd1), mask_fwd1
	250	+ )
	251	+
	252	+ T.copy(
	253	+ reverse_ub[:copy_length_fwd1],
	254	+ mirror_cp_index[out_start_fwd1 + offset_fwd1 : out_start_fwd1 + offset_fwd1 + copy_length_fwd1],
	255	+ )
	256	+ T.copy(
	257	+ ks_ub[:copy_length_fwd1],
	258	+ ks_out[out_start_fwd1 + offset_fwd1 : out_start_fwd1 + offset_fwd1 + copy_length_fwd1],
	259	+ )
	260	+ T.copy(
	261	+ ke_ub[:copy_length_fwd1],
	262	+ ke_out[out_start_fwd1 + offset_fwd1 : out_start_fwd1 + offset_fwd1 + copy_length_fwd1],
	263	+ )
	264	+
	265	+ return _scatter

Asrc/cann_ops_tilelang/attention/prepare_llama_decode_metadata.py+63-0

@@ -0,0 +1,63 @@
	1	+"""Public Ascend interface for LLaMA decode metadata."""
	2	+
	3	+import torch
	4	+
	5	+from cann_ops_tilelang.config import get_num_vector_cores, require_ascend
	6	+
	7	+from .prepare_llama_decode_metadata_kernel import _prepare_llama_decode_metadata_tl_ascend
	8	+
	9	+
	10	+def prepare_llama_decode_metadata(
	11	+ seqlens_q: torch.Tensor,
	12	+ seqlens_k: torch.Tensor,
	13	+ cu_seqlens_q: torch.Tensor,
	14	+ block_tables_in: torch.Tensor,
	15	+ block_tables_offsets: torch.Tensor,
	16	+ is_placeholder: torch.Tensor,
	17	+ positions: torch.Tensor,
	18	+ valid_token_mask: torch.Tensor,
	19	+ gate_force_random: torch.Tensor,
	20	+ fix_routing_mask: torch.Tensor,
	21	+ context_lens: torch.Tensor,
	22	+ slot_mapping_store: torch.Tensor,
	23	+ query_indices: torch.Tensor,
	24	+ block_tables_out: torch.Tensor,
	25	+ *,
	26	+ max_seqlen_q: int,
	27	+ max_num_blocks: int,
	28	+ block_size: int,
	29	+ force_load_balance: bool = False,
	30	+) -> None:
	31	+ """Compute LLaMA decode metadata on the active accelerator via TileLang.
	32	+
	33	+ Standard (non-compressed, non-SWA) attention: each query token maps to a slot in
	34	+ a paged KV cache addressed by its block table. Supports variable-length seqlen_q.
	35	+ block_tables_in is compact 1D with block_tables_offsets for variable-length per-seq
	36	+ block tables; block_tables_out is the per-token expanded [total_q, max_num_blocks] form.
	37	+ """
	38	+ require_ascend()
	39	+ n_seq = seqlens_q.shape[0]
	40	+ num_cores = get_num_vector_cores()
	41	+ long_block_table_threshold = 512
	42	+ balance_sequences = num_cores < n_seq <= 64 or block_tables_in.numel() >= n_seq * long_block_table_threshold
	43	+ average_num_blocks = block_tables_in.numel() // max(n_seq, 1)
	44	+ use_double_buffer = n_seq >= 4 * num_cores or (n_seq > 2 * num_cores and average_num_blocks >= 128)
	45	+ kernel = _prepare_llama_decode_metadata_tl_ascend(
	46	+ max_seqlen_q, max_num_blocks, block_size, force_load_balance, balance_sequences, 2 if use_double_buffer else 1
	47	+ )
	48	+ kernel(
	49	+ seqlens_q,
	50	+ seqlens_k,
	51	+ cu_seqlens_q,
	52	+ block_tables_in,
	53	+ block_tables_offsets,
	54	+ is_placeholder,
	55	+ positions,
	56	+ valid_token_mask,
	57	+ gate_force_random,
	58	+ fix_routing_mask,
	59	+ context_lens,
	60	+ slot_mapping_store,
	61	+ query_indices,
	62	+ block_tables_out,
	63	+ )

Asrc/cann_ops_tilelang/attention/prepare_llama_decode_metadata_kernel.py+108-0

@@ -0,0 +1,108 @@
	1	+"""TileLang Ascend kernel for LLaMA decode metadata."""
	2	+
	3	+import tilelang
	4	+from tilelang.ascend import language as T
	5	+
	6	+from cann_ops_tilelang.config import get_num_vector_cores
	7	+
	8	+
	9	+@tilelang.jit(
	10	+ out_idx=None,
	11	+ target='ascend',
	12	+ pass_configs={
	13	+ tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True,
	14	+ tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
	15	+ },
	16	+)
	17	+def _prepare_llama_decode_metadata_tl_ascend(
	18	+ max_seqlen_q: int,
	19	+ max_num_blocks: int,
	20	+ block_size: int,
	21	+ force_load_balance: bool,
	22	+ balance_sequences: bool = False,
	23	+ num_stages: int = 1,
	24	+) -> None:
	25	+ n_seq = T.dynamic('n_seq')
	26	+ total_q = T.dynamic('total_q')
	27	+ total_blocks = T.dynamic('total_blocks')
	28	+ num_cores = get_num_vector_cores()
	29	+
	30	+ @T.prim_func
	31	+ def _prepare_llama_decode_metadata(
	32	+ seqlens_q: T.Tensor((n_seq,), 'int32'), # type: ignore
	33	+ seqlens_k: T.Tensor((n_seq,), 'int32'), # type: ignore
	34	+ cu_seqlens_q: T.Tensor((n_seq,), 'int32'), # type: ignore
	35	+ block_tables_in: T.Tensor((total_blocks,), 'int32'), # type: ignore
	36	+ block_tables_offsets: T.Tensor((n_seq + 1,), 'int32'), # type: ignore
	37	+ is_placeholder: T.Tensor((n_seq,), T.bool), # type: ignore
	38	+ positions: T.Tensor((total_q,), 'int32'), # type: ignore
	39	+ valid_token_mask: T.Tensor((total_q,), T.bool), # type: ignore
	40	+ gate_force_random: T.Tensor((total_q,), T.bool), # type: ignore
	41	+ fix_routing_mask: T.Tensor((total_q,), T.bool), # type: ignore
	42	+ context_lens: T.Tensor((total_q,), 'int32'), # type: ignore
	43	+ slot_mapping_store: T.Tensor((total_q,), 'int32'), # type: ignore
	44	+ query_indices: T.Tensor((total_q,), 'int32'), # type: ignore
	45	+ block_tables_out: T.Tensor((total_q, max_num_blocks), 'int32'), # type: ignore
	46	+ ):
	47	+ with T.Kernel(num_cores) as pid:
	48	+ block_table_ub = T.alloc_shared((max_num_blocks,), 'int32')
	49	+ if balance_sequences:
	50	+ base_sequences = T.floordiv(n_seq, num_cores)
	51	+ remainder = T.floormod(n_seq, num_cores)
	52	+ seq_begin = pid * base_sequences + T.min(pid, remainder)
	53	+ seq_end = seq_begin + base_sequences + T.if_then_else(pid < remainder, 1, 0)
	54	+ else:
	55	+ sequences_per_core = T.ceildiv(n_seq, num_cores)
	56	+ seq_begin = pid * sequences_per_core
	57	+ seq_end = T.min(seq_begin + sequences_per_core, n_seq)
	58	+ for seq_idx in T.Pipelined(
	59	+ seq_begin,
	60	+ seq_end,
	61	+ num_stages=num_stages,
	62	+ annotations={'multi_buffer_eligible': [block_table_ub]},
	63	+ ):
	64	+ sq = seqlens_q[seq_idx]
	65	+ sk = seqlens_k[seq_idx]
	66	+ is_pad = is_placeholder[seq_idx]
	67	+ q_base = T.if_then_else(seq_idx == 0, 0, cu_seqlens_q[seq_idx - 1])
	68	+ bt_offset = block_tables_offsets[seq_idx]
	69	+ num_blocks = block_tables_offsets[seq_idx + 1] - bt_offset
	70	+
	71	+ T.assume(sq > 0)
	72	+ T.assume(sq <= max_seqlen_q)
	73	+ T.assume(q_base >= 0)
	74	+ T.assume(q_base + sq <= total_q)
	75	+ T.assume(num_blocks > 0)
	76	+ T.assume(num_blocks <= max_num_blocks)
	77	+ T.assume(bt_offset >= 0)
	78	+ T.assume(bt_offset + num_blocks <= total_blocks)
	79	+
	80	+ # Cache each sequence's compact block table once, then reuse it for
	81	+ # both slot lookup and all per-token expanded block-table rows.
	82	+ T.copy(
	83	+ block_tables_in[bt_offset : bt_offset + num_blocks],
	84	+ block_table_ub[:num_blocks],
	85	+ )
	86	+
	87	+ with T.SimtVF(threads=max_seqlen_q):
	88	+ j = T.get_thread_binding()
	89	+ if j < sq:
	90	+ pos = sk - sq + j
	91	+ q_off = q_base + j
	92	+ block_number = block_table_ub[pos // block_size]
	93	+
	94	+ positions[q_off] = pos
	95	+ valid_token_mask[q_off] = T.if_then_else(is_pad, False, True)
	96	+ gate_force_random[q_off] = force_load_balance
	97	+ fix_routing_mask[q_off] = False
	98	+ context_lens[q_off] = pos + 1
	99	+ query_indices[q_off] = T.if_then_else(is_pad, 998244353, seq_idx)
	100	+ slot_mapping_store[q_off] = block_number * block_size + pos % block_size
	101	+
	102	+ for j in T.serial(sq):
	103	+ T.copy(
	104	+ block_table_ub[:num_blocks],
	105	+ block_tables_out[q_base + j, :num_blocks],
	106	+ )
	107	+
	108	+ return _prepare_llama_decode_metadata

Asrc/cann_ops_tilelang/attention/prepare_v33_decode_metadata.py+93-0

@@ -0,0 +1,93 @@
	1	+"""Public Ascend interface for V33 decode metadata."""
	2	+
	3	+import torch
	4	+
	5	+from cann_ops_tilelang.config import get_num_vector_cores, require_ascend
	6	+
	7	+from .prepare_v33_decode_metadata_kernel import _prepare_v33_decode_metadata_tl_ascend
	8	+
	9	+
	10	+def prepare_v33_decode_metadata(
	11	+ seqlens_q: torch.Tensor,
	12	+ seqlens_k: torch.Tensor,
	13	+ cu_seqlens_q: torch.Tensor,
	14	+ block_tables_in: torch.Tensor,
	15	+ block_tables_offsets: torch.Tensor,
	16	+ state_block_idx_in: torch.Tensor,
	17	+ is_placeholder: torch.Tensor,
	18	+ positions: torch.Tensor,
	19	+ valid_token_mask: torch.Tensor,
	20	+ gate_force_random: torch.Tensor,
	21	+ fix_routing_mask: torch.Tensor,
	22	+ context_lens: torch.Tensor,
	23	+ klen: torch.Tensor,
	24	+ kv_ca_index_offset_len_4: torch.Tensor,
	25	+ kv_ca_index_offset_len_128: torch.Tensor,
	26	+ kv_swa_index_offset_len: torch.Tensor,
	27	+ slot_mapping_store_swa: torch.Tensor,
	28	+ query_indices: torch.Tensor,
	29	+ state_block_idx_out: torch.Tensor,
	30	+ block_tables_out: torch.Tensor,
	31	+ positions_c4: torch.Tensor,
	32	+ slot_mapping_store_ca4: torch.Tensor,
	33	+ positions_c128: torch.Tensor,
	34	+ slot_mapping_store_ca128: torch.Tensor,
	35	+ *,
	36	+ max_seqlen_q: int,
	37	+ max_num_blocks: int,
	38	+ window_size: int,
	39	+ swa_block_size: int,
	40	+ block_size: int,
	41	+ force_load_balance: bool = False,
	42	+) -> None:
	43	+ """Compute V33 decode metadata on the active accelerator via TileLang.
	44	+
	45	+ Supports positive variable-length seqlen_q. q_offsets is computed internally from seqlens_q.
	46	+ block_tables_in is compact 1D with block_tables_offsets for variable-length per-seq block tables.
	47	+ """
	48	+ require_ascend()
	49	+ if block_size < 128 or block_size % 128 != 0:
	50	+ raise ValueError(f'block_size must be a positive multiple of 128, got {block_size}')
	51	+ num_sequences = seqlens_q.shape[0]
	52	+ if num_sequences == 0:
	53	+ return
	54	+ num_vec_cores = get_num_vector_cores()
	55	+ average_num_blocks = block_tables_in.numel() // max(num_sequences, 1)
	56	+ use_double_buffer = num_sequences >= 4 * num_vec_cores or (
	57	+ num_sequences > 2 * num_vec_cores and average_num_blocks >= 128
	58	+ )
	59	+ kernel = _prepare_v33_decode_metadata_tl_ascend(
	60	+ max_seqlen_q,
	61	+ max_num_blocks,
	62	+ window_size,
	63	+ swa_block_size,
	64	+ block_size,
	65	+ force_load_balance,
	66	+ 2 if use_double_buffer else 1,
	67	+ )
	68	+ kernel(
	69	+ seqlens_q,
	70	+ seqlens_k,
	71	+ cu_seqlens_q,
	72	+ block_tables_in,
	73	+ block_tables_offsets,
	74	+ state_block_idx_in,
	75	+ is_placeholder,
	76	+ positions,
	77	+ valid_token_mask,
	78	+ gate_force_random,
	79	+ fix_routing_mask,
	80	+ context_lens,
	81	+ klen,
	82	+ kv_ca_index_offset_len_4,
	83	+ kv_ca_index_offset_len_128,
	84	+ kv_swa_index_offset_len,
	85	+ slot_mapping_store_swa,
	86	+ query_indices,
	87	+ state_block_idx_out,
	88	+ block_tables_out,
	89	+ positions_c4,
	90	+ slot_mapping_store_ca4,
	91	+ positions_c128,
	92	+ slot_mapping_store_ca128,
	93	+ )

Asrc/cann_ops_tilelang/attention/prepare_v33_decode_metadata_kernel.py+153-0

@@ -0,0 +1,153 @@
	1	+"""TileLang Ascend kernel for V33 decode metadata."""
	2	+
	3	+import tilelang
	4	+from tilelang.ascend import language as T
	5	+
	6	+from cann_ops_tilelang.config import get_num_vector_cores
	7	+
	8	+
	9	+@tilelang.jit(
	10	+ out_idx=None,
	11	+ target='ascend',
	12	+ pass_configs={
	13	+ tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True,
	14	+ tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
	15	+ },
	16	+)
	17	+def _prepare_v33_decode_metadata_tl_ascend(
	18	+ max_seqlen_q: int,
	19	+ max_num_blocks: int,
	20	+ window_size: int,
	21	+ swa_block_size: int,
	22	+ block_size: int,
	23	+ force_load_balance: bool,
	24	+ num_stages: int = 1,
	25	+) -> None:
	26	+ n_seq = T.dynamic('n_seq')
	27	+ total_q = T.dynamic('total_q')
	28	+ total_blocks = T.dynamic('total_blocks')
	29	+ c4_cbs = block_size // 4
	30	+ c4_tpb = block_size // 4
	31	+ c128_cbs = block_size // 128
	32	+ c128_tpb = block_size // 128
	33	+ num_cores = get_num_vector_cores()
	34	+
	35	+ @T.prim_func
	36	+ def _prepare_v33_decode_metadata(
	37	+ seqlens_q: T.Tensor((n_seq,), 'int32'), # type: ignore
	38	+ seqlens_k: T.Tensor((n_seq,), 'int32'), # type: ignore
	39	+ cu_seqlens_q: T.Tensor((n_seq,), 'int32'), # type: ignore
	40	+ block_tables_in: T.Tensor((total_blocks,), 'int32'), # type: ignore
	41	+ block_tables_offsets: T.Tensor((n_seq + 1,), 'int32'), # type: ignore
	42	+ state_block_idx_in: T.Tensor((n_seq,), 'int32'), # type: ignore
	43	+ is_placeholder: T.Tensor((n_seq,), T.bool), # type: ignore
	44	+ positions: T.Tensor((total_q,), 'int32'), # type: ignore
	45	+ valid_token_mask: T.Tensor((total_q,), T.bool), # type: ignore
	46	+ gate_force_random: T.Tensor((total_q,), T.bool), # type: ignore
	47	+ fix_routing_mask: T.Tensor((total_q,), T.bool), # type: ignore
	48	+ context_lens: T.Tensor((total_q,), 'int32'), # type: ignore
	49	+ klen: T.Tensor((total_q,), 'int32'), # type: ignore
	50	+ kv_ca_index_offset_len_4: T.Tensor((total_q, 2), 'int32'), # type: ignore
	51	+ kv_ca_index_offset_len_128: T.Tensor((total_q, 2), 'int32'), # type: ignore
	52	+ kv_swa_index_offset_len: T.Tensor((total_q, 2), 'int32'), # type: ignore
	53	+ slot_mapping_store_swa: T.Tensor((total_q,), 'int32'), # type: ignore
	54	+ query_indices: T.Tensor((total_q,), 'int32'), # type: ignore
	55	+ state_block_idx_out: T.Tensor((total_q,), 'int32'), # type: ignore
	56	+ block_tables_out: T.Tensor((total_q, max_num_blocks), 'int32'), # type: ignore
	57	+ positions_c4: T.Tensor((total_q,), 'int32'), # type: ignore
	58	+ slot_mapping_store_ca4: T.Tensor((total_q,), 'int32'), # type: ignore
	59	+ positions_c128: T.Tensor((total_q,), 'int32'), # type: ignore
	60	+ slot_mapping_store_ca128: T.Tensor((total_q,), 'int32'), # type: ignore
	61	+ ):
	62	+ with T.Kernel(T.min(n_seq, num_cores)) as pid:
	63	+ block_table_ub = T.alloc_shared((max_num_blocks,), 'int32')
	64	+ sequences_per_core = T.ceildiv(n_seq, num_cores)
	65	+ seq_begin = pid * sequences_per_core
	66	+ seq_end = T.min(seq_begin + sequences_per_core, n_seq)
	67	+ for seq_idx in T.Pipelined(
	68	+ seq_begin,
	69	+ seq_end,
	70	+ num_stages=num_stages,
	71	+ annotations={'multi_buffer_eligible': [block_table_ub]},
	72	+ ):
	73	+ sq = seqlens_q[seq_idx]
	74	+ sk = seqlens_k[seq_idx]
	75	+ q_base = T.if_then_else(seq_idx == 0, 0, cu_seqlens_q[seq_idx - 1])
	76	+ state_id = state_block_idx_in[seq_idx]
	77	+ is_pad = is_placeholder[seq_idx]
	78	+ bt_offset = block_tables_offsets[seq_idx]
	79	+ num_blocks = block_tables_offsets[seq_idx + 1] - bt_offset
	80	+
	81	+ if num_stages > 1:
	82	+ T.assume(sq > 0)
	83	+ T.assume(sq <= max_seqlen_q)
	84	+ T.assume(q_base >= 0)
	85	+ T.assume(q_base + sq <= total_q)
	86	+ T.assume(num_blocks > 0)
	87	+ T.assume(num_blocks <= max_num_blocks)
	88	+ T.assume(bt_offset >= 0)
	89	+ T.assume(bt_offset + num_blocks <= total_blocks)
	90	+
	91	+ # Load the compact table once per sequence. CA lookups and all
	92	+ # expanded per-token rows reuse the same UB-resident data.
	93	+ T.copy(
	94	+ block_tables_in[bt_offset : bt_offset + num_blocks],
	95	+ block_table_ub[:num_blocks],
	96	+ )
	97	+
	98	+ with T.SimtVF(threads=max_seqlen_q):
	99	+ j = T.get_thread_binding()
	100	+ if j < sq:
	101	+ pos = sk - sq + j
	102	+ q_off = q_base + j
	103	+
	104	+ positions[q_off] = pos
	105	+ valid_token_mask[q_off] = T.if_then_else(is_pad, False, True)
	106	+ gate_force_random[q_off] = force_load_balance
	107	+ fix_routing_mask[q_off] = False
	108	+ context_lens[q_off] = pos + 1
	109	+ klen[q_off] = (pos + 1) // 4
	110	+
	111	+ kv_ca_index_offset_len_4[q_off, 0] = q_off * max_num_blocks * c4_tpb
	112	+ kv_ca_index_offset_len_4[q_off, 1] = (pos + 1) // 4
	113	+ kv_ca_index_offset_len_128[q_off, 0] = q_off * max_num_blocks * c128_tpb
	114	+ kv_ca_index_offset_len_128[q_off, 1] = (pos + 1) // 128
	115	+
	116	+ swa_len = T.min(pos + 1, window_size)
	117	+ window_start_pos = pos - swa_len + 1
	118	+ kv_swa_index_offset_len[q_off, 0] = (
	119	+ state_id * swa_block_size + window_start_pos % swa_block_size
	120	+ )
	121	+ kv_swa_index_offset_len[q_off, 1] = swa_len
	122	+ slot_mapping_store_swa[q_off] = state_id * swa_block_size + pos % swa_block_size
	123	+
	124	+ query_indices[q_off] = T.if_then_else(is_pad, 998244353, seq_idx)
	125	+ state_block_idx_out[q_off] = state_id
	126	+
	127	+ c4_idx = pos // 4
	128	+ if (pos + 1) % 4 == 0:
	129	+ positions_c4[q_off] = c4_idx * 4
	130	+ bt_idx = c4_idx // c4_cbs
	131	+ block_number = block_table_ub[bt_idx]
	132	+ slot_mapping_store_ca4[q_off] = block_number * c4_cbs + c4_idx % c4_cbs
	133	+ else:
	134	+ positions_c4[q_off] = 0
	135	+ slot_mapping_store_ca4[q_off] = -1
	136	+
	137	+ c128_idx = pos // 128
	138	+ if (pos + 1) % 128 == 0:
	139	+ positions_c128[q_off] = c128_idx * 128
	140	+ bt_idx = c128_idx // c128_cbs
	141	+ block_number = block_table_ub[bt_idx]
	142	+ slot_mapping_store_ca128[q_off] = block_number * c128_cbs + c128_idx % c128_cbs
	143	+ else:
	144	+ positions_c128[q_off] = 0
	145	+ slot_mapping_store_ca128[q_off] = -1
	146	+
	147	+ for j in T.serial(sq):
	148	+ T.copy(
	149	+ block_table_ub[:num_blocks],
	150	+ block_tables_out[q_base + j, :num_blocks],
	151	+ )
	152	+
	153	+ return _prepare_v33_decode_metadata

Asrc/cann_ops_tilelang/attention/rope.py+147-0

@@ -0,0 +1,147 @@
	1	+"""Public Ascend interface for rotary embeddings."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import torch
	6	+
	7	+from cann_ops_tilelang.config import require_ascend
	8	+
	9	+from .rope_kernel import _rotary_embedding_tl_ascend
	10	+
	11	+
	12	+def rotary_embedding(
	13	+ positions: torch.Tensor, # [seqlen]
	14	+ query: torch.Tensor, # [seqlen, num_heads * head_size] or [seqlen, num_heads, head_size]
	15	+ key: torch.Tensor \| None, # [seqlen, num_kv_heads * head_size] or [seqlen, num_kv_heads, head_size]
	16	+ query_out: torch.Tensor \| None, # The same shape as `query`
	17	+ key_out: torch.Tensor \| None, # The same shape as `key`
	18	+ head_size: int,
	19	+ cos_sin_cache: torch.Tensor, # [seqlen_ro, rot_dim]
	20	+ is_neox: bool,
	21	+ conjugate: bool = False,
	22	+) -> None:
	23	+ require_ascend()
	24	+ assert positions.ndim == 1
	25	+ assert query.ndim == 2 or query.ndim == 3
	26	+ if key is not None:
	27	+ assert key.ndim == 2 or key.ndim == 3
	28	+ if query_out is not None:
	29	+ assert query_out.shape == query.shape
	30	+ if key_out is not None:
	31	+ assert key is not None and key_out.shape == key.shape
	32	+ assert cos_sin_cache.ndim == 2
	33	+
	34	+ assert head_size > 0
	35	+ assert positions.is_contiguous(), 'positions must be contiguous'
	36	+ assert positions.dtype in (torch.int32, torch.int64), 'positions must use int32 or int64'
	37	+ assert query.dtype in (torch.bfloat16, torch.float32), 'query must use bfloat16 or float32'
	38	+ assert cos_sin_cache.dtype == torch.float32, 'cos_sin_cache must use float32'
	39	+ assert cos_sin_cache.stride(-1) == 1, 'cos_sin_cache must be contiguous in its last dimension'
	40	+ assert positions.device == query.device == cos_sin_cache.device, 'inputs must be on the same device'
	41	+
	42	+ seqlen = positions.shape[0]
	43	+ _, rot_dim = cos_sin_cache.shape
	44	+ assert rot_dim > 0 and rot_dim % 2 == 0, 'rotary_dim must be a positive even number'
	45	+ assert rot_dim <= head_size, 'rotary_dim must not exceed head_size'
	46	+
	47	+ assert query.shape[0] == seqlen
	48	+ if key is not None:
	49	+ assert key.shape[0] == seqlen
	50	+
	51	+ if query.ndim == 2:
	52	+ assert query.shape[1] % head_size == 0
	53	+ num_heads = query.shape[1] // head_size
	54	+ query = query.view(seqlen, num_heads, head_size)
	55	+ if query_out is not None:
	56	+ query_out = query_out.view(seqlen, num_heads, head_size)
	57	+ else:
	58	+ assert query.shape[2] == head_size
	59	+ num_heads = query.shape[1]
	60	+ if key is not None:
	61	+ if key.ndim == 2:
	62	+ assert key.shape[1] % head_size == 0
	63	+ num_kv_heads = key.shape[1] // head_size
	64	+ key = key.view(seqlen, num_kv_heads, head_size)
	65	+ if key_out is not None:
	66	+ key_out = key_out.view(seqlen, num_kv_heads, head_size)
	67	+ else:
	68	+ assert key.shape[2] == head_size
	69	+ num_kv_heads = key.shape[1]
	70	+
	71	+ positions_dtype = 'int32'
	72	+ if positions is not None:
	73	+ if positions.dtype == torch.int32:
	74	+ positions_dtype = 'int32'
	75	+ elif positions.dtype == torch.int64:
	76	+ positions_dtype = 'int64'
	77	+ else:
	78	+ raise AssertionError('Unsupported dtype')
	79	+
	80	+ dtype_str = ''
	81	+ if query.dtype == torch.bfloat16:
	82	+ dtype_str = 'bfloat16'
	83	+ elif query.dtype == torch.float:
	84	+ dtype_str = 'float32'
	85	+ else:
	86	+ raise AssertionError('Unsupported dtype')
	87	+ if key is not None:
	88	+ assert key.dtype == query.dtype and key.device == query.device
	89	+ if query_out is not None:
	90	+ assert query_out.dtype == query.dtype, 'query_out must have the same dtype as query'
	91	+ assert query_out.device == query.device
	92	+ if key_out is not None:
	93	+ assert key_out.dtype == key.dtype and key_out.device == key.device
	94	+ for tensor in (query, key, query_out, key_out):
	95	+ assert tensor is None or tensor.stride(-1) == 1
	96	+
	97	+ assert query.device.type == 'npu', 'inputs must be NPU tensors'
	98	+
	99	+ cos = cos_sin_cache[:, : rot_dim // 2]
	100	+ sin = cos_sin_cache[:, rot_dim // 2 :]
	101	+ assert cos.stride(0) == sin.stride(0)
	102	+ query_rope_kernel = _rotary_embedding_tl_ascend(
	103	+ dtype_str,
	104	+ positions_dtype,
	105	+ num_heads,
	106	+ head_size,
	107	+ rot_dim,
	108	+ True,
	109	+ not is_neox,
	110	+ conjugate,
	111	+ cos_sin_cache.stride(0),
	112	+ query.stride(0),
	113	+ query.stride(1),
	114	+ query.stride(0) if query_out is None else query_out.stride(0),
	115	+ query.stride(1) if query_out is None else query_out.stride(1),
	116	+ )
	117	+ query_rope_kernel(
	118	+ query,
	119	+ cos_sin_cache,
	120	+ positions,
	121	+ query if query_out is None else query_out,
	122	+ 0,
	123	+ )
	124	+
	125	+ if key is not None:
	126	+ key_rope_kernel = _rotary_embedding_tl_ascend(
	127	+ dtype_str,
	128	+ positions_dtype,
	129	+ num_kv_heads,
	130	+ head_size,
	131	+ rot_dim,
	132	+ True,
	133	+ not is_neox,
	134	+ conjugate,
	135	+ cos_sin_cache.stride(0),
	136	+ key.stride(0),
	137	+ key.stride(1),
	138	+ key.stride(0) if key_out is None else key_out.stride(0),
	139	+ key.stride(1) if key_out is None else key_out.stride(1),
	140	+ )
	141	+ key_rope_kernel(
	142	+ key,
	143	+ cos_sin_cache,
	144	+ positions,
	145	+ key if key_out is None else key_out,
	146	+ 0,
	147	+ )

Asrc/cann_ops_tilelang/attention/rope_kernel.py+136-0

@@ -0,0 +1,136 @@
	1	+"""TileLang Ascend kernel for rotary embeddings."""
	2	+
	3	+import tilelang
	4	+from tilelang.ascend import language as T
	5	+
	6	+from cann_ops_tilelang.config import get_num_vector_cores
	7	+
	8	+
	9	+@tilelang.jit(
	10	+ out_idx=None,
	11	+ target='ascend',
	12	+ pass_configs={
	13	+ tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True,
	14	+ tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
	15	+ },
	16	+)
	17	+def _rotary_embedding_tl_ascend(
	18	+ dtype: str,
	19	+ positions_dtype: str,
	20	+ nheads: int,
	21	+ head_size: int,
	22	+ rotary_dim: int,
	23	+ has_position_tensor: bool,
	24	+ interleaved: bool,
	25	+ conjugate: bool,
	26	+ cos_sin_stride: int,
	27	+ x_stride_0: int,
	28	+ x_stride_1: int,
	29	+ y_stride_0: int,
	30	+ y_stride_1: int,
	31	+) -> None:
	32	+ seqlen = T.dynamic('seqlen')
	33	+ seqlen_ro = T.dynamic('seqlen_ro')
	34	+ half_d = rotary_dim // 2
	35	+ dtype_bytes = 4 if dtype == 'float32' else 2
	36	+ supported_simd_dtype = dtype in ('bfloat16', 'float32')
	37	+ simd_half_aligned = half_d * dtype_bytes % 32 == 0
	38	+ use_simd_fast_path = supported_simd_dtype and half_d <= 64 and not interleaved and simd_half_aligned
	39	+ simd_mask_pattern = 'PAT_ALL' if half_d >= 64 else f'PAT_VL{half_d}'
	40	+ if use_simd_fast_path:
	41	+ min_padded_head_size = max(head_size, half_d + 64)
	42	+ ub_align_elements = 32 // dtype_bytes
	43	+ padded_head_size = (min_padded_head_size + ub_align_elements - 1) // ub_align_elements * ub_align_elements
	44	+ else:
	45	+ padded_head_size = head_size
	46	+ padded_half_d = (half_d + 63) // 64 * 64 if use_simd_fast_path else half_d
	47	+ tile_elements_per_head = padded_head_size
	48	+ max_heads_per_tile = 64 if use_simd_fast_path else 32
	49	+ heads_per_tile = max(1, min(nheads, max_heads_per_tile, 32 * 1024 // (tile_elements_per_head * dtype_bytes)))
	50	+ while nheads % heads_per_tile != 0:
	51	+ heads_per_tile -= 1
	52	+ head_tiles = nheads // heads_per_tile
	53	+ threads = 256 if heads_per_tile * half_d >= 4096 else 128
	54	+ num_cores = get_num_vector_cores()
	55	+ num_stages = 4 if use_simd_fast_path else 1
	56	+
	57	+ @T.macro
	58	+ def rotate_simd_block(x0_buffer, x1_buffer, cos_sin_buffer, x0_offset, x1_offset, cos_offset, mask):
	59	+ cos_value = T.simd.vld(cos_sin_buffer[cos_offset])
	60	+ sin_value = T.simd.vld(cos_sin_buffer[half_d + cos_offset])
	61	+ for head in T.unroll(heads_per_tile, explicit=True):
	62	+ if dtype == 'bfloat16':
	63	+ x0 = T.simd.vcvt(T.simd.vld(x0_buffer[head, x0_offset], dist='UNPK_B16'), T.float32, mask, part=0)
	64	+ x1 = T.simd.vcvt(T.simd.vld(x1_buffer[head, x1_offset], dist='UNPK_B16'), T.float32, mask, part=0)
	65	+ else:
	66	+ x0 = T.simd.vld(x0_buffer[head, x0_offset])
	67	+ x1 = T.simd.vld(x1_buffer[head, x1_offset])
	68	+ if conjugate:
	69	+ x0_result = T.simd.vadd(T.simd.vmul(x0, cos_value, mask), T.simd.vmul(x1, sin_value, mask), mask)
	70	+ x1_result = T.simd.vsub(T.simd.vmul(x1, cos_value, mask), T.simd.vmul(x0, sin_value, mask), mask)
	71	+ else:
	72	+ x0_result = T.simd.vsub(T.simd.vmul(x0, cos_value, mask), T.simd.vmul(x1, sin_value, mask), mask)
	73	+ x1_result = T.simd.vadd(T.simd.vmul(x0, sin_value, mask), T.simd.vmul(x1, cos_value, mask), mask)
	74	+ if dtype == 'bfloat16':
	75	+ T.simd.vsts(x0_buffer[head, x0_offset], T.simd.vcvt(x0_result, T.bfloat16, mask), mask, 'PK_B32')
	76	+ T.simd.vsts(x1_buffer[head, x1_offset], T.simd.vcvt(x1_result, T.bfloat16, mask), mask, 'PK_B32')
	77	+ else:
	78	+ T.simd.vsts(x0_buffer[head, x0_offset], x0_result, mask)
	79	+ T.simd.vsts(x1_buffer[head, x1_offset], x1_result, mask)
	80	+
	81	+ @T.prim_func
	82	+ def _rotary_embedding_ascend(
	83	+ x: T.StridedTensor((seqlen, nheads, head_size), (x_stride_0, x_stride_1, 1), dtype), # type: ignore
	84	+ cos_sin: T.StridedTensor((seqlen_ro, rotary_dim), (cos_sin_stride, 1), 'float32'), # type: ignore
	85	+ positions: T.Tensor((seqlen,), positions_dtype), # type: ignore
	86	+ y: T.StridedTensor((seqlen, nheads, head_size), (y_stride_0, y_stride_1, 1), dtype), # type: ignore
	87	+ seqlen_offsets: T.int32,
	88	+ ) -> None:
	89	+ with T.Kernel(num_cores) as core_id:
	90	+ x_ub = T.alloc_shared((heads_per_tile, padded_head_size), dtype)
	91	+ cos_sin_ub = T.alloc_shared((padded_half_d * 2,), 'float32')
	92	+ T.annotate_buffer_versions({x_ub: num_stages, cos_sin_ub: num_stages})
	93	+
	94	+ for task_id in T.Persistent(
	95	+ [seqlen * head_tiles],
	96	+ num_cores,
	97	+ core_id,
	98	+ group_size=1,
	99	+ num_stages=num_stages,
	100	+ annotations={'enable_offset': True},
	101	+ ):
	102	+ token_id = task_id // head_tiles
	103	+ head_begin = task_id % head_tiles * heads_per_tile
	104	+ position = T.alloc_var(positions_dtype, seqlen_offsets)
	105	+ if has_position_tensor:
	106	+ position += positions[token_id]
	107	+ T.copy(x[token_id, head_begin : head_begin + heads_per_tile, :head_size], x_ub[:, :head_size])
	108	+ T.copy(cos_sin[position, :rotary_dim], cos_sin_ub[:rotary_dim])
	109	+
	110	+ if use_simd_fast_path:
	111	+ with T.SimdVF():
	112	+ mask = T.simd.pset(32, simd_mask_pattern)
	113	+ rotate_simd_block(x_ub, x_ub, cos_sin_ub, 0, half_d, 0, mask)
	114	+ else:
	115	+ with T.SimtVF(threads=threads):
	116	+ for head, lane in T.Parallel(heads_per_tile, half_d):
	117	+ if interleaved:
	118	+ x0_col = lane * 2
	119	+ x1_col = lane * 2 + 1
	120	+ else:
	121	+ x0_col = lane
	122	+ x1_col = half_d + lane
	123	+ x0 = T.cast(x_ub[head, x0_col], 'float32')
	124	+ x1 = T.cast(x_ub[head, x1_col], 'float32')
	125	+ cos_value = cos_sin_ub[lane]
	126	+ sin_value = cos_sin_ub[half_d + lane]
	127	+ if conjugate:
	128	+ x_ub[head, x0_col] = T.cast(x0 * cos_value + x1 * sin_value, dtype)
	129	+ x_ub[head, x1_col] = T.cast(x1 * cos_value - x0 * sin_value, dtype)
	130	+ else:
	131	+ x_ub[head, x0_col] = T.cast(x0 * cos_value - x1 * sin_value, dtype)
	132	+ x_ub[head, x1_col] = T.cast(x0 * sin_value + x1 * cos_value, dtype)
	133	+
	134	+ T.copy(x_ub[:, :head_size], y[token_id, head_begin : head_begin + heads_per_tile, :head_size])
	135	+
	136	+ return _rotary_embedding_ascend

Asrc/cann_ops_tilelang/attention/state_cache.py+69-0

@@ -0,0 +1,69 @@
	1	+import torch
	2	+
	3	+from cann_ops_tilelang.attention.state_cache_kernel import _load_state_to_cache_tl_ascend, _select_state_slice_tl_ascend
	4	+from cann_ops_tilelang.config import require_ascend
	5	+
	6	+
	7	+def select_state_slice(state_tensor: torch.Tensor, state_cache: torch.Tensor, last_hit_position: int) -> None:
	8	+ require_ascend()
	9	+ _validate_state_tensors(state_tensor, state_cache, last_hit_position)
	10	+ num_layers, block_size_kvstore, feature_dim = state_cache.shape
	11	+ _, block_size_inference, _ = state_tensor.shape
	12	+ contiguous_tokens = (
	13	+ feature_dim > 1 and state_tensor.stride(1) == feature_dim and (state_cache.stride(1) == feature_dim)
	14	+ )
	15	+ kernel = _select_state_slice_tl_ascend(
	16	+ block_size_kvstore, block_size_inference, num_layers, feature_dim, contiguous_tokens
	17	+ )
	18	+ return kernel(state_tensor, state_cache, last_hit_position)
	19	+
	20	+
	21	+def load_state_to_cache(state_tensor: torch.Tensor, state_cache: torch.Tensor, last_hit_position: int) -> None:
	22	+ require_ascend()
	23	+ _validate_state_tensors(state_tensor, state_cache, last_hit_position)
	24	+ return _load_state_to_cache_impl(state_tensor, state_cache, last_hit_position)
	25	+
	26	+
	27	+def _validate_state_tensors(
	28	+ state_tensor: torch.Tensor,
	29	+ state_cache: torch.Tensor,
	30	+ last_hit_position: int,
	31	+) -> None:
	32	+ if state_tensor.ndim != 3 or state_cache.ndim != 3:
	33	+ raise ValueError('state_tensor and state_cache must both be 3D tensors')
	34	+ if state_tensor.dtype != torch.uint8 or state_cache.dtype != torch.uint8:
	35	+ raise TypeError('state_tensor and state_cache must both have dtype torch.uint8')
	36	+ if state_tensor.device != state_cache.device:
	37	+ raise ValueError('state_tensor and state_cache must be on the same device')
	38	+ if state_tensor.device.type != 'npu':
	39	+ raise ValueError(f'state tensors must be on an NPU, got {state_tensor.device}.')
	40	+ num_layers, block_size_kvstore, feature_dim = state_cache.shape
	41	+ tensor_num_layers, block_size_inference, tensor_feature_dim = state_tensor.shape
	42	+ if num_layers <= 0 or block_size_kvstore <= 0 or block_size_inference <= 0 or (feature_dim <= 0):
	43	+ raise ValueError('all state tensor dimensions must be positive')
	44	+ if tensor_num_layers != num_layers or tensor_feature_dim != feature_dim:
	45	+ raise ValueError('state_tensor and state_cache must have matching layer and feature dimensions')
	46	+ for name, tensor in (('state_tensor', state_tensor), ('state_cache', state_cache)):
	47	+ if tensor.stride(2) != 1:
	48	+ raise ValueError(f'{name} must be contiguous in the feature dimension')
	49	+ if tensor.stride(1) < feature_dim or tensor.stride(0) < tensor.shape[1] * tensor.stride(1):
	50	+ raise ValueError(f'{name} must have a non-overlapping row-major layout')
	51	+ if not isinstance(last_hit_position, int) or not 0 <= last_hit_position <= torch.iinfo(torch.int32).max:
	52	+ raise ValueError('last_hit_position must be a non-negative int32 value')
	53	+ return None
	54	+
	55	+
	56	+def _load_state_to_cache_impl(
	57	+ state_tensor: torch.Tensor,
	58	+ state_cache: torch.Tensor,
	59	+ last_hit_position: int,
	60	+) -> None:
	61	+ num_layers, block_size_kvstore, feature_dim = state_cache.shape
	62	+ _, block_size_inference, _ = state_tensor.shape
	63	+ contiguous_tokens = (
	64	+ feature_dim > 1 and state_tensor.stride(1) == feature_dim and state_cache.stride(1) == feature_dim
	65	+ )
	66	+ kernel = _load_state_to_cache_tl_ascend(
	67	+ block_size_kvstore, block_size_inference, num_layers, feature_dim, contiguous_tokens
	68	+ )
	69	+ return kernel(state_tensor, state_cache, last_hit_position)