chore: remove operators implementation-ops-tilelang-AtomGit

Msrc/cann_ops_tilelang/__init__.py+0-31

	@@ -49,28 +49,15 @@ __all__ = [
49	'unpack_from_e2m1fn_x2',	49	'unpack_from_e2m1fn_x2',
50	'apply_token_bitmask_inplace',	50	'apply_token_bitmask_inplace',
51	'attn_rel_h_rel_w',	51	'attn_rel_h_rel_w',
52	- 'batch_find_value_by_index',
53	'bmp_to_patches',	52	'bmp_to_patches',
54	- 'compress_and_update_state_decode',
55	- 'compress_and_update_state_prefill',
56	- 'compute_confidence_error',
57	'compute_seqlens_k_zos',	53	'compute_seqlens_k_zos',
58	'compute_seqlens_q_by_confidence_rates',	54	'compute_seqlens_q_by_confidence_rates',
59	- 'confidence_head_fused',
60	'convert_kv_index',	55	'convert_kv_index',
61	- 'dflash_input_store',
62	- 'find_value_by_index',
63	'fused_sample',	56	'fused_sample',
64	'gather',	57	'gather',
65	'gather_engram_tokens_zos',	58	'gather_engram_tokens_zos',
66	'gather_input_ids_zos',	59	'gather_input_ids_zos',
67	- 'get_accept_result',
68	- 'get_dflash_mainx_slot_mapping',
69	'get_fused_dflash_get_input',	60	'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',	61	'get_logits_cur_by_step',
75	'get_prompt_logprobs_tensor',	62	'get_prompt_logprobs_tensor',
76	'get_sum_smaller_probs',	63	'get_sum_smaller_probs',
	@@ -89,8 +76,6 @@ __all__ = [
89	'random_sampling',	76	'random_sampling',
90	'replace_input_embeds',	77	'replace_input_embeds',
91	'rotary_embedding',	78	'rotary_embedding',
92	- 'sampled_tokens_update',
93	- 'scatter_verified',
94	'select_state_slice',	79	'select_state_slice',
95	'shuffle_image_tokens',	80	'shuffle_image_tokens',
96	'softmax',	81	'softmax',
	@@ -99,17 +84,11 @@ __all__ = [
99	'store_kv_cache',	84	'store_kv_cache',
100	'store_kv_cache_llama',	85	'store_kv_cache_llama',
101	'unbind',	86	'unbind',
102	- 'update_verify_states',
103	]	87	]
104		88
105	from .attention import (	89	from .attention import (
106	convert_kv_index,	90	convert_kv_index,
107	- get_dflash_mainx_slot_mapping,
108	get_fused_dflash_get_input,	91	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,	92	kv_index_cat,
114	load_index_kv_cache,	93	load_index_kv_cache,
115	load_kv_cache,	94	load_kv_cache,
	@@ -124,7 +103,6 @@ from .attention import (
124	store_kv_cache,	103	store_kv_cache,
125	store_kv_cache_llama,	104	store_kv_cache_llama,
126	)	105	)
127	-from .compression import compress_and_update_state_decode, compress_and_update_state_prefill
128	from .engram import (	106	from .engram import (
129	engram_fused_weight,	107	engram_fused_weight,
130	engram_gate_bwd,	108	engram_gate_bwd,
	@@ -185,23 +163,14 @@ from .quant import (
185	)	163	)
186	from .sampling import (	164	from .sampling import (
187	apply_token_bitmask_inplace,	165	apply_token_bitmask_inplace,
188	- batch_find_value_by_index,
189	- compute_confidence_error,
190	compute_seqlens_q_by_confidence_rates,	166	compute_seqlens_q_by_confidence_rates,
191	- confidence_head_fused,
192	- dflash_input_store,
193	- find_value_by_index,
194	fused_sample,	167	fused_sample,
195	- get_accept_result,
196	get_logits_cur_by_step,	168	get_logits_cur_by_step,
197	get_prompt_logprobs_tensor,	169	get_prompt_logprobs_tensor,
198	get_sum_smaller_probs,	170	get_sum_smaller_probs,
199	log_softmax,	171	log_softmax,
200	random_sampling,	172	random_sampling,
201	- sampled_tokens_update,
202	- scatter_verified,
203	softmax,	173	softmax,
204	softmax_and_logsoftmax,	174	softmax_and_logsoftmax,
205	- update_verify_states,
206	)	175	)
207	from .tensor import compute_seqlens_k_zos, gather, gather_engram_tokens_zos, gather_input_ids_zos, unbind	176	from .tensor import compute_seqlens_k_zos, gather, gather_engram_tokens_zos, gather_input_ids_zos, unbind

Msrc/cann_ops_tilelang/attention/__init__.py+1-13

	@@ -2,12 +2,7 @@
2		2
3	__all__ = [	3	__all__ = [
4	'convert_kv_index',	4	'convert_kv_index',
5	- 'get_dflash_mainx_slot_mapping',
6	'get_fused_dflash_get_input',	5	'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',	6	'kv_index_cat',
12	'load_index_kv_cache',	7	'load_index_kv_cache',
13	'load_kv_cache',	8	'load_kv_cache',
	@@ -32,14 +27,7 @@ from .cache import (
32	store_kv_cache_llama,	27	store_kv_cache_llama,
33	)	28	)
34	from .flash_mla_utils import convert_kv_index	29	from .flash_mla_utils import convert_kv_index
35	-from .get_kv_index import (	30	+from .get_kv_index import get_fused_dflash_get_input
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	31	from .kv_index_cat import kv_index_cat
44	from .prepare_cp_mirror import prepare_for_mirror_cp_tilelang	32	from .prepare_cp_mirror import prepare_for_mirror_cp_tilelang
45	from .prepare_llama_decode_metadata import prepare_llama_decode_metadata	33	from .prepare_llama_decode_metadata import prepare_llama_decode_metadata

Msrc/cann_ops_tilelang/attention/get_kv_index.py+1-91

	@@ -1,10 +1,6 @@
1	import torch	1	import torch
2		2
3	-from cann_ops_tilelang.attention.get_kv_index_kernel import (	3	+from cann_ops_tilelang.attention.get_kv_index_kernel import _get_fused_dflash_get_input_tl_ascend
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	4	from cann_ops_tilelang.config import require_ascend
9		5
10		6
	@@ -17,73 +13,6 @@ def _validate_int32_npu_tensor(tensor: torch.Tensor, name: str, ndim: int \| None
17	raise ValueError(f'{name} must be {ndim}D, got shape {tuple(tensor.shape)}.')	13	raise ValueError(f'{name} must be {ndim}D, got shape {tuple(tensor.shape)}.')
18		14
19		15
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(	16	def get_fused_dflash_get_input(
88	slot_mapping_swa: torch.Tensor,	17	slot_mapping_swa: torch.Tensor,
89	positions: torch.Tensor,	18	positions: torch.Tensor,
	@@ -162,22 +91,3 @@ def get_fused_dflash_get_input(
162	gate_force_random,	91	gate_force_random,
163	kv_index_out,	92	kv_index_out,
164	)	93	)
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

Msrc/cann_ops_tilelang/attention/get_kv_index_kernel.py+0-187

	@@ -7,114 +7,6 @@ from tilelang.ascend import language as T
7	from cann_ops_tilelang.config import get_num_vector_cores	7	from cann_ops_tilelang.config import get_num_vector_cores
8		8
9		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:	10	def get_ascend_disable_smem_reuse_config() -> dict:
119	pass_configs = {tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True}	11	pass_configs = {tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True}
120	key = getattr(tilelang.PassConfigKey, 'TL_DISABLE_SHARED_MEMORY_REUSE', None)	12	key = getattr(tilelang.PassConfigKey, 'TL_DISABLE_SHARED_MEMORY_REUSE', None)
	@@ -324,82 +216,3 @@ def _get_fused_dflash_get_input_tl_ascend(
324	)	216	)
325		217
326	return _get_fused_dflash_get_input	218	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

Dsrc/cann_ops_tilelang/compression/__init__.py+0-5

	@@ -1,5 +0,0 @@
1	-"""Compression operators for Ascend."""
2	-
3	-__all__ = ['compress_and_update_state_decode', 'compress_and_update_state_prefill']
4	-
5	-from .compress import compress_and_update_state_decode, compress_and_update_state_prefill

Dsrc/cann_ops_tilelang/compression/compress.py+0-266

	@@ -1,266 +0,0 @@
1	-"""Public Ascend interfaces for state-compression operators."""
2	-
3	-from __future__ import annotations
4	-
5	-import math
6	-
7	-import torch
8	-
9	-from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
10	-
11	-from .compress_kernel import (
12	- _compress_and_update_state_decode_tl_ascend,
13	- _compress_and_update_state_prefill_tl_ascend,
14	-)
15	-
16	-_UB_BUDGET_BYTES = 200 * 1024
17	-
18	-
19	-def _validate_kernel_capacity(
20	- dim: int,
21	- compress_ratio: int,
22	- overlap_ratio: int,
23	- *,
24	- has_ape: bool,
25	- is_prefill: bool,
26	-) -> None:
27	- if dim <= 0 or compress_ratio <= 0 or overlap_ratio <= 0:
28	- raise ValueError('dim, compress_ratio, and overlap_ratio must be positive.')
29	-
30	- compress_tokens = compress_ratio * overlap_ratio
31	- dim_block_base = (128 if compress_tokens <= 64 else 64) if is_prefill else 128
32	- dim_block = math.gcd(dim_block_base, dim)
33	- compute_buffers = (3 if has_ape else 2) if is_prefill else 2
34	- compute_stage_bytes = compute_buffers * compress_tokens * dim_block * 4 + dim_block * 2
35	- store_versions = 3 if is_prefill else 6
36	- store_buffers = 3 if has_ape else 2
37	- store_bytes = store_versions * store_buffers * overlap_ratio * dim * 4
38	- if not is_prefill:
39	- store_bytes += dim * 2
40	-
41	- if compute_stage_bytes > _UB_BUDGET_BYTES:
42	- raise ValueError(
43	- f'compression compute requires {compute_stage_bytes} bytes of UB per stage, '
44	- f'exceeding the {_UB_BUDGET_BYTES}-byte limit.'
45	- )
46	- if store_bytes > _UB_BUDGET_BYTES:
47	- raise ValueError(
48	- f'compression state update requires {store_bytes} bytes of UB, exceeding the {_UB_BUDGET_BYTES}-byte limit.'
49	- )
50	-
51	-
52	-def _validate_common_inputs(
53	- score: torch.Tensor,
54	- latent: torch.Tensor,
55	- ape: torch.Tensor \| None,
56	- positions: torch.Tensor \| None,
57	- state_cache: torch.Tensor,
58	- state_block_idx: torch.Tensor,
59	-) -> None:
60	- tensors = (score, latent, state_cache, state_block_idx)
61	- if not all(isinstance(tensor, torch.Tensor) for tensor in tensors):
62	- raise TypeError('score, latent, state_cache and state_block_idx must be tensors.')
63	- if score.device.type != 'npu' or any(tensor.device != score.device for tensor in tensors[1:]):
64	- raise ValueError('all required inputs must be on the same NPU device.')
65	- if score.dtype != torch.float32 or latent.dtype != torch.float32 or state_cache.dtype != torch.float32:
66	- raise TypeError('score, latent and state_cache must use torch.float32.')
67	- if state_block_idx.dtype != torch.int32:
68	- raise TypeError('state_block_idx must use torch.int32.')
69	- if score.shape != latent.shape or score.ndim != 3:
70	- raise ValueError('score and latent must have the same three-dimensional shape.')
71	- for name, tensor in (('ape', ape), ('positions', positions)):
72	- if tensor is not None and tensor.device != score.device:
73	- raise ValueError(f'{name} must be on {score.device}.')
74	-
75	-
76	-def compress_and_update_state_prefill(
77	- score: torch.Tensor,
78	- latent: torch.Tensor,
79	- ape: torch.Tensor \| None,
80	- positions: torch.Tensor \| None,
81	- state_cache: torch.Tensor,
82	- state_block_idx: torch.Tensor,
83	- compressed_batch_mapping: torch.Tensor,
84	- cu_seqlens_q: torch.Tensor,
85	- cu_seqlens_k: torch.Tensor,
86	- cu_compressed_seqlens_q: torch.Tensor,
87	- mirror_cp_reverse_index: torch.Tensor \| None = None,
88	- compress_ratio: int \| None = None,
89	- *,
90	- bitwise_consistent: bool = False,
91	-) -> torch.Tensor:
92	- """
93	- Args:
94	- score: [total_q \| cp_size * cp_mirror_len, overlap_ratio, dim], fp32
95	- latent: [total_q \| cp_size * cp_mirror_len, overlap_ratio, dim], fp32
96	- ape: [compress_ratio, overlap_ratio, dim], fp32
97	- positions: [total_q], int32
98	- state_cache: [num_state_slots, state_cache_size, 2, overlap_ratio, dim], fp32
99	- state_block_idx: [num_seqs], int32
100	- compressed_batch_mapping: [total_num_compressed_tokens], int32
101	- cu_seqlens_q: [num_seqs + 1], int32
102	- cu_seqlens_k: [num_seqs + 1], int32
103	- cu_compressed_seqlens_q: [num_seqs + 1], int32. Each query must be >=1,
104	- or state_cache will not be updated.
105	- mirror_cp_reverse_index: [total_q], int32, in [0, cp_size * cp_mirror_len)
106	-
107	- Returns:
108	- kv_compressed: [total_num_compressed_tokens, dim], bf16
109	- """
110	- require_ascend()
111	- _validate_common_inputs(score, latent, ape, positions, state_cache, state_block_idx)
112	- _, overlap_ratio, dim = score.shape
113	- if ape is not None:
114	- has_ape = True
115	- compress_ratio = ape.shape[0]
116	- else:
117	- has_ape = False
118	- assert compress_ratio is not None
119	- _validate_kernel_capacity(
120	- dim,
121	- compress_ratio,
122	- overlap_ratio,
123	- has_ape=has_ape,
124	- is_prefill=True,
125	- )
126	- total_num_compressed_tokens = compressed_batch_mapping.shape[0]
127	- kv_compressed = torch.empty(total_num_compressed_tokens, dim, dtype=torch.bfloat16, device=score.device)
128	-
129	- if not isinstance(bitwise_consistent, bool):
130	- raise TypeError('bitwise_consistent must be a bool.')
131	- get_kernel = _compress_and_update_state_prefill_tl_ascend
132	-
133	- compute_kernel = get_kernel(
134	- dim=dim,
135	- compress_ratio=compress_ratio,
136	- overlap_ratio=overlap_ratio,
137	- has_ape=has_ape,
138	- has_index=mirror_cp_reverse_index is not None,
139	- bitwise_consistent=bitwise_consistent,
140	- stage='compute',
141	- )
142	- store_kernel = get_kernel(
143	- dim=dim,
144	- compress_ratio=compress_ratio,
145	- overlap_ratio=overlap_ratio,
146	- has_ape=has_ape,
147	- has_index=mirror_cp_reverse_index is not None,
148	- bitwise_consistent=bitwise_consistent,
149	- stage='store',
150	- )
151	- if should_print_kernel_source():
152	- print(compute_kernel.get_kernel_source())
153	- print(store_kernel.get_kernel_source())
154	- compute_kernel(
155	- score,
156	- latent,
157	- ape,
158	- positions,
159	- state_cache,
160	- state_block_idx,
161	- compressed_batch_mapping,
162	- cu_seqlens_q,
163	- cu_seqlens_k,
164	- cu_compressed_seqlens_q,
165	- mirror_cp_reverse_index,
166	- kv_compressed,
167	- )
168	- store_kernel(
169	- score,
170	- latent,
171	- ape,
172	- positions,
173	- state_cache,
174	- state_block_idx,
175	- cu_seqlens_q,
176	- cu_seqlens_k,
177	- mirror_cp_reverse_index,
178	- )
179	- return kv_compressed
180	-
181	-
182	-def compress_and_update_state_decode(
183	- score: torch.Tensor,
184	- latent: torch.Tensor,
185	- ape: torch.Tensor \| None,
186	- positions: torch.Tensor \| None,
187	- state_cache: torch.Tensor,
188	- state_block_idx: torch.Tensor,
189	- context_lens: torch.Tensor,
190	- max_seqlen_q: int,
191	- compress_ratio: int \| None = None,
192	- *,
193	- bitwise_consistent: bool = False,
194	-) -> torch.Tensor:
195	- """
196	- Args:
197	- score: [total_q, overlap_ratio, dim], fp32
198	- latent: [total_q, overlap_ratio, dim], fp32
199	- ape: [compress_ratio, overlap_ratio, dim], fp32
200	- positions: [total_q], int32
201	- state_cache: [num_state_slots, state_cache_size, 2, overlap_ratio, dim], fp32
202	- state_block_idx: [total_q], int32
203	- context_lens: [total_q], int32
204	-
205	- Returns:
206	- kv_compressed: [total_q, dim], bf16
207	- """
208	- require_ascend()
209	- _validate_common_inputs(score, latent, ape, positions, state_cache, state_block_idx)
210	- total_q, overlap_ratio, dim = score.shape
211	- if ape is not None:
212	- has_ape = True
213	- compress_ratio = ape.shape[0]
214	- else:
215	- has_ape = False
216	- assert compress_ratio is not None
217	- _validate_kernel_capacity(
218	- dim,
219	- compress_ratio,
220	- overlap_ratio,
221	- has_ape=has_ape,
222	- is_prefill=False,
223	- )
224	- assert state_cache.shape[1] >= compress_ratio * overlap_ratio + max_seqlen_q - 1
225	- kv_compressed = torch.empty(total_q, dim, dtype=torch.bfloat16, device=score.device)
226	-
227	- if not isinstance(bitwise_consistent, bool):
228	- raise TypeError('bitwise_consistent must be a bool.')
229	- get_kernel = _compress_and_update_state_decode_tl_ascend
230	-
231	- store_kernel = get_kernel(
232	- dim=dim,
233	- compress_ratio=compress_ratio,
234	- overlap_ratio=overlap_ratio,
235	- has_ape=has_ape,
236	- bitwise_consistent=bitwise_consistent,
237	- stage='store',
238	- )
239	- compute_kernel = get_kernel(
240	- dim=dim,
241	- compress_ratio=compress_ratio,
242	- overlap_ratio=overlap_ratio,
243	- has_ape=has_ape,
244	- bitwise_consistent=bitwise_consistent,
245	- stage='compute',
246	- )
247	- if should_print_kernel_source():
248	- print(store_kernel.get_kernel_source())
249	- print(compute_kernel.get_kernel_source())
250	- store_kernel(
251	- score,
252	- latent,
253	- ape,
254	- positions,
255	- state_cache,
256	- state_block_idx,
257	- context_lens,
258	- kv_compressed,
259	- )
260	- compute_kernel(
261	- state_cache,
262	- state_block_idx,
263	- context_lens,
264	- kv_compressed,
265	- )
266	- return kv_compressed

Dsrc/cann_ops_tilelang/compression/compress_kernel.py+0-476

	@@ -1,476 +0,0 @@
1	-"""TileLang Ascend kernels for state compression."""
2	-
3	-import math
4	-
5	-import tilelang
6	-from tilelang.ascend import language as T
7	-
8	-from cann_ops_tilelang.config import get_num_vector_cores
9	-
10	-
11	-@tilelang.jit(
12	- pass_configs={
13	- tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True,
14	- tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
15	- tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
16	- },
17	-)
18	-def _compress_and_update_state_prefill_tl_ascend(
19	- dim: int,
20	- compress_ratio: int,
21	- overlap_ratio: int,
22	- has_ape: bool,
23	- has_index: bool,
24	- bitwise_consistent: bool,
25	- stage: str,
26	-) -> object:
27	- total_q = T.dynamic('total_q')
28	- score_latent_shape_0 = T.dynamic('score_latent_shape_0') if has_index else total_q
29	- score_stride0 = T.dynamic('score_stride0', dtype='int64')
30	- latent_stride0 = T.dynamic('latent_stride0', dtype='int64')
31	- num_state_slots = T.dynamic('num_state_slots')
32	- state_cache_size = T.dynamic('state_cache_size')
33	- state_cache_stride0 = T.dynamic('state_cache_stride0', dtype='int64')
34	- num_seqs = T.dynamic('num_seqs')
35	- total_c = T.dynamic('total_c')
36	- compress_ratio_0 = compress_ratio * overlap_ratio
37	-
38	- num_cores = get_num_vector_cores()
39	-
40	- if stage == 'store':
41	- blocks_per_seq = 32
42	-
43	- @T.prim_func
44	- def _update_state_prefill(
45	- score: T.StridedTensor(
46	- shape=[score_latent_shape_0, overlap_ratio, dim], strides=[score_stride0, dim, 1], dtype=T.float32
47	- ), # type: ignore
48	- latent: T.StridedTensor(
49	- shape=[score_latent_shape_0, overlap_ratio, dim], strides=[latent_stride0, dim, 1], dtype=T.float32
50	- ), # type: ignore
51	- ape: T.Tensor([compress_ratio, overlap_ratio, dim], T.float32), # type: ignore
52	- positions: T.Tensor([total_q], T.int32), # type: ignore
53	- state_cache: T.StridedTensor(
54	- shape=[num_state_slots, state_cache_size, 2, overlap_ratio, dim],
55	- strides=[state_cache_stride0, 2 * overlap_ratio * dim, overlap_ratio * dim, dim, 1],
56	- dtype=T.float32,
57	- ), # type: ignore
58	- state_block_idx: T.Tensor([num_seqs], T.int32), # type: ignore
59	- cu_seqlens_q: T.Tensor([num_seqs + 1], T.int32), # type: ignore
60	- cu_seqlens_k: T.Tensor([num_seqs + 1], T.int32), # type: ignore
61	- mirror_cp_reverse_index: T.Tensor([total_q], T.int32), # type: ignore
62	- ) -> None:
63	- T.assume(state_cache_stride0 % (2 * overlap_ratio * dim) == 0)
64	- T.assume(score_stride0 % dim == 0)
65	- T.assume(latent_stride0 % dim == 0)
66	- with T.Kernel(num_cores) as pid:
67	- for seq_idx, by in T.Persistent([num_seqs, blocks_per_seq], num_cores, pid, group_size=1):
68	- qs = T.int32(cu_seqlens_q[seq_idx])
69	- qe = T.int32(cu_seqlens_q[seq_idx + 1])
70	- ks = T.int32(cu_seqlens_k[seq_idx])
71	- ke = T.int32(cu_seqlens_k[seq_idx + 1])
72	- klen = ke - ks
73	- qlen = qe - qs
74	- qstart = klen - qlen
75	- state_idx = T.int32(state_block_idx[seq_idx])
76	- T.assume(0 <= state_idx < num_state_slots)
77	-
78	- num_store_tokens = T.min(state_cache_size, qlen)
79	- _slot = T.truncmod(qstart + qlen - 1, state_cache_size)
80	-
81	- score_cache_ub = T.alloc_shared([overlap_ratio, dim], T.float32)
82	- latent_cache_ub = T.alloc_shared([overlap_ratio, dim], T.float32)
83	- if has_ape:
84	- ape_ub = T.alloc_shared([overlap_ratio, dim], T.float32)
85	-
86	- T.annotate_buffer_versions(
87	- {
88	- score_cache_ub: 3,
89	- latent_cache_ub: 3,
90	- **({ape_ub: 3} if has_ape else {}),
91	- }
92	- )
93	-
94	- for lastq in T.serial(by, num_store_tokens, blocks_per_seq):
95	- lastqid = qlen - lastq - 1
96	- idx = T.alloc_var(T.int32)
97	- idx = qs + lastqid
98	- if has_ape:
99	- position_load = positions[idx] % compress_ratio
100	- if has_index:
101	- T.assume(0 <= idx < total_q)
102	- idx = mirror_cp_reverse_index[idx]
103	- cache_idx = T.alloc_var(T.int32)
104	- cache_idx = _slot - lastq
105	- cache_idx = T.Select(cache_idx < 0, cache_idx + state_cache_size, cache_idx)
106	- T.assume(0 <= cache_idx < state_cache_size)
107	-
108	- if has_ape:
109	- T.assume(0 <= position_load < compress_ratio)
110	- T.copy(ape[position_load, 0, 0], ape_ub)
111	- T.copy(score[idx, 0, 0], score_cache_ub)
112	- T.copy(latent[idx, 0, 0], latent_cache_ub)
113	-
114	- if has_ape:
115	- with T.SimtVF(256):
116	- for i, j in T.Parallel(overlap_ratio, dim):
117	- score_cache_ub[i, j] += ape_ub[i, j]
118	-
119	- T.assume_no_conflict(state_cache[state_idx, cache_idx, 0, 0:overlap_ratio, 0:dim])
120	- T.assume_no_conflict(state_cache[state_idx, cache_idx, 1, 0:overlap_ratio, 0:dim])
121	- T.copy(score_cache_ub, state_cache[state_idx, cache_idx, 0, 0, 0])
122	- T.copy(latent_cache_ub, state_cache[state_idx, cache_idx, 1, 0, 0])
123	-
124	- return _update_state_prefill
125	-
126	- elif stage == 'compute':
127	- dim_blk = math.gcd(128 if compress_ratio_0 <= 64 else 64, dim)
128	- n_dim_blks = dim // dim_blk
129	-
130	- _n_buffers = 3 if has_ape else 2
131	- n_stages = min(8, 200 * 1024 // (_n_buffers * compress_ratio_0 * dim_blk * 4 + dim_blk * 2))
132	-
133	- @T.prim_func
134	- def _compress_prefill(
135	- score: T.StridedTensor(
136	- shape=[score_latent_shape_0, overlap_ratio, dim], strides=[score_stride0, dim, 1], dtype=T.float32
137	- ), # type: ignore
138	- latent: T.StridedTensor(
139	- shape=[score_latent_shape_0, overlap_ratio, dim], strides=[latent_stride0, dim, 1], dtype=T.float32
140	- ), # type: ignore
141	- ape: T.Tensor([compress_ratio, overlap_ratio, dim], T.float32), # type: ignore
142	- positions: T.Tensor([total_q], T.int32), # type: ignore
143	- state_cache: T.StridedTensor(
144	- shape=[num_state_slots, state_cache_size, 2, overlap_ratio, dim],
145	- strides=[state_cache_stride0, 2 * overlap_ratio * dim, overlap_ratio * dim, dim, 1],
146	- dtype=T.float32,
147	- ), # type: ignore
148	- state_block_idx: T.Tensor([num_seqs], T.int32), # type: ignore
149	- compressed_batch_mapping: T.Tensor([total_c], T.int32), # type: ignore
150	- cu_seqlens_q: T.Tensor([num_seqs + 1], T.int32), # type: ignore
151	- cu_seqlens_k: T.Tensor([num_seqs + 1], T.int32), # type: ignore
152	- cu_compressed_seqlens_q: T.Tensor([num_seqs + 1], T.int32), # type: ignore
153	- mirror_cp_reverse_index: T.Tensor([total_q], T.int32), # type: ignore
154	- kv_compressed: T.Tensor([total_c, dim], T.bfloat16), # type: ignore
155	- ) -> None:
156	- T.assume(state_cache_stride0 % (2 * overlap_ratio * dim) == 0)
157	- T.assume(score_stride0 % dim == 0)
158	- T.assume(latent_stride0 % dim == 0)
159	- with T.Kernel(num_cores) as pid:
160	- for bx, bdim in T.Persistent([total_c, n_dim_blks], num_cores, pid, group_size=1, num_stages=n_stages):
161	- dim_offset = bdim * dim_blk
162	- seq_id = T.int32(compressed_batch_mapping[bx])
163	- T.assume(0 <= seq_id < num_seqs)
164	- qs = T.int32(cu_seqlens_q[seq_id])
165	- qe = T.int32(cu_seqlens_q[seq_id + 1])
166	- ks = T.int32(cu_seqlens_k[seq_id])
167	- ke = T.int32(cu_seqlens_k[seq_id + 1])
168	- klen = ke - ks
169	- qlen = qe - qs
170	- qstart = klen - qlen
171	- state_idx = T.int32(state_block_idx[seq_id])
172	- T.assume(0 <= state_idx < num_state_slots)
173	-
174	- total_new_blocks = klen // compress_ratio - qstart // compress_ratio
175	- block_id = bx - cu_compressed_seqlens_q[seq_id]
176	-
177	- result_ub = T.alloc_shared([dim_blk], T.bfloat16)
178	-
179	- if block_id >= 0 and block_id < total_new_blocks:
180	- block_start_pos = (qstart // compress_ratio + block_id) * compress_ratio - (
181	- overlap_ratio - 1
182	- ) * compress_ratio
183	- block_start_pos_slot = T.truncmod(block_start_pos, state_cache_size)
184	- T.assume(-state_cache_size < block_start_pos_slot < state_cache_size)
185	-
186	- score_ub = T.alloc_shared([compress_ratio_0, dim_blk], T.float32)
187	- latent_ub = T.alloc_shared([compress_ratio_0, dim_blk], T.float32)
188	- if has_ape:
189	- ape_ub = T.alloc_shared([compress_ratio_0, dim_blk], T.float32)
190	-
191	- T.annotate_buffer_versions(
192	- {
193	- score_ub: n_stages,
194	- latent_ub: n_stages,
195	- result_ub: n_stages,
196	- **({ape_ub: n_stages} if has_ape else {}),
197	- }
198	- )
199	-
200	- for i in T.Serial(compress_ratio_0):
201	- overlap_idx = overlap_ratio - i // compress_ratio - 1
202	- pos = block_start_pos + i
203	-
204	- if pos >= 0 and pos < qstart:
205	- _cache_idx = block_start_pos_slot + i
206	- cache_idx = T.Select(
207	- _cache_idx >= state_cache_size, _cache_idx - state_cache_size, _cache_idx
208	- )
209	- T.assume(0 <= cache_idx < state_cache_size)
210	- T.copy(
211	- state_cache[
212	- state_idx, cache_idx, 0, overlap_idx, dim_offset : dim_offset + dim_blk
213	- ],
214	- score_ub[i, :],
215	- )
216	- T.copy(
217	- state_cache[
218	- state_idx, cache_idx, 1, overlap_idx, dim_offset : dim_offset + dim_blk
219	- ],
220	- latent_ub[i, :],
221	- )
222	- elif pos >= qstart:
223	- orig_idx = pos - qstart + qs
224	- if has_index:
225	- T.assume(0 <= orig_idx < total_q)
226	- load_idx = mirror_cp_reverse_index[orig_idx]
227	- else:
228	- load_idx = orig_idx
229	- T.copy(score[load_idx, overlap_idx, dim_offset : dim_offset + dim_blk], score_ub[i, :])
230	- T.copy(
231	- latent[load_idx, overlap_idx, dim_offset : dim_offset + dim_blk], latent_ub[i, :]
232	- )
233	- if has_ape:
234	- ape_pos = positions[orig_idx] % compress_ratio
235	- T.assume(0 <= ape_pos < compress_ratio)
236	- T.copy(ape[ape_pos, overlap_idx, dim_offset : dim_offset + dim_blk], ape_ub[i, :])
237	-
238	- with T.SimtVF(dim_blk):
239	- # Init: handle invalid positions and APE
240	- for i, j in T.Parallel(compress_ratio_0, dim_blk):
241	- if block_start_pos + i < 0:
242	- score_ub[i, j] = float('-inf')
243	- latent_ub[i, j] = 0
244	-
245	- if has_ape:
246	- for i, j in T.Parallel(compress_ratio_0, dim_blk):
247	- if block_start_pos + i >= qstart:
248	- score_ub[i, j] += ape_ub[i, j]
249	-
250	- # softmax + weighted sum (fused pass 2+3)
251	- max_score = T.alloc_reducer([dim_blk], T.float32, op='max')
252	- sum_score = T.alloc_reducer([dim_blk], T.float32)
253	- result = T.alloc_reducer([dim_blk], T.float32)
254	-
255	- T.fill(max_score, float('-inf'))
256	- for i, j in T.Parallel(compress_ratio_0, dim_blk):
257	- max_score[j] = T.max(max_score[j], score_ub[i, j])
258	- T.finalize_reducer(max_score)
259	-
260	- T.fill(sum_score, 0)
261	- T.fill(result, 0)
262	- for i, j in T.Parallel(compress_ratio_0, dim_blk):
263	- exp_score = T.exp(score_ub[i, j] - max_score[j])
264	- sum_score[j] += exp_score
265	- result[j] += latent_ub[i, j] * exp_score
266	- T.finalize_reducer(sum_score)
267	- T.finalize_reducer(result)
268	-
269	- for j in T.Parallel(dim_blk):
270	- result_ub[j] = result[j] / sum_score[j]
271	- else:
272	- with T.SimtVF(dim_blk):
273	- T.clear(result_ub)
274	-
275	- T.copy(result_ub, kv_compressed[bx, dim_offset : dim_offset + dim_blk])
276	-
277	- return _compress_prefill
278	-
279	- else:
280	- raise ValueError()
281	-
282	-
283	-@tilelang.jit(
284	- pass_configs={
285	- tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True,
286	- tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
287	- tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
288	- },
289	-)
290	-def _compress_and_update_state_decode_tl_ascend(
291	- dim: int,
292	- compress_ratio: int,
293	- overlap_ratio: int,
294	- has_ape: bool,
295	- bitwise_consistent: bool,
296	- stage: str,
297	-) -> object:
298	- total_q = T.dynamic('total_q')
299	- score_stride0 = T.dynamic('score_stride0', dtype='int64')
300	- latent_stride0 = T.dynamic('latent_stride0', dtype='int64')
301	- num_state_slots = T.dynamic('num_state_slots')
302	- state_cache_size = T.dynamic('state_cache_size')
303	- state_cache_stride0 = T.dynamic('state_cache_stride0', dtype='int64')
304	- compress_ratio_0 = compress_ratio * overlap_ratio
305	-
306	- num_cores = get_num_vector_cores()
307	-
308	- if stage == 'store':
309	-
310	- @T.prim_func
311	- def _update_state_decode(
312	- score: T.StridedTensor(
313	- shape=[total_q, overlap_ratio, dim], strides=[score_stride0, dim, 1], dtype=T.float32
314	- ), # type: ignore
315	- latent: T.StridedTensor(
316	- shape=[total_q, overlap_ratio, dim], strides=[latent_stride0, dim, 1], dtype=T.float32
317	- ), # type: ignore
318	- ape: T.Tensor([compress_ratio, overlap_ratio, dim], T.float32), # type: ignore
319	- positions: T.Tensor([total_q], T.int32), # type: ignore
320	- state_cache: T.StridedTensor(
321	- shape=[num_state_slots, state_cache_size, 2, overlap_ratio, dim],
322	- strides=[state_cache_stride0, 2 * overlap_ratio * dim, overlap_ratio * dim, dim, 1],
323	- dtype=T.float32,
324	- ), # type: ignore
325	- state_block_idx: T.Tensor([total_q], T.int32), # type: ignore
326	- context_lens: T.Tensor([total_q], T.int32), # type: ignore
327	- kv_compressed: T.Tensor([total_q, dim], T.bfloat16), # type: ignore
328	- ) -> None:
329	- T.assume(state_cache_stride0 % (2 * overlap_ratio * dim) == 0)
330	- T.assume(score_stride0 % dim == 0)
331	- T.assume(latent_stride0 % dim == 0)
332	- with T.Kernel(num_cores) as pid:
333	- zeros = T.alloc_shared([dim], T.bfloat16)
334	- with T.SimtVF(dim):
335	- for j in T.Parallel(dim):
336	- zeros[j] = 0
337	-
338	- for idx in T.Persistent([total_q], num_cores, pid, group_size=32):
339	- state_idx = T.int32(state_block_idx[idx])
340	- T.assume(0 <= state_idx < num_state_slots)
341	- klen = T.int32(context_lens[idx])
342	- cache_pos = T.uint32(klen - 1) % state_cache_size
343	- T.assume(0 <= cache_pos < state_cache_size)
344	-
345	- score_cache_ub = T.alloc_shared([overlap_ratio, dim], T.float32)
346	- latent_cache_ub = T.alloc_shared([overlap_ratio, dim], T.float32)
347	- if has_ape:
348	- ape_ub = T.alloc_shared([overlap_ratio, dim], T.float32)
349	- position_load = positions[idx] % compress_ratio
350	-
351	- T.annotate_buffer_versions(
352	- {
353	- score_cache_ub: 6,
354	- latent_cache_ub: 6,
355	- **({ape_ub: 6} if has_ape else {}),
356	- }
357	- )
358	-
359	- if has_ape:
360	- T.assume(0 <= position_load < compress_ratio)
361	- T.copy(ape[position_load, 0, 0], ape_ub)
362	- T.copy(score[idx, 0, 0], score_cache_ub)
363	- T.copy(latent[idx, 0, 0], latent_cache_ub)
364	-
365	- if has_ape:
366	- with T.SimtVF(256):
367	- for i, j in T.Parallel(overlap_ratio, dim):
368	- score_cache_ub[i, j] += ape_ub[i, j]
369	-
370	- T.copy(score_cache_ub, state_cache[state_idx, cache_pos, 0, 0, 0])
371	- T.copy(latent_cache_ub, state_cache[state_idx, cache_pos, 1, 0, 0])
372	-
373	- if (klen % compress_ratio) != 0:
374	- T.copy(zeros, kv_compressed[idx, 0])
375	-
376	- return _update_state_decode
377	-
378	- elif stage == 'compute':
379	- dim_blk = math.gcd(128, dim)
380	- n_dim_blks = dim // dim_blk
381	-
382	- n_stages = min(8, 200 * 1024 // (2 * compress_ratio_0 * dim_blk * 4 + dim_blk * 2))
383	-
384	- @T.prim_func
385	- def _compress_decode(
386	- state_cache: T.StridedTensor(
387	- shape=[num_state_slots, state_cache_size, 2, overlap_ratio, dim],
388	- strides=[state_cache_stride0, 2 * overlap_ratio * dim, overlap_ratio * dim, dim, 1],
389	- dtype='float32',
390	- ), # type: ignore
391	- state_block_idx: T.Tensor([total_q], 'int32'), # type: ignore
392	- context_lens: T.Tensor([total_q], 'int32'), # type: ignore
393	- kv_compressed: T.Tensor([total_q, dim], 'bfloat16'), # type: ignore
394	- ) -> None:
395	- T.assume(state_cache_stride0 % (2 * overlap_ratio * dim) == 0)
396	- with T.Kernel(num_cores) as pid:
397	- for bx, bdim in T.Persistent([total_q, n_dim_blks], num_cores, pid, group_size=1, num_stages=n_stages):
398	- dim_offset = bdim * dim_blk
399	-
400	- klen = T.int32(context_lens[bx])
401	- abs_pos = klen - 1
402	- block_abs_idx = abs_pos // compress_ratio
403	- block_start_pos = block_abs_idx * compress_ratio - (overlap_ratio - 1) * compress_ratio
404	-
405	- if (abs_pos % compress_ratio) == compress_ratio - 1:
406	- state_idx = T.int32(state_block_idx[bx])
407	- block_start_pos_slot = T.truncmod(block_start_pos, state_cache_size)
408	- T.assume(-state_cache_size < block_start_pos_slot < state_cache_size)
409	- T.assume(0 <= state_idx < num_state_slots)
410	-
411	- score_ub = T.alloc_shared([compress_ratio_0, dim_blk], T.float32)
412	- latent_ub = T.alloc_shared([compress_ratio_0, dim_blk], T.float32)
413	- result_ub = T.alloc_shared([dim_blk], T.bfloat16)
414	-
415	- T.annotate_buffer_versions(
416	- {
417	- score_ub: n_stages,
418	- latent_ub: n_stages,
419	- result_ub: n_stages,
420	- }
421	- )
422	-
423	- with T.SimtVF(256):
424	- T.fill(score_ub, float('-inf'))
425	- T.fill(latent_ub, 0)
426	-
427	- for i in T.Serial(compress_ratio_0):
428	- overlap_idx = overlap_ratio - i // compress_ratio - 1
429	- if block_start_pos + i >= 0:
430	- _cache_idx = block_start_pos_slot + i
431	- cache_idx = T.Select(
432	- _cache_idx >= state_cache_size, _cache_idx - state_cache_size, _cache_idx
433	- )
434	- T.assume(0 <= cache_idx < state_cache_size)
435	- T.copy(
436	- state_cache[
437	- state_idx, cache_idx, 0, overlap_idx, dim_offset : dim_offset + dim_blk
438	- ],
439	- score_ub[i, :],
440	- )
441	- T.copy(
442	- state_cache[
443	- state_idx, cache_idx, 1, overlap_idx, dim_offset : dim_offset + dim_blk
444	- ],
445	- latent_ub[i, :],
446	- )
447	-
448	- with T.SimtVF(dim_blk):
449	- # softmax + weighted sum (fused pass 2+3)
450	- max_score = T.alloc_reducer([dim_blk], T.float32, op='max')
451	- sum_score = T.alloc_reducer([dim_blk], T.float32)
452	- result = T.alloc_reducer([dim_blk], T.float32)
453	-
454	- T.fill(max_score, float('-inf'))
455	- for i, j in T.Parallel(compress_ratio_0, dim_blk):
456	- max_score[j] = T.max(max_score[j], score_ub[i, j])
457	- T.finalize_reducer(max_score)
458	-
459	- T.fill(sum_score, 0)
460	- T.fill(result, 0)
461	- for i, j in T.Parallel(compress_ratio_0, dim_blk):
462	- exp_score = T.exp(score_ub[i, j] - max_score[j])
463	- sum_score[j] += exp_score
464	- result[j] += latent_ub[i, j] * exp_score
465	- T.finalize_reducer(sum_score)
466	- T.finalize_reducer(result)
467	-
468	- for j in T.Parallel(dim_blk):
469	- result_ub[j] = result[j] / sum_score[j]
470	-
471	- T.copy(result_ub, kv_compressed[bx, dim_offset : dim_offset + dim_blk])
472	-
473	- return _compress_decode
474	-
475	- else:
476	- raise ValueError()

Msrc/cann_ops_tilelang/sampling/__init__.py+1-20

	@@ -2,40 +2,21 @@
2		2
3	__all__ = [	3	__all__ = [
4	'apply_token_bitmask_inplace',	4	'apply_token_bitmask_inplace',
5	- 'batch_find_value_by_index',
6	- 'compute_confidence_error',
7	'compute_seqlens_q_by_confidence_rates',	5	'compute_seqlens_q_by_confidence_rates',
8	- 'confidence_head_fused',
9	- 'dflash_input_store',
10	- 'find_value_by_index',
11	'fused_sample',	6	'fused_sample',
12	- 'get_accept_result',
13	'get_logits_cur_by_step',	7	'get_logits_cur_by_step',
14	'get_prompt_logprobs_tensor',	8	'get_prompt_logprobs_tensor',
15	'get_sum_smaller_probs',	9	'get_sum_smaller_probs',
16	'log_softmax',	10	'log_softmax',
17	'random_sampling',	11	'random_sampling',
18	- 'sampled_tokens_update',
19	- 'scatter_verified',
20	'softmax',	12	'softmax',
21	'softmax_and_logsoftmax',	13	'softmax_and_logsoftmax',
22	- 'update_verify_states',
23	]	14	]
24		15
25	-from .dflash import confidence_head_fused, dflash_input_store
26	from .fused_sample import fused_sample	16	from .fused_sample import fused_sample
27	from .prompt_logprobs import get_prompt_logprobs_tensor	17	from .prompt_logprobs import get_prompt_logprobs_tensor
28	from .random_sampling import random_sampling	18	from .random_sampling import random_sampling
29	-from .sampler import (	19	+from .sampler import get_logits_cur_by_step
30	- batch_find_value_by_index,
31	- compute_confidence_error,
32	- find_value_by_index,
33	- get_accept_result,
34	- get_logits_cur_by_step,
35	- sampled_tokens_update,
36	- scatter_verified,
37	- update_verify_states,
38	-)
39	from .schedule_mtp_varlen import compute_seqlens_q_by_confidence_rates	20	from .schedule_mtp_varlen import compute_seqlens_q_by_confidence_rates
40	from .softmax import log_softmax, softmax, softmax_and_logsoftmax	21	from .softmax import log_softmax, softmax, softmax_and_logsoftmax
41	from .structure_output import apply_token_bitmask_inplace	22	from .structure_output import apply_token_bitmask_inplace

Dsrc/cann_ops_tilelang/sampling/dflash.py+0-158

	@@ -1,158 +0,0 @@
1	-"""Public Ascend interfaces for DFlash helper operators."""
2	-
3	-from __future__ import annotations
4	-
5	-import math
6	-from collections.abc import Sequence
7	-from numbers import Integral, Real
8	-
9	-import torch
10	-
11	-from cann_ops_tilelang.config import get_num_vector_cores, require_ascend
12	-
13	-from .dflash_kernel import _confidence_head_fused_tl_ascend, _dflash_input_store_tl_ascend
14	-
15	-_SAMPLING_EPS = 1e-5
16	-
17	-
18	-def _check_npu_tensor(
19	- name: str,
20	- tensor: torch.Tensor,
21	- *,
22	- dtype: torch.dtype,
23	- device: torch.device \| None = None,
24	-) -> None:
25	- if not isinstance(tensor, torch.Tensor):
26	- raise TypeError(f'{name} must be a torch.Tensor.')
27	- if tensor.device.type != 'npu':
28	- raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')
29	- if device is not None and tensor.device != device:
30	- raise ValueError(f'{name} must be on {device}, got {tensor.device}.')
31	- if tensor.dtype != dtype:
32	- raise TypeError(f'{name} must use {dtype}, got {tensor.dtype}.')
33	-
34	-
35	-def _resolve_temperature(temperature: Real \| Sequence[Real], step_i: int) -> float:
36	- if isinstance(step_i, bool) or not isinstance(step_i, Integral) or step_i < 0:
37	- raise ValueError('step_i must be a non-negative integer.')
38	- value: Real
39	- if isinstance(temperature, Sequence) and not isinstance(temperature, (str, bytes)):
40	- if step_i >= len(temperature):
41	- raise ValueError(f'step_i {step_i} is outside the temperature sequence.')
42	- value = temperature[step_i]
43	- else:
44	- value = temperature # type: ignore[assignment]
45	- if isinstance(value, bool) or not isinstance(value, Real):
46	- raise TypeError('temperature must be a real number or a sequence of real numbers.')
47	- value = float(value)
48	- if not math.isfinite(value):
49	- raise ValueError('temperature must be finite.')
50	- return max(value, _SAMPLING_EPS)
51	-
52	-
53	-def confidence_head_fused(
54	- hidden: torch.Tensor,
55	- weight: torch.Tensor,
56	- markov_emb: torch.Tensor \| None = None,
57	- out: torch.Tensor \| None = None,
58	- step_i: int = 0,
59	- temperature: Real \| Sequence[Real] = 1.0,
60	-) -> torch.Tensor:
61	- """Apply the DFlash confidence head on Ascend.
62	-
63	- ``temperature`` replaces the framework environment setting used by the
64	- original integration and may be a scalar or a per-step sequence.
65	- """
66	-
67	- require_ascend()
68	- _check_npu_tensor('hidden', hidden, dtype=torch.bfloat16)
69	- _check_npu_tensor('weight', weight, dtype=torch.bfloat16, device=hidden.device)
70	- if hidden.ndim != 2:
71	- raise ValueError(f'hidden must be two-dimensional, got shape {tuple(hidden.shape)}.')
72	- if hidden.shape[1] <= 0:
73	- raise ValueError('hidden last dimension must be positive.')
74	- if hidden.stride(1) != 1:
75	- raise ValueError('hidden must have unit inner stride.')
76	- if not weight.is_contiguous():
77	- raise ValueError('weight must be contiguous.')
78	-
79	- num_rows, n_embed = hidden.shape
80	- has_markov = markov_emb is not None
81	- markov_rank = 0
82	- if markov_emb is not None:
83	- _check_npu_tensor('markov_emb', markov_emb, dtype=torch.bfloat16, device=hidden.device)
84	- if markov_emb.ndim != 2 or markov_emb.shape[0] != num_rows:
85	- raise ValueError(f'markov_emb must have shape ({num_rows}, markov_rank).')
86	- markov_rank = markov_emb.shape[1]
87	- if markov_rank <= 0:
88	- raise ValueError('markov_emb last dimension must be positive when provided.')
89	- if not markov_emb.is_contiguous():
90	- raise ValueError('markov_emb must be contiguous.')
91	- total_dim = n_embed + markov_rank
92	- if weight.numel() != total_dim:
93	- raise ValueError(f'weight must contain {total_dim} elements, got {weight.numel()}.')
94	- if total_dim % 64 or total_dim > 11000:
95	- raise ValueError('n_embed + markov_rank must be divisible by 64 and at most 11000.')
96	-
97	- if out is None:
98	- out = torch.empty(num_rows, dtype=torch.float32, device=hidden.device)
99	- else:
100	- _check_npu_tensor('out', out, dtype=torch.float32, device=hidden.device)
101	- if out.shape != (num_rows,):
102	- raise ValueError(f'out must have shape ({num_rows},), got {tuple(out.shape)}.')
103	- if out.stride(0) <= 0:
104	- raise ValueError('out must have a positive stride.')
105	-
106	- b8_fits_ub = 2 * (8 * 2 * total_dim + 64) + 2 * total_dim <= 210000
107	- b8_enough_work = num_rows >= 4 * get_num_vector_cores()
108	- block_rows = 8 if b8_enough_work and b8_fits_ub else 4
109	- kernel = _confidence_head_fused_tl_ascend(
110	- n_embed=n_embed,
111	- markov_rank=markov_rank,
112	- has_markov_emb=has_markov,
113	- block_rows=block_rows,
114	- )
115	- kernel(hidden, markov_emb, weight.view(-1), out, _resolve_temperature(temperature, step_i))
116	- return out
117	-
118	-
119	-def dflash_input_store(
120	- hidden_states: torch.Tensor,
121	- out: torch.Tensor \| None = None,
122	-) -> torch.Tensor:
123	- """Mean-reduce the optional HC dimension and store BF16 DFlash input."""
124	-
125	- require_ascend()
126	- _check_npu_tensor('hidden_states', hidden_states, dtype=torch.bfloat16)
127	- if hidden_states.ndim < 2 or not hidden_states.is_contiguous():
128	- raise ValueError('hidden_states must be a contiguous tensor with at least two dimensions.')
129	- n_embed = hidden_states.shape[-1]
130	- if n_embed <= 0 or n_embed > 8192 or n_embed % 128:
131	- raise ValueError('hidden_states last dimension must be positive, divisible by 128, and at most 8192.')
132	- hc_mult = hidden_states.shape[-2] if hidden_states.ndim > 2 else 1
133	- if hc_mult <= 0:
134	- raise ValueError('hidden_states HC dimension must be positive.')
135	- num_rows = hidden_states.numel() // (hc_mult * n_embed)
136	- hidden_states = hidden_states.view(num_rows, hc_mult, n_embed)
137	- rows_per_task = max(1, min(4, 32768 // n_embed)) if hc_mult == 1 else 1
138	- stage_bytes = 2 * rows_per_task * hc_mult * n_embed
139	- if stage_bytes > 240000:
140	- raise ValueError(
141	- f'dflash_input_store requires {stage_bytes} bytes per UB stage, exceeding the 240000-byte limit.'
142	- )
143	-
144	- if out is None:
145	- out = torch.empty((num_rows, n_embed), dtype=torch.bfloat16, device=hidden_states.device)
146	- else:
147	- _check_npu_tensor('out', out, dtype=torch.bfloat16, device=hidden_states.device)
148	- if out.shape != (num_rows, n_embed):
149	- raise ValueError(f'out must have shape {(num_rows, n_embed)}, got {tuple(out.shape)}.')
150	- if out.stride(1) != 1 or out.stride(0) % 4:
151	- raise ValueError('out must have unit inner stride and a row stride divisible by 4.')
152	-
153	- if num_rows == 0:
154	- return out
155	-
156	- kernel = _dflash_input_store_tl_ascend(n_embed, hc_mult, rows_per_task)
157	- kernel(hidden_states, out)
158	- return out

Dsrc/cann_ops_tilelang/sampling/dflash_kernel.py+0-189

	@@ -1,189 +0,0 @@
1	-"""TileLang Ascend kernels for DFlash helpers."""
2	-
3	-from typing import Any
4	-
5	-import tilelang
6	-from tilelang.ascend import language as T
7	-
8	-from cann_ops_tilelang.config import get_num_vector_cores
9	-
10	-
11	-@tilelang.jit(
12	- pass_configs={
13	- tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True,
14	- tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
15	- },
16	-)
17	-def _confidence_head_fused_tl_ascend(
18	- n_embed: int,
19	- markov_rank: int,
20	- has_markov_emb: bool,
21	- block_rows: int = 4,
22	-) -> Any:
23	- n = T.dynamic('n')
24	- out_stride = T.dynamic('out_stride')
25	- hidden_stride = T.dynamic('hidden_stride')
26	- num_cores = get_num_vector_cores()
27	- total_dim = n_embed + markov_rank
28	- vector_lanes = 64
29	- num_slices_total = total_dim // vector_lanes
30	- block_size = block_rows
31	- ub_size = 210000
32	- num_stages = min(8, (ub_size - 2 * total_dim) // (block_size * 2 * total_dim + 32 * 2))
33	-
34	- assert total_dim % vector_lanes == 0, 'total_dim must be VEC-aligned'
35	- assert total_dim <= 11000, 'unified buffer size insufficient'
36	-
37	- @T.prim_func
38	- def _confidence_head_fused(
39	- x0: T.StridedTensor((n, n_embed), (hidden_stride, 1), T.bfloat16), # type: ignore
40	- x1: T.Tensor((n, markov_rank or 1), T.bfloat16), # type: ignore
41	- weight: T.Tensor((total_dim,), T.bfloat16), # type: ignore
42	- out: T.StridedTensor((n,), (out_stride,), 'float32'), # type: ignore
43	- temp: T.float32,
44	- ) -> None:
45	- with T.Kernel(num_cores) as pid:
46	- w_ub = T.alloc_shared((total_dim,), T.bfloat16)
47	- x_ub = T.alloc_shared((block_size, total_dim), T.bfloat16)
48	- sum_ub = T.alloc_shared((block_size,), T.float32)
49	- T.annotate_buffer_versions({x_ub: num_stages, sum_ub: num_stages * 2})
50	-
51	- T.copy(weight[0:total_dim], w_ub)
52	-
53	- for bx in T.Persistent([T.ceildiv(n, block_size)], num_cores, pid, num_stages=num_stages * 2):
54	- batch_start = bx * block_size
55	- rem = T.min(block_size, n - batch_start)
56	-
57	- T.copy(x0[batch_start : batch_start + rem, :], x_ub[:rem, :n_embed])
58	- if has_markov_emb:
59	- T.copy(x1[batch_start : batch_start + rem, :], x_ub[:rem, n_embed:total_dim])
60	-
61	- with T.SimdVF():
62	- full = T.simd.pset(32)
63	- one = T.simd.pset(32, 'PAT_VL1')
64	- zero_v = T.simd.vdup(0.0, T.float32)
65	-
66	- sum_accs = T.simd.alloc_local((block_size,), T.float32)
67	- sum_accs[0] = zero_v
68	- sum_accs[1] = zero_v
69	- sum_accs[2] = zero_v
70	- sum_accs[3] = zero_v
71	- if block_size == 8:
72	- sum_accs[4] = zero_v
73	- sum_accs[5] = zero_v
74	- sum_accs[6] = zero_v
75	- sum_accs[7] = zero_v
76	-
77	- for i in T.serial(num_slices_total):
78	- off = i * vector_lanes
79	- w_v = T.simd.vcvt(T.simd.vld(w_ub[off], dist='UNPK_B16'), T.float32, part=0)
80	- x_v_0 = T.simd.vcvt(T.simd.vld(x_ub[0, off], dist='UNPK_B16'), T.float32, part=0)
81	- T.simd.vmula(sum_accs[0], w_v, x_v_0, full)
82	- x_v_1 = T.simd.vcvt(T.simd.vld(x_ub[1, off], dist='UNPK_B16'), T.float32, part=0)
83	- T.simd.vmula(sum_accs[1], w_v, x_v_1, full)
84	- x_v_2 = T.simd.vcvt(T.simd.vld(x_ub[2, off], dist='UNPK_B16'), T.float32, part=0)
85	- T.simd.vmula(sum_accs[2], w_v, x_v_2, full)
86	- x_v_3 = T.simd.vcvt(T.simd.vld(x_ub[3, off], dist='UNPK_B16'), T.float32, part=0)
87	- T.simd.vmula(sum_accs[3], w_v, x_v_3, full)
88	- if block_size == 8:
89	- x_v_4 = T.simd.vcvt(T.simd.vld(x_ub[4, off], dist='UNPK_B16'), T.float32, part=0)
90	- T.simd.vmula(sum_accs[4], w_v, x_v_4, full)
91	- x_v_5 = T.simd.vcvt(T.simd.vld(x_ub[5, off], dist='UNPK_B16'), T.float32, part=0)
92	- T.simd.vmula(sum_accs[5], w_v, x_v_5, full)
93	- x_v_6 = T.simd.vcvt(T.simd.vld(x_ub[6, off], dist='UNPK_B16'), T.float32, part=0)
94	- T.simd.vmula(sum_accs[6], w_v, x_v_6, full)
95	- x_v_7 = T.simd.vcvt(T.simd.vld(x_ub[7, off], dist='UNPK_B16'), T.float32, part=0)
96	- T.simd.vmula(sum_accs[7], w_v, x_v_7, full)
97	-
98	- total_0 = T.simd.vcadd(sum_accs[0], full)
99	- T.simd.vsts(sum_ub[0], total_0, one, 'ONEPT_B32')
100	- total_1 = T.simd.vcadd(sum_accs[1], full)
101	- T.simd.vsts(sum_ub[1], total_1, one, 'ONEPT_B32')
102	- total_2 = T.simd.vcadd(sum_accs[2], full)
103	- T.simd.vsts(sum_ub[2], total_2, one, 'ONEPT_B32')
104	- total_3 = T.simd.vcadd(sum_accs[3], full)
105	- T.simd.vsts(sum_ub[3], total_3, one, 'ONEPT_B32')
106	- if block_size == 8:
107	- total_4 = T.simd.vcadd(sum_accs[4], full)
108	- T.simd.vsts(sum_ub[4], total_4, one, 'ONEPT_B32')
109	- total_5 = T.simd.vcadd(sum_accs[5], full)
110	- T.simd.vsts(sum_ub[5], total_5, one, 'ONEPT_B32')
111	- total_6 = T.simd.vcadd(sum_accs[6], full)
112	- T.simd.vsts(sum_ub[6], total_6, one, 'ONEPT_B32')
113	- total_7 = T.simd.vcadd(sum_accs[7], full)
114	- T.simd.vsts(sum_ub[7], total_7, one, 'ONEPT_B32')
115	-
116	- # SimtVF supports 4 threads (threads>4 silently mis-computes):
117	- # each thread covers B/4 consecutive outputs (1 for B=4, 2 for B=8).
118	- with T.SimtVF(threads=4):
119	- j = T.get_thread_binding()
120	- for half in T.serial(block_size // 4):
121	- idx = half * 4 + j
122	- if idx < rem:
123	- raw = sum_ub[idx]
124	- scaled = T.min(T.max(raw / temp, T.cast(-60.0, T.float32)), T.cast(60.0, T.float32))
125	- out[batch_start + idx] = 1.0 / (1.0 + T.exp(-scaled))
126	-
127	- return _confidence_head_fused
128	-
129	-
130	-@tilelang.jit(out_idx=None, target='ascend')
131	-def _dflash_input_store_tl_ascend(
132	- n_embed: int,
133	- hc_mult: int,
134	- rows_per_task: int = 1,
135	-) -> Any:
136	- n = T.dynamic('n')
137	- out_stride = T.dynamic('out_stride')
138	- assert n_embed <= 8192 and n_embed % 128 == 0
139	- ub_size = 240000
140	- dtype = T.bfloat16
141	- block_size = n_embed
142	- row_tile_size = rows_per_task
143	- num_stages = min(16, ub_size // (row_tile_size * block_size * hc_mult * dtype.bytes))
144	- vec_size = 64
145	- num_cores = get_num_vector_cores()
146	-
147	- @T.prim_func
148	- def _dflash_input_store_ascend(
149	- x: T.Tensor[(n, hc_mult, n_embed), dtype],
150	- out: T.StridedTensor[(n, n_embed), (out_stride, 1), dtype],
151	- ) -> None:
152	- num_row_blocks = T.ceildiv(n, row_tile_size)
153	- num_col_blocks = T.ceildiv(n_embed, block_size)
154	- num_blocks = num_row_blocks * num_col_blocks
155	-
156	- with T.Kernel(num_cores) as core_id:
157	- x_ub = T.alloc_shared((row_tile_size * hc_mult, block_size), dtype)
158	- T.annotate_buffer_versions({x_ub: num_stages})
159	-
160	- for pid in T.Persistent([num_blocks], num_cores, core_id, group_size=1, num_stages=num_stages):
161	- rb = pid // num_col_blocks
162	- col = pid % num_col_blocks
163	- col_base = col * block_size
164	- row = rb * row_tile_size
165	- rem = T.min(row_tile_size, n - row)
166	-
167	- if row_tile_size > 1:
168	- T.copy(x[row : row + rem, :, col_base : col_base + block_size], x_ub[: rem * hc_mult, :])
169	- else:
170	- T.copy(x[row, :, col_base : col_base + block_size], x_ub[:, :])
171	-
172	- if hc_mult > 1:
173	- with T.SimdVF():
174	- acc = T.simd.alloc_local((1,), T.float32)
175	- for c in T.serial(block_size // vec_size):
176	- off = c * vec_size
177	- acc[0] = T.simd.vcvt(T.simd.vld(x_ub[0, off], dist='UNPK_B16'), T.float32, part=0)
178	- for j in T.serial(1, hc_mult):
179	- val = T.simd.vcvt(T.simd.vld(x_ub[j, off], dist='UNPK_B16'), T.float32, part=0)
180	- acc[0] = T.simd.vadd(acc[0], val)
181	- acc[0] = T.simd.vmuls(acc[0], T.float32(1.0 / hc_mult))
182	- T.simd.vsts(x_ub[0, off], T.simd.vcvt(acc[0], T.bfloat16), dist='PK_B32')
183	-
184	- if row_tile_size > 1:
185	- T.copy(x_ub[: rem * hc_mult, :], out[row : row + rem, col_base : col_base + block_size])
186	- else:
187	- T.copy(x_ub[0, :], out[row, col_base : col_base + block_size])
188	-
189	- return _dflash_input_store_ascend

Msrc/cann_ops_tilelang/sampling/sampler.py+2-317

	@@ -1,130 +1,10 @@
1	-"""Sampling and speculative-decoding helpers for Ascend."""	1	+"""Gather logits for the current decoding step on Ascend."""
2		2
3	import torch	3	import torch
4		4
5	from cann_ops_tilelang.config import require_ascend	5	from cann_ops_tilelang.config import require_ascend
6		6
7	-from .sampler_kernel import (	7	+from .sampler_kernel import _get_logits_cur_by_step_tl_ascend
8	- _batch_find_value_by_index_tl_ascend,
9	- _compute_confidence_error_tl_ascend,
10	- _find_value_by_index_tl_ascend,
11	- _get_accept_result_tl_ascend,
12	- _get_logits_cur_by_step_tl_ascend,
13	- _sampled_tokens_update_tl_ascend,
14	- _scatter_verified_tl_ascend,
15	- _update_verify_states_tl_ascend,
16	-)
17	-
18	-
19	-def _require_shape(name: str, tensor: torch.Tensor, expected: tuple[int, ...]) -> None:
20	- if tuple(tensor.shape) != expected:
21	- raise ValueError(f'{name} must have shape {expected}, got {tuple(tensor.shape)}')
22	-
23	-
24	-def find_value_by_index(index: torch.Tensor, value: torch.Tensor, q_index: torch.Tensor) -> torch.Tensor:
25	- """Find one value for each requested token index."""
26	- require_ascend()
27	- force_topk = index.shape[-1]
28	- assert index.stride(0) % force_topk == 0 and index.stride(1) % force_topk == 0
29	- assert index.stride(-1) == 1 and value.stride(-1) == 1 and (q_index.stride(-1) == 1)
30	- out = torch.empty((index.shape[0], index.shape[1]), dtype=torch.float32, device=index.device)
31	- assert index.shape[-1] % 512 == 0
32	- get_kernel = _find_value_by_index_tl_ascend
33	- kernel = get_kernel(force_topk)
34	- kernel(index, value, q_index, out)
35	- return out
36	-
37	-
38	-def batchfind_value_by_index_ref(
39	- index: torch.Tensor, value: torch.Tensor, q_index: torch.Tensor, vocab_size: int, use_pdl: bool = False
40	-) -> torch.Tensor:
41	- require_ascend()
42	- num_steps, num_requests = (index.shape[0], index.shape[1])
43	- num_queries = q_index.shape[-1]
44	- index_flat = index.reshape(-1, index.shape[-1])
45	- value_flat = value.reshape(-1, index.shape[-1])
46	- q_index_flat = q_index.reshape(-1, num_queries)
47	- pool = torch.full((num_steps * num_requests, vocab_size), float('-inf'), dtype=torch.float32, device=index.device)
48	- pool.scatter_(1, q_index_flat, float('-inf'))
49	- pool.scatter_(1, index_flat, value_flat)
50	- out_flat = pool.gather(-1, q_index_flat)
51	- return out_flat.view(num_steps, num_requests, num_queries)
52	-
53	-
54	-def batch_find_value_by_index(
55	- index: torch.Tensor, value: torch.Tensor, q_index: torch.Tensor, vocab_size: int, use_pdl: bool = False
56	-) -> torch.Tensor:
57	- """Find values for a batch of requested token indices."""
58	- require_ascend()
59	- assert index.stride(-1) == 1 and value.stride(-1) == 1
60	- ft, qd = (index.shape[-1], q_index.shape[-1])
61	- qcap = None
62	- qcap = 2048 if qd > ft else None
63	- use_shared_memory = (
64	- ft > 0
65	- and ft % 512 == 0
66	- and vocab_size <= 2**20
67	- and (qd <= ft <= 4096 or (qcap is not None and ft in (512, 1024) and (qd <= 2048)))
68	- )
69	- if use_shared_memory:
70	- num_steps, num_requests = (index.shape[0], index.shape[1])
71	- out = torch.empty((num_steps, num_requests, qd), dtype=torch.float32, device=index.device)
72	- kernel = _batch_find_value_by_index_tl_ascend(ft, vocab_size, use_pdl, qcap)
73	- kernel(index, value, q_index, out)
74	- return out
75	- return batchfind_value_by_index_ref(index, value, q_index, vocab_size)
76	-
77	-
78	-def get_accept_result(
79	- accept_result_verify: torch.Tensor,
80	- so_valid: torch.Tensor \| None,
81	- golden_output: torch.Tensor \| None,
82	- last_next_n_token_ids: torch.Tensor,
83	- gold_accepted: torch.Tensor \| None,
84	- is_verify: torch.Tensor,
85	- sample_reject_cur: torch.Tensor,
86	- fresh_sample_final: torch.Tensor,
87	-) -> tuple[torch.Tensor, torch.Tensor]:
88	- require_ascend()
89	- if accept_result_verify.ndim != 2:
90	- raise ValueError('accept_result_verify must be two-dimensional')
91	- m, n = accept_result_verify.shape
92	- for name, tensor in (
93	- ('last_next_n_token_ids', last_next_n_token_ids),
94	- ('sample_reject_cur', sample_reject_cur),
95	- ):
96	- _require_shape(name, tensor, (m, n))
97	- if is_verify.ndim != 2 or is_verify.shape[1] != n or is_verify.shape[0] < m:
98	- raise ValueError('is_verify must have shape (max_seqlen_q, N) with max_seqlen_q >= M')
99	- _require_shape('fresh_sample_final', fresh_sample_final, tuple(is_verify.shape))
100	- if so_valid is not None:
101	- _require_shape('so_valid', so_valid, (m, n))
102	- if (golden_output is None) != (gold_accepted is None):
103	- raise ValueError('golden_output and gold_accepted must be provided together')
104	- if golden_output is not None:
105	- if golden_output.ndim != 2 or golden_output.shape[1] != n or golden_output.shape[0] > m:
106	- raise ValueError('golden_output must have shape (golden_steps, N) with golden_steps <= M')
107	- _require_shape('gold_accepted', gold_accepted, tuple(golden_output.shape))
108	- max_seqlen_q = is_verify.shape[0]
109	- sample_results_cur = torch.empty((max_seqlen_q, n), dtype=torch.int32, device=accept_result_verify.device)
110	- accept_result = torch.empty((max_seqlen_q, n), dtype=torch.bool, device=accept_result_verify.device)
111	- has_so_valid_flag = so_valid is not None
112	- has_golden_output_flag = golden_output is not None
113	- get_kernel = _get_accept_result_tl_ascend
114	- kernel = get_kernel(has_so_valid_flag, has_golden_output_flag)
115	- kernel(
116	- accept_result_verify,
117	- so_valid,
118	- golden_output,
119	- last_next_n_token_ids,
120	- gold_accepted,
121	- is_verify,
122	- sample_reject_cur,
123	- fresh_sample_final,
124	- sample_results_cur,
125	- accept_result,
126	- )
127	- return (sample_results_cur, accept_result)
128		8
129		9
130	def get_logits_cur_by_step(logits: torch.Tensor, ftp: torch.Tensor, has_token: torch.Tensor, step: int) -> torch.Tensor:	10	def get_logits_cur_by_step(logits: torch.Tensor, ftp: torch.Tensor, has_token: torch.Tensor, step: int) -> torch.Tensor:
	@@ -136,198 +16,3 @@ def get_logits_cur_by_step(logits: torch.Tensor, ftp: torch.Tensor, has_token: t
136	kernel = _get_logits_cur_by_step_tl_ascend(vocab_size)	16	kernel = _get_logits_cur_by_step_tl_ascend(vocab_size)
137	kernel(logits, ftp, has_token, out, step)	17	kernel(logits, ftp, has_token, out, step)
138	return out	18	return out
139	-
140	-
141	-def update_verify_states(
142	- accept_result_all: torch.Tensor,
143	- has_token_all: torch.Tensor,
144	- seqlen_q: torch.Tensor,
145	- sample_results_cur_all: torch.Tensor,
146	- cur_step_mtp_value: int,
147	-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
148	- require_ascend()
149	- num_steps, num_requests = accept_result_all.shape
150	- device = accept_result_all.device
151	- accepted_all = torch.empty((num_steps, num_requests), dtype=torch.int32, device=device)
152	- is_final_all = torch.empty((num_steps, num_requests), dtype=torch.bool, device=device)
153	- last_offset = torch.empty(num_requests, dtype=torch.int32, device=device)
154	- num_gen = torch.empty(num_steps, dtype=torch.int32, device=device)
155	- num_gen.zero_()
156	- kernel = _update_verify_states_tl_ascend()
157	- kernel(
158	- accept_result_all,
159	- has_token_all,
160	- seqlen_q,
161	- sample_results_cur_all,
162	- cur_step_mtp_value,
163	- accepted_all,
164	- is_final_all,
165	- last_offset,
166	- num_gen,
167	- )
168	- return (accepted_all, is_final_all, last_offset, num_gen[:cur_step_mtp_value])
169	-
170	-
171	-def scatter_verified(
172	- is_verify: torch.Tensor,
173	- ftp: torch.Tensor,
174	- accepted_tokens: torch.Tensor,
175	- token_logprobs_cur: torch.Tensor,
176	- sum_smaller_probs_cur: torch.Tensor \| None,
177	- last_nextn_token_logprobs: torch.Tensor \| None,
178	- top_logprobs_cur: torch.Tensor \| None,
179	- top_token_ids_cur: torch.Tensor \| None,
180	- total_q: int,
181	-) -> tuple:
182	- """Scatter verified tokens and their optional statistics."""
183	- require_ascend()
184	- if is_verify.ndim != 2:
185	- raise ValueError('is_verify must be two-dimensional')
186	- steps, req_cnt = is_verify.shape
187	- _require_shape('ftp', ftp, (req_cnt,))
188	- _require_shape('accepted_tokens', accepted_tokens, (steps, req_cnt))
189	- _require_shape('token_logprobs_cur', token_logprobs_cur, (steps, req_cnt))
190	- if (sum_smaller_probs_cur is None) != (last_nextn_token_logprobs is None):
191	- raise ValueError('sum_smaller_probs_cur and last_nextn_token_logprobs must be provided together')
192	- if sum_smaller_probs_cur is not None:
193	- _require_shape('sum_smaller_probs_cur', sum_smaller_probs_cur, (steps, req_cnt))
194	- _require_shape('last_nextn_token_logprobs', last_nextn_token_logprobs, (steps, req_cnt))
195	- if (top_logprobs_cur is None) != (top_token_ids_cur is None):
196	- raise ValueError('top_logprobs_cur and top_token_ids_cur must be provided together')
197	- if top_logprobs_cur is not None:
198	- if top_logprobs_cur.ndim != 3 or top_logprobs_cur.shape[:2] != (steps, req_cnt):
199	- raise ValueError('top_logprobs_cur must have shape (steps, req_cnt, topk_n)')
200	- if top_logprobs_cur.shape[2] <= 0:
201	- raise ValueError('topk_n must be positive when top-k tensors are provided')
202	- _require_shape('top_token_ids_cur', top_token_ids_cur, tuple(top_logprobs_cur.shape))
203	- device = is_verify.device
204	- is_rl_or_eval = sum_smaller_probs_cur is not None
205	- has_topk_n = top_logprobs_cur is not None
206	- topk_n = top_logprobs_cur.shape[2] if top_logprobs_cur is not None else 0
207	- topk_stage_bytes = 8 * ((topk_n + 15) // 16 * 16)
208	- if topk_stage_bytes > 240000:
209	- raise ValueError('topk_n exceeds the scatter_verified kernel UB capacity.')
210	- verified_token_ids = torch.empty(total_q, dtype=torch.int32, device=device)
211	- verified_token_logprobs = torch.empty(total_q, dtype=torch.float32, device=device)
212	- verified_sum_smaller_probs = torch.empty(total_q, dtype=torch.float32, device=device) if is_rl_or_eval else None
213	- accepted_token_last_logprobs = torch.empty(total_q, dtype=torch.float32, device=device) if is_rl_or_eval else None
214	- verified_topk_logps = torch.empty(total_q, topk_n, dtype=torch.float32, device=device) if has_topk_n else None
215	- verified_topk_tokens = torch.empty(total_q, topk_n, dtype=torch.int32, device=device) if has_topk_n else None
216	- top_token_ids_cur_dtype = torch.int32 if top_token_ids_cur is None else top_token_ids_cur.dtype
217	- kernel = _scatter_verified_tl_ascend(is_rl_or_eval, has_topk_n, topk_n, top_token_ids_cur_dtype)
218	- kernel(
219	- is_verify,
220	- ftp,
221	- accepted_tokens,
222	- token_logprobs_cur,
223	- sum_smaller_probs_cur,
224	- last_nextn_token_logprobs,
225	- verified_token_ids,
226	- verified_token_logprobs,
227	- verified_sum_smaller_probs,
228	- accepted_token_last_logprobs,
229	- top_logprobs_cur,
230	- top_token_ids_cur,
231	- verified_topk_logps,
232	- verified_topk_tokens,
233	- )
234	- return (
235	- verified_token_ids,
236	- verified_token_logprobs,
237	- verified_sum_smaller_probs,
238	- accepted_token_last_logprobs,
239	- verified_topk_logps,
240	- verified_topk_tokens,
241	- )
242	-
243	-
244	-def sampled_tokens_update(
245	- is_final: torch.Tensor,
246	- sample_results_cur: torch.Tensor,
247	- token_logprobs_cur: torch.Tensor,
248	- sum_smaller_probs_cur: torch.Tensor \| None,
249	- top_logprobs_cur: torch.Tensor \| None,
250	- top_token_ids_cur: torch.Tensor \| None,
251	- token_ids_out: torch.Tensor,
252	- token_logprobs_out: torch.Tensor,
253	- sum_smaller_probs_out: torch.Tensor \| None,
254	- top_logps_out: torch.Tensor \| None,
255	- top_tokens_out: torch.Tensor \| None,
256	-) -> None:
257	- """Scatter the last final step for each request into sampled-token outputs."""
258	- require_ascend()
259	- if is_final.ndim != 2:
260	- raise ValueError('is_final must be two-dimensional')
261	- steps, req_cnt = is_final.shape
262	- _require_shape('sample_results_cur', sample_results_cur, (steps, req_cnt))
263	- _require_shape('token_logprobs_cur', token_logprobs_cur, (steps, req_cnt))
264	- _require_shape('token_ids_out', token_ids_out, (req_cnt,))
265	- _require_shape('token_logprobs_out', token_logprobs_out, (req_cnt,))
266	- if (sum_smaller_probs_cur is None) != (sum_smaller_probs_out is None):
267	- raise ValueError('sum_smaller_probs_cur and sum_smaller_probs_out must be provided together')
268	- if sum_smaller_probs_cur is not None:
269	- _require_shape('sum_smaller_probs_cur', sum_smaller_probs_cur, (steps, req_cnt))
270	- _require_shape('sum_smaller_probs_out', sum_smaller_probs_out, (req_cnt,))
271	- if (top_logprobs_cur is None) != (top_logps_out is None):
272	- raise ValueError('top_logprobs_cur and top_logps_out must be provided together')
273	- if (top_token_ids_cur is None) != (top_tokens_out is None):
274	- raise ValueError('top_token_ids_cur and top_tokens_out must be provided together')
275	- has_sum_smaller = sum_smaller_probs_cur is not None and sum_smaller_probs_out is not None
276	- has_top_logprobs = top_logprobs_cur is not None and top_logps_out is not None
277	- has_top_token_ids = top_token_ids_cur is not None and top_tokens_out is not None
278	- topk_n = top_logprobs_cur.shape[-1] if has_top_logprobs else 0
279	- if has_top_logprobs:
280	- if top_logprobs_cur.ndim != 3 or top_logprobs_cur.shape[:2] != (steps, req_cnt) or topk_n <= 0:
281	- raise ValueError('top_logprobs_cur must have shape (steps, req_cnt, positive topk_n)')
282	- _require_shape('top_logps_out', top_logps_out, (req_cnt, topk_n))
283	- if has_top_token_ids:
284	- if top_token_ids_cur.ndim != 3 or top_token_ids_cur.shape[:2] != (steps, req_cnt):
285	- raise ValueError('top_token_ids_cur must have shape (steps, req_cnt, topk_n)')
286	- token_topk_n = top_token_ids_cur.shape[-1]
287	- if token_topk_n <= 0:
288	- raise ValueError('topk_n must be positive when top-token tensors are provided')
289	- _require_shape('top_tokens_out', top_tokens_out, (req_cnt, token_topk_n))
290	- if has_top_logprobs and token_topk_n != topk_n:
291	- raise ValueError('top-logprob and top-token tensors must use the same top-k size')
292	- topk_n = token_topk_n
293	- stage_bytes = 576 + (256 if has_sum_smaller else 0) + 4 * topk_n * (int(has_top_logprobs) + int(has_top_token_ids))
294	- if stage_bytes > 240000:
295	- raise ValueError('topk_n exceeds the sampled_tokens_update kernel UB capacity.')
296	- top_token_ids_cur_dtype = torch.int32 if top_token_ids_cur is None else top_token_ids_cur.dtype
297	- kernel = _sampled_tokens_update_tl_ascend(
298	- has_sum_smaller,
299	- has_top_logprobs,
300	- has_top_token_ids,
301	- topk_n,
302	- top_token_ids_cur_dtype,
303	- )
304	- kernel(
305	- is_final,
306	- sample_results_cur,
307	- token_logprobs_cur,
308	- sum_smaller_probs_cur,
309	- top_logprobs_cur,
310	- top_token_ids_cur,
311	- token_ids_out,
312	- token_logprobs_out,
313	- sum_smaller_probs_out,
314	- top_logps_out,
315	- top_tokens_out,
316	- )
317	-
318	-
319	-def compute_confidence_error(
320	- last_nextn_token_logprobs: torch.Tensor,
321	- nextn_token_logprobs: torch.Tensor,
322	- confidence: torch.Tensor,
323	- verify_masks: torch.Tensor,
324	-) -> torch.Tensor:
325	- require_ascend()
326	- device = last_nextn_token_logprobs.device
327	- num_steps = confidence.shape[0]
328	- out = torch.empty(1, dtype=torch.float32, device=device)
329	- temp = torch.empty(num_steps, dtype=torch.float32, device=device)
330	- out.zero_()
331	- kernel = _compute_confidence_error_tl_ascend()
332	- kernel(last_nextn_token_logprobs, nextn_token_logprobs, confidence, verify_masks, out, temp)
333	- return out.view(())

Msrc/cann_ops_tilelang/sampling/sampler_kernel.py+2-745

Mtests/attention/test_get_kv_index.py+1-293

	@@ -4,102 +4,12 @@ from typing import Any
4	import pytest	4	import pytest
5	import torch	5	import torch
6		6
7	-from cann_ops_tilelang.attention.get_kv_index import (	7	+from cann_ops_tilelang.attention.get_kv_index import get_fused_dflash_get_input
8	- get_dflash_mainx_slot_mapping,
9	- get_fused_dflash_get_input,
10	- get_kv_ca_index_decode,
11	- get_kv_ca_index_prefill,
12	- get_kv_swa_index_decode,
13	-)
14	from cann_ops_tilelang.config import get_device	8	from cann_ops_tilelang.config import get_device
15		9
16	pytestmark = pytest.mark.npu	10	pytestmark = pytest.mark.npu
17		11
18		12
19	-def get_kv_ca_index_prefill_ref(kv_ca_index_offset_len: torch.Tensor, kv_ca_index_length: int) -> torch.Tensor:
20	- kv_ca_index = torch.full((kv_ca_index_offset_len.shape[0], kv_ca_index_length), -1, dtype=torch.int32, device='npu')
21	- arange = torch.arange(kv_ca_index.shape[1], dtype=torch.int32, device='npu').unsqueeze(0)
22	- kv_ca_index = torch.where(
23	- arange < kv_ca_index_offset_len[:, 1].unsqueeze(1),
24	- kv_ca_index_offset_len[:, 0].unsqueeze(1) + arange,
25	- kv_ca_index,
26	- )
27	- return kv_ca_index
28	-
29	-
30	-def get_kv_ca_index_decode_ref(
31	- kv_ca_index_offset_len: torch.Tensor, kv_ca_index_length: int, token_per_block: int, block_tables: torch.Tensor
32	-) -> torch.Tensor:
33	- kv_ca_index = torch.full((kv_ca_index_offset_len.shape[0], kv_ca_index_length), -1, dtype=torch.int32, device='npu')
34	- arange = torch.arange(kv_ca_index.shape[1], dtype=torch.int32, device='npu').unsqueeze(0)
35	- arange = arange.minimum(kv_ca_index_offset_len[:, 1].unsqueeze(1))
36	- pos = kv_ca_index_offset_len[:, 0].unsqueeze(1) + arange
37	- kv_ca_index = torch.where(
38	- arange < kv_ca_index_offset_len[:, 1].unsqueeze(1),
39	- block_tables.flatten().index_select(0, (pos // token_per_block).flatten().long()).reshape_as(pos)
40	- * token_per_block
41	- + pos % token_per_block,
42	- kv_ca_index,
43	- )
44	- return kv_ca_index
45	-
46	-
47	-def get_kv_swa_index_decode_ref(
48	- kv_swa_index_offset_len: torch.Tensor, window_size: int, swa_block_size: int
49	-) -> torch.Tensor:
50	- kv_swa_index = torch.full((kv_swa_index_offset_len.shape[0], window_size), -1, dtype=torch.int32, device='npu')
51	- offset = kv_swa_index_offset_len[:, 0].unsqueeze(1)
52	- length = kv_swa_index_offset_len[:, 1].unsqueeze(1)
53	- arange = torch.arange(kv_swa_index.shape[1], dtype=torch.int32, device='npu').unsqueeze(0)
54	- kv_swa_index = torch.where(
55	- arange < length, offset // swa_block_size * swa_block_size + (offset + arange) % swa_block_size, kv_swa_index
56	- )
57	- return kv_swa_index
58	-
59	-
60	-@pytest.mark.parametrize(
61	- 'n,index_length,block_size,mode',
62	- [
63	- (1, 128, 64, 'swa_decode'),
64	- (128, 128, 128, 'swa_decode'),
65	- (128, 128, 133, 'swa_decode'),
66	- (1, 8192, 128, 'ca_decode'),
67	- (1, 256, 1, 'ca_decode'),
68	- (128, 8192, 1, 'ca_decode'),
69	- (1, 8192, 32, 'ca_decode'),
70	- (128, 262144, 32, 'ca_decode'),
71	- (1, 256, 0, 'ca_prefill'),
72	- (128, 8192, 0, 'ca_prefill'),
73	- (128, 262144, 0, 'ca_prefill'),
74	- ],
75	-)
76	-def test_kv_index(n: int, index_length: int, block_size: int, mode: str) -> None:
77	- index_offset_len = torch.empty((n, 2), dtype=torch.int32, device='npu')
78	- index_offset_len[:, 0] = torch.randint(0, 2000000000, (n,), dtype=torch.int32, device='npu')
79	- index_offset_len[:, 1] = torch.randint(0, index_length * 2, (n,), dtype=torch.int32, device='npu')
80	- block_tables = torch.randint(0, 10000000, (2100000000,), dtype=torch.int32, device='npu')
81	-
82	- def test(pytorch_func: Callable[..., Any], func: Callable[..., Any], *args: Any) -> None:
83	- pytorch_result = pytorch_func(*args)
84	- result = func(*args)
85	- assert torch.equal(pytorch_result, result), f'{pytorch_result} vs {result}'
86	-
87	- match mode:
88	- case 'ca_decode':
89	- test(
90	- get_kv_ca_index_decode_ref,
91	- get_kv_ca_index_decode,
92	- index_offset_len,
93	- index_length,
94	- block_size,
95	- block_tables,
96	- )
97	- case 'swa_decode':
98	- test(get_kv_swa_index_decode_ref, get_kv_swa_index_decode, index_offset_len, index_length, block_size)
99	- case 'ca_prefill':
100	- test(get_kv_ca_index_prefill_ref, get_kv_ca_index_prefill, index_offset_len, index_length)
101	-
102	-
103	def get_fused_dflash_get_input_ref(	13	def get_fused_dflash_get_input_ref(
104	slot_mapping_swa: torch.Tensor,	14	slot_mapping_swa: torch.Tensor,
105	positions: torch.Tensor,	15	positions: torch.Tensor,
	@@ -237,205 +147,3 @@ def test_fused_dflash_get_input_benchmark(
237	time_us=t_us,	147	time_us=t_us,
238	extras={'bandwidth_gbs': data_bytes / 1000000000.0 / (t_us / 1000000.0)},	148	extras={'bandwidth_gbs': data_bytes / 1000000000.0 / (t_us / 1000000.0)},
239	)	149	)
240	-
241	-
242	-def get_dflash_mainx_slot_mapping_ref(
243	- ftp: torch.Tensor, valid_offsets: torch.Tensor, slot_mapping_swa: torch.Tensor, real_cnt: int
244	-) -> torch.Tensor:
245	- """PyTorch reference implementation for dflash_mainx_slot_mapping."""
246	- total_q = slot_mapping_swa.shape[0]
247	- device = slot_mapping_swa.device
248	- req_idx = torch.zeros(total_q, dtype=torch.long, device=device)
249	- if real_cnt > 1:
250	- req_start_idx = ftp[1:real_cnt].long()
251	- req_idx.scatter_(dim=0, index=req_start_idx, src=torch.ones_like(req_start_idx, dtype=req_idx.dtype))
252	- req_idx = req_idx.cumsum(0)
253	- offset_in_req = torch.arange(total_q, device=device, dtype=torch.long) - ftp[req_idx]
254	- mask = offset_in_req <= valid_offsets[req_idx]
255	- return torch.where(mask, slot_mapping_swa, torch.full_like(slot_mapping_swa, -1))
256	-
257	-
258	-@pytest.mark.parametrize('real_cnt', [7, 17, 157, 233, 297, 777, 1024, 2048, 8192])
259	-def test_dflash_mainx_slot_mapping(real_cnt: int) -> None:
260	- max_seq_len = 8
261	- valid_offsets = torch.randint(0, max_seq_len - 1, (real_cnt,), device=get_device(), dtype=torch.int32)
262	- seq_lens = torch.max(
263	- valid_offsets, torch.randint(1, max_seq_len, (real_cnt,), device=get_device(), dtype=torch.int32)
264	- )
265	- ftp = torch.zeros(real_cnt, device=get_device(), dtype=torch.int32)
266	- if real_cnt > 1:
267	- ftp[1:] = seq_lens[:-1].cumsum(0)
268	- total_q = int(ftp[-1].item() + seq_lens[-1].item())
269	- slot_mapping_swa = torch.randint(0, 100000, (total_q,), device=get_device(), dtype=torch.int32)
270	- valid_offsets = valid_offsets.int()
271	- args = (ftp, valid_offsets, slot_mapping_swa, real_cnt)
272	- pytorch_result = get_dflash_mainx_slot_mapping_ref(*args)
273	- result = get_dflash_mainx_slot_mapping(*args)
274	- assert torch.equal(pytorch_result, result), f'Mismatch!\npytorch:\n{pytorch_result}\ntilelang:\n{result}'
275	-
276	-
277	-@pytest.mark.benchmark
278	-@pytest.mark.parametrize(
279	- 'n,index_length,block_size,mode',
280	- [
281	- (1024, 128, 128, 'swa_decode'),
282	- (1024, 128, 133, 'swa_decode'),
283	- (1024, 8192, 1, 'ca_decode'),
284	- (1024, 8192, 32, 'ca_decode'),
285	- (1024, 262144, 32, 'ca_decode'),
286	- (1024, 256, 0, 'ca_prefill'),
287	- (1024, 8192, 0, 'ca_prefill'),
288	- (1024, 262144, 0, 'ca_prefill'),
289	- ],
290	-)
291	-def test_kv_index_benchmark(
292	- n: int,
293	- index_length: int,
294	- block_size: int,
295	- mode: str,
296	- benchmark_timer: Callable[..., Any],
297	- benchmark_record: Callable[..., Any],
298	-) -> Any:
299	- index_offset_len = torch.empty((n, 2), dtype=torch.int32, device='npu')
300	- index_offset_len[:, 0] = torch.randint(0, 2000000000, (n,), dtype=torch.int32, device='npu')
301	- index_offset_len[:, 1] = torch.randint(0, index_length * 2, (n,), dtype=torch.int32, device='npu')
302	- block_tables = torch.randint(0, 10000000, (2100000000,), dtype=torch.int32, device='npu')
303	- match mode:
304	- case 'ca_decode':
305	- func = get_kv_ca_index_decode
306	- args = (index_offset_len, index_length, block_size, block_tables)
307	- case 'swa_decode':
308	- func = get_kv_swa_index_decode
309	- args = (index_offset_len, index_length, block_size)
310	- case 'ca_prefill':
311	- func = get_kv_ca_index_prefill
312	- args = (index_offset_len, index_length)
313	- func(*args)
314	- torch.npu.synchronize()
315	- t_us = benchmark_timer(lambda: func(*args))
316	- B = 8192 if index_length <= 8192 else 16384
317	- if mode == 'ca_decode':
318	- blocks_per_token = (index_length + B - 1) // B
319	- table_bytes = n * (index_length // block_size + blocks_per_token) * 4
320	- total_bytes = n * index_length * 4 + table_bytes + n * 8
321	- else:
322	- total_bytes = n * index_length * 4 + n * 8
323	- benchmark_record(
324	- kernel='get_kv_index',
325	- operation=mode,
326	- params={'block_size': block_size, 'index_length': index_length, 'n': n},
327	- time_us=t_us,
328	- extras={'bandwidth_gbs': total_bytes / 1000000000.0 / (t_us / 1000000.0)},
329	- )
330	-
331	-
332	-@pytest.mark.benchmark
333	-@pytest.mark.parametrize('real_cnt', [128, 1024, 8192])
334	-def test_dflash_mainx_slot_mapping_benchmark(
335	- real_cnt: int, benchmark_timer: Callable[..., Any], benchmark_record: Callable[..., Any]
336	-) -> Any:
337	- max_seq_len = 8
338	- valid_offsets = torch.randint(0, max_seq_len - 1, (real_cnt,), device=get_device(), dtype=torch.int32)
339	- seq_lens = torch.max(
340	- valid_offsets, torch.randint(1, max_seq_len, (real_cnt,), device=get_device(), dtype=torch.int32)
341	- )
342	- ftp = torch.zeros(real_cnt, device=get_device(), dtype=torch.int32)
343	- if real_cnt > 1:
344	- ftp[1:] = seq_lens[:-1].cumsum(0)
345	- total_q = int(ftp[-1].item() + seq_lens[-1].item())
346	- slot_mapping_swa = torch.randint(0, 100000, (total_q,), device=get_device(), dtype=torch.int32)
347	- valid_offsets = valid_offsets.int()
348	- args = (ftp, valid_offsets, slot_mapping_swa, real_cnt)
349	- t_us = benchmark_timer(lambda: get_dflash_mainx_slot_mapping(*args))
350	- result = get_dflash_mainx_slot_mapping(*args)
351	- data_bytes = result.numel() * result.element_size() + slot_mapping_swa.numel() * slot_mapping_swa.element_size()
352	- benchmark_record(
353	- kernel='get_kv_index',
354	- operation='dflash_mainx_slot_mapping',
355	- params={'real_cnt': real_cnt},
356	- time_us=t_us,
357	- extras={'bandwidth_gbs': data_bytes / 1000000000.0 / (t_us / 1000000.0)},
358	- )
359	-
360	-
361	-@pytest.mark.benchmark
362	-@pytest.mark.parametrize(
363	- 'n,index_length,block_size',
364	- [
365	- (1, 256, 1),
366	- (1024, 256, 1),
367	- (1, 8192, 1),
368	- (1024, 8192, 1),
369	- (1, 8192, 32),
370	- (1024, 8192, 32),
371	- (1, 262144, 32),
372	- (1024, 262144, 32),
373	- ],
374	-)
375	-def test_kv_ca_index_decode_benchmark(
376	- n: int,
377	- index_length: int,
378	- block_size: int,
379	- benchmark_timer: Callable[..., Any],
380	- benchmark_record: Callable[..., Any],
381	-) -> Any:
382	- index_offset_len = torch.zeros((n, 2), dtype=torch.int32, device='npu')
383	- for i in range(n):
384	- index_offset_len[i, 0] = i * index_length
385	- index_offset_len[i, 1] = index_length
386	- block_tables = torch.arange(n * index_length // block_size, dtype=torch.int32, device='npu')
387	- t_us = benchmark_timer(lambda: get_kv_ca_index_decode(index_offset_len, index_length, block_size, block_tables))
388	- total_bytes = n * index_length * 4 + block_tables.numel() * block_tables.element_size() + n * 8
389	- benchmark_record(
390	- kernel='get_kv_index',
391	- operation='ca_index_decode',
392	- params={'block_size': block_size, 'index_length': index_length, 'n': n},
393	- time_us=t_us,
394	- extras={'bandwidth_gbs': total_bytes / 1000000000.0 / (t_us / 1000000.0)},
395	- )
396	-
397	-
398	-@pytest.mark.benchmark
399	-@pytest.mark.parametrize(
400	- 'n,index_length,block_size', [(1, 128, 128), (128, 128, 128), (1024, 128, 128), (128, 128, 133)]
401	-)
402	-def test_kv_swa_index_decode_benchmark(
403	- n: int,
404	- index_length: int,
405	- block_size: int,
406	- benchmark_timer: Callable[..., Any],
407	- benchmark_record: Callable[..., Any],
408	-) -> Any:
409	- index_offset_len = torch.zeros((n, 2), dtype=torch.int32, device='npu')
410	- for i in range(n):
411	- index_offset_len[i, 0] = i * index_length
412	- index_offset_len[i, 1] = index_length
413	- t_us = benchmark_timer(lambda: get_kv_swa_index_decode(index_offset_len, index_length, block_size))
414	- total_bytes = n * index_length * 4 + n * 8
415	- benchmark_record(
416	- kernel='get_kv_index',
417	- operation='swa_index_decode',
418	- params={'block_size': block_size, 'index_length': index_length, 'n': n},
419	- time_us=t_us,
420	- extras={'bandwidth_gbs': total_bytes / 1000000000.0 / (t_us / 1000000.0)},
421	- )
422	-
423	-
424	-@pytest.mark.benchmark
425	-@pytest.mark.parametrize('n,index_length', [(32768, 256), (32768, 8192), (1024, 262144)])
426	-def test_kv_ca_index_prefill_benchmark(
427	- n: int, index_length: int, benchmark_timer: Callable[..., Any], benchmark_record: Callable[..., Any]
428	-) -> Any:
429	- index_offset_len = torch.zeros((n, 2), dtype=torch.int32, device='npu')
430	- for i in range(n):
431	- index_offset_len[i, 0] = i * index_length
432	- index_offset_len[i, 1] = index_length
433	- t_us = benchmark_timer(lambda: get_kv_ca_index_prefill(index_offset_len, index_length))
434	- total_bytes = n * index_length * 4 + n * 8
435	- benchmark_record(
436	- kernel='get_kv_index',
437	- operation='ca_index_prefill',
438	- params={'index_length': index_length, 'n': n},
439	- time_us=t_us,
440	- extras={'bandwidth_gbs': total_bytes / 1000000000.0 / (t_us / 1000000.0)},
441	- )

Dtests/compression/test_compress.py+0-694

	@@ -1,694 +0,0 @@
1	-import random
2	-from collections.abc import Callable
3	-from typing import Any
4	-
5	-import pytest
6	-import torch
7	-
8	-from cann_ops_tilelang.compression.compress import (
9	- compress_and_update_state_decode,
10	- compress_and_update_state_prefill,
11	-)
12	-
13	-pytestmark = pytest.mark.npu
14	-
15	-
16	-@pytest.fixture(autouse=True)
17	-def _restore_default_device():
18	- """Keep compression tests from changing tensor placement in other modules."""
19	- yield
20	- torch.set_default_device('cpu')
21	-
22	-
23	-def assert_close(
24	- actual: torch.Tensor,
25	- expected: torch.Tensor,
26	- name: str = '',
27	- eps: float = 1e-5,
28	-) -> None:
29	- try:
30	- torch.testing.assert_close(actual, expected, rtol=eps, atol=eps)
31	- except AssertionError as error:
32	- raise AssertionError(f'{name}: {error}') from error
33	-
34	-
35	-def compress_and_update_state_prefill_ref(
36	- score: torch.Tensor,
37	- latent: torch.Tensor,
38	- ape: torch.Tensor \| None,
39	- positions: torch.Tensor \| None,
40	- state_cache: torch.Tensor,
41	- state_block_idx: torch.Tensor,
42	- compressed_batch_mapping: torch.Tensor,
43	- cu_seqlens_q: torch.Tensor,
44	- cu_seqlens_k: torch.Tensor,
45	- cu_compressed_seqlens_q: torch.Tensor,
46	- mirror_cp_reverse_index: torch.Tensor \| None = None,
47	- compress_ratio: int \| None = None,
48	-) -> torch.Tensor:
49	- """
50	- Reference implementation of compress_and_update_state_prefill with
51	- mirror_cp_reverse_index support.
52	-
53	- Args:
54	- score: [total_q \| cp_size * cp_mirror_len, overlap_ratio, dim], fp32
55	- latent: [total_q \| cp_size * cp_mirror_len, overlap_ratio, dim], fp32
56	- ape: [compress_ratio, overlap_ratio, dim], fp32, or None
57	- positions: [total_q], int32, or None
58	- state_cache: [num_state_slots, state_cache_size, 2, overlap_ratio, dim], fp32
59	- state_block_idx: [num_seqs], int32
60	- compressed_batch_mapping: [total_num_compressed_tokens], int32
61	- cu_seqlens_q: [num_seqs + 1], int32
62	- cu_seqlens_k: [num_seqs + 1], int32
63	- cu_compressed_seqlens_q: [num_seqs + 1], int32. Each query must be >=1,
64	- or state_cache will not be updated.
65	- mirror_cp_reverse_index: [total_q], int32, in [0, cp_size * cp_mirror_len)
66	-
67	- Returns:
68	- kv_compressed: [total_num_compressed_tokens, dim], bf16
69	- """
70	- has_ape = ape is not None
71	- if has_ape:
72	- assert positions is not None
73	- compress_ratio, overlap_ratio, dim = ape.shape
74	- else:
75	- assert compress_ratio is not None
76	- _, overlap_ratio, dim = score.shape
77	- has_index = mirror_cp_reverse_index is not None
78	- compress_token = compress_ratio * overlap_ratio
79	- _, state_token, _, _, _ = state_cache.shape
80	- total_num_compressed_tokens = compressed_batch_mapping.shape[0]
81	- kv_compressed = torch.zeros(total_num_compressed_tokens, dim, dtype=torch.bfloat16, device=score.device)
82	- num_seq = cu_seqlens_q.size(0) - 1
83	- for i in range(num_seq):
84	- qs = cu_seqlens_q[i]
85	- qe = cu_seqlens_q[i + 1]
86	- ks = cu_seqlens_k[i]
87	- ke = cu_seqlens_k[i + 1]
88	- state_idx = state_block_idx[i]
89	- qlen = qe - qs
90	- klen = ke - ks
91	- qstart = klen - qlen
92	- compressed_start = cu_compressed_seqlens_q[i]
93	- score_slice = torch.empty((compress_token, dim), device=score.device, dtype=score.dtype)
94	- latent_slice = torch.empty((compress_token, dim), device=score.device, dtype=score.dtype)
95	-
96	- for blk_id in range(qstart // compress_ratio, klen // compress_ratio):
97	- score_slice[:, :] = float('-inf')
98	- latent_slice[:, :] = 0
99	-
100	- qid_r = (blk_id + 1) * compress_ratio

Dtests/sampling/test_dflash.py+0-144

	@@ -1,144 +0,0 @@
1	-from collections.abc import Callable
2	-from typing import Any
3	-
4	-import pytest
5	-import torch
6	-
7	-from cann_ops_tilelang.config import get_device
8	-from cann_ops_tilelang.sampling.dflash import confidence_head_fused, dflash_input_store
9	-
10	-pytestmark = pytest.mark.npu
11	-
12	-
13	-def confidence_head_fused_ref(
14	- hidden: torch.Tensor,
15	- weight: torch.Tensor,
16	- markov_emb: torch.Tensor \| None = None,
17	- out: torch.Tensor \| None = None,
18	- temperature: float = 1.0,
19	-) -> torch.Tensor:
20	- features = torch.cat([hidden, markov_emb], dim=-1) if markov_emb is not None else hidden
21	- logits = (features.float() @ weight.float()).squeeze(-1)
22	- scaled = (logits / max(temperature, 1e-5)).clamp(-60.0, 60.0)
23	- result = 1.0 / (1.0 + (-scaled).exp())
24	- if out is None:
25	- return result
26	- out.copy_(result)
27	- return out
28	-
29	-
30	-def dflash_input_store_ref(hidden_states: torch.Tensor, out: torch.Tensor) -> None:
31	- hc_mult = hidden_states.shape[-2] if hidden_states.ndim > 2 else 1
32	- if hc_mult > 1:
33	- hidden_states = hidden_states.mean(dim=-2).to(hidden_states.dtype)
34	- out.copy_(hidden_states)
35	-
36	-
37	-@pytest.mark.parametrize('num_rows', [4097])
38	-@pytest.mark.parametrize('n_embed', [4096, 7168])
39	-@pytest.mark.parametrize('markov_rank', [None, 128, 256, 512])
40	-def test_confidence_head(num_rows: int, n_embed: int, markov_rank: int \| None) -> None:
41	- device = get_device()
42	- total_dim = n_embed + (markov_rank or 0)
43	- hidden = torch.randn(num_rows, n_embed, dtype=torch.bfloat16, device=device) / 10
44	- weight = torch.randn(total_dim, 1, dtype=torch.bfloat16, device=device) / 10
45	- markov_emb = torch.randn(num_rows, markov_rank, dtype=torch.bfloat16, device=device) if markov_rank else None
46	-
47	- def run_kernel() -> torch.Tensor:
48	- out = torch.empty(num_rows * 7, dtype=torch.float32, device=device)[::7]
49	- confidence_head_fused(hidden, weight, markov_emb, out=out, temperature=1.0)
50	- return out.contiguous()
51	-
52	- expected = confidence_head_fused_ref(hidden, weight, markov_emb)
53	- actual = run_kernel()
54	- max_diff = (actual - expected).abs().max().item()
55	- assert max_diff < 1e-4, f'{num_rows=}, {markov_rank=}, {max_diff=}'
56	-
57	- results = [run_kernel() for _ in range(10)]
58	- assert all(torch.equal(results[0], result) for result in results[1:])
59	-
60	-
61	-@pytest.mark.benchmark
62	-@pytest.mark.parametrize('num_rows', [1, 128, 4097])
63	-@pytest.mark.parametrize('n_embed', [4096, 7168])
64	-@pytest.mark.parametrize('markov_rank', [None, 256])
65	-def test_confidence_head_benchmark(
66	- num_rows: int,
67	- n_embed: int,
68	- markov_rank: int \| None,
69	- benchmark_timer: Callable[..., Any],
70	- benchmark_record: Callable[..., Any],
71	-) -> None:
72	- device = get_device()
73	- total_dim = n_embed + (markov_rank or 0)
74	- hidden = torch.randn(num_rows, n_embed, dtype=torch.bfloat16, device=device) / 10
75	- weight = torch.randn(total_dim, 1, dtype=torch.bfloat16, device=device) / 10
76	- markov_emb = torch.randn(num_rows, markov_rank, dtype=torch.bfloat16, device=device) if markov_rank else None
77	-
78	- time_us = benchmark_timer(lambda: confidence_head_fused(hidden, weight, markov_emb, temperature=1.0))
79	- tensors = [hidden, weight] + ([markov_emb] if markov_emb is not None else [])
80	- data_bytes = sum(tensor.numel() * tensor.element_size() for tensor in tensors)
81	- benchmark_record(
82	- kernel='dflash',
83	- operation='confidence_head',
84	- params={'num_rows': num_rows, 'markov_rank': markov_rank or 0, 'n_embed': n_embed},
85	- time_us=time_us,
86	- extras={'bandwidth_gbs': data_bytes / time_us / 1e3},
87	- )
88	-
89	-
90	-@pytest.mark.parametrize('num_rows', [3, 7, 233, 4097, 8198])
91	-@pytest.mark.parametrize('n_embed', [4096, 7168])
92	-@pytest.mark.parametrize('hc_mult', [1, 4])
93	-def test_dflash_input_store(num_rows: int, n_embed: int, hc_mult: int) -> None:
94	- device = get_device()
95	- if hc_mult > 1:
96	- hidden_states = torch.randn(num_rows, hc_mult, n_embed, dtype=torch.bfloat16, device=device)
97	- else:
98	- hidden_states = torch.randn(num_rows, n_embed, dtype=torch.bfloat16, device=device)
99	-
100	- num_layers = 3
101	- buffer_ref = torch.randn(num_rows, n_embed * num_layers, dtype=torch.bfloat16, device=device)
102	- buffer_actual = buffer_ref.clone()
103	- buffer_repeat = buffer_ref.clone()
104	- out_ref = buffer_ref[:, n_embed : 2 * n_embed]
105	- out_actual = buffer_actual[:, n_embed : 2 * n_embed]
106	-
107	- dflash_input_store_ref(hidden_states, out_ref)
108	- dflash_input_store(hidden_states, out_actual)
109	-
110	- max_diff = (buffer_actual.float() - buffer_ref.float()).abs().max().item()
111	- assert max_diff < 2e-3, f'{num_rows=}, {n_embed=}, {hc_mult=}, {max_diff=}'
112	-
113	- dflash_input_store(hidden_states, buffer_repeat[:, n_embed : 2 * n_embed])
114	- assert torch.equal(buffer_actual, buffer_repeat)
115	-
116	-
117	-@pytest.mark.benchmark
118	-@pytest.mark.parametrize('num_rows', [1, 128, 4097])
119	-@pytest.mark.parametrize('n_embed', [4096, 7168])
120	-@pytest.mark.parametrize('hc_mult', [1, 4])
121	-def test_dflash_input_store_benchmark(
122	- num_rows: int,
123	- n_embed: int,
124	- hc_mult: int,
125	- benchmark_timer: Callable[..., Any],
126	- benchmark_record: Callable[..., Any],
127	-) -> None:
128	- device = get_device()
129	- if hc_mult > 1:
130	- hidden_states = torch.randn(num_rows, hc_mult, n_embed, dtype=torch.bfloat16, device=device)
131	- else:
132	- hidden_states = torch.randn(num_rows, n_embed, dtype=torch.bfloat16, device=device)
133	- buffer_actual = torch.randn(num_rows, n_embed * 3, dtype=torch.bfloat16, device=device)
134	- out = buffer_actual[:, n_embed : 2 * n_embed]
135	-
136	- time_us = benchmark_timer(lambda: dflash_input_store(hidden_states, out))
137	- data_bytes = sum(tensor.numel() * tensor.element_size() for tensor in (hidden_states, out))
138	- benchmark_record(
139	- kernel='dflash',
140	- operation='dflash_input_store',
141	- params={'hc_mult': hc_mult, 'num_rows': num_rows, 'n_embed': n_embed},
142	- time_us=time_us,
143	- extras={'bandwidth_gbs': data_bytes / time_us / 1e3},
144	- )

Mtests/sampling/test_sampler.py+1-964