perf(inductor): finalize planned NPU fast launch submission-pytorch-AtomGit

已合并

perf(inductor): finalize planned NPU fast launch submission #45508

Erwinnn创建于 8 天前从jiabiao_o/pytorch： p1-ab-master 合入到 Ascend/pytorch： master

perf(inductor): finalize planned NPU fast launch submission #45508

已合并

Erwinnn创建于 8 天前

从jiabiao_o/pytorch p1-ab-master合入到Ascend/pytorch master

共 4 个文件变更+286-60

Mtorch_npu/_inductor/experimental/python_wrapper_fast_launch/__init__.py+12-2

	@@ -2,11 +2,20 @@ from .types import FastLaunchError
2		2
3		3
4	def __getattr__(name):	4	def __getattr__(name):
5	- if name in ("BoundFastLaunch", "bind_python_wrapper_kernel_fast"):	5	+ if name in (
6	- from .bind import BoundFastLaunch, bind_python_wrapper_kernel_fast	6	+ "BoundFastLaunch",
		7	+ "FinalizedFastLaunch",
		8	+ "bind_python_wrapper_kernel_fast",
		9	+ ):
		10	+ from .bind import (
		11	+ BoundFastLaunch,
		12	+ FinalizedFastLaunch,
		13	+ bind_python_wrapper_kernel_fast,
		14	+ )
7		15
8	return {	16	return {
9	"BoundFastLaunch": BoundFastLaunch,	17	"BoundFastLaunch": BoundFastLaunch,
		18	+ "FinalizedFastLaunch": FinalizedFastLaunch,
10	"bind_python_wrapper_kernel_fast": bind_python_wrapper_kernel_fast,	19	"bind_python_wrapper_kernel_fast": bind_python_wrapper_kernel_fast,
11	}[name]	20	}[name]
12	raise AttributeError(name)	21	raise AttributeError(name)
	@@ -14,6 +23,7 @@ def __getattr__(name):
14		23
15	__all__ = [	24	__all__ = [
16	"BoundFastLaunch",	25	"BoundFastLaunch",
		26	+ "FinalizedFastLaunch",
17	"FastLaunchError",	27	"FastLaunchError",
18	"bind_python_wrapper_kernel_fast",	28	"bind_python_wrapper_kernel_fast",
19	]	29	]

Mtorch_npu/_inductor/experimental/python_wrapper_fast_launch/backend.py+73-19

	@@ -1,5 +1,6 @@
1	from __future__ import annotations	1	from __future__ import annotations
2		2
		3	+import ast
3	from importlib import import_module	4	from importlib import import_module
4	from numbers import Real	5	from numbers import Real
5	from operator import index as operator_index	6	from operator import index as operator_index
	@@ -84,6 +85,17 @@ def _normalize_grid(grid: Any) -> tuple[int, int, int]:
84	return values	85	return values
85		86
86		87
		88	+def _constant_grid(launcher: Any) -> tuple[int, int, int] \| None:
		89	+ exprs = tuple(getattr(launcher, "_npu_fast_launch_grid_exprs", ()) or ())
		90	+ if len(exprs) != 3:
		91	+ return None
		92	+ try:
		93	+ values = tuple(operator_index(ast.literal_eval(expr)) for expr in exprs)
		94	+ except (SyntaxError, ValueError, TypeError):
		95	+ return None
		96	+ return _normalize_grid(values)
		97	+
		98	+
87	def _load_c_extension() -> Any:	99	def _load_c_extension() -> Any:
88	try:	100	try:
89	return import_module("torch_npu._C")	101	return import_module("torch_npu._C")
	@@ -183,10 +195,14 @@ def _validate_runtime_arg_categories(
183	class PlannedFastLaunch:	195	class PlannedFastLaunch:
184	__slots__ = (	196	__slots__ = (
185	"arg_kinds",	197	"arg_kinds",
		198	+ "fixed_args",
186	"get_grid",	199	"get_grid",
187	"launcher",	200	"launcher",
188	"plan",	201	"plan",
		202	+ "runtime_arg_count",
		203	+ "static_grid",
189	"untimed_launch",	204	"untimed_launch",
		205	+ "untimed_static_launch",
190	)	206	)
191		207
192	def __init__(	208	def __init__(
	@@ -195,14 +211,22 @@ class PlannedFastLaunch:
195	launcher: Any,	211	launcher: Any,
196	plan: Any,	212	plan: Any,
197	arg_kinds: tuple[str, ...],	213	arg_kinds: tuple[str, ...],
		214	+ runtime_arg_count: int,
		215	+ fixed_args: tuple[Any, ...],
198	get_grid: Callable[..., Any],	216	get_grid: Callable[..., Any],
199	untimed_launch: Callable[..., Any],	217	untimed_launch: Callable[..., Any],
		218	+ static_grid: tuple[int, int, int] \| None,
		219	+ untimed_static_launch: Callable[..., Any] \| None,
200	) -> None:	220	) -> None:
201	self.launcher = launcher	221	self.launcher = launcher
202	self.plan = plan	222	self.plan = plan
203	self.arg_kinds = arg_kinds	223	self.arg_kinds = arg_kinds
		224	+ self.runtime_arg_count = runtime_arg_count
		225	+ self.fixed_args = fixed_args
204	self.get_grid = get_grid	226	self.get_grid = get_grid
205	self.untimed_launch = untimed_launch	227	self.untimed_launch = untimed_launch
		228	+ self.static_grid = static_grid
		229	+ self.untimed_static_launch = untimed_static_launch
206		230
207	def __call__(	231	def __call__(
208	self,	232	self,
	@@ -210,9 +234,9 @@ class PlannedFastLaunch:
210	*,	234	*,
211	stream: Any,	235	stream: Any,
212	) -> None:	236	) -> None:
213	- if len(args) != len(self.arg_kinds):	237	+ if len(args) != self.runtime_arg_count:
214	raise FastLaunchError(	238	raise FastLaunchError(
215	- f"args_size_mismatch:{len(args)}:{len(self.arg_kinds)}",	239	+ f"args_size_mismatch:{len(args)}:{self.runtime_arg_count}",
216	backend_submitted=False,	240	backend_submitted=False,
217	stable=True,	241	stable=True,
218	)	242	)
	@@ -221,25 +245,32 @@ class PlannedFastLaunch:
221	"stream_is_none",	245	"stream_is_none",
222	backend_submitted=False,	246	backend_submitted=False,
223	)	247	)
224	- try:	248	+ grid = self.static_grid
225	- grid = _normalize_grid(self.get_grid(*args))	249	+ if grid is None:
226	- except FastLaunchError:	250	+ try:
227	- raise	251	+ grid = _normalize_grid(self.get_grid(args, self.fixed_args))
228	- except Exception as exc:	252	+ except FastLaunchError:
229	- raise FastLaunchError(	253	+ raise
230	- f"grid_resolve_error:{type(exc).__name__}",	254	+ except Exception as exc:
231	- backend_submitted=False,	255	+ raise FastLaunchError(
232	- ) from exc	256	+ f"grid_resolve_error:{type(exc).__name__}",
		257	+ backend_submitted=False,
		258	+ ) from exc
233		259
234	try:	260	try:
235	- self.untimed_launch(	261	+ if self.static_grid is not None:
236	- self.plan,	262	+ if self.untimed_static_launch is None:
237	- stream,	263	+ raise RuntimeError("static fast launch entry is unavailable")
238	- grid[0],	264	+ self.untimed_static_launch(self.plan, stream, args)
239	- grid[1],	265	+ else:
240	- grid[2],	266	+ self.untimed_launch(
241	- args,	267	+ self.plan,
242	- )	268	+ stream,
		269	+ grid[0],
		270	+ grid[1],
		271	+ grid[2],
		272	+ args,
		273	+ )
243	except Exception as exc:	274	except Exception as exc:
244	# All recoverable validation is completed before entering C++.	275	# All recoverable validation is completed before entering C++.
245	# Treat errors after the boundary as submitted so fallback can never	276	# Treat errors after the boundary as submitted so fallback can never
	@@ -296,6 +327,13 @@ def build_planned_fast_launch(
296	_validate_callsite_schema(callsite_metadata, arg_kinds, runtime_arg_count)	327	_validate_callsite_schema(callsite_metadata, arg_kinds, runtime_arg_count)
297	if canonical_args is not None:	328	if canonical_args is not None:
298	_validate_runtime_arg_categories(canonical_args, arg_kinds)	329	_validate_runtime_arg_categories(canonical_args, arg_kinds)
		330	+ if canonical_args is None:
		331	+ canonical_args = ()
		332	+ fixed_args = tuple(canonical_args[runtime_arg_count:])
		333	+ if any(
		334	+ kind == "tensor" for kind in arg_kinds[runtime_arg_count:]
		335	+ ):
		336	+ raise FastLaunchPlanUnavailable("fixed_tensor_arg_unsupported")
299		337
300	extension = _load_c_extension()	338	extension = _load_c_extension()
301	make_plan = getattr(extension, "_npu_inductor_make_fast_launch_plan", None)	339	make_plan = getattr(extension, "_npu_inductor_make_fast_launch_plan", None)
	@@ -303,6 +341,15 @@ def build_planned_fast_launch(
303	if not callable(make_plan) or not callable(launch):	341	if not callable(make_plan) or not callable(launch):
304	raise FastLaunchPlanUnavailable("planned_backend_unavailable")	342	raise FastLaunchPlanUnavailable("planned_backend_unavailable")
305		343
		344	+ static_grid = _constant_grid(launcher)
		345	+ static_launch = getattr(
		346	+ extension,
		347	+ "_npu_inductor_fast_launch_static_with_plan",
		348	+ None,
		349	+ )
		350	+ if static_grid is not None and not callable(static_launch):
		351	+ static_grid = None
		352	+
306	enable_simt = bool(getattr(launcher, "_npu_fast_launch_enable_simt", False))	353	enable_simt = bool(getattr(launcher, "_npu_fast_launch_enable_simt", False))
307	shared_mem_dynamic_size = int(	354	shared_mem_dynamic_size = int(
308	getattr(launcher, "_npu_fast_launch_shared_mem_dynamic_size", 0) or 0	355	getattr(launcher, "_npu_fast_launch_shared_mem_dynamic_size", 0) or 0
	@@ -317,6 +364,9 @@ def build_planned_fast_launch(
317	shared_mem_dynamic_size,	364	shared_mem_dynamic_size,
318	is_pure_simt,	365	is_pure_simt,
319	bool(target_support_ffts),	366	bool(target_support_ffts),
		367	+ runtime_arg_count,
		368	+ fixed_args,
		369	+ static_grid or (),
320	)	370	)
321	# The C++ plan owns the stub object; this additional reference owns the	371	# The C++ plan owns the stub object; this additional reference owns the
322	# loaded binary that produced it.	372	# loaded binary that produced it.
	@@ -330,8 +380,12 @@ def build_planned_fast_launch(
330	launcher=launcher,	380	launcher=launcher,
331	plan=plan,	381	plan=plan,
332	arg_kinds=arg_kinds,	382	arg_kinds=arg_kinds,
		383	+ runtime_arg_count=runtime_arg_count,
		384	+ fixed_args=fixed_args,
333	get_grid=get_grid,	385	get_grid=get_grid,
334	untimed_launch=launch,	386	untimed_launch=launch,
		387	+ static_grid=static_grid,
		388	+ untimed_static_launch=static_launch,
335	)	389	)
336		390
337		391

Mtorch_npu/_inductor/experimental/python_wrapper_fast_launch/bind.py+71-5

	@@ -76,15 +76,23 @@ class BoundFastLaunch:
76	"autotuner",	76	"autotuner",
77	"metadata",	77	"metadata",
78	"_direct",	78	"_direct",
		79	+ "_call_slot",
79	"_negative_callable",	80	"_negative_callable",
80	"_negative_launcher",	81	"_negative_launcher",
81	"_static_full_entry_reason",	82	"_static_full_entry_reason",
82	)	83	)
83		84
84	- def __init__(self, autotuner: Any, metadata: dict[str, Any]) -> None:	85	+ def __init__(
		86	+ self,
		87	+ autotuner: Any,
		88	+ metadata: dict[str, Any],
		89	+ *,
		90	+ call_slot: list[Any] \| None = None,
		91	+ ) -> None:
85	self.autotuner = autotuner	92	self.autotuner = autotuner
86	self.metadata = dict(metadata)	93	self.metadata = dict(metadata)
87	self._direct: PlannedFastLaunch \| None = None	94	self._direct: PlannedFastLaunch \| None = None
		95	+ self._call_slot = call_slot
88	self._static_full_entry_reason = _static_full_entry_reason(autotuner)	96	self._static_full_entry_reason = _static_full_entry_reason(autotuner)
89	self._negative_launcher: Any \| None = None	97	self._negative_launcher: Any \| None = None
90	self._negative_callable: Any \| None = None	98	self._negative_callable: Any \| None = None
	@@ -124,6 +132,8 @@ class BoundFastLaunch:
124	self._direct = None	132	self._direct = None
125	self._negative_launcher = launcher	133	self._negative_launcher = launcher
126	self._negative_callable = launcher	134	self._negative_callable = launcher
		135	+ if self._call_slot is not None:
		136	+ self._call_slot[0] = self
127		137
128	def _try_promote(self, args: tuple[Any, ...]) -> bool:	138	def _try_promote(self, args: tuple[Any, ...]) -> bool:
129	launcher = self._stable_launcher()	139	launcher = self._stable_launcher()
	@@ -150,6 +160,8 @@ class BoundFastLaunch:
150	self._install_negative(launcher)	160	self._install_negative(launcher)
151	return False	161	return False
152	self._clear_negative()	162	self._clear_negative()
		163	+ if self._call_slot is not None:
		164	+ self._call_slot[0] = FinalizedFastLaunch(self, self._direct)
153	return True	165	return True
154		166
155	def _fallback(	167	def _fallback(
	@@ -211,9 +223,8 @@ class BoundFastLaunch:
211	args,	223	args,
212	stream=stream,	224	stream=stream,
213	)	225	)
214	- canonical_args = self._canonical_args(args)
215	try:	226	try:
216	- direct(canonical_args, stream=stream)	227	+ direct(args, stream=stream)
217	except FastLaunchError as exc:	228	except FastLaunchError as exc:
218	if exc.backend_submitted:	229	if exc.backend_submitted:
219	raise	230	raise
	@@ -284,6 +295,57 @@ class BoundFastLaunch:
284	return result	295	return result
285		296
286		297
		298	+class FinalizedFastLaunch:
		299	+ """Steady-state entry installed directly into a generated call slot."""
		300	+
		301	+ __slots__ = ("bound", "direct")
		302	+
		303	+ def __init__(self, bound: BoundFastLaunch, direct: PlannedFastLaunch) -> None:
		304	+ self.bound = bound
		305	+ self.direct = direct
		306	+
		307	+ def __call__(
		308	+ self,
		309	+ *args: Any,
		310	+ stream: Any,
		311	+ benchmark_run: bool = False,
		312	+ **kwargs: Any,
		313	+ ) -> Any:
		314	+ if (
		315	+ benchmark_run
		316	+ or kwargs
		317	+ or autograd_profiler._is_profiler_enabled
		318	+ or getattr(self.bound.autotuner, "best_launcher", None)
		319	+ is not self.direct.launcher
		320	+ ):
		321	+ return self.bound(
		322	+ *args,
		323	+ stream=stream,
		324	+ benchmark_run=benchmark_run,
		325	+ **kwargs,
		326	+ )
		327	+ if _launcher_has_active_launch_hooks(self.direct.launcher):
		328	+ return self.bound._hook_fallback(
		329	+ self.direct.launcher,
		330	+ args,
		331	+ stream=stream,
		332	+ )
		333	+ try:
		334	+ self.direct(args, stream=stream)
		335	+ except FastLaunchError as exc:
		336	+ if exc.backend_submitted:
		337	+ raise
		338	+ if exc.stable:
		339	+ self.bound._install_negative(self.direct.launcher)
		340	+ return self.bound._fallback(
		341	+ args,
		342	+ stream=stream,
		343	+ benchmark_run=False,
		344	+ kwargs={},
		345	+ )
		346	+ return None
		347	+
		348	+
287	def bind_python_wrapper_kernel_fast(	349	def bind_python_wrapper_kernel_fast(
288	metadata: dict[str, Any],	350	metadata: dict[str, Any],
289	autotuner: Any,	351	autotuner: Any,
	@@ -293,11 +355,15 @@ def bind_python_wrapper_kernel_fast(
293	bound = (	355	bound = (
294	autotuner.run	356	autotuner.run
295	if _is_grouped_autotuner(autotuner)	357	if _is_grouped_autotuner(autotuner)
296	- else BoundFastLaunch(autotuner, metadata)	358	+ else BoundFastLaunch(autotuner, metadata, call_slot=call_slot)
297	)	359	)
298	if call_slot is not None:	360	if call_slot is not None:
299	call_slot[0] = bound	361	call_slot[0] = bound
300	return bound	362	return bound
301		363
302		364
303	-__all__ = ["BoundFastLaunch", "bind_python_wrapper_kernel_fast"]	365	+__all__ = [
		366	+ "BoundFastLaunch",
		367	+ "FinalizedFastLaunch",
		368	+ "bind_python_wrapper_kernel_fast",
		369	+]

Mtorch_npu/_inductor/experimental/python_wrapper_fast_launch/csrc/bindings.cpp+130-34

	@@ -46,6 +46,7 @@ struct FastLaunchPlan {
46	void* kernelStub = nullptr;	46	void* kernelStub = nullptr;
47	std::vector<FastLaunchArgKind> argKinds;	47	std::vector<FastLaunchArgKind> argKinds;
48	std::vector<FastLaunchArgLayout> argLayouts;	48	std::vector<FastLaunchArgLayout> argLayouts;
		49	+ size_t runtimeArgCount = 0;
49	size_t fftsOffset = 0;	50	size_t fftsOffset = 0;
50	size_t gridOffsets[3] = {0, 0, 0};	51	size_t gridOffsets[3] = {0, 0, 0};
51	size_t packedArgsSize = 0;	52	size_t packedArgsSize = 0;
	@@ -54,6 +55,9 @@ struct FastLaunchPlan {
54	bool isPureSimt = false;	55	bool isPureSimt = false;
55	bool targetSupportFfts = false;	56	bool targetSupportFfts = false;
56	void* fftsAddress = nullptr;	57	void* fftsAddress = nullptr;
		58	+ std::vector<uint8_t> packedArgsTemplate;
		59	+ uint32_t staticBlockNum = 0;
		60	+ bool hasStaticGrid = false;
57	};	61	};
58		62
59	size_t AlignOffset(size_t offset, size_t alignment) {	63	size_t AlignOffset(size_t offset, size_t alignment) {
	@@ -277,23 +281,7 @@ struct PackedLaunch {
277	rtStream_t stream = nullptr;	281	rtStream_t stream = nullptr;
278	};	282	};
279		283
280	-PackedLaunch PackLaunch(	284	+uint32_t ValidateGrid(uint32_t grid0, uint32_t grid1, uint32_t grid2) {
281	- const FastLaunchPlan& plan,
282	- uint64_t streamValue,
283	- uint32_t grid0,
284	- uint32_t grid1,
285	- uint32_t grid2,
286	- const py::sequence& args) {
287	- size_t argCount = static_cast<size_t>(py::len(args));
288	- TORCH_CHECK(
289	- argCount == plan.argKinds.size(),
290	- "fast launch args and arg_kinds size mismatch: ",
291	- argCount,
292	- " vs ",
293	- plan.argKinds.size());
294	- rtStream_t stream = reinterpret_cast<rtStream_t>(streamValue);
295	- TORCH_CHECK(stream != nullptr, "fast launch stream pointer is null");
296	-
297	const uint32_t grid[3] = {grid0, grid1, grid2};	285	const uint32_t grid[3] = {grid0, grid1, grid2};
298	uint64_t blockNum = 1;	286	uint64_t blockNum = 1;
299	for (size_t index = 0; index < 3; ++index) {	287	for (size_t index = 0; index < 3; ++index) {
	@@ -307,30 +295,80 @@ PackedLaunch PackLaunch(
307	blockNum <= std::numeric_limits<uint16_t>::max(),	295	blockNum <= std::numeric_limits<uint16_t>::max(),
308	"fast launch grid product exceeds uint16 max");	296	"fast launch grid product exceeds uint16 max");
309	}	297	}
		298	+ return static_cast<uint32_t>(blockNum);
		299	+}
310		300
311	- PackedLaunch packed;	301	+void WriteGrid(
312	- packed.blockNum = static_cast<uint32_t>(blockNum);	302	+ const FastLaunchPlan& plan,
313	- packed.stream = stream;	303	+ std::vector<uint8_t>& args,
314	- TORCH_INTERNAL_ASSERT(plan.argLayouts.size() == argCount);	304	+ uint32_t grid0,
315	- packed.args.resize(plan.packedArgsSize, 0);	305	+ uint32_t grid1,
316	- if (plan.targetSupportFfts) {	306	+ uint32_t grid2) {
317	- WritePointerAt(packed.args, plan.fftsOffset, plan.fftsAddress);	307	+ const int32_t signedGrid[3] = {
318	- }
319	- for (size_t index = 0; index < argCount; ++index) {
320	- WriteArgAt(packed.args, args[index], plan.argLayouts[index]);
321	- }
322	- int32_t signedGrid[3] = {
323	static_cast<int32_t>(grid0),	308	static_cast<int32_t>(grid0),
324	static_cast<int32_t>(grid1),	309	static_cast<int32_t>(grid1),
325	static_cast<int32_t>(grid2),	310	static_cast<int32_t>(grid2),
326	};	311	};
327	for (size_t index = 0; index < 3; ++index) {	312	for (size_t index = 0; index < 3; ++index) {
328	WriteBytesAt(	313	WriteBytesAt(
329	- packed.args,	314	+ args,
330	plan.gridOffsets[index],	315	plan.gridOffsets[index],
331	&signedGrid[index],	316	&signedGrid[index],
332	sizeof(signedGrid[index]));	317	sizeof(signedGrid[index]));
333	}	318	}
		319	+}
		320	+
		321	+PackedLaunch PackLaunch(
		322	+ const FastLaunchPlan& plan,
		323	+ uint64_t streamValue,
		324	+ uint32_t grid0,
		325	+ uint32_t grid1,
		326	+ uint32_t grid2,
		327	+ const py::sequence& args) {
		328	+ size_t argCount = static_cast<size_t>(py::len(args));
		329	+ TORCH_CHECK(
		330	+ argCount == plan.runtimeArgCount,
		331	+ "fast launch args and arg_kinds size mismatch: ",
		332	+ argCount,
		333	+ " vs ",
		334	+ plan.runtimeArgCount);
		335	+ rtStream_t stream = reinterpret_cast<rtStream_t>(streamValue);
		336	+ TORCH_CHECK(stream != nullptr, "fast launch stream pointer is null");
		337	+
		338	+ PackedLaunch packed;
		339	+ packed.blockNum = ValidateGrid(grid0, grid1, grid2);
		340	+ packed.stream = stream;
		341	+ TORCH_INTERNAL_ASSERT(plan.argLayouts.size() >= argCount);
		342	+ packed.args = plan.packedArgsTemplate;
		343	+ for (size_t index = 0; index < argCount; ++index) {
		344	+ WriteArgAt(packed.args, args[index], plan.argLayouts[index]);
		345	+ }
		346	+ WriteGrid(plan, packed.args, grid0, grid1, grid2);
		347	+ return packed;
		348	+}
		349	+
		350	+PackedLaunch PackStaticLaunch(
		351	+ const FastLaunchPlan& plan,
		352	+ uint64_t streamValue,
		353	+ const py::sequence& args) {
		354	+ TORCH_CHECK(plan.hasStaticGrid, "fast launch plan has no static grid");
		355	+ size_t argCount = static_cast<size_t>(py::len(args));
		356	+ TORCH_CHECK(
		357	+ argCount == plan.runtimeArgCount,
		358	+ "fast launch args and arg_kinds size mismatch: ",
		359	+ argCount,
		360	+ " vs ",
		361	+ plan.runtimeArgCount);
		362	+ rtStream_t stream = reinterpret_cast<rtStream_t>(streamValue);
		363	+ TORCH_CHECK(stream != nullptr, "fast launch stream pointer is null");
		364	+
		365	+ PackedLaunch packed;
		366	+ packed.blockNum = plan.staticBlockNum;
		367	+ packed.stream = stream;
		368	+ packed.args = plan.packedArgsTemplate;
		369	+ for (size_t index = 0; index < argCount; ++index) {
		370	+ WriteArgAt(packed.args, args[index], plan.argLayouts[index]);
		371	+ }
334	return packed;	372	return packed;
335	}	373	}
336		374
	@@ -364,8 +402,9 @@ void SubmitLaunch(const FastLaunchPlan& plan, PackedLaunch packed) {
364	return static_cast<int>(result);	402	return static_cast<int>(result);
365	};	403	};
366		404
367	- at_npu::native::OpCommand command;	405	+ // The launch callable is fully prepared. Reuse the existing OpAPI V2 queue
368	- command.Name(plan.kernelName).SetCustomHandler(std::move(launchCall)).Run();	406	+ // entry instead of rebuilding a generic zero-I/O OpCommand for every hit.
		407	+ at_npu::native::OpCommand::RunOpApiV2(plan.kernelName, launchCall);
369	}	408	}
370		409
371	std::shared_ptr<FastLaunchPlan> MakeFastLaunchPlan(	410	std::shared_ptr<FastLaunchPlan> MakeFastLaunchPlan(
	@@ -375,7 +414,10 @@ std::shared_ptr<FastLaunchPlan> MakeFastLaunchPlan(
375	bool enableSimt,	414	bool enableSimt,
376	uint64_t sharedMemDynamicSize,	415	uint64_t sharedMemDynamicSize,
377	bool isPureSimt,	416	bool isPureSimt,
378	- bool targetSupportFfts) {	417	+ bool targetSupportFfts,
		418	+ size_t runtimeArgCount,
		419	+ const py::sequence& fixedArgs,
		420	+ const std::vector<uint32_t>& staticGrid) {
379	TORCH_CHECK(	421	TORCH_CHECK(
380	sharedMemDynamicSize <= std::numeric_limits<uint32_t>::max(),	422	sharedMemDynamicSize <= std::numeric_limits<uint32_t>::max(),
381	"shared_mem_dynamic_size exceeds uint32 max");	423	"shared_mem_dynamic_size exceeds uint32 max");
	@@ -386,6 +428,17 @@ std::shared_ptr<FastLaunchPlan> MakeFastLaunchPlan(
386	plan->kernelStubOwner = kernelStub;	428	plan->kernelStubOwner = kernelStub;
387	plan->kernelStub = ExtractPointer(kernelStub, "kernel_stub");	429	plan->kernelStub = ExtractPointer(kernelStub, "kernel_stub");
388	plan->argKinds = ParseArgKinds(argKinds);	430	plan->argKinds = ParseArgKinds(argKinds);
		431	+ if (runtimeArgCount == std::numeric_limits<size_t>::max()) {
		432	+ runtimeArgCount = plan->argKinds.size();
		433	+ }
		434	+ TORCH_CHECK(
		435	+ runtimeArgCount <= plan->argKinds.size(),
		436	+ "runtime arg count exceeds fast launch ABI size");
		437	+ TORCH_CHECK(
		438	+ static_cast<size_t>(py::len(fixedArgs)) ==
		439	+ plan->argKinds.size() - runtimeArgCount,
		440	+ "fixed fast launch args do not complete the ABI");
		441	+ plan->runtimeArgCount = runtimeArgCount;
389	plan->enableSimt = enableSimt;	442	plan->enableSimt = enableSimt;
390	plan->sharedMemDynamicSize = sharedMemDynamicSize;	443	plan->sharedMemDynamicSize = sharedMemDynamicSize;
391	plan->isPureSimt = isPureSimt;	444	plan->isPureSimt = isPureSimt;
	@@ -404,6 +457,32 @@ std::shared_ptr<FastLaunchPlan> MakeFastLaunchPlan(
404	plan->fftsAddress = reinterpret_cast<void*>(fftsAddress);	457	plan->fftsAddress = reinterpret_cast<void*>(fftsAddress);
405	}	458	}
406	BuildPackedLayout(*plan);	459	BuildPackedLayout(*plan);
		460	+ plan->packedArgsTemplate.resize(plan->packedArgsSize, 0);
		461	+ if (plan->targetSupportFfts) {
		462	+ WritePointerAt(
		463	+ plan->packedArgsTemplate, plan->fftsOffset, plan->fftsAddress);
		464	+ }
		465	+ for (size_t index = runtimeArgCount; index < plan->argKinds.size(); ++index) {
		466	+ TORCH_CHECK(
		467	+ plan->argKinds[index] != FastLaunchArgKind::Tensor,
		468	+ "fixed tensor fast launch arguments are unsupported");
		469	+ WriteArgAt(
		470	+ plan->packedArgsTemplate,
		471	+ fixedArgs[index - runtimeArgCount],
		472	+ plan->argLayouts[index]);
		473	+ }
		474	+ if (!staticGrid.empty()) {
		475	+ TORCH_CHECK(staticGrid.size() == 3, "static fast launch grid must have rank 3");
		476	+ plan->staticBlockNum =
		477	+ ValidateGrid(staticGrid[0], staticGrid[1], staticGrid[2]);
		478	+ WriteGrid(
		479	+ *plan,
		480	+ plan->packedArgsTemplate,
		481	+ staticGrid[0],
		482	+ staticGrid[1],
		483	+ staticGrid[2]);
		484	+ plan->hasStaticGrid = true;
		485	+ }
407	return plan;	486	return plan;
408	}	487	}
409		488
	@@ -418,6 +497,14 @@ void FastLaunchWithPlan(
418	SubmitLaunch(plan, PackLaunch(plan, stream, grid0, grid1, grid2, args));	497	SubmitLaunch(plan, PackLaunch(plan, stream, grid0, grid1, grid2, args));
419	}	498	}
420		499
		500	+void FastLaunchStaticWithPlan(
		501	+ const std::shared_ptr<FastLaunchPlan>& plan,
		502	+ uint64_t stream,
		503	+ const py::sequence& args) {
		504	+ TORCH_CHECK(plan != nullptr, "fast launch plan is null");
		505	+ SubmitLaunch(plan, PackStaticLaunch(plan, stream, args));
		506	+}
		507	+
421	} // namespace	508	} // namespace
422		509
423	void RegisterNPUFastLaunchBindings(PyObject* module) {	510	void RegisterNPUFastLaunchBindings(PyObject* module) {
	@@ -433,7 +520,10 @@ void RegisterNPUFastLaunchBindings(PyObject* module) {
433	py::arg("enable_simt") = false,	520	py::arg("enable_simt") = false,
434	py::arg("shared_mem_dynamic_size") = 0,	521	py::arg("shared_mem_dynamic_size") = 0,
435	py::arg("is_pure_simt") = false,	522	py::arg("is_pure_simt") = false,
436	- py::arg("target_support_ffts") = false);	523	+ py::arg("target_support_ffts") = false,
		524	+ py::arg("runtime_arg_count") = std::numeric_limits<size_t>::max(),
		525	+ py::arg("fixed_args") = py::tuple(),
		526	+ py::arg("static_grid") = std::vector<uint32_t>());
437	m.def(	527	m.def(
438	"_npu_inductor_fast_launch_with_plan",	528	"_npu_inductor_fast_launch_with_plan",
439	&FastLaunchWithPlan,	529	&FastLaunchWithPlan,
	@@ -443,6 +533,12 @@ void RegisterNPUFastLaunchBindings(PyObject* module) {
443	py::arg("grid_1"),	533	py::arg("grid_1"),
444	py::arg("grid_2"),	534	py::arg("grid_2"),
445	py::arg("args"));	535	py::arg("args"));
		536	+ m.def(
		537	+ "_npu_inductor_fast_launch_static_with_plan",
		538	+ &FastLaunchStaticWithPlan,
		539	+ py::arg("plan"),
		540	+ py::arg("stream"),
		541	+ py::arg("args"));
446	}	542	}
447		543
448	#endif	544	#endif