[Fix] Fix static check errors detected by SPACES-pytorch-AtomGit

M.lintrunner.toml+28-38

	@@ -428,44 +428,34 @@ command = [
428	]	428	]
429	is_formatter = true	429	is_formatter = true
430		430
431	-# [[linter]]	431	+[[linter]]
432	-# code = 'SPACES'	432	+code = 'SPACES'
433	-# include_patterns = ['**']	433	+include_patterns = ['**']
434	-# exclude_patterns = [	434	+exclude_patterns = [
435	-# 'test_upstream/**',	435	+ '/contrib/',
436	-# '/contrib/',	436	+ '*/.diff',
437	-# '*/.diff',	437	+ '*/.patch',
438	-# '*/.patch',	438	+ 'third_party/**',
439	-# 'third_party/**',	439	+ 'aten/src/ATen/native/vulkan/api/vk_mem_alloc.h',
440	-# 'aten/src/ATen/native/vulkan/api/vk_mem_alloc.h',	440	+ 'fb/**',
441	-# 'fb/**',	441	+ '/fb/',
442	-# '/fb/',	442	+ 'test/cpp/jit/upgrader_models/*.ptl',
443	-# 'test/cpp/jit/upgrader_models/*.ptl',	443	+ 'test/cpp/jit/upgrader_models/*.ptl.ff',
444	-# 'test/cpp/jit/upgrader_models/*.ptl.ff',	444	+ '*/.md',
445	-# # NPUGraph logs files	445	+]
446	-# 'torch_npu/_logging/_internal.py',	446	+command = [
447	-# 'torch_npu/csrc/core/npu/NPUGraph.cpp',	447	+ 'python3',
448	-# 'torch_npu/csrc/core/npu/NPUGraph.h',	448	+ 'tools/linter/adapters/grep_linter.py',
449	-# 'torch_npu/csrc/npu/Graph.cpp',	449	+ '--pattern=[[:blank:]]$',
450	-# 'torch_npu/csrc/core/npu/NPUCachingAllocator.cpp',	450	+ '--linter-name=SPACES',
451	-# 'torch_npu/csrc/core/npu/NPUWorkspaceAllocator.cpp',	451	+ '--error-name=trailing spaces',
452	-# 'torch_npu/npu/_graph_tree.py',	452	+ '--replace-pattern=s/[[:blank:]]+$//',
453	-# 'torch_npu/npu/graphs.py',	453	+ """--error-description=\
454	-# 'torch_npu/utils/_graph_tree.py',	454	+ This line has trailing spaces; please remove them.\
455	-# ]	455	+ """,
456	-# command = [	456	+ '--',
457	-# 'python3',	457	+ '@{{PATHSFILE}}'
458	-# 'tools/linter/adapters/grep_linter.py',	458	+]
459	-# '--pattern=[[:blank:]]$',
460	-# '--linter-name=SPACES',
461	-# '--error-name=trailing spaces',
462	-# '--replace-pattern=s/[[:blank:]]+$//',
463	-# """--error-description=\
464	-# This line has trailing spaces; please remove them.\
465	-# """,
466	-# '--',
467	-# '@{{PATHSFILE}}'
468	-# ]
469		459
470	# [[linter]]	460	# [[linter]]
471	# code = 'TABS'	461	# code = 'TABS'

Mbenchmarks/inference/DCNv2/dcnv2.py+4-4

	@@ -144,7 +144,7 @@ def eval_op_prof(model: nn.Module, device, mod, args):
144		144
145	model.eval()	145	model.eval()
146	with torch.no_grad():	146	with torch.no_grad():
147	- res = model(x)	147	+ res = model(x)
148	torch.save(res, pt_path)	148	torch.save(res, pt_path)
149	device_synchronize()	149	device_synchronize()
150		150
	@@ -163,13 +163,13 @@ def eval_op_prof(model: nn.Module, device, mod, args):
163	end_time = time.time()	163	end_time = time.time()
164	step_time_ms = (end_time - start_time) * 1000	164	step_time_ms = (end_time - start_time) * 1000
165	print(f"[{mod}]: Step {i}: {step_time_ms:.4f} ms")	165	print(f"[{mod}]: Step {i}: {step_time_ms:.4f} ms")
166	-	166	+
167	if i >= 10:	167	if i >= 10:
168	execution_times.append(step_time_ms)	168	execution_times.append(step_time_ms)
169		169
170	if args.enable_profiler:	170	if args.enable_profiler:
171	prof.step()	171	prof.step()
172	-	172	+
173	if args.enable_profiler:	173	if args.enable_profiler:
174	prof.stop()	174	prof.stop()
175	if execution_times:	175	if execution_times:
	@@ -177,7 +177,7 @@ def eval_op_prof(model: nn.Module, device, mod, args):
177	print(f"[{mod}]:: Avg over {len(execution_times)} steps: {avg_ms:.4f} ms")	177	print(f"[{mod}]:: Avg over {len(execution_times)} steps: {avg_ms:.4f} ms")
178		178
179	parser = argparse.ArgumentParser(description=MODEL_NAME + " infernece")	179	parser = argparse.ArgumentParser(description=MODEL_NAME + " infernece")
180	-parser.add_argument("--max_steps", type=int, default=200,	180	+parser.add_argument("--max_steps", type=int, default=200,
181	help="Total training steps")	181	help="Total training steps")
182	parser.add_argument("--enable_compile", action="store_true",	182	parser.add_argument("--enable_compile", action="store_true",
183	help="Enable torch.compile and Inductor backend")	183	help="Enable torch.compile and Inductor backend")

Mbenchmarks/inference/DIN/din.py+4-4

	@@ -117,7 +117,7 @@ def eval_op_prof(model: nn.Module, device, mod, args):
117		117
118	model.eval()	118	model.eval()
119	with torch.no_grad():	119	with torch.no_grad():
120	- res = model(x)	120	+ res = model(x)
121	torch.save(res, pt_path)	121	torch.save(res, pt_path)
122	device_synchronize()	122	device_synchronize()
123		123
	@@ -136,13 +136,13 @@ def eval_op_prof(model: nn.Module, device, mod, args):
136	end_time = time.time()	136	end_time = time.time()
137	step_time_ms = (end_time - start_time) * 1000	137	step_time_ms = (end_time - start_time) * 1000
138	print(f"[{mod}]: Step {i}: {step_time_ms:.4f} ms")	138	print(f"[{mod}]: Step {i}: {step_time_ms:.4f} ms")
139	-	139	+
140	if i >= 10:	140	if i >= 10:
141	execution_times.append(step_time_ms)	141	execution_times.append(step_time_ms)
142		142
143	if args.enable_profiler:	143	if args.enable_profiler:
144	prof.step()	144	prof.step()
145	-	145	+
146	if args.enable_profiler:	146	if args.enable_profiler:
147	prof.stop()	147	prof.stop()
148	if execution_times:	148	if execution_times:
	@@ -150,7 +150,7 @@ def eval_op_prof(model: nn.Module, device, mod, args):
150	print(f"[{mod}]:: Avg over {len(execution_times)} steps: {avg_ms:.4f} ms")	150	print(f"[{mod}]:: Avg over {len(execution_times)} steps: {avg_ms:.4f} ms")
151		151
152	parser = argparse.ArgumentParser(description=MODEL_NAME + " infernece")	152	parser = argparse.ArgumentParser(description=MODEL_NAME + " infernece")
153	-parser.add_argument("--max_steps", type=int, default=200,	153	+parser.add_argument("--max_steps", type=int, default=200,
154	help="Total training steps")	154	help="Total training steps")
155	parser.add_argument("--enable_compile", action="store_true",	155	parser.add_argument("--enable_compile", action="store_true",
156	help="Enable torch.compile and Inductor backend")	156	help="Enable torch.compile and Inductor backend")

Mbenchmarks/inference/ETA/eta.py+6-6

	@@ -22,9 +22,9 @@ torch.manual_seed(2024)
22	random.seed(2024)	22	random.seed(2024)
23		23
24	MODEL_NAME = "ETA"	24	MODEL_NAME = "ETA"
25	-STD_DEV = (2 / 512) ** 0.5	25	+STD_DEV = (2 / 512) ** 0.5
26	-EMBEDDING_FEATURE_NUM = 27	26	+EMBEDDING_FEATURE_NUM = 27
27	-TARGET_FIELDS = ["206", "207", "216", "210"]	27	+TARGET_FIELDS = ["206", "207", "216", "210"]
28		28
29	def detect_device_type():	29	def detect_device_type():
30	try:	30	try:
	@@ -103,7 +103,7 @@ def parse_arguments():
103	parser.add_argument("--clear_existing_model", action="store_true", help="")	103	parser.add_argument("--clear_existing_model", action="store_true", help="")
104	parser.add_argument("--task_type", type=str, default="train",	104	parser.add_argument("--task_type", type=str, default="train",
105	choices=["train", "eval", "predict"], help="task type")	105	choices=["train", "eval", "predict"], help="task type")
106	- parser.add_argument("--max_steps", type=int, default=100,	106	+ parser.add_argument("--max_steps", type=int, default=100,
107	help="Total training steps (优先级高于num_epochs，-1表示使用num_epochs控制)")	107	help="Total training steps (优先级高于num_epochs，-1表示使用num_epochs控制)")
108	parser.add_argument('--epoch_num', type=int, default=1, help="Number of epochs")	108	parser.add_argument('--epoch_num', type=int, default=1, help="Number of epochs")
109	parser.add_argument('--train_batch_num', type=int, default=2000, help="Number of train batchs")	109	parser.add_argument('--train_batch_num', type=int, default=2000, help="Number of train batchs")
	@@ -455,7 +455,7 @@ def evaluate_op(model: ETA, dataloader, device, args):
455	model = torch.compile(model, dynamic=False)	455	model = torch.compile(model, dynamic=False)
456	model.eval()	456	model.eval()
457	with torch.no_grad():	457	with torch.no_grad():
458	- model(features)	458	+ model(features)
459	device_synchronize()	459	device_synchronize()
460		460
461	prof = None	461	prof = None
	@@ -480,7 +480,7 @@ def evaluate_op(model: ETA, dataloader, device, args):
480	if args.enable_profiler:	480	if args.enable_profiler:
481	prof.step()	481	prof.step()
482	if args.enable_profiler:	482	if args.enable_profiler:
483	- prof.stop()	483	+ prof.stop()
484	if exec_times:	484	if exec_times:
485	avg_time = sum(exec_times) / len(exec_times)	485	avg_time = sum(exec_times) / len(exec_times)
486	print("Step time consumption statistics (excluding the first 10 steps)")	486	print("Step time consumption statistics (excluding the first 10 steps)")

Mbenchmarks/inference/MMOE/mmoe.py+4-4

	@@ -23,8 +23,8 @@ torch.manual_seed(2024)
23	random.seed(2024)	23	random.seed(2024)
24		24
25	MODEL_NAME = "MMOE"	25	MODEL_NAME = "MMOE"
26	-EMBEDDING_FEATURE_NUM = 23	26	+EMBEDDING_FEATURE_NUM = 23
27	-STD_DEV = (2 / 512) ** 0.5	27	+STD_DEV = (2 / 512) ** 0.5
28		28
29	def detect_device_type():	29	def detect_device_type():
30	try:	30	try:
	@@ -529,7 +529,7 @@ def evaluate_op(model: TorchMmoeModel, te_files, device, args):
529		529
530	model.eval()	530	model.eval()
531	with torch.no_grad():	531	with torch.no_grad():
532	- model(input_sample)	532	+ model(input_sample)
533	device_synchronize()	533	device_synchronize()
534		534
535	prof = None	535	prof = None
	@@ -554,7 +554,7 @@ def evaluate_op(model: TorchMmoeModel, te_files, device, args):
554	if args.enable_profiler:	554	if args.enable_profiler:
555	prof.step()	555	prof.step()
556	if args.enable_profiler:	556	if args.enable_profiler:
557	- prof.stop()	557	+ prof.stop()
558	if exec_times:	558	if exec_times:
559	avg_time = sum(exec_times) / len(exec_times)	559	avg_time = sum(exec_times) / len(exec_times)
560	print("Step time consumption statistics (excluding the first 10 steps)")	560	print("Step time consumption statistics (excluding the first 10 steps)")

Mbenchmarks/llm/baichuan2-7B-Chat/train_baichuan2_7B.py+1-1

	@@ -188,7 +188,7 @@ class Baichuan2Trainer:
188	profiling_save_path = self.args.profiler_save_path + '/' + MODEL_NAME + '/' + mode	188	profiling_save_path = self.args.profiler_save_path + '/' + MODEL_NAME + '/' + mode
189	prof = get_profile(self.args.profiler_start_step, self.args.profiler_end_step, profiling_save_path)	189	prof = get_profile(self.args.profiler_start_step, self.args.profiler_end_step, profiling_save_path)
190		190
191	- timing_callback = TimingCallback(prof, mode)	191	+ timing_callback = TimingCallback(prof, mode)
192		192
193	trainer = BaichuanTrainer(	193	trainer = BaichuanTrainer(
194	model=self.model,	194	model=self.model,

Mbenchmarks/llm/gpt-oss-20B/train_gpt_oss_20B.py+56-56

	@@ -22,10 +22,10 @@ from typing import Dict, List, Optional
22	import argparse	22	import argparse
23	import logging	23	import logging
24		24
25	-sys.path.append(str(Path(__file__).parent.parent))	25	+sys.path.append(str(Path(__file__).parent.parent))
26	from utils.utils import (	26	from utils.utils import (
27	- TimingCallback,	27	+ TimingCallback,
28	- get_profile,	28	+ get_profile,
29	detect_device_type	29	detect_device_type
30	)	30	)
31		31
	@@ -38,10 +38,10 @@ class GPT_OSS_20BTrainer:
38	def __init__(self, args):	38	def __init__(self, args):
39	self.args = args	39	self.args = args
40	self.setup_training()	40	self.setup_training()
41	-	41	+
42	def setup_training(self):	42	def setup_training(self):
43	logger.info(f"Loading model from {self.args.model_path}")	43	logger.info(f"Loading model from {self.args.model_path}")
44	-	44	+
45	self.model = AutoModelForCausalLM.from_pretrained(	45	self.model = AutoModelForCausalLM.from_pretrained(
46	self.args.model_path,	46	self.args.model_path,
47	torch_dtype=torch.bfloat16 if self.args.use_bf16 else torch.float16,	47	torch_dtype=torch.bfloat16 if self.args.use_bf16 else torch.float16,
	@@ -50,30 +50,30 @@ class GPT_OSS_20BTrainer:
50		50
51	logger.info(f"Moving model to {self.args.device_type}...")	51	logger.info(f"Moving model to {self.args.device_type}...")
52	self.model = self.model.to(self.args.device_type)	52	self.model = self.model.to(self.args.device_type)
53	-	53	+
54	self.tokenizer = AutoTokenizer.from_pretrained(	54	self.tokenizer = AutoTokenizer.from_pretrained(
55	self.args.model_path,	55	self.args.model_path,
56	trust_remote_code=True	56	trust_remote_code=True
57	)	57	)
58	-	58	+
59	if self.tokenizer.pad_token is None:	59	if self.tokenizer.pad_token is None:
60	self.tokenizer.pad_token = self.tokenizer.eos_token	60	self.tokenizer.pad_token = self.tokenizer.eos_token
61	self.model.config.pad_token_id = self.tokenizer.eos_token_id	61	self.model.config.pad_token_id = self.tokenizer.eos_token_id
62	-	62	+
63	if self.args.gradient_checkpointing:	63	if self.args.gradient_checkpointing:
64	self.model.gradient_checkpointing_enable()	64	self.model.gradient_checkpointing_enable()
65	self.model.config.use_cache = False	65	self.model.config.use_cache = False
66		66
67	-	67	+
68	def apply_lora(self):	68	def apply_lora(self):
69	if not self.args.use_lora:	69	if not self.args.use_lora:
70	return	70	return
71	-	71	+
72	logger.info("Applying LoRA configuration...")	72	logger.info("Applying LoRA configuration...")
73	-	73	+
74	if self.args.use_4bit:	74	if self.args.use_4bit:
75	self.model = prepare_model_for_kbit_training(self.model)	75	self.model = prepare_model_for_kbit_training(self.model)
76	-	76	+
77	lora_config = LoraConfig(	77	lora_config = LoraConfig(
78	task_type=TaskType.CAUSAL_LM,	78	task_type=TaskType.CAUSAL_LM,
79	r=self.args.lora_r,	79	r=self.args.lora_r,
	@@ -82,38 +82,38 @@ class GPT_OSS_20BTrainer:
82	target_modules=self.get_lora_target_modules(),	82	target_modules=self.get_lora_target_modules(),
83	bias="none",	83	bias="none",
84	)	84	)
85	-	85	+
86	self.model = get_peft_model(self.model, lora_config)	86	self.model = get_peft_model(self.model, lora_config)
87	self.model.print_trainable_parameters()	87	self.model.print_trainable_parameters()
88	-	88	+
89	if self.args.enable_compile:	89	if self.args.enable_compile:
90	self.model = torch.compile(self.model, dynamic=False)	90	self.model = torch.compile(self.model, dynamic=False)
91	-	91	+
92	def get_lora_target_modules(self):	92	def get_lora_target_modules(self):
93	target_modules = [	93	target_modules = [
94	"q_proj", "k_proj", "v_proj", "o_proj",	94	"q_proj", "k_proj", "v_proj", "o_proj",
95	"gate_proj", "up_proj", "down_proj",	95	"gate_proj", "up_proj", "down_proj",
96	]	96	]
97	-	97	+
98	model_modules = set([name for name, _ in self.model.named_modules()])	98	model_modules = set([name for name, _ in self.model.named_modules()])
99	available_modules = [m for m in target_modules if any(m in name for name in model_modules)]	99	available_modules = [m for m in target_modules if any(m in name for name in model_modules)]
100	-	100	+
101	if not available_modules:	101	if not available_modules:
102	available_modules = ["qkv_proj", "dense", "fc1", "fc2"]	102	available_modules = ["qkv_proj", "dense", "fc1", "fc2"]
103	-	103	+
104	logger.info(f"Using LoRA target modules: {available_modules}")	104	logger.info(f"Using LoRA target modules: {available_modules}")
105	return available_modules	105	return available_modules
106	-	106	+
107	def load_and_preprocess_data(self) -> Dataset:	107	def load_and_preprocess_data(self) -> Dataset:
108	logger.info(f"Loading dataset from {self.args.data_path}")	108	logger.info(f"Loading dataset from {self.args.data_path}")
109	-	109	+
110	if self.args.data_path.endswith('.json') or self.args.data_path.endswith('.jsonl'):	110	if self.args.data_path.endswith('.json') or self.args.data_path.endswith('.jsonl'):
111	with open(self.args.data_path, 'r', encoding='utf-8') as f:	111	with open(self.args.data_path, 'r', encoding='utf-8') as f:
112	if self.args.data_path.endswith('.jsonl'):	112	if self.args.data_path.endswith('.jsonl'):
113	data = [json.loads(line) for line in f]	113	data = [json.loads(line) for line in f]
114	else:	114	else:
115	data = json.load(f)	115	data = json.load(f)
116	-	116	+
117	formatted_data = []	117	formatted_data = []
118	for item in data:	118	for item in data:
119	if "conversations" in item:	119	if "conversations" in item:
	@@ -128,11 +128,11 @@ class GPT_OSS_20BTrainer:
128	)	128	)
129	else:	129	else:
130	text = item.get("text", "")	130	text = item.get("text", "")
131	-	131	+
132	formatted_data.append({"text": text})	132	formatted_data.append({"text": text})
133	-	133	+
134	dataset = Dataset.from_list(formatted_data)	134	dataset = Dataset.from_list(formatted_data)
135	-	135	+
136	else:	136	else:
137	try:	137	try:
138	dataset = load_dataset(	138	dataset = load_dataset(
	@@ -145,7 +145,7 @@ class GPT_OSS_20BTrainer:
145	data_files=self.args.data_path,	145	data_files=self.args.data_path,
146	split="train"	146	split="train"
147	)	147	)
148	-	148	+
149	def preprocess_function(examples, tokenizer=self.tokenizer, max_length=self.args.max_length):	149	def preprocess_function(examples, tokenizer=self.tokenizer, max_length=self.args.max_length):
150	tokenized = tokenizer(	150	tokenized = tokenizer(
151	examples["text"],	151	examples["text"],
	@@ -155,45 +155,45 @@ class GPT_OSS_20BTrainer:
155	return_tensors=None,	155	return_tensors=None,
156	return_attention_mask=True,	156	return_attention_mask=True,
157	)	157	)
158	-	158	+
159	import copy	159	import copy
160	tokenized["labels"] = copy.deepcopy(tokenized["input_ids"])	160	tokenized["labels"] = copy.deepcopy(tokenized["input_ids"])
161	-	161	+
162	return tokenized	162	return tokenized
163	-	163	+
164	num_proc = self.args.num_proc if self.args.num_proc > 0 else None	164	num_proc = self.args.num_proc if self.args.num_proc > 0 else None
165	processed_dataset = dataset.map(	165	processed_dataset = dataset.map(
166	preprocess_function,	166	preprocess_function,
167	batched=True,	167	batched=True,
168	remove_columns=dataset.column_names,	168	remove_columns=dataset.column_names,
169	- num_proc=num_proc,	169	+ num_proc=num_proc,
170	load_from_cache_file=not self.args.overwrite_cache	170	load_from_cache_file=not self.args.overwrite_cache
171	)	171	)
172	-	172	+
173	logger.info(f"Dataset size: {len(processed_dataset)}")	173	logger.info(f"Dataset size: {len(processed_dataset)}")
174	return processed_dataset	174	return processed_dataset
175	-	175	+
176	def format_conversation(self, conversations: List[Dict]) -> str:	176	def format_conversation(self, conversations: List[Dict]) -> str:
177	formatted_text = ""	177	formatted_text = ""
178	-	178	+
179	for turn in conversations:	179	for turn in conversations:
180	role = turn.get("from", "").lower()	180	role = turn.get("from", "").lower()
181	content = turn.get("value", "")	181	content = turn.get("value", "")
182	-	182	+
183	if role == "human" or role == "user":	183	if role == "human" or role == "user":
184	formatted_text += f"<\|im_start\|>user\n{content}<\|im_end\|>\n"	184	formatted_text += f"<\|im_start\|>user\n{content}<\|im_end\|>\n"
185	elif role == "gpt" or role == "assistant":	185	elif role == "gpt" or role == "assistant":
186	formatted_text += f"<\|im_start\|>assistant\n{content}<\|im_end\|>\n"	186	formatted_text += f"<\|im_start\|>assistant\n{content}<\|im_end\|>\n"
187	elif role == "system":	187	elif role == "system":
188	formatted_text += f"<\|im_start\|>system\n{content}<\|im_end\|>\n"	188	formatted_text += f"<\|im_start\|>system\n{content}<\|im_end\|>\n"
189	-	189	+
190	return formatted_text	190	return formatted_text
191	-	191	+
192	def create_trainer(self, train_dataset, eval_dataset=None):	192	def create_trainer(self, train_dataset, eval_dataset=None):
193	has_eval = eval_dataset is not None	193	has_eval = eval_dataset is not None
194	eval_strategy = "steps" if has_eval else "no"	194	eval_strategy = "steps" if has_eval else "no"
195	- save_strategy = "steps" if has_eval else "steps"	195	+ save_strategy = "steps" if has_eval else "steps"
196	-	196	+
197	training_args = TrainingArguments(	197	training_args = TrainingArguments(
198	output_dir=self.args.output_dir,	198	output_dir=self.args.output_dir,
199	overwrite_output_dir=True,	199	overwrite_output_dir=True,
	@@ -207,10 +207,10 @@ class GPT_OSS_20BTrainer:
207	logging_steps=self.args.logging_steps,	207	logging_steps=self.args.logging_steps,
208	save_steps=self.args.save_steps,	208	save_steps=self.args.save_steps,
209	save_total_limit=self.args.save_total_limit,	209	save_total_limit=self.args.save_total_limit,
210	- eval_strategy=eval_strategy,	210	+ eval_strategy=eval_strategy,
211	eval_steps=self.args.eval_steps if has_eval else None,	211	eval_steps=self.args.eval_steps if has_eval else None,
212	- save_strategy=save_strategy,	212	+ save_strategy=save_strategy,
213	- load_best_model_at_end=has_eval,	213	+ load_best_model_at_end=has_eval,
214	metric_for_best_model="eval_loss",	214	metric_for_best_model="eval_loss",
215	greater_is_better=False,	215	greater_is_better=False,
216	learning_rate=self.args.learning_rate,	216	learning_rate=self.args.learning_rate,
	@@ -222,7 +222,7 @@ class GPT_OSS_20BTrainer:
222	ddp_find_unused_parameters=False if torch.cuda.device_count() > 1 else None,	222	ddp_find_unused_parameters=False if torch.cuda.device_count() > 1 else None,
223	remove_unused_columns=False,	223	remove_unused_columns=False,
224	)	224	)
225	-	225	+
226	data_collator = DataCollatorForLanguageModeling(	226	data_collator = DataCollatorForLanguageModeling(
227	tokenizer=self.tokenizer,	227	tokenizer=self.tokenizer,
228	mlm=False,	228	mlm=False,
	@@ -235,8 +235,8 @@ class GPT_OSS_20BTrainer:
235	profiling_save_path = self.args.profiler_save_path + '/' + model_name + '/' + mod	235	profiling_save_path = self.args.profiler_save_path + '/' + model_name + '/' + mod
236	prof = get_profile(self.args.profiler_start_step, self.args.profiler_end_step, profiling_save_path)	236	prof = get_profile(self.args.profiler_start_step, self.args.profiler_end_step, profiling_save_path)
237		237
238	- timing_callback = TimingCallback(prof, mod)	238	+ timing_callback = TimingCallback(prof, mod)
239	-	239	+
240	trainer = Trainer(	240	trainer = Trainer(
241	model=self.model,	241	model=self.model,
242	args=training_args,	242	args=training_args,
	@@ -246,17 +246,17 @@ class GPT_OSS_20BTrainer:
246	data_collator=data_collator,	246	data_collator=data_collator,
247	callbacks=[timing_callback],	247	callbacks=[timing_callback],
248	)	248	)
249	-	249	+
250	return trainer	250	return trainer
251	-	251	+
252	def train(self):	252	def train(self):
253	logger.info("Starting training...")	253	logger.info("Starting training...")
254	-	254	+
255	if self.args.use_lora:	255	if self.args.use_lora:
256	self.apply_lora()	256	self.apply_lora()
257	-	257	+
258	dataset = self.load_and_preprocess_data()	258	dataset = self.load_and_preprocess_data()
259	-	259	+
260	if self.args.validation_split > 0:	260	if self.args.validation_split > 0:
261	split_dataset = dataset.train_test_split(	261	split_dataset = dataset.train_test_split(
262	test_size=self.args.validation_split,	262	test_size=self.args.validation_split,
	@@ -267,11 +267,11 @@ class GPT_OSS_20BTrainer:
267	else:	267	else:
268	train_dataset = dataset	268	train_dataset = dataset
269	eval_dataset = None	269	eval_dataset = None
270	-	270	+
271	trainer = self.create_trainer(train_dataset, eval_dataset)	271	trainer = self.create_trainer(train_dataset, eval_dataset)
272	-	272	+
273	train_result = trainer.train()	273	train_result = trainer.train()
274	-	274	+
275	if self.args.enable_compile:	275	if self.args.enable_compile:
276	headers, values = torch._dynamo.utils.compile_times("csv")	276	headers, values = torch._dynamo.utils.compile_times("csv")
277	for header, value in zip(headers, values):	277	for header, value in zip(headers, values):
	@@ -282,14 +282,14 @@ class GPT_OSS_20BTrainer:
282		282
283	trainer.save_model()	283	trainer.save_model()
284	self.tokenizer.save_pretrained(self.args.output_dir)	284	self.tokenizer.save_pretrained(self.args.output_dir)
285	-	285	+
286	metrics = train_result.metrics	286	metrics = train_result.metrics
287	trainer.log_metrics("train", metrics)	287	trainer.log_metrics("train", metrics)
288	trainer.save_metrics("train", metrics)	288	trainer.save_metrics("train", metrics)
289	trainer.save_state()	289	trainer.save_state()
290	-	290	+
291	logger.info(f"Training completed! Model saved to {self.args.output_dir}")	291	logger.info(f"Training completed! Model saved to {self.args.output_dir}")
292	-	292	+
293	return metrics	293	return metrics
294		294
295	def main():	295	def main():
	@@ -310,7 +310,7 @@ def main():
310	help="Learning rate")	310	help="Learning rate")
311	parser.add_argument("--warmup_steps", type=int, default=5,	311	parser.add_argument("--warmup_steps", type=int, default=5,
312	help="Warmup steps")	312	help="Warmup steps")
313	- parser.add_argument("--max_steps", type=int, default=1,	313	+ parser.add_argument("--max_steps", type=int, default=1,
314	help="Total training steps")	314	help="Total training steps")
315	parser.add_argument("--weight_decay", type=float, default=0.01,	315	parser.add_argument("--weight_decay", type=float, default=0.01,
316	help="Weight decay")	316	help="Weight decay")
	@@ -379,7 +379,7 @@ def main():
379		379
380	trainer = GPT_OSS_20BTrainer(args)	380	trainer = GPT_OSS_20BTrainer(args)
381	metrics = trainer.train()	381	metrics = trainer.train()
382	-	382	+
383	print("\n" + "="*50)	383	print("\n" + "="*50)
384	print("Training completed successfully!")	384	print("Training completed successfully!")
385	print(f"Model saved to: {args.output_dir}")	385	print(f"Model saved to: {args.output_dir}")

Mbenchmarks/llm/mamba-codestral-7B/train_mamba2_7B.py+13-13

	@@ -105,7 +105,7 @@ def main():
105	)	105	)
106	if tokenizer.pad_token is None:	106	if tokenizer.pad_token is None:
107	tokenizer.pad_token = tokenizer.eos_token	107	tokenizer.pad_token = tokenizer.eos_token
108	-	108	+
109	print(f"Loading base model from: {args.model_path} (no quantization)")	109	print(f"Loading base model from: {args.model_path} (no quantization)")
110	base_model = AutoModelForCausalLM.from_pretrained(	110	base_model = AutoModelForCausalLM.from_pretrained(
111	args.model_path,	111	args.model_path,
	@@ -113,9 +113,9 @@ def main():
113	trust_remote_code=True,	113	trust_remote_code=True,
114	use_cache=False, # Disable cache to save memory	114	use_cache=False, # Disable cache to save memory
115	)	115	)
116	-	116	+
117	base_model.gradient_checkpointing_enable()	117	base_model.gradient_checkpointing_enable()
118	-	118	+
119	print(f"Configuring LoRA, target modules: {TARGET_MODULES}")	119	print(f"Configuring LoRA, target modules: {TARGET_MODULES}")
120	peft_config = LoraConfig(	120	peft_config = LoraConfig(
121	task_type=TaskType.CAUSAL_LM,	121	task_type=TaskType.CAUSAL_LM,
	@@ -125,13 +125,13 @@ def main():
125	target_modules=TARGET_MODULES,	125	target_modules=TARGET_MODULES,
126	bias="none",	126	bias="none",
127	)	127	)
128	-	128	+
129	print("Attaching LoRA adapters to model...")	129	print("Attaching LoRA adapters to model...")
130	peft_model = get_peft_model(base_model, peft_config)	130	peft_model = get_peft_model(base_model, peft_config)
131	peft_model.print_trainable_parameters() # Print number of trainable parameters	131	peft_model.print_trainable_parameters() # Print number of trainable parameters
132	-	132	+
133	peft_model.to(device)	133	peft_model.to(device)
134	-	134	+
135	print("Preparing dataset...")	135	print("Preparing dataset...")
136	from datasets import load_dataset	136	from datasets import load_dataset
137	dataset = load_dataset(	137	dataset = load_dataset(
	@@ -139,7 +139,7 @@ def main():
139	data_files=args.data_file,	139	data_files=args.data_file,
140	split="train"	140	split="train"
141	)	141	)
142	-	142	+
143	def tokenize_function(examples):	143	def tokenize_function(examples):
144	return tokenizer(	144	return tokenizer(
145	examples["text"],	145	examples["text"],
	@@ -148,7 +148,7 @@ def main():
148	max_length=args.max_seq_length,	148	max_length=args.max_seq_length,
149	return_tensors=None	149	return_tensors=None
150	)	150	)
151	-	151	+
152	train_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])	152	train_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
153	train_dataset = train_dataset.map(lambda x: {"labels": x["input_ids"]}, batched=True)	153	train_dataset = train_dataset.map(lambda x: {"labels": x["input_ids"]}, batched=True)
154		154
	@@ -157,7 +157,7 @@ def main():
157	peft_model,	157	peft_model,
158	dynamic=False	158	dynamic=False
159	)	159	)
160	-	160	+
161	training_args = TrainingArguments(	161	training_args = TrainingArguments(
162	output_dir=args.output_dir,	162	output_dir=args.output_dir,
163	num_train_epochs=args.epochs,	163	num_train_epochs=args.epochs,
	@@ -180,7 +180,7 @@ def main():
180	max_grad_norm=0.3,	180	max_grad_norm=0.3,
181	logging_dir=f"{args.output_dir}/logs",	181	logging_dir=f"{args.output_dir}/logs",
182	)	182	)
183	-	183	+
184	mode = "compile" if args.enable_compile else "eager"	184	mode = "compile" if args.enable_compile else "eager"
185	prof=None	185	prof=None
186		186
	@@ -188,7 +188,7 @@ def main():
188	profiling_save_path = args.profiler_save_path + '/' + MODEL_NAME + '/' + mode	188	profiling_save_path = args.profiler_save_path + '/' + MODEL_NAME + '/' + mode
189	prof = get_profile(args.profiler_start_step, args.profiler_end_step, profiling_save_path)	189	prof = get_profile(args.profiler_start_step, args.profiler_end_step, profiling_save_path)
190		190
191	- timing_callback = TimingCallback(prof, mode)	191	+ timing_callback = TimingCallback(prof, mode)
192		192
193	print("Initializing Trainer...")	193	print("Initializing Trainer...")
194	trainer = Trainer(	194	trainer = Trainer(
	@@ -199,7 +199,7 @@ def main():
199	data_collator=default_data_collator,	199	data_collator=default_data_collator,
200	callbacks=[timing_callback],	200	callbacks=[timing_callback],
201	)	201	)
202	-	202	+
203	print("Starting LoRA training...")	203	print("Starting LoRA training...")
204	trainer.train()	204	trainer.train()
205		205
	@@ -210,7 +210,7 @@ def main():
210	numbers = [float(num.strip()) for num in value.split(',') if num.strip()]	210	numbers = [float(num.strip()) for num in value.split(',') if num.strip()]
211	op_compile_time = sum(numbers)	211	op_compile_time = sum(numbers)
212	print(f"op_compile_time:{op_compile_time * 1e3} ms", )	212	print(f"op_compile_time:{op_compile_time * 1e3} ms", )
213	-	213	+
214	print(f"Saving LoRA weights to {args.output_dir}")	214	print(f"Saving LoRA weights to {args.output_dir}")
215	peft_model.save_pretrained(args.output_dir)	215	peft_model.save_pretrained(args.output_dir)
216	tokenizer.save_pretrained(args.output_dir)	216	tokenizer.save_pretrained(args.output_dir)

Mbenchmarks/llm/sdxl-base-1.0/train_sdxl_base.py+17-17

	@@ -12,9 +12,9 @@ import numpy as np
12	import argparse	12	import argparse
13	import logging	13	import logging
14	from transformers import Trainer, TrainingArguments	14	from transformers import Trainer, TrainingArguments
15	-sys.path.append(str(Path(__file__).parent.parent))	15	+sys.path.append(str(Path(__file__).parent.parent))
16	from utils.utils import (	16	from utils.utils import (
17	- TimingCallback,	17	+ TimingCallback,
18	get_profile,	18	get_profile,
19	detect_device_type	19	detect_device_type
20	)	20	)
	@@ -37,7 +37,7 @@ activations.GEGLU.forward = patched_geglu_forward
37		37
38	def set_seed(seed=42):	38	def set_seed(seed=42):
39	torch.manual_seed(seed)	39	torch.manual_seed(seed)
40	-	40	+
41	def custom_collate_fn(examples):	41	def custom_collate_fn(examples):
42	pixel_values = torch.stack([example["pixel_values"] for example in examples])	42	pixel_values = torch.stack([example["pixel_values"] for example in examples])
43	texts = [example["text"] for example in examples]	43	texts = [example["text"] for example in examples]
	@@ -52,11 +52,11 @@ class SDXLTrainer(Trainer):
52		52
53	def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None, **kwargs):	53	def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None, **kwargs):
54	device = model.device	54	device = model.device
55	- model_dtype = model.dtype	55	+ model_dtype = model.dtype
56	pixel_values = inputs["pixel_values"].to(device, dtype=model_dtype)	56	pixel_values = inputs["pixel_values"].to(device, dtype=model_dtype)
57	prompts = inputs["text"]	57	prompts = inputs["text"]
58	bsz = pixel_values.shape[0]	58	bsz = pixel_values.shape[0]
59	-	59	+
60	with torch.no_grad():	60	with torch.no_grad():
61	prompt_embeds, _, pooled_prompt_embeds, _ = self.pipe.encode_prompt(	61	prompt_embeds, _, pooled_prompt_embeds, _ = self.pipe.encode_prompt(
62	prompt=prompts,	62	prompt=prompts,
	@@ -73,7 +73,7 @@ class SDXLTrainer(Trainer):
73	timesteps = torch.randint(	73	timesteps = torch.randint(
74	0, self.noise_scheduler.config.num_train_timesteps, (bsz,), device=device	74	0, self.noise_scheduler.config.num_train_timesteps, (bsz,), device=device
75	).long()	75	).long()
76	-	76	+
77	noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)	77	noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)
78		78
79	add_time_ids = torch.tensor(	79	add_time_ids = torch.tensor(
	@@ -105,7 +105,7 @@ class SDXLLoRAFineTuner:
105	self.setup_model()	105	self.setup_model()
106		106
107	def setup_model(self):	107	def setup_model(self):
108	- logger.info(f"Loading SDXL model from {self.args.model_path}")	108	+ logger.info(f"Loading SDXL model from {self.args.model_path}")
109	self.pipe = DiffusionPipeline.from_pretrained(	109	self.pipe = DiffusionPipeline.from_pretrained(
110	self.args.model_path,	110	self.args.model_path,
111	torch_dtype=self.dtype,	111	torch_dtype=self.dtype,
	@@ -118,7 +118,7 @@ class SDXLLoRAFineTuner:
118	self.vae = self.pipe.vae	118	self.vae = self.pipe.vae
119	self.text_encoder = self.pipe.text_encoder	119	self.text_encoder = self.pipe.text_encoder
120	self.text_encoder_2 = self.pipe.text_encoder_2	120	self.text_encoder_2 = self.pipe.text_encoder_2
121	-	121	+
122	self.noise_scheduler = DDPMScheduler.from_config(self.pipe.scheduler.config)	122	self.noise_scheduler = DDPMScheduler.from_config(self.pipe.scheduler.config)
123		123
124	self.vae.requires_grad_(False)	124	self.vae.requires_grad_(False)
	@@ -131,7 +131,7 @@ class SDXLLoRAFineTuner:
131	r=self.args.lora_rank,	131	r=self.args.lora_rank,
132	lora_alpha=self.args.lora_rank,	132	lora_alpha=self.args.lora_rank,
133	init_lora_weights="gaussian",	133	init_lora_weights="gaussian",
134	- target_modules=["to_k", "to_q", "to_v", "to_out.0"],	134	+ target_modules=["to_k", "to_q", "to_v", "to_out.0"],
135	)	135	)
136		136
137	self.unet.add_adapter(unet_lora_config)	137	self.unet.add_adapter(unet_lora_config)
	@@ -159,7 +159,7 @@ class SDXLLoRAFineTuner:
159	if image.mode != "RGB":	159	if image.mode != "RGB":
160	image = image.convert("RGB")	160	image = image.convert("RGB")
161	pixel_values.append(transform(image))	161	pixel_values.append(transform(image))
162	-	162	+
163	examples["pixel_values"] = pixel_values	163	examples["pixel_values"] = pixel_values
164	return examples	164	return examples
165		165
	@@ -176,12 +176,12 @@ class SDXLLoRAFineTuner:
176	if self.args.enable_profiler:	176	if self.args.enable_profiler:
177	if not os.path.exists(self.args.profiler_save_path):	177	if not os.path.exists(self.args.profiler_save_path):
178	os.makedirs(self.args.profiler_save_path)	178	os.makedirs(self.args.profiler_save_path)
179	-	179	+
180	profiling_save_path = os.path.join(self.args.profiler_save_path, 'sdxl', mod)	180	profiling_save_path = os.path.join(self.args.profiler_save_path, 'sdxl', mod)
181	-	181	+
182	prof = get_profile(	182	prof = get_profile(
183	- profiler_start_step=self.args.profiler_start_step,	183	+ profiler_start_step=self.args.profiler_start_step,
184	- profiler_end_step=self.args.profiler_end_step,	184	+ profiler_end_step=self.args.profiler_end_step,
185	profiling_save_path=profiling_save_path	185	profiling_save_path=profiling_save_path
186	)	186	)
187		187
	@@ -233,11 +233,11 @@ class SDXLLoRAFineTuner:
233	unet_lora_layers=unet_lora_state_dict,	233	unet_lora_layers=unet_lora_state_dict,
234	safe_serialization=True	234	safe_serialization=True
235	)	235	)
236	-	236	+
237	return train_result.metrics	237	return train_result.metrics
238		238
239	def main():	239	def main():
240	- parser = argparse.ArgumentParser(description="Train Stable Diffusion XL LoRA")	240	+ parser = argparse.ArgumentParser(description="Train Stable Diffusion XL LoRA")
241	parser.add_argument("--model_path", type=str, required=True, help="Path to the pretrained SDXL model")	241	parser.add_argument("--model_path", type=str, required=True, help="Path to the pretrained SDXL model")
242	parser.add_argument("--data_path", type=str, required=True, help="Path to training data (parquet file)")	242	parser.add_argument("--data_path", type=str, required=True, help="Path to training data (parquet file)")
243	parser.add_argument("--output_dir", type=str, default="./sdxl_lora_weights", help="Output directory for LoRA weights")	243	parser.add_argument("--output_dir", type=str, default="./sdxl_lora_weights", help="Output directory for LoRA weights")
	@@ -259,7 +259,7 @@ def main():
259	parser.add_argument("--npu-backend", type=str, default="mlir")	259	parser.add_argument("--npu-backend", type=str, default="mlir")
260	parser.add_argument("--mfusion", action="store_true",	260	parser.add_argument("--mfusion", action="store_true",
261	help="Enable MFusion for graph fusion optimization")	261	help="Enable MFusion for graph fusion optimization")
262	-	262	+
263	args = parser.parse_args()	263	args = parser.parse_args()
264	device_type = detect_device_type()	264	device_type = detect_device_type()
265	os.environ['TORCHINDUCTOR_NPU_BACKEND'] = args.npu_backend	265	os.environ['TORCHINDUCTOR_NPU_BACKEND'] = args.npu_backend

Mbenchmarks/llm/utils/utils.py+9-9

	@@ -4,31 +4,31 @@ import time
4	from transformers import TrainerCallback	4	from transformers import TrainerCallback
5		5
6	class TimingCallback(TrainerCallback):	6	class TimingCallback(TrainerCallback):
7	-	7	+
8	def __init__(self, profiler=None, mod='eager'):	8	def __init__(self, profiler=None, mod='eager'):
9	self.step_start_time = None	9	self.step_start_time = None
10	self.step_times = []	10	self.step_times = []
11	self.epoch_start_time = None	11	self.epoch_start_time = None
12	self.profiler = profiler	12	self.profiler = profiler
13	self.mod = mod	13	self.mod = mod
14	-	14	+
15	def on_epoch_begin(self, args, state, control, **kwargs):	15	def on_epoch_begin(self, args, state, control, **kwargs):
16	self.epoch_start_time = time.time()	16	self.epoch_start_time = time.time()
17	print(f"\n{'='*60}")	17	print(f"\n{'='*60}")
18	print(f"Epoch {state.epoch + 1 if state.epoch else 1} begin")	18	print(f"Epoch {state.epoch + 1 if state.epoch else 1} begin")
19	print(f"{'='*60}")	19	print(f"{'='*60}")
20	-	20	+
21	def on_epoch_end(self, args, state, control, **kwargs):	21	def on_epoch_end(self, args, state, control, **kwargs):
22	epoch_time = time.time() - self.epoch_start_time	22	epoch_time = time.time() - self.epoch_start_time
23	print(f"\n{'='*60}")	23	print(f"\n{'='*60}")
24	print(f"Epoch {int(state.epoch)} ended, time taken: {epoch_time:.2f}s")	24	print(f"Epoch {int(state.epoch)} ended, time taken: {epoch_time:.2f}s")
25	print(f"{'='*60}\n")	25	print(f"{'='*60}\n")
26	-	26	+
27	def on_step_begin(self, args, state, control, **kwargs):	27	def on_step_begin(self, args, state, control, **kwargs):
28	self.step_start_time = time.time()	28	self.step_start_time = time.time()
29	-	29	+
30	def on_step_end(self, args, state, control, **kwargs):	30	def on_step_end(self, args, state, control, **kwargs):
31	- step_time = (time.time() - self.step_start_time) * 1000	31	+ step_time = (time.time() - self.step_start_time) * 1000
32	self.step_times.append(step_time)	32	self.step_times.append(step_time)
33		33
34	if self.profiler is not None:	34	if self.profiler is not None:
	@@ -37,7 +37,7 @@ class TimingCallback(TrainerCallback):
37	except (AssertionError, StopIteration):	37	except (AssertionError, StopIteration):
38	self.profiler = None	38	self.profiler = None
39	print("Profiler done.")	39	print("Profiler done.")
40	-	40	+
41	if torch.cuda.is_available():	41	if torch.cuda.is_available():
42	torch.cuda.synchronize()	42	torch.cuda.synchronize()
43	elif torch.npu.is_available():	43	elif torch.npu.is_available():
	@@ -46,7 +46,7 @@ class TimingCallback(TrainerCallback):
46	print(f"[{self.mod}] step {state.global_step:4d} "	46	print(f"[{self.mod}] step {state.global_step:4d} "
47	f"step_time: {step_time:.3f}ms "	47	f"step_time: {step_time:.3f}ms "
48	f"loss: {state.log_history[-1].get('loss', 'N/A') if state.log_history else 'N/A'}")	48	f"loss: {state.log_history[-1].get('loss', 'N/A') if state.log_history else 'N/A'}")
49	-	49	+
50	def on_train_begin(self, args, state, control, **kwargs):	50	def on_train_begin(self, args, state, control, **kwargs):
51	print("============= training begin =============")	51	print("============= training begin =============")
52	if self.profiler is not None:	52	if self.profiler is not None:
	@@ -59,7 +59,7 @@ class TimingCallback(TrainerCallback):
59	valid_steps = self.step_times[100:] if len(self.step_times) > 100 else self.step_times[10:]	59	valid_steps = self.step_times[100:] if len(self.step_times) > 100 else self.step_times[10:]
60	avg_time = sum(valid_steps) / len(valid_steps)	60	avg_time = sum(valid_steps) / len(valid_steps)
61	total_time = sum(valid_steps)	61	total_time = sum(valid_steps)
62	-	62	+
63	print(f"\n{'='*60}")	63	print(f"\n{'='*60}")
64	print("Step time consumption statistics (keeping only the last 100 steps)")	64	print("Step time consumption statistics (keeping only the last 100 steps)")
65	print(f"[{self.mod}] total step:{len(valid_steps)}"	65	print(f"[{self.mod}] total step:{len(valid_steps)}"

Mbenchmarks/torchbench/extract_log.py+26-26

	@@ -8,7 +8,7 @@ import pandas as pd
8	def extract_log_info(log_file, profile_dir='./profile', output_file='log_analysis.xlsx'):	8	def extract_log_info(log_file, profile_dir='./profile', output_file='log_analysis.xlsx'):
9	"""	9	"""
10	提取日志文件中各模型的训练时间信息和profile数据	10	提取日志文件中各模型的训练时间信息和profile数据
11	-	11	+
12	Args:	12	Args:
13	log_file: 日志文件路径	13	log_file: 日志文件路径
14	profile_dir: profile数据目录路径	14	profile_dir: profile数据目录路径
	@@ -19,14 +19,14 @@ def extract_log_info(log_file, profile_dir='./profile', output_file='log_analysi
19	model_pattern = r'(npu\|cuda)\s+train\s+(\S+)'	19	model_pattern = r'(npu\|cuda)\s+train\s+(\S+)'
20	eager_pattern = r'eager.avg step time:\s([\d\.]+)\s*ms'	20	eager_pattern = r'eager.avg step time:\s([\d\.]+)\s*ms'
21	compile_pattern = r'compile.avg step time:\s([\d\.]+)\s*ms'	21	compile_pattern = r'compile.avg step time:\s([\d\.]+)\s*ms'
22	-	22	+
23	# 模式匹配：算子编译时间	23	# 模式匹配：算子编译时间
24	op_compile_time_pattern = r'op_compile_time:\s([\d\.]+)\sms'	24	op_compile_time_pattern = r'op_compile_time:\s([\d\.]+)\sms'
25	-	25	+
26	# 存储结果	26	# 存储结果
27	data = defaultdict(lambda: {	27	data = defaultdict(lambda: {
28	'accuracy': None, # 修改2: 存储完整的精度校验日志	28	'accuracy': None, # 修改2: 存储完整的精度校验日志
29	- 'eager_E2E_avg_time': None,	29	+ 'eager_E2E_avg_time': None,
30	'compile_E2E_avg_time': None,	30	'compile_E2E_avg_time': None,
31	'op_compile_time': None,	31	'op_compile_time': None,
32	'eager_OP_avg_time': None,	32	'eager_OP_avg_time': None,
	@@ -35,10 +35,10 @@ def extract_log_info(log_file, profile_dir='./profile', output_file='log_analysi
35	current_model = None	35	current_model = None
36	in_compile_block = False	36	in_compile_block = False
37	compile_block_lines = []	37	compile_block_lines = []
38	-	38	+
39	with open(log_file, 'r', encoding='utf-8') as f:	39	with open(log_file, 'r', encoding='utf-8') as f:
40	lines = f.readlines()	40	lines = f.readlines()
41	-	41	+
42	for _, line in enumerate(lines):	42	for _, line in enumerate(lines):
43	# 匹配模型名	43	# 匹配模型名
44	model_match = re.search(model_pattern, line)	44	model_match = re.search(model_pattern, line)
	@@ -50,57 +50,57 @@ def extract_log_info(log_file, profile_dir='./profile', output_file='log_analysi
50	if 'pass_accuracy' in log_line:	50	if 'pass_accuracy' in log_line:
51	data[current_model]['accuracy'] = log_line.strip()	51	data[current_model]['accuracy'] = log_line.strip()
52	break	52	break
53	-	53	+
54	# 提取op_compile_time	54	# 提取op_compile_time
55	for log_line in compile_block_lines:	55	for log_line in compile_block_lines:
56	op_compile_match = re.search(op_compile_time_pattern, log_line)	56	op_compile_match = re.search(op_compile_time_pattern, log_line)
57	if op_compile_match:	57	if op_compile_match:
58	data[current_model]['op_compile_time'] = float(op_compile_match.group(1))	58	data[current_model]['op_compile_time'] = float(op_compile_match.group(1))
59	break	59	break
60	-	60	+
61	# 重置状态	61	# 重置状态
62	current_model = model_match.group(2) # 第二个分组是模型名	62	current_model = model_match.group(2) # 第二个分组是模型名
63	in_compile_block = False	63	in_compile_block = False
64	compile_block_lines = []	64	compile_block_lines = []
65	continue	65	continue
66	-	66	+
67	# 匹配eager模式时间	67	# 匹配eager模式时间
68	if current_model:	68	if current_model:
69	eager_match = re.search(eager_pattern, line)	69	eager_match = re.search(eager_pattern, line)
70	if eager_match:	70	if eager_match:
71	data[current_model]['eager_E2E_avg_time'] = float(eager_match.group(1))	71	data[current_model]['eager_E2E_avg_time'] = float(eager_match.group(1))
72	-	72	+
73	# 匹配compile模式时间	73	# 匹配compile模式时间
74	compile_match = re.search(compile_pattern, line)	74	compile_match = re.search(compile_pattern, line)
75	if compile_match:	75	if compile_match:
76	data[current_model]['compile_E2E_avg_time'] = float(compile_match.group(1))	76	data[current_model]['compile_E2E_avg_time'] = float(compile_match.group(1))
77	in_compile_block = True	77	in_compile_block = True
78	compile_block_lines = [] # 开始收集compile块日志	78	compile_block_lines = [] # 开始收集compile块日志
79	-	79	+
80	# 收集compile块的日志	80	# 收集compile块的日志
81	if current_model and in_compile_block:	81	if current_model and in_compile_block:
82	compile_block_lines.append(line)	82	compile_block_lines.append(line)
83	-	83	+
84	# 处理最后一个模型的compile块日志	84	# 处理最后一个模型的compile块日志
85	if current_model and in_compile_block and compile_block_lines:	85	if current_model and in_compile_block and compile_block_lines:
86	for log_line in compile_block_lines:	86	for log_line in compile_block_lines:
87	if 'pass_accuracy' in log_line:	87	if 'pass_accuracy' in log_line:
88	data[current_model]['accuracy'] = log_line.strip()	88	data[current_model]['accuracy'] = log_line.strip()
89	break	89	break
90	-	90	+
91	for log_line in compile_block_lines:	91	for log_line in compile_block_lines:
92	op_compile_match = re.search(op_compile_time_pattern, log_line)	92	op_compile_match = re.search(op_compile_time_pattern, log_line)
93	if op_compile_match:	93	if op_compile_match:
94	data[current_model]['op_compile_time'] = float(op_compile_match.group(1))	94	data[current_model]['op_compile_time'] = float(op_compile_match.group(1))
95	break	95	break
96	-	96	+
97	# 需求3: 读取profile目录中的step_trace_time.csv文件	97	# 需求3: 读取profile目录中的step_trace_time.csv文件
98	profile_path = Path(profile_dir)	98	profile_path = Path(profile_dir)
99	if profile_path.exists():	99	if profile_path.exists():
100	for model_dir in profile_path.iterdir():	100	for model_dir in profile_path.iterdir():
101	if model_dir.is_dir():	101	if model_dir.is_dir():
102	model_name = model_dir.name	102	model_name = model_dir.name
103	-	103	+
104	# 读取eager模式下的step_trace_time.csv	104	# 读取eager模式下的step_trace_time.csv
105	# 修改3: 自动获取下一级目录	105	# 修改3: 自动获取下一级目录
106	eager_dir = model_dir / 'eager'	106	eager_dir = model_dir / 'eager'
	@@ -119,7 +119,7 @@ def extract_log_info(log_file, profile_dir='./profile', output_file='log_analysi
119	print(f"警告: {eager_csv_path} 中没有Computing列")	119	print(f"警告: {eager_csv_path} 中没有Computing列")
120	except Exception as e:	120	except Exception as e:
121	print(f"读取{eager_csv_path}时出错: {e}")	121	print(f"读取{eager_csv_path}时出错: {e}")
122	-	122	+
123	# 读取compile模式下的step_trace_time.csv	123	# 读取compile模式下的step_trace_time.csv
124	compile_dir = model_dir / 'compile'	124	compile_dir = model_dir / 'compile'
125	if compile_dir.exists() and compile_dir.is_dir():	125	if compile_dir.exists() and compile_dir.is_dir():
	@@ -139,7 +139,7 @@ def extract_log_info(log_file, profile_dir='./profile', output_file='log_analysi
139	print(f"读取{compile_csv_path}时出错: {e}")	139	print(f"读取{compile_csv_path}时出错: {e}")
140	else:	140	else:
141	print(f"警告: profile目录不存在: {profile_dir}")	141	print(f"警告: profile目录不存在: {profile_dir}")
142	-	142	+
143	# 转换为DataFrame	143	# 转换为DataFrame
144	df = pd.DataFrame.from_dict(data, orient='index')	144	df = pd.DataFrame.from_dict(data, orient='index')
145	df.index.name = 'model_name'	145	df.index.name = 'model_name'
	@@ -149,17 +149,17 @@ def extract_log_info(log_file, profile_dir='./profile', output_file='log_analysi
149	# 1. 计算E2E_speed_up_rate = eager_E2E_avg_time / compile_E2E_avg_time	149	# 1. 计算E2E_speed_up_rate = eager_E2E_avg_time / compile_E2E_avg_time
150	# 2. 计算OP_speed_up_rate = eager_OP_avg_time / compile_OP_avg_time	150	# 2. 计算OP_speed_up_rate = eager_OP_avg_time / compile_OP_avg_time
151	df['E2E_speed_up_rate'] = df.apply(	151	df['E2E_speed_up_rate'] = df.apply(
152	- lambda row: row['eager_E2E_avg_time'] / row['compile_E2E_avg_time']	152	+ lambda row: row['eager_E2E_avg_time'] / row['compile_E2E_avg_time']
153	- if row['compile_E2E_avg_time'] and row['compile_E2E_avg_time'] != 0 else None,	153	+ if row['compile_E2E_avg_time'] and row['compile_E2E_avg_time'] != 0 else None,
154	axis=1	154	axis=1
155	)	155	)
156	-	156	+
157	df['OP_speed_up_rate'] = df.apply(	157	df['OP_speed_up_rate'] = df.apply(
158	- lambda row: row['eager_OP_avg_time'] / row['compile_OP_avg_time']	158	+ lambda row: row['eager_OP_avg_time'] / row['compile_OP_avg_time']
159	- if row['compile_OP_avg_time'] and row['compile_OP_avg_time'] != 0 else None,	159	+ if row['compile_OP_avg_time'] and row['compile_OP_avg_time'] != 0 else None,
160	axis=1	160	axis=1
161	)	161	)
162	-	162	+
163	# 重排列顺序，使相关列更清晰	163	# 重排列顺序，使相关列更清晰
164	column_order = [	164	column_order = [
165	'model_name', 'accuracy', 'op_compile_time',	165	'model_name', 'accuracy', 'op_compile_time',
	@@ -168,7 +168,7 @@ def extract_log_info(log_file, profile_dir='./profile', output_file='log_analysi
168	]	168	]
169	existing_columns = [col for col in column_order if col in df.columns]	169	existing_columns = [col for col in column_order if col in df.columns]
170	df = df[existing_columns + [col for col in df.columns if col not in existing_columns]]	170	df = df[existing_columns + [col for col in df.columns if col not in existing_columns]]
171	-	171	+
172	# 保存到Excel	172	# 保存到Excel
173	df.to_excel(output_file, index=False)	173	df.to_excel(output_file, index=False)
174	return df	174	return df
	@@ -179,9 +179,9 @@ def main():
179	parser.add_argument('--log_file', required=True, help='日志文件路径')	179	parser.add_argument('--log_file', required=True, help='日志文件路径')
180	parser.add_argument('--profile_dir', default='./profile', help='profile数据目录路径，默认为./profile')	180	parser.add_argument('--profile_dir', default='./profile', help='profile数据目录路径，默认为./profile')
181	parser.add_argument('--output_file', default='analysis.xlsx', help='输出Excel文件路径，默认为log_analysis.xlsx')	181	parser.add_argument('--output_file', default='analysis.xlsx', help='输出Excel文件路径，默认为log_analysis.xlsx')
182	-	182	+
183	args = parser.parse_args()	183	args = parser.parse_args()
184	-	184	+
185	result = extract_log_info(	185	result = extract_log_info(
186	log_file=args.log_file,	186	log_file=args.log_file,
187	profile_dir=args.profile_dir,	187	profile_dir=args.profile_dir,

Mci/build.sh+1-1

	@@ -46,7 +46,7 @@ function parse_script_args() {
46	export PGO_MODE=1	46	export PGO_MODE=1
47	args_num=$((args_num-1))	47	args_num=$((args_num-1))
48	;;	48	;;
49	- 2)	49	+ 2)
50	export PGO_MODE=2	50	export PGO_MODE=2
51	args_num=$((args_num-1))	51	args_num=$((args_num-1))
52	;;	52	;;

Mexamples/libtorch_hccl/CMakeLists.txt+4-4

	@@ -68,9 +68,9 @@ endif()
68	add_executable(example_allreduce_hccl allreduce_hccl.cpp)	68	add_executable(example_allreduce_hccl allreduce_hccl.cpp)
69		69
70	# 链接库	70	# 链接库
71	-target_link_libraries(example_allreduce_hccl	71	+target_link_libraries(example_allreduce_hccl
72	- -ltorch	72	+ -ltorch
73	- -ltorch_cpu	73	+ -ltorch_cpu
74	- -lc10	74	+ -lc10
75	-ltorch_npu	75	-ltorch_npu
76	)	76	)

Mexamples/libtorch_hccl/allreduce_hccl.cpp+17-17

	@@ -24,64 +24,64 @@ int main(int argc, char** argv)
24	{	24	{
25	int rank = g_rank;	25	int rank = g_rank;
26	int size = g_size;	26	int size = g_size;
27	-	27	+
28	std::cout << "启动 HCCL allreduce 示例: rank=" << rank << ", size=" << size << std::endl;	28	std::cout << "启动 HCCL allreduce 示例: rank=" << rank << ", size=" << size << std::endl;
29	-	29	+
30	// 初始化NPU设备 - 使用npu字符串格式	30	// 初始化NPU设备 - 使用npu字符串格式
31	std::string device_str = "npu:" + std::to_string(rank);	31	std::string device_str = "npu:" + std::to_string(rank);
32	torch_npu::init_npu(device_str);	32	torch_npu::init_npu(device_str);
33	std::cout << "NPU设备 " << rank << " 初始化完成" << std::endl;	33	std::cout << "NPU设备 " << rank << " 初始化完成" << std::endl;
34	-	34	+
35	// 创建FileStore用于进程间通信协调	35	// 创建FileStore用于进程间通信协调
36	auto store = c10::make_intrusive<FileStore>("/tmp/c10d_hccl_example", size);	36	auto store = c10::make_intrusive<FileStore>("/tmp/c10d_hccl_example", size);
37	-	37	+
38	// 创建ProcessGroupHCCL选项	38	// 创建ProcessGroupHCCL选项
39	auto options = ProcessGroupHCCL::Options::create();	39	auto options = ProcessGroupHCCL::Options::create();
40	-	40	+
41	// 创建ProcessGroupHCCL实例	41	// 创建ProcessGroupHCCL实例
42	auto pg = c10::make_intrusive<ProcessGroupHCCL>(store, rank, size, options);	42	auto pg = c10::make_intrusive<ProcessGroupHCCL>(store, rank, size, options);
43	-	43	+
44	// 通过传NPU字符串构造NPU设备	44	// 通过传NPU字符串构造NPU设备
45	auto device = at::Device(device_str);	45	auto device = at::Device(device_str);
46	-	46	+
47	// 创建10个张量用于测试	47	// 创建10个张量用于测试
48	const auto ntensors = 10;	48	const auto ntensors = 10;
49	std::vector<at::Tensor> tensors;	49	std::vector<at::Tensor> tensors;
50	-	50	+
51	for (const auto i : c10::irange(ntensors)) {	51	for (const auto i : c10::irange(ntensors)) {
52	// 在NPU设备上创建全1张量	52	// 在NPU设备上创建全1张量
53	auto x = at::ones({1000, 16 * (i + 1)}, at::TensorOptions(device).dtype(at::kFloat));	53	auto x = at::ones({1000, 16 * (i + 1)}, at::TensorOptions(device).dtype(at::kFloat));
54	tensors.push_back(x);	54	tensors.push_back(x);
55	}	55	}
56	-	56	+
57	std::cout << "在NPU设备 " << rank << " 上创建了 " << ntensors << " 个张量" << std::endl;	57	std::cout << "在NPU设备 " << rank << " 上创建了 " << ntensors << " 个张量" << std::endl;
58	-	58	+
59	// 提交所有allreduce操作	59	// 提交所有allreduce操作
60	std::vector<c10::intrusive_ptr<Work>> pending;	60	std::vector<c10::intrusive_ptr<Work>> pending;
61	for (const auto i : c10::irange(ntensors)) {	61	for (const auto i : c10::irange(ntensors)) {
62	std::vector<at::Tensor> tmp = {tensors[i]};	62	std::vector<at::Tensor> tmp = {tensors[i]};
63	pending.push_back(pg->allreduce(tmp));	63	pending.push_back(pg->allreduce(tmp));
64	}	64	}
65	-	65	+
66	std::cout << "已提交 " << ntensors << " 个allreduce操作" << std::endl;	66	std::cout << "已提交 " << ntensors << " 个allreduce操作" << std::endl;
67	-	67	+
68	// 等待所有操作完成	68	// 等待所有操作完成
69	for (auto& work : pending) {	69	for (auto& work : pending) {
70	work->wait();	70	work->wait();
71	}	71	}
72	-	72	+
73	std::cout << "所有操作已完成!" << std::endl;	73	std::cout << "所有操作已完成!" << std::endl;
74	-	74	+
75	// 验证结果 - 打印前3个张量的第一个元素	75	// 验证结果 - 打印前3个张量的第一个元素
76	for (const auto i : c10::irange(std::min(ntensors, 3))) {	76	for (const auto i : c10::irange(std::min(ntensors, 3))) {
77	auto cpu_tensor = tensors[i].to(at::kCPU);	77	auto cpu_tensor = tensors[i].to(at::kCPU);
78	std::cout << "张量 " << i << " 第一个元素: " << cpu_tensor.data_ptr<float>()[0] << std::endl;	78	std::cout << "张量 " << i << " 第一个元素: " << cpu_tensor.data_ptr<float>()[0] << std::endl;
79	}	79	}
80	-	80	+
81	std::cout << "HCCL allreduce示例运行成功!" << std::endl;	81	std::cout << "HCCL allreduce示例运行成功!" << std::endl;
82	-	82	+
83	// 使用NPU设备结束需进行反初始化	83	// 使用NPU设备结束需进行反初始化
84	torch_npu::finalize_npu();	84	torch_npu::finalize_npu();
85	-	85	+
86	return 0;	86	return 0;
87	}	87	}

Msetup.py+5-5

	@@ -359,12 +359,12 @@ class CPPLibBuild(build_clib, object):
359	if DISABLE_RPC == 'FALSE':	359	if DISABLE_RPC == 'FALSE':
360	if check_tensorpipe_valid(BASE_DIR):	360	if check_tensorpipe_valid(BASE_DIR):
361	cmake_args.append('-DBUILD_TENSORPIPE=on')	361	cmake_args.append('-DBUILD_TENSORPIPE=on')
362	-	362	+
363	if ENABLE_LTO == "TRUE":	363	if ENABLE_LTO == "TRUE":
364	cmake_args.append('-DENABLE_LTO=on')	364	cmake_args.append('-DENABLE_LTO=on')
365	if PGO_MODE != 0:	365	if PGO_MODE != 0:
366	cmake_args.append('-DPGO_MODE=' + str(PGO_MODE))	366	cmake_args.append('-DPGO_MODE=' + str(PGO_MODE))
367	-	367	+
368	if USE_CXX11_ABI:	368	if USE_CXX11_ABI:
369	cmake_args.append('-DGLIBCXX_USE_CXX11_ABI=1')	369	cmake_args.append('-DGLIBCXX_USE_CXX11_ABI=1')
370		370
	@@ -551,11 +551,11 @@ def get_src_py_and_dst():
551	# 按原目录结构复制到目标路径	551	# 按原目录结构复制到目标路径
552	for src in codegen_files:	552	for src in codegen_files:
553	# 仅过滤指定目录下的根级__init__.py	553	# 仅过滤指定目录下的根级__init__.py
554	- if (exclude_root_init is not None and	554	+ if (exclude_root_init is not None and
555	- os.path.basename(src) == '__init__.py' and	555	+ os.path.basename(src) == '__init__.py' and
556	os.path.dirname(src) == exclude_root_init):	556	os.path.dirname(src) == exclude_root_init):
557	continue # 跳过op-plugin/codegen根目录的__init__.py	557	continue # 跳过op-plugin/codegen根目录的__init__.py
558	-	558	+
559	# 计算目标路径（保留原目录层级）	559	# 计算目标路径（保留原目录层级）
560	dst = os.path.join(	560	dst = os.path.join(
561	codegen_dst_dir,	561	codegen_dst_dir,

Mtest/_inductor/test_ascend_graph_pass.py+41-41

Mtest/_inductor/test_cat.py+1-1

	@@ -44,7 +44,7 @@ class TestCat(TestUtils):
44	output_tensor = torch.cat(slices, self.dim)	44	output_tensor = torch.cat(slices, self.dim)
45		45
46	return output_tensor	46	return output_tensor
47	-	47	+
48	@parametrize('shape', [(128, 50, 128)])	48	@parametrize('shape', [(128, 50, 128)])
49	@parametrize('dim', [2])	49	@parametrize('dim', [2])
50	@parametrize('dtype', ['float32', 'bfloat16'])	50	@parametrize('dtype', ['float32', 'bfloat16'])

Mtest/_inductor/test_check_accuracy.py+3-3

	@@ -16,7 +16,7 @@ class TestCheckAccuracy(TestUtils):
16	def test_check_accuracy_1(self):	16	def test_check_accuracy_1(self):
17	count_data_dump = 0	17	count_data_dump = 0
18	count_check_accuracy = 0	18	count_check_accuracy = 0
19	-	19	+
20	def run(x, y):	20	def run(x, y):
21	return F.relu(x) - y	21	return F.relu(x) - y
22		22
	@@ -51,7 +51,7 @@ class TestCheckAccuracy(TestUtils):
51	patch.object(torch_npu._inductor.runtime.triton_heuristics, "check_accuracy_triton", wrap_check_accuracy):	51	patch.object(torch_npu._inductor.runtime.triton_heuristics, "check_accuracy_triton", wrap_check_accuracy):
52	self.assertTrue(torch_npu._inductor.config.dump_fx_graph)	52	self.assertTrue(torch_npu._inductor.config.dump_fx_graph)
53	self.assertTrue(torch_npu._inductor.config.check_accuracy)	53	self.assertTrue(torch_npu._inductor.config.check_accuracy)
54	-	54	+
55	# Try run custom path and make sure no data_dump and check_accuracy is invoked.	55	# Try run custom path and make sure no data_dump and check_accuracy is invoked.
56	torch_npu._inductor.config.dump_fx_graph = False	56	torch_npu._inductor.config.dump_fx_graph = False
57	torch_npu._inductor.config.check_accuracy = False	57	torch_npu._inductor.config.check_accuracy = False
	@@ -66,7 +66,7 @@ class TestCheckAccuracy(TestUtils):
66	self.assertEqual(count_data_dump, 1)	66	self.assertEqual(count_data_dump, 1)
67	self.assertEqual(count_check_accuracy, 1)	67	self.assertEqual(count_check_accuracy, 1)
68	self.assertEqual(z, g)	68	self.assertEqual(z, g)
69	-	69	+
70		70
71	if __name__ == "__main__":	71	if __name__ == "__main__":
72	run_tests()	72	run_tests()

Mtest/_inductor/test_clamp.py+1-1

	@@ -68,7 +68,7 @@ class TestClamp(TestUtils):
68	self.assertEqual(std_result, inductor_result)	68	self.assertEqual(std_result, inductor_result)
69		69
70	@parametrize('shape', TestUtils._pointwise_demo_shapes)	70	@parametrize('shape', TestUtils._pointwise_demo_shapes)
71	- @parametrize('dtype', ['float16', 'float32', 'bfloat16', 'int32', 'int64'])	71	+ @parametrize('dtype', ['float16', 'float32', 'bfloat16', 'int32', 'int64'])
72	def test_pointwise_cases_min_only(self, shape, dtype):	72	def test_pointwise_cases_min_only(self, shape, dtype):
73	min_numel = 0	73	min_numel = 0
74		74

Mtest/_inductor/test_debug_msg.py+1-1

	@@ -14,7 +14,7 @@ os.environ["INDUCTOR_ASCEND_DUMP_FX_GRAPH"] = "1"
14	os.environ["TORCH_COMPILE_DEBUG"] = "1"	14	os.environ["TORCH_COMPILE_DEBUG"] = "1"
15		15
16		16
17	-class TestDebugMsg(TestUtils):	17	+class TestDebugMsg(TestUtils):
18	@parametrize('shape_x', [(32, 512, 64)])	18	@parametrize('shape_x', [(32, 512, 64)])
19	@parametrize('shape_y', [(32, 1, 64)])	19	@parametrize('shape_y', [(32, 1, 64)])
20	@parametrize('dtype', ['float32'])	20	@parametrize('dtype', ['float32'])

Mtest/_inductor/test_dropout_with_checkpoint_recompute.py+8-8

	@@ -6,14 +6,14 @@ from torch.testing._internal.common_utils import (
6	)	6	)
7	from testutils import TestUtils	7	from testutils import TestUtils
8	import torch_npu	8	import torch_npu
9	-	9	+
10	class TestDropoutWithCheckpointRecompute(TestUtils):	10	class TestDropoutWithCheckpointRecompute(TestUtils):
11	def test_dropout_with_checkpoint_recompute(self):	11	def test_dropout_with_checkpoint_recompute(self):
12	device = "npu"	12	device = "npu"
13	-	13	+
14	def gn(x):	14	def gn(x):
15	return torch.sigmoid(torch.dropout(torch.sigmoid(x), p=0.5, train=True))	15	return torch.sigmoid(torch.dropout(torch.sigmoid(x), p=0.5, train=True))
16	-	16	+
17	def fn(x):	17	def fn(x):
18	return checkpoint(	18	return checkpoint(
19	gn,	19	gn,
	@@ -21,22 +21,22 @@ class TestDropoutWithCheckpointRecompute(TestUtils):
21	use_reentrant=False,	21	use_reentrant=False,
22	preserve_rng_state=True,	22	preserve_rng_state=True,
23	)	23	)
24	-	24	+
25	x = torch.randn(4, 4, requires_grad=True, device=device)	25	x = torch.randn(4, 4, requires_grad=True, device=device)
26	-	26	+
27	torch.manual_seed(42)	27	torch.manual_seed(42)
28	eager_out = fn(x)	28	eager_out = fn(x)
29	eager_out.sum().backward()	29	eager_out.sum().backward()
30	eager_grad = x.grad.clone()	30	eager_grad = x.grad.clone()
31	-	31	+
32	x.grad = None	32	x.grad = None
33	-	33	+
34	torch.manual_seed(42)	34	torch.manual_seed(42)
35	compiled_fn = torch.compile(fn, backend="inductor")	35	compiled_fn = torch.compile(fn, backend="inductor")
36	compiled_out = compiled_fn(x)	36	compiled_out = compiled_fn(x)
37	compiled_out.sum().backward()	37	compiled_out.sum().backward()
38	compiled_grad = x.grad.clone()	38	compiled_grad = x.grad.clone()
39	-	39	+
40	self.assertEqual(eager_out, compiled_out)	40	self.assertEqual(eager_out, compiled_out)
41	self.assertEqual(eager_grad, compiled_grad)	41	self.assertEqual(eager_grad, compiled_grad)
42		42

Mtest/_inductor/test_embedding.py+1-1

	@@ -13,7 +13,7 @@ class TestEmbeddingDense(TestUtils):
13	# UT skip, reason: precision fail	13	# UT skip, reason: precision fail
14	# Added to pytorch-disable-tests.json	14	# Added to pytorch-disable-tests.json
15	def test_pointwise_cases(self):	15	def test_pointwise_cases(self):
16	-	16	+
17	arg0 = torch.tensor([[14, 1, 2, 10, 0, 10, 0],	17	arg0 = torch.tensor([[14, 1, 2, 10, 0, 10, 0],
18	[9, 13, 13, 4, 7, 15, 14],	18	[9, 13, 13, 4, 7, 15, 14],
19	[8, 0, 3, 15, 4, 2, 6],	19	[8, 0, 3, 15, 4, 2, 6],

Mtest/_inductor/test_exceptions.py+5-5

	@@ -27,13 +27,13 @@ import torch_npu
27	size_hints={'y0': 16384, 'x1': 32}, tile_hint=TileHint.DEFAULT,	27	size_hints={'y0': 16384, 'x1': 32}, tile_hint=TileHint.DEFAULT,
28	filename=__file__,	28	filename=__file__,
29	triton_meta={'signature': {'in_ptr0': 'fp16', 'in_ptr1': 'fp16', 'out_ptr0': '*fp16', 'y0_numel': 'i32', 'x1_numel': 'i32'},	29	triton_meta={'signature': {'in_ptr0': 'fp16', 'in_ptr1': 'fp16', 'out_ptr0': '*fp16', 'y0_numel': 'i32', 'x1_numel': 'i32'},
30	- 'device': DeviceProperties(type='npu', index=0, multi_processor_count=40, cc='Ascend910B3',	30	+ 'device': DeviceProperties(type='npu', index=0, multi_processor_count=40, cc='Ascend910B3',
31	major=None, regs_per_multiprocessor=None, max_threads_per_multi_processor=None, warp_size=32),	31	major=None, regs_per_multiprocessor=None, max_threads_per_multi_processor=None, warp_size=32),
32	'constants': {}, 'mix_mode': 'aiv'},	32	'constants': {}, 'mix_mode': 'aiv'},
33	- inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_unk_fused_add_0', 'mutated_arg_names': [],	33	+ inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_unk_fused_add_0', 'mutated_arg_names': [],
34	- 'backend_hash': 'bc71dba4086164e7ac2b0779fa861dbf7467f0265d4a57b8f48cf6dda02b150f', 'split_axis': [0],	34	+ 'backend_hash': 'bc71dba4086164e7ac2b0779fa861dbf7467f0265d4a57b8f48cf6dda02b150f', 'split_axis': [0],
35	- 'tiling_axis': [0, 1], 'no_loop_axis': [1], 'axis_names': ['y0', 'x1'], 'low_dims': {1}, 'numof_reduction_axis': 0,	35	+ 'tiling_axis': [0, 1], 'no_loop_axis': [1], 'axis_names': ['y0', 'x1'], 'low_dims': {1}, 'numof_reduction_axis': 0,
36	- 'split_axis_dtype': torch.float16, 'dual_reduction': False, 'traced_graph_hash': 'TRACED_GRAPH_HASH',	36	+ 'split_axis_dtype': torch.float16, 'dual_reduction': False, 'traced_graph_hash': 'TRACED_GRAPH_HASH',
37	'traced_graph_dir': 'TRACED_GRAPH_DIR'},	37	'traced_graph_dir': 'TRACED_GRAPH_DIR'},
38	min_elem_per_thread=0	38	min_elem_per_thread=0
39	)	39	)

Mtest/_inductor/test_inductor_fallback_list.py+6-6

Mtest/_inductor/test_lazy_register.py+1-1

	@@ -9,7 +9,7 @@ import torch_npu
9	@skipIf(torch_npu.utils._dynamo.is_inductor_npu_initialized(), reason="Inductor npu has initialized")	9	@skipIf(torch_npu.utils._dynamo.is_inductor_npu_initialized(), reason="Inductor npu has initialized")
10	class TestLazyRegister(TestUtils):	10	class TestLazyRegister(TestUtils):
11		11
12	-	12	+
13	def test_disable_register_inductor_npu(self):	13	def test_disable_register_inductor_npu(self):
14	torch_npu.utils._dynamo.disable_register_inductor_npu()	14	torch_npu.utils._dynamo.disable_register_inductor_npu()
15		15

Mtest/_inductor/test_mlir_enable.py+1-1

	@@ -6,7 +6,7 @@ import torch_npu
6		6
7		7
8	class TestAdd(TestUtils):	8	class TestAdd(TestUtils):
9	-	9	+
10		10
11	def op_calc(self, first_element, second_element):	11	def op_calc(self, first_element, second_element):
12	result = first_element + second_element	12	result = first_element + second_element

Mtest/_inductor/test_multi_stream.py+4-4

	@@ -80,7 +80,7 @@ class FakeNode:
80		80
81	def mark_run(self):	81	def mark_run(self):
82	pass	82	pass
83	-	83	+
84	def get_nodes(self):	84	def get_nodes(self):
85	return []	85	return []
86		86
	@@ -369,7 +369,7 @@ class TestMultiStreamPass(TestUtils):
369	"""	369	"""
370	fake_graph = FakeGraph(cpp_wrapper=False)	370	fake_graph = FakeGraph(cpp_wrapper=False)
371	node = FakeNode(workspace_size=2048, name="test_02")	371	node = FakeNode(workspace_size=2048, name="test_02")
372	- with V.set_graph_handler(fake_graph):	372	+ with V.set_graph_handler(fake_graph):
373	kernel = self.define_catlass_template_kernel()	373	kernel = self.define_catlass_template_kernel()
374	kernel.call_kernel(	374	kernel.call_kernel(
375	name="test_kernel",	375	name="test_kernel",
	@@ -407,7 +407,7 @@ class TestMultiStreamPass(TestUtils):
407	def render():	407	def render():
408	return "fake_src"	408	return "fake_src"
409	return kernel, render	409	return kernel, render
410	-	410	+
411	template_node.node.make_kernel_render = fake_make_kernel_render	411	template_node.node.make_kernel_render = fake_make_kernel_render
412	with V.set_graph_handler(fake_graph):	412	with V.set_graph_handler(fake_graph):
413	scheduler = self.define_catlass_scheduling()	413	scheduler = self.define_catlass_scheduling()
	@@ -480,7 +480,7 @@ class TestMultiStreamPass(TestUtils):
480	assert kwargs["name"] == "k0"	480	assert kwargs["name"] == "k0"
481	assert kwargs["origin_node"] is origin_node	481	assert kwargs["origin_node"] is origin_node
482		482
483	-	483	+
484	@patch("torch_npu._inductor.codegen.scheduling.is_multi_stream", return_value=False)	484	@patch("torch_npu._inductor.codegen.scheduling.is_multi_stream", return_value=False)
485	@patch("torch_npu._inductor.codegen.scheduling.V")	485	@patch("torch_npu._inductor.codegen.scheduling.V")
486	def test_codegen_node_schedule_single_stream(self, mock_V, mock_multi_stream):	486	def test_codegen_node_schedule_single_stream(self, mock_V, mock_multi_stream):

Mtest/_inductor/test_npu_fusion_attention_pass.py+3-3

	@@ -45,8 +45,8 @@ class TestFusionAttentionUnchangePass(TestUtils):
45	self.assertIsInstance(inductor_result["getitem_4"], torch.Tensor, "The output parameter 'seed' of the npu_fusion_attention_v3 should be Tensor")	45	self.assertIsInstance(inductor_result["getitem_4"], torch.Tensor, "The output parameter 'seed' of the npu_fusion_attention_v3 should be Tensor")
46	self.assertIsInstance(inductor_result["getitem_5"], torch.Tensor, "The output parameter 'offset' of the npu_fusion_attention_v3 should be Tensor")	46	self.assertIsInstance(inductor_result["getitem_5"], torch.Tensor, "The output parameter 'offset' of the npu_fusion_attention_v3 should be Tensor")
47	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	47	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
48	-	48	+
49	-	49	+
50	def test_ut_cases(self):	50	def test_ut_cases(self):
51	primals_1 = torch.randn(2, 8, 16, 64, dtype=torch.float32, device="npu")	51	primals_1 = torch.randn(2, 8, 16, 64, dtype=torch.float32, device="npu")
52	primals_2 = torch.randn(2, 8, 16, 64, dtype=torch.float32, device="npu")	52	primals_2 = torch.randn(2, 8, 16, 64, dtype=torch.float32, device="npu")
	@@ -54,7 +54,7 @@ class TestFusionAttentionUnchangePass(TestUtils):
54	model = FusionAttentionUnchangeModel()	54	model = FusionAttentionUnchangeModel()
55	graph_module = fx.symbolic_trace(model)	55	graph_module = fx.symbolic_trace(model)
56	ShapeProp(graph_module).propagate(primals_1, primals_2, primals_3)	56	ShapeProp(graph_module).propagate(primals_1, primals_2, primals_3)
57	-	57	+
58	# 应用优化 Pass	58	# 应用优化 Pass
59	from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fusion_attention_v3_pass	59	from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fusion_attention_v3_pass
60	fusion_attention_v3_pass(graph_module.graph)	60	fusion_attention_v3_pass(graph_module.graph)

Mtest/_inductor/test_npu_kernel_features.py+5-5

	@@ -10,7 +10,7 @@ class TestNumeList(TestCase):
10	def test_numels(self):	10	def test_numels(self):
11	numel_list = NumelList([2, 3, 4])	11	numel_list = NumelList([2, 3, 4])
12	self.assertEqual(numel_list.numels(), 24)	12	self.assertEqual(numel_list.numels(), 24)
13	-	13	+
14	def test_equality(self):	14	def test_equality(self):
15	numel_list1 = NumelList([2, 3, 4])	15	numel_list1 = NumelList([2, 3, 4])
16	numel_list2 = NumelList([2, 3, 4])	16	numel_list2 = NumelList([2, 3, 4])
	@@ -45,11 +45,11 @@ class TestNumeList(TestCase):
45	numel_list1 = NumelList([2, 3, 5])	45	numel_list1 = NumelList([2, 3, 5])
46	numel_list2 = NumelList([2, 3, 4])	46	numel_list2 = NumelList([2, 3, 4])
47	self.assertTrue(numel_list1 >= numel_list2)	47	self.assertTrue(numel_list1 >= numel_list2)
48	-	48	+
49	def test_modulo(self):	49	def test_modulo(self):
50	numel_list = NumelList([2, 3, 4])	50	numel_list = NumelList([2, 3, 4])
51	self.assertEqual(numel_list % 5, 4)	51	self.assertEqual(numel_list % 5, 4)
52	-	52	+
53	def test_division(self):	53	def test_division(self):
54	numel_list = NumelList([2, 3, 4])	54	numel_list = NumelList([2, 3, 4])
55	self.assertEqual(numel_list / 2, 12.0)	55	self.assertEqual(numel_list / 2, 12.0)
	@@ -63,8 +63,8 @@ class TestNumeList(TestCase):
63	def test_addition(self):	63	def test_addition(self):
64	numel_list = NumelList([2, 3, 4])	64	numel_list = NumelList([2, 3, 4])
65	self.assertEqual(numel_list + 2, 26)	65	self.assertEqual(numel_list + 2, 26)
66	- self.assertEqual(2 + numel_list, 26)	66	+ self.assertEqual(2 + numel_list, 26)
67	-	67	+
68	def test_hash(self):	68	def test_hash(self):
69	# 测试相同内容的hash值相同	69	# 测试相同内容的hash值相同
70	numel_list1 = NumelList([2, 3, 4])	70	numel_list1 = NumelList([2, 3, 4])

Mtest/_inductor/test_opensora_graph1.py+1-1

	@@ -196,7 +196,7 @@ class TestModel(TestUtils):
196	mul_1: "i64[]" = torch.ops.aten.mul.Tensor(primals_3, 2)	196	mul_1: "i64[]" = torch.ops.aten.mul.Tensor(primals_3, 2)
197	mul_2: "i64[]" = torch.ops.aten.mul.Tensor(primals_4, 2)	197	mul_2: "i64[]" = torch.ops.aten.mul.Tensor(primals_4, 2)
198	return [permute, mul, mul_1, mul_2]	198	return [permute, mul, mul_1, mul_2]
199	-	199	+
200	primals_1 = torch.randn((1, 8, 30, 40, 1, 2, 2, 8), device=device_npu, dtype=torch.float32)	200	primals_1 = torch.randn((1, 8, 30, 40, 1, 2, 2, 8), device=device_npu, dtype=torch.float32)
201	primals_2 = torch.tensor((1), device=device_npu, dtype=torch.int64)	201	primals_2 = torch.tensor((1), device=device_npu, dtype=torch.int64)
202	primals_3 = torch.tensor((1), device=device_npu, dtype=torch.int64)	202	primals_3 = torch.tensor((1), device=device_npu, dtype=torch.int64)

Mtest/_inductor/test_rng_prims.py+1-1

	@@ -15,7 +15,7 @@ class TestRNGPrims(TestUtils):
15	def test_default(self):	15	def test_default(self):
16	register_run_and_save_rng_state_op()	16	register_run_and_save_rng_state_op()
17	x = torch.randn(4, 4).to("npu")	17	x = torch.randn(4, 4).to("npu")
18	- args = (x,)	18	+ args = (x,)
19	kwargs = {}	19	kwargs = {}
20	expected_rng_state = torch_npu.npu.get_rng_state()	20	expected_rng_state = torch_npu.npu.get_rng_state()
21	rng_state, out = run_and_save_rng_state(lambda x: x, args, *kwargs)	21	rng_state, out = run_and_save_rng_state(lambda x: x, args, *kwargs)

Mtest/_inductor/test_run_and_save_rng_state.py+1-1

	@@ -40,7 +40,7 @@ class TestRunAndSaveRngState(TestUtils):
40	self.op_calc(like, device, dtype)	40	self.op_calc(like, device, dtype)
41		41
42	self.assertEqual(res1_eager, res2_eager)	42	self.assertEqual(res1_eager, res2_eager)
43	- self.assertTrue(torch.equal(rng_state1_eager, rng_state2_eager))	43	+ self.assertTrue(torch.equal(rng_state1_eager, rng_state2_eager))
44		44
45	instantiate_parametrized_tests(TestRunAndSaveRngState)	45	instantiate_parametrized_tests(TestRunAndSaveRngState)
46		46

Mtest/_inductor/test_where.py+1-1

	@@ -11,7 +11,7 @@ class TestWhere(TestUtils):
11		11
12	@unittest.skip("it takes too long, not supported yet")	12	@unittest.skip("it takes too long, not supported yet")
13	@parametrize('shape', TestUtils._pointwise_demo_shapes)	13	@parametrize('shape', TestUtils._pointwise_demo_shapes)
14	- @parametrize('dtype', ['float16', 'float32', 'bfloat16', 'int32'])	14	+ @parametrize('dtype', ['float16', 'float32', 'bfloat16', 'int32'])
15	def test_pointwise_cases(self, shape, dtype):	15	def test_pointwise_cases(self, shape, dtype):
16	first_element = self._generate_tensor(shape, dtype)	16	first_element = self._generate_tensor(shape, dtype)
17	second_element = self._generate_tensor(shape, dtype)	17	second_element = self._generate_tensor(shape, dtype)

Mtest/allocator/host/test_expandable_host_allocator.py+2-2

	@@ -45,7 +45,7 @@ class TestHostCachingAllocator(TestCase):
45	def setUpClass(cls):	45	def setUpClass(cls):
46	os.environ['PYTORCH_NPU_ALLOC_CONF'] = 'pin_memory_expandable_segments:True'	46	os.environ['PYTORCH_NPU_ALLOC_CONF'] = 'pin_memory_expandable_segments:True'
47		47
48	-	48	+
49	def test_allocate_with_block_cut(self):	49	def test_allocate_with_block_cut(self):
50	# 申请一个64M的tensor, 会预选绑定80M物理内存	50	# 申请一个64M的tensor, 会预选绑定80M物理内存
51	memory_64m = torch.ones([1024, 1024, 16]).pin_memory()	51	memory_64m = torch.ones([1024, 1024, 16]).pin_memory()
	@@ -130,7 +130,7 @@ class TestHostCachingAllocator(TestCase):
130	npu_output = npu_copy_op_exec(npu_input1, cpu_out2)	130	npu_output = npu_copy_op_exec(npu_input1, cpu_out2)
131	self.assertRtolEqual(cpu_output, npu_output)	131	self.assertRtolEqual(cpu_output, npu_output)
132		132
133	-	133	+
134	def test_event_free(self):	134	def test_event_free(self):
135	tensor = torch.ones([1024, 1024, 16]).npu()	135	tensor = torch.ones([1024, 1024, 16]).npu()
136	tensor_cpu = torch.ones([1024, 1024, 16]).pin_memory()	136	tensor_cpu = torch.ones([1024, 1024, 16]).pin_memory()

Mtest/allocator/test_hostallocator.py+10-10

	@@ -24,20 +24,20 @@ def _collect():
24		24
25		25
26	class TestHostCachingAllocatorBasic(TestCase):	26	class TestHostCachingAllocatorBasic(TestCase):
27	- def test_pin_memory_on_non_blocking_copy(self):	27	+ def test_pin_memory_on_non_blocking_copy(self):
28	t_acc = torch.randn(100).to(torch.accelerator.current_accelerator())	28	t_acc = torch.randn(100).to(torch.accelerator.current_accelerator())
29	t_host = t_acc.to("cpu", non_blocking=True)	29	t_host = t_acc.to("cpu", non_blocking=True)
30	torch.accelerator.synchronize()	30	torch.accelerator.synchronize()
31	self.assertTrue(t_host.is_pinned())	31	self.assertTrue(t_host.is_pinned())
32	self.assertEqual(t_acc.cpu(), t_host)	32	self.assertEqual(t_acc.cpu(), t_host)
33	-	33	+
34	def test_pin_memory_reuse(self):	34	def test_pin_memory_reuse(self):
35	t = torch.FloatTensor([1]).pin_memory()	35	t = torch.FloatTensor([1]).pin_memory()
36	ptr = t.data_ptr()	36	ptr = t.data_ptr()
37	del t	37	del t
38	t_new = torch.FloatTensor([1]).pin_memory()	38	t_new = torch.FloatTensor([1]).pin_memory()
39	self.assertEqual(t_new.data_ptr(), ptr)	39	self.assertEqual(t_new.data_ptr(), ptr)
40	-	40	+
41	def test_to_non_blocking(self):	41	def test_to_non_blocking(self):
42	stream = torch_npu.npu.current_stream()	42	stream = torch_npu.npu.current_stream()
43		43
	@@ -59,7 +59,7 @@ class TestHostCachingAllocatorBasic(TestCase):
59	device="npu" if dst == "cpu" else "cpu",	59	device="npu" if dst == "cpu" else "cpu",
60	pin_memory=True if dst == "npu" else False)	60	pin_memory=True if dst == "npu" else False)
61	_test_to_non_blocking(src, try_non_blocking, dst)	61	_test_to_non_blocking(src, try_non_blocking, dst)
62	-	62	+
63	def test_pin_memory_basic(self):	63	def test_pin_memory_basic(self):
64	a = torch.Tensor([1])	64	a = torch.Tensor([1])
65	b = a.pin_memory()	65	b = a.pin_memory()
	@@ -68,7 +68,7 @@ class TestHostCachingAllocatorBasic(TestCase):
68	self.assertTrue(a.data_ptr() != b.data_ptr())	68	self.assertTrue(a.data_ptr() != b.data_ptr())
69	self.assertTrue(b.data_ptr() != c.data_ptr())	69	self.assertTrue(b.data_ptr() != c.data_ptr())
70	self.assertTrue(b.data_ptr() == d.data_ptr())	70	self.assertTrue(b.data_ptr() == d.data_ptr())
71	-	71	+
72	def test_malloc_copykernel(self):	72	def test_malloc_copykernel(self):
73	a = torch.Tensor([1])	73	a = torch.Tensor([1])
74	b = torch.Tensor([1])	74	b = torch.Tensor([1])
	@@ -117,7 +117,7 @@ class TestHostCachingAllocatorBasic(TestCase):
117		117
118	torch.npu.synchronize()	118	torch.npu.synchronize()
119	self.assertTrue(not errs)	119	self.assertTrue(not errs)
120	-	120	+
121	def test_pin_memory_on_views_and_clones(self):	121	def test_pin_memory_on_views_and_clones(self):
122	base = torch.randn(1024, 1024)	122	base = torch.randn(1024, 1024)
123	view = base[:512, :].pin_memory()	123	view = base[:512, :].pin_memory()
	@@ -155,7 +155,7 @@ class CountingDataset(Dataset):
155		155
156	def __getitem__(self, i):	156	def __getitem__(self, i):
157	return i	157	return i
158	-	158	+
159	def __len__(self):	159	def __len__(self):
160	return self.n	160	return self.n
161		161
	@@ -163,7 +163,7 @@ class CountingDataset(Dataset):
163	class DictDataset(Dataset):	163	class DictDataset(Dataset):
164	def __len__(self):	164	def __len__(self):
165	return 4	165	return 4
166	-	166	+
167	def __getitem__(self, ndx):	167	def __getitem__(self, ndx):
168	return {	168	return {
169	'a_tensor': torch.empty(4, 2).fill_(ndx),	169	'a_tensor': torch.empty(4, 2).fill_(ndx),
	@@ -179,7 +179,7 @@ class StringDataset(Dataset):
179		179
180	def __len__(self):	180	def __len__(self):
181	return len(self.s)	181	return len(self.s)
182	-	182	+
183	def __getitem__(self, ndx):	183	def __getitem__(self, ndx):
184	return (self.s[ndx], ndx)	184	return (self.s[ndx], ndx)
185		185
	@@ -194,7 +194,7 @@ class SimpleCustomBatch:
194	self.inp = self.inp.pin_memory()	194	self.inp = self.inp.pin_memory()
195	self.tgt = self.tgt.pin_memory()	195	self.tgt = self.tgt.pin_memory()
196	return self	196	return self
197	-	197	+
198	def is_pinned(self):	198	def is_pinned(self):
199	return self.inp.is_pinned() and self.tgt.is_pinned()	199	return self.inp.is_pinned() and self.tgt.is_pinned()
200		200

Mtest/allocator/test_pluggable_allocator_extensions.py+1-1

	@@ -57,7 +57,7 @@ class TestPluggableAllocator(TestCase):
57	build_directory=cls.build_directory,	57	build_directory=cls.build_directory,
58	verbose=True,	58	verbose=True,
59	)	59	)
60	-	60	+
61	def test_pluggable_allocator(self):	61	def test_pluggable_allocator(self):
62	os_path = os.path.join(TestPluggableAllocator.build_directory, 'pluggable_allocator_extensions.so')	62	os_path = os.path.join(TestPluggableAllocator.build_directory, 'pluggable_allocator_extensions.so')
63	# Load the allocator	63	# Load the allocator

Mtest/custom_ops/test_masked_softmax_with_rel_pos_bias.py+1-1

	@@ -10,7 +10,7 @@ from torch_npu.testing.common_utils import SupportedDevices
10	class TestMaskedSoftmaxWithRelPosBias(TestCase):	10	class TestMaskedSoftmaxWithRelPosBias(TestCase):
11		11
12	def supported_op_exec(self, x, relative_pos_bias, atten_mask):	12	def supported_op_exec(self, x, relative_pos_bias, atten_mask):
13	- # add + add + softmax	13	+ # add + add + softmax
14	y = torch.add(x, atten_mask)	14	y = torch.add(x, atten_mask)
15	y = torch.add(y, relative_pos_bias)	15	y = torch.add(y, relative_pos_bias)
16	softmax_out = torch.nn.functional.softmax(y, dim=-1)	16	softmax_out = torch.nn.functional.softmax(y, dim=-1)

Mtest/custom_ops/test_npu_anti_quant.py+6-6

	@@ -32,12 +32,12 @@ class TestAntiQuant(TestCase):
32	scale = torch.broadcast_to(scale, input_x.shape)	32	scale = torch.broadcast_to(scale, input_x.shape)
33	if offset is None:	33	if offset is None:
34	offset = torch.zeros_like(scale)	34	offset = torch.zeros_like(scale)
35	-	35	+
36	x = input_x.to(torch.float32)	36	x = input_x.to(torch.float32)
37	-	37	+
38	offset_temp = x + offset	38	offset_temp = x + offset
39	output = offset_temp * scale	39	output = offset_temp * scale
40	-	40	+
41	output = output.to(dst_dtype)	41	output = output.to(dst_dtype)
42	return output.cpu().detach()	42	return output.cpu().detach()
43		43
	@@ -57,12 +57,12 @@ class TestAntiQuant(TestCase):
57	[[np.int32, -1, [10, 25]], [np.float32, -1, [200]], [np.float32, -1, [200]], torch.float16, None],	57	[[np.int32, -1, [10, 25]], [np.float32, -1, [200]], [np.float32, -1, [200]], torch.float16, None],
58	[[np.int32, -1, [10, 25]], [np.float32, -1, [200]], [np.float32, -1, [200]], torch.bfloat16, None],	58	[[np.int32, -1, [10, 25]], [np.float32, -1, [200]], [np.float32, -1, [200]], torch.bfloat16, None],
59	]	59	]
60	-	60	+
61	for item in shape_format:	61	for item in shape_format:
62	cpu_input_x, npu_input_x = create_common_tensor(item[0], -127, 127)	62	cpu_input_x, npu_input_x = create_common_tensor(item[0], -127, 127)
63	cpu_scale, npu_scale = create_common_tensor(item[1], -100, 100)	63	cpu_scale, npu_scale = create_common_tensor(item[1], -100, 100)
64	cpu_offset, npu_offset = (None, None) if item[2] is None else create_common_tensor(item[2], -100, 100)	64	cpu_offset, npu_offset = (None, None) if item[2] is None else create_common_tensor(item[2], -100, 100)
65	-	65	+
66	npu_output = self.npu_op_exec(npu_input_x, npu_scale, npu_offset, *item[3:])	66	npu_output = self.npu_op_exec(npu_input_x, npu_scale, npu_offset, *item[3:])
67	custom_output = self.custom_op_exec(cpu_input_x, cpu_scale, cpu_offset, *item[3:])	67	custom_output = self.custom_op_exec(cpu_input_x, cpu_scale, cpu_offset, *item[3:])
68		68
	@@ -77,7 +77,7 @@ class TestAntiQuant(TestCase):
77	shape_format = [	77	shape_format = [
78	[[np.int8, -1, [10, 100]], [np.float32, -1, [100]], [np.float32, -1, [100]], torch.float16, None],	78	[[np.int8, -1, [10, 100]], [np.float32, -1, [100]], [np.float32, -1, [100]], torch.float16, None],
79	]	79	]
80	-	80	+
81	for item in shape_format:	81	for item in shape_format:
82	_, npu_input_x = create_common_tensor(item[0], -127, 127)	82	_, npu_input_x = create_common_tensor(item[0], -127, 127)
83	_, npu_scale = create_common_tensor(item[1], -100, 100)	83	_, npu_scale = create_common_tensor(item[1], -100, 100)

Mtest/custom_ops/test_npu_dtype_cast.py+1-1

	@@ -41,7 +41,7 @@ class TestDtypeCast(TestCase):
41	y = torch_npu.npu_dtype_cast(x, torch.complex128)	41	y = torch_npu.npu_dtype_cast(x, torch.complex128)
42	grad_fn = str(y.grad_fn)	42	grad_fn = str(y.grad_fn)
43	self.assertTrue("NpuDtypeCastBackward" in grad_fn)	43	self.assertTrue("NpuDtypeCastBackward" in grad_fn)
44	-	44	+
45	with self.assertRaisesRegex(RuntimeError, r'grad can be implicitly created'):	45	with self.assertRaisesRegex(RuntimeError, r'grad can be implicitly created'):
46	y.sum().backward()	46	y.sum().backward()
47		47

Mtest/custom_ops/test_npu_stride_copy.py+1-1

	@@ -17,7 +17,7 @@ class TestNpuStrideCopy(TestCase):
17	output = torch_npu.npu_stride_copy(input1, size, stride, storage_offset)	17	output = torch_npu.npu_stride_copy(input1, size, stride, storage_offset)
18	output = output.cpu().numpy()	18	output = output.cpu().numpy()
19	return output	19	return output
20	-	20	+
21	@unittest.skip("skip test_npu_stride_copy now")	21	@unittest.skip("skip test_npu_stride_copy now")
22	def test_npu_stride_copy(self):	22	def test_npu_stride_copy(self):
23	shape_format = [	23	shape_format = [

Mtest/custom_ops/test_resize_.py+1-1

	@@ -21,7 +21,7 @@ class TestResize(TestCase):
21	out_tensor_npu = torch.masked_select(input_data_npu, mask_npu, out=out_tensor_npu)	21	out_tensor_npu = torch.masked_select(input_data_npu, mask_npu, out=out_tensor_npu)
22	out_tensor = torch.masked_select(input_data, mask, out=out_tensor)	22	out_tensor = torch.masked_select(input_data, mask, out=out_tensor)
23	self.assertRtolEqual(out_tensor_npu, out_tensor)	23	self.assertRtolEqual(out_tensor_npu, out_tensor)
24	-	24	+
25	def test_resize_ncdhw(self):	25	def test_resize_ncdhw(self):
26	out_tensor = torch.empty((1, 1, 1, 1, 1), dtype=torch.float16).npu()	26	out_tensor = torch.empty((1, 1, 1, 1, 1), dtype=torch.float16).npu()
27	shape = [25]	27	shape = [25]

Mtest/distributed/elastic/agent/server/test_healthcheckserver_api.py+1-1

	@@ -2,7 +2,7 @@
2	Add validation cases for torch.distributed.elastic.agent.server.health_check_server APIs.	2	Add validation cases for torch.distributed.elastic.agent.server.health_check_server APIs.
3		3
4	1. PyTorch community tests do not cover HealthCheckServer APIs, so this file is added.	4	1. PyTorch community tests do not cover HealthCheckServer APIs, so this file is added.
5	-2. This file validates :	5	+2. This file validates :
6	torch.distributed.elastic.agent.server.health_check_server.HealthCheckServer	6	torch.distributed.elastic.agent.server.health_check_server.HealthCheckServer
7	torch.distributed.elastic.agent.server.health_check_server.HealthCheckServer.start	7	torch.distributed.elastic.agent.server.health_check_server.HealthCheckServer.start
8	torch.distributed.elastic.agent.server.health_check_server.HealthCheckServer.stop	8	torch.distributed.elastic.agent.server.health_check_server.HealthCheckServer.stop

Mtest/distributed/tensor/test_math_ops.py+20-20

	@@ -325,7 +325,7 @@ class TestConv2d(NPUDTensorTestBase):
325	output_tensor = torch_npu.npu_conv2d(input_tensor, weight_tensor, bias, stride, padding, dilation, groups)	325	output_tensor = torch_npu.npu_conv2d(input_tensor, weight_tensor, bias, stride, padding, dilation, groups)
326	output_dtensor = torch_npu.npu_conv2d(input_dtensor, weight_dtensor, bias, stride, padding, dilation, groups)	326	output_dtensor = torch_npu.npu_conv2d(input_dtensor, weight_dtensor, bias, stride, padding, dilation, groups)
327	self.assertEqual(output_dtensor.full_tensor(), output_tensor)	327	self.assertEqual(output_dtensor.full_tensor(), output_tensor)
328	-	328	+
329	@SupportedDevices(['Ascend910B'])	329	@SupportedDevices(['Ascend910B'])
330	@skipIfUnsupportMultiNPU(2)	330	@skipIfUnsupportMultiNPU(2)
331	@with_comms	331	@with_comms
	@@ -431,7 +431,7 @@ class TestConv2d(NPUDTensorTestBase):
431	input_dgrad, weight_dgrad, bias_dgrad = torch_npu.npu_conv2d_backward(input_dtensor, grad_output_dtensor, weight_dtensor, stride, padding, dilation, groups, output_mask)	431	input_dgrad, weight_dgrad, bias_dgrad = torch_npu.npu_conv2d_backward(input_dtensor, grad_output_dtensor, weight_dtensor, stride, padding, dilation, groups, output_mask)
432	self.assertEqual(input_dgrad.full_tensor(), input_grad)	432	self.assertEqual(input_dgrad.full_tensor(), input_grad)
433	self.assertEqual(weight_dgrad.full_tensor(), weight_grad)	433	self.assertEqual(weight_dgrad.full_tensor(), weight_grad)
434	-	434	+
435	@SupportedDevices(['Ascend910B'])	435	@SupportedDevices(['Ascend910B'])
436	@skipIfUnsupportMultiNPU(2)	436	@skipIfUnsupportMultiNPU(2)
437	@with_comms	437	@with_comms
	@@ -461,7 +461,7 @@ class TestConv2d(NPUDTensorTestBase):
461	self.assertEqual(input_dgrad.full_tensor(), input_grad)	461	self.assertEqual(input_dgrad.full_tensor(), input_grad)
462	self.assertEqual(weight_dgrad.full_tensor(), weight_grad)	462	self.assertEqual(weight_dgrad.full_tensor(), weight_grad)
463	self.assertEqual(bias_dgrad.full_tensor(), bias_grad)	463	self.assertEqual(bias_dgrad.full_tensor(), bias_grad)
464	-	464	+
465	@SupportedDevices(['Ascend910B'])	465	@SupportedDevices(['Ascend910B'])
466	@skipIfUnsupportMultiNPU(2)	466	@skipIfUnsupportMultiNPU(2)
467	@with_comms	467	@with_comms
	@@ -490,7 +490,7 @@ class TestConv2d(NPUDTensorTestBase):
490	input_dgrad, weight_dgrad, bias_dgrad = torch_npu.npu_conv2d_backward(input_dtensor, grad_output_dtensor, weight_dtensor, stride, padding, dilation, groups, output_mask)	490	input_dgrad, weight_dgrad, bias_dgrad = torch_npu.npu_conv2d_backward(input_dtensor, grad_output_dtensor, weight_dtensor, stride, padding, dilation, groups, output_mask)
491	self.assertEqual(input_dgrad.full_tensor(), input_grad)	491	self.assertEqual(input_dgrad.full_tensor(), input_grad)
492	self.assertEqual(weight_dgrad.full_tensor(), weight_grad)	492	self.assertEqual(weight_dgrad.full_tensor(), weight_grad)
493	-	493	+
494	@SupportedDevices(['Ascend910B'])	494	@SupportedDevices(['Ascend910B'])
495	@skipIfUnsupportMultiNPU(2)	495	@skipIfUnsupportMultiNPU(2)
496	@with_comms	496	@with_comms
	@@ -572,7 +572,7 @@ class TestGroupedMatmulAdd(NPUDTensorTestBase):
572	torch_npu.npu_grouped_matmul_add_(y, x, weight, group_list, transpose_x=transpose_x, transpose_weight=transpose_weight, group_type=group_type)	572	torch_npu.npu_grouped_matmul_add_(y, x, weight, group_list, transpose_x=transpose_x, transpose_weight=transpose_weight, group_type=group_type)
573	torch_npu.npu_grouped_matmul_add_(y_dtensor, x_dtensor, weight_dtensor, group_list_dtensor, transpose_x=transpose_x, transpose_weight=transpose_weight, group_type=group_type)	573	torch_npu.npu_grouped_matmul_add_(y_dtensor, x_dtensor, weight_dtensor, group_list_dtensor, transpose_x=transpose_x, transpose_weight=transpose_weight, group_type=group_type)
574	self.assertEqual(y_dtensor.full_tensor(), y)	574	self.assertEqual(y_dtensor.full_tensor(), y)
575	-	575	+
576	@SupportedDevices(['Ascend910B'])	576	@SupportedDevices(['Ascend910B'])
577	@skipIfUnsupportMultiNPU(2)	577	@skipIfUnsupportMultiNPU(2)
578	@with_comms	578	@with_comms
	@@ -614,7 +614,7 @@ class TestCrossEntropyLoss(NPUDTensorTestBase):
614	return input_tuple	614	return input_tuple
615	else:	615	else:
616	input_tuple = (x, target, input_dtensor, target_dtensor, mesh)	616	input_tuple = (x, target, input_dtensor, target_dtensor, mesh)
617	-	617	+
618	return input_tuple	618	return input_tuple
619		619
620		620
	@@ -630,7 +630,7 @@ class TestCrossEntropyLoss(NPUDTensorTestBase):
630	self.assertEqual(loss_dtensor.full_tensor(), loss)	630	self.assertEqual(loss_dtensor.full_tensor(), loss)
631	self.assertEqual(log_prob_dtensor.full_tensor(), log_prob)	631	self.assertEqual(log_prob_dtensor.full_tensor(), log_prob)
632		632
633	-	633	+
634	@SupportedDevices(['Ascend910B'])	634	@SupportedDevices(['Ascend910B'])
635	@skipIfUnsupportMultiNPU(2)	635	@skipIfUnsupportMultiNPU(2)
636	@with_comms	636	@with_comms
	@@ -643,7 +643,7 @@ class TestCrossEntropyLoss(NPUDTensorTestBase):
643	self.assertEqual(loss_dtensor.full_tensor(), loss)	643	self.assertEqual(loss_dtensor.full_tensor(), loss)
644	self.assertEqual(log_prob_dtensor.full_tensor(), log_prob)	644	self.assertEqual(log_prob_dtensor.full_tensor(), log_prob)
645		645
646	-	646	+
647	@SupportedDevices(['Ascend910B'])	647	@SupportedDevices(['Ascend910B'])
648	@skipIfUnsupportMultiNPU(2)	648	@skipIfUnsupportMultiNPU(2)
649	@with_comms	649	@with_comms
	@@ -671,7 +671,7 @@ class TestCrossEntropyLoss(NPUDTensorTestBase):
671	self.assertEqual(loss_dtensor.full_tensor(), loss)	671	self.assertEqual(loss_dtensor.full_tensor(), loss)
672	self.assertEqual(log_prob_dtensor.full_tensor(), log_prob)	672	self.assertEqual(log_prob_dtensor.full_tensor(), log_prob)
673		673
674	-	674	+
675	@SupportedDevices(['Ascend910B'])	675	@SupportedDevices(['Ascend910B'])
676	@skipIfUnsupportMultiNPU(2)	676	@skipIfUnsupportMultiNPU(2)
677	@with_comms	677	@with_comms
	@@ -685,14 +685,14 @@ class TestCrossEntropyLoss(NPUDTensorTestBase):
685	loss_dtensor.backward()	685	loss_dtensor.backward()
686	self.assertEqual(input_dtensor.grad.full_tensor(), x.grad)	686	self.assertEqual(input_dtensor.grad.full_tensor(), x.grad)
687		687
688	-	688	+
689	@SupportedDevices(['Ascend910B'])	689	@SupportedDevices(['Ascend910B'])
690	@skipIfUnsupportMultiNPU(2)	690	@skipIfUnsupportMultiNPU(2)
691	@with_comms	691	@with_comms
692	def test_torch_npu_npu_cross_entropy_loss_backward_input_shard0_reduction_is_none(self):	692	def test_torch_npu_npu_cross_entropy_loss_backward_input_shard0_reduction_is_none(self):
693	reductions = ["none", "sum", "mean"]	693	reductions = ["none", "sum", "mean"]
694	x, target, input_dtensor, target_dtensor, mesh = self.generate_data_cross_entropy_loss(8, 8, [Shard(0)], [Shard(0)])	694	x, target, input_dtensor, target_dtensor, mesh = self.generate_data_cross_entropy_loss(8, 8, [Shard(0)], [Shard(0)])
695	-	695	+
696	for re in reductions:	696	for re in reductions:
697	loss, log_prob, _, _ = torch_npu.npu_cross_entropy_loss(x, target, reduction=re)	697	loss, log_prob, _, _ = torch_npu.npu_cross_entropy_loss(x, target, reduction=re)
698	loss_dtensor, log_prob_dtensor, _, _ = torch_npu.npu_cross_entropy_loss(input_dtensor, target_dtensor, reduction=re)	698	loss_dtensor, log_prob_dtensor, _, _ = torch_npu.npu_cross_entropy_loss(input_dtensor, target_dtensor, reduction=re)
	@@ -707,7 +707,7 @@ class TestCrossEntropyLoss(NPUDTensorTestBase):
707	loss.backward()	707	loss.backward()
708	loss_dtensor.backward()	708	loss_dtensor.backward()
709	self.assertEqual(input_dtensor.grad.full_tensor(), x.grad)	709	self.assertEqual(input_dtensor.grad.full_tensor(), x.grad)
710	-	710	+
711	@SupportedDevices(['Ascend910B'])	711	@SupportedDevices(['Ascend910B'])
712	@skipIfUnsupportMultiNPU(2)	712	@skipIfUnsupportMultiNPU(2)
713	@with_comms	713	@with_comms
	@@ -716,7 +716,7 @@ class TestCrossEntropyLoss(NPUDTensorTestBase):
716		716
717	loss, log_prob, _, _ = torch_npu.npu_cross_entropy_loss(x, target, reduction="sum")	717	loss, log_prob, _, _ = torch_npu.npu_cross_entropy_loss(x, target, reduction="sum")
718	loss_dtensor, log_prob_dtensor, _, _ = torch_npu.npu_cross_entropy_loss(input_dtensor, target_dtensor, reduction="sum")	718	loss_dtensor, log_prob_dtensor, _, _ = torch_npu.npu_cross_entropy_loss(input_dtensor, target_dtensor, reduction="sum")
719	-	719	+
720	loss.backward()	720	loss.backward()
721	loss_dtensor.backward()	721	loss_dtensor.backward()
722	self.assertEqual(input_dtensor.grad.full_tensor(), x.grad)	722	self.assertEqual(input_dtensor.grad.full_tensor(), x.grad)
	@@ -741,7 +741,7 @@ class TestRepeatInterleaveSelfInt(NPUDTensorTestBase):
741		741
742	output_dtensor = torch.repeat_interleave(input_dtensor, repeats_value, dim=1)	742	output_dtensor = torch.repeat_interleave(input_dtensor, repeats_value, dim=1)
743	output = torch.repeat_interleave(input_tensor, repeats_value, dim=1)	743	output = torch.repeat_interleave(input_tensor, repeats_value, dim=1)
744	-	744	+
745	self.assertEqual(output_dtensor.full_tensor(), output)	745	self.assertEqual(output_dtensor.full_tensor(), output)
746		746
747	@SupportedDevices(['Ascend910B'])	747	@SupportedDevices(['Ascend910B'])
	@@ -752,7 +752,7 @@ class TestRepeatInterleaveSelfInt(NPUDTensorTestBase):
752		752
753	output_dtensor = torch.repeat_interleave(input_dtensor, repeats_value, dim=1)	753	output_dtensor = torch.repeat_interleave(input_dtensor, repeats_value, dim=1)
754	output = torch.repeat_interleave(input_tensor, repeats_value, dim=1)	754	output = torch.repeat_interleave(input_tensor, repeats_value, dim=1)
755	-	755	+
756	self.assertEqual(output_dtensor.full_tensor(), output)	756	self.assertEqual(output_dtensor.full_tensor(), output)
757		757
758	@SupportedDevices(['Ascend910B'])	758	@SupportedDevices(['Ascend910B'])
	@@ -763,7 +763,7 @@ class TestRepeatInterleaveSelfInt(NPUDTensorTestBase):
763		763
764	output_dtensor = torch.repeat_interleave(input_dtensor, repeats_value, dim=1)	764	output_dtensor = torch.repeat_interleave(input_dtensor, repeats_value, dim=1)
765	output = torch.repeat_interleave(input_tensor, repeats_value, dim=1)	765	output = torch.repeat_interleave(input_tensor, repeats_value, dim=1)
766	-	766	+
767	self.assertEqual(output_dtensor.full_tensor(), output)	767	self.assertEqual(output_dtensor.full_tensor(), output)
768		768
769	@SupportedDevices(['Ascend910B'])	769	@SupportedDevices(['Ascend910B'])
	@@ -774,7 +774,7 @@ class TestRepeatInterleaveSelfInt(NPUDTensorTestBase):
774		774
775	output_dtensor = torch.repeat_interleave(input_dtensor, repeats_value)	775	output_dtensor = torch.repeat_interleave(input_dtensor, repeats_value)
776	output = torch.repeat_interleave(input_tensor, repeats_value)	776	output = torch.repeat_interleave(input_tensor, repeats_value)
777	-	777	+
778	self.assertEqual(output_dtensor.full_tensor(), output)	778	self.assertEqual(output_dtensor.full_tensor(), output)
779		779
780	@SupportedDevices(['Ascend910B'])	780	@SupportedDevices(['Ascend910B'])
	@@ -785,7 +785,7 @@ class TestRepeatInterleaveSelfInt(NPUDTensorTestBase):
785		785
786	output_dtensor = torch.repeat_interleave(input_dtensor, repeats_value, dim=1)	786	output_dtensor = torch.repeat_interleave(input_dtensor, repeats_value, dim=1)
787	output = torch.repeat_interleave(input_tensor, repeats_value, dim=1)	787	output = torch.repeat_interleave(input_tensor, repeats_value, dim=1)
788	-	788	+
789	self.assertEqual(output_dtensor.full_tensor(), output)	789	self.assertEqual(output_dtensor.full_tensor(), output)
790		790
791	@SupportedDevices(['Ascend910B'])	791	@SupportedDevices(['Ascend910B'])
	@@ -796,7 +796,7 @@ class TestRepeatInterleaveSelfInt(NPUDTensorTestBase):
796		796
797	output_dtensor = torch.repeat_interleave(input_dtensor, repeats_value, dim=1)	797	output_dtensor = torch.repeat_interleave(input_dtensor, repeats_value, dim=1)
798	output = torch.repeat_interleave(input_tensor, repeats_value, dim=1)	798	output = torch.repeat_interleave(input_tensor, repeats_value, dim=1)
799	-	799	+
800	self.assertEqual(output_dtensor.full_tensor(), output)	800	self.assertEqual(output_dtensor.full_tensor(), output)
801		801
802	@SupportedDevices(['Ascend910B'])	802	@SupportedDevices(['Ascend910B'])
	@@ -983,7 +983,7 @@ class TestKLDivLoss(NPUDTensorTestBase):
983		983
984	pred = torch.randn(4, 4, 4, device="npu", requires_grad=True)	984	pred = torch.randn(4, 4, 4, device="npu", requires_grad=True)
985	target = torch.randn(4, 4, 4, device="npu")	985	target = torch.randn(4, 4, 4, device="npu")
986	-	986	+
987	# def test_placement_comb(placements1, placements2):	987	# def test_placement_comb(placements1, placements2):
988	pred_dt = distribute_tensor(pred, mesh, [pred_placement])	988	pred_dt = distribute_tensor(pred, mesh, [pred_placement])
989	# pred_dt = pred_dt	989	# pred_dt = pred_dt

Mtest/distributed/test_allgather.py+2-2

	@@ -73,7 +73,7 @@ class HcclAllGatherTestBase(TestCase):
73	gather_tensor = list()	73	gather_tensor = list()
74	for input_tensor in inputlist:	74	for input_tensor in inputlist:
75	gather_tensor.append(torch.empty_like(input_tensor, device="cpu"))	75	gather_tensor.append(torch.empty_like(input_tensor, device="cpu"))
76	-	76	+
77	for i in range(world_size):	77	for i in range(world_size):
78	p = ctx.Process(	78	p = ctx.Process(
79	target=f,	79	target=f,
	@@ -184,7 +184,7 @@ class HcclAllGatherTest(HcclAllGatherTestBase):
184	for _ in range(dim):	184	for _ in range(dim):
185	shape_list.append(randint(1, max_value))	185	shape_list.append(randint(1, max_value))
186	return create_common_tensor([np.float32, format_list[randint(0, 3)], shape_list], -10, 10)	186	return create_common_tensor([np.float32, format_list[randint(0, 3)], shape_list], -10, 10)
187	-	187	+
188	for world_size in ranks:	188	for world_size in ranks:
189	cpu_excepted_result = list()	189	cpu_excepted_result = list()
190	npu_excepted_result = list()	190	npu_excepted_result = list()

Mtest/distributed/test_device.py+2-2

	@@ -15,7 +15,7 @@ class TestDevice(MultiProcessTestCase):
15	@property	15	@property
16	def world_size(self):	16	def world_size(self):
17	return 1	17	return 1
18	-	18	+
19	def test_event_create(self):	19	def test_event_create(self):
20	a = torch.full((3, 4), float(0), device='npu:0')	20	a = torch.full((3, 4), float(0), device='npu:0')
21	e = torch.npu.Event()	21	e = torch.npu.Event()
	@@ -31,7 +31,7 @@ class TestDevice(MultiProcessTestCase):
31	t.start()	31	t.start()
32	t.join()	32	t.join()
33	self.assertEqual(result[0], 1)	33	self.assertEqual(result[0], 1)
34	-	34	+
35	def test_event_isinstance(self):	35	def test_event_isinstance(self):
36	npu_event = torch.npu.Event()	36	npu_event = torch.npu.Event()
37	self.assertIsInstance(npu_event, torch.npu.Event)	37	self.assertIsInstance(npu_event, torch.npu.Event)

Mtest/distributed/test_device_mesh.py+1-1

	@@ -226,7 +226,7 @@ class DeviceMeshTestF(NPUDTensorTestBase):
226	mesh = torch.arange(self.world_size).to(self.rank)	226	mesh = torch.arange(self.world_size).to(self.rank)
227	with self.assertRaises(ValueError):	227	with self.assertRaises(ValueError):
228	device_mesh = DeviceMesh(self.device_type, mesh)	228	device_mesh = DeviceMesh(self.device_type, mesh)
229	-	229	+
230	@skipIfUnsupportMultiNPU(4)	230	@skipIfUnsupportMultiNPU(4)
231	@with_comms	231	@with_comms
232	def test_get_local_rank(self):	232	def test_get_local_rank(self):

Mtest/distributed/test_distributed.py+1-1

	@@ -673,7 +673,7 @@ class _DistTestBase(object):
673	torch.testing.assert_allclose(running_mean, all_input_var.mean(1))	673	torch.testing.assert_allclose(running_mean, all_input_var.mean(1))
674	torch.testing.assert_allclose(running_var.cpu(), all_input_var.cpu().var(1, unbiased=False))	674	torch.testing.assert_allclose(running_var.cpu(), all_input_var.cpu().var(1, unbiased=False))
675		675
676	- # need more 4 device, less 4 divice there may be accuracy issues	676	+ # need more 4 device, less 4 divice there may be accuracy issues
677	@skipIfUnsupportMultiNPU(4)	677	@skipIfUnsupportMultiNPU(4)
678	def test_DistributedDataParallel_SyncBatchNorm_Diff_Input_Sizes_Running_Value(self):	678	def test_DistributedDataParallel_SyncBatchNorm_Diff_Input_Sizes_Running_Value(self):
679	for bk in [True, False]:	679	for bk in [True, False]:

Mtest/distributed/test_flight_recorder.py+1-1

	@@ -907,7 +907,7 @@ class HcclHeartbeatDumpTest(HCCLTraceTestBase):
907	if self.rank == 0:	907	if self.rank == 0:
908	# sleep for heartbeat dump	908	# sleep for heartbeat dump
909	time.sleep(30)	909	time.sleep(30)
910	-	910	+
911	pg.allreduce(a).wait()	911	pg.allreduce(a).wait()
912		912
913	torch.npu.synchronize(device=device)	913	torch.npu.synchronize(device=device)

Mtest/distributed/test_hccl_shared_buffer.py+5-5

	@@ -71,12 +71,12 @@ class HcclSharedBufferTest(TestCase):
71	self.assertEqual(output, expected,	71	self.assertEqual(output, expected,
72	"rank {} world_size {} dtype {} shape {} Expect receive tensor {} but got {}.".format(	72	"rank {} world_size {} dtype {} shape {} Expect receive tensor {} but got {}.".format(
73	rank, world_size, expected.dtype, expected.shape, expected, output))	73	rank, world_size, expected.dtype, expected.shape, expected, output))
74	-	74	+
75	# For case where we want to examine the memory usage of ccl buffer	75	# For case where we want to examine the memory usage of ccl buffer
76	if mem_diff:	76	if mem_diff:
77	for pg in mem_diff:	77	for pg in mem_diff:
78	used_mem, expected_mem = mem_diff[pg]	78	used_mem, expected_mem = mem_diff[pg]
79	- self.assertLess(used_mem, expected_mem,	79	+ self.assertLess(used_mem, expected_mem,
80	f"Expected memory used to be less than {expected_mem} for {pg}, but got {used_mem}")	80	f"Expected memory used to be less than {expected_mem} for {pg}, but got {used_mem}")
81		81
82	for _ in range(world_size):	82	for _ in range(world_size):
	@@ -90,7 +90,7 @@ class HcclSharedBufferTest(TestCase):
90	op1_expected = 0	90	op1_expected = 0
91	for _ in range(world_size):	91	for _ in range(world_size):
92	op1_expected += inputs	92	op1_expected += inputs
93	-	93	+
94	#dist.ReduceOp.AVG	94	#dist.ReduceOp.AVG
95	expected = 0	95	expected = 0
96	for _ in range(world_size):	96	for _ in range(world_size):
	@@ -222,7 +222,7 @@ class HcclSharedBufferTest(TestCase):
222	del pg1	222	del pg1
223	del pg2	223	del pg2
224		224
225	- hccl_config3 = {"hccl_buffer_name": "subSharedBuffer"} #	225	+ hccl_config3 = {"hccl_buffer_name": "subSharedBuffer"} #
226	options.hccl_config = hccl_config3	226	options.hccl_config = hccl_config3
227	pg3 = dist.new_group(backend='hccl', ranks=ranks, pg_options=options)	227	pg3 = dist.new_group(backend='hccl', ranks=ranks, pg_options=options)
228	dist.all_reduce(input1, group=pg3, async_op=True)	228	dist.all_reduce(input1, group=pg3, async_op=True)
	@@ -348,7 +348,7 @@ class HcclSharedBufferTest(TestCase):
348	p2c.get()	348	p2c.get()
349		349
350	return True	350	return True
351	-	351	+
352	@classmethod	352	@classmethod
353	# pylint:disable=huawei-too-many-arguments	353	# pylint:disable=huawei-too-many-arguments
354	def _test_buffer_memory_with_deleted_pg(cls, rank, input1, world_size, init_pg, c2p, p2c, seq):	354	def _test_buffer_memory_with_deleted_pg(cls, rank, input1, world_size, init_pg, c2p, p2c, seq):

Mtest/distributed/test_nslb.py+1-1

	@@ -31,7 +31,7 @@ class HcclNslbTest(TestCase):
31	dist_group.all_reduce(input1)	31	dist_group.all_reduce(input1)
32		32
33	def _test_multiprocess(self, f, init_pg, input1, world_size, nslb_dir):	33	def _test_multiprocess(self, f, init_pg, input1, world_size, nslb_dir):
34	-	34	+
35	ctx = mp.get_context('spawn')	35	ctx = mp.get_context('spawn')
36		36
37	ps = []	37	ps = []

Mtest/distributed/test_options.py+17-17

Mtest/distributed/test_send_recv.py+1-1

	@@ -228,7 +228,7 @@ class HcclSendRecvDistTest(TestCase):
228	HcclSendRecvDistTest._test_send_recv_dist_with_internal_format_and_offset,	228	HcclSendRecvDistTest._test_send_recv_dist_with_internal_format_and_offset,
229	torch.randn(31, 31),	229	torch.randn(31, 31),
230	HcclSendRecvDistTest._init_dist_hccl)	230	HcclSendRecvDistTest._init_dist_hccl)
231	-	231	+
232	@unittest.skip("Temporarily disabled")	232	@unittest.skip("Temporarily disabled")
233	@skipIfUnsupportMultiNPU(4)	233	@skipIfUnsupportMultiNPU(4)
234	def test_send_recv_hccl_dist_with_p2p(self):	234	def test_send_recv_hccl_dist_with_p2p(self):

Mtest/distributed/test_store_api.py+14-14

	@@ -1,7 +1,7 @@
1	"""	1	"""
2	Add validation cases for torch.distributed APIs on NPU:	2	Add validation cases for torch.distributed APIs on NPU:
3	1. test/distributed/test_store.py from PyTorch community lacks sufficient API validations, so this file is added.	3	1. test/distributed/test_store.py from PyTorch community lacks sufficient API validations, so this file is added.
4	-2. This file validates	4	+2. This file validates
5	torch.distributed.FileStore.path	5	torch.distributed.FileStore.path
6	torch.distributed.Store.__init__	6	torch.distributed.Store.__init__
7	torch.distributed.Store.add	7	torch.distributed.Store.add
	@@ -37,14 +37,14 @@ class TestStoreAPIs(TestCase):
37	"""Test if FileStore actually uses the specified path for data exchange."""	37	"""Test if FileStore actually uses the specified path for data exchange."""
38	with tempfile.TemporaryDirectory() as temp_dir:	38	with tempfile.TemporaryDirectory() as temp_dir:
39	filename = os.path.join(temp_dir, "npu_filestore.txt")	39	filename = os.path.join(temp_dir, "npu_filestore.txt")
40	-	40	+
41	store_master = dist.FileStore(filename, 2)	41	store_master = dist.FileStore(filename, 2)
42	self.assertEqual(store_master.path, filename)	42	self.assertEqual(store_master.path, filename)
43	store_master.set("shared_key", "npu_data")	43	store_master.set("shared_key", "npu_data")
44		44
45	store_worker = dist.FileStore(filename, 2)	45	store_worker = dist.FileStore(filename, 2)
46	val = store_worker.get("shared_key")	46	val = store_worker.get("shared_key")
47	-	47	+
48	self.assertEqual(val, b"npu_data")	48	self.assertEqual(val, b"npu_data")
49		49
50	def test_store_init(self):	50	def test_store_init(self):
	@@ -62,22 +62,22 @@ class TestStoreAPIs(TestCase):
62	"""Test the add operation mathematically on a HashStore."""	62	"""Test the add operation mathematically on a HashStore."""
63	store = dist.HashStore()	63	store = dist.HashStore()
64	key = "test_add_key"	64	key = "test_add_key"
65	-	65	+
66	res1 = store.add(key, 5)	66	res1 = store.add(key, 5)
67	self.assertEqual(res1, 5)	67	self.assertEqual(res1, 5)
68	-	68	+
69	res2 = store.add(key, 10)	69	res2 = store.add(key, 10)
70	self.assertEqual(res2, 15)	70	self.assertEqual(res2, 15)
71	-	71	+
72	self.assertEqual(store.get(key), b"15")	72	self.assertEqual(store.get(key), b"15")
73	-	73	+
74	res3 = store.add(key, -3)	74	res3 = store.add(key, -3)
75	self.assertEqual(res3, 12)	75	self.assertEqual(res3, 12)
76		76
77	def test_store_timeout_behavior(self):	77	def test_store_timeout_behavior(self):
78	"""Test if the timeout property actively interrupts blocking operations."""	78	"""Test if the timeout property actively interrupts blocking operations."""
79	store = dist.HashStore()	79	store = dist.HashStore()
80	-	80	+
81	timeout_seconds = 1	81	timeout_seconds = 1
82	test_timeout = timedelta(seconds=timeout_seconds)	82	test_timeout = timedelta(seconds=timeout_seconds)
83	store.set_timeout(test_timeout)	83	store.set_timeout(test_timeout)
	@@ -92,7 +92,7 @@ class TestStoreAPIs(TestCase):
92	"Timeout" in str(context.exception) or "Wait timeout" in str(context.exception),	92	"Timeout" in str(context.exception) or "Wait timeout" in str(context.exception),
93	f"Exception message does not indicate timeout: {context.exception}"	93	f"Exception message does not indicate timeout: {context.exception}"
94	)	94	)
95	-	95	+
96	self.assertTrue(	96	self.assertTrue(
97	0.8 <= elapsed_time <= 2.5,	97	0.8 <= elapsed_time <= 2.5,
98	f"Actual wait time {elapsed_time:.2f}s did not respect the {timeout_seconds}s timeout."	98	f"Actual wait time {elapsed_time:.2f}s did not respect the {timeout_seconds}s timeout."
	@@ -102,19 +102,19 @@ class TestStoreAPIs(TestCase):
102	"""Test TCPStore host and port by establishing an actual Client-Server connection."""	102	"""Test TCPStore host and port by establishing an actual Client-Server connection."""
103	host = "127.0.0.1"	103	host = "127.0.0.1"
104	port = find_free_port()	104	port = find_free_port()
105	-	105	+
106	server_store = dist.TCPStore(	106	server_store = dist.TCPStore(
107	host_name=host,	107	host_name=host,
108	port=port,	108	port=port,
109	world_size=2,	109	world_size=2,
110	is_master=True,	110	is_master=True,
111	timeout=timedelta(seconds=5),	111	timeout=timedelta(seconds=5),
112	- wait_for_workers=False	112	+ wait_for_workers=False
113	)	113	)
114	-	114	+
115	self.assertEqual(server_store.host, host)	115	self.assertEqual(server_store.host, host)
116	self.assertEqual(server_store.port, port)	116	self.assertEqual(server_store.port, port)
117	-	117	+
118	server_store.set("tcp_key", "tcp_value")	118	server_store.set("tcp_key", "tcp_value")
119		119
120	client_store = dist.TCPStore(	120	client_store = dist.TCPStore(
	@@ -124,7 +124,7 @@ class TestStoreAPIs(TestCase):
124	is_master=False,	124	is_master=False,
125	timeout=timedelta(seconds=5)	125	timeout=timedelta(seconds=5)
126	)	126	)
127	-	127	+
128	client_store.wait(["tcp_key"], timedelta(seconds=5))	128	client_store.wait(["tcp_key"], timedelta(seconds=5))
129	val = client_store.get("tcp_key")	129	val = client_store.get("tcp_key")
130	self.assertEqual(val, b"tcp_value")	130	self.assertEqual(val, b"tcp_value")

Mtest/distributed/test_watchdog.py+2-2

	@@ -40,13 +40,13 @@ class ElasticLaunchTest(TestCase):
40	)	40	)
41	except Exception:	41	except Exception:
42	print("Program fail and exit")	42	print("Program fail and exit")
43	-	43	+
44	end_time = time.time()	44	end_time = time.time()
45	excution_time = end_time - start_time	45	excution_time = end_time - start_time
46	if excution_time > 120:	46	if excution_time > 120:
47	print(f"Excution time using time.time(): {excution_time} seconds")	47	print(f"Excution time using time.time(): {excution_time} seconds")
48	raise RuntimeError("Test case fail")	48	raise RuntimeError("Test case fail")
49	-	49	+
50		50
51	if __name__ == "__main__":	51	if __name__ == "__main__":
52	run_tests()	52	run_tests()

Mtest/distributed/test_with_device.py+1-1

	@@ -14,7 +14,7 @@ class TestWithDevice(TestCase):
14	# -int -> ignore, return -1	14	# -int -> ignore, return -1
15	# future exchangeDevice:	15	# future exchangeDevice:
16	# if < std::numeric_limits<c10::DeviceIndex>::min(), raise error	16	# if < std::numeric_limits<c10::DeviceIndex>::min(), raise error
17	- # else, ignore, return -1	17	+ # else, ignore, return -1
18	for i in [-258, -200.8, -128, -128.8, -127.99, -7, -7.88, -1, -0.2]:	18	for i in [-258, -200.8, -128, -128.8, -127.99, -7, -7.88, -1, -0.2]:
19	s = torch.npu.Stream(i)	19	s = torch.npu.Stream(i)
20	self.assertEqual(s.device_index, 1)	20	self.assertEqual(s.device_index, 1)

Mtest/dynamo/test_minifier.py+5-5

	@@ -51,7 +51,7 @@ inner(torch.randn(20, 20).to("{device}"))
51	self._test_after_dynamo(	51	self._test_after_dynamo(
52	"cuda", "relu_compile_error_TESTING_ONLY", "ReluCompileError"	52	"cuda", "relu_compile_error_TESTING_ONLY", "ReluCompileError"
53	)	53	)
54	-	54	+
55	@requires_npu()	55	@requires_npu()
56	def test_after_dynamo_npu_compile_error(self):	56	def test_after_dynamo_npu_compile_error(self):
57	self._test_after_dynamo(	57	self._test_after_dynamo(
	@@ -63,7 +63,7 @@ inner(torch.randn(20, 20).to("{device}"))
63	self._test_after_dynamo(	63	self._test_after_dynamo(
64	"cuda", "relu_runtime_error_TESTING_ONLY", "ReluRuntimeError"	64	"cuda", "relu_runtime_error_TESTING_ONLY", "ReluRuntimeError"
65	)	65	)
66	-	66	+
67	@requires_npu()	67	@requires_npu()
68	def test_after_dynamo_npu_runtime_error(self):	68	def test_after_dynamo_npu_runtime_error(self):
69	self._test_after_dynamo(	69	self._test_after_dynamo(
	@@ -75,7 +75,7 @@ inner(torch.randn(20, 20).to("{device}"))
75	self._test_after_dynamo(	75	self._test_after_dynamo(
76	"cuda", "relu_accuracy_error_TESTING_ONLY", "AccuracyError"	76	"cuda", "relu_accuracy_error_TESTING_ONLY", "AccuracyError"
77	)	77	)
78	-	78	+
79	@requires_npu()	79	@requires_npu()
80	def test_after_dynamo_npu_accuracy_error(self):	80	def test_after_dynamo_npu_accuracy_error(self):
81	self._test_after_dynamo(	81	self._test_after_dynamo(
	@@ -134,7 +134,7 @@ inner(torch.randn(20, 20, requires_grad=True) + 1)
134	self._test_after_dynamo_backend_passes(	134	self._test_after_dynamo_backend_passes(
135	"cuda", "relu_runtime_error_TESTING_ONLY"	135	"cuda", "relu_runtime_error_TESTING_ONLY"
136	)	136	)
137	-	137	+
138	@requires_npu()	138	@requires_npu()
139	def test_after_dynamo_npu_runtime_backend_passes(self):	139	def test_after_dynamo_npu_runtime_backend_passes(self):
140	self._test_after_dynamo_backend_passes(	140	self._test_after_dynamo_backend_passes(
	@@ -146,7 +146,7 @@ inner(torch.randn(20, 20, requires_grad=True) + 1)
146	self._test_after_dynamo_backend_passes(	146	self._test_after_dynamo_backend_passes(
147	"cuda", "relu_accuracy_error_TESTING_ONLY"	147	"cuda", "relu_accuracy_error_TESTING_ONLY"
148	)	148	)
149	-	149	+
150	@requires_npu()	150	@requires_npu()
151	def test_after_dynamo_npu_accuracy_backend_passes(self):	151	def test_after_dynamo_npu_accuracy_backend_passes(self):
152	self._test_after_dynamo_backend_passes(	152	self._test_after_dynamo_backend_passes(

Mtest/dynamo/test_npu_backend.py+1-1

	@@ -14,7 +14,7 @@ class TestNpuBackend(TestCase):
14	eager_result = func(x)	14	eager_result = func(x)
15	dynamo_result = dynamo_func(x)	15	dynamo_result = dynamo_func(x)
16	self.assertEqual(eager_result, dynamo_result)	16	self.assertEqual(eager_result, dynamo_result)
17	-	17	+
18		18
19	if __name__ == "__main__":	19	if __name__ == "__main__":
20	from torch._dynamo.test_case import run_tests	20	from torch._dynamo.test_case import run_tests

Mtest/dynamo/test_npugraph_ex.py+2-2

	@@ -44,14 +44,14 @@ class TestNpuGraphEx(TestCase):
44	def custom_compiler(gm: torch.fx.GraphModule, example_inputs):	44	def custom_compiler(gm: torch.fx.GraphModule, example_inputs):
45	compiled_graph = torch.npu.npugraph_ex.compile_fx(gm, example_inputs)	45	compiled_graph = torch.npu.npugraph_ex.compile_fx(gm, example_inputs)
46	return compiled_graph	46	return compiled_graph
47	-	47	+
48	def custom_compiler_with_options(gm: torch.fx.GraphModule, example_inputs):	48	def custom_compiler_with_options(gm: torch.fx.GraphModule, example_inputs):
49	test_kwargs = {	49	test_kwargs = {
50	"clone_input": False	50	"clone_input": False
51	}	51	}
52	compiled_graph = torch.npu.npugraph_ex.compile_fx(gm, example_inputs, test_kwargs)	52	compiled_graph = torch.npu.npugraph_ex.compile_fx(gm, example_inputs, test_kwargs)
53	return compiled_graph	53	return compiled_graph
54	-	54	+
55	def my_backend(gm: torch.fx.GraphModule, example_inputs):	55	def my_backend(gm: torch.fx.GraphModule, example_inputs):
56	return aot_module_simplified(gm, example_inputs, fw_compiler=custom_compiler)	56	return aot_module_simplified(gm, example_inputs, fw_compiler=custom_compiler)
57		57

Mtest/dynamo/test_subclasses.py+1-1

	@@ -904,7 +904,7 @@ class TestNestedTensor(torch._dynamo.test_case.TestCase):
904	@unittest.skipIf(not torch.cuda.is_available(), "requires cuda")	904	@unittest.skipIf(not torch.cuda.is_available(), "requires cuda")
905	def test_basic_autograd_inductor(self):	905	def test_basic_autograd_inductor(self):
906	self._test_autograd("inductor")	906	self._test_autograd("inductor")
907	-	907	+
908	@unittest.skipIf(not torch.npu.is_available(), "requires npu")	908	@unittest.skipIf(not torch.npu.is_available(), "requires npu")
909	def test_basic_autograd_npu_backend(self):	909	def test_basic_autograd_npu_backend(self):
910	npu_backend = torchair.get_npu_backend()	910	npu_backend = torchair.get_npu_backend()

Mtest/dynamo/test_torchair_no_init.py+4-4

Mtest/nn/test_linear_functions.py+1-1

	@@ -20,7 +20,7 @@ class TestLinearFunctions(TestCase):
20	npu_output = F.linear(npu_input, npu_weight)	20	npu_output = F.linear(npu_input, npu_weight)
21		21
22	self.assertRtolEqual(cpu_output.numpy(), npu_output.cpu().numpy())	22	self.assertRtolEqual(cpu_output.numpy(), npu_output.cpu().numpy())
23	-	23	+
24	@unittest.skip("skip test_bilinear now")	24	@unittest.skip("skip test_bilinear now")
25	def test_bilinear(self):	25	def test_bilinear(self):
26	input1 = torch.randn(10, 30)	26	input1 = torch.randn(10, 30)

Mtest/nn/test_loss_functions.py+1-1

	@@ -279,7 +279,7 @@ class TestLossFunctions(TestCase):
279	npu_output = F.hinge_embedding_loss(npu_input, npu_targets)	279	npu_output = F.hinge_embedding_loss(npu_input, npu_targets)
280		280
281	self.assertRtolEqual(cpu_output.detach().numpy(), npu_output.detach().cpu().numpy())	281	self.assertRtolEqual(cpu_output.detach().numpy(), npu_output.detach().cpu().numpy())
282	-	282	+
283	@unittest.skip("skip test_kl_div now")	283	@unittest.skip("skip test_kl_div now")
284	def test_kl_div(self):	284	def test_kl_div(self):
285	input1 = torch.randn(5, 3)	285	input1 = torch.randn(5, 3)

Mtest/nn/test_modules_api.py+10-10

	@@ -16,7 +16,7 @@ class TestNPUModuleAPIs(TestCase):
16	def test_module_device_consistency(self):	16	def test_module_device_consistency(self):
17	"""验证Module设备迁移后参数/缓冲区设备一致"""	17	"""验证Module设备迁移后参数/缓冲区设备一致"""
18	m = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.BatchNorm1d(20)).to(device)	18	m = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.BatchNorm1d(20)).to(device)
19	-	19	+
20	for p in m.parameters():	20	for p in m.parameters():
21	self.assertEqual(p.device, device)	21	self.assertEqual(p.device, device)
22	for b in m.buffers():	22	for b in m.buffers():
	@@ -29,13 +29,13 @@ class TestNPUModuleAPIs(TestCase):
29	m.dict = nn.ModuleDict({"linear": nn.Linear(10, 20), "base": base})	29	m.dict = nn.ModuleDict({"linear": nn.Linear(10, 20), "base": base})
30	m.list = nn.ModuleList([nn.Linear(20, 30), base])	30	m.list = nn.ModuleList([nn.Linear(20, 30), base])
31	m.to(device)	31	m.to(device)
32	-	32	+
33	sd = m.state_dict()	33	sd = m.state_dict()
34	torch.npu.synchronize()	34	torch.npu.synchronize()
35	-	35	+
36	for v in sd.values():	36	for v in sd.values():
37	self.assertEqual(v.device, device)	37	self.assertEqual(v.device, device)
38	-	38	+
39	base2 = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.BatchNorm1d(20))	39	base2 = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.BatchNorm1d(20))
40	m2 = nn.Module()	40	m2 = nn.Module()
41	m2.dict = nn.ModuleDict({"linear": nn.Linear(10, 20), "base": base2})	41	m2.dict = nn.ModuleDict({"linear": nn.Linear(10, 20), "base": base2})
	@@ -43,20 +43,20 @@ class TestNPUModuleAPIs(TestCase):
43	m2.load_state_dict(sd)	43	m2.load_state_dict(sd)
44	m2.to(device)	44	m2.to(device)
45	torch.npu.synchronize()	45	torch.npu.synchronize()
46	-	46	+
47	for (n1, p1), (n2, p2) in zip(m.named_parameters(), m2.named_parameters()):	47	for (n1, p1), (n2, p2) in zip(m.named_parameters(), m2.named_parameters()):
48	self.assertTrue(torch.allclose(p1, p2))	48	self.assertTrue(torch.allclose(p1, p2))
49		49
50	def test_moduledict_operations(self):	50	def test_moduledict_operations(self):
51	"""验证ModuleDict增删/索引/遍历"""	51	"""验证ModuleDict增删/索引/遍历"""
52	m = nn.ModuleDict({"a": nn.Linear(10, 20).to(device)})	52	m = nn.ModuleDict({"a": nn.Linear(10, 20).to(device)})
53	-	53	+
54	self.assertIn("a", m)	54	self.assertIn("a", m)
55	m["b"] = nn.Linear(20, 30).to(device)	55	m["b"] = nn.Linear(20, 30).to(device)
56	self.assertEqual(m["b"].weight.device, device)	56	self.assertEqual(m["b"].weight.device, device)
57	del m["b"]	57	del m["b"]
58	self.assertNotIn("b", m)	58	self.assertNotIn("b", m)
59	-	59	+
60	for sub in m.values():	60	for sub in m.values():
61	for p in sub.parameters():	61	for p in sub.parameters():
62	self.assertEqual(p.device, device)	62	self.assertEqual(p.device, device)
	@@ -64,15 +64,15 @@ class TestNPUModuleAPIs(TestCase):
64	def test_modulelist_operations(self):	64	def test_modulelist_operations(self):
65	"""验证ModuleList索引/新增/删除/遍历"""	65	"""验证ModuleList索引/新增/删除/遍历"""
66	m = nn.ModuleList([nn.Linear(20, 30).to(device), nn.BatchNorm1d(30).to(device)])	66	m = nn.ModuleList([nn.Linear(20, 30).to(device), nn.BatchNorm1d(30).to(device)])
67	-	67	+
68	self.assertEqual(m[0].weight.device, device)	68	self.assertEqual(m[0].weight.device, device)
69	self.assertEqual(m[1].running_mean.device, device)	69	self.assertEqual(m[1].running_mean.device, device)
70	-	70	+
71	m.append(nn.Linear(30, 40).to(device))	71	m.append(nn.Linear(30, 40).to(device))
72	m.insert(0, nn.Linear(10, 20).to(device))	72	m.insert(0, nn.Linear(10, 20).to(device))
73	m.pop(-1) # 修复：指定索引	73	m.pop(-1) # 修复：指定索引
74	m.pop(0) # 修复：指定索引	74	m.pop(0) # 修复：指定索引
75	-	75	+
76	self.assertEqual(len(m), 2)	76	self.assertEqual(len(m), 2)
77	for sub in m:	77	for sub in m:
78	for p in sub.parameters():	78	for p in sub.parameters():

Mtest/nn/test_nonlinear_activation_functions.py+1-1

	@@ -122,7 +122,7 @@ class TestNonLiACFunctions(TestCase):
122	npu_output = F.glu(npu_input)	122	npu_output = F.glu(npu_input)
123		123
124	self.assertRtolEqual(cpu_output.numpy(), npu_output.cpu().numpy())	124	self.assertRtolEqual(cpu_output.numpy(), npu_output.cpu().numpy())
125	-	125	+
126	@unittest.skip("skip test_gelu now")	126	@unittest.skip("skip test_gelu now")
127	def test_gelu(self):	127	def test_gelu(self):
128	input1 = torch.randn(2)	128	input1 = torch.randn(2)

Mtest/nn/test_recurrent_layers.py+1-1

	@@ -30,7 +30,7 @@ class TestRecurrentLayers(TestCase):
30	rnn = nn.GRU(10, 20, 2).npu()	30	rnn = nn.GRU(10, 20, 2).npu()
31	output, hn = rnn(input1, h0)	31	output, hn = rnn(input1, h0)
32	self.assertEqual(output is not None, True)	32	self.assertEqual(output is not None, True)
33	-	33	+
34	@unittest.skip("Temporarily skipping")	34	@unittest.skip("Temporarily skipping")
35	def test_RNNCell(self):	35	def test_RNNCell(self):
36	input1 = torch.randn(6, 3, 10).npu()	36	input1 = torch.randn(6, 3, 10).npu()

Mtest/nn/test_uninitialized_parameter_cls_to_become.py+2-2

	@@ -9,13 +9,13 @@ torch_npu.npu.set_compile_mode(jit_compile=False)
9	# 修复：将自定义属性设为类属性（确保实例化后必存在）	9	# 修复：将自定义属性设为类属性（确保实例化后必存在）
10	class CustomParameter(torch.nn.Parameter):	10	class CustomParameter(torch.nn.Parameter):
11	custom_attr = "custom_param" # 类属性，所有实例共享，无需__init__赋值	11	custom_attr = "custom_param" # 类属性，所有实例共享，无需__init__赋值
12	-	12	+
13	def __init__(self, data=None, requires_grad=True):	13	def __init__(self, data=None, requires_grad=True):
14	super().__init__(data, requires_grad)	14	super().__init__(data, requires_grad)
15		15
16		16
17	class TestUninitializedParameterClsToBecome(TestCase):	17	class TestUninitializedParameterClsToBecome(TestCase):
18	-	18	+
19	def test_core_functionality_npu(self):	19	def test_core_functionality_npu(self):
20	"""极简验证NPU环境下cls_to_become+materialize核心功能"""	20	"""极简验证NPU环境下cls_to_become+materialize核心功能"""
21	# 1. 创建NPU未初始化参数	21	# 1. 创建NPU未初始化参数

Mtest/nn/test_vision_functions.py+1-1

	@@ -67,7 +67,7 @@ class TestVisionFunctions(TestCase):
67		67
68	def test_affine_grid(self):	68	def test_affine_grid(self):
69	'''	69	'''
70	- Because of the limitation of NPU op, the NPU op will automatically convert the input	70	+ Because of the limitation of NPU op, the NPU op will automatically convert the input
71	fp32 to fp16 for calculation, so the input must be passed data within the representable	71	fp32 to fp16 for calculation, so the input must be passed data within the representable
72	range of fp16.	72	range of fp16.
73	'''	73	'''

Mtest/npu/test_aclgraph_launch_host_func.py+1-1

	@@ -38,7 +38,7 @@ class TestAclgraphLaunchHostFunc(TestCase):
38		38
39	self.capture_stream = torch_npu.npu.Stream()	39	self.capture_stream = torch_npu.npu.Stream()
40	self.graph = torch_npu.npu.NPUGraph()	40	self.graph = torch_npu.npu.NPUGraph()
41	-	41	+
42	torch_npu.npu._subscribe_report(self.capture_stream)	42	torch_npu.npu._subscribe_report(self.capture_stream)
43	a = torch.randn([5, 5]).npu()	43	a = torch.randn([5, 5]).npu()
44	b = torch.randn([5, 5]).npu()	44	b = torch.randn([5, 5]).npu()

Mtest/npu/test_aclgraph_support_blocking.py+3-3

Mtest/npu/test_aclgraph_update.py+14-14

Mtest/npu/test_allocator_multi_thread_prof.py+1-1

	@@ -12,7 +12,7 @@ class RandomDataset(Dataset):
12		12
13	def __len__(self):	13	def __len__(self):
14	return self.len	14	return self.len
15	-	15	+
16	def __getitem__(self, index):	16	def __getitem__(self, index):
17	return self.data[index].clone()	17	return self.data[index].clone()
18		18

Mtest/npu/test_cann_version.py+2-2

	@@ -24,10 +24,10 @@ class TestCANNversion(TestCase):
24	self.assertTrue(is_match, f"The env version is {version_env}. The format of cann version {version} is invalid.")	24	self.assertTrue(is_match, f"The env version is {version_env}. The format of cann version {version} is invalid.")
25	else:	25	else:
26	self.assertTrue(version == "", "When verssion_env < '8.1.RC1', the result of get_cann_version is not right.")	26	self.assertTrue(version == "", "When verssion_env < '8.1.RC1', the result of get_cann_version is not right.")
27	-	27	+
28	version = get_cann_version(module="CAN")	28	version = get_cann_version(module="CAN")
29	self.assertTrue(version == "", "When module is invalid, the result of get_cann_version is not right.")	29	self.assertTrue(version == "", "When module is invalid, the result of get_cann_version is not right.")
30	-	30	+
31	def test_get_driver_version(self):	31	def test_get_driver_version(self):
32	try:	32	try:
33	version = get_cann_version(module="DRIVER")	33	version = get_cann_version(module="DRIVER")

Mtest/npu/test_compatibility.py+5-5

	@@ -12,7 +12,7 @@ import pkgutil
12	import torch	12	import torch
13	from torch.testing._internal.common_utils import TestCase, run_tests	13	from torch.testing._internal.common_utils import TestCase, run_tests
14	from torch._utils_internal import get_file_path_2	14	from torch._utils_internal import get_file_path_2
15	-import torch_npu	15	+import torch_npu
16		16
17		17
18	NOT_IMPORT_LIST = [	18	NOT_IMPORT_LIST = [
	@@ -101,7 +101,7 @@ def is_not_compatibility(base_str, new_str, api_str=None):
101	# case: delete/different default value/different parameter name/different parameter dtype	101	# case: delete/different default value/different parameter name/different parameter dtype
102	if base_diff_params:	102	if base_diff_params:
103	return True	103	return True
104	-	104	+
105	# case: add params	105	# case: add params
106	new_diff_params = set(new_params) - set(base_params)	106	new_diff_params = set(new_params) - set(base_params)
107	# special case	107	# special case
	@@ -260,7 +260,7 @@ class TestPublicApiCompatibility(TestCase):
260	for key, value in base_schema0.items():	260	for key, value in base_schema0.items():
261	if not key.startswith("torch_c_func:") and not key.startswith("torch_npu_public_env:"):	261	if not key.startswith("torch_c_func:") and not key.startswith("torch_npu_public_env:"):
262	base_schema[key] = value	262	base_schema[key] = value
263	-	263	+
264	# load torchair torch_npu_schema.json	264	# load torchair torch_npu_schema.json
265	torchair_schema = {}	265	torchair_schema = {}
266	try:	266	try:
	@@ -270,7 +270,7 @@ class TestPublicApiCompatibility(TestCase):
270	except Exception:	270	except Exception:
271	warnings.warn(	271	warnings.warn(
272	"if you are debugging UT file in clone repo, please recursively update the torchair submodule")	272	"if you are debugging UT file in clone repo, please recursively update the torchair submodule")
273	-	273	+
274	if torchair_schema:	274	if torchair_schema:
275	base_schema.update(torchair_schema)	275	base_schema.update(torchair_schema)
276		276
	@@ -350,7 +350,7 @@ class TestPublicApiCompatibility(TestCase):
350	for func in deleted_apis:	350	for func in deleted_apis:
351	failure_list.append(f"# {func}:")	351	failure_list.append(f"# {func}:")
352	failure_list.append(f" - {func} has been deleted.")	352	failure_list.append(f" - {func} has been deleted.")
353	-	353	+
354	newly_apis = set(now_funcs) - set(base_funcs)	354	newly_apis = set(now_funcs) - set(base_funcs)
355	for func in newly_apis:	355	for func in newly_apis:
356	failure_list.append(f"# {func}:")	356	failure_list.append(f"# {func}:")

Mtest/npu/test_copy.py+15-15

	@@ -14,7 +14,7 @@ class TestCopyKernelMemoryFormat(TestCase):
14	shape_format = [	14	shape_format = [
15	[dtype, 2, shape] for dtype in dtype_list for shape in shape_list	15	[dtype, 2, shape] for dtype in dtype_list for shape in shape_list
16	]	16	]
17	-	17	+
18	for item in shape_format:	18	for item in shape_format:
19	cpu_input, npu_input = create_common_tensor(item, -100, 100)	19	cpu_input, npu_input = create_common_tensor(item, -100, 100)
20	npu_input_copy = cpu_input.npu()	20	npu_input_copy = cpu_input.npu()
	@@ -26,7 +26,7 @@ class TestCopyKernelMemoryFormat(TestCase):
26	shape_format = [	26	shape_format = [
27	[dtype, 2, shape] for dtype in dtype_list for shape in shape_list	27	[dtype, 2, shape] for dtype in dtype_list for shape in shape_list
28	]	28	]
29	-	29	+
30	for item in shape_format:	30	for item in shape_format:
31	cpu_input, npu_input = create_common_tensor(item, -100, 100)	31	cpu_input, npu_input = create_common_tensor(item, -100, 100)
32	cpu_transposed = cpu_input.transpose(-1, -2)	32	cpu_transposed = cpu_input.transpose(-1, -2)
	@@ -40,7 +40,7 @@ class TestCopyKernelMemoryFormat(TestCase):
40	shape_format = [	40	shape_format = [
41	[dtype, 2, shape] for dtype in dtype_list for shape in shape_list	41	[dtype, 2, shape] for dtype in dtype_list for shape in shape_list
42	]	42	]
43	-	43	+
44	for item in shape_format:	44	for item in shape_format:
45	cpu_input, npu_input = create_common_tensor(item, -100, 100)	45	cpu_input, npu_input = create_common_tensor(item, -100, 100)
46	cpu_output = npu_input.cpu()	46	cpu_output = npu_input.cpu()
	@@ -52,7 +52,7 @@ class TestCopyKernelMemoryFormat(TestCase):
52	shape_format = [	52	shape_format = [
53	[dtype, 2, shape] for dtype in dtype_list for shape in shape_list	53	[dtype, 2, shape] for dtype in dtype_list for shape in shape_list
54	]	54	]
55	-	55	+
56	for item in shape_format:	56	for item in shape_format:
57	cpu_input, npu_input = create_common_tensor(item, -100, 100)	57	cpu_input, npu_input = create_common_tensor(item, -100, 100)
58	npu_transposed = npu_input.transpose(-1, -2)	58	npu_transposed = npu_input.transpose(-1, -2)
	@@ -63,14 +63,14 @@ class TestCopyKernelMemoryFormat(TestCase):
63	src_dtype_list = [np.float32, np.float16]	63	src_dtype_list = [np.float32, np.float16]
64	dst_dtype_list = [torch.float16, torch.float32]	64	dst_dtype_list = [torch.float16, torch.float32]
65	shape = [32, 64]	65	shape = [32, 64]
66	-	66	+
67	for src_dtype, dst_dtype in zip(src_dtype_list, dst_dtype_list):	67	for src_dtype, dst_dtype in zip(src_dtype_list, dst_dtype_list):
68	cpu_input = torch.randn(shape, dtype=torch.float32) * 100	68	cpu_input = torch.randn(shape, dtype=torch.float32) * 100
69	cpu_input = cpu_input.to(torch.from_numpy(np.array([])).dtype if src_dtype == np.float32 else torch.float16)	69	cpu_input = cpu_input.to(torch.from_numpy(np.array([])).dtype if src_dtype == np.float32 else torch.float16)
70	-	70	+
71	npu_input = cpu_input.npu()	71	npu_input = cpu_input.npu()
72	npu_output = npu_input.to(dst_dtype)	72	npu_output = npu_input.to(dst_dtype)
73	-	73	+
74	cpu_output = cpu_input.to(dst_dtype)	74	cpu_output = cpu_input.to(dst_dtype)
75	self.assertRtolEqual(npu_output.cpu().numpy(), cpu_output.numpy())	75	self.assertRtolEqual(npu_output.cpu().numpy(), cpu_output.numpy())
76		76
	@@ -80,21 +80,21 @@ class TestCopyKernelMemoryFormat(TestCase):
80	(np.float32, torch.float16),	80	(np.float32, torch.float16),
81	]	81	]
82	shape = [32, 64]	82	shape = [32, 64]
83	-	83	+
84	for src_dtype, dst_dtype in dtype_pairs:	84	for src_dtype, dst_dtype in dtype_pairs:
85	cpu_input, npu_input = create_common_tensor([src_dtype, 0, shape], -100, 100)	85	cpu_input, npu_input = create_common_tensor([src_dtype, 0, shape], -100, 100)
86	cpu_output = npu_input.cpu().to(dst_dtype)	86	cpu_output = npu_input.cpu().to(dst_dtype)
87	-	87	+
88	expected = cpu_input.to(dst_dtype)	88	expected = cpu_input.to(dst_dtype)
89	self.assertRtolEqual(cpu_output.numpy(), expected.numpy())	89	self.assertRtolEqual(cpu_output.numpy(), expected.numpy())
90		90
91	def test_h2d_copy_slice_tensor(self):	91	def test_h2d_copy_slice_tensor(self):
92	shape = [64, 128]	92	shape = [64, 128]
93	cpu_input = torch.randn(shape)	93	cpu_input = torch.randn(shape)
94	-	94	+
95	cpu_slice = cpu_input[10:30, 20:60]	95	cpu_slice = cpu_input[10:30, 20:60]
96	npu_slice = cpu_slice.npu()	96	npu_slice = cpu_slice.npu()
97	-	97	+
98	npu_contiguous = npu_slice.contiguous()	98	npu_contiguous = npu_slice.contiguous()
99	self.assertRtolEqual(npu_contiguous.cpu().numpy(), cpu_slice.contiguous().numpy())	99	self.assertRtolEqual(npu_contiguous.cpu().numpy(), cpu_slice.contiguous().numpy())
100		100
	@@ -102,7 +102,7 @@ class TestCopyKernelMemoryFormat(TestCase):
102	shape = [64, 128]	102	shape = [64, 128]
103	cpu_input = torch.randn(shape)	103	cpu_input = torch.randn(shape)
104	npu_input = cpu_input.npu()	104	npu_input = cpu_input.npu()
105	-	105	+
106	npu_slice = npu_input[10:30, 20:60]	106	npu_slice = npu_input[10:30, 20:60]
107	cpu_slice = npu_slice.cpu()	107	cpu_slice = npu_slice.cpu()
108	self.assertRtolEqual(cpu_slice.numpy(), cpu_input[10:30, 20:60].contiguous().numpy())	108	self.assertRtolEqual(cpu_slice.numpy(), cpu_input[10:30, 20:60].contiguous().numpy())
	@@ -111,7 +111,7 @@ class TestCopyKernelMemoryFormat(TestCase):
111	shape = [1, 64, 1]	111	shape = [1, 64, 1]
112	cpu_input = torch.randn(shape)	112	cpu_input = torch.randn(shape)
113	npu_input = cpu_input.npu()	113	npu_input = cpu_input.npu()
114	-	114	+
115	npu_broadcast = npu_input.expand(4, 64, 128)	115	npu_broadcast = npu_input.expand(4, 64, 128)
116	cpu_output = npu_broadcast.cpu()	116	cpu_output = npu_broadcast.cpu()
117	self.assertRtolEqual(cpu_output.numpy(), cpu_input.expand(4, 64, 128).contiguous().numpy())	117	self.assertRtolEqual(cpu_output.numpy(), cpu_input.expand(4, 64, 128).contiguous().numpy())
	@@ -128,8 +128,8 @@ class TestCopyKernelMemoryFormat(TestCase):
128	shape = [32, 64, 128]	128	shape = [32, 64, 128]
129	cpu_input = torch.randn(shape)	129	cpu_input = torch.randn(shape)
130	npu_input = cpu_input.npu()	130	npu_input = cpu_input.npu()
131	-	131	+
132	- npu_permuted = npu_input.permute(2, 0, 1)	132	+ npu_permuted = npu_input.permute(2, 0, 1)
133	cpu_output = npu_permuted.cpu()	133	cpu_output = npu_permuted.cpu()
134	self.assertRtolEqual(cpu_output.numpy(), cpu_input.permute(2, 0, 1).contiguous().numpy())	134	self.assertRtolEqual(cpu_output.numpy(), cpu_input.permute(2, 0, 1).contiguous().numpy())
135		135

Mtest/npu/test_dlpack.py+17-17

Mtest/npu/test_errcode.py+2-2

	@@ -10,13 +10,13 @@ class TestErrorCode(TestCase):
10	def test_set_per_process_memory_fraction(self):	10	def test_set_per_process_memory_fraction(self):
11	with self.assertRaisesRegex(TypeError, "ERR00002 PTA invalid type"):	11	with self.assertRaisesRegex(TypeError, "ERR00002 PTA invalid type"):
12	torch_npu.npu.set_per_process_memory_fraction(1)	12	torch_npu.npu.set_per_process_memory_fraction(1)
13	-	13	+
14	def test_div(self):	14	def test_div(self):
15	x1 = torch.tensor(1).npu()	15	x1 = torch.tensor(1).npu()
16	x2 = torch.tensor(1).npu()	16	x2 = torch.tensor(1).npu()
17	with self.assertRaisesRegex(RuntimeError, "ERR01001 OPS invalid parameter"):	17	with self.assertRaisesRegex(RuntimeError, "ERR01001 OPS invalid parameter"):
18	torch.div(x1, x2, rounding_mode="test")	18	torch.div(x1, x2, rounding_mode="test")
19	-	19	+
20		20
21	if __name__ == "__main__":	21	if __name__ == "__main__":
22	run_tests()	22	run_tests()

Mtest/npu/test_graph_tree.py+18-18

Mtest/npu/test_mstx.py+1-1

	@@ -56,7 +56,7 @@ class TestMstx(TestCase):
56	self.assertEqual("", self.mark_domain)	56	self.assertEqual("", self.mark_domain)
57	torch_npu.npu.mstx.mark("test", stream=1, domain="test")	57	torch_npu.npu.mstx.mark("test", stream=1, domain="test")
58	self.assertEqual("", self.mark_msg)	58	self.assertEqual("", self.mark_msg)
59	- self.assertEqual("", self.mark_domain)	59	+ self.assertEqual("", self.mark_domain)
60		60
61	# valid inputs	61	# valid inputs
62	torch_npu.npu.mstx.mark("test1")	62	torch_npu.npu.mstx.mark("test1")

Mtest/npu/test_multi_devices_single_process.py+2-2

	@@ -102,7 +102,7 @@ class TestOp(TestCase):
102	output = torch.abs(input1)	102	output = torch.abs(input1)
103	output = output.cpu().numpy()	103	output = output.cpu().numpy()
104	return output	104	return output
105	-	105	+
106	def _test_abs(self, device="npu:1"):	106	def _test_abs(self, device="npu:1"):
107	torch.npu.set_device(0)	107	torch.npu.set_device(0)
108	cpu_input = torch.Tensor([1, -2, -10])	108	cpu_input = torch.Tensor([1, -2, -10])
	@@ -161,7 +161,7 @@ class TestOp(TestCase):
161	scale = 1 / 0.0078125	161	scale = 1 / 0.0078125
162	return torch_npu.npu_prompt_flash_attention(	162	return torch_npu.npu_prompt_flash_attention(
163	query, key, value, num_heads=32, input_layout="BNSD", scale_value=scale, pre_tokens=65535, next_tokens=65535, sparse_mode=0)	163	query, key, value, num_heads=32, input_layout="BNSD", scale_value=scale, pre_tokens=65535, next_tokens=65535, sparse_mode=0)
164	-	164	+
165	@SupportedDevices(['Ascend910B'])	165	@SupportedDevices(['Ascend910B'])
166	def _test_npu_prompt_flash_attention(self, device="npu:1"):	166	def _test_npu_prompt_flash_attention(self, device="npu:1"):
167	torch.npu.set_device(0)	167	torch.npu.set_device(0)

Mtest/npu/test_multi_stream_lazy_reclaim.py+7-7

	@@ -25,7 +25,7 @@ def extract_aclrtQueryEventStatus_count(prof_dir):
25		25
26	Args:	26	Args:
27	prof_dir: str, path to the profiling result directory	27	prof_dir: str, path to the profiling result directory
28	-	28	+
29	Returns:	29	Returns:
30	count: int, call count of aclrtQueryEventStatus, 0 if not found	30	count: int, call count of aclrtQueryEventStatus, 0 if not found
31	"""	31	"""
	@@ -166,12 +166,12 @@ class TestMultiStreamLazyReclaim(TestCase):
166	- lazy reclaim mode: Only queries in the following cases:	166	- lazy reclaim mode: Only queries in the following cases:
167	1. No available memory block found (!block_found)	167	1. No available memory block found (!block_found)
168	2. Event queue exceeds threshold kLazyQuerySize (512)	168	2. Event queue exceeds threshold kLazyQuerySize (512)
169	-	169	+
170	Test Method:	170	Test Method:
171	Use multiprocessing to test in two separate processes:	171	Use multiprocessing to test in two separate processes:
172	- Process 1: Enable multi_stream_lazy_reclaim	172	- Process 1: Enable multi_stream_lazy_reclaim
173	- Process 2: Disable multi_stream_lazy_reclaim	173	- Process 2: Disable multi_stream_lazy_reclaim
174	-	174	+
175	Each process sets environment variables independently to ensure configuration takes effect.	175	Each process sets environment variables independently to ensure configuration takes effect.
176	"""	176	"""
177		177
	@@ -197,7 +197,7 @@ class TestMultiStreamLazyReclaim(TestCase):
197	)	197	)
198		198
199	process.start()	199	process.start()
200	- process.join(timeout=300) #	200	+ process.join(timeout=300) #
201		201
202	if process.is_alive():	202	if process.is_alive():
203	process.terminate()	203	process.terminate()
	@@ -214,7 +214,7 @@ class TestMultiStreamLazyReclaim(TestCase):
214		214
215	eager_counts = results["eager"]	215	eager_counts = results["eager"]
216	lazy_counts = results["lazy"]	216	lazy_counts = results["lazy"]
217	-	217	+
218	# Output comparison results	218	# Output comparison results
219	print(f"\n========== Event Query Count Comparison ==========")	219	print(f"\n========== Event Query Count Comparison ==========")
220	print(f"Eager reclaim (multi_stream_lazy_reclaim:False): {eager_counts}")	220	print(f"Eager reclaim (multi_stream_lazy_reclaim:False): {eager_counts}")
	@@ -223,8 +223,8 @@ class TestMultiStreamLazyReclaim(TestCase):
223	# Core validation: aclrtQueryEventStatus call count in lazy mode must be less than eager mode	223	# Core validation: aclrtQueryEventStatus call count in lazy mode must be less than eager mode
224	# This is direct evidence that multi_stream_lazy_reclaim feature is working	224	# This is direct evidence that multi_stream_lazy_reclaim feature is working
225	self.assertLessEqual(	225	self.assertLessEqual(
226	- lazy_counts,	226	+ lazy_counts,
227	- eager_counts,	227	+ eager_counts,
228	f"Lazy reclaim mode should reduce event queries. "	228	f"Lazy reclaim mode should reduce event queries. "
229	f"Eager: {eager_counts}, Lazy: {lazy_counts}. "	229	f"Eager: {eager_counts}, Lazy: {lazy_counts}. "
230	f"If lazy >= eager, the optimization may not be working."	230	f"If lazy >= eager, the optimization may not be working."

Mtest/npu/test_npu_format.py+2-2

	@@ -5,7 +5,7 @@ from torch_npu._C import _weak_ref_tensor
5		5
6		6
7	class TestNPUFormat(TestCase):	7	class TestNPUFormat(TestCase):
8	-	8	+
9	def test_enum_values(self):	9	def test_enum_values(self):
10	"""test the enumeration value"""	10	"""test the enumeration value"""
11	self.assertEqual(torch_npu.Format.NCHW.value, 0)	11	self.assertEqual(torch_npu.Format.NCHW.value, 0)
	@@ -24,7 +24,7 @@ class TestNPUFormat(TestCase):
24	self.assertEqual(fmt2, torch_npu.Format.NHWC)	24	self.assertEqual(fmt2, torch_npu.Format.NHWC)
25		25
26	torch_npu.npu.config.allow_internal_format = True	26	torch_npu.npu.config.allow_internal_format = True
27	-	27	+
28	out3 = torch_npu.npu_format_cast(tensor, torch_npu.Format.FRACTAL_NZ)	28	out3 = torch_npu.npu_format_cast(tensor, torch_npu.Format.FRACTAL_NZ)
29	fmt3 = torch_npu.get_npu_format(out3)	29	fmt3 = torch_npu.get_npu_format(out3)
30	self.assertEqual(fmt3, torch_npu.Format.FRACTAL_NZ)	30	self.assertEqual(fmt3, torch_npu.Format.FRACTAL_NZ)

Mtest/npu/test_resnet.py+2-2

	@@ -167,7 +167,7 @@ def train(model, criterion, optimizer, epoch):
167	num += 1	167	num += 1
168	yield (torch.randn([128, 3, 224, 224]).npu() + torch.randint(-2, 2, [128, 3, 224, 224]).npu()).cpu(), \	168	yield (torch.randn([128, 3, 224, 224]).npu() + torch.randint(-2, 2, [128, 3, 224, 224]).npu()).cpu(), \
169	torch.randint(1, 1000, [128])	169	torch.randint(1, 1000, [128])
170	-	170	+
171	# switch to train mode	171	# switch to train mode
172	model.train()	172	model.train()
173		173
	@@ -217,7 +217,7 @@ def validate(model, criterion):
217	1,	217	1,
218	[batch_time, losses, top1, top5],	218	[batch_time, losses, top1, top5],
219	prefix='Test: ')	219	prefix='Test: ')
220	-	220	+
221	def fake_val_data(num):	221	def fake_val_data(num):
222	while num < 5:	222	while num < 5:
223	num += 1	223	num += 1

Mtest/npu/test_sanitizer.py+1-1

	@@ -14,7 +14,7 @@ REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
14	PYTORCH_INSTALL_PATH = os.path.dirname(os.path.realpath(torch.__file__))	14	PYTORCH_INSTALL_PATH = os.path.dirname(os.path.realpath(torch.__file__))
15	PYTORCH_NPU_INSTALL_PATH = os.path.dirname(os.path.realpath(torch_npu.__file__))	15	PYTORCH_NPU_INSTALL_PATH = os.path.dirname(os.path.realpath(torch_npu.__file__))
16		16
17	-	17	+
18	class TestSanitizer(TestCase):	18	class TestSanitizer(TestCase):
19	def tearDown(self):	19	def tearDown(self):
20	if sanitizer.npu_sanitizer.dispatch is not None:	20	if sanitizer.npu_sanitizer.dispatch is not None:

Mtest/npu/test_sanitizer_record_stream.py+1-1

	@@ -35,7 +35,7 @@ import os
35	import torch	35	import torch
36	import torch.cuda._sanitizer as csan	36	import torch.cuda._sanitizer as csan
37	import torch.distributed as dist	37	import torch.distributed as dist
38	-	38	+
39	import torch_npu	39	import torch_npu
40	from torch_npu.testing.testcase import TestCase, run_tests	40	from torch_npu.testing.testcase import TestCase, run_tests
41		41

Mtest/npu/test_save_async.py+4-4

	@@ -22,7 +22,7 @@ class TestAsyncSave(TestCase):
22	@classmethod	22	@classmethod
23	def tearDownClass(cls):	23	def tearDownClass(cls):
24	PathManager.remove_path_safety(TestAsyncSave.test_save_path)	24	PathManager.remove_path_safety(TestAsyncSave.test_save_path)
25	-	25	+
26	def wait_for_save_completion(self, file_path, timeout_sec=60, poll_interval_sec=0.5):	26	def wait_for_save_completion(self, file_path, timeout_sec=60, poll_interval_sec=0.5):
27	start_time = time.time()	27	start_time = time.time()
28		28
	@@ -43,13 +43,13 @@ class TestAsyncSave(TestCase):
43	save_tensor = torch.rand(1024, dtype=torch.float32).npu()	43	save_tensor = torch.rand(1024, dtype=torch.float32).npu()
44	async_save_path = os.path.join(TestAsyncSave.test_save_path, "async_save_tensor.pt")	44	async_save_path = os.path.join(TestAsyncSave.test_save_path, "async_save_tensor.pt")
45	torch_npu.utils.save_async(save_tensor, async_save_path)	45	torch_npu.utils.save_async(save_tensor, async_save_path)
46	-	46	+
47	if self.wait_for_save_completion(async_save_path):	47	if self.wait_for_save_completion(async_save_path):
48	tensor_async = torch.load(async_save_path, weights_only=False)	48	tensor_async = torch.load(async_save_path, weights_only=False)
49	self.assertEqual(tensor_async, save_tensor)	49	self.assertEqual(tensor_async, save_tensor)
50	else:	50	else:
51	self.assertTrue(False, f"{async_save_path} is not exist!")	51	self.assertTrue(False, f"{async_save_path} is not exist!")
52	-	52	+
53	def test_save_async(self):	53	def test_save_async(self):
54	loss1 = [1.6099495, 1.6099086, 1.6098710]	54	loss1 = [1.6099495, 1.6099086, 1.6098710]
55	loss2 = []	55	loss2 = []
	@@ -78,7 +78,7 @@ class TestAsyncSave(TestCase):
78	loss.backward()	78	loss.backward()
79		79
80	optimerizer.step()	80	optimerizer.step()
81	-	81	+
82	loss2.append(loss)	82	loss2.append(loss)
83	checkpoint = {	83	checkpoint = {
84	"model": model.state_dict(),	84	"model": model.state_dict(),

Mtest/npu/test_storage.py+1-1

	@@ -203,7 +203,7 @@ class TestStorage(TestCase):
203	self.assertEqual(cpu_res.size(), npu_res.size())	203	self.assertEqual(cpu_res.size(), npu_res.size())
204	self.assertEqual(cpu_res, npu_res.cpu())	204	self.assertEqual(cpu_res, npu_res.cpu())
205	self.assertEqual(cpu_res.tolist(), npu_res.cpu().tolist())	205	self.assertEqual(cpu_res.tolist(), npu_res.cpu().tolist())
206	-	206	+
207	@SupportedDevices(['Ascend910B'])	207	@SupportedDevices(['Ascend910B'])
208	def _test_datatype_cast_complex(cpu_storage, npu_storage):	208	def _test_datatype_cast_complex(cpu_storage, npu_storage):
209	dtypes = [	209	dtypes = [

Mtest/npu/test_swapped_memory_allocator.py+13-13

	@@ -39,13 +39,13 @@ class TestNPUSwappedMemoryAllocator(unittest.TestCase):
39	tensors = []	39	tensors = []
40	for i in range(5):	40	for i in range(5):
41	tensor = torch_npu.empty_with_swapped_memory(	41	tensor = torch_npu.empty_with_swapped_memory(
42	- [256 * (i + 1)],	42	+ [256 * (i + 1)],
43	- dtype=torch.float32,	43	+ dtype=torch.float32,
44	device='npu:0'	44	device='npu:0'
45	)	45	)
46	tensor.fill_(float(i))	46	tensor.fill_(float(i))
47	tensors.append(tensor)	47	tensors.append(tensor)
48	-	48	+
49	for t in tensors:	49	for t in tensors:
50	del t	50	del t
51	del tensors	51	del tensors
	@@ -57,34 +57,34 @@ class TestNPUSwappedMemoryAllocator(unittest.TestCase):
57	def test_02_async_operations_and_release(self):	57	def test_02_async_operations_and_release(self):
58	"""Test 2: Async operations followed by release (verify stream sync works)"""	58	"""Test 2: Async operations followed by release (verify stream sync works)"""
59	tensor = torch_npu.empty_with_swapped_memory(	59	tensor = torch_npu.empty_with_swapped_memory(
60	- [512, 512],	60	+ [512, 512],
61	- dtype=torch.float32,	61	+ dtype=torch.float32,
62	device='npu:0'	62	device='npu:0'
63	)	63	)
64	-	64	+
65	tensor.fill_(1.0)	65	tensor.fill_(1.0)
66	-	66	+
67	for i in range(10):	67	for i in range(10):
68	tensor = tensor * 1.1	68	tensor = tensor * 1.1
69	tensor = tensor + i * 0.1	69	tensor = tensor + i * 0.1
70	-	70	+
71	tensor.sqrt_()	71	tensor.sqrt_()
72	-	72	+
73	expected_tensor = torch.empty([512, 512], dtype=torch.float32, device='npu:0')	73	expected_tensor = torch.empty([512, 512], dtype=torch.float32, device='npu:0')
74	expected_tensor.fill_(1.0)	74	expected_tensor.fill_(1.0)
75	for i in range(10):	75	for i in range(10):
76	expected_tensor = expected_tensor * 1.1	76	expected_tensor = expected_tensor * 1.1
77	expected_tensor = expected_tensor + i * 0.1	77	expected_tensor = expected_tensor + i * 0.1
78	expected_tensor = expected_tensor.sqrt()	78	expected_tensor = expected_tensor.sqrt()
79	-	79	+
80	self.assertTrue(torch.allclose(tensor, expected_tensor, rtol=1e-5, atol=1e-5))	80	self.assertTrue(torch.allclose(tensor, expected_tensor, rtol=1e-5, atol=1e-5))
81	-	81	+
82	weak_ref = weakref.ref(tensor)	82	weak_ref = weakref.ref(tensor)
83	del tensor	83	del tensor
84	gc.collect()	84	gc.collect()
85	-	85	+
86	self.assertIsNone(weak_ref())	86	self.assertIsNone(weak_ref())
87	-	87	+
88	torch.npu.empty_cache()	88	torch.npu.empty_cache()
89		89
90		90

Mtest/npu/test_tensors.py+2-2

	@@ -163,10 +163,10 @@ class TestTensor(TestCase):
163	scalar_input = torch.randn(1).npu()	163	scalar_input = torch.randn(1).npu()
164	bool_scalar = scalar_input.to(torch.bool)	164	bool_scalar = scalar_input.to(torch.bool)
165	self.assertTrue(isinstance(bool_scalar.item(), bool))	165	self.assertTrue(isinstance(bool_scalar.item(), bool))
166	-	166	+
167	half_scalar = scalar_input.to(torch.float16)	167	half_scalar = scalar_input.to(torch.float16)
168	self.assertTrue(isinstance(half_scalar.item(), float))	168	self.assertTrue(isinstance(half_scalar.item(), float))
169	-	169	+
170	bf16_scalar = scalar_input.to(torch.bfloat16)	170	bf16_scalar = scalar_input.to(torch.bfloat16)
171	self.assertTrue(isinstance(bf16_scalar.item(), float))	171	self.assertTrue(isinstance(bf16_scalar.item(), float))
172		172

Mtest/npu/test_torch_npu.py+5-5

Mtest/onnx/onnx_test_common.py+1-1

	@@ -428,7 +428,7 @@ def add_decorate_info(
428	# Skip does not apply to this opset	428	# Skip does not apply to this opset
429	continue	429	continue
430	opinfo = ops_mapping.get((decorate_meta.op_name, decorate_meta.variant_name))	430	opinfo = ops_mapping.get((decorate_meta.op_name, decorate_meta.variant_name))
431	- assert opinfo is not None, (	431	+ assert opinfo is not None, (
432	f"Couldn't find OpInfo for {decorate_meta}. Did you need to specify variant_name?")	432	f"Couldn't find OpInfo for {decorate_meta}. Did you need to specify variant_name?")
433	assert decorate_meta.model_type is None, (	433	assert decorate_meta.model_type is None, (
434	f"Tested op: {decorate_meta.op_name} in wrong position! "	434	f"Tested op: {decorate_meta.op_name} in wrong position! "

Mtest/onnx/test_combined_onnx_ops.py+1-1

	@@ -222,7 +222,7 @@ class TestOnnxOps(TestCase):
222		222
223	torch.npu.config.allow_internal_format = True	223	torch.npu.config.allow_internal_format = True
224	torch.npu.set_compile_mode(jit_compile=True)	224	torch.npu.set_compile_mode(jit_compile=True)
225	-	225	+
226	def export_onnx(onnx_model_name):	226	def export_onnx(onnx_model_name):
227	input_ = torch.rand([1, 128, 4, 14, 14]).npu()	227	input_ = torch.rand([1, 128, 4, 14, 14]).npu()
228	model = Model().to("npu")	228	model = Model().to("npu")

Mtest/onnx/test_wrapper_onnx_ops.py+15-15

	@@ -194,7 +194,7 @@ class TestOnnxOps(TestCase):
194	class Model(torch.nn.Module):	194	class Model(torch.nn.Module):
195	def __init__(self):	195	def __init__(self):
196	super(Model, self).__init__()	196	super(Model, self).__init__()
197	-	197	+
198	def forward(self, x):	198	def forward(self, x):
199	return torch_npu.npu_geglu(x)	199	return torch_npu.npu_geglu(x)
200		200
	@@ -203,7 +203,7 @@ class TestOnnxOps(TestCase):
203	model = Model().to("npu")	203	model = Model().to("npu")
204	model(x)	204	model(x)
205	self.onnx_export(model, x, onnx_model_name, ["input"], ["output1", "output2"])	205	self.onnx_export(model, x, onnx_model_name, ["input"], ["output1", "output2"])
206	-	206	+
207	onnx_model_name = "model_npu_geglu.onnx"	207	onnx_model_name = "model_npu_geglu.onnx"
208	export_onnx(onnx_model_name)	208	export_onnx(onnx_model_name)
209	assert(os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path,	209	assert(os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path,
	@@ -1156,7 +1156,7 @@ class TestOnnxOps(TestCase):
1156		1156
1157	def forward(self, sorted_experts):	1157	def forward(self, sorted_experts):
1158	return torch_npu.npu_moe_compute_expert_tokens(sorted_experts=5)	1158	return torch_npu.npu_moe_compute_expert_tokens(sorted_experts=5)
1159	-	1159	+
1160	def export_onnx(onnx_model_name):	1160	def export_onnx(onnx_model_name):
1161	data = list(range(20))	1161	data = list(range(20))
1162	experts = torch.tensor(data, dtype=torch.int32).npu()	1162	experts = torch.tensor(data, dtype=torch.int32).npu()
	@@ -1167,7 +1167,7 @@ class TestOnnxOps(TestCase):
1167	onnx_model_name = "model_moe_compute_expert_tokens.onnx"	1167	onnx_model_name = "model_moe_compute_expert_tokens.onnx"
1168	export_onnx(onnx_model_name)	1168	export_onnx(onnx_model_name)
1169	assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path,	1169	assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path,
1170	- onnx_model_name)))	1170	+ onnx_model_name)))
1171		1171
1172	@unittest.skip("skip now")	1172	@unittest.skip("skip now")
1173	def test_wrapper_npu_mish(self):	1173	def test_wrapper_npu_mish(self):
	@@ -1226,7 +1226,7 @@ class TestOnnxOps(TestCase):
1226	epsilon = 1e-6	1226	epsilon = 1e-6
1227	x = torch_npu.npu_rms_norm(x, gamma, epsilon)	1227	x = torch_npu.npu_rms_norm(x, gamma, epsilon)
1228	return x	1228	return x
1229	-	1229	+
1230	def export_onnx(onnx_model_name):	1230	def export_onnx(onnx_model_name):
1231	x = torch.rand(10, 1024).uniform_(-3, 3).npu().half()	1231	x = torch.rand(10, 1024).uniform_(-3, 3).npu().half()
1232	gamma = torch.rand(1024).uniform_(-3, 3).npu().half()	1232	gamma = torch.rand(1024).uniform_(-3, 3).npu().half()
	@@ -1249,7 +1249,7 @@ class TestOnnxOps(TestCase):
1249	epsilon = 1e-6	1249	epsilon = 1e-6
1250	x = torch_npu.npu_add_rms_norm(x1, x2, gamma, epsilon)	1250	x = torch_npu.npu_add_rms_norm(x1, x2, gamma, epsilon)
1251	return x	1251	return x
1252	-	1252	+
1253	def export_onnx(onnx_model_name):	1253	def export_onnx(onnx_model_name):
1254	x1 = torch.rand(10, 1024).uniform_(-3, 3).npu().half()	1254	x1 = torch.rand(10, 1024).uniform_(-3, 3).npu().half()
1255	x2 = torch.rand(10, 1024).uniform_(-3, 3).npu().half()	1255	x2 = torch.rand(10, 1024).uniform_(-3, 3).npu().half()
	@@ -1318,7 +1318,7 @@ class TestOnnxOps(TestCase):
1318	def forward(self, input_dummy, smooth_scales_dummy):	1318	def forward(self, input_dummy, smooth_scales_dummy):
1319	output, scale = torch_npu.npu_dynamic_quant(input_dummy, smooth_scales=smooth_scales_dummy)	1319	output, scale = torch_npu.npu_dynamic_quant(input_dummy, smooth_scales=smooth_scales_dummy)
1320	return output, scale	1320	return output, scale
1321	-	1321	+
1322	def export_onnx(onnx_model_name):	1322	def export_onnx(onnx_model_name):
1323	input_dummy = torch.rand(4, 1024, 512).uniform_(-3, 3).npu().to(torch.float16)	1323	input_dummy = torch.rand(4, 1024, 512).uniform_(-3, 3).npu().to(torch.float16)
1324	smooth_scales_dummy = torch.rand(512).uniform_(-3, 3).npu().to(torch.float16)	1324	smooth_scales_dummy = torch.rand(512).uniform_(-3, 3).npu().to(torch.float16)
	@@ -1340,7 +1340,7 @@ class TestOnnxOps(TestCase):
1340	def forward(self, input_dummy, smooth_scales_dummy, group_index_dummy):	1340	def forward(self, input_dummy, smooth_scales_dummy, group_index_dummy):
1341	output, scale = torch_npu.npu_dynamic_quant(input_dummy, smooth_scales=smooth_scales_dummy, group_index=group_index_dummy)	1341	output, scale = torch_npu.npu_dynamic_quant(input_dummy, smooth_scales=smooth_scales_dummy, group_index=group_index_dummy)
1342	return output, scale	1342	return output, scale
1343	-	1343	+
1344	def export_onnx(onnx_model_name):	1344	def export_onnx(onnx_model_name):
1345	input_dummy = torch.rand(4, 1024, 512).uniform_(-3, 3).npu().to(torch.float16)	1345	input_dummy = torch.rand(4, 1024, 512).uniform_(-3, 3).npu().to(torch.float16)
1346	group_num = 10	1346	group_num = 10
	@@ -1370,7 +1370,7 @@ class TestOnnxOps(TestCase):
1370	def forward(self, input_dummy, smooth_scales_dummy, group_index_dummy):	1370	def forward(self, input_dummy, smooth_scales_dummy, group_index_dummy):
1371	output, scale, offset = torch_npu.npu_dynamic_quant_asymmetric(input_dummy, smooth_scales=smooth_scales_dummy, group_index=group_index_dummy)	1371	output, scale, offset = torch_npu.npu_dynamic_quant_asymmetric(input_dummy, smooth_scales=smooth_scales_dummy, group_index=group_index_dummy)
1372	return output, scale, offset	1372	return output, scale, offset
1373	-	1373	+
1374	def export_onnx(onnx_model_name):	1374	def export_onnx(onnx_model_name):
1375	input_dummy = torch.rand(4, 1024, 512).uniform_(-3, 3).npu().to(torch.float16)	1375	input_dummy = torch.rand(4, 1024, 512).uniform_(-3, 3).npu().to(torch.float16)
1376	group_num = 10	1376	group_num = 10
	@@ -1436,7 +1436,7 @@ class TestOnnxOps(TestCase):
1436	assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, onnx_model_name)))	1436	assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, onnx_model_name)))
1437		1437
1438	@unittest.skip("skip now")	1438	@unittest.skip("skip now")
1439	- def test_wrapper_npu_quantize(self):	1439	+ def test_wrapper_npu_quantize(self):
1440	class Model(torch.nn.Module):	1440	class Model(torch.nn.Module):
1441	def __init__(self):	1441	def __init__(self):
1442	super().__init__()	1442	super().__init__()
	@@ -1461,7 +1461,7 @@ class TestOnnxOps(TestCase):
1461		1461
1462	@unittest.skip("skip now")	1462	@unittest.skip("skip now")
1463	@SupportedDevices(['Ascend910B'])	1463	@SupportedDevices(['Ascend910B'])
1464	- def test_wrapper_npu_group_quant(self):	1464	+ def test_wrapper_npu_group_quant(self):
1465	class Model(torch.nn.Module):	1465	class Model(torch.nn.Module):
1466	def __init__(self):	1466	def __init__(self):
1467	super().__init__()	1467	super().__init__()
	@@ -1488,7 +1488,7 @@ class TestOnnxOps(TestCase):
1488		1488
1489	@unittest.skip("skip now")	1489	@unittest.skip("skip now")
1490	@SupportedDevices(['Ascend910B'])	1490	@SupportedDevices(['Ascend910B'])
1491	- def test_wrapper_npu_moe_finalize_routing(self):	1491	+ def test_wrapper_npu_moe_finalize_routing(self):
1492	class Model(torch.nn.Module):	1492	class Model(torch.nn.Module):
1493	def __init__(self):	1493	def __init__(self):
1494	super().__init__()	1494	super().__init__()
	@@ -1521,7 +1521,7 @@ class TestOnnxOps(TestCase):
1521		1521
1522	@unittest.skip("skip now")	1522	@unittest.skip("skip now")
1523	@SupportedDevices(['Ascend910B'])	1523	@SupportedDevices(['Ascend910B'])
1524	- def test_wrapper_npu_moe_finalize_routing_v2(self):	1524	+ def test_wrapper_npu_moe_finalize_routing_v2(self):
1525	class Model(torch.nn.Module):	1525	class Model(torch.nn.Module):
1526	def __init__(self):	1526	def __init__(self):
1527	super().__init__()	1527	super().__init__()
	@@ -1529,7 +1529,7 @@ class TestOnnxOps(TestCase):
1529	def forward(self, expanded_permuted_rows, skip1, skip2_optional, bias, scales,	1529	def forward(self, expanded_permuted_rows, skip1, skip2_optional, bias, scales,
1530	expanded_src_to_dst_row, expert_for_source_row):	1530	expanded_src_to_dst_row, expert_for_source_row):
1531	return torch_npu.npu_moe_finalize_routing(expanded_permuted_rows, skip1, skip2_optional,	1531	return torch_npu.npu_moe_finalize_routing(expanded_permuted_rows, skip1, skip2_optional,
1532	- bias, scales, expanded_src_to_dst_row,	1532	+ bias, scales, expanded_src_to_dst_row,
1533	expert_for_source_row, drop_pad_mode=1)	1533	expert_for_source_row, drop_pad_mode=1)
1534		1534
1535	def export_onnx(onnx_model_name):	1535	def export_onnx(onnx_model_name):
	@@ -1566,7 +1566,7 @@ class TestOnnxOps(TestCase):
1566	return y	1566	return y
1567		1567
1568	def export_onnx(onnx_model_name):	1568	def export_onnx(onnx_model_name):
1569	- x = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]],	1569	+ x = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]],
1570	dtype=torch.float16).npu()	1570	dtype=torch.float16).npu()
1571	model = Model().to("npu")	1571	model = Model().to("npu")
1572	model(x)	1572	model(x)

Mtest/optim/test_fused_optimizers.py+3-3

	@@ -103,7 +103,7 @@ class TestFusedOptim(TestCase):
103	if p.grad is not None:	103	if p.grad is not None:
104	self.assertEqual(p.grad, p_clone.grad)	104	self.assertEqual(p.grad, p_clone.grad)
105	self.assertEqual(p.grad, torch.zeros_like(p.grad))	105	self.assertEqual(p.grad, torch.zeros_like(p.grad))
106	-	106	+
107	def test_step(self):	107	def test_step(self):
108	optim_cases = self._create_optimizer_cases(all_cases=True)	108	optim_cases = self._create_optimizer_cases(all_cases=True)
109	num_iters = 10	109	num_iters = 10
	@@ -160,7 +160,7 @@ class TestFusedOptim(TestCase):
160	for p in m.parameters():	160	for p in m.parameters():
161	if p.grad is not None:	161	if p.grad is not None:
162	self.assertEqual(grads_before_unscale[p] / 128, p.grad)	162	self.assertEqual(grads_before_unscale[p] / 128, p.grad)
163	-	163	+
164	@unittest.skip("Temporarily skipping")	164	@unittest.skip("Temporarily skipping")
165	def test_simple_model_train_dynamic(self):	165	def test_simple_model_train_dynamic(self):
166	model = self._create_simple_model()	166	model = self._create_simple_model()
	@@ -191,7 +191,7 @@ class TestFusedOptim(TestCase):
191	scaler_fused.step(opt_fused)	191	scaler_fused.step(opt_fused)
192	scaler_fused.update()	192	scaler_fused.update()
193	self.assertRtolEqual(loss, loss_fused)	193	self.assertRtolEqual(loss, loss_fused)
194	-	194	+
195	@unittest.skip("Temporarily skipping")	195	@unittest.skip("Temporarily skipping")
196	def test_simple_model_train_static(self):	196	def test_simple_model_train_static(self):
197	model = self._create_simple_model()	197	model = self._create_simple_model()

Mtest/profiler/analysis/prof_bean/test_op_mark_bean.py+1-1

	@@ -12,7 +12,7 @@ class TestOpMarkBean(TestCase):
12	def setUpClass(cls):	12	def setUpClass(cls):
13	super().setUpClass()	13	super().setUpClass()
14	cls.samples = cls.generate_samples()	14	cls.samples = cls.generate_samples()
15	-	15	+
16		16
17	@classmethod	17	@classmethod
18	def generate_samples(cls):	18	def generate_samples(cls):

Mtest/profiler/analysis/prof_common_func/test_constant.py+1-1

	@@ -18,7 +18,7 @@ class TestConstant(TestCase):
18	with self.assertRaises((RuntimeError, TypeError)):	18	with self.assertRaises((RuntimeError, TypeError)):
19	convert_ns2us_float(18789.89)	19	convert_ns2us_float(18789.89)
20	self.assertEqual(str(convert_ns2us_float(1459635878536856678)), str(1459635878536856678 / 1000))	20	self.assertEqual(str(convert_ns2us_float(1459635878536856678)), str(1459635878536856678 / 1000))
21	-	21	+
22	def test_convert_ns2us_str(self):	22	def test_convert_ns2us_str(self):
23	self.assertEqual(convert_ns2us_str(float("inf")), "inf")	23	self.assertEqual(convert_ns2us_str(float("inf")), "inf")
24	with self.assertRaises((RuntimeError, TypeError)):	24	with self.assertRaises((RuntimeError, TypeError)):

Mtest/profiler/analysis/prof_common_func/test_trace_event_manager.py+1-1

	@@ -87,7 +87,7 @@ class TestTraceEventManager(TestCase):
87	"tid": 444, "ts": "4.000", "cat": "fwdbwd"}	87	"tid": 444, "ts": "4.000", "cat": "fwdbwd"}
88	]	88	]
89	self.assertEqual(expect, TraceEventManager.create_fwd_flow(events))	89	self.assertEqual(expect, TraceEventManager.create_fwd_flow(events))
90	-	90	+
91	def test_python_event(self):	91	def test_python_event(self):
92	process_id = random.randint(1, 2**64 - 1)	92	process_id = random.randint(1, 2**64 - 1)
93	thread_id = random.randint(1, 2**64 - 1)	93	thread_id = random.randint(1, 2**64 - 1)

Mtest/profiler/test_profiler_tree.py+0-1

	@@ -733,4 +733,3 @@ class TestProfilerTree(TestCase):
733		733
734	if __name__ == "__main__":	734	if __name__ == "__main__":
735	run_tests()	735	run_tests()
736	-

Mtest/sample_torch_npu_run_store.py+2-2

	@@ -16,7 +16,7 @@ class TorchNpuRunStoreSample:
16	self._init_method = f'parallel://{ip}:{port}'	16	self._init_method = f'parallel://{ip}:{port}'
17	self._timeout = timedelta(minutes=1)	17	self._timeout = timedelta(minutes=1)
18	self._key = 'sample_torch_npu_run_store:test_case_001'	18	self._key = 'sample_torch_npu_run_store:test_case_001'
19	-	19	+
20	rendezvous_iterator = rendezvous(	20	rendezvous_iterator = rendezvous(
21	self._init_method, self._current_rank, self._world_size, timeout=self._timeout	21	self._init_method, self._current_rank, self._world_size, timeout=self._timeout
22	)	22	)
	@@ -41,7 +41,7 @@ class TorchNpuRunStoreSample:
41	if timedelta(seconds=(time.time() - start_time)) > self._timeout:	41	if timedelta(seconds=(time.time() - start_time)) > self._timeout:
42	timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')	42	timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')
43	raise RuntimeError(f'[{timestamp}]rank: {self._current_rank} wait all workers ready timeout')	43	raise RuntimeError(f'[{timestamp}]rank: {self._current_rank} wait all workers ready timeout')
44	-	44	+
45	timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')	45	timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')
46	print(f'[{timestamp}] Rank: {self._current_rank} complete store-based barrier for worker count: {alive_count}')	46	print(f'[{timestamp}] Rank: {self._current_rank} complete store-based barrier for worker count: {alive_count}')
47		47

Mtest/test_autocast.py+2-2

	@@ -289,7 +289,7 @@ class TestAutocastNPUfp32(TestCase):
289	with torch.autocast(device_type=device, dtype=torch.float32):	289	with torch.autocast(device_type=device, dtype=torch.float32):
290	b = torch.mm(a, a)	290	b = torch.mm(a, a)
291	self.assertEqual(b.dtype, torch.float32)	291	self.assertEqual(b.dtype, torch.float32)
292	-	292	+
293	def test_autocast_fp32_when_origin_dtype_is_bfloat16(self):	293	def test_autocast_fp32_when_origin_dtype_is_bfloat16(self):
294	device = "npu"	294	device = "npu"
295	a = torch.rand((8, 8), device=device, dtype=torch.bfloat16)	295	a = torch.rand((8, 8), device=device, dtype=torch.bfloat16)
	@@ -303,7 +303,7 @@ class TestAutocastNPUfp32(TestCase):
303	with torch.autocast(device_type=device, dtype=torch.float32):	303	with torch.autocast(device_type=device, dtype=torch.float32):
304	b = torch.mm(a, a)	304	b = torch.mm(a, a)
305	self.assertEqual(b.dtype, torch.float32)	305	self.assertEqual(b.dtype, torch.float32)
306	-	306	+
307	def test_autocast_fp32_when_disabled(self):	307	def test_autocast_fp32_when_disabled(self):
308	device = "npu"	308	device = "npu"
309	a = torch.rand((8, 8), device=device, dtype=torch.bfloat16)	309	a = torch.rand((8, 8), device=device, dtype=torch.bfloat16)

Mtest/test_binary_ufuncs.py+20-20

Mtest/test_fake_tensor.py+1-1

	@@ -1903,7 +1903,7 @@ class TestQuantMatmul(TestCase):
1903	expect_ret_fp16 = torch.randint(-1, 1, (1, 1, 100), dtype=torch.float16).npu()	1903	expect_ret_fp16 = torch.randint(-1, 1, (1, 1, 100), dtype=torch.float16).npu()
1904	bias_fp32 = torch.randint(-1, 1, (1, 1, 100), dtype=torch.float32).npu()	1904	bias_fp32 = torch.randint(-1, 1, (1, 1, 100), dtype=torch.float32).npu()
1905	pertoken_scale = torch.randn(1, dtype=torch.float32).npu()	1905	pertoken_scale = torch.randn(1, dtype=torch.float32).npu()
1906	- res_fp16 = torch_npu.npu_quant_matmul(x1, x2, scale, offset=None, pertoken_scale=pertoken_scale,	1906	+ res_fp16 = torch_npu.npu_quant_matmul(x1, x2, scale, offset=None, pertoken_scale=pertoken_scale,
1907	bias=bias_fp32, output_dtype=torch.float16)	1907	bias=bias_fp32, output_dtype=torch.float16)
1908	self.assertTrue(expect_ret_fp16.shape == res_fp16.shape)	1908	self.assertTrue(expect_ret_fp16.shape == res_fp16.shape)
1909	self.assertTrue(expect_ret_fp16.dtype == res_fp16.dtype)	1909	self.assertTrue(expect_ret_fp16.dtype == res_fp16.dtype)

Mtest/test_indexing.py+2-2

	@@ -128,7 +128,7 @@ class TestIndexing(TestCase):
128		128
129	self.assertRaises(TypeError, delitem)	129	self.assertRaises(TypeError, delitem)
130		130
131	-	131	+
132	@dtypes(torch.half, torch.double)	132	@dtypes(torch.half, torch.double)
133	def test_advancedindex(self, device, dtype):	133	def test_advancedindex(self, device, dtype):
134	# Tests for Integer Array Indexing, Part I - Purely integer array	134	# Tests for Integer Array Indexing, Part I - Purely integer array
	@@ -876,7 +876,7 @@ class TestIndexing(TestCase):
876		876
877	self.assertEqual(output, input_list)	877	self.assertEqual(output, input_list)
878		878
879	-	879	+
880	def test_index_ind_dtype(self, device):	880	def test_index_ind_dtype(self, device):
881	x = torch.randn(4, 4, device=device)	881	x = torch.randn(4, 4, device=device)
882	ind_long = torch.randint(4, (4,), dtype=torch.long, device=device)	882	ind_long = torch.randint(4, (4,), dtype=torch.long, device=device)

Mtest/test_jit_fuser.py+1-1

	@@ -787,7 +787,7 @@ class TestFuser(JitTestCase):
787	FileCheck.check("FusionGroup").run(str(graph))	787	FileCheck.check("FusionGroup").run(str(graph))
788	except RuntimeError as e:	788	except RuntimeError as e:
789	if 'Failed to compile' in e.args[0]:	789	if 'Failed to compile' in e.args[0]:
790	- warnings.warn('CPU fuser test has failed! This is not a hard failure, '	790	+ warnings.warn('CPU fuser test has failed! This is not a hard failure, '
791	'because the kernels sometimes trigger bugs in compilers '	791	'because the kernels sometimes trigger bugs in compilers '
792	'(most notably GCC 7.2).')	792	'(most notably GCC 7.2).')
793	raise unittest.SkipTest('Failed to compile') from e	793	raise unittest.SkipTest('Failed to compile') from e

Mtest/test_multiprocessing_api.py+15-15

Mtest/test_nestedtensor.py+5-5

Mtest/test_npu_expandable_segments.py+10-10

	@@ -69,14 +69,14 @@ class TestPluggableAllocator(TestCase):
69	)	69	)
70	# Load the allocator	70	# Load the allocator
71	cls.new_alloc = torch_npu.npu.memory.NPUPluggableAllocator(os_path, 'my_malloc', 'my_free')	71	cls.new_alloc = torch_npu.npu.memory.NPUPluggableAllocator(os_path, 'my_malloc', 'my_free')
72	-	72	+
73	def test_pluggable_allocator(self):	73	def test_pluggable_allocator(self):
74	torch.npu.memory._set_allocator_settings("expandable_segments:False")	74	torch.npu.memory._set_allocator_settings("expandable_segments:False")
75	with torch.npu.use_mem_pool(torch.npu.MemPool(TestPluggableAllocator.new_alloc._allocator)):	75	with torch.npu.use_mem_pool(torch.npu.MemPool(TestPluggableAllocator.new_alloc._allocator)):
76	x = torch.empty((7500, 1024, 1024), device="npu")	76	x = torch.empty((7500, 1024, 1024), device="npu")
77	del x	77	del x
78	torch.npu.memory._set_allocator_settings("expandable_segments:True")	78	torch.npu.memory._set_allocator_settings("expandable_segments:True")
79	-	79	+
80	@staticmethod	80	@staticmethod
81	def conv_operation(x):	81	def conv_operation(x):
82	return TestPluggableAllocator.deconv(TestPluggableAllocator.conv(x) + 0.005)	82	return TestPluggableAllocator.deconv(TestPluggableAllocator.conv(x) + 0.005)
	@@ -131,28 +131,28 @@ class TestPluggableAllocator(TestCase):
131	with torch.npu.stream(stream1):	131	with torch.npu.stream(stream1):
132	x1 = self.conv_with_allocator(x1)	132	x1 = self.conv_with_allocator(x1)
133	events[0].record()	133	events[0].record()
134	-	134	+
135	with torch.npu.stream(stream2):	135	with torch.npu.stream(stream2):
136	events[0].wait(stream2)	136	events[0].wait(stream2)
137	x2 = self.conv_operation(x2)	137	x2 = self.conv_operation(x2)
138	events[1].record()	138	events[1].record()
139	-	139	+
140	with torch.npu.stream(stream1):	140	with torch.npu.stream(stream1):
141	events[1].wait(stream1)	141	events[1].wait(stream1)
142	x1 = self.conv_with_allocator(x1)	142	x1 = self.conv_with_allocator(x1)
143	events[2].record()	143	events[2].record()
144	-	144	+
145	with torch.npu.stream(stream2):	145	with torch.npu.stream(stream2):
146	events[2].wait(stream2)	146	events[2].wait(stream2)
147	x2 = self.conv_operation(x2)	147	x2 = self.conv_operation(x2)
148		148
149	torch.npu.synchronize()	149	torch.npu.synchronize()
150	self.assertEqual(x1, x2)	150	self.assertEqual(x1, x2)
151	-	151	+
152	def test_mul_stream_with_threads(self):	152	def test_mul_stream_with_threads(self):
153	input_data = torch.randn(1, 1024, 96, dtype=torch.float32, device="npu")	153	input_data = torch.randn(1, 1024, 96, dtype=torch.float32, device="npu")
154	events = [torch.npu.Event(False, False) for _ in range(3)]	154	events = [torch.npu.Event(False, False) for _ in range(3)]
155	-	155	+
156	def stream_worker(data, stream, event_sequence):	156	def stream_worker(data, stream, event_sequence):
157	"""Generic stream worker function"""	157	"""Generic stream worker function"""
158	with torch.npu.stream(stream):	158	with torch.npu.stream(stream):
	@@ -161,7 +161,7 @@ class TestPluggableAllocator(TestCase):
161	data = operation(data)	161	data = operation(data)
162	events[event_sequence.index((event, operation)) + 1].record()	162	events[event_sequence.index((event, operation)) + 1].record()
163	return data	163	return data
164	-	164	+
165	# Define operation sequences for two streams	165	# Define operation sequences for two streams
166	stream1_ops = [	166	stream1_ops = [
167	(events[0], self.conv_with_allocator),	167	(events[0], self.conv_with_allocator),
	@@ -174,7 +174,7 @@ class TestPluggableAllocator(TestCase):
174		174
175	result_container = {}	175	result_container = {}
176	stream2 = torch.npu.Stream()	176	stream2 = torch.npu.Stream()
177	-	177	+
178	def thread_func():	178	def thread_func():
179	result_container["x2"] = stream_worker(input_data, stream2, stream2_ops)	179	result_container["x2"] = stream_worker(input_data, stream2, stream2_ops)
180		180
	@@ -207,7 +207,7 @@ class TestPluggableAllocator(TestCase):
207	class TestDictDataLoader():	207	class TestDictDataLoader():
208	def __init__(self):	208	def __init__(self):
209	self.dataset = DictDataset()	209	self.dataset = DictDataset()
210	-	210	+
211	def test_memory(self):	211	def test_memory(self):
212	loader = DataLoader(self.dataset, batch_size=2)	212	loader = DataLoader(self.dataset, batch_size=2)
213	for sample in loader:	213	for sample in loader:

Mtest/test_npu_pinned_memory_background_threads.py+1-1

	@@ -38,4 +38,4 @@ class TestPinnedMemoryBackgroundThreads(TestCase):
38	self.copy_tensor(ITERS)	38	self.copy_tensor(ITERS)
39		39
40	if __name__ == '__main__':	40	if __name__ == '__main__':
41	- run_tests()	41	+ run_tests()

Mtest/test_npu_sanitizer.py+2-2

	@@ -160,7 +160,7 @@ class TestArgumentHandler(TestCase):
160		160
161	self.assertEqual({out.data_ptr()}, argument_handler.dataptrs_written)	161	self.assertEqual({out.data_ptr()}, argument_handler.dataptrs_written)
162	self.assertEqual({out.data_ptr()}, argument_handler.outputs)	162	self.assertEqual({out.data_ptr()}, argument_handler.outputs)
163	-	163	+
164	def test_equal_reads_inputs_but_no_tensor_output_written(self):	164	def test_equal_reads_inputs_but_no_tensor_output_written(self):
165	"""Data-reading op with non-tensor output should not record tensor writes."""	165	"""Data-reading op with non-tensor output should not record tensor writes."""
166	equal_func = torch.ops.aten.equal.default	166	equal_func = torch.ops.aten.equal.default
	@@ -177,7 +177,7 @@ class TestArgumentHandler(TestCase):
177	self.assertEqual(set(), argument_handler.outputs)	177	self.assertEqual(set(), argument_handler.outputs)
178	self.assertTrue(isinstance(out, bool))	178	self.assertTrue(isinstance(out, bool))
179		179
180	-	180	+
181	class TestRecordStreamHandler(TestCase):	181	class TestRecordStreamHandler(TestCase):
182	def test_erase_stream_removes_recorded_stream(self):	182	def test_erase_stream_removes_recorded_stream(self):
183	"""Communication eraseStream should clear the matching recorded stream."""	183	"""Communication eraseStream should clear the matching recorded stream."""

Mtest/test_reductions.py+16-16

Mtest/test_sort_and_select.py+2-2

	@@ -6,7 +6,7 @@ import numpy as np
6	import torch	6	import torch
7	from torch import nan	7	from torch import nan
8	from torch.testing import make_tensor	8	from torch.testing import make_tensor
9	-from torch.testing._internal.common_dtype import (all_types, all_types_and, floating_types_and,	9	+from torch.testing._internal.common_dtype import (all_types, all_types_and, floating_types_and,
10	integral_types, _dispatch_dtypes)	10	integral_types, _dispatch_dtypes)
11	from torch.testing._internal.common_utils import (TestCase, run_tests, slowTest, skipIfTorchDynamo)	11	from torch.testing._internal.common_utils import (TestCase, run_tests, slowTest, skipIfTorchDynamo)
12	from torch.testing._internal.common_device_type import \	12	from torch.testing._internal.common_device_type import \
	@@ -17,7 +17,7 @@ import torch_npu
17	import torch_npu.testing	17	import torch_npu.testing
18		18
19	SIZE = 100	19	SIZE = 100
20	-all_types_without_double = _dispatch_dtypes((torch.half, torch.float32, torch.uint8,	20	+all_types_without_double = _dispatch_dtypes((torch.half, torch.float32, torch.uint8,
21	torch.int8, torch.int16, torch.int32, torch.int64))	21	torch.int8, torch.int16, torch.int32, torch.int64))
22		22
23		23

Mtest/test_torch_npu_init.py+1-1

	@@ -453,7 +453,7 @@ class TestTorchNpuBootstrap(TestCase):
453	)	453	)
454	"""	454	"""
455	)	455	)
456	-	456	+
457	def test_08_top_level_unsupported_dtype_compatibility(self):	457	def test_08_top_level_unsupported_dtype_compatibility(self):
458	self._run_python(	458	self._run_python(
459	"""	459	"""

Mtest/trans_contiguous/test_tensorto_preserve_format.py+7-7

Mtest/utils/test_torch_npu_logs_filter.py+1-1

	@@ -19,7 +19,7 @@ class TestTorchNpuLogs(TestCase):
19	os.environ['TORCH_NPU_LOGS'] = self.original_torch_npu_logs	19	os.environ['TORCH_NPU_LOGS'] = self.original_torch_npu_logs
20	else:	20	else:
21	del os.environ['TORCH_NPU_LOGS']	21	del os.environ['TORCH_NPU_LOGS']
22	-	22	+
23	if self.original_torch_npu_logs_filter is not None:	23	if self.original_torch_npu_logs_filter is not None:
24	os.environ['TORCH_NPU_LOGS_FILTER'] = self.original_torch_npu_logs_filter	24	os.environ['TORCH_NPU_LOGS_FILTER'] = self.original_torch_npu_logs_filter
25	else:	25	else:

Mtools/flight_recorder/check_path.py+1-1

	@@ -24,7 +24,7 @@ def type_to_str(value_type):
24	def check_type(value, value_type, param_name="value"):	24	def check_type(value, value_type, param_name="value"):
25	if not isinstance(value, value_type):	25	if not isinstance(value, value_type):
26	raise TypeError('{} must be {}, not {}.'.format(param_name, type_to_str(value_type), type(value).__name__))	26	raise TypeError('{} must be {}, not {}.'.format(param_name, type_to_str(value_type), type(value).__name__))
27	-	27	+
28		28
29	def get_valid_path(path):	29	def get_valid_path(path):
30	check_type(path, str, "path")	30	check_type(path, str, "path")

Mtools/flight_recorder/components/types.py+1-1

	@@ -338,7 +338,7 @@ class Op:
338	def __init__(self, event: dict[Any, Any], memberships: dict[str, set[Any]], pg_name: str):	338	def __init__(self, event: dict[Any, Any], memberships: dict[str, set[Any]], pg_name: str):
339		339
340	frames = event.get("frames")	340	frames = event.get("frames")
341	- if not frames:	341	+ if not frames:
342	raise ValueError(self.MISSING_FRAMES_ERR)	342	raise ValueError(self.MISSING_FRAMES_ERR)
343	first_frame = frames[0] if len(frames) > 0 else None	343	first_frame = frames[0] if len(frames) > 0 else None
344	if not first_frame:	344	if not first_frame:

Mtorch_npu/__init__.py+1-1

	@@ -58,7 +58,7 @@ def _initialize():
58	# 5. final extension barrier and shutdown hook	58	# 5. final extension barrier and shutdown hook
59	_initialize_runtime_lifecycle()	59	_initialize_runtime_lifecycle()
60		60
61	- # 6. optional runtime features	61	+ # 6. optional runtime features
62	_enable_optional_features()	62	_enable_optional_features()
63		63
64		64

Mtorch_npu/_inductor/__init__.py+6-6

	@@ -116,8 +116,8 @@ def _load_triton_backend():
116	from .shape_handling import NPUShapeHandling, patch_shape_handling	116	from .shape_handling import NPUShapeHandling, patch_shape_handling
117	from .utils import patch_get_first_incompatible_cudagraph_node	117	from .utils import patch_get_first_incompatible_cudagraph_node
118		118
119	- from .graph import patch_count_bytes, patch_run_node	119	+ from .graph import patch_count_bytes, patch_run_node
120	- from .autotune_process import patch_tuning_process, patch_tuning_process_pool	120	+ from .autotune_process import patch_tuning_process, patch_tuning_process_pool
121	flex_attention._validate_device = _validate_device	121	flex_attention._validate_device = _validate_device
122		122
123	def _inductor_register_backend_for_device():	123	def _inductor_register_backend_for_device():
	@@ -214,10 +214,10 @@ def _load_triton_backend():
214	patch_get_graph_partition_signature()	214	patch_get_graph_partition_signature()
215	patch_get_optimization_cflags()	215	patch_get_optimization_cflags()
216	patch_extract_read_writes()	216	patch_extract_read_writes()
217	- patch_count_bytes()	217	+ patch_count_bytes()
218	- patch_run_node()	218	+ patch_run_node()
219	- patch_tuning_process()	219	+ patch_tuning_process()
220	- patch_tuning_process_pool()	220	+ patch_tuning_process_pool()
221		221
222	def add_additional_op():	222	def add_additional_op():
223	from torch._inductor.ops_handler import OpsHandler	223	from torch._inductor.ops_handler import OpsHandler

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/codecache.py+8-8

	@@ -55,7 +55,7 @@ def codegen_subgraph_dump(inds, shapes, strides, dtypes, inds2):
55	codes.append(f'args = new_args')	55	codes.append(f'args = new_args')
56	return '\n'.join(codes)	56	return '\n'.join(codes)
57		57
58	-	58	+
59	def _worker_compile(	59	def _worker_compile(
60	kernel, cc: int, device: torch.device, logger_level=None, extra_env=None	60	kernel, cc: int, device: torch.device, logger_level=None, extra_env=None
61	) -> None:	61	) -> None:
	@@ -79,10 +79,10 @@ def _akg_worker_compile(
79		79
80		80
81	def _load_kernel(	81	def _load_kernel(
82	- kernel_name: str,	82	+ kernel_name: str,
83	- source_code: str,	83	+ source_code: str,
84	- no_more_compile=False,	84	+ no_more_compile=False,
85	- suppress_error=False,	85	+ suppress_error=False,
86	kernel_meta=None,	86	kernel_meta=None,
87	extra_env=None) -> ModuleType:	87	extra_env=None) -> ModuleType:
88	if os.getenv("TORCHINDUCTOR_USE_AKG", "0") == "1":	88	if os.getenv("TORCHINDUCTOR_USE_AKG", "0") == "1":
	@@ -144,7 +144,7 @@ class MulitprocessCompileFuture(CodeCacheFuture):
144	errors.append(e)	144	errors.append(e)
145		145
146	if len(errors) < len(self.futures):	146	if len(errors) < len(self.futures):
147	- kernel = self.kernel = _load_kernel(self.kernel_name, self.source_code,	147	+ kernel = self.kernel = _load_kernel(self.kernel_name, self.source_code,
148	no_more_compile=True, suppress_error=True,	148	no_more_compile=True, suppress_error=True,
149	kernel_meta=self.kernel_meta, extra_env=self.extra_env)	149	kernel_meta=self.kernel_meta, extra_env=self.extra_env)
150	elif self.kernel_meta.get('num_outputs', 0): # All compiles fail and auto fallback	150	elif self.kernel_meta.get('num_outputs', 0): # All compiles fail and auto fallback
	@@ -223,7 +223,7 @@ class CustomAsyncCompile(AsyncCompile):
223	pool.ready_future = pool.submit(AsyncCompile._get_ready) # type: ignore[attr-defined]	223	pool.ready_future = pool.submit(AsyncCompile._get_ready) # type: ignore[attr-defined]
224	_pool_set.add(pool)	224	_pool_set.add(pool)
225	return pool	225	return pool
226	-	226	+
227	def mlir(	227	def mlir(
228	self, kernel_name: str, source_code: str, device_str: str = "npu"	228	self, kernel_name: str, source_code: str, device_str: str = "npu"
229	) -> Union[NPUTritonFuture, ModuleType]:	229	) -> Union[NPUTritonFuture, ModuleType]:
	@@ -241,7 +241,7 @@ class CustomAsyncCompile(AsyncCompile):
241	return NPUTritonFuture(kernel_name, source_code, future)	241	return NPUTritonFuture(kernel_name, source_code, future)
242	else:	242	else:
243	return _load_kernel(kernel_name, source_code)	243	return _load_kernel(kernel_name, source_code)
244	-	244	+
245	def mlir_auto_fallback(	245	def mlir_auto_fallback(
246	self, kernel_name: str, source_code: str, kernel_meta: Dict[str, Any]) -> Callable:	246	self, kernel_name: str, source_code: str, kernel_meta: Dict[str, Any]) -> Callable:
247	_compile_start()	247	_compile_start()

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/config.py+8-8

@@ -26,7 +26,7 @@ autotune_fx_fallback = False

26

cache_named_op = False

26

cache_named_op = False

27

28

# NPU_INDUCTOR_FALLBACK_LIST=allfallback forces ops entering the NPU inductor lowering

28

# NPU_INDUCTOR_FALLBACK_LIST=allfallback forces ops entering the NPU inductor lowering

29

-# path to register fallback lowerings, so optimized/fused lowerings are not used.

29

+# path to register fallback lowerings, so optimized/fused lowerings are not used.

30

enable_full_lowering_fallback = os.environ.get("NPU_INDUCTOR_FALLBACK_LIST", "")

30

enable_full_lowering_fallback = os.environ.get("NPU_INDUCTOR_FALLBACK_LIST", "")

31

32

traced_graph_cache = os.environ.get("ANIR_TRACED_GRAPH_CACHE", None)

32

traced_graph_cache = os.environ.get("ANIR_TRACED_GRAPH_CACHE", None)

@@ -49,7 +49,7 @@ def parse_rtol_atol(env_str: str):

49

rtol, atol = None, None

49

rtol, atol = None, None

50

if not env_str.strip():

50

if not env_str.strip():

51

return rtol, atol

51

return rtol, atol

52

-

52

+

53

parts = [p.strip() for p in env_str.split(",") if p.strip()]

53

parts = [p.strip() for p in env_str.split(",") if p.strip()]

54

for part in parts:

54

for part in parts:

55

match = re.match(r"^(rtol|atol)\s*=s\*([0-9.eE+-]+)$", part, re.IGNORECASE)

55

match = re.match(r"^(rtol|atol)\s*=s\*([0-9.eE+-]+)$", part, re.IGNORECASE)

atomgit-bot5月21日

🟠 High Priority

第 55 行 re.match(r"^(rtol|atol)\s*=s*([0-9.eE+-]+)$", ...) 中的 =s* 匹配的是字面量 '=s*'，而非允许等号两侧可选空格的 \s*=\s*。这与错误提示信息 'rtol=1e-6,atol=1e-5' 矛盾，实际合法输入（如 rtol=1e-6）将无法通过匹配，导致 rtol/atol 永远取默认值。虽非本 diff 引入，但属于文件内既存 bug。

改动建议

55	- match = re.match(r"^(rtol\|atol)\s=s\([0-9.eE+-]+)$", part, re.IGNORECASE)
55	+ match = re.match(r"^(rtol\|atol)\s=\s([0-9.eE+-]+)$", part, re.IGNORECASE)

应用建议

Jingwei Huang

5月21日评论：

@@ -59,7 +59,7 @@ def parse_rtol_atol(env_str: str):

59

f"It should be like 'rtol=1e-6,atol=1e-5'. "

59

f"It should be like 'rtol=1e-6,atol=1e-5'. "

60

)

60

)

61

continue

61

continue

62

-

62

+

63

key, value_str = match.groups()

63

key, value_str = match.groups()

64

try:

64

try:

65

value = float(value_str)

65

value = float(value_str)

@@ -138,7 +138,7 @@ fx_subgraph_dump_path: str = os.environ.get("FX_SUBGRAPH_DUMP_PATH", None)

138

compile_mode introductions:

138

compile_mode introductions:

139

"default" refers to the mode of fully compiling with MLIR. Currently, it is not fully supported, but it will be set as the default once the capability matures.

139

"default" refers to the mode of fully compiling with MLIR. Currently, it is not fully supported, but it will be set as the default once the capability matures.

140

"complete_fallback" refers to completely falling back to the eager execution mode of the FX graph, without performing any MLIR compilation. It is primarily used for debugging.

140

"complete_fallback" refers to completely falling back to the eager execution mode of the FX graph, without performing any MLIR compilation. It is primarily used for debugging.

141

-"auto_fallback" refers to automatically falling back to the fx_graph_backend when compilation fails.

141

+"auto_fallback" refers to automatically falling back to the fx_graph_backend when compilation fails.

142

auto_fallback mechanism is designed to provide a fallback strategy when the primary compilation process encounters an issue. It works in conjunction with the fx_graph_backend configuration, allowing for the fallback approach:

142

auto_fallback mechanism is designed to provide a fallback strategy when the primary compilation process encounters an issue. It works in conjunction with the fx_graph_backend configuration, allowing for the fallback approach:

143

Fallback to fx_graph_backend: If the first fallback attempt fails, the system falls back to the fx_graph_backend.

143

Fallback to fx_graph_backend: If the first fallback attempt fails, the system falls back to the fx_graph_backend.

144

If you need further clarification or have other questions, please let me know!

144

If you need further clarification or have other questions, please let me know!

@@ -158,7 +158,7 @@ def _get_compile_mode():

158

block_dim = 48

158

block_dim = 48

159

160

"""

160

"""

161

-support {"off", "include", "exclude"}, to

161

+support {"off", "include", "exclude"}, to

162

"off": No fallback at all.

162

"off": No fallback at all.

163

"include": At compile-time, Aten IR included in FALLBACK_LIST will fall back to aten.

163

"include": At compile-time, Aten IR included in FALLBACK_LIST will fall back to aten.

164

"exclude": At compile-time, Aten IR excluded from GENERATE_LIST will fall back to aten.

164

"exclude": At compile-time, Aten IR excluded from GENERATE_LIST will fall back to aten.

@@ -169,7 +169,7 @@ if enable_full_lowering_fallback.strip()=='allfallback':

169

fallback_to_aten_mode = "all"

169

fallback_to_aten_mode = "all"

170

171

REDUCTION_OPS = [

171

REDUCTION_OPS = [

172

- aten.sum,

172

+ aten.sum,

173

prims.sum,

173

prims.sum,

174

aten.prod,

174

aten.prod,

175

aten.any,

175

aten.any,

@@ -181,12 +181,12 @@ REDUCTION_OPS = [

181

aten.argmax,

181

aten.argmax,

182

aten.argmin,

182

aten.argmin,

183

aten.mean,

183

aten.mean,

184

- aten.var,

184

+ aten.var,

185

prims.var,

185

prims.var,

186

aten.var_mean,

186

aten.var_mean,

187

]

187

]

188

189

-# fall back to aten exclude GENERATE_LIST, all aten IR except

189

+# fall back to aten exclude GENERATE_LIST, all aten IR except

190

POINTWISE_OPS = [

190

POINTWISE_OPS = [

191

aten.mul,

191

aten.mul,

192

aten.add,

192

aten.add,

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/akg.py+1-1

	@@ -13,7 +13,7 @@ class AkgKernel(NpuMetaKernel):
13	def call_kernel(self, name: str, node=None):	13	def call_kernel(self, name: str, node=None):
14	wrapper = V.graph.wrapper_code	14	wrapper = V.graph.wrapper_code
15	call_args = self.get_call_args()	15	call_args = self.get_call_args()
16	-	16	+
17	if len(call_args) > 0:	17	if len(call_args) > 0:
18	wrapper.generate_kernel_call(	18	wrapper.generate_kernel_call(
19	name,	19	name,

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/cpp_wrapper.py+5-5

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/meta_kernel.py+25-25

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/inductor_patch/ir.py+1-1

	@@ -450,7 +450,7 @@ ir.View.create = _patch_view_create
450	@classmethod	450	@classmethod
451	def _patch_sliceview_create(	451	def _patch_sliceview_create(
452	cls, x, dim, start, end, step=1, clamp=True, traced_graph=None, node_name=None	452	cls, x, dim, start, end, step=1, clamp=True, traced_graph=None, node_name=None
453	-):	453	+):
454	step = sympy.expand(step)	454	step = sympy.expand(step)
455	assert isinstance(step, sympy.Expr) or step > 0	455	assert isinstance(step, sympy.Expr) or step > 0
456	try:	456	try:

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/inductor_patch/lowering.py+38-38

	@@ -17,13 +17,13 @@ from collections.abc import Iterable, Sequence
17	from typing import Any, Callable, cast, Optional, TYPE_CHECKING, TypeVar, Union	17	from typing import Any, Callable, cast, Optional, TYPE_CHECKING, TypeVar, Union
18	from typing_extensions import ParamSpec	18	from typing_extensions import ParamSpec
19	from typing import (	19	from typing import (
20	- Any,	20	+ Any,
21	- Callable,	21	+ Callable,
22	- Dict,	22	+ Dict,
23	- List,	23	+ List,
24	- Optional,	24	+ Optional,
25	- Set,	25	+ Set,
26	- Tuple,	26	+ Tuple,
27	Union,	27	Union,
28	)	28	)
29	from unittest.mock import patch	29	from unittest.mock import patch
	@@ -230,7 +230,7 @@ def create_sym_inputs(traced_graph: TracedGraph, size: List[Expr]):
230	traced_graph.sym_nodes.update({s_name: new_node})	230	traced_graph.sym_nodes.update({s_name: new_node})
231		231
232		232
233	-def process_ir_constant(inp: ExpandView) -> Union[TracedGraph, int, float]:	233	+def process_ir_constant(inp: ExpandView) -> Union[TracedGraph, int, float]:
234	skip = False	234	skip = False
235	if isinstance(inp.data, IndexingConstant):	235	if isinstance(inp.data, IndexingConstant):
236	dtype = inp.data.dtype	236	dtype = inp.data.dtype
	@@ -268,13 +268,13 @@ def fetch_graphs(inputs: Optional[List[TensorBox]]):
268	input_graphs.append(inp)	268	input_graphs.append(inp)
269	continue	269	continue
270	if isinstance(inp, ExpandView):	270	if isinstance(inp, ExpandView):
271	- inp, skip = process_ir_constant(inp)	271	+ inp, skip = process_ir_constant(inp)
272	if not skip:	272	if not skip:
273	input_graphs.append(inp)	273	input_graphs.append(inp)
274	continue	274	continue
275	name = inp.get_name()	275	name = inp.get_name()
276	traced_graph = inp.get_traced_graph()	276	traced_graph = inp.get_traced_graph()
277	- if traced_graph is not None:	277	+ if traced_graph is not None:
278	input_graphs.append(traced_graph)	278	input_graphs.append(traced_graph)
279	continue	279	continue
280	traced_graph = TracedGraph()	280	traced_graph = TracedGraph()
	@@ -308,7 +308,7 @@ def merge_traced_graphs(input_graphs: List[TracedGraph], origin_fn, node_name, *
308	exist_nodes[node.name] = new_node	308	exist_nodes[node.name] = new_node
309	if node.name in input_graph.sym_nodes:	309	if node.name in input_graph.sym_nodes:
310	new_graph.sym_nodes.update({node.name: new_node})	310	new_graph.sym_nodes.update({node.name: new_node})
311	-	311	+
312	def parse_args(input_graphs, exist_nodes):	312	def parse_args(input_graphs, exist_nodes):
313	args = []	313	args = []
314	for input_graph in input_graphs:	314	for input_graph in input_graphs:
	@@ -324,7 +324,7 @@ def merge_traced_graphs(input_graphs: List[TracedGraph], origin_fn, node_name, *
324	else:	324	else:
325	args.append(input_graph)	325	args.append(input_graph)
326	return args	326	return args
327	-	327	+
328	num_args = len(input_graphs)	328	num_args = len(input_graphs)
329		329
330	for k, v in kwargs.items():	330	for k, v in kwargs.items():
	@@ -1054,8 +1054,8 @@ def _foreach_map(subgraph, args, *kwargs):
1054		1054
1055	assert all(x is not None for x in outputs)	1055	assert all(x is not None for x in outputs)
1056	return outputs	1056	return outputs
1057	-	1057	+
1058	-	1058	+
1059	@register_lowering(torch.ops.npu.npu_dtype_cast_backward, type_promotion_kind=None)	1059	@register_lowering(torch.ops.npu.npu_dtype_cast_backward, type_promotion_kind=None)
1060	@register_lowering(torch.ops.npu.npu_dtype_cast, type_promotion_kind=None)	1060	@register_lowering(torch.ops.npu.npu_dtype_cast, type_promotion_kind=None)
1061	@register_lowering(torch.ops.npu._npu_dtype_cast_backward, type_promotion_kind=None)	1061	@register_lowering(torch.ops.npu._npu_dtype_cast_backward, type_promotion_kind=None)
	@@ -1110,7 +1110,7 @@ def to_device(x: TensorBox, device: torch.device, *, copy=False, non_blocking=Fa
1110	device = decode_device(device)	1110	device = decode_device(device)
1111	if x.get_device() == device:	1111	if x.get_device() == device:
1112	return clone(x) if copy else x	1112	return clone(x) if copy else x
1113	-	1113	+
1114	input_graphs = fetch_graphs([x, device])	1114	input_graphs = fetch_graphs([x, device])
1115	node_name = f'to_device_{next(node_id)}'	1115	node_name = f'to_device_{next(node_id)}'
1116	new_graph = merge_traced_graphs(input_graphs, aten.to.device, node_name, dtype=src_dtype, copy=copy)	1116	new_graph = merge_traced_graphs(input_graphs, aten.to.device, node_name, dtype=src_dtype, copy=copy)
	@@ -1150,7 +1150,7 @@ def register_pointwise(
1150	if override_fn_when_input_bool is not None:	1150	if override_fn_when_input_bool is not None:
1151	override_fn_when_input_bool = ops_wrapper(override_fn_when_input_bool)	1151	override_fn_when_input_bool = ops_wrapper(override_fn_when_input_bool)
1152		1152
1153	- fn = register_fn_to_aten_fn(fn, aten_fn)	1153	+ fn = register_fn_to_aten_fn(fn, aten_fn)
1154		1154
1155	fn = make_pointwise(	1155	fn = make_pointwise(
1156	fn,	1156	fn,
	@@ -1435,7 +1435,7 @@ def repeat(x, repeats):
1435	input_graphs = fetch_graphs([x, repeats])	1435	input_graphs = fetch_graphs([x, repeats])
1436	node_name = f'repeat_{next(node_id)}'	1436	node_name = f'repeat_{next(node_id)}'
1437	new_graph = merge_traced_graphs(input_graphs, aten.repeat, node_name)	1437	new_graph = merge_traced_graphs(input_graphs, aten.repeat, node_name)
1438	-	1438	+
1439	old_size = list(x.get_size())	1439	old_size = list(x.get_size())
1440	if len(repeats) > len(old_size):	1440	if len(repeats) > len(old_size):
1441	old_size = [sympy.S.One] * (len(repeats) - len(old_size)) + old_size	1441	old_size = [sympy.S.One] * (len(repeats) - len(old_size)) + old_size
	@@ -1516,7 +1516,7 @@ def slice_(x, dim=0, start=0, end=2**63, step=1, clamp=True):
1516	input_graphs = fetch_graphs([x.data])	1516	input_graphs = fetch_graphs([x.data])
1517	node_name = f'slice_{next(node_id)}'	1517	node_name = f'slice_{next(node_id)}'
1518	new_graph = merge_traced_graphs(input_graphs, aten.slice, node_name, dim=dim, start=start, end=end, step=step)	1518	new_graph = merge_traced_graphs(input_graphs, aten.slice, node_name, dim=dim, start=start, end=end, step=step)
1519	-	1519	+
1520	return TensorBox(ir.SliceView.create(x.data, dim, start, end, step, clamp=clamp, traced_graph=new_graph, node_name=node_name))	1520	return TensorBox(ir.SliceView.create(x.data, dim, start, end, step, clamp=clamp, traced_graph=new_graph, node_name=node_name))
1521		1521
1522		1522
	@@ -3253,8 +3253,8 @@ def slice_scatter(x, src, dim=0, start=None, end=None, step=1):
3253	input_graphs = fetch_graphs([x, src])	3253	input_graphs = fetch_graphs([x, src])
3254	node_name = f'slice_scatter_{next(node_id)}'	3254	node_name = f'slice_scatter_{next(node_id)}'
3255	new_graph = merge_traced_graphs(input_graphs, aten.slice_scatter, node_name, \	3255	new_graph = merge_traced_graphs(input_graphs, aten.slice_scatter, node_name, \
3256	- dim=dim,	3256	+ dim=dim,
3257	- start=start,	3257	+ start=start,
3258	end=end,	3258	end=end,
3259	step=step)	3259	step=step)
3260	x_loader = x.make_loader()	3260	x_loader = x.make_loader()
	@@ -3337,7 +3337,7 @@ def tensor(data, *, dtype=None, device=None, layout=None, pin_memory=False):
3337	input_graphs = fetch_graphs([data])	3337	input_graphs = fetch_graphs([data])
3338	node_name = f'tensor_{next(node_id)}'	3338	node_name = f'tensor_{next(node_id)}'
3339	new_graph = merge_traced_graphs(input_graphs, torch.tensor, node_name, \	3339	new_graph = merge_traced_graphs(input_graphs, torch.tensor, node_name, \
3340	- dtype=dtype,	3340	+ dtype=dtype,
3341	device='npu',	3341	device='npu',
3342	pin_memory=False)	3342	pin_memory=False)
3343	if isinstance(_unwrap(data), int):	3343	if isinstance(_unwrap(data), int):
	@@ -3496,7 +3496,7 @@ def _full(fill_value, device, dtype, size):
3496		3496
3497	def inner_fn(index):	3497	def inner_fn(index):
3498	return value_loader([])	3498	return value_loader([])
3499	-	3499	+
3500	node_name = f'full_{next(node_id)}'	3500	node_name = f'full_{next(node_id)}'
3501	# [wtd#18] Use passed-in device param instead of hardcoded 'npu' to avoid device mismatch	3501	# [wtd#18] Use passed-in device param instead of hardcoded 'npu' to avoid device mismatch
3502	new_graph = merge_traced_graphs([size, fill_value], aten.full.default, node_name, \	3502	new_graph = merge_traced_graphs([size, fill_value], aten.full.default, node_name, \
	@@ -3751,8 +3751,8 @@ def embedding(weight, indices, padding_idx=-1, scale_grad_by_freq=False, sparse=
3751	input_graphs = fetch_graphs([weight, indices])	3751	input_graphs = fetch_graphs([weight, indices])
3752	node_name = f'embedding_{next(node_id)}'	3752	node_name = f'embedding_{next(node_id)}'
3753	new_graph = merge_traced_graphs(input_graphs, aten.embedding, node_name, \	3753	new_graph = merge_traced_graphs(input_graphs, aten.embedding, node_name, \
3754	- padding_idx=padding_idx,	3754	+ padding_idx=padding_idx,
3755	- scale_grad_by_freq=scale_grad_by_freq,	3755	+ scale_grad_by_freq=scale_grad_by_freq,
3756	sparse=sparse)	3756	sparse=sparse)
3757		3757
3758	return Pointwise.create(	3758	return Pointwise.create(
	@@ -3933,7 +3933,7 @@ def _unsafe_index(x, indices):
3933	# We cannot have this lowering as a decomposition as it introduces	3933	# We cannot have this lowering as a decomposition as it introduces
3934	# mutation in the graph, which is bad for Aot Autograd. Aot Autograd runs dead	3934	# mutation in the graph, which is bad for Aot Autograd. Aot Autograd runs dead
3935	# code elimination and common subexpression elimination optimizations, which	3935	# code elimination and common subexpression elimination optimizations, which
3936	-# assume graphs to be side-effect free.	3936	+# assume graphs to be side-effect free.
3937	@register_lowering(aten.index_put)	3937	@register_lowering(aten.index_put)
3938	def index_put(x, indices, values, accumulate=False):	3938	def index_put(x, indices, values, accumulate=False):
3939	return index_put_impl_(	3939	return index_put_impl_(
	@@ -6059,7 +6059,7 @@ def should_not_sum(a, b, keepdims):
6059	if not keepdims:	6059	if not keepdims:
6060	for i in unique_indices:	6060	for i in unique_indices:
6061	del a[i]	6061	del a[i]
6062	- return a	6062	+ return a
6063	return a	6063	return a
6064		6064
6065	def make_reduction(reduction_type: ReductionType, override_return_dtype=None):	6065	def make_reduction(reduction_type: ReductionType, override_return_dtype=None):
	@@ -6082,15 +6082,15 @@ def make_reduction(reduction_type: ReductionType, override_return_dtype=None):
6082	node_name = f'reshape_{next(node_id)}'	6082	node_name = f'reshape_{next(node_id)}'
6083	input_graphs = fetch_graphs([x, new_size])	6083	input_graphs = fetch_graphs([x, new_size])
6084	new_graph = merge_traced_graphs(input_graphs, aten.reshape, node_name)	6084	new_graph = merge_traced_graphs(input_graphs, aten.reshape, node_name)
6085	- else:	6085	+ else:
6086	node_name = f'reduction_{next(node_id)}'	6086	node_name = f'reduction_{next(node_id)}'
6087	input_graphs = fetch_graphs([x, axis if axis is not None else list(range(len(x.get_size())))])	6087	input_graphs = fetch_graphs([x, axis if axis is not None else list(range(len(x.get_size())))])
6088	- new_graph = merge_traced_graphs(input_graphs, reduction_type_to_aten_fn[reduction_type],	6088	+ new_graph = merge_traced_graphs(input_graphs, reduction_type_to_aten_fn[reduction_type],
6089	node_name, keepdim=keepdims)	6089	node_name, keepdim=keepdims)
6090	- result = Reduction.create(reduction_type=reduction_type,	6090	+ result = Reduction.create(reduction_type=reduction_type,
6091	- input_node=x,	6091	+ input_node=x,
6092	- node_name=node_name,	6092	+ node_name=node_name,
6093	- traced_graph=new_graph,	6093	+ traced_graph=new_graph,
6094	**kwargs)	6094	**kwargs)
6095	if isinstance(	6095	if isinstance(
6096	result.data.data, # type: ignore[attr-defined]	6096	result.data.data, # type: ignore[attr-defined]
	@@ -6304,7 +6304,7 @@ def pow(a, b):
6304		6304
6305	def fn(idx):	6305	def fn(idx):
6306	return pow_recursive(loader(idx), b, a.get_dtype())	6306	return pow_recursive(loader(idx), b, a.get_dtype())
6307	-	6307	+
6308	input_graphs = fetch_graphs([a, b])	6308	input_graphs = fetch_graphs([a, b])
6309	node_name = f'pointwise_{next(node_id)}'	6309	node_name = f'pointwise_{next(node_id)}'
6310	new_graph = merge_traced_graphs(input_graphs, aten.pow, node_name)	6310	new_graph = merge_traced_graphs(input_graphs, aten.pow, node_name)
	@@ -6426,7 +6426,7 @@ def mul(a, b):
6426	return logical_and(a, b)	6426	return logical_and(a, b)
6427	else:	6427	else:
6428	fn = ops_wrapper(aten.mul.__name__)	6428	fn = ops_wrapper(aten.mul.__name__)
6429	- fn = register_fn_to_aten_fn(fn, aten.mul)	6429	+ fn = register_fn_to_aten_fn(fn, aten.mul)
6430	return make_pointwise(fn)(a, b)	6430	return make_pointwise(fn)(a, b)
6431		6431
6432		6432
	@@ -6534,13 +6534,13 @@ def split_last_continuous(lst):
6534	n = len(lst)	6534	n = len(lst)
6535	if n == 1:	6535	if n == 1:
6536	return lst, []	6536	return lst, []
6537	-	6537	+
6538	i = n - 2	6538	i = n - 2
6539	while i >= 0:	6539	while i >= 0:
6540	if lst[i] + 1 != lst[i + 1]:	6540	if lst[i] + 1 != lst[i + 1]:
6541	break	6541	break
6542	i -= 1	6542	i -= 1
6543	-	6543	+
6544	start = i + 1	6544	start = i + 1
6545	last_part = lst[start:]	6545	last_part = lst[start:]
6546	remaining = lst[:start]	6546	remaining = lst[:start]
	@@ -6564,7 +6564,7 @@ def sum_(x, axis=None, keepdims=False, *, dtype=None):
6564	fn = make_reduction("sum", override_return_dtype=dtype)	6564	fn = make_reduction("sum", override_return_dtype=dtype)
6565	r = fn(x, axis, keepdims, dtype=dtype)	6565	r = fn(x, axis, keepdims, dtype=dtype)
6566	return r	6566	return r
6567	-	6567	+
6568		6568
6569	fallback_cumsum = fallback_handler(aten.cumsum.default)	6569	fallback_cumsum = fallback_handler(aten.cumsum.default)
6570	fallback_cumprod = fallback_handler(aten.cumprod.default)	6570	fallback_cumprod = fallback_handler(aten.cumprod.default)
	@@ -6799,7 +6799,7 @@ def register_pointwise_numeric_ldf64(op):
6799	type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT,	6799	type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT,
6800	use_libdevice_for_f64=True,	6800	use_libdevice_for_f64=True,
6801	)	6801	)
6802	-	6802	+
6803	@register_lowering([aten.neg])	6803	@register_lowering([aten.neg])
6804	def neg(a):	6804	def neg(a):
6805	if a.get_dtype() in (torch.int32, torch.int64):	6805	if a.get_dtype() in (torch.int32, torch.int64):

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/meta_compiler.py+4-4

	@@ -100,10 +100,10 @@ class MetaCompiler:
100	shutil.rmtree(failed_subgraph_dump_path)	100	shutil.rmtree(failed_subgraph_dump_path)
101	shutil.copytree(subgraph_dump_path, failed_subgraph_dump_path)	101	shutil.copytree(subgraph_dump_path, failed_subgraph_dump_path)
102	return failed_subgraph_dump_path	102	return failed_subgraph_dump_path
103	-	103	+
104	def acc_compare_and_dump(self, args, dump_data=True, *kwargs) -> Tuple[Any, bool]:	104	def acc_compare_and_dump(self, args, dump_data=True, *kwargs) -> Tuple[Any, bool]:
105	self.register_fx_fallback(self.kernel_meta)	105	self.register_fx_fallback(self.kernel_meta)
106	-	106	+
107	output, has_acc_error = check_accuracy_mlir(	107	output, has_acc_error = check_accuracy_mlir(
108	*args,	108	*args,
109	kernel_name=self.kernel_name,	109	kernel_name=self.kernel_name,
	@@ -112,7 +112,7 @@ class MetaCompiler:
112	dynamic=self.dynamic,	112	dynamic=self.dynamic,
113	**kwargs	113	**kwargs
114	)	114	)
115	-	115	+
116	if anir_config.fx_subgraph_dump_path:	116	if anir_config.fx_subgraph_dump_path:
117	data = args	117	data = args
118	if dump_data and has_acc_error:	118	if dump_data and has_acc_error:
	@@ -122,7 +122,7 @@ class MetaCompiler:
122	torch.npu.synchronize()	122	torch.npu.synchronize()
123	self.launchers = [self.launchers[0]]	123	self.launchers = [self.launchers[0]]
124	self.is_fallback_kernels = [self.is_fallback_kernels[0]]	124	self.is_fallback_kernels = [self.is_fallback_kernels[0]]
125	-	125	+
126	return output, not has_acc_error	126	return output, not has_acc_error
127		127
128	def compile(self, args, *kwargs):	128	def compile(self, args, *kwargs):

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/mlir_compiler.py+24-24

	@@ -37,9 +37,9 @@ reinterpret_tensor = torch.ops.inductor._reinterpret_tensor
37	global_cache = set()	37	global_cache = set()
38		38
39	class NpuMlirCompiler(MetaCompiler):	39	class NpuMlirCompiler(MetaCompiler):
40	- def __init__(self,	40	+ def __init__(self,
41	- kernel_name: str = '',	41	+ kernel_name: str = '',
42	- multiprocess_compile=False,	42	+ multiprocess_compile=False,
43	no_more_compile=False,	43	no_more_compile=False,
44	kernel_meta=None,	44	kernel_meta=None,
45	autotune=True):	45	autotune=True):
	@@ -127,7 +127,7 @@ class NpuMlirCompiler(MetaCompiler):
127	logger.info(f"[bisheng-compile failed]")	127	logger.info(f"[bisheng-compile failed]")
128	logger.warning(f"Compile error msg: {e.stderr.decode('utf-8')}")	128	logger.warning(f"Compile error msg: {e.stderr.decode('utf-8')}")
129	raise e	129	raise e
130	-	130	+
131	def prepare_launch(self, need_pickle=False):	131	def prepare_launch(self, need_pickle=False):
132	def get_launch_mod(so_path):	132	def get_launch_mod(so_path):
133	spec = importlib.util.spec_from_file_location("__launcher", so_path)	133	spec = importlib.util.spec_from_file_location("__launcher", so_path)
	@@ -169,7 +169,7 @@ class NpuMlirCompiler(MetaCompiler):
169	shutil.copy(cache_mlir_path, os.path.join(anir_config.fx_subgraph_dump_path, \	169	shutil.copy(cache_mlir_path, os.path.join(anir_config.fx_subgraph_dump_path, \
170	str(self.device_index), self.kernel_name))	170	str(self.device_index), self.kernel_name))
171	return cache_mlir_path	171	return cache_mlir_path
172	-	172	+
173	def get_launch_dynamic(self, function, tiling_func, tiling_size):	173	def get_launch_dynamic(self, function, tiling_func, tiling_size):
174	block_dim = anir_config.block_dim	174	block_dim = anir_config.block_dim
175	arg_tiling_device = torch.empty((tiling_size // 8), device='npu', dtype=torch.int64)	175	arg_tiling_device = torch.empty((tiling_size // 8), device='npu', dtype=torch.int64)
	@@ -177,30 +177,30 @@ class NpuMlirCompiler(MetaCompiler):
177	def kernel_call(*args, stream=None):	177	def kernel_call(*args, stream=None):
178	self.launch(block_dim, stream, function, tiling_func, tiling_size, arg_tiling_host, arg_tiling_device, None, None, None, *args)	178	self.launch(block_dim, stream, function, tiling_func, tiling_size, arg_tiling_host, arg_tiling_device, None, None, None, *args)
179	return kernel_call	179	return kernel_call
180	-	180	+
181	def get_launch(self, function):	181	def get_launch(self, function):
182	block_dim = anir_config.block_dim	182	block_dim = anir_config.block_dim
183	def kernel_call(*args, function, stream=None):	183	def kernel_call(*args, function, stream=None):
184	self.launch(block_dim, stream, function, None, None, None, *args)	184	self.launch(block_dim, stream, function, None, None, None, *args)
185		185
186	return functools.partial(kernel_call, function=function)	186	return functools.partial(kernel_call, function=function)
187	-	187	+
188	def get_launch_func(self, cache_kernel_path):	188	def get_launch_func(self, cache_kernel_path):
189	if self.dynamic:	189	if self.dynamic:
190	- function, tiling_func, tiling_size = self.get_host_func_and_tiling_size(self.kernel_name,	190	+ function, tiling_func, tiling_size = self.get_host_func_and_tiling_size(self.kernel_name,
191	- self.kernel_name + '_tiling_function',	191	+ self.kernel_name + '_tiling_function',
192	- self.kernel_name + '_get_tiling_struct_size_function',	192	+ self.kernel_name + '_get_tiling_struct_size_function',
193	cache_kernel_path)	193	cache_kernel_path)
194	return self.get_launch_dynamic(function, tiling_func, tiling_size)	194	return self.get_launch_dynamic(function, tiling_func, tiling_size)
195	else:	195	else:
196	function = load_kernel_binary(self.kernel_name, cache_kernel_path)	196	function = load_kernel_binary(self.kernel_name, cache_kernel_path)
197	return self.get_launch(function)	197	return self.get_launch(function)
198	-	198	+
199	- def register_launcher(self,	199	+ def register_launcher(self,
200	- launcher,	200	+ launcher,
201	- kernel_path=None,	201	+ kernel_path=None,
202	- num_outputs=None,	202	+ num_outputs=None,
203	- disable_dump=False,	203	+ disable_dump=False,
204	auto_fallback=False,	204	auto_fallback=False,
205	is_fallback_kernel=False):	205	is_fallback_kernel=False):
206	if num_outputs:	206	if num_outputs:
	@@ -216,7 +216,7 @@ class NpuMlirCompiler(MetaCompiler):
216	self.fx_subgraph_dump('fallback')	216	self.fx_subgraph_dump('fallback')
217	logger.info(f"register launcher {launcher} {kernel_path} success")	217	logger.info(f"register launcher {launcher} {kernel_path} success")
218		218
219	- def compile_mlir(self,	219	+ def compile_mlir(self,
220	device_info: Tuple[Any],	220	device_info: Tuple[Any],
221	compile_args: List[Any],	221	compile_args: List[Any],
222	logger_level = None) -> Callable[..., None]:	222	logger_level = None) -> Callable[..., None]:
	@@ -237,7 +237,7 @@ class NpuMlirCompiler(MetaCompiler):
237		237
238	logger.info("Start to get cached kernel. Tiling info: " +	238	logger.info("Start to get cached kernel. Tiling info: " +
239	f"tiling_size {tiling_size} ops_reorder {ops_reorder} auto_db {auto_db}")	239	f"tiling_size {tiling_size} ops_reorder {ops_reorder} auto_db {auto_db}")
240	-	240	+
241	if cache_kernel_path is None and self.no_more_compile:	241	if cache_kernel_path is None and self.no_more_compile:
242	raise RuntimeError("Skip compile.")	242	raise RuntimeError("Skip compile.")
243		243
	@@ -248,8 +248,8 @@ class NpuMlirCompiler(MetaCompiler):
248	self.bisheng_compile(named_op_mlir_path, kernel_path, tiling_size=tiling_size,	248	self.bisheng_compile(named_op_mlir_path, kernel_path, tiling_size=tiling_size,
249	ops_reorder=ops_reorder, auto_db=auto_db,	249	ops_reorder=ops_reorder, auto_db=auto_db,
250	extra_command=anir_config.extra_command)	250	extra_command=anir_config.extra_command)
251	-	251	+
252	-	252	+
253	if self.dynamic:	253	if self.dynamic:
254	kernel_path = os.path.join(tmpdir, f"lib{tiling_kernel_name}.so")	254	kernel_path = os.path.join(tmpdir, f"lib{tiling_kernel_name}.so")
255	with open(kernel_path, "rb") as f:	255	with open(kernel_path, "rb") as f:
	@@ -303,7 +303,7 @@ class NpuMlirCompiler(MetaCompiler):
303	self.launch = getattr(mod, "launch")	303	self.launch = getattr(mod, "launch")
304	if self.dynamic:	304	if self.dynamic:
305	self.get_host_func_and_tiling_size = getattr(mod, "get_host_func_and_tiling_size")	305	self.get_host_func_and_tiling_size = getattr(mod, "get_host_func_and_tiling_size")
306	-	306	+
307	launch_func = self.get_launch_func(kernel_path)	307	launch_func = self.get_launch_func(kernel_path)
308	self.register_launcher(launch_func, kernel_path)	308	self.register_launcher(launch_func, kernel_path)
309	return True	309	return True
	@@ -324,7 +324,7 @@ class NpuMlirCompiler(MetaCompiler):
324	and self.kernel_meta.get('is_reduction', False)	324	and self.kernel_meta.get('is_reduction', False)
325	)	325	)
326		326
327	- def precompile(self,	327	+ def precompile(self,
328	device_info: Tuple[Any],	328	device_info: Tuple[Any],
329	suppress_error=False,	329	suppress_error=False,
330	logger_level=None):	330	logger_level=None):
	@@ -403,7 +403,7 @@ class NpuMlirCompiler(MetaCompiler):
403	else clone_preserve_strides(arg) for arg in args[-self.num_outputs:]]	403	else clone_preserve_strides(arg) for arg in args[-self.num_outputs:]]
404	fx_inputs = [clone_preserve_strides(arg) if isinstance(arg, torch.Tensor) else arg for arg in args[:-self.num_outputs]]	404	fx_inputs = [clone_preserve_strides(arg) if isinstance(arg, torch.Tensor) else arg for arg in args[:-self.num_outputs]]
405	fx_inputs = [inp.float() if isinstance(inp, torch.Tensor) and inp.dtype == torch.bfloat16 else inp for inp in fx_inputs]	405	fx_inputs = [inp.float() if isinstance(inp, torch.Tensor) and inp.dtype == torch.bfloat16 else inp for inp in fx_inputs]
406	-	406	+
407	fx_args = fx_inputs + fx_outputs	407	fx_args = fx_inputs + fx_outputs
408	launcher_fx(fx_args, *kwargs)	408	launcher_fx(fx_args, *kwargs)
409		409
	@@ -455,7 +455,7 @@ class NpuMlirCompiler(MetaCompiler):
455	print(f"{self.kernel_name}: Tuning accuracy failed, no valid kernels found, using fallback")	455	print(f"{self.kernel_name}: Tuning accuracy failed, no valid kernels found, using fallback")
456	timings.append([float(1.0), len(self.launchers) - 1])	456	timings.append([float(1.0), len(self.launchers) - 1])
457	return timings	457	return timings
458	-	458	+
459	def autotune_to_one_config(self, args, *kwargs):	459	def autotune_to_one_config(self, args, *kwargs):
460	if any([isinstance(arg, torch.Tensor) and not arg.is_contiguous() for arg in args]):	460	if any([isinstance(arg, torch.Tensor) and not arg.is_contiguous() for arg in args]):
461	print(f'Non contiguous args exists! Kernel name is {self.kernel_name}')	461	print(f'Non contiguous args exists! Kernel name is {self.kernel_name}')

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/npu_decomp.py+5-5

	@@ -81,7 +81,7 @@ def npu_convolution_backward(
81	[output_mask[0], output_mask[1], False],	81	[output_mask[0], output_mask[1], False],
82	)	82	)
83	return (grad_inp, grad_weight, grad_bias)	83	return (grad_inp, grad_weight, grad_bias)
84	-	84	+
85	def npu__softmax_backward_data(	85	def npu__softmax_backward_data(
86	grad_output: torch.Tensor,	86	grad_output: torch.Tensor,
87	output: torch.Tensor,	87	output: torch.Tensor,
	@@ -112,9 +112,9 @@ def npu_rms_norm(
112	output = (x * rsqrt * weight).to(dtype)	112	output = (x * rsqrt * weight).to(dtype)
113	return output, rsqrt	113	return output, rsqrt
114		114
115	-def npu_rms_norm_backward(grad_output: torch.Tensor,	115	+def npu_rms_norm_backward(grad_output: torch.Tensor,
116	- x: torch.Tensor,	116	+ x: torch.Tensor,
117	- weight: torch.Tensor,	117	+ weight: torch.Tensor,
118	rsqrt: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:	118	rsqrt: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
119	dx = (grad_output * weight - x * rsqrt * (grad_output * weight * x * rsqrt).mean(-1, keepdim=True)) * rsqrt	119	dx = (grad_output * weight - x * rsqrt * (grad_output * weight * x * rsqrt).mean(-1, keepdim=True)) * rsqrt
120	dgamma = (grad_output * x * rsqrt).sum(0, keepdim=False)	120	dgamma = (grad_output * x * rsqrt).sum(0, keepdim=False)
	@@ -140,7 +140,7 @@ def npu_swiglu_backward(grad_output, x, dim=-1):
140	def _rotate_half(x: Tensor) -> Tensor:	140	def _rotate_half(x: Tensor) -> Tensor:
141	x1, x2 = torch.chunk(x, 2, dim=-1)	141	x1, x2 = torch.chunk(x, 2, dim=-1)
142	return torch.cat((-x2, x1), dim=-1)	142	return torch.cat((-x2, x1), dim=-1)
143	-	143	+
144	def npu_rotary_mul(t, cos_, sin_):	144	def npu_rotary_mul(t, cos_, sin_):
145	t = (t * cos_) + (_rotate_half(t) * sin_)	145	t = (t * cos_) + (_rotate_half(t) * sin_)
146	return t	146	return t

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/npu_inductor_plugin.py+8-8

	@@ -128,7 +128,7 @@ def disable_implicit_decomposition():
128	op_override.py_kernels.pop(DispatchKey.Autograd)	128	op_override.py_kernels.pop(DispatchKey.Autograd)
129	if DispatchKey.CompositeImplicitAutograd in op_override.py_kernels:	129	if DispatchKey.CompositeImplicitAutograd in op_override.py_kernels:
130	op_override.py_kernels.pop(DispatchKey.CompositeImplicitAutograd)	130	op_override.py_kernels.pop(DispatchKey.CompositeImplicitAutograd)
131	-	131	+
132		132
133	def _patch_run_node(tracer, node, args, kwargs, nnmodule):	133	def _patch_run_node(tracer, node, args, kwargs, nnmodule):
134	op = node.op	134	op = node.op
	@@ -140,7 +140,7 @@ def _patch_run_node(tracer, node, args, kwargs, nnmodule):
140		140
141	try:	141	try:
142	if op == "call_function":	142	if op == "call_function":
143	- # patch start	143	+ # patch start
144	if 'npu.npu_fusion_attention' in str(node.target):	144	if 'npu.npu_fusion_attention' in str(node.target):
145	if 'actual_seq_qlen' in kwargs:	145	if 'actual_seq_qlen' in kwargs:
146	kwargs['actual_seq_qlen'] = list(kwargs['actual_seq_qlen'])	146	kwargs['actual_seq_qlen'] = list(kwargs['actual_seq_qlen'])
	@@ -187,7 +187,7 @@ disable_implicit_decomposition()
187	torch._dynamo.utils.run_node = _patch_run_node	187	torch._dynamo.utils.run_node = _patch_run_node
188		188
189		189
190	-from torch._dynamo.backends import common	190	+from torch._dynamo.backends import common
191	from torch._dynamo.backends.common import AotAutograd	191	from torch._dynamo.backends.common import AotAutograd
192		192
193	def wrap_compiler(fn):	193	def wrap_compiler(fn):
	@@ -212,13 +212,13 @@ def wrap_aot_autograd(fn):
212		212
213	AotAutograd.__call__ = wrap_aot_autograd(AotAutograd.__call__)	213	AotAutograd.__call__ = wrap_aot_autograd(AotAutograd.__call__)
214		214
215	-# recompute last usage for inductor scheduler	215	+# recompute last usage for inductor scheduler
216	from torch._inductor import scheduler	216	from torch._inductor import scheduler
217	from torch._inductor.scheduler import (	217	from torch._inductor.scheduler import (
218	Dep,	218	Dep,
219	WeakDep,	219	WeakDep,
220	- Scheduler,	220	+ Scheduler,
221	- SchedulerNode,	221	+ SchedulerNode,
222	SchedulerBuffer,	222	SchedulerBuffer,
223	FusedSchedulerNode,	223	FusedSchedulerNode,
224	BaseSchedulerNode,	224	BaseSchedulerNode,
	@@ -429,7 +429,7 @@ def patch_transfer_to_npu():
429	_replace_cuda_to_npu_in_kwargs,	429	_replace_cuda_to_npu_in_kwargs,
430	)	430	)
431		431
432	- def new_wrapper_cuda(module, method):	432	+ def new_wrapper_cuda(module, method):
433	src_method = f"_src_{method}"	433	src_method = f"_src_{method}"
434	if hasattr(getattr(module, method), '__wrapped__'):	434	if hasattr(getattr(module, method), '__wrapped__'):
435	src_func = getattr(module, method).__wrapped__	435	src_func = getattr(module, method).__wrapped__
	@@ -438,7 +438,7 @@ def patch_transfer_to_npu():
438		438
439	setattr(module, src_method, src_func)	439	setattr(module, src_method, src_func)
440	fn = getattr(module, src_method)	440	fn = getattr(module, src_method)
441	-	441	+
442	def decorated(args, *kwargs):	442	def decorated(args, *kwargs):
443	replace_int = fn.__name__ in ['to', 'to_empty']	443	replace_int = fn.__name__ in ['to', 'to_empty']
444	if args:	444	if args:

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/npu_lowering.py+9-9

	@@ -21,7 +21,7 @@ def _register_npu_inductor_fallbacks():
21	fallback_set = set()	21	fallback_set = set()
22	fallback_set_exclude = set()	22	fallback_set_exclude = set()
23	env_fallback_list = config.enable_full_lowering_fallback	23	env_fallback_list = config.enable_full_lowering_fallback
24	-	24	+
25	def _resolve_op_from_name(op_name: str):	25	def _resolve_op_from_name(op_name: str):
26	try:	26	try:
27	obj = torch.ops	27	obj = torch.ops
	@@ -31,7 +31,7 @@ def _register_npu_inductor_fallbacks():
31	except AttributeError:	31	except AttributeError:
32	log.warning(f"[npu\|inductor\|lowering\|fallback] invalid identifier name: {op_name}")	32	log.warning(f"[npu\|inductor\|lowering\|fallback] invalid identifier name: {op_name}")
33	return None	33	return None
34	-	34	+
35	if env_fallback_list:	35	if env_fallback_list:
36	for op_name in env_fallback_list.split(','):	36	for op_name in env_fallback_list.split(','):
37	op_name = op_name.strip()	37	op_name = op_name.strip()
	@@ -56,7 +56,7 @@ def _register_npu_inductor_fallbacks():
56	for overload in fn.overloads():	56	for overload in fn.overloads():
57	other_fn = getattr(fn, overload)	57	other_fn = getattr(fn, overload)
58	fallback_set.add(other_fn)	58	fallback_set.add(other_fn)
59	-	59	+
60	def fallback_except_gen_set(gen_set):	60	def fallback_except_gen_set(gen_set):
61	for op in lowering.lowerings:	61	for op in lowering.lowerings:
62	if op not in decomposition.decompositions and op not in gen_set:	62	if op not in decomposition.decompositions and op not in gen_set:
	@@ -70,7 +70,7 @@ def _register_npu_inductor_fallbacks():
70	if isinstance(op, torch._ops.OpOverloadPacket) or \	70	if isinstance(op, torch._ops.OpOverloadPacket) or \
71	isinstance(op, (torch._ops.OpOverload, torch._ops.HigherOrderOperator)):	71	isinstance(op, (torch._ops.OpOverload, torch._ops.HigherOrderOperator)):
72	make_fallback(op)	72	make_fallback(op)
73	-	73	+
74	def enable_full_lowering_fallback():	74	def enable_full_lowering_fallback():
75	ops_to_fallback = list(filter(	75	ops_to_fallback = list(filter(
76	lambda op: op not in decomposition.decompositions and	76	lambda op: op not in decomposition.decompositions and
	@@ -81,15 +81,15 @@ def _register_npu_inductor_fallbacks():
81	make_fallback(op)	81	make_fallback(op)
82		82
83	_fallback_ops_with_meta()	83	_fallback_ops_with_meta()
84	-	84	+
85	if config.fallback_to_aten_mode not in {"off", "include", "exclude"}:	85	if config.fallback_to_aten_mode not in {"off", "include", "exclude"}:
86	raise AssertionError(f"Error! Unsupported fallback_to_aten_mode: {config.fallback_to_aten_mode} was set!")	86	raise AssertionError(f"Error! Unsupported fallback_to_aten_mode: {config.fallback_to_aten_mode} was set!")
87	-	87	+
88	if get_anir_mode() == 'O0':	88	if get_anir_mode() == 'O0':
89	fallback_except_gen_set(gen_set=[])	89	fallback_except_gen_set(gen_set=[])
90	decomposition.decompositions.clear()	90	decomposition.decompositions.clear()
91	- return	91	+ return
92	-	92	+
93	if config.fallback_to_aten_mode == 'include':	93	if config.fallback_to_aten_mode == 'include':
94	fallback_via_fallback_set(fallback_set=fallback_set)	94	fallback_via_fallback_set(fallback_set=fallback_set)
95	elif config.fallback_to_aten_mode == 'exclude':	95	elif config.fallback_to_aten_mode == 'exclude':
	@@ -104,7 +104,7 @@ def get_nested_attr(obj, attr_path, default=None):
104	return reduce(getattr, attr_path.split('.'), obj)	104	return reduce(getattr, attr_path.split('.'), obj)
105	except AttributeError:	105	except AttributeError:
106	return default	106	return default
107	-	107	+
108		108
109	def _fallback_ops_with_meta():	109	def _fallback_ops_with_meta():
110	"""	110	"""

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/npu_patch_deprecated.py+1-2

	@@ -50,7 +50,7 @@ def _patch_add_ephemeral_timeout_for_all_pgs(timeout: timedelta) -> None:
50	devices = pg._device_types	50	devices = pg._device_types
51	if torch.device("npu") in devices:	51	if torch.device("npu") in devices:
52	backend = pg._get_backend(torch.device("npu"))	52	backend = pg._get_backend(torch.device("npu"))
53	-	53	+
54	distributed_c10d._add_ephemeral_timeout_for_all_pgs = _patch_add_ephemeral_timeout_for_all_pgs	54	distributed_c10d._add_ephemeral_timeout_for_all_pgs = _patch_add_ephemeral_timeout_for_all_pgs
55		55
56	if get_anir_mode() == 'O0':	56	if get_anir_mode() == 'O0':
	@@ -58,4 +58,3 @@ if get_anir_mode() == 'O0':
58	def my_silu_backward(grad_out, self):	58	def my_silu_backward(grad_out, self):
59	# use with some caution: this is only really valid to run in the context of proxy tensor tracing	59	# use with some caution: this is only really valid to run in the context of proxy tensor tracing
60	return NotImplemented	60	return NotImplemented
61	-

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/npu_stream.py+89-89

	@@ -60,7 +60,7 @@ class StreamResgistrator:
60	NPU_EVENTS[tag] = event	60	NPU_EVENTS[tag] = event
61		61
62	def npu_set_stream(	62	def npu_set_stream(
63	- dependency: Sequence[torch.Tensor],	63	+ dependency: Sequence[torch.Tensor],
64	stream_tag: str,	64	stream_tag: str,
65	) -> List[torch.Tensor]:	65	) -> List[torch.Tensor]:
66	stream = NPU_STREAMS[stream_tag]	66	stream = NPU_STREAMS[stream_tag]
	@@ -68,7 +68,7 @@ def npu_set_stream(
68	return dependency	68	return dependency
69		69
70	def npu_set_stream_fake(	70	def npu_set_stream_fake(
71	- dependency: Sequence[torch.Tensor],	71	+ dependency: Sequence[torch.Tensor],
72	stream_tag: str,	72	stream_tag: str,
73	) -> List[torch.Tensor]:	73	) -> List[torch.Tensor]:
74	return dependency	74	return dependency
	@@ -82,7 +82,7 @@ direct_register_custom_op(
82	)	82	)
83		83
84	def npu_event_record(	84	def npu_event_record(
85	- dependency: Sequence[torch.Tensor],	85	+ dependency: Sequence[torch.Tensor],
86	event_tag: str,	86	event_tag: str,
87	stream_tag: str	87	stream_tag: str
88	) -> List[torch.Tensor]:	88	) -> List[torch.Tensor]:
	@@ -92,7 +92,7 @@ def npu_event_record(
92	return dependency	92	return dependency
93		93
94	def npu_event_record_fake(	94	def npu_event_record_fake(
95	- dependency: Sequence[torch.Tensor],	95	+ dependency: Sequence[torch.Tensor],
96	event_tag: str,	96	event_tag: str,
97	stream_tag: str	97	stream_tag: str
98	) -> List[torch.Tensor]:	98	) -> List[torch.Tensor]:
	@@ -107,7 +107,7 @@ direct_register_custom_op(
107	)	107	)
108		108
109	def npu_event_wait(	109	def npu_event_wait(
110	- dependency: Sequence[torch.Tensor],	110	+ dependency: Sequence[torch.Tensor],
111	event_tag: str,	111	event_tag: str,
112	) -> List[torch.Tensor]:	112	) -> List[torch.Tensor]:
113	event = NPU_EVENTS[event_tag]	113	event = NPU_EVENTS[event_tag]
	@@ -115,7 +115,7 @@ def npu_event_wait(
115	return dependency	115	return dependency
116		116
117	def npu_event_wait_fake(	117	def npu_event_wait_fake(
118	- dependency: Sequence[torch.Tensor],	118	+ dependency: Sequence[torch.Tensor],
119	event_tag: str,	119	event_tag: str,
120	) -> List[torch.Tensor]:	120	) -> List[torch.Tensor]:
121	return dependency	121	return dependency
	@@ -129,12 +129,12 @@ direct_register_custom_op(
129	)	129	)
130		130
131	def graph_break(	131	def graph_break(
132	- dependency: Sequence[torch.Tensor],	132	+ dependency: Sequence[torch.Tensor],
133	) -> List[torch.Tensor]:	133	) -> List[torch.Tensor]:
134	return dependency	134	return dependency
135		135
136	def graph_break_fake(	136	def graph_break_fake(
137	- dependency: Sequence[torch.Tensor],	137	+ dependency: Sequence[torch.Tensor],
138	) -> List[torch.Tensor]:	138	) -> List[torch.Tensor]:
139	return dependency	139	return dependency
140		140
	@@ -150,7 +150,7 @@ direct_register_custom_op(
150	)	150	)
151		151
152	def npu_wait_stream(	152	def npu_wait_stream(
153	- dependency: Sequence[torch.Tensor],	153	+ dependency: Sequence[torch.Tensor],
154	stream1_tag: str,	154	stream1_tag: str,
155	stream2_tag: str,	155	stream2_tag: str,
156	) -> List[torch.Tensor]:	156	) -> List[torch.Tensor]:
	@@ -160,7 +160,7 @@ def npu_wait_stream(
160	return dependency	160	return dependency
161		161
162	def npu_wait_stream_fake(	162	def npu_wait_stream_fake(
163	- dependency: Sequence[torch.Tensor],	163	+ dependency: Sequence[torch.Tensor],
164	stream1_tag: str,	164	stream1_tag: str,
165	stream2_tag: str,	165	stream2_tag: str,
166	) -> List[torch.Tensor]:	166	) -> List[torch.Tensor]:
	@@ -186,48 +186,48 @@ def graph_break(
186	inductor_npu_lib = Library("inductor_npu", "FRAGMENT") # noqa	186	inductor_npu_lib = Library("inductor_npu", "FRAGMENT") # noqa
187		187
188	def npu_fusion_attention(	188	def npu_fusion_attention(
189	- query: torch.Tensor,	189	+ query: torch.Tensor,
190	- key: torch.Tensor,	190	+ key: torch.Tensor,
191	- value: torch.Tensor,	191	+ value: torch.Tensor,
192	- head_num: int,	192	+ head_num: int,
193	- input_layout: str,	193	+ input_layout: str,
194	pse: Optional[torch.Tensor] = None,	194	pse: Optional[torch.Tensor] = None,
195	padding_mask: Optional[torch.Tensor] = None,	195	padding_mask: Optional[torch.Tensor] = None,
196	atten_mask: Optional[torch.Tensor] = None,	196	atten_mask: Optional[torch.Tensor] = None,
197	- scale: float = 1.0,	197	+ scale: float = 1.0,
198	- keep_prob: float = 1.0,	198	+ keep_prob: float = 1.0,
199	- pre_tockens: int = 2147483647,	199	+ pre_tockens: int = 2147483647,
200	next_tockens: int = 2147483647,	200	next_tockens: int = 2147483647,
201	- inner_precise: int = 0,	201	+ inner_precise: int = 0,
202	- prefix: Optional[torch.Tensor] = None,	202	+ prefix: Optional[torch.Tensor] = None,
203	- actual_seq_qlen: Optional[torch.Tensor] = None,	203	+ actual_seq_qlen: Optional[torch.Tensor] = None,
204	- actual_seq_kvlen: Optional[torch.Tensor] = None,	204	+ actual_seq_kvlen: Optional[torch.Tensor] = None,
205	sparse_mode: int = 0,	205	sparse_mode: int = 0,
206	- gen_mask_parallel: bool = True,	206	+ gen_mask_parallel: bool = True,
207	sync: bool = False	207	sync: bool = False
208	) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:	208	) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
209	prefix = prefix.tolist() if prefix is not None else prefix	209	prefix = prefix.tolist() if prefix is not None else prefix
210	actual_seq_qlen = actual_seq_qlen.tolist() if actual_seq_qlen is not None else actual_seq_qlen	210	actual_seq_qlen = actual_seq_qlen.tolist() if actual_seq_qlen is not None else actual_seq_qlen
211	actual_seq_kvlen = actual_seq_kvlen.tolist() if actual_seq_kvlen is not None else actual_seq_kvlen	211	actual_seq_kvlen = actual_seq_kvlen.tolist() if actual_seq_kvlen is not None else actual_seq_kvlen
212	attention_score, softmax_max, softmax_sum, softmax_out, seed, offset, numels = torch.ops.npu.npu_fusion_attention(	212	attention_score, softmax_max, softmax_sum, softmax_out, seed, offset, numels = torch.ops.npu.npu_fusion_attention(
213	- query,	213	+ query,
214	- key,	214	+ key,
215	- value,	215	+ value,
216	- head_num,	216	+ head_num,
217	- input_layout,	217	+ input_layout,
218	pse=pse,	218	pse=pse,
219	padding_mask=padding_mask,	219	padding_mask=padding_mask,
220	atten_mask=atten_mask,	220	atten_mask=atten_mask,
221	- scale=scale,	221	+ scale=scale,
222	- keep_prob=keep_prob,	222	+ keep_prob=keep_prob,
223	- pre_tockens=pre_tockens,	223	+ pre_tockens=pre_tockens,
224	next_tockens=next_tockens,	224	next_tockens=next_tockens,
225	- inner_precise=inner_precise,	225	+ inner_precise=inner_precise,
226	- prefix=prefix,	226	+ prefix=prefix,
227	- actual_seq_qlen=actual_seq_qlen,	227	+ actual_seq_qlen=actual_seq_qlen,
228	- actual_seq_kvlen=actual_seq_kvlen,	228	+ actual_seq_kvlen=actual_seq_kvlen,
229	sparse_mode=sparse_mode,	229	sparse_mode=sparse_mode,
230	- gen_mask_parallel=gen_mask_parallel,	230	+ gen_mask_parallel=gen_mask_parallel,
231	sync=sync	231	sync=sync
232	)	232	)
233		233
	@@ -238,24 +238,24 @@ def npu_fusion_attention(
238	return attention_score, softmax_max, softmax_sum, softmax_out, seed, offset, numels	238	return attention_score, softmax_max, softmax_sum, softmax_out, seed, offset, numels
239		239
240	def npu_fusion_attention_fake(	240	def npu_fusion_attention_fake(
241	- query: torch.Tensor,	241	+ query: torch.Tensor,
242	- key: torch.Tensor,	242	+ key: torch.Tensor,
243	- value: torch.Tensor,	243	+ value: torch.Tensor,
244	- head_num: int,	244	+ head_num: int,
245	- input_layout: str,	245	+ input_layout: str,
246	pse: Optional[torch.Tensor] = None,	246	pse: Optional[torch.Tensor] = None,
247	padding_mask: Optional[torch.Tensor] = None,	247	padding_mask: Optional[torch.Tensor] = None,
248	atten_mask: Optional[torch.Tensor] = None,	248	atten_mask: Optional[torch.Tensor] = None,
249	- scale: float = 1.0,	249	+ scale: float = 1.0,
250	- keep_prob: float = 1.0,	250	+ keep_prob: float = 1.0,
251	- pre_tockens: int = 2147483647,	251	+ pre_tockens: int = 2147483647,
252	next_tockens: int = 2147483647,	252	next_tockens: int = 2147483647,
253	- inner_precise: int = 0,	253	+ inner_precise: int = 0,
254	- prefix: Optional[torch.Tensor] = None,	254	+ prefix: Optional[torch.Tensor] = None,
255	- actual_seq_qlen: Optional[torch.Tensor] = None,	255	+ actual_seq_qlen: Optional[torch.Tensor] = None,
256	- actual_seq_kvlen: Optional[torch.Tensor] = None,	256	+ actual_seq_kvlen: Optional[torch.Tensor] = None,
257	sparse_mode: int = 0,	257	sparse_mode: int = 0,
258	- gen_mask_parallel: bool = True,	258	+ gen_mask_parallel: bool = True,
259	sync: bool = False	259	sync: bool = False
260	) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:	260	) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
261	B = query.size(0)	261	B = query.size(0)
	@@ -301,32 +301,32 @@ direct_register_custom_op(
301		301
302	def npu_fusion_attention_grad(	302	def npu_fusion_attention_grad(
303	query: torch.Tensor,	303	query: torch.Tensor,
304	- key: torch.Tensor,	304	+ key: torch.Tensor,
305	value: torch.Tensor,	305	value: torch.Tensor,
306	- dy: torch.Tensor,	306	+ dy: torch.Tensor,
307	- head_num: int,	307	+ head_num: int,
308	- input_layout: str,	308	+ input_layout: str,
309	- *,	309	+ *,
310	- pse: Optional[torch.Tensor] = None,	310	+ pse: Optional[torch.Tensor] = None,
311	- padding_mask: Optional[torch.Tensor] = None,	311	+ padding_mask: Optional[torch.Tensor] = None,
312	atten_mask: Optional[torch.Tensor] = None,	312	atten_mask: Optional[torch.Tensor] = None,
313	- softmax_max: Optional[torch.Tensor] = None,	313	+ softmax_max: Optional[torch.Tensor] = None,
314	- softmax_sum: Optional[torch.Tensor] = None,	314	+ softmax_sum: Optional[torch.Tensor] = None,
315	- softmax_in: Optional[torch.Tensor] = None,	315	+ softmax_in: Optional[torch.Tensor] = None,
316	- attention_in: Optional[torch.Tensor] = None,	316	+ attention_in: Optional[torch.Tensor] = None,
317	scale_value: float = 1.0,	317	scale_value: float = 1.0,
318	- keep_prob: float = 1.0,	318	+ keep_prob: float = 1.0,
319	- pre_tockens: int = 2147483647,	319	+ pre_tockens: int = 2147483647,
320	- next_tockens: int = 2147483647,	320	+ next_tockens: int = 2147483647,
321	- inner_precise: int = 0,	321	+ inner_precise: int = 0,
322	- seed: Optional[torch.Tensor] = None,	322	+ seed: Optional[torch.Tensor] = None,
323	offset: Optional[torch.Tensor] = None,	323	offset: Optional[torch.Tensor] = None,
324	- numels: Optional[torch.Tensor] = None,	324	+ numels: Optional[torch.Tensor] = None,
325	- prefix: Optional[torch.Tensor] = None,	325	+ prefix: Optional[torch.Tensor] = None,
326	actual_seq_qlen: Optional[torch.Tensor] = None,	326	actual_seq_qlen: Optional[torch.Tensor] = None,
327	- actual_seq_kvlen: Optional[torch.Tensor] = None,	327	+ actual_seq_kvlen: Optional[torch.Tensor] = None,
328	sparse_mode: int = 0,	328	sparse_mode: int = 0,
329	- gen_mask_parallel: bool = True,	329	+ gen_mask_parallel: bool = True,
330	sync: bool = False	330	sync: bool = False
331	) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:	331	) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
332	prefix = prefix.tolist() if prefix is not None else prefix	332	prefix = prefix.tolist() if prefix is not None else prefix
	@@ -349,32 +349,32 @@ def npu_fusion_attention_grad(
349		349
350	def npu_fusion_attention_grad_fake(	350	def npu_fusion_attention_grad_fake(
351	query: torch.Tensor,	351	query: torch.Tensor,
352	- key: torch.Tensor,	352	+ key: torch.Tensor,
353	value: torch.Tensor,	353	value: torch.Tensor,
354	- dy: torch.Tensor,	354	+ dy: torch.Tensor,
355	- head_num: int,	355	+ head_num: int,
356	- input_layout: str,	356	+ input_layout: str,
357	- *,	357	+ *,
358	- pse: Optional[torch.Tensor] = None,	358	+ pse: Optional[torch.Tensor] = None,
359	- padding_mask: Optional[torch.Tensor] = None,	359	+ padding_mask: Optional[torch.Tensor] = None,
360	atten_mask: Optional[torch.Tensor] = None,	360	atten_mask: Optional[torch.Tensor] = None,
361	- softmax_max: Optional[torch.Tensor] = None,	361	+ softmax_max: Optional[torch.Tensor] = None,
362	- softmax_sum: Optional[torch.Tensor] = None,	362	+ softmax_sum: Optional[torch.Tensor] = None,
363	- softmax_in: Optional[torch.Tensor] = None,	363	+ softmax_in: Optional[torch.Tensor] = None,
364	- attention_in: Optional[torch.Tensor] = None,	364	+ attention_in: Optional[torch.Tensor] = None,
365	scale_value: float = 1.0,	365	scale_value: float = 1.0,
366	- keep_prob: float = 1.0,	366	+ keep_prob: float = 1.0,
367	- pre_tockens: int = 2147483647,	367	+ pre_tockens: int = 2147483647,
368	- next_tockens: int = 2147483647,	368	+ next_tockens: int = 2147483647,
369	- inner_precise: int = 0,	369	+ inner_precise: int = 0,
370	- seed: Optional[torch.Tensor] = None,	370	+ seed: Optional[torch.Tensor] = None,
371	offset: Optional[torch.Tensor] = None,	371	offset: Optional[torch.Tensor] = None,
372	- numels: Optional[torch.Tensor] = None,	372	+ numels: Optional[torch.Tensor] = None,
373	prefix: Optional[torch.Tensor] = None,	373	prefix: Optional[torch.Tensor] = None,
374	actual_seq_qlen: Optional[torch.Tensor] = None,	374	actual_seq_qlen: Optional[torch.Tensor] = None,
375	- actual_seq_kvlen: Optional[torch.Tensor] = None,	375	+ actual_seq_kvlen: Optional[torch.Tensor] = None,
376	sparse_mode: int = 0,	376	sparse_mode: int = 0,
377	- gen_mask_parallel: bool = True,	377	+ gen_mask_parallel: bool = True,
378	sync: bool = False	378	sync: bool = False
379	) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:	379	) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
380	dq = torch.empty_like(query, dtype=query.dtype, device=query.device).contiguous()	380	dq = torch.empty_like(query, dtype=query.dtype, device=query.device).contiguous()
	@@ -436,7 +436,7 @@ class InductorNpuAttentionFunction(torch.autograd.Function):
436	return (	436	return (
437	grad_query, grad_key, grad_value, None, None, grad_pse, None, None, None, None, None, None, None, None, None,	437	grad_query, grad_key, grad_value, None, None, grad_pse, None, None, None, None, None, None, None, None, None,
438	None, None, None, None, None, None, None, None, None, None, None)	438	None, None, None, None, None, None, None, None, None, None, None)
439	-	439	+
440	def inductor_npu_fusion_attention(query, key, value, head_num, input_layout, pse=None, padding_mask=None,	440	def inductor_npu_fusion_attention(query, key, value, head_num, input_layout, pse=None, padding_mask=None,
441	atten_mask=None, scale=1.0, keep_prob=1.0, pre_tockens=2147483647,	441	atten_mask=None, scale=1.0, keep_prob=1.0, pre_tockens=2147483647,
442	next_tockens=2147483647,	442	next_tockens=2147483647,

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/torch_mlir_patch.py+1-1

	@@ -66,7 +66,7 @@ def _patch_import_stateless_graph(
66	range_constraints = {}	66	range_constraints = {}
67	for nd in graph.find_nodes(	67	for nd in graph.find_nodes(
68	op="placeholder"	68	op="placeholder"
69	- ):	69	+ ):
70	if isinstance(nd.meta['val'], torch.Tensor):	70	if isinstance(nd.meta['val'], torch.Tensor):
71	for s in nd.meta['val'].size():	71	for s in nd.meta['val'].size():
72	if isinstance(s, torch.SymInt):	72	if isinstance(s, torch.SymInt):

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/utils.py+55-55

	@@ -68,7 +68,7 @@ def get_device_info(example_inputs) -> Union[Tuple[str, int], None]:
68	for inp in example_inputs:	68	for inp in example_inputs:
69	if isinstance(inp, torch.Tensor):	69	if isinstance(inp, torch.Tensor):
70	return inp.device, inp.device.index	70	return inp.device, inp.device.index
71	-	71	+
72	@functools.lru_cache(None)	72	@functools.lru_cache(None)
73	def _get_ascend_path() -> str:	73	def _get_ascend_path() -> str:
74	path = os.getenv("ASCEND_HOME_PATH", "")	74	path = os.getenv("ASCEND_HOME_PATH", "")
	@@ -102,11 +102,11 @@ def _build_npu_ext(obj_name: str, src_path, src_dir) -> str:
102		102
103	cc_cmd += [f"-I{py_include_dir}"]	103	cc_cmd += [f"-I{py_include_dir}"]
104	torch_npu_root = Path(torch_npu.__file__).resolve().parent	104	torch_npu_root = Path(torch_npu.__file__).resolve().parent
105	-	105	+
106	cpp_common_dir = (	106	cpp_common_dir = (
107	torch_npu_root / "include" / "torch_npu" / "csrc" / "inductor" / "mlir"	107	torch_npu_root / "include" / "torch_npu" / "csrc" / "inductor" / "mlir"
108	)	108	)
109	-	109	+
110	torch_npu_dir = torch_npu_root / "include"	110	torch_npu_dir = torch_npu_root / "include"
111	torch_npu_lib_dir = torch_npu_root / "lib"	111	torch_npu_lib_dir = torch_npu_root / "lib"
112		112
	@@ -158,7 +158,7 @@ import torch._inductor.inductor_prims
158	return model_str	158	return model_str
159		159
160		160
161	-def get_fx_graph_code(code, num_args, method=2, runnable=False, kernel_code='', kernel_name=None):	161	+def get_fx_graph_code(code, num_args, method=2, runnable=False, kernel_code='', kernel_name=None):
162	kernel_header = ''	162	kernel_header = ''
163	kernel_wrapper = ''	163	kernel_wrapper = ''
164	kernel_runner_and_acc_comp = ''	164	kernel_runner_and_acc_comp = ''
	@@ -214,7 +214,7 @@ def get_args():
214	"""	214	"""
215	run_code_template = f"""	215	run_code_template = f"""
216		216
217	-try:	217	+try:
218	args = torch.load(os.path.join(dir_path, "data.pth"))	218	args = torch.load(os.path.join(dir_path, "data.pth"))
219	except Exception as e:	219	except Exception as e:
220	{{{{FAKE_ARGS_PLACEHOLDER}}}}	220	{{{{FAKE_ARGS_PLACEHOLDER}}}}
	@@ -229,16 +229,16 @@ with torch.no_grad():
229	output2 = model(*fx_inputs)	229	output2 = model(*fx_inputs)
230	"""	230	"""
231	code_template = f"""	231	code_template = f"""
232	-import os	232	+import os
233	import torch	233	import torch
234	from torch._inductor.compile_fx import clone_preserve_strides	234	from torch._inductor.compile_fx import clone_preserve_strides
235	from torch._dynamo.testing import rand_strided	235	from torch._dynamo.testing import rand_strided
236	from torch import device	236	from torch import device
237		237
238	import torch_npu	238	import torch_npu
239	-from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import config as npu_config	239	+from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import config as npu_config
240	{kernel_header}	240	{kernel_header}
241	-file_path = os.path.abspath(__file__)	241	+file_path = os.path.abspath(__file__)
242	dir_path = os.path.dirname(file_path)	242	dir_path = os.path.dirname(file_path)
243		243
244	{kernel_code}	244	{kernel_code}
	@@ -250,7 +250,7 @@ class GraphModule(torch.nn.Module):
250	{code}	250	{code}
251	model = GraphModule().npu()	251	model = GraphModule().npu()
252		252
253	-{run_code_template if runnable else transformed_code_template}	253	+{run_code_template if runnable else transformed_code_template}
254	{fx_runner if runnable else ''}	254	{fx_runner if runnable else ''}
255	{kernel_runner_and_acc_comp if runnable else ''}	255	{kernel_runner_and_acc_comp if runnable else ''}
256	"""	256	"""
	@@ -270,28 +270,28 @@ def view_to_reshape(gm: torch.fx.GraphModule):
270	op="call_function", target=torch.ops.aten.view.default	270	op="call_function", target=torch.ops.aten.view.default
271	):	271	):
272	nd.target = torch.ops.aten.reshape.default	272	nd.target = torch.ops.aten.reshape.default
273	-	273	+
274	for nd in gm.graph.find_nodes(	274	for nd in gm.graph.find_nodes(
275	op="call_function", target=torch.ops.aten.div.Tensor	275	op="call_function", target=torch.ops.aten.div.Tensor
276	):	276	):
277	if not (isinstance(nd.args[1], torch.fx.node.Node) and \	277	if not (isinstance(nd.args[1], torch.fx.node.Node) and \
278	isinstance(nd.args[1].meta['val'], torch.Tensor)):	278	isinstance(nd.args[1].meta['val'], torch.Tensor)):
279	nd.target = torch.ops.aten.div.Scalar	279	nd.target = torch.ops.aten.div.Scalar
280	-	280	+
281	for nd in gm.graph.find_nodes(	281	for nd in gm.graph.find_nodes(
282	op="call_function", target=torch.ops.aten.add.Tensor	282	op="call_function", target=torch.ops.aten.add.Tensor
283	):	283	):
284	if not (isinstance(nd.args[1], torch.fx.node.Node) and \	284	if not (isinstance(nd.args[1], torch.fx.node.Node) and \
285	isinstance(nd.args[1].meta['val'], torch.Tensor)):	285	isinstance(nd.args[1].meta['val'], torch.Tensor)):
286	nd.target = torch.ops.aten.add.Scalar	286	nd.target = torch.ops.aten.add.Scalar
287	-	287	+
288	for nd in gm.graph.find_nodes(	288	for nd in gm.graph.find_nodes(
289	op="call_function", target=torch.ops.aten.sub.Tensor	289	op="call_function", target=torch.ops.aten.sub.Tensor
290	- ):	290	+ ):
291	if not (isinstance(nd.args[1], torch.fx.node.Node) and \	291	if not (isinstance(nd.args[1], torch.fx.node.Node) and \
292	isinstance(nd.args[1].meta['val'], torch.Tensor)):	292	isinstance(nd.args[1].meta['val'], torch.Tensor)):
293	nd.target = torch.ops.aten.sub.Scalar	293	nd.target = torch.ops.aten.sub.Scalar
294	-	294	+
295	for nd in gm.graph.find_nodes(	295	for nd in gm.graph.find_nodes(
296	op="call_function", target=torch.ops.aten.mul.Tensor	296	op="call_function", target=torch.ops.aten.mul.Tensor
297	):	297	):
	@@ -303,7 +303,7 @@ def view_to_reshape(gm: torch.fx.GraphModule):
303	op="call_function", target=torch.ops.prims.convert_element_type.default	303	op="call_function", target=torch.ops.prims.convert_element_type.default
304	):	304	):
305	nd.target = torch.ops.npu.npu_dtype_cast.default	305	nd.target = torch.ops.npu.npu_dtype_cast.default
306	-	306	+
307	def npu_cast_to_prim_cast(gm: torch.fx.GraphModule):	307	def npu_cast_to_prim_cast(gm: torch.fx.GraphModule):
308	"""	308	"""
309	Replace npu.npu_dtype_cast ops in the GraphModule to prims.convert_element_type ops.	309	Replace npu.npu_dtype_cast ops in the GraphModule to prims.convert_element_type ops.
	@@ -359,7 +359,7 @@ def npu_optimize_fx_graph(gm: torch.fx.GraphModule):
359	gm.graph.erase_node(nd)	359	gm.graph.erase_node(nd)
360	aten_empty_nodes.remove(node0)	360	aten_empty_nodes.remove(node0)
361	gm.graph.erase_node(node0)	361	gm.graph.erase_node(node0)
362	-	362	+
363	gm.recompile()	363	gm.recompile()
364		364
365		365
	@@ -374,7 +374,7 @@ def fold_expand(gm: torch.fx.GraphModule) -> None:
374		374
375	inp0 = node.args[0] if len(node.args) > 0 else None	375	inp0 = node.args[0] if len(node.args) > 0 else None
376	inp1 = node.args[1] if len(node.args) > 1 else None	376	inp1 = node.args[1] if len(node.args) > 1 else None
377	- if (isinstance(inp0, torch.fx.Node) and inp0.op == 'call_function' and	377	+ if (isinstance(inp0, torch.fx.Node) and inp0.op == 'call_function' and
378	inp0.target == torch.ops.aten.expand.default):	378	inp0.target == torch.ops.aten.expand.default):
379	if len(inp0.args) > 0:	379	if len(inp0.args) > 0:
380	expand_input = inp0.args[0]	380	expand_input = inp0.args[0]
	@@ -382,7 +382,7 @@ def fold_expand(gm: torch.fx.GraphModule) -> None:
382	if len(inp0.users) == 0:	382	if len(inp0.users) == 0:
383	graph.erase_node(inp0)	383	graph.erase_node(inp0)
384	changed = True	384	changed = True
385	- elif (isinstance(inp1, torch.fx.Node) and inp1.op == 'call_function' and	385	+ elif (isinstance(inp1, torch.fx.Node) and inp1.op == 'call_function' and
386	inp1.target == torch.ops.aten.expand.default):	386	inp1.target == torch.ops.aten.expand.default):
387	if len(inp1.args) > 0:	387	if len(inp1.args) > 0:
388	expand_input = inp1.args[0]	388	expand_input = inp1.args[0]
	@@ -393,7 +393,7 @@ def fold_expand(gm: torch.fx.GraphModule) -> None:
393	if changed:	393	if changed:
394	graph.lint()	394	graph.lint()
395	graph.eliminate_dead_code()	395	graph.eliminate_dead_code()
396	-	396	+
397	gm.recompile()	397	gm.recompile()
398		398
399		399
	@@ -466,12 +466,12 @@ class MLIRProcessor:
466	def __init__(self, bisheng_install_path: str = None):	466	def __init__(self, bisheng_install_path: str = None):
467	"""	467	"""
468	初始化MLIR处理器	468	初始化MLIR处理器
469	-	469	+
470	:param bisheng_install_path: Bisheng安装路径，默认从环境变量获取	470	:param bisheng_install_path: Bisheng安装路径，默认从环境变量获取
471	"""	471	"""
472	bisheng_install_path = os.getenv('BISHENG_INSTALL_PATH', '')	472	bisheng_install_path = os.getenv('BISHENG_INSTALL_PATH', '')
473	self.bisheng_torch_mlir_path = os.path.join(bisheng_install_path, "bishengir-opt")	473	self.bisheng_torch_mlir_path = os.path.join(bisheng_install_path, "bishengir-opt")
474	-	474	+
475	def extract_function(self, module: Any) -> Any:	475	def extract_function(self, module: Any) -> Any:
476	"""从MLIR模块中提取主函数并添加标记属性"""	476	"""从MLIR模块中提取主函数并添加标记属性"""
477	with module.context:	477	with module.context:
	@@ -480,20 +480,20 @@ class MLIRProcessor:
480	func.attributes["hacc.placeholder"] = ir.UnitAttr.get(func.context)	480	func.attributes["hacc.placeholder"] = ir.UnitAttr.get(func.context)
481	return func	481	return func
482	raise ValueError("No valid FuncOp found in module")	482	raise ValueError("No valid FuncOp found in module")
483	-	483	+
484	def rebuild_mlir_module(self, module_str: str) -> Any:	484	def rebuild_mlir_module(self, module_str: str) -> Any:
485	"""从字符串重新构建MLIR模块"""	485	"""从字符串重新构建MLIR模块"""
486	with ir.Context() as ctx:	486	with ir.Context() as ctx:
487	ctx.allow_unregistered_dialects = True	487	ctx.allow_unregistered_dialects = True
488	torch_mlir.dialects.torch.register_dialect(ctx)	488	torch_mlir.dialects.torch.register_dialect(ctx)
489	return ir.Module.parse(module_str)	489	return ir.Module.parse(module_str)
490	-	490	+
491	def get_signature(self, func: Any) -> tuple:	491	def get_signature(self, func: Any) -> tuple:
492	"""获取函数的签名信息：类型签名、输出数量和张量维度"""	492	"""获取函数的签名信息：类型签名、输出数量和张量维度"""
493	func_type = func.type	493	func_type = func.type
494	signature = {}	494	signature = {}
495	ranks = []	495	ranks = []
496	-	496	+
497	# 处理输入+输出类型	497	# 处理输入+输出类型
498	for i, tensor_type in enumerate(func_type.inputs + func_type.results):	498	for i, tensor_type in enumerate(func_type.inputs + func_type.results):
499	try: # RankedTensorType	499	try: # RankedTensorType
	@@ -507,17 +507,17 @@ class MLIRProcessor:
507	dim_end = type_str.find(']', dim_start)	507	dim_end = type_str.find(']', dim_start)
508	dim_str = type_str[dim_start:dim_end]	508	dim_str = type_str[dim_start:dim_end]
509	ranks.append(dim_str.count(',') + 1 if dim_str else 0)	509	ranks.append(dim_str.count(',') + 1 if dim_str else 0)
510	-	510	+
511	num_outputs = len(func_type.results)	511	num_outputs = len(func_type.results)
512	return signature, num_outputs, ranks	512	return signature, num_outputs, ranks
513	-	513	+
514	- def process_mlir(self,	514	+ def process_mlir(self,
515	- module: Union[str, Any],	515	+ module: Union[str, Any],
516	- get_sig: bool = True,	516	+ get_sig: bool = True,
517	dynamic: bool = False) -> tuple:	517	dynamic: bool = False) -> tuple:
518	"""	518	"""
519	处理MLIR模块的核心方法	519	处理MLIR模块的核心方法
520	-	520	+
521	:param module: MLIR模块字符串或对象	521	:param module: MLIR模块字符串或对象
522	:param get_sig: 是否获取函数签名	522	:param get_sig: 是否获取函数签名
523	:param dynamic: 是否为动态执行模式	523	:param dynamic: 是否为动态执行模式
	@@ -525,7 +525,7 @@ class MLIRProcessor:
525	"""	525	"""
526	if isinstance(module, str):	526	if isinstance(module, str):
527	module = self.rebuild_mlir_module(module)	527	module = self.rebuild_mlir_module(module)
528	-	528	+
529	func = self.extract_function(module)	529	func = self.extract_function(module)
530	kernel_info = None	530	kernel_info = None
531	func_str = str(func)	531	func_str = str(func)
	@@ -540,29 +540,29 @@ class MLIRProcessor:
540	"ranks": ranks,	540	"ranks": ranks,
541	'kernel_hash': module_hash,	541	'kernel_hash': module_hash,
542	}	542	}
543	-	543	+
544	return func_str, kernel_info	544	return func_str, kernel_info
545	-	545	+
546	def get_named_op_str(self,	546	def get_named_op_str(self,
547	module: Union[str, Any],	547	module: Union[str, Any],
548	kernel_name: str,	548	kernel_name: str,
549	dynamic: bool = False) -> Dict[str, Any]:	549	dynamic: bool = False) -> Dict[str, Any]:
550	"""	550	"""
551	获取命名操作格式的MLIR字符串	551	获取命名操作格式的MLIR字符串
552	-	552	+
553	:param module: MLIR模块字符串或对象	553	:param module: MLIR模块字符串或对象
554	:param kernel_name: 内核名称（用于临时文件）	554	:param kernel_name: 内核名称（用于临时文件）
555	:param dynamic: 是否为动态执行模式	555	:param dynamic: 是否为动态执行模式
556	:return: 包含处理结果和签名字典	556	:return: 包含处理结果和签名字典
557	"""	557	"""
558	func_str, sig_dict = self.process_mlir(module, get_sig=True, dynamic=dynamic)	558	func_str, sig_dict = self.process_mlir(module, get_sig=True, dynamic=dynamic)
559	-	559	+
560	cleaned_func = func_str.replace(	560	cleaned_func = func_str.replace(
561	- '"#hfusion.fusion_kind<PURE_ELEMWISE>"',	561	+ '"#hfusion.fusion_kind<PURE_ELEMWISE>"',
562	'#hfusion.fusion_kind<PURE_ELEMWISE>'	562	'#hfusion.fusion_kind<PURE_ELEMWISE>'
563	)	563	)
564	logger.debug(f"原始Linalg方言MLIR:\n{cleaned_func}")	564	logger.debug(f"原始Linalg方言MLIR:\n{cleaned_func}")
565	-	565	+
566	# 执行转换命令	566	# 执行转换命令
567	with tempfile.TemporaryDirectory() as tmpdir:	567	with tempfile.TemporaryDirectory() as tmpdir:
568	torch_mlir_path = os.path.join(tmpdir, f"{kernel_name}.mlir")	568	torch_mlir_path = os.path.join(tmpdir, f"{kernel_name}.mlir")
	@@ -573,33 +573,33 @@ class MLIRProcessor:
573	"--torch-backend-to-named-op-backend-pipeline="	573	"--torch-backend-to-named-op-backend-pipeline="
574	"\"ensure-no-implicit-broadcast=true\" "	574	"\"ensure-no-implicit-broadcast=true\" "
575	f"{torch_mlir_path}")	575	f"{torch_mlir_path}")
576	-	576	+
577	try:	577	try:
578	result = subprocess.check_output(	578	result = subprocess.check_output(
579	cmd, text=True, shell=True	579	cmd, text=True, shell=True
580	)	580	)
581	# 过滤全局定义并更新函数属性	581	# 过滤全局定义并更新函数属性
582	processed_mlir = "\n".join(	582	processed_mlir = "\n".join(
583	- line for line in result.splitlines()	583	+ line for line in result.splitlines()
584	if "ml_program.global" not in line	584	if "ml_program.global" not in line
585	)	585	)
586	-	586	+
587	# 根据模式设置函数属性	587	# 根据模式设置函数属性
588	- func_attr = ("hacc.entry, hacc.function_kind = #hacc.function_kind<HOST>"	588	+ func_attr = ("hacc.entry, hacc.function_kind = #hacc.function_kind<HOST>"
589	- if dynamic else	589	+ if dynamic else
590	"hacc.entry, hacc.function_kind = #hacc.function_kind<DEVICE>")	590	"hacc.entry, hacc.function_kind = #hacc.function_kind<DEVICE>")
591	processed_mlir = processed_mlir.replace("hacc.placeholder", func_attr)	591	processed_mlir = processed_mlir.replace("hacc.placeholder", func_attr)
592	-	592	+
593	# 应用额外的数据类型处理（需实现mlir_match_and_replace_unsupported_dtypes）	593	# 应用额外的数据类型处理（需实现mlir_match_and_replace_unsupported_dtypes）
594	final_mlir = self._replace_unsupported_dtypes(processed_mlir)	594	final_mlir = self._replace_unsupported_dtypes(processed_mlir)
595	logger.debug(f"转换后的NamedOp方言MLIR:\n{final_mlir}")	595	logger.debug(f"转换后的NamedOp方言MLIR:\n{final_mlir}")
596	-	596	+
597	return final_mlir, sig_dict	597	return final_mlir, sig_dict
598	-	598	+
599	except subprocess.CalledProcessError as e:	599	except subprocess.CalledProcessError as e:
600	logger.error(f"命令执行失败: {cmd}\n错误: {e.output}")	600	logger.error(f"命令执行失败: {cmd}\n错误: {e.output}")
601	raise RuntimeError(f"MLIR转换失败: {e.stderr}") from e	601	raise RuntimeError(f"MLIR转换失败: {e.stderr}") from e
602	-	602	+
603	def _replace_unsupported_dtypes(self, mlir_text: str) -> str:	603	def _replace_unsupported_dtypes(self, mlir_text: str) -> str:
604	"""替换不支持的MLIR数据类型"""	604	"""替换不支持的MLIR数据类型"""
605	pattern1 = r"%(\d+) = arith\.truncf %(\w+) : f64 to bf16"	605	pattern1 = r"%(\d+) = arith\.truncf %(\w+) : f64 to bf16"
	@@ -625,8 +625,8 @@ def mlir_match_and_replace_unsupported_dtypes(mlir_text: str) -> str:
625		625
626		626
627	def to_folder(	627	def to_folder(
628	- gm: torch.fx.GraphModule,	628	+ gm: torch.fx.GraphModule,
629	- folder: Union[str, os.PathLike],	629	+ folder: Union[str, os.PathLike],
630	graph_hash: str,	630	graph_hash: str,
631	module_name: str = "FxModule"):	631	module_name: str = "FxModule"):
632	"""Dumps out module to ``folder`` with ``module_name`` so that it can be	632	"""Dumps out module to ``folder`` with ``module_name`` so that it can be
	@@ -727,27 +727,27 @@ def is_fx_dynamic(graph):
727	def replace_placeholders(file_path: str, replacements: dict, placeholder_format: str = r'\{\{(\w+)\}\}') -> None:	727	def replace_placeholders(file_path: str, replacements: dict, placeholder_format: str = r'\{\{(\w+)\}\}') -> None:
728	"""	728	"""
729	替换文件中的占位符	729	替换文件中的占位符
730	-	730	+
731	:param file_path: 文件路径	731	:param file_path: 文件路径
732	:param replacements: 替换字典，如 {'function_body': 'your _code'}	732	:param replacements: 替换字典，如 {'function_body': 'your _code'}
733	:param placeholder_format: 占位符正则表达式（默认匹配{{xxx}}）	733	:param placeholder_format: 占位符正则表达式（默认匹配{{xxx}}）
734	"""	734	"""
735	with open(file_path, 'r', encoding='utf-8') as f:	735	with open(file_path, 'r', encoding='utf-8') as f:
736	content = f.read()	736	content = f.read()
737	-	737	+
738	pattern = re.compile(placeholder_format)	738	pattern = re.compile(placeholder_format)
739	-	739	+
740	def replacer(match: re.Match) -> str:	740	def replacer(match: re.Match) -> str:
741	placeholder = match.group(1)	741	placeholder = match.group(1)
742	replacement = replacements.get(placeholder, match.group(0))	742	replacement = replacements.get(placeholder, match.group(0))
743	-	743	+
744	line_start = content.rfind('\n', 0, match.start()) + 1	744	line_start = content.rfind('\n', 0, match.start()) + 1
745	indent = re.match(r'^\s*', content[line_start:match.start()]).group(0)	745	indent = re.match(r'^\s*', content[line_start:match.start()]).group(0)
746	-	746	+
747	return '\n'.join([indent + line for line in replacement.split('\n')])	747	return '\n'.join([indent + line for line in replacement.split('\n')])
748	-	748	+
749	new_content = pattern.sub(replacer, content)	749	new_content = pattern.sub(replacer, content)
750	-	750	+
751	with open(file_path, 'w', encoding='utf-8') as f:	751	with open(file_path, 'w', encoding='utf-8') as f:
752	f.write(new_content)	752	f.write(new_content)
753		753

Mtorch_npu/_inductor/codegen/catlass/catlass_kernel.py+1-1

	@@ -472,7 +472,7 @@ class CATLASSTemplateKernel(Kernel):
472	Mock load function for memory planning to optimize allocations properly.	472	Mock load function for memory planning to optimize allocations properly.
473	"""	473	"""
474	return self.create_cse_var(name, bounds=ValueRanges.unknown())	474	return self.create_cse_var(name, bounds=ValueRanges.unknown())
475	-	475	+
476	def store(self, name: str, index: Expr, value: Any, mode: Any = None) -> None:	476	def store(self, name: str, index: Expr, value: Any, mode: Any = None) -> None:
477	"""	477	"""
478	Mock store function for memory planning to optimize allocations properly.	478	Mock store function for memory planning to optimize allocations properly.

Mtorch_npu/_inductor/codegen/catlass/catlass_library/gemm_autotune.py+2-2

	@@ -79,7 +79,7 @@ class TileAutotune:
79	return 1 << (n.bit_length() - 1)	79	return 1 << (n.bit_length() - 1)
80	except AttributeError:	80	except AttributeError:
81	import sympy	81	import sympy
82	-	82	+
83	exp = sympy.floor(sympy.log(n, 2))	83	exp = sympy.floor(sympy.log(n, 2))
84	return 2 ** exp	84	return 2 ** exp
85		85
	@@ -300,7 +300,7 @@ class GemmAutotune:
300	if tile[1].k > tile[0].k:	300	if tile[1].k > tile[0].k:
301	tile[1].k = tile[0].k	301	tile[1].k = tile[0].k
302		302
303	-	303	+
304	def may_adjust_l1_tile_for_bias(self, dtype_size, tile):	304	def may_adjust_l1_tile_for_bias(self, dtype_size, tile):
305	# default l1 stages & size	305	# default l1 stages & size
306	l1_stages = 2	306	l1_stages = 2

Mtorch_npu/_inductor/codegen/catlass/catlass_python_evg.py+1-1

	@@ -51,7 +51,7 @@ class CatlassEVGOpsMixIn:
51	@staticmethod	51	@staticmethod
52	def constant(value: Any, dtype: Any) -> str:	52	def constant(value: Any, dtype: Any) -> str:
53	from catlass_cppgen.common.data_type import DataType	53	from catlass_cppgen.common.data_type import DataType
54	-	54	+
55	element = DataType.from_dtype(dtype).value	55	element = DataType.from_dtype(dtype).value
56	return CatlassEVGOpsMixIn._prefix_bin_op("constant", value, element)	56	return CatlassEVGOpsMixIn._prefix_bin_op("constant", value, element)
57		57

Mtorch_npu/_inductor/codegen/catlass/catlass_template.py+1-1

	@@ -226,7 +226,7 @@ class CATLASSTemplate(KernelTemplate):
226	return ptr	226	return ptr
227	else:	227	else:
228	return f"(uint8_t*)({ptr})"	228	return f"(uint8_t*)({ptr})"
229	-	229	+
230		230
231	def render(self, **kwargs) -> str:	231	def render(self, **kwargs) -> str:
232	raise NotImplementedError	232	raise NotImplementedError

Mtorch_npu/_inductor/codegen/catlass/catlass_utils.py+4-4

Mtorch_npu/_inductor/codegen/catlass/gemm_template.py+6-6

	@@ -54,7 +54,7 @@ PT_EXPORT {{kernel_call_signature}} {
54	{% endif %}	54	{% endif %}
55		55
56	{{op.gen_input_template()}}	56	{{op.gen_input_template()}}
57	-	57	+
58	{{evg_ptr}}	58	{{evg_ptr}}
59		59
60	{{evg_template}}	60	{{evg_template}}
	@@ -299,7 +299,7 @@ class CATLASSGemmTemplate(CATLASSTemplate, ABC):
299	return layout::RowMajor(layout.shape(0), layout.shape(1), RoundUp(layout.shape(1), align));	299	return layout::RowMajor(layout.shape(0), layout.shape(1), RoundUp(layout.shape(1), align));
300	}	300	}
301		301
302	-	302	+
303	layout::ColumnMajor GetWorkspaceLayout(layout::ColumnMajor layout, uint32_t align)	303	layout::ColumnMajor GetWorkspaceLayout(layout::ColumnMajor layout, uint32_t align)
304	{	304	{
305	if (align == 0) {	305	if (align == 0) {
	@@ -308,13 +308,13 @@ class CATLASSGemmTemplate(CATLASSTemplate, ABC):
308	return layout::ColumnMajor(layout.shape(0), layout.shape(1), RoundUp(layout.shape(0), align));	308	return layout::ColumnMajor(layout.shape(0), layout.shape(1), RoundUp(layout.shape(0), align));
309	}	309	}
310		310
311	-	311	+
312	size_t GetWorkspaceLen(layout::RowMajor layout)	312	size_t GetWorkspaceLen(layout::RowMajor layout)
313	{	313	{
314	return layout.shape(0) * layout.stride(0);	314	return layout.shape(0) * layout.stride(0);
315	}	315	}
316		316
317	-	317	+
318	size_t GetWorkspaceLen(layout::ColumnMajor layout)	318	size_t GetWorkspaceLen(layout::ColumnMajor layout)
319	{	319	{
320	return layout.shape(1) * layout.stride(1);	320	return layout.shape(1) * layout.stride(1);
	@@ -383,7 +383,7 @@ class CATLASSGemmTemplate(CATLASSTemplate, ABC):
383	"""	383	"""
384	import catlass_cppgen.catlass.layout as catlass_lib_layout	384	import catlass_cppgen.catlass.layout as catlass_lib_layout
385		385
386	- # bias stride could be (0, 1), which indicates (n,) bias	386	+ # bias stride could be (0, 1), which indicates (n,) bias
387	if len(torch_layout.stride) == 1 or torch_layout.stride[0] == 0:	387	if len(torch_layout.stride) == 1 or torch_layout.stride[0] == 0:
388	return catlass_lib_layout.VectorLayout	388	return catlass_lib_layout.VectorLayout
389		389
	@@ -744,7 +744,7 @@ class CATLASS1xGemmTemplate(CATLASSGemmTemplate):
744	@staticmethod	744	@staticmethod
745	def is_mixed_template(op: "GemmKernelBase") -> bool:	745	def is_mixed_template(op: "GemmKernelBase") -> bool:
746	return getattr(op, "is_mix", False)	746	return getattr(op, "is_mix", False)
747	-	747	+
748	@staticmethod	748	@staticmethod
749	def epilogue_fusion_type(op: "GemmKernelBase") -> int:	749	def epilogue_fusion_type(op: "GemmKernelBase") -> int:
750	fusion_type = 0	750	fusion_type = 0

Mtorch_npu/_inductor/codegen/kernel_analysis.py+3-3

Mtorch_npu/_inductor/codegen/npu/device_op_overrides.py+1-1

	@@ -71,7 +71,7 @@ class NewNPUDeviceOpOverrides(DeviceOpOverrides):
71	load_code = """	71	load_code = """
72	static std::unordered_map<std::string, size_t> registered_names;	72	static std::unordered_map<std::string, size_t> registered_names;
73	static std::unordered_map<std::string, std::unique_ptr<size_t>> func_stubs;	73	static std::unordered_map<std::string, std::unique_ptr<size_t>> func_stubs;
74	-	74	+
75	static inline void * loadKernel(	75	static inline void * loadKernel(
76	std::string filePath,	76	std::string filePath,
77	const std::string &&nameFunc,	77	const std::string &&nameFunc,

Mtorch_npu/_inductor/codegen/split_tiling.py+3-3

	@@ -39,7 +39,7 @@ class SplitTiling:
39	if not self.kernel.golden_var_list:	39	if not self.kernel.golden_var_list:
40	self.kernel.select_golden_varlist()	40	self.kernel.select_golden_varlist()
41	stride_sorted_var_list = list(self.kernel.golden_var_list) if self.kernel.golden_var_list else []	41	stride_sorted_var_list = list(self.kernel.golden_var_list) if self.kernel.golden_var_list else []
42	- reduction_dim_list = []	42	+ reduction_dim_list = []
43	for i, x in enumerate(reversed(stride_sorted_var_list)):	43	for i, x in enumerate(reversed(stride_sorted_var_list)):
44	if x.name[0] == 'r':	44	if x.name[0] == 'r':
45	reduction_dim_list.append(i)	45	reduction_dim_list.append(i)
	@@ -79,7 +79,7 @@ class SplitTiling:
79	self.kernel.split_axis.clear()	79	self.kernel.split_axis.clear()
80		80
81	# total numel exceed aicore or total split axis exceed 3	81	# total numel exceed aicore or total split axis exceed 3
82	- def meet_stop_condition():	82	+ def meet_stop_condition():
83	sv = V.graph.sizevars	83	sv = V.graph.sizevars
84	current_numels = self.total_split_numels(self.kernel.split_axis)	84	current_numels = self.total_split_numels(self.kernel.split_axis)
85	try:	85	try:
	@@ -160,7 +160,7 @@ class SplitTiling:
160	self.kernel.range_tree_nodes[var].is_tiling_axis	160	self.kernel.range_tree_nodes[var].is_tiling_axis
161	for var in self.kernel.reduction_axis_list()	161	for var in self.kernel.reduction_axis_list()
162	) and not self.contiguous_reduction	162	) and not self.contiguous_reduction
163	-	163	+
164	if can_stop():	164	if can_stop():
165	return True	165	return True
166	return False	166	return False

Mtorch_npu/_inductor/codegen/tile_generator.py+3-3

	@@ -65,7 +65,7 @@ class TileGenerator:
65	continue	65	continue
66	if self.axis_name[tiling_axis][0] == "r" and self.persistent_reduction:	66	if self.axis_name[tiling_axis][0] == "r" and self.persistent_reduction:
67	continue	67	continue
68	- self.real_tiling_axis.append(tiling_axis)	68	+ self.real_tiling_axis.append(tiling_axis)
69	self.split_axis_num = len(self.split_axis)	69	self.split_axis_num = len(self.split_axis)
70		70
71	def reset_configs(self):	71	def reset_configs(self):
	@@ -83,7 +83,7 @@ class TileGenerator:
83	continue	83	continue
84	if self.axis_name[tiling_axis][0] == "r" and self.persistent_reduction:	84	if self.axis_name[tiling_axis][0] == "r" and self.persistent_reduction:
85	continue	85	continue
86	- self.real_tiling_axis.append(tiling_axis)	86	+ self.real_tiling_axis.append(tiling_axis)
87	self.split_axis_num = len(self.split_axis)	87	self.split_axis_num = len(self.split_axis)
88		88
89	def calcu_last_split_blocks(self, axis):	89	def calcu_last_split_blocks(self, axis):
	@@ -407,7 +407,7 @@ class TileGenerator:
407	self.tune_multibuffer()	407	self.tune_multibuffer()
408	if self.npu_kernel_type == NPUKernelType.SIMT_ONLY:	408	if self.npu_kernel_type == NPUKernelType.SIMT_ONLY:
409	self.tune_simt_num_warps()	409	self.tune_simt_num_warps()
410	-	410	+
411	def set_kernel_type(self, npu_kernel_type):	411	def set_kernel_type(self, npu_kernel_type):
412	self.npu_kernel_type = npu_kernel_type	412	self.npu_kernel_type = npu_kernel_type
413		413

Mtorch_npu/_inductor/codegen/triton_utils.py+2-2

	@@ -54,7 +54,7 @@ def get_aligned_numel(dtype):
54	def get_indirect_var(node_name):	54	def get_indirect_var(node_name):
55	match = re.compile(r"indirect").search(node_name)	55	match = re.compile(r"indirect").search(node_name)
56	if match is None:	56	if match is None:
57	- return None	57	+ return None
58	return node_name[match.start():]	58	return node_name[match.start():]
59		59
60		60
	@@ -62,5 +62,5 @@ def get_indirect_mem_var(node_name):
62	indirect_mem_pattern = r'index_select\|gather_template\|indexput_template\|scatter_template'	62	indirect_mem_pattern = r'index_select\|gather_template\|indexput_template\|scatter_template'
63	match = re.compile(indirect_mem_pattern).search(node_name)	63	match = re.compile(indirect_mem_pattern).search(node_name)
64	if match is None:	64	if match is None:
65	- return None	65	+ return None
66	return node_name[match.start():]	66	return node_name[match.start():]

Mtorch_npu/_inductor/codegen/wrapper.py+10-10

Mtorch_npu/_inductor/config.py+1-1

	@@ -29,7 +29,7 @@ config.trace.enabled = True
29	config.fallback_random = True	29	config.fallback_random = True
30		30
31	config.graph_partition = False	31	config.graph_partition = False
32	-	32	+
33	config.triton.coalesce_tiling_analysis = False	33	config.triton.coalesce_tiling_analysis = False
34		34
35	device = torch.npu.current_device()	35	device = torch.npu.current_device()

Mtorch_npu/_inductor/cpp_builder.py+4-4

	@@ -29,7 +29,7 @@ def include_paths(npu: bool = False) -> List[str]:
29		29
30	Args:	30	Args:
31	npu: If 'True', includes NPU-specific include paths.	31	npu: If 'True', includes NPU-specific include paths.
32	-	32	+
33	Returns:	33	Returns:
34	A list if include path strings.	34	A list if include path strings.
35	"""	35	"""
	@@ -88,7 +88,7 @@ def get_cpp_torch_device_options(
88	aot_mode: bool = False,	88	aot_mode: bool = False,
89	compile_only: bool = False,	89	compile_only: bool = False,
90	) -> Tuple[List[str], List[str], List[str], List[str], List[str], List[str], List[str]]:	90	) -> Tuple[List[str], List[str], List[str], List[str], List[str], List[str], List[str]]:
91	-	91	+
92	npu = "npu" == device_type	92	npu = "npu" == device_type
93		93
94	definations: List[str] = []	94	definations: List[str] = []
	@@ -145,14 +145,14 @@ def _get_optimization_cflags(
145	cflags.append(f"ffp-contract={config.cpp.enable_floating_point_contract_flag}")	145	cflags.append(f"ffp-contract={config.cpp.enable_floating_point_contract_flag}")
146		146
147	if sys.platform != "darwin":	147	if sys.platform != "darwin":
148	- # on macos, unknown argument: '-fno-tree-loop-vectorize'	148	+ # on macos, unknown argument: '-fno-tree-loop-vectorize'
149	if _is_gcc(cpp_compiler):	149	if _is_gcc(cpp_compiler):
150	cflags.append("fno-tree-loop-vectorize")	150	cflags.append("fno-tree-loop-vectorize")
151	# -march=native is unrecognized option on M1	151	# -march=native is unrecognized option on M1
152	if not config.is_fbcode():	152	if not config.is_fbcode():
153	if platform.machine() == "ppc64le":	153	if platform.machine() == "ppc64le":
154	cflags.append("mcpu=native")	154	cflags.append("mcpu=native")
155	-	155	+
156	return cflags, ldflags	156	return cflags, ldflags
157		157
158		158

Mtorch_npu/_inductor/fx_passes/ascend_custom_passes/__init__.py+1-1

	@@ -23,7 +23,7 @@ def run_register_pre_custom_passes(gm):
23	fn(gm)	23	fn(gm)
24		24
25	log.debug(f"after pre_grad graph optimizer pass, graph is: {gm}")	25	log.debug(f"after pre_grad graph optimizer pass, graph is: {gm}")
26	-	26	+
27		27
28	def run_register_post_custom_passes(gm):	28	def run_register_post_custom_passes(gm):
29	log.debug(f"before post_grad graph optimizer pass, graph is: {gm}")	29	log.debug(f"before post_grad graph optimizer pass, graph is: {gm}")

Mtorch_npu/_inductor/fx_passes/joint_graph.py+1-1

	@@ -11,5 +11,5 @@ def patch_constant_fold_uniform_value():
11	src_func(gm)	11	src_func(gm)
12	if isinstance(gm, torch.fx.GraphModule):	12	if isinstance(gm, torch.fx.GraphModule):
13	gm.graph.eliminate_dead_code()	13	gm.graph.eliminate_dead_code()
14	-	14	+
15	joint_graph.constant_fold_uniform_value = new_constant_fold_uniform_value	15	joint_graph.constant_fold_uniform_value = new_constant_fold_uniform_value

Mtorch_npu/_inductor/fx_passes/parallelism_strategy_base.py+1-1

	@@ -9,7 +9,7 @@ class ParallelStrategyBase(ABC):
9		9
10	def __init__(self):	10	def __init__(self):
11	self.name: str	11	self.name: str
12	-	12	+
13	@abstractmethod	13	@abstractmethod
14	def assign_parallel_groups(self, nodes: List[BaseSchedulerNode]) -> Dict[str, List[BaseSchedulerNode]]:	14	def assign_parallel_groups(self, nodes: List[BaseSchedulerNode]) -> Dict[str, List[BaseSchedulerNode]]:
15	pass	15	pass

Mtorch_npu/_inductor/fx_passes/parallelism_strategy_cv.py+6-6

Mtorch_npu/_inductor/fx_passes/parallelism_strategy_default.py+1-1

	@@ -51,7 +51,7 @@ class DefaultParallelStrategy(ParallelStrategyBase):
51	sorted_preds = sorted(preds, key=lambda x: node_to_idx.get(x, -1))	51	sorted_preds = sorted(preds, key=lambda x: node_to_idx.get(x, -1))
52	groups = []	52	groups = []
53	assigned = set()	53	assigned = set()
54	-	54	+
55	def get_ancestors_node(start):	55	def get_ancestors_node(start):
56	ancestors = set()	56	ancestors = set()
57	stack = [start]	57	stack = [start]

Mtorch_npu/_inductor/fx_passes/parallelism_strategy_framework.py+3-3

	@@ -20,11 +20,11 @@ def register_custom_parallel_strategy(strategy_name: str, strategy: ParallelStra
20		20
21		21
22	class ParallelGroupingStrategy:	22	class ParallelGroupingStrategy:
23	-	23	+
24	def __init__(self):	24	def __init__(self):
25	register_custom_parallel_strategy("default", CVParallelismStrategy)	25	register_custom_parallel_strategy("default", CVParallelismStrategy)
26	-	26	+
27	-	27	+
28	def execute_strategy(self, nodes: List[BaseSchedulerNode]) -> Dict[str, List[BaseSchedulerNode]]:	28	def execute_strategy(self, nodes: List[BaseSchedulerNode]) -> Dict[str, List[BaseSchedulerNode]]:
29	parallel_scheduler_nodes_min = os.environ.get("PARALLEL_SCHEDULER_NODES_MIN", 20)	29	parallel_scheduler_nodes_min = os.environ.get("PARALLEL_SCHEDULER_NODES_MIN", 20)
30	if len(nodes) <= parallel_scheduler_nodes_min:	30	if len(nodes) <= parallel_scheduler_nodes_min:

Mtorch_npu/_inductor/fx_passes/post_grad.py+2-2

	@@ -14,9 +14,9 @@ def patch_pattern_mm_plus_mm():
14	return False	14	return False
15		15
16	pattern = CallFunction(	16	pattern = CallFunction(
17	- aten.add,	17	+ aten.add,
18	CallFunction(aten.mm, KeywordArg("mat1"), KeywordArg("mat2")),	18	CallFunction(aten.mm, KeywordArg("mat1"), KeywordArg("mat2")),
19	- CallFunction(aten.mm, KeywordArg("mat3"), KeywordArg("mat4")),	19	+ CallFunction(aten.mm, KeywordArg("mat3"), KeywordArg("mat4")),
20	extra_check=is_valid_mm_plus_mm	20	extra_check=is_valid_mm_plus_mm
21	)	21	)
22		22

Mtorch_npu/_inductor/fx_passes/utils/fx_pass_level.py+4-4

	@@ -17,7 +17,7 @@ class FxPassLevel(Enum):
17	if isinstance(other, FxPassLevel):	17	if isinstance(other, FxPassLevel):
18	return self.value == other.value	18	return self.value == other.value
19	return NotImplemented	19	return NotImplemented
20	-	20	+
21	def __hash__(self):	21	def __hash__(self):
22	return hash(self.value)	22	return hash(self.value)
23		23
	@@ -36,7 +36,7 @@ class PassType(Enum):
36	if isinstance(other, PassType):	36	if isinstance(other, PassType):
37	return self.value == other.value	37	return self.value == other.value
38	return NotImplemented	38	return NotImplemented
39	-	39	+
40	def __hash__(self):	40	def __hash__(self):
41	return hash(self.value)	41	return hash(self.value)
42		42
	@@ -60,7 +60,7 @@ class ComputeType(Enum):
60	if isinstance(other, ComputeType):	60	if isinstance(other, ComputeType):
61	return self.value == other.value	61	return self.value == other.value
62	return NotImplemented	62	return NotImplemented
63	-	63	+
64	def __hash__(self):	64	def __hash__(self):
65	return hash(self.value)	65	return hash(self.value)
66		66
	@@ -82,6 +82,6 @@ class GroupType(Enum):
82	if isinstance(other, ComputeType):	82	if isinstance(other, ComputeType):
83	return self.value == other.value	83	return self.value == other.value
84	return NotImplemented	84	return NotImplemented
85	-	85	+
86	def __hash__(self):	86	def __hash__(self):
87	return hash(self.value)	87	return hash(self.value)

Mtorch_npu/_inductor/graph.py+1-1

	@@ -404,7 +404,7 @@ def patch_run_node():
404	# Use inner fn as a rough proxy. Good enough.	404	# Use inner fn as a rough proxy. Good enough.
405	if curr.has_large_inner_fn(threshold=100):	405	if curr.has_large_inner_fn(threshold=100):
406	result.realize()	406	result.realize()
407	-	407	+
408	from .config import lowering_axis_count	408	from .config import lowering_axis_count
409	if lowering_axis_count and len(curr.ranges) >= lowering_axis_count:	409	if lowering_axis_count and len(curr.ranges) >= lowering_axis_count:
410	result.realize()	410	result.realize()

Mtorch_npu/_inductor/kernel/flex_attention.py+16-16

Mtorch_npu/_inductor/kernel/mm_grouped.py+6-6

	@@ -125,7 +125,7 @@ def check_catlass_support(
125		125
126	if len(mat_a[0].get_size()) != 2 or len(mat_b[0].get_size()) != 3:	126	if len(mat_a[0].get_size()) != 2 or len(mat_b[0].get_size()) != 3:
127	return False	127	return False
128	-	128	+
129	# Catlass currently not support group mm with bias	129	# Catlass currently not support group mm with bias
130	if bias:	130	if bias:
131	return False	131	return False
	@@ -134,14 +134,14 @@ def check_catlass_support(
134	return False	134	return False
135		135
136	if offset:	136	if offset:
137	- return False	137	+ return False
138		138
139	if group_list is None or not isinstance(group_list, TensorBox):	139	if group_list is None or not isinstance(group_list, TensorBox):
140	return False	140	return False
141		141
142	if group_list.get_size()[0] != mat_b[0].get_size()[0]:	142	if group_list.get_size()[0] != mat_b[0].get_size()[0]:
143	return False	143	return False
144	-	144	+
145	# Catlass only support splitting in m axis	145	# Catlass only support splitting in m axis
146	if group_type is not None and group_type != 0:	146	if group_type is not None and group_type != 0:
147	return False	147	return False
	@@ -181,7 +181,7 @@ def _tuned_grouped_mm_common(
181	output_dtype: Optional[torch.dtype] = None,	181	output_dtype: Optional[torch.dtype] = None,
182	**kwargs,	182	**kwargs,
183	) -> List[TensorBox]:	183	) -> List[TensorBox]:
184	- catlass_compatible = check_catlass_support(mat_a, mat_b, bias, scale, offset,	184	+ catlass_compatible = check_catlass_support(mat_a, mat_b, bias, scale, offset,
185	group_list, group_type, group_list_type, act_type, output_dtype	185	group_list, group_type, group_list_type, act_type, output_dtype
186	)	186	)
187	# not support lowering grouped-mm for cpp_wrapper yet	187	# not support lowering grouped-mm for cpp_wrapper yet
	@@ -237,8 +237,8 @@ def _tuned_grouped_mm_common(
237		237
238	if is_contiguous_input and is_nonzero and use_catlass_template("grouped_mm", layout, m, n, k):	238	if is_contiguous_input and is_nonzero and use_catlass_template("grouped_mm", layout, m, n, k):
239	CATLASS1xGemmTemplate.add_catlass_gemm_choices(	239	CATLASS1xGemmTemplate.add_catlass_gemm_choices(
240	- choices,	240	+ choices,
241	- layout,	241	+ layout,
242	[mat_a, mat_b, bias, offs], # currently catlass does not support grouped mm with bias	242	[mat_a, mat_b, bias, offs], # currently catlass does not support grouped mm with bias
243	)	243	)
244	# debug log	244	# debug log

Mtorch_npu/_inductor/lowering.py+14-14

	@@ -208,7 +208,7 @@ def _register_npu_inductor_fallbacks():
208	log.info(f"[npu\|inductor\|lowering\|fallback] with FALLBACK_LIST, len(lowerings): {len(lowerings)}, "	208	log.info(f"[npu\|inductor\|lowering\|fallback] with FALLBACK_LIST, len(lowerings): {len(lowerings)}, "
209	f"len(FALLBACK_LIST): {len(FALLBACK_LIST)}, make_fallback finished.")	209	f"len(FALLBACK_LIST): {len(FALLBACK_LIST)}, make_fallback finished.")
210	log.info(f"[npu\|inductor\|lowering\|fallback] len(NPU_EXTRA_FALLBACK_LIST): {len(NPU_EXTRA_FALLBACK_LIST)}")	210	log.info(f"[npu\|inductor\|lowering\|fallback] len(NPU_EXTRA_FALLBACK_LIST): {len(NPU_EXTRA_FALLBACK_LIST)}")
211	-	211	+
212	# 把需要overload的op在lowering里删除	212	# 把需要overload的op在lowering里删除
213	overload_op_set = set()	213	overload_op_set = set()
214	_add_overload(LOWERING_OVERRIDE_OP, overload_op_set)	214	_add_overload(LOWERING_OVERRIDE_OP, overload_op_set)
	@@ -484,7 +484,7 @@ def _register_npu_inductor_fallbacks():
484	idx = [gather_idx]	484	idx = [gather_idx]
485	else:	485	else:
486	idx[dim] = gather_idx	486	idx[dim] = gather_idx
487	-	487	+
488	return ops.gather_template(loader_name, x_loader(idx), index_value, gather_idx, int(index_boundary))	488	return ops.gather_template(loader_name, x_loader(idx), index_value, gather_idx, int(index_boundary))
489		489
490	return Pointwise.create(	490	return Pointwise.create(
	@@ -494,7 +494,7 @@ def _register_npu_inductor_fallbacks():
494	ranges=index.get_size(),	494	ranges=index.get_size(),
495	)	495	)
496		496
497	-	497	+
498	def index_put_impl_(self, indices, values, accumulate, check, may_realize=False):	498	def index_put_impl_(self, indices, values, accumulate, check, may_realize=False):
499	if may_realize:	499	if may_realize:
500		500
	@@ -652,7 +652,7 @@ def _register_npu_inductor_fallbacks():
652	if x_ndim == 0:	652	if x_ndim == 0:
653	self = view(self, [])	653	self = view(self, [])
654	return self	654	return self
655	-	655	+
656	# All the indexing decompositions are written in terms of index, index_put, and index_put_	656	# All the indexing decompositions are written in terms of index, index_put, and index_put_
657	# We cannot have this lowering as a decomposition as it introduces	657	# We cannot have this lowering as a decomposition as it introduces
658	# mutation in the graph, which is bad for Aot Autograd. Aot Autograd runs dead	658	# mutation in the graph, which is bad for Aot Autograd. Aot Autograd runs dead
	@@ -951,13 +951,13 @@ def _register_npu_inductor_fallbacks():
951	# Validate input	951	# Validate input
952	if not isinstance(normalized_shape, (list, tuple)):	952	if not isinstance(normalized_shape, (list, tuple)):
953	normalized_shape = (normalized_shape,)	953	normalized_shape = (normalized_shape,)
954	-	954	+
955	normalized_ndim = len(normalized_shape)	955	normalized_ndim = len(normalized_shape)
956	input_shape = x.get_size()	956	input_shape = x.get_size()
957	-	957	+
958	# Calculate reduction dimension indices	958	# Calculate reduction dimension indices
959	reduce_dims = list(range(len(input_shape) - normalized_ndim, len(input_shape)))	959	reduce_dims = list(range(len(input_shape) - normalized_ndim, len(input_shape)))
960	-	960	+
961	# Compute mean and variance	961	# Compute mean and variance
962	var, mean = var_mean_helper_(	962	var, mean = var_mean_helper_(
963	x=x,	963	x=x,
	@@ -966,30 +966,30 @@ def _register_npu_inductor_fallbacks():
966	keepdim=True, # Keep dimensions for broadcasting	966	keepdim=True, # Keep dimensions for broadcasting
967	return_mean=True	967	return_mean=True
968	)	968	)
969	-	969	+
970	# Calculate normalized result (x - mean) / sqrt(var + eps)	970	# Calculate normalized result (x - mean) / sqrt(var + eps)
971	x_normalized = sub(x, mean)	971	x_normalized = sub(x, mean)
972	-	972	+
973	# Add eps to variance	973	# Add eps to variance
974	eps_tensor = ir.IndexingConstant(index=eps, dtype=var.get_dtype(), device=var.get_device())	974	eps_tensor = ir.IndexingConstant(index=eps, dtype=var.get_dtype(), device=var.get_device())
975	eps_tensor = ExpandView.create(eps_tensor, var.get_size())	975	eps_tensor = ExpandView.create(eps_tensor, var.get_size())
976	var_eps = add(var, eps_tensor)	976	var_eps = add(var, eps_tensor)
977	-	977	+
978	# Calculate reciprocal of sqrt(var + eps)	978	# Calculate reciprocal of sqrt(var + eps)
979	inv_std = rsqrt(var_eps) # 1 / sqrt(var + eps)	979	inv_std = rsqrt(var_eps) # 1 / sqrt(var + eps)
980	-	980	+
981	# Normalization	981	# Normalization
982	normalized = mul(x_normalized, inv_std)	982	normalized = mul(x_normalized, inv_std)
983	-	983	+
984	# Apply optional affine transformation (gamma * normalized + beta)	984	# Apply optional affine transformation (gamma * normalized + beta)
985	if weight is not None:	985	if weight is not None:
986	# weight will be broadcast automatically, mul function in lowering supports broadcasting	986	# weight will be broadcast automatically, mul function in lowering supports broadcasting
987	normalized = mul(normalized, weight)	987	normalized = mul(normalized, weight)
988	-	988	+
989	if bias is not None:	989	if bias is not None:
990	# add will be broadcast automatically	990	# add will be broadcast automatically
991	normalized = add(normalized, bias)	991	normalized = add(normalized, bias)
992	-	992	+
993	# native_layer_norm returns three values: output, mean, reciprocal of standard deviation	993	# native_layer_norm returns three values: output, mean, reciprocal of standard deviation
994	return normalized, mean, inv_std	994	return normalized, mean, inv_std
995		995

Mtorch_npu/_inductor/lowering_fallback_list.py+2-2

	@@ -762,7 +762,7 @@ TORCH_NATIVE_FALLBACK_LIST = [
762	aten._thnn_fused_lstm_cell.out,	762	aten._thnn_fused_lstm_cell.out,
763	aten._to_sparse.default,	763	aten._to_sparse.default,
764	aten._to_sparse.out,	764	aten._to_sparse.out,
765	- aten._to_sparse.sparse_dim,	765	+ aten._to_sparse.sparse_dim,
766	aten._to_sparse.sparse_dim_out,	766	aten._to_sparse.sparse_dim_out,
767	aten._trilinear.default,	767	aten._trilinear.default,
768	aten._trilinear.out,	768	aten._trilinear.out,
	@@ -930,7 +930,7 @@ TORCH_NATIVE_FALLBACK_LIST = [
930	aten.special_chebyshev_polynomial_v.out,	930	aten.special_chebyshev_polynomial_v.out,
931	aten.special_chebyshev_polynomial_v.x_scalar,	931	aten.special_chebyshev_polynomial_v.x_scalar,
932	aten.special_chebyshev_polynomial_v.x_scalar_out,	932	aten.special_chebyshev_polynomial_v.x_scalar_out,
933	- aten.special_chebyshev_polynomial_w.default,	933	+ aten.special_chebyshev_polynomial_w.default,
934	aten.special_chebyshev_polynomial_w.n_scalar,	934	aten.special_chebyshev_polynomial_w.n_scalar,
935	aten.special_chebyshev_polynomial_w.n_scalar_out,	935	aten.special_chebyshev_polynomial_w.n_scalar_out,
936	aten.special_chebyshev_polynomial_w.out,	936	aten.special_chebyshev_polynomial_w.out,

Mtorch_npu/_inductor/lowering_fx.py+8-8

	@@ -698,7 +698,7 @@ def create_compile_kwargs(final_kernel, fx_call_args, fx_args):
698	def generate_fx_graph_code(code, kernel_code, kernel_name, compile_kwargs):	698	def generate_fx_graph_code(code, kernel_code, kernel_name, compile_kwargs):
699	code = textwrap.indent(code, ' ')	699	code = textwrap.indent(code, ' ')
700	code_template = f"""	700	code_template = f"""
701	-import os	701	+import os
702	import torch	702	import torch
703	from torch._inductor.compile_fx import clone_preserve_strides	703	from torch._inductor.compile_fx import clone_preserve_strides
704	from torch._dynamo.testing import rand_strided	704	from torch._dynamo.testing import rand_strided
	@@ -759,7 +759,7 @@ def run():
759		759
760	stream0 = get_raw_stream(0)	760	stream0 = get_raw_stream(0)
761		761
762	-	762	+
763	args = torch.load(os.path.join(dir_path, "data.pth"))	763	args = torch.load(os.path.join(dir_path, "data.pth"))
764		764
765	call_inputs_indices = call_args_mapping[:num_inputs]	765	call_inputs_indices = call_args_mapping[:num_inputs]
	@@ -767,7 +767,7 @@ def run():
767		767
768	args = [arg.npu() if isinstance(arg, torch.Tensor) else arg for arg in args]	768	args = [arg.npu() if isinstance(arg, torch.Tensor) else arg for arg in args]
769		769
770	- fx_args = []	770	+ fx_args = []
771	for idx in call_args_mapping:	771	for idx in call_args_mapping:
772	arg = args[idx]	772	arg = args[idx]
773	if isinstance(arg, torch.Tensor):	773	if isinstance(arg, torch.Tensor):
	@@ -785,7 +785,7 @@ def run():
785	out1 = out1.reshape(out2.shape)	785	out1 = out1.reshape(out2.shape)
786	if idx in non_contiguous_indices['outputs']:	786	if idx in non_contiguous_indices['outputs']:
787	out2.copy_(out1)	787	out2.copy_(out1)
788	- else:	788	+ else:
789	out2.data = out1.data	789	out2.data = out1.data
790		790
791	{kernel_name}.run(*args, stream=stream0)	791	{kernel_name}.run(*args, stream=stream0)
	@@ -994,7 +994,7 @@ def _make_reduction_inner(x, *, axis, keepdims, dtype, override_return_dtype):
994	reduction_ranges=reduced_sizes,	994	reduction_ranges=reduced_sizes,
995	)	995	)
996		996
997	-	997	+
998	def dump_fx_graph_code(code, dump_path, traced_graph_hash):	998	def dump_fx_graph_code(code, dump_path, traced_graph_hash):
999	py_path = os.path.join(dump_path, traced_graph_hash + '.py')	999	py_path = os.path.join(dump_path, traced_graph_hash + '.py')
1000	PathManager.check_input_file_path(py_path)	1000	PathManager.check_input_file_path(py_path)
	@@ -1227,7 +1227,7 @@ def _register_npu_inductor_fallbacks_fx(make_reduction):
1227	)(fn)	1227	)(fn)
1228	return fn	1228	return fn
1229		1229
1230	-	1230	+
1231		1231
1232	@register_lowering(aten.where, broadcast=False, type_promotion_kind=None)	1232	@register_lowering(aten.where, broadcast=False, type_promotion_kind=None)
1233	def where(cond, a, b):	1233	def where(cond, a, b):
	@@ -1425,7 +1425,7 @@ def _register_npu_inductor_fallbacks_fx(make_reduction):
1425	@register_lowering(prims.device_put, type_promotion_kind=None)	1425	@register_lowering(prims.device_put, type_promotion_kind=None)
1426	def _device_put(x: TensorBox, device: torch.device, non_blocking=False):	1426	def _device_put(x: TensorBox, device: torch.device, non_blocking=False):
1427	return to_device(x, device, copy=True, non_blocking=non_blocking)	1427	return to_device(x, device, copy=True, non_blocking=non_blocking)
1428	-	1428	+
1429	@register_lowering(aten.repeat)	1429	@register_lowering(aten.repeat)
1430	def repeat(x, repeats):	1430	def repeat(x, repeats):
1431	input_graphs = fetch_graphs([x, repeats])	1431	input_graphs = fetch_graphs([x, repeats])
	@@ -2361,7 +2361,7 @@ def _register_npu_inductor_fallbacks_fx(make_reduction):
2361	*idx[indices_ndim:]	2361	*idx[indices_ndim:]
2362	]	2362	]
2363	return weight_loader(weight_idx)	2363	return weight_loader(weight_idx)
2364	-	2364	+
2365	input_graphs = fetch_graphs([weight, indices])	2365	input_graphs = fetch_graphs([weight, indices])
2366	node_name = f'embedding_{next(node_id)}'	2366	node_name = f'embedding_{next(node_id)}'
2367	new_graph = merge_traced_graphs(input_graphs, torch.ops.aten.embedding.default, node_name, padding_idx=padding_idx, scale_grad_by_freq=scale_grad_by_freq, sparse=sparse)	2367	new_graph = merge_traced_graphs(input_graphs, torch.ops.aten.embedding.default, node_name, padding_idx=padding_idx, scale_grad_by_freq=scale_grad_by_freq, sparse=sparse)

Mtorch_npu/_inductor/lowering_override_list.py+1-1

	@@ -7,7 +7,7 @@ prims = torch.ops.prims
7	npu = torch.ops.npu	7	npu = torch.ops.npu
8		8
9		9
10	-# in lowering.py, we will remove the following op's default lowering function,	10	+# in lowering.py, we will remove the following op's default lowering function,
11	# and register its new override-lowering-function.	11	# and register its new override-lowering-function.
12	LOWERING_OVERRIDE_OP = [	12	LOWERING_OVERRIDE_OP = [
13	aten.cumsum,	13	aten.cumsum,

Mtorch_npu/_inductor/npu_compare.py+6-6

	@@ -26,7 +26,7 @@ def compare_outputs(
26	continue	26	continue
27	if actual.dtype != expected.dtype:	27	if actual.dtype != expected.dtype:
28	expected = expected.to(actual.dtype)	28	expected = expected.to(actual.dtype)
29	-	29	+
30	tol = tolerances.get(actual.dtype, tolerances["default"])	30	tol = tolerances.get(actual.dtype, tolerances["default"])
31	rtol, atol = tol["rtol"], tol["atol"]	31	rtol, atol = tol["rtol"], tol["atol"]
32	matches = torch.isclose(actual, expected, rtol=rtol, atol=atol, equal_nan=True)	32	matches = torch.isclose(actual, expected, rtol=rtol, atol=atol, equal_nan=True)
	@@ -34,7 +34,7 @@ def compare_outputs(
34	_report_mismatch(idx, actual, expected, matches, rtol, atol, kernel_name)	34	_report_mismatch(idx, actual, expected, matches, rtol, atol, kernel_name)
35	failed_indices.append(idx)	35	failed_indices.append(idx)
36	del matches	36	del matches
37	-	37	+
38	return not failed_indices	38	return not failed_indices
39		39
40		40
	@@ -117,7 +117,7 @@ def check_accuracy_triton(args, launcher, grid, stream, inductor_meta, *kwargs
117	arg = args[idx]	117	arg = args[idx]
118	if isinstance(arg, torch.Tensor):	118	if isinstance(arg, torch.Tensor):
119	fx_args.append(clone_for_accuracy(arg))	119	fx_args.append(clone_for_accuracy(arg))
120	-	120	+
121	fx_graph_call(*fx_args)	121	fx_graph_call(*fx_args)
122		122
123	launcher(args, *kwargs, stream=stream)	123	launcher(args, *kwargs, stream=stream)
	@@ -155,7 +155,7 @@ def check_accuracy_mlir(args, kernel_name, launchers, num_outputs, dynamic, *k
155	args_new = args_new + (arg, arg, 0) + arg.size() + arg.stride()	155	args_new = args_new + (arg, arg, 0) + arg.size() + arg.stride()
156	else:	156	else:
157	args_new = args	157	args_new = args
158	-	158	+
159	output = launcher(args_new, *kwargs)	159	output = launcher(args_new, *kwargs)
160	result = compare_outputs(	160	result = compare_outputs(
161	args[num_inputs:],	161	args[num_inputs:],
	@@ -186,8 +186,8 @@ def _load_fx_model(acc_meta):
186	model = Model()	186	model = Model()
187	acc_meta['_fx_model'] = model	187	acc_meta['_fx_model'] = model
188	return model	188	return model
189	-	189	+
190	-	190	+
191	def check_accuracy_dvm(kobj, acc_meta, kernel_name, args):	191	def check_accuracy_dvm(kobj, acc_meta, kernel_name, args):
192	"""Run DVM kernel then compare outputs against FX graph reference."""	192	"""Run DVM kernel then compare outputs against FX graph reference."""
193	from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import config as anir_config	193	from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import config as anir_config

Mtorch_npu/_inductor/select_algorithm.py+13-13

	@@ -71,16 +71,16 @@ class NPUCompileError(CppCompileError):
71		71
72	class NPUTritonTemplate(TritonTemplate):	72	class NPUTritonTemplate(TritonTemplate):
73	"""NPU-specific Triton template for kernel generation.	73	"""NPU-specific Triton template for kernel generation.
74	-	74	+
75	This class extends TritonTemplate to provide NPU-specific optimizations	75	This class extends TritonTemplate to provide NPU-specific optimizations
76	and configurations for Triton kernel generation.	76	and configurations for Triton kernel generation.
77	"""	77	"""
78	-	78	+
79	index_counter = itertools.count()	79	index_counter = itertools.count()
80		80
81	def __init__(self, name: str, grid: Any, source: str, debug: bool = False) -> None:	81	def __init__(self, name: str, grid: Any, source: str, debug: bool = False) -> None:
82	"""Initialize NPU Triton template.	82	"""Initialize NPU Triton template.
83	-	83	+
84	Args:	84	Args:
85	name: Template name for identification	85	name: Template name for identification
86	grid: Grid function for kernel launch configuration	86	grid: Grid function for kernel launch configuration
	@@ -274,11 +274,11 @@ class NPUTritonTemplate(TritonTemplate):
274		274
275	class NPUTritonTemplateKernel(TritonTemplateKernel):	275	class NPUTritonTemplateKernel(TritonTemplateKernel):
276	"""NPU-specific Triton template kernel for code generation.	276	"""NPU-specific Triton template kernel for code generation.
277	-	277	+
278	This class extends TritonTemplateKernel to provide NPU-specific	278	This class extends TritonTemplateKernel to provide NPU-specific
279	kernel generation and compilation functionality.	279	kernel generation and compilation functionality.
280	"""	280	"""
281	-	281	+
282	def __init__(	282	def __init__(
283	self,	283	self,
284	kernel_name: str,	284	kernel_name: str,
	@@ -298,7 +298,7 @@ class NPUTritonTemplateKernel(TritonTemplateKernel):
298	workspace_arg: Optional[Any] = None,	298	workspace_arg: Optional[Any] = None,
299	) -> None:	299	) -> None:
300	"""Initialize NPU Triton template kernel.	300	"""Initialize NPU Triton template kernel.
301	-	301	+
302	Args:	302	Args:
303	kernel_name: Name of the kernel	303	kernel_name: Name of the kernel
304	input_nodes: List of input IR nodes	304	input_nodes: List of input IR nodes
	@@ -336,10 +336,10 @@ class NPUTritonTemplateKernel(TritonTemplateKernel):
336		336
337	def def_kernel(self, *argnames: str) -> str:	337	def def_kernel(self, *argnames: str) -> str:
338	"""Hook called from template code to generate function def and needed args.	338	"""Hook called from template code to generate function def and needed args.
339	-	339	+
340	Args:	340	Args:
341	*argnames: Variable number of argument names	341	*argnames: Variable number of argument names
342	-	342	+
343	Returns:	343	Returns:
344	Render hook key string	344	Render hook key string
345	"""	345	"""
	@@ -360,13 +360,13 @@ class NPUTritonTemplateKernel(TritonTemplateKernel):
360	# Unified processing of all input nodes	360	# Unified processing of all input nodes
361	for idx, input_node in enumerate(self.input_nodes):	361	for idx, input_node in enumerate(self.input_nodes):
362	node_name = input_node.get_name()	362	node_name = input_node.get_name()
363	-	363	+
364	# Skip removed or fused buffers	364	# Skip removed or fused buffers
365	if node_name in V.graph.removed_buffers:	365	if node_name in V.graph.removed_buffers:
366	continue	366	continue
367	if node_name in self.prologue_fused_inputs:	367	if node_name in self.prologue_fused_inputs:
368	continue	368	continue
369	-	369	+
370	# Process prefix args	370	# Process prefix args
371	if idx < self.prefix_args:	371	if idx < self.prefix_args:
372	self.args.input(node_name)	372	self.args.input(node_name)
	@@ -415,7 +415,7 @@ class NPUTritonTemplateKernel(TritonTemplateKernel):
415		415
416	def patch_algorithm_selector() -> None:	416	def patch_algorithm_selector() -> None:
417	"""Patch AlgorithmSelectorCache with NPU-specific implementations.	417	"""Patch AlgorithmSelectorCache with NPU-specific implementations.
418	-	418	+
419	This function replaces the default AlgorithmSelectorCache methods with	419	This function replaces the default AlgorithmSelectorCache methods with
420	NPU-optimized versions that include profiling and benchmarking capabilities	420	NPU-optimized versions that include profiling and benchmarking capabilities
421	specific to NPU hardware.	421	specific to NPU hardware.
	@@ -709,13 +709,13 @@ def patch_algorithm_selector() -> None:
709	input_gen_fns: Optional[Dict[int, Callable[[ir.Buffer], torch.Tensor]]] = None,	709	input_gen_fns: Optional[Dict[int, Callable[[ir.Buffer], torch.Tensor]]] = None,
710	) -> Callable:	710	) -> Callable:
711	"""Create a benchmark function for the given choices.	711	"""Create a benchmark function for the given choices.
712	-	712	+
713	Args:	713	Args:
714	choices: List of choice callers to benchmark	714	choices: List of choice callers to benchmark
715	input_nodes: List of input IR nodes	715	input_nodes: List of input IR nodes
716	layout: Output layout	716	layout: Output layout
717	input_gen_fns: Optional dict mapping arg indices to input generation functions	717	input_gen_fns: Optional dict mapping arg indices to input generation functions
718	-	718	+
719	Returns:	719	Returns:
720	Benchmark function that can be called with choices	720	Benchmark function that can be called with choices
721	"""	721	"""

Mtorch_npu/_inductor/shape_handling.py+43-43

	@@ -17,7 +17,7 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
17	transform_post_fn: Post-processing function to convert tensor lists to structured outputs for transformation (optional).	17	transform_post_fn: Post-processing function to convert tensor lists to structured outputs for transformation (optional).
18	recover_pre_fn: Pre-processing function to convert inputs to tensor lists for recovery (optional).	18	recover_pre_fn: Pre-processing function to convert inputs to tensor lists for recovery (optional).
19	recover_post_fn: Post-processing function tp convert tensor lists to structured outputs for recovery (optional).	19	recover_post_fn: Post-processing function tp convert tensor lists to structured outputs for recovery (optional).
20	-	20	+
21	Each config dictionary in configs supports the following keys:	21	Each config dictionary in configs supports the following keys:
22	- type (str):	22	- type (str):
23	Logical dimension type. Supported values:	23	Logical dimension type. Supported values:
	@@ -41,7 +41,7 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
41	- policy (str):	41	- policy (str):
42	Gear generation strategy. Supported values:	42	Gear generation strategy. Supported values:
43	"TIMES" \| "CUSTOM"	43	"TIMES" \| "CUSTOM"
44	-	44	+
45	If no configs are provided at construction, a default configuration handling batch size on dimension 0 is created.	45	If no configs are provided at construction, a default configuration handling batch size on dimension 0 is created.
46	"""	46	"""
47	def __init__(	47	def __init__(
	@@ -63,7 +63,7 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
63	"TIMES": torch_npu._C.ShapePolicy.TIMES,	63	"TIMES": torch_npu._C.ShapePolicy.TIMES,
64	"CUSTOM": torch_npu._C.ShapePolicy.CUSTOM	64	"CUSTOM": torch_npu._C.ShapePolicy.CUSTOM
65	}	65	}
66	-	66	+
67	# Register processing functions	67	# Register processing functions
68	self.transform_pre_fn = transform_pre_fn	68	self.transform_pre_fn = transform_pre_fn
69	self.transform_post_fn = transform_post_fn	69	self.transform_post_fn = transform_post_fn
	@@ -91,7 +91,7 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
91	return	91	return
92	if len(configs) > 2:	92	if len(configs) > 2:
93	raise ValueError("NPUShapeHandling currently supports only two dimensions.")	93	raise ValueError("NPUShapeHandling currently supports only two dimensions.")
94	-	94	+
95	required_fields = ["type"]	95	required_fields = ["type"]
96	int_list_fields = ["dimensions", "indices", "gears"]	96	int_list_fields = ["dimensions", "indices", "gears"]
97	int_fields = ["min_size", "max_size"]	97	int_fields = ["min_size", "max_size"]
	@@ -107,11 +107,11 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
107	f"Invalid 'type' in config[{i}]: {config['type']}. "	107	f"Invalid 'type' in config[{i}]: {config['type']}. "
108	f"Must be one of: {', '.join(repr(k) for k in self.shape_type_map.keys())}."	108	f"Must be one of: {', '.join(repr(k) for k in self.shape_type_map.keys())}."
109	)	109	)
110	-	110	+
111	for field in int_list_fields:	111	for field in int_list_fields:
112	if field not in config:	112	if field not in config:
113	continue	113	continue
114	-	114	+
115	if field == "dimensions":	115	if field == "dimensions":
116	if isinstance(config[field], int):	116	if isinstance(config[field], int):
117	config[field] = [config[field]]	117	config[field] = [config[field]]
	@@ -121,20 +121,20 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
121		121
122	if not isinstance(config[field], (list, tuple)):	122	if not isinstance(config[field], (list, tuple)):
123	raise ValueError(f"Config {i} {field} must be a list, got {type(config[field])}.")	123	raise ValueError(f"Config {i} {field} must be a list, got {type(config[field])}.")
124	-	124	+
125	for item in config[field]:	125	for item in config[field]:
126	if not isinstance(item, int):	126	if not isinstance(item, int):
127	raise ValueError(f"Config {i} {field} must contain integers, got {type(item)}.")	127	raise ValueError(f"Config {i} {field} must contain integers, got {type(item)}.")
128	-	128	+
129	for field in int_fields:	129	for field in int_fields:
130	if field not in config:	130	if field not in config:
131	continue	131	continue
132	if not isinstance(config[field], int):	132	if not isinstance(config[field], int):
133	raise ValueError(f"Config {i} {field} must be an integer, got {type(config[field])}.")	133	raise ValueError(f"Config {i} {field} must be an integer, got {type(config[field])}.")
134	-	134	+
135	if "value" in config and not isinstance(config["value"], (int, float)):	135	if "value" in config and not isinstance(config["value"], (int, float)):
136	raise ValueError(f"Config {i} 'value' must be a number, got {type(config['value'])}.")	136	raise ValueError(f"Config {i} 'value' must be a number, got {type(config['value'])}.")
137	-	137	+
138	if "policy" in config:	138	if "policy" in config:
139	if not isinstance(config["policy"], str):	139	if not isinstance(config["policy"], str):
140	raise ValueError(f"Config {i} 'policy' must be a str, got {type(config['policy'])}.")	140	raise ValueError(f"Config {i} 'policy' must be a str, got {type(config['policy'])}.")
	@@ -143,7 +143,7 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
143	f"Invalid 'policy' in config[{i}]: {config['policy']}. "	143	f"Invalid 'policy' in config[{i}]: {config['policy']}. "
144	f"Must be one of: {', '.join(repr(k) for k in self.policy_map.keys())}."	144	f"Must be one of: {', '.join(repr(k) for k in self.policy_map.keys())}."
145	)	145	)
146	-	146	+
147	if len(configs) == 2 and configs[0]["type"] == configs[1]["type"]:	147	if len(configs) == 2 and configs[0]["type"] == configs[1]["type"]:
148	raise ValueError("Cannot initialize the same type repeatedly.")	148	raise ValueError("Cannot initialize the same type repeatedly.")
149		149
	@@ -164,8 +164,8 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
164	dimensions = [dimensions[0] for _ in range(len(indices))]	164	dimensions = [dimensions[0] for _ in range(len(indices))]
165	if not dimensions:	165	if not dimensions:
166	dimensions = [1 for _ in range(len(indices))]	166	dimensions = [1 for _ in range(len(indices))]
167	-	167	+
168	-	168	+
169	if len(dimensions) == 0 or len(indices) == 0:	169	if len(dimensions) == 0 or len(indices) == 0:
170	self.delay_init = True	170	self.delay_init = True
171	continue	171	continue
	@@ -183,22 +183,22 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
183	if not dimensions:	183	if not dimensions:
184	dimensions = [0]	184	dimensions = [0]
185	dimensions = [dimensions[0] for _ in range(len(tensors))]	185	dimensions = [dimensions[0] for _ in range(len(tensors))]
186	-	186	+
187	if dimension_type == "SEQLEN":	187	if dimension_type == "SEQLEN":
188	if not dimensions:	188	if not dimensions:
189	dimensions = [1]	189	dimensions = [1]
190	if len(dimensions) == 1:	190	if len(dimensions) == 1:
191	dimensions = [dimensions[0] for _ in range(len(tensors))]	191	dimensions = [dimensions[0] for _ in range(len(tensors))]
192	-	192	+
193	index = 0	193	index = 0
194	indices = []	194	indices = []
195	for dimension, tensor in zip(dimensions, tensors):	195	for dimension, tensor in zip(dimensions, tensors):
196	if tensor.ndim > dimension:	196	if tensor.ndim > dimension:
197	indices.append(index)	197	indices.append(index)
198	index += 1	198	index += 1
199	-	199	+
200	return indices	200	return indices
201	-	201	+
202	def delay_initialize(self, tensors: List[torch.Tensor]):	202	def delay_initialize(self, tensors: List[torch.Tensor]):
203	delay_init_configs = []	203	delay_init_configs = []
204	for config in self.configs:	204	for config in self.configs:
	@@ -206,7 +206,7 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
206	if "indices" not in config or len(config["indices"]) == 0:	206	if "indices" not in config or len(config["indices"]) == 0:
207	init_flag = True	207	init_flag = True
208	config["indices"] = self._construct_indices(tensors, config.get("dimensions", []), config["type"])	208	config["indices"] = self._construct_indices(tensors, config.get("dimensions", []), config["type"])
209	-	209	+
210	if init_flag:	210	if init_flag:
211	delay_init_configs.append(config)	211	delay_init_configs.append(config)
212	if len(delay_init_configs) > 0:	212	if len(delay_init_configs) > 0:
	@@ -244,7 +244,7 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
244	inputs = self.transform_pre_fn(args, *kwargs)	244	inputs = self.transform_pre_fn(args, *kwargs)
245	else:	245	else:
246	inputs, indices, leaves, spec = self._process_inputs(args, kwargs)	246	inputs, indices, leaves, spec = self._process_inputs(args, kwargs)
247	-	247	+
248	# 提取转换前的形状 (inputs 通常是 Tensor 列表)	248	# 提取转换前的形状 (inputs 通常是 Tensor 列表)
249	if logger.isEnabledFor(logging.INFO):	249	if logger.isEnabledFor(logging.INFO):
250	pre_shapes = [self.get_shape_safe(t) for t in inputs]	250	pre_shapes = [self.get_shape_safe(t) for t in inputs]
	@@ -257,16 +257,16 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
257	if logger.isEnabledFor(logging.INFO):	257	if logger.isEnabledFor(logging.INFO):
258	post_shapes = [self.get_shape_safe(t) for t in trans_outputs]	258	post_shapes = [self.get_shape_safe(t) for t in trans_outputs]
259	logger.info(f"> Post-transform content: {post_shapes}")	259	logger.info(f"> Post-transform content: {post_shapes}")
260	-	260	+
261	# 后处理阶段优化：避免嵌套循环	261	# 后处理阶段优化：避免嵌套循环
262	if self.transform_post_fn:	262	if self.transform_post_fn:
263	outputs = self.transform_post_fn(trans_outputs)	263	outputs = self.transform_post_fn(trans_outputs)
264	else:	264	else:
265	outputs = self._recover_inputs(trans_outputs, indices, leaves, spec)	265	outputs = self._recover_inputs(trans_outputs, indices, leaves, spec)
266	-	266	+
267	if not outputs:	267	if not outputs:
268	logger.error(f"CRITICAL: _recover_inputs returned NULL")	268	logger.error(f"CRITICAL: _recover_inputs returned NULL")
269	-	269	+
270	return outputs	270	return outputs
271		271
272	def flatten_to_tensors(self, structure: Any) -> Tuple[List[torch.Tensor], List[int], List[Any], TreeSpec]:	272	def flatten_to_tensors(self, structure: Any) -> Tuple[List[torch.Tensor], List[int], List[Any], TreeSpec]:
	@@ -277,7 +277,7 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
277	if indexed_tensors is not None and len(indexed_tensors) > 0:	277	if indexed_tensors is not None and len(indexed_tensors) > 0:
278	indices, tensors = zip(*indexed_tensors)	278	indices, tensors = zip(*indexed_tensors)
279	return tensors, indices, leaves, spec	279	return tensors, indices, leaves, spec
280	-	280	+
281	def unflatten_from_tensors(	281	def unflatten_from_tensors(
282	self,	282	self,
283	tensors: List[torch.Tensor],	283	tensors: List[torch.Tensor],
	@@ -291,7 +291,7 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
291		291
292	def _process_inputs(self, args: Tuple, kwargs: dict) -> List[torch.Tensor]:	292	def _process_inputs(self, args: Tuple, kwargs: dict) -> List[torch.Tensor]:
293	return self.flatten_to_tensors((args, kwargs))	293	return self.flatten_to_tensors((args, kwargs))
294	-	294	+
295	def _recover_inputs(	295	def _recover_inputs(
296	self,	296	self,
297	transform_res: List[List[torch.Tensor]],	297	transform_res: List[List[torch.Tensor]],
	@@ -303,7 +303,7 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
303	for processd_tensors in transform_res:	303	for processd_tensors in transform_res:
304	res.append(self.unflatten_from_tensors(processd_tensors, indices, list(leaves), spec))	304	res.append(self.unflatten_from_tensors(processd_tensors, indices, list(leaves), spec))
305	return zip(*res)	305	return zip(*res)
306	-	306	+
307	def _process_outputs(	307	def _process_outputs(
308	self,	308	self,
309	outputs_list: List[Any]	309	outputs_list: List[Any]
	@@ -332,10 +332,10 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
332	) -> Any:	332	) -> Any:
333	"""	333	"""
334	Process input groups through recovery pipeline.	334	Process input groups through recovery pipeline.
335	-	335	+
336	Args:	336	Args:
337	groups: List of input data to be processed.	337	groups: List of input data to be processed.
338	-	338	+
339	Returns:	339	Returns:
340	Processed outputs after recovery and postprocessing.	340	Processed outputs after recovery and postprocessing.
341	"""	341	"""
	@@ -344,16 +344,16 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
344	inputs = self.recover_pre_fn(groups)	344	inputs = self.recover_pre_fn(groups)
345	else:	345	else:
346	inputs, indices, leaves, spec = self._process_outputs(groups)	346	inputs, indices, leaves, spec = self._process_outputs(groups)
347	-	347	+
348	# 执行恢复操作	348	# 执行恢复操作
349	re_outputs = self.recover(tensor_groups=inputs)	349	re_outputs = self.recover(tensor_groups=inputs)
350	-	350	+
351	# 后处理：使用自定义函数或默认方法	351	# 后处理：使用自定义函数或默认方法
352	if self.recover_post_fn:	352	if self.recover_post_fn:
353	outputs = self.recover_post_fn(re_outputs)	353	outputs = self.recover_post_fn(re_outputs)
354	else:	354	else:
355	outputs = self._recover_outputs(re_outputs, indices, leaves, spec)	355	outputs = self._recover_outputs(re_outputs, indices, leaves, spec)
356	-	356	+
357	return outputs	357	return outputs
358		358
359		359
	@@ -361,28 +361,28 @@ def unified_copy(data: Any) -> Any:
361	"""	361	"""
362	对输入数据进行安全且统一的深拷贝。	362	对输入数据进行安全且统一的深拷贝。
363	支持PyTorch Tensor、字典、列表等常见数据类型。	363	支持PyTorch Tensor、字典、列表等常见数据类型。
364	-	364	+
365	Args:	365	Args:
366	data: 输入数据，可以是Tensor、dict、list等	366	data: 输入数据，可以是Tensor、dict、list等
367	-	367	+
368	Returns:	368	Returns:
369	数据的独立副本	369	数据的独立副本
370	"""	370	"""
371	if data is None:	371	if data is None:
372	return None	372	return None
373	-	373	+
374	# 处理PyTorch Tensor	374	# 处理PyTorch Tensor
375	if isinstance(data, torch.Tensor):	375	if isinstance(data, torch.Tensor):
376	return data.clone().detach()	376	return data.clone().detach()
377	-	377	+
378	# 处理字典类型	378	# 处理字典类型
379	elif isinstance(data, dict):	379	elif isinstance(data, dict):
380	return {key: unified_copy(value) for key, value in data.items()}	380	return {key: unified_copy(value) for key, value in data.items()}
381	-	381	+
382	# 处理列表类型	382	# 处理列表类型
383	elif isinstance(data, list):	383	elif isinstance(data, list):
384	return [unified_copy(item) for item in data]	384	return [unified_copy(item) for item in data]
385	-	385	+
386	# 处理元组类型	386	# 处理元组类型
387	elif isinstance(data, tuple):	387	elif isinstance(data, tuple):
388	return tuple(unified_copy(item) for item in data)	388	return tuple(unified_copy(item) for item in data)
	@@ -411,21 +411,21 @@ def patch_dynamo_context():
411	"""	411	"""
412	if compiler_config is None or not compiler_config.get("enable_shape_handling", False):	412	if compiler_config is None or not compiler_config.get("enable_shape_handling", False):
413	return False	413	return False
414	-	414	+
415	if not isinstance(callback, CatchErrorsWrapper):	415	if not isinstance(callback, CatchErrorsWrapper):
416	return False	416	return False
417	-	417	+
418	orig_callable = callback._torchdynamo_orig_callable	418	orig_callable = callback._torchdynamo_orig_callable
419	if not isinstance(orig_callable, ConvertFrame):	419	if not isinstance(orig_callable, ConvertFrame):
420	return False	420	return False
421	-	421	+
422	deep_callable = orig_callable._torchdynamo_orig_callable	422	deep_callable = orig_callable._torchdynamo_orig_callable
423	if not isinstance(deep_callable, WrapBackendDebug):	423	if not isinstance(deep_callable, WrapBackendDebug):
424	return False	424	return False
425	-	425	+
426	if getattr(deep_callable, "_compiler_name", None) != "inductor":	426	if getattr(deep_callable, "_compiler_name", None) != "inductor":
427	return False	427	return False
428	-	428	+
429	return True	429	return True
430		430
431	def nothing():	431	def nothing():
	@@ -464,7 +464,7 @@ def patch_dynamo_context():
464	trans_post_fn = function_dict.get("trans_post_fn", None)	464	trans_post_fn = function_dict.get("trans_post_fn", None)
465	re_pre_fn = function_dict.get("re_pre_fn", None)	465	re_pre_fn = function_dict.get("re_pre_fn", None)
466	re_post_fn = function_dict.get("re_post_fn", None)	466	re_post_fn = function_dict.get("re_post_fn", None)
467	-	467	+
468	self.shape_handling = NPUShapeHandling(	468	self.shape_handling = NPUShapeHandling(
469	configs=compiler_config.get("shape_handling_configs"),	469	configs=compiler_config.get("shape_handling_configs"),
470	transform_pre_fn=trans_pre_fn,	470	transform_pre_fn=trans_pre_fn,
	@@ -477,7 +477,7 @@ def patch_dynamo_context():
477	src_fn = src_call(self, fn)	477	src_fn = src_call(self, fn)
478	if isinstance(fn, torch.nn.Module) or inspect.isclass(fn):	478	if isinstance(fn, torch.nn.Module) or inspect.isclass(fn):
479	return src_fn	479	return src_fn
480	-	480	+
481	def new_fn(args, *kwargs):	481	def new_fn(args, *kwargs):
482	if (is_enable_shape_handling(self.callback, compiler_config=self.compiler_config)):	482	if (is_enable_shape_handling(self.callback, compiler_config=self.compiler_config)):
483	new_args, new_kwargs = self.shape_handling.transform_hook(args, *kwargs)	483	new_args, new_kwargs = self.shape_handling.transform_hook(args, *kwargs)

Mtorch_npu/_inductor/tools/aten_op_tool.py+13-13

Mtorch_npu/_inductor/tools/fallback_list_tool.py+1-1

	@@ -12,6 +12,6 @@ def get_npu_fallback_list():
12	return npu_fallback_list	12	return npu_fallback_list
13		13
14		14
15	-if __name__ == "__main__":	15	+if __name__ == "__main__":
16	npu_fallback_list = get_npu_fallback_list()	16	npu_fallback_list = get_npu_fallback_list()
17	write_to_file(sorted(npu_fallback_list), f"npu_fallback_list_{len(npu_fallback_list)}.txt")	17	write_to_file(sorted(npu_fallback_list), f"npu_fallback_list_{len(npu_fallback_list)}.txt")

Mtorch_npu/asd/_silent_fault_data.py+1-1

	@@ -11,7 +11,7 @@ def parse_thresh(env_var_name, default_value, min_value):
11	thresh = [max(int(value), min_value) for value in thresh]	11	thresh = [max(int(value), min_value) for value in thresh]
12	if thresh[0] <= thresh[1]:	12	if thresh[0] <= thresh[1]:
13	thresh = [int(value) for value in default_value.split(",")]	13	thresh = [int(value) for value in default_value.split(",")]
14	-	14	+
15	return thresh	15	return thresh
16		16
17		17

Mtorch_npu/asd/asd.py+8-8

	@@ -55,7 +55,7 @@ class _SilentFaultDetector:
55	self.high_step = torch.tensor(self.min_step, dtype=torch.int32).npu()	55	self.high_step = torch.tensor(self.min_step, dtype=torch.int32).npu()
56	if grad.dtype == torch.float16:	56	if grad.dtype == torch.float16:
57	if not self.set_loss_scale_flag:	57	if not self.set_loss_scale_flag:
58	- return	58	+ return
59	else:	59	else:
60	grad = grad.float() / self.loss_scale	60	grad = grad.float() / self.loss_scale
61		61
	@@ -66,7 +66,7 @@ class _SilentFaultDetector:
66	self.silent_data_dict[idx] = SilentFaultData()	66	self.silent_data_dict[idx] = SilentFaultData()
67		67
68	sfda = self.silent_data_dict[idx]	68	sfda = self.silent_data_dict[idx]
69	-	69	+
70	if self.global_step <= self.min_step:	70	if self.global_step <= self.min_step:
71	self.step += 1	71	self.step += 1
72	self.global_step = self.step // (len(self.silent_data_dict) + 1)	72	self.global_step = self.step // (len(self.silent_data_dict) + 1)
	@@ -427,13 +427,13 @@ class _MatmulSilentCheck:
427	def parameter_filtering(self):	427	def parameter_filtering(self):
428	self.filter_index = (self.filter_index + 1) % self.filter_interval	428	self.filter_index = (self.filter_index + 1) % self.filter_interval
429	return self.filter_index == 0	429	return self.filter_index == 0
430	-	430	+
431	def register_module_hook(self, module, name):	431	def register_module_hook(self, module, name):
432	self.check_stat[name + "_backward"] = {'avg': 0, 'pre_val': 0, 'step': 0, 'none_zero_step': 0}	432	self.check_stat[name + "_backward"] = {'avg': 0, 'pre_val': 0, 'step': 0, 'none_zero_step': 0}
433	hook = partial(self.module_hook, name=name + "_backward")	433	hook = partial(self.module_hook, name=name + "_backward")
434	self.hook_dict[name + "_backward"] = module.register_full_backward_hook(hook)	434	self.hook_dict[name + "_backward"] = module.register_full_backward_hook(hook)
435	self.registered_modules.append(name)	435	self.registered_modules.append(name)
436	-	436	+
437	def module_hook(self, module, grad_input, grad_output, name):	437	def module_hook(self, module, grad_input, grad_output, name):
438	for _, param in module.named_parameters():	438	for _, param in module.named_parameters():
439	if param.dim() >= 2:	439	if param.dim() >= 2:
	@@ -706,7 +706,7 @@ class _MatmulSilentCheck:
706		706
707	while int(self.store.get('counter2').decode()) < world_size and self.checksum_state_thread_running:	707	while int(self.store.get('counter2').decode()) < world_size and self.checksum_state_thread_running:
708	time.sleep(0.1)	708	time.sleep(0.1)
709	-	709	+
710	if self.rank == 0:	710	if self.rank == 0:
711	self.store.add('checksum_state', 0 - global_state)	711	self.store.add('checksum_state', 0 - global_state)
712	self.store.add('counter', 0 - world_size)	712	self.store.add('counter', 0 - world_size)
	@@ -719,12 +719,12 @@ class _MatmulSilentCheck:
719	state['_lock'] = None	719	state['_lock'] = None
720	state['store'] = None	720	state['store'] = None
721	return state	721	return state
722	-	722	+
723	def __setstate(self, state):	723	def __setstate(self, state):
724	self.__dict__.update(state)	724	self.__dict__.update(state)
725	self.store = None	725	self.store = None
726		726
727	- def _startup(self):	727	+ def _startup(self):
728	if not self.check_thread_running:	728	if not self.check_thread_running:
729	self.check_thread_running = True	729	self.check_thread_running = True
730	self.check_thread = threading.Thread(	730	self.check_thread = threading.Thread(
	@@ -826,7 +826,7 @@ def _matmul_silent_check_decorator(func):
826	matmul_check.init_marks[matmul_check.first_module_id] = True	826	matmul_check.init_marks[matmul_check.first_module_id] = True
827		827
828	tmp = func(self, args, *kwargs)	828	tmp = func(self, args, *kwargs)
829	-	829	+
830	if matmul_check.get_matmul_hook_enable():	830	if matmul_check.get_matmul_hook_enable():
831	if hasattr(self, "matmul_check_outer") and self.matmul_check_outer:	831	if hasattr(self, "matmul_check_outer") and self.matmul_check_outer:
832	matmul_check.init_param()	832	matmul_check.init_param()

Mtorch_npu/csrc/aten/ops/CopyFromAndResizeKernelNpu.cpp+1-1

	@@ -8,7 +8,7 @@ at::Tensor NPUNativeFunctions::_copy_from_and_resize(const at::Tensor& self, con
8	{	8	{
9	TORCH_CHECK(dst.defined(), "dst is undefined", OPS_ERROR(ErrCode::NOT_SUPPORT));	9	TORCH_CHECK(dst.defined(), "dst is undefined", OPS_ERROR(ErrCode::NOT_SUPPORT));
10	TORCH_CHECK(self.defined(), "self is undefined", OPS_ERROR(ErrCode::NOT_SUPPORT));	10	TORCH_CHECK(self.defined(), "self is undefined", OPS_ERROR(ErrCode::NOT_SUPPORT));
11	-	11	+
12	if (dst.numel() == 0) {	12	if (dst.numel() == 0) {
13	dst.resize_as_(self);	13	dst.resize_as_(self);
14	}	14	}

Mtorch_npu/csrc/core/npu/GetCANNInfo.cpp+5-5

	@@ -253,7 +253,7 @@ int64_t VersionV2ToNum(std::string versionStr) {
253	parsed = true;	253	parsed = true;
254	}	254	}
255		255
256	- if (!parsed && tokens.size() == tokenNum4) { // ([0-9]+).([0-9]+).([0-9]+)-alpha, ([0-9]+).([0-9]+).([0-9]+)-beta, ([0-9]+).([0-9]+).([0-9]+)-rc,	256	+ if (!parsed && tokens.size() == tokenNum4) { // ([0-9]+).([0-9]+).([0-9]+)-alpha, ([0-9]+).([0-9]+).([0-9]+)-beta, ([0-9]+).([0-9]+).([0-9]+)-rc,
257	// ([0-9]+).([0-9]+).([0-9]+).alpha([0-9]+), ([0-9]+).([0-9]+).([0-9]+).beta([0-9]+), ([0-9]+).([0-9]+).([0-9]+).rc([0-9]+)	257	// ([0-9]+).([0-9]+).([0-9]+).alpha([0-9]+), ([0-9]+).([0-9]+).([0-9]+).beta([0-9]+), ([0-9]+).([0-9]+).([0-9]+).rc([0-9]+)
258	parsed = true;	258	parsed = true;
259	if (tokens[index3] == "alpha") {	259	if (tokens[index3] == "alpha") {
	@@ -469,7 +469,7 @@ std::string GetCANNVersion(const std::string& module)
469	module_version = version.version;	469	module_version = version.version;
470	CANNVersionCache[module] = module_version;	470	CANNVersionCache[module] = module_version;
471	}	471	}
472	-	472	+
473	if (find_module_v2 != pkgNameV2Map.end()) {	473	if (find_module_v2 != pkgNameV2Map.end()) {
474	char versionStr[ACL_PKG_VERSION_MAX_SIZE] = {0};	474	char versionStr[ACL_PKG_VERSION_MAX_SIZE] = {0};
475	aclError retV2 = c10_npu::acl::AclsysGetVersionStr(const_cast<char*>(module.c_str()), versionStr);	475	aclError retV2 = c10_npu::acl::AclsysGetVersionStr(const_cast<char*>(module.c_str()), versionStr);
	@@ -481,7 +481,7 @@ std::string GetCANNVersion(const std::string& module)
481	module_version = versionStr;	481	module_version = versionStr;
482	CANNVersionCache[module] = module_version;	482	CANNVersionCache[module] = module_version;
483	}	483	}
484	-	484	+
485	return module_version;	485	return module_version;
486	}	486	}
487		487
	@@ -492,7 +492,7 @@ bool IsGteCANNVersion(const std::string version, const std::string module)
492	if (module.compare(unsupportedModule) == 0) {	492	if (module.compare(unsupportedModule) == 0) {
493	TORCH_CHECK(false, "When the module is DRIVER, IsGteCANNVersion is not supported. ", PTA_ERROR(ErrCode::VALUE));	493	TORCH_CHECK(false, "When the module is DRIVER, IsGteCANNVersion is not supported. ", PTA_ERROR(ErrCode::VALUE));
494	}	494	}
495	-	495	+
496	std::vector<std::string> tokensVersion = SplitVersionStr(version);	496	std::vector<std::string> tokensVersion = SplitVersionStr(version);
497	std::vector<std::string> tokensBaseVersion = SplitVersionStr(baseVersion);	497	std::vector<std::string> tokensBaseVersion = SplitVersionStr(baseVersion);
498	if (tokensVersion.size() < tokenNum2 \|\| !isDigits(tokensVersion[index0])) {	498	if (tokensVersion.size() < tokenNum2 \|\| !isDigits(tokensVersion[index0])) {
	@@ -507,7 +507,7 @@ bool IsGteCANNVersion(const std::string version, const std::string module)
507		507
508	std::string currentVersion = GetCANNVersion(module);	508	std::string currentVersion = GetCANNVersion(module);
509	std::vector<std::string> tokensCurrentVersion = SplitVersionStr(currentVersion);	509	std::vector<std::string> tokensCurrentVersion = SplitVersionStr(currentVersion);
510	-	510	+
511	int64_t current_num = 0;	511	int64_t current_num = 0;
512	int64_t boundary_num = 0;	512	int64_t boundary_num = 0;
513	bool isInvalid = false;	513	bool isInvalid = false;

Mtorch_npu/csrc/core/npu/NPUCachingAllocator.cpp+1-1

	@@ -3575,7 +3575,7 @@ public:
3575	TORCH_NPU_MEMORY_LOGE("%s", retmsg.c_str());	3575	TORCH_NPU_MEMORY_LOGE("%s", retmsg.c_str());
3576	TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str());	3576	TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str());
3577	}	3577	}
3578	-	3578	+
3579	int device = 0;	3579	int device = 0;
3580	NPU_CHECK_ERROR(c10_npu::GetDevice(&device));	3580	NPU_CHECK_ERROR(c10_npu::GetDevice(&device));
3581	LazySetDevice(device);	3581	LazySetDevice(device);

Mtorch_npu/csrc/core/npu/NPUErrorCodes.h+1-1

	@@ -3,7 +3,7 @@
3	#include <unordered_map>	3	#include <unordered_map>
4		4
5	namespace c10_npu::acl {	5	namespace c10_npu::acl {
6	-	6	+
7	class AclErrorCode {	7	class AclErrorCode {
8	public:	8	public:
9	std::unordered_map<int, std::string> error_code_map = {	9	std::unordered_map<int, std::string> error_code_map = {

Mtorch_npu/csrc/core/npu/NPUPeerToPeerAccess.cpp+1-1

	@@ -50,7 +50,7 @@ bool NpuP2pCtrl::get_p2p_access(int32_t source_dev, int32_t dest_dev, bool& flag
50	// get access source_dev -> dest_dev	50	// get access source_dev -> dest_dev
51	auto &cache_s2d = p2p_access_enabled_cache_[source_dev * num_devices_ + dest_dev];	51	auto &cache_s2d = p2p_access_enabled_cache_[source_dev * num_devices_ + dest_dev];
52	auto &cache_d2s = p2p_access_enabled_cache_[dest_dev * num_devices_ + source_dev];	52	auto &cache_d2s = p2p_access_enabled_cache_[dest_dev * num_devices_ + source_dev];
53	-	53	+
54	if (cache_s2d != P2pStatus::UNKONWN) {	54	if (cache_s2d != P2pStatus::UNKONWN) {
55	return static_cast<bool>(cache_s2d);	55	return static_cast<bool>(cache_s2d);
56	}	56	}

Mtorch_npu/csrc/core/npu/NPUQueue.cpp+1-1

	@@ -620,7 +620,7 @@ void Repository::Enqueue(void *cur_paras)
620	#ifndef BUILD_LIBTORCH	620	#ifndef BUILD_LIBTORCH
621	// double check the current thread hold a Gil lock	621	// double check the current thread hold a Gil lock
622	// and release the GIL to TE op compiler in case the acl thread deadlock.	622	// and release the GIL to TE op compiler in case the acl thread deadlock.
623	- // However, this operator could produce another form of deadlock.	623	+ // However, this operator could produce another form of deadlock.
624	// When thread A deconstract a tensor, it will hold the mutex of deviceCachingAllocator and insert an event into the taskqueue.	624	// When thread A deconstract a tensor, it will hold the mutex of deviceCachingAllocator and insert an event into the taskqueue.
625	// If the taskqueue is full, thead A will run into here and release the GIL.	625	// If the taskqueue is full, thead A will run into here and release the GIL.
626	// Once another thread B get GIL and trigger GC, it may deconstract another tensor	626	// Once another thread B get GIL and trigger GC, it may deconstract another tensor

Mtorch_npu/csrc/core/npu/NPUWorkspaceAllocator.cpp+1-1

	@@ -512,7 +512,7 @@ public:
512	if (c10_npu::option::OptionsManager::CheckForceUncached() &&	512	if (c10_npu::option::OptionsManager::CheckForceUncached() &&
513	(c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None)) {	513	(c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None)) {
514	return &uncached_delete;	514	return &uncached_delete;
515	- }	515	+ }
516	return &local_raw_delete;	516	return &local_raw_delete;
517	}	517	}
518		518

Mtorch_npu/csrc/core/npu/NpuVariables.cpp+1-1

	@@ -50,7 +50,7 @@ void SetSocVersion(const char* const socVersion)
50	SocVersion curSocVersion = SocVersion::UnsupportedSocVersion;	50	SocVersion curSocVersion = SocVersion::UnsupportedSocVersion;
51	std::string inputVersion = socVersion;	51	std::string inputVersion = socVersion;
52	std::string ascend950 = "Ascend950";	52	std::string ascend950 = "Ascend950";
53	-	53	+
54	auto const& iter = socVersionMap.find(socVersion);	54	auto const& iter = socVersionMap.find(socVersion);
55	if (iter != socVersionMap.end()) {	55	if (iter != socVersionMap.end()) {
56	curSocVersion = iter->second;	56	curSocVersion = iter->second;

Mtorch_npu/csrc/core/npu/register/OptionsManager.cpp+1-1

	@@ -520,7 +520,7 @@ uint32_t OptionsManager::GetAclOpInitMode()
520	} else {	520	} else {
521	acl_op_init_mode_ = (buf_val != nullptr) ? strtol(buf_val, nullptr, 10) : 0;	521	acl_op_init_mode_ = (buf_val != nullptr) ? strtol(buf_val, nullptr, 10) : 0;
522	}	522	}
523	-	523	+
524	std::unordered_map<int32_t, std::string> aclOpInitMode = getAclOpInitMode();	524	std::unordered_map<int32_t, std::string> aclOpInitMode = getAclOpInitMode();
525	if (aclOpInitMode.find(acl_op_init_mode_) == aclOpInitMode.end()) {	525	if (aclOpInitMode.find(acl_op_init_mode_) == aclOpInitMode.end()) {
526	if (default_value_acl_mode && isCannVersionGteBase) {	526	if (default_value_acl_mode && isCannVersionGteBase) {

Mtorch_npu/csrc/distributed/HCCLUtils.cpp+1-1

	@@ -112,7 +112,7 @@ bool isSupportHcclCommName()
112		112
113	HCCLComm::HCCLComm(HcclComm hcclComm) : hcclComm_(hcclComm), hcclAsyncErr_(HCCL_SUCCESS),	113	HCCLComm::HCCLComm(HcclComm hcclComm) : hcclComm_(hcclComm), hcclAsyncErr_(HCCL_SUCCESS),
114	hcclCommType(0), p2pPeer(0) {}	114	hcclCommType(0), p2pPeer(0) {}
115	-	115	+
116	HCCLComm::~HCCLComm()	116	HCCLComm::~HCCLComm()
117	{	117	{
118	destroyHcclComm();	118	destroyHcclComm();

Mtorch_npu/csrc/distributed/Init.cpp+2-2

	@@ -210,7 +210,7 @@ PyObject* c10d_npu_init(PyObject* _unused, PyObject* noargs)
210	throw python_error();	210	throw python_error();
211	}	211	}
212	auto torch_npu_C_m = py::handle(torch_npu_C_module).cast<py::module>();	212	auto torch_npu_C_m = py::handle(torch_npu_C_module).cast<py::module>();
213	-	213	+
214	auto m =	214	auto m =
215	torch_npu_C_m.def_submodule("_distributed_c10d", "distributed c10d bindings");	215	torch_npu_C_m.def_submodule("_distributed_c10d", "distributed c10d bindings");
216	auto module = py::handle(m).cast<py::module>();	216	auto module = py::handle(m).cast<py::module>();
	@@ -473,7 +473,7 @@ PyObject* c10d_npu_init(PyObject* _unused, PyObject* noargs)
473	.def_readwrite("hccl_config", &::c10d_npu::ProcessGroupHCCL::Options::hccl_config)	473	.def_readwrite("hccl_config", &::c10d_npu::ProcessGroupHCCL::Options::hccl_config)
474	.def_readwrite("group_id",	474	.def_readwrite("group_id",
475	&::c10d_npu::ProcessGroupHCCL::Options::group_id);	475	&::c10d_npu::ProcessGroupHCCL::Options::group_id);
476	-	476	+
477	// bind for ProcessGroupLCCL	477	// bind for ProcessGroupLCCL
478	auto processGroupLCCL = intrusive_ptr_no_gil_destructor_class_<::c10d_npu::ProcessGroupLCCL>(	478	auto processGroupLCCL = intrusive_ptr_no_gil_destructor_class_<::c10d_npu::ProcessGroupLCCL>(
479	module, "ProcessGroupLCCL", dist.attr("Backend"))	479	module, "ProcessGroupLCCL", dist.attr("Backend"))

Mtorch_npu/csrc/distributed/ParallelTcpServer.cpp+3-3

	@@ -232,7 +232,7 @@ int ParallelTcpServer::CreateSocket(const std::string host, uint16_t port) noexc
232	if (sockFd >= 0) {	232	if (sockFd >= 0) {
233	return sockFd;	233	return sockFd;
234	}	234	}
235	-	235	+
236	sockFd = CreateSocketWithFamily(host, port, AF_INET6);	236	sockFd = CreateSocketWithFamily(host, port, AF_INET6);
237	if (sockFd >= 0) {	237	if (sockFd >= 0) {
238	return sockFd;	238	return sockFd;
	@@ -298,7 +298,7 @@ int ParallelTcpServer::CreateLocalSocket(const std::string &localSocketPath) noe
298	LOG(ERROR) << "local socket path invalid." << errno << " : " << strerror(errno);	298	LOG(ERROR) << "local socket path invalid." << errno << " : " << strerror(errno);
299	return -1;	299	return -1;
300	}	300	}
301	-	301	+
302	struct sockaddr_un servAddr {};	302	struct sockaddr_un servAddr {};
303	servAddr.sun_family = AF_UNIX;	303	servAddr.sun_family = AF_UNIX;
304	servAddr.sun_path[0] = '\0';	304	servAddr.sun_path[0] = '\0';
	@@ -386,7 +386,7 @@ int ParallelTcpServer::SetBlockSocketTimeout(int fd) noexcept
386	LOG(ERROR) << "set block accept timeout failed " << errno << " : " << strerror(errno);	386	LOG(ERROR) << "set block accept timeout failed " << errno << " : " << strerror(errno);
387	return -1;	387	return -1;
388	}	388	}
389	-	389	+
390	return 0;	390	return 0;
391	}	391	}
392		392

Mtorch_npu/csrc/distributed/ProcessGroupHCCL.cpp+13-13

	@@ -154,7 +154,7 @@ HcclReduceOp getHcclReduceOp(const c10d::ReduceOp reduceOp, at::Tensor& input)
154	// represent a bool (see hcclDataType mapping).	154	// represent a bool (see hcclDataType mapping).
155	return HCCL_REDUCE_MAX;	155	return HCCL_REDUCE_MAX;
156	}	156	}
157	-	157	+
158	if (unsupportedOp.find(reduceOp) != unsupportedOp.end()) {	158	if (unsupportedOp.find(reduceOp) != unsupportedOp.end()) {
159	TORCH_CHECK(false,	159	TORCH_CHECK(false,
160	"Cannot use ReduceOp." + unsupportedOp[reduceOp] + " with HCCL",	160	"Cannot use ReduceOp." + unsupportedOp[reduceOp] + " with HCCL",
	@@ -1289,7 +1289,7 @@ void ProcessGroupHCCL::waitForFutureOrTimeout(
1289	void ProcessGroupHCCL::shutdown()	1289	void ProcessGroupHCCL::shutdown()
1290	{	1290	{
1291	LOG(INFO) << logPrefix() << "Starting to destroy process group, flushing operations.";	1291	LOG(INFO) << logPrefix() << "Starting to destroy process group, flushing operations.";
1292	-	1292	+
1293	if (terminateProcessGroup_.exchange(true)) {	1293	if (terminateProcessGroup_.exchange(true)) {
1294	return;	1294	return;
1295	}	1295	}
	@@ -1377,7 +1377,7 @@ void ProcessGroupHCCL::deleteTCPStoreKey()
1377	}	1377	}
1378		1378
1379	TORCH_NPU_HCCL_LOGI("Delete TCP store key success.");	1379	TORCH_NPU_HCCL_LOGI("Delete TCP store key success.");
1380	-	1380	+
1381	TCPStoreKeyList_.clear();	1381	TCPStoreKeyList_.clear();
1382	}	1382	}
1383		1383
	@@ -1951,7 +1951,7 @@ void ProcessGroupHCCL::logWorkEnd(WorkHCCL& work)
1951		1951
1952	storeError_ = !c10d::traceUpdate(store_, traceKeyEnd_, work.seq_, opTypeToString(work.opType_));	1952	storeError_ = !c10d::traceUpdate(store_, traceKeyEnd_, work.seq_, opTypeToString(work.opType_));
1953	}	1953	}
1954	-	1954	+
1955	std::string ProcessGroupHCCL::createLogPrefix() const	1955	std::string ProcessGroupHCCL::createLogPrefix() const
1956	{	1956	{
1957	if (!pg_desc_.empty() && pg_desc_ != "undefined") {	1957	if (!pg_desc_.empty() && pg_desc_ != "undefined") {
	@@ -2040,7 +2040,7 @@ void ProcessGroupHCCL::Watchdog::runLoop()
2040	auto timenow = std::chrono::steady_clock::now();	2040	auto timenow = std::chrono::steady_clock::now();
2041	bool recordflag = false;	2041	bool recordflag = false;
2042	int kThousandMillis = 1000;	2042	int kThousandMillis = 1000;
2043	-	2043	+
2044	while (!pg_->terminateProcessGroup_.load()) {	2044	while (!pg_->terminateProcessGroup_.load()) {
2045	if (status_save_enable) {	2045	if (status_save_enable) {
2046	checkAndMakePath(status_save_path.c_str(), "Open shared directory failed. Please check whether input path is valid.");	2046	checkAndMakePath(status_save_path.c_str(), "Open shared directory failed. Please check whether input path is valid.");
	@@ -2085,7 +2085,7 @@ void ProcessGroupHCCL::Watchdog::runLoop()
2085	TORCH_NPU_HCCL_LOGI("Find FORCE STOP when runloop setDevice.");	2085	TORCH_NPU_HCCL_LOGI("Find FORCE STOP when runloop setDevice.");
2086	}	2086	}
2087	}	2087	}
2088	-	2088	+
2089	// check NCCL errors first	2089	// check NCCL errors first
2090	if (!pg_->terminateProcessGroup_.load()) {	2090	if (!pg_->terminateProcessGroup_.load()) {
2091	work.checkAndSetException();	2091	work.checkAndSetException();
	@@ -3960,7 +3960,7 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::collective(
3960		3960
3961	const std::vector<uint32_t>& ranks = groupRanks();	3961	const std::vector<uint32_t>& ranks = groupRanks();
3962	outfile << "[GLOBAL RANKID]:" << ranks[rank_] << "\n";	3962	outfile << "[GLOBAL RANKID]:" << ranks[rank_] << "\n";
3963	-	3963	+
3964	outfile.close();	3964	outfile.close();
3965	}	3965	}
3966	} else {	3966	} else {
	@@ -4064,7 +4064,7 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::collective(
4064	} else {	4064	} else {
4065	c10_npu::NPUGraph::dec_pending_event_queries();	4065	c10_npu::NPUGraph::dec_pending_event_queries();
4066	}	4066	}
4067	-	4067	+
4068	return work;	4068	return work;
4069	}	4069	}
4070		4070
	@@ -4196,7 +4196,7 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::collectiveCoalesced(
4196		4196
4197	const std::vector<uint32_t>& ranks = groupRanks();	4197	const std::vector<uint32_t>& ranks = groupRanks();
4198	outfile << "[GLOBAL RANKID]:" << ranks[rank_] << "\n";	4198	outfile << "[GLOBAL RANKID]:" << ranks[rank_] << "\n";
4199	-	4199	+
4200	outfile.close();	4200	outfile.close();
4201	}	4201	}
4202	} else {	4202	} else {
	@@ -4281,7 +4281,7 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::collectiveCoalesced(
4281	} else {	4281	} else {
4282	c10_npu::NPUGraph::dec_pending_event_queries();	4282	c10_npu::NPUGraph::dec_pending_event_queries();
4283	}	4283	}
4284	-	4284	+
4285	return work;	4285	return work;
4286	}	4286	}
4287		4287
	@@ -4352,7 +4352,7 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::pointToPoint(
4352	"Got device ", device.index(), " but expected ", coalescedDevice_.index());	4352	"Got device ", device.index(), " but expected ", coalescedDevice_.index());
4353	}	4353	}
4354	}	4354	}
4355	-	4355	+
4356	// Verify communicator consistency	4356	// Verify communicator consistency
4357	if (coalescedComm_ == nullptr) {	4357	if (coalescedComm_ == nullptr) {
4358	coalescedComm_ = hcclComms[0];	4358	coalescedComm_ = hcclComms[0];
	@@ -4452,7 +4452,7 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::pointToPoint(
4452		4452
4453	const std::vector<uint32_t>& ranks = groupRanks();	4453	const std::vector<uint32_t>& ranks = groupRanks();
4454	outfile << "[GLOBAL RANKID]:" << ranks[rank_] << "\n";	4454	outfile << "[GLOBAL RANKID]:" << ranks[rank_] << "\n";
4455	-	4455	+
4456	outfile.close();	4456	outfile.close();
4457	}	4457	}
4458	} else {	4458	} else {
	@@ -4532,7 +4532,7 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::pointToPoint(
4532	// as multi-device per process is deprecated	4532	// as multi-device per process is deprecated
4533	work->numelIn_ = work->numelOut_ = static_cast<size_t>(tensors[i].numel());	4533	work->numelIn_ = work->numelOut_ = static_cast<size_t>(tensors[i].numel());
4534	}	4534	}
4535	-	4535	+
4536	c10_npu::NPUGraph::inc_pending_event_queries();	4536	c10_npu::NPUGraph::inc_pending_event_queries();
4537	if (asyncErrorHandling_ != NoHandling && capture_status == c10_npu::CaptureStatus::None) {	4537	if (asyncErrorHandling_ != NoHandling && capture_status == c10_npu::CaptureStatus::None) {
4538	workEnqueue(work);	4538	workEnqueue(work);

Mtorch_npu/csrc/distributed/ProcessGroupHCCL.hpp+12-12

	@@ -400,7 +400,7 @@ public:
400	// for tests.	400	// for tests.
401	virtual std::exception_ptr checkForHCCLErrors(	401	virtual std::exception_ptr checkForHCCLErrors(
402	const std::vector<std::shared_ptr<HCCLComm>>& hcclComms) const;	402	const std::vector<std::shared_ptr<HCCLComm>>& hcclComms) const;
403	-	403	+
404	friend std::ostream& operator<<(	404	friend std::ostream& operator<<(
405	std::ostream& output,	405	std::ostream& output,
406	const WorkHCCL& workHCCL);	406	const WorkHCCL& workHCCL);
	@@ -452,11 +452,11 @@ public:
452	std::vector<std::pair<c10::weak_intrusive_ptr<c10::StorageImpl>, c10_npu::NPUStream>> recorded_outputs_;	452	std::vector<std::pair<c10::weak_intrusive_ptr<c10::StorageImpl>, c10_npu::NPUStream>> recorded_outputs_;
453		453
454	std::vector<at::Tensor> lazy_destroy_tensors_;	454	std::vector<at::Tensor> lazy_destroy_tensors_;
455	-	455	+
456	// unique id used to tell the trace buffer that this	456	// unique id used to tell the trace buffer that this
457	// work has completed	457	// work has completed
458	c10::optional<uint64_t> trace_id_;	458	c10::optional<uint64_t> trace_id_;
459	-	459	+
460	mutable std::once_flag print_flag;	460	mutable std::once_flag print_flag;
461		461
462	friend class ProcessGroupHCCL;	462	friend class ProcessGroupHCCL;
	@@ -598,7 +598,7 @@ public:
598	{	598	{
599	return std::string(HCCL_BACKEND_NAME);	599	return std::string(HCCL_BACKEND_NAME);
600	}	600	}
601	-	601	+
602	bool supportsCoalescing() const override	602	bool supportsCoalescing() const override
603	{	603	{
604	return true;	604	return true;
	@@ -869,7 +869,7 @@ protected:
869	{	869	{
870	return pg_desc_;	870	return pg_desc_;
871	}	871	}
872	-	872	+
873	void setP2pPeer(int newPeer)	873	void setP2pPeer(int newPeer)
874	{	874	{
875	peer_ = newPeer;	875	peer_ = newPeer;
	@@ -879,7 +879,7 @@ protected:
879	{	879	{
880	return peer_;	880	return peer_;
881	}	881	}
882	-	882	+
883	// In the timeout case and we will dump debug info such as the NCCL flight	883	// In the timeout case and we will dump debug info such as the NCCL flight
884	// recorder to storage. Down the road, if we have more complicated or blocking	884	// recorder to storage. Down the road, if we have more complicated or blocking
885	// operations, we might need to use a side thread to do it.	885	// operations, we might need to use a side thread to do it.
	@@ -895,7 +895,7 @@ protected:
895	// we can dump the debugging information and abort the process.	895	// we can dump the debugging information and abort the process.
896	virtual void heartbeatMonitor();	896	virtual void heartbeatMonitor();
897	#endif	897	#endif
898	-	898	+
899	// Instance of the watchdog thread.	899	// Instance of the watchdog thread.
900	std::unique_ptr<Watchdog> watchdog_;	900	std::unique_ptr<Watchdog> watchdog_;
901	// Function that directly trigger std::abort so that the whole process	901	// Function that directly trigger std::abort so that the whole process
	@@ -982,7 +982,7 @@ protected:
982	std::unordered_map<std::string, std::vector<std::shared_ptr<HCCLComm>>> devHCCLCommMap_;	982	std::unordered_map<std::string, std::vector<std::shared_ptr<HCCLComm>>> devHCCLCommMap_;
983		983
984	std::unordered_set<std::string> reportedErrorComms_;	984	std::unordered_set<std::string> reportedErrorComms_;
985	-	985	+
986	std::unordered_map<int, std::vector<std::string>> p2pSendRecvKeys_;	986	std::unordered_map<int, std::vector<std::string>> p2pSendRecvKeys_;
987		987
988	std::unordered_map<std::string, std::string> devHCCLCommNameMap_;	988	std::unordered_map<std::string, std::string> devHCCLCommNameMap_;
	@@ -1075,7 +1075,7 @@ protected:
1075	// The NPU events used to control task rate to protect streams	1075	// The NPU events used to control task rate to protect streams
1076	std::unordered_map<std::string, std::vector<c10_npu::NPUEvent>>	1076	std::unordered_map<std::string, std::vector<c10_npu::NPUEvent>>
1077	rateCtrlEvents_;	1077	rateCtrlEvents_;
1078	-	1078	+
1079	std::unordered_map<std::string, std::vector<uint64_t>> collectiveCnts_;	1079	std::unordered_map<std::string, std::vector<uint64_t>> collectiveCnts_;
1080		1080
1081	// Device Indexes used for all collectives in this group	1081	// Device Indexes used for all collectives in this group
	@@ -1232,7 +1232,7 @@ private:
1232	Fn fn,	1232	Fn fn,
1233	c10d::OpType opType,	1233	c10d::OpType opType,
1234	bool asyncOp = false);	1234	bool asyncOp = false);
1235	-	1235	+
1236	template <typename Fn, typename PreProcess, typename PostProcess>	1236	template <typename Fn, typename PreProcess, typename PostProcess>
1237	c10::intrusive_ptr<c10d::Work> collective(	1237	c10::intrusive_ptr<c10d::Work> collective(
1238	std::vector<at::Tensor>& input,	1238	std::vector<at::Tensor>& input,
	@@ -1355,7 +1355,7 @@ private:
1355	// Util function to assign timeout to each work.	1355	// Util function to assign timeout to each work.
1356	void assignTimeoutToWork(const c10::intrusive_ptr<ProcessGroupHCCL::WorkHCCL>& work,	1356	void assignTimeoutToWork(const c10::intrusive_ptr<ProcessGroupHCCL::WorkHCCL>& work,
1357	const c10::intrusive_ptr<Options>& option);	1357	const c10::intrusive_ptr<Options>& option);
1358	-	1358	+
1359	void silenceCheck(at::Tensor &input, c10d::OpType opType);	1359	void silenceCheck(at::Tensor &input, c10d::OpType opType);
1360		1360
1361	HcclCommConfig createHcclCommConfigWithOptions();	1361	HcclCommConfig createHcclCommConfigWithOptions();
	@@ -1382,7 +1382,7 @@ private:
1382	c10::optional<at::Tensor> windowMem_;	1382	c10::optional<at::Tensor> windowMem_;
1383		1383
1384	uint32_t cached_aic_num;	1384	uint32_t cached_aic_num;
1385	-	1385	+
1386	uint32_t cached_aiv_num;	1386	uint32_t cached_aiv_num;
1387		1387
1388	};	1388	};

Mtorch_npu/csrc/distributed/StoreClient.cpp+3-3

	@@ -30,7 +30,7 @@
30		30
31	namespace c10d {	31	namespace c10d {
32	namespace torch_npu {	32	namespace torch_npu {
33	-	33	+
34	Client::Client(const std::string host, uint16_t port, const std::chrono::milliseconds timeout) noexcept	34	Client::Client(const std::string host, uint16_t port, const std::chrono::milliseconds timeout) noexcept
35	: host_{ host }, port_{ port }, socketFd_(-1), timeout_{ timeout }	35	: host_{ host }, port_{ port }, socketFd_(-1), timeout_{ timeout }
36	{}	36	{}
	@@ -111,7 +111,7 @@ int Client::Connect() noexcept
111	if (ret >= 0) {	111	if (ret >= 0) {
112	return 0;	112	return 0;
113	}	113	}
114	-	114	+
115	ret = TryConnect(AF_INET6);	115	ret = TryConnect(AF_INET6);
116	if (ret >= 0) {	116	if (ret >= 0) {
117	return 0;	117	return 0;
	@@ -221,7 +221,7 @@ int Client::SyncCall(const StoreMessage &request, StoreMessage &response) noexce
221	if (errno == EINTR) { // interrupted by signal	221	if (errno == EINTR) { // interrupted by signal
222	continue;	222	continue;
223	}	223	}
224	-	224	+
225	LOG(ERROR) << "read data from server(" << host_ << ":" << port_ << ") failed " << errno << " : " <<	225	LOG(ERROR) << "read data from server(" << host_ << ":" << port_ << ") failed " << errno << " : " <<
226	strerror(errno);	226	strerror(errno);
227	return -1;	227	return -1;

Mtorch_npu/csrc/framework/OpCommand.cpp+1-1

	@@ -243,7 +243,7 @@ void OpCommand::RunOpApiV2(const string &op_name, const PROC_FUNC &func, bool sy
243	execParams.customHandler = const_cast<PROC_FUNC*>(&func);	243	execParams.customHandler = const_cast<PROC_FUNC*>(&func);
244		244
245	c10_npu::queue::QueueParas params(c10_npu::queue::EXECUTE_OPAPI_V2, sizeof(ExecuteParasOpApiV2), &execParams);	245	c10_npu::queue::QueueParas params(c10_npu::queue::EXECUTE_OPAPI_V2, sizeof(ExecuteParasOpApiV2), &execParams);
246	-	246	+
247	auto start = std::chrono::steady_clock::now();	247	auto start = std::chrono::steady_clock::now();
248		248
249	c10_npu::enCurrentNPUStream(&params);	249	c10_npu::enCurrentNPUStream(&params);

Mtorch_npu/csrc/framework/OpCommand.h+2-2

	@@ -72,7 +72,7 @@ public:
72	toType);	72	toType);
73	return AddHostTensorInput(cpuTensor, compileType, realDtype, descName);	73	return AddHostTensorInput(cpuTensor, compileType, realDtype, descName);
74	}	74	}
75	-	75	+
76	// IntArrayRef/SmallVector Input, usually hostmemory input, we will do h2d in launch kernel	76	// IntArrayRef/SmallVector Input, usually hostmemory input, we will do h2d in launch kernel
77	OpCommand& Input(const c10::IntArrayRef &dimListRef,	77	OpCommand& Input(const c10::IntArrayRef &dimListRef,
78	at::ScalarType toType = at::kLong,	78	at::ScalarType toType = at::kLong,
	@@ -133,7 +133,7 @@ private:
133	OpCommand& AddTensorInput(at::Tensor &tensor,	133	OpCommand& AddTensorInput(at::Tensor &tensor,
134	at::ScalarType forceScaleType = at::ScalarType::Undefined,	134	at::ScalarType forceScaleType = at::ScalarType::Undefined,
135	const string &descName = "", const string &realData = "");	135	const string &descName = "", const string &realData = "");
136	-	136	+
137	OpCommand& AddTensorInput(const string &str);	137	OpCommand& AddTensorInput(const string &str);
138		138
139	OpCommand& AddHostTensorInput(	139	OpCommand& AddHostTensorInput(

Mtorch_npu/csrc/framework/utils/CalcuOpUtil.cpp+1-1

	@@ -390,7 +390,7 @@ at::ScalarType CalcuOpUtil::ConvertToScalarType(const aclDataType data_type)
390	std::string("aclDataType:") + std::to_string(data_type) + " has not been supported",	390	std::string("aclDataType:") + std::to_string(data_type) + " has not been supported",
391	OPS_ERROR(ErrCode::NOT_SUPPORT))	391	OPS_ERROR(ErrCode::NOT_SUPPORT))
392	}	392	}
393	-	393	+
394	return iter->second;	394	return iter->second;
395	}	395	}
396		396

Mtorch_npu/csrc/inductor/aoti_package/shape_handling.cpp+2-2

	@@ -69,7 +69,7 @@ void THNPShapeHandling_init(PyObject *module)
69	.value("BATCHSIZE", torch::aot_inductor::ShapeType::BATCHSIZE)	69	.value("BATCHSIZE", torch::aot_inductor::ShapeType::BATCHSIZE)
70	.value("SEQLEN", torch::aot_inductor::ShapeType::SEQLEN)	70	.value("SEQLEN", torch::aot_inductor::ShapeType::SEQLEN)
71	.export_values();	71	.export_values();
72	-	72	+
73	py::enum_<torch::aot_inductor::ShapePolicy>(torch_N_m, "ShapePolicy")	73	py::enum_<torch::aot_inductor::ShapePolicy>(torch_N_m, "ShapePolicy")
74	.value("TIMES", torch::aot_inductor::ShapePolicy::TIMES)	74	.value("TIMES", torch::aot_inductor::ShapePolicy::TIMES)
75	.value("CUSTOM", torch::aot_inductor::ShapePolicy::CUSTOM)	75	.value("CUSTOM", torch::aot_inductor::ShapePolicy::CUSTOM)
	@@ -129,7 +129,7 @@ void THNPShapeHandling_init(PyObject *module)
129	py::arg("outputs")	129	py::arg("outputs")
130	)	130	)
131	.def_readwrite("dimension", &torch::aot_inductor::BSShapeOpStrategy::m_dimension);	131	.def_readwrite("dimension", &torch::aot_inductor::BSShapeOpStrategy::m_dimension);
132	-	132	+
133	py::class_<torch::aot_inductor::SeqShapeOpStrategy, PySeqShapeOpStrategy,	133	py::class_<torch::aot_inductor::SeqShapeOpStrategy, PySeqShapeOpStrategy,
134	torch::aot_inductor::ShapeOpStrategyBase,	134	torch::aot_inductor::ShapeOpStrategyBase,
135	std::shared_ptr<torch::aot_inductor::SeqShapeOpStrategy>>(torch_N_m, "_SeqShapeOpStrategy")	135	std::shared_ptr<torch::aot_inductor::SeqShapeOpStrategy>>(torch_N_m, "_SeqShapeOpStrategy")

Mtorch_npu/csrc/inductor/aoti_torch/npu_shape_handling.cpp+6-6

Mtorch_npu/csrc/inductor/mlir/hacl_rt.h+1-1

	@@ -183,7 +183,7 @@ typedef struct tagRtArgsEx {
183	uint16_t tilingDataOffset; // size to tiling data	183	uint16_t tilingDataOffset; // size to tiling data
184	uint16_t hostInputInfoNum; // 0	184	uint16_t hostInputInfoNum; // 0
185	uint8_t hasTiling; // has tiling	185	uint8_t hasTiling; // has tiling
186	- uint8_t isNoNeedH2DCopy; // not need rtKernelLaunchWithFlag copy tiling from host to device	186	+ uint8_t isNoNeedH2DCopy; // not need rtKernelLaunchWithFlag copy tiling from host to device
187	uint8_t reserved[4];	187	uint8_t reserved[4];
188	} rtArgsEx_t;	188	} rtArgsEx_t;
189		189

Mtorch_npu/csrc/npu/Graph.cpp+1-1

	@@ -789,7 +789,7 @@ void TORCH_NPU_API THNPGraph_init(PyObject* module) {
789	helper.processStringArrayOption(key, item.second.cast<std::vector<std::string>>());	789	helper.processStringArrayOption(key, item.second.cast<std::vector<std::string>>());
790	} else if (py::isinstance<py::int_>(item.second)) {	790	} else if (py::isinstance<py::int_>(item.second)) {
791	helper.processInitOption(key, item.second.cast<int>());	791	helper.processInitOption(key, item.second.cast<int>());
792	- }	792	+ }
793	}	793	}
794	}	794	}
795		795

Mtorch_npu/csrc/npu/Stress_detect.cpp+1-1

	@@ -130,7 +130,7 @@ int StressDetector::perform_stress_detect(int deviceid, int mode, int64_t comm)
130		130
131	// Set task parameters	131	// Set task parameters
132	task_in_progress.store(true);	132	task_in_progress.store(true);
133	-	133	+
134	// Allocate workspace memory	134	// Allocate workspace memory
135	workspaceAddr = nullptr;	135	workspaceAddr = nullptr;
136	uint64_t size = 10;	136	uint64_t size = 10;

Mtorch_npu/csrc/npu/Stress_detect.h+3-3

	@@ -20,17 +20,17 @@ private:
20	static void worker_thread();	20	static void worker_thread();
21		21
22	static int transfer_result(int detectResult);	22	static int transfer_result(int detectResult);
23	-	23	+
24	// Thread for handling the stress detection task	24	// Thread for handling the stress detection task
25	static std::thread stress_detect_thread;	25	static std::thread stress_detect_thread;
26		26
27	// Condition variable and mutex to control the thread	27	// Condition variable and mutex to control the thread
28	static std::condition_variable cv;	28	static std::condition_variable cv;
29	static std::mutex mtx;	29	static std::mutex mtx;
30	-	30	+
31	// Flag to indicate if a task is in progress	31	// Flag to indicate if a task is in progress
32	static std::atomic<bool> task_in_progress;	32	static std::atomic<bool> task_in_progress;
33	-	33	+
34	// Flag to signal the thread to stop	34	// Flag to signal the thread to stop
35	static std::atomic<bool> stop_thread;	35	static std::atomic<bool> stop_thread;
36		36

Mtorch_npu/csrc/profiler/profiler_mgr.cpp+2-2

	@@ -95,7 +95,7 @@ void ProfilerMgr::EnableMsProfiler(uint32_t *deviceIdList, uint32_t deviceNum, a
95	if (profConfig_ == nullptr) {	95	if (profConfig_ == nullptr) {
96	profConfig_ = at_npu::native::AclProfilingCreateConfig(deviceIdList, deviceNum, aicMetrics, nullptr, dataTypeConfig);	96	profConfig_ = at_npu::native::AclProfilingCreateConfig(deviceIdList, deviceNum, aicMetrics, nullptr, dataTypeConfig);
97	}	97	}
98	-	98	+
99	if (profConfig_ == nullptr) {	99	if (profConfig_ == nullptr) {
100	ASCEND_LOGE("Create Prof Config failed.");	100	ASCEND_LOGE("Create Prof Config failed.");
101	return;	101	return;
	@@ -275,7 +275,7 @@ void ProfilerMgr::Stop()
275	StopDataReceiver();	275	StopDataReceiver();
276	profile_memory_.store(false);	276	profile_memory_.store(false);
277	}	277	}
278	-	278	+
279	if (npu_trace_.load()) {	279	if (npu_trace_.load()) {
280	at_npu::native::AclProfilingStop(profConfig_);	280	at_npu::native::AclProfilingStop(profConfig_);
281	auto ret = at_npu::native::AclProfilingDestroyConfig(profConfig_);	281	auto ret = at_npu::native::AclProfilingDestroyConfig(profConfig_);

Mtorch_npu/distributed/nn/functional.py+1-1

	@@ -8,7 +8,7 @@ from torch.distributed import ReduceOp
8	def _allgather_base_backward_hccl(ctx, grad_output):	8	def _allgather_base_backward_hccl(ctx, grad_output):
9	"""	9	"""
10	Backward function for _AllGatherBase that supports HCCL backend.	10	Backward function for _AllGatherBase that supports HCCL backend.
11	-	11	+
12	Original PyTorch implementation only supports NCCL backend.	12	Original PyTorch implementation only supports NCCL backend.
13	This version adds HCCL support for NPU devices.	13	This version adds HCCL support for NPU devices.
14	"""	14	"""

Mtorch_npu/distributed/run.py+1-1

	@@ -15,7 +15,7 @@ def parse_args(args):
15	action=env,	15	action=env,
16	type=str,	16	type=str,
17	default="false",	17	default="false",
18	- help="Turn parallel tcpstore tiered optimization, if true, The agent adds a proxy role,"	18	+ help="Turn parallel tcpstore tiered optimization, if true, The agent adds a proxy role,"
19	"the worker on this node will connect to the server through the proxy.",	19	"the worker on this node will connect to the server through the proxy.",
20	)	20	)
21	return parser.parse_args(args)	21	return parser.parse_args(args)

Mtorch_npu/distributed/tensor/_attention.py+3-3

Mtorch_npu/distributed/tensor/_math_ops.py+5-5

	@@ -402,7 +402,7 @@ def custom_npu_conv2d_strategy(x, weight, bias, stride, padding, dilation, group
402	]	402	]
403	)	403	)
404	acceptable_shardings.append(replicate_strategy)	404	acceptable_shardings.append(replicate_strategy)
405	-	405	+
406	# x layout: (N, Ci, Hi, Wi)	406	# x layout: (N, Ci, Hi, Wi)
407	# weight layout: (Co, Ci/groups, Hk, Wk)	407	# weight layout: (Co, Ci/groups, Hk, Wk)
408	# bias layout: (Co)	408	# bias layout: (Co)
	@@ -500,7 +500,7 @@ def custom_grouped_matmul_add__strategy(y, x, weight, group_list, transpose_x=Tr
500		500
501	def is_tensor_evenly_shardable(shape, spec):	501	def is_tensor_evenly_shardable(shape, spec):
502	"""Check if the shape is evenly shardable according to the spec."""	502	"""Check if the shape is evenly shardable according to the spec."""
503	- # verify parameter validity	503	+ # verify parameter validity
504	if not isinstance(spec, DTensorSpec):	504	if not isinstance(spec, DTensorSpec):
505	raise TypeError(	505	raise TypeError(
506	f"Expected 'spec' to be DTensorSpec instance, got {type(spec).__name__} instead."	506	f"Expected 'spec' to be DTensorSpec instance, got {type(spec).__name__} instead."
	@@ -509,7 +509,7 @@ def is_tensor_evenly_shardable(shape, spec):
509	raise ValueError("'shape' must have at least 1 dimension (empty shape is invalid).")	509	raise ValueError("'shape' must have at least 1 dimension (empty shape is invalid).")
510	if len(spec.placements) == 0:	510	if len(spec.placements) == 0:
511	raise ValueError("'spec.placements' cannot be empty (must have at least one placement).")	511	raise ValueError("'spec.placements' cannot be empty (must have at least one placement).")
512	-	512	+
513	# number of shards in each tensor dimension	513	# number of shards in each tensor dimension
514	shards_map = [1] * len(shape)	514	shards_map = [1] * len(shape)
515	for i, placement in enumerate(spec.placements):	515	for i, placement in enumerate(spec.placements):
	@@ -529,14 +529,14 @@ def custom_cross_entropy_loss_sharding(op_schema: OpSchema):
529	single_mesh_dim_strategies = []	529	single_mesh_dim_strategies = []
530		530
531	args_schema = op_schema.args_schema	531	args_schema = op_schema.args_schema
532	-	532	+
533	input_strategy = args_schema[0] if len(args_schema) > 0 else None	533	input_strategy = args_schema[0] if len(args_schema) > 0 else None
534	target_strategy = args_schema[1] if len(args_schema) > 1 else None	534	target_strategy = args_schema[1] if len(args_schema) > 1 else None
535	weight_strategy = args_schema[2] if len(args_schema) > 2 else None	535	weight_strategy = args_schema[2] if len(args_schema) > 2 else None
536	reduction = args_schema[3] if len(args_schema) > 3 else 'mean'	536	reduction = args_schema[3] if len(args_schema) > 3 else 'mean'
537		537
538	mesh = input_strategy.mesh	538	mesh = input_strategy.mesh
539	-	539	+
540	all_replicate: PlacementList = [	540	all_replicate: PlacementList = [
541	Replicate(), # loss	541	Replicate(), # loss
542	Replicate(), # log_prob	542	Replicate(), # log_prob

Mtorch_npu/jit/fusion_pass/fast_gelu.py+1-1

	@@ -14,5 +14,5 @@ def fast_gelu_pass(jit_mod):
14	%out = npu::fast_gelu(%x)	14	%out = npu::fast_gelu(%x)
15	return (%out)	15	return (%out)
16	"""	16	"""
17	-	17	+
18	torch._C._jit_pass_custom_pattern_based_rewrite_graph(pattern, replacement, jit_mod.graph)	18	torch._C._jit_pass_custom_pattern_based_rewrite_graph(pattern, replacement, jit_mod.graph)

Mtorch_npu/jit/register_fusion_pattern.py+1-2

	@@ -7,6 +7,5 @@ __all__ = ["optimize"]
7	def optimize(jit_mod):	7	def optimize(jit_mod):
8	if isinstance(jit_mod, torch.jit.ScriptModule):	8	if isinstance(jit_mod, torch.jit.ScriptModule):
9	torch.jit.optimize_for_inference(jit_mod)	9	torch.jit.optimize_for_inference(jit_mod)
10	-	10	+
11	fast_gelu_pass(jit_mod)	11	fast_gelu_pass(jit_mod)
12	-

Mtorch_npu/npu/__init__.py+4-4

	@@ -155,7 +155,7 @@ from torch_npu._init.common.warning_utils import _should_print_warning
155		155
156	import torch_npu	156	import torch_npu
157	from torch_npu.utils._error_code import ErrCode, pta_error, prof_error	157	from torch_npu.utils._error_code import ErrCode, pta_error, prof_error
158	-from .utils import (obfuscation_initialize, obfuscation_calculate, obfuscation_finalize,	158	+from .utils import (obfuscation_initialize, obfuscation_calculate, obfuscation_finalize,
159	synchronize, set_device, current_device, _get_device_index,	159	synchronize, set_device, current_device, _get_device_index,
160	device, device_of, StreamContext, stream, set_stream, current_stream, default_stream, set_sync_debug_mode,	160	device, device_of, StreamContext, stream, set_stream, current_stream, default_stream, set_sync_debug_mode,
161	get_sync_debug_mode, init_dump, current_blas_handle, is_bf16_supported,	161	get_sync_debug_mode, init_dump, current_blas_handle, is_bf16_supported,
	@@ -510,7 +510,7 @@ _cached_device_capability_env = None
510	def get_device_capability(device=None):	510	def get_device_capability(device=None):
511	r"""Query the minor and major data of device.	511	r"""Query the minor and major data of device.
512		512
513	- This function can be configured via the TORCH_NPU_DEVICE_CAPABILITY environment variable.	513	+ This function can be configured via the TORCH_NPU_DEVICE_CAPABILITY environment variable.
514	The format should be "major.minor", e.g., "9.0" or "8.0".	514	The format should be "major.minor", e.g., "9.0" or "8.0".
515		515
516	.. note::	516	.. note::
	@@ -520,7 +520,7 @@ def get_device_capability(device=None):
520	device (torch.device or int, optional): The device parameter has no practical meaning.	520	device (torch.device or int, optional): The device parameter has no practical meaning.
521		521
522	Returns:	522	Returns:
523	- tuple(int, int): the device capability of the device. Returns the tuple(major, minor) configured via	523	+ tuple(int, int): the device capability of the device. Returns the tuple(major, minor) configured via
524	TORCH_NPU_DEVICE_CAPABILITY, or None if TORCH_NPU_DEVICE_CAPABILITY not configured.	524	TORCH_NPU_DEVICE_CAPABILITY, or None if TORCH_NPU_DEVICE_CAPABILITY not configured.
525		525
526	Example:	526	Example:
	@@ -619,7 +619,7 @@ def _get_deterministic_level():
619	torch_npu.npu.set_deterministic_level(level)	619	torch_npu.npu.set_deterministic_level(level)
620	return level	620	return level
621	if level >= 1 and not torch.are_deterministic_algorithms_enabled():	621	if level >= 1 and not torch.are_deterministic_algorithms_enabled():
622	- level = 0	622	+ level = 0
623	torch_npu.npu.set_deterministic_level(level)	623	torch_npu.npu.set_deterministic_level(level)
624	return level	624	return level
625	return level	625	return level

Mtorch_npu/npu/_fft_plan_cache.py+1-1

	@@ -11,7 +11,7 @@ class NPUFFTPlanCache:
11	if name == "max_size":	11	if name == "max_size":
12	return torch_npu._C._npu_get_fft_plan_cache_max_size()	12	return torch_npu._C._npu_get_fft_plan_cache_max_size()
13	raise AttributeError("Unknown attribute " + name)	13	raise AttributeError("Unknown attribute " + name)
14	-	14	+
15	def __setattr__(self, name, value):	15	def __setattr__(self, name, value):
16	if name == "size":	16	if name == "size":
17	raise RuntimeError(".size is a read-only property showing the number of plans currently in the cache.")	17	raise RuntimeError(".size is a read-only property showing the number of plans currently in the cache.")

Mtorch_npu/npu/_format.py+2-2

	@@ -33,11 +33,11 @@ class Format(IntEnum):
33		33
34	def _apply_npu_format_patch():	34	def _apply_npu_format_patch():
35	orig_get_format = torch_npu.get_npu_format	35	orig_get_format = torch_npu.get_npu_format
36	-	36	+
37	def patched_get_format(tensor):	37	def patched_get_format(tensor):
38	"""get the Format type of tensor"""	38	"""get the Format type of tensor"""
39	format_int = orig_get_format(tensor)	39	format_int = orig_get_format(tensor)
40	return Format(format_int)	40	return Format(format_int)
41	-	41	+
42	torch_npu.get_npu_format = patched_get_format	42	torch_npu.get_npu_format = patched_get_format
43	torch_npu.Format = Format	43	torch_npu.Format = Format

Mtorch_npu/npu/_graph_tree.py+2-2

	@@ -640,7 +640,7 @@ class NPUWarmupNode:
640	s = storage()	640	s = storage()
641	if s is not None:	641	if s is not None:
642	non_npugraph_inps_storage_ptrs.add(s._cdata)	642	non_npugraph_inps_storage_ptrs.add(s._cdata)
643	-	643	+
644	if not len(new_inputs) == 0:	644	if not len(new_inputs) == 0:
645	raise RuntimeError("check len(new_inputs) == 0 fail")	645	raise RuntimeError("check len(new_inputs) == 0 fail")
646		646
	@@ -842,7 +842,7 @@ class NPUGraphNode:
842	)	842	)
843		843
844	self.non_static_input_idx: LevelList[int] = [	844	self.non_static_input_idx: LevelList[int] = [
845	- i	845	+ i
846	for i in range(len(inputs))	846	for i in range(len(inputs))
847	if i not in self.static_input_idxs	847	if i not in self.static_input_idxs
848	]	848	]

Mtorch_npu/npu/aclnn/backends.py+1-1

	@@ -2,7 +2,7 @@ import warnings
2		2
3		3
4	def version():	4	def version():
5	- """Currently, the ACLNN version is not available and does not support it.	5	+ """Currently, the ACLNN version is not available and does not support it.
6	By default, it returns None.	6	By default, it returns None.
7	"""	7	"""
8	warnings.warn("torch.npu.aclnn.version isn't implemented!")	8	warnings.warn("torch.npu.aclnn.version isn't implemented!")

Mtorch_npu/npu/amp/sharded_grad_scaler.py+2-2

	@@ -280,7 +280,7 @@ class _ShardedGradScaler(GradScaler):
280	# Synchronize the detected inf across the ranks	280	# Synchronize the detected inf across the ranks
281	optimizer_state = self._per_optimizer_states[id(optimizer)]	281	optimizer_state = self._per_optimizer_states[id(optimizer)]
282	works = []	282	works = []
283	-	283	+
284	for found_inf in optimizer_state["found_inf_per_device"].values():	284	for found_inf in optimizer_state["found_inf_per_device"].values():
285	if found_inf.device.type == "cpu":	285	if found_inf.device.type == "cpu":
286	found_inf_npu = found_inf.to(self._scale.device)	286	found_inf_npu = found_inf.to(self._scale.device)
	@@ -288,7 +288,7 @@ class _ShardedGradScaler(GradScaler):
288	works.append((work, found_inf, found_inf_npu))	288	works.append((work, found_inf, found_inf_npu))
289	else:	289	else:
290	works.append((dist.all_reduce(found_inf, async_op=True, group=self.process_group), None, None))	290	works.append((dist.all_reduce(found_inf, async_op=True, group=self.process_group), None, None))
291	-	291	+
292	for item in works:	292	for item in works:
293	if item[1] is not None:	293	if item[1] is not None:
294	work, found_inf_cpu, found_inf_npu = item	294	work, found_inf_cpu, found_inf_npu = item

Mtorch_npu/npu/autocast_utils.py+1-1

	@@ -8,7 +8,7 @@ __all__ = ["get_amp_supported_dtype", "is_autocast_enabled", "set_autocast_enabl
8	def get_amp_supported_dtype():	8	def get_amp_supported_dtype():
9	if torch.npu.is_bf16_supported():	9	if torch.npu.is_bf16_supported():
10	return [torch.float16, torch.bfloat16, torch.float32]	10	return [torch.float16, torch.bfloat16, torch.float32]
11	- return [torch.float16, torch.float32]	11	+ return [torch.float16, torch.float32]
12		12
13		13
14	def is_autocast_enabled():	14	def is_autocast_enabled():

Mtorch_npu/npu/graphs.py+12-12

	@@ -229,18 +229,18 @@ def _print_npugraph_tensor_impl(input, tensor_name=None):
229	if device.type == "cpu":	229	if device.type == "cpu":
230	_print_callback_pending(tensor_name, input)	230	_print_callback_pending(tensor_name, input)
231	return	231	return
232	-	232	+
233	if device.type != "npu":	233	if device.type != "npu":
234	return	234	return
235		235
236	device_index = device.index	236	device_index = device.index
237	save_stream = _get_save_tensor_stream(device_index)	237	save_stream = _get_save_tensor_stream(device_index)
238	-	238	+
239	# Record event on the original compute stream before switching	239	# Record event on the original compute stream before switching
240	event1 = torch.npu.Event()	240	event1 = torch.npu.Event()
241	event2 = torch.npu.Event()	241	event2 = torch.npu.Event()
242	event1.record()	242	event1.record()
243	-	243	+
244	with torch.npu.stream(save_stream):	244	with torch.npu.stream(save_stream):
245	# Wait for the original stream to complete before D2H	245	# Wait for the original stream to complete before D2H
246	event1.wait()	246	event1.wait()
	@@ -255,7 +255,7 @@ def _print_npugraph_tensor_impl(input, tensor_name=None):
255	)	255	)
256	# Mark save_stream completion	256	# Mark save_stream completion
257	event2.record()	257	event2.record()
258	-	258	+
259	# Wait for save_stream to complete (back to original stream now)	259	# Wait for save_stream to complete (back to original stream now)
260	event2.wait()	260	event2.wait()
261		261
	@@ -268,19 +268,19 @@ def _save_npugraph_tensor_impl(input, save_path=None, overwrite=False):
268	if device.type == "cpu":	268	if device.type == "cpu":
269	torch.save(input, _build_save_npugraph_tensor_path(save_path, overwrite=overwrite))	269	torch.save(input, _build_save_npugraph_tensor_path(save_path, overwrite=overwrite))
270	return	270	return
271	-	271	+
272	if device.type != "npu":	272	if device.type != "npu":
273	return	273	return
274		274
275	device_index = device.index	275	device_index = device.index
276	save_stream = _get_save_tensor_stream(device_index)	276	save_stream = _get_save_tensor_stream(device_index)
277	final_path = _build_save_npugraph_tensor_path(save_path, device_index, overwrite)	277	final_path = _build_save_npugraph_tensor_path(save_path, device_index, overwrite)
278	-	278	+
279	# Record event on the original compute stream before switching	279	# Record event on the original compute stream before switching
280	event1 = torch.npu.Event()	280	event1 = torch.npu.Event()
281	event2 = torch.npu.Event()	281	event2 = torch.npu.Event()
282	event1.record()	282	event1.record()
283	-	283	+
284	with torch.npu.stream(save_stream):	284	with torch.npu.stream(save_stream):
285	# Wait for the original stream to complete before D2H	285	# Wait for the original stream to complete before D2H
286	event1.wait()	286	event1.wait()
	@@ -295,7 +295,7 @@ def _save_npugraph_tensor_impl(input, save_path=None, overwrite=False):
295	)	295	)
296	# Mark save_stream completion	296	# Mark save_stream completion
297	event2.record()	297	event2.record()
298	-	298	+
299	# Wait for save_stream to complete (back to original stream now)	299	# Wait for save_stream to complete (back to original stream now)
300	event2.wait()	300	event2.wait()
301		301
	@@ -305,19 +305,19 @@ def _save_npugraph_tensor_tensor_list_impl(input, save_path=None, overwrite=Fals
305	if device.type == "cpu":	305	if device.type == "cpu":
306	torch.save(list(input), _build_save_npugraph_tensor_path(save_path, overwrite=overwrite))	306	torch.save(list(input), _build_save_npugraph_tensor_path(save_path, overwrite=overwrite))
307	return	307	return
308	-	308	+
309	if device.type != "npu":	309	if device.type != "npu":
310	return	310	return
311		311
312	device_index = device.index	312	device_index = device.index
313	save_stream = _get_save_tensor_stream(device_index)	313	save_stream = _get_save_tensor_stream(device_index)
314	final_path = _build_save_npugraph_tensor_path(save_path, device_index, overwrite)	314	final_path = _build_save_npugraph_tensor_path(save_path, device_index, overwrite)
315	-	315	+
316	# Record event on the original compute stream before switching	316	# Record event on the original compute stream before switching
317	event1 = torch.npu.Event()	317	event1 = torch.npu.Event()
318	event2 = torch.npu.Event()	318	event2 = torch.npu.Event()
319	event1.record()	319	event1.record()
320	-	320	+
321	with torch.npu.stream(save_stream):	321	with torch.npu.stream(save_stream):
322	# Wait for the original stream to complete before D2H	322	# Wait for the original stream to complete before D2H
323	event1.wait()	323	event1.wait()
	@@ -332,7 +332,7 @@ def _save_npugraph_tensor_tensor_list_impl(input, save_path=None, overwrite=Fals
332	)	332	)
333	# Mark save_stream completion	333	# Mark save_stream completion
334	event2.record()	334	event2.record()
335	-	335	+
336	# Wait for save_stream to complete (back to original stream now)	336	# Wait for save_stream to complete (back to original stream now)
337	event2.wait()	337	event2.wait()
338		338

Mtorch_npu/npu/npugraph_ex/__init__.py+1-1

	@@ -30,5 +30,5 @@ def register_replacement(search_fn: SearchFn, replace_fn: ReplaceFn, example_inp
30	return npugraph_ex.patterns.pattern_pass_manager.register_replacement(search_fn, replace_fn, example_inputs,	30	return npugraph_ex.patterns.pattern_pass_manager.register_replacement(search_fn, replace_fn, example_inputs,
31	trace_fn=trace_fn, extra_check=extra_check,	31	trace_fn=trace_fn, extra_check=extra_check,
32	search_fn_pattern=search_fn_pattern,	32	search_fn_pattern=search_fn_pattern,
33	- scalar_workaround=scalar_workaround,	33	+ scalar_workaround=scalar_workaround,
34	skip_duplicates=skip_duplicates)	34	skip_duplicates=skip_duplicates)

Mtorch_npu/npu/utils.py+1-1

	@@ -22,7 +22,7 @@ __all__ = ["obfuscation_initialize", "obfuscation_finalize", "obfuscation_calcul
22		22
23		23
24	def obfuscation_initialize(hidden_size, tp_rank, cmd, *, data_type=None, model_obf_seed_id=0, data_obf_seed_id=0, thread_num=4, obf_coefficient=1.0):	24	def obfuscation_initialize(hidden_size, tp_rank, cmd, *, data_type=None, model_obf_seed_id=0, data_obf_seed_id=0, thread_num=4, obf_coefficient=1.0):
25	- return torch_npu.obfuscation_initialize(hidden_size, tp_rank, cmd, data_type=data_type, model_obf_seed_id=model_obf_seed_id,	25	+ return torch_npu.obfuscation_initialize(hidden_size, tp_rank, cmd, data_type=data_type, model_obf_seed_id=model_obf_seed_id,
26	data_obf_seed_id=data_obf_seed_id, thread_num=thread_num, obf_coefficient=obf_coefficient)	26	data_obf_seed_id=data_obf_seed_id, thread_num=thread_num, obf_coefficient=obf_coefficient)
27		27
28		28

Mtorch_npu/optim/npu_fused_adadelta.py+7-7

	@@ -71,13 +71,13 @@ class NpuFusedAdadelta(NpuFusedOptimizerBase):
71	if grad.is_sparse:	71	if grad.is_sparse:
72	raise RuntimeError('NpuFusedAdadelta does not support sparse gradients' +	72	raise RuntimeError('NpuFusedAdadelta does not support sparse gradients' +
73	pta_error(ErrCode.NOT_SUPPORT))	73	pta_error(ErrCode.NOT_SUPPORT))
74	-	74	+
75	self._init_param_state(p)	75	self._init_param_state(p)
76	state = self.state[p]	76	state = self.state[p]
77	step_list.append(state['step'])	77	step_list.append(state['step'])
78	square_avg_list.append(state['square_avg'])	78	square_avg_list.append(state['square_avg'])
79	acc_delta_list.append(state['acc_delta'])	79	acc_delta_list.append(state['acc_delta'])
80	-	80	+
81	combined_step = 0	81	combined_step = 0
82	combined_square_avg = None	82	combined_square_avg = None
83	combined_acc_delta = None	83	combined_acc_delta = None
	@@ -86,7 +86,7 @@ class NpuFusedAdadelta(NpuFusedOptimizerBase):
86	combined_step = step_list[0]	86	combined_step = step_list[0]
87	combined_square_avg = npu_combine_tensors(square_avg_list)	87	combined_square_avg = npu_combine_tensors(square_avg_list)
88	combined_acc_delta = npu_combine_tensors(acc_delta_list)	88	combined_acc_delta = npu_combine_tensors(acc_delta_list)
89	-	89	+
90	combined_state = defaultdict(dict)	90	combined_state = defaultdict(dict)
91	combined_state['step'] = combined_step	91	combined_state['step'] = combined_step
92	combined_state['square_avg'] = combined_square_avg	92	combined_state['square_avg'] = combined_square_avg
	@@ -97,12 +97,12 @@ class NpuFusedAdadelta(NpuFusedOptimizerBase):
97	def _maybe_init_combined_states(self):	97	def _maybe_init_combined_states(self):
98	if self.is_states_combined:	98	if self.is_states_combined:
99	return	99	return
100	-	100	+
101	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]	101	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]
102		102
103	for i, _ in enumerate(self.param_groups):	103	for i, _ in enumerate(self.param_groups):
104	self._combine_group_param_states(i)	104	self._combine_group_param_states(i)
105	-	105	+
106	if not all(value is None for value in self.combined_param_states_indexed_by_group):	106	if not all(value is None for value in self.combined_param_states_indexed_by_group):
107	self.is_states_combined = True	107	self.is_states_combined = True
108		108
	@@ -124,8 +124,8 @@ class NpuFusedAdadelta(NpuFusedOptimizerBase):
124	combined_group_grads = self.combined_grads_indexed_by_group[group_index]	124	combined_group_grads = self.combined_grads_indexed_by_group[group_index]
125	combined_group_param_states = self.combined_param_states_indexed_by_group[group_index]	125	combined_group_param_states = self.combined_param_states_indexed_by_group[group_index]
126		126
127	- for combined_param, combined_grad, combined_param_state in zip(combined_group_params,	127	+ for combined_param, combined_grad, combined_param_state in zip(combined_group_params,
128	- combined_group_grads,	128	+ combined_group_grads,
129	combined_group_param_states):	129	combined_group_param_states):
130	if combined_param is None or combined_grad is None:	130	if combined_param is None or combined_grad is None:
131	continue	131	continue

Mtorch_npu/optim/npu_fused_adam.py+6-6

	@@ -103,7 +103,7 @@ class NpuFusedAdam(NpuFusedOptimizerBase):
103	exp_avg_sq_list.append(state['exp_avg_sq'])	103	exp_avg_sq_list.append(state['exp_avg_sq'])
104	if amsgrad:	104	if amsgrad:
105	max_exp_avg_sq_list.append(state['max_exp_avg_sq'])	105	max_exp_avg_sq_list.append(state['max_exp_avg_sq'])
106	-	106	+
107	combined_step = 0	107	combined_step = 0
108	combined_exp_avg = None	108	combined_exp_avg = None
109	combined_exp_avg_sq = None	109	combined_exp_avg_sq = None
	@@ -114,7 +114,7 @@ class NpuFusedAdam(NpuFusedOptimizerBase):
114	combined_exp_avg = npu_combine_tensors(exp_avg_list)	114	combined_exp_avg = npu_combine_tensors(exp_avg_list)
115	combined_exp_avg_sq = npu_combine_tensors(exp_avg_sq_list)	115	combined_exp_avg_sq = npu_combine_tensors(exp_avg_sq_list)
116	combined_max_exp_avg_sq = npu_combine_tensors(max_exp_avg_sq_list)	116	combined_max_exp_avg_sq = npu_combine_tensors(max_exp_avg_sq_list)
117	-	117	+
118	combined_state = defaultdict(dict)	118	combined_state = defaultdict(dict)
119	combined_state['step'] = combined_step	119	combined_state['step'] = combined_step
120	combined_state['exp_avg'] = combined_exp_avg	120	combined_state['exp_avg'] = combined_exp_avg
	@@ -126,12 +126,12 @@ class NpuFusedAdam(NpuFusedOptimizerBase):
126	def _maybe_init_combined_states(self):	126	def _maybe_init_combined_states(self):
127	if self.is_states_combined:	127	if self.is_states_combined:
128	return	128	return
129	-	129	+
130	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]	130	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]
131		131
132	for i, _ in enumerate(self.param_groups):	132	for i, _ in enumerate(self.param_groups):
133	self._combine_group_param_states(i)	133	self._combine_group_param_states(i)
134	-	134	+
135	if not all(value is None for value in self.combined_param_states_indexed_by_group):	135	if not all(value is None for value in self.combined_param_states_indexed_by_group):
136	self.is_states_combined = True	136	self.is_states_combined = True
137		137
	@@ -154,8 +154,8 @@ class NpuFusedAdam(NpuFusedOptimizerBase):
154	combined_group_grads = self.combined_grads_indexed_by_group[group_index]	154	combined_group_grads = self.combined_grads_indexed_by_group[group_index]
155	combined_group_param_states = self.combined_param_states_indexed_by_group[group_index]	155	combined_group_param_states = self.combined_param_states_indexed_by_group[group_index]
156		156
157	- for combined_param, combined_grad, combined_param_state in zip(combined_group_params,	157	+ for combined_param, combined_grad, combined_param_state in zip(combined_group_params,
158	- combined_group_grads,	158	+ combined_group_grads,
159	combined_group_param_states):	159	combined_group_param_states):
160	if combined_param is None or combined_grad is None:	160	if combined_param is None or combined_grad is None:
161	continue	161	continue

Mtorch_npu/optim/npu_fused_adamp.py+3-3

Mtorch_npu/optim/npu_fused_bert_adam.py+2-2

	@@ -149,12 +149,12 @@ class NpuFusedBertAdam(NpuFusedOptimizerBase):
149	def _maybe_init_combined_states(self):	149	def _maybe_init_combined_states(self):
150	if self.is_states_combined:	150	if self.is_states_combined:
151	return	151	return
152	-	152	+
153	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]	153	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]
154		154
155	for i, _ in enumerate(self.param_groups):	155	for i, _ in enumerate(self.param_groups):
156	self._combine_group_param_states(i)	156	self._combine_group_param_states(i)
157	-	157	+
158	if not all(value is None for value in self.combined_param_states_indexed_by_group):	158	if not all(value is None for value in self.combined_param_states_indexed_by_group):
159	self.is_states_combined = True	159	self.is_states_combined = True
160		160

Mtorch_npu/optim/npu_fused_lamb.py+8-8

	@@ -137,13 +137,13 @@ class NpuFusedLamb(NpuFusedOptimizerBase):
137	if grad.is_sparse:	137	if grad.is_sparse:
138	raise RuntimeError('NpuFusedLamb does not support sparse gradients, '	138	raise RuntimeError('NpuFusedLamb does not support sparse gradients, '
139	'please consider SparseAdam instead.' + pta_error(ErrCode.NOT_SUPPORT))	139	'please consider SparseAdam instead.' + pta_error(ErrCode.NOT_SUPPORT))
140	-	140	+
141	self._init_param_state(p)	141	self._init_param_state(p)
142	state = self.state[p]	142	state = self.state[p]
143	step_list.append(state['step'])	143	step_list.append(state['step'])
144	exp_avg_list.append(state['exp_avg'])	144	exp_avg_list.append(state['exp_avg'])
145	exp_avg_sq_list.append(state['exp_avg_sq'])	145	exp_avg_sq_list.append(state['exp_avg_sq'])
146	-	146	+
147	combined_step = 0	147	combined_step = 0
148	combined_exp_avg = None	148	combined_exp_avg = None
149	combined_exp_avg_sq = None	149	combined_exp_avg_sq = None
	@@ -152,7 +152,7 @@ class NpuFusedLamb(NpuFusedOptimizerBase):
152	combined_step = step_list[0]	152	combined_step = step_list[0]
153	combined_exp_avg = npu_combine_tensors(exp_avg_list)	153	combined_exp_avg = npu_combine_tensors(exp_avg_list)
154	combined_exp_avg_sq = npu_combine_tensors(exp_avg_sq_list)	154	combined_exp_avg_sq = npu_combine_tensors(exp_avg_sq_list)
155	-	155	+
156	combined_state = defaultdict(dict)	156	combined_state = defaultdict(dict)
157	combined_state['step'] = combined_step	157	combined_state['step'] = combined_step
158	combined_state['exp_avg'] = combined_exp_avg	158	combined_state['exp_avg'] = combined_exp_avg
	@@ -163,12 +163,12 @@ class NpuFusedLamb(NpuFusedOptimizerBase):
163	def _maybe_init_combined_states(self):	163	def _maybe_init_combined_states(self):
164	if self.is_states_combined:	164	if self.is_states_combined:
165	return	165	return
166	-	166	+
167	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]	167	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]
168		168
169	for i, _ in enumerate(self.param_groups):	169	for i, _ in enumerate(self.param_groups):
170	self._combine_group_param_states(i)	170	self._combine_group_param_states(i)
171	-	171	+
172	if not all(value is None for value in self.combined_param_states_indexed_by_group):	172	if not all(value is None for value in self.combined_param_states_indexed_by_group):
173	self.is_states_combined = True	173	self.is_states_combined = True
174		174
	@@ -237,8 +237,8 @@ class NpuFusedLamb(NpuFusedOptimizerBase):
237	combined_param_pow.copy_(combined_param.pow(2))	237	combined_param_pow.copy_(combined_param.pow(2))
238	combined_adam_step_pow.copy_(adam_step.pow(2))	238	combined_adam_step_pow.copy_(adam_step.pow(2))
239		239
240	- for param_pow, adam_step_pow, trust_ratio in zip(param_pow_list,	240	+ for param_pow, adam_step_pow, trust_ratio in zip(param_pow_list,
241	- adam_step_pow_list,	241	+ adam_step_pow_list,
242	trust_ratio_list):	242	trust_ratio_list):
243	weight_norm = param_pow.sum().sqrt().clamp(0, 10)	243	weight_norm = param_pow.sum().sqrt().clamp(0, 10)
244	adam_norm = adam_step_pow.sum().sqrt()	244	adam_norm = adam_step_pow.sum().sqrt()
	@@ -253,7 +253,7 @@ class NpuFusedLamb(NpuFusedOptimizerBase):
253	def step(self, closure=None):	253	def step(self, closure=None):
254	if not self.is_params_grads_combined:	254	if not self.is_params_grads_combined:
255	self._maybe_init_combined_params_and_grads()	255	self._maybe_init_combined_params_and_grads()
256	-	256	+
257	if not self.is_states_combined:	257	if not self.is_states_combined:
258	self._maybe_init_combined_states()	258	self._maybe_init_combined_states()
259		259

Mtorch_npu/optim/npu_fused_optim_base.py+5-5

	@@ -29,7 +29,7 @@ class NpuFusedOptimizerBase(Optimizer):
29	return	29	return
30		30
31	self.combined_params_indexed_by_group = len(self.param_groups) * [[]]	31	self.combined_params_indexed_by_group = len(self.param_groups) * [[]]
32	- self.combined_grads_indexed_by_group = len(self.param_groups) * [[]]	32	+ self.combined_grads_indexed_by_group = len(self.param_groups) * [[]]
33		33
34	params_list_each_group = []	34	params_list_each_group = []
35	params_size_each_group = []	35	params_size_each_group = []
	@@ -103,9 +103,9 @@ class NpuFusedOptimizerBase(Optimizer):
103		103
104	self.combined_params_indexed_by_group[group_index] = group_combined_params	104	self.combined_params_indexed_by_group[group_index] = group_combined_params
105	self.combined_grads_indexed_by_group[group_index] = group_combined_grads	105	self.combined_grads_indexed_by_group[group_index] = group_combined_grads
106	-	106	+
107	if not all(value is None for value in self.params_all_group_combined):	107	if not all(value is None for value in self.params_all_group_combined):
108	- self.is_params_grads_combined = True	108	+ self.is_params_grads_combined = True
109		109
110	@torch.no_grad()	110	@torch.no_grad()
111	def step(self, closure=None):	111	def step(self, closure=None):
	@@ -115,7 +115,7 @@ class NpuFusedOptimizerBase(Optimizer):
115		115
116	if not self.is_params_grads_combined:	116	if not self.is_params_grads_combined:
117	self._maybe_init_combined_params_and_grads()	117	self._maybe_init_combined_params_and_grads()
118	-	118	+
119	if not self.is_states_combined:	119	if not self.is_states_combined:
120	self._maybe_init_combined_states()	120	self._maybe_init_combined_states()
121		121
	@@ -135,7 +135,7 @@ class NpuFusedOptimizerBase(Optimizer):
135	if not self.is_params_grads_combined:	135	if not self.is_params_grads_combined:
136	super().zero_grad(set_to_none)	136	super().zero_grad(set_to_none)
137	return	137	return
138	-	138	+
139	for grads_combined_one_dtype in self.grads_all_group_combined:	139	for grads_combined_one_dtype in self.grads_all_group_combined:
140	if grads_combined_one_dtype is None:	140	if grads_combined_one_dtype is None:
141	continue	141	continue

Mtorch_npu/optim/npu_fused_rmsprop.py+2-2

	@@ -129,12 +129,12 @@ class NpuFusedRMSprop(NpuFusedOptimizerBase):
129	def _maybe_init_combined_states(self):	129	def _maybe_init_combined_states(self):
130	if self.is_states_combined:	130	if self.is_states_combined:
131	return	131	return
132	-	132	+
133	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]	133	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]
134		134
135	for i, _ in enumerate(self.param_groups):	135	for i, _ in enumerate(self.param_groups):
136	self._combine_group_param_states(i)	136	self._combine_group_param_states(i)
137	-	137	+
138	if not all(value is None for value in self.combined_param_states_indexed_by_group):	138	if not all(value is None for value in self.combined_param_states_indexed_by_group):
139	self.is_states_combined = True	139	self.is_states_combined = True
140		140

Mtorch_npu/optim/npu_fused_rmsprop_tf.py+2-2

	@@ -129,12 +129,12 @@ class NpuFusedRMSpropTF(NpuFusedOptimizerBase):
129	def _maybe_init_combined_states(self):	129	def _maybe_init_combined_states(self):
130	if self.is_states_combined:	130	if self.is_states_combined:
131	return	131	return
132	-	132	+
133	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]	133	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]
134		134
135	for i, _ in enumerate(self.param_groups):	135	for i, _ in enumerate(self.param_groups):
136	self._combine_group_param_states(i)	136	self._combine_group_param_states(i)
137	-	137	+
138	if not all(value is None for value in self.combined_param_states_indexed_by_group):	138	if not all(value is None for value in self.combined_param_states_indexed_by_group):
139	self.is_states_combined = True	139	self.is_states_combined = True
140		140

Mtorch_npu/optim/npu_fused_sgd.py+3-3

	@@ -117,12 +117,12 @@ class NpuFusedSGD(NpuFusedOptimizerBase):
117	def _maybe_init_combined_states(self):	117	def _maybe_init_combined_states(self):
118	if self.is_states_combined:	118	if self.is_states_combined:
119	return	119	return
120	-	120	+
121	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]	121	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]
122		122
123	for i, _ in enumerate(self.param_groups):	123	for i, _ in enumerate(self.param_groups):
124	self._combine_group_param_states(i)	124	self._combine_group_param_states(i)
125	-	125	+
126	if not all(value is None for value in self.combined_param_states_indexed_by_group):	126	if not all(value is None for value in self.combined_param_states_indexed_by_group):
127	self.is_states_combined = True	127	self.is_states_combined = True
128		128
	@@ -163,7 +163,7 @@ class NpuFusedSGD(NpuFusedOptimizerBase):
163		163
164	combined_param_one_dtype.add_(combined_grad_one_dtype,	164	combined_param_one_dtype.add_(combined_grad_one_dtype,
165	alpha=-group['lr'])	165	alpha=-group['lr'])
166	-	166	+
167	def step(self, closure=None):	167	def step(self, closure=None):
168	ret = super().step(closure)	168	ret = super().step(closure)
169	self._momentum_buffer_already_in_state = True	169	self._momentum_buffer_already_in_state = True

Mtorch_npu/profiler/_profiler_action_controller.py+1-1

	@@ -13,7 +13,7 @@ class ProfActionController:
13	self,	13	self,
14	prof,	14	prof,
15	prof_inst: _ProfInterface,	15	prof_inst: _ProfInterface,
16	- on_trace_ready: Optional[Callable[..., Any]] = None,	16	+ on_trace_ready: Optional[Callable[..., Any]] = None,
17	) -> None:	17	) -> None:
18	self.prof = prof	18	self.prof = prof
19	self.prof_inst = prof_inst	19	self.prof_inst = prof_inst

Mtorch_npu/profiler/_profiler_path_creator.py+1-1

	@@ -41,7 +41,7 @@ class ProfPathCreator:
41	PathManager.check_input_directory_path(dir_path)	41	PathManager.check_input_directory_path(dir_path)
42	self._dir_path = dir_path	42	self._dir_path = dir_path
43	elif dir_name is None:	43	elif dir_name is None:
44	- self._dir_path = dir_name	44	+ self._dir_path = dir_name
45	else:	45	else:
46	print_warn_msg("Invalid parameter dir_name, reset it to default.")	46	print_warn_msg("Invalid parameter dir_name, reset it to default.")
47	self._dir_path = None	47	self._dir_path = None

Mtorch_npu/profiler/analysis/prof_bean/_ge_memory_record_bean.py+1-1

	@@ -5,7 +5,7 @@ from ..prof_common_func._constant import convert_us2ns
5		5
6	__all__ = []	6	__all__ = []
7		7
8	-	8	+
9	class GeMemoryRecordBean(CommonBean):	9	class GeMemoryRecordBean(CommonBean):
10		10
11	def __init__(self, data: dict):	11	def __init__(self, data: dict):

Mtorch_npu/profiler/analysis/prof_bean/_ge_op_memory_bean.py+1-1

	@@ -15,7 +15,7 @@ class GeOpMemoryBean(CommonBean):
15		15
16	@property	16	@property
17	def row(self) -> list:	17	def row(self) -> list:
18	- return [self.name, self.size, self.allocation_time, self.release_time, None,	18	+ return [self.name, self.size, self.allocation_time, self.release_time, None,
19	self.dur, None, self.allocation_total_allocated, self.allocation_total_reserved, None,	19	self.dur, None, self.allocation_total_allocated, self.allocation_total_reserved, None,
20	self.release_total_allocated, self.release_total_reserved, None, None, self.device]	20	self.release_total_allocated, self.release_total_reserved, None, None, self.device]
21		21

Mtorch_npu/profiler/analysis/prof_bean/_memory_use_bean.py+1-1

	@@ -117,7 +117,7 @@ class MemoryUseBean(CommonBean):
117	@property	117	@property
118	def data_type(self) -> int:	118	def data_type(self) -> int:
119	return self._data_type	119	return self._data_type
120	-	120	+
121	@property	121	@property
122	def allocator_type(self) -> int:	122	def allocator_type(self) -> int:
123	return self._allocator_type	123	return self._allocator_type

Mtorch_npu/profiler/analysis/prof_bean/_param_tensor_bean.py+4-4

	@@ -27,19 +27,19 @@ class ParamTensorBean:
27	self._key = self._constant_data[0]	27	self._key = self._constant_data[0]
28	self._module_params = None	28	self._module_params = None
29	self._optimizer_params = None	29	self._optimizer_params = None
30	-	30	+
31	module_params = self._origin_data.get(self.TLV_TYPE_DICT.get(Constant.MODULE_PARAM))	31	module_params = self._origin_data.get(self.TLV_TYPE_DICT.get(Constant.MODULE_PARAM))
32	if module_params is not None:	32	if module_params is not None:
33	self._module_params = [param for param in module_params.split('}')]	33	self._module_params = [param for param in module_params.split('}')]
34	-	34	+
35	optimizer_params = self._origin_data.get(self.TLV_TYPE_DICT.get(Constant.OPTIMIZER_PARAM))	35	optimizer_params = self._origin_data.get(self.TLV_TYPE_DICT.get(Constant.OPTIMIZER_PARAM))
36	if optimizer_params is not None:	36	if optimizer_params is not None:
37	self._optimizer_params = [param for param in optimizer_params.split('}')]	37	self._optimizer_params = [param for param in optimizer_params.split('}')]
38	-	38	+
39	@property	39	@property
40	def key(self) -> int:	40	def key(self) -> int:
41	return self._key	41	return self._key
42	-	42	+
43	@property	43	@property
44	def params(self) -> KeyAndParam:	44	def params(self) -> KeyAndParam:
45	return KeyAndParam(self._key, self._module_params, self._optimizer_params)	45	return KeyAndParam(self._key, self._module_params, self._optimizer_params)

Mtorch_npu/profiler/analysis/prof_bean/_torch_op_bean.py+2-2

	@@ -95,7 +95,7 @@ class TorchOpBean:
95	if self._call_stack is None:	95	if self._call_stack is None:
96	self._call_stack = self._origin_data.get(self.TLV_TYPE_DICT.get(Constant.CALL_STACK), "").replace(";", ";\r\n")	96	self._call_stack = self._origin_data.get(self.TLV_TYPE_DICT.get(Constant.CALL_STACK), "").replace(";", ";\r\n")
97	return self._call_stack	97	return self._call_stack
98	-	98	+
99	@property	99	@property
100	def inputs(self):	100	def inputs(self):
101	if self._inputs is None:	101	if self._inputs is None:
	@@ -121,7 +121,7 @@ class TorchOpBean:
121	@property	121	@property
122	def is_torch_op(self):	122	def is_torch_op(self):
123	return True	123	return True
124	-	124	+
125	def _init_timestamps(self):	125	def _init_timestamps(self):
126	profiler_config = ProfilerConfig()	126	profiler_config = ProfilerConfig()
127	start_syscnt = self._constant_data[TorchOpEnum.START_NS.value]	127	start_syscnt = self._constant_data[TorchOpEnum.START_NS.value]

Mtorch_npu/profiler/analysis/prof_common_func/_cann_package_manager.py+2-2

	@@ -11,10 +11,10 @@ def check_msprof_help_output(search_text: str) -> bool:
11	msprof_path = shutil.which("msprof")	11	msprof_path = shutil.which("msprof")
12	if not msprof_path:	12	if not msprof_path:
13	return False	13	return False
14	-	14	+
15	if not ProfilerPathManager.check_path_permission(msprof_path):	15	if not ProfilerPathManager.check_path_permission(msprof_path):
16	return False	16	return False
17	-	17	+
18	completed_process = subprocess.run([msprof_path, "--help"], capture_output=True, shell=False, text=True)	18	completed_process = subprocess.run([msprof_path, "--help"], capture_output=True, shell=False, text=True)
19	if completed_process.returncode != COMMAND_SUCCESS:	19	if completed_process.returncode != COMMAND_SUCCESS:
20	return False	20	return False

Mtorch_npu/profiler/analysis/prof_common_func/_db_manager.py+4-4

	@@ -33,7 +33,7 @@ class DbManager:
33	def create_connect_db(cls, db_path: str) -> tuple:	33	def create_connect_db(cls, db_path: str) -> tuple:
34	"""	34	"""
35	create and connect database	35	create and connect database
36	- """	36	+ """
37	if os.path.exists(db_path):	37	if os.path.exists(db_path):
38	FileManager.check_db_file_vaild(db_path)	38	FileManager.check_db_file_vaild(db_path)
39	try:	39	try:
	@@ -41,7 +41,7 @@ class DbManager:
41	conn = sqlite3.connect(db_path, timeout=2147483, check_same_thread=False)	41	conn = sqlite3.connect(db_path, timeout=2147483, check_same_thread=False)
42	except sqlite3.Error as err:	42	except sqlite3.Error as err:
43	return EmptyClass("emoty conn"), EmptyClass("empty curs")	43	return EmptyClass("emoty conn"), EmptyClass("empty curs")
44	-	44	+
45	try:	45	try:
46	curs = conn.cursor()	46	curs = conn.cursor()
47	os.chmod(db_path, Constant.FILE_AUTHORITY)	47	os.chmod(db_path, Constant.FILE_AUTHORITY)
	@@ -60,7 +60,7 @@ class DbManager:
60	cur.close()	60	cur.close()
61	except sqlite3.Error as err:	61	except sqlite3.Error as err:
62	raise RuntimeError(f"Falied to close db connection cursor") from err	62	raise RuntimeError(f"Falied to close db connection cursor") from err
63	-	63	+
64	try:	64	try:
65	conn.close()	65	conn.close()
66	except sqlite3.Error as err:	66	except sqlite3.Error as err:
	@@ -162,7 +162,7 @@ class DbManager:
162	except sqlite3.Error as err:	162	except sqlite3.Error as err:
163	print_error_msg("SQLite Error: %s" % " ".join(err.args))	163	print_error_msg("SQLite Error: %s" % " ".join(err.args))
164	return []	164	return []
165	-	165	+
166	@classmethod	166	@classmethod
167	def fetch_one_data(cls, cur: sqlite3.Cursor, sql: str) -> list:	167	def fetch_one_data(cls, cur: sqlite3.Cursor, sql: str) -> list:
168	"""	168	"""

Mtorch_npu/profiler/analysis/prof_common_func/_id_manager.py+2-2

	@@ -9,10 +9,10 @@ class Str2IdManager:
9	def __init__(self) -> None:	9	def __init__(self) -> None:
10	self._str_id_map = {}	10	self._str_id_map = {}
11	self._curr_id = 0	11	self._curr_id = 0
12	-	12	+
13	def set_start_id(self, start_id: int):	13	def set_start_id(self, start_id: int):
14	self._curr_id = start_id	14	self._curr_id = start_id
15	-	15	+
16	def get_id_from_str(self, string: str) -> int:	16	def get_id_from_str(self, string: str) -> int:
17	# 先查询；有性能影响的话直接+1，不查询了	17	# 先查询；有性能影响的话直接+1，不查询了
18	if not string:	18	if not string:

Mtorch_npu/profiler/analysis/prof_common_func/_singleton.py+0-1

	@@ -7,4 +7,3 @@ class Singleton(object):
7	if self._cls not in self._instance:	7	if self._cls not in self._instance:
8	self._instance[self._cls] = self._cls()	8	self._instance[self._cls] = self._cls()
9	return self._instance[self._cls]	9	return self._instance[self._cls]
10	-

Mtorch_npu/profiler/analysis/prof_common_func/_time_range_calculator.py+1-1

	@@ -27,7 +27,7 @@ class TimeRange:
27		27
28		28
29	class CommunicationTimeRange(TimeRange):	29	class CommunicationTimeRange(TimeRange):
30	-	30	+
31	def __init__(self):	31	def __init__(self):
32	super().__init__()	32	super().__init__()
33		33

Mtorch_npu/profiler/analysis/prof_config/_parser_config.py+1-1

	@@ -27,7 +27,7 @@ from ..prof_view._memory_view_parser import MemoryViewParser
27	from ..prof_view._integrate_parser import IntegrateParser	27	from ..prof_view._integrate_parser import IntegrateParser
28	from ..prof_view._communication_parser import CommunicationParser	28	from ..prof_view._communication_parser import CommunicationParser
29	from ..prof_view._memory_timeline_parser import MemoryTimelineParser	29	from ..prof_view._memory_timeline_parser import MemoryTimelineParser
30	-from ..prof_view.prof_db_parse._db_parser import DbParser	30	+from ..prof_view.prof_db_parse._db_parser import DbParser
31	from ..prof_view.prepare_parse._fwk_pre_parser import (	31	from ..prof_view.prepare_parse._fwk_pre_parser import (
32	TracePreParser,	32	TracePreParser,
33	TreeBuildParser,	33	TreeBuildParser,

Mtorch_npu/profiler/analysis/prof_parse/_event_tree_parser.py+39-39

	@@ -111,11 +111,11 @@ class _TensorMetadata:
111	self.device_type = device_type	111	self.device_type = device_type
112	self.device_index = device_index	112	self.device_index = device_index
113	self.id = _IDReference(None, None)	113	self.id = _IDReference(None, None)
114	-	114	+
115	@property	115	@property
116	def tensor_id(self) -> Optional[int]:	116	def tensor_id(self) -> Optional[int]:
117	return self.id.tensor_id	117	return self.id.tensor_id
118	-	118	+
119	@property	119	@property
120	def allocation_id(self) -> Optional[int]:	120	def allocation_id(self) -> Optional[int]:
121	return self.id.allocation_id	121	return self.id.allocation_id
	@@ -125,7 +125,7 @@ def parse_tensor_metadata(tensor_str: str) -> Optional[_TensorMetadata]:
125	parts = tensor_str.split(';')	125	parts = tensor_str.split(';')
126	if len(parts) != TensorEnum.NUM_FIELDS.value:	126	if len(parts) != TensorEnum.NUM_FIELDS.value:
127	return None	127	return None
128	-	128	+
129	impl = int(parts[TensorEnum.TENSOR_IMPL.value], BASE_16)	129	impl = int(parts[TensorEnum.TENSOR_IMPL.value], BASE_16)
130	ptr = int(parts[TensorEnum.STORAGE_PTR.value], BASE_16) if parts[TensorEnum.STORAGE_PTR.value] else None	130	ptr = int(parts[TensorEnum.STORAGE_PTR.value], BASE_16) if parts[TensorEnum.STORAGE_PTR.value] else None
131	dtype = parts[TensorEnum.DTYPE.value]	131	dtype = parts[TensorEnum.DTYPE.value]
	@@ -193,7 +193,7 @@ class _ExtraFields_TorchOp:
193	self.scope = bean.scope	193	self.scope = bean.scope
194	self.forward_tid = bean.args.get(Constant.FORWARD_THREAD_ID, -1)	194	self.forward_tid = bean.args.get(Constant.FORWARD_THREAD_ID, -1)
195	self.sequence_num = bean.args.get(Constant.SEQUENCE_NUMBER, -1)	195	self.sequence_num = bean.args.get(Constant.SEQUENCE_NUMBER, -1)
196	-	196	+
197	types_string = bean.args.get(Constant.INPUT_DTYPES, None)	197	types_string = bean.args.get(Constant.INPUT_DTYPES, None)
198	tensors_string = bean.inputs.get(Constant.INPUT_TENSORS, None)	198	tensors_string = bean.inputs.get(Constant.INPUT_TENSORS, None)
199	tensorlists_string = bean.inputs.get(Constant.INPUT_TENSORLISTS, None)	199	tensorlists_string = bean.inputs.get(Constant.INPUT_TENSORLISTS, None)
	@@ -211,11 +211,11 @@ class _ExtraFields_Allocation:
211	self.device_type = bean.device_type	211	self.device_type = bean.device_type
212	self.device_index = bean.device_index	212	self.device_index = bean.device_index
213	self.id = _IDReference(None, None)	213	self.id = _IDReference(None, None)
214	-	214	+
215	@property	215	@property
216	def tensor_id(self) -> Optional[int]:	216	def tensor_id(self) -> Optional[int]:
217	return self.id.tensor_id	217	return self.id.tensor_id
218	-	218	+
219	@property	219	@property
220	def allocation_id(self) -> Optional[int]:	220	def allocation_id(self) -> Optional[int]:
221	return self.id.allocation_id	221	return self.id.allocation_id
	@@ -242,7 +242,7 @@ def parse_module_param(param_str: str) -> Optional[_ModuleParam]:
242	param_list = param_str.strip().split(')')	242	param_list = param_str.strip().split(')')
243	if len(param_list) != ModuleParamEnum.NUM_FIELDS.value:	243	if len(param_list) != ModuleParamEnum.NUM_FIELDS.value:
244	return None	244	return None
245	-	245	+
246	name = param_list[ModuleParamEnum.NAME.value]	246	name = param_list[ModuleParamEnum.NAME.value]
247	tensor = parse_tensor_metadata(param_list[ModuleParamEnum.METADATA.value])	247	tensor = parse_tensor_metadata(param_list[ModuleParamEnum.METADATA.value])
248	grad = (parse_tensor_metadata(param_list[ModuleParamEnum.GRAD.value])	248	grad = (parse_tensor_metadata(param_list[ModuleParamEnum.GRAD.value])
	@@ -261,7 +261,7 @@ def parse_state_param(state_str: str) -> Optional[List[Tuple[str, _TensorMetadat
261	return None	261	return None
262	state_pairs.append(tuple([state_pair[StateParamEnum.NAME.value],	262	state_pairs.append(tuple([state_pair[StateParamEnum.NAME.value],
263	parse_tensor_metadata(state_pair[StateParamEnum.METADATA.value])]))	263	parse_tensor_metadata(state_pair[StateParamEnum.METADATA.value])]))
264	-	264	+
265	return state_pairs	265	return state_pairs
266		266
267		267
	@@ -269,7 +269,7 @@ def parse_optimizer_param(param_str: str) -> Optional[_OptimizerParam]:
269	param_list = param_str.strip().split(')')	269	param_list = param_str.strip().split(')')
270	if len(param_list) != OptimizerParamEnum.NUM_FIELDS.value:	270	if len(param_list) != OptimizerParamEnum.NUM_FIELDS.value:
271	return None	271	return None
272	-	272	+
273	tensor = parse_tensor_metadata(param_list[OptimizerParamEnum.METADATA.value])	273	tensor = parse_tensor_metadata(param_list[OptimizerParamEnum.METADATA.value])
274	grad = (parse_tensor_metadata(param_list[OptimizerParamEnum.GRAD.value])	274	grad = (parse_tensor_metadata(param_list[OptimizerParamEnum.GRAD.value])
275	if param_list[OptimizerParamEnum.GRAD.value] else None)	275	if param_list[OptimizerParamEnum.GRAD.value] else None)
	@@ -316,7 +316,7 @@ class _ProfilerEvent:
316	self.tid = bean.tid	316	self.tid = bean.tid
317	self.start_time_ns = bean.ts	317	self.start_time_ns = bean.ts
318	self.extra_fields = _ExtraFields_PyCall(bean)	318	self.extra_fields = _ExtraFields_PyCall(bean)
319	-	319	+
320	@property	320	@property
321	def name(self) -> str:	321	def name(self) -> str:
322	if self.tag == _EventType.TorchOp:	322	if self.tag == _EventType.TorchOp:
	@@ -326,7 +326,7 @@ class _ProfilerEvent:
326	elif self.tag == _EventType.PyCall:	326	elif self.tag == _EventType.PyCall:
327	return self.extra_fields.name	327	return self.extra_fields.name
328	return ""	328	return ""
329	-	329	+
330	@property	330	@property
331	def end_time_ns(self) -> int:	331	def end_time_ns(self) -> int:
332	if self.tag == _EventType.TorchOp:	332	if self.tag == _EventType.TorchOp:
	@@ -336,7 +336,7 @@ class _ProfilerEvent:
336	elif self.tag == _EventType.PyCall:	336	elif self.tag == _EventType.PyCall:
337	return self.extra_fields.end_time_ns	337	return self.extra_fields.end_time_ns
338	return -1	338	return -1
339	-	339	+
340	def __lt__(self, other: '_ProfilerEvent') -> bool:	340	def __lt__(self, other: '_ProfilerEvent') -> bool:
341	return self.end_time_ns < other.end_time_ns	341	return self.end_time_ns < other.end_time_ns
342		342
	@@ -362,17 +362,17 @@ def push_event(event: _ProfilerEvent,
362	if event.finished:	362	if event.finished:
363	print_error_msg("Error when building tree: the event finished.")	363	print_error_msg("Error when building tree: the event finished.")
364	return False	364	return False
365	-	365	+
366	parent = thread_event.get(event.tid)	366	parent = thread_event.get(event.tid)
367	if parent is None:	367	if parent is None:
368	fwd_tid = event.extra_fields.forward_tid if event.tag == _EventType.TorchOp else 0	368	fwd_tid = event.extra_fields.forward_tid if event.tag == _EventType.TorchOp else 0
369	if fwd_tid:	369	if fwd_tid:
370	parent = thread_event.get(fwd_tid)	370	parent = thread_event.get(fwd_tid)
371	-	371	+
372	if parent is not None:	372	if parent is not None:
373	event.parent = parent	373	event.parent = parent
374	parent.children.append(event)	374	parent.children.append(event)
375	-	375	+
376	if event.end_time_ns > event.start_time_ns:	376	if event.end_time_ns > event.start_time_ns:
377	thread_event[event.tid] = event	377	thread_event[event.tid] = event
378	unfinished_events.put((event.end_time_ns, event))	378	unfinished_events.put((event.end_time_ns, event))
	@@ -380,20 +380,20 @@ def push_event(event: _ProfilerEvent,
380	else:	380	else:
381	if not mark_finished(event):	381	if not mark_finished(event):
382	return False	382	return False
383	-	383	+
384	return True	384	return True
385		385
386		386
387	def pop_event(event: _ProfilerEvent, thread_event: Dict[int, _ProfilerEvent]) -> bool:	387	def pop_event(event: _ProfilerEvent, thread_event: Dict[int, _ProfilerEvent]) -> bool:
388	if event.finished:	388	if event.finished:
389	return True	389	return True
390	-	390	+
391	tid = event.tid	391	tid = event.tid
392	cur_event = thread_event.get(tid)	392	cur_event = thread_event.get(tid)
393	if cur_event is None:	393	if cur_event is None:
394	print_error_msg("Error when building tree: current event is none.")	394	print_error_msg("Error when building tree: current event is none.")
395	return False	395	return False
396	-	396	+
397	while cur_event != event:	397	while cur_event != event:
398	if not mark_finished(cur_event):	398	if not mark_finished(cur_event):
399	return False	399	return False
	@@ -401,13 +401,13 @@ def pop_event(event: _ProfilerEvent, thread_event: Dict[int, _ProfilerEvent]) ->
401	print_error_msg("Error when building tree: current event's parent is None.")	401	print_error_msg("Error when building tree: current event's parent is None.")
402	return False	402	return False
403	cur_event = cur_event.parent	403	cur_event = cur_event.parent
404	-	404	+
405	if not mark_finished(event):	405	if not mark_finished(event):
406	return False	406	return False
407	thread_event.pop(tid, None)	407	thread_event.pop(tid, None)
408	if event.parent:	408	if event.parent:
409	thread_event[tid] = event.parent	409	thread_event[tid] = event.parent
410	-	410	+
411	return True	411	return True
412		412
413		413
	@@ -423,7 +423,7 @@ def build_event_tree(sorted_events: List[_ProfilerEvent]) -> None:
423	return	423	return
424	if not push_event(ev, thread_event, unfinished_events):	424	if not push_event(ev, thread_event, unfinished_events):
425	return	425	return
426	-	426	+
427	# Cleanup remaining exit events.	427	# Cleanup remaining exit events.
428	while not unfinished_events.empty():	428	while not unfinished_events.empty():
429	_, top_event = unfinished_events.get()	429	_, top_event = unfinished_events.get()
	@@ -453,7 +453,7 @@ def get_tensor_info(sorted_events: List[_ProfilerEvent]) -> List[_RawTensorInfo]
453	elif ev.tag == _EventType.PyCall:	453	elif ev.tag == _EventType.PyCall:
454	if ev.extra_fields.key is None or ev.extra_fields.key in seen_pycalls:	454	if ev.extra_fields.key is None or ev.extra_fields.key in seen_pycalls:
455	continue	455	continue
456	-	456	+
457	seen_pycalls.add(ev.extra_fields.key)	457	seen_pycalls.add(ev.extra_fields.key)
458	if ev.extra_fields.module_parameters is not None:	458	if ev.extra_fields.module_parameters is not None:
459	for p in ev.extra_fields.module_parameters:	459	for p in ev.extra_fields.module_parameters:
	@@ -477,17 +477,17 @@ def get_tensor_info(sorted_events: List[_ProfilerEvent]) -> List[_RawTensorInfo]
477	_RawTensorInfo(t.impl, t.ptr, t.device_type, t.device_index, False, t.id)	477	_RawTensorInfo(t.impl, t.ptr, t.device_type, t.device_index, False, t.id)
478	for _, t in p.state	478	for _, t in p.state
479	)	479	)
480	-	480	+
481	return tensors	481	return tensors
482		482
483		483
484	# Assign Allocation ID for each Storage, and ID for each tensor.	484	# Assign Allocation ID for each Storage, and ID for each tensor.
485	-# A tensor has a unique id, but it can have multiple allocation IDs,	485	+# A tensor has a unique id, but it can have multiple allocation IDs,
486	-# because the tensor might use memory multiple times.	486	+# because the tensor might use memory multiple times.
487	def calculate_unique_id(sorted_events: List[_ProfilerEvent]):	487	def calculate_unique_id(sorted_events: List[_ProfilerEvent]):
488	# Step 1: Flatten events to a uniform representation	488	# Step 1: Flatten events to a uniform representation
489	tensors = get_tensor_info(sorted_events)	489	tensors = get_tensor_info(sorted_events)
490	-	490	+
491	# Step 2: Assign Allocation IDs for Storage	491	# Step 2: Assign Allocation IDs for Storage
492	counter: int = 0	492	counter: int = 0
493	storage_map: Dict[Tuple[int, int, int], int] = {}	493	storage_map: Dict[Tuple[int, int, int], int] = {}
	@@ -499,11 +499,11 @@ def calculate_unique_id(sorted_events: List[_ProfilerEvent]):
499	t.id_ref.allocation_id = storage_map[key]	499	t.id_ref.allocation_id = storage_map[key]
500	if t.is_free:	500	if t.is_free:
501	storage_map.pop(key, None)	501	storage_map.pop(key, None)
502	-	502	+
503	# Step 3: Handle allocation events which we cannot prove are for Tensor storage	503	# Step 3: Handle allocation events which we cannot prove are for Tensor storage
504	tensor_set = {t.id_ref.allocation_id for t in tensors if t.impl is not None}	504	tensor_set = {t.id_ref.allocation_id for t in tensors if t.impl is not None}
505	tensors = [t for t in tensors if t.id_ref.allocation_id in tensor_set]	505	tensors = [t for t in tensors if t.id_ref.allocation_id in tensor_set]
506	-	506	+
507	# Step 4: Assign tensor IDs using allocation IDs	507	# Step 4: Assign tensor IDs using allocation IDs
508	id_map: Dict[int, int] = {}	508	id_map: Dict[int, int] = {}
509	counter = 0	509	counter = 0
	@@ -511,7 +511,7 @@ def calculate_unique_id(sorted_events: List[_ProfilerEvent]):
511	if t.id_ref.allocation_id not in id_map:	511	if t.id_ref.allocation_id not in id_map:
512	id_map[t.id_ref.allocation_id] = counter	512	id_map[t.id_ref.allocation_id] = counter
513	counter += 1	513	counter += 1
514	-	514	+
515	# Step 5: Write back to Tensor IDs	515	# Step 5: Write back to Tensor IDs
516	for t in tensors:	516	for t in tensors:
517	if t.id_ref.allocation_id not in id_map:	517	if t.id_ref.allocation_id not in id_map:
	@@ -523,7 +523,7 @@ def calculate_unique_id(sorted_events: List[_ProfilerEvent]):
523	class EventTree:	523	class EventTree:
524	def __init__(self, profiler_path: str):	524	def __init__(self, profiler_path: str):
525	self.profiler_path = profiler_path	525	self.profiler_path = profiler_path
526	-	526	+
527	self.events: List[_ProfilerEvent] = []	527	self.events: List[_ProfilerEvent] = []
528	self.fetch_op_events(FwkFileParser(self.profiler_path))	528	self.fetch_op_events(FwkFileParser(self.profiler_path))
529	self.fetch_allocation_events(FwkFileParser(self.profiler_path))	529	self.fetch_allocation_events(FwkFileParser(self.profiler_path))
	@@ -538,7 +538,7 @@ class EventTree:
538	op_bean_list: List[TorchOpBean] = fwk_file_parser.get_file_data_by_tag(FileTag.TORCH_OP)	538	op_bean_list: List[TorchOpBean] = fwk_file_parser.get_file_data_by_tag(FileTag.TORCH_OP)
539	if not op_bean_list:	539	if not op_bean_list:
540	return	540	return
541	-	541	+
542	op_events = [_ProfilerEvent(op_bean) for op_bean in op_bean_list]	542	op_events = [_ProfilerEvent(op_bean) for op_bean in op_bean_list]
543		543
544	# Connected Autograd info to the top level annotation	544	# Connected Autograd info to the top level annotation
	@@ -548,22 +548,22 @@ class EventTree:
548	and op_events[i].extra_fields.name.startswith("autograd::engine::evaluate_function: ")):	548	and op_events[i].extra_fields.name.startswith("autograd::engine::evaluate_function: ")):
549	op_events[i].extra_fields.sequence_num = op_events[i + 1].extra_fields.sequence_num	549	op_events[i].extra_fields.sequence_num = op_events[i + 1].extra_fields.sequence_num
550	op_events[i].extra_fields.forward_tid = op_events[i + 1].extra_fields.forward_tid	550	op_events[i].extra_fields.forward_tid = op_events[i + 1].extra_fields.forward_tid
551	-	551	+
552	self.events.extend(op_events)	552	self.events.extend(op_events)
553	-	553	+
554	def fetch_allocation_events(self, fwk_file_parser: FwkFileParser) -> None:	554	def fetch_allocation_events(self, fwk_file_parser: FwkFileParser) -> None:
555	mem_bean_list: List[MemoryUseBean] = fwk_file_parser.get_file_data_by_tag(FileTag.MEMORY)	555	mem_bean_list: List[MemoryUseBean] = fwk_file_parser.get_file_data_by_tag(FileTag.MEMORY)
556	if not mem_bean_list:	556	if not mem_bean_list:
557	return	557	return
558	-	558	+
559	mem_events = [	559	mem_events = [
560	_ProfilerEvent(mem_bean)	560	_ProfilerEvent(mem_bean)
561	for mem_bean in mem_bean_list	561	for mem_bean in mem_bean_list
562	if mem_bean.data_type != _AllocEventType.BLOCK_FREE.value	562	if mem_bean.data_type != _AllocEventType.BLOCK_FREE.value
563	]	563	]
564	-	564	+
565	self.events.extend(mem_events)	565	self.events.extend(mem_events)
566	-	566	+
567	def fetch_pycall_events(self, fwk_file_parser: FwkFileParser) -> None:	567	def fetch_pycall_events(self, fwk_file_parser: FwkFileParser) -> None:
568	trace_hash_data = fwk_file_parser.get_file_data_by_tag(FileTag.PYTHON_TRACER_HASH)	568	trace_hash_data = fwk_file_parser.get_file_data_by_tag(FileTag.PYTHON_TRACER_HASH)
569	func_call_data = fwk_file_parser.get_file_data_by_tag(FileTag.PYTHON_TRACER_FUNC)	569	func_call_data = fwk_file_parser.get_file_data_by_tag(FileTag.PYTHON_TRACER_FUNC)
	@@ -573,18 +573,18 @@ class EventTree:
573	pycall_bean_list = python_trace_parser.get_pycall_data()	573	pycall_bean_list = python_trace_parser.get_pycall_data()
574	if not pycall_bean_list:	574	if not pycall_bean_list:
575	return	575	return
576	-	576	+
577	pycall_events = [_ProfilerEvent(pycall_bean) for pycall_bean in pycall_bean_list]	577	pycall_events = [_ProfilerEvent(pycall_bean) for pycall_bean in pycall_bean_list]
578		578
579	self.events.extend(pycall_events)	579	self.events.extend(pycall_events)
580	-	580	+
581	def validate_events(self) -> None:	581	def validate_events(self) -> None:
582	for ev in self.sorted_events:	582	for ev in self.sorted_events:
583	# Check the time of events is right	583	# Check the time of events is right
584	if ev.start_time_ns > ev.end_time_ns:	584	if ev.start_time_ns > ev.end_time_ns:
585	print_error_msg(f"Error in {ev.name}: {ev.start_time_ns} > {ev.end_time_ns}.")	585	print_error_msg(f"Error in {ev.name}: {ev.start_time_ns} > {ev.end_time_ns}.")
586	return	586	return
587	-	587	+
588	# Check the inputs in TorchOp	588	# Check the inputs in TorchOp
589	if ev.tag == _EventType.TorchOp:	589	if ev.tag == _EventType.TorchOp:
590	for i in ev.extra_fields.inputs:	590	for i in ev.extra_fields.inputs:

Mtorch_npu/profiler/analysis/prof_parse/_fwk_cann_relation_parser.py+1-1

	@@ -68,7 +68,7 @@ class FwkCANNRelationParser:
68	if not step_node_list:	68	if not step_node_list:
69	self.logger.warning("Get step range failed, the step node list is empty.")	69	self.logger.warning("Get step range failed, the step node list is empty.")
70	return []	70	return []
71	-	71	+
72	# Gather flow events start time in each step node	72	# Gather flow events start time in each step node
73	if not FwkFileParser(self._profiler_path).has_task_queue_data():	73	if not FwkFileParser(self._profiler_path).has_task_queue_data():
74	acl_start_time_list = sorted(list(kernel_dict.keys()))	74	acl_start_time_list = sorted(list(kernel_dict.keys()))

Mtorch_npu/profiler/analysis/prof_parse/_python_trace_parser.py+7-7

	@@ -10,10 +10,10 @@ __all__ = []
10	MODULE_NAME_DELIMITER = "######"	10	MODULE_NAME_DELIMITER = "######"
11		11
12		12
13	-class TraceTag(Enum):	13	+class TraceTag(Enum):
14	kPy_Call = 0	14	kPy_Call = 0
15	kPy_Return = 1	15	kPy_Return = 1
16	- kC_Call = 2	16	+ kC_Call = 2
17	kC_Return = 3	17	kC_Return = 3
18		18
19		19
	@@ -59,7 +59,7 @@ class PyTraceEvent:
59	@property	59	@property
60	def parent_id(self):	60	def parent_id(self):
61	return self._parent_id	61	return self._parent_id
62	-	62	+
63	@parent_id.setter	63	@parent_id.setter
64	def parent_id(self, parent_id):	64	def parent_id(self, parent_id):
65	self._parent_id = parent_id	65	self._parent_id = parent_id
	@@ -96,7 +96,7 @@ class PyTraceEvent:
96	@property	96	@property
97	def dur(self):	97	def dur(self):
98	return self._end_time - self._start_time	98	return self._end_time - self._start_time
99	-	99	+
100	@property	100	@property
101	def params(self):	101	def params(self):
102	return self._params	102	return self._params
	@@ -152,7 +152,7 @@ class PythonTraceParser:
152	trace_api_data[i] = [event.ts, event.ts + event.dur, contact_2num(event.pid, event.tid), None,	152	trace_api_data[i] = [event.ts, event.ts + event.dur, contact_2num(event.pid, event.tid), None,
153	str2id_manager.get_id_from_str(event.name), None, None, None, None, None, ApiType.PYTHON_TRACE]	153	str2id_manager.get_id_from_str(event.name), None, None, None, None, None, ApiType.PYTHON_TRACE]
154	return trace_api_data	154	return trace_api_data
155	-	155	+
156	def get_pycall_data(self) -> list:	156	def get_pycall_data(self) -> list:
157	self._gen_param_map()	157	self._gen_param_map()
158	return self._gen_python_trace_event_data()	158	return self._gen_python_trace_event_data()
	@@ -225,11 +225,11 @@ class PythonTraceParser:
225		225
226	def _gen_hash_map(self):	226	def _gen_hash_map(self):
227	self._hash_map = {hash_bean.key: hash_bean.value for hash_bean in self._hash_data}	227	self._hash_map = {hash_bean.key: hash_bean.value for hash_bean in self._hash_data}
228	-	228	+
229	def _gen_param_map(self):	229	def _gen_param_map(self):
230	if self._param_data is not None:	230	if self._param_data is not None:
231	self._param_map = {param_bean.key: param_bean.params for param_bean in self._param_data}	231	self._param_map = {param_bean.key: param_bean.params for param_bean in self._param_data}
232	-	232	+
233	@staticmethod	233	@staticmethod
234	def get_module_info_from_value(name: str, module_name_counter: dict, module_uid_map: dict):	234	def get_module_info_from_value(name: str, module_name_counter: dict, module_uid_map: dict):
235	module_name, module_uid = name.split(MODULE_NAME_DELIMITER)	235	module_name, module_uid = name.split(MODULE_NAME_DELIMITER)

Mtorch_npu/profiler/analysis/prof_view/_memory_timeline_parser.py+60-60

Mtorch_npu/profiler/analysis/prof_view/prof_db_parse/_basic_db_parser.py+2-2

	@@ -66,7 +66,7 @@ class BasicDbParser(BaseParser):
66	continue	66	continue
67	return file_path	67	return file_path
68	return ""	68	return ""
69	-	69	+
70	def create_ascend_db(self):	70	def create_ascend_db(self):
71	if not TorchDb().create_connect_db():	71	if not TorchDb().create_connect_db():
72	raise RuntimeError(f"Failed to connect to db file: {TorchDb().get_db_path()}")	72	raise RuntimeError(f"Failed to connect to db file: {TorchDb().get_db_path()}")
	@@ -89,7 +89,7 @@ class BasicDbParser(BaseParser):
89	rank_device_pairs.append([rank_id, device_id])	89	rank_device_pairs.append([rank_id, device_id])
90	TorchDb().insert_data_into_table(DbConstant.TABLE_RANK_DEVICE_MAP,	90	TorchDb().insert_data_into_table(DbConstant.TABLE_RANK_DEVICE_MAP,
91	rank_device_pairs)	91	rank_device_pairs)
92	-	92	+
93	def save_host_info_to_db(self):	93	def save_host_info_to_db(self):
94	if TorchDb().judge_table_exist(DbConstant.TABLE_HOST_INFO):	94	if TorchDb().judge_table_exist(DbConstant.TABLE_HOST_INFO):
95	return	95	return

Mtorch_npu/profiler/analysis/prof_view/prof_db_parse/_communication_db_parser.py+2-2

	@@ -97,7 +97,7 @@ class CommunicationDbParser(CommunicationParser):
97		97
98	def generate_view(self) -> None:	98	def generate_view(self) -> None:
99	self.generate_communication_db()	99	self.generate_communication_db()
100	-	100	+
101	def generate_communication_db(self):	101	def generate_communication_db(self):
102	db_files = CANNFileParser(self._profiler_path).get_file_list_by_type(CANNDataEnum.ANALYSIS_DB)	102	db_files = CANNFileParser(self._profiler_path).get_file_list_by_type(CANNDataEnum.ANALYSIS_DB)
103	if not db_files:	103	if not db_files:
	@@ -186,7 +186,7 @@ class CommunicationDbParser(CommunicationParser):
186	op_info.get(self.BANDWIDTH_GB_S), step, op_type, hccl_op_name	186	op_info.get(self.BANDWIDTH_GB_S), step, op_type, hccl_op_name
187	])	187	])
188	return res_data	188	return res_data
189	-	189	+
190	step_op_dict = {}	190	step_op_dict = {}
191	for data in matrix_data:	191	for data in matrix_data:
192	op_name = \	192	op_name = \

Mtorch_npu/profiler/analysis/prof_view/prof_db_parse/_memory_db_parser.py+4-4

Mtorch_npu/profiler/analysis/prof_view/prof_db_parse/_trace_step_time_db_parser.py+3-3

	@@ -140,7 +140,7 @@ class TraceStepTimeDbParser(BaseParser):
140	return	140	return
141	if TorchDb().judge_table_exist(DbConstant.TABLE_COMPUTE_TASK_INFO):	141	if TorchDb().judge_table_exist(DbConstant.TABLE_COMPUTE_TASK_INFO):
142	sql = """	142	sql = """
143	- SELECT	143	+ SELECT
144	STRING_IDS.value,	144	STRING_IDS.value,
145	task.startNs,	145	task.startNs,
146	task.endNs,	146	task.endNs,
	@@ -163,14 +163,14 @@ class TraceStepTimeDbParser(BaseParser):
163	connectionId	163	connectionId
164	FROM COMMUNICATION_OP c	164	FROM COMMUNICATION_OP c
165	)	165	)
166	- SELECT	166	+ SELECT
167	comm.opName,	167	comm.opName,
168	comm.startNs,	168	comm.startNs,
169	comm.endNs,	169	comm.endNs,
170	t.deviceId	170	t.deviceId
171	FROM comm_info comm	171	FROM comm_info comm
172	JOIN (	172	JOIN (
173	- SELECT	173	+ SELECT
174	connectionId,	174	connectionId,
175	deviceId	175	deviceId
176	FROM TASK	176	FROM TASK

Mtorch_npu/testing/common_utils.py+2-2

	@@ -206,7 +206,7 @@ class SupportedDevices:
206	reason = f"Only run on {repr(self.supported_devices)}, current device is {device_name}."	206	reason = f"Only run on {repr(self.supported_devices)}, current device is {device_name}."
207	raise unittest.SkipTest(reason)	207	raise unittest.SkipTest(reason)
208	return fn(slf, args, *kwargs)	208	return fn(slf, args, *kwargs)
209	-	209	+
210	return dep_fn	210	return dep_fn
211		211
212		212
	@@ -214,7 +214,7 @@ class SkipIfNotGteCANNVersion:
214	def __init__(self, base_version, module="CANN"):	214	def __init__(self, base_version, module="CANN"):
215	self.base_version = base_version	215	self.base_version = base_version
216	self.module = module	216	self.module = module
217	-	217	+
218	def __call__(self, fn):	218	def __call__(self, fn):
219	@wraps(fn)	219	@wraps(fn)
220	def func(slf, args, *kwargs):	220	def func(slf, args, *kwargs):

Mtorch_npu/testing/decorator.py+3-3

	@@ -81,7 +81,7 @@ def gen_ops_testcase(cls, func, name, keys, value, op_info):
81		81
82	def gen_op_input(testcase, func, op_info):	82	def gen_op_input(testcase, func, op_info):
83	data = {	83	data = {
84	- 'dtype': func.dtypes if hasattr(func, "dtypes") else op_info.dtypesIfNPU,	84	+ 'dtype': func.dtypes if hasattr(func, "dtypes") else op_info.dtypesIfNPU,
85	'npu_format': func.formats if hasattr(func, "formats") else op_info.formats	85	'npu_format': func.formats if hasattr(func, "formats") else op_info.formats
86	}	86	}
87		87
	@@ -98,7 +98,7 @@ def instantiate_ops_tests(op_db):
98		98
99	def wrapper(cls):	99	def wrapper(cls):
100	testcases = [x for x in dir(cls) if x.startswith('test_')]	100	testcases = [x for x in dir(cls) if x.startswith('test_')]
101	- for testcase in testcases:	101	+ for testcase in testcases:
102	if hasattr(cls, testcase):	102	if hasattr(cls, testcase):
103	func = getattr(cls, testcase)	103	func = getattr(cls, testcase)
104	for op_info in op_db:	104	for op_info in op_db:
	@@ -112,7 +112,7 @@ def instantiate_ops_tests(op_db):
112	delattr(cls, testcase)	112	delattr(cls, testcase)
113		113
114	return cls	114	return cls
115	-	115	+
116	return wrapper	116	return wrapper
117		117
118		118

Mtorch_npu/testing/testcase.py+4-4

	@@ -154,7 +154,7 @@ class TestCase(expecttest.TestCase):
154	self.assertEqual(tc._values(), t._values())	154	self.assertEqual(tc._values(), t._values())
155		155
156	return tg	156	return tg
157	-	157	+
158	def assertRtolEqual(self, x, y, prec=1.e-4, prec16=1.e-3, auto_trans_dtype=False, message=None):	158	def assertRtolEqual(self, x, y, prec=1.e-4, prec16=1.e-3, auto_trans_dtype=False, message=None):
159		159
160	def _assertRtolEqual(x, y, prec, prec16, message):	160	def _assertRtolEqual(x, y, prec, prec16, message):
	@@ -198,7 +198,7 @@ class TestCase(expecttest.TestCase):
198	self.fail("result error!")	198	self.fail("result error!")
199	return	199	return
200	x = x.detach().cpu().numpy()	200	x = x.detach().cpu().numpy()
201	- y = y.detach().cpu().numpy()	201	+ y = y.detach().cpu().numpy()
202	elif isinstance(x, Number) and isinstance(y, Number):	202	elif isinstance(x, Number) and isinstance(y, Number):
203	x = np.array(x)	203	x = np.array(x)
204	y = np.array(y)	204	y = np.array(y)
	@@ -208,7 +208,7 @@ class TestCase(expecttest.TestCase):
208	self.fail("shape error")	208	self.fail("shape error")
209	if (x.dtype != y.dtype):	209	if (x.dtype != y.dtype):
210	self.fail("dtype error")	210	self.fail("dtype error")
211	- dtype_list = [np.bool_, np.uint16, np.int16, np.int32, np.float16,	211	+ dtype_list = [np.bool_, np.uint16, np.int16, np.int32, np.float16,
212	np.float32, np.int8, np.uint8, np.int64, np.float64]	212	np.float32, np.int8, np.uint8, np.int64, np.float64]
213	if x.dtype not in dtype_list:	213	if x.dtype not in dtype_list:
214	self.fail("required dtype in [np.bool_, np.uint16, np.int16, " +	214	self.fail("required dtype in [np.bool_, np.uint16, np.int16, " +
	@@ -502,7 +502,7 @@ class TestCase(expecttest.TestCase):
502	def run(self, result=None):	502	def run(self, result=None):
503	# run test to precompile operators	503	# run test to precompile operators
504	super(TestCase, self).run(result)	504	super(TestCase, self).run(result)
505	-	505	+
506	if PERF_TEST_ENABLE:	506	if PERF_TEST_ENABLE:
507	performanceResult = TestResult()	507	performanceResult = TestResult()
508	startTime = time.perf_counter()	508	startTime = time.perf_counter()

Mtorch_npu/utils/_graph_tree.py+1-1

	@@ -228,7 +228,7 @@ def npugraphify_impl(
228		228
229	else:	229	else:
230	copy_indices = [	230	copy_indices = [
231	- idx	231	+ idx
232	for idx in range(len(static_inputs))	232	for idx in range(len(static_inputs))
233	if idx not in static_input_idxs	233	if idx not in static_input_idxs
234	]	234	]

Mtorch_npu/utils/_inductor.py+5-5

	@@ -54,8 +54,8 @@ def patch_register_philox_rand():
54	def get_register_philox_rand_patch():	54	def get_register_philox_rand_patch():
55	name = "philox_rand"	55	name = "philox_rand"
56	schema = "(SymInt[] size, Tensor seed, Tensor offset, int[]? stride, Device? device=None, ScalarType? dtype=None) -> (Tensor, Tensor)" # noqa: B950	56	schema = "(SymInt[] size, Tensor seed, Tensor offset, int[]? stride, Device? device=None, ScalarType? dtype=None) -> (Tensor, Tensor)" # noqa: B950
57	-	57	+
58	-	58	+
59	def _philox_rand_meta(	59	def _philox_rand_meta(
60	shape: torch.Size,	60	shape: torch.Size,
61	seed: torch.Tensor,	61	seed: torch.Tensor,
	@@ -71,7 +71,7 @@ def patch_register_philox_rand():
71	offset = philox_rand_offset_meta(shape)	71	offset = philox_rand_offset_meta(shape)
72	return (random_values, offset)	72	return (random_values, offset)
73		73
74	-	74	+
75	def _philox_rand(	75	def _philox_rand(
76	shape: torch.Size,	76	shape: torch.Size,
77	seed: torch.Tensor,	77	seed: torch.Tensor,
	@@ -85,13 +85,13 @@ def patch_register_philox_rand():
85	else:	85	else:
86	devices = [device]	86	devices = [device]
87		87
88	- with torch.random.fork_rng(devices, device_type="npu"):	88	+ with torch.random.fork_rng(devices, device_type="npu"):
89	CUDARngStateHelper.set_torch_state_tensor(seed, offset)	89	CUDARngStateHelper.set_torch_state_tensor(seed, offset)
90	random_values = torch.rand(shape, device=device, dtype=dtype)	90	random_values = torch.rand(shape, device=device, dtype=dtype)
91		91
92	return random_values, philox_rand_offset(shape)	92	return random_values, philox_rand_offset(shape)
93		93
94	-	94	+
95	register_rng_prim(	95	register_rng_prim(
96	name=name,	96	name=name,
97	schema=schema,	97	schema=schema,

Mtorch_npu/utils/_npu_meta_registration.py+1-1

	@@ -87,7 +87,7 @@ def patch_torch_inductor_decompositions():
87	don't accidentally overwrite unrelated inductor decompositions.	87	don't accidentally overwrite unrelated inductor decompositions.
88	'''	88	'''
89	import torch._inductor.decomposition as inductor_decomposition	89	import torch._inductor.decomposition as inductor_decomposition
90	-	90	+
91	for op_overload in inductor_decomp_table:	91	for op_overload in inductor_decomp_table:
92	if op_overload in npu_meta_table:	92	if op_overload in npu_meta_table:
93	inductor_decomposition.decompositions[op_overload] = npu_meta_table[op_overload]	93	inductor_decomposition.decompositions[op_overload] = npu_meta_table[op_overload]

Mtorch_npu/utils/_step.py+5-5

	@@ -37,7 +37,7 @@ class PerfDumpState:
37	if sub_module != module:	37	if sub_module != module:
38	module_list.append(sub_module)	38	module_list.append(sub_module)
39	self.module_dict[module] = module_list	39	self.module_dict[module] = module_list
40	-	40	+
41	def is_child_module(self, module):	41	def is_child_module(self, module):
42	for item in self.module_dict.items():	42	for item in self.module_dict.items():
43	if module in item[1]:	43	if module in item[1]:
	@@ -72,7 +72,7 @@ def _validate_path(path):
72	return True	72	return True
73	else:	73	else:
74	return False	74	return False
75	-	75	+
76		76
77	def _get_perf_dump_path():	77	def _get_perf_dump_path():
78	perf_dump_path = os.environ.get("PERF_DUMP_PATH")	78	perf_dump_path = os.environ.get("PERF_DUMP_PATH")
	@@ -85,7 +85,7 @@ def _get_perf_dump_path():
85	def delete_pref_pt_logs(perf_dump_path, device_id):	85	def delete_pref_pt_logs(perf_dump_path, device_id):
86	log_pattern = os.path.join(perf_dump_path, f"perf_pt__{device_id}.log")	86	log_pattern = os.path.join(perf_dump_path, f"perf_pt__{device_id}.log")
87	log_files = glob.glob(log_pattern)	87	log_files = glob.glob(log_pattern)
88	-	88	+
89	for log_file in log_files:	89	for log_file in log_files:
90	if os.path.islink(log_file):	90	if os.path.islink(log_file):
91	continue	91	continue
	@@ -101,9 +101,9 @@ def _get_uuid():
101		101
102	if master_addr is None or master_port is None:	102	if master_addr is None or master_port is None:
103	return "127.0.0.1_8888"	103	return "127.0.0.1_8888"
104	-	104	+
105	return master_addr + "_" + master_port	105	return master_addr + "_" + master_port
106	-	106	+
107		107
108	def _setup_logger(name, path):	108	def _setup_logger(name, path):
109	logger = logging.getLogger(name)	109	logger = logging.getLogger(name)

Mtorch_npu/utils/dlpack.py+10-10

Mtorch_npu/utils/flops_count.py+3-3

	@@ -7,10 +7,10 @@ __all__ = []
7	class _FlopsCounter:	7	class _FlopsCounter:
8	def __init__(self, ):	8	def __init__(self, ):
9	self.flop_count_instance = torch_npu._C._flops_count._FlopCountContext.GetInstance()	9	self.flop_count_instance = torch_npu._C._flops_count._FlopCountContext.GetInstance()
10	-	10	+
11	def __enter__(self):	11	def __enter__(self):
12	self.count_enable()	12	self.count_enable()
13	-	13	+
14	def __exit__(self):	14	def __exit__(self):
15	self.count_disable()	15	self.count_disable()
16		16
	@@ -20,7 +20,7 @@ class _FlopsCounter:
20	def stop(self):	20	def stop(self):
21	self.flop_count_instance.disable()	21	self.flop_count_instance.disable()
22	self.flop_count_instance.reset()	22	self.flop_count_instance.reset()
23	-	23	+
24	def pause(self):	24	def pause(self):
25	self.flop_count_instance.pause()	25	self.flop_count_instance.pause()
26		26

Mtorch_npu/utils/profiler.py+4-4

	@@ -26,7 +26,7 @@ class Profile(object):
26	save_path: str = "./npu_profiling",	26	save_path: str = "./npu_profiling",
27	profile_type: str = None,	27	profile_type: str = None,
28	use_npu=True,	28	use_npu=True,
29	- record_shape: bool = True,	29	+ record_shape: bool = True,
30	experimental_config: Optional[_ExperimentalConfig] = torch_npu.profiler._ExperimentalConfig(	30	experimental_config: Optional[_ExperimentalConfig] = torch_npu.profiler._ExperimentalConfig(
31	profiler_level=torch_npu.profiler.ProfilerLevel.Level2	31	profiler_level=torch_npu.profiler.ProfilerLevel.Level2
32	),	32	),
	@@ -72,9 +72,9 @@ class Profile(object):
72	raise ValueError("Args '%s' invaild, expect args '%s' ." % (kwargs.keys(), ascend_profiler_args_set) +	72	raise ValueError("Args '%s' invaild, expect args '%s' ." % (kwargs.keys(), ascend_profiler_args_set) +
73	prof_error(ErrCode.VALUE))	73	prof_error(ErrCode.VALUE))
74	self.prof = torch_npu.profiler.profile(	74	self.prof = torch_npu.profiler.profile(
75	- on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(self.save_path),	75	+ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(self.save_path),
76	- experimental_config=self.experimental_config,	76	+ experimental_config=self.experimental_config,
77	- record_shapes=self.record_shape,	77	+ record_shapes=self.record_shape,
78	**kwargs	78	**kwargs
79	)	79	)
80		80

Mtorch_npu/utils/tensor_methods.py+1-1

	@@ -70,7 +70,7 @@ class _NPUTensortypeCache(object):
70	def _npu_type(self, dtype=None, non_blocking=False, **kwargs):	70	def _npu_type(self, dtype=None, non_blocking=False, **kwargs):
71	if dtype is None:	71	if dtype is None:
72	return self.type_raw(dtype, non_blocking, **kwargs)	72	return self.type_raw(dtype, non_blocking, **kwargs)
73	-	73	+
74	_NPUTensortypeCache.tensortype_list_dict_init()	74	_NPUTensortypeCache.tensortype_list_dict_init()
75	if isinstance(dtype, str) and dtype in _NPUTensortypeCache.get_tensortype_dict():	75	if isinstance(dtype, str) and dtype in _NPUTensortypeCache.get_tensortype_dict():
76	tensortype_class = _NPUTensortypeCache.get_tensortype_dict()[dtype]	76	tensortype_class = _NPUTensortypeCache.get_tensortype_dict()[dtype]

Mtorchnpugen/__init__.py+1-1

	@@ -22,6 +22,6 @@ def _write_if_changed_security(self, filename: str, contents: str) -> None:
22		22
23	def apply_codegen_patches():	23	def apply_codegen_patches():
24	torchgen.gen.FileManager._write_if_changed = _write_if_changed_security	24	torchgen.gen.FileManager._write_if_changed = _write_if_changed_security
25	-	25	+
26		26
27	apply_codegen_patches()	27	apply_codegen_patches()

Mtorchnpugen/autograd/gen_autograd_functions.py+1-1

	@@ -47,7 +47,7 @@ def gen_autograd_functions_python(
47	infos,	47	infos,
48	key_fn=lambda info: info.name,	48	key_fn=lambda info: info.name,
49	base_env={	49	base_env={
50	- "generated_comment":	50	+ "generated_comment":
51	f"@ generated from {fm.template_dir_for_comments()}/python_functions.cpp",	51	f"@ generated from {fm.template_dir_for_comments()}/python_functions.cpp",
52	},	52	},
53	env_callable=lambda info: {	53	env_callable=lambda info: {

Mtorchnpugen/autograd/gen_variable_type.py+2-2

	@@ -56,7 +56,7 @@ def gen_variable_type(
56	template_path: str,	56	template_path: str,
57	) -> None:	57	) -> None:
58	"""Generate VariableType.cpp body	58	"""Generate VariableType.cpp body
59	-	59	+
60	Generate variable type definition for torch and npu method here.	60	Generate variable type definition for torch and npu method here.
61	"""	61	"""
62	fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)	62	fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)
	@@ -80,7 +80,7 @@ def gen_variable_type_head(
80	fns_with_diff_infos: List[NativeFunctionWithDifferentiabilityInfo],	80	fns_with_diff_infos: List[NativeFunctionWithDifferentiabilityInfo],
81	template_path: str,	81	template_path: str,
82	) -> None:	82	) -> None:
83	-	83	+
84	"""Generate VariableType.h body	84	"""Generate VariableType.h body
85	"""	85	"""
86	fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)	86	fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)

Mtorchnpugen/autograd/utils.py+1-1

	@@ -46,7 +46,7 @@ def parse_derivatives(
46	# original code logic	46	# original code logic
47	derivatives_path = str(Path(autograd_dir).parents[1].joinpath(	47	derivatives_path = str(Path(autograd_dir).parents[1].joinpath(
48	f'third_party/op-plugin/op_plugin/config/v{VERSION_PART[0]}r{VERSION_PART[1]}/derivatives.yaml'	48	f'third_party/op-plugin/op_plugin/config/v{VERSION_PART[0]}r{VERSION_PART[1]}/derivatives.yaml'
49	- ))	49	+ ))
50		50
51	differentiability_infos, _ = load_derivatives(	51	differentiability_infos, _ = load_derivatives(
52	derivatives_path, native_functions_path, tags_path)	52	derivatives_path, native_functions_path, tags_path)

Mtorchnpugen/codegen_ops_info.py+1-1

	@@ -15,7 +15,7 @@ from torchnpugen.utils import PathManager
15	project_path = Path(os.path.dirname(__file__)).parent	15	project_path = Path(os.path.dirname(__file__)).parent
16	op_plugin_info_path = os.path.realpath(os.path.join(	16	op_plugin_info_path = os.path.realpath(os.path.join(
17	project_path,	17	project_path,
18	- f'third_party/op-plugin/test/test_v{VERSION_PART[0]}r{VERSION_PART[1]}_ops',	18	+ f'third_party/op-plugin/test/test_v{VERSION_PART[0]}r{VERSION_PART[1]}_ops',
19	"unsupported_ops_info.yaml"))	19	"unsupported_ops_info.yaml"))
20	torch_npu_info_path = os.path.realpath(os.path.join(project_path, "test", "unsupported_ops_info.yaml"))	20	torch_npu_info_path = os.path.realpath(os.path.join(project_path, "test", "unsupported_ops_info.yaml"))
21		21

Mtorchnpugen/templates/npu_testing_utils.py+1-1

	@@ -5,7 +5,7 @@ from torch.testing._internal.common_methods_invocations import op_db, python_ref
5	from torch.testing._internal.opinfo.core import DecorateInfo	5	from torch.testing._internal.opinfo.core import DecorateInfo
6		6
7	"""	7	"""
8	-strategy: Due to the restriction of NPU operators.	8	+strategy: Due to the restriction of NPU operators.
9	patch the data classes to avoid unsupported cases.	9	patch the data classes to avoid unsupported cases.
10	"""	10	"""
11		11

	@@ -182,8 +182,8 @@ class TestAscendGraphPass(TestUtils):
182	cast_1 = torch.ops.npu._npu_dtype_cast.default(first_element, torch.int64)	182	cast_1 = torch.ops.npu._npu_dtype_cast.default(first_element, torch.int64)
183	output = torch.ops.aten.relu.default(cast_1)	183	output = torch.ops.aten.relu.default(cast_1)
184	return output	184	return output
185	-	185	+
186	-	186	+
187	@parametrize('shape', [(256, 5)])	187	@parametrize('shape', [(256, 5)])
188	@parametrize('dtype', ['int64'])	188	@parametrize('dtype', ['int64'])
189	def test_cast_standard_compile_cases(self, shape, dtype):	189	def test_cast_standard_compile_cases(self, shape, dtype):
	@@ -192,7 +192,7 @@ class TestAscendGraphPass(TestUtils):
192	compiled_op_calc = torch.compile(self.cast_standard_op_calc, backend="inductor")	192	compiled_op_calc = torch.compile(self.cast_standard_op_calc, backend="inductor")
193	inductor_result = compiled_op_calc(first_element)	193	inductor_result = compiled_op_calc(first_element)
194	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	194	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
195	-	195	+
196		196
197	@parametrize('shape', [(256, 5)])	197	@parametrize('shape', [(256, 5)])
198	@parametrize('dtype', ['int64'])	198	@parametrize('dtype', ['int64'])
	@@ -243,8 +243,8 @@ class TestAscendGraphPass(TestUtils):
243	compiled_op_calc = torch.compile(self.cat_slice_cat_op_calc, backend="inductor")	243	compiled_op_calc = torch.compile(self.cat_slice_cat_op_calc, backend="inductor")
244	inductor_result = compiled_op_calc(first_element)	244	inductor_result = compiled_op_calc(first_element)
245	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	245	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
246	-	246	+
247	-	247	+
248	@parametrize('shape', [(1, 3, 64, 1024)])	248	@parametrize('shape', [(1, 3, 64, 1024)])
249	@parametrize('dtype', ['float32'])	249	@parametrize('dtype', ['float32'])
250	def test_cat_slice_cat_ut_cases(self, shape, dtype):	250	def test_cat_slice_cat_ut_cases(self, shape, dtype):
	@@ -276,8 +276,8 @@ class TestAscendGraphPass(TestUtils):
276	compiled_op_calc = torch.compile(self.fold_add_op_calc, backend="inductor")	276	compiled_op_calc = torch.compile(self.fold_add_op_calc, backend="inductor")
277	inductor_result = compiled_op_calc(first_element)	277	inductor_result = compiled_op_calc(first_element)
278	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	278	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
279	-	279	+
280	-	280	+
281	@parametrize('shape', [(1, 2, 3)])	281	@parametrize('shape', [(1, 2, 3)])
282	@parametrize('dtype', ['float32'])	282	@parametrize('dtype', ['float32'])
283	def test_fold_add_ut_cases(self, shape, dtype):	283	def test_fold_add_ut_cases(self, shape, dtype):
	@@ -316,8 +316,8 @@ class TestAscendGraphPass(TestUtils):
316	compiled_op_calc = torch.compile(self.fold_cat_op_calc, backend="inductor")	316	compiled_op_calc = torch.compile(self.fold_cat_op_calc, backend="inductor")
317	inductor_result = compiled_op_calc(t1, t2, t3, t4, t5)	317	inductor_result = compiled_op_calc(t1, t2, t3, t4, t5)
318	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	318	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
319	-	319	+
320	-	320	+
321	@parametrize('shape', [(2, 4)])	321	@parametrize('shape', [(2, 4)])
322	@parametrize('dtype', ['float32'])	322	@parametrize('dtype', ['float32'])
323	def test_fold_cat_ut_cases(self, shape, dtype):	323	def test_fold_cat_ut_cases(self, shape, dtype):
	@@ -347,14 +347,14 @@ class TestAscendGraphPass(TestUtils):
347	@parametrize('dtype', ['float32'])	347	@parametrize('dtype', ['float32'])
348	def test_fold_clone_compile_cases(self, shape, dtype):	348	def test_fold_clone_compile_cases(self, shape, dtype):
349	t1 = self._generate_tensor(shape, dtype)	349	t1 = self._generate_tensor(shape, dtype)
350	-	350	+
351	std_result = self.fold_clone_op_calc(t1)	351	std_result = self.fold_clone_op_calc(t1)
352		352
353	compiled_op_calc = torch.compile(self.fold_clone_op_calc, backend="inductor")	353	compiled_op_calc = torch.compile(self.fold_clone_op_calc, backend="inductor")
354	inductor_result = compiled_op_calc(t1)	354	inductor_result = compiled_op_calc(t1)
355	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	355	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
356	-	356	+
357	-	357	+
358	@parametrize('shape', [(1, 2, 3)])	358	@parametrize('shape', [(1, 2, 3)])
359	@parametrize('dtype', ['float32'])	359	@parametrize('dtype', ['float32'])
360	def test_fold_clone_ut_cases(self, shape, dtype):	360	def test_fold_clone_ut_cases(self, shape, dtype):
	@@ -385,8 +385,8 @@ class TestAscendGraphPass(TestUtils):
385	compiled_op_calc = torch.compile(self.fold_detach_op_calc, backend="inductor")	385	compiled_op_calc = torch.compile(self.fold_detach_op_calc, backend="inductor")
386	inductor_result = compiled_op_calc(t1)	386	inductor_result = compiled_op_calc(t1)
387	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	387	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
388	-	388	+
389	-	389	+
390	@parametrize('shape', [(3, 3)])	390	@parametrize('shape', [(3, 3)])
391	@parametrize('dtype', ['float32'])	391	@parametrize('dtype', ['float32'])
392	def test_fold_detach_ut_cases(self, shape, dtype):	392	def test_fold_detach_ut_cases(self, shape, dtype):
	@@ -417,8 +417,8 @@ class TestAscendGraphPass(TestUtils):
417	compiled_op_calc = torch.compile(self.fold_div_op_calc, backend="inductor")	417	compiled_op_calc = torch.compile(self.fold_div_op_calc, backend="inductor")
418	inductor_result = compiled_op_calc(t1)	418	inductor_result = compiled_op_calc(t1)
419	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	419	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
420	-	420	+
421	-	421	+
422	@parametrize('shape', [(2, 4, 8)])	422	@parametrize('shape', [(2, 4, 8)])
423	@parametrize('dtype', ['float32'])	423	@parametrize('dtype', ['float32'])
424	def test_fold_div_ut_cases(self, shape, dtype):	424	def test_fold_div_ut_cases(self, shape, dtype):
	@@ -446,14 +446,14 @@ class TestAscendGraphPass(TestUtils):
446	@parametrize('dtype', ['float32'])	446	@parametrize('dtype', ['float32'])
447	def test_fold_expand_compile_cases(self, shape, dtype):	447	def test_fold_expand_compile_cases(self, shape, dtype):
448	t1 = self._generate_tensor(shape, dtype)	448	t1 = self._generate_tensor(shape, dtype)
449	-	449	+
450	std_result = self.fold_expand_op_calc(t1)	450	std_result = self.fold_expand_op_calc(t1)
451		451
452	compiled_op_calc = torch.compile(self.fold_expand_op_calc, backend="inductor")	452	compiled_op_calc = torch.compile(self.fold_expand_op_calc, backend="inductor")
453	inductor_result = compiled_op_calc(t1)	453	inductor_result = compiled_op_calc(t1)
454	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	454	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
455	-	455	+
456	-	456	+
457	@parametrize('shape', [(256, 128, 1)])	457	@parametrize('shape', [(256, 128, 1)])
458	@parametrize('dtype', ['float32'])	458	@parametrize('dtype', ['float32'])
459	def test_fold_expand_ut_cases(self, shape, dtype):	459	def test_fold_expand_ut_cases(self, shape, dtype):
	@@ -484,8 +484,8 @@ class TestAscendGraphPass(TestUtils):
484	compiled_op_calc = torch.compile(self.fold_mul_op_calc, backend="inductor")	484	compiled_op_calc = torch.compile(self.fold_mul_op_calc, backend="inductor")
485	inductor_result = compiled_op_calc(t1)	485	inductor_result = compiled_op_calc(t1)
486	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	486	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
487	-	487	+
488	-	488	+
489	@parametrize('shape', [(1, 2, 3)])	489	@parametrize('shape', [(1, 2, 3)])
490	@parametrize('dtype', ['float32'])	490	@parametrize('dtype', ['float32'])
491	def test_fold_mul_ut_cases(self, shape, dtype):	491	def test_fold_mul_ut_cases(self, shape, dtype):
	@@ -514,8 +514,8 @@ class TestAscendGraphPass(TestUtils):
514	compiled_op_calc = torch.compile(self.fold_reduce_op_calc, backend="inductor")	514	compiled_op_calc = torch.compile(self.fold_reduce_op_calc, backend="inductor")
515	inductor_result = compiled_op_calc(t1)	515	inductor_result = compiled_op_calc(t1)
516	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	516	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
517	-	517	+
518	-	518	+
519	@parametrize('shape', [(128, 1, 64, 1)])	519	@parametrize('shape', [(128, 1, 64, 1)])
520	@parametrize('dtype', ['float32'])	520	@parametrize('dtype', ['float32'])
521	def test_fold_reduce_ut_cases(self, shape, dtype):	521	def test_fold_reduce_ut_cases(self, shape, dtype):
	@@ -564,8 +564,8 @@ class TestAscendGraphPass(TestUtils):
564	compiled_op_calc = torch.compile(self.fold_redundant_op_calc, backend="inductor")	564	compiled_op_calc = torch.compile(self.fold_redundant_op_calc, backend="inductor")
565	inductor_result = compiled_op_calc(arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1)	565	inductor_result = compiled_op_calc(arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1)
566	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	566	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
567	-	567	+
568	-	568	+
569	def test_fold_redundant_ut_cases(self):	569	def test_fold_redundant_ut_cases(self):
570	arg0_1 = torch.randn(289094, 64, dtype=torch.float32)	570	arg0_1 = torch.randn(289094, 64, dtype=torch.float32)
571	arg1_1 = torch.randint(0, 289094, (128,), dtype=torch.int64)	571	arg1_1 = torch.randint(0, 289094, (128,), dtype=torch.int64)
	@@ -600,8 +600,8 @@ class TestAscendGraphPass(TestUtils):
600	compiled_op_calc = torch.compile(self.fold_sink_view_op_calc, backend="inductor")	600	compiled_op_calc = torch.compile(self.fold_sink_view_op_calc, backend="inductor")
601	inductor_result = compiled_op_calc(t1)	601	inductor_result = compiled_op_calc(t1)
602	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	602	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
603	-	603	+
604	-	604	+
605	@parametrize('shape', [(2, 3, 4)])	605	@parametrize('shape', [(2, 3, 4)])
606	@parametrize('dtype', ['float32'])	606	@parametrize('dtype', ['float32'])
607	def test_fold_sink_viewut_cases(self, shape, dtype):	607	def test_fold_sink_viewut_cases(self, shape, dtype):
	@@ -637,8 +637,8 @@ class TestAscendGraphPass(TestUtils):
637	compiled_op_calc = torch.compile(self.fold_slice_op_calc, backend="inductor")	637	compiled_op_calc = torch.compile(self.fold_slice_op_calc, backend="inductor")
638	inductor_result = compiled_op_calc(base, view, t1, t2, t3)	638	inductor_result = compiled_op_calc(base, view, t1, t2, t3)
639	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	639	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
640	-	640	+
641	-	641	+
642	def test_fold_slice_ut_cases(self):	642	def test_fold_slice_ut_cases(self):
643	base = torch.randn(8, 16, 32)	643	base = torch.randn(8, 16, 32)
644	view = torch.ones(8, 16, 32)	644	view = torch.ones(8, 16, 32)
	@@ -671,8 +671,8 @@ class TestAscendGraphPass(TestUtils):
671	compiled_op_calc = torch.compile(self.fold_squeeze_op_calc, backend="inductor")	671	compiled_op_calc = torch.compile(self.fold_squeeze_op_calc, backend="inductor")
672	inductor_result = compiled_op_calc(t1, t2)	672	inductor_result = compiled_op_calc(t1, t2)
673	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	673	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
674	-	674	+
675	-	675	+
676	def test_fold_squeeze_ut_cases(self):	676	def test_fold_squeeze_ut_cases(self):
677	t1 = torch.randn(2, 4)	677	t1 = torch.randn(2, 4)
678	t2 = torch.randn(2, 1, 1, 4)	678	t2 = torch.randn(2, 1, 1, 4)
	@@ -703,8 +703,8 @@ class TestAscendGraphPass(TestUtils):
703	compiled_op_calc = torch.compile(self.fold_sub_op_calc, backend="inductor")	703	compiled_op_calc = torch.compile(self.fold_sub_op_calc, backend="inductor")
704	inductor_result = compiled_op_calc(t1)	704	inductor_result = compiled_op_calc(t1)
705	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	705	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
706	-	706	+
707	-	707	+
708	@parametrize('shape', [(1, 2, 3)])	708	@parametrize('shape', [(1, 2, 3)])
709	@parametrize('dtype', ['float32'])	709	@parametrize('dtype', ['float32'])
710	def test_fold_sub_ut_cases(self, shape, dtype):	710	def test_fold_sub_ut_cases(self, shape, dtype):
	@@ -734,8 +734,8 @@ class TestAscendGraphPass(TestUtils):
734	compiled_op_calc = torch.compile(self.fold_to_copy_op_calc, backend="inductor")	734	compiled_op_calc = torch.compile(self.fold_to_copy_op_calc, backend="inductor")
735	inductor_result = compiled_op_calc(t1)	735	inductor_result = compiled_op_calc(t1)
736	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	736	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
737	-	737	+
738	-	738	+
739	@parametrize('shape', [(2, 4)])	739	@parametrize('shape', [(2, 4)])
740	@parametrize('dtype', ['float32'])	740	@parametrize('dtype', ['float32'])
741	def test_fold_to_copy_ut_cases(self, shape, dtype):	741	def test_fold_to_copy_ut_cases(self, shape, dtype):
	@@ -766,8 +766,8 @@ class TestAscendGraphPass(TestUtils):
766	compiled_op_calc = torch.compile(self.fold_view_op_calc, backend="inductor")	766	compiled_op_calc = torch.compile(self.fold_view_op_calc, backend="inductor")
767	inductor_result = compiled_op_calc(t1, t2)	767	inductor_result = compiled_op_calc(t1, t2)
768	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	768	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
769	-	769	+
770	-	770	+
771	def test_fold_view_ut_cases(self):	771	def test_fold_view_ut_cases(self):
772	t1 = torch.randn(1, 3, 1, 5)	772	t1 = torch.randn(1, 3, 1, 5)
773	t2 = torch.randn(128, 64)	773	t2 = torch.randn(128, 64)
	@@ -795,8 +795,8 @@ class TestAscendGraphPass(TestUtils):
795	compiled_op_calc = torch.compile(self.fold_where_op_calc, backend="inductor")	795	compiled_op_calc = torch.compile(self.fold_where_op_calc, backend="inductor")
796	inductor_result = compiled_op_calc(t1)	796	inductor_result = compiled_op_calc(t1)
797	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	797	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
798	-	798	+
799	-	799	+
800	@parametrize('shape', [(3, 4)])	800	@parametrize('shape', [(3, 4)])
801	@parametrize('dtype', ['float32'])	801	@parametrize('dtype', ['float32'])
802	def test_fold_where_ut_cases(self, shape, dtype):	802	def test_fold_where_ut_cases(self, shape, dtype):
	@@ -827,8 +827,8 @@ class TestAscendGraphPass(TestUtils):
827	compiled_op_calc = torch.compile(self.pad_slice_op_calc, backend="inductor")	827	compiled_op_calc = torch.compile(self.pad_slice_op_calc, backend="inductor")
828	inductor_result = compiled_op_calc(t1)	828	inductor_result = compiled_op_calc(t1)
829	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)	829	self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
830	-	830	+
831	-	831	+
832	@parametrize('shape', [(128, 50, 128)])	832	@parametrize('shape', [(128, 50, 128)])
833	@parametrize('dtype', ['float32'])	833	@parametrize('dtype', ['float32'])
834	def test_pad_slice_ut_cases(self, shape, dtype):	834	def test_pad_slice_ut_cases(self, shape, dtype):

	@@ -7,25 +7,25 @@ import os
7	os.environ["NPU_INDUCTOR_FALLBACK_LIST"] = "aten.div,aten.add.Tensor"	7	os.environ["NPU_INDUCTOR_FALLBACK_LIST"] = "aten.div,aten.add.Tensor"
8		8
9	class TestFallback(TestUtils):	9	class TestFallback(TestUtils):
10	-	10	+
11	def add_op(self, x, y):	11	def add_op(self, x, y):
12	return x / y	12	return x / y
13	-	13	+
14	def test_add_fallback_detection(self):	14	def test_add_fallback_detection(self):
15	-	15	+
16	compiled_add = torch.compile(self.add_op, backend="inductor")	16	compiled_add = torch.compile(self.add_op, backend="inductor")
17	-	17	+
18	x = torch.randn(4, 4, dtype=torch.float32).to("npu")	18	x = torch.randn(4, 4, dtype=torch.float32).to("npu")
19	y = torch.randn(4, 4, dtype=torch.float32).to("npu")	19	y = torch.randn(4, 4, dtype=torch.float32).to("npu")
20		20
21	_ , codes = run_and_get_code(compiled_add, x, y)	21	_ , codes = run_and_get_code(compiled_add, x, y)
22		22
23	self.assertTrue('unk_fused_div' not in codes[0])	23	self.assertTrue('unk_fused_div' not in codes[0])
24	-	24	+
25	def test_add_fallback_detection_mlir(self):	25	def test_add_fallback_detection_mlir(self):
26		26
27	compiled_add = torch.compile(self.add_op, backend="inductor", options={"npu_backend": "mlir"})	27	compiled_add = torch.compile(self.add_op, backend="inductor", options={"npu_backend": "mlir"})
28	-	28	+
29	x = torch.randn(4, 4, dtype=torch.float32).to("npu")	29	x = torch.randn(4, 4, dtype=torch.float32).to("npu")
30	y = torch.randn(4, 4, dtype=torch.float32).to("npu")	30	y = torch.randn(4, 4, dtype=torch.float32).to("npu")
31		31

	@@ -59,7 +59,7 @@ class OptionsTest(TestCase):
59	1. HCCL config is correctly applied to default process group	59	1. HCCL config is correctly applied to default process group
60	2. New process group inherits the correct HCCL configuration	60	2. New process group inherits the correct HCCL configuration
61	3. all_reduce executes successfully on NPU	61	3. all_reduce executes successfully on NPU
62	-	62	+
63	Args:	63	Args:
64	rank (int): Current process rank	64	rank (int): Current process rank
65	ranks (list[int]): List of ranks in the process group	65	ranks (list[int]): List of ranks in the process group
	@@ -86,32 +86,32 @@ class OptionsTest(TestCase):
86	# 4. Move input tensor to target NPU device (specify rank to avoid device conflict)	86	# 4. Move input tensor to target NPU device (specify rank to avoid device conflict)
87	input1 = input1.npu(rank)	87	input1 = input1.npu(rank)
88	test_case.assertEqual(	88	test_case.assertEqual(
89	- input1.device,	89	+ input1.device,
90	- torch.device(f'npu:{rank}'),	90	+ torch.device(f'npu:{rank}'),
91	"Tensor not correctly moved to target NPU device"	91	"Tensor not correctly moved to target NPU device"
92	)	92	)
93		93
94	# 5. Execute all_reduce on default process group and validate configuration	94	# 5. Execute all_reduce on default process group and validate configuration
95	dist.all_reduce(input1)	95	dist.all_reduce(input1)
96	-	96	+
97	# Get default process group's HCCL backend	97	# Get default process group's HCCL backend
98	default_pg = c10d._get_default_group()._get_backend(torch.device(f'npu:{rank}'))	98	default_pg = c10d._get_default_group()._get_backend(torch.device(f'npu:{rank}'))
99	-	99	+
100	# Validate full HCCL configuration	100	# Validate full HCCL configuration
101	test_case.assertEqual(	101	test_case.assertEqual(
102	- default_pg.options.hccl_config,	102	+ default_pg.options.hccl_config,
103	cls.HCCL_DEFAULT_CONFIG,	103	cls.HCCL_DEFAULT_CONFIG,
104	"Default process group HCCL config does not match expected configuration"	104	"Default process group HCCL config does not match expected configuration"
105	)	105	)
106	-	106	+
107	# Validate individual config items (add default value to avoid KeyError)	107	# Validate individual config items (add default value to avoid KeyError)
108	test_case.assertEqual(	108	test_case.assertEqual(
109	- default_pg.options.hccl_config.get("hccl_exec_timeout", -1),	109	+ default_pg.options.hccl_config.get("hccl_exec_timeout", -1),
110	500,	110	500,
111	"hccl_exec_timeout config value mismatch"	111	"hccl_exec_timeout config value mismatch"
112	)	112	)
113	test_case.assertEqual(	113	test_case.assertEqual(
114	- default_pg.options.hccl_config.get("hccl_algo", ""),	114	+ default_pg.options.hccl_config.get("hccl_algo", ""),
115	"allreduce=level0:NA;level1:ring/allgather=level0:NA;level1:H-D_R",	115	"allreduce=level0:NA;level1:ring/allgather=level0:NA;level1:H-D_R",
116	"hccl_algo config value mismatch"	116	"hccl_algo config value mismatch"
117	)	117	)
	@@ -119,22 +119,22 @@ class OptionsTest(TestCase):
119	# 6. Create new process group with same options and validate configuration	119	# 6. Create new process group with same options and validate configuration
120	new_pg = dist.new_group(backend='hccl', ranks=ranks, pg_options=options)	120	new_pg = dist.new_group(backend='hccl', ranks=ranks, pg_options=options)
121	test_case.assertTrue(new_pg is not None, "Failed to create new HCCL process group")	121	test_case.assertTrue(new_pg is not None, "Failed to create new HCCL process group")
122	-	122	+
123	# Validate new process group's HCCL configuration	123	# Validate new process group's HCCL configuration
124	new_pg_backend = new_pg._get_backend(torch.device(f'npu:{rank}'))	124	new_pg_backend = new_pg._get_backend(torch.device(f'npu:{rank}'))
125	test_case.assertEqual(	125	test_case.assertEqual(
126	- new_pg_backend.options.hccl_config,	126	+ new_pg_backend.options.hccl_config,
127	cls.HCCL_DEFAULT_CONFIG,	127	cls.HCCL_DEFAULT_CONFIG,
128	"New process group HCCL config does not match expected configuration"	128	"New process group HCCL config does not match expected configuration"
129	)	129	)
130	-	130	+
131	# Execute all_reduce on new process group	131	# Execute all_reduce on new process group
132	dist.all_reduce(input1, group=new_pg)	132	dist.all_reduce(input1, group=new_pg)
133		133
134	except Exception as e:	134	except Exception as e:
135	# Capture exceptions and mark test as failed	135	# Capture exceptions and mark test as failed
136	test_case.fail(f"Test execution failed with error: {str(e)}")	136	test_case.fail(f"Test execution failed with error: {str(e)}")
137	-	137	+
138	finally:	138	finally:
139	# 7. Clean up resources to prevent memory leaks	139	# 7. Clean up resources to prevent memory leaks
140	# Destroy custom process group if created	140	# Destroy custom process group if created
	@@ -156,7 +156,7 @@ class OptionsTest(TestCase):
156	with test_case.assertRaises(RuntimeError) as cm:	156	with test_case.assertRaises(RuntimeError) as cm:
157	OptionsTest._init_dist_hccl(rank, options, world_size)	157	OptionsTest._init_dist_hccl(rank, options, world_size)
158	dist.all_reduce(input1)	158	dist.all_reduce(input1)
159	-	159	+
160	test_case.assertTrue(error_expect in str(cm.exception),	160	test_case.assertTrue(error_expect in str(cm.exception),
161	f"Expected error messages '{error_expect}' not found in actual error: {str(cm.exception)}")	161	f"Expected error messages '{error_expect}' not found in actual error: {str(cm.exception)}")
162		162
	@@ -207,16 +207,16 @@ class OptionsTest(TestCase):
207	exceed_length = max_length + 100	207	exceed_length = max_length + 100
208	long_algo_str = "a" * exceed_length	208	long_algo_str = "a" * exceed_length
209	hccl_config = {"hccl_algo": long_algo_str}	209	hccl_config = {"hccl_algo": long_algo_str}
210	-	210	+
211	options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options()	211	options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options()
212	options.hccl_config = hccl_config	212	options.hccl_config = hccl_config
213	input1 = input1.npu()	213	input1 = input1.npu()
214	-	214	+
215	test_case = TestCase()	215	test_case = TestCase()
216	try:	216	try:
217	OptionsTest._init_dist_hccl(rank, options, world_size)	217	OptionsTest._init_dist_hccl(rank, options, world_size)
218	dist.all_reduce(input1)	218	dist.all_reduce(input1)
219	-	219	+
220	default_pg = c10d._get_default_group()._get_backend(torch.device('npu'))	220	default_pg = c10d._get_default_group()._get_backend(torch.device('npu'))
221	actual_algo = default_pg.options.hccl_config.get("hccl_algo", "")	221	actual_algo = default_pg.options.hccl_config.get("hccl_algo", "")
222	test_case.assertEqual(len(actual_algo), max_length - 1,	222	test_case.assertEqual(len(actual_algo), max_length - 1,

	@@ -38,7 +38,7 @@ class TestTorchairNoInit(TestCase):
38	for m in sys.modules:	38	for m in sys.modules:
39	if hasattr(sys.modules[m], '_attr_test_hasattr'):	39	if hasattr(sys.modules[m], '_attr_test_hasattr'):
40	setattr(sys.modules[m], '_attr_test_hasattr', 1)	40	setattr(sys.modules[m], '_attr_test_hasattr', 1)
41	-	41	+
42	torchair = sys.modules.get('torchair', None)	42	torchair = sys.modules.get('torchair', None)
43	self.assertTrue(torchair is not None)	43	self.assertTrue(torchair is not None)
44	self.assertTrue(not hasattr(torchair, '_attr_test_hasattr'))	44	self.assertTrue(not hasattr(torchair, '_attr_test_hasattr'))
	@@ -48,13 +48,13 @@ class TestTorchairNoInit(TestCase):
48	for m in sys.modules.values():	48	for m in sys.modules.values():
49	if getattr(m, '__warningregistry__', None):	49	if getattr(m, '__warningregistry__', None):
50	m.__warningregistry__ = {}	50	m.__warningregistry__ = {}
51	-	51	+
52	self._check_torchair_no_init()	52	self._check_torchair_no_init()
53	-	53	+
54	def test_attribute_error(self):	54	def test_attribute_error(self):
55	torchair = sys.modules.get('torchair', None)	55	torchair = sys.modules.get('torchair', None)
56	self.assertTrue(torchair is not None)	56	self.assertTrue(torchair is not None)
57	- with self.assertRaisesRegex(AttributeError,	57	+ with self.assertRaisesRegex(AttributeError,
58	"Try to get torchair's attr `get_npu_backend` before torchair is initialized."):	58	"Try to get torchair's attr `get_npu_backend` before torchair is initialized."):
59	torchair.get_npu_backend()	59	torchair.get_npu_backend()
60	self._check_torchair_no_init()	60	self._check_torchair_no_init()

	@@ -50,7 +50,7 @@ class TestIFAAclgraphUpdateSupportBlocking(TestCase):
50	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535, workspace=workspace,	50	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535, workspace=workspace,
51	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length, out=[output, softmax_lse])	51	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length, out=[output, softmax_lse])
52	handle = torch.npu.graph_task_group_end(stream)	52	handle = torch.npu.graph_task_group_end(stream)
53	-	53	+
54	with torch.npu.stream(update_stream):	54	with torch.npu.stream(update_stream):
55	torch.npu.graph_task_update_begin(update_stream, handle)	55	torch.npu.graph_task_update_begin(update_stream, handle)
56	torch_npu.npu_fused_infer_attention_score.out(	56	torch_npu.npu_fused_infer_attention_score.out(
	@@ -91,7 +91,7 @@ class TestIFAAclgraphUpdateSupportBlocking(TestCase):
91	torch_npu.npu_fused_infer_attention_score.out(	91	torch_npu.npu_fused_infer_attention_score.out(
92	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535, workspace=workspace,	92	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535, workspace=workspace,
93	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length, out=[output, softmax_lse])	93	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length, out=[output, softmax_lse])
94	-	94	+
95	g.update(cpu_update_input=[{"actual_seq_lengths": length_new}])	95	g.update(cpu_update_input=[{"actual_seq_lengths": length_new}])
96	g.replay()	96	g.replay()
97	self.assertEqual(output.cpu(), res_src[0].cpu())	97	self.assertEqual(output.cpu(), res_src[0].cpu())
	@@ -121,7 +121,7 @@ class TestIFAAclgraphUpdateSupportBlocking(TestCase):
121	output, softmax_lse = torch_npu.npu_fused_infer_attention_score(	121	output, softmax_lse = torch_npu.npu_fused_infer_attention_score(
122	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535,	122	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535,
123	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length)	123	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length)
124	-	124	+
125	g.update(cpu_update_input=[{"actual_seq_lengths": length_new}])	125	g.update(cpu_update_input=[{"actual_seq_lengths": length_new}])
126	g.replay()	126	g.replay()
127	self.assertEqual(output.cpu(), res_src[0].cpu())	127	self.assertEqual(output.cpu(), res_src[0].cpu())

	@@ -51,7 +51,7 @@ class TestIFAAclgraphUpdate(TestCase):
51	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535, workspace=workspace,	51	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535, workspace=workspace,
52	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length, out=[output, softmax_lse])	52	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length, out=[output, softmax_lse])
53	handle = torch.npu.graph_task_group_end(stream)	53	handle = torch.npu.graph_task_group_end(stream)
54	-	54	+
55	with torch.npu.stream(update_stream):	55	with torch.npu.stream(update_stream):
56	torch.npu.graph_task_update_begin(update_stream, handle)	56	torch.npu.graph_task_update_begin(update_stream, handle)
57	torch_npu.npu_fused_infer_attention_score.out(	57	torch_npu.npu_fused_infer_attention_score.out(
	@@ -100,7 +100,7 @@ class TestIFAAclgraphUpdate(TestCase):
100	torch_npu.npu_fused_infer_attention_score.out(	100	torch_npu.npu_fused_infer_attention_score.out(
101	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535, workspace=workspace,	101	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535, workspace=workspace,
102	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length, out=[output, softmax_lse])	102	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length, out=[output, softmax_lse])
103	-	103	+
104	g.update(cpu_update_input=[{"actual_seq_lengths": length_new}])	104	g.update(cpu_update_input=[{"actual_seq_lengths": length_new}])
105	g.replay()	105	g.replay()
106	self.assertEqual(output.cpu(), res_src[0].cpu())	106	self.assertEqual(output.cpu(), res_src[0].cpu())
	@@ -130,7 +130,7 @@ class TestIFAAclgraphUpdate(TestCase):
130	output, softmax_lse = torch_npu.npu_fused_infer_attention_score(	130	output, softmax_lse = torch_npu.npu_fused_infer_attention_score(
131	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535,	131	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535,
132	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length)	132	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length)
133	-	133	+
134	g.update(cpu_update_input=[{"actual_seq_lengths": length_new}])	134	g.update(cpu_update_input=[{"actual_seq_lengths": length_new}])
135	g.replay()	135	g.replay()
136	self.assertEqual(output.cpu(), res_src[0].cpu())	136	self.assertEqual(output.cpu(), res_src[0].cpu())
	@@ -164,7 +164,7 @@ class TestIFAAclgraphUpdate(TestCase):
164	torch_npu.npu_fused_infer_attention_score_v2.out(	164	torch_npu.npu_fused_infer_attention_score_v2.out(
165	query, key, value, num_query_heads=32, input_layout="BNSD", softmax_scale=scale, pre_tokens=65535, workspace=workspace,	165	query, key, value, num_query_heads=32, input_layout="BNSD", softmax_scale=scale, pre_tokens=65535, workspace=workspace,
166	next_tokens=65535, return_softmax_lse=False, actual_seq_qlen=length, out=[output, softmax_lse])	166	next_tokens=65535, return_softmax_lse=False, actual_seq_qlen=length, out=[output, softmax_lse])
167	-	167	+
168	g.update(cpu_update_input=[{"actual_seq_lengths": length_new}])	168	g.update(cpu_update_input=[{"actual_seq_lengths": length_new}])
169	g.replay()	169	g.replay()
170	self.assertEqual(output.cpu(), res_src[0].cpu())	170	self.assertEqual(output.cpu(), res_src[0].cpu())
	@@ -193,7 +193,7 @@ class TestIFAAclgraphUpdate(TestCase):
193	output, softmax_lse = torch_npu.npu_fused_infer_attention_score_v2(	193	output, softmax_lse = torch_npu.npu_fused_infer_attention_score_v2(
194	query, key, value, num_query_heads=32, input_layout="BNSD", softmax_scale=scale, pre_tokens=65535,	194	query, key, value, num_query_heads=32, input_layout="BNSD", softmax_scale=scale, pre_tokens=65535,
195	next_tokens=65535, return_softmax_lse=False, actual_seq_qlen=length)	195	next_tokens=65535, return_softmax_lse=False, actual_seq_qlen=length)
196	-	196	+
197	g.update(cpu_update_input=[{"actual_seq_qlen": length_new}])	197	g.update(cpu_update_input=[{"actual_seq_qlen": length_new}])
198	g.replay()	198	g.replay()
199	self.assertEqual(output.cpu(), res_src[0].cpu())	199	self.assertEqual(output.cpu(), res_src[0].cpu())
	@@ -233,7 +233,7 @@ class TestIFAAclgraphUpdate(TestCase):
233	query, key, value, num_query_heads=32, input_layout="BNSD", softmax_scale=scale, pre_tokens=65535, workspace=workspace,	233	query, key, value, num_query_heads=32, input_layout="BNSD", softmax_scale=scale, pre_tokens=65535, workspace=workspace,
234	next_tokens=65535, return_softmax_lse=False, actual_seq_qlen=length, out=[output, softmax_lse])	234	next_tokens=65535, return_softmax_lse=False, actual_seq_qlen=length, out=[output, softmax_lse])
235	handle = torch.npu.graph_task_group_end(stream)	235	handle = torch.npu.graph_task_group_end(stream)
236	-	236	+
237	with torch.npu.stream(update_stream):	237	with torch.npu.stream(update_stream):
238	torch.npu.graph_task_update_begin(update_stream, handle)	238	torch.npu.graph_task_update_begin(update_stream, handle)
239	torch_npu.npu_fused_infer_attention_score_v2.out(	239	torch_npu.npu_fused_infer_attention_score_v2.out(
	@@ -309,7 +309,7 @@ class TestIFAAclgraphUpdate(TestCase):
309	torch.nn.Dropout(p=0.2),	309	torch.nn.Dropout(p=0.2),
310	torch.nn.Linear(H, D_out),	310	torch.nn.Linear(H, D_out),
311	torch.nn.Dropout(p=0.1)).npu()	311	torch.nn.Dropout(p=0.1)).npu()
312	-	312	+
313	static_input = torch.randn(N, D_in, device='npu')	313	static_input = torch.randn(N, D_in, device='npu')
314	s = torch.npu.Stream()	314	s = torch.npu.Stream()
315	s.wait_stream(torch.npu.current_stream())	315	s.wait_stream(torch.npu.current_stream())
	@@ -365,7 +365,7 @@ class TestIFAAclgraphUpdate(TestCase):
365	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535, workspace=workspace,	365	query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535, workspace=workspace,
366	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length, out=[output, softmax_lse])	366	next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length, out=[output, softmax_lse])
367	handle = torch.npu.graph_task_group_end(stream)	367	handle = torch.npu.graph_task_group_end(stream)
368	-	368	+
369	with torch.npu.stream(update_stream):	369	with torch.npu.stream(update_stream):
370	torch.npu.graph_task_update_begin(update_stream, handle)	370	torch.npu.graph_task_update_begin(update_stream, handle)
371	torch_npu.npu_fused_infer_attention_score.out(	371	torch_npu.npu_fused_infer_attention_score.out(
	@@ -412,7 +412,7 @@ class TestIFAAclgraphUpdate(TestCase):
412	query, key, value, num_query_heads=32, input_layout="BNSD", softmax_scale=scale, pre_tokens=65535, workspace=workspace,	412	query, key, value, num_query_heads=32, input_layout="BNSD", softmax_scale=scale, pre_tokens=65535, workspace=workspace,
413	next_tokens=65535, return_softmax_lse=False, actual_seq_qlen=length, out=[output, softmax_lse])	413	next_tokens=65535, return_softmax_lse=False, actual_seq_qlen=length, out=[output, softmax_lse])
414	handle = torch.npu.graph_task_group_end(stream)	414	handle = torch.npu.graph_task_group_end(stream)
415	-	415	+
416	with torch.npu.stream(update_stream):	416	with torch.npu.stream(update_stream):
417	torch.npu.graph_task_update_begin(update_stream, handle)	417	torch.npu.graph_task_update_begin(update_stream, handle)
418	torch_npu.npu_fused_infer_attention_score_v2.out(	418	torch_npu.npu_fused_infer_attention_score_v2.out(
	@@ -503,8 +503,8 @@ class TestPAAclgraphUpdate(TestCase):
503		503
504	# 执行注意力计算	504	# 执行注意力计算
505	out = self.ref_masked_attention(	505	out = self.ref_masked_attention(
506	- params_np.query[i:i + 1],	506	+ params_np.query[i:i + 1],
507	- np.stack(keys),	507	+ np.stack(keys),
508	np.stack(values)	508	np.stack(values)
509	)	509	)
510	output[i] = out.reshape(self.num_heads, -1)	510	output[i] = out.reshape(self.num_heads, -1)
	@@ -536,11 +536,11 @@ class TestPAAclgraphUpdate(TestCase):
536		536
537	def atb_paged_attention(self, params):	537	def atb_paged_attention(self, params):
538	torch_npu._npu_paged_attention(	538	torch_npu._npu_paged_attention(
539	- query=params.query,	539	+ query=params.query,
540	- key_cache=params.key_cache,	540	+ key_cache=params.key_cache,
541	value_cache=params.value_cache,	541	value_cache=params.value_cache,
542	num_kv_heads=self.kv_heads,	542	num_kv_heads=self.kv_heads,
543	- num_heads=self.num_heads,	543	+ num_heads=self.num_heads,
544	scale_value=self.scale,	544	scale_value=self.scale,
545	block_table=params.block_table,	545	block_table=params.block_table,
546	context_lens=params.context_lens,	546	context_lens=params.context_lens,

	@@ -18,13 +18,13 @@ class TestDLPack(TestCase):
18	original = torch.randint(-10, 10, (2, 3, 4), dtype=dtype, device=device)	18	original = torch.randint(-10, 10, (2, 3, 4), dtype=dtype, device=device)
19	else:	19	else:
20	original = torch.randn(2, 3, 4, dtype=dtype, device=device)	20	original = torch.randn(2, 3, 4, dtype=dtype, device=device)
21	-	21	+
22	# Convert to dlpack	22	# Convert to dlpack
23	dlpack_tensor = to_dlpack(original)	23	dlpack_tensor = to_dlpack(original)
24	-	24	+
25	# Convert back to torch_npu tensor	25	# Convert back to torch_npu tensor
26	restored = from_dlpack(dlpack_tensor)	26	restored = from_dlpack(dlpack_tensor)
27	-	27	+
28	# Verify the roundtrip	28	# Verify the roundtrip
29	self.assertEqual(original, restored)	29	self.assertEqual(original, restored)
30	self.assertEqual(original.dtype, restored.dtype)	30	self.assertEqual(original.dtype, restored.dtype)
	@@ -42,13 +42,13 @@ class TestDLPack(TestCase):
42	(2, 2, 2, 2), # 4D tensor	42	(2, 2, 2, 2), # 4D tensor
43	(1, 1, 1, 1, 1) # 5D tensor	43	(1, 1, 1, 1, 1) # 5D tensor
44	]	44	]
45	-	45	+
46	for shape in shapes:	46	for shape in shapes:
47	with self.subTest(shape=shape):	47	with self.subTest(shape=shape):
48	original = torch.randn(shape, dtype=dtype, device=device)	48	original = torch.randn(shape, dtype=dtype, device=device)
49	dlpack_tensor = to_dlpack(original)	49	dlpack_tensor = to_dlpack(original)
50	restored = from_dlpack(dlpack_tensor)	50	restored = from_dlpack(dlpack_tensor)
51	-	51	+
52	self.assertEqual(original, restored)	52	self.assertEqual(original, restored)
53	self.assertEqual(original.shape, restored.shape)	53	self.assertEqual(original.shape, restored.shape)
54		54
	@@ -58,20 +58,20 @@ class TestDLPack(TestCase):
58	# Test contiguous tensor	58	# Test contiguous tensor
59	original_contiguous = torch.randn(4, 4, dtype=dtype, device=device)	59	original_contiguous = torch.randn(4, 4, dtype=dtype, device=device)
60	self.assertTrue(original_contiguous.is_contiguous())	60	self.assertTrue(original_contiguous.is_contiguous())
61	-	61	+
62	dlpack_tensor = to_dlpack(original_contiguous)	62	dlpack_tensor = to_dlpack(original_contiguous)
63	restored = from_dlpack(dlpack_tensor)	63	restored = from_dlpack(dlpack_tensor)
64	-	64	+
65	self.assertEqual(original_contiguous, restored)	65	self.assertEqual(original_contiguous, restored)
66	self.assertTrue(restored.is_contiguous())	66	self.assertTrue(restored.is_contiguous())
67	-	67	+
68	# Test non-contiguous tensor (transpose)	68	# Test non-contiguous tensor (transpose)
69	original_non_contiguous = original_contiguous.t()	69	original_non_contiguous = original_contiguous.t()
70	self.assertFalse(original_non_contiguous.is_contiguous())	70	self.assertFalse(original_non_contiguous.is_contiguous())
71	-	71	+
72	dlpack_tensor = to_dlpack(original_non_contiguous)	72	dlpack_tensor = to_dlpack(original_non_contiguous)
73	restored = from_dlpack(dlpack_tensor)	73	restored = from_dlpack(dlpack_tensor)
74	-	74	+
75	self.assertEqual(original_non_contiguous, restored)	75	self.assertEqual(original_non_contiguous, restored)
76	self.assertEqual(original_non_contiguous.stride(), restored.stride())	76	self.assertEqual(original_non_contiguous.stride(), restored.stride())
77		77
	@@ -80,14 +80,14 @@ class TestDLPack(TestCase):
80	"""Test that dlpack shares memory with original tensor"""	80	"""Test that dlpack shares memory with original tensor"""
81	original = torch.randn(3, 3, dtype=dtype, device=device)	81	original = torch.randn(3, 3, dtype=dtype, device=device)
82	original_data_ptr = original.data_ptr()	82	original_data_ptr = original.data_ptr()
83	-	83	+
84	# Convert to dlpack and back	84	# Convert to dlpack and back
85	dlpack_tensor = to_dlpack(original)	85	dlpack_tensor = to_dlpack(original)
86	restored = from_dlpack(dlpack_tensor)	86	restored = from_dlpack(dlpack_tensor)
87	-	87	+
88	# Check if memory is shared (data_ptr should be the same)	88	# Check if memory is shared (data_ptr should be the same)
89	self.assertEqual(original_data_ptr, restored.data_ptr())	89	self.assertEqual(original_data_ptr, restored.data_ptr())
90	-	90	+
91	# Modify original tensor and check if restored tensor is also modified	91	# Modify original tensor and check if restored tensor is also modified
92	original.fill_(42.0)	92	original.fill_(42.0)
93	self.assertEqual(original, restored)	93	self.assertEqual(original, restored)
	@@ -98,20 +98,20 @@ class TestDLPack(TestCase):
98	original = torch.randn(3, 4, dtype=dtype, device=device)	98	original = torch.randn(3, 4, dtype=dtype, device=device)
99	dlpack_tensor = to_dlpack(original)	99	dlpack_tensor = to_dlpack(original)
100	restored = from_dlpack(dlpack_tensor)	100	restored = from_dlpack(dlpack_tensor)
101	-	101	+
102	self.assertEqual(original, restored)	102	self.assertEqual(original, restored)
103	self.assertEqual(original.dtype, restored.dtype)	103	self.assertEqual(original.dtype, restored.dtype)
104	-	104	+
105	@Dtypes(torch.float)	105	@Dtypes(torch.float)
106	def test_dlpack_cpu(self, dtype, device="cpu"):	106	def test_dlpack_cpu(self, dtype, device="cpu"):
107	"""Test that dlpack shares memory with original cpu tensor"""	107	"""Test that dlpack shares memory with original cpu tensor"""
108	original = torch.randn(3, 3, dtype=dtype, device=device)	108	original = torch.randn(3, 3, dtype=dtype, device=device)
109	original_data_ptr = original.data_ptr()	109	original_data_ptr = original.data_ptr()
110	-	110	+
111	# Convert to dlpack and back	111	# Convert to dlpack and back
112	dlpack_tensor = to_dlpack(original)	112	dlpack_tensor = to_dlpack(original)
113	restored = from_dlpack(dlpack_tensor)	113	restored = from_dlpack(dlpack_tensor)
114	-	114	+
115	# Check if memory is shared (data_ptr should be the same)	115	# Check if memory is shared (data_ptr should be the same)
116	self.assertEqual(original_data_ptr, restored.data_ptr())	116	self.assertEqual(original_data_ptr, restored.data_ptr())
117		117

	@@ -694,15 +694,15 @@ class TestNPUGraphNodeRun(TestCase):
694		694
695		695
696	class TestGetNpugraphSegments(TestCase):	696	class TestGetNpugraphSegments(TestCase):
697	- @patch('torch.npu.memory_snapshot')	697	+ @patch('torch.npu.memory_snapshot')
698	- def test_get_npugraph_segments(self, mock_snapshot):	698	+ def test_get_npugraph_segments(self, mock_snapshot):
699	mock_snapshot.return_value = [	699	mock_snapshot.return_value = [
700	{"segment_pool_id": (0, 1), "address": 1000, "blocks": []},	700	{"segment_pool_id": (0, 1), "address": 1000, "blocks": []},
701	{"segment_pool_id": (0, 0), "address": 2000, "blocks": []},	701	{"segment_pool_id": (0, 0), "address": 2000, "blocks": []},
702	{"segment_pool_id": (0, 1), "address": 3000, "blocks": []},	702	{"segment_pool_id": (0, 1), "address": 3000, "blocks": []},
703	- ]	703	+ ]
704	- result = get_npugraph_segments((0, 1))	704	+ result = get_npugraph_segments((0, 1))
705	- self.assertEqual(len(result), 2)	705	+ self.assertEqual(len(result), 2)
706	mock_snapshot.assert_called_once_with()	706	mock_snapshot.assert_called_once_with()
707		707
708		708
	@@ -919,15 +919,15 @@ class TestNPUGraphTreeManager:
919	manager.npu_graphs_thread_pool = "pool_handle"	919	manager.npu_graphs_thread_pool = "pool_handle"
920	manager.device_index = 0	920	manager.device_index = 0
921	manager.stream = MagicMock()	921	manager.stream = MagicMock()
922	-	922	+
923	# 设置模拟返回值	923	# 设置模拟返回值
924	mock_node_instance = MagicMock()	924	mock_node_instance = MagicMock()
925	mock_node.return_value = mock_node_instance	925	mock_node.return_value = mock_node_instance
926	mock_node_instance.run_first_inputs.return_value = [torch.tensor([1.0])]	926	mock_node_instance.run_first_inputs.return_value = [torch.tensor([1.0])]
927	-	927	+
928	# 执行测试	928	# 执行测试
929	result = manager.record_function([torch.tensor([1.0])], FunctionID(1))	929	result = manager.record_function([torch.tensor([1.0])], FunctionID(1))
930	-	930	+
931	# 验证调用	931	# 验证调用
932	mock_synchronize.assert_any_call()	932	mock_synchronize.assert_any_call()
933	mock_node.assert_called_once_with(	933	mock_node.assert_called_once_with(
	@@ -950,10 +950,10 @@ class TestNPUGraphTreeManager:
950	manager = NPUGraphTreeManager(0)	950	manager = NPUGraphTreeManager(0)
951	mock_node = MagicMock()	951	mock_node = MagicMock()
952	mock_node.run.return_value = [torch.tensor([1.0])]	952	mock_node.run.return_value = [torch.tensor([1.0])]
953	-	953	+
954	# 执行测试	954	# 执行测试
955	result = manager.execute_node(mock_node, [torch.tensor([1.0])])	955	result = manager.execute_node(mock_node, [torch.tensor([1.0])])
956	-	956	+
957	# 验证调用	957	# 验证调用
958	mock_update_gen.assert_called_once_with()	958	mock_update_gen.assert_called_once_with()
959	assert manager.current_node == mock_node	959	assert manager.current_node == mock_node
	@@ -970,15 +970,15 @@ class TestNPUGraphTreeManager:
970	manager.graph = MagicMock()	970	manager.graph = MagicMock()
971	manager.device_index = 0	971	manager.device_index = 0
972	manager.stream = MagicMock()	972	manager.stream = MagicMock()
973	-	973	+
974	# 设置模拟返回值	974	# 设置模拟返回值
975	mock_node_instance = MagicMock()	975	mock_node_instance = MagicMock()
976	mock_warmup_node.return_value = mock_node_instance	976	mock_warmup_node.return_value = mock_node_instance
977	mock_node_instance.run.return_value = [torch.tensor([1.0])]	977	mock_node_instance.run.return_value = [torch.tensor([1.0])]
978	-	978	+
979	# 执行测试	979	# 执行测试
980	result = manager.run_eager([torch.tensor([1.0])], FunctionID(1))	980	result = manager.run_eager([torch.tensor([1.0])], FunctionID(1))
981	-	981	+
982	# 验证调用	982	# 验证调用
983	mock_update_gen.assert_called_once_with()	983	mock_update_gen.assert_called_once_with()
984	mock_warmup_node.assert_called_once_with(	984	mock_warmup_node.assert_called_once_with(
	@@ -1163,7 +1163,7 @@ class TestNPUGraphTreeManager:
1163	):	1163	):
1164	manager = NPUGraphTreeManager(0)	1164	manager = NPUGraphTreeManager(0)
1165	mock_in_new_invocation.return_value = True	1165	mock_in_new_invocation.return_value = True
1166	-	1166	+
1167	mock_node = MagicMock()	1167	mock_node = MagicMock()
1168	mock_node._path_from_root = [MagicMock()]	1168	mock_node._path_from_root = [MagicMock()]
1169	mock_node._path_from_root[0].wrapped_function.id = FunctionID(2)	1169	mock_node._path_from_root[0].wrapped_function.id = FunctionID(2)
	@@ -1179,22 +1179,22 @@ class TestNPUGraphTreeManager:
1179	):	1179	):
1180	manager = NPUGraphTreeManager(0)	1180	manager = NPUGraphTreeManager(0)
1181	mock_in_new_invocation.return_value = True	1181	mock_in_new_invocation.return_value = True
1182	-	1182	+
1183	mock_node1 = MagicMock()	1183	mock_node1 = MagicMock()
1184	mock_node1.wrapped_function.id = FunctionID(1)	1184	mock_node1.wrapped_function.id = FunctionID(1)
1185	mock_node1.parent = MagicMock()	1185	mock_node1.parent = MagicMock()
1186	mock_node1.parent.wrapped_function.id = FunctionID(0)	1186	mock_node1.parent.wrapped_function.id = FunctionID(0)
1187	-	1187	+
1188	mock_node2 = MagicMock()	1188	mock_node2 = MagicMock()
1189	mock_node2.wrapped_function.id = FunctionID(1)	1189	mock_node2.wrapped_function.id = FunctionID(1)
1190	mock_node2.parent = MagicMock()	1190	mock_node2.parent = MagicMock()
1191	mock_node2.parent.wrapped_function.id = FunctionID(0)	1191	mock_node2.parent.wrapped_function.id = FunctionID(0)
1192	-	1192	+
1193	mock_current_node = MagicMock()	1193	mock_current_node = MagicMock()
1194	mock_current_node.wrapped_function.id = FunctionID(1)	1194	mock_current_node.wrapped_function.id = FunctionID(1)
1195	mock_current_node.parent = MagicMock()	1195	mock_current_node.parent = MagicMock()
1196	mock_current_node.parent.wrapped_function.id = FunctionID(0)	1196	mock_current_node.parent.wrapped_function.id = FunctionID(0)
1197	-	1197	+
1198	mock_current_node._path_from_root = [mock_node1, mock_node2]	1198	mock_current_node._path_from_root = [mock_node1, mock_node2]
1199	manager.current_node = mock_current_node	1199	manager.current_node = mock_current_node
1200	manager.check_warn_on_unable_to_start_executing(FunctionID(1))	1200	manager.check_warn_on_unable_to_start_executing(FunctionID(1))

	@@ -95,14 +95,14 @@ class TorchNPUDeviceTestCase(TestCase):
95	os.environ["TORCH_NPU_DEVICE_CAPABILITY"] = "9.0"	95	os.environ["TORCH_NPU_DEVICE_CAPABILITY"] = "9.0"
96	res = torch_npu.npu.get_device_capability()	96	res = torch_npu.npu.get_device_capability()
97	self.assertEqual(res, (9, 0))	97	self.assertEqual(res, (9, 0))
98	-	98	+
99	os.environ["TORCH_NPU_DEVICE_CAPABILITY"] = "8.0"	99	os.environ["TORCH_NPU_DEVICE_CAPABILITY"] = "8.0"
100	res = torch_npu.npu.get_device_capability(0)	100	res = torch_npu.npu.get_device_capability(0)
101	self.assertEqual(res, (8, 0))	101	self.assertEqual(res, (8, 0))
102	device = torch_npu.npu.device("npu")	102	device = torch_npu.npu.device("npu")
103	res = torch_npu.npu.get_device_capability(device)	103	res = torch_npu.npu.get_device_capability(device)
104	self.assertEqual(res, (8, 0))	104	self.assertEqual(res, (8, 0))
105	-	105	+
106	os.environ["TORCH_NPU_DEVICE_CAPABILITY"] = "8.a"	106	os.environ["TORCH_NPU_DEVICE_CAPABILITY"] = "8.a"
107	res = torch_npu.npu.get_device_capability()	107	res = torch_npu.npu.get_device_capability()
108	self.assertEqual(res, None)	108	self.assertEqual(res, None)
	@@ -118,7 +118,7 @@ class TorchNPUDeviceTestCase(TestCase):
118	torch_npu.npu.synchronize()	118	torch_npu.npu.synchronize()
119	after_free_memory, after_total_memory = torch_npu.npu.mem_get_info(0)	119	after_free_memory, after_total_memory = torch_npu.npu.mem_get_info(0)
120	self.assertEqual(before_total_memory, after_total_memory)	120	self.assertEqual(before_total_memory, after_total_memory)
121	-	121	+
122	@unittest.skip("CANN doesn't support now.")	122	@unittest.skip("CANN doesn't support now.")
123	def test_set_device_res_limit(self):	123	def test_set_device_res_limit(self):
124	ans_dict = {'cube_core_num': 12, 'vector_core_num': 24}	124	ans_dict = {'cube_core_num': 12, 'vector_core_num': 24}
	@@ -301,7 +301,7 @@ class TorchNPUApiTestCase(TestCase):
301	end_event.record()	301	end_event.record()
302	res = start_event.elapsed_time(end_event)	302	res = start_event.elapsed_time(end_event)
303	self.assertIsInstance(res, float)	303	self.assertIsInstance(res, float)
304	-	304	+
305	def test_npu_event_recorded_time(self):	305	def test_npu_event_recorded_time(self):
306	event_1 = torch_npu.npu.Event(enable_timing=True)	306	event_1 = torch_npu.npu.Event(enable_timing=True)
307	event_1.record()	307	event_1.record()
	@@ -426,7 +426,7 @@ print(f"{{r1}}, {{r2}}")
426	self.fail("Expected exception not raised for negative timeout value")	426	self.fail("Expected exception not raised for negative timeout value")
427	except Exception as e:	427	except Exception as e:
428	self.assertIn("can't convert negative value to unsigned int", str(e), f"{e}")	428	self.assertIn("can't convert negative value to unsigned int", str(e), f"{e}")
429	-	429	+
430	try:	430	try:
431	torch_npu.npu.set_op_timeout_ms(2**32)	431	torch_npu.npu.set_op_timeout_ms(2**32)
432	self.fail("Expected exception not raised for large timeout value")	432	self.fail("Expected exception not raised for large timeout value")

	@@ -1604,7 +1604,7 @@ class TestBinaryUfuncs(TestCase):
1604	for tensor in tensors:	1604	for tensor in tensors:
1605	self._test_pow(base, tensor)	1605	self._test_pow(base, tensor)
1606		1606
1607	-	1607	+
1608	def test_cuda_tensor_pow_scalar_tensor(self, device):	1608	def test_cuda_tensor_pow_scalar_tensor(self, device):
1609	cuda_tensors = [	1609	cuda_tensors = [
1610	torch.randn((3, 3), device=device),	1610	torch.randn((3, 3), device=device),
	@@ -1618,7 +1618,7 @@ class TestBinaryUfuncs(TestCase):
1618	for base, exp in product(cuda_tensors, scalar_tensors):	1618	for base, exp in product(cuda_tensors, scalar_tensors):
1619	self._test_pow(base, exp)	1619	self._test_pow(base, exp)
1620		1620
1621	-	1621	+
1622	def test_cpu_tensor_pow_cuda_scalar_tensor(self, device):	1622	def test_cpu_tensor_pow_cuda_scalar_tensor(self, device):
1623	cuda_tensors = [	1623	cuda_tensors = [
1624	torch.tensor(5.0, device="privatuse1"),	1624	torch.tensor(5.0, device="privatuse1"),
	@@ -1633,7 +1633,7 @@ class TestBinaryUfuncs(TestCase):
1633	base = torch.tensor(3.0, device="cpu")	1633	base = torch.tensor(3.0, device="cpu")
1634	self._test_pow(base, exp)	1634	self._test_pow(base, exp)
1635		1635
1636	-	1636	+
1637	@dtypes(torch.complex64, torch.complex128)	1637	@dtypes(torch.complex64, torch.complex128)
1638	def test_pow_cuda_complex_extremal_failing(self, device, dtype):	1638	def test_pow_cuda_complex_extremal_failing(self, device, dtype):
1639	t = torch.tensor(complex(-1.0, float("inf")), dtype=dtype, device=device)	1639	t = torch.tensor(complex(-1.0, float("inf")), dtype=dtype, device=device)
	@@ -1642,7 +1642,7 @@ class TestBinaryUfuncs(TestCase):
1642	cpu_out = t.cpu().pow(2)	1642	cpu_out = t.cpu().pow(2)
1643	self.assertEqual(cpu_out, cuda_out)	1643	self.assertEqual(cpu_out, cuda_out)
1644		1644
1645	-	1645	+
1646	@skipIfTorchDynamo()	1646	@skipIfTorchDynamo()
1647	@dtypes(*all_types_and_complex_and(torch.half))	1647	@dtypes(*all_types_and_complex_and(torch.half))
1648	def test_complex_scalar_pow_tensor(self, device, dtype):	1648	def test_complex_scalar_pow_tensor(self, device, dtype):
	@@ -1669,7 +1669,7 @@ class TestBinaryUfuncs(TestCase):
1669	self._test_pow(base, first_exp)	1669	self._test_pow(base, first_exp)
1670	self._test_pow(base, second_exp)	1670	self._test_pow(base, second_exp)
1671		1671
1672	-	1672	+
1673	@skipMeta	1673	@skipMeta
1674	def test_pow_scalar_type_promotion(self, device):	1674	def test_pow_scalar_type_promotion(self, device):
1675	# Test against a scalar and non-scalar input	1675	# Test against a scalar and non-scalar input
	@@ -1837,7 +1837,7 @@ class TestBinaryUfuncs(TestCase):
1837	_scalar_helper(lambda a, b: math.floor(a / b), operator.floordiv)	1837	_scalar_helper(lambda a, b: math.floor(a / b), operator.floordiv)
1838	_scalar_helper(lambda a, b: math.floor(a / b), torch.floor_divide)	1838	_scalar_helper(lambda a, b: math.floor(a / b), torch.floor_divide)
1839		1839
1840	-	1840	+
1841	@skipIfTorchDynamo("Not a suitable test for TorchDynamo")	1841	@skipIfTorchDynamo("Not a suitable test for TorchDynamo")
1842	def test_div_and_floordiv_script_vs_python(self, device):	1842	def test_div_and_floordiv_script_vs_python(self, device):
1843	# Creates jitted functions of two tensors	1843	# Creates jitted functions of two tensors
	@@ -1908,7 +1908,7 @@ class TestBinaryUfuncs(TestCase):
1908	# See issue gh-52387	1908	# See issue gh-52387
1909	self.assertEqual(5 // a, scripted_rfloordiv_scalar(a_t))	1909	self.assertEqual(5 // a, scripted_rfloordiv_scalar(a_t))
1910		1910
1911	-	1911	+
1912	@skipIfTorchDynamo("Not a suitable test for TorchDynamo")	1912	@skipIfTorchDynamo("Not a suitable test for TorchDynamo")
1913	def test_idiv_and_ifloordiv_vs_python(self, device):	1913	def test_idiv_and_ifloordiv_vs_python(self, device):
1914	def _wrapped_idiv_tensor(a, b):	1914	def _wrapped_idiv_tensor(a, b):
	@@ -2239,7 +2239,7 @@ class TestBinaryUfuncs(TestCase):
2239	torch.ones(1, device=device, dtype=dtypes[0]),	2239	torch.ones(1, device=device, dtype=dtypes[0]),
2240	)	2240	)
2241		2241
2242	-	2242	+
2243	def test_maximum_minimum_cross_device(self, device):	2243	def test_maximum_minimum_cross_device(self, device):
2244	a = torch.tensor((1, 2, -1))	2244	a = torch.tensor((1, 2, -1))
2245	b = torch.tensor((3, 0, 4), device=device)	2245	b = torch.tensor((3, 0, 4), device=device)
	@@ -2822,7 +2822,7 @@ class TestBinaryUfuncs(TestCase):
2822	expected = np.hypot(input[0].cpu().numpy(), input[1].cpu().numpy())	2822	expected = np.hypot(input[0].cpu().numpy(), input[1].cpu().numpy())
2823	self.assertEqual(actual, expected, exact_dtype=False)	2823	self.assertEqual(actual, expected, exact_dtype=False)
2824		2824
2825	-	2825	+
2826	@dtypes(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64)	2826	@dtypes(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64)
2827	def test_gcd(self, device, dtype):	2827	def test_gcd(self, device, dtype):
2828	# Tests gcd(0, 0), gcd(0, a) cases	2828	# Tests gcd(0, 0), gcd(0, a) cases
	@@ -2847,7 +2847,7 @@ class TestBinaryUfuncs(TestCase):
2847	expected = np.gcd(a.cpu().numpy(), b.cpu().numpy())	2847	expected = np.gcd(a.cpu().numpy(), b.cpu().numpy())
2848	self.assertEqual(actual, expected)	2848	self.assertEqual(actual, expected)
2849		2849
2850	-	2850	+
2851	@dtypes(torch.int16, torch.int32, torch.int64)	2851	@dtypes(torch.int16, torch.int32, torch.int64)
2852	def test_lcm(self, device, dtype):	2852	def test_lcm(self, device, dtype):
2853	# Tests lcm(0, 0), lcm(0, a) cases	2853	# Tests lcm(0, 0), lcm(0, a) cases
	@@ -2864,7 +2864,7 @@ class TestBinaryUfuncs(TestCase):
2864	expected = np.lcm(a.cpu().numpy(), b.cpu().numpy())	2864	expected = np.lcm(a.cpu().numpy(), b.cpu().numpy())
2865	self.assertEqual(actual, expected, exact_dtype=False)	2865	self.assertEqual(actual, expected, exact_dtype=False)
2866		2866
2867	-	2867	+
2868	@dtypes(torch.float32, torch.float64)	2868	@dtypes(torch.float32, torch.float64)
2869	def test_nextafter(self, device, dtype):	2869	def test_nextafter(self, device, dtype):
2870	# Test special cases	2870	# Test special cases
	@@ -2888,7 +2888,7 @@ class TestBinaryUfuncs(TestCase):
2888	expected = np.nextafter(a.cpu().numpy(), b.cpu().numpy())	2888	expected = np.nextafter(a.cpu().numpy(), b.cpu().numpy())
2889	self.assertEqual(actual, expected, atol=0, rtol=0)	2889	self.assertEqual(actual, expected, atol=0, rtol=0)
2890		2890
2891	-	2891	+
2892	@dtypes(torch.bfloat16)	2892	@dtypes(torch.bfloat16)
2893	def test_nextafter_bfloat16(self, device, dtype):	2893	def test_nextafter_bfloat16(self, device, dtype):
2894	nan = float("nan")	2894	nan = float("nan")
	@@ -3108,7 +3108,7 @@ class TestBinaryUfuncs(TestCase):
3108	) # all casts to complex128 are safe	3108	) # all casts to complex128 are safe
3109	compare_with_numpy_bin_op(torch_op, numpy_op, a, b, out=out)	3109	compare_with_numpy_bin_op(torch_op, numpy_op, a, b, out=out)
3110		3110
3111	-	3111	+
3112	@dtypes(torch.int8, torch.int16, torch.int32, torch.int64)	3112	@dtypes(torch.int8, torch.int16, torch.int32, torch.int64)
3113	def test_signed_shift(self, device, dtype):	3113	def test_signed_shift(self, device, dtype):
3114	"Ensure that signed integer bit shifting works as expected."	3114	"Ensure that signed integer bit shifting works as expected."
	@@ -3124,7 +3124,7 @@ class TestBinaryUfuncs(TestCase):
3124	self.assertEqual(a >> 1, expected_r)	3124	self.assertEqual(a >> 1, expected_r)
3125	self.compare_with_numpy(lambda x: x >> 1, lambda x: np.right_shift(x, 1), a)	3125	self.compare_with_numpy(lambda x: x >> 1, lambda x: np.right_shift(x, 1), a)
3126		3126
3127	-	3127	+
3128	@dtypes(*get_all_int_dtypes())	3128	@dtypes(*get_all_int_dtypes())
3129	def test_shift_limits(self, device, dtype):	3129	def test_shift_limits(self, device, dtype):
3130	"Ensure that integer bit shifting works as expected with out-of-limits shift values."	3130	"Ensure that integer bit shifting works as expected with out-of-limits shift values."
	@@ -3167,7 +3167,7 @@ class TestBinaryUfuncs(TestCase):
3167	exact_dtype=exact_dtype, msg=f">> {shift}"	3167	exact_dtype=exact_dtype, msg=f">> {shift}"
3168	)	3168	)
3169		3169
3170	-	3170	+
3171	@dtypes(	3171	@dtypes(
3172	*list(	3172	*list(
3173	product(	3173	product(
	@@ -3236,7 +3236,7 @@ class TestBinaryUfuncs(TestCase):
3236	):	3236	):
3237	input.heaviside_(values)	3237	input.heaviside_(values)
3238		3238
3239	-	3239	+
3240	def test_heaviside_cross_device(self, device):	3240	def test_heaviside_cross_device(self, device):
3241	x = torch.tensor([-9, 5, 0, 6, -2, 2], device=device)	3241	x = torch.tensor([-9, 5, 0, 6, -2, 2], device=device)
3242	y = torch.tensor(0)	3242	y = torch.tensor(0)
	@@ -3405,7 +3405,7 @@ class TestBinaryUfuncs(TestCase):
3405	expected = start + weight * (end - start)	3405	expected = start + weight * (end - start)
3406	self.assertEqual(expected, actual)	3406	self.assertEqual(expected, actual)
3407		3407
3408	-	3408	+
3409	@dtypes(torch.half, torch.bfloat16)	3409	@dtypes(torch.half, torch.bfloat16)
3410	def test_lerp_lowp(self, device, dtype):	3410	def test_lerp_lowp(self, device, dtype):
3411	xvals = (0.0, -30000.0)	3411	xvals = (0.0, -30000.0)
	@@ -3646,7 +3646,7 @@ class TestBinaryUfuncs(TestCase):
3646	lambda: torch.add(m1, m1, out=m2),	3646	lambda: torch.add(m1, m1, out=m2),
3647	)	3647	)
3648		3648
3649	-	3649	+
3650	def test_addsub_half_tensor(self, device):	3650	def test_addsub_half_tensor(self, device):
3651	x = torch.tensor([60000.0], dtype=torch.half, device=device)	3651	x = torch.tensor([60000.0], dtype=torch.half, device=device)
3652	for op, y, alpha in (	3652	for op, y, alpha in (
	@@ -3879,7 +3879,7 @@ class TestBinaryUfuncs(TestCase):
3879	def test_cumulative_trapezoid(self, device):	3879	def test_cumulative_trapezoid(self, device):
3880		3880
3881	import scipy.integrate	3881	import scipy.integrate
3882	-	3882	+
3883	if hasattr(scipy.integrate, "cumulative_trapezoid"):	3883	if hasattr(scipy.integrate, "cumulative_trapezoid"):
3884	_scipy_cumulative_trapezoid = scipy.integrate.cumulative_trapezoid	3884	_scipy_cumulative_trapezoid = scipy.integrate.cumulative_trapezoid
3885	else: # Older version of SciPy uses a different name	3885	else: # Older version of SciPy uses a different name
	@@ -4352,7 +4352,7 @@ class TestBinaryUfuncs(TestCase):
4352	x = make_tensor((2, 3, 4), dtype=x_dtype, device=device)	4352	x = make_tensor((2, 3, 4), dtype=x_dtype, device=device)
4353	test_helper(x, q)	4353	test_helper(x, q)
4354		4354
4355	-	4355	+
4356	@dtypes(	4356	@dtypes(
4357	torch.chalf,	4357	torch.chalf,
4358	)	4358	)

	@@ -50,22 +50,22 @@ class TestMultiprocessingAPIs(TestCase):
50	mp.set_start_method(method, force=True)	50	mp.set_start_method(method, force=True)
51	current_method = mp.get_start_method()	51	current_method = mp.get_start_method()
52	self.assertEqual(current_method, method)	52	self.assertEqual(current_method, method)
53	-	53	+
54	# Verify that the child process uses the correct start method	54	# Verify that the child process uses the correct start method
55	queue = mp.SimpleQueue()	55	queue = mp.SimpleQueue()
56	process = mp.Process(target=_worker, args=(queue,))	56	process = mp.Process(target=_worker, args=(queue,))
57	process.start()	57	process.start()
58	process.join()	58	process.join()
59	-	59	+
60	# Get the start method from the child process	60	# Get the start method from the child process
61	self.assertFalse(queue.empty(), "Queue should contain the start method")	61	self.assertFalse(queue.empty(), "Queue should contain the start method")
62	child_method = queue.get()	62	child_method = queue.get()
63	- self.assertEqual(child_method, method,	63	+ self.assertEqual(child_method, method,
64	f"Child process should use {method} start method")	64	f"Child process should use {method} start method")
65	-	65	+
66	# Verify that the main process context has not changed	66	# Verify that the main process context has not changed
67	self.assertEqual(mp.get_start_method(), method)	67	self.assertEqual(mp.get_start_method(), method)
68	-	68	+
69		69
70	def test_value(self):	70	def test_value(self):
71	"""Test Value API"""	71	"""Test Value API"""
	@@ -184,49 +184,49 @@ class TestMultiprocessingAPIs(TestCase):
184	"""Test torch.multiprocessing.reductions.init_reductions and reduce_tensor APIs"""	184	"""Test torch.multiprocessing.reductions.init_reductions and reduce_tensor APIs"""
185	# Test init_reductions - verify it doesn't raise any exception	185	# Test init_reductions - verify it doesn't raise any exception
186	mp.reductions.init_reductions()	186	mp.reductions.init_reductions()
187	-	187	+
188	# Test reduce_tensor directly for CPU tensor	188	# Test reduce_tensor directly for CPU tensor
189	# Create a simple CPU tensor	189	# Create a simple CPU tensor
190	tensor = torch.tensor([1, 2, 3, 4])	190	tensor = torch.tensor([1, 2, 3, 4])
191	reduced = mp.reductions.reduce_tensor(tensor)	191	reduced = mp.reductions.reduce_tensor(tensor)
192	-	192	+
193	# Verify the reduced form is a tuple with expected structure	193	# Verify the reduced form is a tuple with expected structure
194	self.assertIsInstance(reduced, tuple)	194	self.assertIsInstance(reduced, tuple)
195	self.assertEqual(len(reduced), 2)	195	self.assertEqual(len(reduced), 2)
196	-	196	+
197	# Try to reconstruct the tensor	197	# Try to reconstruct the tensor
198	constructor, args = reduced	198	constructor, args = reduced
199	reconstructed = constructor(*args)	199	reconstructed = constructor(*args)
200	-	200	+
201	# Verify reconstruction worked	201	# Verify reconstruction worked
202	self.assertTrue(torch.equal(tensor, reconstructed))	202	self.assertTrue(torch.equal(tensor, reconstructed))
203	self.assertEqual(tensor.device, reconstructed.device)	203	self.assertEqual(tensor.device, reconstructed.device)
204	self.assertEqual(tensor.dtype, reconstructed.dtype)	204	self.assertEqual(tensor.dtype, reconstructed.dtype)
205	-	205	+
206	# Test with NPU tensor if available	206	# Test with NPU tensor if available
207	if torch.npu.is_available():	207	if torch.npu.is_available():
208	# Create a simple NPU tensor	208	# Create a simple NPU tensor
209	npu_tensor = torch.tensor([1, 2, 3, 4], device='npu:0')	209	npu_tensor = torch.tensor([1, 2, 3, 4], device='npu:0')
210	-	210	+
211	# Test reduce_tensor for NPU tensor	211	# Test reduce_tensor for NPU tensor
212	reduced_npu = mp.reductions.reduce_tensor(npu_tensor)	212	reduced_npu = mp.reductions.reduce_tensor(npu_tensor)
213	self.assertIsInstance(reduced_npu, tuple)	213	self.assertIsInstance(reduced_npu, tuple)
214	self.assertEqual(len(reduced_npu), 2)	214	self.assertEqual(len(reduced_npu), 2)
215	-	215	+
216	# Verify reconstruction for NPU tensor	216	# Verify reconstruction for NPU tensor
217	constructor_npu, args_npu = reduced_npu	217	constructor_npu, args_npu = reduced_npu
218	reconstructed_npu = constructor_npu(*args_npu)	218	reconstructed_npu = constructor_npu(*args_npu)
219	self.assertTrue(torch.equal(npu_tensor.cpu(), reconstructed_npu.cpu()))	219	self.assertTrue(torch.equal(npu_tensor.cpu(), reconstructed_npu.cpu()))
220	-	220	+
221	def test_reductions_invalid_input(self):	221	def test_reductions_invalid_input(self):
222	"""Test reduction APIs with invalid inputs"""	222	"""Test reduction APIs with invalid inputs"""
223	# Test reduce_tensor with invalid input	223	# Test reduce_tensor with invalid input
224	with self.assertRaises(Exception):	224	with self.assertRaises(Exception):
225	mp.reductions.reduce_tensor(None)	225	mp.reductions.reduce_tensor(None)
226	-	226	+
227	with self.assertRaises(Exception):	227	with self.assertRaises(Exception):
228	mp.reductions.reduce_tensor("not a tensor")	228	mp.reductions.reduce_tensor("not a tensor")
229	-	229	+
230	# Test init_reductions multiple times (should be safe)	230	# Test init_reductions multiple times (should be safe)
231	mp.reductions.init_reductions()	231	mp.reductions.init_reductions()
232	mp.reductions.init_reductions()	232	mp.reductions.init_reductions()

	@@ -84,13 +84,13 @@ class TestNestedTensor(TestCase):
84	self.assertEqual(len(nt_as_list), len(nt_list))	84	self.assertEqual(len(nt_as_list), len(nt_list))
85	self.assertEqual(nt_as_list[0], nt_list[0])	85	self.assertEqual(nt_as_list[0], nt_list[0])
86	self.assertEqual(nt_as_list[1], nt_list[1])	86	self.assertEqual(nt_as_list[1], nt_list[1])
87	-	87	+
88	def test_unbind_and_asnested_int64(self):	88	def test_unbind_and_asnested_int64(self):
89	a = torch.tensor([[1, 2, 3], [4, 5, 6]])	89	a = torch.tensor([[1, 2, 3], [4, 5, 6]])
90	b = torch.tensor([[7, 8], [10, 11]])	90	b = torch.tensor([[7, 8], [10, 11]])
91	self._test_unbind_case(a, b)	91	self._test_unbind_case(a, b)
92	self._test_asnested_case(a, b)	92	self._test_asnested_case(a, b)
93	-	93	+
94	def test_unbind_and_asnested_float32(self):	94	def test_unbind_and_asnested_float32(self):
95	a = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)	95	a = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)
96	b = torch.tensor([[7, 8], [10, 11]], dtype=torch.float32)	96	b = torch.tensor([[7, 8], [10, 11]], dtype=torch.float32)
	@@ -101,7 +101,7 @@ class TestNestedTensor(TestCase):
101	a = torch.tensor([[], []])	101	a = torch.tensor([[], []])
102	b = torch.tensor([[], [], []])	102	b = torch.tensor([[], [], []])
103	self._test_unbind_case(a, b)	103	self._test_unbind_case(a, b)
104	- self._test_asnested_case(a, b)	104	+ self._test_asnested_case(a, b)
105		105
106	def test_default_options_nested_tensor(self):	106	def test_default_options_nested_tensor(self):
107	default_nested_tensor = torch.nested.nested_tensor([], device="npu:0")	107	default_nested_tensor = torch.nested.nested_tensor([], device="npu:0")
	@@ -111,12 +111,12 @@ class TestNestedTensor(TestCase):
111	self.assertEqual(default_nested_tensor.layout, default_tensor.layout)	111	self.assertEqual(default_nested_tensor.layout, default_tensor.layout)
112	self.assertEqual(default_nested_tensor.dim(), default_tensor.dim())	112	self.assertEqual(default_nested_tensor.dim(), default_tensor.dim())
113	self.assertEqual(default_nested_tensor.requires_grad, default_tensor.requires_grad)	113	self.assertEqual(default_nested_tensor.requires_grad, default_tensor.requires_grad)
114	-	114	+
115	def test_nested_tensor_errsize(self):	115	def test_nested_tensor_errsize(self):
116	nt = torch.nested.nested_tensor([torch.tensor([[1, 2, 3], [4, 5, 6]]).npu(), torch.tensor([[7, 8], [10, 11], [12, 13]]).npu()])	116	nt = torch.nested.nested_tensor([torch.tensor([[1, 2, 3], [4, 5, 6]]).npu(), torch.tensor([[7, 8], [10, 11], [12, 13]]).npu()])
117	self.assertEqual(nt.size(0), 2)	117	self.assertEqual(nt.size(0), 2)
118	self.assertRaisesRegex(RuntimeError,	118	self.assertRaisesRegex(RuntimeError,
119	- "Given dimension 1 is irregular and does not have a size",	119	+ "Given dimension 1 is irregular and does not have a size",
120	lambda: nt.size(1),	120	lambda: nt.size(1),
121	)	121	)
122		122

	@@ -993,7 +993,7 @@ class TestReductions(TestCase):
993	# Check whether the returned values are the mode	993	# Check whether the returned values are the mode
994	self.assertTrue((values == v).all().item())	994	self.assertTrue((values == v).all().item())
995		995
996	-	996	+
997	@dtypes(*all_types_and(torch.half, torch.bfloat16))	997	@dtypes(*all_types_and(torch.half, torch.bfloat16))
998	def test_mode_large(self, device, dtype):	998	def test_mode_large(self, device, dtype):
999	# i should be less than (d - 2) / 2	999	# i should be less than (d - 2) / 2
	@@ -1063,7 +1063,7 @@ class TestReductions(TestCase):
1063	test_for_dtypes(torch.int32, torch.int32, torch.float32, indices_err)	1063	test_for_dtypes(torch.int32, torch.int32, torch.float32, indices_err)
1064	test_for_dtypes(torch.float32, torch.float32, torch.float64, indices_err)	1064	test_for_dtypes(torch.float32, torch.float32, torch.float64, indices_err)
1065		1065
1066	-	1066	+
1067	def test_mode_wrong_device(self, device):	1067	def test_mode_wrong_device(self, device):
1068	# CPU Input Tensor	1068	# CPU Input Tensor
1069	x = torch.ones(2)	1069	x = torch.ones(2)
	@@ -1212,7 +1212,7 @@ class TestReductions(TestCase):
1212	@dtypes(torch.float, torch.double)	1212	@dtypes(torch.float, torch.double)
1213	def test_amax(self, device, dtype):	1213	def test_amax(self, device, dtype):
1214	self._test_minmax_helper(torch.amax, np.amax, device, dtype)	1214	self._test_minmax_helper(torch.amax, np.amax, device, dtype)
1215	-	1215	+
1216	@dtypes(torch.float, torch.double)	1216	@dtypes(torch.float, torch.double)
1217	@dtypesIfPRIVATEUSE1(torch.half, torch.float, torch.bfloat16)	1217	@dtypesIfPRIVATEUSE1(torch.half, torch.float, torch.bfloat16)
1218	def test_aminmax(self, device, dtype):	1218	def test_aminmax(self, device, dtype):
	@@ -1463,7 +1463,7 @@ class TestReductions(TestCase):
1463	torch.sum(x, (2, 1), out=res2)	1463	torch.sum(x, (2, 1), out=res2)
1464	self.assertEqual(res1, res2)	1464	self.assertEqual(res1, res2)
1465		1465
1466	-	1466	+
1467	@dtypes(torch.float16, torch.float32)	1467	@dtypes(torch.float16, torch.float32)
1468	def test_prod_gpu(self, device, dtype):	1468	def test_prod_gpu(self, device, dtype):
1469	x = torch.tensor([2, 3, 6, 9, 8], dtype=dtype, device=device)	1469	x = torch.tensor([2, 3, 6, 9, 8], dtype=dtype, device=device)
	@@ -1729,7 +1729,7 @@ class TestReductions(TestCase):
1729	# So we must skip this as well.	1729	# So we must skip this as well.
1730	if dtype == torch.uint8:	1730	if dtype == torch.uint8:
1731	exact_dtype = False	1731	exact_dtype = False
1732	-	1732	+
1733	# Investigate why the output is not close to numpy.	1733	# Investigate why the output is not close to numpy.
1734	atol, rtol = self._get_relaxed_tolerances_for(dtype)	1734	atol, rtol = self._get_relaxed_tolerances_for(dtype)
1735		1735
	@@ -1743,7 +1743,7 @@ class TestReductions(TestCase):
1743	self._test_sum_reduction_vs_numpy(torch.sum, np.sum, device, dtype, with_extremal=True)	1743	self._test_sum_reduction_vs_numpy(torch.sum, np.sum, device, dtype, with_extremal=True)
1744	self._test_sum_reduction_vs_numpy(torch.sum, np.sum, device, dtype, with_keepdim=True)	1744	self._test_sum_reduction_vs_numpy(torch.sum, np.sum, device, dtype, with_keepdim=True)
1745		1745
1746	-	1746	+
1747	@dtypes(*set(all_types_and(torch.half)) - {torch.uint8})	1747	@dtypes(*set(all_types_and(torch.half)) - {torch.uint8})
1748	def test_nansum_vs_numpy(self, device, dtype):	1748	def test_nansum_vs_numpy(self, device, dtype):
1749	self._test_sum_reduction_vs_numpy(torch.nansum, np.nansum, device, dtype)	1749	self._test_sum_reduction_vs_numpy(torch.nansum, np.nansum, device, dtype)
	@@ -1977,7 +1977,7 @@ class TestReductions(TestCase):
1977	op(x, dim=dim)	1977	op(x, dim=dim)
1978		1978
1979	# update this test to comapre against NumPy	1979	# update this test to comapre against NumPy
1980	-	1980	+
1981	def test_var(self, device):	1981	def test_var(self, device):
1982	cpu_tensor = torch.randn(2, 3, 3)	1982	cpu_tensor = torch.randn(2, 3, 3)
1983	device_tensor = cpu_tensor.to(device)	1983	device_tensor = cpu_tensor.to(device)
	@@ -1992,7 +1992,7 @@ class TestReductions(TestCase):
1992	device_tensor = cpu_tensor.to(device)	1992	device_tensor = cpu_tensor.to(device)
1993	self.assertEqual(device_tensor.var(), cpu_tensor.var())	1993	self.assertEqual(device_tensor.var(), cpu_tensor.var())
1994		1994
1995	- # update this test to compare against NumPy	1995	+ # update this test to compare against NumPy
1996	def test_var_large_input(self, device):	1996	def test_var_large_input(self, device):
1997	# Large, not-nice input	1997	# Large, not-nice input
1998	cpu_tensor = torch.randn(2 * 32 * 1024 + 1, 2, 67)	1998	cpu_tensor = torch.randn(2 * 32 * 1024 + 1, 2, 67)
	@@ -2001,7 +2001,7 @@ class TestReductions(TestCase):
2001	self.assertEqual(cpu_tensor.var(2), device_tensor.var(2))	2001	self.assertEqual(cpu_tensor.var(2), device_tensor.var(2))
2002		2002
2003	# update this to compare against NumPy instead of CPU	2003	# update this to compare against NumPy instead of CPU
2004	-	2004	+
2005	@dtypes(torch.double)	2005	@dtypes(torch.double)
2006	def test_sum_noncontig(self, device, dtype):	2006	def test_sum_noncontig(self, device, dtype):
2007	x = torch.randn(1, 75, 57, 20, dtype=dtype, device=device).permute(0, 3, 1, 2)	2007	x = torch.randn(1, 75, 57, 20, dtype=dtype, device=device).permute(0, 3, 1, 2)
	@@ -2046,7 +2046,7 @@ class TestReductions(TestCase):
2046	torch.sum(x, dim=[0], dtype=torch.float32, out=y)	2046	torch.sum(x, dim=[0], dtype=torch.float32, out=y)
2047		2047
2048	# Assert for illegal dtype would not be raised on XLA	2048	# Assert for illegal dtype would not be raised on XLA
2049	-	2049	+
2050	def test_minmax_illegal_dtype(self, device):	2050	def test_minmax_illegal_dtype(self, device):
2051	x = torch.randn(5, 5, dtype=torch.float32, device=device)	2051	x = torch.randn(5, 5, dtype=torch.float32, device=device)
2052	valid_values = torch.empty(5, dtype=torch.float32, device=device)	2052	valid_values = torch.empty(5, dtype=torch.float32, device=device)
	@@ -2266,7 +2266,7 @@ class TestReductions(TestCase):
2266	expected = fn(y, 1, keepdim=False)	2266	expected = fn(y, 1, keepdim=False)
2267	self.assertEqual(x[:, 1], expected, msg=f'{fn_name} with out= kwarg')	2267	self.assertEqual(x[:, 1], expected, msg=f'{fn_name} with out= kwarg')
2268		2268
2269	-	2269	+
2270	@largeTensorTest('10GB')	2270	@largeTensorTest('10GB')
2271	def test_reduction_split(self, device):	2271	def test_reduction_split(self, device):
2272	# Test reduction when there is a 32bit-indexing split	2272	# Test reduction when there is a 32bit-indexing split
	@@ -2275,7 +2275,7 @@ class TestReductions(TestCase):
2275	expect = input_[0] + input_[1] + input_[2] + input_[3] + input_[4]	2275	expect = input_[0] + input_[1] + input_[2] + input_[3] + input_[4]
2276	self.assertEqual(result, expect)	2276	self.assertEqual(result, expect)
2277		2277
2278	-	2278	+
2279	@dtypes(torch.half, torch.float, torch.double, torch.bfloat16)	2279	@dtypes(torch.half, torch.float, torch.double, torch.bfloat16)
2280	def test_reduction_vectorize_along_input_corner(self, device, dtype):	2280	def test_reduction_vectorize_along_input_corner(self, device, dtype):
2281	# 1D case: sum	2281	# 1D case: sum
	@@ -2373,7 +2373,7 @@ class TestReductions(TestCase):
2373	self.assertEqual(xs1[j].item(), size[1] - i)	2373	self.assertEqual(xs1[j].item(), size[1] - i)
2374	self.assertEqual(xs2[j].item(), size[1] - i)	2374	self.assertEqual(xs2[j].item(), size[1] - i)
2375		2375
2376	-	2376	+
2377	@dtypes(torch.half, torch.float, torch.double, torch.bfloat16)	2377	@dtypes(torch.half, torch.float, torch.double, torch.bfloat16)
2378	def test_reduction_vectorize_along_output(self, device, dtype):	2378	def test_reduction_vectorize_along_output(self, device, dtype):
2379	def run_test(input_):	2379	def run_test(input_):
	@@ -2397,7 +2397,7 @@ class TestReductions(TestCase):
2397	run_test(torch.zeros(64, 61, dtype=dtype, device=device))	2397	run_test(torch.zeros(64, 61, dtype=dtype, device=device))
2398	run_test(torch.zeros(64, 1, dtype=dtype, device=device))	2398	run_test(torch.zeros(64, 1, dtype=dtype, device=device))
2399		2399
2400	-	2400	+
2401	def test_argminmax_large_axis(self, device):	2401	def test_argminmax_large_axis(self, device):
2402	# Regression test for gh-32863	2402	# Regression test for gh-32863
2403	x = torch.zeros(2**31, device=device, dtype=torch.int8)	2403	x = torch.zeros(2**31, device=device, dtype=torch.int8)
	@@ -2541,7 +2541,7 @@ class TestReductions(TestCase):
2541	self.assertEqual(a[:, ::2, :].nanmedian(-1)[0], torch.tensor([[0, 4], [6, 10]], device=device))	2541	self.assertEqual(a[:, ::2, :].nanmedian(-1)[0], torch.tensor([[0, 4], [6, 10]], device=device))
2542		2542
2543		2543
2544	-	2544	+
2545	@dtypes(torch.float, torch.double)	2545	@dtypes(torch.float, torch.double)
2546	def test_quantile(self, device, dtype):	2546	def test_quantile(self, device, dtype):
2547	# Generate some random test cases	2547	# Generate some random test cases
	@@ -3571,7 +3571,7 @@ as the input tensor excluding its innermost dimension'):
3571		3571
3572	self.assertEqual(actual, expected, msg, exact_dtype=exact_dtype)	3572	self.assertEqual(actual, expected, msg, exact_dtype=exact_dtype)
3573		3573
3574	-	3574	+
3575	@largeTensorTest("8GB")	3575	@largeTensorTest("8GB")
3576	@dtypes(torch.half, torch.chalf, torch.bfloat16)	3576	@dtypes(torch.half, torch.chalf, torch.bfloat16)
3577	def test_reductions_large_half_tensors(self, device, dtype):	3577	def test_reductions_large_half_tensors(self, device, dtype):

	@@ -70,7 +70,7 @@ class TestTensorToPreserveFormat(TestCase):
70	"""H2D: transpose 2D non-contiguous -> preserve_format preserves stride"""	70	"""H2D: transpose 2D non-contiguous -> preserve_format preserves stride"""
71	cpu_t = torch.randn(4, 6).t()	71	cpu_t = torch.randn(4, 6).t()
72	self.assertFalse(cpu_t.is_contiguous())	72	self.assertFalse(cpu_t.is_contiguous())
73	-	73	+
74		74
75	npu_t = cpu_t.to("npu", memory_format=torch.preserve_format)	75	npu_t = cpu_t.to("npu", memory_format=torch.preserve_format)
76	self._verify_preserve_format_result(cpu_t, npu_t, "H2D-transpose_2d")	76	self._verify_preserve_format_result(cpu_t, npu_t, "H2D-transpose_2d")
	@@ -96,7 +96,7 @@ class TestTensorToPreserveFormat(TestCase):
96	"""H2D: slice dim0 non-contiguous -> preserve_format preserves stride"""	96	"""H2D: slice dim0 non-contiguous -> preserve_format preserves stride"""
97	cpu_t = torch.randn(8, 5)[::2]	97	cpu_t = torch.randn(8, 5)[::2]
98	self.assertFalse(cpu_t.is_contiguous())	98	self.assertFalse(cpu_t.is_contiguous())
99	-	99	+
100		100
101	npu_t = cpu_t.to("npu", memory_format=torch.preserve_format)	101	npu_t = cpu_t.to("npu", memory_format=torch.preserve_format)
102	expec_strides = self._get_dense_strides(cpu_t)	102	expec_strides = self._get_dense_strides(cpu_t)
	@@ -107,7 +107,7 @@ class TestTensorToPreserveFormat(TestCase):
107	"""H2D: slice dim1 non-contiguous -> preserve_format preserves stride"""	107	"""H2D: slice dim1 non-contiguous -> preserve_format preserves stride"""
108	cpu_t = torch.randn(4, 8)[:, ::2]	108	cpu_t = torch.randn(4, 8)[:, ::2]
109	self.assertFalse(cpu_t.is_contiguous())	109	self.assertFalse(cpu_t.is_contiguous())
110	-	110	+
111		111
112	npu_t = cpu_t.to("npu", memory_format=torch.preserve_format)	112	npu_t = cpu_t.to("npu", memory_format=torch.preserve_format)
113	self._verify_preserve_format_result(cpu_t, npu_t, "H2D-slice_dim1")	113	self._verify_preserve_format_result(cpu_t, npu_t, "H2D-slice_dim1")
	@@ -119,14 +119,14 @@ class TestTensorToPreserveFormat(TestCase):
119	"""H2D: narrow non-contiguous -> preserve_format preserves stride"""	119	"""H2D: narrow non-contiguous -> preserve_format preserves stride"""
120	cpu_t = torch.randn(6, 8).narrow(1, 1, 5)	120	cpu_t = torch.randn(6, 8).narrow(1, 1, 5)
121	self.assertFalse(cpu_t.is_contiguous())	121	self.assertFalse(cpu_t.is_contiguous())
122	-	122	+
123		123
124	npu_t = cpu_t.to("npu", memory_format=torch.preserve_format)	124	npu_t = cpu_t.to("npu", memory_format=torch.preserve_format)
125	self._verify_preserve_format_result(cpu_t, npu_t, "H2D-narrow")	125	self._verify_preserve_format_result(cpu_t, npu_t, "H2D-narrow")
126	expec_strides = self._get_dense_strides(cpu_t)	126	expec_strides = self._get_dense_strides(cpu_t)
127	self.assertEqual(expec_strides, npu_t.stride(),	127	self.assertEqual(expec_strides, npu_t.stride(),
128	"H2D-narrow: stride not preserved")	128	"H2D-narrow: stride not preserved")
129	-	129	+
130	def test_h2d_select(self):	130	def test_h2d_select(self):
131	"""H2D: select (reduces dim) non-contiguous -> preserve_format preserves stride"""	131	"""H2D: select (reduces dim) non-contiguous -> preserve_format preserves stride"""
132	cpu_t = torch.randn(3, 5, 4).select(1, 2)	132	cpu_t = torch.randn(3, 5, 4).select(1, 2)
	@@ -164,7 +164,7 @@ class TestTensorToPreserveFormat(TestCase):
164	"""H2D: expand 3D (stride=0) -> preserve_format falls back to suggest_memory_format"""	164	"""H2D: expand 3D (stride=0) -> preserve_format falls back to suggest_memory_format"""
165	cpu_t = torch.randn(2, 1, 4).expand(2, 3, 4)	165	cpu_t = torch.randn(2, 1, 4).expand(2, 3, 4)
166	self.assertFalse(cpu_t.is_contiguous())	166	self.assertFalse(cpu_t.is_contiguous())
167	-	167	+
168		168
169	npu_t = cpu_t.to("npu", memory_format=torch.preserve_format)	169	npu_t = cpu_t.to("npu", memory_format=torch.preserve_format)
170	self._verify_preserve_format_result(cpu_t, npu_t, "H2D-expand_3d")	170	self._verify_preserve_format_result(cpu_t, npu_t, "H2D-expand_3d")
	@@ -174,7 +174,7 @@ class TestTensorToPreserveFormat(TestCase):
174	def test_h2d_as_strided_overlap(self):	174	def test_h2d_as_strided_overlap(self):
175	"""H2D: as_strided with overlap -> preserve_format falls back to suggest_memory_format"""	175	"""H2D: as_strided with overlap -> preserve_format falls back to suggest_memory_format"""
176	cpu_t = torch.randn(12).as_strided((3, 3), (4, 1))	176	cpu_t = torch.randn(12).as_strided((3, 3), (4, 1))
177	-	177	+
178		178
179	npu_t = cpu_t.to("npu", memory_format=torch.preserve_format)	179	npu_t = cpu_t.to("npu", memory_format=torch.preserve_format)
180	self._verify_preserve_format_result(cpu_t, npu_t, "H2D-as_strided_overlap")	180	self._verify_preserve_format_result(cpu_t, npu_t, "H2D-as_strided_overlap")

	@@ -104,18 +104,18 @@ static void _launch(void* func, void* tiling_func, int64_t tiling_size, void* ar
104	// only 1D parallelization is supported for NPU	104	// only 1D parallelization is supported for NPU
105	// Pointer type becomes flattend 1-D Memref tuple: base_ptr, data_ptr, offset, shape, stride	105	// Pointer type becomes flattend 1-D Memref tuple: base_ptr, data_ptr, offset, shape, stride
106	// base_ptr offset shape and stride are not used, arbitrarily set for now	106	// base_ptr offset shape and stride are not used, arbitrarily set for now
107	-	107	+
108	if (tiling_size == 0) {{	108	if (tiling_size == 0) {{
109	auto launch_call = [func, tiling_func, tiling_size, arg_tiling_host, arg_tiling_device, gridX, stream, {', '.join(f"arg{i}" + ("" if "torch." in ty else f", arg_allocate{i}, offset{i}" +(', ' if ranks[i] > 0 else '') + ', '.join(f"sizes{i}_{rank}" for rank in range(ranks[i])) + (', ' if ranks[i] > 0 else '') + ', '.join(f"strides{i}_{rank}" for rank in range(ranks[i]))) for i, ty in signature.items())}]() {{	109	auto launch_call = [func, tiling_func, tiling_size, arg_tiling_host, arg_tiling_device, gridX, stream, {', '.join(f"arg{i}" + ("" if "torch." in ty else f", arg_allocate{i}, offset{i}" +(', ' if ranks[i] > 0 else '') + ', '.join(f"sizes{i}_{rank}" for rank in range(ranks[i])) + (', ' if ranks[i] > 0 else '') + ', '.join(f"strides{i}_{rank}" for rank in range(ranks[i]))) for i, ty in signature.items())}]() {{
110	struct __attribute__((packed)) {{	110	struct __attribute__((packed)) {{
111	-	111	+
112	{' '.join(f'{_ty_to_cpp(ty)} arg{i} __attribute__((aligned({4 if ty[0] != "" and ty[-2:] != "64" else 8}))); ' + ('' if "torch." in ty else f'{_ty_to_cpp(ty)} arg_allocate{i} __attribute__((aligned({4 if ty[0] != "" and ty[-2:] != "64" else 8}))); {_ty_to_cpp(ty)} offset{i} __attribute__((aligned(8))); ' + ' '.join(f'{_ty_to_cpp(ty)} sizes{i}_{rank} __attribute__((aligned(8)));' for rank in range(ranks[i])) + ' ' + ' '.join(f'{_ty_to_cpp(ty)} strides{i}_{rank} __attribute__((aligned(8)));' for rank in range(ranks[i]))) for i, ty in signature.items())}	112	{' '.join(f'{_ty_to_cpp(ty)} arg{i} __attribute__((aligned({4 if ty[0] != "" and ty[-2:] != "64" else 8}))); ' + ('' if "torch." in ty else f'{_ty_to_cpp(ty)} arg_allocate{i} __attribute__((aligned({4 if ty[0] != "" and ty[-2:] != "64" else 8}))); {_ty_to_cpp(ty)} offset{i} __attribute__((aligned(8))); ' + ' '.join(f'{_ty_to_cpp(ty)} sizes{i}_{rank} __attribute__((aligned(8)));' for rank in range(ranks[i])) + ' ' + ' '.join(f'{_ty_to_cpp(ty)} strides{i}_{rank} __attribute__((aligned(8)));' for rank in range(ranks[i]))) for i, ty in signature.items())}
113		113
114	}} args = {{	114	}} args = {{
115	{', '.join(f"static_cast<{_ty_to_cpp(ty)}>(arg{i})" + ("" if "torch." in ty else f", static_cast<{_ty_to_cpp(ty)}>(arg_allocate{i}), static_cast<{_ty_to_cpp(ty)}>(offset{i})"+ (', ' if ranks[i] > 0 else '') + ', '.join(f"static_cast<{_ty_to_cpp(ty)}>(sizes{i}_{rank})" for rank in range(ranks[i])) + (', ' if ranks[i] > 0 else '') + ', '.join(f"static_cast<{_ty_to_cpp(ty)}>(strides{i}_{rank})" for rank in range(ranks[i]))) for i, ty in signature.items())}	115	{', '.join(f"static_cast<{_ty_to_cpp(ty)}>(arg{i})" + ("" if "torch." in ty else f", static_cast<{_ty_to_cpp(ty)}>(arg_allocate{i}), static_cast<{_ty_to_cpp(ty)}>(offset{i})"+ (', ' if ranks[i] > 0 else '') + ', '.join(f"static_cast<{_ty_to_cpp(ty)}>(sizes{i}_{rank})" for rank in range(ranks[i])) + (', ' if ranks[i] > 0 else '') + ', '.join(f"static_cast<{_ty_to_cpp(ty)}>(strides{i}_{rank})" for rank in range(ranks[i]))) for i, ty in signature.items())}
116		116
117	}};	117	}};
118	-	118	+
119	rtError_t ret = common_launch_dyn(const_cast<char>("{kernel_name}"), func, tiling_func, tiling_size, arg_tiling_host, arg_tiling_device, gridX, static_cast<void >(&args), sizeof(args), stream);	119	rtError_t ret = common_launch_dyn(const_cast<char>("{kernel_name}"), func, tiling_func, tiling_size, arg_tiling_host, arg_tiling_device, gridX, static_cast<void >(&args), sizeof(args), stream);
120	return ret;	120	return ret;
121	}};	121	}};
	@@ -128,7 +128,7 @@ static void _launch(void* func, void* tiling_func, int64_t tiling_size, void* ar
128	void* strides_tiling = (void*)1;	128	void* strides_tiling = (void*)1;
129	auto launch_call = [func, tiling_func, tiling_size, arg_tiling_host, arg_tiling_device, gridX, stream, {', '.join(f"arg{i}" + ("" if "torch." in ty else f", arg_allocate{i}, offset{i}" + (', ' if ranks[i] > 0 else '') + ', '.join(f"sizes{i}_{rank}" for rank in range(ranks[i])) + (', ' if ranks[i] > 0 else '') + ', '.join(f"strides{i}_{rank}" for rank in range(ranks[i]))) for i, ty in signature.items())}, key_tiling, offset_tiling, sizes_tiling, strides_tiling]() {{	129	auto launch_call = [func, tiling_func, tiling_size, arg_tiling_host, arg_tiling_device, gridX, stream, {', '.join(f"arg{i}" + ("" if "torch." in ty else f", arg_allocate{i}, offset{i}" + (', ' if ranks[i] > 0 else '') + ', '.join(f"sizes{i}_{rank}" for rank in range(ranks[i])) + (', ' if ranks[i] > 0 else '') + ', '.join(f"strides{i}_{rank}" for rank in range(ranks[i]))) for i, ty in signature.items())}, key_tiling, offset_tiling, sizes_tiling, strides_tiling]() {{
130	struct __attribute__((packed)) {{	130	struct __attribute__((packed)) {{
131	-	131	+
132	{' '.join(f'{_ty_to_cpp(ty)} arg{i} __attribute__((aligned({4 if ty[0] != "" and ty[-2:] != "64" else 8}))); ' + ('' if "torch." in ty else f'{_ty_to_cpp(ty)} arg_allocate{i} __attribute__((aligned({4 if ty[0] != "" and ty[-2:] != "64" else 8}))); {_ty_to_cpp(ty)} offset{i} __attribute__((aligned(8))); ' + ' '.join(f'{_ty_to_cpp(ty)} sizes{i}_{rank} __attribute__((aligned(8)));' for rank in range(ranks[i])) + ' ' + ' '.join(f'{_ty_to_cpp(ty)} strides{i}_{rank} __attribute__((aligned(8)));' for rank in range(ranks[i]))) for i, ty in signature.items())}	132	{' '.join(f'{_ty_to_cpp(ty)} arg{i} __attribute__((aligned({4 if ty[0] != "" and ty[-2:] != "64" else 8}))); ' + ('' if "torch." in ty else f'{_ty_to_cpp(ty)} arg_allocate{i} __attribute__((aligned({4 if ty[0] != "" and ty[-2:] != "64" else 8}))); {_ty_to_cpp(ty)} offset{i} __attribute__((aligned(8))); ' + ' '.join(f'{_ty_to_cpp(ty)} sizes{i}_{rank} __attribute__((aligned(8)));' for rank in range(ranks[i])) + ' ' + ' '.join(f'{_ty_to_cpp(ty)} strides{i}_{rank} __attribute__((aligned(8)));' for rank in range(ranks[i]))) for i, ty in signature.items())}
133		133
134	void* key_tiling __attribute__((aligned(8)));	134	void* key_tiling __attribute__((aligned(8)));
	@@ -143,7 +143,7 @@ static void _launch(void* func, void* tiling_func, int64_t tiling_size, void* ar
143		143
144	(void)(&key_tiling), arg_tiling_host, arg_tiling_device, static_cast<void>(offset_tiling), static_cast<void>(sizes_tiling), static_cast<void>(strides_tiling)	144	(void)(&key_tiling), arg_tiling_host, arg_tiling_device, static_cast<void>(offset_tiling), static_cast<void>(sizes_tiling), static_cast<void>(strides_tiling)
145	}};	145	}};
146	-	146	+
147	rtError_t ret = common_launch_dyn(const_cast<char>("{kernel_name}"), func, tiling_func, tiling_size, arg_tiling_host, arg_tiling_device, gridX, static_cast<void >(&args), sizeof(args), stream);	147	rtError_t ret = common_launch_dyn(const_cast<char>("{kernel_name}"), func, tiling_func, tiling_size, arg_tiling_host, arg_tiling_device, gridX, static_cast<void >(&args), sizeof(args), stream);
148	return ret;	148	return ret;
149	}};	149	}};

	@@ -59,7 +59,7 @@ id_iter = count()
59		59
60		60
61	class NpuTritonKernel(TritonKernel):	61	class NpuTritonKernel(TritonKernel):
62	- def __init__(self,	62	+ def __init__(self,
63	tiling: Dict[str, sympy.Expr],	63	tiling: Dict[str, sympy.Expr],
64	min_elem_per_thread=0,	64	min_elem_per_thread=0,
65	optimize_mask=True,	65	optimize_mask=True,
	@@ -77,7 +77,7 @@ class NpuTritonKernel(TritonKernel):
77	@staticmethod	77	@staticmethod
78	def inductor_meta_common():	78	def inductor_meta_common():
79	return {}	79	return {}
80	-	80	+
81	def call_kernel(self, call_args, name: str):	81	def call_kernel(self, call_args, name: str):
82	wrapper = V.graph.wrapper_code	82	wrapper = V.graph.wrapper_code
83	for call_arg in call_args:	83	for call_arg in call_args:
	@@ -170,14 +170,14 @@ def create_fx_from_snodes_by_traced_graph(snodes: List[scheduler.SchedulerNode],
170		170
171	def runnable_gm(*args):	171	def runnable_gm(*args):
172	return torch.fx.Interpreter(gm).run(*args)	172	return torch.fx.Interpreter(gm).run(*args)
173	- with V.graph.fake_mode:	173	+ with V.graph.fake_mode:
174	gm = make_fx(runnable_gm)(*inputs)	174	gm = make_fx(runnable_gm)(*inputs)
175	view_to_reshape(gm)	175	view_to_reshape(gm)
176	- non_contiguous_indices["outputs"] = [i + num_inputs	176	+ non_contiguous_indices["outputs"] = [i + num_inputs
177	for i, call_output in enumerate(call_outputs)	177	for i, call_output in enumerate(call_outputs)
178	if not V.graph.try_get_buffer(call_output).layout.is_contiguous()]	178	if not V.graph.try_get_buffer(call_output).layout.is_contiguous()]
179	- return (gm, call_args, {"num_outputs": num_outputs,	179	+ return (gm, call_args, {"num_outputs": num_outputs,
180	- "non_contiguous_indices": non_contiguous_indices,	180	+ "non_contiguous_indices": non_contiguous_indices,
181	"mutated_indices": mutated_indices, })	181	"mutated_indices": mutated_indices, })
182		182
183		183
	@@ -192,10 +192,10 @@ class NpuMetaKernel(Kernel):
192	self._gm = gm	192	self._gm = gm
193	self._gm_with_prim_cast = self.build_gm_with_prim_cast(gm)	193	self._gm_with_prim_cast = self.build_gm_with_prim_cast(gm)
194	self._is_dynamic = is_fx_dynamic(self._gm)	194	self._is_dynamic = is_fx_dynamic(self._gm)
195	-	195	+
196	if anir_config.online_acc_comp:	196	if anir_config.online_acc_comp:
197	modify_gm_for_acc_comp(self._gm)	197	modify_gm_for_acc_comp(self._gm)
198	-	198	+
199	self._snodes = snodes	199	self._snodes = snodes
200	self._call_args = call_args	200	self._call_args = call_args
201	self.non_contiguous_indices = non_contiguous_indices	201	self.non_contiguous_indices = non_contiguous_indices
	@@ -216,7 +216,7 @@ class NpuMetaKernel(Kernel):
216	V.graph.device_ops.import_get_raw_stream_as("get_raw_stream")	216	V.graph.device_ops.import_get_raw_stream_as("get_raw_stream")
217	)	217	)
218	)	218	)
219	-	219	+
220	def build_gm_with_prim_cast(self, gm):	220	def build_gm_with_prim_cast(self, gm):
221	return npu_cast_to_prim_cast(gm)	221	return npu_cast_to_prim_cast(gm)
222		222
	@@ -254,12 +254,12 @@ class NpuMetaKernel(Kernel):
254	def call_kernel(self, name: str, node=None):	254	def call_kernel(self, name: str, node=None):
255	wrapper = V.graph.wrapper_code	255	wrapper = V.graph.wrapper_code
256	call_args = self.get_call_args()	256	call_args = self.get_call_args()
257	-	257	+
258	for call_arg in call_args:	258	for call_arg in call_args:
259	if call_arg.startswith('_uwu'):	259	if call_arg.startswith('_uwu'):
260	expression = map_strings_to_operators(call_arg)	260	expression = map_strings_to_operators(call_arg)
261	wrapper.writeline(f'{call_arg} = {expression}')	261	wrapper.writeline(f'{call_arg} = {expression}')
262	-	262	+
263	if len(call_args) > 0:	263	if len(call_args) > 0:
264	wrapper.generate_kernel_call(name, call_args)	264	wrapper.generate_kernel_call(name, call_args)
265		265
	@@ -296,7 +296,7 @@ class NpuMetaScheduling(SIMDScheduling):
296	def define_kernel(self, src_code, mlir_kernel, traced_graph, mode=None):	296	def define_kernel(self, src_code, mlir_kernel, traced_graph, mode=None):
297	if mode is None:	297	if mode is None:
298	mode = anir_config._get_compile_mode()	298	mode = anir_config._get_compile_mode()
299	-	299	+
300	wrapper = V.graph.wrapper_code	300	wrapper = V.graph.wrapper_code
301		301
302	kernel_key = (src_code, tuple(mlir_kernel.non_contiguous_indices))	302	kernel_key = (src_code, tuple(mlir_kernel.non_contiguous_indices))
	@@ -338,9 +338,9 @@ class NpuMetaScheduling(SIMDScheduling):
338	kernel_meta.update(extra_kernel_meta)	338	kernel_meta.update(extra_kernel_meta)
339		339
340	wrapper.src_to_kernel[kernel_key] = kernel_name	340	wrapper.src_to_kernel[kernel_key] = kernel_name
341	-	341	+
342	subs_name = kernel_name if config.triton.unique_kernel_names else f"{self._get_kernel_prefix()}_"	342	subs_name = kernel_name if config.triton.unique_kernel_names else f"{self._get_kernel_prefix()}_"
343	-	343	+
344	compile_wrapper = IndentedBuffer()	344	compile_wrapper = IndentedBuffer()
345	metadata_comment = ""	345	metadata_comment = ""
346		346
	@@ -390,7 +390,7 @@ class NpuMetaScheduling(SIMDScheduling):
390		390
391	def _handle_auto_fallback_mode(self, compile_wrapper, src_code, name, subs_name, meta, wrapper, metadata_comment, mlir_kernel=None):	391	def _handle_auto_fallback_mode(self, compile_wrapper, src_code, name, subs_name, meta, wrapper, metadata_comment, mlir_kernel=None):
392	_basename, _, kernel_path = get_path(code_hash(src_code.strip()), "py")	392	_basename, _, kernel_path = get_path(code_hash(src_code.strip()), "py")
393	-	393	+
394	compile_wrapper.writeline(f"async_compile.{self._get_compile_api()}({subs_name!r}, '''")	394	compile_wrapper.writeline(f"async_compile.{self._get_compile_api()}({subs_name!r}, '''")
395	compile_wrapper.splice(src_code, strip=True)	395	compile_wrapper.splice(src_code, strip=True)
396	compile_wrapper.writeline(f"''', kernel_meta={meta})")	396	compile_wrapper.writeline(f"''', kernel_meta={meta})")
	@@ -399,7 +399,7 @@ class NpuMetaScheduling(SIMDScheduling):
399		399
400	origins, detailed_origins = get_kernel_metadata(mlir_kernel._snodes, wrapper)	400	origins, detailed_origins = get_kernel_metadata(mlir_kernel._snodes, wrapper)
401	metadata_comment += "\n" + origins + "\n" + detailed_origins	401	metadata_comment += "\n" + origins + "\n" + detailed_origins
402	-	402	+
403	wrapper.define_kernel(name, compile_wrapper.getvalue(), metadata_comment)	403	wrapper.define_kernel(name, compile_wrapper.getvalue(), metadata_comment)
404		404
405	if metrics.is_metric_table_enabled("kernel_metadata"):	405	if metrics.is_metric_table_enabled("kernel_metadata"):
	@@ -417,15 +417,15 @@ class NpuMetaScheduling(SIMDScheduling):
417	return "auto_fallback"	417	return "auto_fallback"
418		418
419	def _dump_fx_graph_for_fallback(self, mlir_kernel, device, graph_hash, kernel_name, compile_code):	419	def _dump_fx_graph_for_fallback(self, mlir_kernel, device, graph_hash, kernel_name, compile_code):
420	-	420	+
421	cache_root = os.getenv("TORCHINDUCTOR_CACHE_DIR")	421	cache_root = os.getenv("TORCHINDUCTOR_CACHE_DIR")
422	dump_path = os.path.join(	422	dump_path = os.path.join(
423	- cache_root,	423	+ cache_root,
424	- anir_config.traced_graph_cache or "traced_graph_cache",	424	+ anir_config.traced_graph_cache or "traced_graph_cache",
425	- str(device.index),	425	+ str(device.index),
426	graph_hash	426	graph_hash
427	)	427	)
428	-	428	+
429	if not os.path.exists(dump_path):	429	if not os.path.exists(dump_path):
430	os.makedirs(dump_path, exist_ok=True)	430	os.makedirs(dump_path, exist_ok=True)
431	to_folder(mlir_kernel._gm, dump_path, graph_hash=graph_hash, module_name=graph_hash)	431	to_folder(mlir_kernel._gm, dump_path, graph_hash=graph_hash, module_name=graph_hash)
	@@ -433,11 +433,11 @@ class NpuMetaScheduling(SIMDScheduling):
433	if anir_config.fx_subgraph_dump_path is not None:	433	if anir_config.fx_subgraph_dump_path is not None:
434	subgraph_dump_path = os.path.join(anir_config.fx_subgraph_dump_path, str(device.index), kernel_name)	434	subgraph_dump_path = os.path.join(anir_config.fx_subgraph_dump_path, str(device.index), kernel_name)
435	os.makedirs(subgraph_dump_path, exist_ok=True)	435	os.makedirs(subgraph_dump_path, exist_ok=True)
436	-	436	+
437	num_args = len(mlir_kernel._gm.code.split('forward(', )[1].split(')')[0].split(', ')) - 1	437	num_args = len(mlir_kernel._gm.code.split('forward(', )[1].split(')')[0].split(', ')) - 1
438	fx_graph_code = get_fx_graph_code(mlir_kernel._gm.code, num_args, runnable=False, kernel_code=compile_code, kernel_name=kernel_name)	438	fx_graph_code = get_fx_graph_code(mlir_kernel._gm.code, num_args, runnable=False, kernel_code=compile_code, kernel_name=kernel_name)
439	runnable_fx_graph_code = get_fx_graph_code(mlir_kernel._gm.code, num_args, runnable=True, kernel_code=compile_code, kernel_name=kernel_name)	439	runnable_fx_graph_code = get_fx_graph_code(mlir_kernel._gm.code, num_args, runnable=True, kernel_code=compile_code, kernel_name=kernel_name)
440	-	440	+
441	with open(os.path.join(subgraph_dump_path, f'{kernel_name}.py'), 'w') as f:	441	with open(os.path.join(subgraph_dump_path, f'{kernel_name}.py'), 'w') as f:
442	f.write(fx_graph_code)	442	f.write(fx_graph_code)
443	with open(os.path.join(subgraph_dump_path, f'runnable_{kernel_name}.py'), 'w') as f:	443	with open(os.path.join(subgraph_dump_path, f'runnable_{kernel_name}.py'), 'w') as f:
	@@ -519,10 +519,10 @@ class NpuMetaScheduling(SIMDScheduling):
519		519
520	def codegen_node(self, node: Union[scheduler.SchedulerNode, object]):	520	def codegen_node(self, node: Union[scheduler.SchedulerNode, object]):
521	nodes: List[scheduler.SchedulerNode] = node.get_nodes()	521	nodes: List[scheduler.SchedulerNode] = node.get_nodes()
522	-	522	+
523	_, (numel, rnumel) = max(nodes, key=lambda x: int(x.is_reduction())).group	523	_, (numel, rnumel) = max(nodes, key=lambda x: int(x.is_reduction())).group
524		524
525	node_schedule = self.generate_node_schedule(nodes, numel, rnumel)	525	node_schedule = self.generate_node_schedule(nodes, numel, rnumel)
526	kernel_features = SIMDKernelFeatures(node_schedule, numel, rnumel)	526	kernel_features = SIMDKernelFeatures(node_schedule, numel, rnumel)
527	-	527	+
528	return self.codegen_node_schedule(kernel_features, nodes)	528	return self.codegen_node_schedule(kernel_features, nodes)

	@@ -156,9 +156,9 @@ def _catlass_tensor_from_node_for_bias(node):
156		156
157	# bias node is different to A B input tensors. Even (n,) bias, the shape of this bias at this step,	157	# bias node is different to A B input tensors. Even (n,) bias, the shape of this bias at this step,
158	# should already be the broadcasted shape (m, n); the difference between the (n,) bias and (m, n) bias	158	# should already be the broadcasted shape (m, n); the difference between the (n,) bias and (m, n) bias
159	- # is the stride. (n,) bias stride should be (0, 1), but (m, n) bias stride should not contain any zero.	159	+ # is the stride. (n,) bias stride should be (0, 1), but (m, n) bias stride should not contain any zero.
160		160
161	- if len(node.get_size()) == 1 and len(node.get_stride()) == 1:	161	+ if len(node.get_size()) == 1 and len(node.get_stride()) == 1:
162	shape = tuple(node.get_layout().size)	162	shape = tuple(node.get_layout().size)
163	stride = tuple(node.get_layout().stride)	163	stride = tuple(node.get_layout().stride)
164	elif node.get_stride()[0] == 0:	164	elif node.get_stride()[0] == 0:
	@@ -200,10 +200,10 @@ def _gen_ops_cached(arch: str, op_tensors=None, is_group_mm=False) -> List[Any]:
200	gemm_plan = Gemm(atlas_arch=arch, element_C=element_C, A=op_tensors[0], B=op_tensors[1], Bias=op_tensors[2])	200	gemm_plan = Gemm(atlas_arch=arch, element_C=element_C, A=op_tensors[0], B=op_tensors[1], Bias=op_tensors[2])
201	kernels = gemm_plan.get_kernels()	201	kernels = gemm_plan.get_kernels()
202	else: # group mm	202	else: # group mm
203	- gemm_plan = GroupGemm(atlas_arch=arch, element_C=element_C, A=op_tensors[0], B=op_tensors[1],	203	+ gemm_plan = GroupGemm(atlas_arch=arch, element_C=element_C, A=op_tensors[0], B=op_tensors[1],
204	groupList=op_tensors[2])	204	groupList=op_tensors[2])
205	kernels = gemm_plan.get_kernels()	205	kernels = gemm_plan.get_kernels()
206	-	206	+
207	for kernel in kernels:	207	for kernel in kernels:
208	tilings = generate_configs(arch, kernel)	208	tilings = generate_configs(arch, kernel)
209	if tilings:	209	if tilings:

	@@ -75,7 +75,7 @@ class IndexAnalysis:
75	for key, coeff in self.index.as_coefficients_dict().items()	75	for key, coeff in self.index.as_coefficients_dict().items()
76	if not isinstance(key, sympy.Integer)	76	if not isinstance(key, sympy.Integer)
77	]	77	]
78	- # sort by stride	78	+ # sort by stride
79	self.var_stride.sort(key=lambda x: x[1])	79	self.var_stride.sort(key=lambda x: x[1])
80	# only contains tiling axis var	80	# only contains tiling axis var
81	self.var_list = tuple([x[0] for x in self.var_stride if x[0] in self.tiling_axis])	81	self.var_list = tuple([x[0] for x in self.var_stride if x[0] in self.tiling_axis])
	@@ -325,7 +325,7 @@ class ReductionAnalysis:
325	sizes = self.dense_size_list()	325	sizes = self.dense_size_list()
326	num_red = self.numof_reduction_axis()	326	num_red = self.numof_reduction_axis()
327	is_contig = self.contiguous_reduction	327	is_contig = self.contiguous_reduction
328	-	328	+
329	if num_red > 1:	329	if num_red > 1:
330	if is_contig:	330	if is_contig:
331	result = f"[{', '.join(self.dense_post_reduction_list())}]"	331	result = f"[{', '.join(self.dense_post_reduction_list())}]"
	@@ -333,7 +333,7 @@ class ReductionAnalysis:
333	result = f"[{'* '.join(sizes)}]"	333	result = f"[{'* '.join(sizes)}]"
334	else:	334	else:
335	result = f"[{', '.join(sizes)}]"	335	result = f"[{', '.join(sizes)}]"
336	-	336	+
337	return result	337	return result
338		338
339	def numof_reduction_axis(self):	339	def numof_reduction_axis(self):

	@@ -25,9 +25,9 @@ from ..fx_passes.utils.schedule_node_utils import is_multi_stream
25		25
26	@dataclasses.dataclass	26	@dataclasses.dataclass
27	class NPUMultiOutputLine(MultiOutputLine):	27	class NPUMultiOutputLine(MultiOutputLine):
28	-	28	+
29	multi_stream_intent: str = ""	29	multi_stream_intent: str = ""
30	-	30	+
31	def codegen(self, code: IndentedBuffer) -> None:	31	def codegen(self, code: IndentedBuffer) -> None:
32	def codegen_list_tuple_access(basename, indices): # type: ignore[no-untyped-def]	32	def codegen_list_tuple_access(basename, indices): # type: ignore[no-untyped-def]
33	if len(indices) > 0:	33	if len(indices) > 0:
	@@ -373,8 +373,8 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
373	multi_stream_intent_str = self.get_buffer_define_multi_stream_by_name(new_name)	373	multi_stream_intent_str = self.get_buffer_define_multi_stream_by_name(new_name)
374	return f"{multi_stream_intent_str}{self.declare_maybe_reference}{new_name} = {old_name}{del_line}{self.ending} {self.comment} reuse"	374	return f"{multi_stream_intent_str}{self.declare_maybe_reference}{new_name} = {old_name}{del_line}{self.ending} {self.comment} reuse"
375	return super().codegen_exact_buffer_reuse(old_name, new_name, del_line)	375	return super().codegen_exact_buffer_reuse(old_name, new_name, del_line)
376	-	376	+
377	-	377	+
378	def codegen_deferred_allocation(self, name: str, view: ir.ReinterpretView) -> None:	378	def codegen_deferred_allocation(self, name: str, view: ir.ReinterpretView) -> None:
379	if is_multi_stream():	379	if is_multi_stream():
380	multi_stream_intent_str = self.get_buffer_define_multi_stream_by_name(name)	380	multi_stream_intent_str = self.get_buffer_define_multi_stream_by_name(name)
	@@ -517,8 +517,8 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
517	self.kernel_declarations.getvaluewithlinemap(),	517	self.kernel_declarations.getvaluewithlinemap(),
518	)	518	)
519	return super()._generate(is_inference)	519	return super()._generate(is_inference)
520	-	520	+
521	-	521	+
522	def handle_cross_stream_del_buf(self):	522	def handle_cross_stream_del_buf(self):
523	total_lines = len(self.lines)	523	total_lines = len(self.lines)
524	sub_streams_line_no = self.get_sub_streams_line_no()	524	sub_streams_line_no = self.get_sub_streams_line_no()
	@@ -530,12 +530,12 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
530	tab_value = self.buffer_args_multi_stream_intent[keys[0]]	530	tab_value = self.buffer_args_multi_stream_intent[keys[0]]
531	if idx > sub_stream_line[0] and idx < sub_stream_line[1] and isinstance(line, WrapperLine) and hasattr(line, "node") and line.node.get_name() not in self.buffer_args_multi_stream_intent.keys():	531	if idx > sub_stream_line[0] and idx < sub_stream_line[1] and isinstance(line, WrapperLine) and hasattr(line, "node") and line.node.get_name() not in self.buffer_args_multi_stream_intent.keys():
532	self.buffer_args_multi_stream_intent[line.node.get_name()] = tab_value	532	self.buffer_args_multi_stream_intent[line.node.get_name()] = tab_value
533	-	533	+
534	n = len(sub_streams_line_no)	534	n = len(sub_streams_line_no)
535	for i in range(1, n):	535	for i in range(1, n):
536	prev_end = sub_streams_line_no[i-1][1]	536	prev_end = sub_streams_line_no[i-1][1]
537	curr_start = sub_streams_line_no[i][0]	537	curr_start = sub_streams_line_no[i][0]
538	-	538	+
539	if prev_end < idx < curr_start and isinstance(line, WrapperLine) and hasattr(line, "node") and line.node.get_name() in self.buffer_args_multi_stream_intent.keys():	539	if prev_end < idx < curr_start and isinstance(line, WrapperLine) and hasattr(line, "node") and line.node.get_name() in self.buffer_args_multi_stream_intent.keys():
540	self.buffer_args_multi_stream_intent.pop(line.node.get_name(), None)	540	self.buffer_args_multi_stream_intent.pop(line.node.get_name(), None)
541		541
	@@ -543,8 +543,8 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
543	last_end = sub_streams_line_no[-1][1]	543	last_end = sub_streams_line_no[-1][1]
544	if last_end < idx < total_lines and isinstance(line, WrapperLine) and hasattr(line, "node") and line.node.get_name() in self.buffer_args_multi_stream_intent.keys():	544	if last_end < idx < total_lines and isinstance(line, WrapperLine) and hasattr(line, "node") and line.node.get_name() in self.buffer_args_multi_stream_intent.keys():
545	self.buffer_args_multi_stream_intent.pop(line.node.get_name(), None)	545	self.buffer_args_multi_stream_intent.pop(line.node.get_name(), None)
546	-	546	+
547	-	547	+
548	def get_sub_streams_line_no(self):	548	def get_sub_streams_line_no(self):
549	sub_streams_line_no = []	549	sub_streams_line_no = []
550	i = 0	550	i = 0

	@@ -29,7 +29,7 @@ class CVParallelismStrategy(ParallelStrategyBase):
29	return ComputeType.VECTOR	29	return ComputeType.VECTOR
30		30
31	node_class = node_obj.__class__.__name__.lower()	31	node_class = node_obj.__class__.__name__.lower()
32	-	32	+
33	if "multioutput" in node_class:	33	if "multioutput" in node_class:
34	inputs = getattr(node_obj, 'inputs', [])	34	inputs = getattr(node_obj, 'inputs', [])
35	if inputs and len(inputs) > 0:	35	if inputs and len(inputs) > 0:
	@@ -55,7 +55,7 @@ class CVParallelismStrategy(ParallelStrategyBase):
55	'convolution', 'conv', 'mm', 'addmm', 'bmm', 'matmul', 'linear', 'einsum'	55	'convolution', 'conv', 'mm', 'addmm', 'bmm', 'matmul', 'linear', 'einsum'
56	]):	56	]):
57	return ComputeType.CUBE	57	return ComputeType.CUBE
58	-	58	+
59	if isinstance(node_data, (Pointwise, Reduction)) and hasattr(node_obj, "origins"):	59	if isinstance(node_data, (Pointwise, Reduction)) and hasattr(node_obj, "origins"):
60	origin_ops = [str(o) for o in node_obj.origins]	60	origin_ops = [str(o) for o in node_obj.origins]
61	complex_ops = ['softmax', 'norm', 'layer_norm']	61	complex_ops = ['softmax', 'norm', 'layer_norm']
	@@ -69,7 +69,7 @@ class CVParallelismStrategy(ParallelStrategyBase):
69	'add', 'relu', 'gelu', 'mul', 'bias', 'sigmoid', 'silu'	69	'add', 'relu', 'gelu', 'mul', 'bias', 'sigmoid', 'silu'
70	]):	70	]):
71	return ComputeType.MLP_VECTOR	71	return ComputeType.MLP_VECTOR
72	-	72	+
73	if 'pointwise' in data_class.lower():	73	if 'pointwise' in data_class.lower():
74	return ComputeType.VECTOR	74	return ComputeType.VECTOR
75	return ComputeType.UNKNOWN	75	return ComputeType.UNKNOWN
	@@ -126,11 +126,11 @@ class CVParallelismStrategy(ParallelStrategyBase):
126	sorted_candidates = sorted(candidates, key=lambda x: pos_map.get(x, float('inf')))	126	sorted_candidates = sorted(candidates, key=lambda x: pos_map.get(x, float('inf')))
127	max_segment = []	127	max_segment = []
128	current_segment = [sorted_candidates[0]]	128	current_segment = [sorted_candidates[0]]
129	-	129	+
130	for i in range(1, len(sorted_candidates)):	130	for i in range(1, len(sorted_candidates)):
131	prev_node = sorted_candidates[i-1]	131	prev_node = sorted_candidates[i-1]
132	curr_node = sorted_candidates[i]	132	curr_node = sorted_candidates[i]
133	-	133	+
134	if pos_map[curr_node] == pos_map[prev_node] + 1:	134	if pos_map[curr_node] == pos_map[prev_node] + 1:
135	current_segment.append(curr_node)	135	current_segment.append(curr_node)
136	else:	136	else:
	@@ -185,7 +185,7 @@ class CVParallelismStrategy(ParallelStrategyBase):
185	# 计算每个分组中node节点数量，如果小于最小值则不分组	185	# 计算每个分组中node节点数量，如果小于最小值则不分组
186	vector_group_len = self.calculate_group_len(final_vec)	186	vector_group_len = self.calculate_group_len(final_vec)
187	cube_group_len = self.calculate_group_len(final_cube)	187	cube_group_len = self.calculate_group_len(final_cube)
188	-	188	+
189	sub_nodes_min = os.environ.get("PARALLEL_SCHEDULER_NODES_MIN", 20) // 5	189	sub_nodes_min = os.environ.get("PARALLEL_SCHEDULER_NODES_MIN", 20) // 5
190	if cube_group_len < sub_nodes_min or vector_group_len < sub_nodes_min:	190	if cube_group_len < sub_nodes_min or vector_group_len < sub_nodes_min:
191	invalidate_all = True	191	invalidate_all = True

	@@ -75,52 +75,52 @@ class Mode(Enum):
75	def _get_flex_attention_additional_lowerings():	75	def _get_flex_attention_additional_lowerings():
76	"""	76	"""
77	Get additional lowerings for flex_attention subgraph.	77	Get additional lowerings for flex_attention subgraph.
78	-	78	+
79	These lowerings are used to allow index and bitwise operations to be lowered	79	These lowerings are used to allow index and bitwise operations to be lowered
80	as pointwise ops instead of fallback in the mask_mod subgraph.	80	as pointwise ops instead of fallback in the mask_mod subgraph.
81	"""	81	"""
82	from torch._inductor.lowering import make_pointwise, index_impl	82	from torch._inductor.lowering import make_pointwise, index_impl
83	from torch._inductor.subgraph_lowering import PointwiseSubgraphLowering	83	from torch._inductor.subgraph_lowering import PointwiseSubgraphLowering
84	-	84	+
85	additional_lowerings = {}	85	additional_lowerings = {}
86	-	86	+
87	def index_pointwise(x, indices):	87	def index_pointwise(x, indices):
88	return index_impl(x, indices, check=True)	88	return index_impl(x, indices, check=True)
89	-	89	+
90	additional_lowerings[aten.index] = index_pointwise	90	additional_lowerings[aten.index] = index_pointwise
91	additional_lowerings[aten.index.Tensor] = index_pointwise	91	additional_lowerings[aten.index.Tensor] = index_pointwise
92	-	92	+
93	bitwise_and_fn = make_pointwise(ops.bitwise_and)	93	bitwise_and_fn = make_pointwise(ops.bitwise_and)
94	-	94	+
95	def bitwise_and_tensor(a, b):	95	def bitwise_and_tensor(a, b):
96	return bitwise_and_fn(a, b)	96	return bitwise_and_fn(a, b)
97	-	97	+
98	bitwise_or_fn = make_pointwise(ops.bitwise_or)	98	bitwise_or_fn = make_pointwise(ops.bitwise_or)
99	-	99	+
100	def bitwise_or_tensor(a, b):	100	def bitwise_or_tensor(a, b):
101	return bitwise_or_fn(a, b)	101	return bitwise_or_fn(a, b)
102	-	102	+
103	bitwise_not_fn = make_pointwise(ops.bitwise_not)	103	bitwise_not_fn = make_pointwise(ops.bitwise_not)
104	-	104	+
105	def bitwise_not_default(a):	105	def bitwise_not_default(a):
106	return bitwise_not_fn(a)	106	return bitwise_not_fn(a)
107	-	107	+
108	additional_lowerings[aten.bitwise_and.Tensor] = bitwise_and_tensor	108	additional_lowerings[aten.bitwise_and.Tensor] = bitwise_and_tensor
109	additional_lowerings[aten.bitwise_or.Tensor] = bitwise_or_tensor	109	additional_lowerings[aten.bitwise_or.Tensor] = bitwise_or_tensor
110	additional_lowerings[aten.bitwise_not.default] = bitwise_not_default	110	additional_lowerings[aten.bitwise_not.default] = bitwise_not_default
111	-	111	+
112	return additional_lowerings	112	return additional_lowerings
113		113
114		114
115	def _build_subgraph_buffer_with_additional_lowerings(args, subgraph):	115	def _build_subgraph_buffer_with_additional_lowerings(args, subgraph):
116	"""	116	"""
117	Build subgraph buffer with additional lowerings for flex_attention.	117	Build subgraph buffer with additional lowerings for flex_attention.
118	-	118	+
119	This function creates a PointwiseSubgraphLowering with additional_lowerings	119	This function creates a PointwiseSubgraphLowering with additional_lowerings
120	to handle index and bitwise operations as pointwise ops.	120	to handle index and bitwise operations as pointwise ops.
121	"""	121	"""
122	from torch._inductor.subgraph_lowering import PointwiseSubgraphLowering	122	from torch._inductor.subgraph_lowering import PointwiseSubgraphLowering
123	-	123	+
124	additional_lowerings = _get_flex_attention_additional_lowerings()	124	additional_lowerings = _get_flex_attention_additional_lowerings()
125	pw_subgraph = PointwiseSubgraphLowering(	125	pw_subgraph = PointwiseSubgraphLowering(
126	subgraph.graph_module,	126	subgraph.graph_module,
	@@ -129,7 +129,7 @@ def _build_subgraph_buffer_with_additional_lowerings(args, subgraph):
129	)	129	)
130	with V.set_graph_handler(pw_subgraph):	130	with V.set_graph_handler(pw_subgraph):
131	pw_subgraph.run(*args)	131	pw_subgraph.run(*args)
132	-	132	+
133	def convert_output_node_to_buffer(output_buffer):	133	def convert_output_node_to_buffer(output_buffer):
134	from torch._inductor.ir import ComputedBuffer, FlexibleLayout, StorageBox	134	from torch._inductor.ir import ComputedBuffer, FlexibleLayout, StorageBox
135	if output_buffer is None:	135	if output_buffer is None:
	@@ -154,7 +154,7 @@ def _build_subgraph_buffer_with_additional_lowerings(args, subgraph):
154	data=output_buffer.data.data,	154	data=output_buffer.data.data,
155	)	155	)
156	return subgraph_buffer	156	return subgraph_buffer
157	-	157	+
158	return tree_map(convert_output_node_to_buffer, pw_subgraph.graph_outputs)	158	return tree_map(convert_output_node_to_buffer, pw_subgraph.graph_outputs)
159		159
160		160

	@@ -12,8 +12,8 @@ from torch._ops import OpOverload, OpOverloadPacket
12	def write_to_file(items: Set[str], file_name: str):	12	def write_to_file(items: Set[str], file_name: str):
13	with open(file_name, "w") as f:	13	with open(file_name, "w") as f:
14	for item in items:	14	for item in items:
15	- f.write(f"{item},\n")	15	+ f.write(f"{item},\n")
16	-	16	+
17		17
18	def get_base_packet(op):	18	def get_base_packet(op):
19	if isinstance(op, OpOverload):	19	if isinstance(op, OpOverload):
	@@ -29,32 +29,32 @@ def get_overload_set(inputs, output_set):
29	output_set.add(other_fn)	29	output_set.add(other_fn)
30	else:	30	else:
31	output_set.add(fn)	31	output_set.add(fn)
32	-	32	+
33	-	33	+
34	def get_native_fallbacks_base() -> Set[str]:	34	def get_native_fallbacks_base() -> Set[str]:
35	native_fallbacks = {str(op) for op in lowering.fallbacks}	35	native_fallbacks = {str(op) for op in lowering.fallbacks}
36	write_to_file(sorted(native_fallbacks), f"native_fallbacks_{len(native_fallbacks)}.txt")	36	write_to_file(sorted(native_fallbacks), f"native_fallbacks_{len(native_fallbacks)}.txt")
37	-	37	+
38	native_fallbacks_base = {str(get_base_packet(op)) for op in lowering.fallbacks}	38	native_fallbacks_base = {str(get_base_packet(op)) for op in lowering.fallbacks}
39	write_to_file(sorted(native_fallbacks_base), f"native_fallbacks_base_{len(native_fallbacks_base)}.txt")	39	write_to_file(sorted(native_fallbacks_base), f"native_fallbacks_base_{len(native_fallbacks_base)}.txt")
40	-	40	+
41	print(f"[native torch] len(lowering.fallbacks): {len(lowering.fallbacks)}, len(native_fallbacks_base): {len(native_fallbacks_base)}")	41	print(f"[native torch] len(lowering.fallbacks): {len(lowering.fallbacks)}, len(native_fallbacks_base): {len(native_fallbacks_base)}")
42	return native_fallbacks, native_fallbacks_base	42	return native_fallbacks, native_fallbacks_base
43	-	43	+
44	-	44	+
45	def get_npu_fallbacks_base() -> Set[str]:	45	def get_npu_fallbacks_base() -> Set[str]:
46	import torch_npu._inductor	46	import torch_npu._inductor
47	-	47	+
48	npu_fallbacks = {str(op) for op in lowering.fallbacks}	48	npu_fallbacks = {str(op) for op in lowering.fallbacks}
49	write_to_file(sorted(npu_fallbacks), f"npu_fallbacks_{len(npu_fallbacks)}.txt")	49	write_to_file(sorted(npu_fallbacks), f"npu_fallbacks_{len(npu_fallbacks)}.txt")
50	-	50	+
51	npu_fallbacks_base = {str(get_base_packet(op)) for op in lowering.fallbacks}	51	npu_fallbacks_base = {str(get_base_packet(op)) for op in lowering.fallbacks}
52	- write_to_file(sorted(npu_fallbacks_base), f"npu_fallbacks_base_{len(npu_fallbacks_base)}.txt")	52	+ write_to_file(sorted(npu_fallbacks_base), f"npu_fallbacks_base_{len(npu_fallbacks_base)}.txt")
53	-	53	+
54	print(f"[torch_npu] len(lowering.fallbacks): {len(lowering.fallbacks)}, len(npu_fallbacks_base): {len(npu_fallbacks_base)}")	54	print(f"[torch_npu] len(lowering.fallbacks): {len(lowering.fallbacks)}, len(npu_fallbacks_base): {len(npu_fallbacks_base)}")
55	return npu_fallbacks, npu_fallbacks_base	55	return npu_fallbacks, npu_fallbacks_base
56		56
57	-	57	+
58	def npu_extra_fallbacks_base_diff(native_fallbacks_base, npu_fallbacks_base) -> None:	58	def npu_extra_fallbacks_base_diff(native_fallbacks_base, npu_fallbacks_base) -> None:
59	npu_extra_fallbacks_base = npu_fallbacks_base - native_fallbacks_base	59	npu_extra_fallbacks_base = npu_fallbacks_base - native_fallbacks_base
60	write_to_file(sorted(npu_extra_fallbacks_base), f"npu_extra_fallbacks_base_{len(npu_extra_fallbacks_base)}.txt")	60	write_to_file(sorted(npu_extra_fallbacks_base), f"npu_extra_fallbacks_base_{len(npu_extra_fallbacks_base)}.txt")

	@@ -79,11 +79,11 @@ void BSShapeOpStrategy::InitializeCore(std::vector<int64_t>& gears, int dimensio
79	"Maximum batch size (", gears.back(), ") must be <= ", MAX_BS_GEAR, ".");	79	"Maximum batch size (", gears.back(), ") must be <= ", MAX_BS_GEAR, ".");
80	TORCH_CHECK(std::adjacent_find(gears.begin(), gears.end()) == gears.end(),	80	TORCH_CHECK(std::adjacent_find(gears.begin(), gears.end()) == gears.end(),
81	"Batch size gears must be unique.")	81	"Batch size gears must be unique.")
82	-	82	+
83	TORCH_CHECK(dimension >= 0, "Dimension must be non-negative (got ", dimension, ").");	83	TORCH_CHECK(dimension >= 0, "Dimension must be non-negative (got ", dimension, ").");
84	-	84	+
85	TORCH_CHECK(indices.size() > 0, "At least one tensor index must be provided for transformation.");	85	TORCH_CHECK(indices.size() > 0, "At least one tensor index must be provided for transformation.");
86	-	86	+
87	std::sort(indices.begin(), indices.end());	87	std::sort(indices.begin(), indices.end());
88	TORCH_CHECK(indices.front() >= 0,	88	TORCH_CHECK(indices.front() >= 0,
89	"Tensor index must be non-negative (found negative index: ", indices.front(), ").");	89	"Tensor index must be non-negative (found negative index: ", indices.front(), ").");
	@@ -104,7 +104,7 @@ void SeqShapeOpStrategy::InitializeCore(std::vector<int64_t>& gears, std::vector
104	TORCH_CHECK(!gears.empty(), "At least one sequence gear must be provided.");	104	TORCH_CHECK(!gears.empty(), "At least one sequence gear must be provided.");
105	TORCH_CHECK(gears.size() <= MAX_GEARS_NUM,	105	TORCH_CHECK(gears.size() <= MAX_GEARS_NUM,
106	"Number of sequence length gears (", gears.size(), ") exceeds maximum supported (", MAX_GEARS_NUM, ").");	106	"Number of sequence length gears (", gears.size(), ") exceeds maximum supported (", MAX_GEARS_NUM, ").");
107	-	107	+
108	std::sort(gears.begin(), gears.end());	108	std::sort(gears.begin(), gears.end());
109	TORCH_CHECK(gears.front() >= MIN_SEQ_GEAR,	109	TORCH_CHECK(gears.front() >= MIN_SEQ_GEAR,
110	"Minimum sequence length (", gears.front(), ") must be >= ", MIN_SEQ_GEAR, ".");	110	"Minimum sequence length (", gears.front(), ") must be >= ", MIN_SEQ_GEAR, ".");
	@@ -136,7 +136,7 @@ void SeqShapeOpStrategy::InitializeCore(std::vector<int64_t>& gears, std::vector
136	m_indices.push_back(indices[i]);	136	m_indices.push_back(indices[i]);
137	m_dimensions.push_back(dimensions[i]);	137	m_dimensions.push_back(dimensions[i]);
138	}	138	}
139	-	139	+
140	m_value = value;	140	m_value = value;
141	m_gears = gears;	141	m_gears = gears;
142	m_min_gear = gears.front();	142	m_min_gear = gears.front();
	@@ -383,7 +383,7 @@ void NPUShapeHandling::Initialize(ShapeType type, int64_t min_size, int64_t max_
383	std::vector<int>& dimensions, std::vector<int>& indices, double value)	383	std::vector<int>& dimensions, std::vector<int>& indices, double value)
384	{	384	{
385	m_policy = policy;	385	m_policy = policy;
386	-	386	+
387	std::vector<int64_t> gears;	387	std::vector<int64_t> gears;
388	GenerateGears(min_size, max_size, policy, gears);	388	GenerateGears(min_size, max_size, policy, gears);
389		389

	@@ -426,7 +426,7 @@ def npu_fusion_attention_grad_v3_strategy(query, key, value, dy, head_num, input
426	None, None, None, None, None, # others	426	None, None, None, None, None, # others
427	None if seed is None else Replicate(), # seed	427	None if seed is None else Replicate(), # seed
428	None if offset is None else Replicate(), # offset	428	None if offset is None else Replicate(), # offset
429	- None,	429	+ None,
430	None if actual_seq_qlen is None else Replicate(), # actual_seq_qlen	430	None if actual_seq_qlen is None else Replicate(), # actual_seq_qlen
431	None if actual_seq_kvlen is None else Replicate(), # actual_seq_kvlen	431	None if actual_seq_kvlen is None else Replicate(), # actual_seq_kvlen
432	None, None, None, None, # others	432	None, None, None, None, # others
	@@ -477,7 +477,7 @@ def npu_fusion_attention_grad_v3_strategy(query, key, value, dy, head_num, input
477	None, None, None, None, None, # others	477	None, None, None, None, None, # others
478	None if seed is None else Replicate(), # seed	478	None if seed is None else Replicate(), # seed
479	None if offset is None else Replicate(), # offset	479	None if offset is None else Replicate(), # offset
480	- None,	480	+ None,
481	None if actual_seq_qlen is None else Replicate(), # actual_seq_qlen	481	None if actual_seq_qlen is None else Replicate(), # actual_seq_qlen
482	None if actual_seq_kvlen is None else Replicate(), # actual_seq_kvlen	482	None if actual_seq_kvlen is None else Replicate(), # actual_seq_kvlen
483	None, None, None, None, # others	483	None, None, None, None, # others
	@@ -520,7 +520,7 @@ def npu_fusion_attention_grad_v3_strategy(query, key, value, dy, head_num, input
520	None, None, None, None, None, # others	520	None, None, None, None, None, # others
521	None if seed is None else Replicate(), # seed	521	None if seed is None else Replicate(), # seed
522	None if offset is None else Replicate(), # offset	522	None if offset is None else Replicate(), # offset
523	- None,	523	+ None,
524	None if actual_seq_qlen is None else Replicate(), # actual_seq_qlen	524	None if actual_seq_qlen is None else Replicate(), # actual_seq_qlen
525	None if actual_seq_kvlen is None else Replicate(), # actual_seq_kvlen	525	None if actual_seq_kvlen is None else Replicate(), # actual_seq_kvlen
526	None, None, None, None, # others	526	None, None, None, None, # others

	@@ -176,12 +176,12 @@ class NpuFusedAdamP(NpuFusedOptimizerBase):
176	def _maybe_init_combined_states(self):	176	def _maybe_init_combined_states(self):
177	if self.is_states_combined:	177	if self.is_states_combined:
178	return	178	return
179	-	179	+
180	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]	180	self.combined_param_states_indexed_by_group = len(self.param_groups) * [None]
181		181
182	for i, _ in enumerate(self.param_groups):	182	for i, _ in enumerate(self.param_groups):
183	self._combine_group_param_states(i)	183	self._combine_group_param_states(i)
184	-	184	+
185	if not all(value is None for value in self.combined_param_states_indexed_by_group):	185	if not all(value is None for value in self.combined_param_states_indexed_by_group):
186	self.is_states_combined = True	186	self.is_states_combined = True
187		187
	@@ -260,7 +260,7 @@ class NpuFusedAdamP(NpuFusedOptimizerBase):
260	def step(self, closure=None):	260	def step(self, closure=None):
261	if not self.is_params_grads_combined:	261	if not self.is_params_grads_combined:
262	self._maybe_init_combined_params_and_grads()	262	self._maybe_init_combined_params_and_grads()
263	-	263	+
264	if not self.is_states_combined:	264	if not self.is_states_combined:
265	self._maybe_init_combined_states()	265	self._maybe_init_combined_states()
266		266

	@@ -77,7 +77,7 @@ class MemoryDbParser(BaseParser):
77	pta_ge_record_list[MemoryRecordTableRow.STREAM_PTR.value] = cur_record[MemoryRecordTableRow.STREAM_PTR.value] if cur_record[MemoryRecordTableRow.STREAM_PTR.value] \	77	pta_ge_record_list[MemoryRecordTableRow.STREAM_PTR.value] = cur_record[MemoryRecordTableRow.STREAM_PTR.value] if cur_record[MemoryRecordTableRow.STREAM_PTR.value] \
78	else last_record_data[MemoryRecordTableRow.STREAM_PTR.value]	78	else last_record_data[MemoryRecordTableRow.STREAM_PTR.value]
79	return [cur_record, pta_ge_record_list]	79	return [cur_record, pta_ge_record_list]
80	-	80	+
81	def run(self, deps_data: dict):	81	def run(self, deps_data: dict):
82	self.logger.info("MemoryDbParser start.")	82	self.logger.info("MemoryDbParser start.")
83	try:	83	try:
	@@ -94,7 +94,7 @@ class MemoryDbParser(BaseParser):
94	return Constant.FAIL, None	94	return Constant.FAIL, None
95	self.logger.info("MemoryDbParser finish.")	95	self.logger.info("MemoryDbParser finish.")
96	return Constant.SUCCESS, None	96	return Constant.SUCCESS, None
97	-	97	+
98	def init_db_connect(self):	98	def init_db_connect(self):
99	if not TorchDb().create_connect_db():	99	if not TorchDb().create_connect_db():
100	raise RuntimeError(f"Failed to connect to db file: {TorchDb().get_db_path()}")	100	raise RuntimeError(f"Failed to connect to db file: {TorchDb().get_db_path()}")
	@@ -197,7 +197,7 @@ class MemoryDbParser(BaseParser):
197	memory_bean.total_allocated_for_db, memory_bean.total_reserved_for_db,	197	memory_bean.total_allocated_for_db, memory_bean.total_reserved_for_db,
198	memory_bean.total_active_for_db, memory_bean.stream_ptr,	198	memory_bean.total_active_for_db, memory_bean.stream_ptr,
199	self.device_index if self.device_index != -1 else memory_bean.device_index])	199	self.device_index if self.device_index != -1 else memory_bean.device_index])
200	-	200	+
201	def get_pta_ge_record_list(self):	201	def get_pta_ge_record_list(self):
202	"""	202	"""
203	ge records are to be sorted firstly and pta records are already sorted,	203	ge records are to be sorted firstly and pta records are already sorted,
	@@ -245,7 +245,7 @@ class MemoryDbParser(BaseParser):
245	def save_strings_id(self):	245	def save_strings_id(self):
246	TorchDb().create_table_with_headers(DbConstant.TABLE_STRING_IDS, TableColumnsManager.TableColumns.get(DbConstant.TABLE_STRING_IDS))	246	TorchDb().create_table_with_headers(DbConstant.TABLE_STRING_IDS, TableColumnsManager.TableColumns.get(DbConstant.TABLE_STRING_IDS))
247	TorchDb().insert_data_into_table(DbConstant.TABLE_STRING_IDS, Str2IdManager().get_all_string_2_id_data())	247	TorchDb().insert_data_into_table(DbConstant.TABLE_STRING_IDS, Str2IdManager().get_all_string_2_id_data())
248	-	248	+
249	def save_memory_data_to_db(self):	249	def save_memory_data_to_db(self):
250	self.get_ge_memory_data()	250	self.get_ge_memory_data()
251	self.save_memory_record_data_to_db()	251	self.save_memory_record_data_to_db()

	@@ -21,11 +21,11 @@ def _from_dlpack(ext_tensor) -> 'torch.Tensor':
21	def _apply_dlpack_patch():	21	def _apply_dlpack_patch():
22	"""Patch torch.utils.dlpack and torch.utils to use torch_npu implementation for NPU tensors"""	22	"""Patch torch.utils.dlpack and torch.utils to use torch_npu implementation for NPU tensors"""
23	import torch.utils.dlpack as torch_dlpack	23	import torch.utils.dlpack as torch_dlpack
24	-	24	+
25	# Store original functions	25	# Store original functions
26	_original_to_dlpack = torch_dlpack.to_dlpack	26	_original_to_dlpack = torch_dlpack.to_dlpack
27	_original_from_dlpack = torch_dlpack.from_dlpack	27	_original_from_dlpack = torch_dlpack.from_dlpack
28	-	28	+
29	def create_patched_to_dlpack(module_name):	29	def create_patched_to_dlpack(module_name):
30	"""Create a patched to_dlpack function with proper __module__ attribute"""	30	"""Create a patched to_dlpack function with proper __module__ attribute"""
31	def patched_to_dlpack(tensor):	31	def patched_to_dlpack(tensor):
	@@ -35,7 +35,7 @@ def _apply_dlpack_patch():
35	return _original_to_dlpack(tensor)	35	return _original_to_dlpack(tensor)
36	patched_to_dlpack.__module__ = module_name	36	patched_to_dlpack.__module__ = module_name
37	return patched_to_dlpack	37	return patched_to_dlpack
38	-	38	+
39	def create_patched_from_dlpack(module_name):	39	def create_patched_from_dlpack(module_name):
40	"""Create a patched from_dlpack function with proper __module__ attribute"""	40	"""Create a patched from_dlpack function with proper __module__ attribute"""
41	def patched_from_dlpack(ext_tensor):	41	def patched_from_dlpack(ext_tensor):
	@@ -48,35 +48,35 @@ def _apply_dlpack_patch():
48	return _original_from_dlpack(ext_tensor)	48	return _original_from_dlpack(ext_tensor)
49	patched_from_dlpack.__module__ = module_name	49	patched_from_dlpack.__module__ = module_name
50	return patched_from_dlpack	50	return patched_from_dlpack
51	-	51	+
52	# Apply patches to torch.utils.dlpack	52	# Apply patches to torch.utils.dlpack
53	torch_dlpack.to_dlpack = create_patched_to_dlpack('torch.utils.dlpack')	53	torch_dlpack.to_dlpack = create_patched_to_dlpack('torch.utils.dlpack')
54	torch_dlpack.from_dlpack = create_patched_from_dlpack('torch.utils.dlpack')	54	torch_dlpack.from_dlpack = create_patched_from_dlpack('torch.utils.dlpack')
55	-	55	+
56	# Also patch torch.utils.to_dlpack and torch.utils.from_dlpack if they exist	56	# Also patch torch.utils.to_dlpack and torch.utils.from_dlpack if they exist
57	if hasattr(torch.utils, 'to_dlpack'):	57	if hasattr(torch.utils, 'to_dlpack'):
58	_original_torch_utils_to_dlpack = torch.utils.to_dlpack	58	_original_torch_utils_to_dlpack = torch.utils.to_dlpack
59	torch.utils.to_dlpack = create_patched_to_dlpack('torch.utils')	59	torch.utils.to_dlpack = create_patched_to_dlpack('torch.utils')
60	-	60	+
61	if hasattr(torch.utils, 'from_dlpack'):	61	if hasattr(torch.utils, 'from_dlpack'):
62	_original_torch_utils_from_dlpack = torch.utils.from_dlpack	62	_original_torch_utils_from_dlpack = torch.utils.from_dlpack
63	torch.utils.from_dlpack = create_patched_from_dlpack('torch.utils')	63	torch.utils.from_dlpack = create_patched_from_dlpack('torch.utils')
64	-	64	+
65	# Also patch torch.from_dlpack and torch.to_dlpack if they exist	65	# Also patch torch.from_dlpack and torch.to_dlpack if they exist
66	if hasattr(torch, 'from_dlpack'):	66	if hasattr(torch, 'from_dlpack'):
67	_original_torch_from_dlpack = torch.from_dlpack	67	_original_torch_from_dlpack = torch.from_dlpack
68	torch.from_dlpack = create_patched_from_dlpack('torch')	68	torch.from_dlpack = create_patched_from_dlpack('torch')
69	-	69	+
70	if hasattr(torch, 'to_dlpack'):	70	if hasattr(torch, 'to_dlpack'):
71	_original_torch_to_dlpack = torch.to_dlpack	71	_original_torch_to_dlpack = torch.to_dlpack
72	torch.to_dlpack = create_patched_to_dlpack('torch')	72	torch.to_dlpack = create_patched_to_dlpack('torch')
73	-	73	+
74	# Add to_dlpack to torch.__all__ if it exists, otherwise create it	74	# Add to_dlpack to torch.__all__ if it exists, otherwise create it
75	if not hasattr(torch, '__all__'):	75	if not hasattr(torch, '__all__'):
76	torch.__all__ = []	76	torch.__all__ = []
77	if 'to_dlpack' not in torch.__all__:	77	if 'to_dlpack' not in torch.__all__:
78	torch.__all__.append('to_dlpack')	78	torch.__all__.append('to_dlpack')
79	-	79	+
80	# Also ensure from_dlpack is in torch.__all__ if it exists	80	# Also ensure from_dlpack is in torch.__all__ if it exists
81	if hasattr(torch, 'from_dlpack'):	81	if hasattr(torch, 'from_dlpack'):
82	if not hasattr(torch, '__all__'):	82	if not hasattr(torch, '__all__'):